diff --git a/.gitattributes b/.gitattributes index 604ffa217ded3a6b645774342eb32d3e53b12936..94f9f694539db0808c2714cae37a6ce19edb750c 100644 --- a/.gitattributes +++ b/.gitattributes @@ -232,3 +232,5 @@ workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cpu workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda.so filter=lfs diff=lfs merge=lfs -text workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_cuda_linalg.so filter=lfs diff=lfs merge=lfs -text workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/lib/libtorch_python.so filter=lfs diff=lfs merge=lfs -text +workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/_yaml.cpython-311-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text +workspace/outputs/external_vla_export_maniskill_full_no_images/train.jsonl filter=lfs diff=lfs merge=lfs -text diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/INSTALLER b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/METADATA b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..446cb82e5521ece54e8d24efc9bf14f5f4530276 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/METADATA @@ -0,0 +1,48 @@ +Metadata-Version: 2.2 +Name: torch +Version: 2.12.1+computecanada +Summary: Tensors and Dynamic neural networks in Python with strong GPU acceleration +Description-Content-Type: text/markdown +Keywords: pytorch,machine learning +Classifier: Development Status :: 5 - Production/Stable +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Education +Classifier: Intended Audience :: Science/Research +Classifier: Topic :: Scientific/Engineering +Classifier: Topic :: Scientific/Engineering :: Mathematics +Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence +Classifier: Topic :: Software Development +Classifier: Topic :: Software Development :: Libraries +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Classifier: Programming Language :: C++ +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Author-email: PyTorch Team +Project-URL: Homepage, https://pytorch.org +Project-URL: Repository, https://github.com/pytorch/pytorch +Project-URL: Documentation, https://pytorch.org/docs +Project-URL: Issue Tracker, https://github.com/pytorch/pytorch/issues +Project-URL: Forum, https://discuss.pytorch.org +Requires-Python: >=3.10 +Requires-Dist: filelock +Requires-Dist: typing-extensions>=4.10.0 +Requires-Dist: setuptools<82 +Requires-Dist: sympy>=1.13.3 +Requires-Dist: networkx>=2.5.1 +Requires-Dist: jinja2 +Requires-Dist: fsspec>=0.8.5 +Requires-Dist: optree>=0.13.0; extra == "optree" +Requires-Dist: opt-einsum>=3.3; extra == "opt-einsum" +Requires-Dist: pyyaml; extra == "pyyaml" +Provides-Extra: optree +Provides-Extra: opt-einsum +Provides-Extra: pyyaml +License-File: LICENSE +License-File: NOTICE +Dynamic: license-file + +BSD-3-Clause \ No newline at end of file diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/RECORD b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..7ae1ddfabca66638f8a7f41de29fb843923d1b28 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/RECORD @@ -0,0 +1,14122 @@ +../../../bin/torchrun,sha256=xtDxWRTFNFAQ6Nyvi66oy9VkeAQx3X12VXDTB-EPyE0,271 +functorch/__init__.py,sha256=NAwGN21zq-tccaF-ROtv-VWFoPdb7y9iuAt6Hy6QCtc,1037 +functorch/__pycache__/__init__.cpython-311.pyc,, +functorch/_src/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +functorch/_src/__pycache__/__init__.cpython-311.pyc,, +functorch/_src/aot_autograd/__init__.py,sha256=SGo7gh6XGYcOxTGf5g-R8Y9AF95ICcJbQwQD6rFyrpQ,291 +functorch/_src/aot_autograd/__pycache__/__init__.cpython-311.pyc,, +functorch/_src/eager_transforms/__init__.py,sha256=kX_52fDvSX9YX9OAwo5bjvJtrxyjEBUJ1PueW8xgsuw,291 +functorch/_src/eager_transforms/__pycache__/__init__.cpython-311.pyc,, +functorch/_src/make_functional/__init__.py,sha256=b3y8s3KhtCqFB8lM4Pi48AwuztUt7NBK-VISZNJYYjw,235 +functorch/_src/make_functional/__pycache__/__init__.cpython-311.pyc,, +functorch/_src/vmap/__init__.py,sha256=k8r2Esz6tB5D7U_UA0_BCDaWoOmn8JNVrRqK7nG7_fM,467 +functorch/_src/vmap/__pycache__/__init__.cpython-311.pyc,, +functorch/compile/__init__.py,sha256=fZnNG56VBLfKlXMqX5Rj3tORQYLyxbyoA0rEoEBt3KM,756 +functorch/compile/__pycache__/__init__.cpython-311.pyc,, +functorch/dim/__init__.py,sha256=uQEKMitItfVkGFolejbBM4RhfVUe1pbCiFbC5B9bdgY,53600 +functorch/dim/__pycache__/__init__.cpython-311.pyc,, +functorch/dim/__pycache__/_dim_entry.cpython-311.pyc,, +functorch/dim/__pycache__/_enable_all_layers.cpython-311.pyc,, +functorch/dim/__pycache__/_getsetitem.cpython-311.pyc,, +functorch/dim/__pycache__/_order.cpython-311.pyc,, +functorch/dim/__pycache__/_py_inst_decoder.cpython-311.pyc,, +functorch/dim/__pycache__/_tensor_info.cpython-311.pyc,, +functorch/dim/__pycache__/_wrap.cpython-311.pyc,, +functorch/dim/__pycache__/magic_trace.cpython-311.pyc,, +functorch/dim/__pycache__/op_properties.cpython-311.pyc,, +functorch/dim/__pycache__/wrap_type.cpython-311.pyc,, +functorch/dim/_dim_entry.py,sha256=cceDVoHZemWMtW72AoGJCpZBp_lyOSDCdABJ29RLbmQ,3884 +functorch/dim/_enable_all_layers.py,sha256=xKbVnuuh2MhC7GjZLDGpk_DD-wyMT-uDBBgVU-KSuCs,5220 +functorch/dim/_getsetitem.py,sha256=8dL1J7e2wvUx7xw9r-Ffunn075fjK82CX94RHGtxbgY,19039 +functorch/dim/_order.py,sha256=0n1ew1EHEQZnpxiwC0_k7vkno5-wLiOfRkhnKhYut04,7011 +functorch/dim/_py_inst_decoder.py,sha256=OUIOiKA0zZ2C3b8iCRgD9woShn1VkHNrJRVXneJa40E,2091 +functorch/dim/_tensor_info.py,sha256=eXHojr8bOn-1N5vkWVDIlmpp7o9iwCexXxAP4vXL4A4,2098 +functorch/dim/_wrap.py,sha256=NaV0AQZx0yDRw_37U4jfSS_B_ko9XHUqdocz7XHPu2Q,8385 +functorch/dim/magic_trace.py,sha256=spyJfvOiY6sUnfiDwLGE2Pejv5DyFSnZBmbio0nh1iA,1547 +functorch/dim/op_properties.py,sha256=GpW8Ylgq1YlMilQh6cgNQ66tAiN84WRyEjn_VkbJYFQ,6687 +functorch/dim/wrap_type.py,sha256=haCYf-nAHCw236O1jh8GHjFgmzDNKrltdNvtVG7y7pU,2090 +functorch/einops/__init__.py,sha256=qNdomhBnsKNuNNlGsbqqipT9wYQkcxMuEQJKq4zhry8,59 +functorch/einops/__pycache__/__init__.cpython-311.pyc,, +functorch/einops/__pycache__/_parsing.cpython-311.pyc,, +functorch/einops/__pycache__/rearrange.cpython-311.pyc,, +functorch/einops/_parsing.py,sha256=yjwK7zoSov-VRwJzYZsGDQ-rBFb5RCApZ6Dln9Y2NIE,12208 +functorch/einops/rearrange.py,sha256=_8jfWFNnN16v6C3qeDX8LItay8Iqp7UvfJs3DAjbjis,8123 +functorch/experimental/__init__.py,sha256=qgO0Y4xqrtxucMbH-Q__wtJGvzRk9ckQIqo7E94-ROA,274 +functorch/experimental/__pycache__/__init__.cpython-311.pyc,, +functorch/experimental/__pycache__/control_flow.cpython-311.pyc,, +functorch/experimental/__pycache__/ops.cpython-311.pyc,, +functorch/experimental/control_flow.py,sha256=CyLqEoRLJT5fSnwcTrmw_3LQk9SD54fJElQ4stSD_64,144 +functorch/experimental/ops.py,sha256=kDGcckdoYwOg9fS4JqvZnsIt8Ss9O6RGeasJSyi0OUY,57 +torch-2.12.1+computecanada.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +torch-2.12.1+computecanada.dist-info/METADATA,sha256=Tc7PPsbXv1ASN-PnwNPLetj0Z5Scn3R42VTq0cXoZ8s,1959 +torch-2.12.1+computecanada.dist-info/RECORD,, +torch-2.12.1+computecanada.dist-info/WHEEL,sha256=dCAA691vV7BW_aZWazBJJ5ldVUtr6Pzmdb_zOoyLvl4,99 +torch-2.12.1+computecanada.dist-info/entry_points.txt,sha256=omsXNYssMtT87QYPUINGtY0zX0MwdqjL08dLZcvRnHk,148 +torch-2.12.1+computecanada.dist-info/licenses/LICENSE,sha256=vQGP7viCXogYHITrfjqk6vuPCKINn9bvlIVpYQxKPkM,3464 +torch-2.12.1+computecanada.dist-info/licenses/NOTICE,sha256=wsx78MrsdlLCtGCopHC-oWd_JB5KuOQx3zTPF_Wp_sA,23632 +torch-2.12.1+computecanada.dist-info/top_level.txt,sha256=Hv8AT5yFnStBq__ugW6DV18rPKJVgZB4vVyHcFKQmVw,15 +torch/_C.cpython-311-x86_64-linux-gnu.so,sha256=OCW85N0gMyML7aXWizvi48zX371th6cvJigkasVaJb4,22104 +torch/_C/_VariableFunctions.pyi,sha256=_5kG2AX7ZlFlqQGjzj3xy7AQUTsYKVZ02E9MbymNTaA,1209737 +torch/_C/__init__.pyi,sha256=l4Mq3oHa0ssdoAsNfDTdMI3o6SoVJed2f7-eVeUXwzs,419415 +torch/_C/_acc/__init__.pyi,sha256=YrGID2b-0nfrv2Kqozt-brcNB3Prd_a1PPaa2TFWy5A,568 +torch/_C/_aoti.pyi,sha256=gP8rTuk9mN1ZLsy43cpaitglNYOPbjfGGSW3F1e0wQU,6015 +torch/_C/_autograd.pyi,sha256=fEMVOqEgDDRZ_T6kDrgIB9CkdPqUwVc6mBFw-jJuSLw,5699 +torch/_C/_cpu.pyi,sha256=zLjxjpuOW1VCoe99C3I-HOf7b6FzfARkPO-JktJOCcU,174 +torch/_C/_cudnn.pyi,sha256=m4nyLKmUJTRrcb-e94ghGbQ7D_Y1E7eqzueyH6QeANs,317 +torch/_C/_cusparselt.pyi,sha256=kD30gidilch4jKuIr7-ItTvDuEtt6L7vxL1l56sb_mw,32 +torch/_C/_distributed.pyi,sha256=qXyYUT8V9MPMXhx38eU0YD6CvaqS4CLvWMrEmuDqWj4,598 +torch/_C/_distributed_autograd.pyi,sha256=brb51Q-9UlMWCs4Aoy7ZLHH4hwhYq0acXm_hgfz_5HA,899 +torch/_C/_distributed_c10d.pyi,sha256=Kzxn7RKcqPlqOoRTbS2ReoLMLcET25APwLTIQQ8jYoM,27012 +torch/_C/_distributed_rpc.pyi,sha256=0mlTwOqrF608gU4oUnIPc25nQ4zeJ6LWycHcuyfB3e8,6080 +torch/_C/_distributed_rpc_testing.pyi,sha256=2cUz7bIO5y9fKUzfafx6Uh84QNO10gtP3uJXAXmE8Dk,1017 +torch/_C/_dynamo/__init__.pyi,sha256=Me61v9GA3oBVWBDhoEO8SenUDFzmX_Mdc992qH4Dros,1592 +torch/_C/_dynamo/compiled_autograd.pyi,sha256=zbX2AKKVEX1yzpOsa2xIcoDBFM7YBpKSNPJjEsCOTZU,522 +torch/_C/_dynamo/eval_frame.pyi,sha256=2BZrPmlZQMMb0uA6oJEj_V9Jos6wHmOwZRDYBwD_4RM,3371 +torch/_C/_dynamo/guards.pyi,sha256=doDOQZ9sPDGkJe8tHXzbtmkt5djZlSbo1zRfYU4wJG4,14782 +torch/_C/_export/__init__.pyi,sha256=CUcMBKNDN_31e66pbRuZ9O8ilVR4-RSKAeMFRfRIxS8,257 +torch/_C/_export/pt2_archive_constants.pyi,sha256=8oIuaAUlGzny2thjeDhwuFKAk3WXi1yxnRkIijnDgN8,788 +torch/_C/_functionalization.pyi,sha256=z-v6deG3dIJsP-wV8sIZ022eu18T-yefDV20qN2wAA0,566 +torch/_C/_functions.pyi,sha256=0aapmQmjpa_pFZT_BWhddVB8jZh0DVLDKzXlWqYfz0k,608 +torch/_C/_functorch.pyi,sha256=u6CKxd5xb4JbK71Z8vaRlJp2yD73S0tzApCTgWZWpg8,3835 +torch/_C/_instruction_counter.pyi,sha256=xjwkiHBgXbvN8Iw7UNVEvE2YDJKFUU1_EoeTS4prCV8,109 +torch/_C/_itt.pyi,sha256=6fhhHGYgreXbGka-VtqX9FjjPaSznfOmDHPVC171DII,169 +torch/_C/_jit_tree_views.pyi,sha256=StShF20JIzi3BtE05CREHBAXWXGbmmMo7pWKR2EHuRw,5498 +torch/_C/_lazy.pyi,sha256=h8_ZY1JSJTiGebWekqfMvfFGZdt_g8dqFq90TG0-7hI,1015 +torch/_C/_lazy_ts_backend.pyi,sha256=YfYcEssTgLqOHQqUDYg5ZgbdB_D52mG19VETBmh7yuc,326 +torch/_C/_monitor.pyi,sha256=ND1PrgzrSUMGCBd50Hyz8M-mz2uJRNyFU_RGxSgx3kc,1447 +torch/_C/_nn.pyi,sha256=86kZpQcxqoHBkhynAgBuPUk0ZKQc5xJ0sZtID1N_ah0,9741 +torch/_C/_nvtx.pyi,sha256=MzekiGde65H1P_cJ_zER9kjXVRDW12eNtXUeItZG3WA,380 +torch/_C/_onnx.pyi,sha256=7xxwokJOBBQxnbox3lGLgMo9RkVGVgtfPBtsKcaGaTE,710 +torch/_C/_profiler.pyi,sha256=vwO83UhNl0kEbtahhkqU9JnyRvkNvzMT_wi3gZ5bINg,6301 +torch/_C/_verbose.pyi,sha256=vMdQYMqABMqBFxYykp8_VD0QBaubolczBmCEv-UcA00,134 +torch/_C_flatbuffer/__init__.pyi,sha256=SdmD2nh01AdVYvwMJeo32D5BJW2Ar2pvjQTup0uBIko,544 +torch/_VF.py,sha256=XIfh8pIjzvG0ySWcOScz4E067DFoT-8TkcmAG5oYn94,664 +torch/_VF.pyi,sha256=_5kG2AX7ZlFlqQGjzj3xy7AQUTsYKVZ02E9MbymNTaA,1209737 +torch/__config__.py,sha256=VCuBweOjH807O10Xrumlb_BrtWNTQMYIGjb3aw6T2po,574 +torch/__future__.py,sha256=yk9l_KWsfVIzUBx9cGr-OdtWmb-pI8ZhcROAm3a_FQw,3185 +torch/__init__.py,sha256=2d__S3XUbkx1VyIAo0ZrcCMdBbAxjjisG9EheJFl-0k,109312 +torch/__pycache__/_VF.cpython-311.pyc,, +torch/__pycache__/__config__.cpython-311.pyc,, +torch/__pycache__/__future__.cpython-311.pyc,, +torch/__pycache__/__init__.cpython-311.pyc,, +torch/__pycache__/_appdirs.cpython-311.pyc,, +torch/__pycache__/_classes.cpython-311.pyc,, +torch/__pycache__/_compile.cpython-311.pyc,, +torch/__pycache__/_custom_ops.cpython-311.pyc,, +torch/__pycache__/_environment.cpython-311.pyc,, +torch/__pycache__/_guards.cpython-311.pyc,, +torch/__pycache__/_jit_internal.cpython-311.pyc,, +torch/__pycache__/_linalg_utils.cpython-311.pyc,, +torch/__pycache__/_lobpcg.cpython-311.pyc,, +torch/__pycache__/_lowrank.cpython-311.pyc,, +torch/__pycache__/_meta_registrations.cpython-311.pyc,, +torch/__pycache__/_namedtensor_internals.cpython-311.pyc,, +torch/__pycache__/_opaque_base.cpython-311.pyc,, +torch/__pycache__/_ops.cpython-311.pyc,, +torch/__pycache__/_python_dispatcher.cpython-311.pyc,, +torch/__pycache__/_size_docs.cpython-311.pyc,, +torch/__pycache__/_sources.cpython-311.pyc,, +torch/__pycache__/_storage_docs.cpython-311.pyc,, +torch/__pycache__/_streambase.cpython-311.pyc,, +torch/__pycache__/_tensor.cpython-311.pyc,, +torch/__pycache__/_tensor_docs.cpython-311.pyc,, +torch/__pycache__/_tensor_str.cpython-311.pyc,, +torch/__pycache__/_thread_safe_fork.cpython-311.pyc,, +torch/__pycache__/_torch_docs.cpython-311.pyc,, +torch/__pycache__/_utils.cpython-311.pyc,, +torch/__pycache__/_utils_internal.cpython-311.pyc,, +torch/__pycache__/_vmap_internals.cpython-311.pyc,, +torch/__pycache__/_weights_only_unpickler.cpython-311.pyc,, +torch/__pycache__/functional.cpython-311.pyc,, +torch/__pycache__/hub.cpython-311.pyc,, +torch/__pycache__/library.cpython-311.pyc,, +torch/__pycache__/overrides.cpython-311.pyc,, +torch/__pycache__/quasirandom.cpython-311.pyc,, +torch/__pycache__/random.cpython-311.pyc,, +torch/__pycache__/return_types.cpython-311.pyc,, +torch/__pycache__/serialization.cpython-311.pyc,, +torch/__pycache__/storage.cpython-311.pyc,, +torch/__pycache__/torch_version.cpython-311.pyc,, +torch/__pycache__/types.cpython-311.pyc,, +torch/__pycache__/version.cpython-311.pyc,, +torch/_appdirs.py,sha256=oTDrdoBwt0m-Vb27B8Bu3rfMYN6l7QBHeJepRqO_kL8,26198 +torch/_awaits/__init__.py,sha256=nAHROtkm0lmScDGnnEY5qGqK7W0ylL4NCutPDu3_lF4,1687 +torch/_awaits/__pycache__/__init__.cpython-311.pyc,, +torch/_classes.py,sha256=DbPcMC4o_QhQZeoog84z-kr0Brn9FDnypDHwU-mnZLk,1786 +torch/_compile.py,sha256=ssdSFpLJtFRF2DcQ3gPhYgV40s7ypu6T7Ao_rj3B0nc,1994 +torch/_custom_op/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_custom_op/__pycache__/__init__.cpython-311.pyc,, +torch/_custom_op/__pycache__/autograd.cpython-311.pyc,, +torch/_custom_op/__pycache__/impl.cpython-311.pyc,, +torch/_custom_op/autograd.py,sha256=ZtdjgZGheZ05fEyoB2Ue1f7nG3bikm5OpuMqMMUU1t0,12453 +torch/_custom_op/impl.py,sha256=3KAT0KL9fu-jpi99oL4KzdW-aWwbw-hLf1Wj784F6Zg,27419 +torch/_custom_ops.py,sha256=DyTzqyo_fZ582qmMiyfoLdT9-cAnor776dgqb_GhPIo,12824 +torch/_decomp/__init__.py,sha256=3eDHJ6rhyvJj-drGvr9kR6YQ60xM-PeqkaG7nBrscKM,19768 +torch/_decomp/__pycache__/__init__.cpython-311.pyc,, +torch/_decomp/__pycache__/decompositions.cpython-311.pyc,, +torch/_decomp/__pycache__/decompositions_for_jvp.cpython-311.pyc,, +torch/_decomp/__pycache__/decompositions_for_rng.cpython-311.pyc,, +torch/_decomp/decompositions.py,sha256=wfnK8Lz2AslMdYJkPuYafOPwXWittnKK8ITS1e-Sr_o,195777 +torch/_decomp/decompositions_for_jvp.py,sha256=C1DRcZOhDs87-gqfk26RULAMKx1_XJXzSG3fafEtAFw,11975 +torch/_decomp/decompositions_for_rng.py,sha256=4XNVPJhy02PLLaycLTwCrl_xQ5WRnqEvuvPF-wkURrk,9547 +torch/_dispatch/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_dispatch/__pycache__/__init__.cpython-311.pyc,, +torch/_dispatch/__pycache__/python.cpython-311.pyc,, +torch/_dispatch/python.py,sha256=tlR8LSP57-9WLKmyhISGRJmL-WewqtYe-krGz6x5cl4,7505 +torch/_dynamo/__init__.py,sha256=KSlJ5foy3SHr5Ue62WOwqTr28WdGZvnCRWIFyzFwNhE,7603 +torch/_dynamo/__pycache__/__init__.cpython-311.pyc,, +torch/_dynamo/__pycache__/_trace_wrapped_higher_order_op.cpython-311.pyc,, +torch/_dynamo/__pycache__/aot_compile.cpython-311.pyc,, +torch/_dynamo/__pycache__/aot_compile_types.cpython-311.pyc,, +torch/_dynamo/__pycache__/bytecode_analysis.cpython-311.pyc,, +torch/_dynamo/__pycache__/bytecode_debugger.cpython-311.pyc,, +torch/_dynamo/__pycache__/bytecode_transformation.cpython-311.pyc,, +torch/_dynamo/__pycache__/cache_size.cpython-311.pyc,, +torch/_dynamo/__pycache__/callback.cpython-311.pyc,, +torch/_dynamo/__pycache__/code_context.cpython-311.pyc,, +torch/_dynamo/__pycache__/codegen.cpython-311.pyc,, +torch/_dynamo/__pycache__/compiled_autograd.cpython-311.pyc,, +torch/_dynamo/__pycache__/comprehension_graph_break.cpython-311.pyc,, +torch/_dynamo/__pycache__/comptime.cpython-311.pyc,, +torch/_dynamo/__pycache__/config.cpython-311.pyc,, +torch/_dynamo/__pycache__/convert_frame.cpython-311.pyc,, +torch/_dynamo/__pycache__/create_parameter_op.cpython-311.pyc,, +torch/_dynamo/__pycache__/current_scope_id.cpython-311.pyc,, +torch/_dynamo/__pycache__/dce_extra_outputs.cpython-311.pyc,, +torch/_dynamo/__pycache__/debug_utils.cpython-311.pyc,, +torch/_dynamo/__pycache__/decorators.cpython-311.pyc,, +torch/_dynamo/__pycache__/device_interface.cpython-311.pyc,, +torch/_dynamo/__pycache__/distributed.cpython-311.pyc,, +torch/_dynamo/__pycache__/dynamo_profiler.cpython-311.pyc,, +torch/_dynamo/__pycache__/eval_frame.cpython-311.pyc,, +torch/_dynamo/__pycache__/exc.cpython-311.pyc,, +torch/_dynamo/__pycache__/external_utils.cpython-311.pyc,, +torch/_dynamo/__pycache__/funcname_cache.cpython-311.pyc,, +torch/_dynamo/__pycache__/functional_export.cpython-311.pyc,, +torch/_dynamo/__pycache__/graph_break_hints.cpython-311.pyc,, +torch/_dynamo/__pycache__/graph_bytecode_inputs.cpython-311.pyc,, +torch/_dynamo/__pycache__/graph_deduplication.cpython-311.pyc,, +torch/_dynamo/__pycache__/graph_id_filter.cpython-311.pyc,, +torch/_dynamo/__pycache__/graph_region_tracker.cpython-311.pyc,, +torch/_dynamo/__pycache__/graph_utils.cpython-311.pyc,, +torch/_dynamo/__pycache__/guards.cpython-311.pyc,, +torch/_dynamo/__pycache__/hooks.cpython-311.pyc,, +torch/_dynamo/__pycache__/logging.cpython-311.pyc,, +torch/_dynamo/__pycache__/metrics_context.cpython-311.pyc,, +torch/_dynamo/__pycache__/mutation_guard.cpython-311.pyc,, +torch/_dynamo/__pycache__/output_graph.cpython-311.pyc,, +torch/_dynamo/__pycache__/package.cpython-311.pyc,, +torch/_dynamo/__pycache__/pgo.cpython-311.pyc,, +torch/_dynamo/__pycache__/precompile_context.cpython-311.pyc,, +torch/_dynamo/__pycache__/profiler.cpython-311.pyc,, +torch/_dynamo/__pycache__/replay_record.cpython-311.pyc,, +torch/_dynamo/__pycache__/resume_execution.cpython-311.pyc,, +torch/_dynamo/__pycache__/side_effects.cpython-311.pyc,, +torch/_dynamo/__pycache__/source.cpython-311.pyc,, +torch/_dynamo/__pycache__/symbolic_convert.cpython-311.pyc,, +torch/_dynamo/__pycache__/tensor_version_op.cpython-311.pyc,, +torch/_dynamo/__pycache__/test_case.cpython-311.pyc,, +torch/_dynamo/__pycache__/test_dont_skip_tracing_functions.cpython-311.pyc,, +torch/_dynamo/__pycache__/test_minifier_common.cpython-311.pyc,, +torch/_dynamo/__pycache__/testing.cpython-311.pyc,, +torch/_dynamo/__pycache__/trace_rules.cpython-311.pyc,, +torch/_dynamo/__pycache__/types.cpython-311.pyc,, +torch/_dynamo/__pycache__/utils.cpython-311.pyc,, +torch/_dynamo/_trace_wrapped_higher_order_op.py,sha256=05E245ft09YIqDnPjdeVIna1tIA2NGEJDcV7vuSeFnY,9806 +torch/_dynamo/aot_compile.py,sha256=J3qgsI4aeyEtEVr25gcW0I97hu0qIpj4Jg6vz91SKtU,19671 +torch/_dynamo/aot_compile_types.py,sha256=drBq7dXHLqT2X3buLJeLcXhwlUP-ZffL_4WivNZveGE,8184 +torch/_dynamo/backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_dynamo/backends/__pycache__/__init__.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/common.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/cudagraphs.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/debugging.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/distributed.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/inductor.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/onnxrt.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/registry.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/tensorrt.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/torchxla.cpython-311.pyc,, +torch/_dynamo/backends/__pycache__/tvm.cpython-311.pyc,, +torch/_dynamo/backends/common.py,sha256=Ok3aUAUqICLj7Wd6mx8N7h0SwtVVVLrbka2jMAeu3Dk,6439 +torch/_dynamo/backends/cudagraphs.py,sha256=0z5DwRcwxCS9WPUguV8y1oa8fLFneDpjz1H7k5CQxtQ,10194 +torch/_dynamo/backends/debugging.py,sha256=9ZSUQ0EaOVC_DN610x4weuPrYn4lRkL9kDzGpAFIMP8,25394 +torch/_dynamo/backends/distributed.py,sha256=F_I-dcA1MyVySQhT4IxyTLJFPMWzjcF43r7m04Gmq8g,28490 +torch/_dynamo/backends/inductor.py,sha256=fSjxEorULdmCAxVxMoNcBTvTtzEpGJv-WHUTHvTE7QU,1104 +torch/_dynamo/backends/onnxrt.py,sha256=KMF_RHHh_JSV8W5pNXvtM_EiLcMdBWpfNF5iruUwyIg,1623 +torch/_dynamo/backends/registry.py,sha256=HV_eE6UlsTSgWYskx8sC34PZmHCUhU5dLtw2QsCvp8U,6506 +torch/_dynamo/backends/tensorrt.py,sha256=dxRCaQgNbw8Ly0SgGSOK7At2OANtvCe5mSHpAjvEqYM,383 +torch/_dynamo/backends/torchxla.py,sha256=o75DfoFrS6htTH1SfFx50Nu8xLWhT0hGvzoBnXfHpcU,1551 +torch/_dynamo/backends/tvm.py,sha256=BW7XT2H0wq7StneIE40TtgdH_sqeFa4hAHsq2Yp278I,7203 +torch/_dynamo/bytecode_analysis.py,sha256=keVsj6_DkVRJDeR9Xg1dwOmOJXDbshHdoO-Eeezzst8,9189 +torch/_dynamo/bytecode_debugger.py,sha256=fyPVm7sY0mbS8MrEEJLotYUvyVTylXpjlntK4UzZ-Tg,46481 +torch/_dynamo/bytecode_transformation.py,sha256=0U2PjnsGYOq7o2J2_GaeSi9SlFhRHODh8cvRYe8lPkA,68253 +torch/_dynamo/cache_size.py,sha256=VieF2JDo0IuHQA5Vs9A3QXJ0sBk2tuK9YKFEqIIGUGU,7963 +torch/_dynamo/callback.py,sha256=h_CPUwJLNeMpqVNBnYzsX9aV5Kk9_3gBzhTNh687Kog,5622 +torch/_dynamo/code_context.py,sha256=MXMoP-r7lN4jFHKLUDopun0ZJE98RC1VA0St0geTbuc,1863 +torch/_dynamo/codegen.py,sha256=rNor5x9ZhX8HRpfnXrOxN1iNz1LAqJl2rNdBVPSqooc,29906 +torch/_dynamo/compiled_autograd.py,sha256=cxM7aABQXm-pLlYtOyJWDISU_e5-4Tk97x3e9ae2EBs,64937 +torch/_dynamo/comprehension_graph_break.py,sha256=aEnaK3bKNHSiDceSdsMU98IqahDHmaDx3ymxUXWlf-g,24748 +torch/_dynamo/comptime.py,sha256=qYJNutWbkahnp7w95TW9OcbnGoU-jpw1ry68V5uZbL4,15354 +torch/_dynamo/config.py,sha256=itTc-drddHoO4FUGaWA_PYKegTluuXxn9gksG18DvXQ,40386 +torch/_dynamo/convert_frame.py,sha256=asbSRoqEcbEh_fWpeT94aWBfR44zda9OIBVcfqtFvCg,94367 +torch/_dynamo/create_parameter_op.py,sha256=hHU3LYMbZvgOTK9tVKMFXsgzf4yxHT50_B18xjDW350,2591 +torch/_dynamo/current_scope_id.py,sha256=CD2PnuBjVr2ns7_oNaPg6djM5wa-qNIaJYLFA4uThUo,1433 +torch/_dynamo/dce_extra_outputs.py,sha256=M10Drw4FDORgzG6NgRz_cyVe-6hjHyX7XnsJcxneTdk,7974 +torch/_dynamo/debug_utils.py,sha256=dKy7Wr6eEDvJdMEmGe-UzOIPBQiHYnO1j8HxjNxYGKg,33878 +torch/_dynamo/decorators.py,sha256=53ChXoi6R0-vz_yDIg-x6sI32EaGTFXZsE0Ixh85rMQ,65479 +torch/_dynamo/device_interface.py,sha256=8j3zk_P9iD4TDofk5jnzG4PpRx-RdlASNHwrDTV-9hM,23249 +torch/_dynamo/distributed.py,sha256=kCIL8k3R4XreKtF3LxKSeAQyC8rhoRKSul54Ztqvy-Y,1629 +torch/_dynamo/dynamo_profiler.py,sha256=1KH6CZ3fYA5IHUc3gNUnFGDtiIfn-N5ojP7fBzaO7JU,15231 +torch/_dynamo/eval_frame.py,sha256=5r-20_Ivh5Y1sddwLltKSPBR8kiV_30bxOzDiLO-EQA,105018 +torch/_dynamo/exc.py,sha256=RDDN-0HT7Bzfh9_uMXxQ8l-v-4pSy4qTWagU1AlQMc4,29791 +torch/_dynamo/external_utils.py,sha256=gfcDoEoscYCqysTTh6qYFxbOESGSKiGzmHVrWM2x5VA,9560 +torch/_dynamo/funcname_cache.py,sha256=N_JK8IaQU12eNjMJcuZxdAoOUDp-UYHtnoqnKM-kFkM,2517 +torch/_dynamo/functional_export.py,sha256=8hQ9FPpIOvdE_QUjRPfnXncbAH1uval_C5Bh7Q9Gkmg,39507 +torch/_dynamo/graph_break_hints.py,sha256=iqxvoIxyIrQEVZBAgTy4wwIt7cmn1wGFhTPxTKJjtjw,1657 +torch/_dynamo/graph_break_registry.json,sha256=W6qFGfVDWeOVbYW3QmAM3b78d6whZG5zielpa4l5us4,215929 +torch/_dynamo/graph_bytecode_inputs.py,sha256=nLYNmROoe9yjpJdOYL1qNgBUU9hSh55TYDeFIbmUFh0,3857 +torch/_dynamo/graph_deduplication.py,sha256=VhMHaHkeGAZr6h-22Wj_rUsf2ucn0mEFHIWyV12clEI,23476 +torch/_dynamo/graph_id_filter.py,sha256=Vj9-tVjVxCftBvh5UftAtyO1SZFBwoGYPEgGqQvEbYM,16016 +torch/_dynamo/graph_region_tracker.py,sha256=Zs0HlmBYHECjBV_O2QJlhmtuTK5tYyW6TM22ZPD2LZU,18567 +torch/_dynamo/graph_utils.py,sha256=GtBsKm5KiFeMfNwEhi7eDWEQcAZhHL0revRzgW-1pRw,3687 +torch/_dynamo/guards.py,sha256=gjRvIUFigFL4d6wovL74I0uw7OuDhuupRyKlgxMJKns,215410 +torch/_dynamo/hooks.py,sha256=prkixNTkyo9iZjlpdftr3ang4UR0t9bgdJ3jYoN1Q7k,891 +torch/_dynamo/logging.py,sha256=ltvH1A9WYtIDZkqFODBbBVglBybJpRFDj79gyYjyclM,2215 +torch/_dynamo/metrics_context.py,sha256=6SJ4moK3KY23u_i9Z1ke-DLKra-G9HIzELj1z25_kcY,8888 +torch/_dynamo/mutation_guard.py,sha256=qF01zB224O4SHvAGYEAPIkoOiUVJCzNcgTX3pkOfinU,5111 +torch/_dynamo/output_graph.py,sha256=W3TYFT825KEFsc1gWA6VZ_zdXcU9ZM74pmBg1haCDLQ,195336 +torch/_dynamo/package.py,sha256=kKs-5Ll0poUJfsWpXnhjwn1frJ9SH070v6LccIrI-tw,44020 +torch/_dynamo/pgo.py,sha256=Wjo1IuslvQ2h0t1iOtROJ5gBe4zqNPe12VjfPnnsnG4,38486 +torch/_dynamo/polyfills/__init__.py,sha256=HP74YiSD7K0RifXR-zxAbl-Y-PyWeuFSAqFMFq2JMBI,22235 +torch/_dynamo/polyfills/__pycache__/__init__.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/_collections.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/builtins.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/copy.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/functools.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/fx.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/heapq.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/itertools.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/loader.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/operator.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/os.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/pytree.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/struct.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/sys.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/tensor.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/torch_c_nn.cpython-311.pyc,, +torch/_dynamo/polyfills/__pycache__/traceback.cpython-311.pyc,, +torch/_dynamo/polyfills/_collections.py,sha256=RHkqS80la8AfbmHThq2HnoYN0BQ_HAtELrKDqWs2Bfg,662 +torch/_dynamo/polyfills/builtins.py,sha256=rfm3SBtDmXnQrRAvGYQljOtNKlLD5NE3Xem1lwgdVgg,3824 +torch/_dynamo/polyfills/copy.py,sha256=7q1StrjPFgz0fNtBMEg2gZyYWb7eGFEhgHPqufBlCSc,884 +torch/_dynamo/polyfills/functools.py,sha256=FQxB7yVCzn-IMzkjNDOVKDVYU_7BYvFIEEAt2a15h2s,961 +torch/_dynamo/polyfills/fx.py,sha256=nyQ9cjgpBs8c8KvPwfSXTs6ItDGE1swv2k3-rXhROrU,1350 +torch/_dynamo/polyfills/heapq.py,sha256=xwQLTlK-joqIDFRK7qEToXKn1NB04NmRD9gLWqkVS_E,3291 +torch/_dynamo/polyfills/itertools.py,sha256=EPOtEak1SK-pPXBylS_M3m4t06mgILv8OB_UhNSui50,9809 +torch/_dynamo/polyfills/loader.py,sha256=ZQJl16LmQXmRDjRawjWHslOtGCXBPAYDOY_frrtjsgk,1430 +torch/_dynamo/polyfills/operator.py,sha256=R9SU3Vz9riZBeCWzBjNAGJ5bkEEyl9LTNLqL858VFeA,3455 +torch/_dynamo/polyfills/os.py,sha256=XaGG7-4htiZnKr8cgpMClZRHqP1PFIqtjPGNaAbEJo0,1016 +torch/_dynamo/polyfills/pytree.py,sha256=fP8FrxZYZiMZuMuk9LD_lHZJnZUhMOD1kku0leor8m8,24429 +torch/_dynamo/polyfills/struct.py,sha256=JortoCie90NJhYDWVfmWOHQbyCPPlhEQMCMoAntrrMg,618 +torch/_dynamo/polyfills/sys.py,sha256=haOwDvxbbAK3V48lPa6Gy-T93E5TeNZsTca6IwgVkA0,947 +torch/_dynamo/polyfills/tensor.py,sha256=u5La8g1ucxhI_dzLNZsLCjGeyfesGr-117D-dNzfrME,1404 +torch/_dynamo/polyfills/torch_c_nn.py,sha256=8flz-dvgV45P7mYfEcQ4X5yg_38M8pM6Lpnk4kuDrjM,4097 +torch/_dynamo/polyfills/traceback.py,sha256=Q24-9F_v8VX7ho-xcQZeNGYggkxN3dB0g6dI_BXa-Mc,1127 +torch/_dynamo/precompile_context.py,sha256=wHzjTD4BMmN3IosSebFnKZFNRRv_ANM3Rl5qHy9M10M,7744 +torch/_dynamo/profiler.py,sha256=fgeOGbE0NI_eUOpIlyVmRffhSfFJ5icbZsXp9fRBc0A,5841 +torch/_dynamo/replay_record.py,sha256=YEpUmH6bsgSNTGE2ToVmft5ZA5DKqEuPCyfpeMR1bvQ,4370 +torch/_dynamo/repro/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_dynamo/repro/__pycache__/__init__.cpython-311.pyc,, +torch/_dynamo/repro/__pycache__/after_aot.cpython-311.pyc,, +torch/_dynamo/repro/__pycache__/after_dynamo.cpython-311.pyc,, +torch/_dynamo/repro/__pycache__/aoti.cpython-311.pyc,, +torch/_dynamo/repro/after_aot.py,sha256=E1bFTupD6C13gwVqW9hgLcruYWWbqm8aA6F8G7A22Ks,53714 +torch/_dynamo/repro/after_dynamo.py,sha256=SOGlLKKDY9W_NOHP-HtoPSWpRJiBQNbJ0WOVP7-bCok,22174 +torch/_dynamo/repro/aoti.py,sha256=RgxRkdsx8CuvmtI4Y9JdcGxLnqzVWjfTll22FtInd4k,22388 +torch/_dynamo/resume_execution.py,sha256=gBp6Qp6u0D3WsafHxdGR4Ph3ODYpjy1odm_s5WB6PpU,31059 +torch/_dynamo/side_effects.py,sha256=eK22SRKiS5BAII4UOZ4kCK8jaOBv-TJViEwAn5TxatQ,66218 +torch/_dynamo/source.py,sha256=02Owxkrj147apSkI5iA37SWWxw5xCUg37h-he9psExE,43285 +torch/_dynamo/symbolic_convert.py,sha256=aKVcyjjA4drmXNaKizyEkTLt_D6Uw1gjMP7ep_gQnzw,237477 +torch/_dynamo/tensor_version_op.py,sha256=-Jrhyh-ymNeWkx1AHblVtM81Gga5hyZ7rVQ9WjvnH8Y,2735 +torch/_dynamo/test_case.py,sha256=-pnlqO0CMUE_hUF6MgtWVt4ZZqmYtM0GFPFdHYvoIWs,8368 +torch/_dynamo/test_dont_skip_tracing_functions.py,sha256=v8sCtPbLM_cEgEXvKUvtwoVnclOjJkaDcVDYgBWL5Ww,817 +torch/_dynamo/test_minifier_common.py,sha256=-VjhygJ06RvGqXSICyZdNhR23mHH2kYhCVEEAO9IO44,12640 +torch/_dynamo/testing.py,sha256=urUKAA2_PQnDnYXftpb8W1fpRSKNP7yjMjIx1doxM0w,19302 +torch/_dynamo/trace_rules.py,sha256=bjx8emCNakAnuk3kLCitlh1_nsVOkTOHZU8K_2K5oD0,162486 +torch/_dynamo/types.py,sha256=d98_zxvBonjVMrNmB0FWb7Tujdkz9a-8Gjn7e27UOZA,4142 +torch/_dynamo/utils.py,sha256=5wPYv5thvYW1MNzw80zUEujgmUcsvmMXgt4uI-FhR20,188223 +torch/_dynamo/variables/__init__.py,sha256=fxqFL7d7cJGer_QrZ4pl9905z0FaAd1FhWdflUPjrqw,7630 +torch/_dynamo/variables/__pycache__/__init__.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/base.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/builder.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/builtin.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/constant.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/ctx_manager.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/dicts.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/distributed.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/functions.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/hashable.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/higher_order_ops.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/invoke_subgraph.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/iter.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/lazy.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/lists.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/misc.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/nn_module.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/object_protocol.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/optimizer.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/script_object.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/sdpa.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/sets.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/streams.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/tensor.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/torch.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/torch_function.cpython-311.pyc,, +torch/_dynamo/variables/__pycache__/user_defined.cpython-311.pyc,, +torch/_dynamo/variables/base.py,sha256=zUC6YX446F0RIcd60-G1jN48Jyqw0sneDSgGu8gVp2o,46061 +torch/_dynamo/variables/builder.py,sha256=QHjp3Zto4TZvuyXBdssFRyjKD5bV-G6_jM7MlrSIqco,191678 +torch/_dynamo/variables/builtin.py,sha256=rwh8ATU-T9w2I9ejeS7RuBJVkkiYPZR3bip74Ep8ooo,134882 +torch/_dynamo/variables/constant.py,sha256=YIkxRe5S4IdA7pHgwwlMr1dwFBg8HF3ZH4osjhrTiZ0,18857 +torch/_dynamo/variables/ctx_manager.py,sha256=hQmYk5Ayhu99Kj9j0l_D8DQLN7qdMwvFDNwKUteErHg,60377 +torch/_dynamo/variables/dicts.py,sha256=4_YHyRrHj_vITsO-pVDgtc0ra-q6LTJFG1CC4sSZ1To,57901 +torch/_dynamo/variables/distributed.py,sha256=GsueUEmJyf6rXabi7AogMKnDr1DafXjHnlTxRK0x5u8,8957 +torch/_dynamo/variables/functions.py,sha256=41O7VtgrwAKVEV76uL5aXG7Qyj7j_RNG2Wih7sBFAgQ,140002 +torch/_dynamo/variables/hashable.py,sha256=nn93HnM8yPS9UDu68SSVc4XeqqASXQ7gC49EJpkPpow,5476 +torch/_dynamo/variables/higher_order_ops.py,sha256=QeBxRQ054dAo68uvI3fD1ilPapPOvWHmR5PX0hIvFxQ,222601 +torch/_dynamo/variables/invoke_subgraph.py,sha256=c3iIOCUwrYndPwizg1vIfutzbHuIIEVIsSaStv_eKms,54018 +torch/_dynamo/variables/iter.py,sha256=m_RHXmoZOfwID42J7SyttZ_LqkxVN_519IfBwdTzKV8,23666 +torch/_dynamo/variables/lazy.py,sha256=VSdSv4TxtWvh8RDt3Jl44NGScCTCjvQMEwM3SKRNgsI,10600 +torch/_dynamo/variables/lists.py,sha256=UaKpmBjTpAywhtABZR37rZ8UW-aFwXqIVdq1ZXsc9Pc,65791 +torch/_dynamo/variables/misc.py,sha256=J8DRqVVo6NHGdJd0lCwdF_CZR7V5bxYZeUz69oshYTY,97222 +torch/_dynamo/variables/nn_module.py,sha256=vDwXNztsq7tfFQcBj4LgiaEdbGyTLamTsH5k_pai9gY,60199 +torch/_dynamo/variables/object_protocol.py,sha256=2lugmmULuAmg61o2h-tnAvDZOClgqaVhhjaIrMv7IlA,11384 +torch/_dynamo/variables/optimizer.py,sha256=FHWz85gE2dL9QXJaostyrh1-mYz_pbOXNifqojApbEU,17876 +torch/_dynamo/variables/script_object.py,sha256=-xyQn7oYvUGe9efLwdOJ0FZA8Tckfdob4JorMVJrtMw,21327 +torch/_dynamo/variables/sdpa.py,sha256=TCcZ_4qj7alBMoZiM5yVGIv28oon93PfnwyezWp0_l4,3049 +torch/_dynamo/variables/sets.py,sha256=mWeCsgk8poL7ziGJv6pvPb8RTQgv9PcHkV8SrgbLNps,29347 +torch/_dynamo/variables/streams.py,sha256=riWBzrd3nZy-dN437VWIDzhGunYyIFqS5_4GdnRmG3Y,23796 +torch/_dynamo/variables/tensor.py,sha256=5zU3VA9eUqjkaj-bDMjW5uc0qzhk2xbxjj7_b5Jjdmc,106637 +torch/_dynamo/variables/torch.py,sha256=3HyY3qwiOLCZMx3GjZQNskhceXx0xndalzKVDgrvLzg,149387 +torch/_dynamo/variables/torch_function.py,sha256=Da3eF8D34Z_7S_G9kZleVJ7bpCuxQPSxcoSY4R_ZAys,28632 +torch/_dynamo/variables/user_defined.py,sha256=f5QjX5oimTxyqcpdpsEavy8EEYX6g_r7UHqtWgYLu7g,144845 +torch/_environment.py,sha256=pDUFrq8tC8SqJTir9Q3Vr01UbYCYH-2KEJbleFrV0FQ,133 +torch/_export/__init__.py,sha256=jHjsuXiPAKHY3YZqRsc-ir5VA109_FngyZYnE3yiQJk,6753 +torch/_export/__pycache__/__init__.cpython-311.pyc,, +torch/_export/__pycache__/config.cpython-311.pyc,, +torch/_export/__pycache__/converter.cpython-311.pyc,, +torch/_export/__pycache__/error.cpython-311.pyc,, +torch/_export/__pycache__/non_strict_utils.cpython-311.pyc,, +torch/_export/__pycache__/pass_base.cpython-311.pyc,, +torch/_export/__pycache__/tools.cpython-311.pyc,, +torch/_export/__pycache__/utils.cpython-311.pyc,, +torch/_export/__pycache__/verifier.cpython-311.pyc,, +torch/_export/__pycache__/wrappers.cpython-311.pyc,, +torch/_export/config.py,sha256=eqicgOdgNXZ_txD6gBpZ4bDgJ0SMlGhV2eWpv6CNvdA,1410 +torch/_export/converter.py,sha256=dtBqLdYJX5I4Wppn5UTe2rfPuiAHjW-5-ipUClf07-g,65079 +torch/_export/db/__init__.py,sha256=a3XxW1RcNAPwEVaI2g11hpnJvSHxGUFsGmDAWfDnLP8,206 +torch/_export/db/__pycache__/__init__.cpython-311.pyc,, +torch/_export/db/__pycache__/case.cpython-311.pyc,, +torch/_export/db/__pycache__/gen_example.cpython-311.pyc,, +torch/_export/db/__pycache__/logging.cpython-311.pyc,, +torch/_export/db/case.py,sha256=nHiUtUtBGliPIKjHI07ZBUq6NE6Re4qwIZEsF66ZPpg,5199 +torch/_export/db/examples/__init__.py,sha256=PlziokXeStI1gWELwC7zX-T7ZViONPRitQ1uFWNAxL4,1648 +torch/_export/db/examples/__pycache__/__init__.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/assume_constant_result.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/autograd_function.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/class_method.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/cond_branch_class_method.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/cond_branch_nested_function.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/cond_branch_nonlocal_variables.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/cond_closed_over_variable.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/cond_operands.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/cond_predicate.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/constrain_as_size_example.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/constrain_as_value_example.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/decorator.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/dictionary.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/dynamic_shape_assert.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/dynamic_shape_constructor.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/dynamic_shape_if_guard.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/dynamic_shape_map.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/dynamic_shape_round.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/dynamic_shape_slicing.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/dynamic_shape_view.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/fn_with_kwargs.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/list_contains.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/list_unpack.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/model_attr_mutation.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/nested_function.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/null_context_manager.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/optional_input.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/pytree_flatten.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/scalar_output.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/specialized_attribute.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/static_for_loop.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/static_if.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/tensor_setattr.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/type_reflection_method.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/unsupported_operator.cpython-311.pyc,, +torch/_export/db/examples/__pycache__/user_input_mutation.cpython-311.pyc,, +torch/_export/db/examples/assume_constant_result.py,sha256=MPGQemV_ymD07kXuOVQkizSxTaTlu4vr2bAP3xxE3M8,510 +torch/_export/db/examples/autograd_function.py,sha256=byZs2M2TE3nU9a_UM3LMZw4gf-iFCsAok-wnqxAGah4,652 +torch/_export/db/examples/class_method.py,sha256=KM6pkWdEc1SMybUR-O2hCN1tD9Vl0B22gPI2LgDiZZ4,499 +torch/_export/db/examples/cond_branch_class_method.py,sha256=PpspYrzN3LBl0CIuGjjZon6rGQeQ32YtaIqbpRDbFyc,1327 +torch/_export/db/examples/cond_branch_nested_function.py,sha256=HgN_wnbfzoATNkD7ksQSwdo6PVhiIP5RqneZvGSTxeQ,1302 +torch/_export/db/examples/cond_branch_nonlocal_variables.py,sha256=B_ROqnPdNhhy7zbP9cyqLUl1LgDgAYYGSr8EyNRiy-M,1841 +torch/_export/db/examples/cond_closed_over_variable.py,sha256=iSWSxhJhumU6CsRSzXUeW7YVxbgjf-C05EDc0Foi-a4,547 +torch/_export/db/examples/cond_operands.py,sha256=mBIzaZy_VYUzNWS8aqrTCPwUM1qbUPlJWXX4oWYgJGE,799 +torch/_export/db/examples/cond_predicate.py,sha256=PKpI-GKAdVMDh11EDIpQ3JXvRMXNGuJTKXbGb_uyQjY,663 +torch/_export/db/examples/constrain_as_size_example.py,sha256=13bb0R9cuRWRNYgEFfSe8AkQdgf39zvf5EDvhIVXJhw,515 +torch/_export/db/examples/constrain_as_value_example.py,sha256=XtZLDWMmL1hOZPqtNbOIUYQ6FNX24_jMvu7Uzq9sXDw,571 +torch/_export/db/examples/decorator.py,sha256=U5YTo_25VfE_Hy9uoD4eYN8wKSJGpOG0dAJ1TZjL1i4,480 +torch/_export/db/examples/dictionary.py,sha256=gSkKkS68Vw_cropGQ57i61jqTleFsLECIXP6pZMqLBE,404 +torch/_export/db/examples/dynamic_shape_assert.py,sha256=5ORg04FDKJb0rEmBVlsliyeNDerTuCDANWJbHz6_nr4,478 +torch/_export/db/examples/dynamic_shape_constructor.py,sha256=68CRA-j59c0dPBc5hOBsEBmWpWz0opzovyqGIaTfKuw,396 +torch/_export/db/examples/dynamic_shape_if_guard.py,sha256=VG6YmeK31p2jwLPG_Ci3VWmWUMQnJIp_IuijY_UY-dU,560 +torch/_export/db/examples/dynamic_shape_map.py,sha256=XN8_Et1EVYaxmdnACPkiDAVJZ3jkYRaj9U6mU44ZntM,454 +torch/_export/db/examples/dynamic_shape_round.py,sha256=sMyW43DVXFtZa6jFO08XxVjyw06LcJ9oyuJH5jAsG4Q,525 +torch/_export/db/examples/dynamic_shape_slicing.py,sha256=Rowy_RfPf1kW1l_oCp71YVflt2ivhpP1u_HzTEupxQA,388 +torch/_export/db/examples/dynamic_shape_view.py,sha256=fEuSTQRghtJY8VcBi7qwiZTQXwd9ux8d1PRBYw6Qd6E,444 +torch/_export/db/examples/fn_with_kwargs.py,sha256=fEndpGN__D456vh3NAamoK1KdmtEyvmAsW0gZB0TExo,731 +torch/_export/db/examples/list_contains.py,sha256=lzzkGcg6YoDeFDRh7GlY6BlmQoeA8aVfSFu-1FoiMhk,519 +torch/_export/db/examples/list_unpack.py,sha256=u7tczh6wbR1PoYQxuRPpsSfxgMiGhu4VocA50ZrwKFw,568 +torch/_export/db/examples/model_attr_mutation.py,sha256=wIrxdizyoC1EGByE4THkYXbDPw6Vh86lphvh_OR9kug,628 +torch/_export/db/examples/nested_function.py,sha256=S-MXIlaFn4RGeA4GEFIEHG-pcM9WBznXN3E2k-sDlZo,491 +torch/_export/db/examples/null_context_manager.py,sha256=LM5a8DO8x0A0yaA0ZAhOVybcRIcpHy2fxCH0YoqH_2M,478 +torch/_export/db/examples/optional_input.py,sha256=JHVO8AbT5ihdUrr_ee5kZMuJOnyJZS2EdL9FYdS9Yns,447 +torch/_export/db/examples/pytree_flatten.py,sha256=MFIw36BZAXANZ-iajZi--lbPzf1EyDrcmWxPeQ_VCyA,376 +torch/_export/db/examples/scalar_output.py,sha256=_CMPbP3nbTMGzfiUy3ACOdxxoL8_F1lIEya_JH4HTZ4,543 +torch/_export/db/examples/specialized_attribute.py,sha256=GiKqGiSSPwPIL-QPinKinHNEiQvRsptBx319j7FpRfA,520 +torch/_export/db/examples/static_for_loop.py,sha256=wYtwzouoqqVC-Zsw_kHYe_Fzz-0mPxZ9zGwWrjdv8OQ,385 +torch/_export/db/examples/static_if.py,sha256=uzqQHEytorjGRbxJQQxam6nJZLHcPIPw1uzMfj7KU8w,397 +torch/_export/db/examples/tensor_setattr.py,sha256=LXED12gZiYOvhZXBzThejfW2meNYI8U-8_aIVwfdaqI,337 +torch/_export/db/examples/type_reflection_method.py,sha256=CoXAARKtqSNciuA5V7Tkx0MVhq-h6c4khPRSnxFDjcM,461 +torch/_export/db/examples/unsupported_operator.py,sha256=Y8tqzwbbN14AznoBlavyPBLPYPqQtnp9lx47rvJMQJU,399 +torch/_export/db/examples/user_input_mutation.py,sha256=KJY_YR4w_nEZ66YVMmPsBKv6PfQ7GnQwkV12NJcUCxs,302 +torch/_export/db/gen_example.py,sha256=cyFTffHUtwkSPfhk-5rKaKFA5l56JErtcR-VBPb7t_o,604 +torch/_export/db/logging.py,sha256=96gvkbZUoxwdFkzsSP2ofX9It_S2TiK_fTnfc4Wn8GI,1671 +torch/_export/error.py,sha256=OjvFCTGZVLlBtbTvaL4xDbBAizynfhwO_2fONObaISk,1770 +torch/_export/non_strict_utils.py,sha256=orBE2uwIo_NyZd0T6m8gfD3ITkvv8fu4Fs4bfUys08s,44149 +torch/_export/pass_base.py,sha256=DD-Dl5NagBSyQLZ69VZu6CIX60RTZFtRiDYqsTt8Rnc,19037 +torch/_export/pass_infra/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_export/pass_infra/__pycache__/__init__.cpython-311.pyc,, +torch/_export/pass_infra/__pycache__/node_metadata.cpython-311.pyc,, +torch/_export/pass_infra/__pycache__/proxy_value.cpython-311.pyc,, +torch/_export/pass_infra/node_metadata.py,sha256=fF9lD71OeQ5HlSfR--IzjNMe7s67_N9T5Z1QvhGsa5s,771 +torch/_export/pass_infra/proxy_value.py,sha256=tIkdsWMBGdKHWtEW3vVS_6RqZTAZf7CimcY0OIEF31Y,1456 +torch/_export/passes/__init__.py,sha256=78MzFjtaVabk1z2X12WMIlmtkjYaFmgNJlDCTCepNt8,88 +torch/_export/passes/__pycache__/__init__.cpython-311.pyc,, +torch/_export/passes/__pycache__/_node_metadata_hook.cpython-311.pyc,, +torch/_export/passes/__pycache__/add_runtime_assertions_for_constraints_pass.cpython-311.pyc,, +torch/_export/passes/__pycache__/collect_tracepoints_pass.cpython-311.pyc,, +torch/_export/passes/__pycache__/constant_folding.cpython-311.pyc,, +torch/_export/passes/__pycache__/functionalize_side_effectful_ops_pass.cpython-311.pyc,, +torch/_export/passes/__pycache__/insert_custom_op_guards.cpython-311.pyc,, +torch/_export/passes/__pycache__/lift_constants_pass.cpython-311.pyc,, +torch/_export/passes/__pycache__/remove_runtime_assertions.cpython-311.pyc,, +torch/_export/passes/__pycache__/replace_autocast_with_hop_pass.cpython-311.pyc,, +torch/_export/passes/__pycache__/replace_quantized_ops_with_standard_ops_pass.cpython-311.pyc,, +torch/_export/passes/__pycache__/replace_set_grad_with_hop_pass.cpython-311.pyc,, +torch/_export/passes/__pycache__/replace_view_ops_with_view_copy_ops_pass.cpython-311.pyc,, +torch/_export/passes/__pycache__/replace_with_hop_pass_util.cpython-311.pyc,, +torch/_export/passes/_node_metadata_hook.py,sha256=L9i3p7Vzz53vG8wpirkYMi61F-Uxcn7-mnjMS09n7VU,3435 +torch/_export/passes/add_runtime_assertions_for_constraints_pass.py,sha256=Lrh2xrn5yh7UivXloc0rK4ETIQRS9ssUQyJSxPfeJjk,10300 +torch/_export/passes/collect_tracepoints_pass.py,sha256=3vWcN1V-hl588dQ9estymdTNxivxiex8YLA-7ZKmShs,7215 +torch/_export/passes/constant_folding.py,sha256=qBI327xEx9Eq1iM4RqmEntRfyxnIstvhsIpcm8bdbwc,11290 +torch/_export/passes/functionalize_side_effectful_ops_pass.py,sha256=1fjMfdGrG138Lm7n5KLVkbsbAM9bGxu66BdjAwov42o,3370 +torch/_export/passes/insert_custom_op_guards.py,sha256=89yaS7aO2E-BzxllQWa3HJRmFBKwFbSg1861bAem4CI,3096 +torch/_export/passes/lift_constants_pass.py,sha256=2GYnum31seGBk39TsaqDT9ydKgq8uOjXWFas1v7PK6U,18927 +torch/_export/passes/remove_runtime_assertions.py,sha256=uyvuiO5Qy9AYVfbaOvyOYeBaZ_yUzgoPSsfJ1lDb0-g,1588 +torch/_export/passes/replace_autocast_with_hop_pass.py,sha256=D0uBUEdo2ogaitXuCm8P84kO3e8vMw0wD1cSpkiCzGs,8358 +torch/_export/passes/replace_quantized_ops_with_standard_ops_pass.py,sha256=c4jBJptGwEdfsojm9qQgFcmUUd3IcBUpfbws1WlC0ZQ,27651 +torch/_export/passes/replace_set_grad_with_hop_pass.py,sha256=gY_EdMXcjPmzDrxlxG-eW7JCuCYzYXbqJ9z2lAcaYis,4911 +torch/_export/passes/replace_view_ops_with_view_copy_ops_pass.py,sha256=-8zA-otMVordLl2mWpPcutMJiTF1pZCLKS-xh2_1cjY,2380 +torch/_export/passes/replace_with_hop_pass_util.py,sha256=iVnsVwmMSBBrnz9oU79yfApJgWvFiI9l3K7hj54iulU,8243 +torch/_export/serde/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_export/serde/__pycache__/__init__.cpython-311.pyc,, +torch/_export/serde/__pycache__/dynamic_shapes.cpython-311.pyc,, +torch/_export/serde/__pycache__/schema.cpython-311.pyc,, +torch/_export/serde/__pycache__/schema_check.cpython-311.pyc,, +torch/_export/serde/__pycache__/serialize.cpython-311.pyc,, +torch/_export/serde/__pycache__/union.cpython-311.pyc,, +torch/_export/serde/dynamic_shapes.py,sha256=EgxY4ZeT5UOeSKVo0qCbCBm9BdIbjnQZvjhI0ko0EmQ,11752 +torch/_export/serde/export_schema.thrift,sha256=XXqURSIjrmItA6xEl_Oy98le65Ss4NQsFKtoQ-7n64k,7493 +torch/_export/serde/schema.py,sha256=Jmr2URIz5FGGWArVxraeDQ8-Bo2lOGvZnRlsaI8KvuU,16742 +torch/_export/serde/schema.yaml,sha256=oz48yb7FH8txlwWsE9DehcHcYBNz4JYTtkfq7klp6Ys,10147 +torch/_export/serde/schema_check.py,sha256=3KRNhxvwuKzI2B4b24Ks18_zbSVBWGfSk_dKZzQsUOs,30520 +torch/_export/serde/serialize.py,sha256=cW0hCU42Ub6TO0stN_qflIY2zRkoWg20hj1u69HwMO4,180536 +torch/_export/serde/union.py,sha256=rDgvWbB3546NWLbFc6UZRkHyvDqKF-2F3k-MRw5xWHo,3336 +torch/_export/tools.py,sha256=vLys4J5NoTU1EY6L1Ear24gOZX6mXGG6CraTpxiHyiU,4577 +torch/_export/utils.py,sha256=u5-fAILCbbUPfqSrn2Vpcw1YZLpyGBXHtPmoVpaINvo,60446 +torch/_export/verifier.py,sha256=8SdVOxLVnGKmTa9QQlpKCpfZFwdnNk-tefcl6ETsAQo,22795 +torch/_export/wrappers.py,sha256=UzgIPm59W_k8-oWjmpUsJa7118KRHOPpuIvro-ang0k,13386 +torch/_functorch/__init__.py,sha256=a3XxW1RcNAPwEVaI2g11hpnJvSHxGUFsGmDAWfDnLP8,206 +torch/_functorch/__pycache__/__init__.cpython-311.pyc,, +torch/_functorch/__pycache__/aot_autograd.cpython-311.pyc,, +torch/_functorch/__pycache__/apis.cpython-311.pyc,, +torch/_functorch/__pycache__/autograd_function.cpython-311.pyc,, +torch/_functorch/__pycache__/batch_norm_replacement.cpython-311.pyc,, +torch/_functorch/__pycache__/benchmark_utils.cpython-311.pyc,, +torch/_functorch/__pycache__/compile_utils.cpython-311.pyc,, +torch/_functorch/__pycache__/compilers.cpython-311.pyc,, +torch/_functorch/__pycache__/config.cpython-311.pyc,, +torch/_functorch/__pycache__/deprecated.cpython-311.pyc,, +torch/_functorch/__pycache__/eager_transforms.cpython-311.pyc,, +torch/_functorch/__pycache__/functional_call.cpython-311.pyc,, +torch/_functorch/__pycache__/fx_minifier.cpython-311.pyc,, +torch/_functorch/__pycache__/make_functional.cpython-311.pyc,, +torch/_functorch/__pycache__/partitioners.cpython-311.pyc,, +torch/_functorch/__pycache__/predispatch.cpython-311.pyc,, +torch/_functorch/__pycache__/pyfunctorch.cpython-311.pyc,, +torch/_functorch/__pycache__/python_key.cpython-311.pyc,, +torch/_functorch/__pycache__/pytree_hacks.cpython-311.pyc,, +torch/_functorch/__pycache__/top_operators_github_usage.cpython-311.pyc,, +torch/_functorch/__pycache__/utils.cpython-311.pyc,, +torch/_functorch/__pycache__/vmap.cpython-311.pyc,, +torch/_functorch/_activation_checkpointing/__init__.py,sha256=a3XxW1RcNAPwEVaI2g11hpnJvSHxGUFsGmDAWfDnLP8,206 +torch/_functorch/_activation_checkpointing/__pycache__/__init__.cpython-311.pyc,, +torch/_functorch/_activation_checkpointing/__pycache__/ac_logging_utils.cpython-311.pyc,, +torch/_functorch/_activation_checkpointing/__pycache__/graph_info_provider.cpython-311.pyc,, +torch/_functorch/_activation_checkpointing/__pycache__/knapsack.cpython-311.pyc,, +torch/_functorch/_activation_checkpointing/__pycache__/knapsack_evaluator.cpython-311.pyc,, +torch/_functorch/_activation_checkpointing/__pycache__/remat_using_tags_for_fwd_loss_bwd_graph_pass.cpython-311.pyc,, +torch/_functorch/_activation_checkpointing/ac_logging_utils.py,sha256=3IDGpnkWMtZMp25VDVQC2bo_UaFfqWinVYdK-NIxm4s,7984 +torch/_functorch/_activation_checkpointing/graph_info_provider.py,sha256=agG5nlht2-xRTRCjIXsJOn0S1XmMhXZpXPiHyx-mTDg,12701 +torch/_functorch/_activation_checkpointing/knapsack.py,sha256=eULO3neRqPtg84WsKf57jn4thXQjxZ_oQEUdQ_6_erQ,9260 +torch/_functorch/_activation_checkpointing/knapsack_evaluator.py,sha256=o1E0Rg96tyfcmCumZ7UA1oYE100EXexi5xW2A-ebvcE,12279 +torch/_functorch/_activation_checkpointing/remat_using_tags_for_fwd_loss_bwd_graph_pass.py,sha256=3piETBHgWHOilu0XzNx3TDy-mlGdMNcowThCNxdryz8,7653 +torch/_functorch/_activation_offloading/__init__.py,sha256=a3XxW1RcNAPwEVaI2g11hpnJvSHxGUFsGmDAWfDnLP8,206 +torch/_functorch/_activation_offloading/__pycache__/__init__.cpython-311.pyc,, +torch/_functorch/_activation_offloading/__pycache__/activation_offloading.cpython-311.pyc,, +torch/_functorch/_activation_offloading/__pycache__/offload_ops.cpython-311.pyc,, +torch/_functorch/_activation_offloading/activation_offloading.py,sha256=TVq4akXwZP83WGSx7VyjMrORdKH8A_PLNopQ0SM7hAY,30867 +torch/_functorch/_activation_offloading/offload_ops.py,sha256=VN1tAVU93r0ydKRghS0mPvFi_g-xbP3WTsSZx_hpRj8,6449 +torch/_functorch/_aot_autograd/__init__.py,sha256=a3XxW1RcNAPwEVaI2g11hpnJvSHxGUFsGmDAWfDnLP8,206 +torch/_functorch/_aot_autograd/__pycache__/__init__.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/aot_autograd_result.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/autograd_cache.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/collect_metadata_analysis.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/descriptors.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/frontend_utils.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/functional_utils.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/fx_utils.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/graph_capture.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/graph_capture_wrappers.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/graph_compile.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/indexed_dict.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/input_output_analysis.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/logging_utils.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/runtime_wrappers.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/schemas.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/streams.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/subclass_codegen.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/subclass_parametrization.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/subclass_utils.cpython-311.pyc,, +torch/_functorch/_aot_autograd/__pycache__/utils.cpython-311.pyc,, +torch/_functorch/_aot_autograd/aot_autograd_result.py,sha256=pl2-fA5_1V5_8k2wl9c7mrsz2xh8pTPtWxjT73McS_s,27457 +torch/_functorch/_aot_autograd/autograd_cache.py,sha256=ouUX-FEjoEGNrYxDvVRi26_bxAY7frycdTKUaKemAVo,58933 +torch/_functorch/_aot_autograd/collect_metadata_analysis.py,sha256=bXzDMVFNpfu_oBymxwoElISZY3gw-E-ke9wU9niYnvI,45986 +torch/_functorch/_aot_autograd/descriptors.py,sha256=ETe2aGDaX-uMGHOe6RlAauQWTz0ioUyy3wcvq1YG6no,28153 +torch/_functorch/_aot_autograd/frontend_utils.py,sha256=AnR3PUyTvba2r77mgUpcW8G3ui7OrwjAWeHmIUrXjio,17632 +torch/_functorch/_aot_autograd/functional_utils.py,sha256=8H0OTZgxTEid61fzfHsIVhDtDuFhiWW4ymybcyLsY3U,28692 +torch/_functorch/_aot_autograd/fx_utils.py,sha256=Kj_tJK0kOBxI-eQqB3EiT-q6RYtDZ1M_Pn-pQXMKgWc,12645 +torch/_functorch/_aot_autograd/graph_capture.py,sha256=D8Km-3G6Rnxq7c1Hvn5j35_HN5Nkir70uQJsQiuW2f4,21894 +torch/_functorch/_aot_autograd/graph_capture_wrappers.py,sha256=Dtj1i389JM6Ty2zLBEB1L9W09i4mMOzWaU0WFLbZDrA,70310 +torch/_functorch/_aot_autograd/graph_compile.py,sha256=dBymkjcJb2R4FoYuBvOOZrbQQcSmjkiXQQs04Ju6-7c,103634 +torch/_functorch/_aot_autograd/indexed_dict.py,sha256=tXYzmFA39UlhLc9x_4suTdDXgMMLXsJYhPiocutWjfw,1743 +torch/_functorch/_aot_autograd/input_output_analysis.py,sha256=f8seNBKQCX1EvX-Jldft3krD4e_8zJfnSvvm-lLUGDI,20338 +torch/_functorch/_aot_autograd/logging_utils.py,sha256=u9H_0eLOYgRbyeAVjL6BBq0xJs1s9vbxjqvb0-hgkeA,5515 +torch/_functorch/_aot_autograd/runtime_wrappers.py,sha256=tkpE_lr8a17QclZvYXvSZzqR31ESt2gskp24fDGY794,135866 +torch/_functorch/_aot_autograd/schemas.py,sha256=1qeMUipmmSBGO3ckfzQCPVAUYYgWAMUID_GCbNcAcQ4,63942 +torch/_functorch/_aot_autograd/streams.py,sha256=7TkqSr5jEHW0iAlfu3724WcH3fap_I8ef2wwOvwzvXU,23628 +torch/_functorch/_aot_autograd/subclass_codegen.py,sha256=fq5tWAX6Ep0ltmI21GqvGbu8xn38cIK7k94MzB3oUqk,14474 +torch/_functorch/_aot_autograd/subclass_parametrization.py,sha256=Ti1mD_B-pXVGl9yrp7Fc__IyZfrjNnGX8o9iPEJ4KRo,5728 +torch/_functorch/_aot_autograd/subclass_utils.py,sha256=SHBJXeRbcU8vQToKLN-GD0oINjkXNC_BJnHEwgc67A8,26439 +torch/_functorch/_aot_autograd/utils.py,sha256=M5L2Jqu3C82jkJA6wq4cBEAnj5ZjBvYxiHaTl5A_dHY,32604 +torch/_functorch/aot_autograd.py,sha256=ee1xQ2wFH7VDDsu599EBb5PgG-MVoXxp06rESjSSrvg,78981 +torch/_functorch/apis.py,sha256=UvFNcxTjoBTKknskA1d4KYtIw-HxfLqt4Dd-AJpXapE,19970 +torch/_functorch/autograd_function.py,sha256=7cHMLB3BxmcgMX-_HACqIivS7b3Q5sfs_g3DV3Uel-I,34257 +torch/_functorch/batch_norm_replacement.py,sha256=liTJLEOFiap3ed2fBoxdRUM-9iuw-6985kJSXZnUlNs,857 +torch/_functorch/benchmark_utils.py,sha256=6umZ5dDJO6bRKqa6U7OXhTFXjC9a_OviLMe5emlWIrs,7153 +torch/_functorch/compile_utils.py,sha256=oq_gBsPFNXBvrNmswlPu6C8ywBVeO_bC0oDAT2_qbuU,8463 +torch/_functorch/compilers.py,sha256=bhsaDddxFMHOuK7OFDksaYii48v2d0MDMcrX9xneV-U,16626 +torch/_functorch/config.py,sha256=EmkkJpo3GcQpR7OhP09o8_exVVdHitaaMpiypugCbbk,19233 +torch/_functorch/deprecated.py,sha256=SHvlwZPnuj_0sGtuj2O1WTs0Sb9-ZWvhVOJSxnJQqtU,5731 +torch/_functorch/eager_transforms.py,sha256=XhtVv0UlZeewJ8SLjj4Bgbo0KFKHwro1B0Qewt5BJLc,77114 +torch/_functorch/functional_call.py,sha256=3O5PMf1mu3ZB_BqZ3W62zPH6Hn5_PfmjQ_JKNtnbJAc,10577 +torch/_functorch/fx_minifier.py,sha256=5svylRsJBBNBSZvt5JXgaN2dEpjWmm4qtS3fgtvbhCk,20132 +torch/_functorch/make_functional.py,sha256=PVBP0fo9nphZXaGu-3wMfEdiDjPd05pWbpiqjPcHUno,24492 +torch/_functorch/partitioners.py,sha256=_-FYhRKGW1yUcprOS1odT7q2TsGzFjZ2hG3XEMkb93k,149424 +torch/_functorch/predispatch.py,sha256=F6FAxaAv9ijDXsOudvQyOx51hP9GhrMOrDI4Pn1e6dg,5707 +torch/_functorch/pyfunctorch.py,sha256=orEdNWO3nruuRt9zFhM8lr3aLLG_ANPZNBlTiUGf_Pg,12546 +torch/_functorch/python_key.py,sha256=ZeINkQZ4B4PAdoz-zLMNU72w6cVHzE1nObMikLI0oL4,442 +torch/_functorch/pytree_hacks.py,sha256=dWzTSKcqGhK4zHgmgJbyis2kzB69OH71GP_YCDBDbp0,698 +torch/_functorch/top_operators_github_usage.py,sha256=pn1_T5Grb-I5GjY-8vAMoVuKAPyhOvCVpKDQpLcI_Kk,21568 +torch/_functorch/utils.py,sha256=wAjJNpP0VfMJXCNSZkOpl7vmPb_Fo25aJna8JXyzBk8,1039 +torch/_functorch/vmap.py,sha256=ckC0H8Hbi487GJ5XPt1OjWqdsbPtL5wIVwLaYGIwEo0,19405 +torch/_guards.py,sha256=rRq8DRibXuMQ1s81d4sTLk_uLmuruJaIOAboqE7-Adw,56770 +torch/_higher_order_ops/__init__.py,sha256=ug4B40etVAvaRuHXIIFSwlihNkrDF5folzyXL8qgaig,2795 +torch/_higher_order_ops/__pycache__/__init__.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/_invoke_quant.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/aoti_call_delegate.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/associative_scan.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/auto_functionalize.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/base_hop.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/cond.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/cudagraph_conditional_nodes.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/effects.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/executorch_call_delegate.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/flat_apply.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/flex_attention.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/foreach_map.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/hints_wrap.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/inline_asm_elementwise.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/invoke_leaf_function.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/invoke_subgraph.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/local_map.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/map.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/out_dtype.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/partitioner.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/print.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/run_const_graph.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/scan.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/schema.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/strict_mode.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/torchbind.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/triton_kernel_wrap.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/utils.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/while_loop.cpython-311.pyc,, +torch/_higher_order_ops/__pycache__/wrap.cpython-311.pyc,, +torch/_higher_order_ops/_invoke_quant.py,sha256=igRdMRGVeX-YVwt5-kHJLUGQhtQ11s7LnP0Hq4ypxqU,1739 +torch/_higher_order_ops/aoti_call_delegate.py,sha256=yqIpNSyyk5gV_6BAYysQISovkdxN9inzPX-KOCCDG3M,6172 +torch/_higher_order_ops/associative_scan.py,sha256=rN_jqb6KgbJGyvLoRLgtyC6r0hv-svnmI9Yk0yxEeIk,37196 +torch/_higher_order_ops/auto_functionalize.py,sha256=ph0-sijqv56Z2NPxRYO3eTmJbwNPQtpj_YholnE48AU,40179 +torch/_higher_order_ops/base_hop.py,sha256=bY20-p-EWR1hS9zojIBRDqdoSl8Am6lH7iLnv4DBn-0,10914 +torch/_higher_order_ops/cond.py,sha256=lK4Gv7LBa7__Po_WtHb7_c0U5lvvJkwQOR8GRa_o1eo,30294 +torch/_higher_order_ops/cudagraph_conditional_nodes.py,sha256=ILIw-3hjAYAqUzzIpHF-AvYFbwlIu-M4YQuWOVAd-Es,4277 +torch/_higher_order_ops/effects.py,sha256=0MWVhXGHhp3AgHO5WppouiaB6IiqBBR8U13yCaR7dCA,11827 +torch/_higher_order_ops/executorch_call_delegate.py,sha256=SRnG6PRrZ6MzPvV43Jo_dkn-mFNba7liwfkEQxnwiRY,6148 +torch/_higher_order_ops/flat_apply.py,sha256=IZO-uONbPfB_L1_SMR_wXEQr702tdn1-lcuyrOpCHA8,6323 +torch/_higher_order_ops/flex_attention.py,sha256=u91_ECgC7BTy1gcK69hSUlUvBgrQnrG4mnhCx9PcrCI,51104 +torch/_higher_order_ops/foreach_map.py,sha256=xEvWbG-fWPZaVF1AjLrPnG2bJ1Duwr6Uf0YmVvndJx8,690 +torch/_higher_order_ops/hints_wrap.py,sha256=QL9hly43ZnvGkYFuDX0-pF3G0fGy6Ag1bwRuzbv8EyA,4821 +torch/_higher_order_ops/inline_asm_elementwise.py,sha256=Jc1jAX_b5XzlNt98x8cqvavhalhQlFnYUXB8LYHzAB8,9405 +torch/_higher_order_ops/invoke_leaf_function.py,sha256=cYiTmiBDGPSxUjdjpP1E5FjiqrRd-8GySMV8yng_7Bk,33403 +torch/_higher_order_ops/invoke_subgraph.py,sha256=HjYRcF20mpyC0JFQINGILrCq5tiLRu54RhdIm23Ocgs,43209 +torch/_higher_order_ops/local_map.py,sha256=sEDnYjVHDH-AbtkMEGwXHbdu1KJF6FnpQPX7IaMInuY,24803 +torch/_higher_order_ops/map.py,sha256=3xrpy6HF6uEwM1G3dsOgjASAEAwWtYhZhWo34ou8cXI,11940 +torch/_higher_order_ops/out_dtype.py,sha256=c35Sik50tRAq_v5iP5Ft9RKXhi7gXlJxT3_0ML9vkCg,5612 +torch/_higher_order_ops/partitioner.py,sha256=IF7oOJsjP_PgANEvEvN-Z9xhW4BdxX4Mauezaz8oY30,14279 +torch/_higher_order_ops/passes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_higher_order_ops/passes/__pycache__/__init__.cpython-311.pyc,, +torch/_higher_order_ops/passes/__pycache__/inline_invoke_subgraph.cpython-311.pyc,, +torch/_higher_order_ops/passes/inline_invoke_subgraph.py,sha256=RPkYAKqqabWySa1L1lszw77NKCn-8p_RbhpjUsERI34,2397 +torch/_higher_order_ops/print.py,sha256=Yq1cTAJk-eWLOK2LW2FBgIQ0OqqdRXiH1Jceeiqy3Do,5706 +torch/_higher_order_ops/run_const_graph.py,sha256=qWXm3rsyRsaOiX0P0KzJ1iFUX89hzOb12bBpNKE20IU,2942 +torch/_higher_order_ops/scan.py,sha256=W1nt59x2tORuVz0QcCJlEbBSrk7M1DVzz6FsrvRw7Qs,37547 +torch/_higher_order_ops/schema.py,sha256=7X9iaFjzbYVw4pJAHSN0nwDc4p0m5t85ji3BIZoJJB8,12005 +torch/_higher_order_ops/strict_mode.py,sha256=RAXUZoKR6uXCxIg6AE_wgWFl5UZeNMcVW99tL50cJfU,2920 +torch/_higher_order_ops/torchbind.py,sha256=B2h718_dJujJvJwhdoKrQ8WXjCiXydI-cYdFRqTvkwg,6490 +torch/_higher_order_ops/triton_kernel_wrap.py,sha256=WgC0MnVwAJ3_fmjJS9edXOe7rN9YBa2JOgP2HZDhyW0,98662 +torch/_higher_order_ops/utils.py,sha256=uI_6dzVIyiQXcNg14KtyoOMiBJjzsm-b0pXN99L7A0o,52758 +torch/_higher_order_ops/while_loop.py,sha256=Lap2NB_j3Ywf-W0uV6MqigYeLI78vqL0PTlH112FLrk,37671 +torch/_higher_order_ops/wrap.py,sha256=Ta1NtOOtA6Uuc-B6y1uxEfuWobL4M_ebLhy7h6Fa7VU,20167 +torch/_inductor/__autotune_main__.py,sha256=teBVDb5194YrCJgy05KRE1yAkfDrAfFIA2cIwYFnyRw,913 +torch/_inductor/__init__.py,sha256=Uh5kkTHhfWeZHtCZ8z5TfG3TM7tCOwT9JPHnHLMOURQ,15232 +torch/_inductor/__pycache__/__autotune_main__.cpython-311.pyc,, +torch/_inductor/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/__pycache__/analyze_preserves_zero_mask.cpython-311.pyc,, +torch/_inductor/__pycache__/aoti_eager.cpython-311.pyc,, +torch/_inductor/__pycache__/async_compile.cpython-311.pyc,, +torch/_inductor/__pycache__/augmented_graph_helper.cpython-311.pyc,, +torch/_inductor/__pycache__/autotune_process.cpython-311.pyc,, +torch/_inductor/__pycache__/await_utils.cpython-311.pyc,, +torch/_inductor/__pycache__/bounds.cpython-311.pyc,, +torch/_inductor/__pycache__/cache.cpython-311.pyc,, +torch/_inductor/__pycache__/choices.cpython-311.pyc,, +torch/_inductor/__pycache__/codecache.cpython-311.pyc,, +torch/_inductor/__pycache__/comm_analysis.cpython-311.pyc,, +torch/_inductor/__pycache__/comm_lowering.cpython-311.pyc,, +torch/_inductor/__pycache__/comms.cpython-311.pyc,, +torch/_inductor/__pycache__/comms_debug.cpython-311.pyc,, +torch/_inductor/__pycache__/compile_fx.cpython-311.pyc,, +torch/_inductor/__pycache__/compile_fx_async.cpython-311.pyc,, +torch/_inductor/__pycache__/compile_fx_ext.cpython-311.pyc,, +torch/_inductor/__pycache__/compile_fx_subproc.cpython-311.pyc,, +torch/_inductor/__pycache__/compiler_bisector.cpython-311.pyc,, +torch/_inductor/__pycache__/config.cpython-311.pyc,, +torch/_inductor/__pycache__/config_comms.cpython-311.pyc,, +torch/_inductor/__pycache__/constant_folding.cpython-311.pyc,, +torch/_inductor/__pycache__/cpp_builder.cpython-311.pyc,, +torch/_inductor/__pycache__/cpu_vec_isa.cpython-311.pyc,, +torch/_inductor/__pycache__/cudagraph_trees.cpython-311.pyc,, +torch/_inductor/__pycache__/cudagraph_utils.cpython-311.pyc,, +torch/_inductor/__pycache__/custom_graph_pass.cpython-311.pyc,, +torch/_inductor/__pycache__/debug.cpython-311.pyc,, +torch/_inductor/__pycache__/decomposition.cpython-311.pyc,, +torch/_inductor/__pycache__/dependencies.cpython-311.pyc,, +torch/_inductor/__pycache__/distributed_autotune.cpython-311.pyc,, +torch/_inductor/__pycache__/dtype_propagation.cpython-311.pyc,, +torch/_inductor/__pycache__/exc.cpython-311.pyc,, +torch/_inductor/__pycache__/extern_node_serializer.cpython-311.pyc,, +torch/_inductor/__pycache__/freezing.cpython-311.pyc,, +torch/_inductor/__pycache__/freezing_utils.cpython-311.pyc,, +torch/_inductor/__pycache__/fuzzer.cpython-311.pyc,, +torch/_inductor/__pycache__/fx_utils.cpython-311.pyc,, +torch/_inductor/__pycache__/graph.cpython-311.pyc,, +torch/_inductor/__pycache__/hooks.cpython-311.pyc,, +torch/_inductor/__pycache__/index_propagation.cpython-311.pyc,, +torch/_inductor/__pycache__/inductor_prims.cpython-311.pyc,, +torch/_inductor/__pycache__/invert_expr_analysis.cpython-311.pyc,, +torch/_inductor/__pycache__/ir.cpython-311.pyc,, +torch/_inductor/__pycache__/jagged_lowerings.cpython-311.pyc,, +torch/_inductor/__pycache__/kernel_inputs.cpython-311.pyc,, +torch/_inductor/__pycache__/kernel_template_choice.cpython-311.pyc,, +torch/_inductor/__pycache__/loop_body.cpython-311.pyc,, +torch/_inductor/__pycache__/lowering.cpython-311.pyc,, +torch/_inductor/__pycache__/memory.cpython-311.pyc,, +torch/_inductor/__pycache__/metrics.cpython-311.pyc,, +torch/_inductor/__pycache__/mkldnn_ir.cpython-311.pyc,, +torch/_inductor/__pycache__/mkldnn_lowerings.cpython-311.pyc,, +torch/_inductor/__pycache__/mock_cache.cpython-311.pyc,, +torch/_inductor/__pycache__/ops_handler.cpython-311.pyc,, +torch/_inductor/__pycache__/optimize_indexing.cpython-311.pyc,, +torch/_inductor/__pycache__/output_code.cpython-311.pyc,, +torch/_inductor/__pycache__/pattern_matcher.cpython-311.pyc,, +torch/_inductor/__pycache__/quantized_lowerings.cpython-311.pyc,, +torch/_inductor/__pycache__/remote_cache.cpython-311.pyc,, +torch/_inductor/__pycache__/remote_gemm_autotune_cache.cpython-311.pyc,, +torch/_inductor/__pycache__/rocm_multiarch_utils.cpython-311.pyc,, +torch/_inductor/__pycache__/scheduler.cpython-311.pyc,, +torch/_inductor/__pycache__/select_algorithm.cpython-311.pyc,, +torch/_inductor/__pycache__/shape_propagation.cpython-311.pyc,, +torch/_inductor/__pycache__/sizevars.cpython-311.pyc,, +torch/_inductor/__pycache__/standalone_compile.cpython-311.pyc,, +torch/_inductor/__pycache__/stream_constants.cpython-311.pyc,, +torch/_inductor/__pycache__/stream_utils.cpython-311.pyc,, +torch/_inductor/__pycache__/subgraph_lowering.cpython-311.pyc,, +torch/_inductor/__pycache__/test_case.cpython-311.pyc,, +torch/_inductor/__pycache__/test_operators.cpython-311.pyc,, +torch/_inductor/__pycache__/tiling_utils.cpython-311.pyc,, +torch/_inductor/__pycache__/triton_bundler.cpython-311.pyc,, +torch/_inductor/__pycache__/utils.cpython-311.pyc,, +torch/_inductor/__pycache__/virtualized.cpython-311.pyc,, +torch/_inductor/__pycache__/wrapper_benchmark.cpython-311.pyc,, +torch/_inductor/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/analysis/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/analysis/__pycache__/device_info.cpython-311.pyc,, +torch/_inductor/analysis/__pycache__/profile_analysis.cpython-311.pyc,, +torch/_inductor/analysis/device_info.py,sha256=Jz63QjLCokPSAwksR6sMXoPBloZCDzja3POejAk_4w4,7419 +torch/_inductor/analysis/profile_analysis.py,sha256=eHojV4fdT7cOrgyCVqEW04jYqgyOCohs1i9p5MOIZkQ,27536 +torch/_inductor/analyze_preserves_zero_mask.py,sha256=cU6tCfAX4N0LbZV-KQVnsskwwdy16b2PBVjGnkowx8I,5477 +torch/_inductor/aoti_eager.py,sha256=hI2qXFWl07WbGKoRlV5yLO01sOWXz70e917KB9O7ZJo,12302 +torch/_inductor/async_compile.py,sha256=yybMjYUC05kKYRK4iIuuXoTKjozNpIbQcXL7V7f0GPU,29609 +torch/_inductor/augmented_graph_helper.py,sha256=LWEjisyBfWekwbQGUlPvB1UWQhR0HFbkdtRbX8RT1b8,11130 +torch/_inductor/autoheuristic/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/autoheuristic/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/autoheuristic/__pycache__/autoheuristic.cpython-311.pyc,, +torch/_inductor/autoheuristic/__pycache__/autoheuristic_utils.cpython-311.pyc,, +torch/_inductor/autoheuristic/__pycache__/learned_heuristic_controller.cpython-311.pyc,, +torch/_inductor/autoheuristic/__pycache__/learnedheuristic_interface.cpython-311.pyc,, +torch/_inductor/autoheuristic/artifacts/_MMRankingA100.py,sha256=wGZv4bFCwDlPlavBWdX4J9aKw9wRZaPTnRsqgyZnQog,28038 +torch/_inductor/autoheuristic/artifacts/_MMRankingH100.py,sha256=OQN-QYXZ2yYIZQUvGlEy-AC6fTzUVqgf9i_uxAp3ZdA,30662 +torch/_inductor/autoheuristic/artifacts/_MixedMMA100.py,sha256=odGa8EdsobbeiKoKDrzkoYmDYgICi62YDFtwHE5fzEM,7914 +torch/_inductor/autoheuristic/artifacts/_MixedMMH100.py,sha256=-ASwKTfLGuKkdoPYpszoAJOFhRZzQ5UxA_V9y-f1SXc,7863 +torch/_inductor/autoheuristic/artifacts/_PadMMA100.py,sha256=VJxvqunz6QFV25E4FrmM4O_2N2tB1svXYVgK-MIHIys,13991 +torch/_inductor/autoheuristic/artifacts/_PadMMH200.py,sha256=kwnmibiVDQ1Oh_Dl7nBd9hugqvqoJT1XrVfoFI6WpkY,8501 +torch/_inductor/autoheuristic/artifacts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/autoheuristic/artifacts/__pycache__/_MMRankingA100.cpython-311.pyc,, +torch/_inductor/autoheuristic/artifacts/__pycache__/_MMRankingH100.cpython-311.pyc,, +torch/_inductor/autoheuristic/artifacts/__pycache__/_MixedMMA100.cpython-311.pyc,, +torch/_inductor/autoheuristic/artifacts/__pycache__/_MixedMMH100.cpython-311.pyc,, +torch/_inductor/autoheuristic/artifacts/__pycache__/_PadMMA100.cpython-311.pyc,, +torch/_inductor/autoheuristic/artifacts/__pycache__/_PadMMH200.cpython-311.pyc,, +torch/_inductor/autoheuristic/artifacts/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/autoheuristic/autoheuristic.py,sha256=j1PbPzxk2c_wClAibii75zloeSGFJRzB7X1pQ82DPVU,12043 +torch/_inductor/autoheuristic/autoheuristic_utils.py,sha256=gRMDihFQA_YWc0gzNgxg8zSvgns2IK6oAiAQQ8grJes,11322 +torch/_inductor/autoheuristic/learned_heuristic_controller.py,sha256=cLy1FwX_1ByT1VPbv6UtNjFLF3g6NLR15YAZ0QljL48,4301 +torch/_inductor/autoheuristic/learnedheuristic_interface.py,sha256=_oEC4sIE5l50zhgPutB9O7ccRomvzOaaVmf-GBeC6DI,2641 +torch/_inductor/autotune_process.py,sha256=KDkwz8mSsBDzqjqeqAE42GepnS7Y1gz7BK3zl-RQxhI,49844 +torch/_inductor/await_utils.py,sha256=dSA8QYMzzVtu6YNiKL9FeZHewAL3GNywqxULkTn7jvY,5819 +torch/_inductor/bounds.py,sha256=f9-jBmF0NNgbf1YacT1nAKaYA9n-M1_PfRDOSXksKWU,9802 +torch/_inductor/cache.py,sha256=WCVejf3ofSc2lvXuFrE9BbQrTo1RSMeomWri1yVSizM,14823 +torch/_inductor/choices.py,sha256=_KB--mHn7fDrr_zJxf-8b7ys7_sqcWmTj2aiNCWVy0k,28350 +torch/_inductor/codecache.py,sha256=LHkyDox6yuItjyy-Mfs5BYoFgtk-tPZY4OrV2UXPnn4,184551 +torch/_inductor/codegen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/codegen/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/aoti_hipify_utils.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/block_analysis.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/common.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_bmm_template.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_flex_attention_template.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_gemm_template.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_grouped_gemm_template.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_micro_gemm.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_template.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_template_kernel.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_utils.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_wrapper_cpu.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_wrapper_cpu_array_ref.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_wrapper_gpu.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpp_wrapper_mps.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cpu_device_op_overrides.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/cuda_combined_scheduling.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/custom_extern_kernel_codegen.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/debug_utils.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/halide.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/memory_planning.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/mps.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/mps_device_op_overrides.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/multi_kernel.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/pallas.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/python_wrapper_mtia.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/segmented_tree.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/simd.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/simd_kernel_features.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/subgraph.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/triton.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/triton_combo_kernel.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/triton_split_scan.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/triton_utils.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/wrapper.cpython-311.pyc,, +torch/_inductor/codegen/__pycache__/wrapper_fxir.cpython-311.pyc,, +torch/_inductor/codegen/aoti_hipify_utils.py,sha256=ntqiznqR7z1Oa_GnR_IQcYOOYiM2F08jD-MaDZJ43TM,1440 +torch/_inductor/codegen/aoti_runtime/interface.cpp,sha256=VZl13TORNVPFYO5D-vYyESdD1_40k1pmjjAnjXVsYTo,17914 +torch/_inductor/codegen/block_analysis.py,sha256=L_LpGyiPFMYq1YQRBstzcRhTVIoEyOM5ChkZd-09kM8,8162 +torch/_inductor/codegen/common.py,sha256=-j-X0yzToqMLTIZuDtdEyPKDme_rcg8g5_75q74NLkk,107440 +torch/_inductor/codegen/cpp.py,sha256=5Fj5UKspoMVD3Xr31M7DCsU3TX22rYNCzoLx1j5uxiI,240135 +torch/_inductor/codegen/cpp_bmm_template.py,sha256=8ZiFbL7mci4iIuGk4klK4I1Y0LLIN7Y0OXiqLPs8b_k,9576 +torch/_inductor/codegen/cpp_flex_attention_template.py,sha256=Zn54ncdirUKnnkJDXR4znCNyWEY5oYJKqGb40oXpttc,41366 +torch/_inductor/codegen/cpp_gemm_template.py,sha256=jSil5_Vcm6AQtClK4xvUKxJe8XxC9gYJ7h4mQ8XmV1Y,78393 +torch/_inductor/codegen/cpp_grouped_gemm_template.py,sha256=mKVGfbTGjWa9fW-rVdlbxSIRw8SZfyuEH9M6UwnZH68,21146 +torch/_inductor/codegen/cpp_micro_gemm.py,sha256=JAYc6znHuoLPsdPfU1nOLkei4cNYX1kdU3X5WM-L294,77552 +torch/_inductor/codegen/cpp_template.py,sha256=c_R8X6s5kaNI86UrKe7FlBBqEKSRVpVfkFomNYnoExo,5028 +torch/_inductor/codegen/cpp_template_kernel.py,sha256=dobRQOlUs6bM5jPC8iTrjbcNVSLvQRBaxN3LFDSsIOU,26015 +torch/_inductor/codegen/cpp_utils.py,sha256=PsG7Z84Aain5rFCQ1W2FUWGHKGaz3X_URxINnRM_J-4,28409 +torch/_inductor/codegen/cpp_wrapper_cpu.py,sha256=9c-pWp9Hj3Hg-Q2yRkv9gmQETi2E0MA9VcnVN2HvPA4,134792 +torch/_inductor/codegen/cpp_wrapper_cpu_array_ref.py,sha256=QEK0AH9tl7NwBL_PdTCOsgSQVFP3ofl9q04ib9zYt9c,39018 +torch/_inductor/codegen/cpp_wrapper_gpu.py,sha256=ozjVFGow4o0GLji7LN9TGlgFsROBqgNjXPOOc9YM5uY,55244 +torch/_inductor/codegen/cpp_wrapper_mps.py,sha256=QGEG0Q0z_Va2q45MDRqsIYuvIYciMN1ePQZiRP-ZAdk,12405 +torch/_inductor/codegen/cpu_device_op_overrides.py,sha256=LhucRsCxEnVVckP2YLQy41J7lk-vXZnDNWAVUZKyvaQ,694 +torch/_inductor/codegen/cuda/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/codegen/cuda/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/cuda/__pycache__/compile_utils.cpython-311.pyc,, +torch/_inductor/codegen/cuda/__pycache__/cuda_env.cpython-311.pyc,, +torch/_inductor/codegen/cuda/__pycache__/device_op_overrides.cpython-311.pyc,, +torch/_inductor/codegen/cuda/compile_utils.py,sha256=J5ujGCmRrpzDj-ZV6Jz0cdu4pUpwwxiKlhO-Myvz5QI,9048 +torch/_inductor/codegen/cuda/cuda_env.py,sha256=BB3H-VM9lTs82v4l2YwQbCb2wAhqdhC91V_R4FBQgI8,1354 +torch/_inductor/codegen/cuda/device_op_overrides.py,sha256=bhCRrxuyhmW6A4mJZW42g_JdWkOUK3LwmxMHw55fZkQ,14804 +torch/_inductor/codegen/cuda_combined_scheduling.py,sha256=c6-a0pzgYOodkoSOXUfSugN_NOJdql8mHWQmTbBu0h8,7519 +torch/_inductor/codegen/custom_extern_kernel_codegen.py,sha256=IAdanV4i1K76cG6jPes_-dusUGMfeAm6YNntXO_p3WE,3066 +torch/_inductor/codegen/cutedsl/__init__.py,sha256=LCdU39rzLOP3q27650YlpGE3wo5SRH2IPdvEoybU6rg,164 +torch/_inductor/codegen/cutedsl/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/cutedsl/__pycache__/_cutedsl_utils.cpython-311.pyc,, +torch/_inductor/codegen/cutedsl/__pycache__/cutedsl_kernel.cpython-311.pyc,, +torch/_inductor/codegen/cutedsl/__pycache__/cutedsl_op_overrides.cpython-311.pyc,, +torch/_inductor/codegen/cutedsl/__pycache__/cutedsl_scheduling.cpython-311.pyc,, +torch/_inductor/codegen/cutedsl/__pycache__/cutedsl_template.cpython-311.pyc,, +torch/_inductor/codegen/cutedsl/_cutedsl_utils.py,sha256=1ecTHpFzE9tWjciDwBHMmpWqp3r_fKkpZAQmFkHkWIM,942 +torch/_inductor/codegen/cutedsl/cutedsl_kernel.py,sha256=uecEdr_e0jxgOuY8Dn4L_AvZn10iAq2B8XsOcRclE20,23678 +torch/_inductor/codegen/cutedsl/cutedsl_op_overrides.py,sha256=mJmziCx2Ncamjtrc_xHJ--feRFbOlUml1f8acoWGLcU,19850 +torch/_inductor/codegen/cutedsl/cutedsl_scheduling.py,sha256=jNKkpmy_Bc3u7Mw4d6D3otc1cPLyqWT3j33Mgym0r0Y,5310 +torch/_inductor/codegen/cutedsl/cutedsl_template.py,sha256=yT5mhVuZQXKF4gVNy9_4z2tlYJbIYv1yHnNz4vZGpCM,7611 +torch/_inductor/codegen/cutlass/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/codegen/cutlass/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/__pycache__/cache.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/__pycache__/gemm_template.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/__pycache__/kernel.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/__pycache__/python_evt.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/__pycache__/scheduling.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/__pycache__/serialization.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/__pycache__/template.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/__pycache__/utils.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/cache.py,sha256=HMkOGgI-sIjKves4cf0GMHjJJy0zhiVBuTdSoNc0xbI,3676 +torch/_inductor/codegen/cutlass/gemm_template.py,sha256=68xeUwBKJK5sSU4rbyCeIkNdpa9oSeJ5hxEl8rP2GVA,81309 +torch/_inductor/codegen/cutlass/kernel.py,sha256=3TV0GrfHnd3x_x12l2jGk3A-t1vV-SWEHmtpk22x6oM,24931 +torch/_inductor/codegen/cutlass/lib_extensions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/codegen/cutlass/lib_extensions/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/__pycache__/evt_extensions.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/__pycache__/gemm_operation_extensions.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/cuda/__init__.py,sha256=zIVQSgcv5C12tIsovmz3xGgBqoATUltX2c4-VG2e3tU,83 +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/cuda/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/cuda/__pycache__/cuda.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/cuda/__pycache__/cudart.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/cuda/cuda.py,sha256=MAbpo0K0ESplc0TsHAeuOE0ty5K1SUrROdDZ8sqVvWs,314 +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/cuda/cudart.py,sha256=pkNjquWqOMt8KAnOBkHbIxpgBmdWRZqXZbpQXS7ElK0,305 +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/pydot/__init__.py,sha256=AkdGQWifQLxFzB5RR5EcXNCg_AQX446DLfp-64Q7D4I,54 +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/pydot/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/scipy/__init__.py,sha256=SiRBwZsGLISQfmWy5-pkpAzRopAkCy9kzx8JOjV4YX8,55 +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/scipy/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/scipy/__pycache__/special.cpython-311.pyc,, +torch/_inductor/codegen/cutlass/lib_extensions/cutlass_mock_imports/scipy/special.py,sha256=NcOk2HGGaIsTnV0LPyRYnBXcR3WumPqGw2rKDT4dedk,54 +torch/_inductor/codegen/cutlass/lib_extensions/evt_extensions.py,sha256=28TqEkJrQwhYU7atnvC7xMgSnBpShMmFk9MCYJ06MAY,11442 +torch/_inductor/codegen/cutlass/lib_extensions/gemm_operation_extensions.py,sha256=cr-e7HhnwBELKUm90yBzpWpVSYpuwGcOnuzg2fRpufs,18972 +torch/_inductor/codegen/cutlass/python_evt.py,sha256=146u5IqKnktm17MZ_T52wc3e2jaGxajzPS4hcYo-mxQ,11833 +torch/_inductor/codegen/cutlass/scheduling.py,sha256=fNp8zanHeGAtRdqRi42AVLxGoTzaNiB6Mt5fwB_jKrQ,12037 +torch/_inductor/codegen/cutlass/serialization.py,sha256=tHx1fRCqahw5KpJ3Hyr9q0uh0TCi35e3CQsADzvwuQk,17761 +torch/_inductor/codegen/cutlass/template.py,sha256=Nwz3uShe3eG3PgFVNSIndR2FRFnMp10vdxdpeHcBLi0,13867 +torch/_inductor/codegen/cutlass/utils.py,sha256=DCOme7xnM4l-hJCwV7b7YDDzB0pTZF8KkznBzwnp_Y4,17813 +torch/_inductor/codegen/debug_utils.py,sha256=zIKAzL3oT5LmwDkVNRG5Eyg67NGMgqInb_S6Q6nq3Yc,11365 +torch/_inductor/codegen/halide.py,sha256=yi89Gxcb1NlqgEpYkg1Rx0di_jxcslO6FoLXCk0hkJk,66463 +torch/_inductor/codegen/memory_planning.py,sha256=5KzpVP2BHuiucaz_32Bdq77ZfjdTYotUkHmecY7tmuQ,26578 +torch/_inductor/codegen/mps.py,sha256=2Dq4Ahosdi34NLn9rxvVj6gLj-syfx36azA_io80fhM,47962 +torch/_inductor/codegen/mps_device_op_overrides.py,sha256=AY7QsOH2mELU0HsZz6B6ygxdSGwVnvS7v4ypb8KRXlo,671 +torch/_inductor/codegen/mtia/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/codegen/mtia/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/mtia/__pycache__/device_op_overrides.cpython-311.pyc,, +torch/_inductor/codegen/mtia/device_op_overrides.py,sha256=3pP0peusaHi8fkuf9ZHnwFE88KYjH-0YJcAT07ZCFuA,637 +torch/_inductor/codegen/multi_kernel.py,sha256=G_9V9Ev8IzUDk0HPR2AWunc9_rrs8zzpA6bU9YSrZmc,23355 +torch/_inductor/codegen/nv_universal_gemm/__init__.py,sha256=9D2vBtuP0_NBGKGcFXxB8dK68mquVbiSHe3JqrHbYG0,416 +torch/_inductor/codegen/nv_universal_gemm/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/nv_universal_gemm/__pycache__/kernel_cache.cpython-311.pyc,, +torch/_inductor/codegen/nv_universal_gemm/__pycache__/nv_universal_gemm.cpython-311.pyc,, +torch/_inductor/codegen/nv_universal_gemm/__pycache__/nv_universal_gemm_kernel.cpython-311.pyc,, +torch/_inductor/codegen/nv_universal_gemm/__pycache__/nv_universal_gemm_scheduling.cpython-311.pyc,, +torch/_inductor/codegen/nv_universal_gemm/__pycache__/nv_universal_gemm_utils.cpython-311.pyc,, +torch/_inductor/codegen/nv_universal_gemm/kernel_cache.py,sha256=iuTLqH6jwZHAm5jY7Mf9IOtslHvDhitDjfyKvovxC5M,2717 +torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm.py,sha256=d6JVTr0hMhUphCo4DjNzx5QIOgZhtZaMaEgyqttqGvk,18668 +torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_kernel.py,sha256=hs1HUW-x3euWWhAvMCMYPh3KHdrKn7K1covYwb8Qqro,11240 +torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_scheduling.py,sha256=qhlQFG2z5-4WFisypVELfWXMkoNBpcDEJOe75Bnctys,5185 +torch/_inductor/codegen/nv_universal_gemm/nv_universal_gemm_utils.py,sha256=mNkSqUjrkIXWw0tK6c5kM-dvCz6G6ZSizpoi4hq0c9I,1417 +torch/_inductor/codegen/pallas.py,sha256=K5XIlR-7tbjfR5MKuGIng1j5CrcsKvnBV9FxDPxM1IM,200891 +torch/_inductor/codegen/python_wrapper_mtia.py,sha256=XEIoUYldr_tZ-E1301BGTvgSsari-L2ygLOmouNNCww,988 +torch/_inductor/codegen/rocm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/codegen/rocm/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/ck_conv_template.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/ck_template.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/ck_tile_template.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/ck_tile_universal_gemm_template.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/ck_universal_gemm_template.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/compile_command.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/rocm_benchmark_request.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/rocm_cpp_scheduling.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/rocm_kernel.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/rocm_template.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/rocm_template_buffer.cpython-311.pyc,, +torch/_inductor/codegen/rocm/__pycache__/rocm_utils.cpython-311.pyc,, +torch/_inductor/codegen/rocm/ck_conv_template.py,sha256=j6EuqXegHn8_EKas2uQOZVlg3bJq_YVY-_dNciVuivQ,24936 +torch/_inductor/codegen/rocm/ck_template.py,sha256=eNUthfvZo-H9Ayvg-1yJUfVZFhF3eBJHxESY5oWZyts,3695 +torch/_inductor/codegen/rocm/ck_tile_template.py,sha256=Qm1oW_WPVexswfI8HwL9lgfYUxEwhd4ehmDKz3AOHik,1591 +torch/_inductor/codegen/rocm/ck_tile_universal_gemm_template.py,sha256=xj3bXtb-mEIAVQmp7TkLEGqA1yyL5oMV55y8L2_O2HI,37033 +torch/_inductor/codegen/rocm/ck_universal_gemm_template.py,sha256=xXyV6cUoddgTdEEzhDWSY4N0OCL8R0Z6FDamgcCxyUg,39641 +torch/_inductor/codegen/rocm/compile_command.py,sha256=mK5kZTig4Ld8Sj1PMDdIZJhobmGnQwFJJ-dD_1X7okU,4628 +torch/_inductor/codegen/rocm/rocm_benchmark_request.py,sha256=9Dxy73lx5UId2zNSTf5T9oyeNqzd8rrXFzW9yn6bf-0,5029 +torch/_inductor/codegen/rocm/rocm_cpp_scheduling.py,sha256=MMO-IotoXV99j_kr81raYYs74tazaJLX34nzMTg4N-0,3850 +torch/_inductor/codegen/rocm/rocm_kernel.py,sha256=ngiOD-ERJ3WViyJCV8NVwYM7lkgvSy-Aco-5Eaq38mE,10450 +torch/_inductor/codegen/rocm/rocm_template.py,sha256=N6axOREeEk38YCNkbaJzXXpu2FS9a_Q3L4m7gnEFqSg,6634 +torch/_inductor/codegen/rocm/rocm_template_buffer.py,sha256=UMhXlEMC2acFTLJqJwKx1btyB7uuxOe_Ftr364CzNLw,827 +torch/_inductor/codegen/rocm/rocm_utils.py,sha256=_tYJSiuN0xtpJvqM9xycWNsRRC0UzuhGxIMwp0gIHxw,336 +torch/_inductor/codegen/segmented_tree.py,sha256=BOqdDKVQtRK2h0hc6T6zkj6occ6y79Rzvdaf-xmh60k,8177 +torch/_inductor/codegen/simd.py,sha256=JdS3hRKR4rEKUjcDazg77ndnyQLlS4h7PdXz4H4J0s8,127018 +torch/_inductor/codegen/simd_kernel_features.py,sha256=HwZCDCNRbBqUM1UMxYE7HFxh6JimzdcnSge2XdouA3I,24586 +torch/_inductor/codegen/subgraph.py,sha256=KgixXbwcccviLPq6cBBi68ZN8L3jiilxtgMSjb7pAr8,24203 +torch/_inductor/codegen/triton.py,sha256=gWok7P7xa3Rg90Tp5fZ6IxPJyZMFSkKHRwueuD8czQ0,278256 +torch/_inductor/codegen/triton_combo_kernel.py,sha256=l3JVStWdXtKxlVjjkN1HTCPS_PrC9edNiRcyzaBoPDo,51598 +torch/_inductor/codegen/triton_split_scan.py,sha256=1RZWmgEVdqyf77LkdTi0bP-m3qk3F6it29Sg9NR10P0,7670 +torch/_inductor/codegen/triton_utils.py,sha256=jP-WkQRoWqdfAt9lGrNLmBQ_7O0jfAk9O83aFx5ahis,10894 +torch/_inductor/codegen/wrapper.py,sha256=_oolt_aaZfWXTxF2YiYDzUrv4eu1Nlt9xrzu51l2s3k,174350 +torch/_inductor/codegen/wrapper_fxir.py,sha256=H6Yh5uG9eqWru5Pc8nQfUiEKcRRtq2iX_aTYemMh4F4,46320 +torch/_inductor/codegen/xpu/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/codegen/xpu/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/codegen/xpu/__pycache__/compile_utils.cpython-311.pyc,, +torch/_inductor/codegen/xpu/__pycache__/device_op_overrides.cpython-311.pyc,, +torch/_inductor/codegen/xpu/__pycache__/xpu_combined_scheduling.cpython-311.pyc,, +torch/_inductor/codegen/xpu/__pycache__/xpu_env.cpython-311.pyc,, +torch/_inductor/codegen/xpu/compile_utils.py,sha256=op2hFf0rNcv5o5mlUQZIDHc_DCqBB-aXl8LS0syzqx0,4088 +torch/_inductor/codegen/xpu/device_op_overrides.py,sha256=MMpBqzsxeQgdd9tImLEUqxxVeUWl23wVdDyuUJyFeYk,1849 +torch/_inductor/codegen/xpu/xpu_combined_scheduling.py,sha256=dQ28ZcNMbYLAZyQ6gAeGSotEBwa71XsRhlNXlbAEwxg,4735 +torch/_inductor/codegen/xpu/xpu_env.py,sha256=_PsKVlhLOP229ZgfJkKrz43NGkRglYetn_czpbG1dvw,902 +torch/_inductor/comm_analysis.py,sha256=ZubYn4iYadACiBGuA4fYYMyIdDuqC-2Ui2b-TrR3uNU,18147 +torch/_inductor/comm_lowering.py,sha256=8DxGSHT1KqVMyFTxJhPyzevLnDZT_EG92BsdDlxH_Lc,29000 +torch/_inductor/comms.py,sha256=VCtOAz84LcUP_5RvdBWFDItJCGjNRy0PGBX4yJxXRpo,104964 +torch/_inductor/comms_debug.py,sha256=Y_qM_qrvN4mQhSiAddBPwQCa-KX8EyB1Ho0rDP2mRwk,4214 +torch/_inductor/compile_fx.py,sha256=4r9tayXLmDI8Y0ayHGB1qnI_w-oYzDK-Zp1bRuV1odg,128771 +torch/_inductor/compile_fx_async.py,sha256=s-g98pt2_qiYHGz5RnqargZi1XyglTL9LAUTriAi0Jo,14887 +torch/_inductor/compile_fx_ext.py,sha256=VM1jDkWoLXNkTOMvTzzgqoFu7Qm6GPoUP2f3wP1iWac,24672 +torch/_inductor/compile_fx_subproc.py,sha256=QpkT9JjsaZPk2wwF1hVumjUcn6kD1FK5sGwXVF3B2SQ,3158 +torch/_inductor/compile_worker/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/compile_worker/__main__.py,sha256=Yj1LnCvSJWe_Y8LWarFQLceqCmadcX7RhqP9uBS-zYA,2245 +torch/_inductor/compile_worker/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/compile_worker/__pycache__/__main__.cpython-311.pyc,, +torch/_inductor/compile_worker/__pycache__/subproc_pool.cpython-311.pyc,, +torch/_inductor/compile_worker/__pycache__/timer.cpython-311.pyc,, +torch/_inductor/compile_worker/__pycache__/tracked_process_pool.cpython-311.pyc,, +torch/_inductor/compile_worker/__pycache__/utils.cpython-311.pyc,, +torch/_inductor/compile_worker/subproc_pool.py,sha256=7PpSDYrRnVaD48msK9i86lOuq68gXxshzPei0L8EOV0,17443 +torch/_inductor/compile_worker/timer.py,sha256=8VuP5_NwEN-zt_7vnXNSAiX3RUFUhaUMuxfLg0x16NA,1818 +torch/_inductor/compile_worker/tracked_process_pool.py,sha256=kw1Ur8UzCdTf91lbKnbNe_O4Ua5jQb58UKOINNRdku0,3674 +torch/_inductor/compile_worker/utils.py,sha256=y5eRcA1k1h1PxpEDhpndNKVZngLleEHcOpoYcnDyuKE,1595 +torch/_inductor/compiler_bisector.py,sha256=63lBK4W7iH9dXgkHgV_5GfDicIBVszTazjsZ4jfAqUA,27642 +torch/_inductor/config.py,sha256=EzTT_3mnk7k9pAX56mOp7qoU7q_bhFW9Lj9jwcTl3Bw,109332 +torch/_inductor/config_comms.py,sha256=KwIjE-bJGkaqi3fxiSbNNRYnBsilaUltR022M3ShPQU,3052 +torch/_inductor/constant_folding.py,sha256=1tpPJcLGz35B96SXWm8uoqYidQ3ch2kpoCp8JfmuEug,15639 +torch/_inductor/cpp_builder.py,sha256=r81b6O8tLxgEHA86v5bSDfP5mMzaacecujzjAILmWvk,99972 +torch/_inductor/cpu_vec_isa.py,sha256=sQwNDLO-xsxeXZfK9JvjOyWixd0QZrAF0b53XQtm7vo,18103 +torch/_inductor/cudagraph_trees.py,sha256=KFuhwa2ywVSUd-LCIkSiTp2NWSbf6scrcnnsajYbf08,110992 +torch/_inductor/cudagraph_utils.py,sha256=8a3TmP3Brs1A8ZKbsJ_SzrB9z8OCS5zlm9Z8sPdGyPI,18356 +torch/_inductor/custom_graph_pass.py,sha256=l2UH4rvqnwPr44GgZGuQ0tylMasfJq_mdwpD9T1yfNs,10598 +torch/_inductor/debug.py,sha256=yj0FJJZRo1K6GqDxM6pWwsHw--ciVpedcHR3uQOdnuQ,47127 +torch/_inductor/decomposition.py,sha256=ewkuxCPjgxNjA-uHHix2QSb48k_PUxCRDpvpNhwf8Yo,43711 +torch/_inductor/dependencies.py,sha256=1QNzjFNlcMA7X5-46pdGmqHByN34Tb_YwZIyCgtTR4M,32302 +torch/_inductor/distributed_autotune.py,sha256=eVOhuBuKL-r6SCVUOQ3HN967W-1vvwMP8laiQydzWdM,12303 +torch/_inductor/dtype_propagation.py,sha256=8aW2oMdw4jZ3U3_VH-th0SFJkYr7SahKkWczdLZrGRw,14315 +torch/_inductor/exc.py,sha256=fEd61Ca3fO4Tpd-Rz6f60JuEaKwB0_5y-SOlZ5tHSq0,5002 +torch/_inductor/extern_node_serializer.py,sha256=tvVk5NzYVvqTI0v-k2hQikc8h6EsXazK1NQ2soHZ2yc,830 +torch/_inductor/freezing.py,sha256=Gw98XWkl3-S3LKhbRP3LT7p40YhUG3WEOFjvwKS7IYg,11077 +torch/_inductor/freezing_utils.py,sha256=0kQfN4H2yEI_eNXINqa0ws5Q-mm9EVT6ARcAKrz3P80,1268 +torch/_inductor/fuzzer.py,sha256=xCI5ldbNxKjLXFBmMd2zBfNyC5cgKYHaX8tZ0H_g2tM,37785 +torch/_inductor/fx_passes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/fx_passes/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/apply_gumbel_max_trick.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/b2b_gemm.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/binary_folding.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/bucketing.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/control_dependencies.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/ddp_fusion.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/decompose_mem_bound_mm.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/dedupe_symint_uses.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/efficient_conv_bn_eval.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/freezing_patterns.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/fsdp.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/fuse_attention.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/fusion_regions.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/graph_view.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/group_batch_fusion.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/joint_graph.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/memory_estimator.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/micro_pipeline_tp.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/misc_patterns.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/mkldnn_fusion.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/node_runtime_estimation.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/numeric_utils.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/overlap_manual_scheduling.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/overlap_preserving_bucketer.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/overlap_scheduling.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/pad_mm.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/post_grad.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/pre_grad.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/quantization.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/reduced_atomic_contention.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/reinplace.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/replace_random.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/split_cat.cpython-311.pyc,, +torch/_inductor/fx_passes/__pycache__/spmd_check.cpython-311.pyc,, +torch/_inductor/fx_passes/apply_gumbel_max_trick.py,sha256=DsbKM0ZVUOcTyk1T740CUYSBFtNVCkZMhbLPlY1hzSw,1847 +torch/_inductor/fx_passes/auto_chunker/__init__.py,sha256=pqRDYIVNxTGFK5R54Iw0beoE31PoYA2A8fP4937t87w,2708 +torch/_inductor/fx_passes/auto_chunker/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/fx_passes/auto_chunker/__pycache__/applier.cpython-311.pyc,, +torch/_inductor/fx_passes/auto_chunker/__pycache__/common.cpython-311.pyc,, +torch/_inductor/fx_passes/auto_chunker/__pycache__/core.cpython-311.pyc,, +torch/_inductor/fx_passes/auto_chunker/__pycache__/propagate_scale_by.cpython-311.pyc,, +torch/_inductor/fx_passes/auto_chunker/__pycache__/propagator.cpython-311.pyc,, +torch/_inductor/fx_passes/auto_chunker/__pycache__/utils.cpython-311.pyc,, +torch/_inductor/fx_passes/auto_chunker/applier.py,sha256=4pEY1ilC7QOPgoc5Jx0VH0nTKgqMz1rvAm2UbFRBYpU,17378 +torch/_inductor/fx_passes/auto_chunker/common.py,sha256=jPYW-i1jXvxzJJlPsgbGVb7Ho-m1asA-ebz2-LjgL_o,2240 +torch/_inductor/fx_passes/auto_chunker/core.py,sha256=17BZwtJQ4jTSnC1PFV-mV5HZbLkPnaL6yqj4KVezMDo,6395 +torch/_inductor/fx_passes/auto_chunker/propagate_scale_by.py,sha256=ZaFcUmzBd7Q6XkwfJbjdliE5gQ2WQpWqRTfKJc-8Axs,7667 +torch/_inductor/fx_passes/auto_chunker/propagator.py,sha256=VYmiz8bS1z7f2CpJTwA8I7DvIY1InZdPxGuWTAKD874,30663 +torch/_inductor/fx_passes/auto_chunker/utils.py,sha256=Hc3QnLIMewNdpjBdChw-4jiNCecQSxTJCb66yNu5aw8,5173 +torch/_inductor/fx_passes/b2b_gemm.py,sha256=z6qwi1UQD4r3qv7tbLEcRuw2y0MNsHi5VrW88aL8hxI,25633 +torch/_inductor/fx_passes/binary_folding.py,sha256=GftJvNpuD8uJh7vNYnV4QXLXL4sq5S680D6Y0IYndmE,19830 +torch/_inductor/fx_passes/bucketing.py,sha256=NWfwn16jcSeACjmke-n9rIaxSEUgOA-infuQbj0wBE8,46392 +torch/_inductor/fx_passes/control_dependencies.py,sha256=q1BNaiEkzkIdScUdTuziOvGcG8f36CkpPRNLDKCfOnY,10342 +torch/_inductor/fx_passes/ddp_fusion.py,sha256=trXcn6RdhE3VXpzw6Ti1ik1bA_5CJI3SgrbQhjAeQY8,21116 +torch/_inductor/fx_passes/decompose_mem_bound_mm.py,sha256=0THzgsrRPfs3JAmDFeHGeIQ5zKZL8diQ1RBPc-ent7E,11217 +torch/_inductor/fx_passes/dedupe_symint_uses.py,sha256=Bfuj7IhzY3DuixClyuziJM_jJ_w_zdfxsi7ULJcIxjk,2505 +torch/_inductor/fx_passes/efficient_conv_bn_eval.py,sha256=GtDFXHZ8S9Oo7Ys1iG5sXABxOc6Ey33GnZsAVW9a6LQ,14114 +torch/_inductor/fx_passes/freezing_patterns.py,sha256=OsRRUNr2XbXUet2EazlsjEVTVBgNheMDYCs22dJUXos,9850 +torch/_inductor/fx_passes/fsdp.py,sha256=-pFBpGNSFKAPd2haNXlUFMK-sqbHSeGNonpre4b4Ots,3551 +torch/_inductor/fx_passes/fuse_attention.py,sha256=2i7VsKbHrLVU3Hr6IQXDe7JbYpmGRwOLHWHMTFebvms,44010 +torch/_inductor/fx_passes/fusion_regions.py,sha256=VQZAZP5_DFH3V_JY7kxFBEFlcN5ZXt7Sptb2cqm9Ld8,11590 +torch/_inductor/fx_passes/graph_view.py,sha256=48ragODf9N7apGrt7rbLFq5A4vq4BfzHq6aeosxEtMQ,8992 +torch/_inductor/fx_passes/group_batch_fusion.py,sha256=WTvPM2GNOYzD3fqy1frSIgFxpULW4ynRHD344vDic1Q,59375 +torch/_inductor/fx_passes/joint_graph.py,sha256=zNeA6WvieRo4YRGXCMVPH85AaptzD_Winajb3_0TDnc,40362 +torch/_inductor/fx_passes/memory_estimator.py,sha256=cpPgFUlFoKyeIsqdYauLE9FFdyDxcDuc7K6eS_5BadQ,17146 +torch/_inductor/fx_passes/micro_pipeline_tp.py,sha256=uRr1yn9SLAazplpvicEHB6p2iDeZIk-XLwd2dfFlVLQ,40346 +torch/_inductor/fx_passes/misc_patterns.py,sha256=PoZU_iIc1zW4-33Tpb3wMUxS6Oske1SR4oY3E8ck_to,8362 +torch/_inductor/fx_passes/mkldnn_fusion.py,sha256=ui6VKbaPRnCJ9lEpAdgDPw66PWIXQeq9OY7lPBf7hng,62579 +torch/_inductor/fx_passes/node_runtime_estimation.py,sha256=7DEiOI6-u-wgTjga9V3s9-nNpiMFai6XhPSSOcIVNus,13274 +torch/_inductor/fx_passes/numeric_utils.py,sha256=D-tB9vsKK1V1lR-7UCXBJailxQuHPkkKmuddOYrfnPk,7285 +torch/_inductor/fx_passes/overlap_manual_scheduling.py,sha256=S9hNQS6xlOzwdjCg7jcS25VehqvcHUDK1uJa6aRTp9U,15520 +torch/_inductor/fx_passes/overlap_preserving_bucketer.py,sha256=wPdTYRIY_9VkqJgikuXw24DNBKfFNr1aVmFUfH-bs90,44036 +torch/_inductor/fx_passes/overlap_scheduling.py,sha256=XB8J9dMCAlxBhnilpPNJ8dpcZsTemxugFXQ3dbAoeYc,74361 +torch/_inductor/fx_passes/pad_mm.py,sha256=ZJB62GwOokW_bQSHSL0G0L_Pt95cqbsJgk93EARVqrY,32700 +torch/_inductor/fx_passes/post_grad.py,sha256=tw5MHdH5iSFeLsmNYboLxb_7UbPXrRS5I4lRZj0qBto,73610 +torch/_inductor/fx_passes/pre_grad.py,sha256=2BbuLY9KacQLc8scR3grSmZT1XM7FsSe-FCoHH5Q2Dg,31751 +torch/_inductor/fx_passes/quantization.py,sha256=ozkuab3CNd0zqeknxTl9X_zkpRvtIzjGU6GmcBQGGY4,144005 +torch/_inductor/fx_passes/reduced_atomic_contention.py,sha256=DT0iXVy7Ntb3ul7oYnfftLOP2dikW8VJLF6BV5qgPrw,13477 +torch/_inductor/fx_passes/reinplace.py,sha256=kTd5jdqpp_TVB1KWRu7opxTYwE_6CBkxvS2sWuzNcXo,38142 +torch/_inductor/fx_passes/replace_random.py,sha256=3mRK1hTs7GSs7CXDvponVrQr49pv17bqpBRMCKcgrpw,8352 +torch/_inductor/fx_passes/serialized_patterns/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/fx_passes/serialized_patterns/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_1.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_10.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_11.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_12.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_13.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_14.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_15.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_16.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_17.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_18.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_19.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_2.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_20.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_21.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_22.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_23.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_24.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_25.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_26.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_27.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_28.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_3.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_4.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_5.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_6.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_7.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_8.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/_sfdp_pattern_9.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/addmm_pattern.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/bmm_pattern.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/__pycache__/mm_pattern.cpython-311.pyc,, +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_1.py,sha256=jLk3u9vyyPhzzmJCf82H-lEa_ycI9FGR2F1e14BF7JA,11177 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_10.py,sha256=bQt8DWmSMYkw1UmqS1kb6OHenKl549UFKueC14Cstwc,14566 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_11.py,sha256=YzKAdp1wKREpwh-ma-MPjMKRAImnfIFpWrqbpBtJBPc,13987 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_12.py,sha256=6o07OMAP3P4yuOSrmsW_AsSq9gYeTVD2SkxJh8xQoZc,15257 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_13.py,sha256=yQOZ4yfymHPP_gNISoxq4pOt3xPic8eJY0InMr9hsPg,7862 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_14.py,sha256=p729osgik7UB0ySzKMcthIj9Pk3rqh6OdP6wKZFNqvA,14323 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_15.py,sha256=MdfVHtUJn0ZzgPsDoLIsj7Sla3HnflFCIELfS1uAYy0,16213 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_16.py,sha256=-1tQU2MyC9heRGD6sOibLRER_Mlzx9civVLrREdQSMo,43596 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_17.py,sha256=6MsCMyU-wCAUxWOmGt6fTqmbmTmm1BPeb3QhI8SDQuc,17441 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_18.py,sha256=s_OuvfnHfCbwfWtG9QqMwuRx0uYAbvbPs6swWL58MUc,32880 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_19.py,sha256=80S_fYu8eFi9FhufWN__ZeG3EseD_XMarFB7TFpryhk,14118 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_2.py,sha256=-QeXpOThLf9w6HsKJT5nJdCfyXRK4sgxE6Jy4v87tdw,11187 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_20.py,sha256=AMXQ7encxsqVvXc9a7ehsHSV9_WBEFvsKtABpbpcjeQ,17441 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_21.py,sha256=JuyA3MEoqmEdHBKS8R4snC91h-1wJA4ryoCuuPToMdU,26370 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_22.py,sha256=5uQr9YnaGegg59Hu8JHys5TZ55WQfU0qHYOre7KROOc,27046 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_23.py,sha256=_k2LHqN7_IMqlAtoXicKTu_c5t1orr_VRVRbl2BueaE,26890 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_24.py,sha256=02rc69wQcftSytBxCl5qbEG-MPQEpEi-es3NBNKxKpA,9115 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_25.py,sha256=QoJdK9Pi-atiGTQMFH6lhGnnPA2pQFmz6nLtPFVmkT0,27818 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_26.py,sha256=tFxqio8yOpYhDROKzBZnJ-Z8ID2kOtTUfiU6jVrQplM,28494 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_27.py,sha256=r5fZLxsUTPkVPU4B0SLGX48twqPaOeqIusB3R7Ht2VQ,27838 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_28.py,sha256=V3AZYXt4_0Z2KjVKNsltCtubC_ilZj_qPpqMU9WPnqg,13683 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_3.py,sha256=GAyzPOGSMzkdAxXkk3qpf-8y8GrWY-Wetzpy1aJU4oI,12447 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_4.py,sha256=x3mGjvSYS7hImiBFUpXbySYVw8ZNVjCyqjjTrNIiHsg,12411 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_5.py,sha256=bmujzU6QCO7aYxqtE8J8esBq9lsXr6OJZi8rXpPwsFU,11513 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_6.py,sha256=OitLrZvDQcqBzKGxOEmnvXuIUpf-K946j1WeA4fr7m4,12741 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_7.py,sha256=4E2j9W9BLOzYNw4VdYilr2ItK8ZVvgUoA6lXlfgVdrM,15786 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_8.py,sha256=FQhPAUET_ADKZDAcbT2onqhSoum9Ac1M9TkuvPw4Ni4,14554 +torch/_inductor/fx_passes/serialized_patterns/_sfdp_pattern_9.py,sha256=Ai-262aREFemsZQpk_tA5GWMrMW7xW_m7PbPa_9jcL0,15794 +torch/_inductor/fx_passes/serialized_patterns/addmm_pattern.py,sha256=9cl9L8L2CDu--cjIWRF63Ngl7YX9F38gGXCeEz3r0bY,1858 +torch/_inductor/fx_passes/serialized_patterns/bmm_pattern.py,sha256=6PYPyMz8dOVNYEKIIR67MvfdWL6TlF0xC8VnzsbHBLo,1272 +torch/_inductor/fx_passes/serialized_patterns/mm_pattern.py,sha256=MIExrwuf-ZmoMuOiUU3SENtMW8AjnW2yL9Q-kKsiqs0,1260 +torch/_inductor/fx_passes/split_cat.py,sha256=wX4o70ARFScNJxs3aCoj8SszkJR4T64bDJlVoU0lyUo,121569 +torch/_inductor/fx_passes/spmd_check.py,sha256=Z5nBoH8K6lohP4k9L8RUOdhAWCTHnxQ-e9HVY4EKyfc,7507 +torch/_inductor/fx_utils.py,sha256=70JdxsffpIrbMC11JYTQdgYN91OOA9qN-FzqWTmz7-U,13269 +torch/_inductor/graph.py,sha256=Bm8EnAu6S19qjzPbjb9wLvk0Cp3liFMvQpWstjEBdno,117489 +torch/_inductor/hooks.py,sha256=kUoOn0_wibHB8Cn9iSU427vUlTypKYh2J9DhwkKosEQ,666 +torch/_inductor/index_propagation.py,sha256=_L8UdV4Wl3OB3Oo2vSlF6gD4efpVpmBtS4NkARR1u-s,13280 +torch/_inductor/inductor_prims.py,sha256=EnafmiuPFR3t6fRS9C6JGDb00WOSw09JYJBjhok4HmQ,12482 +torch/_inductor/invert_expr_analysis.py,sha256=Fuwi_Avy3UjQ7kod4t2K-GSv8s8yXIWbtIDpMgYfBNE,6905 +torch/_inductor/ir.py,sha256=sZn9xqWv8dLVcO1xpQMAliEJXetksa7oBUC_ifT8Zew,392581 +torch/_inductor/jagged_lowerings.py,sha256=JtD7PHWQ_7IK25mRiH4NAr6TshXUUmPaesddnGHoMhI,9014 +torch/_inductor/kernel/__init__.py,sha256=4dfIYq5T99Rq8ncmT682RzjtYYLOYEwUVRzIbeyGysA,46 +torch/_inductor/kernel/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/kernel/__pycache__/bmm.cpython-311.pyc,, +torch/_inductor/kernel/__pycache__/conv.cpython-311.pyc,, +torch/_inductor/kernel/__pycache__/custom_op.cpython-311.pyc,, +torch/_inductor/kernel/__pycache__/mm.cpython-311.pyc,, +torch/_inductor/kernel/__pycache__/mm_common.cpython-311.pyc,, +torch/_inductor/kernel/__pycache__/mm_grouped.cpython-311.pyc,, +torch/_inductor/kernel/__pycache__/mm_plus_mm.cpython-311.pyc,, +torch/_inductor/kernel/bmm.py,sha256=5FgL7LcVqGpMpepXW3x-351QUneHz5mfUMakBA-jrdc,10153 +torch/_inductor/kernel/conv.py,sha256=UTDn29ps3latwZf1XquK2wxUM8TAaAzO05hZ8fMhJKs,23386 +torch/_inductor/kernel/custom_op.py,sha256=q3Kfea_XE-HaqwEKIjYdAatrVAl3POZlOIUlT28GIQ8,49422 +torch/_inductor/kernel/flex/__init__.py,sha256=_ca2fmTtYz5clo4uzQy2tbaLmypoau8LwEBJe8o5Kd0,153 +torch/_inductor/kernel/flex/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/kernel/flex/__pycache__/common.cpython-311.pyc,, +torch/_inductor/kernel/flex/__pycache__/flex_attention.cpython-311.pyc,, +torch/_inductor/kernel/flex/__pycache__/flex_cpu.cpython-311.pyc,, +torch/_inductor/kernel/flex/__pycache__/flex_decoding.cpython-311.pyc,, +torch/_inductor/kernel/flex/__pycache__/flex_flash_attention.cpython-311.pyc,, +torch/_inductor/kernel/flex/common.py,sha256=IQsO28smftUPZnYn34Swv0AqhDZ_X2KEYSEpZX2Dftk,11625 +torch/_inductor/kernel/flex/flex_attention.py,sha256=mH4iJnRqWpBCjfQgl9eQuwqh3-i0g73r2s6ekUKKXXU,38596 +torch/_inductor/kernel/flex/flex_cpu.py,sha256=nElaF3mWzOSU7WKQ7XfHmRfFErxgCqk2CaUqqjNMJOI,12226 +torch/_inductor/kernel/flex/flex_decoding.py,sha256=QjvfoTyVsAbTOzoKcgCHa4Cznl828zwErEM9v4ieW3E,15810 +torch/_inductor/kernel/flex/flex_flash_attention.py,sha256=jj8XBWwhJgFUCFId1SjBm_ZukVGf5Zq6_ZmOOpmrvc8,23361 +torch/_inductor/kernel/flex/templates/common.py.jinja,sha256=hXfMKwQFqcHBQFDv_T4kluovQYlwiy42SkYRgKtM_To,7074 +torch/_inductor/kernel/flex/templates/flash_attention.py.jinja,sha256=2DrEMG5o9QoCViz1lyAsTB8-2UbNkfO8xPMSXbi-sRg,2891 +torch/_inductor/kernel/flex/templates/flash_attention_backward.py.jinja,sha256=MoaLIllHNSwCoF3_sCmr35ngyZUv8EFU0OjRhbqUeqc,3648 +torch/_inductor/kernel/flex/templates/flex_attention.py.jinja,sha256=7IhvBAr6HYuaEigdRDhROCiZ-U61Yk2qLmV4IbYkm-0,9307 +torch/_inductor/kernel/flex/templates/flex_backwards.py.jinja,sha256=pfSQGO0LcN3-rCeT5484JwbmoxxwtxDrvgG9FKeNwIU,30142 +torch/_inductor/kernel/flex/templates/flex_decode.py.jinja,sha256=tEjiRfpbTx9z8YPHoBxWEiscNEpkwwEcnBwHkxaQQ6g,11070 +torch/_inductor/kernel/flex/templates/utilities.py.jinja,sha256=0Siotkd0DpHO797fORLKuikl4zg68KG1ZzL-7Y-mm_0,2258 +torch/_inductor/kernel/mm.py,sha256=ztwIfvpQxa7r5olggzxyEm-GbCCfYJ_IzDGsmVxKulg,38873 +torch/_inductor/kernel/mm_common.py,sha256=1cANcf8roYMsUX1QhxEpjXTNMuYE7OZckijfsb_7EXU,8246 +torch/_inductor/kernel/mm_grouped.py,sha256=Ppb6nHSG7YuxxFuPweUoWyF-GXPxsQJpmgt4mXLcqwE,15102 +torch/_inductor/kernel/mm_plus_mm.py,sha256=k48JgNuo-d5ZMp_kEkSGS_GMNhq36lBJpK8X9Drne4M,5903 +torch/_inductor/kernel/templates/cutedsl_mm_grouped.py.jinja,sha256=zq-mpTCpEi-noT_MbWOSqRAcp37vNtRez3HpyJFZKK4,11499 +torch/_inductor/kernel/templates/triton_blackwell_ws_persistent_device_tma_mm.py.jinja,sha256=hIMljiFj5avpzdnm1bW0jzduikleKXrvqLNCsk7ohhw,5135 +torch/_inductor/kernel/templates/triton_bmm.py.jinja,sha256=tm4Cejifgn2mNIDKz5_sGlE9txzwnferYJ3_Y52TX4k,2935 +torch/_inductor/kernel/templates/triton_depthwise_conv.py.jinja,sha256=vniupjpzjkQIr4xuJ8z9IX3Zy_696GYnJ7UHFF1eqM8,2118 +torch/_inductor/kernel/templates/triton_epilogue_scaled_mm.py.jinja,sha256=Ia6sY7G98qTPTLvst1-vlq8X61FotuHjs9MwK3AolaU,6370 +torch/_inductor/kernel/templates/triton_main_loop_scaled_mm.py.jinja,sha256=gK43gOVrS7TS4GmyUfmggvCs6Ypn_GGi1aHLsPq3euA,6035 +torch/_inductor/kernel/templates/triton_mm.py.jinja,sha256=2ZhC5GaAPi0dzRm4-qmSrNJtyctnG466dGjBh_O5AH8,2878 +torch/_inductor/kernel/templates/triton_mm_grouped.py.jinja,sha256=Y4maui6ZW2ZTq_LdpCyCx6ve-_E68fXi-fTlG2HDqXw,12552 +torch/_inductor/kernel/templates/triton_mm_rocm.py.jinja,sha256=jZCVMvz_VGV_RjGyqDnzBRDaSDOlO7VAfQf7y5jVlzc,2815 +torch/_inductor/kernel/templates/triton_persistent_mm.py.jinja,sha256=wpK3T5UTwo7L1bC9u-1_UNDeIjmz8tU9fKB6BHgO02I,3123 +torch/_inductor/kernel/templates/triton_persistent_tma_mm.py.jinja,sha256=DdrQ3xVsQQBoHaZhVjc4bR6MJAPnDqRmL3R97-wJP_g,4645 +torch/_inductor/kernel/vendored_templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/kernel/vendored_templates/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/kernel/vendored_templates/cutedsl/__init__.py,sha256=JL0fKO0h8FafXYiX9rhcnEAHoxhUZhqBR30T7q6GOCI,62 +torch/_inductor/kernel/vendored_templates/cutedsl/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/kernel/vendored_templates/cutedsl/__pycache__/dense_blockscaled_gemm_persistent.cpython-311.pyc,, +torch/_inductor/kernel/vendored_templates/cutedsl/dense_blockscaled_gemm_persistent.py,sha256=n_fzjFvcfSAjDRncv1JJE7_NzY5YBjxYWPK43tWdaFg,85972 +torch/_inductor/kernel/vendored_templates/cutedsl/kernels/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/kernel/vendored_templates/cutedsl/kernels/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/kernel/vendored_templates/cutedsl/kernels/__pycache__/cutedsl_grouped_gemm.cpython-311.pyc,, +torch/_inductor/kernel/vendored_templates/cutedsl/kernels/cutedsl_grouped_gemm.py,sha256=ShDnjUzXbQRBlcPTHqn4hCXbUoZiX9l5PAe8DxuRAXA,94438 +torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/__init__.py,sha256=yMEmzIjn2of5wzG-nlWe_d3Jmg0u8kCFrIhrumDxowg,116 +torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/__pycache__/dense_blockscaled_gemm_kernel.cpython-311.pyc,, +torch/_inductor/kernel/vendored_templates/cutedsl/wrappers/dense_blockscaled_gemm_kernel.py,sha256=tQ-OBAPvadBIifl_BXB7AKl9H1OCmpmqKgmixAntWvs,15600 +torch/_inductor/kernel_inputs.py,sha256=e9TP6EotC3ba6Kh2dQdsM72g3d7I7dcBYb_DwjQ34HQ,11051 +torch/_inductor/kernel_template_choice.py,sha256=4RM1Cn_azFhkXVfoZ0SUnpT4hxgRI82sAXOa1urLuz0,3331 +torch/_inductor/lookup_table/__init__.py,sha256=grqROhCd6DpGiI-5pQ0zDqm9wDL0-7uwQJGQXzqYoKI,877 +torch/_inductor/lookup_table/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/lookup_table/__pycache__/choices.cpython-311.pyc,, +torch/_inductor/lookup_table/choices.py,sha256=hU4qtxuFCkUfsQZIc7uk_x_4lYVGL6SDjTc-K7Ebvn8,16194 +torch/_inductor/loop_body.py,sha256=uuaZCyw_8eWY567IYK5wjwEHiZ5TgeTqu84rx79LwhQ,27149 +torch/_inductor/lowering.py,sha256=DQLf7x0_WwI9rMPx05yXkEP2znat3uSZNhbZJ5RbIy8,287123 +torch/_inductor/memory.py,sha256=x6JWY4lzJj3b9g92VihDqlqQVINi2BHfMTQDuG3wiCE,40727 +torch/_inductor/metrics.py,sha256=zlIvB25KJNaAFxPqb3M85AV1xy2Wvn7UsQPw5H4EDmc,14626 +torch/_inductor/mkldnn_ir.py,sha256=af1RJ3ra6f__Xk4J8oxN7ZLHht7k0DV8x4qt0MwxPng,43442 +torch/_inductor/mkldnn_lowerings.py,sha256=B4y5Gentflfi8b5gflzvJBumzDvU7W3oB8W_Hvp0MBc,56692 +torch/_inductor/mock_cache.py,sha256=k9g2wAteqKUJBY2RfTJoKsVbg9B1KRn94iuW4mqR89M,8565 +torch/_inductor/ops_handler.py,sha256=H1IZs9UOyPys4gqiPKNCEFWvKQcTTVrbx7DJK0HBGf8,37230 +torch/_inductor/optimize_indexing.py,sha256=KdIFrlHwny2NlrQoiMp5G1fm0e-EwZEzzQ0tNXA4nXM,4135 +torch/_inductor/output_code.py,sha256=jf--lvI_k76NjSsZsMZmpw6urEW9ADryq6LsH_76TZI,47950 +torch/_inductor/package/__init__.py,sha256=s1BQDRTTu4baQTZPdR07q8CBO_b5kUpXEp5VawAeQQI,67 +torch/_inductor/package/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/package/__pycache__/build_package.cpython-311.pyc,, +torch/_inductor/package/__pycache__/package.cpython-311.pyc,, +torch/_inductor/package/build_package.py,sha256=ZS3248m9MtGnCGG0YWrne2y7FDTIoF2ZWmKdxvIRzMg,329 +torch/_inductor/package/package.py,sha256=Wy03oSuOkzWpMX6mC8AJv116hh_l4-T6_FJSAiNJvHk,4427 +torch/_inductor/pattern_matcher.py,sha256=obYbJ64KyV8kil1MbkQNjDkyDsCW5OJsWg0IaF67Dn0,89439 +torch/_inductor/quantized_lowerings.py,sha256=X9NMCae-OB7cbBsez7NtEY_DP5ppqh2mKxaa1aSLdI0,5774 +torch/_inductor/remote_cache.py,sha256=03pgIJlBKAgeSTKBW9ipFOdoHT08vHx7dT0lua_tdEY,13301 +torch/_inductor/remote_gemm_autotune_cache.py,sha256=1ixVIU_KN1XNTBGZuJolUfaJWClQJUURc9nghIp-kAk,551 +torch/_inductor/rocm_multiarch_utils.py,sha256=9JvvvTARB7jdMTDngKGyuc8edzKgpEHfo5sumt5dCc4,9049 +torch/_inductor/runtime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_inductor/runtime/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/autotune_cache.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/benchmarking.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/cache_dir_utils.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/compile_tasks.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/coordinate_descent_tuner.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/debug_utils.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/halide_helpers.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/hints.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/proton_utils.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/runtime_utils.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/static_triton_launcher.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/triton_compat.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/triton_helpers.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/triton_heuristics.cpython-311.pyc,, +torch/_inductor/runtime/__pycache__/triton_lazy_compile.cpython-311.pyc,, +torch/_inductor/runtime/autotune_cache.py,sha256=EVBGL_5uyi3eq2J_-s4NYA2WTPHKWoeCP6EMpP9ItqA,24059 +torch/_inductor/runtime/benchmarking.py,sha256=zdIo0PtNm9dwCYVyTbynKK0FtX-by-pyxG01DlcgdUY,21296 +torch/_inductor/runtime/cache_dir_utils.py,sha256=VMMt-78VVHpAE90lDR7cNBLc9pxE3mE5pFuhGnyh1V0,1444 +torch/_inductor/runtime/caching/__init__.py,sha256=fRIP7WF5M9w6agp6eDMrPyjRWL8p_w_3V0liSPKtpKA,730 +torch/_inductor/runtime/caching/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/runtime/caching/__pycache__/config.cpython-311.pyc,, +torch/_inductor/runtime/caching/__pycache__/context.cpython-311.pyc,, +torch/_inductor/runtime/caching/__pycache__/encoders.cpython-311.pyc,, +torch/_inductor/runtime/caching/__pycache__/exceptions.cpython-311.pyc,, +torch/_inductor/runtime/caching/__pycache__/implementations.cpython-311.pyc,, +torch/_inductor/runtime/caching/__pycache__/interfaces.cpython-311.pyc,, +torch/_inductor/runtime/caching/__pycache__/locks.cpython-311.pyc,, +torch/_inductor/runtime/caching/__pycache__/memoizers.cpython-311.pyc,, +torch/_inductor/runtime/caching/__pycache__/utils.cpython-311.pyc,, +torch/_inductor/runtime/caching/config.py,sha256=1bBNlOGURk_hDv_D8zfaFrHBU9wEx2b4zkZm9y23G4s,5902 +torch/_inductor/runtime/caching/context.py,sha256=Wp5yQx2etfH4YK3YsfGkKldmkD1agnW-sYyKYMnrsb4,10485 +torch/_inductor/runtime/caching/encoders.py,sha256=c3PE8XFfkPTlxhjkf4RukF52JxIU88Q0JTolhf8dX58,2137 +torch/_inductor/runtime/caching/exceptions.py,sha256=lTWLVXJkhktVagqgYwNDHOoRoiwmouNexQi9KSQ7elE,3463 +torch/_inductor/runtime/caching/implementations.py,sha256=s524jk8GNBLXun6FYJsvLY9p7enugEZCK70pjTP4D-c,13753 +torch/_inductor/runtime/caching/interfaces.py,sha256=u2CAsqte2SjcxWaiTdaH-R_-uvLz742pIAw_A9evrFI,51853 +torch/_inductor/runtime/caching/locks.py,sha256=MyqtTTQ4d-0-wopXK0weXYlPzDxmT1dbypUS-tT4mhQ,7506 +torch/_inductor/runtime/caching/memoizers.py,sha256=pmnwTMuUK8QHG8GUiikvM1PUN5m-Tggjndy_KdgncGg,177 +torch/_inductor/runtime/caching/utils.py,sha256=PZiNE9MVezRbxW2rpmKV8-W6z2dYbftjub0NY7jxPp4,1661 +torch/_inductor/runtime/compile_tasks.py,sha256=bgJRCQtveWT_wvSMYW7YQ_wNw2SK5onhs7Wd3TES6po,5065 +torch/_inductor/runtime/coordinate_descent_tuner.py,sha256=rFSqCsUNe5dQ_JZREUNZwjmHzZILkSkBiRra2TEKc0c,14237 +torch/_inductor/runtime/debug_utils.py,sha256=bzIPymkfFZyxIsso-rWAj7vtzOKGhX1Kjmq-3UwH6B4,4275 +torch/_inductor/runtime/halide_helpers.py,sha256=E-XakaJ4btm8Dsh3A9q5XK9BJO0YG_kAYRA0de95nfs,4817 +torch/_inductor/runtime/hints.py,sha256=IjQcXvbZ3A6_lOwOfK2_PPvEq-npXFMxlVHlbWL8vOA,7730 +torch/_inductor/runtime/proton_utils.py,sha256=NfcjNIq3mHwe1ov_-rLMgquiNnrFENL7WuZrv5JFE9k,11097 +torch/_inductor/runtime/runtime_utils.py,sha256=xl5-yfGMquZRKXvonoXSKRVwKgG7hPrrTN2YlfveR5U,31005 +torch/_inductor/runtime/static_triton_launcher.py,sha256=WzCrVeOtfra3ARrDoaeUDxq9yuKbdJ8SAXdXE016iOM,14144 +torch/_inductor/runtime/triton_compat.py,sha256=rcc77nXeylZGanMHQGOoGi8BgHZHQAoC7Ek020vzWNc,4600 +torch/_inductor/runtime/triton_helpers.py,sha256=QmKMfkjzvrU0XLvagjddgEtB70qWWpnqIt18ZENzlso,27019 +torch/_inductor/runtime/triton_heuristics.py,sha256=2tBbvoHD7UJx0wH3PLnhOce_QvyEqKr0QF7bZezyzLk,174716 +torch/_inductor/runtime/triton_lazy_compile.py,sha256=XyKu25Sc3qOarPTQsFPAPtNbmTMz8Tn58FZpbAXWJ34,8592 +torch/_inductor/scheduler.py,sha256=mMdgxuvRB3Vrq-u_uEEowIQBQSnv-Fr2ULs8Oo8InF0,319839 +torch/_inductor/script.ld,sha256=pYkdES4_8MJi9MkEF79gr1KFT4iPUvbMT7NlDWP7aec,430 +torch/_inductor/select_algorithm.py,sha256=Zo0SF5MOT1F0RhnXQd5SG1cGxNKqSGZ3KjjKFXQN7Yc,226073 +torch/_inductor/shape_propagation.py,sha256=dbVdk-EL7VzHdLxxTw_CRvGFTPGroslChmXWuw2lmyQ,4376 +torch/_inductor/sizevars.py,sha256=oPj4gtBtcPRoHU9kBxWN35MpNh5sMe9PLuN5BJfh7-k,46466 +torch/_inductor/standalone_compile.py,sha256=7WHokYJUmhJb9T-o4O-bdaUISymtvcZeQAUmjgMuUws,18616 +torch/_inductor/stream_constants.py,sha256=Z19olpHrM08Q64HqJ9g10BtiDoseojEjlHZq5191GDc,179 +torch/_inductor/stream_utils.py,sha256=S231iAfe4YnddpJhnbHYwBgRU2jaFc6WZT6zARVWudY,924 +torch/_inductor/subgraph_lowering.py,sha256=TReHghy7pIPVGlwNvsxI4bCu2gpQea419eSQUhq2pPc,7287 +torch/_inductor/template_heuristics/__init__.py,sha256=XhuhQFWFqOhU3UK6tG0UGJOfxneWi6jeN3RuO36daj0,380 +torch/_inductor/template_heuristics/__pycache__/__init__.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/aten.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/base.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/contiguous_mm.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/cutedsl.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/decompose_k.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/gemm.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/nv_universal_gemm.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/params.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/registry.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/tlx.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/triton.cpython-311.pyc,, +torch/_inductor/template_heuristics/__pycache__/triton_addmm.cpython-311.pyc,, +torch/_inductor/template_heuristics/aten.py,sha256=SFbdj1U9xLQtALGuq8pRMKarNU1Ln8xjbLYIybwcydc,3124 +torch/_inductor/template_heuristics/base.py,sha256=mTxhZaLAmWKVeGjIMeUfNBPOO7liEKJYw8-JRgEio-w,2497 +torch/_inductor/template_heuristics/contiguous_mm.py,sha256=jQg313JcP2DeekcygTozWG9Lkpf69qCoHPVCCTIL77I,2232 +torch/_inductor/template_heuristics/cutedsl.py,sha256=2s0YcvzfF0TMtDGjVPNma0x-tmD-ZCsB127qHVfkdmU,4455 +torch/_inductor/template_heuristics/decompose_k.py,sha256=XDUM5DyYk8ER7WqZ3KK_nkdhi_9lAD88kc-yeRW-a5Q,2380 +torch/_inductor/template_heuristics/gemm.py,sha256=MULIGFXVMRAncb638ftAP57myYNIoPkHXBEyW983N2A,540 +torch/_inductor/template_heuristics/nv_universal_gemm.py,sha256=y_LaJjj31NE8FgQo8DXJ9927H8p_ZzsMud6b6hPVNYY,13603 +torch/_inductor/template_heuristics/params.py,sha256=shqTLFkq9Od8WGcRyXI-PSM7hwSmbKEgeArhsQWOSNk,1257 +torch/_inductor/template_heuristics/registry.py,sha256=57TmEBFeTAKi4cda83zQ8Vnk8OFwfPOUosCX0QRULfo,6835 +torch/_inductor/template_heuristics/tlx.py,sha256=Ecx3Hu78YpAl7p-nQvYL9yAmJYyFd3WY7VfC138ZEHw,160 +torch/_inductor/template_heuristics/triton.py,sha256=sDTrxPv_fJxgtl1evpxweGhQm4fizKvBticGcUhHE1A,115426 +torch/_inductor/template_heuristics/triton_addmm.py,sha256=tqFmxlatf62RUGUxJvARqYANedz9KyoxUtG5cPQh1Uo,1085 +torch/_inductor/test_case.py,sha256=AWBtSqYplhcONz2Ozc2lAbmX_EeNyom9Qr2RcA7z8pY,1444 +torch/_inductor/test_operators.py,sha256=0qiHRxRoouHirpj13eTQJyWqYgzEmvjj2aWasALEV3k,861 +torch/_inductor/tiling_utils.py,sha256=ueG81O13KBOrLSexAZVv6w7h5eTe8ie8qJKbn58Ai6g,28409 +torch/_inductor/triton_bundler.py,sha256=ULG1GhCEddDAMPDhSO3OscniEug_0PhhWLPkvfcpJ_I,17034 +torch/_inductor/utils.py,sha256=0Joi5EJmVASY3oSYPQFabOmLAE2Gv3afmlFhY0YBLng,154410 +torch/_inductor/virtualized.py,sha256=S6sHEQfPjNaf2UoVjOtspKRx-0Wy2dqP8s_RYVEeO8g,16285 +torch/_inductor/wrapper_benchmark.py,sha256=v6HUlj_t0fEwXv4aaDBI8NXIlomq_s40flUzmQiSpkM,16409 +torch/_jit_internal.py,sha256=YvxWyC4KvwAvR3HzCLWgocjFp_PYlyBLHKCBNcFoHCQ,55130 +torch/_lazy/__init__.py,sha256=F0lAsONCnXzChhzI8m2QTyNg7WFrXHak0C38ZsPW2SI,1793 +torch/_lazy/__pycache__/__init__.cpython-311.pyc,, +torch/_lazy/__pycache__/closure.cpython-311.pyc,, +torch/_lazy/__pycache__/computation.cpython-311.pyc,, +torch/_lazy/__pycache__/config.cpython-311.pyc,, +torch/_lazy/__pycache__/debug.cpython-311.pyc,, +torch/_lazy/__pycache__/device_context.cpython-311.pyc,, +torch/_lazy/__pycache__/extract_compiled_graph.cpython-311.pyc,, +torch/_lazy/__pycache__/ir_cache.cpython-311.pyc,, +torch/_lazy/__pycache__/metrics.cpython-311.pyc,, +torch/_lazy/__pycache__/tensor_factory_functions.cpython-311.pyc,, +torch/_lazy/__pycache__/ts_backend.cpython-311.pyc,, +torch/_lazy/closure.py,sha256=QYdZ2qT1zA9x-_2EEvZ76mAvIOSAHpkMCIwfoFjwNms,5724 +torch/_lazy/computation.py,sha256=qGyb2-6Mk9HtK1ekYjC-RaRAFG8sUEAazCoO4BuCM1U,919 +torch/_lazy/config.py,sha256=-3YMuc6MvZef7TCPIX4J-MtrHwXVsrHPKZDvdCQ6iYM,448 +torch/_lazy/debug.py,sha256=LOZyiVd9aFg1xYiUwq-FKA1jMxuGDyU0KcTpor81r1s,738 +torch/_lazy/device_context.py,sha256=RNChVQuKOwclZN0xYy8SRFdYG89NZzNBbVmYTLP07Hc,668 +torch/_lazy/extract_compiled_graph.py,sha256=LiYQHFvb2rnYNxUlvy-uYoCW5cZMpH9SZycMXBKhuH0,8912 +torch/_lazy/ir_cache.py,sha256=HcW7N_L3ff-7_dDnpTsuh0ZQHWU-CSkawOriqQQCoL4,348 +torch/_lazy/metrics.py,sha256=emYDBgpRK4nYV8b5miNCvEB0ses4DQkq4cBTVhGjQ5E,545 +torch/_lazy/tensor_factory_functions.py,sha256=_vbfXc_XMUD6dXa3dP-BHlwmc1nWhnZ5R1SkY9sIZ98,1368 +torch/_lazy/ts_backend.py,sha256=BfAAT0WhImXNRhQp3Pfbp1tLuKuy7UBt4OziWTEYi9o,163 +torch/_library/__init__.py,sha256=nrHN83rXP5crQ0OeQiDyqXNcyoctiW-5dOFVAPArtzs,269 +torch/_library/__pycache__/__init__.cpython-311.pyc,, +torch/_library/__pycache__/_out_variant.cpython-311.pyc,, +torch/_library/__pycache__/autograd.cpython-311.pyc,, +torch/_library/__pycache__/custom_ops.cpython-311.pyc,, +torch/_library/__pycache__/effects.cpython-311.pyc,, +torch/_library/__pycache__/fake_class_registry.cpython-311.pyc,, +torch/_library/__pycache__/fake_impl.cpython-311.pyc,, +torch/_library/__pycache__/fake_profile.cpython-311.pyc,, +torch/_library/__pycache__/infer_schema.cpython-311.pyc,, +torch/_library/__pycache__/opaque_object.cpython-311.pyc,, +torch/_library/__pycache__/simple_registry.cpython-311.pyc,, +torch/_library/__pycache__/triton.cpython-311.pyc,, +torch/_library/__pycache__/utils.cpython-311.pyc,, +torch/_library/_out_variant.py,sha256=YmG-_JK0yKwm6DxcyITJcLSpupSy-YcmK58qIsjp-XQ,6386 +torch/_library/autograd.py,sha256=qurBYsR8RhPautAQLcJJcHb7GOb3kvpJa-VS_4SDFu8,9035 +torch/_library/custom_ops.py,sha256=HKx5Wvo7JhgjUFQTGAaxuySH_m8EeME5C2H5bh8l5p4,39967 +torch/_library/effects.py,sha256=iazijTW8NbQeCY6ME3DCwWZuqpMnnUUK_wsZESKaZvs,2781 +torch/_library/fake_class_registry.py,sha256=Tfk671usmR5dEz-6n_lwbpnKdadWmr853oA7jmMr2G0,20163 +torch/_library/fake_impl.py,sha256=4P-ohlUnBIZGgP8EcvVChi00li70gdTHeQIklrNU21U,8908 +torch/_library/fake_profile.py,sha256=gf2bzo9i-0VlFn6v9aww20CsveVPC-LjQTyRoWXNuuc,11571 +torch/_library/infer_schema.py,sha256=jHFJsfW2U0U55xNA17At_Zr9bZil5Fmpn301x4pxGFs,14328 +torch/_library/opaque_object.py,sha256=OrnUD_aXD1K1LGbv-s_MSjhLWaBB8Ih930tt_aaH2Us,14813 +torch/_library/simple_registry.py,sha256=Tkad72zL0M0DYw3zdtiXJw86mJXVYAz0AyPKNOsPSXk,2933 +torch/_library/triton.py,sha256=tnRnrXxAqPdk4xdVgYw0lD03K1N-0dMjsjr8x8tJSIM,23746 +torch/_library/utils.py,sha256=CfjIRvd4fflagvW6jiEy95h_eWv0t7-HOg9JE9-DLK4,23040 +torch/_linalg_utils.py,sha256=Q2UhuVexGLubUoqrPJZTOn5lZerPgnwkVwjSXBrMqQ0,5199 +torch/_lobpcg.py,sha256=KpttDXkgt-wxgyAkZobHRdf3zyt0e3hjgRBltaukkLw,44233 +torch/_logging/__init__.py,sha256=wTE40MPYVeohdWbp_5z_a3QBDWUYOLV_CSd0MuwP9pc,818 +torch/_logging/__pycache__/__init__.cpython-311.pyc,, +torch/_logging/__pycache__/_internal.cpython-311.pyc,, +torch/_logging/__pycache__/_registrations.cpython-311.pyc,, +torch/_logging/__pycache__/scribe.cpython-311.pyc,, +torch/_logging/__pycache__/structured.cpython-311.pyc,, +torch/_logging/_internal.py,sha256=ACRcyJiLrvgUPkbYsW8R3jtfxO_mowthaqQPkrbsqGk,56369 +torch/_logging/_registrations.py,sha256=-yIpdNUSDTe948alwUhqPogT17mNUpEheZHVjKhvhzk,8913 +torch/_logging/scribe.py,sha256=IUrj5iFNr6FVKHyyC4Uf7cu32YRoN3Zr6E3H-erOGdA,2553 +torch/_logging/structured.py,sha256=LgZTqykz-WIzLI6OI0D354bBqttp8Kcyj8_TuSMBQxw,2909 +torch/_lowrank.py,sha256=eBhxRMhfi9L1h_LVjrzulAJEBHwTKzBQe_AprNFtvUg,10497 +torch/_meta_registrations.py,sha256=eNrbgZoLS2N2jyiptxdJZA2VKDlGetMJ6x1C9XhOxmk,292468 +torch/_namedtensor_internals.py,sha256=CMgdQl-RJDFXZz_TOsuGIrMoxdszASwnIfpYjDkB8as,5276 +torch/_native/__init__.py,sha256=-Gw8ZtGmJbMoqHyjvsGb-_3F86XjzahgjBsdAU19mR8,1650 +torch/_native/__pycache__/__init__.cpython-311.pyc,, +torch/_native/__pycache__/common_utils.cpython-311.pyc,, +torch/_native/__pycache__/cutedsl_utils.cpython-311.pyc,, +torch/_native/__pycache__/dsl_registry.cpython-311.pyc,, +torch/_native/__pycache__/registry.cpython-311.pyc,, +torch/_native/__pycache__/triton_utils.cpython-311.pyc,, +torch/_native/common_utils.py,sha256=zzGHS4LSAg3sNcajkb7joyGc7heGbVEVjglHzBVLKWA,1831 +torch/_native/cutedsl_utils.py,sha256=KTm0UU4w5OoA_yTUBcUQjuhXx5Zoowyc4K9m7X35QpU,3432 +torch/_native/dsl_registry.py,sha256=Ulm54qoIo_HrrNzH_m9IRz9NrgH0eSpE4poxYLB93TM,5166 +torch/_native/ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_native/ops/__pycache__/__init__.cpython-311.pyc,, +torch/_native/registry.py,sha256=UCNAQSGyzsHXZJ3WKkCbj1spKflkGAX67H_oRaIirUU,24763 +torch/_native/triton_utils.py,sha256=8S8ER3Qpx4IBgOMnDb6-hUQUhEHoKMZHeh1I46F5ajs,3335 +torch/_numpy/__init__.py,sha256=_Jxx8XlO6rLDkI0bP7kYY_JLbYI8kVA3eFfjesMk-qg,590 +torch/_numpy/__pycache__/__init__.cpython-311.pyc,, +torch/_numpy/__pycache__/_binary_ufuncs_impl.cpython-311.pyc,, +torch/_numpy/__pycache__/_casting_dicts.cpython-311.pyc,, +torch/_numpy/__pycache__/_dtypes.cpython-311.pyc,, +torch/_numpy/__pycache__/_dtypes_impl.cpython-311.pyc,, +torch/_numpy/__pycache__/_funcs.cpython-311.pyc,, +torch/_numpy/__pycache__/_funcs_impl.cpython-311.pyc,, +torch/_numpy/__pycache__/_getlimits.cpython-311.pyc,, +torch/_numpy/__pycache__/_ndarray.cpython-311.pyc,, +torch/_numpy/__pycache__/_normalizations.cpython-311.pyc,, +torch/_numpy/__pycache__/_reductions_impl.cpython-311.pyc,, +torch/_numpy/__pycache__/_ufuncs.cpython-311.pyc,, +torch/_numpy/__pycache__/_unary_ufuncs_impl.cpython-311.pyc,, +torch/_numpy/__pycache__/_util.cpython-311.pyc,, +torch/_numpy/__pycache__/fft.cpython-311.pyc,, +torch/_numpy/__pycache__/linalg.cpython-311.pyc,, +torch/_numpy/__pycache__/random.cpython-311.pyc,, +torch/_numpy/_binary_ufuncs_impl.py,sha256=h_xCtCyRAdTEB2T_r0AZ25G7u9IM3PRPr1mVv8Cr0qY,1871 +torch/_numpy/_casting_dicts.py,sha256=EmT0GYUz7Ii1iT1prWZgsahqtZZ069I1O0T2pzhSlOE,42478 +torch/_numpy/_dtypes.py,sha256=m3UM7nb7RFXlHSbtbv1GAu0c8AD4I4Uc1Rn40s1iLQ0,10279 +torch/_numpy/_dtypes_impl.py,sha256=878WyQ_5NSSp8T22sBZrWi16ParwHA80OGkgCxcA_FA,6475 +torch/_numpy/_funcs.py,sha256=vT_F_qnfL3CJoZr_R1NEHvETRCnIoEOrNGXVP5o-5aM,2097 +torch/_numpy/_funcs_impl.py,sha256=Q8on4aGY4oO_lOJbLxqXZaPzR4lxiAtanaU4XDOsQ1c,59362 +torch/_numpy/_getlimits.py,sha256=QRUbPnrehERJA9Wk01SdUaF1F2FhksfgpewYN6SmtgE,269 +torch/_numpy/_ndarray.py,sha256=Rf4zVsiY4BluCqquLMD9HJe9st-M7kyvI0sXGKNqIwk,21895 +torch/_numpy/_normalizations.py,sha256=UwbHxScNwVLgYI8UaGI9r6DjoDckG0SoLut5shFhsS0,8980 +torch/_numpy/_reductions_impl.py,sha256=moT-A4XxdcPnlvSfG0hCMS2SWV6VJVO96rZ8TFqz5qI,11843 +torch/_numpy/_ufuncs.py,sha256=HbjZbhF5ZMhSTMa0tGA_8oEdUBd5rj63sfhKpVmM3Us,8283 +torch/_numpy/_unary_ufuncs_impl.py,sha256=IO2kPjYhNoKCyu_9NcMFQF57le1U9VjFmqQFbsWA3Jc,1161 +torch/_numpy/_util.py,sha256=AITuUazGv6PxHIX1xTj6Ujv6HGOOWeQbMjIQLJ1uvGg,7836 +torch/_numpy/fft.py,sha256=lqeN-889bRRT8VjzgAaXoK9aAz8wdQdoiUSVN-oIhRQ,2805 +torch/_numpy/linalg.py,sha256=0YwsdXjndqqifpznxyxzU9ki7aQi5q3vZRrMK41Nsgw,5648 +torch/_numpy/random.py,sha256=i8D7VLbdUosxWruXrUGXG00TcHpGcaio0jCw1r7HTj4,4620 +torch/_numpy/testing/__init__.py,sha256=t5Re9c4lijwKoegYKS74CFviLCfSDSWZMvw78ohKYss,375 +torch/_numpy/testing/__pycache__/__init__.cpython-311.pyc,, +torch/_numpy/testing/__pycache__/utils.cpython-311.pyc,, +torch/_numpy/testing/utils.py,sha256=FDje7mVmZvrLmbrmq3C7xlu5K8PGyRcz_SmwmRrjnM0,77459 +torch/_opaque_base.py,sha256=31QtGBgy-IBMA7ZZmN8D2aAlvW1v7w6cO4mJhYbQqJw,462 +torch/_ops.py,sha256=w9Vh0Tqk3t5FJaXfi0fjilQTVlPRkHEBIRyEPfRnpRk,63940 +torch/_prims/__init__.py,sha256=nann-bsQeMfkqLWdpHNYLXepV4BRAoFyp7sy5GHUgRM,86254 +torch/_prims/__pycache__/__init__.cpython-311.pyc,, +torch/_prims/__pycache__/context.cpython-311.pyc,, +torch/_prims/__pycache__/debug_prims.cpython-311.pyc,, +torch/_prims/__pycache__/executor.cpython-311.pyc,, +torch/_prims/__pycache__/rng_prims.cpython-311.pyc,, +torch/_prims/context.py,sha256=wMn_9XZQqJw7JOoV6d5bCSYXSQQyEvsItLi-HOpfXTY,6248 +torch/_prims/debug_prims.py,sha256=sfqX5M4dBhvQRgMrt6-o8WHTzXjvIUBehRxcwFjjHxY,2250 +torch/_prims/executor.py,sha256=6qVwYXMvV0rAsvRa4hpW3bZ7X4Pem4mzGNRWpjgtij8,1923 +torch/_prims/rng_prims.py,sha256=SqRCxzXJ8Hrur4N5cAjXC9q3sd2byxeaadRgKfzT6N0,18934 +torch/_prims_common/__init__.py,sha256=Ni5pPfWQpiIjABvhgyvnV6NQzpVTQfFFqmAAqknq5d8,73964 +torch/_prims_common/__pycache__/__init__.cpython-311.pyc,, +torch/_prims_common/__pycache__/wrappers.cpython-311.pyc,, +torch/_prims_common/wrappers.py,sha256=Xeso-cgOZ4UShVH13GIPaBc8UxwvkBLmYhSKAM9qJZY,19611 +torch/_python_dispatcher.py,sha256=tuIlb49DqVSQW8fC4wYUrMYfpou5q_-AR_O47jTElgM,7137 +torch/_refs/__init__.py,sha256=oa54b9lFxzFyQInlsN32UJi2zRokhC75066RNOu4XxU,228757 +torch/_refs/__pycache__/__init__.cpython-311.pyc,, +torch/_refs/__pycache__/_conversions.cpython-311.pyc,, +torch/_refs/__pycache__/fft.cpython-311.pyc,, +torch/_refs/_conversions.py,sha256=BkNyamBLOQLlKy6mB4BSdPRjAKP1ghZ5xcyQDkUs-_s,3533 +torch/_refs/fft.py,sha256=zrpE4e0HsVewb43wiumlAO7dZYRvDC34p8S1y76LhpA,17971 +torch/_refs/linalg/__init__.py,sha256=o9_Ajxfaw0lyGeMN3YJrx4l0C7TSaxzffsGZovO7SSg,14497 +torch/_refs/linalg/__pycache__/__init__.cpython-311.pyc,, +torch/_refs/nn/__init__.py,sha256=juqd9PbXs4yg45zMJ7BHAOPQjb7sgEbWE9InBtGZhfo,24 +torch/_refs/nn/__pycache__/__init__.cpython-311.pyc,, +torch/_refs/nn/functional/__init__.py,sha256=5SE96Uca9OTrK6er5okycoxXoAOjdrzd5TPZZzIvvTY,42550 +torch/_refs/nn/functional/__pycache__/__init__.cpython-311.pyc,, +torch/_refs/special/__init__.py,sha256=Njxv8FI1zSzZANrqwvhVe9oD3hZsY2rgsWkHF6vwAuc,7035 +torch/_refs/special/__pycache__/__init__.cpython-311.pyc,, +torch/_size_docs.py,sha256=z2Ckk4Ufq-Ek6Av1-rs1Wcrg6OoJxTCRPYjEjhzG1Xo,900 +torch/_sources.py,sha256=qnTfAX_mWhAQ5cMVJYYWDOK1o1LoepbMllA1FIxIFRY,4404 +torch/_storage_docs.py,sha256=KyxMZ6j9UPSolidRjR7WDu7gJWkYot8TFggfPA1fygI,1335 +torch/_streambase.py,sha256=NAKZfV-BhZ22Jyt_3VQDH1m4QQdIVc4G8VvUBkZ9wJ8,435 +torch/_strobelight/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_strobelight/__pycache__/__init__.cpython-311.pyc,, +torch/_strobelight/__pycache__/cli_function_profiler.cpython-311.pyc,, +torch/_strobelight/__pycache__/compile_time_profiler.cpython-311.pyc,, +torch/_strobelight/cli_function_profiler.py,sha256=srLhCyX1lK-devswH5sih_7_MWTHaiZPFekSXkbcV-M,11776 +torch/_strobelight/compile_time_profiler.py,sha256=LmKB8TtCO1MijfGm94g0NxCN0jE4kBpm1c4dCWJXO8w,7503 +torch/_subclasses/__init__.py,sha256=7NDSoV4xw-GxCRtHasX_MFFLe_S1-1BHRFvoDcgE4J0,375 +torch/_subclasses/__pycache__/__init__.cpython-311.pyc,, +torch/_subclasses/__pycache__/_fake_tensor_utils.cpython-311.pyc,, +torch/_subclasses/__pycache__/fake_impls.cpython-311.pyc,, +torch/_subclasses/__pycache__/fake_tensor.cpython-311.pyc,, +torch/_subclasses/__pycache__/fake_utils.cpython-311.pyc,, +torch/_subclasses/__pycache__/functional_tensor.cpython-311.pyc,, +torch/_subclasses/__pycache__/meta_utils.cpython-311.pyc,, +torch/_subclasses/__pycache__/schema_check_mode.cpython-311.pyc,, +torch/_subclasses/_fake_tensor_utils.py,sha256=EK3B6s3cbyvXpsYLjjmFuJl__wBs6bgnSbOUAVZuIhY,9251 +torch/_subclasses/complex_tensor/__init__.py,sha256=XR7PLtAI6NSJlcrlloNyijyfO-ExZaMyiF0-m-x-0mI,277 +torch/_subclasses/complex_tensor/__pycache__/__init__.cpython-311.pyc,, +torch/_subclasses/complex_tensor/__pycache__/_core.cpython-311.pyc,, +torch/_subclasses/complex_tensor/_core.py,sha256=dGlspjGJHwJ154gc_b_ZsYhbN3xkKtdLdRIJ90HNWgM,5325 +torch/_subclasses/complex_tensor/_ops/__init__.py,sha256=DtdL96mlaWM6COV7bXMMNvXdQC3brGUtZXgni8oVhRo,155 +torch/_subclasses/complex_tensor/_ops/__pycache__/__init__.cpython-311.pyc,, +torch/_subclasses/complex_tensor/_ops/__pycache__/aten.cpython-311.pyc,, +torch/_subclasses/complex_tensor/_ops/__pycache__/common.cpython-311.pyc,, +torch/_subclasses/complex_tensor/_ops/__pycache__/prims.cpython-311.pyc,, +torch/_subclasses/complex_tensor/_ops/aten.py,sha256=tBXyjKcJHymyM5s_NoG1W0us2KHz0ZzZeeWdl_3dHJo,29736 +torch/_subclasses/complex_tensor/_ops/common.py,sha256=GD-DKVo9zOZydAc-xIpsOQWHBrXTlru7UKCVIsbE7LE,10414 +torch/_subclasses/complex_tensor/_ops/prims.py,sha256=qEqDoFU20-woc8B-DE9f-iKvA7-TPBxgQtUTB1QD0h0,864 +torch/_subclasses/fake_impls.py,sha256=Orerc-uBpEyuhWBEcUd0mGYNob53ujg_R9JFTPdx3CU,58723 +torch/_subclasses/fake_tensor.py,sha256=DkE5sPY-4KQVxx_cQBi0I4aU-bhmevK0jGvH-SERGd0,142542 +torch/_subclasses/fake_utils.py,sha256=bqZzCDUCjoCxHb-rgDNFd04XcKFzc79Pkiyr4nupXls,10988 +torch/_subclasses/functional_tensor.py,sha256=pTayA64-jQXHKxCHnrXaNLqGii08ppVzqKzVl7NSCPg,38904 +torch/_subclasses/meta_utils.py,sha256=XlCPeCrbMpfLnCCLQU_x0yUYs9xI66o3iiUi-2h3THc,102715 +torch/_subclasses/schema_check_mode.py,sha256=iI1Lw8IJRUxnkjIr4kBCVcfkq8Jgy_oE8P5Ivo8tPpc,9705 +torch/_tensor.py,sha256=I2zi_uiU1v_YD5JdiaSBXyTpKJcfIMy_rovFsdheyXM,76227 +torch/_tensor_docs.py,sha256=X-nIkVReAgXEGLXcwAiMQSUoWE6zK-xz3SqYIEOcE2A,145539 +torch/_tensor_str.py,sha256=1o0Fm0kBFiYmcyFgupJflsuhzjApCubFU5O2klNZ560,28537 +torch/_thread_safe_fork.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_torch_docs.py,sha256=pYCy4c-HP412NxEguKgcj_QPhIrA6ZHQsAf9BigBzlI,436335 +torch/_utils.py,sha256=NS4m-z0YAnT8wKjfhESV918G_p2McBHPGvzVeoBEPG8,50649 +torch/_utils_internal.py,sha256=Uztc8B3tqEZgCIqq8_6Z3c537JAndINEE2R-Co1wIp8,11704 +torch/_vendor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/_vendor/__pycache__/__init__.cpython-311.pyc,, +torch/_vendor/packaging/__init__.py,sha256=EhCMuCSz60IgQJ93b_4wJyAoHpU9J-uddG4QaMT0Pu4,496 +torch/_vendor/packaging/__pycache__/__init__.cpython-311.pyc,, +torch/_vendor/packaging/__pycache__/_structures.cpython-311.pyc,, +torch/_vendor/packaging/__pycache__/version.cpython-311.pyc,, +torch/_vendor/packaging/_structures.py,sha256=q3eVNmbWJGG_S0Dit_S3Ao8qQqz_5PYTXFAKBZe5yr4,1431 +torch/_vendor/packaging/version.py,sha256=XjRBLNK17UMDgLeP8UHnqwiY3TdSi03xFQURtec211A,16236 +torch/_vmap_internals.py,sha256=snAmSMoKFr4LrkcZqTNTwkt01RRF-06_38WE8792d4s,9425 +torch/_weights_only_unpickler.py,sha256=2f3LFPJV1g1EDhWOAGUBRHQX25VIDN1z9m-xiAuEV4s,22994 +torch/accelerator/__init__.py,sha256=y_AcEMF3K6TyHEjzy9qJXGdonfydCtsJnBhumxFZ2q8,11089 +torch/accelerator/__pycache__/__init__.cpython-311.pyc,, +torch/accelerator/__pycache__/_utils.cpython-311.pyc,, +torch/accelerator/__pycache__/graphs.cpython-311.pyc,, +torch/accelerator/__pycache__/memory.cpython-311.pyc,, +torch/accelerator/_utils.py,sha256=Qw0FIL4j__H_UNaMcaQ5LJAKd4VSF09LWqBrMHxhfTk,936 +torch/accelerator/graphs.py,sha256=8lEdABUDNYwIj-V15ARMIYBDF41Uvf7HjvjlzcgOP5Q,6543 +torch/accelerator/memory.py,sha256=TvT4C1rxIZzdn387wGnjijcTtP499Xa9A7GFdvKLzrI,10947 +torch/amp/__init__.py,sha256=OuEdgK_zzuRWo5_vfM_VhTooVAm1BXTBkUWxLkY6zlU,181 +torch/amp/__pycache__/__init__.cpython-311.pyc,, +torch/amp/__pycache__/autocast_mode.cpython-311.pyc,, +torch/amp/__pycache__/grad_scaler.cpython-311.pyc,, +torch/amp/autocast_mode.py,sha256=fLHaIKbcLOi-NufdjprnWADByudZNE5qzgbuBfL2MOw,22651 +torch/amp/grad_scaler.py,sha256=SaO1OVtjDeVQY5LOjjLW4nFd6oxz5Kd9QaMz_AKWF_Y,31629 +torch/ao/__init__.py,sha256=t-7QOhmWP9oe0Mh2RHAH119yFNo-jVdfOB6VIZyYS5M,678 +torch/ao/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/__init__.py,sha256=4j-eekD5zpSPgRw4Z-6hVhWmeeYaDcD9bar24mvxbVo,834 +torch/ao/nn/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/intrinsic/__init__.py,sha256=yaXY9qRdzLDsAk2wwLcm5x2PKvvwEi9DLbuDUcPYZeA,961 +torch/ao/nn/intrinsic/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/intrinsic/modules/__init__.py,sha256=FlHZWKH0Xumwd5xBvZgyeGkbYehn1oDCUnAqZ_c4WN4,655 +torch/ao/nn/intrinsic/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/intrinsic/modules/__pycache__/fused.cpython-311.pyc,, +torch/ao/nn/intrinsic/modules/fused.py,sha256=SLik04JP3MXGwRND9lW1EjtqHSiFujwzAMmyi0-nxvo,11847 +torch/ao/nn/intrinsic/qat/__init__.py,sha256=M0iylhjuqtPcbO2pAVdDg9n9e1ET359nQ9txOEErmMo,37 +torch/ao/nn/intrinsic/qat/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/intrinsic/qat/modules/__init__.py,sha256=OIDHW4Q4Mjd_Y_eR1UWDLuPOVEIJOJJj5MbbaQRdI1A,547 +torch/ao/nn/intrinsic/qat/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/intrinsic/qat/modules/__pycache__/conv_fused.cpython-311.pyc,, +torch/ao/nn/intrinsic/qat/modules/__pycache__/linear_fused.cpython-311.pyc,, +torch/ao/nn/intrinsic/qat/modules/__pycache__/linear_relu.cpython-311.pyc,, +torch/ao/nn/intrinsic/qat/modules/conv_fused.py,sha256=FHbEconqBIJUp7ymPaaMFN2H70ffSZk-ckNApY2kNzw,31861 +torch/ao/nn/intrinsic/qat/modules/linear_fused.py,sha256=iyJQ5Xir6QVzVpxImwYaValWbGEk3r09VKxDReWFiT4,6905 +torch/ao/nn/intrinsic/qat/modules/linear_relu.py,sha256=bZKak53c9vQNd_7Xnh10xeXyVZ8ymO6vVhch1TOPeeM,2184 +torch/ao/nn/intrinsic/quantized/__init__.py,sha256=AyOx7RavGaehtAWX6dEUBQcrl23h2_S3Q9ZAWYimCLQ,236 +torch/ao/nn/intrinsic/quantized/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/intrinsic/quantized/dynamic/__init__.py,sha256=M0iylhjuqtPcbO2pAVdDg9n9e1ET359nQ9txOEErmMo,37 +torch/ao/nn/intrinsic/quantized/dynamic/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/intrinsic/quantized/dynamic/modules/__init__.py,sha256=kbhh2dL33XSuVEAZNtSHMmt1papX7az6WlVALevMgZA,70 +torch/ao/nn/intrinsic/quantized/dynamic/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/intrinsic/quantized/dynamic/modules/__pycache__/linear_relu.cpython-311.pyc,, +torch/ao/nn/intrinsic/quantized/dynamic/modules/linear_relu.py,sha256=JGHgYcHBFkcj_65pA9iuF9QrORrbKC7LdlT4sJxwGKc,2223 +torch/ao/nn/intrinsic/quantized/modules/__init__.py,sha256=8aofpYeXIA8WQaao5pvOjWtLnf4rmL9xlJO9hZ8aras,409 +torch/ao/nn/intrinsic/quantized/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/intrinsic/quantized/modules/__pycache__/bn_relu.cpython-311.pyc,, +torch/ao/nn/intrinsic/quantized/modules/__pycache__/conv_add.cpython-311.pyc,, +torch/ao/nn/intrinsic/quantized/modules/__pycache__/conv_relu.cpython-311.pyc,, +torch/ao/nn/intrinsic/quantized/modules/__pycache__/linear_relu.cpython-311.pyc,, +torch/ao/nn/intrinsic/quantized/modules/bn_relu.py,sha256=5z1wesnwC49jVjLIWagOdmRW1s-Zhm41tS9tgL19tl8,3647 +torch/ao/nn/intrinsic/quantized/modules/conv_add.py,sha256=EX4p7WEeFsAkZ_4lqAi5rHaqJpBcz0Gj7i8A8j1Kaug,4821 +torch/ao/nn/intrinsic/quantized/modules/conv_relu.py,sha256=_aqBbU6SW8ddej_MH5u36vibx-QBWjS84VTbsckycBw,9575 +torch/ao/nn/intrinsic/quantized/modules/linear_relu.py,sha256=XpXdOcXMDYHN0cxE_bB02rTIZTv6EPjXagWnFxVOi8M,7144 +torch/ao/nn/qat/__init__.py,sha256=M0iylhjuqtPcbO2pAVdDg9n9e1ET359nQ9txOEErmMo,37 +torch/ao/nn/qat/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/qat/dynamic/__init__.py,sha256=M0iylhjuqtPcbO2pAVdDg9n9e1ET359nQ9txOEErmMo,37 +torch/ao/nn/qat/dynamic/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/qat/dynamic/modules/__init__.py,sha256=_5hfV0E5b71TgYVQPrPF9QZNrnw5XPuZppEwj0QItaM,50 +torch/ao/nn/qat/dynamic/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/qat/dynamic/modules/__pycache__/linear.cpython-311.pyc,, +torch/ao/nn/qat/dynamic/modules/linear.py,sha256=XJJdSaclI87qBmSywWkcxmfqxqnav5M3nBtzN1ntgm0,1215 +torch/ao/nn/qat/modules/__init__.py,sha256=RmKZ7d0ds96EIRSqAGGTrDwuXVWerE_S_NCWuEuX_20,228 +torch/ao/nn/qat/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/qat/modules/__pycache__/conv.cpython-311.pyc,, +torch/ao/nn/qat/modules/__pycache__/embedding_ops.cpython-311.pyc,, +torch/ao/nn/qat/modules/__pycache__/linear.cpython-311.pyc,, +torch/ao/nn/qat/modules/conv.py,sha256=p4Q_jIVoUv6Zwx6M7H1c7uCecwXjPeBEj_YFX2Pb-74,10082 +torch/ao/nn/qat/modules/embedding_ops.py,sha256=N29BQOpG7ZPrlB1H94lXEcIyXijb_N45iPMuKJNwqto,8345 +torch/ao/nn/qat/modules/linear.py,sha256=h_s7-_2lQfjymOVDJrgVVWwHicyWmcPqLCYsoYlRldY,3216 +torch/ao/nn/quantizable/__init__.py,sha256=M0iylhjuqtPcbO2pAVdDg9n9e1ET359nQ9txOEErmMo,37 +torch/ao/nn/quantizable/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/quantizable/modules/__init__.py,sha256=N6niR-YsfqqP_5a8ksCJ8BL-Ol1Mch_412g-I4Z8VeQ,145 +torch/ao/nn/quantizable/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/quantizable/modules/__pycache__/activation.cpython-311.pyc,, +torch/ao/nn/quantizable/modules/__pycache__/rnn.cpython-311.pyc,, +torch/ao/nn/quantizable/modules/activation.py,sha256=XSZxfdqIWp3f2I_gl3TMDBpMl8e3LWPJ55CspHM3wLg,26225 +torch/ao/nn/quantizable/modules/rnn.py,sha256=k03thvONMgb7YVMko8CaPaAM76BYnb-GhpMTGGvvSVQ,22482 +torch/ao/nn/quantized/__init__.py,sha256=NtHecsYVtg9e1IhTCdvoXSV-oQEXRQnt54dZlScPo6I,686 +torch/ao/nn/quantized/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/quantized/__pycache__/functional.cpython-311.pyc,, +torch/ao/nn/quantized/dynamic/__init__.py,sha256=M0iylhjuqtPcbO2pAVdDg9n9e1ET359nQ9txOEErmMo,37 +torch/ao/nn/quantized/dynamic/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/quantized/dynamic/modules/__init__.py,sha256=VzgogVUFAj-yE8a5Vmjiq5TV3SP51t0wXslabLVPITw,413 +torch/ao/nn/quantized/dynamic/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/quantized/dynamic/modules/__pycache__/conv.cpython-311.pyc,, +torch/ao/nn/quantized/dynamic/modules/__pycache__/linear.cpython-311.pyc,, +torch/ao/nn/quantized/dynamic/modules/__pycache__/rnn.cpython-311.pyc,, +torch/ao/nn/quantized/dynamic/modules/conv.py,sha256=osw4iXPTe1aVSk1t5oG8ak7enJNCi25Y20mo4vh-JI8,18413 +torch/ao/nn/quantized/dynamic/modules/linear.py,sha256=2mI1DjB2kVWl57shRpRli0wYMkB7ygqO1G45U4FBM3I,6495 +torch/ao/nn/quantized/dynamic/modules/rnn.py,sha256=oyCtf8urGTQVlfx5DGxIb89SSsOMHb7wXgvnC7yCv1E,51849 +torch/ao/nn/quantized/functional.py,sha256=fbcf0vbzbYB3cwJYIQab2MtXtfdTClU6HcxWtUJIc0U,29641 +torch/ao/nn/quantized/modules/__init__.py,sha256=Y56ItuRywBPBUWIROd3zxkrmrSr438EK-N71MT5i0Go,4661 +torch/ao/nn/quantized/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/activation.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/batchnorm.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/conv.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/dropout.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/embedding_ops.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/functional_modules.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/linear.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/normalization.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/rnn.cpython-311.pyc,, +torch/ao/nn/quantized/modules/__pycache__/utils.cpython-311.pyc,, +torch/ao/nn/quantized/modules/activation.py,sha256=C0GSs6WWHXAM4YKlZmxiYrarAonJIyjsNXqAVtkxht4,11670 +torch/ao/nn/quantized/modules/batchnorm.py,sha256=waZCeOapEFHC2pdFSFn1GYXcsEiW_uJ9nMX6xox0_fQ,4427 +torch/ao/nn/quantized/modules/conv.py,sha256=oADn2oxCFNbHlyS2k0hfpXsxfqmqpYFqSr4X-kdPdng,43809 +torch/ao/nn/quantized/modules/dropout.py,sha256=5nPYMBmjOyEzOZj0lz_CtUmJASXrrEr63YJnVytPTQs,806 +torch/ao/nn/quantized/modules/embedding_ops.py,sha256=_eOkKTzu3Fa-b0nfK7TBPs2rcLlk4N3dLfzdmdzSHqg,15238 +torch/ao/nn/quantized/modules/functional_modules.py,sha256=DCA34ufixknSkyGVPySPmTcSRTb_MVLjQishnAA-Zjo,9309 +torch/ao/nn/quantized/modules/linear.py,sha256=adIyE-xeyBz14S75Yqizy7nUM4v8jJplEsdAGQ3CUVI,13689 +torch/ao/nn/quantized/modules/normalization.py,sha256=zFTgY6ABpezWNjWEoHnl09qgLGlNJVrxxznhRecYP5g,9539 +torch/ao/nn/quantized/modules/rnn.py,sha256=gH1nq8M2l5Ie2nPBaY8ap91OmPlDAYSwjc_RU98vbdM,2026 +torch/ao/nn/quantized/modules/utils.py,sha256=qD-dHMUNgtgb-6WdnSA5jvR7q--wej0CCq1KLgxpVPo,4693 +torch/ao/nn/quantized/reference/__init__.py,sha256=hu3hyozMJtkc86Bu1v8tmyayuxAdeDSv0sUiOeNAr_s,284 +torch/ao/nn/quantized/reference/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/quantized/reference/modules/__init__.py,sha256=YS8MYKlqBPoA2UEBcS-MgbOzZ8_HDp-7ET4BC7KtCBk,494 +torch/ao/nn/quantized/reference/modules/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/quantized/reference/modules/__pycache__/conv.cpython-311.pyc,, +torch/ao/nn/quantized/reference/modules/__pycache__/linear.cpython-311.pyc,, +torch/ao/nn/quantized/reference/modules/__pycache__/rnn.cpython-311.pyc,, +torch/ao/nn/quantized/reference/modules/__pycache__/sparse.cpython-311.pyc,, +torch/ao/nn/quantized/reference/modules/__pycache__/utils.cpython-311.pyc,, +torch/ao/nn/quantized/reference/modules/conv.py,sha256=XxZfpM20-QJOjj6n9w6bIeSD_0LGmyl-py1uTTZIOww,16049 +torch/ao/nn/quantized/reference/modules/linear.py,sha256=oxH2CSsE_asNpObp2CUQjDsqsgSWDHeMHK9p2EF4834,2254 +torch/ao/nn/quantized/reference/modules/rnn.py,sha256=9mo3UKI7w-n1iW-ruqWQRycEgQwTH4HjoUMwuTybUio,29981 +torch/ao/nn/quantized/reference/modules/sparse.py,sha256=glgaM9jAtZt9FIqrCBJel32K0KDtk8p9k1K02F4Jcos,4683 +torch/ao/nn/quantized/reference/modules/utils.py,sha256=aOqOOq75IqlUgcioarrOqDGyx0i06KYi_jt6a1wD9Mc,15604 +torch/ao/nn/sparse/__init__.py,sha256=PfB-tgPOelyV_0eb_ipJK4LzPHwB5Z-wfJXeE_O3AK4,24 +torch/ao/nn/sparse/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/sparse/quantized/__init__.py,sha256=xO8RdXEcj7ggZh7-f2bReVBRO-F6PNBZjb5DBPwV52w,168 +torch/ao/nn/sparse/quantized/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/sparse/quantized/__pycache__/linear.cpython-311.pyc,, +torch/ao/nn/sparse/quantized/__pycache__/utils.cpython-311.pyc,, +torch/ao/nn/sparse/quantized/dynamic/__init__.py,sha256=lYDGtZ8rNR56jJLk6s5WOJ65E7qqQ_LFICpuD76NWYI,57 +torch/ao/nn/sparse/quantized/dynamic/__pycache__/__init__.cpython-311.pyc,, +torch/ao/nn/sparse/quantized/dynamic/__pycache__/linear.cpython-311.pyc,, +torch/ao/nn/sparse/quantized/dynamic/linear.py,sha256=ZV7BlimowKVu6jDBsDQecdlOI4JcBMwLT4X4WE650uE,6956 +torch/ao/nn/sparse/quantized/linear.py,sha256=-kpbqePTgr4lfjNKfAOYQm_7-3S9iEBRLEBD2emE_hs,10150 +torch/ao/nn/sparse/quantized/utils.py,sha256=fU0jKT2k7i8On6fgx1QEFxyzY2_3XuqjcH9Qw-PD3As,2234 +torch/ao/ns/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/ao/ns/__pycache__/__init__.cpython-311.pyc,, +torch/ao/ns/__pycache__/_numeric_suite.cpython-311.pyc,, +torch/ao/ns/__pycache__/_numeric_suite_fx.cpython-311.pyc,, +torch/ao/ns/_numeric_suite.py,sha256=fPy3eN33wImLp8FFODnH3byzxukzA-GJXQjP0hRnASc,20065 +torch/ao/ns/_numeric_suite_fx.py,sha256=FB23FPdq3xeOoV-K27ow2equHsQ3WZ-cCk3QH59h2wg,41390 +torch/ao/ns/fx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/ao/ns/fx/__pycache__/__init__.cpython-311.pyc,, +torch/ao/ns/fx/__pycache__/graph_matcher.cpython-311.pyc,, +torch/ao/ns/fx/__pycache__/graph_passes.cpython-311.pyc,, +torch/ao/ns/fx/__pycache__/mappings.cpython-311.pyc,, +torch/ao/ns/fx/__pycache__/n_shadows_utils.cpython-311.pyc,, +torch/ao/ns/fx/__pycache__/ns_types.cpython-311.pyc,, +torch/ao/ns/fx/__pycache__/pattern_utils.cpython-311.pyc,, +torch/ao/ns/fx/__pycache__/qconfig_multi_mapping.cpython-311.pyc,, +torch/ao/ns/fx/__pycache__/utils.cpython-311.pyc,, +torch/ao/ns/fx/__pycache__/weight_utils.cpython-311.pyc,, +torch/ao/ns/fx/graph_matcher.py,sha256=_1DKursxX_YdJTq3fPwmeSN7ZxHp7UGJYlLTBSZbv9U,20047 +torch/ao/ns/fx/graph_passes.py,sha256=HoAnevrFbZIjr5YLKK_i8wdt1_seAOzEDZTyQd-gTkM,45739 +torch/ao/ns/fx/mappings.py,sha256=I3ZnL73L2BTNVhhI9cR7GL-mgjjLx18f5n76XoFBt2o,18446 +torch/ao/ns/fx/n_shadows_utils.py,sha256=XXDjcij7No1I11JMxjbOJ6LQVFDHh_Qs5eGpHYoNhD8,53187 +torch/ao/ns/fx/ns_types.py,sha256=8BL1eChSs_jwmCENUoe8xGmdQp5wS8qQEm3GHKa_T2U,2341 +torch/ao/ns/fx/pattern_utils.py,sha256=Uw4gFW1G6zAfUfCQvqj_U0yfJDbd2RnBKk2v6zm2dss,8653 +torch/ao/ns/fx/qconfig_multi_mapping.py,sha256=q4SF9lmy89wQ-YFxHJgurdwfiKBn0hRPVk1ZWCvE0nA,10197 +torch/ao/ns/fx/utils.py,sha256=tjviJVCEB9vIJP8p9eujNoJtXrR60uWpINLzTJNy8do,22761 +torch/ao/ns/fx/weight_utils.py,sha256=mkhs5U5qrV_MIKelp_OAOaljhb1_kabt1NF_0gdVGik,12786 +torch/ao/pruning/__init__.py,sha256=dp2CA7CO_FRHTE7E0Ft495uR5IlUi8HsnG8ujLsz6OA,640 +torch/ao/pruning/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/__pycache__/_mappings.cpython-311.pyc,, +torch/ao/pruning/_experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/ao/pruning/_experimental/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/_experimental/activation_sparsifier/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/ao/pruning/_experimental/activation_sparsifier/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/_experimental/activation_sparsifier/__pycache__/activation_sparsifier.cpython-311.pyc,, +torch/ao/pruning/_experimental/activation_sparsifier/activation_sparsifier.py,sha256=Vb54teEd0at8qH0G0j-cBlsc1DwzlhJt4lSmD_BCXhs,19309 +torch/ao/pruning/_experimental/data_scheduler/__init__.py,sha256=q_95mAMpHldGWwLAgYS-F07ReGKDLIAZLsP7BtlmgTE,92 +torch/ao/pruning/_experimental/data_scheduler/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_scheduler/__pycache__/base_data_scheduler.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_scheduler/base_data_scheduler.py,sha256=sNQEoYKRbglXHjHbQh6Q2msIDK_fYB1py8IrExc_cbo,7733 +torch/ao/pruning/_experimental/data_sparsifier/__init__.py,sha256=9ktAif-dttGBmiw745tN2WtUTOTZch0E_HA4cE1bfTo,174 +torch/ao/pruning/_experimental/data_sparsifier/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_sparsifier/__pycache__/base_data_sparsifier.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_sparsifier/__pycache__/data_norm_sparsifier.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_sparsifier/__pycache__/quantization_utils.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_sparsifier/base_data_sparsifier.py,sha256=GIe1rn6fEjtcQ7XMptkDVPSqTx0LnNbc8wnAbejCeIE,13552 +torch/ao/pruning/_experimental/data_sparsifier/data_norm_sparsifier.py,sha256=e4It7UrOu1FyQif3ijmDYGU_D5HscmW9Wdu8qbbT1PE,7778 +torch/ao/pruning/_experimental/data_sparsifier/lightning/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/ao/pruning/_experimental/data_sparsifier/lightning/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/__pycache__/_data_sparstity_utils.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/__pycache__/data_sparsity.cpython-311.pyc,, +torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/_data_sparstity_utils.py,sha256=Vtb__BJ_qrCy7TBisG4BM8qNXzDYrl-xNDVLbHCsYIY,1636 +torch/ao/pruning/_experimental/data_sparsifier/lightning/callbacks/data_sparsity.py,sha256=-AMNv2c1CakFrnm4F6Vs-PyvnN21CdCBi7rICTrn_e4,6595 +torch/ao/pruning/_experimental/data_sparsifier/quantization_utils.py,sha256=9fkvGYqQbyzAPFf_CJjC2ChOuvDevlGOFvlqeR4_z9E,6207 +torch/ao/pruning/_experimental/pruner/FPGM_pruner.py,sha256=m75j6P43OX2ICxugVq1Zdqlsz8q5GAaOR8zSN7MEzTk,3425 +torch/ao/pruning/_experimental/pruner/__init__.py,sha256=abfUoG48Kc6ddQZmFAkYjQH5J8Hae8xJiqxz89MTgDE,260 +torch/ao/pruning/_experimental/pruner/__pycache__/FPGM_pruner.cpython-311.pyc,, +torch/ao/pruning/_experimental/pruner/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/_experimental/pruner/__pycache__/base_structured_sparsifier.cpython-311.pyc,, +torch/ao/pruning/_experimental/pruner/__pycache__/lstm_saliency_pruner.cpython-311.pyc,, +torch/ao/pruning/_experimental/pruner/__pycache__/match_utils.cpython-311.pyc,, +torch/ao/pruning/_experimental/pruner/__pycache__/parametrization.cpython-311.pyc,, +torch/ao/pruning/_experimental/pruner/__pycache__/prune_functions.cpython-311.pyc,, +torch/ao/pruning/_experimental/pruner/__pycache__/saliency_pruner.cpython-311.pyc,, +torch/ao/pruning/_experimental/pruner/base_structured_sparsifier.py,sha256=8ThuyadcbX6SNuzO0gEMK7okMCR2IWxziue0xzuTY-Q,11029 +torch/ao/pruning/_experimental/pruner/lstm_saliency_pruner.py,sha256=nBliDOjjzo42r3g6yj5q9qq4tbSLg2e_bwc1KiHA9L4,2197 +torch/ao/pruning/_experimental/pruner/match_utils.py,sha256=PLY0zuh1dWLCtLwmhAgXv0h5WDzCjXJZMM7UNRFNuO4,1949 +torch/ao/pruning/_experimental/pruner/parametrization.py,sha256=79BNtvk1z3aQL2WauM0h-PCjhJtxJCZD0UiyoJm9gh4,2047 +torch/ao/pruning/_experimental/pruner/prune_functions.py,sha256=O5lbpjsVE2rgilgB2-mVRlHmU1M5OhSJFtbp6KtJwyI,19285 +torch/ao/pruning/_experimental/pruner/saliency_pruner.py,sha256=9_q30_e_oArjSWcZ16UEx0HRcBfxCqcJTrFGrpMatrg,1536 +torch/ao/pruning/_mappings.py,sha256=O9aJePj2X0nV76KTcSU93V7GA8dxEzfr5G9QvjyWOqQ,597 +torch/ao/pruning/scheduler/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/ao/pruning/scheduler/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/scheduler/__pycache__/base_scheduler.cpython-311.pyc,, +torch/ao/pruning/scheduler/__pycache__/cubic_scheduler.cpython-311.pyc,, +torch/ao/pruning/scheduler/__pycache__/lambda_scheduler.cpython-311.pyc,, +torch/ao/pruning/scheduler/base_scheduler.py,sha256=IipBI2JjzC5ah4ocvY5mfNjWfJMLAyzWgjyIR6BmxNM,6625 +torch/ao/pruning/scheduler/cubic_scheduler.py,sha256=9j4ifJT9lavz1pCgulfgpHgHaasvx_KkqNwqCE1_NC8,3875 +torch/ao/pruning/scheduler/lambda_scheduler.py,sha256=ztFUy3RecZ5JhkDfOUg1JNVd9qnm59LqdUxoKumVrvQ,2416 +torch/ao/pruning/sparsifier/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/ao/pruning/sparsifier/__pycache__/__init__.cpython-311.pyc,, +torch/ao/pruning/sparsifier/__pycache__/base_sparsifier.cpython-311.pyc,, +torch/ao/pruning/sparsifier/__pycache__/nearly_diagonal_sparsifier.cpython-311.pyc,, +torch/ao/pruning/sparsifier/__pycache__/utils.cpython-311.pyc,, +torch/ao/pruning/sparsifier/__pycache__/weight_norm_sparsifier.cpython-311.pyc,, +torch/ao/pruning/sparsifier/base_sparsifier.py,sha256=iQuwcy8Vqbsp70TC01C0thf3TT1Z3yXnx11VH6tItwQ,13962 +torch/ao/pruning/sparsifier/nearly_diagonal_sparsifier.py,sha256=vKZbWhsxPzMDfk6ZntgIBncW81Vh19DisumzDWDhLUE,2263 +torch/ao/pruning/sparsifier/utils.py,sha256=fxTRyQBCjxVI6D-YEUDytd7hTDyVJb5MHvKM1M9J3w0,4997 +torch/ao/pruning/sparsifier/weight_norm_sparsifier.py,sha256=gIivd_Kyop1c_g5T9hWsYoI6CWC_Mcq1qTZO1P4EKzc,9351 +torch/ao/quantization/__init__.py,sha256=ZDEZfDok59UAWYObnbOqXylLbxU9aBYk7UMFhhH4UHw,6532 +torch/ao/quantization/__pycache__/__init__.cpython-311.pyc,, +torch/ao/quantization/__pycache__/_correct_bias.cpython-311.pyc,, +torch/ao/quantization/__pycache__/_equalize.cpython-311.pyc,, +torch/ao/quantization/__pycache__/_learnable_fake_quantize.cpython-311.pyc,, +torch/ao/quantization/__pycache__/fake_quantize.cpython-311.pyc,, +torch/ao/quantization/__pycache__/fuse_modules.cpython-311.pyc,, +torch/ao/quantization/__pycache__/fuser_method_mappings.cpython-311.pyc,, +torch/ao/quantization/__pycache__/observer.cpython-311.pyc,, +torch/ao/quantization/__pycache__/qconfig.cpython-311.pyc,, +torch/ao/quantization/__pycache__/qconfig_mapping.cpython-311.pyc,, +torch/ao/quantization/__pycache__/quant_type.cpython-311.pyc,, +torch/ao/quantization/__pycache__/quantization_mappings.cpython-311.pyc,, +torch/ao/quantization/__pycache__/quantize.cpython-311.pyc,, +torch/ao/quantization/__pycache__/quantize_fx.cpython-311.pyc,, +torch/ao/quantization/__pycache__/quantize_jit.cpython-311.pyc,, +torch/ao/quantization/__pycache__/stubs.cpython-311.pyc,, +torch/ao/quantization/__pycache__/utils.cpython-311.pyc,, +torch/ao/quantization/_correct_bias.py,sha256=tTf6RanT0gwmTkNfJrPGDCfcf8yHcpyGJ5HUm7UiB-Y,5415 +torch/ao/quantization/_equalize.py,sha256=b3HzjqrX-40SC3NXv_CevxRePFd5SQhLHyYmQ6UthwI,9493 +torch/ao/quantization/_learnable_fake_quantize.py,sha256=o0VkrbUsv7uMr0jDMfga2pr-xU5XLcszvtEQDiv6h9I,7959 +torch/ao/quantization/backend_config/__init__.py,sha256=lnyxe_DTaYkdoSLhxNlUcMPlck2SAngmI8GNcYDHI4M,915 +torch/ao/quantization/backend_config/__pycache__/__init__.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/_common_operator_config_utils.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/_qnnpack_pt2e.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/backend_config.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/executorch.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/fbgemm.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/native.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/onednn.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/qnnpack.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/tensorrt.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/utils.cpython-311.pyc,, +torch/ao/quantization/backend_config/__pycache__/x86.cpython-311.pyc,, +torch/ao/quantization/backend_config/_common_operator_config_utils.py,sha256=EE_d7xO_Gsh4Y__yPOaxnguZItAUSE8WNvcaCJL98J4,27501 +torch/ao/quantization/backend_config/_qnnpack_pt2e.py,sha256=Z0CgzNQw0gA1n96DRblor08FXUtWpEoVf7vYb9IEx5c,6431 +torch/ao/quantization/backend_config/backend_config.py,sha256=I65caYp3oiracD9g0HKRU9b39OxD2iTryccrjy5VeMY,31509 +torch/ao/quantization/backend_config/executorch.py,sha256=t72escVnr3ztTYDpiyHIJlOWTwHgE1m4wn6fziIxyP4,16924 +torch/ao/quantization/backend_config/fbgemm.py,sha256=dCsGDn-PbASye1CPf_9poNM6eOMN10_NUFTQat95-0Q,4208 +torch/ao/quantization/backend_config/native.py,sha256=xKBjjC-uZ8qZsbHluftqC6Tn9vo_6ivcsVYFWLFFQtQ,8242 +torch/ao/quantization/backend_config/onednn.py,sha256=fHj_5RoNjs52M3tN2BFi6H5omY9D8HEBAAZYnp0yZOk,19108 +torch/ao/quantization/backend_config/qnnpack.py,sha256=PkPTzCS-Q3p8MiUq6mytJlFt-Zox7AO2kxbTg5Qx-AU,5400 +torch/ao/quantization/backend_config/tensorrt.py,sha256=Drb7vD3gWVGnSa41C6rYVR7po0qA420Nu3hoOH483-A,3021 +torch/ao/quantization/backend_config/utils.py,sha256=sUEWKDDUsW4R0AbqOpnAn4CK7HnWe2BwtBdlZtRx40A,12582 +torch/ao/quantization/backend_config/x86.py,sha256=4geWLr6mkFhSEYiXFXpkB8NegS05EXhv2fYHTpajNyg,3869 +torch/ao/quantization/fake_quantize.py,sha256=Tkn1yPz8PimU64XyONud5_-nXHDWAgE2j8bqQ-4ptXk,23465 +torch/ao/quantization/fuse_modules.py,sha256=fBhqELU4N8hFcKa2SxS31xPx3rkssxjS2Yr60C_4vAM,6794 +torch/ao/quantization/fuser_method_mappings.py,sha256=vO3IWudszvfv7pHcbAG70X0FMb8XSX27PTfM6Q3PK7k,10951 +torch/ao/quantization/fx/__init__.py,sha256=65h6iR_5XARcYpGS8A2qRza_2y7cvqCowqkwwziYz8M,81 +torch/ao/quantization/fx/__pycache__/__init__.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/_decomposed.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/_equalize.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/_lower_to_native_backend.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/convert.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/custom_config.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/fuse.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/fuse_handler.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/graph_module.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/lower_to_fbgemm.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/lower_to_qnnpack.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/lstm_utils.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/match_utils.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/pattern_utils.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/prepare.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/qconfig_mapping_utils.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/quantize_handler.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/tracer.cpython-311.pyc,, +torch/ao/quantization/fx/__pycache__/utils.cpython-311.pyc,, +torch/ao/quantization/fx/_decomposed.py,sha256=OCdbeG8OiFwo0cU2KEc-xJKGiSi8dWBPI3Yudmuax3A,44280 +torch/ao/quantization/fx/_equalize.py,sha256=3LCxKha2HfEAPRoeOrkjgqBkmZ9hivdmr8aefSQFKuI,40829 +torch/ao/quantization/fx/_lower_to_native_backend.py,sha256=odKq3dU3Mz0Jy6G8GpyWIqzdMELM-I8mDFUVNkXbNCs,57173 +torch/ao/quantization/fx/_model_report/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/ao/quantization/fx/_model_report/__pycache__/__init__.cpython-311.pyc,, +torch/ao/quantization/fx/_model_report/__pycache__/detector.cpython-311.pyc,, +torch/ao/quantization/fx/_model_report/__pycache__/model_report.cpython-311.pyc,, +torch/ao/quantization/fx/_model_report/__pycache__/model_report_observer.cpython-311.pyc,, +torch/ao/quantization/fx/_model_report/__pycache__/model_report_visualizer.cpython-311.pyc,, +torch/ao/quantization/fx/_model_report/detector.py,sha256=1rVf-mlV-AD1QBke-Obabkvm4xLpD6YZnOmhwRZYuWE,76891 +torch/ao/quantization/fx/_model_report/model_report.py,sha256=LjyEonNh1kyqu_sdaqp0aP8xorb7wNQjweEyVxUp2ao,29740 +torch/ao/quantization/fx/_model_report/model_report_observer.py,sha256=BWuxgjbh5Egg1xxq9eIQt3v8gDbxM28AQmhG0N1jUG0,12084 +torch/ao/quantization/fx/_model_report/model_report_visualizer.py,sha256=qnUOY2kvQEsAlyGAQqRtGftdeUR451b2pwuMQEyMuaA,32667 +torch/ao/quantization/fx/convert.py,sha256=Nki8juH7qIKH6STrOoaqUYlPiFz6rfPQeL4AIXRvT-I,60154 +torch/ao/quantization/fx/custom_config.py,sha256=XXczvULZT7A6u8av5PMmgy94viHeOslT73lceU_8N1s,21815 +torch/ao/quantization/fx/fuse.py,sha256=pMxYdgQBYHgNrk5AUOiaaUXyJryX0NeoFJjRrJNtAbs,7396 +torch/ao/quantization/fx/fuse_handler.py,sha256=KlI4Gk2e5bc_8Whfkib45iWLffWnPSU2qvU0QA3-umw,4704 +torch/ao/quantization/fx/graph_module.py,sha256=QG9sVe1kfWGg5o9cyOiUaPrvplbwbWubGYTyGxCwG24,6607 +torch/ao/quantization/fx/lower_to_fbgemm.py,sha256=ERyoustxB5yjBVinb4V2_YTZF-76c7NaPxaVP2l_Pbs,602 +torch/ao/quantization/fx/lower_to_qnnpack.py,sha256=pSLRhkHuxXdbu7a1yEzQPbomFrepOvfiKWpg4HNqEYw,527 +torch/ao/quantization/fx/lstm_utils.py,sha256=dFESL358seHZX4XU46mGv57PouGbmJEh67Iid-cldBA,10544 +torch/ao/quantization/fx/match_utils.py,sha256=fKW4f3YavhquBHtPWdpp6aIp6AtZx17ubkRyeN29Bus,8963 +torch/ao/quantization/fx/pattern_utils.py,sha256=aMM87O63tKrgnLtuAIPEPNXApbzTxL4SIOvyWZTofJI,3668 +torch/ao/quantization/fx/prepare.py,sha256=YUNZAz-YPOL7JPs2GfOCQal-5oZ2oh7zlcwY-0fXndM,87716 +torch/ao/quantization/fx/qconfig_mapping_utils.py,sha256=yGl3Bpb6nGqM9WswYDKlCK2kaxpJ6znTXmcxC3T5scU,15425 +torch/ao/quantization/fx/quantize_handler.py,sha256=jFTXAjdLTYb1MGic0nZDngYmEfHSh_sTbjDCJGorGUU,7325 +torch/ao/quantization/fx/tracer.py,sha256=2haqkImWeOrUUr7V_b_bgYwuIhU6GFHKEY4edqmUJ4M,1737 +torch/ao/quantization/fx/utils.py,sha256=2WqHkDdbJowbL2MHguYHWdHJaAChsA0WQ2wE5ybPFUU,38827 +torch/ao/quantization/observer.py,sha256=NpxUO3XP6rRKcMLTOzL5R8_yCML8fyTB47x6wdj5zzM,80452 +torch/ao/quantization/qconfig.py,sha256=mA_SDFrqWN4MMzOXkdblA-VLOxZeG4igXClVALeG06w,24632 +torch/ao/quantization/qconfig_mapping.py,sha256=dYrgvV5RU3ub3qqCqEI2zEpU4nxbXA6dVsX8Awuxz6o,14845 +torch/ao/quantization/quant_type.py,sha256=fY7JDEr_PzBFymxy9fQmnRIFKiQoL2QoKJtxZKSJsoc,760 +torch/ao/quantization/quantization_mappings.py,sha256=IMePDAbFLx-0MDc2kF7_hAa1SSiPG7bOBzRimlD4yPA,13865 +torch/ao/quantization/quantize.py,sha256=uEm9aKJ1fku7Up2W7-DH2ThUL3LF1WFdbMaucCHlHto,31396 +torch/ao/quantization/quantize_fx.py,sha256=nZk4g-86EdrlvaR2uWYl9_Z-7wANLSHKJsbFDjYfy0A,32464 +torch/ao/quantization/quantize_jit.py,sha256=nW1j383ayR-XgMZO98WuITK9CPe7B18mfuOQD6YWc6U,14699 +torch/ao/quantization/stubs.py,sha256=V-AJl3upvDL1qiTDdTBT7kJfdQ4rwFJT_UG1-A0oEP8,2282 +torch/ao/quantization/utils.py,sha256=7gX90uUbXlMD9Fl7luRQIywjMYVMrz864x5pjeBEcmE,29599 +torch/autograd/__init__.py,sha256=f5TkaU4-MR59bb72No-FA_G-IeXfR8oLqzce0qa42uI,26863 +torch/autograd/__pycache__/__init__.cpython-311.pyc,, +torch/autograd/__pycache__/anomaly_mode.cpython-311.pyc,, +torch/autograd/__pycache__/forward_ad.cpython-311.pyc,, +torch/autograd/__pycache__/function.cpython-311.pyc,, +torch/autograd/__pycache__/functional.cpython-311.pyc,, +torch/autograd/__pycache__/grad_mode.cpython-311.pyc,, +torch/autograd/__pycache__/gradcheck.cpython-311.pyc,, +torch/autograd/__pycache__/graph.cpython-311.pyc,, +torch/autograd/__pycache__/profiler.cpython-311.pyc,, +torch/autograd/__pycache__/profiler_legacy.cpython-311.pyc,, +torch/autograd/__pycache__/profiler_util.cpython-311.pyc,, +torch/autograd/__pycache__/variable.cpython-311.pyc,, +torch/autograd/_functions/__init__.py,sha256=sdKJj6Dia1vNRSolyufzRKdn5qHkgCBCUI3OBm-PGW0,36 +torch/autograd/_functions/__pycache__/__init__.cpython-311.pyc,, +torch/autograd/_functions/__pycache__/tensor.cpython-311.pyc,, +torch/autograd/_functions/__pycache__/utils.cpython-311.pyc,, +torch/autograd/_functions/tensor.py,sha256=8SJkduO10ftxiPMbTlr_l8Vb9rluVnuZWQNaVBpxW0E,2496 +torch/autograd/_functions/utils.py,sha256=XSfqpnmJKm2LuoVc2FVFms4wV3hJxMhFrfbCAytX2JY,753 +torch/autograd/anomaly_mode.py,sha256=HNYmekIC9V3GRtrFVCvXm9c432KFmUlIb6fpi48EY4M,4964 +torch/autograd/forward_ad.py,sha256=eS1kUD2wnLKUVLDroJwfRhOg6xiTr3jbtRfjHUmARUE,7626 +torch/autograd/function.py,sha256=AHa7EaZ3MyQYdTxBuMznxVkqJ5njj6sgMqMo3BbeMmU,34018 +torch/autograd/functional.py,sha256=mOyYtay1nlQ6ojqb8yNpS-l3V7oTgob4csV4NRMoGZ8,53417 +torch/autograd/grad_mode.py,sha256=Ri-3CW75uslJ1USqCK5cUFcdRieahdQWVpaUHNgN4kQ,14650 +torch/autograd/gradcheck.py,sha256=frfC51Wn5FbIIp8nYHKir6r0ap0kekH696zKMJ0FtSQ,92837 +torch/autograd/graph.py,sha256=puEliVe387-hgryq4ZwVu7ht7RpemOuZzymlnoajKqA,32486 +torch/autograd/profiler.py,sha256=hXnhvMRePluGmocRreG7Yxi0gkU90zOwAfb5QuQkVhc,53280 +torch/autograd/profiler_legacy.py,sha256=7MLdb4wNGuoalcvk-Ohj1dGVGUauo5OweEYLGmyTgeQ,12210 +torch/autograd/profiler_util.py,sha256=u918Y--E_DiDW63IC7nAXn05KBWVWvlZh07zI6xnPIo,61153 +torch/autograd/variable.py,sha256=N0cAiO8ZPZ6rtRtU-hJYH_6NVt-zun16CmO_KrZV7-0,391 +torch/backends/__init__.py,sha256=wNSGGZPEL2dEb6YLC25-l0WumRcFK31vCuvdoYk237c,3622 +torch/backends/__pycache__/__init__.cpython-311.pyc,, +torch/backends/_coreml/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/backends/_coreml/__pycache__/__init__.cpython-311.pyc,, +torch/backends/_coreml/__pycache__/preprocess.cpython-311.pyc,, +torch/backends/_coreml/preprocess.py,sha256=Nzzwas90UYSckQQuyjMWDFuzlva73YNGQd4JUCS3VKc,4508 +torch/backends/_nnapi/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/backends/_nnapi/__pycache__/__init__.cpython-311.pyc,, +torch/backends/_nnapi/__pycache__/prepare.cpython-311.pyc,, +torch/backends/_nnapi/__pycache__/serializer.cpython-311.pyc,, +torch/backends/_nnapi/prepare.py,sha256=2t_-EEWffVa7PatU-5BZ_ZiqPemCBjwBBmhbBidE-Oc,6975 +torch/backends/_nnapi/serializer.py,sha256=81wuXWGDt4JyHCvPc9ZSRHN_GXCVh8rmz1cO7B3Aa1o,102138 +torch/backends/cpu/__init__.py,sha256=V3lRML9bHlNvfE5iOHJP9kN-P4WvEN23SeGQLLDYEnM,314 +torch/backends/cpu/__pycache__/__init__.cpython-311.pyc,, +torch/backends/cuda/__init__.py,sha256=guZqAdZ0VpLaT6KnImBAWnPu8cZ4KjXirGE4K1-THyU,21862 +torch/backends/cuda/__pycache__/__init__.cpython-311.pyc,, +torch/backends/cudnn/__init__.py,sha256=wjtpVoOfEK5ikTwZnkzWsKytZIGWZeaEmQN1TRhcGkg,8221 +torch/backends/cudnn/__pycache__/__init__.cpython-311.pyc,, +torch/backends/cudnn/__pycache__/rnn.cpython-311.pyc,, +torch/backends/cudnn/rnn.py,sha256=vduXdLChmzWOdSAsLkDkYqqVY0S9roOegVbvB1n-bJs,4138 +torch/backends/cusparselt/__init__.py,sha256=LCf7CAhlOka2FofcBZP5gh05-3bzf4VtxmwPA9IFdek,1273 +torch/backends/cusparselt/__pycache__/__init__.cpython-311.pyc,, +torch/backends/kleidiai/__init__.py,sha256=V5pt-IFG6zqGlgzrLQVDOakig9WD1QhZbl1AsfkDgTU,162 +torch/backends/kleidiai/__pycache__/__init__.cpython-311.pyc,, +torch/backends/mha/__init__.py,sha256=tUmUrm0Z3I5dnDnhSbuDgZifGDemdgPrdQdsokvVssA,718 +torch/backends/mha/__pycache__/__init__.cpython-311.pyc,, +torch/backends/miopen/__init__.py,sha256=T9aQ5QNnKy9jBHXRZZQD3E3sQ_lbPFYkYx5gX7VRGTU,1208 +torch/backends/miopen/__pycache__/__init__.cpython-311.pyc,, +torch/backends/mkl/__init__.py,sha256=knXrTfpLjP_371pp4V3Qzinf3nBw8MWrRqaLvp5taEA,1823 +torch/backends/mkl/__pycache__/__init__.cpython-311.pyc,, +torch/backends/mkldnn/__init__.py,sha256=V0rvGEXMkw59IzPQitu99ayW9nVbGaTBVx6nLZ0JyDo,4388 +torch/backends/mkldnn/__pycache__/__init__.cpython-311.pyc,, +torch/backends/mps/__init__.py,sha256=863xI79hPErZAR9o6Z0YH0BmTpfMW7T7gVmERd3RJa4,2128 +torch/backends/mps/__pycache__/__init__.cpython-311.pyc,, +torch/backends/nnpack/__init__.py,sha256=6cApdM4SVAbns4TuVAuFI9sXskhu7piYVsmP4rqpdOA,837 +torch/backends/nnpack/__pycache__/__init__.cpython-311.pyc,, +torch/backends/openmp/__init__.py,sha256=h6ebEMpGQavTuGZ3eJVK4JMUMMUgTYcrmvHxh_Fcgy8,157 +torch/backends/openmp/__pycache__/__init__.cpython-311.pyc,, +torch/backends/opt_einsum/__init__.py,sha256=xdL9w-v1UteyrJPFvVheYlxNqeTB7Z2HLpanp5mgMkw,3836 +torch/backends/opt_einsum/__pycache__/__init__.cpython-311.pyc,, +torch/backends/python_native/__init__.py,sha256=dQ5u4VKiujx-t5WG9ld6IC6lkQzdm7rrz82eYM0jI8Y,12662 +torch/backends/python_native/__pycache__/__init__.cpython-311.pyc,, +torch/backends/quantized/__init__.py,sha256=vD2Ds81mKtQGrMrAYy1FAo_OPfyzvhDObF276kT_t-U,1861 +torch/backends/quantized/__pycache__/__init__.cpython-311.pyc,, +torch/backends/xeon/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/backends/xeon/__pycache__/__init__.cpython-311.pyc,, +torch/backends/xeon/__pycache__/run_cpu.cpython-311.pyc,, +torch/backends/xeon/run_cpu.py,sha256=c6UHx3K7-mxVbztS2uSsJvKf9c_OfMtojfvMZKveb-k,38299 +torch/backends/xnnpack/__init__.py,sha256=BNL7CNyTlaTNxu4JVDE8VZgiu16mRO6JpZeQ0Ns3Nr8,702 +torch/backends/xnnpack/__pycache__/__init__.cpython-311.pyc,, +torch/bin/torch_shm_manager,sha256=P5af_RQ-1u7rf6KM0eQ2fpk5xrIrDNSEbjyTWqusv7A,37944 +torch/compiler/__init__.py,sha256=3ISB6Z3B8gKs1k9ULYgyFTLj0DDVx5wtCGZuasJTpx8,29395 +torch/compiler/__pycache__/__init__.cpython-311.pyc,, +torch/compiler/__pycache__/_cache.cpython-311.pyc,, +torch/compiler/__pycache__/config.cpython-311.pyc,, +torch/compiler/_cache.py,sha256=djksgtegrgGz2uFcUprLGh2SD7Blz1D8kMD10eRH58g,11090 +torch/compiler/config.py,sha256=G9dJL-rJ_4EoIizuPgb--_poTDafU_lThKv4NznA204,10724 +torch/contrib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/contrib/__pycache__/__init__.cpython-311.pyc,, +torch/contrib/__pycache__/_tensorboard_vis.cpython-311.pyc,, +torch/contrib/_tensorboard_vis.py,sha256=NHCnPKg_OGbn9EMjqkTVzALEsShQh2N2-sFsBIDGwEo,5875 +torch/cpu/__init__.py,sha256=LR_4YHxgTCLu9emPFRevb6xaAn4s8wn_ha8i6UDVaJ4,6865 +torch/cpu/__pycache__/__init__.cpython-311.pyc,, +torch/cpu/amp/__init__.py,sha256=Q_Rmzfz39fTI2sYsf0hrKJysAR3d2HHMb6KjZlKQurs,103 +torch/cpu/amp/__pycache__/__init__.cpython-311.pyc,, +torch/cpu/amp/__pycache__/autocast_mode.cpython-311.pyc,, +torch/cpu/amp/__pycache__/grad_scaler.cpython-311.pyc,, +torch/cpu/amp/autocast_mode.py,sha256=XC8Kac0XKSERFrp8Oc3W1G77_cOqrRuG_zOWKjua85I,2180 +torch/cpu/amp/grad_scaler.py,sha256=hT5iBOSV285sW4X7s00Mfg7Pi_eJ-sKCS5BT1qSrz-U,958 +torch/csrc/inductor/aoti_runtime/model.h,sha256=bY1OyO-gjO0hChRQiyIP_7lhLEBX4y30zjzC6zD8G7g,2257 +torch/cuda/__init__.py,sha256=7WyRvll7L8G57JvgWhLH2zj2j51Aw1J8GRLnuMXBSXk,70493 +torch/cuda/__pycache__/__init__.cpython-311.pyc,, +torch/cuda/__pycache__/_annotate_cuda_graph_trace.cpython-311.pyc,, +torch/cuda/__pycache__/_device_limits.cpython-311.pyc,, +torch/cuda/__pycache__/_gpu_trace.cpython-311.pyc,, +torch/cuda/__pycache__/_graph_annotations.cpython-311.pyc,, +torch/cuda/__pycache__/_memory_viz.cpython-311.pyc,, +torch/cuda/__pycache__/_pin_memory_utils.cpython-311.pyc,, +torch/cuda/__pycache__/_sanitizer.cpython-311.pyc,, +torch/cuda/__pycache__/_utils.cpython-311.pyc,, +torch/cuda/__pycache__/comm.cpython-311.pyc,, +torch/cuda/__pycache__/gds.cpython-311.pyc,, +torch/cuda/__pycache__/graphs.cpython-311.pyc,, +torch/cuda/__pycache__/green_contexts.cpython-311.pyc,, +torch/cuda/__pycache__/jiterator.cpython-311.pyc,, +torch/cuda/__pycache__/memory.cpython-311.pyc,, +torch/cuda/__pycache__/nccl.cpython-311.pyc,, +torch/cuda/__pycache__/nvtx.cpython-311.pyc,, +torch/cuda/__pycache__/profiler.cpython-311.pyc,, +torch/cuda/__pycache__/random.cpython-311.pyc,, +torch/cuda/__pycache__/sparse.cpython-311.pyc,, +torch/cuda/__pycache__/streams.cpython-311.pyc,, +torch/cuda/__pycache__/tunable.cpython-311.pyc,, +torch/cuda/_annotate_cuda_graph_trace.py,sha256=7Exb4ZKPl9-56L9-zSWv0KsduGbGDcUC7mb4KVbRZ0k,12767 +torch/cuda/_device_limits.py,sha256=fYfWl5kv31rsLXJs8ASkz3AYrpnC1O2phMQy8zIB-aw,5496 +torch/cuda/_gpu_trace.py,sha256=OMYsQlXjF7PvEOGCaMWw__K_iwQYbilVdRV6xoJZvDI,2386 +torch/cuda/_graph_annotations.py,sha256=VVPAFZmqQIR86xa4ilI7d4SX0GVURWH3xp25cyv6pfM,15045 +torch/cuda/_memory_viz.py,sha256=LU_40JI2LWLf_T8m-lWEgyMM2QBGVqllQbbax413k74,28686 +torch/cuda/_pin_memory_utils.py,sha256=uk6Tnea02s5jWmKIdD2JX7p4W4H-74xpKNf7ZryTFqs,773 +torch/cuda/_sanitizer.py,sha256=fJc3BYVYlk2B4frXL_KpBEA7LjD77pS2IddKqw5aZbQ,24514 +torch/cuda/_utils.py,sha256=TBFaEOY_5IOvDlHsRNnizYvPDVWcs3K4KI98sEVepQk,21018 +torch/cuda/amp/__init__.py,sha256=cyLIhZXcuKMKb54GjtdhC0GvA1r018SNByosz5S0gFg,298 +torch/cuda/amp/__pycache__/__init__.cpython-311.pyc,, +torch/cuda/amp/__pycache__/autocast_mode.cpython-311.pyc,, +torch/cuda/amp/__pycache__/common.cpython-311.pyc,, +torch/cuda/amp/__pycache__/grad_scaler.cpython-311.pyc,, +torch/cuda/amp/autocast_mode.py,sha256=sVhM2WjWcFXHOp-_6l5ycFZHOUfSsqS5ee0SsBISgYc,3477 +torch/cuda/amp/common.py,sha256=rPnt6OLsROsahVbMyX5w5Fl3R4G6VUiSl5uioKuDEFg,230 +torch/cuda/amp/grad_scaler.py,sha256=WTG3BGVvo4ifizDYTNTq18lRp_CkqVbYS7Kv5N9n6ag,1073 +torch/cuda/comm.py,sha256=DjabyiIf7WmKgD84HmP4_XYd3Cau_7X_E7dcW1UAUjg,344 +torch/cuda/gds.py,sha256=qFygo1HhUPuysVg4Rqhvklqyhn2brl1U0t27GbMVe1A,6361 +torch/cuda/graphs.py,sha256=8IlmR41W0fF51hpEtDF_ag5NAXmk5EHwNPo0qAHKdQw,29872 +torch/cuda/green_contexts.py,sha256=gEIpCN5dkNtTC5UbFE6sTMIoqeuGFh7y65-fCoLfQbU,3622 +torch/cuda/jiterator.py,sha256=UNAoi692RWMgNPmXBaNhEHyKgoWyGZl3_c1CHvgnSaY,6991 +torch/cuda/memory.py,sha256=5tondwGljSt8alD_5WgKu0o1-bj2EMYKBKJw8nWHRuQ,52420 +torch/cuda/nccl.py,sha256=SrAgsI9xG6pd5hxMEtFZkeEK7fgY0RCgk-NdHovmuvo,4535 +torch/cuda/nvtx.py,sha256=iAOj9LvGDv47y6bctZPrfyp6OeAyg4mqP17TtvXIrvg,3705 +torch/cuda/profiler.py,sha256=aYzsY7dpmV-vGvwKl3uUPWhCC0VDiP02FmNMNkGbkAU,1230 +torch/cuda/random.py,sha256=a7N7kfaMoKKFSulZ0AQObsImnLOJK2G6-mMNLLjHhRo,5400 +torch/cuda/sparse.py,sha256=H912FRisikGM9SSVDQQq-YX5TNIilBYBZs1AwmIv3GQ,67 +torch/cuda/streams.py,sha256=OzTtCLZ88uxBHhpIOzGDW7_4ZSyRWgjyoOTBQONC3Zw,10672 +torch/cuda/tunable.py,sha256=tzo2wbwbTe7J5fTVa95lC9P5BWK4ujsSM9K1FlaCLDc,32063 +torch/distributed/__init__.py,sha256=3edNghejxOMdliGU_olHuPV3lGvAA48uv6UuGtPaJjs,5685 +torch/distributed/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/__pycache__/_checkpointable.cpython-311.pyc,, +torch/distributed/__pycache__/_composable_state.cpython-311.pyc,, +torch/distributed/__pycache__/_dist2.cpython-311.pyc,, +torch/distributed/__pycache__/_functional_collectives.cpython-311.pyc,, +torch/distributed/__pycache__/_functional_collectives_impl.cpython-311.pyc,, +torch/distributed/__pycache__/_mesh_layout.cpython-311.pyc,, +torch/distributed/__pycache__/_meta_registrations.cpython-311.pyc,, +torch/distributed/__pycache__/_serialization.cpython-311.pyc,, +torch/distributed/__pycache__/_state_dict_utils.cpython-311.pyc,, +torch/distributed/__pycache__/argparse_util.cpython-311.pyc,, +torch/distributed/__pycache__/c10d_logger.cpython-311.pyc,, +torch/distributed/__pycache__/collective_utils.cpython-311.pyc,, +torch/distributed/__pycache__/config.cpython-311.pyc,, +torch/distributed/__pycache__/constants.cpython-311.pyc,, +torch/distributed/__pycache__/device_mesh.cpython-311.pyc,, +torch/distributed/__pycache__/distributed_c10d.cpython-311.pyc,, +torch/distributed/__pycache__/launch.cpython-311.pyc,, +torch/distributed/__pycache__/logging_handlers.cpython-311.pyc,, +torch/distributed/__pycache__/remote_device.cpython-311.pyc,, +torch/distributed/__pycache__/rendezvous.cpython-311.pyc,, +torch/distributed/__pycache__/run.cpython-311.pyc,, +torch/distributed/__pycache__/utils.cpython-311.pyc,, +torch/distributed/_checkpointable.py,sha256=VqtUDzJDcTeMUFOguRIo_AvLRXg4n4U73-frKynBqo8,1305 +torch/distributed/_composable/__init__.py,sha256=SUm5fu8Tg24Upb5iHCV849oZSrrtSacaG-7T7A5L2qA,125 +torch/distributed/_composable/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_composable/__pycache__/checkpoint_activation.cpython-311.pyc,, +torch/distributed/_composable/__pycache__/contract.cpython-311.pyc,, +torch/distributed/_composable/__pycache__/replicate.cpython-311.pyc,, +torch/distributed/_composable/__pycache__/replicate_with_fsdp.cpython-311.pyc,, +torch/distributed/_composable/checkpoint_activation.py,sha256=ooriyRjD7aeGu_EKO0zHwpYCIJQxT10D5WjnhwHvWac,4819 +torch/distributed/_composable/contract.py,sha256=qfezUWwTqhPcTeyjjEXn_XjYhyRXaQaLxIhcrj6xwDQ,10447 +torch/distributed/_composable/fsdp/__init__.py,sha256=PAE0ViGvjHO65XbIe-CM5BgRx9KNsVL4xU7QRUyqL9M,169 +torch/distributed/_composable/fsdp/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_composable/fsdp/__pycache__/fully_shard.cpython-311.pyc,, +torch/distributed/_composable/fsdp/fully_shard.py,sha256=LggyTExDS-LXmmZeCoSp-GvlMKOmRulWv6k4aPyYsWM,240 +torch/distributed/_composable/replicate.py,sha256=sWxlvMYTkK2DkEPFgiAzrioJH0-tk6r2WTyw7FmeTcY,9350 +torch/distributed/_composable/replicate_with_fsdp.py,sha256=nMYPrR5CrWJsFH6maoEeWQMwZUIh4Y065ZqodwtNfV0,6485 +torch/distributed/_composable_state.py,sha256=rOKLQZE5s1TbgdfB4M9kd94AaTAmnnw-AmyqL5fxpjQ,1403 +torch/distributed/_dist2.py,sha256=us0LVkzxrelZO7rlZhV4vYf5HTNgHKenJGcAJefh_40,4835 +torch/distributed/_functional_collectives.py,sha256=rTGuWOycrrhHVP_deZl4ctrIPLT-rcczEgsJHJNmyVo,68214 +torch/distributed/_functional_collectives_impl.py,sha256=gwLRcqtq9NGUBkXDhYElw6PpDVtDmKkWew5Bt9ZJZsc,3790 +torch/distributed/_local_tensor/__init__.py,sha256=97TbeCcz5IlYgUho4WYiEPLFNdnE5Mi10XfC95flumw,81709 +torch/distributed/_local_tensor/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_local_tensor/__pycache__/_c10d.cpython-311.pyc,, +torch/distributed/_local_tensor/_c10d.py,sha256=WlDeRmyk63tOPPAfT1tcgp4gHP23T632e7IUwzIVhO0,41562 +torch/distributed/_mesh_layout.py,sha256=mxfQss0pyrV2P_9zBQO9C99TtK_f4Rz_m2wwxaCAfzY,13132 +torch/distributed/_meta_registrations.py,sha256=dfpshci3NHj3PTMJYtytDuPvE61BFUnzAWfeyfF3r1g,2395 +torch/distributed/_ops/__init__.py,sha256=jLQo34kgGOSuumAbVsa8rQRJtt8LZKnXvqpTvpsNSjo,52 +torch/distributed/_ops/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_ops/__pycache__/device_mesh.cpython-311.pyc,, +torch/distributed/_ops/device_mesh.py,sha256=82Fy4vP7lggJbCideZaQhdUZKfWYFhCAN_bpKuVsko0,3942 +torch/distributed/_pycute/__init__.py,sha256=7Rjy64t4ztEB5BNIFIjFPaYxv1afzff1HcKLBtnPYx8,2456 +torch/distributed/_pycute/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_pycute/__pycache__/int_tuple.cpython-311.pyc,, +torch/distributed/_pycute/__pycache__/layout.cpython-311.pyc,, +torch/distributed/_pycute/__pycache__/typing.cpython-311.pyc,, +torch/distributed/_pycute/int_tuple.py,sha256=HZegxSKHMl86jE7BRIMVPV83qdq5FQrvxp-xu3apPlU,10304 +torch/distributed/_pycute/layout.py,sha256=Zdp-gGy6U2C5H4LZHUCFXz3J_pllw--7xhPyu_bozYk,18452 +torch/distributed/_pycute/typing.py,sha256=0kzCgBGBRRt1hfAACBnwFdC9g11Ujh8dozJ5GAd9tDw,2062 +torch/distributed/_serialization.py,sha256=nXtpB29C_HiU3WCeWcEISJCAdN0FbOFNAyx3YxjUwaw,4584 +torch/distributed/_shard/__init__.py,sha256=9mrjpti4iifWb8jnGoZypHH0ovdZg3A2KQnxg6BACKI,87 +torch/distributed/_shard/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_shard/__pycache__/_utils.cpython-311.pyc,, +torch/distributed/_shard/__pycache__/api.cpython-311.pyc,, +torch/distributed/_shard/__pycache__/common_op_utils.cpython-311.pyc,, +torch/distributed/_shard/__pycache__/metadata.cpython-311.pyc,, +torch/distributed/_shard/__pycache__/op_registry_utils.cpython-311.pyc,, +torch/distributed/_shard/__pycache__/sharder.cpython-311.pyc,, +torch/distributed/_shard/_utils.py,sha256=ra_iuevX3gIRVWOz9tql3hLtLRm0JXzeREiL9vwVm14,1070 +torch/distributed/_shard/api.py,sha256=hsDn-zXLsKOR-_exVqTGkILCa-B3lSzU8u_FYaHQd2M,12365 +torch/distributed/_shard/checkpoint/__init__.py,sha256=kILV09-2Glss17itgmv1SbBq7hp2Y7HLUugCZRGYCFE,584 +torch/distributed/_shard/checkpoint/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_shard/common_op_utils.py,sha256=kj53BqGnob9GRuoRgzn1MXOAmtFRvy_7nzXdJA_1Qrg,2148 +torch/distributed/_shard/metadata.py,sha256=krJyeG9SY2KiJtVxUq4th4qDFUR4WbW5j01VA9RO8bA,2168 +torch/distributed/_shard/op_registry_utils.py,sha256=aK7aBTzgR298AvchxXB1VNdeIT6UCO5emnOoOadjM-I,1031 +torch/distributed/_shard/sharded_optim/__init__.py,sha256=23n5rJ4CKIZb1TwXlI7cX6DgOURqKO5eiM7PJbyi9aY,1863 +torch/distributed/_shard/sharded_optim/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_shard/sharded_optim/__pycache__/api.cpython-311.pyc,, +torch/distributed/_shard/sharded_optim/api.py,sha256=azG0U2MZOpi5qY2TyEy0qyFgNemadga_bSOu1ooL438,4264 +torch/distributed/_shard/sharded_tensor/__init__.py,sha256=uatjwQJIrccr9STR1IqUkrY_hcLNJAf3RPvvbOmY8P8,19247 +torch/distributed/_shard/sharded_tensor/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/__pycache__/api.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/__pycache__/logger.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/__pycache__/logging_handlers.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/__pycache__/metadata.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/__pycache__/reshard.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/__pycache__/shard.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/__pycache__/utils.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/_ops/__init__.py,sha256=NW4TXyw5RGFCbWfikCuxxCOrWuVIUovX06jI8PgcqNM,498 +torch/distributed/_shard/sharded_tensor/_ops/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/_ops/__pycache__/_common.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/_ops/__pycache__/binary_cmp.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/_ops/__pycache__/init.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/_ops/__pycache__/misc_ops.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/_ops/__pycache__/tensor_ops.cpython-311.pyc,, +torch/distributed/_shard/sharded_tensor/_ops/_common.py,sha256=KuLeI0PylCkKm5_pnK61Qo3YQhuiFte_VwYpMxjdVG4,4276 +torch/distributed/_shard/sharded_tensor/_ops/binary_cmp.py,sha256=IQ9RiKLLcj4MsJQLb-PCxVCbDfDmAh06FIqLjoV4kS8,2736 +torch/distributed/_shard/sharded_tensor/_ops/init.py,sha256=TdtFNapDXAcB_qL2G7MqbSKfOewocl_MwNtxbsDY5HU,6070 +torch/distributed/_shard/sharded_tensor/_ops/misc_ops.py,sha256=lp_SXuhFUagN_5bxfmZVeqY0SCMhgnGeLGxLq38qvsg,497 +torch/distributed/_shard/sharded_tensor/_ops/tensor_ops.py,sha256=76Uf_goUHNZPpyfVk6nnNmFnmhm9brOR08O11DMaIZ4,7813 +torch/distributed/_shard/sharded_tensor/api.py,sha256=Jdn4hua557q-qdMt21XdHORpQ1Qxe8a8r1nCntm5WWA,55182 +torch/distributed/_shard/sharded_tensor/logger.py,sha256=yny64ZlZ-RzlS5DcX-SAipWqk4xcR2AQO4yl4Vf8B50,1084 +torch/distributed/_shard/sharded_tensor/logging_handlers.py,sha256=nXuxHnOEhAHqqDPhTk_-aLYIVHoZpQEYW_X4N1Qg0_Q,359 +torch/distributed/_shard/sharded_tensor/metadata.py,sha256=92isIH0O2TqPy3Li1CG00R0Ca1BZ40gGM2GHKkJraUA,2998 +torch/distributed/_shard/sharded_tensor/reshard.py,sha256=p_h7m8AxDJ1S2Ei86zcD31DQ89PhpzGOoWEcFU8_TeU,10726 +torch/distributed/_shard/sharded_tensor/shard.py,sha256=zNfIL2FIGBNwPUFcrfusf0E1R-YO04ZoOdcvY6DpXbA,2361 +torch/distributed/_shard/sharded_tensor/utils.py,sha256=hsWa2xhcZ2PAOzfu66VtZvXR0eW_6zVI2BDTpDpdC5c,11781 +torch/distributed/_shard/sharder.py,sha256=nErlvDuMU29IUB6oWB295FaXYj8H42FuDQJXdepvQ1E,901 +torch/distributed/_shard/sharding_plan/__init__.py,sha256=9w-j8bY8VJ0dGT_ueJR2yTfzuech180Xkq0XELn1tKs,47 +torch/distributed/_shard/sharding_plan/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_shard/sharding_plan/__pycache__/api.cpython-311.pyc,, +torch/distributed/_shard/sharding_plan/api.py,sha256=k3rOiiRTSOf7SVjoScGmP6-AtSOoYpRBLaOjOAdywE4,3602 +torch/distributed/_shard/sharding_spec/__init__.py,sha256=hih75sKTSXPxdyUOS24Vl6K3Q3ht-74l33uAjM4I2Kw,291 +torch/distributed/_shard/sharding_spec/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_shard/sharding_spec/__pycache__/_internals.cpython-311.pyc,, +torch/distributed/_shard/sharding_spec/__pycache__/api.cpython-311.pyc,, +torch/distributed/_shard/sharding_spec/__pycache__/chunk_sharding_spec.cpython-311.pyc,, +torch/distributed/_shard/sharding_spec/_internals.py,sha256=VZz-fES-f5tcJhWJ1IK4A_NAmOf06wFl8NwA6GBnk-0,8517 +torch/distributed/_shard/sharding_spec/api.py,sha256=fdg-lx3naoZ1m4d8NuMYRcpXHPeFiDn-QEAUdjxFV4U,9848 +torch/distributed/_shard/sharding_spec/chunk_sharding_spec.py,sha256=ImU70gOX7JHSuNN59DSj_t_bXARlJzZ144MQW7cvIig,9368 +torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/_common.cpython-311.pyc,, +torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/embedding.cpython-311.pyc,, +torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/__pycache__/embedding_bag.cpython-311.pyc,, +torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/_common.py,sha256=pU0FdT98noqNnNMxs9H4hBrNi_wom1SZgm4UjDGPUXg,13121 +torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding.py,sha256=vCSV8WXu8DSO9BQon8hNE8mf4xiqbTAvNv6o_3jpj-U,11209 +torch/distributed/_shard/sharding_spec/chunk_sharding_spec_ops/embedding_bag.py,sha256=jFU9qfivU0p5awLa6YjSVYiaBy8nsmvdZ4RCxVRDBHk,18319 +torch/distributed/_sharded_tensor/__init__.py,sha256=sL4DKv-CLF-9k_iqL2U689tEIJCSKgeWLsKNtjoqjdg,617 +torch/distributed/_sharded_tensor/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_sharding_spec/__init__.py,sha256=PR7htIt_o1lIdWq_dvH3aXMHA22fj59dOpLT15vx3zE,646 +torch/distributed/_sharding_spec/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_state_dict_utils.py,sha256=U1OJN7xo04iZeDI-oFDa9clARD9TN9r2unHkTy3ynZM,29743 +torch/distributed/_symmetric_memory/__init__.py,sha256=zaqOHbZivBJ1k1DWX1X9BX6fNLcxxS0kAzqvj1elGSo,82245 +torch/distributed/_symmetric_memory/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_symmetric_memory/__pycache__/_nvshmem_triton.cpython-311.pyc,, +torch/distributed/_symmetric_memory/_nvshmem_triton.py,sha256=Wc-y6gAAGAdFTM0Ff5lY4FKCLJZmVCe9fDN5N1XSvaA,46962 +torch/distributed/_tensor/__init__.py,sha256=6PYDt4AcQDMnCvbcTdOSNy1z-3kiftGBIGxXnJWT0PM,969 +torch/distributed/_tensor/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_tensor/__pycache__/api.cpython-311.pyc,, +torch/distributed/_tensor/__pycache__/placement_types.cpython-311.pyc,, +torch/distributed/_tensor/api.py,sha256=ErftfbUVGCDavz7FfrQ0rEVlfYC0XmcTiWyjb4IZ1BY,300 +torch/distributed/_tensor/placement_types.py,sha256=pvi_dHBNCsZl1C-tAM9sdfYOiKb0Jwr3SqyAWMKgfo8,384 +torch/distributed/_tools/__init__.py,sha256=DYVOrVO_6vBZuEuW8EaKdxsR8bFVoE16R2DiGfKo6JM,327 +torch/distributed/_tools/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/common_utils.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/fake_collectives.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/fsdp2_mem_tracker.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/ilp_utils.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/mem_tracker.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/memory_tracker.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/mod_tracker.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/runtime_estimator.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/sac_estimator.cpython-311.pyc,, +torch/distributed/_tools/__pycache__/sac_ilp.cpython-311.pyc,, +torch/distributed/_tools/common_utils.py,sha256=4l1wP19s0oZrNIIliVqH37sZ5Y9qtNKDYuvlia3vvSw,1556 +torch/distributed/_tools/fake_collectives.py,sha256=7FVey50Xy6sZ0NIYJkXKyGDbxgx7b9ZqcQ50KpAIyb4,10051 +torch/distributed/_tools/fsdp2_mem_tracker.py,sha256=HZM0gICkmUBagWsi7wBpZ4_NcjNUis2vSICDhrIVre0,25233 +torch/distributed/_tools/ilp_utils.py,sha256=xmsBK7RtUTyKVgKMbMJ2_pVxS7kLmFT3QX8tc-CL6rY,10114 +torch/distributed/_tools/mem_tracker.py,sha256=k-g2IocvPxs1oQHFZ5tMs2ePQOfj4bUHaHKHcUZ_zJw,42476 +torch/distributed/_tools/memory_tracker.py,sha256=Et-MJb93th5_MVTmNwtqcY3FN21ZJ1-YIAeqmA8FPpU,11718 +torch/distributed/_tools/mod_tracker.py,sha256=GGqQCn4sX5J9s8ReLPMC4CTpBVgrYB7yXOEamu3LxKw,10267 +torch/distributed/_tools/runtime_estimator.py,sha256=NQHFZRDIwYIMPLg-NtD2gAlPpq1nE7TOBauEmIq6CtI,16849 +torch/distributed/_tools/sac_estimator.py,sha256=tk9rLjEURHgSChuARvgJA5CdumxI5wrz6N34XVRPT1Q,42504 +torch/distributed/_tools/sac_ilp.py,sha256=gXQAz3ul17XO-ZtY6ZSs0OJpGSNqZyVDP_ikpRyRB8U,11287 +torch/distributed/algorithms/__init__.py,sha256=e5HJMsdSDkPOxNWYUh9YvO26GDODG5PBedTP2-Y16Nw,43 +torch/distributed/algorithms/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/algorithms/__pycache__/join.cpython-311.pyc,, +torch/distributed/algorithms/_checkpoint/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/algorithms/_checkpoint/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/algorithms/_checkpoint/__pycache__/checkpoint_wrapper.cpython-311.pyc,, +torch/distributed/algorithms/_checkpoint/checkpoint_wrapper.py,sha256=0CETclY_Mg34hjDdVTq6l_7PPoco9Cxq_E20G2Rn8dE,12001 +torch/distributed/algorithms/_comm_hooks/__init__.py,sha256=LkI4VBMJ6_1KGG3Nz0bPxGFOO-ZjfBj-38FNROlJt8Q,131 +torch/distributed/algorithms/_comm_hooks/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/algorithms/_comm_hooks/__pycache__/default_hooks.cpython-311.pyc,, +torch/distributed/algorithms/_comm_hooks/default_hooks.py,sha256=lAqlZiFeVsIH_mPUtOEHp-KXRxILRhuqsqNs2R6P0SQ,7616 +torch/distributed/algorithms/_optimizer_overlap/__init__.py,sha256=zuKlfE0DcQCZm0av9HrJZfXH6R5AzV2K3xrEAAgoVsk,52 +torch/distributed/algorithms/_optimizer_overlap/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/algorithms/_optimizer_overlap/__pycache__/optimizer_overlap.cpython-311.pyc,, +torch/distributed/algorithms/_optimizer_overlap/optimizer_overlap.py,sha256=zkw6dsk0wrK9Nj4-IYQU0q2yH2hzIL2tag5W9ccz9Oc,3753 +torch/distributed/algorithms/_quantization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/algorithms/_quantization/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/algorithms/_quantization/__pycache__/quantization.cpython-311.pyc,, +torch/distributed/algorithms/_quantization/quantization.py,sha256=jbL_YAS9g8PVEjCk1hgi0b3A992tKevudfyndlGYWiA,5646 +torch/distributed/algorithms/ddp_comm_hooks/__init__.py,sha256=hd0jCD96_JI6-8hHm2K5Sr8In4nD8a3P95i5CgsdPUs,4210 +torch/distributed/algorithms/ddp_comm_hooks/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/algorithms/ddp_comm_hooks/__pycache__/ddp_zero_hook.cpython-311.pyc,, +torch/distributed/algorithms/ddp_comm_hooks/__pycache__/debugging_hooks.cpython-311.pyc,, +torch/distributed/algorithms/ddp_comm_hooks/__pycache__/default_hooks.cpython-311.pyc,, +torch/distributed/algorithms/ddp_comm_hooks/__pycache__/mixed_precision_hooks.cpython-311.pyc,, +torch/distributed/algorithms/ddp_comm_hooks/__pycache__/optimizer_overlap_hooks.cpython-311.pyc,, +torch/distributed/algorithms/ddp_comm_hooks/__pycache__/post_localSGD_hook.cpython-311.pyc,, +torch/distributed/algorithms/ddp_comm_hooks/__pycache__/powerSGD_hook.cpython-311.pyc,, +torch/distributed/algorithms/ddp_comm_hooks/__pycache__/quantization_hooks.cpython-311.pyc,, +torch/distributed/algorithms/ddp_comm_hooks/ddp_zero_hook.py,sha256=JJyb96z4iT3Z8e04ljj8BzCK58Remp_roKXxwGP1u6o,19840 +torch/distributed/algorithms/ddp_comm_hooks/debugging_hooks.py,sha256=X49n0X3TnX650I-DeU-XtL_ssKjdYsn0RYyhB2Ndcms,1115 +torch/distributed/algorithms/ddp_comm_hooks/default_hooks.py,sha256=_v7QSzzDWZmzI3SrRu40IPdClkVGfTQrCsISFVy0Kus,7962 +torch/distributed/algorithms/ddp_comm_hooks/mixed_precision_hooks.py,sha256=sc2NnDmQKmSNYyZ6A83-s9DgYBV-YTfmMdrzurNouag,3254 +torch/distributed/algorithms/ddp_comm_hooks/optimizer_overlap_hooks.py,sha256=LXDz72NQAeDRBRvTAIbnb1pdoyG3jUxTHAr5tfWPZn8,6152 +torch/distributed/algorithms/ddp_comm_hooks/post_localSGD_hook.py,sha256=MPpCW3JHvFIqfC3frNnzuATOtlLiML6hXEgM6C8ZeGk,5150 +torch/distributed/algorithms/ddp_comm_hooks/powerSGD_hook.py,sha256=iAt6WrcK3SLvNpSa63rW2BkqBcn88qFRP_Kvk1dTn0w,40590 +torch/distributed/algorithms/ddp_comm_hooks/quantization_hooks.py,sha256=jDQwXM6McOwPUH5S6Sj9H-hV2rqT-NG0c-InuuyBLsM,8318 +torch/distributed/algorithms/join.py,sha256=r29MXh8j7EhEYR7CDPdk5eiqLYh03d3HGU5p_ldscMA,13520 +torch/distributed/algorithms/model_averaging/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/algorithms/model_averaging/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/algorithms/model_averaging/__pycache__/averagers.cpython-311.pyc,, +torch/distributed/algorithms/model_averaging/__pycache__/hierarchical_model_averager.cpython-311.pyc,, +torch/distributed/algorithms/model_averaging/__pycache__/utils.cpython-311.pyc,, +torch/distributed/algorithms/model_averaging/averagers.py,sha256=sNQvYipPY_JNWeF_zqPUCFXcYR8WLHt4ohqbKQdjkdQ,5417 +torch/distributed/algorithms/model_averaging/hierarchical_model_averager.py,sha256=ILK-Ra0ExrJMV1yhygp2dJsBxnBrOykdbJPf044Xx3o,9772 +torch/distributed/algorithms/model_averaging/utils.py,sha256=_H_E5K-EQv2p5RbFiVa8LH7zjds4IlKiwl6pYMv2DAA,3087 +torch/distributed/argparse_util.py,sha256=NFxjw2asYt06aFffcu6rhJpC3FFthwKLu9wvWEm_qQU,3903 +torch/distributed/autograd/__init__.py,sha256=o1YwfjRin8NfYkCjzcdABFT0bGRGAiuYmSLEwmbmeJ4,1931 +torch/distributed/autograd/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/c10d_logger.py,sha256=t7WGbk08OEfiTz_kiPGHy9kKyUlVokqY3BDw2WyZtkk,3185 +torch/distributed/checkpoint/__init__.py,sha256=OCW4GkIntMFO-DI2AK2jx4IH41jZsA338tqZ2NRBo30,826 +torch/distributed/checkpoint/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_async_executor.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_async_process_executor.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_async_thread_executor.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_checkpointer.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_consolidate_hf_safetensors.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_dedup_save_plans.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_dedup_tensors.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_extension.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_fsspec_filesystem.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_hf_utils.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_nested_dict.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_pg_transport.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_sharded_tensor_utils.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_state_dict_stager.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_storage_utils.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_traverse.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/_version.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/api.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/default_planner.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/filesystem.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/format_utils.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/hf_storage.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/logger.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/logging_handlers.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/metadata.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/optimizer.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/planner.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/planner_helpers.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/quantized_hf_storage.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/resharding.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/staging.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/state_dict.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/state_dict_loader.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/state_dict_saver.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/stateful.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/storage.cpython-311.pyc,, +torch/distributed/checkpoint/__pycache__/utils.cpython-311.pyc,, +torch/distributed/checkpoint/_async_executor.py,sha256=tTpnGTOG1U8BiROCkS9JB_UPv7E6ZvLWFkOboqisMTw,1154 +torch/distributed/checkpoint/_async_process_executor.py,sha256=mtdiwRF7b4u3R3h0W5ju7XG4KsSfD00wDwXAkUu5yLs,18045 +torch/distributed/checkpoint/_async_thread_executor.py,sha256=ojp5WJc383zu1R5B1GqgXdK9acA-ctCph12HAJsRMkA,2401 +torch/distributed/checkpoint/_checkpointer.py,sha256=bK2cKf324BRJmhz9DtCa3AmRWIs4B-aCOGErMUF-UCo,3783 +torch/distributed/checkpoint/_consolidate_hf_safetensors.py,sha256=BI-up1pLOlr4rVUzhic6mbyF6kBhHZJ62WdPcFmn0WU,28448 +torch/distributed/checkpoint/_dedup_save_plans.py,sha256=QP__wSP0I0jglKEyPZ_JGXrgHIKSrU6c6Gow8vQTZkk,2754 +torch/distributed/checkpoint/_dedup_tensors.py,sha256=MU2RaM-qHxvn2I-LmNTgvw1cvGwoqItiCHxhb5vNrRM,1991 +torch/distributed/checkpoint/_experimental/__init__.py,sha256=ukQ1Zu9_BedHvl_JQrJhsWKwcJ8nxEYh4_AzvevbT9k,1761 +torch/distributed/checkpoint/_experimental/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/barriers.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/builder.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/checkpoint_process.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/checkpoint_reader.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/checkpoint_writer.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/checkpointer.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/config.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/staging.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/types.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/__pycache__/utils.cpython-311.pyc,, +torch/distributed/checkpoint/_experimental/barriers.py,sha256=57Tgfop5ZTxmx6fBaNNwOjBiFAICmNT0KWnHoFm0paQ,9134 +torch/distributed/checkpoint/_experimental/builder.py,sha256=tn8g361poE45V-G5bLhD-0-bBx8iGWx97z7sSqyi41k,6101 +torch/distributed/checkpoint/_experimental/checkpoint_process.py,sha256=sOjG5jg-50aiJfEd6yHj6CtnmGod6FsDTHhkrejttPY,13049 +torch/distributed/checkpoint/_experimental/checkpoint_reader.py,sha256=Dw2xF4Fk5aO9gNNw1fHVNO6Jj9JT3F30KzS0jJUAl_I,8679 +torch/distributed/checkpoint/_experimental/checkpoint_writer.py,sha256=ntVXEIcTqiH74yjE3c_v7_j_DlwMr5dpdVZ3wnRZ4ls,5251 +torch/distributed/checkpoint/_experimental/checkpointer.py,sha256=EIDf-R3iJCqs6X1BxMJZfWjEuSg8GKdqWqBPqGjKRcI,11917 +torch/distributed/checkpoint/_experimental/config.py,sha256=81BcsUGoE59ibhFW0OQ05Hi4zmqJ2-v9tITPJqMQDTM,1528 +torch/distributed/checkpoint/_experimental/staging.py,sha256=xn8oX3VNmh16icfCCDqmvPV1n8t4dCBjQEkoXIe8-Jg,8120 +torch/distributed/checkpoint/_experimental/types.py,sha256=8BI2ouC6nlgXa3nKOjcDPVDZaLwDb__S_f1G79hR9Eg,759 +torch/distributed/checkpoint/_experimental/utils.py,sha256=rIaJLYRF03P1FPH57lo17JUD-b8Kd0NhHer1JdpEZa0,1305 +torch/distributed/checkpoint/_extension.py,sha256=m0DPQCZV1yn9ssXcW4xonhyM1AHSCImJJykLk1CLF_Y,7737 +torch/distributed/checkpoint/_fsspec_filesystem.py,sha256=dLJqJtESqkjrA7_hnI5XUDWu1fyRsMfnk7LeNyoUBIM,5473 +torch/distributed/checkpoint/_hf_utils.py,sha256=RoH7XHPrRd23_XcnlvobWlhPgKUrtcv7y1ifd9B7_Vw,2819 +torch/distributed/checkpoint/_nested_dict.py,sha256=IPtO0mCEYB9xT0Il5g4GW978lubbCFCj51Qo9mTDJWo,2273 +torch/distributed/checkpoint/_pg_transport.py,sha256=Dor7iQYkLmCk9PAveO3Jk_5TbKEzdP4GnmAoy9LCyLI,13268 +torch/distributed/checkpoint/_sharded_tensor_utils.py,sha256=HZOdtN7SH2D79rgdm2g7BDEacEwTxkAJWZhNtSVv6O0,4144 +torch/distributed/checkpoint/_state_dict_stager.py,sha256=oSmI5zXC2Aava3bHtmKu-C_P1l-j6HWgtUuqSQmuAM4,18893 +torch/distributed/checkpoint/_storage_utils.py,sha256=adK0EQsZ7U9w-zh2qskmXkn27pmsP1uV5phseLurw68,1362 +torch/distributed/checkpoint/_traverse.py,sha256=f77M9DnlLh8m_FsUa7av5H8DwJgrxSIJBcehD9PTvQ4,7197 +torch/distributed/checkpoint/_version.py,sha256=wjFuEFA3gH625gFQMVRpt55yBAUObGUHe_1iuM8xMVw,90 +torch/distributed/checkpoint/api.py,sha256=iBZAi6ydUSefDNOLzBDZs0GK7Ji5BFHaNfS5lSqbepE,1799 +torch/distributed/checkpoint/default_planner.py,sha256=Nl7FWqD9_iLdXK6NyJKfWri2TR8niT9-IqbTfgRkkXg,28386 +torch/distributed/checkpoint/filesystem.py,sha256=AJtRXQXb5jSEmwdLGNBBuBywY-9p13qo7MWFefOUVYw,36737 +torch/distributed/checkpoint/format_utils.py,sha256=Cp4LvqDD9mUgsjmfsXIpQuyZbaJrwFEBMphPn1yq0ZI,10733 +torch/distributed/checkpoint/hf_storage.py,sha256=vxbIuOfjLjFxJzjhtHKiSr_D9TuMDzB_sii4X9WKmGc,15383 +torch/distributed/checkpoint/logger.py,sha256=N0xKPxzuyAH2g0TyPI_ij48U7z1HvgTXgfWJxcZOO5s,3635 +torch/distributed/checkpoint/logging_handlers.py,sha256=Ky7OdEi5JzxTwc2u34Mq-tX5S-Y0MynrHAc1xe2nCEk,220 +torch/distributed/checkpoint/metadata.py,sha256=YjGicfDEgIBi1EmC9zC9-b6DQ5-tYnYtChpK0hXeHdk,5573 +torch/distributed/checkpoint/optimizer.py,sha256=8f9p3ML6870yxRsLgRHmGACxgjxvcT_PqSFzwVp3SN4,13342 +torch/distributed/checkpoint/planner.py,sha256=cnJzCR83A2U5Tnrrr8dzHgEX35_N_wmOMqmp8vhon3E,16225 +torch/distributed/checkpoint/planner_helpers.py,sha256=Sqa7m_QdJNXM0y2h3ne0ciH9qgBxVgFc1sooo-COHKU,18775 +torch/distributed/checkpoint/quantized_hf_storage.py,sha256=7B0_-qjPnSU0Qt5TdZVPJFX_KL4sbEyHmNJX-00b38o,20168 +torch/distributed/checkpoint/resharding.py,sha256=1vF4sTAE5WaZaundzKVOOom7CgKohwF52TUQUoWpVxw,2275 +torch/distributed/checkpoint/staging.py,sha256=vd9Ajd0Lf4s4_vyYxzWp1ke7MPJtcSlCtdkmesNuaOo,19331 +torch/distributed/checkpoint/state_dict.py,sha256=IqbiOVbFdhTi8Z6SIsRFCAl59MZWsju_COR6axA7HgA,62502 +torch/distributed/checkpoint/state_dict_loader.py,sha256=rkVzKxEJGP3H99X9FT4H7QGjjx2pE802NEDCDwNQh1Q,14898 +torch/distributed/checkpoint/state_dict_saver.py,sha256=-wmQr3ZH5Gkc00dYTxu8ih5Btg4VY21DBjFoGTazqq8,18475 +torch/distributed/checkpoint/stateful.py,sha256=1cJTGrO_TPFeSguyQhPYCw87ph_QjJRafyAVgx9u7bI,1061 +torch/distributed/checkpoint/storage.py,sha256=xMa-GpSd3X43Yo0yab9hR7Uucato_m1XHWh-JRe-cCs,9792 +torch/distributed/checkpoint/utils.py,sha256=7q28dr36EjaT9VHwxQu4Htm5Jieq0EoAEM1GkEsA_1w,16557 +torch/distributed/collective_utils.py,sha256=r-4K4WK0asp9jg9mWEOrSnPfh5vu_dXeTjyxyZ8uxkg,11835 +torch/distributed/config.py,sha256=XRvD6RN7R38niYpzcSYMSqPagn05QzCuCVLzFRX7TrA,999 +torch/distributed/constants.py,sha256=V6gthYXMG3nf08ObvYceJaNu43yLAVCZ6LXhCGAaKMo,1198 +torch/distributed/debug/__init__.py,sha256=SOf7zuLqqiEGzlMGZoLUxWYfD51-A6DhPXOY5Ubzap0,5243 +torch/distributed/debug/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/debug/__pycache__/_debug_handlers.cpython-311.pyc,, +torch/distributed/debug/__pycache__/_frontend.cpython-311.pyc,, +torch/distributed/debug/__pycache__/_handlers.cpython-311.pyc,, +torch/distributed/debug/__pycache__/_store.cpython-311.pyc,, +torch/distributed/debug/_debug_handlers.py,sha256=rTSedQqONf_XufeD6TubE_miCCODzKsFfYPG4nmEa6M,20159 +torch/distributed/debug/_frontend.py,sha256=9IkNq2lu69IHSSD6mCRtI7uNVja2fO68kGOq_xnP9U8,14592 +torch/distributed/debug/_handlers.py,sha256=k9NoN59EZUgvdL41w2qfXRvSgn_WHk_tp8HNEgSq1qQ,776 +torch/distributed/debug/_store.py,sha256=mjjNaR1pizABV2Q8w9Ap9COIewtpt76VyIiwJMXj7gM,532 +torch/distributed/device_mesh.py,sha256=y1sGO23JXql6PXA1sW_B_K5GBrHfSwQViPxX4BzF1HQ,80466 +torch/distributed/distributed_c10d.py,sha256=wDSIozDWbZsg-5SA9sOxu_umDprq3xiz7KHFypP4hyg,258283 +torch/distributed/elastic/__init__.py,sha256=QMPilFK2QkDj895i4fqrtE5bzyO55GJkV4t4EM_1Qzw,3654 +torch/distributed/elastic/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/__pycache__/control_plane.cpython-311.pyc,, +torch/distributed/elastic/agent/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/elastic/agent/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/agent/server/__init__.py,sha256=p6OmZ3UC2RN276H3JLUZm0ygZ1GBWRcOOq4OQwozTo4,1401 +torch/distributed/elastic/agent/server/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/agent/server/__pycache__/api.cpython-311.pyc,, +torch/distributed/elastic/agent/server/__pycache__/health_check_server.cpython-311.pyc,, +torch/distributed/elastic/agent/server/__pycache__/local_elastic_agent.cpython-311.pyc,, +torch/distributed/elastic/agent/server/api.py,sha256=gYSvoCKJfGIt6f8bdlBOIhBVD6dyzLwvukQnnJE7wf4,39518 +torch/distributed/elastic/agent/server/health_check_server.py,sha256=-HrqCGGbZMklUL6Y0lWt6iwZEVNx7GXknKnB8KLCTtg,1688 +torch/distributed/elastic/agent/server/local_elastic_agent.py,sha256=35TKpJveVtxI_xKsRbjtXFq6rN05TOHJIExUXomq64Q,21840 +torch/distributed/elastic/control_plane.py,sha256=EeYAL-0jl46KYrj16DDO8JVUybqqu7Jgwx7bm28VWmI,1181 +torch/distributed/elastic/events/__init__.py,sha256=-JZ75BaqO7Lz3cHReIQwq8o1xTd4cXRHWVGmsZubTv4,5367 +torch/distributed/elastic/events/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/events/__pycache__/api.cpython-311.pyc,, +torch/distributed/elastic/events/__pycache__/handlers.cpython-311.pyc,, +torch/distributed/elastic/events/api.py,sha256=7o8OinQr8d4hXR6tp_kVNn2EPeVTXfj3dm478Fwtyfw,3310 +torch/distributed/elastic/events/handlers.py,sha256=GVwoHQQFcpFeDW6eobtpDADBX8eIxYCN0tsiFO2pcDQ,556 +torch/distributed/elastic/metrics/__init__.py,sha256=HPZvYUDGNnxv9xqVW0xib9OBj1aBCYejIM4PembVACQ,4899 +torch/distributed/elastic/metrics/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/metrics/__pycache__/api.cpython-311.pyc,, +torch/distributed/elastic/metrics/api.py,sha256=KGdw_W79IGlKYN7GM_qYTZ_0dGjO5CthWLxhgBY1d4E,5751 +torch/distributed/elastic/multiprocessing/__init__.py,sha256=5vx85lNgWFLillVNsaUwj8a6vISu48sVhRKlIiBlEBw,8698 +torch/distributed/elastic/multiprocessing/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/__pycache__/api.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/__pycache__/redirects.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/__pycache__/tail_log.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/api.py,sha256=38kfqq8u9exdysFJXIQn3-FQfX5-0SpGGO396Mclu5M,38701 +torch/distributed/elastic/multiprocessing/errors/__init__.py,sha256=MFL2mlCVFrJoaGOCjSImqJgJspQnTcxt7zohwIRRhQU,14942 +torch/distributed/elastic/multiprocessing/errors/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/errors/__pycache__/error_handler.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/errors/__pycache__/handlers.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/errors/error_handler.py,sha256=HhOFjIQlCdKYwyt1CGXczy98P6Pg7sR0KN1HSarAtSs,6675 +torch/distributed/elastic/multiprocessing/errors/handlers.py,sha256=ea-3iQUTQRDp_YAUhk-3hZ2HQyCPqz_2iiL5_xfD9Q8,464 +torch/distributed/elastic/multiprocessing/redirects.py,sha256=cRBfC8-pbXXRtjp98R5hQ3S2EukqjfhYBNJziQoxAb0,7644 +torch/distributed/elastic/multiprocessing/subprocess_handler/__init__.py,sha256=pz2OzruD8OBIFFxiI1_TLy0Bd8LDBKNLowzW1_jZUSs,523 +torch/distributed/elastic/multiprocessing/subprocess_handler/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/subprocess_handler/__pycache__/handlers.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/subprocess_handler/__pycache__/subprocess_handler.cpython-311.pyc,, +torch/distributed/elastic/multiprocessing/subprocess_handler/handlers.py,sha256=iqetloWZn4hjEoGUlefmHE1RCAP2DF1ht2MZ3Pu96aA,849 +torch/distributed/elastic/multiprocessing/subprocess_handler/subprocess_handler.py,sha256=Y7d7nzs9j4XoiwcMFU2Sl1oFRJvBeteiVCXzy41yLW4,2731 +torch/distributed/elastic/multiprocessing/tail_log.py,sha256=MxlX53zbt3db9NDq5Vt4vnAIiyIEZqieSEl8cOXBLHU,5273 +torch/distributed/elastic/rendezvous/__init__.py,sha256=tcH-Q6v-bdJTWunq4JT-crJ7TxTh9jxEmdzt3vtRBbI,6269 +torch/distributed/elastic/rendezvous/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/_etcd_stub.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/api.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/c10d_rendezvous_backend.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/dynamic_rendezvous.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/etcd_rendezvous.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/etcd_rendezvous_backend.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/etcd_server.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/etcd_store.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/registry.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/static_tcp_rendezvous.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/__pycache__/utils.cpython-311.pyc,, +torch/distributed/elastic/rendezvous/_etcd_stub.py,sha256=r_XAqXKOwrTjmltbMyiKUPVsHtMe7sylfnzdVYvaVek,1998 +torch/distributed/elastic/rendezvous/api.py,sha256=VJD9NR3yEtrLM5AOj5avFe6kXDZVCoRS4nUqB9DH4eg,13084 +torch/distributed/elastic/rendezvous/c10d_rendezvous_backend.py,sha256=3ltlhI2UYHFH0bHBqm__GTLU9OcXEhpAMvaZKRU0VzA,10738 +torch/distributed/elastic/rendezvous/dynamic_rendezvous.py,sha256=JgAWH0yvDYEDsipAB3SHiApjHkH1N3ZcaoShuuvXokc,49432 +torch/distributed/elastic/rendezvous/etcd_rendezvous.py,sha256=BiGR5wMY92EJxxdhRpYCZEZL6CoC9LwhXzKyGxD7vwU,43586 +torch/distributed/elastic/rendezvous/etcd_rendezvous_backend.py,sha256=4AruXChEmCCkVjonf27LHI9pb5o2mO7anC9xy_k1GPU,7348 +torch/distributed/elastic/rendezvous/etcd_server.py,sha256=XegMH-Yixe654S_GMn1quvJzsdz8-pKmgJvkshcLGuM,8410 +torch/distributed/elastic/rendezvous/etcd_store.py,sha256=CjxMeTxKCGTXZ7QwNiH8kLmZLU2LzGv1AbPWsUJpj6Y,7258 +torch/distributed/elastic/rendezvous/registry.py,sha256=uhizddahfTJgSBCIFnfbYR8-f_dl72ORAp2kl52WMHk,2964 +torch/distributed/elastic/rendezvous/static_tcp_rendezvous.py,sha256=Vni_9RgotHveIpmEOwUGcGbScwsx-mh9VwtHFr-XNPs,3652 +torch/distributed/elastic/rendezvous/utils.py,sha256=TP-dMF-YoQHP0FETNj_QEpUzo_EhyMSacZYvBqlmDts,8373 +torch/distributed/elastic/timer/__init__.py,sha256=IQ79GKrQX_UKab_Lbp0uqMuLgb2QzDQgd9EX7NJO5aU,1750 +torch/distributed/elastic/timer/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/timer/__pycache__/api.cpython-311.pyc,, +torch/distributed/elastic/timer/__pycache__/debug_info_logging.cpython-311.pyc,, +torch/distributed/elastic/timer/__pycache__/file_based_local_timer.cpython-311.pyc,, +torch/distributed/elastic/timer/__pycache__/local_timer.cpython-311.pyc,, +torch/distributed/elastic/timer/api.py,sha256=Cx5-79KRjjTub8YiPsIUxqTRoRyecOguChKSADUAz58,9558 +torch/distributed/elastic/timer/debug_info_logging.py,sha256=6mix9Su3bo9BVTM_QXNz75c-WAbDLL8jNbojKbjdEsE,610 +torch/distributed/elastic/timer/file_based_local_timer.py,sha256=OXtzd_N19NA27YJ986lupwk4b9dRB5O5Lklzfr5VepM,16389 +torch/distributed/elastic/timer/local_timer.py,sha256=yaOwmn9FXPTMJ57ubJ0K7TYkWkAeCqY4WF9pMpjrXMg,4282 +torch/distributed/elastic/utils/__init__.py,sha256=ztvtzgzf5pR1ixbbYfY0lGx85VhyhJuHH8er7-XNVpE,318 +torch/distributed/elastic/utils/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/utils/__pycache__/api.cpython-311.pyc,, +torch/distributed/elastic/utils/__pycache__/distributed.cpython-311.pyc,, +torch/distributed/elastic/utils/__pycache__/log_level.cpython-311.pyc,, +torch/distributed/elastic/utils/__pycache__/logging.cpython-311.pyc,, +torch/distributed/elastic/utils/__pycache__/store.cpython-311.pyc,, +torch/distributed/elastic/utils/api.py,sha256=QMd6BGMe0lVXYnOSMLinlB3Hh_EJn2s6R3Cre2K2l0s,1705 +torch/distributed/elastic/utils/data/__init__.py,sha256=tF96JUuxZmWRkvGOzP7tbZT6s8CKNz29F6O_OYTcL3o,372 +torch/distributed/elastic/utils/data/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/elastic/utils/data/__pycache__/cycling_iterator.cpython-311.pyc,, +torch/distributed/elastic/utils/data/__pycache__/elastic_distributed_sampler.cpython-311.pyc,, +torch/distributed/elastic/utils/data/cycling_iterator.py,sha256=T74d2dY-bCXqIYc0kHK0EMCFYAO54c-5otSzMHusDtQ,1643 +torch/distributed/elastic/utils/data/elastic_distributed_sampler.py,sha256=-rTvCe6GL1Ys15687nvWEgUXsF-vaIFhTV_al5o9Rzk,3074 +torch/distributed/elastic/utils/distributed.py,sha256=AOinZvHARbf5sGXBN7Ic9mI7cPmiNJ4Y4UgmZrTOJaE,5892 +torch/distributed/elastic/utils/log_level.py,sha256=1pOw0YanV5aPrMGwqkATk8OTXm6iLF2iuUYKbUf1b6Q,339 +torch/distributed/elastic/utils/logging.py,sha256=UoUWU9dJZBAlYRI833CcdGPSXtTeXt7cQ7pvIJtcFxM,2291 +torch/distributed/elastic/utils/store.py,sha256=3lCBVsuMXBMCmL3HbqAEGCj3IKcR37vVBA2CfcbtsKA,7275 +torch/distributed/flight_recorder/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/flight_recorder/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/flight_recorder/__pycache__/fr_trace.cpython-311.pyc,, +torch/distributed/flight_recorder/components/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/flight_recorder/components/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/flight_recorder/components/__pycache__/builder.cpython-311.pyc,, +torch/distributed/flight_recorder/components/__pycache__/config_manager.cpython-311.pyc,, +torch/distributed/flight_recorder/components/__pycache__/fr_logger.cpython-311.pyc,, +torch/distributed/flight_recorder/components/__pycache__/loader.cpython-311.pyc,, +torch/distributed/flight_recorder/components/__pycache__/types.cpython-311.pyc,, +torch/distributed/flight_recorder/components/__pycache__/utils.cpython-311.pyc,, +torch/distributed/flight_recorder/components/builder.py,sha256=lOfiJg7ygRoZNoBOmC61FLIlOpUWFxmvgckMtRWJlzA,18442 +torch/distributed/flight_recorder/components/config_manager.py,sha256=PmiQ3g-8Av1cSEM5OTktUEoJOPI5sUVUnU_q5E2LXrw,3994 +torch/distributed/flight_recorder/components/fr_logger.py,sha256=DLnrDoW3Pgy0oBUjGpKWo_LCmiWQuJ59MyWlsv9Zknw,1528 +torch/distributed/flight_recorder/components/loader.py,sha256=Awm2L6-Dzz4dyejonfCaKNsuup_PrJ5uuAfSgj-mCwc,2959 +torch/distributed/flight_recorder/components/types.py,sha256=fstyxFlMe4yZaPbrxTDaGhW4TJVJWLg41QmUNOvUaSI,23358 +torch/distributed/flight_recorder/components/utils.py,sha256=FkbiGI2FYcbqGMpE9GfCkxr1j5bZmwNnJ31TLznDh6Q,29500 +torch/distributed/flight_recorder/fr_trace.py,sha256=xxVNSrInrHPW7FUw4OkQ8zm2L_9AOosx4jfOtCGGChw,3066 +torch/distributed/fsdp/__init__.py,sha256=O_Vc-R5chY3IH5f9-TEA6JsGeqQovCJhFIwv-abthkM,1948 +torch/distributed/fsdp/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_common_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_debug_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_dynamo_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_exec_order_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_flat_param.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_fsdp_extensions.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_init_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_limiter_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_optim_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_runtime_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_shard_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_state_dict_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_trace_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_traversal_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_unshard_param_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/_wrap_utils.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/api.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/fully_sharded_data_parallel.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/sharded_grad_scaler.cpython-311.pyc,, +torch/distributed/fsdp/__pycache__/wrap.cpython-311.pyc,, +torch/distributed/fsdp/_common_utils.py,sha256=spEi9t3f5IgXikFTJxqay2DN_MAImXALXXRm_N60mPA,29281 +torch/distributed/fsdp/_debug_utils.py,sha256=1SmIdIGZ42D8y-OWBbbTXCTvh7hf82h0jn2SbRquSc4,5818 +torch/distributed/fsdp/_dynamo_utils.py,sha256=_YtXy6KD9FUp--G1z7xpA4rrcSWFnF0bMkRtcDptpeg,2631 +torch/distributed/fsdp/_exec_order_utils.py,sha256=BrCXrx1FXh3s6cxKE2QBi26Q-jTGdYDobiXLIOwvq1U,16089 +torch/distributed/fsdp/_flat_param.py,sha256=rT6fKOQ8opBEZ90HvfKAt_MjppDl7UK8Ar4aMLLNA6g,125974 +torch/distributed/fsdp/_fsdp_extensions.py,sha256=iGFirKzsyf3OEeKV8v2fB9dkkmBP8GlSS-vg6N92AOI,4987 +torch/distributed/fsdp/_fully_shard/__init__.py,sha256=uXSibxJe0zL41WR-oTOtpzPSk6Wxq6UTSc6ge0kYO38,489 +torch/distributed/fsdp/_fully_shard/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/fsdp/_fully_shard/__pycache__/_fsdp_api.cpython-311.pyc,, +torch/distributed/fsdp/_fully_shard/__pycache__/_fsdp_collectives.cpython-311.pyc,, +torch/distributed/fsdp/_fully_shard/__pycache__/_fsdp_common.cpython-311.pyc,, +torch/distributed/fsdp/_fully_shard/__pycache__/_fsdp_init.cpython-311.pyc,, +torch/distributed/fsdp/_fully_shard/__pycache__/_fsdp_param.cpython-311.pyc,, +torch/distributed/fsdp/_fully_shard/__pycache__/_fsdp_param_group.cpython-311.pyc,, +torch/distributed/fsdp/_fully_shard/__pycache__/_fsdp_state.cpython-311.pyc,, +torch/distributed/fsdp/_fully_shard/__pycache__/_fully_shard.cpython-311.pyc,, +torch/distributed/fsdp/_fully_shard/_fsdp_api.py,sha256=FA9J_OEpKK7gfiB6ttEveEQfZrIRcx9H0k-lOkM2ss4,6918 +torch/distributed/fsdp/_fully_shard/_fsdp_collectives.py,sha256=R9g7vHpdOtAQsVo_ntsV8G2kMlO9_HIcJoPOZJ_s1rk,32670 +torch/distributed/fsdp/_fully_shard/_fsdp_common.py,sha256=MWPf1XC6fDceISiz70JMHupdKrH33YmiXcQuqrDZDBQ,6886 +torch/distributed/fsdp/_fully_shard/_fsdp_init.py,sha256=oRlEG-VjJxPzCDaI15xyUDUC0eQEUrPEehu7Pzg-k4g,20751 +torch/distributed/fsdp/_fully_shard/_fsdp_param.py,sha256=UcYEsZfrc2uAzibyFloi1s1UBzI4u_sb136pHT09qsM,45313 +torch/distributed/fsdp/_fully_shard/_fsdp_param_group.py,sha256=5kYsoLFVcGsB1xp_nMq8h9lNPFLKZht579gKLdt6tnc,41088 +torch/distributed/fsdp/_fully_shard/_fsdp_state.py,sha256=xu-mq7Ajeq46DMfH3EXyaQKcrhV2ocxkaX2xmoLjMys,21730 +torch/distributed/fsdp/_fully_shard/_fully_shard.py,sha256=o2hQ8g8Fe1R8afVtX3adRz9-UquyLgVJ8vpYXAwkCkk,36703 +torch/distributed/fsdp/_init_utils.py,sha256=ay9gaBk2AMx49wSVO1jAhqLyPSononalYNep_Ou2-k0,48488 +torch/distributed/fsdp/_limiter_utils.py,sha256=Uer7diNEpPlQG9ijdEwkh8K0d9iJl2texk_ArIz7euk,1052 +torch/distributed/fsdp/_optim_utils.py,sha256=YKrL2dQhs4H_XBhjtAqUSmG8_p_gl6VOaSPfGYhJs_M,89828 +torch/distributed/fsdp/_runtime_utils.py,sha256=6axOB-QKPjEsfBdsUzggHSAvjw5f75znFnWPJEtoa-o,66905 +torch/distributed/fsdp/_shard_utils.py,sha256=MMZ5-X4Yg8KWMoqlPebVMyXo0lzoiUUayi2iKevoYak,4828 +torch/distributed/fsdp/_state_dict_utils.py,sha256=xOnTccogJteGH79Rjz50DhiQoVI-teeawVe13OchtCw,34737 +torch/distributed/fsdp/_trace_utils.py,sha256=JLa_T89INq4QDKVN1w9qQcaXtTW2-8HXqkEcxtpEUBI,10798 +torch/distributed/fsdp/_traversal_utils.py,sha256=haGRJUg9EyPuRJjUyUWfBDaRmEdJ7VOaOPi722DQ2Xg,4610 +torch/distributed/fsdp/_unshard_param_utils.py,sha256=hG8RJ0F-04VVqzmcOAf7c4upY6iapif3lOsimcjXSpQ,11663 +torch/distributed/fsdp/_wrap_utils.py,sha256=uhD6UktLckRjpzXOWlpI-mLsiB0qSuyHh6cU_75EShQ,10946 +torch/distributed/fsdp/api.py,sha256=QzyR5LX5OC6GdiPvZFq_SJbF8-UD4i07X-l_0-WSJBQ,18938 +torch/distributed/fsdp/fully_sharded_data_parallel.py,sha256=FTPIj3VJztHNVGFfe_Ke1d0De002iYt-bP5ZT4IKcs8,100777 +torch/distributed/fsdp/sharded_grad_scaler.py,sha256=5Whb-weywbMd7Fbpj8oAv5JY_cubShTiZribl12TNA0,17057 +torch/distributed/fsdp/wrap.py,sha256=bFgUTmDiPzwF3sWm_5QpqowOVJ0uRK-PXvtBEuhtSqo,23110 +torch/distributed/launch.py,sha256=e7R-D1Wy1X98YSTVQESxOkDBA3UwCkrzpEzraXovYYc,7595 +torch/distributed/launcher/__init__.py,sha256=V0J5Y0SVG1ybhGeHnjarMc4dcOgyhRsHpnP7lBMyB48,349 +torch/distributed/launcher/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/launcher/__pycache__/api.cpython-311.pyc,, +torch/distributed/launcher/api.py,sha256=HGMP2TvEZSjc1OSv88e-5eWihbceqkuXF7dhjsSpMtU,16162 +torch/distributed/logging_handlers.py,sha256=nXuxHnOEhAHqqDPhTk_-aLYIVHoZpQEYW_X4N1Qg0_Q,359 +torch/distributed/nn/__init__.py,sha256=Ta4l0Jr2gOK7eQR-e32stWaDKPOlV_Kii0SExydXk_Q,145 +torch/distributed/nn/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/nn/__pycache__/functional.cpython-311.pyc,, +torch/distributed/nn/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/nn/api/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/nn/api/__pycache__/remote_module.cpython-311.pyc,, +torch/distributed/nn/api/remote_module.py,sha256=HgXhfAA32lLaByw1PLfy-jF5LPOE5HnODDnvgD6VFFg,31378 +torch/distributed/nn/functional.py,sha256=OikdkDC7Qg-fBwJn8KuBz-WN327Gj0ThtnuQrW_uytI,18370 +torch/distributed/nn/jit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/nn/jit/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/nn/jit/__pycache__/instantiator.cpython-311.pyc,, +torch/distributed/nn/jit/instantiator.py,sha256=IhEEPfRwITTS2oQipmVr_w1So9sAgOSsPgi0ff-XrCg,6198 +torch/distributed/nn/jit/templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/distributed/nn/jit/templates/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/nn/jit/templates/__pycache__/remote_module_template.cpython-311.pyc,, +torch/distributed/nn/jit/templates/remote_module_template.py,sha256=D_GFa70DG1KY1GwGwFAqvCC9gE9CwZSe1BJ-jHc2YMw,3463 +torch/distributed/optim/__init__.py,sha256=HxJ6GIWc3b6uAZWNh45oePXlQLAdX9Kmu3PByAOvlh4,1440 +torch/distributed/optim/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/optim/__pycache__/_deprecation_warning.cpython-311.pyc,, +torch/distributed/optim/__pycache__/apply_optimizer_in_backward.cpython-311.pyc,, +torch/distributed/optim/__pycache__/functional_adadelta.cpython-311.pyc,, +torch/distributed/optim/__pycache__/functional_adagrad.cpython-311.pyc,, +torch/distributed/optim/__pycache__/functional_adam.cpython-311.pyc,, +torch/distributed/optim/__pycache__/functional_adamax.cpython-311.pyc,, +torch/distributed/optim/__pycache__/functional_adamw.cpython-311.pyc,, +torch/distributed/optim/__pycache__/functional_rmsprop.cpython-311.pyc,, +torch/distributed/optim/__pycache__/functional_rprop.cpython-311.pyc,, +torch/distributed/optim/__pycache__/functional_sgd.cpython-311.pyc,, +torch/distributed/optim/__pycache__/named_optimizer.cpython-311.pyc,, +torch/distributed/optim/__pycache__/optimizer.cpython-311.pyc,, +torch/distributed/optim/__pycache__/post_localSGD_optimizer.cpython-311.pyc,, +torch/distributed/optim/__pycache__/utils.cpython-311.pyc,, +torch/distributed/optim/__pycache__/zero_redundancy_optimizer.cpython-311.pyc,, +torch/distributed/optim/_deprecation_warning.py,sha256=4lon9fWNrTNS4ARfi5Dag5tvohwNyPMflHBl9VFd_0E,547 +torch/distributed/optim/apply_optimizer_in_backward.py,sha256=uTNZylxBaBVQK1teYyPJ7R45-t8aqG6joeGT3qyi0AQ,5214 +torch/distributed/optim/functional_adadelta.py,sha256=rnI34B6w5-ISY_21uC2CgfoZVWVbovWOUsOAwFpQBqc,3923 +torch/distributed/optim/functional_adagrad.py,sha256=5qM9YJdGHQmYK9dhGYFCFZ50r-6PoqXI2pe9b9xbgTM,4273 +torch/distributed/optim/functional_adam.py,sha256=o0feTLaXh_Lx_DneB5Vy8vX4R0iO3lw1-Dk-dBeRWro,7391 +torch/distributed/optim/functional_adamax.py,sha256=OAc-Snni8vYQq5sJWcC8aV-dnrOA26i8zOS6nCsGGCg,4620 +torch/distributed/optim/functional_adamw.py,sha256=BJMztbCxW3lGdaBP-9AifkX2cHNnyskmMGbazMfKuFE,7511 +torch/distributed/optim/functional_rmsprop.py,sha256=2FLOhJNuA4wh2vnlb1mPGrll4dJ9zSjdAtImk-3dqno,4653 +torch/distributed/optim/functional_rprop.py,sha256=avQYdiVv0UaTcRGsPTFoeGrvSM1OMpsColMd06zJPdE,3807 +torch/distributed/optim/functional_sgd.py,sha256=dpVzD0a-GSIMG_iqwMvqDUyKOnsvi4sT9vjii0KHlxs,5898 +torch/distributed/optim/named_optimizer.py,sha256=n1T-Chv8j87UIP7-iVVVZJ1DTbsja7DmjA4EhjRCtrM,14184 +torch/distributed/optim/optimizer.py,sha256=PGqrPW0HJZmQPm6dT1b6V1_xmyVIk1yFB7DVfRb6w4Y,9756 +torch/distributed/optim/post_localSGD_optimizer.py,sha256=MT7Kno6EvcmnjAiQGJdsgOnKXpHXqm9Sm9QnHzV-EP4,4498 +torch/distributed/optim/utils.py,sha256=JhzIndnTdt0CMJiEblAbg-tOdVDhFv6tMk00_WjBPMI,2238 +torch/distributed/optim/zero_redundancy_optimizer.py,sha256=85DkD6UeM2VJtzN72HZ4HxBrqv3YetdjYnD9l43WUA8,73748 +torch/distributed/optim/zero_redundancy_optimizer.pyi,sha256=I-G_ysSHrYjja75Vn5DPbuSCtyE4g_Pg_Mp28EBA0Js,2852 +torch/distributed/pipelining/_IR.py,sha256=yTBOTevhHqv4BKtPVPE079fcFZmOgKhQR0Mggf7EGZ8,50442 +torch/distributed/pipelining/__init__.py,sha256=ywlBKSDSFkgi5wBXXS02nJ_srC1l2ACoCkI1PtDjcDY,689 +torch/distributed/pipelining/__pycache__/_IR.cpython-311.pyc,, +torch/distributed/pipelining/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/pipelining/__pycache__/_backward.cpython-311.pyc,, +torch/distributed/pipelining/__pycache__/_debug.cpython-311.pyc,, +torch/distributed/pipelining/__pycache__/_schedule_visualizer.cpython-311.pyc,, +torch/distributed/pipelining/__pycache__/_unflatten.cpython-311.pyc,, +torch/distributed/pipelining/__pycache__/_utils.cpython-311.pyc,, +torch/distributed/pipelining/__pycache__/microbatch.cpython-311.pyc,, +torch/distributed/pipelining/__pycache__/schedules.cpython-311.pyc,, +torch/distributed/pipelining/__pycache__/stage.cpython-311.pyc,, +torch/distributed/pipelining/_backward.py,sha256=EwG65PCIyUijL-dFK28uJ4NErsSsQuU5-pl5HCecvok,18250 +torch/distributed/pipelining/_debug.py,sha256=dDUjgp_-XnUEsM4I-zj5-yQGxcapHbBb3JcRSdRuCqc,608 +torch/distributed/pipelining/_schedule_visualizer.py,sha256=0RxVafpG3HMO-4zGSSeieDE27BfYW4P9dn7uF3Kynp8,17193 +torch/distributed/pipelining/_unflatten.py,sha256=GUpLGbhqxbAMjcY9aDxandxUd3Rva1aXs4BRApx9Gj8,952 +torch/distributed/pipelining/_utils.py,sha256=FabRwgUdvT6e_Xci7YXxzZBBF699xoyFZwrTJfirr6U,32412 +torch/distributed/pipelining/microbatch.py,sha256=Q3PykJO2THHN9vhD0Bo1tANe5wmD7eHelA8hioJ4tqg,21760 +torch/distributed/pipelining/schedules.py,sha256=mLDI9kFOome137q19nznODMZXWnTTXsxnhe8vutxEWY,147424 +torch/distributed/pipelining/stage.py,sha256=pXOosshTfhGBEtxzdAM-tqqC9DNt_V2p0kyDEKNGeh8,93744 +torch/distributed/remote_device.py,sha256=FaHtxJp4aLyH63gBkigD-p3lJ67NZCgWkYEUNPKYvok,4545 +torch/distributed/rendezvous.py,sha256=jI1J6mN7mns8k7o-DNoGXxCsx4zYJukdEozRlV2zXFg,10273 +torch/distributed/rpc/__init__.py,sha256=2aiCTn524Vsvkdr0QsKBTCpT_tyQCpWCGxqRcBScPcg,9993 +torch/distributed/rpc/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/rpc/__pycache__/_utils.cpython-311.pyc,, +torch/distributed/rpc/__pycache__/api.cpython-311.pyc,, +torch/distributed/rpc/__pycache__/backend_registry.cpython-311.pyc,, +torch/distributed/rpc/__pycache__/constants.cpython-311.pyc,, +torch/distributed/rpc/__pycache__/functions.cpython-311.pyc,, +torch/distributed/rpc/__pycache__/internal.cpython-311.pyc,, +torch/distributed/rpc/__pycache__/options.cpython-311.pyc,, +torch/distributed/rpc/__pycache__/rref_proxy.cpython-311.pyc,, +torch/distributed/rpc/__pycache__/server_process_global_profiler.cpython-311.pyc,, +torch/distributed/rpc/_testing/__init__.py,sha256=xSU_1OK5D0vaFTGSxPlzX6Ad4CEvsZv7DIJYY1d2Unk,479 +torch/distributed/rpc/_testing/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/rpc/_testing/__pycache__/faulty_agent_backend_registry.cpython-311.pyc,, +torch/distributed/rpc/_testing/faulty_agent_backend_registry.py,sha256=KmeZ2KHDwSWJB3mdfA06hdaabt25vKkkUU1h8SFVk1o,1639 +torch/distributed/rpc/_utils.py,sha256=tcq0ADOGMmd7OkKYFY7apu6c44VbgFLxCLOCh-aiyvU,1649 +torch/distributed/rpc/api.py,sha256=c6v_xV1cHtHJ8TLaA-oD9tqal_Vm38TwgiHfTcloltk,37293 +torch/distributed/rpc/backend_registry.py,sha256=QUc-0QIFAETaBhCtY0Bl_COcyrTnXl8O9p7jbrHQNek,16256 +torch/distributed/rpc/constants.py,sha256=AFzuxIi9gIVxb9nyZNZ6h70Is-DUMsCHZ-WEO3MOMP8,804 +torch/distributed/rpc/functions.py,sha256=9sE_wYiqlfUmb6AFBFPmOq6PApwYKxosNOnteoy_0Oc,7272 +torch/distributed/rpc/internal.py,sha256=gEsK3c38SKriXoE0Uy3BMtar6_smcS_xc3_st5PYeLc,11149 +torch/distributed/rpc/options.py,sha256=uSupCGgnrJNCf3E0eIaH4hySXM7vT-1gslMx6k3ceXw,7218 +torch/distributed/rpc/rref_proxy.py,sha256=QRFHO3EvUay19HkZgUkf1uNLDQgp1Gmop8z4Ics5Aq4,2705 +torch/distributed/rpc/server_process_global_profiler.py,sha256=NpQiYNSv10xOWQbtvbYrg3Q2nXdE6BuHTz82BkUyCak,8477 +torch/distributed/run.py,sha256=mQpVMEGD10f-heG6hFdy8ANbWTkSfigfDbXS7fj8aS0,36223 +torch/distributed/tensor/__init__.py,sha256=PpV4NhDxXs8qAYkPWAjSTYf0qMtb7vg3o_JaoTxP2tc,2439 +torch/distributed/tensor/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_api.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_collective_utils.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_decompositions.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_dispatch.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_dtensor_spec.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_nonlinear_redux.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_op_schema.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_random.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_redistribute.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_sharding_prop.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_shards_wrapper.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_tp_conv.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/_utils.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/device_mesh.cpython-311.pyc,, +torch/distributed/tensor/__pycache__/placement_types.cpython-311.pyc,, +torch/distributed/tensor/_api.py,sha256=Xi5OxhFlIKgB6UbDoLZZxYnGFG3PhahgM_rjLDEJg2E,66489 +torch/distributed/tensor/_collective_utils.py,sha256=ejKzKJqabPzUpRvCZTX9LEIeO-L5BQlNXuw2p0pKVPM,21783 +torch/distributed/tensor/_decompositions.py,sha256=ZC5J-15bnsFLYs4o0wjV3CDAyc7rX2y77slpdu1m_60,13116 +torch/distributed/tensor/_dispatch.py,sha256=VyEHpnzKeIi_eXSRcW_88iG1uKAWE5uaEZFUvOGm5f4,35212 +torch/distributed/tensor/_dtensor_spec.py,sha256=UHXDnendOhTVnL4AutIq2sF4nNYvDdTvKiZmg0DVBn4,33717 +torch/distributed/tensor/_nonlinear_redux.py,sha256=KJtZlRx_YPPnborI1-bR7oJdsZfFk9Nqsg1LqrUb-eQ,10519 +torch/distributed/tensor/_op_schema.py,sha256=8ims9jnbaGY8xYjzdul1A6-TC7Mo_SO6c_cckijVjR4,27782 +torch/distributed/tensor/_ops/__init__.py,sha256=bQDReWqPLgNAfjBJC18_Q8Yg9wTj_ZlIut6CnQ_KBiY,380 +torch/distributed/tensor/_ops/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_common_rules.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_conv_ops.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_einsum_strategy.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_embedding_ops.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_mask_buffer.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_math_ops.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_matrix_ops.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_pointwise_ops.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_random_ops.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_tensor_ops.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/_view_ops.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/single_dim_strategy.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/strategy_validation.cpython-311.pyc,, +torch/distributed/tensor/_ops/__pycache__/utils.cpython-311.pyc,, +torch/distributed/tensor/_ops/_common_rules.py,sha256=2ChHkqXUJXulPsPHZWv1gGgvWXdSwXxTGRLLPhHN2CY,11998 +torch/distributed/tensor/_ops/_conv_ops.py,sha256=I9JZp9xOoErMSak1BjUzuiMmVeWuMzJvLNF7siToxhI,6642 +torch/distributed/tensor/_ops/_einsum_strategy.py,sha256=gGM6LUYNEqXHN_slfFCXu1zt6tXkTQHJpQcIdVYLANc,7691 +torch/distributed/tensor/_ops/_embedding_ops.py,sha256=4tXUEBYgKNdEPRfcbnH_HjriDJHWupiTXQSkXchBbOc,4370 +torch/distributed/tensor/_ops/_mask_buffer.py,sha256=c5sml3lk7LOlKaNNcH9VIU0uVC2qAT9GPcyBY-OhYEU,1562 +torch/distributed/tensor/_ops/_math_ops.py,sha256=BpvI570pwfgF80-VXFiBA96-RIAbeocbeSoL9lRrMWc,64387 +torch/distributed/tensor/_ops/_matrix_ops.py,sha256=8vy6KNafkJ3kGKiwnaa5218z7Dc6guJMDxVYtdmHpIc,50084 +torch/distributed/tensor/_ops/_pointwise_ops.py,sha256=slREdDyMcHzzCqo-J6my2FCln2MJ7ZPspKRmJEs6c9s,19983 +torch/distributed/tensor/_ops/_random_ops.py,sha256=fAR1vpt9qxPn1bnB5EGEYpXXi8qZBTpGhxPypoh7H6A,1286 +torch/distributed/tensor/_ops/_tensor_ops.py,sha256=jAbxxw3fifykz49dznuBxbWhpR9wnF0DAQwVQiIGFT4,62545 +torch/distributed/tensor/_ops/_view_ops.py,sha256=2pKw7UXqMjSO0Te5K-2-6ea0p86OxlXJryKZeQxTfGg,56485 +torch/distributed/tensor/_ops/single_dim_strategy.py,sha256=sY_rZlyXZDvwD34aEycMsrlbGjNqCpMIQVtox8ZTU58,45004 +torch/distributed/tensor/_ops/strategy_validation.py,sha256=fMJvuqLTCRUCab3yLnRv8XKq0E5PTugjIVFDcQjlnMA,84711 +torch/distributed/tensor/_ops/utils.py,sha256=I2Hd8FfOmCwqWeCecu0jmOeYDBcbeJQe8DiZYZzogYk,29063 +torch/distributed/tensor/_random.py,sha256=GSsjvutUjYZVMcCDjcS1JgWEnEkQLt13Zfew5SdQVxs,20376 +torch/distributed/tensor/_redistribute.py,sha256=PWmARKFvN78GLXOkzRahHquIjEh8zxqr8TAH8Uq_jiw,84177 +torch/distributed/tensor/_sharding_prop.py,sha256=f3Hn31EPZmGRX72jhN5zq2ycWyeChjNgKDA_5re4qVw,47938 +torch/distributed/tensor/_shards_wrapper.py,sha256=EjVQNKFDmDbD1N-aFENgLjqmoaTfOLmLbyvtvEW7ccE,13592 +torch/distributed/tensor/_tp_conv.py,sha256=VJz75AZypeKHdSOPM8c65VF4XVfF_PaJmcg095vBbio,11370 +torch/distributed/tensor/_utils.py,sha256=BvU5G1pFBMdSbL4ittorB9f8dIW06JBaF07kCKfBbAI,20937 +torch/distributed/tensor/debug/__init__.py,sha256=1fkt2_pt90BuM0T1H6DWeWO12P_ivEOik7x4SIvrJFY,2446 +torch/distributed/tensor/debug/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/tensor/debug/__pycache__/_comm_mode.cpython-311.pyc,, +torch/distributed/tensor/debug/__pycache__/_op_coverage.cpython-311.pyc,, +torch/distributed/tensor/debug/__pycache__/_visualize_sharding.cpython-311.pyc,, +torch/distributed/tensor/debug/_comm_mode.py,sha256=zQytrcRSpo_OWwlowu0OQGxz3BzDRilCkm-WEdqMggc,29316 +torch/distributed/tensor/debug/_op_coverage.py,sha256=JVC581pNJZyM6fMwyhzXNLMxC4mXJWQ6y2Qpjxd2JRg,3168 +torch/distributed/tensor/debug/_visualize_sharding.py,sha256=P2rXbtLlKqci4CcKe_yIlWZCN2sGwR0vg624dLylSEk,7562 +torch/distributed/tensor/device_mesh.py,sha256=MONPC_i1XphcY2E0AAte1MW668BTlUGugMSQzfJqzik,190 +torch/distributed/tensor/experimental/__init__.py,sha256=9I70AFPtyjGyKxAuaBsRfDxvHyCeM5a7o7q3X1yUe5g,1424 +torch/distributed/tensor/experimental/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/tensor/experimental/__pycache__/_attention.cpython-311.pyc,, +torch/distributed/tensor/experimental/__pycache__/_func_map.cpython-311.pyc,, +torch/distributed/tensor/experimental/__pycache__/_register_sharding.cpython-311.pyc,, +torch/distributed/tensor/experimental/__pycache__/_tp_transform.cpython-311.pyc,, +torch/distributed/tensor/experimental/_attention.py,sha256=Hz6IsHinuWEYMWAMGzQMuVEYAHiJF3VqLQjIbysMBtw,1229 +torch/distributed/tensor/experimental/_context_parallel/__init__.py,sha256=6yKCIGx9a4EB26ST6wrFUihdq8SPHvOj8wPgehrHeb0,1130 +torch/distributed/tensor/experimental/_context_parallel/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/tensor/experimental/_context_parallel/__pycache__/_attention.cpython-311.pyc,, +torch/distributed/tensor/experimental/_context_parallel/__pycache__/_cp_custom_ops.cpython-311.pyc,, +torch/distributed/tensor/experimental/_context_parallel/__pycache__/_load_balancer.cpython-311.pyc,, +torch/distributed/tensor/experimental/_context_parallel/__pycache__/_sharding_rules.cpython-311.pyc,, +torch/distributed/tensor/experimental/_context_parallel/_attention.py,sha256=GE9GL8sw3e1SnVVH66xtIbCER2NqUxo5JMgIodYtbBY,62230 +torch/distributed/tensor/experimental/_context_parallel/_cp_custom_ops.py,sha256=xNWTO29HqyGhpcnheOMuujOh4q7zlCeAbVlFPNjQB_8,2969 +torch/distributed/tensor/experimental/_context_parallel/_load_balancer.py,sha256=-N0YgLLtKtZSJh8pwnHdLxG61rIqzNZSzZ0El1zLxzU,22511 +torch/distributed/tensor/experimental/_context_parallel/_sharding_rules.py,sha256=MmvXLhj1PmPqmfmBQfQ6KEDJmvrl18IwgryzI4a-W_A,13740 +torch/distributed/tensor/experimental/_func_map.py,sha256=fDjc4ZTho7dUyfZh6KLls_TsnJRb-8EdxQiCnAQO6ZQ,12422 +torch/distributed/tensor/experimental/_register_sharding.py,sha256=9R8PjVMEesw-JpIuGcFruwR3gh5T9BxQHkOKjVvWo3U,5636 +torch/distributed/tensor/experimental/_tp_transform.py,sha256=rtodmb-7vFY1XuNr4WRjVYiRLWsoWyTWoYbvd-qxEeE,21142 +torch/distributed/tensor/parallel/__init__.py,sha256=_rDiCQ0mhllWr5HSDkOV47G4Z-WAZ3L1PCyzLhjgXeU,643 +torch/distributed/tensor/parallel/__pycache__/__init__.cpython-311.pyc,, +torch/distributed/tensor/parallel/__pycache__/_data_parallel_utils.cpython-311.pyc,, +torch/distributed/tensor/parallel/__pycache__/api.cpython-311.pyc,, +torch/distributed/tensor/parallel/__pycache__/ddp.cpython-311.pyc,, +torch/distributed/tensor/parallel/__pycache__/fsdp.cpython-311.pyc,, +torch/distributed/tensor/parallel/__pycache__/input_reshard.cpython-311.pyc,, +torch/distributed/tensor/parallel/__pycache__/loss.cpython-311.pyc,, +torch/distributed/tensor/parallel/__pycache__/style.cpython-311.pyc,, +torch/distributed/tensor/parallel/_data_parallel_utils.py,sha256=toFgFzlwhtjqPA6hOJjRFW4ZYjV77KZpmu2TssdBaPw,1498 +torch/distributed/tensor/parallel/api.py,sha256=h8TIf56coMFuKfOfgDAKxTfMOzmBZ3izgfEMwvCsoFE,6476 +torch/distributed/tensor/parallel/ddp.py,sha256=ES69uXY8nKFKp7je9RQ7n8tP40GwcUekLGY00qAmL3c,3756 +torch/distributed/tensor/parallel/fsdp.py,sha256=STxhc5IdKZAYZWVghX_YLwVxfmToNfh-3gSOi372wa8,14044 +torch/distributed/tensor/parallel/input_reshard.py,sha256=ZJTv6n0S0qNspaHvEGFbF_9HKuEIBL-xY0XwLNcL2dw,3813 +torch/distributed/tensor/parallel/loss.py,sha256=e2qsNa3QJGJ2kPyKk0AQqP6kzVkRfVVcgMgcpNB7MUw,18741 +torch/distributed/tensor/parallel/style.py,sha256=hhcKLwaetx7HaTlcZGe8X0bAj0l4jAP_gQ6xsxixIzQ,37639 +torch/distributed/tensor/placement_types.py,sha256=WAdyzNSCmqvc203nAH0v0WTtoB1ZtzezuPxhmt_8grY,61780 +torch/distributed/utils.py,sha256=bcr0u2E70WLteT4C-z56bTdPsv1cPr-0I68wz4Hp_Rc,13729 +torch/distributions/__init__.py,sha256=FdWgzsV7zpnl3LMsEv73LOOxNCSkciXh9JGuxpiCyDU,6111 +torch/distributions/__pycache__/__init__.cpython-311.pyc,, +torch/distributions/__pycache__/bernoulli.cpython-311.pyc,, +torch/distributions/__pycache__/beta.cpython-311.pyc,, +torch/distributions/__pycache__/binomial.cpython-311.pyc,, +torch/distributions/__pycache__/categorical.cpython-311.pyc,, +torch/distributions/__pycache__/cauchy.cpython-311.pyc,, +torch/distributions/__pycache__/chi2.cpython-311.pyc,, +torch/distributions/__pycache__/constraint_registry.cpython-311.pyc,, +torch/distributions/__pycache__/constraints.cpython-311.pyc,, +torch/distributions/__pycache__/continuous_bernoulli.cpython-311.pyc,, +torch/distributions/__pycache__/dirichlet.cpython-311.pyc,, +torch/distributions/__pycache__/distribution.cpython-311.pyc,, +torch/distributions/__pycache__/exp_family.cpython-311.pyc,, +torch/distributions/__pycache__/exponential.cpython-311.pyc,, +torch/distributions/__pycache__/fishersnedecor.cpython-311.pyc,, +torch/distributions/__pycache__/gamma.cpython-311.pyc,, +torch/distributions/__pycache__/generalized_pareto.cpython-311.pyc,, +torch/distributions/__pycache__/geometric.cpython-311.pyc,, +torch/distributions/__pycache__/gumbel.cpython-311.pyc,, +torch/distributions/__pycache__/half_cauchy.cpython-311.pyc,, +torch/distributions/__pycache__/half_normal.cpython-311.pyc,, +torch/distributions/__pycache__/independent.cpython-311.pyc,, +torch/distributions/__pycache__/inverse_gamma.cpython-311.pyc,, +torch/distributions/__pycache__/kl.cpython-311.pyc,, +torch/distributions/__pycache__/kumaraswamy.cpython-311.pyc,, +torch/distributions/__pycache__/laplace.cpython-311.pyc,, +torch/distributions/__pycache__/lkj_cholesky.cpython-311.pyc,, +torch/distributions/__pycache__/log_normal.cpython-311.pyc,, +torch/distributions/__pycache__/logistic_normal.cpython-311.pyc,, +torch/distributions/__pycache__/lowrank_multivariate_normal.cpython-311.pyc,, +torch/distributions/__pycache__/mixture_same_family.cpython-311.pyc,, +torch/distributions/__pycache__/multinomial.cpython-311.pyc,, +torch/distributions/__pycache__/multivariate_normal.cpython-311.pyc,, +torch/distributions/__pycache__/negative_binomial.cpython-311.pyc,, +torch/distributions/__pycache__/normal.cpython-311.pyc,, +torch/distributions/__pycache__/one_hot_categorical.cpython-311.pyc,, +torch/distributions/__pycache__/pareto.cpython-311.pyc,, +torch/distributions/__pycache__/poisson.cpython-311.pyc,, +torch/distributions/__pycache__/relaxed_bernoulli.cpython-311.pyc,, +torch/distributions/__pycache__/relaxed_categorical.cpython-311.pyc,, +torch/distributions/__pycache__/studentT.cpython-311.pyc,, +torch/distributions/__pycache__/transformed_distribution.cpython-311.pyc,, +torch/distributions/__pycache__/transforms.cpython-311.pyc,, +torch/distributions/__pycache__/uniform.cpython-311.pyc,, +torch/distributions/__pycache__/utils.cpython-311.pyc,, +torch/distributions/__pycache__/von_mises.cpython-311.pyc,, +torch/distributions/__pycache__/weibull.cpython-311.pyc,, +torch/distributions/__pycache__/wishart.cpython-311.pyc,, +torch/distributions/bernoulli.py,sha256=yU1Sf9xML4FNaN1YnNYhjKPFXQyM-2bT0w6ni1C-bEc,4762 +torch/distributions/beta.py,sha256=X7bUYY6OSy-SxJNWCXRA3A1mWs0Hn3gfGuxkhC-PhXU,4000 +torch/distributions/binomial.py,sha256=cXO378OmHHclqL_YKnBfpjH9dzCO279RFk7iiXJJ6SQ,6481 +torch/distributions/categorical.py,sha256=eU-zgx7Yp3cyWp00lNMzfr9V5Y_cFKQXLXKM3BCwvbQ,6277 +torch/distributions/cauchy.py,sha256=ZubvZlaBqfin-Yg2Qlq_R4VA4jcyEg9gLLzMm3Lsox8,3159 +torch/distributions/chi2.py,sha256=BP_QvY2zY7xvDx-M_hMmmB_Ral0eUyYd-Ak8B9W7AqI,1109 +torch/distributions/constraint_registry.py,sha256=Z49QrxE5eJqi7_84YM73khNXMNVGhM1jqb-ovWBKT-A,10306 +torch/distributions/constraints.py,sha256=RvzbDaVIrAehxkxk61Z4BmOfPi77CxBGpW66pTxccNs,22767 +torch/distributions/continuous_bernoulli.py,sha256=m-gMA7LVbPhqFaZNpt91opQh-VwpBX2BbdE7a5g1Ffg,9298 +torch/distributions/dirichlet.py,sha256=swXYGrzZlnmYktGw0W6Q3Tx3AATbvU5La0HOzXFHO7k,4576 +torch/distributions/distribution.py,sha256=qnxc2TqPpoVQOp5dxClje01wK3TNVi7M0-3iBsEtge8,12695 +torch/distributions/exp_family.py,sha256=qr4IKnvdwg0_6NWulD8TrtaCIRBD5lr6x4KJOuhEmXE,2454 +torch/distributions/exponential.py,sha256=Gggg_b0qOTBaYL1sFMrJ8PeQeNxGZI7NMTlEmKefAMQ,2725 +torch/distributions/fishersnedecor.py,sha256=Rv_OyjkelKUrjdtAuVk8eMRG_cy8qiKGjXNgEENarA0,3665 +torch/distributions/gamma.py,sha256=67kn1C2YRwybKAyRt6ayIcQqYrGRavmJQqZCINlAoiE,3932 +torch/distributions/generalized_pareto.py,sha256=wUgZyItbYV9m3HKBNdPzH7QVf9_0cJgr528XWZLKNrY,5831 +torch/distributions/geometric.py,sha256=RBngoQtIAckbYTSn73Id7edOLTGhdANttBsrt7YkKTI,5221 +torch/distributions/gumbel.py,sha256=62tPcwq8-w9HzSsTcCh127OGDsaPWi7NHKBBXtas_ko,3002 +torch/distributions/half_cauchy.py,sha256=Rk0IdrzwWDdh2nBVZbbxQNN4rhdiXtXSOl9qCCv6q1w,2638 +torch/distributions/half_normal.py,sha256=e9xjQP1zD7oynxyJNLVa0zk1AGrtbgQgxAAWgbrKOuM,2403 +torch/distributions/independent.py,sha256=UWRk4sveW6xR64vbrNUINOiMnCWI2oSmrcSGomExnz4,5052 +torch/distributions/inverse_gamma.py,sha256=2pRcikm37uzzG8VGm0yPbDZSdWYsEPUnpLSMWwZ84-A,2771 +torch/distributions/kl.py,sha256=_mJjMlKkxwlz-snAr9UKpDQgmuMa8JcjJdNPwLJxPn4,31818 +torch/distributions/kumaraswamy.py,sha256=AhGhMxIrpq0-RyIc09A9k0YQIPwL7A9eLMYcYTxFFms,3685 +torch/distributions/laplace.py,sha256=WqwZxvW8Ur5exRGF7iu2fj_CYIgO_Ct4Vra8UH5vP7w,3490 +torch/distributions/lkj_cholesky.py,sha256=ghOGpnV_zOYgqKMAIzJSL7OJEu_2n3jrXITUDPTj6M4,6550 +torch/distributions/log_normal.py,sha256=q2OmSbopn3A8xDlSjLBignD_QFDv7eTW5Ov0mQU9r28,2224 +torch/distributions/logistic_normal.py,sha256=m7cwrUWy-EYlXB1MhpAVEQlS2IZmp4m_k-yObEsU83Q,2240 +torch/distributions/lowrank_multivariate_normal.py,sha256=qnWA_8pcNxBTVl1nocJ2h9y82mc7APDtmtSsKwKpPIU,10178 +torch/distributions/mixture_same_family.py,sha256=bCd9ELluDLQ5jKMX53CJWVb6BQN9LDrg9Lvdda5xu_A,8733 +torch/distributions/multinomial.py,sha256=hK7Kpynlr5Uk3hP2cbsoDZSBmcTE2fnpY9JCTUjcxZ4,5723 +torch/distributions/multivariate_normal.py,sha256=Ediy83Dtof22tsbhEVcSVLkbLWZVMwigfRXEb9k7QR0,11319 +torch/distributions/negative_binomial.py,sha256=KFFBDY-EXa8ylxEX2FQFV6yP-4MbjW4zMNtJG5_-6NU,5171 +torch/distributions/normal.py,sha256=iP5y20oT_yDAPcDAPiITu4lnYf2TNtkKEwPO4oWAEoI,3890 +torch/distributions/one_hot_categorical.py,sha256=-PcTluwbf9UKI6m2C7OlEI4jVEIS9_uV96yjUvU6lf4,5037 +torch/distributions/pareto.py,sha256=7XmKjKbqJEOJ60IB3VoJJnvD5bOlCkwUx0yhZyc1Fqk,2522 +torch/distributions/poisson.py,sha256=ZYOXnjAC_4mQZdd3u61NDpeEf6aAS7egRa6Sc8A1vU8,2477 +torch/distributions/relaxed_bernoulli.py,sha256=Hf6XfCvXqLP3qOLnSwXOItiahS754sEvj7hjwTCMsVE,6121 +torch/distributions/relaxed_categorical.py,sha256=XhE9kz80C5X4Rjny64uPqfSS7ubWTLJm-4G_uA8M8yA,5752 +torch/distributions/studentT.py,sha256=4GT8RrqckofZWFp80AMnzmZvBLKLslH1n0XdSJN1fMM,4131 +torch/distributions/transformed_distribution.py,sha256=X3TnplWYfr3-xJN1gdz91wFGZQpGvzsTjGFCEoTvSgc,9021 +torch/distributions/transforms.py,sha256=6t7TtmRfeisxCZ819dUCmvL3hLwZgnCsKD6a1FS_-nI,46084 +torch/distributions/uniform.py,sha256=Fb75o9UR1cfSb8UP04ah6OXHppJnsLOahVkJDIem2Yw,3348 +torch/distributions/utils.py,sha256=9EeEGgl3z6xnu4D1_D-EgvaAMUt9yPeJ0SEjy6JD260,7986 +torch/distributions/von_mises.py,sha256=jgwiV2-J3BzHNmgscm6XODymFuPX4A4QfpHCh8mQucQ,6507 +torch/distributions/weibull.py,sha256=3IMQkEvqGZ1cq3Z4pqixK6_7QKnDO51UQOsTXlLNA9o,3417 +torch/distributions/wishart.py,sha256=mvcYtuDosgkyCR2CqCWdiIR2HdIs3LscED3-x_Egvg8,14004 +torch/export/__init__.py,sha256=WmOCeZ_lXsOoZcc6cGNI3M9uu4HhG8A6SFQg_tZXxqo,19089 +torch/export/__pycache__/__init__.cpython-311.pyc,, +torch/export/__pycache__/_draft_export.cpython-311.pyc,, +torch/export/__pycache__/_leakage_detection_utils.cpython-311.pyc,, +torch/export/__pycache__/_patches.cpython-311.pyc,, +torch/export/__pycache__/_remove_auto_functionalized_pass.cpython-311.pyc,, +torch/export/__pycache__/_remove_effect_tokens_pass.cpython-311.pyc,, +torch/export/__pycache__/_safeguard.cpython-311.pyc,, +torch/export/__pycache__/_state_dict_utils.cpython-311.pyc,, +torch/export/__pycache__/_swap.cpython-311.pyc,, +torch/export/__pycache__/_trace.cpython-311.pyc,, +torch/export/__pycache__/_tree_utils.cpython-311.pyc,, +torch/export/__pycache__/_unlift.cpython-311.pyc,, +torch/export/__pycache__/_wrapper_utils.cpython-311.pyc,, +torch/export/__pycache__/custom_obj.cpython-311.pyc,, +torch/export/__pycache__/custom_ops.cpython-311.pyc,, +torch/export/__pycache__/decomp_utils.cpython-311.pyc,, +torch/export/__pycache__/dynamic_shapes.cpython-311.pyc,, +torch/export/__pycache__/exported_program.cpython-311.pyc,, +torch/export/__pycache__/graph_signature.cpython-311.pyc,, +torch/export/__pycache__/unflatten.cpython-311.pyc,, +torch/export/_draft_export.py,sha256=e0kipj3WGxPkoxmDfc8aVyUqT2oN8E7dA7nOOoyxHF8,20706 +torch/export/_leakage_detection_utils.py,sha256=cFO_EtruWPRAksV5lS_Je9_FSREizRK6opqkHtEjYfc,3564 +torch/export/_patches.py,sha256=_NKRIISeKTjHCEhL5uD-B5ekQJqb3Mhtb6_bD0fPQfs,10774 +torch/export/_remove_auto_functionalized_pass.py,sha256=tqaXBjtGuXcsmaq3quv444K-79sIejf_cYecxqqp410,2111 +torch/export/_remove_effect_tokens_pass.py,sha256=qtodS8DssfsXbSt8F3SBoWLXmdtpWjAf8tKYiPWl1sY,9195 +torch/export/_safeguard.py,sha256=48fjMGTWwJ26j0B6q4jZG6mOcsIO1SXDbOrgflDESuk,2098 +torch/export/_state_dict_utils.py,sha256=iYYPNWxP2lR13pFWelSIJiNt2PSB3B5CTvvnfOBgkYw,7601 +torch/export/_swap.py,sha256=1s-_RkvEwtUfh99b2H0Tai0NTQxYMrjsHT4x8DNEPm0,19099 +torch/export/_trace.py,sha256=fuln6wFjs9jcw2kgyZlCk03ZEj0IXp6033pySmqAqrE,103513 +torch/export/_tree_utils.py,sha256=PFHm-v6FAcxMb5neZeiPlKH_hjW1OW39Q8S_QwmbuHc,2563 +torch/export/_unlift.py,sha256=slBUeYRcopHwkh4sOqBYq6uFwkzp6EA2MUpEztejtAo,33800 +torch/export/_wrapper_utils.py,sha256=pAHvH2i0I5mn42Cpl4oyTBQzJz0rZjBi2siRQr2khOo,271 +torch/export/custom_obj.py,sha256=H6j0qawn-qiMlmafZDdjAuOOQY7YvQ7pbhiFW71hOMY,296 +torch/export/custom_ops.py,sha256=oaS0ITKDAS05pOsGXJ5o7BefsqWSL0gDYN6MNY-vbMc,2077 +torch/export/decomp_utils.py,sha256=3WGV86fP5QSFIe31u9jyG63N-mhnsLxS5V2mtqRBK5c,5753 +torch/export/dynamic_shapes.py,sha256=mKIt4ibaWzehIhyKTvxf-RyyIiy83OgB8nrnmm29etA,55882 +torch/export/experimental/__init__.py,sha256=XkiG6Gp0dH90QMdwgHGxE_bI7X7PQYiiKQCGS92be4A,17890 +torch/export/experimental/__pycache__/__init__.cpython-311.pyc,, +torch/export/experimental/__pycache__/_utils.cpython-311.pyc,, +torch/export/experimental/_utils.py,sha256=A17QWFUiisPCAWnuD7VGS2b8CSkcvbxwUa0oRPt0lOo,8880 +torch/export/exported_program.py,sha256=8_OaRBkgVMX8WaQQL4nFvL-wRJuwpgRCMC8epiJUTF8,70176 +torch/export/graph_signature.py,sha256=4b4l6Alm0VusSBuaBtxmCVLediCwveFKvip88zV50IQ,27547 +torch/export/passes/__init__.py,sha256=6vyfrZbM2pi3XGyRYPu-J9p9tknELWJV4JfkJs1dqP8,3337 +torch/export/passes/__pycache__/__init__.cpython-311.pyc,, +torch/export/pt2_archive/__init__.py,sha256=vYIe2lq0nWwY77PVEsuIo_0fEWuAeX5E95FhGtXmwZ8,144 +torch/export/pt2_archive/__pycache__/__init__.cpython-311.pyc,, +torch/export/pt2_archive/__pycache__/_package.cpython-311.pyc,, +torch/export/pt2_archive/__pycache__/_package_weights.cpython-311.pyc,, +torch/export/pt2_archive/__pycache__/constants.cpython-311.pyc,, +torch/export/pt2_archive/_package.py,sha256=swalZPPTflPrXjECZxTRuVu1MmYZ3PGldtLRvaNusps,43416 +torch/export/pt2_archive/_package_weights.py,sha256=UdaUHBaCJKLOuWkDuZc5wgVy1-OORZ-tHtkCa8knZFE,8135 +torch/export/pt2_archive/constants.py,sha256=I03jOUF2310XCUciV5QRi-QW3XbcvgUFWBuMWo99cMI,1765 +torch/export/unflatten.py,sha256=LI6zJM6ypsgup677ZNg4pfzeaehJUggrBA6gde4iAW0,75875 +torch/fft/__init__.py,sha256=WqJ1Z1KMnbxX-6KdwopMcr7m-j6R9E0wd-cFlgf9VQ8,55337 +torch/fft/__pycache__/__init__.cpython-311.pyc,, +torch/func/__init__.py,sha256=sUjka9l7rIWURFlItERpQgoeiNfWjMi083ISlouT1cg,656 +torch/func/__pycache__/__init__.cpython-311.pyc,, +torch/func/__pycache__/_random.cpython-311.pyc,, +torch/func/_random.py,sha256=yflUP74RHLvaMdJEgfMFgwwNAWV100nSlstTxYPboiU,9979 +torch/functional.py,sha256=wCYHjWOR2aRcww83f6x9N14G_xS2FZZohHcEe69BLB8,88759 +torch/futures/__init__.py,sha256=FRibgVxq5rx2VBAycXQhBrV7eNdjHQkbDdnF2zkYfiM,14480 +torch/futures/__pycache__/__init__.cpython-311.pyc,, +torch/fx/__init__.py,sha256=u-GEqNYTjXXKCezS9wIYxqfRYBM8EWUbdYpaTqUu6eI,4162 +torch/fx/__pycache__/__init__.cpython-311.pyc,, +torch/fx/__pycache__/_compatibility.cpython-311.pyc,, +torch/fx/__pycache__/_graph_pickler.cpython-311.pyc,, +torch/fx/__pycache__/_lazy_graph_module.cpython-311.pyc,, +torch/fx/__pycache__/_pytree.cpython-311.pyc,, +torch/fx/__pycache__/_symbolic_trace.cpython-311.pyc,, +torch/fx/__pycache__/_utils.cpython-311.pyc,, +torch/fx/__pycache__/annotate.cpython-311.pyc,, +torch/fx/__pycache__/config.cpython-311.pyc,, +torch/fx/__pycache__/graph.cpython-311.pyc,, +torch/fx/__pycache__/graph_module.cpython-311.pyc,, +torch/fx/__pycache__/immutable_collections.cpython-311.pyc,, +torch/fx/__pycache__/interpreter.cpython-311.pyc,, +torch/fx/__pycache__/node.cpython-311.pyc,, +torch/fx/__pycache__/operator_schemas.cpython-311.pyc,, +torch/fx/__pycache__/proxy.cpython-311.pyc,, +torch/fx/__pycache__/subgraph_rewriter.cpython-311.pyc,, +torch/fx/__pycache__/tensor_type.cpython-311.pyc,, +torch/fx/__pycache__/traceback.cpython-311.pyc,, +torch/fx/_compatibility.py,sha256=R37yvvOmOJ5KVkKkJm4RsVzGbhAKyoMDFba2KRpcQ-o,1110 +torch/fx/_graph_pickler.py,sha256=OV5nnm0o90x6ZPC23fCo0ozLEucRKKsQ5PHJaMlZRX8,34801 +torch/fx/_lazy_graph_module.py,sha256=F2zYNnT1w-l5gaVq9qLRuWpYRmTToR-zNR4ai_TzpzE,7358 +torch/fx/_pytree.py,sha256=yJdeitYr930buUWLRIZuTCitjhaDRskJlzV05YDcx50,3762 +torch/fx/_symbolic_trace.py,sha256=6huoFQrGAtZOK-G7yw4zvRMky_9CDJdOJ1imCSovhoo,54990 +torch/fx/_utils.py,sha256=Ep7Ebir2dP7NhowZva1ci5sE1fdAxh5tPLbteJeZNvY,1883 +torch/fx/annotate.py,sha256=Uek_AOziex8JDoBdvg2UMZcmwBglDRx6Xg8cEAwhX9I,1296 +torch/fx/config.py,sha256=L7Vy7NuLuhWYjWj8BVjra-uHIKdVCgOuTwyh3n6_fdI,609 +torch/fx/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/fx/experimental/__pycache__/__init__.cpython-311.pyc,, +torch/fx/experimental/__pycache__/_backward_state.cpython-311.pyc,, +torch/fx/experimental/__pycache__/_config.cpython-311.pyc,, +torch/fx/experimental/__pycache__/_constant_symnode.cpython-311.pyc,, +torch/fx/experimental/__pycache__/_dynamism.cpython-311.pyc,, +torch/fx/experimental/__pycache__/_size_hinting.cpython-311.pyc,, +torch/fx/experimental/__pycache__/accelerator_partitioner.cpython-311.pyc,, +torch/fx/experimental/__pycache__/const_fold.cpython-311.pyc,, +torch/fx/experimental/__pycache__/debug.cpython-311.pyc,, +torch/fx/experimental/__pycache__/graph_gradual_typechecker.cpython-311.pyc,, +torch/fx/experimental/__pycache__/merge_matmul.cpython-311.pyc,, +torch/fx/experimental/__pycache__/meta_tracer.cpython-311.pyc,, +torch/fx/experimental/__pycache__/normalize.cpython-311.pyc,, +torch/fx/experimental/__pycache__/optimization.cpython-311.pyc,, +torch/fx/experimental/__pycache__/partitioner_utils.cpython-311.pyc,, +torch/fx/experimental/__pycache__/proxy_tensor.cpython-311.pyc,, +torch/fx/experimental/__pycache__/recording.cpython-311.pyc,, +torch/fx/experimental/__pycache__/refinement_types.cpython-311.pyc,, +torch/fx/experimental/__pycache__/rewriter.cpython-311.pyc,, +torch/fx/experimental/__pycache__/schema_type_annotation.cpython-311.pyc,, +torch/fx/experimental/__pycache__/sym_node.cpython-311.pyc,, +torch/fx/experimental/__pycache__/symbolic_shapes.cpython-311.pyc,, +torch/fx/experimental/__pycache__/unify_refinements.cpython-311.pyc,, +torch/fx/experimental/__pycache__/validator.cpython-311.pyc,, +torch/fx/experimental/_backward_state.py,sha256=TzC9Uin0ccyk7oG5z7HQPhRP6I0_6-HC28DOXpzUZzE,967 +torch/fx/experimental/_config.py,sha256=s0KFJ2slOpigJzfvohivabnJmZFLIrRSao3Ge2vbgQw,6040 +torch/fx/experimental/_constant_symnode.py,sha256=yQhfLfIcIY0HKznOrhlwb7SN5HqVSqA_gTGkY5JW1UM,1751 +torch/fx/experimental/_dynamism.py,sha256=Yw_M7sAkQbHk35l95i2Rsdhg-NuVDcHCvmktxMtXIHk,4515 +torch/fx/experimental/_size_hinting.py,sha256=rDPx0faznyp4qyuidtjdhb7NIEgpILZJ0eJvH8wSvus,16337 +torch/fx/experimental/accelerator_partitioner.py,sha256=N3JKTaZh8c-pK1h1yTPBjQsPAzgXSKaE3ETFsxKq2DE,47953 +torch/fx/experimental/const_fold.py,sha256=e1xIHGwfk2d622kpYr4hEpXLEFobN6-vV2zphXoLthM,16035 +torch/fx/experimental/debug.py,sha256=i3xyFB3WsaE3k09YInW_WUMTPi9mpNNbfTbB-7jFmvs,811 +torch/fx/experimental/graph_gradual_typechecker.py,sha256=fn6Omc0KCpdHvQKH81Tpb0sJCOWO_nv2Dq3yQJiZSfk,35772 +torch/fx/experimental/merge_matmul.py,sha256=j9LGpyTAhxEfJciK2U_wcOhjYowzoueGsAzMA-iikGw,6096 +torch/fx/experimental/meta_tracer.py,sha256=UdQ3nIArYBNjvT4kOCHA7aUwSrUC10uxFsW5UdiaFKA,11138 +torch/fx/experimental/migrate_gradual_types/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/fx/experimental/migrate_gradual_types/__pycache__/__init__.cpython-311.pyc,, +torch/fx/experimental/migrate_gradual_types/__pycache__/constraint.cpython-311.pyc,, +torch/fx/experimental/migrate_gradual_types/__pycache__/constraint_generator.cpython-311.pyc,, +torch/fx/experimental/migrate_gradual_types/__pycache__/constraint_transformation.cpython-311.pyc,, +torch/fx/experimental/migrate_gradual_types/__pycache__/operation.cpython-311.pyc,, +torch/fx/experimental/migrate_gradual_types/__pycache__/transform_to_z3.cpython-311.pyc,, +torch/fx/experimental/migrate_gradual_types/__pycache__/util.cpython-311.pyc,, +torch/fx/experimental/migrate_gradual_types/__pycache__/z3_types.cpython-311.pyc,, +torch/fx/experimental/migrate_gradual_types/constraint.py,sha256=Ayq_LWRTnFtjiiWS1TG63G0UTJ8D-s5R3P3QvR5wxZ4,20836 +torch/fx/experimental/migrate_gradual_types/constraint_generator.py,sha256=uOPBc9B27M479s32uV3Kmb32l-lrqkTI8heaDYoK7mc,62331 +torch/fx/experimental/migrate_gradual_types/constraint_transformation.py,sha256=WjUlxkoyqHgpWQDHBJZVhvoDI3tyULmv_OnHJ6DJ8R0,46141 +torch/fx/experimental/migrate_gradual_types/operation.py,sha256=v7haImk8KCWq9XBmzu2MWoApiGGQsQY4xF-VChtRBVI,287 +torch/fx/experimental/migrate_gradual_types/transform_to_z3.py,sha256=gv0a_tG2THM3H3H0tPSI1KTs2gljAyXvLXCw8ut8HUg,20113 +torch/fx/experimental/migrate_gradual_types/util.py,sha256=7FpQPVuVeocjt7ZC8UYSiCES-Qe-D2HQzYWSTvaR5lY,1485 +torch/fx/experimental/migrate_gradual_types/z3_types.py,sha256=9CzjT7bWhU2CUNwZrQcdI7PDY8haxTrzGBRw_V7X7dk,807 +torch/fx/experimental/normalize.py,sha256=JMVnUQiiyvmyXpQdItb-7sTtVkqKkAKfX1TMolUXRpY,5795 +torch/fx/experimental/optimization.py,sha256=IJ5z0GKvS8YsDdQrUFO_8p6C7FiBWpkwc69ODFm4ByY,18443 +torch/fx/experimental/partitioner_utils.py,sha256=OlitD91AGdurF5-P7hITygyWArX5vL9iL95iZlwkGXA,12295 +torch/fx/experimental/proxy_tensor.py,sha256=taSn4i_wm8oSx4ZIKWBU9dBcbUF73vCCfG38rqmvSiY,118570 +torch/fx/experimental/recording.py,sha256=dkWCY4HVyc2TsC1sVX7VoxSU4SaUwWS2TMgEBJFa4-o,21219 +torch/fx/experimental/refinement_types.py,sha256=888K8p43hadBWu9MPg4lrPeDIF8g9XYRnmI5nnwFYpw,451 +torch/fx/experimental/rewriter.py,sha256=1VezunMnlBoRWT5qRiILS-J4RXgLQCPBU2jX2luTAJQ,5687 +torch/fx/experimental/schema_type_annotation.py,sha256=rs5bSlVu4sNMDcNFXpVU2fyk5VS39GgGOuZLvxrIEx4,5802 +torch/fx/experimental/sym_node.py,sha256=JVsAEGOLvp2lij_rPCVFDfociNpRtgK0ihxDPArnV1s,71202 +torch/fx/experimental/symbolic_shapes.py,sha256=Kn-TNCaGGxk6NDPiDYjHSJi5_QT-LpyGEo9_gPIvbho,365640 +torch/fx/experimental/unification/__init__.py,sha256=zJTtn828T13Feto4TKixj88dtdoRz5_TjoseKssyzeo,196 +torch/fx/experimental/unification/__pycache__/__init__.cpython-311.pyc,, +torch/fx/experimental/unification/__pycache__/core.cpython-311.pyc,, +torch/fx/experimental/unification/__pycache__/dispatch.cpython-311.pyc,, +torch/fx/experimental/unification/__pycache__/match.cpython-311.pyc,, +torch/fx/experimental/unification/__pycache__/more.cpython-311.pyc,, +torch/fx/experimental/unification/__pycache__/unification_tools.cpython-311.pyc,, +torch/fx/experimental/unification/__pycache__/utils.cpython-311.pyc,, +torch/fx/experimental/unification/__pycache__/variable.cpython-311.pyc,, +torch/fx/experimental/unification/core.py,sha256=TsI4QAfsoffC1cjq1moVtbVobw-P8rJlDUQ0hiMAAfM,2780 +torch/fx/experimental/unification/dispatch.py,sha256=9oSoC5tHuRoDMwgjgI7NvGVo-aW9qRrgz6SBNEnx6Lc,207 +torch/fx/experimental/unification/match.py,sha256=nIPkwl-CJykPeZ9s97iMksw4l3OV2eOIBxVaIYBLD5s,3414 +torch/fx/experimental/unification/more.py,sha256=5JZ2WU4mbfG-a9WkuFzNNap28R5-aXVfCM5PxqbdPAc,3158 +torch/fx/experimental/unification/multipledispatch/__init__.py,sha256=xjEjlwV6TTwkgteOREoXUJwMTiNnpJN3YfBALC9eslU,139 +torch/fx/experimental/unification/multipledispatch/__pycache__/__init__.cpython-311.pyc,, +torch/fx/experimental/unification/multipledispatch/__pycache__/conflict.cpython-311.pyc,, +torch/fx/experimental/unification/multipledispatch/__pycache__/core.cpython-311.pyc,, +torch/fx/experimental/unification/multipledispatch/__pycache__/dispatcher.cpython-311.pyc,, +torch/fx/experimental/unification/multipledispatch/__pycache__/utils.cpython-311.pyc,, +torch/fx/experimental/unification/multipledispatch/__pycache__/variadic.cpython-311.pyc,, +torch/fx/experimental/unification/multipledispatch/conflict.py,sha256=mvp3SkUZ0fKpCuCATDiFZmB_1euN3RsfoLvya3jfrBg,4556 +torch/fx/experimental/unification/multipledispatch/core.py,sha256=FS6cUF9-0llkG63x7hUCeUDdXVcjDnN_kTA3ngfniT4,2859 +torch/fx/experimental/unification/multipledispatch/dispatcher.py,sha256=11tXE2QB0pJLEvv4o5K0mpPDpstMoyTc1YdwB0ZWsAw,14012 +torch/fx/experimental/unification/multipledispatch/utils.py,sha256=BHnpSEp0idrQU2H1V3SKorvS_yf5WoWlZzD1C2laz7U,3958 +torch/fx/experimental/unification/multipledispatch/variadic.py,sha256=irgdEYhgS2XLxgdZZLPrApQPiB7c2wh9jZt4J_yu4Xs,2962 +torch/fx/experimental/unification/unification_tools.py,sha256=uz1BytKNE5rzW327q9PWOZUhAeybmBOFPDGwRy-OOeo,10567 +torch/fx/experimental/unification/utils.py,sha256=WiCStj6-1r0Wpnh7wsEeg2163W47nBJwXa6qDPuk-ow,3129 +torch/fx/experimental/unification/variable.py,sha256=mwytnm5NDShbcqIc6Jw3ARu5qCVPErP0WmV1HYw_exw,2057 +torch/fx/experimental/unify_refinements.py,sha256=60MsjQaDJfFvjcb4z4pchIllfuHcS5Mt3BPpes5oP6Y,3130 +torch/fx/experimental/validator.py,sha256=Knxbn2W2bzflJwWa20drEl641hUCIzxm-_Ps0wcUjTg,35328 +torch/fx/graph.py,sha256=jaHeKmUg7LUlyOFOIPM4yfFXB8LIvE9MQlddZJx0Lzk,93020 +torch/fx/graph_module.py,sha256=9JrwY-boft8D1w6CbEbApO2sCZWp8kxhNn8YRcXfh6w,47314 +torch/fx/immutable_collections.py,sha256=9OeeV5ZPLcswvuNEo7xEepwdUJsW0JYdIwaNNYCOKG4,3269 +torch/fx/interpreter.py,sha256=KJWVqlLkFFpRkZtJAiyegE1WJlVZ6LTuW2siNecQ_iU,26485 +torch/fx/node.py,sha256=kDMiTQirpFSyPpCFgQugAwfEj3z9fiOQFE45w2ByPLE,35645 +torch/fx/operator_schemas.py,sha256=VSfSKUmgAdINLiy5aXWqg-gsDeJeqDte9VY-B7ELP0c,24005 +torch/fx/passes/__init__.py,sha256=qXzO3_W60MDLQ1l4Rymt3ZwAxtbmkQ7Plsx9kQjWIm8,263 +torch/fx/passes/__pycache__/__init__.cpython-311.pyc,, +torch/fx/passes/__pycache__/_tensorify_python_scalars.cpython-311.pyc,, +torch/fx/passes/__pycache__/annotate_getitem_nodes.cpython-311.pyc,, +torch/fx/passes/__pycache__/fake_tensor_prop.cpython-311.pyc,, +torch/fx/passes/__pycache__/graph_drawer.cpython-311.pyc,, +torch/fx/passes/__pycache__/graph_manipulation.cpython-311.pyc,, +torch/fx/passes/__pycache__/graph_transform_observer.cpython-311.pyc,, +torch/fx/passes/__pycache__/net_min_base.cpython-311.pyc,, +torch/fx/passes/__pycache__/operator_support.cpython-311.pyc,, +torch/fx/passes/__pycache__/param_fetch.cpython-311.pyc,, +torch/fx/passes/__pycache__/pass_manager.cpython-311.pyc,, +torch/fx/passes/__pycache__/regional_inductor.cpython-311.pyc,, +torch/fx/passes/__pycache__/regional_inductor_invoke_subgraph.cpython-311.pyc,, +torch/fx/passes/__pycache__/reinplace.cpython-311.pyc,, +torch/fx/passes/__pycache__/runtime_assert.cpython-311.pyc,, +torch/fx/passes/__pycache__/shape_prop.cpython-311.pyc,, +torch/fx/passes/__pycache__/split_module.cpython-311.pyc,, +torch/fx/passes/__pycache__/split_utils.cpython-311.pyc,, +torch/fx/passes/__pycache__/splitter_base.cpython-311.pyc,, +torch/fx/passes/__pycache__/tools_common.cpython-311.pyc,, +torch/fx/passes/_tensorify_python_scalars.py,sha256=XVt3XHW-z_uBSsWcq5jJjoD88A2qMfLk_EKGvE-OTac,17902 +torch/fx/passes/annotate_getitem_nodes.py,sha256=8uE8EHJnlxC_BJmSTOl_r-ixz-8GGtJKc_xipmVQYt8,3561 +torch/fx/passes/backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/fx/passes/backends/__pycache__/__init__.cpython-311.pyc,, +torch/fx/passes/backends/__pycache__/cudagraphs.cpython-311.pyc,, +torch/fx/passes/backends/cudagraphs.py,sha256=iRy6rWuLOPITB_-gkgJmK8eGYluYIz7ZGjYaO4x41fc,2079 +torch/fx/passes/dialect/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/fx/passes/dialect/__pycache__/__init__.cpython-311.pyc,, +torch/fx/passes/dialect/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/fx/passes/dialect/common/__pycache__/__init__.cpython-311.pyc,, +torch/fx/passes/dialect/common/__pycache__/cse_pass.cpython-311.pyc,, +torch/fx/passes/dialect/common/cse_pass.py,sha256=CBRTo0Hw3MOIzh65C-f-l-eJ3nWt4IBgUHPg_klizyA,5248 +torch/fx/passes/fake_tensor_prop.py,sha256=U2hOUtbxk513ixyAdhqRVSPi4-nWiWcDXL_q2Erhj9Q,4463 +torch/fx/passes/graph_drawer.py,sha256=39Ldfv-6jqoOZi4A2cbxg1S8OWWGy6BR78fHWSu6owI,19463 +torch/fx/passes/graph_manipulation.py,sha256=GvMbuMQS1VbGOAW5evE9Ykv39nnxMwb1uwKlsPMCSxk,4114 +torch/fx/passes/graph_transform_observer.py,sha256=HpFPWK8G_AtuF052epvJ8MpyI1OCzDtwmlnciAE8fmA,8536 +torch/fx/passes/infra/__init__.py,sha256=UuBORluunUu4B8DKb0OlqkNIm71XUW_ZD_utIvqnBLs,27 +torch/fx/passes/infra/__pycache__/__init__.cpython-311.pyc,, +torch/fx/passes/infra/__pycache__/partitioner.cpython-311.pyc,, +torch/fx/passes/infra/__pycache__/pass_base.cpython-311.pyc,, +torch/fx/passes/infra/__pycache__/pass_manager.cpython-311.pyc,, +torch/fx/passes/infra/partitioner.py,sha256=wIS3UHWHTUX1kfOqi0NAupe2VVTOMv4RAu0D4rcQ23U,18066 +torch/fx/passes/infra/pass_base.py,sha256=m4Ml6B2wFnMn3WXHE0ayJ3W-Bn9me1x0wvSUWbaqyRA,2507 +torch/fx/passes/infra/pass_manager.py,sha256=kav0ykYHWq71GiTAg3r_HngsboEr-29K8tl3DTNlru8,10393 +torch/fx/passes/net_min_base.py,sha256=rDCEFxDCXSNcU5NrdCw9bfxN8aVzHbspCGRZ5obCOzM,36975 +torch/fx/passes/operator_support.py,sha256=zczkPh0eYCUwov4Sn1EkB6ABQwBK703Gjj_WbKf__tE,7702 +torch/fx/passes/param_fetch.py,sha256=6gOQFJ-MRJF6g5sp8Typ9j6pBR9ibsVlJ50xjM2af4g,3741 +torch/fx/passes/pass_manager.py,sha256=NCA5OAY-5mW-lQWBjqfPQkB54pfet3WyRD17YGRSBVo,7112 +torch/fx/passes/regional_inductor.py,sha256=LIbuiUknLmbWUXfSvoZ6F9JmeKEOARp2UiTAjoZZz88,10605 +torch/fx/passes/regional_inductor_invoke_subgraph.py,sha256=v4gT1ax4NVAp195BrhBsW0x9Tn1Z2tFW5unJakxeFmw,6222 +torch/fx/passes/reinplace.py,sha256=irI-uQVXN6FJKSu3lImDtvUkigM88qtB3wIsksNpMwM,35628 +torch/fx/passes/runtime_assert.py,sha256=_A-AFI5qnhj5HA6S7gpF_UwHZIjbmK2mLvCo95wHOxU,31447 +torch/fx/passes/shape_prop.py,sha256=XS1gpmCthWNGYGnrr3p73qB7l56X3MXaFllNWp3IAZ4,8313 +torch/fx/passes/split_module.py,sha256=RwhUqX4LKhvsxMZmD94drZLd_D-mE8WlpPsTN8BTVQg,29029 +torch/fx/passes/split_utils.py,sha256=OJZEK0OZK0yJEKFQglyduiyylJNB7F7Ts0S-S_kEL9o,20210 +torch/fx/passes/splitter_base.py,sha256=hYyT-LBQD9LJsN_D-5wTZOgw5hCT6otBvPkEjAqB8tE,44592 +torch/fx/passes/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/fx/passes/tests/__pycache__/__init__.cpython-311.pyc,, +torch/fx/passes/tests/__pycache__/_test_split_utils.cpython-311.pyc,, +torch/fx/passes/tests/__pycache__/test_pass_manager.cpython-311.pyc,, +torch/fx/passes/tests/_test_split_utils.py,sha256=WybDpQFnUmPDHIjIIdZSMtuivIhyskR1nOpW4_fgozw,22497 +torch/fx/passes/tests/test_pass_manager.py,sha256=_KIwhlzOTY0XQb2WeaYYOFuBUKDtLFRqYeSxHVrS76E,1873 +torch/fx/passes/tools_common.py,sha256=DQ2__ux4f6AC5KzmDAIo2kysaFW2dgQ5e5PGPurCHtQ,14152 +torch/fx/passes/utils/__init__.py,sha256=6xOYtGu1w_nVvuaxvBoQkSRtxkiIkPh-MzuXjmNC1hI,74 +torch/fx/passes/utils/__pycache__/__init__.cpython-311.pyc,, +torch/fx/passes/utils/__pycache__/common.cpython-311.pyc,, +torch/fx/passes/utils/__pycache__/fuser_utils.cpython-311.pyc,, +torch/fx/passes/utils/__pycache__/matcher_utils.cpython-311.pyc,, +torch/fx/passes/utils/__pycache__/matcher_with_name_node_map_utils.cpython-311.pyc,, +torch/fx/passes/utils/__pycache__/source_matcher_utils.cpython-311.pyc,, +torch/fx/passes/utils/common.py,sha256=z7nJRQbMOvXqQ4_ndwHwLM4l70Q8EfHLeYSefE5sdvU,3185 +torch/fx/passes/utils/fuser_utils.py,sha256=yvh8SAFtU-Yv0fXZjT96d1ZqrQuD8urYewJi__XIi3c,11199 +torch/fx/passes/utils/matcher_utils.py,sha256=0vlo5_Ot3id2o0xfMxx7dCFWfPpBtd6TygG3w2AnfwU,17681 +torch/fx/passes/utils/matcher_with_name_node_map_utils.py,sha256=wtKkIKjnW3fduXdkIRNiIA4zZ1GZVTpSpEPcGGN3NIk,4358 +torch/fx/passes/utils/source_matcher_utils.py,sha256=uTDD1LsizcSe3FnyaYOlPaLAztc4u1K4RXEDa8tJR74,7312 +torch/fx/proxy.py,sha256=nvXNw2bSVDz938XTUvoQe1NSEA7HIV6HNDqMEtIQnCE,36137 +torch/fx/subgraph_rewriter.py,sha256=JJY756A_75NKdZMvShpcFW2zFJqtjxplDzBl1Oi76HY,16808 +torch/fx/tensor_type.py,sha256=vb3E5_qLGIm2sFQfu3EtsxpIRl3SSQ3Z7lQC7HBppZk,3458 +torch/fx/traceback.py,sha256=m_wZNmKb_jv4AMyjlKDLG8I530KwquBpJUyZFLsITYw,20486 +torch/hub.py,sha256=rsemIf1w8RKV0YMH0kYpPmPnpXjIB60elCUe75ZIXvo,33647 +torch/include/ATen/ATen.h,sha256=BHC8UJw-Kkgse6RoCnFPsJL1ZTodth1VYbFstbIiByc,1361 +torch/include/ATen/AccumulateType.h,sha256=kzgekl85-Pf2id9qB_Uo2y5QyXxMGNcamXG-yvueh0Y,6365 +torch/include/ATen/ArrayRef.h,sha256=ofa5-yWJikSN5a0VsjvX-G2qByUSKsCE7z2Hdd1zZjE,298 +torch/include/ATen/Backend.h,sha256=rAIfvXr2KgiCtq5uAiBapYC1r3SLtupFjrmbpeXPhC4,297 +torch/include/ATen/Backtrace.h,sha256=yKbvQ16ihAHBkGQGKh-6pjm6jK-6cp6tjKAhs-3pQHI,300 +torch/include/ATen/BlasBackend.h,sha256=dM74cin5ObOdOP67BUREoUocYhMV043FeI0vom10vu4,1371 +torch/include/ATen/CPUApplyUtils.h,sha256=HImZBT4CCuqV47EmefXlaFdeuoVFOBt8Ey0UQoeNu1g,10872 +torch/include/ATen/CPUFixedAllocator.h,sha256=-z7Tz9exkjw8w2Lq-yjCgOK23L5YbPW4usEjJ2QEPEg,1136 +torch/include/ATen/CPUFunctions.h,sha256=UbQ7FfyyuDj2oE7kHAhxaGD9LFUdleSeWwUHpH795fo,2208 +torch/include/ATen/CPUFunctions_inl.h,sha256=wNOp4CYqsHfnmMFuTWah3yHX6qnpu0-jTH5dvhQpTBE,27816 +torch/include/ATen/CPUGeneratorImpl.h,sha256=qogO5E0lj9VQz6QVdoaNKSoFFrGagDBJbXVpEnO80Rg,1781 +torch/include/ATen/CUDAFunctions.h,sha256=nwdEjF4wGGMImssj3hmzSOQ_kz4ZKSeT4WQ8UdiBj98,2209 +torch/include/ATen/CUDAFunctions_inl.h,sha256=TLfsai3knyVk4loKvnZypKYDvgrePA2Qj57qxKsUWjk,33868 +torch/include/ATen/CachedTensorUtils.h,sha256=Rfo1BCyQrhFdY2tO7o2beQD8LxU-7pyZWrdjtTGIacw,1261 +torch/include/ATen/CollapseDims.h,sha256=dWXH85U4p4jsam0cS6pw0m9xunZ34HDLbF-0NR9eCQk,2814 +torch/include/ATen/CompositeExplicitAutogradFunctions.h,sha256=n7GTKOgLj3NDa9ELmt5PBi2AsqMZP26HUPD-YpGEqHQ,2230 +torch/include/ATen/CompositeExplicitAutogradFunctions_inl.h,sha256=l_yanwyYWK0VlpUAXZGPF0CyYMVCPCdeZKFu0ZQMeDc,41518 +torch/include/ATen/CompositeExplicitAutogradNonFunctionalFunctions.h,sha256=5fa_cZ7Y3np6u7c94LzhMe5g1F_aPAcJVDvNxL1NZcY,2243 +torch/include/ATen/CompositeExplicitAutogradNonFunctionalFunctions_inl.h,sha256=YdYab88X8blUNj8L7fuOPqdk7pWinI4wMoA_xCpwmuc,26775 +torch/include/ATen/CompositeImplicitAutogradFunctions.h,sha256=wNw8oZ66wh3uuYpxU9X-e3Xd9umlWxbhOWKqja8oE2U,2230 +torch/include/ATen/CompositeImplicitAutogradFunctions_inl.h,sha256=zC7mIJFgGqz-5-t4nJJQq5J-ugNmJDYMZoJfBqhIk2Q,35601 +torch/include/ATen/CompositeImplicitAutogradNestedTensorFunctions.h,sha256=LtEbQKDO2Ez1IEVsjvy4GT2ZzDPOaWhFmIeHimLtIPs,2242 +torch/include/ATen/CompositeImplicitAutogradNestedTensorFunctions_inl.h,sha256=ViwELt6UAz5OKrc24yLg7-VTwdrqRge4iq7K6JseUfU,1380 +torch/include/ATen/Config.h,sha256=7gV7Z3O9A5pb1NOou6cG0GnfDIXyO5CUvuMvTqEaVqw,1055 +torch/include/ATen/Context.h,sha256=NcxvO7ZnY-PwnJo65A92E2axXOGJUZF0wz4CZWQkw7E,25579 +torch/include/ATen/DLConvertor.h,sha256=-LqU122QRzQjYP2f5PVJ86w55oxIKNpz9SGnnbPH51c,3168 +torch/include/ATen/DTensorState.h,sha256=5Tf_4vlPJcvhauK-lyQU9FI9nicxD_7DLDbA1nukmsE,1270 +torch/include/ATen/Device.h,sha256=vefXE-ivuf98uv1Px06EGn9MycHCV350yNT9g3JCxXY,296 +torch/include/ATen/DeviceAccelerator.h,sha256=5Z6wZLmAOnQybiXzifqxxkdn5DV32H349h2JT4ELXRM,4983 +torch/include/ATen/DeviceGuard.h,sha256=JMlpTXKWyXFaIIj_ZehQgfM5bP2EKdg4VnPbHAxy0JI,1419 +torch/include/ATen/DimVector.h,sha256=6Xv8E-rdaF4g-zRjm8aVO2AF41P61MpNjUbH-7OhL00,300 +torch/include/ATen/Dimname.h,sha256=KGwDTru2t_tk24ewGxmJexXVeYAOdQOkv23W1yXBE64,285 +torch/include/ATen/Dispatch.h,sha256=0cQd6O-2y6OwcdhzcnIBE_HYY_HVn7-ns21NuzkcUgY,39011 +torch/include/ATen/Dispatch_v2.h,sha256=spV91MM7IE6FwA1gQ8KsHRBs13SK8opA1Zj28BvaBk4,58876 +torch/include/ATen/DynamicLibrary.h,sha256=578X2i8u8535Av4VGp0XYwSPY-_rbvt7B0q0QaxTYeY,933 +torch/include/ATen/EmptyTensor.h,sha256=ZL1Pz5KThapI0aPLX5VW37HLtGP5ifznxv3seu736Wc,4946 +torch/include/ATen/ExpandBase.h,sha256=51GgMs8O36068o3xhE-2NjvOTV0KwVzyDbWOglc8GCI,1168 +torch/include/ATen/ExpandUtils.h,sha256=7eWtQgp1V6m7yauLzdXwY0LxBrbXxHvap8KYmFk8jAg,16935 +torch/include/ATen/Formatting.h,sha256=a7uH36sKyuNzq8kFdXO-214shUmRw8KmcnMt8VOTMJw,288 +torch/include/ATen/FuncTorchTLS.h,sha256=Fl4PmCDf7Sf6JM2ApHD5jbSEZBuG_ZCG_Buk_y2h8F0,2070 +torch/include/ATen/FunctionalStorageImpl.h,sha256=GD9QhDinZ0CRwSvD0U2bO-6UMOuXDPpOJfJmiF8cFdk,9905 +torch/include/ATen/FunctionalTensorWrapper.h,sha256=1sFB2PgE5pk9rqmKhmRuA-G_RcFeJ0strxwekBYq8Nc,18661 +torch/include/ATen/FunctionalizeFallbackKernel.h,sha256=LWHVmU55UcNZEyK4w1paUBkA1hqTkUSbGny0hoi_GMo,2082 +torch/include/ATen/Functions.h,sha256=Bwble7JPY07bs0qUboY0xUCEQ8UiCeX_kQXLGWq9kYQ,56234 +torch/include/ATen/Generator.h,sha256=HOTLCyhKpa4uAF4taTguJPTKvRhBXooFbVKKEomg6OM,300 +torch/include/ATen/InferSize.h,sha256=RTa90ICq20CJZ-pUshTdcGLtgWMeYEi1-HRnc7WQugo,4182 +torch/include/ATen/InitialTensorOptions.h,sha256=AjYaGyhq4NcQ3eyNnj6ctamEghynIVm1guG3VFczdOQ,693 +torch/include/ATen/Layout.h,sha256=tRw-433FrW3h68PD2DQO95UHaUt8yLMdMHiydQrCv20,296 +torch/include/ATen/LegacyBatchedFallback.h,sha256=wegTT0gYFlyoX8I7IuXa7H7o8nGHxZTjI7UadTUSd64,1228 +torch/include/ATen/LegacyBatchedTensorImpl.h,sha256=fS03y5jlojJCq3738MG4YhK2ocIfH6yTdv2Q92rxax8,5817 +torch/include/ATen/LegacyVmapMode.h,sha256=lnhuI05JjTfS7uL4feKP-iLZ5C0eLja6i3WysAYOO6U,1181 +torch/include/ATen/LegacyVmapTransforms.h,sha256=b2ZRmNr3IfHKM-mTVpjYF5ti5yyUKrLZNUh5dg-7G1E,8068 +torch/include/ATen/LinalgBackend.h,sha256=eBFKPJs-zcaW1GEVbojXthxShQ404cfTV6f5VqjEWmA,973 +torch/include/ATen/MapAllocator.h,sha256=0dB-x0W0HKVhxBcTy_O6qEkIfl-YWuF4M67yKH7rLlg,3942 +torch/include/ATen/MatrixRef.h,sha256=CSsbtMITAlDwQQ7nG9REtvsZsb1Ft4xcd3s08NXiQos,3260 +torch/include/ATen/MemoryOverlap.h,sha256=E06jBGvemBeQzG5f83ZyGbyUKUiwypNqLkhkBB3M_eQ,1541 +torch/include/ATen/MetaFunctions.h,sha256=F4QJuE1qvgTMD_pxybgCc3UQH3c7GqdjNB3BdRFvMsE,2209 +torch/include/ATen/MetaFunctions_inl.h,sha256=FN3AJWSeQ8td2LcXCEB-k1clqCSMfEo__c77-3gcxjo,16440 +torch/include/ATen/MethodOperators.h,sha256=o-x3zs_o4h1VQ5k8PFJzwHrj0lIBCjVKw-4qL55riHQ,16000 +torch/include/ATen/NamedTensor.h,sha256=pWHd3I9W1OQTcP7xqgD7y8M_xLuT8escT3VngDS3TGY,289 +torch/include/ATen/NamedTensorUtils.h,sha256=9Pa61OV7pxXtSasv9kKpzvjN-F0KSrhoH98s9EopgqE,7009 +torch/include/ATen/NativeFunctions.h,sha256=LWIAXQPecYfBcBsjLcSxaDVisA2gEq1Vg9fz4jmRkhQ,62162 +torch/include/ATen/NativeMetaFunctions.h,sha256=41q3_2hFOlBi9fxI4oVzhqtsBo0V_KVZe7nrpHeVbvE,58623 +torch/include/ATen/NestedTensorImpl.h,sha256=Ut705YrlQpqGLCseQFvHVIF2zK2p1-0x504sT4NM4YU,10514 +torch/include/ATen/NumericUtils.h,sha256=DFP-P8EXkAhpX3xIzZljpzG8BdNxYZYc17KMTlMcY9o,5391 +torch/include/ATen/OpMathType.h,sha256=DZeNTy6bHsuldLJLLlnqHd_5atCSYt3BXBw1m_ba1cw,1835 +torch/include/ATen/OpaqueTensorImpl.h,sha256=-DEfs1JdJ9WEs6rpfXGJLs6atOM24dUKf7JOGWl7gbU,7069 +torch/include/ATen/Operators.h,sha256=KfabKVPG5XPDljzFvQs_ZQhwyRB05nHrfZyGaA7jJM4,60181 +torch/include/ATen/PTThreadPool.h,sha256=B2QRG9BWowrxxpNvjj39vsVk3pFP54MkJJ4BAJR3mvQ,645 +torch/include/ATen/PadNd.h,sha256=UkJhEHwFKN-TZcN2i6XJEDprHoSSqUJo46diM8SttB4,380 +torch/include/ATen/Parallel-inl.h,sha256=vFqIxc3KSSdtx6Z_dkECPkgMy0AbMMhTvAd75nFvnZE,2547 +torch/include/ATen/Parallel.h,sha256=7tfhOczECJB8NB9Lc2kL5p1KVbQQKyMGksueVtBxvOk,5039 +torch/include/ATen/ParallelFuture.h,sha256=4vx_WX3OIIbFlQ4jkn44Kf3waTXNnEyd48iquBO5qVw,553 +torch/include/ATen/ParallelNative.h,sha256=p3C7kunTT9v2KvAOt6aWd1lh9p9VFPmHDBiPGxlhQvg,546 +torch/include/ATen/ParallelOpenMP.h,sha256=W6LXF3k0MoE21pu3Q7MlyI3D03-2dxHyU4c0HslE5ns,1537 +torch/include/ATen/PythonTorchFunctionTLS.h,sha256=wGyBA6K8qfWh_owJBAbdlcZSArbf0Ff9aVOlD6f25TY,1502 +torch/include/ATen/ROCmFABackend.h,sha256=1mY6byW5IpdUXVGZugzREWnh6QqX75B6OZfIeFReEDU,977 +torch/include/ATen/RedispatchFunctions.h,sha256=j55-OR8L9JVlihXi_IgR4kV-c03SvC7UKvB4FayID8w,2271610 +torch/include/ATen/RegistrationDeclarations.h,sha256=wgir-jUS9jh_nG7KopQZbTlF3q_5GdpN1FtaG_HKR2E,936287 +torch/include/ATen/SDPBackend.h,sha256=cmXomX1RjUAY9uCh2-bSn9FzRpMjFMvUF6bWJBIUczE,507 +torch/include/ATen/SavedTensorHooks.h,sha256=c3B9NplJ5PeOU3u3por0HtEbClY_tijozi0EwK1eyx8,2740 +torch/include/ATen/Scalar.h,sha256=LSNWNrn_jf_qZSVOFXeONdYs5wdSE0uRzKfNX78c_V8,298 +torch/include/ATen/ScalarOps.h,sha256=71v1_d2-kJqm1LD1IzlsFgAY4mufIgiryxmBo3-qaBQ,1849 +torch/include/ATen/ScalarType.h,sha256=SfIfa3E8pBOkFD19qgu43TLOGKptlPIOZrTa4NmD5Pc,383 +torch/include/ATen/SequenceNumber.h,sha256=3lZRqokrZ29LTAwEAFVUso6WrxBO_ZE3rTvvN2DKqy8,587 +torch/include/ATen/SmallVector.h,sha256=gNWDIANGgC7StyDTrqVTMbA2jQSV1AWiwVx397q8n8o,301 +torch/include/ATen/SparseCsrTensorImpl.h,sha256=utiY7n547k3fbK0jeN4kiQdzBCX2dhVUzS2pPEZp_NI,7421 +torch/include/ATen/SparseCsrTensorUtils.h,sha256=p7oJc8xJwel9HIi-t1TsydwyoYwKGoxTby8gkGGoxmM,18043 +torch/include/ATen/SparseTensorImpl.h,sha256=hwtP016bri4KpgnUyu5pueXUDh3DlOAYGUckKim9MC4,15673 +torch/include/ATen/Storage.h,sha256=X1c-p8qXYDFlMsVH5Kw_dgncGyQrrKuzyA0PaY9lY-o,297 +torch/include/ATen/StorageUtils.h,sha256=_VKiR9j8cAdZAHUyT5FA0C9E7HfTDzQxOCv41tMRDNk,1562 +torch/include/ATen/Tensor.h,sha256=C6FD9cacNwOuxJD4_JqjHg34qxaDUZMXWq2GZvm2Yqg,298 +torch/include/ATen/TensorAccessor.h,sha256=zASg1wYR2DpIcYfZnGqgEd9QG8RNWhzJT0mcE7D-WWA,305 +torch/include/ATen/TensorGeometry.h,sha256=cOYlRlEK5XQy5DryIrAAdrJ7uFg0nKnUMROreNDeWc4,4810 +torch/include/ATen/TensorIndexing.h,sha256=x0pJf2Z_cC3wxn00F3QjPGT8fXo7xCm2vCnlA3yPtKo,25356 +torch/include/ATen/TensorIterator.h,sha256=BqdFbEK1OOLkBCn_sZTprRt4ZzUJKBSp4kL9jUCHdxI,39240 +torch/include/ATen/TensorIteratorInternal.h,sha256=N-f8g2bYnUwKBFz_M4gaPKC-7u7SOGjz8Waa7QT3P8U,2176 +torch/include/ATen/TensorMeta.h,sha256=FvY4aUQbKUy3jLV8uOL5whccUfE337DbpxewhG3CEHU,5288 +torch/include/ATen/TensorNames.h,sha256=JXurFfj4LhYDhrpqS9qlPPci4bA-jaleQgAl8fx286A,2825 +torch/include/ATen/TensorOperators.h,sha256=rFW5uIPy2FvyDhV_Cv-SX4HluU7mkR-4sbt2kGF313g,2745 +torch/include/ATen/TensorOptions.h,sha256=58KP7XEfy87cJVZ8fXD1QPvhTMLTYtcV0DUiJSo5jgs,303 +torch/include/ATen/TensorSubclassLikeUtils.h,sha256=vYN0jOsm3jMh_yNKtB3GZzbnImJa-LAthHJbhRKlkVg,3484 +torch/include/ATen/TensorUtils.h,sha256=_-KPHNuzJz-5-HdisWWhK2zbk7CoPgSEQJpymCM0kms,6212 +torch/include/ATen/ThreadLocalPythonObjects.h,sha256=ucUPlRvWSUvfbXwa8qAyrkr1nwAfgqy9veUFsKvxVhg,861 +torch/include/ATen/ThreadLocalState.h,sha256=BapXkyAgHfntv0Dya1DIvjfwPHRO3DLYIHlWzTi8oM4,4608 +torch/include/ATen/TracerMode.h,sha256=zrQHJTz242lMURs7jQp5cIEMYIi-r0w0GL4LjJiqV7E,5763 +torch/include/ATen/TypeDefault.h,sha256=gItr5F3ooMk0KcBhYl07hHAelS1_Hn69maVJ871RNF8,697 +torch/include/ATen/Utils.h,sha256=bLoH6TfXKju5mLSMXCWSTGEPAppiheEEdly9jiVMgzw,3826 +torch/include/ATen/Version.h,sha256=gwJrgxcrcmw3WBLQ7zRzxhNoKmlF6_6nVDbaV2dYZV0,638 +torch/include/ATen/VmapGeneratedPlumbing.h,sha256=kqrn-vxR65jrXv843gLyzZ-lwknxNAv_EL6PLnm9upI,1784532 +torch/include/ATen/WrapDimUtils.h,sha256=ICeanLwAPsgdXpzErqbh3Fh8lE7str_veIpRSERuGwg,5109 +torch/include/ATen/WrapDimUtilsMulti.h,sha256=ktOVH3vp7JM5jKUy_bWisOamR-aomgs06AlRuxAg2IM,1329 +torch/include/ATen/accelerator/Graph.h,sha256=Vb1H2gPcEBbLcMVWdMbhGfIP4pqaynqwWxJICniosd0,2035 +torch/include/ATen/autocast_mode.h,sha256=ThHbFEK1oRrnocsfJ9-H-pshIciZeXHOnGmb-b5vK8M,41495 +torch/include/ATen/ceil_div.h,sha256=ijWlLD2OVoG2d0qgrxvh9wlxYiLthVaZfLssulHBvhM,751 +torch/include/ATen/code_template.h,sha256=Dx-x6o01aL6MmdiezcVb_Uc-qFcb2b2WOBZgcJwrC8c,7205 +torch/include/ATen/core/ATenGeneral.h,sha256=_O9V6rXClA-PYtHzn1XvolQJnfw-eIOk-HprG_O-On4,299 +torch/include/ATen/core/ATenOpList.h,sha256=bbA9vLBtbI6ExK7ZNoVXC_jIuNlASPvIZWBFEE-xU6k,500 +torch/include/ATen/core/ATen_fwd.h,sha256=yPzzTcmZAkRJGr9raXER0AURpGXjBMy8znRvjlYu3VQ,1278 +torch/include/ATen/core/ATen_pch.h,sha256=TuhCmPdxNxZkN97r2MbONRpWfezpI2y59auOAIjzfl4,5304 +torch/include/ATen/core/Array.h,sha256=kVpLnzc2-eUkkMLK4KglsvpF7a7dwDFi9VFh6YKMEWU,1388 +torch/include/ATen/core/Backtrace.h,sha256=y1M82xbaY2gwML--qDHL1JbuyseD6P5UX1w3kqkcf8Q,313 +torch/include/ATen/core/CachingHostAllocator.h,sha256=iB9oBTohWPOgAmCWm0U88jSwuAZZVQDY-MRk33rK66E,41530 +torch/include/ATen/core/CheckMemoryFormat.h,sha256=V9IvV2yN0x7ogLqw33ygmriPS-u7V91cGN9AhjqIVI8,1044 +torch/include/ATen/core/DeprecatedTypeProperties.h,sha256=dhzNJbU70lKzF-bktvSMcGVBI0JjNXhQmP6RfcVQkTw,4133 +torch/include/ATen/core/DeprecatedTypePropertiesRegistry.h,sha256=CAh7FOjK___7yRlcthPcfZJ-3PvF4LsmXsRv-YVXskw,1098 +torch/include/ATen/core/Dict.h,sha256=ep-GA0dVgLgqBeyc6bMKZpUiBV-t6gLiTdAdK5SZ6NE,13530 +torch/include/ATen/core/Dict_inl.h,sha256=FzQ-5rlxuxfEjID0s3trK3NAi2217-yUnBCHYnLRfrE,7716 +torch/include/ATen/core/DimVector.h,sha256=LVBWqI_2nHqX39SBwtuFY9hTYs1l1FZNq2UV0VMHx5k,532 +torch/include/ATen/core/Dimname.h,sha256=T_w3U9xt8ZwhDg2zhus7svNQdKzSthsJF82uVQdFmhU,1421 +torch/include/ATen/core/DistributionsHelper.h,sha256=UeSZW5eA-Cj0f7jRg0b3PBCt5-In_AbHMP9DseleR7o,9172 +torch/include/ATen/core/Formatting.h,sha256=yBqgBa7fjx1aLMk8vyKnDZd3F7fLFy2BUN3bwX6xh1I,947 +torch/include/ATen/core/Generator.h,sha256=wD7l-ASTdD4PZvvGcS7-31zpj3bAyO_C_Lj98s13kUg,6618 +torch/include/ATen/core/GeneratorForPrivateuseone.h,sha256=xRETR5JniORd3NzvcUNJnQygihErGRFB0417-wcOplw,1335 +torch/include/ATen/core/GraphImplInterface.h,sha256=-azhinQlE8Dk-13kzviLHQh5J7fZDNB2NCw-mJP_o5Y,4365 +torch/include/ATen/core/IListRef.h,sha256=bzt0QXkaVtN3tdLbfWF7-Cb0dAyl4XdlckN7U_AeXkQ,21284 +torch/include/ATen/core/IListRef_inl.h,sha256=yKTQ1K-WqJ6sXJ346fLeWY3KFyNtH2pUzvDW3mULyaQ,6460 +torch/include/ATen/core/LegacyTypeDispatch.h,sha256=KI_IJGWgHAAbcBcRQl9jvTleBMbkXSgPWaY9EReKEJI,5111 +torch/include/ATen/core/List.h,sha256=mdmo9faTv3jym9eSvvx95U-VGml4idh2VAXeVXaiA9w,16332 +torch/include/ATen/core/List_inl.h,sha256=O9HZbTr7CKntmYoMMxU77lkilYzg6hxl5h6zyxYeK7g,11006 +torch/include/ATen/core/MT19937RNGEngine.h,sha256=jfMpQi0pyWXxNWtRHKom8PcvdNuXsZh1Cee6-pT1dGc,6764 +torch/include/ATen/core/NamedTensor.h,sha256=drMtfa-rZFsy8xZ4rGJgQi3ivpVvclPyyV3xadVO324,5530 +torch/include/ATen/core/NestedIntSymNodeImpl.h,sha256=GE_nsax-X8ELeyu2ZYBA5gCOVEotOuRK9Z6RtPB3UPk,6214 +torch/include/ATen/core/PhiloxRNGEngine.h,sha256=YBxZwCRgyY1ci3yinbmwoFoQDFT88UmfSXYkomus-7I,7984 +torch/include/ATen/core/PythonFallbackKernel.h,sha256=H3uBekJXfQC--9X4ydMjPggK2DuWUEhpBErCOjewedw,1348 +torch/include/ATen/core/PythonOpRegistrationTrampoline.h,sha256=M4T2s3fROpzCOXJnft8tHdRGkp6g16-kC53DqNLopJw,860 +torch/include/ATen/core/QuantizerBase.h,sha256=wQAV8ECaiwW9UdEkOAKv-z_XpZGI5fpS1JdW686c_IA,2941 +torch/include/ATen/core/Range.h,sha256=zzBaUyO3CcuLvUga7QKP4H_gWeJ8wE7Ew2c1B3OedJw,672 +torch/include/ATen/core/Reduction.h,sha256=QAvMKcHeef39irdGfa4jxAmljpQe05HOxRSh3Lt4JSk,653 +torch/include/ATen/core/Scalar.h,sha256=NcGp-BNXkybOBVLuWjEFHe8G3UnAY4y2JV5BeAUaCzQ,283 +torch/include/ATen/core/ScalarType.h,sha256=MSQSYWx0IuOIFuzEBBKyPeVjYlR0xMDm72dGgs7VFqI,287 +torch/include/ATen/core/Tensor.h,sha256=YVN7gz5Ty4KrToW4GOivewFC_ISHApmYayNXwD-x_fA,2657 +torch/include/ATen/core/TensorAccessor.h,sha256=1YyEt1lMIJy4ttt48lBm3RhHOSRL27C9CSKZGA5fSwQ,2816 +torch/include/ATen/core/TensorBase.h,sha256=DrJ9IErVB5svyM1pBq3l-hV18dXAWkMkYFH9dKJWrfI,39959 +torch/include/ATen/core/TensorBody.h,sha256=eIvImuBOZuyCT2e7W3OdEMp4HhTR4dPQXrPGh0koVic,296703 +torch/include/ATen/core/TorchDispatchUtils.h,sha256=iQ3i4dTXgFFX7uEXkFnhJlN_Kmx2xvOkfEzZKZdZxMk,738 +torch/include/ATen/core/TransformationHelper.h,sha256=jv-KmUraeyjA8O2MZWTSMhYN54shG6SeSRRbxYYlkmA,7085 +torch/include/ATen/core/UndefinedTensorImpl.h,sha256=SRf1XtiwBIMCQhRqppRAfKU3prcLBdFyUrqnJnNKkFc,296 +torch/include/ATen/core/UnsafeFromTH.h,sha256=XnaiDLsTu8NRwGHbONwvtqcDupaTxr62L5F_2NqclUU,962 +torch/include/ATen/core/VariableHooksInterface.h,sha256=7B7OleaLaYjYDQblChSk02xpr8sc6E7j61_PsQ1atjg,3973 +torch/include/ATen/core/Variadic.h,sha256=SrRuKG67BcShAjwck5fso_vUVSrSoQJQaxBpNniHIEE,2635 +torch/include/ATen/core/alias_info.h,sha256=L1ASjM52he4SEZWb-cHB1Of-wnnF2PDb1SiR40iL2g4,4836 +torch/include/ATen/core/aten_interned_strings.h,sha256=iaviCL88CXrttikyDUT6zcTkwy0x2KbwanZt5TdOhT8,57670 +torch/include/ATen/core/blob.h,sha256=HzPJuyyhpAzgjQZNYNS_pYVZqes10TmvWp6sBe68z5c,5498 +torch/include/ATen/core/boxing/BoxedKernel.h,sha256=bXjLkqAiI_J4U61C5Vqq1hnhW1TcO7gaEzwdD0MJpCs,8555 +torch/include/ATen/core/boxing/BoxedKernel_impl.h,sha256=vtenHg6JkQ2_6WEjDM9K3fOly6MWdbiRbV9aKdx04ZA,3530 +torch/include/ATen/core/boxing/KernelFunction.h,sha256=YljsN0rP9AM0cRWkgciZ4B-MChv1VCfCPriKTFv3rC4,10544 +torch/include/ATen/core/boxing/KernelFunction_impl.h,sha256=qdnInxyKkHiF5M0Mz-eeVvqLS_GErsMlf8cXLw7Ei8A,14059 +torch/include/ATen/core/boxing/OperatorKernel.h,sha256=QHTWNctJNO9c1k45cElj6L_vWzgwR-_UeK62y1Emip0,945 +torch/include/ATen/core/boxing/impl/WrapFunctionIntoFunctor.h,sha256=qwWa2ySrY88KlQoDVMu9TBcHcz-hbbEZzqd_oqR70hk,1582 +torch/include/ATen/core/boxing/impl/WrapFunctionIntoRuntimeFunctor.h,sha256=jlmCXnBvefaVbKn_VVpku6ICZMeP0mW6O2CXbJmcxiE,1686 +torch/include/ATen/core/boxing/impl/boxing.h,sha256=pzo7YY_riut-_W1DyzJ6xcF_M4ECGajhsbVEjhfcfIQ,13841 +torch/include/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h,sha256=jBrh52cB0YWdRoIRxc6iqC-LPYRU0MLkYMAIxkpaYt0,32519 +torch/include/ATen/core/boxing/impl/test_helpers.h,sha256=nz9XM9HqNklaxmKZqqvhdhKlteXJ0suYExAhwA8zyfk,4655 +torch/include/ATen/core/builtin_function.h,sha256=hcvw2w83YPTxT7bfKkVdfdOKfpLVb6CEoFNdJgFjxWI,2331 +torch/include/ATen/core/class_type.h,sha256=WX2rG5KT-B-UHuw-YiJsOWitQQZSHs9-q4odcIqb8XY,14351 +torch/include/ATen/core/custom_class.h,sha256=g593lYHGtGMEi7wUutuC2qAom6hX0TmbyXJm05PxLsc,998 +torch/include/ATen/core/dispatch/CppSignature.h,sha256=GHUTygJ3aRssV9XcRTt28XyBKq_zaZ5aXLLsOWkKJ2g,2636 +torch/include/ATen/core/dispatch/DispatchKeyExtractor.h,sha256=ZJiK1NFO03FQZVWhT7cE_4-LHGCqSuXb8Mc-FY2lkL4,11211 +torch/include/ATen/core/dispatch/Dispatcher.h,sha256=pO4kCX8NIyHzWiVKuwbwEiPq7LqXSC7y8maG6P_v-mc,35046 +torch/include/ATen/core/dispatch/ObservedOperators.h,sha256=eKNoomteAbCsaQ7Qofff0buj44X4yGyrAOlAd794UFU,583 +torch/include/ATen/core/dispatch/OperatorEntry.h,sha256=B6FWoiRbFz6KE5g4Ky4mPdZvVgL2j4FKEJVx_YQy-pU,13434 +torch/include/ATen/core/dispatch/OperatorOptions.h,sha256=Ibh7cTueyJ3FNcDS86MMprRq3uNlZO19f2yURfIu-gc,1140 +torch/include/ATen/core/dispatch/RegistrationHandleRAII.h,sha256=5sWLlO1mx9tV__IwcZcUqftnxvjJ2_CA8pmG-r6WTkk,1131 +torch/include/ATen/core/dynamic_type.h,sha256=6a7Tkj4OVKepnB8ixw5e1Pl-MZ6BZZGPPlRaFtHhZ0o,11155 +torch/include/ATen/core/enum_tag.h,sha256=DJDKXIXYyPTMVxuwHaAz_asYcI4bMuOtzKOr0O43zUw,362 +torch/include/ATen/core/enum_type.h,sha256=2si4vrOCltW7rJjlxYiuODIBy5CYgXBved82o05vhLk,3183 +torch/include/ATen/core/function.h,sha256=MF6WVKSjPmoBUVLtZbQWYfHAyxNTxo_j6iboqbzcqQE,3764 +torch/include/ATen/core/function_schema.h,sha256=6dJ9IyguRRV9NN0-TTflGizmvnKu8eby90c7V9v8UQU,24368 +torch/include/ATen/core/function_schema_inl.h,sha256=pgUnnBp9VgVUFWUtZKlFg3eLFV1mz34d5E1WyKibNIQ,2325 +torch/include/ATen/core/functional.h,sha256=2YeMi9lGkXDPJTQtr9AtxErPP72ioMMLf0v5jj0TqRQ,1718 +torch/include/ATen/core/grad_mode.h,sha256=v4cKv2y2EFgAqEH2yI_D1PdoPs5_neoLiddCCwrvqrM,464 +torch/include/ATen/core/interned_strings.h,sha256=ad0z7aOSN3HBijiIywa4g9fopPMaXg-mtcbUHnv914A,13623 +torch/include/ATen/core/interned_strings_class.h,sha256=rxTLcElTBOppL-iUUQitRcC-RebaQp3Ch0--85iS9Z8,976 +torch/include/ATen/core/ivalue.h,sha256=suZdJ9UILnO2BkWjWxADsufZmoQTbCHiR1l9IW9WhW8,53213 +torch/include/ATen/core/ivalue_inl.h,sha256=8ZbpZrHlbi2sDamV_Ij5YKkK5qVUtU_TLQfSiElXvQ0,88368 +torch/include/ATen/core/ivalue_to.h,sha256=OPCshnu8WClCKdXRzuNPhnwi9dtcH7HVB7NZQf2_cLo,1010 +torch/include/ATen/core/jit_type.h,sha256=a9zWqtZfMgK845jeedkqrUqDie5rkzjTZyzBnUcySSg,72426 +torch/include/ATen/core/jit_type_base.h,sha256=_ZbemJj0czwJ2Rzmo8aDnL9oceAmsp5q0QHMfYw4hiQ,19409 +torch/include/ATen/core/op_registration/adaption.h,sha256=kRQiFmEFwArdI_l_F7MpAg7AzdWGanMvunJZe4hcfqQ,3470 +torch/include/ATen/core/op_registration/infer_schema.h,sha256=5IaNPUJuTtzbv-Gl591VPWrDZuaRw_IkRcsltPIY3Aw,6992 +torch/include/ATen/core/op_registration/op_allowlist.h,sha256=BKCMs9ZStEtcOJ6dlRxJWnzM6DKyknZ0l-9rgjOdh-o,6559 +torch/include/ATen/core/op_registration/op_registration.h,sha256=-4NqKDvpmQiGuNZdFFRh60u9FhvVurjVacHvH-27Z6Y,28800 +torch/include/ATen/core/operator_name.h,sha256=9R8G2y2OpcQwf0ndN7VPWfe9E34pxB1l0XChElvkWCM,3315 +torch/include/ATen/core/qualified_name.h,sha256=wMsJ1QY3tiLOSe04oAx-vfELuZXraI5Lpy6UQUEnXJQ,4627 +torch/include/ATen/core/rref_interface.h,sha256=mhAcftY4AYad0H3q2vuRqqZBOKu_m0yqzC-z33yl3cQ,1462 +torch/include/ATen/core/stack.h,sha256=WVUVRJwhYl8mPcNwXVxd410k53ZA-mzlOy-hNQbLBuc,6436 +torch/include/ATen/core/symbol.h,sha256=6skSgLomwIw_70vw7ikZdNWQDBvQk9zmcdz-_RBxHVo,6145 +torch/include/ATen/core/type_factory.h,sha256=R9o5Y2pdvrg3gOP-2sv82vEqqDPiMMyp1B7smfLir3Q,3501 +torch/include/ATen/core/type_ptr.h,sha256=Mby1oP4rzL07obzNnL0DD5reEEnPfBsJJI3hto8NlzQ,1483 +torch/include/ATen/core/typeid.h,sha256=eMggcfZfdVXFBl32XD_y3SL-xntu8-3vsW2gaorUV0Q,283 +torch/include/ATen/cpp_custom_type_hack.h,sha256=uRYZgQGKQRkEh6F9SnGGWDNyBU-oxzlP2vo5wQEU4io,5684 +torch/include/ATen/cpu/FlushDenormal.h,sha256=9gYBrIZWnFSzBKy4SpTo7QPFzSHi5WSkGNrBbzLjMjM,791 +torch/include/ATen/cpu/Utils.h,sha256=9fzOipawtA_FNN5Ejy9sqTzawnJM29NCKo-qBL1_Fhg,922 +torch/include/ATen/cpu/vec/functional.h,sha256=AWcimK14Ju73ML7A0rEXLtY4C9kH2qDeN0jLCdVpklg,356 +torch/include/ATen/cpu/vec/functional_base.h,sha256=6k0KTdb6veYbBaGMjEqXa9l65Xev3uEnQCyQl_VpzqY,15739 +torch/include/ATen/cpu/vec/functional_bfloat16.h,sha256=bUDJvUshdQsBQL3q-TzY9FXE8EZrRQHJHmIYtCB0fvk,25665 +torch/include/ATen/cpu/vec/intrinsics.h,sha256=etCWzMEatbbky2AF29OWWN1owMc00LGkw50bh3CEBOQ,303 +torch/include/ATen/cpu/vec/sve/sve_helper.h,sha256=93OjzvSBUpLvypKeiO9ZBadwEpkA1512aLJRLODntf0,3429 +torch/include/ATen/cpu/vec/sve/vec_bfloat16.h,sha256=SIehL_6yLmOsnoq5CcTEBp-3Ir-cKCFRhwS6fFuu4Y4,20827 +torch/include/ATen/cpu/vec/sve/vec_common_sve.h,sha256=lNA_x6FVIXPKkEXj2uiJ4z5l-By6SOupeT1XFNMISzc,8687 +torch/include/ATen/cpu/vec/sve/vec_double.h,sha256=EuC1-cuLGBjnhB8n8Ig_x5rAggD1_eKDoWzokXUrC2g,19956 +torch/include/ATen/cpu/vec/sve/vec_float.h,sha256=ysHvUnlHhaS_pwp7c4g5opGRsqdl-hmErEaki7Uedbs,27621 +torch/include/ATen/cpu/vec/sve/vec_int.h,sha256=5P8ODDVF5qMQrs1fwwo7kwoZONDznb9qSPwHpC2haV8,28230 +torch/include/ATen/cpu/vec/sve/vec_qint.h,sha256=xYLhGKQMz2TZl-9NJYcKlRYE6nkRMQ_LWdd23y8MXco,20303 +torch/include/ATen/cpu/vec/vec.h,sha256=GagNXvVOqjRWeA6kvDUFj8C7U4mEpGvD0GTsPGzaGeo,1603 +torch/include/ATen/cpu/vec/vec128/vec128.h,sha256=nXYcpEQuQW3pwtUxYIS6RUGeXFZSuG0pJI2L_TRy5Ns,779 +torch/include/ATen/cpu/vec/vec128/vec128_bfloat16_neon.h,sha256=N--NSD9XRqRBt_cI81hqxznxLIhe-kdx90kBQR6WzSs,25989 +torch/include/ATen/cpu/vec/vec128/vec128_convert.h,sha256=jOLFwVl4yw_QNBFMkhMc5aSceBw7jYnXmRgQ-QWe_Ew,12870 +torch/include/ATen/cpu/vec/vec128/vec128_double_neon.h,sha256=yKkClDbtj0BtocR6Jd68h2Qa8_53APHIKaq7nA4lNJI,18325 +torch/include/ATen/cpu/vec/vec128/vec128_float_neon.h,sha256=oVmFBTIKFVuSVy8JKystyCuRu46UGnXj0Uek2wC0Uhg,23228 +torch/include/ATen/cpu/vec/vec128/vec128_half_neon.h,sha256=hRsr4yxvVIPRjYS6FdJFXagrBS8dw6mCBBwGzEF-UQ4,20633 +torch/include/ATen/cpu/vec/vec128/vec128_int_aarch64.h,sha256=TDUVGA-OI12ERVDI7n5lCYTt8mChYbPLTqANk32Wmnw,32389 +torch/include/ATen/cpu/vec/vec128/vec128_reduced_precision_common_neon.h,sha256=VTyCdjQtf17cd5yUhOZSd44oTL32EtajY8mq3zBZIx8,10176 +torch/include/ATen/cpu/vec/vec128/vec128_uint_aarch64.h,sha256=51VwJSzDxBGbsMnEy9mRvADkMYEweNK7dwEoalfndhs,21095 +torch/include/ATen/cpu/vec/vec256/vec256.h,sha256=2LpUHPcDBYMV-cjB8qd7y0mkVTzsJAVsv-h1EljZgZ0,12974 +torch/include/ATen/cpu/vec/vec256/vec256_16bit_float.h,sha256=qi_RnATEaWQGHOqDldKSYTgn53Vyxuj5XRBdXNK8PmM,28702 +torch/include/ATen/cpu/vec/vec256/vec256_bfloat16.h,sha256=ozcJO1Jp3RNJ8Vo97IxhE_GvnhFxSmabYnMTmmf3IkI,9061 +torch/include/ATen/cpu/vec/vec256/vec256_complex_double.h,sha256=lceQ0Eo7PHIHq5VP4Pi6XHr6Uv5oxJ2iegF54Pez7sE,19830 +torch/include/ATen/cpu/vec/vec256/vec256_complex_float.h,sha256=J8nG_j0sWMUEX_xsXooWVobFfoVVVvvsLUOtwkEilNM,21867 +torch/include/ATen/cpu/vec/vec256/vec256_convert.h,sha256=sKvyirrkm4QhenxzQZYEatuzIOp1RIomJ4h8VSwo0ho,11287 +torch/include/ATen/cpu/vec/vec256/vec256_double.h,sha256=ocwFIIOxGpuPej65VNSY4fCfc0XYu-66VmPCWU5f5Zw,15832 +torch/include/ATen/cpu/vec/vec256/vec256_float.h,sha256=Ys1pFXJnAbe2LeoAwWQG6xeSwUDBl64aEuTTW5xtgd0,27555 +torch/include/ATen/cpu/vec/vec256/vec256_half.h,sha256=MKSjE6wiMvHRxr_Ci_prSCJ2eeonRFPS5IK8WNUs7Jg,8643 +torch/include/ATen/cpu/vec/vec256/vec256_int.h,sha256=fUGI0xrSGpC59pUgZh_0OMyiztXGYemk0yFDb4ao8uI,65745 +torch/include/ATen/cpu/vec/vec256/vec256_mask.h,sha256=W4mmnWT4Eh7nI6t1JziXE3-3SP9XKEpAg5CbBlUN3B8,9026 +torch/include/ATen/cpu/vec/vec256/vec256_qint.h,sha256=iB2J4szv4AKpkq5ANwVI9jB1UCCBSiruJ_i4Fw4cfaM,51098 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_bfloat16_vsx.h,sha256=PvlK-Gb3-0MFQylPhaG-6ho72E0wN_-1gzOwbT6Xu1s,2391 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_common_vsx.h,sha256=H-2vPZiGmNo-0PJmRrqOpMU4_oCzthYuwAgCGwxJThk,8395 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_complex_double_vsx.h,sha256=4-QuTuE6f7edBTD16E_GWAQaLaVHwq5J7Jqf8oH7I-E,21936 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_complex_float_vsx.h,sha256=a7NIslKk8zsXM8MFr5D8CdqjtWFSXdCec03-HjNTKf0,24625 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_double_vsx.h,sha256=zoZrV8is_U-4pghjBqfAsa07gtjZvG7Z_EETZrXFAvg,16539 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_float_vsx.h,sha256=M9xzffcMLDeGTFEWBLJikAXIjzZPO72gZtuh0T3TdEA,17089 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_int16_vsx.h,sha256=1jy4dhqJwivpN21EvwOCnVdUnz_X_eCXB-dwj20PS1Y,13874 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_int32_vsx.h,sha256=k3kEBDt6jcCRFPR9yBmSJuSaikUJj6_QMiFI0EllZo0,11689 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_int64_vsx.h,sha256=YDixVuNdXttFOUxJGqG2RbRar_BlKhcr0MKx30DH7Fg,9958 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_mask_vsx.h,sha256=CQJIsE6joHaFwbtDr3ePX85j4NSj8U_SIGOmEFHZ4X0,1980 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_qint32_vsx.h,sha256=WDOkNhlzeahJ-jOpPgILGUeyo1h8K-wCflugVGJiV8w,10319 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_qint8_vsx.h,sha256=8aqlZqkHgBMOjTgaheDZfsd9VZL-sjXYn0GsFmMqD-c,17672 +torch/include/ATen/cpu/vec/vec256/vsx/vec256_quint8_vsx.h,sha256=-13l08O0IoVfrjUeL_pyDScyCWiM6a_TiyaaHFLBglI,18369 +torch/include/ATen/cpu/vec/vec256/vsx/vsx_helpers.h,sha256=uoMPqd68K1_R4jlS6vPKq245iVj-Kw_BkfQBlqdeVDs,22756 +torch/include/ATen/cpu/vec/vec256/zarch/vec256_zarch.h,sha256=41RaDQLMscyGLfoVEdNLjvk2hXcaGEP556QpStSieU8,102279 +torch/include/ATen/cpu/vec/vec512/vec512.h,sha256=6aEoKS1gqdUSyrHAdIm6VqRuUJgtvfpfHnFAxkI9N6M,11390 +torch/include/ATen/cpu/vec/vec512/vec512_bfloat16.h,sha256=RsSO9F7SP913aHUJunUvLkmDBFtF8MBlCF0bK_p9S9k,66745 +torch/include/ATen/cpu/vec/vec512/vec512_complex_double.h,sha256=DJW_DgozjXVDRkUWLV56NAn6CFIvWwtyU1HnYnfC2x4,24307 +torch/include/ATen/cpu/vec/vec512/vec512_complex_float.h,sha256=MzE3jbTqotDvaIQLAcTa2FqfO39BzhglEPZn7uaz6Pg,44800 +torch/include/ATen/cpu/vec/vec512/vec512_convert.h,sha256=r2BxXhwG0dzxGl7NXklUh9wU8pK2-ORQdJSAkUqzMaM,12177 +torch/include/ATen/cpu/vec/vec512/vec512_double.h,sha256=Lyka1WQbrdC94P6zivLmDnpdnKRAPeSwK0e-MLccnPk,17362 +torch/include/ATen/cpu/vec/vec512/vec512_float.h,sha256=AVDn7rsHpSNfGc5cdZNorDw7-T9l_0WWMqMLz5nViWw,30372 +torch/include/ATen/cpu/vec/vec512/vec512_float8.h,sha256=V_vCGQpIL8_83i3yZYjau6d-oEzs2ooaqqPgjQAeq1c,24814 +torch/include/ATen/cpu/vec/vec512/vec512_int.h,sha256=yDzLgRxx9Ej3m-Ovt_R8tU4bELhVPEFNFIfl5XEehn8,59208 +torch/include/ATen/cpu/vec/vec512/vec512_mask.h,sha256=KC5eBIuuC14cSChSVe41Z6FtFqqRn8OVCeG0nMY-9y0,12552 +torch/include/ATen/cpu/vec/vec512/vec512_qint.h,sha256=_TCyuFUzKfY56MvLwe2sf_MgJBWbEyJ2AlnpmGK2vcg,52929 +torch/include/ATen/cpu/vec/vec_base.h,sha256=pBudEEQ6Ez4PVIEKRSXWU9sRemHhUVOVAElRt2tv6B4,49547 +torch/include/ATen/cpu/vec/vec_convert.h,sha256=SBBa53Pjg7qbNnDKdI5D_qpxPhuzUL7x9BWGZxJR7MY,4704 +torch/include/ATen/cpu/vec/vec_half.h,sha256=PW1nLS9MdKp_smD5_xNW7hcy0gTc9nHIwG9TrtWGT88,3748 +torch/include/ATen/cpu/vec/vec_mask.h,sha256=Tyn3J6n-w2V7KTZd9R18j_2dEl8b9f3HaHcjqAdvLBs,10584 +torch/include/ATen/cpu/vec/vec_n.h,sha256=65PbUFbRex448sk8dXUH2MvQ6SZDIZClzjLEH_FPJ0k,13810 +torch/include/ATen/cpu/vec/vec_quant.h,sha256=b4rBHPDN4KwkTkoM0LtoiiLDhtdF8dvY9B19YwXu6uk,8891 +torch/include/ATen/cpu/vml.h,sha256=_biUSzcgtpOx-atj0MS8QOCFlQxHZDU2hrg7GQWgT0w,6327 +torch/include/ATen/cuda/ATenCUDAGeneral.h,sha256=Y3g0kFj0PlB9p2jUAA4XS22IQ3_R6tSGs6FEiTUBbRk,444 +torch/include/ATen/cuda/ApplyGridUtils.cuh,sha256=Y-V0zrTY1oL4SOLJzAHxz7wDiZ1hT_6H3mdi53c21Ss,1563 +torch/include/ATen/cuda/AsmUtils.cuh,sha256=hmNSZxEXEFJmvOHvSwJ9sekEhb73IB2gLhIAZjqxq20,3648 +torch/include/ATen/cuda/Atomic.cuh,sha256=gxnFD8uCEokJbD88OgqGa3NGSB7VQIx3H17BRbRzajo,27057 +torch/include/ATen/cuda/CUDAApplyUtils.cuh,sha256=2FO8fHZKLVe3V-IXmHU1pudlu8VpcExKOcAtzC5QSbo,20597 +torch/include/ATen/cuda/CUDABlas.h,sha256=T-hxIvXGWAYg7Jx5d47Xbhc66MNqsUEpxD8MXrqMWrU,15513 +torch/include/ATen/cuda/CUDAConfig.h,sha256=UsHlR5o8JmAWZX7OCdDYXPjcQhYzgvMOHk8Ykk4nKV0,1076 +torch/include/ATen/cuda/CUDAContext.h,sha256=FVxhveWjVQ87qecv__fmXrFj1HVowBTGI_uBUUhEQ9A,492 +torch/include/ATen/cuda/CUDAContextLight.h,sha256=p4peSa7dOenwaBAGoGT0cDzGi-ByXlAiuKVuffBMhG4,3504 +torch/include/ATen/cuda/CUDADataType.h,sha256=9LRtoarR6hBpMhKhYnd2oCgMX33oTgVG2Cno6gPAvag,3089 +torch/include/ATen/cuda/CUDADevice.h,sha256=ibOUWACsMvCwJkj8Vs-JaSAacLjfePU25ikYXuE0h6s,786 +torch/include/ATen/cuda/CUDAEvent.h,sha256=1n0unf9uzX68TriMIOlSv3rtmMDaCr0khuMi3pInToA,440 +torch/include/ATen/cuda/CUDAGeneratorImpl.h,sha256=wStSUGQriB7R8gZH2dgLDv6CLfr4UdQkVhXVx4mZyw8,7183 +torch/include/ATen/cuda/CUDAGraph.h,sha256=gc4vhbNyA4_K2rls55VYkaLwfwriI1I7dBuJ_Wd-jZ0,6636 +torch/include/ATen/cuda/CUDAGraphsUtils.cuh,sha256=Wahj6otCRgDBglGPuVQrqwy5iirfotsxHRcmSZaj1zo,2310 +torch/include/ATen/cuda/CUDAGreenContext.h,sha256=GGkSSC4Z8LSIXslmnwO46rV0TGYa2t9N7JMVSudgchI,1977 +torch/include/ATen/cuda/CUDASparse.h,sha256=sRn70ExXg9HNm9pSH9f-Xm2u6jxbBZFGWyyyEdHJMTQ,1035 +torch/include/ATen/cuda/CUDASparseBlas.h,sha256=NW9_yX3WvUtCA5A-d5OZrK3Z_d5hlCRXNYncjhG7TkE,12945 +torch/include/ATen/cuda/CUDASparseDescriptors.h,sha256=6GqtUUckWOZojdaAkoqon3QLkifGy1cfZQmz9iOdBzE,8034 +torch/include/ATen/cuda/CUDATensorMethods.cuh,sha256=0XUAaYRzKUbJnw8-QsCk9huRWC_eBNqPI5rgw8m2Ecw,524 +torch/include/ATen/cuda/CUDAUtils.h,sha256=aqOrT9ciEgvqxTs6AIJ1g8AUYpuEkYDGbRDFrJHk3CY,670 +torch/include/ATen/cuda/CachingHostAllocator.h,sha256=Cubx1tn7320HVu04x0uI7uJP-UkxVma6N9Q8--xnqSM,3284 +torch/include/ATen/cuda/DeviceUtils.cuh,sha256=ckhE-i-e55aZrBmh6A0KMa0KJf_d4yan1AZ-YZCwlF8,3510 +torch/include/ATen/cuda/EmptyTensor.h,sha256=ic1E7hVwYA_iNMXuG8cFUN0j213xDHgv1wN9gsaXUFA,1460 +torch/include/ATen/cuda/Exceptions.h,sha256=gFwW4-ehCXlB0lLW6Xtv9XVTdEHG__alKBmWmJSdiIY,12842 +torch/include/ATen/cuda/MemPool.h,sha256=wGz7JyS4bl1bObFNXCwhLL7rFhNCSKXQ7p7AZaJNZtc,1539 +torch/include/ATen/cuda/NumericLimits.cuh,sha256=UED-0HPXq3bA_mqroXCxFCrUBW7n4aLTiZ2xjrsJ15w,1597 +torch/include/ATen/cuda/PeerToPeerAccess.h,sha256=3-oEtfUZnYnanTvgS1xcb96lMZzWdxgB9cLvRRehZQc,1515 +torch/include/ATen/cuda/PhiloxCudaState.h,sha256=x1KYpK2i45uMEWt3uidetR8Vyp-0vsGkcpMQdqEE_2k,339 +torch/include/ATen/cuda/PhiloxUtils.cuh,sha256=42vYVkx5qDsakAMXakmFM3zEXgzxAooYG9j59NT_dJ8,349 +torch/include/ATen/cuda/PinnedMemoryAllocator.h,sha256=ZYJfIg5C4V3bISgVLBvk5Go_xeEcOuiu84XyZLyY2lI,477 +torch/include/ATen/cuda/ScanUtils.cuh,sha256=2z2Hqz5wfzxSwSpbzDfeuT9F4Zji2Ju3RtyCwNiggO8,2281 +torch/include/ATen/cuda/Sleep.h,sha256=GFH-xrg91qwm84BnIyb8V4-rwzBLAiTjbwaLsU3N5ws,573 +torch/include/ATen/cuda/StatelessPhilox4x32.cuh,sha256=y4xOVqat8H5jjWfUva1JpWh40i4QUnLL5Lz9yp86Y8Q,2359 +torch/include/ATen/cuda/ThrustAllocator.h,sha256=vuoo-4-S71Xy4C5lfrVttB6UChIlYxlr3Wi6inoYuMo,759 +torch/include/ATen/cuda/cub-RadixSortPairs.cuh,sha256=7ZRAJC-Wvk63eaa7XLVl5qsiOMWWrdi909yoyvWWMdY,2459 +torch/include/ATen/cuda/cub.cuh,sha256=q7KqvMMSDDJbYqsyeNMY0BidRKL6hM7uf9G7-FWT95w,21494 +torch/include/ATen/cuda/cub.h,sha256=D0lN7cr8YsLXyx83xbz7CYBCwmo1J0tcKV1dSTo_JyA,4016 +torch/include/ATen/cuda/cub_definitions.cuh,sha256=2Rg0L2a6O8__ZQc50KeC4F_5ww411XIRGusCtbRdMDE,1073 +torch/include/ATen/cuda/detail/BLASConstants.h,sha256=goHy09UIQi7XLieR3E3tfjcpzgpwxSkIhpjtDhMlUHc,460 +torch/include/ATen/cuda/detail/CUDAHooks.h,sha256=mdCiITTzHnJ8-mxedSglmjUXDj6rQ215g6GwXEYefmA,3343 +torch/include/ATen/cuda/detail/DeviceThreadHandles.h,sha256=yW1GbtigusLf7FP5EP8wUlm_dSmlyr6I3N0-m9tL7uc,9066 +torch/include/ATen/cuda/detail/IndexUtils.cuh,sha256=S-kBfp6T_yGnnfsFv-H78_iK5pHPBLi7ZLyv8XiFOII,1983 +torch/include/ATen/cuda/detail/IntegerDivider.cuh,sha256=qORRBvLTJxI54_r7S1_Yn7HftysizNWtKKSRysEHzVs,4273 +torch/include/ATen/cuda/detail/KernelUtils.h,sha256=I7_ldw6KRcc2hGcAv3Jqz1vARlHZbTMVfOdJMJla88w,1789 +torch/include/ATen/cuda/detail/LazyNVRTC.h,sha256=YvqkMO6N5dTfYQaSB2B5iv7kZaCV-dfeD4uPTu63NVg,474 +torch/include/ATen/cuda/detail/OffsetCalculator.cuh,sha256=Q8O-0N9eBPTnSDrmMQ7xvFdMvyKeB_Ig7vrigO3dux4,5112 +torch/include/ATen/cuda/detail/PhiloxCudaStateRaw.cuh,sha256=1bwPiGKT62D7Vs4llEQl4WTT75aWa3NfgNY7HBRyNvU,1612 +torch/include/ATen/cuda/detail/TensorInfo.cuh,sha256=KJwJRq_CHfJHX7RzpQQZshnm27T86F1k9Ry-QEOUz20,3494 +torch/include/ATen/cuda/detail/UnpackRaw.cuh,sha256=j5i7r21kCrg6SKQf_BBP_uDrii4xIup4IX4Z2pKLXno,1993 +torch/include/ATen/cuda/jiterator.h,sha256=gO2gKosxQYi4aLSupXH99e-TkBKSqifzCQNLdI1ltA8,1214 +torch/include/ATen/cuda/jiterator_impl.h,sha256=O9kiGZeeCGzohlucr92WunWpjzUZavj81oYc47IzOaY,7355 +torch/include/ATen/cuda/llvm_jit_strings.h,sha256=kJpasSteT1vKc5ujFRyDGa7whcIKnLwaOxyxtRKtkRw,682 +torch/include/ATen/cuda/tunable/GemmCommon.h,sha256=wqYi5EIY7r6D4b_dgV-9dvmYodImZk4ATJ9A1VLMoNk,23633 +torch/include/ATen/cuda/tunable/GemmHipblaslt.h,sha256=QpViCtbMOtqHpThPXLJDxq2teljj56Vv9nHR2eJc_Nw,21988 +torch/include/ATen/cuda/tunable/GemmRocblas.h,sha256=_oyMmgdqH6HSfoLisS13F30uh5GAE-HIziKemA0irh4,10804 +torch/include/ATen/cuda/tunable/StreamTimer.h,sha256=g5V6mzx3oFvoFkYiRatDxEFGjrdSGoo2B0yJakz9IxM,1299 +torch/include/ATen/cuda/tunable/Tunable.h,sha256=W9NutBaBfl36r1ShX9bxffOU52RK464cxAsNnfQrv7Y,8587 +torch/include/ATen/cuda/tunable/TunableGemm.h,sha256=7PhUhYJzrVTkVHfeNSQ9p-6qSGPDmD2yIuKtFSVRlcs,9886 +torch/include/ATen/cuda/tunable/TunableOp.h,sha256=hg5dLMfsD1wPpbcRLXAo_iwL7v11VXcdLM3ClVJv80c,14988 +torch/include/ATen/cudnn/Descriptors.h,sha256=DmWD2pCzqLqxsfPCGmc4rDK6sCgmFkZPKaBcQDYEaVw,15335 +torch/include/ATen/cudnn/Handle.h,sha256=pGkiq0ZzcZLgbiGQ8ybHUvnVZmoPyLhjgsha4JNSkYY,447 +torch/include/ATen/cudnn/Handles.h,sha256=jmhakLkgV8AdF07lLAQbwsmcEeRRDi8fNcYhtYSI77g,298 +torch/include/ATen/cudnn/Types.h,sha256=HfZhdzuBc-DZ8L0ketUnMViu39Ct7jTkhhglPgzGQ2M,564 +torch/include/ATen/cudnn/Utils.h,sha256=y3Hdkk0UfdPIl8rj9NRpz-5jq11m8NJy8wK16GeHFn0,849 +torch/include/ATen/cudnn/cudnn-wrapper.h,sha256=QTBGaaa7_h2SB3HmG4U8LjN4AOx7393HWo70zI-TfGY,794 +torch/include/ATen/detail/AcceleratorHooksInterface.h,sha256=9bytpnF9OcNzUfqXbIjV66b9BQ_S25C0ByyOE6Eychc,3248 +torch/include/ATen/detail/CUDAHooksInterface.h,sha256=IBHZitHjaYSD4WYJCAo0lKtztz06K_kxBOlxkorR7qw,8976 +torch/include/ATen/detail/FunctionTraits.h,sha256=s_ett-iw2Qflvg9_x-aJL8VRXZS5ck2t5o-U7FZPg88,3316 +torch/include/ATen/detail/HIPHooksInterface.h,sha256=r_g_MvwN_IUeRWn4LjHMJ_xjFEfH-088fmkSNYHSyoM,2258 +torch/include/ATen/detail/HPUHooksInterface.h,sha256=B5GS9btNsCRr17mytiHpK5sANHxZ61T9lFIUurboGps,1708 +torch/include/ATen/detail/IPUHooksInterface.h,sha256=7Z6JZQ202rhDoCQebYV8J6kdFj_81N8eCn_XpCdN_I8,1498 +torch/include/ATen/detail/MAIAHooksInterface.h,sha256=WRf1-O8_Qkw-yKRbwzXRx473kCTjxofmWMQALzH_RtI,1529 +torch/include/ATen/detail/MPSHooksInterface.h,sha256=jXs5PiR5OxrRMHGd2r_8KQjNvU-b4TfasAH_4zhv4xs,3994 +torch/include/ATen/detail/MTIAHooksInterface.h,sha256=X3n25nhMJCINp694ET0KdaYmVphIvNmgW1XwSF8u_z8,5970 +torch/include/ATen/detail/PrivateUse1HooksInterface.h,sha256=OSy0j1RhVNngvxjGtO86N6lnMnQxQt4YnQwjdr5EOGU,2743 +torch/include/ATen/detail/XLAHooksInterface.h,sha256=4QzNbPhXB0BpjgCCgZw5zBsLrf8MPMv9CPZ1lTu9lwA,2677 +torch/include/ATen/detail/XPUHooksInterface.h,sha256=7viqhQ830QSZwoltxKoW1eJ8NxwbzGiAHWpxKSoQyEA,2924 +torch/include/ATen/div_rtn.h,sha256=4nvTpsu7JOVXbz2gkH7EqdYfREUw8zg3yjJLEo08IDc,465 +torch/include/ATen/dlpack.h,sha256=YtCs_UQXAPm5cO8taUvehijM1KNrC10xFZJ2g1D2jXM,23516 +torch/include/ATen/functorch/ADInterpreters.h,sha256=SimV299hZtyUMJekW7ytVKpvPW686YhuLctUjvG5Zro,1814 +torch/include/ATen/functorch/BatchRulesHelper.h,sha256=l553Ne5NOm--LGw3lBRvqTwF-cEIFsnkH16dsFcYOQw,19053 +torch/include/ATen/functorch/BatchedFallback.h,sha256=0ECOodFmirkxr_ztE34K_EToFxZ75h3pltUjnXawP1Q,3693 +torch/include/ATen/functorch/BatchedTensorImpl.h,sha256=G2e-GWothDMfdvbafrTKzTNoKE33oqxyioqi6va4G-0,6783 +torch/include/ATen/functorch/BatchingMetaprogramming.h,sha256=R8l4RZdmjefQfMPNMG1jvo3pB7RK0UVTw6AXR53ln00,5244 +torch/include/ATen/functorch/DynamicLayer.h,sha256=NMxlke8eCf9T2KSNdWT_fbifb884srNDKk6-bPdHupQ,5938 +torch/include/ATen/functorch/FunctionalizeInterpreter.h,sha256=3jU-mTn8FDBqjyd1gUfW0dAE8ACSPFDx2dAbbgArX0c,1161 +torch/include/ATen/functorch/Interpreter.h,sha256=8o8SL3qrcxjdu7zwoTfowAj-iseudhvTufnjNkj-3dg,14967 +torch/include/ATen/functorch/LegacyVmapTransforms.h,sha256=nHwpaVy-psnxBrDnAbcxedCc2iWjQ75QjUWVosg-ZWw,8505 +torch/include/ATen/functorch/Macros.h,sha256=hqjfrYvy4mqhY3MXdsNKQUXzb2y8LrkdWgColdke2bE,304 +torch/include/ATen/functorch/PlumbingHelper.h,sha256=qRDYK1BuwfrE1vr3QWGKEhSPeosHrq0lZqKOlirt_TQ,3108 +torch/include/ATen/functorch/TensorWrapper.h,sha256=kghomRRLftFFipY5vp6zZpCYAA6k1SyJOpHoXAYPYzY,4279 +torch/include/ATen/functorch/VmapInterpreter.h,sha256=egzDZVrqImPBju3_cVquUj4xKR6OSeNjKd5Gnt7dP6A,1211 +torch/include/ATen/hip/impl/HIPAllocatorMasqueradingAsCUDA.h,sha256=P71dPrMT71EerhduHcHb83JY3FYWlyyZ8Q6AYNAFDFU,7616 +torch/include/ATen/hip/impl/HIPCachingAllocatorMasqueradingAsCUDA.h,sha256=uTzJYZqczGk1L9rW-WgWYjaaIXWrXi5RdRwaRQpg8jY,2809 +torch/include/ATen/hip/impl/HIPGuardImplMasqueradingAsCUDA.h,sha256=CvUM0wbuQFlB43HQOmsYL0DmOxH75-1Hz8j8ksIjrRo,4072 +torch/include/ATen/hip/impl/HIPStreamMasqueradingAsCUDA.h,sha256=fFzzF1C-DC5hyT8e5T0kQITqNfiEfPhPxxj-0flXP2M,1577 +torch/include/ATen/jit_macros.h,sha256=M6xEVYL7q98hfLdEBiTuCjtgR-pkGmclxU8L09Qntik,484 +torch/include/ATen/jiterator_macros.h,sha256=g-vRy4Cy6UmhGVXErSsGIM_vcAtxu40vmcLA_6xPF2I,1760 +torch/include/ATen/metal/Context.h,sha256=3Sf57yULLhzLS_JYMrGeIaC99E4_0V9-G4VzeV6YjH8,936 +torch/include/ATen/miopen/Descriptors.h,sha256=Rm4ItrthpyKpBEhdu1Tzx533qgZanIF2zBw4ffYYZYI,8238 +torch/include/ATen/miopen/Exceptions.h,sha256=Yicb0z4mTJD6FSfGhKUXELkeTYOXByxGC8fJX7NPCN8,1350 +torch/include/ATen/miopen/Handle.h,sha256=N79o8e6G7CH_Lno1cT_hWY6K_P6DhZwH3h0E-jF8-uw,443 +torch/include/ATen/miopen/Types.h,sha256=sP-Ig7kjm3eN6W_My-JSBEE1yRYpaOHwctYa7m-PMbQ,524 +torch/include/ATen/miopen/Utils.h,sha256=U17pahwfbSBWMSUnA2d--b1SeHbtRjrGTSGHHLuPUvU,655 +torch/include/ATen/miopen/miopen-wrapper.h,sha256=Y1cYYpIt-bcyZLNVSgjo2Ye1OsIKfzJHs4xFP6TPX3s,1050 +torch/include/ATen/mps/EmptyTensor.h,sha256=_fvSflyj38etC3e3sL1pWP2OsuDAEsY3AwaHlY1vd98,1009 +torch/include/ATen/mps/IndexKernels.h,sha256=34eAxG3mfF6MwY1sdF-4ZeNUZV6ap1SxVGEmoosEv8I,9730 +torch/include/ATen/mps/MPSAllocator.h,sha256=91Kkt2Jv5bJEA3e20zK0Th3KpqGtYqYYhCnlJ6u9LSI,18772 +torch/include/ATen/mps/MPSAllocatorInterface.h,sha256=o9q4Zw8_Z2LLILArBvmh-pi1S4Tv-2PKN8j3slwPBoM,2944 +torch/include/ATen/mps/MPSDevice.h,sha256=EE8D7xXZ97ePt9RfALBi5ter7NMIm7rL_mX5X709t54,2145 +torch/include/ATen/mps/MPSEvent.h,sha256=dtNUJogYAMLEbWIm0Ek-pKHfA-N1xbzBwjbS3L7HRUo,3822 +torch/include/ATen/mps/MPSGeneratorImpl.h,sha256=Moeo7-MXtLeO1d3JsIOKcittj89_V1Xnk7q7Re5FHTo,1816 +torch/include/ATen/mps/MPSGuardImpl.h,sha256=CLj6s6jdwRDAF5jzsGYRK-vzWzQ8DN-NE2_YpP2eBFs,6248 +torch/include/ATen/mps/MPSHooks.h,sha256=Ouma6-LIm1nYmd7vjtJFM-acP7Oh6P2OkWy5mrCad1A,2556 +torch/include/ATen/mps/MPSProfiler.h,sha256=FEz5ZCb97uL3Eer8IQtXSa-qXV4kItq71_ZAnr9JyxU,16844 +torch/include/ATen/mps/MPSStream.h,sha256=4M_oNkmq4B_JOXiFhTdmGAV0nqk6AqrAzRZ-uvMnAJU,4980 +torch/include/ATen/native/Activation.h,sha256=5jUBzS9MrhXbg9i0cc0yWS_mG0CTYOd8FyELEqpSj9E,3769 +torch/include/ATen/native/AdaptivePooling.h,sha256=gbGBydyKyO833hfzVuf6O09uJrEPEAuNVLALm1GHfvg,2682 +torch/include/ATen/native/AmpKernels.h,sha256=NF0ckllzP6u0cuybr-a7mHBEF7iRAmdq8OtuMoomSiE,871 +torch/include/ATen/native/BatchLinearAlgebra.h,sha256=vc3DbmSXZkrkr9QfLaYUnNiENr_wLyJwzovyqaga8c0,10509 +torch/include/ATen/native/BinaryOps.h,sha256=jbHv8lK91cmQDQsQNNgw8eNTPdbed1SjvFqZsCNFxKM,6312 +torch/include/ATen/native/BucketizationUtils.h,sha256=DNCdyntNrq3f8qbdWbcjYJEERDuBIKUeWpIlKNryyjg,8036 +torch/include/ATen/native/CPUBlas.h,sha256=JzMm7HtekP9q4HPAnv3A22UJ_0COj2O1MD97o4rxAWQ,9747 +torch/include/ATen/native/CPUFallback.h,sha256=db7ZH0Xvlxq2yFGBkdYqp-6uKZ2cKfT45H-7KklaGs0,2667 +torch/include/ATen/native/CanUse32BitIndexMath.h,sha256=8u7Fm_XJ5AqSN_h2G2jxyAlYtpxGDNR-Hdufhwu5p54,496 +torch/include/ATen/native/ComplexHelper.h,sha256=PCu5VSO3tP_NiKQ2r3WAWVnfPDxlI8JRVp2EdcijzzM,4577 +torch/include/ATen/native/CompositeRandomAccessor.h,sha256=-jQf0lb1feMo8S_HqBbpp2jY3y8ORmKNjNnT2syptUo,1130 +torch/include/ATen/native/CompositeRandomAccessorCommon.h,sha256=7edXzVXw7rpfekC5_PRaCIzexfRYEJxUm8WXUZBsmlE,6987 +torch/include/ATen/native/ConvUtils.h,sha256=gE1fYnnI9sVIiB_5g1iZyXqo-46N5Zu0QRJcFh6EQNU,21498 +torch/include/ATen/native/ConvolutionMM3d.h,sha256=G0D6p90oXpyeb2mhp_fovuRaIBWX37n6jN6q0Nxc8Ao,594 +torch/include/ATen/native/Copy.h,sha256=9mJ--HBtsJ4-gjBbG2uMgjHNN8_cOL3fEFkL-kSCQtQ,626 +torch/include/ATen/native/Cross.h,sha256=7yIZE8n6_-zcjM494HQmI2n7v7ncpJIbqx_-yzpc-gs,513 +torch/include/ATen/native/DilatedConvolutionUtils.h,sha256=GidEeuor8wOenGeOxSXQN4Tn6_DNOHjfCEmsHwoIzRY,6656 +torch/include/ATen/native/DispatchStub.h,sha256=w24Fr4JUcA5X9zspynbvrEZLNTgMqwNCks99a_hc4oU,16205 +torch/include/ATen/native/Distance.h,sha256=z1vYpapV_W6PyHABYumrYd57ggsR1gK7vBPlf89PzTk,974 +torch/include/ATen/native/DistributionTemplates.h,sha256=eDSxPmwC-NrVQuM-A54b6xTzMh3wLCt5kECtllrO_6M,18746 +torch/include/ATen/native/Distributions.h,sha256=ct735m3wdATGBJSQQ9UdlzmJJTLLZvN1gLq3a_T_huQ,21551 +torch/include/ATen/native/EmbeddingBag.h,sha256=dlik8mXJIeXv57OzTxGs2DpWYg3CzfPs_xMRDFCp26Y,5482 +torch/include/ATen/native/Fill.h,sha256=4RBR3xh3iiYpk0SU93MXSVuh03u0PRpQNraXcU4T7UI,650 +torch/include/ATen/native/ForeachUtils.h,sha256=hQT2Yo_3BVhtcxBc9dx9gMTZK0IFzhItDTf-Ri6eAIc,15316 +torch/include/ATen/native/FractionalMaxPooling.h,sha256=NIrym6OL5oKgXOuG-jWS5Hd-37hKOypZUbXOloWqz3Y,2415 +torch/include/ATen/native/FunctionOfAMatrixUtils.h,sha256=TBbECD73rwVVPvMePLeIxZhTyQJYTvLs7p2om2EkmMQ,642 +torch/include/ATen/native/FusedAdagrad.h,sha256=xYmXlb-ICVMYCI8EVdfZ6CgoD9AbVJSn_Zaz4Gg8V2k,749 +torch/include/ATen/native/FusedAdam.h,sha256=9f7Q9Gv9DbsLRwsnhxK5Evy8_DoFMmU_MuPEQAK9OVQ,937 +torch/include/ATen/native/FusedSGD.h,sha256=m_wkHWws_x4aplWRxIS-8nWyflrUGrGpJVMSqHMFLDA,770 +torch/include/ATen/native/Gelu.h,sha256=_xxqp_o3ocmwVdwYiyJ_Zr_Ak8fTqfRQmGr80GkUlAs,1097 +torch/include/ATen/native/GridSampler.h,sha256=RAbPa6nbPgTD-8qO-e3Dol64-8WiiNYFyqOYxvQB_s0,10661 +torch/include/ATen/native/GridSamplerUtils.h,sha256=9uAjJvZ3lzV6hfBlcCgpocCjIyaYC0LRyHt7l--bMO4,3959 +torch/include/ATen/native/GroupedMMUtils.h,sha256=F5B8kFCOYg7P5ycPQETfAo6Ir1MDXMBYWAsFaIgnSOM,7611 +torch/include/ATen/native/Histogram.h,sha256=81LzK2uszgc5Hk35lUpf7Ocuibm6nmO7hSDZdOQZ-yk,999 +torch/include/ATen/native/IndexKernel.h,sha256=z0pBirBIRp3tNWAqtvFsYlTj0tuVQt-LQqLB3yBm09U,1957 +torch/include/ATen/native/IndexingUtils.h,sha256=OjtrR29cRwi9-4fH0BjD2RFVfcpSScQWt7pABGjl3S4,6594 +torch/include/ATen/native/Lerp.h,sha256=dM_FRmp1GtU8hHyJzcsPhYAxi9SldMaWMQnPzvNyaB0,1715 +torch/include/ATen/native/LinearAlgebra.h,sha256=UzXQ5Os6eK8E_TU2EAJpUOHjVEVpmEw6hPP4lCY5BuM,553 +torch/include/ATen/native/LinearAlgebraUtils.h,sha256=k5um8omLdExad2wIuFvE4tVraBNS8plbaeWaFqeU8_M,26719 +torch/include/ATen/native/LossMulti.h,sha256=bE__nG_H1ux0ygwp8OfxzM6DWydfHPxSscPHZxYiahM,2526 +torch/include/ATen/native/Math.h,sha256=TagZWGU2KFGoEEaKv3pvgnWnmvedgTK5_AV5hbAWM04,142344 +torch/include/ATen/native/MathBitFallThroughLists.h,sha256=LV6aBXlLS3r6JeU0EGHrzFQ5STahu2NUlS_vREsH0EE,4390 +torch/include/ATen/native/MathBitsFallback.h,sha256=qNxqOQoXjCyZ7FXrgAxb5HwhEigG_Gm13QQB7wLLSIo,7572 +torch/include/ATen/native/MaxPooling.h,sha256=9ArUuAHT2x-sECepMiaFZ2VvcQoTENfr6aGPCERXbpg,3522 +torch/include/ATen/native/NonEmptyUtils.h,sha256=OaBD1Vyz4PTpn3Us5BQ3MZ9A7xrFiRuxWs9e_kVBQGg,853 +torch/include/ATen/native/NonSymbolicBC.h,sha256=biHlvGKtt6Ui93Ju32MmURhSEc2npmkjX-7JGsxW6oI,3130 +torch/include/ATen/native/Normalization.h,sha256=AQ5IVOcofut6ZE2l9SedS-XVD9kxQK1qOzLud2HjNi0,807 +torch/include/ATen/native/Padding.h,sha256=KGtXh2yuMsGoD_cTGRnRkz-rRdcIwwxGJGSyLGDdKf4,2315 +torch/include/ATen/native/PixelShuffle.h,sha256=LtPpsV4uTDfmEDnDNL6vxDqV8gzFa8XtCQ5pYIf4g5o,2211 +torch/include/ATen/native/PointwiseOps.h,sha256=QcuNw-z_rvR3fz5fLw9U-w9Am-akjoCIefj5doVF9bA,1035 +torch/include/ATen/native/Pool.h,sha256=gTWkL_odUrvxSDhyoX7dblHMu4DYsDlTXS_84RiVL3Y,13508 +torch/include/ATen/native/Pow.h,sha256=HAOQzKdHYhaHKFgigE4SjB6svSZSNdldUkv1Z-5WGqk,1903 +torch/include/ATen/native/RNN.h,sha256=9DZwrm7U_KLV0y8nrdXQAs7PIlEtH763L8iG6wKlmuE,2758 +torch/include/ATen/native/RangeFactories.h,sha256=GT82Hg0DPHlPWboZA-e2myfK4wyGHaBXcPWU4Wu3HLY,608 +torch/include/ATen/native/RangeUtils.h,sha256=U5qvswfBKI4rYQUoD_sUPLAWN9FT1wWHoXDs1495iKI,2356 +torch/include/ATen/native/ReduceAllOps.h,sha256=StNDQgM49wEdfdroZvpTcll7h0hrIkqb76U0lnya_kA,651 +torch/include/ATen/native/ReduceOps.h,sha256=H0Ifl1k2tQ1d2z2tEHloAmfyKO2GJgqqjse9qy8bEzs,2108 +torch/include/ATen/native/ReduceOpsUtils.h,sha256=6BVLtVZ56urcbQbmSR9dkVuHImpbsxwYUYJ6PeTJq4E,16691 +torch/include/ATen/native/ReductionType.h,sha256=-84o3NlrAKdzIHf2vu5wgDikNy9-UvowgkMs9XnZINY,1393 +torch/include/ATen/native/Repeat.h,sha256=e-6_LlR3AydXMSaL57bvclfYNfMhE6BELxJQggDrtmY,1731 +torch/include/ATen/native/Resize.h,sha256=T_EyA73ldgodOYTz6mct1xidMJSUhNwcGufn1TQmsfA,9819 +torch/include/ATen/native/ResizeCommon.h,sha256=DWUOSk4e5RUodW3aJ1_a9bqsg47N2QB3BC_TN4PHd4M,2745 +torch/include/ATen/native/ScaledBlasUtils.h,sha256=TOdBbX8NrHvkEuRce6dQiNj1WevqvxpKU3-2VqvUflA,3630 +torch/include/ATen/native/ScatterGatherChecks.h,sha256=jh1aGzDCQQvyGmo1BqMbi2PDujCYWLkgqqSWE7emjQE,3990 +torch/include/ATen/native/SegmentReduce.h,sha256=KHETESlR7O3jxUcHvj5CD4MjAwzM8VMNM-Wv8RHfqDM,1519 +torch/include/ATen/native/SharedReduceOps.h,sha256=daa62aPH9HVyq4Jk_WZIdpfip3u5-6Gw1VXZZ6rBi14,16470 +torch/include/ATen/native/SobolEngineOpsUtils.h,sha256=pNNPAQN0x0Br2jQHikIWd2xht9D0yGU8DUg_n7pQ6IA,2089 +torch/include/ATen/native/Sorting.h,sha256=YthtMuMvduf7oy_qqfjL3hbxZH9nNbbg4c91v2BF4aU,870 +torch/include/ATen/native/SortingUtils.h,sha256=bV3C_qbVQTfKY84nXQGcCDL3xUrzprsaviBKx3vf8r4,2926 +torch/include/ATen/native/SparseTensorUtils.h,sha256=HQVciWWBvDldwO7g-RSKPldE3uOujADedTbqVuK45aI,6754 +torch/include/ATen/native/SpectralOpsUtils.h,sha256=i78AR7qegJD9FC_CerTuu3dXYOMK65No7cUxDqfMBSU,3536 +torch/include/ATen/native/StridedRandomAccessor.h,sha256=PZKVO993ozr5g9IsTE5FZlDwn4yjA3kXX5VAIAFU7zU,7089 +torch/include/ATen/native/TensorAdvancedIndexing.h,sha256=wUXDM_I6v7TYYKPzg-USJF6o4mA9WzILx1M7cFjLlG8,3151 +torch/include/ATen/native/TensorAdvancedIndexingUtils.h,sha256=b1l9x8rZl9tFS877RrQUWwwN_Y77Fo0j0NHzT2tfuB4,3378 +torch/include/ATen/native/TensorCompare.h,sha256=xYZgiLO_WTB8nrPNBsDMsTsjR9akehbFLGZe8YTNGNw,1772 +torch/include/ATen/native/TensorConversions.h,sha256=0MeFKSwGxMwjeE7kPyHTMmVog1un0SCl5hv7akwlTJM,1092 +torch/include/ATen/native/TensorDimApply.h,sha256=xCVtk1TTObjwmO4cBByaHLS16i-phVragc01PCx5WwQ,2081 +torch/include/ATen/native/TensorFactories.h,sha256=iMrkaCXDAL_3PlPg3hd9tzJjpMQUtKVCL9IH_c9TQPg,5652 +torch/include/ATen/native/TensorIterator.h,sha256=EGjHBhlzhM9qtyOLFTPl9Cion2Tk55izMcG7hZMIaJk,300 +torch/include/ATen/native/TensorIteratorDynamicCasting.h,sha256=gfkNvB2B9s6IFHANJR0hLWUzORb24pCEVIJmbVoJhZY,2064 +torch/include/ATen/native/TensorProperties.h,sha256=BQlo_WTYnIIFSy6PQtwSmBOcdp4BTwhzvhRkAPn6JpA,452 +torch/include/ATen/native/TensorShape.h,sha256=94C3OMr9UObeXNtrNnn3OCqcI3AHdblYgPQcpSn3Kjw,4851 +torch/include/ATen/native/TensorTransformations.h,sha256=cOkKfad87XienVW5PnqiQNCkdyQq-bSoIxo4t9BY35Q,1181 +torch/include/ATen/native/TopKImpl.h,sha256=H_JLqHjMs4FlEbo0YJ1yRyJTQsaqYXQLUZF8jKeUB6s,3713 +torch/include/ATen/native/TransposeType.h,sha256=fdgHe1LZjUL7SyFS5UojXJVJaPXJ1Ffo-rmm6y-CsKc,935 +torch/include/ATen/native/TriangularOpsUtils.h,sha256=sGk-__VAdO7bFuSwPEnIEYnVPEzOiR8ynfOqljMF3Ws,2256 +torch/include/ATen/native/TypeProperties.h,sha256=MJATK6vHjuCCgS8N2wde2QxDYsG0YQM_VqN8nXeMVO4,912 +torch/include/ATen/native/UnaryOps.h,sha256=8Pu91eWMl8D5PSJRA7zVp8r2XbER_LMRpAKlJ68VbM4,5669 +torch/include/ATen/native/Unfold2d.h,sha256=r5hQGu1j7sNzIaZzrn0WybL9HOtXfnXM2lXTWXas81k,1233 +torch/include/ATen/native/Unfold3d.h,sha256=9H3NktUNEnXLlSDxJQ1MsFSp9E_Ea1Xq5QiJa-jqouU,1127 +torch/include/ATen/native/UnfoldBackward.h,sha256=53PYK7KVjOHfnnDd5s5idG9soaE5y4ql5qWYilyYjC8,3347 +torch/include/ATen/native/UpSample.h,sha256=U6SwWAXuyRcGQrIh5U5cgx-WN6nybpcIU5NYcJMhr3c,19748 +torch/include/ATen/native/ao_sparse/quantized/cpu/fbgemm_utils.h,sha256=bn3kx0zdKPRUTi1ATGIOnylaPwbH33_pFPyd3qShkHU,3234 +torch/include/ATen/native/ao_sparse/quantized/cpu/packed_params.h,sha256=wUv3YyCo_WkTQfM6o5r8VidmyAUisf_WixBEfELkvJQ,2995 +torch/include/ATen/native/ao_sparse/quantized/cpu/qnnpack_utils.h,sha256=1rAzuMEU4AdDuEvherVtCc4hK8bMoNysZNH_4-71b94,3496 +torch/include/ATen/native/batch_norm.h,sha256=DMbmfhic-MpRaWOjcvtgds0yUN9MnY6jr6hhZn3BTIE,1679 +torch/include/ATen/native/cpu/AtomicAddFloat.h,sha256=u2eTBDVfo1ZegYDR-3eZF-6wrTDy_5qTonS6M-alqao,1096 +torch/include/ATen/native/cpu/CatKernel.h,sha256=xgya3MDFZuSx6GvrxfEYCwzjlqhjPgwzkK9wdPvf51U,560 +torch/include/ATen/native/cpu/ChannelShuffleKernel.h,sha256=_HVvj3sFQKRmPnsn2KYuwYIrzWqOZkHU06wnXVWsd90,540 +torch/include/ATen/native/cpu/CopyKernel.h,sha256=5RtfJ5tLE8jp2azBmsxJV0xCsexHxAzmne3Uc9Ly2zw,566 +torch/include/ATen/native/cpu/DepthwiseConvKernel.h,sha256=ogyYOVvDHyWg6ACz0jTZesJs07A5pyUXdWW08UrCZ5c,724 +torch/include/ATen/native/cpu/DistributionTemplates.h,sha256=3OdfgDag2-7YI_eygiRWc7dHguevGCdV_Wur7ncIMxs,18024 +torch/include/ATen/native/cpu/Elu.h,sha256=7hsZrAQh4SmgD_4eEv8E1Mkv5PkW1IkyCEXflgByRzM,3120 +torch/include/ATen/native/cpu/Gelu.h,sha256=gLK-AVzZIUwC_hPb5mXyomsuxNkbLljS_gMf8mmbkok,3358 +torch/include/ATen/native/cpu/GridSamplerKernel.h,sha256=NcAOuHU33S2mQT1pxM-o9MZJcvMacy43GQieH4YDfCI,1077 +torch/include/ATen/native/cpu/IndexKernelUtils.h,sha256=XUCtKIuJ4g75rN_GZekuUlgyFRfqZBr5yzAFlS1Kfx0,3181 +torch/include/ATen/native/cpu/Intrinsics.h,sha256=Is00ayXygoN_okRnxej4mDi0UCHphme37N7fVKmicvs,1466 +torch/include/ATen/native/cpu/IsContiguous.h,sha256=g8a8_i1hFP8Gu34DA9IC1TuSNYL54KsNeEiB1BctEAI,2623 +torch/include/ATen/native/cpu/LogAddExp.h,sha256=z2avx_FN6ANXdonid9GMF6JpwdevZ6SzCRQlynkANl4,2700 +torch/include/ATen/native/cpu/LogSoftmaxKernelImpl.h,sha256=6yR5jf2tWdYR03lf3-B0E7qr4mm5AlX10rBKORGK_tY,13404 +torch/include/ATen/native/cpu/Loops.h,sha256=JbAFu0_vDIzVA2irFOum9g6AHiPm83F2_W_OVfMvBDM,15402 +torch/include/ATen/native/cpu/MaxUnpoolKernel.h,sha256=gRdiFBOyWHNBjsDJclEudywwpCPhLOfMBYqKUaK6DK8,560 +torch/include/ATen/native/cpu/PixelShuffleKernel.h,sha256=gYm4HRFVMcBhRCObqeGH8mZWRfj0Yeha5YKZlyy5eoI,574 +torch/include/ATen/native/cpu/Reduce.h,sha256=9yhsH6l-lll7hyiWip9jnR0X8ZE8ZJJvjkhe1XDfhU8,12324 +torch/include/ATen/native/cpu/ReduceUtils.h,sha256=bT7DEvWx_0YHMvXhqVRN01YaQ8Lm9h4D7A8kKGdIjGI,8963 +torch/include/ATen/native/cpu/ReducedPrecisionFloatGemvFastPathKernel.h,sha256=vuSKm5PeIz56Ko3ctjEvg7Gm1Aqr0euypN-cqk_VYLc,1374 +torch/include/ATen/native/cpu/SampledAddmmKernel.h,sha256=hG9d1W8jvG-36slz4ApkeDEJrPu9n7o_5_9H9ba-X5s,576 +torch/include/ATen/native/cpu/SerialStackImpl.h,sha256=Sq8zzXe7u9c3o6sLY0VBHCNK92YORvU_GvbNt-NVKQA,5710 +torch/include/ATen/native/cpu/SoftmaxKernel.h,sha256=Ll2YpX52fims5NsYPrbT8ygfZp5wk7CwC9-7lGEYpTY,1189 +torch/include/ATen/native/cpu/SpmmReduceKernel.h,sha256=85iB6aCW_bH7v5iLwb_ofq6idbS-rkdktLx0lW5PNmg,1606 +torch/include/ATen/native/cpu/StackKernel.h,sha256=IaY88AYCdMcb0c3TWdpZVaGsTTqDOc4gDj3HDNFWr3E,562 +torch/include/ATen/native/cpu/UpSampleKernelAVXAntialias.h,sha256=LMvbkoyLh6gyLxT6J9IZvAMYuDRI67o1zc2HBbi-0sI,58404 +torch/include/ATen/native/cpu/UpSampleKernelNEONAntialias.h,sha256=9vZLT_LXwi2l1JTvo6Cds6f4ntEkeQHqxxub04KcahY,16293 +torch/include/ATen/native/cpu/WeightNormKernel.h,sha256=GzLqqFhn6W_bf0MRcBPOP8kl6qGCejNbdtWMytinYEc,804 +torch/include/ATen/native/cpu/avx_mathfun.h,sha256=5xPW_GSg4DSxPyz0s7nUgfy0qw6iqSUVTuEob-Um0Ho,17702 +torch/include/ATen/native/cpu/int_mm_kernel.h,sha256=VGOuXJXQ8zVU1GSBXkH-bFVrK-8TwjPuiBLxv9cocHA,1359 +torch/include/ATen/native/cpu/mixed_data_type.h,sha256=WAxkpUy7WJoR8cwYBENYWN_o5NicWAe6UaYq518p0AA,1662 +torch/include/ATen/native/cpu/moments_utils.h,sha256=AiVFmQj3IOdLTWgU5LYXNAQrPX2yL8jJ2YqIntESbf8,7162 +torch/include/ATen/native/cpu/utils.h,sha256=A0dq7dHUPvH9NwdZnFO5buqJbzpCMkVuG40rm8n8ZKM,7770 +torch/include/ATen/native/cpu/zmath.h,sha256=pXuvamS1M8HhO81LfeVy9zLi9-X4VniQadUpAHYCoMg,6876 +torch/include/ATen/native/cuda/Activation.h,sha256=os9YOpv6tVDWxq8FsjvMEf_ip-aj-g-6XZXK7XSWb7c,790 +torch/include/ATen/native/cuda/BinaryInternal.h,sha256=3c2TDjMEUA7aOoj_LHQM17CHB49TN7ajr13BPlLU6yU,1441 +torch/include/ATen/native/cuda/CUDAJitLoops.cuh,sha256=UHtIF3HcMsFINgAfPY1HwhhMNDaxgxGUmHzRShDVz2k,12131 +torch/include/ATen/native/cuda/CUDALoops.cuh,sha256=Emx5WIkKBDfnnGkM6r0R5GdVJxJI9uxUXvbfhjuN8ic,40879 +torch/include/ATen/native/cuda/CompositeRandomAccessor.h,sha256=ufMChPnW_wYVMaUT6MmtLvu_oX96wCMdXDyeGNGY0ms,1201 +torch/include/ATen/native/cuda/Copy.h,sha256=mcsIfjzE4YvB0opzHQzwysrCdlasG9f5PeZ4OVCIp4M,408 +torch/include/ATen/native/cuda/CuFFTPlanCache.h,sha256=uKMnLjAkwtkpkvrhIqye3PEGy4-4-DlILsjioYFP8gM,18182 +torch/include/ATen/native/cuda/CuFFTUtils.h,sha256=Fay9wPMjepT8wcVEFxa0BcsyD7eGrF8LeCxCL7VqjKk,2185 +torch/include/ATen/native/cuda/DeviceAddCmulCdiv.cuh,sha256=y_1B3yB1E43ZzHIWZiMK_d46CAavOcOoK_HBAs8Mscg,1282 +torch/include/ATen/native/cuda/DeviceSqrt.cuh,sha256=LXvb7RnB_WDeLmmRDAlc5y-I6bWJg_YBv5wxn1e6Qv4,839 +torch/include/ATen/native/cuda/DistributionTemplates.h,sha256=1VQCK26LYYgMoCx7Vy2G7yZqZJygRMZhOEuJvFm8nHg,29502 +torch/include/ATen/native/cuda/Distributions.h,sha256=sWc_6Qaj8BXxIjWCfmi7rsdLwvAq5jx1vicAqE6KXeA,895 +torch/include/ATen/native/cuda/EmbeddingBackwardKernel.cuh,sha256=TUm8dx58lh7P86KCI_wfliCtu8DKm4ybcAu9540NcpA,809 +torch/include/ATen/native/cuda/ForeachFunctors.cuh,sha256=4FHSnFIbdbBXjOQYwSquUhibdIQXhpms6YMW7z206f8,30310 +torch/include/ATen/native/cuda/ForeachMinMaxFunctors.cuh,sha256=-qBpnBW-iMbqxIWPIsXk2bEUwH91DV3vtUTVui8fzKw,680 +torch/include/ATen/native/cuda/GridSampler.cuh,sha256=wiPXUzbFGiQXF60jGAlqigEYMwrNEreGvVrlqC88a68,11229 +torch/include/ATen/native/cuda/GridSampler.h,sha256=P0YJuHG5fCMvbqJvHEYP4K8bn4i_vOvJ1NEDHwUNjxQ,1399 +torch/include/ATen/native/cuda/GroupMM.h,sha256=-L4d4RsZiACoeXzFuHkQFmvaxi8PuNx3OlILtXY-Sfo,581 +torch/include/ATen/native/cuda/GroupMMCommon.cuh,sha256=_mQeCqW2nyRjPOlompxxJKxN5wyh8WJWUhfJFRTO530,6706 +torch/include/ATen/native/cuda/IndexKernel.h,sha256=aotVt_oAfWzDL5AjtcEfCVp2lHop4KzwtQSR3SoXXvE,592 +torch/include/ATen/native/cuda/IndexKernelUtils.h,sha256=VtIHfQrdq92aKX-ET9NvR33XN4jO0r2M3M_MwHC6AZI,2160 +torch/include/ATen/native/cuda/JitLoops.cuh,sha256=PT5ek97JOjp45YfLpfHxZz0fGwzRPb6JUqirw-a5kXQ,7161 +torch/include/ATen/native/cuda/KernelUtils.cuh,sha256=DZKhQH_ABVE__BpthKqqC-zOlWIoHrooxfkHomVmEAU,15356 +torch/include/ATen/native/cuda/LaunchUtils.h,sha256=Pbewby1CzWNCP8etqYoWS9pNYmeRrIPILujKxTywS3Y,536 +torch/include/ATen/native/cuda/Loops.cuh,sha256=FZR4RYhNe20cHnC0OQCIujb4-J2iqOhp0tNbCU5DUEU,12338 +torch/include/ATen/native/cuda/Math.cuh,sha256=BiYbgQ8jK41prA-xRUv4xotrVlQ4a6eFovJhMgxbtVY,123364 +torch/include/ATen/native/cuda/MemoryAccess.cuh,sha256=YVJgkgH8dDmE4dHNx8Z5StQtCp-9jYhwrIs0AudAn2Q,22533 +torch/include/ATen/native/cuda/MiscUtils.h,sha256=L73h_i3i-iP9aoPPfLP1PUSmpBsTQ6ED4GiBboqkVI0,1188 +torch/include/ATen/native/cuda/MultiTensorApply.cuh,sha256=uN3LY8MTefHKLDDBR_Fp7cTqGAgsZfyJqjcgfJsd1Uk,15257 +torch/include/ATen/native/cuda/Normalization.cuh,sha256=gkq0oihV01mGKLQ0IWd8gsMyS6FemN3o-c4eDjJou_U,76332 +torch/include/ATen/native/cuda/PersistentSoftmax.cuh,sha256=dD0LAFWcskRkd3VYEQyGlglkaW-DcbCfW_9fcPAWv8Y,20485 +torch/include/ATen/native/cuda/Pow.cuh,sha256=OQysVDoNrRMJqiCd-d457gjAyrDY3xjgNdAQh_Vo0RI,2408 +torch/include/ATen/native/cuda/Randperm.cuh,sha256=rMZPMF3X_t5foRCbZIjUNg7qKDxmEDtN1BBPyc8mQ-I,2362 +torch/include/ATen/native/cuda/Reduce.cuh,sha256=Ev8iOMnUHamRay_9-3YXEclSDCGOZdeyOImIiWoV9kE,50301 +torch/include/ATen/native/cuda/ReduceOps.h,sha256=1S_nThLRx9CAF-TatBDunTqEal1ciIepo525hWzBDcU,803 +torch/include/ATen/native/cuda/Resize.h,sha256=OrAewKGX34CeDGNxaF_6ja4Z4Ld7UdqTWXwlxM8Mic4,1798 +torch/include/ATen/native/cuda/RowwiseScaledMM.h,sha256=w1t47IiYnTc9w2lGdhWRA4Kg-TIaCKBaE0Y0rjmMl1w,623 +torch/include/ATen/native/cuda/ScaledGroupMM.h,sha256=Ig4NQervFHI5rwjR6e8VZeYhz7lIr5WrEWx0uKwSNWg,668 +torch/include/ATen/native/cuda/ScanKernels.h,sha256=vFFvo1hDdUGcZ6pBJSMPuMq3d3HSvyESdQep7KOmWxU,1033 +torch/include/ATen/native/cuda/ScanUtils.cuh,sha256=0yPz3LXhdf7thAKU5n6CMA_dbava1RYHL1SXurED0gY,21650 +torch/include/ATen/native/cuda/Sort.h,sha256=Ak6nIQWuEAAmsUVIIvp7XEa8-pDKNOPBe8lY9g8J-c0,655 +torch/include/ATen/native/cuda/SortStable.h,sha256=0bVzhFzFllv-yHxlYiIrJiUMB1vkAGkuchYwUOOfI1c,693 +torch/include/ATen/native/cuda/SortUtils.cuh,sha256=_VR_2RT0Wvt46x_rHEFf61CTPDhi1xai9s1RSwXITGQ,13466 +torch/include/ATen/native/cuda/Sorting.h,sha256=CWfKEDadWdjAy-osIW9yRNCPtpm0ifR9V0Mfy2kTrf4,650 +torch/include/ATen/native/cuda/SortingCommon.cuh,sha256=7mvhis3CPtIu6_tyCiyi_In6u-ozg-AwKWmq6Aj86HM,5710 +torch/include/ATen/native/cuda/SortingRadixSelect.cuh,sha256=1AmvW02mVhUSeOyhhoTaVaCDIpJyDqBCaohelsNTK9Q,42341 +torch/include/ATen/native/cuda/TensorModeKernel.cuh,sha256=TLm-0Mxr3GZ2ktHf8kK8JapE0nmU0JrlLfXUeLYlpQc,14683 +torch/include/ATen/native/cuda/TensorModeKernel.h,sha256=xxQTk4ZtcEQG4EA6jVkigxUAFeL5rAgBVFgjm6vBW9g,673 +torch/include/ATen/native/cuda/TensorTopK.h,sha256=Yz2XLDh4W77eya4wUV70dJ3IqTm-r2PbN_oTOK4O67c,508 +torch/include/ATen/native/cuda/UniqueCub.cuh,sha256=5UcInlJD8-4H0pDNDw7XBTR2Lgcsf24rKZ-4p2R8VAQ,556 +torch/include/ATen/native/cuda/UpSample.cuh,sha256=ilf5bsh3nfgd421zzMPS77zv8yrCYkb3A5G_S317VXw,11892 +torch/include/ATen/native/cuda/block_reduce.cuh,sha256=rcTxhFDkP28XZafyxQfGWqHIgRecsCtHQZFA_ss4ye0,4733 +torch/include/ATen/native/cuda/cuBlasCommonArgs.h,sha256=eTaQRfMTPe7bdAnJ7ddHKXwkFbqMRtBc2He29qd2CXM,7119 +torch/include/ATen/native/cuda/cutlass_common.cuh,sha256=hVn6wIrFmMVY_DwRvPdNKpLe82ss50Do8p99w_DVXbA,1531 +torch/include/ATen/native/cuda/fused_adagrad_impl.cuh,sha256=LbnP4NqPrpXCPup_huZIMvHz-_fZt_Mgktil69vpySg,1099 +torch/include/ATen/native/cuda/fused_adagrad_utils.cuh,sha256=_GIkuYrr76j0vaT90n3b5UANfPWKtVVxjtAOO9Tbwjk,4476 +torch/include/ATen/native/cuda/fused_adam_amsgrad_impl.cuh,sha256=wnDK1DTBn9kwV3Kzybo0t9b_WEFhxg806saEN3T3wLE,1283 +torch/include/ATen/native/cuda/fused_adam_impl.cuh,sha256=33CbBaJsUWSO9nO5mvJN9XhD0sngVwemzdueL_2AQ9g,1195 +torch/include/ATen/native/cuda/fused_adam_utils.cuh,sha256=uYaphOgnWTSHuIFKdWqVSPGgwv_LvsJlswpdkFHiJeQ,19839 +torch/include/ATen/native/cuda/fused_adamw_amsgrad_impl.cuh,sha256=y00SBNCAsVlsZ1yNMvoxQVfIJyDRr0j_oK_QeKrwBK4,1285 +torch/include/ATen/native/cuda/fused_adamw_impl.cuh,sha256=NEimgkgX2RGXMO_KGfGCQ37GPrHi7NXdGeSbp-4dtLU,1197 +torch/include/ATen/native/cuda/im2col.cuh,sha256=GK-g822eCAPDyguDVMb6RIxm2jzAmbTi0d7Loj5Z_h8,9954 +torch/include/ATen/native/cuda/jit_utils.h,sha256=03BRLVBw9yC2PPzbaUB0MqAO5IYCnvCwbi1YFQKfc-w,7365 +torch/include/ATen/native/cuda/reduction_template.cuh,sha256=_q76j_Z_lW_B3k2hXwJHGnGnj8HBefSlNitUWNh661I,22057 +torch/include/ATen/native/cuda/thread_constants.h,sha256=YeD3M6Aacp-WA2zYS3hkVzHRJlNPPsyXKj6R08N0FtI,914 +torch/include/ATen/native/cuda/vol2col.cuh,sha256=Pm5y-TyCSSlUZCkom13iE85Efr9V141hhYcshxE3JBo,8347 +torch/include/ATen/native/group_norm.h,sha256=3u_A32TeaFVatZXYQChtc-98oSaaBZBC_m253cOhEYI,1159 +torch/include/ATen/native/hip/bgemm_kernels/bgemm_kernel_collection.h,sha256=3QNiGsAXKaS5wl9jlOUdp43uqc7ejrc2Fa4B0imnqfg,3708 +torch/include/ATen/native/hip/bgemm_kernels/bgemm_kernel_template.h,sha256=QTcgcIQ40ciKvIZjhBBXhyV_79qOTwuIYa0YRJR4rxs,6209 +torch/include/ATen/native/hip/ck_bgemm.h,sha256=2YM3vQrL0cznzCtrLD8jdyj1YfvbkZzWL4rKOIQm0sg,699 +torch/include/ATen/native/hip/ck_gemm.h,sha256=j5WmuO-ZeBqgGgdwMbXwZ7lgDhuPsw5zRGgQWzlHMSg,933 +torch/include/ATen/native/hip/ck_gemm_template.h,sha256=ia8G2fb8ZeWQrQfBRXELE80h-9ZM6i4qychghyB4GKg,15834 +torch/include/ATen/native/hip/ck_group_gemm.h,sha256=zdAmCdcA5YFBDElcsRJuh5RtKnW8vfOxFOcwzH_4vYM,643 +torch/include/ATen/native/hip/ck_types.h,sha256=afMlwkSm6I1DhOS6v5jDSwOal7C1g0Zl73nNmeiajew,1849 +torch/include/ATen/native/im2col.h,sha256=fMatUfeOB6WmPopF-fxFy_dATyP8C43su7DJaqveaXs,5491 +torch/include/ATen/native/im2col_shape_check.h,sha256=3HJPR3_e-amDwi0sZr2ylKbFxdoIRGNP9vPT8bW6Vg4,7589 +torch/include/ATen/native/kleidiai/kai_kernels.h,sha256=_bbIJKjF-_rBeWJPMTBXV4TylwpVPt1MwEEynm3V_gg,1346 +torch/include/ATen/native/kleidiai/kai_pack.h,sha256=x8qdEciDczhNbVYxzI0hAqIt3bxj3peMxDM5X7AnVRw,3009 +torch/include/ATen/native/kleidiai/kai_ukernel_interface.h,sha256=Pnj6ZOTnfrEbIvx2Q2epTrWwqimzaUbVqgYZR35Ib5A,8379 +torch/include/ATen/native/layer_norm.h,sha256=nkD5VO_N5SxQ-8CyTcIi0OZ7UZytMrmrQkToqUOkXz0,4874 +torch/include/ATen/native/mps/Copy.h,sha256=cWZYhk8NCIlRVberjIBYriUyhpU7to3gVc5qp4hJLW8,549 +torch/include/ATen/native/mps/MPSGraphSequoiaOps.h,sha256=WH7-8aQhNo5GK5bK-dY9xSSMg50s5GzJPp7Rhyfzq0o,2124 +torch/include/ATen/native/mps/MetalShaderLibrary.h,sha256=Xj0Vn8-kQ26LxPFXzPt95jnlmRhm_NUrARJ0oowg0ms,6882 +torch/include/ATen/native/mps/OperationUtils.h,sha256=CYPVzVxfaycR8tlWJI0r9r78_AB4yr22O_98QWkm3Jg,29961 +torch/include/ATen/native/mps/TensorFactory.h,sha256=8h9YbcxYmei_ZpCZo2km9xCxUsI8c-07wSqyerPFXmk,1258 +torch/include/ATen/native/mps/kernels/Activation.h,sha256=leCfxgI79PkwRDayn8hE4PkfBPRxajfRZ1KrAU-OKDE,461 +torch/include/ATen/native/mps/kernels/EmbeddingBag.h,sha256=iqVw27waBNcFgZq8_q5Wo0JSc0sq6kzF0D1MqAf40UM,1622 +torch/include/ATen/native/mps/kernels/GridSampler.h,sha256=FhbSBk-7k3Q6kWFHQSs8Cp3KoOw-aVXkY_2_jUPGhYk,737 +torch/include/ATen/native/mps/kernels/Indexing.h,sha256=AeqM1zzjdJ7sfsnqLm3k735x4hOnovtMRCaQl7YpHk4,632 +torch/include/ATen/native/mps/kernels/LinearAlgebra.h,sha256=0XZfcQL4T5iB-Tn1bNDKJJ4_ZvYcIg3OeUdUTjNMXSo,801 +torch/include/ATen/native/mps/kernels/Pooling.h,sha256=AqZYSkn5btqhWq_CnsS9yGB9_S17BI7RRPKxxEJMKw4,2586 +torch/include/ATen/native/mps/kernels/ReduceOps.h,sha256=CZZ2lxLfjdATI7BocG0zsvhdgz3yUXC6wCLcKo7_l4k,679 +torch/include/ATen/native/mps/kernels/SamplingHelpers.h,sha256=3ne2EzzkM7GbJBWvfxTl6eZsG1UyXXV8h7DpykdtG34,1519 +torch/include/ATen/native/mps/kernels/Shape.h,sha256=a7AyZzapWbkD_liMf400jvYyJLDr8UKP57DRUv0whsA,811 +torch/include/ATen/native/mps/kernels/TensorCompare.h,sha256=aSdaemSF2GiRF_7CALJZ4yktIgof63-Tah6zLnI_lV4,338 +torch/include/ATen/native/mps/kernels/UpSample.h,sha256=XUU1fisZof9f7SZVxCuKMy3CEjFITiiqkqMEjKnGja4,615 +torch/include/ATen/native/mps/operations/BinaryKernel.h,sha256=Ek5fMGOxQ3_wtcTk5xOBsJO4GCJJanbdHMcGzF0mkpM,514 +torch/include/ATen/native/mps/operations/FusedAdamAmsgradKernelImpl.h,sha256=G1xwWT0LLYyUs8aHXn7Up8sxahjir4DgGsyHqaRIXLo,1235 +torch/include/ATen/native/mps/operations/FusedAdamKernelImpl.h,sha256=OkmiF5mPsfFZToRmusrwik-P_lyvo49HskPCRq28hYM,1142 +torch/include/ATen/native/mps/operations/FusedAdamWAmsgradKernelImpl.h,sha256=Ho9opeG56HSH8yC8Hb7-wWYH9kW402dQI4NR5CeEPcI,1224 +torch/include/ATen/native/mps/operations/FusedAdamWKernelImpl.h,sha256=xXf1MVaDz2m_RUB6EsuEEYWwfWaV6Eev_cgKOdc_iCs,1145 +torch/include/ATen/native/mps/operations/MultiTensorApply.h,sha256=rZt4ixFlCLtxNMGESV0QqlkTWtQPjOM68ZXX8fjPRt4,16278 +torch/include/ATen/native/mps/operations/ScanKernel.h,sha256=5yN_lLVWwor6jk1Cx8ey-zQUkIFvdExF_pn3EIunWfg,456 +torch/include/ATen/native/nested/NestedTensorBinaryOps.h,sha256=NMoLtp9nf45FWGIvt451kQdx1lHmfiNRh0AtEdxEa6E,667 +torch/include/ATen/native/nested/NestedTensorMath.h,sha256=bQj336uDnNYyflProhEgl8eh2FVEPAZFBUIQG9k4uo4,2972 +torch/include/ATen/native/nested/NestedTensorTransformerFunctions.h,sha256=lZY9I9b3J06iT2alJeCCEqGR_hfbvVho2l0Pr5Nk3E8,3086 +torch/include/ATen/native/nested/NestedTensorTransformerUtils.h,sha256=a4vfhIgp242JzWWPVW4U4ctPGPwgSdO27Rgu8pGveYE,1652 +torch/include/ATen/native/nested/NestedTensorUtils.h,sha256=6-cVbZmOmDqVaPHD7eN0Oqb9UZzPQt0yZIzpN0hhgoc,16084 +torch/include/ATen/native/quantized/AffineQuantizer.h,sha256=4HhLJpNCAXkWvnW09Alk4_QCBa1igqV4X66QyNk54Us,3969 +torch/include/ATen/native/quantized/AffineQuantizerBase.h,sha256=UDX0a_eLF_a2JDVrSf-Y-VGgEk8-Ne25Z2KqmPBUl3o,1747 +torch/include/ATen/native/quantized/ConvUtils.h,sha256=8F1-pAtLWECGODsMOQGp8WcmeEns8ZxwuIO9Ilfvuss,2494 +torch/include/ATen/native/quantized/Copy.h,sha256=P3SjvRGc1TR2-lQ1TmGd0l1BzlB1NVEzm6C8LrtMXWA,410 +torch/include/ATen/native/quantized/FakeQuantAffine.h,sha256=mneUnXSFk6uTqtr8vDp8u0ARXM_qCoM2Cq0PLqe2iiI,2041 +torch/include/ATen/native/quantized/IndexKernel.h,sha256=7n0KvKky8K1dARuMLS0LrxMVfmXepTQ6qzT1VEvcdnE,834 +torch/include/ATen/native/quantized/PackedParams.h,sha256=TjRb8G6SvN_lDqEh3Gjf-lAbtJ1hmNYWJ_tWWC02g9A,4855 +torch/include/ATen/native/quantized/cpu/ACLUtils.h,sha256=gP4f8Vojx_EBJBcJ2-yjmZXegLGxJ9wYtwKpPEHxmpY,8141 +torch/include/ATen/native/quantized/cpu/BinaryOps.h,sha256=OvT3b3V7CcNXnB6iZp-GK8tT7eR9VKBiGsZ4wQSjTEI,422 +torch/include/ATen/native/quantized/cpu/EmbeddingPackedParams.h,sha256=9StMsjybpzchN6Uo0DqF4Tq7TTuQhGUDSrvmN4IR_nc,1175 +torch/include/ATen/native/quantized/cpu/OnednnUtils.h,sha256=eIce-wlwIWkk2OoGEsTqUFlgrRTVRsvWvBxnwlpVU6s,15513 +torch/include/ATen/native/quantized/cpu/QnnpackUtils.h,sha256=sCMzRkF9QMaDiMU_0IbfdWNbfd_XNe7gHV3oB33oMVY,17663 +torch/include/ATen/native/quantized/cpu/QuantUtils.h,sha256=FGCsFT_VNSDEd6C_UcokTYSuMtw1FvgSd6efw_wHfww,8501 +torch/include/ATen/native/quantized/cpu/QuantizedOps.h,sha256=UFaInkdktyPSpCbbcxkmfhca2B1jxF0baSx9XZtDxQY,8864 +torch/include/ATen/native/quantized/cpu/RuyUtils.h,sha256=F7tDQoOD_kZeAHhjbQA1BuluvmbT2brbp5J7XAKYuOE,592 +torch/include/ATen/native/quantized/cpu/XnnpackUtils.h,sha256=KijeLMwXrpDfgJlZgyEUYdvnBl5RvLzQLg0HSr-qhMQ,14359 +torch/include/ATen/native/quantized/cpu/conv_serialization.h,sha256=u33Rn54djfJ0l1E7Hi8PFsxpiRGNwoXSGr9Bz26Rszk,13142 +torch/include/ATen/native/quantized/cpu/fbgemm_utils.h,sha256=NqmpZpLfewpnwH7avgftbcP8nzmF9fLgUuEsKywqm4A,12278 +torch/include/ATen/native/quantized/cpu/init_qnnpack.h,sha256=-_t7Ge0F62TBT3YIrsEWds7bYYFjdtb-GCO0p1Xkb30,375 +torch/include/ATen/native/quantized/cpu/qconv.h,sha256=-nk0VPofdj0D8WiJFD88F_ZQvXSbRnldUx_SNvZETao,3809 +torch/include/ATen/native/quantized/cpu/qembeddingbag.h,sha256=IHc-UVtGXTW1ZOsCIAnCwUecwqtFYLzaafki8a7D5o0,1269 +torch/include/ATen/native/quantized/cpu/qembeddingbag_prepack.h,sha256=R7xEP6Xcg3gVU4ZqOMG3XWxr0kS5UE3i5lk6B4woLs0,626 +torch/include/ATen/native/quantized/cpu/qlinear.h,sha256=lzjM6vKSEiXGSRhe9V9Q3qIhp1dPBm1hsC4VNl2Pq84,1932 +torch/include/ATen/native/quantized/cudnn/utils.h,sha256=qrdJrTM4Dq5k1j6gbRYS8yfUaghFasD6VOdu0JMLHho,10732 +torch/include/ATen/native/quantized/library.h,sha256=68goeCvNPxaYgPVXKxldAUpXQVdtObbGtl50lJwqDZ8,443 +torch/include/ATen/native/transformers/attention.h,sha256=M9kUVQO91aN83n5pJWaHRc0JgxK7eseS6vV58nwiOFI,2520 +torch/include/ATen/native/transformers/cuda/flash_attn/flash_api.h,sha256=FHSqediqbOOZhNvDwKrZ26KDq4GfA0Nw5c-3CYBSTuc,5608 +torch/include/ATen/native/transformers/cuda/flash_attn/static_switch.h,sha256=oh3hsokaOuamc3t3wnjbWVaw-3jMwU7FdizCHDw2hFQ,4021 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/debug_utils.h,sha256=hxFI0_GO_WCTXvNcbXDw6JLd1ZdRV7SczH8B4x7mt_8,10508 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/epilogue/epilogue_pipelined.h,sha256=oBygv63GlhEVVSGPJX_SA6ClKGdnjg9tKPLD0wRlirw,22541 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/epilogue/epilogue_rescale_output.h,sha256=aaF13g1a-GpjSSnkWlGSmbdyi6I3OCVoiREPLEEn_Vc,8036 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/epilogue/epilogue_thread_apply_logsumexp.h,sha256=EiwYb10Y3xRNcSyJYe7B5F8e5fq9JJKkBfMAiNesjrY,6233 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/gemm/custom_mma.h,sha256=ODUdD2YzVUUDgp0bUyIW6a7IfKJ5xBrX8VGgiTzD-t8,2745 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/gemm/custom_mma_base.h,sha256=CZnj6_RVJOuf0L2qqjhBg3610A4ZOsqpTnsu6TlVEHo,6495 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/gemm/custom_mma_multistage.h,sha256=_kvJ_lQvprby69IWo9YgYtlqIM02u95wg-FACAZWqjY,27535 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/gemm/custom_mma_pipelined.h,sha256=j0leQxb_tkfXkUy6nmpGo4Un_j2W9iIvXe9a2BHmgcI,14399 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/gemm/find_default_mma.h,sha256=u2YDMC6WhwawOznVd-hUHEz87Tbn_vxTdcYztSPmnvk,5430 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/gemm/mma_accum_lambda_iterator.h,sha256=Yl5F-XCv26rf4XfSVXcEwAiqBbGMXGYHlvMkb-3AuvE,12604 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/gemm/mma_from_smem.h,sha256=MQKS3DmL_g3dyCKrwyx-N_vHb9c5vRau2pR9YGF-ew8,69054 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/gemm_kernel_utils.h,sha256=c8HRDeKpXFmf_pYPv1hcPWXETI7G1k5oZrmAS9J74kw,8862 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/iterators/default_warp_iterator_from_smem.h,sha256=_QDBg3mpYTLvkm6QmMPvCnlXKH-6q5UBgxjFnjDbeDs,6081 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/iterators/epilogue_predicated_tile_iterator.h,sha256=ji1fnnW6YcwUF0-PZ46XcVX-PiOmhc6V8GLk_TM383Q,24110 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/iterators/make_residual_last.h,sha256=5TwEgMXnapiI8dXgSD2qKaZaTHdTmu4k5OEiIl8Iaok,1904 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/iterators/predicated_tile_access_iterator_residual_last.h,sha256=t1ZDqOd92TNY7poZz8aWKgxcEYxo5UQZ6Lr_trMdpKU,64720 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/iterators/predicated_tile_iterator_residual_last.h,sha256=gp7am6w0DCJEhboAUVzw5hXDASCxHIh-6bIOWxERn_A,64753 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/iterators/transpose_warp_iterator.h,sha256=a4X84U9_G20Fc_XFD3fCkfjvzo2a7ako4AS44pj8ngE,1215 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/iterators/warp_iterator_from_smem.h,sha256=xh58eUP15rAToajIZdPtEldFhQaezl8YbF3vKyHW4rw,10309 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/kernel_backward.h,sha256=fF5Gh_8XIbqv4jRg3q2IXS7PXK7dy1yf11J-039MUEM,100289 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/kernel_forward.h,sha256=GwwC8VHauUblm43-BHo7Cm-u5UgdnjeyYQrhHCoXTK8,54368 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassB.h,sha256=fE6YB0JhXGXoaybYqbJG_kyv1wtGkarMIcCQhUrV800,95653 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/kernels/cutlassF.h,sha256=xCBo8Ic7gnlqjXZhcFMIK2d40oZMRhpXY_h4zS-fDNs,26210 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/pytorch_utils.h,sha256=elA-QEGbaRvCdvmzJBgImCcDhyGyOb4f1shhYMiQ-as,1214 +torch/include/ATen/native/transformers/cuda/mem_eff_attention/transform/tile_smem_loader.h,sha256=tf0MqYkM9JIaMj88yQfIWgzaBMfaIPKMAojAtJgze_E,2406 +torch/include/ATen/native/transformers/cuda/sdp_utils.h,sha256=HbGU4wsXmqHb5zdoQqYLTQ8FxmzjDBn-3FVA0u2Dc1g,881 +torch/include/ATen/native/transformers/hip/aotriton_adapter.h,sha256=GQ5kXAMLeiUYPBla9QHPiGCTCWEuR2B7hlju1lyp1no,6112 +torch/include/ATen/native/transformers/hip/aotriton_versions.h,sha256=-cM4CU3Of2lVciNyxfkGy6mCVq7Jo3qB6SLfAhnNvuU,697 +torch/include/ATen/native/transformers/hip/flash_attn/ck/me_ck_api.h,sha256=2Au0JqJZmYLlAovY6rD_-qFrnp06USJIMe58wDfDRak,2021 +torch/include/ATen/native/transformers/hip/flash_attn/flash_api.h,sha256=U4mNQNV-HAUAEd1Cirn5VB8qMkbdOcylZ6LiWXOGpz0,21108 +torch/include/ATen/native/transformers/hip/gemm_kernel_utils.h,sha256=2r4RXOKhLO517Ex4Vrt1Asa2xU54DqV6mi2t7uChnoY,1729 +torch/include/ATen/native/transformers/sdp_utils.h,sha256=g0V9b6B0GnpezL8TMQpO6a49pwUeKx8onydtOM08Nk0,3471 +torch/include/ATen/native/transformers/sdp_utils_cpp.h,sha256=0hSyQX8zGHEJeDh_OoGx-_hvLaBE6FK_p5uFNf1BmjE,18169 +torch/include/ATen/native/transformers/xpu/sdp_utils.h,sha256=TJP33v_Umxffcwa04qphK2N93jnnpWKMDp9hujoTO4c,754 +torch/include/ATen/native/utils/Factory.h,sha256=hulmOFEIX_WWiqQwSH-RwO_eeOS0yN5GiKK6mCeBPlo,741 +torch/include/ATen/native/utils/ParamUtils.h,sha256=ih8hZmb-uy6FgF43gZJJkP280JiXnh56tEzH1QYm2kM,1452 +torch/include/ATen/native/utils/ParamsHash.h,sha256=vQO1aftMAZP6KS9_10QogPZ3W9vljb3FFUz0wW3lkOo,2794 +torch/include/ATen/native/verbose_wrapper.h,sha256=tSL4joW3TAZdsqba7Y7Cauyq0Jp3G9C4116lGmdTs6c,447 +torch/include/ATen/native/vol2col.h,sha256=zAA7LrKdvrCGsuitLlcWXHMD7V75Gh0C21r6-wW17oM,3809 +torch/include/ATen/ops/_adaptive_avg_pool2d.h,sha256=VKv9hXeA4KAtFD4NCgbXz8wfdejTzHII8vcP5_D4oJA,4395 +torch/include/ATen/ops/_adaptive_avg_pool2d_backward.h,sha256=8Ubr9PuqM_27gN2Yyk03vq3RrO8Hwovqoy2VVBrxIuc,1706 +torch/include/ATen/ops/_adaptive_avg_pool2d_backward_compositeexplicitautograd_dispatch.h,sha256=mLkVi1HCGh9nSBkKESDUAqtfRaZ0Z-00rW8-UgIq554,1231 +torch/include/ATen/ops/_adaptive_avg_pool2d_backward_cpu_dispatch.h,sha256=OsC_c6UUjet4H_zYi7IIZw0sKrzC0V9ztHGk9pjFhnQ,1029 +torch/include/ATen/ops/_adaptive_avg_pool2d_backward_cuda_dispatch.h,sha256=6QcvbVmZky7Q0hbelPJyv3I10D1R0b0aaVfPp7bNaww,1031 +torch/include/ATen/ops/_adaptive_avg_pool2d_backward_native.h,sha256=r7HHxy8BF7Ou5o3Z3lS0FGx3j-7zaNeMUlVzix4fHHw,1034 +torch/include/ATen/ops/_adaptive_avg_pool2d_backward_ops.h,sha256=BhB-sEfb_DaxKIKckF_68PuVI-wCY_fOjgGMYehQG-s,2147 +torch/include/ATen/ops/_adaptive_avg_pool2d_compositeexplicitautograd_dispatch.h,sha256=k7uR9x4eqtcgWoutyVO4IR9WdCp45_k52Zjb0rLgbLQ,1472 +torch/include/ATen/ops/_adaptive_avg_pool2d_cpu_dispatch.h,sha256=lp8Ne9a25VTEcrtpsYruG9qlsSKA9PlJ320oGoEubdY,1125 +torch/include/ATen/ops/_adaptive_avg_pool2d_cuda_dispatch.h,sha256=NuLYus9z6EKbte63g72zIR_z5zHzz0VaSggJ45L_54o,1127 +torch/include/ATen/ops/_adaptive_avg_pool2d_native.h,sha256=yQB5NsFghSasLHbkDz-XawSCqyHzjQVxamL0UhK5ehI,1230 +torch/include/ATen/ops/_adaptive_avg_pool2d_ops.h,sha256=99HZh1i09c0SA5KV64nk-VTPK4vaeIEOzjW5jY6-BBg,2105 +torch/include/ATen/ops/_adaptive_avg_pool3d.h,sha256=v2a_8P7iuJ06UvQKIWKYI5_HxhaIOSnfo_RlAO_gUu8,4395 +torch/include/ATen/ops/_adaptive_avg_pool3d_backward.h,sha256=3cVh_GUjUbgLPJLRTW2e3Ta97b6bggisWhNN2k79_FM,1706 +torch/include/ATen/ops/_adaptive_avg_pool3d_backward_compositeexplicitautograd_dispatch.h,sha256=43DpFO1vxssk3yj7QYLBfdY6-jrjwVZmjrmBgPmc9GI,1231 +torch/include/ATen/ops/_adaptive_avg_pool3d_backward_cpu_dispatch.h,sha256=UM0xvlbTYxwnsDKzEOqaosV9s3cNUIWYt29gwQqOnbs,1029 +torch/include/ATen/ops/_adaptive_avg_pool3d_backward_cuda_dispatch.h,sha256=_350PXjjuL0GNI753yx8iKuZwKQ2Uh9c6zLWjih2uXw,1031 +torch/include/ATen/ops/_adaptive_avg_pool3d_backward_native.h,sha256=KCwuX9lPBO2jfosSnOiPqNJ9o8uXg8GSPXKtEDUXnGY,1034 +torch/include/ATen/ops/_adaptive_avg_pool3d_backward_ops.h,sha256=4K6JCalwtIBYas68HT8fgERDxVShrbLhhZbTP791-co,2147 +torch/include/ATen/ops/_adaptive_avg_pool3d_compositeexplicitautograd_dispatch.h,sha256=t-Iy_nYO5oxgcoAEZtM8bUZymZwR8aeBdVWHimIff38,1472 +torch/include/ATen/ops/_adaptive_avg_pool3d_cpu_dispatch.h,sha256=uJjRa0Koz3qaFx-e9SlGZ3io2is7smRWro2zH5KSxK8,1125 +torch/include/ATen/ops/_adaptive_avg_pool3d_cuda_dispatch.h,sha256=3iRhAMZ1qyofZpU8uIIKN5adBjekHe6upXAwr55SgTY,1127 +torch/include/ATen/ops/_adaptive_avg_pool3d_native.h,sha256=XtGD7YLyX-gz9KqSHVRAhk2D8PlBytuEg332Xg0APTM,1119 +torch/include/ATen/ops/_adaptive_avg_pool3d_ops.h,sha256=EXAipK5kpqItDp0H81q9hS8TnA7goCYh-iwBeZErOyM,2105 +torch/include/ATen/ops/_add_batch_dim.h,sha256=VIz0gGHwr_Sye4buLvGsa2nTXtotVQNl1qne_Bub-1E,1001 +torch/include/ATen/ops/_add_batch_dim_compositeimplicitautograd_dispatch.h,sha256=ZfoHh_My_mxOKUlNz8nNDG1aFL_oWJiLtMwwV_Xay-A,1060 +torch/include/ATen/ops/_add_batch_dim_native.h,sha256=h2Kk9wSX8FFzZCjJh76H6TyJg9gXcP_bQIcpkn8T2zs,772 +torch/include/ATen/ops/_add_batch_dim_ops.h,sha256=0y0o7r6MA1vnR77n2LsW5ukEaHIqxx-uzeE-ZMD8DFs,1347 +torch/include/ATen/ops/_add_relu.h,sha256=vnceR2sZ-FMjyE00C_I_vwd_SKNHuKrtT3XjEuFloaY,3035 +torch/include/ATen/ops/_add_relu_compositeexplicitautograd_dispatch.h,sha256=hUNyJx_nMiZwYeb5i75UP5uLcN9UFLIjxsx1jVHl0Oc,1233 +torch/include/ATen/ops/_add_relu_cpu_dispatch.h,sha256=REC2x8KKagTEtAwbtfHN4RtIjXIONz5eelHVd86rG-I,1627 +torch/include/ATen/ops/_add_relu_meta_dispatch.h,sha256=JLWkg0N4MNaLrhV7J0PNevRi3P2XVcA56XSNDZRLKfc,1138 +torch/include/ATen/ops/_add_relu_native.h,sha256=ptS5iBg9uD4tzaRHGaOtFCIEC4eYGx8mwqoAyn_5PBA,1382 +torch/include/ATen/ops/_add_relu_ops.h,sha256=0iMn3D4ZwCoKsi_9CDvGIpkfZOiLnApPTDQ8hDFnJt8,5094 +torch/include/ATen/ops/_addmm_activation.h,sha256=NFXY6zE8pZJ53l0hwyQTFPJPbbD9fVS4iNZNQUwCywQ,2101 +torch/include/ATen/ops/_addmm_activation_compositeexplicitautogradnonfunctional_dispatch.h,sha256=zc4M5jyoTUOOvCUdQ7DFlFzvsy8hd-u8sHt0ahifELo,1181 +torch/include/ATen/ops/_addmm_activation_cpu_dispatch.h,sha256=QFoZ5pSn68aAgyxyCXBq6CGRCEhBVIESOCloETpWfg4,1532 +torch/include/ATen/ops/_addmm_activation_cuda_dispatch.h,sha256=x-ArIRqgPH9MaWVhIWPt2PPgTuYJucA2NjkQwLa5Bnk,1534 +torch/include/ATen/ops/_addmm_activation_meta.h,sha256=zQs0DP6jEuXc131wioxiuLSx7f57u2GB6xYySqzErbQ,948 +torch/include/ATen/ops/_addmm_activation_meta_dispatch.h,sha256=LJ-qSa6F7dtiale1qdWtblQDpyNQ0qCohC7DI0bcUZw,1534 +torch/include/ATen/ops/_addmm_activation_native.h,sha256=9KrUGdrLmXNecFwyiGaXDeU9XMBlAA0ostjW9jnQmHo,1285 +torch/include/ATen/ops/_addmm_activation_ops.h,sha256=cDCyYLEx9WwDg1ymxPYEUvB_FdmXqphuzTD9Af138-U,2662 +torch/include/ATen/ops/_aminmax.h,sha256=_JKA35j3Hw6HhCJqbcqHKACzDK-dEaerEL_nvQnLbAU,2504 +torch/include/ATen/ops/_aminmax_compositeexplicitautograd_dispatch.h,sha256=-UB2ENjtdZhPT8vTj4fd-XWhNI444dgmGccSVJtNxEk,1534 +torch/include/ATen/ops/_aminmax_cpu_dispatch.h,sha256=JfDl8NSKZ_KuT8advdH0RUiGvkQroIbrypcYO0zoBhM,1115 +torch/include/ATen/ops/_aminmax_cuda_dispatch.h,sha256=cbjKuLphT7N66jabChNFqo2oWbbFKAyg8t5LIiePLf0,1117 +torch/include/ATen/ops/_aminmax_native.h,sha256=M48NN_OtbMGf9bYDJEXlsHyBDWfrMm4edOtFaMR5YpI,1160 +torch/include/ATen/ops/_aminmax_ops.h,sha256=xW39vDxpi23hwdnGEszwL5oRNUY8CNFnJxh2ijS83_M,3638 +torch/include/ATen/ops/_amp_foreach_non_finite_check_and_unscale.h,sha256=e8CuktcHI7bn87uBRzuJtg5WQSO76KpjBrWzHOlA_Tg,2377 +torch/include/ATen/ops/_amp_foreach_non_finite_check_and_unscale_compositeexplicitautograd_dispatch.h,sha256=VXLsaOjSRF2B2xQph85eMtcePFfWzGhVZukvnTjG2II,1464 +torch/include/ATen/ops/_amp_foreach_non_finite_check_and_unscale_cpu_dispatch.h,sha256=TfYb0wqSIp7wJvieB9xHFAGsFhlKi8nGgRTk2JVCSAE,1054 +torch/include/ATen/ops/_amp_foreach_non_finite_check_and_unscale_cuda_dispatch.h,sha256=mGRzwu0aolPlbO2Nho0ibcuLp86DH54hsB38PyH92w0,1056 +torch/include/ATen/ops/_amp_foreach_non_finite_check_and_unscale_native.h,sha256=X5k13v-pQX-wkr3HSkBRbrPuplybZ5W8OjVlEwkaVsU,1295 +torch/include/ATen/ops/_amp_foreach_non_finite_check_and_unscale_ops.h,sha256=PH3IpIBYawlq9GeRIIMUxOGm4_q2yQE32TxLvHrY0vM,3262 +torch/include/ATen/ops/_amp_update_scale.h,sha256=m5NSz43QaTDRqAyEvKOya71WqTwK5mdUloeH0puemx8,3018 +torch/include/ATen/ops/_amp_update_scale_compositeexplicitautograd_dispatch.h,sha256=q-KqXU0UiXyiHAvsA91rq8ZUYka3gtrHxWi6mKve1iI,1662 +torch/include/ATen/ops/_amp_update_scale_cpu_dispatch.h,sha256=7N15m6iQbPylF_iCg-KTrs-v0WvXI2LydzfinzIM90o,1123 +torch/include/ATen/ops/_amp_update_scale_cuda_dispatch.h,sha256=yNcvckOdIf88VDGvYUV0ZfBbJ6sTEq5Wzwl4RgAGPGc,1125 +torch/include/ATen/ops/_amp_update_scale_meta_dispatch.h,sha256=8VDbSMIA53C_Wex35_MHAGxgyHoTnzk7MY2Udclp54s,1125 +torch/include/ATen/ops/_amp_update_scale_native.h,sha256=-3IT4l-Se22BUw52HnDiT7knBPUR9Fkiu3rGcidtXv4,1558 +torch/include/ATen/ops/_amp_update_scale_ops.h,sha256=7pt7ckvF4x80pmTtcZI0sF_l0llEWq6D16HCnxkf_4Y,3903 +torch/include/ATen/ops/_assert_async.h,sha256=NU64JdJs9kO2GaC3hDE8oJJAqkzPU3rjWDeLGaq-I4c,1120 +torch/include/ATen/ops/_assert_async_cpu_dispatch.h,sha256=8TWiV4wTxBGib2X21PnqM6xcyyBKpTATUiCspwkXo9I,1059 +torch/include/ATen/ops/_assert_async_cuda_dispatch.h,sha256=hb89YSZIwi1MDNAWXDaIaNCKKopZnmOlF_Q4mWu8GPo,1061 +torch/include/ATen/ops/_assert_async_native.h,sha256=Suv-5Ls1siFxFXjhIIHMjl10UXTM5Wxx143C2mQ3FcQ,980 +torch/include/ATen/ops/_assert_async_ops.h,sha256=qlkmiPuJ7jmlQPWenFdY2WGtWYGki3WjPkeWI8OtFsw,1812 +torch/include/ATen/ops/_assert_scalar.h,sha256=cuff-IVFPf2VkAJ-hpG9D8MRCmrc8Dq3vQjr4oISYIg,970 +torch/include/ATen/ops/_assert_scalar_compositeexplicitautograd_dispatch.h,sha256=TtEYQMZ69qKusuxsf-aeOjJUP31rCsguWCOkI5AV4pI,1049 +torch/include/ATen/ops/_assert_scalar_native.h,sha256=XTx6IF8VVQ-cq9oZCDlLwmPbg4BNjtSGiaRPrnCVS1I,761 +torch/include/ATen/ops/_assert_scalar_ops.h,sha256=jU1JJNn4_3UGgtEJnUqtzM-nQquVmUM-MMmTXlnoSno,1305 +torch/include/ATen/ops/_assert_tensor_metadata.h,sha256=cXDhhq656r4YDYdFinLYOn9eu3Z13H_kloP5GEg0XvQ,3272 +torch/include/ATen/ops/_assert_tensor_metadata_compositeexplicitautograd_dispatch.h,sha256=qQkLHdbYoBWq3jM7pYTD9Zu_nDffN68TEL23cDp0Pm8,1597 +torch/include/ATen/ops/_assert_tensor_metadata_meta_dispatch.h,sha256=pA9v-ov8VAO5wUQ-grfgzxy7x638CTqDyIN11LuBexw,1555 +torch/include/ATen/ops/_assert_tensor_metadata_native.h,sha256=7h-X_m5RsMAC3QPX7fyrEVio1fFmIA_6X4cRPluwHoE,1314 +torch/include/ATen/ops/_assert_tensor_metadata_ops.h,sha256=t4gkCXNhLgte-i0jz5ATfwzzZinlnVBq0YPS0Xt4XsY,1849 +torch/include/ATen/ops/_autocast_to_full_precision.h,sha256=3F4Q2TewdBXsSFSg6MTL4AVbvLlq4-mqQZ6qElUvQh8,778 +torch/include/ATen/ops/_autocast_to_full_precision_compositeimplicitautograd_dispatch.h,sha256=ZV6rFidwsaHXWOWQNPPkpkYZbGny9fPDeoEvsbo1lhY,1076 +torch/include/ATen/ops/_autocast_to_full_precision_native.h,sha256=5G8Jeejaliid0n7xxs8i3zeyQN-NonRM44IfjD8pJSc,788 +torch/include/ATen/ops/_autocast_to_full_precision_ops.h,sha256=_wouMqdvfFJL4-TrK7eLXcyTCri_u33NpQmUmT2a_es,1403 +torch/include/ATen/ops/_autocast_to_reduced_precision.h,sha256=cgqPSA-IAKKYs_xsrri3-UburVPSvZJWxhfeO8POkWE,781 +torch/include/ATen/ops/_autocast_to_reduced_precision_compositeimplicitautograd_dispatch.h,sha256=i13BWqXPk8hFtzXQNZg7YsSmj9gMzwc-34_l8O5GvdU,1132 +torch/include/ATen/ops/_autocast_to_reduced_precision_native.h,sha256=MFscclOIoT5hAlHdZ2ewIyd0gCzFbz8icgT5pzqO7qY,844 +torch/include/ATen/ops/_autocast_to_reduced_precision_ops.h,sha256=G-f_alHvPgQjuPz5O3EAWBzITSTrcPilnaPzMxyOMvo,1595 +torch/include/ATen/ops/_backward.h,sha256=BW07CJoOliQbp9F42ephO8cNvgorce3PNcOs3rPNA48,760 +torch/include/ATen/ops/_backward_compositeimplicitautograd_dispatch.h,sha256=jA9epTU1M9w1Evm03JmVCoJaQuF7rnacxhpvOWKXGHM,1163 +torch/include/ATen/ops/_backward_native.h,sha256=-AbG5eqnAl3djIXSiRej_apsXFGRYUb1nEyP45S3GkE,875 +torch/include/ATen/ops/_backward_ops.h,sha256=xm_Nte-p1bAnRZZgtubHW9C2jNRb7Z7iY20P-4VS4tw,1618 +torch/include/ATen/ops/_batch_norm_impl_index.h,sha256=f-cwCC2OOt3JCgQglZmh0PQMzRl7XlKMGViMMs5Od78,1509 +torch/include/ATen/ops/_batch_norm_impl_index_backward.h,sha256=7Bp5xP1DZIC75Gs_rXx5mCQMYr2Cmwcgldagep6c4FQ,1777 +torch/include/ATen/ops/_batch_norm_impl_index_backward_compositeimplicitautograd_dispatch.h,sha256=2BJgAze5zGVOeP4HQqulXhKipOJ0pKA13qPze50RmoM,1470 +torch/include/ATen/ops/_batch_norm_impl_index_backward_native.h,sha256=xUnZjUpfRL6POftgbtMH2t9VI03OUtkMN0ssS-vdtsI,1182 +torch/include/ATen/ops/_batch_norm_impl_index_backward_ops.h,sha256=2r5mtBrLVtKMpAjvylJWAIIhT438oBVYVcvV6F7hO3M,2673 +torch/include/ATen/ops/_batch_norm_impl_index_compositeimplicitautograd_dispatch.h,sha256=Ewr-XGnFSikwRYpyGOZk7frC9eYHd4dNAsgXqX5ZCPc,1339 +torch/include/ATen/ops/_batch_norm_impl_index_native.h,sha256=tAoprCXpkL45Y-SL5mOVMYXH6dRuodGSvpdwP-0TZzY,1051 +torch/include/ATen/ops/_batch_norm_impl_index_ops.h,sha256=PBqZ6dNJtPu5W8llYfCmbzPpCBmY12o2ULUWiO2TWVs,2267 +torch/include/ATen/ops/_batch_norm_no_update.h,sha256=0q26fQ4TkGIFBOQ-1rXmiqQhbQKFK1qD0XF-tunPhWs,3063 +torch/include/ATen/ops/_batch_norm_no_update_compositeexplicitautograd_dispatch.h,sha256=HAxsh33mQ4H9sQSN7DQqyKXJvjvEC02aGTyOQb9JxQo,2134 +torch/include/ATen/ops/_batch_norm_no_update_native.h,sha256=rAUyAzEl9xsqjUaEHtQ3raTYnUNaBlBziQWwW-MC8fM,1426 +torch/include/ATen/ops/_batch_norm_no_update_ops.h,sha256=TILf3FK6lm-ASfHLmEy1XnWDAcrz05Y0JDzBKQX48_E,3855 +torch/include/ATen/ops/_batch_norm_with_update.h,sha256=N5fMIXovPnpBDuIGeCqG_4o0mjgtiqr4bGe_8VV9lqQ,3747 +torch/include/ATen/ops/_batch_norm_with_update_compositeexplicitautograd_dispatch.h,sha256=aEJAQBBW8f33iIXBeiuGCLxYYxrhCnVmSaOlwETCWTU,1296 +torch/include/ATen/ops/_batch_norm_with_update_cpu_dispatch.h,sha256=wuN1YbE3ujiNhjHf5YYDGJztIOEiklZpOSWGthNdPQ0,1986 +torch/include/ATen/ops/_batch_norm_with_update_cuda_dispatch.h,sha256=Rp_AgdR8ymgAexMLxz4rq35snL7UXDSmsz7WIX60QBA,1988 +torch/include/ATen/ops/_batch_norm_with_update_native.h,sha256=5mzx_xP-eoAOx77l3X399JlsntboLI2ZD9DdsDikwB8,2672 +torch/include/ATen/ops/_batch_norm_with_update_ops.h,sha256=IVRgqFcs55zMdv8aTuF35iEkMOQhKnaLih7S75i3Q_U,5102 +torch/include/ATen/ops/_cast_Byte.h,sha256=iUEt6J2gNyETwvNPufNb3QXf1B0Lk82uI9IbBYJre0c,971 +torch/include/ATen/ops/_cast_Byte_compositeimplicitautograd_dispatch.h,sha256=GAMd54TpuYLJqla4jkgSncD36gAk9jrZC5oF4gAmJ9U,1047 +torch/include/ATen/ops/_cast_Byte_native.h,sha256=r3Cz2B1FnuwiG3OYBrNMwqrI21-O0E8P1Tt9NuMC8-4,759 +torch/include/ATen/ops/_cast_Byte_ops.h,sha256=uw1XDJr8AEx3Y9cKkB2TRJxVpV3dCegCuVaYUL_7_Fw,1292 +torch/include/ATen/ops/_cast_Char.h,sha256=W8CCak9jEOtwzrrxHEHjmRA4ZRF0pXJSEIMCuKhUddg,971 +torch/include/ATen/ops/_cast_Char_compositeimplicitautograd_dispatch.h,sha256=QCN2swYd6xNMS2rFKyre_D4tdZ8Ys-e5dych1VyHHL4,1047 +torch/include/ATen/ops/_cast_Char_native.h,sha256=tmhgYowjNN_Z4h8JL9DeVce9zq5xtWwSu8Qt6KD_cEc,759 +torch/include/ATen/ops/_cast_Char_ops.h,sha256=IJJww6lMEhtzi7AaE9KJ8oXfwqjJlx8SSKSP2u2vDwk,1292 +torch/include/ATen/ops/_cast_Double.h,sha256=qfyoNBLnGsT6T-w7U_5d8o_v3so49U-E7yn4du8mnFg,979 +torch/include/ATen/ops/_cast_Double_compositeimplicitautograd_dispatch.h,sha256=MSiiUBlIJjbi-9Vi1i9mbrco6KKsh1zXG8hHemVUQlU,1049 +torch/include/ATen/ops/_cast_Double_native.h,sha256=5UQdbXBOdrut4ZBXlcICD5zplBm1YQoT8HHHLrM1nTU,761 +torch/include/ATen/ops/_cast_Double_ops.h,sha256=ZsbiOn9mzLE8db7ubpoVpWvtCFoKqc1aFNUfAz2_n6o,1298 +torch/include/ATen/ops/_cast_Float.h,sha256=tATJSdY3Alki6EbaIOZROvlWrOICxXbSiSM-BqJh5f4,975 +torch/include/ATen/ops/_cast_Float_compositeimplicitautograd_dispatch.h,sha256=OWeqPDgm426C6d3TjMJpuU6XYqk46qV9mu9TCcwHGQI,1048 +torch/include/ATen/ops/_cast_Float_native.h,sha256=pdAGU-5aIjQv8NzKIeXlTU1Y6o3A8aMIW2zOAlNGQs4,760 +torch/include/ATen/ops/_cast_Float_ops.h,sha256=p5kNK97whfzBZuCso0s-l_sYOy9eXR4TZ_Kfqmk5cOg,1295 +torch/include/ATen/ops/_cast_Half.h,sha256=TaqI_D3YLGhwLX6zkNz8Ioo7E988i_L0aBNx8KhbFbw,971 +torch/include/ATen/ops/_cast_Half_compositeimplicitautograd_dispatch.h,sha256=Daf1_C6LfpVL7nr16k_oatRAk_iupE4LU84-HNBVKu0,1047 +torch/include/ATen/ops/_cast_Half_native.h,sha256=koynLsIimWLnvFTbh_OUtbb-j4Sk_YN2W7_p1erjdP4,759 +torch/include/ATen/ops/_cast_Half_ops.h,sha256=Cnp7hvpq5cCmyGEbYXY_V7i10SIirxxF3nAeLIKUWKo,1292 +torch/include/ATen/ops/_cast_Int.h,sha256=_i0h3Hcwxovwfnb80Y7_SxsU-fMCV-d9uZM3QdabIDs,967 +torch/include/ATen/ops/_cast_Int_compositeimplicitautograd_dispatch.h,sha256=qyJwe6vAd06IDjhj3IFdRj_csh4jLT6vnzllcbvuCt4,1046 +torch/include/ATen/ops/_cast_Int_native.h,sha256=fU7d7ep-emIz4Wj8auNcBgjgRCgm68UDvguodGoYCVE,758 +torch/include/ATen/ops/_cast_Int_ops.h,sha256=Gnq8WunEMadHatpIRAb972UB3aw1f891pBHKujTTOnU,1289 +torch/include/ATen/ops/_cast_Long.h,sha256=VZweppx4HFvooEZpt3K-8-PsRyEPSV5PP98OugsNjtY,971 +torch/include/ATen/ops/_cast_Long_compositeimplicitautograd_dispatch.h,sha256=JDWTNIDQ_BB4MFnSOuudgxHb4dF1Uq2DepBI4KCUk_U,1047 +torch/include/ATen/ops/_cast_Long_native.h,sha256=8YBfwMAvxCeQXyXDgdP3-Ix8XZHzZokD135VxjOijMc,759 +torch/include/ATen/ops/_cast_Long_ops.h,sha256=QszCLEyc2Ez62bHrt0JlEUfxibnaUZMs-2TCLkfHPQY,1292 +torch/include/ATen/ops/_cast_Short.h,sha256=GgjfVe2S0RgVpZgHo6Kau1qbwNiE9DfWFkA0l0pavfc,975 +torch/include/ATen/ops/_cast_Short_compositeimplicitautograd_dispatch.h,sha256=VLVdL13yRlMSxSJZyG-1Xr4Ytct4rqObuz-G3_xUXsE,1048 +torch/include/ATen/ops/_cast_Short_native.h,sha256=It3CgornyGV7rnn9TZzwh00Jy6m7ToEubFrsyNbLf-Y,760 +torch/include/ATen/ops/_cast_Short_ops.h,sha256=79_EaptGAYz4X7YfPzd-QC1hzm9RG8gT2t2dh8RW08E,1295 +torch/include/ATen/ops/_cdist_backward.h,sha256=U-V3i5Lizw0Gc5xDZxygsa7p3uwjgB7syKcgFuLwdRQ,1806 +torch/include/ATen/ops/_cdist_backward_compositeexplicitautograd_dispatch.h,sha256=TVX87qwBiRVjqROxup6FZTlkJq9GiUlhqNyemwqice8,1303 +torch/include/ATen/ops/_cdist_backward_cpu_dispatch.h,sha256=AYYZyPELdw-24gl9NrNETvN4OslT2ebXduzBDOVfsnw,1065 +torch/include/ATen/ops/_cdist_backward_cuda_dispatch.h,sha256=7FulepCGFLuBBIKMSIPlT7SrYdjuN9TaBhT1xhLE3VM,1067 +torch/include/ATen/ops/_cdist_backward_native.h,sha256=XvmWKGrIkXSx1vZ8je_a7eGkhoPIAnDjOEGrdxe-NyM,990 +torch/include/ATen/ops/_cdist_backward_ops.h,sha256=WaGGYd8GQdyFA9A_R4nmOo4ERAHK08yKA2el3sfd-_Y,2409 +torch/include/ATen/ops/_cdist_forward.h,sha256=KZ6DmmHw3jPvtB4VfYYGT0nMshTKcveu0LbPa4N4GQY,1739 +torch/include/ATen/ops/_cdist_forward_compositeexplicitautograd_dispatch.h,sha256=UgQS0ZUGuf1ats_zIuSZdP3_-LyY31AeS6zWPhn3SIo,1277 +torch/include/ATen/ops/_cdist_forward_cpu_dispatch.h,sha256=8FlF-PXLzmcNcycd3w3fYKgjg-vri4Rx_KHkqp2-F3A,1052 +torch/include/ATen/ops/_cdist_forward_cuda_dispatch.h,sha256=Kf6jo1ZPTPWiAQLmrzV5gOeuPMDhBOTAIVKvK0YCMNo,1054 +torch/include/ATen/ops/_cdist_forward_native.h,sha256=-MG8DF1QwsiAjRyiqCDwul9HJlMYbwxgkJVsLashk44,964 +torch/include/ATen/ops/_cdist_forward_ops.h,sha256=hzMkV3q3BzgWMwPV5mqS4S4ylU2M5-IpmHVxhFx_uCk,2311 +torch/include/ATen/ops/_cholesky_solve_helper.h,sha256=WSC4_9F4Bbap4noFpACrORmdYLxiG-CJJOpoACP1uD8,1639 +torch/include/ATen/ops/_cholesky_solve_helper_compositeexplicitautograd_dispatch.h,sha256=BO5o1jS9T-kkw9ZdcO_9rer3ha3YLXSBz2ag7CUWRJ8,1221 +torch/include/ATen/ops/_cholesky_solve_helper_cpu_dispatch.h,sha256=ObB-pLnODn4w45B_mfDplqdqj01b62iItFJ8z3Y4Y1E,1024 +torch/include/ATen/ops/_cholesky_solve_helper_cuda_dispatch.h,sha256=zOBtWBt10BPpa_6jylzVjFCReQDaGfK-CcSTx6wSTmE,1026 +torch/include/ATen/ops/_cholesky_solve_helper_native.h,sha256=yeiQt30yfDz94JoDHqzSy9eDJgLhKzOVg2ILDAIhVcQ,1021 +torch/include/ATen/ops/_cholesky_solve_helper_ops.h,sha256=xJqu_YSMiUtyBKhaDL5w0uF0v6HjtRKbWwi_ft9c7uo,2129 +torch/include/ATen/ops/_choose_qparams_per_tensor.h,sha256=B8FYNi4YNh5HVA6orsR7bBHngPmiHsrofnaFvjXWP0A,1059 +torch/include/ATen/ops/_choose_qparams_per_tensor_compositeimplicitautograd_dispatch.h,sha256=RJKzWLBY7okEhWlnEb1F_nx4kuHt8dlzufG-niQ_3gg,1081 +torch/include/ATen/ops/_choose_qparams_per_tensor_native.h,sha256=vAQvomoawvPD0d7jpIjwNbbhVrGXSwToK_e5ZZ2Ygv8,793 +torch/include/ATen/ops/_choose_qparams_per_tensor_ops.h,sha256=y57Z7uuQUQF2cG_UBZPLi7hmTNHorPvE2gBnYEXrAm4,1400 +torch/include/ATen/ops/_chunk_cat.h,sha256=T9pTGfyJy52r7bgPJlpOsRVmbBKWO9Su_1OV8l8zYYE,1567 +torch/include/ATen/ops/_chunk_cat_compositeexplicitautograd_dispatch.h,sha256=epexgL_T4kVC7Hi8cFbmsmtE08AaK4NQn4NtdUxIw2s,1283 +torch/include/ATen/ops/_chunk_cat_cuda_dispatch.h,sha256=VTrX11jLlT5B9zTPYTD49ZLwQ9LC3N-fHc_tcDTBb8k,1241 +torch/include/ATen/ops/_chunk_cat_native.h,sha256=e7J_6_g0BvWrt93V3-yid2-3RQe_kgVHu-pyxaSG7Ic,1094 +torch/include/ATen/ops/_chunk_cat_ops.h,sha256=CWggP0xIMY8ssvXBfyIo7DACCgC6CSpq7Nblp-4s0-k,2041 +torch/include/ATen/ops/_coalesce.h,sha256=N0BSrvTZjla5WLys7UM5t0KkaOIMEXZju3kVyWhdpDI,1311 +torch/include/ATen/ops/_coalesce_compositeexplicitautograd_dispatch.h,sha256=Hwc44Lu6xmDsoNlFToxCyooWSqO4qkEXeBIBO7Npqu4,1127 +torch/include/ATen/ops/_coalesce_native.h,sha256=IZzOMWVpLqWR3S1R_pxXZaeiOsgvqQ-WmWg6UcQW0Ng,894 +torch/include/ATen/ops/_coalesce_ops.h,sha256=182XiGHaPEqjvqZVVpVh-7041I6k-_04PJCIOAg3cw0,1819 +torch/include/ATen/ops/_coalesced.h,sha256=6upz2I_Cz6AAj-noGFu9fytzrasEoHmcEFfSG68eisQ,1450 +torch/include/ATen/ops/_coalesced_compositeexplicitautograd_dispatch.h,sha256=bQAOq--qjPQXpU458gmtCIjRv--1tm4XQUigqo6vqbI,1235 +torch/include/ATen/ops/_coalesced_meta_dispatch.h,sha256=uoZYnRD-za2B4bk6wcPbPORdR8_OGloko6sYx-a18Zk,993 +torch/include/ATen/ops/_coalesced_native.h,sha256=dbrXtbdzfRU3cu3oNWd6CYKSzkDdsIDMgg7i-UTKlHo,926 +torch/include/ATen/ops/_coalesced_ops.h,sha256=hgewfg3wKaUot6UiGPjWg78VhXXwOO8MhmJikYCrg-8,2498 +torch/include/ATen/ops/_compute_linear_combination.h,sha256=8wGZlGnNfIWISMCvVzlL7ROCjl47pdwCqTQemYqLTEc,1704 +torch/include/ATen/ops/_compute_linear_combination_cpu_dispatch.h,sha256=Ajtd3CjthKwJ7V3qtszHnv0FuF_hlxYnDNSu-Iqs7jo,1296 +torch/include/ATen/ops/_compute_linear_combination_cuda_dispatch.h,sha256=n3HM0-oJNiu5YSsbZ6YG8I8VgTIH4JjhZ8jE1yyBvcQ,1298 +torch/include/ATen/ops/_compute_linear_combination_native.h,sha256=za2Y3ia5OKIs36aqoDm1ag4Dbsjb3C8kt3s_Z30Gnno,918 +torch/include/ATen/ops/_compute_linear_combination_ops.h,sha256=5jNsiFzBxgyO9_odc7t1rhyTzuud2mtonfsG9oThVXg,2147 +torch/include/ATen/ops/_conj.h,sha256=Vqqr4_E1xNeHHWsjX6zD5djPF0Xax7K-aDugkARvXVw,893 +torch/include/ATen/ops/_conj_compositeexplicitautograd_dispatch.h,sha256=4OMPve0UexbWAY6l-772oFDMRBkiSbsdRzvioc4IXBk,1017 +torch/include/ATen/ops/_conj_copy.h,sha256=YhalmUTfwScRx01_Ywimk54DfAryQyWWdrFWJN9XKms,1321 +torch/include/ATen/ops/_conj_copy_compositeexplicitautograd_dispatch.h,sha256=EJc9ItHLNPpT3LSrqgX7PVEfvD6qT3zYODWbmMar30Y,1129 +torch/include/ATen/ops/_conj_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=CsDc8Mfo_u67sRUvY9K4tWqghj8hhltsiyODNkkw_Pg,1048 +torch/include/ATen/ops/_conj_copy_native.h,sha256=9f_Okj59riaDKjCf-0kYFu0o4xb_IsFG7A48r6Iq_b0,816 +torch/include/ATen/ops/_conj_copy_ops.h,sha256=D3FDjoHVKA4ftBmATHt-ZX4q3dfiHcUWgiM2sn0OHRE,1825 +torch/include/ATen/ops/_conj_native.h,sha256=cSdPVDZuBSqhVVlgYyL7uBIhg1d95B-J8PhC-wEPbbI,729 +torch/include/ATen/ops/_conj_ops.h,sha256=5FxQr6pCdhZw0fHu-A0LEMvh4q6CVTY0A9zi7dKJWj4,1214 +torch/include/ATen/ops/_conj_physical.h,sha256=UAzAku0g6WDKOs19vYlnNCOWjiRO6phZsMpVC7Xp9Jo,1361 +torch/include/ATen/ops/_conj_physical_compositeexplicitautograd_dispatch.h,sha256=M8f53piexBsVfRHMrzQE3sIzHbVfBXYAlfg-fcDcUiU,1199 +torch/include/ATen/ops/_conj_physical_native.h,sha256=hGW_VDV9utF3Y4DB4SP4FVqEstX82rQJVOfvgrKsQoQ,896 +torch/include/ATen/ops/_conj_physical_ops.h,sha256=121AJUJBjTeTuw4qPQfhhXC5lH8SWW9tRPhKtblECNs,1849 +torch/include/ATen/ops/_conv_depthwise2d.h,sha256=C0OEfvrPKN24-7iz1PVVb-nbajfwmtko7py41yx4Q04,7578 +torch/include/ATen/ops/_conv_depthwise2d_cuda_dispatch.h,sha256=9Uulq3grbRQKSzhRLMAyfJ6r32iAn7E9rg_F2ECCP_M,2517 +torch/include/ATen/ops/_conv_depthwise2d_native.h,sha256=TL4CMRMhGgk8vvIoeeKpkvZ9sJiN9ifBMBwdQwbVumQ,1184 +torch/include/ATen/ops/_conv_depthwise2d_ops.h,sha256=yYdFvBUWT4TWShQLvbDHSMmqCcj0ajzYv6aMTUt-XnA,3123 +torch/include/ATen/ops/_convert_indices_from_coo_to_csr.h,sha256=y6F5QIzN4lmg88sOsLH-AkL7Orbx-j-EhYl2MRTojfw,1793 +torch/include/ATen/ops/_convert_indices_from_coo_to_csr_compositeexplicitautogradnonfunctional_dispatch.h,sha256=VmRCUDmcdWiwHnyO89Q0-0KFFSKTTxjXHa1jED7qj3k,1106 +torch/include/ATen/ops/_convert_indices_from_coo_to_csr_cpu_dispatch.h,sha256=5ZLNV4BXw96cIfwPGyO9CguDVHrq521NOnFEn_zjgiI,1311 +torch/include/ATen/ops/_convert_indices_from_coo_to_csr_cuda_dispatch.h,sha256=uiCx_mveYp1ax36n22hj-r4_tc7CQWCJq9NnG_mNpfo,1313 +torch/include/ATen/ops/_convert_indices_from_coo_to_csr_meta.h,sha256=rvq_cGMN7RQoAGSrq9n6DoukTAlHCW6sTY2CNpDxcmY,877 +torch/include/ATen/ops/_convert_indices_from_coo_to_csr_meta_dispatch.h,sha256=B0k5XWC6aDpEwJ77xLG2u7jnHFyPZRALFHIJKTqjWls,1313 +torch/include/ATen/ops/_convert_indices_from_coo_to_csr_native.h,sha256=QWU5T6PNeQBX1BthUMVammhscFQaN9cyzlAAZBqtqcg,1204 +torch/include/ATen/ops/_convert_indices_from_coo_to_csr_ops.h,sha256=Sm5n8A2hp3wijgyR5nWiJ3eSbVPVRz4VijSRvWU0Ldw,2174 +torch/include/ATen/ops/_convert_indices_from_csr_to_coo.h,sha256=OD3B0dwhrRbq1itJQCDaJgWFXs-9pMgW-DdYKGdMK8c,2129 +torch/include/ATen/ops/_convert_indices_from_csr_to_coo_compositeexplicitautogradnonfunctional_dispatch.h,sha256=OFI4r1IM6gYL2h4JCpPY6IRP9lQeJXqbLbQYdicR6Io,1154 +torch/include/ATen/ops/_convert_indices_from_csr_to_coo_cpu_dispatch.h,sha256=PtqQQ5ubZTicTR3kKiaq4JaZXmsJF4hDfLaYh0HQDAg,1449 +torch/include/ATen/ops/_convert_indices_from_csr_to_coo_cuda_dispatch.h,sha256=G8ScbUetSbvFVF8KdWrbfxHITiJRT6RTbmLGRbKOZFY,1451 +torch/include/ATen/ops/_convert_indices_from_csr_to_coo_meta.h,sha256=pxEZUmYTv3XycK7htukyU_bdTD60FgF07sqFbZsmfV0,919 +torch/include/ATen/ops/_convert_indices_from_csr_to_coo_meta_dispatch.h,sha256=Vm331Z7hvxirfJqdEL6u-yqpsFexuH-rFxiljX3tykE,1451 +torch/include/ATen/ops/_convert_indices_from_csr_to_coo_native.h,sha256=-LZJCtnRrdBLa_fh8xvNfTBLf7zt5h_d-yKlWMN-rWk,1288 +torch/include/ATen/ops/_convert_indices_from_csr_to_coo_ops.h,sha256=q43U78gVKNtouUgPVdjU2Nc_0N8ltIc4HVbTowJEI3k,2456 +torch/include/ATen/ops/_convert_weight_to_int4pack.h,sha256=clAs63Ym8Z6fS3RV8yWCkVEK8mDwiMObipkHT7vMPkI,1026 +torch/include/ATen/ops/_convert_weight_to_int4pack_cuda_dispatch.h,sha256=gJmUQtj4btzKA71uQALoSxcB5SwejJG2Q3xO52a-ckI,1018 +torch/include/ATen/ops/_convert_weight_to_int4pack_for_cpu.h,sha256=PJ2yIS77nZIuJqeA3BcjL34HWF4ysq_Y6tGAPuF9rQ0,1058 +torch/include/ATen/ops/_convert_weight_to_int4pack_for_cpu_cpu_dispatch.h,sha256=crWsYSxqmoCOgj5VOWoUkp0BrW78R2p415tGkKKqbNw,1024 +torch/include/ATen/ops/_convert_weight_to_int4pack_for_cpu_native.h,sha256=fofPdIl8fBJmdddCQzyaWX5VdRTwHnmRwubMte4H8Iw,776 +torch/include/ATen/ops/_convert_weight_to_int4pack_for_cpu_ops.h,sha256=Gxq8mfMMczVVXCyf6VqX4ZQRSXQkW7gHRRc66fxleME,1366 +torch/include/ATen/ops/_convert_weight_to_int4pack_native.h,sha256=sHaYxqPr2SWJb74zbpqL2ECuKgYS8DM0duAQvoOZoQ4,777 +torch/include/ATen/ops/_convert_weight_to_int4pack_ops.h,sha256=Y2bhIULr9GBS96v9m3HKWvZ2pVT1--siycwbmwVLJsc,1342 +torch/include/ATen/ops/_convolution.h,sha256=3auCvvpmqx0E3g1ByQB-sVS1IW-ZD0PL6iNAv0ZMjQw,13232 +torch/include/ATen/ops/_convolution_compositeexplicitautograd_dispatch.h,sha256=RvzdIzZAg8olya3lONwzCC9NhoYy-tIOMqUOV4UZ2nY,3201 +torch/include/ATen/ops/_convolution_compositeimplicitautograd_dispatch.h,sha256=n4Gk9OcpOYERvsmJ7rj3OEujPyFnkf7ROxCQM60UPKA,1639 +torch/include/ATen/ops/_convolution_double_backward.h,sha256=_d8lQqr7QeWHcW49ZgkF2UUO0ZAnt21semHW6hj0Hac,4289 +torch/include/ATen/ops/_convolution_double_backward_compositeimplicitautograd_dispatch.h,sha256=3MP4QuWcx5N9tIwlsCMWFPizxWU_9tOB6EqZPke0rcY,1909 +torch/include/ATen/ops/_convolution_double_backward_native.h,sha256=CL3sWzZYuj-thybC1EDA25szJppDWRK1kiT44Dad_TU,1135 +torch/include/ATen/ops/_convolution_double_backward_ops.h,sha256=Yf4rrQMZlDleevdEdS2ANOAslAl1CXyrKw23BTbNpsM,2610 +torch/include/ATen/ops/_convolution_mode.h,sha256=6nDUoGLHi7vHwEA6IdaZNOvGZo4oM-hoL4u0NB5_To8,2710 +torch/include/ATen/ops/_convolution_mode_compositeimplicitautograd_dispatch.h,sha256=I-SFmVpM5P64X59qw_lyf9ORMJxhijOyg_2fwCp8xhU,1437 +torch/include/ATen/ops/_convolution_mode_native.h,sha256=mrxLJmT8HQZMGokfQ2KF3eI2tOhfyOgUS7bvYKEhDo0,922 +torch/include/ATen/ops/_convolution_mode_ops.h,sha256=jXbkXUt0XhqpcV1ICWfbrbyIjjkGxIEHuqiJDuPr3nA,1816 +torch/include/ATen/ops/_convolution_native.h,sha256=6mrBYsun3KF8GVe1iQQ5SAIzeWW7_fRe3rAxle4Rv0E,1739 +torch/include/ATen/ops/_convolution_ops.h,sha256=Bd3M7Ww3z4Kxa8QKSsmLE4ZRjxkwccGlGCEX8kL2Vxk,5327 +torch/include/ATen/ops/_copy_from.h,sha256=pDJW1q2oJddNVYBbH6Khi-D7WoSraBu1TGAVyUSd30g,1630 +torch/include/ATen/ops/_copy_from_and_resize.h,sha256=8udbMt8r-DoM-Y6YFXn86kUXPUDM7CWC5ScKD4Hu-i0,1554 +torch/include/ATen/ops/_copy_from_and_resize_compositeexplicitautograd_dispatch.h,sha256=szedXNusDjSX5t-Cfxz24hTN54tB7dZjQM_33L22_bU,1199 +torch/include/ATen/ops/_copy_from_and_resize_native.h,sha256=PaLHyphbbxbd2eQ2tttTnUOiHrbrk69slap1McuzECk,793 +torch/include/ATen/ops/_copy_from_and_resize_ops.h,sha256=Q02ezptAzTTRSotMbaOPosW5f4tywYkVh8AaNce9hyw,2051 +torch/include/ATen/ops/_copy_from_compositeexplicitautograd_dispatch.h,sha256=i7_6r8WmqkVWCVKVT3UEKUZgZsfOVs2xZq2lYzVih-M,1221 +torch/include/ATen/ops/_copy_from_native.h,sha256=pt_tBlRzrTiWX2DFeXj5EOUMYVpmAE7HQ-tJTaiCfbE,801 +torch/include/ATen/ops/_copy_from_ops.h,sha256=Tsegewq3PeolieA3ynHK8vmbye625Pqivg9xkEifdeo,2123 +torch/include/ATen/ops/_cslt_compress.h,sha256=qTrF7uvH9q7PXTCMyHWuk1qL6qTgeTmZR3cJMLaJOX0,926 +torch/include/ATen/ops/_cslt_compress_cuda_dispatch.h,sha256=-4ZSi06rU-L55Wu-8LxXTGr1u2MRItxo1eTvYFqIUUw,985 +torch/include/ATen/ops/_cslt_compress_native.h,sha256=BVEnYUmbQpcKk4i1_zxkR6YUlJT4vigw2U2qIOs5UFI,739 +torch/include/ATen/ops/_cslt_compress_ops.h,sha256=D7Nqx2ZSGn8r36tFVtnLFiSwPIqSZY-blp9NN_1IqKs,1238 +torch/include/ATen/ops/_cslt_sparse_mm.h,sha256=rA6OUpZpjv6ARg30nrq9Lh4nQkmvro_qW7-n-EqXWnU,1463 +torch/include/ATen/ops/_cslt_sparse_mm_cuda_dispatch.h,sha256=V8DiN2dfGhFs20FPkM3fUg1M3RVeMvwku0E0O9cMX_w,1261 +torch/include/ATen/ops/_cslt_sparse_mm_native.h,sha256=AeTRHqCUFWaIyYf_fjTt465KRij9hzYkBiFCgG1HCis,1015 +torch/include/ATen/ops/_cslt_sparse_mm_ops.h,sha256=shSaNCBa5K1zQL-Lpy77l9BLIO_-TLnhOU-Pwr3TFIo,2052 +torch/include/ATen/ops/_cslt_sparse_mm_search.h,sha256=6GtENsF_a2eXdkQdYHcgrM3j8JwI5SGFofE3qwZYTNI,1342 +torch/include/ATen/ops/_cslt_sparse_mm_search_cuda_dispatch.h,sha256=dJGjzjCsXvBaKotlwl4dh59LIpwvvMatDzMAMeYj_uE,1203 +torch/include/ATen/ops/_cslt_sparse_mm_search_native.h,sha256=lEuNQVGKAsgM_Ciii5MKDfEmUbZQMOwPRrNCQ1QoxBo,957 +torch/include/ATen/ops/_cslt_sparse_mm_search_ops.h,sha256=zDzWYINV3cnFXCAdvJ0uuS4u1niPwzkAtbhp01aN4XM,1874 +torch/include/ATen/ops/_ctc_loss.h,sha256=hxqHNrx2XCzsh66yks_dGkvo4bEA4FcAcGzq118kASg,4181 +torch/include/ATen/ops/_ctc_loss_backward.h,sha256=u1xL7bKDbN5I5eTm8rUzI6zyuUbN_me-rLhhfQ2SydU,3564 +torch/include/ATen/ops/_ctc_loss_backward_compositeexplicitautograd_dispatch.h,sha256=2uAknJCSIhVo3l3yqkkkZVwdg1DNGrJum5d75uG5iPA,1601 +torch/include/ATen/ops/_ctc_loss_backward_cpu_dispatch.h,sha256=tQvxXUJqOC1khhpyuUIGPUWnQ6SWxecCYjIwQOVtZc8,1520 +torch/include/ATen/ops/_ctc_loss_backward_cuda_dispatch.h,sha256=1hqPL0Pc8KUXJqwBpqzgAV7dRyMuNTflUjKV8j4PC40,1522 +torch/include/ATen/ops/_ctc_loss_backward_native.h,sha256=juUH-nqtsgW7cIcVMVXgXkCAnAhofVIWRQBn5Kkko-I,1900 +torch/include/ATen/ops/_ctc_loss_backward_ops.h,sha256=YccJeFoP1gQ4g4FYGH2SgrzWNndJGCUWtJSklpt_ZZc,4663 +torch/include/ATen/ops/_ctc_loss_compositeexplicitautograd_dispatch.h,sha256=J_4ShbyWtk3RubHjuPFhcXKC9VWBxiyFCAVF7kTSnC8,2030 +torch/include/ATen/ops/_ctc_loss_cpu_dispatch.h,sha256=KRPTeCeJ9nWdgK-FIsvmqvtJLCmb2JQEzJF4ccKHRTw,1368 +torch/include/ATen/ops/_ctc_loss_cuda_dispatch.h,sha256=E_GQ51VdoSHuIIKCyK7LxSCgzWZJRzoQY9-jnVB6DPk,1370 +torch/include/ATen/ops/_ctc_loss_meta_dispatch.h,sha256=W73IJbx1BT4OC9yeF9BaqQcoBMVRPh5mOs3xHts0BsU,1143 +torch/include/ATen/ops/_ctc_loss_native.h,sha256=fNlM1OaZdffvlgiTUrFeLjTKVGOjCrBij_u3E6PTKCE,2113 +torch/include/ATen/ops/_ctc_loss_ops.h,sha256=0Xd4cwR12xsf0ltk6FvnLMLqv3Sbk0WCrG3w-M-55jw,5248 +torch/include/ATen/ops/_cudnn_attention_backward.h,sha256=OzaadgOnD2VVZiUKo3Lyt_iqSDQic3fJLhYMGZD31Dg,4508 +torch/include/ATen/ops/_cudnn_attention_backward_cuda_dispatch.h,sha256=q6zIR38mHzaY9eH71ozTOv8-WLEtkgvGBzKwi1qG1LA,1955 +torch/include/ATen/ops/_cudnn_attention_backward_native.h,sha256=Pft_AcyOQgxgw521KMUZCuWryx9HfL91NwHt5YH9LSw,1185 +torch/include/ATen/ops/_cudnn_attention_backward_ops.h,sha256=lVenysf5Um5_aGWv2dOruaDNlyRa9qupWbQdQvW2w2w,2698 +torch/include/ATen/ops/_cudnn_attention_forward.h,sha256=5_S-T8gLDVuW-caVmoMd9-GZC3_FvCILyPiCOVujBfY,4813 +torch/include/ATen/ops/_cudnn_attention_forward_cuda_dispatch.h,sha256=Hjy-36-6VPpw3tkNwDfyZQ5yOcmme4XglCA85gaCTwg,2023 +torch/include/ATen/ops/_cudnn_attention_forward_native.h,sha256=VihKiFlWz6FNxMtdCANtMMemSWBI1jrUpbNteZL1Vdc,1219 +torch/include/ATen/ops/_cudnn_attention_forward_ops.h,sha256=s9FVwOIk1ivS3zb_N3JIR1mWZyQSJJz9OHl03avptGM,2882 +torch/include/ATen/ops/_cudnn_ctc_loss.h,sha256=nAgADceZjhVsA8XgoLe_q4TXhMHovuw3J8dwKdKNc3E,3189 +torch/include/ATen/ops/_cudnn_ctc_loss_compositeexplicitautograd_dispatch.h,sha256=QUF-5of97XiGsDqL_h7KL5v2ctGfv33TFhJDb5du2ig,1535 +torch/include/ATen/ops/_cudnn_ctc_loss_cuda_dispatch.h,sha256=43FN93b7zLxCw3N6jX4fuJWnBH9mTsKpSGRBS69isVg,1406 +torch/include/ATen/ops/_cudnn_ctc_loss_native.h,sha256=t8ayfKbh2QizjHxt0cC6XVyqJel7eCDlYs7GUfHowhg,1452 +torch/include/ATen/ops/_cudnn_ctc_loss_ops.h,sha256=KCpHIHOnTUt7cdvEnBctpoosuhCRLo94PbFqkgpDXsw,4235 +torch/include/ATen/ops/_cudnn_init_dropout_state.h,sha256=hJyo-naI6KsTKGSZa7rnxrdCfReBPfUrqLydQB7XyaI,2503 +torch/include/ATen/ops/_cudnn_init_dropout_state_compositeexplicitautograd_dispatch.h,sha256=yJLLK4yefNBK9RGsZHCpqOQhl2sO2FvSrrC6JWPvwQ0,1209 +torch/include/ATen/ops/_cudnn_init_dropout_state_cuda_dispatch.h,sha256=L4EaPrI__6ShcU5YaHpCAmg258C3ztxYCJu8crGAx7o,1290 +torch/include/ATen/ops/_cudnn_init_dropout_state_native.h,sha256=R9LqorAehoTFWS2kVkzOlacz9BLhdw-5S6oTnlIcjgw,1053 +torch/include/ATen/ops/_cudnn_init_dropout_state_ops.h,sha256=h7vcoHuTjB8Oqx25E9AVZgFtQ42SeoQOF_weLiUgtmk,2582 +torch/include/ATen/ops/_cudnn_rnn.h,sha256=_9h24m4R8CcUzV-oHQLwwFYs4Q95c8n7hMLN2QhSqJA,13597 +torch/include/ATen/ops/_cudnn_rnn_backward.h,sha256=6MPc7MQ7WB-Ypc3lQMlyPhK_vXZmMJKXBqwOON6jFnA,16712 +torch/include/ATen/ops/_cudnn_rnn_backward_compositeexplicitautograd_dispatch.h,sha256=CFMxzWsbRdZ9YCESUaleDp_ayJ9MTKN9lzbkbJh4fNI,3936 +torch/include/ATen/ops/_cudnn_rnn_backward_cuda_dispatch.h,sha256=N7ysay-6ViX1RxiO16TF4se-Xd6fQwPFN9tCh2srPMM,2379 +torch/include/ATen/ops/_cudnn_rnn_backward_native.h,sha256=ArkPkbugVJfW10oONGQ4fye0_h8ZBNfdAREow9YnJpw,2147 +torch/include/ATen/ops/_cudnn_rnn_backward_ops.h,sha256=dGt8jwIWS-bugalOABRzrWhHWJsPjI0t76u0XD-ow68,6215 +torch/include/ATen/ops/_cudnn_rnn_compositeexplicitautograd_dispatch.h,sha256=QaQv-hyg1U9y1OD0MHC3KIE7CfeSWm5tWQItG93kXyU,3420 +torch/include/ATen/ops/_cudnn_rnn_cuda_dispatch.h,sha256=oOwpaV-RHiGDcH1m4bxFuzCb5FNP7ZOF75eG11f9LdU,1931 +torch/include/ATen/ops/_cudnn_rnn_flatten_weight.h,sha256=og2KovZnaRrnbEBQCeLPCrS6k3QrX8BqdyzGofik1vw,7838 +torch/include/ATen/ops/_cudnn_rnn_flatten_weight_compositeexplicitautograd_dispatch.h,sha256=rn5EvbHxUswehYnNUjm4pYYU6zHyKsSZ-fgv_zuRgvg,2024 +torch/include/ATen/ops/_cudnn_rnn_flatten_weight_cuda_dispatch.h,sha256=w3407R86rGJiE-XL5WyNj_LtIp9fm06qw9Mc5Ml8ofc,1403 +torch/include/ATen/ops/_cudnn_rnn_flatten_weight_native.h,sha256=9284SeZWDMCOh2OfqWpZ-CpBv5St7tRtN7ls-Hp94Os,1181 +torch/include/ATen/ops/_cudnn_rnn_flatten_weight_ops.h,sha256=4cFcdJbq7Kuz_C6gweag5kOBVaANNuw4T-NkvzVC7GY,3041 +torch/include/ATen/ops/_cudnn_rnn_native.h,sha256=jL4LxWlkeT7ys_3vy-kp2fOXnqYSzjLCgnlhPn3RGXo,1794 +torch/include/ATen/ops/_cudnn_rnn_ops.h,sha256=OsVPj2AqRBfeDWxugUUFiHMQsu-yhkqm6zjLDgNMSVU,5128 +torch/include/ATen/ops/_cufft_clear_plan_cache.h,sha256=ZZC4WBBuxAITnBhD3UGB2VPGk4ia-1DFGbpZLo35VKg,975 +torch/include/ATen/ops/_cufft_clear_plan_cache_compositeimplicitautograd_dispatch.h,sha256=wOA17ABLp4_Lu5R6ILOdAtI5OmsloBxhosD3WlKWcrc,1034 +torch/include/ATen/ops/_cufft_clear_plan_cache_native.h,sha256=U4chjMHkmSw1chMIPO4BhjQIEkHCHjST8CG39bYnQB4,746 +torch/include/ATen/ops/_cufft_clear_plan_cache_ops.h,sha256=gGAH1J-EfcXJ48t_8DdFIycDDb2AGMSDnBsSsLgQI2Y,1260 +torch/include/ATen/ops/_cufft_get_plan_cache_max_size.h,sha256=ymozapyrEoZqfcAXp5J7jTrIw_LLlaTlfhncSVwIP4g,1007 +torch/include/ATen/ops/_cufft_get_plan_cache_max_size_compositeimplicitautograd_dispatch.h,sha256=q4bkX5P29uU9IzWHj-cuI6568WwK9euGiKAzf1kOOtE,1044 +torch/include/ATen/ops/_cufft_get_plan_cache_max_size_native.h,sha256=TjXJSwsOzS6vRXxtyjVIqgla6PzeCkQ5SF_IH6GmepU,756 +torch/include/ATen/ops/_cufft_get_plan_cache_max_size_ops.h,sha256=verPE2PINzX8oiC3hknhMJr2RZZPeVkF3izyGwu97QI,1291 +torch/include/ATen/ops/_cufft_get_plan_cache_size.h,sha256=-EZ0PTwuwqcKBzjmPNwhs9QTxm5DDcCXqX9UJhnjfEo,991 +torch/include/ATen/ops/_cufft_get_plan_cache_size_compositeimplicitautograd_dispatch.h,sha256=W_Wt4a9r0eAvQL0OC2ffXdx452RU3VzAdgiAIdcNhB0,1040 +torch/include/ATen/ops/_cufft_get_plan_cache_size_native.h,sha256=dSzhi1Y4jkQDuh-6EkfJrrx_qyq7-IhhErI5eqa9KKs,752 +torch/include/ATen/ops/_cufft_get_plan_cache_size_ops.h,sha256=ZC1gL0T01MAIPte9nPZaXdyxpJf-XiVNowkwwI6DPqA,1279 +torch/include/ATen/ops/_cufft_set_plan_cache_max_size.h,sha256=ys2B7AIeGCnoyZYwz-YzO3vrJrm4R7yUW9RRDT1XIq4,1045 +torch/include/ATen/ops/_cufft_set_plan_cache_max_size_compositeimplicitautograd_dispatch.h,sha256=OVP6gALGbiO_BxkxjJc4sxI7E76r93VeNTw4cDe1WyY,1059 +torch/include/ATen/ops/_cufft_set_plan_cache_max_size_native.h,sha256=FL4dZmSLhg_7xhR_G5oPKtnMytQg2ABKzeI6-F9SDBo,771 +torch/include/ATen/ops/_cufft_set_plan_cache_max_size_ops.h,sha256=qBoYdkhCWS5ENVpJkHdkRygcmWd9Mphs0Xb5CQFyshg,1340 +torch/include/ATen/ops/_cummax_helper.h,sha256=It9F836ziJSYsDsAu3h5vhtS96tV00nKbAXmVNudlZk,1039 +torch/include/ATen/ops/_cummax_helper_cpu_dispatch.h,sha256=IVGL2jYSDP2txe08RD30GyVr-Joo4ufuyLS8rAGqgco,1032 +torch/include/ATen/ops/_cummax_helper_cuda_dispatch.h,sha256=UO02t4HF8bfgXL1HiyN4RZfmFOFLhjDzC8aUVGSAbC0,1034 +torch/include/ATen/ops/_cummax_helper_native.h,sha256=_VR9YcSzmE9UpdqgSPKylv6EPl26ktrXnd-wF3ncR30,907 +torch/include/ATen/ops/_cummax_helper_ops.h,sha256=oLOa9MUc5Q6huIoOqhNlIhwBotGQESnN_wanw74ZOho,1410 +torch/include/ATen/ops/_cummin_helper.h,sha256=yCIFnbdBIL9rlO4WMGLvLGky01t-W6bdl6CR9MVUgyo,1039 +torch/include/ATen/ops/_cummin_helper_cpu_dispatch.h,sha256=pnqYF6bxThbn_3IHjfnKIN74MR_RtOKO9oEmaKar9ms,1032 +torch/include/ATen/ops/_cummin_helper_cuda_dispatch.h,sha256=nefInjgToYo5_bj36-LOtSx9T37hpZMeQ2SI8rKjmLE,1034 +torch/include/ATen/ops/_cummin_helper_native.h,sha256=k8EbJNzgjSXf6StEOvgBtLEiXnpfpYbaPZHT_wlEpD4,907 +torch/include/ATen/ops/_cummin_helper_ops.h,sha256=jKAW2klA7MjFUFLwEHVRx513RWunsmuks3YaTxNUmhU,1410 +torch/include/ATen/ops/_debug_has_internal_overlap.h,sha256=odlifVaXZdMqmkcQJM9kDLSfKG2ZBFhOJm8vJeQvIpE,969 +torch/include/ATen/ops/_debug_has_internal_overlap_compositeimplicitautograd_dispatch.h,sha256=zo2r-ziwiqKyQ7JB07Ow84GWmPfAg1Qvll8Ks_d2mEc,1036 +torch/include/ATen/ops/_debug_has_internal_overlap_native.h,sha256=uGLy0kkJHzgaU-PulXTxQFvXvU9cwpPqqYFpXFFtwBQ,748 +torch/include/ATen/ops/_debug_has_internal_overlap_ops.h,sha256=s7zczZwzKw_Gm3OtriKUFNXbjwytoVgLTC3CYJ6Pww8,1262 +torch/include/ATen/ops/_dimI.h,sha256=Cx62L7zqnJ3TIcWzT69IrOdxL_-hkuSRqRjgnSLDbeY,756 +torch/include/ATen/ops/_dimI_native.h,sha256=F-t3aPIhLO10H32AqEO27rLoWvHcJ512pSK09Lwqf0Q,738 +torch/include/ATen/ops/_dimI_ops.h,sha256=4vxt9zoyzka_YA9uy_anw9wmijp3z26o8PuIHPurx14,1196 +torch/include/ATen/ops/_dimV.h,sha256=IF5xJLSAwkwPmFgMeiU6WJSbbfKRYDx97L4VipjK3oE,756 +torch/include/ATen/ops/_dimV_native.h,sha256=U0yalBhDzUVuQ-E1QnVSH2smn3nJnlRO9BQKsIhxh2g,737 +torch/include/ATen/ops/_dimV_ops.h,sha256=LNnX6fTC8jG_IaO9S9h9gvy9Dh_6Z0jVzdPcPdW7VEw,1196 +torch/include/ATen/ops/_dim_arange.h,sha256=50Va0IxTAjs9UPuOA2qHugOTRDEws_xgbPkmmtagkXo,938 +torch/include/ATen/ops/_dim_arange_compositeimplicitautograd_dispatch.h,sha256=ojMc5z0W0_9nRgI4KCcTUl4Qp_tqlvG5W3lZhiBQoIk,1036 +torch/include/ATen/ops/_dim_arange_native.h,sha256=-3e4A0sXzfZB-y8gvGbG1S2hR2cCd2Jkk2IfPX5LmZw,748 +torch/include/ATen/ops/_dim_arange_ops.h,sha256=CEremvxrEYml_NqzZWqzxEEfxKgxzf4Zdfsh6rt-XZw,1270 +torch/include/ATen/ops/_dirichlet_grad.h,sha256=qVuIXEe_nijryVN3K2HAk-1VC_OuNvXOL5IVka31vuU,1626 +torch/include/ATen/ops/_dirichlet_grad_compositeexplicitautograd_dispatch.h,sha256=H65Fllj3pol3go4cBOxdqHp83gWtZcuQg5sVoIO-_XU,1237 +torch/include/ATen/ops/_dirichlet_grad_cpu_dispatch.h,sha256=eav3KrAAS2FUim-nebCTnIUlOTi6NDujp_H4d6Lyrtg,1032 +torch/include/ATen/ops/_dirichlet_grad_cuda_dispatch.h,sha256=UzHzrCRJz4w4X1V9ZA31goGMmICizge_g3a-bqLSmmk,1034 +torch/include/ATen/ops/_dirichlet_grad_native.h,sha256=aTVnxksJA3NbWO4yqmcYulnZtBHcS5XaOdvnheFCJN0,1045 +torch/include/ATen/ops/_dirichlet_grad_ops.h,sha256=mb88CbtgAAbY3N4Z4CsjQosY1eFQA6xtnmT4_izMyds,2181 +torch/include/ATen/ops/_dyn_quant_matmul_4bit.h,sha256=CuW3ucyBZkGAhBw5J0hkKjYpPNeiwcWlgJjwCaLBuJ8,1179 +torch/include/ATen/ops/_dyn_quant_matmul_4bit_cpu_dispatch.h,sha256=Jjg6SbeD7EUpzv8XCiOPghwl4Z_sABrMi1DhE6bN1qw,1087 +torch/include/ATen/ops/_dyn_quant_matmul_4bit_native.h,sha256=_FtksKJhpqqFkrluXxrkEo4WcTD5I6-oOm95b_lcSgE,847 +torch/include/ATen/ops/_dyn_quant_matmul_4bit_ops.h,sha256=G0eRfu340uhLQkecwQKJukV7-awcYgov1QWmvh3BHmA,1573 +torch/include/ATen/ops/_dyn_quant_pack_4bit_weight.h,sha256=OWObLMDe7Ux7PbzSrIWLOpNxLgQ7CIdUS9KiUwKy7Vs,1267 +torch/include/ATen/ops/_dyn_quant_pack_4bit_weight_cpu_dispatch.h,sha256=ayVeu3na6MVdsyTDNMipiAX2i6sSh0sb95M6tL9otNs,1136 +torch/include/ATen/ops/_dyn_quant_pack_4bit_weight_native.h,sha256=iKQkYVRsHNGiZcNE4VxWBhw041_sMsPnITh-tPiJrCs,896 +torch/include/ATen/ops/_dyn_quant_pack_4bit_weight_ops.h,sha256=M7LD3PzL9U0gzWKlCzjkaPjTdMF9Uyn01j1rrDZh7LY,1729 +torch/include/ATen/ops/_efficient_attention_backward.h,sha256=MVnKPjlpEf1wKxVC9sxhBJEE-mJEO0uVARpmiWSZE0U,6238 +torch/include/ATen/ops/_efficient_attention_backward_cuda_dispatch.h,sha256=XUMj2TNH42Yvw8D9qbMWrfsO2BtLCPNYCTbaUpwAxKA,2475 +torch/include/ATen/ops/_efficient_attention_backward_native.h,sha256=z-V_tH8B6yhH9-V3kG3ICw_SVEHq9ELxt1I3Py5q_54,1445 +torch/include/ATen/ops/_efficient_attention_backward_ops.h,sha256=bFYFkY5NCtNsdoGpWbuiF7_IMoywCFzvku0fRp4pJrk,3420 +torch/include/ATen/ops/_efficient_attention_forward.h,sha256=cHYPjuU2GxQWwCC__rLYYf3B6l7TySf6lL--kSlQ5fs,5659 +torch/include/ATen/ops/_efficient_attention_forward_cuda_dispatch.h,sha256=9tQLaJ59jwwgUFhcu94KMcLh1gccpfUL6jRqcsSVEbQ,2219 +torch/include/ATen/ops/_efficient_attention_forward_native.h,sha256=BZIphbOtsVo6Q6Au_hR3uw9-B8bAdNVRr3XhXmaPWGY,1317 +torch/include/ATen/ops/_efficient_attention_forward_ops.h,sha256=TFeiO38vvhRifGl0mHDKdrwgeSLnsBaUSUXXCcax1RQ,3128 +torch/include/ATen/ops/_efficientzerotensor.h,sha256=9Q3R0P_EnmdX27hy_SMN5Rmb6c7BI7RbmF90YUy1Ve8,6294 +torch/include/ATen/ops/_efficientzerotensor_compositeexplicitautograd_dispatch.h,sha256=62zWmeMSVAC7xEYI0lbcK_OcGfjpwdV2JoUGjJLAeJA,1344 +torch/include/ATen/ops/_efficientzerotensor_cpu_dispatch.h,sha256=eFrsnXh8Ehz0nrZV-_u6lemsOo6FKIyv0tGjh4KIrTM,1552 +torch/include/ATen/ops/_efficientzerotensor_cuda_dispatch.h,sha256=E0jAaKF067A7ehNcSxkc_u5AXoDKs4JtEtCH1z-P4Tk,1554 +torch/include/ATen/ops/_efficientzerotensor_meta_dispatch.h,sha256=vdpSqeBXNIufJ0E2f8-rv8TuvBJWVN9dXT3tRC09HFk,1554 +torch/include/ATen/ops/_efficientzerotensor_native.h,sha256=5AnXWOnBkqQTlT8ELj0pIkiilGQpymcEveELZK9ojZI,1463 +torch/include/ATen/ops/_efficientzerotensor_ops.h,sha256=x01MQTGXpxy5CDpFEVs7x8AWz334sWa5kO2BFvEysmU,2391 +torch/include/ATen/ops/_embedding_bag.h,sha256=-GZ6pdJ7VgIklM6PjzNeDYFvkzqJYWGf-eTzh4h6I-s,3432 +torch/include/ATen/ops/_embedding_bag_backward.h,sha256=NJ5Sw_KxRhU2G5PQ9K6vzVdOF50ckh515nyxreWBdpI,3804 +torch/include/ATen/ops/_embedding_bag_backward_cpu_dispatch.h,sha256=xrvnd9GImT7-v44gRFBCb6n6xZq6obYA53ChV4_hOsM,1683 +torch/include/ATen/ops/_embedding_bag_backward_cuda_dispatch.h,sha256=r6CUqnPsatki51BarUvtr6xAWm8WcNFQbUC8PaCSxw8,1685 +torch/include/ATen/ops/_embedding_bag_backward_native.h,sha256=Rw7xGFq0385Y4ggVKNRLcG0lj-biH2ROXCNLdYHUuVY,1063 +torch/include/ATen/ops/_embedding_bag_backward_ops.h,sha256=blLpnpU2cq60urzf8zbibegGVMcb6Qt0RKCzzdpT6IM,2262 +torch/include/ATen/ops/_embedding_bag_compositeexplicitautograd_dispatch.h,sha256=FNXwLTAmUkwgJGN9YQF5fB1hLRRc7vWY-QIGCo_GjvQ,1811 +torch/include/ATen/ops/_embedding_bag_cpu_dispatch.h,sha256=ZhCQc6e_4m0iLXZ6i_5YK7qfInVIBFYxbvX7fYmfUbc,1268 +torch/include/ATen/ops/_embedding_bag_cuda_dispatch.h,sha256=Sl9BH3udbP7b-wuLZEMULarvWIlfGJbsMbnnvQvWJak,1270 +torch/include/ATen/ops/_embedding_bag_dense_backward.h,sha256=MJSUI3hNDGkQ4jOCT-wmy4KkvYuvOVxripcHn6CcrlI,9546 +torch/include/ATen/ops/_embedding_bag_dense_backward_compositeexplicitautograd_dispatch.h,sha256=GhgIbw42JGRTvNuZGxpzbVKzEw9VD0vPclbgKcM0l38,2442 +torch/include/ATen/ops/_embedding_bag_dense_backward_cpu_dispatch.h,sha256=TCeNyY4_qFGMftZ_qLcmDm1f1O1IP1G3HcovYrIqXNo,1613 +torch/include/ATen/ops/_embedding_bag_dense_backward_cuda_dispatch.h,sha256=_PbwxBjfEh-qskR82LUAzeIQZ5wLYj_N04g5ldWG7MM,1615 +torch/include/ATen/ops/_embedding_bag_dense_backward_native.h,sha256=Hd7SnZ--9ooRW_85oifDn6OxZjUh4V2v4k_8fXTRW0U,1740 +torch/include/ATen/ops/_embedding_bag_dense_backward_ops.h,sha256=FZ4L2Od9jrAEolD3GpLTayvlq6ZdodFC6ztaHaJOn4s,3665 +torch/include/ATen/ops/_embedding_bag_forward_only.h,sha256=GPYXLnAS0l-eWm6kC1GDRuP88GOSR0N85gvs9SlfPkI,3562 +torch/include/ATen/ops/_embedding_bag_forward_only_compositeexplicitautograd_dispatch.h,sha256=_gJsOl0wDQhQvNaINsHtYxiFiD8POIelOzsR9jp8WpQ,1837 +torch/include/ATen/ops/_embedding_bag_forward_only_cpu_dispatch.h,sha256=C-RFxQqYC3qxaAw_L3nI-5Meq7XEkog7f1aFZVWB3MI,1281 +torch/include/ATen/ops/_embedding_bag_forward_only_cuda_dispatch.h,sha256=0JYuvVYoH0XmTPHYi8uZBvtyGd_6nrEsh0F61cF8MDA,1283 +torch/include/ATen/ops/_embedding_bag_forward_only_native.h,sha256=mhpf3kXa4tEPOH-RFd3Ghd9WNCdDx4kcCDBBx8YGEWQ,1830 +torch/include/ATen/ops/_embedding_bag_forward_only_ops.h,sha256=xSvOM0TI-v7qDhQGdncW_XafU7OE1luBjcSqf4-YbX0,3941 +torch/include/ATen/ops/_embedding_bag_native.h,sha256=pzuBiDXZZG54qIbFmylElJwE-B1ipJqh9IcNTUURBVY,1791 +torch/include/ATen/ops/_embedding_bag_ops.h,sha256=3fSeu07m5HCNrXE_N-qWmgsumeUElpDwArGBXU8754g,3863 +torch/include/ATen/ops/_embedding_bag_per_sample_weights_backward.h,sha256=tdZxxvkkSwc2Jb7S9YvFKILki_tiMQ9wpJCGQvR7TC4,2553 +torch/include/ATen/ops/_embedding_bag_per_sample_weights_backward_compositeexplicitautograd_dispatch.h,sha256=ZLQt0xFlv_ahOMYd_m8hy_G417OIYgaPnPTKe6DsxY8,1494 +torch/include/ATen/ops/_embedding_bag_per_sample_weights_backward_cpu_dispatch.h,sha256=EL34PhlVUOCs10Mg9rVA04VaiUPeCXcjvV_ghLw-lMs,1162 +torch/include/ATen/ops/_embedding_bag_per_sample_weights_backward_cuda_dispatch.h,sha256=XDlOD1I-ZDH096Opw9HobMj-Groz_QlL9cde2gLoXRA,1164 +torch/include/ATen/ops/_embedding_bag_per_sample_weights_backward_native.h,sha256=c-uiha-2DiY0L0xGUm-6rT93mbeSbmhpqdyRV4Eezrs,1432 +torch/include/ATen/ops/_embedding_bag_per_sample_weights_backward_ops.h,sha256=s7sXZWM8OnhCiEAT2yZQynwZe8eYjV4OyHOLzEF9uGo,3001 +torch/include/ATen/ops/_embedding_bag_sparse_backward.h,sha256=NTiAVTGLW4GPfjzQK9AfaNx3f5GVREv4PXvhbsxgQx0,3511 +torch/include/ATen/ops/_embedding_bag_sparse_backward_compositeimplicitautograd_dispatch.h,sha256=IFx_v4S7J_8ckaT1Arct3PXxxCcsKcaJcP1B5x8VDJQ,1643 +torch/include/ATen/ops/_embedding_bag_sparse_backward_native.h,sha256=sTHY4G8VfeBT3T4umm7iMH9BqZvUdMd2RFjDSrhzZsk,1021 +torch/include/ATen/ops/_embedding_bag_sparse_backward_ops.h,sha256=oEmN4eprf_CxlAA5mkyEE2mNe3Oiq1ymUMBUY7Vujq0,2122 +torch/include/ATen/ops/_empty_affine_quantized.h,sha256=4LyPSMoy-FwF-uZCo7IQJABNhyhdV7qsP5CVtQ6T7lg,9473 +torch/include/ATen/ops/_empty_affine_quantized_compositeexplicitautograd_dispatch.h,sha256=q3f5beVH5kQqLp_jxFejar6YtvYHWZV10Jh3nsVCv84,1756 +torch/include/ATen/ops/_empty_affine_quantized_cpu_dispatch.h,sha256=GBzcPHvSRpDo25VOic0dD7fy8Gpe_ERxdYLzOslP5Zc,1964 +torch/include/ATen/ops/_empty_affine_quantized_native.h,sha256=NYKnX5S6yRHvhIgfNe2grdf14YXnw8iPOqUJrrq2QT0,1564 +torch/include/ATen/ops/_empty_affine_quantized_ops.h,sha256=DtaVTfee2CL_sHPG0k0J_UxU_iaclVymshoveaKdSzc,3005 +torch/include/ATen/ops/_empty_per_channel_affine_quantized.h,sha256=tCRs5RrCO7pd_KqtE87N1Im4nXUrUKg3u414CtRTVV0,10781 +torch/include/ATen/ops/_empty_per_channel_affine_quantized_compositeexplicitautograd_dispatch.h,sha256=eYsRCX0CCUEpUWIi7AYslCxR7D9lqgMyRvWmsUZ88XI,1952 +torch/include/ATen/ops/_empty_per_channel_affine_quantized_cpu_dispatch.h,sha256=PgAWsJGHWAp4Rk-WU6PbQgQ029gWZVJDfEOzo3ol3So,2160 +torch/include/ATen/ops/_empty_per_channel_affine_quantized_native.h,sha256=ReCG3ckv-uPpPWsZiqxRukdL-Rhslgof9X6OYAGz_x4,1709 +torch/include/ATen/ops/_empty_per_channel_affine_quantized_ops.h,sha256=c5QBjpab_cKDm7VxIb_s4fh9K5ht0xvu5i1hwSw2OlM,3321 +torch/include/ATen/ops/_euclidean_dist.h,sha256=XPELByV84weaRT5wb2NnFU19LlGbWMhDhNyVjltiY74,1467 +torch/include/ATen/ops/_euclidean_dist_compositeexplicitautograd_dispatch.h,sha256=st9pYJzsV0ja8IEEuPtQGB3nnpkYDPgxd6eaPdsn1Iw,1265 +torch/include/ATen/ops/_euclidean_dist_native.h,sha256=OPNt5RmKW-2qrLTNOLrY57bcgPiifNmmXJ-3RuDGYuQ,868 +torch/include/ATen/ops/_euclidean_dist_ops.h,sha256=J9XjRcmvZEkoR_slwFGZgXPIWlyyLWxmiWSLbyA3sa0,1997 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine.h,sha256=ZJGPuaSavLVzgpOG7gxbdNchhiQNE9_ybk001tPctyw,2514 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_backward.h,sha256=hHw5MALf5AB_O3X5YoobKltAnfJmqdUI9SL_IEYiRac,1462 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_backward_cpu_dispatch.h,sha256=YxYavfm4XyFlyePswco2f5Jv_3waaIFguDToWDam0o8,1214 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_backward_cuda_dispatch.h,sha256=QpFbTiLiFR_VHOiUive07zO323W5fSD0sOk8D7Uh2rM,1216 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_backward_native.h,sha256=9D9FXJ8vykwyoQScen8J6R1KSYBeLOiwtFswsbZWAB4,970 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_backward_ops.h,sha256=MpHWhN5M0HzQ188h0NERVGPgwaLVsl2Y3k7OAb-i9A0,1987 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_compositeexplicitautograd_dispatch.h,sha256=VhI2U3Ezz9nTRmCcqGHAAtyby-7qwZ0ZWiU070ZD9vA,1457 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_cpu_dispatch.h,sha256=v4N9BDpFenkmLT-gMiEWhRdtI7pWlg_EsNK8is4-vHo,1144 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_cuda_dispatch.h,sha256=-vrAKuz3HHpzjyZLT-OsPNlwaZDUBGUkPqJMYKbo8IA,1146 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_native.h,sha256=57dwJvkMruko0G3Aw767Faz3fbFktVo4Cobvc-jd_ls,1144 +torch/include/ATen/ops/_fake_quantize_learnable_per_channel_affine_ops.h,sha256=Mivsy0rLR2NDv5xVkRTfbNSP1f5Tia7_D7a1y2FAe84,2881 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine.h,sha256=as4jcG_npOdTWLIlwIIsydqi5TyG6zMn1DEipQKYEJE,2414 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_backward.h,sha256=HB5GrByPw-ayrBxSsorhb8BRCI7jgG_pj63Qsfp4EGo,1428 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_backward_cpu_dispatch.h,sha256=LW4QWW9iiHe84LkOBxYDrlcIwi1G41TP6M-ING6yQS8,1199 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_backward_cuda_dispatch.h,sha256=Ezc77DMKXV9E0XAAIIIC_3xtOcaYbze52pWt1jrl2Vk,1201 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_backward_native.h,sha256=Cxv-XMyYZoTVIX7eDCMfEH_JVu16X9VgEW2MrJx7tCY,955 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_backward_ops.h,sha256=x0fQggJ-WcY4Jx_5fH1Im1if3UqZq8wKtgMr3jedp90,1937 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_compositeexplicitautograd_dispatch.h,sha256=RyY2XkPaq556mq6vm51XmtVByEuGrMYH0aTo7n3EDlY,1427 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_cpu_dispatch.h,sha256=pEZLfN2mtqQ2xzHQFMOeT5mDWrUE_yijT0_a1YRsmdk,1129 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_cuda_dispatch.h,sha256=yn9e9wBdeYLLD3JBld7-RypHxXW4LejZeSouq_F_6Ws,1131 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_native.h,sha256=GWdKLVqgzmKl_QzULxxobEDuqMGvX47ICYgdDK5cunM,1114 +torch/include/ATen/ops/_fake_quantize_learnable_per_tensor_affine_ops.h,sha256=vfCXW13QszoJtoHlLSdMTBQlLnIitFw1DdrlNWJ4SGc,2781 +torch/include/ATen/ops/_fake_quantize_per_tensor_affine_cachemask_tensor_qparams.h,sha256=dDW6dGNEMXPhWkzcxVcZgacp_GbraVSgZE3gwhhaCtA,2865 +torch/include/ATen/ops/_fake_quantize_per_tensor_affine_cachemask_tensor_qparams_compositeexplicitautograd_dispatch.h,sha256=TKvt0C08H8uhvYzRdYaPMINTx0hz7LChgCKwE3uD0uM,1585 +torch/include/ATen/ops/_fake_quantize_per_tensor_affine_cachemask_tensor_qparams_cpu_dispatch.h,sha256=mSrS2uFM9DlRW6OrhUjPyHyHtx519hIo12zfSvbdMEY,1184 +torch/include/ATen/ops/_fake_quantize_per_tensor_affine_cachemask_tensor_qparams_cuda_dispatch.h,sha256=AOR5lDHkbWNWGgxR2G5fCjCGhbEdyMeNtayR3JR-h14,1186 +torch/include/ATen/ops/_fake_quantize_per_tensor_affine_cachemask_tensor_qparams_native.h,sha256=JypsTHDOmK1eDCYhHLQT0vBQ5mHi5j0USc-ZpGMK8jA,1250 +torch/include/ATen/ops/_fake_quantize_per_tensor_affine_cachemask_tensor_qparams_ops.h,sha256=pG5BsRaOxwDB1YdahNx8-w_rwak209X6ooimtK_HExQ,3243 +torch/include/ATen/ops/_fft_c2c.h,sha256=t9eZZrQJSWQPbKAasRDof6LD9jiahSboAyG4LxUeaIA,4707 +torch/include/ATen/ops/_fft_c2c_cpu_dispatch.h,sha256=9MK-rivL7BM3Zpvk35ztxOkjuiJb3MrjspFB5VAqvFI,1735 +torch/include/ATen/ops/_fft_c2c_cuda_dispatch.h,sha256=YfG7ep5cn9YFVJZdCWF95AFcwujFaOH9iKG2M5eiaYA,1737 +torch/include/ATen/ops/_fft_c2c_native.h,sha256=mWU7grU5WQgaUbbhpJ0gRLQpOB7TN17wTleZVROJUP0,1200 +torch/include/ATen/ops/_fft_c2c_ops.h,sha256=hgnrV6Gl4N7Yq-EjG8lhd79E1hMHRQqtlVkxAzEqPE8,2227 +torch/include/ATen/ops/_fft_c2r.h,sha256=sMbYev1IASIs-bpqknrFwXOjZgTIHYRoYfCPz9i9TkQ,4761 +torch/include/ATen/ops/_fft_c2r_cpu_dispatch.h,sha256=JrQQ6k3VhhGqj8a7SHnuLEclmmOOomueaWeuUxnaT4E,1789 +torch/include/ATen/ops/_fft_c2r_cuda_dispatch.h,sha256=LaxvuxAHWJxvtcee1REoNsi9g-pP7T7pDVaEo6Vx3lg,1791 +torch/include/ATen/ops/_fft_c2r_native.h,sha256=OXdFhtOZ7mKVQ8E_jXWrwK8ukH31bAnIsrXj1-Mz3Wo,1236 +torch/include/ATen/ops/_fft_c2r_ops.h,sha256=si7oZH1F5gfVyhf6vQ4vX66xo7MaltbkLLfvRvMzp48,2279 +torch/include/ATen/ops/_fft_r2c.h,sha256=sPc9kbqy7k2AlQ9wieVzTIT1jtaKHP8H2JbbRNOcmFU,1703 +torch/include/ATen/ops/_fft_r2c_cpu_dispatch.h,sha256=IcILr3PMvEioM6-NY3phWeUnG9MmZvHqub5LN1N9J08,1314 +torch/include/ATen/ops/_fft_r2c_cuda_dispatch.h,sha256=3eIssSqvb3L5rMO18fMIGBo_wtvnC4TU-8_6iWUh8Qg,1316 +torch/include/ATen/ops/_fft_r2c_native.h,sha256=vyvmL3RnX-5GrSdifgKasIg_kllEKzSzbxGMwMKaOes,1204 +torch/include/ATen/ops/_fft_r2c_ops.h,sha256=QnjqS5dgWxprVfU0bsmhYCFG4t_8Mr0aT43IWEKh0VM,2203 +torch/include/ATen/ops/_fill_mem_eff_dropout_mask.h,sha256=ui78K-IOF5CCQQcP_NBPdlp9R3-TXYpFeQ3jMZFKhDo,1090 +torch/include/ATen/ops/_fill_mem_eff_dropout_mask_cuda_dispatch.h,sha256=z6FZ0tvGQ_0PhcfB-nF0Ycwte1CclFSfxc11OCVdClE,1041 +torch/include/ATen/ops/_fill_mem_eff_dropout_mask_meta_dispatch.h,sha256=yq7ecW0HOo7zoVhL8QsdA9FFirnhITPag-Urtw18XxI,1041 +torch/include/ATen/ops/_fill_mem_eff_dropout_mask_native.h,sha256=-ku1Qi-aNyOn4mejeWQZ7czydboWWf7Wljv0HX8U1Sw,795 +torch/include/ATen/ops/_fill_mem_eff_dropout_mask_ops.h,sha256=lzJy6QpsfC-nUQDq5LWIVBgzQ7_hHJZUCcJ6t1owJuQ,1431 +torch/include/ATen/ops/_flash_attention_backward.h,sha256=mL4YZFPTHDLs-RIe7uo6LSYDaKZ3Cl9cgrFencVqBhs,5314 +torch/include/ATen/ops/_flash_attention_backward_cuda_dispatch.h,sha256=RAv8Vk5LcMUknovqNroXI5JnN-M_p6zceEtotMD8iIU,2119 +torch/include/ATen/ops/_flash_attention_backward_native.h,sha256=0m3B_3vJ45JneH-4XbFxQgYuyIt-PW8bhq0v-vHzDwY,1263 +torch/include/ATen/ops/_flash_attention_backward_ops.h,sha256=Rjg6IlgQFkpVa2pfe-bxIHVD7AZ9iGLDzIIKCKTx9aw,2886 +torch/include/ATen/ops/_flash_attention_forward.h,sha256=ODLgqoQLjEfgUkLd13Ln87l0pSe7E-D9oY6Ras2wXFI,11934 +torch/include/ATen/ops/_flash_attention_forward_cuda_dispatch.h,sha256=rh-TZV1QZIayHqg8R1cBUDRUDvFnzzlW9jwbeGHQ7WM,3974 +torch/include/ATen/ops/_flash_attention_forward_native.h,sha256=aU8DmvthZn8j6DmkKUO_PzqpRBYUEg7W3-ilShO4518,2189 +torch/include/ATen/ops/_flash_attention_forward_no_dropout_inplace.h,sha256=pJApA3sEKn8V70Lfcj6x4tCfn6mfSQyeMF9abW5ch50,6178 +torch/include/ATen/ops/_flash_attention_forward_no_dropout_inplace_cuda_dispatch.h,sha256=n110ANswz4Qy7o1sfzPBYPFS7nTsibcQsLHAt5ow0mE,2369 +torch/include/ATen/ops/_flash_attention_forward_no_dropout_inplace_native.h,sha256=gfo3SPKcH10DbKTLenWuFllqYpoZtjXrww20osOutb8,1388 +torch/include/ATen/ops/_flash_attention_forward_no_dropout_inplace_ops.h,sha256=LB0pSUQ-uyGy2rHM1qtStEqOYzvaC2HwVN8ks_6o2S8,3222 +torch/include/ATen/ops/_flash_attention_forward_ops.h,sha256=b7YHts8y7DsJ1ezIiQcz5_aTtFA1JqV2BZL7MoLNTRc,6179 +torch/include/ATen/ops/_foobar.h,sha256=OHMlqMu7z_V4xehj__H86Qlz94UKxH6zwuwNszYpieQ,1621 +torch/include/ATen/ops/_foobar_compositeexplicitautograd_dispatch.h,sha256=EX9OMnCcQbR3_-Elc5NgubLUL6nCDlYTgqDDjQiu6-8,1204 +torch/include/ATen/ops/_foobar_cpu_dispatch.h,sha256=reCzEZ9MnpHtxpygrwLN2iqudL9AA20ne7mKiWGIMUg,1023 +torch/include/ATen/ops/_foobar_native.h,sha256=0GVP3RRBvPBHFUlZrABKnqKztv3wL4DCRR6Xwaem_tA,890 +torch/include/ATen/ops/_foobar_ops.h,sha256=YrlYfACiQX7WOTqaaHNWCVODX8GJbPZte52UtzMtisE,2074 +torch/include/ATen/ops/_foreach_abs.h,sha256=H863vI7g-07E_lUN5pDSIGAeZmL6RGz6pp6rPORYeoE,1476 +torch/include/ATen/ops/_foreach_abs_compositeexplicitautograd_dispatch.h,sha256=sqUV-WF32oTQGLusgnH26uIVUK409aiJ22pzGkMXG3k,1235 +torch/include/ATen/ops/_foreach_abs_cuda_dispatch.h,sha256=cnAXMVAXFd8FgOkgZt3xAb8mJIRQ6v6jJ8OcL9okebE,1044 +torch/include/ATen/ops/_foreach_abs_native.h,sha256=cHEe4iXe0JX6QUwPlWOxjK7CFkB_EraPu3dn0zjbXUw,1038 +torch/include/ATen/ops/_foreach_abs_ops.h,sha256=_pTuqBnMrUVhGDTD0pInq3870ObN1TSdpJbtJF6DFT8,2333 +torch/include/ATen/ops/_foreach_acos.h,sha256=mm2QP92RDiVeMpNegL19Ghu0dPkGRQ-0BEPO7RL3o9U,1489 +torch/include/ATen/ops/_foreach_acos_compositeexplicitautograd_dispatch.h,sha256=uSrz49WDHwSExLYI3Pl6BZrivhIstkeKPtaTwZdgzlg,1239 +torch/include/ATen/ops/_foreach_acos_cuda_dispatch.h,sha256=AEVicWLKkERE-cwbkK138wLCqgEO76tMNXnl9L0lt50,1046 +torch/include/ATen/ops/_foreach_acos_native.h,sha256=85Orh43GvNVmePYlaG2q6Jppw6gnCsZzOvkZketB7LI,1043 +torch/include/ATen/ops/_foreach_acos_ops.h,sha256=XeKTKWYH07s8SElwwn_7f29XlYONtQS0wIzfhaBMce8,2342 +torch/include/ATen/ops/_foreach_add.h,sha256=r7_hg7RPm74YE095Fqgx3A0-dvhZLlTcu_4khQPzkYQ,5086 +torch/include/ATen/ops/_foreach_add_compositeexplicitautograd_dispatch.h,sha256=KDGfuajqUsUqiEyltzNlke8Q-765o7X3wh7j_mRFSiw,2704 +torch/include/ATen/ops/_foreach_add_cuda_dispatch.h,sha256=Lmycadk-Ljhl3Jp_UEH7JgHwQpxg_Hxnlj8dAYTB5Hk,1740 +torch/include/ATen/ops/_foreach_add_native.h,sha256=uDNupOwbM9fNIKA_9x74-9faOJLadBAeFEMW777wFTE,3207 +torch/include/ATen/ops/_foreach_add_ops.h,sha256=PVe2D0-WNu9HUmPKoLqXGhLjjlqQixOm5HJz4ocIODo,9090 +torch/include/ATen/ops/_foreach_addcdiv.h,sha256=QdqMWlv6rWisZoQS9CvaqclJUE57ALaX3VPlFtXcbnI,5134 +torch/include/ATen/ops/_foreach_addcdiv_compositeexplicitautograd_dispatch.h,sha256=L_iSOvJkfX6mGoYADBsptgFWVPtT4n4lWQjYEkY3RQw,2759 +torch/include/ATen/ops/_foreach_addcdiv_cuda_dispatch.h,sha256=Zoe0yN5WPdWYMUhrl-v8AQWNGfdMQpVFSje1mT6cjwc,1780 +torch/include/ATen/ops/_foreach_addcdiv_native.h,sha256=gtnxGvk5WWVgewh9zNN7XafmUo873UtiKr-SZE6s7zk,3115 +torch/include/ATen/ops/_foreach_addcdiv_ops.h,sha256=DPPJbw81OAROJtSFLUtoEf3qoW3pgAyXHPM8ObMk8l8,8256 +torch/include/ATen/ops/_foreach_addcmul.h,sha256=NRegWHyCzqk_CywgC5NYaVIdHtc7pAHu4qeu4Zc-7ao,5134 +torch/include/ATen/ops/_foreach_addcmul_compositeexplicitautograd_dispatch.h,sha256=vZduFS7dP9mR-KEiNl9U_p1W73mswv081Y_a6uedi1U,2759 +torch/include/ATen/ops/_foreach_addcmul_cuda_dispatch.h,sha256=sFZfz5V_mKVm3te-V9JIDEmg8BGSVtqgYoY_GACuj74,1780 +torch/include/ATen/ops/_foreach_addcmul_native.h,sha256=TMluaUjl6WJSdGnfc7_AfCVzAqY3b7jKJni5MLuxbSk,3115 +torch/include/ATen/ops/_foreach_addcmul_ops.h,sha256=qIfClJLemXQcb25qQmoFHcAiN-kR-0EgFrT3X-tgjAc,8256 +torch/include/ATen/ops/_foreach_asin.h,sha256=AZ1lCR5Ly7qkoJdkFvDCUYlLmumyUc1cZXOtSlsDbAg,1489 +torch/include/ATen/ops/_foreach_asin_compositeexplicitautograd_dispatch.h,sha256=MXVA17N5xbhb_Fv1yBOfk_3pq-TbvKlHCct5XqbtgtM,1239 +torch/include/ATen/ops/_foreach_asin_cuda_dispatch.h,sha256=ycthJAtsYQ78QvhDddarb9P8pBYnXYSffz4xFEprFxs,1046 +torch/include/ATen/ops/_foreach_asin_native.h,sha256=xQnFUvkaVhBrqmajLDp2jeXmsz8slaSfrz1mw40phXc,1043 +torch/include/ATen/ops/_foreach_asin_ops.h,sha256=zia9gGDBjxTixjApN7K5XpcYQ-x3vLMOICtMBZYal1E,2342 +torch/include/ATen/ops/_foreach_atan.h,sha256=ly38qb46pBtJezAfJwZFTzQ4qHWlAIubPFMu1inaBAA,1489 +torch/include/ATen/ops/_foreach_atan_compositeexplicitautograd_dispatch.h,sha256=RsAVr5mF-RBdMFzWfmApC5HPq5J3trFPFIM-4OrHbMA,1239 +torch/include/ATen/ops/_foreach_atan_cuda_dispatch.h,sha256=3FRgTZh7vbGRUHOuKDV92XTj4uyCsN0ujg7mj6rrxhk,1046 +torch/include/ATen/ops/_foreach_atan_native.h,sha256=A70eF0bj1fPT9WZCVaDjB-PKleXIOd-rVLXVQuP2uoM,1043 +torch/include/ATen/ops/_foreach_atan_ops.h,sha256=vj9Jw3ulFiad6XNtXs6Bi4vZ4OeOriH2kEj1-uogXaE,2342 +torch/include/ATen/ops/_foreach_ceil.h,sha256=ylomgp_WfsCcSnGg5_uY0FxlMiKm-dwO5OXqwH9ISOw,1489 +torch/include/ATen/ops/_foreach_ceil_compositeexplicitautograd_dispatch.h,sha256=Sgfo4Qje1tT-mCeKExLibph112avchPLsYOO6-CBp1M,1239 +torch/include/ATen/ops/_foreach_ceil_cuda_dispatch.h,sha256=U28q_9kLOCWLewRuyozP3rveOjIml1PyAaJ-AagJZEk,1046 +torch/include/ATen/ops/_foreach_ceil_native.h,sha256=2ttSNiQU5mEDR4Pr20fADYyAK99ftB6A2oPfWS-gBhY,1043 +torch/include/ATen/ops/_foreach_ceil_ops.h,sha256=nTB0u2b404gCrZooJtHkk9WFFLuqZw3EDu9ovr9YkIg,2342 +torch/include/ATen/ops/_foreach_clamp_max.h,sha256=a_BW0FS6OZsX_Yibdamiuhz2xChrHZ7DnOBq7W4_-bM,3934 +torch/include/ATen/ops/_foreach_clamp_max_compositeexplicitautograd_dispatch.h,sha256=QHEHN5oNSJO-E0R7sYNh5V_D6Pqg9XfgI5Jfq5z_1-g,2181 +torch/include/ATen/ops/_foreach_clamp_max_cuda_dispatch.h,sha256=e2y109ojK3XojCHM2fVsqBOzXMlLfs4ihR3XdMMY5VA,1490 +torch/include/ATen/ops/_foreach_clamp_max_native.h,sha256=AAwI5fVIlYOOXIqcs9l8fVnNPjYr-Sz93w4VYvwtg4k,2466 +torch/include/ATen/ops/_foreach_clamp_max_ops.h,sha256=kaJdwnosGzpEEBLEGNLx4mrpOdLmDe7HfnFr0XcXS_A,6771 +torch/include/ATen/ops/_foreach_clamp_min.h,sha256=25dLVXkKbBf11Bvq9-GT2e1Bsfo-XPgDHlnHZr5CMl0,3934 +torch/include/ATen/ops/_foreach_clamp_min_compositeexplicitautograd_dispatch.h,sha256=94KOoii2md9z6kpDv5S2-xlymZFn2SYbEXXNb-klVRk,2181 +torch/include/ATen/ops/_foreach_clamp_min_cuda_dispatch.h,sha256=ENO8mCycIXkuVtCfAhqz8M6leW9HXX3PDjxARxzuwso,1490 +torch/include/ATen/ops/_foreach_clamp_min_native.h,sha256=MPxdF1l-17q5a5_I6wGSL_mM_jxo8hGop2Ae-pf1drs,2466 +torch/include/ATen/ops/_foreach_clamp_min_ops.h,sha256=1qloDHuRonStj0x3gSbCwihO88zzSztUJvonsMRY3q4,6771 +torch/include/ATen/ops/_foreach_clone.h,sha256=HLc2MEtcovJQjY34w_1tyd0QVrUo0KUTRDGZDTUrU5g,1675 +torch/include/ATen/ops/_foreach_clone_compositeexplicitautograd_dispatch.h,sha256=JytNs_w0-4iO8hVrEfwpWvaj5QDiQQoTUORVFLkbFuw,1367 +torch/include/ATen/ops/_foreach_clone_cuda_dispatch.h,sha256=Fz2a3wqdZEB7a-psnIEhwtVQYqkuh4JoMrGVkNWMoEc,1059 +torch/include/ATen/ops/_foreach_clone_native.h,sha256=JOOmm5xJIS5icSANI-KyvY5J2_ZlS4Km_5EHsGE0uOs,1097 +torch/include/ATen/ops/_foreach_clone_ops.h,sha256=SgNgQqAP1Rl9v8jD9EznDaKrMwUvBVtS3p4Fwq6wtFc,2189 +torch/include/ATen/ops/_foreach_copy.h,sha256=Pk8drkCVEwKYrKsIFHTTnZ__UM-0m3kCenUguFpRz8I,1904 +torch/include/ATen/ops/_foreach_copy_compositeexplicitautograd_dispatch.h,sha256=XyUYvtKdm_QVO04rJFQ1T5oePEkv3Nhus0ii8TN5eqQ,1413 +torch/include/ATen/ops/_foreach_copy_cuda_dispatch.h,sha256=0P4FLaXGMQR7RxKYIcQ2gdEDQiZK2-uEu-CdMHGDNUs,1019 +torch/include/ATen/ops/_foreach_copy_native.h,sha256=iEqLLaHxqf2n7WPXIcO81FoXDlA86j8SKrAaET8kKcE,1147 +torch/include/ATen/ops/_foreach_copy_ops.h,sha256=IidO_pI1qQMx3fOAEzFGOR3Aby52ZhFXmYpQovS3zHc,2768 +torch/include/ATen/ops/_foreach_cos.h,sha256=mm-lR8hp6b5T36aK_DY_UxXpDCRs_3t6M6UkPc1Hp7A,1476 +torch/include/ATen/ops/_foreach_cos_compositeexplicitautograd_dispatch.h,sha256=NZ8xb3dOpyrV2Lr66YJG_8nD8x5urxwj-Pxje7HDH9k,1235 +torch/include/ATen/ops/_foreach_cos_cuda_dispatch.h,sha256=ljip2w81ew2OU03AQQUJ_OACs5PviAwic5qZSFsZuic,1044 +torch/include/ATen/ops/_foreach_cos_native.h,sha256=BXp97JJxUDwVJJQRKhFZ3juoOmiGCS_llHzWtGw8U6I,1038 +torch/include/ATen/ops/_foreach_cos_ops.h,sha256=RTVTtHJzHIsJcOLGdLbELjZOFsk0YSb23vVz6mPqxCw,2333 +torch/include/ATen/ops/_foreach_cosh.h,sha256=iV57AAPWDrqjDevEuddigjyr-XGisCr9ZzQPtfNzFas,1489 +torch/include/ATen/ops/_foreach_cosh_compositeexplicitautograd_dispatch.h,sha256=EdOnWqAsggtDIyNned5ApVuvSZ0TfsgN1wxjOAdQ_3Y,1239 +torch/include/ATen/ops/_foreach_cosh_cuda_dispatch.h,sha256=HOYIU5vfN7me-70xM7oygwC-eDrYPQTs_7mpL7W_Cy4,1046 +torch/include/ATen/ops/_foreach_cosh_native.h,sha256=elnAgkI38rrJiUK3epdWCdd3GaWOQreIClYT9NKyaXs,1043 +torch/include/ATen/ops/_foreach_cosh_ops.h,sha256=5l8mNInBPwZIzBMKAv-HOqrl2e5pDz0fHj55U5g83oA,2342 +torch/include/ATen/ops/_foreach_div.h,sha256=7tL_jzBZYVATEXYxAg0uhzf0REpuwZ-nr_seV4Tu8wo,4670 +torch/include/ATen/ops/_foreach_div_compositeexplicitautograd_dispatch.h,sha256=o23HPAV0dtqoJPBkR9UU7RQkVlsPUS9JFf0XIzcvCOY,2484 +torch/include/ATen/ops/_foreach_div_cuda_dispatch.h,sha256=2WEVSD6Ed2X3J6Kc6S3kqA4iWGBMmxr1Ywp2g0yxaEw,1628 +torch/include/ATen/ops/_foreach_div_native.h,sha256=r9iNzytjGfIHJGu-bllTYwU_9vzv_WD_9zyHAtGR9fc,2931 +torch/include/ATen/ops/_foreach_div_ops.h,sha256=2ePX4pNpCpgj_wDF7rvJsG0b3q4VoO59NI1WLMBGwSk,8550 +torch/include/ATen/ops/_foreach_erf.h,sha256=38gYf2ysFed8ObLjQVPr0SSRAUP908l6dEInrooYr2c,1476 +torch/include/ATen/ops/_foreach_erf_compositeexplicitautograd_dispatch.h,sha256=h4bCyke1qDQ-mJKocKU-YPZR4WkllMqNg4Fc13wvRuo,1235 +torch/include/ATen/ops/_foreach_erf_cuda_dispatch.h,sha256=lKyNusBh-1YZaCxM0u1Nhkbr4DBtoQCGjiRvNO9HCJk,1044 +torch/include/ATen/ops/_foreach_erf_native.h,sha256=2ayNv2c2WZrPYtkQgxOlX0SAbf6SfVUvd4CqHe5yQYI,1038 +torch/include/ATen/ops/_foreach_erf_ops.h,sha256=QTae0NNBZIoBVPKrPFAxCqStOmRQdU9dCOjmUVts-qw,2333 +torch/include/ATen/ops/_foreach_erfc.h,sha256=Od5r6R7pJQk9hg3uE29In2eCBcKWMD9yvcGY4CvbA_k,1489 +torch/include/ATen/ops/_foreach_erfc_compositeexplicitautograd_dispatch.h,sha256=14iV6dmI5YIyq4p03Nb-cXUEyM7CPvNylnJJy94BrdA,1239 +torch/include/ATen/ops/_foreach_erfc_cuda_dispatch.h,sha256=z_R1wxaVb0MPM3mx-543rITbyIxRM4vgDSsWNPn89to,1046 +torch/include/ATen/ops/_foreach_erfc_native.h,sha256=ddQnOPmpzhDhClzOjANPs6umcQbafJH1fhZH5prBuws,1043 +torch/include/ATen/ops/_foreach_erfc_ops.h,sha256=OdnSirkkA0W-cIbte2iEnHx3FVdSdGNgcLJQDqfA424,2342 +torch/include/ATen/ops/_foreach_exp.h,sha256=Vd0LWgNOjrenImwVhuJpuITyNr55VN_HSNtwB4lWS80,1476 +torch/include/ATen/ops/_foreach_exp_compositeexplicitautograd_dispatch.h,sha256=bPe0PDDhV6uJEl7WjZzRRxcBaA1MU7nA2eDOqJLqjq8,1235 +torch/include/ATen/ops/_foreach_exp_cuda_dispatch.h,sha256=Qt8Ll-En7IogavZKZeiuILHbZODYP7dS62gRPBAUlDo,1044 +torch/include/ATen/ops/_foreach_exp_native.h,sha256=177rueOWuMS_umUOb_BVByeqzQ0Q6mMTmjyiEGZ3HFg,1038 +torch/include/ATen/ops/_foreach_exp_ops.h,sha256=VaQaevTmB8qKgwTwT8rjVuiPuBikpneCrtYMVUC2pSg,2333 +torch/include/ATen/ops/_foreach_expm1.h,sha256=zuvcpuiReWSzhGlTElAVX3-JJK6FDnrgwnl8x7o4k4Q,1502 +torch/include/ATen/ops/_foreach_expm1_compositeexplicitautograd_dispatch.h,sha256=1Q0qBTfhFLEwMTOyfpK5P9zO3jizX5ZUc3u-hEmHLWM,1243 +torch/include/ATen/ops/_foreach_expm1_cuda_dispatch.h,sha256=l-4ROJhUnVdM8zMF4A83BuqbJYdSrzyzySzZM91GQQk,1048 +torch/include/ATen/ops/_foreach_expm1_native.h,sha256=WRbWVaF5d2ZITLWpck7hpTSIN7jlRUHjgZqgRMXoiI4,1048 +torch/include/ATen/ops/_foreach_expm1_ops.h,sha256=iX071mcf9zuE_kZMmhHgPsJBpt-MztQ03yhqmktuiiQ,2351 +torch/include/ATen/ops/_foreach_floor.h,sha256=Hm7ephKM9iyKPGm6pjjM0Gre6AowCt2yXOfe577wNOQ,1502 +torch/include/ATen/ops/_foreach_floor_compositeexplicitautograd_dispatch.h,sha256=A9B6UJ-4FXFEMvSw3LBE4whsScVNTEVEIVB1SpWPmDE,1243 +torch/include/ATen/ops/_foreach_floor_cuda_dispatch.h,sha256=jnloWHKWy4z3qtHnFx8EfUbEuOI08wRSeTgQXxHdDsQ,1048 +torch/include/ATen/ops/_foreach_floor_native.h,sha256=pDEWkpjrZSFDo6YHvJdOUCxHDxMfIUCM50AtfmVS_yM,1048 +torch/include/ATen/ops/_foreach_floor_ops.h,sha256=qVdoWm1oEkVltP6jddpF2_Wl4zrb_msDd2OFNLJndbk,2351 +torch/include/ATen/ops/_foreach_frac.h,sha256=rOl9T4TYEZsLsB-oKszqGWWjS8t8XTnYLskhbaC1iJ0,1489 +torch/include/ATen/ops/_foreach_frac_compositeexplicitautograd_dispatch.h,sha256=FPwx5X7NPsDzKtqHe4oSlHPaw17hXZTeshkKsqjP1fk,1239 +torch/include/ATen/ops/_foreach_frac_cuda_dispatch.h,sha256=TwHTX5rcTITTdaaTZoImvSnsbIcSK1A9K3W6NR2Dn-I,1046 +torch/include/ATen/ops/_foreach_frac_native.h,sha256=6MxCgTD3_51OdyNmJTjfZf_x0CvAT4i2LhEiKj6J4NQ,1043 +torch/include/ATen/ops/_foreach_frac_ops.h,sha256=42VbxkwtYQy_qSpJfZ7KXqGWJrn2DmpRBZMkR3PRQKc,2342 +torch/include/ATen/ops/_foreach_lerp.h,sha256=AZtgGbGwOQetDEeTmc7k_lQCSQspkg09CThinZKkWu4,4409 +torch/include/ATen/ops/_foreach_lerp_compositeexplicitautograd_dispatch.h,sha256=Lvhy_KhTgnQ9Ue7ovMj8W6J7EffJ95bmmyYv6uvRnCw,2425 +torch/include/ATen/ops/_foreach_lerp_cuda_dispatch.h,sha256=5cy7EWDwxXs8KZyAEg6P-mDnbrCZ6bS0iPpr09cTx5g,1612 +torch/include/ATen/ops/_foreach_lerp_native.h,sha256=aZOu9iAB40q1Dg_1JqmF2SwNOGphBK8FXMJx0Xp4OBI,2719 +torch/include/ATen/ops/_foreach_lerp_ops.h,sha256=vM8Wfx03t87nRzmynTp_zLCESE_3wDNMiVdaVk3I6Xs,7410 +torch/include/ATen/ops/_foreach_lgamma.h,sha256=4iHn8HKc7QwX9c_8uxRWUNUVHQ4gyGwHKCbtKQ5Ev1s,1515 +torch/include/ATen/ops/_foreach_lgamma_compositeexplicitautograd_dispatch.h,sha256=OReicr7ulgPTvspPJ1Di2piOD0BaxesLoa0h6_yueLs,1247 +torch/include/ATen/ops/_foreach_lgamma_cuda_dispatch.h,sha256=PHVTDRpFsU_g2xDZMkSr3tfqMd7VAD06GN2TiuNHnUI,1050 +torch/include/ATen/ops/_foreach_lgamma_native.h,sha256=JX4E9KlG5tyw0OqsL4VROsSq3VorJIK4pwUrnbgpMgU,1053 +torch/include/ATen/ops/_foreach_lgamma_ops.h,sha256=2LZwYWbvbiGPUOFsWrkgcFCUtw16N1SSnN0C5vGnvMM,2360 +torch/include/ATen/ops/_foreach_log.h,sha256=STsM_GXWdN-7Mk2N0WfbCsqkD_JTwU3AcPt7wA0Y9A4,1476 +torch/include/ATen/ops/_foreach_log10.h,sha256=qm2nxgsfzw82xgRX1dCUL_7Efa8M-B8vCikcV-kyNxA,1502 +torch/include/ATen/ops/_foreach_log10_compositeexplicitautograd_dispatch.h,sha256=PrPQgoqCgTqNQTvZDhn8TL_OMhD9Ir51zJiNTiMO2Nk,1243 +torch/include/ATen/ops/_foreach_log10_cuda_dispatch.h,sha256=hFhB8DaVRBMpGl6-LYzcW6MYH2PE0lldhXoRcPMkzGU,1048 +torch/include/ATen/ops/_foreach_log10_native.h,sha256=3fb7wnkKIwolFx-Ccjl8vLXhxK687_HYofdRq7e4Iu4,1048 +torch/include/ATen/ops/_foreach_log10_ops.h,sha256=yinh1BipcB_2KqB_nK8sIfIJLN6WSYyfrsMjQSNaYC8,2351 +torch/include/ATen/ops/_foreach_log1p.h,sha256=C9-W3fYHFRh7dSmMH0Pyvngc1GbsJI4inm1wYEr3Jh8,1502 +torch/include/ATen/ops/_foreach_log1p_compositeexplicitautograd_dispatch.h,sha256=hMvNDcEimHAbuipN1B6pIqYBkTZHttmCR47aVBHR03k,1243 +torch/include/ATen/ops/_foreach_log1p_cuda_dispatch.h,sha256=WnaJJRWgOz1AE73IttiO5PFL6X32zSTYrvEhAaQSBao,1048 +torch/include/ATen/ops/_foreach_log1p_native.h,sha256=WlxuK1Es7YX3X6gB33AapX6ENjnSywifTn7CKOP4IUo,1048 +torch/include/ATen/ops/_foreach_log1p_ops.h,sha256=c7qCg3AEFkLjVBqFuwaRNVznXrwAkXMEV-ER8cPqF1s,2351 +torch/include/ATen/ops/_foreach_log2.h,sha256=V1jZHttgoka1gKa9L92tB0eeh52xWYHG0djnuH_eGHg,1489 +torch/include/ATen/ops/_foreach_log2_compositeexplicitautograd_dispatch.h,sha256=OdGW9UwPYRwEtA9t6bHAAImzW_TOJT693SIFXXc-bEE,1239 +torch/include/ATen/ops/_foreach_log2_cuda_dispatch.h,sha256=HiuoNRGvHtxoa_KSujZQdVvAl7KyetiEEoFVpwDdzqU,1046 +torch/include/ATen/ops/_foreach_log2_native.h,sha256=IKBnVWLN5mQKdAyWEplc4qSlnXuzcAhsBaBaoeu8Xq8,1043 +torch/include/ATen/ops/_foreach_log2_ops.h,sha256=oSwZVFS6K9mT9XlaLlyyORrdZwVqo_Jk3xIGaiMRMqU,2342 +torch/include/ATen/ops/_foreach_log_compositeexplicitautograd_dispatch.h,sha256=uV7YAY7SYPr19EfPIxwiQdK4TaDo_WFgHeEbXnCHC-g,1235 +torch/include/ATen/ops/_foreach_log_cuda_dispatch.h,sha256=LloA_WZ00FnhlE-Kv2eSbYzHKZw7wuvD868_MMMxSpY,1044 +torch/include/ATen/ops/_foreach_log_native.h,sha256=nSaNq7XU5aoS6NmM9Rm10fxV819GGhqsAXcMAomGjxs,1038 +torch/include/ATen/ops/_foreach_log_ops.h,sha256=5Z1RG81knsJ4gCECyv7LqIiaBsqdy8718HRv-0EOns8,2333 +torch/include/ATen/ops/_foreach_max.h,sha256=a0wbuTvV2pvD9TNq-V4OgtBBWjLdX77JiEFKU18DxtU,1328 +torch/include/ATen/ops/_foreach_max_compositeexplicitautograd_dispatch.h,sha256=a8VHMeSCtgmpP5AP9rPGP3lzdzThPOixu4XBxbxArHA,1184 +torch/include/ATen/ops/_foreach_max_cuda_dispatch.h,sha256=QeiCSqUg8WbtaQEYR2B9OY9mb4ENmOyCEXvKYXAgxpM,993 +torch/include/ATen/ops/_foreach_max_native.h,sha256=4S_VvDyRgMW7xhpIXonJecOYZIpzINDYlqFHAfRBRbo,914 +torch/include/ATen/ops/_foreach_max_ops.h,sha256=0EWp4GXeELrHig2xuH8iF_aCIEPLaMTcVlmcfmvhP7c,1840 +torch/include/ATen/ops/_foreach_maximum.h,sha256=Mp5ZcMzfg_gl0_7teXFg_ojdSbn1IRqGwBSS5EBkuGg,3860 +torch/include/ATen/ops/_foreach_maximum_compositeexplicitautograd_dispatch.h,sha256=W-Gf1WH1rstCKwKHtvwWXkbilUnHmKszC3CX06jajFY,2157 +torch/include/ATen/ops/_foreach_maximum_cuda_dispatch.h,sha256=uo3AoRyJzybbQGmZ_fHk0xC8X6nmZWEtMGIMbjq9TsQ,1478 +torch/include/ATen/ops/_foreach_maximum_native.h,sha256=k_7sTfeOZm4ReqeWuX2Lcju13fPOZv099Oxig3bhoNw,2460 +torch/include/ATen/ops/_foreach_maximum_ops.h,sha256=EItwNAP5KlKCl66TQvgUVnhtlbNTbULvImqMUeOvZAg,6717 +torch/include/ATen/ops/_foreach_minimum.h,sha256=RffEwCJCcpuaT_dJU3AA46FJEHe7XtibFx11xGlI_jM,3860 +torch/include/ATen/ops/_foreach_minimum_compositeexplicitautograd_dispatch.h,sha256=qAP5EzKdl3nrptdjNsluWgy5SLMctgxN--8GQUaKsVU,2157 +torch/include/ATen/ops/_foreach_minimum_cuda_dispatch.h,sha256=t_XFkxwi2sGv4KFL5fr_-DgOQQOWlfAIyGZ_qUd58ZE,1478 +torch/include/ATen/ops/_foreach_minimum_native.h,sha256=k3kHt0HHhA7xfYZtELtOWpb7rAOVTmdbRjumVffQ9qQ,2460 +torch/include/ATen/ops/_foreach_minimum_ops.h,sha256=zFiblbn0l07Wt8Rj9rVPYbSycPOyx2qQOZqpTWcGzoQ,6717 +torch/include/ATen/ops/_foreach_mul.h,sha256=9ajVCQvTXsBcfeSSZhjtyGegYyCkIE5WR2A1RzqRLnA,4670 +torch/include/ATen/ops/_foreach_mul_compositeexplicitautograd_dispatch.h,sha256=TNAKI_Km56tWjXvYtXWPTwEazRcCxYmB1hUVxsqcNGY,2484 +torch/include/ATen/ops/_foreach_mul_cuda_dispatch.h,sha256=kRUdpZ0dLpF5-CCZOFZ6ep7jGsshIdXQRGdhvvYcw7E,1628 +torch/include/ATen/ops/_foreach_mul_native.h,sha256=Drr3pJ_BqXrFLdsDlM59M2Hl0zj4SWBNT7DG-IPRD5o,2931 +torch/include/ATen/ops/_foreach_mul_ops.h,sha256=p2YHSMBZgghntRmjcqc5LXy7mkB-YYdk5WZRp-DY2Qw,8550 +torch/include/ATen/ops/_foreach_neg.h,sha256=f2x-pvcxiYFg4vG4NwlktlEmTWk3YVv4Cl-6qJ8yva0,1476 +torch/include/ATen/ops/_foreach_neg_compositeexplicitautograd_dispatch.h,sha256=XywtYFHU-u8yArCKnE4vfjnVSIVun_5O1wOEGP_MVDQ,1235 +torch/include/ATen/ops/_foreach_neg_cuda_dispatch.h,sha256=LPVYNaE-gBkG-rPj6SfEA6GMq2HyV9aSqUG715zh27Q,1044 +torch/include/ATen/ops/_foreach_neg_native.h,sha256=uLsNgz0OpcMZs9a8xdvLQQqNZyOS1nZB0uSed5uAxFo,1038 +torch/include/ATen/ops/_foreach_neg_ops.h,sha256=xpcJUIDkpQDJ1MO6-_Je5_iU4qXUmNib_fUwU_ft_eI,2333 +torch/include/ATen/ops/_foreach_norm.h,sha256=W-1qQrogPdbIAnCBm1dV4LZzsJ1ga055wnIZWlyWGkA,1753 +torch/include/ATen/ops/_foreach_norm_compositeexplicitautograd_dispatch.h,sha256=4mJ1ag9PjJRAyVvh8Xnk1W6_KAKIWAPW4UaM8-Oq4OE,1410 +torch/include/ATen/ops/_foreach_norm_cuda_dispatch.h,sha256=VdTC_m4xcvd4LWKbu6RMXE2cPpRna40LstMqA7Gbz0w,1074 +torch/include/ATen/ops/_foreach_norm_native.h,sha256=-L3O42cEqG5eD0131Z3pgRxThv1SmUtpfkhGTJQgu6o,1147 +torch/include/ATen/ops/_foreach_norm_ops.h,sha256=N-jKDjNpmoRnGm4jDOgtRQJPz9o78QXnYu-3M2DYoq4,2321 +torch/include/ATen/ops/_foreach_pow.h,sha256=-lKiqGD6SqjLc2BOE9j5vRtzupA04WUdi3CEYneRudk,4040 +torch/include/ATen/ops/_foreach_pow_compositeexplicitautograd_dispatch.h,sha256=GTRYb9tsFgcu3Q3PvwAvPVn5askjon87TOCDCv3xx6o,2233 +torch/include/ATen/ops/_foreach_pow_cuda_dispatch.h,sha256=22HIk8rUtUf9jMnsTFb6MMl3hoCA9V265ezAxQMU5tI,1566 +torch/include/ATen/ops/_foreach_pow_native.h,sha256=_RhdTs4jarBm9fcDJjJdQohp8YAGs_1_oll_W6nLBJY,2652 +torch/include/ATen/ops/_foreach_pow_ops.h,sha256=iphu4rfr_pz7fukCdZgLII1Kk-bofPyl5g5M_cFd_qE,7360 +torch/include/ATen/ops/_foreach_powsum.h,sha256=F9SaXsEZwRKP0A0oaJoU7AIHlPnN0I2NaYWqPinAi5M,1773 +torch/include/ATen/ops/_foreach_powsum_compositeexplicitautograd_dispatch.h,sha256=0-C3mgbUrxDgztpICH1IA4jnPEorpefubFsmfk9Ktsc,1416 +torch/include/ATen/ops/_foreach_powsum_cuda_dispatch.h,sha256=CrGd_oUTmn2j6Wd8H9uqiUPOSkrQX4AbLmlhMu2JK18,1076 +torch/include/ATen/ops/_foreach_powsum_native.h,sha256=agrvpx5O4etWwH7PbLl4zYFjLZl8rhNCZd5NMLiHCZM,1153 +torch/include/ATen/ops/_foreach_powsum_ops.h,sha256=VKKf-rmtlwEQdXqDkMW--IfACrMP-eK6cHIxxKJO1YE,2333 +torch/include/ATen/ops/_foreach_reciprocal.h,sha256=zrp2-TrifV9HJLgA3bRNtgZ2tLJu22OHGXxz-bzZYQI,1567 +torch/include/ATen/ops/_foreach_reciprocal_compositeexplicitautograd_dispatch.h,sha256=YuUJdFd9vvITYxj4yG-26R7_fU-_elBZpyJxcUaF7G4,1263 +torch/include/ATen/ops/_foreach_reciprocal_cuda_dispatch.h,sha256=g1ZQKR08zNfyoaHTmZf6sp49AYmN94urcTpAEzUhya0,1058 +torch/include/ATen/ops/_foreach_reciprocal_native.h,sha256=_FZB-q2w_tZADF5FQlupTiQXt6Echx6E0luOzUIgZ7o,1073 +torch/include/ATen/ops/_foreach_reciprocal_ops.h,sha256=dgcEk2YLh4OeeQLXrHwu67bYWbCWlw7cPJLxrTnyKR0,2396 +torch/include/ATen/ops/_foreach_round.h,sha256=HO8iLDuX1cW-uLKq4kTmLzP--J_I_XgBgy7ItM19fk0,1502 +torch/include/ATen/ops/_foreach_round_compositeexplicitautograd_dispatch.h,sha256=r5ovquR7UP-MNtEm9SQZQ9Hd-5W_pKTS8jEYSMLPlyo,1243 +torch/include/ATen/ops/_foreach_round_cuda_dispatch.h,sha256=H1_XxtOTnBIQEg0pWjipGH3sMwJT6a4MO7pjwuneqRo,1048 +torch/include/ATen/ops/_foreach_round_native.h,sha256=358Bq56dXVg-R8BIGK-mW2Sklg3OfEr9Idr0V0mpfQE,1048 +torch/include/ATen/ops/_foreach_round_ops.h,sha256=2rogSekI7iktty6J_FT1v5hfGpXz829rtQZOXUGQD-M,2351 +torch/include/ATen/ops/_foreach_rsqrt.h,sha256=JV8xRQv5qWIa_yFkKdtmAWB4vRTjJvCWo9BpLA7lcNg,1502 +torch/include/ATen/ops/_foreach_rsqrt_compositeexplicitautograd_dispatch.h,sha256=LVMtWef5OtXMniZwpf9yfaOl9JprNMqpWv4iWT0_TkU,1243 +torch/include/ATen/ops/_foreach_rsqrt_cuda_dispatch.h,sha256=IbgBwe_VNVo-_muQ32MCrcQfpQQuzvA7L_FkTqmoTDg,1048 +torch/include/ATen/ops/_foreach_rsqrt_native.h,sha256=jOdnUZW0xZ8relwXMg9wk3FXHG0dQzPf3LrZz_0N6jI,1048 +torch/include/ATen/ops/_foreach_rsqrt_ops.h,sha256=sjeT527-oIAxOCjr_7Kui-MA48HvYttPFuMztn2cSPI,2351 +torch/include/ATen/ops/_foreach_sigmoid.h,sha256=Di4wPgfzzXTJUIiaGYAJX4basJJfqvELEA4KOIgrfVk,1528 +torch/include/ATen/ops/_foreach_sigmoid_compositeexplicitautograd_dispatch.h,sha256=6WbOVDA-iu9pUFKR1N5Qlbh7BUChar5ztZo_cXcGWMs,1251 +torch/include/ATen/ops/_foreach_sigmoid_cuda_dispatch.h,sha256=myuRlVqZK9iKugcTgT0Fio_ueR9hirPVTxjuDX26fnI,1052 +torch/include/ATen/ops/_foreach_sigmoid_native.h,sha256=UnmVWKMao1vjORYmmdC1IiPQId_gevnufam_Gsy407o,1058 +torch/include/ATen/ops/_foreach_sigmoid_ops.h,sha256=CYhginFUzrmqTgcQhfOTmxCwoZrRCFr1xBDCzKPlpZg,2369 +torch/include/ATen/ops/_foreach_sign.h,sha256=4-ZaVDCZx7gONguObU9Q68EwkPpWTA3iVtHuR_DC530,1489 +torch/include/ATen/ops/_foreach_sign_compositeexplicitautograd_dispatch.h,sha256=cedUDCe5TEh-T1JzUlgf6Gc_aMoOFrhXTVBJ2qSbjbw,1239 +torch/include/ATen/ops/_foreach_sign_cuda_dispatch.h,sha256=aRNfBqg3kvHyQC-n8sQIVs1kBkbsNkALxdZ-vKLwPOc,1046 +torch/include/ATen/ops/_foreach_sign_native.h,sha256=n8Tzx2_rSx9QL-ktFIz0dHkLp96paPcUEp8RcdKE4YI,1043 +torch/include/ATen/ops/_foreach_sign_ops.h,sha256=y64JnpNYl_T7swUVMmHc3XcmxGHtWEvdULoSXAazSwg,2342 +torch/include/ATen/ops/_foreach_sin.h,sha256=lF-jF7TQJqcPLDw9VUf4ljOJ890UkaZBpSL094q7gUg,1476 +torch/include/ATen/ops/_foreach_sin_compositeexplicitautograd_dispatch.h,sha256=1uCQUdqUZLZxhAzQwg1sRlI5aisMdSoeem7m3rUWkeM,1235 +torch/include/ATen/ops/_foreach_sin_cuda_dispatch.h,sha256=LA3hhqnAjtlNtIYuCfDz56qaQWoe6k4QXmSqJt8QE6U,1044 +torch/include/ATen/ops/_foreach_sin_native.h,sha256=3YAAZVODEwgTFPY3EntwnmgBrUD6xN69iguJ1lnhOQQ,1038 +torch/include/ATen/ops/_foreach_sin_ops.h,sha256=x8yLT729lln4-6XjWg2hN3pF1wkDIl6fxD0TlvtQ5-I,2333 +torch/include/ATen/ops/_foreach_sinh.h,sha256=EZtu8CFHTc92qEo89KsENz9gWM1Gvwt38lpbVrqzm08,1489 +torch/include/ATen/ops/_foreach_sinh_compositeexplicitautograd_dispatch.h,sha256=DqYfQBn55qXnoUBwiaiW0T8dptpB_NeTRAaXICp8PtM,1239 +torch/include/ATen/ops/_foreach_sinh_cuda_dispatch.h,sha256=ZR9siY3t5T0yGh8baMRDyDfem1QXOv41MfYjsWW5GB4,1046 +torch/include/ATen/ops/_foreach_sinh_native.h,sha256=PcA16H8nhg9eLcQlAKH0AT037Gl4n1VDUfWOtwPAFv4,1043 +torch/include/ATen/ops/_foreach_sinh_ops.h,sha256=AyD3tG1OHOVosobFXw7cj5e55z1J-HL2FTOV2hEe3xg,2342 +torch/include/ATen/ops/_foreach_sqrt.h,sha256=QGJyYtGx5_pK38TJ4kAaTW00i3cOD5ALRluZijG42yg,1489 +torch/include/ATen/ops/_foreach_sqrt_compositeexplicitautograd_dispatch.h,sha256=KCDYeq0eRSsJrdNBsKg-xgjw4TWs0Z4w59G7NcIFxZk,1239 +torch/include/ATen/ops/_foreach_sqrt_cuda_dispatch.h,sha256=qn6bQS82-PJid6yksU19GEupBm4JhB6z8YFsHEYVdCE,1046 +torch/include/ATen/ops/_foreach_sqrt_native.h,sha256=Hu0L_QWZk0rO30Ox9ASrVHWAtWgkYLGcR3UrPIl44-s,1043 +torch/include/ATen/ops/_foreach_sqrt_ops.h,sha256=dM_VJreHXO5U6z9b55UnyguccJ4Hf7OHTIFtuYnkTG0,2342 +torch/include/ATen/ops/_foreach_sub.h,sha256=_EbUgHyvVYfGyVoRClzEcRmVITxetnpP56YrMyJKbsI,3920 +torch/include/ATen/ops/_foreach_sub_compositeexplicitautograd_dispatch.h,sha256=hOKkh-18w6BOFIEA4SEoODYkGC5O4cAduN_w_jyOZ3o,2219 +torch/include/ATen/ops/_foreach_sub_cuda_dispatch.h,sha256=F5NRnKE-5ErGvpGjkhKb_XVEZAlgb231zS-phJDyM-g,1510 +torch/include/ATen/ops/_foreach_sub_native.h,sha256=lvR_Nmh6qAlRXLLEwGWxEEM64fmL0dhryFGdYx9Lm6I,2514 +torch/include/ATen/ops/_foreach_sub_ops.h,sha256=_cKVdFUp75SwFKHF1JWlxwvkhbwnMzH_WUkq4okQO2g,6879 +torch/include/ATen/ops/_foreach_tan.h,sha256=_7I0j7mDxBRR3kSaq75qmFwr8vPmoN-JUbTHvVJMTVM,1476 +torch/include/ATen/ops/_foreach_tan_compositeexplicitautograd_dispatch.h,sha256=lLY4NItXa8VBF5uabl7kZuknRSiqHNkGUXbLMGIXuDY,1235 +torch/include/ATen/ops/_foreach_tan_cuda_dispatch.h,sha256=ONrhGzHeYoee8g8XYyPAffn4sTaKIwEcZhrlWM9TunU,1044 +torch/include/ATen/ops/_foreach_tan_native.h,sha256=N5Ohc8bIxv-RvId_TFlq7bd0WbthQE8iv-oI7OxRt5M,1038 +torch/include/ATen/ops/_foreach_tan_ops.h,sha256=9SjxsWf2S-D3MPwnY-m7QSUbAxWXRq08JbvMSzmmRns,2333 +torch/include/ATen/ops/_foreach_tanh.h,sha256=-JOTdt54l6vCKVREf_O9QaPnFYjT_Si3Eagb6Wq35wI,1489 +torch/include/ATen/ops/_foreach_tanh_compositeexplicitautograd_dispatch.h,sha256=yStqmJpT4WVLZVMl1j_TFBWalCd___Owouoc2BNNoXo,1239 +torch/include/ATen/ops/_foreach_tanh_cuda_dispatch.h,sha256=q5k2GcVxPZBFxVDrcgnY46_B08v6nGnGSdv9pEG2Uhs,1046 +torch/include/ATen/ops/_foreach_tanh_native.h,sha256=O9X9kSNDxjuwt1m-lJaeXuEdnvrPa5_PhsPqnYLNAbk,1043 +torch/include/ATen/ops/_foreach_tanh_ops.h,sha256=Y0m_yTJt3l4XBpSABsPgvxb8dHHtrM3W6KvP5h1_cFQ,2342 +torch/include/ATen/ops/_foreach_trunc.h,sha256=H-_8HPF6s-GcQyWa1UIsYGpGwujl4frkPNgXeUfydGY,1502 +torch/include/ATen/ops/_foreach_trunc_compositeexplicitautograd_dispatch.h,sha256=CggimwSCuqV-garTGnMa-8J6cCV88wRb9BF45RbIovg,1243 +torch/include/ATen/ops/_foreach_trunc_cuda_dispatch.h,sha256=sJSlCjErUBZnEN7xlegy38O1v2x45nDmBOp6xUjb1CM,1048 +torch/include/ATen/ops/_foreach_trunc_native.h,sha256=u9-q21wCzFGTxS7pA-hMXqJI1Z5pjl4dnM_csv3lAJo,1048 +torch/include/ATen/ops/_foreach_trunc_ops.h,sha256=7Ic71FFdspzGbqwphxfdp8XhjQQ3d1Awlob1qsl_ab4,2351 +torch/include/ATen/ops/_foreach_zero.h,sha256=7aYbFms3Uq2a_hGDBss2HnSASvdYSbPvB7BGYf9trBs,1498 +torch/include/ATen/ops/_foreach_zero_compositeexplicitautograd_dispatch.h,sha256=_Go_fWuNKmEnWtwKqJaJ0bzx3zSjrZy2INDJgOhUkRI,1239 +torch/include/ATen/ops/_foreach_zero_cuda_dispatch.h,sha256=qV4RIuX7YDd2QnaWNeFI3V2__GzzfBMzOGI-3dI4F7s,974 +torch/include/ATen/ops/_foreach_zero_native.h,sha256=v_AuTLkkmt3Rf3XGSsmBbnZ2OfhsSCoYBwZGneO5zVM,949 +torch/include/ATen/ops/_foreach_zero_ops.h,sha256=mW19R0-S7SvWKnCNi2iKOuiSblnRYjLq9F6-UxtFHS8,2351 +torch/include/ATen/ops/_functional_assert_async.h,sha256=W6HDTE5JmdmJijszzDMKXVcoTwNzH620DjIEck0cpZ8,1087 +torch/include/ATen/ops/_functional_assert_async_cpu_dispatch.h,sha256=nhBcOZ4gZZckPfaUCYVo-tS5oCb_rapoJbI-L63C2yI,1051 +torch/include/ATen/ops/_functional_assert_async_native.h,sha256=BEZoN2QPpAhRxhHvq5ed_foEIeKnRCLi0t3bv9okUTA,815 +torch/include/ATen/ops/_functional_assert_async_ops.h,sha256=Lt0WTqU8kGmMGdaMUxycYAj6EY5D10SMsWKBtUaVPGo,1466 +torch/include/ATen/ops/_functional_assert_scalar.h,sha256=wXUJEP5VjMC-O_0mRidnwFg-H7dfQSJLLi9n3KqNecw,1083 +torch/include/ATen/ops/_functional_assert_scalar_compositeexplicitautograd_dispatch.h,sha256=kB5oy61fM8HprD-QAuiwnabEolmjhAPtdAbBblfCDys,1096 +torch/include/ATen/ops/_functional_assert_scalar_native.h,sha256=y_MoX99UnZvdpCzWghOAI3tzKZJfPKmMay0ZYK40tEY,808 +torch/include/ATen/ops/_functional_assert_scalar_ops.h,sha256=UUrhCEyPGBa2po3N5Zamd4DHZ5b9wMi7coca3KrqpQU,1458 +torch/include/ATen/ops/_functional_sym_constrain_range.h,sha256=GUUOhmT0c_ukEp6TAHVcoLwJYeAHrpR7PQydxspTL3E,1140 +torch/include/ATen/ops/_functional_sym_constrain_range_compositeexplicitautograd_dispatch.h,sha256=_Sj45uWlcBgpSP6ZL5HduisY0SJFfQH8rp7ObaBJssY,1133 +torch/include/ATen/ops/_functional_sym_constrain_range_for_size.h,sha256=Je4Y86981gJGLUhJXcsLXL4Vf8opJliaWtkf9e2t3GQ,1176 +torch/include/ATen/ops/_functional_sym_constrain_range_for_size_compositeexplicitautograd_dispatch.h,sha256=OuZ9Bj3usdVqxugoYrje7i32ICHq8oQvkWfstcW6vHU,1142 +torch/include/ATen/ops/_functional_sym_constrain_range_for_size_native.h,sha256=dVKydlWzUtCkLlIigDI1-x0H0lQgz4q7W0nFtkR3-Nk,854 +torch/include/ATen/ops/_functional_sym_constrain_range_for_size_ops.h,sha256=EGIIdkAl_8JeO_I47n18PvjWqzJ01ocid82_fHMrPoY,1603 +torch/include/ATen/ops/_functional_sym_constrain_range_native.h,sha256=bx_QtLkDAJ3hgEDXJh9LcyjfaW8H5mepoI4115yxLMI,845 +torch/include/ATen/ops/_functional_sym_constrain_range_ops.h,sha256=r-3S0vkU-qb3qt2D93aWomcy3p6c1dg5X63BYVWQSSk,1576 +torch/include/ATen/ops/_fused_adagrad.h,sha256=m_RygKv8bERhR8fC1smEKkvcfVDcm0SGWxDjIdL1bAE,6984 +torch/include/ATen/ops/_fused_adagrad_compositeexplicitautograd_dispatch.h,sha256=GKMJr9vxK7Jx3_h48dovgRo_r-kqYAMu3mYfrDGn5oY,3122 +torch/include/ATen/ops/_fused_adagrad_cpu_dispatch.h,sha256=K-5jeNGchNTjQar7CyrpkIH7wolF2XzQISR5JYFYcNI,1546 +torch/include/ATen/ops/_fused_adagrad_cuda_dispatch.h,sha256=Ssu2ujMcrLLilFm_KTZFRbshE9kN1BSPhG1VhVMak6M,1548 +torch/include/ATen/ops/_fused_adagrad_native.h,sha256=Rd2-cIuFk4iAZh8c5SwtSq1U7fZ1EiZVtnnn_YZl1pk,3468 +torch/include/ATen/ops/_fused_adagrad_ops.h,sha256=uqpMtJzCSfY3RU9KPVzTzLpq2pDXXJ-HCw9JVo9MB98,9870 +torch/include/ATen/ops/_fused_adam.h,sha256=4gu7b8llgCs-MKBqdfhT2Gi4eK2uZXeiV-3x6vxZk-I,8657 +torch/include/ATen/ops/_fused_adam_compositeexplicitautograd_dispatch.h,sha256=MkeYXsYf3P7p1ql2q88jJe44WuQbs8fh_0Wa74Syh38,3680 +torch/include/ATen/ops/_fused_adam_cpu_dispatch.h,sha256=gAqHh0CTO7D0BDpnkIwMonn9Zgc1At8562mfd0YCctg,1706 +torch/include/ATen/ops/_fused_adam_cuda_dispatch.h,sha256=AnSRyMEDkoCayIRY_1QTbU9HPS35WtoFRTSeka8sqUE,1708 +torch/include/ATen/ops/_fused_adam_native.h,sha256=KMRl1U7-wuGu1Oa8MxYhhA8VKtUMRoFkCf1MqDJJWOg,4186 +torch/include/ATen/ops/_fused_adam_ops.h,sha256=0er_m6NFg3gGlz4AfvbmNgzq7CWrNdlafTL1mN6XAlc,11848 +torch/include/ATen/ops/_fused_adamw.h,sha256=gjcLKnseiuOg1DjEldDdxeumsWBJdmcmt3KW4iS1_2M,8682 +torch/include/ATen/ops/_fused_adamw_compositeexplicitautograd_dispatch.h,sha256=pRd-OIed7QKpcRGw_26CtP7sAx5Hq4nHkemonbGbJew,3686 +torch/include/ATen/ops/_fused_adamw_cpu_dispatch.h,sha256=T3-O1luEVQZE-7Oi8CyvV0lUrmPgep_UduNsyTMzyNg,1708 +torch/include/ATen/ops/_fused_adamw_cuda_dispatch.h,sha256=uf92VyZ13criF2iWdLabMYGg_uTrmCxQb8-r0ovRLQ4,1710 +torch/include/ATen/ops/_fused_adamw_native.h,sha256=kestFLIiQQnc-1OBlgZsqVlA_roPQInMAwKk_Co1gIw,4194 +torch/include/ATen/ops/_fused_adamw_ops.h,sha256=SFRU5tI9ZXYpEThmlfQzOXUkQ0AOAGfKr7gbwrYXWQ4,11866 +torch/include/ATen/ops/_fused_dropout.h,sha256=NZmoa_Pzhz_doB-Mw6O_9PFAacxLXXIsMFRjEgCcFIM,1904 +torch/include/ATen/ops/_fused_dropout_compositeexplicitautograd_dispatch.h,sha256=_2L2YkqjQNGbJsAGg1K5f11MvRoKqpEATL6jTTdrA0A,1350 +torch/include/ATen/ops/_fused_dropout_cuda_dispatch.h,sha256=Meiu-cfclgDAa-Kw0Sh9Wefnld2FYAlMQok0uaArnzk,1076 +torch/include/ATen/ops/_fused_dropout_native.h,sha256=BOHrkLwgDv_FAdEVQnPXk4SkzEhYEH3h-ubGm3TLKWQ,1019 +torch/include/ATen/ops/_fused_dropout_ops.h,sha256=7djoXYEL90MSOB0PEs89tonh7QZO4hYuEwDIkHDnc4k,2461 +torch/include/ATen/ops/_fused_moving_avg_obs_fq_helper.h,sha256=NVf_adVLWAUqxqNu7AnxvjEDM_DGAZhsI32GlNDiIiQ,5029 +torch/include/ATen/ops/_fused_moving_avg_obs_fq_helper_compositeexplicitautograd_dispatch.h,sha256=TA3YjyVb7hYBuZ0c3G_7AW11HfJfSDOifddUMdK7xDs,2339 +torch/include/ATen/ops/_fused_moving_avg_obs_fq_helper_cpu_dispatch.h,sha256=1R_GDqoo9w3O_9Iozqo2WsSdAAq0M8RN_59lW3LJCe8,1325 +torch/include/ATen/ops/_fused_moving_avg_obs_fq_helper_cuda_dispatch.h,sha256=7QOJ3Oe2OmUakqOaPCbVsZs7IzHh8yDM_Jwqs20Z5DI,1327 +torch/include/ATen/ops/_fused_moving_avg_obs_fq_helper_native.h,sha256=yO2HLpfvuro99T8ygD8_BJ6PBz1javZPxP4q5uc7fQ8,2418 +torch/include/ATen/ops/_fused_moving_avg_obs_fq_helper_ops.h,sha256=V6ywNUALozA4Q_D5puIT3P_r8hgwmJOmmYLJUwPvdaI,6128 +torch/include/ATen/ops/_fused_rms_norm.h,sha256=GOt2xf3p0Nd53N83eXAowzyXj2dI-1mEP4MQ2n07wfQ,1155 +torch/include/ATen/ops/_fused_rms_norm_backward.h,sha256=Nuf4HhveJzXXyqs3IV65tt8ivgpsL1O5SgAWYcnUg3Y,1313 +torch/include/ATen/ops/_fused_rms_norm_backward_cuda_dispatch.h,sha256=57IUFwQsfPEfaKU4sPn9Fz4tUvIzUnHrNEsFgfH7l8M,1186 +torch/include/ATen/ops/_fused_rms_norm_backward_native.h,sha256=ZZNH4XLvDUU7L--pgo4qascNVAnTBCkw6ZptrVBmWhQ,945 +torch/include/ATen/ops/_fused_rms_norm_backward_ops.h,sha256=n_I3yJ0jlX65-DYDGjjVqnPESnxmc0QEdVr8bLr5e6E,1892 +torch/include/ATen/ops/_fused_rms_norm_compositeimplicitautograd_dispatch.h,sha256=mt-FKhSXc4s3WrDg-YIV4RVjsc7in3tT4Q2Ju7At7fE,1160 +torch/include/ATen/ops/_fused_rms_norm_cuda_dispatch.h,sha256=aq7kf92EtEoo5v_AZn8lO1hU1-erwRw6Nj_1lvt1OI4,1118 +torch/include/ATen/ops/_fused_rms_norm_native.h,sha256=_I8UfSoAXZtYXYTnlzwpD64VIcIWBZmm6I59DHSpk1w,1076 +torch/include/ATen/ops/_fused_rms_norm_ops.h,sha256=2XT_ijz0-BLZpO2E_NHDk-cH-DYf7g3Ih7yutXDP_gk,1671 +torch/include/ATen/ops/_fused_sdp_choice.h,sha256=_HEFRv_VRlLRIAoAyRVHjVSUT1FgYIpbNQCCHhz7hV4,1347 +torch/include/ATen/ops/_fused_sdp_choice_cpu_dispatch.h,sha256=aORjujP3gbUp4GtzJ8i7Nks5x4_BZFBqyaGT-C3sW1M,1196 +torch/include/ATen/ops/_fused_sdp_choice_cuda_dispatch.h,sha256=Cu9tZ73xM32Rd5xXD89LUophfZ-EhWZyQVDoUqG3eoA,1198 +torch/include/ATen/ops/_fused_sdp_choice_meta_dispatch.h,sha256=liX_lECLD9EhJbEdSxZXwUEqoc8SVq5yE0AsT6H-WVE,1198 +torch/include/ATen/ops/_fused_sdp_choice_native.h,sha256=VB43OsQZq_zFD2XcsIWclu77TLeUqnpfyppCP-gePaw,1518 +torch/include/ATen/ops/_fused_sdp_choice_ops.h,sha256=3IThPWmUDTbjFeFgtiiErqcdIPv_L4WGtXRFnBYygS0,1853 +torch/include/ATen/ops/_fused_sgd.h,sha256=92RH6gah9UQ13vQk6Ivw1OWA7YfhxWI4-APYavBwEBA,7480 +torch/include/ATen/ops/_fused_sgd_compositeexplicitautograd_dispatch.h,sha256=s3wzfP7iGYLfWyUbicFriGOC2_DK81VJz2NPunVB71g,3210 +torch/include/ATen/ops/_fused_sgd_cpu_dispatch.h,sha256=rPLpZowdauA6q7UbgoKb2FzHK5AMN9kqUxZmLsclxCU,1584 +torch/include/ATen/ops/_fused_sgd_cuda_dispatch.h,sha256=kJv-9IWn3Vx-CT_y8SYTDcg7xvpfKYhGwl9I7AV2Snk,1586 +torch/include/ATen/ops/_fused_sgd_native.h,sha256=sJgU4NMnCQAozx9UTDqJ74befOYo00Sjyw9-vJP7JF0,3594 +torch/include/ATen/ops/_fused_sgd_ops.h,sha256=RU2ZpSyTrED0K7FO5oFc2ailR951pwMA3E1_PcaIRxM,10132 +torch/include/ATen/ops/_fw_primal.h,sha256=eqnSVoaz_S1dGjcflCN7C3SDn4wMk84GlECSG7h7X7Q,761 +torch/include/ATen/ops/_fw_primal_compositeexplicitautograd_dispatch.h,sha256=byVDwgcrFjUY_7-vwtuO5cM3HGhNKnyxF3Cdpq-LvIk,1037 +torch/include/ATen/ops/_fw_primal_copy.h,sha256=EPEcdpjDqsfuOYmrSXlzF59IWbxdrdTUkcPxQV9pc2g,1470 +torch/include/ATen/ops/_fw_primal_copy_compositeexplicitautograd_dispatch.h,sha256=xljGBzOGRheaqvv8CsMovYmard6RU1bfuEElvX0fx7I,1169 +torch/include/ATen/ops/_fw_primal_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=eWCHGSGN23_AouEUgK55Bx8bmY7S2StCo3BAvcy3tpU,1068 +torch/include/ATen/ops/_fw_primal_copy_native.h,sha256=3COu1iUg5bUftO7xZFq7IjBnHo-MIPfKPoxsPkhjK8s,856 +torch/include/ATen/ops/_fw_primal_copy_ops.h,sha256=BqPn0y-ejD9W_kihe-AMQUNNtMnZRiFkHttyO-iUurc,1955 +torch/include/ATen/ops/_fw_primal_native.h,sha256=1hgOMQ8BKAJvouhw1jmoQBjPirudECnFhFovpKzGKtg,749 +torch/include/ATen/ops/_fw_primal_ops.h,sha256=xlvOAkCdGyh5vDHfiWgt4HkytX3h4P-3ICdZf13iVQw,1279 +torch/include/ATen/ops/_gather_sparse_backward.h,sha256=OsMhYE-vtK8MmzizCLwyv_I1wcUrCPPyLrRQd-TvO7o,1077 +torch/include/ATen/ops/_gather_sparse_backward_compositeimplicitautograd_dispatch.h,sha256=WlxYG619OdDjb6h7qXMC84sINUC11CFXhwintaKzxS8,1099 +torch/include/ATen/ops/_gather_sparse_backward_native.h,sha256=Q5mHqsWbalZR7Qz4lMR85b1f8NBzQoqpo8rXEPFQqy0,811 +torch/include/ATen/ops/_gather_sparse_backward_ops.h,sha256=4P-fVu4UPm9K6xXA4ujhTfuAyCDy7rl04P2CsNq-0hw,1475 +torch/include/ATen/ops/_grid_sampler_2d_cpu_fallback.h,sha256=3AGo7iv7NQ-mVBF_e04qljzyNgJ0Gp6aiY4R0AtLDZY,2195 +torch/include/ATen/ops/_grid_sampler_2d_cpu_fallback_backward.h,sha256=JjkgzSbPqWNM6510ku26ajgU6il0TtHpP7ZzvIdW3Z0,1347 +torch/include/ATen/ops/_grid_sampler_2d_cpu_fallback_backward_compositeexplicitautograd_dispatch.h,sha256=BF5HPp-pY2tsVUqI19Q60iQR9JkwrS02gcNYA3TBjQc,1203 +torch/include/ATen/ops/_grid_sampler_2d_cpu_fallback_backward_native.h,sha256=dinXnigeucoYbz-iQ29Rvx-DuTWHhy4KJQEj8keo3XA,915 +torch/include/ATen/ops/_grid_sampler_2d_cpu_fallback_backward_ops.h,sha256=0Fn7oGVShuglrhzNjfscLsc83xiJiondgoJMCE3Td0o,1808 +torch/include/ATen/ops/_grid_sampler_2d_cpu_fallback_compositeexplicitautograd_dispatch.h,sha256=4xPbK_VCXWXbku8LGhy_EQ8rHjF-qXGVxQWm5iBeuEA,1532 +torch/include/ATen/ops/_grid_sampler_2d_cpu_fallback_native.h,sha256=F2bv-kkWS8LtG985EK7-rFfv7dzPnQkKoVsY43xWOBM,1046 +torch/include/ATen/ops/_grid_sampler_2d_cpu_fallback_ops.h,sha256=b5BtmpGqmN7zGhZCp4H_Tlm92EzWt4Nwf3BfHYTfm20,2563 +torch/include/ATen/ops/_grouped_mm.h,sha256=Oknxhi_RW3urjNu3gLQSruaE2B6MhgYU1l7UgEJXDYU,1192 +torch/include/ATen/ops/_grouped_mm_compositeexplicitautograd_dispatch.h,sha256=-GLNHx9IRx864cXNkMEl-tE5VyOnHiEWYdoS7njmAtQ,1196 +torch/include/ATen/ops/_grouped_mm_cuda_dispatch.h,sha256=vrJzHv9jE17zK1iXs7LvmZj6b39NVtmv6qRRKUEST5U,1154 +torch/include/ATen/ops/_grouped_mm_native.h,sha256=KihZtmkZ6QCB0XANpp5Ek3Kqyo8vwGzLdPWonsUKAe4,1145 +torch/include/ATen/ops/_grouped_mm_ops.h,sha256=WHjBMKMfSyAJuDK-we8TZVcti6doLlEiGVvIcJkAAdA,1736 +torch/include/ATen/ops/_has_compatible_shallow_copy_type.h,sha256=JvzVYninQkma7cCLj7A9Ovm-6gB0INlSFDtXPMGAv18,1035 +torch/include/ATen/ops/_has_compatible_shallow_copy_type_compositeimplicitautograd_dispatch.h,sha256=A1ugji9dhSaxQBwA62uN1QkZakN4DBbZyPC99cEmhYE,1064 +torch/include/ATen/ops/_has_compatible_shallow_copy_type_native.h,sha256=lPEEEyzCgCD6u7MJyB4hxR96aOWkYv1jW5h1H7eFQm0,776 +torch/include/ATen/ops/_has_compatible_shallow_copy_type_ops.h,sha256=y1SQrcX2Hhz77UqvYsir5XBF8hvj2jDTG20N_qN7oA8,1355 +torch/include/ATen/ops/_has_same_storage_numel.h,sha256=JsIte5GuSkCizr-dzT3Y-qpPEnTnpmKt1czMJTT8wMU,998 +torch/include/ATen/ops/_has_same_storage_numel_compositeexplicitautograd_dispatch.h,sha256=yk48siOtnOPaKdXKj8frEKZ671d0YARYf3HvYu_vjiw,1055 +torch/include/ATen/ops/_has_same_storage_numel_native.h,sha256=XBykH_i69Vo2mSYsdX5vqR5K97Jrj39_FQTs1amdzjM,767 +torch/include/ATen/ops/_has_same_storage_numel_ops.h,sha256=gDmTLhUD91aM48v_e6_xkqbHbq8LRFNjAMvqidR9Z4A,1328 +torch/include/ATen/ops/_histogramdd_bin_edges.h,sha256=zGcwlj0CPYmN-MjKFr1NvKIRaurObaVTlzv_AW8Grkc,2172 +torch/include/ATen/ops/_histogramdd_bin_edges_compositeexplicitautograd_dispatch.h,sha256=xC4THPV4Q5X53wg7diS--cg7g35rkBDjhVJXBel6hto,1415 +torch/include/ATen/ops/_histogramdd_bin_edges_cpu_dispatch.h,sha256=iBRVd-tGILfHi0Ya6o7Y3Bo_SHVArwotMbFiYwQnsIs,1154 +torch/include/ATen/ops/_histogramdd_bin_edges_native.h,sha256=zUCxrJaZfkmMbXncGENEN2ob2wger14G4MMgZbyHy0k,1122 +torch/include/ATen/ops/_histogramdd_bin_edges_ops.h,sha256=drxSaAGdAH93INVP456MKlYT8EpE1_sxKZ0_GO85Q9E,2769 +torch/include/ATen/ops/_histogramdd_from_bin_cts.h,sha256=In6rhp37etcyN2SGp370g7BIeER6eXnamicm9NP9RpI,2209 +torch/include/ATen/ops/_histogramdd_from_bin_cts_compositeexplicitautograd_dispatch.h,sha256=DRf3RHIugrMTu2KEQ26ub3wiX9X5v91yeZIMOrFObZ8,1433 +torch/include/ATen/ops/_histogramdd_from_bin_cts_cpu_dispatch.h,sha256=hyTkeFAqjh8eAYDp1tytyedRukkstmq6EemWf6Rq0e4,1142 +torch/include/ATen/ops/_histogramdd_from_bin_cts_native.h,sha256=D4Jr4I1_Jw1YUx5El3AKhrRoD6LZQ_fo1Vuz2-bt4Po,1107 +torch/include/ATen/ops/_histogramdd_from_bin_cts_ops.h,sha256=t-ryxmvxK6SgkDNCfat7NNECgqdOpfzcaY53d23BtYQ,2764 +torch/include/ATen/ops/_histogramdd_from_bin_tensors.h,sha256=XXi3LjXFtO86-EVHXkjK1vedEgDRcs_dtfoOX_7wa2o,2006 +torch/include/ATen/ops/_histogramdd_from_bin_tensors_compositeexplicitautograd_dispatch.h,sha256=qC6e9ppTsrpaTPH5_w38avOMRc8nOcl-Wni-nPLihPg,1334 +torch/include/ATen/ops/_histogramdd_from_bin_tensors_cpu_dispatch.h,sha256=5U3Eq_9aVTfxpTcFux73HC2PSSSM8XgRcVQkfw2I2I4,1085 +torch/include/ATen/ops/_histogramdd_from_bin_tensors_native.h,sha256=CWTma5YlLtHjOW6UNldw6SADVVI9lwVVFZQftca_Peg,1004 +torch/include/ATen/ops/_histogramdd_from_bin_tensors_ops.h,sha256=CSj5COc2SiQNRP28-Ijn_rB0FVylF_0mEh7BKf_8EaQ,2488 +torch/include/ATen/ops/_index_put_impl.h,sha256=yLwFi2tr1vKi3ldSW9AAPMy-pI7HaGMtrPOjZ0FdSgU,2480 +torch/include/ATen/ops/_index_put_impl_compositeexplicitautograd_dispatch.h,sha256=ZdxPYVQN11B2vzRQnVuKHZGjfZX4S6frTHbs-oHjCHM,1565 +torch/include/ATen/ops/_index_put_impl_cpu_dispatch.h,sha256=Eckm0NE-gWbQ0Zn44L4d8Y2Dz3qT3u8X19hh-cDUaC0,1105 +torch/include/ATen/ops/_index_put_impl_cuda_dispatch.h,sha256=51gizcWyfrNX9rKTxHlpJPG0b9kwWzllj7WQtu2pR40,1107 +torch/include/ATen/ops/_index_put_impl_meta_dispatch.h,sha256=uL4nByEuc-KdpaN84KzyUQeO1oMu7NDNdTIAaX5m4Ag,1107 +torch/include/ATen/ops/_index_put_impl_native.h,sha256=HAe0daZdssepvdISf4C51B4-MEfH2Qv_BMtmE46FTMM,1648 +torch/include/ATen/ops/_index_put_impl_ops.h,sha256=-EiIDOjLsULTswoDTDw0D3YuDhis1dA9iLnCFuvUKUw,3527 +torch/include/ATen/ops/_indices.h,sha256=B9bgDB6Hd_qIyHiyX0M52n0QoxAxMcbyZZ17MwT1ahE,759 +torch/include/ATen/ops/_indices_copy.h,sha256=jTYE9cJGJHvQK6ykNMDVU9vdNg1L4HEgzZiYEShcmJk,1351 +torch/include/ATen/ops/_indices_copy_compositeexplicitautograd_dispatch.h,sha256=QgR3pNQjcVbUAgB1S4D1qmC8Pk681CP90LAnVBY7UJc,1135 +torch/include/ATen/ops/_indices_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=2nZ49beNBblP71sAv06BwTiGddig-K0EKRsHlye4Fqc,1051 +torch/include/ATen/ops/_indices_copy_native.h,sha256=lvpN2CAOPXdpfRVS7VW2EzhaCtDL-tLCnSGtE2R4xuM,822 +torch/include/ATen/ops/_indices_copy_ops.h,sha256=5ABJhQinJjUS0SGVDeLJ7Mq72a_BRn1CFjtAObVCYlw,1843 +torch/include/ATen/ops/_indices_native.h,sha256=Bviu9rgkpNQsO54VIvqXwxB03TNrnL1rl-cPc6weVQA,739 +torch/include/ATen/ops/_indices_ops.h,sha256=XEzlLxcR26JtHM4DNJl1AdOL6v5sKTxqOrThUaXwkS4,1223 +torch/include/ATen/ops/_int_mm.h,sha256=vP6kGVBO0FozEkwDY2wr5MoZND_axEJ1PmSpihIJakw,1423 +torch/include/ATen/ops/_int_mm_cpu_dispatch.h,sha256=O-kSdVrNEN3K2hAUmU1M3FJTQ1oEvYVDdL9JOSciC-8,1209 +torch/include/ATen/ops/_int_mm_cuda_dispatch.h,sha256=Upp3HbSSNqd9Dfo1XsegSfge29olCIZiOMyz0Zti9i4,1211 +torch/include/ATen/ops/_int_mm_native.h,sha256=qvUnlJjpTFYJNQzC-JpEo8l-GD2wBdjG0L0r9ZgIzgA,1062 +torch/include/ATen/ops/_int_mm_ops.h,sha256=uzON3HOkVfn-VaHs7Uit0oBm5ZK_vJzNAQaVZiy9-lU,1973 +torch/include/ATen/ops/_is_all_true.h,sha256=7HwG7ClGkY-kMz_dV5rFMnI_mnvuB2Ut2uGGjdV4v2k,915 +torch/include/ATen/ops/_is_all_true_compositeexplicitautograd_dispatch.h,sha256=8dGqPyH1aZXRypFevFfh5i_ku2N7L7Ty0M6QlxCXGsk,1024 +torch/include/ATen/ops/_is_all_true_native.h,sha256=WxZLnnRxQKJCuCvlHPhdGVROUJdc3s8n4ZfGMXRW8B0,736 +torch/include/ATen/ops/_is_all_true_ops.h,sha256=oMy7qgZVxPPzYslsvladgqsUWZG29yn9URUI7ODGLgM,1229 +torch/include/ATen/ops/_is_any_true.h,sha256=KZni9CoGBeA812_m_Sy1sLh79a0gHtt_1tb1G6bfMV8,915 +torch/include/ATen/ops/_is_any_true_compositeexplicitautograd_dispatch.h,sha256=GDp3leGCp9Pk18PoJPnehQHpgo4wEulydPMBKr13DSc,1024 +torch/include/ATen/ops/_is_any_true_native.h,sha256=Qfw9XCWtC3ds_cSjW3HEyHJ7GNjAymuospkoEPs9FhE,736 +torch/include/ATen/ops/_is_any_true_ops.h,sha256=mccte8AhGJ8imAiIRogXRVg_Ae8Ov9g_DQ_8ehDAQaQ,1229 +torch/include/ATen/ops/_is_zerotensor.h,sha256=z1izUiqfNhXX6O-8TUfS5F2dWcfz1uqmbvN97yqN4FI,926 +torch/include/ATen/ops/_is_zerotensor_compositeimplicitautograd_dispatch.h,sha256=RQ2qriVDEWDI96lqyzc_dw4Zlk-dKRjGnNwjEQ7uF08,1020 +torch/include/ATen/ops/_is_zerotensor_native.h,sha256=8M3z83uvLl285t7mhTmotaobg5jq1ePFP8saUIIzlUs,732 +torch/include/ATen/ops/_is_zerotensor_ops.h,sha256=aK9Vwvof9XrPNNaPCy7G1Lk6RWe6TiSfDNOwLjTe1xE,1215 +torch/include/ATen/ops/_jagged_to_padded_dense_forward.h,sha256=y14xK8qzPSQHFtF0waSlBczdikPWmpYMgzlqhtl7Wu8,2374 +torch/include/ATen/ops/_jagged_to_padded_dense_forward_cpu_dispatch.h,sha256=62WMJ-oBtP_AI_9A3tGiF8Hqr1fYQNrceLxb-HOwPV4,1251 +torch/include/ATen/ops/_jagged_to_padded_dense_forward_cuda_dispatch.h,sha256=tGGkmrpliFrP_qc8kimo2rDV6sRw17SnYSKPTpJpGw4,1253 +torch/include/ATen/ops/_jagged_to_padded_dense_forward_native.h,sha256=QdOdnQGTpkkL8uJTOPbcT0eMkkD-eWsRET8LPshguo0,1007 +torch/include/ATen/ops/_jagged_to_padded_dense_forward_ops.h,sha256=_EqMYg0vBmrseAg--3FthFdEx4FOamOi8oxbkCk3RLU,1560 +torch/include/ATen/ops/_lazy_clone.h,sha256=y_BJD4U-ki-BqnrvgqTca4LYQi1dCiRmpwKs5s3WtB8,911 +torch/include/ATen/ops/_lazy_clone_compositeexplicitautograd_dispatch.h,sha256=3Uc8hTM0lw3nM9r3SN1rbsFpr7-szP3ZuIcYjXipa1w,1023 +torch/include/ATen/ops/_lazy_clone_native.h,sha256=-KVgQ1ouT7TjUWaTcFBeJbk0GqhM1vLT-quBbIVfbHM,735 +torch/include/ATen/ops/_lazy_clone_ops.h,sha256=Opv7LelkqSvBnwXj8cAYFSQP9pjaLJCppzpIKe4y_k8,1226 +torch/include/ATen/ops/_linalg_check_errors.h,sha256=nN8ew-8Jee5DaXtUGNb6tZtD7v5FK1pvPjlUq2FlojU,1034 +torch/include/ATen/ops/_linalg_check_errors_compositeexplicitautograd_dispatch.h,sha256=Jn6cmdryebetzoNUxd2GbA6Clx_uEDFYhLBVaHjj-fQ,1069 +torch/include/ATen/ops/_linalg_check_errors_native.h,sha256=apqCF9Nk8pWEzLbsNaEoPhhaUgselsBBp1gH34S1f04,781 +torch/include/ATen/ops/_linalg_check_errors_ops.h,sha256=AU7t6aeSziKIFe4qyTGpRMk4bowjBWHuzFdlbOOyzR0,1374 +torch/include/ATen/ops/_linalg_det.h,sha256=7OAdq-cnH1tOx2qHxRBm4FxoUNJpbU5keNnrYBjTDc4,1739 +torch/include/ATen/ops/_linalg_det_compositeexplicitautogradnonfunctional_dispatch.h,sha256=bKs74yufkRN8TSAEprisvYxHfAdNEb0ZT4oCoeKEWeI,1082 +torch/include/ATen/ops/_linalg_det_cpu_dispatch.h,sha256=orTSclzAkSEuVAIKfJZSMP5F7BVvLx_Lh7ESxOHiN5A,1335 +torch/include/ATen/ops/_linalg_det_cuda_dispatch.h,sha256=tv0iyI66MeotWwWbUi0kMVDhSAXpiPdczKG5cFGGtTQ,1337 +torch/include/ATen/ops/_linalg_det_meta.h,sha256=f1gRAY-pVAFngq2-9844Utz8ck9F_clnXeYbVSVXPXU,823 +torch/include/ATen/ops/_linalg_det_meta_dispatch.h,sha256=VkszfbK7XhaY8_omQxtXhr4S_-OdCCoTbhHb9cajuEY,1337 +torch/include/ATen/ops/_linalg_det_native.h,sha256=dJtiMdRi6W9cxhYUccW6XUJ5ZhWZqlXFFtd5N1tT67c,915 +torch/include/ATen/ops/_linalg_det_ops.h,sha256=UexIu5Ngk5t5hCtvGPzunB_0TFO9ALM_46HiDS7pGDc,2275 +torch/include/ATen/ops/_linalg_eigh.h,sha256=pRgxxSyESWhwCzqE3vy7Eincef9iAhU0pYUOhBiT4KU,2009 +torch/include/ATen/ops/_linalg_eigh_compositeexplicitautogradnonfunctional_dispatch.h,sha256=beXmcF-FRIQUi4G_ATVYBpUHlG8iHdKKXIpvoABg4FM,1120 +torch/include/ATen/ops/_linalg_eigh_cpu_dispatch.h,sha256=oXCjYSY9MDYZvrQg9r_LYGw6yO_u9uU8bu7Ut6-sQJ4,1424 +torch/include/ATen/ops/_linalg_eigh_cuda_dispatch.h,sha256=3ECkpLJwdCXaklrZu6m3_xeQa0TOJceBphGOvvIPpdc,1426 +torch/include/ATen/ops/_linalg_eigh_meta.h,sha256=IVjTv1swP6esVMlYBJYmkjSTFD4lJnJVkOMchtHdfMg,863 +torch/include/ATen/ops/_linalg_eigh_meta_dispatch.h,sha256=6BXt1HovirmqvMYO5-hjX9aVyZv8jPKKx6PTH3UZ_Gw,1426 +torch/include/ATen/ops/_linalg_eigh_native.h,sha256=TYq7E-8EwVBrZHz_sKPIV38qg8yUxZ7Sb2N3y466QxY,945 +torch/include/ATen/ops/_linalg_eigh_ops.h,sha256=leg8vw1gKVwoen9nI0HKAYYGtYIfKiCbTQg8A-FjB6E,2468 +torch/include/ATen/ops/_linalg_eigvals.h,sha256=T9GN8LIoVrqq071U2NlrNf81IZ7ucgp8yvtb9Zea5VI,927 +torch/include/ATen/ops/_linalg_eigvals_cpu_dispatch.h,sha256=G8NZm4VDaffheSe3GIrnOCa8yzthMvhZJnmmBc39H28,983 +torch/include/ATen/ops/_linalg_eigvals_cuda_dispatch.h,sha256=lTah4XbL8EqrVXrhmHMMrXMSw3pKObt38mgstUQrzPM,985 +torch/include/ATen/ops/_linalg_eigvals_native.h,sha256=uWniTMiY4vVBx3049rzStty0vBDSi48eCFV-_rUh6VY,739 +torch/include/ATen/ops/_linalg_eigvals_ops.h,sha256=m45gYnwjVooaMSq_u-m814LxQ_GNVHERlDC6fy1WVZ4,1238 +torch/include/ATen/ops/_linalg_slogdet.h,sha256=3sAUHeQZBLIvdEfexYshSjyQxSdKQuApc5A34pAyunc,1966 +torch/include/ATen/ops/_linalg_slogdet_compositeexplicitautogradnonfunctional_dispatch.h,sha256=U-YcUlJxCBR6QfUG7p1cAURY7P_giSJdfUsFh77rYIU,1097 +torch/include/ATen/ops/_linalg_slogdet_cpu_dispatch.h,sha256=IusECkYHKucwoYHVxUzncB5YNROOjVJl--glLSadwPE,1428 +torch/include/ATen/ops/_linalg_slogdet_cuda_dispatch.h,sha256=rx0PrV1M5h-NS9e13oMgjq5EjI43VwOJG475LTFP0OM,1430 +torch/include/ATen/ops/_linalg_slogdet_meta.h,sha256=uyrjqpym_MBMCBulrbwZmwt4KCylcEdLI2Iz_XyqTbY,827 +torch/include/ATen/ops/_linalg_slogdet_meta_dispatch.h,sha256=Wk7bilFLzFAQURTXnLsHlHs2dQZAR0ErPIxmIjNKxeM,1430 +torch/include/ATen/ops/_linalg_slogdet_native.h,sha256=UeJc2hENJ6S862TDnP6SnQm_8lGEorqharLmWPSVVIg,955 +torch/include/ATen/ops/_linalg_slogdet_ops.h,sha256=sALoNjfl9hSjX5NWMzJkdBv560NgMP_ZR9STrWlEBdQ,2479 +torch/include/ATen/ops/_linalg_solve_ex.h,sha256=dFYgvtzzgEDJQJy1w_isoXUR900CGivuSdjlkTgUALs,2360 +torch/include/ATen/ops/_linalg_solve_ex_compositeexplicitautogradnonfunctional_dispatch.h,sha256=xRoO0ptz9eb3VL6KLzuMokk9mKwuN2qc39mgWYX4gaQ,1161 +torch/include/ATen/ops/_linalg_solve_ex_cpu_dispatch.h,sha256=oajvVMzPYPDSWaZNCaPEKH7TjQolwRnG2c62BfXVMXs,1603 +torch/include/ATen/ops/_linalg_solve_ex_cuda_dispatch.h,sha256=XZzsQrJYUa5NpXGFjAFAwV_oTB3pfhRMV93CXjvrMs8,1605 +torch/include/ATen/ops/_linalg_solve_ex_meta.h,sha256=zjTWANQOPPa5oUklsIQv4DLWCWb0s3X9mUcLWq2C_gE,880 +torch/include/ATen/ops/_linalg_solve_ex_meta_dispatch.h,sha256=xyp25fmuR2ZxonzHqqRMLpaNXdDnewcf4F4FsWdCCEM,1605 +torch/include/ATen/ops/_linalg_solve_ex_native.h,sha256=7Bo6VLU6GULd681DXVVz8vHDQLofuCD2nruWDHYcVUY,1007 +torch/include/ATen/ops/_linalg_solve_ex_ops.h,sha256=dW2Z_nCry3rR65vH3ZuN5BqYTyrhHna5wQVMIFbHw2Q,2853 +torch/include/ATen/ops/_linalg_svd.h,sha256=rNbYo18rgipCI1IRuGBhrnJ3zEZ0kX8amTfIyGuEUVE,2224 +torch/include/ATen/ops/_linalg_svd_compositeexplicitautogradnonfunctional_dispatch.h,sha256=nzl8wivTY3HCfFtAqobddkPz2H7UJhpkXnfrOJwO7HY,1187 +torch/include/ATen/ops/_linalg_svd_cpu_dispatch.h,sha256=OW1qhWXr2EKS4p6R3SnE831kldRJSHMfJUudywrO79c,1604 +torch/include/ATen/ops/_linalg_svd_cuda_dispatch.h,sha256=grQTlXIQU9rWvdsiBc5y3JieZzwdS-NLJMwaWgRhMLw,1606 +torch/include/ATen/ops/_linalg_svd_meta.h,sha256=PAJddROtT27cyo1bPzNE3E15nrqliYQRCE44Pd2wXKg,902 +torch/include/ATen/ops/_linalg_svd_meta_dispatch.h,sha256=OqF4rnwzKOJUq1Uo8ZQwe6prqyhwo8EWwVs-0e9vK-s,1606 +torch/include/ATen/ops/_linalg_svd_native.h,sha256=BWyvEaKXH5ehBTbEYQiulTqrXNzKaTTYB7Pojg6tlAs,984 +torch/include/ATen/ops/_linalg_svd_ops.h,sha256=AT7NxPanqKGUyycJ0dGS8t8yo63SmMwRb9KEqiLcmcs,2755 +torch/include/ATen/ops/_local_scalar_dense.h,sha256=YdQi4kyJSjnkOxmDgn58MVsByVP34zT6kbhxdOhUxqo,943 +torch/include/ATen/ops/_local_scalar_dense_cpu_dispatch.h,sha256=B4RkqORNP85-ldpilHuVkSnglFod9vSkXvlXBcyAS1k,987 +torch/include/ATen/ops/_local_scalar_dense_cuda_dispatch.h,sha256=SJLgmAc1nEEO8_JOc0LNeCjakVubV2gxI0cGkcX8Gz0,989 +torch/include/ATen/ops/_local_scalar_dense_native.h,sha256=IGmKiqohSHMVzGSGZI48oyDqilpmkGxWKTQbV7Vzwu4,819 +torch/include/ATen/ops/_local_scalar_dense_ops.h,sha256=JocN89AMOnwG7MlkC8N_P5zJjXqMjAial5wlvp4IfRw,1250 +torch/include/ATen/ops/_log_softmax.h,sha256=Vc17KWDatvvzl_jFAiDtSib5rKPaNJOScIWSRAEJboM,1587 +torch/include/ATen/ops/_log_softmax_backward_data.h,sha256=cENGT2LAJCsZuaV3qTr2DgI8kxQixrl9BRo1nH9VTzg,1970 +torch/include/ATen/ops/_log_softmax_backward_data_compositeexplicitautogradnonfunctional_dispatch.h,sha256=31ElJCAOKNJOT9j6ScR1aADbeWtcGR_rZX8-wWLI-Ho,1139 +torch/include/ATen/ops/_log_softmax_backward_data_cpu_dispatch.h,sha256=Da92Ve5wa9zL1poYggLwMNMafOjZOwxsNnVhOqlwj2A,1416 +torch/include/ATen/ops/_log_softmax_backward_data_cuda_dispatch.h,sha256=kq12AV2s0b0g51RsYhVE_59tGSfi_axRN_GK_InDxxc,1418 +torch/include/ATen/ops/_log_softmax_backward_data_meta.h,sha256=RevOxSYwpFwekZx6a4t38aH-7nTFT_eb8SZ_AAVCD5w,916 +torch/include/ATen/ops/_log_softmax_backward_data_meta_dispatch.h,sha256=GxbPLRJwDak0NB8zPNjxu1gjf0X3xrUrqOSoWv9PssM,1418 +torch/include/ATen/ops/_log_softmax_backward_data_native.h,sha256=54HTvaEO-PloYuHc4NsYNc5G67oJuHbV9k2M1pKfqT0,1238 +torch/include/ATen/ops/_log_softmax_backward_data_ops.h,sha256=7AR2cy_pkPWl0H5q1t3uv-rLLlZZSvD9QfITzxrhgb0,2421 +torch/include/ATen/ops/_log_softmax_compositeexplicitautogradnonfunctional_dispatch.h,sha256=bfKJtzt_jZ_HHW_KjGEeoMFMM-cMAjz6-rbCTtufgzI,1083 +torch/include/ATen/ops/_log_softmax_cpu_dispatch.h,sha256=iGpInGk53v4qMD1ti-0Lsptov0pjY8tWmPVq6q3M83w,1248 +torch/include/ATen/ops/_log_softmax_cuda_dispatch.h,sha256=VC_ulzLlh443M41Ba5uOVDD4LK3Hf2lG1S2ULps3jdc,1250 +torch/include/ATen/ops/_log_softmax_meta.h,sha256=cQDb6O1X-OksB9qIQMj3iMrFkyRevOVCCsP_W0No4gM,860 +torch/include/ATen/ops/_log_softmax_meta_dispatch.h,sha256=xl3kcn8NY_8mPCD3z2bfxS_BEEKBIYD71LyUORkKS1A,1250 +torch/include/ATen/ops/_log_softmax_native.h,sha256=laZ31Pmq6KA045SMNjHp2Xi2wl5TBs-oqDQ9TPv_0g8,1094 +torch/include/ATen/ops/_log_softmax_ops.h,sha256=dcDZPT41LhVi8xoPJwsx6vIMrG2zYUyq5c5o-ars9n0,2057 +torch/include/ATen/ops/_logcumsumexp.h,sha256=k1kTuCmZWazs-DrmDS-NWlpflZKZ2r_O-CCGxk6LMg8,1432 +torch/include/ATen/ops/_logcumsumexp_cpu_dispatch.h,sha256=xOQdURtU8lDYB9qoa5IWk5c2dFWyIAvQb8z0EHIHLr8,1191 +torch/include/ATen/ops/_logcumsumexp_cuda_dispatch.h,sha256=TfCloLY4vg-Xun6dMYv31tNQdfuLsw66VpWmga77WUE,1193 +torch/include/ATen/ops/_logcumsumexp_native.h,sha256=Rm-xxbppAG4ZkMUZbcR5wj7W4clKvmggfLHFr9ljDjs,1038 +torch/include/ATen/ops/_logcumsumexp_ops.h,sha256=eyCNGjqtjok3vbXjJcw436KisudTv8C6jEZeDXu0-p0,1931 +torch/include/ATen/ops/_lstm_mps.h,sha256=Q8q4tEYEna7gGmLQZWkymWXC6Y1Fk6xVI9RanVHS1F0,3198 +torch/include/ATen/ops/_lstm_mps_compositeexplicitautograd_dispatch.h,sha256=veskHXK-6LMES7KxDpiCVjpQG_xA4VhcGS-3XCX8VAs,1769 +torch/include/ATen/ops/_lstm_mps_native.h,sha256=-si0UsLrZdgtcgTosZhKkcRVjfKoHQ5TJgcegPdp2Pw,1078 +torch/include/ATen/ops/_lstm_mps_ops.h,sha256=OAP9-bU8ddOeLRuDHIavhyyeK8pMXYCJY7d-eArn8HU,3703 +torch/include/ATen/ops/_lu_with_info.h,sha256=bAzbZk274x74UBS70O7zYWFpM0tZLjcidN1hkvQBsWI,1091 +torch/include/ATen/ops/_lu_with_info_compositeimplicitautograd_dispatch.h,sha256=pprMtCFldmT7wZ5VLhzDukxs74S7ZP9ton1aaZ8vioY,1102 +torch/include/ATen/ops/_lu_with_info_native.h,sha256=4UmSNjUC3z80AMPEsBspYd4gDQd3bokM3mlUvaf9u14,814 +torch/include/ATen/ops/_lu_with_info_ops.h,sha256=EPoJ9siquxaVIpnE8CpuXCg6yDjSJziMK-9sneuQ_M0,1488 +torch/include/ATen/ops/_make_dep_token.h,sha256=L4RMR-RO4zC7Aq0VldcN5WoMQZZ9YdCfo4XR7rW4Zvk,1797 +torch/include/ATen/ops/_make_dep_token_cpu_dispatch.h,sha256=aSM4FB2peO1e0nita16J7I9ON3qOWJd62N4OJCsufXw,1284 +torch/include/ATen/ops/_make_dep_token_native.h,sha256=f-1ezQBz8VLD9t3hDX3gv2zX0CVA8bGvCqyr0txeNTI,939 +torch/include/ATen/ops/_make_dep_token_ops.h,sha256=CjR3gH5mAK7DvPeMaYcaD9U9m-du5F4YfcuA5NyLxhQ,1818 +torch/include/ATen/ops/_make_dual.h,sha256=kShiyoMAx4XfV722yNcvJRqVkrL9dSZG9eDj0O-JIOk,1005 +torch/include/ATen/ops/_make_dual_compositeexplicitautograd_dispatch.h,sha256=_5LzYcN2OtDkOnxlHaGTUs8pQU1vYMPSiklhLw_QZrQ,1067 +torch/include/ATen/ops/_make_dual_copy.h,sha256=gvng4RVbzjFRg2cbfFIxoe5bDZfDYEpfsGZYSToTzQY,1647 +torch/include/ATen/ops/_make_dual_copy_compositeexplicitautograd_dispatch.h,sha256=iL8tCYcu0SelFb5FIs6kQZpsWiMZvvAimWeqBojnb8I,1229 +torch/include/ATen/ops/_make_dual_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=QciZdYPYQHzO5ITk-3_BNgLFiVlAPpcCga1uV0fc3KI,1098 +torch/include/ATen/ops/_make_dual_copy_native.h,sha256=46emE78H2drJ1kU2hAr0a7niFnkSyCcrSGHljHSGXtk,916 +torch/include/ATen/ops/_make_dual_copy_ops.h,sha256=F21kvgkZxmJkJINPnF1duFWK4YJo56QxGqZNc57-TUY,2151 +torch/include/ATen/ops/_make_dual_native.h,sha256=PmcU0bylU7qk6ArXzWZsAa0wEeDDCgjrKe1BjKtzRlM,779 +torch/include/ATen/ops/_make_dual_ops.h,sha256=_-9UOxoXCUdh_Pu7PgNIUH05cCqz4-wVt26rUANLl0M,1377 +torch/include/ATen/ops/_make_per_channel_quantized_tensor.h,sha256=Nlm-3RbfYmYBuzosO9Doo7whGxir74q4q6lZf5j2LI8,1978 +torch/include/ATen/ops/_make_per_channel_quantized_tensor_compositeexplicitautograd_dispatch.h,sha256=ub6HPt7XX-MKCXTfAL70BWI7Scs9k8hXoSGKj2Yxjf0,1319 +torch/include/ATen/ops/_make_per_channel_quantized_tensor_cpu_dispatch.h,sha256=SCf4HJ2Im114u_Bq1wzvvyl29lBfoKkllpOhsJb9zSg,1073 +torch/include/ATen/ops/_make_per_channel_quantized_tensor_cuda_dispatch.h,sha256=6mLl-TQM8YIUriwhJoGVhe7rKe1uq7L6NtjRY3lEL2A,1075 +torch/include/ATen/ops/_make_per_channel_quantized_tensor_native.h,sha256=ZfAJxbGdz1eD4iiaZ8Casp1yI5Tu39cYiM_OypYMOzM,1166 +torch/include/ATen/ops/_make_per_channel_quantized_tensor_ops.h,sha256=HsnKiZ4jnftpDi4LYEiOKELDKsZTqgbpnSYBqyoB8CE,2437 +torch/include/ATen/ops/_make_per_tensor_quantized_tensor.h,sha256=eW6yd2Yu7gndltfIQVWPar92bhA1lXi-C7rFcdnygKQ,1797 +torch/include/ATen/ops/_make_per_tensor_quantized_tensor_compositeexplicitautograd_dispatch.h,sha256=mancFkP3TRUbu4vM2Eeems9-OoZD-RPbSOGhHy-SVJY,1243 +torch/include/ATen/ops/_make_per_tensor_quantized_tensor_cpu_dispatch.h,sha256=MZXM5SWZCJVWO8jjfZp1_isKyp83bzP3CovSLx0UZug,1035 +torch/include/ATen/ops/_make_per_tensor_quantized_tensor_cuda_dispatch.h,sha256=Sacy14ZrAHsZTeWpscCnTS34ZHfUZ8jwSOdm3BtxI1Q,1037 +torch/include/ATen/ops/_make_per_tensor_quantized_tensor_native.h,sha256=odWNKoJcIz6RAw76S-YuHKb_GR3zy7EVLGxspffxwNg,1052 +torch/include/ATen/ops/_make_per_tensor_quantized_tensor_ops.h,sha256=lbdp-9kHyX5-N53dtXG8bRDj-expygmvKS8JENwNN74,2191 +torch/include/ATen/ops/_masked_scale.h,sha256=2l0PqbGWQO_CpqhyEpXltfmyAa99oBH4qXiTHG2iS8Y,1585 +torch/include/ATen/ops/_masked_scale_compositeexplicitautograd_dispatch.h,sha256=CbCAXOjrQSx4NE904hBOGVejXfLxIj4j1Xt2OWXoyts,1213 +torch/include/ATen/ops/_masked_scale_cuda_dispatch.h,sha256=09MMY41P8TkQESTMppFHDy99lCGM_H0LsyI_-UdWH14,1022 +torch/include/ATen/ops/_masked_scale_native.h,sha256=glMoa-DUeZjO4l-zd7wbpNbP6VzzHtOkGJaH-H2kT4Q,904 +torch/include/ATen/ops/_masked_scale_ops.h,sha256=KJnnCB5izfAv0SEdBb09XryasK1bkqo38fQFiApS9n0,2107 +torch/include/ATen/ops/_masked_softmax.h,sha256=hif_h_cXz_HLZrm5MBR0-8CJsYw3UR0ocuMm0ElIbmQ,1917 +torch/include/ATen/ops/_masked_softmax_backward.h,sha256=dNeFW72xVQiuOaOHhVRGpcV6PeGYSo2fC6yrmqd9cI0,1986 +torch/include/ATen/ops/_masked_softmax_backward_compositeexplicitautograd_dispatch.h,sha256=QYfH6D1cLPzb1TWcxCI1Qq7y_qlDKcb8jpsXeVyMuec,1350 +torch/include/ATen/ops/_masked_softmax_backward_cpu_dispatch.h,sha256=O-F-uzYX3b3Y-5ArdAcAAU6R2CAbOwdNhkMiHNw0b9k,1096 +torch/include/ATen/ops/_masked_softmax_backward_cuda_dispatch.h,sha256=hN3CgsvM4rxDBeMzkpIp00PS_EZK8YIygocw5UqE98s,1098 +torch/include/ATen/ops/_masked_softmax_backward_native.h,sha256=-scXyc1NA9QFNW1mUqDpCpFYkORstxazi29e4q5BQUs,1220 +torch/include/ATen/ops/_masked_softmax_backward_ops.h,sha256=7F_9wVT4_BfGX_TAHF-tCiWO91qBhWU1_p9JbWeiD9c,2497 +torch/include/ATen/ops/_masked_softmax_compositeexplicitautograd_dispatch.h,sha256=-zRxlmj8qQPkIAxQJk4B6gLxRDXej_d6M0ofM9HLV6Q,1351 +torch/include/ATen/ops/_masked_softmax_cpu_dispatch.h,sha256=8J5e6cVoL56S5v5ZHxY1oSSPsS877OCydHQhc8jA8-o,1104 +torch/include/ATen/ops/_masked_softmax_cuda_dispatch.h,sha256=fwISicEL5mnjA3w1UmCrGocbDbHA5b0UdWcWN_Emvuo,1106 +torch/include/ATen/ops/_masked_softmax_native.h,sha256=7CvF2U2YXTC0anqzpI14R8TEFoeEF8_82RVOBwJU2jc,1229 +torch/include/ATen/ops/_masked_softmax_ops.h,sha256=aoI36gcmhfoHhy-uRUtH7PucjBf-bE2oy5K70ZEpoh8,2461 +torch/include/ATen/ops/_mixed_dtypes_linear.h,sha256=7qVLBSmVjZQbsqjAyHfn26Kw0kIx8p9qoYajUjqnz0I,1215 +torch/include/ATen/ops/_mixed_dtypes_linear_cuda_dispatch.h,sha256=DnWTWYlSCjRjcxY7DjwmCHYW_U2j_nDZvRfqN-bsHwM,1150 +torch/include/ATen/ops/_mixed_dtypes_linear_native.h,sha256=ZM_2YNgYIAQDjIlwDJx1ZHJSk6_Lb6h9tZlVJBCiw08,904 +torch/include/ATen/ops/_mixed_dtypes_linear_ops.h,sha256=sZd6bOLHXTUUyVqG0NSCyvI5uJeVY7oVbMxcC19MjTA,1723 +torch/include/ATen/ops/_mkldnn_reshape.h,sha256=X6ifDS83CzjgdFdaQgb-YB4cyW-NKV4yfXmI-gMWiTs,1500 +torch/include/ATen/ops/_mkldnn_reshape_compositeexplicitautograd_dispatch.h,sha256=9fdIyoMF2So9lhmv0XYnz-pFXiSiBNL0V80YjrO3Wmw,1185 +torch/include/ATen/ops/_mkldnn_reshape_native.h,sha256=dYQO5_IxIAOTPpNozsd4pCi4_EWjXuNbl_76qb_Cjj4,871 +torch/include/ATen/ops/_mkldnn_reshape_ops.h,sha256=a2FMs-xL1Ff3Yuz95cpOfkBGB6kghL3daVbxwBHbXLs,2007 +torch/include/ATen/ops/_mkldnn_transpose.h,sha256=DLH5rxwrtcKSUvD1rop6x4vjybnbMC8YWHCyIISGuNs,1806 +torch/include/ATen/ops/_mkldnn_transpose_compositeexplicitautograd_dispatch.h,sha256=fdvNDAUeLrJW7BYvsAhCiSr0l53PhnLaWmYm-YZmq1I,1199 +torch/include/ATen/ops/_mkldnn_transpose_meta_dispatch.h,sha256=aAWHv7QMsJGJRadz1hoLQ-nOhjn4RD8udT5PbthVMrI,1012 +torch/include/ATen/ops/_mkldnn_transpose_native.h,sha256=S-2LxbWOFrjLANO3X1vxNF1A_lHli4snraD3PKlMMV0,974 +torch/include/ATen/ops/_mkldnn_transpose_ops.h,sha256=OyJQOpMbEOla-IvUjltnFnHT-56teDq9HkYvaxclG5Y,2681 +torch/include/ATen/ops/_mps_convolution.h,sha256=WOXceZcaOz2AurYWbUn3DPCkUWDpzMx8q0aOjmf3GCc,7109 +torch/include/ATen/ops/_mps_convolution_compositeexplicitautograd_dispatch.h,sha256=Hpf6-vaCasZA5SuiH_xCYnxVPgvgvJkRugS-fhfe-z0,2004 +torch/include/ATen/ops/_mps_convolution_native.h,sha256=N7hCaItIiY-LZeJJMNfu9yUnVeNBKshnUzeC70tdD70,947 +torch/include/ATen/ops/_mps_convolution_ops.h,sha256=SxGn8w9cM6HVfUDKBMER99ecpTcUBD4eTgRmeGo4J_o,3027 +torch/include/ATen/ops/_mps_convolution_transpose.h,sha256=7rDbr402p4OQrfzU7K-vpbwFtVzMQ1NRR5aj50f1T8M,7665 +torch/include/ATen/ops/_mps_convolution_transpose_compositeexplicitautograd_dispatch.h,sha256=yev_85hfQeZSFMql1Lm_Y-nV_LKXPsCwUYsRoaa6NPQ,2012 +torch/include/ATen/ops/_mps_convolution_transpose_native.h,sha256=sCoee66jIuAgDSTmqebXxxrkB4CmpFrlEcSNB2sAiPY,951 +torch/include/ATen/ops/_mps_convolution_transpose_ops.h,sha256=8XIIfGfZR_x1tamDqpVJ0oFR0pEmjQpqtYsJa_cEZys,3053 +torch/include/ATen/ops/_native_batch_norm_legit.h,sha256=Piv9Qzj77TBu3ZZYn4C-x4Lfp2C1RCTjZa2-c-XqZKw,5629 +torch/include/ATen/ops/_native_batch_norm_legit_compositeexplicitautograd_dispatch.h,sha256=LZoJ3rEVP7f0kzlpCsLBN-ETh__BgILrCDGKovFWYOY,1301 +torch/include/ATen/ops/_native_batch_norm_legit_cpu_dispatch.h,sha256=9bBXrY0yYBwv9vSjsOvxWRIxLKIbT585xYC9JPEUU7k,2827 +torch/include/ATen/ops/_native_batch_norm_legit_cuda_dispatch.h,sha256=uP7iwdONIPLFU9MA5Kev2JD5LzvDaTv--HDZz5dFVKI,2829 +torch/include/ATen/ops/_native_batch_norm_legit_native.h,sha256=dLHhP24L4Pb8wKx269ju3URq5dlwm43yAdp6GLrREbA,4005 +torch/include/ATen/ops/_native_batch_norm_legit_no_training.h,sha256=40_UY_kSvZd8eWfDtghDl9CAZHcZHFWgPAoDjKuz-2I,2952 +torch/include/ATen/ops/_native_batch_norm_legit_no_training_compositeexplicitautograd_dispatch.h,sha256=4mexEc_nSHC7KN5aSO6m8Utqf2AcIJi0RpjIxRaAoo8,2002 +torch/include/ATen/ops/_native_batch_norm_legit_no_training_native.h,sha256=ZeA9PDWZcIbwgBO2Y34Z1e48yqJnqrpN9GouLwlc1tk,1338 +torch/include/ATen/ops/_native_batch_norm_legit_no_training_ops.h,sha256=iWJJLKTjtfey9lbDR97TTqVMhURqazBwJi9CRMvkcHU,3576 +torch/include/ATen/ops/_native_batch_norm_legit_ops.h,sha256=Ig2LpSR-assOsIkywsYFSuHSxliMa0YEvLXnZnluREA,7612 +torch/include/ATen/ops/_native_multi_head_attention.h,sha256=7ChUggNS4MauHuDhmoA6AXG-7RRGDrE5eYWdGetW55M,3843 +torch/include/ATen/ops/_native_multi_head_attention_compositeexplicitautograd_dispatch.h,sha256=2DPz1ZIw_wXaqqXHLIcauzHkRrXoDrIAFICgBYA4egI,1955 +torch/include/ATen/ops/_native_multi_head_attention_cpu_dispatch.h,sha256=xn53Go1cIV3oZOGa521HlsXAwiNFzh7opb8Zl3ptv5A,1383 +torch/include/ATen/ops/_native_multi_head_attention_cuda_dispatch.h,sha256=iaiZ3v6VXOPpS9v8fE_WbkYsucI3Z_CO8mJmCVSqeCs,1385 +torch/include/ATen/ops/_native_multi_head_attention_native.h,sha256=0y0nK_orB9mfl36OyifNQbFZuY70LZdJkcA4MZPycCA,2090 +torch/include/ATen/ops/_native_multi_head_attention_ops.h,sha256=Z5sIbRr50PEYV5B19-mwB2D9A2Y_FS6dtkyc2US6nrY,4399 +torch/include/ATen/ops/_neg_view.h,sha256=p7o-UzQA_Dn32bEOK__65k0ibZ13Iqa6hLOStPG33H4,909 +torch/include/ATen/ops/_neg_view_compositeexplicitautograd_dispatch.h,sha256=Ex6ZBQqWqarYoFbdcZpsrXs9PQRJg-xWwEuEISKpwL4,1021 +torch/include/ATen/ops/_neg_view_copy.h,sha256=sWVQ2X9Zy0CBuDCybbd6xYyCOvPX67t_CXC9XTK-m7c,1361 +torch/include/ATen/ops/_neg_view_copy_compositeexplicitautograd_dispatch.h,sha256=pDPtjQhbYYa-yTf1PMUz-u61Kqq4HtkMZ2LHlBXhejc,1137 +torch/include/ATen/ops/_neg_view_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=078FA_mJpLZmv5GTNNTO6unMJ9y85yQv-JAA-MrcmhY,1052 +torch/include/ATen/ops/_neg_view_copy_native.h,sha256=fvwXd21CjdtNsUMNaURx3po0Fvp48gW7zdRV5ekONGs,824 +torch/include/ATen/ops/_neg_view_copy_ops.h,sha256=eMEN_ebhUr1DTYPmJr6zd3K74mBHdEXIT5T7evPdEx0,1849 +torch/include/ATen/ops/_neg_view_native.h,sha256=BD2ddBGe-P4PZ5ZMQ5LjKxRP335jrmYBxOnd5ouQBw8,733 +torch/include/ATen/ops/_neg_view_ops.h,sha256=bljzYKV-uoFUsGm3B9tief60GpoGlKjwRfqV1Ul3E_U,1226 +torch/include/ATen/ops/_nested_compute_contiguous_strides_offsets.h,sha256=eJz39LIIOscRhVZYQtZibgoDz9o6kBguyX4UBL2h52g,1091 +torch/include/ATen/ops/_nested_compute_contiguous_strides_offsets_cpu_dispatch.h,sha256=E0AQGIcTB7jxSksViD8VaeGGXwemIWzkE6Y8gE_pHfE,1042 +torch/include/ATen/ops/_nested_compute_contiguous_strides_offsets_cuda_dispatch.h,sha256=p2rDWVcVbQpc8-pm6k9o2ElJFmgIXgDG_lPkD5yjc2k,1044 +torch/include/ATen/ops/_nested_compute_contiguous_strides_offsets_native.h,sha256=S8gaFsTR150w-fNgrTBcibUh6Px8wsllYhKbhO4b4eI,798 +torch/include/ATen/ops/_nested_compute_contiguous_strides_offsets_ops.h,sha256=zUlSz4sUYEkhQW3mH1HZUFN8S4wmIZPCBhVkjWfhT3I,1425 +torch/include/ATen/ops/_nested_from_padded.h,sha256=qbPDTatCYiUIFUGOOgWIQdxL3n-2402UY2uEI8blr20,1990 +torch/include/ATen/ops/_nested_from_padded_and_nested_example.h,sha256=MwsHQ-RjZyiOPzdWVv0NaFX3Jan7ghsaS9a2iPyheS4,1805 +torch/include/ATen/ops/_nested_from_padded_and_nested_example_compositeexplicitautograd_dispatch.h,sha256=EEfOuSpqFC1OI_b_rc6PnWXYVnOV_5ZYYZmQ8ztvkJo,1251 +torch/include/ATen/ops/_nested_from_padded_and_nested_example_native.h,sha256=GClSqEA2CKgxbj5tbWc-M3dKrzyvvnwH76RAM8qoUNM,943 +torch/include/ATen/ops/_nested_from_padded_and_nested_example_ops.h,sha256=OW0dR3yqWtc6kohw6esBAA_1ZRfFikNsl3meXbhynRI,2207 +torch/include/ATen/ops/_nested_from_padded_compositeexplicitautograd_dispatch.h,sha256=fhCMm5O1pPZAuNxiflhVEcUpBh-yrnww-dyXu2KJpsc,1299 +torch/include/ATen/ops/_nested_from_padded_cpu_dispatch.h,sha256=y259uoH0wavd5fhAeGUHPvwNfAnTqAPUWfqNIMWxi5g,1066 +torch/include/ATen/ops/_nested_from_padded_cuda_dispatch.h,sha256=mKaHoc_6LuuywBKrMAwFllxj8b_KS7Tm69Z1T-aruFI,1068 +torch/include/ATen/ops/_nested_from_padded_native.h,sha256=GB5KLo-ARcVxt3GNUk8KaA0IbuLt95QTkOUihO4p-C8,1143 +torch/include/ATen/ops/_nested_from_padded_ops.h,sha256=w1O1u_hnp4q4i28qdp6gX9QCBXa_Peo__phlB-rbVoc,2357 +torch/include/ATen/ops/_nested_from_padded_tensor.h,sha256=vL_QabkoP93d0va5WPoo9uKY8CL7kv2IiPXOWg6pJDI,3191 +torch/include/ATen/ops/_nested_from_padded_tensor_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_nested_from_padded_tensor_ops.h,sha256=UG33LDfo_u6hJbCqaVR6Vtz___z-ocwZ-z2ulXiFLsc,1960 +torch/include/ATen/ops/_nested_get_jagged_dummy.h,sha256=u0t9Dp8KGzMZt6c2BrnuuflNqBzdGtX0pJ3DYrh77XY,960 +torch/include/ATen/ops/_nested_get_jagged_dummy_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_nested_get_jagged_dummy_ops.h,sha256=txHoQqMUEOrTWo620X4dXvh9tgSeFmJ51xshJJZ_bsQ,1262 +torch/include/ATen/ops/_nested_get_lengths.h,sha256=0TiBIJyyCdg5tExL3WnB18xRet0P1gaji7c1fFejr3c,943 +torch/include/ATen/ops/_nested_get_lengths_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_nested_get_lengths_ops.h,sha256=D43NKE_z8wB0-j8jozy_HyeiLyzoxbsgz2RcNTihqag,1250 +torch/include/ATen/ops/_nested_get_max_seqlen.h,sha256=nzZbB18DiKvmKWAw79alg6_ikm0QXkqpy-chSj87FWU,955 +torch/include/ATen/ops/_nested_get_max_seqlen_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_nested_get_max_seqlen_ops.h,sha256=Z73wu4YKRc05gORZcuuj_pV_8_o0BMmOgw28lkBGsgE,1259 +torch/include/ATen/ops/_nested_get_min_seqlen.h,sha256=syxLMrzaRafldlb72DYVtRNuX9hj61-BNZMA0mhJVWw,955 +torch/include/ATen/ops/_nested_get_min_seqlen_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_nested_get_min_seqlen_ops.h,sha256=22dLbSFfiEMotNaEDqq_vr_HYWLDBxzVVYenVLx_DVY,1259 +torch/include/ATen/ops/_nested_get_offsets.h,sha256=Zq67eFwT3BN5qC-tFC8lNg53m_bmUOdf3xt1lkSXUmo,943 +torch/include/ATen/ops/_nested_get_offsets_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_nested_get_offsets_ops.h,sha256=Lixu-4IOsN0o9Ng_rEVezkVkd8pIn1weKPfqZ8s-mTI,1250 +torch/include/ATen/ops/_nested_get_ragged_idx.h,sha256=4m0HMSdv6asF1mFwBx3otAqvv8T8rnomhgCFZykdEXk,949 +torch/include/ATen/ops/_nested_get_ragged_idx_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_nested_get_ragged_idx_ops.h,sha256=nrxTrL-h-PqazlXSFu5kHzUEXzA0v93QBoahEAk5q7k,1247 +torch/include/ATen/ops/_nested_get_values.h,sha256=r_ni-_UlF_M3fXAc0E5A3qPQ1bQzNHpi1KDFm9j_WwE,945 +torch/include/ATen/ops/_nested_get_values_copy.h,sha256=IGi1nzacRvkXiAbKdOth_axdtF6mMgcRdo-NGmabOqQ,1451 +torch/include/ATen/ops/_nested_get_values_copy_compositeexplicitautograd_dispatch.h,sha256=wHk_loD_qpswu7pBICrJs5G8ZE74eU3iAUhufPGzSXY,1155 +torch/include/ATen/ops/_nested_get_values_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=1l3lpyfq0q9IwynINogr7fsDGU3XWb5XZcNmJ_XyZU4,1061 +torch/include/ATen/ops/_nested_get_values_copy_native.h,sha256=g6noUV2MhXFSbDsHqL5gcjIFLyYCs5xhtJqIzMXAxnM,842 +torch/include/ATen/ops/_nested_get_values_copy_ops.h,sha256=sYjF9AxHizL4ZLmsCpUC5FaeFWvqz0DPV6jMLoxoJLY,1903 +torch/include/ATen/ops/_nested_get_values_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_nested_get_values_ops.h,sha256=qKjIL9j2liwOY93UXYDjiYu7ctUuQa_f-WYWC1ZPhkE,1253 +torch/include/ATen/ops/_nested_select_backward.h,sha256=tcQkdS4RMMmOHU92vqMTT1WSxSMvVGXsX7AzZonwKt0,2046 +torch/include/ATen/ops/_nested_select_backward_native.h,sha256=SF90xp5J-pGQ53Eet8pI1h4mwryadrqCQf59b4jXlyI,818 +torch/include/ATen/ops/_nested_select_backward_ops.h,sha256=Lcb-XM1bIG4DVKkG_3EqdJII19NEtibzhU0bx7BtODE,1475 +torch/include/ATen/ops/_nested_sum_backward.h,sha256=6uiixQaF-kPYwuP9QRfPzP0vgTy1FetkCzLLspewfcE,1087 +torch/include/ATen/ops/_nested_sum_backward_native.h,sha256=ZWqtD1k7li6LxgVdjoh-iMwXG5dbcHLtjwKxV2EJs7U,822 +torch/include/ATen/ops/_nested_sum_backward_ops.h,sha256=RIx-yj5qGlPyOz35NA-Ibv0qzh8emnw8PK4yJqP8rYI,1486 +torch/include/ATen/ops/_nested_tensor_from_mask.h,sha256=6uMwnpvqUOrTkupzZwhGz1W_D2sbf2Hj_BgoD6umLm0,1729 +torch/include/ATen/ops/_nested_tensor_from_mask_compositeexplicitautograd_dispatch.h,sha256=cJP7mz1eINDE5jdsVk0JssHY6F05QBlJ-6ovE5e1XzI,1240 +torch/include/ATen/ops/_nested_tensor_from_mask_cpu_dispatch.h,sha256=d7FQymcdFUYborzyVmSU34CIj-aoAR15ngVMLx2ISzg,1036 +torch/include/ATen/ops/_nested_tensor_from_mask_cuda_dispatch.h,sha256=tr2JMVq7Wuu0Fktij3qAoxu0GIXW26qkyX0J3HcJH_k,1038 +torch/include/ATen/ops/_nested_tensor_from_mask_left_aligned.h,sha256=m-8mI_HgKjNJPbU9cw4aT_kTxez4ELY5TrjlyQcfUKM,1042 +torch/include/ATen/ops/_nested_tensor_from_mask_left_aligned_cpu_dispatch.h,sha256=B8SdezwPYv43KzjKyLIJrfp7yNPhwKggJeC9PbDN9Zk,1021 +torch/include/ATen/ops/_nested_tensor_from_mask_left_aligned_cuda_dispatch.h,sha256=tEgesJtiiPNZ0jsWuGT7HpGiSxETG7ka-SvNHjGhCw4,1023 +torch/include/ATen/ops/_nested_tensor_from_mask_left_aligned_native.h,sha256=mcULN7xoyoZ2TRsZtF7D1Mwh61SfuceE-kY3KedFK04,789 +torch/include/ATen/ops/_nested_tensor_from_mask_left_aligned_ops.h,sha256=qymThN3IylZFHM4tIz0yg5vBDMIOevQ1MzTSMnYyvDg,1358 +torch/include/ATen/ops/_nested_tensor_from_mask_native.h,sha256=CsMBroXHC5BVXFNxPZG43qgR0ZY3T86-fSRErRK6JuA,939 +torch/include/ATen/ops/_nested_tensor_from_mask_ops.h,sha256=ocK5ESPILoqEWMpW79BrN7YJs_4X1-25VcHIlBE5RLY,2181 +torch/include/ATen/ops/_nested_tensor_from_tensor_list.h,sha256=iXinPceFxSheZpC773jdPI2VAeRFkfeUBgjySEGCsqU,2452 +torch/include/ATen/ops/_nested_tensor_from_tensor_list_compositeexplicitautograd_dispatch.h,sha256=S9NnU5P1toJ9CCt3w3fkPo7un8mtgwZLqj5lzAzPGtw,1793 +torch/include/ATen/ops/_nested_tensor_from_tensor_list_native.h,sha256=4beD9th12gVvJSNrIFCT4A-xRSpIHXLCHwzUy_0weAg,1200 +torch/include/ATen/ops/_nested_tensor_from_tensor_list_ops.h,sha256=p-m3FksJlwvyvMW2UpmgcFyXEvB2sM4ZBO25uKibD8I,2917 +torch/include/ATen/ops/_nested_tensor_size.h,sha256=TOcJ4zjcHmWDlL96sl-E1Ic7Hvuv1lbXr6GZUqAspO0,1237 +torch/include/ATen/ops/_nested_tensor_size_compositeexplicitautograd_dispatch.h,sha256=Ys0d2nB5aY_diK5HbZgPLSjz-jfcLxA8ejl4oeEd2_A,1147 +torch/include/ATen/ops/_nested_tensor_size_native.h,sha256=Dve5Q3AsGvOC-FEzImBACLfgCczuqQ0La-edLOMGPpk,834 +torch/include/ATen/ops/_nested_tensor_size_ops.h,sha256=ZZYqACzwLvOARnQagYpWJ31uuAwIYja7Y0dhLsH1g1U,1879 +torch/include/ATen/ops/_nested_tensor_softmax_with_shape.h,sha256=RXMJtTEl-8ngq4LRo-Gpq7oJ2VrDGtFx9sbxLoZp4Wo,1046 +torch/include/ATen/ops/_nested_tensor_softmax_with_shape_native.h,sha256=NnTWSLC0aXIy2R0rq6vYqoDHS0ncz1nJiMbcVQlELAA,885 +torch/include/ATen/ops/_nested_tensor_softmax_with_shape_ops.h,sha256=m4Js9TZNsoCMa5g6DIzUoN2iPy-s9DQVtBMxDhQ_bl4,1378 +torch/include/ATen/ops/_nested_tensor_storage_offsets.h,sha256=gYtJHftdPs9l-t0VVoV22usLk1vDWvQaJmzgDyEjXuQ,1314 +torch/include/ATen/ops/_nested_tensor_storage_offsets_compositeexplicitautograd_dispatch.h,sha256=DDs78gyojwTRBzMgZqoCUna68UU17j3dLyKbpK08doQ,1169 +torch/include/ATen/ops/_nested_tensor_storage_offsets_native.h,sha256=9wWVxOIb5VwFQbWZa1BfPMh73UME_IwCF4XTLj5bVmY,856 +torch/include/ATen/ops/_nested_tensor_storage_offsets_ops.h,sha256=4JYMSYdsHWrJq3WobN89kGdr94iymE3tHhWs6P_GBdw,1945 +torch/include/ATen/ops/_nested_tensor_strides.h,sha256=T7UZLFrXgE8wWkfmK6snp7cD7iANCSbUasVvwrTJHAw,1258 +torch/include/ATen/ops/_nested_tensor_strides_compositeexplicitautograd_dispatch.h,sha256=HYJ_89DTAS-K2hnivmLYl7bf1pzqqSfwHZVH7bXsYIo,1153 +torch/include/ATen/ops/_nested_tensor_strides_native.h,sha256=M1JxPb1uml6NQmMfBmuKWmezNo8Q-xCf_7sHegs2NzU,840 +torch/include/ATen/ops/_nested_tensor_strides_ops.h,sha256=yyO8U40gG65-21vidpDZ0a6idrSwP92gtQmhoQbpFcQ,1897 +torch/include/ATen/ops/_nested_view_from_buffer.h,sha256=2F05MotFQrLut8-9WXxn-f6YD2ae6x1GAIq3D-C6f1s,1161 +torch/include/ATen/ops/_nested_view_from_buffer_copy.h,sha256=-cezMkB4XzvtENbv8YAfAJhkpEqE4V5lZlr-JKzwgmM,2087 +torch/include/ATen/ops/_nested_view_from_buffer_copy_compositeexplicitautograd_dispatch.h,sha256=dUk1mGiyIxE3pSNl1uDtfVvRzmhTHqlHGs7jSF_ivfE,1357 +torch/include/ATen/ops/_nested_view_from_buffer_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=NSg5ojU4IYoRSRHjGwrnteUVoRIsxKuWY8upXAoGkuY,1162 +torch/include/ATen/ops/_nested_view_from_buffer_copy_native.h,sha256=4h7sHjtAMEWcZJW5tOntNhhu5p4rPXJivsqmI1-yr40,1044 +torch/include/ATen/ops/_nested_view_from_buffer_copy_ops.h,sha256=mp3RQDugyeEmnfoDvDQeftoQr7H6iPaFgj1aSnqLXkY,2557 +torch/include/ATen/ops/_nested_view_from_buffer_cpu_dispatch.h,sha256=noEMqKx5oBcIg4JfjD4A-R6LgS4nxaj9yGINfrV66Ds,1087 +torch/include/ATen/ops/_nested_view_from_buffer_cuda_dispatch.h,sha256=0SOfAqLY0jjyhrxvHlRzTmJhsB9201XQqeXQw5YDH-U,1089 +torch/include/ATen/ops/_nested_view_from_buffer_native.h,sha256=yK0-7yVCHR0z2C6BwDTfC6LKsX8YicPMcrHa-iGtBsE,843 +torch/include/ATen/ops/_nested_view_from_buffer_ops.h,sha256=cC0HgQLdMlsJ2Erv1xEJc-zyA8hOmzmGwHH9tcpKLjQ,1580 +torch/include/ATen/ops/_nested_view_from_jagged.h,sha256=HBSJCZQSB2m22oJFdIcd0WHOWWX7vf9A1P7LdFp89ho,1376 +torch/include/ATen/ops/_nested_view_from_jagged_copy.h,sha256=EIfzBk1_wizg39QzGP_gTfenokRA7FSO3tnA8nLblvE,2721 +torch/include/ATen/ops/_nested_view_from_jagged_copy_compositeexplicitautograd_dispatch.h,sha256=oneV_O0u2la4ZjSa7f8oXiLAtcnkZSMtnbc1gmkuCD8,1608 +torch/include/ATen/ops/_nested_view_from_jagged_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=n1oP1N01iBEYfOVatzEKEI19OXpr7KoQmnD8I45CDB4,1293 +torch/include/ATen/ops/_nested_view_from_jagged_copy_native.h,sha256=NA8Z06s0On7_YpOLNaY3pYc-m3oJPR53L-5Exsp6dgk,1295 +torch/include/ATen/ops/_nested_view_from_jagged_copy_ops.h,sha256=ds2eLXsd6Lm2bIBUGKQtS4yK_GmKCud3gLEji3GmI7c,3359 +torch/include/ATen/ops/_nested_view_from_jagged_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_nested_view_from_jagged_ops.h,sha256=2J2NFfl9yby6482uSund9MMbmUjwvhRa4Q4dHYTUpf0,1981 +torch/include/ATen/ops/_new_zeros_with_same_feature_meta.h,sha256=dIlT3WxZdEP3gsDxcoY2TIHhPClg_TsnFXqAxQS6vnM,1930 +torch/include/ATen/ops/_new_zeros_with_same_feature_meta_compositeexplicitautograd_dispatch.h,sha256=qNdZue4vTN725RTTbSBVeCYes1i_8XRBhgTZa295VAg,1425 +torch/include/ATen/ops/_new_zeros_with_same_feature_meta_native.h,sha256=21dyGksOg9JVquqobYotzQQKTPQ8hDKdxmMJiB_dDx8,974 +torch/include/ATen/ops/_new_zeros_with_same_feature_meta_ops.h,sha256=FiTwAT05czcmKiUQqT_J33iwX39vzd0ex15KAVmIhOM,2326 +torch/include/ATen/ops/_nnpack_available.h,sha256=oV6goa4QfSK2ZnP0YkUj-htxuPMbpkTqCZ6ev0aal2U,889 +torch/include/ATen/ops/_nnpack_available_compositeimplicitautograd_dispatch.h,sha256=nwl_jrYLOQJNF_Nf4Bsde1Y9YjwmLe_z32DdmxZa4jc,1000 +torch/include/ATen/ops/_nnpack_available_native.h,sha256=OJdl_W9li1i9IT0aq4sSRqlRqTVeTZqp3-5dwOd-zsw,712 +torch/include/ATen/ops/_nnpack_available_ops.h,sha256=cPDJhQBUZgS5-bmdylCEzFvkuLOqJ4fwXHqpiqUhhho,1147 +torch/include/ATen/ops/_nnpack_spatial_convolution.h,sha256=2xPWYvPZyNLoJUjU4ILOwD17DDSyJAo8RAmpl0OcIS0,6444 +torch/include/ATen/ops/_nnpack_spatial_convolution_compositeexplicitautograd_dispatch.h,sha256=qhSqoJQDRGtq7APTeOT8T0-GVJ3oinWkJaccLj6S3v4,2305 +torch/include/ATen/ops/_nnpack_spatial_convolution_native.h,sha256=EJBcEcar35qUtvANNl2NHnNJNh9fHkAo_6ea94bCbHY,1105 +torch/include/ATen/ops/_nnpack_spatial_convolution_ops.h,sha256=mBA7eDU5ysJdVpDZVUKoH9-84A_ioFb-dfrRo8dywrM,2771 +torch/include/ATen/ops/_nnz.h,sha256=1lhtBgMlJLNkXzhrmiBSMKlTdO_ElbjumibHEQWwPFM,755 +torch/include/ATen/ops/_nnz_native.h,sha256=MghWvuHpTXpKNlv5O9WMKJPd8eF2EKjHDRR6jGsz2aY,792 +torch/include/ATen/ops/_nnz_ops.h,sha256=tO21m1z6ovn14GY2KpuCjYsidQY9oROtax81ef13zys,1193 +torch/include/ATen/ops/_pack_padded_sequence.h,sha256=VAzBB43qgvNRqR_MWqqA_f1bhhAnobizbcVWmTwoZWc,1953 +torch/include/ATen/ops/_pack_padded_sequence_backward.h,sha256=aHHcQJlth0ZXzfhTvezDUoOJ1CXMnlDjtxQVQg9rZ5A,2331 +torch/include/ATen/ops/_pack_padded_sequence_backward_compositeimplicitautograd_dispatch.h,sha256=D3Edg-_Dm_2YabwbbCPUEOruFbb5DYMAHsHwVIfpEaQ,1287 +torch/include/ATen/ops/_pack_padded_sequence_backward_native.h,sha256=OyFkYMzbnQcIzlLBfSUASwTjAEKjsRWByc3ZY7BW-MU,843 +torch/include/ATen/ops/_pack_padded_sequence_backward_ops.h,sha256=XneV5fcPXUjXcs3x_Ls03vu3Fl3EJX_sxgmbBSX8VDk,1553 +torch/include/ATen/ops/_pack_padded_sequence_compositeexplicitautograd_dispatch.h,sha256=hOG9jRikNJXBqYLi71HXSvdohCRAEIlT54JGVbEIgvg,1480 +torch/include/ATen/ops/_pack_padded_sequence_native.h,sha256=PI7qNyx_QLPnYEnHrgy7WRp1tjFxrmQiaPwI12S2L8k,1004 +torch/include/ATen/ops/_pack_padded_sequence_ops.h,sha256=JabqY2LKAhiaeudwRjm0VSXAtS0N1jHBr9H5dISK_zU,2453 +torch/include/ATen/ops/_pad_circular.h,sha256=o5ucFDdsd1JISu64r-rPDXKIjxgQlbY-Ciuqs7vuUYI,1694 +torch/include/ATen/ops/_pad_circular_compositeimplicitautograd_dispatch.h,sha256=AgoYwD5vcDecPJI4HGsTU1KRZet5m-zlyIDwA5IKdxk,1139 +torch/include/ATen/ops/_pad_circular_native.h,sha256=vnkhNRGw30N893xxiI_dAqNu1Inov2PE9MqGL2DQU0o,769 +torch/include/ATen/ops/_pad_circular_ops.h,sha256=cN5RmjIHW6BPUNvE3DmXRIzUb_u7wn4mt_EpS62cb8A,1317 +torch/include/ATen/ops/_pad_enum.h,sha256=cheoBM9BGnVmhgPYVw0nqgtpmF795qOqqthOmAWheI4,2000 +torch/include/ATen/ops/_pad_enum_compositeimplicitautograd_dispatch.h,sha256=_85vUmAnX-QvMUCH9z0K49jPHJm7bJfg8uuH9ydPU34,1251 +torch/include/ATen/ops/_pad_enum_native.h,sha256=NnBsHt8oH95te3tKnLz1hbJi1EJY4IqEKI42Hzl5_6Q,825 +torch/include/ATen/ops/_pad_enum_ops.h,sha256=cImMByfi56IM9G0cpBewab1Z-9kh5M29PJ-7fqyMUyM,1458 +torch/include/ATen/ops/_pad_packed_sequence.h,sha256=_wyjVZx-ohpIF2dD9iTbPnrszIFcIhAFhzTgmnOBmLk,1221 +torch/include/ATen/ops/_pad_packed_sequence_compositeimplicitautograd_dispatch.h,sha256=knDvQ9I0JD2ZNbwn0nmSYqzGd4nPHxUviC3OLaCxCq4,1163 +torch/include/ATen/ops/_pad_packed_sequence_native.h,sha256=cKQKcsCkqSX3BEQ96PdJUus9m3WADoiw6bnKCacpiEs,875 +torch/include/ATen/ops/_pad_packed_sequence_ops.h,sha256=13sTUuiUZaS4U1LFjCx9Ymskk00jqCe3CvNEs31KHiA,1683 +torch/include/ATen/ops/_padded_dense_to_jagged_forward.h,sha256=lYAKhkELvNpDruD83Ylvzf57qJPjFCyzuVHDkSpaP00,2312 +torch/include/ATen/ops/_padded_dense_to_jagged_forward_cpu_dispatch.h,sha256=rF9of2-Su4dYaVHQg80s5oEq5Y3psbsVepw623FUOuc,1237 +torch/include/ATen/ops/_padded_dense_to_jagged_forward_cuda_dispatch.h,sha256=yT4E2vcxrIlp4tBy1ixfMDksUgGz9ERgVSCr7EqfBSE,1239 +torch/include/ATen/ops/_padded_dense_to_jagged_forward_native.h,sha256=VNMKk3iUGeb9gzTbloNRnDs4ExnvkjFUBaBBXpgXYto,997 +torch/include/ATen/ops/_padded_dense_to_jagged_forward_ops.h,sha256=PfoDrDSwS2kDzs-xqQVtbDLXM-IO42fXupICwitz_Vw,1499 +torch/include/ATen/ops/_pdist_backward.h,sha256=k6WF9BWQWVzCgrc8Dw-TI6U3YYONJoRjgnemtM8T0gY,1710 +torch/include/ATen/ops/_pdist_backward_compositeexplicitautograd_dispatch.h,sha256=eOJkEYOZ6zCO1o0174UKCoWq7lhmgtlVODrCkBUHeO4,1261 +torch/include/ATen/ops/_pdist_backward_cpu_dispatch.h,sha256=9N1r5E0pWVef9Z9hhQA3P_qXEcqlFfe1HbunCR6gahE,1044 +torch/include/ATen/ops/_pdist_backward_cuda_dispatch.h,sha256=9y5Jx36VQPDqbnFCX5CAdr0vv_R7wTaWREZjTep2qNw,1046 +torch/include/ATen/ops/_pdist_backward_native.h,sha256=BanagFCshbXiAIhQG29lQT2abDt0p2qwsHrwWn4EgBU,948 +torch/include/ATen/ops/_pdist_backward_ops.h,sha256=oskEoGnOyyQBYbGdKsqR_r2CnRcUF2h2z5DU1UaVMw0,2267 +torch/include/ATen/ops/_pdist_forward.h,sha256=SAhfuJiC_sSJ4CBTJjbYNuGiKtAd3FMUqCG4JFMyq3w,1437 +torch/include/ATen/ops/_pdist_forward_compositeexplicitautograd_dispatch.h,sha256=oZkc-g21inrOivY3kypbp_-5LOQ08eX_iz4YLl-8gD0,1159 +torch/include/ATen/ops/_pdist_forward_cpu_dispatch.h,sha256=w_y6YucqwE2ESCC3CVXXkpAUbNCFZT1i-RF4bUMRZCs,994 +torch/include/ATen/ops/_pdist_forward_cuda_dispatch.h,sha256=BbKWrdPU5JUJWpMKX0Hx-O9MgzFoIuwsxdoh64XXecY,996 +torch/include/ATen/ops/_pdist_forward_native.h,sha256=DJV2ETpnGpK2KTwEBSkI5QKrelauXuUvVJ73Zm0jdhQ,846 +torch/include/ATen/ops/_pdist_forward_ops.h,sha256=gZA3SmdXe3Z1WQ0ezUuK2ns26vBcCP_JhZuJFYysidA,1927 +torch/include/ATen/ops/_philox_key_fold_in.h,sha256=dAVEw1Uqa-sNY9ci14yh8ik6w5czXoYMXRJQUpVokh0,970 +torch/include/ATen/ops/_philox_key_fold_in_cuda_dispatch.h,sha256=rQSoZ31FTWJ-NPlMGJf9WpwrD-GOmAzl2kkONNtMkRc,1002 +torch/include/ATen/ops/_philox_key_fold_in_native.h,sha256=FvMEiLmE4GXp9kToDHole35fxjo1kHdY_fxfO_xqlFA,761 +torch/include/ATen/ops/_philox_key_fold_in_ops.h,sha256=Z4wOGMPfH0F4E8H5ym9ZXGKJLug_fweZRVVdAnWW7zQ,1294 +torch/include/ATen/ops/_philox_key_split.h,sha256=NyOXbkQ_Oz5qC_OZUVBzARiw80aFd2NPYtiG4BKloP4,980 +torch/include/ATen/ops/_philox_key_split_cuda_dispatch.h,sha256=hJlklv_LxpbMyPgJGgCMI0SyoBPUsLTbisz7tvP_bW4,1006 +torch/include/ATen/ops/_philox_key_split_native.h,sha256=O_AQR6yvr-wrbLUJ5yBzfqkhUivhnnMHoaH6V4tGMp8,765 +torch/include/ATen/ops/_philox_key_split_ops.h,sha256=szHQlgJ6LZMlzi-UvIybgPhVbSnhYqKmxGOVwyCerxU,1306 +torch/include/ATen/ops/_philox_normal.h,sha256=-WtAEdtt1l5hOgT1wwXA24rZ6wchCmG-wcS7mjKO-aE,1955 +torch/include/ATen/ops/_philox_normal_compositeexplicitautograd_dispatch.h,sha256=u8EsO8EftjTQyh8ClN1jGKaoRfKcyaZn4L404wb0nT8,1354 +torch/include/ATen/ops/_philox_normal_cuda_dispatch.h,sha256=4nIRL22APCCfXqyBvY_eOLISnPpcVyHCVUqy1cILgxU,1034 +torch/include/ATen/ops/_philox_normal_meta_dispatch.h,sha256=hP2HbUnUblle_WQXWPVb6VEy9K_XcVMd018t_QqoDT0,1034 +torch/include/ATen/ops/_philox_normal_native.h,sha256=1yxcfFnHTeVa9Y676u4fDVxcVEfcTfxllQs0SNk6980,1043 +torch/include/ATen/ops/_philox_normal_ops.h,sha256=C0CU-zUx8hOPSyjWCLV1flIhxssTY5faeUia7lJB7SY,2891 +torch/include/ATen/ops/_philox_uniform.h,sha256=PjZdF0-IymC79edE4V-meHEs1yIpl5yw3ZAAuF02YVo,1968 +torch/include/ATen/ops/_philox_uniform_compositeexplicitautograd_dispatch.h,sha256=dKc7YLdYx-GAUu92hPZQCDjujFglN9mq31TcRWPREn8,1357 +torch/include/ATen/ops/_philox_uniform_cuda_dispatch.h,sha256=V0BFcdl1BKxSNG80SjWf24xK0x5UBvvWHyhPXjaN-sk,1035 +torch/include/ATen/ops/_philox_uniform_meta_dispatch.h,sha256=FOiOLxx2ZY3LxUITZITYESVHtXzS_cY6vvP9NNHPB98,1035 +torch/include/ATen/ops/_philox_uniform_native.h,sha256=fX8SFaPSaIlbImunXX6sKigyEEJkeJGnwxkPn1RxTE4,1046 +torch/include/ATen/ops/_philox_uniform_ops.h,sha256=DhO5HlCHjPagpCc1KWhZsdGoH4uqILICeMTuHCCAxPw,2900 +torch/include/ATen/ops/_pin_memory.h,sha256=f48BJC7C9RGchaYMBreKLAAnOWt7lUZ8tO528RCEJGY,1556 +torch/include/ATen/ops/_pin_memory_compositeexplicitautograd_dispatch.h,sha256=2pR4TFh4mLg8JIbNPjmZlnouvrlVFNx79nqrOtGINnI,1328 +torch/include/ATen/ops/_pin_memory_native.h,sha256=GVWmoFeYHgQsUV2vKkkBY_0PB_-J6o8eOCfizizp3VM,1271 +torch/include/ATen/ops/_pin_memory_ops.h,sha256=DRrtsUcd52KFy3fBCF0bRCLTbRw2o7xClb5sl6z6JsY,2075 +torch/include/ATen/ops/_prelu_kernel.h,sha256=t9cVX2B3sJE5L8U6qDuF226d9qHbcgw4Cb6c1nhwyLk,969 +torch/include/ATen/ops/_prelu_kernel_backward.h,sha256=houKbWxWPT83Ni-v4qOVAdvvXhO48VFyPJZRT0lzHxQ,1105 +torch/include/ATen/ops/_prelu_kernel_backward_cpu_dispatch.h,sha256=8Pov0cexygKx5Hjc9FwHN-I_qyOl1sZA8ZkEOuCVtHI,1074 +torch/include/ATen/ops/_prelu_kernel_backward_cuda_dispatch.h,sha256=3OnRK62T-H-oEtFdfGcAAJXZ1wKn1qetfaSzz-5bns0,1076 +torch/include/ATen/ops/_prelu_kernel_backward_native.h,sha256=GESj8stQzse6y6RE_CecxT1-CUfie7y1QStXmDp5G3s,983 +torch/include/ATen/ops/_prelu_kernel_backward_ops.h,sha256=qN2pj58EOAhRhqdDMpwzQ5pHJVSN7acrlzt_Dl2LfjQ,1537 +torch/include/ATen/ops/_prelu_kernel_cpu_dispatch.h,sha256=804eCC9spcycAV9wPsigRLg2YuANH_hxF4i4PaTMsgg,1008 +torch/include/ATen/ops/_prelu_kernel_cuda_dispatch.h,sha256=SmM1Ll432iY1i285LPsdEfb0ciCa2bj6Haug1fSzy0Y,1010 +torch/include/ATen/ops/_prelu_kernel_native.h,sha256=6iZiOIuIG_fdMZjmYbAw2SCGEjABiYosFpKO19kpFj8,953 +torch/include/ATen/ops/_prelu_kernel_ops.h,sha256=3EI2ZjI9LjgB4eSXrv-og899JycdyKr04_mlVnlPRuc,1321 +torch/include/ATen/ops/_print.h,sha256=5OFA77nJYf8LkC-rSZ2KGkTQZiXBpmbUdR-3MJOto_o,867 +torch/include/ATen/ops/_print_compositeexplicitautograd_dispatch.h,sha256=zl_WKOQf4zPyxUUhBPK68ujxSzFQIQzG8_CC8Hs42Yk,1007 +torch/include/ATen/ops/_print_native.h,sha256=eCprfTdpT5-nMOcNXP-_tWir1um52cHjaGAawJSO3rs,719 +torch/include/ATen/ops/_print_ops.h,sha256=zkvwZygubiiO8XzYUrStY762cfrUzRsOITK0HYviVZ8,1171 +torch/include/ATen/ops/_propagate_xla_data.h,sha256=6Jo_dbF13X7oQr7NmtU7rjd3IYrgAsOdIJmUWX5SxVs,986 +torch/include/ATen/ops/_propagate_xla_data_compositeimplicitautograd_dispatch.h,sha256=ps64PnUTXuRhjL1hILg9AehyosHnrB_EAhN3-3l4OYg,1053 +torch/include/ATen/ops/_propagate_xla_data_native.h,sha256=gZ88nHqN9B6w97UnbbrRRWhy6II8UDb4s_7hPSVkgm8,765 +torch/include/ATen/ops/_propagate_xla_data_ops.h,sha256=3xnHo-pSDPbfE1QYDMoeEMilDEttUmapQl6B7VluJwk,1320 +torch/include/ATen/ops/_remove_batch_dim.h,sha256=eOpL9u5q-Qz_kB6lgOYu1BOK3eWqYSgyX5a2TdQyVZ0,1960 +torch/include/ATen/ops/_remove_batch_dim_compositeimplicitautograd_dispatch.h,sha256=Xwe8xXM4f8DLcRc9q6sthSqyWmMQGDDgDtwYBVUKczY,1209 +torch/include/ATen/ops/_remove_batch_dim_native.h,sha256=ZNniP5PB9mIjrWMf4h3ZEeP--tmkM8_8tnsPI4Vi9NE,793 +torch/include/ATen/ops/_remove_batch_dim_ops.h,sha256=6znZrY45-YMP4FoV5AdOvTD0qEyk1ByZXAUCMzxNj_U,1430 +torch/include/ATen/ops/_reshape_alias.h,sha256=ReViPo3V-LiTtuD5qSw0i-6xitFPNs64mawUXIXKj88,1949 +torch/include/ATen/ops/_reshape_alias_copy.h,sha256=7hrQF2mFmBilYx_spJuzzLi6quOqmDoArsCSgeSY2do,4814 +torch/include/ATen/ops/_reshape_alias_copy_compositeexplicitautograd_dispatch.h,sha256=0-YZVkeLhJYOD5Ba-ZNI9ppRMfp-rz2VDJDB4WIOpAw,1544 +torch/include/ATen/ops/_reshape_alias_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=N8GKbgxyKQt2sYVRnmQQ3TYOjUNDVyNMLvNZzlLmJYY,1231 +torch/include/ATen/ops/_reshape_alias_copy_native.h,sha256=tph-eWL2VYNjTkbDx0lVyK5ZQ1llXI81LbR5vxEzanI,956 +torch/include/ATen/ops/_reshape_alias_copy_ops.h,sha256=JWjYFRQPyuI0TRDTVVUDJrYGKcbxzvtLEXwtNFeGPuI,2243 +torch/include/ATen/ops/_reshape_alias_cpu_dispatch.h,sha256=V9z_QhGXRKUN7O8IL5w4hJWLiXmcZRh8d-MAqZlN0MA,1151 +torch/include/ATen/ops/_reshape_alias_cuda_dispatch.h,sha256=-t-_vz8cqBndjmsF7hFa1bC3V-pLEeUmyL5DcUI9IHg,1153 +torch/include/ATen/ops/_reshape_alias_meta_dispatch.h,sha256=XaGj8uswEiurDazL7HzNwj3QSnBpF43xCU9y9ipWzaY,1153 +torch/include/ATen/ops/_reshape_alias_native.h,sha256=GR4Y6u8DRhmQl6EzXx8xgaxSIrQNWyrVc8Aj8izm19E,784 +torch/include/ATen/ops/_reshape_alias_ops.h,sha256=jZGqQc2izhQrLdWoE__8IBkEmFXgHr5x0P7M5o3SbTg,1423 +torch/include/ATen/ops/_reshape_copy.h,sha256=CetYqNpV4vOlT0tGu-beT1uO0mvK-IqpV9vNA-pxZIw,1704 +torch/include/ATen/ops/_reshape_copy_compositeexplicitautograd_dispatch.h,sha256=YovmUvVX4I97WjWOe9dWtHJRH3GAWkxYjEz6blhZ53w,1141 +torch/include/ATen/ops/_reshape_copy_native.h,sha256=vGML376As41Qx1hxEY6YyqqjVxB94S7SErJmh8gd89E,770 +torch/include/ATen/ops/_reshape_copy_ops.h,sha256=NAV-P_6a_8neA5jIY8btg4xNywiDtMigJMDd3kdJWoM,1320 +torch/include/ATen/ops/_reshape_from_tensor.h,sha256=MAiAplBHMN_UkymV0Qjp-MN-HagklaV9h1xyf_HdKqA,994 +torch/include/ATen/ops/_reshape_from_tensor_compositeimplicitautograd_dispatch.h,sha256=JaOPbYWXJFOPlQc1zUwbZcFkz63X2QFd3kqiva_CfGE,1058 +torch/include/ATen/ops/_reshape_from_tensor_native.h,sha256=v4_cg7r-qz16qgmhShPYtQNEP2QZfVWR0W1CSpZPmzo,770 +torch/include/ATen/ops/_reshape_from_tensor_ops.h,sha256=qISJkKrMyBVq9dKuKV2oloyN1k7ozg2GcARSZwdE9Vs,1339 +torch/include/ATen/ops/_resize_output.h,sha256=0Y0DX1HE1azWUJQaPFJRV4tO4iCxPd00KTCd3w47BYs,5650 +torch/include/ATen/ops/_resize_output_compositeexplicitautograd_dispatch.h,sha256=hznXLriUm5o0ojq_EvKhDkxqQeJIYFabGrlhTHOXGOE,1761 +torch/include/ATen/ops/_resize_output_meta_dispatch.h,sha256=BgJf9sp67uXVRJ0ZQ7Rum0JPveFLiOUEPyT434D7gW0,1157 +torch/include/ATen/ops/_resize_output_native.h,sha256=Dkl0hwr8J_vUe6jDsly4GPIj8HNoIUwgWjjAjGOXdzs,1052 +torch/include/ATen/ops/_resize_output_ops.h,sha256=lFyaKqF0X2wzcxtrcvnX-skHuE2MZA-LzNXQxapE4p4,2903 +torch/include/ATen/ops/_rowwise_prune.h,sha256=gOJxK_UJhxkXINCe9w3xTXvV7A32g3tNgxe1RlF10SY,1112 +torch/include/ATen/ops/_rowwise_prune_compositeimplicitautograd_dispatch.h,sha256=Kbn76FF543PcePQ23cmOZhI2BwonNQjRxD3gH2eJv6I,1119 +torch/include/ATen/ops/_rowwise_prune_native.h,sha256=bgvfHcLPZv9OZ9MFYb9YjBHuASKfMxExJp-n4eT5zRk,831 +torch/include/ATen/ops/_rowwise_prune_ops.h,sha256=SLCR51rc_G-VTbvmP9OH6gP6a7f8oliwhUVIwG1KmKY,1544 +torch/include/ATen/ops/_safe_softmax.h,sha256=fXBpzi7mkNyn0oeq9-Q65ltkcWZ_njA-WZRYklnWfhQ,1031 +torch/include/ATen/ops/_safe_softmax_compositeexplicitautograd_dispatch.h,sha256=Ig2Qvf97RxcSTFzsDJwmYMDpqCPWatXHIbgDN_6L52g,1092 +torch/include/ATen/ops/_safe_softmax_native.h,sha256=IUiXfibyQdI4ukkw2yWcyKLDZrctDq-1RBQfEHJF-ck,804 +torch/include/ATen/ops/_safe_softmax_ops.h,sha256=jKYKZgg7IBEQ4J0BO6JaUeyHwicu5aUlKI79sB9d9_Y,1411 +torch/include/ATen/ops/_sample_dirichlet.h,sha256=uCfuSXMIHuxUw6cyz9bqnf4KnWxE9WfOGks7fJ7Y8Uw,1661 +torch/include/ATen/ops/_sample_dirichlet_compositeexplicitautograd_dispatch.h,sha256=3jq_U_Qz7HZtPmhQDbWf6Zf0ykxJv2zWXWM8RtsPFEQ,1242 +torch/include/ATen/ops/_sample_dirichlet_cpu_dispatch.h,sha256=ASQ3z19-oFmygGnHZ5vzHGHD7rLtIjOfyOWVGxUKMFY,1042 +torch/include/ATen/ops/_sample_dirichlet_cuda_dispatch.h,sha256=olV0fpUQUcXm2xT_R7LVj87KbP-oeqESo9V7HaM6QyY,1044 +torch/include/ATen/ops/_sample_dirichlet_native.h,sha256=cFMCLxvkjS_PSAEKWiRGad_6rd8xEfVd143wXI5Id8w,1050 +torch/include/ATen/ops/_sample_dirichlet_ops.h,sha256=X6HJ3f1MPN9r22ZQ3zp3nlHmaZ7ELG4lDNhR4Uk4_8k,2153 +torch/include/ATen/ops/_saturate_weight_to_fp16.h,sha256=RNha5FkTQxxR0UXIlNCuSS4z2XANaMJG11LFvBFrecM,969 +torch/include/ATen/ops/_saturate_weight_to_fp16_compositeimplicitautograd_dispatch.h,sha256=yS0SNFXQz2jxuwe7Zt87JFq9UqXKA16TPX34rhwXJxE,1038 +torch/include/ATen/ops/_saturate_weight_to_fp16_native.h,sha256=AdxN_fk1juTYK7zMrCiNNvL8lLnXycQSf7zN60ZWap0,750 +torch/include/ATen/ops/_saturate_weight_to_fp16_ops.h,sha256=5T9Rc8BlPbtjuUSzLc6ysXENetHb8C1wB7xrc-WZ1L4,1271 +torch/include/ATen/ops/_scaled_dot_product_attention_math.h,sha256=_TYa8ktH1PHzHcwtRwN6qBpxbZmjQ-ajgPMjKqBuijY,1550 +torch/include/ATen/ops/_scaled_dot_product_attention_math_compositeimplicitautograd_dispatch.h,sha256=_eEt8qvdNJ35HTyWqYqYbFPHr9plP-zWGlaq5zooseQ,1338 +torch/include/ATen/ops/_scaled_dot_product_attention_math_for_mps.h,sha256=g7hR0QgcE_sOreA6O-kj-3j0_5jNvp8IAXfDviJB5Aw,1582 +torch/include/ATen/ops/_scaled_dot_product_attention_math_for_mps_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_scaled_dot_product_attention_math_for_mps_ops.h,sha256=JeGfvW6cgTJhIweu16yLwFYmA-HH5lirnEggUYpkI9o,2189 +torch/include/ATen/ops/_scaled_dot_product_attention_math_native.h,sha256=lrUtfvmGJ2lI1ZTPQ7PzjPtuhqdQk0Ep24jcgXN4OCg,1050 +torch/include/ATen/ops/_scaled_dot_product_attention_math_ops.h,sha256=Hxchl300QXavAAQyZLgpZNZWyF63rOTqILjce5u3Ypc,2165 +torch/include/ATen/ops/_scaled_dot_product_cudnn_attention.h,sha256=Y_5_NeGXfvohk235UKUVZgRFs7HZnK_3b0T9tHfLEKg,1769 +torch/include/ATen/ops/_scaled_dot_product_cudnn_attention_backward.h,sha256=guWO3YXvUDY_PY3ZgF4zA6SaBelbZjioKtmCwjUdAes,4717 +torch/include/ATen/ops/_scaled_dot_product_cudnn_attention_backward_cuda_dispatch.h,sha256=oZA8H0WIJDscfG5VIjMP9dcWSeaYYeSIFWrOZ4tewbk,1993 +torch/include/ATen/ops/_scaled_dot_product_cudnn_attention_backward_native.h,sha256=Tk4j3AvE4DqNIIsGZQV16lk9adYcXpkWVQSBSZzMOdA,1755 +torch/include/ATen/ops/_scaled_dot_product_cudnn_attention_backward_ops.h,sha256=DzEJ_0MEY6U3ltKq8oKS4od907SuAFpJrlmkh0WI5_M,2755 +torch/include/ATen/ops/_scaled_dot_product_cudnn_attention_cuda_dispatch.h,sha256=ja7ieHhj_P7B5Tyf4RlWQyiV3aMKL9WZHYcWIZcD-OY,1352 +torch/include/ATen/ops/_scaled_dot_product_cudnn_attention_native.h,sha256=JTe1uzm1cKP2SMKTEJvSD1VXRR7Vxbpkb1DjALSo9j0,1559 +torch/include/ATen/ops/_scaled_dot_product_cudnn_attention_ops.h,sha256=ztRblSKZWgG-JzkKd59jzGOzt_wF_6uaNPKJT9Az4H8,2485 +torch/include/ATen/ops/_scaled_dot_product_efficient_attention.h,sha256=oIzzTb7AxKgxUT6h8fhUsrKV8tTHYWIQ9VXy9bY53xA,1562 +torch/include/ATen/ops/_scaled_dot_product_efficient_attention_backward.h,sha256=8iG2kcH3RzKLKUvgyMJhTU0khR8p-kq4Kd-F-MrABys,1833 +torch/include/ATen/ops/_scaled_dot_product_efficient_attention_backward_cuda_dispatch.h,sha256=xt_dgIQETh8Xd00wCk7flwo0HiZ8-ycdmXKNdSorp3Q,1420 +torch/include/ATen/ops/_scaled_dot_product_efficient_attention_backward_native.h,sha256=UiE9SlKp51xMit_8HeDuMIzo5AiUVLGVJyS0DOGKoXU,1179 +torch/include/ATen/ops/_scaled_dot_product_efficient_attention_backward_ops.h,sha256=jdMQ0XZWQ490YpHrer26ZyP6XEFvtAiE0ED_cuFQ5WQ,2614 +torch/include/ATen/ops/_scaled_dot_product_efficient_attention_cuda_dispatch.h,sha256=XR3vVssxe6h1rAzus4mu8tg6pmU-dtQmtD0Vc8m1tK8,1269 +torch/include/ATen/ops/_scaled_dot_product_efficient_attention_native.h,sha256=zjbVX1j4eayG1cpQ7U96v4TVa45laC1126gvCyEiDAA,1393 +torch/include/ATen/ops/_scaled_dot_product_efficient_attention_ops.h,sha256=Qp5BxaHjlxCeU0pIMoqyupIiyMUfTC37zp1CogwKEnY,2155 +torch/include/ATen/ops/_scaled_dot_product_flash_attention.h,sha256=cngjBw61Uvy4WM0h8MZV-9e7Pwz7Z4U2juloMZh9wk8,2692 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_backward.h,sha256=bfTRMLwYVILUmWXMqgGry3eoSmHqU0C5HOwsysS8IPA,4579 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_backward_cuda_dispatch.h,sha256=xpG07K1M8QZugD2zuSibkHp5Rvtg2Cg4VeP1e3xiDaM,1933 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_backward_native.h,sha256=VMOnC4uilfT2ycz5d26S5qAm1e8nbpY6bi1AXNdepCw,1684 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_backward_ops.h,sha256=XUVZ39QSXd3DGzoEPKCu9YYk8-gW7YwNSfV_fT7CHgc,2688 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_cuda_dispatch.h,sha256=xS9MO3W0qgRmGBZnRZwFaIUbLwYHiC-zBgdbnzTusus,1779 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_for_cpu.h,sha256=XvYvOuB4Lc9LBRUKStYjSZNnjS9BhBYbfRx9Zps05iY,1451 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_for_cpu_backward.h,sha256=T44fSiswgcEhzGmnnblItTTHKUx1knaOGm1HB09xIG0,1656 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_for_cpu_backward_cpu_dispatch.h,sha256=xRJSuu7OdMOh3p61Iy5K8yMHZdYfXZKPX4KDRedGSxs,1320 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_for_cpu_backward_native.h,sha256=Ok5e_LibHBAyw4MWT0Fw4jbfVz2LUQN8_4RP63Mvxz0,1072 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_for_cpu_backward_ops.h,sha256=BsXy05iGQr3gywL-66xw4zu6ECoziyfFLxadvy1UDVA,2327 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_for_cpu_cpu_dispatch.h,sha256=ubvOoqbhAHoxXMRC-8zy3W2DkBLlPdOgBpDbHPR4tK8,1227 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_for_cpu_native.h,sha256=t4iu9-OSrOW62l-RGutiOUGASk6zAV_ir75QSK5UnJs,979 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_for_cpu_ops.h,sha256=mtmXBmdPYb0uqfXjdvh1Zd9_UZSKEsMC36W9CYonL8g,1982 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_native.h,sha256=MEFq7loFsToVtuVMtof3GauDBEJ-0NOSByUL_7aSqJM,1929 +torch/include/ATen/ops/_scaled_dot_product_flash_attention_ops.h,sha256=FLqMF0IZ07A8kT6dKfaz3E8vhFtqh_IibFi5IKBZ5QA,4258 +torch/include/ATen/ops/_scaled_dot_product_fused_attention_overrideable.h,sha256=OX16nWo0XgylrrHdN43oHx9A5_i8GU_6zkxyfbrcRkw,1759 +torch/include/ATen/ops/_scaled_dot_product_fused_attention_overrideable_backward.h,sha256=dTHaB2E3ejSBjRUmMOmrrz2gMbTMrR2zSGeC2QNpF08,5282 +torch/include/ATen/ops/_scaled_dot_product_fused_attention_overrideable_backward_compositeexplicitautograd_dispatch.h,sha256=x4_rX0eySBtka0fHxPb5wpxCnvLRS8nBQ1zSdvU3rm8,2159 +torch/include/ATen/ops/_scaled_dot_product_fused_attention_overrideable_backward_native.h,sha256=R8HD1CtZxVrA876Ss3FGgFNyNfx7Euo8bPg-09tY-u0,1266 +torch/include/ATen/ops/_scaled_dot_product_fused_attention_overrideable_backward_ops.h,sha256=n7dvP0-SJPSva12sveflpCDsZEEZDjYfS8B3kEK8wsA,3004 +torch/include/ATen/ops/_scaled_dot_product_fused_attention_overrideable_compositeexplicitautograd_dispatch.h,sha256=gNB-hIMsciN7KHcvN8mLle7y_jBWKbj4s2l5gFRqLRU,1385 +torch/include/ATen/ops/_scaled_dot_product_fused_attention_overrideable_native.h,sha256=7rUEqX__mv3Gigo808FQLL926uKdLbnFc3l4MsOhl98,1097 +torch/include/ATen/ops/_scaled_dot_product_fused_attention_overrideable_ops.h,sha256=aL877R44j4ObBaTMB1coDQ-wmTybIuE1EQbSSlojI2A,2448 +torch/include/ATen/ops/_scaled_grouped_mm.h,sha256=5fH8ry71RiDFW-8OfBXBPvo6iTisN0VgfdFtKWx2t2Q,1490 +torch/include/ATen/ops/_scaled_grouped_mm_cuda_dispatch.h,sha256=pOmC0lABNzTDimislWqwtXfApjj7iHlZb2HysadRKgE,1297 +torch/include/ATen/ops/_scaled_grouped_mm_native.h,sha256=JyBtRzlUDxKfBRSalre6Qo4g53irZ8FaTdg-He2-iRA,1056 +torch/include/ATen/ops/_scaled_grouped_mm_ops.h,sha256=jQZgnj4Pp8iSdGG1Qlh-egj7vLZL_cdr1qMvJRLc7Fo,2180 +torch/include/ATen/ops/_scaled_grouped_mm_v2.h,sha256=trd4A9GAer1lzzPXfdudscoBUSvqlxyRpAE5CPta21g,1697 +torch/include/ATen/ops/_scaled_grouped_mm_v2_cuda_dispatch.h,sha256=LuQcyCfWlwKAgJZtuWwJKpKGUq-XJP-DFc7qhTc4--0,1381 +torch/include/ATen/ops/_scaled_grouped_mm_v2_native.h,sha256=aPa3n-jxYTLKlgTe6fyls8oPCSgW-4tP6s4jbwkWpoQ,1140 +torch/include/ATen/ops/_scaled_grouped_mm_v2_ops.h,sha256=Qg7xYjRsP1NzsGl_ZTyCJG8jMSMNhZim0uMzvzFo2T8,2460 +torch/include/ATen/ops/_scaled_mm.h,sha256=SxYr8M0K-PvNWV0rq1SXxqk6r4B4RfwlaiTL0R1LELU,2737 +torch/include/ATen/ops/_scaled_mm_cpu_dispatch.h,sha256=uXKAuos4jLgSi2NYYH4Ql1aVrDHC0o3qOcgSw56ahtY,1908 +torch/include/ATen/ops/_scaled_mm_cuda_dispatch.h,sha256=TmK3LinNsD1INvIPaVWtVvRweI6PmDHptC2f5n0aeqg,1910 +torch/include/ATen/ops/_scaled_mm_native.h,sha256=G2dQN2ulY42Ts9dRuI89IfcuErWjSt1NsdtLjqp3Teg,1976 +torch/include/ATen/ops/_scaled_mm_ops.h,sha256=QGCS_zvD2Dk2P0lCYPBuRDhAc8W5dTnckSnwMkgiVCQ,3411 +torch/include/ATen/ops/_scaled_mm_v2.h,sha256=uDzVdSWYu1m8grLHpY-aIUueeqM9C5dbsM2zPbQjN9A,3286 +torch/include/ATen/ops/_scaled_mm_v2_cpu_dispatch.h,sha256=h0bNzzmbkViof6KexaXiu4jEW50ktCRpNd3K_LV3TuE,2124 +torch/include/ATen/ops/_scaled_mm_v2_cuda_dispatch.h,sha256=MzYexNGWBwhumd6ja18Fb7Ow5yBuUbK-fMzNjVAK27M,2126 +torch/include/ATen/ops/_scaled_mm_v2_native.h,sha256=GpyI3mFA3jx0V-7pRJHkxHLJvsFUHGyVzgmCOiRuFNg,2276 +torch/include/ATen/ops/_scaled_mm_v2_ops.h,sha256=JtTwMckuG0DEn2qZkmGHkAajEhcS2v1807LThlCpOWQ,3951 +torch/include/ATen/ops/_segment_reduce_backward.h,sha256=vWDpFfJnjyTOSmc2nHlmlO5A8UfRQ2kadsMQKuQdpes,2707 +torch/include/ATen/ops/_segment_reduce_backward_compositeexplicitautograd_dispatch.h,sha256=QmJhzSRJYKQhDspir9ZHg88cb5TbVhkhm44orgLHXmE,1632 +torch/include/ATen/ops/_segment_reduce_backward_cpu_dispatch.h,sha256=fBCvjzY1-eqrongBoCAnyaTux29OtpQWttWOV4tHzx0,1241 +torch/include/ATen/ops/_segment_reduce_backward_cuda_dispatch.h,sha256=Eg8VaPpl5O1S5zDpAQCTsKuGvi6gIu_K7pogSCJo1jY,1243 +torch/include/ATen/ops/_segment_reduce_backward_native.h,sha256=eY7cybZJSDck8JxrPYdsgQYD2UxiJ_ang_06pW95i4s,1326 +torch/include/ATen/ops/_segment_reduce_backward_ops.h,sha256=XCDcxrC0_5e4K7IkF6a-O-0DNkABncuRD-X0B1kRPQE,3408 +torch/include/ATen/ops/_shape_as_tensor.h,sha256=5cBcV9Y3mIgrz9FYwsv6P9AFPV0NYa_32s3PXwJAqUk,931 +torch/include/ATen/ops/_shape_as_tensor_compositeimplicitautograd_dispatch.h,sha256=LkHLx596h8aTto0QYiS5Bp1ZRgkUzAd2ivtQx_ba1Yw,1028 +torch/include/ATen/ops/_shape_as_tensor_native.h,sha256=6_jJzgkp8kEN5q18QSiV9y98CXgCm0DDsrhNRwRjOiI,740 +torch/include/ATen/ops/_shape_as_tensor_ops.h,sha256=Q_WY76hISVwIcpo9qukzHnyNpmSCViT2j5LXCuJgo9I,1241 +torch/include/ATen/ops/_slow_conv2d_backward.h,sha256=q5cCUlgUROJMkEVsWUt4B61fstGH8vA5h_LdG7IIrlc,14784 +torch/include/ATen/ops/_slow_conv2d_backward_compositeexplicitautograd_dispatch.h,sha256=1eq9KUm7R8OMCQAgJj0AaBCVyrG27_WrgUD1Ii_VEbQ,2376 +torch/include/ATen/ops/_slow_conv2d_backward_cpu_dispatch.h,sha256=9ZG5qi3HQfBISIRCoba20QYy_7MeZJlt1LdlWY39Ogg,2839 +torch/include/ATen/ops/_slow_conv2d_backward_cuda_dispatch.h,sha256=qE01tDtOcLqdmU3Md2B1t90Edp1Oa5Qh8WCAUkhTeEA,2841 +torch/include/ATen/ops/_slow_conv2d_backward_native.h,sha256=v3MfyZWyD2h8AYRIT7gQfHrx3K7B3RD4UFxLEj9uM50,2270 +torch/include/ATen/ops/_slow_conv2d_backward_ops.h,sha256=xJk_txYe9lmIZIrP-NKKij__-fsPnPNvzdn_cDnqyAU,5115 +torch/include/ATen/ops/_slow_conv2d_forward.h,sha256=qxkRvjqxn357DQSz78KkkT6w69lrrg23pAxZpazZxRk,7035 +torch/include/ATen/ops/_slow_conv2d_forward_cpu_dispatch.h,sha256=KRL0Jud3OYaSu610QTiFpgyGZ42a4FS_hiaEyVJV39s,2377 +torch/include/ATen/ops/_slow_conv2d_forward_cuda_dispatch.h,sha256=fgfhOSXJp5OZSrJDT6LSIIWHKq0esaOSX3YB6dFd1Wc,2379 +torch/include/ATen/ops/_slow_conv2d_forward_native.h,sha256=5HTfK0oOpvJJ6I_WUDtyIbUCc9RNmsI20pzu7wzhhvQ,1604 +torch/include/ATen/ops/_slow_conv2d_forward_ops.h,sha256=IyiW8shg7orMvtokL5slUchXv9rbR62TD9Hadn39jVM,2957 +torch/include/ATen/ops/_sobol_engine_draw.h,sha256=lvsUWXNIgr9r56OOVJFq9CRz7vpxsdHJFPBPQ2zMzyI,1227 +torch/include/ATen/ops/_sobol_engine_draw_compositeimplicitautograd_dispatch.h,sha256=65Ym5KrYaQNz4TdyxF0GgBpxPb2Xy3BMJDGiaFwMnH8,1179 +torch/include/ATen/ops/_sobol_engine_draw_native.h,sha256=vNo9jtij6VeNmN9dXuzmknkeph_0BRsgmvxT5p5_oiU,891 +torch/include/ATen/ops/_sobol_engine_draw_ops.h,sha256=1aXyvdNp-nUZgxAz8qNLs4N-YvacCKfNeQiPyKwhfbU,1740 +torch/include/ATen/ops/_sobol_engine_ff.h,sha256=ZH5vRLbXj81-ahRDCQq6c16lg489WdzUyage0k1zV8U,1123 +torch/include/ATen/ops/_sobol_engine_ff_compositeimplicitautograd_dispatch.h,sha256=iF1A81S53AIupGsoa7Kru6LC-lpWRYnHU6WEVccLUyM,1109 +torch/include/ATen/ops/_sobol_engine_ff_native.h,sha256=cNK56cUosSNE6BK1JZEbWK8YpzdBPwnVD_dY7u9wV70,821 +torch/include/ATen/ops/_sobol_engine_ff_ops.h,sha256=3GLaQlauLSyTMAXDjAK0jx8RuALpvL7MntnEMbee_b4,1515 +torch/include/ATen/ops/_sobol_engine_initialize_state.h,sha256=_PdWJ2ZbqLToQuWJUF8zC23aFA4daA5SKUW7EQz0k0o,1039 +torch/include/ATen/ops/_sobol_engine_initialize_state_compositeimplicitautograd_dispatch.h,sha256=u148FflLZtFN_nfTPzAhpxQkSNPV32vU6bCKtRQHMMU,1058 +torch/include/ATen/ops/_sobol_engine_initialize_state_native.h,sha256=vJA4CfAXV9boGHvHdbcKcSf1QjrUWMOB94rQiQ8H-fU,770 +torch/include/ATen/ops/_sobol_engine_initialize_state_ops.h,sha256=vexYzpW2yMRfwXqHVqriXgb4IXlb1x7CrO31Nh-LNws,1344 +torch/include/ATen/ops/_sobol_engine_scramble.h,sha256=DQOI4cEmgNk-4engAmYN-O6crpFPGhltET9GDqpWdtQ,1048 +torch/include/ATen/ops/_sobol_engine_scramble_compositeimplicitautograd_dispatch.h,sha256=e7NWRZCxyhN7YsebZOzdId9XIYNT9BL1Ft5bhKT1EfM,1074 +torch/include/ATen/ops/_sobol_engine_scramble_native.h,sha256=wD6RyT-6CfRE626PmzOxucBmDbkK0OGJ2v7nmpTa8OM,786 +torch/include/ATen/ops/_sobol_engine_scramble_ops.h,sha256=lq8Abkv4aJJL2qA8y8lmXlMtwb9YBMg4tA6pnbkKIRE,1400 +torch/include/ATen/ops/_softmax.h,sha256=8lOsfawl-QWdBIWspiWtWMOWalcHFMQIasvUolQ7y8g,1547 +torch/include/ATen/ops/_softmax_backward_data.h,sha256=dFj_eQOb2zHjd9U9iQMFqHgDSXLF76Gmf2PBLVjuhV8,1972 +torch/include/ATen/ops/_softmax_backward_data_compositeexplicitautogradnonfunctional_dispatch.h,sha256=dpoxWsAEphwDD8laI2JQkrOOCFBTFwL8n49Nk2M4Z18,1135 +torch/include/ATen/ops/_softmax_backward_data_cpu_dispatch.h,sha256=Z47zOaFUgXtAtlvonugAgVKfrSV5mPH6B-nNhb9EzSM,1418 +torch/include/ATen/ops/_softmax_backward_data_cuda_dispatch.h,sha256=wGhaqWx7Z8wi1lmnBac_dwzXI0SvYcQwlu6RWs8aKYo,1420 +torch/include/ATen/ops/_softmax_backward_data_meta.h,sha256=cxDyRkbKF_4azg6yjdUoYvK0AxQ7fYlq1DqkB8tdPsg,912 +torch/include/ATen/ops/_softmax_backward_data_meta_dispatch.h,sha256=HdM-irR7kbKN9l4yENrmJzCDeWAc1tC5WqscbeNzIrA,1420 +torch/include/ATen/ops/_softmax_backward_data_native.h,sha256=o3mB1NEY5NboLqp1H7rGSH9MWQT4UPgsYPgSNgk_a5s,1378 +torch/include/ATen/ops/_softmax_backward_data_ops.h,sha256=0cxJZhGOtOUpTUkxJoq6-up6hwmeB-6l9XBUqP9MPU8,2418 +torch/include/ATen/ops/_softmax_compositeexplicitautogradnonfunctional_dispatch.h,sha256=pyg2QGOmd5bf_CcTCmF3bLiNiXwtdVp0ZiWtFhXiOYU,1079 +torch/include/ATen/ops/_softmax_cpu_dispatch.h,sha256=nWnn5RAdOuSOdxRWKovDc78M4dQv-G4ekxnf3-QS_3I,1236 +torch/include/ATen/ops/_softmax_cuda_dispatch.h,sha256=eyU0svBZ3lzVw-usu6SlQJtiFJcM6IjTHstX-YIzHCQ,1238 +torch/include/ATen/ops/_softmax_meta.h,sha256=aRU51tHchS-GOToBhfgBbLDa-1RKq2GdDdhmf0543Q4,856 +torch/include/ATen/ops/_softmax_meta_dispatch.h,sha256=obxSUmWx_vMfnR_vZOy7cVKmQk4gZKmcPcnawXJg5WM,1238 +torch/include/ATen/ops/_softmax_native.h,sha256=9hrvUOrNewBAEBVymKc2hfp1qT7_7YdvBF1VUMgQm5Q,1264 +torch/include/ATen/ops/_softmax_ops.h,sha256=e0_lj7QnXdXxtVHvqMbHaF3JDuC6cV4uB3tXuZS4c4Q,2033 +torch/include/ATen/ops/_sparse_addmm.h,sha256=JEdDAUeHzFcYB1FpwpQQW7o49AtinJCvrRhW0p31F1U,1911 +torch/include/ATen/ops/_sparse_addmm_compositeexplicitautograd_dispatch.h,sha256=51DXy9yeutts-LVrlp4J2IDDQHPCFi9iaNKFgzgZ65A,1507 +torch/include/ATen/ops/_sparse_addmm_native.h,sha256=drjrZr7Caqn2xmXATLFVBFFPbTVgfuPmsLe4MdZTckM,1028 +torch/include/ATen/ops/_sparse_addmm_ops.h,sha256=mhHRDikq1oFt8Ru8xFa1sf94bcPCSdS5ZL2rwcwLFqw,2524 +torch/include/ATen/ops/_sparse_broadcast_to.h,sha256=yUS1PDFUk1Jb0PDUX3g0ot-M5bMp4nIVbP7qrASbWLM,993 +torch/include/ATen/ops/_sparse_broadcast_to_copy.h,sha256=nrKzCAJzPSeZKvTnDO5eavcSGyIL8BJczKJ9T6_Z5Ro,1591 +torch/include/ATen/ops/_sparse_broadcast_to_copy_compositeexplicitautograd_dispatch.h,sha256=zrtDHs6SaISFfJSasbJJV4ziG-sX5WA2SC_m9JtgurQ,1203 +torch/include/ATen/ops/_sparse_broadcast_to_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=TK2VUzQzMHc8E4KsaTaHvSxp1RNvHI4t0SmfdLpoIMY,1085 +torch/include/ATen/ops/_sparse_broadcast_to_copy_native.h,sha256=2iik3L_kA3bawcG0SVqdXovWDGR03FV6OhJjwaVhF-M,890 +torch/include/ATen/ops/_sparse_broadcast_to_copy_ops.h,sha256=em22CsWUQz8V9_ZJAQC0JLka6szWPGsFAQQKLfRdogs,2061 +torch/include/ATen/ops/_sparse_broadcast_to_native.h,sha256=eZPbTAYymLlhoPzk_1pkDGu062cygSzL95GLXRwy2Hg,765 +torch/include/ATen/ops/_sparse_broadcast_to_ops.h,sha256=Kz8GL2O53DLp9cBtOiqUacRnuzchb_jVxa0zlC2N_8Q,1332 +torch/include/ATen/ops/_sparse_bsc_tensor_unsafe.h,sha256=A2gN1HzWjMa_tMcO-72B6tadPtBAbuKeQYuSxbFYKKs,2034 +torch/include/ATen/ops/_sparse_bsc_tensor_unsafe_compositeimplicitautograd_dispatch.h,sha256=uYmltbo0r8Y1eE94IBKOTcwmUXZF47oCl8MZLB7JvFE,1463 +torch/include/ATen/ops/_sparse_bsc_tensor_unsafe_native.h,sha256=sKWgsCui-VOs1KC5YqP8hgH_FLcQtJv_BYAbkPJgtdI,995 +torch/include/ATen/ops/_sparse_bsc_tensor_unsafe_ops.h,sha256=mfpeNSPJRWupTJUgnasOtDtzGrAi-qUIOYlmp37vAQQ,2054 +torch/include/ATen/ops/_sparse_bsr_tensor_unsafe.h,sha256=Q9xkWvd4fFu_2-FIejybvXsjrJXAD55JebRKdeEhp9M,2034 +torch/include/ATen/ops/_sparse_bsr_tensor_unsafe_compositeimplicitautograd_dispatch.h,sha256=1a0LwgVW-jOhsbCIDot7Yguay54rGphf8YMvuYbESIM,1463 +torch/include/ATen/ops/_sparse_bsr_tensor_unsafe_native.h,sha256=7SGI_yU64V1K-U7CLYx51A1zl8cO5VqpEi2f1uwxdsI,995 +torch/include/ATen/ops/_sparse_bsr_tensor_unsafe_ops.h,sha256=GPvZEdldYhgn0Z5DOYhk6gdJY_lsyRH2AXv4QKcRfik,2054 +torch/include/ATen/ops/_sparse_compressed_tensor_unsafe.h,sha256=zeEkVnGhQg2wxwNl3vZrpR5DEgZ0yAO44k0m48QSGVM,5848 +torch/include/ATen/ops/_sparse_compressed_tensor_unsafe_compositeimplicitautograd_dispatch.h,sha256=yapsu8aA6yNbvHAAa2AtfwEfHEd2u8mvvjpDA8C9vs0,2044 +torch/include/ATen/ops/_sparse_compressed_tensor_unsafe_native.h,sha256=4QGn8qwVFR8jgs2yGbRbaW6HFocgS-fR2-1oBwFcRew,1021 +torch/include/ATen/ops/_sparse_compressed_tensor_unsafe_ops.h,sha256=ya8ZkcT2VTETU0zGDLZu5qbABidrTNbthlALCRGt3tw,2114 +torch/include/ATen/ops/_sparse_compressed_tensor_with_dims.h,sha256=Txw7x1xRKi8Hg-gomGnwVHwQWqvIGyYxKEI0WdiMi_E,2121 +torch/include/ATen/ops/_sparse_compressed_tensor_with_dims_compositeexplicitautograd_dispatch.h,sha256=gU9X8no1v8UPgW0K9rsQUPmt3edA9z72-OPAtS2hvXM,1470 +torch/include/ATen/ops/_sparse_compressed_tensor_with_dims_native.h,sha256=wSjtN6NMkx-wabA6DmMJAVWjAbfzmjtS_gxXddpNIVw,999 +torch/include/ATen/ops/_sparse_compressed_tensor_with_dims_ops.h,sha256=Qxbr_UV6rG8_rtVPEts_kgMANSmr4tmGQDL3mn-Nrr0,2075 +torch/include/ATen/ops/_sparse_coo_tensor_unsafe.h,sha256=MRqgWKE7p1XjAbXSi385DNQQ5fgc9Rsg0YChCfu9HYI,5561 +torch/include/ATen/ops/_sparse_coo_tensor_unsafe_compositeimplicitautograd_dispatch.h,sha256=I4COQMvvNAqKCnBG4sviRtpuGGJrdqa9b0vQCoYTqOY,2010 +torch/include/ATen/ops/_sparse_coo_tensor_unsafe_native.h,sha256=PNN84LpBpZgp6l-QwsweQ554iugLVaASvufw5Cw-lwg,1020 +torch/include/ATen/ops/_sparse_coo_tensor_unsafe_ops.h,sha256=ZCYYGDivCKCjXlFhxeORHFG6B7RkR_oB4xJBSL_usYU,2070 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims.h,sha256=hP2mu_mF1NZKJ-UuGoGfyn2Lr258zx-tz8tS816-1uQ,2574 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims_and_tensors.h,sha256=_ud5eZHzzm3A3miiomMbVNKn8K5a4Qnp7DDV25tuhJ4,10866 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims_and_tensors_compositeexplicitautograd_dispatch.h,sha256=7fhZ8UmuXz4UqCy5TzyLaH2zZf0k7ygbrg_8p3NAZwI,1974 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims_and_tensors_meta_dispatch.h,sha256=BfX_2JnsWN4_zjB_7kWtdr_79ODjb0GZc1AZWbGrP9Q,2178 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims_and_tensors_native.h,sha256=kuOmLICVh3tXoIgtMFgnItiLHO4bnRZNj8wP1UithMU,1315 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims_and_tensors_ops.h,sha256=jeL9qWLxP-JsVCLIM_-cnXAeS6QGTju45_9x7i_QAmE,3368 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims_compositeexplicitautograd_dispatch.h,sha256=unUeTIogxeXNsACfFRqOLY8FqwezsiXsltiVw5ZF7Qs,1237 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims_meta_dispatch.h,sha256=8vuvVso7W4x_SsjvxQTS9cfdjPojVurwcPfJq0Z_v7A,1318 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims_native.h,sha256=8UBpjmPNTLwa0Lq3DUpKNUVAlYcWb1LE419r4EeXxFE,1073 +torch/include/ATen/ops/_sparse_coo_tensor_with_dims_ops.h,sha256=6dA9JReQiQRoNhZKx2VA_Vdr6cIxnyM14U599u4mi-c,2664 +torch/include/ATen/ops/_sparse_csc_tensor_unsafe.h,sha256=Kd6aJgw73c6W2TiVTKZdlPmcjgk9WW2JKmWuQ_8dQ8g,2034 +torch/include/ATen/ops/_sparse_csc_tensor_unsafe_compositeimplicitautograd_dispatch.h,sha256=0ACSp5RERfN0O-p7GjUKcx88IXgqkSruQZWamVqoeGM,1463 +torch/include/ATen/ops/_sparse_csc_tensor_unsafe_native.h,sha256=rkp4pNmus9mptxFaMQaGivfx2zQpUgtWqL0w_zeLuhQ,995 +torch/include/ATen/ops/_sparse_csc_tensor_unsafe_ops.h,sha256=tPrEXo5lJ7nL9RBUeXphUvkFsVGM5U89HrauLz-44Kc,2054 +torch/include/ATen/ops/_sparse_csr_prod.h,sha256=pU8xfUzU3ciWeMOAUupWigKMM1GYmbokfEKV9CM5lPI,1939 +torch/include/ATen/ops/_sparse_csr_prod_compositeexplicitautograd_dispatch.h,sha256=eCYI3JZAipAMwj2RsFQC8EqUQNLQPwya8t2zqMYQyp0,1310 +torch/include/ATen/ops/_sparse_csr_prod_native.h,sha256=KbO2KlLvXp0HxNJAHYg7ovb1zYTW3THJapmYtDbOaCw,1175 +torch/include/ATen/ops/_sparse_csr_prod_ops.h,sha256=tTq7zKGz20lG_I0EI-Oj6arlHoxdlz-YKMWcweOzHTw,2443 +torch/include/ATen/ops/_sparse_csr_sum.h,sha256=W70LlX-1oMqF84EaH5ApjGZIz6gOlYRjI5AXq5KRXIk,1929 +torch/include/ATen/ops/_sparse_csr_sum_compositeexplicitautograd_dispatch.h,sha256=RusXksQYxdWOExp_15c-T3C6WlwaARbR07ofnInjrXc,1308 +torch/include/ATen/ops/_sparse_csr_sum_native.h,sha256=_2Dl4mymkm4810qo8Om1rrZu6XUk290OZR6t4PWtvR0,1172 +torch/include/ATen/ops/_sparse_csr_sum_ops.h,sha256=zjfamLMz6ydl-Qf7DY7bPel2HsBsFkSSXzURZ11vLZM,2437 +torch/include/ATen/ops/_sparse_csr_tensor_unsafe.h,sha256=IPaUHDUl_KojGaczSaxcQd73zbvta1yoYfGAR9q5Sto,2034 +torch/include/ATen/ops/_sparse_csr_tensor_unsafe_compositeimplicitautograd_dispatch.h,sha256=IEcOH_XXfEZ6lolx2HD9_1j-3zs-eOQ3X24x4y6Hnvo,1463 +torch/include/ATen/ops/_sparse_csr_tensor_unsafe_native.h,sha256=d9V63axTfdAKYN3bCqeIqp2eynwOuF2xreUxD8tQrQI,995 +torch/include/ATen/ops/_sparse_csr_tensor_unsafe_ops.h,sha256=77KzwoW5hNfLc7yh0e-YLVSDCwWqdSCBaMuUSuh62Hs,2054 +torch/include/ATen/ops/_sparse_log_softmax.h,sha256=Cidw4F-GJcjHpioYr4ZVIipdtKqc2o0Jlh2irqyq0G8,2264 +torch/include/ATen/ops/_sparse_log_softmax_backward_data.h,sha256=7Y6uM2Ddf2V6uRfrioAA7BkzOfVWDooc_64GAUlIk9U,1977 +torch/include/ATen/ops/_sparse_log_softmax_backward_data_compositeexplicitautograd_dispatch.h,sha256=eWodcdmSG8PSk2uLYpjqbLaPreyMekhzk2dmqgLlsJ8,1319 +torch/include/ATen/ops/_sparse_log_softmax_backward_data_native.h,sha256=Hub86o313x5WOi7dafisMLg0-IgP_Ultz5valJx8NG0,1156 +torch/include/ATen/ops/_sparse_log_softmax_backward_data_ops.h,sha256=HSRDJqXo8zG7wTkli4MAtqgIT8Ny3hHty7Hnw4QcuyM,2437 +torch/include/ATen/ops/_sparse_log_softmax_compositeexplicitautograd_dispatch.h,sha256=QgIjZskOPK4xGBj8mff7SMwovE27mRT_-gE5RKOojwc,1213 +torch/include/ATen/ops/_sparse_log_softmax_compositeimplicitautograd_dispatch.h,sha256=LpomyYXIbISeqSuBSi--f4Fn2iHaaN_F-9d3YHgYCf0,1236 +torch/include/ATen/ops/_sparse_log_softmax_native.h,sha256=1oNEdQyExba9-5OALr21g3AewyGirzDwNoroGPrn5MA,1279 +torch/include/ATen/ops/_sparse_log_softmax_ops.h,sha256=nA27sg0wvjV1MPCqcq3mbclATKVjyVlu7V12eIbM70Y,3588 +torch/include/ATen/ops/_sparse_mask_projection.h,sha256=FoaHGbwWWMLd02YR_RsQ3yqEJohr6f9PNP3b3EaO_Pw,1511 +torch/include/ATen/ops/_sparse_mask_projection_compositeexplicitautograd_dispatch.h,sha256=U3o3tSyDMITR31LY46IPXQgAdVkS98_lZiIAY7Pklws,1261 +torch/include/ATen/ops/_sparse_mask_projection_native.h,sha256=cVjQfwmWBehkIZc6Q-Rgozct7cailHO8MpP51hXocyQ,947 +torch/include/ATen/ops/_sparse_mask_projection_ops.h,sha256=KI3xUidyM47D9He9FWuDdg2JbLgm39yj1iausnnuIaE,2243 +torch/include/ATen/ops/_sparse_mm.h,sha256=HfkW16IjYIL4egSq8DCoIJj4R8dHfae3TdzD0FmS6_U,1219 +torch/include/ATen/ops/_sparse_mm_compositeimplicitautograd_dispatch.h,sha256=pU1WH4vYdcJZDi3bIQkyqWOpyrCHHM-wTGRRmYdkJ2M,1161 +torch/include/ATen/ops/_sparse_mm_native.h,sha256=A2A3t-d4mUSGtoJ6lIQexV9irH3ZOia7zrGpCRYwT68,873 +torch/include/ATen/ops/_sparse_mm_ops.h,sha256=v7rKiTxWlJ5RKTwD4BbF9ZZkhziRHn7hQ6IPJcamWO8,2019 +torch/include/ATen/ops/_sparse_mm_reduce_impl.h,sha256=uEifMo5AFiy9XuMYA6BuptEKcAYEa5jC4JQ9gn0q7qU,1082 +torch/include/ATen/ops/_sparse_mm_reduce_impl_backward.h,sha256=BB279Lss-5J8I29XESfPgsIWmQ_MI9yKnDCrArRbvIY,1298 +torch/include/ATen/ops/_sparse_mm_reduce_impl_backward_native.h,sha256=e7awKjdhzjSh9-0J61H9Ecm5W90We5N0J1C_aB3e3Us,938 +torch/include/ATen/ops/_sparse_mm_reduce_impl_backward_ops.h,sha256=YWitpUyea5R-OzrypovoEThEmzVOgB4PKBWrLzyEq5Q,1838 +torch/include/ATen/ops/_sparse_mm_reduce_impl_native.h,sha256=7yWu1afm8uZ8lv31dISjHoNka8KQHAWnrxkH5-mCXmo,837 +torch/include/ATen/ops/_sparse_mm_reduce_impl_ops.h,sha256=VAJrRP31GkwuurZs5i-L9mjyZmROsfnaoaooJgBA7bQ,1510 +torch/include/ATen/ops/_sparse_semi_structured_addmm.h,sha256=mCAqiZsnMFaULuBWihG1UxK1vCrkLIT-hOxvkVIIYqg,1332 +torch/include/ATen/ops/_sparse_semi_structured_addmm_cuda_dispatch.h,sha256=IOywnxpQBMJoCa2urDcUPWimsB3hPYZGbiyrxIcnsnc,1193 +torch/include/ATen/ops/_sparse_semi_structured_addmm_native.h,sha256=TslEn0gAcV2rOaKPVcTLsQ6DDg7vt-kNyKOgyhG1QR0,947 +torch/include/ATen/ops/_sparse_semi_structured_addmm_ops.h,sha256=O1v-3RFtr9XdIQHik-2OtIcXqPofytf6_Di_5d4g9es,1870 +torch/include/ATen/ops/_sparse_semi_structured_apply.h,sha256=SuyU6NZUWh8rv9AYoxOQlPejy83FzqZBn5F8eujIejs,1089 +torch/include/ATen/ops/_sparse_semi_structured_apply_cuda_dispatch.h,sha256=b8k4KvRC697kVzTOOUV7XvezeYk14uQh4uhBo8kCOgE,1058 +torch/include/ATen/ops/_sparse_semi_structured_apply_dense.h,sha256=4urKlxLsz-b8WtuTI5utUMO_AwHY1AcSTuIxjhPgkm4,1078 +torch/include/ATen/ops/_sparse_semi_structured_apply_dense_cuda_dispatch.h,sha256=U0PZwukT3WhYuwyiM8MOunIMER_MziaNW-yWuVZ7lUg,1039 +torch/include/ATen/ops/_sparse_semi_structured_apply_dense_native.h,sha256=M80uZ_e0ykiJGUQyQBDMieoE8R5YJuJJ7tGScCuQaAQ,793 +torch/include/ATen/ops/_sparse_semi_structured_apply_dense_ops.h,sha256=ryjbK6f1ah6Uk5yjRzQY8bmbq9UEJfxPTjx5of3hlp4,1408 +torch/include/ATen/ops/_sparse_semi_structured_apply_native.h,sha256=Ra5TuyofUEyBj1a1s4jeOKVFzx-CI5hwq9p2EzY7KfI,812 +torch/include/ATen/ops/_sparse_semi_structured_apply_ops.h,sha256=J_4tvx4orqWUwopZC5GuV6F6UWcJAWjMf5l4eEPEhKQ,1475 +torch/include/ATen/ops/_sparse_semi_structured_linear.h,sha256=RitS-VqQD6z9Kr8nO42HG3l_v0c4Uz5WqwKyI2W8ehw,1349 +torch/include/ATen/ops/_sparse_semi_structured_linear_cuda_dispatch.h,sha256=bizinDywN1rMa_SP9kAMRC_SljifVzdrrTCthcdmY5c,1217 +torch/include/ATen/ops/_sparse_semi_structured_linear_native.h,sha256=AwJ9JnoSQz__M2c9Q6YL5XFrzBVCG4r8r7-7e7FhKb4,971 +torch/include/ATen/ops/_sparse_semi_structured_linear_ops.h,sha256=14GM-GjwN7Zqtww5JKnbL8H3TEqi59ZbWOggHiP1JVM,1897 +torch/include/ATen/ops/_sparse_semi_structured_mm.h,sha256=zl3wgkIqzZ4dQDqPtn9v1iUfFhRyE2QmgCTSXIGrsvE,1174 +torch/include/ATen/ops/_sparse_semi_structured_mm_cuda_dispatch.h,sha256=3VmapmQJgJWeQGHCwvLI-xCJOW6nMrjvahmQHIIb3vg,1109 +torch/include/ATen/ops/_sparse_semi_structured_mm_native.h,sha256=nneNw3v6IAgBaxt8raRGoo82MTgB2mWItRFKXCmJIog,863 +torch/include/ATen/ops/_sparse_semi_structured_mm_ops.h,sha256=-8B9GHe_4N3pP42VzJjZufl4HusSwT78YXnEZ9zCp1A,1602 +torch/include/ATen/ops/_sparse_semi_structured_tile.h,sha256=WAfC5WFpJGW59fNNubWF5iUg7cTpyGgYS9z7n7u1_gY,1193 +torch/include/ATen/ops/_sparse_semi_structured_tile_cuda_dispatch.h,sha256=4OJKQ5l2KGlX-o7Kdn_Tc6nSbXy4DRCDJRErH4GYcSI,1111 +torch/include/ATen/ops/_sparse_semi_structured_tile_native.h,sha256=olsupM3GOEAIkrO3_Zni3BtkuvcjPL0OJjsVdyKQZw8,865 +torch/include/ATen/ops/_sparse_semi_structured_tile_ops.h,sha256=u-IukdlhaQ894uXGIL0ZV2IISghVkYIsP_zUjuYbVYg,1647 +torch/include/ATen/ops/_sparse_softmax.h,sha256=x4p-0H_87O6IeohGjZHqjEc-s3SEC0co_ctjgkRqu44,2200 +torch/include/ATen/ops/_sparse_softmax_backward_data.h,sha256=iVMeIXD1SJsTHRE9zOKH35Q9WrellOb9KnwdXivDuQ8,1937 +torch/include/ATen/ops/_sparse_softmax_backward_data_compositeexplicitautograd_dispatch.h,sha256=Un5JWG1L6j8YqUGaaCe-YfaWjoxqXAYnq0GJeN00Bao,1311 +torch/include/ATen/ops/_sparse_softmax_backward_data_native.h,sha256=p6hXfyhQKvSFxLoQYX48jVol2csu6VNXBJqersAwlkQ,1144 +torch/include/ATen/ops/_sparse_softmax_backward_data_ops.h,sha256=w43m6xczEcQjkDpYzsaVNaB5vjplF86eKiqGfpbxR7M,2413 +torch/include/ATen/ops/_sparse_softmax_compositeexplicitautograd_dispatch.h,sha256=ciQS89KIPR6-GRA0Y5ugtxaDFy4o6LEUND6oE5FCEHE,1205 +torch/include/ATen/ops/_sparse_softmax_compositeimplicitautograd_dispatch.h,sha256=9TbKUv06GXMvLgD9acuC3iHraUw9A2fcKG93yMyBZWg,1228 +torch/include/ATen/ops/_sparse_softmax_native.h,sha256=dtNHO5qh-9fYQYgH8eMpUY9zgM8M0tt-RRY6Nh0y1j4,1259 +torch/include/ATen/ops/_sparse_softmax_ops.h,sha256=NarX-lC5H05UcniMJ83zB7H9vJ6_Ola3DFIPyM-iQfM,3540 +torch/include/ATen/ops/_sparse_sparse_matmul.h,sha256=e2UucJi5VgQvBpl6Z9H7PmWSY3GTZbwaMPsfCNj_G4M,1572 +torch/include/ATen/ops/_sparse_sparse_matmul_compositeexplicitautograd_dispatch.h,sha256=72-49i__SIo0HQpV63v1MvsppLp73vgcuyNm8Yxqf8M,1203 +torch/include/ATen/ops/_sparse_sparse_matmul_native.h,sha256=Yq-3LjgMSuExeMng8lLDVnI7yn4kx675rylGc1iQsm0,992 +torch/include/ATen/ops/_sparse_sparse_matmul_ops.h,sha256=zwknDCvBhb8lDmRBGEC3PuS6OP3qq5vleDMLJDY3gPc,2063 +torch/include/ATen/ops/_sparse_sum.h,sha256=QZjgSW1nbi3r0TLn_YRsqL833ogLtQU7AblcckvTOG0,2089 +torch/include/ATen/ops/_sparse_sum_backward.h,sha256=SdhxYiZir_ccdH-hpWfXkozbdBUeRAJ2rpwB4fCOd-w,1664 +torch/include/ATen/ops/_sparse_sum_backward_compositeexplicitautograd_dispatch.h,sha256=qkaIlTWUwrK7Rd44f-6MJDB0TkjK-olpfE7QklU3eUg,1241 +torch/include/ATen/ops/_sparse_sum_backward_native.h,sha256=Jv_VV7b-3QGFFShAxJmumHWplEl8FVFcVtymA4dlJz4,1051 +torch/include/ATen/ops/_sparse_sum_backward_ops.h,sha256=G9D7nr5P3o52X5PcGFIcgKbAlHK-cQu2vMr5ewi_NUg,2191 +torch/include/ATen/ops/_sparse_sum_compositeexplicitautograd_dispatch.h,sha256=0VSZOM5OGo3iijux1aAU5DxbWaW38Svf690IYp3Uwhg,1253 +torch/include/ATen/ops/_sparse_sum_compositeimplicitautograd_dispatch.h,sha256=S5_wOPtSALWkZEeJoGEi38eg50coYP5xFEHTGlvR3zQ,1206 +torch/include/ATen/ops/_sparse_sum_native.h,sha256=wK1grkBCuzNTQ_nm5RbcYeqwsF0zscc7hMNG0LCCykg,1106 +torch/include/ATen/ops/_sparse_sum_ops.h,sha256=T146SrkDcTzZVRAivQ2dl92JG9rr0jCvaKXEyXy4tpk,3820 +torch/include/ATen/ops/_spdiags.h,sha256=hmokoL_XkEDpK5gqTYIR7F-njd0WGNR3F4j7PZh730A,1859 +torch/include/ATen/ops/_spdiags_compositeexplicitautograd_dispatch.h,sha256=UbRPm531USYxmiw8A6rBnVSBhOwqB-UUo7o-XYTIgOA,1324 +torch/include/ATen/ops/_spdiags_cpu_dispatch.h,sha256=JdlH1ntLN6NxWtH1HwEoovXXnVi9kC3yBN4zXfnEie4,1083 +torch/include/ATen/ops/_spdiags_native.h,sha256=Zm3V1MJSf5vev3dLBMg-Lb8baiBBmCz9HA5-Ozr4kTc,1010 +torch/include/ATen/ops/_spdiags_ops.h,sha256=_0X9qyNcNn6HPfX3HWViH7rNP7N0XtNGTzar-OHd7Y0,2423 +torch/include/ATen/ops/_spsolve.h,sha256=P_jbKsoRCLSqlchfnfoMPhZI8tzBRKO1AjK-WnCiHW0,966 +torch/include/ATen/ops/_spsolve_native.h,sha256=KWTsfcG47pJvToZ2kSRpmAV_qSrQ8QXTAhWzTl-J9RM,783 +torch/include/ATen/ops/_spsolve_ops.h,sha256=jk6LRaRv16UHadUJAm4RDl8yowK5BpXW5V4eyN_MWvM,1329 +torch/include/ATen/ops/_stack.h,sha256=6ZidL51Sq5pJIuRpTyp0zMl0OZ30SC7McXC4e9gHUlI,1393 +torch/include/ATen/ops/_stack_compositeexplicitautograd_dispatch.h,sha256=d6FPj-EEKnwaE00xkZ36li_dF8yLkRtErLtUbsPnZLI,1215 +torch/include/ATen/ops/_stack_cpu_dispatch.h,sha256=tyj0nLjDPr6P3UOw47BKKuiGA1CgDqXe5u2m3Fz9Vvw,1171 +torch/include/ATen/ops/_stack_native.h,sha256=Cy_GV1hYym7KG5_6MCD1UgTnGhd3wu6jhk-sFpAcRlc,1000 +torch/include/ATen/ops/_stack_ops.h,sha256=h3pO7_QHRrZRQatSHEpisnT1iAl6PP6gN5FMyIgXWJ0,1891 +torch/include/ATen/ops/_standard_gamma.h,sha256=-KzniOFGogDHzok-kuEPStCYKJo6cfI0WwQmI8KfD2Q,1641 +torch/include/ATen/ops/_standard_gamma_compositeexplicitautograd_dispatch.h,sha256=-n5jp-Qoj30yUEdqFu-ieCVNgEr7Mn8TU_Ai6cfmIhU,1238 +torch/include/ATen/ops/_standard_gamma_cpu_dispatch.h,sha256=oDGhf689aXLHSqhFx3s6tlQ7RSu3WmoneVjupHT4vac,1040 +torch/include/ATen/ops/_standard_gamma_cuda_dispatch.h,sha256=mWGYhroqAyXC8pwiBnlfyNjnfHDewfQAa651IHEFH7w,1042 +torch/include/ATen/ops/_standard_gamma_grad.h,sha256=C2PxB5rU9p3ckbRX5euZ-fzhiA0pdW0cBG34R8LJ7e4,1571 +torch/include/ATen/ops/_standard_gamma_grad_compositeexplicitautograd_dispatch.h,sha256=8uydEf72uD9UdImqI0-JBxraOE8gFdCnCKITFMvf_zE,1203 +torch/include/ATen/ops/_standard_gamma_grad_cpu_dispatch.h,sha256=vqE9ir-H-pAeVk8lHrP2rU73W8_jQEgqRBVBpDa9gnc,1015 +torch/include/ATen/ops/_standard_gamma_grad_cuda_dispatch.h,sha256=qz_2wfki-yfGQ8jWCb4WdF3OyF1OvwI3fRMpsAOzipg,1017 +torch/include/ATen/ops/_standard_gamma_grad_native.h,sha256=M19WbHdhSB5G4Z2-gzhszcjAhkYDfWhkn7H9QizKeQI,994 +torch/include/ATen/ops/_standard_gamma_grad_ops.h,sha256=Du9ZuWG4BpLoSmlxroDtEJXQFu2VjEuiZUPI3ezqVLY,2063 +torch/include/ATen/ops/_standard_gamma_native.h,sha256=Y6hwkP3RjNYUh3YDqloQwRBmBMIACBoTos9US49wBDI,1040 +torch/include/ATen/ops/_standard_gamma_ops.h,sha256=HtyoxMF8Pu6N6XmQpuCEGERWu827L_ifOpUdAzYbmJY,2141 +torch/include/ATen/ops/_test_ambiguous_defaults.h,sha256=iiPtWn9ulkBfqYYjQoFRq5nYGChMY1kPRWlTX7UFO3w,1273 +torch/include/ATen/ops/_test_ambiguous_defaults_compositeimplicitautograd_dispatch.h,sha256=L75HozizZ7Iya6kQXRP4oe2tOXhBBEzVGtxWpG_2S-A,1167 +torch/include/ATen/ops/_test_ambiguous_defaults_native.h,sha256=jr9B7bYvx6C3fy6hCVfrusEHGaiJY8ypy-8YmlJbeiE,885 +torch/include/ATen/ops/_test_ambiguous_defaults_ops.h,sha256=cMVW6fX2gumn6g7HKde7zAQ8ItLL2pyf5IadLjPVe80,2026 +torch/include/ATen/ops/_test_autograd_multiple_dispatch.h,sha256=3yh8AVcE3X3A4NJkZWo_vdwBWcfchWejgoywny1Bn4k,1865 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_compositeexplicitautograd_dispatch.h,sha256=gq5orbrQFiF69pNB4VkDTrlQSPLJ6CDzIN1GiqwnVh8,1253 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_compositeimplicitautograd_dispatch.h,sha256=MX2T8vzFVdHkrW-a-WMNrJsQsx_r1qA5Hw3qusRZVKo,1052 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_native.h,sha256=eOQfDPVueJiTNQfR8Qcn9sdD1n4hqCuhE7Ki2ZDathw,981 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_ops.h,sha256=jne9YO_AT_O75rbKU2y3XcHntXssxtD4a_iX5wXAcR8,2662 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_view.h,sha256=Lqdwo5wPaxS7u9bTKJfuj8FynwHH4xTKCc8LAlC1ZWM,1021 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_view_compositeexplicitautograd_dispatch.h,sha256=cVsywZBhxk1JNeSKXmGNulTde20iGUnybrUH0s9wuvY,1049 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_view_copy.h,sha256=qOD7aMjHdtabPyI-x0zWhILvpN57IJhRklFjpHKLRMA,1641 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_view_copy_compositeexplicitautograd_dispatch.h,sha256=7xqOvc8zJt2JCc9eQIHglbttOHjtWTX3ywMnZzjICnE,1193 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_view_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=44Lo7Aj3C5QTkXrrhQFYiy0WsnHdHeHAa-saQqbFU9o,1080 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_view_copy_native.h,sha256=2UH3GOMfrNRo-ITvG_1_NA3nrebhI6clR1AHL2shGBM,880 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_view_copy_ops.h,sha256=wfc2QKwk1HSg-gpLUfOSUAJcsMkpkbhugH8OCJcGFfM,2017 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_view_native.h,sha256=zv9UT9OBO1RQyrthAffhkskiGIeXEtKGyoL4eQMfbS4,761 +torch/include/ATen/ops/_test_autograd_multiple_dispatch_view_ops.h,sha256=T2mbQs-DO5fEn904Ds_5y4Gqu-nhreSAJ9rG6tXyh_g,1310 +torch/include/ATen/ops/_test_check_tensor.h,sha256=RkhX2TkiDgaLZ23G8D5UtXnsmA7I-M6xn6H4_c0WEpU,939 +torch/include/ATen/ops/_test_check_tensor_compositeimplicitautograd_dispatch.h,sha256=yetb_9AbjPBYcafcVg7_KLX_-Oe4OK6RsQktBD_vYKI,1030 +torch/include/ATen/ops/_test_check_tensor_native.h,sha256=Yr56mwGyLjxYV-FYiLhe92eT--XPV4bRw-ymI-SIIdk,742 +torch/include/ATen/ops/_test_check_tensor_ops.h,sha256=zzSMYPS5AMfj57zE9RYHwAb31uGsb6M7VsXWNNk3nhw,1247 +torch/include/ATen/ops/_test_functorch_fallback.h,sha256=hZlXO84x0tHo1VDpvbnrwGzOJdjlxdkSk1yB2adNKo8,1602 +torch/include/ATen/ops/_test_functorch_fallback_compositeexplicitautograd_dispatch.h,sha256=35xKz2Q8FEMcI2OsqYXl6hdotm65t7y8jrjxfM4DXMQ,1209 +torch/include/ATen/ops/_test_functorch_fallback_cpu_dispatch.h,sha256=k5A4RAW4tCQsaDua4yjdqqkMFy3XdLz7G5DXaFF1bMk,1018 +torch/include/ATen/ops/_test_functorch_fallback_native.h,sha256=6JYXhXAJCjITbx1w1IcukfQ6vBVaMDusHr5otHGJ1HA,896 +torch/include/ATen/ops/_test_functorch_fallback_ops.h,sha256=Y_gPiQRrV787GGrwuWrNV8y5Z6-nu-HdNBvFwYx66O0,2081 +torch/include/ATen/ops/_test_optional_filled_intlist.h,sha256=lYrO4oOhnypBvMYCqJnWPdsZ_MfvjFdr7w_r1f5dAtw,1706 +torch/include/ATen/ops/_test_optional_filled_intlist_compositeexplicitautograd_dispatch.h,sha256=MUUnG2zgq8yUXrpEoTNQkso2xEZ-9B7b_5Q8CrTXt0Y,1237 +torch/include/ATen/ops/_test_optional_filled_intlist_cpu_dispatch.h,sha256=lcLe8-9tfyzFmtV_ZacLxjhQ2UXojJwXHN5zWYJbhH4,1032 +torch/include/ATen/ops/_test_optional_filled_intlist_native.h,sha256=VUQVnTJz7lR_WBl4RnYL7Fv9XcVd4KY8bhYDFRImlCY,917 +torch/include/ATen/ops/_test_optional_filled_intlist_ops.h,sha256=wXPxlWCQJcbVWrCYD6qXK12l5oYsl48ZnjscjclBSoE,2167 +torch/include/ATen/ops/_test_optional_floatlist.h,sha256=NXjBKt0Ih4z9s0wNF2V3kGgKHc1BixIeO1Byx7r_sEs,1701 +torch/include/ATen/ops/_test_optional_floatlist_compositeexplicitautograd_dispatch.h,sha256=Mk7xzLP7Zau1TS-9MadiV1KLCK2RO2XunxtXmsJtjxw,1255 +torch/include/ATen/ops/_test_optional_floatlist_cpu_dispatch.h,sha256=2p4O-2D29o6ZuHMiAuXmb2A93uSmJXR3MQHCbPoAwyY,1041 +torch/include/ATen/ops/_test_optional_floatlist_native.h,sha256=lU9leB1FcvJ1v-j10qhOj-2_OpHXO23DOrV3fxmvEY0,942 +torch/include/ATen/ops/_test_optional_floatlist_ops.h,sha256=CTwGUUkBgbbIJ3ovoE3-B-zBBzDkcWG0NJ4bVCR4AUw,2223 +torch/include/ATen/ops/_test_optional_intlist.h,sha256=gsoYLS-wG97IkuYqVsUt5fWPLFrzQtTXt6y1S_BLJig,1633 +torch/include/ATen/ops/_test_optional_intlist_compositeexplicitautograd_dispatch.h,sha256=c_I6ss6CQH0sL-I7ASwqDCvVSFjACOqD0fHh4eQ2kYA,1223 +torch/include/ATen/ops/_test_optional_intlist_cpu_dispatch.h,sha256=O1mM_o7BAe9mwbORP304rlJ2UJB_ssm4VBlJcz4NiuA,1025 +torch/include/ATen/ops/_test_optional_intlist_native.h,sha256=Csx-clXQ8K1etHdDBh2jDn9zTrAlWtQKqDa44cS5Xks,910 +torch/include/ATen/ops/_test_optional_intlist_ops.h,sha256=Zp7jqQikbc1qUyBHBib5TreRymYF8YYkDid_LiPJ9es,2123 +torch/include/ATen/ops/_test_parallel_materialize.h,sha256=8iVOsGVdf7hlUmfgJzVieIiCvonMVl2B6asSror2RLY,1083 +torch/include/ATen/ops/_test_parallel_materialize_compositeexplicitautograd_dispatch.h,sha256=gG0Lzh7Gtz32PWB771DD2VEe-8YgKc8tlm-GiWCiQZ8,1083 +torch/include/ATen/ops/_test_parallel_materialize_native.h,sha256=RlO4B6CrscoSEQhpCho5wtrpQpwMHzcVw8T7B7zeH_U,795 +torch/include/ATen/ops/_test_parallel_materialize_ops.h,sha256=Wwc16ZTWV4Ltk7Ps3PD7f0gjuwCorXFhexScEqkB2Yg,1405 +torch/include/ATen/ops/_test_serialization_subcmul.h,sha256=c94gI8FK_cSpWBUlJNAcPS3_SCm06L1ac3rD43Ja4MY,1073 +torch/include/ATen/ops/_test_serialization_subcmul_compositeimplicitautograd_dispatch.h,sha256=9F9Zx2vOVVizB6CWR37s1FBsriyxLc0xkham2u5N-_M,1093 +torch/include/ATen/ops/_test_serialization_subcmul_native.h,sha256=Fph5WdDjpuPj6Gr8ag6unMEXoNln49Xb9P02LYbA_y8,805 +torch/include/ATen/ops/_test_serialization_subcmul_ops.h,sha256=JcV9JJwbX32tMC-NdQcDEZno-0Wwww_qqxCAtF-uowo,1448 +torch/include/ATen/ops/_test_string_default.h,sha256=MkDq9X7C6LNIqbaAdigNAoYCpGNUrJC3VEozLxfjZEU,1042 +torch/include/ATen/ops/_test_string_default_compositeimplicitautograd_dispatch.h,sha256=95twlmJzcdPqZNp3dBO9zwCJNrX1A3Vany6gkB9LAI4,1089 +torch/include/ATen/ops/_test_string_default_native.h,sha256=YuUCUnOdzdV9Xg0BpzcI_TwuHE0PemIROwDwqltWonA,801 +torch/include/ATen/ops/_test_string_default_ops.h,sha256=xzUdQK1Q5WpI8N2RabosI5krcQOnu_c6Zsut3x13eTc,1412 +torch/include/ATen/ops/_test_warn_in_autograd.h,sha256=Gi-ntGUk4XePMOqdAfNVYU5czjLksh7dIt-4G0pyBwI,1441 +torch/include/ATen/ops/_test_warn_in_autograd_compositeexplicitautograd_dispatch.h,sha256=_Q3jsnVyp3c4sa7YFe8rPaJCis0AUfFuvMKey9In47o,1223 +torch/include/ATen/ops/_test_warn_in_autograd_native.h,sha256=fRS5fd_8Bze8ZmdQoU8IEJW5KaixpzBNDoU2qzfcJ34,840 +torch/include/ATen/ops/_test_warn_in_autograd_ops.h,sha256=BVaQph_lXY3xiDNy94G4Nkme-xVp5SG8NDyEE6v_fVA,1897 +torch/include/ATen/ops/_thnn_differentiable_gru_cell_backward.h,sha256=3rYNygXOxdIInxRTzgx8lxcbvglR5ps9fz43xyew62Y,1454 +torch/include/ATen/ops/_thnn_differentiable_gru_cell_backward_compositeimplicitautograd_dispatch.h,sha256=ayMB6RaUjgaYGXhbNROL_k75g78xr2IgH7JSBFFiMOU,1296 +torch/include/ATen/ops/_thnn_differentiable_gru_cell_backward_native.h,sha256=9bQRwlW-KZPiezKunCyHLaHWGOIYqZWf_VNm--Rlxc8,1008 +torch/include/ATen/ops/_thnn_differentiable_gru_cell_backward_ops.h,sha256=LkUlqgdeziGo26TdMXNhpVg9cbiMtn8LRGV5upwJ3jo,2121 +torch/include/ATen/ops/_thnn_differentiable_lstm_cell_backward.h,sha256=7ORzQuyn76mxppmfJ4z3yJJnb6wcRYq46XlHcTGcVBA,1585 +torch/include/ATen/ops/_thnn_differentiable_lstm_cell_backward_compositeimplicitautograd_dispatch.h,sha256=ZAhJe97_XioS_mr3Hr_GwSe_-QWA7ZJZOq5MuD2fyo4,1382 +torch/include/ATen/ops/_thnn_differentiable_lstm_cell_backward_native.h,sha256=eTaSLHJcoGl-GTkoog70UD51p_Svj9jfWFre0jnEvBs,1094 +torch/include/ATen/ops/_thnn_differentiable_lstm_cell_backward_ops.h,sha256=ulRHgb1PdTcg54seN0gP8t4DGMIjn0ysJKFxLBmDkgQ,2397 +torch/include/ATen/ops/_thnn_fused_gru_cell.h,sha256=TZD6mm3HwOcvsxUdScXu_Y0tLmodM1OMx4jKtRnIjIQ,2540 +torch/include/ATen/ops/_thnn_fused_gru_cell_backward.h,sha256=blgtOG-qOt3OOghh7pwb2RSlz_sMOI_eXh5RwoYYb7s,2501 +torch/include/ATen/ops/_thnn_fused_gru_cell_backward_compositeexplicitautograd_dispatch.h,sha256=CVLGfjVwdhBsSElbfsDe1IGXwJxJRTgT8K5Y8ImQpgI,1549 +torch/include/ATen/ops/_thnn_fused_gru_cell_backward_cuda_dispatch.h,sha256=m3ry3x_6zrZZcbN2QSjXo0skv_ZW-LRkHsC60ObZm14,1105 +torch/include/ATen/ops/_thnn_fused_gru_cell_backward_native.h,sha256=tUbVl5QyCKVvr9srb8mClH9AspZuFZW2NAOoxVvEq_E,1156 +torch/include/ATen/ops/_thnn_fused_gru_cell_backward_ops.h,sha256=VvluDNLP-aMG7nGBsV-Auf9tgwlnm_32aSvftzOjX7c,2990 +torch/include/ATen/ops/_thnn_fused_gru_cell_compositeexplicitautograd_dispatch.h,sha256=QaTe26tIyQkScWIvOQWpRW0gfPEwuZksM9m0A-ht2BE,1569 +torch/include/ATen/ops/_thnn_fused_gru_cell_cuda_dispatch.h,sha256=9PFDuacMncKGt5iniW7ji3MsJM68lcwatlRC7C-Tbj8,1181 +torch/include/ATen/ops/_thnn_fused_gru_cell_native.h,sha256=k_zNAKP-upd44yiioQz5Epf8z9_fLbqImJszy0dp0-w,1239 +torch/include/ATen/ops/_thnn_fused_gru_cell_ops.h,sha256=ZwXU1Hq2yTYjS1dm2txZ7YS8y1_LRw4m4e_ReuciEUo,3185 +torch/include/ATen/ops/_thnn_fused_lstm_cell.h,sha256=LxpRssFRX03DOkSdHv0SZOGI4nopHg5EYlKM97DhgRE,2703 +torch/include/ATen/ops/_thnn_fused_lstm_cell_backward.h,sha256=eRZxbhFrEx0Ct27bS0u-KIzIzeu2kwqvqwiwQQSxlw4,1352 +torch/include/ATen/ops/_thnn_fused_lstm_cell_backward_compositeimplicitautograd_dispatch.h,sha256=2Qm_7BsP8P17VYxji3AmEIamfK5m5N958eOyrQUm9mU,1256 +torch/include/ATen/ops/_thnn_fused_lstm_cell_backward_impl.h,sha256=YZJrnHzKvchwNtxFV9fi9m5V4pZkyKYXEyIRXnQoa2s,2750 +torch/include/ATen/ops/_thnn_fused_lstm_cell_backward_impl_compositeexplicitautograd_dispatch.h,sha256=xS_8bA-RVA0KGTO9Lpu1vjW1U_6lzvoo4GzIlHBl7kc,1649 +torch/include/ATen/ops/_thnn_fused_lstm_cell_backward_impl_cuda_dispatch.h,sha256=eWBlJPHyzh_5_onvqDhMiNj1Bz-c1dABLgIk-2QTxvg,1197 +torch/include/ATen/ops/_thnn_fused_lstm_cell_backward_impl_native.h,sha256=tBAOyw72oi8MjdYIRj3gqx87Iy6XvIgY4FGbzs9VrSw,1298 +torch/include/ATen/ops/_thnn_fused_lstm_cell_backward_impl_ops.h,sha256=lhESqcBvYOsV79aRx-L4OKfcpU48E733VEZ-jz5wQ1A,3404 +torch/include/ATen/ops/_thnn_fused_lstm_cell_backward_native.h,sha256=2S0LBDdGQbJtyEnmkfdMei9yoYZjZtP1eVUJePMTOPY,968 +torch/include/ATen/ops/_thnn_fused_lstm_cell_backward_ops.h,sha256=Oa1K65JHI53azCXF8H5-ftwtoeHDrVtMALEKshT25-U,1999 +torch/include/ATen/ops/_thnn_fused_lstm_cell_compositeexplicitautograd_dispatch.h,sha256=HLDN--t8Pycpt6BoojAvbwYsZvtaoJ3J-ulT_vLh6q0,1635 +torch/include/ATen/ops/_thnn_fused_lstm_cell_cuda_dispatch.h,sha256=Qm_T4EXAMKlq1wwn9B5DFutmMs-i4L0NzqlntnXPGCk,1193 +torch/include/ATen/ops/_thnn_fused_lstm_cell_native.h,sha256=quhv1CVF6ROjcYItTKfh13Iu6YPWthVR4Pylo4S4udw,1284 +torch/include/ATen/ops/_thnn_fused_lstm_cell_ops.h,sha256=W4mz7BrNZlFyxnFI4LdNr7OaB-jKjz1qtR8ahMD7t_I,3352 +torch/include/ATen/ops/_to_copy.h,sha256=tpoNP066ht4v1ErhbFf_qGK0lK5Ad1ipxZr0Q3QGVbw,2693 +torch/include/ATen/ops/_to_copy_compositeexplicitautograd_dispatch.h,sha256=0VzVL9Z4dtt2GLIhhkXobAL60ZpyKD-1-GN_qRV2oGA,1726 +torch/include/ATen/ops/_to_copy_native.h,sha256=x4RJ5pROivBwy4eRFl0c5RQ3PGBuV1RVteadk3oUmas,1435 +torch/include/ATen/ops/_to_copy_ops.h,sha256=G907Zm67PpZzv82bOD9-MLdE-tw-yGjHx2Fa0OB6SKY,2781 +torch/include/ATen/ops/_to_cpu.h,sha256=0Swg1FZ0RY7Kf6ayBwx_cQqjsQJAAGvFR6YY4thXuGE,919 +torch/include/ATen/ops/_to_cpu_compositeimplicitautograd_dispatch.h,sha256=_bCRoXU70kTb2yrOdMj_5WQrq1K_WHf8LztBbzsCtfA,1033 +torch/include/ATen/ops/_to_cpu_native.h,sha256=h1lwJcSbsMEP0MtTWF5AGf2WJUTdjW8JO2V8wu5FUDg,745 +torch/include/ATen/ops/_to_cpu_ops.h,sha256=ozgmX-jT6ggk0lujxiAEe-AzVUg1f4VMgS8y7_Sku7g,1260 +torch/include/ATen/ops/_to_dense.h,sha256=uaEPGLX0IyPdRx7aYFq83qKGWaj64MvdZsr_jpRdnsc,1481 +torch/include/ATen/ops/_to_dense_compositeexplicitautograd_dispatch.h,sha256=_JJBbaQa-a5ercbt_y9DMLvQM-K8FkH6W_4jFluHeHg,1305 +torch/include/ATen/ops/_to_dense_native.h,sha256=srizXVjFy42LtfMK8m2vqOyXZYtOu-rlNAhghTHG5QA,1343 +torch/include/ATen/ops/_to_dense_ops.h,sha256=khbFC-w5ssmlPkXCvsmD6tSZNJenPHTsSgu1UPes5TQ,2323 +torch/include/ATen/ops/_to_sparse.h,sha256=1IAf1F-luOLE5uSH7YaUy1jPIOyd9OSfBEcYj6TNOyU,2179 +torch/include/ATen/ops/_to_sparse_bsc.h,sha256=J6FTBASFvRx4mdOtYH0FF3vuJ80NsAtd-ZkLrCEIARk,1465 +torch/include/ATen/ops/_to_sparse_bsc_compositeexplicitautograd_dispatch.h,sha256=AQNyoGIMeKiRZGkQzmQ7LVPXUFk6X0-3GAb47jcLhao,1278 +torch/include/ATen/ops/_to_sparse_bsc_cpu_dispatch.h,sha256=nROADU_Frn7O_8u1Iu48fmotZeYZVNLm-WXVbY8jYyc,1060 +torch/include/ATen/ops/_to_sparse_bsc_cuda_dispatch.h,sha256=Qn2ZNHW8anH9YgGgWNawz-B7hQrERv5hh6gs8HZVj-8,1062 +torch/include/ATen/ops/_to_sparse_bsc_native.h,sha256=TY4fGH2XwXC0UjBffwntD27_mApnvaYsVihn6XTKKBw,1270 +torch/include/ATen/ops/_to_sparse_bsc_ops.h,sha256=cG5kZny2-Kh9PZWM8zaFIDRODha722Dgy3dR0yM3Ydo,2265 +torch/include/ATen/ops/_to_sparse_bsr.h,sha256=8iYLbWbqd12yNYgr9u-jVO-WXm_6X0yLAwDHNdm8qBo,1465 +torch/include/ATen/ops/_to_sparse_bsr_compositeexplicitautograd_dispatch.h,sha256=90sqsU_SxEzydbt5QaC_V-8hQSSR0KPa-IdyLvVhKjg,1278 +torch/include/ATen/ops/_to_sparse_bsr_cpu_dispatch.h,sha256=6dzkt81WGdIFGyyqj64S7fj4-IcpDk0uoxgkztMdOVw,1060 +torch/include/ATen/ops/_to_sparse_bsr_cuda_dispatch.h,sha256=aKKHNV294ySAQY9Bna5d6udKdAkjnmq9B9I9dy6es2A,1062 +torch/include/ATen/ops/_to_sparse_bsr_native.h,sha256=hQy8QH0uZJqVFr6VJCaKAIeAvLPR2n9P7_yUBQ3FoBs,1270 +torch/include/ATen/ops/_to_sparse_bsr_ops.h,sha256=WuONWRG747jmI6BMOY6-ufe9l04m5L17RaELH1n1tzY,2265 +torch/include/ATen/ops/_to_sparse_compositeexplicitautograd_dispatch.h,sha256=sfJ_wKTZAgsF_lVjYiooHFJUdvHG3q0nlBJGjB_n-wQ,1593 +torch/include/ATen/ops/_to_sparse_cpu_dispatch.h,sha256=cLcYe7nMrOqL_mTXloboizJ9RNs6levVCwhgd5IUWpg,1208 +torch/include/ATen/ops/_to_sparse_csc.h,sha256=Cuopo1zMED5QszPJ-9OZ9uxmVU4EvvNLVpkDKOdk3xU,1353 +torch/include/ATen/ops/_to_sparse_csc_compositeexplicitautograd_dispatch.h,sha256=5dU_F4SqaULfV7jl_AdgMdv1rOTEl6gPkwdsF9BwZdc,1224 +torch/include/ATen/ops/_to_sparse_csc_cpu_dispatch.h,sha256=2PzIQy0Q9EH7dYJe6keoapENxX2D1mEMQi3KWv4fDXM,1033 +torch/include/ATen/ops/_to_sparse_csc_cuda_dispatch.h,sha256=e2uVuonpP3_8CRxlDRijgd1xX3KofGSBBJbUaZpcOXA,1035 +torch/include/ATen/ops/_to_sparse_csc_native.h,sha256=T_2e5QnDuSD2ZOtWQkLo7AjMsjGsplZlfb2j89H7SuI,1162 +torch/include/ATen/ops/_to_sparse_csc_ops.h,sha256=1X4xl2FBIwbOrhk15871vJlqa09_llcfnyOyDdG1eYQ,2087 +torch/include/ATen/ops/_to_sparse_csr.h,sha256=BihHWEBdsjLwa6jCfArqraQF6d0Slu9IqAufcpU9O84,1353 +torch/include/ATen/ops/_to_sparse_csr_compositeexplicitautograd_dispatch.h,sha256=f3QFjUvMYGxxEpc2Z37rA-NsCwJbws2TlhfZ_UMRCdk,1224 +torch/include/ATen/ops/_to_sparse_csr_cpu_dispatch.h,sha256=8QZRGZwGCHjgcvukd2Zi2GlVWDSeMhEde1URTj_lo1I,1033 +torch/include/ATen/ops/_to_sparse_csr_cuda_dispatch.h,sha256=1i8E88LM7yiwPBAi4bxRE-hRm1fyv7XqAeyk-kOwKI0,1035 +torch/include/ATen/ops/_to_sparse_csr_native.h,sha256=od03sYzRJRdZF22Vh9R9s6HTSLdtA_hSeQ-LLR55dbA,1162 +torch/include/ATen/ops/_to_sparse_csr_ops.h,sha256=qHO0iyPyco3ICqifeOe6hakLF_Ek3MqjYN7zTgCiFhg,2087 +torch/include/ATen/ops/_to_sparse_cuda_dispatch.h,sha256=cxiAoklxa_3LEnihP4IG2-1LG2EB35Pq9sKau072sXo,1210 +torch/include/ATen/ops/_to_sparse_native.h,sha256=9nTzl9JW3h_nhq56dwM3ETMIcDsQLRyWeVWHw-2_Ais,1906 +torch/include/ATen/ops/_to_sparse_ops.h,sha256=zRMEiSFRTxDJgF1Fhn9hqBhUv1rCE4tCapCrTQdJfVw,3857 +torch/include/ATen/ops/_to_sparse_semi_structured.h,sha256=wgJmbxJ4oQKHSmT7pvY1j4QUebK8SAKXOHW1nPnqFmo,1009 +torch/include/ATen/ops/_to_sparse_semi_structured_cuda_dispatch.h,sha256=EZrmc-Tq5Oec1TnAyXkxhY2FNqrzLbU9tBsc7e9GImE,1022 +torch/include/ATen/ops/_to_sparse_semi_structured_native.h,sha256=oZItHPByU6fk34Z9nwdiAVSjrb5gYG_KRGK7beK7kMI,776 +torch/include/ATen/ops/_to_sparse_semi_structured_ops.h,sha256=Uxoqo5uvFQqCPts6cMS69VOv-WTIS5wGIb0aAr_uwR4,1359 +torch/include/ATen/ops/_transform_bias_rescale_qkv.h,sha256=6gku0rcEDkC1811k_xCUS3DzOwastelNpTNw5HezcZk,2145 +torch/include/ATen/ops/_transform_bias_rescale_qkv_compositeexplicitautograd_dispatch.h,sha256=hLKILXQ4udHttEAMQgGC3RUqpkuDVbLz14RjqOBa6o0,1415 +torch/include/ATen/ops/_transform_bias_rescale_qkv_cpu_dispatch.h,sha256=ud1g1jadambd9o168X2eksnRRLz8TKYSKjukRrg93cw,1078 +torch/include/ATen/ops/_transform_bias_rescale_qkv_cuda_dispatch.h,sha256=4dCJGHvqgqZi023TzGXtCe8rqmu8pQWyb3ydPbMMj_k,1080 +torch/include/ATen/ops/_transform_bias_rescale_qkv_native.h,sha256=IpsnK4ozG2dIG3cDnOUFKbx7c8biTRAUP29srJ7LiRw,1224 +torch/include/ATen/ops/_transform_bias_rescale_qkv_ops.h,sha256=OYTQsL9dq4ZJVzjDlB7LWHam5oFOR5gqDEIhCklvZtU,2648 +torch/include/ATen/ops/_transformer_encoder_layer_fwd.h,sha256=4kflBQOqZjWzW7qkn7tBHdzEDOubiCmu7DKS37oWXhg,4887 +torch/include/ATen/ops/_transformer_encoder_layer_fwd_compositeexplicitautograd_dispatch.h,sha256=Jp3KMyhoPvItq_BKT7XJ7uzSGIFXXn0YKc89L1kXFyA,2269 +torch/include/ATen/ops/_transformer_encoder_layer_fwd_cpu_dispatch.h,sha256=SK2XOuopetC6pxsn44jftP8esI1uKOPvyfvo7L6ZJfY,1557 +torch/include/ATen/ops/_transformer_encoder_layer_fwd_cuda_dispatch.h,sha256=hKTWBHC2a1SQRPs6abcc1Nj08VQC5yr5GQf2lpK8xRc,1559 +torch/include/ATen/ops/_transformer_encoder_layer_fwd_native.h,sha256=s-q8chcwY-4Dbvr5HDvW-YbXOz5efV01awE3YCi1sdU,1959 +torch/include/ATen/ops/_transformer_encoder_layer_fwd_ops.h,sha256=id1Zw8x670qVh-l-xhdm__VaHiSt6fmpfQBudSLEo64,5491 +torch/include/ATen/ops/_trilinear.h,sha256=srDh0Ggs6TYyhcqmWUirhlIALcokJia28GXtV0CDbdM,2264 +torch/include/ATen/ops/_trilinear_compositeexplicitautograd_dispatch.h,sha256=-fWp8km2cB31JautRCWfBk4ziUZzk33Jyyn9x4wq5a8,1457 +torch/include/ATen/ops/_trilinear_compositeexplicitautogradnonfunctional_dispatch.h,sha256=-eVbZWsDJBrLonyTNiuTLPJyd-cYG2BQ2kh81VbK1RE,1213 +torch/include/ATen/ops/_trilinear_native.h,sha256=hExdKVtIQFbxyhh8WgrWe6knk10FaTPavRW3uJ68Dcs,1144 +torch/include/ATen/ops/_trilinear_ops.h,sha256=wkuix8odQE0_HHk53OGXfa_Orpzc1tFXv3OLLRH1jgQ,2905 +torch/include/ATen/ops/_triton_multi_head_attention.h,sha256=kXIu7XGcUz9MpoW82eZ3IX1Y_RP2o4gXqGu0XehFeJM,2968 +torch/include/ATen/ops/_triton_multi_head_attention_compositeexplicitautograd_dispatch.h,sha256=5WCY6MZVIEno9HAF_UPvHpXCvUb9x3q9oA53LuEYak4,1672 +torch/include/ATen/ops/_triton_multi_head_attention_cuda_dispatch.h,sha256=abR5-NvbQ0DYHmyFShW-O9ur1-g4m7fo3TrF09TNJs8,1253 +torch/include/ATen/ops/_triton_multi_head_attention_native.h,sha256=6bJpLjY1CJ92P22NMwFvWOoRBoT6iDtj-QToKX2steY,1358 +torch/include/ATen/ops/_triton_multi_head_attention_ops.h,sha256=LH_nwn_cYh-nwvRM4l6xS8l_zgS01vwwpGtbaq_x3cU,3589 +torch/include/ATen/ops/_triton_scaled_dot_attention.h,sha256=fWD5zErvCpc4k1wijK8hSwEALPOYZk94jlUJlbFkHoc,1842 +torch/include/ATen/ops/_triton_scaled_dot_attention_compositeexplicitautograd_dispatch.h,sha256=cA-65pJuUIAV5i9ugLzDyYp4dryi2ByMlY4cRQ7VGXg,1287 +torch/include/ATen/ops/_triton_scaled_dot_attention_cuda_dispatch.h,sha256=VIVFOLyBgTKQ_1S1edrFncvLSExGTa9-9QK4jkKelUY,1061 +torch/include/ATen/ops/_triton_scaled_dot_attention_native.h,sha256=CtseK7CfJ3vGzP8rodW6CKcW_1YKrFMGwwNtJKnVoAk,973 +torch/include/ATen/ops/_triton_scaled_dot_attention_ops.h,sha256=LENUJ9C2kXqyQPysqeaAP7LqT_O7fDbCBr4v_R9Sftk,2341 +torch/include/ATen/ops/_unique.h,sha256=rsqLx0mhRbgCX2eYi5x2u4ogvYKmnK0_h0nnFHX0PhA,1829 +torch/include/ATen/ops/_unique2.h,sha256=QwDxpD1oOcskM_fTg3tbuGlkkrZIPvLgo-6UP6XowI0,2187 +torch/include/ATen/ops/_unique2_compositeexplicitautograd_dispatch.h,sha256=ySHQRbtkGisvdRleUXa3F_MI9hX7MhCJgbHuXXJgI_0,1408 +torch/include/ATen/ops/_unique2_cpu_dispatch.h,sha256=hZ-roODay6Lq9rAsKZTs11H4ai7X7OGctTcx06k_r_Q,1083 +torch/include/ATen/ops/_unique2_cuda_dispatch.h,sha256=YivjIGnQa8Nj2Ld1nNYfAVr00yrJCjfPV4yVhxLzmIY,1085 +torch/include/ATen/ops/_unique2_native.h,sha256=WYltVX55P4VrjpPBBXRv9bqnW-AYANdlVM5KxybtjN8,1224 +torch/include/ATen/ops/_unique2_ops.h,sha256=Sb4UohpniqR5WQoyUYSCj3FtMJtmzD7h-FjIipYTDtA,2620 +torch/include/ATen/ops/_unique_compositeexplicitautograd_dispatch.h,sha256=vxE9L9FkrbDtegjsLEdF7CDfNyTNa_8ZBWxggW90US0,1296 +torch/include/ATen/ops/_unique_cpu_dispatch.h,sha256=8YaUqvCqyUP5xGnlUe6Wmao1DIbiXq5SV4LmQBd7y00,1045 +torch/include/ATen/ops/_unique_cuda_dispatch.h,sha256=FP1taNR0AfrBLjgkl-j8NoKGmO4QPlDKpyi21QT9wm8,1047 +torch/include/ATen/ops/_unique_native.h,sha256=79Q54mE5ALgnBI_P6cyfW9hYrLWRKgRk8TDNWeS-1NE,1095 +torch/include/ATen/ops/_unique_ops.h,sha256=GX19M2iokyDKfBigxLPy3mHIwrF4X7sUtH_AOhqQ4uE,2309 +torch/include/ATen/ops/_unpack_dual.h,sha256=brWF2l6JfeVlIMeiYBdEMpu4Dj2f9oMm1zLrklsGVDw,1004 +torch/include/ATen/ops/_unpack_dual_compositeimplicitautograd_dispatch.h,sha256=4kc9r9ThBnwM-H1-_mVIB326iYiQxV6g_4YzjhVQsF4,1064 +torch/include/ATen/ops/_unpack_dual_native.h,sha256=ed7Z3N_dwDUgyufmf69qtur7Ow-_mBUUqBMUqFO9aFc,776 +torch/include/ATen/ops/_unpack_dual_ops.h,sha256=5LTmk3aKWbUGeRPImUbQ6_ZxIfePGU3kZ14kmesYPSY,1385 +torch/include/ATen/ops/_unsafe_index.h,sha256=hmqradtOd00tUFx_L8MQzr4eZLLGTCeblnAECPeiX6U,1017 +torch/include/ATen/ops/_unsafe_index_compositeexplicitautograd_dispatch.h,sha256=8pacFQIDwJ02Td3ocszbl3Zy4D01GjTpEM_gMkZ6SLM,1081 +torch/include/ATen/ops/_unsafe_index_native.h,sha256=Z5b21NnM8Yg6C-QULzb2tE-dEKuBMHZ11BMwF0tCmt4,793 +torch/include/ATen/ops/_unsafe_index_ops.h,sha256=_TBdjzHJw73BVk474H2JhvohhhkpZLxWl65Qtc77Vuc,1431 +torch/include/ATen/ops/_unsafe_index_put.h,sha256=TRH3v7-MlIboiSh-WvlrttMdbFR6TVyZtgwl4kfFGio,1127 +torch/include/ATen/ops/_unsafe_index_put_compositeexplicitautograd_dispatch.h,sha256=PTSv-olEDp985V0lzpffHLUSkCVmBWro9uS8D7IJg8U,1135 +torch/include/ATen/ops/_unsafe_index_put_native.h,sha256=2fECDmIRQAO--I-8QzIqjNy-yvea7vgE_HEDVFpqZOk,847 +torch/include/ATen/ops/_unsafe_index_put_ops.h,sha256=6S5CIlpNtu4DcRIgnw0hNJVZbJ2E3-Tfvke42bMokXY,1575 +torch/include/ATen/ops/_unsafe_masked_index.h,sha256=N8gUW1UuTzmxWboI8wZYHlkbeIACGO-7i-W5aZhH76w,1119 +torch/include/ATen/ops/_unsafe_masked_index_compositeexplicitautograd_dispatch.h,sha256=jMuot6odiiXSaeB20ZHxQeS34Odw3XQZ0VV9lsZ7nlU,1138 +torch/include/ATen/ops/_unsafe_masked_index_native.h,sha256=iWj1FFAc62GwnTFgUdUc-dUndV6UhB3vocdAoSpEEBA,850 +torch/include/ATen/ops/_unsafe_masked_index_ops.h,sha256=zZhROJWH0z5m9fi8DdPrjOj1qB1ZhMFGifP2BDVo7J8,1598 +torch/include/ATen/ops/_unsafe_masked_index_put_accumulate.h,sha256=6-c5FqqWIJarYlkKSdjhoBgxGWDjRVQYOLyh1LidYBU,1185 +torch/include/ATen/ops/_unsafe_masked_index_put_accumulate_compositeexplicitautograd_dispatch.h,sha256=tixFFt6nN3TLQQ-ACd8ac7hseUSCbuBK27X5kZ1WjZ0,1155 +torch/include/ATen/ops/_unsafe_masked_index_put_accumulate_native.h,sha256=4zwlnCAYJ8H97YQoi1FfPr9mbGzdMdkYqp49DIBLKJ4,867 +torch/include/ATen/ops/_unsafe_masked_index_put_accumulate_ops.h,sha256=MjWX6KKvWZGUCIF6k9iznV36TvqOp3n-bfGR6I6--24,1649 +torch/include/ATen/ops/_unsafe_view.h,sha256=iWocRbucWFBOCvCcDyVRFzdpt780gnL0ZE8OMLVToUM,3931 +torch/include/ATen/ops/_unsafe_view_compositeexplicitautograd_dispatch.h,sha256=lvA9HWSfa-DVbSxecuti6sloK6r-_tY8ZUk6LxLr67Y,1587 +torch/include/ATen/ops/_unsafe_view_native.h,sha256=krtBJoLo-CkfQxBi_z_0XcEPtcgPXZ4zUD6xOwGRNnk,875 +torch/include/ATen/ops/_unsafe_view_ops.h,sha256=YIoUMtZGrQVXEgtjbmwUjgl032oEsPCogr3-zyBaSt4,2013 +torch/include/ATen/ops/_upsample_bicubic2d_aa.h,sha256=a505MLcsXUXioGqzgMy7l3APw7YNO_3sHPSo5YKuoxc,8356 +torch/include/ATen/ops/_upsample_bicubic2d_aa_backward.h,sha256=i4pmLCRcgFdIs9ye1EvL_M4yBuOSk4Wyjy9NtzLBV2s,8062 +torch/include/ATen/ops/_upsample_bicubic2d_aa_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=W_9xAAMDhMnpUwjqGAX9kTGebwert8vbZEE9sVCAMfE,1527 +torch/include/ATen/ops/_upsample_bicubic2d_aa_backward_cpu_dispatch.h,sha256=TtDQHCGRM__yxkFoNJh1jBq3JaCQYPi_xNWWbLV5fxk,2597 +torch/include/ATen/ops/_upsample_bicubic2d_aa_backward_cuda_dispatch.h,sha256=UCFFR7P6iSvpMhugfqZJi_NRIqcRPFjPnekrKvP2WDo,2599 +torch/include/ATen/ops/_upsample_bicubic2d_aa_backward_meta.h,sha256=LwU2FU6gP1lN86wNuhWadqMpmT7z3uWVEZHU_3HWeKs,1010 +torch/include/ATen/ops/_upsample_bicubic2d_aa_backward_meta_dispatch.h,sha256=pqUCYptw_A2NGFTe9Vgl0G2XtNqnb5WmkllW1reN7TU,2599 +torch/include/ATen/ops/_upsample_bicubic2d_aa_backward_native.h,sha256=nJg7zl7NRrZfZrMt1hzK1kAX75DRGc3QI-Ojm6FCmyc,1467 +torch/include/ATen/ops/_upsample_bicubic2d_aa_backward_ops.h,sha256=LY0YbBxzz_WUzrRGVEzaruV1RZBmxApBz2uVp83z6pY,3061 +torch/include/ATen/ops/_upsample_bicubic2d_aa_compositeexplicitautogradnonfunctional_dispatch.h,sha256=3S6UhJjgiPhYABBzXEl62OQ3HHeba_mEEsSX2uOcrZM,1435 +torch/include/ATen/ops/_upsample_bicubic2d_aa_compositeimplicitautograd_dispatch.h,sha256=Xup1yTCrYPBYmJxMAeE91mtVJt4e1SeABmrAdrKJ0Z0,1336 +torch/include/ATen/ops/_upsample_bicubic2d_aa_cpu_dispatch.h,sha256=jylbU4Hm4HjVVRwMlZH_qELnUgbrjns52k6l6tmvrgY,2293 +torch/include/ATen/ops/_upsample_bicubic2d_aa_cuda_dispatch.h,sha256=2YO3KWkm4JQ2V0B84HW_7PNAIsP2wWp-XMxZsn4sTOc,2295 +torch/include/ATen/ops/_upsample_bicubic2d_aa_meta.h,sha256=dzJ5dBmumhp4WsR_ofKMaWU3-KGRdB1SoJQvTtfAONY,960 +torch/include/ATen/ops/_upsample_bicubic2d_aa_meta_dispatch.h,sha256=v56I6MFZ0Q_ndz7_kGi-xHEdxKxbdwGY0Mwa6sdd_OQ,2295 +torch/include/ATen/ops/_upsample_bicubic2d_aa_native.h,sha256=cwYKmsxsYHhN_tA5zwi8oyUZLzTajfmWOfke0lABlS4,1507 +torch/include/ATen/ops/_upsample_bicubic2d_aa_ops.h,sha256=yMbSWB2p_4jdtxnSvdnCiP01z5nxhkwJm232VzW8Lag,3637 +torch/include/ATen/ops/_upsample_bilinear2d_aa.h,sha256=YvzGar8UtH6sIUpcloHX_lQb4HsVE2_zFBm_Ko_xbCo,8397 +torch/include/ATen/ops/_upsample_bilinear2d_aa_backward.h,sha256=cZWyZZfbkDSn5PWQ7FbaHWBJ_PJKZIFX7MohHAI844s,8093 +torch/include/ATen/ops/_upsample_bilinear2d_aa_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=rIXZ4EhlE-XK2JHk6OG2z_sLKk71_iKan_I07THrftQ,1529 +torch/include/ATen/ops/_upsample_bilinear2d_aa_backward_cpu_dispatch.h,sha256=Kl2BoO9w2E8WCYbT3UeIZtQDwSK3pOeMI8IuyN0MOCA,2603 +torch/include/ATen/ops/_upsample_bilinear2d_aa_backward_cuda_dispatch.h,sha256=VKS6OSPTdcwuWJy6VkiPFUj38RKUCGfXZ6YScMaoC7M,2605 +torch/include/ATen/ops/_upsample_bilinear2d_aa_backward_meta.h,sha256=Fnw_6m4lpSvYZvGJgvwCCD4IDcR_SDW8uM63n83pjbA,1011 +torch/include/ATen/ops/_upsample_bilinear2d_aa_backward_meta_dispatch.h,sha256=rhk4ZnZncfcz6FSKlbtXAW9CQX0bepl5Cu15xPVnnzc,2605 +torch/include/ATen/ops/_upsample_bilinear2d_aa_backward_native.h,sha256=MXh0vjrR6uzhoA33v4awCddfT-j7N_U3MCJ1iYjGLBY,1472 +torch/include/ATen/ops/_upsample_bilinear2d_aa_backward_ops.h,sha256=qyb7xOIH1D6H_rR_BIbxKPuFyYSkW049vH4JfAfO--A,3067 +torch/include/ATen/ops/_upsample_bilinear2d_aa_compositeexplicitautogradnonfunctional_dispatch.h,sha256=K0Z7mqEwvw-q9bhOIR_pzbtIRy-QNwHhxIPRZElWlvs,1437 +torch/include/ATen/ops/_upsample_bilinear2d_aa_compositeimplicitautograd_dispatch.h,sha256=bKKUnvrWzy8IkleB5q0HzrRLgLaND0Du0ganG-Ew4Qg,1338 +torch/include/ATen/ops/_upsample_bilinear2d_aa_cpu_dispatch.h,sha256=q5nBGrSbHYNHErUc1LpPvC24LVUXpB_yDDh1WWWgjls,2299 +torch/include/ATen/ops/_upsample_bilinear2d_aa_cuda_dispatch.h,sha256=6AkwSQliMB4kXSdDgyevsVU2wENF-M7n37y5vn1W45c,2301 +torch/include/ATen/ops/_upsample_bilinear2d_aa_meta.h,sha256=qWLGU_eX5jO6ohH75JVgF2vLiZ1YoBz3AGuGnTUHI-E,961 +torch/include/ATen/ops/_upsample_bilinear2d_aa_meta_dispatch.h,sha256=bzfgxinFjQy3f8HFlGaFdWQvpNkJe5Jtis2a3folihQ,2301 +torch/include/ATen/ops/_upsample_bilinear2d_aa_native.h,sha256=n9Z_xgi8LzuWD4ssW1ImaOjWN0XTrjs7kXfdSEtC4ns,1513 +torch/include/ATen/ops/_upsample_bilinear2d_aa_ops.h,sha256=0Olc6R8uKhj1PQPfeGTr1mPaEqY73oddDf9r3q0OK2A,3646 +torch/include/ATen/ops/_upsample_lanczos2d_aa.h,sha256=EtdVT90xRaFZP1jhUyVlFT7vusQpg-1Cs-Y8r-e68EE,8356 +torch/include/ATen/ops/_upsample_lanczos2d_aa_backward.h,sha256=vwzVw6R4aCo4OROzlHroocM6PNso6TtsjK96boDgVKA,8062 +torch/include/ATen/ops/_upsample_lanczos2d_aa_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Blwhobd3QxlLyZzspmNWVUi9IOn6bi6ipw0yWsTED64,1527 +torch/include/ATen/ops/_upsample_lanczos2d_aa_backward_cpu_dispatch.h,sha256=yuyU1PdjzryOh3gbQdZAHGe8na0F8CfytE9vXK3FdkE,2597 +torch/include/ATen/ops/_upsample_lanczos2d_aa_backward_meta.h,sha256=fSDUXMUXoBvE3OWByfeo-wbiPBPmBYVu5kmOu7NIHzc,1010 +torch/include/ATen/ops/_upsample_lanczos2d_aa_backward_meta_dispatch.h,sha256=Ri1W55ihmwmKCQXgGN-m45-UdLnVFeaphcH235zh7Oo,2599 +torch/include/ATen/ops/_upsample_lanczos2d_aa_backward_native.h,sha256=ZKra-SpzkJ8FoDBeuAHirywk5suoESb8ijlDSMYaNc4,1100 +torch/include/ATen/ops/_upsample_lanczos2d_aa_backward_ops.h,sha256=kcltuThUhEodDY6BdPRKSf3gBEprI5wNDMb0b9Jxky0,3061 +torch/include/ATen/ops/_upsample_lanczos2d_aa_compositeexplicitautogradnonfunctional_dispatch.h,sha256=azfrP5YrON77tfmAA5uQkwz1nRDmKqdJyynXRdJeVfg,1435 +torch/include/ATen/ops/_upsample_lanczos2d_aa_compositeimplicitautograd_dispatch.h,sha256=dnv70_Y4TJW0V_f191LxcxN3t13l_EkLGWplmxmfRYQ,1336 +torch/include/ATen/ops/_upsample_lanczos2d_aa_cpu_dispatch.h,sha256=tHbYWrfOukkAdLZzNdsA0U0rtq5-191l9eWdD1AlC_I,2293 +torch/include/ATen/ops/_upsample_lanczos2d_aa_meta.h,sha256=4xssoiPxfJN_r4MHgnq_SEnptKu8kZgPX5e7xbgot5c,960 +torch/include/ATen/ops/_upsample_lanczos2d_aa_meta_dispatch.h,sha256=mjWNq0MYalKLXvJDha_MvQUDtB54U496VpK2zUkk5oU,2295 +torch/include/ATen/ops/_upsample_lanczos2d_aa_native.h,sha256=2xZmc5fNq4WgjdyhoXpgErSsdXfqLI5N2bJSc2UnU1M,1206 +torch/include/ATen/ops/_upsample_lanczos2d_aa_ops.h,sha256=Wj7o5WgPorlNNQcZVTdYiQGrtNFao8laD6MQJqQFsq8,3637 +torch/include/ATen/ops/_upsample_nearest_exact1d.h,sha256=-VzqLi1f5OqNPNXBfRCuAf9XITS7vhPEfdeZ7aLd98E,6919 +torch/include/ATen/ops/_upsample_nearest_exact1d_backward.h,sha256=Zq8vPZFESYz-2YhSXApOYSfaHwdp6DpBGbWLxHMN-BQ,6775 +torch/include/ATen/ops/_upsample_nearest_exact1d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=JtPQnN0MBhW7jVGVRgfe5z-b1t6UBqf19FFPvGCWkMY,1391 +torch/include/ATen/ops/_upsample_nearest_exact1d_backward_cpu_dispatch.h,sha256=Buxozcm713Vzuo_DLXTmGXnG5OjutSNfQvN6k0-USLI,2219 +torch/include/ATen/ops/_upsample_nearest_exact1d_backward_cuda_dispatch.h,sha256=qR0U3H7WZp4fnRXpPimht7hB_f-OfVvCLLwf7QNXDnU,2221 +torch/include/ATen/ops/_upsample_nearest_exact1d_backward_meta.h,sha256=nHyhPyOPQrpHAT3AAMXyxPI3GCmStmCzyt0-9flq1y0,957 +torch/include/ATen/ops/_upsample_nearest_exact1d_backward_meta_dispatch.h,sha256=sARQXSfllsDLVuTw46DzReAbeAmlobxVZVX0X9LmY30,2221 +torch/include/ATen/ops/_upsample_nearest_exact1d_backward_native.h,sha256=KUxRIl74AANLQcB7b1pdLmsJYMF1apieVw5RMcWadfs,1370 +torch/include/ATen/ops/_upsample_nearest_exact1d_backward_ops.h,sha256=2SpIfpezDzwgkFl1R3B98PIrAxOv2NrIxlEj4Ph3Zmc,2705 +torch/include/ATen/ops/_upsample_nearest_exact1d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=KtdVloCczJfolOClZkGjukZ5YuQkEqJ3LD3NPx757zQ,1299 +torch/include/ATen/ops/_upsample_nearest_exact1d_compositeimplicitautograd_dispatch.h,sha256=hPRQy2ofENs0SetSTeaZl42mhucTQj1KkSVN4MT8_Jo,1302 +torch/include/ATen/ops/_upsample_nearest_exact1d_cpu_dispatch.h,sha256=TtTBUGPwY7AV0cEe7uVJs9QHf0NTEDpbjo8nZ13HKvU,1915 +torch/include/ATen/ops/_upsample_nearest_exact1d_cuda_dispatch.h,sha256=qi43BUOSbTNMhxMbq8F_DrfvdtI8ZH4XLyWNI8k8A8U,1917 +torch/include/ATen/ops/_upsample_nearest_exact1d_meta.h,sha256=7HSjwa0I-NQB936fGElGvp-RIWmus3QHSZP3WoktRIE,907 +torch/include/ATen/ops/_upsample_nearest_exact1d_meta_dispatch.h,sha256=layj9A_7u3EDqLaMhJJ-1cP2VhppbzJc0ckjUSGG-18,1917 +torch/include/ATen/ops/_upsample_nearest_exact1d_native.h,sha256=lazql9Sdgxqv21aKEv-2IHrNEubItERKEIDuTfVzjcQ,1393 +torch/include/ATen/ops/_upsample_nearest_exact1d_ops.h,sha256=GKrlFdXOLwJkKJ9U9-wtV5ZCkkwJYy0pLhzy28IXB9Q,3224 +torch/include/ATen/ops/_upsample_nearest_exact2d.h,sha256=er7MXi7yI46xSpkIflPtVPGrKczXdwgn9oM_ImftD6E,7759 +torch/include/ATen/ops/_upsample_nearest_exact2d_backward.h,sha256=jxZdHt20lbOtnX28tVHGaSs1lGK8tEa5snhilG7tNIo,7615 +torch/include/ATen/ops/_upsample_nearest_exact2d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=wfxLzb2DI9s3lNrP29rUltdVjqliI4NRxBDJgj-a7Lc,1493 +torch/include/ATen/ops/_upsample_nearest_exact2d_backward_cpu_dispatch.h,sha256=JvQgfoNa1J9_bK7AWSOFWDYKf5l4Hx1Y0uuAN8pGlIM,2495 +torch/include/ATen/ops/_upsample_nearest_exact2d_backward_cuda_dispatch.h,sha256=VPvkSeiBjHlQk3437ZQnVPUX-kjAm9HDYAHVXyrleR4,2497 +torch/include/ATen/ops/_upsample_nearest_exact2d_backward_meta.h,sha256=zDN18dX0bcQ1-NV0zSCao4VWoDwd4PXWFktT-OO9390,993 +torch/include/ATen/ops/_upsample_nearest_exact2d_backward_meta_dispatch.h,sha256=STtiQQ6Zg8fbszcGtR9AnLtQsEri0qsDSwtBG7Uww_Y,2497 +torch/include/ATen/ops/_upsample_nearest_exact2d_backward_native.h,sha256=NAWvamyBca1bbaz6iQSZ2C_Ho5_hxaLI1d7xL8nwdGA,1442 +torch/include/ATen/ops/_upsample_nearest_exact2d_backward_ops.h,sha256=Fj1buERxJHULu7Ljz3qJ56rRHmTiEuNAS2ttF3YftJo,2947 +torch/include/ATen/ops/_upsample_nearest_exact2d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=55uGnP6Vas5Z94bMqZocP4Wddx2nKALWJz8SeB7_Ffg,1401 +torch/include/ATen/ops/_upsample_nearest_exact2d_compositeimplicitautograd_dispatch.h,sha256=_8ZULi7SYU1W3OIgdwoDaZi_qMUegD3Tbmp2XzoB928,1302 +torch/include/ATen/ops/_upsample_nearest_exact2d_cpu_dispatch.h,sha256=NJMT6BPIlENl8qnpiACaEhRyfPy0M8tELzG0U0rlU7M,2191 +torch/include/ATen/ops/_upsample_nearest_exact2d_cuda_dispatch.h,sha256=RZ38hgBJK0z7wOYxBNarekhjhOtw8jejV-BHsNY-BRw,2193 +torch/include/ATen/ops/_upsample_nearest_exact2d_meta.h,sha256=wOy0A87SiFip_E-IU5svRcy1-jJRAiJSNOm0ClCnPX4,943 +torch/include/ATen/ops/_upsample_nearest_exact2d_meta_dispatch.h,sha256=a82hPBXPZkFGzC2AO_lg8pyX63aw7_NtgNYOha90ElY,2193 +torch/include/ATen/ops/_upsample_nearest_exact2d_native.h,sha256=ztbEqaJR1Un3378p5ZrAC6RdrqD80DnBX4mm489hnZw,1679 +torch/include/ATen/ops/_upsample_nearest_exact2d_ops.h,sha256=9_ETL_S3FxEx2z0f_j6YASRF11_S-8RVXrRovn3fr7g,3466 +torch/include/ATen/ops/_upsample_nearest_exact3d.h,sha256=m66uZfhFrU3Q68rrx9i0IruJxF8MTD6XLsGbTNVqaN0,8539 +torch/include/ATen/ops/_upsample_nearest_exact3d_backward.h,sha256=gGAEvwTsKbDJJbfZN1L0N_hRKjRyInUMOtd917ewm0w,8395 +torch/include/ATen/ops/_upsample_nearest_exact3d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Dj45oW2SKqkOxcJifeL6KGdICQcfrKEVHyHlrmAvMp8,1591 +torch/include/ATen/ops/_upsample_nearest_exact3d_backward_cpu_dispatch.h,sha256=FeOrO-KUXod0-XZTDYL-FgLvIWcETT_fD6eB5oL0ygc,2759 +torch/include/ATen/ops/_upsample_nearest_exact3d_backward_cuda_dispatch.h,sha256=BxgfmFKa-sW4-NGVGyFWk6OzIkRrvo8pulnJlv_nBDk,2761 +torch/include/ATen/ops/_upsample_nearest_exact3d_backward_meta.h,sha256=tOsyGMLkR7kQ-txLR6A8uLGNv79_FXPTsmgD-uRC0iM,1027 +torch/include/ATen/ops/_upsample_nearest_exact3d_backward_meta_dispatch.h,sha256=WfTcn_mx8KFRPJHqlGvZRRse3Q7rL1-v7D_P36ePOgE,2761 +torch/include/ATen/ops/_upsample_nearest_exact3d_backward_native.h,sha256=7WPaxijaJ-pw_aeSoTvTtUeY1NF8p0edcCjAeOW_CFo,1510 +torch/include/ATen/ops/_upsample_nearest_exact3d_backward_ops.h,sha256=V7irNu74WJGlj0iSCwZXaiQJW6ts0UrAkAn-SBjjIM0,3177 +torch/include/ATen/ops/_upsample_nearest_exact3d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ukjik3qvwnmcUQuR8P7ppo3b_m9BQVENFbcQmsWE6lY,1499 +torch/include/ATen/ops/_upsample_nearest_exact3d_compositeimplicitautograd_dispatch.h,sha256=Wd79kYkAQsrl4jAvkPAq2itfTImnu2KllGfhCqXVwkU,1302 +torch/include/ATen/ops/_upsample_nearest_exact3d_cpu_dispatch.h,sha256=9d89SQwDQqPlRrDT9qN90GoN2T8o0_j6CSDYt_OxfFs,2455 +torch/include/ATen/ops/_upsample_nearest_exact3d_cuda_dispatch.h,sha256=Mnkl-UneGBFuU3AdDmq4AcP6mz12zKO7eTlPGtyLFoI,2457 +torch/include/ATen/ops/_upsample_nearest_exact3d_meta.h,sha256=qfVa4ROO2O5uOalXdUaT7JzJCQ2_DtykNhTH08761SQ,977 +torch/include/ATen/ops/_upsample_nearest_exact3d_meta_dispatch.h,sha256=H4x4ehY2jkrfO4jy__o7ykqrLVqLFaD0Az3dh9KGtHg,2457 +torch/include/ATen/ops/_upsample_nearest_exact3d_native.h,sha256=M34L2Ck0Hd_otuL9gGPwHUbmde57WSn2dYUWYH80R88,1796 +torch/include/ATen/ops/_upsample_nearest_exact3d_ops.h,sha256=ymyECr6drkzmluGFbjcpR0z5ncBwAkkJBr9rYq1UdWQ,3696 +torch/include/ATen/ops/_use_cudnn_ctc_loss.h,sha256=Q6dtRzf423FRHCbPF8Ys8oHqMOAqFA-sam0N-F3sfbA,1599 +torch/include/ATen/ops/_use_cudnn_ctc_loss_cuda_dispatch.h,sha256=DTkZf4JozwSePZGvbikJ4a8Mv3iv-aQMJW4A5h0DAAk,1272 +torch/include/ATen/ops/_use_cudnn_ctc_loss_native.h,sha256=a_UfLMuo9166t4hvHmvqkejgWnEmui2gK_ysLKpC5FQ,1033 +torch/include/ATen/ops/_use_cudnn_ctc_loss_ops.h,sha256=xAg02IGRYu-Y1u7keUOjEcNWi3gfhJMwvBY1LSqAZos,2509 +torch/include/ATen/ops/_use_cudnn_rnn_flatten_weight.h,sha256=KMjHrpDN8_LM_RniE2Gab1A-vBRa0Pp83X_KH_SKzrU,937 +torch/include/ATen/ops/_use_cudnn_rnn_flatten_weight_compositeimplicitautograd_dispatch.h,sha256=hlgWIOKVPFJ-wFRJ1gDJovQu8nIKlresHJ0jgHMxXhQ,1012 +torch/include/ATen/ops/_use_cudnn_rnn_flatten_weight_native.h,sha256=wxTQ21j3wTM5xQBgfydOXm8C0MLwH5uIALNtKH8W2E4,724 +torch/include/ATen/ops/_use_cudnn_rnn_flatten_weight_ops.h,sha256=xw7TKU8qsfBLCqOXIvNRwtZ-6twudxr48W1dTkO8X4s,1183 +torch/include/ATen/ops/_use_miopen_ctc_loss.h,sha256=ch2oyRNMzvsKwkRvsh93F2uc1sqycsribmR6dnK7d5U,1606 +torch/include/ATen/ops/_use_miopen_ctc_loss_cuda_dispatch.h,sha256=vLFjbLKRx_FfSr6xUQ6y1MRAU2ygHiQbnYxdlO1gFmU,1274 +torch/include/ATen/ops/_use_miopen_ctc_loss_native.h,sha256=12GjqucV_vjKINQ-woJcj_Hvu0tCsVb5XjwQhMnMayI,1035 +torch/include/ATen/ops/_use_miopen_ctc_loss_ops.h,sha256=avNAoNEX9JoZRC-1iIddUv3xnyBacpVTGgn4F9eR08A,2515 +torch/include/ATen/ops/_validate_compressed_sparse_indices.h,sha256=TwKTJlw2cRBvVIIroFatmHZDMfHUMUp9mSAbea6XjbM,1207 +torch/include/ATen/ops/_validate_compressed_sparse_indices_cpu_dispatch.h,sha256=p1WvNhq856wg39g-58tQwA8HHYbAblibVLYAbrLSZp0,1091 +torch/include/ATen/ops/_validate_compressed_sparse_indices_cuda_dispatch.h,sha256=BSNli22tTDuxG3itkZcpGXQsc41kBP1GvoDSyDyAdUg,1093 +torch/include/ATen/ops/_validate_compressed_sparse_indices_native.h,sha256=glqBe3VCHW3vTd2bUCH3MkkQ3sWnFG62F2tsWpZtZmk,1027 +torch/include/ATen/ops/_validate_compressed_sparse_indices_ops.h,sha256=Zf_fMspJD9BqD8LMV_AEuOwJ2-eAQoUR_jhSf2AxH_Q,1587 +torch/include/ATen/ops/_validate_sparse_bsc_tensor_args.h,sha256=ixNPc5dm8_i606gp5fbUIo5W29jBmXSaCwYVTPl6crw,1257 +torch/include/ATen/ops/_validate_sparse_bsc_tensor_args_compositeimplicitautograd_dispatch.h,sha256=lcsWlj5pzZdSoVh9dkh8MH_NIA986axGqPhLq-vpvWw,1179 +torch/include/ATen/ops/_validate_sparse_bsc_tensor_args_native.h,sha256=zWxbahDHV3c5zRGuM7FVHn5gSYmi2uxiwkd0hL0vt7Y,891 +torch/include/ATen/ops/_validate_sparse_bsc_tensor_args_ops.h,sha256=kgwjjPBD9jTiudWr0GONJbkV32YC2_SiiCBPsotmm0I,1680 +torch/include/ATen/ops/_validate_sparse_bsr_tensor_args.h,sha256=rXkSNbt5k2zbytjH_lRPx7TNspy-8IjedrY3ep-HVJM,1257 +torch/include/ATen/ops/_validate_sparse_bsr_tensor_args_compositeimplicitautograd_dispatch.h,sha256=xD8GBwPs5j4_Cs-rVtEjIeWdrehYRMSJlNNLe1nOGbc,1179 +torch/include/ATen/ops/_validate_sparse_bsr_tensor_args_native.h,sha256=_88gen0XPNtZEEWEJcNuUWyLVJLPa4Y1tNbz01p5njo,891 +torch/include/ATen/ops/_validate_sparse_bsr_tensor_args_ops.h,sha256=YVuveutg3GJ6mZZLQ6xMQwD95SIOPAAHMPcCc3I8FnY,1680 +torch/include/ATen/ops/_validate_sparse_compressed_tensor_args.h,sha256=7YGke7COaJqCpj8Qk8S8mK7ka1BaLxe5IQgNIb7HCN0,1351 +torch/include/ATen/ops/_validate_sparse_compressed_tensor_args_compositeimplicitautograd_dispatch.h,sha256=EVH8XCKguYPWBXgX_gaOp5ocDYRs57nAZwJL8bVAsS8,1213 +torch/include/ATen/ops/_validate_sparse_compressed_tensor_args_native.h,sha256=B4pvQ7iPzNOZsEH_N2sP0eygQtFbYcLOcGHrnPnnEaQ,925 +torch/include/ATen/ops/_validate_sparse_compressed_tensor_args_ops.h,sha256=fWlNj3YlEEBZXM3LVJ72Bi2_LAlgqeAO8Flp7tiM5mc,1790 +torch/include/ATen/ops/_validate_sparse_coo_tensor_args.h,sha256=WBRIzoKsgrWjaQTDW6NyTQrAnytE9d45A4jRHZ7iXQI,1267 +torch/include/ATen/ops/_validate_sparse_coo_tensor_args_compositeimplicitautograd_dispatch.h,sha256=3NEmCQCZ683DROT3XMRaI_8sRWgUDHFQEll02M6U7y0,1193 +torch/include/ATen/ops/_validate_sparse_coo_tensor_args_native.h,sha256=q_hWs9fNeUoMFyz4OaTWbjiA5iIOfXdaf4L6BJUwT64,905 +torch/include/ATen/ops/_validate_sparse_coo_tensor_args_ops.h,sha256=oA_FWScGvbW334yN9b83bWmXkRk1ZbtzFwBVNyxf9Bc,1681 +torch/include/ATen/ops/_validate_sparse_csc_tensor_args.h,sha256=WgPHJOtv2MP7vY1fvHRBLXkcvnDLNzw-wMIc1FPr5bs,1257 +torch/include/ATen/ops/_validate_sparse_csc_tensor_args_compositeimplicitautograd_dispatch.h,sha256=VBjHhDxZM4B6DYRiCRIoeMVLWGD1-Eh4pUp1v4jOsE4,1179 +torch/include/ATen/ops/_validate_sparse_csc_tensor_args_native.h,sha256=Srfpvfb_RG1a93X9VAwcHIXW3ooV16eIU1kkwrmjgFU,891 +torch/include/ATen/ops/_validate_sparse_csc_tensor_args_ops.h,sha256=3zxD2w3kY9usfrrNeNbEsRiIV8Ady8I9gkULhYZmzTM,1680 +torch/include/ATen/ops/_validate_sparse_csr_tensor_args.h,sha256=VnlhJ_bOBPeUMxkyT-mVmQSmVAkVJ97EJCmdg3vslTI,1257 +torch/include/ATen/ops/_validate_sparse_csr_tensor_args_compositeimplicitautograd_dispatch.h,sha256=6RoNg5U44a6AYbBq7QTfNC927ES9AII-SN5nwYv_TVk,1179 +torch/include/ATen/ops/_validate_sparse_csr_tensor_args_native.h,sha256=6kRnFWMOs3eJ0yJDIk5nlV_deUIKo9F-0lkQwjZWlyU,891 +torch/include/ATen/ops/_validate_sparse_csr_tensor_args_ops.h,sha256=-whUikyf2-mb0uyqAiEEkNSj1cmmVqGhjUpjzdCXDjQ,1680 +torch/include/ATen/ops/_values.h,sha256=Dg_QWAMq-LLw-BhgaT0IJp42KzKqfOD81eeqesFgXus,758 +torch/include/ATen/ops/_values_copy.h,sha256=wbSEhync8x7_GE35tOco5AgT9W9FCgdjIiPZyZNDBuw,1341 +torch/include/ATen/ops/_values_copy_compositeexplicitautograd_dispatch.h,sha256=XN9fghPHAE3SITptDJ7XTbo676rHuPAO7my_yvFjs2A,1133 +torch/include/ATen/ops/_values_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=_xVpqjkTVrMuk1MbbVj1mPbtO2B5If_1Ez6OFje2ZG0,1050 +torch/include/ATen/ops/_values_copy_native.h,sha256=_8AdpssN8ZEmmVoSg3yrVehBMp39YNT539UPryAXlOg,820 +torch/include/ATen/ops/_values_copy_ops.h,sha256=C4A9ULsTUU2kTyMCF42pIvHszsVU0_Wwt2QdV_JBSjA,1837 +torch/include/ATen/ops/_values_native.h,sha256=dKcw-lt74N_w3u5aLwIgacV_LBzklGQY-GZpHJCQmX4,738 +torch/include/ATen/ops/_values_ops.h,sha256=NMvnlMuOvESmPZgDq_csK17NdxQzvlaIns4BJr9NsY4,1220 +torch/include/ATen/ops/_version.h,sha256=LNcq-X-EzC_6_HcNV7LvEz0t8Egc7os0P1J6ZUkwsGA,759 +torch/include/ATen/ops/_version_compositeimplicitautograd_dispatch.h,sha256=3YwoJcapKyeRV4v4Z7QFIig-jjWT6OsPREw2Zi5pQCc,1017 +torch/include/ATen/ops/_version_native.h,sha256=AX6eWzsaxd6erMgw9UbE6SrzjuvUGXMmxYu8A8dbRzU,729 +torch/include/ATen/ops/_version_ops.h,sha256=tQVkbnf_X2Z3_CxZcTOEw1b-tk-avcZDSYt5FyCDfw4,1205 +torch/include/ATen/ops/_weight_int4pack_mm.h,sha256=d_A5j_NLmHnTjOdrms94oL8u0ZmL7jCu_YnsOjC-TA8,1109 +torch/include/ATen/ops/_weight_int4pack_mm_cuda_dispatch.h,sha256=rq1R-KrOuyzeYOgNK785Up7aocrOoD2XN3bJtzgki2c,1069 +torch/include/ATen/ops/_weight_int4pack_mm_for_cpu.h,sha256=mpPE3zMiTIFstbHvw5CSTf3LcEQrycU4CiyoDeDqhrM,1141 +torch/include/ATen/ops/_weight_int4pack_mm_for_cpu_cpu_dispatch.h,sha256=bPXPjUFsE-0wQjKNlVj7bCDbV6rCSqM2QhKo3H5BVvA,1075 +torch/include/ATen/ops/_weight_int4pack_mm_for_cpu_native.h,sha256=DspzQu0zLyAFgMSlI25j86NPkvHlbw6pnQkZ5u8FdiU,827 +torch/include/ATen/ops/_weight_int4pack_mm_for_cpu_ops.h,sha256=zk_5TdC4rAjBtS6w8pAJBcX8Py_6gppYg1dbVtHZjIc,1535 +torch/include/ATen/ops/_weight_int4pack_mm_native.h,sha256=zzeyLnTm-SVLwjo76XijeXR6p0Kos6tLtudEKLt2ZrA,828 +torch/include/ATen/ops/_weight_int4pack_mm_ops.h,sha256=9gMEf6DWByJoj2Ig2k8bYDgeXeX9MabhviaTykwwGHY,1511 +torch/include/ATen/ops/_weight_int4pack_mm_with_scales_and_zeros.h,sha256=J6Zbv4lrP2njY1W7JrVFAkw9Lc8ihp720tjgkUc7wC8,1223 +torch/include/ATen/ops/_weight_int4pack_mm_with_scales_and_zeros_native.h,sha256=0SAoE4aPf5nZIVqujS9jFASmof4GFAJst2Lcyq_eobA,676 +torch/include/ATen/ops/_weight_int4pack_mm_with_scales_and_zeros_ops.h,sha256=ZM2T5f15r7-ZOq090mOZBErf2FvBmym0qEv9w0h-Y_U,1642 +torch/include/ATen/ops/_weight_int8pack_mm.h,sha256=ANwF-STyZXSbLgt0vmXCBGkjb7aIUHtryjP7JK7Adjg,1037 +torch/include/ATen/ops/_weight_int8pack_mm_cpu_dispatch.h,sha256=faxXveYKBUFxP8qpcOPRaIneL_gWFgDa1z9rae5FAyk,1039 +torch/include/ATen/ops/_weight_int8pack_mm_cuda_dispatch.h,sha256=LeE0AoQYkbCt1JNp5Xf3FupaDETMQ7l8Ka_9LArA-BI,1041 +torch/include/ATen/ops/_weight_int8pack_mm_native.h,sha256=Mb0vQ-eedFjP67QVoXtXHhTm3ZFj0bSlqOfgNxNLuG8,923 +torch/include/ATen/ops/_weight_int8pack_mm_ops.h,sha256=txoGYIYKRfBf2DNsW2hdJqx6JqcMlzNjh35kat3hf2o,1422 +torch/include/ATen/ops/_weight_norm.h,sha256=53Igt22qD3j3vJeRsVrCpxiIZlFbperaR7czMYPssMs,972 +torch/include/ATen/ops/_weight_norm_compositeimplicitautograd_dispatch.h,sha256=EVdypz618yzvtYHohBCqJE2OlxVPuGSRyrJ-nsTrS6g,1058 +torch/include/ATen/ops/_weight_norm_differentiable_backward.h,sha256=a9yqCBnDttIWMYdaIBPWnwpywN3s6mUbWvGOT2NEryk,1250 +torch/include/ATen/ops/_weight_norm_differentiable_backward_compositeimplicitautograd_dispatch.h,sha256=BiIwoIpFNwoeHMOJPEzkLR1ri5MCg9eNI7jAq1jfzX4,1176 +torch/include/ATen/ops/_weight_norm_differentiable_backward_native.h,sha256=mueSgaJg2vLKmTHMv9FSZFIwH34pOh5YgyGxEQFmEQo,888 +torch/include/ATen/ops/_weight_norm_differentiable_backward_ops.h,sha256=qqGc2wqeuT1u3ftQBdze5Ko07UmJt9JKqULbYF9Hs5o,1724 +torch/include/ATen/ops/_weight_norm_interface.h,sha256=Q-7cHN7MqPwEtti8gWDpYBpHu9g-HiWqUyBmnhHtFhk,1817 +torch/include/ATen/ops/_weight_norm_interface_backward.h,sha256=3w5avTEfb0RWLOa2qstBuIMvi3X3P1zEb501hKaEJr8,2350 +torch/include/ATen/ops/_weight_norm_interface_backward_compositeexplicitautograd_dispatch.h,sha256=d5D5t0o1-dt4GhMzM-jx0gxnck6nbNJjWRpeSSs3Uzg,1471 +torch/include/ATen/ops/_weight_norm_interface_backward_cpu_dispatch.h,sha256=-l5pfXMZyBR-Yi-ipEv5cXkoeMclTRj-x3WGdy2KQuM,1127 +torch/include/ATen/ops/_weight_norm_interface_backward_cuda_dispatch.h,sha256=C3pxxD-DNIeMfMhLAYpG7J4xnSutPHkvM6ec-vUx924,1129 +torch/include/ATen/ops/_weight_norm_interface_backward_native.h,sha256=jaHptF8p7gC9RuuIfOYVoEro6iX-5iJnQYYP8YpWdWo,1330 +torch/include/ATen/ops/_weight_norm_interface_backward_ops.h,sha256=YXsymQ-zpFhpT9xk1hpM2xjumvYoe--TtT7PkmRKbok,2879 +torch/include/ATen/ops/_weight_norm_interface_compositeexplicitautograd_dispatch.h,sha256=x-gifq-hW3-GI-MZUJswEn19VEKcgzVaaK2rOLSUTzg,1313 +torch/include/ATen/ops/_weight_norm_interface_cpu_dispatch.h,sha256=_4qkNcvXOOjgV8kPfndwpkQZsDa_B9UsJpmy8cDB1o4,1049 +torch/include/ATen/ops/_weight_norm_interface_cuda_dispatch.h,sha256=hGm-oFnLcqherEvgsxkp-pjs_PAy-UTri8M07oVALeo,1051 +torch/include/ATen/ops/_weight_norm_interface_native.h,sha256=jWUabmzr8wmyFCdwYst9umGl6KtmQ0oOK_Z3jEagzTo,1094 +torch/include/ATen/ops/_weight_norm_interface_ops.h,sha256=alXGbBsVF2vnMqwIJrPX8D6xrojG3ALVrpdW4haOjuo,2371 +torch/include/ATen/ops/_weight_norm_native.h,sha256=6BGgUKA-xq2m-mMd6ExG4STPlCTk7TpSwOSEWwAxT18,770 +torch/include/ATen/ops/_weight_norm_ops.h,sha256=db1Y0FTDHTNULHav8b3IWpUhQoyYQTJtB4UaDqwIKhg,1340 +torch/include/ATen/ops/_wrapped_linear_prepack.h,sha256=JjNAgMROK2W7Qkr_IGUZQWmaBP-KLXBzaiONj7_ycxs,1160 +torch/include/ATen/ops/_wrapped_linear_prepack_compositeimplicitautograd_dispatch.h,sha256=rpHde5mnq6ksUC6ZRTEyT5uDPDDBpMSjT7_lv2dC_Us,1133 +torch/include/ATen/ops/_wrapped_linear_prepack_native.h,sha256=Ch7JDo7iqRnV0pX-Bzmq1pyLwXUKJ6lmQ1F5mp7h56o,845 +torch/include/ATen/ops/_wrapped_linear_prepack_ops.h,sha256=eBuyNNMhjQwZbFRtYmHmCkGpRDWh2sYe1LWp-UCPFzo,1580 +torch/include/ATen/ops/_wrapped_quantized_linear_prepacked.h,sha256=53a_dA5maOraTx47ZtLHBs-St7Ur3bZokrUSSTLrh3Q,1428 +torch/include/ATen/ops/_wrapped_quantized_linear_prepacked_compositeimplicitautograd_dispatch.h,sha256=gCB1ZGWw48KsWTdoBhjz6DSQKbsyQejH5p4aHuYsuu8,1243 +torch/include/ATen/ops/_wrapped_quantized_linear_prepacked_native.h,sha256=inuj7a2dL8DF3pmNozXomSTy3UQE97mXAFMWidu-se4,955 +torch/include/ATen/ops/_wrapped_quantized_linear_prepacked_ops.h,sha256=NEwrsYRTuJ_41uiIf4u1eor20BKTEp_c1_OZs9I4q6U,1931 +torch/include/ATen/ops/abs.h,sha256=VXFWROeF1b0aR89jjuqNJyWSNyzJnUb5JZAlsqS8oxU,1384 +torch/include/ATen/ops/abs_compositeexplicitautograd_dispatch.h,sha256=u6L9CwtSRZRzdS0_c790fSGbDXtENGMgGsE_rlYnQb4,1063 +torch/include/ATen/ops/abs_cpu_dispatch.h,sha256=RmBoDt2pf_ZcrjtiDibUtue7LckDTWp98829UkSK3ns,1071 +torch/include/ATen/ops/abs_cuda_dispatch.h,sha256=slcECZpytQQvdRamusQOZzoBvjZ56olHX8gqorrQON4,1073 +torch/include/ATen/ops/abs_native.h,sha256=S94QPOwCbKpx8-4VSnKs8OOfitkXbV6mulYerIYthOk,1377 +torch/include/ATen/ops/abs_ops.h,sha256=dpwWlKCtkJjwVFLldthcEVScXEOeZHLyIQtjqpZtJkY,2273 +torch/include/ATen/ops/absolute.h,sha256=q2xeb2HznUk2cnzOUXLXMhALhiEJ8CrDZR3er0pLvxs,1301 +torch/include/ATen/ops/absolute_compositeimplicitautograd_dispatch.h,sha256=gavmbjy7CJUy6Tv0EoL0734v59nz6zQTP76qTc2LudQ,1234 +torch/include/ATen/ops/absolute_native.h,sha256=4BnamnLyV9FNQnkwMfKYCmScXFpaPb4PKBx7zUGVAkE,865 +torch/include/ATen/ops/absolute_ops.h,sha256=tr6kE0CsXtkedRNoX_aGbeFrRSVSNZQjLZIu3pEiGqY,2318 +torch/include/ATen/ops/acos.h,sha256=SKpJhnPpujmVHi1iJ8KgEI5f9SoHiogDcq06Dn_NQYo,1397 +torch/include/ATen/ops/acos_compositeexplicitautogradnonfunctional_dispatch.h,sha256=M-Fc06_5nf9VQr_jpO3OXp-xwLypFGBrSoCLogCDrTs,1091 +torch/include/ATen/ops/acos_cpu_dispatch.h,sha256=JCidOgttCzsb0OsfTZyxeRpPy4wsUZ4crLs8ya4hQ84,1174 +torch/include/ATen/ops/acos_cuda_dispatch.h,sha256=Z7NDPlrOc_7fN0G2W0k0pmsWoHz5hDhqMiSZWS5DsLc,1176 +torch/include/ATen/ops/acos_meta.h,sha256=Fxbx8TWXXStWsWioRtg8As5GSOuLHtrc7WBAka3NPB8,819 +torch/include/ATen/ops/acos_meta_dispatch.h,sha256=A9w0xY-3Fx7VKf4ZoOkiaaUQuQ72f1W7S9v0jRftezI,1176 +torch/include/ATen/ops/acos_native.h,sha256=7VpMi1YjhJ_Terwk7Ozi6rGD15Bpp1kR2R7OUPik8I8,844 +torch/include/ATen/ops/acos_ops.h,sha256=P64P2WrWDRjw_e3qz1PyU7g3B_PUBAjn0kQdADkLf-o,2282 +torch/include/ATen/ops/acosh.h,sha256=EX0TixTnxegkFnOGVG4Mjr5FJBkVpDHeGBuJai9KIXw,1410 +torch/include/ATen/ops/acosh_compositeexplicitautogradnonfunctional_dispatch.h,sha256=zRKBMgPIxXPdf9oeMtxk05VxLI-QohO6Z67VRqk-UMY,1093 +torch/include/ATen/ops/acosh_cpu_dispatch.h,sha256=BSsz3pCVkbC6NQ3PgGYFoYX_2umRiNqAI4M9n1yYStQ,1178 +torch/include/ATen/ops/acosh_cuda_dispatch.h,sha256=HhiwzW9WovfcXxFCFD9wrTiP-0dXlxCigU4I2RJVPlo,1180 +torch/include/ATen/ops/acosh_meta.h,sha256=MnBmRufU_rF0A3deLJM88POAlsF4z_JhKMuDBiYVnek,820 +torch/include/ATen/ops/acosh_meta_dispatch.h,sha256=y1V6AsJD6J6BedlPJWSlkiL2Oin6szguVi95nfrLZAw,1180 +torch/include/ATen/ops/acosh_native.h,sha256=hEINhUyNBvZhIviefJ6HMI2TcxEsd4sRQmu2vLeRLKY,847 +torch/include/ATen/ops/acosh_ops.h,sha256=KcQtwP5odYPQ4jrNejm01e9EiYXmR5UN7upRvqB0zK4,2291 +torch/include/ATen/ops/adaptive_avg_pool1d.h,sha256=cvsqAqTtZ40owD4ztXLtC6ABdzkHkDTNvb53QTSUJNs,1597 +torch/include/ATen/ops/adaptive_avg_pool1d_compositeexplicitautograd_dispatch.h,sha256=wdimpWVjiyPUmC_x7svgZwIJGpCJurjbBOZn5dTyXIM,1205 +torch/include/ATen/ops/adaptive_avg_pool1d_compositeimplicitautograd_dispatch.h,sha256=TMUh3JrcxOCWurevgh_SOltOCLlDATQznB0i3qee6E8,1060 +torch/include/ATen/ops/adaptive_avg_pool1d_native.h,sha256=bVW8pIThIaO_ztDhjqKlt1KWiwljT5g9rp-JxKlzO0M,892 +torch/include/ATen/ops/adaptive_avg_pool1d_ops.h,sha256=TabCm5eD_1CLwISDZW6ZfS2uCvqrQmw1UcmV04Sxqj4,2069 +torch/include/ATen/ops/adaptive_avg_pool2d.h,sha256=vYkvJHL1O-psigyoX4Ir6lJcPGnmsZN6Pr-vmz2OUKs,4364 +torch/include/ATen/ops/adaptive_avg_pool2d_compositeimplicitautograd_dispatch.h,sha256=_iB9nsHDQLelULrwxKstVF3Ymt4ODy2BbnxGdG8VE9A,1167 +torch/include/ATen/ops/adaptive_avg_pool2d_cpu_dispatch.h,sha256=riHjvmY7FcYlICc2G8qKU-sDQxvDk1eZemC9uu0C-Qw,1424 +torch/include/ATen/ops/adaptive_avg_pool2d_cuda_dispatch.h,sha256=FySDy4KR6dR0LWAYqVxj2irJCXdcp-zQiXYl1nNo_ZU,1426 +torch/include/ATen/ops/adaptive_avg_pool2d_native.h,sha256=dmjqNiS0zu2YJZ8R_td-SNgxsezNOL-7ctKmvuCmO_I,1164 +torch/include/ATen/ops/adaptive_avg_pool2d_ops.h,sha256=hONaDBIPqUT0goP3QsLrcJJGydx-xDGuc22NNUQ6SCE,2099 +torch/include/ATen/ops/adaptive_avg_pool3d.h,sha256=vQGmOt4e0xO65j1tdX2KGcOnNh10fg1QgKCqu3LOcR0,4364 +torch/include/ATen/ops/adaptive_avg_pool3d_backward.h,sha256=kGBO5I4IU6XB-W4vxWIRKwlavAZaPH6Nrag7BmdAvqY,1500 +torch/include/ATen/ops/adaptive_avg_pool3d_backward_cpu_dispatch.h,sha256=crqhUCUQ0tf9lsX27HkFtxNMBsKaqCvq_EpzXmqGv-k,1199 +torch/include/ATen/ops/adaptive_avg_pool3d_backward_cuda_dispatch.h,sha256=2NRZ9pY3aRj0-c7PizQweWOJ1eLvR0RH4o23XpEZBGQ,1201 +torch/include/ATen/ops/adaptive_avg_pool3d_backward_native.h,sha256=lPeOAGo5WFgDzkcfTEcuWqUEGviiJdQFKVeE1NEo9hs,963 +torch/include/ATen/ops/adaptive_avg_pool3d_backward_ops.h,sha256=7o-jdsz-2aZOk6k4G-tYJ_BRALs3B71SkLpdiacSvTs,1513 +torch/include/ATen/ops/adaptive_avg_pool3d_compositeimplicitautograd_dispatch.h,sha256=29Uyj16XTh4IwPFrPq2NSBRg8IG9cszYK9yeV67ocjU,1167 +torch/include/ATen/ops/adaptive_avg_pool3d_cpu_dispatch.h,sha256=tECjefOtr6qP2_4V8HlafzZVB4lS6V34BOjp4V9u9N0,1424 +torch/include/ATen/ops/adaptive_avg_pool3d_cuda_dispatch.h,sha256=uaEyHDUA4sYOflg37tVYk9AcJ9KvWLd3chP_oBJfO1A,1426 +torch/include/ATen/ops/adaptive_avg_pool3d_native.h,sha256=4-iFS6x1dsVC9lUN6z7es7Vo0GbBlYSHRAMpbvMN2R4,1166 +torch/include/ATen/ops/adaptive_avg_pool3d_ops.h,sha256=XdGoF4Ed5SCSzDhol6XP1uwdl8KUZZUx4jA6ggXDcqA,2099 +torch/include/ATen/ops/adaptive_max_pool1d.h,sha256=-0nUlAjaMhmDW-8jJpDB6cf9p6vyrhjx58xawtfU43U,1040 +torch/include/ATen/ops/adaptive_max_pool1d_compositeimplicitautograd_dispatch.h,sha256=PIEUSJlaiP2HPbXk1qSqzv8j577spTaGT8Pz7WbWG0Q,1085 +torch/include/ATen/ops/adaptive_max_pool1d_native.h,sha256=v_03nADvw8uSvoQCSX-gw4ltLXjtJZA1uk0yMGyR9mE,797 +torch/include/ATen/ops/adaptive_max_pool1d_ops.h,sha256=FIsp8cU6UcWzRINYPbIFJvfWejjb-9oaPRFrlvW7M0c,1430 +torch/include/ATen/ops/adaptive_max_pool2d.h,sha256=DiqjLsIJP2Sm1WmHm5OGkLz0IecNxWJSMHCTQiIh0y8,1816 +torch/include/ATen/ops/adaptive_max_pool2d_backward.h,sha256=PiQpByD_7uKpDggvX--9GtypF0eHpNOmEC-BEcnnV5o,1925 +torch/include/ATen/ops/adaptive_max_pool2d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=02oRBkplBUc9eEdS5j-h0juQzUNdS_J01mddjEs818g,1126 +torch/include/ATen/ops/adaptive_max_pool2d_backward_cpu_dispatch.h,sha256=tqLtKqYZlr-6p6WjqW46Dt_arVQl_U6zPGwE1zSvZ0g,1391 +torch/include/ATen/ops/adaptive_max_pool2d_backward_cuda_dispatch.h,sha256=3moJIuUczYODEL_Qcps0EB28xisuG0nwPAI5miNEQc4,1393 +torch/include/ATen/ops/adaptive_max_pool2d_backward_meta.h,sha256=3DADcZ9jFtPW7FQ_gcHg1baAi8wdMy7ZwqAwZ3mzEkw,903 +torch/include/ATen/ops/adaptive_max_pool2d_backward_meta_dispatch.h,sha256=tEVGBbPrA0EAWhvd3ffFfsngrmNkCHMBBRwrHzDRrdM,1393 +torch/include/ATen/ops/adaptive_max_pool2d_backward_native.h,sha256=LUJYEPD6xlYduE-N7s0qBWsy1KR00S5dpTKcwPpYcu0,1244 +torch/include/ATen/ops/adaptive_max_pool2d_backward_ops.h,sha256=mSmZWYiPjdV_GdOkfpzMTezz2pLHzPizO3DlxIHCVdI,2367 +torch/include/ATen/ops/adaptive_max_pool2d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=k-1oMt7cse1IXiRnJ9KLPJBlRPigDdjplHljnU4kYqU,1111 +torch/include/ATen/ops/adaptive_max_pool2d_cpu_dispatch.h,sha256=DXINe5KXCxRT-x7fDSshhcNxdrkgjzr4il3-L9Mwdyg,1380 +torch/include/ATen/ops/adaptive_max_pool2d_cuda_dispatch.h,sha256=yBpiXu_emEYAnhpdn9yNO-3Hdme2iJIOIfwtfT8VZdQ,1382 +torch/include/ATen/ops/adaptive_max_pool2d_meta.h,sha256=9JRYrYuvXxUtNN4tFkBDE8zsKV7Nl6KTqdZDEj7K6SM,863 +torch/include/ATen/ops/adaptive_max_pool2d_meta_dispatch.h,sha256=yhk7E9WvMoedpFCOZ4xM0oUDLDsUIHIQn3wv4yucfdo,1382 +torch/include/ATen/ops/adaptive_max_pool2d_native.h,sha256=k5cP57Q1tt01OqRHiAxFt0EaC_JKImrSDeLHo-mtCdQ,1179 +torch/include/ATen/ops/adaptive_max_pool2d_ops.h,sha256=LbM5aT7NC94rDKKsMuRi7NXGs_7aDT2okScHfZfXY6o,2327 +torch/include/ATen/ops/adaptive_max_pool3d.h,sha256=Hoeb545SJ6Uc8yaV11NvBaNUBX377iyuJwlfhFv_7TY,1816 +torch/include/ATen/ops/adaptive_max_pool3d_backward.h,sha256=3h2u6Az-qQ26JJ-2OINMhyWx0vglIIK-pjtJxGx850E,1925 +torch/include/ATen/ops/adaptive_max_pool3d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=kk3oXJyrZ7JmF3G8Y5BxJ7bffPijQNzyMFfsCjeBelY,1126 +torch/include/ATen/ops/adaptive_max_pool3d_backward_cpu_dispatch.h,sha256=UFXGqbwSgFOxRQVX2vFQLNx6ea_DvBEgbYqQ9T1nSus,1391 +torch/include/ATen/ops/adaptive_max_pool3d_backward_cuda_dispatch.h,sha256=ulInN7kRNb-B-HPP2ksr8_cro7zMj_fa-r0tzC3kaxk,1393 +torch/include/ATen/ops/adaptive_max_pool3d_backward_meta.h,sha256=7gBriWeTokAlpa21zWII6FMSD9dHOQIA9TY2B5XI180,903 +torch/include/ATen/ops/adaptive_max_pool3d_backward_meta_dispatch.h,sha256=cTBBNkqu42TwIcacV2Y0BlDg1cVgxSUX-UCCsgt7mK0,1393 +torch/include/ATen/ops/adaptive_max_pool3d_backward_native.h,sha256=irD9zifg2Fdlte7xCUDj-BmhHsejs4EljVeHdaafxWU,1244 +torch/include/ATen/ops/adaptive_max_pool3d_backward_ops.h,sha256=euHKvHjqqo0rLIznV7Em3VyzpBQyKkuSZJh7VYR1MmA,2367 +torch/include/ATen/ops/adaptive_max_pool3d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=k8Qg5NkGA1UWpJOjTcUznOLF8BocpL6OIC1kJEXQqEw,1111 +torch/include/ATen/ops/adaptive_max_pool3d_cpu_dispatch.h,sha256=u71JjIEhqWKqMuMJcS3FLtqL9vrIaAs75BLRbJ2TGd4,1380 +torch/include/ATen/ops/adaptive_max_pool3d_cuda_dispatch.h,sha256=_5eSr9wwtG341WSKVqGf6l4PkTmXUXe8v9ZAYg-T9FE,1382 +torch/include/ATen/ops/adaptive_max_pool3d_meta.h,sha256=q46EVVKttpu6Qyof7xlg3Py6vn7Tp5QOh11oZozy-0A,863 +torch/include/ATen/ops/adaptive_max_pool3d_meta_dispatch.h,sha256=MGRsJ_ions_nZXwhfjkcd6xeMqwVPMa41eLko-rfwYs,1382 +torch/include/ATen/ops/adaptive_max_pool3d_native.h,sha256=ML41yts1pIlfeq7LjmrDj4YIQtZzD6Uv_srUMiuZ3bI,1179 +torch/include/ATen/ops/adaptive_max_pool3d_ops.h,sha256=tGRlbBitSgWSrWl4J--_O3yb-Hv1BFCY6IaGCCsKf2o,2327 +torch/include/ATen/ops/add.h,sha256=m6lrA7B-rSG51Io7oG5_ewnfC1aBLFkCgFU1_4dFE3w,2392 +torch/include/ATen/ops/add_compositeexplicitautograd_dispatch.h,sha256=D03vfZ9tGlCEE4nkGdALemm9hU9Au0nppYI0OAhbIbs,1428 +torch/include/ATen/ops/add_compositeexplicitautogradnonfunctional_dispatch.h,sha256=-75RV2Oijnk1s42Hlx-A3JXJCvw2MW3F294vw9cqo2g,1197 +torch/include/ATen/ops/add_cpu_dispatch.h,sha256=Q9exi82uR6zlensj46mqIWfbZCslgKKm3u52lwys0p8,1384 +torch/include/ATen/ops/add_cuda_dispatch.h,sha256=CghG0_M5za1UyzcdMz0xzQCxfS-eNs2R5i8GRNZnnok,1386 +torch/include/ATen/ops/add_meta.h,sha256=6ezeJtIKFAdP-GOBpTAzxHUGGMbSZyDziwlfUtlkWd0,877 +torch/include/ATen/ops/add_meta_dispatch.h,sha256=XsCV2u0nuzkf-emhWiMxI-FgQf-oenCWcYG4U9sm8GM,1386 +torch/include/ATen/ops/add_native.h,sha256=-3Gfdb9dKvIeQSJhiuoJqZE-Otuxgj6UODoim2rXU0Y,3192 +torch/include/ATen/ops/add_ops.h,sha256=uhVYiPvpY3UOEiy7L5RcWXDTg6iDIdaQ7VDj5p098Ow,4986 +torch/include/ATen/ops/addbmm.h,sha256=YpB0pzkc5Kf1QsFE3WgTzxs0Dw1jwrhAjpFEQC_q1Vo,1877 +torch/include/ATen/ops/addbmm_cpu_dispatch.h,sha256=ii329H4Wgj75g1Mls-ibCF2RYJoK026mPulo4PQsRZo,1614 +torch/include/ATen/ops/addbmm_cuda_dispatch.h,sha256=B3jfKpPKNkEp0Sqy6v82AFcRmM7ZJ70Ko1WC1ClzGPM,1616 +torch/include/ATen/ops/addbmm_meta_dispatch.h,sha256=pgTGN84OfzusKgeTOEPUWUNn1GNd_cie5aYwTZ1U0RU,1082 +torch/include/ATen/ops/addbmm_native.h,sha256=vf3QakRHDJ-bLAi_aTQi8vLear21WvoOMuir5mdRp30,1182 +torch/include/ATen/ops/addbmm_ops.h,sha256=GoyU1hlrfdLEPZNWZ_zuBek8g_OpzGcFEL9aRg9mPTs,3359 +torch/include/ATen/ops/addcdiv.h,sha256=t1U1iGjBm28GpDL8V9A24UtbXEEjl_BGFDNsb8Py4yA,1763 +torch/include/ATen/ops/addcdiv_compositeexplicitautogradnonfunctional_dispatch.h,sha256=0kYJpDH1sM06As6VjjhAJnMHX9SNfDV7ktYkNuIoEPY,1265 +torch/include/ATen/ops/addcdiv_cpu_dispatch.h,sha256=ckKtNft32z9OhHOwC8CbMgcQXz2JVyS7V3cvEaVqgLM,1520 +torch/include/ATen/ops/addcdiv_cuda_dispatch.h,sha256=U_BZgItyHrVbYYqaVaeepNMfRYGadNQ5bsTnuEttLeo,1522 +torch/include/ATen/ops/addcdiv_meta.h,sha256=M1U0yXPKvh6GmvmVZdR66gr3JlZudAnAUPJy3i7v-aQ,904 +torch/include/ATen/ops/addcdiv_meta_dispatch.h,sha256=xob8iHLMv6Byfcd-1QiECDW_QVXf-H4JWX_c4-VRgsQ,1522 +torch/include/ATen/ops/addcdiv_native.h,sha256=ciBlAlh3nJgu9kk-hJCu8tUjO5pu1ttKXtitnBo62Uk,935 +torch/include/ATen/ops/addcdiv_ops.h,sha256=3lU7IxvAddBY5j4U07glO6C1eisnnrHFkR7m3_p0BFU,3131 +torch/include/ATen/ops/addcmul.h,sha256=d6yA0MFFrCTnUQAE2cNdgLiI-tVKSm2uWfT0jssBzk4,1763 +torch/include/ATen/ops/addcmul_compositeexplicitautogradnonfunctional_dispatch.h,sha256=nPohaadyYLB8gTbpFwHWOsldaKht-G4NDuksGQW0qOc,1265 +torch/include/ATen/ops/addcmul_cpu_dispatch.h,sha256=9lbRsf9NnypuUx2ANm4AOdlhWYutQZczs_JckYPK2x4,1520 +torch/include/ATen/ops/addcmul_cuda_dispatch.h,sha256=yoHtIwKQ2X_W0NR5vIkzDDIyJQBiAUgtRoW1_z11_h0,1522 +torch/include/ATen/ops/addcmul_meta.h,sha256=4dkSP08fKvzhfnbDkfHUiD_ALxoeWCkIhDb6SDMKM18,904 +torch/include/ATen/ops/addcmul_meta_dispatch.h,sha256=DxXVOvhZtAvjnSOckDnltKHW-6li068UyqlvsXDgA4E,1522 +torch/include/ATen/ops/addcmul_native.h,sha256=B0l0RAVuKH25aZCMwMs4U5lgdwKvGsNeUbl8-ictIR0,935 +torch/include/ATen/ops/addcmul_ops.h,sha256=_SHuJnwpd2P8EBoYAVGUw1O1H8_WKQiclKOUp_osUFA,3131 +torch/include/ATen/ops/addmm.h,sha256=-o4XEvaGbcMOUVC3fYmNtzRbpuenK_BozOSu3PxuW1U,3120 +torch/include/ATen/ops/addmm_compositeexplicitautogradnonfunctional_dispatch.h,sha256=pTq70OPuEuU2y-2EKy9itDkiamtnIcJXpldE-cSq7Cg,1303 +torch/include/ATen/ops/addmm_cpu_dispatch.h,sha256=XcQ2YnhQ30RBEWx1aBAh41Lb_3SV9j9uhiyayjpXZsk,1594 +torch/include/ATen/ops/addmm_cuda_dispatch.h,sha256=Hwi-qjy0MF__LENK4XIzryX-LY4vmJAWoNtG3A1dX6U,2193 +torch/include/ATen/ops/addmm_meta.h,sha256=yROoxshtWW22uqJlz3fvkp6jg3BBMl8J0OJs68DyrV8,921 +torch/include/ATen/ops/addmm_meta_dispatch.h,sha256=-ZMxQ-JE76mSB7gDQYhIu5Du3m0H0fqAy0jolBsLTj4,1596 +torch/include/ATen/ops/addmm_native.h,sha256=vs-YcOXmXcStBbkisyf_efeJpB1M2xlcsRjA-re5z-0,3283 +torch/include/ATen/ops/addmm_ops.h,sha256=qdE3qyCMnWHrwvO8NzKjjO8pw5erGAvLL26ekb-oZjU,5294 +torch/include/ATen/ops/addmv.h,sha256=3nzVDqeH7GXa-NnDQkx_vHNOs2b4WuYsdjBqmVJQ7Eg,2136 +torch/include/ATen/ops/addmv_compositeexplicitautogradnonfunctional_dispatch.h,sha256=T3ypXKUu0OQprtEXQcG35Bw-iAzOeR29MX3m6g4Yeis,1299 +torch/include/ATen/ops/addmv_cpu_dispatch.h,sha256=VZ35EDvsSpxLG1k6Y8mFtPYc7ISEVPKriZQ7vUwBIC8,1586 +torch/include/ATen/ops/addmv_cuda_dispatch.h,sha256=pqfco4hy60f_7EfoHTJUfPeRG6iVjCkXFc4Evpklxlc,1588 +torch/include/ATen/ops/addmv_meta.h,sha256=cYVBtR_dkXbUJi3dZHqEO8VVBv2j27bOpujSmhcR8bc,919 +torch/include/ATen/ops/addmv_meta_dispatch.h,sha256=uZMhLPU3Vy8NPsDA0n6hwQxKIkp4npvjjA1h9D_pxgk,1588 +torch/include/ATen/ops/addmv_native.h,sha256=okfZBptq9X2rQEZKhgoXF3mcoWUq0qwQ3ERStbCa3Ag,1586 +torch/include/ATen/ops/addmv_ops.h,sha256=an0RD3s7ZFVaeySL_Ec0l-ihYH2yYYe6XtjH1rz3pSc,3296 +torch/include/ATen/ops/addr.h,sha256=dEn7LHRNJGIVldThT92l4K2TAxKn6vVBkS1_3qN6KL8,1821 +torch/include/ATen/ops/addr_compositeexplicitautograd_dispatch.h,sha256=P0K5QJPwOkuCZAUsKQXCUXEWStMRUu0AoDQYVKjalUg,1634 +torch/include/ATen/ops/addr_cpu_dispatch.h,sha256=7HaSSdqfzjDzt_T3xHvNPTSbnBRjl6GKl-_YBrOV3Io,1436 +torch/include/ATen/ops/addr_cuda_dispatch.h,sha256=PFOzyufX3Vavp9-R49R1d_Z99qP_0lBQOm4z8cd2Yek,1438 +torch/include/ATen/ops/addr_native.h,sha256=YqPn6BzOANF1u-djnS2VITqI_oAtUUhMqbN3idzP31I,1508 +torch/include/ATen/ops/addr_ops.h,sha256=QjiEGQos_q_JBj9ZnX70RfcwqqFen_BaCnKPUDL0gaM,3305 +torch/include/ATen/ops/adjoint.h,sha256=F4s5cN51xoXVAZySno7UItx-OTulfrZJIghCx_mZ97Y,901 +torch/include/ATen/ops/adjoint_compositeimplicitautograd_dispatch.h,sha256=vOv3srLAcfcrhWeXgTJ6fXg-kxIPqf7ZgbWejFdS034,1019 +torch/include/ATen/ops/adjoint_native.h,sha256=iTV4PJEUIZ8tECa8n9uownsHvZv_ncuFC8Kbp4j-smc,731 +torch/include/ATen/ops/adjoint_ops.h,sha256=iECvks7drZWCFkU7gMQUkXgZYeimSaCn_rST02Hkh9o,1220 +torch/include/ATen/ops/affine_grid_generator.h,sha256=8rQXkwWs772jDr86bPgLP5KFHabnJ17K10PEsomFOMA,4780 +torch/include/ATen/ops/affine_grid_generator_backward.h,sha256=MQV2eisa_JgMsrrEiwjB-L30h1ae0hYFUajcScZsnDw,2071 +torch/include/ATen/ops/affine_grid_generator_backward_compositeimplicitautograd_dispatch.h,sha256=4aS922-zBlehaRkf2-6AWrXOUz-0Vp9RHOgoF8bMfCE,1215 +torch/include/ATen/ops/affine_grid_generator_backward_native.h,sha256=C1VkRkqkdbms7sTqtq-7P-e1YhxPJZdHDgNITz6nzZg,796 +torch/include/ATen/ops/affine_grid_generator_backward_ops.h,sha256=3yvW_Zy50JcrWZggc7FuuVZUvTPU5Kekv0DRj5Qbk7c,1437 +torch/include/ATen/ops/affine_grid_generator_compositeexplicitautograd_dispatch.h,sha256=TpVik6hkMQWZKeLI5K8gvFXdBUfI9b4rZNUmKB95iUg,1767 +torch/include/ATen/ops/affine_grid_generator_native.h,sha256=gE5EaMMPudhWonJpB3E56HWW9tGqsWixqfwqs98G0aw,935 +torch/include/ATen/ops/affine_grid_generator_ops.h,sha256=WoGEgeJqPsnjgBfU6maVP-JF0ZfECVVkbPVGtphQ8kY,2205 +torch/include/ATen/ops/alias.h,sha256=C3LMBrHP_HTyd1AOfdjTV-1scg7ofP821Id2Kg3Jeyw,893 +torch/include/ATen/ops/alias_compositeexplicitautograd_dispatch.h,sha256=nGzzFW756BOmAxwNfYDAyVWU8x5o-E10KRFVQeqscGc,1017 +torch/include/ATen/ops/alias_copy.h,sha256=5vy0sl19ExUPPGiUBOEPK00khwpGrWFpVZdAh3O1AS0,1321 +torch/include/ATen/ops/alias_copy_compositeexplicitautograd_dispatch.h,sha256=cri_CZ9jyxStuFPaEOx0uRvIfg941j_QZrwWG_UNrA8,1129 +torch/include/ATen/ops/alias_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=9psWzdjcAxS9J7YS7Ovo3XwRqG9_r1QDIm4ZHwHMcbM,1048 +torch/include/ATen/ops/alias_copy_native.h,sha256=OdmggpAqopsw1afrHW4T1duNA3yrlKS0k8WVGmTJZFc,816 +torch/include/ATen/ops/alias_copy_ops.h,sha256=njs9obBY3RTj2MBiVLuEp_lVGCVaryS0Y-P7LpuY2aI,1825 +torch/include/ATen/ops/alias_native.h,sha256=e1VY3X5nTr-DwVXWC5FBOB0bDBWXPb5EEycaN37n9MA,789 +torch/include/ATen/ops/alias_ops.h,sha256=NYc_8dfK1oLYtC7uzl1HTSoFsJFIdd1EjXDFmhAjjEY,1214 +torch/include/ATen/ops/align_as.h,sha256=ujMiBDHo7VkIPyI8BTsxJQv-mF6WSHE8vzdUPvE6MRA,759 +torch/include/ATen/ops/align_as_compositeimplicitautograd_dispatch.h,sha256=b0Rrzlo1h8_RmUFCWArMX6OY0aOa7LdKRn2UjCJB-Vs,1046 +torch/include/ATen/ops/align_as_native.h,sha256=CdxT_2m488GjcGE7Pw3OxNGKVnuLmF_aH2E2HtTlmFs,758 +torch/include/ATen/ops/align_as_ops.h,sha256=YGBFsxFJxH6RQVtEmgwW4mu4roRl5HZXVdGlSz5gwaY,1303 +torch/include/ATen/ops/align_tensors.h,sha256=88_U4CVD5klXL14gov16_gufkvKZ5jwK-yvuL_j3Iv0,943 +torch/include/ATen/ops/align_tensors_compositeimplicitautograd_dispatch.h,sha256=6Sh7qf2gZsgjUmhwtye50w8M36tuClWiBC3gHNANBHg,1039 +torch/include/ATen/ops/align_tensors_native.h,sha256=BGzfMifCGiaNaBhtnu24C7PKxoBR_F-ZewIW_7InUaE,751 +torch/include/ATen/ops/align_tensors_ops.h,sha256=YJt6RE5TsyUh4OtYr1J5nCdAv_wZcHvWr26s_hVgGh8,1278 +torch/include/ATen/ops/align_to.h,sha256=kQ_ykm1cX048R4vl3hE12USWqvOyPG6LYM3KuAvVuSc,759 +torch/include/ATen/ops/align_to_compositeimplicitautograd_dispatch.h,sha256=ZK99o-bZ4zXlrxhYyXS0Y5FV2dhkugdCLzyJgeAYev0,1144 +torch/include/ATen/ops/align_to_native.h,sha256=NDLpM31kErwhRhXS_vL3FRgDVedpKGVDI-_Ql3f7bPM,856 +torch/include/ATen/ops/align_to_ops.h,sha256=Fxzd4I0jK5Cc89C8rJOUa-hzA2E_VtLkoI_rZqrXZgY,2004 +torch/include/ATen/ops/all.h,sha256=m-OA6Lk9lJxyq4NuVhBf_hY_hUq7jKoWGJJj-_KjhOc,3609 +torch/include/ATen/ops/all_compositeexplicitautograd_dispatch.h,sha256=mmbBiM2Zn5e1VPEby000TDr3kAHlXskmwYoIGij1P6I,1307 +torch/include/ATen/ops/all_compositeexplicitautogradnonfunctional_dispatch.h,sha256=5nlH9Qbc_MKWl2lGEcNmR6tP2b95ePZ31QKWDc7-8e8,1225 +torch/include/ATen/ops/all_compositeimplicitautograd_dispatch.h,sha256=hal_NnwWNXoVR2_s58r_-0c2OST6-Do9GluNodd-13I,1271 +torch/include/ATen/ops/all_cpu_dispatch.h,sha256=K2NbhOus9oXaAOUqXuvuiqp6jcHm3HR0h9iAnfLsMyA,1760 +torch/include/ATen/ops/all_cuda_dispatch.h,sha256=_XYF8bg24XGb1XozeaZ98VitehNLIK4oMQ436Ww5Xbo,1762 +torch/include/ATen/ops/all_meta.h,sha256=tfRrW8h6Too1ySoip6viegLP981E_miX23dHTx0ve28,1111 +torch/include/ATen/ops/all_meta_dispatch.h,sha256=JNgmdwnAbTf4tFuJjJfrZKm3jszdsrw1lCVSFS0yZec,1762 +torch/include/ATen/ops/all_native.h,sha256=d-_Gde2BxCqp-F6v-MRnQc8pxYAcEkJMLymO-fJ0png,1749 +torch/include/ATen/ops/all_ops.h,sha256=jAkvm6soh5cS-3jMXXm9xum6nZ9wi6WeQybxNgo8G2A,5830 +torch/include/ATen/ops/allclose.h,sha256=HnudG9IsIjn4-Rr0MO_hA4JfKw34snPxpvCy8Dlwt2o,1079 +torch/include/ATen/ops/allclose_compositeexplicitautograd_dispatch.h,sha256=9A3NW8DQFDHQRRwxI3Kg0IyQ1zhOY_iZRfC8JXr7Pgg,1100 +torch/include/ATen/ops/allclose_native.h,sha256=0TQiTzJCl0Di97b54c73YhPcwfT-TfQRid9JkNys3XQ,812 +torch/include/ATen/ops/allclose_ops.h,sha256=n68rqZCqEUG8tA7P36BpP1F0lgbzQsCrYwL6puKJjCg,1447 +torch/include/ATen/ops/alpha_dropout.h,sha256=N4BN3lvZSnWZOmMbGF7_PyQmPqCDxvvbh8vSs-Nz1ew,1191 +torch/include/ATen/ops/alpha_dropout_compositeimplicitautograd_dispatch.h,sha256=5xAG7bhcYGrOWWpNDCxflqmKO3sryvhiQ6Jp_n1T7sc,1128 +torch/include/ATen/ops/alpha_dropout_native.h,sha256=PosbfR95kB-Bn1ikwRDnjRlJjpQjLW6tkPGe9ki-l6M,840 +torch/include/ATen/ops/alpha_dropout_ops.h,sha256=9gk_Nx2RTfTLjXqUYxCabKLdILs16hSkWi2Yp_a9WYs,1913 +torch/include/ATen/ops/amax.h,sha256=WqDBBZstAAC-a32OWFFx_6YrOhs40w6VpK6BIq0CxQU,1531 +torch/include/ATen/ops/amax_compositeexplicitautogradnonfunctional_dispatch.h,sha256=epYea8UhHq6-xL290YLgWWk_pyUgz87Aym0c5HFPhz4,1086 +torch/include/ATen/ops/amax_cpu_dispatch.h,sha256=BTQMNI7KiMPfEpISmuPHGy3qG3XxGiuXhQRECRWWPsA,1248 +torch/include/ATen/ops/amax_cuda_dispatch.h,sha256=D7_kizvYrmC1qyyEU6fv0vRXhnmoyy_5bLH1lfDyLvQ,1250 +torch/include/ATen/ops/amax_meta.h,sha256=-kS1mONJ60VjoruggHXlOmvRGT--FOtrVgfVBi0nlqo,854 +torch/include/ATen/ops/amax_meta_dispatch.h,sha256=ahiWzmV4oILNaQJPgZMQ-E7mB1DS46l--naNxGGszHE,1250 +torch/include/ATen/ops/amax_native.h,sha256=pfyy_S3FTHq1pR5jhsf15c6Gk_jJlShnrUI-XwjTNgI,879 +torch/include/ATen/ops/amax_ops.h,sha256=WbUAGnxlfXCl-NZf9DVwy43WTPRBTkvPriOcKUTnSVQ,2045 +torch/include/ATen/ops/amin.h,sha256=ci0xVbQeMyy0giupqRD1R5MyIjTBUl4IZ8g5SCQqhF0,1531 +torch/include/ATen/ops/amin_compositeexplicitautogradnonfunctional_dispatch.h,sha256=WqXwuPaYDZMkJCM9W96AspBK8PpOrbj895feSVO0XvI,1086 +torch/include/ATen/ops/amin_cpu_dispatch.h,sha256=rVHPbkrq9bFHRHkFjZ0zQDQCHTnzAwuzYdTpFk3ukQ8,1248 +torch/include/ATen/ops/amin_cuda_dispatch.h,sha256=baxW91d_C_P56g7YqexLeyJHapim4YvyVz1qFhb-7oQ,1250 +torch/include/ATen/ops/amin_meta.h,sha256=ZHKQgvQLDtk1CkfpQWYp13MfWTyJtTFfuslxboLJej4,854 +torch/include/ATen/ops/amin_meta_dispatch.h,sha256=5Rd9iLZN05jRS1YT8lwAMZ_cj1foGHkfkwe30AAo0vE,1250 +torch/include/ATen/ops/amin_native.h,sha256=pwtGuBWvIgDPKHZUVMZPAMQZrAFmHS0nTKwr0P-xdG4,879 +torch/include/ATen/ops/amin_ops.h,sha256=rMxfVThw5U8NONM7hfX93GsBJBfyqN84QrB9OE_gPeQ,2045 +torch/include/ATen/ops/aminmax.h,sha256=bQGvNkG0savYjSdImC9eitYWHeLzMMD4OjbrqmRIh-g,1834 +torch/include/ATen/ops/aminmax_compositeexplicitautogradnonfunctional_dispatch.h,sha256=KAR7lCRKROBN5L9yD2T7b-o2OYr_zuunTQrf6MLBJ70,1135 +torch/include/ATen/ops/aminmax_cpu_dispatch.h,sha256=SpKfPBzl44ZC0BAC6RoR8NnHZJ7JR8cfjVjKRZRfZhA,1423 +torch/include/ATen/ops/aminmax_cuda_dispatch.h,sha256=oVDKenyb6eoDfj2-8_b7a-qowJqbcmYFIRjDaYZC6SY,1425 +torch/include/ATen/ops/aminmax_meta.h,sha256=G6d7rD8shu4SYoLNhwV4DG_5i2nNTsHyy4Kbt2UwEeU,866 +torch/include/ATen/ops/aminmax_meta_dispatch.h,sha256=BgMJAQt7Yi686_dG5Zs3KkbXr0UdnEZzIKOiKNc9zy0,1425 +torch/include/ATen/ops/aminmax_native.h,sha256=7IJhFuuP_DjdQyoJ2wHeJ2XG-eIBclTK4O3YjnBj1dg,921 +torch/include/ATen/ops/aminmax_ops.h,sha256=QISYeuZishhhlckBsmrtmYOLszmsrzd5KOXEVtO6o_o,2382 +torch/include/ATen/ops/and.h,sha256=AfdCY8DiDqbF-enw8ZCxtrQC4IVPH8R_Gv3JXBLz8_Y,1151 +torch/include/ATen/ops/and_compositeimplicitautograd_dispatch.h,sha256=5g_sYbFMv32u9xp5FWoc3sT7_-m0L1Z-s6UUL7HnroQ,1282 +torch/include/ATen/ops/and_native.h,sha256=-sX5_rHtYxeNREVIGa7Fl1wXYOKsYyW-CA0LOwnXMmE,994 +torch/include/ATen/ops/and_ops.h,sha256=J3Cr8oBu3R_BNrpnGPESqx55LJOJhOzYdXaXJ8ghzsU,3145 +torch/include/ATen/ops/angle.h,sha256=90wdYFa_q_k42MD_g8ctgyOf6BcfgLC5wyJjxDLYxtI,1271 +torch/include/ATen/ops/angle_cpu_dispatch.h,sha256=bnibbXmwA48hSTJfpselkDimsEj5UlRjL5YrvfBIyJw,1128 +torch/include/ATen/ops/angle_cuda_dispatch.h,sha256=KIM5oQrYRaEzbQ8PHwlQ23xDs2ZDbeliL69SchyjStw,1130 +torch/include/ATen/ops/angle_native.h,sha256=Xuzdky7qeAuYU1TS-aB1L8hvShvV0xw5zL-CBIRQgQg,958 +torch/include/ATen/ops/angle_ops.h,sha256=g80gxwCtt2_jc2y3R9PSMto3QHQVzHEX_0RIxm8XcVU,1795 +torch/include/ATen/ops/any.h,sha256=O2OM8fdyEedhvg4FwUt5bvwi5QsHznX-mU3EAgtb5-s,3609 +torch/include/ATen/ops/any_compositeexplicitautograd_dispatch.h,sha256=SJslF-IN092tN2SYcxQep8zmJWChwoexIAXRRkHEyng,1307 +torch/include/ATen/ops/any_compositeexplicitautogradnonfunctional_dispatch.h,sha256=XcZKZhHKk1oE_QmoqC5gvnQGNCXmKGzD7eXYN_pWAmE,1225 +torch/include/ATen/ops/any_compositeimplicitautograd_dispatch.h,sha256=9lbOK5LwRhI7dn-hJ-kloXNgW5XW5wdkST0sE2ujQWc,1271 +torch/include/ATen/ops/any_cpu_dispatch.h,sha256=sWT1UNvum9ZQW_QyhyNYxvSCT9JwGZDcR_velMkv8PQ,1760 +torch/include/ATen/ops/any_cuda_dispatch.h,sha256=RWeH9RF17J_I79QNlH9Pi50w-Rdz0Zhij_c8H4ZtOQQ,1762 +torch/include/ATen/ops/any_meta.h,sha256=wlYXdI-imjoS6AUmQ--hgq1uA9ZcKdzH5zKgkddHMuM,1111 +torch/include/ATen/ops/any_meta_dispatch.h,sha256=jQSiHAePU_687iSWqz9WotdfwUzwmVBfMhbm6NazRng,1762 +torch/include/ATen/ops/any_native.h,sha256=uR5vqKxAz6fp6AVVZEmWN0QExeUoTwWfwCDfhqSfGoA,1710 +torch/include/ATen/ops/any_ops.h,sha256=0EiCLcqDPkeGyRKb7WReZrvAMdtRMPTXYvBbeA4UV90,5830 +torch/include/ATen/ops/arange.h,sha256=ZC56YE4iOmBl6n-ac5x6eGYfeDWHAi9jTlEJ8N1QcqM,4425 +torch/include/ATen/ops/arange_compositeexplicitautograd_dispatch.h,sha256=BXxVE0ecp-tENMd-6TJ2liaRx90s-INb8clb_Pkmk7o,2116 +torch/include/ATen/ops/arange_cpu_dispatch.h,sha256=yftx4XApoiS2TGAtjlfhmlrpomD3pLK1hDD9_9bjTD8,1177 +torch/include/ATen/ops/arange_cuda_dispatch.h,sha256=O2d8hakjkuO08SEZrMmkiw_vHXSLjvbOIvP9G7X9J4s,1179 +torch/include/ATen/ops/arange_meta_dispatch.h,sha256=_WYLjKaJDqjXrA4MIWZ1VFIWRm791JGWY-X011MXXis,1179 +torch/include/ATen/ops/arange_native.h,sha256=AQB_g8uGZ36kn1cPsL3qyf-g7LzO_dlUhOLaJvCr5QY,1723 +torch/include/ATen/ops/arange_ops.h,sha256=mOYjd_YPPA904rmA3RyJjgPriDhogeFFezEo5lwjuqw,5359 +torch/include/ATen/ops/arccos.h,sha256=5Uk3Ps5cTYJOoI1uRfF4vGlfQf-7bexPSi4N-wanxos,1423 +torch/include/ATen/ops/arccos_compositeimplicitautograd_dispatch.h,sha256=20FTmdmQdnBpiD61wcZCla33mPnWhRHtsPBPX1XgoyM,1226 +torch/include/ATen/ops/arccos_native.h,sha256=yErsbOBMKzEgOZKqmt1u_XVGeYpUzYV6JI3xS8zDI6k,859 +torch/include/ATen/ops/arccos_ops.h,sha256=MtwZnA2kMzcF0qVWDGWqTw-xgF7JthGbchWCGR8PgGk,2300 +torch/include/ATen/ops/arccosh.h,sha256=p9En4OL2rdTbXa0tsA5KOuYvKbNBNhoWLC4pDv8Okc4,1436 +torch/include/ATen/ops/arccosh_compositeimplicitautograd_dispatch.h,sha256=ogslVweuObutgL4Pq66k-oWOZzC-oSLmoHaRmpScV5s,1230 +torch/include/ATen/ops/arccosh_native.h,sha256=vfT1T5VA205RoNphHt6a90IhfZ_dYC63Es25p9FeIE4,862 +torch/include/ATen/ops/arccosh_ops.h,sha256=qBgF84vdlKjvlWtKOHbtwJJeqL9WJuv39lLFYGkHbzk,2309 +torch/include/ATen/ops/arcsin.h,sha256=x2yuqImoVOx-osPoQKeprHFPaaVFgff5hO4SdAv6Y3U,1423 +torch/include/ATen/ops/arcsin_compositeimplicitautograd_dispatch.h,sha256=qqSFPhFEf5tpr9-8oFGdNyJohZRbvBf70g9OlDwavoU,1226 +torch/include/ATen/ops/arcsin_native.h,sha256=uMPZE0fDu53Nb6BneOe6tZY1Iue_Andh1TqafqK43RA,859 +torch/include/ATen/ops/arcsin_ops.h,sha256=WPA8_X5cnnxl8dJUkoB4bZtc7c6vroLjMj9L9CA0RtU,2300 +torch/include/ATen/ops/arcsinh.h,sha256=MDyYdl7ux4Z0_2sjMCWaSSj3jX8HhYmGdp_lZT3IpoA,1436 +torch/include/ATen/ops/arcsinh_compositeimplicitautograd_dispatch.h,sha256=vZu0a0jk-yLlFlmKWD9LrHvk1EDh5EDJ1HfSSsa_hho,1230 +torch/include/ATen/ops/arcsinh_native.h,sha256=JhmmU6Kl1vaghGW3C7f37BPIdo-GbvPM-2GQG8XR9qU,862 +torch/include/ATen/ops/arcsinh_ops.h,sha256=2e2tGcaiwr8f1LXAbdvNZtGeP4GFkMyA3xIqdFNiTGE,2309 +torch/include/ATen/ops/arctan.h,sha256=zb-WgJZLEAvPJLFeIdSNjSY90nJac5mr1GiHkNC6z7Q,1423 +torch/include/ATen/ops/arctan2.h,sha256=FYrvrRpCoh2pCPi9CsnlpaL8OdhUzdDRyXvgjvNa5W4,1432 +torch/include/ATen/ops/arctan2_compositeimplicitautograd_dispatch.h,sha256=fypc2XKT61xM0GQ8-dzVUjqfEXrPV_Y6KNfoHJctfVM,1334 +torch/include/ATen/ops/arctan2_native.h,sha256=Kwvcy1iWGYA3VuIx92CLaxyvhvdcVJnc4oHOA7OzdGs,940 +torch/include/ATen/ops/arctan2_ops.h,sha256=LwpOsawft6yiyzY5brvOu1tcervVBFkq5LW3Pmvy5c4,2567 +torch/include/ATen/ops/arctan_compositeimplicitautograd_dispatch.h,sha256=1MALW4LAvSK6hc3RgXlDDWWZF6CN_zAGeabYbAh-W2Y,1226 +torch/include/ATen/ops/arctan_native.h,sha256=IU-c4I1YdT6Iyofe4Jq4TJzSCtMCFHud4Wlgy5301T0,859 +torch/include/ATen/ops/arctan_ops.h,sha256=6La376X5EkAzig226WdvdaOD46s13NdsDdXQgh259u8,2300 +torch/include/ATen/ops/arctanh.h,sha256=nnaFGSuGPTVOMcdIJ5wiFaS_wa8sx5_K353TAvn8D48,1436 +torch/include/ATen/ops/arctanh_compositeimplicitautograd_dispatch.h,sha256=UuImT6ryokuenV8F2mvr1vsm52UKviVbEv6es7wX35w,1230 +torch/include/ATen/ops/arctanh_native.h,sha256=iDY4ajSbtFqagCIGfHWZOpGiLQZjNyOsd9aJsVxF1Jc,862 +torch/include/ATen/ops/arctanh_ops.h,sha256=3C5UrKxlNwi-FnM4P_bUJUC_EVxLp8LzcwuJld-tzEo,2309 +torch/include/ATen/ops/argmax.h,sha256=BOaWkUGCQ52vEFFOgGnphdWc7haL2syExSmeW6xwk0I,1602 +torch/include/ATen/ops/argmax_compositeexplicitautogradnonfunctional_dispatch.h,sha256=_Gy_M1Wd_sCmkTpPo06kWPX0JzcE3Oi1kLE8QMGO3vQ,1109 +torch/include/ATen/ops/argmax_cpu_dispatch.h,sha256=5VQgtrUwgkvLY5nJrTJmHsI6PZVor7PbJ-xmrgJI3TQ,1305 +torch/include/ATen/ops/argmax_cuda_dispatch.h,sha256=BiENtvgFCJ19UW_6hBSjJbaQ9iixl9fgAqHZAttY1Ao,1307 +torch/include/ATen/ops/argmax_meta.h,sha256=ylVsv-yozGfX0-4dvaqNLl6NLeesW78A4oVJVbdZaNY,865 +torch/include/ATen/ops/argmax_meta_dispatch.h,sha256=9LS6-qecZPjCzF0ZR-MLKbrOXSwxvc9LRaYUZo6ZAEg,1307 +torch/include/ATen/ops/argmax_native.h,sha256=i98fmx-73yBjZuats3G2x7tF9NEVlLvtq27HthxjRPg,894 +torch/include/ATen/ops/argmax_ops.h,sha256=WwP7S_xzRIzayYstRW49xEbULtm3at6ZJosHkP-rKDA,2111 +torch/include/ATen/ops/argmin.h,sha256=hWI6tUKRkEXSxpn_gUxSFTsMK7dRlPqdu5eh-KJHKH4,1602 +torch/include/ATen/ops/argmin_compositeexplicitautogradnonfunctional_dispatch.h,sha256=OErO2xFMI74nlrcYGv6TypvHuSrqmklo6Aq9ugaQK_A,1109 +torch/include/ATen/ops/argmin_cpu_dispatch.h,sha256=yIM31vEpJ-sPDyBMFN_6XAHTG2nM2EBJrAbwVHz4uIU,1305 +torch/include/ATen/ops/argmin_cuda_dispatch.h,sha256=9mKzbp3MenIDMgiXWxYd0ntzta0h_J6AIENiwT9Y1Wc,1307 +torch/include/ATen/ops/argmin_meta.h,sha256=0hZA7tEq_TxBwj30Hy01j6-tswYRnpJNTJFyaMY1xCw,865 +torch/include/ATen/ops/argmin_meta_dispatch.h,sha256=9TZOAvEJz7ZFyvS_28u89EIaAsaTvL5WCMcJv1cfTCk,1307 +torch/include/ATen/ops/argmin_native.h,sha256=F9KZtYuNzEypuWlaIM2naQtYfeGV4oDKBiyqWwyTI1w,894 +torch/include/ATen/ops/argmin_ops.h,sha256=3osfrvcAKFrAXmGVNf_tG2Pwoz8nEMnf6ld_giRjQIs,2111 +torch/include/ATen/ops/argsort.h,sha256=6GwFmV2vXCr4HiJuKO8MFz4QR08iqvO99X3E4opoe10,2178 +torch/include/ATen/ops/argsort_compositeimplicitautograd_dispatch.h,sha256=C8IzcVGUjSwf-We7uwDkbaFU_XgGd1k65Oms5R6uL8A,1514 +torch/include/ATen/ops/argsort_native.h,sha256=n4WWFfxBzNNrRmbneNWmokbqfksdOzZrYGOGlwtZc-A,1094 +torch/include/ATen/ops/argsort_ops.h,sha256=7btlypVACdG56oGqkvOUEMRjGG8Sz5DObjPIbdH6Hlw,3423 +torch/include/ATen/ops/argwhere.h,sha256=Nq_Zqfh23SZ6-Ti-hxFCc0-j-CXcmyPQ6HM2GPPqTs4,899 +torch/include/ATen/ops/argwhere_compositeimplicitautograd_dispatch.h,sha256=AB3PDtUWP1Xn6X5WFaZknKUi7lLm_XvWGl4bGEzfID8,1020 +torch/include/ATen/ops/argwhere_native.h,sha256=gmts4IxllA2N6bVI24zT6IMq-7oFetY-rbGzpNq0POg,732 +torch/include/ATen/ops/argwhere_ops.h,sha256=nA-QJ7XKGbtt8YD3wXA2CnQhS5DxN9UzNmC5SRarPSA,1217 +torch/include/ATen/ops/as_strided.h,sha256=wUrhKUFhKNPg8zaIGq3CGfXUaUUGG8PtGPk54_KZr_Q,4132 +torch/include/ATen/ops/as_strided_compositeexplicitautogradnonfunctional_dispatch.h,sha256=AhnQg49Y4YLSokWXqOfLP8iP1nZ_SiSfPtkOhV3DaO8,1347 +torch/include/ATen/ops/as_strided_copy.h,sha256=CtKlIupFJhHj213e82Fp6ONHTrkV0tnDS6KU3x99HXo,6184 +torch/include/ATen/ops/as_strided_copy_compositeexplicitautograd_dispatch.h,sha256=5kBkiuBTg8sYOXYztjk-V8_R2gUejpgrA6D1CDnxOsQ,1730 +torch/include/ATen/ops/as_strided_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=0cUJbTv-Qw0U5zDFtvSpABbEgkT1-hHIFJW88-fYs7A,1339 +torch/include/ATen/ops/as_strided_copy_native.h,sha256=c2Pp2Wb6cKoamLB9oE9gT2dMAhLzBgo4DXGG76dOS54,1053 +torch/include/ATen/ops/as_strided_copy_ops.h,sha256=LgZ-m1QHzwbzFzMtMKntMrOLNuYfZ5zBLNGRw9W_0cI,2517 +torch/include/ATen/ops/as_strided_cpu_dispatch.h,sha256=U4-kH_sTgIBunN5YXokn_WprFxJXYZylhC1Q2ZWVZmM,1259 +torch/include/ATen/ops/as_strided_cuda_dispatch.h,sha256=eDfgl_rwHq8ugvWhJA3QrSZMflMmgXVmSMtzHZ5k7MI,1261 +torch/include/ATen/ops/as_strided_meta_dispatch.h,sha256=rLGfb9kI-PMVeQNuP4TVl-6XghaAFDxurm6olzH8Ae0,1261 +torch/include/ATen/ops/as_strided_native.h,sha256=fC7uGifTcEzKKbwRdNoHsjYHSkolg24wzsrAHXPtpN8,1402 +torch/include/ATen/ops/as_strided_ops.h,sha256=GoIWysvt1VPdLBm6ogoVxJw3nXMjxSK3nK9a5FFw_hc,2438 +torch/include/ATen/ops/as_strided_scatter.h,sha256=YLigoF9Kw-XMdVICCakVhFmRwLCJ7qE923tCwt2IDPE,6697 +torch/include/ATen/ops/as_strided_scatter_compositeexplicitautograd_dispatch.h,sha256=1upHXfTvZhSrMP8rWZ21dF1Hd6Qepqa7aB6w6YCWD2o,1838 +torch/include/ATen/ops/as_strided_scatter_compositeexplicitautogradnonfunctional_dispatch.h,sha256=J8uzr66aZPTstDIErZHKVCFF7Qzjku6P7Bb4xwIF9gU,1393 +torch/include/ATen/ops/as_strided_scatter_native.h,sha256=aulN12IHKBwp0In-BvNX1E5idPdIsCfzWkjdV3fPBdU,1107 +torch/include/ATen/ops/as_strided_scatter_ops.h,sha256=C0VF4z8kW6pJtdVK35FlPeXB3t0io1XigWBSvkNAudo,2695 +torch/include/ATen/ops/asin.h,sha256=6dZJKY4qGci-jldWDnOKRWrA-a4swl5A04YnEX59IpE,1397 +torch/include/ATen/ops/asin_compositeexplicitautogradnonfunctional_dispatch.h,sha256=pS1KONtDptODIyWb24gDwd6zYWxqq3UU8ce95217Z-I,1091 +torch/include/ATen/ops/asin_cpu_dispatch.h,sha256=TU_cp8cqXXLd0TnQcWdAdyIop5VfjeywdzIvHB6MNLM,1174 +torch/include/ATen/ops/asin_cuda_dispatch.h,sha256=sYh8Y4VSCZDrIeXoD9Ugb1TRiPyImmdL1GvB7KkU83E,1176 +torch/include/ATen/ops/asin_meta.h,sha256=O5kgNxFEMt2e7Y1xEIYhIGO2B3kKa9gbGwGBEVdS05k,819 +torch/include/ATen/ops/asin_meta_dispatch.h,sha256=gGj8XqdzL4KOQ_M1KGbS8hyHLJ52fVo9zliBkBgvKZU,1176 +torch/include/ATen/ops/asin_native.h,sha256=eyknEruGD7pH_F_RW-r7nhqQm-rYZRosKAchUKsA48Q,1252 +torch/include/ATen/ops/asin_ops.h,sha256=RGX5MrYJ1ypFCCveF2MMiISTd_1bGdYfR4Cty0VpU-w,2282 +torch/include/ATen/ops/asinh.h,sha256=B5iPezw69GwhGn9zthCnlgb-PR4wCqx-KM4etx_tpoQ,1410 +torch/include/ATen/ops/asinh_compositeexplicitautogradnonfunctional_dispatch.h,sha256=tA0Ci9Wp9pTHmyQu5KdZht4A8wwVm6GAoR9vFh5kph8,1093 +torch/include/ATen/ops/asinh_cpu_dispatch.h,sha256=gL3KygVUN0VpdOvH3QmVGjCpbnS0LetxoK5MZnpcqSY,1178 +torch/include/ATen/ops/asinh_cuda_dispatch.h,sha256=fcgoxE_B54nYsLDiVp0sWBnwF5E7-KjctvdapzFUynw,1180 +torch/include/ATen/ops/asinh_meta.h,sha256=uu_oPJwLAO80zNWX-YNSHXAuLgGkozoqqfsVMdLkR44,820 +torch/include/ATen/ops/asinh_meta_dispatch.h,sha256=7VJ0tjC6w2_Fa9lMuzpIHSa05LKKMpJlXSUJQq5VcJ0,1180 +torch/include/ATen/ops/asinh_native.h,sha256=3z81vK6b55PFPiN-UB4fxmBLi6dpvh4wVs96P3oAVT4,1261 +torch/include/ATen/ops/asinh_ops.h,sha256=7EfRswjdCqmUf5Tr3ZiGP4IsoiRKuX07HnQOXYRbYNY,2291 +torch/include/ATen/ops/atan.h,sha256=-SJqz5oGrT9HvRgKsmS8SqS7IXD5ol32w8215ZW6WJE,1397 +torch/include/ATen/ops/atan2.h,sha256=1zEyo-XJ40uIovW8ywDpT2DMvDTwYgkwNjn0kNOva7U,1412 +torch/include/ATen/ops/atan2_compositeexplicitautogradnonfunctional_dispatch.h,sha256=692Lrb7ojzHqKgGwo_Stt68QD-C4BEomPM7d-PAHjSM,1145 +torch/include/ATen/ops/atan2_cpu_dispatch.h,sha256=cj6dDror4xzOtToQm5Vw-gCM7K3R9vp4YYaUSjoWknY,1282 +torch/include/ATen/ops/atan2_cuda_dispatch.h,sha256=Ni9e9SPjoVHutlIv8hL0Bc9UKNw4a1j5fxmjbsNEh3o,1284 +torch/include/ATen/ops/atan2_meta.h,sha256=P--U8dOOFN8U_j-FDkO5Ph8h0iS8BomIupQnzvPrr28,846 +torch/include/ATen/ops/atan2_meta_dispatch.h,sha256=Wgbo19_on-VDMBS3UE9ucOdZ9Q0svlvqboVI7awKWDQ,1284 +torch/include/ATen/ops/atan2_native.h,sha256=vMivDfcTyRxV2AuT3HVuqrA70wfPqxN5PAi25nxhQ7A,873 +torch/include/ATen/ops/atan2_ops.h,sha256=Z4Y95GL1FWpzjsDe7khNOhJ7_4_a_2_kNJBOvinm42c,2549 +torch/include/ATen/ops/atan_compositeexplicitautogradnonfunctional_dispatch.h,sha256=F53co3TTVHf_6VdjI5ot-pBoph23Am70UpCTe3r9BxM,1091 +torch/include/ATen/ops/atan_cpu_dispatch.h,sha256=pZl0lgc8um-cj36fSxWUhfXZwf2aiISgexPVxa1TYlo,1174 +torch/include/ATen/ops/atan_cuda_dispatch.h,sha256=BUOsp0_a4krYjeNuuawq-LdYGuPKFhj4sOQJtacMozI,1176 +torch/include/ATen/ops/atan_meta.h,sha256=mLf6Rm97Uuj2PRlSXevn70Bf3Sc4rA6HY0IenQiOTzk,819 +torch/include/ATen/ops/atan_meta_dispatch.h,sha256=23awb09Suck2tJK7_13e8AXn8qOvO05hhW_gw0XX2YI,1176 +torch/include/ATen/ops/atan_native.h,sha256=F48TyFjJBGas9TBbyqFrD_V8Veiz44nDAAL8CJ3Nuu4,1252 +torch/include/ATen/ops/atan_ops.h,sha256=oERuhew2uNIkvGwFI_NxIiJIYdCNMG00D6HJZhTVF0I,2282 +torch/include/ATen/ops/atanh.h,sha256=VZJaifaBgMMw0WKMvYy2zMUlRlBxHGzDtwjTZLD9GWk,1410 +torch/include/ATen/ops/atanh_compositeexplicitautogradnonfunctional_dispatch.h,sha256=aoVfbRzWJ8yo7LkQjQAgodxiZHNSR75YyZnNECFw_08,1093 +torch/include/ATen/ops/atanh_cpu_dispatch.h,sha256=lGWIjCsDc6wnyHc0GHeKb6nKJcdYZQzMG8rvMpyK2tk,1178 +torch/include/ATen/ops/atanh_cuda_dispatch.h,sha256=Ynl5-b8IA0bOpT7DJXhM6PApF2MiaT8TJw6KgckJcGU,1180 +torch/include/ATen/ops/atanh_meta.h,sha256=jOwwN_UxjCWHJ_WuQ7NuGP-RcXrr5i-ezrM9147aIlo,820 +torch/include/ATen/ops/atanh_meta_dispatch.h,sha256=HoqJGAxHmTwxED5N7KpFlWLH4_0T4UVJWzf5o77PB9A,1180 +torch/include/ATen/ops/atanh_native.h,sha256=wKNo6q05YLH7nEO0AV_PpcffS4Yv5oa9BBAZMYjUUy8,1261 +torch/include/ATen/ops/atanh_ops.h,sha256=qXSEXc2cRG8bRy-Ks48xbqQcH1bxxrpeJl0yJa8428E,2291 +torch/include/ATen/ops/atleast_1d.h,sha256=Y2h7FBWO4_ZdnxwkH2n7CFeDPKQFttDw1JT1qK63XVM,1096 +torch/include/ATen/ops/atleast_1d_compositeimplicitautograd_dispatch.h,sha256=cSV_5QNtgEvLGPl67GDLC0vs23zqrIOo2JHxTilNUAM,1094 +torch/include/ATen/ops/atleast_1d_native.h,sha256=qGWk-JixUS2GW_mSwuEnsqIzEdAsIwFAeUu28fNzbWE,806 +torch/include/ATen/ops/atleast_1d_ops.h,sha256=yzKSzeCx6HkC2RHkTc1B6Yp42UC7S8ljja8wdvs8Ziw,1807 +torch/include/ATen/ops/atleast_2d.h,sha256=mzq2LFo5UJjh140DDHNcKEPzV3hX0MSzZWzvNyaMQ2E,1096 +torch/include/ATen/ops/atleast_2d_compositeimplicitautograd_dispatch.h,sha256=9HA8LiiHCm4baPn3MBQhTYLnJXpQUZec_uTB9whnfy8,1094 +torch/include/ATen/ops/atleast_2d_native.h,sha256=ID2YC10s5sZaHUjsDbj2apIB8COPweEREHXMzxsHMuU,806 +torch/include/ATen/ops/atleast_2d_ops.h,sha256=nYVgyh6lO1GPgRVQEeWyfL7TzmhHkD_lat8xpYhrPQc,1807 +torch/include/ATen/ops/atleast_3d.h,sha256=fTcDHZ3EI-auZ_BSG1FUBFheWNq8N4uHnHpRlYf6-sQ,1096 +torch/include/ATen/ops/atleast_3d_compositeimplicitautograd_dispatch.h,sha256=ymjqBMIDFWQaeQy9tBa8J5MjUZnHuKAg3EmO9Bl7D1w,1094 +torch/include/ATen/ops/atleast_3d_native.h,sha256=2DY9kGlI_Zwzo47BV_q4Gjgz-SXcrqTme-wVCEXUdxM,806 +torch/include/ATen/ops/atleast_3d_ops.h,sha256=Rx66_U0cioX8NSPQ4s1xSGGxrCXwENLe0v_ItxEiZY8,1807 +torch/include/ATen/ops/avg_pool1d.h,sha256=GAS9xY3dzNIJ2GPehFWg_oU42MV_XoBeR71nLMDwIfY,2208 +torch/include/ATen/ops/avg_pool1d_compositeexplicitautograd_dispatch.h,sha256=3-Ah5DXwbjPddbzmx0AnsjNLXaN_XsiIfUmtZHT5sw0,1381 +torch/include/ATen/ops/avg_pool1d_compositeimplicitautograd_dispatch.h,sha256=po0kt9Zee6y434D9BmV8e7s_V6QxThEtmDcTdL1jYjo,1156 +torch/include/ATen/ops/avg_pool1d_native.h,sha256=OhtFUdUNC5uIqQUVeJYovj8LTIt12vTSdHQDSJrHHCA,1068 +torch/include/ATen/ops/avg_pool1d_ops.h,sha256=fF3y7aNqFX9a7QNPGYA1CFrNiKRZJRbKqk2Js_56g7E,2637 +torch/include/ATen/ops/avg_pool2d.h,sha256=zzctGWwcWrZFEL6KIf59HWwuyMvCGhv3W883Y6ES3cM,2505 +torch/include/ATen/ops/avg_pool2d_backward.h,sha256=Jsxj5sW6zD3gIhziA-EKNByq9Xm2ibUQ8HaKGic1XuU,2735 +torch/include/ATen/ops/avg_pool2d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=G-nAYYYS6LghwifGh3ZuLISuI02wvCK4qDo7pSNZ9Nk,1250 +torch/include/ATen/ops/avg_pool2d_backward_cpu_dispatch.h,sha256=66mB_AehRWKNj4-vSKo6lHbV2E1fkaTarGmrKhE5QY0,1763 +torch/include/ATen/ops/avg_pool2d_backward_cuda_dispatch.h,sha256=6dPIVB2jOmbjMsIg-IdmC_6dlCLdxLkAHwgDTnxkaNQ,1765 +torch/include/ATen/ops/avg_pool2d_backward_meta.h,sha256=T_nfx1dVnS9aBFrp4msWpjudKDDkKxJ-EI4UNPcxMcg,1027 +torch/include/ATen/ops/avg_pool2d_backward_meta_dispatch.h,sha256=2CfFr-CY9XsEBtx9hDHRTZgm3nz3GM3K7ujJa9VU7qY,1765 +torch/include/ATen/ops/avg_pool2d_backward_native.h,sha256=Pd1ltV_pBxAXjKaixWNRWmcBg9UGcS4JEfBX_Ou3DfE,2030 +torch/include/ATen/ops/avg_pool2d_backward_ops.h,sha256=TP5KSfqcZcJoZk-f7ystcV2VEACZsLzFtGun10ihoZA,3179 +torch/include/ATen/ops/avg_pool2d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=SjBe7jvi6SroirbD0Jgbfoxe_FCw1LAgYuMOu1p9G2E,1240 +torch/include/ATen/ops/avg_pool2d_cpu_dispatch.h,sha256=Fr0c3sjjXUVGXK6k1RylsLoGM1oO6kZurn_tJgK1NdU,1688 +torch/include/ATen/ops/avg_pool2d_cuda_dispatch.h,sha256=htK8eLQ1_aj6bbomIN9NoxRas8Iwbw1AF9yXPqiB_1I,1690 +torch/include/ATen/ops/avg_pool2d_meta.h,sha256=jhOe0UDiJonIpDovNKbSfNX-A-eR9QWStrF0_caOPic,3696 +torch/include/ATen/ops/avg_pool2d_meta_dispatch.h,sha256=RLSQM1VAuzAM1SZ19RVU05Hhi3FQEo88jA-wN9B4Uds,1690 +torch/include/ATen/ops/avg_pool2d_native.h,sha256=CRBYguBG8G4n6y_AnB8YlVjVIEHW1eS5HLFxMdf_3FY,2109 +torch/include/ATen/ops/avg_pool2d_ops.h,sha256=z3mHD10BR9MEjIUDfzofY7MZmkK-WTiuLnDw7dShu9I,2917 +torch/include/ATen/ops/avg_pool3d.h,sha256=pOHOCZfdU02EbBNYScxKtIAaqvVu3Rx8t1HYD29Hk1g,2505 +torch/include/ATen/ops/avg_pool3d_backward.h,sha256=J-JoplellOM4YMYJGH6XxYOtaYjn2TjtPDs_q4TbJiM,2735 +torch/include/ATen/ops/avg_pool3d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=fHCPVhseY42rw4HNB1wfPzrsUVRhnN4c0sQlReSBVp8,1250 +torch/include/ATen/ops/avg_pool3d_backward_cpu_dispatch.h,sha256=QXjzdsDhsrTJsY47iuBJOlS_xxlpic---PuvW_5Cmzk,1763 +torch/include/ATen/ops/avg_pool3d_backward_cuda_dispatch.h,sha256=O3HTIehTjwReRGRDq9eYQCJvnx_danLT_YIj064XMIQ,1765 +torch/include/ATen/ops/avg_pool3d_backward_meta.h,sha256=t-guFOuk1OKRykPPza4QV1ErOITzDteuXCd0JT2lKuo,1027 +torch/include/ATen/ops/avg_pool3d_backward_meta_dispatch.h,sha256=L4y34jkq0CHscSAi-D_lJ3aWEmivEx__lFLPDFxFTa0,1765 +torch/include/ATen/ops/avg_pool3d_backward_native.h,sha256=_FEs7qPPaDLXjIe4cscLmZfE-qXNCahKr-wS_VLbL44,2030 +torch/include/ATen/ops/avg_pool3d_backward_ops.h,sha256=lJ6-lxJ4XpdVKiuvP_n01E3TLrUYdA_gyu5u9dKZ_7A,3179 +torch/include/ATen/ops/avg_pool3d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=vr0DzAinimkzlf6qWuYVlIWFdyuZoVLj71SQd7w1CNM,1240 +torch/include/ATen/ops/avg_pool3d_cpu_dispatch.h,sha256=0omjaNll7JkYCMVZZIvIYThYRp-7n6Uyo3m8lbkugjQ,1688 +torch/include/ATen/ops/avg_pool3d_cuda_dispatch.h,sha256=p4qMEwwGfom8p-C5OGZ55Wk1gDx2vh3RbeM7RNzSSYw,1690 +torch/include/ATen/ops/avg_pool3d_meta.h,sha256=oS2j0t1wsZ-Ggs9qPdGdoiyrevuKP0hN_NhDaL1X5do,986 +torch/include/ATen/ops/avg_pool3d_meta_dispatch.h,sha256=ajLQpnjRtIxPSMz9YjyEOIQCs9yOx7_yx-x8zi6OfIM,1690 +torch/include/ATen/ops/avg_pool3d_native.h,sha256=tciz5tF8OwHbrcSHh3VXlrlbt_4cC9_F1khH0tVuZ54,2113 +torch/include/ATen/ops/avg_pool3d_ops.h,sha256=2UuwBhneqHoU32UPwBQnFfsa-in1XOTjHJxN-7GspCs,2917 +torch/include/ATen/ops/baddbmm.h,sha256=rnbZi1R9Lny3CABJZt_U6fCTTAwcJeo7YAKwvnz8J6U,3230 +torch/include/ATen/ops/baddbmm_compositeexplicitautogradnonfunctional_dispatch.h,sha256=wKXOhmtTSl_bkkls4IUVy6ph3VeMvyRGsvphh4TKilk,1315 +torch/include/ATen/ops/baddbmm_cpu_dispatch.h,sha256=SdfMUkrYPByfx6zkDeIYS-CprSjrKNFbKaTc53Yg1BU,1618 +torch/include/ATen/ops/baddbmm_cuda_dispatch.h,sha256=TXpPDMF5qXu5nw7dVAEenyx_mZ7rvz6aj5CNednlUSw,2235 +torch/include/ATen/ops/baddbmm_meta.h,sha256=jgBmEsntIO-92RLc42kBP987RlQo2ELCjkLCOUiTdnU,927 +torch/include/ATen/ops/baddbmm_meta_dispatch.h,sha256=7t1cJ8xpyTSsnRapiMq1TDjO1Zp8XdPatDTQPw_0Q08,1620 +torch/include/ATen/ops/baddbmm_native.h,sha256=t8PaKQDsNURwvL74Oq6ZhU7kXsiadqFBcK-47GvtAQM,1839 +torch/include/ATen/ops/baddbmm_ops.h,sha256=Q8cSb3m4DefUVeY7wVu2Ea-eOaOcGYgtpwnWMRlo1iI,5384 +torch/include/ATen/ops/bartlett_window.h,sha256=EcWHAx2pNqicPUPWA39vER6JEhKbKMtmfT6uWXD8oQw,3671 +torch/include/ATen/ops/bartlett_window_compositeexplicitautograd_dispatch.h,sha256=iVkQn8joRZw1k6pYXNCfjVvATo-qwQl236TUCHkr3Mw,1960 +torch/include/ATen/ops/bartlett_window_native.h,sha256=ithhbQXMutpR9w3u1R10znm3L-_cL4yPzdGAQUjhl-4,1321 +torch/include/ATen/ops/bartlett_window_ops.h,sha256=KeRzUE3yz-ab79mrzWAW159X6WhUncZeGfwUwHkNweo,4110 +torch/include/ATen/ops/batch_norm.h,sha256=LFi4DCZm_lZ-UEAPcQh1C3Z1vAInCV-3EClwTTFChPs,1375 +torch/include/ATen/ops/batch_norm_backward.h,sha256=Ix3krX6gPn_6ka1GoNXze4hniInBaI4y9jnTnI2hzQM,1609 +torch/include/ATen/ops/batch_norm_backward_cpu_dispatch.h,sha256=Thw7yMxSfgAW4rRmDZmVICF6ttmojf3mGzCE-ZQP-zI,1359 +torch/include/ATen/ops/batch_norm_backward_cuda_dispatch.h,sha256=A92xSMyMeqTSTaUaSe10oxFNvpkhIRzr5FzG_SEjiE8,1361 +torch/include/ATen/ops/batch_norm_backward_elemt.h,sha256=vc2wMSx4kgE0EEVcOkwCkCnNkiHQPMs21SWn6qR9zrI,2611 +torch/include/ATen/ops/batch_norm_backward_elemt_compositeexplicitautograd_dispatch.h,sha256=3otyyE6qzQGikvHxuA-dhTPaVNezovDCqT2U0SCWzfQ,1579 +torch/include/ATen/ops/batch_norm_backward_elemt_cuda_dispatch.h,sha256=155d-60cp5mMW4t9StIIrkrVoIVvFX3OnB_DFHpKnz8,1205 +torch/include/ATen/ops/batch_norm_backward_elemt_native.h,sha256=wsILFVrlVTfv48hHnBQ7tzx8JnYGyYe3fKzwuF64Jtg,1271 +torch/include/ATen/ops/batch_norm_backward_elemt_ops.h,sha256=iREb6jye293XQtvBJYCObr9KeNdHeNEUyaSMDvQZBWc,3289 +torch/include/ATen/ops/batch_norm_backward_native.h,sha256=58zP0Y852N-k4-KpC-7WH5S5CuDJKCGe6oOsbo3iEYc,2024 +torch/include/ATen/ops/batch_norm_backward_ops.h,sha256=15IQ8uCSE4hesE4WJxKnAX6-aNQXPQhq6wioZEjryGA,2466 +torch/include/ATen/ops/batch_norm_backward_reduce.h,sha256=5i2EX80H-ybANCg6RP9_bvZ1Bp8g_Y2E2qp9HUFNnGQ,2990 +torch/include/ATen/ops/batch_norm_backward_reduce_compositeexplicitautograd_dispatch.h,sha256=9aI10F1lQ7fQopytP3-WPmMLRLzqd1R5CWTjlL5ZcQ0,1719 +torch/include/ATen/ops/batch_norm_backward_reduce_cuda_dispatch.h,sha256=QOIQ-1ZFgwUHyZVg3CV07IEd1foVlGquKuRkXFDqy2I,1211 +torch/include/ATen/ops/batch_norm_backward_reduce_native.h,sha256=163855V2OnH39p6_u0coTSPxTTYwLRxbjTxtJLVvdes,1347 +torch/include/ATen/ops/batch_norm_backward_reduce_ops.h,sha256=5GG1OEWLAh-nRB_e6XJynAAGKyZi3eT4dgaurfFqeq4,3605 +torch/include/ATen/ops/batch_norm_compositeimplicitautograd_dispatch.h,sha256=U97fC9tk8J6mw9qaq0nR2WpZKMFR9CX7qQ3Ehyue0_8,1272 +torch/include/ATen/ops/batch_norm_elemt.h,sha256=CqYc72LHO7iq5yi9FCm9mxKanWXtWAjZV1e8LV3pEdg,2146 +torch/include/ATen/ops/batch_norm_elemt_cuda_dispatch.h,sha256=-gOrrAWEUE1LDZBh7opbmsiNbb50BzoxyO32DtF8zvU,1616 +torch/include/ATen/ops/batch_norm_elemt_native.h,sha256=JHtJh340ZK_GgvwJxJk8jb2fEQyBMWlUpmzgRDn3Q1Q,1140 +torch/include/ATen/ops/batch_norm_elemt_ops.h,sha256=KOgVoKZe-NddNNcvxiQBgFknh8cHmB9b6TZhUw0aGK8,2849 +torch/include/ATen/ops/batch_norm_gather_stats.h,sha256=J90Pi9BhzMyv1q21l1ObqIHVGOro2_AQeKN8EYdTL7s,2768 +torch/include/ATen/ops/batch_norm_gather_stats_compositeexplicitautograd_dispatch.h,sha256=CU-WZCQ4eCkXfpnE8dko9fneQRz3ta--7zpkjjhNp_I,1641 +torch/include/ATen/ops/batch_norm_gather_stats_cuda_dispatch.h,sha256=erMlizQQFYtz4lXSWwjrBq7U8luJgR3Wpjv9axE2G8U,1214 +torch/include/ATen/ops/batch_norm_gather_stats_native.h,sha256=-QiU22S_AOUFrCVclH9gFJ6_Q4DoUdi5Cs9s_VT6kUw,1311 +torch/include/ATen/ops/batch_norm_gather_stats_ops.h,sha256=B_F2AJG5mlBP8nNk5Ld9qwlwy-keUwUj4QGpz0ch5Cs,3437 +torch/include/ATen/ops/batch_norm_gather_stats_with_counts.h,sha256=0SZm219jWI0J3QmMrfKxOvBoD12d3tCx17Fc2mu3cvI,2939 +torch/include/ATen/ops/batch_norm_gather_stats_with_counts_compositeexplicitautograd_dispatch.h,sha256=rl9fglm_nFsxlkOCWgUPsqhcWY1mZTWHsGMcn1AnPD4,1689 +torch/include/ATen/ops/batch_norm_gather_stats_with_counts_cuda_dispatch.h,sha256=gHAx_WiKr6R69eNX2uKMeRzD4d3_kjBkC8oqnJU35vE,1238 +torch/include/ATen/ops/batch_norm_gather_stats_with_counts_native.h,sha256=O2znHltp1VMqbAwnZp8MbYThS5KFS6KCXUK1O1hHTmM,1359 +torch/include/ATen/ops/batch_norm_gather_stats_with_counts_ops.h,sha256=kWTLHmVV5JlB1DNZDiicaEyh7mijEVNaHT-kDIN6vtk,3587 +torch/include/ATen/ops/batch_norm_native.h,sha256=DhlN-5nJdZ8t6jRw-9QLGov-SoFCZRfR4yFzfJx1Tbs,984 +torch/include/ATen/ops/batch_norm_ops.h,sha256=S5H5cMyVDvNoG-FAsv3s8cN2QrLsQAwLMn4VaD3WzSc,2035 +torch/include/ATen/ops/batch_norm_stats.h,sha256=a1O45Ft0a3M5C3-7JyG3bRxSLBiXfSDb131rZ326YdY,1681 +torch/include/ATen/ops/batch_norm_stats_compositeexplicitautograd_dispatch.h,sha256=oWcMS63gF6_rWFEzCU32hhgeHZajls6rzDZVALJPUms,1261 +torch/include/ATen/ops/batch_norm_stats_cuda_dispatch.h,sha256=tUArcNzKAK512laBnEdIAoCTR7bWF3BT7e1OXcOv3f4,1024 +torch/include/ATen/ops/batch_norm_stats_native.h,sha256=-rmoHPcCKDPEEZJ4OkOSTl5SACBflyTLo8-gJFALTYY,931 +torch/include/ATen/ops/batch_norm_stats_ops.h,sha256=pD6tFonqozI0Ge_tkD3E6CaPQ8AyhJbB3h_vb0Pg-OY,2205 +torch/include/ATen/ops/batch_norm_update_stats.h,sha256=Z3iMOEL_uxncsA0aqmeDn6MY05UHrwEByzSdMj_ypaY,2303 +torch/include/ATen/ops/batch_norm_update_stats_compositeexplicitautograd_dispatch.h,sha256=gzVw2m52ywuj0Hi16OVNGRo6NPLow6R4nWt4aoydZO0,1483 +torch/include/ATen/ops/batch_norm_update_stats_cpu_dispatch.h,sha256=G110gtSbrjf54bO-HUaUwGkCh3tUJdlvgEp8GYJAyYM,1133 +torch/include/ATen/ops/batch_norm_update_stats_cuda_dispatch.h,sha256=UFmJ6x_bWNh63cVfSgkeTLILdf90Yu4XMyoStq8qWts,1135 +torch/include/ATen/ops/batch_norm_update_stats_native.h,sha256=1rjoRTg3jwuVNsODK9Vl6GN1WgPfCZueOMtLhZTBHic,1370 +torch/include/ATen/ops/batch_norm_update_stats_ops.h,sha256=qN7CflwLSXAjnQAQqKReZxoBvpUdOrqHVcBAzp1DACc,2907 +torch/include/ATen/ops/bernoulli.h,sha256=B0xXToy5r6yPYDIEaCyTZ8pNlOBTiS0Ohr-Vp-BD-I8,3487 +torch/include/ATen/ops/bernoulli_compositeexplicitautograd_dispatch.h,sha256=qxWLxWJOqrGTmc_lFudOAPrinIS7QTRr11Bn9vz4plk,1806 +torch/include/ATen/ops/bernoulli_compositeexplicitautogradnonfunctional_dispatch.h,sha256=PQU6gOYC_HNaDLhlJCLCdKlCidKHTTfqVFTBTXyvNH4,1114 +torch/include/ATen/ops/bernoulli_cpu_dispatch.h,sha256=yibi4VzrAVN1Kp1uzxIpVFMTx4dX1KehEcoQO3JByKc,1440 +torch/include/ATen/ops/bernoulli_cuda_dispatch.h,sha256=ucwS4BbWGFLnAhZ5aBEFyPkDLorVC7k3INu3EkAIXwk,1442 +torch/include/ATen/ops/bernoulli_meta_dispatch.h,sha256=pOz3yuIbZZnP0_w5R1dMoBQQgMrTgw4rqawhVJkS91w,1180 +torch/include/ATen/ops/bernoulli_native.h,sha256=X2BYZoCeV6AG8qbHA1aSxwQou7woAT9ydOJXWAbv9C8,1722 +torch/include/ATen/ops/bernoulli_ops.h,sha256=S8DmmYaRvo1xRC6KeB6npo-kji0tts5qwXY-_Ew19co,6652 +torch/include/ATen/ops/bilinear.h,sha256=FYNT3GohTwwoKL5ODZP0HG7ZJ-TdclUxSf9Yx-o9Uqo,1075 +torch/include/ATen/ops/bilinear_compositeimplicitautograd_dispatch.h,sha256=e3YXu_hWDmorApn1kuLBqWX1KDMCsWDo6_KEkhKgXvs,1121 +torch/include/ATen/ops/bilinear_native.h,sha256=PBrU7BXA-8tvkWY1LKRfRfuFKEie3lifBOafIMQHroc,833 +torch/include/ATen/ops/bilinear_ops.h,sha256=qjaEWwoCBghOTo9Gw0kMcOgLPF8Eak-qtCFjHkIcAaM,1541 +torch/include/ATen/ops/binary_cross_entropy.h,sha256=DCF7LF1HHd19oSDts25qagv9asrXXQElT37AUCqq2ag,1986 +torch/include/ATen/ops/binary_cross_entropy_backward.h,sha256=7RCy4uoQy9PZI2gYJj539lEHPTMv6V1qZ8nm22JejbM,2341 +torch/include/ATen/ops/binary_cross_entropy_backward_cpu_dispatch.h,sha256=JSyYOMegYHYYFqkmCnBzqgT3gLI17XpeHBhWwX6K6Mk,1626 +torch/include/ATen/ops/binary_cross_entropy_backward_cuda_dispatch.h,sha256=wIr-nVSozbeGZwfBpeW34WPF6ZqNcCGR6S91OfV1Xdk,1628 +torch/include/ATen/ops/binary_cross_entropy_backward_native.h,sha256=x4spYwPmveyMQjtEcNWGVrhJYDnKj1B_PCd3ZHtyPsk,1598 +torch/include/ATen/ops/binary_cross_entropy_backward_ops.h,sha256=0xjgIB9_TGuvGTG0jb3qzZffgcpbWuZdcEhqbi53Dsw,2793 +torch/include/ATen/ops/binary_cross_entropy_cpu_dispatch.h,sha256=xxQMdCXt2GZ6D-SgSBiIsu8QIu8tCxH_re28py0ivIc,1489 +torch/include/ATen/ops/binary_cross_entropy_cuda_dispatch.h,sha256=P-ZweLmp1dWSj_VmTv0J4h9KiTwLzA40vVj66cWwXZc,1491 +torch/include/ATen/ops/binary_cross_entropy_native.h,sha256=KgYba0LJtDl7gmw7xJa2MAn1aK2o8u9CVbPVPJp-pFI,1420 +torch/include/ATen/ops/binary_cross_entropy_ops.h,sha256=yVypxWGMmZcwVs81oKaX8Za_tcr68Doq6jC0uVQTXhk,2489 +torch/include/ATen/ops/binary_cross_entropy_with_logits.h,sha256=NpkbQVY1TE5K24xXPY5Lq9TUFS_M8lhue9ufwF-2nx0,2367 +torch/include/ATen/ops/binary_cross_entropy_with_logits_compositeexplicitautograd_dispatch.h,sha256=UQQwtNrIusKed8oifuTq0C8JFJ1lshe8ASKWuPbyFCY,1719 +torch/include/ATen/ops/binary_cross_entropy_with_logits_native.h,sha256=WUGJDzPdDGnCRMD5d8M7nAEIN40etar8foL7UHgFLLE,1162 +torch/include/ATen/ops/binary_cross_entropy_with_logits_ops.h,sha256=2CZZIU74najADjMReDtGY9XurD0FBK74cfQzwNfPN8o,2877 +torch/include/ATen/ops/bincount.h,sha256=zpEwE5AOGA4kzDzu1dc8A7bKgyN45hQjPM6f53wCOSk,4525 +torch/include/ATen/ops/bincount_compositeexplicitautograd_dispatch.h,sha256=qVY8LA1oNcZrro7ZbNTyNLmaMKjibaJ2J85p0R_v8Ms,1574 +torch/include/ATen/ops/bincount_cpu_dispatch.h,sha256=PLY2VvXZPzpup8wyL-JLOQwd9ncivzLQlVqnNs1IabI,1181 +torch/include/ATen/ops/bincount_cuda_dispatch.h,sha256=rQlaxDtxTnMup732yi_5SLFsN0Rz6GB9sRi62UfCyw0,1183 +torch/include/ATen/ops/bincount_native.h,sha256=iqtTp__75jB88uAV04Acup-SkcUFitUXM9K8pS0usuU,1092 +torch/include/ATen/ops/bincount_ops.h,sha256=4X8DRBzXZK6zO65z3PtS75En1VoUfACGNcCK6h33tXQ,2269 +torch/include/ATen/ops/binomial.h,sha256=qGDK6P_X0RzjHwQoO6Cen9B3iqqERlGx7DpXC4Qv2eo,1712 +torch/include/ATen/ops/binomial_compositeexplicitautograd_dispatch.h,sha256=F6-55607MVn2ZoHz8BpOVlDgI1NZkVL1vS4DTfe6A3k,1276 +torch/include/ATen/ops/binomial_cpu_dispatch.h,sha256=ML9WXhZWcO58zDFJOXbuGJk8aLpFMyvv55rdFJRpSDs,1059 +torch/include/ATen/ops/binomial_cuda_dispatch.h,sha256=Rzlyj5xMrm1-2msFyDbdLK0wqonWUhfHGMB4e9_PfMM,1061 +torch/include/ATen/ops/binomial_native.h,sha256=1EMlINS04wgfd1TivQTJsdae4cPhSa0ZASOMrXamUwQ,1117 +torch/include/ATen/ops/binomial_ops.h,sha256=IqYFvyiwEWMPvt1X0_hpc6RMxHoF2RLSyitA3UHKnnY,2271 +torch/include/ATen/ops/bitwise_and.h,sha256=CLnDto9LHTOSQO_mfNQ4rLCw-jdiyCdIb7ujwWaTR0c,3062 +torch/include/ATen/ops/bitwise_and_compositeexplicitautograd_dispatch.h,sha256=5oAHtdrUIRMRgxbKxPsMU_QdpQXnkR8jOrpTn7PfWYM,1654 +torch/include/ATen/ops/bitwise_and_compositeexplicitautogradnonfunctional_dispatch.h,sha256=RRqH3_LwHfq38YDnF4m-8Pwgvv4W_kY3ZC5vkLX-MwY,1157 +torch/include/ATen/ops/bitwise_and_cpu_dispatch.h,sha256=CubvluiIpMCfUj3HlJVgWHipVb5xNjX6PnsLD7Ul0eY,1306 +torch/include/ATen/ops/bitwise_and_cuda_dispatch.h,sha256=30pdrrJIB6LlajhqfCamdSn9AYIHPr2PSPzowCpkX54,1308 +torch/include/ATen/ops/bitwise_and_meta.h,sha256=U8K8WmOvuYdT9S39Gzy9YDYf5fzueWCR789gHqAxvSA,859 +torch/include/ATen/ops/bitwise_and_meta_dispatch.h,sha256=JIzvnJoxHWeY0n6GQ_EWqhkJWcCcvxN5Skz5H2Pjgm4,1308 +torch/include/ATen/ops/bitwise_and_native.h,sha256=U96dGy4rnd3ND4UKMty-61sG3SvfaGF6HLa74-pXVCE,1382 +torch/include/ATen/ops/bitwise_and_ops.h,sha256=4W_wWJAAPAmPg0YyjzMJSi-4SxAAjfl0nBxmq1BKEFg,5992 +torch/include/ATen/ops/bitwise_left_shift.h,sha256=oMcxJpA2P9yuXrRjTCooE_C6LoFlH_0MbNGBj7MStdQ,3300 +torch/include/ATen/ops/bitwise_left_shift_compositeexplicitautograd_dispatch.h,sha256=AKp5xmpr7fAKQZ8__vhqh9s1Wehd684iPPzVcRDjDSg,1703 +torch/include/ATen/ops/bitwise_left_shift_compositeexplicitautogradnonfunctional_dispatch.h,sha256=1t3_VdKZMfYTFuVndXxcJ7EHhxSdlfNsnavpyvzv8Zs,1171 +torch/include/ATen/ops/bitwise_left_shift_cpu_dispatch.h,sha256=UWDDYZT-kPNxL55iUgllZtfTMKc_0LtAGs3bU3fZHSA,1334 +torch/include/ATen/ops/bitwise_left_shift_cuda_dispatch.h,sha256=79OydsP3XTYGf5AFmSpE0hhq-4uUDKdl-fLo71Squ44,1336 +torch/include/ATen/ops/bitwise_left_shift_meta.h,sha256=MULPWugoHVhjYNdFq7jsFmgdaqkIP3HradZ6FTWADUU,866 +torch/include/ATen/ops/bitwise_left_shift_meta_dispatch.h,sha256=B3Qahb2LvZP-FIkMshMfJrH-q_rO-2G2GER8Fnibnoc,1336 +torch/include/ATen/ops/bitwise_left_shift_native.h,sha256=AmYhQHzJb6JBwkh8EBTWxEDQJrKbAIlFGMLIPcnzWfM,1438 +torch/include/ATen/ops/bitwise_left_shift_ops.h,sha256=DIOuexY7sm3VkHrdPorurBKRDCNeUeE9wMWKl0pveDA,6223 +torch/include/ATen/ops/bitwise_not.h,sha256=3cQohMDKPaXDYE0JkJp8RkAlqbc9_khCNaPLrk2hqVQ,1331 +torch/include/ATen/ops/bitwise_not_compositeexplicitautogradnonfunctional_dispatch.h,sha256=U_0i_TRsnmBcy-c8o6sjQrsMG3ayYfCbUOMLx4m0Yt8,1105 +torch/include/ATen/ops/bitwise_not_cpu_dispatch.h,sha256=P86aJOqHlqGt4BcgvxMJfbIEdlfIteDMDOKIZNaUuNU,1202 +torch/include/ATen/ops/bitwise_not_cuda_dispatch.h,sha256=wenskXlaBh2-u9nvfLAbB-InMrvQUitGMJPZcM5LzrY,1204 +torch/include/ATen/ops/bitwise_not_meta.h,sha256=-kKMzUjLBcqaRZe9yZcsxz_ylVHP1hjiPmtQI8_4pws,826 +torch/include/ATen/ops/bitwise_not_meta_dispatch.h,sha256=NlWbzqh8U4bKHuiGjfKy0rLktQnCIbi9PgZka3TX5lo,1204 +torch/include/ATen/ops/bitwise_not_native.h,sha256=rX-_LvhGQ8sW_yUBNOQvAdDdfFkVn-zNFVbu6BGveiQ,865 +torch/include/ATen/ops/bitwise_not_ops.h,sha256=t9kYlpM7xHdcHcudtWbjUn714J0tgUwJQD_PLYUdOoo,2345 +torch/include/ATen/ops/bitwise_or.h,sha256=t5hFHyg1yXrz3hx3HMUka6gk3mxkOaUJDFaocZcqFiM,3034 +torch/include/ATen/ops/bitwise_or_compositeexplicitautograd_dispatch.h,sha256=dvye_yu_Ssr7ka2cU031i9MqerD0RIYxtM7rb6Eos0U,1647 +torch/include/ATen/ops/bitwise_or_compositeexplicitautogradnonfunctional_dispatch.h,sha256=6m83NqXR_FlMI9EYhP2hsQLSDe8AFHC3k6sFf1Ro1EQ,1155 +torch/include/ATen/ops/bitwise_or_cpu_dispatch.h,sha256=xroX9NVDO-bNGw74l-UeHWWQwIjzQ2cb4VZILCtK48I,1302 +torch/include/ATen/ops/bitwise_or_cuda_dispatch.h,sha256=XwYo5OPNzZMi0A-6PTLlqg6aU-f_zuFz4OimrrVChX4,1304 +torch/include/ATen/ops/bitwise_or_meta.h,sha256=t7r4cEqy1i1T8TyxVLOVmsYSGk13RK2xr9SJh-oPdWo,858 +torch/include/ATen/ops/bitwise_or_meta_dispatch.h,sha256=ViNKm8QcwlpYwSF0g2Y2GmF3RAEv0WZ-YjBf6frxAh8,1304 +torch/include/ATen/ops/bitwise_or_native.h,sha256=X1-J9UPXTLgPQYSgGxSxkijM-3jl1-6aQNNyvVotTdo,1374 +torch/include/ATen/ops/bitwise_or_ops.h,sha256=Ec9AlsHaY_UzJDwnu102BDYWTHDaatla1QyYVacz2Ks,5968 +torch/include/ATen/ops/bitwise_right_shift.h,sha256=ZSRJoCZoRYRgpghFYRdc2ynvw8S8FgVu6MNUynu7Oqs,3328 +torch/include/ATen/ops/bitwise_right_shift_compositeexplicitautograd_dispatch.h,sha256=65j4-CNKmDKmTKpvpsaP663aegfkYud7fIMTR21gjXU,1710 +torch/include/ATen/ops/bitwise_right_shift_compositeexplicitautogradnonfunctional_dispatch.h,sha256=6cKQhbj-sKPCs6dsGJKWjHel-hA6jF01IEU2KXt-n2A,1173 +torch/include/ATen/ops/bitwise_right_shift_cpu_dispatch.h,sha256=i5g1QTRgkgxhL2SknTA5MOKDpCCJnkzp6LuJuAl5vao,1338 +torch/include/ATen/ops/bitwise_right_shift_cuda_dispatch.h,sha256=fEmaX4tcc5fb_zfE_G_ApB5FgDmfr7mCh8kznfxGQxs,1340 +torch/include/ATen/ops/bitwise_right_shift_meta.h,sha256=VZv8tNQCJepr_JI6ErlaXPUOtQyqyuHTx8kLAU94as0,867 +torch/include/ATen/ops/bitwise_right_shift_meta_dispatch.h,sha256=RP3QitT4gOr9VZ82o2R0kx285mWCyNR6Ymp42-sSmYs,1340 +torch/include/ATen/ops/bitwise_right_shift_native.h,sha256=j6qgg5CKTZ8M_vFc64mae_DoVm_hcJfri-w9GUvGsMc,1446 +torch/include/ATen/ops/bitwise_right_shift_ops.h,sha256=QYuDcJsQQPGKQ66FB2cQWf2TGAH-hvKnikXaF1Q_nFg,6247 +torch/include/ATen/ops/bitwise_xor.h,sha256=AqyGL6wXyy2ugK8oDDR-OVRlAlrg9FiO54i8RCcayjY,3062 +torch/include/ATen/ops/bitwise_xor_compositeexplicitautograd_dispatch.h,sha256=N6ShDRUxa5WrdUCAoqKYSWkbflXXFGFnAZwl6B2GsrE,1654 +torch/include/ATen/ops/bitwise_xor_compositeexplicitautogradnonfunctional_dispatch.h,sha256=38eTFTo0RP6Sj7Em-5-cxzOykWaYfmWOgnxF8uOkqvY,1157 +torch/include/ATen/ops/bitwise_xor_cpu_dispatch.h,sha256=kPj6XBUAhHpK01y98m4fB7hODoyFGgR4gSPoZsHRVh4,1306 +torch/include/ATen/ops/bitwise_xor_cuda_dispatch.h,sha256=2jJSLVWY0UgT7nYG3hXZcaE49KXBeQ43HyIZiqAOb8s,1308 +torch/include/ATen/ops/bitwise_xor_meta.h,sha256=4JR5wmov_3T8yyuoelc_Y927wRZVq4c7dego6s5P9rI,859 +torch/include/ATen/ops/bitwise_xor_meta_dispatch.h,sha256=o4doeAJae5xyHDJaCJxmM2LGW4TsAgCuAihFAwdYzEw,1308 +torch/include/ATen/ops/bitwise_xor_native.h,sha256=JTZ548ZF52GpEkAiIoJmjB3H4DoeaBT73M5OBfmaU7M,1382 +torch/include/ATen/ops/bitwise_xor_ops.h,sha256=MnweymakQHl3y9b0RU7X1aO8LGl7JTJFqINA43pcu9A,5992 +torch/include/ATen/ops/blackman_window.h,sha256=PJ4BiS7fkTgWl0apKCorR4iwPBkShSCs0QEFt5kexFs,3671 +torch/include/ATen/ops/blackman_window_compositeexplicitautograd_dispatch.h,sha256=aPYzSx0LmDokm3rPloJblKuSdXW-sGOpFDR5HrKd2LA,1960 +torch/include/ATen/ops/blackman_window_native.h,sha256=FNQyxSNQSBVpAh5irIOeAa-izJXkWs3PgKkKBo6YU1g,1321 +torch/include/ATen/ops/blackman_window_ops.h,sha256=-iUzN9oqi0AfgJo6L-vi0glRSaN4IH_n4ZfWcRe0Jb8,4110 +torch/include/ATen/ops/block_diag.h,sha256=Cep1xA7-_a9eLh5NrLfoSlgXpaNHsJDcHIL-6kd12nE,1342 +torch/include/ATen/ops/block_diag_compositeexplicitautograd_dispatch.h,sha256=ZYMRhzAt-9RVShw9E_oW5sfjANvNj0-7HQqrrOBZ-ng,1184 +torch/include/ATen/ops/block_diag_native.h,sha256=rit_2uwtA7HJ_ZAyp1aRuNPVZnN5VIvIXUvLBwaV_Ks,814 +torch/include/ATen/ops/block_diag_ops.h,sha256=JMzfvH6Nbr51898ww9jFzdlxng6FJTcE34ajAI3KZyQ,1823 +torch/include/ATen/ops/bmm.h,sha256=YE5pesQ8-LICeGtu5rwUCKfZwJrQe0KLVYULDNltfrM,2226 +torch/include/ATen/ops/bmm_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ewnx2XAt0IcPK_wkWKKtmepuyJ2aYwIqVbdMneXRpYc,1066 +torch/include/ATen/ops/bmm_cpu_dispatch.h,sha256=H_7IrlmbZpy0GEbX5D2ROi7gSRuqB5AUMDN3O_yZH0o,1197 +torch/include/ATen/ops/bmm_cuda_dispatch.h,sha256=qpStrLwVgPekY8YmvyzPVFm6gGQbH7u4-TdxZ96Uhhg,1554 +torch/include/ATen/ops/bmm_meta.h,sha256=7vvi-eZOa_KSehwvGe2JoCtBORXG2OumHWSTDMDbwoc,843 +torch/include/ATen/ops/bmm_meta_dispatch.h,sha256=bA6Lf-9SSmaQA8NmTSZqHB6bNhYhSV37MpZsE8yJHI0,1199 +torch/include/ATen/ops/bmm_native.h,sha256=3jbVY0aFvauu3JwXCapJCqtc28pf6rm7Y7CUEeHNAXc,1972 +torch/include/ATen/ops/bmm_ops.h,sha256=y0VtT8Uvpj2vsyQ0C2qU0095ePztLntxtK5yZU3x0nA,3402 +torch/include/ATen/ops/broadcast_tensors.h,sha256=JoP7HiuHIYoyXqbDxyjAmJ8Q857aaRJIVR5lEGq2Pyw,959 +torch/include/ATen/ops/broadcast_tensors_compositeimplicitautograd_dispatch.h,sha256=OXMeBBMcXQHsk-Ft353kmJQnuYteb3FxJ0zexnHBp44,1043 +torch/include/ATen/ops/broadcast_tensors_native.h,sha256=q0T5QQa0bFlY6CHaNlkQ22Br6E0vvykNQ-kErS7XY8k,755 +torch/include/ATen/ops/broadcast_tensors_ops.h,sha256=dm2WeOnCxxmONgbID734Lk8PK4qz67MVNAiBsY93qLA,1290 +torch/include/ATen/ops/broadcast_to.h,sha256=xUAhNiHus5FE66lcTbtap3rtky1qIwBtDZ1z4H_iPs8,1705 +torch/include/ATen/ops/broadcast_to_compositeimplicitautograd_dispatch.h,sha256=jMK3yT5Q-9W3CXEKHBSW_oTUraOWY849XK44BqMrD3E,1139 +torch/include/ATen/ops/broadcast_to_native.h,sha256=KDSLk201EBwRnexlAVYjm0XrpsofzAS1-rsg2jmbaGU,769 +torch/include/ATen/ops/broadcast_to_ops.h,sha256=H8AQD-XoKRg2aFd7rAnaXzRWSjSgr5tQjQcV0R4g0Qw,1323 +torch/include/ATen/ops/bucketize.h,sha256=CrAisMsWIsQPwD999GLAFf-r2yR5AX-gGx1kCj1LSXA,2889 +torch/include/ATen/ops/bucketize_compositeexplicitautograd_dispatch.h,sha256=8AlSkC7dBTgN1hyAKqwHvsjYII_msXh0auzx2ba_R4E,1257 +torch/include/ATen/ops/bucketize_cpu_dispatch.h,sha256=bU8i1pR8phE1FOITliaAwa_PmLVq4zqAMxMSqTAOYF4,1469 +torch/include/ATen/ops/bucketize_cuda_dispatch.h,sha256=wZHCh8wert8awZey9TbYhR5KPxwQRD_qGJSzmlxAzJo,1471 +torch/include/ATen/ops/bucketize_native.h,sha256=OCaJ_Xq6Y6kszqdO4-TT9ipjaOq7QxD-7PVVU4dLPmk,1642 +torch/include/ATen/ops/bucketize_ops.h,sha256=5CVTumNf05JifGSCsVRUdHTG5abJVg7pVnXEf7Y4sPo,3851 +torch/include/ATen/ops/can_cast.h,sha256=71vNOsb8uJtHpnn8w4U2gWVHXsOmW3TzP70vMjqiHw8,932 +torch/include/ATen/ops/can_cast_compositeimplicitautograd_dispatch.h,sha256=xwwIEuliU09wyo5A-m2BKy7rWrZ8xiXX-B1__ccpa7E,1030 +torch/include/ATen/ops/can_cast_native.h,sha256=AgwnlPvwSosVqt49JSHq2G0mmxKBinXrV9-IftPYESY,742 +torch/include/ATen/ops/can_cast_ops.h,sha256=xIXV1rmHMZXczKGXp_fXbXhVyB1QBda5u-GTStDu_Y4,1261 +torch/include/ATen/ops/cartesian_prod.h,sha256=M5-Q2uCjPP-0X7ag8A8Ps2qX2uWH2I4hv5y7A_o80hE,930 +torch/include/ATen/ops/cartesian_prod_compositeimplicitautograd_dispatch.h,sha256=vd3MNkCWCabbAC4Cdj1--ThVG8Uz6Zn8Rt0VWXmCYe8,1025 +torch/include/ATen/ops/cartesian_prod_native.h,sha256=mmBRBCffrZfT7yuk1DyKe3FH0UFaozFDG5JAS4Ga7Wc,737 +torch/include/ATen/ops/cartesian_prod_ops.h,sha256=j90wDadDEV--q9-S1sp7zyafjew60LESBHBfmK3Y95Y,1234 +torch/include/ATen/ops/cat.h,sha256=nRl_af61gCbEfIyWfX6Y-klna6QU4RHRKUsnJ9c4uIc,2059 +torch/include/ATen/ops/cat_compositeexplicitautogradnonfunctional_dispatch.h,sha256=alQw12v-X-6Hi696dM_nBlgJT1uj6eJnIwG_vY-rqDM,1067 +torch/include/ATen/ops/cat_compositeimplicitautograd_dispatch.h,sha256=pZFMdtlLNJOERvwAFai56LM2KeEJIxUT4ChLUd1W20Q,1214 +torch/include/ATen/ops/cat_cpu_dispatch.h,sha256=t3bHA5XTrHVOZ9zkq-oLJPHcp5I9diRfHQT6Zs_kZ6o,1198 +torch/include/ATen/ops/cat_cuda_dispatch.h,sha256=4SEMu4cZhcLSm3BSz8jTQ1HvEc8ogEPhKzkIvBM2q7E,1200 +torch/include/ATen/ops/cat_meta.h,sha256=VlMv2ne_DNklqcrQGoRQTDQWKtn1AEDQKXuhF-FeggQ,5060 +torch/include/ATen/ops/cat_meta_dispatch.h,sha256=KPwkwoPyp7FKJvBT4cK12Cly9zvleXTJQJ3EKb31svg,1200 +torch/include/ATen/ops/cat_native.h,sha256=NwRR27UW8qv2dTRhpz5z_Msl_LAKUPKHzcCOJQR6oWs,1805 +torch/include/ATen/ops/cat_ops.h,sha256=BoUzPbrjxNw3HhYU4yy-mZuWuOQND9W8NEPN_i55VDE,3170 +torch/include/ATen/ops/cauchy.h,sha256=BdUIyMufrVKZdcdtycSa5c9yxR7pOq8TjyLP1Q3i0Cc,1787 +torch/include/ATen/ops/cauchy_compositeexplicitautograd_dispatch.h,sha256=MqWVRW9raVK2Px8H7X2QoPtllCZwieulT5zG8OOF8Fw,1426 +torch/include/ATen/ops/cauchy_cpu_dispatch.h,sha256=Lq83Czdk3o1OptTQIevDfVfUF12ILII85d--DAeozQw,1061 +torch/include/ATen/ops/cauchy_cuda_dispatch.h,sha256=p27VKHIPdJ4bO3aaPj55O_-ARbPny9DWmiS46I-xkT0,1063 +torch/include/ATen/ops/cauchy_meta_dispatch.h,sha256=4nxXRVBGeTk_9C4WebFxgl6ukrnojYBjYbXB0f--tbY,1063 +torch/include/ATen/ops/cauchy_native.h,sha256=tWtRqn3-8LBgzbjZfkN9MGdTidUqXy6bGR9FnOyA1Ho,1110 +torch/include/ATen/ops/cauchy_ops.h,sha256=3x6VNCI7BVbuYnPxl4zqzhI_z_VXlR273MVzE7JEpC0,3050 +torch/include/ATen/ops/ccol_indices.h,sha256=MrORAXn8d13FtSy6pNyJrlwrOeybly0_xSA1sGH1OpY,763 +torch/include/ATen/ops/ccol_indices_compositeexplicitautograd_dispatch.h,sha256=cDDpnCkXfXjVfC_ozwJkFNDgpZcOs19C5OWe82xUPno,1024 +torch/include/ATen/ops/ccol_indices_copy.h,sha256=17AXM43zKwwmSuw0oAbMTGDKbDgq73kMXZgnFDrkFF4,1391 +torch/include/ATen/ops/ccol_indices_copy_compositeexplicitautograd_dispatch.h,sha256=iZ8H30ey1RQ9XyY0zbF-Y3CdVZP8oHSQwabAZjAeHZQ,1143 +torch/include/ATen/ops/ccol_indices_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=WX3Aj0RNJpz1tPdGsplgKmiqpO1N6V_okzh2LpW68_M,1055 +torch/include/ATen/ops/ccol_indices_copy_native.h,sha256=jWiE8EgNtDtJkLewzxo53wHTG9oGkudcAuM5FtGUUKU,830 +torch/include/ATen/ops/ccol_indices_copy_ops.h,sha256=ne8hlbcD2XuXFCA5D21a5KW6KjF3xxaBEYewX1sIOek,1867 +torch/include/ATen/ops/ccol_indices_native.h,sha256=JF974sn-ATcW-hgdihkDhoSgCSjK_P5MX3K5W6Ks2Dc,815 +torch/include/ATen/ops/ccol_indices_ops.h,sha256=at44EkIDqFfngxWzvbo7_3EzRD-tg4qRcdNOXEHaVt4,1235 +torch/include/ATen/ops/cdist.h,sha256=iL5d45RRMRSkOQTRE4ZDwBT6BRNfUsWlpN8VQmaa9W4,1037 +torch/include/ATen/ops/cdist_compositeimplicitautograd_dispatch.h,sha256=pBXsj8yog3cQ9cfkrKLFFqg6HiQ8nj1I5JKrsJ2E5Rg,1104 +torch/include/ATen/ops/cdist_native.h,sha256=flat0EyF91YkuiDS38M5UZZPK8qcwN4TInyrxrVkZ5U,816 +torch/include/ATen/ops/cdist_ops.h,sha256=m4G5nSXG6nGcyTMb5JXQ-ykT5LdHG6Ke0Bo3RH4B1_8,1446 +torch/include/ATen/ops/ceil.h,sha256=MZlrN_eZCy8l01wqQWMtGh-6oTCXJF-JH57tFdSgzIs,1397 +torch/include/ATen/ops/ceil_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ESqG0tWQwMmnZ3VPWGke2DPIYIW8TFHIQlP8CrOzoBI,1091 +torch/include/ATen/ops/ceil_cpu_dispatch.h,sha256=XcPCA7xm67Jgp7aKc8RSuSKvLdXPJdnY81ZfYoBDK6U,1174 +torch/include/ATen/ops/ceil_cuda_dispatch.h,sha256=8sLQ4KG3gKTn2SMs6fR2oKSiryfD4mGiEl97Y6Yr1HM,1176 +torch/include/ATen/ops/ceil_meta.h,sha256=iUeIRF6NOMnJU8tC7mBfd1ptLryj0Ge0YF90mLXaV9c,819 +torch/include/ATen/ops/ceil_meta_dispatch.h,sha256=LS39p9XGcg8ZEtX_HwFbqH8MmXgVeurZJNrfkGNENyk,1176 +torch/include/ATen/ops/ceil_native.h,sha256=8V3OfSrNH2aPRls_WtJn4E8G18pzj6l_KdT_SQ3sptg,1252 +torch/include/ATen/ops/ceil_ops.h,sha256=rH-B9IMGmb8ZTyd-T-Ezxem7XUCacUOY4zukuApKVa8,2282 +torch/include/ATen/ops/celu.h,sha256=YI9cBTvXttj5BnX1b3eRKOS4z0J0QMBWKa4yozJNG3o,1613 +torch/include/ATen/ops/celu_compositeexplicitautograd_dispatch.h,sha256=e6fi2FW8mfJSISx9q9IYcnFXR7xyjNShRgczwIl85lU,1334 +torch/include/ATen/ops/celu_native.h,sha256=4DIKLYcIq-5U7kBZH5FdvS1lToV53FCKiE8IFljyg78,939 +torch/include/ATen/ops/celu_ops.h,sha256=E-MukROShG-br6H3QkH8pqDitFUVRmgaIwbPcksPJ2k,2552 +torch/include/ATen/ops/chain_matmul.h,sha256=zyyGGVzI2eN_iHIAPrpU3cFYApy4PP6r4kVXR1YJFSA,1371 +torch/include/ATen/ops/chain_matmul_compositeimplicitautograd_dispatch.h,sha256=LZDPowO23rjc0c1-zkARFLHnq8TZnmK5YzKp_SPESmY,1193 +torch/include/ATen/ops/chain_matmul_native.h,sha256=rMlxporPAKKMdkMI7ohRLq-MZlvqrJNGDNmdfJRnHLc,820 +torch/include/ATen/ops/chain_matmul_ops.h,sha256=HoBLWIluXSfDr7c508vxt0pbzLxdwcAcPe9rITZ-uzM,1841 +torch/include/ATen/ops/chalf.h,sha256=73gO9KgAAnuO-m_4QZfBXC7EpTagK4ezsyObERuXS4o,756 +torch/include/ATen/ops/chalf_compositeimplicitautograd_dispatch.h,sha256=ffBzHy3md47Touge699tOB8jB1odq_fwGQR7TTOm3CU,1081 +torch/include/ATen/ops/chalf_native.h,sha256=YOYoP1QPha3ctTVzyByWHpYvcMca7HYuA7QsWPJisbo,793 +torch/include/ATen/ops/chalf_ops.h,sha256=QnERIcjkeMe80gWW3UAaIJLJFvGTxyNlaiXZm0qPNYk,1378 +torch/include/ATen/ops/channel_shuffle.h,sha256=nxFwNSgOS_oGczkr7i5tPzbWSMTQWhJsQugHvMw7v-E,3820 +torch/include/ATen/ops/channel_shuffle_compositeexplicitautograd_dispatch.h,sha256=_zPnnmChxjcyqxS_D4pQIQ2G7mjn2ZlPAEPc5-DwjTM,1400 +torch/include/ATen/ops/channel_shuffle_cpu_dispatch.h,sha256=7DXG7HBcAQgWgPV7O0ock0eQEsuXdPhtiwSSGg2GSbU,1089 +torch/include/ATen/ops/channel_shuffle_cuda_dispatch.h,sha256=IazzCupurAWVn-TO-RJsPzer3J2FohryAIf3gh4YhOI,1091 +torch/include/ATen/ops/channel_shuffle_native.h,sha256=T6ybPlxvYtuK6GU_lPCskPxhVVHgTF4jU6izLumjuIU,962 +torch/include/ATen/ops/channel_shuffle_ops.h,sha256=2RZPliQ-ZeEOcizh8KEaDDRfQhfNcydx_3tQzb5ks4Q,1991 +torch/include/ATen/ops/cholesky.h,sha256=Uaij77awWE2Ot4geSe-QcHgyH05pkp9pP3er_EBbHm0,1424 +torch/include/ATen/ops/cholesky_cpu_dispatch.h,sha256=4sqoTMBPLYatr7pJyhRnoA6bCukT0GifU1SyfE0dHSQ,1185 +torch/include/ATen/ops/cholesky_cuda_dispatch.h,sha256=ybDa9JiGC8o1lNcHsRWvA2RIIxrniNSlgCKdmPW6krY,1187 +torch/include/ATen/ops/cholesky_inverse.h,sha256=vQctG1lXwosH6Zue4rG_axx8Dnus5A6tsyuqPRzES8g,1504 +torch/include/ATen/ops/cholesky_inverse_cpu_dispatch.h,sha256=K6-_o3HxbnZ-xZd7qmKhpr_-PKaK-dxONirdojaj38U,1209 +torch/include/ATen/ops/cholesky_inverse_cuda_dispatch.h,sha256=t1QWaKsQLaECeDi9b9XC4dRhAjQ1V4ut6JA-oVDSAm4,1211 +torch/include/ATen/ops/cholesky_inverse_native.h,sha256=vg-HrtoYqPxfNYZ1YKRWMx9mbcwThGQ5PwTZ7DydjbA,858 +torch/include/ATen/ops/cholesky_inverse_ops.h,sha256=_B8IvLKeTQKb3L5hixxtPFCBRj1eoa7nnGaEz5H4yy8,1957 +torch/include/ATen/ops/cholesky_native.h,sha256=cRO2WioemDedhUwnyY1tpwgUqzgZUJ-LvbC_QD_qH0M,842 +torch/include/ATen/ops/cholesky_ops.h,sha256=kHUGK6p30wpjYid1E-arHZMeep63hvJMTXdZ6HLveNc,1909 +torch/include/ATen/ops/cholesky_solve.h,sha256=NxFAUNS0fwuBwiK0l96D5uAFJc6EO-unoGvZL9Rd3jo,1634 +torch/include/ATen/ops/cholesky_solve_compositeexplicitautograd_dispatch.h,sha256=bWQ2HdHF6fm8dU0BeOxtZdTmwuDsWJShRpXkD8zv21g,1328 +torch/include/ATen/ops/cholesky_solve_native.h,sha256=sF3hXLhlwzX1DcCIn4ci52JFuHayrINlQRJKUChGLAM,908 +torch/include/ATen/ops/cholesky_solve_ops.h,sha256=FOc6SIqOTkge-qjztG8h_NowBx3XMhlIBQhI1uoUz04,2123 +torch/include/ATen/ops/choose_qparams_optimized.h,sha256=9gYjgTdh05Ja3Z47oBpaZSsAROZAfoABhRiShk4Zep4,1149 +torch/include/ATen/ops/choose_qparams_optimized_compositeimplicitautograd_dispatch.h,sha256=jzMyOWmFIVMzBYvpHS0UBsVpssmK9IzQdOy5v6OXv88,1126 +torch/include/ATen/ops/choose_qparams_optimized_native.h,sha256=oyMpO6bW5uEQI-YGpGHcqaDFWeZy4-JqyEeFfrmcPh4,838 +torch/include/ATen/ops/choose_qparams_optimized_ops.h,sha256=L_s7Etg6S2ajIJcsnrZHkMf_ZLeKmhipiA_2qk6cZFY,1567 +torch/include/ATen/ops/chunk.h,sha256=di92a3s7psLitFXkeDnjxIeeUUQxbjqpCfho9HHbP44,982 +torch/include/ATen/ops/chunk_compositeimplicitautograd_dispatch.h,sha256=kBhiQa73y2mOFJgygD2kNvVFcO2xQx-TzXRfUW1vRAA,1063 +torch/include/ATen/ops/chunk_native.h,sha256=w1NwgHvt4j-bVWJBQnPz2NDP7a_SPyIQaKJNpNan67E,888 +torch/include/ATen/ops/chunk_ops.h,sha256=I0tHF_LYU3PRK4QQ_XwofsxhrZ7mLUDMQbGuRKG0xpA,1365 +torch/include/ATen/ops/clamp.h,sha256=UjXuzQ0rQ5D9rX0eOACfToNSlFfzc0XIMs23eua1DkM,3208 +torch/include/ATen/ops/clamp_compositeexplicitautogradnonfunctional_dispatch.h,sha256=n6YpO8bwivt-zEz-fIB8mUlqPQe0Y2nO7eJfYO03xLk,1566 +torch/include/ATen/ops/clamp_cpu_dispatch.h,sha256=PfM0xYbc1KvJ_02tJV5di2Z90d468QWCwInaoWNSGAE,2155 +torch/include/ATen/ops/clamp_cuda_dispatch.h,sha256=PQJ_ue1XnHv_apNUTV6haBEirNZbM6HxkgyUejRXnOQ,2157 +torch/include/ATen/ops/clamp_max.h,sha256=o3wVpB2NXwKWEQK2t6eeQ_-3JVdtyzZs8OjF77kNCNA,2549 +torch/include/ATen/ops/clamp_max_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ie0z3Tye2m4vSMwWGngxVCd3y1HYipBORB0Oc75o5b4,1308 +torch/include/ATen/ops/clamp_max_cpu_dispatch.h,sha256=At_nCxGYaphAWyp_w-Tmsk8bJjeXUGcAv3REtTidFyI,1660 +torch/include/ATen/ops/clamp_max_cuda_dispatch.h,sha256=DVRXnKqGD3zERN8HnZPwmMLgy_7MciGlTMpHTM00jQc,1662 +torch/include/ATen/ops/clamp_max_meta.h,sha256=9hHVBEOfmV_ij_ebRX8xIOhMZ7vrBXJFbpJC5fiRJUE,992 +torch/include/ATen/ops/clamp_max_meta_dispatch.h,sha256=EDbmTNOE2EGTzQBXWxnwzhkxnmfUyGfn44FCt6Hay9M,1662 +torch/include/ATen/ops/clamp_max_native.h,sha256=dtVsGAXFhlmUj6Ygzim_D2wFjgcNr7j27e8zSvjN6Ik,1068 +torch/include/ATen/ops/clamp_max_ops.h,sha256=aoyJor9rPU_AAubt8zg1xITFEZ1PQr_2VU5uDlIDlbo,4484 +torch/include/ATen/ops/clamp_meta.h,sha256=Ds3LVlaysfz2ifWHtz9CeVlbtJjKIUTw7V8YxLFCeMk,1044 +torch/include/ATen/ops/clamp_meta_dispatch.h,sha256=0SejF_mlB6RvOOF5fw4FAIE6VVURJ-UqPwyAfdIPOu8,2157 +torch/include/ATen/ops/clamp_min.h,sha256=BnxkOAUPJ22OCQxJXd7JiqIj69_oPRVR100zX_cD7S4,2549 +torch/include/ATen/ops/clamp_min_compositeexplicitautogradnonfunctional_dispatch.h,sha256=YTVWNqV2kSnRehIgJgudxzq0bCXdExG4ygEy5oHphBQ,1308 +torch/include/ATen/ops/clamp_min_cpu_dispatch.h,sha256=TuWQLiy1LK56pDowJMELZGTlKfrrSFncswIM5sEHnd0,1660 +torch/include/ATen/ops/clamp_min_cuda_dispatch.h,sha256=3vXAxTpXxhzYhEBWOlQq60Mhz3DtRBiY5qx22SWHKg0,1662 +torch/include/ATen/ops/clamp_min_meta.h,sha256=U2F4AF2a9goEe0GdBUX98HAJLra1YUcM8qNZd-BdVDg,992 +torch/include/ATen/ops/clamp_min_meta_dispatch.h,sha256=AXd_XLt2rSTQymwAUXpUzQPOdcl9v9qsMovL0TCDIJE,1662 +torch/include/ATen/ops/clamp_min_native.h,sha256=GOGe1C_U2EPdS5vVcbQGXg77kTpYcFP53_tdRXC7vzc,1068 +torch/include/ATen/ops/clamp_min_ops.h,sha256=FL5-taE7AEWGR-yxjBROrNhxE_5pCVFyRXHp1Ftclc0,4484 +torch/include/ATen/ops/clamp_native.h,sha256=mEKETr2TkqT2rEaLlc7uLaVEvELBoI7Dr_P1coUxwU4,1287 +torch/include/ATen/ops/clamp_ops.h,sha256=H3sSF5XmbjkphNtc3my4rYbCGKbySr4YG_5RiFSw0R8,5576 +torch/include/ATen/ops/clip.h,sha256=3V16PBHRjwbL9wzdqltDmelX5wGauKlWQh3O6VUwckw,3183 +torch/include/ATen/ops/clip_compositeimplicitautograd_dispatch.h,sha256=oxTmHj_LAts0BA-AeXK-I5eaVjKHKR7_YVAPifoxRg4,2191 +torch/include/ATen/ops/clip_native.h,sha256=g02MmUoCVfZm9yT_tgU4fFFuhURIkYGc56sl1xanMwE,1594 +torch/include/ATen/ops/clip_ops.h,sha256=BMe1Eug03Rd_tT5LNVHWLUhsYJnGelGggU8BVmpjZtk,5558 +torch/include/ATen/ops/clone.h,sha256=Je3kkQsjxF8s81wb04JH5JdTuHyEyNiEm5WVIF83lwY,1598 +torch/include/ATen/ops/clone_compositeexplicitautograd_dispatch.h,sha256=rsm9snOz_SL9yq2hyPNB5_yAZfy9uubUmtpaQYjo-l0,1349 +torch/include/ATen/ops/clone_native.h,sha256=E4-AIzJAnFDZpb3BjJM_OarNr98sONJCB8zEz9uuG30,1553 +torch/include/ATen/ops/clone_ops.h,sha256=ACo6tfe7mU5hFS10nsTEH-ncIFl72a3yoMfDy7f_BWQ,2132 +torch/include/ATen/ops/coalesce.h,sha256=HUkBjCZQkbHlMqAEYxB10EAA23R78zYeuV1PNzXMmo4,759 +torch/include/ATen/ops/coalesce_compositeimplicitautograd_dispatch.h,sha256=6Fj3vK8f-ZEBSu5coJDVDp_WBAGncjA3MLz8zhwnzI4,1020 +torch/include/ATen/ops/coalesce_native.h,sha256=yNALkZe7lLS0KsWXyhWjZ5QH3cQequU5QqccUw-pkKc,732 +torch/include/ATen/ops/coalesce_ops.h,sha256=Auzn3fDPZT_VNhDZXrCO2RyfKUF9N2PcSAKFdqh3GY0,1223 +torch/include/ATen/ops/col2im.h,sha256=XUy5OcXDjvF4ksnYcAxHnm2j0cyo9-PgE_f92tJAiPY,6097 +torch/include/ATen/ops/col2im_cpu_dispatch.h,sha256=ACKydLEcP4Bs1gr0xyGkYJku9YYXKZh-j8uXG-EUohY,2173 +torch/include/ATen/ops/col2im_cuda_dispatch.h,sha256=8RPo9Y6xqy1jF7b5o179NQfTpjt-eaGUh4DWhz--6Jg,2175 +torch/include/ATen/ops/col2im_native.h,sha256=D2HQeIPCJdCp2sKqGEeXVRfdzre7V_Dq9biCP2JG45w,1490 +torch/include/ATen/ops/col2im_ops.h,sha256=Fg8Tn2LvYSSg3v-p1acri_GpILMVkwPqFWcDCMdcP8Q,2709 +torch/include/ATen/ops/col_indices.h,sha256=UhTewJI5Tpahhah9dPtO7BU67xBGt7LIkK59l4AaPZI,762 +torch/include/ATen/ops/col_indices_compositeexplicitautograd_dispatch.h,sha256=LS3kwnTww5m3Efh1zoVKFlQXwdmEZDfKR7I27Uac150,1023 +torch/include/ATen/ops/col_indices_copy.h,sha256=qbGGfEICE83iwHB8puIyOP42da9aJLGz9LRH_qmd2MI,1381 +torch/include/ATen/ops/col_indices_copy_compositeexplicitautograd_dispatch.h,sha256=KpAD3-aDYpp0EbXjcajvLqRFOFrbiDHLtbp173yN04Q,1141 +torch/include/ATen/ops/col_indices_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=kyvz9dAp3ANmLT_wZ7U8HFop_HMNhWEwTTTOlK0NmJU,1054 +torch/include/ATen/ops/col_indices_copy_native.h,sha256=jG5QInHHTj01b5Czda3_Hmy--j0lcq9EawEm5Mc0GZw,828 +torch/include/ATen/ops/col_indices_copy_ops.h,sha256=uITUMZ_j1lJ0_ae2EU9yhcC6DagGLwH5dMm0PwcWrH4,1861 +torch/include/ATen/ops/col_indices_native.h,sha256=tfghsI3m7QDgmCfVBJW8qLXwmWnaMZduzq_tYDtpKAU,813 +torch/include/ATen/ops/col_indices_ops.h,sha256=j0PeK4q8gFkUj4u-0l1QSlR5Fj1wsXtedcIhUNo04BM,1232 +torch/include/ATen/ops/column_stack.h,sha256=KDwlnOJFAZu7H1oMOPBr0ODH8fSKCmtfwJi6Xqu_EJg,1362 +torch/include/ATen/ops/column_stack_compositeimplicitautograd_dispatch.h,sha256=Vyun0VsVMi824UCKMMajcF6N2T2y8xEl8iYDcWDXQ5E,1190 +torch/include/ATen/ops/column_stack_native.h,sha256=BaW-d9f7JNXtnMfpYwhMdpyTqlrr-tC2dGGevCFHqiY,818 +torch/include/ATen/ops/column_stack_ops.h,sha256=eiVJjA_LdFlQYYxVRpp5SAGb0oZEFOPnjNTCBcAXsaQ,1835 +torch/include/ATen/ops/combinations.h,sha256=ro260zW4yd2SzHmHruOTgoay7DpW9u9K7-fD5GvMVs4,1016 +torch/include/ATen/ops/combinations_compositeimplicitautograd_dispatch.h,sha256=Rw5CyWsTL-m7flEDcRaS8BRQ0JklxENDIwU5V-ZpAbQ,1066 +torch/include/ATen/ops/combinations_native.h,sha256=R9nmXYkuCtUn9glitGGX5MSkU2Ot91vIwov39_aIkLg,778 +torch/include/ATen/ops/combinations_ops.h,sha256=p6GZTECy97kWxlSuyCk6mTecSovhW2AKmWQRJ1ncz5g,1350 +torch/include/ATen/ops/complex.h,sha256=DXkvUGeeVO_Z_Js2XQbFOt_w8LKwWHEF5z789oj_SIQ,1423 +torch/include/ATen/ops/complex_compositeexplicitautograd_dispatch.h,sha256=1dLCA_gDslLar-c3caO9G5ZEsvptzhwAFSpCxdiIq8M,1044 +torch/include/ATen/ops/complex_cpu_dispatch.h,sha256=6B04wdf3hbB0r2ojUWNR9sN5uIibNcYrpJj4Y2LEjkg,1129 +torch/include/ATen/ops/complex_cuda_dispatch.h,sha256=QCWuckNTjNDS0gttroxxLeXtFMLwAXevmTNHJewh_ns,1131 +torch/include/ATen/ops/complex_native.h,sha256=g4fyHkZ3AL7DgNSNl5twLCJevcr7FHYrIg1XG0XJXiQ,860 +torch/include/ATen/ops/complex_ops.h,sha256=XCIxK4uKqqPWgk9A-nhb8lu8zgo1xrjT3ippEfS0_-8,1973 +torch/include/ATen/ops/concat.h,sha256=4v_e0sFmmk7uIoe9W2N-wotSLRUkf0RZaqhLArDe-g4,2080 +torch/include/ATen/ops/concat_compositeimplicitautograd_dispatch.h,sha256=CJlt-qaz3Oz0z8agLvnDfpk7jMb015-DFwF_OnPEsLk,1474 +torch/include/ATen/ops/concat_native.h,sha256=XUHM-P5zogIM1J03CAmsphRkk9WDeVkSasBZaKvD7fc,998 +torch/include/ATen/ops/concat_ops.h,sha256=KUrbZ3dmCW3omeERK2wdHGX1eSpPl6QpTKW9pcjNd4Y,3134 +torch/include/ATen/ops/concatenate.h,sha256=sLzqg5d2Jn1H-tI0GpOw6jYH8es-DBdpmCPCqlh1xck,2175 +torch/include/ATen/ops/concatenate_compositeimplicitautograd_dispatch.h,sha256=MwaIQqoOt7WrfcHLrUU5s54t9M7v86DgrUCEM2upRLM,1504 +torch/include/ATen/ops/concatenate_native.h,sha256=I6OCO3zr5e4sZPkP8xD3Pxoy6fE7xJTsGSn3ixWm-eg,1018 +torch/include/ATen/ops/concatenate_ops.h,sha256=B64nL8t96vzEc9BIQpk30QZDdMjJQKxkTDKmDzAGHkY,3194 +torch/include/ATen/ops/conj.h,sha256=z5HYNIiMQFMMuPuSqQT3hNTA35iCEdiObeJkl8isvQ0,900 +torch/include/ATen/ops/conj_compositeimplicitautograd_dispatch.h,sha256=xaucZb3n9arlmBnnY8PS0G3m-DjQ_jM9QTk-ICKiICQ,1016 +torch/include/ATen/ops/conj_native.h,sha256=lKC1hFYihoj38yKKuTuvZAH1Smr7RK5v_wzRRI7rz2w,728 +torch/include/ATen/ops/conj_ops.h,sha256=dPIJq3gmsnkhqRFOotc2nRHDGMUdRG7MsDv7OznDAps,1211 +torch/include/ATen/ops/conj_physical.h,sha256=UB4Rsu_3y73FI7AFwvtADdzEu1H0pXZ6aX1vVeJAGfU,1514 +torch/include/ATen/ops/conj_physical_compositeexplicitautograd_dispatch.h,sha256=wBzpRJexVfkRRQK6PJ-dAqK3FJsCcXq1BB-WLgy7w98,1022 +torch/include/ATen/ops/conj_physical_compositeimplicitautograd_dispatch.h,sha256=ZFAdnLbq8GILHntlz1E44zUvFS11mk5u8DaeGmp2oyc,1025 +torch/include/ATen/ops/conj_physical_cpu_dispatch.h,sha256=xeWO_OvjgFG6uXC1wAYCCwtfSDooRiI_v96_zX_o3rY,1091 +torch/include/ATen/ops/conj_physical_cuda_dispatch.h,sha256=GOXaMFwmxZoPe-MCpSozemHfBp5ZNVCVhJh0VyOQt3M,1093 +torch/include/ATen/ops/conj_physical_native.h,sha256=GJDM-LZM6olX1S3WgwBcaqbIGxOL9etcgUtedQavEmM,1137 +torch/include/ATen/ops/conj_physical_ops.h,sha256=Kfej-QviSIPW9R16f7dozlD66lCIIphtQxEBTk2HjyA,2363 +torch/include/ATen/ops/constant_pad_nd.h,sha256=tky6JL56A-gbgo0p08GyhUSF_P-aktrLK3wlu88FEmQ,4502 +torch/include/ATen/ops/constant_pad_nd_compositeexplicitautograd_dispatch.h,sha256=9dBNnQvGLW-1W8iIqfjNHf1tv0L8G5Dn1AHkAPQFI-0,1763 +torch/include/ATen/ops/constant_pad_nd_native.h,sha256=IgjF-9RxQY95X5DMQ4d2yXEQYV8av2zFqfrYMpZtQBI,933 +torch/include/ATen/ops/constant_pad_nd_ops.h,sha256=iJqA0wExUBh8R-eRVniRzwZH_F9kYOkg3gM0GM0EHqU,2201 +torch/include/ATen/ops/contiguous.h,sha256=sEXG0CvdIqvTB4SJcgg_LVUyRL_hy4PmR7tyYPDQlao,761 +torch/include/ATen/ops/contiguous_compositeimplicitautograd_dispatch.h,sha256=Um986pl2vq-sZXPZNft5CIPVOk0uZiYKyXQzxExmkPA,1084 +torch/include/ATen/ops/contiguous_native.h,sha256=BC66UB4__Kuarfp-TDKdf49OET07QmJPxj9NlUuUrs4,796 +torch/include/ATen/ops/contiguous_ops.h,sha256=SU_Ieb_PAYowdvOn6gr3jrf36A_c_dJa9jcz-0sXVV8,1360 +torch/include/ATen/ops/conv1d.h,sha256=FpB3TnHXPpjvGvz6mKCcDnzuQXDZHtWwzQmLrPNzzJY,4774 +torch/include/ATen/ops/conv1d_compositeimplicitautograd_dispatch.h,sha256=E1kVxz-w09ZpFgZRC7gDQk8fYOXt1idJpnH8x8wakQ8,1950 +torch/include/ATen/ops/conv1d_native.h,sha256=pVJgyZkymCYvP_cvZhow6Ot5IXrQk8bXMdGTObdI-7c,1250 +torch/include/ATen/ops/conv1d_ops.h,sha256=umiymGU3Y8XrvBdzLpu5E_07LxclbKbhZknidJEjd5E,2931 +torch/include/ATen/ops/conv2d.h,sha256=CKXkJEMnUUPuiqVO4CwHGJbN-1-L1eCoMAefZDz_UP4,4774 +torch/include/ATen/ops/conv2d_compositeimplicitautograd_dispatch.h,sha256=HqsZ03dpQHCBFob6YuHiHWiW5irldJAXo3S4RspeqXE,1950 +torch/include/ATen/ops/conv2d_native.h,sha256=DDiIEflW4T-JigaF-dVM6mNtHu-YQNGmPWMH1-WKAF8,1250 +torch/include/ATen/ops/conv2d_ops.h,sha256=vrRBzu0Iy0kg3zkWxf1APHpZEn5xAk7Nntp-sYfUjQU,2931 +torch/include/ATen/ops/conv3d.h,sha256=e-2yGUt_74wueDE8ypyZ3er--3dxfxyZJ-60478MU-o,4774 +torch/include/ATen/ops/conv3d_compositeimplicitautograd_dispatch.h,sha256=rzC71YCme2QyW4gBsClhLYLSy5fRN7EA3lbSikpRQXs,1950 +torch/include/ATen/ops/conv3d_native.h,sha256=yXo626lnS_W5cPKADvePTtw87q4-Ejg3P6kVJGJ25g4,1250 +torch/include/ATen/ops/conv3d_ops.h,sha256=QBbDlg0mIWC68zezA2tchNdVjkWQIhqS0xpr-9OYd1s,2931 +torch/include/ATen/ops/conv_depthwise3d.h,sha256=xBLhCJAeR4u9tFpFhcPLj6-FnONlqYAt1Tkbl_FOh3w,7547 +torch/include/ATen/ops/conv_depthwise3d_compositeexplicitautograd_dispatch.h,sha256=PFn8ZwTtl5-p1t0srlrT0adDmMAlC87osa393cUTHAA,2056 +torch/include/ATen/ops/conv_depthwise3d_cuda_dispatch.h,sha256=oIvIzaDEozbkqL5gN-2XPYW9I5rL0DUHKXqpbXxsz2M,1419 +torch/include/ATen/ops/conv_depthwise3d_native.h,sha256=nqfOk9kjlW_n2WZNrNM8XsbxKHPHbtqIntfEykq-1k4,1202 +torch/include/ATen/ops/conv_depthwise3d_ops.h,sha256=1uSsEuesUy3mi7aRrvjv7M-DgFKd-1W9SA8VoN1rAQ8,3117 +torch/include/ATen/ops/conv_tbc.h,sha256=6YlWDXf20HAJFZ71ZTgAuWHFX3AOsrd7DhbcL8sRhlU,1674 +torch/include/ATen/ops/conv_tbc_backward.h,sha256=anLSiBIum1LcjCJw2y_w2jX0FoGGkwd_PrHWQ4JWcwI,1157 +torch/include/ATen/ops/conv_tbc_backward_compositeimplicitautograd_dispatch.h,sha256=J-VS2dIXFgdAzkkBI-TYIKE6qz-BI6YtH9w73bJCwpQ,1156 +torch/include/ATen/ops/conv_tbc_backward_native.h,sha256=lC4eZS70oacOuU4ruvttcQat0qFIQosT8bm99PETjpU,868 +torch/include/ATen/ops/conv_tbc_backward_ops.h,sha256=w3uk8PGyiWD7nsGBfhLONWrXR21pcS8VEU6dnCgFZAg,1672 +torch/include/ATen/ops/conv_tbc_compositeexplicitautograd_dispatch.h,sha256=_Ojd5y8COZimt6BdZ3zkOS_nNA-i_8b1H48avSE_r0g,1380 +torch/include/ATen/ops/conv_tbc_native.h,sha256=6SG095g4Gs4K1xEY-SBin3Ep-LLLNhrwfuERYWVE-5c,944 +torch/include/ATen/ops/conv_tbc_ops.h,sha256=XiCJlEdMJaHkU1vccCVQJsn_JBy9uNJOq_fbkffABUE,2249 +torch/include/ATen/ops/conv_transpose1d.h,sha256=JAhJir3ep1UJhFG3LZ5QAOQcUKCtquxGnkALkFBaHPo,3261 +torch/include/ATen/ops/conv_transpose1d_compositeimplicitautograd_dispatch.h,sha256=VBdDHmztns6YRTS7rRUtiSCzokMS2-aXHlzAPPwFG-E,1583 +torch/include/ATen/ops/conv_transpose1d_native.h,sha256=u07kTTj_qQjn36zgYJLoZGoPlxJUVd82HwWkKzjj8aI,1025 +torch/include/ATen/ops/conv_transpose1d_ops.h,sha256=IktNQ2WYCMxL4LNlBZDSBzh0UJcY0y2jGs1sYYHxQaE,1964 +torch/include/ATen/ops/conv_transpose2d.h,sha256=uvcepvf_-rKEnkFNHsp0fMAEn8OaHOLuwEt_q2vd_0s,3297 +torch/include/ATen/ops/conv_transpose2d_compositeimplicitautograd_dispatch.h,sha256=40UtjH50BGmUiNtwXPug-CrrJwIklUNrl9UHig74oos,1583 +torch/include/ATen/ops/conv_transpose2d_native.h,sha256=BdemX3ANVwnMwKqpZRyoKUxFHbv8bNuvWxHCzNxmKas,1025 +torch/include/ATen/ops/conv_transpose2d_ops.h,sha256=OXLGzbUYWL6yMHzrdTQ5bfEpA4dxAq3fL743GdnckWE,1981 +torch/include/ATen/ops/conv_transpose3d.h,sha256=vyMBb7wZLXd1rVyHijWNrcttWiZ178_vscrKmI69q_A,3297 +torch/include/ATen/ops/conv_transpose3d_compositeimplicitautograd_dispatch.h,sha256=JTWuU6fcvfl6wFZ_m6PFWihHTgvs3hIcvxIv6XjH8A4,1583 +torch/include/ATen/ops/conv_transpose3d_native.h,sha256=6Y4J-kElFHC9b-LhJyTOZDEJZNUu8uyVuiINAACXHIw,1025 +torch/include/ATen/ops/conv_transpose3d_ops.h,sha256=XwwUK4P9gUFVU3qKqk3dMHjy07xtScGfoYPQMCZ7fXo,1981 +torch/include/ATen/ops/convolution.h,sha256=T_AT68SB3_vUgd3p9xaF-YLIWzqZ5nqy1OLG-s5xfJE,8340 +torch/include/ATen/ops/convolution_backward.h,sha256=B2ni9MyJiW9OTZtKzR7tdAXDtnUUhMYTkC0G1MUzu7U,11741 +torch/include/ATen/ops/convolution_backward_compositeexplicitautograd_dispatch.h,sha256=nlvEzclraYLJFFH6D6OYFeXqrJdPV2oa68MVOVC15hQ,3568 +torch/include/ATen/ops/convolution_backward_cuda_dispatch.h,sha256=UydXMamLCUDCfix1fivLE82ANm3ofVAa0MlrdJbktIU,1700 +torch/include/ATen/ops/convolution_backward_native.h,sha256=qte5LkqG4oSH7g-idIJwa9QkN55YpnugwJQYCYlRfAg,1521 +torch/include/ATen/ops/convolution_backward_ops.h,sha256=dDOBUtbPTcFSE-BY1C0_IWYPZCMS_OwXk-84qaAS8Mo,4236 +torch/include/ATen/ops/convolution_backward_overrideable.h,sha256=gQJ26cofeCQ5dbAH7_o8oVZZzpsr7VYfsdsPebqwNwA,10938 +torch/include/ATen/ops/convolution_backward_overrideable_compositeexplicitautograd_dispatch.h,sha256=c5J7o6Eu1ftfXXjQ-FcuBIJjsBE6BpJhoydpt_vT3gA,3421 +torch/include/ATen/ops/convolution_backward_overrideable_native.h,sha256=WjMIlhmaN1hdVb6igqyq7s5x5Y5oKkH_74dFInSPP9I,1472 +torch/include/ATen/ops/convolution_backward_overrideable_ops.h,sha256=6KzPByXqTqDZSvawFy6Od8gX0s4Wl3n1NwxyqeEXVUg,4091 +torch/include/ATen/ops/convolution_compositeexplicitautograd_dispatch.h,sha256=4wxazokwjYrBemMdVsseg6JsIoc58kcGpA6cqBegfjc,2757 +torch/include/ATen/ops/convolution_native.h,sha256=4HREAKXNdZe2v_Ht1dB60YKm_snHDWDR6SUyBqzqXZ0,1265 +torch/include/ATen/ops/convolution_ops.h,sha256=GGT5WoKn2TZczrNuUVr9L-n6jHRAx-S0xddAFRfAzvw,3353 +torch/include/ATen/ops/convolution_overrideable.h,sha256=Z_DiL3m_Bp3eBcFPqmDNLymMxtHywMvaUw32RjzpHfQ,8743 +torch/include/ATen/ops/convolution_overrideable_compositeexplicitautograd_dispatch.h,sha256=3woHMh0AGzgT7FiKb6M7cjP-vs9sIJthLFmcMemd_d4,2835 +torch/include/ATen/ops/convolution_overrideable_native.h,sha256=dQosXkrBwVzKIXmgITJNWwNEE6F3S3NSC6rFY_lxNhM,1291 +torch/include/ATen/ops/convolution_overrideable_ops.h,sha256=WrPe-SRsIwV5Dn5IzXFy2ijPQEoYf80A4MTOEaYEDCo,3431 +torch/include/ATen/ops/copy.h,sha256=9IFZrggWuRygYxK_DSlxTHdwkdqiGtPLJNkrw0j5lVA,1570 +torch/include/ATen/ops/copy_compositeexplicitautograd_dispatch.h,sha256=yd4c9qsm49ZfduuiiKCOuag6evx12urQWhjes38lePU,1307 +torch/include/ATen/ops/copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=PiE9IE-N7gy_T3bBlnj2iam2gKwBY68AdtfwReCceV4,1091 +torch/include/ATen/ops/copy_meta_dispatch.h,sha256=9CXIF_9FkwBHGzTEqSC19DSviTx4QpMMTukhNF5O3T8,1023 +torch/include/ATen/ops/copy_native.h,sha256=pAw4pEXEnyPyLWCnkddlH1FWsxn3h-IOy2ozLvR0oZw,1539 +torch/include/ATen/ops/copy_ops.h,sha256=SqoiiThKFhs1-Cuu5IPyBd0BnBpjY2Nknti1PWXith4,2729 +torch/include/ATen/ops/copy_sparse_to_sparse.h,sha256=EsPfiSSI2z0HPiaS0EkuvsYSyiLzLDCjNEmB0Dw3Ns4,2032 +torch/include/ATen/ops/copy_sparse_to_sparse_compositeexplicitautograd_dispatch.h,sha256=oSGqNJnxA3jU4cTLLEzKkieWoi6-2eEF19PPYFKaZC0,1361 +torch/include/ATen/ops/copy_sparse_to_sparse_meta_dispatch.h,sha256=5Y3Vl2EFwH6MlqATX2rKhHVctCoaOkfURmjivNgSfg4,1037 +torch/include/ATen/ops/copy_sparse_to_sparse_native.h,sha256=Yb0mw17GEVjAVTFhkXlzfoPW57ACBqfYjkbkFDDoP78,1035 +torch/include/ATen/ops/copy_sparse_to_sparse_ops.h,sha256=UrJKUc6ySxZLQ4WQPexNwRCR4vJ1smxTBSQbJwgD1wE,2882 +torch/include/ATen/ops/copysign.h,sha256=me1H4Am8APHMzqf5ICqVJKms-ffXw4nNf0kIbfzvpiY,2182 +torch/include/ATen/ops/copysign_compositeexplicitautograd_dispatch.h,sha256=sFzlfsS4RrTYgSn9LBJ2eAUPEFbANmM5lcZp1YSzfrs,1338 +torch/include/ATen/ops/copysign_compositeexplicitautogradnonfunctional_dispatch.h,sha256=8-BnzmjOy6TU_KbXL5zju4t4Sdqzxf88GJX6uJhrEcU,1151 +torch/include/ATen/ops/copysign_cpu_dispatch.h,sha256=wOvjvuKqOj6_6FaRXEpnbr_3rI-1NrxhTLNhwl2yRkc,1294 +torch/include/ATen/ops/copysign_cuda_dispatch.h,sha256=o_FPMuq37qTS1TlZA-I0fF5NDE8vMQhZVULMsopyCJE,1296 +torch/include/ATen/ops/copysign_meta.h,sha256=7EMbc8F9wZ1FCQqbI7xMzb8MlueVGJXkt7afAKQ01Y0,856 +torch/include/ATen/ops/copysign_meta_dispatch.h,sha256=-7RHCxwfC0Rr32etOAx6tjdoSCiPU9unf06IUZbvLI0,1296 +torch/include/ATen/ops/copysign_native.h,sha256=GRfSWimfGAZmxrxMhEezdcR0FartkSZ9v_7CbRN01OE,1156 +torch/include/ATen/ops/copysign_ops.h,sha256=zdWMGeETcznuIDgz8OtzfU-IyHoENPt0z08XFmAi1rw,4542 +torch/include/ATen/ops/corrcoef.h,sha256=0qjbqFk4uQ5FgzxfNnpTr2_YfzaVsIeTI8NeMgXMjc8,899 +torch/include/ATen/ops/corrcoef_compositeimplicitautograd_dispatch.h,sha256=0i-bG2qHZBXnRcmf_R7-BQohGRMost700wsgjzTNuVk,1020 +torch/include/ATen/ops/corrcoef_native.h,sha256=MdueadJ39JBThzKhfUMeCE7gCGM5XTqFI_Fu4p64TdE,732 +torch/include/ATen/ops/corrcoef_ops.h,sha256=zy8DAi7nkPpTtb2f5sEeFn6iFYTkxRCo-AlXBQUf28s,1217 +torch/include/ATen/ops/cos.h,sha256=BHXeo3u6ALuH8dzM7zRDZlGJsGUsHc6OshSHth871Cg,1384 +torch/include/ATen/ops/cos_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Ll50kBkGvgqgl24hIZ5xTHPALioiuXnc8q-lgWyIjWg,1089 +torch/include/ATen/ops/cos_cpu_dispatch.h,sha256=68QbnS7-UCBSOPAMNcVZofvZ1x3qSYnwaHvNxQwVqtY,1170 +torch/include/ATen/ops/cos_cuda_dispatch.h,sha256=wAbf4O9Gl9TUxFndRA-IK1cv3PPocMVZJiIYnj3pzoo,1172 +torch/include/ATen/ops/cos_meta.h,sha256=lazB6zUAyI2CFce3snNR6UUypB6lYrx9P3-m4EaDH2M,818 +torch/include/ATen/ops/cos_meta_dispatch.h,sha256=4jpRTWRwv6SmK0F6exDBc0oLYPX-im5Zq51mo2W-Sgo,1172 +torch/include/ATen/ops/cos_native.h,sha256=aj0-ZwAgwKGtPsN-gP5HetE349YF-PwvW82y2pGXAbM,905 +torch/include/ATen/ops/cos_ops.h,sha256=f5fFhL-FgY5qa0W6PYH7eHqsYawB_s4lV457J9SsLvQ,2273 +torch/include/ATen/ops/cosh.h,sha256=qZu69Tj_PTz4iPbqA2oa-uxM8WB_mXNUk_DbO9gIsEY,1397 +torch/include/ATen/ops/cosh_compositeexplicitautogradnonfunctional_dispatch.h,sha256=SLP7Oa7IR71jJ2JwqaAI3zSMivmUiODLFPUI3nPelEI,1091 +torch/include/ATen/ops/cosh_cpu_dispatch.h,sha256=0Uay4rVvvdOSKXmawx5qGTAHqJ9FLVV1ntCqW_34Aj0,1174 +torch/include/ATen/ops/cosh_cuda_dispatch.h,sha256=F0X6Ly-_2Q24o_KC13Vebq_kHTMeGmIdFd_LMLK3iUA,1176 +torch/include/ATen/ops/cosh_meta.h,sha256=UGgAuh2ebNPHQWZGHGidT_IxgWGB3BIoXP8AUndFkLE,819 +torch/include/ATen/ops/cosh_meta_dispatch.h,sha256=Dk3RySp1Qs4DgArnkBgrJcMFE4uNVwP60JXJhuKlBGo,1176 +torch/include/ATen/ops/cosh_native.h,sha256=dw0IVylLDazt4dU1falh6sUiAw-v9XKxjHk0Iiug1hM,844 +torch/include/ATen/ops/cosh_ops.h,sha256=YKb6TtpeD-WdYmI6CtI1FlSahTzblrKtdudzNm-vW-M,2282 +torch/include/ATen/ops/cosine_embedding_loss.h,sha256=WsHFO6_Xp_raI8dSSpIWFXv73K86qqIjsbtorcWYPPU,1172 +torch/include/ATen/ops/cosine_embedding_loss_compositeimplicitautograd_dispatch.h,sha256=mSH80gavA9CmFwPvsCFwXgxxxW-b6WuWtopR736e0pY,1147 +torch/include/ATen/ops/cosine_embedding_loss_native.h,sha256=bA5uanPyztttUAZuu_kQHr6-O-m7-luR40_uzyeCI98,859 +torch/include/ATen/ops/cosine_embedding_loss_ops.h,sha256=BQCnGAH1BETytucpsr6nchDR6nrRSPm1QsWqpB8gKqc,1563 +torch/include/ATen/ops/cosine_similarity.h,sha256=37E2HMqSTVJCgxsA7_-RnFbYIoA8qheETENxKf5Ry3Y,1038 +torch/include/ATen/ops/cosine_similarity_compositeimplicitautograd_dispatch.h,sha256=5KP-xURniNVMjkIq0tsuZv-VkUgouZRwE4WBGI0OIOo,1083 +torch/include/ATen/ops/cosine_similarity_native.h,sha256=U4EMGswPevu4XGODLx2W1Qfgqbs0GyY6yu1aLxU0b3g,795 +torch/include/ATen/ops/cosine_similarity_ops.h,sha256=Pz5lbF22TgynzdWhrht5IgYct6y65eDsr9GLXQkZiNQ,1410 +torch/include/ATen/ops/count_nonzero.h,sha256=d0w7bunieMNtbp0lzkaDOBEGwPcJcd_v8myRR8VRkGg,2302 +torch/include/ATen/ops/count_nonzero_compositeexplicitautograd_dispatch.h,sha256=7Rlu6hmK_-0ETXzt0wQBe0RXt37eypt4o68JURX6Ia0,1529 +torch/include/ATen/ops/count_nonzero_cpu_dispatch.h,sha256=SsWKGRDePzawMDoAHdKefvCxAsDUC2fvKt2zm-4sKhk,1002 +torch/include/ATen/ops/count_nonzero_cuda_dispatch.h,sha256=Aot9cdQZU_lr-VCsOvbdr3GmqHeXuixRvO4UUKmHiMM,1004 +torch/include/ATen/ops/count_nonzero_native.h,sha256=V0vbET8kcow8J4Z4XsbtvjuuCmbYIBPw6yk6BVwki6c,1188 +torch/include/ATen/ops/count_nonzero_ops.h,sha256=6APJKICqqKebl97mLbq9pz02pw2G3q3Sy1nW9A8IKhc,3388 +torch/include/ATen/ops/cov.h,sha256=vPAA-2SYmIN3QdQXagu4yNKVJzVPLnt343LOkS0DKf0,1098 +torch/include/ATen/ops/cov_compositeimplicitautograd_dispatch.h,sha256=L0TfAbwQ8q46g5PQqOhKRx3MxJm-vS0oRTJAJco4PiI,1135 +torch/include/ATen/ops/cov_native.h,sha256=uvTMkhEtAx5aw3VWZPA8zFGmMylhIpRqcV2dqughjQ8,847 +torch/include/ATen/ops/cov_ops.h,sha256=wKqKIKd01yoh6bIAUXjkyuv3Rn9SXZqWMPex22wIWPg,1576 +torch/include/ATen/ops/cross.h,sha256=BHADy1TIXHnGLOD66PtDynfbo1ObOb_ILEcakBt4OMo,1592 +torch/include/ATen/ops/cross_compositeimplicitautograd_dispatch.h,sha256=-cBQ1nVjOc0ZPfYlhqwL3n2kup-T34TAHEQ6z_dn4Qs,1370 +torch/include/ATen/ops/cross_entropy_loss.h,sha256=ddynqf3sj80UXPUgN7-bXATfpqHBsSRJ15FPltg46Fo,2687 +torch/include/ATen/ops/cross_entropy_loss_compositeimplicitautograd_dispatch.h,sha256=Hnb8T__Ojjjt6Igs4KSo8Kj6Q4ufebe70mmJ52obmss,1443 +torch/include/ATen/ops/cross_entropy_loss_native.h,sha256=eEcoLJeUZmwKNVPu-8gOczEcqWoere8z5L7T8jNMMMI,921 +torch/include/ATen/ops/cross_entropy_loss_ops.h,sha256=CKzht_utb3Zt1CQBk42RQi6l9oz9KOKjo3HidAt7ncI,1723 +torch/include/ATen/ops/cross_native.h,sha256=ope4SGodbjJuFmIZjwzmr3q8XboI-cA9d-0EtIiunys,933 +torch/include/ATen/ops/cross_ops.h,sha256=VayqhwkW5gQMHFLVYymi-6S6b4OjmZM5q7MiMaynyB4,2169 +torch/include/ATen/ops/crow_indices.h,sha256=RZuzgJu8NV_D4320naG0VvM_jJEA73hu2KNvonoEpXg,763 +torch/include/ATen/ops/crow_indices_compositeexplicitautograd_dispatch.h,sha256=PWht7xnhUGSw3mzQ-K2_hRCv89j9TVpthsXkYEWqY2w,1024 +torch/include/ATen/ops/crow_indices_copy.h,sha256=jd5mUmoep5sh5-GWLwSQjDKFuDPkug6i3eYV5XObpww,1391 +torch/include/ATen/ops/crow_indices_copy_compositeexplicitautograd_dispatch.h,sha256=DiOgETwSXBAShC_ID8rvhMwICs88Gdf7ozmjcP0uLIs,1143 +torch/include/ATen/ops/crow_indices_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=xz-nFw_N1MLsjf3PT3SKP7B_AmUCjhI421Ty0uUCjCI,1055 +torch/include/ATen/ops/crow_indices_copy_native.h,sha256=HhbESces5CbHQqYVNKrE9lG-KI3JnBbvyh86E9Iesjk,830 +torch/include/ATen/ops/crow_indices_copy_ops.h,sha256=I8yrMGgpxGOQwEuSf98XJY7wMg033Dm-HwB75wtlGmA,1867 +torch/include/ATen/ops/crow_indices_native.h,sha256=EP6wm-BBTVITZGoensi-Uq0kzCF5aIFsdaO_P9AHx5U,815 +torch/include/ATen/ops/crow_indices_ops.h,sha256=4L8f2oAOgIdZv3GlOsK7YaBwZEyipj8Yhbod6CMjPyE,1235 +torch/include/ATen/ops/ctc_loss.h,sha256=B_h1qP3Cy7TJ-PCwTsT2dMju0mB7do5v9TM1woOCmug,1836 +torch/include/ATen/ops/ctc_loss_compositeimplicitautograd_dispatch.h,sha256=wgQuO0jY9Aj03r8w7CoS2OKrfclyHBWWUadBtc3RMKA,1438 +torch/include/ATen/ops/ctc_loss_native.h,sha256=Ah7xAT4JZDyXuEvFN3Z8i-YrgnhgTkAkSYoy0i9GK1M,1150 +torch/include/ATen/ops/ctc_loss_ops.h,sha256=K09pvc3nrEiAP5xedKaIkpcy1_Eh9jKL3h-jGi-ZT48,2788 +torch/include/ATen/ops/cudnn_affine_grid_generator.h,sha256=qxS-5E1sDi-9xQ6DFCNG4-roAIt0CHsWGZWTDx66CZw,1757 +torch/include/ATen/ops/cudnn_affine_grid_generator_backward.h,sha256=IIwidjfjH-3k0fBHxTUJeX6IRaQZ2QIhByImcL0HnU4,1844 +torch/include/ATen/ops/cudnn_affine_grid_generator_backward_compositeexplicitautograd_dispatch.h,sha256=vkKeJj1-IyQcqUz2E9DSH5gO9jOZWUqPLn1t9ZwGaz8,1269 +torch/include/ATen/ops/cudnn_affine_grid_generator_backward_cuda_dispatch.h,sha256=hk-ssJZpfqR4cyGxPLf1MqRx38sAbmYYdHbIghBb0RQ,1050 +torch/include/ATen/ops/cudnn_affine_grid_generator_backward_native.h,sha256=H8C9TxsJeROesJJ7ZEhKuioDZ_Qcp8E8q4oBcYUswYY,956 +torch/include/ATen/ops/cudnn_affine_grid_generator_backward_ops.h,sha256=7yXOrlMRZjkGrAz2I63ZZUgUTY9e0QNe_0K3hncpEzM,2296 +torch/include/ATen/ops/cudnn_affine_grid_generator_compositeexplicitautograd_dispatch.h,sha256=1HLOOgxvkoaJEGYh8mbi-9GY8pZYgkNsd5e5gpbDKyQ,1253 +torch/include/ATen/ops/cudnn_affine_grid_generator_cuda_dispatch.h,sha256=eFpG6tgTaDq5Esr_KvKZEl0YWqIeOxvebgnlzSt-LOE,1042 +torch/include/ATen/ops/cudnn_affine_grid_generator_native.h,sha256=kWvXFSZE5cRYz9bKVetxtWMZJ0JbZ0XqFAPORde4G1Y,948 +torch/include/ATen/ops/cudnn_affine_grid_generator_ops.h,sha256=UKpG7zlY03h5Esrv1Hu_Q1ETJcwTyZDSlOFemhOEmUs,2242 +torch/include/ATen/ops/cudnn_batch_norm.h,sha256=c4nml8LM5eZnNPaROEAGZfyK39D_mJHjwxgPNDVZKDQ,3277 +torch/include/ATen/ops/cudnn_batch_norm_backward.h,sha256=B7MlxYbauHW5DZVRDfESsNlZ8esRt8SDDcOODubrmJo,3469 +torch/include/ATen/ops/cudnn_batch_norm_backward_compositeexplicitautograd_dispatch.h,sha256=yzySpEs_VsDkbRFX6tPFO6utLsVJ7JEC-yIMszdADYA,1919 +torch/include/ATen/ops/cudnn_batch_norm_backward_cuda_dispatch.h,sha256=viw1pA-IaiW3akgmWzYKC4ckqbY1e6f4i83-PK6Idg4,1332 +torch/include/ATen/ops/cudnn_batch_norm_backward_native.h,sha256=ih9R5zG1pNfA_yjr8tSC4AYmGB_lvAhFNYXGioIvzwQ,1563 +torch/include/ATen/ops/cudnn_batch_norm_backward_ops.h,sha256=RHozn6S92sE4hOw6u0nqC7Z6pEQannh7LBDm8ox7k9Y,4268 +torch/include/ATen/ops/cudnn_batch_norm_cuda_dispatch.h,sha256=1eONKh-tnpRZ45CfOBR84l8qwqElzE95zmh3_nMxKlw,2137 +torch/include/ATen/ops/cudnn_batch_norm_native.h,sha256=NKoxNq9YjXgOQI77f1GglH1yRD_JTtX4YEycpzr3XJY,1456 +torch/include/ATen/ops/cudnn_batch_norm_ops.h,sha256=SEVagF9WREs3GhSeNbHCKJ1eXrOQzl7YUbM40OUgmzw,3955 +torch/include/ATen/ops/cudnn_convolution.h,sha256=OM0kl2IlUCdQp1UTBGrMuCeJHQGRF1k9xdknFDH-JsE,7890 +torch/include/ATen/ops/cudnn_convolution_add_relu.h,sha256=Vzsd9CatvFsFIntpxIRs4ovWOlggrWNc0Y3ZL3AELSE,8469 +torch/include/ATen/ops/cudnn_convolution_add_relu_compositeexplicitautograd_dispatch.h,sha256=lR4_Itws4_cAovSsVXj3Hle9aZ4rkCb9PW7XYkpd-Y8,2304 +torch/include/ATen/ops/cudnn_convolution_add_relu_cuda_dispatch.h,sha256=DKESKHV8F_wgcsQbeWq8p6lmVUGB6r6ZJfQjzL9TtIw,1543 +torch/include/ATen/ops/cudnn_convolution_add_relu_native.h,sha256=9WL0i3Y8HOO-htpEZmjPgiCA34yWaNQ2KONeI2tsVd8,1321 +torch/include/ATen/ops/cudnn_convolution_add_relu_ops.h,sha256=iJyCBEidjokS7w_k1JYljxgdBHVj0pahch3kdh--ous,3511 +torch/include/ATen/ops/cudnn_convolution_cuda_dispatch.h,sha256=ONsQTYmdbG564qqWELM-2qTX3qGnRGiburJx6Js8OMU,2505 +torch/include/ATen/ops/cudnn_convolution_native.h,sha256=GS1yg7BaKzX4p5EFU6eeXsxXridPLQNzgDiRUnRVqDA,1172 +torch/include/ATen/ops/cudnn_convolution_ops.h,sha256=6oWC8M_5hEGtUCMgkTeNDPM0uP8Sn4c_EFjqyxQ0Ew8,3117 +torch/include/ATen/ops/cudnn_convolution_relu.h,sha256=5_pSqRBIx7YBSrREXYJnxvYOyCgQ2hzeXv-o9m3jL_8,7295 +torch/include/ATen/ops/cudnn_convolution_relu_compositeexplicitautograd_dispatch.h,sha256=aKstaYPvziQRw-RoN8CyLULQSvSpfhbsPG1SKTyVu-w,2028 +torch/include/ATen/ops/cudnn_convolution_relu_cuda_dispatch.h,sha256=xcHzaS67d-UWObaIwpywtm2fuz9-8CTLNpBLukJVcBk,1405 +torch/include/ATen/ops/cudnn_convolution_relu_native.h,sha256=2euB-QaPSv_l0hlaBQfelA32UbESfahXJBaXY5j9d5Q,1183 +torch/include/ATen/ops/cudnn_convolution_relu_ops.h,sha256=0P1PLtK6-AWtMCv6IqsxPqiEyXUNP-JU-3bAJpMZ_KA,3063 +torch/include/ATen/ops/cudnn_convolution_transpose.h,sha256=DJY4r-QKnjUxWIpL4CQhpIRj-VyV5N2EWPHmr4x52lg,9106 +torch/include/ATen/ops/cudnn_convolution_transpose_compositeexplicitautograd_dispatch.h,sha256=7ObqRzFhPJ_ALLtE1gJ8I-JHNkrYjS2o8nkcgC6wikk,2228 +torch/include/ATen/ops/cudnn_convolution_transpose_cuda_dispatch.h,sha256=CTYVKHl6FE0Cf_jitF2OKRWhOaHLSQAVf09SdHD1EB4,1505 +torch/include/ATen/ops/cudnn_convolution_transpose_native.h,sha256=V5gpemtkwcfYcfogCadv7J_PGrwX4IhiDfOZc9Jw9TI,1283 +torch/include/ATen/ops/cudnn_convolution_transpose_ops.h,sha256=3XgexstO1s-xy7Dxq7X8LQ_MWuEqPo7lrWjC_LSkchs,3413 +torch/include/ATen/ops/cudnn_grid_sampler.h,sha256=Lc6JCvM5hznRbdLAQIfB8-nUaIYMG85X2zAvsHH9pfM,1540 +torch/include/ATen/ops/cudnn_grid_sampler_backward.h,sha256=B-UN2i5ZSvBUXpy_NyY9ii6LZREOjbMqEvi9DAyMV5o,2045 +torch/include/ATen/ops/cudnn_grid_sampler_backward_compositeexplicitautograd_dispatch.h,sha256=LMu2qZEe3O6-1CpDfnXTp8w6YRQpFibWn4WhPbXj9b0,1371 +torch/include/ATen/ops/cudnn_grid_sampler_backward_cuda_dispatch.h,sha256=zO6G7ZtXGAbLJd6mAFby9ZmweX5DM22NrjrwEdbnyZs,1079 +torch/include/ATen/ops/cudnn_grid_sampler_backward_native.h,sha256=T3DS2ZipefEWGs609-pFifzjqDH3TdHcqGEOZb9ap4c,1036 +torch/include/ATen/ops/cudnn_grid_sampler_backward_ops.h,sha256=qseTid59T9rZRLmSLvO6Tx1rafE3y7bZFO1sAmmOG6k,2573 +torch/include/ATen/ops/cudnn_grid_sampler_compositeexplicitautograd_dispatch.h,sha256=F9LJfrknJEeczFmAo46rExBd0-dhe76QM3HvVlLZouI,1195 +torch/include/ATen/ops/cudnn_grid_sampler_cuda_dispatch.h,sha256=MRAnv5uytQfRxyCMXO2tCjuEO36a-PMLA1RFRKPajkM,1013 +torch/include/ATen/ops/cudnn_grid_sampler_native.h,sha256=ZhL3WUeU1RXO3Fn_sweRJYhaJtEfoMF_DIg2wq0AvAk,890 +torch/include/ATen/ops/cudnn_grid_sampler_ops.h,sha256=E3zLe7WmfhmNmmGc-Y34gJodz8AOkKv1s0w9qdYZh9E,2046 +torch/include/ATen/ops/cudnn_is_acceptable.h,sha256=poqK07lhDhmj2DzjVoE0wm__UjQJePEk6OzOQumq2t0,935 +torch/include/ATen/ops/cudnn_is_acceptable_compositeimplicitautograd_dispatch.h,sha256=h_JyK8XS2h8o2v76O-_Z2feOC2n_FF5bIT7MjpTDrgw,1025 +torch/include/ATen/ops/cudnn_is_acceptable_native.h,sha256=ROICuxcj7oHyAQ-45H_iBdZm9CFv5D5I2UJ5YdM4MdM,737 +torch/include/ATen/ops/cudnn_is_acceptable_ops.h,sha256=_N9AGGY17yLQSKzz4XF05CE1MvHdq8iG7Sap97bm_1U,1230 +torch/include/ATen/ops/cummax.h,sha256=xNek60u6PCDIKIwun6OSsoXzS2hdvBuUqLQv32VrvIQ,2604 +torch/include/ATen/ops/cummax_compositeexplicitautograd_dispatch.h,sha256=s80yG1PVRFVODBszuUhVVccF9fBDI3f4z71ZDi1k8YI,1343 +torch/include/ATen/ops/cummax_compositeimplicitautograd_dispatch.h,sha256=HUw2Nlg5d8De9AlK3X1UcJjekbhSS48s3uZ4I8O6VHQ,1355 +torch/include/ATen/ops/cummax_native.h,sha256=EX8JU35S6qGm-7NE3nQWgFOOSl8ZrDI4s7u8NpM8OpY,1154 +torch/include/ATen/ops/cummax_ops.h,sha256=eOV12t2C5CF7RZIWlpDPOV_702nz5lCYtPcG3Iej854,3740 +torch/include/ATen/ops/cummaxmin_backward.h,sha256=q0AaydzvsTclefcXvylLDRkNLXO6oVAp1sITjNubO5Q,1066 +torch/include/ATen/ops/cummaxmin_backward_compositeimplicitautograd_dispatch.h,sha256=C0GX2WFeO1HG5bGWJtiqzb4-j1o80Ri-p_FedPNvH48,1097 +torch/include/ATen/ops/cummaxmin_backward_native.h,sha256=QORI2zgh1-BLiAFLhTJTnZzT8TxM1-p0k9H-UXd7Xjw,809 +torch/include/ATen/ops/cummaxmin_backward_ops.h,sha256=Jb65S8JKd0otaKmK3-4tyzvezAAiQPsb074Oq7_RqFQ,1469 +torch/include/ATen/ops/cummin.h,sha256=omgTlgkRz0Ei0HuYA24m5Nm6OX5hbEFzOvXq2WWgZk0,2604 +torch/include/ATen/ops/cummin_compositeexplicitautograd_dispatch.h,sha256=64aTbWHfROi9jAvlBp4_hYiG4tWEh5BTcR-4WqVzBl0,1343 +torch/include/ATen/ops/cummin_compositeimplicitautograd_dispatch.h,sha256=P4gIx9UOKGTYVswH8qoG5ezMOCYwxZTZFU-GmhKYTj8,1355 +torch/include/ATen/ops/cummin_native.h,sha256=ziocBbqCy8RnE0m7pceECqoR-ZMwhH77t_k6r1mFKIw,1154 +torch/include/ATen/ops/cummin_ops.h,sha256=QAVRzYbauU6IzhH8a5gkrzC_JTsC8QT4IeFe896IXNo,3740 +torch/include/ATen/ops/cumprod.h,sha256=-UMBZFohCzhbuAqyqF56TrXhOU376C4I2xEL1Nov0k0,2545 +torch/include/ATen/ops/cumprod_backward.h,sha256=KBjrfpKaNuGG5Npm5c4UZhUQScQVe_ocO4e3vL1fT90,1055 +torch/include/ATen/ops/cumprod_backward_compositeimplicitautograd_dispatch.h,sha256=DBkR8KzDXzHkriO3ntXkvZnqHej6PHuOGNpDVb6iPFs,1094 +torch/include/ATen/ops/cumprod_backward_native.h,sha256=IiUsG4O4ABrUWp7sV24j0O1pNlaJe288wAIum8VlNPo,806 +torch/include/ATen/ops/cumprod_backward_ops.h,sha256=BoUg4yUpDoXpHS9hsxkXrWdeb3baxPgdQn7Tgot6wUY,1460 +torch/include/ATen/ops/cumprod_compositeexplicitautogradnonfunctional_dispatch.h,sha256=K7Dr2h-pj8gQ06oCMTyjl61pT3GQ8iK1L18ygaViqMQ,1231 +torch/include/ATen/ops/cumprod_compositeimplicitautograd_dispatch.h,sha256=6BNna1gWmkDW_s82W2YRot5X3cNE4ikZII9_QvEd8PA,1499 +torch/include/ATen/ops/cumprod_cpu_dispatch.h,sha256=j2ThUdvWFBBuXRjCk3QebhDv7ZcCFfP-_XGjIlzDaMc,1439 +torch/include/ATen/ops/cumprod_cuda_dispatch.h,sha256=igJVnpvNugE4nDkOh7RBY_aDdEYhlc_Px3KmUC2Lqm0,1441 +torch/include/ATen/ops/cumprod_meta.h,sha256=Ph2Fu5rq4FV-B79uyVrXi0HIIWD6MfnWUhoe4znyvaY,874 +torch/include/ATen/ops/cumprod_meta_dispatch.h,sha256=FXG-qAiGvdzXpObdB2S9Qn9UdQNBscLsWBLP1cV8N9M,1441 +torch/include/ATen/ops/cumprod_native.h,sha256=95rDMcg3fg0nahEs9SQfVPrsN80p9gS23GMKUnsMCRY,1289 +torch/include/ATen/ops/cumprod_ops.h,sha256=hMFsrotr514_fo502uwrQF6iqxlQNd9Z8LfNlUAY0a8,5111 +torch/include/ATen/ops/cumsum.h,sha256=mrBtYxJFFZLWzTVHxNGeFSeE6CcowiBIZs87PYUL4S0,2526 +torch/include/ATen/ops/cumsum_compositeexplicitautogradnonfunctional_dispatch.h,sha256=As4Nwe_0GgX05iwvzLabtdNPBHtHJuEQQzKaR-zrP8U,1229 +torch/include/ATen/ops/cumsum_compositeimplicitautograd_dispatch.h,sha256=EFZbzczwSqAGQHAd3aEncMlheg5ZSyO9I1T6Zi1Lod0,1495 +torch/include/ATen/ops/cumsum_cpu_dispatch.h,sha256=vu2UZIYr27EWOBByKWsYdj9K1F8n6sKW5V1pmx19Uwk,1435 +torch/include/ATen/ops/cumsum_cuda_dispatch.h,sha256=R0s2fgK3eCJBwO8ebBdreEuD8m0dobmnxpkxntgq-mE,1437 +torch/include/ATen/ops/cumsum_meta.h,sha256=4vSVZtKdalwlGLc3owEnUdk2YNgBRHRys7DHNpIjLC8,873 +torch/include/ATen/ops/cumsum_meta_dispatch.h,sha256=9IfEv-67XlD-v7dmaHBLxQ0Gr6yf6mFG93-GTOC7UU8,1437 +torch/include/ATen/ops/cumsum_native.h,sha256=dE2UQajQABC8PvVBPSn0x2VbnHqbQAeQsaKmRpvzLJE,1283 +torch/include/ATen/ops/cumsum_ops.h,sha256=5NKUPWuwpSnn4mceD1zXHD2wn2CHD-Of7ATpu-AmrEM,5093 +torch/include/ATen/ops/cumulative_trapezoid.h,sha256=ziem09D79D1butfFv8Dd4rrrxUqg1QOftKRbS2IFxuw,1265 +torch/include/ATen/ops/cumulative_trapezoid_compositeimplicitautograd_dispatch.h,sha256=hJLSc1F0xE0LQk8LlNNMCiYDrH0vnGQbACmr4Ak_g00,1173 +torch/include/ATen/ops/cumulative_trapezoid_native.h,sha256=gyy97gUefEHb2J_VQL3tGJpPkSOghFaTUKzxCVoLxic,885 +torch/include/ATen/ops/cumulative_trapezoid_ops.h,sha256=zPV2nS3-C7ytI4KS2CiNVnTU-kikBVsrTr660BKcPOY,2043 +torch/include/ATen/ops/data.h,sha256=TV23nqvF1-IUUNP8Rmp5CcKR_YdrzbPi7AtQaLLjBnY,755 +torch/include/ATen/ops/data_compositeimplicitautograd_dispatch.h,sha256=OHGdDIOt4M8aPru4Tf3UgK9A2bPLMWeDY63rJ4bBxo0,1016 +torch/include/ATen/ops/data_native.h,sha256=LW4mOj4nzjGJNchppy6q01d355x5oGBCLMie-CXKBvw,728 +torch/include/ATen/ops/data_ops.h,sha256=b9RUO92occsbRmvb0XlairjvyCY-wk8ECIQvq5wINJQ,1205 +torch/include/ATen/ops/deg2rad.h,sha256=ZDutqzy9rtUbTTJ_lqfxRw4XTTj5_F4FLgUJPMUZAO8,1436 +torch/include/ATen/ops/deg2rad_compositeexplicitautograd_dispatch.h,sha256=8Uit0Z4HDWkgc-OQowCpChdfJUXjtpDi6NdK5_pPvCg,1230 +torch/include/ATen/ops/deg2rad_native.h,sha256=TYnm5nEVQLaRJ11TOs4SnS7BnctLmGOjAPjQVys4Xuo,1288 +torch/include/ATen/ops/deg2rad_ops.h,sha256=RpoxGAWWtoK8_jw8DrOI91KF9niO7uwLJ7hmupsLv-I,2309 +torch/include/ATen/ops/dense_dim.h,sha256=h5QI8jNXIGVA88k7k5D2oUHuthCG1McLB1P6eAE00Dc,760 +torch/include/ATen/ops/dense_dim_compositeexplicitautograd_dispatch.h,sha256=H-LxjmR3ABztgjTdR_htNeSZe3fjMhyzTlLpUOz3drA,1018 +torch/include/ATen/ops/dense_dim_native.h,sha256=1ZnKfv1v9SnXl__S0ZjtarYR2KpimTnrr0h5rBIJllY,864 +torch/include/ATen/ops/dense_dim_ops.h,sha256=ZL0S4zwUp52mQKjMpojFeNwxP_b7E_bWxzvjNv9h884,1208 +torch/include/ATen/ops/dequantize.h,sha256=urVj_PRq46lUNzGXdu8y9IcZnovkqcw8vhKHMu9N1Fo,1974 +torch/include/ATen/ops/dequantize_compositeexplicitautograd_dispatch.h,sha256=BSNJCJYJNbXCq39AvKkVxtxedSYmZwQw4vbJYy8iqDU,1280 +torch/include/ATen/ops/dequantize_cpu_dispatch.h,sha256=SUHNM2pB13Nk3kI6K81Vfn61-mmBl8PVutFckk0umY8,978 +torch/include/ATen/ops/dequantize_cuda_dispatch.h,sha256=PLGC5IPiaHdh4HI0epknETUmTBTYQP_EQqtVp7Jf3Kk,980 +torch/include/ATen/ops/dequantize_native.h,sha256=6HCXsQcLRuejIf7Vozo1xZt0e8vSMuS9AKJSROkCdME,1078 +torch/include/ATen/ops/dequantize_ops.h,sha256=prcEi9M_4YF_JnBaeG824qsnub4Gw-Vh8tItaCLJikY,3036 +torch/include/ATen/ops/det.h,sha256=jczO8csSP5hHd11SdUiFCJX94uL2GKgZLYTnd3q6pIs,879 +torch/include/ATen/ops/det_compositeimplicitautograd_dispatch.h,sha256=B2uhUQDffTv-bkxIXXZpWk-UKf6vgLlkcAuxpNHadjw,1015 +torch/include/ATen/ops/det_native.h,sha256=-DhbUfDpzsubnx0HjK0_xKS0kK-8Ons9i3lCum9pREU,727 +torch/include/ATen/ops/det_ops.h,sha256=m1XcnujJooJ7cFga14-mv83VieVjDwnOPrSz9dL7mhY,1202 +torch/include/ATen/ops/detach.h,sha256=7WkxgBHIDTHVUYbwXHYNcanAP776PW__GMkds4AM1uA,1039 +torch/include/ATen/ops/detach_compositeexplicitautograd_dispatch.h,sha256=0vTOlltJ85a8L6IqxnwTz4F36xKiLmKOjw7HS8XxZOg,1069 +torch/include/ATen/ops/detach_copy.h,sha256=1Iyoh6pazDN5gNzj1F6oCJvTBY1kshy6sqM0qnj2uA4,1331 +torch/include/ATen/ops/detach_copy_compositeexplicitautograd_dispatch.h,sha256=ICrGZ8XGvGoWEzTfqjyQgXobodtlZ81DIsmJKXI_MX8,1131 +torch/include/ATen/ops/detach_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=SqKzf2xzbYt_R2_4bf-jsxZIIoOfpqHdslVVwqh6L8o,1049 +torch/include/ATen/ops/detach_copy_native.h,sha256=rJjufj_HIzh38Y4eMmK98WEFXjLv7S3Kdy4eniJJI5Y,818 +torch/include/ATen/ops/detach_copy_ops.h,sha256=ZHTi9XtoZfEAto5TX4A_WTQRWFp1AtA12EgCkbb3kNo,1831 +torch/include/ATen/ops/detach_native.h,sha256=leB-3tOPgzMfyqYj4aGZuq7hp2Psz4aUVQhfzFZOJ00,781 +torch/include/ATen/ops/detach_ops.h,sha256=TRFHRQie_ydJaWDVkzXFZrVuSlYK7vwkmIDY4-Gc3PQ,1716 +torch/include/ATen/ops/diag.h,sha256=npmLuxTu6iSN-q1w4T0QLZmIHGF4zry37tB3bDgd9qQ,1397 +torch/include/ATen/ops/diag_compositeimplicitautograd_dispatch.h,sha256=VnMMOs_BI1uTOeu7m7lh5x1mVi7v9Gni7jy3fZyHw_w,1227 +torch/include/ATen/ops/diag_embed.h,sha256=qYt_UU4bidVK-zNgBNq9RiB_ghVi6ONrNLPHuk2jk5I,1649 +torch/include/ATen/ops/diag_embed_compositeexplicitautograd_dispatch.h,sha256=lWUVeAO3kJs06hNiOGTIBftwhy-uWMU41hA_hW9wlw8,1225 +torch/include/ATen/ops/diag_embed_compositeexplicitautogradnonfunctional_dispatch.h,sha256=IBZZGoM1qXtWPAN8OpZDFNX1PgbaW5c2-Mtk2QPnA8o,1100 +torch/include/ATen/ops/diag_embed_native.h,sha256=5epC0Z5oYa7thdrn_OZVoyS-S3Rag79_5CXlS7DcGkY,912 +torch/include/ATen/ops/diag_embed_ops.h,sha256=hxVLWRPTT5lwZi0xKFsM5SJLlZa6mKNeMPLnJus1wSQ,2135 +torch/include/ATen/ops/diag_native.h,sha256=LR16KF71HKv_7sgex2BFoKu3v26FI34V_IK-PiWWRVU,842 +torch/include/ATen/ops/diag_ops.h,sha256=r0Z5JtUofeOuOZKPJVZeBykbZdZI42W7oWCr_dXSna8,1911 +torch/include/ATen/ops/diagflat.h,sha256=Dxb_iNbtMfseol9iZBJRoV1UTYcHMyOrcffWOoGcHwM,939 +torch/include/ATen/ops/diagflat_compositeimplicitautograd_dispatch.h,sha256=P0FVDLPD5l-3Gf3xSB7mQ03-K41Q0zb54tT0k7SG5dY,1038 +torch/include/ATen/ops/diagflat_native.h,sha256=pDg3-n6t4VZAqR2TlWL5fx2kK4RBpR7ANIxJ0I5LvPI,750 +torch/include/ATen/ops/diagflat_ops.h,sha256=aoNV5awRbRyCI2Ok1lwQKjHVUPh94k5y19petTX9k70,1272 +torch/include/ATen/ops/diagonal.h,sha256=6-vDs-uCrpMnS-7oWZ8imYJW_aThLURgP-ujXhg-oEs,1339 +torch/include/ATen/ops/diagonal_backward.h,sha256=R37XNX6jEepjng_SulC4iMUfY8fCgY8mjHXao9bICYU,5466 +torch/include/ATen/ops/diagonal_backward_compositeexplicitautograd_dispatch.h,sha256=LwsJiHRPnjPxXWcy95F1lNBk6mmmmYq7bfXULHsQuqQ,1965 +torch/include/ATen/ops/diagonal_backward_native.h,sha256=bd22EjrPPLJ51zcpz67YtvLO2JgEAeF1UpGaqiI1aFQ,1012 +torch/include/ATen/ops/diagonal_backward_ops.h,sha256=lsq8k7Zr8m9CWv9oYSMP9AynamgyDHk_JvbOK7DFJO8,2421 +torch/include/ATen/ops/diagonal_compositeexplicitautograd_dispatch.h,sha256=QyJQsA5HARvowGz6Rm2oM71NzjmDnrt7wJVxIqzs7xw,1070 +torch/include/ATen/ops/diagonal_compositeimplicitautograd_dispatch.h,sha256=b24IjHUIJpRXlZ_6t5zUEf-k9RCJ-Ob9jUj7vAE7rsI,1094 +torch/include/ATen/ops/diagonal_copy.h,sha256=98BOPZ80hP-N8GN5D9XO3p--leVmeJT3Uec8LOYQqnA,1669 +torch/include/ATen/ops/diagonal_copy_compositeexplicitautograd_dispatch.h,sha256=p1DGiIxdSy75uRfRkcqW-zhnMDqO3Bj_m9vNerBn5Wc,1229 +torch/include/ATen/ops/diagonal_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Wz83kIAFeFTpV11LuGwxBJ58gSLAMwTe2szGjlxVeEA,1101 +torch/include/ATen/ops/diagonal_copy_native.h,sha256=5V6V3_SAwOYZ5OQXpKkNQoOV5HGIjVoLzZAfmADFMv0,916 +torch/include/ATen/ops/diagonal_copy_ops.h,sha256=Q6jQ2EjDHV9Q32YuryQUrniST1rf3i3lujyM06nyL1g,2149 +torch/include/ATen/ops/diagonal_native.h,sha256=Rc9q1mXHuUWP4crSIlY0Ko98k8LR5qVjziSjKBDIrsQ,912 +torch/include/ATen/ops/diagonal_ops.h,sha256=TgagXCDFfapg2BMJ8UMyW1nown5b0dzmvZWoFzYT_oQ,2164 +torch/include/ATen/ops/diagonal_scatter.h,sha256=m_avua31mmIK4kxco4j3AXu6iXmqpL84iN4EWVZdW_A,1822 +torch/include/ATen/ops/diagonal_scatter_compositeexplicitautograd_dispatch.h,sha256=QDoMJTb076wKchX3cqBwuCRojIHOkRmBpFgifyhqREQ,1283 +torch/include/ATen/ops/diagonal_scatter_compositeexplicitautogradnonfunctional_dispatch.h,sha256=_zmPeLNcYUZ4cZ01acqw5eP44_qqiveJx9YU3nkM8dw,1128 +torch/include/ATen/ops/diagonal_scatter_native.h,sha256=MScxR2seL2a8CklIYTdz-VPgCRs9CTfIAIUnb2JtQFM,970 +torch/include/ATen/ops/diagonal_scatter_ops.h,sha256=J159wcm0Wb_W0EggqN3rvjU5yIrXqPVzK6_pS45NmPE,2327 +torch/include/ATen/ops/diff.h,sha256=G3JPIXwoFJjLgpUi33QMVdJOTmBA_lQk5T2OaHWFsPw,1889 +torch/include/ATen/ops/diff_compositeimplicitautograd_dispatch.h,sha256=EH4Zb8rN_56BwG9DsNeZyYWJfXj2vassCz7XmpTPVLg,1530 +torch/include/ATen/ops/diff_native.h,sha256=zJ0D9Ytoa_08JgZ-1cyZv02MOlBdn1BISAm6gr_oXnk,1041 +torch/include/ATen/ops/diff_ops.h,sha256=P_mHceTuhC5IE5vHS6SWuPQrxCRFQtisNRrGFZNKkTw,2553 +torch/include/ATen/ops/digamma.h,sha256=__YyiHLpeo6QlblwJo4IUD60lcTMbGxko3H2XwLHeKI,1291 +torch/include/ATen/ops/digamma_compositeexplicitautogradnonfunctional_dispatch.h,sha256=G1-lQ6E0FmVzwBgMSqVCIr7m-qV6lRwyVAx0IdGIik8,1097 +torch/include/ATen/ops/digamma_cpu_dispatch.h,sha256=J7Es2jdFZCwQs6_gjnrD5yac_coy376ZaUw6qUEEjsc,1186 +torch/include/ATen/ops/digamma_cuda_dispatch.h,sha256=ugmHMoobUQcb3W6zZEDQ4nAf91FLhiI7YvcS5Te8YL4,1188 +torch/include/ATen/ops/digamma_meta.h,sha256=mqYuU1CysSjg-o9lwrW9xZAf3VrhhcPTgXmOfsmOAo8,822 +torch/include/ATen/ops/digamma_meta_dispatch.h,sha256=M3rPChCNkJ6JjaAMghOc_OHLC9exDerG4809iT1cHsw,1188 +torch/include/ATen/ops/digamma_native.h,sha256=3C0xpjLprAZxmanX1e0cUk15E_29NaA8f1rgmZY8qyM,853 +torch/include/ATen/ops/digamma_ops.h,sha256=z2qiKgEnoXZ81a6BHX5macsYkyUprHoo5DwsPq8QWDo,2309 +torch/include/ATen/ops/dim.h,sha256=aboV7ScjyD2KAcF9OG27768Vp4HoOVBf-so6D-tnfGE,754 +torch/include/ATen/ops/dim_compositeimplicitautograd_dispatch.h,sha256=4Y1uEI-FnXSpa42kLWqRn1uLo7WMpMw-HpyzUXo0G5M,1012 +torch/include/ATen/ops/dim_native.h,sha256=oCGLqnVXDk3rFYnowkZaf8YyNBpK9UuMVReX0Httfzg,724 +torch/include/ATen/ops/dim_ops.h,sha256=-iFYOeX0iOLFg_6ohwwtAXbAMxlDFT-nig_QCgCj-6E,1190 +torch/include/ATen/ops/dist.h,sha256=RGfDnzA9c06_snD5B5DhHlo0dpEZcyobnZH6_S3kFPM,1517 +torch/include/ATen/ops/dist_compositeexplicitautograd_dispatch.h,sha256=i6QhidUOeNQpn5-eJFRCrOjEZ_lA3zMkon68-9OKCu4,1317 +torch/include/ATen/ops/dist_native.h,sha256=IriwHkiPCd5ZJNnA6ZDn64re-u1W84z2CeKHgtOn_m8,902 +torch/include/ATen/ops/dist_ops.h,sha256=oHqPzPGY-M2GWydMjiWRU4U_p47URTilWMZAihWAVIo,2113 +torch/include/ATen/ops/div.h,sha256=DM_ytex7IWnzl8KuItLoyy2Iff4wpryYEB6ekFpH6Ps,3991 +torch/include/ATen/ops/div_compositeexplicitautograd_dispatch.h,sha256=jqRoRmx5r68_3GDXNgnVNWiZ0Igl5V6KktibSbOLO2c,1868 +torch/include/ATen/ops/div_compositeexplicitautogradnonfunctional_dispatch.h,sha256=whSTYQZxAHeOVeapdR0GCY6mG5CuYtsZt6MZTtnaIr0,1390 +torch/include/ATen/ops/div_cpu_dispatch.h,sha256=310bHR6uRfb7GgZ4y69_-h4RyatMsODeRPzRmPFIgHU,1824 +torch/include/ATen/ops/div_cuda_dispatch.h,sha256=e3YOYIl4SCPQIwGi78HTChuTcwl5En5rPWwwUKzaxV4,1826 +torch/include/ATen/ops/div_meta.h,sha256=j_0xV0dF_uWPUFh9WUV7RJwbhLq09UWc5IF9mP7g8aA,1045 +torch/include/ATen/ops/div_meta_dispatch.h,sha256=fJJ4Lt4auGV1_SrUGQyhIHXk-eB4ZZKitKnAUJtYYcg,1826 +torch/include/ATen/ops/div_native.h,sha256=mhfqiXM_l4P4u9d3om-GpDzznxnxQvGBpt8joKe_ufE,2762 +torch/include/ATen/ops/div_ops.h,sha256=2dfhhisJIVcZYwXf8wR4ZhGI_L-ogD8NL0rV2ge5pP4,9213 +torch/include/ATen/ops/divide.h,sha256=Y9V-QfK_BEnEKZlwJglmf2iWHQ7V_vC1h7Sn4qkL5WI,2890 +torch/include/ATen/ops/divide_compositeimplicitautograd_dispatch.h,sha256=Oa-uiYExQf4TBSs2v6CbHKZNQZsD5i6TrqhSsICYX2M,2304 +torch/include/ATen/ops/divide_native.h,sha256=-FSAo6PLmfpQm69GPFMG58oVjuYhfcm-rLrREW0L2P8,1757 +torch/include/ATen/ops/divide_ops.h,sha256=F-8TYrRsXip1BRhhRVor6gyiBOcm8Urf1w1mOyyAYU4,7759 +torch/include/ATen/ops/dot.h,sha256=hDxawQL2lna5PrlIaufZrCbACF_rcoHJDth90DShQ8g,1401 +torch/include/ATen/ops/dot_compositeexplicitautograd_dispatch.h,sha256=wcw5KQn5pJ72LPBKQXNRWmGgaORWWWsDfA4QRHxqZno,1169 +torch/include/ATen/ops/dot_cpu_dispatch.h,sha256=cjLAj9j4idJZ9iuw7vyhaNVvjgOyPjzMFyvyRU_T_tM,998 +torch/include/ATen/ops/dot_cuda_dispatch.h,sha256=fpqckH--kSA6f0xnufNmbmRjmUWLiTfuFu_WD85g7fA,1000 +torch/include/ATen/ops/dot_native.h,sha256=KqLQjyiLykDN-wb8n6Nuq_yZA1RdoaVJ0tu--Hs14RQ,939 +torch/include/ATen/ops/dot_ops.h,sha256=DHXMBOESq_NdyzU-0UY7jd91vP6PuaEx9LKOHP24G14,1961 +torch/include/ATen/ops/dropout.h,sha256=ffgo8SUCRGpVNc1wGajUUor3J5nsZ2m9mchvT6T2ogI,1149 +torch/include/ATen/ops/dropout_compositeimplicitautograd_dispatch.h,sha256=A62lu5eajcqjz-PvdfrsJrLLQI1urzSB0lQnzsLnWm0,1116 +torch/include/ATen/ops/dropout_native.h,sha256=UCpMFQHq58nA8Tfkw-lyshJSg9Li0qf43OXi1LP33RE,828 +torch/include/ATen/ops/dropout_ops.h,sha256=nBlyQIs8qckruwwQap5zz979xW18ur9vrJoC4RjnVQk,1877 +torch/include/ATen/ops/dsplit.h,sha256=K71kl82VElRd99-JOukwK7SRVpnNmpB_aFGfiiKX5Yg,1193 +torch/include/ATen/ops/dsplit_compositeimplicitautograd_dispatch.h,sha256=-RGb4ceDOGIuKJjEXHXdJzsZZYvx5Ytj7wzlxl72kpc,1145 +torch/include/ATen/ops/dsplit_native.h,sha256=gKwPnQFs_neaSi-DShY2MpSBpvswLhhrm8_V05uzYOM,857 +torch/include/ATen/ops/dsplit_ops.h,sha256=thGZzVGK40bWJCEQXi_mB7XAhsyj_B9pz-DxgJM7D2o,1996 +torch/include/ATen/ops/dstack.h,sha256=tvO9rs5SNNrs0T4AujeqXpZ-TPwrkY4nZ9N_6_LANTI,1302 +torch/include/ATen/ops/dstack_compositeimplicitautograd_dispatch.h,sha256=I8O2hHEFfPmP8xAZAU2ICtXjcSw0JvYK71FZ0CG-B1A,1172 +torch/include/ATen/ops/dstack_native.h,sha256=wgq-TUR8kydVd-srbZWHlTdUPwICQs_3vBt4QQMuHGM,806 +torch/include/ATen/ops/dstack_ops.h,sha256=jE3VRqZZ-N6m2g3eQtpxtIdVQk9kMousI1FlWlj7CPE,1799 +torch/include/ATen/ops/einsum.h,sha256=D3a0HC_yNt_12sMPesuQqdCFX_qI6TouIIjSyJmfCig,1021 +torch/include/ATen/ops/einsum_compositeimplicitautograd_dispatch.h,sha256=pb_pi68lRsZQvqT4hkC2QrNiyA4vrC1EpRYccqHGULY,1089 +torch/include/ATen/ops/einsum_native.h,sha256=S14xJR_nmLjzUYMwCtdzrsI0SIBa0mjB5xHKaMlu1N4,801 +torch/include/ATen/ops/einsum_ops.h,sha256=1YmRg95kg8CaFOn1jISe3ZHgF3IA23rq6OD9o2ASzaQ,1402 +torch/include/ATen/ops/elu.h,sha256=geP02olb1-qY7fAmw56HE35zcyXcOLnI9AUlmhmhJHU,2062 +torch/include/ATen/ops/elu_backward.h,sha256=5nUKmINRGnl8vqHR05Xu4yHQMtZZReAJwKQqXm6--6I,2302 +torch/include/ATen/ops/elu_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=BCKv8Rl1OWVDYagC-1Ce1thb3MkztI1mk5O0NXq_waM,1192 +torch/include/ATen/ops/elu_backward_cpu_dispatch.h,sha256=FOA61fDZ73rBKQIqclj8UEH8oV3FLLTjxlnbBCUhx5M,1589 +torch/include/ATen/ops/elu_backward_cuda_dispatch.h,sha256=A15wtBV05y3mNuwdVARSaARNt7Xyb6vIKRAIIL9N594,1591 +torch/include/ATen/ops/elu_backward_meta.h,sha256=cuMqof_dpN44dU6kAG0camgQUCHh-PqgSUum7xknge0,969 +torch/include/ATen/ops/elu_backward_meta_dispatch.h,sha256=HgrJZjt5hf0G0lDNBxESqstcDBYQwAb09Ul9LjY7Cyw,1591 +torch/include/ATen/ops/elu_backward_native.h,sha256=L58MWiCBc1hxBd90tMo11M7syJdSFBgt-hiYOXobBLc,1017 +torch/include/ATen/ops/elu_backward_ops.h,sha256=jOfVQW8jUwDzeVvBumeVBcBg4ltpjW3__dX50X_JeDM,2807 +torch/include/ATen/ops/elu_compositeexplicitautogradnonfunctional_dispatch.h,sha256=baUpx2rkErCKs_oLurmv4GUxlsXX6BL1sxFIJ7fGSho,1269 +torch/include/ATen/ops/elu_cpu_dispatch.h,sha256=y9DShreyq2Vw-O82O5qUCh1BI03Wnbt-6hIlgF9xgIY,1524 +torch/include/ATen/ops/elu_cuda_dispatch.h,sha256=ME9OBP8EdSjqbthwgq06jctUj3_zounYAInQnmhAmq0,1526 +torch/include/ATen/ops/elu_meta.h,sha256=UZgQceU3Qt27H796eHvvYw47hAqprBNGFtJUb_DEsjM,902 +torch/include/ATen/ops/elu_meta_dispatch.h,sha256=6u1fbb1-P9ryRfHAEVTxMnYhra0aBQpcci8lY2Zazog,1526 +torch/include/ATen/ops/elu_native.h,sha256=qm2ltdXeKmStN9l2brxvoc0Q398Oaiq3W3TKY3-kVq8,925 +torch/include/ATen/ops/elu_ops.h,sha256=B2AV3R8ImCvKtQB_N_xiAaIU0pc-kSRtxuVoWu8esCY,3119 +torch/include/ATen/ops/embedding.h,sha256=VXZlgOEHobJWEqmR4MndMdOwj_acjSToAl-8dDkmUKY,5614 +torch/include/ATen/ops/embedding_backward.h,sha256=ZbyfA5bhJwGnCs3I0qY--Y-ygsqMHpMf52BEQ6eUPM4,2445 +torch/include/ATen/ops/embedding_backward_compositeimplicitautograd_dispatch.h,sha256=9oLX0S9dGk-bgqN6DwfBXQwjtfyeozsjT5nvPMyoXww,1327 +torch/include/ATen/ops/embedding_backward_native.h,sha256=cobHiJZ2dlOJ2lPx9OFXXErT5lQJb9eL8tAfmtM_ew8,865 +torch/include/ATen/ops/embedding_backward_ops.h,sha256=XSknOoabtnW_quarg8o8Uc5YUPwU6nfOju3gIJ-JRVQ,1631 +torch/include/ATen/ops/embedding_bag.h,sha256=bvZE-ccv4K2Q5dS6MVyUL6YX0x6wrCfn3x3kF3AKVTs,2208 +torch/include/ATen/ops/embedding_bag_compositeimplicitautograd_dispatch.h,sha256=z2Hlk63FbSyf_et9yE5x2yN3PjW1G2wtEf86UYBmWB4,1625 +torch/include/ATen/ops/embedding_bag_native.h,sha256=_Nxd1dFBAPzB3Fe1pFWITFVD_BtUpLtXMWpMgTmoTts,1337 +torch/include/ATen/ops/embedding_bag_ops.h,sha256=6S9H9Ws95KiDe_Z_ab_lus-_pQ3olNtXAz0bI4vszeM,3515 +torch/include/ATen/ops/embedding_compositeexplicitautograd_dispatch.h,sha256=X9wEeoa3zC3sKdFDFHn7Ktc1qxnPq9VgBD44J7sCqaE,2031 +torch/include/ATen/ops/embedding_dense_backward.h,sha256=V0VD1_9FWwJmZa_BjQqfRdD44WGhZih9K3SEsy7oDu4,6241 +torch/include/ATen/ops/embedding_dense_backward_compositeexplicitautograd_dispatch.h,sha256=ua4FPXRvYb1BJFIZMXjuIq5fSYCQMJrq5yyxTMyJzrE,1788 +torch/include/ATen/ops/embedding_dense_backward_cpu_dispatch.h,sha256=zwa51u0ncbal4FAZsJouJ8cKGSF5si1YcGgrRMkT5nw,1283 +torch/include/ATen/ops/embedding_dense_backward_cuda_dispatch.h,sha256=j-r5ABK_xdDlm1twqGrocOvBiMa2nFr9Ne5puL8qiK8,1285 +torch/include/ATen/ops/embedding_dense_backward_native.h,sha256=xCx4lj1WCt_Z2Ks1nWmD2Pufx9RRKiVmuoUtdQLbpPM,1246 +torch/include/ATen/ops/embedding_dense_backward_ops.h,sha256=9NSiVZo4QX-9F1-DdAZwVsezhi5-uRAwo6bZslqw4X8,2629 +torch/include/ATen/ops/embedding_native.h,sha256=cjL2FfhDAtpYbWUazK9kzTbEVmgkWfIKAgC-lUpGuts,1203 +torch/include/ATen/ops/embedding_ops.h,sha256=iJAQUOuP03GCc3tjvpZ3fwjc_2NNlYYcvBvChA4bzHI,2463 +torch/include/ATen/ops/embedding_renorm.h,sha256=jGEnaPdejB8jSpmVd6QtFHUuv9IRLfq5zqhhDe28yjc,2121 +torch/include/ATen/ops/embedding_renorm_compositeexplicitautograd_dispatch.h,sha256=kt2IrehVKKhjrr9dWCQGVEGJgy-a8wCcjrT8O9SoHhA,1394 +torch/include/ATen/ops/embedding_renorm_cpu_dispatch.h,sha256=OILyb8p8mrUfJiSbU3DHG_doMQSf0SgHfeQJddDb0ok,1044 +torch/include/ATen/ops/embedding_renorm_cuda_dispatch.h,sha256=RGTLkfRRY7F6y_nLAjoP9PzKE1FtWJr_Vt1TcwkwSig,1046 +torch/include/ATen/ops/embedding_renorm_meta_dispatch.h,sha256=eeBiLc_7hoapRnvjc6EogQaO60COpr_1qpHt4swUBXg,1046 +torch/include/ATen/ops/embedding_renorm_native.h,sha256=080IuuA0PMfhi5J-0WNEcl4CihtKNzAO72sZdxMX7Ac,1211 +torch/include/ATen/ops/embedding_renorm_ops.h,sha256=XwwEveXxTvrNz7EGMp4EWYN3D7P5Wd94fCjKI5ZWuVM,3023 +torch/include/ATen/ops/embedding_sparse_backward.h,sha256=jrSdzjOKHcULM1hFOpfIVRrJ2_ueWRJOQz45L2RbvU8,1192 +torch/include/ATen/ops/embedding_sparse_backward_compositeimplicitautograd_dispatch.h,sha256=9pQ4jqFBnXneOWDfRPTKHkzDwNAYYdL7awWbdT_qa_E,1132 +torch/include/ATen/ops/embedding_sparse_backward_native.h,sha256=AE0MzTSqnSWk3zAlE4oxgzXYluvm-ZADy9ks6Ur67EU,844 +torch/include/ATen/ops/embedding_sparse_backward_ops.h,sha256=ibDyI8UpizWRB3HYjf-gQHYxd1WgeuQzrzs-oIaKbZk,1577 +torch/include/ATen/ops/empty.h,sha256=E_sO_rri6_vgvkYqw5oeWrMg3XN45r37Q8L-yx4lxn0,9482 +torch/include/ATen/ops/empty_compositeexplicitautograd_dispatch.h,sha256=oZ8R_jfxK65sS99AlwCCE1nEIwHNO7QzNpLWvpD1flo,1774 +torch/include/ATen/ops/empty_compositeimplicitautograd_dispatch.h,sha256=ku65s8JHL5Ow3Eq8U2jVvRVstMeKVFyhKu7v0YGJEAQ,1510 +torch/include/ATen/ops/empty_cpu_dispatch.h,sha256=6iQ5rpUED09D-jLOu8bqfM1s8hp0bY-4dYGs1ExTQ2M,1718 +torch/include/ATen/ops/empty_cuda_dispatch.h,sha256=6NipKv-tGlF9wZZABIsPVhrQtJFCqNIRzRPYLRxdKuI,1720 +torch/include/ATen/ops/empty_like.h,sha256=dnLoJqk_Olujv8LnpLM_15zZThPdvXu5wmNVjPOT_p0,2475 +torch/include/ATen/ops/empty_like_compositeexplicitautograd_dispatch.h,sha256=JSHN3P1bs6ta2dTb_TlvXnhNMsuZt9GKbnMyMw5UUX8,1646 +torch/include/ATen/ops/empty_like_native.h,sha256=myr7KZbvEauIshBb8LiAIdZua6uSI3-gMipFg7Yao18,2241 +torch/include/ATen/ops/empty_like_ops.h,sha256=vUKmhPn6RopZy2Bo98BlJxctrdvOxX9GXt2fVocj6gs,2655 +torch/include/ATen/ops/empty_meta_dispatch.h,sha256=-46psyLd8P9957SHSpXmFWhEq1JS3_EG55gpz6RpPy4,1720 +torch/include/ATen/ops/empty_native.h,sha256=uNW3rx-w2NXk5skO9JlahinwyFPl7cffE1hJdarDRnc,3846 +torch/include/ATen/ops/empty_ops.h,sha256=mmlimGY-RTfRaXZJAKqK9wLhLYWlBzW6yHnhWrIuvcc,4869 +torch/include/ATen/ops/empty_permuted.h,sha256=tdRLCvKQbDIQ9GTob8i1q8TXDjt_tZB8ulb7MisqO7U,7032 +torch/include/ATen/ops/empty_permuted_compositeexplicitautograd_dispatch.h,sha256=1tmPQnli_yinJV-r-3zJqjq5cwXt1Fs3uGFKVHr7u0U,2192 +torch/include/ATen/ops/empty_permuted_native.h,sha256=1xFuLCm83kiTQSTU3xLRXFnTbhmWcRg7aySoVKN9DyA,1063 +torch/include/ATen/ops/empty_permuted_ops.h,sha256=Yx1pw0EGRJukhgFuex89NyrjNIUMwvPryMoRt2Onq-4,2567 +torch/include/ATen/ops/empty_quantized.h,sha256=W1u-gjSEPeRiwPUE70HYsAkXaIFw59h4xP_GU1NZjIg,2736 +torch/include/ATen/ops/empty_quantized_compositeexplicitautograd_dispatch.h,sha256=YafgEwi-ZJin0TdyrlgkMKpBfu6aQgO6LmrxmkkwF9I,1302 +torch/include/ATen/ops/empty_quantized_native.h,sha256=EsAB1we283m5Eo1DgNrG89KF5ZloS3WoIeBnSEgN7Bo,1146 +torch/include/ATen/ops/empty_quantized_ops.h,sha256=AMPzbmaQufqeQJKC6AbSdkw1Z61UoM3HQCej_0bcYik,2849 +torch/include/ATen/ops/empty_strided.h,sha256=8aqvFYjwEcBdsXVacmcbC8WcCVQqSraI3p02A_6kCo8,6895 +torch/include/ATen/ops/empty_strided_compositeexplicitautograd_dispatch.h,sha256=xhSMpKVmVfiQaNwGIcq7WV3yNAXLQKJd2mRcuEMjV-U,1420 +torch/include/ATen/ops/empty_strided_cpu_dispatch.h,sha256=06o1S6KGVG33CehvBYBu9eT1BoKBYScNv0ArzzybI-k,1628 +torch/include/ATen/ops/empty_strided_cuda_dispatch.h,sha256=ngvPLf8nUvMdTDpE1iF_mS7-kRrULjDlAEvNg5e1X-M,1630 +torch/include/ATen/ops/empty_strided_meta_dispatch.h,sha256=K7IC5qx22Pntiin98KXzfm7tNTbrG5GvCvx9T5DL8HQ,1630 +torch/include/ATen/ops/empty_strided_native.h,sha256=hrdn1rzoZyGehfnzWkk6kQPNbpYMWdSi-P4LaWqpYOY,1800 +torch/include/ATen/ops/empty_strided_ops.h,sha256=bUoaB54xIOOSfO4VKnICAQI3ILwqTOwyUkkLtI9cZEM,2537 +torch/include/ATen/ops/eq.h,sha256=jbiBR_T-cDpsD_icbEZsNdl76sLKZiD83Fj_KQj18Ic,2096 +torch/include/ATen/ops/eq_compositeexplicitautogradnonfunctional_dispatch.h,sha256=_h16vxOSj6JM0iPdBfhjY2g5eaMsVhx3iJ4-3KL5qvk,1288 +torch/include/ATen/ops/eq_cpu_dispatch.h,sha256=p3xdTFcaS2uu3jyvNw6PDDQ_CKrThfqb3rJ0PWNm97M,1620 +torch/include/ATen/ops/eq_cuda_dispatch.h,sha256=YcfuDTYBTe-Hp3ifsK1JtR1QRZbcpZlqGICqTJKWytQ,1622 +torch/include/ATen/ops/eq_meta.h,sha256=ZfxJdGriWEPf5VOGd0VfLWtopp31JMoZBMXg3BnFDbI,989 +torch/include/ATen/ops/eq_meta_dispatch.h,sha256=m-AF4j2HXN6-RJETYIx65S81_kbNroV-4I-QKfx6eek,1622 +torch/include/ATen/ops/eq_native.h,sha256=h9HJF-2Fr053_TtC8GjGw_NdIErC3xqqQeTeQpfqtJI,1639 +torch/include/ATen/ops/eq_ops.h,sha256=YJ0XCXuKoChB6NUDpNW5gDRmOqey8dlRJxsjlattBew,4455 +torch/include/ATen/ops/equal.h,sha256=B1qej1C9ErRx4W4mir8puHicgCT3XMhYSczYO81J1oA,926 +torch/include/ATen/ops/equal_cpu_dispatch.h,sha256=Kg3CwVywunYYUcqiE0sMTyAlIfvnyRSegOb_ejq-kEo,993 +torch/include/ATen/ops/equal_cuda_dispatch.h,sha256=WnEg9gXMy9HzNX_oYdffDqOGlAPLzv7Gn6uQSbZ4Suo,995 +torch/include/ATen/ops/equal_native.h,sha256=4u0z-oBAbVo0uom0BT1tUsNrI3g6hYzDNf2RtkXcPik,918 +torch/include/ATen/ops/equal_ops.h,sha256=OWQeEjGANNFsG2LVJoA0parq-cu2y2zjODl4n9UD4D0,1274 +torch/include/ATen/ops/erf.h,sha256=-Jyqt7S3hLZnbic_e6IpZfj-piJJZXrY1culRJybpJQ,1384 +torch/include/ATen/ops/erf_compositeexplicitautogradnonfunctional_dispatch.h,sha256=aaY4RsKmcDhRjQTQs0w9zjNPv6qFqpxRYNN6IxDptvU,1089 +torch/include/ATen/ops/erf_cpu_dispatch.h,sha256=pcW8HmhwMlLqOCPpxsmZOuNIrzMf_QTazQgHj19zzKs,1170 +torch/include/ATen/ops/erf_cuda_dispatch.h,sha256=wgKacr_pKctv8amK93v334UHHWaXACXOZBYGLLB-F18,1172 +torch/include/ATen/ops/erf_meta.h,sha256=GG8-dyt4ovvx0SOpBNL5I7cVw93ZBveooufTRzt84jc,818 +torch/include/ATen/ops/erf_meta_dispatch.h,sha256=xqqYrV5LfF_eh6FY2OR_6afI6o9teIMZAts5AAuApgQ,1172 +torch/include/ATen/ops/erf_native.h,sha256=-jKnypdbIR7rnGX9Zp4EDhxcz-z36BGsvzbHm22kRiM,1243 +torch/include/ATen/ops/erf_ops.h,sha256=M0CoEgvMBcX80kAjytdJmkkv59q6feHWuVWE_Vr_PoI,2273 +torch/include/ATen/ops/erfc.h,sha256=3fhKgdJE6SHBN74Ge9LVUUa_WjlftLWAvlSLRKZNDNA,1397 +torch/include/ATen/ops/erfc_compositeexplicitautogradnonfunctional_dispatch.h,sha256=p_QN-mNWvtVZQh6pEfIbgX_nW-6mGIj0soUEK9Ig2ro,1091 +torch/include/ATen/ops/erfc_cpu_dispatch.h,sha256=7MqFKcNcp8zalFMN8Qskn-8mzIIiI6IQNAXaMxntCAM,1174 +torch/include/ATen/ops/erfc_cuda_dispatch.h,sha256=TJYzfkTR7Crd0rm7ZjWdEja8GGLp_EUWrjF4nKtxHMY,1176 +torch/include/ATen/ops/erfc_meta.h,sha256=7MRi-D7QH5s4rFHEdjLZYOs0WvTxp94BfXToDSNzgS8,819 +torch/include/ATen/ops/erfc_meta_dispatch.h,sha256=-MIAF7ORvLm00QM75MGflyBr0anj-A-xUkxdPk4OVz0,1176 +torch/include/ATen/ops/erfc_native.h,sha256=vukzl9L6b2yLrCEjp1orpBtu3TjDVO5NEYkUDVOxoTQ,844 +torch/include/ATen/ops/erfc_ops.h,sha256=dGs_VfG_sh0lT0khVc2X8pCrFU9HWkex_6e_aQ7tWqg,2282 +torch/include/ATen/ops/erfinv.h,sha256=XFrToqDxEennZ_A4nBCwjwdgSI19sNJ2FuSF4PDOkfE,1281 +torch/include/ATen/ops/erfinv_compositeexplicitautogradnonfunctional_dispatch.h,sha256=spa_KYP8Us4sENmQpwUdBS9Zx8T8MloytqkyZjWZTYA,1095 +torch/include/ATen/ops/erfinv_cpu_dispatch.h,sha256=eQ-0qRYrh06TYGl33Hj9irVngUjhoIBnDTuwropGGxU,1182 +torch/include/ATen/ops/erfinv_cuda_dispatch.h,sha256=arizfANIM1LTL2IKSNb_bP4-CPqM516v6tF5sJWPrb0,1184 +torch/include/ATen/ops/erfinv_meta.h,sha256=Xq4DfrgcaGXtEbApeKQlLIn_kWfB2x63VxWT7LVgMO4,821 +torch/include/ATen/ops/erfinv_meta_dispatch.h,sha256=2oRmI8jvH_tNQ4-CVa1RgYRfuB9VJdHLetIggkvyxWg,1184 +torch/include/ATen/ops/erfinv_native.h,sha256=i7KumcRxPO-o6fwymkGjYMJxJGln9c_10wSMlGFBq6M,1270 +torch/include/ATen/ops/erfinv_ops.h,sha256=SbPjRYY32oaURpFSk3HoqbQmhHHuR35GuvobiwSYMiA,2300 +torch/include/ATen/ops/exp.h,sha256=DC3Tw7Fm5uilEIeumDm1s2RZOLEHrj_iFuLSbNefE_8,1384 +torch/include/ATen/ops/exp2.h,sha256=W3FGcnwLlqr6NJJFOrOKBlYp82PYKkQnmOYZurehTkc,1397 +torch/include/ATen/ops/exp2_compositeexplicitautogradnonfunctional_dispatch.h,sha256=o50NiBKbX8vfSwSkpUtRIqRsOBh0EK109jpza1-IUcA,1091 +torch/include/ATen/ops/exp2_cpu_dispatch.h,sha256=abYoMbZCXR-1lnszZHf-GGvJiLSw_eOz0k9Wohk1MiM,1174 +torch/include/ATen/ops/exp2_cuda_dispatch.h,sha256=_hGNy4JRC_CzMnFnnUPTLn6zmSfoyk7ulKAVcKxnpFE,1176 +torch/include/ATen/ops/exp2_meta.h,sha256=m8gtBBUgkIkQ1cK4FJ44mx9oNPruOeRXnlaoB26kAWA,819 +torch/include/ATen/ops/exp2_meta_dispatch.h,sha256=pt6cs-uwHTI707A9v09_22pVmmNyQSWnslJC6coMyb4,1176 +torch/include/ATen/ops/exp2_native.h,sha256=7WlENldCopf4GFyv7eCRX-_zDfT3MNaEWcA0JY3h2qI,844 +torch/include/ATen/ops/exp2_ops.h,sha256=5z-tGauNj5uHCel0uYFs6Zex4C2P0mLb0PEDb2wGSZ4,2282 +torch/include/ATen/ops/exp_compositeexplicitautogradnonfunctional_dispatch.h,sha256=KpM3kWXfdW1TTuujUTgis-c98Ze-YUCcVr2WeExCYbg,1089 +torch/include/ATen/ops/exp_cpu_dispatch.h,sha256=Tjpdz2NwhGzA77Mn9DlZdq-TjWVML05JZcPOfnSFHl8,1170 +torch/include/ATen/ops/exp_cuda_dispatch.h,sha256=rjPRw1Zx-Gglf6dT77T66aNJW2yt4l5GeeEto10RD9M,1172 +torch/include/ATen/ops/exp_meta.h,sha256=hmbpqASi6dexagN1P2vkI2OfMppWcolzuhxP7nOfIZQ,818 +torch/include/ATen/ops/exp_meta_dispatch.h,sha256=Vwt198tZdsIJIGKMrMtEMuzHW382JpWqePasNFEu8yw,1172 +torch/include/ATen/ops/exp_native.h,sha256=o4S7V-PU4DYI96hyvJ0KkY0ibEV8i2Hfvv_5ZkgiCA0,841 +torch/include/ATen/ops/exp_ops.h,sha256=cmoOBrkS_1J6bJuRx0vegON4k6cJeo8FfmLoq7Xf2jg,2273 +torch/include/ATen/ops/expand.h,sha256=PezukEin25kZBEXY2tKbl0WvLI3IdYrD3fFrultfHwg,1298 +torch/include/ATen/ops/expand_as.h,sha256=ChCHxBu0I_f237DgEbRZYA_5pWtki-Cx_OPIttdwf-Q,760 +torch/include/ATen/ops/expand_as_compositeimplicitautograd_dispatch.h,sha256=ob0JbBH0eN3_hV1RCXCYDKrrB1zhfa6pQZ5ApsjNjTU,1047 +torch/include/ATen/ops/expand_as_native.h,sha256=MW0wIq_wR6hEtJieGgqozbsEIDz4vadJ_WiY36vSx_Q,759 +torch/include/ATen/ops/expand_as_ops.h,sha256=vM-KWBm9FoEa2Frv-wGv3C2bnGUOLKCXYaxH91orFgc,1312 +torch/include/ATen/ops/expand_compositeexplicitautograd_dispatch.h,sha256=eIYq4nY7QZvYsS3q1QKdxBbALr4PhBfQo8VLB1qxUUc,1169 +torch/include/ATen/ops/expand_copy.h,sha256=jgOt7AwlivhVa8gKcO9qCD3QIuNhuIzoDo8WQSUqDGw,4380 +torch/include/ATen/ops/expand_copy_compositeexplicitautograd_dispatch.h,sha256=YzTyn-w0g6jMkg9QUIJjeGPUiYWoKnVV_30Oe_y9DDQ,1480 +torch/include/ATen/ops/expand_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=0lXoxSKePgFRf_W3yvgDm-85pvldj6wFCoMbOcKWK8A,1205 +torch/include/ATen/ops/expand_copy_native.h,sha256=Lshhf3DWKcEPDbHNVoQxxEcQNJx_ArLPU_QIaP_dSw8,920 +torch/include/ATen/ops/expand_copy_ops.h,sha256=u0cn7LFUFg1H4hEKAH-E-m5TNBa8y65lltLPEi2uSyY,2124 +torch/include/ATen/ops/expand_native.h,sha256=uviKJ_7U7-kOxAkT_1GA0iX2qHwwsIWEdOxSCuZjMEU,773 +torch/include/ATen/ops/expand_ops.h,sha256=dNeK11hk_5eNaU_TplXdMnhMsRi0d5KIaNqUAlvCtr0,1365 +torch/include/ATen/ops/expm1.h,sha256=921qXV22Du8cn9ILwPQkgE-mWpWu-jT1qjmL5aIAzCA,1410 +torch/include/ATen/ops/expm1_compositeexplicitautogradnonfunctional_dispatch.h,sha256=XJKsf7e-R577kGHz084hu1hr9Fsbcm1uYgNaMvKIzAQ,1093 +torch/include/ATen/ops/expm1_cpu_dispatch.h,sha256=oxCmUv1cPWq_taq58ZpKvXCAbKgH7yHtCooYm3yEV-o,1178 +torch/include/ATen/ops/expm1_cuda_dispatch.h,sha256=mE-5U-nqsWPO5BukyZx2W-_dW_aBUJ-P81pTMNemYyo,1180 +torch/include/ATen/ops/expm1_meta.h,sha256=adHkytCmZBxNFqWxUUaVIgyYkr3SaOasjWyTNDYydLE,820 +torch/include/ATen/ops/expm1_meta_dispatch.h,sha256=oEhwLhzasqFzgGi1C9wXLQ_2SFjDJeprT3UspH3PUek,1180 +torch/include/ATen/ops/expm1_native.h,sha256=4XLabeEQ77eh92h4o4s9fgfhxoeq7vTCC8XMh8U9itc,1261 +torch/include/ATen/ops/expm1_ops.h,sha256=hQSO3un2pSI1DvzFBJ949uBvUa234T5_ASfKKaF6NMI,2291 +torch/include/ATen/ops/exponential.h,sha256=OLwz1tMRIyNAQ_Wg3iMZ3cOsVcEgWxosYDCW5n0NKeU,1716 +torch/include/ATen/ops/exponential_compositeexplicitautograd_dispatch.h,sha256=2JGiyMW_TduZm3inMUgXq6yWIvvkon_lRPCBVePMywo,1392 +torch/include/ATen/ops/exponential_cpu_dispatch.h,sha256=Kc2RT6ty52LwMMZwILvYmBUwEy0J8001xFQeC_WOlkc,1049 +torch/include/ATen/ops/exponential_cuda_dispatch.h,sha256=HS-hxBRa1-E6n-z8-Adpn5cEzgmi3-3ihuhRNFPCnOI,1051 +torch/include/ATen/ops/exponential_meta_dispatch.h,sha256=WjauTK8PSgtlf3mwMuAt5SPfVCw9pIcUcODauuR3k1U,1051 +torch/include/ATen/ops/exponential_native.h,sha256=w-2ggg6SMNnqKT0FmKt8huOkKYznDk1rXd5UfjM_bWQ,1076 +torch/include/ATen/ops/exponential_ops.h,sha256=aynzxfKhVSLcVEpOYLhIh5bXY3afU2tGX-zuWrmJ4-c,2933 +torch/include/ATen/ops/eye.h,sha256=yvN9tnXWSxjpZq2umTRF5DIJHXGkbGzONnNiSWDt2Ds,9881 +torch/include/ATen/ops/eye_compositeexplicitautograd_dispatch.h,sha256=YwxF-xhjKV1WLRvRQF7_NYbQTp9i-b0bffZ_QV7nxrQ,2056 +torch/include/ATen/ops/eye_cpu_dispatch.h,sha256=OrZnWPbjBiqa82qIrHrS_YCKSEucuaXYMdQV8bULpzI,1508 +torch/include/ATen/ops/eye_cuda_dispatch.h,sha256=ZUIk0lYNrHpl0xgmx2NCSDDYL9dV2TtuRWQaq33SGAg,1510 +torch/include/ATen/ops/eye_meta_dispatch.h,sha256=mJ-ddSEcIUuRyiJjMeQ71z8WQLjjPzfhOv5eFNKB-fY,1510 +torch/include/ATen/ops/eye_native.h,sha256=3qjHBew7DSKUXihFMsWlIlRtBi9ra-oDY1xMfT-DSYM,1359 +torch/include/ATen/ops/eye_ops.h,sha256=WgmenuNAH2ip3udCtF5chVOHuGXmoTMs-h1xbQO9od0,3844 +torch/include/ATen/ops/fake_quantize_per_channel_affine.h,sha256=PeVPisP3elUtykaZbx00M3xm01AycS81z9Fnr5zjKKI,1224 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask.h,sha256=D2O0dx7Cukpmuul5h7i6OrUwhnwSD2MIL7hcBJCDqKk,2547 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask_backward.h,sha256=LnI91W7oOlfDtsd6K5Mlv4QQk3BIUWZ9g8lBy4DCAVA,1115 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask_backward_compositeimplicitautograd_dispatch.h,sha256=jdjlS_IgLV7IiLS6CzJVg5yVmFXocPgDRKofSiP_FIU,1088 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask_backward_native.h,sha256=1-zBhn4n0FDmCZ3YAVogddgbjgBp0cLuN6W7trh_K9M,800 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask_backward_ops.h,sha256=LZli9OZ8vxXIERdK4E4kPL0Ah-mQgnNTXLc_-YN8z9c,1429 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask_compositeexplicitautograd_dispatch.h,sha256=UwhphtOYbfLX2CelW1IBh4EnjOnEyR-dIXR4Ou8Qp38,1505 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask_cpu_dispatch.h,sha256=3hfbe3Y8Az_t4_sASQwHgUg8Yd43iGEZJhdjA-istRA,1144 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask_cuda_dispatch.h,sha256=wwWUOeNnNg5SBE-KJ9VymmoRuchqro6UnKz0rCiYmSA,1146 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask_native.h,sha256=yZpy2PcC2hHILwWx4_BTZhQ6DHQnpWMBE3RzfmZksCg,1170 +torch/include/ATen/ops/fake_quantize_per_channel_affine_cachemask_ops.h,sha256=QdVnqbPrqO1NUyyLa4QUPwVSAmnYB8tTb94Srd-YF8s,2997 +torch/include/ATen/ops/fake_quantize_per_channel_affine_compositeimplicitautograd_dispatch.h,sha256=CO_4xzd6HXq2C6SAfH5By6OJmoxrxOb7dOTGauq56QA,1153 +torch/include/ATen/ops/fake_quantize_per_channel_affine_native.h,sha256=Yxyu38RSBHgjvPjv2fU2bMi0V3cTW38AHem7ZaLJsLQ,865 +torch/include/ATen/ops/fake_quantize_per_channel_affine_ops.h,sha256=tosCjZgzGyneSPAuvTpW-PetIxEunfPa7UwhzDqY52w,1647 +torch/include/ATen/ops/fake_quantize_per_tensor_affine.h,sha256=yglQ2559s2zf8mxDb6DnxDSj20AXtZwfwCngiF0cPzQ,1602 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask.h,sha256=_YWL2kvwwND0TPq5KoUqLsUuGANW2HLFhwaET3KwIg0,2366 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask_backward.h,sha256=R4aEAs5vmG4FmAhqknIIJu9B1ilDxjRtDZ0YjI96c5w,1111 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask_backward_compositeimplicitautograd_dispatch.h,sha256=WY_f8NPlOyaO270NjTpFTK5d5SEm-aRYcDJG7VPHd8Q,1087 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask_backward_native.h,sha256=XlFsAxkCXUh29GZybxIlk55oqXPN10wmKNCaVNYEkS0,799 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask_backward_ops.h,sha256=tAX9NvU4zbLm1jUJgjzPLPw1Onjm3f9i5Rt9cYRAE6M,1426 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask_compositeexplicitautograd_dispatch.h,sha256=aJxGKx6U62atFKH5aHvKtSydWD5YbB2YWBUdPXlsQu0,1429 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask_cpu_dispatch.h,sha256=xt97c30q8oLo2XMPT22BENcJ3vPzvnH-GKjHLtrS9eM,1106 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask_cuda_dispatch.h,sha256=JVIQzZ96Lzpi-f50iuT8vH7D7-ekOjOU5P7IX3wWTJ0,1108 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask_native.h,sha256=I3rSDY1c4tIFFo6AFzb4ZnBAsazUsmhruf5-N7Fzopg,1094 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_cachemask_ops.h,sha256=jolIWSHPIp96Otyjkxt3XITxRAQpvX0DDPG01Vr2Q60,2751 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_compositeimplicitautograd_dispatch.h,sha256=GQWq-eTUpibA5EZTKtw13KuNkYidU7dXaR2g5Ko28B8,1289 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_native.h,sha256=JBEBVI4-af_-ADdPVDmzEkfTe15hDDZeLG3KwXu3K7Y,1001 +torch/include/ATen/ops/fake_quantize_per_tensor_affine_ops.h,sha256=R35j5L7sJ1AMaVawwjwDp_M1MAS9AFphSC2SYpExYXI,2454 +torch/include/ATen/ops/fbgemm_linear_fp16_weight.h,sha256=rEJWd7iWaoeuNbb31IDvESxaZ_ze2ZPgyGizjvDAqA4,1451 +torch/include/ATen/ops/fbgemm_linear_fp16_weight_compositeimplicitautograd_dispatch.h,sha256=NnxbWuFM_n8flaF5VOtoRU_omvQUbqaWTXd3tq1f6Og,1251 +torch/include/ATen/ops/fbgemm_linear_fp16_weight_fp32_activation.h,sha256=pizXa-j2MbghR0K6Za_-eucQ2xjPRlLGccjnIeWVNVU,1599 +torch/include/ATen/ops/fbgemm_linear_fp16_weight_fp32_activation_compositeimplicitautograd_dispatch.h,sha256=2qHvTxaPDzUmmi3BC2OznzjK7Mf6wc7HpmBWqoOwxX8,1317 +torch/include/ATen/ops/fbgemm_linear_fp16_weight_fp32_activation_native.h,sha256=ALLHF_CWsZ5rABccJNWztgMUckCi0Xda8HfTkXyk2wc,1029 +torch/include/ATen/ops/fbgemm_linear_fp16_weight_fp32_activation_ops.h,sha256=wFHnXSSAfS9d8QMkn63l1zGbxosfgSKt-5UsxBDY9OY,2503 +torch/include/ATen/ops/fbgemm_linear_fp16_weight_native.h,sha256=fKIKK4Ikfjr7OQPxuqSEt9b3EjgPlWUFTiC0JJq1Agg,963 +torch/include/ATen/ops/fbgemm_linear_fp16_weight_ops.h,sha256=jKAw5AFTYNk6j6QfpMOroz1z2wewahT_EtvPev0EAAc,2303 +torch/include/ATen/ops/fbgemm_linear_int8_weight.h,sha256=2SLO1CDjGYvWl3TlPPrNsH5q7PjduZUUNEAWiOHSn_s,1330 +torch/include/ATen/ops/fbgemm_linear_int8_weight_compositeimplicitautograd_dispatch.h,sha256=WTyPo1IuzZbEiGo6txngpU-jMWVRnzIOkHUYrnbBUeA,1220 +torch/include/ATen/ops/fbgemm_linear_int8_weight_fp32_activation.h,sha256=3yrw-E6DlS_AEMsQcEIJdxRr19rZYsT4mn2vwt54BLk,1394 +torch/include/ATen/ops/fbgemm_linear_int8_weight_fp32_activation_compositeimplicitautograd_dispatch.h,sha256=eq6lQUt9oqADBCsCb_FZhepUpuhmFxft1akZiorbArU,1236 +torch/include/ATen/ops/fbgemm_linear_int8_weight_fp32_activation_native.h,sha256=pXBHxCRXh0Dv04wHSSDu2dYQ6pup07KfI2Eb9gb3UoQ,948 +torch/include/ATen/ops/fbgemm_linear_int8_weight_fp32_activation_ops.h,sha256=D4AhBZBISv9YbV95mVFJG0ow467D_0C-nI5KESWwgIw,1913 +torch/include/ATen/ops/fbgemm_linear_int8_weight_native.h,sha256=gP2h6COUhM4enFgI8TPhqcOmH9Zq2mYGxxKMZzfVDd8,932 +torch/include/ATen/ops/fbgemm_linear_int8_weight_ops.h,sha256=NSFL6EZyxesFHNY04vo7bOkgecjjEBYcfk0ZZJChNfE,1865 +torch/include/ATen/ops/fbgemm_linear_quantize_weight.h,sha256=r4NZmsvcL46eoxeVmNbTLSupO2O92kdjta_2RSQ8kUM,1048 +torch/include/ATen/ops/fbgemm_linear_quantize_weight_compositeimplicitautograd_dispatch.h,sha256=jrhuVbpaRgfX7n9Cn4F7ZiUMKMIo35R-okv4S29SPeM,1082 +torch/include/ATen/ops/fbgemm_linear_quantize_weight_native.h,sha256=Y7HhnR60kNBYDX29jtHrKMmwaMA8M8AniZwlGL669-g,794 +torch/include/ATen/ops/fbgemm_linear_quantize_weight_ops.h,sha256=ETt7KOzq-R5BtT2m41ZymB_3JIuPxdZsYOvFYsPxiTo,1425 +torch/include/ATen/ops/fbgemm_pack_gemm_matrix_fp16.h,sha256=ZH7OkWkVCTT0-5fQkHv7jI9ATQqaJ_8TvILfdtgOBYw,982 +torch/include/ATen/ops/fbgemm_pack_gemm_matrix_fp16_compositeimplicitautograd_dispatch.h,sha256=_hB1ElIij5AXi_N4Sy_hZRC5MBD7ndun_fVOPz87gP4,1041 +torch/include/ATen/ops/fbgemm_pack_gemm_matrix_fp16_native.h,sha256=8wpyhAiBln9Taj_sqfvJARfVDU2TTUNt5fb5lPvcauo,753 +torch/include/ATen/ops/fbgemm_pack_gemm_matrix_fp16_ops.h,sha256=vnke88L_XH0eZI4aCt7Gu5J1mpQIy4A7dRwZJJRkOpE,1280 +torch/include/ATen/ops/fbgemm_pack_quantized_matrix.h,sha256=AJ2lH-Vzh9Ye_DXfgV5PbrxOCXIZP2ixXvoYIhppiNA,1234 +torch/include/ATen/ops/fbgemm_pack_quantized_matrix_compositeimplicitautograd_dispatch.h,sha256=kCSvlteZQETsVtytyoGRA63NnS-8Pi9NnHxflgX-9cw,1140 +torch/include/ATen/ops/fbgemm_pack_quantized_matrix_native.h,sha256=1nEER8XyhgNr7VX-hRKWZy15mC9h92XOxqTFY4laQSo,852 +torch/include/ATen/ops/fbgemm_pack_quantized_matrix_ops.h,sha256=P9SAYE65ukiBn2y3JF_X4VqrQvsaUwa7eKghq6kNgVU,1933 +torch/include/ATen/ops/feature_alpha_dropout.h,sha256=IEOzoHvVUT665YiE5OltVEEPAfe0VvXSpr6ePW_hwnM,1247 +torch/include/ATen/ops/feature_alpha_dropout_compositeimplicitautograd_dispatch.h,sha256=7MNUIzUIGo9sxRNW51zRFWz71zKqynUZ1syimBzLjAc,1144 +torch/include/ATen/ops/feature_alpha_dropout_native.h,sha256=wxBM9GcRAQkrxu4mApXJknw__KcTf0-E0bNW2dPmvgE,856 +torch/include/ATen/ops/feature_alpha_dropout_ops.h,sha256=tZG4WA1FEFeRKh9PYfY2CSoBn1Q7_TbE-KMTTt_Q308,1961 +torch/include/ATen/ops/feature_dropout.h,sha256=LtdOsP_y4Wddw1fv60eQWaAt9knJzUrPk1OY143IM_o,1205 +torch/include/ATen/ops/feature_dropout_compositeimplicitautograd_dispatch.h,sha256=F1PqkcJwP3bCsmi7KoSCsQAf63AXEZ0dtLa8uFN2wZg,1132 +torch/include/ATen/ops/feature_dropout_native.h,sha256=lN63OrljGMVRMNqx75hZUHlnag9ksf6en-_gb2SCdn4,844 +torch/include/ATen/ops/feature_dropout_ops.h,sha256=CbEAQ33pnNUJrrNlhI_WtvcxtT83J10ixvUMD50FuBc,1925 +torch/include/ATen/ops/fft_fft.h,sha256=M11Bpu6vaLV4EATPMTVOrKmc2ow6bj7EQP3BHzQ_oWQ,5276 +torch/include/ATen/ops/fft_fft2.h,sha256=0wGgXCjXj1dgDBWaRKal1tURy_KuYNC781_uHMCbCHo,5569 +torch/include/ATen/ops/fft_fft2_compositeimplicitautograd_dispatch.h,sha256=VvqQFq-lPF8DjMk3UWxgvXPy0Wj5r8b9JZ-obTng26M,2108 +torch/include/ATen/ops/fft_fft2_native.h,sha256=DlXa0AHqJ9hFlh8ZgGTyl0aWRNi8U3KYsCa6OgHidlI,1046 +torch/include/ATen/ops/fft_fft2_ops.h,sha256=3H6YGe9ney9ZSc3IEOi0s6_cROL5OmLhs1J8XQgGNrw,2447 +torch/include/ATen/ops/fft_fft_compositeimplicitautograd_dispatch.h,sha256=GKBoWJn3LX6pI8xu03uO3t3LC7WG73XUcggG6H3X6eY,2043 +torch/include/ATen/ops/fft_fft_native.h,sha256=OddNp963aPDRrczthntRUIzSYZ8G_bMbO-KiDrBdwz0,1027 +torch/include/ATen/ops/fft_fft_ops.h,sha256=2dGAHOMeDl7jSYdweA4heUOAyZuUKEinx1dimYmbWsE,2383 +torch/include/ATen/ops/fft_fftfreq.h,sha256=NWlzS0IjAhh8GpcAFQUpilkG7BiKRw_i5b5gjLcvfFg,2018 +torch/include/ATen/ops/fft_fftfreq_compositeexplicitautograd_dispatch.h,sha256=h6v3goHYS8qamyTY5hxLZhkayR8LntViOgxNprp9umA,1416 +torch/include/ATen/ops/fft_fftfreq_native.h,sha256=J0Ii11CbYNVM_unyCRA_w6_QxSKp8aFspiYKsG5XAfs,971 +torch/include/ATen/ops/fft_fftfreq_ops.h,sha256=X7uK1Ndy6eCHNRUjLq8GaCblOYYyVz7Dc92q1o2GUgs,2319 +torch/include/ATen/ops/fft_fftn.h,sha256=UeQOT38Q76hI-5Ew_utnNhk4uTgIItOHt6IDuzK2N7Q,5709 +torch/include/ATen/ops/fft_fftn_compositeimplicitautograd_dispatch.h,sha256=JhEq2CWFST3wpAvyYLp5SZz69By_sMJqa4Q8R7ORX2s,2184 +torch/include/ATen/ops/fft_fftn_native.h,sha256=p0c7bE0iNWEtS4vBIa0BHkCcoKE-7HanIsdgJzBgO9g,1069 +torch/include/ATen/ops/fft_fftn_ops.h,sha256=S9c5O5I3TD_niIx8KO8p4j6Sodrek1oeNDvGnkiEG0Y,2491 +torch/include/ATen/ops/fft_fftshift.h,sha256=yQ_WiuirsiiizCi5okKsFKPcZeqbYe2w7VGrqmjtreg,982 +torch/include/ATen/ops/fft_fftshift_compositeimplicitautograd_dispatch.h,sha256=Y6LA_tyxazzn5yP8ykkFy9oHDXD686CO-XgF3dgTUDw,1068 +torch/include/ATen/ops/fft_fftshift_native.h,sha256=MOa0EAZlLvy7FiwenLfs4c1PMtdf5Cz2tA_NYLmu9jI,780 +torch/include/ATen/ops/fft_fftshift_ops.h,sha256=My5Mjv45q_Auc_lePE-Vw4l6e0uuASnrkpYOxAwgGlM,1330 +torch/include/ATen/ops/fft_hfft.h,sha256=gTfDc8GGlPEeHdDeb2_di2dLfE2M1d1OVLRQ6LNnEzM,5307 +torch/include/ATen/ops/fft_hfft2.h,sha256=uy1oF3Zc4ycrNRgDXqJSaQVFJWa_fG_hl9SYpeJZuXk,5600 +torch/include/ATen/ops/fft_hfft2_compositeimplicitautograd_dispatch.h,sha256=GcOG-F5AlVA03Uu-opJlJDW_Hua8RkXqfVylXBqhDiQ,2114 +torch/include/ATen/ops/fft_hfft2_native.h,sha256=qTENpXwSHQnjhuBL9d4Lv1Ydg4-w7yHJ9le-nAuApJ4,1048 +torch/include/ATen/ops/fft_hfft2_ops.h,sha256=dz_0r8pezT9mTLexp5iWR8VFjDg_So8YP5ZTvluF65k,2453 +torch/include/ATen/ops/fft_hfft_compositeimplicitautograd_dispatch.h,sha256=Yr3snNmW-G92XJftSlhmumUdmXZSk2TXlbUvVEGHz-A,2049 +torch/include/ATen/ops/fft_hfft_native.h,sha256=uQ05CzOpegxwzk7QbaGHVXyNvtwkJ9QBHFwR29zx9-U,1029 +torch/include/ATen/ops/fft_hfft_ops.h,sha256=pBP291CrOQnonQ8o0WZgd98rriqbT-G9jK0zcdKh9lg,2389 +torch/include/ATen/ops/fft_hfftn.h,sha256=yqW73zTcfcuKqAFCxZ-3d0m-cm-K6jVWqdWlmjAZYVk,5740 +torch/include/ATen/ops/fft_hfftn_compositeimplicitautograd_dispatch.h,sha256=oozUau26mclGYm0yL-J5_RKK5KQNyKDnJuK5Sd65KbE,2190 +torch/include/ATen/ops/fft_hfftn_native.h,sha256=zk72mLw0fLEOxKqfPHmLZvnPrIA51UEKMtvrRQp1woY,1071 +torch/include/ATen/ops/fft_hfftn_ops.h,sha256=rD61ulpbfZEAbxE9lRhhsSzp0xiiOTCwWtxcRsrHSas,2497 +torch/include/ATen/ops/fft_ifft.h,sha256=CbpOXQoczX4zacu2GqtHFGFCXFeLZI8RD1h9x33iOiM,5307 +torch/include/ATen/ops/fft_ifft2.h,sha256=B_UaNZEP4ixm0VaCN1yePeAVryTUZXcyuwxzaL2jgdA,5600 +torch/include/ATen/ops/fft_ifft2_compositeimplicitautograd_dispatch.h,sha256=RWkirtkMD5ZhYC0gXpFclnEKw8bC2qaIE0gpmp3X-GY,2114 +torch/include/ATen/ops/fft_ifft2_native.h,sha256=tXygv1YIADTYK9-xbin8XjibW7M62Myet6uE1wde1Ok,1048 +torch/include/ATen/ops/fft_ifft2_ops.h,sha256=Z0OWp9W9p_q5swvflnI5biPsi62cSghPZTien9XBplg,2453 +torch/include/ATen/ops/fft_ifft_compositeimplicitautograd_dispatch.h,sha256=1sqOD_6xaYPhhFPgc8uERKGbPi3WIF0h0L7NkjxXXw0,2049 +torch/include/ATen/ops/fft_ifft_native.h,sha256=Xs-2NrL6Lfy_OTCMETyGSuBsATYfY2CpQE1GWPWOGtE,1029 +torch/include/ATen/ops/fft_ifft_ops.h,sha256=jHOMhQd4W0PvTorj5jEJUSHW2s6PlyAU7BBnYGXGycE,2389 +torch/include/ATen/ops/fft_ifftn.h,sha256=MV2J_BWthgclOzYjGQP0iBQjVJNhUOoXgC1wIhtqwMo,5740 +torch/include/ATen/ops/fft_ifftn_compositeimplicitautograd_dispatch.h,sha256=ylI3iuR3GdhmE7OUP-3vMiLc2xy8qt8Na79CyTGKRlY,2190 +torch/include/ATen/ops/fft_ifftn_native.h,sha256=fh33OfkdYJ61_frIPaQImVbqek7FY8D4a3W0ARC5uRw,1071 +torch/include/ATen/ops/fft_ifftn_ops.h,sha256=_5hcv18jtNgpsKqCaviRENrdMRjs6vNtg61ygHEHAUs,2497 +torch/include/ATen/ops/fft_ifftshift.h,sha256=Ia5BXMcCS0jTGN3PQm35eOxnjwCQYlStRwAdUnSyT50,986 +torch/include/ATen/ops/fft_ifftshift_compositeimplicitautograd_dispatch.h,sha256=psnZTfM0wpG-8FXJ-Ajd8s9pDH7_uM818XGhynoK-hc,1069 +torch/include/ATen/ops/fft_ifftshift_native.h,sha256=n2aP4vhYXVr4OQ4EG70qAx2F0UAOy8G15XimYXg2a_g,781 +torch/include/ATen/ops/fft_ifftshift_ops.h,sha256=gGh28wwcZPah6MfPPCGbKYzMQvm6hfEqpZ5Qg9QwPOM,1333 +torch/include/ATen/ops/fft_ihfft.h,sha256=OSCim3A5EPy74f8ZGWMMB7cWdiSSSvmQ8i3RH8kGkBw,5338 +torch/include/ATen/ops/fft_ihfft2.h,sha256=BpmKJ-bzeumbdNrFSjmJQaAqM8fNNaCjxbWTdumoSTY,5631 +torch/include/ATen/ops/fft_ihfft2_compositeimplicitautograd_dispatch.h,sha256=NN8-Y6NQ47k98G8-1-mIr3DHpUqc8DPyFG5PW740aAE,2120 +torch/include/ATen/ops/fft_ihfft2_native.h,sha256=3JxULIKOLGbPlkT3mOmimmIIlgzeurjRzsg1MBy4Scg,1050 +torch/include/ATen/ops/fft_ihfft2_ops.h,sha256=H_7TwM9v8CddZDHkn2poXz_gQlKRUf_GEi_4n1Js89k,2459 +torch/include/ATen/ops/fft_ihfft_compositeimplicitautograd_dispatch.h,sha256=sI6YpbPpDngZwybOK03ajz72yY_HDOrdmSAipS2je5I,2055 +torch/include/ATen/ops/fft_ihfft_native.h,sha256=kStcTUMX7krIJjKp260Ix7hh67c121L2XVwAB1xBXeY,1031 +torch/include/ATen/ops/fft_ihfft_ops.h,sha256=k_nqM-o8GS7fYTE40xoDwMKuuQWtCLP8Ssr5w6jg9VY,2395 +torch/include/ATen/ops/fft_ihfftn.h,sha256=hhY_ITrEHKw6zWtg5QdueMuoEN-b4u8zbSy4BI2E43w,5771 +torch/include/ATen/ops/fft_ihfftn_compositeimplicitautograd_dispatch.h,sha256=ElvcxXpbHGHOgSaSFF1av_x_RCMxB_1n1Dz2m74Hlp0,2196 +torch/include/ATen/ops/fft_ihfftn_native.h,sha256=tK0B8zcF85H_r5c5Oo-XnBKr-l7-2wtBxdYCPLgWSx8,1073 +torch/include/ATen/ops/fft_ihfftn_ops.h,sha256=t6K6N_ejJj318Ht9oEBBIFw5Q1-4JAZhu2ayFBcWzEo,2503 +torch/include/ATen/ops/fft_irfft.h,sha256=Aug9ws-ifnVir2-p-QDcnv2HX3AZe6ahT8dghIlUAxA,5338 +torch/include/ATen/ops/fft_irfft2.h,sha256=b4DOMevddnkpxUtXtP3YeaEKcxdRB-flFZD3Ac105WE,5631 +torch/include/ATen/ops/fft_irfft2_compositeimplicitautograd_dispatch.h,sha256=0NzWTteqIn9oCPu7ZW0s-9C6dY9zTmWJcImUMwDkxDI,2120 +torch/include/ATen/ops/fft_irfft2_native.h,sha256=TnQiGSfnTlthusoJFtD_TzffgH5t4UGQIYohO7m8d58,1050 +torch/include/ATen/ops/fft_irfft2_ops.h,sha256=y4HRA8aiKneNL30jIm_0ZBExOC_chI_NNffnTZcGfSA,2459 +torch/include/ATen/ops/fft_irfft_compositeimplicitautograd_dispatch.h,sha256=LgUaE8M280ewQ238ImMNsjaXNlYkBy8BUyHKv88FWXc,2055 +torch/include/ATen/ops/fft_irfft_native.h,sha256=ZYZNyw5MFe8VwDmTDZt7y9D5_LoHEgp4gDDibDtCu-4,1031 +torch/include/ATen/ops/fft_irfft_ops.h,sha256=yyS0WTWqrQ3PEyYgm3YVyU1Z0d1hGCcgLeBZjTkzVHU,2395 +torch/include/ATen/ops/fft_irfftn.h,sha256=fWQ8jaKmicIqAjEXC30fXf7P0qz5PwhU_usqEDpiz_Q,5771 +torch/include/ATen/ops/fft_irfftn_compositeimplicitautograd_dispatch.h,sha256=JUbrqKQPxFnIh1S5X5dOL-FOGYdi8qJa8d_83XBBbwY,2196 +torch/include/ATen/ops/fft_irfftn_native.h,sha256=3rEB4LstsS-A4wNbZ6rZm9IlmPoJsLACu4W9nvGWM6E,1073 +torch/include/ATen/ops/fft_irfftn_ops.h,sha256=tqcyE97ohdbQfBv1yxn0DOjaslFmrFpYTEkcJx_zfi0,2503 +torch/include/ATen/ops/fft_rfft.h,sha256=wXupBC0vkiy3byz_UxgGpcfJ49p4wzX0ZutwYwFyTsc,5307 +torch/include/ATen/ops/fft_rfft2.h,sha256=5aL8vROWEYD5FYWKaLT5uyt2DePnMGpCkTPKt5ejnvU,5600 +torch/include/ATen/ops/fft_rfft2_compositeimplicitautograd_dispatch.h,sha256=1fSxwMhyPrzWghUSmCFaREtLYjs0MlSkbmOvplQhD3E,2114 +torch/include/ATen/ops/fft_rfft2_native.h,sha256=KnRTOTNzW8GqHghack_T-f2LA5F6mAfWbRyEVuMaogU,1048 +torch/include/ATen/ops/fft_rfft2_ops.h,sha256=81kX9nrqx0d_agv4AvbuxB6i3ubWUbqk2OWqs9FCSi8,2453 +torch/include/ATen/ops/fft_rfft_compositeimplicitautograd_dispatch.h,sha256=boLg9FJ5a9I2ZISLHmN1eAWIxcZQS9Jeb-5Gr-h3BrY,2049 +torch/include/ATen/ops/fft_rfft_native.h,sha256=yhX8ogoSu1cfodcYXMsNGJEKL6mAG6t4Xy8jOhbpPDc,1029 +torch/include/ATen/ops/fft_rfft_ops.h,sha256=gRsO4HlemDDbRRJXMMsrg0-ei9OLs9v2qpCXOeU8BFs,2389 +torch/include/ATen/ops/fft_rfftfreq.h,sha256=V0wZvl3hHrSJ38A_0OD-R1kHiCZfeN8FlISXNHaky20,2031 +torch/include/ATen/ops/fft_rfftfreq_compositeexplicitautograd_dispatch.h,sha256=86fPq1shdnB9jGLvDkZnTOn3wz-O-XQHCPxMDkrph0M,1420 +torch/include/ATen/ops/fft_rfftfreq_native.h,sha256=9trMZVfdreHMz5oh6EwAQNlF14RAlWu0XmTC7ulbTEA,973 +torch/include/ATen/ops/fft_rfftfreq_ops.h,sha256=mH0kfDrx70NFYBylozAISjXicttZRRqU_CFqsKU8eME,2325 +torch/include/ATen/ops/fft_rfftn.h,sha256=Los8e2Z1JWQXuLJ4YHiXEvETTen6XU8oDx-LMA20py4,5740 +torch/include/ATen/ops/fft_rfftn_compositeimplicitautograd_dispatch.h,sha256=q6d8x-DPk2utIvXzErWwAWHZfEWQdOOIX9yz7fr_dA0,2190 +torch/include/ATen/ops/fft_rfftn_native.h,sha256=g-dG_BIECIrM2j4otNzH64paS7-S3TAcf3fmyaXvfdI,1071 +torch/include/ATen/ops/fft_rfftn_ops.h,sha256=NyIpFN6IuvWLEb8x5cGSWKdI1rIATr8oNzdsVOigJSI,2497 +torch/include/ATen/ops/fill.h,sha256=CmYsDrpGMo3LC8A3qWCDjquXPN3puo1sqmj_exlmWxE,2528 +torch/include/ATen/ops/fill_compositeexplicitautograd_dispatch.h,sha256=SMbE8E7Cur9L5SPnzHDkMXAQ8hmGrFN2cfEJbDW-rUA,1530 +torch/include/ATen/ops/fill_cpu_dispatch.h,sha256=1wxfvPbjfmKG1p-p6TT8P3xDS83j7D0EfCZ2GXU1Bxk,1070 +torch/include/ATen/ops/fill_cuda_dispatch.h,sha256=mdrA6gVYyEgelG8S-td_a3IhKlptb0Co8ofIe6eBpKo,1072 +torch/include/ATen/ops/fill_diagonal.h,sha256=6MOHN0GTTTQmcnt3tLf275O65XOV1MIH94mZIOVn48k,764 +torch/include/ATen/ops/fill_diagonal_compositeimplicitautograd_dispatch.h,sha256=eOL7prkbALjYa1TH2EV00C0JIvg-WIxO-KIOrtmUXq4,1070 +torch/include/ATen/ops/fill_diagonal_native.h,sha256=fo5-lcVXdX2LcokflqCnXhXkrrcgiuOQ9paqL6eRrVM,782 +torch/include/ATen/ops/fill_diagonal_ops.h,sha256=g-mqRusbFaGWt0qGHHkHT2mUqO5Ndu1z5OAShSJ7piI,1377 +torch/include/ATen/ops/fill_meta_dispatch.h,sha256=15gVrQvkc8Cvpm3MSV3oSuKzDh-SYebjWZ-P-ym7F-w,1072 +torch/include/ATen/ops/fill_native.h,sha256=FINbaiTTgMSkz9m_PievKrWX2z-vdOFMr9uOhaMVnmE,1780 +torch/include/ATen/ops/fill_ops.h,sha256=08a4ZKO6VWXusZm89fzd6AhFln0xEhVwLVwGPTIIhcI,4491 +torch/include/ATen/ops/fix.h,sha256=LjhIoTc-Pc11lqgSHnFhdOG8MeMXOC0xhot8PsahYPs,1384 +torch/include/ATen/ops/fix_compositeimplicitautograd_dispatch.h,sha256=Ww9OtVDxfWUgpOcdmHXD3OPLs7RksbZKq9UTE2mn3xk,1214 +torch/include/ATen/ops/fix_native.h,sha256=DUkMUeoWv4Y_zvX6P07csWwKK9NgpwhXGbc1PmuQvaU,850 +torch/include/ATen/ops/fix_ops.h,sha256=L6-wobHT_ndsDPXDNa4S15uycQ5gQKtruC2NHLsi6ms,2273 +torch/include/ATen/ops/flatten.h,sha256=g50AuLO58HKIYSPalb6lTbZvG0lZhX6JhFAuC_3MQck,1894 +torch/include/ATen/ops/flatten_compositeimplicitautograd_dispatch.h,sha256=qjL8cQD86Ti-HIfj7asT6dLP3seAEjbCGdxK4WDQ7SM,1390 +torch/include/ATen/ops/flatten_dense_tensors.h,sha256=hQzF31WjAfioGMzTst8m5Yh5oFpeKTO3_gXx3HZYETY,958 +torch/include/ATen/ops/flatten_dense_tensors_compositeimplicitautograd_dispatch.h,sha256=mwmLr9TQmOa2YUcyrSvtCOpDhXoiz3WmtmLcV0EsoBk,1032 +torch/include/ATen/ops/flatten_dense_tensors_native.h,sha256=baRJPS0D3sDDeiAtG-FLl46Y5e9Sp4-jrnFLeyaEu_s,744 +torch/include/ATen/ops/flatten_dense_tensors_ops.h,sha256=-wk6ErOd0adX6xxzSyjdFaGqdxmpEOUHZQacUUrCCS0,1255 +torch/include/ATen/ops/flatten_native.h,sha256=Hco0dxlohS1AwgCPY_NmzSj0m3nPbVXhEE_1Rn2UvYk,1102 +torch/include/ATen/ops/flatten_ops.h,sha256=hkiJwHh6Fp0TIh--jrUEv4Zg7a-PNpxr5N7fTJSctOk,3574 +torch/include/ATen/ops/flip.h,sha256=2rRZdFRuX7vbvvMxjsg8-EVYifzDso60Rcp7t6nxqNE,1381 +torch/include/ATen/ops/flip_compositeexplicitautograd_dispatch.h,sha256=otnkQYBXa0NOI33GtBMW7JtD41-RDX8JaIMwjxQYrDk,1161 +torch/include/ATen/ops/flip_cpu_dispatch.h,sha256=6BVlvIbh_YM5pGjTHgewv5hG-QDMzxsHaX69N4NIhPo,994 +torch/include/ATen/ops/flip_cuda_dispatch.h,sha256=j3ZoleuxEb3q8S7ZUCWApDgOLmhLkkCnvRj_AgcC7AU,996 +torch/include/ATen/ops/flip_native.h,sha256=5Cfw7aoDe0kaqApnodI4w0eIU6INYofELgwwy_w9fTM,848 +torch/include/ATen/ops/flip_ops.h,sha256=Ytnfwlev3m8eCrDoTt02QeD4ICwb6R7KgOBeM6jjHaE,1935 +torch/include/ATen/ops/fliplr.h,sha256=kVQMLQ7M1sBGTDI-XB-S9z2rnYiVMt9qcuHZ3t01Ryo,891 +torch/include/ATen/ops/fliplr_compositeimplicitautograd_dispatch.h,sha256=1ZXDrZTe_PEYlSKfLQfsS1Ya-418EjrMoSX0VbRaho4,1018 +torch/include/ATen/ops/fliplr_native.h,sha256=rMJyOgWZgN8B-nblfri1vyYCpy7dvidNPe6YwyHdjiE,730 +torch/include/ATen/ops/fliplr_ops.h,sha256=-AE_Av2caiEDa6pMpsglj17cFhGIlrJdxnvR1W4sAN0,1211 +torch/include/ATen/ops/flipud.h,sha256=EU_INCEcxy8U-uv8TNN9nf5KNsJ3Y43eDo0Gsdh0MrQ,891 +torch/include/ATen/ops/flipud_compositeimplicitautograd_dispatch.h,sha256=3j-jfDVZvyoOqh-EyQsZ5V6kXzcNbhsG9_X-at2Gpm4,1018 +torch/include/ATen/ops/flipud_native.h,sha256=4hbi2pH6fdLMEbSSfVqWETvRrGVuErvXrMIAG2aN7yE,730 +torch/include/ATen/ops/flipud_ops.h,sha256=H5M9oEFOLn9b9vwSGYplNthValqHCVq9NCmn42dxcWw,1211 +torch/include/ATen/ops/float_power.h,sha256=p_WfNyBweg7YRlGF20luBXFYm4VJDhJy0Y6CNnoT8aE,3185 +torch/include/ATen/ops/float_power_compositeimplicitautograd_dispatch.h,sha256=3g36PswNAx6RgKDLaVZ5BQTK5E21v5F0eHigQyLxL6A,2073 +torch/include/ATen/ops/float_power_native.h,sha256=CNKKLNUxh57vqzCDX4zBtGrJN86iqZ82XpkcUTlmqpE,1446 +torch/include/ATen/ops/float_power_ops.h,sha256=ylPgKexYH-oXFhNQGj3_5AUSPRtou8oRGjHb_YkJlOc,6106 +torch/include/ATen/ops/floor.h,sha256=jeMqZhwHjHFAxDoRfBWEaPC0TcKTcFT8jeMrtUoq_ug,1410 +torch/include/ATen/ops/floor_compositeexplicitautogradnonfunctional_dispatch.h,sha256=KxcARUIkRIxefK_mieTzzFPE765ud5D_zQicBWxAdzk,1093 +torch/include/ATen/ops/floor_cpu_dispatch.h,sha256=vZD7gChTT-a8EH1fIsbOcuuvqxbyHvbJSi1rVQUwyGI,1178 +torch/include/ATen/ops/floor_cuda_dispatch.h,sha256=RpDy36nL6nYMGswyyKUKTgAEAvsdkzmTVPmGIPTUuRc,1180 +torch/include/ATen/ops/floor_divide.h,sha256=-G13tDwJbvESjQQ3bGAf9P-c1pKS-USjrGwEBaEaZHk,2244 +torch/include/ATen/ops/floor_divide_compositeexplicitautograd_dispatch.h,sha256=JVHa8paW7S8RGEhigBT39u77dfc-AK4DTfXxrejnkoE,1354 +torch/include/ATen/ops/floor_divide_cpu_dispatch.h,sha256=YGvvcwlAyxmCf3UolSOfV_9IyW2srv8aWFztXZO5OnQ,1310 +torch/include/ATen/ops/floor_divide_cuda_dispatch.h,sha256=5W1PuMzmsZOp6mlaHR7bMgiLaqiX4TbN1dJFyR5yK54,1312 +torch/include/ATen/ops/floor_divide_meta_dispatch.h,sha256=K7dTzgvaaCm9ijjhuaYZuj2PF4bXM2KHtWcCcSDGuTk,1005 +torch/include/ATen/ops/floor_divide_native.h,sha256=7tfQOpKTR5IoBSFqwHUsXCB3gc4yjtgIMtjbLb1JzmI,1549 +torch/include/ATen/ops/floor_divide_ops.h,sha256=bhIBBESa_XKcsbbnVj1LHoIVZqJ0QTS-KAVxAmLJuyU,4594 +torch/include/ATen/ops/floor_meta.h,sha256=i86x47d9H7yJsclw1Ti-HE2jVhJMPsfdfH08Gixtzpw,820 +torch/include/ATen/ops/floor_meta_dispatch.h,sha256=puYl619XWp4Ps0IM6Sk_8QKbR6CzA2CnmZrg116SHzc,1180 +torch/include/ATen/ops/floor_native.h,sha256=L19rirdAt1vQNq3RsGaRWw5rMWxBdRh6_K8_uxeVBNk,1261 +torch/include/ATen/ops/floor_ops.h,sha256=-9ZqNwqPXO37mPBu5n35ffp7t6unQJG4hThKjOJ0FMU,2291 +torch/include/ATen/ops/fmax.h,sha256=C8xrycKj0rAkQ87es-fFKDiGCrtik3J1nAJRNysfkzs,1402 +torch/include/ATen/ops/fmax_compositeexplicitautogradnonfunctional_dispatch.h,sha256=CJ4yQK6B8tc8VJHsS0cr4T9K2kCAXnnnXS3yMevzGgs,1068 +torch/include/ATen/ops/fmax_cpu_dispatch.h,sha256=Mj5TBFWU6BdTslprzpiCPrreCmtbT1SFJy2UqIRlNko,1203 +torch/include/ATen/ops/fmax_cuda_dispatch.h,sha256=1_swZzGFZu_YHxzk5MYHG2Nl6Cd5Q_6e5o93bACEklU,1205 +torch/include/ATen/ops/fmax_meta.h,sha256=VQzFhCpOHumFB7J7h6a3yQDQUcvabsjF5jjkoNmgSFo,845 +torch/include/ATen/ops/fmax_meta_dispatch.h,sha256=88I5nqvOuHgZM4s7wBNK30jluZ7tIf-YjiMvPiY8W14,1205 +torch/include/ATen/ops/fmax_native.h,sha256=7Zuy1wo9ieWjUOcqKIBx4poJn3iztBodsSpmIzYYUaQ,870 +torch/include/ATen/ops/fmax_ops.h,sha256=grtHjIIueHpaM77vrwhy-i6VXOFkmzRePcNkgCjfzto,1961 +torch/include/ATen/ops/fmin.h,sha256=LtR-fLULsZd9lKA4ryRFCOToxLuzv2ZjpaxdE1RODbs,1402 +torch/include/ATen/ops/fmin_compositeexplicitautogradnonfunctional_dispatch.h,sha256=-PFMakY72hswvDN7x6rpEuleI_zLy2SvFFuaXLAsDkE,1068 +torch/include/ATen/ops/fmin_cpu_dispatch.h,sha256=RSsjBPS7j0xoAlc3lxKir-zy1tCkBG1IR1I-0P9md3Q,1203 +torch/include/ATen/ops/fmin_cuda_dispatch.h,sha256=PBUqaAsBs6gdczqr4Tv0IXy0NcAuVJVTTmHE8OgOGfQ,1205 +torch/include/ATen/ops/fmin_meta.h,sha256=asGDdC4BJCNclr5liYoPogohKHxW9eqf3yJVZixFzbU,845 +torch/include/ATen/ops/fmin_meta_dispatch.h,sha256=DOztqHYtB_wPICD7KnZbuReubjU3sy5ieecsZ6PpWhs,1205 +torch/include/ATen/ops/fmin_native.h,sha256=WapxsiIYYKqXK1m8pY3-a62LIYsvMYkC7aSGK_kHhcw,870 +torch/include/ATen/ops/fmin_ops.h,sha256=b4LZPNU1-K5_2U5C9ZJr3ObpfFygWBxz_1ONVngawB8,1961 +torch/include/ATen/ops/fmod.h,sha256=OJxMPap8iFBFybff9t5kryaO-lTRXttTAoYfPRttyaE,2134 +torch/include/ATen/ops/fmod_compositeexplicitautograd_dispatch.h,sha256=_bAyJn5u-v39eSmgWTlEsZx50J2ll9IrtYbKdMr7Ve4,1322 +torch/include/ATen/ops/fmod_compositeexplicitautogradnonfunctional_dispatch.h,sha256=RqmMhxykawrRTo-nrCWkUMrj5TVFuxJzQzg6k_TTGMc,1143 +torch/include/ATen/ops/fmod_cpu_dispatch.h,sha256=Lza0K2I5e1p84KQOP53lrcyjUwx1r_feAIipwK05Cv4,1278 +torch/include/ATen/ops/fmod_cuda_dispatch.h,sha256=s3cTmJQC_4n2DR9xlGaVy4-s4vzLKimIx8D-ImjJWPM,1280 +torch/include/ATen/ops/fmod_meta.h,sha256=C1Jpy-wydYgr99CDE0fXxP9iNmRheNY3CXvsAWw10iE,852 +torch/include/ATen/ops/fmod_meta_dispatch.h,sha256=6nX20SaIF9N8vy_VX9ZQbniey4W87tJ0k1-WLw-ffNE,1280 +torch/include/ATen/ops/fmod_native.h,sha256=4fIEEWKrMzfDtfi4df2IIYjbeHQ7R0wmAOjMEhYoKU4,1132 +torch/include/ATen/ops/fmod_ops.h,sha256=iX2alweqSvcLvAqq62fWjrah-wlqfaTBbV-ZfFKkD-8,4491 +torch/include/ATen/ops/frac.h,sha256=qY580vdXLvpcuzvQbieYB-4i9Lbd6jMVSY6o4uqSD5A,1397 +torch/include/ATen/ops/frac_compositeexplicitautogradnonfunctional_dispatch.h,sha256=qsVHJtG7pueSSM9RgwvveigUPOWU1xM-inMzPAMD68c,1091 +torch/include/ATen/ops/frac_cpu_dispatch.h,sha256=sU0VMCJ_gaPO-FlYYxIj4jux4OXk5nX9XIK2O3mEyRI,1174 +torch/include/ATen/ops/frac_cuda_dispatch.h,sha256=t5re6MKQvGtvB08VMd8C28ByXL9gXDV6rKQTOw0tuzA,1176 +torch/include/ATen/ops/frac_meta.h,sha256=RYWiYd5NTHdmBY9uoE5fmgTDiepOSVDP9id6dG69fuU,819 +torch/include/ATen/ops/frac_meta_dispatch.h,sha256=wK11eaXrEGI14x5Yqq8qso6ldJmWE2gTIzJFQCX3kUk,1176 +torch/include/ATen/ops/frac_native.h,sha256=aIR5V8yg_K1pZ6F1y3pJ5uE0qnmgRxEqyoXLbq_Q-F0,1252 +torch/include/ATen/ops/frac_ops.h,sha256=Qt3H4DnwZKE-AxfZJJgapP5JpU5nFLof_u8rF5MR8QU,2282 +torch/include/ATen/ops/fractional_max_pool2d.h,sha256=cxNTlUUqx8QyneQ_C2QI--bHxr5smaJOHak4VAzFHTk,2274 +torch/include/ATen/ops/fractional_max_pool2d_backward.h,sha256=SqwUjv4KyMMlEHjZMfvdubt9HBxaSVsmnEs8eUuv1Ck,2317 +torch/include/ATen/ops/fractional_max_pool2d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=rpaTwPZScDBQR_r9yQr0DPY27AoxlBOovbTpve4FZBU,1186 +torch/include/ATen/ops/fractional_max_pool2d_backward_cpu_dispatch.h,sha256=SYpAqGh2YyEpR-HkhNlkkWSUICwFI3Eu8Ps6TJWEdu4,1571 +torch/include/ATen/ops/fractional_max_pool2d_backward_cuda_dispatch.h,sha256=aB4QDIhPLBRSLrxnzDJUNx_4yY6PuiFHF2OdG2ZWkjM,1573 +torch/include/ATen/ops/fractional_max_pool2d_backward_meta.h,sha256=GThhHLOO1b8mGaEGyOtYnGaLDaNkPvhZnIWbTYEtE64,963 +torch/include/ATen/ops/fractional_max_pool2d_backward_meta_dispatch.h,sha256=tF3BQc3Pxozu662IB6gROAUHMPW0-hffLMKymd1ptTc,1573 +torch/include/ATen/ops/fractional_max_pool2d_backward_native.h,sha256=8Dc6nr3djwjesHHa_V8u-U1oQqxV2BijQZfzUTJ_hrs,1362 +torch/include/ATen/ops/fractional_max_pool2d_backward_ops.h,sha256=HURVUw_JKGc8QFFa8EF7MOoi_GqXzu8KK3bcJkFvkAM,2759 +torch/include/ATen/ops/fractional_max_pool2d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Vfyaj2bpg5DLXulVAp7pfr7yiRt-ulmf8wNy6nUwXNc,1177 +torch/include/ATen/ops/fractional_max_pool2d_cpu_dispatch.h,sha256=0V-fmkeZH-RDpFZGdyaLkVwHHRsJCOKuptZRCrf89WI,1584 +torch/include/ATen/ops/fractional_max_pool2d_cuda_dispatch.h,sha256=CKZmmpH90_hn8hZ_UiB2mj2UtU1FB9oyaMsE-mbwfEg,1586 +torch/include/ATen/ops/fractional_max_pool2d_meta.h,sha256=Bgr5f2-LrgPvrFR5CMZaH2RFrP_TaHB9vrmfEU34ELA,929 +torch/include/ATen/ops/fractional_max_pool2d_meta_dispatch.h,sha256=PagdMeRMM2tkhaJMGriTXL3MVTMhuD_roMJrWMCpkX8,1586 +torch/include/ATen/ops/fractional_max_pool2d_native.h,sha256=Mt9PbQxk2PwIiEGBDk5-suNTccnKuVMx2Yw6QitWkJw,1323 +torch/include/ATen/ops/fractional_max_pool2d_ops.h,sha256=fw9DKaasKi6l2PGaXmrLoiHfpzIfbzDJQcoMf4lsEko,2773 +torch/include/ATen/ops/fractional_max_pool3d.h,sha256=p0ZIF4OZoYV34lnO1_UPci5s7HGp7h7VRJDQOmGxTnk,2274 +torch/include/ATen/ops/fractional_max_pool3d_backward.h,sha256=_CFhrWIpHXBduLpEVb1XWlhrA0mnzo6-gAfJggi21og,2317 +torch/include/ATen/ops/fractional_max_pool3d_backward_cpu_dispatch.h,sha256=fH1a9hlmo0qyXDwQAYWxK-Ozwd_SpTO1kOse4UiNz7s,1571 +torch/include/ATen/ops/fractional_max_pool3d_backward_cuda_dispatch.h,sha256=XFl-vsQnMLVpeDhRO--AzJhcwwHrE3sDrIdulurmEGo,1573 +torch/include/ATen/ops/fractional_max_pool3d_backward_native.h,sha256=URU5O3duKKLPz3liV-EsY5IWgypU4eyewVpa8ArL4oE,1540 +torch/include/ATen/ops/fractional_max_pool3d_backward_ops.h,sha256=HahbbIEB9on2PBL3aS_JhMd3NZPCSmJ-j1FtzeKYoFo,2759 +torch/include/ATen/ops/fractional_max_pool3d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=mvXd0szGn6Np92Y4GeCqYtnyi6araZHJF0gTYfN9AJM,1177 +torch/include/ATen/ops/fractional_max_pool3d_cpu_dispatch.h,sha256=C6Hx4MhTBY58O_L_eMaMepRMH4ZNbOEdTT52LldCo5o,1584 +torch/include/ATen/ops/fractional_max_pool3d_cuda_dispatch.h,sha256=NwWnlhnKMC3x7yQEJQgObe6u-JNZBuXhqCXKn7w6UG8,1586 +torch/include/ATen/ops/fractional_max_pool3d_meta.h,sha256=uHzdLcs_itpvkm-8Axv1V9Ni8d9RKaMY9lHoWmpd1HI,9861 +torch/include/ATen/ops/fractional_max_pool3d_meta_dispatch.h,sha256=clnwBmvmZ6tCt9iFceJ0lv8xFyoBImN60bC2zJGW_yk,1586 +torch/include/ATen/ops/fractional_max_pool3d_native.h,sha256=jJZjNstcckPWOIX-0Su_Jr7Ue385hk5j0pShh-G-_AY,1593 +torch/include/ATen/ops/fractional_max_pool3d_ops.h,sha256=kUKtr4xY5oeknFc50d8NIgE10iWejSuGP2k9azrIyKQ,2773 +torch/include/ATen/ops/frexp.h,sha256=lsto9Fz9Gce5RkHWxWlVFrVgFF6tck8DGv1rljq0SE4,1622 +torch/include/ATen/ops/frexp_compositeexplicitautograd_dispatch.h,sha256=gkewRnZ-TlaXiFdePQuPkUNctjxs68glpFBG9Fiz-4g,1042 +torch/include/ATen/ops/frexp_cpu_dispatch.h,sha256=uZsXm8jR29ztFjYcZf2IZuy_Umx91ik2h937U5X_g8M,1185 +torch/include/ATen/ops/frexp_cuda_dispatch.h,sha256=DAB94PU4tagA3yOQ7uzuhwMGhev5UoPj3YWkHs8tc1Y,1187 +torch/include/ATen/ops/frexp_native.h,sha256=R2xmWQm9i6MsNwOyLyzrw_eZBsuScIgQfNhOS5jRg6M,886 +torch/include/ATen/ops/frexp_ops.h,sha256=4-DwHjVWJpHZdDQMIOXEBSg4Xrk1Bukn5_xwc9n_dB4,2148 +torch/include/ATen/ops/frobenius_norm.h,sha256=Ljau3LSqZpAOClM4xU0XN58hokwbFGxJH47WLhZ3QwM,1624 +torch/include/ATen/ops/frobenius_norm_compositeimplicitautograd_dispatch.h,sha256=Qrou2XNU8ZQyJDsEwnmOpHJfjvj5ic8JJMKl4Nk5LHk,1316 +torch/include/ATen/ops/frobenius_norm_native.h,sha256=SXauc9EZYRCpU4N7KpgWUuvAosUxSBngI9OmPGj7NbE,900 +torch/include/ATen/ops/frobenius_norm_ops.h,sha256=jiB3kfUvIrnswqUF_PR1EGz0klvn6OQ-g4wgljRBdvw,2110 +torch/include/ATen/ops/from_blob.h,sha256=2aU1uljA--xA1E2lUlY1OwTen3mWY878Q6swjXnx8qE,4408 +torch/include/ATen/ops/from_file.h,sha256=bKqQNCw9ThLSVYi6NPa17nvK8_o1Y1km4P92XQc6TIE,2462 +torch/include/ATen/ops/from_file_compositeexplicitautograd_dispatch.h,sha256=g-J2jY9VRLBVbHsL06rKqkB4jJ8MnnWxW_f3vyuvPJE,1270 +torch/include/ATen/ops/from_file_cpu_dispatch.h,sha256=W33xVoEQr2c8e6G7-4IZ01eywWEads4iAWQZ552oR80,1352 +torch/include/ATen/ops/from_file_native.h,sha256=ytF2a3bDdutQvN4glY23DBoCzTomT-EYrSmNkTvWCMY,1114 +torch/include/ATen/ops/from_file_ops.h,sha256=cSxTyW4huzL-pBY0un9o3B4egw2Vp1qf64BgGoyFf-A,2727 +torch/include/ATen/ops/full.h,sha256=W_KL7PH07h_yCHFMPDIlz_p982ZDDhvDIWD2dikO6Uc,8171 +torch/include/ATen/ops/full_compositeexplicitautograd_dispatch.h,sha256=oZEQbovNU2DNRy7FymH-pRt_43maViEhQlA7RiKLQwA,2800 +torch/include/ATen/ops/full_like.h,sha256=_kjVJ5ccsE9Gj0zHyUDVMpV0CY78_Il9MjfIslic6eM,2710 +torch/include/ATen/ops/full_like_compositeexplicitautograd_dispatch.h,sha256=GMOt6ayOgfcK4CtV0FHYawSaUhU1n5JeUmxo1LoUTN0,1766 +torch/include/ATen/ops/full_like_native.h,sha256=l-CUK1VokH7LWXbacXCizZy4FG2wiscYILZvMEt_ccQ,1146 +torch/include/ATen/ops/full_like_ops.h,sha256=PS9-gP2lqGFuicySJKcThE3ZMa3ENX1IlMGDqMRpGps,2851 +torch/include/ATen/ops/full_native.h,sha256=tvuSZfFsBVPYxv-PJIZL5mpdTZZVEwb2C8kX6C1_XDA,1444 +torch/include/ATen/ops/full_ops.h,sha256=cik1svV1fhpgqELsRW6vLavooCo2h0r6e1ZquPTQBhU,4552 +torch/include/ATen/ops/fused_moving_avg_obs_fake_quant.h,sha256=ZqRhBq4mXANLBZtz2KuAIV2_QtzB-Xlt-YwS3AHBMow,1706 +torch/include/ATen/ops/fused_moving_avg_obs_fake_quant_compositeimplicitautograd_dispatch.h,sha256=qpyMZX0q8jPA7hTwZI-y3XIojUoAbU3UzAr-1_XWbjY,1344 +torch/include/ATen/ops/fused_moving_avg_obs_fake_quant_native.h,sha256=fvXyZM3zFbY58f15vtpVh3-EuJgpRDKd2_CytQzgWjU,1056 +torch/include/ATen/ops/fused_moving_avg_obs_fake_quant_ops.h,sha256=_u3kd5Q0WSt7_BVRH1ovEjmYBITPJYNydRFtydgAoHU,2263 +torch/include/ATen/ops/gather.h,sha256=HoyFKag3EMrGK-ODSeaDl0hccTeeT8JhK10hT7IQ0H4,2682 +torch/include/ATen/ops/gather_backward.h,sha256=vhNYOZXNLn0I-JJZ0toD1CYW7TCAZVF9fO32ZGP8gNM,1094 +torch/include/ATen/ops/gather_backward_compositeimplicitautograd_dispatch.h,sha256=RIjP3JnSvef_v51S3I7RbmIw1m6XbkzMJ-l8dG0UANI,1109 +torch/include/ATen/ops/gather_backward_native.h,sha256=o30n_u-LiFTT8u9Vhr3RciyQxXme77elKGBUGvXebjY,821 +torch/include/ATen/ops/gather_backward_ops.h,sha256=aY4hINjO0WQw3XEXoQte0jhnziR0AbXyRpApZGIPomc,1511 +torch/include/ATen/ops/gather_compositeexplicitautogradnonfunctional_dispatch.h,sha256=eWri3P8Ujanm1cdoCnjxia8VZ0rtjT_Qv7Mg6RMIaCY,1107 +torch/include/ATen/ops/gather_compositeimplicitautograd_dispatch.h,sha256=mNPXnlEtLQkca1ezUK3rN35_n5SN1084YPPTqSKlt2A,1370 +torch/include/ATen/ops/gather_cpu_dispatch.h,sha256=pdk3G3NyqS0ljfRv4T7LGu2ekeEkZUJNAWmJBN5AirA,1314 +torch/include/ATen/ops/gather_cuda_dispatch.h,sha256=rbd2ri9hG86gHYxywmMe4pW1F6VDY1R5rKP3Ntzf0_w,1316 +torch/include/ATen/ops/gather_meta.h,sha256=n_-1fl1QfWcfJG7Ha7I1sR83QUesxW4jtNB2PFhf44o,878 +torch/include/ATen/ops/gather_meta_dispatch.h,sha256=T9x3pcLxQVU31e221aLuWswMc17AJ6M-wYk9iMe1rqI,1316 +torch/include/ATen/ops/gather_native.h,sha256=8zkVxZejNyL99eRrcm_EAQIN-xLbf8GNd1xxxzjdzMU,1167 +torch/include/ATen/ops/gather_ops.h,sha256=QrS-Yz7U0XSLa_l_rsy4TAmcKYhAmqDEl3naORDf_IU,3760 +torch/include/ATen/ops/gcd.h,sha256=dsJg-9nh9KIfD26ioieG0GlQA4fFDZJMH8vxmJvjADs,1572 +torch/include/ATen/ops/gcd_compositeexplicitautogradnonfunctional_dispatch.h,sha256=oum-ltk-OvL0q5wDAQssdewVcvt3Jf46rj4cEHw_B5Y,1141 +torch/include/ATen/ops/gcd_cpu_dispatch.h,sha256=YK7nrMwISnF5m07aNmxgYQ0XDxdbVXAiR73X67EX6pg,1274 +torch/include/ATen/ops/gcd_cuda_dispatch.h,sha256=YHFX1WM-B93lVqieNC8ptr8fjGzPfUOPtFNWWqfFH_g,1276 +torch/include/ATen/ops/gcd_meta.h,sha256=GAs9aVubl_vMTZV_0LVjsYYf2lBRvsc3HCy79uhiBgc,844 +torch/include/ATen/ops/gcd_meta_dispatch.h,sha256=pU93t0zI3ktd2heasFWFwkLiwsZ3HHUNgg-R16AM6no,1276 +torch/include/ATen/ops/gcd_native.h,sha256=YQ6UQjhVSHJHfnnd_BfjB6ogsx9V9R2d6MYbHOlF1MA,867 +torch/include/ATen/ops/gcd_ops.h,sha256=s08wXStZ8sm8pl6oMVzHflZ5I5zT7iYeIce5JRQEuDs,2531 +torch/include/ATen/ops/ge.h,sha256=4htS61aHFxyeVyTApFAgWJOFO11-uEqUSiMwwIFzjIE,2096 +torch/include/ATen/ops/ge_compositeexplicitautogradnonfunctional_dispatch.h,sha256=vNmAOi4mf-a5MVqKnxGC-Tc00ta6QWlhpGjYq6sB_gE,1288 +torch/include/ATen/ops/ge_cpu_dispatch.h,sha256=xbip_u5vJr9F9tFnMT748yBohuPtG3S8ThskqDzoxRQ,1620 +torch/include/ATen/ops/ge_cuda_dispatch.h,sha256=MFLhSA7zXyN8veoR9FPTUI3q0OiwSF8f0LQZFzJVZnI,1622 +torch/include/ATen/ops/ge_meta.h,sha256=fAfB6lqROadWzN2r2lTzK1fZyujjjJXeLaoIamqWGgg,989 +torch/include/ATen/ops/ge_meta_dispatch.h,sha256=Cr8asEW5IqeKgiRRZF-JiqvVd47Er2YYHh2M5jF_Lg8,1622 +torch/include/ATen/ops/ge_native.h,sha256=b7LAzyFNtMrv8sQqnp_KZwzXGUS1IcXmOXm8fSxhSjo,1549 +torch/include/ATen/ops/ge_ops.h,sha256=202W2Q8w1E64PoHu_jg8u-S9QoCnS00ZedbtcPkkGcU,4455 +torch/include/ATen/ops/gelu.h,sha256=vl3Y21MXm1J-df-sllHwcHtPA83d9gtj1muB97J8Q7g,1692 +torch/include/ATen/ops/gelu_backward.h,sha256=YBxYppBY7N9ZXurLcZn_Ji0dyfJeercJTCyGmvGmCuM,1834 +torch/include/ATen/ops/gelu_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=8Idi-LFD4qMlVtuQZiVB-IrGwK2CcRFKCfHB9Othp78,1120 +torch/include/ATen/ops/gelu_backward_cpu_dispatch.h,sha256=H9ARpmU4_n1AVt5akEWJbQnc1eZa-1jOhROHawB87t4,1366 +torch/include/ATen/ops/gelu_backward_cuda_dispatch.h,sha256=nOxRijUP-6M9bTBnN2oSiVeOKHKt6Rz07GpDqFPCRe4,1368 +torch/include/ATen/ops/gelu_backward_meta.h,sha256=4RHWdiRlzsQCA0cQ9Lge1qR4f-DiE4v6FLg_3ZkGctA,890 +torch/include/ATen/ops/gelu_backward_meta_dispatch.h,sha256=Ixqp-Mls_KEbiyZPDhyYsjYVUMgvGzZhoFKAztx_OQE,1368 +torch/include/ATen/ops/gelu_backward_native.h,sha256=MM-rIHQUSwnxaT2iEnBZEx63xfWo_QyVNkGsyUHU2bk,1448 +torch/include/ATen/ops/gelu_backward_ops.h,sha256=Y-8dgBy0WiWrNg9mavdz13DoSHezCrcDO4QGFqy1blc,2300 +torch/include/ATen/ops/gelu_compositeexplicitautogradnonfunctional_dispatch.h,sha256=deL4YyKRMLXJn68NNqcq_HROSaioghWGTpHvqvXOgNg,1165 +torch/include/ATen/ops/gelu_cpu_dispatch.h,sha256=meuSCIb_unDFHsI37OwVqytfn4eIfWo96LYY2-pLmfM,1315 +torch/include/ATen/ops/gelu_cuda_dispatch.h,sha256=KU-wjmLNzE3nIdHL5IIgT7p1oiktGufk6Dufkku99EA,1317 +torch/include/ATen/ops/gelu_meta.h,sha256=AqCyPaeBZNyao5LrSca2hRVWHuHbbo032346AdsPUbI,849 +torch/include/ATen/ops/gelu_meta_dispatch.h,sha256=FSUWNU62DTEJ5hzCvEEmFZGJZFD2Dzk-lRGbUwtpl90,1317 +torch/include/ATen/ops/gelu_native.h,sha256=wMSEfMwO4_1839qER50KUy8ZKUv4A5Nh1iTpfKFWoxQ,1654 +torch/include/ATen/ops/gelu_ops.h,sha256=zqy2KdqC6XSc2RySVb2mXlWEFPDjGGVl8F_qyg7hzO0,2594 +torch/include/ATen/ops/geometric.h,sha256=v-NdFjyWYWsDemX-jxI-GvoZknYeSOccGec2DG4UKvo,1650 +torch/include/ATen/ops/geometric_compositeexplicitautograd_dispatch.h,sha256=jzHhEsFd2wexc4_aIcMUc4natsMyXo3-NzLvh3PQ0Yk,1370 +torch/include/ATen/ops/geometric_cpu_dispatch.h,sha256=p8VfEpw-FOc6FXnzwiH5tLRQvTpwiKbcUYDf0Kz2g5I,1041 +torch/include/ATen/ops/geometric_cuda_dispatch.h,sha256=YvAXhmiYbpz4fh-jdrLk4EYhmqbxKge8qYSgp9QYnaw,1043 +torch/include/ATen/ops/geometric_meta_dispatch.h,sha256=IqeOw2Vyr7JK3YjG2vIE56O6AuELZSjtYx1o8mepSYk,1043 +torch/include/ATen/ops/geometric_native.h,sha256=wKAgHccAmE7iCxQjDMHrTpF6n_b9t6wO_BWa2zyjIzg,1054 +torch/include/ATen/ops/geometric_ops.h,sha256=4z9ZC0JG-Hw-4OS_PXk4edpTinIXgecYP56LyLOLCUc,2873 +torch/include/ATen/ops/geqrf.h,sha256=6ZAaK9W5gnqMT9-I2Fdl9-JLJrsbV0LUzKgrlZ-zoSc,1464 +torch/include/ATen/ops/geqrf_cpu_dispatch.h,sha256=IsUEeLjkD00l83l8vyZaZs8oZA6Ip25LbrOajqnnnVw,1239 +torch/include/ATen/ops/geqrf_cuda_dispatch.h,sha256=0nm5aDSkk3AAAL8FNP0gyDb3XOCUfKTUMDM5rOPMJ9U,1241 +torch/include/ATen/ops/geqrf_native.h,sha256=yB-UfHfTqHoeJw269yOhunxLYbA6_0cy6O8S41uS2nE,874 +torch/include/ATen/ops/geqrf_ops.h,sha256=0yW4k4sox_n0uGhnU45NgXemLoJ4MnX-4VZWJZsTVZ0,2041 +torch/include/ATen/ops/ger.h,sha256=aiCGNNheihEpYD_ixYKdBJXsD4TZGaIdqhpnBdKkGNM,1383 +torch/include/ATen/ops/ger_compositeimplicitautograd_dispatch.h,sha256=k2Ba7OG0dhYHYu3IAgHQB3WJyv2TmefGdxzutyWgsLQ,1241 +torch/include/ATen/ops/ger_native.h,sha256=fNB_f1Oml7qe0qB9yU3xYrkChJWt-gVBCvW9rMgxi8c,852 +torch/include/ATen/ops/ger_ops.h,sha256=1Sbhv-ZOGTHWPG0GkkRoud1josXU5xgE8CIUivXF-Ag,1949 +torch/include/ATen/ops/get_device.h,sha256=MsH4wu4BZaWrA9iOyQmP3KgORASKy8MsrcRmIS17gSM,761 +torch/include/ATen/ops/get_device_compositeimplicitautograd_dispatch.h,sha256=DKpOXXSfikV88MXECZ64lnPjmrQlzsDSzsO2brpox3k,1019 +torch/include/ATen/ops/get_device_native.h,sha256=Ag6pFJKCRD0ESAaESUf3KxohBntW6ZMF3wrHyghbQg8,731 +torch/include/ATen/ops/get_device_ops.h,sha256=ZUJHuQpApmn_7InBW-c3hr9HuzfleGEAaDcYQC9bVIs,1211 +torch/include/ATen/ops/glu.h,sha256=xn0SzjqWKMSRqOXgkEXZkztH6YBsVXshaT3shsjVnyQ,1347 +torch/include/ATen/ops/glu_backward.h,sha256=FDTTv0J9QqlFWqKyBW8kPRwEOBWNSjTp9Ayr47yaWjo,1687 +torch/include/ATen/ops/glu_backward_cpu_dispatch.h,sha256=ZIm8KMmVGqMso3kn4ucINZT-2nzTUdlueJGkx_-9e3o,1298 +torch/include/ATen/ops/glu_backward_cuda_dispatch.h,sha256=mEvAf08XBZ3gmCg0RQLwh5aVQBy1P0Hvifa-SmFsbKA,1300 +torch/include/ATen/ops/glu_backward_jvp.h,sha256=AA_MOKMZATomnaCmwl_-gvhfvFpHdYbXTKZKCFZzNFk,2044 +torch/include/ATen/ops/glu_backward_jvp_compositeexplicitautograd_dispatch.h,sha256=AVJNiZgQ7gp6nbhx1ILkJPrPCsy_1-DdPOc1Z-Tlca0,1379 +torch/include/ATen/ops/glu_backward_jvp_cpu_dispatch.h,sha256=w_tDy61KI-6nWWXtphNOqHUcfkd0hla4kUFgmh5yKsU,1103 +torch/include/ATen/ops/glu_backward_jvp_cuda_dispatch.h,sha256=s2sT-xSUKiK2e6UQIptp0pRsLQGDKwYwDv0SitEOmaw,1105 +torch/include/ATen/ops/glu_backward_jvp_native.h,sha256=iMM04n5Kdh0mmlZTb_rox4cSxUO3gJoAjhSjGwuzqAU,1066 +torch/include/ATen/ops/glu_backward_jvp_ops.h,sha256=EqE0X-U2NgSzR6L2MV9IdChg4FPiaVIHjdQkyRhmIqk,2649 +torch/include/ATen/ops/glu_backward_native.h,sha256=XaSHWK6I6N5Q47po4eEY3jj1J11F0u2snLXmuJmYFYs,1176 +torch/include/ATen/ops/glu_backward_ops.h,sha256=udmCrV_gTqRNDaflJyzextzYl9LgskM9Cvh_pOtB-bE,2175 +torch/include/ATen/ops/glu_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Yz7Mrf3IHsvsiO5NT6_DmVJBS1JU398e4vYXZYn6jUQ,1057 +torch/include/ATen/ops/glu_cpu_dispatch.h,sha256=qi_gH8FERbPmwq8KvZz8ob4OWtT_XuB5ZnOSEYPjFIA,1167 +torch/include/ATen/ops/glu_cuda_dispatch.h,sha256=qfT23KGaitJ2Vh11Fu5ozJHlaja7J_Acz5WrtYeovnw,1169 +torch/include/ATen/ops/glu_jvp.h,sha256=RSiGu2RjyHuSW4hDqYLmpkZDQrRj3EhyWu7H8frmF50,1582 +torch/include/ATen/ops/glu_jvp_compositeexplicitautograd_dispatch.h,sha256=YlOVRhc2-P9vV-EzvVHbl1316AGHx9BcbmhAvLhIpf4,1237 +torch/include/ATen/ops/glu_jvp_cpu_dispatch.h,sha256=8SnhMSnKuLy4Z8tbliRzHhTy4Ievwl-YcXS28CaeL-Y,1032 +torch/include/ATen/ops/glu_jvp_cuda_dispatch.h,sha256=aLTW1uPTxdicNR4svAd0I93v1w0og7qc8OH73NK2Q_s,1034 +torch/include/ATen/ops/glu_jvp_native.h,sha256=MeRS6rff1ToGTBWRe1E4ise6oy4RQsZWTP459cZV9Z0,924 +torch/include/ATen/ops/glu_jvp_ops.h,sha256=C8dsRFQ2rGbvhpDBmdr6oC3eReP8LFYW3Ws1eEToLw0,2191 +torch/include/ATen/ops/glu_meta.h,sha256=Ok6PBNZYx3hOC9dqFgZogvhNADAei90MVASTOKUaufw,831 +torch/include/ATen/ops/glu_meta_dispatch.h,sha256=AFOk0jhFDlkAE6h4NT_lJRnmvJIJmImjNls6grPf35s,1169 +torch/include/ATen/ops/glu_native.h,sha256=ZXDHHkGOYagfi644sIYWcXZuDGiLufip_1Hq5iHTR-M,854 +torch/include/ATen/ops/glu_ops.h,sha256=yh0Fh4Wb1-nt-Q27EP5CraohK0njCIN0DOTPmOpwLvo,1877 +torch/include/ATen/ops/gradient.h,sha256=9hh831UQ5KfvlOURZ6l4NjhzCiHST0yoiYOitDG6V64,3126 +torch/include/ATen/ops/gradient_compositeimplicitautograd_dispatch.h,sha256=nYcD48S_G7LNvQL1jMrFkaloxsWl4vgNs8SORDKUudI,2038 +torch/include/ATen/ops/gradient_native.h,sha256=nUnhXOeZz9y9ATk94XY_Wl6FWhpU8r9LuRu3y_3eEIk,1750 +torch/include/ATen/ops/gradient_ops.h,sha256=lTSLz2M9MBBGm56StlK6gU0UX8wMQA-Klp4fEvJJlo8,6526 +torch/include/ATen/ops/greater.h,sha256=qzqSySHLDwlyQhs3XrLFo06ps-WXx25DC9gP-BTPuY0,2191 +torch/include/ATen/ops/greater_compositeimplicitautograd_dispatch.h,sha256=sDCDMOnWIaAG3MWPb4whbh9CQDeiUS4NCK-t-hQkH2U,1704 +torch/include/ATen/ops/greater_equal.h,sha256=FWtJKJRjXKhC99ZS7fTmT3P08XOe0SZBXdsEMx6XFtQ,2305 +torch/include/ATen/ops/greater_equal_compositeimplicitautograd_dispatch.h,sha256=wf_64kCUUS1oa_dDclA6qk7LyXi6_a27DKmPcC8A7sc,1752 +torch/include/ATen/ops/greater_equal_native.h,sha256=Sszs-lNJkw1SUpjn9jua7bMjV1Nduif3TgrY-1GraBQ,1240 +torch/include/ATen/ops/greater_equal_ops.h,sha256=qmq2Ka_uYDqy-XL2MBoEcHLoPUFqvcEyG7kjg8zrzDE,4653 +torch/include/ATen/ops/greater_native.h,sha256=8GxNQvriWPABq6B1OM12x60CI-cmrLitz2p9aTwZpgM,1204 +torch/include/ATen/ops/greater_ops.h,sha256=C4M5YqFs43Ju2Ay7ePW4LK3YueOgapSNKHZxIPzdH_M,4545 +torch/include/ATen/ops/grid_sampler.h,sha256=LD9Ok2tqYoC5UPndDsTDnRyZO7kV0SBPDrD31jMgsjc,1143 +torch/include/ATen/ops/grid_sampler_2d.h,sha256=KFT1UlwovtC6NUTMgFGd6XW-uK61yiVhfaGbPiNa6ZE,2055 +torch/include/ATen/ops/grid_sampler_2d_backward.h,sha256=9Lv8qDvd8Bs1N60DP9_jKZvlQ8NQIof29CwCVoFCK4E,2751 +torch/include/ATen/ops/grid_sampler_2d_backward_compositeexplicitautograd_dispatch.h,sha256=KRIKJ-bn2negsyXTkrx2OTMTIC2vVayJTMxzsUKP7FM,1575 +torch/include/ATen/ops/grid_sampler_2d_backward_cpu_dispatch.h,sha256=ovSy72TICQtbm61xcM1qIfr9cVf3F5dE00hytTmNzvQ,1179 +torch/include/ATen/ops/grid_sampler_2d_backward_cuda_dispatch.h,sha256=14EWRnAqnJUSnSNCVETcvb_go7vLowutSLnF-OUHvck,1181 +torch/include/ATen/ops/grid_sampler_2d_backward_native.h,sha256=gLn0Lob5L-hqhBCE2d-lVwku7nzqQ5lxQX2dgCawl9A,1508 +torch/include/ATen/ops/grid_sampler_2d_backward_ops.h,sha256=XHs90TkJr7_mbgbHChSBBI_sdavYysNiinLwJoLQwns,3215 +torch/include/ATen/ops/grid_sampler_2d_compositeexplicitautograd_dispatch.h,sha256=ljQBVSmyOCZTAYennI08xroLp26C6On8ocqUOJzZzPk,1331 +torch/include/ATen/ops/grid_sampler_2d_cpu_dispatch.h,sha256=7zxZD2jMAtfveSPa3epi4WnqFjH9tps2CwNuxq7rZXA,1079 +torch/include/ATen/ops/grid_sampler_2d_cuda_dispatch.h,sha256=NCfAlqdXkjOpw9ISeIEtjz-KbpOqBBzb04zO_1xX2lw,1081 +torch/include/ATen/ops/grid_sampler_2d_native.h,sha256=5gFDR7dx3Yvh0PO9C9R5IkHztXGyF0MsxQpojsPhJJA,1186 +torch/include/ATen/ops/grid_sampler_2d_ops.h,sha256=_NmGwc4slgnnObZbBLVeD6w6gJuyxlmDwKW-F9iyxsw,2479 +torch/include/ATen/ops/grid_sampler_3d.h,sha256=sDeCvln2qf4qu3bQEvwZh5wXBx_k00eCNIB51dj3DvQ,2055 +torch/include/ATen/ops/grid_sampler_3d_backward.h,sha256=0XcBmRklJYsctj6awz7xdSxt_4X1hPtu4NkJWDzIu78,2751 +torch/include/ATen/ops/grid_sampler_3d_backward_compositeexplicitautograd_dispatch.h,sha256=7Erkhvn7LJQSSmR7qOXhc6BxdNxzAyOCO46qh5OW3Es,1575 +torch/include/ATen/ops/grid_sampler_3d_backward_cpu_dispatch.h,sha256=dtRsgau5lqKu-7Z69N88k09DmVFtWq4t50ux_gk-Aeo,1179 +torch/include/ATen/ops/grid_sampler_3d_backward_cuda_dispatch.h,sha256=LYe6znNZcaEPH9BpN3_kvjBIQlWRRpVtDq28nlMtrQQ,1181 +torch/include/ATen/ops/grid_sampler_3d_backward_native.h,sha256=D4VyUGxZiSmozgbbWbcALARzDytjejRYHsW4oI7rlgI,1508 +torch/include/ATen/ops/grid_sampler_3d_backward_ops.h,sha256=JdK9XZFKywFz2ze3HKYPLAxw9VsD0HgwkkOr8z0BnnE,3215 +torch/include/ATen/ops/grid_sampler_3d_compositeexplicitautograd_dispatch.h,sha256=HeB1Ca2MKf_Oqvpmdtu8EYmphTTBjQ3ehkCZ25rezbY,1331 +torch/include/ATen/ops/grid_sampler_3d_cpu_dispatch.h,sha256=FeM8WVIFus_5un3A9iGLu4NYclnRBHW8Oan-_MdbRZE,1079 +torch/include/ATen/ops/grid_sampler_3d_cuda_dispatch.h,sha256=XJlepHkCgPp-8UDDDlVKMRmWbazDy5nj3BZ1W-aPwbE,1081 +torch/include/ATen/ops/grid_sampler_3d_native.h,sha256=hFch_qdpZRmdh7cSux82tpjuGnutBQF-ct1EpF-XBQY,1186 +torch/include/ATen/ops/grid_sampler_3d_ops.h,sha256=3NHrNc3GrHZj2SZWggPDYVJn69oPEzvfnO3dcBuVizc,2479 +torch/include/ATen/ops/grid_sampler_compositeimplicitautograd_dispatch.h,sha256=OVVInODeVkoinPvRhmfEGPhdm0MkglexEBmBo7XBHSc,1120 +torch/include/ATen/ops/grid_sampler_native.h,sha256=uc3jYmBAyf4Wnw43hqNSYGX3f0S7YWu1ZjMZNh1Wimw,832 +torch/include/ATen/ops/grid_sampler_ops.h,sha256=WQy9BjniNDggIcitPQYXxRxrzUvjSLnp8rWUK-AOqNc,1541 +torch/include/ATen/ops/group_norm.h,sha256=Md9aGrPMKZJxgZ6O4Cmbi4XnvFEIV7aZeFxzsGREL0Q,1209 +torch/include/ATen/ops/group_norm_compositeimplicitautograd_dispatch.h,sha256=lxV6G2liF6LMf1yOhMbqTspk-dGXlAq5fP0y04HfDI0,1178 +torch/include/ATen/ops/group_norm_native.h,sha256=_Uo0D_CATGofe4S-iq_mqjd3969iBH_YGcPqOtkzq0c,890 +torch/include/ATen/ops/group_norm_ops.h,sha256=gKvbxsHXoQj_O7QZf4IVWH8juswLLQ5bzRvNxWBIkyQ,1697 +torch/include/ATen/ops/gru.h,sha256=K5NeoaUxqLzFQZQzG2Xeh0_evEmjwUxlzLktEdUtKkg,1826 +torch/include/ATen/ops/gru_cell.h,sha256=v2dgfe5obQ4tHFCo27x5XsjlBXQYVCwd6qBp8dVdz68,1168 +torch/include/ATen/ops/gru_cell_compositeimplicitautograd_dispatch.h,sha256=GULA8vDuuN9ZgBPDtOJP3x7gYMRHZJoitcLodWUYP3A,1184 +torch/include/ATen/ops/gru_cell_native.h,sha256=6ijjmxKg2vi4tfAxpBDLpsgMI9v2igfzSuL2_Ho3oBs,896 +torch/include/ATen/ops/gru_cell_ops.h,sha256=onqMGqg1LXKYYDNIv_4ayhhAEV489qhltn78iWE2kaE,1743 +torch/include/ATen/ops/gru_compositeimplicitautograd_dispatch.h,sha256=h_6VkDT1dhwlgOmg_rN4AJlVLstpnhzhVkHXeGko5d8,1429 +torch/include/ATen/ops/gru_native.h,sha256=XD2FpZ6--_yx-aE9wMnHuLNVfK2g0zqvVk5QiKXGh3I,1141 +torch/include/ATen/ops/gru_ops.h,sha256=W-YdrUAbeT9JAF2c8dLkOeleDdOQKvep9J226KpOQ84,2943 +torch/include/ATen/ops/gt.h,sha256=CmCWKASvzFZ9M1VJovX8uH34Q526vStvj7hDpVWMZgo,2096 +torch/include/ATen/ops/gt_compositeexplicitautogradnonfunctional_dispatch.h,sha256=NSZDdJ_G1bGiNn3pe-YYUbwWx5zI86lf7-mWW-fn-SE,1288 +torch/include/ATen/ops/gt_cpu_dispatch.h,sha256=lMFDI_bII8HyfD9cY6VX1JIyXO7hHxaK1jJkRpx329c,1620 +torch/include/ATen/ops/gt_cuda_dispatch.h,sha256=d3D8iRz0SM8oSCihqvy3EPh6s1_0oeYyD9knOA_qYVU,1622 +torch/include/ATen/ops/gt_meta.h,sha256=bjIYECmoBhmaMKTls3_ljcpWmBuvurLZiFbvqLndtek,989 +torch/include/ATen/ops/gt_meta_dispatch.h,sha256=_LF7A-xng55zcgEHKFP2_HAWTE5ZY7Iuyl53rLm-w7c,1622 +torch/include/ATen/ops/gt_native.h,sha256=ACWqY-pW4dQ3dasiX2NnvQvPvwMO9xaJ53--v4E-Eh0,1549 +torch/include/ATen/ops/gt_ops.h,sha256=2pB6HoSEhQ0RMXjEhDShDBlX1mrkl0civmeJEuJXdfw,4455 +torch/include/ATen/ops/hamming_window.h,sha256=FyyJIGMLRb0NL4nCw42h2DdpqAGteTTlP4mzLrSlvmU,7292 +torch/include/ATen/ops/hamming_window_compositeexplicitautograd_dispatch.h,sha256=EBQPtI7JcWXfDpF6NHCTA095VvsQgwtgCRGIkv-qwLw,3164 +torch/include/ATen/ops/hamming_window_native.h,sha256=oaRl9aVs9pZm1iBJj3cKw3b4l47ZEx1rcksszVTZUqU,2096 +torch/include/ATen/ops/hamming_window_ops.h,sha256=6axxbCRGscGsuapARD1v4xHb2Jz5h2uw55f6nALGnp0,8030 +torch/include/ATen/ops/hann_window.h,sha256=3CQtp35kyUTx8q_F9KeKJpuq_xVZqbD5yNr3k8VOl4k,3571 +torch/include/ATen/ops/hann_window_compositeexplicitautograd_dispatch.h,sha256=M1TLXUuEciKp4-0K562sGpOG6nQvUx0vtBI_MEZK-hU,1928 +torch/include/ATen/ops/hann_window_native.h,sha256=Qsluvu6NvnYd0NkMQJ36QlRYYYDFeKmFV8QhY7w2na8,1305 +torch/include/ATen/ops/hann_window_ops.h,sha256=ClgbvDRxSz05StaGb6GECqzYTGJG8ZdRL05Qm9Jtdy8,4062 +torch/include/ATen/ops/hardshrink.h,sha256=v4AkR70pPWqzOCN8pt9DX2dLyD4KoO9-8_H_UFloU4I,1482 +torch/include/ATen/ops/hardshrink_backward.h,sha256=p3DL3FVmb-ua_7YSraqXUGNvaW6HdCRC8H0y6l5ZpUQ,1790 +torch/include/ATen/ops/hardshrink_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=I9qcMLH-e9Kb-s9sxK2uru8CsAVWUx30BV50Mh35vUY,1112 +torch/include/ATen/ops/hardshrink_backward_cpu_dispatch.h,sha256=X9QDbS8L03BSI0oa_y-TFEUrALfym3ch3cTSWHdTm_g,1349 +torch/include/ATen/ops/hardshrink_backward_cuda_dispatch.h,sha256=0Kg3R8hd1Yizq2fa0Yjjm1j4AxZa74sh2JkNOinMWJs,1351 +torch/include/ATen/ops/hardshrink_backward_meta.h,sha256=1PfH8EQsbe3KqDn-vskQuSV9bAyiwW6r8hARZRXoeX4,889 +torch/include/ATen/ops/hardshrink_backward_meta_dispatch.h,sha256=KSzzTmveGAsorNNOB-oDFLMr9OxmM8XoVdG2Fq7pW0E,1351 +torch/include/ATen/ops/hardshrink_backward_native.h,sha256=AhuOadw4gl7-6_LlpdHp4TLbxft1EoGRqeo13mTcT0g,951 +torch/include/ATen/ops/hardshrink_backward_ops.h,sha256=0cbS2qO66bPfhJnzmYRLa-W7QcagR0hArQrE4wCeGIU,2283 +torch/include/ATen/ops/hardshrink_compositeexplicitautogradnonfunctional_dispatch.h,sha256=P3Q_Ju7TDpOPJTqKJDkAgRiWapwgqvUVsZa6tpbqGSM,1078 +torch/include/ATen/ops/hardshrink_cpu_dispatch.h,sha256=ZcNTVnbczXeP2HiOuR_JZdJMJ7zV79JJABaCR-ZnRF4,1229 +torch/include/ATen/ops/hardshrink_cuda_dispatch.h,sha256=n8au5-nwnKrQnv8RoJLRhTb-qTo-4aynVDPn-19LwUA,1231 +torch/include/ATen/ops/hardshrink_meta.h,sha256=zyKrCYPogrTpLd5ZueUm2U46TPtWLb6M3mfMBgCWRF8,851 +torch/include/ATen/ops/hardshrink_meta_dispatch.h,sha256=_DhDjqnuKEwdzWnPZfEdinPOQ3aB_bM-4yFNSZaudz8,1231 +torch/include/ATen/ops/hardshrink_native.h,sha256=0yNAMTnSz0sCob1mdHjwoeai0N07khAcBwECpncwsUM,888 +torch/include/ATen/ops/hardshrink_ops.h,sha256=WUzSkoeSpr27j3SrAHcGXXGAUMGZM8wb3g-FNoFf5dE,2005 +torch/include/ATen/ops/hardsigmoid.h,sha256=brGWlSoNJ9FGBbvh8V-WLZDDaIgu5DKwprCDmVFm1jc,1488 +torch/include/ATen/ops/hardsigmoid_backward.h,sha256=X3Wagbgdt2lSSN7i-CdsDY3AvN4ZtDP-CV_2w0T_iL4,1686 +torch/include/ATen/ops/hardsigmoid_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=qluB9yj4ThqAkpe6MeczqcpZjOrkAMQnVHY9m0bXHG4,1090 +torch/include/ATen/ops/hardsigmoid_backward_cpu_dispatch.h,sha256=hTMduBEVsEW8pBV6Gwa0NkXt1SZQHeU8QQlal0JBH6Y,1283 +torch/include/ATen/ops/hardsigmoid_backward_cuda_dispatch.h,sha256=PAXvzFC3ydQQ-uCNB1qLP8vGSu8jJupN7_XgnMPXbXU,1285 +torch/include/ATen/ops/hardsigmoid_backward_meta.h,sha256=o6buS1fDFRYWLxwJUd-E-3vJKYuAiKR94cIVoFNP6Q8,867 +torch/include/ATen/ops/hardsigmoid_backward_meta_dispatch.h,sha256=ao6zS2EQFNkk8WY8jo-PULEpKT0dvIyn5DAuBoGnB6o,1285 +torch/include/ATen/ops/hardsigmoid_backward_native.h,sha256=5h2m6JbX4oXG930D2ZMUyJKalHVvDZoas8H8vyULBr0,931 +torch/include/ATen/ops/hardsigmoid_backward_ops.h,sha256=_w_zQos8l7nLKs3ZhFXY9OvpWhU2HCbLRC-nK9cBOrw,2135 +torch/include/ATen/ops/hardsigmoid_compositeexplicitautogradnonfunctional_dispatch.h,sha256=uJpvTur1T6EnTaXyetF_ZBViT6eGWGcYCV04wQhtvbU,1105 +torch/include/ATen/ops/hardsigmoid_cpu_dispatch.h,sha256=x69Rk2-jGgjyhZaCaYDhPs5jZOvx0b8Ts3p19y_RBNI,1202 +torch/include/ATen/ops/hardsigmoid_cuda_dispatch.h,sha256=pXw3--3Fg91280E_qeWb6azJ1hyrweYcuukW_DAyg44,1204 +torch/include/ATen/ops/hardsigmoid_meta.h,sha256=0AxhGV39dK8dV8CELTGTA0AHxz655gkCO6iF32ONMqs,826 +torch/include/ATen/ops/hardsigmoid_meta_dispatch.h,sha256=S4LJUEpOBb0Lq1MLYRVjjOnZa4RsHLETO6ErJ6MYgaA,1204 +torch/include/ATen/ops/hardsigmoid_native.h,sha256=Yhj06Bd0NqZzQetCgd3Zn4O2gftz3izt61E_h6VgfNc,1035 +torch/include/ATen/ops/hardsigmoid_ops.h,sha256=nIBmfiMLJrIBJzlIt89ZC0MMJmF2w9eEi2mMye4EGwE,2345 +torch/include/ATen/ops/hardswish.h,sha256=NPdkaxJnLdXrETxx2fBQX5OnU_0S2qGFdHA0LCqEObc,1462 +torch/include/ATen/ops/hardswish_backward.h,sha256=F6J4C7OlRIbDjOnuuE5UJ8YXmbyaxWm2tDVg1HvkiDs,1596 +torch/include/ATen/ops/hardswish_backward_compositeexplicitautograd_dispatch.h,sha256=8tuWr39kiTWe3wdcMjXvILHqpdl1L3Qke0wPVZ1ZYKs,1209 +torch/include/ATen/ops/hardswish_backward_cpu_dispatch.h,sha256=_YVgkurVjT3Y2lPgkQ2NqymLKIEgEJE81tAh_J8VJBM,1018 +torch/include/ATen/ops/hardswish_backward_cuda_dispatch.h,sha256=udWkDTXwXGsZoQ6YfABnD3cHKJb1VI5oPfd9PsF5rzQ,1020 +torch/include/ATen/ops/hardswish_backward_native.h,sha256=n-H0PPZYyd4TWl-7YwhnNmYqUmD_3UHgEIDPGFjklIc,896 +torch/include/ATen/ops/hardswish_backward_ops.h,sha256=96LUi3Y6HM9X_HT1o_Y-Fmh7fVWjgKmgzlISx5BAilI,2081 +torch/include/ATen/ops/hardswish_cpu_dispatch.h,sha256=L48BaSexR4UQMbno6As1lft5DRf8Q4lAMk83ZXmrU1A,1194 +torch/include/ATen/ops/hardswish_cuda_dispatch.h,sha256=ox8utXQY3pXoeCEAMXdVf6jLM_3Rrs2v6K8xv7YJGgk,1196 +torch/include/ATen/ops/hardswish_meta_dispatch.h,sha256=81KVnWt6VqCOHNs9S1rlMMXgsIrNMpFrNbz8z0Lml2o,976 +torch/include/ATen/ops/hardswish_native.h,sha256=MzFXLh3k-boA_Wym8j5qtnNwH6OKsathU5Z7UDtbxhQ,868 +torch/include/ATen/ops/hardswish_ops.h,sha256=HCzdXInIg9ei4WhwwC6hTq2v0myYNiPnPdDrWYQCQHE,2327 +torch/include/ATen/ops/hardtanh.h,sha256=5rmBGsJC8zoAkynQxTfn5nc03bn0oAR5mPcnyW9_PUU,1908 +torch/include/ATen/ops/hardtanh_backward.h,sha256=Kz64IMoPLNp9EEonAP8N0vGjaXiyHlaN-oEO2XHXhMw,1974 +torch/include/ATen/ops/hardtanh_backward_cpu_dispatch.h,sha256=2Xrp4AK41iEA1jKvO5qyn5fVpy1WHUfRWpQ1U9FaI_s,1442 +torch/include/ATen/ops/hardtanh_backward_cuda_dispatch.h,sha256=ZjHrYw7iJOIFKIs2DPuhIBbdaZEMpdFsUlT0tyJOiiU,1444 +torch/include/ATen/ops/hardtanh_backward_native.h,sha256=uy80hqPZkoirL2SjgwQDeVmbQB-Wo2V8s_b9LHu8MLc,1013 +torch/include/ATen/ops/hardtanh_backward_ops.h,sha256=IDNN8AnLNnmR7_CHqEuviAqvgrhi5U-Mz7fxNAM4Yuo,2485 +torch/include/ATen/ops/hardtanh_cpu_dispatch.h,sha256=zO3URwZcWv4XyM4QsxPNqHrAih2EuCHs-GlyhSvP0Vw,1429 +torch/include/ATen/ops/hardtanh_cuda_dispatch.h,sha256=W9k31RCJohNIeEZDI0e6sMEZ2CH8CgWYAM_glAtYboU,1431 +torch/include/ATen/ops/hardtanh_meta_dispatch.h,sha256=tsJCqg01OXhqqRG9kmu5_mmvDDGKyRV45pr9W4BSPuc,1036 +torch/include/ATen/ops/hardtanh_native.h,sha256=pFL0LakwmvhAFO7sGAKiYH9uA5kRzcG9AeS7F3_9ZKM,1452 +torch/include/ATen/ops/hardtanh_ops.h,sha256=FcpiSvBIGhhV54OBSN3x3AAJjx-yfVU0xHRARweHOK8,2885 +torch/include/ATen/ops/hash_tensor.h,sha256=_K4GSgr0ODEiMh4GMFLeW2owM-NRo4fhC0keJRGU_2w,1704 +torch/include/ATen/ops/hash_tensor_compositeexplicitautogradnonfunctional_dispatch.h,sha256=MqDuhbX4RaVYeHO7LgC7OpuOO2v925RejDd0BuLSKYA,1109 +torch/include/ATen/ops/hash_tensor_cpu_dispatch.h,sha256=FCA-al2SwDsuPLI-PrdDEjqhUht2ry5tLebE39L78gE,1315 +torch/include/ATen/ops/hash_tensor_cuda_dispatch.h,sha256=JAjQoyC32yC-RIeMxKdRC8xSKs01JD8VjqguRQK7Ja0,1317 +torch/include/ATen/ops/hash_tensor_meta.h,sha256=2MWgy_rO0XR1LeE4j12FLFCBZJeUUL0aR0xQp6bMzNw,875 +torch/include/ATen/ops/hash_tensor_meta_dispatch.h,sha256=Hnp7ON7_tV9RASK_j_5E5dEWZT8p82Xk2UtHNw6UujA,1317 +torch/include/ATen/ops/hash_tensor_native.h,sha256=V8qQURPUIIUhRH33kNQBCguJ8WvSubLbTluEUhCIWCM,914 +torch/include/ATen/ops/hash_tensor_ops.h,sha256=zLlQSJSaE1GgMnoYESOhHoU8hoGgrIpQcSWKGjEwNUE,2188 +torch/include/ATen/ops/heaviside.h,sha256=TXM03TPE2h_EfIMsfzPnFgl5XPTxy84XUUiM4X9eEx0,1461 +torch/include/ATen/ops/heaviside_compositeexplicitautogradnonfunctional_dispatch.h,sha256=c4algCWF9CNAZONHoUIIP1c5cO6jqpbqKx3j_D1MZOw,1155 +torch/include/ATen/ops/heaviside_cpu_dispatch.h,sha256=QGlf3kg5No0qcI64dVycjcg3pZl1uR1xknBjb1wya6g,1302 +torch/include/ATen/ops/heaviside_cuda_dispatch.h,sha256=6mLKuAZm5I70niFPB2-hDuocuF_BgqlwV35p7uo-tDk,1304 +torch/include/ATen/ops/heaviside_meta.h,sha256=Pr8gCF6nhjYmeyHXc_-9vO6hOG3-o5ukuN2U2wZz-UI,851 +torch/include/ATen/ops/heaviside_meta_dispatch.h,sha256=peVOKxrFNDG48JtzTv6MZrzbO6ygmSLE-inqDbMm9fs,1304 +torch/include/ATen/ops/heaviside_native.h,sha256=uuIlBXYGmy_eQK3Z3xdMhsCu0KEv2EvL7DCDvR2tgPk,886 +torch/include/ATen/ops/heaviside_ops.h,sha256=EixNKEMB3zUxCf_gIXxMV9vvcTTEushbLBuFHS2oxoY,2594 +torch/include/ATen/ops/hinge_embedding_loss.h,sha256=_Kawd8P0i12m5HQx9iRRDA0NvWRJ79xyIISdiCYUyfE,1112 +torch/include/ATen/ops/hinge_embedding_loss_compositeimplicitautograd_dispatch.h,sha256=Xv9X7yRwKQY26ZAvgqzQ2aYM2eF3gPnb_fKe55a4Kzw,1117 +torch/include/ATen/ops/hinge_embedding_loss_native.h,sha256=EoQ7EUDr1Vu5A8HJMimkH54GsjKsbjDicnPjXexLmkQ,829 +torch/include/ATen/ops/hinge_embedding_loss_ops.h,sha256=DdQ6rRnYLHCZar4jV4Xm5m8IvjD-flP-MOVXRKf5wrc,1465 +torch/include/ATen/ops/histc.h,sha256=LdQcPudL0z9WLrekyWU5CrB1DACJvOFqvh5pe7aZbfM,1647 +torch/include/ATen/ops/histc_cpu_dispatch.h,sha256=3NQFjnAH51yWXZzUS8oNXdqwioScqWDYr7kUpSBUZ5Y,1330 +torch/include/ATen/ops/histc_cuda_dispatch.h,sha256=DLzOiDUyznsDIfv8G78IczxTpfLBULMQltq9SFbxWnc,1332 +torch/include/ATen/ops/histc_native.h,sha256=I1dQ-0oPklqqJkoDaBn1G3pqhzwAXDeUdgYFr5dVzgc,1232 +torch/include/ATen/ops/histc_ops.h,sha256=VX1uIOarE6DBY9TprVXtQ5UmToVM7y3XRhWJcgq0uTM,2225 +torch/include/ATen/ops/histogram.h,sha256=vkrSfGhiv1ArjerlM7hFskf-onWjArTVFaARrElBXAQ,3770 +torch/include/ATen/ops/histogram_cpu_dispatch.h,sha256=C4ThtlN9QObPuHOo2M_WGXB49_kSPRGrXCccEjTWOEg,2292 +torch/include/ATen/ops/histogram_native.h,sha256=mc5bJ-TWXkfPuCZNMvEjeVPkWf_E-4Y1MxlmNdATr7c,1543 +torch/include/ATen/ops/histogram_ops.h,sha256=sKFyJ2Dcku75kEZTjZP3ch993dfEOoal1EbskgAmypw,5007 +torch/include/ATen/ops/histogramdd.h,sha256=rOVwTMrfQ5XXTNLr7hotG5DQuVrV0Li76o_txoh26lA,2213 +torch/include/ATen/ops/histogramdd_compositeimplicitautograd_dispatch.h,sha256=wyFoT-gG_BZo2dlYYo3spLX49NXaJZSvgYAJkysfEYI,1699 +torch/include/ATen/ops/histogramdd_native.h,sha256=jLM7bzMKckl6hJvhw0ndJYRVQCygPIqVxyQFZsqHCuU,1411 +torch/include/ATen/ops/histogramdd_ops.h,sha256=R1MLGSvVi4jeaZPAfqDZYd4OdEceW8L9i2Okf8MjOn0,4013 +torch/include/ATen/ops/hsplit.h,sha256=Mlvg3xhYQVUXea_iUc-NOwaapbw77Jrt4f1QNkmVvwM,1193 +torch/include/ATen/ops/hsplit_compositeimplicitautograd_dispatch.h,sha256=uWtCOsxhkkIK76VYFomSbOhMnYW_XP7EwhtVFNO6QW4,1145 +torch/include/ATen/ops/hsplit_native.h,sha256=0yatecVih8I6n2TYNAyaA0xpJT1gIUrYe76LcQRvCT4,857 +torch/include/ATen/ops/hsplit_ops.h,sha256=O5J5ncjApdn1pSvWw504Egex2BNyu65YK_0i5GlN88g,1996 +torch/include/ATen/ops/hspmm.h,sha256=DprT2f7mcA2TB0dUj-l5G3oG_kbkeW5ZILwHUTeXuj8,1403 +torch/include/ATen/ops/hspmm_native.h,sha256=df_HGep2ajOKN2N-UFCGySP8Z8bjk3RX1t_MqWyyxgA,1082 +torch/include/ATen/ops/hspmm_ops.h,sha256=sayAxG2HdU8CSo_JK7u0o67oQUYnO6muOCehBVShemU,1961 +torch/include/ATen/ops/hstack.h,sha256=eJUcmXDPA7vqdQ7W3DOCFkByI5K0sDhEvlUlF7cEeEY,1302 +torch/include/ATen/ops/hstack_compositeimplicitautograd_dispatch.h,sha256=WteoJOs6qqr1IVYK_XnGAd09GKLkMfnmwGQHqzwoKd0,1172 +torch/include/ATen/ops/hstack_native.h,sha256=b-EUGlkrPtoBpRrAlbOTpZlSxjhl27YjqjbYXPKrgiA,806 +torch/include/ATen/ops/hstack_ops.h,sha256=qwxIOi_nuIw-EiAm3UdAvbHmg053DP1w8esyvZJSnlY,1799 +torch/include/ATen/ops/huber_loss.h,sha256=xuS_FrxHsefkpAy0T4JI1zuRbI0aOCtuUovE-8Hz0Fs,1783 +torch/include/ATen/ops/huber_loss_backward.h,sha256=tByeJ-n6vEsGLj7_tVYzICda7yswkuDMyADv3IBtVbA,2035 +torch/include/ATen/ops/huber_loss_backward_compositeexplicitautograd_dispatch.h,sha256=w7-u4rSO8MCq9-6ZUYMAtge0bQMjSAxMxKjCKl00bfY,1123 +torch/include/ATen/ops/huber_loss_backward_cpu_dispatch.h,sha256=V0lZElxOBAVdlRhiOTcm-G2UmdO2VbQTyK14CS3m-GA,1301 +torch/include/ATen/ops/huber_loss_backward_cuda_dispatch.h,sha256=LvATj2EzgSmHki_eSM61dn3HRGVUkIMs-qPfYESWjPY,1303 +torch/include/ATen/ops/huber_loss_backward_native.h,sha256=qTi5r3kUpZpgqB9vJ9t7acUbrOTmYQHb6fspdtzrRcU,1025 +torch/include/ATen/ops/huber_loss_backward_ops.h,sha256=H9TuUDadZPvckCgLUM4Wl-zfNF6Uo_oAat2gJKbH3JQ,2508 +torch/include/ATen/ops/huber_loss_cpu_dispatch.h,sha256=9YMqjOvBln1BtfPcQi-aSZap0h__5zT0Ucn5eNNlA2o,1371 +torch/include/ATen/ops/huber_loss_cuda_dispatch.h,sha256=f2t9RJ2awh7zbuSY6umX_YLZn_Sh01p6O7KncwaJuQI,1373 +torch/include/ATen/ops/huber_loss_native.h,sha256=6-4iN6CvQJnvZnVpRnz7uAnuwVrH-o6UbMFQhlFbYNE,960 +torch/include/ATen/ops/huber_loss_ops.h,sha256=ercboX9fOr2Yx08ce3xhl2tR7QU1IN-_iXoHElXVgHM,2243 +torch/include/ATen/ops/hypot.h,sha256=5iYlUAxQ2Zt73k0WuiSm30VNENUr0AYAVXJXa9da7TQ,1412 +torch/include/ATen/ops/hypot_compositeexplicitautogradnonfunctional_dispatch.h,sha256=K_6T-1CIyiYGrgGYK8_PdBCHhRIpBobe8pAYGafUMrc,1145 +torch/include/ATen/ops/hypot_cpu_dispatch.h,sha256=7HrOioVBGgCt-sPEJwbVSH6JckExCVUETQUKNx6vpGw,1282 +torch/include/ATen/ops/hypot_cuda_dispatch.h,sha256=KJb-6w1R4PvF3uXbHCus486X-qPcKXxtz0OLOtmd_fI,1284 +torch/include/ATen/ops/hypot_meta.h,sha256=KC6TtMhXm2Pqa4gQdaLtggrmhmf3-yH3yonQR-f5Za0,846 +torch/include/ATen/ops/hypot_meta_dispatch.h,sha256=EWNVRRAaWolzV1giKH86II3WhBBJF7YITw9_QB92GSA,1284 +torch/include/ATen/ops/hypot_native.h,sha256=yuJpaBkuaonG--6q5CZl6Uayv7uju6FNKYgiLmKwcTw,873 +torch/include/ATen/ops/hypot_ops.h,sha256=Q_e801uV1XQ13hhWU2hilMn0qlJUI0k7xyV27PydXIc,2549 +torch/include/ATen/ops/i0.h,sha256=T308xBS1wjB9K2ToInJxKKuiLg9XsBx0O11oYwoR9oU,1371 +torch/include/ATen/ops/i0_compositeexplicitautogradnonfunctional_dispatch.h,sha256=_Foe9ReQfNB-9e0lCBD5ETjO52mbMFT5kuRbRBGzSrw,1087 +torch/include/ATen/ops/i0_cpu_dispatch.h,sha256=QdV1ZOv885ouzdDR49LHHE5w7soGzfVzqu6yPH7W9SA,1166 +torch/include/ATen/ops/i0_cuda_dispatch.h,sha256=QTG2-hCydLIWXLpOaTTRuBwYd_9aySwEYIZBE0_yer4,1168 +torch/include/ATen/ops/i0_meta.h,sha256=vtqfBzUpUf78Iq9ld7QuJnsVkfFpTz3ejg8pwM12A68,817 +torch/include/ATen/ops/i0_meta_dispatch.h,sha256=wmvTwEIDLHTN30hVhlb10vQPH01o-J3sEn0Kj57UiOw,1168 +torch/include/ATen/ops/i0_native.h,sha256=85ouRhxa1L79w1WhYpG_FsYpNNrYMrI3QK-RKgZWZxQ,838 +torch/include/ATen/ops/i0_ops.h,sha256=wkVudm7WUExUpKAvrdjAIafKB8rI2sgZlUGVbj9XFGU,2264 +torch/include/ATen/ops/igamma.h,sha256=st-BaNICd7J466iGqtXudFTbOZVup9nht46aNLIyzrM,1422 +torch/include/ATen/ops/igamma_compositeexplicitautogradnonfunctional_dispatch.h,sha256=dIg5allbrVYb8KjUtkIEh2pJ-2_Hwqmlu0SGEzGQIvw,1147 +torch/include/ATen/ops/igamma_cpu_dispatch.h,sha256=XIKq35mDfWRPJaVaAn8951UHuHO97WLGf_4wMmPNI58,1286 +torch/include/ATen/ops/igamma_cuda_dispatch.h,sha256=mEFj2OyIq2pP1XeyAb6eMBClIUDrSy6EP9zsOQGICGY,1288 +torch/include/ATen/ops/igamma_meta.h,sha256=X2EIw4iTtlNlzneM1z-jgy1-pUc9GB8nSVHxAqQ_RIg,847 +torch/include/ATen/ops/igamma_meta_dispatch.h,sha256=xWPYQr1HANT8cVLarzRQmUkrDt24xxmii5_1iIqepYk,1288 +torch/include/ATen/ops/igamma_native.h,sha256=XMblXiG9VQulfksc0YqxIUlc_Z0QzWRioTJBXzhBaUc,876 +torch/include/ATen/ops/igamma_ops.h,sha256=xRlr6FgAOtUxKMauMMfwaxMycecNHkDJsiw9u1166ko,2558 +torch/include/ATen/ops/igammac.h,sha256=tsdDcPjfq57i5aIaPZRSUiNRegjdMI7MtWnvdGrEM3I,1432 +torch/include/ATen/ops/igammac_compositeexplicitautogradnonfunctional_dispatch.h,sha256=8D80aYXXuQz7BoZ5vJyJtAyGUBPoapzpUFyRqYBEWiY,1149 +torch/include/ATen/ops/igammac_cpu_dispatch.h,sha256=dgMYmRnMYDaUVQIXKwGLUd9okMTS_oelXl1qOWj7j3o,1290 +torch/include/ATen/ops/igammac_cuda_dispatch.h,sha256=8gSJnFxGgHAfFAAOmEelH8iX8aKtzB78e9NH-JdHflI,1292 +torch/include/ATen/ops/igammac_meta.h,sha256=dfACs6_XYbhf0jNRo-oZk_Bb9_Lu1efXP8MEIZDF8KU,848 +torch/include/ATen/ops/igammac_meta_dispatch.h,sha256=Y0H23tNOs3fLKDho1qmxvJrqMJaytNRqplKeZ_scQ8A,1292 +torch/include/ATen/ops/igammac_native.h,sha256=QzogQqr8kTcBc8TXvUU7-RUim4sf606CzzehPNwwJds,879 +torch/include/ATen/ops/igammac_ops.h,sha256=ARAh0VHU2m2G4LpBXGJRGNkwLFk7q0b-n_qNV_aLZeA,2567 +torch/include/ATen/ops/im2col.h,sha256=YlDUTDLoz2ZBKGl_lC46ral5zaaW17pqC9UTpKTEPug,1917 +torch/include/ATen/ops/im2col_cpu_dispatch.h,sha256=uOJ22eXzcIHH07RKU0m9CXDhduXH3Q70ULeIBOPenS0,1443 +torch/include/ATen/ops/im2col_cuda_dispatch.h,sha256=fhvxSuBrl5d5pOQlT7tpmnrRAJg6cRj7lUQexSx-NO4,1445 +torch/include/ATen/ops/im2col_native.h,sha256=ePgBbkJNl730DBnNJLCqw8elpcy9CRER-hBLuUX1BGc,1374 +torch/include/ATen/ops/im2col_ops.h,sha256=ua99MTF4_a3rc6UEhyp9ObxNoAuR_pbgATSbMQUVlBw,2489 +torch/include/ATen/ops/imag.h,sha256=yFZjfMxPblNPb9HlmZFIrTa-3XCMUX6dHJ2DbkB_8L0,889 +torch/include/ATen/ops/imag_compositeimplicitautograd_dispatch.h,sha256=SxQ-Wyw8WPQzptOBMIlPwXYIMZaFNPZLh4cK4ITrjp0,1016 +torch/include/ATen/ops/imag_native.h,sha256=WXorn8w1rdDP3XkJB848Apkh9rOrPzJj2Cg2PHjaIXs,728 +torch/include/ATen/ops/imag_ops.h,sha256=HElfQDb9uXvnYsH1XBoRnuPjMXKftZVrmKoSlXLcgLY,1211 +torch/include/ATen/ops/index.h,sha256=ruhPu8n7JyeaFk421zcfKupeioQd9FiOUv5By8f5mac,1565 +torch/include/ATen/ops/index_add.h,sha256=ZgtlPdPxvHtO5SQYW_wdcH5pzBsXfgdZ3du6WcH3Fc4,2183 +torch/include/ATen/ops/index_add_compositeexplicitautogradnonfunctional_dispatch.h,sha256=z9BTErtRCGf-rYCAaTs4VbEJF23IY-0LJPuLNXNCIVY,1289 +torch/include/ATen/ops/index_add_compositeimplicitautograd_dispatch.h,sha256=mBhbVWXroKcxrrZRDEFCGGrqXFaUh8iVWOMvYU2yiGg,1119 +torch/include/ATen/ops/index_add_cpu_dispatch.h,sha256=2F4JpX2QZjRsOI1JTmtaWXSJ81UbZ4INd03JaQlR_6g,1568 +torch/include/ATen/ops/index_add_cuda_dispatch.h,sha256=Tui8AsWwqVONrg3sor30TT9HKJ2ijGRzjDlFwb7tdb0,1570 +torch/include/ATen/ops/index_add_meta.h,sha256=DQBoGG5Y0wfzpXygggb_vAdLf0zT-eYjOTwv971Z8yU,1358 +torch/include/ATen/ops/index_add_meta_dispatch.h,sha256=YcuZ2okuVdtOVShtA6pC3f4ZuvMfCT70TnLdUIP1xqg,1570 +torch/include/ATen/ops/index_add_native.h,sha256=uKn99o38xbqpo7bz4faI7iivPwyj5mRRYADwu8oLzgI,1354 +torch/include/ATen/ops/index_add_ops.h,sha256=ki6nGowWvapAWS-D-6J-CqBJLntrNDb1llqnPTw8yWw,4112 +torch/include/ATen/ops/index_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ZV_DNR4d_ncrtDwB4RYYw1SimIxzje_0B4ktoZ_S-aY,1099 +torch/include/ATen/ops/index_copy.h,sha256=LqoYqVMVtJAuGimQRhGK_qTrCvRbtPLSmUADntRo0bc,1988 +torch/include/ATen/ops/index_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Bae7fhjweIgjzIZLkFynMrGpwni2JadYu4wspwOq4vo,1235 +torch/include/ATen/ops/index_copy_compositeimplicitautograd_dispatch.h,sha256=7cdaGi5MIuqxqEZO0HVwoTidcJmwTJ4d3DKFZ1U6ZN8,1217 +torch/include/ATen/ops/index_copy_cpu_dispatch.h,sha256=bE1T6LKEhVBpeBxVioeXMy6KnvhhtjuwGWw5Tweqk3w,1462 +torch/include/ATen/ops/index_copy_cuda_dispatch.h,sha256=ZhQFZ0qKg4h2K7zGCsFbiMq6-b4-4fcMSrx4m8CliGE,1464 +torch/include/ATen/ops/index_copy_meta.h,sha256=YWuC5O3xUPp38-j_0LjReV7VdkFWa6f6XZjIwnhRGuo,1333 +torch/include/ATen/ops/index_copy_meta_dispatch.h,sha256=lBnai1XrXcLEBNwD1rB13GFfW1YQRtY26mA6fgWYnfY,1464 +torch/include/ATen/ops/index_copy_native.h,sha256=4Iq-6mnRqfiDdqgP3G0W9nhK7OxQr4f-T6aCFx0HesI,1181 +torch/include/ATen/ops/index_copy_ops.h,sha256=tr0zwWTqD4xG0Z2x_Xuxm9ChV6U6Lw2VbkwkM6Q3Wjc,4532 +torch/include/ATen/ops/index_cpu_dispatch.h,sha256=RX6W_Jj268ioEBNpscHjqnK5srueJTkiCK8qw9XgZ4k,1296 +torch/include/ATen/ops/index_cuda_dispatch.h,sha256=LlS5NdMJKzo7u3C_dhihdlpLHI7rin7sHK0lD2a41S8,1298 +torch/include/ATen/ops/index_fill.h,sha256=NqFKcxAylfrv_gEriv2XQ35MAiDa4qcnbob0dHKKgDc,3352 +torch/include/ATen/ops/index_fill_compositeexplicitautograd_dispatch.h,sha256=Dcyf5OE5no45rVa-wSKjiQKjpEF47fy9eDLzHShCsG4,1800 +torch/include/ATen/ops/index_fill_compositeimplicitautograd_dispatch.h,sha256=8xu8vsMBY5K-3LioE8MT1_hg5Zh-YCTB_XI6HfvYy_A,1466 +torch/include/ATen/ops/index_fill_cpu_dispatch.h,sha256=0DpqbB7yHf_ophHYJhNxeuiqoQtdRG87rxwtW1lk62E,1160 +torch/include/ATen/ops/index_fill_cuda_dispatch.h,sha256=pgayRlrQ_f3zHnL4tB7yl4EPVAS-ReO1rTnJ2qEUWN0,1162 +torch/include/ATen/ops/index_fill_meta_dispatch.h,sha256=jEj5x-xgAGmEg3CRIBlrcXaorGVctBk_PBRBGHpMSxA,1162 +torch/include/ATen/ops/index_fill_native.h,sha256=uu1nCP2DlrQC7hZWUeoyDNHrv2DPv93SrZEG-hytOf0,1980 +torch/include/ATen/ops/index_fill_ops.h,sha256=JpcJJK6rF4_0k1isJdIvzB4lw6YvhHowAoSvvWnlyOo,8601 +torch/include/ATen/ops/index_meta.h,sha256=5S_h461qMH4jxskr3Fq_tAhItYCbriTDeA8Z-bFsaEQ,1724 +torch/include/ATen/ops/index_meta_dispatch.h,sha256=XrwxDDDRynUi1n8ta0AaYNqHhr1oduDbL21k4ytb2t4,1298 +torch/include/ATen/ops/index_native.h,sha256=09Lmtqaqa085HGjGWdW8AWGAoMN6DScig4af2Qx7lVs,1017 +torch/include/ATen/ops/index_ops.h,sha256=bwyezNPt29E34-RsJEJaZ8gsYxpPH_0DB36ytz_xz5w,2194 +torch/include/ATen/ops/index_put.h,sha256=Ks-6-jJ2eEoIAMJuwQSMfYCQ2BD-scLrmO2oEJuj3hI,2224 +torch/include/ATen/ops/index_put_compositeexplicitautograd_dispatch.h,sha256=mB5eByzJoPYsqozpQdwRznS7nQDjMRbvVBXCdb55LlU,1656 +torch/include/ATen/ops/index_put_native.h,sha256=oxXCba3BBaIXCVie-cbReO9NUREKQxYdyCs8Sy3IEJw,1180 +torch/include/ATen/ops/index_put_ops.h,sha256=hMUbiJ6pgrcc9G4wvOAbxUR1aFm729E_SzTBqjpvcik,3320 +torch/include/ATen/ops/index_reduce.h,sha256=qamOHdKOCypJx9Q9mG0CFV0MYBFRyhNT3OZ9qG6jdb0,2032 +torch/include/ATen/ops/index_reduce_compositeexplicitautogradnonfunctional_dispatch.h,sha256=TwWtHcnGfcnEVTeeYSj-X58QCKxW4-vKFw_70vQfBc8,1337 +torch/include/ATen/ops/index_reduce_cpu_dispatch.h,sha256=OwK25FXCWKylR0Q7nnM2s3smGdtwmx7pjD8nUq_-LF8,1661 +torch/include/ATen/ops/index_reduce_cuda_dispatch.h,sha256=9B3tVtyEAIKQG6BX12fsfWkvjBat_F8WLuVe1qkadRs,1663 +torch/include/ATen/ops/index_reduce_meta.h,sha256=yNr1Q1yztfsj21p-uf7536-OLQG_xiZaHn-vTfbnmYI,1379 +torch/include/ATen/ops/index_reduce_meta_dispatch.h,sha256=f0PHMzb771qSC1a-Mpe-4vqvJuODpN8LN58EmXrsWYw,1663 +torch/include/ATen/ops/index_reduce_native.h,sha256=93nmh87vvzHNLo0w0jXX8I_b6Yz9YfAuuyqWHwSFlZc,1250 +torch/include/ATen/ops/index_reduce_ops.h,sha256=bbqTJ0RXFEHFX8303ac8e0raRGNP4YKM9mQFhSXsSKg,3461 +torch/include/ATen/ops/index_select.h,sha256=FiMwDIdXzfSAzdDfnnSYHAJOaqZH8Qg4pdvMVlGWnOo,2436 +torch/include/ATen/ops/index_select_backward.h,sha256=slOc28DdgfNC_ubjBAmKz8yPjdS-jmQbQyrHWULP2-w,2102 +torch/include/ATen/ops/index_select_backward_compositeimplicitautograd_dispatch.h,sha256=jLGWp0uhJCSVaT2Dr9O-Eb69R6HVB4y1ghSepfMboYk,1247 +torch/include/ATen/ops/index_select_backward_native.h,sha256=gffJaodvfLdw70boye5nj0NQfgp5ddCdix0Hf5kcWI8,823 +torch/include/ATen/ops/index_select_backward_ops.h,sha256=8ME8cNRChaXiYNEv59CQ7-f8iPzfTOLas9IEI8IXsH4,1492 +torch/include/ATen/ops/index_select_compositeimplicitautograd_dispatch.h,sha256=rcy84dSmOJx15-QEvjVkrOv8XqDSe5fafG4lYF7aRso,1322 +torch/include/ATen/ops/index_select_cpu_dispatch.h,sha256=S7WsUdWsPKFvV4ZqB8rgCt_s5VoXXqlvTGadHgLZ8EY,1266 +torch/include/ATen/ops/index_select_cuda_dispatch.h,sha256=daEAg9AsS9_4NWm-lDsIppklprjd0l0_WURns5eAw28,1268 +torch/include/ATen/ops/index_select_native.h,sha256=G-_P1tWs31OBVpHYNF8dX2JWenAE3bK-8E1FuxPKUV0,1819 +torch/include/ATen/ops/index_select_ops.h,sha256=vi3v3gXx9WQ9G6RGymp3_66dKED2zZqBP8Qp5e8-ZaE,3562 +torch/include/ATen/ops/indices.h,sha256=biiLfebFduEOWdRsTTQD2akmfFG9VaEQkYMJaB0BFGM,758 +torch/include/ATen/ops/indices_compositeexplicitautograd_dispatch.h,sha256=GpSLJ6mt8tofwWsMxD2B3_4eE1qmLzWKQQywRVXrGxc,1019 +torch/include/ATen/ops/indices_copy.h,sha256=_p8Qd4y6N1_LMJnpHoDqe6K1xgdOC17PBkmNHcvjFms,1341 +torch/include/ATen/ops/indices_copy_compositeexplicitautograd_dispatch.h,sha256=oG5_Mhjt7Pc1hO8cMaCv_TF9jF8q_xRxkiYdnLJ4cs0,1133 +torch/include/ATen/ops/indices_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=oUW-_Gg9sZICspCkDSGLWgXrpORttFhhkykL7E9ejTg,1050 +torch/include/ATen/ops/indices_copy_native.h,sha256=Jhrp4gw8IkIbBSYiCA72FjEDgurR-O55c00rCo3dZps,820 +torch/include/ATen/ops/indices_copy_ops.h,sha256=BX1WBcoe3Q09d8AiiTwewazDEUnR0fNW8GKLcipvXxw,1837 +torch/include/ATen/ops/indices_native.h,sha256=bQmbcba8oUdpMoQM1qhnLTj1uwcM0nAxXRLPsfVGma0,801 +torch/include/ATen/ops/indices_ops.h,sha256=2twPl1-5RUgcCS5G-7XOgwBJoNNRb0Q0WGj98BM-On8,1220 +torch/include/ATen/ops/infinitely_differentiable_gelu_backward.h,sha256=6CL9ib1hNQQstISRcupwbnz1Kb53D-hIe8uR4X2Hn9U,1067 +torch/include/ATen/ops/infinitely_differentiable_gelu_backward_compositeimplicitautograd_dispatch.h,sha256=P0swEBGplvXoxxCPjL4J402pAfL6iwPoGpo2FrBsQxw,1076 +torch/include/ATen/ops/infinitely_differentiable_gelu_backward_native.h,sha256=Q6Xy5dnCV9Ey-5P0vnjC8Z6ayRZ_aAVluCrM1yGxMHE,788 +torch/include/ATen/ops/infinitely_differentiable_gelu_backward_ops.h,sha256=1ej29AVmtzPcIK2KqYr02DeEVjnw215g8q7tcMoacdA,1393 +torch/include/ATen/ops/inner.h,sha256=6VRopWehatz55WV_ppIbtt1CdsaK56oJ3pQpdnlojNo,1412 +torch/include/ATen/ops/inner_compositeimplicitautograd_dispatch.h,sha256=xQifvGdPeEXR5dQyZzdc7z6MfMJq_k8Qn4Lm0WpEi6M,1250 +torch/include/ATen/ops/inner_native.h,sha256=3lz0YaoHG9Cq4LwZvlSyNdQZFA5oCuTdm-OpezMcN2o,858 +torch/include/ATen/ops/inner_ops.h,sha256=WVHK0JxecEEYBCiq3tHrbrZ4eUORBh31ZnNXS1g-HVU,1967 +torch/include/ATen/ops/instance_norm.h,sha256=V1tCUnvCqvoSUY18d1pyWHT5svNsELNipScxmjCyy44,1408 +torch/include/ATen/ops/instance_norm_compositeimplicitautograd_dispatch.h,sha256=u3rK6Q5D0koH9sDa14e_qE84teW3luBNOl8jxw9YC-A,1282 +torch/include/ATen/ops/instance_norm_native.h,sha256=70-KXX4cnVxleoMFw3yMG1sktxtfIrO8QrhHVtsOgTo,994 +torch/include/ATen/ops/instance_norm_ops.h,sha256=IYdXEjZohxdjpnMDzxn-nc6lQgGyqJzlZAqUeZDoocE,2065 +torch/include/ATen/ops/int_repr.h,sha256=ycMZZTaTzKFhkPObOnDrFzauNruY1VxCBJ4_KWH-ijg,1301 +torch/include/ATen/ops/int_repr_compositeexplicitautograd_dispatch.h,sha256=fIIwBnoGsN5y1Lhg_8R90GgTBSFKrTocroCzS4z89iY,1125 +torch/include/ATen/ops/int_repr_native.h,sha256=EUKLTh4zazD3MuOy4IzYPz0e141Y4g1ryASQvggv8HM,897 +torch/include/ATen/ops/int_repr_ops.h,sha256=0uupzEn212-x-NvvVhYjQnvK9QP8v1_3aYZJFPujVAQ,1813 +torch/include/ATen/ops/inverse.h,sha256=pRWZUdhSC8wuz0E8rujCAznVF9RRRGcHtjHKJLz0PdU,1291 +torch/include/ATen/ops/inverse_compositeimplicitautograd_dispatch.h,sha256=sa_HVkfj9FO-zSrIGwXs-PReI11tSTFfKUW4aXemRgY,1178 +torch/include/ATen/ops/inverse_native.h,sha256=PH9-A5moRWzDHPcMU3uLx5KR15Zhh5oixC9hLwKu1_o,810 +torch/include/ATen/ops/inverse_ops.h,sha256=F3OKmZrWLuu7IeCtPZvXWvqFjVkpHsEdulCDQONj3AA,1807 +torch/include/ATen/ops/is_coalesced.h,sha256=ZW1-5Et3yx3yR42xodhOWhSHBf9u9xsx8_21e7gR2gA,763 +torch/include/ATen/ops/is_coalesced_compositeexplicitautograd_dispatch.h,sha256=hUlIkuWnqTJQhVbUDdNjY5jTT-MTBvaxFsKyP0OSipk,1018 +torch/include/ATen/ops/is_coalesced_native.h,sha256=l3N5JaXgAqKZQw4xe-Ei1Q7dTaRCWDjJfzyu0HVQaeo,799 +torch/include/ATen/ops/is_coalesced_ops.h,sha256=6mhh4K0KKdQEQdYlTG81pVx5L54FsyGqcv8AoQclrDw,1209 +torch/include/ATen/ops/is_complex.h,sha256=xI4aDfwm0GBVP-RjfC8VJETrSS1aMa1hWgH3i3lGGt4,910 +torch/include/ATen/ops/is_complex_compositeimplicitautograd_dispatch.h,sha256=QjufQ2YdaNIIpSfVy8q1PCdPXTiYKrHJnXl66JIJXSo,1016 +torch/include/ATen/ops/is_complex_native.h,sha256=K0z1XMOTkS1dnoPlj9KLX8WBHf0slUJ0y5SnzCgKrpw,728 +torch/include/ATen/ops/is_complex_ops.h,sha256=Rh8ByZgdXyF0GUP74ZKDa1B82HNFq8D4He6x5FxODp8,1203 +torch/include/ATen/ops/is_conj.h,sha256=YwTCcx5souEsObQYIFNDtdvbzaQbFc4DR_e8z1rw_e4,898 +torch/include/ATen/ops/is_conj_compositeimplicitautograd_dispatch.h,sha256=jj6s52CvUa2vzY8QP-aAzTCtCCenqw_IfCqtOSQKNbs,1013 +torch/include/ATen/ops/is_conj_native.h,sha256=j7u5Rb8BGwYcdsF9s8utwE4gxjg4v8igJj1_bJrzaZA,725 +torch/include/ATen/ops/is_conj_ops.h,sha256=VfXZw3JdCfKItAhvFEFLqoakVcYZ9BxWZLYqqaD-XhU,1194 +torch/include/ATen/ops/is_contiguous.h,sha256=v8QjGuWBzeWKWt70HKmTU5QH9FauY7KsCvlA1kV2JQ4,764 +torch/include/ATen/ops/is_contiguous_compositeimplicitautograd_dispatch.h,sha256=W7ygfJNr1Arx-22dK6QQ86xZa8HouLTRCBxzks0XFXw,1106 +torch/include/ATen/ops/is_contiguous_native.h,sha256=UKYScvs3zkAW_sgelr6gOexaGmljwJFE6JcXgpY1kyM,818 +torch/include/ATen/ops/is_contiguous_ops.h,sha256=sknPt-kpX65jEMoU6a2Pc2oWC1kWgsXOlutM1ClLs3c,1864 +torch/include/ATen/ops/is_distributed.h,sha256=imL-wiVkVCUHTGLBaJY39k1jnnUmdmn_xuVHe_BvXtc,915 +torch/include/ATen/ops/is_distributed_compositeimplicitautograd_dispatch.h,sha256=61V8EBwrJcd2h3csREc6vFPc_Y2CTfIFgkgmobEcrA4,1020 +torch/include/ATen/ops/is_distributed_native.h,sha256=mXFsndnWH3Q2eqkuKEIHGkT3PBXYnznwoVUNFcjIM2s,732 +torch/include/ATen/ops/is_distributed_ops.h,sha256=2tPcfGq5kkcBz3j8D1wv3hnoLYrpDPbKT27vAhsmJK0,1215 +torch/include/ATen/ops/is_floating_point.h,sha256=4IADbedwoa2UqqHiQeEQUL2uFFU-MUEME2QPMhivF1k,938 +torch/include/ATen/ops/is_floating_point_compositeimplicitautograd_dispatch.h,sha256=H59X-DeBoOABV4H6mqpXGnyXVAsxrZ-r1wqebXfhaG4,1023 +torch/include/ATen/ops/is_floating_point_native.h,sha256=yYLKOKzHOrFyPwuOXJi1TuzDcbH-ADjJ3gH0yc93Q_0,735 +torch/include/ATen/ops/is_floating_point_ops.h,sha256=C2s13gNQvqiM6hJiJ07rYwqMRRZKL75YX0o2UutF7Z8,1224 +torch/include/ATen/ops/is_inference.h,sha256=MrBFh043o5hdUsAtcM0NRa5uMb6JDy-mh4iCu4TlKHU,918 +torch/include/ATen/ops/is_inference_compositeimplicitautograd_dispatch.h,sha256=-6wpdSiNuSlOn8e8_GU7NztI15BJrqXvw_U6HVrCraE,1018 +torch/include/ATen/ops/is_inference_native.h,sha256=raRGdsoG9cXjOhou_72jDySWuYgvuOa9LEazG760ekE,730 +torch/include/ATen/ops/is_inference_ops.h,sha256=_1MgjhX7cIwaxuapmIDft_bT1M1AyVfnHwSo8ezJKWk,1209 +torch/include/ATen/ops/is_leaf.h,sha256=6p9jCL1uRs_CR0zyAvD0sNQ-p9VSV0ndQWnK_2ikYfk,758 +torch/include/ATen/ops/is_leaf_compositeimplicitautograd_dispatch.h,sha256=Hkn_EvMg2BcCzh8zvq87y-cC7FzvTyo4sPsiQMsOSvg,1013 +torch/include/ATen/ops/is_leaf_native.h,sha256=gPGiIQPdnneaHuU-dv6PKdOTh0lyjTs9GwFrJC14nrc,725 +torch/include/ATen/ops/is_leaf_ops.h,sha256=A2BFI1ADvX3VxxFRP8C9--vgkm0_BnG0RMSVcBjUdXI,1194 +torch/include/ATen/ops/is_neg.h,sha256=2NqwBkbwXrU_u3JxxYA9TesF2_E52hM6QhwoTl1sPEE,894 +torch/include/ATen/ops/is_neg_compositeimplicitautograd_dispatch.h,sha256=p2wU9Nuzx-3XJ-2Lh_CCTKdDN4RPWcwk-73akS3batg,1012 +torch/include/ATen/ops/is_neg_native.h,sha256=nV9-jkWR_jB9zqUZlW2f0Lm-y4lstQMl1pVmrWJi0vw,724 +torch/include/ATen/ops/is_neg_ops.h,sha256=Lj-KPMw9zkqyspOJH2maU7D4v9MDWNl8CfEYpU59BIc,1191 +torch/include/ATen/ops/is_nonzero.h,sha256=du320jUzSUkWBI59XkWtRENJnqAfjvYXK-s2TWuLO8c,899 +torch/include/ATen/ops/is_nonzero_compositeimplicitautograd_dispatch.h,sha256=YPhPhFiHvnixH4F2ZxHCGLSqhO0eaKK30NSHZq8XDrU,1016 +torch/include/ATen/ops/is_nonzero_native.h,sha256=fA9ymsxDYnv3bmYoUMj45W3yTi2lQ4MoNeBf2gsGzm0,728 +torch/include/ATen/ops/is_nonzero_ops.h,sha256=akpknsqDvhur33b7DYwMh2QsD_aeqix6lfVXt6T5HpQ,1203 +torch/include/ATen/ops/is_pinned.h,sha256=RFCwbemUy4l5PqvKtJZyjB42LtOaEvSo9WMQNaKHhuc,760 +torch/include/ATen/ops/is_pinned_compositeexplicitautograd_dispatch.h,sha256=mEirh2dRfyF7Y7cVxTq-CqiKsAclQQ2cq4ir7G9i2U8,1066 +torch/include/ATen/ops/is_pinned_native.h,sha256=EoZbd05jbb4Ly-BNV00lBF_IiAeccF2SzCQQCpJHGIc,1011 +torch/include/ATen/ops/is_pinned_ops.h,sha256=8r601Z8R2ohjtOXRJrcsmHx0hiufqlD6GjoPWFQiS3o,1322 +torch/include/ATen/ops/is_same_size.h,sha256=YY6JTmKvWO4vRcmqywTOhKd1yEQ82IKWutIKXEkJyhg,954 +torch/include/ATen/ops/is_same_size_compositeexplicitautograd_dispatch.h,sha256=4yVtWXLhSiNxmB-ZMbj6_jf6WHfbhcpf9Rz6nhARD5E,1044 +torch/include/ATen/ops/is_same_size_native.h,sha256=IgDqoTbfvWejDdZhPaPE1bHCtbt6AQzB5p20Jz6-EkE,843 +torch/include/ATen/ops/is_same_size_ops.h,sha256=ublGvFsLWwlhVA09TaL84h_eGcpG413qRfQ4O2GDXyk,1295 +torch/include/ATen/ops/is_set_to.h,sha256=-52DcqLN9voUTp2BcCDm6gYrjzlCbBRq1ZYVzqEF40M,760 +torch/include/ATen/ops/is_set_to_cpu_dispatch.h,sha256=_l_1fXjCdi-Bv1R426pp2bqKhQSCyEle_3fkrKKgAT0,998 +torch/include/ATen/ops/is_set_to_cuda_dispatch.h,sha256=vL6qgiFhkYwg9Sc21o4HKK9fIRiDq3CMpvJNTdXB29E,1000 +torch/include/ATen/ops/is_set_to_native.h,sha256=VcC0hPObWRnbl1HI7DrbjrCGmCjdWDadPp3YWpa4GKk,754 +torch/include/ATen/ops/is_set_to_ops.h,sha256=C2vPudmTX__mrja4NovmyJXjJ9iFxq80fehzM1TBIxE,1289 +torch/include/ATen/ops/is_signed.h,sha256=tcQ35vAA00pLtgAdgqkpbrHyS8gnRmwWU_jiG3Xus5Y,906 +torch/include/ATen/ops/is_signed_compositeimplicitautograd_dispatch.h,sha256=sgJXH_giAQDghHd8wrsmMfI2WF4pwk9lNIu99p_Ayho,1015 +torch/include/ATen/ops/is_signed_native.h,sha256=2pdzOvBDzuV0muVIwbYhYvhq2p6LUJeEAkBQByLJqk4,727 +torch/include/ATen/ops/is_signed_ops.h,sha256=2CSJ_2wRsX-VDHoERaYVy-ZmWCONG3JY05gOqoNBOaI,1200 +torch/include/ATen/ops/is_vulkan_available.h,sha256=mFdqi6hWCHywLPZarfzpqGtewj3LDxIlBcs83i_KrXI,897 +torch/include/ATen/ops/is_vulkan_available_compositeimplicitautograd_dispatch.h,sha256=tVfzrVDeJ0ZF317z4WyUUWuHUq-gu50Y9-DReBL7oss,1002 +torch/include/ATen/ops/is_vulkan_available_native.h,sha256=XgHjC0aiYvMBQmlkGPEThqaRIbj1z0jlPtGvxQoljVw,714 +torch/include/ATen/ops/is_vulkan_available_ops.h,sha256=PIsTRqyHVjFIUhxKdkkTpgQAvij8FvSl0u_EAPwK8Nc,1153 +torch/include/ATen/ops/isclose.h,sha256=-8PBlEtOMJEJOjkAbsxbIg9oEWXeELkjpv6b9iDa2DA,1083 +torch/include/ATen/ops/isclose_compositeimplicitautograd_dispatch.h,sha256=Rld3SOMcLpjog3VPgNI2Oc7WCYjJdGC5WvpG4IyUMH8,1105 +torch/include/ATen/ops/isclose_native.h,sha256=Rshvo3tkn7ojoINHclQOIbYVq9m3t-i2la2YfL9OpA8,817 +torch/include/ATen/ops/isclose_ops.h,sha256=ivRXVQ4o1feQR5ChrobS6zk9S2htBsAYOzNEsYL8YD4,1464 +torch/include/ATen/ops/isfinite.h,sha256=5ToWsv0-Q-_5EChjF9Eeo-cb79RbjbG3oGY24R8R54I,899 +torch/include/ATen/ops/isfinite_compositeimplicitautograd_dispatch.h,sha256=GZiqnUJSOFcdIyuHt0r6YzcV5Bllx1uLdLdMjonh6a0,1020 +torch/include/ATen/ops/isfinite_native.h,sha256=_u-MSvWeA8nF0Jk-_MZmf85GSjFjpC-jGQYQ0ayWIYI,732 +torch/include/ATen/ops/isfinite_ops.h,sha256=vaZ-U87s4wAkB63r8zwQWIqdRiyYwnRFXiOp39cL7kU,1217 +torch/include/ATen/ops/isin.h,sha256=WhazdlD_Opdn7LuG4A09qTWcZlO-XfCEyVTItdSRXDE,4246 +torch/include/ATen/ops/isin_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Qe13jK0y2LvbuO2G-sAH7XbcIltKxbnsJvfMKu_KIQE,1393 +torch/include/ATen/ops/isin_cpu_dispatch.h,sha256=lsT6oOW3z6kmqL5xtPsNRHjFnjzwHHrpV7p7ygmDTQs,2240 +torch/include/ATen/ops/isin_cuda_dispatch.h,sha256=SVubfGLawoH2ppQVDuZVRP05OMYFw3DyreMAm8uPH6I,2242 +torch/include/ATen/ops/isin_meta.h,sha256=fbztCmAgpbgUoxooa5MU6YVS0BbxjIIJlDn0ldMoxMQ,1288 +torch/include/ATen/ops/isin_meta_dispatch.h,sha256=D5ktQImj-v5Hd7-lDFv1deZmvty-sqff66SI_bkIIuo,2242 +torch/include/ATen/ops/isin_native.h,sha256=xB-rEmOmPEjUHnSEP_WCg3RXHORzaOypd_hzwegecb0,1413 +torch/include/ATen/ops/isin_ops.h,sha256=rfm5pWfeuu9hFdUnIutfMl_XmuhyRF_VXsvliZADHWI,5661 +torch/include/ATen/ops/isinf.h,sha256=gypDpC03jA1B4vik9uXFcWd6Uc42CbXRv2LlPHJoWt4,1271 +torch/include/ATen/ops/isinf_compositeexplicitautograd_dispatch.h,sha256=bKqdE6Vai3sIcM0CcfNHkR_6EANwEl38ZOhrJU8yDGc,1172 +torch/include/ATen/ops/isinf_native.h,sha256=TI-grvT35kUJe1hWMmb7XXsvrasd92TZ1_qjzoV_SIc,1061 +torch/include/ATen/ops/isinf_ops.h,sha256=HuWOPetJp37DYsCR60zkssPOeEiE8McCUO2xQvdlqhI,1795 +torch/include/ATen/ops/isnan.h,sha256=77TyLPRerfi28Mfbz7ACkrNvx7Sg87c0lwghJ8iSQpo,1271 +torch/include/ATen/ops/isnan_compositeexplicitautograd_dispatch.h,sha256=AwqiUdjPi7-aruCVfPahUaYe1hIKIRaku5godcXPMRw,1119 +torch/include/ATen/ops/isnan_cpu_dispatch.h,sha256=HVJtU8n3sKIjGr9OQghjM6lU9Z23ek-bNTIEgBKagVU,973 +torch/include/ATen/ops/isnan_cuda_dispatch.h,sha256=vctyu7gU3TY_2JQWYu8FpMCz-CEmh1Cp9HaYrlSAW0k,975 +torch/include/ATen/ops/isnan_native.h,sha256=WsDMzgMMIzxGCDKUE4bLAHq2x9U04vdTUMtrF5qnZL4,996 +torch/include/ATen/ops/isnan_ops.h,sha256=XdC1BkfY6aRhRnQeNxFUGojZafWgsCPLT8Uf-Ibg734,1795 +torch/include/ATen/ops/isneginf.h,sha256=iof4In0F0j-c-1_Xw-Kd3w7pjUIfiO7yNK40Rh4vSoE,1301 +torch/include/ATen/ops/isneginf_compositeexplicitautogradnonfunctional_dispatch.h,sha256=QReQGn0CbxrgaB7o_0v2lSN-PzxSuc1J8ILeeCZkhko,1046 +torch/include/ATen/ops/isneginf_cpu_dispatch.h,sha256=-2QaPawkZP2C31pkXogocXs8tJZ5HIExgoulDvZo5ks,1137 +torch/include/ATen/ops/isneginf_cuda_dispatch.h,sha256=e9Xsl_tsFmQHxCcKtKSPdhbXFSXtVLwSeLL4AlqNpFc,1139 +torch/include/ATen/ops/isneginf_meta.h,sha256=N5qQE5KonmrbHTh9hHxjAPceG7FU8obggG7zEaCh230,823 +torch/include/ATen/ops/isneginf_meta_dispatch.h,sha256=4Dx6mrypjgCIqfRd48sL1cdipVgA5NCfMlgNgqAMEdo,1139 +torch/include/ATen/ops/isneginf_native.h,sha256=1UagMP1SQ01JycJa6oKi8EA4Iocawh2Xlh5YkSJFteI,1233 +torch/include/ATen/ops/isneginf_ops.h,sha256=ei81kyj0f0fIe5KvTjdcJKyFZIkKN_GmzWb0O-igO7M,1813 +torch/include/ATen/ops/isposinf.h,sha256=WLRNW7GNpOR0DyC9ciFiA7YuQfrgQE1634q69c4fJKE,1301 +torch/include/ATen/ops/isposinf_compositeexplicitautogradnonfunctional_dispatch.h,sha256=pPubQ6FYnJ3i_teq9IguUtx8Znb9U95vBGR9-NejlB0,1046 +torch/include/ATen/ops/isposinf_cpu_dispatch.h,sha256=WzV1F_przwXLoEPvWjlz55gads9NLRxvhjXZ_K9wo_I,1137 +torch/include/ATen/ops/isposinf_cuda_dispatch.h,sha256=BptRc-BkZbX2kHsOt5EQqOapZsVl1zmjT7rFynWAKt0,1139 +torch/include/ATen/ops/isposinf_meta.h,sha256=Wbt43lJCcs_ICXWt646GlBkaMj5eJAhFJPAXK9DiyIE,823 +torch/include/ATen/ops/isposinf_meta_dispatch.h,sha256=MRlwyYcBnX6hBgNGQ_7NHkP7iy9lYQFsIaeZB0CEWYM,1139 +torch/include/ATen/ops/isposinf_native.h,sha256=fSWxoxXKNPGGYwfvse_3qQd-dVDAR4qAgFaIMnUqhKo,1233 +torch/include/ATen/ops/isposinf_ops.h,sha256=hp8P_SA6voqbfw6g4LApS7r99ZFyytRoccXrVT-JoTs,1813 +torch/include/ATen/ops/isreal.h,sha256=bGTe-VSHg9TftSCKJtiVjsFrvT2sEgoRrp6LM2FzcvM,891 +torch/include/ATen/ops/isreal_compositeimplicitautograd_dispatch.h,sha256=Jc353OyqBXgmJLYR8IUyw6IWVyIS2UFgnnVX3mJE3zA,1018 +torch/include/ATen/ops/isreal_native.h,sha256=ptkmkB4sShi0e7AagNjvClwYbMHUPHGhPDDz6mErQGA,730 +torch/include/ATen/ops/isreal_ops.h,sha256=5uby9MgEcYx71TAjx7CvWjT3Lai9A8zY4fvMDf9myFc,1211 +torch/include/ATen/ops/istft.h,sha256=Yzjv3QVopjdAt_LowWswshpTDlamT9L5y9bHDfgU24U,1492 +torch/include/ATen/ops/istft_compositeimplicitautograd_dispatch.h,sha256=9Tgl-nGsoUCxkHi8-VZbpdXTB28LZWyFaDI11XuXlqc,1346 +torch/include/ATen/ops/istft_native.h,sha256=MfJ1rPM5bgvL9iMACSA2iz1r8w_PY975VkOmRQuZOBs,1058 +torch/include/ATen/ops/istft_ops.h,sha256=PRxhyMKBPsJoht904kHDJvf3JOE1Vwc1UipHAe24r10,2054 +torch/include/ATen/ops/item.h,sha256=M_CBOXGPchUuvPzFYsD9vuEZWbSuz2Yz3MTjJP12g_o,755 +torch/include/ATen/ops/item_compositeimplicitautograd_dispatch.h,sha256=fU__fqXJsjSZ3HH9vdcEtRedSPJAEStwRwMy2uwjEOs,1016 +torch/include/ATen/ops/item_native.h,sha256=4cjkD8rHfcJVhL8a2GnyYOoA8lXJHGjZthYsHIPr0qc,728 +torch/include/ATen/ops/item_ops.h,sha256=vtZU6i98N49S9LvT761jg5601yg5_axXBufL-LBm_sI,1205 +torch/include/ATen/ops/kaiser_window.h,sha256=isilAV00c5MRgJC2VG1mN5sU4IPYtgqfELXT_ci1CE8,5258 +torch/include/ATen/ops/kaiser_window_compositeexplicitautograd_dispatch.h,sha256=7wg_rjpQCyD8YQFWsSGSPvwmKYqIoG9I8XNWhvcuPac,2516 +torch/include/ATen/ops/kaiser_window_native.h,sha256=sg6GGF0y2guMwldtdB7FcMcMaEo8L-4QfuBCBgc1nBQ,1673 +torch/include/ATen/ops/kaiser_window_ops.h,sha256=TeSIVHxwedxYocx9-wkk9t5g7INbJ3zjrlpeFKN0s_w,5919 +torch/include/ATen/ops/kl_div.h,sha256=83hIlDx7C9fFDZCcfa4OaE8tdEm08nPO6LBLTqlGR_M,1072 +torch/include/ATen/ops/kl_div_compositeimplicitautograd_dispatch.h,sha256=zxBof7Nf3LHnWwrJT_dD9PjnE9rNn5XTUZBsM7QZBXM,1107 +torch/include/ATen/ops/kl_div_native.h,sha256=cjKgIbUNHO_EAl12r4jzmEIR3AfgLAlkK5OxVF0z6Vw,819 +torch/include/ATen/ops/kl_div_ops.h,sha256=a6zf7oZSAU2a47zBWwslq7Dg8dNKph8-uJ8koYxyv1M,1433 +torch/include/ATen/ops/kron.h,sha256=kGIVCkAc6abFhxEUOVWjRq1c_6Pcib3xNF5vgeXymm4,1402 +torch/include/ATen/ops/kron_compositeimplicitautograd_dispatch.h,sha256=XV6Iwbadg4M0_rgZh8ExX4L23OKBfRFiiVq32CT2wI4,1247 +torch/include/ATen/ops/kron_native.h,sha256=9h3bqZZqoKLBKk0N7Yq5r1AxPhB67mGS8DXF6H-3pT0,856 +torch/include/ATen/ops/kron_ops.h,sha256=-YgQbZfQp2vvuCB6BcHFAFBIK11Hfxkp2FHlXYhHYao,1961 +torch/include/ATen/ops/kthvalue.h,sha256=ws_ihJD0AU2UjLSSFHY6akceyssMkje5sEyzonBhNZ0,9610 +torch/include/ATen/ops/kthvalue_compositeexplicitautograd_dispatch.h,sha256=bdHYW7tMi3O0HdKwF6-9G37cnB5qj_V7tjKiPjQa1Z4,1231 +torch/include/ATen/ops/kthvalue_compositeimplicitautograd_dispatch.h,sha256=8FrB7d28o1PKm9W_fAf5-9KbIGWU5d6cl3jmDFNgnRw,1965 +torch/include/ATen/ops/kthvalue_cpu_dispatch.h,sha256=WZVenAAr02TXBwC5sBBkJE3GbJL7-_firfsc6WRN9Aw,1642 +torch/include/ATen/ops/kthvalue_cuda_dispatch.h,sha256=Da2qPRZxqN8OZ3xS_kh1Y-9LQR1wRJ2GPSdFR6nK7wU,1644 +torch/include/ATen/ops/kthvalue_native.h,sha256=02RBJcs0AH5y_GYUh69GatE2zXwcuGtPdcRy_69adSQ,1456 +torch/include/ATen/ops/kthvalue_ops.h,sha256=B8rUwMGSNULR6zCwpXzVDn6uJBhyagdCaFkmfpAyml8,4207 +torch/include/ATen/ops/l1_loss.h,sha256=k1a9jf8A0YzQgXdpluFd0O4EqdfU2As-eknk0BkWPHk,1015 +torch/include/ATen/ops/l1_loss_compositeimplicitautograd_dispatch.h,sha256=AJymPUvlIR9MqtkSSNmnBKROoCRuweUa0F0UVKvqUdY,1085 +torch/include/ATen/ops/l1_loss_native.h,sha256=xyUtLW6_QQsdggIV_nhcmF_Pj1tqjXWh_z3Xmn5hB7k,797 +torch/include/ATen/ops/l1_loss_ops.h,sha256=qv0S-bGah_B1cU7oPah8XTK-c__YwAheKsmOKW7LCC8,1370 +torch/include/ATen/ops/layer_norm.h,sha256=qVPK32H5XXf_JjhxsW4DmAzqIFEZ_ZsX133huSLVf4U,2631 +torch/include/ATen/ops/layer_norm_compositeimplicitautograd_dispatch.h,sha256=jI01jMp18uguzxhw3LifyDkgrpBhS7qQ28Na4p6Ljjc,1429 +torch/include/ATen/ops/layer_norm_native.h,sha256=mnvTtwXqu1xB1HSXdj1I2FJPPXjUyHxnsvIrlbg9Z78,914 +torch/include/ATen/ops/layer_norm_ops.h,sha256=Xy7bSxRDT9m1fQrBtD4uwyziOMWfe7y1_K7DwtpT16w,1753 +torch/include/ATen/ops/lcm.h,sha256=4vThhCCCFB0tx2r6-tvPP49HXMdrJmlvd5BvSbI_E2k,1572 +torch/include/ATen/ops/lcm_compositeexplicitautogradnonfunctional_dispatch.h,sha256=3dujOiSGprHvzvMstZdM4wIKAYKKuQ-m8zEUkbzobog,1141 +torch/include/ATen/ops/lcm_cpu_dispatch.h,sha256=eKKeVSnNRjLDImX0rpK8hTYkFzI7Bb5ZfzT7n0jYY3A,1274 +torch/include/ATen/ops/lcm_cuda_dispatch.h,sha256=cENvljCsOCi2BgBcnjgKjzZGiKr5gpRWrXKqbcok88c,1276 +torch/include/ATen/ops/lcm_meta.h,sha256=5-fPw3Pgip9aa6zU-H3ke42xNGdmAz--9OGisNRcha0,844 +torch/include/ATen/ops/lcm_meta_dispatch.h,sha256=srquQgSVi-NxvdsW1LnMzmgtt2Ubjr-LMny8YhleJYM,1276 +torch/include/ATen/ops/lcm_native.h,sha256=ljxcPeT1HmO5qmweFB9lufXTP6NG_0uOwwGIXdfboXI,867 +torch/include/ATen/ops/lcm_ops.h,sha256=FU9_QoC_vG-FXaKBxkUcv2SmwNalCX24vP6XUJPGFFs,2531 +torch/include/ATen/ops/ldexp.h,sha256=SqGIh2-eKtyWRlm-QfPI7yl3noMYe8cTjy3a0ZE6YTE,1612 +torch/include/ATen/ops/ldexp_compositeexplicitautograd_dispatch.h,sha256=I3ji3YpqEojxSz60SXO5Oxj693OPWlbKRUzwElUhKiY,1326 +torch/include/ATen/ops/ldexp_native.h,sha256=8X_ttcVRv3r3JTpp3m_JdYvi6UUkT6-0D8ypAuji6WU,934 +torch/include/ATen/ops/ldexp_ops.h,sha256=Fcae5btGXCVQ-aYkdfUFScWUHpXy_z4luYTRewvy29k,2569 +torch/include/ATen/ops/le.h,sha256=19VZOqVVuWxuKBwioVNIVpWeL7HI_0fNmO_enauTTvs,2096 +torch/include/ATen/ops/le_compositeexplicitautogradnonfunctional_dispatch.h,sha256=WSQI30NcgnnDYsBmmdBYaa0qdGq7hEZ2XH9PIvVm3RU,1288 +torch/include/ATen/ops/le_cpu_dispatch.h,sha256=ESdwdZ3dR2FS4sD9JYnUGjh3zhxUG_ynCPIJ-9sPW6w,1620 +torch/include/ATen/ops/le_cuda_dispatch.h,sha256=zq2_YhS10bXpTb7o3XOvDdTUvxnetTVf-JBKAzNN0is,1622 +torch/include/ATen/ops/le_meta.h,sha256=9sSCUK1F71kS4jHyOwwn7oMk-UhZguRKBHoGO05iwvQ,989 +torch/include/ATen/ops/le_meta_dispatch.h,sha256=IOQmMLQohyiMFEdUcw03fFAlUwgZ3Yk0IH3cyGW5FkM,1622 +torch/include/ATen/ops/le_native.h,sha256=nRxwPXb-WqrQFwoHNFZ_hqiTTqfJ-cVqBPfLYDtctEk,1459 +torch/include/ATen/ops/le_ops.h,sha256=ZMz2UPN6CtD4weuREmGUFjrsmZFowczPflwuw94e_FU,4455 +torch/include/ATen/ops/leaky_relu.h,sha256=-CLGc7m1qfcfpvDdtFRLDm7qUsf7OOHdp6UHu3onKFo,1806 +torch/include/ATen/ops/leaky_relu_backward.h,sha256=gZf4_ZyUhppyzlqBu-tFD45mz46BYBLKA27HISIBPoA,2072 +torch/include/ATen/ops/leaky_relu_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=r6aY2_i1WLNMMUuK6_Y4hXKzSTLasOdKt4zAQ3oyP20,1145 +torch/include/ATen/ops/leaky_relu_backward_cpu_dispatch.h,sha256=hI3h3YAGs3XuW7oXyTubFwBqnFa6JE5RNuPrnCEGFD8,1448 +torch/include/ATen/ops/leaky_relu_backward_cuda_dispatch.h,sha256=vqre75Zx7jRUvG8cg1mHGhbwA50AmoNuJ5QftbMnfZI,1450 +torch/include/ATen/ops/leaky_relu_backward_meta.h,sha256=-NZNhWUdiS1eDHXWd4ZCDFDH04W_7QFBAXr87tAX3u8,922 +torch/include/ATen/ops/leaky_relu_backward_meta_dispatch.h,sha256=XEoRZ1Ah3rleU0AS2uJDkU-ELg3V2-4GQv77ijnpB4w,1450 +torch/include/ATen/ops/leaky_relu_backward_native.h,sha256=P9eNYOd-Jhtw5sZBi4pbTqPQjWM6GQKkuaPBkfWMZ7c,984 +torch/include/ATen/ops/leaky_relu_backward_ops.h,sha256=ZQLlgoUVr2-9c6mzru1hB0sPtWpwsKL5UBvl8QzB0Fo,2493 +torch/include/ATen/ops/leaky_relu_compositeexplicitautogradnonfunctional_dispatch.h,sha256=oTTwf3vV_rlfRVs6vqgveLPJ9UPvtM2_Axu6IKk-XsM,1183 +torch/include/ATen/ops/leaky_relu_cpu_dispatch.h,sha256=PVYVcXH0vmGIebh7pni8H69p1BoDB7ojYC88sAitCVw,1353 +torch/include/ATen/ops/leaky_relu_cuda_dispatch.h,sha256=2gDaaoN0VFGtuXUdq9xehwaMWWe9DuAITOrCKSL1YiM,1355 +torch/include/ATen/ops/leaky_relu_meta.h,sha256=8geIt0FuoZ_2Xs2ZNMdJr4apgtf4VGCbEPFomrSk7bA,860 +torch/include/ATen/ops/leaky_relu_meta_dispatch.h,sha256=_JaiWSEnfjscrGVN23NJc5sJ1j6DUSyViOtETKBxF4U,1355 +torch/include/ATen/ops/leaky_relu_native.h,sha256=H23MHGgMYV3cFkA7W1HVH8cAi6Xr6Yj3geORjK0rYEg,1249 +torch/include/ATen/ops/leaky_relu_ops.h,sha256=Ac0ctgrObbYc3jAADZei4CCIHNdQJaBffyEd0gJOqXc,2690 +torch/include/ATen/ops/lerp.h,sha256=cV1A3CgmYqbIxUBW6j7qf_AHb60f8mK-DQ5HCzeRV2I,2398 +torch/include/ATen/ops/lerp_compositeexplicitautogradnonfunctional_dispatch.h,sha256=tbwm2obIQWAM_mqcKMOw4uVjgk7daIvtlKFmWXZ2MHg,1396 +torch/include/ATen/ops/lerp_cpu_dispatch.h,sha256=j0dcjawLFY1Cn0qZWWZ8o3nB42Fa2THm2VvkvEESsUI,1836 +torch/include/ATen/ops/lerp_cuda_dispatch.h,sha256=VeWnSEjz4ezKT70LvY_9YCgGowLXA6XvR5X9U4fIWHo,1838 +torch/include/ATen/ops/lerp_meta.h,sha256=FWska9KUKFzIcLfW7nShwiKmkQjW5tL196GxQb_QfUg,1043 +torch/include/ATen/ops/lerp_meta_dispatch.h,sha256=ZjL6XVFv0pQTzQTBlTPMMdtUbqE002ZHRBJJEiZj5OM,1838 +torch/include/ATen/ops/lerp_native.h,sha256=2B-Qh3h5JSwM5znXFtez_DOOZjNtFl9xwQpQW7214ZA,1103 +torch/include/ATen/ops/lerp_ops.h,sha256=vlztqvQSgLaVewjEby7Ry0FFPGiS0-K_ehNAbFVWKdE,4989 +torch/include/ATen/ops/less.h,sha256=TaK3EUbj54zdC6t5fEKctgzBl4syqn0A2P_PoKFr1ds,2134 +torch/include/ATen/ops/less_compositeimplicitautograd_dispatch.h,sha256=daH3ip2_gyq5VLdE_UskyunU04o6XEFZbVr04E1ob84,1680 +torch/include/ATen/ops/less_equal.h,sha256=yJBNuQ2UY6t2pCDVSN8Vmw-2FP1F50qEHHLPbzv4vf4,2248 +torch/include/ATen/ops/less_equal_compositeimplicitautograd_dispatch.h,sha256=ShkgZRScfhJgLSs5uHQPVc5azAAdzqITZV9kmzMhDTQ,1728 +torch/include/ATen/ops/less_equal_native.h,sha256=zr9b49YifeXmmdTjF3yIMZg-lvnRClPzZBneiCUDmU8,1222 +torch/include/ATen/ops/less_equal_ops.h,sha256=IJjNdOIJbN49mbSoo00PVio8hp0KlO41JVCi4r1C0t8,4599 +torch/include/ATen/ops/less_native.h,sha256=4r3FR6IRUXdRtsgi37yF8cyvSVnG9eLQjKdaXPmcskA,1186 +torch/include/ATen/ops/less_ops.h,sha256=KgbZ94sD8_7MLN7mxFwj_7Z-zwWTd3UcavGgPYMsPs4,4491 +torch/include/ATen/ops/lgamma.h,sha256=AvNeL1aLSbccWnwoTaOWG1wSZy20Q2yMhXwR0DvLLrQ,1281 +torch/include/ATen/ops/lgamma_compositeexplicitautogradnonfunctional_dispatch.h,sha256=fs_JZBCBJfDMBuDjtLcmBS8Wa0yX6hThZ5Vm21sTmV4,1095 +torch/include/ATen/ops/lgamma_cpu_dispatch.h,sha256=zT_rfENZynhi80uTsSezmONosqKulvCxUcjE5ctUD4A,1182 +torch/include/ATen/ops/lgamma_cuda_dispatch.h,sha256=IkOFjnafoj6Tw9wCfguphdAn6L3dyD3UhNrzqaMWEJM,1184 +torch/include/ATen/ops/lgamma_meta.h,sha256=MQ44SPZgk1x1G1mZxS7KDt4rePmZ46KBxS0JZwMnu7I,821 +torch/include/ATen/ops/lgamma_meta_dispatch.h,sha256=CimaogXH3Yd3d2-J4eBfo3GVntPJA7WVgT-6qCDjFQ0,1184 +torch/include/ATen/ops/lgamma_native.h,sha256=hpsGgP8cYDAXUBR1CkajzAe1Zsf0sKlUHiIZ2ctsFWo,850 +torch/include/ATen/ops/lgamma_ops.h,sha256=vzjaqK6UdW3uSxZ75IVattiTYJ84CIUr-uNHUwtUSHI,2300 +torch/include/ATen/ops/lift.h,sha256=lmY_MR-QmdqfcxtzuQrSvmQ5WgZbuzTV5tBMb0cll8s,1261 +torch/include/ATen/ops/lift_compositeexplicitautograd_dispatch.h,sha256=a_onejlEgkMEhPoN1okFy6-odh39wP-2M-7Zf6BISeE,1169 +torch/include/ATen/ops/lift_fresh.h,sha256=od7cAqJbhdrShvmOPiYF55vFWhcVHKwaVhE5zdvW6ZM,913 +torch/include/ATen/ops/lift_fresh_compositeexplicitautograd_dispatch.h,sha256=4RHH_A_PHR8pUF-FAKeTRhtCgjwGmaD1komeePICyF4,1022 +torch/include/ATen/ops/lift_fresh_copy.h,sha256=I-NX-JmAVvUbTb2nqekp2ePTfmeKXGwuJjaW21i5dks,1371 +torch/include/ATen/ops/lift_fresh_copy_compositeexplicitautograd_dispatch.h,sha256=R9PtKEeb_4YPfSE9_IUpEZvFvcRlI1-oa_HYqtYuBP4,1139 +torch/include/ATen/ops/lift_fresh_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=XUCNNl5fEbVM5JkgHaO3LggULTUeqNMRcjeZnwAndWo,1053 +torch/include/ATen/ops/lift_fresh_copy_native.h,sha256=SeFk8VKnHQy2uYpRoFXHX85skS91wm4EE7973ZFMvEY,826 +torch/include/ATen/ops/lift_fresh_copy_ops.h,sha256=c6owk1ZThAwX817C443bW3K8YrRK4o0VzMoZjWBhO8s,1855 +torch/include/ATen/ops/lift_fresh_native.h,sha256=m2p0jxvIHFSzhQw-kqbqH-efyYMhZyytY-HGmOFnm88,734 +torch/include/ATen/ops/lift_fresh_ops.h,sha256=ATFEMUKzTc267tpBs_MnUtqxoqkXPQu0ryy5FYgE1Do,1229 +torch/include/ATen/ops/lift_native.h,sha256=wn7cGshZrMEBJTDC-KERf9bKp_AUTsmqkQHkZhC7G2I,804 +torch/include/ATen/ops/lift_ops.h,sha256=AikcYovQ71v93GeX4einlF9HYEWw3fI8v6RdeNqSBC4,1789 +torch/include/ATen/ops/linalg__powsum.h,sha256=74f_pFuwe_QXKww_dSHggG7aHm7KlmKJgnqT3dK8Njk,1172 +torch/include/ATen/ops/linalg__powsum_compositeexplicitautograd_dispatch.h,sha256=GKVHdtgZcKhsGGBqKIcWu8hHg3URBlmMoxZXdhnqbpA,1170 +torch/include/ATen/ops/linalg__powsum_cpu_dispatch.h,sha256=QmOSXyNk39nbBmbbVf7rxD9xtNmLXAqTFxXYQf7W0h8,1126 +torch/include/ATen/ops/linalg__powsum_cuda_dispatch.h,sha256=o0QSnlS4y_6J8DYL3v0fYvkRnRApf2Io0Z6PNn_oqyQ,1128 +torch/include/ATen/ops/linalg__powsum_native.h,sha256=7F81_bGfeYKjhxyz69NgWqLrpKm-TsMK4sjm0lvucbI,1093 +torch/include/ATen/ops/linalg__powsum_ops.h,sha256=krWNNuDD1WeiIWh6fvQy3WoC3NV3xH8svmBA6PkBsd0,1610 +torch/include/ATen/ops/linalg_cholesky.h,sha256=x2zwvMMkQ74N2YJrMp1h47S1rFgpD3FB-i8V1jiGd08,1497 +torch/include/ATen/ops/linalg_cholesky_compositeimplicitautograd_dispatch.h,sha256=LMDISHf4MZn8uscele95Z0Xw2sm85S_BYuhSLK_aVyk,1250 +torch/include/ATen/ops/linalg_cholesky_ex.h,sha256=660E56_DFGHmvbzYvBS1r22D2eooSADJnQll-jbqLFY,1915 +torch/include/ATen/ops/linalg_cholesky_ex_compositeexplicitautogradnonfunctional_dispatch.h,sha256=uyD3dRPDXzKPgwLxUVoBNMzbqAaUCluyH_JVxpfqsUk,1124 +torch/include/ATen/ops/linalg_cholesky_ex_cpu_dispatch.h,sha256=WyrhGxEdJ8Vs5vNaTroswiem7anacRyqFsT15qyVWmA,1397 +torch/include/ATen/ops/linalg_cholesky_ex_cuda_dispatch.h,sha256=uv_dPKUBOhXV9ipX5-AoCKLSovHNB9_GMpEPgbT00ME,1399 +torch/include/ATen/ops/linalg_cholesky_ex_meta.h,sha256=aogo9AD55KY1d0nzb-wXYO3ZULg9bqzSfA09H66jyZ4,864 +torch/include/ATen/ops/linalg_cholesky_ex_meta_dispatch.h,sha256=A-SeRYJ8OKryvZi4D_dW0cfm2IEhk8tKT-VXxIAWT-g,1399 +torch/include/ATen/ops/linalg_cholesky_ex_native.h,sha256=kbVierApm4ykLBN-1dq7VBlUtCVouKJ_lU9vNVP6dUQ,940 +torch/include/ATen/ops/linalg_cholesky_ex_ops.h,sha256=UWg7OO9Qq-6ip_aVHBvF5QM73HhSTsELWHriDxmsA7g,2361 +torch/include/ATen/ops/linalg_cholesky_native.h,sha256=Rjz1xOS_dGlqRGaxttRBlvJq_yUjJQDiryhfV3PUoE8,856 +torch/include/ATen/ops/linalg_cholesky_ops.h,sha256=0ibWvh2RPlVi7RyMRJ13mtcU5M-2MSLSPZzZt_WalgI,1954 +torch/include/ATen/ops/linalg_cond.h,sha256=zGdqpDlj8BNXjHwXiK8gF-Mz3EQuweJA8eMuri0tb88,2231 +torch/include/ATen/ops/linalg_cond_compositeimplicitautograd_dispatch.h,sha256=3_GyYscnwGetxAr6-T5NBnbCh9fnJSArX-W9eGZ8bPI,1623 +torch/include/ATen/ops/linalg_cond_native.h,sha256=ewDMhWJUUb5uO21sFw7ZsvtFEig_Sfq6ma6EcO4yEyw,1093 +torch/include/ATen/ops/linalg_cond_ops.h,sha256=k02o--ursXip0jZj-5kOICrFXsImL7l5xntF4xadGwk,3378 +torch/include/ATen/ops/linalg_cross.h,sha256=BbmgQx0LzzSvYANpDkbFadBZY-9Ox7JR5gRs9Oa9XWA,1581 +torch/include/ATen/ops/linalg_cross_compositeexplicitautogradnonfunctional_dispatch.h,sha256=xt98a9At8gD2JkkQJfpcJlYd21mw8eJ3v7g-j7-CMcY,1092 +torch/include/ATen/ops/linalg_cross_cpu_dispatch.h,sha256=-8dgrBXaUh5aDEnNRF_oBYxloIaSNng7R-bjoOGvA4I,1272 +torch/include/ATen/ops/linalg_cross_cuda_dispatch.h,sha256=rkY-sP_r8FWL0gvrY2q-4KeMxLE-oOW2QIgsHQOH4m4,1274 +torch/include/ATen/ops/linalg_cross_meta.h,sha256=R-Ld2d4O9bSk9_RXniR1VvyOKgsTzCt9WCb0JFzBpK8,866 +torch/include/ATen/ops/linalg_cross_meta_dispatch.h,sha256=0Zo0GtBYRYqL3oYCSYgpQI0pd_AbWSMy78sp0frtWPo,1274 +torch/include/ATen/ops/linalg_cross_native.h,sha256=lr9JliLKilkv8uRHTBTAmwt0-CzDderS5ug4eagC6VM,1020 +torch/include/ATen/ops/linalg_cross_ops.h,sha256=Et4fJoTSuBZDXS5H9zsJVtcUMwo4WnOf4LbQ4Vtukko,2106 +torch/include/ATen/ops/linalg_det.h,sha256=qOK9s0G_CQoyJUZ85TnBLysPVUqyzkDD6iUpdCT-fnQ,1294 +torch/include/ATen/ops/linalg_det_compositeimplicitautograd_dispatch.h,sha256=mE314h14Oj4dg4T6k69saIUyfvGCLom4VgpaVQB8vg4,1178 +torch/include/ATen/ops/linalg_det_native.h,sha256=nTJzSpZexmzHdPdMsj-xl9sGon6bBRPGxB2z7LdH9_g,810 +torch/include/ATen/ops/linalg_det_ops.h,sha256=KBEZsoISoHT5ybUkzaT9fDt47yRYOcG-C1OYEDxkzAU,1807 +torch/include/ATen/ops/linalg_diagonal.h,sha256=wKXG2ajQJG6SmQPZUHng63UFEMU61zGiS302fdzQMZo,1039 +torch/include/ATen/ops/linalg_diagonal_compositeimplicitautograd_dispatch.h,sha256=eCSgUovp3SN85fxFJaq7mYRbRouM8YHfgH4OcfHluJY,1076 +torch/include/ATen/ops/linalg_diagonal_native.h,sha256=pWV1suH6pSeTEXv4BfBsIxzBzU5wwkSi4-Wr_O0odeE,788 +torch/include/ATen/ops/linalg_diagonal_ops.h,sha256=Ry15c7gyN9b1EPM-0yZngaixLvHL0o38xGqeGIKtNhc,1393 +torch/include/ATen/ops/linalg_eig.h,sha256=jaBjrOkEfvjMLvDormk8NGAm41nRAEiw9pKj8QzgReI,1693 +torch/include/ATen/ops/linalg_eig_cpu_dispatch.h,sha256=VG0p-dx3oQOitdXiZrd9ZDxroLod4dnnAGiTd2kZiA4,1292 +torch/include/ATen/ops/linalg_eig_cuda_dispatch.h,sha256=h8vsyYjznpcFbLXB6-O3b0WeJY_QdwR-ooc1SzNhnuQ,1294 +torch/include/ATen/ops/linalg_eig_native.h,sha256=nIo1ReBx0tMiMjzlKRnbDAdLxR64OtTYRepdqvEE8lU,903 +torch/include/ATen/ops/linalg_eig_ops.h,sha256=v6tM3JQcGHsxpzJaO-M0GXoX1Xm1pqrP4HA_gYB_6sU,2172 +torch/include/ATen/ops/linalg_eigh.h,sha256=NBxfedX7HJWbeWdBU5DqSE_euNcqxhriE0nQVOLihgA,1802 +torch/include/ATen/ops/linalg_eigh_compositeimplicitautograd_dispatch.h,sha256=QWKJfNgbZHJYtOhbwjjtN6FRExRCq8cp6YGtSO5Fhaw,1398 +torch/include/ATen/ops/linalg_eigh_native.h,sha256=ZxqLYU12YUJYizF1NrxBHT5QOKRtM7KsWlSDEXpnzck,946 +torch/include/ATen/ops/linalg_eigh_ops.h,sha256=CcP4ynUIivnPkxmiig4lZjwErfuPtvSLwZ2oq1A3ljw,2323 +torch/include/ATen/ops/linalg_eigvals.h,sha256=1IkJDMb9dHcHy0-IOoV77TMKS1NfXt6ZFpwbjQOlFrw,1361 +torch/include/ATen/ops/linalg_eigvals_compositeimplicitautograd_dispatch.h,sha256=XoYonDKjvxPUHb2TuW4nsgIuegIRKCEpPA-YwyII4Hs,1026 +torch/include/ATen/ops/linalg_eigvals_cpu_dispatch.h,sha256=mTtZO6JlXXDPwIXaOzWbiyv3qHvprrsojF_ZJyunWbE,1093 +torch/include/ATen/ops/linalg_eigvals_cuda_dispatch.h,sha256=ekzpzUEHwrzThPMKZ0AB2_gwHROim6xvHyl7XcbyzGs,1095 +torch/include/ATen/ops/linalg_eigvals_native.h,sha256=ohHvZ06lKIpLcX1YLmDwys88UJjz9K3DwEEEspKTdcM,824 +torch/include/ATen/ops/linalg_eigvals_ops.h,sha256=yAHMkUQRb85BUYmrv_RaA_dMzC3yVATfXPYTlIB0Rp4,1849 +torch/include/ATen/ops/linalg_eigvalsh.h,sha256=jqVwJXVQAcWVR7cORErnsVI9iiPuFf8NynZSeRDx12s,1508 +torch/include/ATen/ops/linalg_eigvalsh_compositeimplicitautograd_dispatch.h,sha256=7f4qG5c3NfkbJNrtooeqjdjAWrr-EnKIlfx011Wzh_o,1279 +torch/include/ATen/ops/linalg_eigvalsh_native.h,sha256=NJusBBvCGAU_PZwH3Z_wVCkdq6a7TgYsos_WaUP_JVg,876 +torch/include/ATen/ops/linalg_eigvalsh_ops.h,sha256=K7YFYF0Qr1YsFWgCLqgLN8pP_1yAct6YKczovCSOh-g,2015 +torch/include/ATen/ops/linalg_householder_product.h,sha256=_YQRA0JQcCGhB4jQ0UnPtE-m_YZwGLudZUEwlyvADCQ,1613 +torch/include/ATen/ops/linalg_householder_product_cpu_dispatch.h,sha256=Y3TJaqf0h8Q6B7YnOPcn8CzyrOfQd49QV57eLjzRSzg,1266 +torch/include/ATen/ops/linalg_householder_product_cuda_dispatch.h,sha256=tyoiewXg9sp2qOjMQXOa4vqGfd7XGMPe6_LBnapuotM,1268 +torch/include/ATen/ops/linalg_householder_product_native.h,sha256=9bAjC4_sOxUF-eT91KfJtkgaTZ0L6MuL4kVxVdohJYI,898 +torch/include/ATen/ops/linalg_householder_product_ops.h,sha256=wMCzogDmFH9kjJvt5vzo4A6W8Q-Yem7gyBeaC7BtAPA,2087 +torch/include/ATen/ops/linalg_inv.h,sha256=9PUeK0zgGfDG9f-nKo8A2Wxphg7xjtVZTVkGGCkGVTM,1294 +torch/include/ATen/ops/linalg_inv_compositeimplicitautograd_dispatch.h,sha256=LstNcxPDx2Z6cuQLXVWCX7W1BG17VrxazcCsWQx0s8k,1178 +torch/include/ATen/ops/linalg_inv_ex.h,sha256=Xfue3sdq82mJTE5R6lzb18WVT0FOU4cKxu6fvFz1I5w,1793 +torch/include/ATen/ops/linalg_inv_ex_compositeexplicitautogradnonfunctional_dispatch.h,sha256=yjymCW-QdMSxSBOk4AepXkSHoIgyfOQnlll-Q-17H1s,1098 +torch/include/ATen/ops/linalg_inv_ex_cpu_dispatch.h,sha256=dXmY7mwsvfWMAe-mFvA91ufplaCuIbQe9YRfAC3_4EA,1337 +torch/include/ATen/ops/linalg_inv_ex_cuda_dispatch.h,sha256=KeeSs7xa2AGbpDZKcrTU3KCmOhIEee78eBed3SyHeQw,1339 +torch/include/ATen/ops/linalg_inv_ex_meta.h,sha256=w8BfQSfy0hDvQ8QcGdvig64mxn9rpjHxP2BR4uHpsOg,844 +torch/include/ATen/ops/linalg_inv_ex_meta_dispatch.h,sha256=abX504tv_-z3s6cYtZLqSgfFHXY8ok3sYOPEApiREkU,1339 +torch/include/ATen/ops/linalg_inv_ex_native.h,sha256=gFVCRLMrsllOp6vnA6nwhDiwnEh7Yn85fkV7oUBymmM,916 +torch/include/ATen/ops/linalg_inv_ex_ops.h,sha256=vbG2bmvSK7Zfw_RgNo4YbvpDMhQ93gfIBJDqv288vJg,2265 +torch/include/ATen/ops/linalg_inv_native.h,sha256=nrfTfn7Q1zFyHo50pg-kicygmC47FsOFpKPuZpZsKDY,810 +torch/include/ATen/ops/linalg_inv_ops.h,sha256=RVm98LlBc0ItgQ_JHveKBPwX7jm7vvoYyazpx4R6Qfw,1807 +torch/include/ATen/ops/linalg_ldl_factor.h,sha256=EzrfKYeOzJzeBfcG8tStlV3eUSEalOJSVf6iwon6Mx4,1790 +torch/include/ATen/ops/linalg_ldl_factor_compositeimplicitautograd_dispatch.h,sha256=dnUji-9watS2s40fCSf7uyxzF9Qj-CbF_erWzwQVAtE,1387 +torch/include/ATen/ops/linalg_ldl_factor_ex.h,sha256=E0JQO1G42QfQF4MXdKQx3YffYZLZWEDBQLPR8Kj3FZ4,2174 +torch/include/ATen/ops/linalg_ldl_factor_ex_compositeexplicitautogradnonfunctional_dispatch.h,sha256=AwCTPCD-AVZIpxQGkMHVpi04jK9-nCukuW-sW25KPmI,1141 +torch/include/ATen/ops/linalg_ldl_factor_ex_cpu_dispatch.h,sha256=8W8js6HFdHKPPXA5xbHSUzu-XBKwgg7XDzQ9UmtHrAk,1496 +torch/include/ATen/ops/linalg_ldl_factor_ex_cuda_dispatch.h,sha256=c3Zus3yVBoVz_RJRtFJoA_aKVzzDOJUyVBhCkuJtDw0,1498 +torch/include/ATen/ops/linalg_ldl_factor_ex_meta.h,sha256=qWKY_-QeCakfQsSnb0I-YuuNb-3JVhp2v63NyLvCvt4,870 +torch/include/ATen/ops/linalg_ldl_factor_ex_meta_dispatch.h,sha256=mpjY4aZQ8EjBPYavL05r1AzaNR-DXvOMy5DtvU4IL5w,1498 +torch/include/ATen/ops/linalg_ldl_factor_ex_native.h,sha256=gOwCZMo-OJfs7dgYSsYs41O6yv_oJ2hfc-8xd5cSAlg,978 +torch/include/ATen/ops/linalg_ldl_factor_ex_ops.h,sha256=CPzuGRUzKFjz92TO-sqjwC4LMDsWu3IKZnuRRaweD28,2589 +torch/include/ATen/ops/linalg_ldl_factor_native.h,sha256=_CV7CynjnQ7rCiFo0i0roM4-wNeHY_QlXJfaPud9UMU,940 +torch/include/ATen/ops/linalg_ldl_factor_ops.h,sha256=i5piUdrVQTCQbSS_w6qkOXPubV8mqbq4oeGp4CpjBTk,2262 +torch/include/ATen/ops/linalg_ldl_solve.h,sha256=Z7_SISkO0sw0G2SY5XEvmIAydaNKz8f2iaFMjKt1Fu4,1780 +torch/include/ATen/ops/linalg_ldl_solve_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Zo0HfLwVZCAx-86f1jMQuIYLmeA1TaRza0s-gFNpq98,1123 +torch/include/ATen/ops/linalg_ldl_solve_cpu_dispatch.h,sha256=eKVhnapxsYIIjvQbB2CAGhmR5K98PihMWocLg0Wd6wA,1362 +torch/include/ATen/ops/linalg_ldl_solve_cuda_dispatch.h,sha256=TZij5CTle16QLNRcLFNX_MjLrpqsR-OONf6jQ3qKMqU,1364 +torch/include/ATen/ops/linalg_ldl_solve_meta.h,sha256=vHrPcz1qtglcguq-c7pCMLX0ZnyxUnvosqWaOblpig4,894 +torch/include/ATen/ops/linalg_ldl_solve_meta_dispatch.h,sha256=jmJ7AooOk3MTO3cJJsbol5pYU3zvgr6RX4RzTyQGOVo,1364 +torch/include/ATen/ops/linalg_ldl_solve_native.h,sha256=fzR_UhEPiVaarKrxYscyJxWGRA32M8lPXZuO76tjNNc,943 +torch/include/ATen/ops/linalg_ldl_solve_ops.h,sha256=TYInDRUnvnZIC5IXefj0rx6jttCNMZ09B2je9ccTzHw,2298 +torch/include/ATen/ops/linalg_lstsq.h,sha256=eY4mZrc_PyJwdFT-D1T5n-zQiDGx9xmMs3oxBcwpP-o,2637 +torch/include/ATen/ops/linalg_lstsq_compositeexplicitautograd_dispatch.h,sha256=B491D3aN6MFFogIspibuGTQbqzr0OGWLH8is4ORUBSQ,1196 +torch/include/ATen/ops/linalg_lstsq_cpu_dispatch.h,sha256=ZCslv7V9EZ30Db7c9eWqSFCNBE4c_PFnpPtovqmqxdY,1571 +torch/include/ATen/ops/linalg_lstsq_cuda_dispatch.h,sha256=_sCwbXjfWwI_FYckj8xaKPzT1YRGlljEKeWZhQE7_0k,1573 +torch/include/ATen/ops/linalg_lstsq_native.h,sha256=vx2Zvdp1evPIfD1rNpFT8cI7wa254916IUURVknu4ZY,1218 +torch/include/ATen/ops/linalg_lstsq_ops.h,sha256=UWW9eLKVG2mzUWy4QmV6526ADkBYtrv6u7S7tv--vxE,3188 +torch/include/ATen/ops/linalg_lu.h,sha256=Ox2sr3s_BwLU6L1PWsNs5PHn-wscZhT6ONynY9SVtdQ,1729 +torch/include/ATen/ops/linalg_lu_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Tq4lVIBJLeV7q6mjQIc46s30oXlxKZCBj0s7h1M3U-s,1097 +torch/include/ATen/ops/linalg_lu_cpu_dispatch.h,sha256=ayco25_68R37zckSxHNJ-CsRQZBjyL5N-4-KSe2XkPw,1353 +torch/include/ATen/ops/linalg_lu_cuda_dispatch.h,sha256=WZG78J8v6u_Aaq1_Mzjo_dHYLTsRy0S_wmNcm2l2bNI,1355 +torch/include/ATen/ops/linalg_lu_factor.h,sha256=7ZD9WZHPN9ZXS9Sr_kBWGcgAmKv0hLAkbQwkukHaIuk,1712 +torch/include/ATen/ops/linalg_lu_factor_compositeimplicitautograd_dispatch.h,sha256=KFVa1qke9OCtvMJaKyopeUsw6K0q0_6lmDyiw5mrN9M,1361 +torch/include/ATen/ops/linalg_lu_factor_ex.h,sha256=y51x13dmX2afQoWE5qZ9A-YEg2C0bhXxvs30FsKVyO8,2096 +torch/include/ATen/ops/linalg_lu_factor_ex_compositeexplicitautogradnonfunctional_dispatch.h,sha256=QL4qz14k-9TP090aKKj1mufOdhmwKqbtlZnZGax_RG0,1132 +torch/include/ATen/ops/linalg_lu_factor_ex_cpu_dispatch.h,sha256=9j7hoizb5y3qZceZEa_waJs54gHx1C7xyHzJ4wBKr1Y,1470 +torch/include/ATen/ops/linalg_lu_factor_ex_cuda_dispatch.h,sha256=fIPD1N0nRKmOHIEZd-tRn6muw6e2jhv5PcD2kOc1aOQ,1472 +torch/include/ATen/ops/linalg_lu_factor_ex_meta.h,sha256=2SSMymduI4Pbj-1aB6xaOfNDZCpco10dYV-PBi6-ZJc,862 +torch/include/ATen/ops/linalg_lu_factor_ex_meta_dispatch.h,sha256=0_GZk_Y9IZHemr3xihAyX1yTC_vjj3SxeXB-4KgN3Z0,1472 +torch/include/ATen/ops/linalg_lu_factor_ex_native.h,sha256=_ArtkrzB2B8YYm_XwWvGqbHjLrraNu1-i__51HtPfIw,968 +torch/include/ATen/ops/linalg_lu_factor_ex_ops.h,sha256=2JxIPt4ZoAxtjOo4EjcpzVpSrIPqDIpoeqjjHVrPKew,2539 +torch/include/ATen/ops/linalg_lu_factor_native.h,sha256=XiDD3sEhizMJR9uKhni6Kvx1u6dypaypS-L9lRct5oo,923 +torch/include/ATen/ops/linalg_lu_factor_ops.h,sha256=5Lm4G-J8d6IEdm8xtu-bCIJE4Yuqt--8kboJBlZpbqs,2212 +torch/include/ATen/ops/linalg_lu_meta.h,sha256=zQJlysRbWPPFF__tlnFnl9kkmW5WBGv2SbyRBFklzhM,833 +torch/include/ATen/ops/linalg_lu_meta_dispatch.h,sha256=teAv6JQ3QiWoUY6rzWKwAXXs5HwzLFl_J1Z2y_lHrPo,1355 +torch/include/ATen/ops/linalg_lu_native.h,sha256=1l3QXIiwLNyRVduX7_ujPAkp-hnGMk7yXZab8RQBdh0,910 +torch/include/ATen/ops/linalg_lu_ops.h,sha256=46LiBQx1CmdoKV-ll82UMuiuzHWaH-lfgDwzCbxsBO4,2296 +torch/include/ATen/ops/linalg_lu_solve.h,sha256=s1_0al2-sDQ7ZmDj6vPk2MIF-J1iLu7OCbwTpSlfYII,1861 +torch/include/ATen/ops/linalg_lu_solve_compositeexplicitautogradnonfunctional_dispatch.h,sha256=v7y6QWty3XAkrZMN_nepVlDnGKntB2kKehqeGcgwBCU,1136 +torch/include/ATen/ops/linalg_lu_solve_cpu_dispatch.h,sha256=FtLX1ZAmhK_4pz1mnYdSyvhA_Gr43G8INhn4XcxDRqE,1396 +torch/include/ATen/ops/linalg_lu_solve_cuda_dispatch.h,sha256=kvAOyYE9MZqy1E347QKHBVyfJyoipJbdgc_nbF6PcJA,1398 +torch/include/ATen/ops/linalg_lu_solve_meta.h,sha256=JS3ZLUUNAIMFeP_T6AU4hzqwj0WyKlbiEjUpRQlnWcI,902 +torch/include/ATen/ops/linalg_lu_solve_meta_dispatch.h,sha256=MOpZJy2-nh3WaFLsh2QxoIdA2nKAtF8DdX-IIialTt8,1398 +torch/include/ATen/ops/linalg_lu_solve_native.h,sha256=_umB08vNaOgmvAyfCOYp1KSQK78NRWa8QEIlTS70gc0,949 +torch/include/ATen/ops/linalg_lu_solve_ops.h,sha256=89z2Ns2hVyUQzBV6QQUoxMDAG_rOPDh5ARH0sa18EKQ,2368 +torch/include/ATen/ops/linalg_matmul.h,sha256=eOYG8-QlmbI_uaE6SGqhenBLKDmX9Q9f4Ce0vAXlJio,1492 +torch/include/ATen/ops/linalg_matmul_compositeimplicitautograd_dispatch.h,sha256=kzYfzY65To2KgFilD5_gD4nVdvU81WfQhCSki-Yfltg,1274 +torch/include/ATen/ops/linalg_matmul_native.h,sha256=UXufPQrmjThziUKEtvx9lXNG2I0r4gp1ZxrQ4VXTQdU,874 +torch/include/ATen/ops/linalg_matmul_ops.h,sha256=XPY_relY5iZw7nDkuYZFqOv9IxZvLSR7ZiPORaE-5g4,2015 +torch/include/ATen/ops/linalg_matrix_exp.h,sha256=xXa4s2RxaJLxYJGa8TQvp11CeH6ThBjqTRuw0GIs2Io,1391 +torch/include/ATen/ops/linalg_matrix_exp_compositeexplicitautograd_dispatch.h,sha256=QRezsk1kM8wswEnsxWzcgqwMWBlggLtu1QIQJpjWcsQ,1143 +torch/include/ATen/ops/linalg_matrix_exp_cpu_dispatch.h,sha256=jetPXq9tvVH4WTrepAt9jvzeXvcrS-_LlnXb5wpzarU,985 +torch/include/ATen/ops/linalg_matrix_exp_cuda_dispatch.h,sha256=1hf255FmbFLjMl2JSZUmaQKAZwbNw2n5jXUQb6XFNLA,987 +torch/include/ATen/ops/linalg_matrix_exp_native.h,sha256=a-RGdUhRZXmEJ7i-mtZTo03NqRs3v2SbUNb9rWQVmwo,830 +torch/include/ATen/ops/linalg_matrix_exp_ops.h,sha256=CuzE6x655TUQHb6rHK23MWx6A3VNkol9aFYGH5D1NJY,1867 +torch/include/ATen/ops/linalg_matrix_norm.h,sha256=ICvhd6q6AJ4apbyzUoWgO-b3tnATPu29nhDJY3pa6b4,3413 +torch/include/ATen/ops/linalg_matrix_norm_compositeimplicitautograd_dispatch.h,sha256=z5DGd6RezpXzK8uo3SKIQO_tVWWm2FZQH2q_PLN4fpI,2168 +torch/include/ATen/ops/linalg_matrix_norm_native.h,sha256=v_V6AOtwH3Jdb7Ixxb_t2forVGbC1mE4Oa_RgsEGfBc,1440 +torch/include/ATen/ops/linalg_matrix_norm_ops.h,sha256=ehm8hiuWzIMJpllNWmyVvaayp-QVDVSLISiO0jmxrNY,4470 +torch/include/ATen/ops/linalg_matrix_power.h,sha256=UE6UvgBroiH8qT0BoSMJmCDqItymllHQ2KqxyqJzWnE,1474 +torch/include/ATen/ops/linalg_matrix_power_compositeimplicitautograd_dispatch.h,sha256=A5Zs98TOmZh2msBLhlxsm09U9J5-olKVyva7oJ8e5sc,1247 +torch/include/ATen/ops/linalg_matrix_power_native.h,sha256=VwZlIwFeyqH2ROab7KuAJC4hLF7q0sNZi81e7ZPvEZw,856 +torch/include/ATen/ops/linalg_matrix_power_ops.h,sha256=S9FYF2bcdHikNOcmF7anNWasIK-crgp8vqdL3th1HZI,1955 +torch/include/ATen/ops/linalg_matrix_rank.h,sha256=fHKhs-LorBMnPveQCT74RrDW9LIw3BAp_GoFtuBxD5w,5169 +torch/include/ATen/ops/linalg_matrix_rank_compositeimplicitautograd_dispatch.h,sha256=J3S1rCeAGL0kmob6mbyddfhueroVLgDSzm3kIZBWVqA,2750 +torch/include/ATen/ops/linalg_matrix_rank_native.h,sha256=Voe1sNBuYYihdIL8YIrtiSW0abvKhTaUttX0U93MFAA,1852 +torch/include/ATen/ops/linalg_matrix_rank_ops.h,sha256=VXWDoNbJATQx8h3Ip697D-CA59pgZRcbCEbi-RBPKyA,7336 +torch/include/ATen/ops/linalg_multi_dot.h,sha256=p5BCH44cK8ELNiec5_5yxqS1DkUw5vxGLKXQNf-xAHw,1402 +torch/include/ATen/ops/linalg_multi_dot_compositeimplicitautograd_dispatch.h,sha256=qiA2kx_ccqIyVnS2VCYLJWQ8NBa3InKYrYtyKfNm81Q,1202 +torch/include/ATen/ops/linalg_multi_dot_native.h,sha256=xDnloi3IcF_ut5X5tEOrcrNTWlB3Ld156rNMmLHbBug,826 +torch/include/ATen/ops/linalg_multi_dot_ops.h,sha256=YEwkVAY3pGETxszliueBS1DnRxxDVpIpgaeDwiFy1Jo,1859 +torch/include/ATen/ops/linalg_norm.h,sha256=INu8d_T6WFiAqx0Nu8XFZwXY6QWQ7RRMlO63-4xsqBY,3419 +torch/include/ATen/ops/linalg_norm_compositeimplicitautograd_dispatch.h,sha256=djX0meLaJeEgX3Vc9mU-PjnPJV4J2GmyCbkZtdYKPqc,2271 +torch/include/ATen/ops/linalg_norm_native.h,sha256=Qxg47-0c6o-0Enorar2HNnuswbW0BH46SDO_5Aeh8Pc,1501 +torch/include/ATen/ops/linalg_norm_ops.h,sha256=qfBoUdh0FCphEq4LQbyiS7ZJjKt39PG7buyLrWqw4xM,4580 +torch/include/ATen/ops/linalg_pinv.h,sha256=gBjCw4aUu-O9BMHvAsFmu3G2wVueKVGcZGJ0T-lXZpo,4940 +torch/include/ATen/ops/linalg_pinv_compositeexplicitautograd_dispatch.h,sha256=j83aplXZMtVlMrMVMSXHyFWv3IqF_qlHgnQPJloYlc8,1343 +torch/include/ATen/ops/linalg_pinv_compositeexplicitautogradnonfunctional_dispatch.h,sha256=hImXM_P3Dd7r3HXHUm8Wzqh-l4SEId7batxFDv85SJQ,1161 +torch/include/ATen/ops/linalg_pinv_compositeimplicitautograd_dispatch.h,sha256=1Xm7rzihAQJVokNf7pX4XD42YQQYNLuU30su1cFTTVE,2122 +torch/include/ATen/ops/linalg_pinv_native.h,sha256=mOZUTRApqEuHS1QeZf2NoPCf2E3a2foZcbQvjmhoLU0,1800 +torch/include/ATen/ops/linalg_pinv_ops.h,sha256=jEwnEHDNalnZJzQl1VIMKS6wQkzSPXI2vaD3lCi4nr0,7192 +torch/include/ATen/ops/linalg_qr.h,sha256=IH9kK8F0E7p09s28JPf03SgiRDTIoEXI9C4OQAdUxgQ,1634 +torch/include/ATen/ops/linalg_qr_compositeexplicitautogradnonfunctional_dispatch.h,sha256=-pHxs2ZtnyA-9Sj8Bii637c2_XRszVIPdiUYFKyEgdU,1102 +torch/include/ATen/ops/linalg_qr_cpu_dispatch.h,sha256=pmMyWWRMpFqDaE9ICBB5ztmS81oG9-FGdLiDZSqPTaI,1327 +torch/include/ATen/ops/linalg_qr_cuda_dispatch.h,sha256=oQs7_EPm2OhmCsaAqTxxmpxoH_LzHazTyT7kia_UD9Q,1329 +torch/include/ATen/ops/linalg_qr_meta.h,sha256=nHSyUEM8110G2goBuN-iZ6r8UiNR-wBEOUiedr14wnA,844 +torch/include/ATen/ops/linalg_qr_meta_dispatch.h,sha256=Gus9Cts7sDp73V7FeL0cZhkEqKkm-jTHMz-d4Do8z9E,1329 +torch/include/ATen/ops/linalg_qr_native.h,sha256=ry61jq-f_KBYg9gltrJLkSNdizqONALvxM7P4VLUZ1k,899 +torch/include/ATen/ops/linalg_qr_ops.h,sha256=wzy9NVttrcs_wNKitpQsAY5_C8Ab4mNJGFwLVSigTBA,2211 +torch/include/ATen/ops/linalg_slogdet.h,sha256=XO4kgsyj0-0HJkJN-0KSkAbQ9cxppi_d5KjBHHsfKw0,1616 +torch/include/ATen/ops/linalg_slogdet_compositeimplicitautograd_dispatch.h,sha256=NN65qjD1tDmbQI3mJvGs0hOBlMFuxLl8OHY5Eeft5d4,1319 +torch/include/ATen/ops/linalg_slogdet_native.h,sha256=M4vdriYbYCcjojReIOkvH5P8k0cvGS2a2OZ30F2yxJY,895 +torch/include/ATen/ops/linalg_slogdet_ops.h,sha256=5fDrKWdRfRMY18TpJel4sEFNR24rSNdQy9YspdwIZrI,2128 +torch/include/ATen/ops/linalg_solve.h,sha256=sUOUj_9oPI-PY58G75x8GWOjASVUwip5oxwrdHqWU_8,1531 +torch/include/ATen/ops/linalg_solve_compositeimplicitautograd_dispatch.h,sha256=3LU0jn8OooVJ3uKXSczqG7PJ1s-hQya-WHAQwXlAkk0,1293 +torch/include/ATen/ops/linalg_solve_ex.h,sha256=gUkvsKLqO1EO1XgmhrCAMju7NT4cbE4AUTp5vCwYUoM,2002 +torch/include/ATen/ops/linalg_solve_ex_compositeimplicitautograd_dispatch.h,sha256=bcGxED3HeyJNQ2J8epnmsk1c8vPSemCoGAdzyVGH1WE,1494 +torch/include/ATen/ops/linalg_solve_ex_native.h,sha256=YPRTBgzu-9ZQEccObZfzrJuoO-9mT3vKJdT9DcPGJsU,1009 +torch/include/ATen/ops/linalg_solve_ex_ops.h,sha256=DXTGw_jrSiGm-lrKaIul9sJuvAMnraIXabxe0vFo1yM,2496 +torch/include/ATen/ops/linalg_solve_native.h,sha256=HUJfU8CstRWvhKu2zQJ5R4C84xnO9rWfSUHM4vswnHs,885 +torch/include/ATen/ops/linalg_solve_ops.h,sha256=5FADMlo6TYZT9dGO_4805y8vgzJ9-1Mm440Af5wmtFo,2058 +torch/include/ATen/ops/linalg_solve_triangular.h,sha256=JzFF64F-6VrHIepTJlZTHY4GrrWnz2kQX66wtEeu09c,1956 +torch/include/ATen/ops/linalg_solve_triangular_cpu_dispatch.h,sha256=JEBt6K0WNu4YW4h0yQIsaLDYIr94ekZ3fKTSzZMk2Hg,1399 +torch/include/ATen/ops/linalg_solve_triangular_cuda_dispatch.h,sha256=aAktMfLDxFse-f2aWYHGtxqASEItPGtPU9TuPyoT7UM,1401 +torch/include/ATen/ops/linalg_solve_triangular_native.h,sha256=c3bWQyFFyAFcRjL0bAHoe7jyPLe32k16T_sv4nndMYQ,983 +torch/include/ATen/ops/linalg_solve_triangular_ops.h,sha256=ZAlFA_ivtqSIeDk2RgcPkJE4-fB8kOxhq7MeirZqWUc,2370 +torch/include/ATen/ops/linalg_svd.h,sha256=CidG3-IJrXcLQPgi2-t_YaZ8qQJJ65FROJz5ud3y9HI,2046 +torch/include/ATen/ops/linalg_svd_compositeimplicitautograd_dispatch.h,sha256=_EDo0ux1ze03tcyrc2K8fEbYYvZKmmAv6m2u25M46ms,1582 +torch/include/ATen/ops/linalg_svd_native.h,sha256=mOMa98wfy9Xq4s44Ao-FG_C9htaEjAPMoDd7D5Rgrxk,1061 +torch/include/ATen/ops/linalg_svd_ops.h,sha256=Uv2nw3Rspc3oqMZfM8_FEvwY7PtF6lpTdIxkojuJ8gc,2623 +torch/include/ATen/ops/linalg_svdvals.h,sha256=S0AFo3h6AhdEXCeHdhZ42FxugMY6KBl7jtGn8pZdhfw,1571 +torch/include/ATen/ops/linalg_svdvals_compositeimplicitautograd_dispatch.h,sha256=Mgux0yIJjerTMtzcJ814vyvxymvyeITJ6KTB4u39wVY,1346 +torch/include/ATen/ops/linalg_svdvals_native.h,sha256=AKYe99jOnVZeCZc333JZrT48FFy96wpA569Y_OugxbE,917 +torch/include/ATen/ops/linalg_svdvals_ops.h,sha256=jQlGE87mPqGk-_jVOCaHIbCm58cA46MAKuKJFSzklgk,2108 +torch/include/ATen/ops/linalg_tensorinv.h,sha256=PThoNEJjzHGnjw5I6C-XccrgpMFEELXXGPbwDOCPH20,1472 +torch/include/ATen/ops/linalg_tensorinv_compositeimplicitautograd_dispatch.h,sha256=QhWibJvDLyBfd4qA28NKk5umE-iR792cKFEFtLbrn6E,1248 +torch/include/ATen/ops/linalg_tensorinv_native.h,sha256=oeZOezZ1d9X5m_KlusAuhilxXR0QyljgT2sN1tsoAJ4,856 +torch/include/ATen/ops/linalg_tensorinv_ops.h,sha256=zmvYevJ_1O4pVgbDvnCM7gU-3dbBUOlnAijxWlLPjqU,1953 +torch/include/ATen/ops/linalg_tensorsolve.h,sha256=oRTz-l08ExGr_K57BEa0vHsNeGG-LcFri8Hj_OHCBB0,1734 +torch/include/ATen/ops/linalg_tensorsolve_compositeimplicitautograd_dispatch.h,sha256=6OKTucGsojuwtMlMAIIyPojiA9ECRiKf1-Yi6eMZG0o,1409 +torch/include/ATen/ops/linalg_tensorsolve_native.h,sha256=fk8xoByzegrJCLbVYRSmA3mfi0vPSrcc9WuQ3FhAxVg,959 +torch/include/ATen/ops/linalg_tensorsolve_ops.h,sha256=QdoyVipbqJhvPRpUNf1RbPGYn6qR80TqfVcw_o8vBGM,2251 +torch/include/ATen/ops/linalg_vander.h,sha256=RW-7ePVWoNSJ5PozO8h8iBUp0Y8LPaiPPLEXBI0tlBI,1840 +torch/include/ATen/ops/linalg_vander_compositeimplicitautograd_dispatch.h,sha256=woFTLPNBQmRVvkkeoeTzT3-5uAnhNGkGgdy-UbnuKTY,1177 +torch/include/ATen/ops/linalg_vander_native.h,sha256=atoJuiM2bAJ8fOBHzUmxIhGL-z7WnWXeNzVjS-ubRrU,788 +torch/include/ATen/ops/linalg_vander_ops.h,sha256=Kixnr_Bq9TL7cSfBdmnuFQfn7phMrHXZX5oGrubTKDg,1336 +torch/include/ATen/ops/linalg_vecdot.h,sha256=zkjF6Tieo07wla-4Z1xRyJ6mZBuTjqLgKRFmkqW6WcY,1528 +torch/include/ATen/ops/linalg_vecdot_compositeimplicitautograd_dispatch.h,sha256=4jAn2-6RtA94O64tB91XkksEFKY0L4Ut-SL0EeOFn9g,1298 +torch/include/ATen/ops/linalg_vecdot_native.h,sha256=3fPm8ubqoX-3gBH0o_8YLH47tkC150uPcv64hiB78RQ,889 +torch/include/ATen/ops/linalg_vecdot_ops.h,sha256=egLqi5iLfcW78erbzWsK405bddpzLMmHAFe996emisU,2070 +torch/include/ATen/ops/linalg_vector_norm.h,sha256=Dix3RyQ-1gRkFV2KhNjZKQpzqUaWBWRF-9F-M9ImGxw,2104 +torch/include/ATen/ops/linalg_vector_norm_compositeexplicitautogradnonfunctional_dispatch.h,sha256=o4BycAW0E8tM-TrBSdNsgYWe7wi0Zp-qwORyke0owdM,1200 +torch/include/ATen/ops/linalg_vector_norm_cpu_dispatch.h,sha256=F06vJc6SBYGF3_LbKlbR-U2NZIIrhWwoBx2KTdHNjfQ,1561 +torch/include/ATen/ops/linalg_vector_norm_cuda_dispatch.h,sha256=DeF-8Qh8mf0XKPZaD7CXo0K4HJzibnzPt_69skj-uik,1563 +torch/include/ATen/ops/linalg_vector_norm_meta.h,sha256=ZjZ6-e2CIuMO6PzDhKefYoQi4oP69NVRORi27QRAl5c,939 +torch/include/ATen/ops/linalg_vector_norm_meta_dispatch.h,sha256=CzOlrzt3tB3XD1H8w0ERW5vegGsI-A5KegDkk2A6D8o,1563 +torch/include/ATen/ops/linalg_vector_norm_native.h,sha256=H_gHPDhI8rrJAKsfoQl-7x_Oc__QO3n_rc0VvYmN_ug,992 +torch/include/ATen/ops/linalg_vector_norm_ops.h,sha256=ehT6-r6FW3vu-Aok7YLNDcGxBSwajiTzW-SoCPhT8BA,2620 +torch/include/ATen/ops/linear.h,sha256=01NokxKu32cxtcMGbjYmlFkzme3PWvCnLNZiCoxkyME,1647 +torch/include/ATen/ops/linear_backward.h,sha256=TicG3L8MI7vLcRyjr7x50_Z2T9GyWZG7aI60caiNWOU,2280 +torch/include/ATen/ops/linear_backward_compositeexplicitautograd_dispatch.h,sha256=H-Eu6Ft5dzxlkOvy25Jg540UHnapKK51EqtUQD3xQKc,1483 +torch/include/ATen/ops/linear_backward_native.h,sha256=kxFDMmtqwnTr7iCUkVGK98G_gJ-7VD_Ys_riaMZJqgg,1134 +torch/include/ATen/ops/linear_backward_ops.h,sha256=eUS9SWQ3SjgiDF5IgyxSoQW3_waxtXMQC34muJPsJ58,2876 +torch/include/ATen/ops/linear_compositeexplicitautograd_dispatch.h,sha256=e58pUTv6ng4qvxUfbFHJ8NaeeXXAtlW5mT5PNdAytl4,1264 +torch/include/ATen/ops/linear_compositeimplicitautograd_dispatch.h,sha256=_TA-g4WbvagtJcXO10D0Z4HaEsK--wnpew1n4FWpJ6E,1091 +torch/include/ATen/ops/linear_native.h,sha256=FJdzdwYLF3XyFiWNJUT91UhnXaczVUcKwbczvgLIJu4,1085 +torch/include/ATen/ops/linear_ops.h,sha256=TNKnjv89VAemh4VL4Lm_sV5Wzabc318StGONbhHLqYo,2265 +torch/include/ATen/ops/linspace.h,sha256=BjrAfCrVYYRPXZg54Ax8-FRkqvvyQym0yrmdffZ4KSs,7130 +torch/include/ATen/ops/linspace_compositeexplicitautograd_dispatch.h,sha256=rAyH894irmhtzNcNcWzcjXpmWjEmfwpE8pbsLW55glo,3155 +torch/include/ATen/ops/linspace_cpu_dispatch.h,sha256=ecsxyxSojt-uO_HAdeqgEDN1D_XxXbIV_RXQx5bIdqI,1161 +torch/include/ATen/ops/linspace_cuda_dispatch.h,sha256=iYJWAcuTUZ2FcOPhfak1qXfktf_maQKFdrgW_t77PiA,1163 +torch/include/ATen/ops/linspace_meta_dispatch.h,sha256=Y4-rZ7EqCL6Z8ILzIP3F-q2K5LJbEV7Fw-0ta9VUWgc,1163 +torch/include/ATen/ops/linspace_native.h,sha256=qp1DUTgf4KUxE1P03pUKgY49E-H3g0pAYRrE7gN6Mxc,2293 +torch/include/ATen/ops/linspace_ops.h,sha256=O4MzxVxSQwJcVnzanBn9gj2B05-ObeInQye4wg_HxWM,8416 +torch/include/ATen/ops/log.h,sha256=xOQaF50tq9lcrCmH4Iezr9Kr3JAb97F5_U0SIpJQHqM,1384 +torch/include/ATen/ops/log10.h,sha256=xryDeKUv0QNiJSy_cJtbl_0FYskLAv975IScBqGoCoI,1410 +torch/include/ATen/ops/log10_compositeexplicitautogradnonfunctional_dispatch.h,sha256=qkkqkBU48LN26KruF22zWeghO3Rogaua_9ULnXB0s9A,1093 +torch/include/ATen/ops/log10_cpu_dispatch.h,sha256=mkNZWZsAT1kCIi0MkXbsZ2Q1XU7mR9VXscNgA6lYDoo,1178 +torch/include/ATen/ops/log10_cuda_dispatch.h,sha256=PAdt6UUCFcqogXnE6i05JLGH5xMpIbF0MnYe59fo1HI,1180 +torch/include/ATen/ops/log10_meta.h,sha256=4LIY7g43NNA1e2WfwVZH7UJ8TabNy_UIAm4BH4ZYCws,820 +torch/include/ATen/ops/log10_meta_dispatch.h,sha256=NlEefBvTdHBwyK-oUARwFZ7QVd0ysRwTuP9Gn4k-oro,1180 +torch/include/ATen/ops/log10_native.h,sha256=MJKOcod__cxaAmd2JpyrzXNwu5kzI9yH2X-tnOSqUOA,847 +torch/include/ATen/ops/log10_ops.h,sha256=BUejzlx8zbqNhCQvWYsxH4wydyj7JUcUqyzpyKpTzt0,2291 +torch/include/ATen/ops/log1p.h,sha256=OG2kPlZlVEhZHiE_fN4S2tONFJyB0KuRATbYN-gCcgA,1410 +torch/include/ATen/ops/log1p_compositeexplicitautogradnonfunctional_dispatch.h,sha256=aa38aHhxPV2nntvQV92r3FkRXc-UJbjGJXWxx-ncSxk,1093 +torch/include/ATen/ops/log1p_cpu_dispatch.h,sha256=MjaXv2WalOgfiP4QungEnD_63_945CZQsCp5sjPDbSI,1178 +torch/include/ATen/ops/log1p_cuda_dispatch.h,sha256=d1qSIy_FVDhYK8kOTuzg2laXOuVTh_JCTi71SiqEX0I,1180 +torch/include/ATen/ops/log1p_meta.h,sha256=vviGLx641ST0YzpoIDcYsRiiEu0JX3jx7vRKgXbssBU,820 +torch/include/ATen/ops/log1p_meta_dispatch.h,sha256=17Sj8Bc8Xv4VvZ2ILmhKidfrm7IpYPqWZFA6yMZHEoQ,1180 +torch/include/ATen/ops/log1p_native.h,sha256=dDrZEZpCtZCwGng_GCKs46VdmfAZE_vH8FcmNhy1f_w,1261 +torch/include/ATen/ops/log1p_ops.h,sha256=kfZm2qspmY-5YLLxLZMXxHovyIVvq8eSAevcChMi6xo,2291 +torch/include/ATen/ops/log2.h,sha256=ZsUCAC7Yg2u6glmX7m-BJ-X5J72XX0tGnAFnFB5zJgc,1397 +torch/include/ATen/ops/log2_compositeexplicitautogradnonfunctional_dispatch.h,sha256=nK5nvLJ0Gry9P6VWStGGRcWQquvrIxlqG7hpj41L748,1091 +torch/include/ATen/ops/log2_cpu_dispatch.h,sha256=uQof6NP-SS_yRQwZBqQc6Z6UR4EV8iILc7G-dBOkFHc,1174 +torch/include/ATen/ops/log2_cuda_dispatch.h,sha256=OUyzi3_I4ZdjFuzXGreORuwih3OXy8efssRn1-Uz5c0,1176 +torch/include/ATen/ops/log2_meta.h,sha256=NwOlwUGSMgyyR_fbcF5lNStLUUycMnTnZy4_r5fwclk,819 +torch/include/ATen/ops/log2_meta_dispatch.h,sha256=-mn18o1vbJF2mWu1AnDBldzC9k7q38QwsdaVEGAus0E,1176 +torch/include/ATen/ops/log2_native.h,sha256=QYS0_A6U9hHTAOB5sZpONdQzQRnEh4QhA5yC5eRZQw8,844 +torch/include/ATen/ops/log2_ops.h,sha256=RN65vluzw9xI2MFtIosJHmPDxAOs-k-p3LvyR4nMUFk,2282 +torch/include/ATen/ops/log_compositeexplicitautogradnonfunctional_dispatch.h,sha256=zawq8hY1L3QuF7G4KL7SXA88VGWA-7fx8z1_fy8MJ-w,1089 +torch/include/ATen/ops/log_cpu_dispatch.h,sha256=fnLB3WJHkN8-27Ca6-NDf3CFqI7mL2-Es0WOY_B0E3M,1170 +torch/include/ATen/ops/log_cuda_dispatch.h,sha256=ZUC-UzXJziTuBICEueFvHfaud3nN5WzPY_PAuDCRiV0,1172 +torch/include/ATen/ops/log_meta.h,sha256=c-pI2Ni5Zc2twOwAo8csq84GSdaJ0QxQ-9aIb86hGls,818 +torch/include/ATen/ops/log_meta_dispatch.h,sha256=qmMW14zAGzjcQb750Rji2dy7C1lEdGE-yC3JxXWib7U,1172 +torch/include/ATen/ops/log_native.h,sha256=CkQiyOgWfrNIdEwMOPWwoYm3FYOLbHklSyd73N3qF_A,841 +torch/include/ATen/ops/log_normal.h,sha256=7MYdqRD2ks9mxwuCzlEGZHbb4TZvmx6ngIw1v3S_vzY,1791 +torch/include/ATen/ops/log_normal_compositeexplicitautograd_dispatch.h,sha256=xcfQNYzJlZkUGfGXqQyhsenc3ACPBwi8SCT98vYLZeQ,1426 +torch/include/ATen/ops/log_normal_cpu_dispatch.h,sha256=HqJnk6grZmcU7YL1cuiQctsgy6ks7N1pj32YjEyrpVM,1061 +torch/include/ATen/ops/log_normal_cuda_dispatch.h,sha256=wylWaji9QReD9PbjajHB-yICpOFKObmNcsEoSCv0Alo,1063 +torch/include/ATen/ops/log_normal_meta_dispatch.h,sha256=EC3fxPiksSCkiUd9o09GWnGOMFRAZTwpcsIL2aHc8Vs,1063 +torch/include/ATen/ops/log_normal_native.h,sha256=ahg-_k9aYRzJ4vd_4zslgmZTMUPmgAjKjtARjxU1Wkc,1110 +torch/include/ATen/ops/log_normal_ops.h,sha256=IpzBUdapLU0rOHJbdYaboTZ6ZptPpQSRNBouuDic_0k,3050 +torch/include/ATen/ops/log_ops.h,sha256=uaEkHCLFuu4ioZqF7PTlP3aZY-OBRq6GZrOAHj1iBls,2273 +torch/include/ATen/ops/log_sigmoid.h,sha256=tVpIKd73NIWqD4WFU2SFd4lp6GmUMdpjcloSniL9bTg,1331 +torch/include/ATen/ops/log_sigmoid_backward.h,sha256=5AxC2EmORABYo5zlXFO8XExJiJ2iMD4W6M5MnFmFZcE,1836 +torch/include/ATen/ops/log_sigmoid_backward_cpu_dispatch.h,sha256=yq_N4Kh9f-Hh-ZyOazRlSBooIh9dAc83fZp4YUb_Mt8,1364 +torch/include/ATen/ops/log_sigmoid_backward_cuda_dispatch.h,sha256=g_us_wU6_cseacHJ3sFkitMbjgNYRyzjqbJUHy770Ko,1366 +torch/include/ATen/ops/log_sigmoid_backward_native.h,sha256=OnL7aSMp4vlMAOxRo-Kfru4ztbjVHYpudt5RKTxoMLc,1264 +torch/include/ATen/ops/log_sigmoid_backward_ops.h,sha256=AO9wm-Ebyj9nyqliZHIDO48FF5qU3aZzZuC7wgKIUOA,2313 +torch/include/ATen/ops/log_sigmoid_compositeimplicitautograd_dispatch.h,sha256=2-zUIbNN1BEWZOB82kS6eFkhkWCKK1zMsOxtLBEgqKA,1190 +torch/include/ATen/ops/log_sigmoid_forward.h,sha256=_gnULzh92JXHEARxbVfmwtTvdi9KWKtG01gbUDEmz6A,1668 +torch/include/ATen/ops/log_sigmoid_forward_cpu_dispatch.h,sha256=yv12S1BJqsaW5Dk9VGbQhWUutgpV6A-LMxInOvBSnfQ,1297 +torch/include/ATen/ops/log_sigmoid_forward_cuda_dispatch.h,sha256=WdjG0Y_VPCkQD58wXcr4uOFXHGzykjdRc6iHMw7nIAI,1299 +torch/include/ATen/ops/log_sigmoid_forward_native.h,sha256=uZs8jJh4isJq70gcNjOoF52L0CHoF4xAb40dXIRZJSM,1162 +torch/include/ATen/ops/log_sigmoid_forward_ops.h,sha256=K_XnvihHFQ-aOsc4nR6PXMeFVvHs93RIRF4UNskHES8,2166 +torch/include/ATen/ops/log_sigmoid_native.h,sha256=HV3C-nCOCAj37EZTFWn72fXz6cSOVUC_rWcjudv3FVM,818 +torch/include/ATen/ops/log_sigmoid_ops.h,sha256=mRyB4WkSGivY20MVGHgPh4aAULcAKFlrvcVPtTCpJ3s,1831 +torch/include/ATen/ops/log_softmax.h,sha256=NWcOTaYoowJWETz-OjaMfoDf1fJfxRqY2ZfJVvc6SMs,1965 +torch/include/ATen/ops/log_softmax_compositeexplicitautograd_dispatch.h,sha256=rLrrXZUUVt1mWYuls5uw_Z9RDm0u3cD2lIXzby1mU0Y,1250 +torch/include/ATen/ops/log_softmax_compositeimplicitautograd_dispatch.h,sha256=fhZiHcucjudl__eb151V74yfgbZL8x_XWpWZHSzNg44,1220 +torch/include/ATen/ops/log_softmax_native.h,sha256=QuHrs3QaCyT7CiKBjlRaCVbj3X0UpbjVzxnqtqIPIBU,1067 +torch/include/ATen/ops/log_softmax_ops.h,sha256=6kv1Hs7kekXNsBQCZVHXMAoh4_CKFguzfAd9LJgBQ0I,2948 +torch/include/ATen/ops/logaddexp.h,sha256=Y23kfc0RhyR8cai_JsadKidS5bp6hAm3_Ac_9zDPkvo,1452 +torch/include/ATen/ops/logaddexp2.h,sha256=KjwZji3NwDB2tR3W0Kaxe0frkMQ-TPTFmVLPNxXxNzI,1462 +torch/include/ATen/ops/logaddexp2_compositeexplicitautogradnonfunctional_dispatch.h,sha256=EaLUi41_axPiHd000OopXD6iYHkWUDIcYHKxhYj7fGY,1074 +torch/include/ATen/ops/logaddexp2_cpu_dispatch.h,sha256=G04mFynA5zME8xlybq10pYzb2jnrNV2ornDT6rgcBpc,1221 +torch/include/ATen/ops/logaddexp2_cuda_dispatch.h,sha256=3UdZTRjoamcAfQb0-8_b7zNO1GSqoOL8gFe_Np5u1cY,1223 +torch/include/ATen/ops/logaddexp2_meta.h,sha256=i2HyOPegYvVizWcEDCTt_u9n0eg_rvW5R31cfCRK1Ks,851 +torch/include/ATen/ops/logaddexp2_meta_dispatch.h,sha256=Ib0YXv2uCXdBVI1s6PLOg3VgS5C-uuy0pDtawjY-yds,1223 +torch/include/ATen/ops/logaddexp2_native.h,sha256=cxMMc0-hNa2RDQaUvBnLqk-Xs4xFsPK0HMtVcy6uulI,888 +torch/include/ATen/ops/logaddexp2_ops.h,sha256=r4VqUdZSZOVjfnVE1AKMNU0QUuL4s1Uhm9_o6Fbx_SM,1997 +torch/include/ATen/ops/logaddexp_compositeexplicitautogradnonfunctional_dispatch.h,sha256=vUOyXE6CcAB_xCQdVzIED5KKCBNL0Ju9C9dRUQ04Aj8,1073 +torch/include/ATen/ops/logaddexp_cpu_dispatch.h,sha256=Eg5PnJ4Aaa9vXQIAJ0o5SuCUmIDbiEUQJ_3IIqSYdok,1218 +torch/include/ATen/ops/logaddexp_cuda_dispatch.h,sha256=M0YdN8nDynFuhGoSEXkCmJy44IaETALR6PcieefLfDY,1220 +torch/include/ATen/ops/logaddexp_meta.h,sha256=reW3tJD45oPKBCDql38j3TsNBoAo-ffoHk1Oa8dMwD8,850 +torch/include/ATen/ops/logaddexp_meta_dispatch.h,sha256=dryZqZMD_SzE3uMF8aejRvZolCN0W7T9CrdIphHyadM,1220 +torch/include/ATen/ops/logaddexp_native.h,sha256=BbCD3W3B0-bQG-SRCQtmublLZnCwtSkSuMFUrOshESk,885 +torch/include/ATen/ops/logaddexp_ops.h,sha256=Dc2Gy_EAUGzTV0QYwxTE9-x1Qa_B-RXed7O-zT27Hws,1991 +torch/include/ATen/ops/logcumsumexp.h,sha256=DwjeITmT-NmoRrI7x6Hx3nYAkm49afn9mNRSaDJh0nU,2154 +torch/include/ATen/ops/logcumsumexp_compositeexplicitautograd_dispatch.h,sha256=ImURomyp9ubZsRe0xarU7Gw6y4qymFxtmJ6Of3iHImU,1232 +torch/include/ATen/ops/logcumsumexp_compositeimplicitautograd_dispatch.h,sha256=Mzx3lbL33xZ0DdLmDzpxAk10e_niKn-wKhbkimwaYkI,1244 +torch/include/ATen/ops/logcumsumexp_native.h,sha256=kAl-xZn1GF3s6Ezq5eQzFKCQ-55OYphU1BRACB9ewfw,1024 +torch/include/ATen/ops/logcumsumexp_ops.h,sha256=DKCI2I09a7wBdnyXE7KvZKkIyzY9eB5bchFI1jMbUCM,3218 +torch/include/ATen/ops/logdet.h,sha256=diUYhDXqTDjqOEgB5X892AXTqPf9sVz2g-nIrpv5XIU,891 +torch/include/ATen/ops/logdet_compositeimplicitautograd_dispatch.h,sha256=-XAH5Ayc1S1rHJwzB2G3D-jZ430HoRTbCkL1n6fHGVE,1018 +torch/include/ATen/ops/logdet_native.h,sha256=OvgfbgSbGdUeBN02pu35DFQAsvrLIaFhm2GdUiSm5jg,730 +torch/include/ATen/ops/logdet_ops.h,sha256=D_7lXSDkbmrfaltIFapPGWqR-hui2jUHl7GsoUylAB8,1211 +torch/include/ATen/ops/logical_and.h,sha256=oE5sRLTXqACgwoWBbxkxsfyeHnJXAxoUsDJT_kqtOQk,1472 +torch/include/ATen/ops/logical_and_compositeexplicitautograd_dispatch.h,sha256=-ap6pSytHxAZOfWT-XzvpJpN2wOUIrdGEjltZ1nqLqM,1131 +torch/include/ATen/ops/logical_and_cpu_dispatch.h,sha256=47tcL1vx1A0FAkwssZSj_R3tAHBAO4PfIuGryWLnjug,1139 +torch/include/ATen/ops/logical_and_cuda_dispatch.h,sha256=d1E6se5-fvbfxyIMhs3xKbBiIjMEIXRB3NhH25ROaec,1141 +torch/include/ATen/ops/logical_and_native.h,sha256=QSE7EjoZ0ExfYN_1clQn6daqYwZcK_Vh5xbmV2jC_mI,952 +torch/include/ATen/ops/logical_and_ops.h,sha256=VwuoqFUW0FunjlWd6Eoww4o8o5Un6tM4Aqzxo-kQ7HE,2603 +torch/include/ATen/ops/logical_not.h,sha256=GBdu7OodZwXyH63pgarjjeYyh8WtMydf48-tECE3Yv0,1331 +torch/include/ATen/ops/logical_not_compositeexplicitautograd_dispatch.h,sha256=i66uZ3cf3D8aCC-ZTD5c6F3k5DvSqxDfULKqBF8r3P0,1079 +torch/include/ATen/ops/logical_not_cpu_dispatch.h,sha256=8Zlgk1tND-AtMAzbYt7VXkp0jRFdG7J3C447PSWql_g,1087 +torch/include/ATen/ops/logical_not_cuda_dispatch.h,sha256=vO8Bse_bBA6rGfIXQQ2BG4iQrWaqO7JoJfDhQf8o19s,1089 +torch/include/ATen/ops/logical_not_native.h,sha256=rd3DYC22BGoowBmbtHiXAgDGQuVNXK7lgIhxhm8gAjE,1015 +torch/include/ATen/ops/logical_not_ops.h,sha256=AkLvfs8EVfmrnUfoDkQ82GIgSjOPBVjuB4rhcx4hNp4,2345 +torch/include/ATen/ops/logical_or.h,sha256=0Ho-YmstSA4Q0GtdqL-BNbiYfuHgZorBkvEk6WIW9UY,1462 +torch/include/ATen/ops/logical_or_compositeexplicitautograd_dispatch.h,sha256=3nWr8pcF_hPjd1-n65-AG_j7_NqIepZl8hWmyvQ01G8,1129 +torch/include/ATen/ops/logical_or_cpu_dispatch.h,sha256=3XstS9SpRsKJuK0LqzoBTAEMQyd5gBT1eo-1wjdm_rU,1137 +torch/include/ATen/ops/logical_or_cuda_dispatch.h,sha256=SZsKMRpZoFJcq2GZlRLaxy1k2Llvwpt608FA57QIgIU,1139 +torch/include/ATen/ops/logical_or_native.h,sha256=8WfC05fxpX2EuVNBBBpWlJ0j-wymLxTpxC1JweYG5S8,949 +torch/include/ATen/ops/logical_or_ops.h,sha256=ioTtybEg5NkT22K0DzLaJYMNV64VTcgjimW-0mgGv-I,2594 +torch/include/ATen/ops/logical_xor.h,sha256=CXdQxCrTEMnxUlE8RA93715We4Id-9NZ24B4wp8JO1U,1472 +torch/include/ATen/ops/logical_xor_compositeexplicitautograd_dispatch.h,sha256=GUv4fv2Xnt1Z6SbK5pgZRSQzwETK-bR34MCkSPZP2rk,1131 +torch/include/ATen/ops/logical_xor_cpu_dispatch.h,sha256=X6BtNUAyK-m6GOhFzjiwiUicHQnzpTmuxN--seWHKrs,1139 +torch/include/ATen/ops/logical_xor_cuda_dispatch.h,sha256=vZbd3ADpAYjcj4HP3ZFdy6Yc3ESWf0H0QxM5jrM6xPA,1141 +torch/include/ATen/ops/logical_xor_native.h,sha256=1uvTjT_q0DOFeeqR-AgOGfBVjV5nBipiV8ypBUxiZiU,952 +torch/include/ATen/ops/logical_xor_ops.h,sha256=Qqypf8jFj67tubYNZAua1wn6oW7lZohGwZ0vMwev3OM,2603 +torch/include/ATen/ops/logit.h,sha256=Ya3qHmsYy_OX5twva9sTNpFcrrGGMZpTsXg_tVgsTTs,1659 +torch/include/ATen/ops/logit_backward.h,sha256=NMoU12-Tw-XziicLqWfJ3F4wwp_1uG9PdJ8fYBOMJQ8,1809 +torch/include/ATen/ops/logit_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=S0FI3VQ1qWU1XvWu5HiL9nMRHwo1hTk-ZgiybXmHMG0,1128 +torch/include/ATen/ops/logit_backward_cpu_dispatch.h,sha256=joDDCtCW-Cok28YwckZx_bfxtFO3fNc0jd3yfc_Hhig,1382 +torch/include/ATen/ops/logit_backward_cuda_dispatch.h,sha256=VHxTok09L8JtelQHq4yY0hY7n6u4Iw20Vz-1koVJxPw,1384 +torch/include/ATen/ops/logit_backward_meta.h,sha256=xmlACX-lpM7_I6fHbqEDVWeqXvFilAzR7t5p4jvl2C8,890 +torch/include/ATen/ops/logit_backward_meta_dispatch.h,sha256=0mhXLbSgSq4Iq3lFzrmVXM4gEs91MnP-OzUDp5z6YTc,1384 +torch/include/ATen/ops/logit_backward_native.h,sha256=9ncuT0W0PIMVnFwCy2QyHEmHP5APvU-zhJaUGM3nJ7A,942 +torch/include/ATen/ops/logit_backward_ops.h,sha256=E4ibf5FOu3JchDuBiWuxUXNcVzimjiE18TX4XY-DjB8,2299 +torch/include/ATen/ops/logit_cpu_dispatch.h,sha256=Px9fLWzg908fJHurvJmlSZaPTnOAPdogryOOMp3qcms,1339 +torch/include/ATen/ops/logit_cuda_dispatch.h,sha256=PiDG8Dj7lASH51KKuOuamr94-POT4J742UTbZbV3JjA,1341 +torch/include/ATen/ops/logit_meta_dispatch.h,sha256=DW5CU2-g0nAS2sRx9asRkViHLKGu15BTV5ZEY64VJkI,1016 +torch/include/ATen/ops/logit_native.h,sha256=lAsyt4rp34A4QvW9JwSrbdMIqQG-CXc_JZBGs8Mw4HE,973 +torch/include/ATen/ops/logit_ops.h,sha256=Rlz34du2wCxNHVXgMNomT6jj2860ik6yoLsdhXqo_Bs,2591 +torch/include/ATen/ops/logspace.h,sha256=i4wqGhOhK8bepCr1uErG8T_2cWMSyc4BBPnxvnWMfPs,7746 +torch/include/ATen/ops/logspace_compositeexplicitautograd_dispatch.h,sha256=r9cf7Zmdsdzy61-mo9qKQExkWo1qXE95U9JHKbwkhQw,3372 +torch/include/ATen/ops/logspace_cpu_dispatch.h,sha256=eg58ehs1Da6YoWoDN37ByZnhKP7n92STIoXJG-BMz08,1192 +torch/include/ATen/ops/logspace_cuda_dispatch.h,sha256=W26u2y7qcGAajB_JCKaoRD-IFKVCagzVcIs1ihQ7xEE,1194 +torch/include/ATen/ops/logspace_meta_dispatch.h,sha256=p3ISKkvaODPme-4JiaEGWNZI1nm9caGnTX3ZqPMSq34,1194 +torch/include/ATen/ops/logspace_native.h,sha256=Ie7eZy96iyE7boinKgMjfKJriZcTvmtMh2-dPlZRunQ,2430 +torch/include/ATen/ops/logspace_ops.h,sha256=K15L8st9R2uRhiIr2z305XoOby4haqNsc0ukrUlFtgM,8824 +torch/include/ATen/ops/logsumexp.h,sha256=bWYKMa8IY0X2mJBRi4tK03K6DXIKSwUHTlG49MveqUo,2421 +torch/include/ATen/ops/logsumexp_compositeexplicitautograd_dispatch.h,sha256=fdxdeJb5GBSk5RzDEeZlA8B9BA7Mk1u_PXFFYn2JkcU,1062 +torch/include/ATen/ops/logsumexp_compositeexplicitautogradnonfunctional_dispatch.h,sha256=0dBxOOGT9MZ4xPbeCggbXbKzIuMAqbXfXH5bRQKp7VU,1229 +torch/include/ATen/ops/logsumexp_compositeimplicitautograd_dispatch.h,sha256=62Dl-pwDNMJAU6qoqIqiQTqAFslFiQgZO-xqKr2tF6w,1301 +torch/include/ATen/ops/logsumexp_native.h,sha256=YiPAf73pqPllJJnoEoGc5Zdrxo15DIZtzduIgaXKlsU,1104 +torch/include/ATen/ops/logsumexp_ops.h,sha256=CFv6i2aJqSwoN9ztLh30xiyiGwwbx4YdcA06ZPTomkQ,3470 +torch/include/ATen/ops/lshift.h,sha256=RDWWXhX-4pvyqB6tnF1e91V4bNd1Acqkp3ZM-Z1Kvms,2244 +torch/include/ATen/ops/lshift_compositeexplicitautograd_dispatch.h,sha256=xfEwuM5w3MrqHP-24tJvDwgm_32HoBimYDy8kVfYBuI,1398 +torch/include/ATen/ops/lshift_cpu_dispatch.h,sha256=AJdR9AJbOqOJi7XPdvfUzPDGl4-LMFZ-CmpWWRZUOYo,1250 +torch/include/ATen/ops/lshift_cuda_dispatch.h,sha256=_VtjUtAbtfQOwsKToJB9Z6yBBzJ3zGCkz8eqwbbf1Kc,1252 +torch/include/ATen/ops/lshift_meta_dispatch.h,sha256=fHcInyTf_S7qmmTDYen43KQBkhrGoJvJNOdKZ_nRe3U,1084 +torch/include/ATen/ops/lshift_native.h,sha256=GnfuVfDpmlVLB_GdXO_5X8PvnrT6KmpBqijJjfvV8Gk,1236 +torch/include/ATen/ops/lshift_ops.h,sha256=ExSeUk_d7EoqKRPjrPDsz7A6W3C6czi5vXLtMDfI1Gs,4599 +torch/include/ATen/ops/lstm.h,sha256=npSUDH9EE88Tb45L0k1qjAa5NHFfH2RSuiFF5pRVECQ,1867 +torch/include/ATen/ops/lstm_cell.h,sha256=9ChhGCyLelrwsuysnDQRXfMa7c2-ccAcpkqbH7I1EII,1205 +torch/include/ATen/ops/lstm_cell_compositeimplicitautograd_dispatch.h,sha256=x9Jum4tHatBrTYcpOjvKcPD07wKG78jXG-pbem9QT3U,1206 +torch/include/ATen/ops/lstm_cell_native.h,sha256=qZsF9RM6SH_V3-nNpOjCBtfHiXiz1qm2fb-tY0QGuqU,918 +torch/include/ATen/ops/lstm_cell_ops.h,sha256=Ri-TRhMoB0bVDdDJeBR2uvzJrWUJ4CvKSzCuk0kcF4c,1821 +torch/include/ATen/ops/lstm_compositeimplicitautograd_dispatch.h,sha256=oo1bsfEuBqdy5Xdw8mf-WqpkZaGO5tuO4RcjXfqDoro,1445 +torch/include/ATen/ops/lstm_mps_backward.h,sha256=nq72DjKGxZ-m7Uscrn_C1Eo9uLWp3zFiVYVQSlr6g94,3929 +torch/include/ATen/ops/lstm_mps_backward_compositeexplicitautograd_dispatch.h,sha256=S7mc68XXvZuqszph4aR_aVKPXXqxpPDuJuzvj3ioKfM,1967 +torch/include/ATen/ops/lstm_mps_backward_native.h,sha256=Wlag0V4q0C_u6eEwlKToBo6HVDJrBQo3GEkGIOAiiEk,1177 +torch/include/ATen/ops/lstm_mps_backward_ops.h,sha256=MEalpRoQvChUtS9wdMlx5D4umvciq4Rm6QsfcB_2eBs,4688 +torch/include/ATen/ops/lstm_native.h,sha256=JJ1efTBefLuAICX2t0CqMOF-UyKU8lRaiW2PD65yGKc,1157 +torch/include/ATen/ops/lstm_ops.h,sha256=9cUtntGblA0Lfk6XMkp6adhPD5XqFI7kRX3gNPWg1kY,3011 +torch/include/ATen/ops/lt.h,sha256=pJM_QUThaUSVrMwYAiWvP-hvpCZwBwi0_o15X4OLdFc,2096 +torch/include/ATen/ops/lt_compositeexplicitautogradnonfunctional_dispatch.h,sha256=0ql3-9FWLwNopQf502-k391RocBo62upxdmYTh1lrM0,1288 +torch/include/ATen/ops/lt_cpu_dispatch.h,sha256=Gkayne__OEdU-jOP2_Qj0HzgNPwfkK6Hm_QDk_sFryQ,1620 +torch/include/ATen/ops/lt_cuda_dispatch.h,sha256=JVjXzeuGAJnRxAxnojF2XRg3kXKWSCGo65cHHsec8vU,1622 +torch/include/ATen/ops/lt_meta.h,sha256=tLPOYvSg5YXSet2D-MKkmqkVRIgV3jnVoF6NDcJISfE,989 +torch/include/ATen/ops/lt_meta_dispatch.h,sha256=P0d_sAgOfKDAger7m_1dYvw5-BlljzGpqOkZGPoA-Fg,1622 +torch/include/ATen/ops/lt_native.h,sha256=u8FZvO_Oz7sj3-SxkV1H5_ggoKb1-XjwfNOG2qz78WQ,1459 +torch/include/ATen/ops/lt_ops.h,sha256=QpVL1mJUMd3HclW1dzIE3o_MLOBJzyUdgPjlS4WVw8k,4455 +torch/include/ATen/ops/lu_solve.h,sha256=gUMqtLx25YQdMV_TYceuOI31LZ3X08l5urnieFbckj4,1637 +torch/include/ATen/ops/lu_solve_compositeimplicitautograd_dispatch.h,sha256=NiRiMBRYPIlFRtu2iPtjqUPECt6qMQjynjIktLecwKk,1355 +torch/include/ATen/ops/lu_solve_native.h,sha256=HEthPORZztgozQYyxQQgYmvMgFYM05vnxx2FR6AiRTM,928 +torch/include/ATen/ops/lu_solve_ops.h,sha256=5CiFj94sZeb7yyTtRGTdmuH6bDwROReuAvqbCGQHtso,2193 +torch/include/ATen/ops/lu_unpack.h,sha256=sf9UQ9XXv6OECLmAZCdANrSwCpER9AU3rapkKfyqYB0,2201 +torch/include/ATen/ops/lu_unpack_compositeexplicitautogradnonfunctional_dispatch.h,sha256=c0m1UHGkOIpgqENEwrJ2DAPySp0aEcP1_9sdtV2mvOY,1164 +torch/include/ATen/ops/lu_unpack_cpu_dispatch.h,sha256=UjFauGJ1LEh-1DUzBCBho03FN9XZM4uPL9QziahKmwA,1549 +torch/include/ATen/ops/lu_unpack_cuda_dispatch.h,sha256=-4hA4j37aCfZWWdPkZwNhxyEt-8y05bXsSicyeZ88vw,1551 +torch/include/ATen/ops/lu_unpack_meta.h,sha256=lW8JM5YQAkeZvANUEadeJJ-oVRk2_FoSamer5Xj3Ab8,895 +torch/include/ATen/ops/lu_unpack_meta_dispatch.h,sha256=nlVYF2pvH3HaLoImdzGdH3Wp_BhHs2nUlJ_DeBq2N3s,1551 +torch/include/ATen/ops/lu_unpack_native.h,sha256=Cul8PAQB06yLf4kK4QbpGsOLCsCKfwUcWqH8EKCA4oI,972 +torch/include/ATen/ops/lu_unpack_ops.h,sha256=YKpNT47IfgZneVIfWkuzBB4qH259_8tguWBFJtIyIu8,2703 +torch/include/ATen/ops/mH.h,sha256=pbNrcfKRz-DtptslD6_oOhJh2dQ5_zJP9f1TxcbtwcE,753 +torch/include/ATen/ops/mH_compositeimplicitautograd_dispatch.h,sha256=5c7p1m62IDlFqSdC-EL-1DJYg3QO6d6mRKcOdWGwcYs,1014 +torch/include/ATen/ops/mH_native.h,sha256=grLNKNYYBR6simKgZ_OgIR03IycJfmB05YYQRfOpKaA,726 +torch/include/ATen/ops/mH_ops.h,sha256=DzFMYyTHKnfMWAoprPBlwTkTSap-CNMmwA-a7bXPuoU,1205 +torch/include/ATen/ops/mT.h,sha256=asgp9Kiafnhz0RwllEkh2US0r-y7kxtQfR8WNjcG0-w,753 +torch/include/ATen/ops/mT_compositeimplicitautograd_dispatch.h,sha256=0pOR0xayTRUR4ah9abHsQ9E8-P71Udq0ZjMhueyKPZ8,1014 +torch/include/ATen/ops/mT_native.h,sha256=G62vhGJK_48ZVkW1uArLuBD2IjsdzUIuZS_V342y1K0,726 +torch/include/ATen/ops/mT_ops.h,sha256=G9gruuYoK9zVidxmZIktforSmwMJGdgTkacUUf22M9g,1205 +torch/include/ATen/ops/margin_ranking_loss.h,sha256=PFLMchhpdbTzDMXkLVD8oi8tyGc-WaSOd4YXxn5TwVg,1164 +torch/include/ATen/ops/margin_ranking_loss_compositeimplicitautograd_dispatch.h,sha256=_-8nLE6Hx_cIGR1rCD67voOMYmep-s3Nr9V6g6f3hsk,1145 +torch/include/ATen/ops/margin_ranking_loss_native.h,sha256=d1WrPBpG_A0Lj9uuzEGjIx6FYHOffTfNkms1kHZp32k,857 +torch/include/ATen/ops/margin_ranking_loss_ops.h,sha256=gK-5f3inySwwtdGgjxiFf5KrZfKww-BfCqVt73t2vcA,1557 +torch/include/ATen/ops/masked_fill.h,sha256=orT5D32DYTLYqOPvh037zzMDUwLj0QvPNQLx_KQAjaQ,2531 +torch/include/ATen/ops/masked_fill_compositeexplicitautograd_dispatch.h,sha256=JOPu-_EcrDfa4WROicheidtzasK15rKeEHGMi8RxJZ0,1722 +torch/include/ATen/ops/masked_fill_cpu_dispatch.h,sha256=PiQxtYy71JJWagCkfWRI-DuAjCjI0ezG4wvU_hxfk2k,1134 +torch/include/ATen/ops/masked_fill_cuda_dispatch.h,sha256=ZuhxwuLPH46NygQwFnRa6_EBCoVOJJZ3-mn5atykTeo,1136 +torch/include/ATen/ops/masked_fill_meta_dispatch.h,sha256=o9tTd4hsl8qOKlA5aaWUUGeIpCaR2y6QMZynrKyCdjU,1136 +torch/include/ATen/ops/masked_fill_native.h,sha256=gGVcfz43gSuLsAFr5mNcjtlAGQh-7K7lrE1YOHxg1VQ,2233 +torch/include/ATen/ops/masked_fill_ops.h,sha256=8iEbhJTqeLCTswUGNpFB7HJePwdTYMip9xsXTOYikoo,5115 +torch/include/ATen/ops/masked_scatter.h,sha256=a9U2m2OIIhvA3gZTrUHWzMKyDMKp1OushwVL6-Y_KyE,1643 +torch/include/ATen/ops/masked_scatter_backward.h,sha256=7sN_0PMkcnXlQFlY16zKzCPa-g7ZvQ3_B8W7flMV22c,2044 +torch/include/ATen/ops/masked_scatter_backward_compositeexplicitautograd_dispatch.h,sha256=yLAVLFUAzozW97x7RpoAiD-SUxj20A4D-IaM5mS4uh0,1227 +torch/include/ATen/ops/masked_scatter_backward_native.h,sha256=PM8ZAvzFXjEuS7ZPjP4NuzEd7K5fbUkfRpNKmX_69Ws,813 +torch/include/ATen/ops/masked_scatter_backward_ops.h,sha256=ZFIWw7eIDsz7EHFgkXXW5BbQPVYnb0oPGuJzyKJ8l4o,1457 +torch/include/ATen/ops/masked_scatter_compositeexplicitautograd_dispatch.h,sha256=Fs9RcwUfI7nij7Yl6TPtXjsaIfBppxCDpQI0whYygs4,1355 +torch/include/ATen/ops/masked_scatter_cpu_dispatch.h,sha256=74enrbznF6YmkbMJc3uVuEMOgq7hrRMDufsQQs-B89E,1031 +torch/include/ATen/ops/masked_scatter_cuda_dispatch.h,sha256=ywRIoNAQ6WHmUm5XdHikKKGQgk50ooaJq5nORaqO188,1033 +torch/include/ATen/ops/masked_scatter_meta_dispatch.h,sha256=eglldQYSep8hKmPwy3r49fS_bTzgqbot1f-jS1AIqUI,1033 +torch/include/ATen/ops/masked_scatter_native.h,sha256=RD0R9pqIpzfNY_IUPu4ilMtIvxuNLpn9gJgAzkYqC9s,1159 +torch/include/ATen/ops/masked_scatter_ops.h,sha256=zR_uyGmrNBoclDVhOOczAqkmvSyunWaKYsx4MA_jx5M,2888 +torch/include/ATen/ops/masked_select.h,sha256=35nj-D4EHuABztDY3A2jq0meqZxh6BROXY1_Yadc5oE,1483 +torch/include/ATen/ops/masked_select_backward.h,sha256=Rn1ubEOG7XpYhgA9GGVCk31RJA2bAIEX4jIvHD-mPpM,1046 +torch/include/ATen/ops/masked_select_backward_compositeimplicitautograd_dispatch.h,sha256=kAYmFlfvkQC3w7byIJXLuKYayV-q2a0yorrMX6JJers,1085 +torch/include/ATen/ops/masked_select_backward_native.h,sha256=yJnZ8KsWdHwn9TG_cM7Tzcd-1-zniIKKj5SubCraEc4,797 +torch/include/ATen/ops/masked_select_backward_ops.h,sha256=SqZLGB0h1AXNltai_muPXsVmnDJ4rUP2jHgtVn43Ikw,1428 +torch/include/ATen/ops/masked_select_cpu_dispatch.h,sha256=CVOrGgxffjmefDE4ol_cbs0WVjQymCTFlnfzC5BXcIs,1227 +torch/include/ATen/ops/masked_select_cuda_dispatch.h,sha256=tputscIuPHtyViFhoUY0i7LIbkfpSRjOlSZ7GMmIgJQ,1229 +torch/include/ATen/ops/masked_select_native.h,sha256=kl28b0787rMCrLFBg1pdptB-xp3B9hMDtSMZla6nw7E,1086 +torch/include/ATen/ops/masked_select_ops.h,sha256=6JxuLlGclmcjBkpU37azlpm-3Onh2sldEB4xpaKrbuA,2009 +torch/include/ATen/ops/matmul.h,sha256=jaovUb9XTwqKE27kbctsdJD0gqmAzIJZBasYMp0gPAw,1422 +torch/include/ATen/ops/matmul_backward.h,sha256=ullOJa7sUxUo2jxR1MVlt8Cm9x19Zgof9nyhQjKuzuQ,1992 +torch/include/ATen/ops/matmul_backward_compositeexplicitautograd_dispatch.h,sha256=zjWedAfyxYniZthVD21CtctrVy7nPSriYAN_OobwrlU,1389 +torch/include/ATen/ops/matmul_backward_native.h,sha256=nNmp-C5gxNU7xMfkpuASxuOIH_uAwbKJUq6QNaIv5CQ,1061 +torch/include/ATen/ops/matmul_backward_ops.h,sha256=2nKXYY2acX_eQ7xX-r3wLkG-Vni9Obpuk4fGJwmEKAQ,2625 +torch/include/ATen/ops/matmul_compositeimplicitautograd_dispatch.h,sha256=Q5RMaHujLUx5jnNrv0t9E-Cs23SBFWZozSiFjvFrAyg,1253 +torch/include/ATen/ops/matmul_native.h,sha256=Zz9dnwdj3Uwio5KhA3KhRnQdqHGYtHiSbFHr_7qxBQo,1058 +torch/include/ATen/ops/matmul_ops.h,sha256=UwYR5wSUYHT6IAcgrpAreyoIjBOBb72Bd4P8LDRf5r8,1973 +torch/include/ATen/ops/matrix_H.h,sha256=sZ73InVoI6GVyw7j-twzKeJTL97Gv5qSghZAps97PLA,759 +torch/include/ATen/ops/matrix_H_compositeimplicitautograd_dispatch.h,sha256=JPmH-207ngk4S-UHHR25ePZyzDh6XHpGuVWHhYWKZVg,1020 +torch/include/ATen/ops/matrix_H_native.h,sha256=ARYclGs5nIcEvL8YJ4lo1Buw28Ls4H2RdOv92yD8Zyw,732 +torch/include/ATen/ops/matrix_H_ops.h,sha256=fmyA9eNELzU30vgLZNdgVD3jWAUkVyGOCszpxh1SI7A,1223 +torch/include/ATen/ops/matrix_exp.h,sha256=IrzitVtO0-BND4KUrFOvCq8_Knex_3EAdCeiCqKMyfY,907 +torch/include/ATen/ops/matrix_exp_backward.h,sha256=YrGWftbIT2b5dIiKfrpRfZAiSXxM8MtBGGy4UVOwGd8,987 +torch/include/ATen/ops/matrix_exp_backward_compositeimplicitautograd_dispatch.h,sha256=cUZHMzZN2wCjwrfPc-OadcSORBXJ3-QstLtIKdkvS7E,1056 +torch/include/ATen/ops/matrix_exp_backward_native.h,sha256=iR-evb-MdJb2m8qjJ-Tgvu-SAuBtFqANn8LnRd_sKRQ,768 +torch/include/ATen/ops/matrix_exp_backward_ops.h,sha256=n6GP8TRN8cJaiZ-v4Phmc736aUIuNLowspayB_FQR1I,1333 +torch/include/ATen/ops/matrix_exp_compositeimplicitautograd_dispatch.h,sha256=4BtDf0tsCCvsbvo5OV5X3lPaefvaNna3yq_Ss8wSuSw,1022 +torch/include/ATen/ops/matrix_exp_native.h,sha256=tXFiZgSZWGJL_pXmMN4NmCtbq7n-_wffVqWW8bSIgpE,734 +torch/include/ATen/ops/matrix_exp_ops.h,sha256=IowmHVhsp9I5Y_EyxObyGYuQeN_e5UjHmAi_7urG4ew,1223 +torch/include/ATen/ops/matrix_power.h,sha256=jh5c8AoNqfTsA4pT_bh-FXd2RJU4iQTthmUi0Mm4mzc,1404 +torch/include/ATen/ops/matrix_power_compositeimplicitautograd_dispatch.h,sha256=KylM42hSOfvhT8ubTXCqFzOL8HcMaQJ_72S0ODedcd4,1226 +torch/include/ATen/ops/matrix_power_native.h,sha256=f9blx8Wi0WJVRgZ4mntZZ6Pk7nz0Q_E4ee1hrXG2le4,842 +torch/include/ATen/ops/matrix_power_ops.h,sha256=-oxUv7NCbE62rMAXBr6ow-C2QCpje87cccjY-LyPk5I,1913 +torch/include/ATen/ops/max.h,sha256=u6W-s70uWyYj5XCHSogMJXAKJnaLPR3gCVawinI7AOo,4038 +torch/include/ATen/ops/max_compositeexplicitautogradnonfunctional_dispatch.h,sha256=HXAIBO5ZdfZ-NErR0BDkLAowo3MIanL3FkNR0Li9J_A,1099 +torch/include/ATen/ops/max_compositeimplicitautograd_dispatch.h,sha256=MWm7p7t5tS4e8-l3xC_n775KvuYLhZ6bxAeUfdxT-CI,1680 +torch/include/ATen/ops/max_cpu_dispatch.h,sha256=5eqgzJt7R4bFmmsHH325G0r2VSWpg_YQzXpxJEoZSaQ,1546 +torch/include/ATen/ops/max_cuda_dispatch.h,sha256=47h3noNEiY9uY_gZduq0P0Jv2S_OHNWQcH-aYgKd6KI,1548 +torch/include/ATen/ops/max_meta.h,sha256=BNDgonEkZb_vWkzEq7PnTm9a7qpMbEIXEKA4LZ3Bs9M,1291 +torch/include/ATen/ops/max_meta_dispatch.h,sha256=8PIbjPimZqMXco05kwJkH57gH-GPESaB-vYIFnRVJMw,1346 +torch/include/ATen/ops/max_native.h,sha256=58jeM65cL9sUxeWhhOCRT8qK0_TyILCfcz6mqmcfmR0,1750 +torch/include/ATen/ops/max_ops.h,sha256=kilNHwJo9Q5R-Q7Dm5G6naIXKeI6c3pNY1Z83ywHfw4,6306 +torch/include/ATen/ops/max_pool1d.h,sha256=3nLsOLETqlqJHxeayvmmEo7M9G6eQW46kMp9AyxsVxE,1188 +torch/include/ATen/ops/max_pool1d_compositeimplicitautograd_dispatch.h,sha256=Owf8WVpvZsd6SZx-K8LtitQzbFPfP3x4AxVA5WUWnP8,1155 +torch/include/ATen/ops/max_pool1d_native.h,sha256=PIlVY8KDWNiTtwDT1RrJ-Z1hygsq6W6DIEDYvvXcdOU,867 +torch/include/ATen/ops/max_pool1d_ops.h,sha256=HDcdY2PyLvDhF14kR-px9cCTgh8yANh4fJq45f5RODU,1634 +torch/include/ATen/ops/max_pool1d_with_indices.h,sha256=Tp0jZ0OcxG3IjVyobYp6KvcZoD76R_k_RUcskzAbuSA,1275 +torch/include/ATen/ops/max_pool1d_with_indices_compositeimplicitautograd_dispatch.h,sha256=J4DwKwdv8lDc_UQN2CIUCe6BeEk0RgkGCowCZs-Tpi8,1193 +torch/include/ATen/ops/max_pool1d_with_indices_native.h,sha256=SpLGau3sXnn92j7aZ7blLbQ7GLhPM3dzNmY3nO30vAA,905 +torch/include/ATen/ops/max_pool1d_with_indices_ops.h,sha256=tkZs4MuIebTLJ-w1HUFLVfFqGkwVxZvB3X63D6Cta3Q,1758 +torch/include/ATen/ops/max_pool2d.h,sha256=73-9YM7nwbkayq09YNV1lTWLIe8abeWFNYWldNZOr6k,1188 +torch/include/ATen/ops/max_pool2d_backward.h,sha256=HOeg3hpbmfYolMD7Z_7eH_rHzbedMVnz_I3EV5-G-Us,2436 +torch/include/ATen/ops/max_pool2d_backward_compositeexplicitautograd_dispatch.h,sha256=7aFxvT8NWsyoOSje32j0Qjp2BvYRB583zInkYNPWIFA,1464 +torch/include/ATen/ops/max_pool2d_backward_native.h,sha256=lDGs7P8EAQje5FoajByzdOQsRurBYis060R2RjQB10g,919 +torch/include/ATen/ops/max_pool2d_backward_ops.h,sha256=v_1jOxKJEyL70E_lGOri-42XLY8WHjxrb1us1iduAPA,2909 +torch/include/ATen/ops/max_pool2d_compositeimplicitautograd_dispatch.h,sha256=4bRNk_juojizScxd8SSSYrMtp_R4S231zdVWE8GvVgU,1155 +torch/include/ATen/ops/max_pool2d_native.h,sha256=F9EUWLoNlKvNEMxiTKQCjFPKEedoy6gBz-I39DXBVt0,867 +torch/include/ATen/ops/max_pool2d_ops.h,sha256=_Qnsd_7AQNZNL5SPzC9S0u2ntaDZDeG3IruAGNTKsQs,1634 +torch/include/ATen/ops/max_pool2d_with_indices.h,sha256=wJ_eYaqc9gkjm8KOFWjc4nlZG0yykKBgBL5YrF5WA8w,2500 +torch/include/ATen/ops/max_pool2d_with_indices_backward.h,sha256=p3siom_wNOWepDPBXYxuD2q1fKF6LSlLuU5YuxV1W9g,2730 +torch/include/ATen/ops/max_pool2d_with_indices_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=3EQunEvLzt83_Mo9wuoTg11wBWMTSj4AysdKDy9TgQc,1250 +torch/include/ATen/ops/max_pool2d_with_indices_backward_cpu_dispatch.h,sha256=n0Qk3hxtHaigNSCh0Gc-cX3fsPUALgCiEWGFSNphA74,1763 +torch/include/ATen/ops/max_pool2d_with_indices_backward_cuda_dispatch.h,sha256=C4VeX5BS7hQprkWg3e8m7tIUZJvpCHr8TGhsFU9e0Xw,1765 +torch/include/ATen/ops/max_pool2d_with_indices_backward_meta.h,sha256=1H63oFP4OxZPMUaXcvBZn1ne2kzKQWvsSYNK54JruvI,1027 +torch/include/ATen/ops/max_pool2d_with_indices_backward_meta_dispatch.h,sha256=KlHDoAP_PPOhUBRFlXWU5IlCxQLQDloYnhDIg9JcwfI,1765 +torch/include/ATen/ops/max_pool2d_with_indices_backward_native.h,sha256=0cLzUSly6lHWsQewAdDZPwHv5VmpgE24z5FXyCrQrZI,1504 +torch/include/ATen/ops/max_pool2d_with_indices_backward_ops.h,sha256=7-QLg3POJ2k9gW9pYrxoxgOm1dApJdV98jFWSqPRRGw,3187 +torch/include/ATen/ops/max_pool2d_with_indices_compositeexplicitautogradnonfunctional_dispatch.h,sha256=fRT-3r6eZWGGrB3wTqldhlfYjQzylGwRgbvmgT5-hpg,1219 +torch/include/ATen/ops/max_pool2d_with_indices_cpu_dispatch.h,sha256=7bNB-17Ht8XSHAwI1bVVIACEUGy35qJpi1_ZWZL9TM4,1691 +torch/include/ATen/ops/max_pool2d_with_indices_cuda_dispatch.h,sha256=bax5HkuBGFlPKpwLSk6PjRHoIzGkT7gTnw3zPXmXOqc,1693 +torch/include/ATen/ops/max_pool2d_with_indices_meta.h,sha256=K-dBuTS5cjtPfm0exSzqbhED72MZ9UDcbHT-Uymp7GY,958 +torch/include/ATen/ops/max_pool2d_with_indices_meta_dispatch.h,sha256=BVXVjT4iuhk-b5UyH22LYBLGmi01iLgAlbj9pg_4gng,1693 +torch/include/ATen/ops/max_pool2d_with_indices_native.h,sha256=71QmJhhxp_LBOXYk3NHvNX6tewGGlxuoCDterf0PAjc,1381 +torch/include/ATen/ops/max_pool2d_with_indices_ops.h,sha256=UJ-qIvjbH03fo3D05uMjkoj_edto82xV0pIp7aG8uxw,2983 +torch/include/ATen/ops/max_pool3d.h,sha256=gWR3wqBoVSV5JTE4zo8qHQ_8tcCSa9YsfKOGGmWoqf0,1188 +torch/include/ATen/ops/max_pool3d_compositeimplicitautograd_dispatch.h,sha256=uWPnZeRFJCRvisbsPGmVA_MM36OHxRUARFTaWooAcxU,1155 +torch/include/ATen/ops/max_pool3d_native.h,sha256=zwc6Kjp7HG2YfbCk6FL8t7ddJ1glXQcOt_aREmansVI,867 +torch/include/ATen/ops/max_pool3d_ops.h,sha256=MyQ9hSsLCk_bOrZmfphMceF1sCO41SCuQH_FBlXePwM,1634 +torch/include/ATen/ops/max_pool3d_with_indices.h,sha256=bbRLAlAWoOemhxBjvJAA6yKVEt4JjvP7gTPB9oN2ULM,2500 +torch/include/ATen/ops/max_pool3d_with_indices_backward.h,sha256=eT533Cy4gRumLT5B4HbZoZ9Y8unad6tU9b-Bu78XAT8,2730 +torch/include/ATen/ops/max_pool3d_with_indices_backward_cpu_dispatch.h,sha256=XDYGz30WNKWS4IHAzQsgh5mWFWCjAtLJs1BNyuMVLf8,1763 +torch/include/ATen/ops/max_pool3d_with_indices_backward_cuda_dispatch.h,sha256=vRwyLOQtrbJXPV2o8NWKXrRVS6NrsFKvB1JOCM5y7gA,1765 +torch/include/ATen/ops/max_pool3d_with_indices_backward_native.h,sha256=h354nw1axCC7D4mjmN7aVQaA3cDND52aIKW0CilX7QU,1796 +torch/include/ATen/ops/max_pool3d_with_indices_backward_ops.h,sha256=sjgnQZt-grP561wvS9woaJ8irMck6ERjDmuQwG-B5SQ,3187 +torch/include/ATen/ops/max_pool3d_with_indices_cpu_dispatch.h,sha256=IXn52oYMyeg2jTnpT7dbfMKR21rnyCImAsrDZOcOAOc,1691 +torch/include/ATen/ops/max_pool3d_with_indices_cuda_dispatch.h,sha256=hkXwUEl7iI9etz5kMSMuUkmhbxdJrm6WtjQ6nQFiw1M,1693 +torch/include/ATen/ops/max_pool3d_with_indices_native.h,sha256=fZDueaRIdMVuNSttagl4eCUBIOPRE3yVCW42GpipEiM,1680 +torch/include/ATen/ops/max_pool3d_with_indices_ops.h,sha256=XAtkzJTtnbXmMGkyh0zexvRJ_f6ExcB8iYl6JGZB290,2983 +torch/include/ATen/ops/max_unpool2d.h,sha256=SXqQHj_pfv8qhMqfdvC4_ugC6dzdQVJ84MUz9HhgT48,4687 +torch/include/ATen/ops/max_unpool2d_cpu_dispatch.h,sha256=kFQZz9kAjLOaAPAZvkrRHCW8QToKLGQKlnErtReyxZU,1753 +torch/include/ATen/ops/max_unpool2d_cuda_dispatch.h,sha256=cnEIQEG4PmJQtJd1rsagUC-QNzZoZyaAob62xOhu0aw,1755 +torch/include/ATen/ops/max_unpool2d_native.h,sha256=RMDXXGJUZ7TJ-xtvMgXWXOKJZZdP8FuAhppC1wYjJ6Q,1254 +torch/include/ATen/ops/max_unpool2d_ops.h,sha256=o4Mr5UYs9sCNU9cXNPWzMNEbhlgZu_ktyhPlSL3T5AM,2241 +torch/include/ATen/ops/max_unpool3d.h,sha256=VfwcR-ZKW3i81SNiV8c5kIFybRG-ZW1LDVfiMWkpKFE,5665 +torch/include/ATen/ops/max_unpool3d_cpu_dispatch.h,sha256=o2SA5rvkgBRyoYOeoQFMwtXjW8sFKBN65R-eAOIk2ek,2047 +torch/include/ATen/ops/max_unpool3d_cuda_dispatch.h,sha256=uQoqVl7msihpXZ4KC4_2o6fv8kOn3fWamV39jboB1_Q,2049 +torch/include/ATen/ops/max_unpool3d_native.h,sha256=IZ7PFVKVBmMqDe2yagvU4LkmVTnBD4IJYGhb0VrJzPE,1450 +torch/include/ATen/ops/max_unpool3d_ops.h,sha256=gLtHwsWUnr2dNVOIxcDDKGJGvAFRKWi6s3R0-3_rBuo,2567 +torch/include/ATen/ops/maximum.h,sha256=zT515n3Qi4kj6tSiLH8DMN5hWsLANZttowF3GxaO1Js,1432 +torch/include/ATen/ops/maximum_compositeexplicitautogradnonfunctional_dispatch.h,sha256=4flFPhuJ9-ev43Hm6d_oU4jC6c1IzKjoSd5Yl_JMfNU,1071 +torch/include/ATen/ops/maximum_cpu_dispatch.h,sha256=8TAv6qXCmqcxHNHQrDLzcCqWF0dMqxZPyN7t-E763eo,1212 +torch/include/ATen/ops/maximum_cuda_dispatch.h,sha256=1kw9IfbFRVgQlmE2XT6FonYyyp4xfeB9zGpUCgo1O7w,1214 +torch/include/ATen/ops/maximum_meta.h,sha256=OQssgHUq5hhgOZP-Fold8yuif3_5Q90zCvxC8WFo_uI,848 +torch/include/ATen/ops/maximum_meta_dispatch.h,sha256=iVnmEUJwht-AQ1nhTnyq65uKHzqfvcIABznPC_mHBpc,1214 +torch/include/ATen/ops/maximum_native.h,sha256=jPnmvvnlaxQdnHcJG1evFw7gliZdtJrIdg_BLr8wCkM,879 +torch/include/ATen/ops/maximum_ops.h,sha256=M7mjEkAzeqiRQ2rYhvx5k0Y4trrGzWHcdYcOJoYIY9U,1979 +torch/include/ATen/ops/mean.h,sha256=M3XzIkrNYLjyqaLcnqM9lePTKEVeG3UFPSwCHDPC9dc,3629 +torch/include/ATen/ops/mean_compositeexplicitautograd_dispatch.h,sha256=kONrR20k3MH6fuRoctEvnhCCDzYvoeF0DmrKhdM0jRI,1316 +torch/include/ATen/ops/mean_compositeexplicitautogradnonfunctional_dispatch.h,sha256=PTi35X_c5u6Ufv808yBTPyQq3cb_FigOAv7Bd1kXuAI,1145 +torch/include/ATen/ops/mean_compositeimplicitautograd_dispatch.h,sha256=nlvGAcXdxh6YZ-eUZwGhd4rFyusCiqCPUsCo_NPShNU,1433 +torch/include/ATen/ops/mean_cpu_dispatch.h,sha256=Wexg4e0TlQsSCgKLDs08EjWV682boy6u5d0Xvx6WJbE,1413 +torch/include/ATen/ops/mean_cuda_dispatch.h,sha256=ZAAwc7aSRD_meR_sT3D8yzOFxF-FhcPMflHuIzp0fyk,1415 +torch/include/ATen/ops/mean_meta.h,sha256=cVslIWJxSriwnjCQ8pi3tNaYC6gPhl8dP68ggE5B5zk,905 +torch/include/ATen/ops/mean_meta_dispatch.h,sha256=C4KAHwNCxwZbMJLaKii2r4tn_jglUcWAW0i08tLGIzU,1415 +torch/include/ATen/ops/mean_native.h,sha256=BsA6hVwfBogM3KdUdhb88e0V8-Ez1VvCnPagzE0QEuw,1795 +torch/include/ATen/ops/mean_ops.h,sha256=oQiKOVLKFEwKhRWrnZ1Lc65KZbNI9cPUuK5HUsPtG9g,5398 +torch/include/ATen/ops/median.h,sha256=-CpiD6nMX9An_sNDwZ-Eqj2uCM8iRNVAHjr09QCFQXE,3471 +torch/include/ATen/ops/median_compositeexplicitautograd_dispatch.h,sha256=TI_cc_DVPuPHPcM0OnmGt1yZUIwIt8YkpfYfkCZz-IY,1233 +torch/include/ATen/ops/median_compositeimplicitautograd_dispatch.h,sha256=sLFHsWid9Ja80-AujVohrOw8TXIqC8PblhCPD1vohsI,1409 +torch/include/ATen/ops/median_cpu_dispatch.h,sha256=hlYLTs9fsev6FhdQ_4TbjWCZKCR05_Yg6oIpNKYZ3ME,1295 +torch/include/ATen/ops/median_cuda_dispatch.h,sha256=1qEEeIs8xcWESIdEKkYEk3lVH2LhUejoMhFb2T2OTC8,1297 +torch/include/ATen/ops/median_native.h,sha256=mnKsszWgai55f99JzhFBUX_KpDxYcSliAUXZMIZeZLc,1583 +torch/include/ATen/ops/median_ops.h,sha256=pylvwOb-SfpxwF3pe4VkIkXJPwLiETHv_BwIPqh_9O0,5099 +torch/include/ATen/ops/meshgrid.h,sha256=1ltTaB7X83B4AEsR6imhOHMJSzg_frPMqgNQCEzbQ3g,1160 +torch/include/ATen/ops/meshgrid_compositeimplicitautograd_dispatch.h,sha256=HxALCti4vu7xjxQt-ZOw3rgvMUO1kYOElv3A7DIKO8M,1131 +torch/include/ATen/ops/meshgrid_native.h,sha256=LBo7w0u98i3b7Lez9p6l_5stVlmyjqozMECwFulp6bc,843 +torch/include/ATen/ops/meshgrid_ops.h,sha256=1f1FUppy8pFB_V9LTzPl4omXhz7W1zRwrxB_Ov2ktFc,1930 +torch/include/ATen/ops/min.h,sha256=cpQaPcuKLF49h-qZbcaL83Zl5Vs73NrRoYeqLkbSta8,4050 +torch/include/ATen/ops/min_compositeexplicitautogradnonfunctional_dispatch.h,sha256=qyB7Evfb7kH0qzYBPggbSLSYYzh5mBYWfN2vTlRavuE,1099 +torch/include/ATen/ops/min_compositeimplicitautograd_dispatch.h,sha256=k_92397f52BLK-eYf1AFKqV7dKRiEyAJVLYjK7jex6A,1682 +torch/include/ATen/ops/min_cpu_dispatch.h,sha256=Gws1QIdqhoEhHohdeUJfIIJZ7_W6wNUAqdSYJJWC3aM,1548 +torch/include/ATen/ops/min_cuda_dispatch.h,sha256=uGwQzJQZOD43VrYLsrcdyWBlNRPjqSzPJCQ1Zj_fN-M,1550 +torch/include/ATen/ops/min_meta.h,sha256=L7c2brEm4cYFh4twYgUbGGwpyPSrmtB0qfzMR0tsJbU,1291 +torch/include/ATen/ops/min_meta_dispatch.h,sha256=Gx5ifqW6UYhdbvNaAROn7uTitXGrn6PJYyL9bbcEZmI,1348 +torch/include/ATen/ops/min_native.h,sha256=o-SyJHCO-j7xrNsXHcBi1HptilupVD5cTOsbc6dLT68,1752 +torch/include/ATen/ops/min_ops.h,sha256=WuJV8gZXdE239MPWCjsN9kKnOHREbOUofDrOskNSICE,6312 +torch/include/ATen/ops/minimum.h,sha256=FUDt-lh6oFoSDYPj1iAXt9sLE8Tbw15X9djlvu4nTyA,1432 +torch/include/ATen/ops/minimum_compositeexplicitautogradnonfunctional_dispatch.h,sha256=mB-yGMQp0Hcn6Bc7qYvgekts81y9jC-_m1ji_QmvU7Y,1071 +torch/include/ATen/ops/minimum_cpu_dispatch.h,sha256=3PHqMoGnakLvjSxPT2nSyVFk60XeL3I1PZHBttblF7U,1212 +torch/include/ATen/ops/minimum_cuda_dispatch.h,sha256=UycuTLvdhR0yBFccCvyujK-4RyK6pi9yfvJdi47bk9I,1214 +torch/include/ATen/ops/minimum_meta.h,sha256=7G43dicDfmmT7_vpOelCnfB93-EnL75tFUC7fu5vfRU,848 +torch/include/ATen/ops/minimum_meta_dispatch.h,sha256=DL-1qyViExwJqCumuADsyDpIcUlXr4rsalR3MtbadfA,1214 +torch/include/ATen/ops/minimum_native.h,sha256=VtLsZ2TFJlzItjCO0JJFNzopkUbFD2_Wgt-sEwePb7Q,879 +torch/include/ATen/ops/minimum_ops.h,sha256=wxxlHeTRW7dEzAmQOJyf-rZzPrqAL8SnRLjGOOoobjg,1979 +torch/include/ATen/ops/miopen_batch_norm.h,sha256=DzyIq51M_3hkcQ1OM_rRvbUdH8RO8oC-xatNRbCL9AU,3134 +torch/include/ATen/ops/miopen_batch_norm_backward.h,sha256=3yDrXJGrJ7knlKoeJWAJGkjgO7f6QXkIDDmvZt54PTY,3275 +torch/include/ATen/ops/miopen_batch_norm_backward_compositeexplicitautograd_dispatch.h,sha256=96k4AY0M1CvRNRIa02hOjCn69tXjTHxyfiQaZpRKLvo,1855 +torch/include/ATen/ops/miopen_batch_norm_backward_cuda_dispatch.h,sha256=pO6zsV44b4eCxlWbAzx5E4_KoKJ3vfekfOyvFPIfwJs,1300 +torch/include/ATen/ops/miopen_batch_norm_backward_native.h,sha256=BBkH9I9wCSsaXEUExxh8AvH0jetqr5MIdWxsEHskhro,1499 +torch/include/ATen/ops/miopen_batch_norm_backward_ops.h,sha256=WaFU082P-coTiq6fnzyA3iezG4qFY2m9zGG77BHh9W8,4060 +torch/include/ATen/ops/miopen_batch_norm_compositeexplicitautograd_dispatch.h,sha256=rXefiC1_CAOnPW5S2r3LAEzRYG1WADt-jkF9HwifNvo,1771 +torch/include/ATen/ops/miopen_batch_norm_cuda_dispatch.h,sha256=8eOzvAf3888L7tq9Gzrn2wefADtHv8FyVw6XpxfNfjY,1258 +torch/include/ATen/ops/miopen_batch_norm_native.h,sha256=FJc31fcwy8OKBiqQ0knk2HuymsbdztO_rvNdg-5pIWc,1415 +torch/include/ATen/ops/miopen_batch_norm_ops.h,sha256=DobfZ0Tu_CkFVQBLr1mxMhiy4WdMpxuy2NmMPRsJlh0,3800 +torch/include/ATen/ops/miopen_convolution.h,sha256=vjoINAGOchonGjck1FZKRpJFUYRz4xUFBsv5G9d4b80,8131 +torch/include/ATen/ops/miopen_convolution_add_relu.h,sha256=F0y251lqFpkSQqef62fejd0fADpCGO5ZKBx-p4-QKHw,3226 +torch/include/ATen/ops/miopen_convolution_add_relu_cuda_dispatch.h,sha256=bu5l12qrnHxIQVwNveklWAoBAo03dppr-7nyemlB8wM,1545 +torch/include/ATen/ops/miopen_convolution_add_relu_native.h,sha256=gpPtRnqa0sH34uJ_M_uCSg6hm15hViZwGFwJpvJPZkU,976 +torch/include/ATen/ops/miopen_convolution_add_relu_ops.h,sha256=seInneUD30AZD3gHY01ukWrXIsyfmaA8L-FuUfppcDE,2069 +torch/include/ATen/ops/miopen_convolution_compositeexplicitautograd_dispatch.h,sha256=lPwmtmEHqU3aVSYlLsTpwQEwbNt2hZU25ABprebKKCo,2156 +torch/include/ATen/ops/miopen_convolution_cuda_dispatch.h,sha256=1HVsX8rlHcDcqVP_7KHpM2uu8bJICXQKWoGOlg5RZII,1469 +torch/include/ATen/ops/miopen_convolution_native.h,sha256=muxW44VLb-MeR-SXgECLJEU-LdXQpHM7DvJQhSButyE,1247 +torch/include/ATen/ops/miopen_convolution_ops.h,sha256=wNEFCcc4tyCAx_lWcIXiDWxiWMWKgp_fuUrLnzyu1lE,3279 +torch/include/ATen/ops/miopen_convolution_relu.h,sha256=ZmSOqx-_t7rIjCkWBBTJDbpbByfXkDCeP2i74ZVEB-4,2832 +torch/include/ATen/ops/miopen_convolution_relu_cuda_dispatch.h,sha256=dweI5j-729KqjFBMzvU6C0I8hRMEBWBy_MQP-zAS03k,1407 +torch/include/ATen/ops/miopen_convolution_relu_native.h,sha256=g4vie8NZArMRsAPoAYFt4lvGZ-tajxyc8ySJf6ETe7Y,907 +torch/include/ATen/ops/miopen_convolution_relu_ops.h,sha256=6pWF6RDYmS9-uxdWvjw8jRlfLDs-2qfNI699kIhcEPI,1845 +torch/include/ATen/ops/miopen_convolution_transpose.h,sha256=YDtjkwR_s7KHHQbj3isZwUzUwkwoOkm6nL35t_xtlV8,9347 +torch/include/ATen/ops/miopen_convolution_transpose_compositeexplicitautograd_dispatch.h,sha256=eYIHDkImHd-OHjf8zEtMIdwTeVyV4hq0z_Wsu6JiQt4,2332 +torch/include/ATen/ops/miopen_convolution_transpose_cuda_dispatch.h,sha256=iKQTuQMhQEdqdfslvCpsIzfDXKxcO2uMoXALoIeE9Ws,1557 +torch/include/ATen/ops/miopen_convolution_transpose_native.h,sha256=J9phkFMTErsYNOWG81I87uWBYtxXZyTB9ov9zveonjY,1335 +torch/include/ATen/ops/miopen_convolution_transpose_ops.h,sha256=Yvsp-RkddHerw5uNhoy8vyOHsYXPehBaDDjos8jwYW4,3575 +torch/include/ATen/ops/miopen_ctc_loss.h,sha256=odgd1sX6iSsNiFMYvoKDOVc2rJldKfBVV7PQeXB2UTQ,3189 +torch/include/ATen/ops/miopen_ctc_loss_compositeexplicitautograd_dispatch.h,sha256=v2Pe8eN6UGFUB-3KFgJShNB74GhSJ9IiFlSdS_oSLv4,1535 +torch/include/ATen/ops/miopen_ctc_loss_cuda_dispatch.h,sha256=c0_BQjXyFgGTgx0vkuOHaMANpmJW5iIggZ1v8Mjw1Bg,1406 +torch/include/ATen/ops/miopen_ctc_loss_native.h,sha256=oE8H3iITSdnswboq2r0pMt50fEJ8TT9wKmvc7jyekEQ,1452 +torch/include/ATen/ops/miopen_ctc_loss_ops.h,sha256=ew_NNjkxVpI1nw5wTykRzhi8h0dvcAiGohVHd33yL3Q,4235 +torch/include/ATen/ops/miopen_depthwise_convolution.h,sha256=z4zSRFUpNSv2s9cWEzxzhCgX79-aZ670McIPZsFhdgw,8441 +torch/include/ATen/ops/miopen_depthwise_convolution_compositeexplicitautograd_dispatch.h,sha256=1Kgn6-PIhsoXkDB5TXE8Itp3qaZBf9wRiwoF1h1S_E4,2196 +torch/include/ATen/ops/miopen_depthwise_convolution_cuda_dispatch.h,sha256=9DpqV4lwcaMN5eQ_V0RWrJm1DwlYb_KQx5FbceSe1e4,1489 +torch/include/ATen/ops/miopen_depthwise_convolution_native.h,sha256=TdogTUSlfr4N7EZfpzgneSH7UVftQOmpp7Oaoq-jS34,1267 +torch/include/ATen/ops/miopen_depthwise_convolution_ops.h,sha256=qiJwgD-5YiFL3NCsYbqh7FYJxBFflB-TCF_Pbes0hos,3339 +torch/include/ATen/ops/miopen_rnn.h,sha256=BV3xDqKXPpUpsVqkp9Iwv6rF_cxXVufO_LkRdr9Ae-c,3964 +torch/include/ATen/ops/miopen_rnn_backward.h,sha256=zfGDCwt7FjFUgZ3dfUhkYzDUhBte-KBIxHRDjF5sC6c,5086 +torch/include/ATen/ops/miopen_rnn_backward_compositeexplicitautograd_dispatch.h,sha256=uLUpCNO2POWiLhKPezkl3M56qmdG1Bzru6MdAowaX2U,2393 +torch/include/ATen/ops/miopen_rnn_backward_cuda_dispatch.h,sha256=5FF8zP9_dtYcDRPMUtQ5ZlS_b3AFrw55KXIa9mnXYv8,1622 +torch/include/ATen/ops/miopen_rnn_backward_native.h,sha256=blqOkjdqVyyBhhXMQLTfEaMUFZpE7kOkTozpNGVpAMc,2090 +torch/include/ATen/ops/miopen_rnn_backward_ops.h,sha256=lno9-oinrU5jZBY-vqSC6GeJLW6zis3XCSkHwFeX05Q,6001 +torch/include/ATen/ops/miopen_rnn_compositeexplicitautograd_dispatch.h,sha256=CY8RbdBm-kZILZUbTfozNnihIPqdWbNuMtjdp95SQFA,2039 +torch/include/ATen/ops/miopen_rnn_cuda_dispatch.h,sha256=A-dCGCveAi_2BouFnROF1uIXAsaP-oxoprLqYiCChS8,1350 +torch/include/ATen/ops/miopen_rnn_native.h,sha256=dsPfWLw9OurpYZ8YkPeLWUXbIGk7WtGVbOmf1wfrnPc,1641 +torch/include/ATen/ops/miopen_rnn_ops.h,sha256=WO9BMnhKxvr2Qt7-6OsgiyM6T33tTZ9IUQugDRHogio,4608 +torch/include/ATen/ops/mish.h,sha256=AFjehP8XrwGy8WKYdbziRBQ60Ncia0qB64XsfQpN9Gw,1397 +torch/include/ATen/ops/mish_backward.h,sha256=oazQNMN-gWnv-0Q_yHWd78lrwqeV0GhcvZCJ8q9PLp8,984 +torch/include/ATen/ops/mish_backward_compositeimplicitautograd_dispatch.h,sha256=6UTmKYnWd8vGflhzzvoeboxLsAvia9ruf90SF3hlfg8,1057 +torch/include/ATen/ops/mish_backward_cpu_dispatch.h,sha256=zxh0po5wcWMfIbKAuqSjY1mAXLmKzJRqwGC-Mp81-Dc,1013 +torch/include/ATen/ops/mish_backward_cuda_dispatch.h,sha256=57F_iF1FtW6fnutPijCTzlfqWYHphv6fiIpO9o5BAvk,1015 +torch/include/ATen/ops/mish_backward_native.h,sha256=tkq9H10HLcBUS4VUfFt37tymSdUI8y3bqgblB1uQG7U,867 +torch/include/ATen/ops/mish_backward_ops.h,sha256=ONPY7WZadEAeyIKY_RAi_oReVWhFQKDiQPDpIJB7tZE,1336 +torch/include/ATen/ops/mish_compositeexplicitautogradnonfunctional_dispatch.h,sha256=RxMl1vUDxTwFc2ZV5YVCNQYPhX2_PPMle2Mn5s35EF8,1091 +torch/include/ATen/ops/mish_cpu_dispatch.h,sha256=630vF_vnjCAZvghVjOoGupknZRAJigMVse4idBc6fXI,1174 +torch/include/ATen/ops/mish_cuda_dispatch.h,sha256=hGPq3lVyfahqnfmRtosdkVMn3JeTFGffDVs5fWKX-Dg,1176 +torch/include/ATen/ops/mish_meta.h,sha256=lvJ_U-U9AP5A75rq3IY5MlFGDY3HI_P2j5BgAd6qXG8,819 +torch/include/ATen/ops/mish_meta_dispatch.h,sha256=ByUkAEkUbVfDpf3Vzpdt73oOeVjKih7XqM2edYQ9IdE,1176 +torch/include/ATen/ops/mish_native.h,sha256=1fryZdZzRlotHDfMaUEDGcRFcFWhpcfht6jPVv2U874,844 +torch/include/ATen/ops/mish_ops.h,sha256=2ThhdTYtU8Waljdw8Lu5rxoOppjWy0MBWav3Ugsjgeg,2282 +torch/include/ATen/ops/mkldnn_adaptive_avg_pool2d.h,sha256=OiItyJQkLXpYwzV0Z36fhIZKdsp0ah0cVlc3oSeVlVA,1667 +torch/include/ATen/ops/mkldnn_adaptive_avg_pool2d_backward.h,sha256=KGy2UoFkafW_HiVdiC7Zimv6XYGi1TpFcqHq9nU4jvQ,1766 +torch/include/ATen/ops/mkldnn_adaptive_avg_pool2d_backward_compositeexplicitautograd_dispatch.h,sha256=Bev411GqygfeQ9aCzGS562-oTYa7944R_gLt0aGl2WM,1243 +torch/include/ATen/ops/mkldnn_adaptive_avg_pool2d_backward_native.h,sha256=inVq12MBAeuvAlVZ6FSMVfu_RDOmYvc_uPU-oIK0Sko,930 +torch/include/ATen/ops/mkldnn_adaptive_avg_pool2d_backward_ops.h,sha256=RepvNEOE21bnny-bPJOKhFEAs22KX9Qqu6P4BPg2jf8,2183 +torch/include/ATen/ops/mkldnn_adaptive_avg_pool2d_native.h,sha256=-N6Fuga_L1PvnQiCYlSxn3M-JUnIe2FyHJ-Tb3i0xw4,906 +torch/include/ATen/ops/mkldnn_adaptive_avg_pool2d_ops.h,sha256=okUu1qwLD6H-Gp6mK8L1so4CQv_vJN6doIEH_AV6pnk,2111 +torch/include/ATen/ops/mkldnn_convolution.h,sha256=pIzmpo745a8fcMCEbNwiJm3MTCh1hYHsm45ukK8weAM,7171 +torch/include/ATen/ops/mkldnn_convolution_compositeexplicitautograd_dispatch.h,sha256=HhkZadR3zIsJq3gEdBAL-p-tgKgwFNouBH2CWslNLFI,2487 +torch/include/ATen/ops/mkldnn_convolution_native.h,sha256=jMTdW5qKqbDaD_FHR-5hXb6685MxmW4dJPhjVeYmkVI,1175 +torch/include/ATen/ops/mkldnn_convolution_ops.h,sha256=t89eTIbf15QKipz7_loD4bNrdpT74CK-uyu1QQkaMTw,3039 +torch/include/ATen/ops/mkldnn_linear.h,sha256=UiM5a86XJt1Ot63bRKeQCo0cnEUbW7Iuelh7Ugz-CAc,1708 +torch/include/ATen/ops/mkldnn_linear_backward.h,sha256=cJ9Y07b1178QxQaEOWOhshOhVUIPFbUcm5lOIPNuiuk,2350 +torch/include/ATen/ops/mkldnn_linear_backward_compositeexplicitautograd_dispatch.h,sha256=3Yzy_I744W-26cxZ4yNiUXo2lSvQvOC04eTuYXX1pfA,1497 +torch/include/ATen/ops/mkldnn_linear_backward_input.h,sha256=3Xcp-aDoKyQ1kMWOKzlWrLTX2g7gBsV0LD3RclxYBqY,1888 +torch/include/ATen/ops/mkldnn_linear_backward_input_compositeexplicitautograd_dispatch.h,sha256=kvUTeh7owlnWP8i9R1MQIR_bHtO_wtNdz1L7TIhQY-I,1289 +torch/include/ATen/ops/mkldnn_linear_backward_input_native.h,sha256=plcBsrmnt728jmzYItaThBI1coLwIpLuE_9RoHXH1mQ,976 +torch/include/ATen/ops/mkldnn_linear_backward_input_ops.h,sha256=fi5QTmYklcRXSw8xN9xcP9xS1RjFKrZs6HqGITOf7T8,2335 +torch/include/ATen/ops/mkldnn_linear_backward_native.h,sha256=whQoWTcVfM74ULBrKr64eTwRSMLpJDIfNjKb7Vz_SCo,1141 +torch/include/ATen/ops/mkldnn_linear_backward_ops.h,sha256=xo_LFk3xoEQA5B0umrPspLRiz-nos7I7aclanXRcdIg,2918 +torch/include/ATen/ops/mkldnn_linear_backward_weights.h,sha256=izqEOKt_f-NmHVlCXW1XgK67Fy2g5R6ZhTEKlISxG6c,2238 +torch/include/ATen/ops/mkldnn_linear_backward_weights_compositeexplicitautograd_dispatch.h,sha256=MnJpoOPV_4s7qiyKUF7O2wIC-30t0stazGCjtSJiHp8,1421 +torch/include/ATen/ops/mkldnn_linear_backward_weights_native.h,sha256=hG2ejd5PYL_HIur_5dO5NCsnkSfNkGZ8AxTkwcGmcm8,1086 +torch/include/ATen/ops/mkldnn_linear_backward_weights_ops.h,sha256=DaKBbwTViHdci2X6H1YbiqYNxjJ8RbCSpY5q7HZe5LM,2715 +torch/include/ATen/ops/mkldnn_linear_compositeexplicitautograd_dispatch.h,sha256=kTPQD0X28lzN7o66r11O_ION_rhcMx9zEiZNTJJuY1M,1276 +torch/include/ATen/ops/mkldnn_linear_native.h,sha256=RLAmNgG_uENDJ_RmfVXjKoqnrH6nE1L3y5F5cvrLb2A,963 +torch/include/ATen/ops/mkldnn_linear_ops.h,sha256=N3Xpcx4bljhcupDc2qFdMnBwZjGUO1-MeT78gaBamBk,2301 +torch/include/ATen/ops/mkldnn_max_pool2d.h,sha256=fNgCUQg7ULU50vnCnVqNIomoUzjbmXl__3-xW4RjuEA,2221 +torch/include/ATen/ops/mkldnn_max_pool2d_backward.h,sha256=53HX_XRwFfb0GCZINWIR5GPzpvI5GS0w43MQWHfEmT8,2665 +torch/include/ATen/ops/mkldnn_max_pool2d_backward_compositeexplicitautograd_dispatch.h,sha256=HhVf22eibCb626ro9LDAH3qtoh1xviyqmeEiBxSjH2s,1534 +torch/include/ATen/ops/mkldnn_max_pool2d_backward_native.h,sha256=ruNDMgY1Ea2q2j38DQl5wZE03pVW9051J4c_RVQlwzE,1221 +torch/include/ATen/ops/mkldnn_max_pool2d_backward_ops.h,sha256=NMVMqNgBTG-MCfVjqWhUak0fsmv2JoVnLWV4pBBtyB8,3135 +torch/include/ATen/ops/mkldnn_max_pool2d_compositeexplicitautograd_dispatch.h,sha256=qOqArPvdIBU1u9fotY3cNKV1d5m5irMGCf_wZl9Qr2g,1396 +torch/include/ATen/ops/mkldnn_max_pool2d_native.h,sha256=Xmn9-8UskhsdtG5P8AUTGykOtIa0ti9B92zHBxx6svI,1083 +torch/include/ATen/ops/mkldnn_max_pool2d_ops.h,sha256=3kffAJ09pk_bQJf_-hRxRiJfEuO2xS6p8prJ8M-kFYI,2689 +torch/include/ATen/ops/mkldnn_max_pool3d.h,sha256=ao7vjFT9PzCO-Vdg4rQIUyzOwZ9hXuIalLVMyS0XWds,2221 +torch/include/ATen/ops/mkldnn_max_pool3d_backward.h,sha256=EPHk_OTMyPHZbaUYtMaF15mM1ZQc2owByXRJpNbLxzU,2665 +torch/include/ATen/ops/mkldnn_max_pool3d_backward_compositeexplicitautograd_dispatch.h,sha256=mbQDgWhRUs4X6vYRhq7iHMYpU4dC06g6qvOA8k7jZOw,1534 +torch/include/ATen/ops/mkldnn_max_pool3d_backward_native.h,sha256=3RVD3zo8L0-RbDMvHizjHgo_vRJUCqjA1GRg-F2qqsc,1221 +torch/include/ATen/ops/mkldnn_max_pool3d_backward_ops.h,sha256=Up52WC47eV5BL4SISUm9CATSjCXxyc7cAWCVgGW3GYQ,3135 +torch/include/ATen/ops/mkldnn_max_pool3d_compositeexplicitautograd_dispatch.h,sha256=-IfwNPIIoYSRWtMeRA2QfzegD-9m5vnWc0Fvk2wp-w4,1396 +torch/include/ATen/ops/mkldnn_max_pool3d_native.h,sha256=rqHag1R7l0YGdqDI8o9cj5rq5LrwGri9QfbgSTTmp_I,1083 +torch/include/ATen/ops/mkldnn_max_pool3d_ops.h,sha256=YPAkfEPUJKqdqzwIG7ZdJDo3AmAedN1uFUPiuy0hI9Y,2689 +torch/include/ATen/ops/mkldnn_reorder_conv2d_weight.h,sha256=fVwB0JWSJESu6XhZNw6qQfAehXyQw4e9Ds6_-DVTMxA,8019 +torch/include/ATen/ops/mkldnn_reorder_conv2d_weight_compositeexplicitautograd_dispatch.h,sha256=meQ4Sp8H4y9_y-vxEe0r9pvrG8STS4EAsLt7DMSGuFE,2011 +torch/include/ATen/ops/mkldnn_reorder_conv2d_weight_native.h,sha256=5QMuRkxBh49EydtkSEj_-Bio1OYoM3qFexDWHzCg-Rg,1155 +torch/include/ATen/ops/mkldnn_reorder_conv2d_weight_ops.h,sha256=j2lxKZol50uzH0Z7v5X5FNbR6_V2XgOFxAPwj9XalyY,2939 +torch/include/ATen/ops/mkldnn_reorder_conv3d_weight.h,sha256=C1l7JStT2V26C0cCgGmVC_0V8BnZPOJokR1UW6mcwr4,8019 +torch/include/ATen/ops/mkldnn_reorder_conv3d_weight_compositeexplicitautograd_dispatch.h,sha256=sqchHmDYYowUUREpxPtvdJFL-upuEj8NoabFHOcU46I,2011 +torch/include/ATen/ops/mkldnn_reorder_conv3d_weight_native.h,sha256=j5bvlXT9O-6sh-yIcaL394gCS3ZFR9a7IzEt4xTambo,1155 +torch/include/ATen/ops/mkldnn_reorder_conv3d_weight_ops.h,sha256=PsQ6HMj_uBqtUOOu4xLfn3mp9RkhOLQiaLfspdyxRH8,2939 +torch/include/ATen/ops/mkldnn_rnn_layer.h,sha256=wGploXOeZ6vPc_0S5YUiTxTs7buQ8QxThMjiMpnz9KI,4009 +torch/include/ATen/ops/mkldnn_rnn_layer_backward.h,sha256=cJUWpRoPXWPdNFmbOOrl0O4tb472A12DVVhWfpEcsaU,5833 +torch/include/ATen/ops/mkldnn_rnn_layer_backward_compositeexplicitautograd_dispatch.h,sha256=p8HQltvfWMF7-sTlxAJETLq0-DjWXa7Dsdu_lfbuvA0,2719 +torch/include/ATen/ops/mkldnn_rnn_layer_backward_cpu_dispatch.h,sha256=yCHGaeZT94CRjIDg_XriGHKWkYNanwzngUZHMco75ws,1646 +torch/include/ATen/ops/mkldnn_rnn_layer_backward_native.h,sha256=TIDqXj0LwGAg1KqYP_IlBAkRP7EnHsCbfMFJ27d9NDM,2279 +torch/include/ATen/ops/mkldnn_rnn_layer_backward_ops.h,sha256=Ot6BYMP6RGYpB0y31LdKKFSWBkp1bCeeJzgHTZ5K6dU,6728 +torch/include/ATen/ops/mkldnn_rnn_layer_compositeexplicitautograd_dispatch.h,sha256=OxME3sMYUcTJpf5jdvLflrSajkdKDoObkKHqyegiJds,2015 +torch/include/ATen/ops/mkldnn_rnn_layer_cpu_dispatch.h,sha256=GS_3o9m7W9VP3jZG2Be9DK156BPIXRyszLL7bo-5xNw,1357 +torch/include/ATen/ops/mkldnn_rnn_layer_native.h,sha256=LhohnHmsehGFnqdd5NTYqfm4heTgsZPteThbd5VT254,1638 +torch/include/ATen/ops/mkldnn_rnn_layer_ops.h,sha256=7EKfwpETw9JMidSN_3u5Xkw_foMBUFPyp4iYw0VWrbc,4591 +torch/include/ATen/ops/mm.h,sha256=3gNlWwyefOpkbyUgTZKoHB6S5_mAZwutJ5fp334sHdc,2207 +torch/include/ATen/ops/mm_compositeexplicitautogradnonfunctional_dispatch.h,sha256=_MUXd_6IczxzeMOszOY0SItU3wsVwZkUPOtKvkckOrE,1065 +torch/include/ATen/ops/mm_cpu_dispatch.h,sha256=E1VTfGRHoiGA7E2iwW9jmhBxv73CufYRFxqEZO0vXVc,1194 +torch/include/ATen/ops/mm_cuda_dispatch.h,sha256=iIkQ6zOSmo7KgRLdt6sFaL9qvbYluS8ZcRq4u4_eK-0,1548 +torch/include/ATen/ops/mm_meta.h,sha256=3eB7FnXI07W3vtgPCWR8LJbrmjXYFVWMZ2y8YMiCMFQ,842 +torch/include/ATen/ops/mm_meta_dispatch.h,sha256=eKK-VTqzq3n1yuISj8NlRiEblX0-g6yYTMUTMaAjaMo,1196 +torch/include/ATen/ops/mm_native.h,sha256=cr4QHsey_2OvORLJyOtK93q66XFFYHYlvcxOsfHUCa4,1668 +torch/include/ATen/ops/mm_ops.h,sha256=SysbA_e3XcMy8NxS3ja3LHkFQ1NSqwEwd-q46dXLXbo,3390 +torch/include/ATen/ops/mode.h,sha256=K0r37tfo88C9N7MS95ZKKusaBABu6wUvisNH-zFcouw,2875 +torch/include/ATen/ops/mode_compositeexplicitautograd_dispatch.h,sha256=-FdiEkK6Xejmz39XPsA21HbUncnuflaV_c-a1oM_Lkw,1284 +torch/include/ATen/ops/mode_compositeimplicitautograd_dispatch.h,sha256=tqrqnENPxHuFSa--VvVL686kXXJ9_oteJiXFo1Nn7AA,1403 +torch/include/ATen/ops/mode_cpu_dispatch.h,sha256=7BzFAZ5jENU5rwlGINI0U7hcDS67GDrF3C16w70D62o,1033 +torch/include/ATen/ops/mode_cuda_dispatch.h,sha256=I6zdftYd9saNZLaUEmA8hOkoJ9aiVapbFYAqz1XiXlM,1035 +torch/include/ATen/ops/mode_native.h,sha256=RT0DGNc4coq-H5R_cdHlXh21Nes0P9hYXCugmzjQqmI,1217 +torch/include/ATen/ops/mode_ops.h,sha256=t4LoES896Rj2UH88eIu2MhaMPyJfDPxrqQy3fHxgs8g,3947 +torch/include/ATen/ops/moveaxis.h,sha256=fbQy1NErRcQmsqxP6jlnxPPFuW0bT-JDI6Xlba1enNo,1270 +torch/include/ATen/ops/moveaxis_compositeimplicitautograd_dispatch.h,sha256=oALMgumCd1Utd5uQNYiJkNYyN7Xg7D6pqJzJ6FIyJG8,1166 +torch/include/ATen/ops/moveaxis_native.h,sha256=Rh0Rgl31ziUhSP_dm-D1FwFwHWzCv67lV0IdGsBkAaU,878 +torch/include/ATen/ops/moveaxis_ops.h,sha256=UX-3CtHhkBFbCAhm01o--SPaGohLglOitnmNrGeq3L4,2063 +torch/include/ATen/ops/movedim.h,sha256=ElS6Y9ZTZB9Z2I-lf8OD0q1zIyxgYQHOnTU5VYfUk34,1263 +torch/include/ATen/ops/movedim_compositeimplicitautograd_dispatch.h,sha256=YNdI1Qtu_bMb5B_B18IzMybgfuaU1wKQH-2477ETnoM,1164 +torch/include/ATen/ops/movedim_native.h,sha256=8irGWfcmRuCxodVwVfhdbdCfhlYLwwti3YuJuJnojXE,876 +torch/include/ATen/ops/movedim_ops.h,sha256=nsHVJU7dm7b0SVbReuHdq3Sfm_GqwE6f2ko07J8JE2g,2057 +torch/include/ATen/ops/mps_convolution_backward.h,sha256=ILVW6FoagQq9pk34CsUGhC4C0gyu0h9uz6Ae2Au02mU,9207 +torch/include/ATen/ops/mps_convolution_backward_compositeexplicitautograd_dispatch.h,sha256=amGbfVKGH_mLcmKN4R2GPZ1C4NELV4oj-PYx-sOocCk,2448 +torch/include/ATen/ops/mps_convolution_backward_native.h,sha256=lJjtscC8hACvHQ7SJ0672UHoXBeh8Fi4pabq-IZCYYw,1058 +torch/include/ATen/ops/mps_convolution_backward_ops.h,sha256=Gvx5GoBoL0SicHjCQat4yzDd46jd--tbh9O8KZAzfDI,3648 +torch/include/ATen/ops/mps_convolution_transpose_backward.h,sha256=8QkQuKhkLcBK_-DFg1RUjvh0zjQKGQFDebxUYfsbwH4,9943 +torch/include/ATen/ops/mps_convolution_transpose_backward_compositeexplicitautograd_dispatch.h,sha256=-pthlgSODIFkPj79UlPXXMRYN0NQ3BMfKSkFlNCFyK0,2496 +torch/include/ATen/ops/mps_convolution_transpose_backward_native.h,sha256=GuM9pOUMyLIs2Q7OTvgmx4VAVpQ9yHkviYF-emb-low,1072 +torch/include/ATen/ops/mps_convolution_transpose_backward_ops.h,sha256=JjPIwJQNbg9lY14m1BDmHi5z5JjULPTKidnV1HeJr-k,3783 +torch/include/ATen/ops/mse_loss.h,sha256=d3zyA-zWvT5r8_Obc7or1234nxnrrAJDb7Ea4al2_TQ,1641 +torch/include/ATen/ops/mse_loss_backward.h,sha256=NVivTBkXJdidUr_xqCXXP5SCIIz0HJdHXhqGc8PVcCc,1941 +torch/include/ATen/ops/mse_loss_backward_cpu_dispatch.h,sha256=ebajBXon_udoiNayYNZ3p0cN_KWrvUchrXhpghY0BHA,1412 +torch/include/ATen/ops/mse_loss_backward_cuda_dispatch.h,sha256=DdKdbBrWwQaQw93x92Et3zBIIZEf0tkrGgGz7IfXySA,1414 +torch/include/ATen/ops/mse_loss_backward_native.h,sha256=okIkz67pz_mWVoIBCF-LMmnsUObLPHooqPf9BankCUk,993 +torch/include/ATen/ops/mse_loss_backward_ops.h,sha256=WZ0O3LggZbxqFR06P_kXMg2YLvVMkwe0nnNGbBZW7MA,2419 +torch/include/ATen/ops/mse_loss_compositeexplicitautogradnonfunctional_dispatch.h,sha256=hbBPDxztvKYcZdKCEl-pJJjH9Hf5rcBFQFwJRVCrpCA,1112 +torch/include/ATen/ops/mse_loss_cpu_dispatch.h,sha256=6s3t4cg5mmsiZNa28QwHNZRzsWX_ckwydJKzdJmATTY,1315 +torch/include/ATen/ops/mse_loss_cuda_dispatch.h,sha256=PsTQSwQTq_F2IAqw4hAszQ_AN6ctpLzoVTMHoR6202E,1317 +torch/include/ATen/ops/mse_loss_meta.h,sha256=C8ZhXv5IYxm5m9MQHktQbbXPho9ozjcdS37ah49biMs,869 +torch/include/ATen/ops/mse_loss_meta_dispatch.h,sha256=UhdxJ4Je1IoM4_G52jWpE89UwquiT9afR3F8mde8wwc,1317 +torch/include/ATen/ops/mse_loss_native.h,sha256=YznOoXHkG8guJ576fRAUMCXav_mpG-AIFdsvdXQ4Sic,902 +torch/include/ATen/ops/mse_loss_ops.h,sha256=JpykFUJmAhS89WSlaLx4kIcm6UTzd0srKwPs7CZwTng,2125 +torch/include/ATen/ops/msort.h,sha256=R592S5elfhm5sm3LM-Ym4RJbmUmXZpnVHflq1bd0ho8,1271 +torch/include/ATen/ops/msort_compositeimplicitautograd_dispatch.h,sha256=mOzWk66yCjthOLUDBZ-i9fMDW9yXVQXDSY-0CKzfw4E,1172 +torch/include/ATen/ops/msort_native.h,sha256=toLQqY21bbV_-VJXFpMZNAtmN0CuhzuhiBKYuTIxWGQ,806 +torch/include/ATen/ops/msort_ops.h,sha256=EO5GAjX9CsRVCP7hZ9UDKhnUXPC-2Mz19WLmuJRL31A,1795 +torch/include/ATen/ops/mul.h,sha256=TZrRwMayTY6in94-ahIX4mbX3pcVI6liiqLZCME1L2w,2087 +torch/include/ATen/ops/mul_compositeexplicitautograd_dispatch.h,sha256=uLldtxYv7j97v18gB3YN493P5158pyvQ-2b8zQkciYs,1318 +torch/include/ATen/ops/mul_compositeexplicitautogradnonfunctional_dispatch.h,sha256=AiX2Jrqfzzx9O0pK5B4-2rH7v9xT9fhiyBF9165VToQ,1141 +torch/include/ATen/ops/mul_cpu_dispatch.h,sha256=6oyNCIxgTRYBXKEBDcwCRptSXUxxYfKZtm4LQD0Y5ws,1274 +torch/include/ATen/ops/mul_cuda_dispatch.h,sha256=Ha6FThK3eHVax9shiTb23M4uSfyT5z45hRvBhJikbek,1276 +torch/include/ATen/ops/mul_meta.h,sha256=a7uN2fgODrWiJKGAKOQiNjWeecyHeyET7dneTJRa0Qs,851 +torch/include/ATen/ops/mul_meta_dispatch.h,sha256=qzRm7kDSkrVTlxLCnW59qaXEtYtUxLTjITrunuzGGXc,1276 +torch/include/ATen/ops/mul_native.h,sha256=K_RpZO2msxl445HwfQs3DAvox5oTBhwXyCxwqRTZ_p8,2738 +torch/include/ATen/ops/mul_ops.h,sha256=qpqda27m5uZCWIJ0BWeILqyN6aJ6UshoMxi8-BEqnDg,4452 +torch/include/ATen/ops/multi_margin_loss.h,sha256=zAzeoHQ6x1C6B3eR8pGA14gNagPun9Emg5DxSGcRTcc,2231 +torch/include/ATen/ops/multi_margin_loss_backward.h,sha256=FtBhkrO0_EqNyqcs3OTn6L9pHf-5PvwgdEpfCUGodrc,2566 +torch/include/ATen/ops/multi_margin_loss_backward_cpu_dispatch.h,sha256=qJWmxPBhZv_8xep8KUlVu_Q-cb0aX6SZvowp9oxyD1U,1764 +torch/include/ATen/ops/multi_margin_loss_backward_cuda_dispatch.h,sha256=NXJUzr1RV5P_acKjWIj9NxmrQhezXnl1h9rCEx_mPd8,1766 +torch/include/ATen/ops/multi_margin_loss_backward_native.h,sha256=0EDRGdU4awnD9q7-5BB9gJzNnqlDLs9yoh1p4c2k6D8,1782 +torch/include/ATen/ops/multi_margin_loss_backward_ops.h,sha256=O0tSVWVWWu9PASH9k36yzpCM1GmaCnp5UdNkgFIVZXw,3101 +torch/include/ATen/ops/multi_margin_loss_cpu_dispatch.h,sha256=Dgflnkt7O25XqWjpPySBqP-4A0cXco4jm7-zFSFhHms,1635 +torch/include/ATen/ops/multi_margin_loss_cuda_dispatch.h,sha256=_BYLFxzYACB7TdtPX7oLosEYuE5iSykV7-k7Ktbjm08,1637 +torch/include/ATen/ops/multi_margin_loss_native.h,sha256=gGPXZYEu9m0ZUnX3SBaeZci1mynFLxCRhBjjlvRY9-4,1612 +torch/include/ATen/ops/multi_margin_loss_ops.h,sha256=ChKiNLCwCD2uFo1oqXBVdln9rQIHEtCdRlUunBF1oJg,2805 +torch/include/ATen/ops/multilabel_margin_loss.h,sha256=WNfTo7hf2gqNzwKw35RG29_tftWP_HzCQMdVI3OYw38,1781 +torch/include/ATen/ops/multilabel_margin_loss_backward.h,sha256=hFto8jPx8i_q6HxE8MUA2dW04ZEM7kFDgv9XyMiKygw,2258 +torch/include/ATen/ops/multilabel_margin_loss_backward_cpu_dispatch.h,sha256=wwSj0wSG79lTipGiY0hPhXmhK2bru7Z3k2bPGOHV8Aw,1544 +torch/include/ATen/ops/multilabel_margin_loss_backward_cuda_dispatch.h,sha256=2Y3lNW2kPka02S0DMC4wgcUXuBVCUHdqZywtoHXsdQE,1546 +torch/include/ATen/ops/multilabel_margin_loss_backward_native.h,sha256=fYxqS1rBi8YqZfenbF-OF0tE0TfHxN7qcCt3io87MyM,1504 +torch/include/ATen/ops/multilabel_margin_loss_backward_ops.h,sha256=8QqxxS2IowTgWMIkzDiBoJj04XA-H-fzxV-KqTOaJaw,2699 +torch/include/ATen/ops/multilabel_margin_loss_compositeimplicitautograd_dispatch.h,sha256=cL8fTJkEPKN8uHaQofRsH1k3OlQiJsI9LrgrLXf3A_o,1401 +torch/include/ATen/ops/multilabel_margin_loss_forward.h,sha256=r9d3GsLRUPIRki9n6u3xjqzbwJHyJGxEJbkAOBOGERY,2084 +torch/include/ATen/ops/multilabel_margin_loss_forward_cpu_dispatch.h,sha256=e8Zf_U3W6Sn6YZZ7tSh8MMYy_WxE_Um1IZf6ZDN03-0,1474 +torch/include/ATen/ops/multilabel_margin_loss_forward_cuda_dispatch.h,sha256=v8PslDRuNIVBrBNl03IOHmn8QhHGm4FfIM2mFlWveSM,1476 +torch/include/ATen/ops/multilabel_margin_loss_forward_native.h,sha256=XYxRT66GEfQPjjbBEIxslKXDrEm2MGAm5B4XYZVfUFY,1396 +torch/include/ATen/ops/multilabel_margin_loss_forward_ops.h,sha256=ezfZOZzgCnHBedA8vqW1xxUa9Y0-JLW8B7plDUx6Wg4,2546 +torch/include/ATen/ops/multilabel_margin_loss_native.h,sha256=TqzMIPglGvX_BKlbnromzyoIy8NqxqGmvQHlM5ndl_E,952 +torch/include/ATen/ops/multilabel_margin_loss_ops.h,sha256=Odq5_5tZlfKn1FJG-th3mg4goOKNTeImWWMF5pH4Cjk,2209 +torch/include/ATen/ops/multinomial.h,sha256=yvGZ3Wfvx-LNB3W8L9-5u5eyjSQ5BzQpLLFa205MKrE,5334 +torch/include/ATen/ops/multinomial_cpu_dispatch.h,sha256=cBhHKfdfaGhCIt-EjI1DT4ae154qRvGaAUj9lPwy_WI,1975 +torch/include/ATen/ops/multinomial_cuda_dispatch.h,sha256=Y4xIh2j0CsQ7yhArjERsDsFVXu4dcFKlryACsEvVj0U,1977 +torch/include/ATen/ops/multinomial_native.h,sha256=q1EGnmrsM6gOX7h5qM-rNptb477UMAVz5oYbHxzoBw4,1001 +torch/include/ATen/ops/multinomial_ops.h,sha256=RtShwTYLoVTcg_9rkW51h57cJL8sJynVRpSa9pNAwWc,2418 +torch/include/ATen/ops/multiply.h,sha256=u8mzBJ5Y9R3MET5lkEibB_MP2iG73DcjivoHlspktQw,1658 +torch/include/ATen/ops/multiply_compositeimplicitautograd_dispatch.h,sha256=GSLt3bpNDxmaMsnHZH4CZhHTW9MyyUdkzFBfrsuD3aE,1499 +torch/include/ATen/ops/multiply_native.h,sha256=v6rewbsVV-JL7sZjykUMcDvAIQnD_6ET2oAQyFUm-YY,1104 +torch/include/ATen/ops/multiply_ops.h,sha256=p5HiVpqSsFWuJHmvbqVr4MxwpZdllJZxvXeGpF2qpTg,3839 +torch/include/ATen/ops/mv.h,sha256=X5lykqbK0XT2Z85XK8U02Ay4h6qHYqY3uKX3gB96vzs,1364 +torch/include/ATen/ops/mv_compositeexplicitautograd_dispatch.h,sha256=BXYidWGL-TNbaMqTT1u4EBPI32vylMdZ7wZVLcsTvko,1235 +torch/include/ATen/ops/mv_native.h,sha256=OfOjmJbxiJqxu0JGLfVarAuelBgw9C20VodUFB3nhgo,929 +torch/include/ATen/ops/mv_ops.h,sha256=51P5nk0LLnKoVXGffMFaNiz0FVb_JhQojRYDIoXT_J8,1937 +torch/include/ATen/ops/mvlgamma.h,sha256=m6JOUnsBe4QKCDIsGsKlyp5UawEym_DSjfGRfENX-0Y,1364 +torch/include/ATen/ops/mvlgamma_compositeexplicitautograd_dispatch.h,sha256=GeeB4-qTTs3kSrXH5VqpHJJ2HEy7P3VIoNBCVzocBX8,1095 +torch/include/ATen/ops/mvlgamma_cpu_dispatch.h,sha256=o2rq5NmL0mJvZhzHJ6m3xNuZa8Uo8M1Ur_tJtnQ_6TI,1103 +torch/include/ATen/ops/mvlgamma_cuda_dispatch.h,sha256=rLiquDuthXbayEo3J9zhwcE4eBF2LnqqVeUrNnlPpX0,1105 +torch/include/ATen/ops/mvlgamma_native.h,sha256=ywyVxp8jv8SekLF2dsQnnbNP9gfQHcKTmrM1dpLgNs0,898 +torch/include/ATen/ops/mvlgamma_ops.h,sha256=SXURNS8K5isyJSCqiZZJrTw2JcvN0ni-wendKK8Y-fY,2432 +torch/include/ATen/ops/nan_to_num.h,sha256=MPSKQINhS8_B0jQck-66UhFTp3nOQp94rILzEygMxBw,2294 +torch/include/ATen/ops/nan_to_num_compositeexplicitautograd_dispatch.h,sha256=R4PKDMtv1LVDuHnUARlMWuH15UPh3SixT2tzZCSRsuA,1353 +torch/include/ATen/ops/nan_to_num_cpu_dispatch.h,sha256=4qOBMWh7zintG-vOen7HkpVSaV2OKECVDUEQ219TEHc,1316 +torch/include/ATen/ops/nan_to_num_cuda_dispatch.h,sha256=XHJ8w5f12wdTTUBxPqPPP-b78WxDRwc6L2hX0sVJ-18,1318 +torch/include/ATen/ops/nan_to_num_native.h,sha256=u8sucp0X8lQ7hnrAAKVnpugMobhVCpQ7MOSd_W9aaf0,1825 +torch/include/ATen/ops/nan_to_num_ops.h,sha256=WpftTchVJLbDXSUSJVRCri0qPERPbTs1x1-tgpcGBtA,3290 +torch/include/ATen/ops/nanmean.h,sha256=-JVLwwWhWP1yOhrqJCRAzH-15A7FTqndtuHmnQ1UQDE,1861 +torch/include/ATen/ops/nanmean_compositeimplicitautograd_dispatch.h,sha256=9Y1VMhVvEFFr9zfVNGk4w74SBr6sn6ExicUq4fiQbDg,1496 +torch/include/ATen/ops/nanmean_native.h,sha256=g0eHuBY2qam-D3O07mm0OWDzWQhV076PzBcw_J5m9Lg,1010 +torch/include/ATen/ops/nanmean_ops.h,sha256=gIePLM1Z9BTiywnmpk2k0bQ3vgGEtlTFOXC2coD0yWY,2390 +torch/include/ATen/ops/nanmedian.h,sha256=ISIhvgJmblN-3LJAWTCRxsyTjkufx2a7UpUxyce8XD4,3555 +torch/include/ATen/ops/nanmedian_compositeexplicitautograd_dispatch.h,sha256=1zGv1Zk2NM-4UI14Yl1jX7C-6MoBE37MlkZVRGqOU_Q,1242 +torch/include/ATen/ops/nanmedian_compositeimplicitautograd_dispatch.h,sha256=1c6RG5O7Oa31FBnTzwZLjKE0Li-HX5shhAVTyZOghuI,1418 +torch/include/ATen/ops/nanmedian_cpu_dispatch.h,sha256=9gb-4bWyN5MFMflxSj05m8_6_T1HmsAKBukOUsm9Lpw,1304 +torch/include/ATen/ops/nanmedian_cuda_dispatch.h,sha256=wCuL0EdBMWFjx0WcfQ3fpi5sAav43Q-eAcBU4t-KeP0,1306 +torch/include/ATen/ops/nanmedian_native.h,sha256=gAgOtslf6DblxztciLz7qaVRMvwgSi5PkwuM1CNHzpI,1607 +torch/include/ATen/ops/nanmedian_ops.h,sha256=wzk9sATnW8XUVkBDxxgjD-RyCi9LtiCkYb9NMMnxvVU,5153 +torch/include/ATen/ops/nanquantile.h,sha256=f5NBFtj_EfyLbL7-mK43lrRcWY10cE2O63HiIRMbZ-g,3248 +torch/include/ATen/ops/nanquantile_compositeimplicitautograd_dispatch.h,sha256=Pz830ni4hTtpT-FLXiBqzIUOeinF2g_CKEtE0DqcLoU,2088 +torch/include/ATen/ops/nanquantile_native.h,sha256=MJNMphCjf5I8S_EWrKzVWIyZnkXwL2pbYPgvl26JZHg,1388 +torch/include/ATen/ops/nanquantile_ops.h,sha256=tHuI4Aj676CrB9ovOv0ZiDzLAu0NX4UntD_TtCHpcpI,4280 +torch/include/ATen/ops/nansum.h,sha256=Le-a87c-7MTQzve36XxzZvULVqKSHtCsJ5fxjIims1A,1851 +torch/include/ATen/ops/nansum_cpu_dispatch.h,sha256=qaNgMG6w3CVYSdlICbrB_KvRYNivfduMdjzuHYb1hEc,1449 +torch/include/ATen/ops/nansum_cuda_dispatch.h,sha256=pBHoKGkarU8RtlgFOVVbTvkhVOoVs8xzkzEiFo5r-qc,1451 +torch/include/ATen/ops/nansum_native.h,sha256=jnP2Wa0OIpnQMQfbDEwzdBeA1gmrU9J2byCeHjN-87A,1008 +torch/include/ATen/ops/nansum_ops.h,sha256=l6PwBLAa2MJ83EEuPhuurIAFrY77fFTKb8BCV-ztNXk,2384 +torch/include/ATen/ops/narrow.h,sha256=tuPziGHNvbVggBDD_R2drMWatAj8NieDTfntd2SRedc,2892 +torch/include/ATen/ops/narrow_compositeimplicitautograd_dispatch.h,sha256=D0mlQLjVH8CQk6aYRXFqQ20W66qdIqcVIizSs1ngVJM,1404 +torch/include/ATen/ops/narrow_copy.h,sha256=IwqGzoGV5yFBz28doHMLvk22YWkeB4Gu8iOCQKWLCHk,4338 +torch/include/ATen/ops/narrow_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=TjYOqzvt_fEejTu3SaiZKI52mNryhBT4QzZIe445kB0,1211 +torch/include/ATen/ops/narrow_copy_cpu_dispatch.h,sha256=amjkFzLO1lz7sigD2tmN0WNW6_VPkLLVIbQDMqXwUw0,1681 +torch/include/ATen/ops/narrow_copy_native.h,sha256=XMmRZtyirmLUUOrlSpe9CE6yPMeM469k7t6ltUdbIfk,1160 +torch/include/ATen/ops/narrow_copy_ops.h,sha256=jjUhA5tpoxF7Quqcyo4DkhcQK5UmqQEZuCU967abvQs,2185 +torch/include/ATen/ops/narrow_native.h,sha256=pZnXu6F69N2fBKtevRgAe2vIxzVJG_7KclFwmtqvbyA,1036 +torch/include/ATen/ops/narrow_ops.h,sha256=K5PuDOQxlSQjieSPovjSP1uEsN3Yws6tKwCGr3WotUI,2118 +torch/include/ATen/ops/native_batch_norm.h,sha256=XUgIXrw1RoKY3vx35RZIwGvZpXusT4D399Vnu9KhM8M,3056 +torch/include/ATen/ops/native_batch_norm_backward.h,sha256=pucNYvGdf9x5SJ72wKSW8bfN0kGc0HB2Hs-DUGWRjOE,3590 +torch/include/ATen/ops/native_batch_norm_backward_compositeexplicitautograd_dispatch.h,sha256=LhkCiqvo6E4mImXqvrpH7ZtojM9kccVKEhINFWUNzIM,1973 +torch/include/ATen/ops/native_batch_norm_backward_cpu_dispatch.h,sha256=RcS2pqc3Kr1dazG2uLkx7r5M7xPzUkkMtcnKWbhJktI,1357 +torch/include/ATen/ops/native_batch_norm_backward_cuda_dispatch.h,sha256=g8eBrQ81B97joeVOptxD9XM7CGOPIT8qV30VDOXnKMc,1359 +torch/include/ATen/ops/native_batch_norm_backward_native.h,sha256=5_LkJ_t8x6En2KvGn5aS2zx_pTIOs75_e6729wbrusU,2486 +torch/include/ATen/ops/native_batch_norm_backward_ops.h,sha256=GUY7V6JVOCmd8SSvHRmLKhWTOGRYQDs87P8AO1Cuw4A,4446 +torch/include/ATen/ops/native_batch_norm_cpu_dispatch.h,sha256=_NVyNLSH2Ift5Dr5582qLMvdRy9bhdaCdKXFXsnwo5I,2070 +torch/include/ATen/ops/native_batch_norm_cuda_dispatch.h,sha256=7F4Bk-SMY5kSDvp05k7IEjImDzypOz5tqdWg0OG2ZmA,2072 +torch/include/ATen/ops/native_batch_norm_native.h,sha256=aoOAbODrWN0-rlWrfddua7Wz0vvED4HXEtvhQEDTsMw,2477 +torch/include/ATen/ops/native_batch_norm_ops.h,sha256=X7udntMza0uIr23VXtNcILRPdLMDrFmXXqwiN6tzWMg,3805 +torch/include/ATen/ops/native_channel_shuffle.h,sha256=KzCSZyniWqrSsNMTdGfByce8K9JgLY7WhxNPcD0pMtc,1735 +torch/include/ATen/ops/native_channel_shuffle_compositeimplicitautograd_dispatch.h,sha256=9Xh4jBZmd6E0OBeRSLDZ37ZhKibq22mRBRD7Gt-7buQ,1147 +torch/include/ATen/ops/native_channel_shuffle_cpu_dispatch.h,sha256=dKAC3JYJ9q5FwUTGOyCmTV5duzDIvOoZM_rIn2YqHHc,1103 +torch/include/ATen/ops/native_channel_shuffle_native.h,sha256=LIQO13Xn_htgZOHy1KHEdJJ_6CnIEe3NJdGI3Lh4LOU,843 +torch/include/ATen/ops/native_channel_shuffle_ops.h,sha256=qnDguQrATE0aiHSOR1tiLNM4wVCXJPMctNCy5TKTs68,1327 +torch/include/ATen/ops/native_dropout.h,sha256=vRtZxe8ojAtMMDCV4LcQjPfP8ZypmCfhAtMe3lX4b_k,1790 +torch/include/ATen/ops/native_dropout_backward.h,sha256=GSHGwW3U4mw8Vai0US3oXiDOgO6nkx0UfuTp0xyGvew,1748 +torch/include/ATen/ops/native_dropout_backward_compositeexplicitautograd_dispatch.h,sha256=mPyJ2bGcfRahWboVQ4PuW1RDSU6KCH7IWQTiFiaHmmI,1247 +torch/include/ATen/ops/native_dropout_backward_cpu_dispatch.h,sha256=jt9SR67LMZFNfc2JIVcrOeWcGR8vpbWGovaE965joZM,1037 +torch/include/ATen/ops/native_dropout_backward_cuda_dispatch.h,sha256=jrDqigpz8h1pqAB13KKddu5FsqeKgtIvDIXBWr5ug-A,1039 +torch/include/ATen/ops/native_dropout_backward_native.h,sha256=tuwhBnk7fDrLBq9eIslnKkmNOvUAn3oDB5-GMqYCd2U,1056 +torch/include/ATen/ops/native_dropout_backward_ops.h,sha256=jy1Bfvp9usroxSyM4WVSZUdu0FUSijoERNSFu-P8k0o,2209 +torch/include/ATen/ops/native_dropout_compositeexplicitautograd_dispatch.h,sha256=cW1XJ_vsyVIfY6cMk-PVAQwkl-PWDJx8IMb4QLdV4zM,1311 +torch/include/ATen/ops/native_dropout_cpu_dispatch.h,sha256=vHGXLLMd63nUYGLYI5_og9mrlPqCk4iDh7KqdCCcKX8,1047 +torch/include/ATen/ops/native_dropout_cuda_dispatch.h,sha256=YBDcHuWeBgTwcmd4g1MtjgL2-jkpOIL7JLS5XAGGDfE,1049 +torch/include/ATen/ops/native_dropout_native.h,sha256=Zk0P4LfqwJYmtFDr6oMDFp0POfwgOFNGLM8dib-GprE,1246 +torch/include/ATen/ops/native_dropout_ops.h,sha256=7wUJF9w9LAS0EzuSPbm9Z2HICZXw4klTTXGMoZQUb2Q,2369 +torch/include/ATen/ops/native_group_norm.h,sha256=7Dv2tYVuhWdWP8XbWeTls3NT7_4HW7y7Q_s5Lp2PZ2M,7466 +torch/include/ATen/ops/native_group_norm_backward.h,sha256=0K6VC4WDzHiueA7fq7pBQMa9X9huRQq7ZCx-k3F7zmA,8975 +torch/include/ATen/ops/native_group_norm_backward_compositeexplicitautograd_dispatch.h,sha256=L7cyjYzKqo5z2fJdL9t7CuGZDxgUyp_YtezTj-g6Tb4,2544 +torch/include/ATen/ops/native_group_norm_backward_cpu_dispatch.h,sha256=rC0SorxLN8--x-yGwrKF_TAHiniPKNxIZHoIPd7wSlg,1575 +torch/include/ATen/ops/native_group_norm_backward_cuda_dispatch.h,sha256=Dua1HplB0jVW6brQG0E_e8u7W9DIXwy5pZSR0YJusy4,1577 +torch/include/ATen/ops/native_group_norm_backward_native.h,sha256=XLEUw9MN42ZFOlniHyzHQQf2qtfhJcg1f48aVGlxOCs,1398 +torch/include/ATen/ops/native_group_norm_backward_ops.h,sha256=0LqN48ShBI9Jj7m5-Qww4PIIt4TGNWYhxBIqqc7oUg4,3796 +torch/include/ATen/ops/native_group_norm_compositeexplicitautograd_dispatch.h,sha256=o5PD18tXtO_shBUaAZlDSClpkbj7Hd-E-c7I_nWpJ2Y,2791 +torch/include/ATen/ops/native_group_norm_cpu_dispatch.h,sha256=ni4C5QGUWcIBg_ivqiTUwqd5vkOFPRagzeU2UTu4uL0,1439 +torch/include/ATen/ops/native_group_norm_cuda_dispatch.h,sha256=iwHv89VEkK55zgCDB9yoAWLRpnL5S9zW0VQsx_9IMhk,1441 +torch/include/ATen/ops/native_group_norm_native.h,sha256=yvXkzMk4OoLKV0dqYWqT4PVkS-op3SKIKlgRvcArNis,1510 +torch/include/ATen/ops/native_group_norm_ops.h,sha256=wYzSuV5XSXDasQyasalXJtAnis4cFVZNncaMj1OIpGM,3354 +torch/include/ATen/ops/native_layer_norm.h,sha256=RngddZoFs1Am2SlzoPjbteqLqZuplYCf20-vBgTtk6U,7286 +torch/include/ATen/ops/native_layer_norm_backward.h,sha256=ESHIilI2Bk2QCeyid48ef4f1wncUWzxeKY1w0vC5JXQ,9455 +torch/include/ATen/ops/native_layer_norm_backward_compositeexplicitautograd_dispatch.h,sha256=harJcPHECBPNEGAuVXvdUbfakSyxsw1ZaL6V345xZpI,2632 +torch/include/ATen/ops/native_layer_norm_backward_cpu_dispatch.h,sha256=BSWs45sTtz8k5LUT-6ciWwBf5VpeKXefKaKEgXekrOk,1619 +torch/include/ATen/ops/native_layer_norm_backward_cuda_dispatch.h,sha256=2B3QtvismkQFb2kjxRJybBcdpJGT3NiZl3RP6M2R87U,1621 +torch/include/ATen/ops/native_layer_norm_backward_native.h,sha256=s-0NgBzl23wiY2D_Isj22KojrxGcZ7DPPm8DPQ88h84,2125 +torch/include/ATen/ops/native_layer_norm_backward_ops.h,sha256=u2aoVn7s-csznVdfxumDmZxDrP-sk99bFYseRUigUig,3884 +torch/include/ATen/ops/native_layer_norm_compositeexplicitautograd_dispatch.h,sha256=Gv6S46OolGSlhSw4SMoVD8yt0txehrxWqXuAqNhU4Gg,2671 +torch/include/ATen/ops/native_layer_norm_cpu_dispatch.h,sha256=IF72vS4T2v16qCBsSWG1FsevKusABWQ4jKENZjiPY5o,1399 +torch/include/ATen/ops/native_layer_norm_cuda_dispatch.h,sha256=YzrvBXZ40HJSwLH2q0XiJLjXV6JCyirWFiX_sa6fqUk,1401 +torch/include/ATen/ops/native_layer_norm_native.h,sha256=-aPAmyEZPpDlfW5WcLDq_4pCuh9bnnBni64opcS9DEw,1924 +torch/include/ATen/ops/native_layer_norm_ops.h,sha256=yvV9yx6aP5P0cBvw3M7Fo566ZZramn59D8VZycRRDDw,3172 +torch/include/ATen/ops/native_norm.h,sha256=NmEcCyc71ydiFy1gZRImOUNB88NPiPPtRp5xVc4aXwk,2715 +torch/include/ATen/ops/native_norm_compositeexplicitautograd_dispatch.h,sha256=gVJRjs55h9DUjt1ZsHYvtyEf9OK_RF-JnZSJWSjvrM8,1570 +torch/include/ATen/ops/native_norm_native.h,sha256=Z5mssLD42nSd-4pTM53U7NZCgI683FHeDWoBOKCdvyg,1252 +torch/include/ATen/ops/native_norm_ops.h,sha256=w7HE1Qo-TOPIcFjqqNmowEOItPnrvrY-4w0FAP6BSkQ,3972 +torch/include/ATen/ops/ne.h,sha256=-DdzR1To3wcnzjWk71AD5gXbGGsbiqwSzBeducMFZQk,2096 +torch/include/ATen/ops/ne_compositeexplicitautogradnonfunctional_dispatch.h,sha256=HPaKzP9g-5lPN4YCJDcrA4Nlo5dwWXtl0CmFzxdMyEs,1288 +torch/include/ATen/ops/ne_cpu_dispatch.h,sha256=DoK7b4q2wTo3VIzT5MDA1zIVXLq-pd-nlc-aTLq5p5w,1620 +torch/include/ATen/ops/ne_cuda_dispatch.h,sha256=KeZkxA4X-5uLA5llipDn9dzgY0_CXhHuBF6zEjtwZ2c,1622 +torch/include/ATen/ops/ne_meta.h,sha256=vi02e6F5EIBfBBM1hprGjzkmFYWyydJ84RZvZQcLxo8,989 +torch/include/ATen/ops/ne_meta_dispatch.h,sha256=ElrVXrJagst6mxw4gHAwYXpERTC0eWda908NCEommb0,1622 +torch/include/ATen/ops/ne_native.h,sha256=l8QcQikr_q6Z4XelASLloRheKJH8m1q43m9gbX77-Z0,1459 +torch/include/ATen/ops/ne_ops.h,sha256=QlQ-L6e4pew1Z0OZHW5tV5Vd1121AM8mgnGOKiRDQdM,4455 +torch/include/ATen/ops/neg.h,sha256=3dil3e93bYN1znPwnJcqO_b4Peu4wlLkQ8GupsBgXLg,1384 +torch/include/ATen/ops/neg_compositeexplicitautogradnonfunctional_dispatch.h,sha256=QKZyMJQouMIyjcrLfVbW0vCY79b6Fd1BMN63doqjOj8,1089 +torch/include/ATen/ops/neg_cpu_dispatch.h,sha256=hVkvXUwRKuGzVIxHtdx2NsUn7ZmLNyqaF4KqpqmCruY,1170 +torch/include/ATen/ops/neg_cuda_dispatch.h,sha256=6ykp4X0Fl0oMxoiyxORWZVCt3GjztdWJ5GMRGP2Jc0M,1172 +torch/include/ATen/ops/neg_meta.h,sha256=57nq1B29KTXV2XD3RxU7dei3e10Va_X11WBeK-1LQjU,818 +torch/include/ATen/ops/neg_meta_dispatch.h,sha256=DoHY3wKyUrOYLrxC1xypYFYhNxMk_ztuy82J6HhZxvc,1172 +torch/include/ATen/ops/neg_native.h,sha256=JV7eREAa3hMs2RlTkwIM8rYTzjokJPAXMicr1SDnNmM,1368 +torch/include/ATen/ops/neg_ops.h,sha256=hvRtHi2LAVleZ5nOC1m6Mb9XIXh4yo2FNYih1ypJWkw,2273 +torch/include/ATen/ops/negative.h,sha256=8Oug72iYvfRcQuDGOKu5bhBOAzCE4pjV4GxLJw0KSRE,1449 +torch/include/ATen/ops/negative_compositeimplicitautograd_dispatch.h,sha256=agjot52602rW0eBazmt6PnGfzapX05YhdxAEtDb-F20,1234 +torch/include/ATen/ops/negative_native.h,sha256=ldt_g8Rtmyuxt4F4aAJ7cywNdKFh_7MhixuNZr_8jU8,865 +torch/include/ATen/ops/negative_ops.h,sha256=xxg6662VE6dSvnyLTaW2xAfrLbU-T9Y5rRE7nr31GFI,2318 +torch/include/ATen/ops/nested_to_padded_tensor.h,sha256=9MigTsKd9zLcSJvor8X4vzkgRhHh-HowWPcbvfeJJwU,1089 +torch/include/ATen/ops/nested_to_padded_tensor_compositeimplicitautograd_dispatch.h,sha256=B01U5BmlHOcvRzMpx3ou91HkfUIHbe4s_oxMqtjnwTU,1103 +torch/include/ATen/ops/nested_to_padded_tensor_native.h,sha256=EMWQlBojardN6cqNN_rZRcVM0Tz7GFTejEvh5hZbJKE,815 +torch/include/ATen/ops/nested_to_padded_tensor_ops.h,sha256=StRExLhgnm7MaQFTpOkcqQCv_AL_s3lBzRznzoitvxY,1441 +torch/include/ATen/ops/new_empty.h,sha256=UiSxVfxrKLKmN3PVd4zKMNFgOQ_rqziGUk0RhvPKh0M,4589 +torch/include/ATen/ops/new_empty_compositeexplicitautograd_dispatch.h,sha256=xOTBvta-cgz_rqYN7b_p0ICF_TxpJRnDzdStyaBjD8s,2088 +torch/include/ATen/ops/new_empty_native.h,sha256=5q3d43fyXuBo8GMv-gIP_G1oGA2oeNR2QNIZJ37qnMY,1037 +torch/include/ATen/ops/new_empty_ops.h,sha256=L_hnffy4GHwE0WfB60wUwWUPMcqnFts2KkCn4lKK3kM,2491 +torch/include/ATen/ops/new_empty_strided.h,sha256=Z1527DKyNyM1SARoC4BSyigjFDQxiBQxM735eehdOz4,5453 +torch/include/ATen/ops/new_empty_strided_compositeexplicitautograd_dispatch.h,sha256=D7_L84y0YmVCk-NLb7nxXU1D1UXQKttOOGM4wwXifnM,1536 +torch/include/ATen/ops/new_empty_strided_compositeexplicitautogradnonfunctional_dispatch.h,sha256=_jRCpvWQ0erhDCUzI6W5Qxad9xfLOx6HJ_VtdNKCclU,1814 +torch/include/ATen/ops/new_empty_strided_native.h,sha256=uG55t8f2jXqnPB4CMyf99Hw1O9oONHKYR0MIAy3K864,1109 +torch/include/ATen/ops/new_empty_strided_ops.h,sha256=PTmQ3eRrhhoSXwdubya7bNLJYPkfZp9yD2BsuwtF4Og,2727 +torch/include/ATen/ops/new_full.h,sha256=JPCK3LhVB2xej1K61CqCgca2Jr30l-Fpdqu_MQEd4W8,5152 +torch/include/ATen/ops/new_full_compositeexplicitautograd_dispatch.h,sha256=wokdK2NxHVo5lQWW810YkxpgR7WMVOMR-IUCIOE0phA,2328 +torch/include/ATen/ops/new_full_native.h,sha256=UUNJocjpxhhAkUtFWr43dcPT-Lhghm80hpyY0QfXc_I,1086 +torch/include/ATen/ops/new_full_ops.h,sha256=4rOLCQmEU3NhhU8oaHpht5OJYfsUxXzxgPqOT_8mwSw,2687 +torch/include/ATen/ops/new_ones.h,sha256=J0OLEZJMBU8Z61SxNoBcThNfJ0CbdIW-Bpdhq3gqki8,4560 +torch/include/ATen/ops/new_ones_compositeexplicitautograd_dispatch.h,sha256=snVoFVkayeE3O9VTx9Q9t0LNy4YFemtxmpJPVe8xqmM,2080 +torch/include/ATen/ops/new_ones_native.h,sha256=nggNLllVKJIO54JFgPyYBz7A2GR7Tke7VommY7NxZ1k,1024 +torch/include/ATen/ops/new_ones_ops.h,sha256=H1Ap5v0zRaopmBasWy1aZRsWG1NdSk3GxEaH3xWFmBc,2485 +torch/include/ATen/ops/new_zeros.h,sha256=AO_RrsVoT8m6th_-P_nEO7TLno_D3Bw0cvW2Amy6wJk,4589 +torch/include/ATen/ops/new_zeros_compositeexplicitautograd_dispatch.h,sha256=Vm6wDPLdJpQ5_080UJWMXPfruElK9Zi_CZRU8eQHodo,2088 +torch/include/ATen/ops/new_zeros_native.h,sha256=YA8SmTDS5Aw3GYOu7DSsbtx8Wu_15wx7IMM3K81z3tM,1026 +torch/include/ATen/ops/new_zeros_ops.h,sha256=cr7vaGPIm_01hN9NxyW9T52Tn9TM7B0IZqjYUnUPTWA,2491 +torch/include/ATen/ops/nextafter.h,sha256=-TblIa9sAkmMrIH07FumvBgni-eYY1rS2yLkVWXFCO4,1452 +torch/include/ATen/ops/nextafter_compositeexplicitautogradnonfunctional_dispatch.h,sha256=CqlDNXL3s0xWbeHpWRnQYw6n4jWPL35aiAtpzqQHtnw,1153 +torch/include/ATen/ops/nextafter_cpu_dispatch.h,sha256=_9AQ9leL9FwtC6_RIkZN148aKPLGdfbC-YGt_qRwd7c,1298 +torch/include/ATen/ops/nextafter_cuda_dispatch.h,sha256=mLsfFLwWS102x7BB439S16n34AP7g12iCKpVY8pCod0,1300 +torch/include/ATen/ops/nextafter_meta.h,sha256=xuWLppoQXGI27BFIeInsk-J-1VLl8EKyIffHEbXCca0,850 +torch/include/ATen/ops/nextafter_meta_dispatch.h,sha256=K1oQxUp-7-_2BOnPnDYr-tXSSQzUBe6bfQO1EA-4Xjo,1300 +torch/include/ATen/ops/nextafter_native.h,sha256=Qx4IZw_tAf7cZUM10MRusRG80CbMZU0iZxIQrjG9iqQ,885 +torch/include/ATen/ops/nextafter_ops.h,sha256=p59BwVz8Fyp9iUaK3W4zUgcjdQWBNmTKAzbdMlGPHyI,2585 +torch/include/ATen/ops/nll_loss.h,sha256=Eh2wlGFDgU1czBhePODMcYtHt5Hn66k9UDEoh4EdjsQ,5777 +torch/include/ATen/ops/nll_loss2d.h,sha256=hv1ov3XmTex5-V66gysrONmfDsBL7pFBmk3CgqPCBCI,5839 +torch/include/ATen/ops/nll_loss2d_backward.h,sha256=tfHW2Q6IdqdsU62Z5_4d-9vE4H7yw9cEQixkL6S38TY,7378 +torch/include/ATen/ops/nll_loss2d_backward_cpu_dispatch.h,sha256=LLSPFElfKtTdSj0jmY9h_KavQLG7yuvU3Qf4_mFkNZ4,2543 +torch/include/ATen/ops/nll_loss2d_backward_cuda_dispatch.h,sha256=a3OgFiSNm25AGITfttM0DoAdUm_oEWOMhGxaxH7hjrE,2545 +torch/include/ATen/ops/nll_loss2d_backward_native.h,sha256=2hDdI1qkBxcizGsxvSGGsCFmZCKiJFdvA_LudGhWwRc,1732 +torch/include/ATen/ops/nll_loss2d_backward_ops.h,sha256=4ut3Fk3y8gWxaFMyFeSJy9EIuMSfHc0uSh6Gcxe80AE,3099 +torch/include/ATen/ops/nll_loss2d_compositeimplicitautograd_dispatch.h,sha256=-Ao-XeVoPaI0VX8L26yBYfdruti516SNFeq8OA9L0wM,2227 +torch/include/ATen/ops/nll_loss2d_forward.h,sha256=b7pJLGZmQI_6wwJr2ZeKvyIhqb-Bo-SmFHiQ6aMjjew,6729 +torch/include/ATen/ops/nll_loss2d_forward_cpu_dispatch.h,sha256=OJBAxo5dXQrlb7BmG4NfCFZ40Gti7g7-ReUGIKYOzVU,2397 +torch/include/ATen/ops/nll_loss2d_forward_cuda_dispatch.h,sha256=e6f819-h_Hw9f2pxWg0TyrRBlKncj-uPESZWH1RN_ZI,2399 +torch/include/ATen/ops/nll_loss2d_forward_native.h,sha256=J2JTuC3PVSC5MYGQ2_5jH3DYPaOrDj6HLqgLe0H1FFw,1618 +torch/include/ATen/ops/nll_loss2d_forward_ops.h,sha256=y_pUpQe32hfrO6ho1FdK01ZZAw7syVPNmelmBIrBjnE,2940 +torch/include/ATen/ops/nll_loss2d_native.h,sha256=CmBsNyPXCY5lg02NcEwRaMq-yVa-ZQhZGSpb8yeKDXk,1079 +torch/include/ATen/ops/nll_loss2d_ops.h,sha256=RwtIuO3O8xqbzfDdHLxp20NIbTvTs2cUecvMkw0yRTU,2611 +torch/include/ATen/ops/nll_loss_backward.h,sha256=6ehcxwFbbZZo_AMiLwottcix32UIGg1y3SfD_rzHzkI,7316 +torch/include/ATen/ops/nll_loss_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=U-tbjWgkvoWMJ8TtcNd_aUsTdffbgDuUFFqJ3KaMSG4,1485 +torch/include/ATen/ops/nll_loss_backward_cpu_dispatch.h,sha256=n8F8w36MpzcTheQbfzcCAIrM95XnORr06Nvy3zO6bY4,2531 +torch/include/ATen/ops/nll_loss_backward_cuda_dispatch.h,sha256=XkkLWKJs_Ai-M7gbdy8zvT5Q4nFMboNuauwmUC3exSk,2533 +torch/include/ATen/ops/nll_loss_backward_meta.h,sha256=W6hjRyONtqptasAonOE2uf4dW0VKKqPl442zHWGk_4g,995 +torch/include/ATen/ops/nll_loss_backward_meta_dispatch.h,sha256=BADsEwi0ZRGeEwJjiVKlKYbYO6ZbB8NDgglbkYUViKM,2533 +torch/include/ATen/ops/nll_loss_backward_native.h,sha256=EhRGO42dPRbuhgo-hRol6SB57OX2Bgw_U67URlpD73o,1395 +torch/include/ATen/ops/nll_loss_backward_ops.h,sha256=RFOLfAi0pl7qFh81NEp-erLW24pfUit-0XdNzxnHYHI,3087 +torch/include/ATen/ops/nll_loss_compositeimplicitautograd_dispatch.h,sha256=0m16YPlIckHH5kJaO5Ox44mWeNRJr8Kwi5ob79e8WBY,2215 +torch/include/ATen/ops/nll_loss_forward.h,sha256=65Mv2fysKAeynCm64upMcuTVL2nu864JVcINk2kZlak,6667 +torch/include/ATen/ops/nll_loss_forward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=8s8oNWzpYA9gOthksx4kDck7aIF6fwztb1ZHJ6Xu86E,1403 +torch/include/ATen/ops/nll_loss_forward_cpu_dispatch.h,sha256=qrjsNekoUDZG6WwPBrhPfMGwghGfQ9yoIOjEixixXJg,2385 +torch/include/ATen/ops/nll_loss_forward_cuda_dispatch.h,sha256=uol3mOU5iOBl5J-PK7eRAD9S3L-lHyPyEvuqblKar3A,2387 +torch/include/ATen/ops/nll_loss_forward_meta.h,sha256=2ZnRaukgVnOPU9GObu_o6O5fZ5XOLNBfvyoqfIyd4ow,929 +torch/include/ATen/ops/nll_loss_forward_meta_dispatch.h,sha256=C497g9x9nVIdVMPYGLKA7I4jtm5gowu9C9hlA5AUTgI,2387 +torch/include/ATen/ops/nll_loss_forward_native.h,sha256=z4g2_bimglUW7JcrbPrHvJ1_bFAlr-zS3GluAS_O0tM,1318 +torch/include/ATen/ops/nll_loss_forward_ops.h,sha256=zvdDJhijxbRzLtELepLrcYyWNS7rrEjJDcAsmh-xJzM,2928 +torch/include/ATen/ops/nll_loss_native.h,sha256=Nqkdacz598yRVl6ORucTIljHbTK0AShFLVdXm-NhiW4,1075 +torch/include/ATen/ops/nll_loss_nd.h,sha256=nc2oGZGxx_m3sjuXvteHhQ6hYmrtc8LKGgAOPPCkctI,2376 +torch/include/ATen/ops/nll_loss_nd_compositeimplicitautograd_dispatch.h,sha256=MJ80xvsCHxI0gnbKddLIq4ZnsY-GkGwwgG_ggs0i3Zo,1373 +torch/include/ATen/ops/nll_loss_nd_native.h,sha256=J5MyTLIgAGDek7eToZpPQx7MpH8pn5qymRvM5gkK-zw,886 +torch/include/ATen/ops/nll_loss_nd_ops.h,sha256=_65h8l1sc3aYwkedTUgKPvSmANass22xfLrXUAE0OnA,1619 +torch/include/ATen/ops/nll_loss_ops.h,sha256=uEAvyzxh7IIpKOo3WvcWoU0l7UEL5qVHx-64_CARvBc,2599 +torch/include/ATen/ops/nonzero.h,sha256=ysYJ5d6QrJY0rDdWgy8wVgykPLLK0e5CcakBdyB0pko,1291 +torch/include/ATen/ops/nonzero_cpu_dispatch.h,sha256=wCdSTwF5T2xedb1J24bVixd7LefZ6AyRN9vak-jGzUg,1134 +torch/include/ATen/ops/nonzero_cuda_dispatch.h,sha256=0eFxM-bOyLWHyYK9eWT8S0m9rD6baPFA3IYAY2F4mh8,1136 +torch/include/ATen/ops/nonzero_native.h,sha256=yxQ2ql_VknxsnQcJFKRIqc26wuc0iNiUbV7e3KNUpZg,962 +torch/include/ATen/ops/nonzero_numpy.h,sha256=fUSQKqON8495ueRHDw-A_MhGkdQMJCwnOc_MYxBlDHA,936 +torch/include/ATen/ops/nonzero_numpy_compositeimplicitautograd_dispatch.h,sha256=_Wh8rlMiWAJ-OkAn0r2tW1B6Kn_oUI8sMuoUho5wJ4o,1040 +torch/include/ATen/ops/nonzero_numpy_native.h,sha256=W5GHvlWkvnA10lbm7dy7zrE3Dd4wBfEDf7N4Hu8PadY,752 +torch/include/ATen/ops/nonzero_numpy_ops.h,sha256=Ef-Mf9PYa3iFU_rJTsJQgoD--odaa1E_CBQbBVVmS2w,1279 +torch/include/ATen/ops/nonzero_ops.h,sha256=Krm-E4Xzq1EYX39aGU_V2yW06SQ6RXhVZkO_ebcZ1VM,1807 +torch/include/ATen/ops/nonzero_static.h,sha256=C75Toksa3S0iRwBf9CM9n5L5AtJAJjWW2r1OJGyE928,4257 +torch/include/ATen/ops/nonzero_static_cpu_dispatch.h,sha256=CYqoWnNE1O1CGUtYgKup6PO8vSP4fAGf3lBm-dupCw4,1639 +torch/include/ATen/ops/nonzero_static_cuda_dispatch.h,sha256=3H28CXLAYVNpgNpUch_dLEOBWwXKNhbt_Ll1U2no-JU,1641 +torch/include/ATen/ops/nonzero_static_native.h,sha256=EvFNaBExDso_SQT7tkDr4lcOTXH_DVRjRNBCUbY6t2E,1132 +torch/include/ATen/ops/nonzero_static_ops.h,sha256=WCUf3ImI-rDUpy4rfMLvL-aSLHTnWXZtEzmIwpBDwuE,2112 +torch/include/ATen/ops/norm.h,sha256=5tatF-Kg2GNeL9zA3AnjguexQZhmWQIXvGnbQWXMrPw,6516 +torch/include/ATen/ops/norm_compositeexplicitautograd_dispatch.h,sha256=xqzeGOMtHBfiKn3wIp7rvG8jp-MpRRXs7TU3jS5IkZo,1627 +torch/include/ATen/ops/norm_compositeexplicitautogradnonfunctional_dispatch.h,sha256=3JKsNt_7X8ux7dSqqTEMPuh-9vH8QlfIKs1aoMn-xnE,1270 +torch/include/ATen/ops/norm_compositeimplicitautograd_dispatch.h,sha256=aenIraO0p-7zp9tldlBMwqDBJpY69Njf6OVbmn9HmVQ,1896 +torch/include/ATen/ops/norm_cpu_dispatch.h,sha256=LgN0McgHf_6bzB2u6uHOnGv1ah40xrB8ZdiZLh9LZBY,1852 +torch/include/ATen/ops/norm_cuda_dispatch.h,sha256=rAFtAGElwqki2OiJb2asKKwgAwz1l2jbjoP1qTK2wEs,1854 +torch/include/ATen/ops/norm_except_dim.h,sha256=35bSg2rX8OIMvvuLo6vyRSwKnLKZ099NKSuCAC_rl_4,980 +torch/include/ATen/ops/norm_except_dim_compositeimplicitautograd_dispatch.h,sha256=Z_IAQ-Cj2aNGHGz30vg3ctp4gYcizsjgbT_21dS5E1o,1054 +torch/include/ATen/ops/norm_except_dim_native.h,sha256=qujaPdGjvf8DBRqJTxj79FtxkW8SmnsRPhSa1ImC9YI,766 +torch/include/ATen/ops/norm_except_dim_ops.h,sha256=JX8m-36vd_Xgk92QfC9vxOAdX9OWPkOivEzngKN4yr4,1321 +torch/include/ATen/ops/norm_meta.h,sha256=3OBPCJ6S5yWI0qE9fcm6TXzzeIDnCem1SXkVb7cXw_k,1103 +torch/include/ATen/ops/norm_meta_dispatch.h,sha256=zl1Huec01S-uNhlfj5_vZ71t8rOqRy8fL7jxNsmrFxo,1854 +torch/include/ATen/ops/norm_native.h,sha256=uUwnklId78ipnCBrRhk8zctPPTDs1nRUd20MyZK-tq8,2512 +torch/include/ATen/ops/norm_ops.h,sha256=SB3pJ1tpv7Ya-wLUFFJ4XVs9SZ1ogUTsnUHpl7IxhXU,10504 +torch/include/ATen/ops/normal.h,sha256=wkYf5hdo-VrjhO66g3kHnTu644sZk8oPxxsXzq18QQg,11915 +torch/include/ATen/ops/normal_compositeexplicitautograd_dispatch.h,sha256=hMReygAhIkZH7A3f7Dua6IJ-Q8Oq3E1GC7SRlBszYHI,2921 +torch/include/ATen/ops/normal_cpu_dispatch.h,sha256=RTaM3OHZ5v4TUjLhIGIlaRlVZh8jR8SQjDnA7toHaWg,2306 +torch/include/ATen/ops/normal_cuda_dispatch.h,sha256=hH7zjJTN14P9HMEDPYPtl-OO4_J6pA0IgnbJUZJBqZA,2308 +torch/include/ATen/ops/normal_meta_dispatch.h,sha256=pJ1LMXAEJjz3e8Pg9fSrP3_zWmhcO1erT6L1j69yPVE,2308 +torch/include/ATen/ops/normal_native.h,sha256=nbYl4zEHBkI9leCmqzREc0hP5rULthXRIAb6lbNecTI,3587 +torch/include/ATen/ops/normal_ops.h,sha256=YzA_97ITiVltXJVBa25hS6Pt-hJqvMwwoVrorybo7-s,9888 +torch/include/ATen/ops/not_equal.h,sha256=46rGJui3-kGy_OYmbURITdHsiW9621Aikoh4V1_YkmY,2229 +torch/include/ATen/ops/not_equal_compositeimplicitautograd_dispatch.h,sha256=c9a6mmiEZbgdSpgtbxzwNIATQP5ONSITT-akpvDUiSE,1720 +torch/include/ATen/ops/not_equal_native.h,sha256=s3itfrPp3hKOjBWjWII9-VggE7nxLfnDAORJwQ9WeQo,1216 +torch/include/ATen/ops/not_equal_ops.h,sha256=I-_K8wBN4m4rVhYRByEm8R_ehuHpSCep46lriNkeUIk,4581 +torch/include/ATen/ops/nuclear_norm.h,sha256=VNkWZmqmBZiPEuTHAMJlctWYXR5whhzMtnIcoN3TLMM,2340 +torch/include/ATen/ops/nuclear_norm_compositeimplicitautograd_dispatch.h,sha256=e6v8IqQTUU5woGQ-DlD-Qc6nmWlsbecajbCBQWW6myc,1593 +torch/include/ATen/ops/nuclear_norm_native.h,sha256=pe37nwIvfGzR9Omdl_HLeCl3RjsHZhLtG4V56Wx2aSI,1074 +torch/include/ATen/ops/nuclear_norm_ops.h,sha256=83jT_oJjjYFZ4uT0eYegJEoRtZgGIk9CNwWWZs91bA4,3344 +torch/include/ATen/ops/numel.h,sha256=YiV_IOzZj82f--7CL3z9jx9xDlTUBiM2FFuzCrqUzlw,756 +torch/include/ATen/ops/numel_compositeimplicitautograd_dispatch.h,sha256=2osIzi2Uy3MnHRCu9apwFIYhVXl00uedvx_EqjWvBrA,1014 +torch/include/ATen/ops/numel_native.h,sha256=Ql9Lku3Xf7LgaTnlGcOEzAdeua8c5l5SUwVZtaSKMuI,726 +torch/include/ATen/ops/numel_ops.h,sha256=aFau1zEB2cDfP2IrR2FB2xgQxZ_eJdvwelfrndmXiH8,1196 +torch/include/ATen/ops/numpy_T.h,sha256=32PcuB7WnL62gX0HBE1-5nBWHxCMWuYbU3LAHXZgXyk,758 +torch/include/ATen/ops/numpy_T_compositeimplicitautograd_dispatch.h,sha256=Pv83dKFQ96GgSO9O-dC8zeX90xRZwj5VU_BOl3zSgs0,1019 +torch/include/ATen/ops/numpy_T_native.h,sha256=1JjT7rDgNqehTMlbmFtA0WZCKwgWPZnMuXBqnzLUYR4,731 +torch/include/ATen/ops/numpy_T_ops.h,sha256=hUtDgK-fnM2g_vu-grneQHERiSXCiKV-UxxzjAsKFSw,1220 +torch/include/ATen/ops/one_hot.h,sha256=G557pLCY1UF3S4DntHrQ5iTFO-S1RLESiKKObPdmelw,952 +torch/include/ATen/ops/one_hot_compositeimplicitautograd_dispatch.h,sha256=MkTITT2rvxCjx7CJgcLkqXzCdC-aSYCcDXK12hFPVHM,1043 +torch/include/ATen/ops/one_hot_native.h,sha256=dhrhm1td1sqyj-H8v-g37DlzZ1lf6dWOtH0TGXIim0g,755 +torch/include/ATen/ops/one_hot_ops.h,sha256=0g6hYLhvPvq13oA-Nqg6ngqeDN-T70fGSm3ALe8D9YE,1285 +torch/include/ATen/ops/ones.h,sha256=W_g22DwoMDQ8yNcEAYcUv6DXn2JgTQtYeEvH1qM5V94,7083 +torch/include/ATen/ops/ones_compositeexplicitautograd_dispatch.h,sha256=bcSohEEpEy1V9CHcPLGhQuMisyHX3z3dYAIpgqM7hL8,2428 +torch/include/ATen/ops/ones_like.h,sha256=65po7sEMvdbiZD9nBi4gnq61JaK9mu_310T7-5oN1gQ,2462 +torch/include/ATen/ops/ones_like_compositeexplicitautograd_dispatch.h,sha256=FRsNhP7kllx7X6xwHFxoXDwvahwfhV77yIESTbgusqI,1642 +torch/include/ATen/ops/ones_like_native.h,sha256=ga7jIKFXUX1eqWOdJ9uPvmaYeDjmCMVBZzo1iQnY544,1084 +torch/include/ATen/ops/ones_like_ops.h,sha256=PfK109IN9n1F6DYtWTPA5pKDJEtp_efyVRzEjQXwRDk,2649 +torch/include/ATen/ops/ones_native.h,sha256=HtIyUUspwQSg26Vy8q8f-YEg5CrYOALfyXmU9Whp7oM,1320 +torch/include/ATen/ops/ones_ops.h,sha256=i94zNKxsqQtLOSqiofYVWJDgHKzaGO0kUAN2ZR9HU9E,4148 +torch/include/ATen/ops/or.h,sha256=wEbGL_7iJQ3LyZpFvyXJWb3hsocq34iKEDGlwmJ7l_o,1144 +torch/include/ATen/ops/or_compositeimplicitautograd_dispatch.h,sha256=dWCpHwpRVFMabWAlCBRK1h-A2ouBdvjCh0Ft4GZ2ClM,1278 +torch/include/ATen/ops/or_native.h,sha256=cQ8L_VYxWkgN52P03bKe6-CGdwj_YBsUIKoaTXfb284,990 +torch/include/ATen/ops/or_ops.h,sha256=30lZvc6fitKLpwgplZEc-fsZ0m_pwA0mtvKA_WAkyv4,3133 +torch/include/ATen/ops/orgqr.h,sha256=_kgKEfoZqGKEpRsbqsOyFG_JH6qSvOPhyiW_iOpyiLs,1421 +torch/include/ATen/ops/orgqr_compositeimplicitautograd_dispatch.h,sha256=DnaoWNhTh0ess9dYIMAPAFyq5ckOVkCKwbuSljXGbhI,1253 +torch/include/ATen/ops/orgqr_native.h,sha256=qdmmLM50RmgcxE0C0tTOn0GhDAupY9Q2thoASCrEc6o,860 +torch/include/ATen/ops/orgqr_ops.h,sha256=uBj1UWLK13j1hDC11DP8TZ8ec_M5VFcJXdPafVei-xg,1973 +torch/include/ATen/ops/ormqr.h,sha256=pGd-p5wEcmMR9Kl5Cqqby1WApBo0ngRD_Sf-EqMYRiA,1839 +torch/include/ATen/ops/ormqr_cpu_dispatch.h,sha256=gnSCAjNid_Bp6K2QMSdvBBiRxp4GEPIPrl0M5yZx60g,1393 +torch/include/ATen/ops/ormqr_cuda_dispatch.h,sha256=xP3oHdel9gPmAZ3kMNZziGqzJFo1-6DXAbkQV4jQ1QI,1395 +torch/include/ATen/ops/ormqr_native.h,sha256=alRGzVKLWzNDmwgLxOjphebjq-LvmXZKwxMyezQmbeE,979 +torch/include/ATen/ops/ormqr_ops.h,sha256=FrmoBlDTz8h8BuTuxPh77KechsoAX-v9A2YRx4jj4RA,2359 +torch/include/ATen/ops/outer.h,sha256=lwZApKb_eLqz6Q3oRMtBnQn9CdcDrqheK5mjfw9yPY4,1403 +torch/include/ATen/ops/outer_compositeimplicitautograd_dispatch.h,sha256=vm1KAYvuhjAMGNiGGpuh6nKrp7WIhXR1tzxpQdJeQM4,1247 +torch/include/ATen/ops/outer_native.h,sha256=Y_3kwwRs1la_4cLm6EZ5TfTef7GwYLtgsC5L_-jPEfc,856 +torch/include/ATen/ops/outer_ops.h,sha256=W-585u8uC4AapAs-Gn61_vN5vm31L_rEAr3pbo6hIpo,1961 +torch/include/ATen/ops/output_nr.h,sha256=Qh4E28AaxT7dPXJ85KP-iGp9PIsmw6BfwqCArG3fBGA,760 +torch/include/ATen/ops/output_nr_compositeimplicitautograd_dispatch.h,sha256=UdpJISwTc4233OyFr3eeRwHujCAg5w6kvENi1lkcjMM,1018 +torch/include/ATen/ops/output_nr_native.h,sha256=G7JHbr5l9ho9Bh3NZ3aS8PD6_DIEHMpoxEbnViUzfmY,730 +torch/include/ATen/ops/output_nr_ops.h,sha256=QTcEcdaOf2O6rUkF5n67O2WyFVi1P9lTsBgv1BZrPcs,1208 +torch/include/ATen/ops/pad.h,sha256=x12iCHoGpQF8TydCXZrE3R8jMUVaxnBlU7cJWSsD9Ws,2036 +torch/include/ATen/ops/pad_compositeimplicitautograd_dispatch.h,sha256=SSC5us6Z9D0wXhTxZDjMU5RNAzJXLBb7iLgVna-PlBE,1279 +torch/include/ATen/ops/pad_native.h,sha256=OUY8EBQoFD792LKHSvbwxaKa15D1fNmbIIVjA07Btg4,839 +torch/include/ATen/ops/pad_ops.h,sha256=-oFcHxD7Esw_4xy_ZP_dcqXkjYwbHxtQjI-0UvChf2k,1480 +torch/include/ATen/ops/pad_sequence.h,sha256=KylRCD1P9tCM41-4_i5bL_I2_I6paYzgf6NrRPIc4X4,1134 +torch/include/ATen/ops/pad_sequence_compositeimplicitautograd_dispatch.h,sha256=ldIyjPXz5f4L5Hy1WcnLX-ivA0uRLp4mljl2MRMMTRo,1114 +torch/include/ATen/ops/pad_sequence_native.h,sha256=j5nMc02Rh9sKhL2meN7AZZ0qMjbkS3STCWUteg_3NM8,826 +torch/include/ATen/ops/pad_sequence_ops.h,sha256=kJMRW6YMTB1ZALIVvEmsJwtmVy7-FAwjfLKRymlj2bU,1485 +torch/include/ATen/ops/pairwise_distance.h,sha256=LwNs-fuATmgAxZgHg2lo7nOsRl_xLaPlNGWfAmi1A1M,1082 +torch/include/ATen/ops/pairwise_distance_compositeimplicitautograd_dispatch.h,sha256=9Ca5HKvuSz01zCiFSik3MA8A9M5lBMIEeOwOeAaN9v4,1100 +torch/include/ATen/ops/pairwise_distance_native.h,sha256=HcociC6WWIZbbawY0WYUTNif4G6gc9JQUKiqp_sL-Cc,812 +torch/include/ATen/ops/pairwise_distance_ops.h,sha256=uMbfQQuLz7KI2pzt4tRdZ-ryGTI01r7FZDRSp44V7Z0,1457 +torch/include/ATen/ops/pdist.h,sha256=dD8ZU0Ee2h8FiB9xo5AvnHnr0rI-ckHxaci6eVkvxws,913 +torch/include/ATen/ops/pdist_compositeimplicitautograd_dispatch.h,sha256=eS3GeRxyjlaM4bBc0cmFwQ-L7u61GVTSNEQvhBjcJxs,1029 +torch/include/ATen/ops/pdist_native.h,sha256=tPQIZmKuiPax2Gpy6DknEZchsVHAFe_ivFWZ5yWaRUI,741 +torch/include/ATen/ops/pdist_ops.h,sha256=hl98K9tMth4ixw9nWXSEdwE8huBcC3kiCxWrFpTdjR0,1247 +torch/include/ATen/ops/permute.h,sha256=ZWyQUTnuu1qiPs3H036jYXcgjhJlGNEaaDPkPB8_8PY,941 +torch/include/ATen/ops/permute_compositeexplicitautograd_dispatch.h,sha256=5B-SH_D-nq4rlLYnS8WF8qDkvRbq8pytlQoIlqFMzV0,1041 +torch/include/ATen/ops/permute_copy.h,sha256=Ia321sUCvMi0WSHBq4etdjrQ1MMSbWS_upPLXkeom64,1461 +torch/include/ATen/ops/permute_copy_compositeexplicitautograd_dispatch.h,sha256=YdDg3DKHc5Pfhq8q1zrDSbg-jphY9yY-H6l7BX9uTL4,1177 +torch/include/ATen/ops/permute_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=9Dx9MCSR9a4fews9rdzRBK1oVCap-AL_l0o5_nPszy8,1072 +torch/include/ATen/ops/permute_copy_native.h,sha256=nG1UwTZQYOM0ZxdaCzYYRgp_m3wfB8iaTu7hwUUnP8c,864 +torch/include/ATen/ops/permute_copy_ops.h,sha256=9h2dMXRmfP4mF20v64FdDcbvRTI2YlSr7J1gxH_zq0s,1983 +torch/include/ATen/ops/permute_native.h,sha256=YN2K44AhN2B799YS23MEykktktlZhh1TQEq6m59eWCw,841 +torch/include/ATen/ops/permute_ops.h,sha256=aCjn1K4Y4GKEHWACWIDrFWlbYfHhcElkOK-RCZtPrYc,1293 +torch/include/ATen/ops/pin_memory.h,sha256=m-iQal92NcI7l6icfTfOBPBhaKVwHStutoVUSutEBI4,761 +torch/include/ATen/ops/pin_memory_compositeimplicitautograd_dispatch.h,sha256=I0DqFATtVQ8hR3ugq07fjqSEyf79xyWLgRAuM5U3ciE,1073 +torch/include/ATen/ops/pin_memory_native.h,sha256=35VkVuEljLoSWYqsVAWieLuW0odmC8RrbEFPoMpBPEI,785 +torch/include/ATen/ops/pin_memory_ops.h,sha256=KyU4ieqwAoRrbO3gSX59ah-JJ0zLOL4Ha1FeuxOi9mA,1351 +torch/include/ATen/ops/pinverse.h,sha256=fGcxnWh-hAQJOQwVeRgtcOkatt8AK28R3T1Vt_lsC7s,945 +torch/include/ATen/ops/pinverse_compositeimplicitautograd_dispatch.h,sha256=oqZCmG7bp8UZlEI0YZ_1gz3bRnI52jlrvOaTGKpG1kY,1040 +torch/include/ATen/ops/pinverse_native.h,sha256=ReZC8u40pfr5BHqxX4jVcd-ON3XGkt5fXmAbx9HJq2c,752 +torch/include/ATen/ops/pinverse_ops.h,sha256=lZ3vBriEXftZgXf660MmOA69XSERE_gN98oBb7XaxkQ,1272 +torch/include/ATen/ops/pixel_shuffle.h,sha256=ejKZdtcIDd9tLZlboJsDLKqXQcynS7U5j0eA9RJvNUY,1531 +torch/include/ATen/ops/pixel_shuffle_compositeexplicitautograd_dispatch.h,sha256=S7DC1zmaKpTiBRJOegN88D3UBewRgKBZ4L_x3NVzulw,1183 +torch/include/ATen/ops/pixel_shuffle_compositeexplicitautogradnonfunctional_dispatch.h,sha256=JfZQBrraEp2Vdr0CcuYPRViZ-lYJchGAKbpo2zrOTAg,1075 +torch/include/ATen/ops/pixel_shuffle_cpu_dispatch.h,sha256=SlKVkKIuKa8jGWOoDLnGbjQA2mKe022T5yw_EvACe_w,1005 +torch/include/ATen/ops/pixel_shuffle_native.h,sha256=ZQxeEzBz7z9tnCXsMxuNXv2dtsObmxrbTcNiaSU8pOQ,964 +torch/include/ATen/ops/pixel_shuffle_ops.h,sha256=wuEfBzTwjjf1ZlEWbF8roXLidJ7lsN4T2s8jds8x4ws,1997 +torch/include/ATen/ops/pixel_unshuffle.h,sha256=xmjgeTFjGetuyVGSBV5brLeWRxACilEXbM55spkJvr8,1569 +torch/include/ATen/ops/pixel_unshuffle_compositeexplicitautograd_dispatch.h,sha256=5xBbUDPtX6lLAnG5QaJR9uUjHXcH_1xNaudLfTFb9Wg,1191 +torch/include/ATen/ops/pixel_unshuffle_compositeexplicitautogradnonfunctional_dispatch.h,sha256=2m-mlFy51O12WjhngoxbycfDWuM2cH3BEsZhhO36erM,1079 +torch/include/ATen/ops/pixel_unshuffle_cpu_dispatch.h,sha256=DHgDfhmZ0PUqqe4RqFAx8RpgM8-HXhgcgj4xZNB1n5Y,1009 +torch/include/ATen/ops/pixel_unshuffle_native.h,sha256=oDHVRRKOKKKWCIt0uALzDy0DGJ5pOi5frp-P9reDfWI,976 +torch/include/ATen/ops/pixel_unshuffle_ops.h,sha256=laVhiKWL91sz_nMff4xXljVg3dvYY8IeseXkpLQvNjs,2021 +torch/include/ATen/ops/poisson.h,sha256=AyVek-ij1xNMT9a9gUyQUoyvkGzhmHHNBT9Q0EBqFLw,1561 +torch/include/ATen/ops/poisson_compositeexplicitautograd_dispatch.h,sha256=giBT2SV-N5P33J3vNwThSlNTINe9TGWaIAsC_XDTRC0,1222 +torch/include/ATen/ops/poisson_cpu_dispatch.h,sha256=y0d3Jb4SOCUV5h-HN5pbARqb2vIk6M9slt2MmBdl_tU,1032 +torch/include/ATen/ops/poisson_cuda_dispatch.h,sha256=OdeJuLd5ml7cXSAuTcMhoFqSdTWei07HvY2Us8f5tmo,1034 +torch/include/ATen/ops/poisson_native.h,sha256=ic7tTHq7bgHilkfQzHURCff5Zhw_WbO2zhbQj4oyPME,1036 +torch/include/ATen/ops/poisson_nll_loss.h,sha256=HksHExpx5JmGBYBcOn1zWKyzqnxTx3YEuOXXOz3YY8A,1128 +torch/include/ATen/ops/poisson_nll_loss_compositeimplicitautograd_dispatch.h,sha256=0SlpWtXrli-3drB73FPD_8Wv5TaKlGSXQs2ahkdAHbY,1114 +torch/include/ATen/ops/poisson_nll_loss_native.h,sha256=0hkcCb32MF0gnA2A_O-M2wSAyNv9VUKde2PhblQfvyw,826 +torch/include/ATen/ops/poisson_nll_loss_ops.h,sha256=mr6YHsfVWWT6aPoMooV3DDpRCgiqC1j5PCjgHl1D7q4,1531 +torch/include/ATen/ops/poisson_ops.h,sha256=fs0y73ogeDIqjpInglErUu87M0mex9hPA-c7rPiwju8,2093 +torch/include/ATen/ops/polar.h,sha256=59nwGgeNPIvJ0x3_djeP2Px-1mHrlpyTuUv7CZLcN3o,1403 +torch/include/ATen/ops/polar_compositeexplicitautograd_dispatch.h,sha256=8jC9x-50uJzs4RlQYTTyFI5GqmB1euAJsBFviP46aKw,1042 +torch/include/ATen/ops/polar_cpu_dispatch.h,sha256=HJPL7YAIEvTamcN4zE_CY5mMamLuBAEQtdkvSnP2Pi8,1125 +torch/include/ATen/ops/polar_cuda_dispatch.h,sha256=uQpP6S5i6jmY0_MUenlQsd2o17qbAmlwXuCQRcmTlWI,1127 +torch/include/ATen/ops/polar_native.h,sha256=PEz-OBSeFuiw1jx5AtblKBAZ2f8meYUHoIMEJTa7EVI,856 +torch/include/ATen/ops/polar_ops.h,sha256=g6fpYt16-ercG4hBhs3iWCSQV2QJbJMmJU7ypF90jTo,1961 +torch/include/ATen/ops/polygamma.h,sha256=-3ikoI7C6MD1aOysM7Eg3ciKQVVDJil5BBvrQSAgrME,1374 +torch/include/ATen/ops/polygamma_compositeexplicitautograd_dispatch.h,sha256=ATyYl5quZzXVmMHddMycxu24v6q1ilz6lgcH2YeDZr8,1029 +torch/include/ATen/ops/polygamma_compositeexplicitautogradnonfunctional_dispatch.h,sha256=FRgHWy67jyCpxRL5eQxPoT4sEAo1OLw0t5fyw7j04_c,1058 +torch/include/ATen/ops/polygamma_cpu_dispatch.h,sha256=fPQ7AR0Mhj8hgT3fUnce-yzrjhvT4AQS1deAkvVXdxw,1173 +torch/include/ATen/ops/polygamma_cuda_dispatch.h,sha256=qhFeGZ0lh4FsTDApAZNrLH53-Iw3_1XkxjUtZl0Vy70,1175 +torch/include/ATen/ops/polygamma_meta.h,sha256=dWzsOLaB2VHnnRmqoP6YfnaOY0o7CJuYw1rFjmZ0M5Q,835 +torch/include/ATen/ops/polygamma_meta_dispatch.h,sha256=Jepz4eowtRhljOQqo8qdoJR1QNSVk7Bt33mhEzZ6Fi8,1175 +torch/include/ATen/ops/polygamma_native.h,sha256=DuBh1KOv1nZAtvYFvl3ZtWn-TnNvNvHX3Ndqo4-hFYM,935 +torch/include/ATen/ops/polygamma_ops.h,sha256=PcoteqJd5A6x0Rsj_1Aajnuf7x3bEokXrWkX0rRkRgM,2441 +torch/include/ATen/ops/positive.h,sha256=Mu66OlqZj3X8Fjlj8X0VfTWnZkzWfQBe-6O1EYAHh8g,905 +torch/include/ATen/ops/positive_compositeimplicitautograd_dispatch.h,sha256=QGsRJWN5RhpOrKn1jnHFQnC6wzW4k9CbiSulfubPbj0,1020 +torch/include/ATen/ops/positive_native.h,sha256=hkiqdyQTI31dieW0rysixXq2-LeJnVYn7E5ZKU4boYM,732 +torch/include/ATen/ops/positive_ops.h,sha256=isaCy0KrT9J4uhltWqTJAPQFPxS6PVrfMAhQO81x9F0,1223 +torch/include/ATen/ops/pow.h,sha256=-EK328jhUomchm4GJz8g-lKw2AONa01SQvo9b_O0HXM,2961 +torch/include/ATen/ops/pow_compositeexplicitautogradnonfunctional_dispatch.h,sha256=nljLUc8gm-mxAzjoArePxFZ_1pC97m0XAwXmbyodLgw,1384 +torch/include/ATen/ops/pow_cpu_dispatch.h,sha256=ArgZui8B6sIA5oigxgmsvRU3eZVqfHtaV4p4rn7KmWo,1941 +torch/include/ATen/ops/pow_cuda_dispatch.h,sha256=fqH9JOOy9E5S3UiGAvf5B8tyTh67Fz_dxlg4wzn3ENM,1943 +torch/include/ATen/ops/pow_meta.h,sha256=DGxa8NeFcOfZMxDtSVJVD3jcAQAldAC7mRnkQSiWX3c,1154 +torch/include/ATen/ops/pow_meta_dispatch.h,sha256=OfUX0PQZOMAK1kcpFgCZ9Mtb6aksFaiwG0EXpcUaGHU,1943 +torch/include/ATen/ops/pow_native.h,sha256=I4_ogEWBQmWgIr8gz6xcKS2X5l-IHQdlVJlu0eP5IE0,1480 +torch/include/ATen/ops/pow_ops.h,sha256=yXOFN_7rgdtJcR9a3sv003I5ifbaAz-sAOkrAuUWgAE,5914 +torch/include/ATen/ops/prelu.h,sha256=TJ-kYkNkD_T7VzAySkx7gE5O3rfPUfrPQOn_neOuTXg,937 +torch/include/ATen/ops/prelu_compositeimplicitautograd_dispatch.h,sha256=WKV2r00u93ai8fCd38QsPzg4b5n5wMMPpmxzMXNG01M,1044 +torch/include/ATen/ops/prelu_native.h,sha256=0NBrld2TiiRUTivI5ijBy-cRI2EqSZNmrmFU0paivyg,756 +torch/include/ATen/ops/prelu_ops.h,sha256=3sXLxty2j3z60eIx_i_PoP1RKZlWIiofLunih73JtKM,1297 +torch/include/ATen/ops/prod.h,sha256=pYSiMEQ3BTDvCsqsz_-SAQ_VCE8lHKBKQnhVA3Lq-is,3560 +torch/include/ATen/ops/prod_compositeexplicitautograd_dispatch.h,sha256=jHIedNnCNCGj4Z4SVFipP9ik6kv7FjpnKLHFKGbQx2Y,1210 +torch/include/ATen/ops/prod_compositeexplicitautogradnonfunctional_dispatch.h,sha256=XGARH5XY2QagoAuQ0-vyae-Ji4p4vAR7kOCCYnY6OUA,1129 +torch/include/ATen/ops/prod_compositeimplicitautograd_dispatch.h,sha256=X_ABSbuV7Nm16kQNjWFhaG7KT4aNtIdrX5a4StYpwYs,1421 +torch/include/ATen/ops/prod_cpu_dispatch.h,sha256=fgnYuoD44QvJhTYIe3eJEbuQhGwmgfzrhwPOKWH9HrU,1471 +torch/include/ATen/ops/prod_cuda_dispatch.h,sha256=DmsR0vXUVOuhwxqSBxA_ooTMfaK2asLzvIgo4WM232E,1473 +torch/include/ATen/ops/prod_meta.h,sha256=7VOwKynXEdDYYQ52QslmexZjK9YMCiJrbUfQiiT8NGI,893 +torch/include/ATen/ops/prod_meta_dispatch.h,sha256=t5SZZU0jGkzxTL8hbtKh5Er9vQJen8JC6vael0AGwTA,1367 +torch/include/ATen/ops/prod_native.h,sha256=G1DKLXvmTOATl9eyB9vuA2BdasqmlI7DariPgW71bao,1428 +torch/include/ATen/ops/prod_ops.h,sha256=0Z1tRhWneaq0r99tJJCtVLf9oR3k0rTDdFqxqBLEJL0,5282 +torch/include/ATen/ops/promote_types.h,sha256=yY5J_y9-UDUVg8AOwSjiwBDGk5hgDE-JavgC76jnEiI,977 +torch/include/ATen/ops/promote_types_compositeimplicitautograd_dispatch.h,sha256=Arja6ix0IM_1Ky4cOYBq6FVReHhYa9j0qFg7g8h2bpk,1048 +torch/include/ATen/ops/promote_types_native.h,sha256=HKlBJqK7kookmuQO9vRLxiZb0fr9gEaxmw9lIVxcChQ,760 +torch/include/ATen/ops/promote_types_ops.h,sha256=5AtQQLqc9Whd8GGIpyoc9XsZ4X683vzN835rcZwP0yw,1321 +torch/include/ATen/ops/put.h,sha256=zKxCVw0WuiIJjmPoavVgDbOZEXltqlVr4BWb9tu-hio,1710 +torch/include/ATen/ops/put_compositeexplicitautograd_dispatch.h,sha256=2GFpgaKRAgLzybNMzLnYQ_WpukiljfYr225tc_nE-Bk,1388 +torch/include/ATen/ops/put_cpu_dispatch.h,sha256=AaAzeeYizV4QcCcQNkLj2FqllfIUk5VTlbRmIgXQVi0,1044 +torch/include/ATen/ops/put_cuda_dispatch.h,sha256=DrVD92DdMWmTY2GcKw9V1w5f_Sjsl-rOA4vSuxFF-g8,1046 +torch/include/ATen/ops/put_meta_dispatch.h,sha256=M-jUHofnB-CI49NOTKfmQxz2vjDQJGhfWRxuNUyEJVs,1046 +torch/include/ATen/ops/put_native.h,sha256=OjL_Ozw2av95qgHF1aRfyLj-iMrYinDJ969jmqWlzFg,1072 +torch/include/ATen/ops/put_ops.h,sha256=Xk3ZLM2f8WFluntexeEMJz0WgkCx2S6taNEX8tfHJ3U,2987 +torch/include/ATen/ops/q_per_channel_axis.h,sha256=tYmmRertATakSQ9hD0NCmrW1CJFDu3HNG87ejkqfQpg,933 +torch/include/ATen/ops/q_per_channel_axis_native.h,sha256=f8rqlsUs7RTqUg8bdZTRoL9RV0LtmDgh1OJuiQo8NgM,739 +torch/include/ATen/ops/q_per_channel_axis_ops.h,sha256=pLhS6kbovOUGcBtfg3lqUD-Ugz5NthS5d7wCslKtUec,1235 +torch/include/ATen/ops/q_per_channel_scales.h,sha256=-T265OFu7YwezrloNbyJTrTHXPO7mUXiFiyyNuHrBiM,1421 +torch/include/ATen/ops/q_per_channel_scales_compositeexplicitautograd_dispatch.h,sha256=z_bZB6-nBgom99A-USUFqdOatep1B1jgPhpi2nvPALs,1149 +torch/include/ATen/ops/q_per_channel_scales_native.h,sha256=E5Vczsy39kBk1ZNCYtsoL4DeldjeBkgeNUhzu4TPKUE,836 +torch/include/ATen/ops/q_per_channel_scales_ops.h,sha256=o9qMF9Yq5W8CjsILpC_xY7szEKkWOj5j5iQ7Ef3CjDQ,1885 +torch/include/ATen/ops/q_per_channel_zero_points.h,sha256=xHCcn_bGZ2_YVnjRvL-Cogw6uWFnMoDtC_cGvBSUI3k,1471 +torch/include/ATen/ops/q_per_channel_zero_points_compositeexplicitautograd_dispatch.h,sha256=JDxc3roywf2B4blrB1jaCBATSM6uPllYPU_yRb9v1qw,1159 +torch/include/ATen/ops/q_per_channel_zero_points_native.h,sha256=7f7m6gecEdHNK0nQ-TyBN8AyXwKuZXqLRq4jpLNK6Wo,846 +torch/include/ATen/ops/q_per_channel_zero_points_ops.h,sha256=tqAAxJkMirXR2Xs3mKs2Knex_9ghHri7f874XtWgpJE,1915 +torch/include/ATen/ops/q_scale.h,sha256=1WggPLvpVzf5OYqrqh2nBsRDbq4FMn1H90YNGObIHAI,890 +torch/include/ATen/ops/q_scale_native.h,sha256=Q27fWRo48tf2NhBOYCqXXokNRpAI0xhUgys-P4dHvRY,733 +torch/include/ATen/ops/q_scale_ops.h,sha256=0J-KBz3FzFk572hu2XWKm_E298eC_myyss_Zo0U5vm4,1201 +torch/include/ATen/ops/q_zero_point.h,sha256=vBfIWKFuQayfrECh7jDFjqyCRfm2KielCAMksdjorGk,909 +torch/include/ATen/ops/q_zero_point_native.h,sha256=7B--Bv75JI_LuWHoDIUxAH64If8KUlUWBcRxWol2hJ4,739 +torch/include/ATen/ops/q_zero_point_ops.h,sha256=RJG0IdLFtyI78ImD19zkfxI0Y_UbP2BSYaFBZ9KW_1w,1217 +torch/include/ATen/ops/qr.h,sha256=mCt0cKsXHRusEDxGRr8wvtB5lrLmUoT5fItqpTC5CpE,1525 +torch/include/ATen/ops/qr_compositeimplicitautograd_dispatch.h,sha256=zWs4CYrAxqFkIpDSxsPcDPMvyv6NcU1EjryPbulso0I,1313 +torch/include/ATen/ops/qr_native.h,sha256=GCy2id16vZSRa1p4ZtW9J2Ig5fF4vjZf_Nkls1FqQv0,893 +torch/include/ATen/ops/qr_ops.h,sha256=iKD_9MJtqDyLkqBeZ79YVXBABH9n6qeHf6MWbi0DgEI,2101 +torch/include/ATen/ops/qscheme.h,sha256=qcrePpz_LIFeYvz4JYdodUaorRagEgJgCUSjDu25Gg8,758 +torch/include/ATen/ops/qscheme_native.h,sha256=SP2sqF4szvKIHYX4TXbTcHXBDdsln61BIt_XBe_djMs,738 +torch/include/ATen/ops/qscheme_ops.h,sha256=x-z57214V3Ban8AUYHvV2mNbY9btFN05qDY4gBwAK7c,1218 +torch/include/ATen/ops/quantile.h,sha256=s6m5pSEXB6R8AJg18UMVW2eXawwuwj27_sDrPon3-Q0,3191 +torch/include/ATen/ops/quantile_compositeimplicitautograd_dispatch.h,sha256=jdbRH7Dwb_bd5V9hlOeImfGB14hndR1LHSSk4m0WVSw,2070 +torch/include/ATen/ops/quantile_native.h,sha256=wdI-OpgxXzH-xqC9AVB9dMmdMX5W0gZh9fTst_A59Uc,1376 +torch/include/ATen/ops/quantile_ops.h,sha256=Oy7PlxJUK-eUXLOAnU9MvHqPjNUt_c_oF-UfvRzKoHs,4244 +torch/include/ATen/ops/quantize_per_channel.h,sha256=WkuQCnJuDOOs4BmgJMLsQk_Ex5bY5w60YHKs4hnT3Do,1997 +torch/include/ATen/ops/quantize_per_channel_compositeexplicitautograd_dispatch.h,sha256=qQvHSUOuGiavBC-9iC1mWW5FwwrBoCLtjO0_cUlSSGM,1339 +torch/include/ATen/ops/quantize_per_channel_cpu_dispatch.h,sha256=CbP2HmOAy-BZJMi9eo6gpwDaGGYNJZKWHarAYkVgJcc,1083 +torch/include/ATen/ops/quantize_per_channel_cuda_dispatch.h,sha256=WP3Gw_P3OZsAqDSxmYBEeOnplwmGF-g_PR5Sz5iwNvo,1085 +torch/include/ATen/ops/quantize_per_channel_native.h,sha256=tPlLsPG6noHvkeBoC0B97llxAgJYQB5UbWZu3KubZ3s,1026 +torch/include/ATen/ops/quantize_per_channel_ops.h,sha256=fZoP1IKOPh9OMXq6GuOPf9ejI9uP6PhVl81jGKWWoec,2521 +torch/include/ATen/ops/quantize_per_tensor.h,sha256=34mT0kVQ1JnapUS61DJq-uYQab0tQWydyF4LArMwMYM,4188 +torch/include/ATen/ops/quantize_per_tensor_compositeexplicitautograd_dispatch.h,sha256=T6j9LfbP9VXbuHkfuFuL9vG97uMovGA_hE2urK9K0BA,1931 +torch/include/ATen/ops/quantize_per_tensor_cpu_dispatch.h,sha256=Uu-tbhTQNnuUcgrgcgIIWcN4dfiap_hUTavhPcFWGNQ,1351 +torch/include/ATen/ops/quantize_per_tensor_cuda_dispatch.h,sha256=hiY6FAhmHLPETKxsL7_wR-Z8CGq3zaR89X4EsKTJAT4,1191 +torch/include/ATen/ops/quantize_per_tensor_dynamic.h,sha256=Zt8SdrwOVmrgmVL1sfNB_9VWDBkQTdoDw8ShRhGkJjY,1788 +torch/include/ATen/ops/quantize_per_tensor_dynamic_compositeexplicitautograd_dispatch.h,sha256=HmjSm2p6Ddv9oMM8_4C-7YFvDDZg7-IGy517mm1Pp7k,1245 +torch/include/ATen/ops/quantize_per_tensor_dynamic_cpu_dispatch.h,sha256=qq0AH2mKj5Vn8MaPfs51FRIWCbxbsPgxJo4drDSyc7w,1036 +torch/include/ATen/ops/quantize_per_tensor_dynamic_cuda_dispatch.h,sha256=qWO5oerFtBm33iFtCZ2WQ0B89r9TTGWZO3zW7PiatHM,1038 +torch/include/ATen/ops/quantize_per_tensor_dynamic_native.h,sha256=WJ7lt_Hf8USh63iLtAp7YrZZMoDPBgRMDJUk8g0g5hY,932 +torch/include/ATen/ops/quantize_per_tensor_dynamic_ops.h,sha256=ZHjRfSmQFUEJlwOKGmN_C8SRS3JtetbtG2MjaZyzWKY,2209 +torch/include/ATen/ops/quantize_per_tensor_native.h,sha256=--AnfYaBoBXKJlYlVnVghxzTn058ZQDE_U0zKFBqq8k,1636 +torch/include/ATen/ops/quantize_per_tensor_ops.h,sha256=qWdPhurwMFX7yGfE4_mbh0JpkIIdBa9GGM0KHqrVF-k,5828 +torch/include/ATen/ops/quantized_batch_norm.h,sha256=Re9Hjr6PLCg_ymAZ_oHEI_RA8oKGo1Bp0ESdsprarMc,2531 +torch/include/ATen/ops/quantized_batch_norm_compositeexplicitautograd_dispatch.h,sha256=EY8ZuhILv9_i9PZTu2Jv639KiMibAD8kWp3psKPBkS0,1541 +torch/include/ATen/ops/quantized_batch_norm_native.h,sha256=AYd2nnzy6w8Jx4hoDtermw6O7LCJUy8tgFpp9zsNKVY,1228 +torch/include/ATen/ops/quantized_batch_norm_ops.h,sha256=U6jCfgRrK90aQh0SeXz_xTpTScOEQ2ThczQm4xY0ZMw,3167 +torch/include/ATen/ops/quantized_gru_cell.h,sha256=8ISaZ1OYgDObb2MSq5rsGI-69vbB3a-cWOo--cT67LQ,1676 +torch/include/ATen/ops/quantized_gru_cell_compositeimplicitautograd_dispatch.h,sha256=inhkG_wQnqgQ0ENC82U787h78mjhkbCZSqCxKGdpp6U,1410 +torch/include/ATen/ops/quantized_gru_cell_native.h,sha256=dGkcKSYwVdr3Bsvmp5U5iYJJl4pH8q97QlfXxjA0Gxo,1122 +torch/include/ATen/ops/quantized_gru_cell_ops.h,sha256=bSk8YkuU5wy8JAVDkyCm9ZOR9-bvjsmxCxfQYeHROZU,2491 +torch/include/ATen/ops/quantized_lstm_cell.h,sha256=j2NYTrMBolo_FqzY_lWnepWrX4cSi8QD5jqCyrNIqHM,1713 +torch/include/ATen/ops/quantized_lstm_cell_compositeimplicitautograd_dispatch.h,sha256=QyGEzK8m2VWlcmCDj5DJzyT2mdSOibzk6P3RwTRUVbc,1432 +torch/include/ATen/ops/quantized_lstm_cell_native.h,sha256=T7B7MUP9Hn9pSO0ByVV435idnp3QXbDlNKqRFwmKkAI,1144 +torch/include/ATen/ops/quantized_lstm_cell_ops.h,sha256=B4T9AQvpTv6Tc5g-kVyQ5HgMDeSZ_9TEmt2fjpa_ZqU,2569 +torch/include/ATen/ops/quantized_max_pool1d.h,sha256=GEwGRhrOwBbdodISt4HD8eSH0edViDNGM_yW4wtbBx0,2251 +torch/include/ATen/ops/quantized_max_pool1d_compositeexplicitautograd_dispatch.h,sha256=A07VGjkfwp7xAhVci5x03iVnfVCunf3wqAoFIhxlMpI,1402 +torch/include/ATen/ops/quantized_max_pool1d_native.h,sha256=HK392qLIUpe26ee_RXalccEenbusO2kl8TFxAel8V24,1089 +torch/include/ATen/ops/quantized_max_pool1d_ops.h,sha256=3nuPSb-TZBKFDdjXwl0c9vpJagvhz1vGAdo31wAetls,2707 +torch/include/ATen/ops/quantized_max_pool2d.h,sha256=n7Cw1nFwh0WTxGN5qXJFRvwtLWdIxZqD0716BrcwwVs,2251 +torch/include/ATen/ops/quantized_max_pool2d_compositeexplicitautograd_dispatch.h,sha256=1Ea3qbbHJ-FISD_PYCnTGCJNkw2-ySXciHjVVApPkho,1402 +torch/include/ATen/ops/quantized_max_pool2d_native.h,sha256=tWicl9hnkZ3lXOSHi2wpgAf6iHgB1XE02UEwTQozduI,1296 +torch/include/ATen/ops/quantized_max_pool2d_ops.h,sha256=1tbBoarO2v8UzDbjb9vJ8jbuYFkOWIzjI16dEfuUw6E,2707 +torch/include/ATen/ops/quantized_max_pool3d.h,sha256=SSEbwvl8pS4IaKXZKKmKrnhlf8BXRG0dSeyFN_WgOZo,2251 +torch/include/ATen/ops/quantized_max_pool3d_compositeexplicitautograd_dispatch.h,sha256=ALRVG5vck8nIrNpWrhYtbcXgD20U-mmhAO8T6sdNcVA,1402 +torch/include/ATen/ops/quantized_max_pool3d_native.h,sha256=fEQVDh31mm_JouOd9CqCYWV_9BPc4GZ-quoKuUYhVqg,1089 +torch/include/ATen/ops/quantized_max_pool3d_ops.h,sha256=9LLe7DDVeTUKGsa4mpSkL-ZM9mGNYU3DzoJLdwg9znM,2707 +torch/include/ATen/ops/quantized_rnn_relu_cell.h,sha256=rtqOVAZc1gYada_24zATRm7e_3C7yYU4PGHYguBRcp4,1696 +torch/include/ATen/ops/quantized_rnn_relu_cell_compositeimplicitautograd_dispatch.h,sha256=HjwjPjX80DSmtz6_iYjLB1xA8pwEKeD9pclOGZjLd4M,1415 +torch/include/ATen/ops/quantized_rnn_relu_cell_native.h,sha256=uQVATTp_Twcm76vKglNu7dVU5Spdb9M6yr1Rb4bFC44,1127 +torch/include/ATen/ops/quantized_rnn_relu_cell_ops.h,sha256=VpevJOyReHyE4-um0Nm9LCjng9tsnzovO8KdAf5wy6g,2506 +torch/include/ATen/ops/quantized_rnn_tanh_cell.h,sha256=zxN3OwhGY48MJqpaHdpt4PTb31Dnqp1GkPcOcH6QVZM,1696 +torch/include/ATen/ops/quantized_rnn_tanh_cell_compositeimplicitautograd_dispatch.h,sha256=JhXqJjT8JQnielLh15vQOXTtZ4epL4YeSnI5OckxWFk,1415 +torch/include/ATen/ops/quantized_rnn_tanh_cell_native.h,sha256=7kWFjiDJHPBbc0--3kWTfj88aJB02riZOKjUECdUQWw,1127 +torch/include/ATen/ops/quantized_rnn_tanh_cell_ops.h,sha256=BQBhu6OKofQ-HJATJhVQ-v-mzqWAultISlxCTR3jb_I,2506 +torch/include/ATen/ops/rad2deg.h,sha256=L20HGM_cphhO_glBxFJlHMucSwBjK5mNsfFKIrxsdIg,1436 +torch/include/ATen/ops/rad2deg_compositeexplicitautograd_dispatch.h,sha256=3o0DspkJr7-95oR08TysYEWhkltVTU0Peup0hGdbRA0,1230 +torch/include/ATen/ops/rad2deg_native.h,sha256=FONDrgrUReE-UKb82F9QzNUzzLBDvUQ6EWOsVzD-gUU,1288 +torch/include/ATen/ops/rad2deg_ops.h,sha256=D_UDLDxHI0i9q0-ClarEjV3fffayYJL2OYJr_lGHH3U,2309 +torch/include/ATen/ops/rand.h,sha256=zHCjxsoVWVtSTtEt0XGR6itYNSa-pe7zmpNfAxpbFXQ,25018 +torch/include/ATen/ops/rand_compositeexplicitautograd_dispatch.h,sha256=uRCJlSl2K-tBt8dJv8WCVAtifm2tr-TVAP20C1MeLNM,5328 +torch/include/ATen/ops/rand_compositeimplicitautograd_dispatch.h,sha256=XOf4FXZWAbTzlUzB7cskBVUO_kEpRoAbqKa4NJHKv4U,1448 +torch/include/ATen/ops/rand_like.h,sha256=30z_hfw4t-yuh9iOHMGhHdIRVh5o9_-y_k_-a2e2RW8,4545 +torch/include/ATen/ops/rand_like_compositeexplicitautograd_dispatch.h,sha256=CvN0dq9YBvH6nUxFxlmRwmhEM0a1G-jiEQcUxr3H0ZQ,2488 +torch/include/ATen/ops/rand_like_native.h,sha256=HT1gOkkv_I0a7Bw8Bnq_-Y6ItReGxDTJiQ6nHoUbaUY,1586 +torch/include/ATen/ops/rand_like_ops.h,sha256=iSJJNV3bo6L95OOlMbuQS64lyPyaHWfcGUwzGO0xSXs,4922 +torch/include/ATen/ops/rand_native.h,sha256=23lfmNilug7y-WH9sXgBuVghF0ThTqp0_jZW4EgcnL4,2169 +torch/include/ATen/ops/rand_ops.h,sha256=gymq5wORTLIScd5fO8kMDNxKIkQHhJYD3kB4azoIG9k,8346 +torch/include/ATen/ops/randint.h,sha256=njWb3PkPu81pzt887dvG3XW-DT4AlHNPcjCvXz443BM,26261 +torch/include/ATen/ops/randint_compositeexplicitautograd_dispatch.h,sha256=vwKaTEFxcnvddIxK70wozO8bUsuJjazMa6yk2zFmOFY,6076 +torch/include/ATen/ops/randint_like.h,sha256=Zw_Vxq7aWmWJ9TvohMCmSUEKT2th60gkChh0KfgfNPY,37484 +torch/include/ATen/ops/randint_like_compositeexplicitautograd_dispatch.h,sha256=fONL8wew8dwgBwAzQHEq02udoQSfrzoIYtG2Way030s,9768 +torch/include/ATen/ops/randint_like_native.h,sha256=Hkr1PAIv-kmenkiuqmEGSbf1O5ThaBPfyCqLZm2BaZw,3792 +torch/include/ATen/ops/randint_like_ops.h,sha256=g2x8YQ8Of013u6o5pomIRbqj4mV1idUX0eD8NeP6Y1c,14718 +torch/include/ATen/ops/randint_native.h,sha256=_9fXPK6KxvBkYw692glbOoaJU1RBjjtYAM65UB3tgpk,2148 +torch/include/ATen/ops/randint_ops.h,sha256=NqJbnqsvSDhE21wK77NcrOF4wzYmypsWHptPd_fcR4M,8568 +torch/include/ATen/ops/randn.h,sha256=YW2rOBbSn-y3ze22aip8QgsSmPrn9OKLKSB5O0YiPWw,25179 +torch/include/ATen/ops/randn_compositeexplicitautograd_dispatch.h,sha256=k_EK7MSjVD7IsTVrzGSsJZzcsmRQh1Q8PReR8xthfxY,5036 +torch/include/ATen/ops/randn_compositeimplicitautograd_dispatch.h,sha256=1VYIrrAHNebbAnXhTKgSFOOzXlOoIosEeM_FrLtgNMk,1772 +torch/include/ATen/ops/randn_like.h,sha256=Bye3xG69AGctJD_t7_PzlYw-pqeCc2M10gpJrucG_AQ,4570 +torch/include/ATen/ops/randn_like_compositeexplicitautograd_dispatch.h,sha256=jqF_nztvdBs8m1hOoNbMIiqI7mxhxHFRy8sZX7tIYOw,2496 +torch/include/ATen/ops/randn_like_compositeimplicitautogradnestedtensor_dispatch.h,sha256=hMw9z46joIOu_8W6Y1_UDw-FOjzP4e0Y_5IiX7kM6cI,1880 +torch/include/ATen/ops/randn_like_native.h,sha256=gNfG4KdGjhnrtEf5SAeBiubZybXyfVGnlWL8sQDtpKY,1590 +torch/include/ATen/ops/randn_like_ops.h,sha256=467VLa2wSHK1s5eTsKY1F8vVveWUxK3VSmsJKmSfY6E,4934 +torch/include/ATen/ops/randn_native.h,sha256=WUJ3PU-lUMGvLdyXUQ1ih2WBu4N3puONoSiF06Dn_0Y,2177 +torch/include/ATen/ops/randn_ops.h,sha256=iLDYCcvpaa3KSYc_K52Aesh-BF1PhtsQJlKQR26F0Vk,8370 +torch/include/ATen/ops/random.h,sha256=4FaCZPsIRPxvaSP6BFBi2rYYQvgCBQXEzHUvDxyQgIE,3486 +torch/include/ATen/ops/random_compositeexplicitautograd_dispatch.h,sha256=RQARzbrmjpSbGq3ib7KJ0_49yrdbesOOKdkE86uHSmY,2230 +torch/include/ATen/ops/random_cpu_dispatch.h,sha256=5y06dYxMprv45Tg-Kl9suUtdQKD9TNo1ypAFXbxVEEw,1299 +torch/include/ATen/ops/random_cuda_dispatch.h,sha256=HQn61lwk7EluM2Z6rjylHyw6anPQSzZ5ronaozuefvA,1301 +torch/include/ATen/ops/random_meta_dispatch.h,sha256=mMiibvwYr8fUAoOjwrgZWGHm93ADVrWCuTjfHbC7hjI,1301 +torch/include/ATen/ops/random_native.h,sha256=it_BC0AlkZolxUUFDUC1odw0jUzQM8GtLHV5G2BjrKg,2260 +torch/include/ATen/ops/random_ops.h,sha256=Xqn52y11O6ku4NeEHbo4Tv8MoyoZ02GXN6Ivx9KKbQ4,7394 +torch/include/ATen/ops/randperm.h,sha256=glg9AJ3ylTt-pF6t0UBwdjYccM0cL_sVOM4xT8bq38E,11222 +torch/include/ATen/ops/randperm_compositeexplicitautograd_dispatch.h,sha256=y8aHVtjduxYUUjxN30RqfNShWOO_gONRU0Ih-8SYRJY,2528 +torch/include/ATen/ops/randperm_cpu_dispatch.h,sha256=CcsZ6VtbBGPvQLacoJpXYrVXkq31BKKegwEZ8yO9iRQ,1376 +torch/include/ATen/ops/randperm_cuda_dispatch.h,sha256=QJFO2xE4S14tv3dWY1-5PZHwwO0BnhC0BQG4UQYHpcU,1378 +torch/include/ATen/ops/randperm_native.h,sha256=n40UpAUEqfRjX_ERSgRyT8PcCtwwxn7KPUSvJ765A3s,1407 +torch/include/ATen/ops/randperm_ops.h,sha256=E5OPegkLeLj8SfKq5FMlgGbnL05-kMLH5B46VZxJ4n8,4122 +torch/include/ATen/ops/range.h,sha256=9jRvfb8oVPjUmarWv2-MwFkkQfUrZiw2B0G6k0F2P7o,3647 +torch/include/ATen/ops/range_compositeexplicitautograd_dispatch.h,sha256=fti55kETwhtEhgdNFUg6-qg8LdMXrBPtSeEIzzBF5bA,1883 +torch/include/ATen/ops/range_cpu_dispatch.h,sha256=NgaZ6eDxMoGvcMbTu8AOqhgXWo2y8n_Az88YySObO0U,1175 +torch/include/ATen/ops/range_cuda_dispatch.h,sha256=D1RSUE1eeOT0Z3gFgdpHzgSaBpPNGkbV_lAqr9Js4QU,1177 +torch/include/ATen/ops/range_meta_dispatch.h,sha256=ONjT2hjR5t6jpv1gQHWhO2Tl5x47ZOVBHabhw30HwDE,1177 +torch/include/ATen/ops/range_native.h,sha256=JowoZgl4hPVAnthjug-c7kzXhyxYDqmLTNfGlBmnJDE,1542 +torch/include/ATen/ops/range_ops.h,sha256=5jWNaXTdww0V-68Tit-7fYNq4MaQK7zJAn_HtEsDoFk,4390 +torch/include/ATen/ops/ravel.h,sha256=Us73mVpadnS2DLxLyr_dM9xEkReX8whvTchohGWoEfI,893 +torch/include/ATen/ops/ravel_compositeimplicitautograd_dispatch.h,sha256=AC0K5pnZnU2vsFpHbeXwalZGpbzvTPrSMNMQRb6BNSM,1017 +torch/include/ATen/ops/ravel_native.h,sha256=eAwJ7usg91zpuHC5DZugKJKEG6Yq3095MBcKclfDM3g,729 +torch/include/ATen/ops/ravel_ops.h,sha256=QJe1iqrrX9mfUm9J0RO8mlyF391AzjqV6Dwox0RhJAU,1214 +torch/include/ATen/ops/real.h,sha256=m0C_rkM1UXEnQIPUDJNALKA-JshQphvqg_ePzmbz0ok,889 +torch/include/ATen/ops/real_compositeimplicitautograd_dispatch.h,sha256=rCcPKi6vWpwyZzllVskusj328A7fJPZ2fDWMzD7IOEw,1016 +torch/include/ATen/ops/real_native.h,sha256=1trlUHuin7b5xra45F4tKGv0yd6lkfri3KmEfHoHKi8,728 +torch/include/ATen/ops/real_ops.h,sha256=uSC3nf7sYq8TzgRdU7tZOWMXDvkX_GXHRHu2M2GN8KU,1211 +torch/include/ATen/ops/reciprocal.h,sha256=-jPdvliiLlx75gZjTXlxE1cgJGzd7uHwaJztSt2uKQY,1475 +torch/include/ATen/ops/reciprocal_compositeexplicitautogradnonfunctional_dispatch.h,sha256=MzlFb7uHD_H2hvkP36WTwOeP4KhlqSReNipxxKkEook,1103 +torch/include/ATen/ops/reciprocal_cpu_dispatch.h,sha256=884lW1jzF7CnpFMGCjjWv3LNhoX_aPBNbCNndYJtoTc,1198 +torch/include/ATen/ops/reciprocal_cuda_dispatch.h,sha256=DVAjuichl9gnssKsfWiKIf9j1boRPAojCxXGc6IvQzI,1200 +torch/include/ATen/ops/reciprocal_meta.h,sha256=Lec4KaLukDNry0qos2trLid9M2ktfWOqmJiv7p-o18w,825 +torch/include/ATen/ops/reciprocal_meta_dispatch.h,sha256=Ckqjj75BMxtx9YUibhR2fqybidHIm-RopIxakiczAYk,1200 +torch/include/ATen/ops/reciprocal_native.h,sha256=0XFprc6k6yhPWw2BlDPJ3_rjQ2PNbwDAGIpULZ0zOhY,862 +torch/include/ATen/ops/reciprocal_ops.h,sha256=qcX9FWoThII9ldlSbMcKmms_LRrfRUYJC64Skl0VPcA,2336 +torch/include/ATen/ops/record_stream.h,sha256=RIphFBPfEIjMJvONie68kSEHYP6Kw9ecpzGsF4MHrXY,764 +torch/include/ATen/ops/record_stream_cuda_dispatch.h,sha256=8bjzIFiPwQP-wQcK4VjOV_nmeWozPVV4W1fStlmmT2M,985 +torch/include/ATen/ops/record_stream_native.h,sha256=1hwm6nRo6gwRXqb_96AdOdWBtyO0Sy_We-pDb_g96vY,744 +torch/include/ATen/ops/record_stream_ops.h,sha256=LGYJDo0kliPFdM-Rqdyd__hPI35vTr49u8ei_qSUlgk,1246 +torch/include/ATen/ops/refine_names.h,sha256=sf-67YC6sUgVAg7x_DZ0esEvcUKJQ68YsyiCTfjEgvs,763 +torch/include/ATen/ops/refine_names_compositeimplicitautograd_dispatch.h,sha256=ARFmA7K2LRy9I4vXnSDrdPIP5RqMBvGf0zzjqZMaoP0,1047 +torch/include/ATen/ops/refine_names_native.h,sha256=qwI21kukpYJEiSEjWbCyOkxhH-ncOKZiMEUjybF74_M,759 +torch/include/ATen/ops/refine_names_ops.h,sha256=UzqH3C_r4xOrr9JPAtz9CX_X_mLGeq1ywyEssN4lmbQ,1315 +torch/include/ATen/ops/reflection_pad1d.h,sha256=UY80qaP1xwTwjnTzCbyJlzP80PenYTeEVMgAQJ2-7hM,4151 +torch/include/ATen/ops/reflection_pad1d_backward.h,sha256=l_ePIfxYV_9pSp5LZdOn6e9p-K7T07cM5le4f02jdM0,5314 +torch/include/ATen/ops/reflection_pad1d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=GhYCZL8SysRsZuyOLG0JCeQBjqrGmqMZFZSjnNCPXAA,1261 +torch/include/ATen/ops/reflection_pad1d_backward_cpu_dispatch.h,sha256=KEEVFzJvAZpp5pl_0OxQKWJ-ImpZhRHwUArtOyJV2bQ,1859 +torch/include/ATen/ops/reflection_pad1d_backward_cuda_dispatch.h,sha256=ZMPaRICpZcX5vQ2w3Q5cJxXovmYu75HOoVMilyYOON8,1861 +torch/include/ATen/ops/reflection_pad1d_backward_meta.h,sha256=HfZKDgWXyER7MSi-zWZs5MCiPSlr9iZd67HMpOeJr3k,903 +torch/include/ATen/ops/reflection_pad1d_backward_meta_dispatch.h,sha256=-qrCLdT1ByDqg3ep310A2cmL-y_VtHc58MLctuiiUwc,1861 +torch/include/ATen/ops/reflection_pad1d_backward_native.h,sha256=v4GA8PqAwxNvhhmAAmJmKxAn0qlBLpmFP5o_btw_Czc,1235 +torch/include/ATen/ops/reflection_pad1d_backward_ops.h,sha256=C-Ot0cNPNt731FYDBlGyYFqz_PAA-lJOYLU_bcAtasU,2361 +torch/include/ATen/ops/reflection_pad1d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Y3HT-6Bl8eno5h7t4mmJpx7sYBmI4MJqckzwHh2WwWM,1179 +torch/include/ATen/ops/reflection_pad1d_cpu_dispatch.h,sha256=dmhfwd8xYTgFlPdvXssBPFiSa2PyAqcxwvbSqZzTsPg,1585 +torch/include/ATen/ops/reflection_pad1d_cuda_dispatch.h,sha256=W-98_P50XzYFfIRx9e56e6I7sa8hPnZKi06p-1p4TO0,1587 +torch/include/ATen/ops/reflection_pad1d_meta.h,sha256=NrvYNVehvn2rINRc-ruhBsjWqv-VMHi4GZttwDzNvNg,862 +torch/include/ATen/ops/reflection_pad1d_meta_dispatch.h,sha256=Cv-E_PFMykunJOg9uAm06A5P-j_W6jjFOuMv0AD3pQ8,1587 +torch/include/ATen/ops/reflection_pad1d_native.h,sha256=1Tq0P9chLbvTfXY8NUTCj1GF3s70XvcqRbWM8niAIAk,1239 +torch/include/ATen/ops/reflection_pad1d_ops.h,sha256=mw1r3HxWJ8vRDaq47s7IjFL_5HOxGoDeom-QPeuBw2k,2057 +torch/include/ATen/ops/reflection_pad2d.h,sha256=U2BRBYRegUdF68QuSKXuYMImup7W1rpEUYL3K3rAnHw,4151 +torch/include/ATen/ops/reflection_pad2d_backward.h,sha256=Bt0OJNNkuezIEjx8TRqFmcoVtvQR01X2gGicQrNxnQA,5314 +torch/include/ATen/ops/reflection_pad2d_backward_cpu_dispatch.h,sha256=pNJcHalJIeC8dMO1avdOn8V_l4yRdj88rB0n_KBT-c4,1859 +torch/include/ATen/ops/reflection_pad2d_backward_cuda_dispatch.h,sha256=NLhnKGeptuWQHOz28zr8esQEW_Cr7CU0jEJbFPPCOZA,1861 +torch/include/ATen/ops/reflection_pad2d_backward_native.h,sha256=Kb54ajQxtyrF8DQyXw1jkHYs68l44yT8BuCU1jfG84k,1276 +torch/include/ATen/ops/reflection_pad2d_backward_ops.h,sha256=KlqPoRnuoo0wqoUJyA-Z8nxDrbZyaKVIymyj3V6z20o,2361 +torch/include/ATen/ops/reflection_pad2d_cpu_dispatch.h,sha256=QlzzQ90f5NZevfBtPakBCJ_Xjnins5CeqlqZh9wUuoQ,1585 +torch/include/ATen/ops/reflection_pad2d_cuda_dispatch.h,sha256=429oLBFEnPJ_szzCa4HqO7mq-XIbq3Q90FO98Epd-Fk,1587 +torch/include/ATen/ops/reflection_pad2d_native.h,sha256=6WnpsEvizH8uJTHRJ22doRG2krJbERDIFnDv6S0gWbg,1201 +torch/include/ATen/ops/reflection_pad2d_ops.h,sha256=3wQ4XITmP_quFNIijmfJIXM1sVA59JAVsxOzgUGu4vc,2057 +torch/include/ATen/ops/reflection_pad3d.h,sha256=vcv1E1zsthf-Cl8inuyzbQh7y2QeOmWQ3SNVQ4jdZQw,4151 +torch/include/ATen/ops/reflection_pad3d_backward.h,sha256=2_x1myub93IM9lkKd1-Iee_GMbaJsxXiNqlEGFi-o_4,5314 +torch/include/ATen/ops/reflection_pad3d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=2SIEA00zT0m4Wou2NODQAsjQHp_SR1w6A6eZWFUk-1A,1261 +torch/include/ATen/ops/reflection_pad3d_backward_cpu_dispatch.h,sha256=9ZtGHIQF9UZ8Z2tlVlLaMAxVZD_jlY626savtIYWAsk,1859 +torch/include/ATen/ops/reflection_pad3d_backward_cuda_dispatch.h,sha256=9lxLgRM5r31_jSMzE7OEtjDkB66EyZYTm4gaXdqDfqw,1861 +torch/include/ATen/ops/reflection_pad3d_backward_meta.h,sha256=5WDCmDFE-DcReKo6EQvv7_fFPi1l85ayObZyUWs6Cvs,903 +torch/include/ATen/ops/reflection_pad3d_backward_meta_dispatch.h,sha256=octYefl8wyWJTLYppgLCvifFttsVTnGBEC57kPpUEdU,1861 +torch/include/ATen/ops/reflection_pad3d_backward_native.h,sha256=WvVpG4NWP0LAykxK3LmSODOI3qGLYlPsXl2md2BNdeo,1235 +torch/include/ATen/ops/reflection_pad3d_backward_ops.h,sha256=zozrXunXGe8dcJkZccp8MmmMeVolQHR2JyjWtSRiMoQ,2361 +torch/include/ATen/ops/reflection_pad3d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=aEKFUkbofzrEEOrprIQ86Au9hHwu8B0yTqHz6dhageM,1179 +torch/include/ATen/ops/reflection_pad3d_cpu_dispatch.h,sha256=W23FfB5owb-t2-alhkZ5DsrEB0MfRFLfqN4EqquUZXQ,1585 +torch/include/ATen/ops/reflection_pad3d_cuda_dispatch.h,sha256=IiNHvHicpEu8C_FJNQXrzHb2EMEzjeOhpL7LaWHnD08,1587 +torch/include/ATen/ops/reflection_pad3d_meta.h,sha256=Tbc3adBTs32m_HsQbEmyr5CtN4jNHl88OpAoHWpk-10,862 +torch/include/ATen/ops/reflection_pad3d_meta_dispatch.h,sha256=wOnhEYo9ZRnHwRw06pRKGYNfaK8Tv4J6b7SfAxNe9Do,1587 +torch/include/ATen/ops/reflection_pad3d_native.h,sha256=Ew3mEzMDoMkyh1mQT_p4ClHNFkbrEq-TwgqkoIya4J0,1112 +torch/include/ATen/ops/reflection_pad3d_ops.h,sha256=VPAS7x6khaZsrlqsDYwSiLpJ17mvXh1oBFKn4tru80Q,2057 +torch/include/ATen/ops/relu.h,sha256=M3Ode7artRDDTJwUeP9z0Y7vuic5Bn-5B7XPDxeozG4,1397 +torch/include/ATen/ops/relu6.h,sha256=ICGrAtEkkJq-pTGwPKXwXXhaAoCAcVdgfKP8DBaIQUI,1026 +torch/include/ATen/ops/relu6_compositeimplicitautograd_dispatch.h,sha256=EbYthn2TW62VG-_1Jxn2M9QrFDiPCv97xqk8PbQhcHA,1067 +torch/include/ATen/ops/relu6_native.h,sha256=ugDbunXLVeeYsmMEhE2Ud8AE_82blKXS59JT4JwjG64,779 +torch/include/ATen/ops/relu6_ops.h,sha256=5iT80HKY2kHHZ2oIsA8v8vpXYBdy-hoG_DYZ4rr9E2c,1704 +torch/include/ATen/ops/relu_compositeexplicitautograd_dispatch.h,sha256=cC0-1b2VlMWkpFmygY8D_6reYLz8zBFsWG87sS8Cc7A,1117 +torch/include/ATen/ops/relu_cpu_dispatch.h,sha256=Qb49RqVhCTNvYT3ZLhrubXoTe9H7BzH5QXxXdUoGASQ,1021 +torch/include/ATen/ops/relu_cuda_dispatch.h,sha256=pucxk11bsv6dE1lSMY4LLcNLQzzBlZm2g5QvzS7coO4,1023 +torch/include/ATen/ops/relu_meta_dispatch.h,sha256=xRDQYzYoXHu1QT1uXCKU8QBLRMY9l-LUYQdo_cdQCtk,971 +torch/include/ATen/ops/relu_native.h,sha256=oU2xN_7HRLREco2-tF5P1Ke_1UIP4zkoi83rwepV7sE,1593 +torch/include/ATen/ops/relu_ops.h,sha256=DxhQlofLOB1ibY4VFqHroBkh990x0YUJOzOw1KYqHt8,2282 +torch/include/ATen/ops/remainder.h,sha256=Osj0cmtCbG9_xjdA0Ea05djtjQtCqIqwdGMvBgTpIeA,3006 +torch/include/ATen/ops/remainder_compositeexplicitautograd_dispatch.h,sha256=f7qvcs7gqzV7SJAtoroztt3vY28spb9RvfJhK7boqVA,1557 +torch/include/ATen/ops/remainder_compositeexplicitautogradnonfunctional_dispatch.h,sha256=b8GiBfPhw-i-M4lqSEhj6Rk1kurCUp3ReqVk6rNO_l0,1153 +torch/include/ATen/ops/remainder_cpu_dispatch.h,sha256=8gsBJAMsOZsYbC5W-SIfPSiXXHOI-lp8fA0Lxko-F8U,1381 +torch/include/ATen/ops/remainder_cuda_dispatch.h,sha256=1BaJazlSaoBWqeX7onZpZBpiNE0JXkWh0qVTuYmfrAc,1383 +torch/include/ATen/ops/remainder_meta.h,sha256=6Z_ncBROTEATuHdn3NOZREIk7costpBZo_GJVz8WdQw,857 +torch/include/ATen/ops/remainder_meta_dispatch.h,sha256=Rw0QVtedtDa3RNVwFUqohTNJYutFnGEnD4toZrPMITU,1300 +torch/include/ATen/ops/remainder_native.h,sha256=m9IlMYru23jQPFB4oC5otTHGGvpprGOlxpGFHQF_r7g,1366 +torch/include/ATen/ops/remainder_ops.h,sha256=1zDq8kbwdkTk9tOAhhtTk6QbbiP-XpdQsASuhwlwvMo,5944 +torch/include/ATen/ops/rename.h,sha256=g0mwlwiGwpvsXzd4htrYeKhMW8zieeoglU62zn0JeLU,757 +torch/include/ATen/ops/rename_compositeimplicitautograd_dispatch.h,sha256=Im_kTstnUc9eOiK3GlDYSHXhZMqqXkNQ6zXVZ0Sb65Q,1149 +torch/include/ATen/ops/rename_native.h,sha256=WbcAfMw1LzfVjxi0eW8i9lhsyZm3wzLfCngslDr17YA,861 +torch/include/ATen/ops/rename_ops.h,sha256=40j68r-kGNb7Qy1BIOPYPwj7Gf677Mw1YG4NPZAwhtc,1980 +torch/include/ATen/ops/renorm.h,sha256=iVTnx6dEppCbuP10wa4yu-F4ZLQ03aXRerV4SEBXNNw,1626 +torch/include/ATen/ops/renorm_compositeexplicitautogradnonfunctional_dispatch.h,sha256=eBFw2KmDCM8MhgwMklqnqOTRLvvHLR7fygVGDrMR8yo,1221 +torch/include/ATen/ops/renorm_cpu_dispatch.h,sha256=Q34m2YPI2fwe0vvH7sw4ct5HsCJBxgSvPEKHBBN_ZYw,1434 +torch/include/ATen/ops/renorm_cuda_dispatch.h,sha256=JNc-OZogwv61mKLJEwaZ237zpUJ9xe7Vh-sSJ6WnbM8,1436 +torch/include/ATen/ops/renorm_meta.h,sha256=SuPJShyv9iubL3yXxNhrNfdogFrjQL1LhF0HWZbl3eM,884 +torch/include/ATen/ops/renorm_meta_dispatch.h,sha256=M9L9XchacR90S91h9CJ93WwttKC_TpMCHO4krZEcKqo,1436 +torch/include/ATen/ops/renorm_native.h,sha256=gPXoREEzazMNB0Yj63y5QcQop7Qaclekz_fhgq_KU-g,913 +torch/include/ATen/ops/renorm_ops.h,sha256=l0n_35vl_Nh_0Z8w0nWxmoA_YwS4T9qbQWSbAh5c4-Q,2930 +torch/include/ATen/ops/repeat.h,sha256=YYm-DhJ_s26M5-w_FbF0oafvhEegOI8VMH9oqZCbNSk,3426 +torch/include/ATen/ops/repeat_compositeexplicitautograd_dispatch.h,sha256=PugdJDXjwh-4gHRkyn-J1ahorV8kc3YwxbrJ6ZOI-v0,1569 +torch/include/ATen/ops/repeat_interleave.h,sha256=Y-4qkS3MMNsGpSnZ4bLt4R5FDuTeMpK1ppEOz8PmuBA,8532 +torch/include/ATen/ops/repeat_interleave_compositeexplicitautograd_dispatch.h,sha256=EtEDxoHAP6q01p8VAyHx5XRzsh9ueo36219v_9ziwII,1538 +torch/include/ATen/ops/repeat_interleave_compositeimplicitautograd_dispatch.h,sha256=ZsyN3yZkPLg2GnR_qF9u8b3jxhAS752bvG54-d5Y7I0,1732 +torch/include/ATen/ops/repeat_interleave_cpu_dispatch.h,sha256=shAe6RFMhoFaCY9Xv5wHxvAhkT1VX-JbvNDPIm8vA7I,1173 +torch/include/ATen/ops/repeat_interleave_cuda_dispatch.h,sha256=yorsUGfSYSBJej10a4hbmllPv5a1-nCk8otaSYC0Ado,1175 +torch/include/ATen/ops/repeat_interleave_native.h,sha256=-Yl0xyzyqzPMN-acziJ9jmD4PUTjx_akkHsQwJ6imTI,1472 +torch/include/ATen/ops/repeat_interleave_ops.h,sha256=iKDfyew8-Ujy-HcIbsyomZ7ESI6pNQ3TScMywdGjK34,3987 +torch/include/ATen/ops/repeat_native.h,sha256=VNKlb3YGWgZpyHAj6cEJLXofU94obelDdw2K-Bp5JZk,869 +torch/include/ATen/ops/repeat_ops.h,sha256=_InN0wG9mFg54N1wraWlHODJcxR-Ut7kl8cvpN_bEnA,1995 +torch/include/ATen/ops/replication_pad1d.h,sha256=Cwwr5gTh7z_vO1S8VrmsuhHt1IErOJN-esgl55p2ML8,4182 +torch/include/ATen/ops/replication_pad1d_backward.h,sha256=Sog_edg4WD0M3KGBGov--52l9aNyeX_VqGzEvelMLwc,5345 +torch/include/ATen/ops/replication_pad1d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ptIHUmH-yJIfW2e2hh36zD1davS_sXX-sPclDErSj_0,1263 +torch/include/ATen/ops/replication_pad1d_backward_cpu_dispatch.h,sha256=aKJrgYCMFXNpCDbq8N5y9ritvnfj1o9p9ZQ_ouu7f10,1865 +torch/include/ATen/ops/replication_pad1d_backward_cuda_dispatch.h,sha256=mN-qmRr3cAuPD4waigfdpyjkSTd3tBzJO0xyuRBUyYI,1867 +torch/include/ATen/ops/replication_pad1d_backward_meta.h,sha256=D-aHDcGfrkMnfrRSBI1FGCpxhKp85hpKMg_drn6IYNQ,904 +torch/include/ATen/ops/replication_pad1d_backward_meta_dispatch.h,sha256=mLR6otC481JtqHcJ40W98U0IdDQlyAXKr3B_tetuSn4,1867 +torch/include/ATen/ops/replication_pad1d_backward_native.h,sha256=ojNbOUubaj-hmrVF5xx0XKY43oa2QzYdkuqO15n35Ok,1240 +torch/include/ATen/ops/replication_pad1d_backward_ops.h,sha256=WOUl6RE2fUjCO5UqRxgLWW9QSaWboGc7d234WisbuLo,2367 +torch/include/ATen/ops/replication_pad1d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=_sg7u1j6ORuRstakmnogDMQCKqbpf6PSF4Pl3_xKHSQ,1181 +torch/include/ATen/ops/replication_pad1d_cpu_dispatch.h,sha256=s0uH7ECJUBJVPN7DKz2Z7HAHK62qn8qJh91dXkxyF08,1591 +torch/include/ATen/ops/replication_pad1d_cuda_dispatch.h,sha256=hfJsQXlPy_uYizytPiJAm_u6KRma0N52OMxkXtkIgX8,1593 +torch/include/ATen/ops/replication_pad1d_meta.h,sha256=LbHofGuZTUhPrNbUN0SouXhlEs6oeKUYH2LfdD37qNA,863 +torch/include/ATen/ops/replication_pad1d_meta_dispatch.h,sha256=N2RC2miI99rOtgYEUdPyUvFDGPsiokvVBul5GbiqGLk,1593 +torch/include/ATen/ops/replication_pad1d_native.h,sha256=Hu2sj4UXDLX1gY-uiDB1wfVxpzTXquuhDk_OOAUJ7o0,1117 +torch/include/ATen/ops/replication_pad1d_ops.h,sha256=u-1lNbPLVpEmS9Xl83bmFiPhih_9e1GbFTQ2G8PIhYk,2063 +torch/include/ATen/ops/replication_pad2d.h,sha256=j-SSJ4Co5-s8jfGCJVoPPF74PYrxg8FLfNzVgd8Qpho,4182 +torch/include/ATen/ops/replication_pad2d_backward.h,sha256=eMFa6-htInS6xdgFvthcDlkqQkhTyU0xdVllVjpg3H0,5345 +torch/include/ATen/ops/replication_pad2d_backward_cpu_dispatch.h,sha256=pMkRopxXrFVDzY2k5t13uvUZfZmb61KgdZiypMrLvHk,1865 +torch/include/ATen/ops/replication_pad2d_backward_cuda_dispatch.h,sha256=Ro8fxz3Kjpto5hkY9KTa4KXN5l9rx-ZZ3rupR11ivmY,1867 +torch/include/ATen/ops/replication_pad2d_backward_native.h,sha256=78qaewhRA62UU6_r_UnovcXOCdnkYMmXuRqj5Wrk06s,1280 +torch/include/ATen/ops/replication_pad2d_backward_ops.h,sha256=QvaIkBT1o1G1qu_uGU6Zc84aMIKr18yBsH9C01aNc9A,2367 +torch/include/ATen/ops/replication_pad2d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=6ZGMXo1s7QCVxgCqQlNgPsEMYZeoTl_cWSyQkyIJsyA,1181 +torch/include/ATen/ops/replication_pad2d_cpu_dispatch.h,sha256=IOb1LrLtTR9xpkLblAJwymYRIOCJFF7cUrdupn9gZgQ,1591 +torch/include/ATen/ops/replication_pad2d_cuda_dispatch.h,sha256=nIOrAOck57QzDRYox84FLcJR0p_iPWMpQ2meGQiCOBw,1593 +torch/include/ATen/ops/replication_pad2d_meta.h,sha256=HXnrb9UY0UTBQVM4dIZH4WBD9EHOUdA7XLLMv2gmUz8,863 +torch/include/ATen/ops/replication_pad2d_meta_dispatch.h,sha256=LYKHPDh8ZTHI4a3LdcGn9AS6D-l9_TsHPdmQue-78pA,1593 +torch/include/ATen/ops/replication_pad2d_native.h,sha256=D_tAYx7prU2vov8QqKdzkovxZpanVNUF-OcWCQdX4BU,1117 +torch/include/ATen/ops/replication_pad2d_ops.h,sha256=-HeNAHmbzIaGMe850EvzYTHE4T8gy0KuW94c9Rd8pKI,2063 +torch/include/ATen/ops/replication_pad3d.h,sha256=9Myy1sd0Ss5uhCbMRyYzEboEiWhcmzNgDD4LT-mBM1I,4182 +torch/include/ATen/ops/replication_pad3d_backward.h,sha256=QL0tSpeBNIdIuF6Y0fa_lASmmsjsYnhKRV0PVnCUZKE,5345 +torch/include/ATen/ops/replication_pad3d_backward_cpu_dispatch.h,sha256=kESHBTy6QmdNspXp9gXz-dvtbCyCpDqCQjVjwTOE_Xc,1865 +torch/include/ATen/ops/replication_pad3d_backward_cuda_dispatch.h,sha256=pdMrLi4PG-uwoMf1l_amc-48jemiPTj81SOn8drlkrw,1867 +torch/include/ATen/ops/replication_pad3d_backward_native.h,sha256=g3VedJAwiddBH0gzEruWx6t9U0-Fb7LAxzSM_0Www2k,1280 +torch/include/ATen/ops/replication_pad3d_backward_ops.h,sha256=soIB8cO5lGaJ40ZgK_IznBlyNUordqRoDUgUu3PCjBQ,2367 +torch/include/ATen/ops/replication_pad3d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=KTBVCkRZRdU716CPOZDlrkihHntPQSZZ-VOQs2qx_Bc,1181 +torch/include/ATen/ops/replication_pad3d_cpu_dispatch.h,sha256=VOAZ6H-eFE5cnkYAkoth8OCBB_on_W2f20DK_6IKECs,1591 +torch/include/ATen/ops/replication_pad3d_cuda_dispatch.h,sha256=0Lid9SfKQrp3v-t1yE8fZ_W8eVo_URlCezWPKJOVdoY,1593 +torch/include/ATen/ops/replication_pad3d_meta.h,sha256=DrvZTdrbcwjLdghvk91yc4GKbYVTu5QJaj_WAgy4F-0,863 +torch/include/ATen/ops/replication_pad3d_meta_dispatch.h,sha256=jZ3CXGKCFnjCeDJDjW4kpYcSxJMmUZ-nO-rohuvsbTA,1593 +torch/include/ATen/ops/replication_pad3d_native.h,sha256=iNTK6Kb9ieuCPauTpy2XL3wrX8jqC3nwdM8qZRn6pmk,1117 +torch/include/ATen/ops/replication_pad3d_ops.h,sha256=h2IX1Stq3tUdVJVlDmy2Ohv70EtHIuPPZiXAwwIdFpI,2063 +torch/include/ATen/ops/requires_grad.h,sha256=iFzT1FldNNBKBzludNNejfGosc0lzv9YwmIzn9kp2y0,764 +torch/include/ATen/ops/requires_grad_compositeimplicitautograd_dispatch.h,sha256=qpBzni24u10t6eWB8pr6BiXwnVSH3VkhLCBpMF5xkH8,1047 +torch/include/ATen/ops/requires_grad_native.h,sha256=61peM-jUV1s8fnwW5-dexOp26T3pzTrqdYTOblB8cw0,759 +torch/include/ATen/ops/requires_grad_ops.h,sha256=RpnooLnkLyku83_od4jMsoeqb4f_0cBPkrBfvJkbas4,1302 +torch/include/ATen/ops/reshape.h,sha256=BFPfub9tIH_UvNfI2Xa9M7vSnIKis5aOSYGagW-i07Y,1660 +torch/include/ATen/ops/reshape_as.h,sha256=DyGu6KKfYZvSfk7pTmKSKCb_2Nzf1wsYU4STx56gnRk,761 +torch/include/ATen/ops/reshape_as_compositeimplicitautograd_dispatch.h,sha256=j60g1mY6selAp2GEywCligywln4mX81FSS-rOorA59A,1048 +torch/include/ATen/ops/reshape_as_compositeimplicitautogradnestedtensor_dispatch.h,sha256=HVZk13u_G59WiO8A6VQHYtfNSWGUQRy8bJOb_DYyS9k,1072 +torch/include/ATen/ops/reshape_as_native.h,sha256=-2G3gv7vc5iWNGpXsoSL6ji-hLvMgXsyL6PxxJeHdVA,851 +torch/include/ATen/ops/reshape_as_ops.h,sha256=jN_y295lfb9P4x_FqYcPRMFpYjpESJ4xOESDhluvkpU,1315 +torch/include/ATen/ops/reshape_compositeimplicitautograd_dispatch.h,sha256=o3mwg7Q83CjaBB5j5GstvmonxCr4fOR8Jiz2r5iahF8,1131 +torch/include/ATen/ops/reshape_compositeimplicitautogradnestedtensor_dispatch.h,sha256=DTOxteltG44ExFj1M14Fhg0PkqzSOKlwIvNET474ba8,1155 +torch/include/ATen/ops/reshape_native.h,sha256=uvmBO9ee0wTP5egkqawAgEiiIYpc9Kbdve6I4IWoFP8,861 +torch/include/ATen/ops/reshape_ops.h,sha256=wwnp5E3vhULEg872z7mVOcbvAd3Al4YBztROjdUIREg,1311 +torch/include/ATen/ops/resize.h,sha256=IYWBHzf7PLdQ_wCeEIPuCH5PwLYnJNLkmA8N88TARQ0,5597 +torch/include/ATen/ops/resize_as.h,sha256=xamaWvIy-wwV1DZdmPECPFyezAN6lA81U8XgE_Dn3P8,2213 +torch/include/ATen/ops/resize_as_compositeexplicitautograd_dispatch.h,sha256=3mgsgd8Mtnt9TRAkS_-Ps0GIFYBX8p_k0JxS4TVMIrg,1647 +torch/include/ATen/ops/resize_as_native.h,sha256=sAGg-38LAahm5nCb37WZuMTedvMX0NjCY9LTYHWiOtY,1168 +torch/include/ATen/ops/resize_as_ops.h,sha256=jwl7hsDw3suTcTHvT65bOv-2D3D7SJyekIa1hgSJ_5k,3227 +torch/include/ATen/ops/resize_as_sparse.h,sha256=M18TmIwKH7-oPvmWBcBLdDe4-mP7nbEZZ25IAqmm2BM,1861 +torch/include/ATen/ops/resize_as_sparse_compositeexplicitautograd_dispatch.h,sha256=0wGvXBB0qLc-jaiQMrN0MJKKh38WyYErTsJiNrpgm5w,1328 +torch/include/ATen/ops/resize_as_sparse_meta_dispatch.h,sha256=kmEBxHWIW5pGQSiyH99RxSZWcbnN5vnF1JzqTIzyNFQ,1028 +torch/include/ATen/ops/resize_as_sparse_native.h,sha256=yCAnsL6Ip1bw_54pvNDvALeqxoGsoIndHCFlUxzsLrs,1129 +torch/include/ATen/ops/resize_as_sparse_ops.h,sha256=o2QWVMV_tzrk7reFnpKJJq0Dyh-gsfXeg7wP6Eayw_U,2783 +torch/include/ATen/ops/resize_compositeexplicitautograd_dispatch.h,sha256=wHQiylJRlA3iFcCxep4HlvZnnEc13Lxi3vtrtrbHC1U,1953 +torch/include/ATen/ops/resize_cpu_dispatch.h,sha256=H1lLuHJ9Axfc2vaT1LwLn5_0ouV3D_ff1IqMwBSXF34,1229 +torch/include/ATen/ops/resize_cuda_dispatch.h,sha256=MaC2QglQnPogkHx2fkeSYD2LeMCT7o1Mm_OZICes6cM,1231 +torch/include/ATen/ops/resize_meta_dispatch.h,sha256=nPcAMuZrV8bxIwPtyEN8lnS2kF1msjAEd9suelYjMFI,1231 +torch/include/ATen/ops/resize_native.h,sha256=oppUtU3hm5OC9gCoiCMgpDR4NEJxwq4kR8xCHODdar4,1785 +torch/include/ATen/ops/resize_ops.h,sha256=DFHPGOw2Md-c1xWrfdPMIcIfRq6osG0kfyoEyGHrBIw,3143 +torch/include/ATen/ops/resolve_conj.h,sha256=pZtPlWnWRIbY8W-7LCBdfLbe_sTs-Bo8KsKyjp2_uKc,921 +torch/include/ATen/ops/resolve_conj_compositeimplicitautograd_dispatch.h,sha256=MMWJp6bkyNCnSl8beG7msh9bueBVklGsW3s8fIkXQFw,1024 +torch/include/ATen/ops/resolve_conj_native.h,sha256=scu8vXMMNhYOyHvKRB5HK2oowzOOh2SoS71jCo7AtCw,736 +torch/include/ATen/ops/resolve_conj_ops.h,sha256=WNfXTJeGcOYue_QDzUz5EZaps7ed5O08SzR1wVWm67w,1235 +torch/include/ATen/ops/resolve_neg.h,sha256=9UrAcYJ5Aqk-N3a_pRN0GDU_8fyATM1LyWT7IM_dDyo,917 +torch/include/ATen/ops/resolve_neg_compositeimplicitautograd_dispatch.h,sha256=vMcqr1afLV42uFe_dT1SX1BVLdd3yA7I7Hsq9QPVkHE,1023 +torch/include/ATen/ops/resolve_neg_native.h,sha256=tyQGWZmAnkl8KI7bdKtrol10brhZ_LpAAGg0PN1KKZ0,735 +torch/include/ATen/ops/resolve_neg_ops.h,sha256=p-s6jTrwrCdaHs3CLZX2AqAjAqtrDD_VvDxk30NG18U,1232 +torch/include/ATen/ops/result_type.h,sha256=3in95uv0aS1mYcr59Cwnsr0mhxPshdDR9WqGN57QvMg,1701 +torch/include/ATen/ops/result_type_compositeimplicitautograd_dispatch.h,sha256=vxGsMgAasADhF1SGQlAJDMSZA48UmdFxErNn56u7KWk,1332 +torch/include/ATen/ops/result_type_native.h,sha256=UMIeRIFjJLjuBE4c_f3BJxiZl00lSfadXDxp0ql2BOc,1044 +torch/include/ATen/ops/result_type_ops.h,sha256=WcXMZ86xsqfnLoZp4ndfKEW64AKrop5u3HZh6KqE5CE,3337 +torch/include/ATen/ops/retain_grad.h,sha256=bUNZdCHtNiwc-z-dw5xsHy8G3vPB7JMIo5AOQt4IOBY,762 +torch/include/ATen/ops/retain_grad_compositeimplicitautograd_dispatch.h,sha256=4K4BImGp96Lihs7BA32LkVKIFKfcgQHIdMXmA-5zuNc,1011 +torch/include/ATen/ops/retain_grad_native.h,sha256=96yVKd7RmXZLe7G_b52YchMTiEDCQmknNbAEh7raJ2I,723 +torch/include/ATen/ops/retain_grad_ops.h,sha256=GhFUc4kf8X2UO_5tQxBeeHeh-WoKRBGJrQEWb4ZgLi0,1190 +torch/include/ATen/ops/retains_grad.h,sha256=5BkCy3-WJP2_U_c8cWOEJlUHjCvt9CzJJoMmQSt_fno,763 +torch/include/ATen/ops/retains_grad_compositeimplicitautograd_dispatch.h,sha256=ZRoDPXbg-uJRZMQ6WpvW_DSLLiHe4yk8RJGOFiT8zEE,1018 +torch/include/ATen/ops/retains_grad_native.h,sha256=hb-WUlpilDpbV1YIOEvArLprzHaKglof0WkX9PAo5sw,730 +torch/include/ATen/ops/retains_grad_ops.h,sha256=V_SQL20VIOniAonJ5l-IluUsFw2r23EVy8I90Oh1YG4,1209 +torch/include/ATen/ops/rms_norm.h,sha256=xJQba4oZBg-DWe7WRJ1f4jW_fqgdDGntvmap-BUF8rc,2271 +torch/include/ATen/ops/rms_norm_compositeimplicitautograd_dispatch.h,sha256=b8SYQoMbJd7sAo8mPdTWc7hbJdVCIVv_tMe8NzpxZZY,1339 +torch/include/ATen/ops/rms_norm_native.h,sha256=-JNHx8yW9UgXOio0PInQcoe6kAjdZBuplaySltW1Vn8,869 +torch/include/ATen/ops/rms_norm_ops.h,sha256=7bCTm0b043x4416xi5mI0s5zVYdmtvhStOLskTTS0us,1590 +torch/include/ATen/ops/rnn_relu.h,sha256=QEwxcBWyy6jWgWgwJMM9iRS1af0lcL7QSTKPXjO4IBU,1861 +torch/include/ATen/ops/rnn_relu_cell.h,sha256=m2yH1O2fDnqCVtf0mIUt47I1ymuvvu07MCulmFZewec,1188 +torch/include/ATen/ops/rnn_relu_cell_compositeimplicitautograd_dispatch.h,sha256=HinWxB7fWsAmQ9q2ydL9LDUmIefTXlBWbsOYuSehX-Y,1189 +torch/include/ATen/ops/rnn_relu_cell_native.h,sha256=9tzEeeaTW5EdDefnrthG6D6dq2xlmusSo0dNNZ1yLUM,901 +torch/include/ATen/ops/rnn_relu_cell_ops.h,sha256=1pR-1h8NaRY99OV9wYFTEGMGkUoydlr-oxt0h03Wykc,1758 +torch/include/ATen/ops/rnn_relu_compositeimplicitautograd_dispatch.h,sha256=TQrk-NR5CG-FU-SNlYZ0r0epL0VC4ZB-X7ZX8MPWiM8,1439 +torch/include/ATen/ops/rnn_relu_native.h,sha256=O6OcVbUzxbAEgY0PcmNcG1F5mRab6C-dM6H45eI5R2c,1151 +torch/include/ATen/ops/rnn_relu_ops.h,sha256=Wj3AbagwQNpv344CwEUp95syFEldV00LaPi7gPemGCI,2973 +torch/include/ATen/ops/rnn_tanh.h,sha256=GQnk-5thc9RdtoST57kLlV3o5HYZVJ_xon0mrNN0Fm0,1861 +torch/include/ATen/ops/rnn_tanh_cell.h,sha256=DoD4LDfI5C95r33lHQ36YSDEbqcOel0mNacaYJtQB28,1188 +torch/include/ATen/ops/rnn_tanh_cell_compositeimplicitautograd_dispatch.h,sha256=Mu9nq6TFukphH6uQMph1-cKgTGpc7qJFDLrKMa5Uzgs,1189 +torch/include/ATen/ops/rnn_tanh_cell_native.h,sha256=mKZ2mwWX8-6DVdR_hnIMjPvhsJXmXiKK0BAF3pAY-go,901 +torch/include/ATen/ops/rnn_tanh_cell_ops.h,sha256=L_jDi_0p8Y2CcYg6eUveLWoHUTiJkMsUMqNR2mSJf9o,1758 +torch/include/ATen/ops/rnn_tanh_compositeimplicitautograd_dispatch.h,sha256=S7pqJVofnSkzdoC-pN5B1zElBUMLSQQdnOWN93ZZvww,1439 +torch/include/ATen/ops/rnn_tanh_native.h,sha256=1uWu_uDUskwUF1W9m7M-kQX5qNmTe9eeV2h6Axw8NI0,1151 +torch/include/ATen/ops/rnn_tanh_ops.h,sha256=6eaas8njrf4ooq9cCUqurt8Fpw26fudKnRrwZZawv4E,2973 +torch/include/ATen/ops/roll.h,sha256=Y5QhgZvOxlvx4u9Hi-KPgLmax2qCKt6mO0KlU-U8Vj8,4205 +torch/include/ATen/ops/roll_compositeexplicitautograd_dispatch.h,sha256=lXipB92Mg4riTG1_vXgMb14_tv_js1a3Antu0Aw5G5E,1482 +torch/include/ATen/ops/roll_cpu_dispatch.h,sha256=YVrguTeonBDLSvVixTaSajS-b9eZbsz15908AWn6os4,1133 +torch/include/ATen/ops/roll_cuda_dispatch.h,sha256=EX1-bNw6GhCaeZqibckT2xpbBgwunfowqzeLr_pGb5k,1135 +torch/include/ATen/ops/roll_native.h,sha256=eks9eTtMlvMlpmE_T1gs6wfWm8PyMTUaMrHPoNmSiO8,1016 +torch/include/ATen/ops/roll_ops.h,sha256=vP_G10mPReE77cpEpUE_y6msmHX-xDj9neHSzyh1nOI,2133 +torch/include/ATen/ops/rot90.h,sha256=m_0rYw6u0QxttASHHds_GEjtcGHq7x8Ho0jRFcSlDX8,1494 +torch/include/ATen/ops/rot90_compositeexplicitautograd_dispatch.h,sha256=pz3FgSORR4LO07R0ELZQ3Cv4XYamZttaFkhXA88dF18,1287 +torch/include/ATen/ops/rot90_native.h,sha256=WufEaaZDE7bgw5qnZismipAghFY-zog_2DzCJ3oxO-4,880 +torch/include/ATen/ops/rot90_ops.h,sha256=4gT7zNK_a5qH3WOGo0h6wwKPYfhzB7Ntf6H8l-zSSz8,2033 +torch/include/ATen/ops/round.h,sha256=ghp7jbFWm-Y7wav-ZBZNOq95EgUAiZMpMrp3Hv1y1BI,2311 +torch/include/ATen/ops/round_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Tu4Zl9DIWzGI_ztyt1j2yM6br_KIqqKWX7GSO4i6UJ4,1232 +torch/include/ATen/ops/round_cpu_dispatch.h,sha256=tMkOcMxju_dDaFdbJW1EWXR11auZHTaeCkvuceCtnfs,1508 +torch/include/ATen/ops/round_cuda_dispatch.h,sha256=0pzUNDYaNiL_w6o3MddOMqmo1onwXwdkD765bFSmANA,1510 +torch/include/ATen/ops/round_meta.h,sha256=T5eY_el-mxFZ8Vn2ZtO7IUiZb_NzjrEM6cOSOSGuh9A,956 +torch/include/ATen/ops/round_meta_dispatch.h,sha256=otEmONNKe41ZrBz-Dz8bwzYtjAbljw0unDWxrYODoIY,1510 +torch/include/ATen/ops/round_native.h,sha256=kjQSMNYB90LPkt8ArmPqpV5Q68Zpy9bM6D96HFQgNTM,1436 +torch/include/ATen/ops/round_ops.h,sha256=b6Izjd81rSRPBzJgJDv9QSS0LXegndLwDe_vDCJ8Ce0,4133 +torch/include/ATen/ops/row_indices.h,sha256=o9nZn0NvcMqpha-mVkSSFQdCjYSYEO1C-I0xwqiz02U,762 +torch/include/ATen/ops/row_indices_compositeexplicitautograd_dispatch.h,sha256=mtV3jVdtuO9axQOfyeLiy3Au9QIe_jl1iimzHp7JaF8,1023 +torch/include/ATen/ops/row_indices_copy.h,sha256=tb2QHVhgwWbze6HsImcA58qUsxFIkCrkq_ln3ucOtZw,1381 +torch/include/ATen/ops/row_indices_copy_compositeexplicitautograd_dispatch.h,sha256=-Ew0ig5VhEkc_yYjQ7s8TgYd-54Uh_iZZHSCAqYskQw,1141 +torch/include/ATen/ops/row_indices_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=veNw2jtcIlGpDwoX29vr_pW107ChnThpQPe0gzDrjTY,1054 +torch/include/ATen/ops/row_indices_copy_native.h,sha256=fEdIiytcms4CWzKthLPmKbUi_etE0CrpyRrOwm5I-gg,828 +torch/include/ATen/ops/row_indices_copy_ops.h,sha256=1xb1uDlmJyZ_CV9eeg_THDuh4IZ_96Roe-wJmjMWIEs,1861 +torch/include/ATen/ops/row_indices_native.h,sha256=R8dRIhV-lr2seMEWjXNZynBxTGkDm-ZakAy4eN8LPM4,813 +torch/include/ATen/ops/row_indices_ops.h,sha256=iQHlckPo28PSy1Ly4nI3MrCAQL4WBXsiJLNYylN0IC0,1232 +torch/include/ATen/ops/row_stack.h,sha256=kahBlM4-q49wxRz99mVKn2LbDm-q-aGrP_q1isDOxSw,1332 +torch/include/ATen/ops/row_stack_compositeimplicitautograd_dispatch.h,sha256=D72AykzujZss5C2Ufhj5pWUzXuPAdYO0BlWGmO8gpS4,1181 +torch/include/ATen/ops/row_stack_native.h,sha256=dSsMLMxRIbmvVeX33oFT7Ibu1OVEXW4fK3-evmrZg1Q,812 +torch/include/ATen/ops/row_stack_ops.h,sha256=1VxWuN0TRse0lG3olG1jxstxOesdKqeb6q6ulEKx9N4,1817 +torch/include/ATen/ops/rrelu.h,sha256=bmqBOt7onCiqT-jJMOlDRaGZSLwmUw5eOjvkGVmS3Q8,1608 +torch/include/ATen/ops/rrelu_compositeimplicitautograd_dispatch.h,sha256=OMGTnlclDSiT25pyd63aS7dCaIqDZ7crdnFEQhkaPBk,1377 +torch/include/ATen/ops/rrelu_native.h,sha256=xDrAzGdJGd2BMx5qgZMaBDNYj-6VF1N9qHV7_wrX-_c,1089 +torch/include/ATen/ops/rrelu_ops.h,sha256=Q7PmJDYM-ZAnNZ61T_Ye5tGbBDovGW7D1MiYLBU611k,2498 +torch/include/ATen/ops/rrelu_with_noise.h,sha256=K3n5q1eXTQniw0MA3h6yIpmZNEkejKDznV6Q-oe2llU,3432 +torch/include/ATen/ops/rrelu_with_noise_backward.h,sha256=aSKgzR4QSAj9arm4aevxRGeUZ9BidID3G5Bqtjg2uS4,2383 +torch/include/ATen/ops/rrelu_with_noise_backward_compositeexplicitautograd_dispatch.h,sha256=wl4evcR0LIk12O4OdsoKJrT6LKg10lBZU6rN09vQmG4,1670 +torch/include/ATen/ops/rrelu_with_noise_backward_native.h,sha256=cwl7MgEcztpDjf3BfD5meAwXSfTkL-KyDTIVrhpEyXY,1138 +torch/include/ATen/ops/rrelu_with_noise_backward_ops.h,sha256=pGlPN34EfFSjtDK5pQQjABz7M2FI0xWUw6bAHSqD7eY,2879 +torch/include/ATen/ops/rrelu_with_noise_compositeexplicitautograd_dispatch.h,sha256=ZAoQgJua52TFwVXkkJ3ftELM30NRSFAGNxi-V_4WTaI,1245 +torch/include/ATen/ops/rrelu_with_noise_cpu_dispatch.h,sha256=GRwXfGnQTHJwQ79aj2bAm-mk58eCMZSFMKwxvVqUfBk,1876 +torch/include/ATen/ops/rrelu_with_noise_cuda_dispatch.h,sha256=ptBOwhoYw24F0GvgPVtLoqeHj80CZG2LUKHAEot40yY,1878 +torch/include/ATen/ops/rrelu_with_noise_meta_dispatch.h,sha256=zc8JPzRWVS1JkKjyFHBoJqwvFGtlzweiH7GmtM7m4mo,1158 +torch/include/ATen/ops/rrelu_with_noise_native.h,sha256=YHn8WQ_ZZysGw9Zq9pyjm_KBN-y4HjDx6roxUpg7UGU,2368 +torch/include/ATen/ops/rrelu_with_noise_ops.h,sha256=2jng3KR15CsuF46_jRONyE5BesNXudh6I23W_kum2Uo,4938 +torch/include/ATen/ops/rshift.h,sha256=qgsM7uK6gUvmCleprWnI_BrGeobjfViMBTqdKDFF7mM,2244 +torch/include/ATen/ops/rshift_compositeexplicitautograd_dispatch.h,sha256=0pTXVKXMTK-sFdHkvh1V7jlM3BKh39AgLwispQMN_ZU,1398 +torch/include/ATen/ops/rshift_cpu_dispatch.h,sha256=KNyNENoUSFPfXprgo2mgwU6RS71LMaTsCOvGk2cLwls,1250 +torch/include/ATen/ops/rshift_cuda_dispatch.h,sha256=zYQm2xOyZ1eE9HwhPbVsJ54Y-3lbFw3SiuuGR7tvzJw,1252 +torch/include/ATen/ops/rshift_meta_dispatch.h,sha256=hq5LEIdT8yPOTEGq7bu5xVZO17GO-PF2JAT2B1zqkIM,1084 +torch/include/ATen/ops/rshift_native.h,sha256=0hBLz_e1qVx8YYuZ3c9tO9Ey3laBDZYIvmbgAaQ1UD8,1236 +torch/include/ATen/ops/rshift_ops.h,sha256=p-dePtjLTXw3-ioSq1NbQp42nT9LL_vZ2e8JSxHtpEk,4599 +torch/include/ATen/ops/rsqrt.h,sha256=OQFow5aN1yF6srkgYkDSWWpGEVcI7-uxi268AFIILRU,1410 +torch/include/ATen/ops/rsqrt_compositeexplicitautogradnonfunctional_dispatch.h,sha256=UiHL4fE5ezjXaTuOAWDxAqw-x4srIQ6ppdr_2X00D-8,1093 +torch/include/ATen/ops/rsqrt_cpu_dispatch.h,sha256=F6Zx-gU1GmJ-2as9JNgN1RUPyIs9BE0Yor_3miY0xjw,1178 +torch/include/ATen/ops/rsqrt_cuda_dispatch.h,sha256=2Na13kIC_KOGHPLtBxKL0cxSBRDgsxmf0i2nNcCEWa0,1180 +torch/include/ATen/ops/rsqrt_meta.h,sha256=msSUjsSMNzMEfZH_GKynmUNPEEHMY0nFBBhQ527sZZ0,820 +torch/include/ATen/ops/rsqrt_meta_dispatch.h,sha256=MiXXXnAS36rqNIoO4Lt51_gpDfawZWGjnSt-my6Qh64,1180 +torch/include/ATen/ops/rsqrt_native.h,sha256=4UOMKgClGpukdqxafE-JeBz3wa03BdgkuDlI01n7_aQ,847 +torch/include/ATen/ops/rsqrt_ops.h,sha256=_rdRzTuj4kE9WygHNeDX7LRr5KJb4_dYnURaGssfxJs,2291 +torch/include/ATen/ops/rsub.h,sha256=IbYZi_ubsf6FH1i804LEOOKmt7yvcZOnKwoJZcBKLY8,2439 +torch/include/ATen/ops/rsub_compositeexplicitautograd_dispatch.h,sha256=kVIxMNqzlUNAkV1X-2n6WW_aInOB7KVgbiOLiPPcIUg,1588 +torch/include/ATen/ops/rsub_cpu_dispatch.h,sha256=WQp_YNNrsLUY_2mNrfy2V8eoFvAV18AQ2nJztOUOiAk,1026 +torch/include/ATen/ops/rsub_cuda_dispatch.h,sha256=ag2judvZL_iuuzMcxenZF8pCg6g8KkuB3Mf50ZcYKso,1028 +torch/include/ATen/ops/rsub_native.h,sha256=vaigEO2IFHNl9du9eE1g7CSgkVcwLz_bT9JT5-ddKW0,1158 +torch/include/ATen/ops/rsub_ops.h,sha256=1hjUvdhgmehU8O8mqYV_vwYJyz8PXqcfDPma4RpLzPw,3648 +torch/include/ATen/ops/scalar_tensor.h,sha256=OG5uP7JLZye0yZIpHiPKqXNQNxwhyrc-njGV_SebmsY,1988 +torch/include/ATen/ops/scalar_tensor_compositeexplicitautograd_dispatch.h,sha256=NEEwV6P3W7QDmCDNEcV3Pu1eEa_0GcQS-oS8GD6MQwA,1420 +torch/include/ATen/ops/scalar_tensor_native.h,sha256=8Vb3HPzGeMNycXG3KovDQg2ZNheVixBR0l5Ivq79kDA,973 +torch/include/ATen/ops/scalar_tensor_ops.h,sha256=6-fhFAPbBk_VMVGcIgSx0aGSfdRbMmVDpRnkVCEhK_s,2321 +torch/include/ATen/ops/scaled_dot_product_attention.h,sha256=waNWxgRDuKZlPXiraeJqlUSVWNR4CE6HPiLy3pWp8fM,1397 +torch/include/ATen/ops/scaled_dot_product_attention_compositeimplicitautograd_dispatch.h,sha256=m1QUiPPEAAuJ8fJa4G74wRM96nRJtHGQXqjElJVLJdM,1254 +torch/include/ATen/ops/scaled_dot_product_attention_native.h,sha256=-L6xIvtJGKSTlfbQ-_NCkxKS75urLu77cXNyqwQLco0,966 +torch/include/ATen/ops/scaled_dot_product_attention_ops.h,sha256=ufGybKU04q7S5f4s0RswWHnt-ii3Hw0iCkgx2iiwYLY,1898 +torch/include/ATen/ops/scatter.h,sha256=sdzIICo2_j-JGIKGGeTpYdLkqf39kLyb9HkIbM6Amu8,5345 +torch/include/ATen/ops/scatter_add.h,sha256=dSZMGL9-hg9SDbmM4xs7eJzGBgNyh_0q3WXdsbi6tKs,1965 +torch/include/ATen/ops/scatter_add_compositeexplicitautogradnonfunctional_dispatch.h,sha256=5QyA_y1ln2sUIYJSdEbZQdI4tg-E5fFYgBfzr0knB28,1231 +torch/include/ATen/ops/scatter_add_compositeimplicitautograd_dispatch.h,sha256=pBJk34oXQO0dzY7T9W51br_nCMVcXaZQS1_XvXixaTo,1090 +torch/include/ATen/ops/scatter_add_cpu_dispatch.h,sha256=X7irNUJK6RBdC206W5o-pB9E1PsNR36xBmjLdFi0z7E,1454 +torch/include/ATen/ops/scatter_add_cuda_dispatch.h,sha256=aQScy9LFhdAPcrOjjv1FEPhhF1qCdCKxgI4mSyhOFRs,1456 +torch/include/ATen/ops/scatter_add_meta.h,sha256=cr-h3LflBxtjom5ugQ-SrVsgEliCdDGN2CsULqLz_jk,889 +torch/include/ATen/ops/scatter_add_meta_dispatch.h,sha256=KIASOmQRROY77Asek8YviiEP_t7zLQLg5Or9Gu8VS-g,1456 +torch/include/ATen/ops/scatter_add_native.h,sha256=tPTsOea8SUAqp_edIoIxsNT4BdkPIwp51FN1bZuulno,1050 +torch/include/ATen/ops/scatter_add_ops.h,sha256=FWJbuhcaW5Jqn4a1d2UN2XVbMnFiNXKhNDd2grO1IZE,3739 +torch/include/ATen/ops/scatter_compositeexplicitautogradnonfunctional_dispatch.h,sha256=guN5dRoEwYDMXLAfNzG3S9PfLjlpcDHEqMaw9qrLyP4,2030 +torch/include/ATen/ops/scatter_compositeimplicitautograd_dispatch.h,sha256=EDVn0F0AKgz1zhn7bsabLu3KqytEkKk3hRWqBKHeQrU,1210 +torch/include/ATen/ops/scatter_cpu_dispatch.h,sha256=Y-XCrhhA6_ylmW20J6kXU8_jWV_fifxLx_XQll68s4Y,3208 +torch/include/ATen/ops/scatter_cuda_dispatch.h,sha256=ptJgXfz2dRr7KHv4CV4un3x2CFNo1S89uJXbSOA5edw,3210 +torch/include/ATen/ops/scatter_meta.h,sha256=7-4b5p9SohXf3RNiPewk1bzcDZgnNY0UvHKOppOC-08,1491 +torch/include/ATen/ops/scatter_meta_dispatch.h,sha256=3gB9rZ8AXoUMcuVj6NXBcdB9b7lodZz9lcxzC1WnLD8,3210 +torch/include/ATen/ops/scatter_native.h,sha256=I9iSeph4xJWDXrfOPRTR0p1LtK0qeVCV08SucLHmk-Q,1894 +torch/include/ATen/ops/scatter_ops.h,sha256=9IANj8Bb0-cHRHOGIPUa4d1xkDxtaUbQj4XJMtnR2xg,11953 +torch/include/ATen/ops/scatter_reduce.h,sha256=P-qfQKnAe65txT9YYrmLvOgbkCDT7dWOj7OE1g0zg4I,2049 +torch/include/ATen/ops/scatter_reduce_compositeexplicitautogradnonfunctional_dispatch.h,sha256=x01ccrTho4U1MrLcMYKny2QfECBNbVrh7YBVZ-N9vjY,1335 +torch/include/ATen/ops/scatter_reduce_cpu_dispatch.h,sha256=0-cysGg9b3TwwYTCeauwbAXJJWFo9UjUXwKvfE1QnEE,1657 +torch/include/ATen/ops/scatter_reduce_cuda_dispatch.h,sha256=WvvmrBH0tbxbpLvuc1RzhhPH9v-WoQJvuqCDGBHN5wo,1659 +torch/include/ATen/ops/scatter_reduce_meta.h,sha256=LpqD2aXtxtd1J2ixFlFIG-1PKMmBkQNVsJ0Yfb-yBmY,940 +torch/include/ATen/ops/scatter_reduce_meta_dispatch.h,sha256=W5sJQVT7t4x-hAc_iKGwLe7RLW_US5s9JbuBEIEl67E,1659 +torch/include/ATen/ops/scatter_reduce_native.h,sha256=uYOl2oVbcR8-xo84rhrkoENo4xfx9L45iU_zSM0ot3Y,985 +torch/include/ATen/ops/scatter_reduce_ops.h,sha256=1wOgqs-MKtFuuA85TpyGL6_1WVleuqAYk3D3TQS5UAk,3486 +torch/include/ATen/ops/searchsorted.h,sha256=qJGnimAqzADjhFBSgxsRJledjno2u9WN82hbuW-q1BA,3918 +torch/include/ATen/ops/searchsorted_cpu_dispatch.h,sha256=T97oFww58m-OXUlIM7uaI7qv_aoPrR_TYVtUoQoxn-4,2386 +torch/include/ATen/ops/searchsorted_cuda_dispatch.h,sha256=10EBjqxA5a_0kR88yGRiKzIlMrlzu5sMLHB2uAVauD0,2388 +torch/include/ATen/ops/searchsorted_native.h,sha256=6SLQb-IVU4qv0_WLfUZnURCVYjquRgvcaqrq_KsFik8,2592 +torch/include/ATen/ops/searchsorted_ops.h,sha256=3kymaFbHoUBaSBCWySg0X7hv0zfuR6tTSFQ7NTvs968,5055 +torch/include/ATen/ops/segment_reduce.h,sha256=Yfyvw55KNGD1tkizeu4VZfMMfnMenjjcLqWXTZEmetk,2691 +torch/include/ATen/ops/segment_reduce_compositeexplicitautograd_dispatch.h,sha256=3ydRv6-fP1P6jGGFJBwgRedGM5pD0CDqeSMPAmlKt-c,1633 +torch/include/ATen/ops/segment_reduce_cpu_dispatch.h,sha256=13ol0tPxvXoHOpjJGOrv19nWbu8L5iXYWALF1xBx1ZU,1246 +torch/include/ATen/ops/segment_reduce_cuda_dispatch.h,sha256=3DdKpRlDliNNJTSlli0bKf-JXkybxhtQMli4OzMncNU,1248 +torch/include/ATen/ops/segment_reduce_native.h,sha256=yRJSHwAb512BW7hwPCWUuqQTKy6a5q7uEd4bi0JesB8,1327 +torch/include/ATen/ops/segment_reduce_ops.h,sha256=RoJVJZiDCl5woNaVpPzoN96Rc4Hwn30jBDtTSYAu6YU,3404 +torch/include/ATen/ops/select.h,sha256=4NDeO7llMICBO0ErO2swMT1dXwcRypTbdNwMx_VRAVo,1900 +torch/include/ATen/ops/select_backward.h,sha256=cci8wzIzE612qTyx51S-i9iBXmfiIabaGRLzli7-2LQ,5086 +torch/include/ATen/ops/select_backward_compositeexplicitautograd_dispatch.h,sha256=qiRlEyCwlxoiWePtVX8wXWfFofL5xSw9l6R7_XdCNaI,1600 +torch/include/ATen/ops/select_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=bo2r3jXA7T0W8xoPFnoExLzQi0Auf5AJA12avIxiQQg,1259 +torch/include/ATen/ops/select_backward_native.h,sha256=r4PKzswJB3QMYVVjnC_4HqNXTKtnZ6_0NVhob5k20Fw,984 +torch/include/ATen/ops/select_backward_ops.h,sha256=cA9rrppl0G2Zt4jzgm4IzR1MDJa0S8Y3sz2qdsl6ggM,2333 +torch/include/ATen/ops/select_compositeexplicitautograd_dispatch.h,sha256=CAriz-Wukhboa005Sfq88-VvfH0tx0t7BAxScrHEVaw,1139 +torch/include/ATen/ops/select_compositeimplicitautograd_dispatch.h,sha256=ryAQThoipuBYg2vOfc-jzJ-mOV-kx6QpdM5qHf5oGew,1050 +torch/include/ATen/ops/select_copy.h,sha256=JyqOVVZm0cavwhKgWZsxh9IO1nNg4VSpJSA115O2IgI,4008 +torch/include/ATen/ops/select_copy_compositeexplicitautograd_dispatch.h,sha256=mWQmENydby4TzvZ1N6Ex5eICLfBd1xlycLDUItAO_94,1432 +torch/include/ATen/ops/select_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ccRmau6MIaPbC51GEzDV88h0AtX2W9kuXJ2fklK_8Dc,1175 +torch/include/ATen/ops/select_copy_native.h,sha256=mxgkE48Do2VvXPUlbBH2VYd1j5mK5-vSiqKZ5bQGCA8,998 +torch/include/ATen/ops/select_copy_ops.h,sha256=JogO8p6_WduqfmP6xYMNabbVSsd-FXHAqG8uZ5dqIZg,2072 +torch/include/ATen/ops/select_native.h,sha256=GY1e4nDfXUWJ0GXNX-PcFYX8e9WRJNYMmGReG3KMi3Q,1037 +torch/include/ATen/ops/select_ops.h,sha256=M4STX5Y_8_pZYksjzEZWJCvQrPmR3uwLS-18dghBFRw,1976 +torch/include/ATen/ops/select_scatter.h,sha256=MrfI6oY9dVBUpXXU5QQLCje6C1lIcu_5abqT9CEhD-s,4449 +torch/include/ATen/ops/select_scatter_compositeexplicitautograd_dispatch.h,sha256=iTgT4fnr8dLTPKuSoivNDKOmqyLppKjTtoiznQyUu5w,1540 +torch/include/ATen/ops/select_scatter_compositeexplicitautogradnonfunctional_dispatch.h,sha256=LverrgwrzZMwqGJFH8UbKWauyxMaKVjR8BxPZMHjpFQ,1229 +torch/include/ATen/ops/select_scatter_native.h,sha256=T0FhLnC4OBmDaCUv4yk8xpsSJEfDVxfulmDn7VFlAek,950 +torch/include/ATen/ops/select_scatter_ops.h,sha256=s0gb8NqCH54P0wjpuPtmssRNMmaobAoipSbc_RkOknA,2227 +torch/include/ATen/ops/selu.h,sha256=ITghWDoBQ0vNIfM9nNc9_2TOXm0TPiLVjy4mhPPXX-w,1019 +torch/include/ATen/ops/selu_compositeimplicitautograd_dispatch.h,sha256=6gtCF172XP9dT1GnJ1A2D86Qlenm-TKtYwBSHrWWMmI,1065 +torch/include/ATen/ops/selu_native.h,sha256=fpiZZm6VxLKIQKxIlIXjWOJ7IiH3dLN9U47s4SWLmlM,777 +torch/include/ATen/ops/selu_ops.h,sha256=VNhMqvGJzPtkVtZVHjjXUVK_Nwrb4ssYGr5sgOiC0NE,1698 +torch/include/ATen/ops/set.h,sha256=_qV0ht3c1durpm9fSaDU_3S772ue-Kj5Jljt-Rtp1cU,9464 +torch/include/ATen/ops/set_compositeexplicitautograd_dispatch.h,sha256=9yomqD2b979Bo_uFmFxfv8MbRaYIHd70vhPWDgcU5ZQ,2724 +torch/include/ATen/ops/set_compositeimplicitautograd_dispatch.h,sha256=czaR2sRMEhrINebzvDiK50UDESWdtrRuzeCRwzI6pzU,1279 +torch/include/ATen/ops/set_cpu_dispatch.h,sha256=eAwk6b3skUvO9iJNcbeee6W5GwP93tzWg6Sqx3Q31fo,1412 +torch/include/ATen/ops/set_cuda_dispatch.h,sha256=rIXh27UjeR8VY2qym57p1fkdAUr0WHSLF5JE5Wnp29s,1414 +torch/include/ATen/ops/set_data.h,sha256=Zg11Enb8IekJV53Ci3a4iVLEwOkjT88xeRgLVgBdE3c,759 +torch/include/ATen/ops/set_data_compositeimplicitautograd_dispatch.h,sha256=Bwj-q11Y3fou0GuTETplXF4gC-yhcOJGuU6wanVXK8k,1037 +torch/include/ATen/ops/set_data_native.h,sha256=V85wfKgmMB6O5PXded7OW6hFL2Rycko_O71fDfI_X7c,749 +torch/include/ATen/ops/set_data_ops.h,sha256=dGdEs3cjwKAqdqeAf7M_fXJyEd6y-V9O7qHnHg6NV1M,1276 +torch/include/ATen/ops/set_meta_dispatch.h,sha256=RHyyh5XyHYoSOC7cFhrmeLVoI141jHzTZ_kGiNd1UME,1414 +torch/include/ATen/ops/set_native.h,sha256=JLpKJJpwJ59-A4rUtmEfj-1Xz-aC27DJPTiYv25LnZA,2668 +torch/include/ATen/ops/set_ops.h,sha256=Uk30BQHMCfVq2HmYSIY3YF2YZ9P9eHiHzNhOht55F9I,9938 +torch/include/ATen/ops/sgn.h,sha256=kI6oR1iPTFNCY0ytOgSu4dn06C0ikMp_MVhsWSvRBco,1251 +torch/include/ATen/ops/sgn_compositeexplicitautogradnonfunctional_dispatch.h,sha256=1TObnSBX3qT5dUptHoUXXP3VJszaggpyL7kEF06jZtU,1089 +torch/include/ATen/ops/sgn_cpu_dispatch.h,sha256=7tUfCjhN2aSAyvvuK6fimSVAOpjA0Z4hxZumJsEHbgs,1170 +torch/include/ATen/ops/sgn_cuda_dispatch.h,sha256=ZC17B79tEvNFWdp-1dIcBY620WMcYOxde9HPhDn1K3Q,1172 +torch/include/ATen/ops/sgn_meta.h,sha256=3v8kBuKKtLjcuSFRxBoX1GLjbNcDhD_TTKU-wsL4S4E,818 +torch/include/ATen/ops/sgn_meta_dispatch.h,sha256=3fqrjwCr3QJw8-qgTksL3KUsfG-0R_-yyIflCO1FP5M,1172 +torch/include/ATen/ops/sgn_native.h,sha256=eFs7GLEH00cRgnesHp4alyc2I7_4dYNiMyGSsjZCsmg,1368 +torch/include/ATen/ops/sgn_ops.h,sha256=uJ35vi2kPGZHPr7Zijsl3sNMCJzryCCYpc5H9zPfKvk,2273 +torch/include/ATen/ops/sigmoid.h,sha256=SP35S1-8mE6rellGA63DV4rlaaRjvEWcCZY3ZMRQI5w,1436 +torch/include/ATen/ops/sigmoid_backward.h,sha256=YkZuVSkqxLqmvWEqhJeu6XJmb8EnOtx3QRAoETKYofo,1664 +torch/include/ATen/ops/sigmoid_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=5URGXYFbUjtr7vBzO6FJs_vxzWcHkcmG0KNNNAAfCjA,1088 +torch/include/ATen/ops/sigmoid_backward_cpu_dispatch.h,sha256=vfip83mtbtU6AGKa7dyh63CNYgQHSO37Upvi2d7xIy8,1277 +torch/include/ATen/ops/sigmoid_backward_cuda_dispatch.h,sha256=mV8F5s0GKYeh0Ec6WQIbBJwFae9Rq6fzWrFjU8xSX58,1279 +torch/include/ATen/ops/sigmoid_backward_meta.h,sha256=vamPqk9RH5Oak3b-df8qn9U_K5piC2fGV8W8IvokxRY,865 +torch/include/ATen/ops/sigmoid_backward_meta_dispatch.h,sha256=y2vfQ1HFlEYyF15M-LHxMIGEd1oDFAJMH56jmRFAR6I,1279 +torch/include/ATen/ops/sigmoid_backward_native.h,sha256=7N1vDAB-dTGTDxvi7FpHZn8_OAP4YrdrW80iPhdrhG4,921 +torch/include/ATen/ops/sigmoid_backward_ops.h,sha256=J7EHByjZEKEeIsR9YInQi38LJJ6UG7i6OfHMqnIsZ4Y,2123 +torch/include/ATen/ops/sigmoid_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ircFvP3RbayrEdzzGJdj4isQC23LAH8uwBXNIwg7jSY,1097 +torch/include/ATen/ops/sigmoid_cpu_dispatch.h,sha256=uaMyqYr9JlH_IfkxrspT8jnJ-OVbGBKq_vzLCvZQW9I,1186 +torch/include/ATen/ops/sigmoid_cuda_dispatch.h,sha256=jbvBcLX-LB7czGp-dsc2hWTsGaKnWHGZ_uEy1sHWBno,1188 +torch/include/ATen/ops/sigmoid_meta.h,sha256=8SIac2c0v2xgH4po4wwEHNsgnWqTfSHHNAuYZidRxjk,822 +torch/include/ATen/ops/sigmoid_meta_dispatch.h,sha256=Kr5wm70cOY81sZIkCFmZo_RjaDDqUkvi0sMpLOK4Qv8,1188 +torch/include/ATen/ops/sigmoid_native.h,sha256=EPPFWLxMsbulfwgzI6f-OlnWA8jSy_bPo8gCvS-XFjs,1043 +torch/include/ATen/ops/sigmoid_ops.h,sha256=YMmoS_8uX_2NoYZJ4v2R_5vrZsOzrCTOKE29waccmDM,2309 +torch/include/ATen/ops/sign.h,sha256=zLFbcn3EpuP4EOceElsn7uBz4ywjCz4EJWk-WYbrNic,1261 +torch/include/ATen/ops/sign_compositeexplicitautogradnonfunctional_dispatch.h,sha256=u5_Tb2wgdJ-C-sV5BFAgNQfUczgODEyXahPQhxBLIYw,1091 +torch/include/ATen/ops/sign_cpu_dispatch.h,sha256=iDAbHMqmYggcQ1KZj_DR-4WkgDqOcJk-jyhnjVVrEws,1174 +torch/include/ATen/ops/sign_cuda_dispatch.h,sha256=06dl16Gr9zFlRtGGYGeWShdymHEGHUtamk6cx-cuCjc,1176 +torch/include/ATen/ops/sign_meta.h,sha256=a30zithYb093NjF5Pz0OHllx4usPxeWKrm4IPSgqWt4,819 +torch/include/ATen/ops/sign_meta_dispatch.h,sha256=qSV8Gel4hdVFhVal3JBLuWkGjTA8mjuRbvUf9CXIQ6E,1176 +torch/include/ATen/ops/sign_native.h,sha256=8tE4gCWZM7qoTbZ-c2krE4yL9a0SM89No3vebyi3F2w,1252 +torch/include/ATen/ops/sign_ops.h,sha256=8ZYl8YDO8ez6jpH538cwcTk4CFdsv2fEpBufQJy_C3s,2282 +torch/include/ATen/ops/signbit.h,sha256=8raUbE9A3BSRJ1i91JumIYD8wnNqKsalZvAJISz8Ewc,1291 +torch/include/ATen/ops/signbit_compositeexplicitautogradnonfunctional_dispatch.h,sha256=dKG0zQYQNc-gO1mp9pQSPaE5pB1GSQWxbW2vHAEvLLk,1045 +torch/include/ATen/ops/signbit_cpu_dispatch.h,sha256=ghacY1-BhVYCzdqth8Pd9Rp2h7IbWhaPlXC4EdwdkQg,1134 +torch/include/ATen/ops/signbit_cuda_dispatch.h,sha256=n-hXRnEcBB4D8ZsmkCR2obOslV5X9XpUDObYx0d8yIY,1136 +torch/include/ATen/ops/signbit_meta.h,sha256=z1DB1eigJJgdsalJngbGiWQ1j5GEQR7rkVCw_VcSKRY,822 +torch/include/ATen/ops/signbit_meta_dispatch.h,sha256=XreuC_tXOtOxF35Kj9p1M_ZxSwkciN3VixF8FEzzePk,1136 +torch/include/ATen/ops/signbit_native.h,sha256=lbVBTjMgal7sDsWmxHsfA2cxfwRyTYmJMJAXqkKdH3I,1157 +torch/include/ATen/ops/signbit_ops.h,sha256=o6yV8bSrInlsFurvNs-4i2l4blVYNKN1I7dIEL2jCUM,1807 +torch/include/ATen/ops/silu.h,sha256=1ucL0i4qRjjU0EcV1Jebc10looOuoj0NtGdeqbbi9X8,1397 +torch/include/ATen/ops/silu_backward.h,sha256=5MjcViFALzdxwZ19Tz0etmOC5C63uakz-vTjB3oPtTg,1616 +torch/include/ATen/ops/silu_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=MP0iPaXAmhRYKNEClxJ68u1zihsLs4U3qvfuCMYVTZE,1083 +torch/include/ATen/ops/silu_backward_compositeimplicitautograd_dispatch.h,sha256=uXCIoxGDuqVQRkHtamFdqZ5SASD36cPXTNoR68HwC1c,1057 +torch/include/ATen/ops/silu_backward_cpu_dispatch.h,sha256=hSk2vnrP-7FZnWaN5HwwpVilpHL1GKzAfLVMhP-8Sv8,1262 +torch/include/ATen/ops/silu_backward_cuda_dispatch.h,sha256=WW31UOYV-SvGSDTDRBREnjpSQnuQcR3hN-KE5MOEdqs,1264 +torch/include/ATen/ops/silu_backward_meta.h,sha256=Igtb9j_htA0cAiY78vRKc9504ZWL4rVirWXJkvZ-Lpc,860 +torch/include/ATen/ops/silu_backward_meta_dispatch.h,sha256=yv0qMaKhtyzKfb25pSCwwrCAqQxQzK7GFZ-TrPfxyj8,1264 +torch/include/ATen/ops/silu_backward_native.h,sha256=I2QaNP-pcT7itvu3RHRtmM8vcDbdFfNzsCXKCctyjwM,1108 +torch/include/ATen/ops/silu_backward_ops.h,sha256=pvoDjdylieWb14IdwhPZNFEZBLnkSogPu2IEml82Tns,2093 +torch/include/ATen/ops/silu_compositeexplicitautogradnonfunctional_dispatch.h,sha256=oIErYPa3AcRltGRKOeD957NGfUcPKHUjobCRn83yjYI,1091 +torch/include/ATen/ops/silu_cpu_dispatch.h,sha256=CR0lHgiTUIoIaKbFcBbIZjid8MPVghctLameXKtnNRg,1174 +torch/include/ATen/ops/silu_cuda_dispatch.h,sha256=oDfWNaebBAHaZLZr0trZb1fkKUXo2nSLpGw-_mDzrC0,1176 +torch/include/ATen/ops/silu_meta.h,sha256=MLkU3umju5GEIjiRvxfWyNjdiumdDka8aQq_2dCvDzs,819 +torch/include/ATen/ops/silu_meta_dispatch.h,sha256=53NEmZGzUUAcLKWJU_EYs2LW9xb0y1ua3ZeO2-xZSww,1176 +torch/include/ATen/ops/silu_native.h,sha256=qsvgPSg_Bo0jiIKvkZkMVtkkRfiilz54fscc8Ra6W5Y,971 +torch/include/ATen/ops/silu_ops.h,sha256=Hu17I0HyVv8hwpYXH5goasMa_swSnHwsZBv1OVhu_JU,2282 +torch/include/ATen/ops/sin.h,sha256=VYOnRrwlpzHP6xcD7uUFyMSbZp59nyEZ-1ffZJ96nqc,1384 +torch/include/ATen/ops/sin_compositeexplicitautogradnonfunctional_dispatch.h,sha256=my9QX9Z0GJfSw6Xw7SWHcUTYnrGcMvBE3-rKDnBNYno,1089 +torch/include/ATen/ops/sin_cpu_dispatch.h,sha256=CqBqUhuEbzir-8zZn2ittMKzg_C5ZdYdH3pDVDYsrj8,1170 +torch/include/ATen/ops/sin_cuda_dispatch.h,sha256=QdcgB7HLsRXnhfvj9bldaChoJQGMtV5YV5nkRrdqE0Q,1172 +torch/include/ATen/ops/sin_meta.h,sha256=h2k4PBLyLEaFfNnu4pFC4KvdKjq-b7fRHPnEYjSCaRM,818 +torch/include/ATen/ops/sin_meta_dispatch.h,sha256=bt00eNUqL72JRlh1GjxhnaA84NV7qDqX3NCM6m8nUuk,1172 +torch/include/ATen/ops/sin_native.h,sha256=QDG03Vdp1XpJf6nXO8COvnEm-OJWlMGlwk5ETRzYfjs,1307 +torch/include/ATen/ops/sin_ops.h,sha256=YbcZLqGv3LB_ha9iWdbT1vKvnZQTt3xk4FvvtjPjqDw,2273 +torch/include/ATen/ops/sinc.h,sha256=gH9yZTjoDcWJVNiig2cu2b1Oa601fsFb4LB7SiZXcMs,1397 +torch/include/ATen/ops/sinc_compositeexplicitautogradnonfunctional_dispatch.h,sha256=a0V7eutuY0WYLPqvm3SZumQzzJMG5mq991RYJoi5fnI,1091 +torch/include/ATen/ops/sinc_cpu_dispatch.h,sha256=FqWcZU8S_9w4P_aD4wWNZnvIn-r5TNVfJRg9_uAXyRE,1174 +torch/include/ATen/ops/sinc_cuda_dispatch.h,sha256=Q2G6re58qKB5Jjab2aOlxm-HHqRfpBkWqmgjH9ZnbGg,1176 +torch/include/ATen/ops/sinc_meta.h,sha256=j9PY7loM_Nq7YCm8cBROx73OKduh5rkX4r6jrrdfutY,819 +torch/include/ATen/ops/sinc_meta_dispatch.h,sha256=XcvMLSoKHE_bF6K_6QeIMoI84mNHs8-MGPKw-p413Ag,1176 +torch/include/ATen/ops/sinc_native.h,sha256=O6CyE7AEOqrP5qPLy0zHtvIwxoxMwLtMKDJE0_epAmQ,844 +torch/include/ATen/ops/sinc_ops.h,sha256=cREVeGspVDvP2pEingUp_uKwnAtCJVCn1_2spDvEYz4,2282 +torch/include/ATen/ops/sinh.h,sha256=JaBiBLWfTNW9ymS_YhLn2otw8DZziUkmqFnRrD37we8,1397 +torch/include/ATen/ops/sinh_compositeexplicitautogradnonfunctional_dispatch.h,sha256=F2lmt1N9UdgHiJ9g3CCQL1rvCV7w2hawN3ZijklOaIA,1091 +torch/include/ATen/ops/sinh_cpu_dispatch.h,sha256=Vc8OM0gMxhCOnPLJ_BiScfIoR6bZgWWOddo9DkIbTo0,1174 +torch/include/ATen/ops/sinh_cuda_dispatch.h,sha256=gnXdbzuNB30NiU22rTVMjRcJeY4c68RG63NGtAn0KN0,1176 +torch/include/ATen/ops/sinh_meta.h,sha256=EdjpUvQdtsjsPQb-efPAC14oL8BiOGdw8YcQXhkd0hw,819 +torch/include/ATen/ops/sinh_meta_dispatch.h,sha256=V4YQ_gr9S2v-SuwDe6nuvSP92YVhLz6VztHoV9sanqo,1176 +torch/include/ATen/ops/sinh_native.h,sha256=GQnPzrnjsWmHDq7hh89T3jRaHzawpeDWqcCawXUddA0,1252 +torch/include/ATen/ops/sinh_ops.h,sha256=WlRPU9UJXG4pjuxCY9Ja4Dl5Cu9qC3GOWyryIlhp8KE,2282 +torch/include/ATen/ops/size.h,sha256=VYkVzP-wQGpe1B_40VDtkWRpudg1Jgcs1ZbrJe9_5Wo,1097 +torch/include/ATen/ops/size_compositeimplicitautograd_dispatch.h,sha256=bmRy-OVh4YR3bjBt0z7-Oe0GfDwEaW2SppJoKQTmCj0,1092 +torch/include/ATen/ops/size_native.h,sha256=nm8d1ID3ZrB0ljZVMFV-k5hXlz7DjyE0iVAejedca_o,804 +torch/include/ATen/ops/size_ops.h,sha256=QLx9pZDU7hUG4-XLDrcHsNY9Bb_kTRiJuSt3xXcMYcU,1813 +torch/include/ATen/ops/slice.h,sha256=X3WJmttO7XXf7ua4AaCgRSE8FIl_GMA5ptFR8RLvvPk,2492 +torch/include/ATen/ops/slice_backward.h,sha256=GpdqDAgKSp54FTTBi8HRiuFvGGwSVdv2CnqV7JEoTgE,5709 +torch/include/ATen/ops/slice_backward_compositeexplicitautograd_dispatch.h,sha256=RpGYSUzCKh07LstQtjoGxfxNLyv3C0NkWdGK0sgMScg,2049 +torch/include/ATen/ops/slice_backward_native.h,sha256=KY3r9Z7wjQxxIus--_TT479reZ-aaZS38G1WWZb7dC8,1029 +torch/include/ATen/ops/slice_backward_ops.h,sha256=QbhELvlTKKXG3df5jTKW3XZSezthuX-elMDq4inpOc4,2569 +torch/include/ATen/ops/slice_compositeexplicitautograd_dispatch.h,sha256=8jjPzsvgjt_oHadl1LbNE6vI478sy7PjMCBbK-E8FLg,1335 +torch/include/ATen/ops/slice_copy.h,sha256=nxFp0KSYh9Ka8uzX0UaLnHo1hENaou-4CzJJ1lCDeRc,6325 +torch/include/ATen/ops/slice_copy_compositeexplicitautograd_dispatch.h,sha256=0gYggcdX4SCLlWeTx-GaHWvnz-HaqsnTQVmig8Ix2zg,1756 +torch/include/ATen/ops/slice_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=94wkloCmIt-De3qh4ldU2YG8BhblhwTfM1Xz3aF0wkI,1371 +torch/include/ATen/ops/slice_copy_native.h,sha256=tgvL13FGp9soWzh-Kw9ovvaeb00rjXal309VeJs7tnU,1080 +torch/include/ATen/ops/slice_copy_ops.h,sha256=jXcClfPIo6i1EumEyqFpeD1_CvBE1nbL78bHJ4gZR9c,2562 +torch/include/ATen/ops/slice_inverse.h,sha256=gN3zDtQEohd7f_qVDiS_Q0imhOJYOYmFFRTNKZFuIPE,2678 +torch/include/ATen/ops/slice_inverse_compositeexplicitautograd_dispatch.h,sha256=1psWKg-U8ecGPszm2urUWasfj88GHc5jpwtYjjdQggs,1399 +torch/include/ATen/ops/slice_inverse_native.h,sha256=EpyMQfrnx3hPvj7qnVLdhKOjK5aimsIHZLrAbuBrqGE,903 +torch/include/ATen/ops/slice_inverse_ops.h,sha256=DHAPiBUvL8T5jB16VSVXe8U86CSXqVptykL3X9SjxZ8,1666 +torch/include/ATen/ops/slice_native.h,sha256=qnDEVIbZH9TXgYUH9lnqJwgdnr0F6vgvFQlyTQfBEgY,852 +torch/include/ATen/ops/slice_ops.h,sha256=RAN3IqSjVMNBRRCEqNGfBccOp6CN0bLJXQH24k8UQCI,1582 +torch/include/ATen/ops/slice_scatter.h,sha256=SlpeXRVfbj2xgUVW_9YwvvyTkD-UIBOckl9C5rg0c8o,6712 +torch/include/ATen/ops/slice_scatter_compositeexplicitautograd_dispatch.h,sha256=OvWSqk5oo8H-rj4RgTKLXsN-kflm4FmUOp2c-XKjv74,1864 +torch/include/ATen/ops/slice_scatter_compositeexplicitautogradnonfunctional_dispatch.h,sha256=aUOOtrX4K7KQa4C4b01CEPu_-CIusJMSqO2wHCapDoU,1425 +torch/include/ATen/ops/slice_scatter_native.h,sha256=JSEqPhVSsF1NPe-53VOzh7yXYaXZtitplPd5TZUN42U,1101 +torch/include/ATen/ops/slice_scatter_ops.h,sha256=zW4jt_cZAGF2CQBZvBXnBoPOq-AK9P6kJxxG-yv3S7Q,2699 +torch/include/ATen/ops/slogdet.h,sha256=u9oduL3X6cpCcMxZalrbv2_S0_58oasXE8uxJaCxG0w,1573 +torch/include/ATen/ops/slogdet_compositeimplicitautograd_dispatch.h,sha256=Kb4VUK4gDwVMEoix36rvOEYSyt6dXkuoZ6IfcAYB_Ws,1307 +torch/include/ATen/ops/slogdet_native.h,sha256=7Y9Izwd6biFtIvvxDcXNr2TrkaFRM4fjqtnHxKBgxCU,887 +torch/include/ATen/ops/slogdet_ops.h,sha256=pKpMExaaBcY03iw0BQtx3DOWh8UgQp0HoAFKvJf1an0,2104 +torch/include/ATen/ops/slow_conv3d.h,sha256=OZoNbrkN_gL182n9MyhFFMuv3t9GXx2ZNju9UsdcDG4,6874 +torch/include/ATen/ops/slow_conv3d_compositeimplicitautograd_dispatch.h,sha256=H0Z08mDe2Jkgwn4PnaQ1ie-Z2wciDc6ECK0BvGPfQ40,2435 +torch/include/ATen/ops/slow_conv3d_forward.h,sha256=VrQQ6n4Hq3aCh4ZpgLuQZZNl5MfswzX_tNOFJbTyhxk,7004 +torch/include/ATen/ops/slow_conv3d_forward_cpu_dispatch.h,sha256=gZqnu5dFJXoy4EUBipdZMjMpDGUoVrXTh-8mz5JJ_Lk,2371 +torch/include/ATen/ops/slow_conv3d_forward_native.h,sha256=FNTkZWQwM16huCCTrtDaKmQIn93nEUGG1GKyFMti5rw,1139 +torch/include/ATen/ops/slow_conv3d_forward_ops.h,sha256=U928nbkJMrwniysQqsBIM2gjKbPHv0ipWsk0QJHAr3g,2951 +torch/include/ATen/ops/slow_conv3d_native.h,sha256=fSxIZv3A-zCossM_6m0It9k37avz6ND-1wzLuE2VzcM,1119 +torch/include/ATen/ops/slow_conv3d_ops.h,sha256=Tc10lplvojHfysDv67205tskm9LJVDGiODuEZ0f_ZGs,2903 +torch/include/ATen/ops/slow_conv_dilated2d.h,sha256=2xGzI6_LY_sG5g61-zG03h2dijWv_BmihzFVegVyzos,7934 +torch/include/ATen/ops/slow_conv_dilated2d_compositeexplicitautograd_dispatch.h,sha256=eW0V1iuUNwGhDlO3CJ-vEP81hkWZlWLaGHBDGtVdd8M,2125 +torch/include/ATen/ops/slow_conv_dilated2d_cpu_dispatch.h,sha256=bcu8OP8ZxzG99WtcymLHOMz-BLmG3t3SEZCT85hqJq0,1480 +torch/include/ATen/ops/slow_conv_dilated2d_cuda_dispatch.h,sha256=-fU1Pd3GUbBQTi4Ol1yn5is461DcbYj4ugrSpche0AA,1482 +torch/include/ATen/ops/slow_conv_dilated2d_native.h,sha256=vhFalh5Lb62KV2v5xPOkttPkp9h-MqqOMeTyxL1NdFM,1470 +torch/include/ATen/ops/slow_conv_dilated2d_ops.h,sha256=folH8I1PfOrJbYO_R0tz_q6IsQAWUYOwKcOrA6EyQJo,3157 +torch/include/ATen/ops/slow_conv_dilated3d.h,sha256=vEsBA0Fg_xnOK-illpUiGIZwb2cNUh4fYVhpPEaMvt8,7934 +torch/include/ATen/ops/slow_conv_dilated3d_compositeexplicitautograd_dispatch.h,sha256=kwYOH6GDrl-sc_vqKmdqhTxN5qV7qygO_JvDZCUt2K4,2125 +torch/include/ATen/ops/slow_conv_dilated3d_cpu_dispatch.h,sha256=1bp9ypNUL-Lz0SvZOW6TEaWN7__iEoRVXyu4sK3srZg,1480 +torch/include/ATen/ops/slow_conv_dilated3d_cuda_dispatch.h,sha256=uyv01ofzOANsiQiaQ1rc6vmO56KET_jQ-sLtOlXAMCU,1482 +torch/include/ATen/ops/slow_conv_dilated3d_native.h,sha256=NrkBD1o6Ga4YaYVxz8nOuTN8lbj7ovR3YCfD-cO8S64,1470 +torch/include/ATen/ops/slow_conv_dilated3d_ops.h,sha256=jd4tLqH0kfHRC_6lNtyp2P8dQ_clYhSAD1OjCnwlqPQ,3157 +torch/include/ATen/ops/slow_conv_transpose2d.h,sha256=SVs973WWW6Q-h2bZNlkdUOmEgSn-Fsv8-gHN_fWLaxI,8988 +torch/include/ATen/ops/slow_conv_transpose2d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=J1hOvGYXsLzKzDdnzdoP8YKMeKcn-N6XPDTGVCAsYDw,1639 +torch/include/ATen/ops/slow_conv_transpose2d_cpu_dispatch.h,sha256=K9yn_pb8AyIiSMAsc8BL1HaH9-ykjApc-tYndujko9M,2891 +torch/include/ATen/ops/slow_conv_transpose2d_cuda_dispatch.h,sha256=eVWBLJgSkrSMgEhXPhPxV8pfQO3buuS9gbun4_bt4RY,2893 +torch/include/ATen/ops/slow_conv_transpose2d_meta.h,sha256=iX0YK83sD5p9ellF_a3jB8cJUGmGNsIlmh9puoHznhw,1057 +torch/include/ATen/ops/slow_conv_transpose2d_meta_dispatch.h,sha256=TRonKVL6lXb8gM9sGg4hgkkqJC9lPlhL6MdMq6EMmIw,2893 +torch/include/ATen/ops/slow_conv_transpose2d_native.h,sha256=7rBuN9bebLF8-WCXYJYzx5ftVENwpgA05BrBxCNtlBc,1531 +torch/include/ATen/ops/slow_conv_transpose2d_ops.h,sha256=VIFrXgHwfughWbtaRcsYsLTU2F5j0nOU5CYYjVN1zTs,3411 +torch/include/ATen/ops/slow_conv_transpose3d.h,sha256=RXYwCSccUHOx7McihEqSi2p8uaJbuLWf0LRqpPGhIDI,8988 +torch/include/ATen/ops/slow_conv_transpose3d_cpu_dispatch.h,sha256=Px24CcNxo_HBbsmLK7GQsgfRPmchN2F8oGKjq7-gzhU,2891 +torch/include/ATen/ops/slow_conv_transpose3d_cuda_dispatch.h,sha256=5ynb_SdN9KQjAekTz_nkNiCUuIO0A8W4b_xkw7yXti8,2893 +torch/include/ATen/ops/slow_conv_transpose3d_native.h,sha256=z-T_X56A0kJb-uXuIhqLQcjCzCO0ppUz2Devwgty960,1860 +torch/include/ATen/ops/slow_conv_transpose3d_ops.h,sha256=sZK3QG9km8Cst30-IhMH_yqG4X8M5z02zNlRkQZVD04,3411 +torch/include/ATen/ops/smm.h,sha256=Dui1ezudAT1W44Zs94M6FR8WQSsUywfTHoB-SM12Ngs,923 +torch/include/ATen/ops/smm_compositeimplicitautograd_dispatch.h,sha256=XlRdtT8yCOJSm4CyMIop1P7VgaaO9r1CDwgJRQklqx0,1040 +torch/include/ATen/ops/smm_native.h,sha256=tDrFVit4KTnZYr635CttAd2qWgDT7xmNKTgcJCJRcHU,752 +torch/include/ATen/ops/smm_ops.h,sha256=14swLIWERBM6NAMQwdFomyahG24FM-9Gq8cILWRn6e4,1285 +torch/include/ATen/ops/smooth_l1_loss.h,sha256=JOGdfN3B-1z2e_8zEfZZeKjcUiiPr1b2ZNgSAqpWq-w,1814 +torch/include/ATen/ops/smooth_l1_loss_backward.h,sha256=eufX-PHT1sCRCKTunVX1ue7GPdtg5yx2DLsTulUwQNQ,2094 +torch/include/ATen/ops/smooth_l1_loss_backward_compositeexplicitautograd_dispatch.h,sha256=y9A5ksE8lf08yZymzdncwLCvf7vFh8fyIXH2FRKt2kw,1126 +torch/include/ATen/ops/smooth_l1_loss_backward_cpu_dispatch.h,sha256=a0e_9YvepTOINLcAtsok59lpIBJS-c0Z_XWwbKIPNpc,1307 +torch/include/ATen/ops/smooth_l1_loss_backward_cuda_dispatch.h,sha256=JSOTJUpMwZra6xCRFUJjTONZ4hwXcvTr5HNPliNh2yQ,1309 +torch/include/ATen/ops/smooth_l1_loss_backward_native.h,sha256=r3GY4U4CwQSermAUBn6RaYMmpaGvjpUfwy_31syy_9Y,1031 +torch/include/ATen/ops/smooth_l1_loss_backward_ops.h,sha256=eoLzKkTLqzw9IZdH0xcx5lVzFOPSzF-Dy1JRzl35mnY,2547 +torch/include/ATen/ops/smooth_l1_loss_compositeexplicitautogradnonfunctional_dispatch.h,sha256=-L04JUEMe6wsMHRv241ZJhaXwtaqBIr0Co-r5y8QsnY,1135 +torch/include/ATen/ops/smooth_l1_loss_cpu_dispatch.h,sha256=3c3QzHrOx9JvJJLsYRfqDIHG8i44gob57PvwepgKd1Y,1380 +torch/include/ATen/ops/smooth_l1_loss_cuda_dispatch.h,sha256=yI-Nd_JKwbQlyDZ51YOwKks64pnWxpakipmbuikXr68,1382 +torch/include/ATen/ops/smooth_l1_loss_meta.h,sha256=C6n1qKgEkn6cgUDfPk-_Erz0X3SQ9PcPtCksogfkPnA,888 +torch/include/ATen/ops/smooth_l1_loss_meta_dispatch.h,sha256=K8Q-lz04zW-1-eaKDvPq09NEZ4vnf4rowvJUZRvQWD4,1382 +torch/include/ATen/ops/smooth_l1_loss_native.h,sha256=adYdv1CprSjGEQ1fx48TPw7lqeMefBapxcWkG5rJOqc,933 +torch/include/ATen/ops/smooth_l1_loss_ops.h,sha256=qhrgdXZ9O0BicOx_VB_s3Q7thBMU_6s-0e40tKz5rXk,2261 +torch/include/ATen/ops/soft_margin_loss.h,sha256=CZAe8fC7c_mJnnGzXWvwqbbOmOWpz0p01-t4UUJQ1Kk,1721 +torch/include/ATen/ops/soft_margin_loss_backward.h,sha256=eeA8hZ0h3t54Gxp_VmLAkDDflVSJwl47vYApIF-idEc,2021 +torch/include/ATen/ops/soft_margin_loss_backward_compositeexplicitautograd_dispatch.h,sha256=3OaYNQflmlW96u_ZciQxXY8BXkYwXI3e--1eog6m1bo,1480 +torch/include/ATen/ops/soft_margin_loss_backward_native.h,sha256=SxPOek224M2BJLXoHNSxuPett-ASRlUIXKQ2taczb0k,1009 +torch/include/ATen/ops/soft_margin_loss_backward_ops.h,sha256=hyUZAyh77FP-AFp_8-I0id_zBWcZgZ4YrkvdPUCGhW8,2467 +torch/include/ATen/ops/soft_margin_loss_compositeexplicitautograd_dispatch.h,sha256=Js7yrTmAHxFSFCZxPyuoneQl9P3WJ8fomyCd_vJT86c,1383 +torch/include/ATen/ops/soft_margin_loss_native.h,sha256=iCUdsYRPTNMxsHC75KKEEMqTPIuwssgtzBTVlsWdQo0,940 +torch/include/ATen/ops/soft_margin_loss_ops.h,sha256=mDRg8Vc3Ic01gguVIEqVOX56VjCvY6DxMEz3t87i3-M,2173 +torch/include/ATen/ops/softmax.h,sha256=iNGYNDc8FbuD5VuP2ddBwwiiYBy1Xp0f8verRnkJgX0,1913 +torch/include/ATen/ops/softmax_compositeexplicitautograd_dispatch.h,sha256=46O1w5mWQI4WfjyXW_LaSSslPtX5rSHGqLpCKv-GaiU,1242 +torch/include/ATen/ops/softmax_compositeimplicitautograd_dispatch.h,sha256=CfuMA4IKcrU3-5stoOf43c58a_C5FHYz-cRhLWK_dI8,1212 +torch/include/ATen/ops/softmax_native.h,sha256=eyNTsogd-qWw_ZcAmFUxsI22Yi4ow2-BQ1MnKu-f2eg,1055 +torch/include/ATen/ops/softmax_ops.h,sha256=sRuPE3LIlxD0Owgn50vsqMjP1-LgLRTaSEvq3a0Sm48,2912 +torch/include/ATen/ops/softplus.h,sha256=WQIQXvBbDhRlo4u65s-jYwyECshKvU81Fml1iQiaJEY,1635 +torch/include/ATen/ops/softplus_backward.h,sha256=9As8u9ZqK__KD8_pIVdcex9fQBS1Zi7JqEoVOex6W7k,1965 +torch/include/ATen/ops/softplus_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=LQ8WcPOKD4cD44gqWEFcje3oXCIECs4KFxcBTiShShE,1142 +torch/include/ATen/ops/softplus_backward_cpu_dispatch.h,sha256=Zt4pnDVgwBxakHTojjHGKjL4grPdqwffo6IqcklUfhU,1439 +torch/include/ATen/ops/softplus_backward_cuda_dispatch.h,sha256=U9oL4QJOGoGgLIOdyTtHfhsufaYSwZiYuoyVYZa5548,1441 +torch/include/ATen/ops/softplus_backward_meta.h,sha256=6PbJJT2bgTZiBQto4lXwCs_MUwGsJGRqLgH7Z5e-Ef4,919 +torch/include/ATen/ops/softplus_backward_meta_dispatch.h,sha256=ux1k-QqXJKovpHEt6a-BPrwubeHmgw9ed5GBNh1cIbU,1441 +torch/include/ATen/ops/softplus_backward_native.h,sha256=k_o_c8uao-4x1zfiQfuryDAhD905XQa9hpiN7YkpRE4,977 +torch/include/ATen/ops/softplus_backward_ops.h,sha256=Zl7idp0GUkxgm3WMX_xiWX1YTgoO7FFn9xjb_SXbgPI,2479 +torch/include/ATen/ops/softplus_compositeexplicitautogradnonfunctional_dispatch.h,sha256=9mDOyaBcDTyZfiemmw16ypp4H00RgZGGJv1juV98F3k,1106 +torch/include/ATen/ops/softplus_cpu_dispatch.h,sha256=FwPzQNvjcJx-PETQgnAW8Zml0UBv0bh46oyRgD2q6w4,1312 +torch/include/ATen/ops/softplus_cuda_dispatch.h,sha256=Q7fMWJhIUxPDapjCZ3CuNh2nnwkYChGboYY-Xcz4p_4,1314 +torch/include/ATen/ops/softplus_meta.h,sha256=jbv4pH2VEBB6pAPungWbj6b55x-_XMIu2caz8fR1jGE,878 +torch/include/ATen/ops/softplus_meta_dispatch.h,sha256=yivAyKcxg0IEvF9y6JhdhvN17dzUEdVacrzLhThaexs,1314 +torch/include/ATen/ops/softplus_native.h,sha256=xfZVTk3RXxANJJhqrrApwtt9rd1OKkVQIvI6kS786ms,911 +torch/include/ATen/ops/softplus_ops.h,sha256=qMnql2vSo5vo9PFFEokiHdk1ldro-eWTeFMZ3119bIQ,2185 +torch/include/ATen/ops/softshrink.h,sha256=C9bDTHTegZaDVevI-vAfwJ_TLYGBHvIt7i_kRymLeXA,1482 +torch/include/ATen/ops/softshrink_backward.h,sha256=5GRklJt08Migh6qcpLSbNUqoO54fvnkK3remBjvZtUY,1817 +torch/include/ATen/ops/softshrink_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=bqsdxPOTcEVD4175ZG5qW9VCt6JC15X-dDcH-Ca8Oog,1115 +torch/include/ATen/ops/softshrink_backward_cpu_dispatch.h,sha256=gi9MA86VSsCuFJFMYi4eiwnAOT_n4xoXKx7Mw8y9Upw,1358 +torch/include/ATen/ops/softshrink_backward_cuda_dispatch.h,sha256=nFUNjr5hu2C5ShxXYseOJrLMCICzVr1SCaoRhUguPLE,1360 +torch/include/ATen/ops/softshrink_backward_meta.h,sha256=8EgIn4GdwK7CVVhBiqDm6NAp8qIfljHPiW-TIdk2MuU,892 +torch/include/ATen/ops/softshrink_backward_meta_dispatch.h,sha256=JrsiIYyR5mHtn8OvF5-qXRMocv6OCrDNi99yj_Jo1GE,1360 +torch/include/ATen/ops/softshrink_backward_native.h,sha256=E4MDvgfsjcwcMibfyxvQNiAvXuvp0FYl31KszEae6v4,954 +torch/include/ATen/ops/softshrink_backward_ops.h,sha256=SFQ7HDNOLe60i3gZneAF6nh74tfqwFPr4jDqOYuNINQ,2301 +torch/include/ATen/ops/softshrink_compositeexplicitautogradnonfunctional_dispatch.h,sha256=iBM_HHMHQiPW4p3oEtXXaPpceftMHk7vjBrkfD46DM0,1078 +torch/include/ATen/ops/softshrink_cpu_dispatch.h,sha256=DoVo-tzk7JlQsifxjhiOkF901xBtikSoibEgoz2KA7w,1229 +torch/include/ATen/ops/softshrink_cuda_dispatch.h,sha256=P9IX9CrgoMv1m0nSYl8ACTyjF8OoRbtE2qYxLBa2k68,1231 +torch/include/ATen/ops/softshrink_meta.h,sha256=AEGQC0SvrnRCD5E_TFvQsWbU3t0tVBDZzkprl-tjReQ,851 +torch/include/ATen/ops/softshrink_meta_dispatch.h,sha256=tnlenm1In7uWQtjMZvbBOd77ur7Xrf10-FR7163Z7VE,1231 +torch/include/ATen/ops/softshrink_native.h,sha256=2pANYqaGxU4zr6TKP6YcbQrcnr7BtC6qUey3_wFiAck,888 +torch/include/ATen/ops/softshrink_ops.h,sha256=91weszDfgIxUG3wn70b48P4FKnVtx54vD8QsCbKPN4U,2005 +torch/include/ATen/ops/sort.h,sha256=ESA8JESHU_BCLflxZ0MtDH7I3yq6IbJqfjkx6uHzd00,5530 +torch/include/ATen/ops/sort_compositeexplicitautograd_dispatch.h,sha256=8Xx7fUNc--WuE4piimGTB331SOZTOR7vGiAvlkajakA,1406 +torch/include/ATen/ops/sort_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ivpaf11vy5Bf1NoFrBeci8DW0H3lsvF6UFzuiIV3Nn4,1136 +torch/include/ATen/ops/sort_compositeimplicitautograd_dispatch.h,sha256=31H-pvXwx19p3b7pxU6B0DC6h9b1O42ja9Injltfyu8,1950 +torch/include/ATen/ops/sort_cpu_dispatch.h,sha256=rF9nyrcKpw0PPWumEMFTqc35NwaWJghNK4-9b4uKXDI,1452 +torch/include/ATen/ops/sort_cuda_dispatch.h,sha256=u9l87r60ejdIphyGWwXgVvoNZyRAEnAfthBq475wOlE,1454 +torch/include/ATen/ops/sort_meta.h,sha256=2Il2d2Mj1wVlJaQzGfcvjzYRO6pI5IvLkMeE8noY6KU,886 +torch/include/ATen/ops/sort_meta_dispatch.h,sha256=d6jio7spm_vOSB4nIvBjN71NXoN1moaLEiLU4vRmAjc,1454 +torch/include/ATen/ops/sort_native.h,sha256=4uWKXRcWDOs0sI36fqS70tAghg2iWU1lcvZvfvFSyg4,2008 +torch/include/ATen/ops/sort_ops.h,sha256=aGbkEvPafv6mwviIEKxWSJ9LSt-s__m4Y3F2QpY0ZgQ,7750 +torch/include/ATen/ops/sparse_bsc_tensor.h,sha256=Fq0Q5dPjRDVWVUvh2JjgvfUOzS9r6NUiUBD4Pq6V40A,3247 +torch/include/ATen/ops/sparse_bsc_tensor_compositeimplicitautograd_dispatch.h,sha256=e8f6GdRY7Oh4PQ-v2Fe7rJuGxmMZz8zZTrRQBv0WnTw,1880 +torch/include/ATen/ops/sparse_bsc_tensor_native.h,sha256=IsYVUPPf-2MDqOP7jrFI7Tmtawr0XV1GykGWtXZChcU,1276 +torch/include/ATen/ops/sparse_bsc_tensor_ops.h,sha256=VfWF4b2sIJRnw_UPjGs0d5QGCA_aNu3RGneTgjOG1mg,3381 +torch/include/ATen/ops/sparse_bsr_tensor.h,sha256=1vPZQdI5qudigZklILkvKGty_PgtvuxMQmpv56cxAhk,3247 +torch/include/ATen/ops/sparse_bsr_tensor_compositeimplicitautograd_dispatch.h,sha256=5fTZhnMAAz60Ih0Lna0iVcZ4QzjB6Gb81YQ-ALy_QTc,1880 +torch/include/ATen/ops/sparse_bsr_tensor_native.h,sha256=EhqWdef8wDUPF-XP9C-Q2Cuybyh89tSTJiGUeUezZmc,1276 +torch/include/ATen/ops/sparse_bsr_tensor_ops.h,sha256=hoM9Bic1cevjUsERjCabCvAAxTHf2JZwU5qnSL_rxQg,3381 +torch/include/ATen/ops/sparse_compressed_tensor.h,sha256=5QTL5V4WctniLTXYBiw-gtQJV8qRgLCYPfF2eez7Z0A,7224 +torch/include/ATen/ops/sparse_compressed_tensor_compositeexplicitautograd_dispatch.h,sha256=48KUixU8BXcDw2mAYtabEbDusE63-09-OBf2SnmEOtw,2472 +torch/include/ATen/ops/sparse_compressed_tensor_native.h,sha256=TEtoaO_V6UoMuNuKW3Qb-bdKdI_4EpEGaauFF8sWGu4,1306 +torch/include/ATen/ops/sparse_compressed_tensor_ops.h,sha256=6vpQ90aBC29ZtTdpBVAzE1mMQmqYyy19tPKGtYZRQWU,3498 +torch/include/ATen/ops/sparse_coo_tensor.h,sha256=cjDkL0aEKrNzzBd_KzXMVlBH3XS2GryDR85LLUNdxrQ,4535 +torch/include/ATen/ops/sparse_coo_tensor_compositeexplicitautograd_dispatch.h,sha256=FlSOVJCBBrWtjtlxhjwWQYbTN5VkmYKjoqYzaChmbHs,1433 +torch/include/ATen/ops/sparse_coo_tensor_compositeimplicitautograd_dispatch.h,sha256=cCCF-yLpuwa-V8dwSilqE_TjvthfAtcC1rwkBmIcH9M,1912 +torch/include/ATen/ops/sparse_coo_tensor_native.h,sha256=UXWH223UfzGtuJVVTOFYPCkWt8eYj9kIKV37jvBsvoQ,1614 +torch/include/ATen/ops/sparse_coo_tensor_ops.h,sha256=XpFMmYyR4qvIVh63VFSf7zVvpCfUe2DZTqlRx9VzMCY,5001 +torch/include/ATen/ops/sparse_csc_tensor.h,sha256=f3_js_8uoVbtDaxRtTzPWSLRzpMgpzjeENPpS_uvUcI,3247 +torch/include/ATen/ops/sparse_csc_tensor_compositeimplicitautograd_dispatch.h,sha256=j6KpXBjMBO4Upp7j9Z3JeAHOEGBnoG_6jvUF7a2_FpI,1880 +torch/include/ATen/ops/sparse_csc_tensor_native.h,sha256=H_0c5twzqVfwswcBVgQJ_NKggrhehFm_FOKkEyucChw,1276 +torch/include/ATen/ops/sparse_csc_tensor_ops.h,sha256=yM2g-NHHkMkebFXxTv43mfZ4HiNGkrW8CXaXRLUyHD8,3381 +torch/include/ATen/ops/sparse_csr_tensor.h,sha256=-Yh3gwqcjTcQVTlvwFafYcrE8GfVrWo3auTSQ2eRujk,3247 +torch/include/ATen/ops/sparse_csr_tensor_compositeimplicitautograd_dispatch.h,sha256=3CxhedXXEzGJF8ZTVxqwfFABmmsHea_h2ej1PqIh8xo,1880 +torch/include/ATen/ops/sparse_csr_tensor_native.h,sha256=P8G6aSqmi6fN-HZypmvazbyQn7_6VJXEDhAiFd3pH2s,1276 +torch/include/ATen/ops/sparse_csr_tensor_ops.h,sha256=ikJq-toeMTKrP0Yk9FVsKYCSSUSST4_TKueeP5FwosU,3381 +torch/include/ATen/ops/sparse_dim.h,sha256=GUvojgxOi9n7IQ8Wi7UkGt6-Vy3q7VZGQ1Gt3UmDM9k,761 +torch/include/ATen/ops/sparse_dim_compositeexplicitautograd_dispatch.h,sha256=iH0LCbBzWZb2tkGAz38GnpTgyWHNynjM5ad804GB8OU,1019 +torch/include/ATen/ops/sparse_dim_native.h,sha256=1Ht7WFWJlQouChzvR-tpHdYwpSrT6mm3SgYHRDLil-4,867 +torch/include/ATen/ops/sparse_dim_ops.h,sha256=gw79EMAX-HUG6Uv6NJqoZlRe8L6fdcc1O8CfF9EbWOY,1211 +torch/include/ATen/ops/sparse_mask.h,sha256=jZaxKGZqNzClx9hM1bkoG2s5RrBdJQhoKxbwYE2YUoQ,1269 +torch/include/ATen/ops/sparse_mask_compositeexplicitautograd_dispatch.h,sha256=55Y4UopDswrn4kWu6-2pASWEj-gumemWVMbGRTn7ekI,1181 +torch/include/ATen/ops/sparse_mask_native.h,sha256=eWPWLddFWPP4iE4hVzAH7UdIsjfN9sn2iPTCmR06Gk8,970 +torch/include/ATen/ops/sparse_mask_ops.h,sha256=Z9MV78AADtPk8-KYNxA7j0uotThmpnSg2VYzfjjpLbg,1997 +torch/include/ATen/ops/sparse_resize.h,sha256=bz1iI69PZNC89k5Ul7BwyLh8LRB1kczqG-9momSDgHI,1774 +torch/include/ATen/ops/sparse_resize_and_clear.h,sha256=f0Vy0hPeSo-NqjQjLoZL9bvwwqQn7BmM7CSKeA8oytw,1874 +torch/include/ATen/ops/sparse_resize_and_clear_compositeexplicitautograd_dispatch.h,sha256=WEMQ-GAMCq6Q1sR0xufvpL6rgD8eV7s9sYlhQvS-pJQ,1433 +torch/include/ATen/ops/sparse_resize_and_clear_meta_dispatch.h,sha256=ikbGTRc-SyOBmlYHFtbGeq8YcBkErfoxVwuVyG1sYks,1063 +torch/include/ATen/ops/sparse_resize_and_clear_native.h,sha256=1mGBXIXhnLk2gFwwT079F1Kjy7WUQQK4XOhwfTBbbXU,1117 +torch/include/ATen/ops/sparse_resize_and_clear_ops.h,sha256=up0xhBcF7J0hiyXoqTw_81mp9rKG2UE0Zh4IwCzi8UI,3125 +torch/include/ATen/ops/sparse_resize_compositeexplicitautograd_dispatch.h,sha256=BuIKFbO0XExqms6fNDLhZEKw2pN6gY6yegkV6awfdRU,1403 +torch/include/ATen/ops/sparse_resize_meta_dispatch.h,sha256=l0OXu6u-LnLweEeVji8E6J0rT6spqX0VqB0q7_NuiTA,1053 +torch/include/ATen/ops/sparse_resize_native.h,sha256=T9U2eEtUQg1oyaezslYuG1gwPyCv7-0PJj4MA-tlYbU,1087 +torch/include/ATen/ops/sparse_resize_ops.h,sha256=brrwab60RKJet_ZWIioSgYtmiGS757zEgesQqBKYioo,3035 +torch/include/ATen/ops/sparse_sampled_addmm.h,sha256=QAR_bQg-IN3V7WiLaAh9spJZjWgkzjD8cU1kDpJ2rz0,1981 +torch/include/ATen/ops/sparse_sampled_addmm_native.h,sha256=OpOTKN_7ZhOm8ZVBPk7nE_THrb-Pip66TGwT0c6maAc,1470 +torch/include/ATen/ops/sparse_sampled_addmm_ops.h,sha256=8W-K5jXUeghvwqy-_fUhaK3qPFEH0EuCIYqPFgmyfsQ,2566 +torch/include/ATen/ops/special_airy_ai.h,sha256=MKkP6VJEU4iFHFjposTGGlKu7LCD8Np59-bskHFFwuI,1344 +torch/include/ATen/ops/special_airy_ai_compositeexplicitautogradnonfunctional_dispatch.h,sha256=-7G3ug6FYHXseEgf8P60d3r4hbqULhdh4fe5kh6lYLE,1050 +torch/include/ATen/ops/special_airy_ai_cpu_dispatch.h,sha256=sa-UrcPW1OkJNcpvsazihoJH0V57xypf-jBAi8SGrGI,1149 +torch/include/ATen/ops/special_airy_ai_cuda_dispatch.h,sha256=TD-g4V8Q8qdcZNo2rh1ehEB3QVA7HgAfNUQA5dpvFrc,1151 +torch/include/ATen/ops/special_airy_ai_meta.h,sha256=sUmRfedCbiCVHeuyAG18AVwK5OggeeGRzIXfkEmROCw,827 +torch/include/ATen/ops/special_airy_ai_meta_dispatch.h,sha256=pJnE5QeYLKSEOyHYHYrQ3AlEhzoO9Om_mICEKB3hqCI,1151 +torch/include/ATen/ops/special_airy_ai_native.h,sha256=AF6HaJucbem1_Ke-ltgZNAL0gmA_e9yEzjxsayHFyMs,874 +torch/include/ATen/ops/special_airy_ai_ops.h,sha256=10MxwI2KYSeBIaeSo0DD0UUu0bsBhls1b_cL_gdeOKg,1837 +torch/include/ATen/ops/special_bessel_j0.h,sha256=cY7uADPldNTI3uLCYE8C9orzXvuRiqfpZxm4K56uv9g,1391 +torch/include/ATen/ops/special_bessel_j0_compositeexplicitautogradnonfunctional_dispatch.h,sha256=HIJDQ_TVxHKQ-4GFnDPkMjR03LtzfXbQ3vZTZEqHD90,1055 +torch/include/ATen/ops/special_bessel_j0_cpu_dispatch.h,sha256=6qL6BLh_BRu1cV_OoG1a8P_651ruCAJj7NEwONhuX_k,1164 +torch/include/ATen/ops/special_bessel_j0_cuda_dispatch.h,sha256=IwGeBMGBpWVHg_lHt_Y6zhKw_3BaH6TQXtKCl8B_hnI,1166 +torch/include/ATen/ops/special_bessel_j0_meta.h,sha256=IbpcQ8vd1PvtJbeGDJAs9_CWYgrkJuunXdhiIXG9LVs,832 +torch/include/ATen/ops/special_bessel_j0_meta_dispatch.h,sha256=rMj_YaFJCeRaNFBG8eE095I_QfZYUicteaMwpNvLJqY,1166 +torch/include/ATen/ops/special_bessel_j0_native.h,sha256=_3lUH5OTltPTBmqaHfoRtgJyz6pcD0Cu1l4WzWW9DiY,883 +torch/include/ATen/ops/special_bessel_j0_ops.h,sha256=yCLyfSfnWVrkjXmm7JfPTjF8eQdusLXUIl_5MsiR01M,1867 +torch/include/ATen/ops/special_bessel_j1.h,sha256=nVrhWo9tZ2LeAQ70ffygC3chiz3W-UXGTMMt_DoL37s,1391 +torch/include/ATen/ops/special_bessel_j1_compositeexplicitautogradnonfunctional_dispatch.h,sha256=7MEMk4haPBys3RRxygKoxsZC8cpAamzNDkROlNjn3MY,1055 +torch/include/ATen/ops/special_bessel_j1_cpu_dispatch.h,sha256=8vQMh612vheyQP3rt9LC3hTFtieNPplBsvKhGEK_wTc,1164 +torch/include/ATen/ops/special_bessel_j1_cuda_dispatch.h,sha256=D6aWFlKew1y9-SbVKnfilVYe2G3bjlifKk1OQl-hTj0,1166 +torch/include/ATen/ops/special_bessel_j1_meta.h,sha256=KvBfYUmDK005GRvpqmvxXURg25y3TFjPsL1O3sOW894,832 +torch/include/ATen/ops/special_bessel_j1_meta_dispatch.h,sha256=w_PCyg1v9J6j4TiNLuGFSUZ51Vai8dIYlRM_jCkBvyQ,1166 +torch/include/ATen/ops/special_bessel_j1_native.h,sha256=siQYFLzHqHaKH2uYCuXPVf1v_pGHYJB_pibnqgR8lPU,883 +torch/include/ATen/ops/special_bessel_j1_ops.h,sha256=crKMbapmMiBb1-qmEoMuYYlJWmxPkRlNwsArYw70Bng,1867 +torch/include/ATen/ops/special_bessel_y0.h,sha256=SmTm_yx9QYAxjtq_4iczvFLGMEGUXI-O0FTgO42pUv8,1391 +torch/include/ATen/ops/special_bessel_y0_compositeexplicitautogradnonfunctional_dispatch.h,sha256=w35AhLsOmXvfkCebPmtHiiGOsUo-IhotKwYA-jxstKA,1055 +torch/include/ATen/ops/special_bessel_y0_cpu_dispatch.h,sha256=ickI7Kqy82X31ISx14wv2g9MP6ZGhmKL5a87Sm8QE18,1164 +torch/include/ATen/ops/special_bessel_y0_cuda_dispatch.h,sha256=q4SgZymb03Pd5ecz20hJO1bTDN5pEg3B3H08B3NzgEc,1166 +torch/include/ATen/ops/special_bessel_y0_meta.h,sha256=BSwyOGNAaXLVDiUv5LHFHTlB734hy1WPDzkwDAFIdyc,832 +torch/include/ATen/ops/special_bessel_y0_meta_dispatch.h,sha256=jqRF-zM8B4RiMIC1ODKcuA7QE4q5YBRLqRG5XEl61po,1166 +torch/include/ATen/ops/special_bessel_y0_native.h,sha256=gBxl6iFgYp6uj1RH5I3xaeYOoJAKoQqz5IIQpK3Egi8,883 +torch/include/ATen/ops/special_bessel_y0_ops.h,sha256=s66Fr0lhZ6CfE-Ka-ROmQ4qsfEKeop00ZEbDnMhcpcI,1867 +torch/include/ATen/ops/special_bessel_y1.h,sha256=0PBbb0eMuooE_M7YaOPuPBpeaAfd8vTQ3jRkZ7TCrJU,1391 +torch/include/ATen/ops/special_bessel_y1_compositeexplicitautogradnonfunctional_dispatch.h,sha256=AoWAHK0WCK7Ns3OcwF61dGjTBZU0xQbCfXn3sBshAVc,1055 +torch/include/ATen/ops/special_bessel_y1_cpu_dispatch.h,sha256=fP8NP8yaL3pO_-NRGNNm02yK5kSURwAwmeOooibzLXY,1164 +torch/include/ATen/ops/special_bessel_y1_cuda_dispatch.h,sha256=EmWRwwYmJISV9rO9JREOFdQpJKeTtLr66-2fn8jWsFw,1166 +torch/include/ATen/ops/special_bessel_y1_meta.h,sha256=RuyvQGHmMXkDaZlSXZIlAukCsFuePn5wL7YQWEeeSYE,832 +torch/include/ATen/ops/special_bessel_y1_meta_dispatch.h,sha256=2NTcNGJODqG0lUbSyyNqYZXtAvF7WjbPhTBMrb56yAg,1166 +torch/include/ATen/ops/special_bessel_y1_native.h,sha256=5xEs4HoS_dDjU-yHR63qKcaUAnd5mH9_HBQ0ZhrHCqI,883 +torch/include/ATen/ops/special_bessel_y1_ops.h,sha256=lDIiNDqfNjlVD24Gzytqzq6ks8RAam0oe2NR9L8Z1Lg,1867 +torch/include/ATen/ops/special_chebyshev_polynomial_t.h,sha256=iJGF4l99S8Urp06GTtce3lwSDZX8Jj18R08VXBWhpE0,3345 +torch/include/ATen/ops/special_chebyshev_polynomial_t_compositeexplicitautograd_dispatch.h,sha256=CMWRtks5CwTapR2nAyka1rrj-5s0-vrtfz4CdM0oABU,1644 +torch/include/ATen/ops/special_chebyshev_polynomial_t_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Bmq0VrGdwio8dZGLdtByZcVDl2GNeSPz61p66vIp2tI,1087 +torch/include/ATen/ops/special_chebyshev_polynomial_t_cpu_dispatch.h,sha256=eo2T020-_H2Ek824jblyj4B57JqQJ_SzCWgz0UGPLiI,1260 +torch/include/ATen/ops/special_chebyshev_polynomial_t_cuda_dispatch.h,sha256=2n-MRqFfvjR4O4_Xor5KzvIszpJrgysdxkSOe3YlkD4,1262 +torch/include/ATen/ops/special_chebyshev_polynomial_t_meta.h,sha256=FPPWmsHFYQgDMM4fWpd-rFwPONlgVKW3uIZKV7WP75E,864 +torch/include/ATen/ops/special_chebyshev_polynomial_t_meta_dispatch.h,sha256=G8CpwZ9WH_CbQEgdNFyw_wmW4WtnlwkKCbfRA1tot7A,1262 +torch/include/ATen/ops/special_chebyshev_polynomial_t_native.h,sha256=kpgpZx6RkIf7cHkZ69SMkygKbb9mKEuMmqPnp_-h1MQ,1377 +torch/include/ATen/ops/special_chebyshev_polynomial_t_ops.h,sha256=spsBX-QUpZCMDuWGP20bprmbWsS5t3gHPcBscxUTR-4,4909 +torch/include/ATen/ops/special_chebyshev_polynomial_u.h,sha256=_LJaTeQhRUBaAoK9pNF7wns8rIitcn49PWqde4vIpdk,3345 +torch/include/ATen/ops/special_chebyshev_polynomial_u_compositeexplicitautograd_dispatch.h,sha256=KOHk9UHSNaQOOR4Ryzb2jixdl1mgICS6I89ekfAHxY0,1644 +torch/include/ATen/ops/special_chebyshev_polynomial_u_compositeexplicitautogradnonfunctional_dispatch.h,sha256=0BDs0Mm1RoqSSZldsN9P5Q7lIzfCDmK4KH7LY7ra8EU,1087 +torch/include/ATen/ops/special_chebyshev_polynomial_u_cpu_dispatch.h,sha256=chx-U_NRh_mloUaOW1qN9D4kXD0HTzYw-kdZzPlPyfM,1260 +torch/include/ATen/ops/special_chebyshev_polynomial_u_cuda_dispatch.h,sha256=0VSdEVxNAdk7HB-A714mozml-qtZ_51myydxTXudWVM,1262 +torch/include/ATen/ops/special_chebyshev_polynomial_u_meta.h,sha256=TGOqS3FKLG4Vif_r_D_w7c_8QtgB8o89zvxo6CPenbw,864 +torch/include/ATen/ops/special_chebyshev_polynomial_u_meta_dispatch.h,sha256=hCecgooNwQgrJp2YfibkeRYH7ING33ybLI1SLB0fOH8,1262 +torch/include/ATen/ops/special_chebyshev_polynomial_u_native.h,sha256=_LwdIr0IWI0KfEQB9wonDevxQbxXr6DRB8VFX_e-u0I,1377 +torch/include/ATen/ops/special_chebyshev_polynomial_u_ops.h,sha256=4DdlYHI2-DjgcE5K7KVvg1hYIxj6zi1IUrt_YaoayI4,4909 +torch/include/ATen/ops/special_chebyshev_polynomial_v.h,sha256=91vb59p9yLOKzhsukhKnkFhRk5Df5rtZL-6OvM6f0tE,3345 +torch/include/ATen/ops/special_chebyshev_polynomial_v_compositeexplicitautograd_dispatch.h,sha256=ttyA0gAEONgzWWs0F1AbImF290mwayU-5ieBoKnC0i0,1644 +torch/include/ATen/ops/special_chebyshev_polynomial_v_compositeexplicitautogradnonfunctional_dispatch.h,sha256=6G6Xa0eZKnfZy2FJNHdB6bpJFI9jC9WABYxTefl-liU,1087 +torch/include/ATen/ops/special_chebyshev_polynomial_v_cpu_dispatch.h,sha256=v2b9fl6B-VQIZVSjpwnZYnoYgvTF_-OL2NoQEtkQDv0,1260 +torch/include/ATen/ops/special_chebyshev_polynomial_v_cuda_dispatch.h,sha256=btPy-fatG8R6ybQEAU2uLerQjFtY8b9nh5LXKZ8IJlE,1262 +torch/include/ATen/ops/special_chebyshev_polynomial_v_meta.h,sha256=oi1pEYpPL_VpaTx0a1JlZh1UBcSn77DhOnEELJ0n5nM,864 +torch/include/ATen/ops/special_chebyshev_polynomial_v_meta_dispatch.h,sha256=h13wIrsCL_pt37YfBwq0jNyhDlNTbYLmR9Godh1vvf4,1262 +torch/include/ATen/ops/special_chebyshev_polynomial_v_native.h,sha256=mWSEAWqIHXOSg7nh-l53Xw8qwUlt5QDHzva4xQPhvCc,1377 +torch/include/ATen/ops/special_chebyshev_polynomial_v_ops.h,sha256=u0CYvF7eELw7eb5ANWnRe_JLacS_3OAOG404OOfUG0A,4909 +torch/include/ATen/ops/special_chebyshev_polynomial_w.h,sha256=vQ7ypfNy-yQhzbrpAGxSWA5-maTLR66RLjlvZ_rOyg0,3345 +torch/include/ATen/ops/special_chebyshev_polynomial_w_compositeexplicitautograd_dispatch.h,sha256=0s1D99p5-PVxVdU3DywLSnQt44wXisrYhCOYd4UaRTE,1644 +torch/include/ATen/ops/special_chebyshev_polynomial_w_compositeexplicitautogradnonfunctional_dispatch.h,sha256=NAzztlUcGLfdfVnN3Tbif9Pqz7EpBTrsJLzsLgzuF5s,1087 +torch/include/ATen/ops/special_chebyshev_polynomial_w_cpu_dispatch.h,sha256=sDHtOKJhoeC4ipgLe51V8YfeWL8bsWa3wrVLuST1uds,1260 +torch/include/ATen/ops/special_chebyshev_polynomial_w_cuda_dispatch.h,sha256=YaYUbG6Bn__zFCVC_1DciuyAcHf4qdkF9M68sPDhyTE,1262 +torch/include/ATen/ops/special_chebyshev_polynomial_w_meta.h,sha256=WO2H2QN4UVuYehuFEbAAt0fck2VV4KMu5smBmLhwW0E,864 +torch/include/ATen/ops/special_chebyshev_polynomial_w_meta_dispatch.h,sha256=fHuWQYnPFPp6P4HZgcbfURgYtIODBLoNKCFEzFF9nbQ,1262 +torch/include/ATen/ops/special_chebyshev_polynomial_w_native.h,sha256=ez7ESxhZfbbvubQmb1UQMdMblupq78iYn4O0_9SJu-M,1377 +torch/include/ATen/ops/special_chebyshev_polynomial_w_ops.h,sha256=HSdSOCjrJr5jHqJDeHGthKiiBna3lk5gid8rWLw27qc,4909 +torch/include/ATen/ops/special_digamma.h,sha256=9IJO14dA23fjx-JjlQWLEgAJy6LaYJJ7FQVU09QPg4A,1371 +torch/include/ATen/ops/special_digamma_compositeimplicitautograd_dispatch.h,sha256=4CuO5JDdCwqjq8xRfcNnpno08uyuoArqvShdeAOSklo,1202 +torch/include/ATen/ops/special_digamma_native.h,sha256=kgf2-xl0f0i_nbdlEU0iRKgwH-0l1HuGR330VIJBWso,826 +torch/include/ATen/ops/special_digamma_ops.h,sha256=ekbnyGkNy4pj9HO_wWzRx9Txp_BzZ_nzi_3y6mlq1S4,1855 +torch/include/ATen/ops/special_entr.h,sha256=RXPAAE9uYwxy-UJocNQCWLLc9atFtJVO3W2T9mf_3f4,1341 +torch/include/ATen/ops/special_entr_compositeexplicitautogradnonfunctional_dispatch.h,sha256=knUtBKaUL3jgFe8VKdhtqadvBESqolJ7cSWZYol8d10,1050 +torch/include/ATen/ops/special_entr_cpu_dispatch.h,sha256=bGXo7Ya-ymChtAEcH_Tep-GsGdxqgNS8YXv1AG0M94I,1149 +torch/include/ATen/ops/special_entr_cuda_dispatch.h,sha256=_AVglvrj-gf7TWtAWwVZsXBjxTArP1v158B2TZ4fMsk,1151 +torch/include/ATen/ops/special_entr_meta.h,sha256=wBqjCkTvuIB-TAqXalSb_2wGmTAEiXfrjI8xakvdFqc,827 +torch/include/ATen/ops/special_entr_meta_dispatch.h,sha256=_r-uhmQ4PRZTDPlkgszC90cFvjsPn7mGC_9lVPMClFs,1151 +torch/include/ATen/ops/special_entr_native.h,sha256=0HzQxu5DaiVcwGQcpdeags8RfvbnDsKWnicadCO_QVM,868 +torch/include/ATen/ops/special_entr_ops.h,sha256=PiVB22Vof6Vd70QSzSePmyCLWhysEKQU_tnMpaIkbvk,1837 +torch/include/ATen/ops/special_erf.h,sha256=6ahoHHfgVYsJP4ZGm63Ho0Os8up1znPnk2YpMGMFpuY,1331 +torch/include/ATen/ops/special_erf_compositeimplicitautograd_dispatch.h,sha256=h_bYrWWNMg7_241ad2NZb_y7Ln2TLL_f9zTFdsNUhUw,1190 +torch/include/ATen/ops/special_erf_native.h,sha256=Cjg0Tb1VfnDty0YBy9s7gw0jxq8DpXPqGyvW6EwPDuA,818 +torch/include/ATen/ops/special_erf_ops.h,sha256=hLz2KLgKijNjhSEws118zNS2YT3eqwwcWT4ubnOrxgk,1831 +torch/include/ATen/ops/special_erfc.h,sha256=MVlB88tp5SSSttvm5TQKm3a8rn5k3WGXLUt1pnrlbnI,1341 +torch/include/ATen/ops/special_erfc_compositeimplicitautograd_dispatch.h,sha256=37Lk32_Q6tBZ8ZPxUGk3dO4cV8Ys2gIxOeTgYdKYcYA,1193 +torch/include/ATen/ops/special_erfc_native.h,sha256=NFP0XORc10VRACdJMgXm6bVYPgngrqHjvzO8kUcKiu0,820 +torch/include/ATen/ops/special_erfc_ops.h,sha256=PZKxoYNF1Wv1_vo2LRTJOL7SNYGiNPXQo25tUHsDlBo,1837 +torch/include/ATen/ops/special_erfcx.h,sha256=45fWlZW2faNHPKXRAWwpQ0Ur77sY45NyYC1ul6tEt7w,1351 +torch/include/ATen/ops/special_erfcx_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Yv9ss35FvKIafiJxl8f-7CIoEARS5hQdfovaEztRTho,1051 +torch/include/ATen/ops/special_erfcx_cpu_dispatch.h,sha256=pvT9oVRajlxj7xyXpqFw_XmyfHeNUgoRXEqS_YD3rwM,1152 +torch/include/ATen/ops/special_erfcx_cuda_dispatch.h,sha256=ePdPBCLtPLd9MVaXceOYMAJcY0L6cZnNP48GN5LJs-Q,1154 +torch/include/ATen/ops/special_erfcx_meta.h,sha256=ChdzGzHgGQYG7Sz-k2VjbIx3UtpOQ3SBtKYizWQwYk0,828 +torch/include/ATen/ops/special_erfcx_meta_dispatch.h,sha256=8urEMFpSBgIXw0FaFboczc_vz6IprXBOwP0qD6_Pj5I,1154 +torch/include/ATen/ops/special_erfcx_native.h,sha256=Uh1nffGM6CwIUAEboudSMC2XGzh9WPPhbuo5LWtYybQ,871 +torch/include/ATen/ops/special_erfcx_ops.h,sha256=NF6B2cH7iAeei9wg9DYSacT60kNTdGe4Hf5NAswGnrA,1843 +torch/include/ATen/ops/special_erfinv.h,sha256=L2pSvMBg0m6tXuQ9hEEY-NL8IM1cpBau5NKOm7Yt0a0,1361 +torch/include/ATen/ops/special_erfinv_compositeimplicitautograd_dispatch.h,sha256=1gGaBYei-MK-ev5MsnyafoRSpU3vDeVXj6Sr8otNhPI,1199 +torch/include/ATen/ops/special_erfinv_native.h,sha256=ZZkxIDlrvSqIc7SZS2YKSjsnPGjJPivRYeI_MmSwaks,824 +torch/include/ATen/ops/special_erfinv_ops.h,sha256=LaGbcWJ7CMbeg5ZVqAIMA6MU6Gvc6i3-EEiBHIiqBCs,1849 +torch/include/ATen/ops/special_exp2.h,sha256=D0DhA2YQUqSNQqtlMl2oN44ske7fwCwbhQOKYtkMUFQ,1341 +torch/include/ATen/ops/special_exp2_compositeimplicitautograd_dispatch.h,sha256=8wjnmWcfzFRetLhRnCT2_9o5NRHiKDbi6HvKJDvg4N0,1193 +torch/include/ATen/ops/special_exp2_native.h,sha256=OQCoRcLfVHT07wVAZbZhyEN7oPLmDHiHOT6wpHLHyKQ,820 +torch/include/ATen/ops/special_exp2_ops.h,sha256=31MAdisy6DXWqIl2JN_HDUOTJVAjD13pFjMKXRSE3kg,1837 +torch/include/ATen/ops/special_expit.h,sha256=IfAeCldygOH2FBAFRgfvDhucQBYfcX0qWobtqY7dX2w,1351 +torch/include/ATen/ops/special_expit_compositeimplicitautograd_dispatch.h,sha256=GZ1IKb-EED6Htf5FqrPq-y6TVSBTy3jCAuywLAAYIMc,1196 +torch/include/ATen/ops/special_expit_native.h,sha256=VYjTdnRAoF48b05TV4tYn9ZnUF7nNEab22vuTJFq_sE,822 +torch/include/ATen/ops/special_expit_ops.h,sha256=NPV4sSVF010UxdJFgqU5olskTFuofJwjUkDz1rBxeHE,1843 +torch/include/ATen/ops/special_expm1.h,sha256=dInqsPU1Kyy5aj-32n96F07T14RJaHbMEImcKovOVVo,1351 +torch/include/ATen/ops/special_expm1_compositeimplicitautograd_dispatch.h,sha256=cqRJMp8nKMWCMJfMLW-BugIXP2_gYe_Fhd_sA0il-yU,1196 +torch/include/ATen/ops/special_expm1_native.h,sha256=tak2dFsNCadHVg3bZZ7FbnMW2VfqkdsigPmvICsNioo,822 +torch/include/ATen/ops/special_expm1_ops.h,sha256=xyqoNipDGNlYFymWoe6T6oC8rqGS6ImTlmcVwRxmh1s,1843 +torch/include/ATen/ops/special_gammainc.h,sha256=__o3k6dr-78wbXqejnCBWHQNx8h4gPcP3Wye0PEAUsQ,1522 +torch/include/ATen/ops/special_gammainc_compositeimplicitautograd_dispatch.h,sha256=9hMxf7TTAllEV1TlCouWvgHssP-KmoIOcnjYfArDeoU,1283 +torch/include/ATen/ops/special_gammainc_native.h,sha256=PuV1vZKpTgSxD6QPF6UhO7mpW6ifjWoZs3Q_tHJ3XZE,880 +torch/include/ATen/ops/special_gammainc_ops.h,sha256=pKoAX30_vcVhWjdZ0GbizFHoClBwUTzA5ToRWXVF2Fg,2033 +torch/include/ATen/ops/special_gammaincc.h,sha256=WpUl5Q3REdKvyukcwfg875csnXeOAjRtEySkx5ctl0k,1532 +torch/include/ATen/ops/special_gammaincc_compositeimplicitautograd_dispatch.h,sha256=X3onM9dOzUkb2XIhW8WfeXvnKwx0lAWDPHuhba9xlOI,1286 +torch/include/ATen/ops/special_gammaincc_native.h,sha256=uID8FzXSVHOE4ARBXCGtJv3Bcd1dvBiiXn8mLjXPrtk,882 +torch/include/ATen/ops/special_gammaincc_ops.h,sha256=3gLk5c5mGMVKhjeg0_TSOnMQ5FYza_0Tw3xmB2oXPKo,2039 +torch/include/ATen/ops/special_gammaln.h,sha256=DgU7Gs_u8yTomOfanFbOSjEPPfCPVALTCK4BZyTP1M4,1371 +torch/include/ATen/ops/special_gammaln_compositeimplicitautograd_dispatch.h,sha256=JuBT08-gBvUxQS0jkaxDu9WBgEX9YbEZJN0sObuhU14,1202 +torch/include/ATen/ops/special_gammaln_native.h,sha256=UxjjIJOOdbx88Yrl9nQyd0mdMvx-GGUk1tYZHwOAvH8,826 +torch/include/ATen/ops/special_gammaln_ops.h,sha256=mjU0VPWENCLHfRQxea-E7NAAd-DYdWrpyBO4XLi4ly0,1855 +torch/include/ATen/ops/special_hermite_polynomial_h.h,sha256=1NcmEnYt7W87F02aTb-W6dMbl_f8Ee371ffTBAB7Wa4,3289 +torch/include/ATen/ops/special_hermite_polynomial_h_compositeexplicitautograd_dispatch.h,sha256=cFds92tXmhWqwRhtJSTwRYNLMRmO-QfetMiydZlvazU,1632 +torch/include/ATen/ops/special_hermite_polynomial_h_compositeexplicitautogradnonfunctional_dispatch.h,sha256=HwIIwovuYv2jywNU4I6uJ-LOiXpJPxoaXYRn8abyhpY,1085 +torch/include/ATen/ops/special_hermite_polynomial_h_cpu_dispatch.h,sha256=9YN5hMQ2lkY3KiYP3ToUchuVZDK3UTf85MKxRDNu2sg,1254 +torch/include/ATen/ops/special_hermite_polynomial_h_cuda_dispatch.h,sha256=JIVcIsBYUnkyzwGl6qf7xYl6MK0Qd1VUNidvnSsbf6E,1256 +torch/include/ATen/ops/special_hermite_polynomial_h_meta.h,sha256=laM2E7R2zuw-IYj0HpOvS6N7kOZtX4zAujxbnNyi0EU,862 +torch/include/ATen/ops/special_hermite_polynomial_h_meta_dispatch.h,sha256=RF5oPI6jtmudoPVwYhL2DRo4jDdoEN7M7l6ZOF47Vck,1256 +torch/include/ATen/ops/special_hermite_polynomial_h_native.h,sha256=0gCHnV2iFGDuCelDUEZiIpLIdwiMhU1elq2KYWaLROo,1363 +torch/include/ATen/ops/special_hermite_polynomial_h_ops.h,sha256=jegH2yawZXiXoaQj1seWQDYrmh0WXS1jm6zsvLQ-OrI,4873 +torch/include/ATen/ops/special_hermite_polynomial_he.h,sha256=46rcDhsEvNR0jH0cbhyWYHJI9w5S8jR5taDv830rV_I,3317 +torch/include/ATen/ops/special_hermite_polynomial_he_compositeexplicitautograd_dispatch.h,sha256=JzgfiQf3R48UNUP1HoirMIlbCl_4125mqte-dJnCy1Q,1638 +torch/include/ATen/ops/special_hermite_polynomial_he_compositeexplicitautogradnonfunctional_dispatch.h,sha256=dgWcI8LCfw34FGsDBoiNaJJ-5IcUqv4zD-wlQqWEt0E,1086 +torch/include/ATen/ops/special_hermite_polynomial_he_cpu_dispatch.h,sha256=pS5w3_VAR3_7RPS8_nrV71hAZ6g5GSpwBGbCDQB3DMs,1257 +torch/include/ATen/ops/special_hermite_polynomial_he_cuda_dispatch.h,sha256=kjlAmbgG7LnND8GldBD6u9-UmjIun6RvwwgOXuYrpCc,1259 +torch/include/ATen/ops/special_hermite_polynomial_he_meta.h,sha256=oA8GWkZK93PKJ9QWgO1tC7ER1r1CZoDcsqGr1phVJ9k,863 +torch/include/ATen/ops/special_hermite_polynomial_he_meta_dispatch.h,sha256=jzQcA_qmq-mhz8dO-GUtLOrzKsgqO_UcX3DKPkk8pps,1259 +torch/include/ATen/ops/special_hermite_polynomial_he_native.h,sha256=JohrF2JLG3acHHMzOFGKIo1_a5DL04t2Q2Tp5xaz4lE,1370 +torch/include/ATen/ops/special_hermite_polynomial_he_ops.h,sha256=xgUZrso2yg2rCUqagiDSemof1B99ctX9BpHQV0vWr8k,4891 +torch/include/ATen/ops/special_i0.h,sha256=V8tw2tchIuEULXe0-eUbBSDpZXzxRiV9l2EU9GwU0dE,1321 +torch/include/ATen/ops/special_i0_compositeimplicitautograd_dispatch.h,sha256=3n55QHcEK9bGMLjEHjSVv8V2h_sAVDumKo3rPaEax5o,1187 +torch/include/ATen/ops/special_i0_native.h,sha256=5abiuijh7_P0oyFiXWZBOihUZ3vvWD2bjo8mTFEA4wc,816 +torch/include/ATen/ops/special_i0_ops.h,sha256=7-vH2Zn11RKKv8mC1-xxKIxFvsy1Dy-qQcgKfiXAzh0,1825 +torch/include/ATen/ops/special_i0e.h,sha256=AiDmRthvu28UkJ0uhGemNw039NQR9-HiINojgXUqWH0,1331 +torch/include/ATen/ops/special_i0e_compositeexplicitautogradnonfunctional_dispatch.h,sha256=3VrisaWO3mBL0fxjkqtcOAc7cs-Okwrhn6xgSaCORE0,1049 +torch/include/ATen/ops/special_i0e_cpu_dispatch.h,sha256=GxfnZlWwmNQPQeiXQQS8DD7_huUMoQ2R-bh52uuHXbA,1146 +torch/include/ATen/ops/special_i0e_cuda_dispatch.h,sha256=J5JkIaLlmUprurVjNnfhyg6uMMxmJ_o_jLWm7oyJaSU,1148 +torch/include/ATen/ops/special_i0e_meta.h,sha256=0ExsHhz9Fj-jFueA2lELSakT2bhTuFDjhVEwfUJpqZQ,826 +torch/include/ATen/ops/special_i0e_meta_dispatch.h,sha256=xuX65J-6Chc1qSh-34w2mpj7m2y9cmwDJVItIxCzesc,1148 +torch/include/ATen/ops/special_i0e_native.h,sha256=ptlotzB-Vc3z1xFHPqXAckuZ_mzOuiFk5QXwjCFLzdY,865 +torch/include/ATen/ops/special_i0e_ops.h,sha256=gveA6iHza7QALWZdZMOdl-H1axSyWDVKKWVX2R_Gkx8,1831 +torch/include/ATen/ops/special_i1.h,sha256=dWPvLqcPNANXxEZNr-JY_IKZ0r38aOPXPCNhkuhhDx0,1321 +torch/include/ATen/ops/special_i1_compositeexplicitautogradnonfunctional_dispatch.h,sha256=mrt6tRPObeVF63ubo-KZ69AZKBVWkuH-O24Uk14cD7Y,1048 +torch/include/ATen/ops/special_i1_cpu_dispatch.h,sha256=Y0wgEAJ4hvwEKVCVyv_xWaERHzvuylMER_WtI6vMTE0,1143 +torch/include/ATen/ops/special_i1_cuda_dispatch.h,sha256=m84PlCbbrlpf-3ChcQw9vZtmzufJJoCFNzwkYIpfceM,1145 +torch/include/ATen/ops/special_i1_meta.h,sha256=PU1wCD1BF47vsuIdtXczfrPv2aUpSEBr23i_ef0BrXI,825 +torch/include/ATen/ops/special_i1_meta_dispatch.h,sha256=Bjnpp5U9UiAODTkNfP1-xQ1EvMNbqbnbHWQbvFKYnVo,1145 +torch/include/ATen/ops/special_i1_native.h,sha256=NRjnBeA-UryUnTOZtgYTF0Q4ObEoYyOJAzY9OIRrkq4,862 +torch/include/ATen/ops/special_i1_ops.h,sha256=WXASfh66RmlI8GyeD3hXSgVNATgEfDeubOg64iQlaQI,1825 +torch/include/ATen/ops/special_i1e.h,sha256=zn4Tli6lsMaAm5R0Y53cp447QfP-Z1RRsAZDrvBWCsg,1331 +torch/include/ATen/ops/special_i1e_compositeexplicitautogradnonfunctional_dispatch.h,sha256=BqEOEMu8q64I2yD7LQF2PGUNioRq3lXqjoMHPVMyEBY,1049 +torch/include/ATen/ops/special_i1e_cpu_dispatch.h,sha256=ceBZ37IDl4vchgou6Uj1h2RPl8VKKMXIs0Idw2FQlUw,1146 +torch/include/ATen/ops/special_i1e_cuda_dispatch.h,sha256=k61rhKBWQb8XUujD_LLyDgKcaJHh2bNu1My-D47iZjk,1148 +torch/include/ATen/ops/special_i1e_meta.h,sha256=cV0GwP6QTlVIB_h3VusBJqsE4k-BXxEaHKJKTjNWBCI,826 +torch/include/ATen/ops/special_i1e_meta_dispatch.h,sha256=uVYSQsf-K4Wt4llywZfO1sOixWqszA_I83qqpVwwoqk,1148 +torch/include/ATen/ops/special_i1e_native.h,sha256=vNDtzED4QN7Y43gNq4OU66LQjsIcYWhw3zzL_M9n-Pw,865 +torch/include/ATen/ops/special_i1e_ops.h,sha256=qtyK-ZSeuvVxmILE1RFetSpnCLH_zTVkWOAcZufdPaU,1831 +torch/include/ATen/ops/special_laguerre_polynomial_l.h,sha256=d3qDQmFVD4NiwGpe2VHZu1f4GgeGNU1ebFW5V1-hHG4,3317 +torch/include/ATen/ops/special_laguerre_polynomial_l_compositeexplicitautograd_dispatch.h,sha256=zQa9y6C3lwWXXlrfc-vtIISa0p7HnfILyox9ITSYnvI,1638 +torch/include/ATen/ops/special_laguerre_polynomial_l_compositeexplicitautogradnonfunctional_dispatch.h,sha256=YRwNru4kV_zBgbxdiqJLTwQkdyjJD2drJR_PGivcneA,1086 +torch/include/ATen/ops/special_laguerre_polynomial_l_cpu_dispatch.h,sha256=TtzDW7eUPeLIfpbE2BDExH7h2HgfnQG2rFTlESs5VzY,1257 +torch/include/ATen/ops/special_laguerre_polynomial_l_cuda_dispatch.h,sha256=AsEF7pujBT2EyIg026177YUg5vsiIsfYpMz60tbjLcI,1259 +torch/include/ATen/ops/special_laguerre_polynomial_l_meta.h,sha256=cPV4gb9h11sEhaNqMgHvUJpzUmNtQMERLGfEVllpwo0,863 +torch/include/ATen/ops/special_laguerre_polynomial_l_meta_dispatch.h,sha256=hYAaJxNUztIQAvtQrC_r-3NUoQCpsMvVQKwoBaok-4g,1259 +torch/include/ATen/ops/special_laguerre_polynomial_l_native.h,sha256=9OgmPzBHY3dLkoXAzZsXTYThGs6hGBswHuZhswq-imE,1370 +torch/include/ATen/ops/special_laguerre_polynomial_l_ops.h,sha256=Whwoe1elsEIF11L8FHyCea3a316uWL_xfaDJs2s0eFA,4891 +torch/include/ATen/ops/special_legendre_polynomial_p.h,sha256=wzi2qKGHy3zT1Rvw54QCUc0IIY_ERim55eXySATJizs,3317 +torch/include/ATen/ops/special_legendre_polynomial_p_compositeexplicitautograd_dispatch.h,sha256=Vj0qNn2yK5a85WiguTBC9_DQsh9fNE0kRrYQwXaDQiM,1638 +torch/include/ATen/ops/special_legendre_polynomial_p_compositeexplicitautogradnonfunctional_dispatch.h,sha256=zQn1sTcWMqgwcA2B1EroYjdYXkmFkbuAnzQksEPyXJI,1086 +torch/include/ATen/ops/special_legendre_polynomial_p_cpu_dispatch.h,sha256=oW1fsBQmec4oQINts47Kd9-o6u73zvvl7Ri0nhkgHEw,1257 +torch/include/ATen/ops/special_legendre_polynomial_p_cuda_dispatch.h,sha256=GZxWhCAJDDYWDiKakhSHXda3tvasNIhEQ6DqIPPnACE,1259 +torch/include/ATen/ops/special_legendre_polynomial_p_meta.h,sha256=dcgTa9JH6Lq8W-FZtpAWNv5nuGixKPwhl5R5ekk3K60,863 +torch/include/ATen/ops/special_legendre_polynomial_p_meta_dispatch.h,sha256=Zc6J-UVkbZxBgzvCrlBJ06ZZWtaSSgsIdpSYZYTFmAI,1259 +torch/include/ATen/ops/special_legendre_polynomial_p_native.h,sha256=-L4lBwhaYdTC_XU-n8KdYoQUnCjRM-sbzbSZyZq3sfM,1370 +torch/include/ATen/ops/special_legendre_polynomial_p_ops.h,sha256=Z5LNX6Vc2y9nbAxepUAXxnmZf0EgSzVvUPwkhOjbz9A,4891 +torch/include/ATen/ops/special_log1p.h,sha256=LPDXi2sQG8TPR_yN9SqU03RNqlbMPXG6C-0xxnuj7yc,1351 +torch/include/ATen/ops/special_log1p_compositeimplicitautograd_dispatch.h,sha256=Q-LTYFipCB19V1MOGxmYDHPqS4o1ZYbAeS0NClW4LtU,1196 +torch/include/ATen/ops/special_log1p_native.h,sha256=N6wzYMzGfga4iN7JT9XMQpai5GqToK81wJeOC_MJYPo,822 +torch/include/ATen/ops/special_log1p_ops.h,sha256=UBu24OIQdUJssex3NHCv89cbVW1PibAnlXqhCyJ1szM,1843 +torch/include/ATen/ops/special_log_ndtr.h,sha256=XR1el6-ZDMs3EHkt3mRQT4VOB58hEBmhsTInGQ6pteQ,1381 +torch/include/ATen/ops/special_log_ndtr_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Cb6-4Q8srZvKfKO7dd5eMrprniAdjMNFIrXoxyWbMDY,1054 +torch/include/ATen/ops/special_log_ndtr_cpu_dispatch.h,sha256=0B29Mq593nEQSQO9cRDjSHDYLXnJDcdkKv8cTAjldL8,1161 +torch/include/ATen/ops/special_log_ndtr_cuda_dispatch.h,sha256=LyRmaZ-gHZfFQL5Ic8bGBXRWWX1gG1SqV9GdyrUWCPM,1163 +torch/include/ATen/ops/special_log_ndtr_meta.h,sha256=LCNGFzhV8GoLJfw0SF5q-BVTQDaMiiZ3SEe5n37jf84,831 +torch/include/ATen/ops/special_log_ndtr_meta_dispatch.h,sha256=ATZOolnKmlYQNhenxRpD79_MtkwK3dsK57hzwns27hU,1163 +torch/include/ATen/ops/special_log_ndtr_native.h,sha256=4aI5f1hpDS_CICKlpJPx1GEDQkLR4yxeTG9a7YtHZyY,880 +torch/include/ATen/ops/special_log_ndtr_ops.h,sha256=B7Za6IEfmGfoK1sZHJR1zjarb0uBNWMgxQQUb7qNxwI,1861 +torch/include/ATen/ops/special_log_softmax.h,sha256=vr-KtHqKWHfhnYdaaA4ebVvAGCEcufCRDkKKWGPk26w,1058 +torch/include/ATen/ops/special_log_softmax_compositeimplicitautograd_dispatch.h,sha256=5QPmA_MAk8jPxLQJa9ukDxRf1mN_97u4cRmAwvt-Xzw,1098 +torch/include/ATen/ops/special_log_softmax_native.h,sha256=UtpvAU0SmzIPYd3EQDSIzJRU4n06yfXbsm7jSML04Fo,810 +torch/include/ATen/ops/special_log_softmax_ops.h,sha256=j0O-9aUTYCQL0bQlyGGBdWbukAU-rquC6s85QWkH2mM,1432 +torch/include/ATen/ops/special_logit.h,sha256=gpYJ0ahwOwnb7o1lIc_ouvD5lC7wf5WeqGTv3aSqXWw,1534 +torch/include/ATen/ops/special_logit_compositeimplicitautograd_dispatch.h,sha256=HY7VaQ-yErMjHLHtW9X5w1gExhkXqXuRQenepXLdb0w,1313 +torch/include/ATen/ops/special_logit_native.h,sha256=OpLB3EcTobPPvvF2LfEEOH7vshqr5gRS1Uf6iPmkIG0,895 +torch/include/ATen/ops/special_logit_ops.h,sha256=gwqQnm_wGQteN38TO1ud6f7XE7pr09hhCUhXzv5AvpY,2043 +torch/include/ATen/ops/special_logsumexp.h,sha256=wIufsM4U3zr00vKAt5UzNpKFxnE3-klgLz7LRzCK4vE,1646 +torch/include/ATen/ops/special_logsumexp_compositeimplicitautograd_dispatch.h,sha256=WOjfZnAfdgKygIcPw4Vjo-8cSwxOuT81YWzC1cAy0JY,1325 +torch/include/ATen/ops/special_logsumexp_native.h,sha256=K1TJKA-YX-3EiR0DMKqTxT-kXIuKwkjtFu2k2JjC0Qo,906 +torch/include/ATen/ops/special_logsumexp_ops.h,sha256=ZPuJ59IX93_RRn2x94J2QlBlCmgqbcgULlOKP5U6Xc8,2117 +torch/include/ATen/ops/special_modified_bessel_i0.h,sha256=knoeapUhorXt-clWZTWSxGQAQNa_XeFz_TvJp2l8xdA,1481 +torch/include/ATen/ops/special_modified_bessel_i0_compositeexplicitautogradnonfunctional_dispatch.h,sha256=zN5sbALb8wtJO5uH8_vhfoOazkeCZRH83ExIrlf35sE,1064 +torch/include/ATen/ops/special_modified_bessel_i0_cpu_dispatch.h,sha256=OFoN8Jn2hEmbQtHxIshi2_oWQCqhohJdStcyu6SSDDM,1191 +torch/include/ATen/ops/special_modified_bessel_i0_cuda_dispatch.h,sha256=bfAgVIY-Fi3AF2E6M9iHTO2uaJBoiaIjr_LIcZAOtOA,1193 +torch/include/ATen/ops/special_modified_bessel_i0_meta.h,sha256=CCzP7Zk68Ap_5MMWNrlVOJ66pQBzSDLLZKlHm2EOWsI,841 +torch/include/ATen/ops/special_modified_bessel_i0_meta_dispatch.h,sha256=BSPF5owQRqjZ0z7xV9mBZg3QjiZcBQzwmHmHFYnwTwY,1193 +torch/include/ATen/ops/special_modified_bessel_i0_native.h,sha256=Ln4J-88nEPZHjcZFRIxchk3L7v7DUqyUb8Y_AZKFI6c,910 +torch/include/ATen/ops/special_modified_bessel_i0_ops.h,sha256=Ywdb-YXrqtferIyt5TIanWNKEdyak4g0NdbgH_h43uI,1921 +torch/include/ATen/ops/special_modified_bessel_i1.h,sha256=LYNc8-LLpZMkKiHq9WyTD2RkWZLx6kdbzHe8SRa956M,1481 +torch/include/ATen/ops/special_modified_bessel_i1_compositeexplicitautogradnonfunctional_dispatch.h,sha256=edVq2hBXhza-n_CY65P6XurYwd4i5FqLTQmcnNB2PMw,1064 +torch/include/ATen/ops/special_modified_bessel_i1_cpu_dispatch.h,sha256=700lq2SYTuB9__dI2JBdNpNPMFqdBkddWEbt_kbr7dU,1191 +torch/include/ATen/ops/special_modified_bessel_i1_cuda_dispatch.h,sha256=1X2WFpfQUrHo2DuJ6ZEykJfIKwEXiFTdZnzRM1Pk5d0,1193 +torch/include/ATen/ops/special_modified_bessel_i1_meta.h,sha256=m-kw-FW7K2nJXHpR1f8z8nu1_4TzJM01IyCfNQOOcYA,841 +torch/include/ATen/ops/special_modified_bessel_i1_meta_dispatch.h,sha256=xv7lWaPMJMDHRHEHY-cDsNOvF24b_bhkncsvIg4ShJU,1193 +torch/include/ATen/ops/special_modified_bessel_i1_native.h,sha256=vtha295SK-CFPY2KV9P6rjaJUkQiDklqQLm5MopMCho,910 +torch/include/ATen/ops/special_modified_bessel_i1_ops.h,sha256=cFTDzpIYQD-svJB1P_7zVuq2pwUw6ZKfM8eKUEJrt20,1921 +torch/include/ATen/ops/special_modified_bessel_k0.h,sha256=w3DZNj4M2QXWytNNhdbhiGzvxLXS_HKWrftC-PBwJIU,1481 +torch/include/ATen/ops/special_modified_bessel_k0_compositeexplicitautogradnonfunctional_dispatch.h,sha256=-YR4C87jsSDyzHVfWcVVBkCMQHE-e83qc6Hfa9DpmLY,1064 +torch/include/ATen/ops/special_modified_bessel_k0_cpu_dispatch.h,sha256=jul6H6LZuQh0zRiIykd20XPRkID8f-yXiYJfS0U4Lq8,1191 +torch/include/ATen/ops/special_modified_bessel_k0_cuda_dispatch.h,sha256=4BhmrN79HBpfcaT4nhUuDx62R_l7KJ24bjXuEDAhVVE,1193 +torch/include/ATen/ops/special_modified_bessel_k0_meta.h,sha256=djU-vyM0PaGUQKqm0iauAQ24S666LwxvOxR2T5XHGkU,841 +torch/include/ATen/ops/special_modified_bessel_k0_meta_dispatch.h,sha256=4xmjPpiHA82rC8v1V9rR5SbZoXDLjSZXIUhVUmEIy0s,1193 +torch/include/ATen/ops/special_modified_bessel_k0_native.h,sha256=RweHONzDneJRDmzTYHCA776LRSmDGKe3KzGuLHt4WEg,910 +torch/include/ATen/ops/special_modified_bessel_k0_ops.h,sha256=9gEKY_KhiEh4qROOWzU4hU2QWxNlMQSXJQ7jMW6I4Mo,1921 +torch/include/ATen/ops/special_modified_bessel_k1.h,sha256=WHGLvxJNqlUsdFumKWQWmkSjHKM3V0KcoP90_2akjLc,1481 +torch/include/ATen/ops/special_modified_bessel_k1_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Ke0JlqZASX_3fQtiRYEDC2LIuHn8t4LC-dnZR9sJnOQ,1064 +torch/include/ATen/ops/special_modified_bessel_k1_cpu_dispatch.h,sha256=1xVFdxaV_AaRuCdL53ihEYl0Nx7aHZFbUk9-PrOwJ04,1191 +torch/include/ATen/ops/special_modified_bessel_k1_cuda_dispatch.h,sha256=lL2r5k57gBhN4jDG8m2_LCfZ9l04_4EM2hhOjn6G250,1193 +torch/include/ATen/ops/special_modified_bessel_k1_meta.h,sha256=rFNnu4hAJf5Rhr1bjzHfKMS2OOa-N0Ah1E0sVOuoOMM,841 +torch/include/ATen/ops/special_modified_bessel_k1_meta_dispatch.h,sha256=bb7IKKRR0Fy2IoDalCLWeDngB5fnW1fjCcmjP5K8D8E,1193 +torch/include/ATen/ops/special_modified_bessel_k1_native.h,sha256=lOP2I55dypMpYcOmoFiTULxK6ibd7TRLqCBheYUA8dY,910 +torch/include/ATen/ops/special_modified_bessel_k1_ops.h,sha256=2aPnRmz4tco0_cmOS9bm_79t6fmsUlwh1UoVYG32tPY,1921 +torch/include/ATen/ops/special_multigammaln.h,sha256=5-oT_sOI3A7RLJ8Zs0bcXaOVhwbVkkkVEG5UWao4EjI,1484 +torch/include/ATen/ops/special_multigammaln_compositeimplicitautograd_dispatch.h,sha256=uld9uDo49UCpDyagZG4mlxK7FVpzSufj5yMyNmO07Zc,1250 +torch/include/ATen/ops/special_multigammaln_native.h,sha256=4WzSysSSJoc7iDSYo2V39XBmAqNcz2ExUkgoO4AeZHQ,858 +torch/include/ATen/ops/special_multigammaln_ops.h,sha256=wN8q-OLlqTGc3gtRngZoryH7e7syZB0MluHdDFgqJPc,1961 +torch/include/ATen/ops/special_ndtr.h,sha256=TwOvt_h2VXF6RjM3lOKXFWnd_XK_ixdcou7CZ1WlQ7g,1341 +torch/include/ATen/ops/special_ndtr_compositeimplicitautograd_dispatch.h,sha256=4AzhWFvzzCDurhkl7qT1t-Lv6FEJNzFPTkmPXhtgKuo,1193 +torch/include/ATen/ops/special_ndtr_native.h,sha256=7n6ckBQlZNBQgQevT_2HAXsOmkeeIaLX24M6meFZMiU,820 +torch/include/ATen/ops/special_ndtr_ops.h,sha256=ECbZKHayF9gUav_65OuRSsFcetu4zAX6SNDQAwoKOxs,1837 +torch/include/ATen/ops/special_ndtri.h,sha256=7eEf4wXf32gnqanUZNI0GQLtY4giTgVTTZwJ8E6YK1g,1351 +torch/include/ATen/ops/special_ndtri_compositeexplicitautogradnonfunctional_dispatch.h,sha256=wS2tqh0_aNn5bENe4jR-_Gh_eVfYztArjuNhkciBwj0,1051 +torch/include/ATen/ops/special_ndtri_cpu_dispatch.h,sha256=5fwmcmLEIIHmnZYrcuYBAw8ia8hK466mIcbwdazM5SQ,1152 +torch/include/ATen/ops/special_ndtri_cuda_dispatch.h,sha256=Hs3kGUiEWB2dr6IkKCI5CtNhFg6wxhH21iWK8RTro7o,1154 +torch/include/ATen/ops/special_ndtri_meta.h,sha256=rbewT1A25RjrO7iwl49Q83K7xEtzQllwHJZ8_SLoXq0,828 +torch/include/ATen/ops/special_ndtri_meta_dispatch.h,sha256=rxpurJZ1RQig6W7UIfLwsCmPxlK4wODsI58zatIt34k,1154 +torch/include/ATen/ops/special_ndtri_native.h,sha256=q40c0r_GWdrWUY3HK2Af8o-H7n1lU0NP6p8EAlzvosU,871 +torch/include/ATen/ops/special_ndtri_ops.h,sha256=qqwHxNYaohm8p0TZsRGrOU0rojVd6n6SBCPgXkXWPs4,1843 +torch/include/ATen/ops/special_polygamma.h,sha256=1lPSFppARN3icGVkuNu0m4789zo0DbdX39iXb8OZl88,1454 +torch/include/ATen/ops/special_polygamma_compositeimplicitautograd_dispatch.h,sha256=PeFmbLr--44xjBVcOMqDn6JdeZQf8OJpceSkVXNFcM0,1241 +torch/include/ATen/ops/special_polygamma_native.h,sha256=XPvf-qqb7m2n5MrIrCKQ4BUGjT7EAjlmYHxv_yWPElg,852 +torch/include/ATen/ops/special_polygamma_ops.h,sha256=o2NQYtpB24-ax8lxWt79o4WYHNdqIM7-sOBn43DIqXg,1943 +torch/include/ATen/ops/special_psi.h,sha256=z5maXOO47rGifdAj0CUITXd4lrWtsB8nEtBZqttXR7s,1331 +torch/include/ATen/ops/special_psi_compositeimplicitautograd_dispatch.h,sha256=UL-f5ArY0GTW192qCaf6ODiz-suZpe41wGwBD4RzFsI,1190 +torch/include/ATen/ops/special_psi_native.h,sha256=H972O_EYK50jv1DCS8bg3oAKTpZsJU9HoLsXGBMtYn8,818 +torch/include/ATen/ops/special_psi_ops.h,sha256=hEbXrboP7tMtyv2fqu9L6oscumkw8oCMLyED_pynmZk,1831 +torch/include/ATen/ops/special_round.h,sha256=LSeWpi2zt6i1d5Y_bxyxwkYdKAiimyMuXWjNef2qpTc,1490 +torch/include/ATen/ops/special_round_compositeimplicitautograd_dispatch.h,sha256=P9DIwiDJLamjrbWHwlV49mWbgCLZ0ruXtBl7AGLqweg,1254 +torch/include/ATen/ops/special_round_native.h,sha256=rdEGjONdqI0_jYbctuOdJuSV2JR7QI6leBN9jfTcpCQ,860 +torch/include/ATen/ops/special_round_ops.h,sha256=aGe4qpLO-gPxwFgBSGc363mmApE7DEaFb9zER4fcpHc,1968 +torch/include/ATen/ops/special_scaled_modified_bessel_k0.h,sha256=4ZnaANlB9QM_8w_6tq6OBqTRHWVuiLFInEWKnr0mAQo,1524 +torch/include/ATen/ops/special_scaled_modified_bessel_k0_compositeexplicitautogradnonfunctional_dispatch.h,sha256=D9kguWQ80eEXykSlnGbt7Lw0dZBU3HiTGH8sK9ehme0,1068 +torch/include/ATen/ops/special_scaled_modified_bessel_k0_cpu_dispatch.h,sha256=JgAuDaa9SCTxTJr_SGPNh9seUzRXiMbd4-CVPTMeK_4,1203 +torch/include/ATen/ops/special_scaled_modified_bessel_k0_cuda_dispatch.h,sha256=Jf9eyp8Km9DDpE00ZZGV-iYJJG9pmBgwdxiUa4uGKvU,1205 +torch/include/ATen/ops/special_scaled_modified_bessel_k0_meta.h,sha256=04t1_VYCB4AvbwCdE-HENJ3OKyg0SLKXjFZg3Qz0nsI,845 +torch/include/ATen/ops/special_scaled_modified_bessel_k0_meta_dispatch.h,sha256=1yqdbYu7kAg9p8ZO5WZVBagOIBhWlEHkRMjUwiycdpA,1205 +torch/include/ATen/ops/special_scaled_modified_bessel_k0_native.h,sha256=QVyqiyEvsVKzWCARohDRKZEJq46-0RwmIjy5LV7tfQk,928 +torch/include/ATen/ops/special_scaled_modified_bessel_k0_ops.h,sha256=Yz1gVR8-_r3jS-nvs9ysquOcdOjF51jGYQ9XJiR9CxQ,1945 +torch/include/ATen/ops/special_scaled_modified_bessel_k1.h,sha256=0DppqITfQEDPVBH3iHhRt5J7YYh38AEcJ24VHlS_KpA,1524 +torch/include/ATen/ops/special_scaled_modified_bessel_k1_compositeexplicitautogradnonfunctional_dispatch.h,sha256=0Bwg7ihJNPfAg6xuXW3HRSP1llDm-UWM3o8TlYWcsDI,1068 +torch/include/ATen/ops/special_scaled_modified_bessel_k1_cpu_dispatch.h,sha256=wRWzuki8KAGIsNDkIe6tC0Mq2sirhrau8nU3donxxDo,1203 +torch/include/ATen/ops/special_scaled_modified_bessel_k1_cuda_dispatch.h,sha256=91sqdZEwpJeaZD5_cAJ_VFMUts-SYpvk43ilGsKPleg,1205 +torch/include/ATen/ops/special_scaled_modified_bessel_k1_meta.h,sha256=cGzZ3rl2BD13c6_i_-1GV4i_vahIMk6rvzrQciR04xs,845 +torch/include/ATen/ops/special_scaled_modified_bessel_k1_meta_dispatch.h,sha256=f670AxWiy5hJAnPnB-YvfHK3XTmy7WX1Ch1sl3l3ba8,1205 +torch/include/ATen/ops/special_scaled_modified_bessel_k1_native.h,sha256=O08EcyUqzMRqvk_2mJQQ-TcFXHQs_ACsJHGqrIJbSdY,928 +torch/include/ATen/ops/special_scaled_modified_bessel_k1_ops.h,sha256=X25S02XWxL6zQg56Z7TisEppmyFQMe7P_gqR5lDioB4,1945 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_t.h,sha256=ljKCNRxW188l5nUbaFQcNmNJi7rfQDZWWgEsugQlSsk,3569 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_t_compositeexplicitautograd_dispatch.h,sha256=hT-urNAhc23HedzYvTPvNMe6FV6nGewVVCknfnxCGR4,1692 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_t_compositeexplicitautogradnonfunctional_dispatch.h,sha256=FBcQ2_MAkTQ4GxVK-Jpm2dcGCRJLATmkJyiFBU6Z3aY,1095 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_t_cpu_dispatch.h,sha256=kYNe0g_QyqQ0ZRiASlgcBbiB7Ral9aDKiD0xBYoVgmU,1284 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_t_cuda_dispatch.h,sha256=7qMHGP8PAOV7kGB3JUOBlUFFE6Ayz0AfoHPvdtmfE1M,1286 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_t_meta.h,sha256=8vLjnKRVT97Ex69AeXEGLZUGfFBRCE7hoIvyORj38RY,872 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_t_meta_dispatch.h,sha256=2nNtiOioNw79gFBXNCzqnFDQvq6ZRyrJ7LfOR3b0s90,1286 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_t_native.h,sha256=qQKQbMOjZElascXKM15L1C8fAaejVB3IV30VnG19xIA,1433 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_t_ops.h,sha256=twC2i7DUv5FC1RrGtFkOggMKOxq7JAS6ZJax6JIJZhI,5053 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_u.h,sha256=ju6vHaM7VPDOTJk6W506jeIBuRf0iktOt06XQ5IY7kM,3569 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_u_compositeexplicitautograd_dispatch.h,sha256=erSZLuGNebN7EzySC8-VYxVbHQf93WJvH_d7T5iOu_c,1692 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_u_compositeexplicitautogradnonfunctional_dispatch.h,sha256=RoB8FFeH-J3XZNvTTsglRqsDADPDLIq5GDzGuM_gzGo,1095 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_u_cpu_dispatch.h,sha256=g8_Z6krJVJQNstTNJNijF3GQdKZARo1WTguxB_CDwng,1284 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_u_cuda_dispatch.h,sha256=yWso1KUiwqNI59jAQ8v1duvDoZ03ZWzPNRvhKT-hkfg,1286 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_u_meta.h,sha256=simmZuiwwJoYRmz3h1SIORzzETDoiO-fmm_n8OmPPAw,872 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_u_meta_dispatch.h,sha256=sXj6lMy5fBtEPfCgGcHRYMqTirdhtHTfmvFyzOSLPxs,1286 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_u_native.h,sha256=v8fP-RASCpne0Szoh-3yUDZw6SXjBljSXckN76QiOtQ,1433 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_u_ops.h,sha256=6GiH4BZR6aykUCXz6aA1A4mqwx11remD88j7uwvJTMk,5053 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_v.h,sha256=T1Uc9ohsRc6kHXLMUJa9SmfPo6uZYdrVojcZTwkexjY,3569 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_v_compositeexplicitautograd_dispatch.h,sha256=XcD7eBiPzVfk9r6EBHMlXC5faZ8EohXgzn4TfCIv1l4,1692 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_v_compositeexplicitautogradnonfunctional_dispatch.h,sha256=rJChMcBy3YoQWsh7gh8hvfmHHh1ahHmfhR6eHn-md9M,1095 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_v_cpu_dispatch.h,sha256=fu155XJiSAWEEeAWNRdo5roPIOQwdB-b63FJY_3Xo9Y,1284 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_v_cuda_dispatch.h,sha256=gj5Z3y1VlGxaVDd0PDpptN-DVKrSQKkm7hWtueactmE,1286 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_v_meta.h,sha256=rP2jkzYcfocz1SEoXn8Pt2lkM1-Th0q7k7Nzl6cJdaE,872 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_v_meta_dispatch.h,sha256=aQQX7nojyf0ZmsOJDUxf6dFum41iv4KXWqMS4gpH0zE,1286 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_v_native.h,sha256=O4S0TrTPXCMPsIxkDiE34l53YoHDtT2vwIbCdqhFpZM,1433 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_v_ops.h,sha256=zGGnAYj3zMz-vhTzespj_kHofk5fbE9NxSoWBaO38Do,5053 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_w.h,sha256=ff062I2wZpyX52CxO-pkYpTzw2Peie82O5KqpBVJS8c,3569 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_w_compositeexplicitautograd_dispatch.h,sha256=DApDCH3UpaMJsyQuwryYyy1NDa221YHOdUOgdMRxaJY,1692 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_w_compositeexplicitautogradnonfunctional_dispatch.h,sha256=7n2brL__o-6WS5_SensRHYt_TxGeConmwwjgKg4zIZE,1095 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_w_cpu_dispatch.h,sha256=7cc7E6xiCRB3cVxZ3hyX01fIX81BiVsoHax4Vxx4rmE,1284 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_w_cuda_dispatch.h,sha256=xuIeusRTrGasn34S6vTrGP4dbq9QYrEw-mX6GO9_XQc,1286 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_w_meta.h,sha256=Rxu3uSszXcJKusEPc3vxX4Dd9-AVi0x1LrZud9GGOws,872 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_w_meta_dispatch.h,sha256=y6AmcVGU1cxrOz0iUh7L7oNW8cBfUo4HEgEOW04j03M,1286 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_w_native.h,sha256=uGZMGpKB_brZE0CK2YUriaCRS_uYT8O-92nAmCYqSe0,1433 +torch/include/ATen/ops/special_shifted_chebyshev_polynomial_w_ops.h,sha256=RTMWrVt-MOKVVt-9QeUlLNNgYX_abQx7w0sW7WMjEFw,5053 +torch/include/ATen/ops/special_sinc.h,sha256=dOxn27pAuQVLCHW1SJRhk1HJ6hBWs_5WT_zZQP9gNFM,1341 +torch/include/ATen/ops/special_sinc_compositeimplicitautograd_dispatch.h,sha256=75w5kc3uR8KfGxppEBOqgT7K9PEyU9X6AMsHARzgKC4,1193 +torch/include/ATen/ops/special_sinc_native.h,sha256=3mvrXOhGYB6Ar3lV_sWbxPXt-hQM8WQdN00KiW4Xuss,820 +torch/include/ATen/ops/special_sinc_ops.h,sha256=cSC2g_mGiFif25xqeaM1yRs02TWoJTgoLEqWE1V4NLM,1837 +torch/include/ATen/ops/special_softmax.h,sha256=6Qy2obWyDMKmaSWh9N5syX_PN2kgh_-3Kf2nH9NMFM4,1039 +torch/include/ATen/ops/special_softmax_compositeimplicitautograd_dispatch.h,sha256=0mrweuUJko9utkNIE8OzM6L1sDE5KKeY2H4FBQJ7ejM,1094 +torch/include/ATen/ops/special_softmax_native.h,sha256=BxuGqwf1hSQTFLkJ7YBAMYt_gS7F6g0z_ahXCK4lCvc,806 +torch/include/ATen/ops/special_softmax_ops.h,sha256=V98P9CgFRDrFSnixv9hjSuKa7-ivJKvr3YhM94d909o,1417 +torch/include/ATen/ops/special_spherical_bessel_j0.h,sha256=scTpjGJChq4LQH4ubXAOFMBnT6iKzuRX_btYCikyiEE,1464 +torch/include/ATen/ops/special_spherical_bessel_j0_compositeexplicitautogradnonfunctional_dispatch.h,sha256=n1mUx9qXkEFPLwz-pY-ovV0LgJw5JiIwEnejoVhbfGk,1062 +torch/include/ATen/ops/special_spherical_bessel_j0_cpu_dispatch.h,sha256=NODpgKK53LjDW3cDcvjYwlqxPdVBJbBP12dJUlFe1o0,1185 +torch/include/ATen/ops/special_spherical_bessel_j0_cuda_dispatch.h,sha256=3lOLvfVHp9hEGbFYVebfSqXN2RFMBye-Uj3FJHW9Kbw,1187 +torch/include/ATen/ops/special_spherical_bessel_j0_meta.h,sha256=3EyGu4mbyQVBNQa-PkjrxqKyRcbvgSaKztpCwpZScW4,839 +torch/include/ATen/ops/special_spherical_bessel_j0_meta_dispatch.h,sha256=s1x2F0hDZ9NmSi5Kce4hqw4F3OeOew2Zpsyj21zEZ30,1187 +torch/include/ATen/ops/special_spherical_bessel_j0_native.h,sha256=az0P7BJ_srVRHHjfSK5xvbH_7JfHshU87Zlwxw3J6_w,910 +torch/include/ATen/ops/special_spherical_bessel_j0_ops.h,sha256=u4_SBV8IIPTBY2OuJL5pHsH_AueDbN0sHwkCh12Kvdo,1909 +torch/include/ATen/ops/special_xlog1py.h,sha256=QXHUqpho2S9qj5VClFkgKWaliN2_xHSRW3ouCfNfH84,3156 +torch/include/ATen/ops/special_xlog1py_compositeexplicitautograd_dispatch.h,sha256=9kXNHPGFW6QPtOyeqO2WO-SM-d2navaInKr6y-_90iU,1596 +torch/include/ATen/ops/special_xlog1py_compositeexplicitautogradnonfunctional_dispatch.h,sha256=fcgMuHBAkbu1AEWo4zhPnj1zilnRoRfn7T81OdwRZRg,1079 +torch/include/ATen/ops/special_xlog1py_cpu_dispatch.h,sha256=xKxT92Kg9oTC9mp6gEj2Ia5Rx-OYjA4lWtBOj_E9T4U,1236 +torch/include/ATen/ops/special_xlog1py_cuda_dispatch.h,sha256=J4ivdy1aINjOTHKjoY8GNFcJmV5UpwKe53Ujzktu5QM,1238 +torch/include/ATen/ops/special_xlog1py_meta.h,sha256=KSW_Fz5ntUbbWNZfRtJk24aukTiAX3DkFxFNWVCdbXM,856 +torch/include/ATen/ops/special_xlog1py_meta_dispatch.h,sha256=DLf72TaQKCe069MpA-RMkLLnk2JfXOzRWIoWwACxjTQ,1238 +torch/include/ATen/ops/special_xlog1py_native.h,sha256=7I3bFhzwzIXngHFpLSmgJKTLxLyG1PQWuYom3o_Zevo,1307 +torch/include/ATen/ops/special_xlog1py_ops.h,sha256=b6nFz4fQJXlEhK7VHnQGBlTE0_mP4dMekTbXDHba-rw,4807 +torch/include/ATen/ops/special_xlogy.h,sha256=_NkYG7wNcU94VglkRRZylm38keDjpPH_P1IIU96gEJc,3100 +torch/include/ATen/ops/special_xlogy_compositeimplicitautograd_dispatch.h,sha256=3jH_FojbOq6DtlE5lIpLea3dcUxRT2aoo01NkQoKX6U,1894 +torch/include/ATen/ops/special_xlogy_native.h,sha256=zDnl0dVk9R-hOZxbng4CLymi3nKvJ2lcsMZ5riWsWe8,1270 +torch/include/ATen/ops/special_xlogy_ops.h,sha256=70D6zPiUYT6YdDVVDXMDLHE2Jpdj-4J8LeOiG007q30,4771 +torch/include/ATen/ops/special_zeta.h,sha256=GAN2VCvs76HQN1A5Actk366tjvh8G2GzVzf5Kqr49EA,3072 +torch/include/ATen/ops/special_zeta_compositeexplicitautograd_dispatch.h,sha256=fcjmQTDPjOOUQ-I59FSzCI9GtY4o39LMoGTsnED9wvI,1578 +torch/include/ATen/ops/special_zeta_compositeexplicitautogradnonfunctional_dispatch.h,sha256=raEWaxWLaFFex2wJE1NwZU3oV-YFEtydxr-FV5TrTqU,1076 +torch/include/ATen/ops/special_zeta_cpu_dispatch.h,sha256=AcdCL3lAfGsTXZeoXTBfsbinHFQ2U8cwqGO-BVGwvsk,1227 +torch/include/ATen/ops/special_zeta_cuda_dispatch.h,sha256=GfP7--qjelRQzl3Dny9FcOygD0uhEAm8U32ogSdpWjs,1229 +torch/include/ATen/ops/special_zeta_meta.h,sha256=UTxC2DpPoIdoxxsSugIxJwxRBN3dviPN2tHhDD8V0OY,853 +torch/include/ATen/ops/special_zeta_meta_dispatch.h,sha256=LDzTFXcEXSpfAZrMQ7fDnnWuTU4xRRi8P-AueDculEI,1229 +torch/include/ATen/ops/special_zeta_native.h,sha256=mvlzamNCD1AsMquw0L4-UrnLl0IhEhV8jU9dAVptnO8,1286 +torch/include/ATen/ops/special_zeta_ops.h,sha256=ghtOSy01hilA7BoAeAfLy7NK26oxMm6sFmCXEvvCgSY,4753 +torch/include/ATen/ops/split.h,sha256=h5J5YN-6o8LcLUmncFHcDpfiuzyAngeyOIlIClI1g7w,2963 +torch/include/ATen/ops/split_compositeexplicitautograd_dispatch.h,sha256=MncTnMHOMD8wa3XnKCNSeSuiYF4lItTm5aSIgvjpE-E,1181 +torch/include/ATen/ops/split_compositeimplicitautograd_dispatch.h,sha256=xP5zeap10kz2sSHFLfj70bI_aKbhqtGfrtamMcZDYxo,1197 +torch/include/ATen/ops/split_copy.h,sha256=T1d_ziTAqDnfBvHukgeWH8NIBcIgJlItMr0GHwJ0e_M,4201 +torch/include/ATen/ops/split_copy_compositeexplicitautograd_dispatch.h,sha256=gLkjnJtvIgSJajFHwzRzL64lSxKNnx16hLdgdq6XSIA,1428 +torch/include/ATen/ops/split_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=LEb8v_K8qSzv7oDTPcjWpMxFkiffvea0pTH-_4Xa9oE,1217 +torch/include/ATen/ops/split_copy_native.h,sha256=JgOOj-cNRBT41m5tqLnV_BC-8JOnUH0z_RRTKMAd7KQ,918 +torch/include/ATen/ops/split_copy_ops.h,sha256=y0BIomWBnL_ZLMz5F8lzCFiLdGwvflPF-vM87EY_s1I,2141 +torch/include/ATen/ops/split_native.h,sha256=M-8204CbYMMty0O5Xl9Aq0gUXtiFBNtmEZvX4qdP1UE,901 +torch/include/ATen/ops/split_ops.h,sha256=S7hFWxOs5imD6jB2m2X5Kgw9h71BGsmmJO8arRX6Ezg,2136 +torch/include/ATen/ops/split_with_sizes.h,sha256=1rDs3gjWNK924zsgpO5RpQUcb0cbwjy7Y40Q_wZbw9k,1995 +torch/include/ATen/ops/split_with_sizes_compositeexplicitautograd_dispatch.h,sha256=_Q3dRajwCooY93QH0UAGeqO9ceut9_-QzK2KtZ8AtQU,1221 +torch/include/ATen/ops/split_with_sizes_copy.h,sha256=K_pUUDGP_7rZgK-DeYCW-w7wNxAA8uGgkDrfxFrXz_A,4710 +torch/include/ATen/ops/split_with_sizes_copy_compositeexplicitautograd_dispatch.h,sha256=Ad1oZ19zNZrfioXk1-v6_PeqLy_Mtf-nnpvIQBJQh8o,1508 +torch/include/ATen/ops/split_with_sizes_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=L9uJk97OK3OvRibglcpQ7i5yn3rdLQTsmJsbOfeiqyk,1257 +torch/include/ATen/ops/split_with_sizes_copy_cuda_dispatch.h,sha256=QZROOLbq2U5Ek4ceGkKlHlEfSMPfLhkl6qcy7p6zX-8,1466 +torch/include/ATen/ops/split_with_sizes_copy_native.h,sha256=M8ttUaXpkR2CtuTKAIn1VTmkPqk-odK5QDAnpVA7ao4,1078 +torch/include/ATen/ops/split_with_sizes_copy_ops.h,sha256=NJe6o-MdJX1_iqlSRKmvyhRFK6OE0H5Lsg2PfgQBh9g,2224 +torch/include/ATen/ops/split_with_sizes_native.h,sha256=xqY1T_Z5nhj29HtYCj1bELvjmKP0M-O3QxyT4EfiGeU,929 +torch/include/ATen/ops/split_with_sizes_ops.h,sha256=LRoERWxGmxLOZonierXl-g2P-1wsULymSvfy3Ymi_TI,1454 +torch/include/ATen/ops/sqrt.h,sha256=XHOet260Go_8JBKbHvsAmSvVO8VYY0R-U7f-pQlfb7U,1397 +torch/include/ATen/ops/sqrt_compositeexplicitautogradnonfunctional_dispatch.h,sha256=uQGaFwBmDDOgAO2SYg66pCnD_GXHt4ggTYh5ywNphR0,1091 +torch/include/ATen/ops/sqrt_cpu_dispatch.h,sha256=Fjy3pPCURmAFmAkWN3cm__ae41FJk5S1Yc1ToCqTkio,1174 +torch/include/ATen/ops/sqrt_cuda_dispatch.h,sha256=ypIukgmtYEFqQPBZ3Rr7XYp9wR3UNdtUaMdok92W7vQ,1176 +torch/include/ATen/ops/sqrt_meta.h,sha256=903aTi2X5JhJN9H2DbHxHH3rq7iUoR8W9qWSmTrfcNo,819 +torch/include/ATen/ops/sqrt_meta_dispatch.h,sha256=qGLw8iJGQ7_9a2Rxw37al70Kh0fSFlINCVYQGMTKtiY,1176 +torch/include/ATen/ops/sqrt_native.h,sha256=_fnb5p5W8TeN3EHSMDeL3BvUlbvYaApZSbbPxaxbcbY,1317 +torch/include/ATen/ops/sqrt_ops.h,sha256=kpPjUDtHcGFVCngSvVx1TJ3LkIf4tJkopYpcCyQ-aPI,2282 +torch/include/ATen/ops/square.h,sha256=e4KvE34EgrDnOnTd6m9EfNbsN9EKtBEMV_-Z0AwFwNM,1423 +torch/include/ATen/ops/square_compositeimplicitautograd_dispatch.h,sha256=xo0-A41iPUchBZ6E1Jk3qQuzVypJs2pQmMZHup0_Xzs,1226 +torch/include/ATen/ops/square_native.h,sha256=lKsznyzKFd9EbQjzR8Ge2RkYTPK7YqttZ1Psp9uPq5I,859 +torch/include/ATen/ops/square_ops.h,sha256=dJQh2m70S1BraIARQ4IcsDQHUh1tcJehact_DyEoWEA,2300 +torch/include/ATen/ops/squeeze.h,sha256=g3gXbhI_caJqHdx3Au6NHOKNySusUnawN104FuUJ5Ls,1466 +torch/include/ATen/ops/squeeze_compositeexplicitautograd_dispatch.h,sha256=ZSb1yOyzqe5O_D35hGbhCsDNYuUUqn0DXRMSAkiI2C4,1353 +torch/include/ATen/ops/squeeze_compositeimplicitautograd_dispatch.h,sha256=JVHvTW-l-udf6fAuRXEMJTTFnExQg1fC_iTpfDG6Xbs,1105 +torch/include/ATen/ops/squeeze_copy.h,sha256=IKc8Z-e94OyS-YqJ-q4u0MakENZ2Kb58c3Jr2hv4_ZY,2745 +torch/include/ATen/ops/squeeze_copy_compositeexplicitautograd_dispatch.h,sha256=7nDV0C-X_qZkEhNEWYl2sAjoIxqrSGnbmpdRWm6omeY,1539 +torch/include/ATen/ops/squeeze_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ZsKfXMl5q6az4lo7xuMbL8J8XTZhNehfHjyM3AiTjNA,1204 +torch/include/ATen/ops/squeeze_copy_native.h,sha256=ICFWFY5LWrqAgt4taqvLKW91IxCekEhhFVucCwDlb2o,1194 +torch/include/ATen/ops/squeeze_copy_ops.h,sha256=0hEkbhzOho930mXLaR3ahNFgiSE4bbcyZWpomRDeSd8,4369 +torch/include/ATen/ops/squeeze_native.h,sha256=zH47J6V4oltQD24ko5toAMF3Gv2LjTXxpmS6eKN8Z14,1663 +torch/include/ATen/ops/squeeze_ops.h,sha256=sM6XFSX9RmWSc_-FUr4yxRGxqS9mA5t-opf6al_hXCI,5199 +torch/include/ATen/ops/sspaddmm.h,sha256=GCf00WUgI5klLqG9v4biTgi7xQgcRMmMjnh5XNtRom0,1861 +torch/include/ATen/ops/sspaddmm_compositeimplicitautograd_dispatch.h,sha256=Gqt36TJVgrsynX4Bu5jhgrpCoIcXDfD7y5BJw6UCAqA,1125 +torch/include/ATen/ops/sspaddmm_cpu_dispatch.h,sha256=V718TF-MF5Z2R6ABWt6L91UDLoyoQ26geGbvu6cmg8s,1287 +torch/include/ATen/ops/sspaddmm_cuda_dispatch.h,sha256=WSt5AH5f5s0zcN-pvdG-d5-F4NkGB0VOWgz-t3VCazU,1289 +torch/include/ATen/ops/sspaddmm_native.h,sha256=fX2tMqiLDgc0JecPQVJqqvhaPFN2h_mtWYSnpi1pGEc,1603 +torch/include/ATen/ops/sspaddmm_ops.h,sha256=qEu8LzxebW-oNipAhSdLtDtLJejE45P_eEkVejICDyU,2494 +torch/include/ATen/ops/stack.h,sha256=aMZdsxNT_JVDmSOHo8hmkcJGOnpOj091zDnkf_X4Cfs,1383 +torch/include/ATen/ops/stack_compositeexplicitautograd_dispatch.h,sha256=ZvwP0AOFPMgNQaNxuwQZXYpa6xPEq6HNJc8YD81w4_g,1212 +torch/include/ATen/ops/stack_native.h,sha256=6A10LdJLtHKKarlThokYrAWVHNPK7o6hJwgpGHNmTKw,832 +torch/include/ATen/ops/stack_ops.h,sha256=gn4mp0uSj5_0qhMZoPXDCVhvU_KapoUfdT9OLY25og8,1885 +torch/include/ATen/ops/std.h,sha256=Sn5EBPoguw5p7unsLVcSd-2D1bjvxp8N-_IA6G9QKDk,5125 +torch/include/ATen/ops/std_compositeimplicitautograd_dispatch.h,sha256=vh-g1_c14sj0iLWYrqSbLBoow8M5uyorqz5le_UdyiU,2275 +torch/include/ATen/ops/std_cpu_dispatch.h,sha256=rOfugf4oSD9EcGLjyNQJ8RVW-HKh8WyyBQAKbr4EL8M,1467 +torch/include/ATen/ops/std_cuda_dispatch.h,sha256=yrOm16B3Kzeg_zSI60BHl37n-dDsWcpnuo6JuyR_6fE,1469 +torch/include/ATen/ops/std_mean.h,sha256=OLB9DRQvFCmebTSoTVdF_ebcwZYNfh8NIaObTzhPAmE,3471 +torch/include/ATen/ops/std_mean_compositeexplicitautograd_dispatch.h,sha256=MGcesqjaQ_c87syWBMhvYTE26nxIV-8fOEQKS8iuN44,1437 +torch/include/ATen/ops/std_mean_compositeimplicitautograd_dispatch.h,sha256=BblMcekBck5mtdvUcXZPNT3J4W934QxfS05MZwmAIRI,1527 +torch/include/ATen/ops/std_mean_cpu_dispatch.h,sha256=_rbTV-XV9by8EEazPrj_t_0P-1sIIrtOVpJ7PUxQfEU,1128 +torch/include/ATen/ops/std_mean_cuda_dispatch.h,sha256=dGs3PuoL2P6U_B0oQiOs5KUOr7KjYk8c5HSOSSMxA2Y,1130 +torch/include/ATen/ops/std_mean_native.h,sha256=iSULCV6JEz-Bt0jGz8RA2_NLIc8qyRD0lHnyy88B-Nc,1691 +torch/include/ATen/ops/std_mean_ops.h,sha256=oN5hV1aMzT93kGWGdks0y7a69VdvWl0UJm6kh43wg9A,5950 +torch/include/ATen/ops/std_native.h,sha256=W4VuPglZYrMxCUVgCSiEmnvaXIfAWhEC0ndrSrTlk1k,2266 +torch/include/ATen/ops/std_ops.h,sha256=2mjyC0TWbl2mfA3In0vJx1hc8vTe2m7DUQvbrs7w87A,7758 +torch/include/ATen/ops/stft.h,sha256=pUAXR5Z9uP9mfVWWuM2zXVYu4kX66kjOWW0RIZ6fEKo,2326 +torch/include/ATen/ops/stft_compositeimplicitautograd_dispatch.h,sha256=fq-SB2m1ZGwvor63KYlNYaUpOmdt-ka16-0riYisvcs,1770 +torch/include/ATen/ops/stft_native.h,sha256=xPgk-E3809JgwkfSLVYg5CTZLxtae7BTyCnBpGTb4SY,1521 +torch/include/ATen/ops/stft_ops.h,sha256=AO0RvgR5c_Uj9E2qv6Bgb9RX09AM7oa4TkkyK8n-Sf0,3599 +torch/include/ATen/ops/storage_offset.h,sha256=LwVkr21lne9Tvhv4VpkhPBmJ4R3EqqsnHh2RucP0fbo,765 +torch/include/ATen/ops/storage_offset_compositeimplicitautograd_dispatch.h,sha256=6QBhFnrSRgeTyixmlNKMUmy3B7iJE8U4StaV9jQ2DqA,1023 +torch/include/ATen/ops/storage_offset_native.h,sha256=SeXnb_KzBlSpPXPOj_Mc-QKCYTwlNH9nyTOFpWRyOyw,735 +torch/include/ATen/ops/storage_offset_ops.h,sha256=rnrw3CplU-F-BquJnOSfY1Ohs5mhRlkllz7SBE-3q5w,1223 +torch/include/ATen/ops/stride.h,sha256=sInTcrRrdIimjiCjUtzDIhlTCSfzu6_sXBED-u_VbVo,1111 +torch/include/ATen/ops/stride_compositeimplicitautograd_dispatch.h,sha256=ZFVhKk4rHIj2XHrd_U1d0zhGPEcjSK_HX8fInlhRSAQ,1096 +torch/include/ATen/ops/stride_native.h,sha256=Q7fwpCln3h4T-GapAHqDP2i8tVvpA9nei7x265bY1Tc,808 +torch/include/ATen/ops/stride_ops.h,sha256=jrkWUwA5907rbqMhhVbutYeWfBvBqSK_7KUJMLXyFlU,1825 +torch/include/ATen/ops/sub.h,sha256=MYqpBoJLT6QWdCXNB-soiiYtjcoMpuMjuYDht_fuFQ8,2392 +torch/include/ATen/ops/sub_compositeexplicitautograd_dispatch.h,sha256=f8GQBnN6cbZTLfhmkvqs78JiUfgN5TUnFiMelG9M25A,1428 +torch/include/ATen/ops/sub_compositeexplicitautogradnonfunctional_dispatch.h,sha256=tSjvcBgF5VHlZo6mbEdM-tKUum6YSLhG3wJg-fJ-Uag,1197 +torch/include/ATen/ops/sub_cpu_dispatch.h,sha256=wlywFWikE1X9FAEVaocFH4deo6I857FfFOoKSmk0AKc,1384 +torch/include/ATen/ops/sub_cuda_dispatch.h,sha256=3_CQhhSPJB8Xtxl32zVdmPLWNPcnaSbaO28m4BBxl58,1386 +torch/include/ATen/ops/sub_meta.h,sha256=QWUlFb4sqT61UPmwMyuIfVQbrfDC2OnLLZ_mWuBtv1g,877 +torch/include/ATen/ops/sub_meta_dispatch.h,sha256=gEke6lfMhGgzXCmUirOAVAYw3nuBOEsdASP0_aAVzzs,1386 +torch/include/ATen/ops/sub_native.h,sha256=g-aT-d_XtN9_3eygYLCyvCTF3ZFYk936RqNlbbE_Hrs,1837 +torch/include/ATen/ops/sub_ops.h,sha256=TZPEReaPghV7_6Zk-V3pMCRR_TR5sgXVhSH5KhA0_cs,4986 +torch/include/ATen/ops/subtract.h,sha256=Q45eKrr1fKZAxjez13K74pSMaQVzco5Y4PDww9jCPHw,1863 +torch/include/ATen/ops/subtract_compositeimplicitautograd_dispatch.h,sha256=Fln9p4a7kdCr8BIq2Jv_ejV2U_kemdIMgEv3yHvZmrI,1665 +torch/include/ATen/ops/subtract_native.h,sha256=SXUK8e_f5VMfHFm0DwiTfUN0ViNiuIt2GsdAnIhN7OE,1242 +torch/include/ATen/ops/subtract_ops.h,sha256=MLfovAPIBCPyMlCUxXO7O7i-ehB05M7IBdOpezPRUuc,4285 +torch/include/ATen/ops/sum.h,sha256=2HWkcxc0bZOHg4P8SMRm4pXqCNUNio8k-G_85pFVtzY,3661 +torch/include/ATen/ops/sum_compositeexplicitautograd_dispatch.h,sha256=Zqeei8aGYMkjMdGYlzIj2a_8LLXe3Y4TNiOpBvNGCbk,1313 +torch/include/ATen/ops/sum_compositeexplicitautogradnonfunctional_dispatch.h,sha256=6OMDNRMBC_1bdtDFzODgwkObyGRUC4m88xaMPRHjZXA,1144 +torch/include/ATen/ops/sum_compositeimplicitautograd_dispatch.h,sha256=_cHp8x4eMmzDEBrfK9XMDLQnuKMHH-miW4BJFmUqVrY,1430 +torch/include/ATen/ops/sum_cpu_dispatch.h,sha256=tnvEV4VVVjQvFS2Y8B7u17g2G8zSLbUhe3dFvYJ1quY,1410 +torch/include/ATen/ops/sum_cuda_dispatch.h,sha256=HTtoGmZhrv1YuW2EFhTViJYX3AFtIpLZoXjlI2yN9Mk,1412 +torch/include/ATen/ops/sum_meta.h,sha256=Wiq_jbANgPZv7XqhNrCZPfcAbW8p6ae3i2bxz-B0mY8,912 +torch/include/ATen/ops/sum_meta_dispatch.h,sha256=kuoLlkGprcuL23_1S7q32G6SwsdNC5YoK0tq8LQkiIw,1412 +torch/include/ATen/ops/sum_native.h,sha256=pm8BdxAUUBy85R6c3IhkG94YgCL3vIRWPjoM0DKhmUA,2179 +torch/include/ATen/ops/sum_ops.h,sha256=hoBVoK1v5BP2ej5EttPnTHRij-XWyERO8dNnD7Bfkfg,5446 +torch/include/ATen/ops/sum_to_size.h,sha256=LodmqMNMI_iyxirxS661UHUDbgcwUgNd5ee4stDg9BQ,1261 +torch/include/ATen/ops/sum_to_size_compositeimplicitautograd_dispatch.h,sha256=jr0I3ZTISd2BSrRxpisHDYrSw-TfpyeNt-AqT8JyI7k,1137 +torch/include/ATen/ops/sum_to_size_native.h,sha256=fMQvE6pcoo60o4GG2mK-pVKHUAbpdIbj7FL1cEXy-6E,768 +torch/include/ATen/ops/sum_to_size_ops.h,sha256=MpgsrwhZGkGR0t82hvmLWx2OyShelBCxwmX0d7mKdFA,1314 +torch/include/ATen/ops/svd.h,sha256=ZmGkzEtE2hejqYdbyzFBwfXWGcfSVDiZ4RPDHkC9Z2U,1839 +torch/include/ATen/ops/svd_compositeimplicitautograd_dispatch.h,sha256=86lNN8QpkMUKK4ywF4lTbcilYohIoU3zjJY3PO1MDtw,1446 +torch/include/ATen/ops/svd_native.h,sha256=29sB2pVpERq3myRjUcxcQsQGdr9zjIBFq6CfPIHLJ5c,974 +torch/include/ATen/ops/svd_ops.h,sha256=ShDjnRQF7WjasPTQdUZ2kxmkV2E1tgJ-wXKkMv3qkow,2387 +torch/include/ATen/ops/swapaxes.h,sha256=29t8aYD9AadJpL94IIqTAbeLCtMiD0CpdM9Qz-CYreM,971 +torch/include/ATen/ops/swapaxes_compositeimplicitautograd_dispatch.h,sha256=hB1PurvK8WFYBgDQGSAJYFq58jIPRo-stsWi4VcMQnY,1133 +torch/include/ATen/ops/swapaxes_native.h,sha256=F80B2j2UN0lh7Wz49aqG8XrQRjZ7O-7FTZ7IJq1k9ds,845 +torch/include/ATen/ops/swapaxes_ops.h,sha256=0WvJ4_MIi3tSxw7FqV6XHIiAy5aAZf3TTO82hXy0V20,1928 +torch/include/ATen/ops/swapdims.h,sha256=1kYVp-XKCe_sdPfkoW5mBgh-I3WuzRytJlK6ctolnlw,965 +torch/include/ATen/ops/swapdims_compositeimplicitautograd_dispatch.h,sha256=V997x7yfEy5GAgLmlEWpmCy2GZr8C3qxPgpa3pldPq0,1129 +torch/include/ATen/ops/swapdims_native.h,sha256=5oK3MGVh6NzNMr8OyLolHAhKWV-wAJbXrAzDjvpGfHc,841 +torch/include/ATen/ops/swapdims_ops.h,sha256=nOzW_mchJnSlWG0zziJb0xqgGqHTBj0Tzi6qJVFUw4s,1916 +torch/include/ATen/ops/sym_constrain_range.h,sha256=920VqtO_BcyYZBTg9lFKZDlBe-RvK5PDMhL0Ba6bieo,1066 +torch/include/ATen/ops/sym_constrain_range_compositeexplicitautograd_dispatch.h,sha256=TvCSeIOQyBnbOIBTssN9TaRzOQU-It_r0SfUM3-Wt3o,1115 +torch/include/ATen/ops/sym_constrain_range_for_size.h,sha256=22zDubFHrdSL9v8EH0CXfhpLhYrxmx3GCDZmbhG0STY,1102 +torch/include/ATen/ops/sym_constrain_range_for_size_compositeexplicitautograd_dispatch.h,sha256=hcZ-2ee1dcMBOJYP6p4nAmysKvGxbL_1ympEzo0SEE4,1124 +torch/include/ATen/ops/sym_constrain_range_for_size_native.h,sha256=V9Pu25AdeESmbhN5X_Sk61QlmEsPHveq86Trl6UX-8o,836 +torch/include/ATen/ops/sym_constrain_range_for_size_ops.h,sha256=6Vp1eZJOedqk7QXKg1-UH1QB59M-HHSQyA6NOX6zCJ8,1460 +torch/include/ATen/ops/sym_constrain_range_native.h,sha256=nhPlVm12i5YGm5Tn8y__6TqwvMGDfgl3snli34V7SCw,827 +torch/include/ATen/ops/sym_constrain_range_ops.h,sha256=xx-2cHXUgNvFb1UX_xACI572JLqDO3hWMazv9wpsePU,1433 +torch/include/ATen/ops/sym_is_contiguous.h,sha256=545hxC-VtPox5g4DivnSCriq5LqfQwK1i6NuoVL4Urg,1072 +torch/include/ATen/ops/sym_is_contiguous_compositeimplicitautograd_dispatch.h,sha256=EtqpN5fJ44t_34qhmGV-C0Q1JeA_DBIeazqCKsRW_OQ,1093 +torch/include/ATen/ops/sym_is_contiguous_native.h,sha256=3RShIYGXpoNhpAQl1Zoe28DfJQ8zcHqpIaIwjQUDlsA,805 +torch/include/ATen/ops/sym_is_contiguous_ops.h,sha256=71V5VljVovxsHPIJ8TLXgT9fG_IZm2waAuajYfdQUkU,1379 +torch/include/ATen/ops/sym_numel.h,sha256=U7TIu_SZShOaHN_U0CHLIjLNPjab3dbBLluJhIZ-AMw,915 +torch/include/ATen/ops/sym_numel_compositeimplicitautograd_dispatch.h,sha256=WJtd4KNS4NlRaE_Mf0OKqoEiDwBsSk-_M1wdDvp6MDE,1022 +torch/include/ATen/ops/sym_numel_native.h,sha256=uWIG12AsDJhTVvdtD_dBf3312hpWVAE5WKSH7Fd0wCw,734 +torch/include/ATen/ops/sym_numel_ops.h,sha256=Amr4Dm9Vyd1aj_hA43G3oEJPB_S9v9X4KdpQHhnibq8,1223 +torch/include/ATen/ops/sym_size.h,sha256=YMBukl2X2EejU3pBn0vq8SBP_RZJf9xaHZxFj4BH2sE,946 +torch/include/ATen/ops/sym_size_compositeimplicitautograd_dispatch.h,sha256=vbLHrJfuxmDBvetS4XZ0PAJveyv5dKl_C5Qz-5elcVY,1034 +torch/include/ATen/ops/sym_size_native.h,sha256=edxQty447HlHRaNToij0mdamySiHpHOIRaxKEwPSorU,746 +torch/include/ATen/ops/sym_size_ops.h,sha256=EdhRVhPryJC_WzSDkzQ0Ey2FIuCmKst9EwFKDbnlDXY,1275 +torch/include/ATen/ops/sym_storage_offset.h,sha256=ErnY_L9r4RDI7JVm-vbNfh24IcPr9KRAFV3Q-DSXc9g,951 +torch/include/ATen/ops/sym_storage_offset_compositeimplicitautograd_dispatch.h,sha256=E6cKLs_DkINY2sQ7QOQEZLDsH3_g2zykHyIIR-GWTHk,1031 +torch/include/ATen/ops/sym_storage_offset_native.h,sha256=3cr2k0TMkIsXHRH0bUpGD5F6JTrCQjg0iEysW4j3MqQ,743 +torch/include/ATen/ops/sym_storage_offset_ops.h,sha256=rnQF8QUvu3D6iHbsD7UvCvLFQsU8ZlAtnVi_UHCEdGs,1250 +torch/include/ATen/ops/sym_stride.h,sha256=tuCUaUTJ3sIRmH1EHh5cOkp9mKlO_Y35JuVp4gQ38SU,954 +torch/include/ATen/ops/sym_stride_compositeimplicitautograd_dispatch.h,sha256=jYCNOtxkAX8QMBEWhaNhJ2L84392g_3-DDLuyv6v0Ik,1036 +torch/include/ATen/ops/sym_stride_native.h,sha256=oiBaRgEY-s-uEXhvmQCGeyjYTzAaXLuVwsn9jNHE7Y0,748 +torch/include/ATen/ops/sym_stride_ops.h,sha256=niaNctuWgckNaqbYcNKZepDU-dV_i00ZibxxBjPOjZY,1281 +torch/include/ATen/ops/t.h,sha256=u9g5X1-el2ibTTDSTEEWH8a5U7aoLdD9Cazdok_vtFc,877 +torch/include/ATen/ops/t_compositeexplicitautograd_dispatch.h,sha256=HGXpPcJ8yoJIbS8RXC669bcIkQDeKD80eT4irFyJQHI,1059 +torch/include/ATen/ops/t_copy.h,sha256=7JteRF2KvB6-IYXJv9cce4ThAbdZXuhV5VqDOWVlrA4,1281 +torch/include/ATen/ops/t_copy_compositeexplicitautograd_dispatch.h,sha256=m-nVK9W2tzdV9yjvZ7VIISNtTZpcDbFSL-nLz3uyRUA,1121 +torch/include/ATen/ops/t_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=SkPqxSVglOra5Q-HrMBxMRkx8HTe6le1rc_YRG2NAjk,1044 +torch/include/ATen/ops/t_copy_native.h,sha256=qgQDpbr__9ImbYmJTadCNgdN7_S2TtqRZHRdKkuXlD0,808 +torch/include/ATen/ops/t_copy_ops.h,sha256=iHF3QUMZ9_okqOLGirzFBEHOpXYbaDy473l5lkblOgc,1801 +torch/include/ATen/ops/t_native.h,sha256=Npm4nLUEvm-88NgmL1BaTVZ95i5SQDvukPAhQbtZlRY,771 +torch/include/ATen/ops/t_ops.h,sha256=XXWy-LxZUlJF70KheMCktUoIJrT9TqhHRph-9hEFvOM,1686 +torch/include/ATen/ops/take.h,sha256=h9GX0DT2urGGw2U86S0Mw556saj3tlVAkSsTtP4Won0,1402 +torch/include/ATen/ops/take_along_dim.h,sha256=XOHQnR2X8oeOETCx_bM6flt64BixVLjhzbX72xhiJaI,1700 +torch/include/ATen/ops/take_along_dim_compositeimplicitautograd_dispatch.h,sha256=Gfq5Pm4N_K6TazsL81p6LGq-MT14H_Ywbl4ewL_OV3M,1403 +torch/include/ATen/ops/take_along_dim_native.h,sha256=_uspp1VN-9qn0Ijv4x_Het0ndOdCONiPIA7BeEsEgGI,955 +torch/include/ATen/ops/take_along_dim_ops.h,sha256=0EaGtFKQaK9GVN6nJMkBd67BwfwHUi9RTTMbaWUgZgA,2235 +torch/include/ATen/ops/take_cpu_dispatch.h,sha256=7FWXQJj4mjC5DDzvCcyFWIrKRILuFTnJIzg5HCIDWYU,1203 +torch/include/ATen/ops/take_cuda_dispatch.h,sha256=W7pCDCHcS0Zeqwk0U0ZEtQO3S4TGrclVP-jpOjPiPvs,1205 +torch/include/ATen/ops/take_native.h,sha256=PtbTYdgmS0DmEMgiQDkfCIG9ZvEeZ6Tkmk4SIYBK7TE,856 +torch/include/ATen/ops/take_ops.h,sha256=eUuVYWUGMPVcDXv91rDuCGo1CbSL54xkQ93NR6hmtCw,1961 +torch/include/ATen/ops/tan.h,sha256=Qg8l1Z-sIoWsg7ZVx93-RgtyNRlZ_-V5kV7y3vOAdSo,1384 +torch/include/ATen/ops/tan_compositeexplicitautogradnonfunctional_dispatch.h,sha256=w4l5wEuQHx4oK9UYiCKkhhvljve0ZQJTukuJ8u4it1I,1089 +torch/include/ATen/ops/tan_cpu_dispatch.h,sha256=trzZfZn-aDVkyhVio_K0AKNJR8KEULZzks0eBhQY-gY,1170 +torch/include/ATen/ops/tan_cuda_dispatch.h,sha256=cV7wjgxx4DrJjIAH9MP7R6Z48DVU6zcYSrd3iEjIlrA,1172 +torch/include/ATen/ops/tan_meta.h,sha256=lhLoE-k21uIROQU6KgNiSfx6sEiPYmELsqqKSbP8Cis,818 +torch/include/ATen/ops/tan_meta_dispatch.h,sha256=yZY_4gYPtWeSF2ScksoS-7-R006Oy1e_IWdauAHMEO8,1172 +torch/include/ATen/ops/tan_native.h,sha256=Hs6jUWODlLEQGErWtVVAywRsxHXxT62xbiCoRzHhhCY,1243 +torch/include/ATen/ops/tan_ops.h,sha256=UqayTHgIMHv66ZTJ1Jgm9rqZ4BrAEU4yfc2EKXJnHdw,2273 +torch/include/ATen/ops/tanh.h,sha256=kYp_0mwFqISveSTqRtTKYBfRdqVaG5ocIldAaU6HQkg,1397 +torch/include/ATen/ops/tanh_backward.h,sha256=ezRTu-KK2nhpBGZ7GicCvGZZjUktKxWZ7ny3yKiaJak,1634 +torch/include/ATen/ops/tanh_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=gJVoxe-WXmZeREyVz-QfYCERYWAivd2YdiNb-PqvR5g,1085 +torch/include/ATen/ops/tanh_backward_cpu_dispatch.h,sha256=e-mpAqy1J4hfXuf2vOJKuvAbnRqNZP4qJu8c_hfvLbw,1268 +torch/include/ATen/ops/tanh_backward_cuda_dispatch.h,sha256=nufDEwmP5T_896-UkyO4dVCW679J0Q6BiHUDMae2G9w,1270 +torch/include/ATen/ops/tanh_backward_meta.h,sha256=TAXzJe2rbGsOtSbf5XTKFoSHTnVlglT80tKmvJIgVPE,862 +torch/include/ATen/ops/tanh_backward_meta_dispatch.h,sha256=yJf7-YwZXPFeDJe_Bz3PHMLL0HPAYDy3KhypxUY8WPI,1270 +torch/include/ATen/ops/tanh_backward_native.h,sha256=DUJO-_IQirCBHVwmL1M0SKTyR_vRPvWD7FeZrfkOwVo,912 +torch/include/ATen/ops/tanh_backward_ops.h,sha256=BS3ybMiHGVtkPfYFZ0CjO__3IRbW9MXs0zfb01XUMOM,2105 +torch/include/ATen/ops/tanh_compositeexplicitautogradnonfunctional_dispatch.h,sha256=aqKJ1EwOsNGD36l-9LXJXkUnPuRN3TC8eb38yjJHkbQ,1091 +torch/include/ATen/ops/tanh_cpu_dispatch.h,sha256=2SsHyPlvZlX4-mIwk-6USAQKZWgeSfLQ_xTFHK2N1L4,1174 +torch/include/ATen/ops/tanh_cuda_dispatch.h,sha256=hKlzb5VLTVOu6BQ7iQMr6dC5EV-WHFbrKkCJA--vSZ8,1176 +torch/include/ATen/ops/tanh_meta.h,sha256=hkBWNzudzO0DCxSe9tnMWdKkOaSGaHYNFpgUjmmEaIA,819 +torch/include/ATen/ops/tanh_meta_dispatch.h,sha256=ykbeNysj288-4T20qVbMWKqs7T-ttRkV3pFyGNytRjc,1176 +torch/include/ATen/ops/tanh_native.h,sha256=kueQZXBZYfIHkwT8DXIsVow6iDGYDhYaAyecTk0u8wU,1560 +torch/include/ATen/ops/tanh_ops.h,sha256=7UuSsFu0a67RTSALWkV_L1uK1zWWxzKYf1oYDc_ZEXc,2282 +torch/include/ATen/ops/tensor.h,sha256=bCjxWtrcaqT6yjX9fpCSrf8rjUqUPYy-GyFTcJATMZ8,1885 +torch/include/ATen/ops/tensor_split.h,sha256=efT_6_tZiIzuLhijUBrGBAIzGkeFy3AWC0_ji_iCVb0,3460 +torch/include/ATen/ops/tensor_split_compositeimplicitautograd_dispatch.h,sha256=z7Nj8nXQB2YdKOkxakeSAtLLLVJyjauhXgqj5bYLlUw,1569 +torch/include/ATen/ops/tensor_split_native.h,sha256=k28PvuZEpNl3WEazefvMr_fe4BZP0n0V2Of13lXfNcY,1075 +torch/include/ATen/ops/tensor_split_ops.h,sha256=sYJ_EBA-YWX-oQzGW_1F8XSU0HvG6KhXoeQUSTTmY7A,3026 +torch/include/ATen/ops/tensordot.h,sha256=QPBJ_4TojGG5yd0p4N-MHzRiLmmbVLCRdrneu5-SAzM,1791 +torch/include/ATen/ops/tensordot_compositeimplicitautograd_dispatch.h,sha256=tAuEPyNR5kz_zRUhD2tgienUrhF59yhVOsN-ukwJsEo,1427 +torch/include/ATen/ops/tensordot_native.h,sha256=VwWQpfG0Ed3MlG4d1tnbl4wECAlwZ4a2wY5VJyTJ6Q8,976 +torch/include/ATen/ops/tensordot_ops.h,sha256=kbv4o4q_OKRIKxIunsKAs6c1prBwj4bbuD8_n1o7wW0,2349 +torch/include/ATen/ops/thnn_conv2d.h,sha256=I6YSsJNiruV2dYqPhLaRmcb4RkYboe-rSoGAW1D7Sqg,6874 +torch/include/ATen/ops/thnn_conv2d_compositeimplicitautograd_dispatch.h,sha256=a2rbACf0HJVGCfxfmDeAKMr9IjcTBwDSH8KysBwxAt4,2435 +torch/include/ATen/ops/thnn_conv2d_native.h,sha256=daD5ULTl-Nz2ftBRfic8sMaBTh5n7IIajxjZ7KXK9uk,1119 +torch/include/ATen/ops/thnn_conv2d_ops.h,sha256=OkQbjfqI4nUo-3eyGqa6SV8I4iKv3H9rTtynvlWDgLE,2903 +torch/include/ATen/ops/threshold.h,sha256=sAc16VXqsKz9I28SSH19oEVP2L_YZ1midyk-hq9_7t0,1886 +torch/include/ATen/ops/threshold_backward.h,sha256=DNi6T1wdRP38VeX0DlU4YfVhNtsgATfyshtynOlkXHM,1843 +torch/include/ATen/ops/threshold_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=N8Pm9v2Yqvy4bA8Kq6Dg7lf0abc-HDCnjV7CmZZpy9I,1118 +torch/include/ATen/ops/threshold_backward_cpu_dispatch.h,sha256=u57AxRobSuYCdoRkgMFDtVu35g9w63sshTpTsDauDmI,1367 +torch/include/ATen/ops/threshold_backward_cuda_dispatch.h,sha256=j4KfV8-UwIgMVJqVLpmczXkBdCbiaYp7jCZpd0Jw0P0,1369 +torch/include/ATen/ops/threshold_backward_meta.h,sha256=6J4aA7eeiPrJ4f61hZp7axEzRP8gYCDLHgszmtfq6eA,895 +torch/include/ATen/ops/threshold_backward_meta_dispatch.h,sha256=EvDNFutPr3Q8m2vACRHh08kM3voarQ70AInP6CAxQmA,1369 +torch/include/ATen/ops/threshold_backward_native.h,sha256=xIfOi-zOnZfD2m3bUQ9RQGwLp65OlBM5Qt7CviWDkZI,1845 +torch/include/ATen/ops/threshold_backward_ops.h,sha256=nt6VB43rGB7Sqq5WdVXO1cGVHNay6mSzhITohPmQlao,2319 +torch/include/ATen/ops/threshold_compositeexplicitautogradnonfunctional_dispatch.h,sha256=u_TGVLSgM00QuS3dX7or-r6ioFl_EME2RYoHqzFxAJA,1213 +torch/include/ATen/ops/threshold_cpu_dispatch.h,sha256=MDKpKlo2S5MlKBI06NM9Z-8sClebfx4l-I3k751sCfM,1418 +torch/include/ATen/ops/threshold_cuda_dispatch.h,sha256=jH50Z8yjlC3826l4ey52VkGCHjbJUSgjLZlxEHIq32Q,1420 +torch/include/ATen/ops/threshold_meta.h,sha256=GlsZ8nQLei8OHJVb5wDJOT5OrcA-ZvlOXVAXZZN2TW8,880 +torch/include/ATen/ops/threshold_meta_dispatch.h,sha256=No_lE_Wzu9U-tuplhJJfHaDj2V24xOxmIJSHd6Rl3yU,1420 +torch/include/ATen/ops/threshold_native.h,sha256=0e-KNb9oVqqsI1jsuR0UItVc6hsmtZXO0N9qLUzjubY,1042 +torch/include/ATen/ops/threshold_ops.h,sha256=B9Gw65He73UbJoWuxNsjbsQo_TuMIT9bi-daG5xIWIM,2879 +torch/include/ATen/ops/tile.h,sha256=GL5i3V9uVv8n0DRi5wMZAdmybXdupU5TE3b6mRIioQw,1605 +torch/include/ATen/ops/tile_compositeimplicitautograd_dispatch.h,sha256=EMeOEP8hHt3rmptFkE2CpFe0BM0BfaWeR_FNCTEMHQM,1123 +torch/include/ATen/ops/tile_native.h,sha256=j6v1wvEv0nqUwUJmMgYiXmrtw3MH6pKoi8wR4pEiTk8,761 +torch/include/ATen/ops/tile_ops.h,sha256=tll5NeidjQL9VeIZnMjXkHxtHXJZwHLuptuMlMl0hzQ,1293 +torch/include/ATen/ops/to.h,sha256=uk72uhzpd6xE_3F6SvmPNmULI62M3k_ShTOKOA_rFnE,753 +torch/include/ATen/ops/to_compositeimplicitautograd_dispatch.h,sha256=yQSNN-RKv0JwetkpoiVkI0pPHKNSoArSQuUvFumnn2w,1981 +torch/include/ATen/ops/to_dense.h,sha256=fHv8fHMwxJGM2YVgQbJP6_23jOWA5SFyLuP913EAOuE,759 +torch/include/ATen/ops/to_dense_backward.h,sha256=UMW-FWpJkh0tMXDqgJ8QxmlzcZiuG3pY6sxlEqBt8O0,1069 +torch/include/ATen/ops/to_dense_backward_compositeimplicitautograd_dispatch.h,sha256=LkbNxyMzNtk_lZVf7z2W86kfCh7uNP9_6W7unfF9Z2g,1105 +torch/include/ATen/ops/to_dense_backward_native.h,sha256=sak20s57gMHKoxbh2Ne2ZV51bgu2BeH1rsQZNgFab4w,817 +torch/include/ATen/ops/to_dense_backward_ops.h,sha256=1oUiWlfJGgPdKtLlxjHXXNadQFHOi5GzMp0X6vlXJek,1447 +torch/include/ATen/ops/to_dense_compositeimplicitautograd_dispatch.h,sha256=O7IQ3pznnIwtMXSNfXDoAJ4YQX2SxDWPUSuuWYFbAsU,1124 +torch/include/ATen/ops/to_dense_native.h,sha256=tLsx6TKoPjONACfnWgE2-oYdXwBHOl0rcHpBA8qF1nM,836 +torch/include/ATen/ops/to_dense_ops.h,sha256=wc95r1ro0NS9XGXQpbPIkbHsqe6BcYSjhn5ys1grDXs,1472 +torch/include/ATen/ops/to_mkldnn.h,sha256=JHw_YSLFe7KqbePdPnoUz6ixm2fYvJ9wX3qcHQf8tsY,1322 +torch/include/ATen/ops/to_mkldnn_backward.h,sha256=9jg9UojKNEQ5F4j7T125Q0G8__l5IIFFtjjwFJRiZ0c,986 +torch/include/ATen/ops/to_mkldnn_backward_compositeimplicitautograd_dispatch.h,sha256=AeolIabuJhnL-KIJNzvifupaSFe3Ez7FnU3lWzxLE-0,1056 +torch/include/ATen/ops/to_mkldnn_backward_native.h,sha256=-W4ODuWCpk6TQaJhvjsaj_73frS1xSghmtO32KGseA0,768 +torch/include/ATen/ops/to_mkldnn_backward_ops.h,sha256=fSjcp3eQnzz7U26tuKqdyuFqvdx72duQXkLykDQlMdY,1333 +torch/include/ATen/ops/to_mkldnn_compositeexplicitautograd_dispatch.h,sha256=JXldsfZXgtVaBBfRBIqFCiLkLMeGkGrtEm1adVgiofE,1220 +torch/include/ATen/ops/to_mkldnn_cpu_dispatch.h,sha256=kSek53SiILCJynzCJaO26d7yMmjZUym7sty89vLjBoE,1031 +torch/include/ATen/ops/to_mkldnn_native.h,sha256=gjxvDLEhNhuNCE5oFBI5Y86Aa6iHVc6Qo8hcub_7LbQ,913 +torch/include/ATen/ops/to_mkldnn_ops.h,sha256=I6ekRQWBZuT4vTDnS4A7DVSxPpC8vtWmIUEy7p1nbD4,2089 +torch/include/ATen/ops/to_native.h,sha256=varHgkKEtnW0GNu7oenaL3VJoWqEGhrdh269h0OyccQ,1546 +torch/include/ATen/ops/to_ops.h,sha256=iCZ0RF6M7hDvQ_DgtLHG0_z7hixbb1Gq67ObZGpxVdE,4706 +torch/include/ATen/ops/to_padded_tensor.h,sha256=fkYjvDqjUpX1dUugGLwT-HD2ulP_1KY3IgBPyXu54yc,4601 +torch/include/ATen/ops/to_padded_tensor_compositeexplicitautograd_dispatch.h,sha256=Uus528eR89tjNORMGNb7utCX81nWIqmQ-oedY5zPdYg,1580 +torch/include/ATen/ops/to_padded_tensor_native.h,sha256=hjsUyz5UrWxh0cPj7RLcil6VP81OOkfdSm1hc3Fi0y0,1130 +torch/include/ATen/ops/to_padded_tensor_ops.h,sha256=vxjx21fIaaCt5FnKvpmLKUF82JL2aBe2bU4mek4DPGA,2243 +torch/include/ATen/ops/to_sparse.h,sha256=Go8Ymb444cQPIoI0KpK7M_XqOKUqct5DvkcVvoSPheg,760 +torch/include/ATen/ops/to_sparse_bsc.h,sha256=84KDYvlR5z0LycGO2TLL6apJ7Le15v0AbZ3Q-SXOIoc,764 +torch/include/ATen/ops/to_sparse_bsc_compositeimplicitautograd_dispatch.h,sha256=vUxLgeTpjKdKGkPaQ6RrKrqEpeQpWQolwf11y1pQKWk,1103 +torch/include/ATen/ops/to_sparse_bsc_native.h,sha256=KaGHyVPlJ2O4SSmuQmFJYj0seiak1nFNWImOdoSiurQ,815 +torch/include/ATen/ops/to_sparse_bsc_ops.h,sha256=_LAnxrmEMAkz5Eh_XbzA5ktyCSgGhBvnpC69mG-4PDQ,1440 +torch/include/ATen/ops/to_sparse_bsr.h,sha256=7Ku9qtqm-THc32tRHnmOrUwyLRNeBt1bv0aKtaxf64w,764 +torch/include/ATen/ops/to_sparse_bsr_compositeimplicitautograd_dispatch.h,sha256=bUBegCGsWMbrMyPXhq782MqQz3awVEhSx5Hm0Ajxy3g,1103 +torch/include/ATen/ops/to_sparse_bsr_native.h,sha256=yDGMOHQh9vZFkWxhZAnm4rRAFWrHfTkWND5yFZfOTBM,815 +torch/include/ATen/ops/to_sparse_bsr_ops.h,sha256=JMsERtHsi5h17jZUlS14g0Qi6w45uQ3uCXmoWE3RxZA,1440 +torch/include/ATen/ops/to_sparse_compositeimplicitautograd_dispatch.h,sha256=T00rrQPvK4bqSQgAH-p9baMMbkD6gR-zri8UuBzJyzo,1250 +torch/include/ATen/ops/to_sparse_csc.h,sha256=0JHiT62bxrfu1td4bD9vq41GRkUFJ_sOs7pO46NWoyM,764 +torch/include/ATen/ops/to_sparse_csc_compositeimplicitautograd_dispatch.h,sha256=xIuO2y2kWSL_MtrRa8WCigMx34yd5G6zYZ5EcdhHycw,1076 +torch/include/ATen/ops/to_sparse_csc_native.h,sha256=llKNNxMTC2Vcf2Haurb7_HeEt7Dh42FGUNWNcARBvGY,788 +torch/include/ATen/ops/to_sparse_csc_ops.h,sha256=7WoXLai8dIuBkdMDUDh58Sdv3di80fO3NTYIXT43yVY,1351 +torch/include/ATen/ops/to_sparse_csr.h,sha256=L7QN-fW21xFE2qVQ4f8n5is1Qq_ti8h6a4TXiTCHxcs,764 +torch/include/ATen/ops/to_sparse_csr_compositeimplicitautograd_dispatch.h,sha256=UFV3BsJtY7yuZm-O4QExW5JEBBkxQA02TSZAMOQBBh8,1076 +torch/include/ATen/ops/to_sparse_csr_native.h,sha256=OKZtEXULBHoeJ2OcdzlvF5tKskiI3sbMRJr_Ipyj3QI,788 +torch/include/ATen/ops/to_sparse_csr_ops.h,sha256=n0LzU3IGgbGQi8wMdIiUFCRBt_MJFNmWauHl-3h_2Wo,1351 +torch/include/ATen/ops/to_sparse_native.h,sha256=rB39ccyiEwXsHa_DD_h1OaGs3QEiS30Gr_0KscDs9lc,962 +torch/include/ATen/ops/to_sparse_ops.h,sha256=V5l617qAoin5ADySrMoTGfEdm2USv9JcpDdEk68GyrY,2189 +torch/include/ATen/ops/topk.h,sha256=DJFAcnv7bdIkcMnRbxcY8VtX5d4AvvtBK5UwqibhSuA,5377 +torch/include/ATen/ops/topk_compositeexplicitautogradnonfunctional_dispatch.h,sha256=CYzlMPJGjTKBB9dpmfPiiG6DWUjQRNLBgqEYlqNenjA,1283 +torch/include/ATen/ops/topk_cpu_dispatch.h,sha256=lVrszZjuylFDrJEwodgBK5lVaABxUWmMkssMAACatLQ,1979 +torch/include/ATen/ops/topk_cuda_dispatch.h,sha256=tKCkxpc-E4HHeHfKCfBjNkSe8DYJjsR1nW4rsdVzB4U,1981 +torch/include/ATen/ops/topk_meta.h,sha256=22NNNTMoDh4WMXbJazzj8TlmY0QcN-6oA716xi5cVy4,870 +torch/include/ATen/ops/topk_meta_dispatch.h,sha256=CgUhpAtUtlHoGNsSNKp6CGbadxu-5w0nSGrlbUpnfVY,1981 +torch/include/ATen/ops/topk_native.h,sha256=foxXZmSLN-_8Nasgb2ohgS36lggMu9eJtiVo1p4xY6o,1309 +torch/include/ATen/ops/topk_ops.h,sha256=QWFXbKlTOSx8w_ph6MNtBMirXL87eaSga1hwtTXo5iE,2501 +torch/include/ATen/ops/trace.h,sha256=y7S-OMePDpM17m7LdhEYdeneyMnwi6nhtLQ9hlHMsCo,1271 +torch/include/ATen/ops/trace_backward.h,sha256=MbFkq3U4ipeFbMy2dvn2yzo5Ni10h_cUbSyjfmPdQUc,1725 +torch/include/ATen/ops/trace_backward_compositeimplicitautograd_dispatch.h,sha256=JLhJmSb-EBVNz9nC4jCfykCOqjGqWUWxHvE5zzWI88U,1145 +torch/include/ATen/ops/trace_backward_native.h,sha256=FNJMyBxi91KTBWKBjFOnuM4Px4Lu-Z-BCl1QzrWEgn4,772 +torch/include/ATen/ops/trace_backward_ops.h,sha256=rnEitzif5VYNzCjrsFkSzrv0_taYgeQruQu1jdczF6w,1326 +torch/include/ATen/ops/trace_compositeexplicitautograd_dispatch.h,sha256=v-i6rApLCtmAh022Kug2W4c0awkU27NP-LYg3V8kMWw,1119 +torch/include/ATen/ops/trace_cpu_dispatch.h,sha256=OkPFiwZ7GUBARcii2l94kaA46nw06bqKWGuUOtRy8rI,973 +torch/include/ATen/ops/trace_cuda_dispatch.h,sha256=NrIWj2X4lqIT9Mx7hGZ1QvydlPAS5bwiWtczfbKdJ_s,975 +torch/include/ATen/ops/trace_native.h,sha256=RT6_vEEwGTEwNGaAAVkeY2RWMjp88vrsbCh2cXzlCG4,868 +torch/include/ATen/ops/trace_ops.h,sha256=LDBJR55dDX_FTAECQAllgWuTPatlBv7YQ2iEoiAQTBY,1795 +torch/include/ATen/ops/transpose.h,sha256=gEkAxHY9nguD89H-zzKQkRFtrMfp9KdIidvmcGFfjpA,1219 +torch/include/ATen/ops/transpose_compositeexplicitautograd_dispatch.h,sha256=vDpiK0QPLh2y9JRskDzMwiXyYI1JetrPf39_7SZT2C8,1131 +torch/include/ATen/ops/transpose_compositeimplicitautograd_dispatch.h,sha256=T_-uy500xbidVW-JkgrSK-chVKa6FSSHtz2x93iByt8,1057 +torch/include/ATen/ops/transpose_copy.h,sha256=_67SbEIGUcegzQb2e52MNwyszBcgUyYEHh0D_d0zTWw,1565 +torch/include/ATen/ops/transpose_copy_compositeexplicitautograd_dispatch.h,sha256=u5QsMpPNQBQB-mUpVbgz2tbli4jEkw9hBvkpXiJ8ea8,1193 +torch/include/ATen/ops/transpose_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=dMqvVp1kjSzPme50lEK_JE6LHgmyCMIYfQUfERR0wss,1080 +torch/include/ATen/ops/transpose_copy_native.h,sha256=dLXUeM8t6rQSMXrYCkPYoLbRucA9If3xvcZdNAKKJ1E,888 +torch/include/ATen/ops/transpose_copy_ops.h,sha256=POhI_Bsu47uoB3iepiP3oivv9bsd9OSxgoakRz11dKg,2060 +torch/include/ATen/ops/transpose_native.h,sha256=ldkrpGYUE8qqPJG7PkqpWwDZNE_GJMfB6e8KcXdSFcU,1028 +torch/include/ATen/ops/transpose_ops.h,sha256=4Si_blynot-eEXgiHIDpNzjKVOJVxs-DlwEeNNT3S4Y,2597 +torch/include/ATen/ops/trapezoid.h,sha256=Df62Q_BticehSk2JPI_Y7jpdI7GVwFgmEebOA-a021I,1188 +torch/include/ATen/ops/trapezoid_compositeimplicitautograd_dispatch.h,sha256=RI2CHkvX_RtpVqVuz4EYJ-Eu38OnBscCuYp81lie4F4,1151 +torch/include/ATen/ops/trapezoid_native.h,sha256=_Six84Wz_6APyeC7ddBwylfx7IBJyRKuLlCrDntcUgw,863 +torch/include/ATen/ops/trapezoid_ops.h,sha256=MmahfIp4fXXA_FfGDyaTEo_SQJzMNTTVwh4TgDc9C54,1977 +torch/include/ATen/ops/trapz.h,sha256=WNqjA_Xx39vRHK9HhUEADozdHaOr4J_OEHNSKm8mXSg,1147 +torch/include/ATen/ops/trapz_compositeimplicitautograd_dispatch.h,sha256=9ETB9eJ1rYs3-mgf9NjoQPQlk4AfUS6mt13sGloXV4A,1131 +torch/include/ATen/ops/trapz_native.h,sha256=1Bvfk_A3x8mo47m1z_SFITgVgZ9AD1JF0EWtjP0RX7g,843 +torch/include/ATen/ops/trapz_ops.h,sha256=sqyX0dG-AFxA3FVUi8XMe4iYgCRv-6sVtppzyVVx11A,1916 +torch/include/ATen/ops/triangular_solve.h,sha256=zgJ7wCvEXOM8wEheYnJ06apNR1zLBh7lbZT--0UNQQQ,2205 +torch/include/ATen/ops/triangular_solve_compositeexplicitautogradnonfunctional_dispatch.h,sha256=wmtrooultk359PieXB44iyV173Iw7dLqpg39s7LEo8c,1166 +torch/include/ATen/ops/triangular_solve_cpu_dispatch.h,sha256=tyVcfo7LH1sv2FankB_agR-zwgOAsxJsi_dm1WTVm2s,1512 +torch/include/ATen/ops/triangular_solve_cuda_dispatch.h,sha256=sIFedQ68QDvyD8aNWEL41affhMjj-js90iotDcXj27Y,1514 +torch/include/ATen/ops/triangular_solve_meta.h,sha256=3TYkaBDAGCRSRd1IqVJtEItR4YKgAAp1oqOTlEhIptg,901 +torch/include/ATen/ops/triangular_solve_meta_dispatch.h,sha256=sIrwXPY6kzI3AJ92b71JQewFoVgg2wjH1ffsbT94xU0,1514 +torch/include/ATen/ops/triangular_solve_native.h,sha256=Y_3GozTlRh3MyK-DVgEUBJEZVRfZXiRC63xzHzcPWSY,1399 +torch/include/ATen/ops/triangular_solve_ops.h,sha256=APK_qWnjrnif-Z91iPbNBBs2bqMIZJZwrwrkOgdA3ks,2651 +torch/include/ATen/ops/tril.h,sha256=hWmTyK0ebCcfdty_GI816MPKp_Mrqh9LRLA9dM0NPXU,4013 +torch/include/ATen/ops/tril_compositeexplicitautogradnonfunctional_dispatch.h,sha256=7MonXhoFn1QmPWqWLxNSDUZsfVAlscyPLYT87l9QZ28,1294 +torch/include/ATen/ops/tril_cpu_dispatch.h,sha256=FhClPP2DOzfhqPToYaRUkTaFyayyY5yuDhmzy8IFUQ0,1628 +torch/include/ATen/ops/tril_cuda_dispatch.h,sha256=Pw0fYbCZciBsFDCgweji_lkpBa-CV6JQ8FHw9sZOj9Q,1630 +torch/include/ATen/ops/tril_indices.h,sha256=DNkN-Gx0RZYYcyW6Xnc1m2mZGHdCux2_tG9DKC32JW0,2214 +torch/include/ATen/ops/tril_indices_compositeexplicitautograd_dispatch.h,sha256=YS6sTcx5bKhzyeINbU15uSx6CmKYyC96KeZgZhTT-Tk,1169 +torch/include/ATen/ops/tril_indices_cpu_dispatch.h,sha256=N-VfIgdiwoIGmEbdGNmJ_9uCYowLpzw3dCXAXDQ3AD8,1258 +torch/include/ATen/ops/tril_indices_cuda_dispatch.h,sha256=yUk8qARfloPH57aczIlY3uCtAX7w76JJEa5eo3MGreY,1260 +torch/include/ATen/ops/tril_indices_native.h,sha256=4CWekTKKNvTQEAm6RI0kizHvZ5WKOmIrAe6_L0pZsHk,1258 +torch/include/ATen/ops/tril_indices_ops.h,sha256=jolBeqTHLaYt5_HzhRTMLV73zKLDIuDYTvtmRF35yZw,2453 +torch/include/ATen/ops/tril_meta.h,sha256=6SRAPZRNOUvDonbFzhzM84POk9SztKD2OCyKwREqx_A,837 +torch/include/ATen/ops/tril_meta_dispatch.h,sha256=ef-uTzwH0AuCdz80PtIgi_N0aHLypfl4_KmVeK4qNeE,1630 +torch/include/ATen/ops/tril_native.h,sha256=UtWCWB2fyB5PN5osPg3HRqZXCYayhy_uuAEHIEMNACo,1018 +torch/include/ATen/ops/tril_ops.h,sha256=LDJVlJQxta0IcDYuTS9hQHfxIqj0jo_nMvuf0pd5RaU,2510 +torch/include/ATen/ops/triplet_margin_loss.h,sha256=Hmr39QHwiLBtms6GCYvRVgjgFxDwQvY_6VGxjiNu-PI,1282 +torch/include/ATen/ops/triplet_margin_loss_compositeimplicitautograd_dispatch.h,sha256=ygsQoa7Hv103vD93JaTKTNXWFaoffYFJjGRTfPUI7zo,1196 +torch/include/ATen/ops/triplet_margin_loss_native.h,sha256=mdhEb5IFPQzZV_gWKkJ-QDP90LWdU2GW53zAnuVpvI0,908 +torch/include/ATen/ops/triplet_margin_loss_ops.h,sha256=Hvfe_kINZBMt05RhExIbXY0h-fimQnFv9NtYSUY14T8,1702 +torch/include/ATen/ops/triu.h,sha256=cW4-55KxSo6a4Phy2pikk1DQqGjy4b5mmkRO8-pEsB8,4013 +torch/include/ATen/ops/triu_compositeexplicitautogradnonfunctional_dispatch.h,sha256=IQv_0mjONwZaqwbzN3klw-xDAoE6Jx6aTo43aAeLCSc,1294 +torch/include/ATen/ops/triu_cpu_dispatch.h,sha256=TfDoEX__-JW8wSnZk7sMJx-eFb_9qbOaqvI67F5Pit0,1628 +torch/include/ATen/ops/triu_cuda_dispatch.h,sha256=klxQIokRO6heIi5uHKb0pQQoEdS8B72BxXFp0wzRlFg,1630 +torch/include/ATen/ops/triu_indices.h,sha256=-b8__dGDHUuy1yEdrP2FwjVjhsCTfqLdXO79C4TE3kw,2214 +torch/include/ATen/ops/triu_indices_compositeexplicitautograd_dispatch.h,sha256=p99twpuIiW-I4JYsLnOBEm7oTj0HE6mSapRo6Vr9kgw,1169 +torch/include/ATen/ops/triu_indices_cpu_dispatch.h,sha256=ThiFwj8tvlUiwO8Rpv9JnthVQBtzklleI4sLO61njn4,1258 +torch/include/ATen/ops/triu_indices_cuda_dispatch.h,sha256=4RbNsfGAwm-8eptEXusVFfyxGr9BI1CvlHbyetbuRWQ,1260 +torch/include/ATen/ops/triu_indices_native.h,sha256=vZgKik7LpKla_EYEaNIj91QS62d_SiiBZYD97-FvxAw,1258 +torch/include/ATen/ops/triu_indices_ops.h,sha256=UzN2jCOt_7mfG6MS6-sjQ-Q_R2DiCeXxkAqHpfcV4LI,2453 +torch/include/ATen/ops/triu_meta.h,sha256=R3oRI2ROWLR9TN8lKeeriPQSRvT1KIVvDCArvI5x-t0,837 +torch/include/ATen/ops/triu_meta_dispatch.h,sha256=c0DWcTujwcTJ1guoyN7tsYv1QUUVUAINdlrk3Vx0eJw,1630 +torch/include/ATen/ops/triu_native.h,sha256=NJzfQvtGCktCcMroztM33hdPBseg73Jtk5EtNbevXiI,1018 +torch/include/ATen/ops/triu_ops.h,sha256=xzUvf4tfQnyJ6ooE5yKZp4GSq6w-WrdT20t1IZDTdR4,2510 +torch/include/ATen/ops/true_divide.h,sha256=vVM0u8Jt-tnib5lGLU34nX8m7_KGsLqJgQTRQ-iKUQQ,1697 +torch/include/ATen/ops/true_divide_compositeimplicitautograd_dispatch.h,sha256=GTK2kBA9vEY5dlfzn6PiPi8AZy9NgyztGDf976AQXM8,1517 +torch/include/ATen/ops/true_divide_native.h,sha256=HjUGdggLMAiR72SvFqQK62VkWjd-sW30C5nR6av81EA,1119 +torch/include/ATen/ops/true_divide_ops.h,sha256=izXjsaT_Z7c4PutaawTzj61mGAMkyAOhk9TcrelF0e4,3884 +torch/include/ATen/ops/trunc.h,sha256=bThHrJg42Lc4jBItMj1ALdy3fba__G0GWaekARXavO4,1410 +torch/include/ATen/ops/trunc_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ezkCR1YK0N0ehZyEr3AJ0yepOssUvxq6o7mfdK-8DrY,1093 +torch/include/ATen/ops/trunc_cpu_dispatch.h,sha256=XtMNanhlVlfQehDBd-iMP1o2JUas2Ui9Q4AStZe1MnA,1178 +torch/include/ATen/ops/trunc_cuda_dispatch.h,sha256=WnmnDuQoHLYgRRYqhIswwtHq_ypKfDO21q3QVjhAo7M,1180 +torch/include/ATen/ops/trunc_meta.h,sha256=ePhwtjbH5Y9x2TakgYNupYJSqdVtCDpTDOjczOSS3pY,820 +torch/include/ATen/ops/trunc_meta_dispatch.h,sha256=7fixKYBsxpqVDmrtvnBc3OiPP89HBlL6A8VdgCdC7kw,1180 +torch/include/ATen/ops/trunc_native.h,sha256=1mUL_UoI3Ez1Uro5gTZqKsb3boeZlafXdypYbO93HC8,1261 +torch/include/ATen/ops/trunc_ops.h,sha256=N-k-omO5sUAma2FQPIJuEdoz38AkXHgt66fpDtYue8I,2291 +torch/include/ATen/ops/type_as.h,sha256=AwdVDvYR91dnzV2NBKn3jmaFZw3E1GQFUYiHbz2oT3g,758 +torch/include/ATen/ops/type_as_compositeimplicitautograd_dispatch.h,sha256=dv8uyoCMjTCYB5RtdBBuj2oeTDCJBd6lqxvXMCbjW8Y,1045 +torch/include/ATen/ops/type_as_native.h,sha256=v2LPDmsYu4iQrg66e88k11xbntoyJ14YFfyJmwBTg7Y,757 +torch/include/ATen/ops/type_as_ops.h,sha256=OioOs7XMmXtex2G7ERpJo1-uxjvO62JKxlZoRU9Lll8,1300 +torch/include/ATen/ops/unbind.h,sha256=Us1fppseabi-x5bkZQAWw8Up3CP6joOhr2tAap8wf0o,1172 +torch/include/ATen/ops/unbind_compositeexplicitautograd_dispatch.h,sha256=1AUAnZzYjo8pXCoxPwyvOb_rhy_YBKk2Ppv5BvzxwcI,1048 +torch/include/ATen/ops/unbind_compositeimplicitautograd_dispatch.h,sha256=VZms1q-XD70XSJatP9jiyjnP3gSNDYwvofA9zu7MLTs,1050 +torch/include/ATen/ops/unbind_copy.h,sha256=odoE7jy3xo_wzzF511YrzFRLq2bTvx1UZVPTD2GhoqM,1439 +torch/include/ATen/ops/unbind_copy_compositeexplicitautograd_dispatch.h,sha256=SCiFQGUcY305ijwIzFe_HMbD5aOVW4x7At7HtuqUicA,1147 +torch/include/ATen/ops/unbind_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=WzJCqQShrinmnjWHlQb4ECZ5RZPu6cxcwKEKw7YVzkY,1079 +torch/include/ATen/ops/unbind_copy_native.h,sha256=o4DwnRwMMa-rrmh_KAYsBbEvKLvv0T6y0711c8OwAJ0,863 +torch/include/ATen/ops/unbind_copy_ops.h,sha256=zCODlW7yla9Ef1r2SyaeWyJZu0wCadw1TbgM2eQILSo,1969 +torch/include/ATen/ops/unbind_native.h,sha256=a96IIfvlEmJBsrG26jl-yOqlBBBhidVaUX8iasLjYFc,943 +torch/include/ATen/ops/unbind_ops.h,sha256=8nZXbe4nLJa79FLWC74hLInAxrRbDuzw4HOqjagw-us,1967 +torch/include/ATen/ops/unflatten.h,sha256=XY0zghsNhL4xIFsW-0qta3fhq5EdOjhwLfnh6DtI9kM,3035 +torch/include/ATen/ops/unflatten_compositeimplicitautograd_dispatch.h,sha256=q3pmv37xLSmzONSyxLm5pIOXb-EGLqRmYMqXirOCpIQ,1412 +torch/include/ATen/ops/unflatten_dense_tensors.h,sha256=hPGFTdnsuy3OOLY9swgY3XBM0l_6I88z5qtUBnA-J_c,1027 +torch/include/ATen/ops/unflatten_dense_tensors_compositeimplicitautograd_dispatch.h,sha256=6n0V_CTfUS0ApPHEaKQH128kwn15EnCOFM6BTktjuzc,1074 +torch/include/ATen/ops/unflatten_dense_tensors_native.h,sha256=g1B_XmzUdt_8FAdIAPpGhj097yd7Gu4_UQpO82UwuRM,786 +torch/include/ATen/ops/unflatten_dense_tensors_ops.h,sha256=V0ocHe-XxJUUnNo-MYjP2hUTZE1uf-6WmdP-BdjFkcI,1391 +torch/include/ATen/ops/unflatten_native.h,sha256=7FTdxEzeKUI5vW5JvsE7vGlQ_Rj55LHF6Tha9iyi7PQ,919 +torch/include/ATen/ops/unflatten_ops.h,sha256=NtYccRRzkcys92u6DLFiLCoF2YHaWGK__rRivOiFmXM,2141 +torch/include/ATen/ops/unfold.h,sha256=s3T99PCBveeZExVaiyiTCPCZl0eEwzNqIKcvNXHGEd4,757 +torch/include/ATen/ops/unfold_backward.h,sha256=6LLF2qOMzqe_M8-sOPDTgT5EYiqzpyGxiDkYPwIu95k,5194 +torch/include/ATen/ops/unfold_backward_compositeexplicitautograd_dispatch.h,sha256=unmx5aSR3mko0ZWHt3ohc3o-kF4Rl9pjgHA_usng708,1628 +torch/include/ATen/ops/unfold_backward_cpu_dispatch.h,sha256=8hNRxDa0O11AYJ21tLQWSF8BexFIYwSHcjTxh6KcrzU,1203 +torch/include/ATen/ops/unfold_backward_cuda_dispatch.h,sha256=0QcM2lEahMqldeLaehuMMi57Xldi6VUoongCt-IT1Rk,1205 +torch/include/ATen/ops/unfold_backward_native.h,sha256=j8JFbKSh-UNSJeSH-ax9Yqsa_0sCadYmwky1lcCVdII,983 +torch/include/ATen/ops/unfold_backward_ops.h,sha256=2ZJwAeZKszpYe0ZFJiqJfqpzJ8nMzQt8e9NRM0OCj6A,2367 +torch/include/ATen/ops/unfold_copy.h,sha256=a2ab_KuD3jBa6cikEArKh_e6lOMsc4HngtRjpSF_Ou0,1646 +torch/include/ATen/ops/unfold_copy_compositeexplicitautograd_dispatch.h,sha256=Tu02oAKCNR3swxznEOAuLhY4Bi-gxH_RlQOEfUpYSrc,1225 +torch/include/ATen/ops/unfold_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=bHCFrbr20k_OQl1EBusSBL6EzSEYkL9RfZldhqyo7Gc,1096 +torch/include/ATen/ops/unfold_copy_native.h,sha256=cTUzPy2baCqqZOYXFFasTflV2E6vcQRjyztT6d_zxgw,912 +torch/include/ATen/ops/unfold_copy_ops.h,sha256=F0Lc3nSJkbVCsq9Cnd0hlUwaZTlIG-Ekr_72Cry20yY,2143 +torch/include/ATen/ops/unfold_cpu_dispatch.h,sha256=LQzs-kZ6aZ6HBSaj6ro4Okh7kwfXQ4_fN_Rk6PwPBbY,1021 +torch/include/ATen/ops/unfold_cuda_dispatch.h,sha256=cEc1Mn6UZnpLXwsnhT-cTiXkBzY5_LwL6YQfx86K5aU,1023 +torch/include/ATen/ops/unfold_meta_dispatch.h,sha256=ZENsp1SbM7bDsRO9Ta2Wn5Nitg5_np-wt8e7MgxNSMw,1023 +torch/include/ATen/ops/unfold_native.h,sha256=r4RtD3L7ZsPvzxbOSQZbFnMUQhJoEQfzbUs7MWBq3CQ,777 +torch/include/ATen/ops/unfold_ops.h,sha256=trzss8tbLxPJIbEhC6MvtwSk58KkFNFRJ_950-bGi7s,1373 +torch/include/ATen/ops/uniform.h,sha256=vudakx8sdSljf7Hee_qgm72Im0ZXRdl3-sx4RDa2bQg,1752 +torch/include/ATen/ops/uniform_compositeexplicitautograd_dispatch.h,sha256=Osq3PHPGpq97Xc0pML9YQ_knpVpFrdVnF0K15t-y2Q4,1414 +torch/include/ATen/ops/uniform_cpu_dispatch.h,sha256=Jv2-6RGQtimjkBYTviWVF0wHhuxsi-ShApxxcon9Ikk,1057 +torch/include/ATen/ops/uniform_cuda_dispatch.h,sha256=S0hqmznLSPxHUUY0cILYLPUAEKFpVSKsSXBXjbb8hWI,1059 +torch/include/ATen/ops/uniform_meta_dispatch.h,sha256=OHslAJlbhpCL0_r1rjvDT6k4FQ4CCKOaVgme73q3udA,1059 +torch/include/ATen/ops/uniform_native.h,sha256=OICD2Mu4nG6XIO1fj-I0qLgqRqKk83XhSXquaWVyzJ0,1240 +torch/include/ATen/ops/uniform_ops.h,sha256=BEzEvPH6roudbfOz0R1sRYWdYdsRPpo8TTjwA2eDWpI,3014 +torch/include/ATen/ops/unique_consecutive.h,sha256=EgYnf30mTLscYKb5YATv8HmR1tHyDnEGlfpjxxTshsE,2340 +torch/include/ATen/ops/unique_consecutive_compositeexplicitautograd_dispatch.h,sha256=ip2_E_cZrTnt8M7M4kvTSGsCzIBuCc_w88Vr6axSLzA,1472 +torch/include/ATen/ops/unique_consecutive_cpu_dispatch.h,sha256=LCX03ZhiWrdgeacTAfMpJoyfKbVzCNHJQ-jaRU1hrvg,1120 +torch/include/ATen/ops/unique_consecutive_cuda_dispatch.h,sha256=cC6Zkl1Fa5Q2Dm8FhQk9hRFeUHrTjM-LPUqQxaQ0CPU,1122 +torch/include/ATen/ops/unique_consecutive_native.h,sha256=xdwhZQKoVU6D38bQhqV5bPvxplPqkEzSRL9WGF1nj30,1325 +torch/include/ATen/ops/unique_consecutive_ops.h,sha256=DUo_TOLEHqv0YCigXDl3eE9gaQfpFffgV5OHgco8UHw,2782 +torch/include/ATen/ops/unique_dim.h,sha256=6s8jKULkqYrRBHKdgiB2vVPaGewbAq2o-LnHDCYS048,2288 +torch/include/ATen/ops/unique_dim_compositeexplicitautograd_dispatch.h,sha256=shtRtoF3VSv1wVzB5OkzvJc5S6I0CyvGsB6l60UuJ5I,1438 +torch/include/ATen/ops/unique_dim_consecutive.h,sha256=LYfEJUSF6-Kx4qPz9udEZt5kbtcxldi07VaqQO4DAUU,2281 +torch/include/ATen/ops/unique_dim_consecutive_compositeexplicitautograd_dispatch.h,sha256=MEAiIko-yvrWz_A3MlNFo_aKhr0APGTBysqQOX-9PIc,1431 +torch/include/ATen/ops/unique_dim_consecutive_cpu_dispatch.h,sha256=VkJHGj9AUjpULH0san-ZJcKw_m0KGs8JV4A9VYHt6bQ,1092 +torch/include/ATen/ops/unique_dim_consecutive_cuda_dispatch.h,sha256=9OdxQ2m3pjc1yyX8bI0kDTZUZrWkFYuqa9vyuzMii6U,1094 +torch/include/ATen/ops/unique_dim_consecutive_native.h,sha256=338Q6ZNjaZKtKvYqvxa2TPXxAy5zKjLAEijXg0wT7MU,1256 +torch/include/ATen/ops/unique_dim_consecutive_ops.h,sha256=0uBeRnSK6KCTZ7pf-gkmVNbAqdSYNu3qRf34o65CM1k,2692 +torch/include/ATen/ops/unique_dim_cpu_dispatch.h,sha256=EvZku3JTmuxuOy5amEBWHZnEv-QieaAtOYM-AuKvScA,1098 +torch/include/ATen/ops/unique_dim_cuda_dispatch.h,sha256=gtyuZJM_xaIwkERvmt2SaM39tUKUr76jIa5cF73mJbk,1100 +torch/include/ATen/ops/unique_dim_native.h,sha256=QRPo2F-tnfoUStfotsY-fQX1nO8cpPVejMtbKDnOdEc,1269 +torch/include/ATen/ops/unique_dim_ops.h,sha256=0AbjIDh-58vw9vJfDGKpEnotyhJO5Jepl1tHXKqeVTI,2720 +torch/include/ATen/ops/unsafe_chunk.h,sha256=ZwioqB5WtrkuEaABDz2_4y6naqIfbOQHnWUOp4bEEXY,999 +torch/include/ATen/ops/unsafe_chunk_compositeimplicitautograd_dispatch.h,sha256=nnNOSfnzJzhvD1NySrLD2APoNW_BqT87WicybE9OGIg,1070 +torch/include/ATen/ops/unsafe_chunk_native.h,sha256=58WRzG0jkJgeDpmTETv-J9_kVe9igwAp3d5WRyYnqCk,782 +torch/include/ATen/ops/unsafe_chunk_ops.h,sha256=NZgRwR_gfq1kcG71bgDIXGPEZaXwvokyWKcJA50bx-4,1375 +torch/include/ATen/ops/unsafe_split.h,sha256=udX3V42Z25IjjIMaZI9wsbvpNBJCg4L5FIvEyuFmeYc,4263 +torch/include/ATen/ops/unsafe_split_compositeexplicitautograd_dispatch.h,sha256=tXaEXZQXb7YC16OGSkQPaa3WhTEv71hEYzFurkzwL7w,1667 +torch/include/ATen/ops/unsafe_split_native.h,sha256=n8TRnYCFC8VcNJfGMgcKCqLlixX5Iq0oQh48KUnYZR8,915 +torch/include/ATen/ops/unsafe_split_ops.h,sha256=9VewGYHzwA-s5VNVryD-dkh7I0XBfH_ImwPU_Up4k4s,2153 +torch/include/ATen/ops/unsafe_split_with_sizes.h,sha256=n2TgDEfYBGpwBc4lfO1LbxPffPwJ_lJE1HsAdSgsuYs,4772 +torch/include/ATen/ops/unsafe_split_with_sizes_compositeexplicitautograd_dispatch.h,sha256=iGy0DN6Fz0l7gitayqAv43f85zBohiqo-3j4ZobyJ0U,1787 +torch/include/ATen/ops/unsafe_split_with_sizes_native.h,sha256=bP9WJI3IexJJkAnDkC2s1gp_eUEErmG-ZWdQiAkHFKc,948 +torch/include/ATen/ops/unsafe_split_with_sizes_ops.h,sha256=i4k-PjlXR3V4FfJO7NRofwsZPeHJY_aPGJ6fIxQc-gc,2236 +torch/include/ATen/ops/unsqueeze.h,sha256=uaqaxIIvhFyw46gXSJqKZsH8oWTrDeezZPKH__mV43k,936 +torch/include/ATen/ops/unsqueeze_compositeexplicitautograd_dispatch.h,sha256=y6D0SxsL3aQH-JA_1ATfZAdlFfwzBP0MwZN9v_9KU-o,1101 +torch/include/ATen/ops/unsqueeze_copy.h,sha256=ulF74zjjITbg2oNi4JH5U7MAA7TWmL7k8o0-GpZGbBw,1442 +torch/include/ATen/ops/unsqueeze_copy_compositeexplicitautograd_dispatch.h,sha256=RXBdoSDxk5-n4IXqOauFoR74IYVoY5ViYX-UWy1sjjQ,1163 +torch/include/ATen/ops/unsqueeze_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=bpxg01nCb_nYLIvqY-jKUeojiSpP6-aAj-SsRs3ZakQ,1065 +torch/include/ATen/ops/unsqueeze_copy_native.h,sha256=kcmh6t_O3JHpSSc7xlnv-luwW_heskS80hCt4LeguVc,850 +torch/include/ATen/ops/unsqueeze_copy_ops.h,sha256=fzNr2J_vu0N5HDHX1t-ciwWzjqGePxu4b44Q1D7PIo0,1937 +torch/include/ATen/ops/unsqueeze_native.h,sha256=jKux_Juj458LEfBG_P61YATAA06JEpEE72FcGPFpqkc,1047 +torch/include/ATen/ops/unsqueeze_ops.h,sha256=Jn3hsOhjPmCk9lqmg8scStR2-mIAFVunQiB18nsqjxc,1822 +torch/include/ATen/ops/upsample_bicubic2d.h,sha256=VFG7Sjs3cg54poGpR31zQKO57-cW4YXMRB2f7GUOi9Y,8192 +torch/include/ATen/ops/upsample_bicubic2d_backward.h,sha256=xohEMXqS3OFZ-BqCGBZS8qjXSkOuDpgTAmvDhoCdDCw,7938 +torch/include/ATen/ops/upsample_bicubic2d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=g0d07D8UAM7dis5G1vaj2eSQFbLukIsthdqS7OlLg9A,1519 +torch/include/ATen/ops/upsample_bicubic2d_backward_cpu_dispatch.h,sha256=9Yqn5LLtryH7gKMl3mXRydYQZKxuZIt2W6BiX6cgG70,2573 +torch/include/ATen/ops/upsample_bicubic2d_backward_cuda_dispatch.h,sha256=DYKpeD_DCL4vXznMLKD6CFhukQpIlTd3FbWTEsVbA68,2575 +torch/include/ATen/ops/upsample_bicubic2d_backward_meta.h,sha256=TwzAviaFoWdnSq8wRpPHi8kGcjJogfR8Ojki6MuxuR8,1006 +torch/include/ATen/ops/upsample_bicubic2d_backward_meta_dispatch.h,sha256=S0RysAHoVnXgESmOAdZmg6k548HPetSwXwzOl_hAuTM,2575 +torch/include/ATen/ops/upsample_bicubic2d_backward_native.h,sha256=ZodaA34-MTFpHkguqLRrbY6UosObu6brVHxx7Plw5tg,1447 +torch/include/ATen/ops/upsample_bicubic2d_backward_ops.h,sha256=YSda_a4UhAkCJNqe4KNEiFdPdc-0TYWAwFlQsnNxXJU,3037 +torch/include/ATen/ops/upsample_bicubic2d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=GJRI66OCMJ0ohJusXTX7lTKXYJK9wN1Ik0BkGVt7lJ4,1427 +torch/include/ATen/ops/upsample_bicubic2d_compositeimplicitautograd_dispatch.h,sha256=g_28nb7pK3wbDQFt-lpLbsksbGGyJMw7s4OmbEGewqw,1328 +torch/include/ATen/ops/upsample_bicubic2d_cpu_dispatch.h,sha256=jB8LfeRwokVGHN6A_e6pg7lBvU8dImc5Hv5EJIL4bP4,2269 +torch/include/ATen/ops/upsample_bicubic2d_cuda_dispatch.h,sha256=tUdJWbV0oKP06DyMQkE3eq7eZ1sI6ZdBhJrKZhw-qvU,2271 +torch/include/ATen/ops/upsample_bicubic2d_meta.h,sha256=Ob_5soB-Ow95Dl99VPXlgwM6avIL59jX_3sq2xKAY6U,956 +torch/include/ATen/ops/upsample_bicubic2d_meta_dispatch.h,sha256=ev5daEgNn9PCJ8etxyp3RaLullGOchlFS-xCy5RXESo,2271 +torch/include/ATen/ops/upsample_bicubic2d_native.h,sha256=a09SH9SdyDfsvgQnTRo60-xVq2QlhFSojRRPFmLX22w,1483 +torch/include/ATen/ops/upsample_bicubic2d_ops.h,sha256=s65OEj73Dz_xkDeg6iKQh9ptYKQ-Ddke1ltVVBzvPUM,3601 +torch/include/ATen/ops/upsample_bilinear2d.h,sha256=2SjYLWupIwtLpSWtct1TGNtx65JWveHxQLSSoTG-oA4,12147 +torch/include/ATen/ops/upsample_bilinear2d_backward.h,sha256=BbykxVTg5wsqn9REo43MSFuDkqG0BDCtpOOY5K2uNLM,7969 +torch/include/ATen/ops/upsample_bilinear2d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=u8y3eHcsdzIvQBxvWPvSSzt0hnCesnTHVPYgxNiQVU0,1521 +torch/include/ATen/ops/upsample_bilinear2d_backward_cpu_dispatch.h,sha256=xpWQkhlw7ixr_NWPqmk5Jg36HV-GWHDPpSJ3Z-pScGs,2579 +torch/include/ATen/ops/upsample_bilinear2d_backward_cuda_dispatch.h,sha256=vwxus4eIHVgo7cBz-bxyHzNOc-cgGTKX3hKUh3B-z0s,2581 +torch/include/ATen/ops/upsample_bilinear2d_backward_meta.h,sha256=y3g_dTcr6vGo26ehNPK3Viptne_AbZGIMGbdt_omaMk,1007 +torch/include/ATen/ops/upsample_bilinear2d_backward_meta_dispatch.h,sha256=P6oJdzjsbmsEOttA3_xDm-np4qqylMsA1dY5PCYD69w,2581 +torch/include/ATen/ops/upsample_bilinear2d_backward_native.h,sha256=pvItM3mfrh4ib5MqVl4Sl1l-dA1XbeKIj71TZoiYxTw,1452 +torch/include/ATen/ops/upsample_bilinear2d_backward_ops.h,sha256=ZSZmqVZdp2Fd41msrdIc_U8DQanU0rHyQcIymzc3Qzk,3043 +torch/include/ATen/ops/upsample_bilinear2d_compositeexplicitautograd_dispatch.h,sha256=AFonp7HYJdKstea9nOO5dp9aued3qjazrSJbOAhhjp0,1794 +torch/include/ATen/ops/upsample_bilinear2d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=Wb9fcvUUqRmxlwgQMtw-z_ON9VWuaQCcPKh_hm74q-Y,1429 +torch/include/ATen/ops/upsample_bilinear2d_compositeimplicitautograd_dispatch.h,sha256=IGmpa4Y6UZTWjvDszfzLQPliXi7zCAzSxkyFcR9Pqf8,1330 +torch/include/ATen/ops/upsample_bilinear2d_cpu_dispatch.h,sha256=kC3Q4xHH3X59GQfKMrVja2E-LnkKmx2PFyZn6Kd4r7Q,2275 +torch/include/ATen/ops/upsample_bilinear2d_cuda_dispatch.h,sha256=MdoExLBsocpwAYjWpXFecH2ZGprjNASpl1ccclWOEt8,2277 +torch/include/ATen/ops/upsample_bilinear2d_meta.h,sha256=jgHIEKvRT6etXOtk7OTvnnB3qUULzJMfgHKJ2bu2bjI,957 +torch/include/ATen/ops/upsample_bilinear2d_meta_dispatch.h,sha256=Py2wNqAxmOarNBPOUk4W818VNqMl8FzTrIQ4rN_M1PU,2277 +torch/include/ATen/ops/upsample_bilinear2d_native.h,sha256=hKGce_EcyoFRXyhw9gz0UDiBpFV2csYx11n6FzQHlMw,1933 +torch/include/ATen/ops/upsample_bilinear2d_ops.h,sha256=cREMKi70-JOeYBAFCjZ1gk5BBXL3f_jAMys_Rxy_4LA,4620 +torch/include/ATen/ops/upsample_linear1d.h,sha256=XMfYzr3D0mfrne1CaQwm2A1II08h7qsfRtbOYaklu88,7311 +torch/include/ATen/ops/upsample_linear1d_backward.h,sha256=kwy5bGep8i7ziVowLguD8P4GigD-Xcl3m6koFjv8E6Y,7067 +torch/include/ATen/ops/upsample_linear1d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=lBSE3kmc4BB7CaxFf3MOqmSUMYlHdnfHg698xqP6OwE,1415 +torch/include/ATen/ops/upsample_linear1d_backward_cpu_dispatch.h,sha256=2fB5y1fQw9ACU-SNsur0W-94i_YqVd5T_r-_Afd4Z8Q,2291 +torch/include/ATen/ops/upsample_linear1d_backward_cuda_dispatch.h,sha256=kere2JAlUoC7O7nq92I9qGSlJbtlOky-L6Y9UjbYaTU,2293 +torch/include/ATen/ops/upsample_linear1d_backward_meta.h,sha256=oRXml8iFunMdUKG369sS5DKGmxhtvGndfqhGwkyL3Hw,969 +torch/include/ATen/ops/upsample_linear1d_backward_meta_dispatch.h,sha256=NG-EBlKmzPyIL8wGGoA74jQ3Uypuu0FL3pPvZTGbmPc,2293 +torch/include/ATen/ops/upsample_linear1d_backward_native.h,sha256=vKSHMAWZ6kus7Co20yq9kunrtMQ8X4nW6elQp-D_MIk,1370 +torch/include/ATen/ops/upsample_linear1d_backward_ops.h,sha256=ATodPxp3U30nOJa2tu9j0uI47vxEh8WhnL7WulKxcak,2789 +torch/include/ATen/ops/upsample_linear1d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=ivwdDgfcGzY1iFx_Pp4G-OIh3g2pAtPLzMMJrsGUfNg,1323 +torch/include/ATen/ops/upsample_linear1d_compositeimplicitautograd_dispatch.h,sha256=yI8L8BBuy_JfEo6UVla9bGV8sfvkIqsaEHGWhB-61dY,1326 +torch/include/ATen/ops/upsample_linear1d_cpu_dispatch.h,sha256=pwsgd1HnHL8zQjq31vcQ5j1VDngHavSVIiOucldbXnY,1987 +torch/include/ATen/ops/upsample_linear1d_cuda_dispatch.h,sha256=dRAjjgUEanXTessojXyjjm4jCA2Cl5vp2OeCRa-Bp2I,1989 +torch/include/ATen/ops/upsample_linear1d_meta.h,sha256=1kWxbA_DjYmeRVPQwmiNoYZdvLVsFNpjMKeEa0KZhos,919 +torch/include/ATen/ops/upsample_linear1d_meta_dispatch.h,sha256=dLS_D0AmqPfBU44_nAd1LYMH9e5fUp6JZgsbiDF9D6U,1989 +torch/include/ATen/ops/upsample_linear1d_native.h,sha256=tT4KwDEmOa5zZpgud45S_pUeU6Ptk-gTcJawtNXR564,1405 +torch/include/ATen/ops/upsample_linear1d_ops.h,sha256=wOXKdObt3qKdzKRIakouFHBD4Z0e7P5l8NCXGyDOF6g,3350 +torch/include/ATen/ops/upsample_nearest1d.h,sha256=jBdVNdzc--py5Z2lGXWt-P-J9YIW7hIhVKDu_crTAhY,6632 +torch/include/ATen/ops/upsample_nearest1d_backward.h,sha256=OEyC1qpBn47Gp95vs7cOILx1IlbBK_tCIccXfBeScr0,6558 +torch/include/ATen/ops/upsample_nearest1d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=mdHB-STpHD5mkRhMpFCMRb6xJMwOPal8eI7d9ZriKK8,1377 +torch/include/ATen/ops/upsample_nearest1d_backward_cpu_dispatch.h,sha256=fcRZb_Fju5v35-UqY7TIc9SMYkw4ljXieLLN-hHrwiA,2177 +torch/include/ATen/ops/upsample_nearest1d_backward_cuda_dispatch.h,sha256=5NDlCuj0EwWJ8UCM-IUypq7-7chnqD5TOwju3_GAKF0,2179 +torch/include/ATen/ops/upsample_nearest1d_backward_meta.h,sha256=U2PLncfKNwHzR0oLm2UDjnCrQu7m6mv1q2AwANDuYXE,950 +torch/include/ATen/ops/upsample_nearest1d_backward_meta_dispatch.h,sha256=dMH9nFTQtTzQLwkesGH-KEsa4SrlmZ_UddUXpBYBP_c,2179 +torch/include/ATen/ops/upsample_nearest1d_backward_native.h,sha256=Aql3slzmg9XS5WdNR0euFGdEVxkNOWhTjKgWePBTW50,1335 +torch/include/ATen/ops/upsample_nearest1d_backward_ops.h,sha256=jsS-FUQ4BmB_fpbyJKLg_ClJ__JBhBEVmnfoBgSJXfM,2663 +torch/include/ATen/ops/upsample_nearest1d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=fGFQAHPhtd06nLi8_fZ5Bq58a47af-q9BEKFtpB1w8w,1285 +torch/include/ATen/ops/upsample_nearest1d_compositeimplicitautograd_dispatch.h,sha256=sumWZGvf75FfJIxi_npiYxXr-SWbK_sXSAUkY4WsVXE,1288 +torch/include/ATen/ops/upsample_nearest1d_cpu_dispatch.h,sha256=w0GpkzU97FINPKcdDjfRXCIpJstKN-pG7c0h_KGPFiM,1873 +torch/include/ATen/ops/upsample_nearest1d_cuda_dispatch.h,sha256=3YSirPZNlB9W4lXB-0iaoZE9ibiTzU-8fsP_lmmrhtI,1875 +torch/include/ATen/ops/upsample_nearest1d_meta.h,sha256=4BPjT8S8sENwPkHN-FYkpDrZEpUwHzJU1Y5eUPs_HNQ,900 +torch/include/ATen/ops/upsample_nearest1d_meta_dispatch.h,sha256=eZAeY7L0vrCO4Pg4TtbZCwrz8L_cWqX6VxqUQt-Tw08,1875 +torch/include/ATen/ops/upsample_nearest1d_native.h,sha256=7Usz2bcZf-ZO9wADjpkIEw5jQQOTk-y_AV2PLh7CwU0,1351 +torch/include/ATen/ops/upsample_nearest1d_ops.h,sha256=1mIGCtCOKUHJavVYyIc2Gx_Vlo-RnrPnc5EZtzX95H0,3161 +torch/include/ATen/ops/upsample_nearest2d.h,sha256=O_5ffpixJga9VIeobNTWWdgiKoF5thaefVCBkvPwxRM,11006 +torch/include/ATen/ops/upsample_nearest2d_backward.h,sha256=udedsCys3u4JReu1l9pqCImHk3HH-obKMEfm6aGzUk4,7398 +torch/include/ATen/ops/upsample_nearest2d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=1MzQcxYgVgdMzf3YAJpyY67spU0xOCyhG-tW3JEqoGE,1479 +torch/include/ATen/ops/upsample_nearest2d_backward_cpu_dispatch.h,sha256=GlbG2EorM5DXATkDUSRQVf0JKKFga_d_rKGojVaDKq8,2453 +torch/include/ATen/ops/upsample_nearest2d_backward_cuda_dispatch.h,sha256=HpAVwdP44f7xWDJOKPxjxVRu-sww8w4L3UnTb0bLUJ0,2455 +torch/include/ATen/ops/upsample_nearest2d_backward_meta.h,sha256=zyGFvrmYXIU2sd5LW9nlg6Pw0Rgo-Lw3seMpjUh9MZo,986 +torch/include/ATen/ops/upsample_nearest2d_backward_meta_dispatch.h,sha256=u_hHiKGlewIkNWzp5GPx9tmnD_O0QQKBzTwApN31ZYo,2455 +torch/include/ATen/ops/upsample_nearest2d_backward_native.h,sha256=pyprzLb6jzVT8DaBvFXab0I__HFVmRfRmzAQCycjNqI,1407 +torch/include/ATen/ops/upsample_nearest2d_backward_ops.h,sha256=yElwaXsJSwUpF-IcYvIljlE10-2_ng4ucZRoGrpVjvM,2905 +torch/include/ATen/ops/upsample_nearest2d_compositeexplicitautograd_dispatch.h,sha256=aFScUxyx5xG6b-zDAf5yIBkQ_ufI2GwOUos4yy616DI,1710 +torch/include/ATen/ops/upsample_nearest2d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=hziL-cvrzOJ6PF8iAzrkr66Eil6xVh2pcSzGuzgzgXk,1387 +torch/include/ATen/ops/upsample_nearest2d_compositeimplicitautograd_dispatch.h,sha256=r_Lh330YzJ2sz137JqJ3_OaELvP31_i59R4_SdGaMMU,1288 +torch/include/ATen/ops/upsample_nearest2d_cpu_dispatch.h,sha256=KLrwfa_aQDjnGJ9xBZueZRDXRrT9Q6Zow-rlvwygA5A,2149 +torch/include/ATen/ops/upsample_nearest2d_cuda_dispatch.h,sha256=uj1k2MBSila5Jc5rIXAvskiD5_K8Nk2e_KNt_QfCPTQ,2151 +torch/include/ATen/ops/upsample_nearest2d_meta.h,sha256=MjEloyM-rcqtkXZ-b1IijXxoEsk7WJnRgEbqWDtwkhM,936 +torch/include/ATen/ops/upsample_nearest2d_meta_dispatch.h,sha256=XKMsaV-fW3k6EykUytRxEDS3Z1X69tZQEuNNOxdM5o0,2151 +torch/include/ATen/ops/upsample_nearest2d_native.h,sha256=iFcnobq73QduuYOSFsj-hRwA_7MzLDN1hASdgH5MigA,1825 +torch/include/ATen/ops/upsample_nearest2d_ops.h,sha256=NX_114ByvM5Iu1Zx3GNfolzNBIiTz5LYd-cOKosTreY,4344 +torch/include/ATen/ops/upsample_nearest3d.h,sha256=QmD2cR5kUTXnmT7a3iTlQwnLOehttMBxX-MlTXwHuts,8252 +torch/include/ATen/ops/upsample_nearest3d_backward.h,sha256=txFD-Av9RCMbP0tgZmPkRLKkcWytQdv_Bqp8SYy_sJY,8178 +torch/include/ATen/ops/upsample_nearest3d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=lzmFuH3ecfnpV7CXnOOsFHoSWhJz0VrH7TuERSz7TCg,1577 +torch/include/ATen/ops/upsample_nearest3d_backward_cpu_dispatch.h,sha256=EqnWiEwPQUa5TuP1DPJwb9IUXxOYGbRYo8xj3K24Z3M,2717 +torch/include/ATen/ops/upsample_nearest3d_backward_cuda_dispatch.h,sha256=AJudIp8RCyXaePWGfY9FJcrmJnsxuZfGLwnh3BKVaC0,2719 +torch/include/ATen/ops/upsample_nearest3d_backward_meta.h,sha256=Rg7bRlZPktxhgg6IXQTXRu4Jp8VDfAtEFC9tQ-D8ikg,1020 +torch/include/ATen/ops/upsample_nearest3d_backward_meta_dispatch.h,sha256=jODEUU2iUNxrZHZj26sZ5FvXTSaBCvX9ESV0GKkbxc4,2719 +torch/include/ATen/ops/upsample_nearest3d_backward_native.h,sha256=ZSnOPvFhS3Y8AFObBHa6YD3_bLMFvK6H2hQKwU60akg,1475 +torch/include/ATen/ops/upsample_nearest3d_backward_ops.h,sha256=ITV4Vn-y5-0EarAIpEFD5N3YeD0InGtWNloqG6NVqeQ,3135 +torch/include/ATen/ops/upsample_nearest3d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=1qJ4luQVnGE6NRtOrz5_OvpdVBmdLwCbBZEulOiRF0k,1485 +torch/include/ATen/ops/upsample_nearest3d_compositeimplicitautograd_dispatch.h,sha256=us4SZKx54BNw8qn5HVT9715KEPu0mK2HSyg4Lzg7m7M,1288 +torch/include/ATen/ops/upsample_nearest3d_cpu_dispatch.h,sha256=6HPW0Dkw3gUPtuUUNGFRTJ94rrc8bZYLGEOplBu77tc,2413 +torch/include/ATen/ops/upsample_nearest3d_cuda_dispatch.h,sha256=Diy1FaKuyB3s9OFDay_lNb8YIgRdItKN7A85oNybFYw,2415 +torch/include/ATen/ops/upsample_nearest3d_meta.h,sha256=mESNA1u323tXos4lHI0Xb-QJe93khPp8rvKnyK1R99c,970 +torch/include/ATen/ops/upsample_nearest3d_meta_dispatch.h,sha256=hIHn0s-OaLYF43OpCBq2t-18wxgNaZjMo70pOJUlkfA,2415 +torch/include/ATen/ops/upsample_nearest3d_native.h,sha256=eiQEhIMx8axNFVzjzmmlhy6becX4VzrDGfwhsnsV1jw,1747 +torch/include/ATen/ops/upsample_nearest3d_ops.h,sha256=2EHfxmEjRNwLffTVHyq6LCkN3OgU_Ww7Xw_DqVQa81I,3633 +torch/include/ATen/ops/upsample_trilinear3d.h,sha256=3W9E6kvtVPHhfnVXQZ9lW3hDgvba3wyjdunYW2XjcZg,9054 +torch/include/ATen/ops/upsample_trilinear3d_backward.h,sha256=SqPf9FNBGRBAk-Zp7-fqkcv8DGnvfD3zfC9DNJQgoms,8780 +torch/include/ATen/ops/upsample_trilinear3d_backward_compositeexplicitautogradnonfunctional_dispatch.h,sha256=AP2kXqEo9wi_XHOFXxzuRfoCQA9pf6gaj88T8-5Z4-c,1621 +torch/include/ATen/ops/upsample_trilinear3d_backward_cpu_dispatch.h,sha256=s0iSQflxCP1hF5pjgR_fDo4LjupLwO2m13LMnPIdHP0,2849 +torch/include/ATen/ops/upsample_trilinear3d_backward_cuda_dispatch.h,sha256=GkTpTsqlRe6z51pXzsRpyztFTRE_71TVDVEN4Jz2WEE,2851 +torch/include/ATen/ops/upsample_trilinear3d_backward_meta.h,sha256=TKGxdkvWo-MGz0oYBtu9NzRbiM0xhJ7oEBApujHZIxI,1042 +torch/include/ATen/ops/upsample_trilinear3d_backward_meta_dispatch.h,sha256=kc9pT8jE0XxHPZELuAaNLLCnopoUV76uD_07KnrzTsA,2851 +torch/include/ATen/ops/upsample_trilinear3d_backward_native.h,sha256=DsDe4hdKCvN2mwABULGcANebhdbb8qK-f1_P3VxFSn0,1525 +torch/include/ATen/ops/upsample_trilinear3d_backward_ops.h,sha256=TRWVmmPmMxw39RsgVgnuiy3SOkrR082bo4Tpa0cMPJ0,3279 +torch/include/ATen/ops/upsample_trilinear3d_compositeexplicitautogradnonfunctional_dispatch.h,sha256=8hW2jm5ozyz78lqhe543EgtIDFRDUx1cOtgRQvKR5ls,1529 +torch/include/ATen/ops/upsample_trilinear3d_compositeimplicitautograd_dispatch.h,sha256=zreAYtRkrA6BjIBcrDqzedxVegRyLV0QMvM6DSeDG1M,1332 +torch/include/ATen/ops/upsample_trilinear3d_cpu_dispatch.h,sha256=zu6J7ocKvCaw8UsQY9T4iu2JwEIsjs3do0_aB1eRpSM,2545 +torch/include/ATen/ops/upsample_trilinear3d_cuda_dispatch.h,sha256=izNpTuxeT3tjQeY6TjjbdVEtxBk0kGyrZMLi8f6YuZM,2547 +torch/include/ATen/ops/upsample_trilinear3d_meta.h,sha256=V2VD_cRq_XmREMLDYHb0XIl9aTGMPLgoUiLqnd_heIE,992 +torch/include/ATen/ops/upsample_trilinear3d_meta_dispatch.h,sha256=30aaUlfLBQZMJ0D7beom4Y1JryH2vLnjikuuf1leQ2I,2547 +torch/include/ATen/ops/upsample_trilinear3d_native.h,sha256=In-rSLbiKtoo5vAbaD-9Wf9abS826W5354l6X6WG0AY,1563 +torch/include/ATen/ops/upsample_trilinear3d_ops.h,sha256=TuDc7fNnX1B6qxAUrKbMxMKqoJo2DTO7rgCYXOi4_9Q,3849 +torch/include/ATen/ops/value_selecting_reduction_backward.h,sha256=50iUpIEwpCTGhy81ykc_3hqc-t6sy6LQXxjyOFbacCI,2335 +torch/include/ATen/ops/value_selecting_reduction_backward_compositeimplicitautograd_dispatch.h,sha256=XZOBQxgOmxOTGbXly8T2DYQTmU5aWD2bc0dFmehGHso,1295 +torch/include/ATen/ops/value_selecting_reduction_backward_native.h,sha256=mVQH2Irl_BLOPJLOCx63QNTNSgI7HYlkPem9DN9sxf0,1025 +torch/include/ATen/ops/value_selecting_reduction_backward_ops.h,sha256=DL7brvLTfri9Paf5ZUfCIaFotzYcaZ2GVwSPFiwH120,1570 +torch/include/ATen/ops/values.h,sha256=aHiEbW13ImT9Kxo2-9wG7ERpwYIQFXHmBIfnIXuHg3g,757 +torch/include/ATen/ops/values_compositeexplicitautograd_dispatch.h,sha256=NfxBe-IKRKHgsn-e6BK3h_gcKPJOrwikkx6sZQkNeso,1018 +torch/include/ATen/ops/values_copy.h,sha256=8ScsN7bZKLseag-JbeMqiQ_ZXR3f1Uau-sQJUX-Z_dg,1331 +torch/include/ATen/ops/values_copy_compositeexplicitautograd_dispatch.h,sha256=-_kQv813ADijVdRCK4oDIlLCs-p4lap3kAxzIGtywLY,1131 +torch/include/ATen/ops/values_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=G7qH4UrcePdCuKcFyg-2Rp0jz2RvZlc9a6EIudFUVOA,1049 +torch/include/ATen/ops/values_copy_native.h,sha256=w9zwAvkNe16eAq1rdQQx1r4uEFC9SVh8OxvYp3CB_yw,818 +torch/include/ATen/ops/values_copy_ops.h,sha256=j_LgxXl5pGa1ZN1KJhogl4nufnr9LzGWL3MrYZf1pHA,1831 +torch/include/ATen/ops/values_native.h,sha256=XL8VFOu16JXk1Cr-GXJSLcVqALOOnizI8baLftH-pZw,925 +torch/include/ATen/ops/values_ops.h,sha256=K0NvQylY7UsEzgdCJWriEkc8UZbT1GPkilN1WLG8JzM,1217 +torch/include/ATen/ops/vander.h,sha256=nHXvQ6uH_FJh9g4JeV4nEBseWwJvT3cvX2e4PGFCYWA,999 +torch/include/ATen/ops/vander_compositeimplicitautograd_dispatch.h,sha256=uXoUUm1vx2D0GgP6KAldUXrrIrtjwrxyJh5SYdNzdLo,1081 +torch/include/ATen/ops/vander_native.h,sha256=btiDecUDd6ABR57r8V5PcmQ9RVdY-BHQnBxDbixdyQg,793 +torch/include/ATen/ops/vander_ops.h,sha256=c4C80epl2oE36jNNfI2WJxpztELmNLlVjvJuoN5MHxQ,1360 +torch/include/ATen/ops/var.h,sha256=8TgH44ilLXBzoGuP-jK14DCZ7C1AJ_r7OfrhYFUaEn0,5125 +torch/include/ATen/ops/var_compositeimplicitautograd_dispatch.h,sha256=lneD8C9FsxcAhXtZGSQp3eXZkEAQrvlciZuPlYB9PuU,2275 +torch/include/ATen/ops/var_cpu_dispatch.h,sha256=xbTiFw-dPZ7yOjUcMeeyw9_moM-xRlVVjJ873Nqd1xs,1467 +torch/include/ATen/ops/var_cuda_dispatch.h,sha256=uzyd8yLSrd0yP0b1MHAUSqV2sMES69nzvHLQ9IPW2Fw,1469 +torch/include/ATen/ops/var_mean.h,sha256=tHOPj3nbkkxrrtNjFPUIREWcI2TrbsEIUrm_POi6aIA,3471 +torch/include/ATen/ops/var_mean_compositeexplicitautograd_dispatch.h,sha256=nJCMUsGiTT1vPPZY1S4j2FbynhrMKzGwPb2r4f3Wth8,1437 +torch/include/ATen/ops/var_mean_compositeimplicitautograd_dispatch.h,sha256=ZydlSZl7_itNZD4KsOSy9Ht391qMQd9415IdI--clWY,1527 +torch/include/ATen/ops/var_mean_cpu_dispatch.h,sha256=vYztBG8dlhkkV9Qm0sHdfKesjSBeUwLlAyXWcQ1sR3I,1128 +torch/include/ATen/ops/var_mean_cuda_dispatch.h,sha256=A_4opqHf8q6ilPM5jdkmV8hiQY_UDHbQAZ-t6JFFCmI,1130 +torch/include/ATen/ops/var_mean_native.h,sha256=SqDxuvF1Avwgp6uhfdNxFFOSSvonnYEN5-kzOju0d-s,1691 +torch/include/ATen/ops/var_mean_ops.h,sha256=iMjv5G18ETat0FsLdKBK8ag75tLdnC__ROYIVaN5zVE,5950 +torch/include/ATen/ops/var_native.h,sha256=HjJ-lf9zLKTWY9mjMs-MKNXUASmz5aqTHcOzXwXRNek,1894 +torch/include/ATen/ops/var_ops.h,sha256=q4GJL23srEMokn0Ey5xu8Wq66IIrBgRPPU4g4i03xeI,7758 +torch/include/ATen/ops/vdot.h,sha256=JicLDIkN5IfnoN87enhlobiJviGMJe6C5n1G2wd37UM,1402 +torch/include/ATen/ops/vdot_compositeexplicitautograd_dispatch.h,sha256=l6lc0wGToDZwZPEiSl20UKMYJGxcJi-gUt_xYnHsYjs,1169 +torch/include/ATen/ops/vdot_cpu_dispatch.h,sha256=YHIEGfbDjVlUP_jU-rTJnkolITQ-hkjz_CP1nZDn5JU,998 +torch/include/ATen/ops/vdot_cuda_dispatch.h,sha256=GaZyEFeqgMVWQocjbsYrjrwnfmsi8gdy18XQPDWpcGs,1000 +torch/include/ATen/ops/vdot_native.h,sha256=bxNBtL1qS4gZUG34QntoryniL9Zu2gDzoXQ1LrAhsPA,939 +torch/include/ATen/ops/vdot_ops.h,sha256=9T724eQRJCw7rQYrYOL8ODZ8qU7Ye8qZXgw2HdqdKr8,1961 +torch/include/ATen/ops/view.h,sha256=GGtONjhHkbAAZoG7YkCTLRMSjKO4PJHnyOu4SxGspo8,1226 +torch/include/ATen/ops/view_as.h,sha256=xDxUn24gipOhvcTNyePXEVhT0zQlx7dFF3IySwo45eM,758 +torch/include/ATen/ops/view_as_complex.h,sha256=u7picSMflYrIWKCWQ814OjXBYBFW55tjDY6mHcGzNAk,933 +torch/include/ATen/ops/view_as_complex_copy.h,sha256=1NDBuYRizSL7R0wpMrZehqqnMhwRSID6-I8wdTXQyTs,1421 +torch/include/ATen/ops/view_as_complex_copy_compositeexplicitautograd_dispatch.h,sha256=7YwDxhmZ4enLPGl5UONKstWHBRVHXAv2UxJ5FFcVang,1149 +torch/include/ATen/ops/view_as_complex_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=nIIJh5cxokGirJpguEn88TDnecBG1HxI9RLdSjleJC4,1058 +torch/include/ATen/ops/view_as_complex_copy_native.h,sha256=t_0Bq63ZfO63adSp1mNQ6makFBr2UGGZe4ZKJ8yjG7g,836 +torch/include/ATen/ops/view_as_complex_copy_ops.h,sha256=0rLVSPL32E60Qcjpu-38TfnY_-Jgh7VdxSMxA3Yrt2g,1885 +torch/include/ATen/ops/view_as_complex_cpu_dispatch.h,sha256=3jMO8MKkA30YDaXsRtIsJ7bLvD2hnABjOeCQUUKWpog,983 +torch/include/ATen/ops/view_as_complex_cuda_dispatch.h,sha256=yAopiXBJRt86Ot8cCjrxSTkXRphkp9rQDoe4qQ-i3CU,985 +torch/include/ATen/ops/view_as_complex_meta_dispatch.h,sha256=bnfEFb_v0CGfAxp7lS-4B454LowAle6x6XMdZDk61cI,985 +torch/include/ATen/ops/view_as_complex_native.h,sha256=m_crbI2Lcc2yoM8S4QWDHmZeUyq3clI7S4986T1N61I,809 +torch/include/ATen/ops/view_as_complex_ops.h,sha256=9R6piRXR4G7RDALjgA2UWTU_6wyLo_fwGxmo8V96h9M,1244 +torch/include/ATen/ops/view_as_compositeimplicitautograd_dispatch.h,sha256=C5cdp5DbhwwtcNkmqpICCGuT646zDlRD03zgeIPE8WY,1045 +torch/include/ATen/ops/view_as_native.h,sha256=_wWLqJC__0CWRID3JDYSnqRHczhbQauaX2fDwfMgp-k,757 +torch/include/ATen/ops/view_as_ops.h,sha256=zpiLCubplyWHWYexQuGPI8MyzJZZ2mKkOHm4fw_X02U,1306 +torch/include/ATen/ops/view_as_real.h,sha256=5ayVHt5QlvTfSHaJ1Hz8-4dmx-M_VKWim-WVyA1F9yc,921 +torch/include/ATen/ops/view_as_real_copy.h,sha256=7IQEl-c6a8QJ1wN0jaSyYKcSSGU1-ySygrXWRiyIlNQ,1391 +torch/include/ATen/ops/view_as_real_copy_compositeexplicitautograd_dispatch.h,sha256=7AyOa60Jva43fN_S0GYrJoLyIJt4H0i6bQqpweM67tM,1143 +torch/include/ATen/ops/view_as_real_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=pxYeLcSTKRfJIz1bUIy2ppanoNUMT-Hd2ZVdZ1v0TDo,1055 +torch/include/ATen/ops/view_as_real_copy_native.h,sha256=0WF3X8s2Hqdr_twooGlPfpzwRCMpc6IFsH0YJcXXjVo,830 +torch/include/ATen/ops/view_as_real_copy_ops.h,sha256=1QEGWkzDQPwNNeSuezLkRT278PeaCCkA2bKkVcMx9xg,1867 +torch/include/ATen/ops/view_as_real_cpu_dispatch.h,sha256=IIDz6-o9AGK1JjNbVtGvYxbGGP0cFJ8paH12DRBXxFo,980 +torch/include/ATen/ops/view_as_real_cuda_dispatch.h,sha256=a8jBi4gQ5a-fkvIIVCOLp-vJ4CwbwgZBkX_ZBrgseJs,982 +torch/include/ATen/ops/view_as_real_meta_dispatch.h,sha256=jLFinJQF4v1-HgEsoVUN4-iSY3tTjbHlUG4PMjPtfyw,982 +torch/include/ATen/ops/view_as_real_native.h,sha256=rC2H9uYQki8uZdDLuSdBQXUFoMtsrzyxlx76QQuTWEU,803 +torch/include/ATen/ops/view_as_real_ops.h,sha256=ea3qjXjO9-xonKQR3sG6-kZ-gTzrZxBBLvxmZvFCefg,1235 +torch/include/ATen/ops/view_compositeexplicitautograd_dispatch.h,sha256=YfVf_zRhMJVZwtdi0360E_lsDJZmo0C9jqRX7Ayfq6U,1038 +torch/include/ATen/ops/view_copy.h,sha256=peQjHdsgZV7JG_uG6MABHdu6S-TxxOFhTbD0cxrezYs,4567 +torch/include/ATen/ops/view_copy_compositeexplicitautograd_dispatch.h,sha256=aYEAnUUQxskAc4fziN8xtieCgMmVXicygQUn2l1aWr0,1607 +torch/include/ATen/ops/view_copy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=MvDZROlbaG5ER5LTuq5REUHVGm8kBwwxLP2FVKYX4vU,1238 +torch/include/ATen/ops/view_copy_native.h,sha256=YiUsJYzg0-BgZw7QmPTc7wrqeZJalTBLkngAk_jkVsc,1074 +torch/include/ATen/ops/view_copy_ops.h,sha256=Y2jtVsII1DLIQn31Bcdm57xlnlJOVb4nIoM7-GkwUCQ,3294 +torch/include/ATen/ops/view_cpu_dispatch.h,sha256=n0nHJcsPfSrGba8oDYuE8b93v5BumHDMVt6A-N_enr0,1079 +torch/include/ATen/ops/view_cuda_dispatch.h,sha256=2u7dNdDKeuSEg4aeLaSOOYvfc40dCXn4MINdLMuyLUg,1081 +torch/include/ATen/ops/view_meta_dispatch.h,sha256=vAOyA1AJ5XD0tXbygiTYrRa-uEJX35_SACr1gNTx4lg,1081 +torch/include/ATen/ops/view_native.h,sha256=Cs4X4WNX_eatnOHi2Hhfjmm3Ko6yQLFJcJuNXDq9N4k,992 +torch/include/ATen/ops/view_ops.h,sha256=SwwCiaWM53F4d3WML2FpoU9FK4ZupDcCcWYJkC-eP9g,1894 +torch/include/ATen/ops/vsplit.h,sha256=KM71VxI6WvjpHrNSaX6Otlz5tKSh3YyP5tvdKBySMMQ,1193 +torch/include/ATen/ops/vsplit_compositeimplicitautograd_dispatch.h,sha256=t52jUPFnDHNxt8JdBEq-3SSE1O9tGfzsOczonUdybOE,1145 +torch/include/ATen/ops/vsplit_native.h,sha256=OIWh43o4DGB9m7Y4bT4EcE6T8vsSo8QV1MdA9RGIUXo,857 +torch/include/ATen/ops/vsplit_ops.h,sha256=u1eSYptxW_q-ruySzD6rhOKni5O-9Kzuv_I5r2tworM,1996 +torch/include/ATen/ops/vstack.h,sha256=9-oSnsmcpHnF7OQmLdMbNGUw1F-DFWDS0MwG49lVZMk,1302 +torch/include/ATen/ops/vstack_compositeimplicitautograd_dispatch.h,sha256=emO4YscThcSAWMKGxGuJosB94u8muTcTjLW98iK6xpY,1172 +torch/include/ATen/ops/vstack_native.h,sha256=me3PDh8hR7Z7IBZcjd2zWU_5Q4leydqeEAY4ewtC16s,806 +torch/include/ATen/ops/vstack_ops.h,sha256=WzGCcbXgd_UTVjQZf5ozkAS13a0bTNWfNict4SXTI-0,1799 +torch/include/ATen/ops/where.h,sha256=e5IQB-LZoBA2mv24F8lCpdGtg5mZLp5IN-75vqLyhS0,2557 +torch/include/ATen/ops/where_compositeimplicitautograd_dispatch.h,sha256=peRwiVwihUlDxQwojEO1mXQO7Ha2XaIaAQ9wqwzcu3w,1364 +torch/include/ATen/ops/where_cpu_dispatch.h,sha256=S0qQQ16Fii0CA3i09OAUYWAY9hrRcudP9OGJV5WWgjQ,1296 +torch/include/ATen/ops/where_cuda_dispatch.h,sha256=4EoxCwDtpsOS4qD5vqgW0mB_wPlZOWrwdG3MMCrx-Sk,1298 +torch/include/ATen/ops/where_native.h,sha256=EdIsoAf_mD76R_OaeXYQ-hzkFPS5RckLT9hvMXmjcLE,1591 +torch/include/ATen/ops/where_ops.h,sha256=4UmqmbPgxA9Ow7EEpldJUbrUgm-JowsOwhAy-K6DdPk,4881 +torch/include/ATen/ops/xlogy.h,sha256=TuCqeYsAVCOenJnWAo8TaphxJlOMowdc7RhMNiOCUug,3318 +torch/include/ATen/ops/xlogy_compositeexplicitautograd_dispatch.h,sha256=iVS_KwkXYj3d4WVWMOLPWnAEWouswJnKsjq736Ito-I,1612 +torch/include/ATen/ops/xlogy_compositeexplicitautogradnonfunctional_dispatch.h,sha256=893tL9CGZya8TpO9sRY0zYtCaipyLDE7CIc_Xpf5C1Q,1145 +torch/include/ATen/ops/xlogy_cpu_dispatch.h,sha256=rLLrWLK8GlMaWaFOjXxFMfLPgxyi3Y8fmDsg-5ApNy8,1282 +torch/include/ATen/ops/xlogy_cuda_dispatch.h,sha256=WbCQeTbvAWTPrGUZ4ahFEM8bjrc9ETJJAqslOU6ADAI,1284 +torch/include/ATen/ops/xlogy_meta.h,sha256=XwY96Fur0W4J-FkBBzSLX7vkPd2hIn0CSfoSxZSA1l8,853 +torch/include/ATen/ops/xlogy_meta_dispatch.h,sha256=XfwHSu5uRYR8XerF7Pj7jia6SWcP7ILRR_82gZC5PoU,1284 +torch/include/ATen/ops/xlogy_native.h,sha256=dxKz3BI0ip-r2VMHmdv4Z5M1dT5a1UvRzEHrNL3GQk8,1320 +torch/include/ATen/ops/xlogy_ops.h,sha256=8qWfUHgkOfm_ltEQLeT0mFjr-fyY9nc2O6raghk8g1I,5881 +torch/include/ATen/ops/xor.h,sha256=UIoqw1zoIaprTXMXWtpL9qk0Zs6lERKjzaPsosIrUKM,1151 +torch/include/ATen/ops/xor_compositeimplicitautograd_dispatch.h,sha256=8R1svVhfc2qFTWM-Fj8wooTqKBD6pYAJXgyS_g73A74,1282 +torch/include/ATen/ops/xor_native.h,sha256=gM5dHlixX5ocP3UuTv8teVyONjrhFj8lkBciivnrQ7s,994 +torch/include/ATen/ops/xor_ops.h,sha256=KNyI_wAKSbxhcl7EmkaPR65P_u2wSNxFq2J2ic-Sx6s,3145 +torch/include/ATen/ops/zero.h,sha256=l1nRjtPSq1AQr6AEwYAKTUeywvg004HZYB-NT3LjQ6c,1397 +torch/include/ATen/ops/zero_compositeexplicitautograd_dispatch.h,sha256=RwOxpmH7vFAarQXVZkuYdBdybZy_YhEfUwdLeSwmlnk,1169 +torch/include/ATen/ops/zero_cpu_dispatch.h,sha256=--hGYlPpPqTv9nVIBWSLlCq1eY4F17dhTx6e-vUKy3c,969 +torch/include/ATen/ops/zero_cuda_dispatch.h,sha256=5yduzMQuacGwEV7MWqKhXMWXF2gjbG3FkvNFS-3Y7Gk,971 +torch/include/ATen/ops/zero_meta_dispatch.h,sha256=zj1PwiKraybZ-vSRnryY-eXt37kxO9M8dtG5HaQW2w4,971 +torch/include/ATen/ops/zero_native.h,sha256=zIxO5KqSUxuJcMoAC6Yrw_prZsmxTz6NJGb30lMjmnw,1135 +torch/include/ATen/ops/zero_ops.h,sha256=SH7IZbIrl1-rELgxW4Rl2LIg23hSNCmLMJwGHwy2gL8,2282 +torch/include/ATen/ops/zeros.h,sha256=28Fpk86-P5URArDY3iPweMXVtcc0iJjqRwR1Amr70Ik,7136 +torch/include/ATen/ops/zeros_compositeexplicitautograd_dispatch.h,sha256=flCpRhzYMiT2a3IDtKWxyh4i3954alOAmf0TBqWc6wQ,2440 +torch/include/ATen/ops/zeros_like.h,sha256=jdw3IDyDg_k5_j_t9agBCmStJpL0bfP-PlvJhxqCbpI,2475 +torch/include/ATen/ops/zeros_like_compositeexplicitautograd_dispatch.h,sha256=Zef8fFttnNwB08ztQGWxwQ149IgEewYmpKtgLhLYyu8,1646 +torch/include/ATen/ops/zeros_like_compositeimplicitautogradnestedtensor_dispatch.h,sha256=TmVm2MpPmpHuSrfeqTFuPJ5aVEBnQYMgQauXMJSnbcw,1392 +torch/include/ATen/ops/zeros_like_native.h,sha256=jUbqnlj7UU-v1U3lWjQSFxciVMD7uo-YWOL3JLHcnL8,1086 +torch/include/ATen/ops/zeros_like_ops.h,sha256=FkMXkSDU7uv1jAT_teyuPyr5S1GMgbLVkIkp_i3wSFM,2655 +torch/include/ATen/ops/zeros_native.h,sha256=hGGjFyBZargs2pXY_DrBUfIOehVEMHS_7Xcm9uTd5Vk,1416 +torch/include/ATen/ops/zeros_ops.h,sha256=6OzY2ObtV9NggSvRmLF8aBRlkwkhCKgUKBT3CZB26D8,4160 +torch/include/ATen/quantized/QTensorImpl.h,sha256=YCONez30ZR878R0GgzfJazkS44z8xQBYXrIqGM2JG5s,4275 +torch/include/ATen/quantized/Quantizer.h,sha256=11qYFvRL3cOpn03oX2jj-McVtmqvxCTVp70_i2HgokY,9486 +torch/include/ATen/record_function.h,sha256=dcBOVXeL6voE_KcAcjxKUgGkVtI-JZIFzpN-vsrOEX0,24519 +torch/include/ATen/xpu/CachingHostAllocator.h,sha256=gONgsOfhC26F5q2kJ650FDaMpCQT1EvK3OHy-365fDg,1611 +torch/include/ATen/xpu/EmptyTensor.h,sha256=rOqGXEQjRbuhyjk0MVDF3ys9HbnwSjMzbOwtUq7nGCI,1413 +torch/include/ATen/xpu/PeerToPeerAccess.h,sha256=asP_YLMYgbcCoC3deEigGTYHVMKMOCQhO2s1F8sQYeI,954 +torch/include/ATen/xpu/PhiloxXpuState.h,sha256=1XsgWS7eFn9uyFMrhp9KcWESdjg_TMwP0LCYdS8dWWg,1261 +torch/include/ATen/xpu/PinnedMemoryAllocator.h,sha256=pQ7YKcs5YKKEwiGf5yUQROWjKj5WXtt2GOs-dZg1nIA,500 +torch/include/ATen/xpu/XPUContext.h,sha256=2mMuW1pieqh-HETTKnZBhNdwC1fv-A-taQbVquPzvNE,791 +torch/include/ATen/xpu/XPUDevice.h,sha256=rSn5aU7Pn8Pdmig-ckK2xwXXUpH1Q9rtzGJI63Dhsv8,521 +torch/include/ATen/xpu/XPUEvent.h,sha256=jdUIC50oEIdXyYz1r3J2-gQoYg_HIYl9iIzNVooOszI,330 +torch/include/ATen/xpu/XPUGeneratorImpl.h,sha256=n8gJT7tSo_GSdzMzdaLktqU1EpmHLTi2Bk0tj9CssGQ,2912 +torch/include/ATen/xpu/XPUGraph.h,sha256=1aTwdrHbLCGCA00JpuLq4txTfI6bqASHvmYhVnoP8BU,3557 +torch/include/ATen/xpu/XPUGraphsUtils.h,sha256=Z6abmhD97-catJA8HzmeuHUDPIVg3Bt74WSBNLEMiR4,810 +torch/include/ATen/xpu/XPUScaledBlas.h,sha256=CMnT00Kf9FJX1P3mN2VkZOb6dMFo5FVsEvxjecrLA6U,2639 +torch/include/ATen/xpu/XPUUtils.h,sha256=kpMUbXqmDwFezM13YxjumUw91b-3AKe-DNgWQE78IGo,676 +torch/include/ATen/xpu/detail/LazyLevelZero.h,sha256=7O6n_KlFmEixug90cyphDAx6Ht-zO9c9lDhdne9836w,485 +torch/include/ATen/xpu/detail/XPUHooks.h,sha256=Xl3mIs2Mq8HZuNWRjfporgN9POVL2j0Bhrh2SGEhxOg,1488 +torch/include/THC/THCAtomics.cuh,sha256=H75MAoZpXZIKRrPHUPUk-ZQaky8eTN8UYEcI8-BsfdE,372 +torch/include/THC/THCDeviceUtils.cuh,sha256=U2DABMB1OvzbnGeK5XZgJaHx3rt_Ox9gobPJgzm_cCY,332 +torch/include/advisor-annotate.h,sha256=MgjrWkMfFXXb2B7pJlf4xO-drRhiQUPw5t36P264E-g,23930 +torch/include/c10/core/Allocator.h,sha256=qoy_QZsb6KEIXWvDztG0K6E_JAjeYyQj8LbXCRk6UM4,14571 +torch/include/c10/core/AllocatorConfig.h,sha256=MlMZS9VX7iExnBbZ2l2mJYHkGvENUKF6J6BCKpmZrNU,15070 +torch/include/c10/core/AutogradState.h,sha256=RoJraKpzqdovlTTR6UpNBIaQW1aAwwOHogS4RO__zDM,2268 +torch/include/c10/core/Backend.h,sha256=jUqyE50qhbELtKPwwVryHE_7Aj6M9YXT1PEUbO59uy8,12187 +torch/include/c10/core/CPUAllocator.h,sha256=_tQd7SAIJ5k0s7Inu0P5L-o7SSs7diwk2j71sLX98Lw,1953 +torch/include/c10/core/CachingDeviceAllocator.h,sha256=oWD12jnCBnQKGhSUJDApvu_rLhDqVHGb8SGHSigbUnM,9510 +torch/include/c10/core/CompileTimeFunctionPointer.h,sha256=ea6PBQnezWHDIAcwtLYbIb-iS_nRuNcF8aE5lS4dgYM,1954 +torch/include/c10/core/ConstantSymNodeImpl.h,sha256=S0UAsSIydJwtWNl5lp759zfM77quImnhBQY7PMGPUhE,3363 +torch/include/c10/core/Contiguity.h,sha256=OcYWLOKaMJuo7dRL5Yw3piqMJ_a5706yUPGr1y47Zr0,10039 +torch/include/c10/core/CopyBytes.h,sha256=MawPcjyx1VMe1f9kp2RIx2BjAs_lI8jlK4BuppekD6I,1597 +torch/include/c10/core/DefaultDtype.h,sha256=KZKv7UUxi8v82O_Rk4gLa8yQ7QvkYq2HYUmxinwch3I,648 +torch/include/c10/core/DefaultTensorOptions.h,sha256=HgF3PO38zQ9lgd4RxZ2D5wWhobFPof5Sao6cxb_XJ_8,1318 +torch/include/c10/core/Device.h,sha256=xOo02kJjwgwFhsBck2TOjH3cg97SsNwbo0NItlqkxr8,7120 +torch/include/c10/core/DeviceArray.h,sha256=22EPv2g9iAoMplJyqgcPJkRJ35OwJAp3LOv_6RFg724,937 +torch/include/c10/core/DeviceCapability.h,sha256=sw7aVyoj2M0skKdyJQoauOguTqBcB9yYKvJwRDRmyoY,2534 +torch/include/c10/core/DeviceGuard.h,sha256=Gms0KqpunmWig_7lmwO7jrGu9f9vSESDEbsIQcBz5FI,8062 +torch/include/c10/core/DeviceType.h,sha256=FSeENqyP_wy2oz9TTWCVIn7UAWq0p2v4v9tMqMTCKlc,1055 +torch/include/c10/core/DispatchKey.h,sha256=C3OAST7bz92qGy7aKDdhRZ9hZ0aEBgemBFJ2FmmR8s4,32690 +torch/include/c10/core/DispatchKeySet.h,sha256=g6mlJJHWxlCYR8yzsAim-z0L4gJBSoABSVbVsdDPgqg,42040 +torch/include/c10/core/DynamicCast.h,sha256=Rak9MDhBjaSbRxDI-yqQnB8f55gwpX0OM3n4yUuX-Lc,4868 +torch/include/c10/core/Event.h,sha256=UadXfNbB3Rga50tKW1BbTvfB-0_3vmscx6ta3JbwHxk,4710 +torch/include/c10/core/GeneratorImpl.h,sha256=pXrJKhbpEkzfzW2ZmattDRTK-o4US_0VNP-LSEGcczg,4191 +torch/include/c10/core/GradMode.h,sha256=z9_siQQtwK4RxWAE2eHUyt20IIUgGpdfGLp_LbMzKjE,1915 +torch/include/c10/core/InferenceMode.h,sha256=WiggR4JpCIisvvfKIhts7QMvJBsUvG3uksyGQBAkUZw,4017 +torch/include/c10/core/Layout.h,sha256=WZRk5uQ1-c6RjGjk0TGDwMIWZrgNSxMoXZZxIqQ4TUE,1887 +torch/include/c10/core/MemoryFormat.h,sha256=XI0mFyt3vpwyPJ-urnb0vmdrW2W1EOeoEJuYyKbdGeY,8843 +torch/include/c10/core/OptionalRef.h,sha256=VGHaFtquoi3P9FKVWvCmaEmliVViTAy04X7PLWl-IXc,775 +torch/include/c10/core/PyHandleCache.h,sha256=FouTl09n8S7PrgCeNrIYlsUJefsAuddQkYKc5aJveNI,3355 +torch/include/c10/core/QEngine.h,sha256=0XnrQ3HHh00H39dSdMOnWxDt2mC-2F8gUbRd5mtDopM,1264 +torch/include/c10/core/QScheme.h,sha256=Hxao6hBmnsPo4mDykSRoaHGQV1BKndR9W_qvIoD9P5o,1934 +torch/include/c10/core/RefcountedDeleter.h,sha256=z9vp7IFlO65iZE_GqMSfYYX5tXSLOla7T0jgecY96UU,2505 +torch/include/c10/core/SafePyObject.h,sha256=C4eUbVovN5DSCny3CpaMuV0iD65i3urfiM25v9Ce5c4,3975 +torch/include/c10/core/Scalar.h,sha256=0XcaSo3YRCBTAH2PSq7JT99WXwesnsMjthZj3kWLDDw,14411 +torch/include/c10/core/ScalarType.h,sha256=dXc5-TfXfEyB0pxmXPwxXmONPCxR0gGFee6Bd7knZoQ,9704 +torch/include/c10/core/ScalarTypeToTypeMeta.h,sha256=2ZAPKZkxiEGr1pwqGpnxjqF-ISfk5b4jk6zf18IW_bw,1611 +torch/include/c10/core/Storage.h,sha256=nEHPb13MWuBqzgsdGVd9j9yIb0xteDeoRsSIrkkV97I,7997 +torch/include/c10/core/StorageImpl.h,sha256=sVXhSlIrjrFunrYWqLzlqFUpT6Mjy4aLa8NLa-oCNx4,12990 +torch/include/c10/core/Stream.h,sha256=xmk3T7nP5C62JKrPf3fmLAT2C0Z7JOgkGV8BgDDRmnQ,7032 +torch/include/c10/core/StreamGuard.h,sha256=lSXnEpjYzbhTv5LS9PbBpCHFY90SDmbf-pYfDPIqosk,6842 +torch/include/c10/core/SymBool.h,sha256=rdRarYzPMHq3OvsNyg05BbQMt1dDkqiuBZw7dKkVO7M,4892 +torch/include/c10/core/SymFloat.h,sha256=mhRHhTtXNNcyr3B5sDwclWDOLzI5_HOkA64fBNFN_iE,3891 +torch/include/c10/core/SymInt.h,sha256=zREY9U92LNHcuc7xIhBpDa54wGq1lMGQDdrzkhFMIn4,17653 +torch/include/c10/core/SymIntArrayRef.h,sha256=psgPBcmm_8F1Gr_Z9GrUo94nF7LmvkHhs54bP6yYME4,3530 +torch/include/c10/core/SymNodeImpl.h,sha256=e43eQx5LX4kt6RxU2aDGfaFcl89Lpnl-fcfg33lFjZ8,7598 +torch/include/c10/core/SymbolicShapeMeta.h,sha256=VAbG9Yrm0GZUDbOJfIar8mF3Ulfozh0izsS8FAMnePk,7997 +torch/include/c10/core/TensorImpl.h,sha256=8PxRh1CyMW2g57r4v53mGdJglIFoVSZmAo4eMAAm1l8,116916 +torch/include/c10/core/TensorOptions.h,sha256=qFFAmNKqdaCMXoRMge0hy4DjxOz2SAUAhcfjWFxl_ts,27288 +torch/include/c10/core/UndefinedTensorImpl.h,sha256=YPHuqeThSMX-hlQXH3_1ax1hejSGvexTJQeULXYJs4s,1487 +torch/include/c10/core/WrapDimMinimal.h,sha256=GzEltFCQfxV_PknyCE2aTtan-Uie-0AtUDNPO76f6ZE,1614 +torch/include/c10/core/alignment.h,sha256=CyLRjpw3Jm1j6I9sVCPc4F1wWXOfo1wcEzSYuezXhNo,1132 +torch/include/c10/core/impl/COW.h,sha256=KPS_fVIdl7S7-n4OZAbTbv16aoV_1pSDCXgn6n9uRLg,1310 +torch/include/c10/core/impl/COWDeleter.h,sha256=XnD0YAt5Euep8_WazaoqIs2sxEQpyunbVwLgAoTtGqg,2352 +torch/include/c10/core/impl/DeviceGuardImplInterface.h,sha256=AjqzbFsdfRiujQbfZfHIeSBOXsD0ksIq5dp9jdd1OAs,15803 +torch/include/c10/core/impl/FakeGuardImpl.h,sha256=AO7r_xeS46X0I_63CpAmTlaB7BTNdyo6aYVudvpYBac,3436 +torch/include/c10/core/impl/GPUTrace.h,sha256=h3gsM8i7vgjyb8ypN9JpTtxbcKr5xGRB05IML4lffZY,1128 +torch/include/c10/core/impl/HermeticPyObjectTLS.h,sha256=3ZAGs8gcdy01ZnRPqN9i6SrN_SJUmkXLi4F-FhGqSR4,2800 +torch/include/c10/core/impl/InlineDeviceGuard.h,sha256=R8ooBA35wnukVZ5_Cl4eo19Jn4azaxVeUngCTTS66sI,16165 +torch/include/c10/core/impl/InlineEvent.h,sha256=oSHUnxokPTjgAtJsjrcL4iwYaJogtdguP4Ny8bCubns,4562 +torch/include/c10/core/impl/InlineStreamGuard.h,sha256=y0-b1g1el6WQGK1mHXMwcniDstsC9m8Q6HIUL_K4lWc,10168 +torch/include/c10/core/impl/LocalDispatchKeySet.h,sha256=uPWPtzGIfkJqvALMpp-1kFfQh_Ng7owMmmZGXNt9PCU,6819 +torch/include/c10/core/impl/PyInterpreter.h,sha256=Imdw1o8r0S5Pb1NiNyEK_FpB0StVgyyTl9dV67pa52g,11003 +torch/include/c10/core/impl/PyInterpreterHooks.h,sha256=_6a3hwA1xa1ADfvedrfYZs9A4gaCNdmFHABb8J5_azo,1370 +torch/include/c10/core/impl/PyObjectSlot.h,sha256=uHrktIqLM_6T6wUQurLruZI0CzQOKqsoHk9OETxAq9g,2080 +torch/include/c10/core/impl/PythonDispatcherTLS.h,sha256=iPTFvUJXiDcNL3U0fL9DwbnRBvHKA98Sj1mdxG_esWE,1094 +torch/include/c10/core/impl/SizesAndStrides.h,sha256=3Lk4MQ_G9Yxd466J0Y63J4gznQJvjSTNcTCgn3CEzx0,9010 +torch/include/c10/core/impl/TorchDispatchModeTLS.h,sha256=cc-jimOtCaHndjJJC4CXNc0ZxiJbojjZfbDSnGPbi0w,2535 +torch/include/c10/core/impl/VirtualGuardImpl.h,sha256=5Ds4eZ78mrIaBTos3XwggFbCLqM5wIwjxq-23EsuGhM,3879 +torch/include/c10/core/impl/alloc_cpu.h,sha256=r1kaVQbsvfDznXZhbyfzCd5V1iUhM5cb-i__TG8p7Ik,926 +torch/include/c10/core/thread_pool.h,sha256=I3B_S9Ug2qpxoMRzI6kpxTEJ2QfdE2qQ8_LGEglKpSo,3251 +torch/include/c10/cuda/CUDAAlgorithm.h,sha256=lsQzc-6okqaGCA8Jbeh2Z2nMRzzeMqPrjP8R7_R4XSc,1295 +torch/include/c10/cuda/CUDAAllocatorConfig.h,sha256=vChljCcd_cyTlfnWFxrrb6qqbTVdmsZzDgTU_QOvc6g,8797 +torch/include/c10/cuda/CUDACachingAllocator.h,sha256=_kOT1CwWgMRPHxfMg2DCUWYbFw2zEgHH5JAA-ARZ96w,16482 +torch/include/c10/cuda/CUDADeviceAssertion.h,sha256=8MMMTcwJdHHovtI4mLZQMWwvin9qrDfow5PCTBj9Zq8,4416 +torch/include/c10/cuda/CUDADeviceAssertionHost.h,sha256=86Lu3bQ7p4NJcHwlr_8mD1cThCvr4dWu5hArkd4C2kk,6912 +torch/include/c10/cuda/CUDAEvent.h,sha256=wFTPunNcllPfJIVjwBvLVsXHXXCCjrENWR2e1HBFDz8,12609 +torch/include/c10/cuda/CUDAException.h,sha256=xz5j1Mude9IsuNUReszYRnbTvPPY9ChAu3I-vA9dmSo,4776 +torch/include/c10/cuda/CUDAFunctions.h,sha256=tzta019dB33TZD1HGdR0u4lgp_G416OvCH3U7vJ0Nek,5357 +torch/include/c10/cuda/CUDAGraphsC10Utils.h,sha256=2X18n30MqsLggVVUQ71PWVMphOC69wq8ABoegun0n3Y,4087 +torch/include/c10/cuda/CUDAGuard.h,sha256=pogRsenjYuqUa0rxFCuv4PQVtvnC3fbO8eMDF-llypQ,11599 +torch/include/c10/cuda/CUDAMacros.h,sha256=vKcRKJirbLzl6keZ7h9Nq0GHVBG4k8_xD_F6_DRo-rs,1733 +torch/include/c10/cuda/CUDAMathCompat.h,sha256=-ycGywH4SB7-AFYszcthbtzN3tYeRLC59wbhveUS8Gc,3800 +torch/include/c10/cuda/CUDAMiscFunctions.h,sha256=bAYoH_BsfE0drBHbPwbQY1BKpQvzAqIp47Ttp3ixSQQ,682 +torch/include/c10/cuda/CUDAStream.h,sha256=pfD_N_mGiyV2surx1fIXprZ_ShdnPVvIm9aCPkwsMrg,11217 +torch/include/c10/cuda/PeerToPeerAccess.h,sha256=AAm8L22rGCpiQIXDJNNSPg4aOvm-DUQ9hoa5AG5Cq_E,2035 +torch/include/c10/cuda/driver_api.h,sha256=dzKlqIfRIf44TwIRt4Vj4IOdfFPbcUfZQ0LORghPDJk,7216 +torch/include/c10/cuda/impl/CUDAGuardImpl.h,sha256=EABPMCjkrNmerJ3H-5-KwjZQPyaxZN0cmJGRyCOgQ4w,9841 +torch/include/c10/cuda/impl/CUDATest.h,sha256=BTPl4FJNqWy_3-YjnTBSSP51zAn9mZNRExbII4A6ldY,368 +torch/include/c10/cuda/impl/cuda_cmake_macros.h,sha256=b4GPHGIedCcUlEmdyIUJSueqx6cxYglCPACGG4tBD0Q,448 +torch/include/c10/macros/Export.h,sha256=aKgtX7IoPp-XfETT47kLUzjtrnOMHwXoApOWvfEHm-k,298 +torch/include/c10/macros/Macros.h,sha256=wH4q51OHw80iQFJqqiA9fdomPzKQyNh16vYhQ_FDjkY,298 +torch/include/c10/macros/cmake_macros.h,sha256=1C1RQ6S24kR8rzJDnxSjruaF9pbsx4bfcQjFvsNVKe8,541 +torch/include/c10/metal/atomic.h,sha256=Lui3bBpSuKN6OZ85LZ1Z7LUQRPr9QY3fE2yajqFh36U,8838 +torch/include/c10/metal/common.h,sha256=FrHfAJXhYz8fKzsaJK2_SWTr2WifWv4wy9Y59jaN9UE,1574 +torch/include/c10/metal/error.h,sha256=A9oV-mQYa2KoAp6JqIeI4IA4xdHB8WLo0os6xcXbPos,2662 +torch/include/c10/metal/expm1f.h,sha256=W1tWWSsW05yGNpKRElkDGV2p-AMXWMa_NPhRTEkdN-s,3850 +torch/include/c10/metal/igamma.h,sha256=qqyazrvnsloK5hiJX3kibHJXEwbAHDBxSWfzwmMBcdE,29348 +torch/include/c10/metal/indexing.h,sha256=Uz68ioCEHXegsVtcE4yFgRJUSkcEWmoemwfXJMdt94o,56577 +torch/include/c10/metal/random.h,sha256=_bOmpTuZn1t5JM9iVCvOYWIttk5D8FyK6tjhBcrIT50,2471 +torch/include/c10/metal/reduction_utils.h,sha256=pn6Y5d7N1Qa-Al6YbHirW5RusI4itTtSCj_7fRKLDho,11229 +torch/include/c10/metal/special_math.h,sha256=dkMH8AEiimc1dJ-nEm8-BqaHmvl_b4QJrp8tfMZO0xs,88430 +torch/include/c10/metal/utils.h,sha256=ClfTqoz-pyzXk2ExgLXfX6EqcO7YYMK7eWIWtWVbxFM,13510 +torch/include/c10/mobile/CPUCachingAllocator.h,sha256=siV7RrY8N_OTHsB5pv_i_vLjwP42pTXwrGyXYSJDqnk,4419 +torch/include/c10/mobile/CPUProfilingAllocator.h,sha256=oncuerKj9rl0s1psb-coVn62gPeuacZ9bDRlkRIyS8I,4990 +torch/include/c10/test/util/Macros.h,sha256=zDj8MR87ceOXepi-1GwaSYEnZTKQpcsCivD-tTmzd2k,430 +torch/include/c10/test/util/complex_math_test_common.h,sha256=WybEVfYt5z6FHzF8BQTu3bhpdmizcJYweZWK6T6SUAU,22218 +torch/include/c10/test/util/complex_test_common.h,sha256=iQ67ozmpgtTR3kdLPpWYFjxSV1eQzlXzoQSKOGv5suU,20382 +torch/include/c10/util/AbortHandler.h,sha256=tpnYao6yYLJanEvxZbNYmuWGY0nk9xi-VOaCwPk7k7M,2446 +torch/include/c10/util/AlignOf.h,sha256=2DAJK0-8JTaWbnw0tZdUvyAafOIrnkldhnWVB2oOIpw,5160 +torch/include/c10/util/ApproximateClock.h,sha256=fhQloge2zvy7Px5s4P4BmlWhOzjlkh6XT353zYq-2qM,4053 +torch/include/c10/util/Array.h,sha256=H7UEVhnbTgTUx1NSbaGvr6RcwKM2c2rouVamgDnL2qM,704 +torch/include/c10/util/ArrayRef.h,sha256=Hn57lAMvVMwEAy5gqpEEE4Y1IR3evlLis9DAm-0HyHY,10252 +torch/include/c10/util/BFloat16-inl.h,sha256=tANwC7RCx2Q4TQpZxebutCOlc7XIiQ3rzCiC9GALxWk,298 +torch/include/c10/util/BFloat16-math.h,sha256=rfHxMYHrE90RgFQb6fRdPhlwO8X-gveGM_5XS0At72o,8771 +torch/include/c10/util/BFloat16.h,sha256=tANwC7RCx2Q4TQpZxebutCOlc7XIiQ3rzCiC9GALxWk,298 +torch/include/c10/util/Backtrace.h,sha256=99hh6frZffmL9B6uyPlKMEC_DVrSZ90iBnf0buYsBMA,1051 +torch/include/c10/util/Bitset.h,sha256=viuJC1vkqCvyuQAakEIe7gH1jc4sI6qJOjEmbQb61YM,3674 +torch/include/c10/util/C++17.h,sha256=knpXYMQApc3HoERm4dG22AydBp82DW9eAFbD_bdCVAA,1178 +torch/include/c10/util/CallOnce.h,sha256=MpdyEek7ttrfoxeAKlfihi81e8j2yfiWUlKxsKQbMM0,2289 +torch/include/c10/util/ConstexprCrc.h,sha256=KxnQ5FBry6GVPm0732nDRE82DIq2g7DWfN-LmHcLfoc,6849 +torch/include/c10/util/DeadlockDetection.h,sha256=quzUTfDHPvS_6mttyRt-dDe_9kyGSmnuL-eI_-yx5BU,2471 +torch/include/c10/util/Deprecated.h,sha256=9xSX2Bd3ZgvKHV3qLA28orYjoEi4OlQW4upS17Ik644,313 +torch/include/c10/util/DimVector.h,sha256=CMuEFgqGmK-gM6sVX6LKXCBrbE3KIjy1PkMBlYIW2SQ,698 +torch/include/c10/util/DynamicCounter.h,sha256=9CUviUDJgTs5bkUecWKED1k03u15M_AEQNVonkxFGh4,1571 +torch/include/c10/util/Enumerate.h,sha256=lF-iYhMrwqAgSJR8OFrJN767aNDdhoQ0cTmy6XvbGeU,4255 +torch/include/c10/util/Exception.h,sha256=t2nUJ1mMZw1Iefi0UytSwEKl3BVdW7TbaUtp-nrelwM,33847 +torch/include/c10/util/ExclusivelyOwned.h,sha256=wf_sQd-NXJvDWYzGc8vpY26EOrnNIiPPx4gF3tJBW5c,4718 +torch/include/c10/util/ExclusivelyOwnedTensorTraits.h,sha256=H9-sjDESpRHPJ9PNJ5wHsOznreGa6Q5AZJwIhzZER2s,2474 +torch/include/c10/util/FbcodeMaps.h,sha256=SvWo5BYW6hq-FW5gA-SaXpa2SWxs9wzkKR2erty_Prc,982 +torch/include/c10/util/FileSystem.h,sha256=qa1WnkhdQWOnmAdw6VGPMuRRFC7T7PZZjzTu9r_m6CY,891 +torch/include/c10/util/Flags.h,sha256=LICXlXDamM7sC-CLBxoacYqUvHZpu15JlnWiDb5vGh4,11154 +torch/include/c10/util/Float4_e2m1fn_x2.h,sha256=I5T3I6VvHz10-h6Au7k7hoQPZ0HGgj-QDDA3LmAEJOQ,306 +torch/include/c10/util/Float8_e4m3fn-inl.h,sha256=OUgD3KdzlsxbY4vU88NObVdrSSf2LZktRGUmnUQJsnk,303 +torch/include/c10/util/Float8_e4m3fn.h,sha256=OUgD3KdzlsxbY4vU88NObVdrSSf2LZktRGUmnUQJsnk,303 +torch/include/c10/util/Float8_e4m3fnuz-inl.h,sha256=dU4G7OuItSLI6hzDUQ3akWwsnurqoqmqiguPZqIR8ms,305 +torch/include/c10/util/Float8_e4m3fnuz.h,sha256=dU4G7OuItSLI6hzDUQ3akWwsnurqoqmqiguPZqIR8ms,305 +torch/include/c10/util/Float8_e5m2-inl.h,sha256=iqc-2y8RBLJiqt3vVigCgalkl5Xnv3knpiZ0fLBZLKs,301 +torch/include/c10/util/Float8_e5m2.h,sha256=iqc-2y8RBLJiqt3vVigCgalkl5Xnv3knpiZ0fLBZLKs,301 +torch/include/c10/util/Float8_e5m2fnuz-inl.h,sha256=bSW6u8wgewkReTQ3LAYMfRBsl9iVBlZttCu6ZtbzXtk,305 +torch/include/c10/util/Float8_e5m2fnuz.h,sha256=bSW6u8wgewkReTQ3LAYMfRBsl9iVBlZttCu6ZtbzXtk,305 +torch/include/c10/util/Float8_e8m0fnu-inl.h,sha256=QesGPWulELXJsE1_ymZwYgNS1LACzW1FRM5VIlIyrbU,304 +torch/include/c10/util/Float8_e8m0fnu.h,sha256=QesGPWulELXJsE1_ymZwYgNS1LACzW1FRM5VIlIyrbU,304 +torch/include/c10/util/FunctionRef.h,sha256=fyQ8_KLA5NIOXrYy34ytVaSHMgTm0kuGA9B_T4JnUTs,2584 +torch/include/c10/util/Gauge.h,sha256=kZMzAF3p4sit2L0JEX2gjttCzHPHBRaJyac6FC5YK0U,1446 +torch/include/c10/util/Half-inl.h,sha256=aQVHyZtMaFu4m6195VXdn0DNNsfDuBc4gLu9Jmb82SQ,294 +torch/include/c10/util/Half.h,sha256=WCt2hbTquUvDKVMV-07qwyw0pwhYQeb_V9xDYJeInps,555 +torch/include/c10/util/IdWrapper.h,sha256=XrOlnZhjCy8yK7PGWfS7452TFzEFjQvexjlwkWR1xzE,2590 +torch/include/c10/util/IntrusiveList.h,sha256=U1jvYlB76ulrnrgACg1WihddOs1vtEFq-RInhlPUGF4,4749 +torch/include/c10/util/Lazy.h,sha256=cqlGY0dfmPapAy_WV7Lih5j4U1q0Tobe__JaTW2VjXc,3077 +torch/include/c10/util/LeftRight.h,sha256=GKovBValnILIaGvYNlIhkZ7C0Z1gX0sMFW-qTxXK-kM,7463 +torch/include/c10/util/Load.h,sha256=UojwOytDL4SaAR_KCqRyqPTKXzhlN5RJcwbINCnoqtc,1178 +torch/include/c10/util/Logging.h,sha256=OEyhjoez2US21O9_P0KSQpUgJ8OjCWcWa35NNgWxmSc,14354 +torch/include/c10/util/MathConstants.h,sha256=1mpiCU8wwQfIepeIQznB-39Y3vE5sYBCPloMGZCdx4E,3944 +torch/include/c10/util/MaybeOwned.h,sha256=fClAWJ30LKxRPQEYhunWbKIiZFLhOs8R1wpxULDGug4,7431 +torch/include/c10/util/Metaprogramming.h,sha256=8zzmy2y_5eEUR4lMAPsfBcTadQusC2506eA4Zf7JJ28,305 +torch/include/c10/util/NetworkFlow.h,sha256=EghGK6X4dvknoqYKFIOsPBPiS-Tj-z43CMhDObpOswU,1391 +torch/include/c10/util/Optional.h,sha256=GzsvJjlMuP8KRtSRH7FTMJF6fHZo7k497eXst1y7u2c,2249 +torch/include/c10/util/OptionalArrayRef.h,sha256=BWL-GFUet5QKN9Wi_jaRupiNAgWuxG4iMza-51f06Y4,7438 +torch/include/c10/util/ParallelGuard.h,sha256=uV2vMWwnlXgS7s082IZruJiv6OkLAYMMD3Ty39kIVwc,627 +torch/include/c10/util/Registry.h,sha256=HbWzeygfVQq4hX5gJ_NsPVHlf6wRqrUu20Wn3fZH8LQ,13514 +torch/include/c10/util/ScopeExit.h,sha256=t3YVD8cpprBt4WXFuHi1jHT9BJ8DhFSyGtQZ17CEF3U,1513 +torch/include/c10/util/Semaphore.h,sha256=6MAXrwYvAWZAoVFl6NnLntQ4ugHabFOBdwRkqYzF0O0,1992 +torch/include/c10/util/SmallBuffer.h,sha256=tIOQ9OvQaBuOI6xYm0FNHIx4BsGODO4GEQiUCPV9hDM,2016 +torch/include/c10/util/SmallVector.h,sha256=Ti6EP3OL3y8t0qvtMozmC3y6cbM7U6z6_YBg3t9G-o8,49323 +torch/include/c10/util/StringUtil.h,sha256=Yzl0fdy6UpGMOQnuS8nPW9mWlH8cO1PuWGAfwfVZ2kw,7300 +torch/include/c10/util/Synchronized.h,sha256=2tW_9s7EXK40MUtpwv3-YJjA7FQwWllpaxOZLR0ar6o,2179 +torch/include/c10/util/ThreadLocal.h,sha256=wCti5HZb8cPA89t_y1OpPuw7DHj98QT9YfM9slvgxqI,4274 +torch/include/c10/util/ThreadLocalDebugInfo.h,sha256=X3IiE4KqBH1ftqZxGu2pYU47Ev7WzjTGuJlTFvfd-z0,2919 +torch/include/c10/util/Type.h,sha256=T1XX9iuFVOaeV2BNjLjgt5Hp1Aku0Q-Ecl_PolibcmU,900 +torch/include/c10/util/TypeCast.h,sha256=fSNsybFcY0MkrAxthVh-wdmrDYg3niv8glx9BYurJI4,6916 +torch/include/c10/util/TypeIndex.h,sha256=dbzYTSLXTBMex7BW93K3W5ABC6_2-o7vB-1XbjWIAAY,4947 +torch/include/c10/util/TypeList.h,sha256=zkNZb9AcFUJktCgDJSi0K2OKn1w8K_UvhVMPT5nhWx8,298 +torch/include/c10/util/TypeSafeSignMath.h,sha256=HvEviNWsHTA3-frX0xDhi-n0TGTNrgwC4GzSRy0dxHU,306 +torch/include/c10/util/TypeTraits.h,sha256=7EvZei8Q0T6N4WzrKKxXrKYtqLiBaU9GEh1l2BE37_Y,300 +torch/include/c10/util/Unicode.h,sha256=Ua4k-_JzTSImffi1DR5qlhijico5wN8hDI5pCl9sRB4,549 +torch/include/c10/util/UniqueVoidPtr.h,sha256=S5MaEA0td3mAog2tT9EaBqNhx_XneoasbaSEBRkTf0s,4835 +torch/include/c10/util/Unroll.h,sha256=UPOfiyO4HaMYDiu9w0X48H1cC1noEQmvQbfh6tnReRc,1097 +torch/include/c10/util/WaitCounter.h,sha256=3GaSw6xL0LaXhwA9_Rfi5WJSU3lFiXgHM_x3xwIRL1M,3050 +torch/include/c10/util/WaitCounterDynamicBackend.h,sha256=MsfBpIa5pRwQzILRp5Pg6AmlgPvnXXUOfdiz1RE_SYY,916 +torch/include/c10/util/accumulate.h,sha256=vqbcE_ZUjeZ4eC_9x0dAFDinmn-TWl0jp6ZEHNg3Crg,4286 +torch/include/c10/util/bit_cast.h,sha256=sxSghPDly-p8OH6gmTpfHz45ALOYU3Q_qSSForY8N3k,298 +torch/include/c10/util/bits.h,sha256=NvnFKQ-qXIxRN0UJsSzWvlkCWhAaVwJwOQ_saK6-00U,294 +torch/include/c10/util/complex.h,sha256=Fp2VuG8-5FJP0A-ppPtz8hC0OtPl0BibZPYr5I29nI8,2354 +torch/include/c10/util/complex_math.h,sha256=hFwpXlHE9juLMn5i_yyz9x-alNFHG5hMv7XrIdna-xg,14357 +torch/include/c10/util/complex_utils.h,sha256=uh8DvEyYr1jMQ06vgnx3bS6YNB7zV4sMoyC8pIA7hkY,1331 +torch/include/c10/util/copysign.h,sha256=4fFRt3OSrkCMfGUGbbcqHpCkEs0pw8_uKGOsJGr11Z0,1086 +torch/include/c10/util/env.h,sha256=zuLIDHeM7sxVohvXwl6-aoNwcTbLXozkjUzO5Ugl1Y0,1148 +torch/include/c10/util/error.h,sha256=0PShMfIU8VkM8fME8HC_gAkBTfPsP9Op4eeS_k6VpX4,459 +torch/include/c10/util/flat_hash_map.h,sha256=RrZriUT2NETrS8tCQeh7mET1Nf4Z7-jZWK34UCU6wlI,62262 +torch/include/c10/util/floating_point_utils.h,sha256=l8yiPUHN9Za9CCKY9MUFvPUMmgP6mwfuetzTN5TobgY,310 +torch/include/c10/util/generic_math.h,sha256=ezfY9osScXISPCzKsoVKsVmk-r1zyaWkQKNbda3Hdt8,3379 +torch/include/c10/util/hash.h,sha256=0FzKQcvZ7Cdn29QZRmbdGVSpcUUrG-gn8J5MLucmmPM,11360 +torch/include/c10/util/int128.h,sha256=lGDT-DCXohEx5riIIeri_91JleTaxyiy8SmZtWuQUz4,12577 +torch/include/c10/util/intrusive_ptr.h,sha256=lGJXxgGiea_guUwgI6PQ7PDr6WtHV3aMRJzDAbSfHaQ,46544 +torch/include/c10/util/irange.h,sha256=UzaE6wBFzo7QL4IjDnvnZ2zy1j8AITXz2wRpg91G_sc,3690 +torch/include/c10/util/llvmMathExtras.h,sha256=qKPCbOKabdjuq1_lTAeV49yAzvVHR8d8HA_m6wW2B2g,29717 +torch/include/c10/util/logging_common.h,sha256=tvOqN0oPLTJRZ47uh5xapHMgZaYWLiENqzOsadOaR4A,2053 +torch/include/c10/util/logging_is_google_glog.h,sha256=5_XA51tYu4EgZVQDycJLO-Osjy_jcicQOf_ifTUQzyo,3235 +torch/include/c10/util/logging_is_not_google_glog.h,sha256=D2V1ZZyuk4rteBLKehK8hrnb389H6cpgv_q1Hpj3Qno,6459 +torch/include/c10/util/numa.h,sha256=ThNH5iMdklhU1kgFOdr9IlHryVIOKKdcwexg83E0mJE,967 +torch/include/c10/util/order_preserving_flat_hash_map.h,sha256=O0szYNijg27FvK7v2zLo0lQo7ltmD69U6VL69pCtQRo,65783 +torch/include/c10/util/overflows.h,sha256=oHbFZO3daQDjJgw1skcYrF7rFUhSkXw8MVskwWOPu9U,3721 +torch/include/c10/util/overloaded.h,sha256=FHGcSpssgR5hnXJJ-MIn9wI2Mqb8FVYgTVNNcDlQ7d0,981 +torch/include/c10/util/python_stub.h,sha256=JqFk4mKCyu0chr2QF5a1IqXXbi9HPecVk_VhA2TxtCA,310 +torch/include/c10/util/qint32.h,sha256=ZxbaslMWGDjHB8VlJu25dTRcWWwklUFkUh6T9b8HHsk,296 +torch/include/c10/util/qint8.h,sha256=8IYNTICfJYSwwplrhGH9VGvjyFW56zDXd59Mh_3l6S4,295 +torch/include/c10/util/quint2x4.h,sha256=zcnQCEqSMrTjovsd5FdRF_KMIumwqhT5zwsHMyRcVHI,298 +torch/include/c10/util/quint4x2.h,sha256=IlMXyvQf4QqSO601TKk69qCMcfiSV8EKGKKQixbqFrI,298 +torch/include/c10/util/quint8.h,sha256=aydMTZmqEB4N7-EM5mRUhTZYJIZpSMJkiH_khO5Qcgg,296 +torch/include/c10/util/safe_numerics.h,sha256=g8o-ABoIf_kpszh3kVm6xLtHiqlzzbfNT2Evx8mcLlg,3347 +torch/include/c10/util/signal_handler.h,sha256=gv0fEnjAovuN0bL101qlZCFKTpj0VvXha_7pG8BG__s,4066 +torch/include/c10/util/sparse_bitset.h,sha256=tLRVQ8J_dDTkoIw6lVU3ZbIdcBW9ZHltKf5l0XaONQc,26931 +torch/include/c10/util/ssize.h,sha256=3LVqLbLSSldl1SP7d45-5ov2KxHnluC7_z7Pf_qkFXA,1623 +torch/include/c10/util/static_tracepoint.h,sha256=rED_IEBTyD9MieY029u4UdzsFnLaXVpAA22wplohBbE,1330 +torch/include/c10/util/static_tracepoint_elfx86.h,sha256=uT3xse24frgvxeuyHh8K34U6YcynWti_y9lacxCOZFo,7927 +torch/include/c10/util/strides.h,sha256=ijl5aJ28SyH_L44u77n1FCOs_kwdm96oYnHFo4IjPC8,884 +torch/include/c10/util/string_utils.h,sha256=FZ-mqgraU_uZqPS1m6iaJ0xD00cuCQ50iycgYya2GBY,700 +torch/include/c10/util/string_view.h,sha256=CQhocLd5DudSxVaIG-5MmgTd1xc1sXxQShpyutXBi0k,18654 +torch/include/c10/util/strong_type.h,sha256=riwLN2kJ1cUnuZ8Ow9jDy1Ak_hDFbXHD4jsfLq5eB6A,35793 +torch/include/c10/util/tempfile.h,sha256=QzgZNO5g3E4TW2g16iL6DwmhkdFOUDrSfMKRHCrYMK8,3014 +torch/include/c10/util/thread_name.h,sha256=5EZrohns_m4Tpb9e7lzI11nm6lGBd9wX5C4eA4l9ix4,440 +torch/include/c10/util/typeid.h,sha256=f4AlUA2mZoQUVarFlDgtCkHW3fg5fhcaBRQZdbwxj48,23584 +torch/include/c10/util/win32-headers.h,sha256=_EBVLD8V16450_L4ha7WnXQvO6a2zP_FXILRLogCLPE,1112 +torch/include/c10/xpu/PeerToPeerAccess.h,sha256=opZ60OwHD1qnRupWAY5zKmuEmAeQP63NK0xDe-XOTIc,741 +torch/include/c10/xpu/XPUCachingAllocator.h,sha256=LOCZ3fL59NWYO1cdYMX3kcia_GCXz-kYkFEM1AKIYYI,3828 +torch/include/c10/xpu/XPUDeviceProp.h,sha256=vgC3s2i1LBZVm0pQEwX8LVQD0DlXMOrf2_ZoZpiWZNo,14014 +torch/include/c10/xpu/XPUEvent.h,sha256=f9OVsmLdmvRoXT-NT4pIbcjBou5GWU2ToJEK9ROfBe0,5640 +torch/include/c10/xpu/XPUException.h,sha256=LYBaOSJ4EGshvXPMZQ6CvER0NrWBJy-Hr9hWnWuOfeI,728 +torch/include/c10/xpu/XPUFunctions.h,sha256=d6UF_lw5NTPHSSngOGDVP4JOpKEWg5M5Qk94MXYts_Q,1523 +torch/include/c10/xpu/XPUGraphsC10Utils.h,sha256=jky2gKO_NFZWGy9peL_Krn37HAHUhXuc6seXUr2wYi4,1350 +torch/include/c10/xpu/XPUMacros.h,sha256=r6IGtxNCSJg1ccaOrhyZFS-NfTUmGbWhH39Ce9r6o7I,1124 +torch/include/c10/xpu/XPUStream.h,sha256=RsB35OmcbQJbdmeAx8n4rPokuz3P-PtizW3RDY2MJs0,7133 +torch/include/c10/xpu/impl/XPUGuardImpl.h,sha256=LHK-E6J8XI3wyYodT_iV9Xlad3FCY3D9DhrX1XR8uhM,8915 +torch/include/c10/xpu/test/impl/XPUTest.h,sha256=U4e0J34PGyBX6Tao-X3X85TLgF3sVNpA1xq7oA5Mb6c,720 +torch/include/caffe2/core/common.h,sha256=s1y5zswxMN52p4NyHDZT9IFP4KqgahRgGA8tHVr2Gcw,1613 +torch/include/caffe2/core/macros.h,sha256=_ZsNGYIgK8Oe1hhF04e131CyNM62PTOsZg-Zimg9bKM,3987 +torch/include/caffe2/core/timer.h,sha256=C6gSn8UqhEI1yr7isd_-xGwmYgoIJKhaoIAvah0J12A,1472 +torch/include/caffe2/perfkernels/common.h,sha256=V-FFQ9vAOsaaG4QE08h7Zeuxh6H4m_umwDc8j7FOUuo,6403 +torch/include/caffe2/perfkernels/embedding_lookup_idx.h,sha256=v6rEhVi3EdlyELZwq6W_FZSPS0cdv0emL-0Q-z4kbpU,1928 +torch/include/caffe2/serialize/crc_alt.h,sha256=9-npt1fZ3vAS3LwE9c-sxmOhpgYokW1WvDJ7q-k3iIE,75756 +torch/include/caffe2/serialize/file_adapter.h,sha256=pAP8jaI0BTcjQB2EfPZ76kpsNBQ9G3LPwrvO7nZTchQ,1096 +torch/include/caffe2/serialize/in_memory_adapter.h,sha256=QXZ31tffGCe7LEOd_GWfnlPpneX4VtI-S3BZ-XlH3A0,895 +torch/include/caffe2/serialize/inline_container.h,sha256=ZHsmkcz1SQA5fit_U6iDEkCroGRaQqug0TWOOj-jtMY,10761 +torch/include/caffe2/serialize/istream_adapter.h,sha256=-Y6HdewRC9Uhayv_q9h85R8PPO66a_qIolt1eR0tVlU,899 +torch/include/caffe2/serialize/read_adapter_interface.h,sha256=X-TyU4I-IxYbqQz0lyByvWJFYcAey8tPaX7xx0DESYA,786 +torch/include/caffe2/serialize/versions.h,sha256=nnNzQCd-EYVwQwfLfghke-iDM5X2FKIaVS9wyeJmpKQ,6878 +torch/include/caffe2/utils/fixed_divisor.h,sha256=Zvgc9CidZoQfI7j9QkTqrQaanXf5iZcsvxCJfnJE2bA,3786 +torch/include/caffe2/utils/proto_wrap.h,sha256=1QVXbf-ihQ7qhv9aNvO46oyVXUIYnQrrEi2jc3N9SO8,1514 +torch/include/caffe2/utils/string_utils.h,sha256=OZypdSsHGVDnqhxYJsh_yQXaW0FBlpoEEzpTLlxZieo,1460 +torch/include/caffe2/utils/threadpool/ThreadPool.h,sha256=k8zVVNB-237ElHKofNIJsBPktFzgjp3DX05Dn7j6thg,2630 +torch/include/caffe2/utils/threadpool/ThreadPoolCommon.h,sha256=41zihSOOaucnvGexl1jRBNHowqQDxcwqRPd8mv3PHSg,931 +torch/include/caffe2/utils/threadpool/WorkersPool.h,sha256=Vl9nyM5Fu6CZctEvXOetuqyiAXCoRbl9fSFqiSN8azU,11981 +torch/include/caffe2/utils/threadpool/pthreadpool-cpp.h,sha256=wdhFwxXXDwaAJ-Msy1LkBZD0KA7SEfXZ7EegAf3MpAo,1718 +torch/include/caffe2/utils/threadpool/pthreadpool.h,sha256=J5PnQfqWRF3ETKYJS_vcS95UxeaGIwIA5rmx8T5_spQ,6598 +torch/include/caffe2/utils/threadpool/thread_pool_guard.h,sha256=Bv3NtbE_UJ1KT7Z3Wb0ylMy59G2LOZV09ZWhG4Fr_X4,821 +torch/include/clog.h,sha256=LmhU0B5GdvZxhbdibrfyD6I-wBeQU4ZSuN_mL5O3ZJ8,5154 +torch/include/cpuinfo.h,sha256=2Xqq7HM7ppg34uH7bVY14LTG0H0RJYvgz8dE2HfcYHg,57043 +torch/include/dnnl.h,sha256=-npiSEIICX2ZJy2baWpMX4uOZQpgXfh1A_IZdGSBp8A,1080 +torch/include/dnnl.hpp,sha256=YPZAJVUfe_dEFAyBWc7Zl1dIWZCKQAvnRezxMjIhkSE,1088 +torch/include/dnnl_config.h,sha256=0LRW_9dJNTWARLqhK-nlCgbnWGxWQBszbSoVyRXZcnY,1108 +torch/include/dnnl_debug.h,sha256=XLjSNgFUn83ayTuqCZHZmn_Yvq4mjwaU5A-czblJTc4,1104 +torch/include/dnnl_ocl.h,sha256=03__vBMCxQj81y5cqXR4lF3yoDlZ2Sd6UQ2GJIKCqHQ,1096 +torch/include/dnnl_ocl.hpp,sha256=Oe6xYIdgeoptQbGC9Z8fBKd-iZRLbVZkVpwPxXgnrC4,1104 +torch/include/dnnl_sycl.h,sha256=QuFHJ9y0PVBk7wE5Ad87DUY8w1ut_0284HA3Azr4OSI,1100 +torch/include/dnnl_sycl.hpp,sha256=qyZJxdsGray5-Qb5GlsreeRFni2Ryj4dwGozwFJBkFk,1108 +torch/include/dnnl_sycl_types.h,sha256=cQOxuhg809O4LuehF2b01m66cjXDd11IYcEub_-txJw,1124 +torch/include/dnnl_threadpool.h,sha256=N_AYimWiSinfGRfaJO7QO-wwUSp656InQMsxoupxsvg,1124 +torch/include/dnnl_threadpool.hpp,sha256=FTZXdpXDy7T-HHTtZ4-R8NFuhl8n_61HKmXHqJLZd8I,1132 +torch/include/dnnl_threadpool_iface.hpp,sha256=5H7Q_E09_nyCmWPAg6KAr2_Vs7VyGDbX4pI6Eq9op0M,1156 +torch/include/dnnl_types.h,sha256=342AMEBUhQUC605L8na27vlAj9dChtGiZ645-f4kbfY,1104 +torch/include/dnnl_version.h,sha256=9R16ZzMZoFGazgp5nucYaHprBXG7Qqm2J67EHYUDZ68,1112 +torch/include/fbgemm/ConvUtils.h,sha256=cvaDHju_twkawy_GwHTNWOojzNqVxr34lCcGYWijFEU,6430 +torch/include/fbgemm/Fbgemm.h,sha256=XNuwhJzEdEyRaRRw_atSDUrshEnEPs7Murx5SHXyl3o,44257 +torch/include/fbgemm/FbgemmBuild.h,sha256=8BY1QpKWGNZcIargdcXtUAAVWnSpmz76aVwljBT_qfs,3567 +torch/include/fbgemm/FbgemmConvert.h,sha256=Xr47oOcirsIyLd1F9JRmMV0mDmghMIfVtZtoexGaQ1Q,4842 +torch/include/fbgemm/FbgemmEmbedding.h,sha256=OI1oO4qX1g_I1_tD4HOLGW0fttZCh6oJz0eBoW7Q0SE,12435 +torch/include/fbgemm/FbgemmFP16.h,sha256=qKdkHPTdodaaUQJwNjqp1UhG_WiPqempHwatsmaibvM,1597 +torch/include/fbgemm/FbgemmFP32.h,sha256=QOLrz1KxYQOC4wZbNPvyA_wbEr-CC4ZorFDlW_2JO7U,1398 +torch/include/fbgemm/FbgemmFPCommon.h,sha256=LPEP74Uk4DY9SqzN0vXRFGQR4Hj5H2LRoX8xy1WGjdg,9892 +torch/include/fbgemm/FbgemmI64.h,sha256=oUACFs1tWPDyd3pUO3os_j7tzTH7avHVRRL0OiZFVQ0,835 +torch/include/fbgemm/FbgemmI8DepthwiseAvx2.h,sha256=QumE9n8azzsmzGyx4z754ZYKcRNE-5z8uEIGB40RHm8,3795 +torch/include/fbgemm/FbgemmI8DirectconvAvx2.h,sha256=4oKMoeKXFALz6NSZwZDfNUc3c8Y4S4aXHdcTu22NBU8,2059 +torch/include/fbgemm/FbgemmI8Spmdm.h,sha256=tSDTxQLYGA0p-6NYet_oyL3ybCXjmlKf-XjihebThzI,3691 +torch/include/fbgemm/FbgemmPackMatrixB.h,sha256=GkHBT42qVCRPlSEM6wAd8GQxmcu_wLbRbmyH9h0Cs5o,9485 +torch/include/fbgemm/FbgemmSparse.h,sha256=eFxR4UQOaDQUpKlvw7FQIvvmMYIHl7BqPJf0JM92TF8,6642 +torch/include/fbgemm/FloatConversion.h,sha256=d2lOVk1l9sw1Ja3PSIe4_Yve7I-xBjmMGtEuv8Pd9Ao,11644 +torch/include/fbgemm/OutputProcessing-inl.h,sha256=_Q2OPwlpusdlH2V4SZYNOAm6WcptEuVT5Du8dPuXw0c,10579 +torch/include/fbgemm/PackingTraits-inl.h,sha256=pCHjfXyFkEMc5jnCg0-mO4RY-yLmKQpaJ7cnZo5ai2k,20905 +torch/include/fbgemm/QuantUtils.h,sha256=Ey5B3N1Hg85xU80qR0mxSGrmVXjhb2JoakQiEz0rezo,12920 +torch/include/fbgemm/QuantUtilsAvx2.h,sha256=vMZtI59MCQDHxM2QSMwTI9rP-x8-pC9-Di7QPineG6E,5313 +torch/include/fbgemm/QuantUtilsAvx512.h,sha256=bae15lfRDn8teH7NRXbINDIcHxu84LaG_5Oy-ojy8Q0,1464 +torch/include/fbgemm/QuantUtilsNeon.h,sha256=Wo6kEORiB1q--4U9n6P1vloWu5a2lja01e7oXZdrEc4,1463 +torch/include/fbgemm/SimdUtils.h,sha256=tEsPQiR7vh5V1kPYWH4Az1KeDsI99Z-5uhTjfL7gGKk,3209 +torch/include/fbgemm/Types.h,sha256=LMt3p_xD1ig08FYvzzDw4aC6Ood_SnE4fC7HHRGe_W0,800 +torch/include/fbgemm/Utils.h,sha256=x4ZJ2EF5i_UD_ZLP_UkC_OmlDJYx4RGGQ2jur3TjQkM,14636 +torch/include/fbgemm/UtilsAvx2.h,sha256=JZdUxlO1AggRAecTpyn-eSmyolFvLg8V5RbaBk8NpfA,2536 +torch/include/fbgemm/spmmUtils.h,sha256=yYztyD1mvv1xwSUEg6yP0GdPwH2Li9ogYcniEzT2P8I,1584 +torch/include/fbgemm/spmmUtilsAvx2.h,sha256=xC5hZAJdIEJJWp9dj51GpL9OXEww7dToPQBQ3f0yCkw,1376 +torch/include/fmt/args.h,sha256=DceIEI6i6AajzaHQ3Z0Ye95nHSS9vBtBN3rqs4Bx7R4,7453 +torch/include/fmt/base.h,sha256=vkhZh12cbDVEJMMgUo1aV-cpqSX3Ob7OZDbtHEPGVbI,105285 +torch/include/fmt/chrono.h,sha256=8xDR25gWjdvLhEEcdZp8a2ERFC-ikBTGV3xHjemlXL8,77708 +torch/include/fmt/color.h,sha256=psoxTkLeJR7c8DGpSBFs28uD00KE6sJOcHyouf3Zazw,24399 +torch/include/fmt/compile.h,sha256=aFQtB8Uu6C_LUv_iwlL2tKK6iwpO2p0me4Ol3ZagJao,20747 +torch/include/fmt/core.h,sha256=Fon8wYxr5EmpFEP62k7O0QhukP0z9N8E9PGfx-Ygs0A,441 +torch/include/fmt/format-inl.h,sha256=2rX8BILADI7NRM4WMvBU4BOf7jjaigbGc1_8LJnyqQ4,81262 +torch/include/fmt/format.h,sha256=C4PH44EiUxGOi1qdqLnp0zmTSIcX6IpvJWLjjWslm60,162920 +torch/include/fmt/os.h,sha256=9EaHZUNGSZHEmhPALGAlMKDWrJik47HROopcRhjSYM8,13114 +torch/include/fmt/ostream.h,sha256=TwItPiJ9TDezk5XYfFDL7C78XJZG3Px7YP1L6iUBFws,5281 +torch/include/fmt/printf.h,sha256=_6aU3tIvrI_cL9lLpsi6IM3HguW2jvFUC6qJwD9Cxtg,20245 +torch/include/fmt/ranges.h,sha256=eWBKZ53Smvp8H0qigABJ0G_IWxWXEkQRtCoiTvMkYwU,28380 +torch/include/fmt/std.h,sha256=embCKeLfMWOCMXmcLRPgbRS-t0Vk_n95u7xP8RCumms,23786 +torch/include/fmt/xchar.h,sha256=C_ZqGoQboP4kYOlGeHJFgmeVPt7AZFxSlNukFoeYRuA,13352 +torch/include/fp16.h,sha256=qOihPuIFJc7Ka_UrtrCjHjXFWHKZO5sm4MyRgoa6hUw,395 +torch/include/fp16/bitcasts.h,sha256=j98oDlFXdzNmpyi2G416tRL7K3bl2_LORCQnytZ6_cg,2407 +torch/include/fp16/fp16.h,sha256=O-KcphLcDRvpwLmJsl-3FqxqRoTuOvBYIVb0yPBbNmA,22256 +torch/include/fp16/psimd.h,sha256=2L-pHlaH1_t1SUpq0LOvGqKA3zScoC0hHUYhpYwEgbE,7099 +torch/include/fxdiv.h,sha256=tGAjCriH6CaOaYaGoVsuSr56y4GhUyC2NeVsdM1jKY0,13403 +torch/include/ittnotify-zca.h,sha256=B0LFdu1bVWrH3uE6bD37JkUy6GmUpapjNhT6vH1a8eA,4007 +torch/include/ittnotify.h,sha256=ubI9PiSnV95_u6xdsevUFQOOXNuzkp_TyWfMeI3Tl2k,199788 +torch/include/jitprofiling.h,sha256=KUCBurzd7j_1CDBWlBm6LGD5ccDunyQJw8SRfSX8wsQ,29872 +torch/include/kineto/AbstractConfig.h,sha256=SZtvJkH3sXFFnvlIqQfSiu91tkAuf2K547z8lEXJdiA,4132 +torch/include/kineto/ActivityProfilerInterface.h,sha256=j8dVmzepEFATjxGnQxmuDGX7lPoXkPwRP4S2F34qyt8,4146 +torch/include/kineto/ActivityTraceInterface.h,sha256=GqIs_ETbsFlvokpmirXS2yXi2eDmzodm9vu_4Wf8GKQ,863 +torch/include/kineto/ActivityType.h,sha256=BYH_pmjwtjaSXGtK9UC6kIxNQpPVRfsxwfZb8OruEsQ,4664 +torch/include/kineto/ClientInterface.h,sha256=YByMXn1wrv6ejl60Xz4aBLTWQkSFHs0zyoy0iCZ3cYU,919 +torch/include/kineto/Config.h,sha256=WOyRNM_4wyyyW50XAxgZqe-vXxNppoiYFpyckzOxiUo,16930 +torch/include/kineto/GenericTraceActivity.h,sha256=hUo-NM0HHzfrl2KHNkjUk3Zx7C0xcqM9iYYd2x1cxM4,4704 +torch/include/kineto/IActivityProfiler.h,sha256=DBN4UBzu2Y5buCQ3YJTWcFaNQFx__r49ZjZ6IPU6ifg,5723 +torch/include/kineto/ILoggerObserver.h,sha256=osHMpksk-Y6CttCP9nm7S76a98DWDpXgCEWAHSBlYXc,2205 +torch/include/kineto/ITraceActivity.h,sha256=OHr3NKCpve5ueI6DSm20PyZxtpJKDO2yCyvwQcU_zy0,2911 +torch/include/kineto/LoggingAPI.h,sha256=lguI_Ou7FauSHElTPiWO9L1813UTWuqTdV1hikv5IQE,600 +torch/include/kineto/ThreadUtil.h,sha256=pl7vMR9bCinR4wEs9j9dqjKYmqYNegZqcPIxxPWboHU,1179 +torch/include/kineto/TraceSpan.h,sha256=jZ5d58NqsAWrVd40t0r_kxXbfe9d19UFZQVP60PFT5s,1206 +torch/include/kineto/libkineto.h,sha256=FjisHFO3fHX-Y3JHhpc_69gnqBa--IWnFN0v5v8FTuc,4133 +torch/include/kineto/output_base.h,sha256=EZsfcBdpekr2YyqJCLYhx8yO7UC00YuVqqJAIeC4MZw,2491 +torch/include/kineto/time_since_epoch.h,sha256=UEGKzykCccgf73A5aAYSn-KpxQ7lIlEl1mRd-LlXKq8,749 +torch/include/legacy/ittnotify.h,sha256=WyHVB1ifIqOIvr7OvKeiXHoLi1FRGE_LuT5aGPf-2fM,37628 +torch/include/libittnotify.h,sha256=zA1MPO6UURWLg9TsFl2vb_Pivo3Sg6PsNZXcr31P91Q,823 +torch/include/libshm.h,sha256=NzZuufgNX2-Y5q_APKnvqt8Cj_f3Fbqts07hk9tME6c,1459 +torch/include/nnpack.h,sha256=Y9NhEeYb9-epCjY-a9X0e6xzyM55mXQAoJKR2LQ1Xts,33337 +torch/include/oneapi/dnnl/dnnl.h,sha256=XfmCDVNXJDxR3-v1FwCeD0gcXuEZQYliCtWZc3rrbGk,197127 +torch/include/oneapi/dnnl/dnnl.hpp,sha256=-V7T3-Ek960wpHJNDIdhSJ-GYJhj_bn9dfjg1wkeL78,662256 +torch/include/oneapi/dnnl/dnnl_common.h,sha256=dbUffm_vzMEsvBBdQm4QWiJcKx7dRZD0IL6BQ34MOaw,5905 +torch/include/oneapi/dnnl/dnnl_common.hpp,sha256=wGf9qyugz_u3Pxgl-Nu4ggeUeYdsyc-5Y5xklfP9gNQ,16045 +torch/include/oneapi/dnnl/dnnl_common_types.h,sha256=rqrRru0Lm1r8aQdvGttzFlzFB_Zt_t0-mwZbFHoXZKA,8503 +torch/include/oneapi/dnnl/dnnl_config.h,sha256=jw-6EKEpoe6b0XfqvSAMAoJmfr5xGRImmbBA1Gx1ZXw,6556 +torch/include/oneapi/dnnl/dnnl_debug.h,sha256=CyO3SZ5Xl83SZ8gBfr-I564GgMhZ1JaEm3WXjjXuUwk,2597 +torch/include/oneapi/dnnl/dnnl_graph.h,sha256=SsWXMGvP337gv4zIqpC9E0HnT0VzJPlNFgJvPDFEM64,34541 +torch/include/oneapi/dnnl/dnnl_graph.hpp,sha256=iQqmxTkhlan3Lnw6p7oT0PL7kOKDR9DCUR7P4Lzlaxs,68794 +torch/include/oneapi/dnnl/dnnl_graph_ocl.h,sha256=JBnC4cE6ewlQcKoy7YckAZKAigytXRZNE9WuSuFrcrA,6028 +torch/include/oneapi/dnnl/dnnl_graph_ocl.hpp,sha256=vZSl091yjQ1Msa_rcWBiOlIzgZMVVkn-9xIca7pUw7g,5682 +torch/include/oneapi/dnnl/dnnl_graph_sycl.h,sha256=KY3JaYUBbJFrpvxZJRoekT4kV2x4vCIT6NPm6hpWxoc,3902 +torch/include/oneapi/dnnl/dnnl_graph_sycl.hpp,sha256=1wJHH9FIBZ0KLZuiaPeahsCQ3gEDfuqUDQq1XZd3u2A,4586 +torch/include/oneapi/dnnl/dnnl_graph_types.h,sha256=cSx0EeE2nfF-ox1jkGUL1ncUu4hbewYzT85XcLtsKnM,17908 +torch/include/oneapi/dnnl/dnnl_ocl.h,sha256=Obblsh4Hdxe9fJt7bO-7She9QSiFouKpqnpThGqstdM,11884 +torch/include/oneapi/dnnl/dnnl_ocl.hpp,sha256=hGcbJAPhWM9wx3ylJYoda2HukFI2WUpF7vxSj9tOluU,15571 +torch/include/oneapi/dnnl/dnnl_ocl_types.h,sha256=gzwyeH3ueR6be7RYPtbl4fDNDBnL4ZQiCPXJ0m2HN_Q,1592 +torch/include/oneapi/dnnl/dnnl_sycl.h,sha256=5DhTc-YRz9XF1b-24xBhlh-rrI-bEn1nvOG7DsXGFdA,8562 +torch/include/oneapi/dnnl/dnnl_sycl.hpp,sha256=rmfH8ZRAUHL00kItlGhG3SUsikCxwGsn8hafe3TN4_k,13094 +torch/include/oneapi/dnnl/dnnl_sycl_types.h,sha256=ZUBXIn5gEby1GqtBY8Lm3oqWlIUGeJPu995dC9gvRvc,1599 +torch/include/oneapi/dnnl/dnnl_threadpool.h,sha256=pbXjaWA3NAlbtYyQCzR7KyrgcRvCBDiWfRPL3Agro7I,4773 +torch/include/oneapi/dnnl/dnnl_threadpool.hpp,sha256=wmTX4Asi90ryCGSnDIF5enXRRqcFoWsOZn_l0uTzIiI,4576 +torch/include/oneapi/dnnl/dnnl_threadpool_iface.hpp,sha256=LStQLhzz7sIQMEj5crZ7LO1K2DwiauVUJCj32wzi_c0,2752 +torch/include/oneapi/dnnl/dnnl_types.h,sha256=QlsVXmR8_lQkm2Bm1wr9D6qaBQD5jQKWGvJQX6FuqGE,101688 +torch/include/oneapi/dnnl/dnnl_ukernel.h,sha256=I1iZhzwsd7p32MtbRrHJwlQYKPAwi8Y45gggjwl06Ys,14037 +torch/include/oneapi/dnnl/dnnl_ukernel.hpp,sha256=O8dcx-ypO-jCsGPrDLZIwZRhgCWeAG7cmN8UDKayW6s,17933 +torch/include/oneapi/dnnl/dnnl_ukernel_types.h,sha256=afoOCsxR7bdbi-nmOMcyHhM0Hgn7xCK3EV-m_EZjWBM,2957 +torch/include/oneapi/dnnl/dnnl_version.h,sha256=NrFHY6pY5TFXvzKqUzZ-FIf5s0OHN5EYVh8sUAL7TMw,1262 +torch/include/oneapi/dnnl/dnnl_version_hash.h,sha256=N3FVjeoQcfKZ1vuSe0KW60cN6uup0HF7KWh7ziJ4dN0,1500 +torch/include/psimd.h,sha256=Zz7YehRUM8gHMbAgELTyZqAx1Vo01wJFgN7-XCnaIWY,45758 +torch/include/pthreadpool.h,sha256=WCxI36to5mSi14ZZsoJ_pNnJsNF6ZJVPBZzzcxeCiio,99582 +torch/include/pybind11/attr.h,sha256=1ems9luMNWI0iQLODXwL6eUwqt2oVo3DIYyMiJf8DsI,26477 +torch/include/pybind11/buffer_info.h,sha256=Yvlrn5CAGjdaW7-I7Yjwnq2MvOfVOcAPbq0tWlmbFGQ,8032 +torch/include/pybind11/cast.h,sha256=-CAlLyvCjg5uHz8TmHJ5OUBieX6gUjkBsCQOmg_3xvU,93493 +torch/include/pybind11/chrono.h,sha256=uPfk6B6By7PVNisjsK9POSemaaK067KBFpj9tqpgbSA,8867 +torch/include/pybind11/common.h,sha256=gPyvFv4QoeZXIScQJFUQ4ZorDCNaj7rmauNmyXKryR4,374 +torch/include/pybind11/complex.h,sha256=Ed9RnnUAzCpNKbsohjlK9wiuVMJqTkt534gmpackgLo,2350 +torch/include/pybind11/conduit/pybind11_conduit_v1.h,sha256=S7fBLOEGYLtDKNH3P0RLK2yA5dsxkLfIzrZJAEH-sT0,4407 +torch/include/pybind11/conduit/pybind11_platform_abi_id.h,sha256=xVIPJrQU5t5Fdkuvym-5Uq7hG3h0UuY9HaI-T4UdcZ0,4245 +torch/include/pybind11/conduit/wrap_include_python_h.h,sha256=0i5dyYjj5r-d_b6_2RM2Al4-bw8pk8FWXSPKpRs53u8,2323 +torch/include/pybind11/critical_section.h,sha256=taXAwGTd6qdN0xQ2fFyx54rwtvcxKHtGohSmj0Horwg,1864 +torch/include/pybind11/detail/class.h,sha256=lWdbLzuLWT13FXg3mpGyk3fdCr7YZV_ElK3nvzN03Xw,31490 +torch/include/pybind11/detail/common.h,sha256=5un8HFrbh9Rwl46zrVuORl0PgNFUpItcKMKtVgXSOR0,60170 +torch/include/pybind11/detail/cpp_conduit.h,sha256=Ky4trNogphamYB3mPZRiUqdj-Kn7N-pj6JFHTD3HIwg,2807 +torch/include/pybind11/detail/descr.h,sha256=Tx_PmytruYwTbHD-P-nFabeVeN36QiiFtexnW_fSqD8,8315 +torch/include/pybind11/detail/dynamic_raw_ptr_cast_if_possible.h,sha256=ohZyWzmgYl6C-hVY-M95L2TlJ-o7FWVLhBUzTUJ6Nh0,1430 +torch/include/pybind11/detail/exception_translation.h,sha256=I0Kcq942M9fBqdDcHxC_xi5IGbz9MpyuYTBBXCDSZzw,2847 +torch/include/pybind11/detail/function_record_pyobject.h,sha256=gLKBGhjjZgJyb_Ko_I9UdA93cCYDovJnR8t_mR8naWk,7992 +torch/include/pybind11/detail/init.h,sha256=wAe1Tcq6zTe2fVTGNwD5dopnGggeH1ITf4lj0k4x82w,23327 +torch/include/pybind11/detail/internals.h,sha256=SX3GA-OVNL_hzYUugiiFybNRycomSkhU0pRhLmhxyB0,32455 +torch/include/pybind11/detail/native_enum_data.h,sha256=UbLzcm8Sr8qSIZDZBe23NDQaCG9mv8nOKm3ihpP3JRU,7887 +torch/include/pybind11/detail/pybind11_namespace_macros.h,sha256=HkDg8EaXiYYo40-GVgLB89tu59MfOSpndItiaKmNVxU,3755 +torch/include/pybind11/detail/struct_smart_holder.h,sha256=TyFmoD7GEJPJL7F0eKSsicZVsiVXVSZUMTwSX-s50_0,15326 +torch/include/pybind11/detail/type_caster_base.h,sha256=HMapjYM-G_s6xwWhVTasZIeRHf0ih80XVCDEO6R_8OA,66447 +torch/include/pybind11/detail/typeid.h,sha256=rNavm6ne8gSq8kf9ktNchKrp4LPzvoDkiSffOjKOHg4,1879 +torch/include/pybind11/detail/using_smart_holder.h,sha256=Sw21_eubAl4zLd5vSJeOQrvAM8Qh8axAziuRt0VDQJ4,794 +torch/include/pybind11/detail/value_and_holder.h,sha256=Jj8DQG_7b4GJWUR68lvnzyMSzVQmR6Yq4raU2MhaYLM,3931 +torch/include/pybind11/eigen.h,sha256=i-T90hORMeIsLSOlZKHIWJnZ1Zr_dsXPlk6HOsUcAsk,570 +torch/include/pybind11/eigen/common.h,sha256=RsVa00Nmbrnu-dUnXHM0RPwElcDIWdFr8GUHAeRXtKM,632 +torch/include/pybind11/eigen/matrix.h,sha256=3l6hl3j4Vf5GYiN5eSUet5Fdjn8kYjYAAb_AaTRI2Lk,32844 +torch/include/pybind11/eigen/tensor.h,sha256=hPyEjJ3wVl9ijofydLvW2KYtSGAY19TWpgrWo43wshc,18914 +torch/include/pybind11/embed.h,sha256=Jiig-eB4X29PHsAETSRW9QcMb_Uhl3dI5DeZ0dbTuD4,12890 +torch/include/pybind11/eval.h,sha256=mKpLWiEMihIwfszdLSZ0AR4Er4K-vrkqQvVts3D3T1Q,5116 +torch/include/pybind11/functional.h,sha256=HtoVWGymQtRk8pXMm1XkeX9iyhxeQwRt-cYWCOybtgw,5598 +torch/include/pybind11/gil.h,sha256=slDhtOmhPnygMDRcn4Ow4SEmbB3_AKOPVVWOLSzWBk8,6997 +torch/include/pybind11/gil_safe_call_once.h,sha256=INXGowIf5Pkrj3Dx2-3k8PTtNLpE1KGj0PeuWAwSEng,4308 +torch/include/pybind11/gil_simple.h,sha256=IaDqehOobdxm9vNti3nrgpYwmWyYhq_2hL5RAfMnU9o,1439 +torch/include/pybind11/iostream.h,sha256=8GiODvb8XDo8l8MbkV9doUMQXKcSJGHIY57HUMUfRMk,9116 +torch/include/pybind11/native_enum.h,sha256=_PeH6t6Nto5n_OzPseMoAdfOYSTcM0gCb3JQG1dzGlg,2731 +torch/include/pybind11/numpy.h,sha256=IE5oasvNX5E-IxRBIjeJ_mWOlu6vUosZX83qS65DBP0,91831 +torch/include/pybind11/operators.h,sha256=ngkisot1ufStnPaqrTa5dDdPVpa9iIfo0xI14HHKjkY,9357 +torch/include/pybind11/options.h,sha256=8ZETtW7ki2zoCP0H75SER_KC3bK8v6pyNmWCtXQAh30,2988 +torch/include/pybind11/pybind11.h,sha256=5Rco-aXAtEbWrr-BAfc9hsJFAA3sL1eCL59pfWRZ7-A,160933 +torch/include/pybind11/pytypes.h,sha256=hZtv1K0eSdK0wp7BSltdmLtI61xM0EfqbsXgNfqxmOo,102452 +torch/include/pybind11/stl.h,sha256=2ShB4sVCmoLQEXwQ4TtpdNlfVNjC8Oe3o_ueXqHwryg,24996 +torch/include/pybind11/stl/filesystem.h,sha256=JRytMDkseOyHEUr5_QmQ5Aq3h0KPuFYm05_MBAv9gYc,4327 +torch/include/pybind11/stl_bind.h,sha256=ueu4PBjrmNE8eeYCO_199Ao7UaxvaekIst4X-mTg3YI,29900 +torch/include/pybind11/subinterpreter.h,sha256=jjLvW8JasrqA0TEN4oVv-F3RdXaX3AV2khapPOMXw9M,11721 +torch/include/pybind11/trampoline_self_life_support.h,sha256=n9awxwSLjSCDn_iI0scbyFj-2dFmQ9VDH7oMQUZ_GoY,2927 +torch/include/pybind11/type_caster_pyobject_ptr.h,sha256=L2FBWjW4QZB8tvWspGtLJ5risRpwuYaTDd0lxacPqF4,2183 +torch/include/pybind11/typing.h,sha256=XbhaqcvlrIHdl95G0S9AeFhVP8jE9PK3UrX9oBevUUo,9145 +torch/include/pybind11/warnings.h,sha256=OCaRnwHWlndKlsTAPeMCeG9k27T19vS3hSq-_CdTQIU,2622 +torch/include/qnnpack_func.h,sha256=TtDF5mhpFcTbzvLUYlHV33KG2BkVZELMfdJa8dKN6eo,4400 +torch/include/tensorpipe/channel/basic/factory.h,sha256=ymil_LBc8GItwEFwHfBAuQNp46y_7T3Ll4qnvJ3kSvg,717 +torch/include/tensorpipe/channel/cma/factory.h,sha256=MwKFZJViyVujWnLo4RyfncH50bJSJx7J_gx7hoMxbnA,713 +torch/include/tensorpipe/channel/context.h,sha256=wNHghkQ66kK3UmlgjQVTjzLN1mFIVDQM_pqp6IwHmP0,3955 +torch/include/tensorpipe/channel/cuda_basic/factory.h,sha256=7bOoST16TKhE2dniZEw2pRLYtkU7PoIKETA_y55fPcs,762 +torch/include/tensorpipe/channel/cuda_gdr/error.h,sha256=fn1e-JAtlHlwMPoSzG8sdzqUufoef6R5traX3Kj47AU,893 +torch/include/tensorpipe/channel/cuda_gdr/factory.h,sha256=jk0MGTTp0yOnjlmo-OYIFB-AqSHRC3kWW_1IvWreRFo,846 +torch/include/tensorpipe/channel/cuda_ipc/factory.h,sha256=DaoYze9SwRtxeKBxUfNocEFetD8iC8I8O91zjHxNpWI,723 +torch/include/tensorpipe/channel/cuda_xth/factory.h,sha256=rU3qvRAN8IdjqKRpjKf-VNN5hp1hus4SwCYwPVTd7FM,723 +torch/include/tensorpipe/channel/error.h,sha256=HpiFTA5aFd7omH5ffWJ_QzRVfmzk4GyE3mEZA_KN3Ho,1032 +torch/include/tensorpipe/channel/mpt/factory.h,sha256=4NYX8S-rdFBR9GDWjgVwF0fuE2k1WhTs7-z8XKYJAiE,900 +torch/include/tensorpipe/channel/xth/factory.h,sha256=oq9LF99KTTer6ZkP2QW2u6a7x4aOBXO3fIR4iTbI0iE,713 +torch/include/tensorpipe/common/buffer.h,sha256=vZhzGJbP7wGrwFQQaY0dnYf_PG70tPEqOaX7D8kvf6Q,3726 +torch/include/tensorpipe/common/cpu_buffer.h,sha256=dWjm6FWlfNu-Gye6ARmHgWxWwOJYEAoHvgo2FwR1AII,695 +torch/include/tensorpipe/common/cuda_buffer.h,sha256=f2pUmSj2Zu6cvVh_f4kOt0S1xwrysdazM7xtvmWuBuc,722 +torch/include/tensorpipe/common/device.h,sha256=qyWkJ6NhKKGPqDSzXdDxzHyUkNlYX3OMM-wxTHlgACs,1903 +torch/include/tensorpipe/common/error.h,sha256=sUmM9D9IZBstFv8aV27ZLAxDQ2d194w39-auvOR3ZCc,3382 +torch/include/tensorpipe/common/optional.h,sha256=Cg8ztuIvxC8ckWyM1giXZhoIsYCFGIopRQ8wQd3f17k,382 +torch/include/tensorpipe/config.h,sha256=bAyHYIbpM0M0TMPiTcaFvcWnFCfDodz_rqGZNi_kZck,605 +torch/include/tensorpipe/config_cuda.h,sha256=2sXFPeoSb2bwgUfnhtcwL9QZJbyRxbbdUWwYLZVb9CA,573 +torch/include/tensorpipe/core/context.h,sha256=eUPwjLU3Voc-puVW8YZ03-utlUh1E2pmLrzpmgV8wZQ,2722 +torch/include/tensorpipe/core/error.h,sha256=czqGb_5R6zz5hcj1BrXZN3D04d8ANt49nbEjpsc5sow,1215 +torch/include/tensorpipe/core/listener.h,sha256=NoKRYhODFRTb2oifps5cxRdYoIe6l5pLqhS_m5xHMFQ,3182 +torch/include/tensorpipe/core/message.h,sha256=NylQKLbpf-IKhyNOzOQBUe2j4pVXjJ9Zl0CBruK8hp0,3139 +torch/include/tensorpipe/core/pipe.h,sha256=FzxISVeZVDxVxgLNyz8N4rcqg-En5JYJK3jMDAdDNOM,3226 +torch/include/tensorpipe/tensorpipe.h,sha256=5qmmXu8SyMq_sREiWZ5GljM-glz-pl4bJSeMaRQLh6U,1703 +torch/include/tensorpipe/tensorpipe_cuda.h,sha256=vHBuinUAmAWIB97PgDRZjclLIJ2-sjmoKkGXwG7iMSg,958 +torch/include/tensorpipe/transport/context.h,sha256=dD4NNcqRe1jvzHAVPYlL_koSf0oC6q3u-ZXIGyqHoQQ,2899 +torch/include/tensorpipe/transport/error.h,sha256=h86Nl4wdVPCBcB11mL630XYqSNzAnavmB2oOU3WwVMw,1173 +torch/include/tensorpipe/transport/ibv/error.h,sha256=iVqKZZU6-Mwxek1SaGohndVYPtL19Ei5fbv8X2b27oU,1174 +torch/include/tensorpipe/transport/ibv/factory.h,sha256=uAAKRzbF6MdDTikr3gCTKXJxGS5sfBjyHcZ4oOCrU9Y,719 +torch/include/tensorpipe/transport/ibv/utility.h,sha256=Yi-VEfXhUG39nK8PMXZT9_jEEO5o3JamDb55besOXw4,823 +torch/include/tensorpipe/transport/shm/factory.h,sha256=HcaaTa5gnGqT0Qz2LlsLttU9DYkLYD2XJ5IE2m_ftfQ,719 +torch/include/tensorpipe/transport/uv/error.h,sha256=QsPw4R0VuiMJPiU4bp2QIYw9A40Uyb-F3J5efL88yaY,970 +torch/include/tensorpipe/transport/uv/factory.h,sha256=D3i_rOQ2F31pDo6iivssNX9q2GlmPcxRQLx3kv-UYTM,717 +torch/include/tensorpipe/transport/uv/utility.h,sha256=5Z_QG64aU9n-yMSMKepKBlUI90htLwHBIcZhJrn92Mg,1303 +torch/include/torch/csrc/CudaIPCTypes.h,sha256=qsqVbMeUqs5FMT02n7hsDxpjxLnBrXnjbn6xgAEmTh4,3651 +torch/include/torch/csrc/DataLoader.h,sha256=0lLHCpmBQ9wnuXLOIH_Ok-t1EQHsX9-IgRPmrYTzfL8,476 +torch/include/torch/csrc/Device.h,sha256=WMNOcYpUnj5oDV05ABo2xLUG4AdmJ14WuB7wRy1rJao,739 +torch/include/torch/csrc/DeviceAccelerator.h,sha256=kK0O_HBZNVYKHKEB7zCW9c2IB03enTyLR_Q427Kp47Y,430 +torch/include/torch/csrc/Dtype.h,sha256=u_dUqcRhIC1kenlr-nnrrz0T0I0XhB_48aKqt7Z3shw,1090 +torch/include/torch/csrc/DynamicTypes.h,sha256=LFNtoH9z3wftoGGyS0lNCgmC4vEAFss3S8-mz3UONWY,1316 +torch/include/torch/csrc/Event.h,sha256=ZaiYri4BhtfCgKsXKUFrnj9s0-LxNd4IvdgDggmV6mc,840 +torch/include/torch/csrc/Exceptions.h,sha256=tUo98K-t1Fv4XuZK2icCWw7TNiWyW3R5_5sIRHME4tI,16256 +torch/include/torch/csrc/Export.h,sha256=cqQQF3wDc-5-WHNjcnuu6fEGxNBDZvqGV7xfx1C5fiM,411 +torch/include/torch/csrc/Generator.h,sha256=m2Bqihf3Sc0bWnI4unsiH51o8YmYhGuCtkzATUDVktg,1310 +torch/include/torch/csrc/Layout.h,sha256=NB5zkGBd7Grp7USYFjF1zYqkXJH6FNEMrwPuCRUdyc8,840 +torch/include/torch/csrc/MemoryFormat.h,sha256=leW7WY12ZeRCjgcbTGhRw1hHEyo5MjP0mzY0B-9Z75s,936 +torch/include/torch/csrc/Module.h,sha256=xPZ6iAnFOpb0kJKIkvSamoVaepAvN9nIp8qDnL1vfpw,355 +torch/include/torch/csrc/PyInterpreter.h,sha256=FwV8NJX2vHVvAgSH9pQKqq66z67jwZjciKcipXFbyjg,647 +torch/include/torch/csrc/PyInterpreterHooks.h,sha256=z8Ur8_WiaAWFDFbfT9zRh6w1FNrLxGwcD_9cTIuch8k,649 +torch/include/torch/csrc/QScheme.h,sha256=X-6XDVHnS9nPbqA3sbH4iulEe38QurprXBRIXF5puRQ,862 +torch/include/torch/csrc/Size.h,sha256=GG35TADU4Fb5LsxfMYG2S1li_9unhRLFcfMiQ41KiLY,730 +torch/include/torch/csrc/Storage.h,sha256=2_h8Ndwai_c3HhWZTeEoBEUVcj6vg63VOS6m-jy9JXY,1701 +torch/include/torch/csrc/StorageMethods.h,sha256=wdY_jcfOWT5HwMR2634Q7id1n4COsYf-hJiJxPUPr_Y,386 +torch/include/torch/csrc/StorageSharing.h,sha256=FDdRsnDvPPBWeBpuSztbEEC3UwZkprcmSIcK0BeC65I,393 +torch/include/torch/csrc/Stream.h,sha256=T_frqzmuEVKzU21HrUoJDUiDCaZNoF1vd5UeF77EmCA,963 +torch/include/torch/csrc/THConcat.h,sha256=t4GpgRJv-bY0ccAE0AgICWQNamhXi3g_EDkJOWVpA6I,945 +torch/include/torch/csrc/THP.h,sha256=5gweWqAcY2BZqsPj1bp0at5P-XYyEWN4MM5L2qRBolo,763 +torch/include/torch/csrc/TypeInfo.h,sha256=r7tPAL63xDEPFo8x6uvHtgODFHXjhXJDzRFAUulNFsU,819 +torch/include/torch/csrc/Types.h,sha256=sEsjWCg6HiCzV10np0kjdl6YWmjcpU0biDbJzBQGQ5c,417 +torch/include/torch/csrc/acc/Module.h,sha256=dFVIrVAvxN9Uk4QybsLW1QQZsincqYLEbXsi9fg8FD4,428 +torch/include/torch/csrc/api/include/torch/all.h,sha256=d_uD0PDavFCyhOyDnOdk5wLAVQE9NI7iSBXqTcfiprk,818 +torch/include/torch/csrc/api/include/torch/arg.h,sha256=BL9aGQ0jh03qjO-ktgWHFxTcukzzbsVZlat4G-lSl6s,1681 +torch/include/torch/csrc/api/include/torch/autograd.h,sha256=c2Z8tuuRXj_DwlwBBYHIjDjaHbDjGcA-L6QqrPNy9Yg,426 +torch/include/torch/csrc/api/include/torch/cuda.h,sha256=iEh4xDeeRRMgYfLIkDNu9Q8crqvLJ2CbNTRj6MsNd1s,987 +torch/include/torch/csrc/api/include/torch/data.h,sha256=LQstdyTLY96Iw3-alKqvTgzfeK7McbI5ae83zGoh41E,551 +torch/include/torch/csrc/api/include/torch/data/dataloader.h,sha256=LKWckxETIn2rdHCN_belMApFqkzs-dQfjvZNvUWrek4,2151 +torch/include/torch/csrc/api/include/torch/data/dataloader/base.h,sha256=jBcbswB5JpHaJJ6gn1opeXmZzLF4QUn4ySBcWEQ1XD0,9475 +torch/include/torch/csrc/api/include/torch/data/dataloader/stateful.h,sha256=EGKQPKMzA7uvIuT8jGGoSALVLggpHhZRXT0VuYEEVXU,2589 +torch/include/torch/csrc/api/include/torch/data/dataloader/stateless.h,sha256=HByZb45MUrHOb0q-aLcj9aPq5JnpPByCMtlu2An3-Sg,3000 +torch/include/torch/csrc/api/include/torch/data/dataloader_options.h,sha256=PbZabfCzPYj6WCCX1vUkybqdiYPdk4uBGwF-h7tPnns,2451 +torch/include/torch/csrc/api/include/torch/data/datasets.h,sha256=tir-STDUqaajiGy5-x_bwDoOc4PwlvmwadAITiB82a8,543 +torch/include/torch/csrc/api/include/torch/data/datasets/base.h,sha256=Cjw4p8Vo-rca61JdGvm9Gvum6eYjV9G65WJtIH6UrcU,3430 +torch/include/torch/csrc/api/include/torch/data/datasets/chunk.h,sha256=s8v5sRT9ZulWFXlp-WNfi9eXKnYO9Bzz2hovW-54Aw0,19430 +torch/include/torch/csrc/api/include/torch/data/datasets/map.h,sha256=DAwZ6ieQyoHNZrkHBaGEmTbYPt4XtDeS4XISGl6OMEM,4326 +torch/include/torch/csrc/api/include/torch/data/datasets/mnist.h,sha256=YRPvU01q1ywMPAQjrjBaNo6KUzSgW9uvuNPEpfDtQaE,1483 +torch/include/torch/csrc/api/include/torch/data/datasets/shared.h,sha256=dLHvC69hde_HijNDTWAJ9jYMTk_OAX9GelxNjr_goMc,2849 +torch/include/torch/csrc/api/include/torch/data/datasets/stateful.h,sha256=ltL0s1nqziimtGBnUpEJffAeInVa2KwWe0jxuTcvcMc,2488 +torch/include/torch/csrc/api/include/torch/data/datasets/tensor.h,sha256=5GaBd5nOL3lvJw5po94pIRkGZBv1OFC_Am8RX7gUXAA,1185 +torch/include/torch/csrc/api/include/torch/data/detail/data_shuttle.h,sha256=crrKlaIf0Y7nvvN1qiU7vNdthyS7wH5YTQEwE7IUE0w,2839 +torch/include/torch/csrc/api/include/torch/data/detail/queue.h,sha256=3PbwSIlInczgyEu4kzi1TNK9RPA2Ipc6QI8QT5PmDhg,2709 +torch/include/torch/csrc/api/include/torch/data/detail/sequencers.h,sha256=408hk6loMFpzN9QVZdbyGsv0mQ-m63m_86sUawhH77U,4694 +torch/include/torch/csrc/api/include/torch/data/example.h,sha256=s_XXP010mJyyotdSSY0lIUsWJgyqCALdfXgiPj9OJNA,1543 +torch/include/torch/csrc/api/include/torch/data/iterator.h,sha256=ztvL0ZIktgE9uNVqS32HjQ0yRwckT0h_A3KS3k1EZQI,5559 +torch/include/torch/csrc/api/include/torch/data/samplers.h,sha256=EJcqg_L6zv1LbFAlpkytbkGTZA6pGGIAbDc1y_jeDUQ,572 +torch/include/torch/csrc/api/include/torch/data/samplers/base.h,sha256=roqyYGAFtgoLNPaVbQo0FlQ-f7Hyeg1IYUJ1eCMV5Xg,1420 +torch/include/torch/csrc/api/include/torch/data/samplers/custom_batch_request.h,sha256=Go9M3TfetFP9DqO_VLJh3Q7IHWz67v7opPOTFLQMUrQ,760 +torch/include/torch/csrc/api/include/torch/data/samplers/distributed.h,sha256=TsFvLnW8X24dJArLqcXrxDvmAAn658vReUk8m2YG56s,4332 +torch/include/torch/csrc/api/include/torch/data/samplers/random.h,sha256=PBl9BsQzkJ4iVIWCHxyWjz33mpP6PKAXNRpIFuL6Bjs,1716 +torch/include/torch/csrc/api/include/torch/data/samplers/sequential.h,sha256=NDKQhS8I9qqAZJj789_KnUBMkIZjOijEUe3uKJcGGR4,1448 +torch/include/torch/csrc/api/include/torch/data/samplers/serialize.h,sha256=nubpsLpxYasgAUxYhHQI3F0aXeVO71iyVftZAQKUMsk,911 +torch/include/torch/csrc/api/include/torch/data/samplers/stream.h,sha256=1KrRV4D5s25Zyk8TG_2hlGA6gjWFm2pESGO52UHw1jg,2227 +torch/include/torch/csrc/api/include/torch/data/transforms.h,sha256=-7PZ3J4Sl8IGwJkCtsa3gwyRIRwAakqh0XYCjLIRr30,476 +torch/include/torch/csrc/api/include/torch/data/transforms/base.h,sha256=zy50TQRf0TfljUSw97T-uBt3JOsqpTX29QOT_Tyydns,1833 +torch/include/torch/csrc/api/include/torch/data/transforms/collate.h,sha256=6ixnhJX-HgRWGYOxW7TlRt0jfTHEJ_f9k9M6SuoXVoM,1317 +torch/include/torch/csrc/api/include/torch/data/transforms/lambda.h,sha256=oRuX2N9-kb-ri0kl-nXUWFA2Z-ndJ03W7p34b78mpL0,1913 +torch/include/torch/csrc/api/include/torch/data/transforms/stack.h,sha256=e8RMeA2_J1Ee0CBnJVpvGC6oR1YUJzZJlhRb7lgvCUU,1628 +torch/include/torch/csrc/api/include/torch/data/transforms/tensor.h,sha256=uT8NLQGuIK7cAnu3hqZ1MwE9P5DZwKjlt6dayOzrMjU,2677 +torch/include/torch/csrc/api/include/torch/data/worker_exception.h,sha256=cReMfofGMxB08LBOyvVQj_u7E4IqMN9W6ZTPgw7jI2M,1373 +torch/include/torch/csrc/api/include/torch/detail/TensorDataContainer.h,sha256=Gf65fgWWtL0d8WZ-sRFJ8Bw95DTSgS6RR0m_Y_9ym7c,13206 +torch/include/torch/csrc/api/include/torch/detail/static.h,sha256=slXrhsYFnm2u6o3WMwIwn2olNGwC68waixn1TGf8Xlw,2347 +torch/include/torch/csrc/api/include/torch/enum.h,sha256=4dLXAs573lv-RY-KA9SoVS-fZjsMPTtpQjsIysif-ZM,7704 +torch/include/torch/csrc/api/include/torch/expanding_array.h,sha256=mJgJDI08qavPtouHQwo9ZzD_Al8oIFnDZdwkKYwXOBI,6927 +torch/include/torch/csrc/api/include/torch/fft.h,sha256=UhE-r-y8A3qsxDf8xSu1zXKXfOkXBMzM1Jy7mMDrzzE,12358 +torch/include/torch/csrc/api/include/torch/imethod.h,sha256=9jhzS5dOBFwnbq-1_a7C-coPinIrrkssmZfagWPpyHA,1994 +torch/include/torch/csrc/api/include/torch/jit.h,sha256=fVuU6JSqhzY-GsSoOZf4VHgKtnkHbGFo8KPeOd9kTEE,1142 +torch/include/torch/csrc/api/include/torch/mps.h,sha256=HrGBge68iJlcBjFVYOCQ0v6V-_kIpAiJ_pwwPAro_E0,1448 +torch/include/torch/csrc/api/include/torch/nativert/ModelRunnerHandle.h,sha256=pM8dpUHKtzlOjWp8Dhf-6GRCt_AFhfdbqDWSTJyy5DI,1580 +torch/include/torch/csrc/api/include/torch/nested.h,sha256=VjlIslutEY3O4RQUM652WoldCaH-nNMa_x0dUd5Hj0o,3027 +torch/include/torch/csrc/api/include/torch/nn.h,sha256=Y6_mVoiLIHyzr4RRQ5LYIA9MxJOmTPQn2tFbvrYXcZk,505 +torch/include/torch/csrc/api/include/torch/nn/cloneable.h,sha256=hYgWBoO_CdiI7TIkKvKQBc8uzgqpx7_o--jrV_cq-sU,4155 +torch/include/torch/csrc/api/include/torch/nn/functional.h,sha256=Yebih9a14mooUre0-lyB7AbuSM3zQI54o5hy8j5qXyA,896 +torch/include/torch/csrc/api/include/torch/nn/functional/activation.h,sha256=AanEgOQHeLicHOwTzZizciTofcBC-Jj3st0BB_JcKRo,30070 +torch/include/torch/csrc/api/include/torch/nn/functional/batchnorm.h,sha256=thPydHdHkAH9GiVKLJ6rBiCnjeyxNh1jKiGG5V5Ntc4,2269 +torch/include/torch/csrc/api/include/torch/nn/functional/conv.h,sha256=Tls9WI9J-iSMRhYGmTkzrPoQZwW7bvk-N7MxU1JfIbM,8373 +torch/include/torch/csrc/api/include/torch/nn/functional/distance.h,sha256=rEP1qHLYXatcmYEMAWWHsjVAumgwF1IGiRBk5c-NbrI,2753 +torch/include/torch/csrc/api/include/torch/nn/functional/dropout.h,sha256=3sJbPO0E8TaifvfoJyTZxTC0KDGi6xydUTVT2r4rSLo,6790 +torch/include/torch/csrc/api/include/torch/nn/functional/embedding.h,sha256=-n4IK-tZPGJmAmdxWyT-gqmzk0clttDfsm-HAhrIciU,6613 +torch/include/torch/csrc/api/include/torch/nn/functional/fold.h,sha256=lLbiMxGeCA1dJP3PIAUVOAfLQAzvM1f-Zw2x5Hp2gjs,2990 +torch/include/torch/csrc/api/include/torch/nn/functional/instancenorm.h,sha256=zIIIVJG2Bzo7FZyED4D0SWWIV6WovO7cbPEOryUohX8,1809 +torch/include/torch/csrc/api/include/torch/nn/functional/linear.h,sha256=ugg28fvjDzk2lgYne1UVM4csQdQLUAGda3ezZuHAkBo,1015 +torch/include/torch/csrc/api/include/torch/nn/functional/loss.h,sha256=Z5tmLVYSQt7Zp7IL_4crceVtUL88rIIBdaXnLhZDT10,32095 +torch/include/torch/csrc/api/include/torch/nn/functional/normalization.h,sha256=s5cxNP6BF1UiO40t8oBV1dIT_h3UiRH_zyjYKFH-cSs,6225 +torch/include/torch/csrc/api/include/torch/nn/functional/padding.h,sha256=5gYWsnKjj7ofCPHkjYPpJ-SynLBJud8EAjaDijQCNTg,1928 +torch/include/torch/csrc/api/include/torch/nn/functional/pixelshuffle.h,sha256=pjM2FRB74Vw2dZWre034KB4LsFPUP-BBoclCa9mYTOU,1547 +torch/include/torch/csrc/api/include/torch/nn/functional/pooling.h,sha256=v_5wJJYLo5db9KDb7SyS0wd7ikfzWwMkb5eVcoCrZqU,35664 +torch/include/torch/csrc/api/include/torch/nn/functional/upsampling.h,sha256=gbZ6HgiDysvMmz9IqOeZQwGT5LB23IYRfxqOVg5dLNU,10982 +torch/include/torch/csrc/api/include/torch/nn/functional/vision.h,sha256=HjnFEeYjeBGo7xEZupgryWYaeprbBRpUS0uDO7HI2cw,3833 +torch/include/torch/csrc/api/include/torch/nn/init.h,sha256=TQ4dS3iK_g_lZSA6H3QRjWHCU8LD3bAMdrV4mT1qv9M,5173 +torch/include/torch/csrc/api/include/torch/nn/module.h,sha256=6ADqnUhYJwgDBuolods-IVU0Ja8kDKltRtmlHdzerAs,27242 +torch/include/torch/csrc/api/include/torch/nn/modules.h,sha256=UZxQ4tAgNJ5MB824qRZUVwVBRscPY9h3fslYV3b0H8Q,1543 +torch/include/torch/csrc/api/include/torch/nn/modules/_functions.h,sha256=lfmzcEi-3tq_8YKuvNcrCtrPC3UKVtoOEy0G3AqlGoY,910 +torch/include/torch/csrc/api/include/torch/nn/modules/activation.h,sha256=03i8F0M9vuuxzcf3BqAzwpn-GEaEBMBdWCvDmDib0TY,30571 +torch/include/torch/csrc/api/include/torch/nn/modules/adaptive.h,sha256=gzmhHN6jsfv38PJUZMB5RnSkYpFufVcy1hH_dTq14IM,3771 +torch/include/torch/csrc/api/include/torch/nn/modules/batchnorm.h,sha256=v2fGhpWLroZu-HNmj0GqBEv3CBeabjj0KHKzI5BnXOo,8479 +torch/include/torch/csrc/api/include/torch/nn/modules/common.h,sha256=jVT0e01SXVxuvGpfauX-sjoeFZJImpZ-HvtEqhg0XzU,4602 +torch/include/torch/csrc/api/include/torch/nn/modules/container/any.h,sha256=F01wtQcvmgLe3S0o3iYNVpeIA4OcXAV8eytwvqA2KBA,13712 +torch/include/torch/csrc/api/include/torch/nn/modules/container/any_module_holder.h,sha256=S-IMNkIUdSLtUpRD9NSSD_JnjDuXgwO5QUNqAzzQSUo,5250 +torch/include/torch/csrc/api/include/torch/nn/modules/container/any_value.h,sha256=4WmDto4MkFYybDAfv95JVHZ3fURvguwuWV890TzyHTU,4392 +torch/include/torch/csrc/api/include/torch/nn/modules/container/functional.h,sha256=EGxMBeX0zHtKL6wWek3-3FGhVMlG6Z1Ku5ncqfZikwo,3597 +torch/include/torch/csrc/api/include/torch/nn/modules/container/moduledict.h,sha256=iBorAWRt_qNevxSW_1N74gS-5VCrz3lY8_mG1fRJH74,8676 +torch/include/torch/csrc/api/include/torch/nn/modules/container/modulelist.h,sha256=J4PxwHw4-05Ok-zWqJAyV6QBnGNjHmGQHQ6PvSuR3TI,9196 +torch/include/torch/csrc/api/include/torch/nn/modules/container/named_any.h,sha256=b3TYOyXtnx1PtDUfP03XZDnhdmj5F5xTcfq7n_9PhSc,2696 +torch/include/torch/csrc/api/include/torch/nn/modules/container/parameterdict.h,sha256=22O3j4R_oC6MJAmE9iWNffx9Pb6atCVK0-GD6yQDJv0,4711 +torch/include/torch/csrc/api/include/torch/nn/modules/container/parameterlist.h,sha256=OVlHpI4NxcpAL4JbpDazpF4wlzd0NU78q7z8kLDOEA8,5812 +torch/include/torch/csrc/api/include/torch/nn/modules/container/sequential.h,sha256=bb75yxy9DuRtuhb_U3aYDp1G-fYPoyaJRm3PeWMm2XU,13988 +torch/include/torch/csrc/api/include/torch/nn/modules/conv.h,sha256=EpGiG2mprFb7oocnUc7N-r1BERpm2tllnqZNCD1RRjA,16513 +torch/include/torch/csrc/api/include/torch/nn/modules/distance.h,sha256=nql9Y19QBsW032P6_RVhGmfq3NYtvKcLIqSjc-5IM8w,3310 +torch/include/torch/csrc/api/include/torch/nn/modules/dropout.h,sha256=Alr2vWdC2vQ_P3vAFwXqvGYoG9ENxpmCQhzoR8sN2Xo,6654 +torch/include/torch/csrc/api/include/torch/nn/modules/embedding.h,sha256=nuJtlx1iXIalox9jWBJViUIvofZ2JgbiAFve00alXMI,6309 +torch/include/torch/csrc/api/include/torch/nn/modules/fold.h,sha256=czUUxQhErjSun_9F0cQ3cbYjTJsWKzyT8eB_AluM1IE,3076 +torch/include/torch/csrc/api/include/torch/nn/modules/instancenorm.h,sha256=brdUwQ8jDoyuLbS5y8KxnM7NJyKnWs6mJ7IukdHLhtU,5768 +torch/include/torch/csrc/api/include/torch/nn/modules/linear.h,sha256=kgodTIxGeZOe7VUHCSMkzCNsAxWqgv8C6JXPyHAXSkE,7739 +torch/include/torch/csrc/api/include/torch/nn/modules/loss.h,sha256=QVsU0rm3VxGWPYcH3uS61ed_oJm6s6yKx_8vDOoZ90I,31230 +torch/include/torch/csrc/api/include/torch/nn/modules/normalization.h,sha256=GSvC2nI74qLTUt18BXboGZlfd-SyispJO09d3_6orLA,7185 +torch/include/torch/csrc/api/include/torch/nn/modules/padding.h,sha256=g3RTBTVaJz-8Byd0QaDp7GCXihIh1Sr7ZhjKjKo_90I,14600 +torch/include/torch/csrc/api/include/torch/nn/modules/pixelshuffle.h,sha256=BDmKhAvvlh3XMIBDJCcLwxKZJGcsGj8A0z6nk2ZE8Sw,3363 +torch/include/torch/csrc/api/include/torch/nn/modules/pooling.h,sha256=uikBjDJjHp3v9Ge5WWgF150NIEGEMtAgBgkZjIBxx-s,29894 +torch/include/torch/csrc/api/include/torch/nn/modules/rnn.h,sha256=BF92Lb2iX0HigEyM0mYpms2Yylp4zteTqnvPoLn48K8,13702 +torch/include/torch/csrc/api/include/torch/nn/modules/transformer.h,sha256=VfumjzK-K-ynNtgVZ3kNR81AX17LnQw--a69DTBZoS8,5576 +torch/include/torch/csrc/api/include/torch/nn/modules/transformercoder.h,sha256=DMi4gf-7bhJMIN2Mmo-hcxT0P3Vs55tQJwTnhAFF6nI,5457 +torch/include/torch/csrc/api/include/torch/nn/modules/transformerlayer.h,sha256=wOu5mDfVXfrC6OWeud6vlp_ERqmsCf2hGmOfxIIeLZ8,6661 +torch/include/torch/csrc/api/include/torch/nn/modules/upsampling.h,sha256=4-8Y2-yx6CYi6FMar3jY3OBgJy3scUubRVieChU0-K8,1873 +torch/include/torch/csrc/api/include/torch/nn/modules/utils.h,sha256=EFeKPqk8CAjyY5KvvaWiDFArSVy_MWuHJ4ryA-QbXMM,1654 +torch/include/torch/csrc/api/include/torch/nn/options.h,sha256=sxG7alldUUJoWyfJrvSTi-0B72_NQXn1Mop-JfTHceo,899 +torch/include/torch/csrc/api/include/torch/nn/options/activation.h,sha256=uhyNnM7WIuoju5zQxwdcz33Ci_kMAfKH98poA5aNa1w,19238 +torch/include/torch/csrc/api/include/torch/nn/options/adaptive.h,sha256=sjOvq3GgknXt1PkX9MFkU7qjOSyGjtQYzuRc6JovEJM,1311 +torch/include/torch/csrc/api/include/torch/nn/options/batchnorm.h,sha256=tSndtNVHDgkdoHeiVnzO8wO6sbb9axlSOgQEEN56AoE,2991 +torch/include/torch/csrc/api/include/torch/nn/options/conv.h,sha256=lQXIbZwhpTt95rCaE0VmHDO7pi0fU8pSlGNUDOFdpdI,13648 +torch/include/torch/csrc/api/include/torch/nn/options/distance.h,sha256=IFWwIffyDJ1Lry00dfccYOt_TOGpK_q4X-Mjbxi6JJA,2243 +torch/include/torch/csrc/api/include/torch/nn/options/dropout.h,sha256=6bJy2Iqz0OXfGSDUX23n2dj2CHTwQ0qfCZQm6nlhx70,3299 +torch/include/torch/csrc/api/include/torch/nn/options/embedding.h,sha256=JUsnVNC9P-Sg9rlUOc-6oKb--hPgKsXrklmfOjBbBGo,11852 +torch/include/torch/csrc/api/include/torch/nn/options/fold.h,sha256=USrVVGHqN90-CLvmnQveb-S2MigbiCwZPfCXSNqW34Q,3167 +torch/include/torch/csrc/api/include/torch/nn/options/instancenorm.h,sha256=tN7AKMLz_tX7u0mvLA2C0JTeATVDiq52g_IwsuXy-pY,2506 +torch/include/torch/csrc/api/include/torch/nn/options/linear.h,sha256=zQ0gjydeaArty4-8PqXP__vE-Sf_JSLS3NtS2QEXJk8,3033 +torch/include/torch/csrc/api/include/torch/nn/options/loss.h,sha256=Mtu5B5AaslcrDi6xa7tn50BCtxaxYAjK-epFgIbU8X4,26909 +torch/include/torch/csrc/api/include/torch/nn/options/normalization.h,sha256=LLi4Snl31dy_KfV4m33KbW5Z2Fe5wjyT3pYZRVi34Uc,5731 +torch/include/torch/csrc/api/include/torch/nn/options/padding.h,sha256=tbU2Tvhs-LnyxWJqtOyyb7TGDHycaGcc-l0zIXWWgQo,7089 +torch/include/torch/csrc/api/include/torch/nn/options/pixelshuffle.h,sha256=m5j2SlTRE7BJijFx1RmSloo70YWbPDTlJVm7L3YPSEo,1886 +torch/include/torch/csrc/api/include/torch/nn/options/pooling.h,sha256=rK_4UzIrHnL1BvJ0oX3P5CP3nSJ5HaYg7BasPFd7cEI,17961 +torch/include/torch/csrc/api/include/torch/nn/options/rnn.h,sha256=5qTRgvl4Cq3gTOcQex8Jc_JSPYCmVBewOmxadwkbGyc,8449 +torch/include/torch/csrc/api/include/torch/nn/options/transformer.h,sha256=bvMimjBCIRN0vmA66AZHdKNBNx8Ib0kIen4SODfFhV0,2068 +torch/include/torch/csrc/api/include/torch/nn/options/transformercoder.h,sha256=KyMKPP5LOnqggsVAEh6yKd6vT-i-hpX3gwHhbOVsIgc,2573 +torch/include/torch/csrc/api/include/torch/nn/options/transformerlayer.h,sha256=vbt-2FO0LdboCf6g0zvQdYi8RHHcefezgSGKeJp0V6w,2313 +torch/include/torch/csrc/api/include/torch/nn/options/upsampling.h,sha256=-rGCWi7NB8ZDo4gC8AyluCZboU0CCCSEs4rhvoG2xoo,4379 +torch/include/torch/csrc/api/include/torch/nn/options/vision.h,sha256=8T2d6HnqLAMZWVrjeHdPrG8hwm8-QG_rJO4araZCx7s,1340 +torch/include/torch/csrc/api/include/torch/nn/parallel/data_parallel.h,sha256=cgf5YDCNfv8pmMAtOIym3De_kc1B1IUz23wW9SpdpUo,11452 +torch/include/torch/csrc/api/include/torch/nn/pimpl-inl.h,sha256=_7gjiWDY99dqBzC7khwLF-7g1OKmsunrtTuTMDLhhro,3478 +torch/include/torch/csrc/api/include/torch/nn/pimpl.h,sha256=wzvuQ_6KJXp5O49oIIGNZMsMcnxQyhqj8JHTLOWGjc4,6869 +torch/include/torch/csrc/api/include/torch/nn/utils.h,sha256=eZzeszpy667cwp6FJ523gX_puBwH-2OMvLeqGQPuFEI,385 +torch/include/torch/csrc/api/include/torch/nn/utils/clip_grad.h,sha256=fOyCc1u-hCi6OekXaJ6ZMqGgpcX-6lNE5Y-j1ZsuLs8,5092 +torch/include/torch/csrc/api/include/torch/nn/utils/convert_parameters.h,sha256=cYaGluG4Ra0DFAqtbjo1bqD4INGuQA7G8OHo2IFuoUg,2630 +torch/include/torch/csrc/api/include/torch/nn/utils/rnn.h,sha256=1fO253Y8XJjMdccTDcKD8jLSCbqHFr0nMnL-hIRK3Aw,13047 +torch/include/torch/csrc/api/include/torch/optim.h,sha256=dKbqqbG93ik70Mqw60l65oGLZl9N-rWQVhI_TQ07K38,648 +torch/include/torch/csrc/api/include/torch/optim/adagrad.h,sha256=nAVNhd73uYBYL2ZeTCLU3vjEbfJgut34RBhdP-g_yWk,3451 +torch/include/torch/csrc/api/include/torch/optim/adam.h,sha256=PJvoyY_mO-OsMuS2gPjKrZnwakTVLmBTEXR3sfvq2e0,3129 +torch/include/torch/csrc/api/include/torch/optim/adamw.h,sha256=KkrUzV21r36SWjpZhgB5t4v5f_PuUO9LhSFoPvpOfd0,3149 +torch/include/torch/csrc/api/include/torch/optim/lbfgs.h,sha256=sTD9wwF_XrGtm-2aUr4yyfq7k92Nc-J_lLwge0vbJaM,3684 +torch/include/torch/csrc/api/include/torch/optim/optimizer.h,sha256=mbpFHQYn2qtqqeG9KEPO8UKO9nt3TUOPUruEiEnnPoQ,16402 +torch/include/torch/csrc/api/include/torch/optim/rmsprop.h,sha256=srRMwObMUT9gcndgw8Zr0JBtkb-8Wjz-iX5SRTXh_8g,3142 +torch/include/torch/csrc/api/include/torch/optim/schedulers/lr_scheduler.h,sha256=-7Z-GPuYZWCf8TjI5J87q0jR9xi3Opa5Xbayvy3MEE0,1373 +torch/include/torch/csrc/api/include/torch/optim/schedulers/reduce_on_plateau_scheduler.h,sha256=-rWp5O4vbLWCWhv8Ng5-Pi6AowyEYnKRBrmpDaYqtE0,1634 +torch/include/torch/csrc/api/include/torch/optim/schedulers/step_lr.h,sha256=AR6ANTwsCyEVEnfg4_uNiBKM_zgNPEXNlCHpsJcSIWM,653 +torch/include/torch/csrc/api/include/torch/optim/serialize.h,sha256=n4P_G-sQqX3RKGRz0tDGoI0IQLqK6T-j9u-Co--1GjY,12927 +torch/include/torch/csrc/api/include/torch/optim/sgd.h,sha256=mbnkMFtvgOYUWSS25noZtCL1U0-7_UBE7t-mxrnYAjU,2876 +torch/include/torch/csrc/api/include/torch/ordered_dict.h,sha256=qbkJy2mDZxqt5rkMLhaVYE-6Mlo5wEidHY8LcgQoJX8,16522 +torch/include/torch/csrc/api/include/torch/python.h,sha256=xNAoE1udfG8OcxH-Ob5ZXcm95WfpEuYaQIywN-mTzbI,10174 +torch/include/torch/csrc/api/include/torch/python/init.h,sha256=DI9siWXORGUKh2T9AousnlS3ss2XOLsXmC2ci2fzdm4,458 +torch/include/torch/csrc/api/include/torch/serialize.h,sha256=cJYBTLdpqP6ZhbZD6XPBwPjNG27lofBkgk_AzjonEZE,5498 +torch/include/torch/csrc/api/include/torch/serialize/archive.h,sha256=bgmYgV3iccrvTK1uih12iixI6ZmuXNlwN2yOiA5JgpE,355 +torch/include/torch/csrc/api/include/torch/serialize/input-archive.h,sha256=94hT0mn8WsZMPf6jDzY0GECw3vOeK_oZTWIiVGDWLys,4209 +torch/include/torch/csrc/api/include/torch/serialize/output-archive.h,sha256=JV4yJh7pHvUejgJT3Ks0ByQv8yy8oBnyM1Krc123AoI,2543 +torch/include/torch/csrc/api/include/torch/serialize/tensor.h,sha256=0zHxiYy2sLUsm0I33QgdUI9wOVukkU_toXnQ4TiAk50,686 +torch/include/torch/csrc/api/include/torch/sparse.h,sha256=Gtx29zexv7Kz-O1GmCSsESZHS9Abtd1INKsc21GmqNM,291 +torch/include/torch/csrc/api/include/torch/special.h,sha256=oH3xzdNxoyHZnP1xjxyG-XISWpdfBrO1-VhzrSHHz_8,38539 +torch/include/torch/csrc/api/include/torch/torch.h,sha256=Pl63ncpcYFuxfHZlmVPbj5_dJ4E6pnRAY8G9QbBd-rw,408 +torch/include/torch/csrc/api/include/torch/types.h,sha256=N00lXLA_-FDh8z1MOBB_ygEfKoxieN3KEOQwy-ACX_4,2642 +torch/include/torch/csrc/api/include/torch/utils.h,sha256=GB1CqG8ZKtQmwCv3I7fOMvdlmngE3RG1FG9XNDRNiIo,3839 +torch/include/torch/csrc/api/include/torch/version.h,sha256=e6jwvr0_ZyDbVQT9pfn4HX7RLbLsPjcFFDnijTct2Fs,292 +torch/include/torch/csrc/api/include/torch/xpu.h,sha256=AMYSk9Y1uA92ZQ16t9uXRxEPjtPo5ScFS3KhaWnVKxM,856 +torch/include/torch/csrc/autograd/FunctionsManual.h,sha256=sRpzMUnlb1quOR75qzEPDHr2rLc0IoLnn2YO3OncQUM,35155 +torch/include/torch/csrc/autograd/InferenceMode.h,sha256=IIsNk73vBgJGreu_P72DNwmvvUzynAuFphP3tEAf-eg,410 +torch/include/torch/csrc/autograd/VariableTypeUtils.h,sha256=N4a7JuuQ3O5pg4dAXvqy6O2c9NI9156cOxoE91Y3nKo,14821 +torch/include/torch/csrc/autograd/anomaly_mode.h,sha256=WUXJ0uJlIF2bueS1UEZUyWkj0GUegPoYcj1aiIn5ah0,1979 +torch/include/torch/csrc/autograd/autograd.h,sha256=FFXo-LL0wK9pD3yC5foBkVHE_w7yK1HulOGpzj91ES4,5563 +torch/include/torch/csrc/autograd/autograd_not_implemented_fallback.h,sha256=tCJe7yOLcR7RL_3UxToohetXaL4IV4-AhirhV0-ICFc,1396 +torch/include/torch/csrc/autograd/cpp_hook.h,sha256=EU4UlP4ia7rBL1MFvTwoQvow-FsZU1q87npXD99hLBc,1213 +torch/include/torch/csrc/autograd/custom_function.h,sha256=-ZSPfYRlRrpjl1FfFsCtlblAp-XLeH1ns_fEZtkOmW0,21502 +torch/include/torch/csrc/autograd/edge.h,sha256=3ExDHbB7rd61-6oJ4xVs3WCAMqBwfFYvMdtfgTHVapo,1869 +torch/include/torch/csrc/autograd/engine.h,sha256=YCx302poAtiqT6wSn4b9_W4zN0UXy7S1PxcRhQ2JLuY,11070 +torch/include/torch/csrc/autograd/forward_grad.h,sha256=EvkQ6z7r0OJOFsZOFj3Fj8HOHCVJJxZCzXjpVVlDOqg,9191 +torch/include/torch/csrc/autograd/function.h,sha256=Li4oaACOgNBgRpOJTWhlxBC1xrgWkO1PuFku1_kMuEc,30970 +torch/include/torch/csrc/autograd/function_hook.h,sha256=_mhR01mCC8NZPe4VcGVDCWBrluemLCvKjxyLZ_lhwEU,2506 +torch/include/torch/csrc/autograd/functions/accumulate_grad.h,sha256=1FZmJFc4b4CYTcbBgr-gygTNu66ej74TQLqWI0ann0Q,15229 +torch/include/torch/csrc/autograd/functions/basic_ops.h,sha256=wbS5BgGMoGw3C5UKodlHsXVq_a1aM-sX7rjv94Sbro8,3995 +torch/include/torch/csrc/autograd/functions/comm.h,sha256=FNF3iPcVZJLNbnnL6O5iK39EULAsbdjPuPYQ6J9zNGQ,1451 +torch/include/torch/csrc/autograd/functions/pybind.h,sha256=eMdYn3nYpRUpn2gG6xMtyCUKr37Tdcd8yPNF-mdo6eU,634 +torch/include/torch/csrc/autograd/functions/tensor.h,sha256=C91kezIt12_arV9fzIP7xqgDjJ40bki3_gva7m_m0xE,7523 +torch/include/torch/csrc/autograd/functions/utils.h,sha256=dVZArpaTbZE27304zREYGBfLyzZ3IY8U-CT2ojRU8vo,3484 +torch/include/torch/csrc/autograd/generated/Functions.h,sha256=pMURf4A3L2CiWq_Pf6eL6TPgPKe_SK0CDJfEJmqJMww,525960 +torch/include/torch/csrc/autograd/generated/VariableType.h,sha256=xu17YMWup0QnMT6IdfIN5ro4ft9ekdDJGMm9cNyQzFk,1759 +torch/include/torch/csrc/autograd/generated/ViewFuncs.h,sha256=yPvzfQIRFWg-9GPcG1JBBIIW6pZ150ihOgyrQdf9ofY,37577 +torch/include/torch/csrc/autograd/generated/python_functions.h,sha256=cqFoR0VPydNMFLNZE1gQLyNgKqVlPl-TXqk8dEdo3Dw,1145 +torch/include/torch/csrc/autograd/generated/python_return_types.h,sha256=CvTXlC3oUkpxoNUGArs7fYZGAoJpOVDO6UDpQfRgjt8,4316 +torch/include/torch/csrc/autograd/generated/variable_factories.h,sha256=EWalUuVIP-5o0htCkXwdvVZccPRIb59rnsN8UDxsPxQ,60271 +torch/include/torch/csrc/autograd/grad_mode.h,sha256=49QHebGFym_CUOczJEtR9N64aq9mdZjk7o9scjMT6XU,464 +torch/include/torch/csrc/autograd/graph_task.h,sha256=bJIY4z8Jo4lYaLBitBbm_yKJGXoVvbPlDpKsOaQgvho,9616 +torch/include/torch/csrc/autograd/input_buffer.h,sha256=5rDYazxZzTRe4xTxL9H0bM4rOVLy-pMEpvPbAtEDwt0,2266 +torch/include/torch/csrc/autograd/input_metadata.h,sha256=8BE5MVxErQyJ-uqmLYTR7gHVMaGeqMQbIbazSUmqeRw,4022 +torch/include/torch/csrc/autograd/jit_decomp_interface.h,sha256=cQ3nucJtKJIUvLQCHo5vz3mem5As_-Yb0XZIiN9FQkM,2082 +torch/include/torch/csrc/autograd/profiler.h,sha256=B4BypHNHlVX9vBCnvfwl4yxogmg-pWrKUhQ3SKceoMU,366 +torch/include/torch/csrc/autograd/profiler_kineto.h,sha256=O5Kbc7dLwhihltugoYjk9x1my-W_t8dGnwwxzE_lBBs,9222 +torch/include/torch/csrc/autograd/profiler_legacy.h,sha256=ybnnQqJ55ZyBMj9-X65QC6LBQGHl1a-aGaAYrdziuro,10956 +torch/include/torch/csrc/autograd/profiler_python.h,sha256=w_uymVzw7mSJshQulmwIbbypnob4XkFFcRdr_pLmKeQ,338 +torch/include/torch/csrc/autograd/python_anomaly_mode.h,sha256=U0rQ5UcXxtf5r-77lmLEb60r3DanCUqnbbt-7NICuts,1440 +torch/include/torch/csrc/autograd/python_autograd.h,sha256=vU-tC6gZ3BY11lhixIp-rPmPB67doBTiX8eRV7_ogp8,677 +torch/include/torch/csrc/autograd/python_cpp_function.h,sha256=Hzc5J7cVNNffCylrbR3L82F-A6hI5v8IAY2pyWIpxpE,5498 +torch/include/torch/csrc/autograd/python_engine.h,sha256=lpZEd6U06q3aGIjXXzMNzuaSsOEUgCqlqjxLUSOzcE0,1510 +torch/include/torch/csrc/autograd/python_enum_tag.h,sha256=DWE7qeWc8UxiHD2TgGvRtkmCNiihI6phVvHIOE76vH8,374 +torch/include/torch/csrc/autograd/python_fft_functions.h,sha256=rCqcUNevR9OzMD2NhlUKC3WWinruFbqKYo13oUS5GLA,389 +torch/include/torch/csrc/autograd/python_function.h,sha256=vc6dO6g4n67DQWM_MNSWXnZsLhcfdYKXthcMvh3hW9k,5759 +torch/include/torch/csrc/autograd/python_hook.h,sha256=OdHmJXgXfy5jadWaKCHREp80MInG0q2WSR39921_v7g,2321 +torch/include/torch/csrc/autograd/python_legacy_variable.h,sha256=tNV-tr0AfucX_7ooFijF1lEWN3Tw5VT0myraeUQcVB8,511 +torch/include/torch/csrc/autograd/python_linalg_functions.h,sha256=8wRQxErZ873uJ5RptU29fcGT-OJYlaIy3B8hW8g6J4o,392 +torch/include/torch/csrc/autograd/python_nested_functions.h,sha256=4ipseJnANXuARoQRTVboHeqICoDEDKzeVTb0Y7PyDMU,462 +torch/include/torch/csrc/autograd/python_nn_functions.h,sha256=EKy9pNl3ybOxBu0JZTA8W9auzbtCnBqjctIS40LZOgI,383 +torch/include/torch/csrc/autograd/python_saved_variable_hooks.h,sha256=83uwA7LUqaCznTyUST-mlRNzBPt0MA8RUH_jI2RgfEg,1333 +torch/include/torch/csrc/autograd/python_sparse_functions.h,sha256=-BmwGly0XH3XsNvJbbSFkj0taRxazvFz-XRoT79BO4o,392 +torch/include/torch/csrc/autograd/python_special_functions.h,sha256=rxQgJ6E7VDMf73wk-DSnSwj5-j4pQoatW--g9h8HYeM,392 +torch/include/torch/csrc/autograd/python_torch_functions.h,sha256=ZgViZedLSuehprslQelMPBEjtTJ_xjLKXNcLLZRb2gY,920 +torch/include/torch/csrc/autograd/python_variable.h,sha256=VWKfw-DXVrDFKbSZZyNrBLrX2sqbcbYD284Bm-DKJtQ,4112 +torch/include/torch/csrc/autograd/python_variable_indexing.h,sha256=EdGDkaPMSb6PRj1dTG2tyiqWpDS3Hrrkrmb4csUejs0,3006 +torch/include/torch/csrc/autograd/record_function_ops.h,sha256=FuLjqyuwTmxPp4xP94yxkLp_XfgHcR_H2txjyFqEVzo,1189 +torch/include/torch/csrc/autograd/saved_variable.h,sha256=R4T7rg-Ou9m00Al9IJTHXDhOJkeUpvwDGFuKnaTry3Q,5418 +torch/include/torch/csrc/autograd/saved_variable_hooks.h,sha256=kwnUYEiaui84rDo3h3apvW3G2PXCk-qlB_7bpY213sE,795 +torch/include/torch/csrc/autograd/symbolic.h,sha256=QkadzHJ6l6JTk1lnaBv6t0IQw5L2YQin5b9HKLlZiuQ,554 +torch/include/torch/csrc/autograd/utils/error_messages.h,sha256=dtBSGmA8U46QFfsSfSDsP1mP_cBjbYWO-zHACUQphdw,749 +torch/include/torch/csrc/autograd/utils/grad_layout_contract.h,sha256=-m8k_tjDVBdAlQ6Hq794tYxSILM6R_NWMvKFZN90-QU,3168 +torch/include/torch/csrc/autograd/utils/lambda_post_hook.h,sha256=LmkQcAnvDxXP6ISLWFrmT_qiTeBJsxQvWplpdDqet7o,1658 +torch/include/torch/csrc/autograd/utils/python_arg_parsing.h,sha256=KYRRUVW4HID-8dXLLK8ejR3Q0dRJxK4ORjqsE_RLPBo,1629 +torch/include/torch/csrc/autograd/utils/warnings.h,sha256=vMaKfEyqoSsAqrwKjRVnypKd4wvWDT6h1FYsxwNttUk,837 +torch/include/torch/csrc/autograd/utils/wrap_outputs.h,sha256=voVEs04uTnBr3Cfg7XYBKrPDuNE7UTuWCYwvspbfeao,4091 +torch/include/torch/csrc/autograd/variable.h,sha256=bxSbmVIGP6YnwtOe4MS4uyxBEYHe8U7rRKqzD5H8i9A,42899 +torch/include/torch/csrc/autograd/variable_info.h,sha256=x_powXIvA0FoDX1ZfprEiC7UknRztkLKAfKWpvCDpx4,862 +torch/include/torch/csrc/copy_utils.h,sha256=L27LKL-iro2NvWObZ5WKTH-1Csq304z0U1N3gW26N6k,1674 +torch/include/torch/csrc/cpu/Module.h,sha256=zGx7rIRNPxhfuZdJND23s8ueyqhYxEANUwyMs6XBJA8,393 +torch/include/torch/csrc/cuda/CUDAPluggableAllocator.h,sha256=EzQBqqHZz3tIGfnL4fUsh8WrqVYBtlHeG9Pwew3ts40,6913 +torch/include/torch/csrc/cuda/Event.h,sha256=KLa6boqJoG359kbJUQ05SEnt_jfDTgSTZJx9zT9b10I,687 +torch/include/torch/csrc/cuda/GdsFile.h,sha256=_sbk174mXjBzcCce8rvrChCh_0Yyk0e-hIM3VzanxX8,446 +torch/include/torch/csrc/cuda/Module.h,sha256=mFzWaRz5WRtYv2VaovHWXTdsafKM3KK1pzQhvtyBawk,738 +torch/include/torch/csrc/cuda/Stream.h,sha256=pOltSeNx7OTGfnVQBh1V3USsJwtMfkvYuBzzBjhfhS8,758 +torch/include/torch/csrc/cuda/THCP.h,sha256=824p3D8G5rEcJ8_mqta6VVtehwmQrX8r2LR630cDITM,477 +torch/include/torch/csrc/cuda/comm.h,sha256=X39M4z7rL1-bsydkOa7ctawby7a9c_rVh_vgZUmtxGE,1768 +torch/include/torch/csrc/cuda/device_set.h,sha256=WphcH70-YDeVx1chXwuIzm6uQZRdDuuGptqpD120_lA,439 +torch/include/torch/csrc/cuda/memory_snapshot.h,sha256=9AqVSIqCiXeFCZpP-bqojEarcxnESUKl1llyIdcML1w,1363 +torch/include/torch/csrc/cuda/nccl.h,sha256=oSBApz0v7VDQuSQvqTdmsWtT3Wee-cCegWX73uOI3ko,6125 +torch/include/torch/csrc/cuda/python_comm.h,sha256=3UpS2iokMzxsYGpslUbV5JnGWopGe3zjdTHOp5IzgT0,425 +torch/include/torch/csrc/cuda/python_nccl.h,sha256=438GoxIfBg-mTsXGvv2hMY9mKpclUEGX9M7ZV-9p978,936 +torch/include/torch/csrc/cuda/utils.h,sha256=wlakPJDSnG_QUpIK6f8LvWw0r7y5OQc4w1hei9exg3Y,469 +torch/include/torch/csrc/distributed/Placement.h,sha256=Ljbn5O4sm3CF4fBgshX4YRiQqzO6vy9Sk6n_ZCrXlaI,2831 +torch/include/torch/csrc/distributed/autograd/autograd.h,sha256=3rkmIlVBwPEIBrhQOqZZMJlMDEuVAersaMGqVP7IA3I,1890 +torch/include/torch/csrc/distributed/autograd/context/container.h,sha256=1QFZ6vefKZmxD8Zc9-EBvjCWkjHhfZ9UxYvrRF2-yAc,6638 +torch/include/torch/csrc/distributed/autograd/context/context.h,sha256=WxXzVZWmQqRmkqVTQbuCCyp64KYya1e0OWdMZlIXJkw,6868 +torch/include/torch/csrc/distributed/autograd/engine/dist_engine.h,sha256=oUAdwQuhseLPr6qTwc-ZzPFi11zfDycWPcgWWwP-mHM,7672 +torch/include/torch/csrc/distributed/autograd/functions/recvrpc_backward.h,sha256=yUhTCVz-MeNdBaXoMDHwEzuL-qmFra3_mVjBjkIOopM,1916 +torch/include/torch/csrc/distributed/autograd/functions/sendrpc_backward.h,sha256=ztfoFVFG6KzHKU4Zz3IzpfkshRcOh-awivGQapGF_i8,1577 +torch/include/torch/csrc/distributed/autograd/python_autograd.h,sha256=F9cPdI4F4m7HPMIs_fBea6QMpX0J5kUyD3EDmHHX1Xg,428 +torch/include/torch/csrc/distributed/autograd/rpc_messages/autograd_metadata.h,sha256=VdgRx9Eaj6Jy3s55wgTpRQVpYoX1WqNuHutcs8cZbj8,954 +torch/include/torch/csrc/distributed/autograd/rpc_messages/cleanup_autograd_context_req.h,sha256=QeJbAf00azhBpdrAUaCBXyrT1KHFng6rIeK37d14594,1090 +torch/include/torch/csrc/distributed/autograd/rpc_messages/cleanup_autograd_context_resp.h,sha256=-CIxYg-7UQ52v3om34tmUM0r1nxWHX35BhQqGAXWnZE,915 +torch/include/torch/csrc/distributed/autograd/rpc_messages/propagate_gradients_req.h,sha256=cCl7oeWz1PbpB80JTecbjugSF3SMstl2tPW1mIrrrK4,1502 +torch/include/torch/csrc/distributed/autograd/rpc_messages/propagate_gradients_resp.h,sha256=G_2iVpZVZSecDjhQA-EslD5swLZWHCJ1k1MB5_UvNnM,1008 +torch/include/torch/csrc/distributed/autograd/rpc_messages/rpc_with_autograd.h,sha256=wWIg4p62dn0Kq4BSNnZz0a-CDv0cs09VBtL1Z53W3aw,3760 +torch/include/torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_req.h,sha256=FZqHLNTHs8uMCbRd_fHTjrf_S5dWJ_MwPHK87W3XiME,2764 +torch/include/torch/csrc/distributed/autograd/rpc_messages/rpc_with_profiling_resp.h,sha256=7VpaqDe6N23Qk-vOyJ2J9VQUzrTpHUiiaU9NvR22T2Q,2724 +torch/include/torch/csrc/distributed/autograd/rpc_messages/rref_backward_req.h,sha256=E9P-Y0xyK68cinxepkUq2uKl73qzssChvNtTYZ4eRaA,1444 +torch/include/torch/csrc/distributed/autograd/rpc_messages/rref_backward_resp.h,sha256=yxE9UFlGQauit2LcWyCkFgy8Uw31HXPRheKJFB7Eqvs,762 +torch/include/torch/csrc/distributed/autograd/utils.h,sha256=oRUE0r6ccxMxc3SPd1ViAexoV2c5xsxmyS4YAp2HWW4,2901 +torch/include/torch/csrc/distributed/c10d/Backend.hpp,sha256=sjOtQ9QVRVI36M7utsn9MNNV6WFMqe-pi-deRCAeZ0Q,17528 +torch/include/torch/csrc/distributed/c10d/Backoff.hpp,sha256=DnLFbWRCa3IhP5E29sSnNpsHBO9oJat1Ek4W1hDjIAw,1305 +torch/include/torch/csrc/distributed/c10d/FakeProcessGroup.hpp,sha256=TWpNTw3rjB4mVQMztkfRxR_vaxlLPYk7kpD7hzpjkRU,8491 +torch/include/torch/csrc/distributed/c10d/FileStore.hpp,sha256=iK7cGCElJxpUzi4-Xv4e1mN_6RSEHvoYYcWvynJw640,1842 +torch/include/torch/csrc/distributed/c10d/FlightRecorder.hpp,sha256=SHskWoCb2xDGd0lMB61b25JXH-uKuCDEI1Y9kAj4uW8,12413 +torch/include/torch/csrc/distributed/c10d/FlightRecorderDetail.hpp,sha256=Ef7lHkiZP21eOtfXnHPnT1d_hiXwlmZWFNx4YA8SG0g,21246 +torch/include/torch/csrc/distributed/c10d/Functional.hpp,sha256=710n8pvHltshQs2MvTdO3atzOT5MmwN4FjNTG4h5OYE,5378 +torch/include/torch/csrc/distributed/c10d/GlooDeviceFactory.hpp,sha256=6MsvopYzTV3oGjWqYoDRXQTbeaNCwg7F2CalioXwYeM,1094 +torch/include/torch/csrc/distributed/c10d/GroupRegistry.hpp,sha256=LFpx87v7dvk11uGm6nS4cxkllN4G6gBHSJL6eouJB44,823 +torch/include/torch/csrc/distributed/c10d/HashStore.hpp,sha256=wLEhb20eHMpuzUCLVCx0q4Ne4bzqa1g_YYupgq9Lgxk,2487 +torch/include/torch/csrc/distributed/c10d/NCCLUtils.hpp,sha256=u9mXKVcJvomhTuKMrEmb-JTZUk_SG74qjV2mlBWvWe0,17594 +torch/include/torch/csrc/distributed/c10d/NanCheck.hpp,sha256=7IMJ1X-zeKEMAb4vsPmHB6lRr7NmtkudyjsJyRn5_iw,562 +torch/include/torch/csrc/distributed/c10d/ParamCommsUtils.hpp,sha256=Q84ODQ1pyLO56l2DRoadR0ZOk5pYJ7Vs9_PcSSPxSXE,11813 +torch/include/torch/csrc/distributed/c10d/PrefixStore.hpp,sha256=6ccfMIhqe3OGCWTut0PKGoU43e6uukpMFTu1qpOgioU,2525 +torch/include/torch/csrc/distributed/c10d/ProcessGroup.hpp,sha256=uf7nnFkqazR_ElHojtv2uXNZPLrDMrF4zHvCZxKW8UY,37565 +torch/include/torch/csrc/distributed/c10d/ProcessGroupGloo.hpp,sha256=Piyr9IV-QoMIsqssrXB2Gy1tjwAGKlfWxqvPspLuC1A,17433 +torch/include/torch/csrc/distributed/c10d/ProcessGroupGlooDetail.hpp,sha256=1ENHPv3UM0Mb8bOLWo8lDvfGFamLVZP7EUf4zg4F4u8,22229 +torch/include/torch/csrc/distributed/c10d/ProcessGroupMPI.hpp,sha256=G0Ejy5zjR7xzIleHvbf5c0JsKa_MJ9lceq4bElwabww,9035 +torch/include/torch/csrc/distributed/c10d/ProcessGroupNCCL.hpp,sha256=FOVweakKcFsFEOuBc26T8CZf1up7opkn8AMvH198BVw,56529 +torch/include/torch/csrc/distributed/c10d/ProcessGroupUCC.hpp,sha256=4Fi1Bzw3nBo9zoU92LvtMDpM52gsf7Pm7fmmQPGFCq4,11518 +torch/include/torch/csrc/distributed/c10d/ProcessGroupWrapper.hpp,sha256=VzObSfFjlZIje1NzmVVrAyfR3ez8CtB6tIpMw2KWy1g,8064 +torch/include/torch/csrc/distributed/c10d/PyProcessGroup.hpp,sha256=g3aEN4SyDwE6NaMJfPxgdwpQtRyIKqwFqTsOb-xzYSQ,11295 +torch/include/torch/csrc/distributed/c10d/RankLocal.hpp,sha256=xpv_I2vmc3Y5k6FjIwb3clFuqrWshbe8hXP1rxFXSHI,3180 +torch/include/torch/csrc/distributed/c10d/Store.hpp,sha256=T_7JBhy99kexlNmRPM99eM6iK-hSiT9QXFgaoi-2SRk,5199 +torch/include/torch/csrc/distributed/c10d/TCPStore.hpp,sha256=nhT0-x9ExKKMQdnRoAGHentZG9tyCWuSE9guGRlglY0,5486 +torch/include/torch/csrc/distributed/c10d/TCPStoreBackend.hpp,sha256=EOhlch6kDsQ0lPQuua6CZqHxO7h8DltbXU8tAkaj6eg,1843 +torch/include/torch/csrc/distributed/c10d/TraceUtils.h,sha256=itfJMHemy2q7yvi_PZlaGcj5vOy3p_U3WqGDh_8BPb8,9340 +torch/include/torch/csrc/distributed/c10d/Types.hpp,sha256=wepA8YoV7H_xAkuD7A4HpuzUCWFDIViuia81C8vg5IY,5340 +torch/include/torch/csrc/distributed/c10d/UCCTracing.hpp,sha256=xzw8ckF13C97z9ZzJZJruxZPW1COo96OuuV2_1RD9DA,2555 +torch/include/torch/csrc/distributed/c10d/UCCUtils.hpp,sha256=1yEkP1SJnvJQCCtLG7jKUqNAYDlV0mjkJcnxcLRYm0c,6614 +torch/include/torch/csrc/distributed/c10d/UnixSockUtils.hpp,sha256=RBTSux9rLp20yj8jF_k5c3rguQBp_o9aWy4aohgsyQU,804 +torch/include/torch/csrc/distributed/c10d/Utils.hpp,sha256=UZuyr33mSYtSKl9SuVfOIO2WH-g6oaGeLBIBUbzCFvU,24041 +torch/include/torch/csrc/distributed/c10d/WinSockUtils.hpp,sha256=-DWh4gPRTjDxX8zfhyE9Bu9YXyhUAGciZjCQsctmlkU,795 +torch/include/torch/csrc/distributed/c10d/Work.hpp,sha256=dhsrOZ7xkcdM_BHJN9bX0D4bKhdWhD3cSuy4Rujy95g,6266 +torch/include/torch/csrc/distributed/c10d/c10d.h,sha256=98aqdnQ7MT6I3vRDCjmrF2zpwhmlbjjebZHyaDdUz3k,420 +torch/include/torch/csrc/distributed/c10d/comm.hpp,sha256=MB5AenuguWwMeIKABucAzpwS7T4lNA85YgSKkxtbxfo,4677 +torch/include/torch/csrc/distributed/c10d/control_collectives/ControlCollectives.hpp,sha256=IbknaNbjdaGHyXPzJVP3JaDMOPB5LF0-0pzHVdeGzMA,1959 +torch/include/torch/csrc/distributed/c10d/control_collectives/StoreCollectives.hpp,sha256=nQjw3mAaL0nwCBTfvvtlM9QO31os-yHHc74gzM_k-eE,2195 +torch/include/torch/csrc/distributed/c10d/control_plane/Handlers.hpp,sha256=vG-xauGcy-87UZbzThZHGzwo6wsw25Z79HWhUFZLkGs,2570 +torch/include/torch/csrc/distributed/c10d/control_plane/WaitCounterHandler.hpp,sha256=XveaRQ5aVYkHecZ2Z4a_3HRejCU0233sxJimG1xtLbk,567 +torch/include/torch/csrc/distributed/c10d/control_plane/WorkerServer.hpp,sha256=xxqg78hobK_Z0E587GIjWeOOsEaiENHKsZsnUeRnlso,902 +torch/include/torch/csrc/distributed/c10d/cuda/CUDAEventCache.hpp,sha256=q4lCWhFi99XcNnFdk9A9YhfUtJlQRfSmr6RykF4HvUs,998 +torch/include/torch/csrc/distributed/c10d/cuda/StreamBlock.hpp,sha256=rbq6FcheBGluDgn9NEU4-nuXN8ItfSgVhy8NBZb5C_k,1117 +torch/include/torch/csrc/distributed/c10d/cuda/utils.hpp,sha256=p-DYuAHZo9VjAbXM5FvxqTXSl47rPhPAiWttX5qNprs,484 +torch/include/torch/csrc/distributed/c10d/debug.h,sha256=ebhH_PUCoe7bakjer98f7OMMkrvIkefj3-jdzbwOb9Y,858 +torch/include/torch/csrc/distributed/c10d/default_comm_hooks.hpp,sha256=i9cBVCXiFfq8wAmCUMFnnG4lhvlWyoeaQP7B5GZotmQ,1999 +torch/include/torch/csrc/distributed/c10d/error.h,sha256=VVxbcGxJKtxoYWiji9uEhc3UwMonbEEQ7W5_zMvj_N0,1544 +torch/include/torch/csrc/distributed/c10d/exception.h,sha256=NDYECXJFZdExOQEW-vt3dXvUntTtgP-pvV-UP2uER8s,1600 +torch/include/torch/csrc/distributed/c10d/logger.hpp,sha256=tkHtVcM3BcBv3b4ZyrvuuYEyIo6m-9W2dEtgEWM2ZQM,6686 +torch/include/torch/csrc/distributed/c10d/logging.h,sha256=kQSwLqHtBQGg00Ow6Be3fg8HPUrHCHcAC7mLU_xF2-I,2084 +torch/include/torch/csrc/distributed/c10d/python_callback_work.hpp,sha256=6Gt8x8CATCK7aJiAFygtXfN9dx7RlINra7pJ77ibpeE,1030 +torch/include/torch/csrc/distributed/c10d/python_comm_hook.h,sha256=KZ2FDYZSeCvTUUdakbwM5TgK7MfS8NVfpundzm1-Zkw,1327 +torch/include/torch/csrc/distributed/c10d/quantization/quantization.h,sha256=Md-HOaJ9-DpVUffMQnBRJ812FMA9jX_JfFVvVhQ_LG4,709 +torch/include/torch/csrc/distributed/c10d/quantization/quantization_gpu.h,sha256=jtpDYbQ9RTOF5AeWmm9zi2e13NCIziUX_rvcHJRUl5E,711 +torch/include/torch/csrc/distributed/c10d/quantization/quantization_utils.h,sha256=S40yJ9Z3sN81eQkoCt1BPol-D41J8RRYI8m9WMmY1LU,1510 +torch/include/torch/csrc/distributed/c10d/reducer.hpp,sha256=BuojgLCbBYFghmafxHRC5hw_Ks1d53uhbKe1NJ78pYk,27238 +torch/include/torch/csrc/distributed/c10d/reducer_timer.hpp,sha256=TD4Lpxk-65H8hy9A1eDI60X8pX4U0zqc_EYLrsJrkC0,2638 +torch/include/torch/csrc/distributed/c10d/sequence_num.hpp,sha256=otUETjq8RGaGOOWeqsCjPAq-MrjJ48rNWSXtlcX563Q,1958 +torch/include/torch/csrc/distributed/c10d/socket.h,sha256=HpAenewC5doWDVOyzdVBkmrO-Q40dpnS_kAPzIMZGxM,2714 +torch/include/torch/csrc/distributed/c10d/socket_fmt.h,sha256=ugcrbys6fYksAsFKph3TptWa_b3MDR-0hqL34nZfgxs,956 +torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemory.hpp,sha256=iU5JmU2pVOjL2YaAiTELvBPBlrH-hhBnA1RZpm6-YaY,5140 +torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryTypes.hpp,sha256=VppiFczA3XetbXUOlgejQ21fnyHfsX_dvphy1_1Qoow,1602 +torch/include/torch/csrc/distributed/c10d/symm_mem/CUDASymmetricMemoryUtils.hpp,sha256=eZar06nd5Y50-NAlhQp7YTKDiKxwIMNrNomNnTWRC-4,3243 +torch/include/torch/csrc/distributed/c10d/symm_mem/DMAConnectivity.hpp,sha256=UA_H2-MUPXBRVjruzND7ieJ_h3MsgrzThydPDEYBu3M,1531 +torch/include/torch/csrc/distributed/c10d/symm_mem/NCCLSymmetricMemory.hpp,sha256=VXjJYvHlqNXyowbk7fOtgJBrKcfrpc1JgcFNbyq_SAk,1668 +torch/include/torch/csrc/distributed/c10d/symm_mem/SymmetricMemory.hpp,sha256=GYC-_x8EB6HiohJnA0yut-8jU41fhrp0qicsgij_FB8,8900 +torch/include/torch/csrc/distributed/c10d/symm_mem/env.hpp,sha256=-4nJOjv1iHpWQ8TU3lI2ZgWBcVTLHxHL9dSt1v2jgEk,669 +torch/include/torch/csrc/distributed/c10d/symm_mem/intra_node_comm.hpp,sha256=j17CAZFd8mBUo3OW-AWpgrLjK1u4UQfyJOQLWJ7iMhk,2639 +torch/include/torch/csrc/distributed/c10d/symm_mem/macros.hpp,sha256=9_ey3ngus8M1TwEJI1M10rGWKg1_xMrIbL33jrjv7cA,1343 +torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_dev_cap.hpp,sha256=Tcty_i0pN42WdaE_UGveAFx2nJJLfTNMr7JOIuquhHU,770 +torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_devcomm_manager.hpp,sha256=IjzTO5pm4Z3b9ROixckKIaxkW2cuCaYKVWsm59prf6I,10765 +torch/include/torch/csrc/distributed/c10d/symm_mem/nccl_extension.hpp,sha256=HfIyhjHlbT0QyUM4OUK3eBtJRjtzU084S46cU1Ze2JA,1271 +torch/include/torch/csrc/distributed/c10d/symm_mem/nvshmem_extension.hpp,sha256=wyX4rg9mxjkFbFE7fjvTjqzyeTeQBGkp26ATzqmWxuc,2682 +torch/include/torch/csrc/distributed/c10d/symm_mem/nvshmem_team_manager.hpp,sha256=eTkIWMsk2lctKuH9b6qKFTC1E8Cs4rWqbVDW4Un7EdI,6722 +torch/include/torch/csrc/distributed/python_placement.h,sha256=87wvbwBCZkwHvXHZ12w_RIc8nwlR5GgxZVNmuZQACLY,422 +torch/include/torch/csrc/distributed/rpc/agent_utils.h,sha256=-WmN5L_KWnHQPDVm8G8lnvyl8CVP3itL_71MxSoD930,1883 +torch/include/torch/csrc/distributed/rpc/message.h,sha256=4J0JYMLIRqpvzFsG8zwjv49Biok9AtJWttFcW60vj8M,7887 +torch/include/torch/csrc/distributed/rpc/metrics/RpcMetricsHandler.h,sha256=qBhyhGPXtUqD16i-UP-zGC-FcH3heCX4lnFEZJcC3fA,1822 +torch/include/torch/csrc/distributed/rpc/profiler/remote_profiler_manager.h,sha256=57nLwfYlLfg_4Q82xj6teBpsEI-VtpQmrHikpXOJsOo,2480 +torch/include/torch/csrc/distributed/rpc/profiler/server_process_global_profiler.h,sha256=4LyG-7M-7dnoS7AxoEJ2A_C6ibEFh7XHYgZFOPsx7u8,4653 +torch/include/torch/csrc/distributed/rpc/py_rref.h,sha256=19R-hMHZZ4wwDCYxoGQZGHBmaDSt4F7imiImR61r8QI,3219 +torch/include/torch/csrc/distributed/rpc/python_call.h,sha256=ib1wexw1L3-FANmVfR5z8k6k8bJWDX_LUDHhhc2SIPc,1065 +torch/include/torch/csrc/distributed/rpc/python_functions.h,sha256=udBVDNxsN7qlVKV6RYNWfga9atMEP64wz0B5h9WfVTk,2511 +torch/include/torch/csrc/distributed/rpc/python_remote_call.h,sha256=L6V5BiyzPhuj0W_-0CfbkUlJzcO2ZBbFW4c9mpMrvuQ,1532 +torch/include/torch/csrc/distributed/rpc/python_resp.h,sha256=Lm4TlIG9M46Qn7uXjSbtGHg5heXeKzFzn_Ccuo6lY_M,875 +torch/include/torch/csrc/distributed/rpc/python_rpc_handler.h,sha256=dWlsGBFT9zHB2Hr_BLxSlvM15xWDKjhNogVyjOH4ur4,5215 +torch/include/torch/csrc/distributed/rpc/request_callback.h,sha256=zJnOxrCIzxFWwuia5XvgvIL--WflFgHSPP337JzGcvk,1480 +torch/include/torch/csrc/distributed/rpc/request_callback_impl.h,sha256=ljQoxGJulsdA1jbSXmcKpAOA_fhbSItRx2wCLMAJ9NI,2336 +torch/include/torch/csrc/distributed/rpc/request_callback_no_python.h,sha256=y5wSLgvzWSU9p-8ZATnzUKbRliw29xWVccbSbpf341k,4158 +torch/include/torch/csrc/distributed/rpc/rpc.h,sha256=AZ88ciAR9Kpv7jQjNO9hWvrgFkTMuyssBTr1HaJjBQQ,418 +torch/include/torch/csrc/distributed/rpc/rpc_agent.h,sha256=mlgNrzyiR68jkMDOuMbFxURwYwGahAEtB3fyZXlNf80,13898 +torch/include/torch/csrc/distributed/rpc/rpc_command_base.h,sha256=v3MU4z486fYnio_EQ5lZmWz_l8dcEe6tD-vCq4jUJMU,932 +torch/include/torch/csrc/distributed/rpc/rref_context.h,sha256=bpPQk3p5wEayzyYiusB27j9zOFfNqeExlDOSToUIAjU,16033 +torch/include/torch/csrc/distributed/rpc/rref_impl.h,sha256=y-j6pPBLe85lWMkwrUSkvCGEczuMD7lnhkZrV1AUfEk,16707 +torch/include/torch/csrc/distributed/rpc/rref_proto.h,sha256=SWbK563pyfM4ErGDXn8csprLuyWMhKQDhBj3YACMZMk,5611 +torch/include/torch/csrc/distributed/rpc/script_call.h,sha256=6YqWqBglQVy9d6ekl0oaficMWr6RpsltysOUMwARXIM,2740 +torch/include/torch/csrc/distributed/rpc/script_remote_call.h,sha256=dmWKw9zMvhAr47cx4wUZfBdcKJ4v-fWOs35CY9KwCCY,1947 +torch/include/torch/csrc/distributed/rpc/script_resp.h,sha256=g-NlV9plQD0iRb4Oe-RdMcUROnVR1KVWxDiW3xi2REU,903 +torch/include/torch/csrc/distributed/rpc/tensorpipe_agent.h,sha256=uku_0SKx9IE_kVuCGltuNcLOVpI8t8u7jbDgwl0Lr_E,17894 +torch/include/torch/csrc/distributed/rpc/tensorpipe_utils.h,sha256=jiKArNodJh9qrWWApKos2cQFYYYwQlpdGXF-P9_0OJg,4995 +torch/include/torch/csrc/distributed/rpc/testing/faulty_tensorpipe_agent.h,sha256=knNysmZUsPfvEdMnV8n5RWVkxXH6yC4QEfBuOWUi79Y,4040 +torch/include/torch/csrc/distributed/rpc/testing/testing.h,sha256=9Z7ubO59UGH4QyrBqzqDuSR4LzuiVJzcV0teEMAa5s0,436 +torch/include/torch/csrc/distributed/rpc/torchscript_functions.h,sha256=h16g5rg-iUPYVNzmfovPMZlq4HJXf-yQZ-QKy4X5Dwo,1910 +torch/include/torch/csrc/distributed/rpc/types.h,sha256=3za2zUBLFSjbS4mMrpcWzfG2xPASKvCJz_mmRIv20jM,2464 +torch/include/torch/csrc/distributed/rpc/unpickled_python_call.h,sha256=lQy-y_rF5i3E1gAh_gH9byjIGMMoWFIi7qFMf2AZjN0,1642 +torch/include/torch/csrc/distributed/rpc/unpickled_python_remote_call.h,sha256=AVUQz2-ZGzECQgZl2Ocpvnw8blFvEIAfnc2KOpMnn_0,1461 +torch/include/torch/csrc/distributed/rpc/utils.h,sha256=2CGx4yhtg62GZHCmuPp2DAN0p-GTGA-6jQ7SpkqMBOM,4044 +torch/include/torch/csrc/dynamo/cache_entry.h,sha256=to35fpvYfvTCGY4eiHWFsq9s2HPZKX6cob4AbaKXM8o,3062 +torch/include/torch/csrc/dynamo/compiled_autograd.h,sha256=MtnR3UA_0TNLb9o0S3FtIxy6Nx2X88gnvrqbdDw70T0,51322 +torch/include/torch/csrc/dynamo/cpp_shim.h,sha256=PikXznGvcUJZ5GSh4D6Fs8Pc28kk6BUcB1WzUsDi22E,611 +torch/include/torch/csrc/dynamo/cpython_defs.h,sha256=Axt2oLc0fsLl5VqxCR-5rRyBLHhSpnt-LF8O-4dTYWA,1174 +torch/include/torch/csrc/dynamo/cpython_includes.h,sha256=d53ADPd7Sk_FzfNKvtgPigmzUWJ6H6AZpzu5SqZGkrs,1934 +torch/include/torch/csrc/dynamo/debug_macros.h,sha256=M6MYIcD0iNODFf9uia3dM4XLxiBs0b54HG4JgFh1zHk,3782 +torch/include/torch/csrc/dynamo/eval_frame.h,sha256=hfnluYBSDER9JJmyFq3j4FZHKmdCVDpI62E1ZsOZmjU,1831 +torch/include/torch/csrc/dynamo/eval_frame_cpp.h,sha256=d7NPXFHNt0w0aNoG4qUz9-Q_9swuszwI0gn0Ok7NGNQ,1430 +torch/include/torch/csrc/dynamo/extra_state.h,sha256=UoFwBzF-jO-p_9qFsFHZAK9fTgJtnRnsdnwb3jRy2zo,6682 +torch/include/torch/csrc/dynamo/framelocals_mapping.h,sha256=1wY8fe28hVm3r1oYVnVtPGWQeJdrim7wP1XuuzpKPPQ,3035 +torch/include/torch/csrc/dynamo/guards.h,sha256=V4h9-BvLVxq1Vp9SBERm1p8py73LM89m22TMsM7LkMk,4202 +torch/include/torch/csrc/dynamo/init.h,sha256=2zowU12pO7ujkPnOise78YFjJq85AilpavNxRc_Fe3g,441 +torch/include/torch/csrc/dynamo/python_compiled_autograd.h,sha256=pnKkQpXyekQfk5B4vtzCgxNWHUjKgwp-4J2lPrOfpyg,469 +torch/include/torch/csrc/dynamo/stackref_bridge.h,sha256=VDxP9T4sWOW0QHU3HAWhvX8LV5dMX0PiXbimAa_Dbgo,618 +torch/include/torch/csrc/dynamo/utils.h,sha256=QhXCgdMkEzSha-lH8EJKO7NnsKg5OVEaWu_gWG8X9do,772 +torch/include/torch/csrc/export/example_upgraders.h,sha256=-mACzEc9kbq3weBPH9VqyGkZPjRdWUt374q6Z3SnuUM,761 +torch/include/torch/csrc/export/pt2_archive_constants.h,sha256=Cyyy9SEmMFmOfl310PkG2rvteCBxXzac27YPyTTcF28,5089 +torch/include/torch/csrc/export/pybind.h,sha256=fI_Sphc3RTmoEAd8w5bUsVyY1ZFBaJTvF6mQ_7ls-Q4,396 +torch/include/torch/csrc/export/upgrader.h,sha256=Hpg96dnZyxdzYIfRUblO2D_dwbmVmsIdQ80VbYLD8qk,5603 +torch/include/torch/csrc/functionalization/Module.h,sha256=5ZEw4c7qn1PtmLTX7QH_bjxWsQdgoAZvpUgUeUwmcHE,1308 +torch/include/torch/csrc/functorch/init.h,sha256=oGmErAJcLTYIjo-LgFLpDI934rG-kDnH22cu0jcoKFE,360 +torch/include/torch/csrc/fx/node.h,sha256=82SLiN-PrEtXKlVBK2sECGpVi7SbFbTpLUPn4_LXcu8,384 +torch/include/torch/csrc/inductor/aoti_eager/kernel_holder.h,sha256=cLB3oS4-UnP5kK5mvWlW1gv0gdXuWRl_LAywuwP05Wc,4995 +torch/include/torch/csrc/inductor/aoti_eager/kernel_meta_info.h,sha256=VRE8Nx4ZG4HsKXOymZ4YZsp6-xDXNxOBPtAnhfUp2hs,6396 +torch/include/torch/csrc/inductor/aoti_include/array_ref.h,sha256=SFBN5XNKtGL8M4EOlZh4wKWDEfeFawPurJrz7Q8TuyQ,555 +torch/include/torch/csrc/inductor/aoti_include/common.h,sha256=pk3x6OttsruIOjeDW1AfSKAzTewv0PV84F_6JhIMK3g,670 +torch/include/torch/csrc/inductor/aoti_include/cpu.h,sha256=cxf5uxsTPXTmzx9Tal01sNbQMLpGIV71OQsLQAMtFzU,386 +torch/include/torch/csrc/inductor/aoti_include/cuda.h,sha256=e5T6P_3cdnOAKaSOpNv8JXHp2pNfCsDGqNNvqjv2kdo,387 +torch/include/torch/csrc/inductor/aoti_include/mps.h,sha256=tZd4iKD9EXjrnZ5y_umoJ1rX3wMCbN1F7stgzTniRSg,386 +torch/include/torch/csrc/inductor/aoti_include/xpu.h,sha256=7zsKLgDuFFnmOpsQ18ZNC28pb8-aKaW2b9TnJeEQDv8,386 +torch/include/torch/csrc/inductor/aoti_package/model_package_loader.h,sha256=SCBejTgGLIgcsYctnHOPNRbyP6FrMl3AZz66IhqWjFs,2129 +torch/include/torch/csrc/inductor/aoti_package/pybind.h,sha256=-QbSC0FRBUlK06XOsjhMkqaoLotj8abv4QS6AgEu4nk,403 +torch/include/torch/csrc/inductor/aoti_runner/model_container_runner.h,sha256=GZdHfXt6I4hI154I5TJO76v09zVEnvI66IJW39lUBUM,5837 +torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_cpu.h,sha256=z6F1PWz7f0CK5aZSNnuy24o_5x-zHLDtkD8AFrQ5l9k,733 +torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_cuda.h,sha256=N-sbQlgKK07d01VyelmnwKoTfiufb-R_hpzeFqHV_U8,1385 +torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_mps.h,sha256=Q8DXtywT5LCB62h1PjXAg2RGvhIafrROAncYfWpBURs,710 +torch/include/torch/csrc/inductor/aoti_runner/model_container_runner_xpu.h,sha256=JrBydhSlEk9MOITI33NQg_Yf2hi8dBepGPUlnIoU64k,1470 +torch/include/torch/csrc/inductor/aoti_runner/pybind.h,sha256=sllydAFs1uLRcUJzr6HhfDKdFR-U6_gdquQWoVLQL30,402 +torch/include/torch/csrc/inductor/aoti_runtime/arrayref_tensor.h,sha256=IuuybEMszMIr5k-OiQmJH5u7wtrQB_hlbfyngJSqMDM,6535 +torch/include/torch/csrc/inductor/aoti_runtime/arrayref_tensor_conversion.h,sha256=nlpKJgBu60U_duWCvEOvv8fj3Re-lM7yAWxJZXQzXHQ,3629 +torch/include/torch/csrc/inductor/aoti_runtime/constant_type.h,sha256=kw7DVntwNZI20Tx0yksFG7tZPHrpJuN7gZmqF09LHl4,781 +torch/include/torch/csrc/inductor/aoti_runtime/device_utils.h,sha256=c1H1BS8sIFDLh3UUhBQRCxz7vHykMXVebnYn9q510LI,2574 +torch/include/torch/csrc/inductor/aoti_runtime/interface.h,sha256=RjIM6ILH9OPm2iOiYNAO3lffIZOamb7xiCbyswbmJaM,14231 +torch/include/torch/csrc/inductor/aoti_runtime/kernel_context_tls.h,sha256=0qROm7sfcfo80fxSatPjddnUxZX0YbFCCkvpBO2qzUo,4229 +torch/include/torch/csrc/inductor/aoti_runtime/mini_array_ref.h,sha256=5DbABsGXNt8dA4tLwQRDqUq4ccRlXDXzwXjdxAbsB7k,4980 +torch/include/torch/csrc/inductor/aoti_runtime/model.h,sha256=0crcjSzh8t3hYaOqlKZZMmIQc8jwTdpkIaf9jKl_V8M,2511 +torch/include/torch/csrc/inductor/aoti_runtime/model_base.h,sha256=Q0I_vw0n5bJi0diej1Oa90EMr_NkK0xo3bguK-cfSAM,32684 +torch/include/torch/csrc/inductor/aoti_runtime/model_container.h,sha256=xkSWN-EdioaP7sxO_XE3fAAhL3SVnWkuz6TefstA1Bo,30143 +torch/include/torch/csrc/inductor/aoti_runtime/scalar_to_tensor.h,sha256=1_eNE1F_kl-XWpwsgc33m1p07_lNK1iasp1cRVQt5NI,1844 +torch/include/torch/csrc/inductor/aoti_runtime/sycl_runtime_wrappers.h,sha256=aPA6XbVCaa3AXHFGeJMfYbwv6umbP9BG_LkJsD7Vuqg,6549 +torch/include/torch/csrc/inductor/aoti_runtime/thread_local.h,sha256=kQpIjL_9JLeu-B5C2_PYsv-BhIh8NueCEVupbBdCJyU,4606 +torch/include/torch/csrc/inductor/aoti_runtime/utils.h,sha256=EWudxCMNKZgkYcw_QIc_XWMFBMl2wsGYLOQe_fhzHD4,14357 +torch/include/torch/csrc/inductor/aoti_runtime/utils_cuda.h,sha256=iP8TDGVTQKcAQH5GscGF21TcRPiUggE887Pl_j_n5i8,2079 +torch/include/torch/csrc/inductor/aoti_runtime/utils_xpu.h,sha256=OxYab5unrLFCiIAbqy7iZf5e0_p6IF-YY7NU1VEEHlk,1969 +torch/include/torch/csrc/inductor/aoti_torch/c/macros.h,sha256=dbD2szZ8zECHPYsKCOC8lT4WJj7CnbkKvlGrsBmUl5w,2321 +torch/include/torch/csrc/inductor/aoti_torch/c/shim.h,sha256=Wj8WdV3v2r872yamiwWg6DlxFh5vdgjBxkhy98mFWGM,26047 +torch/include/torch/csrc/inductor/aoti_torch/c/shim_cpu.h,sha256=tx_K_GPKYJ2_erq5WfWx4Cj_U__Ze28tvea1kvPBrrI,7490 +torch/include/torch/csrc/inductor/aoti_torch/c/shim_deprecated.h,sha256=84fTR6qD0UVNXY9ammn7K1Fp5Rj9mvKHjCkWtIOEmCg,6551 +torch/include/torch/csrc/inductor/aoti_torch/c/shim_mps.h,sha256=1zx5xWfQsLU7HfIVMfzotvFDJmJjAlqWVHcnAavB04Y,3189 +torch/include/torch/csrc/inductor/aoti_torch/c/shim_xpu.h,sha256=Psqiql3DOXCZaXLv7GwoKCcQi4B4Ihg4t49PvH8YXCY,5982 +torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_aten.h,sha256=xG_AhAJbz_mAHwzAMSu0uhcgFD98KUfUrmdHoQ_j6zw,1888 +torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_cpu.h,sha256=E42uHG7p1YzJCCLHr48TnemLPLqKn04a--yWqrBh8WU,32605 +torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_cuda.h,sha256=LQdtcYdUO3vHIW2QCFVpaI1vntYkIhKBh2VKPzR-8Yg,40709 +torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_mps.h,sha256=VCJRZrNRnzpJkpBMYwvGzSnPOleBsv3Pb9M_DM8-o0U,27301 +torch/include/torch/csrc/inductor/aoti_torch/generated/c_shim_xpu.h,sha256=LfngaHjbEr5-mTAuHwlW1fgWE_iFh-0RByjgR3Gxs-A,14788 +torch/include/torch/csrc/inductor/aoti_torch/generated_enum_converters.h,sha256=9Soz7Yzln75OEmVwfVWdy91867ZHgqvFBfS0Qt-Ne8s,4624 +torch/include/torch/csrc/inductor/aoti_torch/mkldnn_tensor.h,sha256=ZdeZ09ReNHprPAM-vXNUXrFL3uupfCbuwVaOQS7ivLs,625 +torch/include/torch/csrc/inductor/aoti_torch/oss_proxy_executor.h,sha256=P_YnlpyPprSTfrsebAjxnoiXW4irLumVbKISuEaaAu4,4893 +torch/include/torch/csrc/inductor/aoti_torch/proxy_executor.h,sha256=qIXJIglBUusFnK8VloIMmFur3Z0j8VHzXjVhcHSWxdU,1111 +torch/include/torch/csrc/inductor/aoti_torch/tensor_converter.h,sha256=WdeIfmMjhpTrqPXmWQ09BKeBKJIZ53UIobYOzefW3nk,1207 +torch/include/torch/csrc/inductor/aoti_torch/utils.h,sha256=NkLCyR_WcJo1l4IHElRaIv9w8QFQKmaTOT8OM0vGa-o,7825 +torch/include/torch/csrc/inductor/array_ref_impl.h,sha256=r424waXg_g_yI3xTAXohMc4yxGWqAEChQCtm6zGbVZw,3260 +torch/include/torch/csrc/inductor/cpp_prefix.h,sha256=x1eojnd8jJsXHUVhiNlHBsByfOoNuIAOcU-KvlGuIdc,39688 +torch/include/torch/csrc/inductor/cpp_wrapper/array_ref.h,sha256=XKMFQCveJevHKVCHmyiepJa2JBZTCMFrU0kE8lJmsnE,554 +torch/include/torch/csrc/inductor/cpp_wrapper/common.h,sha256=3ivk7VbPOyMp_OkR85OwrzrHm9b8zNik_ID5Coepx5Y,2688 +torch/include/torch/csrc/inductor/cpp_wrapper/cpu.h,sha256=DsDdKA8bpepv9yp-59zY3HoQa55zmmdgZtRS4y4JXuU,385 +torch/include/torch/csrc/inductor/cpp_wrapper/cuda.h,sha256=rp2P98HM-V1bA6Jqb-ya1P9EsEbdXX35QVjlCLl1bjw,599 +torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/cpu.h,sha256=rOgmA6ZaUHX1Zvj-Tr3qj8Oc_8BYCMmXPE5GAZhj9ZQ,333 +torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/cuda.h,sha256=2x0n7z6ZZ1t0Bo5xEhcyCnoaH0GsvweIfR0fBN_oXZ0,391 +torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/mps.h,sha256=wRwIY3fjx11KMZZuPGffr3sFJilZ39DDGJ9M4UV5_3w,388 +torch/include/torch/csrc/inductor/cpp_wrapper/device_internal/xpu.h,sha256=SulOHp5rrEkUIyWH0wjhPK1ezDtmiX00u3tAnusq1XQ,457 +torch/include/torch/csrc/inductor/cpp_wrapper/lazy_triton_compile.h,sha256=vI4Qw9bqUQ97NU0Q6xbJnZXg_oKNJaLpMy2lFgU9OFo,7959 +torch/include/torch/csrc/inductor/cpp_wrapper/mps.h,sha256=1BbioeTM6zhVciGbOrNx2SHBn7tw_9-z6bBV-QZ8WqA,385 +torch/include/torch/csrc/inductor/cpp_wrapper/xpu.h,sha256=wStJ3gsDAGB6qcUl_X0ihO9pW6noIi9h04THwMJLEO4,450 +torch/include/torch/csrc/inductor/inductor_ops.h,sha256=lln5wsSTK3alK8h3ooWn5VEjKM4m3IewPDdzw9_Xc0w,1349 +torch/include/torch/csrc/inductor/static_launcher/cuda.h,sha256=qw5_d7GMXzu5vomhEkCOmwGNzX5MjU-uwCAyhFQC9x0,450 +torch/include/torch/csrc/inductor/static_launcher/xpu.h,sha256=ONDk0ocLnY-EgVdgqSnO1OxWjMpK-k-f2eOAwIgW27Y,382 +torch/include/torch/csrc/instruction_counter/Module.h,sha256=Ml9ZG1I-Y6RPLlwR-rnjGTsz77UfWljboJcIEhfjgRw,425 +torch/include/torch/csrc/itt.h,sha256=Hh_KZ49HvlacfiBSGHKCIOxr8lJ7gwwAYpFMGxyvgTI,435 +torch/include/torch/csrc/itt_wrapper.h,sha256=AiJkfyZXHdf_XHd5s6hkf4dxtDYCpUS14iFxmzKuSvA,574 +torch/include/torch/csrc/jit/api/compilation_unit.h,sha256=hUIkZcGmufGNvFGijvHyETh9F7G6Rw7Xh6YJ3t2W-ho,11950 +torch/include/torch/csrc/jit/api/function_impl.h,sha256=4yKX3Evd3u0UwmzAM2Cad1mpScm9JtlNfJE6ri36waY,5995 +torch/include/torch/csrc/jit/api/method.h,sha256=eHbmuglj-sEN05sX-njsaznd-xnGRsV06NyD8CkD6vY,2676 +torch/include/torch/csrc/jit/api/module.h,sha256=A-UjaUfOofIqr71ps4ImjsqEVMWKh5DAcrurB1NUcXA,23751 +torch/include/torch/csrc/jit/api/object.h,sha256=26RTyq72lWmX5CFc5jC3ZTpd5Mc8itPqGyeT0FmaBpo,6337 +torch/include/torch/csrc/jit/backends/backend.h,sha256=Jp65JTXDs4J2ThN_jLXjZZnSL7zK6VghpFSnkqFw9s8,4086 +torch/include/torch/csrc/jit/backends/backend_debug_handler.h,sha256=_QT_zsuAF_buHIuago5qA2sNd8REDXq8ylIALk8gM2s,6586 +torch/include/torch/csrc/jit/backends/backend_debug_info.h,sha256=1XFkmDPhdGlyjLpFLx0x_FIIIcnMw32BZk8gg7S-4j4,2567 +torch/include/torch/csrc/jit/backends/backend_detail.h,sha256=wdFD4DqTtzpkP9gDSiv-7UMirUM6wDDGB7H-HR1QPwY,1333 +torch/include/torch/csrc/jit/backends/backend_exception.h,sha256=zLusJ6DuLALB6i6bVPoGJ5uJXa0AynE1t1wRpEb2ajE,2390 +torch/include/torch/csrc/jit/backends/backend_init.h,sha256=RAFmi0TqARbl-ubYaKTCIL8TNqAGm1qFMPCY3Xx9L30,506 +torch/include/torch/csrc/jit/backends/backend_interface.h,sha256=MX68QIi9QUduProJUDhmZRYrI0v70iFW6T-aIoG_l1o,1413 +torch/include/torch/csrc/jit/backends/backend_preprocess.h,sha256=VqNFjRIHvSl7dzAUIL7By46cYMTxFH2jXRhbK62PHYM,667 +torch/include/torch/csrc/jit/backends/backend_resolver.h,sha256=CYBglxD9roKA4sfXp_54q4ruhOKFu16OWd3Lta8dWGw,506 +torch/include/torch/csrc/jit/backends/coreml/cpp/context.h,sha256=tethBGnuL4WyugUIwuYJ0enleu3p0ayUni8k8fzPNb4,692 +torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLCompiler.h,sha256=reCeVEraZWxNme4d0cnKMjywXuSfMWGE6ZSgTkjEBBc,787 +torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLExecutor.h,sha256=xwIlHZRCswsW26xZ7Ql_fjUTjLBt42ZfNlNrJzf9BuM,668 +torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLFeatureProvider.h,sha256=DkokTYTWbUAlpImXcjhqy7a11H7A5Ty0hJ4fxXmMjtY,604 +torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLModelWrapper.h,sha256=wHqsv6LISy7Xy9B6SkOQQMdSdWBuQIypvi_SHRjpSGA,1176 +torch/include/torch/csrc/jit/backends/coreml/objc/PTMCoreMLTensorSpec.h,sha256=5gsPmrVB6h5S29xNnRgVWAiYrwrKK-51ArG_vJ5J5Mg,902 +torch/include/torch/csrc/jit/backends/xnnpack/compiler/xnn_compiler.h,sha256=ADmstNgU9OkZ5aV11bfLXfUexBpDMyTAIhtzz7o03IY,1036 +torch/include/torch/csrc/jit/backends/xnnpack/executor/xnn_executor.h,sha256=iuAnIUo0qIJ1kY9Uuvj5Nmh2fPzRlvrKcjTdE3EE7Tw,1821 +torch/include/torch/csrc/jit/backends/xnnpack/serialization/serializer.h,sha256=tEkmXQjr76j9w_xXGTJDQf6oqeNXGm11EJ-fhVUq904,3033 +torch/include/torch/csrc/jit/backends/xnnpack/xnnpack_graph_builder.h,sha256=2WQ7XQcCZvv619Q-3Lc4aKqwhJEU9eAeOnSz0-C3oZY,3489 +torch/include/torch/csrc/jit/codegen/cuda/interface.h,sha256=n_0AGEGG0qhy6rCnWJpRk3NwlY7fWXNmxSnms0BrWzY,2118 +torch/include/torch/csrc/jit/codegen/fuser/arg_spec.h,sha256=kEbcFpRmYoos45ISwGCEj7xMeR5yTrTqtRVzkVa-22k,1607 +torch/include/torch/csrc/jit/codegen/fuser/codegen.h,sha256=iUhuS0ExBWQ3dJgW3rXixOfs9wr_Hv4W8_1wBVTB4Jk,999 +torch/include/torch/csrc/jit/codegen/fuser/compiler.h,sha256=XlK31j1-g_YBuW4eNveTsIJhsfLhfS5S75BwtV1-5gY,2088 +torch/include/torch/csrc/jit/codegen/fuser/cpu/fused_kernel.h,sha256=4GspYxscZ_D7qaNhqpb3vOdrp8zOk4kDJrwTftge9SQ,1253 +torch/include/torch/csrc/jit/codegen/fuser/cpu/resource_strings.h,sha256=-mHwdYh9bS9Qp0Rjwjhn1slJ5Dsba7WpaoSyEOnNr1g,2532 +torch/include/torch/csrc/jit/codegen/fuser/cpu/temp_file.h,sha256=fjVmQwk3FzmMZaR-3uVOAWBUnS0Karbb2SRRNFdvuQk,3145 +torch/include/torch/csrc/jit/codegen/fuser/cuda/fused_kernel.h,sha256=QYz6Vg2SZAolMNCXOm97jCKznd5YVGxi1jUf_rKaAfc,1765 +torch/include/torch/csrc/jit/codegen/fuser/cuda/resource_strings.h,sha256=uWSW1FCXRu-v37o5IqTmsXcUkXPlFMJSOrZLKXiueSg,11175 +torch/include/torch/csrc/jit/codegen/fuser/executor.h,sha256=B7r0NLU-F8utHGO8CRIuCjp204lK1Ch8_bgGgTU3Rws,754 +torch/include/torch/csrc/jit/codegen/fuser/fallback.h,sha256=8caRHyCb87UPKCrvhOIiQ4rwfWYeut2LUL2qHWrxIE0,428 +torch/include/torch/csrc/jit/codegen/fuser/fused_kernel.h,sha256=aisFlx4O0Ur4aqwSMKKeexzBtMU7cz_ZnBJJbL9TCkc,3559 +torch/include/torch/csrc/jit/codegen/fuser/interface.h,sha256=S-tgQYiHs8Yq8aQ6bS45VKTkY9dwS1TZmLEM8O-OSEA,1976 +torch/include/torch/csrc/jit/codegen/fuser/kernel_cache.h,sha256=RRzEWG3xtahznaAG4DjhdrbmXY1Jp6Ndmv8T_hopvb0,1248 +torch/include/torch/csrc/jit/codegen/fuser/kernel_spec.h,sha256=BivQtrvupHG9BAEjCZCXrH71hc_zuojVtWLUjZKh_LA,4655 +torch/include/torch/csrc/jit/codegen/fuser/partition_desc.h,sha256=IddHbNKQdgPgY4ESijYgYkoPamv5Rl5l0KQpgyqU5Bk,2009 +torch/include/torch/csrc/jit/codegen/fuser/tensor_desc.h,sha256=ATarPR55x29jHUU-RznxLnDmHiIax1Hwo35ymAAfT1w,2955 +torch/include/torch/csrc/jit/codegen/fuser/tensor_info.h,sha256=dzENGXbxvP1X-sHViShL8mjKYCtZyCCY9TAaI7EWAoE,790 +torch/include/torch/csrc/jit/codegen/onednn/LlgaTensorImpl.h,sha256=PzwwFiuS1WR8uguOyLfScgiTq8rPY4qwyJwlu6mkZNQ,7944 +torch/include/torch/csrc/jit/codegen/onednn/decompose_silu.h,sha256=p_G-mIPLgMY8jSGFSk0x3Lku8FaVjyluOluZUlo1On4,442 +torch/include/torch/csrc/jit/codegen/onednn/defer_size_check.h,sha256=pE1__qzRiyQxN2cJbvvqWFWKjWPYBo790Ledlv98MhY,436 +torch/include/torch/csrc/jit/codegen/onednn/graph_fuser.h,sha256=CR2zC1oIpjsRncDjPGgK6SeQR2ZGVGG1iuja1UKwiqY,1470 +torch/include/torch/csrc/jit/codegen/onednn/graph_helper.h,sha256=LGewDjvOWFT38EcIBi_JOnHV_I9zOQcBwzFwgcGC7yc,2718 +torch/include/torch/csrc/jit/codegen/onednn/guard_shape.h,sha256=j20gjfQhLQYJqfBBe8L4AnpzjfP7gFbSkd1Bu5DxAg8,438 +torch/include/torch/csrc/jit/codegen/onednn/interface.h,sha256=h47yQa2csCXocYQVkvTHsPb7iyJJjEnqQDfvHc3bXdU,1652 +torch/include/torch/csrc/jit/codegen/onednn/kernel.h,sha256=RDURTHKr-SJMSACT3S2pV2RnC-8P4yCU-V3cu6U1b2s,2949 +torch/include/torch/csrc/jit/codegen/onednn/layout_propagation.h,sha256=A-K9-r745pQok43Myo_xHK9kmxsf5sut9U8ZolAykvc,443 +torch/include/torch/csrc/jit/codegen/onednn/operator.h,sha256=vX10EMXfm62FHTWVTb284rQGQF7vpdeiznLJmkFgPb0,4176 +torch/include/torch/csrc/jit/codegen/onednn/prepare_binary.h,sha256=u2XAd1M6J7MiGAqN_QH8ktm0Y_ppCJ89YPAOEmnqFC0,753 +torch/include/torch/csrc/jit/cuda/cuda.h,sha256=dRaA7N_y1gPx6TIxr-DgVimq5hrhKrsXwWrKxRYx_4s,5411 +torch/include/torch/csrc/jit/frontend/builtin_functions.h,sha256=IZBvfhC3BrCaGLAUUVWmqchlP7aNvLonuPTxlBskFXY,469 +torch/include/torch/csrc/jit/frontend/canonicalize_modified_loop.h,sha256=jsuIx-jMYBiRGMHonSgVsKio2zrmbgjVRqeDiqOVxqw,541 +torch/include/torch/csrc/jit/frontend/concrete_module_type.h,sha256=GCxfEpftkhM8Z1sP5vTssOJ3eIg-ZViluiZEPmZkQZk,9291 +torch/include/torch/csrc/jit/frontend/convert_to_ssa.h,sha256=3zWfz9auCqI1QD7FxwBBjL2_UMe6E3p3G4_A90JG0J4,556 +torch/include/torch/csrc/jit/frontend/edit_distance.h,sha256=kBIcPiahPtaKkOxWxALNLIi6fjke52sFJqNoII-lmbE,483 +torch/include/torch/csrc/jit/frontend/error_report.h,sha256=NUq0IuGQRclkqoYl683WnEpv5aIU79t2XtQzBMR_sEI,2494 +torch/include/torch/csrc/jit/frontend/exit_transforms.h,sha256=AHUnyTGs7X9qs68WAce4PaeNaYnZU-ZiUXq9N6n4p-c,447 +torch/include/torch/csrc/jit/frontend/function_schema_parser.h,sha256=A2AeqcgaCJv99Meo_KUk7IdUEvd09tdJkDExivRiij8,1058 +torch/include/torch/csrc/jit/frontend/inline_loop_condition.h,sha256=4tqiJQvvtb7knAxeGN7_2XV47pJwoTZO2qLqPpw4GS0,581 +torch/include/torch/csrc/jit/frontend/ir_emitter.h,sha256=0sEUfd9NIyIK_zS4st3_79jowETO89MCwy0-fPnIRCU,770 +torch/include/torch/csrc/jit/frontend/lexer.h,sha256=ybWb6X5xGC4OwghexEU_fVM_tnGF5QnmPGbKfTq4seU,19212 +torch/include/torch/csrc/jit/frontend/mini_environment.h,sha256=u87AA6rHKGLu_MNWaKLtkIkjb_IIqytdpI24T5Qxeug,1629 +torch/include/torch/csrc/jit/frontend/name_mangler.h,sha256=PRM_e4MlsIKA1wowjkp5vHEGNBTOaB0X5K85wHKKszA,884 +torch/include/torch/csrc/jit/frontend/parse_string_literal.h,sha256=Fznb2iPJ4zUS5wUiMAjIY6q4O91g_VyE93DL5sdDbuw,2548 +torch/include/torch/csrc/jit/frontend/parser.h,sha256=_VFpRbn8WXPLBhYDz70syzCXeSXU0SRYAfe8ZaKvvhI,903 +torch/include/torch/csrc/jit/frontend/parser_constants.h,sha256=0oL1Mg44XC6hXBAUVj99yv1wylrXakght95EY8u7aHw,405 +torch/include/torch/csrc/jit/frontend/resolver.h,sha256=nHTt6zBkmJoMOlSKazSzRYwWajRE-DE-d50_vInSpOQ,2211 +torch/include/torch/csrc/jit/frontend/schema_matching.h,sha256=5duuKpN_Iu8dwscoEsjSlxNmAshtYyyLEqMtuQEYwrI,2362 +torch/include/torch/csrc/jit/frontend/schema_type_parser.h,sha256=Vq5C0npJSoWJYb-qJ2KlbNIgqrvLphPdOUKpM2Ruix4,1657 +torch/include/torch/csrc/jit/frontend/script_type_parser.h,sha256=J0NaC5d95lNulX8d1gN7A94c2U44hrh0OI85Oo83KTI,1834 +torch/include/torch/csrc/jit/frontend/source_range.h,sha256=2RKUq0fZGS76QMDILHOltvJjxLE1dGS5C7IWc8oRJ0M,17167 +torch/include/torch/csrc/jit/frontend/source_ref.h,sha256=62PZejPu4EM0IIIzvCc0CLbi-M--HGcAfpGXtgKvYVo,1542 +torch/include/torch/csrc/jit/frontend/strtod.h,sha256=zBQEJYtFWY-Fvdb6nsS13gggDCnZo8UFQDxv90QZv4M,470 +torch/include/torch/csrc/jit/frontend/sugared_value.h,sha256=I4HFJFg5JjLT3lGlE1M5pOKH2ROrg9ArDwIoqANJbRE,28516 +torch/include/torch/csrc/jit/frontend/tracer.h,sha256=-wUprSpD2qxaVJuMOsYFeOpZH15_GqWtJGREqucHsI0,13055 +torch/include/torch/csrc/jit/frontend/tree.h,sha256=INdyEOZij5mHcPQ0dA7weulkHlxPpH4EfcLipDEUDdg,6807 +torch/include/torch/csrc/jit/frontend/tree_views.h,sha256=jUVuM_ORa1Fw2JpabQ437wr9bUTbnKjrBJ1SHWRAdXg,37454 +torch/include/torch/csrc/jit/frontend/versioned_symbols.h,sha256=QQs29ujosMvluxrGO2hAW7BFNDKEk2d9F_nTO3Lmmig,849 +torch/include/torch/csrc/jit/ir/alias_analysis.h,sha256=aW6ia8q2mJjNqzru3uhzz8YWxa2EUL69KDjhDVbUTvM,14545 +torch/include/torch/csrc/jit/ir/attributes.h,sha256=aI-bBxIrJVjRBKMvsHbqeQNPVLBlVKzqMMZ2gNNacsA,5048 +torch/include/torch/csrc/jit/ir/constants.h,sha256=6fLeQ2WAKA5g2gw8Zs-R7XkAXkXk4YXgYUcDwUohf5Y,2262 +torch/include/torch/csrc/jit/ir/graph_node_list.h,sha256=rgOcMU9i8Cl87ZuhS0sSSjeKVHx1lDbAercWE9y7NS0,6608 +torch/include/torch/csrc/jit/ir/graph_utils.h,sha256=f7t71HGhSNU94EvgZvL67Ns4uelB56jQHlK9Nfh0xW0,758 +torch/include/torch/csrc/jit/ir/ir.h,sha256=LUSxB6I-STtgTmW3jX54Kao3GxSjN9o4U8Kw_dWhkzY,54043 +torch/include/torch/csrc/jit/ir/ir_views.h,sha256=Wb8vo6TOqKNyz9JpfZhOxT_EtBCwK5eI2XZwqp8JsCI,4878 +torch/include/torch/csrc/jit/ir/irparser.h,sha256=QR383rzKoILJH2nvRV99BfSLecPlQ0JmNhRUUkdSBY4,1354 +torch/include/torch/csrc/jit/ir/named_value.h,sha256=82WrC4LrBN4TKGqdHMEXgO-n50yfYiUtR2h-Bg30Ny4,2650 +torch/include/torch/csrc/jit/ir/node_hashing.h,sha256=DR0v95_n2vYSlFBwPiJ8ITtCCexa-Yl8QFI3stx6mjM,519 +torch/include/torch/csrc/jit/ir/scope.h,sha256=2kze-TeAMP3hBFE-kK8jyJSYeHR9OeTY7UhFG262FD4,7412 +torch/include/torch/csrc/jit/ir/subgraph_matcher.h,sha256=WImU1qLKIOT_5C7nMn58JP_v5ApuyrZN-uSMEysC4BU,3381 +torch/include/torch/csrc/jit/ir/type_hashing.h,sha256=8TAouNEGWvMPma9IVod1oqIYHuIggQaFsq8_Zq71yYQ,688 +torch/include/torch/csrc/jit/jit_log.h,sha256=syz7hm-O9pgJaPrw4-5q9GNwPCqfNoR6-idW4okMeOg,5060 +torch/include/torch/csrc/jit/jit_opt_limit.h,sha256=l-n0og_7nXcLMg_ae44PbLXrziAZY1ISo3PmTzCVsww,1635 +torch/include/torch/csrc/jit/mobile/code.h,sha256=E-rv5f8zS38MQixts0_MOe7Ur3sGHLM0Hih0rMcgj5I,1332 +torch/include/torch/csrc/jit/mobile/compatibility/backport.h,sha256=my50F4l06bIP3ScMBpCyHLLkzKVI6R5Teno6DO5MVGw,893 +torch/include/torch/csrc/jit/mobile/compatibility/backport_manager.h,sha256=RiXxwVhpwUkte-nxkn5-ChdWNlo21NliH1XFW25R-Rk,1413 +torch/include/torch/csrc/jit/mobile/compatibility/model_compatibility.h,sha256=dywbWpJUEvz_CiepxNHFc0oATDKtWOJwblheeOvIaFU,3864 +torch/include/torch/csrc/jit/mobile/compatibility/runtime_compatibility.h,sha256=z4MRUqp0Zv4ldMRb6BODecdKaoPZvB8MNeMtwces6P8,1450 +torch/include/torch/csrc/jit/mobile/debug_info.h,sha256=zWiX3q7hHzoNvtlkyeqhm5wi5o-VPa9W41Cb7W3DQ74,2460 +torch/include/torch/csrc/jit/mobile/file_format.h,sha256=dhaybsAFdjlU06j0_UY4SCiV4WUl6lJhG_q-akFl6oo,6925 +torch/include/torch/csrc/jit/mobile/flatbuffer_loader.h,sha256=ZemcyJvz1xrViIzqyU9FOpEFKSfs0Zj5KLZFizGbwOk,5028 +torch/include/torch/csrc/jit/mobile/frame.h,sha256=WiI8ygemRG2Cs8kC-XhXSsUEsW4_OPVg7fxVpAB58Sk,1065 +torch/include/torch/csrc/jit/mobile/function.h,sha256=fs89swfvXdN8o1rkHr1YJ7WplNp2GfnzoUbp8qld1lE,3195 +torch/include/torch/csrc/jit/mobile/import.h,sha256=yr-EklN5mCQAZ4UFDu17e1r4Kjje5dpLkhOv3LxYRWk,4094 +torch/include/torch/csrc/jit/mobile/import_data.h,sha256=kqh0Z_dfBXiDhA4cVAURyAn5tmKpjAz71FpwcIj8dzA,1249 +torch/include/torch/csrc/jit/mobile/import_export_common.h,sha256=dLR011fk_eI2oPyJ_s-rSPt0EqhKUHpYiDlBXjfQejE,731 +torch/include/torch/csrc/jit/mobile/interpreter.h,sha256=JUerPsUZ5IqflpfMtbtPOCIW58lD2ZwmAyVXmJqnY5s,901 +torch/include/torch/csrc/jit/mobile/method.h,sha256=oYavtszbUZH2NAe0pM7wkKQtAyr1edu6qwn5PiWq7mo,1078 +torch/include/torch/csrc/jit/mobile/model_tracer/BuildFeatureTracer.h,sha256=K-OXlFb1vxQ9yYmUh9iGcluYaMWrrFpvTyaryo9wX04,1215 +torch/include/torch/csrc/jit/mobile/model_tracer/CustomClassTracer.h,sha256=ZFqxc7Hukxo_Mb4g7ikCxcPvRObAqcfqUxPG-IvQ4_E,1213 +torch/include/torch/csrc/jit/mobile/model_tracer/KernelDTypeTracer.h,sha256=kP9reVHEFkMFMnMr1LR_6CD-352jdMl9dDlCAci3tQM,1466 +torch/include/torch/csrc/jit/mobile/model_tracer/MobileModelRunner.h,sha256=Xfd9FoGNn9_R3oAZaDN4YC35lQzutLYIHIgJyuGztvg,5340 +torch/include/torch/csrc/jit/mobile/model_tracer/OperatorCallTracer.h,sha256=KFHgSJ8NOBDwzrswmf6dNCCW2L2SF8OoqG42NnryF4M,1172 +torch/include/torch/csrc/jit/mobile/model_tracer/TensorUtils.h,sha256=9KJss0Al81LTcQKjbUZJclVKvIqn-0OYrF8rQN2XotA,652 +torch/include/torch/csrc/jit/mobile/model_tracer/TracerRunner.h,sha256=wj-4TNfuyr4Q0yXX5AGbaEI3C7GaGc5o7RTorH3gvHE,1361 +torch/include/torch/csrc/jit/mobile/module.h,sha256=AjMKsMnliGQWiNEY2q6rBJoGg1xq6IvMF5y_QC6_v84,6175 +torch/include/torch/csrc/jit/mobile/nnc/aot_compiler.h,sha256=4QtJt2yncfiqbmu-TdqZXukUJSUoyG299S3SyHwuYMY,884 +torch/include/torch/csrc/jit/mobile/nnc/context.h,sha256=gFDTG6b9go-Oc2gUseF8DNsfvFtdVTlQh1nrzn1PAlA,6870 +torch/include/torch/csrc/jit/mobile/nnc/registry.h,sha256=q_Ba_oX7lIXuiH0Eu6mXvLT1zFEW8X2BhAhGW7EYFv0,1414 +torch/include/torch/csrc/jit/mobile/observer.h,sha256=5nKATYXvy7peg2rKDFro8wGL76rlx9hIJoCURvo3xLo,4080 +torch/include/torch/csrc/jit/mobile/parse_bytecode.h,sha256=o27oJpF6gCfKzSi8p8m5Gu72LfZXckuzJ6lweJLnCuU,994 +torch/include/torch/csrc/jit/mobile/parse_operators.h,sha256=av7eaS1-HSzMpE57EnjVBcBrr0AzWPyCrKpp3UJW1tU,963 +torch/include/torch/csrc/jit/mobile/prim_ops_registery.h,sha256=nhqJNmc8RtbuOJbMfog5XBknbBnw8lukCYsEnGFVnS4,849 +torch/include/torch/csrc/jit/mobile/profiler_edge.h,sha256=K7vw8Lu4GZ22V9PcAVj31heshmiSJQvtZP-k3WYjuhk,4742 +torch/include/torch/csrc/jit/mobile/promoted_prim_ops.h,sha256=3lp0XyINie-YHJQbH9bAr1dqcS2L7e5LXcIiLWwDXBA,1297 +torch/include/torch/csrc/jit/mobile/quantization.h,sha256=5yAbohsf1X8ysN-Fdi74gDpcSn2M4JeqBkR-SF1dSYM,1492 +torch/include/torch/csrc/jit/mobile/register_ops_common_utils.h,sha256=qPbFCTpOJiEo2mbQFMdhgA5VYHEg6pTZVntHoTO0Hzs,1941 +torch/include/torch/csrc/jit/mobile/train/export_data.h,sha256=VjByQpzxac0K4d5RSjDZEf7tJXkitv728Sr4ZhxR_jI,1848 +torch/include/torch/csrc/jit/mobile/train/optim/sgd.h,sha256=LQc_NMk3A0J43eJVLbrGme_tf2lonFAUz5xTH_3XoOM,4585 +torch/include/torch/csrc/jit/mobile/train/random.h,sha256=l8ExBi_8l5sQInNEL4XyvgIh5eUWH6H_xci33saSzmc,1773 +torch/include/torch/csrc/jit/mobile/train/sequential.h,sha256=SN3bNiMvPnwhSvmrp9k__pvHli_Xzc2QfFQgdsgN8qk,1478 +torch/include/torch/csrc/jit/mobile/type_parser.h,sha256=sapP6ZQle_VaZCI1hieMjxZTz5kGQ5-E6_1dju27KkI,1697 +torch/include/torch/csrc/jit/mobile/upgrader_mobile.h,sha256=CfF3MVWCIOkQWEbeJj25Ez2jhxR-OzplLf65Oujqwek,1159 +torch/include/torch/csrc/jit/operator_upgraders/upgraders.h,sha256=tBLNg_Aw07Vox9rPhruU1aoMdpadPMCuwdsyDnbuNOc,1638 +torch/include/torch/csrc/jit/operator_upgraders/upgraders_entry.h,sha256=7oBZBxLSNsQZ7E0HLUL2WPbEBtxElBydUtEjAOEvuEQ,776 +torch/include/torch/csrc/jit/operator_upgraders/utils.h,sha256=oD6TyMldIa8D9ygMZ7fmgjOtQTl2hK6Gsmi00gWy-po,1947 +torch/include/torch/csrc/jit/operator_upgraders/version_map.h,sha256=Oq0QWD-3F83cJ1_VmSkv_C8zJW-7Ndwpdfuvwe7fFaI,1141 +torch/include/torch/csrc/jit/passes/add_if_then_else.h,sha256=f0ShcVoByKs00Ld2YqZYXvquWhtqIikV2zornZZkhIc,417 +torch/include/torch/csrc/jit/passes/annotate_warns.h,sha256=02JWzxMLmy_BfOtgIgCmMeA6MMJWOqh1xm3VH4wAdi8,421 +torch/include/torch/csrc/jit/passes/autocast.h,sha256=_qPHpJ7hB1wzcvYWEuzEuEbmD69eteFsNq2IAVQhTEM,496 +torch/include/torch/csrc/jit/passes/bailout_graph.h,sha256=-E0TEJ6zyidaxR7FZTRymeOpVOlGqF2MsJPiQpaZXcs,1344 +torch/include/torch/csrc/jit/passes/batch_mm.h,sha256=DyxWCGNbD3UeqUbNxlUb4mGeKD1rayabjZcO3SVClZ0,385 +torch/include/torch/csrc/jit/passes/canonicalize.h,sha256=v7t65TCLRkIsGRnDhSa2Vlk65tMBGpJ_kMwERKX7NNw,721 +torch/include/torch/csrc/jit/passes/canonicalize_graph_fuser_ops.h,sha256=xN1GwxxzunnJ8gd-M6MtsKR4iy4_zPRWpCyme7EWCV4,399 +torch/include/torch/csrc/jit/passes/check_strict_fusion.h,sha256=MsxAQiu72jflrxZqFGFv8u3-M98BokJpck5ZruquxgY,420 +torch/include/torch/csrc/jit/passes/clear_profiling.h,sha256=5tAIUwLoKnKFfTK-jTGE622dvzxCuguLeqSElDg3wyI,722 +torch/include/torch/csrc/jit/passes/clear_undefinedness.h,sha256=jnc_bdQUATcOY3Jz7fMHoYaBX71k3IIZfqB_NYI8UEs,1104 +torch/include/torch/csrc/jit/passes/common_subexpression_elimination.h,sha256=rofGn5UT4KhFColNDpbnqiXeL5goOP4U0NQutKRNOUU,416 +torch/include/torch/csrc/jit/passes/concat_opt.h,sha256=YuSp1PFcUmxy_dLmEhVlhM_YcoFb2ZA9tQXIJaP0lYo,779 +torch/include/torch/csrc/jit/passes/constant_pooling.h,sha256=En3nfQkAGG3_470VppGHE65MINttbQgqMV2H_Tw6IfQ,399 +torch/include/torch/csrc/jit/passes/constant_propagation.h,sha256=p35UK_K-X2j2XImbtpgrGFpUNO4SXkRxf2m4htvJ_Rw,1543 +torch/include/torch/csrc/jit/passes/create_autodiff_subgraphs.h,sha256=_mlHdBQFWeKNgLAXa2DuUkBNWDekvjaVYYkN8KiTALA,767 +torch/include/torch/csrc/jit/passes/create_functional_graphs.h,sha256=5_fBFCv2Fa6Ns-7LhXrmZBtOL7hlKhIjjTODem5uzQ0,538 +torch/include/torch/csrc/jit/passes/dbr_quantization/remove_redundant_aliases.h,sha256=cMY-yjhfl6j62rOMFYw2usBgHUqA_pemFtc9tjsP4Co,594 +torch/include/torch/csrc/jit/passes/dead_code_elimination.h,sha256=Egpe15Aw5mI7eETml6LIPzIMmhTsEi6m0swVL5CpeGg,1813 +torch/include/torch/csrc/jit/passes/decompose_ops.h,sha256=B9qYNXZWlkJI5Q-RqSPM0UymrH7_748Rd5IWKtwTVQs,390 +torch/include/torch/csrc/jit/passes/device_type_analysis.h,sha256=NFvPZxkTRmVQ6H8iM0hcF_ksf6qiqS38FYAsZQkbdoM,496 +torch/include/torch/csrc/jit/passes/dtype_analysis.h,sha256=GCEVeveWihqGgJphM-DTpPlhNIvcUk3qGYUjUUlvZ1g,643 +torch/include/torch/csrc/jit/passes/eliminate_no_ops.h,sha256=sSidjyAbmaWMv9g-Y_GybSiJOw06wcwHabBXSZ5cuGA,746 +torch/include/torch/csrc/jit/passes/erase_number_types.h,sha256=n37Zgicm7KfOrz5bMFo-oSfjg1V8EXCWdEVwStYT9xM,1042 +torch/include/torch/csrc/jit/passes/fixup_trace_scope_blocks.h,sha256=PRjV9r6yj3s-N1hU1A2DYQSqU2g4I-BbRVRlBNrBh8w,1902 +torch/include/torch/csrc/jit/passes/fold_conv_bn.h,sha256=rGgsRsYYCyvl4oZ3eiHEPFl0NRnhO7sidzFHY-1QR2Q,1224 +torch/include/torch/csrc/jit/passes/fold_linear_bn.h,sha256=o3V0vkB4kiML7iN7KSvGYX20ta-cozpNT2o-9ofvwd8,920 +torch/include/torch/csrc/jit/passes/freeze_module.h,sha256=zsbw3P9Y4vImNTSrHU3a3LCgbtPCC51qJ7q5ArWdkmM,1473 +torch/include/torch/csrc/jit/passes/frozen_concat_linear.h,sha256=R_dFKDX09qS1051w0BQg4QEuy1urm8Rd2nt-6Qhi0f8,506 +torch/include/torch/csrc/jit/passes/frozen_conv_add_relu_fusion.h,sha256=oJ3U7PeMHAuuGK3LDJRwWAVQHU0KRSuzI8-PHM37dKU,558 +torch/include/torch/csrc/jit/passes/frozen_conv_folding.h,sha256=gqsDYKhyF7MI_OpRx7a922ULvGWGpoa4zTW70HFLVxA,1101 +torch/include/torch/csrc/jit/passes/frozen_graph_optimizations.h,sha256=aaRe66_HHQh2uYaDpGH8Od5NzJKDbMqtTdK9I5Nklt0,695 +torch/include/torch/csrc/jit/passes/frozen_linear_folding.h,sha256=YSh0apny-0M41_JuI7Tri2Z1kQf3UZWfrzaKubdir0M,599 +torch/include/torch/csrc/jit/passes/frozen_linear_transpose.h,sha256=c_Dw0bRU5R1w0-nmiPEauljSJgvi88FD3mUgsiCEiJY,515 +torch/include/torch/csrc/jit/passes/frozen_ops_to_mkldnn.h,sha256=cle2FR8ZgXrj4gEaFmqDIcMTg0TNEg3w4L_GTtiTtMs,643 +torch/include/torch/csrc/jit/passes/fuse_linear.h,sha256=VzB0Mw8-Lu4zS1JRwncRxmUMiU9fFOu80orgNp3Y6Ew,997 +torch/include/torch/csrc/jit/passes/fuse_relu.h,sha256=KcwM0AHZQFtb25kONC0f8oW3XYlUBArSZ2Y0oOtQSMo,502 +torch/include/torch/csrc/jit/passes/graph_fuser.h,sha256=nSYOAoLK2QWkUGSCaQ4zI3iTxT6QzFJnRHxrEIuTQGI,1480 +torch/include/torch/csrc/jit/passes/graph_rewrite_helper.h,sha256=WoXBBh9VgNk-MIqbifhEGiX4kh4LU2qs5odFQGXtwsc,1989 +torch/include/torch/csrc/jit/passes/guard_elimination.h,sha256=Q4EZMWjp79bdW-X45G1IdG4TFj1tl6jMKVVNIPXYCUQ,605 +torch/include/torch/csrc/jit/passes/hoist_conv_packed_params.h,sha256=CdAvrvCgDfL7zvtwgHG0AKqrWOdKe_nqbbwaJcOQRlY,440 +torch/include/torch/csrc/jit/passes/inline_autodiff_subgraphs.h,sha256=VZxi6tzq_hzPn3EIzIg2ZtvXBJ6yGfAZc5b7ATbX_ug,504 +torch/include/torch/csrc/jit/passes/inline_fork_wait.h,sha256=OYvgI7lfPVwGsxgnzilm2OJuhkGt6Ltc__35dn0-Owg,776 +torch/include/torch/csrc/jit/passes/inline_forked_closures.h,sha256=8e3yZCxzuikeldUCGUF3XTgAfkeHh6wEyz-KBCulR9k,456 +torch/include/torch/csrc/jit/passes/inliner.h,sha256=ToYyx5gJ37QcV0oTvzDEkDMVJUnb8GYXzk-4hThoCzs,483 +torch/include/torch/csrc/jit/passes/inplace_check.h,sha256=sOI32FKTl3O3s0tHYlg1k_Sn77sEst2-hWueyAj3i2A,390 +torch/include/torch/csrc/jit/passes/insert_guards.h,sha256=AIb_PhMXGon7heho-glI929uIC1Im0meEjOgRZLiIyI,668 +torch/include/torch/csrc/jit/passes/integer_value_refinement.h,sha256=CLQ2zGwDUAv_Y281ot29mm2GgduyL0vK3eujms0Md1Q,463 +torch/include/torch/csrc/jit/passes/lift_closures.h,sha256=td4UuyZdea9dFkD6W7KJODsmMQtFzUloB7ZRJY_zEsc,451 +torch/include/torch/csrc/jit/passes/liveness.h,sha256=Vxa4INJpMYKITv3m8bqOFZfJzDbfDZAfLtMLewxzuUE,879 +torch/include/torch/csrc/jit/passes/loop_unrolling.h,sha256=asj2phSEZGmoy9emUuDzgkc26M2J3kVk5nc6VfTUG-8,1235 +torch/include/torch/csrc/jit/passes/lower_grad_of.h,sha256=QYS_pCq01yGmynQVkT9l0_v1PgmUrWTMs3j5vdWn1KY,577 +torch/include/torch/csrc/jit/passes/lower_graph.h,sha256=rWUkw2SbpYBkvHYsCSnDrqeHXD75Qu2CNPjJJtC8QZ8,979 +torch/include/torch/csrc/jit/passes/lower_tuples.h,sha256=6LHRDOHVsmYyg1ciIF0PdXDX-yf0Wnj42YeXrMISSWQ,895 +torch/include/torch/csrc/jit/passes/metal_rewrite.h,sha256=Aq9hJ1g_b3rQKMBRm5BboxDHVKXCSKYDZ2A_duZ8LQE,835 +torch/include/torch/csrc/jit/passes/mkldnn_rewrite.h,sha256=3VXK9g-dLkqo4FCUmH8MwloWrtA3KtinJYCNeXgD-uA,859 +torch/include/torch/csrc/jit/passes/mobile_optimizer_type.h,sha256=15fTYpsDTiBMHsuAwOxPsTjXIZL37PtFWtZKRgFRn6M,491 +torch/include/torch/csrc/jit/passes/normalize_ops.h,sha256=lXj48gpVeLtNH-4uxdtIKYTCCMgAXEcXwm0pkzDgCU0,765 +torch/include/torch/csrc/jit/passes/onednn_graph_fuser.h,sha256=GBjQKvwNo03pCg9YKeSRbXSn-6QlqFFcvzfLbiU3gfY,1653 +torch/include/torch/csrc/jit/passes/onnx.h,sha256=_UkiWGPLHo91XTx9I-0bagdDtkLgzuNjaI1N9ZLYSYk,1086 +torch/include/torch/csrc/jit/passes/onnx/cast_all_constant_to_floating.h,sha256=dAvOYnKqFVtEHNOkfClRH3eNj7MbSKFcam_kBj4sGwM,471 +torch/include/torch/csrc/jit/passes/onnx/constant_fold.h,sha256=_XQyhSXo2-qnPkq340iCigEeHpih1PVrHls2iSjgm20,947 +torch/include/torch/csrc/jit/passes/onnx/constant_map.h,sha256=5pZuHzoqzXePcJsmaPYCqfEKKnU3sH4uoNwh82NatdY,4676 +torch/include/torch/csrc/jit/passes/onnx/deduplicate_initializers.h,sha256=w1w5ieTjPU34qymJgA_2JGWjoL2bmjx1A7KYQDSzPJ0,501 +torch/include/torch/csrc/jit/passes/onnx/eliminate_unused_items.h,sha256=krk079cOT4rOH11yfKlxBXyD7zx-KH47s6D22eUfaPk,597 +torch/include/torch/csrc/jit/passes/onnx/eval_peephole.h,sha256=-hDeDvDyI7kkdWzqeANeNLyjZ-XhSNCEtFkPrso0_4I,474 +torch/include/torch/csrc/jit/passes/onnx/fixup_onnx_controlflow.h,sha256=QhZj4aOPAdB3mVdIra74SH8ZIDRis6MMDnagak07f-o,475 +torch/include/torch/csrc/jit/passes/onnx/function_extraction.h,sha256=IrN-zKkMc117HwUOPd8rpLrScYQg0BWjVhsKfL3KpNQ,2505 +torch/include/torch/csrc/jit/passes/onnx/function_substitution.h,sha256=UE5MXmhwcKcnThcaJ44yHVQhUJEOYH1Te1cu8ua6dLQ,389 +torch/include/torch/csrc/jit/passes/onnx/helper.h,sha256=RSnKwjRHMLYSn9MBghSauvoIFIEppnUgSwoDodezT9U,2259 +torch/include/torch/csrc/jit/passes/onnx/list_model_parameters.h,sha256=Bjjdu1GfYUC_XzaoNeDd9uFxxI3BtRqHUBeyQsYTFh4,493 +torch/include/torch/csrc/jit/passes/onnx/naming.h,sha256=xARbuBP5w5yHBxa1SqvHvsIgsgSfzjYCLkbYIyOCKKA,1014 +torch/include/torch/csrc/jit/passes/onnx/onnx_log.h,sha256=k7Q3msE4ebSIbatK2SozmuuHhAN6PmDtG5o7na41lgw,838 +torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/autograd_function_process.h,sha256=TzA91rqMukunWZhNt7243zU9XfOm183P5c8xD4y3QoY,429 +torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/common.h,sha256=SwwoOcou_ddWJ-ZZaq9dGO4DUnQVkZ0H2fH-R3je_u8,590 +torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/pattern_conversion.h,sha256=yWht1SKQmZcuKsB72qwwKXcIBtlvzWA53D_S-em-OtI,2349 +torch/include/torch/csrc/jit/passes/onnx/pattern_conversion/pattern_encapsulation.h,sha256=5NhWadoRDuNm59FWVQ_E1vTuD5XUJ9aLvg_7offUD58,1616 +torch/include/torch/csrc/jit/passes/onnx/peephole.h,sha256=CE7EBRjxIdulDb9RKZKCZzDY8_i6ahhbP4Ly1tpep1E,467 +torch/include/torch/csrc/jit/passes/onnx/prepare_division_for_onnx.h,sha256=BMLs0MXEEuJAzJbrSF7RsC0nUBDYEdb1y59QJzDk9gk,705 +torch/include/torch/csrc/jit/passes/onnx/preprocess_for_onnx.h,sha256=YippKvsxqhHWv_oqgC2ZW-RnI-jROrOlQ6V93sCAZOg,409 +torch/include/torch/csrc/jit/passes/onnx/remove_inplace_ops_for_onnx.h,sha256=-hSGx58ku7bXvK-mb020ASWIUnCX2q3W-82DP8zgaPI,455 +torch/include/torch/csrc/jit/passes/onnx/scalar_type_analysis.h,sha256=74SMht0iHAbNST6nalhUdYH4g3maNy3l7HFSQQvGwcQ,534 +torch/include/torch/csrc/jit/passes/onnx/shape_type_inference.h,sha256=SlGyecscNEDgdxVBblR94i56eM2W-I5nHs3FKVXSPB0,4209 +torch/include/torch/csrc/jit/passes/onnx/unpack_quantized_weights.h,sha256=DkcIIjXB_6VqcDWmss6830dM5ICTHG3pgQhqQMOfVuo,681 +torch/include/torch/csrc/jit/passes/pass_manager.h,sha256=LvOMTVmjxz-XJqZKWj9p6k0UwMlRTualwlZQ9Kar8Hs,4803 +torch/include/torch/csrc/jit/passes/peephole.h,sha256=Mo7rrtoHcaYiQXLtAuzLTrq6havHuGUk6S8EI7TwjWc,736 +torch/include/torch/csrc/jit/passes/peephole_alias_sensitive.h,sha256=Sf5RJwOqKFe4iciAVLkp1FrJpjShsvbjVHL5By-XDQs,664 +torch/include/torch/csrc/jit/passes/peephole_dict_idioms.h,sha256=V1wOXOIN92YLg3pFp1SYwO3zuCxe_c7Iwq9coIatUZU,1229 +torch/include/torch/csrc/jit/passes/peephole_list_idioms.h,sha256=y6BBOysIaiX9LXvuefeadCFj2p0s4sko4-GsQYPnxnI,2232 +torch/include/torch/csrc/jit/passes/peephole_non_tensor.h,sha256=skm22Pz51el-KVG0k2jh9gR8wlDOKfcR5NDQPmsHaRE,571 +torch/include/torch/csrc/jit/passes/prepack_folding.h,sha256=kX30fJNLOVaYaUVY2f1iQk8VySYqlOXv3gi6m5vQy8k,587 +torch/include/torch/csrc/jit/passes/quantization/dedup_module_uses.h,sha256=cLkTsvEGzKSDL1wyuvvvBBagD7DnVXlXTcwj3pTF-NE,1055 +torch/include/torch/csrc/jit/passes/quantization/finalize.h,sha256=-QWRnlHJuRD-VOMy4dhOXHl76pcvNMJHfG3HaKWWSI4,2550 +torch/include/torch/csrc/jit/passes/quantization/fusion_passes.h,sha256=7MHhqzRUYn7PztFKJ-wXW57JLKNQHmiJ6M06Ybi0VAY,420 +torch/include/torch/csrc/jit/passes/quantization/helper.h,sha256=MMFMaD1PNwvUSsdkkfdf2eK8XKpYuFfzysXyiEp33U0,7713 +torch/include/torch/csrc/jit/passes/quantization/insert_observers.h,sha256=k4Z0aSWRs2a2AiHnJPddEr10qb8XtqaQS4VTIRmICiA,2580 +torch/include/torch/csrc/jit/passes/quantization/insert_quant_dequant.h,sha256=sv9uH-qDkwIiS6PfV9KyTuw7SN_xZo40mAGHL7-TXME,1679 +torch/include/torch/csrc/jit/passes/quantization/quantization_patterns.h,sha256=qEZPI2jlOnTnl6wf-wxjG11-_RZm_UvPPU687m4oze0,53615 +torch/include/torch/csrc/jit/passes/quantization/quantization_type.h,sha256=1MjTQ7coNAARQy0W3o_hvfMv-3Wl5VhThENqY99fiBc,587 +torch/include/torch/csrc/jit/passes/quantization/register_packed_params.h,sha256=qPhixrMBoFVCr8GWDqs-v4jIiDDai3ytqbv23aMeDWI,743 +torch/include/torch/csrc/jit/passes/refine_tuple_types.h,sha256=97LywtAOszRtA9MQscOX8gLmB2pZzXQZfEz9C4wtaUg,496 +torch/include/torch/csrc/jit/passes/remove_dropout.h,sha256=ANNXHbSIIcRmZpBcoC0ZlsPAsGgJoobsWCGrLK-gMKI,509 +torch/include/torch/csrc/jit/passes/remove_exceptions.h,sha256=BSf7-h68dGHJ4OCbxFH5G53zLAnqgcQDCnwjb0kvnbM,1183 +torch/include/torch/csrc/jit/passes/remove_expands.h,sha256=_wMfo-Lk2Kkru2UWb7wxY5SjZd9ch-_DVvaTdVXoAjY,397 +torch/include/torch/csrc/jit/passes/remove_inplace_ops.h,sha256=ZeX2AINr3DaNyLJT4mGcwIhT2nwwNi-TjFpVD5_tMSQ,525 +torch/include/torch/csrc/jit/passes/remove_mutation.h,sha256=v_hHpZwvYKxAo3S7AcIURHmz-lI3Zv5YXMnuTM-8wFo,2891 +torch/include/torch/csrc/jit/passes/remove_redundant_profiles.h,sha256=zcRHNfhQ46hRnX5G1vNJJGdNf5ln9O2zFZDmf1REyFM,491 +torch/include/torch/csrc/jit/passes/replacement_of_old_operators.h,sha256=Vt8IwExtBSQL8GgsGUMuAonB3K7ZjfkBzRCHUQOv5L4,688 +torch/include/torch/csrc/jit/passes/requires_grad_analysis.h,sha256=NS7EqB2A-1reB4lsvoL4Yp-sxnWfQFNyr7SIaHzhmyU,475 +torch/include/torch/csrc/jit/passes/restore_mutation.h,sha256=eOvMa99BGqwPfZPgiS0QOq5U0BQJAcNHygQHZzcm2BA,2055 +torch/include/torch/csrc/jit/passes/shape_analysis.h,sha256=lmNddX5YF6phTtpKEEj2EuaShFHmG4THQyB7Td0qHbI,1331 +torch/include/torch/csrc/jit/passes/specialize_autogradzero.h,sha256=21kPaHTi-xVHnofwW5P0sK73y-1XsKnl8lMT0SsSy64,885 +torch/include/torch/csrc/jit/passes/subgraph_rewrite.h,sha256=wMJHB3HCrYpqgBex6Kys9oWZ7OB6xt0nfIOK5pE_PSE,4313 +torch/include/torch/csrc/jit/passes/symbolic_shape_analysis.h,sha256=TOMd0js5rYIWV05AFpf-PJq9OYXUoaFBOl707VvPZvk,2332 +torch/include/torch/csrc/jit/passes/symbolic_shape_cache.h,sha256=sXd7qzIrY6LAcp9iudjXKO23O9J4iS2hMVpQw3QS47I,1828 +torch/include/torch/csrc/jit/passes/symbolic_shape_runtime_fusion.h,sha256=DLMEVq55DQ3b0DYl-XywA0aS1Av7IQdAHC5LubhTa00,2593 +torch/include/torch/csrc/jit/passes/tensorexpr_fuser.h,sha256=BZSULniWR5TvBLOyiryThuYXdissuTKArdcJ66z0VZ0,2953 +torch/include/torch/csrc/jit/passes/update_differentiable_graph_requires_grad.h,sha256=hwcyT7I6dlZWCmLqedbalUm73IvOfAKMhXlZwUZdWcM,969 +torch/include/torch/csrc/jit/passes/utils/check_alias_annotation.h,sha256=hBYq1dwM4MiGPCY1e89K2uxovD2u01KNkHu-c4A4-4s,841 +torch/include/torch/csrc/jit/passes/utils/memory_dag.h,sha256=PgqnB_pKQxz3le9erTTycryFQ9q1YQTz4eri-0wna5o,6645 +torch/include/torch/csrc/jit/passes/utils/op_registry.h,sha256=2egFx0dFliiCY71fhR1HkJKYJrZEA0v4t-0tAITc65w,1262 +torch/include/torch/csrc/jit/passes/utils/optimization_utils.h,sha256=AiONhxLLO2m3l_G7mcAX1xZigva1RrBJHvUoa78ITNw,473 +torch/include/torch/csrc/jit/passes/utils/subgraph_utils.h,sha256=NTZJ2TZ3LpMsje665S60ulrvfG54IKWMn4PHR0x9Rtw,2614 +torch/include/torch/csrc/jit/passes/value_refinement_utils.h,sha256=nDJhO_jc3-arA7JW5WdnWn-s5EKgx4sMEAxsRGXQSlQ,2861 +torch/include/torch/csrc/jit/passes/variadic_ops.h,sha256=gB8X2YrJJvpMxajBSMk36y4Wb-qGMed2vKIKiDBpzE8,1139 +torch/include/torch/csrc/jit/passes/vulkan_rewrite.h,sha256=2am84jm42Je81wqkg8oNJVsSLY39hhxK7nygP7eCLiI,927 +torch/include/torch/csrc/jit/passes/xnnpack_rewrite.h,sha256=yl8n6WNwyx2L4_aH4hUE1UQrbAWO23sZ3qVB_oZLwB4,1049 +torch/include/torch/csrc/jit/python/init.h,sha256=vNYIEI_xN8J5Mt06W0zqxyNmpg9sVxwgpW4zkp037QI,397 +torch/include/torch/csrc/jit/python/module_python.h,sha256=igT4tAZwKBYHmEtHSUyXp8SM_cpiCsWBbJ2y-u5xO_I,2127 +torch/include/torch/csrc/jit/python/opaque_obj.h,sha256=kZy2BWyyma9BVgtF4Owx_mJ0Aj3CPLn-5eh12pf0CE4,3135 +torch/include/torch/csrc/jit/python/pybind.h,sha256=oRXJuaMzSL5txY51K_YchGAQ5AKXuFO6v0JL3EQ3o1c,8207 +torch/include/torch/csrc/jit/python/pybind_utils.h,sha256=yffkkHhrR6jUeVO6KZhaNoNpsmhjEXigvngUqV2qpBU,46161 +torch/include/torch/csrc/jit/python/python_arg_flatten.h,sha256=T_5P6SeD4L0HwgiF1EgBaOJTns53ObnKlrWo8uuSwic,3781 +torch/include/torch/csrc/jit/python/python_custom_class.h,sha256=ld9_AGKC_-TTIhJ8Arb18NskKGvOhmKM5Mk4hb9gqRw,666 +torch/include/torch/csrc/jit/python/python_dict.h,sha256=wnSiklovZ7ltUylHMtu0XTnxSGZeDvQbUYUgP5H86TA,3637 +torch/include/torch/csrc/jit/python/python_ir.h,sha256=WvU_pTbdoTY5yQQ9wJIv6f4ixvr2pfKZ8bAYW7M9Vws,1947 +torch/include/torch/csrc/jit/python/python_ivalue.h,sha256=RoKOjA5RuPZwble4J_LxmG3vg8Epbd6rSDCnia1nEkg,3936 +torch/include/torch/csrc/jit/python/python_list.h,sha256=bR11jlcisx1BfNLiUy_sZt4EChGw3Yzymm6PvdRe7Bw,5748 +torch/include/torch/csrc/jit/python/python_sugared_value.h,sha256=OWOjMwm4H3Mb9lr1hRca7uvAAIIy2jR3TUvfli7-wfk,11590 +torch/include/torch/csrc/jit/python/python_tracer.h,sha256=4T989L04dccgGxgHBXghLodx48imnb7nUfBxO4GiieU,1472 +torch/include/torch/csrc/jit/python/python_tree_views.h,sha256=RPY1Ogr9vgVgg_hAE3ISyF_xYAWp9If4kdwmcaE9TuE,404 +torch/include/torch/csrc/jit/python/script_init.h,sha256=XgA-WXGN1XrJwXmirEGVqc5UtY_KWNb4fUQX-Thfj2E,406 +torch/include/torch/csrc/jit/python/update_graph_executor_opt.h,sha256=xdKRIbEiJewN_qDogGpAFIuyhaoeX-X0n2_4nB6bAwg,439 +torch/include/torch/csrc/jit/python/utf8_decoding_ignore.h,sha256=B6blNz_2z3iECmAe7j7YLkDvvI8C1qS_3XOQE_Xqz14,433 +torch/include/torch/csrc/jit/resource_guard.h,sha256=SlgztVvYnO7gH67Es8lAgD36-otV1Z4Bwha1Lbf3yfQ,694 +torch/include/torch/csrc/jit/runtime/argument_spec.h,sha256=XL7F12AAl0FXOoMXra-nTqUw_K3TeXtmh_FFH1XeCU4,16638 +torch/include/torch/csrc/jit/runtime/autodiff.h,sha256=gk9xNh2dJ_7kotDK_PPTRN7X7GfS1U5o45FqkUpuUYw,4184 +torch/include/torch/csrc/jit/runtime/calculate_necessary_args.h,sha256=VHcSKHiUe_SXS2K7hyc0uG-4OgZQFjZcnzKEgmUz3LU,2535 +torch/include/torch/csrc/jit/runtime/custom_operator.h,sha256=h1APxS6uYOLL9pKJ-xoAcoRb8m0gwaD2lD2S4OOBodc,1308 +torch/include/torch/csrc/jit/runtime/decomposition_registry.h,sha256=rqzNfixCRRvSIACDilIpQSyJ2bcFI0DMJ34B1F92LkQ,1299 +torch/include/torch/csrc/jit/runtime/decomposition_registry_util.h,sha256=x-tYdDgODVbKlfH9CdiePOmbQokrbO3ceTS2Bb_aAtQ,515 +torch/include/torch/csrc/jit/runtime/exception_message.h,sha256=BslHUofsN8BPUqzFUOSbCR4K_vn-6jpuoqMVGyILQhw,868 +torch/include/torch/csrc/jit/runtime/graph_executor.h,sha256=Uf_YQ650NhnpoYN5ZKp78F6NWlmgrANPrs2SPFvkQp0,5237 +torch/include/torch/csrc/jit/runtime/graph_executor_impl.h,sha256=z7RvxYg1R1CgVKtoxxu-cDBhybJddJhgF4fv4voCSVQ,4617 +torch/include/torch/csrc/jit/runtime/graph_iterator.h,sha256=jCz577Ln9PoIwzyx0FEzkHmdqav5fqYC9nJ-_nUNPGE,5192 +torch/include/torch/csrc/jit/runtime/instruction.h,sha256=29MNt2M4kKwP_BhDfpzVJTLAsgDIdUJnUtg4Cft0s7s,5798 +torch/include/torch/csrc/jit/runtime/interpreter.h,sha256=Ae4OGN2Iy6Y2UDGsYHqn6qu2gaLr-K5nbeGU7kHXqtM,5245 +torch/include/torch/csrc/jit/runtime/interpreter/can_emit_inline.h,sha256=y9NZJmjJOWpkZq_E_dI1c2TixH-NoSmaWqzjqmPenuc,4208 +torch/include/torch/csrc/jit/runtime/interpreter/code_impl.h,sha256=SuENU6qmE-6Cq5_9QNudU-pMAXkgRAhdXFSZgK2SDyY,33553 +torch/include/torch/csrc/jit/runtime/interpreter/frame.h,sha256=VeercdUuV2EqaQPgEg_4bANdttPbyxOWtq0cHVlxl-I,1364 +torch/include/torch/csrc/jit/runtime/interpreter/preprocess_graph.h,sha256=nTwhXzmPZROEwL6x3Xa7DSHHelnHkxbUnRZjjQmwIek,655 +torch/include/torch/csrc/jit/runtime/jit_exception.h,sha256=Zc8QecTVAQX8NJmO5jyOlbhiMI85f6cHAPrnrsj-QJc,1423 +torch/include/torch/csrc/jit/runtime/jit_trace.h,sha256=QOfLSGozBn6GWvSU8lZP2ei3H341v5_c56RaQKHo-Do,461 +torch/include/torch/csrc/jit/runtime/logging.h,sha256=dw5NkNpPi_XUNYg7PZ7Y7iBC5ziQull3pL8OsMiU80o,2868 +torch/include/torch/csrc/jit/runtime/operator.h,sha256=GuDfqdJ4QpIcWGQnFn_qO0-P4rGIIl1s9t97AR2fFNA,11872 +torch/include/torch/csrc/jit/runtime/operator_options.h,sha256=aKyUKJcXMUY2coB-r6w7Gn5MqM6Oi_gagAUreCBz3fs,418 +torch/include/torch/csrc/jit/runtime/print_handler.h,sha256=Y8I3uAqbeTiRz0PGKPPuddwV3K-x5TsFUaiHsQx3Tx4,562 +torch/include/torch/csrc/jit/runtime/profiling_graph_executor_impl.h,sha256=VDgNi6zLZCugkS7NMOd_vG4AD-FhLLLNqUJlxF4gDx0,3453 +torch/include/torch/csrc/jit/runtime/profiling_record.h,sha256=Pl8FRDdC3xcnXC0QY_YqxKsXMYAcR0twk3_bDjgNQF8,8811 +torch/include/torch/csrc/jit/runtime/register_ops_utils.h,sha256=q6U1f5ag2mbImjCdDX89RTtn039ONUbsWDhgZcnpUV8,42878 +torch/include/torch/csrc/jit/runtime/script_profile.h,sha256=N1X3tzFIej3-eVHceltWERDxdGVPZVSCLILcwva1q3I,2890 +torch/include/torch/csrc/jit/runtime/serialized_shape_function_registry.h,sha256=mcuHossgV1v7dKRzyGnfhZFyhXsTzDnaWN6Y9HT9Mjs,610 +torch/include/torch/csrc/jit/runtime/shape_function_registry.h,sha256=uoL-f3SUaAQuiukKBOihRcRjbGI4RW7UYHcnDV3cN6o,497 +torch/include/torch/csrc/jit/runtime/simple_graph_executor_impl.h,sha256=99zYDUdxqtTqzauiZUd8DFazx9gFybNDk3kXzadfvPM,1011 +torch/include/torch/csrc/jit/runtime/slice_indices_adjust.h,sha256=n3jrzM_G9_3ccL-w6g17Klgv7m7OZx8Nt4ueF-AYh6s,1036 +torch/include/torch/csrc/jit/runtime/static/ProcessedNodeInputs.h,sha256=0VsG_EWU4RlKMytyPWszfKkN9r1tsUguema93szWs20,6657 +torch/include/torch/csrc/jit/runtime/static/fusion.h,sha256=tbkTq6d876UIXZJUIwG1E9y0Y9G51l38RPRvcDOehC8,561 +torch/include/torch/csrc/jit/runtime/static/impl.h,sha256=ji25jEShUmncR86XQMCPUz4OPSoYXPtJfF60npcw14A,36161 +torch/include/torch/csrc/jit/runtime/static/init.h,sha256=aNaEAVBJYoov2lPcnGzHAgfMYph6jxyAOUfRrxrQK-c,403 +torch/include/torch/csrc/jit/runtime/static/memory_planner.h,sha256=J4ta170Vh1I_hSfwvvzKef8mCqWVPmo2BvPU0J3CEjo,10180 +torch/include/torch/csrc/jit/runtime/static/ops.h,sha256=5I8G87YwBcJ--zBW4vWSqV4BCYa_k_U7w7GRbvduQeU,5824 +torch/include/torch/csrc/jit/runtime/static/passes.h,sha256=Wu3GXnn1CiiP90I5bgtelZQYEn3jFL6arF5X6PsE1MM,3938 +torch/include/torch/csrc/jit/runtime/static/processed_node_wrapper.h,sha256=My4mV4TP1f4HOTVOPj3RQZoKXFIX2haaHuuXnH7Wo9A,6849 +torch/include/torch/csrc/jit/runtime/static/static_method.h,sha256=I9CqnYsJgIzLyVJz-US2O8M0_O6F6J8DgzBOkihoTKY,1578 +torch/include/torch/csrc/jit/runtime/static/te_wrapper.h,sha256=MKHuU07WRx2w7L6AbXVEe9k9PVAUvgmHwZXz_kp4x4s,1399 +torch/include/torch/csrc/jit/runtime/symbolic_script.h,sha256=g3o1hzSledXLC_qGXaW00L5ceN_5f9GkUcsHVj0oh7A,816 +torch/include/torch/csrc/jit/runtime/symbolic_shape_registry.h,sha256=OhMa9p-mfJyzj-iQGgK97aIowSuoWM1_aFwYp5JN76c,3056 +torch/include/torch/csrc/jit/runtime/symbolic_shape_registry_util.h,sha256=x2o-VKti7VMmmQqOIINueRcjcddhyRBLSAhZyYTkFGQ,605 +torch/include/torch/csrc/jit/runtime/vararg_functions.h,sha256=4C3J9b3VP4jaUBAGH3ZF9CRc2WRAfPP0JdhMOiAiUWs,1401 +torch/include/torch/csrc/jit/runtime/variable_tensor_list.h,sha256=fEV6nglRgQkc_2-Gag2kf9lwJ_z4FV2Q5M7xOjUZs-w,781 +torch/include/torch/csrc/jit/serialization/callstack_debug_info_serialization.h,sha256=QcsfMxPT_pqRc7EASChDRGo24Qs4FWRoXsHqcdDdBSw,2858 +torch/include/torch/csrc/jit/serialization/export.h,sha256=X6mqoZw41PetqOCOxaeZyaKcsLcJ8qYnx7T2bOFaaM4,11743 +torch/include/torch/csrc/jit/serialization/export_bytecode.h,sha256=K5fnIyMk2cOEvawaXbJyMwhwPkzwAnCjswZdjiK9Sgo,1622 +torch/include/torch/csrc/jit/serialization/flatbuffer_serializer.h,sha256=zVHXB_S7tz5eQ9_ams7Zs_rfLWZsP_2dSMBNC1ttAKY,3297 +torch/include/torch/csrc/jit/serialization/flatbuffer_serializer_jit.h,sha256=n-sfcJ-3xmrwI5jpCFXoQvSaFEV-gp8tVB_yphdoAXc,426 +torch/include/torch/csrc/jit/serialization/import.h,sha256=CBVFVUcoTm_zaWm7uoqeCR_Zd0Mt-9ViWQ2RZYF-B88,5085 +torch/include/torch/csrc/jit/serialization/import_export_constants.h,sha256=4oXjn6hyyJNAWk5su46kkyjStM189G05kktdT7N43g0,899 +torch/include/torch/csrc/jit/serialization/import_export_functions.h,sha256=LwcXblZza7xkyq0Rka34tHdsd_mTGmyT4tJpoEIsavs,643 +torch/include/torch/csrc/jit/serialization/import_export_helpers.h,sha256=4eXPN8wLr2YGcnCxglW8a3uQP4JarUu6znXfw2Jutt4,905 +torch/include/torch/csrc/jit/serialization/import_read.h,sha256=A5y-7aXZN_AsR_Xs5TeAUFDzu3DKU6pGwuPiv9nCKKo,1086 +torch/include/torch/csrc/jit/serialization/import_source.h,sha256=CsfRD-msOa4zp7bsm3GFRueWXoNKllKtqBnc12E_kQ4,3678 +torch/include/torch/csrc/jit/serialization/mobile_bytecode_generated.h,sha256=lsTXYABe5eo_pIVoaEkB6afoTkB2GZ-m1t4_vPHng5g,98868 +torch/include/torch/csrc/jit/serialization/onnx.h,sha256=qJhopSKUjvWlH2obLaaYCTMrW9W2qnZJVAjN3PNLEQ4,767 +torch/include/torch/csrc/jit/serialization/pickle.h,sha256=pgfxjFeY2B8EbGlNOh1jGnlAN27A9C8yCCpp9RN3AdA,4944 +torch/include/torch/csrc/jit/serialization/pickler.h,sha256=irAEuSyKkGnm_noEteCownc4zn_IhkGUOf2Cr8M3JKY,6589 +torch/include/torch/csrc/jit/serialization/pickler_helper.h,sha256=JP45K7oJ56Carbn1DE4qVTiZpdU4Xp_kbzXgarec2_M,7204 +torch/include/torch/csrc/jit/serialization/python_print.h,sha256=N-gq3PVba8pfhIElWd3BStwWTrR0VpgfM_fOyYZhY6U,1563 +torch/include/torch/csrc/jit/serialization/source_range_serialization.h,sha256=7g9X73noS0b5SL75D4qkmSh9Y5b4eFSaII8TEps4a88,1920 +torch/include/torch/csrc/jit/serialization/source_range_serialization_impl.h,sha256=N3-0_YKngFCYPWVg-vLqjo4NIoJZIHDwubKSEIpWFt8,934 +torch/include/torch/csrc/jit/serialization/storage_context.h,sha256=ndqEnG2nx8IYCScZEGgfSfnK98O8LmTkFsi6nqzGpy8,2741 +torch/include/torch/csrc/jit/serialization/type_name_uniquer.h,sha256=WeGv4RvmdKHQWZkT8cxnS0LXKLkYMVTsfLiGVFFFkfs,1008 +torch/include/torch/csrc/jit/serialization/unpickler.h,sha256=wF1HWI6CT-NqrxfAgrdJBpyBkxIfgGD8p1TLxYhaygA,7938 +torch/include/torch/csrc/jit/tensorexpr/analysis.h,sha256=2ujRn9DusgJSx48wfTBcafF_OmRqI2ZUhSWEON_QUo8,9236 +torch/include/torch/csrc/jit/tensorexpr/block_codegen.h,sha256=rWfhdJdSRI5_Yci_08AdMpAo9DxYL_IGeIxcTFB1qEI,4544 +torch/include/torch/csrc/jit/tensorexpr/bounds_inference.h,sha256=X7rqLiu36NbZ8xyWLHI51dPOxAF0FEB2ZngozgS6FBo,2424 +torch/include/torch/csrc/jit/tensorexpr/bounds_overlap.h,sha256=hRRgaDYY0YQjP09GU4iKkkxcMurzQAUF0u-hD4M4WBg,4725 +torch/include/torch/csrc/jit/tensorexpr/codegen.h,sha256=Gl-PYYGtvgAeIGnXzfIoDlxgaENonD3YbA95eHqwDsQ,7820 +torch/include/torch/csrc/jit/tensorexpr/cpp_codegen.h,sha256=vql6RehQDmUSOGYtaNMfYCnOLkDpUJQbFOIu8kbpB-8,2706 +torch/include/torch/csrc/jit/tensorexpr/cpp_intrinsics.h,sha256=f2nVKYuxGRr0dRpmYBC2s2HouEh_UrCkirRB084kxgI,887 +torch/include/torch/csrc/jit/tensorexpr/cuda_codegen.h,sha256=1tNyoT2slJt8I0eUCv3hvk2vWHCvkgFoIFuC6cpe8l0,8488 +torch/include/torch/csrc/jit/tensorexpr/cuda_random.h,sha256=xNgNZ6ulyIg_RXjBcuLSywKlfFPlfLbYooFQ7o7Igqg,2846 +torch/include/torch/csrc/jit/tensorexpr/eval.h,sha256=di5-_tg9Wxz_zZO1_TfSzETK-BjVWIhrjjHU1AqB0Ps,10097 +torch/include/torch/csrc/jit/tensorexpr/exceptions.h,sha256=h4tscEytrsVBp0FeDu7dQ1a70SyqvhE2zK52gBNLE-c,3475 +torch/include/torch/csrc/jit/tensorexpr/expr.h,sha256=kwHPQLgNNggfMmRuhW-nG_Kx70oBKGCKLvp-pCwyUTU,14435 +torch/include/torch/csrc/jit/tensorexpr/external_functions.h,sha256=TbX1FLAD27y78pp0ibNQbkAtkrRedl819pSEnsCrkkM,3687 +torch/include/torch/csrc/jit/tensorexpr/external_functions_core.h,sha256=-XniYhj7ULLZaIJXLxOmVi5R8M6QoPsppkZ792Xe-rU,707 +torch/include/torch/csrc/jit/tensorexpr/external_functions_registry.h,sha256=pFmM_5RCB22qFInr1mwWoX6hT4vyDDiOMNflfcxlJLE,2560 +torch/include/torch/csrc/jit/tensorexpr/fwd_decls.h,sha256=QoA0Vbh_h-33YrgqSpyn0XvheH9Ubfmy1SQZOkqRrKo,3244 +torch/include/torch/csrc/jit/tensorexpr/graph_opt.h,sha256=nztWGJXqgy2d55nJrA_O1rHuw_32tfUVTDHuQAD_t_Y,4688 +torch/include/torch/csrc/jit/tensorexpr/half_support.h,sha256=9Y-LYGyA0Xqs3NNqzhw4Id-6x2eds0c3AzxgCJ042eY,6143 +torch/include/torch/csrc/jit/tensorexpr/hash_provider.h,sha256=_M-CQWG_NeTYPfXFq-vuEo3gtAZ0Nqq0VFljlF-u7Rs,7803 +torch/include/torch/csrc/jit/tensorexpr/intrinsic_symbols.h,sha256=qc_HAjNnjTPricyApZe-ah89dmIydG-_IohjaC3QOKc,674 +torch/include/torch/csrc/jit/tensorexpr/ir.h,sha256=5AXa-hM0bsST1tTzs4WhBAB-BjJTGgE-6cX8kza-IgQ,23233 +torch/include/torch/csrc/jit/tensorexpr/ir_cloner.h,sha256=88KpZ7EwH5QwR694Lo5gWCj_T8aFTEGXeckCtLbxa9w,2597 +torch/include/torch/csrc/jit/tensorexpr/ir_mutator.h,sha256=xIIYd3jGDJwGQNYNJNGY4zTxONONq1WKhrE4AZ2bzJM,2623 +torch/include/torch/csrc/jit/tensorexpr/ir_printer.h,sha256=i9QX1RBE-T_EsYICKb9FmzdRLhVbPwcC21bbV3110HI,4515 +torch/include/torch/csrc/jit/tensorexpr/ir_simplifier.h,sha256=bfo3nS5BWEJIm5P0d1GWxuLFnpUUSjx2Z0tY-YmGscM,15511 +torch/include/torch/csrc/jit/tensorexpr/ir_verifier.h,sha256=3BIXofQLDRbC4f4oXyxtB8hPZmr58EN6TY6bye-3qaI,1567 +torch/include/torch/csrc/jit/tensorexpr/ir_visitor.h,sha256=yGeW-NEbk9Sr3JWyN-BpTqxWAIIkHFeXVJiA3FwYGoM,2441 +torch/include/torch/csrc/jit/tensorexpr/kernel.h,sha256=9WGYY5qqKt2evd4NCe8sK0fEXFjB-5cQgQUkE9QSV_k,13646 +torch/include/torch/csrc/jit/tensorexpr/llvm_codegen.h,sha256=iNG4AcV8blDZASUGlAHJfsKNp8NwPj_STQD6XhXuJ7E,4093 +torch/include/torch/csrc/jit/tensorexpr/llvm_jit.h,sha256=ti8Y7XRaZTbRofvOim-pe8htkhIK7WiApV5MpaOow2c,2274 +torch/include/torch/csrc/jit/tensorexpr/loopnest.h,sha256=K_rimui2rklVgDOUNCgJeQB_HwobFaabBd-sDAtwc44,22032 +torch/include/torch/csrc/jit/tensorexpr/loopnest_randomization.h,sha256=QYpdnfuqv9YlA0shpYVSoSC3ZdbBEU1k334eMSpl7X4,563 +torch/include/torch/csrc/jit/tensorexpr/lowerings.h,sha256=uLQ7cEVTUanhRoUj70V77d63qmEtEQ9NxOqnYmYAK-c,1534 +torch/include/torch/csrc/jit/tensorexpr/mem_dependency_checker.h,sha256=uZGsYXWIela0_0tBuKxr5Nj6fgEoJdouzDv4--ofZ-Y,13716 +torch/include/torch/csrc/jit/tensorexpr/operators/conv2d.h,sha256=NKC3IMS43R68jwounzCjcpK5kMXbyi_zHXJ_TEeYutI,3147 +torch/include/torch/csrc/jit/tensorexpr/operators/matmul.h,sha256=1_W_nubAKxRtQElE8cIbHUIMnYUpzGh0S-ubZDcUmpE,857 +torch/include/torch/csrc/jit/tensorexpr/operators/misc.h,sha256=Keqv3pTo-YL1k6Sk1uPn-WmQVKd9k5HOIv1yWKXtGfw,3542 +torch/include/torch/csrc/jit/tensorexpr/operators/norm.h,sha256=iP-nTSjXKVrcIJ0BhRfPi0Ureik5eHY4o-D6GIQST_Y,627 +torch/include/torch/csrc/jit/tensorexpr/operators/operators.h,sha256=_Ry3BgjJ7jD_7tip9TcAr456eHg-MBLFOI5SHIBh0vU,725 +torch/include/torch/csrc/jit/tensorexpr/operators/pointwise.h,sha256=8C9DJI3SpUov7VAGWsLgxd0QJFCvVan4zDiUBZVqkVo,3413 +torch/include/torch/csrc/jit/tensorexpr/operators/quantization.h,sha256=ELez7-v-cjZ4jo1uGghV51V34zyyNn9fNJhqnATQNrg,5543 +torch/include/torch/csrc/jit/tensorexpr/operators/reduction.h,sha256=o5Gm21oXdYuDycCYgzDJmTNYx1ORWT9FBj21QvWlDbc,1359 +torch/include/torch/csrc/jit/tensorexpr/operators/softmax.h,sha256=E5SPmORIGJGA7T-WsbQ17XGeOtEAP40S2RA6qXewEZ4,575 +torch/include/torch/csrc/jit/tensorexpr/reduction.h,sha256=IlW0-XCNfWWtBw3uIdtu79KxRpU6MvvF8Gu2Pv38UOI,9120 +torch/include/torch/csrc/jit/tensorexpr/registerizer.h,sha256=2p119qQdHGQP8HErMZfbA4la9YIuxsKyLGc2432tEvM,12773 +torch/include/torch/csrc/jit/tensorexpr/stmt.h,sha256=KX9ydeQE8JsiupZEwuNoTgJMTZCfuo02Ozif6Oc3Uq0,24074 +torch/include/torch/csrc/jit/tensorexpr/tensor.h,sha256=uzw4CIqoSprQabliELSqixZV38OH-1gNYWce6a58PMQ,10753 +torch/include/torch/csrc/jit/tensorexpr/tensorexpr_init.h,sha256=g29feo9fFx5y7g8-FTKgkyLz-067Vk1UgEbGJPhBIJ8,497 +torch/include/torch/csrc/jit/tensorexpr/types.h,sha256=avrlOuBF5XvP8gfas5K0H-Ef4wNyHa2U4WJcP-gGC60,4534 +torch/include/torch/csrc/jit/tensorexpr/unique_name_manager.h,sha256=SRV0LH8tZ03D2QOWlDf_UypE8KW8OY1taliuZ98eNA8,1079 +torch/include/torch/csrc/jit/tensorexpr/var_substitutor.h,sha256=kyXzHdMPDXrWW6qexm2NdLuZDoOckM047EG7gDpSAyk,1914 +torch/include/torch/csrc/jit/testing/catch_utils.hpp,sha256=9J4-d643UOXpEEozt3uj9IjBEOiYzO8hQl6MSrJvwWA,607 +torch/include/torch/csrc/jit/testing/file_check.h,sha256=kmx_9WCtQPYXNBEY16EtQa7hPG5_hdpQ9oUUQRDRcHQ,2797 +torch/include/torch/csrc/jit/testing/hooks_for_testing.h,sha256=qW-z-AW6Bh843JWoG1EaMaGtWnsCIBhSJ_s2BeZ6HOA,832 +torch/include/torch/csrc/lazy/backend/backend_data.h,sha256=KFer511g6xsjax-CTfHhesSd6iVyK-I9Lu0tm0OLfWA,1434 +torch/include/torch/csrc/lazy/backend/backend_device.h,sha256=FHbt6vo6-M7P6RPmTEHsn-wci86hDzaUKpJtokkDzmo,3165 +torch/include/torch/csrc/lazy/backend/backend_interface.h,sha256=LDrHg3jqgn--27fk9LjLLPvFah-BC5CAJeNyGfpQ0i8,5057 +torch/include/torch/csrc/lazy/backend/lowering_context.h,sha256=ByjZ8kTR6xJQv3djqkb8-FVoTwU7KiaMhQod0sCxbnI,3492 +torch/include/torch/csrc/lazy/core/cache.h,sha256=NPX66NO6k1kTnOuxuoXhvqDKeSggJgYdX9ycmXZS_0s,3979 +torch/include/torch/csrc/lazy/core/config.h,sha256=NTCfI1DLjGKJ_CenoDuhGgOGFA1CNiRKQIPikR1UOTk,1165 +torch/include/torch/csrc/lazy/core/debug_util.h,sha256=WofYQtKMOilg_NYSJJGOVpp0-j620_Xow_uq4OKOHdg,1528 +torch/include/torch/csrc/lazy/core/dynamic_ir.h,sha256=r5M9Jz8i2QslpW347zzyDieFodJ2RKDqeouU9CLhr_o,1651 +torch/include/torch/csrc/lazy/core/hash.h,sha256=JqcW9fbKY4LHGuK2n-TKi3FXzIlOMW7hEftR_o5Lvzo,7853 +torch/include/torch/csrc/lazy/core/helpers.h,sha256=spbVctMT_ExB7IAJ_I41c8683CfcuSND9YfrQiHRogY,2477 +torch/include/torch/csrc/lazy/core/internal_ops/ltc_ops.h,sha256=u7eSsjmKarTO6HoXAQmGB6ovkv4heVKITcGW9z1CdB0,1704 +torch/include/torch/csrc/lazy/core/ir.h,sha256=x-wp2s94Pd5-Ap5BuHPNSY1KLDBqhGmRmJJdypD6_P0,8165 +torch/include/torch/csrc/lazy/core/ir_builder.h,sha256=LFFuUi2tQnknn1vAc3NCMbqYsZkeQ2nfNTgslIKwOEo,4963 +torch/include/torch/csrc/lazy/core/ir_dump_util.h,sha256=wTnx_SqhzRO0Iqj_l0XOZyi2VlG5pxl9c-4x7ZW8Gh8,918 +torch/include/torch/csrc/lazy/core/ir_metadata.h,sha256=4gKlybk5kMnedrHtuxn2VFPw_GA9froaCjFyz74eTmU,1577 +torch/include/torch/csrc/lazy/core/ir_util.h,sha256=N8jZgssPrh2RquBZ1MPnCn76GLaRwP7flOdYfoNbpNQ,1620 +torch/include/torch/csrc/lazy/core/lazy_graph_executor.h,sha256=vZ0LzZKUQDKZW9aHmDTnX2VG9CPcqVK6O-WwjW5gG18,15285 +torch/include/torch/csrc/lazy/core/metrics.h,sha256=wlwOLg_QznDs-jpVtpLZUXH1o3y8Hcp7DM1MGgqcb0M,8505 +torch/include/torch/csrc/lazy/core/multi_wait.h,sha256=g_DRC2nbclyo9wPg59GPkARRRkCblYJJNqIoVIxGYn0,2003 +torch/include/torch/csrc/lazy/core/ops/arithmetic_ir_ops.h,sha256=tm0XXxCP76MrNTvigLQ1ndJk6KuqoMSqcutoDtn2iMU,635 +torch/include/torch/csrc/lazy/core/ops/utils.h,sha256=GImlZp7FQGsmM8HqgcEIAvJaugEUSi1Ti4QrJd1X_R4,1239 +torch/include/torch/csrc/lazy/core/permutation_util.h,sha256=Q9MUTbUEkCcM9J5muU9oNYCgBRZmcC2bdldsBM-N5mY,1506 +torch/include/torch/csrc/lazy/core/shape.h,sha256=hkUJcwT7We7i78k6tn1IJYqxBs6ZbQOFu704igBeqQg,2269 +torch/include/torch/csrc/lazy/core/shape_inference.h,sha256=T4svTZ8aGNx-uMpB-YyO0JzP2brFqyAt-Zuo64zMOX8,15656 +torch/include/torch/csrc/lazy/core/tensor.h,sha256=cM7lEIXw4D6wbYk1GetFdqNqai4NIKJXOaz8ZYmSWrc,10134 +torch/include/torch/csrc/lazy/core/tensor_impl.h,sha256=dYQ-izxXl2wfr3OfEkUi-b_FKG-i-fhxuWuPp8CLkGY,2170 +torch/include/torch/csrc/lazy/core/tensor_util.h,sha256=PrHDUfoMAsIrCQMMM2EmwU13pZBR7kaupc0wdWS2upc,2790 +torch/include/torch/csrc/lazy/core/thread_pool.h,sha256=K_THTX12OXQHK-6-Hz9rKnkZOZ-LiGzgRf85XjTGm14,1046 +torch/include/torch/csrc/lazy/core/trie.h,sha256=BSxr1ajdACgUhkgPOat-UgRU54-HPj6YdWFoSvv8juw,2446 +torch/include/torch/csrc/lazy/core/unique.h,sha256=I9Bp4qmGEQL4iBw8gIVUFf6yTE25bmc4T_EOL6aP7eE,1427 +torch/include/torch/csrc/lazy/core/util.h,sha256=HpMCdADPJ5pddLtVHMxCFU1Wz3cRbK1rRzRR-unOjNw,3127 +torch/include/torch/csrc/lazy/generated/LazyIr.h,sha256=_sMPAn8432z8zxwyCin1a6Ry-eIAwz6xeaI1kbtSXZI,296076 +torch/include/torch/csrc/lazy/generated/LazyNativeFunctions.h,sha256=1-U98BOSCutmRtKFomtrqB18Aawfk-C5O4q_IYwVx6E,23104 +torch/include/torch/csrc/lazy/generated/LazyNonNativeIr.h,sha256=gGlt2iV9CIh2ZhzFATnhnRO7jwtPSTDxDw48Op5k4Gk,4197 +torch/include/torch/csrc/lazy/python/init.h,sha256=NZYCX2lL9lwbHUnsXw5YgWxETS4lqMJn5ZPCUuz1_nQ,478 +torch/include/torch/csrc/lazy/python/python_util.h,sha256=gWzcbBoLIC17LCUycSWsDWXr1ZnqzQjvFibyIfpaA_0,569 +torch/include/torch/csrc/lazy/ts_backend/config.h,sha256=JPAgsEYiOlNHlNWphPlOEQ77bRs_cvscW90lKvdXqKc,460 +torch/include/torch/csrc/lazy/ts_backend/dynamic_ir.h,sha256=cyUKC1y86Db3LhHl85BrRoQDvtcSZsPsywGw2Z0AhCE,2610 +torch/include/torch/csrc/lazy/ts_backend/ir_builder.h,sha256=R-54wutbGLkaVoEjX71KPDenJHRvCQdX_67CxuVOlaI,2645 +torch/include/torch/csrc/lazy/ts_backend/ops/device_data.h,sha256=llMQNQtVtsJNjTASKtIsTqHpOZXvkA6_WgZQo6c2FFA,1550 +torch/include/torch/csrc/lazy/ts_backend/ops/generic.h,sha256=YfpQxCCAq4_QoJEy886VfiMFn4LhD8tzd2FPdnZOWzs,1682 +torch/include/torch/csrc/lazy/ts_backend/ops/to_copy.h,sha256=QP6Bl2xw3qUxqY8wzdAm0bJJeprhXBX_SQpFBRGFz9c,4333 +torch/include/torch/csrc/lazy/ts_backend/tensor_aten_ops.h,sha256=zjZsB7ks-66zjdNXNjB1yzaaJfgbbLHcYZwGSFCIXuU,777 +torch/include/torch/csrc/lazy/ts_backend/ts_autograd_functions.h,sha256=PF6QsWXwmBJ8as2Bi1S4D6A-QwkNUVdsBEyTnCxsA0A,880 +torch/include/torch/csrc/lazy/ts_backend/ts_backend_impl.h,sha256=RmrFgmvb6YpwJ9OfalU3ykyLbB36Y-QH57WZUkixh68,1494 +torch/include/torch/csrc/lazy/ts_backend/ts_eager_fallback.h,sha256=zzuu9qy_l3DK__izeiDULKw8BXVHH6_3_VrwD9QZ7Nw,946 +torch/include/torch/csrc/lazy/ts_backend/ts_lowering_context.h,sha256=JLTgyK1_rv0JURMXdpkxypmoeJU9XS1t83_kOZcgEn8,4766 +torch/include/torch/csrc/lazy/ts_backend/ts_node.h,sha256=svuUBEC01SrGzKuSktuyePEGf5LHLij6dzDiSgo8HsE,3605 +torch/include/torch/csrc/lazy/ts_backend/ts_node_lowering.h,sha256=Ae0WcTTYggl-ltP4TJhF4CY_fWLsw6oiL5WW2DNhJWQ,720 +torch/include/torch/csrc/monitor/counters.h,sha256=1M8nPXP8kvJh07G84-8Er_Yxw4aCOx4s13hvkkGT0ME,8097 +torch/include/torch/csrc/monitor/events.h,sha256=gMjVvEPEKxzWiDz_836HQxHw780l01DOxKGZjDw0HXA,2936 +torch/include/torch/csrc/monitor/python_init.h,sha256=ffmFgZ0KuXD-55U1Cl5L5NILNxO1GmFy55Giy78OCKA,381 +torch/include/torch/csrc/mps/Module.h,sha256=Apu7gii0dgPwg9bKgKnTjByEbqlIGalnAcMe_xO8okc,427 +torch/include/torch/csrc/mtia/Module.h,sha256=YFWOdviupcuSSKOYYRkRc0WQvAzgTDm2jFgYg1fMg1M,432 +torch/include/torch/csrc/mtia/profiler/MTIAMemoryProfiler.h,sha256=u6o2GVJoyjszRJ-PtMxVhaqDIjYFPhyKGaxKPM7hUs4,801 +torch/include/torch/csrc/multiprocessing/init.h,sha256=MDM3byavkbVz0s27gK41WWSm2QXYY_D5CzdeQZum3-g,422 +torch/include/torch/csrc/onnx/back_compat.h,sha256=NdCcjbbsoMV9cfLR7VgZG8ZY6Xx3QW67FEPpfBX19VE,1278 +torch/include/torch/csrc/onnx/init.h,sha256=nofikk1oLIWEWuDOtCxajPECVUwk1wfwxT32f-xeuGY,400 +torch/include/torch/csrc/onnx/onnx.h,sha256=iB8JqbKgTU0LEsC-MyZ9CzQMm-vDTwObYrg2JDfyBcU,759 +torch/include/torch/csrc/profiler/api.h,sha256=xGirYXerLVr29MBGKasBpHlty_2R_ubR6PO3tKX1O1I,764 +torch/include/torch/csrc/profiler/collection.h,sha256=9ukrKO6LKrc3sXCce4pdMyetDVy6RF-QZV4KMXEC6BE,21477 +torch/include/torch/csrc/profiler/combined_traceback.h,sha256=dKxMUihbgdnMbjtHvMQ8yUpUesaEC5UyRiNijtQMigo,4158 +torch/include/torch/csrc/profiler/containers.h,sha256=WDZSvZ25EhAZxa6QLBLacg7GG6U3ar_tzJf9sBB7aHU,6248 +torch/include/torch/csrc/profiler/data_flow.h,sha256=q2ojqcszbAWhZip6MuCo2C29sArfBFA5niAMTLuH94A,3880 +torch/include/torch/csrc/profiler/events.h,sha256=QVTGHVeO3B_dz8xeUEEH1EUpK_QxGWFcc2f7-VmXsyk,1290 +torch/include/torch/csrc/profiler/kineto_client_interface.h,sha256=k6Da71DcfA4CuBppIXQHFWx_MzxBxa1DKkQJH8wXBAs,500 +torch/include/torch/csrc/profiler/kineto_shim.h,sha256=ud1EfYiSxQ-K5YDMCEqQcHR4FwO2ZaaZxS0-7gTQki4,4343 +torch/include/torch/csrc/profiler/orchestration/observer.h,sha256=JONGVhaHu3nMVicqi-8SHR-CUKoSBQRpmyUYbHpWMig,7774 +torch/include/torch/csrc/profiler/orchestration/python_tracer.h,sha256=xtE70933sen95GAWc3N2hCTimRsoHukfVRmFBeAjmDw,2558 +torch/include/torch/csrc/profiler/orchestration/vulkan.h,sha256=JnVjekrFW8Syv2J0DIRh9esqtM6UrgjlSeIafZB3m5Q,1030 +torch/include/torch/csrc/profiler/perf-inl.h,sha256=P3W5T5ZGJO6H6kyGVG6jdJ6IZDz6KN7-j8YhZY_HMng,1658 +torch/include/torch/csrc/profiler/perf.h,sha256=dX35-crTvUn84O5tX42GJFFEyTjvsmLeC3tXHd-LW8I,2847 +torch/include/torch/csrc/profiler/python/combined_traceback.h,sha256=_a1nMBwi4_f_CXlvXoLDI_FrL4lWKhkmb2m6e46q8Qo,1246 +torch/include/torch/csrc/profiler/python/init.h,sha256=ZxXqiKbTO0_Ym1-doIZBz9UgqpeLrJR5U6hSfnnSa4s,1252 +torch/include/torch/csrc/profiler/python/pybind.h,sha256=6yCJymuZHstk5mhrEKwMtL9CHtFIc1yZ0-4U5DokRzs,1515 +torch/include/torch/csrc/profiler/standalone/execution_trace_observer.h,sha256=XEu_BQoKycd7zvUbK-8qo6MsoTZiDlGTyQB4QlS6ndw,882 +torch/include/torch/csrc/profiler/standalone/itt_observer.h,sha256=e1m64SgR2WYIjNv8SpkugKFihDL4L0PoYwnZaGPjXpg,478 +torch/include/torch/csrc/profiler/standalone/nvtx_observer.h,sha256=C6vGiw36g4tlZxw_gYFxGwzHkIPdp1FHKJdbCEICyoM,479 +torch/include/torch/csrc/profiler/standalone/privateuse1_observer.h,sha256=7lrsKkoS165BrnTYC1BWwu-qiOTtq2Y4jFwxVTwTcXE,1675 +torch/include/torch/csrc/profiler/standalone/privateuse1_profiler.h,sha256=B-sdxsFd7UKd8EQgSnQK8JbCrg58SRYY7uNXFlMoGM0,4012 +torch/include/torch/csrc/profiler/stubs/base.h,sha256=F7fUVssVS6_IjHPZb5thWw9vro26s8ykVwsazBds5hE,1939 +torch/include/torch/csrc/profiler/unwind/action.h,sha256=9LFKtggxTENmyAtkeJyB-mNJkt6Sl_SpG6_js89wJHc,2233 +torch/include/torch/csrc/profiler/unwind/communicate.h,sha256=9Hw2syApVGUsDyKQiAKXmzWwu9G2xsuIEJGpf-UTFzw,2506 +torch/include/torch/csrc/profiler/unwind/debug_info.h,sha256=_5sgKQQ3zaBPE5Qwo9-RtMmKgFXyd9SKvL9QvPWI6qU,9354 +torch/include/torch/csrc/profiler/unwind/dwarf_enums.h,sha256=FiB4IwZLquC7J5QSlDCZeQSa7_YutAutZX4wi1Umkqc,1397 +torch/include/torch/csrc/profiler/unwind/dwarf_symbolize_enums.h,sha256=02cDTcepBwtkBB2xl6j5kRR3439c7zZv5KDsZs5Moi4,4912 +torch/include/torch/csrc/profiler/unwind/eh_frame_hdr.h,sha256=NGXHq0_cbbCQ5O9GNSexDktFO2IOiXCbIgPPaLl_NQw,2934 +torch/include/torch/csrc/profiler/unwind/fast_symbolizer.h,sha256=tBeWMSZVhAUxibcIYVB9iqwU1DV9pFI_qqCeO8ywAMM,3557 +torch/include/torch/csrc/profiler/unwind/fde.h,sha256=wjsRFhNtKODHk5UPYns-nZHI4fXQmamGNonaeneZW8I,12874 +torch/include/torch/csrc/profiler/unwind/lexer.h,sha256=qteUGXS78lP8Q6EJU2JanNwUSMqRnkUx7GZBs6Vfws4,4135 +torch/include/torch/csrc/profiler/unwind/line_number_program.h,sha256=93WN5uZJIAdg3zdESlMPDOALhphl8HZKUxvs-Vc2UKI,11079 +torch/include/torch/csrc/profiler/unwind/mem_file.h,sha256=1wXn3uYzPJ80v5-XOfmhAncaubiw_jjv-bt686idwpQ,4843 +torch/include/torch/csrc/profiler/unwind/range_table.h,sha256=2fRNVYzs6KNiFkvveN9awqa5oNIxwcxIi8s2F52-0BU,2363 +torch/include/torch/csrc/profiler/unwind/sections.h,sha256=i02YUf5QMuuzJgdfAifWcIOZl8P3Ip4S5SrnqkZKdug,3922 +torch/include/torch/csrc/profiler/unwind/unwind.h,sha256=xaMH9xCsVyhCHotSw6e73aqNRz8d-oCduQlP4vSGtFQ,1387 +torch/include/torch/csrc/profiler/unwind/unwind_error.h,sha256=6SFMi81YnkNMvPRi4_ar_4PMplmX2VJmAP0sVGJZSYI,1152 +torch/include/torch/csrc/profiler/unwind/unwinder.h,sha256=I1rnvUYCiIbqTUfkuY6LECFlOXnzAdTZkgqxU8P3rl4,2672 +torch/include/torch/csrc/profiler/util.h,sha256=8cpmGim5ThzCK07BImF0Kt7m-BnnnPMxxrd6PRmFuZ8,7191 +torch/include/torch/csrc/python_dimname.h,sha256=myFoabBPfQqnBLUwqLepYasPlIhxdBrnHLf4HJxOSxM,468 +torch/include/torch/csrc/python_headers.h,sha256=W1ASyJDtRGr7YTipLtmOGZ6N2gzkNjM_8ju-vWUwi-w,903 +torch/include/torch/csrc/serialization.h,sha256=3_Zm2oq56ytJX1nTZXV1xGdum9BkuUusbG739Tty8wE,935 +torch/include/torch/csrc/shim_conversion_utils.h,sha256=WzsYcnAgrlZQ5lbDSZK7e3hTM4vAoVoQbIDPnYFlrQ4,926 +torch/include/torch/csrc/stable/accelerator.h,sha256=zaFRdnbdCYUTQ21uay8jr8ExHgensZ09ZlfZa-ZQuXo,3154 +torch/include/torch/csrc/stable/c/shim.h,sha256=YQsMH-MMKQKQktSaxumSGLWLpOj4jNXm8vYBMuJQocA,8067 +torch/include/torch/csrc/stable/device.h,sha256=c127B_czfZfuAQikoKN5kuiOj3eK-zWs5aMJ09mz0-8,101 +torch/include/torch/csrc/stable/device_inl.h,sha256=9qoxIAI9UQb9TRi9JWDsWqcPZgH-nkuCD7r-BjC2BrI,1349 +torch/include/torch/csrc/stable/device_struct.h,sha256=ItDv-YgxfSNzDo_IZf9vvA6_vlbC5mfqDJ2YafDii5w,5111 +torch/include/torch/csrc/stable/library.h,sha256=3WCxJThVGS8T66fN2A6BSDOH5kUuLaW7B28rqbWH6no,13092 +torch/include/torch/csrc/stable/macros.h,sha256=vild052zPBgRQSSnc5CPpUrKJV5F81-z9fxUIrc-lKw,1183 +torch/include/torch/csrc/stable/ops.h,sha256=jjAe-zFmTonbYvDDLyB4tqGYhUUbcdTDhC-5zkoFpFA,42768 +torch/include/torch/csrc/stable/stableivalue_conversions.h,sha256=hTVWcPL6rHbAxeX0-ZVdZ9sHKG1LgVGP4hi6YN8tT5c,39867 +torch/include/torch/csrc/stable/tensor.h,sha256=MIwb6Kq9H9Ry_wh9uij65SEUpw-JYddJ4UI9YFMvf6M,101 +torch/include/torch/csrc/stable/tensor_inl.h,sha256=YyQEmbrqJoZcXaS9QuvJ2uHqBWKaCQ2B1VOABGuE7SI,3172 +torch/include/torch/csrc/stable/tensor_struct.h,sha256=hhi_jLgU7fETnM78kxTlPqEsUpS-ZNu6w4falbumbKo,14093 +torch/include/torch/csrc/stable/version.h,sha256=9YK-Wki3fwlaev9WmRCSWFUAo9Xr_FM3hsXnzPC8Z7I,1183 +torch/include/torch/csrc/tensor/python_tensor.h,sha256=W51E9SLaphq9d3dR2CqRLTiOnNrh8q2-hWuFPK4UGDM,1454 +torch/include/torch/csrc/utils.h,sha256=MEEk2pvmH89d8q08Yr31hjOXp9A3xi9pLlHWuUoEDkE,4884 +torch/include/torch/csrc/utils/byte_order.h,sha256=BhoYR327QBjuW7_X6MfbdW2ZcV2bi5x4XOsNuHJGDuw,2508 +torch/include/torch/csrc/utils/cpp_stacktraces.h,sha256=rMp2epRbiAYUnwzeEFtDKUgh4qQtsJfNcJb97fVluDY,484 +torch/include/torch/csrc/utils/cuda_enabled.h,sha256=tHkNXPv0rKNUg_ld-gA61f_7jYZ7AgUf1SkayPhxVhI,424 +torch/include/torch/csrc/utils/device_lazy_init.h,sha256=Lq2sY2b3o9L6sOsLVW3UJ3MsZ-_2NxjRJw4AYLDWF1I,3068 +torch/include/torch/csrc/utils/disable_torch_function.h,sha256=ntEGSYFVYb73xt4ar8X9M03q7ErmDb_KjuEbmqL5zcI,2164 +torch/include/torch/csrc/utils/generated_serialization_types.h,sha256=myz8tbNb5F1DHkDfSBZWNqADq8Nzj6fggrBvupED-kg,125101 +torch/include/torch/csrc/utils/init.h,sha256=yFogmRgeaTycae64CMu1uQsaqrZ28T0XiBmcHTRLcqc,447 +torch/include/torch/csrc/utils/invalid_arguments.h,sha256=OYo7VRpe44aNaGGvVJIL7xPSI3rBFkGUHtRGYQyEEX0,556 +torch/include/torch/csrc/utils/nested.h,sha256=eBsEENIGSfa5kjay2Xmc3AIgUDRRN4D9g5A5nKDsJ-o,560 +torch/include/torch/csrc/utils/numpy_stub.h,sha256=iSr881WC_2tdoB55OTK2oUk2yXL8OscXS5oL45TOY3A,653 +torch/include/torch/csrc/utils/object_ptr.h,sha256=SqtlMKOikYchibyiqOvsV8MNUDvsEiS_4_X5r7QM7Kg,2254 +torch/include/torch/csrc/utils/out_types.h,sha256=mwq_iEveRZqZR9NfUx-XdC4Kr-ltwSVXsi7i27v8RrE,574 +torch/include/torch/csrc/utils/pybind.h,sha256=xa8X7QCUioJ5WC0L87bw1-ADNlK0MH7l40ii1CLD7fA,13429 +torch/include/torch/csrc/utils/pycfunction_helpers.h,sha256=OaKMH7DhZUQeC6EjF-XjrOK4OW1vfZLP6BUzXQ41-EY,1061 +torch/include/torch/csrc/utils/pyobject_preservation.h,sha256=rUW1oCr5l4bbkZ199tXlkDktOVuiC6rooCkzjF58yOc,982 +torch/include/torch/csrc/utils/python_arg_parser.h,sha256=GCsd09IwgHfWGrfhtudpQxzpc0IYry2gPJMZs0UJ3dY,43970 +torch/include/torch/csrc/utils/python_compat.h,sha256=Vr5VaAfrlhP9EJw5DDAa38PuiZAThz_SNOVB0TPe9w4,1222 +torch/include/torch/csrc/utils/python_dispatch.h,sha256=jqjX8Y8ebUUoXE6UZOybFRzMLM5VR94uZmu4fYty7Fg,651 +torch/include/torch/csrc/utils/python_numbers.h,sha256=1jrABqT3xL7BQzUwi5oNfjscA0Z-kxKUWEtano1BxZw,6383 +torch/include/torch/csrc/utils/python_raii.h,sha256=JLoYiMGFICxaIMoH6u5e4vp6SGE4W2N_EpTxD78NlV4,2912 +torch/include/torch/csrc/utils/python_scalars.h,sha256=JLGvlzGLJHFu1PLBKyxEcAc8OVsYHKNcVlhKG-4WkUg,6282 +torch/include/torch/csrc/utils/python_strings.h,sha256=lg0THu6tMRvIsihohsa49N9EpO3BBIt6ogFLv6fegCQ,4857 +torch/include/torch/csrc/utils/python_stub.h,sha256=JqFk4mKCyu0chr2QF5a1IqXXbi9HPecVk_VhA2TxtCA,310 +torch/include/torch/csrc/utils/python_symnode.h,sha256=TsvNCunai6IiBOL_wAcNpqeWMSh_ftOmtxzT-DPzd4c,10375 +torch/include/torch/csrc/utils/python_torch_function_mode.h,sha256=DxJo0hwTAG1hkq5LWJgXK58NUHiS4CDAbsrsBK_SkVY,1094 +torch/include/torch/csrc/utils/python_tuples.h,sha256=I8rMtebgRbe52cgReGUGwmmNk2HBsukIZaGg3V4VjkA,980 +torch/include/torch/csrc/utils/pythoncapi_compat.h,sha256=9ZKnncRIDr_22WnehcTVL4mLVJ9S1t3_sO9KC_Ch7D8,72696 +torch/include/torch/csrc/utils/schema_info.h,sha256=kAn1bspmNcHr91si2o4ei6k3fuQ1-NvtoT-ISRENK6M,4043 +torch/include/torch/csrc/utils/structseq.h,sha256=2G-lkmF_SYF2ECFSJwOLA6HBYFctjEwcgt924a5jfkw,395 +torch/include/torch/csrc/utils/tensor_apply.h,sha256=utGr0N-b-NAYkc68JEZcyeKrNncROANSHTpEQjidt4k,682 +torch/include/torch/csrc/utils/tensor_dtypes.h,sha256=cfAjJqa_c3rgNnj3ZxtnCRK-Ozyaykiwy9TfeXbilpU,496 +torch/include/torch/csrc/utils/tensor_flatten.h,sha256=-JgUD7iNFqUTEhe6bDougZZkdibYt3QuRmN7CPJBD-g,2999 +torch/include/torch/csrc/utils/tensor_layouts.h,sha256=9iqdB6L0Fzm8LXUjOKgSg8JDxxPNAvc6p_YG_YCCqOc,323 +torch/include/torch/csrc/utils/tensor_list.h,sha256=_dWalILXx6uIho9CQRwN-RW7hiVOikXfvR1Vu7geUYo,421 +torch/include/torch/csrc/utils/tensor_memoryformats.h,sha256=nV8q2ZOPt1auubOHrtFZvobiPTsqvo2UUT9ohUAi0qQ,600 +torch/include/torch/csrc/utils/tensor_new.h,sha256=3p24CSRxCljCviIdlOehwuz7jRXW1wWIFfI7epvhuI4,4300 +torch/include/torch/csrc/utils/tensor_numpy.h,sha256=CKuofC4dLDOzDQd7lBuoTWcagNQp--s7Dp_TrwVb6OQ,1129 +torch/include/torch/csrc/utils/tensor_qschemes.h,sha256=yOncI33ZbiK-oPvZEnvWapNXE0wRSqzZNit4zyCCOz0,428 +torch/include/torch/csrc/utils/tensor_types.h,sha256=z8Fq8XAQUE-NWt-AeOelujtTbpzllUyHm3H4D-Xd3EM,923 +torch/include/torch/csrc/utils/throughput_benchmark-inl.h,sha256=JmcH6qIqFot8fp8cLSVF6RIyEMdbGNyCxjUYZY5dP28,6590 +torch/include/torch/csrc/utils/throughput_benchmark.h,sha256=WcS0hu1H0O6aqlflKnQyuwo9V15e9sv4NR-zA0x2WkU,7191 +torch/include/torch/csrc/utils/torch_dispatch_mode.h,sha256=c-D7EjkxhFiKGeEzuGWzAbE-Iq8rVsqmrklnwamj8G4,2529 +torch/include/torch/csrc/utils/variadic.h,sha256=gbXQ7bGebPOoVsGoLjc4RycltMc1BYLTKcUv_VxUGw8,3630 +torch/include/torch/csrc/utils/verbose.h,sha256=YMrRuUzOz8RuhXUsyIH2uzVkHpXGp6iWrV1mgVwrpqY,392 +torch/include/torch/csrc/xpu/Event.h,sha256=DUja-x8C-w3C3xD_lzTVDY61kgx1PVLn1zWi8m758MQ,623 +torch/include/torch/csrc/xpu/Module.h,sha256=ITvFMk9MWU03WLLy-BOZmpoJnIVF0f70KlIO3dtoTTw,430 +torch/include/torch/csrc/xpu/Stream.h,sha256=6lPNJMuXZqw9O-_PNTb5EY9fXplRF0Jj0JfbV0xXJ7g,691 +torch/include/torch/csrc/xpu/XPUPluggableAllocator.h,sha256=ZAKYjogIIzPBamhYGB5ewStINvWxF1lRt8dCbiXCVpQ,3124 +torch/include/torch/csrc/xpu/memory_snapshot.h,sha256=VSJFwIr4rNRd2iY6UrwQAaacCT3DCbt52jQl4UBdN5o,730 +torch/include/torch/custom_class.h,sha256=tTZxsho386NSvF4bRp4L-fNlXiWtlAfWN-NFESIEFTs,20076 +torch/include/torch/custom_class_detail.h,sha256=BIE1P3Ie7e5Ma5Ii9I1_xp5rlR_ZyY95ADZ68LYGoUU,8086 +torch/include/torch/extension.h,sha256=2juzxgbeXcd8wDHl9fCm-8iqxJ8YnqutefZANPuWnkg,467 +torch/include/torch/headeronly/core/DeviceType.h,sha256=3VX6zWVEIf5ChLy6EUKdXEOLMB0I8uIO4cI6Pc66CS4,4460 +torch/include/torch/headeronly/core/Dispatch.h,sha256=cdTOahmNMCFJ_1qKM62sYWYRzooKV7_3nfUY_XxgJ34,3825 +torch/include/torch/headeronly/core/Dispatch_v2.h,sha256=CzUSaaTunLFhT8KmZ3MUNLpcB-AbeIc69-XfWa-J60g,32823 +torch/include/torch/headeronly/core/Layout.h,sha256=HFoNirxYWVYbAcPO_UtnWv3lhAwY37RKczzFKgf3wKw,948 +torch/include/torch/headeronly/core/MemoryFormat.h,sha256=K7Q7XLCpUnKaFrI0k5D1hGOQK5X7obWQAQkF7a1PVH0,1206 +torch/include/torch/headeronly/core/ScalarType.h,sha256=C7Lrm8rgVAqcvDfYmR-lvdl6VBIHZ9cIeneMOiKk9b4,17784 +torch/include/torch/headeronly/core/TensorAccessor.h,sha256=j_0qsKM8oV-CQJGO5YXi6fX9v-bn9LqT4R3tG5dlnuw,13907 +torch/include/torch/headeronly/core/enum_tag.h,sha256=99O4BSvG6nPAIsWXkytyMd2LPB75xx7ZFO9u9UiBdoI,787 +torch/include/torch/headeronly/cpu/vec/intrinsics.h,sha256=PV2_UwjmjaVTKAc1PjPerISePzJnTLKf5iynUNFmnHM,1917 +torch/include/torch/headeronly/cpu/vec/vec_half.h,sha256=kXa2ne1je1uWnwec2732Zx8jV1gNFlRa7yw1k2X5D9o,1740 +torch/include/torch/headeronly/macros/Export.h,sha256=bC2_b5PrzlYBQdkmBhe2MH2nEOADscHqdLuiBVBQcO8,5599 +torch/include/torch/headeronly/macros/Macros.h,sha256=jiu1NgIaK8gExY_IC8o4HXREyQCciXqmFmKK4i0VJp4,27451 +torch/include/torch/headeronly/macros/cmake_macros.h,sha256=GFBE7BEJHYCFrGffrJidYZkCT-srrcDCzMM65W5ZjGQ,445 +torch/include/torch/headeronly/util/BFloat16.h,sha256=zVdwWziwVmCEAigHHrDubaxG81DRvfoRSW7KZDeQvmk,15277 +torch/include/torch/headeronly/util/Deprecated.h,sha256=tyNcwEnjhdHs3LbgKssdAQqInoDJXNhgcSg1z2sj3nc,3579 +torch/include/torch/headeronly/util/Exception.h,sha256=IJVuC3dAGEfh0M5pap7hFGyRixEgOQjwfR83CFsdSUk,3422 +torch/include/torch/headeronly/util/Float4_e2m1fn_x2.h,sha256=G734w0W4oOe35KET9Tap746EWyTe7G6SbMmoNHvyt0g,1326 +torch/include/torch/headeronly/util/Float8_e4m3fn.h,sha256=C1vYaZvcThxJzUg9ZMjrDKpxhrxut8TrzLReI4Kmk84,18606 +torch/include/torch/headeronly/util/Float8_e4m3fnuz.h,sha256=hHD8n_B5ZVtktOqv6jxhtdMteaJt-5JVyNehETKcVo4,14815 +torch/include/torch/headeronly/util/Float8_e5m2.h,sha256=YCEB1vgf5rR42qYQc6VoXz46wecBtlfGl5RbXSRrBeM,14852 +torch/include/torch/headeronly/util/Float8_e5m2fnuz.h,sha256=87PtkT3hft5BedgwSL39-4KA4_KRP_dMFWs51XE0nic,14898 +torch/include/torch/headeronly/util/Float8_e8m0fnu.h,sha256=UtF9d4bS2fnTUFA_UHHK7Y3pG1it99XKo3ZiaA67OsE,6938 +torch/include/torch/headeronly/util/Float8_fnuz_cvt.h,sha256=aVJimoCmZgfahBmkRR40USG1eSeAng5hzhclZTa5TDM,1920 +torch/include/torch/headeronly/util/Half.h,sha256=r2_DJws_bzPgHK1jNzj8v8T2SZXU9sEsik97PPC2hLo,28235 +torch/include/torch/headeronly/util/HeaderOnlyArrayRef.h,sha256=f_p0joOnbl3MTpJuOEefhvRPCg_3L1AuF0d8UiQKqzo,7712 +torch/include/torch/headeronly/util/Metaprogramming.h,sha256=jPSAnrndcYsql4SqkbcffUCImefhQJxQmlNfd112EX0,7444 +torch/include/torch/headeronly/util/TypeList.h,sha256=bbsfasQ9L24ryDshgykojdyL98jbAcFdjr0ygLu4xI4,17804 +torch/include/torch/headeronly/util/TypeSafeSignMath.h,sha256=HjY_mtA-ngu32v3y_U_Kl1SS1RZlW_OE1mjyFXMGzgo,4585 +torch/include/torch/headeronly/util/TypeTraits.h,sha256=C4S0eYZ8fQZnHQAKZcnfMCBE2PrZ8m0W8djVag2_nJE,5698 +torch/include/torch/headeronly/util/bit_cast.h,sha256=-YBNY7PSadYyIoF3IyB-_RvY19evGIVh-Zt9Vi8liT8,1426 +torch/include/torch/headeronly/util/bits.h,sha256=K1SQupBTltkxUVAH28IXO8VQqNf1eLkYqGGMKKT7gmc,1644 +torch/include/torch/headeronly/util/complex.h,sha256=DBoJapV0ySW-w8QqKy3muh1YuxlunJC_4KU_Yx83Hv0,18096 +torch/include/torch/headeronly/util/floating_point_utils.h,sha256=o36InAmLzUt4OsrTbNeM2dNwCT0xbjmkOqYwewTZ1h0,1047 +torch/include/torch/headeronly/util/qint32.h,sha256=NjQBdVMVQoI1SxNNrtgEOOScvOLpgwPrTHIAmIlMSJY,434 +torch/include/torch/headeronly/util/qint8.h,sha256=JIN-Vao78wIx0308UT98URVQZdmXfYX3oTuzqlOi--o,586 +torch/include/torch/headeronly/util/quint2x4.h,sha256=vSoOsu-raTav0QIldQWUhYp9bY88d_HD3Kn_VXdRm9M,483 +torch/include/torch/headeronly/util/quint4x2.h,sha256=Pu82ZJKu21wqi232YMuAmtTWtq7DKaIek56Lr9u-N1c,483 +torch/include/torch/headeronly/util/quint8.h,sha256=2W7tsPA88e_obpE9TqWAeprIXjjZkyqwOqvX6TlGz-s,435 +torch/include/torch/headeronly/util/shim_utils.h,sha256=m0wt7-KIUyyPwP0oD2Bl1TsfLiCMEbVU4ZhsE4hnpkg,1028 +torch/include/torch/headeronly/version.h,sha256=AjZIef0mQ1XP--aXymXLCtlAzdGSJkUdtTHKEYkRaYc,818 +torch/include/torch/library.h,sha256=aacH18q2lzMY84pdeHjuukS2usY1dOQUkdUP_GOGP_U,42695 +torch/include/torch/script.h,sha256=e7ma1b60NplDJbbjKdE0vS98ttRT1GzWH4xHD6jn_MA,723 +torch/jit/__init__.py,sha256=SiT48q3MakNgr5HAMtznarbUnw8UjXXqWJ9oJMHp624,8964 +torch/jit/__pycache__/__init__.cpython-311.pyc,, +torch/jit/__pycache__/_async.cpython-311.pyc,, +torch/jit/__pycache__/_await.cpython-311.pyc,, +torch/jit/__pycache__/_builtins.cpython-311.pyc,, +torch/jit/__pycache__/_check.cpython-311.pyc,, +torch/jit/__pycache__/_dataclass_impls.cpython-311.pyc,, +torch/jit/__pycache__/_decomposition_utils.cpython-311.pyc,, +torch/jit/__pycache__/_decompositions.cpython-311.pyc,, +torch/jit/__pycache__/_freeze.cpython-311.pyc,, +torch/jit/__pycache__/_fuser.cpython-311.pyc,, +torch/jit/__pycache__/_ir_utils.cpython-311.pyc,, +torch/jit/__pycache__/_logging.cpython-311.pyc,, +torch/jit/__pycache__/_monkeytype_config.cpython-311.pyc,, +torch/jit/__pycache__/_pickle.cpython-311.pyc,, +torch/jit/__pycache__/_recursive.cpython-311.pyc,, +torch/jit/__pycache__/_script.cpython-311.pyc,, +torch/jit/__pycache__/_serialization.cpython-311.pyc,, +torch/jit/__pycache__/_shape_functions.cpython-311.pyc,, +torch/jit/__pycache__/_state.cpython-311.pyc,, +torch/jit/__pycache__/_trace.cpython-311.pyc,, +torch/jit/__pycache__/annotations.cpython-311.pyc,, +torch/jit/__pycache__/frontend.cpython-311.pyc,, +torch/jit/__pycache__/generate_bytecode.cpython-311.pyc,, +torch/jit/__pycache__/quantized.cpython-311.pyc,, +torch/jit/__pycache__/supported_ops.cpython-311.pyc,, +torch/jit/__pycache__/unsupported_tensor_ops.cpython-311.pyc,, +torch/jit/_async.py,sha256=VZ-IFUaQy4Lxj_4Op69RSpWoIpdyC8lpqcthWHcptvw,4304 +torch/jit/_await.py,sha256=UG6mM5L4tF7NRDT5BfFX4njoM1B_EhdoMT59KsyHF-8,852 +torch/jit/_builtins.py,sha256=PE1J9wwL6Xi2hBRbC6NKtRCm8B30QDSO0wQbC-P22sI,6776 +torch/jit/_check.py,sha256=pAon5PwX0F0aENhczukFdu8R59Sqob-on9KnEMr9RDE,9772 +torch/jit/_dataclass_impls.py,sha256=3xo2DKMaUZv44Ngw1VYETZ4mqbwPTyX29U1VCItl6Co,6681 +torch/jit/_decomposition_utils.py,sha256=IOugHESe09GNFdtx4geGCppHwGUzp4EPskQXoU3PqG8,510 +torch/jit/_decompositions.py,sha256=Zrc3x4XHhVo9M2je3NfGysiTmL0i8xPaLwTuC-R32vU,4659 +torch/jit/_freeze.py,sha256=hfysa-VuzO1YStEQmu1bSBmEVOYl2ARYWLNLvEEKknw,9969 +torch/jit/_fuser.py,sha256=WK0-Wb-4jxYWBekc7jR9detA9_0-FSO7s_vl0QhbAJo,7617 +torch/jit/_ir_utils.py,sha256=QskhtXHHEmqXbfFBA_R8JdMr7OCdIqqZd2JuLEeT1Kw,838 +torch/jit/_logging.py,sha256=jWaacYoHjd_Cao_1hzyPKYiADS7cb7tLQc1gRkZblUQ,257 +torch/jit/_monkeytype_config.py,sha256=54Y0yNZHjeG84imUFnNk08cl-nWJiB1eSwE19tQ7uU0,7337 +torch/jit/_passes/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/jit/_passes/__pycache__/__init__.cpython-311.pyc,, +torch/jit/_passes/__pycache__/_property_propagation.cpython-311.pyc,, +torch/jit/_passes/_property_propagation.py,sha256=4L1s-f_Ngp4Abz0Mf5t-BtQlEnyIWV2GPya6N2Om7qE,1438 +torch/jit/_pickle.py,sha256=lV929M0juZlRKP8-OwClLpLwxIb7W-74r055yqNc8wQ,1130 +torch/jit/_recursive.py,sha256=KsdaIhHyWG22jCUOV8l0Q34Vfk6K10jFj_mQgxKyNJs,42939 +torch/jit/_script.py,sha256=1M5Byj3XUAaRyXtLkeKjDcZFHQ7Wme2KnOlNRq8DnLU,68126 +torch/jit/_script.pyi,sha256=8LgMHIC1_beIEUEsOF0TvA38-oaRPY0WeocOA1-aoSk,9863 +torch/jit/_serialization.py,sha256=-sJ5EwAc7nr-193XbDmPQ1RkliQshkf_kfQSr0u1vpE,10698 +torch/jit/_shape_functions.py,sha256=_rBTifMTiYR-IwdERfcIfCIzd1KB7HfbKoaxu7FZFhE,52703 +torch/jit/_state.py,sha256=lOiv73EtiQebxiOG1o_0rMnZ1KiKVWkvDikBPLGWjTg,3880 +torch/jit/_trace.py,sha256=9pbDMRyMaa_N0EB-jKCriqnTACqD_xgYj_3p0GS0e6w,56369 +torch/jit/annotations.py,sha256=AMy0ZNG8eOwjwJtydIkX_ifOMJxyfE3isD-FyLCp5X8,18238 +torch/jit/frontend.py,sha256=-TteSBNjJesmL79r-tDOeEkDfYcnzb4dtRhIFkg6NPs,45377 +torch/jit/generate_bytecode.py,sha256=e-Ltj08oGvI5x-Q3SWgnU0v21B-YQafnhYDj3XJ_C2c,1042 +torch/jit/mobile/__init__.py,sha256=1S-ZgAQBSjkLiKnObxqWkxg50Pw2J_GNsX7R1BvHpkY,8601 +torch/jit/mobile/__pycache__/__init__.cpython-311.pyc,, +torch/jit/quantized.py,sha256=_qdmUEm74_VB2z0Dipf61uuvyefSkNAz0kLSPGxFMmA,3193 +torch/jit/supported_ops.py,sha256=Pe-f7iivLwVZF6P2z21xmLkCo9TuNw9S3EUdcu60uRA,10264 +torch/jit/unsupported_tensor_ops.py,sha256=JYLe90X6vb3BBzNB9xKY8VIRKtk2-LmsK-UXUbXMqYI,1995 +torch/lib/libc10.so,sha256=1yxtpEL9rxtLORay77nFsbo1ZPW1Tu-z6TGZB9vfzvw,1423328 +torch/lib/libc10_cuda.so,sha256=jdQRLQK51Hnz2BlV75RTzgTDa5PCGqutRukpY4X5MdM,589080 +torch/lib/libcaffe2_nvrtc.so,sha256=TJ3-yMDyx-UCnuvcueA9hqyiBqobYl3sznvCW1QTWTM,16608 +torch/lib/libshm.so,sha256=mkdkLTytt2vMZoK1ec4N8_G7yPxzBM6dYIJSDnq66yc,47960 +torch/lib/libshm/alloc_info.h,sha256=7SdVmCqiLa2mfKg-4DaB9rdpJnStPCmgXASpbjqtzBU,104 +torch/lib/libshm/err.h,sha256=j5cpgd5gS8_WgroKTyMLvDHDiXSe0GzbwE-1Sy88mA4,1258 +torch/lib/libshm/libshm.h,sha256=_6sOfY_vSO6mECr6XNIChm7r06uVO2fQRW-CZDOf5gs,1205 +torch/lib/libshm/socket.h,sha256=b_IuucKssZ6iHI_coGieeWw4aePAX9eYHly8l1BNyCs,4322 +torch/lib/libshm_windows/libshm.h,sha256=9MByK00M3an2HH2S8gFjZMUJsFJb96VRYE190Z-ALFo,779 +torch/lib/libtorch.so,sha256=ojCVnbgvxrUi3Nl0B8AAuqWz113BeeDzTA1UKxaY-68,14416 +torch/lib/libtorch_cpu.so,sha256=u30WqmVXSTU47J9BAw1uKcs3ED5duauzEhM2X5KF7_M,320251880 +torch/lib/libtorch_cuda.so,sha256=67lYvloIyt7GE66ah_HcPYI5CbVqMZdYSC7I4EORv3s,539185872 +torch/lib/libtorch_cuda_linalg.so,sha256=A8QD2rdnyXR-QG5hIkx__m3ee70Xn7JBT2qlGkYJAEs,574312 +torch/lib/libtorch_global_deps.so,sha256=cdf7A2HJz_3pn2pWkKaWepjIVYlYdizMb5rDu3GKcnM,14584 +torch/lib/libtorch_python.so,sha256=3-NnB2P1s0Dh2zIU1WUctZmAPd0vGMgHBdGqZc3rmmg,31291720 +torch/library.py,sha256=W0cO8XQwEkXnJbUPBGZAlmDzk3mfHlKDwxQZ3dG6gSI,68855 +torch/linalg/__init__.py,sha256=po-2QVSq2_zESTq8rl03s2ommvKgancirhYUAsZCDN4,116471 +torch/linalg/__pycache__/__init__.cpython-311.pyc,, +torch/masked/__init__.py,sha256=WljBC5z-aMQ3EB29GNw5BVnONrx14zh1Cgrz7fqJ-1E,928 +torch/masked/__pycache__/__init__.cpython-311.pyc,, +torch/masked/__pycache__/_docs.cpython-311.pyc,, +torch/masked/__pycache__/_ops.cpython-311.pyc,, +torch/masked/_docs.py,sha256=Jr-iD8MG_4F6Ip_yxboE8Or1ms-bo3ay_6pdggbfV44,49468 +torch/masked/_ops.py,sha256=TwGwveRib62Ad81KuZUr1NpOT25Ekxz-Fus7DvmQ4uY,67692 +torch/masked/maskedtensor/__init__.py,sha256=yRXNgVmvSu6WET5dYVroutE7gtiXkN1JuBwogvcTvsY,344 +torch/masked/maskedtensor/__pycache__/__init__.cpython-311.pyc,, +torch/masked/maskedtensor/__pycache__/_ops_refs.cpython-311.pyc,, +torch/masked/maskedtensor/__pycache__/binary.cpython-311.pyc,, +torch/masked/maskedtensor/__pycache__/core.cpython-311.pyc,, +torch/masked/maskedtensor/__pycache__/creation.cpython-311.pyc,, +torch/masked/maskedtensor/__pycache__/passthrough.cpython-311.pyc,, +torch/masked/maskedtensor/__pycache__/reductions.cpython-311.pyc,, +torch/masked/maskedtensor/__pycache__/unary.cpython-311.pyc,, +torch/masked/maskedtensor/_ops_refs.py,sha256=iUIohR1D-kldMvpsf7fjs9do2tAThFcfA6PRwAkD-YA,18131 +torch/masked/maskedtensor/binary.py,sha256=8cYTvBPO3P-8I99WQrVr1oLMDrXsfcWPwU-HMu6I4Dk,5496 +torch/masked/maskedtensor/core.py,sha256=5m6EdCtOXQfu76SnF5jzU9JdYIi2aREPq9CmaBq_0_Y,13059 +torch/masked/maskedtensor/creation.py,sha256=cuKrXLyknX42q3rbZOmrjmqgrFudey8GUh5Fmpu5YOs,605 +torch/masked/maskedtensor/passthrough.py,sha256=kEDm_BwgAN2erALo2NF8SLWX8J9jrAuySd9upW15-40,1447 +torch/masked/maskedtensor/reductions.py,sha256=lpGZNoL58xnxvhKntt2YMeZZUqcrGAzwijWSVhaRAmg,5599 +torch/masked/maskedtensor/unary.py,sha256=F3sO8CbzFYZhpTZxOx4nobOb7RRMgtx_gbY0IRbAXmA,4197 +torch/monitor/__init__.py,sha256=09VJnyTA2zwN2clbG0o6pV3Yf4YdzW2BospVTxNa4WU,1278 +torch/monitor/__pycache__/__init__.cpython-311.pyc,, +torch/mps/__init__.py,sha256=0IajD_CjncDWF3NfilyGVNilYZ-hYJYxzo_Mr0J_SUw,6239 +torch/mps/__pycache__/__init__.cpython-311.pyc,, +torch/mps/__pycache__/event.cpython-311.pyc,, +torch/mps/__pycache__/profiler.cpython-311.pyc,, +torch/mps/event.py,sha256=tobbRc_mwlY-8W0QlYxxQQg75sPO-25nLzR9Pw9Ix7Q,1730 +torch/mps/profiler.py,sha256=TaunF3Y7c_HW26Y7ECPS6MRUFl-D93ArcSPgKT-hMyw,3553 +torch/mtia/__init__.py,sha256=uE91Ju0ptGP6Dtkguil2HPo-1GYEGng6OCOYZB_g8ew,18991 +torch/mtia/__pycache__/__init__.cpython-311.pyc,, +torch/mtia/__pycache__/_utils.cpython-311.pyc,, +torch/mtia/__pycache__/memory.cpython-311.pyc,, +torch/mtia/__pycache__/mtia_graph.cpython-311.pyc,, +torch/mtia/_utils.py,sha256=b8D_IGip8ER-8kUIfs_tglz_Sc8CG2-_4YZ2pZHB7d4,1916 +torch/mtia/memory.py,sha256=_NQP8NoU3C8ytbaPFMN0tVWONd9Lu4gLh0Xs9wnpxss,2233 +torch/mtia/mtia_graph.py,sha256=tOqt7wSx2HlxHnE7dHmZefXOAMQSEMY_ljBik8Gi-vM,2656 +torch/multiprocessing/__init__.py,sha256=5FklmQB_ltILmXpEkY1kHsEvxVnhMW_fGGXLCn7Qm7U,3552 +torch/multiprocessing/__pycache__/__init__.cpython-311.pyc,, +torch/multiprocessing/__pycache__/_atfork.cpython-311.pyc,, +torch/multiprocessing/__pycache__/pool.cpython-311.pyc,, +torch/multiprocessing/__pycache__/queue.cpython-311.pyc,, +torch/multiprocessing/__pycache__/reductions.cpython-311.pyc,, +torch/multiprocessing/__pycache__/spawn.cpython-311.pyc,, +torch/multiprocessing/_atfork.py,sha256=NDXE2HR06ENj_1bGrQumxLdTGSDOlP-2bzeNwij7ajA,790 +torch/multiprocessing/pool.py,sha256=HQtd_V3oWKnkpQRrR4cutVFdlwxWwgCiUVPYSQXaU7M,1743 +torch/multiprocessing/queue.py,sha256=OT9KTbNOpqjhun_ZQTXizodE1UOwrpSbJhYRKtAyNmE,1477 +torch/multiprocessing/reductions.py,sha256=njCm-1F9r3ITWCKArQ_x47cL5_s4mEJnUYPuddhVWTE,23096 +torch/multiprocessing/spawn.py,sha256=lqqwiBGMGAwdqQftBObfOvpyDVgqdyuTResjEfPuCj4,12863 +torch/nativert/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/nativert/__pycache__/__init__.cpython-311.pyc,, +torch/nativert/backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/nativert/backends/__pycache__/__init__.cpython-311.pyc,, +torch/nativert/backends/__pycache__/_lower_utils.cpython-311.pyc,, +torch/nativert/backends/__pycache__/_lowered_aoti_module.cpython-311.pyc,, +torch/nativert/backends/_lower_utils.py,sha256=2S0KcY-NadyeIEPGYIG6vDfG9H3nb7ac7SKBwpStX6g,3614 +torch/nativert/backends/_lowered_aoti_module.py,sha256=qBtnm7S4xagZczsqAdqcqmpLHwz-Z65tMvenrbSP3Fc,890 +torch/nested/__init__.py,sha256=00i4yf9AOOzVCXviAWuCaZHtCqvwi5ZPV4oepzvDMoY,22307 +torch/nested/__pycache__/__init__.cpython-311.pyc,, +torch/nested/_internal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/nested/_internal/__pycache__/__init__.cpython-311.pyc,, +torch/nested/_internal/__pycache__/nested_int.cpython-311.pyc,, +torch/nested/_internal/__pycache__/nested_tensor.cpython-311.pyc,, +torch/nested/_internal/__pycache__/ops.cpython-311.pyc,, +torch/nested/_internal/__pycache__/sdpa.cpython-311.pyc,, +torch/nested/_internal/nested_int.py,sha256=9FoOZfUv3dcUrDOuRsrje6zORe23oqGnALGxEMsoYDc,3319 +torch/nested/_internal/nested_tensor.py,sha256=0EE3eFBny5uu9YBgA2Hz6eH896nBu630Rki1vQpfrTM,25964 +torch/nested/_internal/ops.py,sha256=oxMmPXIysp2pTVuDLEmjZ5Z_10ydhUUWwzG2Xn8TlSQ,97287 +torch/nested/_internal/sdpa.py,sha256=4vD9zF7SJHn1Hb_0ZVgyy_c7z3t-6rBEQgCzE3pNDTo,34645 +torch/nn/__init__.py,sha256=SGakM4nNGjI1YZwNintPfgNVTV5FrVKhiFXbq5Db8Pk,2425 +torch/nn/__pycache__/__init__.cpython-311.pyc,, +torch/nn/__pycache__/_reduction.cpython-311.pyc,, +torch/nn/__pycache__/common_types.cpython-311.pyc,, +torch/nn/__pycache__/cpp.cpython-311.pyc,, +torch/nn/__pycache__/functional.cpython-311.pyc,, +torch/nn/__pycache__/grad.cpython-311.pyc,, +torch/nn/__pycache__/init.cpython-311.pyc,, +torch/nn/__pycache__/parameter.cpython-311.pyc,, +torch/nn/_reduction.py,sha256=ndJ4VxoNYRrvspg35qJtWdkGi0Cu2XsMv-sa0rBDt-s,1626 +torch/nn/attention/__init__.py,sha256=VuELb5ZcwFDbeC3U3Ecgl8mwLsW1_jqyyLBAVcCwu-A,7256 +torch/nn/attention/__pycache__/__init__.cpython-311.pyc,, +torch/nn/attention/__pycache__/_fa3.cpython-311.pyc,, +torch/nn/attention/__pycache__/_fa4.cpython-311.pyc,, +torch/nn/attention/__pycache__/_registry.cpython-311.pyc,, +torch/nn/attention/__pycache__/_utils.cpython-311.pyc,, +torch/nn/attention/__pycache__/bias.cpython-311.pyc,, +torch/nn/attention/__pycache__/flex_attention.cpython-311.pyc,, +torch/nn/attention/__pycache__/varlen.cpython-311.pyc,, +torch/nn/attention/_fa3.py,sha256=8jIppH8SAQL19e6JM9kZzM20j-iVm6Kbntrz9mHJO7g,21137 +torch/nn/attention/_fa4.py,sha256=ZV6V9u4abWZBHCqrF99RTdkmaj0tJl1JX1fhwYcENOo,14657 +torch/nn/attention/_registry.py,sha256=gD85V4lu2cwtJD6U91RttNy1oZZA0vslni0NwrAvJfY,4463 +torch/nn/attention/_utils.py,sha256=AElRsYvBfPxbPT2xyHq32GTwk_5NgSd1ucR3Jn05CMg,2088 +torch/nn/attention/bias.py,sha256=xIREKtYVtoV0C86cfKNtJyeY74eh2MSKc23mjDqIqHk,13782 +torch/nn/attention/experimental/__init__.py,sha256=VXpUxbEIcm46WvCIT1rgWY3vSN7V7UMT42FqE7UR-9M,111 +torch/nn/attention/experimental/__pycache__/__init__.cpython-311.pyc,, +torch/nn/attention/experimental/__pycache__/_paged_attention.cpython-311.pyc,, +torch/nn/attention/experimental/__pycache__/_scaled_dot_product_attention_quantized.cpython-311.pyc,, +torch/nn/attention/experimental/_paged_attention.py,sha256=r9zNXBmL718JyeAqKWPBU_DRUxyg9gGA1B4J_8Db448,13216 +torch/nn/attention/experimental/_scaled_dot_product_attention_quantized.py,sha256=RwS2dWSK_Nv85Yct-tZ9RWuANeXCrNDNH85R2oWb8Ko,5788 +torch/nn/attention/flex_attention.py,sha256=CclDbi_MRQmTT2B1YiTZxEAIoiT86kDMK1aISUC-GJE,80633 +torch/nn/attention/varlen.py,sha256=L1OE4LyM43HQChwJ04rQGVFwCXmOfLNDT1bPS5-jUeo,22757 +torch/nn/backends/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/nn/backends/__pycache__/__init__.cpython-311.pyc,, +torch/nn/backends/__pycache__/thnn.cpython-311.pyc,, +torch/nn/backends/thnn.py,sha256=DIn1LJwiPPLheO-ePySpWoHesiuc7NUEcfM4X-hyE2Y,154 +torch/nn/common_types.py,sha256=6nu95uJFHkfNYQrVRxzgRWSF-4tVNSzIF7A1h3HWiMI,2089 +torch/nn/cpp.py,sha256=qiuQgs3CT0t-Vy1FcJsrm08DQSHnMcxG09AWIf2hu_A,3100 +torch/nn/functional.py,sha256=J0kxhu4i-BG1U-MdnIBNTUZxbRvmLQNNcxU39m8n7xk,263236 +torch/nn/functional.pyi,sha256=MpdKSw93GzORsnnhrztyXt655YhuqUi_9J7zzay6cj8,27812 +torch/nn/grad.py,sha256=xEFG-ZnkxxjBMJajUSoNPqlHLiFFFAZ3wXBvn2kCRYc,9910 +torch/nn/init.py,sha256=6rv8st2GV8ibatr2tTGKO1PP4DlKRFxjMWNHlAtPFPA,26625 +torch/nn/intrinsic/__init__.py,sha256=mLJtyO2jadna-iZ8mCcW_NoEMq4mJxr3BqaJDy0jBhE,697 +torch/nn/intrinsic/__pycache__/__init__.cpython-311.pyc,, +torch/nn/intrinsic/modules/__init__.py,sha256=bqi9D10UvQcO_cimMFdlI5HjLEvFOp7Y0hT6g3WOSMQ,517 +torch/nn/intrinsic/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/intrinsic/modules/__pycache__/fused.cpython-311.pyc,, +torch/nn/intrinsic/modules/fused.py,sha256=Q6ta4U7JgPoSVs66FPll7Hjhjl_pXcO3pFWKGnBzspk,563 +torch/nn/intrinsic/qat/__init__.py,sha256=7rFxYP3UVmm4gwSlC7aBqXoJF1s0Q4aIdJZuW7ST8wI,59 +torch/nn/intrinsic/qat/__pycache__/__init__.cpython-311.pyc,, +torch/nn/intrinsic/qat/modules/__init__.py,sha256=n9s1HTHyr4Ao5_38irURngF2ZPh7pTJy9eos_oIz3Bs,637 +torch/nn/intrinsic/qat/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/intrinsic/qat/modules/__pycache__/conv_fused.cpython-311.pyc,, +torch/nn/intrinsic/qat/modules/__pycache__/linear_fused.cpython-311.pyc,, +torch/nn/intrinsic/qat/modules/__pycache__/linear_relu.cpython-311.pyc,, +torch/nn/intrinsic/qat/modules/conv_fused.py,sha256=YKNHs_Fh3tfSoWmUAgd5APZwerDHpxPyY8KAXXDdMak,833 +torch/nn/intrinsic/qat/modules/linear_fused.py,sha256=qKnc8PyXqHPM--vfXmlD7JQoR5I480vCcYNZu-52-mM,434 +torch/nn/intrinsic/qat/modules/linear_relu.py,sha256=8740qDh1H-F03wWtKB6tQhPFSf2M2Cowyb91sfs02Ww,434 +torch/nn/intrinsic/quantized/__init__.py,sha256=Elo1b9Ro7sj3MsWn3xys-1rGWSe21SANc4H6Ai3x25c,336 +torch/nn/intrinsic/quantized/__pycache__/__init__.cpython-311.pyc,, +torch/nn/intrinsic/quantized/dynamic/__init__.py,sha256=Ov25o_5XouCKIqpzK42xcw8dAwVd02tGZqS4XL8BOIM,73 +torch/nn/intrinsic/quantized/dynamic/__pycache__/__init__.cpython-311.pyc,, +torch/nn/intrinsic/quantized/dynamic/modules/__init__.py,sha256=acHjDVSQZsjlTsbUoN4pI-971UJRfjavwMJRJuEOOBk,114 +torch/nn/intrinsic/quantized/dynamic/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/intrinsic/quantized/dynamic/modules/__pycache__/linear_relu.cpython-311.pyc,, +torch/nn/intrinsic/quantized/dynamic/modules/linear_relu.py,sha256=z509nJpUu_QTwDUjE_ON9Fbu7n5GfMnOzWju35bEXRk,97 +torch/nn/intrinsic/quantized/modules/__init__.py,sha256=6WHqSE0pA7qaW20FVPWggCar8uawjVKABKVeNdxjNlo,379 +torch/nn/intrinsic/quantized/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/intrinsic/quantized/modules/__pycache__/bn_relu.cpython-311.pyc,, +torch/nn/intrinsic/quantized/modules/__pycache__/conv_relu.cpython-311.pyc,, +torch/nn/intrinsic/quantized/modules/__pycache__/linear_relu.cpython-311.pyc,, +torch/nn/intrinsic/quantized/modules/bn_relu.py,sha256=3t9doT82fRh1SytKSivqQgxVfDP41KuvL-R9hz8oSb8,111 +torch/nn/intrinsic/quantized/modules/conv_relu.py,sha256=QWzacsMW7FtkasYvsZ9WqRExBjaMqsxHQHD5tYsZBFA,149 +torch/nn/intrinsic/quantized/modules/linear_relu.py,sha256=Zuj3I1owPTdoovG1-BoogRjn8JX1rMWM-WXYc-QFneU,89 +torch/nn/modules/__init__.py,sha256=Eef2WuOfFW36I0RRotha0XO8cU3NW9P6t7F-ibNeJpU,6542 +torch/nn/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/modules/__pycache__/_functions.cpython-311.pyc,, +torch/nn/modules/__pycache__/activation.cpython-311.pyc,, +torch/nn/modules/__pycache__/adaptive.cpython-311.pyc,, +torch/nn/modules/__pycache__/batchnorm.cpython-311.pyc,, +torch/nn/modules/__pycache__/channelshuffle.cpython-311.pyc,, +torch/nn/modules/__pycache__/container.cpython-311.pyc,, +torch/nn/modules/__pycache__/conv.cpython-311.pyc,, +torch/nn/modules/__pycache__/distance.cpython-311.pyc,, +torch/nn/modules/__pycache__/dropout.cpython-311.pyc,, +torch/nn/modules/__pycache__/flatten.cpython-311.pyc,, +torch/nn/modules/__pycache__/fold.cpython-311.pyc,, +torch/nn/modules/__pycache__/instancenorm.cpython-311.pyc,, +torch/nn/modules/__pycache__/lazy.cpython-311.pyc,, +torch/nn/modules/__pycache__/linear.cpython-311.pyc,, +torch/nn/modules/__pycache__/loss.cpython-311.pyc,, +torch/nn/modules/__pycache__/module.cpython-311.pyc,, +torch/nn/modules/__pycache__/normalization.cpython-311.pyc,, +torch/nn/modules/__pycache__/padding.cpython-311.pyc,, +torch/nn/modules/__pycache__/pixelshuffle.cpython-311.pyc,, +torch/nn/modules/__pycache__/pooling.cpython-311.pyc,, +torch/nn/modules/__pycache__/rnn.cpython-311.pyc,, +torch/nn/modules/__pycache__/sparse.cpython-311.pyc,, +torch/nn/modules/__pycache__/transformer.cpython-311.pyc,, +torch/nn/modules/__pycache__/upsampling.cpython-311.pyc,, +torch/nn/modules/__pycache__/utils.cpython-311.pyc,, +torch/nn/modules/_functions.py,sha256=C_QAxSaXW57WZxE8ItmvnxA5Nu_-f6e1ftpr4bHhv5U,12180 +torch/nn/modules/activation.py,sha256=icd06dDW_Aaj2MvwYDbzLg_H-9GLgENrbUYsOY1NGos,62129 +torch/nn/modules/adaptive.py,sha256=_fGwVPaYb2GHaz_Uv1klKHGnt7MTH1y4LwFD7Ia21I8,12606 +torch/nn/modules/batchnorm.py,sha256=leSFa4m_BatB3HHeYyaNPMYsfBG44iyNzyjU_yaE0ao,41176 +torch/nn/modules/channelshuffle.py,sha256=1K2PoerBn3WvtN3gfTyLAp6sjMVhADPjgI0eV5XD0WU,1682 +torch/nn/modules/container.py,sha256=BHi5SkO1BydsxNzYXPvf9_yDgH8wIm9XWqrGnUFrZYo,37394 +torch/nn/modules/conv.py,sha256=_NETMHEjVK4i9ukov3KZT7HI9u0wso1YvjxkrMmaRKQ,79715 +torch/nn/modules/distance.py,sha256=z7zUT5-kLio0kqt6r3EwNgDpmPDcoRC6QasilD_yTsk,3372 +torch/nn/modules/dropout.py,sha256=8avb3aCV6e4LuWcB6a_XDbyiApi9MZCqaLedDsmMR_c,11517 +torch/nn/modules/flatten.py,sha256=ogOVF9rao5nMuDXr2bLASB2QzQ4BIWhR3Z_t-byMpL8,5760 +torch/nn/modules/fold.py,sha256=o58ibuN2AmkR3o2Qg9ACBwx8iAFTlhaEuyGpucZxr14,13227 +torch/nn/modules/instancenorm.py,sha256=9JRMWmcS7n5VUJhNtuQCWFaClZx2vO85ie_Fs7aFKdI,21516 +torch/nn/modules/lazy.py,sha256=yAx_4T2XYmaBfaH9fKr9tAeFLA6-aUW2DJtYUvPausI,11674 +torch/nn/modules/linear.py,sha256=GWGtGh9lDhBLPH940Zry0hv5d-evbD_fJCCq1WPJ4gI,12376 +torch/nn/modules/loss.py,sha256=tDSg-3QlINXXMpPE8JYEYg043XmdOX87FwI5tId_W-Y,97637 +torch/nn/modules/module.py,sha256=VqH6OfMyI447rYf0MZA6h1llq--SVHr-5jVsE5DytxM,127992 +torch/nn/modules/normalization.py,sha256=lihLY6IIjH03hxbVEU61tqHActR5EWFw_P0oTFfoS08,15924 +torch/nn/modules/padding.py,sha256=Nyvsav0h0dJ1rtUt6i5c0jnPtfImRlId567eusJcHNk,31660 +torch/nn/modules/pixelshuffle.py,sha256=N4rqCH-2DvBKLG4ZpJLHUNsv8Ylx2uQHm7VxRop4CWw,3948 +torch/nn/modules/pooling.py,sha256=embOujctPnv03alIcbm3MLjQwr-fb1BRh0kXu13KyCY,61114 +torch/nn/modules/rnn.py,sha256=RbpIiDgRnNo2X6ME82ODNbWExMS22jU6Q9TzaNqPDmc,76242 +torch/nn/modules/sparse.py,sha256=SJe1QCw15JyDDtruHqzBscP4At3I5930ESpxx4vqgfU,24262 +torch/nn/modules/transformer.py,sha256=9s1lgkFAIrw-jA1HhsH4XrRoYL8CcU17rwUWMT04hGo,52209 +torch/nn/modules/upsampling.py,sha256=xqIhg4bfAfGsJ3rp8fMm-yeI8P_nKParbOdUGCkALlk,11611 +torch/nn/modules/utils.py,sha256=kS7sezb3cU9-B9swphiw4LaO0kzOlegUliZTHuFtKIk,2601 +torch/nn/parallel/__init__.py,sha256=z43J5FJ5qqe_aVSla6yG0qZ19t_Y9pX3g3BzYJvVuSs,760 +torch/nn/parallel/__pycache__/__init__.cpython-311.pyc,, +torch/nn/parallel/__pycache__/_functions.cpython-311.pyc,, +torch/nn/parallel/__pycache__/comm.cpython-311.pyc,, +torch/nn/parallel/__pycache__/data_parallel.cpython-311.pyc,, +torch/nn/parallel/__pycache__/distributed.cpython-311.pyc,, +torch/nn/parallel/__pycache__/parallel_apply.cpython-311.pyc,, +torch/nn/parallel/__pycache__/replicate.cpython-311.pyc,, +torch/nn/parallel/__pycache__/scatter_gather.cpython-311.pyc,, +torch/nn/parallel/_functions.py,sha256=HHi7gH1S1mwDXLC2ivD3w3e_s3AqO3qgHKelO2mzMOk,6226 +torch/nn/parallel/comm.py,sha256=ogPGR7BU6F8EWlHYpAVjVGdBJrr6uNsfTIVgiZdhAW8,11105 +torch/nn/parallel/data_parallel.py,sha256=dCC7VAilFp53b5FaJIiObG1SBdZkPwM9F3JIhNakKtM,11938 +torch/nn/parallel/distributed.py,sha256=JotDQkRWl_kIxF3O5R2ZCoauIk4R5Xrbw1JCfFQEelM,117758 +torch/nn/parallel/parallel_apply.py,sha256=tQQAIHez6tVvF3-Zgm17cYFi7MV_JGIZYlzf0rIEqWI,5115 +torch/nn/parallel/replicate.py,sha256=esgxFGR67Ni3bQCYcDlKMhbq9KXgYzDE5XTZrKLKRGY,7259 +torch/nn/parallel/scatter_gather.py,sha256=kPoOGvZe5zY2JBncFvOVqmHNYUZnCQUDXIopYmYKP0Y,5550 +torch/nn/parameter.py,sha256=WlPMxw0QjaXkCarcpAFMJrmC4OjQxT4VW85xz3PYysE,12322 +torch/nn/parameter.pyi,sha256=nWmOqB7CJFdB0wabxODC0yZI5j0N20MhsHb_twxD3Ac,1118 +torch/nn/qat/__init__.py,sha256=YHWepHV69h8eIIZ09WgM96wRwYTKAdZ5IXJG4ajhsHI,365 +torch/nn/qat/__pycache__/__init__.cpython-311.pyc,, +torch/nn/qat/dynamic/__init__.py,sha256=3v0msXCbfI7q_EYZ2WUPfuW0IxLRr2s8ICycY6XraDQ,208 +torch/nn/qat/dynamic/__pycache__/__init__.cpython-311.pyc,, +torch/nn/qat/dynamic/modules/__init__.py,sha256=xnBpUudYrmz9leiWC2UaMEGbDX2eO--nNluoaO_V-LE,78 +torch/nn/qat/dynamic/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/qat/dynamic/modules/__pycache__/linear.cpython-311.pyc,, +torch/nn/qat/dynamic/modules/linear.py,sha256=io9A-B4H0RMvMRYuKeVSxRQPsBhdCy5WKfvr5QtKHuo,416 +torch/nn/qat/modules/__init__.py,sha256=xHBwOWC71Dzb9MLPvLkMMdcDlr39rf-dejMilOtshHM,501 +torch/nn/qat/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/qat/modules/__pycache__/conv.cpython-311.pyc,, +torch/nn/qat/modules/__pycache__/embedding_ops.cpython-311.pyc,, +torch/nn/qat/modules/__pycache__/linear.cpython-311.pyc,, +torch/nn/qat/modules/conv.py,sha256=hOAGSJ7oaQTf_9DWXPWrGGsQ9oPLASiCZaeHX4X3BMI,406 +torch/nn/qat/modules/embedding_ops.py,sha256=1BP3HyU1Xkr8yQBH1RNqv2rwZM9Uv6Wl6X5-DM_Jlls,458 +torch/nn/qat/modules/linear.py,sha256=cFw8lAcvlWfV__h6uSDi9w4I7snqlWGVAyfY7nnHiMA,392 +torch/nn/quantizable/__init__.py,sha256=O3C0GktZB_Ut6pNTL17JbCzJ2h2hCIK3KH8sVnhSuDs,57 +torch/nn/quantizable/__pycache__/__init__.cpython-311.pyc,, +torch/nn/quantizable/modules/__init__.py,sha256=gr03bjx4iwgg1n8zCu3B1UpgNqhb6UxxSVvwO5894V4,207 +torch/nn/quantizable/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/quantizable/modules/__pycache__/activation.cpython-311.pyc,, +torch/nn/quantizable/modules/__pycache__/rnn.cpython-311.pyc,, +torch/nn/quantizable/modules/activation.py,sha256=m8gKPli1OoBlwAng4E2ndP9K3ov2t6D7pN1Zj84nv0g,440 +torch/nn/quantizable/modules/rnn.py,sha256=ymfl_ayTkuUunsNt9186sFOhtapsxsTHVtY_QlKFH9A,429 +torch/nn/quantized/__init__.py,sha256=OTMoZbhpm8nI1oMVMuB8tt4JvpDN3MsIoYww2nUJfQM,771 +torch/nn/quantized/__pycache__/__init__.cpython-311.pyc,, +torch/nn/quantized/__pycache__/functional.cpython-311.pyc,, +torch/nn/quantized/_reference/__init__.py,sha256=VrWZh6FzvyvAi744iiFBcYarBX3Akn6kTarYzEIX0_w,66 +torch/nn/quantized/_reference/__pycache__/__init__.cpython-311.pyc,, +torch/nn/quantized/_reference/modules/__init__.py,sha256=yjgG2MRfrWs-u8l1CCILBoPTa-FBORFZh9i86vYKdy8,1018 +torch/nn/quantized/_reference/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/quantized/_reference/modules/__pycache__/conv.cpython-311.pyc,, +torch/nn/quantized/_reference/modules/__pycache__/linear.cpython-311.pyc,, +torch/nn/quantized/_reference/modules/__pycache__/rnn.cpython-311.pyc,, +torch/nn/quantized/_reference/modules/__pycache__/sparse.cpython-311.pyc,, +torch/nn/quantized/_reference/modules/__pycache__/utils.cpython-311.pyc,, +torch/nn/quantized/_reference/modules/conv.py,sha256=FclYIsMCbXH4erKiBDintHNWDsplOru0QY9LDToSL-s,579 +torch/nn/quantized/_reference/modules/linear.py,sha256=Cu-oS3rGRNtFQX8C8yBiLIeORYswqNnFbVw3ds6PMKQ,450 +torch/nn/quantized/_reference/modules/rnn.py,sha256=A6kV2QXCij5mpFFUTWBQwf6ynKQNJj0VMJbvHqkXa-Y,524 +torch/nn/quantized/_reference/modules/sparse.py,sha256=gSNXDYXlx1IoSrRN0i_pJHJt1znBH3sAPvDpivdfask,467 +torch/nn/quantized/_reference/modules/utils.py,sha256=nBSMp0qUV2OBmuCJYKGmpY6JtRinMoHhwPT2B8VbY08,590 +torch/nn/quantized/dynamic/__init__.py,sha256=lN8JfmyIwtMWNsOD0hWGgKMVPUXuUb5cu6QOvIlEJxg,58 +torch/nn/quantized/dynamic/__pycache__/__init__.cpython-311.pyc,, +torch/nn/quantized/dynamic/modules/__init__.py,sha256=ORZYUZYDqtgat8r4CkTj2GRbnR4y2CxoK3hNUT1AFEU,993 +torch/nn/quantized/dynamic/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/quantized/dynamic/modules/__pycache__/conv.cpython-311.pyc,, +torch/nn/quantized/dynamic/modules/__pycache__/linear.cpython-311.pyc,, +torch/nn/quantized/dynamic/modules/__pycache__/rnn.cpython-311.pyc,, +torch/nn/quantized/dynamic/modules/conv.py,sha256=bq0IjYdFIpdJbsQkIeF1LDHZOC8wvGInVohV5KwqqI0,669 +torch/nn/quantized/dynamic/modules/linear.py,sha256=67W7Cw26T5oIiE4n3RJykC6PX5VypOJU3lQcNP6mohg,448 +torch/nn/quantized/dynamic/modules/rnn.py,sha256=vWAb3OQKyHWubLC5U-qIGP-GgYq-GFpvltZCDXjogrA,740 +torch/nn/quantized/functional.py,sha256=YLagkINzH1ksH0zcEEizTr0Wdu8gvziqPkx-elbUkJE,276 +torch/nn/quantized/modules/__init__.py,sha256=oNogz5eH2YEYfHjRkYBEqp_XD2GduDqc9StknEGWx40,2101 +torch/nn/quantized/modules/__pycache__/__init__.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/activation.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/batchnorm.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/conv.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/dropout.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/embedding_ops.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/functional_modules.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/linear.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/normalization.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/rnn.cpython-311.pyc,, +torch/nn/quantized/modules/__pycache__/utils.cpython-311.pyc,, +torch/nn/quantized/modules/activation.py,sha256=52tIo-6xSFg2QIuNL-Q7AKs2IiDBCcGFIvWNupA3fjA,528 +torch/nn/quantized/modules/batchnorm.py,sha256=OpxwQosX9Pv53csZ4ZXKhFOkya_tXMbtndDMSETJKV8,437 +torch/nn/quantized/modules/conv.py,sha256=3JBXb4gZOKXHIyWnnPEVxlb-WRDSMiNEUvuofCSNsSw,666 +torch/nn/quantized/modules/dropout.py,sha256=qdWqDycnaodhJDr_EyIQCiaSg4xK5dqEZ21IBIsC_8M,442 +torch/nn/quantized/modules/embedding_ops.py,sha256=smsmXs2h7t2y6WbToNuJ060heMQbZOXOfkNbdwOqVO4,547 +torch/nn/quantized/modules/functional_modules.py,sha256=bk60spaa4_sUCDXyh4GOh937FVIwNCH-l93lAPi39-Q,554 +torch/nn/quantized/modules/linear.py,sha256=zsZr5Rb4z4XJvTjc-hF9cqG4uqpsbkHOc3jJOxcSM4A,481 +torch/nn/quantized/modules/normalization.py,sha256=vUyWsXC7CFIm0oohnco73FD3Kz9xMGHeJKq3dNKZxKQ,626 +torch/nn/quantized/modules/rnn.py,sha256=NWI5IuH_qaC1YD1ROeoUbodbN3yS0MrZKOhXLmGfPDo,411 +torch/nn/quantized/modules/utils.py,sha256=ZkyzQS1nmLaK5Nvq2eKtSha7mHG3HyM08_5ueUc9FpQ,539 +torch/nn/utils/__init__.py,sha256=IDrUDb4E1qZkMUEfsT0eN0zC9YJyHjbJc0oxtgosE0c,1274 +torch/nn/utils/__pycache__/__init__.cpython-311.pyc,, +torch/nn/utils/__pycache__/_deprecation_utils.cpython-311.pyc,, +torch/nn/utils/__pycache__/_named_member_accessor.cpython-311.pyc,, +torch/nn/utils/__pycache__/_per_sample_grad.cpython-311.pyc,, +torch/nn/utils/__pycache__/clip_grad.cpython-311.pyc,, +torch/nn/utils/__pycache__/convert_parameters.cpython-311.pyc,, +torch/nn/utils/__pycache__/fusion.cpython-311.pyc,, +torch/nn/utils/__pycache__/init.cpython-311.pyc,, +torch/nn/utils/__pycache__/memory_format.cpython-311.pyc,, +torch/nn/utils/__pycache__/parametrizations.cpython-311.pyc,, +torch/nn/utils/__pycache__/parametrize.cpython-311.pyc,, +torch/nn/utils/__pycache__/prune.cpython-311.pyc,, +torch/nn/utils/__pycache__/rnn.cpython-311.pyc,, +torch/nn/utils/__pycache__/spectral_norm.cpython-311.pyc,, +torch/nn/utils/__pycache__/stateless.cpython-311.pyc,, +torch/nn/utils/__pycache__/weight_norm.cpython-311.pyc,, +torch/nn/utils/_deprecation_utils.py,sha256=KB2l-yIWZAA9OI0lT0f2s5WXeFvwvydY4EP0GGDDags,1698 +torch/nn/utils/_expanded_weights/__init__.py,sha256=PHIe9H3l0iuImQ6z3WqMkWIS37FDQVQ16vMo0a08mtA,452 +torch/nn/utils/_expanded_weights/__pycache__/__init__.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/__pycache__/conv_expanded_weights.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/__pycache__/conv_utils.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/__pycache__/embedding_expanded_weights.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/__pycache__/expanded_weights_impl.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/__pycache__/expanded_weights_utils.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/__pycache__/group_norm_expanded_weights.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/__pycache__/instance_norm_expanded_weights.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/__pycache__/layer_norm_expanded_weights.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/__pycache__/linear_expanded_weights.cpython-311.pyc,, +torch/nn/utils/_expanded_weights/conv_expanded_weights.py,sha256=5Kr60dZ4gzcsX-uPJYWw6_9-Afoqu5QgVAmUt-smiDA,2925 +torch/nn/utils/_expanded_weights/conv_utils.py,sha256=iuAChFyYOsVpVcQtlv2Ux2JKuANQfvFdgidAIlgGjeE,10815 +torch/nn/utils/_expanded_weights/embedding_expanded_weights.py,sha256=kVjENz7zj9_3fIiF9rNgNbUx7Z78UNSEJRhO9dYjryc,3073 +torch/nn/utils/_expanded_weights/expanded_weights_impl.py,sha256=z7QPT00z7IxUOpO3Fad3ot-2dPK4YI93D4WHM4aHO_8,6408 +torch/nn/utils/_expanded_weights/expanded_weights_utils.py,sha256=l0Kjs3qjXsjwYxVf86g23uI_UufloirP79V4tBp3ALo,7612 +torch/nn/utils/_expanded_weights/group_norm_expanded_weights.py,sha256=Uw86QFNavKYlaNSnbSbBH0ieyQSh3jyChz9fwbKxoO4,3514 +torch/nn/utils/_expanded_weights/instance_norm_expanded_weights.py,sha256=TEs5rHPpUAet5KUtF1_RgyKr3kbZRKY6kEqsqP5mk5M,3772 +torch/nn/utils/_expanded_weights/layer_norm_expanded_weights.py,sha256=X1r7CwjyK-j03S9J6axQPGJ7DQr-yg1mWUNQ3Rip9UE,3289 +torch/nn/utils/_expanded_weights/linear_expanded_weights.py,sha256=6BI2LGjLpH29uFg5ktMxIW7AmxrA83uBPs62YmZOqww,2259 +torch/nn/utils/_named_member_accessor.py,sha256=m5zwqxzn6O2vwfz9lFZhYwewQLQZHVYPJvTlqbM5WgA,14470 +torch/nn/utils/_per_sample_grad.py,sha256=r2XSvB-eGfw7g-TOhHWT_v_w2K4oI3qaK7pHxm9S_iw,5775 +torch/nn/utils/clip_grad.py,sha256=XjDawX6wMRqxoS5G7SSE3stmFlQnZE0mAdE1cKsO5Nc,11465 +torch/nn/utils/convert_parameters.py,sha256=xiTn9HykDVWMwcR_Kx4zTJrTWxtLlfPv3oWYxeSUku4,3213 +torch/nn/utils/fusion.py,sha256=qbwJpB-sZ4tIAvBcfWDb61Hpla6ERwUAHb5Dx8fxEo4,6753 +torch/nn/utils/init.py,sha256=eA5Y0Ti3Qjhd-pCKEx9MpLHjJuz4bA3XNw6W0lC55XE,2250 +torch/nn/utils/memory_format.py,sha256=fYxW9eaJswQSSDeH1Yn2TVC0C_ExJMKWumdhgl5jiEc,8206 +torch/nn/utils/parametrizations.py,sha256=DSh9GFR0WmH0HdvPoeL-CtVAR0FHb6VyLy3IAR3BNTk,26044 +torch/nn/utils/parametrize.py,sha256=VqB5QsjuESWbnYilV--gF-6k96h_UuPjUa7hr03BdGk,39211 +torch/nn/utils/prune.py,sha256=s1NmGgt7i3NHyrCWozahELKk6nJ35vnhZuQRwMcEP5A,58753 +torch/nn/utils/rnn.py,sha256=dwyonizGne8qGfUuzqkbWb2uLeG0FZs-qDaSdu-TT_4,23131 +torch/nn/utils/spectral_norm.py,sha256=PxYIagdAIGLXL9QO-YRhYgamRzS1MjALbZH_QlJFFIA,14708 +torch/nn/utils/stateless.py,sha256=3rQJAtpopS35ovIp8bCuoIBuz-xWr_0CFigLDwqlw7I,11568 +torch/nn/utils/weight_norm.py,sha256=PZ7TbdZ2zo76SmvY2zAvzUsJ8Dt8YiXVfgtMDByEdgE,5883 +torch/numa/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/numa/__pycache__/__init__.cpython-311.pyc,, +torch/numa/__pycache__/binding.cpython-311.pyc,, +torch/numa/binding.py,sha256=4wnOn3P9D_geFINY2VrJcQjWmfb_CozJaZOtjWlWZKU,24799 +torch/onnx/__init__.py,sha256=ishI0e9uBbh9r6C34WieMCbUaGyi3f-U2kdvX-Vk0qU,15003 +torch/onnx/__pycache__/__init__.cpython-311.pyc,, +torch/onnx/__pycache__/_constants.cpython-311.pyc,, +torch/onnx/__pycache__/_flags.cpython-311.pyc,, +torch/onnx/__pycache__/errors.cpython-311.pyc,, +torch/onnx/__pycache__/operators.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_helper.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset10.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset11.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset12.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset13.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset14.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset15.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset16.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset17.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset18.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset19.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset20.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset7.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset8.cpython-311.pyc,, +torch/onnx/__pycache__/symbolic_opset9.cpython-311.pyc,, +torch/onnx/__pycache__/testing.cpython-311.pyc,, +torch/onnx/__pycache__/utils.cpython-311.pyc,, +torch/onnx/__pycache__/verification.cpython-311.pyc,, +torch/onnx/_constants.py,sha256=KsVyVjgGFLE_G3mr8-ilc3YY6_0zxbKHZsw1OVbQOUo,531 +torch/onnx/_flags.py,sha256=7gg9rPep-Oy_cuneXABiEItasVTVvsp5PJ7cfc_6KLA,1554 +torch/onnx/_internal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/onnx/_internal/__pycache__/__init__.cpython-311.pyc,, +torch/onnx/_internal/__pycache__/_lazy_import.cpython-311.pyc,, +torch/onnx/_internal/_lazy_import.py,sha256=RWu_Nk7U6sIybFCMEMuPG4YrJiZ0LX1tP56a7--61xo,1213 +torch/onnx/_internal/exporter/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/onnx/_internal/exporter/__pycache__/__init__.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_analysis.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_building.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_capture_strategies.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_compat.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_constants.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_core.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_decomp.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_dispatching.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_dynamic_shapes.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_errors.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_exportable_module.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_flags.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_fx_passes.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_input_observer.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_ir_passes.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_isolated.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_onnx_program.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_registration.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_reporting.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_schemas.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_tensors.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_testing.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_type_casting.cpython-311.pyc,, +torch/onnx/_internal/exporter/__pycache__/_verification.cpython-311.pyc,, +torch/onnx/_internal/exporter/_analysis.py,sha256=ZhMKGTs8voROxN0x8azRngSVnqibbRvYyKQkRMib1bY,9969 +torch/onnx/_internal/exporter/_building.py,sha256=bwDvC5ULGGfnpxXy_jdVOwfzqXRRQtppaaxOJS6svZc,29538 +torch/onnx/_internal/exporter/_capture_strategies.py,sha256=upT9bFJMij1tzPOfmM2B0hG0DIWvAJkEe_9UBxpW40k,10542 +torch/onnx/_internal/exporter/_compat.py,sha256=4HBcn7Ifmd7OdvngRS3dMTSyYAsfrHnE0Dru5_D8FQA,7552 +torch/onnx/_internal/exporter/_constants.py,sha256=y8ST-ZGRhbGrA_d-oM2xapBMa0hnFcEgLbLxt8pcMn4,378 +torch/onnx/_internal/exporter/_core.py,sha256=FwrMBhs7rH7imeYkQfu18VnB0KBoJPWM30o75RH1bw4,70360 +torch/onnx/_internal/exporter/_decomp.py,sha256=mVFktC8pOoFyQnl6UvhpDbHlNZtSvB8Hg5jNcui6eFQ,2938 +torch/onnx/_internal/exporter/_dispatching.py,sha256=_4rGW0B6NRWiPEAGg4TgxvzSgW2BlVkqGizNPdynuI8,2070 +torch/onnx/_internal/exporter/_dynamic_shapes.py,sha256=djj0TF_hAucwvAXd1ZSCeJLFu8gFb567Z7viGQj6ADA,14431 +torch/onnx/_internal/exporter/_errors.py,sha256=VziOIvyap9A_UDRP41HkC_s0J-R13dbCrVLIusF-6A8,535 +torch/onnx/_internal/exporter/_exportable_module.py,sha256=ymWXCMLYKiMruQAF3duzJRArqpp4yH19zTUYKaWEOvc,6807 +torch/onnx/_internal/exporter/_flags.py,sha256=0vyPbWUoBb_BUpGJpR5JRKZJxlXfk312XsdrbZsD9JM,794 +torch/onnx/_internal/exporter/_fx_passes.py,sha256=RJDceCGcbBr_lGpVbQGuSReBRbOTkt7uatRNQtNr9w0,1919 +torch/onnx/_internal/exporter/_input_observer.py,sha256=8Lko6FzmvjaFPUapBos7qLqAkSGV1T749rhz8LH3ij8,51666 +torch/onnx/_internal/exporter/_ir_passes.py,sha256=9XFQqrx0mg9EQfzyMWaCS88zD8lBX0tTgKZjsGiMl3Q,6289 +torch/onnx/_internal/exporter/_isolated.py,sha256=4LAKM1QeQ1bn6WePE3NyAS2o8RhE5lLy-lk8ArV-8uU,1869 +torch/onnx/_internal/exporter/_onnx_program.py,sha256=iynL3xhJ9JXU5ZJr91GKSPpLGec6IXqHEAKbfgSMLEk,19606 +torch/onnx/_internal/exporter/_registration.py,sha256=jG9jxd8gPPFGnCJX_d-aqX5VWHWCubYhMsvTA-D4q-w,12422 +torch/onnx/_internal/exporter/_reporting.py,sha256=gTtMzru1mYtL48DRk-8vGmyrCDQm3PxrBsBwKWU1O5s,7411 +torch/onnx/_internal/exporter/_schemas.py,sha256=mGvy0H7IgrHRJ4dOZv_yywXpAKJydeFJefteEc8hIug,12280 +torch/onnx/_internal/exporter/_tensors.py,sha256=qxJ6FpCQZgA0KOMr_7CvrxQD-88b52prQk9QQ6GeYyE,2567 +torch/onnx/_internal/exporter/_testing.py,sha256=vn4KD6-Cd7pD7NUVEWJv8M_dvCeCWobEucKpjvZQPVg,3993 +torch/onnx/_internal/exporter/_torchlib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/onnx/_internal/exporter/_torchlib/__pycache__/__init__.cpython-311.pyc,, +torch/onnx/_internal/exporter/_torchlib/__pycache__/_tensor_typing.cpython-311.pyc,, +torch/onnx/_internal/exporter/_torchlib/__pycache__/_torchlib_registry.cpython-311.pyc,, +torch/onnx/_internal/exporter/_torchlib/_tensor_typing.py,sha256=DfaNR_WzfujCOX0PuD5QeE9hTI86ihF9g-GeRA7k0NQ,1953 +torch/onnx/_internal/exporter/_torchlib/_torchlib_registry.py,sha256=BG8oUTNKK-FVAo_s2Re8g15oJo2m1QVcSF-MiqlBmig,2886 +torch/onnx/_internal/exporter/_torchlib/ops/__init__.py,sha256=bgyPtdiM9s1nE_1ZdR_6-0y0Tea9o0Qiyts2wWU7EM4,180 +torch/onnx/_internal/exporter/_torchlib/ops/__pycache__/__init__.cpython-311.pyc,, +torch/onnx/_internal/exporter/_torchlib/ops/__pycache__/core.cpython-311.pyc,, +torch/onnx/_internal/exporter/_torchlib/ops/__pycache__/hop.cpython-311.pyc,, +torch/onnx/_internal/exporter/_torchlib/ops/__pycache__/nn.cpython-311.pyc,, +torch/onnx/_internal/exporter/_torchlib/ops/__pycache__/symbolic.cpython-311.pyc,, +torch/onnx/_internal/exporter/_torchlib/ops/__pycache__/symops.cpython-311.pyc,, +torch/onnx/_internal/exporter/_torchlib/ops/core.py,sha256=NMSMQLta-VKSWasAcaUQ7h03h62B0WO3cZC2OsN7U0k,1568 +torch/onnx/_internal/exporter/_torchlib/ops/hop.py,sha256=umnwD9LvYT4YXcapTpDzW6guvVq2e3mwOWMROVdK7nM,12837 +torch/onnx/_internal/exporter/_torchlib/ops/nn.py,sha256=lX90Y-wHiWqpVPaVWS5Y1KSZuG9FdWyKuZtIgrS6rqs,14044 +torch/onnx/_internal/exporter/_torchlib/ops/symbolic.py,sha256=ZR470gARvhh1UMMzL3oH7PhaPp9qp7Em2UaEB3wTNgk,4773 +torch/onnx/_internal/exporter/_torchlib/ops/symops.py,sha256=T82DeVN_Y1fmqrlgHZaKw2ZXk6gD5UCXGFh6YGepvfA,1796 +torch/onnx/_internal/exporter/_type_casting.py,sha256=ZfJSaAagtIWGV_0BTDOnjOtnrZRwoS1f2rKuedkQMbw,1313 +torch/onnx/_internal/exporter/_verification.py,sha256=Vo3MCIcV0ALzSICc_PG5buKbG10HvKUdef1vwjX6HHc,12535 +torch/onnx/_internal/fx/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/onnx/_internal/fx/__pycache__/__init__.cpython-311.pyc,, +torch/onnx/_internal/fx/__pycache__/_pass.cpython-311.pyc,, +torch/onnx/_internal/fx/__pycache__/type_utils.cpython-311.pyc,, +torch/onnx/_internal/fx/_pass.py,sha256=ijrOE4Clp5xEKvTid1WJPDwa2QWqkC5FKMu_FOrZrj0,8638 +torch/onnx/_internal/fx/passes/__init__.py,sha256=mImXwcuJd6IfP3-3RkuLStx9i0ykJK8dhozuZNwRdpA,91 +torch/onnx/_internal/fx/passes/__pycache__/__init__.cpython-311.pyc,, +torch/onnx/_internal/fx/passes/__pycache__/type_promotion.cpython-311.pyc,, +torch/onnx/_internal/fx/passes/type_promotion.py,sha256=PGWMKXdLb-bCd66yXQZ3z-T7QiGhoyAP_TpdKbIJp44,65278 +torch/onnx/_internal/fx/type_utils.py,sha256=XblXgNj16vBEpGEd5Obl6D1XBuAnSeAFjfeFp0VtRbk,945 +torch/onnx/_internal/torchscript_exporter/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/onnx/_internal/torchscript_exporter/__pycache__/__init__.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/_experimental.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/_globals.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/_type_utils.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/jit_utils.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/onnx_proto_utils.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/registration.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_helper.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset10.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset11.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset12.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset13.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset14.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset15.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset16.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset17.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset18.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset19.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset20.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset7.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset8.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/symbolic_opset9.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/utils.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/__pycache__/verification.cpython-311.pyc,, +torch/onnx/_internal/torchscript_exporter/_experimental.py,sha256=sNZuVCJZztlSpM8VxZmdXgd8bOcsQo8C_EqF7quG2Cs,968 +torch/onnx/_internal/torchscript_exporter/_globals.py,sha256=dEltGsC4MJTeVUCaAYm3-L6xpjaNPZkGBNwG7jfM47Q,2801 +torch/onnx/_internal/torchscript_exporter/_type_utils.py,sha256=-2e69dp1sBkq9UHRyjbnZaG-cLnVBvffE4tG1AZtOdM,14021 +torch/onnx/_internal/torchscript_exporter/jit_utils.py,sha256=LSmDihegunKT5OUXa920hLxCboOFC_EXZnHD_K2g7KQ,14170 +torch/onnx/_internal/torchscript_exporter/onnx_proto_utils.py,sha256=Iz_QXJoN8MUWTWswjK2BGinni2tzgTXkr3md3lZD6NA,9358 +torch/onnx/_internal/torchscript_exporter/registration.py,sha256=oHZfHRVQEB9bKa3tFa8sQ0Bt9GrjWha49GwTdtS-6OU,11135 +torch/onnx/_internal/torchscript_exporter/symbolic_helper.py,sha256=KFWSRxrt5eFdmAGEdpdoo3SWN3rr7BgWKhaZSLwb3NI,86588 +torch/onnx/_internal/torchscript_exporter/symbolic_opset10.py,sha256=wwyoZDc4LQ5xv8jJYD83kQ6HXHlQSo0l4WBs1EMv87s,37593 +torch/onnx/_internal/torchscript_exporter/symbolic_opset11.py,sha256=RSWXkus5f1uURfj7UXh0LIV6seZTLPhgiesWBUEzz8U,53594 +torch/onnx/_internal/torchscript_exporter/symbolic_opset12.py,sha256=5W_QJD6V7MeCtiRLcN8QBmZT24JKWEuB32wJAOojoxM,15807 +torch/onnx/_internal/torchscript_exporter/symbolic_opset13.py,sha256=4PXy5wy51WWspaoyPfwMG3GEDdDfm7Z5Cw9q6KO2mJE,41510 +torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py,sha256=AUeXVfoV8Lo2SeYdPR-62vyKTkwk6goUZV0vSw0KlfQ,10426 +torch/onnx/_internal/torchscript_exporter/symbolic_opset15.py,sha256=yUBs1fDhvIx7-aWghD1ZtBbYwlVODGJrb_Tk1HQPEFQ,2892 +torch/onnx/_internal/torchscript_exporter/symbolic_opset16.py,sha256=QRJIIAf0VwR4NseYTKn5K3TGNJ8sBNYsgdbQBfxG0oY,6456 +torch/onnx/_internal/torchscript_exporter/symbolic_opset17.py,sha256=RPcutTe7lEEEVHLTZXgg9334wTkGJvcif4FpNbD1Ua4,8356 +torch/onnx/_internal/torchscript_exporter/symbolic_opset18.py,sha256=nAbSdAcm2QnUj0yYK0nUs6RNjiSQuP_R85I79YDzlOQ,8247 +torch/onnx/_internal/torchscript_exporter/symbolic_opset19.py,sha256=mdWDL5ju6R87V0WlsPnwFidl5EwyPeOb4roJkXLW1fI,536 +torch/onnx/_internal/torchscript_exporter/symbolic_opset20.py,sha256=LUCfXezm0zhOyDiSHhuYJac_y9DoIgTtAzgIJ-UtkOw,2462 +torch/onnx/_internal/torchscript_exporter/symbolic_opset7.py,sha256=d71ZOvg4TYXfi9AP3F3YxtUrqrTXxkcXV6Qj0aSQAI0,2192 +torch/onnx/_internal/torchscript_exporter/symbolic_opset8.py,sha256=jR_c6xodWrKtmz8CaV_1zS4Ywhr-xG4DKBOzRajmETY,15081 +torch/onnx/_internal/torchscript_exporter/symbolic_opset9.py,sha256=cYQqe2Ecffim_oj_hZJey7Edj7OG5mQxj0WjQDwmrRY,225996 +torch/onnx/_internal/torchscript_exporter/utils.py,sha256=xoufcmFNox5eLjo8HFOBR2LWAWYHEbFOCIQTbWgRWtE,76841 +torch/onnx/_internal/torchscript_exporter/verification.py,sha256=ysbe_khOrCWnqkuWhuLA-mGwswtfZp2JYm7tSPuHbxQ,19156 +torch/onnx/errors.py,sha256=7Svwc5kOBcoGNz4d7rZ0Cdh217IkmlDADQZOeFL9flg,3478 +torch/onnx/operators.py,sha256=fS94DP2lYBZwI8qpxN6VWHWuTrvD7VWueC8eonqRbnE,1185 +torch/onnx/ops/__init__.py,sha256=OLQtAPGNTE3dZM_4u3KOTedP6OodYoP_2kzIYdliJuE,20939 +torch/onnx/ops/__pycache__/__init__.cpython-311.pyc,, +torch/onnx/ops/__pycache__/_dtype_mappings.cpython-311.pyc,, +torch/onnx/ops/__pycache__/_impl.cpython-311.pyc,, +torch/onnx/ops/__pycache__/_symbolic_impl.cpython-311.pyc,, +torch/onnx/ops/_dtype_mappings.py,sha256=T2_Rwd7HR5CzErvrphymlsCk0kRZhauZdhZbi0JZb3Y,836 +torch/onnx/ops/_impl.py,sha256=vndhmlw6dYLBIf8Te79er5ifyFVWCMq89f-pkeunjE8,19910 +torch/onnx/ops/_symbolic_impl.py,sha256=_zZyJ1o5mI9mAij49CE6ElyoLQStoG4va-WUeqaLaso,11503 +torch/onnx/symbolic_helper.py,sha256=LIHrpW5b175KG4GTm6gWc9WISbYYmgkKEd-hXXRGyGw,222 +torch/onnx/symbolic_opset10.py,sha256=SYoifeNPe7fnfUQ1o7oOfsfJTBkRjsgI-spAJ4zqZfU,325 +torch/onnx/symbolic_opset11.py,sha256=HER3r7zS8st1M2PCCZdyUkDp4ImPaAp4h5-yNjg8arM,224 +torch/onnx/symbolic_opset12.py,sha256=xaQpIzjkwHUiPIsJ3PTDfh-9_DH8ko_Ryve2_oZv4JI,224 +torch/onnx/symbolic_opset13.py,sha256=WqXlv_pfF46qr5GdBgwGQjPYF9kJtuk9EXMyQeQppfM,224 +torch/onnx/symbolic_opset14.py,sha256=nTk5SFQne8C41KTD8ZNyBF5vQfF_YZvg2qXmpA75QYA,224 +torch/onnx/symbolic_opset15.py,sha256=UnyYwoNDZ4wsiSjLzTQug9TabHmEEkXXzt53IPH-aB0,224 +torch/onnx/symbolic_opset16.py,sha256=lo4OY1dQwuJfL52MUh7zuQDdVhyxBJOnsJTTQDTWsnQ,224 +torch/onnx/symbolic_opset17.py,sha256=IdVLH_n6JA3TttyJ6NI1HFYpDdwTChRB5ZjSd7szl7w,224 +torch/onnx/symbolic_opset18.py,sha256=0yKeke5KAewf1P536eOWLondGHKk-kPp1etMqtCyyek,224 +torch/onnx/symbolic_opset19.py,sha256=sUngszvMGJzrGCyrzTuZgSA-aOY4G9-niC9JuwCsVuQ,224 +torch/onnx/symbolic_opset20.py,sha256=rz2fMG4kDcGK1j0qTEb77TwYXBhFVtS_JuIVjp0KMpk,224 +torch/onnx/symbolic_opset7.py,sha256=M27NKKqPnRtfFaW-M5MKFxQ7n0chpvTLN66xwuSyuOA,222 +torch/onnx/symbolic_opset8.py,sha256=8yMePMZ6epwV-RxEKja0OHDGddthTfL0zc7eBoOb-ys,222 +torch/onnx/symbolic_opset9.py,sha256=jAXYpzkr9EBN7tpGVw3JCJhuZN2yWmHVmty_3HgdEeQ,391 +torch/onnx/testing.py,sha256=QbAECaWoSos_a1XwM_KBh63qvQVajvlnMDHr0qUnNO4,219 +torch/onnx/utils.py,sha256=UASN7loGVDomvR9PayxZDTgvq1SRyQC2C7-gcBlSFRI,203 +torch/onnx/verification.py,sha256=L6GdYIC4gPAgZu4uXJmib7vG7u_1bkMMnXmee70zveU,343 +torch/optim/__init__.py,sha256=IJRMuKRPeA_LG1xGkftt7lCJtIUcZumTdvnYTHckfRA,2205 +torch/optim/__pycache__/__init__.cpython-311.pyc,, +torch/optim/__pycache__/_adafactor.cpython-311.pyc,, +torch/optim/__pycache__/_functional.cpython-311.pyc,, +torch/optim/__pycache__/_muon.cpython-311.pyc,, +torch/optim/__pycache__/adadelta.cpython-311.pyc,, +torch/optim/__pycache__/adagrad.cpython-311.pyc,, +torch/optim/__pycache__/adam.cpython-311.pyc,, +torch/optim/__pycache__/adamax.cpython-311.pyc,, +torch/optim/__pycache__/adamw.cpython-311.pyc,, +torch/optim/__pycache__/asgd.cpython-311.pyc,, +torch/optim/__pycache__/lbfgs.cpython-311.pyc,, +torch/optim/__pycache__/lr_scheduler.cpython-311.pyc,, +torch/optim/__pycache__/nadam.cpython-311.pyc,, +torch/optim/__pycache__/optimizer.cpython-311.pyc,, +torch/optim/__pycache__/radam.cpython-311.pyc,, +torch/optim/__pycache__/rmsprop.cpython-311.pyc,, +torch/optim/__pycache__/rprop.cpython-311.pyc,, +torch/optim/__pycache__/sgd.cpython-311.pyc,, +torch/optim/__pycache__/sparse_adam.cpython-311.pyc,, +torch/optim/__pycache__/swa_utils.cpython-311.pyc,, +torch/optim/_adafactor.py,sha256=tvzdaK9HKn6-09J06ElNNfAvvBIGNFkdKPoS4hJji3s,28512 +torch/optim/_functional.py,sha256=kYDxfTcm5jVdomzxYI38DtdvmbyvtHJT3jG9b1cuxLc,3250 +torch/optim/_multi_tensor/__init__.py,sha256=DaZ72GF7H5n46jPJ_mAiz67xOsutFFE16KJsB05ylns,1027 +torch/optim/_multi_tensor/__init__.pyi,sha256=Q4IMQ5dfkrDDq7ooUZubB4Okb7KHVhWpjx5HUF3f7hY,537 +torch/optim/_multi_tensor/__pycache__/__init__.cpython-311.pyc,, +torch/optim/_muon.py,sha256=TQA6ui0Mf8wkh1hFgC5F7bRHWlGgN-LKqntGw5RMDa4,14444 +torch/optim/adadelta.py,sha256=G7oYwi1Wo3fGNsri2fg4QOQxjxmRMfgkBHEuBZ3WYuM,16851 +torch/optim/adagrad.py,sha256=jIExkXD4UPgYQVplP__J5I4uekVjNmUFLJ1VuGUe8S4,22923 +torch/optim/adam.py,sha256=fvqFCqtzQu0ehnXPD3p33wOKtTJJZDSwp_KqaGbLOlw,40124 +torch/optim/adamax.py,sha256=HxXC6PK7pLaEfYJCyOrQcbnX3x7lrdHUQS8qe0vJhv8,17487 +torch/optim/adamw.py,sha256=0POi2ImkO2XtSBdZAPY0u1okTIzXmj8RF9yFLdrUsek,7376 +torch/optim/asgd.py,sha256=ChP8X1v1QLpkoN2fP66JTzrQHTCGRnLndbo2wj47Kio,16448 +torch/optim/lbfgs.py,sha256=obzQpGUK-UfXTXO9xZVYYMf87shCUHhODegsN-oY-E8,19702 +torch/optim/lr_scheduler.py,sha256=-NcFj-dSJWq6NOONeCNsIKJTH6L3pEO5oerVcJ93xxs,103402 +torch/optim/nadam.py,sha256=uVGHwApSmfPJaZUQJ9r706Xm6DqT73BR_u6_Sd4WPm8,26721 +torch/optim/optimizer.py,sha256=RosWCALluXXFhwNaYJC9eXztBUkf-3xs4xTdgfr3eD8,51385 +torch/optim/radam.py,sha256=DIwXjs4xI8niSrZ9Kp4Lj8bEv8LYbVfLLewFGC_MjAA,24927 +torch/optim/rmsprop.py,sha256=ETRPrSwF8pmMskZVgQOTBpciKOmtl8xwpu7yyxKZkrk,20590 +torch/optim/rprop.py,sha256=Mq6ZHRyNMJmeev-zcfsHvN3n_MboCjB10ZnEafsP70o,17695 +torch/optim/sgd.py,sha256=4JVTY0FlUNCqlNUWW_AwA5zez-vMYQjNOk8rInPTRoc,20296 +torch/optim/sparse_adam.py,sha256=M9uMYpTVzM2Jy7WOjLYWxGLTX7RTAh44lYN1fohPoQg,8046 +torch/optim/swa_utils.py,sha256=DLBth9zB5pb-RoJKz2hCRx3JCpFycSvk4u_MfiIUvew,23184 +torch/overrides.py,sha256=pDJDIuoaoGrKAyav72s1oEnS3DcyvaDWOSPuJADv_tU,108122 +torch/package/__init__.py,sha256=ZLLvoviHHErV-XQZagde2I4cuNDK49dFLRGUC5oyOFc,388 +torch/package/__pycache__/__init__.cpython-311.pyc,, +torch/package/__pycache__/_digraph.cpython-311.pyc,, +torch/package/__pycache__/_directory_reader.cpython-311.pyc,, +torch/package/__pycache__/_importlib.cpython-311.pyc,, +torch/package/__pycache__/_mangling.cpython-311.pyc,, +torch/package/__pycache__/_mock.cpython-311.pyc,, +torch/package/__pycache__/_package_pickler.cpython-311.pyc,, +torch/package/__pycache__/_package_unpickler.cpython-311.pyc,, +torch/package/__pycache__/_stdlib.cpython-311.pyc,, +torch/package/__pycache__/file_structure_representation.cpython-311.pyc,, +torch/package/__pycache__/find_file_dependencies.cpython-311.pyc,, +torch/package/__pycache__/glob_group.cpython-311.pyc,, +torch/package/__pycache__/importer.cpython-311.pyc,, +torch/package/__pycache__/package_exporter.cpython-311.pyc,, +torch/package/__pycache__/package_importer.cpython-311.pyc,, +torch/package/_digraph.py,sha256=sFyjCdLOgMvWt5Z8I5f_iliJqdnSRahZF58FPpv4DLA,5630 +torch/package/_directory_reader.py,sha256=4CzVDaz_zvXlYrgPfIJe5PGxyBHpFhMBGNiJvOeGOLM,1915 +torch/package/_importlib.py,sha256=chdvY9cde326LrzXj8VnjV_3eNqrMlhuPwhuEO0BvaQ,2998 +torch/package/_mangling.py,sha256=ultKR57Zsk2oEnpwCYGA2XHyA3dBp406yjUxFgHcdCs,1948 +torch/package/_mock.py,sha256=n4L1d-hoa5XO7kyAlxbYMjeZWa2vwQ_-7Ih_dZfH0SA,2866 +torch/package/_package_pickler.py,sha256=aJu-wsZUxM7AFr1VUJI6qClABvKKjoVaKDzHRW8QNoA,5785 +torch/package/_package_unpickler.py,sha256=hBxGPJDOF1wWmGVAfvN4FMCmm1o1Qt4pDKPmCdSdH-k,992 +torch/package/_stdlib.py,sha256=ttgTOJID6PsBLAyeNQ_tMaNXVF7tNEWUcH2_GIwI7WA,552 +torch/package/analyze/__init__.py,sha256=RtjmM0jmYQwfuv9mQoKgZQZBij-GQ4cFcJX7_-aihDg,130 +torch/package/analyze/__pycache__/__init__.cpython-311.pyc,, +torch/package/analyze/__pycache__/find_first_use_of_broken_modules.cpython-311.pyc,, +torch/package/analyze/__pycache__/is_from_package.cpython-311.pyc,, +torch/package/analyze/__pycache__/trace_dependencies.cpython-311.pyc,, +torch/package/analyze/find_first_use_of_broken_modules.py,sha256=0nn7RnQBbM_OvsyNz7AWHR1IL95wMs768O8vdMD2s7s,1114 +torch/package/analyze/is_from_package.py,sha256=5RN5zZTnmZiMxadQGK3OVYl38Q37nYS92FDVXrpQVEw,404 +torch/package/analyze/trace_dependencies.py,sha256=PHXKGtmWdCy7VczGvfVcV7AwFhFqfToPYVcTI5nSxec,2235 +torch/package/file_structure_representation.py,sha256=uJkvIJanCyW5QFFLGNWg5ICzo812ot_PuGeRGVnSmWE,4734 +torch/package/find_file_dependencies.py,sha256=-47nZhmmsTN92nmLp09KAmymnSzOAiXf5f65-umr_jc,3948 +torch/package/glob_group.py,sha256=aQo1kEtLj8JyyGzxSp3SNMrT1VhMoYZaUWdyNcfY3wQ,3658 +torch/package/importer.py,sha256=9SAhhM7FbU3W5F9FYgL8vZItnbrVC3c3lh9Z25ikquM,10239 +torch/package/package_exporter.py,sha256=EQGBZJxHhs5F3u4Vvcs9bAliRRO_o2_HInRxIyBGHao,52018 +torch/package/package_importer.py,sha256=yxJI1ZnEoF4P8oaRPCjiC-a_TPaWmtYizB8bEvyqqiA,32437 +torch/profiler/__init__.py,sha256=wkycjPawFAc2hKygHPFfcAfdkKKII5tR_Tp4GpGF0tw,1726 +torch/profiler/__pycache__/__init__.cpython-311.pyc,, +torch/profiler/__pycache__/_memory_profiler.cpython-311.pyc,, +torch/profiler/__pycache__/_pattern_matcher.cpython-311.pyc,, +torch/profiler/__pycache__/_utils.cpython-311.pyc,, +torch/profiler/__pycache__/itt.cpython-311.pyc,, +torch/profiler/__pycache__/profiler.cpython-311.pyc,, +torch/profiler/__pycache__/python_tracer.cpython-311.pyc,, +torch/profiler/_memory_profiler.py,sha256=yQXgBSdJaLvNDPnbDJNhzEh4n02nmCxx8Uo1Cmo2Md4,49566 +torch/profiler/_pattern_matcher.py,sha256=DJXdhBfr2k6DuW8ptOE9P0hANYrediNYwhLKi3ekeLw,25676 +torch/profiler/_utils.py,sha256=UaIPAW9xJxFSEulosABnLajpDWQluemlZPagEsWrXPg,22162 +torch/profiler/itt.py,sha256=GwDQYUY-qkdGn6Yg5VuS52aih12iJWLrPxdE3EMEXeg,1830 +torch/profiler/profiler.py,sha256=KRwwEOIg7ehlbkHVXAoMEPjGmXn4-AGrHCi_QOOyFBU,51818 +torch/profiler/python_tracer.py,sha256=mY1c9wqNHjZPmbJCrxNzduQiXUz2FM3wqYLsPwsEhtI,543 +torch/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/quantization/__init__.py,sha256=qdViHYX73iqc5xiNB9pKD6Ybzp33Wa2Xjv8j0DDA1eA,2810 +torch/quantization/__pycache__/__init__.cpython-311.pyc,, +torch/quantization/__pycache__/_numeric_suite.cpython-311.pyc,, +torch/quantization/__pycache__/_numeric_suite_fx.cpython-311.pyc,, +torch/quantization/__pycache__/_quantized_conversions.cpython-311.pyc,, +torch/quantization/__pycache__/fake_quantize.cpython-311.pyc,, +torch/quantization/__pycache__/fuse_modules.cpython-311.pyc,, +torch/quantization/__pycache__/fuser_method_mappings.cpython-311.pyc,, +torch/quantization/__pycache__/observer.cpython-311.pyc,, +torch/quantization/__pycache__/qconfig.cpython-311.pyc,, +torch/quantization/__pycache__/quant_type.cpython-311.pyc,, +torch/quantization/__pycache__/quantization_mappings.cpython-311.pyc,, +torch/quantization/__pycache__/quantize.cpython-311.pyc,, +torch/quantization/__pycache__/quantize_fx.cpython-311.pyc,, +torch/quantization/__pycache__/quantize_jit.cpython-311.pyc,, +torch/quantization/__pycache__/stubs.cpython-311.pyc,, +torch/quantization/__pycache__/utils.cpython-311.pyc,, +torch/quantization/_numeric_suite.py,sha256=kGqWAUJhc0DKg9RjXhiKRx7dxfmX-M5Y1w9jAEnU0Z0,779 +torch/quantization/_numeric_suite_fx.py,sha256=LkEYZQLt_CRTe4d_9EAvCzpuKKHTmTjYoXLXcawjN-I,752 +torch/quantization/_quantized_conversions.py,sha256=OM1zSUmz0HKz_aeVfQf92ZG2xiTFqAVWpQMas5gy8AQ,5204 +torch/quantization/fake_quantize.py,sha256=AZes9LhE_KB-0Q_nV0Qy5kSOY0N9cvv7NPEpdjTwBz4,1015 +torch/quantization/fuse_modules.py,sha256=qdHvwLDuYEu2zJaAKtztSbw157nYTS5rJLCKBbBLPRA,732 +torch/quantization/fuser_method_mappings.py,sha256=QT2XYwG6LRioeQPh0lrmZ2awt05ejXODVO3S1BFgu-I,512 +torch/quantization/fx/__init__.py,sha256=s3Wh6JJbmUuv-pk4yieKjDpsHZRUkEqW04jRMT9iJPs,594 +torch/quantization/fx/__pycache__/__init__.cpython-311.pyc,, +torch/quantization/fx/__pycache__/_equalize.cpython-311.pyc,, +torch/quantization/fx/__pycache__/convert.cpython-311.pyc,, +torch/quantization/fx/__pycache__/fuse.cpython-311.pyc,, +torch/quantization/fx/__pycache__/fusion_patterns.cpython-311.pyc,, +torch/quantization/fx/__pycache__/graph_module.cpython-311.pyc,, +torch/quantization/fx/__pycache__/match_utils.cpython-311.pyc,, +torch/quantization/fx/__pycache__/pattern_utils.cpython-311.pyc,, +torch/quantization/fx/__pycache__/prepare.cpython-311.pyc,, +torch/quantization/fx/__pycache__/quantization_patterns.cpython-311.pyc,, +torch/quantization/fx/__pycache__/quantization_types.cpython-311.pyc,, +torch/quantization/fx/__pycache__/utils.cpython-311.pyc,, +torch/quantization/fx/_equalize.py,sha256=lqBIInk3Gq4t2kdsaO0vpkjpCXriv5WQAgCP2uKlbak,1251 +torch/quantization/fx/convert.py,sha256=hQUg5B8P8CByLc7JIHeBEm15VHyuxPWxv2zFgo5oYmc,387 +torch/quantization/fx/fuse.py,sha256=xKT6NKY-nOHbiUAhvd6UyFLOWaN7IzTXzQkQQRM1OCM,381 +torch/quantization/fx/fusion_patterns.py,sha256=qSu3LbPeOO8dbO14vvA-yHKvoflh8F4Ta6JLCYOfmxU,416 +torch/quantization/fx/graph_module.py,sha256=Q2bphJkn60j97Wvy_gAOol-sAx5VKavZujO5Z1ChBmg,574 +torch/quantization/fx/match_utils.py,sha256=qFqoYdkrjA4jB2Figx-06RzeiN3KGs5WYt68TRqRYKs,457 +torch/quantization/fx/pattern_utils.py,sha256=8YjiPNSE0n9A6pLrlsv92foYAx37EWl93C5tHbxKrUQ,1299 +torch/quantization/fx/prepare.py,sha256=k6hE4OfbcI8r-csnlRP9gd8E4Hbak6NIO-q4E2_TPTg,387 +torch/quantization/fx/quantization_patterns.py,sha256=Qws7uHQqdnaF9nK_UPx_ACyQiREN7D40F9QqcuFafdQ,2088 +torch/quantization/fx/quantization_types.py,sha256=FC78rHJfePvTdJfTmXfKCBUXtz2mbF1SK7s3401zhZY,396 +torch/quantization/fx/utils.py,sha256=3KZFWzVZPxCpzxlNUxv3xigcpmYg2AmtaOyJfKtfuCM,724 +torch/quantization/observer.py,sha256=J8sHSdsC9i6OqY4_OpzrTOqeHS1fZF_VmN132CZQWsY,1079 +torch/quantization/qconfig.py,sha256=NT4W-h_HrXKX06TnuRxrvfnv5Kwho6Xa4vd7ESA7RHs,944 +torch/quantization/quant_type.py,sha256=wwPL8MSv-G6QV5LogvkL4xRJ28Ido0yA7kPzh2dsIo0,399 +torch/quantization/quantization_mappings.py,sha256=F4DwF6qNFg8h4xhKvp8nBGuLLr5VwNjNjRoSTWFebdI,1148 +torch/quantization/quantize.py,sha256=g8mA0kRJIO2z63PuUcawc9bovat-uVjeUxIKFwvLb-I,804 +torch/quantization/quantize_fx.py,sha256=hBC2pUobd9oeYUliLKlIYE2dQ-6BQvdxg2xTyvkFLa8,736 +torch/quantization/quantize_jit.py,sha256=JNw2_M1d6EITYzYrXlrAsUac_4tX-IDMuAO_7prxhk8,714 +torch/quantization/stubs.py,sha256=ukEM_vZ34I3UdBnhFPn5GFL713MwFlEysXUAfYGshnA,392 +torch/quantization/utils.py,sha256=kOpHHmJ602vEybccn81hVLOHwY0CQ2YdnG9n9kOVcQM,833 +torch/quasirandom.py,sha256=APrygrXAGBNtwTyf4v_gqJtT6AvOukzyCkQQf7l1p7k,7905 +torch/random.py,sha256=uouV1Ty8gpCkI3QX-mBEFqteAe2aoe6ttVKK6Tpch7I,9515 +torch/return_types.py,sha256=n2BeQf08b5YJcfOkm6JSqVUPDcy5To831qwGVNWiFf8,1537 +torch/return_types.pyi,sha256=hHUj4bIDTq9gprXJamZ5xR0U5vLVBqQRSOwE0uIg8G4,17837 +torch/serialization.py,sha256=gfKlt1TsxveU7DTlEw9SFLcAx-h5uAwLXxtJZPU1KcE,87310 +torch/share/cmake/ATen/ATenConfig.cmake,sha256=xNnEU2M076jtMAEujFyD6PIVK1CH7_VmvMiQO9hOfjU,299 +torch/share/cmake/Caffe2/Caffe2Config.cmake,sha256=QPGaT88ucHtrXtUIwYwMYFW-zPaKyMv29oA1AT8YD4s,5347 +torch/share/cmake/Caffe2/Caffe2Targets-release.cmake,sha256=ktjB35_r3hByVk9vsV7aPUzb0SYhq7ek8pW53WaOZZg,2580 +torch/share/cmake/Caffe2/Caffe2Targets.cmake,sha256=Mh98tDdDHTWQxIR8vr9WYxqGO44C0P4VHeMbjx_chBI,7413 +torch/share/cmake/Caffe2/FindCUDAToolkit.cmake,sha256=kl99c2bqe1LwRcg7k8VkwqL99pQ9CF9vTwG-kfKIu84,38768 +torch/share/cmake/Caffe2/FindCUDSS.cmake,sha256=54i8T6FKzYP1KMtHTXaUhsuCuXHogM8SDadMAUNZT8I,2698 +torch/share/cmake/Caffe2/FindCUSPARSELT.cmake,sha256=cjNafMBEzBGv0kZAQqMxdpku0Vyf2Yrn3vJ_OVHxdUQ,3068 +torch/share/cmake/Caffe2/FindSYCLToolkit.cmake,sha256=w2vDQMm9VF71HM4v7XmJK8DTdiE3Q_d91rmjB8nE-kQ,4547 +torch/share/cmake/Caffe2/Modules_CUDA_fix/FindCUDA.cmake,sha256=uH9LCtcsv8VgeuFEyyMTugFUiILpcgM3oEnJZka4AB0,459 +torch/share/cmake/Caffe2/Modules_CUDA_fix/FindCUDNN.cmake,sha256=NKwIx_LRJ3uu9oJC9_apdlRooZ-uR_izK1_deByfEQo,3085 +torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/CMakeInitializeConfigs.cmake,sha256=v1O1FBKmJWk1h-Zbh8M6qVlP4OqVsmCcYD8zwdVfb6Q,1657 +torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA.cmake,sha256=vz7pTfsrZCUqr1wtAtEP7gkiTKex3dDaAK4AZLj2hRQ,86626 +torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/make2cmake.cmake,sha256=_KLZxL3AhZehZKubThy4o2C_gEH5mm4h3kMpxLHgajU,3925 +torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/parse_cubin.cmake,sha256=h3Ka8c-mmE2Majl0s8342faZWVPAaDbO1rM_4m--YzA,3439 +torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/run_nvcc.cmake,sha256=xCYQQ4GWx5BsAWYZZEOo79e_Pu_pK7S_pWM-6XNS-Tw,11813 +torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake,sha256=PUypM4YCR0F9MGhlKdZf42uuJ75p7QscDUXYjE0jm5o,11478 +torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindPackageMessage.cmake,sha256=ToKFxPt7HSmEA014cFkMZl79quM2gpF7tmcP8h1BuYs,1564 +torch/share/cmake/Caffe2/public/LoadHIP.cmake,sha256=m5IZN3iFDIH2dYYkm_JusVe0bM3tmSdibD2GPKIn7AI,11583 +torch/share/cmake/Caffe2/public/cuda.cmake,sha256=d81xNZF8GCHeT7xcVxptrewQWF1QXIS1BvyX75EL-kM,14952 +torch/share/cmake/Caffe2/public/gflags.cmake,sha256=YrTkm-nQX6N3mh3tpY9a-i1_s_lAlifW0V_UZdAHS-M,2620 +torch/share/cmake/Caffe2/public/glog.cmake,sha256=zy1mZaicXNUHxSG7eOZ4ZVnbgk5qN-IuG_-jXUi98nk,2320 +torch/share/cmake/Caffe2/public/mkl.cmake,sha256=y1E6axKvT-AvgfmJrVMptYDVJ5svt723jNuQYIJp_kg,1317 +torch/share/cmake/Caffe2/public/mkldnn.cmake,sha256=7D8oS35genLaLdSPfgh3foWXQMsa2bUsQnB5WKfKPAA,444 +torch/share/cmake/Caffe2/public/protobuf.cmake,sha256=weW3OuHBIqIt0KpWRWzAnzSCu-AJ9naUfkjxcvaRl5c,4003 +torch/share/cmake/Caffe2/public/utils.cmake,sha256=1UMBDrswArGXYU6AViKYbIAGjIDwIhbnZyvIrvltNBo,23878 +torch/share/cmake/Caffe2/public/xpu.cmake,sha256=ZIqtXMW4xoEMW0e534gK8EnfdEvy3e5LpXHK5ZK2bag,1551 +torch/share/cmake/Tensorpipe/TensorpipeTargets-release.cmake,sha256=F61imoXwiRatw97-TefIPmRwlsVTdvAXItQDGEFGGPI,1803 +torch/share/cmake/Tensorpipe/TensorpipeTargets.cmake,sha256=Q38LErNHgRz5BhUrCMqs2L3eNClmq1llrh7KZMCyIyI,4837 +torch/share/cmake/Torch/TorchConfig.cmake,sha256=hwK5dIaH-1nA2uJRibCKluPw0fqYDm_22MSHYmSv4z4,4945 +torch/share/cmake/Torch/TorchConfigVersion.cmake,sha256=9UVQsRf7gk6XRd1ZLOcVvWiI5-59S3XWiw7yXWgrPow,367 +torch/signal/__init__.py,sha256=UvInN-gZ-pXXT9wvdSjrESHuSVNz5qEcnLMRzycbXiU,46 +torch/signal/__pycache__/__init__.cpython-311.pyc,, +torch/signal/windows/__init__.py,sha256=pqkf3lWIkwGpjDn_njgkrpqzYYCBpC25zEK7iD4FGaI,383 +torch/signal/windows/__pycache__/__init__.cpython-311.pyc,, +torch/signal/windows/__pycache__/windows.cpython-311.pyc,, +torch/signal/windows/windows.py,sha256=qJi7w3QzoL4Ime9gKazKAaErMj-QreMgDHGr1LCz5cI,22720 +torch/sparse/__init__.py,sha256=0n3D6V_sF2pYzuthQaPy8qx-JBmgtX4-5UEEgDb0xIY,26965 +torch/sparse/__pycache__/__init__.cpython-311.pyc,, +torch/sparse/__pycache__/_semi_structured_conversions.cpython-311.pyc,, +torch/sparse/__pycache__/_semi_structured_ops.cpython-311.pyc,, +torch/sparse/__pycache__/_triton_ops.cpython-311.pyc,, +torch/sparse/__pycache__/_triton_ops_meta.cpython-311.pyc,, +torch/sparse/__pycache__/semi_structured.cpython-311.pyc,, +torch/sparse/_semi_structured_conversions.py,sha256=F1nRbKXb3PClzvZ26vvU8Ud69elhsYqPyN2DocHGco4,14094 +torch/sparse/_semi_structured_ops.py,sha256=6uj2Fvn1JYyBfT2sF9uDI7EzXivvmKhFkxAxEeicJhM,10612 +torch/sparse/_triton_ops.py,sha256=9oBcxE_L_5NDrGZkEi7HdbrYtgzSEvZcfc8XFKUxMjI,89779 +torch/sparse/_triton_ops_meta.py,sha256=y16nK_DKMcm8yH-jzC_HpVbBL30xwIexsxwkvrLoD8w,501362 +torch/sparse/semi_structured.py,sha256=MwHbXxtUtbj01gz_fBfGRsxFQNvO1g0qrUHeYd106Fw,33260 +torch/special/__init__.py,sha256=vua7SYF8ZjpL6Gpf_M6VCrVW1b-WHfJqfx4-KTNMZgU,32035 +torch/special/__pycache__/__init__.cpython-311.pyc,, +torch/storage.py,sha256=-ut6ECHI8BugF6-igb8OlN6vVfye5f50Ly8JzBZGf-A,52636 +torch/testing/__init__.py,sha256=KWGXAtM1wGYjenz5lGf2H16r3-qI1ZmOVsHuxHn7S2k,219 +torch/testing/__pycache__/__init__.cpython-311.pyc,, +torch/testing/__pycache__/_comparison.cpython-311.pyc,, +torch/testing/__pycache__/_creation.cpython-311.pyc,, +torch/testing/__pycache__/_utils.cpython-311.pyc,, +torch/testing/_comparison.py,sha256=-SuWgM5IXl3CUnTUZXL1Y8lQ7rJYMRdjMrWz4SCnJuI,67574 +torch/testing/_creation.py,sha256=t0WsWWWwuWIKu8S7fTSUeEryLxogRishdWJWEXFvoPg,12156 +torch/testing/_internal/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/__pycache__/autocast_test_lists.cpython-311.pyc,, +torch/testing/_internal/__pycache__/autograd_function_db.cpython-311.pyc,, +torch/testing/_internal/__pycache__/check_kernel_launches.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_cuda.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_device_type.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_dist_composable.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_distributed.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_dtype.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_fsdp.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_jit.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_methods_invocations.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_mkldnn.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_modules.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_mps.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_nn.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_ops_unbacked.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_optimizers.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_pruning.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_quantization.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_quantized.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_subclass.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_utils.cpython-311.pyc,, +torch/testing/_internal/__pycache__/common_xpu.cpython-311.pyc,, +torch/testing/_internal/__pycache__/composite_compliance.cpython-311.pyc,, +torch/testing/_internal/__pycache__/custom_op_db.cpython-311.pyc,, +torch/testing/_internal/__pycache__/custom_tensor.cpython-311.pyc,, +torch/testing/_internal/__pycache__/dist_utils.cpython-311.pyc,, +torch/testing/_internal/__pycache__/dynamo_pytree_test_utils.cpython-311.pyc,, +torch/testing/_internal/__pycache__/dynamo_test_failures.cpython-311.pyc,, +torch/testing/_internal/__pycache__/fake_config_module.cpython-311.pyc,, +torch/testing/_internal/__pycache__/fake_config_module2.cpython-311.pyc,, +torch/testing/_internal/__pycache__/fake_config_module3.cpython-311.pyc,, +torch/testing/_internal/__pycache__/hop_db.cpython-311.pyc,, +torch/testing/_internal/__pycache__/hypothesis_utils.cpython-311.pyc,, +torch/testing/_internal/__pycache__/inductor_utils.cpython-311.pyc,, +torch/testing/_internal/__pycache__/jit_metaprogramming_utils.cpython-311.pyc,, +torch/testing/_internal/__pycache__/jit_utils.cpython-311.pyc,, +torch/testing/_internal/__pycache__/logging_tensor.cpython-311.pyc,, +torch/testing/_internal/__pycache__/logging_utils.cpython-311.pyc,, +torch/testing/_internal/__pycache__/quantization_torch_package_models.cpython-311.pyc,, +torch/testing/_internal/__pycache__/static_module.cpython-311.pyc,, +torch/testing/_internal/__pycache__/subclasses.cpython-311.pyc,, +torch/testing/_internal/__pycache__/torchbind_impls.cpython-311.pyc,, +torch/testing/_internal/__pycache__/triton_utils.cpython-311.pyc,, +torch/testing/_internal/__pycache__/two_tensor.cpython-311.pyc,, +torch/testing/_internal/autocast_test_lists.py,sha256=UebIEDySNOBGqMvy6_35gb-631DGLb-uMUnjF5CWNys,28483 +torch/testing/_internal/autograd_function_db.py,sha256=mR_aqHOk6Ib7aZXSW_Y32mfWe7e9iymYw0sJAchdAzY,19755 +torch/testing/_internal/check_kernel_launches.py,sha256=KDW-fjae41C_aetbWyQGeXsEBWpv0cQaM9S57AdbQPo,6027 +torch/testing/_internal/codegen/__init__.py,sha256=8QLhisbHub6VJl6egijnrOPKK5QNAe5FJhfcxEelj4Y,22 +torch/testing/_internal/codegen/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/common_cuda.py,sha256=wXBvxEIx0V8r7iVMwlbAlyRQAvcY8RjGhWam8ZVcxVo,20749 +torch/testing/_internal/common_device_type.py,sha256=u9xE282GZS0BSUZ0tj4Wrk6GDzxA_73UgX1PqDqeEjM,78103 +torch/testing/_internal/common_dist_composable.py,sha256=IwtTq8pjXlXDotDZE_Fpy1o7tkpAP_WnNcx05GpJ3YA,3577 +torch/testing/_internal/common_distributed.py,sha256=5xPiTvkpUCVhMel193UUfbNJJRcaZKroUBVtGdyePSs,77548 +torch/testing/_internal/common_dtype.py,sha256=FTpLgdDdpjIZ1qNYDhBTmvumoryYYJe6dZA5eWf5Jtc,5589 +torch/testing/_internal/common_fsdp.py,sha256=3i8dvE-ig8w8I6vsE4txw7FTadKFFLtJhWfuWdnmiOE,63927 +torch/testing/_internal/common_jit.py,sha256=3TEnvapmxqgf_O31FoBrIJgcGbRp_uERaDJyut39zaU,15823 +torch/testing/_internal/common_methods_invocations.py,sha256=0BXkOmL5n_veijApVH7oVLgUJbpv2G9ot9hJhJoh_gs,1330190 +torch/testing/_internal/common_mkldnn.py,sha256=b0E0qdrCqWWWhonniOmNFT7_aE8DwdDB12cqjuvwt4Y,3895 +torch/testing/_internal/common_modules.py,sha256=bUS5wK5jEn5d5CnD6pUOBrTBenXoXjccUS-_4k0TOBU,227778 +torch/testing/_internal/common_mps.py,sha256=uSpkldJorzjVPH8mQ0AjkBVjX3mVmTIZuASgh4-C8yk,36324 +torch/testing/_internal/common_nn.py,sha256=1BAGkKwPvMZ87G9aLDAu6j8z9-b_6c1_O3TcTCbjEQ0,173860 +torch/testing/_internal/common_ops_unbacked.py,sha256=wSVpv7UEzojYtAk1bmfJQK5RelSo06CWKa0BNxU3zm0,8139 +torch/testing/_internal/common_optimizers.py,sha256=PINJdcjRDV0I2smt1yjFYytFvjc07LiXlDEjKQDfKrw,86428 +torch/testing/_internal/common_pruning.py,sha256=nYwQmrHvNQpRmIkuxK7yYki2CwmaMVJKsnIW_Z-xIw8,13655 +torch/testing/_internal/common_quantization.py,sha256=POrIhsraTvmhDxTl11gfuTzAgiATDhdkbvsfwmSR0gQ,113562 +torch/testing/_internal/common_quantized.py,sha256=DFpHK9zUlWbl81_HN623IAeKTEZ0ZDbSp_A0xJcU37o,26229 +torch/testing/_internal/common_subclass.py,sha256=DnQmbpIJwQulcL2_LKY-0Eecns3UTAmPGjnjPP9_Aek,12277 +torch/testing/_internal/common_utils.py,sha256=Mipwm2PVFvethSbXSNE1zRkf_CCu0SAF5ge4t7rsIsM,254687 +torch/testing/_internal/common_xpu.py,sha256=5MToq646Q7bYRLYlBnvkP-AZNaoQWJfRqvSNGDj_vug,2265 +torch/testing/_internal/composite_compliance.py,sha256=QyoLKwh7grzb8MqdhxamA7X5Bvo_6YqFAnSAHWf1JXM,26584 +torch/testing/_internal/custom_op_db.py,sha256=bJ1_wqRX4UrohsLp29Z3qKJO5M4WRgoXlgpIKss3wgI,20991 +torch/testing/_internal/custom_tensor.py,sha256=f1Yc5nNAanrdQ5gCflY6a-Vs10xZXWYRB8jd5fEKzO4,5337 +torch/testing/_internal/data/__init__.py,sha256=8QLhisbHub6VJl6egijnrOPKK5QNAe5FJhfcxEelj4Y,22 +torch/testing/_internal/data/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/data/__pycache__/network1.cpython-311.pyc,, +torch/testing/_internal/data/__pycache__/network2.cpython-311.pyc,, +torch/testing/_internal/data/network1.py,sha256=ksE5iUCq6Hhp8772pRhqpUna414huKdlzvh92oQgxTc,169 +torch/testing/_internal/data/network2.py,sha256=FlOrR6LuubeOhZHKURvkoWqIGHlvMNX2SyoZ22XPvfg,199 +torch/testing/_internal/dist_utils.py,sha256=WBlzvl4sb7bXl8T_Dv_-LDgZtUp4bdGCFfCNpbDjWck,7255 +torch/testing/_internal/distributed/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/distributed/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/distributed/__pycache__/checkpoint_utils.cpython-311.pyc,, +torch/testing/_internal/distributed/__pycache__/common_state_dict.cpython-311.pyc,, +torch/testing/_internal/distributed/__pycache__/ddp_under_dist_autograd_test.cpython-311.pyc,, +torch/testing/_internal/distributed/__pycache__/distributed_test.cpython-311.pyc,, +torch/testing/_internal/distributed/__pycache__/distributed_utils.cpython-311.pyc,, +torch/testing/_internal/distributed/__pycache__/fake_pg.cpython-311.pyc,, +torch/testing/_internal/distributed/__pycache__/multi_threaded_pg.cpython-311.pyc,, +torch/testing/_internal/distributed/__pycache__/rpc_utils.cpython-311.pyc,, +torch/testing/_internal/distributed/_shard/__init__.py,sha256=mIQlHRV-d0Zr3TOwF0Vr4fBPuPk4uNR8rzKPy1EhlPw,27 +torch/testing/_internal/distributed/_shard/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/distributed/_shard/__pycache__/test_common.cpython-311.pyc,, +torch/testing/_internal/distributed/_shard/sharded_tensor/__init__.py,sha256=h9RFxi9o1Jukm8ez0X7m1JNc9_zEcvBPNHNtODTMIPA,3212 +torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/_test_ops_common.cpython-311.pyc,, +torch/testing/_internal/distributed/_shard/sharded_tensor/__pycache__/_test_st_common.cpython-311.pyc,, +torch/testing/_internal/distributed/_shard/sharded_tensor/_test_ops_common.py,sha256=Ngk6Ynz4RSlUmj5oDTB0d-1jr4hl7DiVvyGy1IZefcw,4011 +torch/testing/_internal/distributed/_shard/sharded_tensor/_test_st_common.py,sha256=Qys7mhe_-Ew8m14WqTGlJ1VJf_gINSZyraRfyDNP44w,1618 +torch/testing/_internal/distributed/_shard/test_common.py,sha256=0siUUMgcT2SqfCOGQvb4_k-_r2yHoOALNG0kM1OTGwc,1219 +torch/testing/_internal/distributed/_tensor/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/distributed/_tensor/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/distributed/_tensor/__pycache__/common_dtensor.cpython-311.pyc,, +torch/testing/_internal/distributed/_tensor/common_dtensor.py,sha256=SK8hNvUEYAvY80ea54htmIVF06JrX819Yr1PzJBC21o,52552 +torch/testing/_internal/distributed/checkpoint_utils.py,sha256=xxtwj393h3AM_1tT4CTCdB0qd_zyYyCc4xvALUjOFDc,6267 +torch/testing/_internal/distributed/common_state_dict.py,sha256=HU0cY-TeqC7QE6BFuQd_GnLF-H5Up0_IY-y9AFw1Ui4,6725 +torch/testing/_internal/distributed/ddp_under_dist_autograd_test.py,sha256=jU3edCICiCyIKvbKM3HSreOzUApjTfUFBefCzJ0wt-c,27326 +torch/testing/_internal/distributed/distributed_test.py,sha256=j59ak0RnBJB3AXHdfrSR8CPCNgG_Vsp0XENMhHB-k1g,455010 +torch/testing/_internal/distributed/distributed_utils.py,sha256=rZ606dkmJjUfgKsgyRMTYOZATiaXwryb6GRcsbQ3wdg,1932 +torch/testing/_internal/distributed/fake_pg.py,sha256=Rw_i-rRfacMra5a7Pg7cWf_YG7YEDCUV3S2LldZ8lL8,1150 +torch/testing/_internal/distributed/multi_threaded_pg.py,sha256=V8fEjM0S09R2gT5y4YVAfTg837EyWCiArD_U6BSrPKM,22866 +torch/testing/_internal/distributed/nn/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/distributed/nn/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/distributed/nn/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/distributed/nn/api/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/distributed/nn/api/__pycache__/remote_module_test.cpython-311.pyc,, +torch/testing/_internal/distributed/nn/api/remote_module_test.py,sha256=GohOmmJgmhgehHlcwU55daRtg1zl1NN447dORb_An9E,29751 +torch/testing/_internal/distributed/rpc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/distributed/rpc/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/__pycache__/dist_autograd_test.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/__pycache__/dist_optimizer_test.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/__pycache__/faulty_agent_rpc_test.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/__pycache__/faulty_rpc_agent_test_fixture.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/__pycache__/rpc_agent_test_fixture.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/__pycache__/rpc_test.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/__pycache__/tensorpipe_rpc_agent_test_fixture.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/dist_autograd_test.py,sha256=m75fbca1a1AKYej-xilTAphhZtTjlwRJKOnNevpH05E,107635 +torch/testing/_internal/distributed/rpc/dist_optimizer_test.py,sha256=HibDFt3xUTuOmm2mXOAnEzCctb48WKGzNPWZ-BRr-hk,10574 +torch/testing/_internal/distributed/rpc/examples/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/distributed/rpc/examples/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/examples/__pycache__/parameter_server_test.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/examples/__pycache__/reinforcement_learning_rpc_test.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/examples/parameter_server_test.py,sha256=_sm9R6jjnAkl49tk_dOH9dRZIdDEmuAmOy4DUkMa9ck,4568 +torch/testing/_internal/distributed/rpc/examples/reinforcement_learning_rpc_test.py,sha256=JDYkhzz7ZsK-H-IVEPic9Sxu7ZLv0taf1REfAnWgTbM,9268 +torch/testing/_internal/distributed/rpc/faulty_agent_rpc_test.py,sha256=zcYRImWXoIg0qmVWpkPlTKA_t0neQPYyi9YXm8IdgTo,14239 +torch/testing/_internal/distributed/rpc/faulty_rpc_agent_test_fixture.py,sha256=ziVR9Z4tDhdG0omKEvGVfOXDN2I10eeXCcMYJ6G_vYk,2142 +torch/testing/_internal/distributed/rpc/jit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/distributed/rpc/jit/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/jit/__pycache__/dist_autograd_test.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/jit/__pycache__/rpc_test.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/jit/__pycache__/rpc_test_faulty.cpython-311.pyc,, +torch/testing/_internal/distributed/rpc/jit/dist_autograd_test.py,sha256=bzq_ajaMCiVEFFmI1IsrK_bRKneEDMyyKOkXuiuq_ZE,4172 +torch/testing/_internal/distributed/rpc/jit/rpc_test.py,sha256=5Z1TtfUjWPqMWS6hlGq53PvW34Cs5A6WTZ3iMxoKTNw,46869 +torch/testing/_internal/distributed/rpc/jit/rpc_test_faulty.py,sha256=8TjJOzAvAlbE2qy3izpDBUHnUzog6r7jjEPciYoLs7U,7974 +torch/testing/_internal/distributed/rpc/rpc_agent_test_fixture.py,sha256=y5P27BCxk82IPTVOzov4-Ojr3_tykIOGlLlSEwWEv_w,1874 +torch/testing/_internal/distributed/rpc/rpc_test.py,sha256=kOBdxxXTM9FWJ7V4JNQUBikliwSgmrcMoftKwFu-Sr4,225756 +torch/testing/_internal/distributed/rpc/tensorpipe_rpc_agent_test_fixture.py,sha256=6RTlsS0sKBh2iZqvLm_ieqIQpjSuwqlYc7qdgjb96YM,974 +torch/testing/_internal/distributed/rpc_utils.py,sha256=_ytL9c2DUzSMzTVyoOTKg8UJVlINNdy0R4mtFXBR0rI,6615 +torch/testing/_internal/dynamo_pytree_test_utils.py,sha256=5zPqhGQDcqd4Ynshj8cS_TVz_syUcPa_nKWDAKCQ0lY,1059 +torch/testing/_internal/dynamo_test_failures.py,sha256=RHMRhAL0qKlVfHftaisckwbED_660cYHzyxCP4ANHVI,5440 +torch/testing/_internal/fake_config_module.py,sha256=py4Y9FaDJ7vdrG2xgJGVkuySNFyN1fe_pXoSgKlJ3xw,1559 +torch/testing/_internal/fake_config_module2.py,sha256=gk8fSMjiMktUYrbdWi-Yvm4aH6TeOBa-7nnB5SDMZwQ,343 +torch/testing/_internal/fake_config_module3.py,sha256=z6yJU2RkB7XBD22EDgPgRo5VVtr-iDR-U0HLdmMF8JU,220 +torch/testing/_internal/generated/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/generated/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/generated/__pycache__/annotated_fn_args.cpython-311.pyc,, +torch/testing/_internal/generated/annotated_fn_args.py,sha256=_CFByfk35O5U_s3tnvTw7Thm5tbWMaqG7PGu1oIlk_M,561508 +torch/testing/_internal/hop_db.py,sha256=uPTK9DAnN4ygDlbZCzP44QOZ1YkvxYKyCNKskqcCdjA,21951 +torch/testing/_internal/hypothesis_utils.py,sha256=KzYeCpETyJ78U5bd7KJYHr9aRhkeS86JUHKEt2t2CHs,15161 +torch/testing/_internal/inductor_utils.py,sha256=oP6jzu4-l0Uwo9IGxjGQYAeg_J0N_UZpcso4EXSuLqI,15029 +torch/testing/_internal/jit_metaprogramming_utils.py,sha256=z1YykYJIqnC2Ik2EepGoB_Hicux88m2Wytze3LX769o,34155 +torch/testing/_internal/jit_utils.py,sha256=cGn3Orv7a1b5sfYwOtJEwfljyQb2UEXYZHxGFHwx_BY,33964 +torch/testing/_internal/logging_tensor.py,sha256=d91-JZrIa-zByXCSVLRuWjbXViUhjXFkqxwV4NwPieQ,6615 +torch/testing/_internal/logging_utils.py,sha256=hHrY_eSidr5n7N602cfaa_rp9tuksIEhNq54g3y-NfI,8240 +torch/testing/_internal/opinfo/__init__.py,sha256=6PWvlARagSjyrZW5xCOA5YGQjLqwaijc36TFX4UrU9g,116 +torch/testing/_internal/opinfo/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/opinfo/__pycache__/core.cpython-311.pyc,, +torch/testing/_internal/opinfo/__pycache__/refs.cpython-311.pyc,, +torch/testing/_internal/opinfo/__pycache__/utils.cpython-311.pyc,, +torch/testing/_internal/opinfo/core.py,sha256=AAf-jgR1cIu-f-J0g6yjceqYXUwHX0hje9o6u0Q6IoM,127681 +torch/testing/_internal/opinfo/definitions/__init__.py,sha256=8IK85dB9S_tx3aiyuUwfSYLKgoGRaBmVnajBKn__Vqk,452 +torch/testing/_internal/opinfo/definitions/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/opinfo/definitions/__pycache__/_masked.cpython-311.pyc,, +torch/testing/_internal/opinfo/definitions/__pycache__/fft.cpython-311.pyc,, +torch/testing/_internal/opinfo/definitions/__pycache__/linalg.cpython-311.pyc,, +torch/testing/_internal/opinfo/definitions/__pycache__/nested.cpython-311.pyc,, +torch/testing/_internal/opinfo/definitions/__pycache__/signal.cpython-311.pyc,, +torch/testing/_internal/opinfo/definitions/__pycache__/sparse.cpython-311.pyc,, +torch/testing/_internal/opinfo/definitions/__pycache__/special.cpython-311.pyc,, +torch/testing/_internal/opinfo/definitions/_masked.py,sha256=SudppCxRYJbWo3AzBEGpJ2o3GniJNsgpjpGf4hgbqj8,52992 +torch/testing/_internal/opinfo/definitions/fft.py,sha256=6rjSHPTRcvEA4Sk2ErE5WiJTsk5oKmtYCXqliY221-Y,34972 +torch/testing/_internal/opinfo/definitions/linalg.py,sha256=bdd9jOzQhKhbYVuiBqpbnqhj6SnTwNx8ph3OfaHc1_Y,103844 +torch/testing/_internal/opinfo/definitions/nested.py,sha256=NBYahzMS7eMKqlQlDfK02tQEFzIERQYsr5wMt9zUPlw,60962 +torch/testing/_internal/opinfo/definitions/signal.py,sha256=LM2_KWXES706JWWxUyf0oPbuz75WaFZ_fpR_UM0AsCY,15351 +torch/testing/_internal/opinfo/definitions/sparse.py,sha256=Qu8nC0I6uqNZp9rTsRvKE2651Tiyj5e5Ldqa1SD3iH8,33990 +torch/testing/_internal/opinfo/definitions/special.py,sha256=X-hckhKa6M0oaZaCE5GlNdtmUlQERbkWlfmm9UQ47-8,38344 +torch/testing/_internal/opinfo/refs.py,sha256=8AK4TpjI2I_4RAirGNEEgUqdOFX3OTRKZzYRzL2FhF8,8601 +torch/testing/_internal/opinfo/utils.py,sha256=h4MlSwpGT0NImvYqQE9M77ZBJC4ZGXbb5JaFTHKDxME,9042 +torch/testing/_internal/optests/__init__.py,sha256=o-8t0Cva860Tcw0Ig_-rBr6-9Uuc9bDRWRMG6FgKzbA,372 +torch/testing/_internal/optests/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/optests/__pycache__/aot_autograd.cpython-311.pyc,, +torch/testing/_internal/optests/__pycache__/autograd_registration.cpython-311.pyc,, +torch/testing/_internal/optests/__pycache__/fake_tensor.cpython-311.pyc,, +torch/testing/_internal/optests/__pycache__/generate_tests.cpython-311.pyc,, +torch/testing/_internal/optests/__pycache__/make_fx.cpython-311.pyc,, +torch/testing/_internal/optests/aot_autograd.py,sha256=-kidAO1i5dDRsdYFthkXhibAV6-w4R_zKnb6cvhmwVs,7123 +torch/testing/_internal/optests/autograd_registration.py,sha256=d1hcQqT58qY_re5xsDl7BL17EraIrvAHvcsTjvrKcQM,5850 +torch/testing/_internal/optests/fake_tensor.py,sha256=WuT0PGbogjTGPJeRYo2JiQYoTzj_iZET56L1zkv0W4w,257 +torch/testing/_internal/optests/generate_tests.py,sha256=03ddn1BtUlpG0bO2JI86FVXY5k1eEAp8Png6NuafHeE,32403 +torch/testing/_internal/optests/make_fx.py,sha256=cERDt-I8VgMfwqddFNkbnkWnBWPZZegU5Yo0ijAtF5E,3265 +torch/testing/_internal/py312_intrinsics.py,sha256=fYgzb7zJXmlxR7JP5sp9RyxwMo3L5XubgD8M7vagwfE,854 +torch/testing/_internal/quantization_torch_package_models.py,sha256=DnChjrnLG54TWRhINvfTxOFmNZt5WflDBJDpam4fj9w,951 +torch/testing/_internal/static_module.py,sha256=bwmblZ7N3UtKuO6UEJDUUJdqdhxxIpCwvmrJEESauOk,893 +torch/testing/_internal/subclasses.py,sha256=8zo_8ospVkQFG66o92R4So-doE1pwmMzUrQtQV5TBNo,2717 +torch/testing/_internal/test_module/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/testing/_internal/test_module/__pycache__/__init__.cpython-311.pyc,, +torch/testing/_internal/test_module/__pycache__/future_div.cpython-311.pyc,, +torch/testing/_internal/test_module/__pycache__/no_future_div.cpython-311.pyc,, +torch/testing/_internal/test_module/future_div.py,sha256=298hLJlLz2QCJ80OVQXQM6CgN71nbkFMNnkesgzSnY8,114 +torch/testing/_internal/test_module/no_future_div.py,sha256=sksxzWFUupRBbSThS69P366IwnWBNGJe9EqQx8m5EKM,145 +torch/testing/_internal/torchbind_impls.py,sha256=8fCTOVfS50ZUumG6lMv0AtA-EQp8MCAFQDolhKkDuho,5655 +torch/testing/_internal/triton_utils.py,sha256=GhJ-drULNUX-RXuG_0V0NUhgCPCqFMEbuXvYYjz9OSA,34222 +torch/testing/_internal/two_tensor.py,sha256=OW6MrKArI4emVGfWA3OC6wVU8HAKDZlFyBkA33yiZso,4451 +torch/testing/_utils.py,sha256=qfU2CiPMoNiej2eD4hqMJDAxMT8tviWzxM0HJqxiQ9M,2519 +torch/torch_version.py,sha256=MheXhDr4r0O0mDO_Eaxz7vYLl0HRWr8hnQciPgx4QuI,2530 +torch/types.py,sha256=p9HIFrfGdCQDv7Ws3JwV-v-k1ogXEzrhudBV3zJfzWg,3642 +torch/utils/__init__.py,sha256=zMoIT_1VcmllcYuDNcUEfhrwT5iksaEIlf0mzinvnqM,3891 +torch/utils/__pycache__/__init__.cpython-311.pyc,, +torch/utils/__pycache__/_appending_byte_serializer.cpython-311.pyc,, +torch/utils/__pycache__/_config_module.cpython-311.pyc,, +torch/utils/__pycache__/_content_store.cpython-311.pyc,, +torch/utils/__pycache__/_contextlib.cpython-311.pyc,, +torch/utils/__pycache__/_cpp_embed_headers.cpython-311.pyc,, +torch/utils/__pycache__/_cpp_extension_versioner.cpython-311.pyc,, +torch/utils/__pycache__/_cxx_pytree.cpython-311.pyc,, +torch/utils/__pycache__/_device.cpython-311.pyc,, +torch/utils/__pycache__/_dtype_abbrs.cpython-311.pyc,, +torch/utils/__pycache__/_exposed_in.cpython-311.pyc,, +torch/utils/__pycache__/_filelock.cpython-311.pyc,, +torch/utils/__pycache__/_foreach_utils.cpython-311.pyc,, +torch/utils/__pycache__/_functools.cpython-311.pyc,, +torch/utils/__pycache__/_get_clean_triton.cpython-311.pyc,, +torch/utils/__pycache__/_helion.cpython-311.pyc,, +torch/utils/__pycache__/_import_utils.cpython-311.pyc,, +torch/utils/__pycache__/_inspect.cpython-311.pyc,, +torch/utils/__pycache__/_mode_utils.cpython-311.pyc,, +torch/utils/__pycache__/_ordered_set.cpython-311.pyc,, +torch/utils/__pycache__/_pallas.cpython-311.pyc,, +torch/utils/__pycache__/_python_dispatch.cpython-311.pyc,, +torch/utils/__pycache__/_pytree.cpython-311.pyc,, +torch/utils/__pycache__/_runtime_estimation.cpython-311.pyc,, +torch/utils/__pycache__/_stats.cpython-311.pyc,, +torch/utils/__pycache__/_thunk.cpython-311.pyc,, +torch/utils/__pycache__/_traceback.cpython-311.pyc,, +torch/utils/__pycache__/_triton.cpython-311.pyc,, +torch/utils/__pycache__/_typing_utils.cpython-311.pyc,, +torch/utils/__pycache__/_zip.cpython-311.pyc,, +torch/utils/__pycache__/backend_registration.cpython-311.pyc,, +torch/utils/__pycache__/bundled_inputs.cpython-311.pyc,, +torch/utils/__pycache__/checkpoint.cpython-311.pyc,, +torch/utils/__pycache__/collect_env.cpython-311.pyc,, +torch/utils/__pycache__/cpp_backtrace.cpython-311.pyc,, +torch/utils/__pycache__/cpp_extension.cpython-311.pyc,, +torch/utils/__pycache__/deterministic.cpython-311.pyc,, +torch/utils/__pycache__/dlpack.cpython-311.pyc,, +torch/utils/__pycache__/file_baton.cpython-311.pyc,, +torch/utils/__pycache__/flop_counter.cpython-311.pyc,, +torch/utils/__pycache__/hooks.cpython-311.pyc,, +torch/utils/__pycache__/mkldnn.cpython-311.pyc,, +torch/utils/__pycache__/mobile_optimizer.cpython-311.pyc,, +torch/utils/__pycache__/model_zoo.cpython-311.pyc,, +torch/utils/__pycache__/module_tracker.cpython-311.pyc,, +torch/utils/__pycache__/show_pickle.cpython-311.pyc,, +torch/utils/__pycache__/throughput_benchmark.cpython-311.pyc,, +torch/utils/__pycache__/weak.cpython-311.pyc,, +torch/utils/_appending_byte_serializer.py,sha256=bs9USsaM-2edEomwNuRARE3qpqM6yMvauob3_EChH4s,4064 +torch/utils/_config_module.py,sha256=k3Fnh6B2XPwapOMJkdrwgMCdF3HJmeFn6VbKnzQa5oQ,35007 +torch/utils/_config_typing.pyi,sha256=ZSNBjlnywGtVn4IxDkS0azK0t3q9Go71xETCA7VTN-E,1333 +torch/utils/_content_store.py,sha256=oe3EG2xt1HjHby1Qx-QXwN4ny004UhStbwtm-ImE79o,9327 +torch/utils/_contextlib.py,sha256=8iwhoDfEsmpDERITj6ZqUQjF6c_kyRaecywdFep6hAY,6402 +torch/utils/_cpp_embed_headers.py,sha256=JEXbg6pJrEuEyc6wTsDYVAJds2U66qZ15-Ik4W1tYdw,1758 +torch/utils/_cpp_extension_versioner.py,sha256=YHjYXX88iLDyZSuhNPJ3O4LS0LkNiqc2HMCW9GJmMSs,1946 +torch/utils/_cxx_pytree.py,sha256=s0mUyHE20HLOZnwiCmXA0TQfK18zBXFZdE6q7BqXMP8,39540 +torch/utils/_debug_mode/__init__.py,sha256=-UAxBbdHX814m-ssq_tP2v9y0zioXiN-d13nBqOnHEU,671 +torch/utils/_debug_mode/__pycache__/__init__.cpython-311.pyc,, +torch/utils/_debug_mode/__pycache__/_calls.cpython-311.pyc,, +torch/utils/_debug_mode/__pycache__/_mode.cpython-311.pyc,, +torch/utils/_debug_mode/__pycache__/_utils.cpython-311.pyc,, +torch/utils/_debug_mode/_calls.py,sha256=on8lxXNnam0qfxMRJejYuH_u3VxTNJnkm1ASULDCN1s,11234 +torch/utils/_debug_mode/_mode.py,sha256=_-oUbJtj6leUMtpAgUk92o4sIzmxKDXwos6owS4gEhs,40727 +torch/utils/_debug_mode/_utils.py,sha256=LTUfYLN6NfcTPHjIOhb36U4JJzEB8e8Q6TQRBhenwnY,8547 +torch/utils/_device.py,sha256=VxuBQa9VHC2KcLpQcXI_clqZpDELIKqwmEUQYsQgy1w,4491 +torch/utils/_dtype_abbrs.py,sha256=L-w4-UnjmD5QeuGbwRqilJ0iCc9Q1XsTjQUvygG0GE0,755 +torch/utils/_exposed_in.py,sha256=BdZbKFvTNdNkns_bTnsscKtH-o-OyJf5WOZS_mZlcXs,720 +torch/utils/_filelock.py,sha256=pzI0dRIiOSru8SdJ_IYI0d3FJqj2PZ0atCQoiGMgTTc,1530 +torch/utils/_foreach_utils.py,sha256=BWAnBP6D5TSRloDeNja7WfPM6Hz4fArfq7TEY_UmZUE,2409 +torch/utils/_functools.py,sha256=-eVphbmHHOzIph1xNgIGf9RtiCAcfENk8qMDajzhzcY,1527 +torch/utils/_get_clean_triton.py,sha256=pGrb9Ec40ZYcUhZLzIcdMiXJECinZp9P2VXvKHW_s34,6965 +torch/utils/_helion.py,sha256=JcmLdzWOPs6yfikYYZDIzWUSejX4gF598lAo9x2Sbcg,364 +torch/utils/_import_utils.py,sha256=WKkO4UjWrjJ-Kw_eeveZUMVsvcDc43YB5IxXf2Icy2o,1174 +torch/utils/_inspect.py,sha256=FFgrCL8pPZGhceoQ3rztoc3fSt4mqECQhjUykvERcgo,3073 +torch/utils/_mode_utils.py,sha256=sYjEYYBLwkP2g45Acigb32vqjNru3RgCJ-Efd1ol0A8,255 +torch/utils/_ordered_set.py,sha256=XCMvtn58N4_UbhnInekig5dDmwn2ZpltuuF3tIdFVxY,5619 +torch/utils/_pallas.py,sha256=AoI32P2NUpnKjndIjz4VafXP5C7SNO4kGo-DIAc0JzM,3340 +torch/utils/_python_dispatch.py,sha256=pGGpRuD_C8U0jbp_5fDo8FUn86dDbhiZeKrfSqPATG8,39506 +torch/utils/_pytree.py,sha256=IMQIV-vubPth9RPhjBtNphaXWDJd7Zn3X9m69s7xNP0,78685 +torch/utils/_runtime_estimation.py,sha256=vCDeombgBoMvReUn3PxKnOo6JbY1XdYjbj39-UIVUsU,4366 +torch/utils/_stats.py,sha256=6tFYMMxr0QGrEZg9qI1Whq-t1uizO3fneCjmjaWyxDA,1041 +torch/utils/_strobelight/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/utils/_strobelight/__pycache__/__init__.cpython-311.pyc,, +torch/utils/_strobelight/__pycache__/cli_function_profiler.cpython-311.pyc,, +torch/utils/_strobelight/cli_function_profiler.py,sha256=g7ZjL4WWrNcd_KtjWQMAyQE41haPAJ0_0TGOZL-vULg,11360 +torch/utils/_sympy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/utils/_sympy/__pycache__/__init__.cpython-311.pyc,, +torch/utils/_sympy/__pycache__/functions.cpython-311.pyc,, +torch/utils/_sympy/__pycache__/interp.cpython-311.pyc,, +torch/utils/_sympy/__pycache__/numbers.cpython-311.pyc,, +torch/utils/_sympy/__pycache__/printers.cpython-311.pyc,, +torch/utils/_sympy/__pycache__/reference.cpython-311.pyc,, +torch/utils/_sympy/__pycache__/singleton_int.cpython-311.pyc,, +torch/utils/_sympy/__pycache__/solve.cpython-311.pyc,, +torch/utils/_sympy/__pycache__/symbol.cpython-311.pyc,, +torch/utils/_sympy/__pycache__/value_ranges.cpython-311.pyc,, +torch/utils/_sympy/functions.py,sha256=fxoHfiT14YBPEWOZlW2bZFjBscYJs3x7vuYU3YUKpYo,54176 +torch/utils/_sympy/interp.py,sha256=Z4G1PT0oma4lwmNGw_mzyJq8XX_JqLsdG9UR_SSA70o,7200 +torch/utils/_sympy/numbers.py,sha256=ZYLHqe6LAl5kgztdANtt35vpB0mTaXQtAXSYPEJ2ONs,12049 +torch/utils/_sympy/printers.py,sha256=zzd2-laYlBCZ7dNWgKirXZK-k6juCljB3TPBxQPnhYY,28988 +torch/utils/_sympy/reference.py,sha256=mVpcw_8mf8gGroKqxhREZzanAEoXKgIIc0iGhUsALkQ,14625 +torch/utils/_sympy/singleton_int.py,sha256=K9fu2DlCBq5cXXQHzo1aJKbA11bxaIzuXptEZ96OP1A,2975 +torch/utils/_sympy/solve.py,sha256=TKC7vPm_YgXV3IPW4YqfwK1GT7RyQVNdhgX7_ojYpKM,6377 +torch/utils/_sympy/symbol.py,sha256=-rU79KpJLIfxT7jAXUlEp3M_OvXj83win7XHVJ3IMc8,3737 +torch/utils/_sympy/value_ranges.py,sha256=o0XioJIreXJwtC9pm80ULW4fPFeXYN06IerqoFqSIwA,39291 +torch/utils/_thunk.py,sha256=m5ck10KTsFbRFnH0qyv9_exIsG8ujIeLozlGLEhVC8Y,644 +torch/utils/_traceback.py,sha256=jLTvYPVDTxdTWs03vZGrDzTs_cVydzEPOt_-Nu664mc,10274 +torch/utils/_triton.py,sha256=u_YHtyb0ICbHSKWoycba4crfs2L_ntwiYQ2cdU0LNn0,6019 +torch/utils/_typing_utils.py,sha256=ez9L9GrSvYbXCC5EQp81cRC1gkUNlUfyrR5I6WtK8VQ,2501 +torch/utils/_zip.py,sha256=eBoEuB5X-ARtblldfoAU5JN-lADeWkRE59wtuFKiFPc,2509 +torch/utils/backcompat/__init__.py,sha256=R3Sq5nrFTb96s7Cs-gNFpXfjDtyVahc-dbqXaEVqUnA,678 +torch/utils/backcompat/__pycache__/__init__.cpython-311.pyc,, +torch/utils/backend_registration.py,sha256=gn1wRf18-IG2KxrLp1h1dT4Su41OSSYbwiroLKtkHQg,22031 +torch/utils/benchmark/__init__.py,sha256=VMZoFTt8YGaScu3-5L1uWwMOED8FRYoEBAX37RFJ3c0,411 +torch/utils/benchmark/__pycache__/__init__.cpython-311.pyc,, +torch/utils/benchmark/examples/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/utils/benchmark/examples/__pycache__/__init__.cpython-311.pyc,, +torch/utils/benchmark/examples/__pycache__/compare.cpython-311.pyc,, +torch/utils/benchmark/examples/__pycache__/fuzzer.cpython-311.pyc,, +torch/utils/benchmark/examples/__pycache__/op_benchmark.cpython-311.pyc,, +torch/utils/benchmark/examples/__pycache__/simple_timeit.cpython-311.pyc,, +torch/utils/benchmark/examples/__pycache__/spectral_ops_fuzz_test.cpython-311.pyc,, +torch/utils/benchmark/examples/compare.py,sha256=vCmmcvWlXuf3hLDdZprpNCeV4Dv-7Qn6fTHH9Mqna24,2931 +torch/utils/benchmark/examples/fuzzer.py,sha256=5e2AwlYhk30tRXL4pmkiP4a6wtiHtgm-6A5m5DRM-zg,2658 +torch/utils/benchmark/examples/op_benchmark.py,sha256=C9lRKNCPijSkKiPjo6-3lu_wVQwt-4oc4z91qAeTuLc,4412 +torch/utils/benchmark/examples/simple_timeit.py,sha256=ioYkChBgAPHwz5gWqTfUFTCiqqk6CRud3Lob1tFe9uM,541 +torch/utils/benchmark/examples/spectral_ops_fuzz_test.py,sha256=Z4r34jAqH9rrDVdnZBj8yMaHigcYrQ2N6iS0O3JDlPU,4787 +torch/utils/benchmark/op_fuzzers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/utils/benchmark/op_fuzzers/__pycache__/__init__.cpython-311.pyc,, +torch/utils/benchmark/op_fuzzers/__pycache__/binary.cpython-311.pyc,, +torch/utils/benchmark/op_fuzzers/__pycache__/sparse_binary.cpython-311.pyc,, +torch/utils/benchmark/op_fuzzers/__pycache__/sparse_unary.cpython-311.pyc,, +torch/utils/benchmark/op_fuzzers/__pycache__/spectral.cpython-311.pyc,, +torch/utils/benchmark/op_fuzzers/__pycache__/unary.cpython-311.pyc,, +torch/utils/benchmark/op_fuzzers/binary.py,sha256=_A--hlOm0zcmwm3JoB_4ggFfNYAHzjFtlK7a3MHmkbI,4144 +torch/utils/benchmark/op_fuzzers/sparse_binary.py,sha256=FfyrCWiRHqxAzK27kp_gqP-zb4V31aT3N9drKRD8URU,4226 +torch/utils/benchmark/op_fuzzers/sparse_unary.py,sha256=TB__9Vo1Q9n6Awp07zMbGj_0nmUb54LCP3YyBRjU_7w,3491 +torch/utils/benchmark/op_fuzzers/spectral.py,sha256=u3ydR5HcQ3DwN0N0IIX9cYRiWDFm32sEw8nCkXa2djU,3632 +torch/utils/benchmark/op_fuzzers/unary.py,sha256=xvApHM3-2cP997UF5gotcTDfUas2tUUkzeUByLgKWV8,3154 +torch/utils/benchmark/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/utils/benchmark/utils/__pycache__/__init__.cpython-311.pyc,, +torch/utils/benchmark/utils/__pycache__/_stubs.cpython-311.pyc,, +torch/utils/benchmark/utils/__pycache__/common.cpython-311.pyc,, +torch/utils/benchmark/utils/__pycache__/compare.cpython-311.pyc,, +torch/utils/benchmark/utils/__pycache__/compile.cpython-311.pyc,, +torch/utils/benchmark/utils/__pycache__/cpp_jit.cpython-311.pyc,, +torch/utils/benchmark/utils/__pycache__/fuzzer.cpython-311.pyc,, +torch/utils/benchmark/utils/__pycache__/sparse_fuzzer.cpython-311.pyc,, +torch/utils/benchmark/utils/__pycache__/timer.cpython-311.pyc,, +torch/utils/benchmark/utils/_stubs.py,sha256=u29cftbrzC0fc9gq-6o78yApgPHNMwBKoXqY4QEiQ1M,1026 +torch/utils/benchmark/utils/common.py,sha256=u7toTEFGVRe4DgX4juDUVNIuzpxUD-k2Y2--lkXA_Yg,13774 +torch/utils/benchmark/utils/compare.py,sha256=vGP9rkUnbQLIXDPtZXDaia5pBet7Y8ACLaidC11wQIE,13473 +torch/utils/benchmark/utils/compile.py,sha256=Wx3JjogowjvrpFbXL5SaFBs-HgLSPazi0HjlxYd8xZk,7675 +torch/utils/benchmark/utils/cpp_jit.py,sha256=9sTV4KXXjMy6-8gVCLzFk2EFRq_Qfg39v2MS-KOw3Fc,6976 +torch/utils/benchmark/utils/fuzzer.py,sha256=I1fpGS3SxZ6evdJIYtAUCDLYuZnHe6N9Y9f_NuH01Ps,18778 +torch/utils/benchmark/utils/sparse_fuzzer.py,sha256=Zfd1iQ5zp1BhgGGYVWNo92FEVaao7j2Xlh0U-bH8TVs,5213 +torch/utils/benchmark/utils/timeit_template.cpp,sha256=Wzz-o6Yjgq3tkmUTxeRldQxrsETh48T2hwL8xbOkRSg,1009 +torch/utils/benchmark/utils/timer.py,sha256=cJM5ENvXJrtngSzVDJ_XDqknJLCDJYU5wBIk8ylpnDg,20805 +torch/utils/benchmark/utils/valgrind_wrapper/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/utils/benchmark/utils/valgrind_wrapper/__pycache__/__init__.cpython-311.pyc,, +torch/utils/benchmark/utils/valgrind_wrapper/__pycache__/timer_interface.cpython-311.pyc,, +torch/utils/benchmark/utils/valgrind_wrapper/callgrind.h,sha256=wK1NVdRImF_4WVLlQrXkufunvE0qLr5dw50v86KPIo0,5744 +torch/utils/benchmark/utils/valgrind_wrapper/compat_bindings.cpp,sha256=ysuof0blt-4g76St1LUrdlZLiaCBKjwjwZXE0L4nI74,813 +torch/utils/benchmark/utils/valgrind_wrapper/timer_callgrind_template.cpp,sha256=ILVnffXHThuakEj3hzKv4usmt8rICdXpSTBeWKz7jU4,1676 +torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py,sha256=o9Fj5P4e3FKEvWy5LXzqDmpFTufLGVYYAFQpgtikiGs,37717 +torch/utils/benchmark/utils/valgrind_wrapper/valgrind.h,sha256=8MpV41sjwR0bIML04pxlLIjVuGhRtWdy1Kmtax4jFLI,422653 +torch/utils/bundled_inputs.py,sha256=f5W5KXT65LbYYiG25rMrMswLL5tPpbkMCpcL57O7evs,22701 +torch/utils/checkpoint.py,sha256=xSjFtN12lQ9DmE5Rr4DV_xKACGYcvU8a55MH0qmGVC0,73669 +torch/utils/collect_env.py,sha256=Bf_D9qgNyNBNIbdQx5xv6AngcDLOMXRhq7v7vBotKE0,31107 +torch/utils/cpp_backtrace.py,sha256=GxSqoJwxCh9nbgm7EJS9KRJ0Mn5UwIwJUje-0vC4C-Y,483 +torch/utils/cpp_extension.py,sha256=r2bPyw3DZpyKlJyU_fv56F8ayfmnkAKRReITCbNCjN8,142190 +torch/utils/data/__init__.py,sha256=atpjQDmxEyRtlfR9vDgGuJNxGQJq3H8g5drwHgkAK9A,1701 +torch/utils/data/__pycache__/__init__.cpython-311.pyc,, +torch/utils/data/__pycache__/backward_compatibility.cpython-311.pyc,, +torch/utils/data/__pycache__/dataloader.cpython-311.pyc,, +torch/utils/data/__pycache__/dataset.cpython-311.pyc,, +torch/utils/data/__pycache__/distributed.cpython-311.pyc,, +torch/utils/data/__pycache__/graph.cpython-311.pyc,, +torch/utils/data/__pycache__/graph_settings.cpython-311.pyc,, +torch/utils/data/__pycache__/sampler.cpython-311.pyc,, +torch/utils/data/_utils/__init__.py,sha256=9-7TzN_myxBjyvWvbQPPgZLmWAJ89O3AH1AKzyORONA,1606 +torch/utils/data/_utils/__pycache__/__init__.cpython-311.pyc,, +torch/utils/data/_utils/__pycache__/collate.cpython-311.pyc,, +torch/utils/data/_utils/__pycache__/fetch.cpython-311.pyc,, +torch/utils/data/_utils/__pycache__/pin_memory.cpython-311.pyc,, +torch/utils/data/_utils/__pycache__/signal_handling.cpython-311.pyc,, +torch/utils/data/_utils/__pycache__/worker.cpython-311.pyc,, +torch/utils/data/_utils/collate.py,sha256=_gv2PGubmzXO9I1ZJl3dZ8fMNiIGQ7yZwXGuJvwhlTg,15938 +torch/utils/data/_utils/fetch.py,sha256=3F1nYcG5DGCtTDd_LhHLnYOC4pzjz0w8BTkJiUy9E7U,2010 +torch/utils/data/_utils/pin_memory.py,sha256=DEJxMASS5TqKu19qIKYoo0EnbR9MLvSoAQTxTfR-2P8,4143 +torch/utils/data/_utils/signal_handling.py,sha256=4QRnZFkfMb3bgaptzCh24hgbwDITJ-32Tl5lgEZusNc,3261 +torch/utils/data/_utils/worker.py,sha256=OBA8-42rVcp9gbftkfmQqAxM7gxizGIvtqLrp5BIlHA,14829 +torch/utils/data/backward_compatibility.py,sha256=hEdTcktD0OuPmwt2eK_F4UkslR-n6D7wd14VSfG1DI4,317 +torch/utils/data/dataloader.py,sha256=cJidL2_9iRqKTMzDMTxsJA432p5QDiWG7vZW0mteyjQ,80463 +torch/utils/data/datapipes/__init__.py,sha256=OIYy5fRjiWdLc8SRrug0pOjrXouG0I-f2G8QQB0ynto,88 +torch/utils/data/datapipes/__pycache__/__init__.cpython-311.pyc,, +torch/utils/data/datapipes/__pycache__/_decorator.cpython-311.pyc,, +torch/utils/data/datapipes/__pycache__/_hook_iterator.cpython-311.pyc,, +torch/utils/data/datapipes/__pycache__/_typing.cpython-311.pyc,, +torch/utils/data/datapipes/__pycache__/datapipe.cpython-311.pyc,, +torch/utils/data/datapipes/__pycache__/gen_pyi.cpython-311.pyc,, +torch/utils/data/datapipes/_decorator.py,sha256=OhlEdmcjdqhVTfTc4FO2gIiJxQUMbcI-ug-5Nrc0Rig,7757 +torch/utils/data/datapipes/_hook_iterator.py,sha256=SciExxOz5XAb6PawByXG9MfuLdIN6XntMV_esMuqyU4,11972 +torch/utils/data/datapipes/_typing.py,sha256=GFWrRncLHo6KV65rIp2flGVntUIlz1foz3tawNRk-a0,16553 +torch/utils/data/datapipes/dataframe/__init__.py,sha256=j6gnc1vHYL9elliHkQsCRh9vaKbFRwHa0oMj6AIU2SY,378 +torch/utils/data/datapipes/dataframe/__pycache__/__init__.cpython-311.pyc,, +torch/utils/data/datapipes/dataframe/__pycache__/dataframe_wrapper.cpython-311.pyc,, +torch/utils/data/datapipes/dataframe/__pycache__/dataframes.cpython-311.pyc,, +torch/utils/data/datapipes/dataframe/__pycache__/datapipes.cpython-311.pyc,, +torch/utils/data/datapipes/dataframe/__pycache__/structures.cpython-311.pyc,, +torch/utils/data/datapipes/dataframe/dataframe_wrapper.py,sha256=9_V6cyBVauL0RpEehdWMCroAn1TGqKcM-M7ggGG7spg,3288 +torch/utils/data/datapipes/dataframe/dataframes.py,sha256=WdqoDNSjPnWV1RohBLKf9opZtCBSpw9Cp20_9XDMPDw,14034 +torch/utils/data/datapipes/dataframe/datapipes.py,sha256=yga6OgqN06IdQDPuCYd8T4tr-IzuyMttwYiIkZdM7t0,4653 +torch/utils/data/datapipes/dataframe/structures.py,sha256=7NKwwvZI2CtjpjvA1eJmyFQcbet7cYXdaxSMTFPZdx0,662 +torch/utils/data/datapipes/datapipe.py,sha256=VVX5I_Ulqt3NYtUsPNNROFNalRsflL1zVu8XjIhXE9g,17268 +torch/utils/data/datapipes/datapipe.pyi,sha256=kx0WXMNJNO74UvU-hs5iiaqijR-DucetLwH4As7sNNY,32730 +torch/utils/data/datapipes/gen_pyi.py,sha256=FR1brWWZ7IfnbsUblKApGwseZ51eORLfooAsf7lxEFE,11798 +torch/utils/data/datapipes/iter/__init__.py,sha256=Mh3NOCQIx3v7WrLqstLcFBuLDz4GRM-cSaBf-kJ6uFY,1862 +torch/utils/data/datapipes/iter/__pycache__/__init__.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/callable.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/combinatorics.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/combining.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/filelister.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/fileopener.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/grouping.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/routeddecoder.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/selecting.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/sharding.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/streamreader.cpython-311.pyc,, +torch/utils/data/datapipes/iter/__pycache__/utils.cpython-311.pyc,, +torch/utils/data/datapipes/iter/callable.py,sha256=vGa_hics1TsiXLsCkFS6gQFlC8mj8E_hA8UMRNh9j8c,9114 +torch/utils/data/datapipes/iter/combinatorics.py,sha256=ZmBXzNg1aVZFFoc_Gib_SkVyNnBzFtPNzRaZEyLB7l0,6557 +torch/utils/data/datapipes/iter/combining.py,sha256=Bp9lg-u3H9LZBaWCrXQOTx9aLVcZPCn9wtGhT14qnbs,28128 +torch/utils/data/datapipes/iter/filelister.py,sha256=pd59x2Ls-g1SpKBm4T5hgln0LHSqCVO8L82rfNVI-VY,2554 +torch/utils/data/datapipes/iter/fileopener.py,sha256=Qz0RAtlb8o2A9-YIMGcRVU2T6RztItmNCuqcGcvPQeA,2889 +torch/utils/data/datapipes/iter/grouping.py,sha256=KvqX3GwQoCEnvdy8_TAsqxLUaVjkBOPwVvjpP7346sA,12390 +torch/utils/data/datapipes/iter/routeddecoder.py,sha256=tdXrrUemDgjMWYxinV6bUeUFWxcujUc-UpqMVnHOW4o,2731 +torch/utils/data/datapipes/iter/selecting.py,sha256=LXpA0_L7TlTGkc5FlZpdzkxEaKqLUAEDoDFnAW3y4Ro,3308 +torch/utils/data/datapipes/iter/sharding.py,sha256=MRUQ2z-E4kn5cezM8PEHBGDu6y320a0IDyhiHp-pIJg,3657 +torch/utils/data/datapipes/iter/streamreader.py,sha256=TWVLV5IUepvqMWcNkEo3cBNbjbNYecXoXRpJ9GKfRGU,1537 +torch/utils/data/datapipes/iter/utils.py,sha256=y4CTzQ6oasBdi4xihoTYopUZ8_4P5eVoOCfk6s_aMSM,2109 +torch/utils/data/datapipes/map/__init__.py,sha256=6weaVqSBgL8DhjqKLXA2mpYjkh8djBYGEl8zwP4RXcY,714 +torch/utils/data/datapipes/map/__pycache__/__init__.cpython-311.pyc,, +torch/utils/data/datapipes/map/__pycache__/callable.cpython-311.pyc,, +torch/utils/data/datapipes/map/__pycache__/combinatorics.cpython-311.pyc,, +torch/utils/data/datapipes/map/__pycache__/combining.cpython-311.pyc,, +torch/utils/data/datapipes/map/__pycache__/grouping.cpython-311.pyc,, +torch/utils/data/datapipes/map/__pycache__/utils.cpython-311.pyc,, +torch/utils/data/datapipes/map/callable.py,sha256=lUC_SHHytKC2yJVi60rQdwxwrJGtxwoRCOkxIYio0k8,1933 +torch/utils/data/datapipes/map/combinatorics.py,sha256=7z15zvemzsOwnZvs9KaFs6WOW65nsj6V9eGHuQT3HOE,4267 +torch/utils/data/datapipes/map/combining.py,sha256=VwrqONcr9r5wGZqqhczQJRWmwFI2PKidi7T4wFg8zo0,3989 +torch/utils/data/datapipes/map/grouping.py,sha256=k7x6DBhQct1D869iIs7Ej3bzyzALO6mV-RsIl2ore2w,2531 +torch/utils/data/datapipes/map/utils.py,sha256=5oq8xhZHdJjrswR2Zzgef32cdSTMrgvvfrSC8f8EKLo,1813 +torch/utils/data/datapipes/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/utils/data/datapipes/utils/__pycache__/__init__.cpython-311.pyc,, +torch/utils/data/datapipes/utils/__pycache__/common.cpython-311.pyc,, +torch/utils/data/datapipes/utils/__pycache__/decoder.cpython-311.pyc,, +torch/utils/data/datapipes/utils/__pycache__/snapshot.cpython-311.pyc,, +torch/utils/data/datapipes/utils/common.py,sha256=Hf_J9RJpQaoahbDYqujoFZR3oVXbfbo73xYvlfpZEcM,13850 +torch/utils/data/datapipes/utils/decoder.py,sha256=BOfLfpyctXU8jDKKpvm0_tSOoXTpGdWfv_4iHwrbEH0,12121 +torch/utils/data/datapipes/utils/snapshot.py,sha256=tk8G2YDUdmEjIOMlLxuqOScYD3bxFyHUmyt_rvZA3Ck,3145 +torch/utils/data/dataset.py,sha256=M1m7rQ01kezp__7bzfaLM950RvLZpn88wF5F8T3vsUo,20952 +torch/utils/data/distributed.py,sha256=lYZvlTxNcfW5ysXW2vvqZBwU3nYMymmi3d45_3aNc7I,6451 +torch/utils/data/graph.py,sha256=duXC7TKPMyJ3sOJhZ9H6PK7RbvIs1Ab0B4ZZUjQ70D8,5918 +torch/utils/data/graph_settings.py,sha256=ImWDygTKAWFF9-wy0-RkgYh8_0l6HLiEJT3uhQK8jWM,5556 +torch/utils/data/sampler.py,sha256=gX-BikiX8MElaaqI7blRw9mRLMILSkfWK_87cMrfGsU,12738 +torch/utils/deterministic.py,sha256=aDwP89FjYTIT-3g2RtWBoagN59_JbqE2BZbgV4HttMU,611 +torch/utils/dlpack.py,sha256=xf66jbvrVfzzOypKLVw6Wgo1vD7MXq0K57cnJRoCbGI,8462 +torch/utils/file_baton.py,sha256=B67yGDzUCaDx6LZYmCXgcciL9TucXJrHZtz7JywZC50,2093 +torch/utils/flop_counter.py,sha256=bYXKB168sIttJkB8UWXo5VTOEkvN4QbgtsS0YFPxiyo,37814 +torch/utils/hipify/__init__.py,sha256=Jzb_RfgvXCrm_SQ4AfeGVi1N36YybxnM5mpyxrnihgI,33 +torch/utils/hipify/__pycache__/__init__.cpython-311.pyc,, +torch/utils/hipify/__pycache__/constants.cpython-311.pyc,, +torch/utils/hipify/__pycache__/cuda_to_hip_mappings.cpython-311.pyc,, +torch/utils/hipify/__pycache__/hipify_python.cpython-311.pyc,, +torch/utils/hipify/__pycache__/version.cpython-311.pyc,, +torch/utils/hipify/constants.py,sha256=rT3bQVod7V-F3Rc6fk9ix9ZEwZ8muEvb89OvRF9XGK8,1300 +torch/utils/hipify/cuda_to_hip_mappings.py,sha256=zTCq_nGEVU2YCXX-eGhjrwP-zgV91MsLiy_7oU_tuoo,204591 +torch/utils/hipify/hipify_python.py,sha256=KhSkwdgeHcw9N0g8w2rzRqTVGOr_P1-YRpWFo1XHoKY,46703 +torch/utils/hipify/version.py,sha256=hwOJuFUEKtsgiOt4l6op9s-hviovM6TK1uCdtRNWY4E,22 +torch/utils/hooks.py,sha256=v6WbXTcuC6b6PNw-Q6xG1mUeDonYmD57-9cRf-2rBH4,10165 +torch/utils/jit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/utils/jit/__pycache__/__init__.cpython-311.pyc,, +torch/utils/jit/__pycache__/log_extract.cpython-311.pyc,, +torch/utils/jit/log_extract.py,sha256=ow9gm6RPIXXwgGYOVx63uEreA31vzKY-pbaWTjfUR-E,4194 +torch/utils/mkldnn.py,sha256=OynlvKEDzAFEV5zLLqUKO61n-P_lExeBtuF_PEGW5qk,8244 +torch/utils/mobile_optimizer.py,sha256=lV9qq0E-O7M3AcVBnNdRJEYaQkEZwz7Tl7R6W1qo60s,6414 +torch/utils/model_dump/__init__.py,sha256=IZfRjacOoMdDjh4CnxvwHBY-cemXHVHm-w9bVCdCgJ8,19311 +torch/utils/model_dump/__main__.py,sha256=jYGPuoI11jzWgGKtK1-E550XLfPkrbHq_BRUbLusK-A,79 +torch/utils/model_dump/__pycache__/__init__.cpython-311.pyc,, +torch/utils/model_dump/__pycache__/__main__.cpython-311.pyc,, +torch/utils/model_dump/code.js,sha256=70w_JAT7N8dkWHrpQsA1enZCJK7VJOIy23ukbbbXQAg,19251 +torch/utils/model_dump/htm.mjs,sha256=m-psDFjVL3_BzcZnYkiRT8AIYhFhCY5ERrh_LP4sEH4,1230 +torch/utils/model_dump/preact.mjs,sha256=005yDhrtmGbeMNCyf7SMhNktZT4VisnDde1R-TMD5gk,10078 +torch/utils/model_dump/skeleton.html,sha256=vq4r1yFKZEchXwQmky1zpj3q65MCPDBD7wDt6SqzJGg,384 +torch/utils/model_zoo.py,sha256=o2NC-XaU8fqDXbUBQDv7thQBld_LOK-Ko63GS4x8Iyg,117 +torch/utils/module_tracker.py,sha256=pt9-CzmJYVRSg7U91RQCtU5rRBVGRmro0a9biTeKjCg,5434 +torch/utils/serialization/__init__.py,sha256=urhmq5QqBImvzUxNXw7VkCHllOx41xEyNenyZOWKYQQ,21 +torch/utils/serialization/__pycache__/__init__.cpython-311.pyc,, +torch/utils/serialization/__pycache__/config.cpython-311.pyc,, +torch/utils/serialization/config.py,sha256=01NhAREofDqt4RBG0NS2TRodKXUEyRwpmYw0dnBELdk,654 +torch/utils/show_pickle.py,sha256=gQRZpzpYj1zBCHMxehPmyqvC-9pIobj4-mhYH7cmkvo,5459 +torch/utils/tensorboard/__init__.py,sha256=ZYXvCXkbnXG7fJB_0aD2m1tdKIjo39LFzMqN7fku94w,480 +torch/utils/tensorboard/__pycache__/__init__.cpython-311.pyc,, +torch/utils/tensorboard/__pycache__/_convert_np.cpython-311.pyc,, +torch/utils/tensorboard/__pycache__/_embedding.cpython-311.pyc,, +torch/utils/tensorboard/__pycache__/_onnx_graph.cpython-311.pyc,, +torch/utils/tensorboard/__pycache__/_proto_graph.cpython-311.pyc,, +torch/utils/tensorboard/__pycache__/_pytorch_graph.cpython-311.pyc,, +torch/utils/tensorboard/__pycache__/_utils.cpython-311.pyc,, +torch/utils/tensorboard/__pycache__/summary.cpython-311.pyc,, +torch/utils/tensorboard/__pycache__/writer.cpython-311.pyc,, +torch/utils/tensorboard/_convert_np.py,sha256=Wcss2Qu8wlM9GHrn6m0AVXK8o8wyhsHdoGtxOjuZSPU,882 +torch/utils/tensorboard/_embedding.py,sha256=z6sP-bBDNOsAm-wJfa9GZ66pU9sJUzHKDiomNZ4BcOY,3281 +torch/utils/tensorboard/_onnx_graph.py,sha256=KMnPjhnuLWLp3jJEnaOhUUV-yIk_4Z9S4x8l2A_mko8,1886 +torch/utils/tensorboard/_proto_graph.py,sha256=WMZNB-Vxc_Px-QxHVJFXQW_GaOIqLR3NT4WdZisWIFQ,1964 +torch/utils/tensorboard/_pytorch_graph.py,sha256=4sY0kGhULwVkGwOKBEIz97GaQcO7IBrgEHnXUp6v9dQ,13843 +torch/utils/tensorboard/_utils.py,sha256=mDiLcRZC_ZBWCzinLD5tTd2Vc_G5tOx4qAjwunQKdfQ,4307 +torch/utils/tensorboard/summary.py,sha256=w7RpsrRwFE_PNpbSUKa-UJd15OHUSJXP5ZQ8k4VXOc0,34719 +torch/utils/tensorboard/writer.py,sha256=0P_p8HODAG9Mm7xlH9WqkyRFMlUVBxsYodZNP1M74n8,47343 +torch/utils/throughput_benchmark.py,sha256=zlDMe1eHSkg2aoXXUiRppGgYGb4IILIIPL2QHHxGxqI,6625 +torch/utils/viz/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torch/utils/viz/__pycache__/__init__.cpython-311.pyc,, +torch/utils/viz/__pycache__/_cycles.cpython-311.pyc,, +torch/utils/viz/_cycles.py,sha256=zHEnEeFxpLC6CPyEVrCuELSwhLeyRlUD3Ek51w_XvHo,16980 +torch/utils/weak.py,sha256=37kIF_DXnasHEwz_OvuTguZaeLFRckXV5Gkf1jz7m2M,12227 +torch/version.py,sha256=Vt-rnx-cOJsuhs81rsAKfbwHQHRuX-FhXDMW4nTKnfQ,311 +torch/xpu/__init__.py,sha256=eaYT1DpSBNCPq7tt0odBEHWF0uF1Iy1NtW5QOPmRYkY,21702 +torch/xpu/__pycache__/__init__.cpython-311.pyc,, +torch/xpu/__pycache__/_gpu_trace.cpython-311.pyc,, +torch/xpu/__pycache__/_utils.cpython-311.pyc,, +torch/xpu/__pycache__/graphs.cpython-311.pyc,, +torch/xpu/__pycache__/memory.cpython-311.pyc,, +torch/xpu/__pycache__/random.cpython-311.pyc,, +torch/xpu/__pycache__/streams.cpython-311.pyc,, +torch/xpu/_gpu_trace.py,sha256=xVF2Wike9f8uqlxtwDshiyTFbduMGq0HupoSLb47_70,2364 +torch/xpu/_utils.py,sha256=JTETOkI3J7tp4SFKV3ukRlvQDO-6Qg850KkJ8wO6gh0,1591 +torch/xpu/graphs.py,sha256=ApQ-sFToGMfblGyiB0csOwOcK-VvwITCBI1mRDJAGgI,22224 +torch/xpu/memory.py,sha256=MazVMWXsezQlQftHGg62cnDQ9CWi-ZmAbUKIjmdu71M,26192 +torch/xpu/random.py,sha256=dnrrJobYRQJd-6-JcuyoNmbkw7IsMsJM-GnigqsfnlY,5282 +torch/xpu/streams.py,sha256=iRsHzdjqjqSFmXd7ehELgR88yjwo0JSugBXfUptX7ts,6305 +torchgen/__init__.py,sha256=iirTpG38WcCsNMhEbi1dg7_jad6ptk_uzZ-BzaGBFyU,348 +torchgen/__pycache__/__init__.cpython-311.pyc,, +torchgen/__pycache__/code_template.cpython-311.pyc,, +torchgen/__pycache__/context.cpython-311.pyc,, +torchgen/__pycache__/gen.cpython-311.pyc,, +torchgen/__pycache__/gen_aoti_c_shim.cpython-311.pyc,, +torchgen/__pycache__/gen_backend_stubs.cpython-311.pyc,, +torchgen/__pycache__/gen_functionalization_type.cpython-311.pyc,, +torchgen/__pycache__/gen_lazy_tensor.cpython-311.pyc,, +torchgen/__pycache__/gen_schema_utils.cpython-311.pyc,, +torchgen/__pycache__/gen_vmap_plumbing.cpython-311.pyc,, +torchgen/__pycache__/local.cpython-311.pyc,, +torchgen/__pycache__/model.cpython-311.pyc,, +torchgen/__pycache__/native_function_generation.cpython-311.pyc,, +torchgen/__pycache__/utils.cpython-311.pyc,, +torchgen/__pycache__/yaml_utils.cpython-311.pyc,, +torchgen/aoti/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torchgen/aoti/__pycache__/__init__.cpython-311.pyc,, +torchgen/aoti/__pycache__/fallback_ops.cpython-311.pyc,, +torchgen/aoti/fallback_ops.py,sha256=yBxBb0ULVatAkRCR3x9hWvdIWHoFTE-7yY6EBEbD05c,8844 +torchgen/api/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torchgen/api/__pycache__/__init__.cpython-311.pyc,, +torchgen/api/__pycache__/autograd.cpython-311.pyc,, +torchgen/api/__pycache__/cpp.cpython-311.pyc,, +torchgen/api/__pycache__/dispatcher.cpython-311.pyc,, +torchgen/api/__pycache__/functionalization.cpython-311.pyc,, +torchgen/api/__pycache__/lazy.cpython-311.pyc,, +torchgen/api/__pycache__/meta.cpython-311.pyc,, +torchgen/api/__pycache__/native.cpython-311.pyc,, +torchgen/api/__pycache__/python.cpython-311.pyc,, +torchgen/api/__pycache__/structured.cpython-311.pyc,, +torchgen/api/__pycache__/translate.cpython-311.pyc,, +torchgen/api/__pycache__/ufunc.cpython-311.pyc,, +torchgen/api/__pycache__/unboxing.cpython-311.pyc,, +torchgen/api/autograd.py,sha256=y6vasrnA5-DSoZjrOJ2PY3yCb-GlsIWoU7wcO22IV-Y,39666 +torchgen/api/cpp.py,sha256=OUm9vRzwhgVZu83wro4qyiJB-OBdIT6xlQlKEM88JHw,16542 +torchgen/api/dispatcher.py,sha256=gPrg9tJ0ob1ErXWoGXJDGZTg0MbDCivTALKE7kjITU4,3479 +torchgen/api/functionalization.py,sha256=e5Jv65MTJiQ5LpYUU57Yf7mLRVmCVBhDS4kywZy1l2k,8349 +torchgen/api/lazy.py,sha256=ztV8Uhy6Ke5ArAVNNBREW7e045jaFMXGY8qgVHQ4DPU,17141 +torchgen/api/meta.py,sha256=zJYzviYI2gY9V9yPUyiZ_fShcdm_5LbigqGUH_WfaWw,483 +torchgen/api/native.py,sha256=xpye5SVepNMaHSIFzIfcrG-6BkDxCB3fPdQUPcilf5w,5205 +torchgen/api/python.py,sha256=fksUYTYPBT8U1Nu4cTnTKiw1lBwdncFnGPrm8AsNtwk,59907 +torchgen/api/structured.py,sha256=MuPt1MpFhHtZxXtUY5KaeccGY0Q_T5RHHXAFqgrdQVY,6115 +torchgen/api/translate.py,sha256=KBoD8K6_1e4I1ftVf2hDAoYri4zILxrs-3ZziTqYN2s,19297 +torchgen/api/types/__init__.py,sha256=bQ29sz_GNJGgqoDUsE6i_AcYZq81pOs1ATXQJQhUhLY,144 +torchgen/api/types/__pycache__/__init__.cpython-311.pyc,, +torchgen/api/types/__pycache__/signatures.cpython-311.pyc,, +torchgen/api/types/__pycache__/types.cpython-311.pyc,, +torchgen/api/types/__pycache__/types_base.cpython-311.pyc,, +torchgen/api/types/signatures.py,sha256=hNWM3qWCfyk635BOjyIJm6Fam-pOlr22zgHb_ltYh7k,12884 +torchgen/api/types/types.py,sha256=gRiCF9DO_x0ARN1faZ9iYx0IRYTzyCywAwhWXqYzcuo,6210 +torchgen/api/types/types_base.py,sha256=v8Ks3hCBZHDWZ-K7tVlNIowqV1wLKnBX-bgITDGzH7E,7171 +torchgen/api/ufunc.py,sha256=lkjsc9zBKYWFScj0dlV0J6fT67kz_doD7PDc2lhEENk,6810 +torchgen/api/unboxing.py,sha256=_JIvPzZF-LoRDkenskzx4JVTqiAit5-TTXk0F-pOjaU,9381 +torchgen/code_template.py,sha256=3WKyggNeuDRmFKwOiw9dTDJLCpjqhiudfAU2Nw7BGxc,3265 +torchgen/context.py,sha256=TESgCeRyBcLVD3ixrHShmeqQQKV9MSfkuiyXze_JH6o,4024 +torchgen/dest/__init__.py,sha256=qECRwrljRjK-kMdBqfc9X8JPVVPU6XlUuWNEdFD9u0w,805 +torchgen/dest/__pycache__/__init__.cpython-311.pyc,, +torchgen/dest/__pycache__/lazy_ir.cpython-311.pyc,, +torchgen/dest/__pycache__/lazy_ts_lowering.cpython-311.pyc,, +torchgen/dest/__pycache__/native_functions.cpython-311.pyc,, +torchgen/dest/__pycache__/register_dispatch_key.cpython-311.pyc,, +torchgen/dest/__pycache__/ufunc.cpython-311.pyc,, +torchgen/dest/lazy_ir.py,sha256=p5ql6rb5bDupsKodcnFGKsz_NKAVRbAOwahdPSIX-NQ,29693 +torchgen/dest/lazy_ts_lowering.py,sha256=9QUvL_Z-mGqRfhtK_X2RSEHlGmFOqDr_J8cYJrALUG4,1831 +torchgen/dest/native_functions.py,sha256=I4SK08v3bRwTfFJP7CKtqP3yqmR6Bxt1MqR8g5h_LPg,3171 +torchgen/dest/register_dispatch_key.py,sha256=KJ4aLf02TYMpYaq_aTGcAF1qnDUORO3E2rMe7U_knZM,42367 +torchgen/dest/ufunc.py,sha256=GyHbO3xBJYg9cGZhTfALevujG0RxMjbcnn4kDju8Tak,18204 +torchgen/gen.py,sha256=BeFPSkyDW80qZMuNxfbGhAtgUwTsczRz6G55yTMv0NM,116890 +torchgen/gen_aoti_c_shim.py,sha256=hmoxP8rvRnC8RZNO1du6wzg2MZPdwzAMT4OZI0_ab-s,28941 +torchgen/gen_backend_stubs.py,sha256=vYp7Jszke3ueM7RDpVl6EDkVpl0hxbEofCgfWhrOzQw,23340 +torchgen/gen_functionalization_type.py,sha256=xTxrqM0sY9Y7a6eXDUm2kU_UrUo2VLzg-w58AN2VVeY,50815 +torchgen/gen_lazy_tensor.py,sha256=GAzbJNqdBcYq5eQoTDOvEcYUXsIbwVy1FAO4pEfXZ0g,23205 +torchgen/gen_schema_utils.py,sha256=-wvIzm5ESeeBqpucXXbFetYjtdm2RJtQ-UXB5rSyW4Q,3350 +torchgen/gen_vmap_plumbing.py,sha256=jcAo7Kx9hX6DWD-9xrorY473o033QE1jxpDBxWOWxKo,9523 +torchgen/local.py,sha256=NPwuRTtAXVEfXOkinBU5trBAyXCa5xdZv2zjQBmmBYg,2227 +torchgen/model.py,sha256=3fwSao7nOxF28wL4sFhPfvp6kSj12fw014IAa_YNW5o,124032 +torchgen/native_function_generation.py,sha256=QF3dV4NTLrOgTjMdQSdGniO3f2NPciMFi_D2U0GJJ70,30568 +torchgen/operator_versions/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torchgen/operator_versions/__pycache__/__init__.cpython-311.pyc,, +torchgen/operator_versions/__pycache__/gen_mobile_upgraders.cpython-311.pyc,, +torchgen/operator_versions/__pycache__/gen_mobile_upgraders_constant.cpython-311.pyc,, +torchgen/operator_versions/gen_mobile_upgraders.py,sha256=C00LckTVFV839j3qTSnyElxe6p5nQOTSWh8YIjUOx68,12583 +torchgen/operator_versions/gen_mobile_upgraders_constant.py,sha256=C-U6rHQybm_FTcxsz27RMgJDj464NOLhlOzVjSjEn0w,243 +torchgen/packaged/ATen/native/native_functions.yaml,sha256=u1xT3tJszGJQpzrAVwEZgPmgyZwXLgk7KKPwKcVZyxE,625883 +torchgen/packaged/ATen/native/tags.yaml,sha256=TxFArqXG1qfFBUpixXqbVr3w2RLPUGoTnrae2i1wp3M,5820 +torchgen/packaged/ATen/templates/ATenOpList.cpp,sha256=YobnhIm91ECCc6uYD2uDOrvFM4WqutKQbQ5x_Fh_5IE,1059 +torchgen/packaged/ATen/templates/CompositeViewCopyKernels.cpp,sha256=H64AHoCBB7MJIECAHNzki8NiTVPND0hy2vJ-KqiSF2c,2077 +torchgen/packaged/ATen/templates/DispatchKeyFunction.h,sha256=npUU8WpU76sZv8oqUQqBpcV_QHT6RW9j42EVTSA6pvA,702 +torchgen/packaged/ATen/templates/DispatchKeyFunctions.h,sha256=KlLfLZSYEG_7miq0fD8yuFtgVtm2mO4CQ9WiT6xoEpY,1937 +torchgen/packaged/ATen/templates/DispatchKeyFunctions_inl.h,sha256=nlAU0xWHQqRZn7JNV163YgAv5wwlykwsikMyjzZeZkI,824 +torchgen/packaged/ATen/templates/DispatchKeyNativeFunctions.cpp,sha256=DYjxJmYQ5Yegq1XW4hakqjvZo5skNWi2f5ZNLPUyQ4c,184 +torchgen/packaged/ATen/templates/DispatchKeyNativeFunctions.h,sha256=e8lUNJZ4jt0uMHjio6HOupVpMv83DtAaFBudAG6pEDw,384 +torchgen/packaged/ATen/templates/Function.h,sha256=Z2OimxsHx42itqLZE04H7-RiAPb7UJRzNZQvmVUW6n4,519 +torchgen/packaged/ATen/templates/FunctionalInverses.h,sha256=ikZPokOnPXrImLa9EbaKGPMgB4NsyIH1cuUSgmLyejQ,797 +torchgen/packaged/ATen/templates/Functions.cpp,sha256=OBZ0Sa3RsXJGI0bBJlxY3FopVvT_76QP9vjWXm-URN8,3107 +torchgen/packaged/ATen/templates/Functions.h,sha256=W_xrKMuoXaRhohnPldbA0OVqkAi_NYid-pZ9zJQAa8U,4637 +torchgen/packaged/ATen/templates/LazyIr.h,sha256=-aVOIdjB719-6SLBsUURho6x0xdKyXsN6h5QP76HTd8,585 +torchgen/packaged/ATen/templates/LazyNonNativeIr.h,sha256=KQbjyZ0Q8qK8JcgqAaF-M-ZhvRE5UhTj8JIHyceNK9Y,178 +torchgen/packaged/ATen/templates/MethodOperators.h,sha256=ifirYleNPll8bjo_OYaC8jEuLWIJlP0Asy0xjtdGfQo,830 +torchgen/packaged/ATen/templates/NativeFunction.h,sha256=C8rQosVZF3aESwI4IvZitd60sdGov6kvClkDAtnx3Oc,366 +torchgen/packaged/ATen/templates/NativeFunctions.h,sha256=N351coNgM9F_jpdma5OK3DvlNUw2oY51iBXKe5dK81Y,1149 +torchgen/packaged/ATen/templates/NativeMetaFunction.h,sha256=C0alnIY6J-5Mlca9_0ocI7bnFKo5jQ_QTuLzo8iCUDQ,452 +torchgen/packaged/ATen/templates/NativeMetaFunctions.h,sha256=mIOwmpkQY9zYolxUXK39c4nRDyG32vBsqEeSDo76p_k,306 +torchgen/packaged/ATen/templates/Operator.h,sha256=5TqYc2DVGzB2gwuIwC8p4D6KnSDIylKXScnjwLfpdu0,448 +torchgen/packaged/ATen/templates/Operators.cpp,sha256=cjfjkIMtfc8n1w0TDy_JJDJq0DK2cT9DfhkZw3YnTWM,347 +torchgen/packaged/ATen/templates/Operators.h,sha256=oU939CI59Drfg2QlKfkCdU6yVQBVi0y9Ia_kQ0rFC5k,3200 +torchgen/packaged/ATen/templates/RedispatchFunctions.cpp,sha256=pNhfp3gMBw4km2c_4EfeF6ge3DGZi8xtGCnkmjzfpi0,307 +torchgen/packaged/ATen/templates/RedispatchFunctions.h,sha256=HFJ8SLBmg2LNRbsp9MdhHBGGztTGhcjzdWCgYPx5f7c,882 +torchgen/packaged/ATen/templates/RegisterBackendSelect.cpp,sha256=Op-RV7_8UnF_dbxO-hZ8X-7DBnFHsP9s11cM0JkJOWY,752 +torchgen/packaged/ATen/templates/RegisterCodegenUnboxedKernels.cpp,sha256=tux2wSt9RalK0o5AfliXJB3N4diTi-88PwNHOobDbmM,1119 +torchgen/packaged/ATen/templates/RegisterDispatchDefinitions.ini,sha256=hnJyQ-BCtiix4FaiwW1l3Zaq95M8C0lFTwbySsiWMDs,455 +torchgen/packaged/ATen/templates/RegisterDispatchKey.cpp,sha256=VBl-G3QyZ2UG-guGMf_2AEKFqq9zn-EAThFTjIUbh50,1428 +torchgen/packaged/ATen/templates/RegisterFunctionalization.cpp,sha256=XwSLhkDBdltLpwLt6tju8BXB0JLWK1MDPb01Gv3agCs,3461 +torchgen/packaged/ATen/templates/RegisterSchema.cpp,sha256=HoLTDNwRhe8xJJucgd6qCkqTglPtxLYCRGCD_4-S0j0,383 +torchgen/packaged/ATen/templates/RegistrationDeclarations.h,sha256=KImic_ILyhxavbGxVna-Ascf--okibalZJlK44a5dic,160 +torchgen/packaged/ATen/templates/TensorBody.h,sha256=g0bpNECPnvZPuwja9K1mvdpAD2o8OtKS11oNvYcgLe8,29327 +torchgen/packaged/ATen/templates/TensorMethods.cpp,sha256=ujhPuoXBOy23ZNJqFYqKbPwwLXoFK3ubnWK9jAfZ9w8,1824 +torchgen/packaged/ATen/templates/UfuncCPU.cpp,sha256=LrnISndBkXtdugvOWeRk9ZGYQlztFI5yqytoaZiKOQk,445 +torchgen/packaged/ATen/templates/UfuncCPUKernel.cpp,sha256=paz66F7U6E9e2X-rpbxlVDGcxevcXIOialaEqAaoArc,350 +torchgen/packaged/ATen/templates/UfuncCUDA.cu,sha256=HOBz8yO4QFxxmX_6gCF7L8MJvrGPwGQjoE-qDf8kF9Y,494 +torchgen/packaged/ATen/templates/UnboxingFunctions.cpp,sha256=wwdlYUaaCjXwhlaoqieeO-3fOqoQSBj62j6Is4n-UKY,709 +torchgen/packaged/ATen/templates/UnboxingFunctions.h,sha256=bcs4ET0LLtzs7nSSWhKA8jJzongib5CGlNA5yaExfKw,1026 +torchgen/packaged/ATen/templates/ViewMetaClasses.cpp,sha256=_8f0OdKZjfhK7LzODYmnmwaXjyie-LRnC-rVHxTv-ZM,346 +torchgen/packaged/ATen/templates/ViewMetaClasses.h,sha256=MhYaW9BQoS8iplXAz9IVu6_SY4JMG4eHP-ERqSXsdD0,233 +torchgen/packaged/ATen/templates/ViewMetaClassesPythonBinding.cpp,sha256=ATEHc06ioTG0d02quoCUFyBgaYR3FRV6SX-tkiLiFdY,292 +torchgen/packaged/ATen/templates/aten_interned_strings.h,sha256=_FM2jXAhATj9GZ66dXXUPx72q2uYnzT8iN4hkTz0rmI,805 +torchgen/packaged/ATen/templates/enum_tag.h,sha256=2f71l3hhsetytOYMpzhSywY-vpjF_uiHZgwZFJ8eG2I,389 +torchgen/packaged/autograd/BUILD.bazel,sha256=Jd76gG6LQlmmEKK9IYlDbkVSmc8L5bTvGjG159L3rJA,104 +torchgen/packaged/autograd/README.md,sha256=hGiUzBaCs0wzBhEXB3vkWUXU4Lima6y5_wPKjAoKQ-Q,147 +torchgen/packaged/autograd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torchgen/packaged/autograd/__pycache__/__init__.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/context.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/gen_annotated_fn_args.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/gen_autograd.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/gen_autograd_functions.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/gen_inplace_or_view_type.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/gen_python_functions.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/gen_trace_type.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/gen_variable_factories.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/gen_variable_type.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/gen_view_funcs.cpython-311.pyc,, +torchgen/packaged/autograd/__pycache__/load_derivatives.cpython-311.pyc,, +torchgen/packaged/autograd/build.bzl,sha256=mByyWKdnGG70dVc0E-YFaqRlPzikan1ZAJDtUNvfD_Y,491 +torchgen/packaged/autograd/context.py,sha256=aCykxnX--PJ_Tl4CW2GOYAbW25o26ZyirEobqc7ozOk,954 +torchgen/packaged/autograd/deprecated.yaml,sha256=UbtajzWo89jf1Q5sw8JNJIZT5s-iDBagnrGJthtDS2Q,6250 +torchgen/packaged/autograd/derivatives.yaml,sha256=dmujUIJuwH-7wnyEurd1HpdL4Tu491NbqHJa2d4afVc,185992 +torchgen/packaged/autograd/gen_annotated_fn_args.py,sha256=7SpCWSDgEj5pBYtWqsPAseC6Cfvyi_hUt-ST7BSF9hY,4476 +torchgen/packaged/autograd/gen_autograd.py,sha256=p6C9Pjmm1DyvB5QnFPqLs4EnWSR82cSpX975BipgJT8,4617 +torchgen/packaged/autograd/gen_autograd_functions.py,sha256=riiNmrDzNxm3-oEXVBImIwrsES90LGw_u8iHqj_dGzo,38452 +torchgen/packaged/autograd/gen_inplace_or_view_type.py,sha256=dmhrvrqn6j1bfGQGvKUbgRgbEYeDlfu5RC5nHhCP1-I,22914 +torchgen/packaged/autograd/gen_python_functions.py,sha256=bmNMEMmSEo9UphPlsDBM08I9-6NTNYbY2ifq3AZqFjU,46668 +torchgen/packaged/autograd/gen_trace_type.py,sha256=xBpXuGpepoTgpwwZNZUwYE2YRNUNROjxh6sRCoon_eY,19167 +torchgen/packaged/autograd/gen_variable_factories.py,sha256=a0Q5N1zGSVxfES_cnq88tjwOB9P79piRqchTrhWCNPk,4479 +torchgen/packaged/autograd/gen_variable_type.py,sha256=0eIQm5xkvp0x-5hhQpQ5-B604kmbtV_sEy4U_bwRrpI,87933 +torchgen/packaged/autograd/gen_view_funcs.py,sha256=LSNYEpzkNZXXFtM2pYpfCMw_QNAF5jm9RTOUMIFwJUM,11817 +torchgen/packaged/autograd/load_derivatives.py,sha256=dgj54kXO2bn0OT1unZxw0v0ArgfY6lciK4iQ5b_z7gw,41570 +torchgen/packaged/autograd/templates/ADInplaceOrViewType.cpp,sha256=6juXEdMJaxhS1nv6bqUmx-cJLWPyc4Yb3teAeTyEBG4,790 +torchgen/packaged/autograd/templates/Functions.cpp,sha256=j8Ga18okz6DI1ulIL5BErJIskHKQZPhuuNlj5RcgUuE,1478 +torchgen/packaged/autograd/templates/Functions.h,sha256=9vjETYlNoQKCJVRNBmIvBoj6T2GA9-c16T5FPVmWpa0,1577 +torchgen/packaged/autograd/templates/TraceType.cpp,sha256=aqACTyrT05ElIiBLJYkfSgxNjwhum6_QUPDAtXwoKqo,695 +torchgen/packaged/autograd/templates/VariableType.cpp,sha256=AG80s3CwhESp-zaewyfHmwmL0cOxOXZJcZxqyn2cIBs,2259 +torchgen/packaged/autograd/templates/VariableType.h,sha256=DVuAf8_kgQCcnZsVBFxSwwJTNZbFo01C5dVFhcROmq4,1467 +torchgen/packaged/autograd/templates/ViewFuncs.cpp,sha256=oas1Pw6wWyAfRy2uh5K7fGD3qR60LD2Oiv-B5OW6Tvc,269 +torchgen/packaged/autograd/templates/ViewFuncs.h,sha256=7RoIEE9NQ6jWdg3SfdbRsQxACKdscsrnZN45JTJER5I,498 +torchgen/packaged/autograd/templates/annotated_fn_args.py.in,sha256=gRgF9BZmylhyfXrVSAVjg9Y4TUYpq1FgYVd3_hk_9no,199 +torchgen/packaged/autograd/templates/python_enum_tag.cpp,sha256=2cTLq6vaU-qjAXLrChZAnCmpVqA4oDvZLOsBY3GWmDA,495 +torchgen/packaged/autograd/templates/python_fft_functions.cpp,sha256=g8Ub5jh_YhgeourCXjB2zlH15Lm_oGwLa37IjibMoEc,1951 +torchgen/packaged/autograd/templates/python_functions.cpp,sha256=b20LJcBbbQHQ4hTt-2h__DIuk8uHksD4ZVD5csfM43Q,1121 +torchgen/packaged/autograd/templates/python_functions.h,sha256=WX52FzntisoprdObQvGXrzMhehgWZ2rIfnHzgNpX-5U,345 +torchgen/packaged/autograd/templates/python_linalg_functions.cpp,sha256=3NyK1tMOYwAsnxyb1k04RPiN9MQkriblY3GnPVOrfSY,1614 +torchgen/packaged/autograd/templates/python_nested_functions.cpp,sha256=2wivRfIBY8pkU9mqHVQfqnsIff6K1Al-BTJyRFuUeAw,2029 +torchgen/packaged/autograd/templates/python_nn_functions.cpp,sha256=8AyIozHVM21KEaVuOPgT-zPfPGj6hpNMC5SiwqUixW0,3501 +torchgen/packaged/autograd/templates/python_return_types.cpp,sha256=GB75OiT-5X3rmgXZlc30MB67qMlvqkzhrDIfIj1mZp4,1219 +torchgen/packaged/autograd/templates/python_return_types.h,sha256=ZDLPH-bxSjCpeyWLhU6kZsuAxa2pnCgTE57YpxXmDmQ,198 +torchgen/packaged/autograd/templates/python_sparse_functions.cpp,sha256=NW_L2mF7ZrR2FvXtDE7iNfw0BM-wTyUzBa_C1B5RZ7A,1551 +torchgen/packaged/autograd/templates/python_special_functions.cpp,sha256=JqcAExUPCnbzUujY2PIDL5Gap-ZKlltKiZmDTleO_8s,1972 +torchgen/packaged/autograd/templates/python_torch_functions.cpp,sha256=QEgUPJbi5TunkkQog-o6zoIjpLFZ6VFHQYxKe7BUAiw,2601 +torchgen/packaged/autograd/templates/python_variable_methods.cpp,sha256=udmhz_G4hUGZjlUlEj70d1wZZ7b1z8e3TIQejK8tc3U,54000 +torchgen/packaged/autograd/templates/variable_factories.h,sha256=g-WOn2XuHBsLNGUk-emw0-KcrNsosDjqAu4NjPJCFA4,5637 +torchgen/selective_build/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torchgen/selective_build/__pycache__/__init__.cpython-311.pyc,, +torchgen/selective_build/__pycache__/operator.cpython-311.pyc,, +torchgen/selective_build/__pycache__/selector.cpython-311.pyc,, +torchgen/selective_build/operator.py,sha256=rA-TkO7eMUnHHysNngoP2fMKJcwRw1wvSeH7Vqq5av8,7209 +torchgen/selective_build/selector.py,sha256=GiA_q-Q_qwyc5s7MzUFkt3f8EwUtxMscYDWtpyo9Qyk,13813 +torchgen/static_runtime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +torchgen/static_runtime/__pycache__/__init__.cpython-311.pyc,, +torchgen/static_runtime/__pycache__/config.cpython-311.pyc,, +torchgen/static_runtime/__pycache__/gen_static_runtime_ops.cpython-311.pyc,, +torchgen/static_runtime/__pycache__/generator.cpython-311.pyc,, +torchgen/static_runtime/config.py,sha256=ZGckbYpoJpi9vmLxVu0yMnYUqNTGhRHWd_2DrBN8oR4,14546 +torchgen/static_runtime/gen_static_runtime_ops.py,sha256=eoYNdqEXbgEG5YMmYeKEMp8hAU9__jrmZ4j_n6DrDpA,7395 +torchgen/static_runtime/generator.py,sha256=pdiU4FSZATPTovtEbfzrXdk4GdkjLAZL2wzOwPCWoXs,29538 +torchgen/utils.py,sha256=R05E4aFW5ZlNRulAzqh0WWHGtB_mb7j9SZybSauizaU,16962 +torchgen/yaml_utils.py,sha256=cQ9aPicmhKYSNfgbTZc4hrg0GPKKxddzZt5FdMHEKEE,1116 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/WHEEL b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..3b96bae3ba84a6075769d2101f804bc5896460cb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: wheelfile 0.0.8 +Root-Is-Purelib: true +Tag: cp311-cp311-linux_x86_64 + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/entry_points.txt b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/entry_points.txt new file mode 100644 index 0000000000000000000000000000000000000000..5851bbc6ac14dfa78834f2dfa02553ce80777dc4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/entry_points.txt @@ -0,0 +1,5 @@ +[console_scripts] +torchrun = torch.distributed.run:main + +[torchrun.logs_specs] +default = torch.distributed.elastic.multiprocessing:DefaultLogsSpecs diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/licenses/LICENSE b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..c23172f7aff0254e4f0f163fb2e6e355cbaec5f4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/licenses/LICENSE @@ -0,0 +1,84 @@ +From PyTorch: + +Copyright (c) 2016- Facebook, Inc (Adam Paszke) +Copyright (c) 2014- Facebook, Inc (Soumith Chintala) +Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert) +Copyright (c) 2012-2014 Deepmind Technologies (Koray Kavukcuoglu) +Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu) +Copyright (c) 2011-2013 NYU (Clement Farabet) +Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston) +Copyright (c) 2006 Idiap Research Institute (Samy Bengio) +Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz) + +From Caffe2: + +Copyright (c) 2016-present, Facebook Inc. All rights reserved. + +All contributions by Facebook: +Copyright (c) 2016 Facebook Inc. + +All contributions by Google: +Copyright (c) 2015 Google Inc. +All rights reserved. + +All contributions by Yangqing Jia: +Copyright (c) 2015 Yangqing Jia +All rights reserved. + +All contributions by Kakao Brain: +Copyright 2019-2020 Kakao Brain + +All contributions by Cruise LLC: +Copyright (c) 2022 Cruise LLC. +All rights reserved. + +All contributions by Tri Dao: +Copyright (c) 2024 Tri Dao. +All rights reserved. + +All contributions by Arm: +Copyright (c) 2021, 2023-2025 Arm Limited and/or its affiliates + +All contributions from Caffe: +Copyright(c) 2013, 2014, 2015, the respective contributors +All rights reserved. + +All other contributions: +Copyright(c) 2015, 2016 the respective contributors +All rights reserved. + +Caffe2 uses a copyright model similar to Caffe: each contributor holds +copyright over their contributions to Caffe2. The project versioning records +all such contribution and copyright details. If a contributor wants to further +mark their specific copyright on a particular contribution, they should +indicate their copyright solely in the commit message of the change when it is +committed. + +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + +2. Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + +3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America + and IDIAP Research Institute nor the names of its contributors may be + used to endorse or promote products derived from this software without + specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE +LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR +CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF +SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS +INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN +CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) +ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE +POSSIBILITY OF SUCH DAMAGE. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/licenses/NOTICE b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/licenses/NOTICE new file mode 100644 index 0000000000000000000000000000000000000000..6effb8b5d70709f90835f2f5d646352fd77b6943 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/licenses/NOTICE @@ -0,0 +1,456 @@ +======================================================================= +Software under third_party +======================================================================= +Software libraries under third_party are provided as github submodule +links, and their content is not part of the Caffe2 codebase. Their +licences can be found under the respective software repositories. + +======================================================================= +Earlier BSD License +======================================================================= +Early development of Caffe2 in 2015 and early 2016 is licensed under the +BSD license. The license is attached below: + +All contributions by Facebook: +Copyright (c) 2016 Facebook Inc. + +All contributions by Google: +Copyright (c) 2015 Google Inc. +All rights reserved. + +All contributions by Yangqing Jia: +Copyright (c) 2015 Yangqing Jia +All rights reserved. + +All contributions by Kakao Brain: +Copyright 2019-2020 Kakao Brain + +All other contributions: +Copyright(c) 2015, 2016 the respective contributors +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +======================================================================= +Caffe's BSD License +======================================================================= +Some parts of the caffe2 code is derived from the original Caffe code, which is +created by Yangqing Jia and is now a BSD-licensed open-source project. The Caffe +license is as follows: + +COPYRIGHT + +All contributions by the University of California: +Copyright (c) 2014, The Regents of the University of California (Regents) +All rights reserved. + +All other contributions: +Copyright (c) 2014, the respective contributors +All rights reserved. + +Caffe uses a shared copyright model: each contributor holds copyright over +their contributions to Caffe. The project versioning records all such +contribution and copyright details. If a contributor wants to further mark +their specific copyright on a particular contribution, they should indicate +their copyright solely in the commit message of the change when it is +committed. + +LICENSE + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +1. Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. +2. Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR +ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +CONTRIBUTION AGREEMENT + +By contributing to the BVLC/caffe repository through pull-request, comment, +or otherwise, the contributor releases their content to the +license and copyright terms herein. + +======================================================================= +Caffe2's Apache License +======================================================================= + +This repo contains Caffe2 code, which was previously licensed under +Apache License Version 2.0: + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +======================================================================= +Cephes's 3-Clause BSD License +======================================================================= + +Code derived from implementations in the Cephes Math Library should mention +its derivation and reference the following license: + + 3-Clause BSD License for the Cephes Math Library + Copyright (c) 2018, Steven Moshier + All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions are met: + + * Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + * Redistributions in binary form must reproduce the above copyright + notice, this list of conditions and the following disclaimer in the + documentation and/or other materials provided with the distribution. + + * Neither the name of the nor the + names of its contributors may be used to endorse or promote products + derived from this software without specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND + ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED + WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE + DISCLAIMED. IN NO EVENT SHALL Steven Moshier BE LIABLE FOR ANY + DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES + (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; + LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND + ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + +======================================================================= +SciPy's 3-Clause BSD License +======================================================================= + +Code derived from implementations in SciPy should mention its derivation +and reference the following license: + + Copyright (c) 2001-2002 Enthought, Inc. 2003-2019, SciPy Developers. + All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions + are met: + + 1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + 2. Redistributions in binary form must reproduce the above + copyright notice, this list of conditions and the following + disclaimer in the documentation and/or other materials provided + with the distribution. + + 3. Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived + from this software without specific prior written permission. + + THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS + "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT + LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR + A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT + OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, + SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT + LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, + DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY + THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT + (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE + OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + +======================================================================= +Boost's 1.0 Software License +======================================================================= + +Code derived from implementations in Boost 1.0 should mention its +derivation and reference the following license: + + Boost Software License - Version 1.0 - August 17th, 2003 + + Permission is hereby granted, free of charge, to any person or organization + obtaining a copy of the software and accompanying documentation covered by + this license (the "Software") to use, reproduce, display, distribute, + execute, and transmit the Software, and to prepare derivative works of the + Software, and to permit third-parties to whom the Software is furnished to + do so, all subject to the following: + + The copyright notices in the Software and this entire statement, including + the above license grant, this restriction and the following disclaimer, + must be included in all copies of the Software, in whole or in part, and + all derivative works of the Software, unless such copies or derivative + works are solely in the form of machine-executable object code generated by + a source language processor. + + THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR + IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, + FITNESS FOR A PARTICULAR PURPOSE, TITLE AND NON-INFRINGEMENT. IN NO EVENT + SHALL THE COPYRIGHT HOLDERS OR ANYONE DISTRIBUTING THE SOFTWARE BE LIABLE + FOR ANY DAMAGES OR OTHER LIABILITY, WHETHER IN CONTRACT, TORT OR OTHERWISE, + ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER + DEALINGS IN THE SOFTWARE. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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 + + http://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. + +======================================================================= +PILLOW-SIMD Software License +======================================================================= + +Code derived from implementations in PILLOW-SIMD should mention its derivation +and reference the following license: + + The Python Imaging Library (PIL) is + + Copyright © 1997-2011 by Secret Labs AB + Copyright © 1995-2011 by Fredrik Lundh + + Pillow is the friendly PIL fork. It is + + Copyright © 2010-2022 by Alex Clark and contributors + + Like PIL, Pillow is licensed under the open source HPND License: + + By obtaining, using, and/or copying this software and/or its associated + documentation, you agree that you have read, understood, and will comply + with the following terms and conditions: + + Permission to use, copy, modify, and distribute this software and its + associated documentation for any purpose and without fee is hereby granted, + provided that the above copyright notice appears in all copies, and that + both that copyright notice and this permission notice appear in supporting + documentation, and that the name of Secret Labs AB or the author not be + used in advertising or publicity pertaining to distribution of the software + without specific, written prior permission. + + SECRET LABS AB AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH REGARD TO THIS + SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS. + IN NO EVENT SHALL SECRET LABS AB OR THE AUTHOR BE LIABLE FOR ANY SPECIAL, + INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM + LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE + OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR + PERFORMANCE OF THIS SOFTWARE. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/top_level.txt b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..3e40c4547cdefb9b810ff0c8f952fdead98563c6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch-2.12.1+computecanada.dist-info/top_level.txt @@ -0,0 +1,2 @@ +torch +torchgen diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/activation.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/activation.py new file mode 100644 index 0000000000000000000000000000000000000000..d85162ef35c7cd6a399f2a73a9a6b8f3c1154cd9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/activation.py @@ -0,0 +1,20 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.activation import ( + ELU, + Hardswish, + LeakyReLU, + MultiheadAttention, + PReLU, + ReLU6, + Sigmoid, + Softmax, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/batchnorm.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/batchnorm.py new file mode 100644 index 0000000000000000000000000000000000000000..8489cdb596ef44ce15d79530c80a9c7ea512e975 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/batchnorm.py @@ -0,0 +1,11 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.batchnorm import BatchNorm2d, BatchNorm3d diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/conv.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/conv.py new file mode 100644 index 0000000000000000000000000000000000000000..9c77b534ff6f6bd28466dda1c16ed219e48c1d73 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/conv.py @@ -0,0 +1,29 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.conv import ( + _reverse_repeat_padding, + Conv1d, + Conv2d, + Conv3d, + ConvTranspose1d, + ConvTranspose2d, + ConvTranspose3d, +) + + +__all__ = [ + "Conv1d", + "Conv2d", + "Conv3d", + "ConvTranspose1d", + "ConvTranspose2d", + "ConvTranspose3d", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/dropout.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/dropout.py new file mode 100644 index 0000000000000000000000000000000000000000..32a7a22d558670cc4ae9a963240badd314ed6d5c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/dropout.py @@ -0,0 +1,14 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.dropout import Dropout + + +__all__ = ["Dropout"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/embedding_ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/embedding_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..d25f8bea7e378023a8eb3ece75a5fb9a23163529 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/embedding_ops.py @@ -0,0 +1,18 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.embedding_ops import ( + Embedding, + EmbeddingBag, + EmbeddingPackedParams, +) + + +__all__ = ["EmbeddingPackedParams", "Embedding", "EmbeddingBag"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/functional_modules.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/functional_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..efe1b38ce3ea4adbae55595d86c2787d7c1f7284 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/functional_modules.py @@ -0,0 +1,18 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.functional_modules import ( + FloatFunctional, + FXFloatFunctional, + QFunctional, +) + + +__all__ = ["FloatFunctional", "FXFloatFunctional", "QFunctional"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/linear.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/linear.py new file mode 100644 index 0000000000000000000000000000000000000000..e9ba5a5c12f82915db53d81a7b9e5a1c0e530e98 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/linear.py @@ -0,0 +1,14 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.linear import Linear, LinearPackedParams + + +__all__ = ["LinearPackedParams", "Linear"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/normalization.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/normalization.py new file mode 100644 index 0000000000000000000000000000000000000000..85462cc365344b004c91ff9c02879477d50041f5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/normalization.py @@ -0,0 +1,26 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.normalization import ( + GroupNorm, + InstanceNorm1d, + InstanceNorm2d, + InstanceNorm3d, + LayerNorm, +) + + +__all__ = [ + "LayerNorm", + "GroupNorm", + "InstanceNorm1d", + "InstanceNorm2d", + "InstanceNorm3d", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/rnn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/rnn.py new file mode 100644 index 0000000000000000000000000000000000000000..a1a0076d13bc4e3ee29e9b3e410171d20e8e9a65 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/rnn.py @@ -0,0 +1,11 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.rnn import LSTM diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ea333af04ca49138a3b3ed35020654d4dad5ffe9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/quantized/modules/utils.py @@ -0,0 +1,17 @@ +# flake8: noqa: F401 +r"""Quantized Modules. + +This file is in the process of migration to `torch/ao/nn/quantized`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate file under the `torch/ao/nn/quantized/modules`, +while adding an import statement here. +""" + +from torch.ao.nn.quantized.modules.utils import ( + _hide_packed_params_repr, + _ntuple_from_first, + _pair_from_first, + _quantize_weight, + WeightedQuantizedModule, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..84145da93f7bd1a38590e38345184be9eaa4d7bf --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/__init__.py @@ -0,0 +1,48 @@ +from . import parametrizations, parametrize, rnn, stateless +from .clip_grad import ( + _clip_grads_with_norm_ as clip_grads_with_norm_, + _get_total_norm as get_total_norm, + clip_grad_norm, + clip_grad_norm_, + clip_grad_value_, +) +from .convert_parameters import parameters_to_vector, vector_to_parameters +from .fusion import ( + fuse_conv_bn_eval, + fuse_conv_bn_weights, + fuse_linear_bn_eval, + fuse_linear_bn_weights, +) +from .init import skip_init +from .memory_format import ( + convert_conv2d_weight_memory_format, + convert_conv3d_weight_memory_format, +) +from .spectral_norm import remove_spectral_norm, spectral_norm +from .weight_norm import remove_weight_norm, weight_norm + + +__all__ = [ + "clip_grad_norm", + "clip_grad_norm_", + "clip_grads_with_norm_", + "clip_grad_value_", + "convert_conv2d_weight_memory_format", + "convert_conv3d_weight_memory_format", + "fuse_conv_bn_eval", + "fuse_conv_bn_weights", + "fuse_linear_bn_eval", + "fuse_linear_bn_weights", + "get_total_norm", + "parameters_to_vector", + "parametrizations", + "parametrize", + "remove_spectral_norm", + "remove_weight_norm", + "rnn", + "skip_init", + "spectral_norm", + "stateless", + "vector_to_parameters", + "weight_norm", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_deprecation_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_deprecation_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a25b647307900e42b11d1cdafc8d9f8785d1a620 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_deprecation_utils.py @@ -0,0 +1,53 @@ +import importlib +import warnings +from collections.abc import Callable + + +_MESSAGE_TEMPLATE = ( + r"Usage of '{old_location}' is deprecated; please use '{new_location}' instead." +) + + +def lazy_deprecated_import( + all: list[str], + old_module: str, + new_module: str, +) -> Callable: + r"""Import utility to lazily import deprecated packages / modules / functional. + + The old_module and new_module are also used in the deprecation warning defined + by the `_MESSAGE_TEMPLATE`. + + Args: + all: The list of the functions that are imported. Generally, the module's + __all__ list of the module. + old_module: Old module location + new_module: New module location / Migrated location + + Returns: + Callable to assign to the `__getattr__` + + Usage: + + # In the `torch/nn/quantized/functional.py` + from torch.nn.utils._deprecation_utils import lazy_deprecated_import + _MIGRATED_TO = "torch.ao.nn.quantized.functional" + __getattr__ = lazy_deprecated_import( + all=__all__, + old_module=__name__, + new_module=_MIGRATED_TO) + """ + warning_message = _MESSAGE_TEMPLATE.format( + old_location=old_module, new_location=new_module + ) + + def getattr_dunder(name: str) -> None: + if name in all: + # We are using the "RuntimeWarning" to make sure it is not + # ignored by default. + warnings.warn(warning_message, RuntimeWarning, stacklevel=2) + package = importlib.import_module(new_module) + return getattr(package, name) + raise AttributeError(f"Module {new_module!r} has no attribute {name!r}.") + + return getattr_dunder diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8a0eaf86bdbeacfe7d4e7cbd50daf11385955d7d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/__init__.py @@ -0,0 +1,10 @@ +from .conv_expanded_weights import ConvPerSampleGrad +from .embedding_expanded_weights import EmbeddingPerSampleGrad +from .expanded_weights_impl import ExpandedWeight +from .group_norm_expanded_weights import GroupNormPerSampleGrad +from .instance_norm_expanded_weights import InstanceNormPerSampleGrad +from .layer_norm_expanded_weights import LayerNormPerSampleGrad +from .linear_expanded_weights import LinearPerSampleGrad + + +__all__ = ["ExpandedWeight"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/conv_expanded_weights.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/conv_expanded_weights.py new file mode 100644 index 0000000000000000000000000000000000000000..da7d8f3dfabb8c82f01d76c27ffbe1d3974473cf --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/conv_expanded_weights.py @@ -0,0 +1,82 @@ +from collections.abc import Callable +from typing import Any, TypeVar +from typing_extensions import ParamSpec + +import torch +import torch.nn.functional as F + + +_P = ParamSpec("_P") +_R = TypeVar("_R") + +from .conv_utils import ( + conv_args_and_kwargs, + conv_backward, + conv_input_for_string_padding, + conv_picker, +) +from .expanded_weights_impl import ExpandedWeight, implements_per_sample_grads +from .expanded_weights_utils import forward_helper + + +@implements_per_sample_grads(F.conv1d) +@implements_per_sample_grads(F.conv2d) +@implements_per_sample_grads(F.conv3d) +class ConvPerSampleGrad(torch.autograd.Function): + @staticmethod + # pyrefly: ignore [bad-override] + def forward( + ctx: Any, + kwarg_names: list[str], + conv_fn: Callable[_P, _R], + *expanded_args_and_kwargs: Any, + ) -> torch.Tensor: + expanded_args, expanded_kwargs = conv_args_and_kwargs( + kwarg_names, expanded_args_and_kwargs + ) + orig_input = expanded_args[0] + was_same_padding = expanded_kwargs["padding"] == "same" + + if isinstance(expanded_kwargs["padding"], str): + # if padding is a string, we'll do the necessary padding (slowly) using F.pad + kernel_size = expanded_args[1].shape[2:] + padding, dilation = expanded_kwargs["padding"], expanded_kwargs["dilation"] + input = conv_input_for_string_padding( + conv_fn, padding, expanded_args[0], dilation, kernel_size + ) + expanded_args = (input, expanded_args[1]) + # since we've already done the padding, don't need any more + expanded_kwargs["padding"] = 0 + + output = forward_helper(conv_fn, expanded_args, expanded_kwargs) + input, weight = expanded_args + batched_dim_size = conv_picker(conv_fn, 3, 4, 5) + if input.dim() != batched_dim_size: + raise RuntimeError( + f"Expanded Weights only support convolution with batched input, got {conv_fn} with an" + f"unbatched input of dim {input.dim()}, expected input of dim {batched_dim_size}" + ) + + # pyrefly: ignore [invalid-type-var] + ctx.conv_fn = conv_fn + + ctx.batch_size = orig_input.shape[0] + ctx.input_required_grad = orig_input.requires_grad + ctx.orig_input_shape = orig_input.shape + ctx.was_same_padding = was_same_padding + ctx.stride, ctx.padding = expanded_kwargs["stride"], expanded_kwargs["padding"] + ctx.dilation, ctx.groups = ( + expanded_kwargs["dilation"], + expanded_kwargs["groups"], + ) + + if isinstance(weight, ExpandedWeight): + ctx.input = input + ctx.weight = weight + ctx.bias = expanded_kwargs["bias"] + + return output + + @staticmethod + def backward(ctx: Any, *grad_outputs: Any) -> Any: + return conv_backward(ctx.conv_fn, ctx, grad_outputs[0]) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/conv_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/conv_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a9cf2c5c780571d1a5fdb8fe90ede2c107c649cd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/conv_utils.py @@ -0,0 +1,355 @@ +# mypy: allow-untyped-defs + +import math + +import torch +import torch.nn.functional as F + +from .expanded_weights_utils import ( + set_grad_sample_if_exists, + unpack_expanded_weight_or_tensor, +) + + +THRESHOLD = 32 + + +def conv_picker(func, conv1dOpt, conv2dOpt, conv3dOpt): + if func is F.conv1d: + return conv1dOpt + if func is F.conv2d: + return conv2dOpt + else: + if func is not F.conv3d: + raise AssertionError( + f"Expected func to be F.conv1d, F.conv2d, or F.conv3d, got {func}" + ) + return conv3dOpt + + +def conv_args_and_kwargs(kwarg_names, expanded_args_and_kwargs): + args = expanded_args_and_kwargs[: len(expanded_args_and_kwargs) - len(kwarg_names)] + kwargs = expanded_args_and_kwargs[ + len(expanded_args_and_kwargs) - len(kwarg_names) : + ] + kwargs = dict(zip(kwarg_names, kwargs, strict=True)) + + return conv_normalizer(*args, **kwargs) + + +def conv_normalizer( + input, + weight, + bias=None, + stride=1, + padding=0, + dilation=1, + groups=1, +): + return (input, weight), { + "bias": bias, + "stride": stride, + "padding": padding, + "dilation": dilation, + "groups": groups, + } + + +def conv_input_for_string_padding(func, padding_style, input, dilation, kernel_size): + if padding_style == "valid": + return input + else: + padding = int_padding_for_string_padding( + func, padding_style, dilation, kernel_size + ) + return F.pad(input, padding) + + +def int_padding_for_string_padding(func, padding_style, dilation, kernel_size): + def get_dilation(i): + return dilation[i] if isinstance(dilation, tuple) else dilation + + if padding_style == "same": + padding: list[int] = [] + # F.pad needs the padding in reverse order from what conv expects + for i in range(conv_picker(func, 0, 1, 2), -1, -1): + padding += conv_padding_for_same(get_dilation(i), kernel_size[i]) + return padding + elif padding_style == "valid": + return conv_picker(func, 2, 4, 6) * (0,) + else: + raise RuntimeError( + f"got padding type of {padding_style}, only accept 'same' or 'valid'" + ) + + +def conv_padding_for_same(dilation, kernel_size): + total_pad = dilation * (kernel_size - 1) + left_pad = total_pad // 2 + right_pad = total_pad - left_pad + return left_pad, right_pad + + +def conv_backward(func, ctx, grad_output): + def weight_grad_sample(weight): + if batch_size < THRESHOLD and groups == 1: + return conv_group_weight_grad_sample( + ctx.input, + grad_output, + weight_shape, + stride, + padding, + dilation, + batch_size, + func, + ) + else: + return conv_unfold_weight_grad_sample( + ctx.input, + grad_output, + weight_shape, + kernel_size, + stride, + padding, + dilation, + groups, + func, + ) + + def expand(param): + if isinstance(param, int): + return conv_picker(func, (param,), (param, param), (param, param, param)) + else: + return param + + def calc_total_padding(func, was_same, padding, dilation, kernel_size): + if was_same: + all_padding = int_padding_for_string_padding( + func, "same", dilation, kernel_size + ) + # F.pad needs the padding in reverse order from what conv expects + total_padding = tuple( + all_padding[i] + all_padding[i - 1] + for i in range(len(all_padding) - 1, -1, -2) + ) + return total_padding + else: + return tuple(2 * pad for pad in padding) + + weight_shape = ctx.weight.shape + stride, padding, dilation, groups = ( + expand(ctx.stride), + expand(ctx.padding), + expand(ctx.dilation), + ctx.groups, + ) + + kernel_size = [weight_shape[i] for i in range(2, conv_picker(func, 3, 4, 5))] + + batch_size = ctx.batch_size + results: list[torch.Tensor | None] = [] + results.append(None) # for kwarg names + results.append(None) # for op reference + + # "same" padding may give uneven padding on either side so we need to separate the "padding" attr and total padding + total_padding = calc_total_padding( + func, ctx.was_same_padding, padding, dilation, kernel_size + ) + + if ctx.input_required_grad: + output_padding = [] + input_dims = conv_picker(func, 1, 2, 3) + for i in range(input_dims): + input_dim = ctx.orig_input_shape[2 + i] + output_padding.append( + ( + total_padding[i] + + input_dim + - (kernel_size[i] * dilation[i] - dilation[i] + 1) + ) + % stride[i] + ) + weight_ = unpack_expanded_weight_or_tensor(ctx.weight) + transpose_func = conv_picker( + func, F.conv_transpose1d, F.conv_transpose2d, F.conv_transpose3d + ) + out = transpose_func( + grad_output, + weight_, + None, + stride, + padding, + tuple(output_padding), + groups, + dilation, + ) + + if ctx.was_same_padding: + for i in range(len(total_padding)): + out = torch.narrow( + out, 2 + i, total_padding[i] // 2, ctx.orig_input_shape[2 + i] + ) + + results.append(out) + else: + results.append(None) + # weight and bias don't compute batched gradients; no other arguments are differentiable + results = results + [None] * 6 + + # set grad_sample field for weight and bias with per sample gradients + set_grad_sample_if_exists(ctx.weight, weight_grad_sample) + set_grad_sample_if_exists( + ctx.bias, lambda _: grad_output.reshape(*grad_output.shape[:2], -1).sum(dim=2) + ) + return tuple(results) + + +def conv_unfold_weight_grad_sample( + input, + grad_output, + weight_shape, + kernel_size, + stride, + padding, + dilation, + groups, + func, +): + n = input.shape[0] + in_channels = input.shape[1] + + unfold_func = conv_picker( + func, + lambda: F.unfold( + input.unsqueeze(-2), + kernel_size=(1, kernel_size[0]), + dilation=(1, dilation[0]), + padding=(0, padding[0]), + stride=(1, stride[0]), + ), + lambda: F.unfold( + input, kernel_size, dilation=dilation, padding=padding, stride=stride + ), + lambda: unfold3d(input, kernel_size, padding, stride, dilation), + ) + + input = unfold_func() + grad_output = grad_output.reshape(n, -1, input.shape[-1]) + + # n=batch_sz; o=num_out_channels; p=(num_in_channels/groups)*kernel_sz + weight_grad_sample = torch.einsum("noq,npq->nop", grad_output, input) + # rearrange the above tensor and extract diagonals. + + weight_grad_sample = weight_grad_sample.view( + n, + groups, + -1, + groups, + int(in_channels / groups), + math.prod(kernel_size), + ) + weight_grad_sample = torch.einsum( + "ngrg...->ngr...", weight_grad_sample + ).contiguous() + shape = [n] + list(weight_shape) + weight_grad_sample = weight_grad_sample.view(shape) + return weight_grad_sample + + +def conv_group_weight_grad_sample( + input, + grad_output, + weight_shape, + stride, + padding, + dilation, + batch_size, + func, +): + I = input.shape[1] + O = grad_output.shape[1] + + input_ = input.transpose(0, 1) + grad_output_ = grad_output.view( + grad_output.shape[0] * grad_output.shape[1], 1, *grad_output.shape[2:] + ) + + weight_grad_sample = func( + input_, + grad_output_, + None, + stride=dilation, + padding=padding, + dilation=stride, + groups=batch_size, + ) + input_dims = conv_picker(func, 3, 4, 5) + for i in range(2, input_dims): + weight_grad_sample = weight_grad_sample.narrow(i, 0, weight_shape[i]) + weight_grad_sample = weight_grad_sample.view( + I, batch_size, O, *weight_grad_sample.shape[2:] + ) + weight_grad_sample = weight_grad_sample.movedim(0, 2) + return weight_grad_sample + + +def unfold3d( + tensor, + kernel_size, + padding, + stride, + dilation, +): + r""" + Extract sliding local blocks from an batched input tensor. + + :class:`torch.nn.Unfold` only supports 4D inputs (batched image-like tensors). + This method implements the same action for 5D inputs + Args: + tensor: An input tensor of shape ``(B, C, D, H, W)``. + kernel_size: the size of the sliding blocks + padding: implicit zero padding to be added on both sides of input + stride: the stride of the sliding blocks in the input spatial dimensions + dilation: the spacing between the kernel points. + Returns: + A tensor of shape ``(B, C * math.prod(kernel_size), L)``, where L - output spatial dimensions. + See :class:`torch.nn.Unfold` for more details + Example: + >>> # xdoctest: +SKIP + >>> B, C, D, H, W = 3, 4, 5, 6, 7 + >>> tensor = torch.arange(1, B * C * D * H * W + 1.0).view(B, C, D, H, W) + >>> unfold3d(tensor, kernel_size=2, padding=0, stride=1).shape + torch.Size([3, 32, 120]) + """ + + if len(tensor.shape) != 5: + raise ValueError( + f"Input tensor must be of the shape [B, C, D, H, W]. Got{tensor.shape}" + ) + + if dilation != (1, 1, 1): + raise NotImplementedError(f"dilation={dilation} not supported.") + + batch_size, channels, _, _, _ = tensor.shape + + # Input shape: (B, C, D, H, W) + tensor = F.pad( + tensor, (padding[2], padding[2], padding[1], padding[1], padding[0], padding[0]) + ) + # Output shape: (B, C, D+2*padding[2], H+2*padding[1], W+2*padding[0]) + + tensor = tensor.unfold(dimension=2, size=kernel_size[0], step=stride[0]) + tensor = tensor.unfold(dimension=3, size=kernel_size[1], step=stride[1]) + tensor = tensor.unfold(dimension=4, size=kernel_size[2], step=stride[2]) + # Output shape: (B, C, D_out, H_out, W_out, kernel_size[0], kernel_size[1], kernel_size[2]) + # For D_out, H_out, W_out definitions see :class:`torch.nn.Unfold` + + tensor = tensor.permute(0, 2, 3, 4, 1, 5, 6, 7) + # Output shape: (B, D_out, H_out, W_out, C, kernel_size[0], kernel_size[1], kernel_size[2]) + + tensor = tensor.reshape( + batch_size, -1, channels * math.prod(kernel_size) + ).transpose(1, 2) + # Output shape: (B, D_out * H_out * W_out, C * kernel_size[0] * kernel_size[1] * kernel_size[2] + + return tensor diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/embedding_expanded_weights.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/embedding_expanded_weights.py new file mode 100644 index 0000000000000000000000000000000000000000..74350b88b5407a01dac92270f8471cc0e37a99c5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/embedding_expanded_weights.py @@ -0,0 +1,88 @@ +from typing import Any + +import torch +import torch.nn.functional as F + +from .expanded_weights_impl import implements_per_sample_grads +from .expanded_weights_utils import ( + forward_helper, + set_grad_sample_if_exists, + standard_kwargs, +) + + +@implements_per_sample_grads(F.embedding) +class EmbeddingPerSampleGrad(torch.autograd.Function): + @staticmethod + # pyrefly: ignore [bad-override] + def forward( + ctx: Any, kwarg_names: list[str], _: Any, *expanded_args_and_kwargs: Any + ) -> torch.Tensor: + expanded_args, expanded_kwargs = standard_kwargs( + kwarg_names, expanded_args_and_kwargs + ) + if len(expanded_args[0].shape) == 1: + raise RuntimeError( + f"Expanded Weights needs an input with a batch size, got a 1D tensor, {expanded_args[0]}" + ) + output = forward_helper(F.embedding, expanded_args, expanded_kwargs) + ctx.input, ctx.weight = expanded_args + ctx.padding_idx, ctx.scale_grad_by_freq = ( + expanded_kwargs["padding_idx"], + expanded_kwargs["scale_grad_by_freq"], + ) + ctx.sparse = expanded_kwargs["sparse"] + return output + + @staticmethod + # pyrefly: ignore [bad-override] + def backward( + ctx: Any, grad_output: torch.Tensor + ) -> tuple[torch.Tensor | None, ...]: + input, weight = ctx.input, ctx.weight + padding_idx, scale_grad_by_freq, sparse = ( + ctx.padding_idx, + ctx.scale_grad_by_freq, + ctx.sparse, + ) + + def weight_per_sample_grad(weight: torch.Tensor) -> torch.Tensor: + batch_size = input.shape[0] + embedding_dim = weight.shape[1] + index = ( + input.unsqueeze(-1) + .expand(*input.shape, embedding_dim) + .reshape(batch_size, -1, embedding_dim) + ) + grad_sample = torch.zeros( # type: ignore[attr-defined] + batch_size, *weight.shape, device=weight.device, dtype=grad_output.dtype + ) + return grad_sample.scatter_add_( + 1, index, grad_output.reshape(batch_size, -1, embedding_dim) + ) + + results: list[torch.Tensor | None] = [] + results.append(None) # for kwarg names + results.append(None) # for op reference + + if input.requires_grad: + bw_fn = torch.ops.aten.embedding_backward + results.append( + bw_fn( + grad_output, + input, + weight.shape[0], + padding_idx, + scale_grad_by_freq, + sparse, + ) + ) + else: + results.append(None) + + # weight doesn't compute batched gradients; no other arguments are differentiable (2 not saved from forward) + results = results + [None] * 6 + + # set grad_sample field for weight with per sample gradients + set_grad_sample_if_exists(weight, weight_per_sample_grad) + return tuple(results) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/expanded_weights_impl.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/expanded_weights_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..f3a5999f392c3a28386b0f1a3d191b5a36fa3e7d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/expanded_weights_impl.py @@ -0,0 +1,186 @@ +# mypy: allow-untyped-defs +import functools +from collections.abc import Callable +from contextlib import contextmanager + +import torch +from torch._decomp import decomposition_table +from torch.utils._pytree import tree_map_only + + +HANDLED_FUNCTIONS: dict[Callable, torch.autograd.Function] = {} + +aten = torch._ops.ops.aten +# __torch_function__ runs before the pydispatcher so we need to manually use the same +# decompositions indexed by their torch equivalent +expanded_weights_rnn_decomps = { + # func: (input_decomp, data_decomp) + torch.rnn_relu: ( + decomposition_table[aten.rnn_relu.input], + decomposition_table[aten.rnn_relu.data], + ), + torch.rnn_tanh: ( + decomposition_table[aten.rnn_tanh.input], + decomposition_table[aten.rnn_tanh.data], + ), + torch.lstm: ( + decomposition_table[aten.lstm.input], + decomposition_table[aten.lstm.data], + ), + torch.gru: ( + decomposition_table[aten.gru.input], + decomposition_table[aten.gru.data], + ), +} + + +# all of the RNN decomps run linear with the batch dimension second, even if batch_first was set +@contextmanager +def batch_second(args, kwargs): + def set_batch_second(ew) -> None: + ew.set_batch_first(False) + + def reset_batch_first(ew) -> None: + ew.set_batch_first(True) + + tree_map_only(ExpandedWeight, set_batch_second, args) + tree_map_only(ExpandedWeight, set_batch_second, kwargs) + try: + yield + finally: + tree_map_only(ExpandedWeight, reset_batch_first, args) + tree_map_only(ExpandedWeight, reset_batch_first, kwargs) + + +# to support packed sequences, we need to allow for smaller batches. Expanded weights represents the largest batch +@contextmanager +def allow_smaller_batches(args, kwargs): + def allow(ew) -> None: + ew.set_allow_smaller_batches(True) + + def reset(ew) -> None: + ew.set_allow_smaller_batches(False) + + tree_map_only(ExpandedWeight, allow, args) + tree_map_only(ExpandedWeight, allow, kwargs) + try: + yield + finally: + tree_map_only(ExpandedWeight, reset, args) + tree_map_only(ExpandedWeight, reset, kwargs) + + +@contextmanager +def setup_rnn(use_input_variant, args, kwargs): + with ( + batch_second(args, kwargs) + if use_input_variant + else allow_smaller_batches(args, kwargs) + ): + yield + + +def implements_per_sample_grads(torch_function): + @functools.wraps(torch_function) + def decorator(autograd_func): + HANDLED_FUNCTIONS[torch_function] = autograd_func + return autograd_func + + return decorator + + +# ExpandedWeight represents a weight (parameter) Tensor that has an expanded +# batch dimension. Operations on the ExpandedWeight Tensor act exactly like +# those without an expanded batch dimension but a call to .backward() populates +# the original (unexpanded) tensor with per-sample-gradients for in the grad_sample field +# +# ExpandedWeight has a fallback that always fails since we cannot know what the batch +# dimension of the input tensor is and therefore cannot know if this is a valid call +# +# This is a __torch_function__ object but it could have also been a Tensor Extension +# with a dispatch key. +# +# Needs to be a tensor subclass to allow reparameterization +class ExpandedWeight(torch.Tensor): + def __init__(self, orig_weight, batch_size, loss_reduction) -> None: + self.batch_size = batch_size + self.batch_first = True + self.allow_smaller_batches = False + self.orig_weight = orig_weight + self.loss_reduction = loss_reduction + + handled_functions = HANDLED_FUNCTIONS + + def __new__(cls, orig_weight, batch_size, loss_reduction): + if not isinstance(orig_weight, torch.Tensor): + raise RuntimeError( + f"Can only make Expanded Weights of Tensors, got {type(orig_weight).__name__}" + ) + if not orig_weight.requires_grad: + raise RuntimeError( + "Can only build ExpandedWeights objects of tensors that require_grad" + ) + ret = torch.Tensor._make_subclass(cls, orig_weight, True) + return ret + + @classmethod + def __torch_function__(cls, func, _, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + if func in expanded_weights_rnn_decomps: + # in aten, choosing the input or data variants is done by parsing logic. This mimics some of that + decomp_opts = expanded_weights_rnn_decomps[func] + use_input_variant = isinstance( + # pyrefly: ignore [bad-index] + args[2], + list, + ) # data variant uses a list here + decomp = decomp_opts[0] if use_input_variant else decomp_opts[1] + + if decomp is not None: + with setup_rnn(use_input_variant, args, kwargs): + return decomp(*args, **kwargs) + if func is torch._cudnn_rnn_flatten_weight: + # since we aren't using the fused cuda kernels for RNNs, don't do this + return + if func in cls.handled_functions: + return cls.handled_functions[func].apply( + tuple(kwargs.keys()), func, *(args + tuple(kwargs.values())) + ) + # We cannot use a fallback here because we do not know the batch dimension for any regular tensor inputs, + # i.e. torch.add(torch.Tensor, ExpandedWeight) + raise RuntimeError( + f"Expanded Weights encountered but cannot handle function {func.__name__}" + ) + + @property + def dtype(self): # type: ignore[override] + return self.orig_weight.dtype + + @property + def data(self): # type: ignore[override] + return self.orig_weight.data + + @property + def shape(self): # type: ignore[override] + return self.orig_weight.shape + + @property + def device(self): # type: ignore[override] + return self.orig_weight.device + + @property + def is_cuda(self): # type: ignore[override] + return self.orig_weight.is_cuda + + def data_ptr(self): + return self.orig_weight.data_ptr() + + def get_device(self): + return self.orig_weight.get_device() + + def set_allow_smaller_batches(self, is_allow_smaller_batches) -> None: + self.allow_smaller_batches = is_allow_smaller_batches + + def set_batch_first(self, is_batch_first=True) -> None: + self.batch_first = is_batch_first diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/expanded_weights_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/expanded_weights_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..97b9729d8a5d64d523dc24d4ea06d9d298f73bd9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/expanded_weights_utils.py @@ -0,0 +1,188 @@ +# mypy: allow-untyped-defs + +import torch + +from .expanded_weights_impl import ExpandedWeight + + +def is_batch_first(expanded_args_and_kwargs): + batch_first = None + # pyrefly: ignore [bad-assignment] + for arg in expanded_args_and_kwargs: + if not isinstance(arg, ExpandedWeight): + continue + + if not batch_first: + batch_first = arg.batch_first + elif arg.batch_first != batch_first: + raise RuntimeError( + "Got conflicting batch_first arguments in the same layer" + ) + return batch_first + + +def standard_kwargs(kwarg_names, expanded_args): + r"""Separate args and kwargs from `__torch_function__`s that standardize kwargs. + + Most `__torch_function__`s standardize the kwargs that they give, so this will separate + the args and kwargs they pass. Functions that don't are linear and convAND. + """ + kwarg_values = expanded_args[len(expanded_args) - len(kwarg_names) :] + expanded_args_without_kwargs = expanded_args[ + : len(expanded_args) - len(kwarg_names) + ] + expanded_kwargs = dict(zip(kwarg_names, kwarg_values, strict=True)) + return expanded_args_without_kwargs, expanded_kwargs + + +def forward_helper(func, expanded_args, expanded_kwargs): + r"""Compute the forward pass for a function that has expanded weight(s) passed to it. + + It will run the forward pass where all ExpandedWeights are their original + weight. It runs checks on the given arguments and detaches the outputs. + + .. note:: First argument in :attr:`expanded_args` must be the input with the batch + dimension as the first element of the shape + + .. note:: :attr:`func` must return a Tensor or tuple of Tensors + + Args: + func: The function to be called + expanded_args: Arguments to be passed to :attr:`func`. Will include arguments + that need to be unpacked because they are ExpandedWeights + expanded_kwargs: Keyword arguments to be passed to :attr:`func`. + Similar to :attr:`expanded_args`. + """ + unexpanded_args, unexpanded_kwargs = _check_and_unexpand_args( + func, expanded_args, expanded_kwargs + ) + return func(*unexpanded_args, **unexpanded_kwargs) + + +def _check_and_unexpand_args(func, expanded_args, expanded_kwargs): + # input must be the first argument passed + input = expanded_args[0] + if isinstance(input, ExpandedWeight): + raise RuntimeError( + "Expanded Weights do not support inputs that are also ExpandedWeights. " + f"Input must be a Tensor, got {type(input).__name__} in function {func.__name__}" + ) + if not isinstance(input, torch.Tensor): + raise RuntimeError( + "Expanded Weights requires a Tensor as the first input to get the batch dimension, " + f"got {type(input).__name__} in function {func.__name__}" + ) + if len(input.shape) == 0: + raise RuntimeError( + f"Expanded Weights requires a batch dimension but got an input of size 0 in function {func.__name__}" + ) + if input.shape[0] == 0: + raise RuntimeError( + "0 is not a valid batch size for Expanded Weights but got input tensor of " + f"{input} in function {func.__name__}" + ) + for arg in expanded_args + tuple(expanded_kwargs.values()): + if not isinstance(arg, ExpandedWeight): + continue + batch_size = input.shape[0] if arg.batch_first else input.shape[1] + if (arg.allow_smaller_batches and batch_size > arg.batch_size) or ( + not arg.allow_smaller_batches and arg.batch_size != batch_size + ): + raise RuntimeError( + "Expected ExpandedWeights to have batch size matching input but got " + f"input batch size of {batch_size} with ExpandedWeight of batch size {arg.batch_size}" + ) + + loss_reduction: str | None = None + for arg in expanded_args + tuple(expanded_kwargs.values()): + if isinstance(arg, ExpandedWeight): + if loss_reduction is None: + loss_reduction = arg.loss_reduction + elif loss_reduction != arg.loss_reduction: + raise RuntimeError( + "Expected ExpandedWeights to all have the same loss_reduction argument but got one" + f"with {loss_reduction} and one with {arg.loss_reduction}" + ) + + unexpanded_args = tuple( + arg.orig_weight if isinstance(arg, ExpandedWeight) else arg + for arg in expanded_args + ) + unexpanded_kwargs = { + name: arg.orig_weight if isinstance(arg, ExpandedWeight) else arg + for (name, arg) in expanded_kwargs.items() + } + return unexpanded_args, unexpanded_kwargs + + +def maybe_scale_by_batch_size(grad_sample, expanded_weight): + if expanded_weight.loss_reduction == "mean": + return grad_sample * expanded_weight.batch_size + else: + return grad_sample + + +def set_grad_sample_if_exists(maybe_expanded_weight, per_sample_grad_fn) -> None: + unpacked = unpack_expanded_weight_or_tensor(maybe_expanded_weight) + if isinstance(maybe_expanded_weight, ExpandedWeight): + grad_sample_contribution = maybe_scale_by_batch_size( + per_sample_grad_fn(unpacked), maybe_expanded_weight + ) + + if maybe_expanded_weight.batch_size > grad_sample_contribution.shape[0]: + # this only passes the other checks if the arg allows smaller batch sizes + intermediate = torch.zeros( + maybe_expanded_weight.batch_size, + *grad_sample_contribution.shape[1:], + dtype=grad_sample_contribution.dtype, + device=grad_sample_contribution.device, + ) + intermediate[: grad_sample_contribution.shape[0]] = grad_sample_contribution + grad_sample_contribution = intermediate + + if hasattr(unpacked, "grad_sample") and unpacked.grad_sample is not None: + unpacked.grad_sample = unpacked.grad_sample + grad_sample_contribution + else: + unpacked.grad_sample = grad_sample_contribution + + +def unpack_expanded_weight_or_tensor(maybe_expanded_weight, func=lambda x: x): + if isinstance(maybe_expanded_weight, ExpandedWeight): + orig_weight = maybe_expanded_weight.orig_weight + return func(orig_weight) + elif ( + isinstance(maybe_expanded_weight, torch.Tensor) + and not maybe_expanded_weight.requires_grad + ): + return func(maybe_expanded_weight) + elif isinstance(maybe_expanded_weight, torch.Tensor): + raise RuntimeError( + "ExpandedWeights currently does not support a mixture of ExpandedWeight parameters " + "and normal Parameters. Please file and issue with pytorch/pytorch" + ) + + +def sum_over_all_but_batch_and_last_n( + tensor: torch.Tensor, + n_dims: int, +) -> torch.Tensor: + r""" + Calculate the sum over all dimensions, except the first (batch dimension), and excluding the last n_dims. + + This function will ignore the first dimension and it will + not aggregate over the last n_dims dimensions. + Args: + tensor: An input tensor of shape ``(B, ..., X[n_dims-1])``. + n_dims: Number of dimensions to keep. + Example: + >>> tensor = torch.ones(1, 2, 3, 4, 5) + >>> sum_over_all_but_batch_and_last_n(tensor, n_dims=2).shape + torch.Size([1, 4, 5]) + Returns: + A tensor of shape ``(B, ..., X[n_dims-1])`` + """ + if tensor.dim() == n_dims + 1: + return tensor + else: + dims = list(range(1, tensor.dim() - n_dims)) + return tensor.sum(dim=dims) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/group_norm_expanded_weights.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/group_norm_expanded_weights.py new file mode 100644 index 0000000000000000000000000000000000000000..704560792dd49d513d233fef429d605fbc49ac42 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/group_norm_expanded_weights.py @@ -0,0 +1,106 @@ +# mypy: allow-untyped-defs +import operator +from functools import reduce + +import torch +import torch.nn.functional as F + +from .expanded_weights_impl import ExpandedWeight, implements_per_sample_grads +from .expanded_weights_utils import ( + forward_helper, + set_grad_sample_if_exists, + standard_kwargs, + unpack_expanded_weight_or_tensor, +) + + +@implements_per_sample_grads(F.group_norm) +class GroupNormPerSampleGrad(torch.autograd.Function): + @staticmethod + # pyrefly: ignore [bad-override] + def forward(ctx, kwarg_names, _, *expanded_args_and_kwargs): + expanded_args, expanded_kwargs = standard_kwargs( + kwarg_names, expanded_args_and_kwargs + ) + input, num_groups = expanded_args + N = input.shape[0] + C = input.shape[1] + HxW = reduce(operator.mul, input.shape[2:], 1) + weight, bias, eps = ( + expanded_kwargs["weight"], + expanded_kwargs["bias"], + expanded_kwargs["eps"], + ) + output, mean, rstd = forward_helper( + torch.native_group_norm, + (input, weight, bias, N, C, HxW, num_groups, eps), + {}, + ) + ctx.input, ctx.num_groups = input, num_groups + ctx.weight, ctx.eps = weight, eps + ctx.mean, ctx.rstd = mean, rstd + if isinstance(bias, ExpandedWeight): + ctx.bias = bias + if input.requires_grad and isinstance(weight, ExpandedWeight): + ctx.weight = weight + return output + + @staticmethod + # pyrefly: ignore [bad-override] + def backward(ctx, grad_output): + input, num_groups = ctx.input, ctx.num_groups + weight, bias, eps = ctx.weight, ctx.bias, ctx.eps + mean, rstd = ctx.mean, ctx.rstd + + results: list[torch.Tensor | None] = [] + results.append(None) # for kwarg names + results.append(None) # for op reference + + if input.requires_grad: + weight_c = unpack_expanded_weight_or_tensor( + weight, lambda t: t.contiguous() + ) + input_c = input.contiguous() + grad_output_c = ( + grad_output.contiguous() if grad_output is not None else None + ) + N = input.shape[0] + C = input.shape[1] + HxW = 1 + for s in input.shape[2:]: + HxW *= s + bw_fn = torch.ops.aten.native_group_norm_backward + results.append( + bw_fn( + grad_output_c, + input_c, + mean, + rstd, + weight_c, + N, + C, + HxW, + num_groups, + (True, False, False), + )[0] + ) + else: + results.append(None) + + # weight and bias don't compute batched gradients; no other arguments are differentiable + results = results + [None] * 4 + + # set grad_sample field for weight and bias with per sample gradients + if hasattr(ctx, "weight"): + set_grad_sample_if_exists( + weight, + lambda _: torch.einsum( + "ni...->ni", + F.group_norm(input, num_groups, eps=eps) * grad_output, + ), + ) + if hasattr(ctx, "bias"): + set_grad_sample_if_exists( + bias, lambda _: torch.einsum("ni...->ni", grad_output) + ) + return tuple(results) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/instance_norm_expanded_weights.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/instance_norm_expanded_weights.py new file mode 100644 index 0000000000000000000000000000000000000000..8a5237cb4e32472204679618b9800b4b85062cc1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/instance_norm_expanded_weights.py @@ -0,0 +1,101 @@ +# mypy: allow-untyped-defs +from functools import partial + +import torch +import torch.nn.functional as F + +from .expanded_weights_impl import implements_per_sample_grads +from .expanded_weights_utils import ( + forward_helper, + set_grad_sample_if_exists, + standard_kwargs, + unpack_expanded_weight_or_tensor, +) + + +@implements_per_sample_grads(F.instance_norm) +class InstanceNormPerSampleGrad(torch.autograd.Function): + @staticmethod + # pyrefly: ignore [bad-override] + def forward(ctx, kwarg_names, _, *expanded_args_and_kwargs): + instance_norm = partial(torch.instance_norm, cudnn_enabled=True) + expanded_args, expanded_kwargs = standard_kwargs( + kwarg_names, expanded_args_and_kwargs + ) + output = forward_helper(instance_norm, expanded_args, expanded_kwargs) + ctx.input = expanded_args[0] + ctx.running_mean, ctx.running_var = ( + expanded_kwargs["running_mean"], + expanded_kwargs["running_var"], + ) + ctx.weight, ctx.bias, ctx.eps = ( + expanded_kwargs["weight"], + expanded_kwargs["bias"], + expanded_kwargs["eps"], + ) + return output + + @staticmethod + # pyrefly: ignore [bad-override] + def backward(ctx, grad_output): + input, running_mean, running_var = ctx.input, ctx.running_mean, ctx.running_var + weight, bias, eps = ctx.weight, ctx.bias, ctx.eps + + results: list[torch.Tensor | None] = [] + results.append(None) # for kwarg names + results.append(None) # for op reference + if input.requires_grad: + b = input.shape[0] + c = input.shape[1] + new_shape = (1, b * c, *input.shape[2:]) + + weight_ = unpack_expanded_weight_or_tensor( + weight, lambda orig_weight: orig_weight.repeat(b) + ) + running_mean_ = running_mean.repeat(b) if running_mean is not None else None + running_var_ = running_var.repeat(b) if running_var is not None else None + input_reshaped = input.contiguous().view(new_shape) + grad_output_reshaped = grad_output.contiguous().view(new_shape) + mean = torch.mean( + input_reshaped, (0,) + tuple(range(2, input.dim())), False + ) + var = torch.var( + input_reshaped, + (0,) + tuple(range(2, input.dim())), + keepdim=False, + unbiased=False, + ) + rstd = 1 / torch.sqrt(var + eps) + + # must use native batch norm since it supports all inputs. This may have used cuda or openmi during the forward but + # it didn't save the metadata, so we don't know during the backward + res = torch.ops.aten.native_batch_norm_backward( + grad_output_reshaped, + input_reshaped, + weight_, + running_mean_, + running_var_, + mean, + rstd, + True, + eps, + (True, False, False), + ) + results.append(res[0].reshape(input.shape)) + else: + results.append(None) + + # weight and bias don't compute batched gradients; no other arguments are differentiable (2 are not saved from the forward) + results = results + [None] * 7 + + # set grad_sample field for weight and bias with per sample gradients + set_grad_sample_if_exists( + weight, + lambda _: torch.einsum( + "ni...->ni", F.instance_norm(input, eps=eps) * grad_output + ), + ) + set_grad_sample_if_exists( + bias, lambda _: torch.einsum("ni...->ni", grad_output) + ) + return tuple(results) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/layer_norm_expanded_weights.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/layer_norm_expanded_weights.py new file mode 100644 index 0000000000000000000000000000000000000000..705253861dbd0c0e16300ecc96561554c8ac6d60 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/layer_norm_expanded_weights.py @@ -0,0 +1,88 @@ +# mypy: allow-untyped-defs + +import torch +import torch.nn.functional as F + +from .expanded_weights_impl import ExpandedWeight, implements_per_sample_grads +from .expanded_weights_utils import ( + forward_helper, + set_grad_sample_if_exists, + standard_kwargs, + sum_over_all_but_batch_and_last_n, + unpack_expanded_weight_or_tensor, +) + + +@implements_per_sample_grads(F.layer_norm) +class LayerNormPerSampleGrad(torch.autograd.Function): + @staticmethod + # pyrefly: ignore [bad-override] + def forward(ctx, kwarg_names, _, *expanded_args_and_kwargs): + expanded_args, expanded_kwargs = standard_kwargs( + kwarg_names, expanded_args_and_kwargs + ) + input = expanded_args[0] + normalized_shape = expanded_args[1] + if len(input.shape) <= len(normalized_shape): + raise RuntimeError( + "Expanded Weights: Layer norm should not normalize over batch dimension for per sample gradient" + f"computations but got that normalized shape, {normalized_shape}, matched input shape." + ) + output, mean, rstd = forward_helper( + torch.native_layer_norm, expanded_args, expanded_kwargs + ) + ctx.args = expanded_args + + if input.requires_grad or isinstance(expanded_kwargs["weight"], ExpandedWeight): + ctx.weight = expanded_kwargs["weight"] + if input.requires_grad or isinstance(expanded_kwargs["bias"], ExpandedWeight): + ctx.bias = expanded_kwargs["bias"] + ctx.eps = expanded_kwargs["eps"] + ctx.mean, ctx.rstd = mean, rstd + return output + + @staticmethod + # pyrefly: ignore [bad-override] + def backward(ctx, grad_output): + def weight_per_sample_grad(weight): + return sum_over_all_but_batch_and_last_n( + F.layer_norm(input, normalized_shape, eps=ctx.eps) * grad_output, + weight.dim(), + ) + + input, normalized_shape = ctx.args + mean, rstd = ctx.mean, ctx.rstd + + results: list[torch.Tensor | None] = [] + results.append(None) # for kwarg names + results.append(None) # for op reference + if input.requires_grad: + weight_ = unpack_expanded_weight_or_tensor(ctx.weight) + bias_ = unpack_expanded_weight_or_tensor(ctx.bias) + results.append( + torch.ops.aten.native_layer_norm_backward( + grad_output, + input, + normalized_shape, + mean, + rstd, + weight_, + bias_, + (True, False, False), + )[0] + ) + else: + results.append(None) + + # weight and bias don't compute batched gradients; no other arguments are differentiable + results = results + [None] * 4 + + # set grad_sample field for weight and bias with per sample gradients + if hasattr(ctx, "weight"): + set_grad_sample_if_exists(ctx.weight, weight_per_sample_grad) + if hasattr(ctx, "bias"): + set_grad_sample_if_exists( + ctx.bias, + lambda bias: sum_over_all_but_batch_and_last_n(grad_output, bias.dim()), + ) + return tuple(results) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/linear_expanded_weights.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/linear_expanded_weights.py new file mode 100644 index 0000000000000000000000000000000000000000..2cd6b96f58bd614e1004de0ce939cdb90a85dc67 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_expanded_weights/linear_expanded_weights.py @@ -0,0 +1,63 @@ +# mypy: allow-untyped-defs + +import torch +import torch.nn.functional as F + +from .expanded_weights_impl import implements_per_sample_grads +from .expanded_weights_utils import ( + forward_helper, + is_batch_first, + set_grad_sample_if_exists, + unpack_expanded_weight_or_tensor, +) + + +@implements_per_sample_grads(F.linear) +class LinearPerSampleGrad(torch.autograd.Function): + @staticmethod + # pyrefly: ignore [bad-override] + def forward(ctx, _, __, *expanded_args_and_kwargs): + if len(expanded_args_and_kwargs[0].shape) <= 1: + raise RuntimeError( + "Input does not have a batch dimension. Expanded Weights expected input " + f"of at least rank 2, got of rank {len(expanded_args_and_kwargs[0].shape)}" + ) + expanded_kwargs = { + "bias": expanded_args_and_kwargs[2] + if len(expanded_args_and_kwargs) == 3 + else None + } + expanded_args = expanded_args_and_kwargs[:2] + ctx.batch_first = is_batch_first(expanded_args_and_kwargs) + output = forward_helper(F.linear, expanded_args, expanded_kwargs) + ctx.args = expanded_args + ctx.kwargs = expanded_kwargs + return output + + @staticmethod + # pyrefly: ignore [bad-override] + def backward(ctx, grad_output): + input, weight = ctx.args + bias = ctx.kwargs["bias"] + results: list[torch.Tensor | None] = [] + results.append(None) # for kwarg_names + results.append(None) # for op reference + + if input.requires_grad: + results.append(grad_output.matmul(unpack_expanded_weight_or_tensor(weight))) + else: + results.append(None) + results.extend([None] * 2) # weight and bias don't compute batched gradients + + if not ctx.batch_first: + grad_output = grad_output.transpose(0, 1) + input = input.transpose(0, 1) + + # weight and bias get their grad_sample fields set directly if they exist + set_grad_sample_if_exists( + weight, lambda _: torch.einsum("n...i,n...j->nij", grad_output, input) + ) + set_grad_sample_if_exists( + bias, lambda _: torch.einsum("n...k->nk", grad_output) + ) + return tuple(results) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_named_member_accessor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_named_member_accessor.py new file mode 100644 index 0000000000000000000000000000000000000000..35c695b77094bb26ea50f36120187508920d8c37 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_named_member_accessor.py @@ -0,0 +1,381 @@ +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from collections.abc import Iterable + +import torch + + +_MISSING: torch.Tensor = object() # type: ignore[assignment] + + +def set_tensor(module: "torch.nn.Module", name: str, tensor: torch.Tensor) -> None: + if not isinstance(module, torch.nn.Module): + raise TypeError(f"{module} is not an instance of torch.nn.Module") + if not isinstance(tensor, torch.Tensor) and tensor is not None: + raise TypeError(f"{tensor} is not an instance of torch.Tensor") + if "." in name: + raise KeyError('tensor name can\'t contain "."') + if name == "": + raise KeyError('tensor name can\'t be empty string ""') + if name in module._parameters: + module._parameters[name] = tensor # type: ignore[assignment] + elif name in module._buffers: + module._buffers[name] = tensor + else: + setattr(module, name, tensor) + + +def swap_tensor( + module: "torch.nn.Module", + name: str, + tensor: torch.Tensor, + allow_missing: bool = False, +) -> torch.Tensor: + if not isinstance(module, torch.nn.Module): + raise TypeError(f"{module} is not an instance of torch.nn.Module") + if ( + tensor is not _MISSING + and not isinstance(tensor, torch.Tensor) + and tensor is not None + ): + raise TypeError(f"{tensor} is not an instance of torch.Tensor") + if "." in name: + raise KeyError('tensor name can\'t contain "."') + if name == "": + raise KeyError('tensor name can\'t be empty string ""') + + orig_tensor: torch.Tensor + if name in module._parameters: + orig_tensor = module._parameters[name] # type: ignore[assignment] + if tensor is not _MISSING: + module._parameters[name] = tensor # type: ignore[assignment] + else: + del module._parameters[name] + elif name in module._buffers: + orig_tensor = module._buffers[name] # type: ignore[assignment] + if tensor is not _MISSING: + module._buffers[name] = tensor + else: + del module._buffers[name] + else: + if hasattr(module, name): + orig_tensor = getattr(module, name) + else: + if not allow_missing: + raise AttributeError(f"{module._get_name()} has no attribute `{name}`") + orig_tensor = _MISSING + if ( + orig_tensor is not _MISSING + and not isinstance(orig_tensor, torch.Tensor) + and orig_tensor is not None + ): + raise TypeError( + f"attribute `{name}`: {orig_tensor} is not an instance of torch.Tensor" + ) + if tensor is not _MISSING: + setattr(module, name, tensor) + elif hasattr(module, name): + delattr(module, name) + # pyrefly: ignore [bad-return] + return orig_tensor + + +def swap_submodule( + module: "torch.nn.Module", + name: str, + submodule: "torch.nn.Module", +) -> "torch.nn.Module": + if not isinstance(module, torch.nn.Module): + raise TypeError(f"{module} is not an instance of torch.nn.Module") + if not isinstance(submodule, torch.nn.Module): + raise TypeError(f"{submodule} is not an instance of torch.nn.Module") + if "." in name: + raise KeyError('submodule name can\'t contain "."') + if name == "": + raise KeyError('submodule name can\'t be empty string ""') + if name not in module._modules: + raise KeyError(f"submodule {name} does not exist") + + orig_submodule = module._modules[name] + if not isinstance(orig_submodule, torch.nn.Module): + raise TypeError(f"{name} attribute is not an instance of torch.nn.Module") + module._modules[name] = submodule + return orig_submodule + + +class NamedMemberAccessor: + """ + A class that provides a way to access the submodules and parameters/buffers of a module. + + It provides caching mechanism to speed up submodule lookups. + This is useful for functional programming to manipulate the module state. + """ + + def __init__(self, module: "torch.nn.Module") -> None: + self.module = module + self.memo: dict[str, torch.nn.Module] = {} + + # Nested attribute access + + def get_submodule(self, name: str) -> "torch.nn.Module": + """ + Return the submodule specified by the given path. + + For example, to get the submodule mod.layer1.conv1, + use accessor.get_submodule("layer1.conv1") + + Compare to mod.get_submodule("layer1.conv1"), this method will cache the + intermediate submodule access to speed up future lookups. + """ + if not name: + return self.module + + if name in self.memo: + return self.memo[name] + else: + prefix, dot, attr = name.rpartition(".") + if dot: + module = self.get_submodule(prefix) + else: + module = self.module + try: + submodule = getattr(module, attr) + except AttributeError as ex: + raise AttributeError( + f"{module._get_name()} has no attribute `{attr}`" + ) from ex + if not isinstance(submodule, torch.nn.Module): + raise TypeError( + f"submodule `{name}`: {submodule} is not an instance of torch.nn.Module" + ) + self.memo[name] = submodule + return submodule + + def swap_submodule(self, path: str, value: "torch.nn.Module") -> "torch.nn.Module": + """ + Swap the submodule specified by the given ``path`` to ``value``. + + For example, to swap the attribute mod.layer1.conv1 use + ``accessor.swap_submodule("layer1.conv1", conv2)``. + """ + prefix, _, attr = path.rpartition(".") + return swap_submodule(self.get_submodule(prefix), attr, value) + + def get_tensor(self, name: str) -> torch.Tensor: + """ + Get the tensor specified by the given path to value. + + For example, to get the attribute mod.layer1.conv1.weight, + use accessor.get_tensor('layer1.conv1.weight') + + Compare to mod.get_parameter("layer1.conv1.weight"), this method will + cache the intermediate submodule access to speed up future lookups. + """ + prefix, _, attr = name.rpartition(".") + submodule = self.get_submodule(prefix) + try: + tensor = getattr(submodule, attr) + except AttributeError as ex: + raise AttributeError( + f"{submodule._get_name()} has no attribute `{name}`" + ) from ex + if not isinstance(tensor, torch.Tensor) and tensor is not None: + raise TypeError(f"{tensor} is not an instance of torch.Tensor") + return tensor # type: ignore[return-value] + + def set_tensor(self, name: str, value: torch.Tensor) -> None: + """ + Set the attribute specified by the given path to value. + + For example, to set the attribute mod.layer1.conv1.weight, + use accessor.set_tensor("layer1.conv1.weight", value) + """ + prefix, _, attr = name.rpartition(".") + set_tensor(self.get_submodule(prefix), attr, value) + + def del_tensor(self, name: str) -> None: + """ + Delete the attribute specified by the given path. + + For example, to delete the attribute mod.layer1.conv1.weight, + use accessor.del_tensor("layer1.conv1.weight") + """ + prefix, _, attr = name.rpartition(".") + submodule = self.get_submodule(prefix) + try: + delattr(submodule, attr) + except AttributeError as ex: + raise AttributeError( + f"{submodule._get_name()} has no attribute `{name}`" + ) from ex + + def swap_tensor( + self, name: str, value: torch.Tensor, allow_missing: bool = False + ) -> torch.Tensor: + """ + Swap the attribute specified by the given path to value. + + For example, to swap the attribute mod.layer1.conv1.weight, + use accessor.swap_tensor("layer1.conv1.weight", value) + """ + prefix, _, attr = name.rpartition(".") + return swap_tensor( + self.get_submodule(prefix), attr, value, allow_missing=allow_missing + ) + + # Batched operations + + def get_tensors(self, names: Iterable[str]) -> list[torch.Tensor]: + """ + Get the tensors specified by the given paths. + + For example, to get the attributes mod.layer1.conv1.weight and + mod.layer1.conv1.bias, use accessor.get_tensors(["layer1.conv1.weight", + "layer1.conv1.bias"]) + """ + return [self.get_tensor(name) for name in names] + + def set_tensors(self, names: Iterable[str], values: Iterable[torch.Tensor]) -> None: + """ + Set the attributes specified by the given paths to values. + + For example, to set the attributes mod.layer1.conv1.weight and + mod.layer1.conv1.bias, use accessor.set_tensors(["layer1.conv1.weight", + "layer1.conv1.bias"], [weight, bias]) + """ + if not isinstance(names, (list, tuple)): + names = list(names) + if not isinstance(values, (list, tuple)): + values = list(values) + if len(names) != len(values): + raise AssertionError( + f"names and values must have the same length, " + f"got {len(names)} names and {len(values)} values" + ) + + for name, value in zip(names, values, strict=True): + self.set_tensor(name, value) + + def set_tensors_dict(self, named_tensors: dict[str, torch.Tensor]) -> None: + """ + Set the attributes specified by the given paths to values. + + For example, to set the attributes mod.layer1.conv1.weight and + mod.layer1.conv1.bias, use accessor.set_tensors_dict({ + "layer1.conv1.weight": weight, + "layer1.conv1.bias": bias, + }) + """ + for name, value in named_tensors.items(): + self.set_tensor(name, value) + + def del_tensors(self, names: Iterable[str]) -> None: + """ + Delete the attributes specified by the given paths. + + For example, to delete the attributes mod.layer1.conv1.weight and + mod.layer1.conv1.bias, use accessor.del_tensors(["layer1.conv1.weight", + "layer1.conv1.bias"]) + """ + for name in names: + self.del_tensor(name) + + def swap_tensors( + self, + names: Iterable[str], + values: Iterable[torch.Tensor], + allow_missing: bool = False, + ) -> list[torch.Tensor]: + """ + Swap the attributes specified by the given paths to values. + + For example, to swap the attributes mod.layer1.conv1.weight and + mod.layer1.conv1.bias, use accessor.swap_tensors(["layer1.conv1.weight", + "layer1.conv1.bias"], [weight, bias]) + """ + if not isinstance(names, (list, tuple)): + names = list(names) + if not isinstance(values, (list, tuple)): + values = list(values) + if len(names) != len(values): + raise AssertionError( + f"names and values must have the same length, " + f"got {len(names)} names and {len(values)} values" + ) + + return [ + self.swap_tensor(name, value, allow_missing=allow_missing) + for name, value in zip(names, values, strict=True) + ] + + def swap_tensors_dict( + self, named_tensors: dict[str, torch.Tensor], allow_missing: bool = False + ) -> tuple[dict[str, torch.Tensor], list[str]]: + """ + Swap the attributes specified by the given paths to values. + + For example, to swap the attributes mod.layer1.conv1.weight and + mod.layer1.conv1.bias, use accessor.swap_tensors_dict({ + "layer1.conv1.weight": weight, + "layer1.conv1.bias": bias, + }) + """ + orig_named_tensors = {} + missing_keys = [] + try: + for name, tensor in named_tensors.items(): + orig_tensor = self.swap_tensor(name, tensor, allow_missing=True) + if orig_tensor is _MISSING: + missing_keys.append(name) + orig_named_tensors[name] = orig_tensor + except Exception: + # Swap back if any exception occurs + for name, orig_tensor in orig_named_tensors.items(): + self.swap_tensor(name, orig_tensor, allow_missing=True) + raise + if missing_keys and not allow_missing: + # Swap back if any key is missing when allow_missing is False + for name, orig_tensor in orig_named_tensors.items(): + self.swap_tensor(name, orig_tensor, allow_missing=True) + raise RuntimeError(f"Missing key(s): {', '.join(map(repr, missing_keys))}.") + return orig_named_tensors, missing_keys + + def check_keys(self, keys: Iterable[str]) -> tuple[list[str], list[str]]: + """Check that the given keys are valid.""" + keys = set(keys) + valid_keys = {name for name, _ in self.named_tensors(remove_duplicate=False)} + missing_keys = valid_keys - keys + unexpected_keys = keys - valid_keys + return sorted(missing_keys), sorted(unexpected_keys) + + # Shortcut methods + + def named_parameters( + self, + remove_duplicate: bool = True, + ) -> Iterable[tuple[str, torch.Tensor]]: + """Iterate over all the parameters in the module.""" + yield from self.module.named_parameters(remove_duplicate=remove_duplicate) + + def named_buffers( + self, + remove_duplicate: bool = True, + ) -> Iterable[tuple[str, torch.Tensor]]: + """Iterate over all the buffers in the module.""" + yield from self.module.named_buffers(remove_duplicate=remove_duplicate) + + def named_tensors( + self, + remove_duplicate: bool = True, + ) -> Iterable[tuple[str, torch.Tensor]]: + """Iterate over all the tensors in the module.""" + yield from self.module.named_parameters(remove_duplicate=remove_duplicate) + yield from self.module.named_buffers(remove_duplicate=remove_duplicate) + + def named_modules( + self, + remove_duplicate: bool = True, + ) -> Iterable[tuple[str, "torch.nn.Module"]]: + """Iterate over all the modules in the module.""" + yield from self.module.named_modules(remove_duplicate=remove_duplicate) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_per_sample_grad.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_per_sample_grad.py new file mode 100644 index 0000000000000000000000000000000000000000..2eae0865845eec9c426c5cc3b7bff1b11b5b1230 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/_per_sample_grad.py @@ -0,0 +1,126 @@ +# mypy: allow-untyped-defs +import functools + +import torch +from torch.nn.utils._expanded_weights.expanded_weights_impl import ExpandedWeight +from torch.utils import _pytree as pytree + + +# dependency on `functional_call` means that this can't be exposed in utils +# without creating circular dependency +def call_for_per_sample_grads( + module, + *, + batch_size=None, + loss_reduction="sum", + batch_first=True, +): + r""" + Return a forward function for a module, populating grad_sample with per sample gradients on backward invocation. + + Args: + module: The ``nn.Module`` to get per sample gradients with respect to. All trainable + parameters will compute per sample gradients, located in a ``grad_sample`` + field when ``backward`` is invoked + batch_size: The batch size of the input. If None is passed, all tensor arguments in args and kwargs must have + the same batch size, which is the size of the first dimension. Otherwise, it must be passed manually. + Default: None + loss_reduction: Indicates if the loss reduction (for aggregating the gradients) is a sum or a mean operation. If + "mean", per sample gradients will be scaled by the batch size to offset the crossbatch interaction from + running mean across a batch. Must be "mean" or "sum". Default: "sum" + batch_first: Indicates if the batch dimension is the first dimension. If True, the batch dimension is the first + dimension. If False, it's the second dimension. Default: True. + + Examples:: + >>> # xdoctest: +SKIP + >>> model = nn.Linear(4, 3) + >>> batched_input = torch.randn(5, 4) # batch size of 5 + >>> res = call_for_per_sample_grads(model)(batched_input).sum() + >>> res.backward() + >>> assert model.weight.shape == (3, 4) + >>> assert model.weight.grad_sample.shape == (5, 3, 4) + >>> assert model.weight.grad is None + >>> assert model.bias.shape == (3,) + >>> assert model.bias.grad_sample.shape == (5, 3) + >>> assert model.bias.grad is None + + An example using "mean" loss reduction. The grad_sample fields will be scaled by batch_size from what they would be + if we ran the same code with loss_reduction="sum". This is because the mean at the end will scale all + grad_outputs by 1 / batch_size from cross batch interaction. + >>> model = nn.Linear(4, 3) + >>> batched_input = torch.randn(5, 4) # batch size of 5 + >>> res = call_for_per_sample_grads(model, 5, loss_reduction="mean")( + ... batched_input + ... ).mean() + >>> res.backward() + + Note:: + Does not work with any `nn.RNN`, including `nn.GRU` or `nn.LSTM`. Please use custom + rewrites that wrap an `nn.Linear` module. See Opacus for an example + """ + + def maybe_build_expanded_weight(og_tensor, batch_size): + if og_tensor.requires_grad: + return ExpandedWeight(og_tensor, batch_size, loss_reduction) + else: + return og_tensor + + def compute_batch_size(*args, **kwargs): + args_and_kwargs = pytree.arg_tree_leaves(*args, **kwargs) + batch_size = None + for arg in args_and_kwargs: + if not isinstance(arg, torch.Tensor): + continue + + arg_batch_size = arg.shape[0] if batch_first else arg.shape[1] + if batch_size is not None and batch_size != arg_batch_size: + raise RuntimeError( + "When computing batch size, found at least one input with batch size " + f"{batch_size} and one with batch size {arg_batch_size}. Please specify it " + "explicitly using the batch size kwarg in call_for_per_sample_grads" + ) + batch_size = arg_batch_size + if batch_size is None: + raise RuntimeError( + "Unable to find a tensor in the passed args and kwargs. They may not be pytree-able " + "and so ExpandedWeights cannot compute the batch size from the inputs. Please specify " + "it explicitly" + ) + return batch_size + + if loss_reduction not in ["sum", "mean"]: + raise RuntimeError( + f"Expected loss_reduction argument to be sum or mean, got {loss_reduction}" + ) + + if not isinstance(module, torch.nn.Module): + raise RuntimeError( + f"Module passed must be nn.Module, got {type(module).__name__}" + ) + if not (batch_size is None or isinstance(batch_size, int)): + raise RuntimeError( + f"Batch size passed must be None or an integer, got {type(batch_size).__name__}" + ) + if batch_size is not None and batch_size < 1: + raise RuntimeError(f"Batch size must be positive, got {batch_size}") + for weight in module.parameters(): + if hasattr(weight, "grad_sample") and weight.grad_sample is not None: # type: ignore[attr-defined] + raise RuntimeError( + "Current Expanded Weights accumulates the gradients, which will be incorrect for multiple " + f"calls without clearing gradients. Please clear out the grad_sample parameter of {weight} or " + "post an issue to pytorch/pytorch to prioritize correct behavior" + ) + + @functools.wraps(module.forward) + def wrapper(*args, **kwargs): + wrapper_batch_size = batch_size + if wrapper_batch_size is None: + wrapper_batch_size = compute_batch_size(*args, **kwargs) + + params = { + name: maybe_build_expanded_weight(value, wrapper_batch_size) + for (name, value) in module.named_parameters() + } + return torch.func.functional_call(module, params, args, kwargs) + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/clip_grad.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/clip_grad.py new file mode 100644 index 0000000000000000000000000000000000000000..c7b2d6fcddfcae934c65826166ac0028f83989a5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/clip_grad.py @@ -0,0 +1,298 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +import functools +import types +import typing +import warnings +from collections.abc import Callable +from typing import cast, TypeAlias, TypeVar +from typing_extensions import deprecated, ParamSpec + +import torch +from torch import Tensor +from torch.utils._foreach_utils import ( + _device_has_foreach_support, + _group_tensors_by_device_and_dtype, + _has_foreach_support, +) + + +__all__: list[str] = [ + "clip_grad_norm", + "clip_grad_norm_", + "clip_grad_value_", +] + + +_tensor_or_tensors: TypeAlias = torch.Tensor | typing.Iterable[torch.Tensor] # noqa: PYI042 + +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +def _no_grad(func: Callable[_P, _R]) -> Callable[_P, _R]: + """ + This wrapper is needed to avoid a circular import when using @torch.no_grad on the exposed functions + clip_grad_norm_ and clip_grad_value_ themselves. + """ + + def _no_grad_wrapper(*args, **kwargs): + with torch.no_grad(): + return func(*args, **kwargs) + + functools.update_wrapper(_no_grad_wrapper, func) + # pyrefly: ignore [bad-return] + return _no_grad_wrapper + + +@_no_grad +def _get_total_norm( + tensors: _tensor_or_tensors, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: bool | None = None, +) -> torch.Tensor: + r"""Compute the norm of an iterable of tensors. + + The norm is computed over the norms of the individual tensors, as if the norms of + the individual tensors were concatenated into a single vector. + + Args: + tensors (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will be normalized + norm_type (float): type of the used p-norm. Can be ``'inf'`` for + infinity norm. + error_if_nonfinite (bool): if True, an error is thrown if the total + norm of :attr:`tensors` is ``nan``, ``inf``, or ``-inf``. + Default: ``False`` + foreach (bool): use the faster foreach-based implementation. + If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently + fall back to the slow implementation for other device types. + Default: ``None`` + + Returns: + Total norm of the tensors (viewed as a single vector). + """ + if isinstance(tensors, torch.Tensor): + tensors = [tensors] + else: + tensors = list(tensors) + norm_type = float(norm_type) + if len(tensors) == 0: + return torch.tensor(0.0) + first_device = tensors[0].device + grouped_tensors: dict[ + tuple[torch.device, torch.dtype], tuple[list[list[Tensor]], list[int]] + ] = _group_tensors_by_device_and_dtype( # pyrefly: ignore [bad-assignment] + [tensors] # type: ignore[list-item] + ) # type: ignore[assignment] + + norms: list[Tensor] = [] + for (device, _), ([device_tensors], _) in grouped_tensors.items(): + if (foreach is None and _has_foreach_support(device_tensors, device)) or ( + foreach and _device_has_foreach_support(device) + ): + norms.extend(torch._foreach_norm(device_tensors, norm_type)) + elif foreach: + raise RuntimeError( + f"foreach=True was passed, but can't use the foreach API on {device.type} tensors" + ) + else: + norms.extend( + [torch.linalg.vector_norm(g, norm_type) for g in device_tensors] + ) + + total_norm = torch.linalg.vector_norm( + torch.stack([norm.to(first_device) for norm in norms]), norm_type + ) + + if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()): + raise RuntimeError( + f"The total norm of order {norm_type} for gradients from " + "`parameters` is non-finite, so it cannot be clipped. To disable " + "this error and scale the gradients by the non-finite norm anyway, " + "set `error_if_nonfinite=False`" + ) + return total_norm + + +@_no_grad +def _clip_grads_with_norm_( + parameters: _tensor_or_tensors, + max_norm: float, + total_norm: torch.Tensor, + foreach: bool | None = None, +) -> None: + r"""Scale the gradients of an iterable of parameters given a pre-calculated total norm and desired max norm. + + The gradients will be scaled by the following calculation + + .. math:: + grad = grad * \min(\frac{max\_norm}{total\_norm + 1e-6}, 1) + + Gradients are modified in-place. + + Note: The scale coefficient is clamped to a maximum of 1.0 to prevent gradient amplification. + This ensures that gradients are only scaled down when the total norm exceeds max_norm. + + This function is equivalent to :func:`torch.nn.utils.clip_grad_norm_` with a pre-calculated + total norm. + + Args: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + max_norm (float): max norm of the gradients + total_norm (Tensor): total norm of the gradients to use for clipping + foreach (bool): use the faster foreach-based implementation. + If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently + fall back to the slow implementation for other device types. + Default: ``None`` + + Returns: + None + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + grads = [p.grad for p in parameters if p.grad is not None] + max_norm = float(max_norm) + if len(grads) == 0: + return + grouped_grads: dict[ + tuple[torch.device, torch.dtype], tuple[list[list[Tensor]], list[int]] + ] = _group_tensors_by_device_and_dtype([grads]) # type: ignore[assignment] + + clip_coef = max_norm / (total_norm + 1e-6) + # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so + # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization + # when the gradients do not reside in CPU memory. + clip_coef_clamped = torch.clamp(clip_coef, max=1.0) + for (device, _), ([device_grads], _) in grouped_grads.items(): + if (foreach is None and _has_foreach_support(device_grads, device)) or ( + foreach and _device_has_foreach_support(device) + ): + torch._foreach_mul_(device_grads, clip_coef_clamped.to(device)) + elif foreach: + raise RuntimeError( + f"foreach=True was passed, but can't use the foreach API on {device.type} tensors" + ) + else: + clip_coef_clamped_device = clip_coef_clamped.to(device) + for g in device_grads: + g.mul_(clip_coef_clamped_device) + + +@_no_grad +def clip_grad_norm_( + parameters: _tensor_or_tensors, + max_norm: float, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: bool | None = None, +) -> torch.Tensor: + r"""Clip the gradient norm of an iterable of parameters. + + The norm is computed over the norms of the individual gradients of all parameters, + as if the norms of the individual gradients were concatenated into a single vector. + Gradients are modified in-place. + + This function is equivalent to :func:`torch.nn.utils.get_total_norm` followed by + :func:`torch.nn.utils.clip_grads_with_norm_` with the ``total_norm`` returned by ``get_total_norm``. + + Args: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + max_norm (float): max norm of the gradients + norm_type (float, optional): type of the used p-norm. Can be ``'inf'`` for + infinity norm. Default: 2.0 + error_if_nonfinite (bool, optional): if True, an error is thrown if the total + norm of the gradients from :attr:`parameters` is ``nan``, + ``inf``, or ``-inf``. Default: False + foreach (bool, optional): use the faster foreach-based implementation. + If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently + fall back to the slow implementation for other device types. + Default: ``None`` + + Returns: + Total norm of the parameter gradients (viewed as a single vector). + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + else: + is_generator = isinstance(parameters, types.GeneratorType) + # prevent generators from being exhausted + parameters = list(parameters) + if is_generator and len(parameters) == 0: + warnings.warn( + "`parameters` is an empty generator, no gradient clipping will occur.", + stacklevel=3, + ) + grads = [p.grad for p in parameters if p.grad is not None] + total_norm = _get_total_norm(grads, norm_type, error_if_nonfinite, foreach) + _clip_grads_with_norm_(parameters, max_norm, total_norm, foreach) + return total_norm + + +@deprecated( + "`torch.nn.utils.clip_grad_norm` is now deprecated " + "in favor of `torch.nn.utils.clip_grad_norm_`.", + category=FutureWarning, +) +def clip_grad_norm( + parameters: _tensor_or_tensors, + max_norm: float, + norm_type: float = 2.0, + error_if_nonfinite: bool = False, + foreach: bool | None = None, +) -> torch.Tensor: + r"""Clip the gradient norm of an iterable of parameters. + + .. warning:: + This method is now deprecated in favor of + :func:`torch.nn.utils.clip_grad_norm_`. + """ + return clip_grad_norm_(parameters, max_norm, norm_type, error_if_nonfinite, foreach) + + +@_no_grad +def clip_grad_value_( + parameters: _tensor_or_tensors, + clip_value: float, + foreach: bool | None = None, +) -> None: + r"""Clip the gradients of an iterable of parameters at specified value. + + Gradients are modified in-place. + + Args: + parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a + single Tensor that will have gradients normalized + clip_value (float): maximum allowed value of the gradients. + The gradients are clipped in the range + :math:`\left[\text{-clip\_value}, \text{clip\_value}\right]` + foreach (bool, optional): use the faster foreach-based implementation + If ``None``, use the foreach implementation for CUDA and CPU native tensors and + silently fall back to the slow implementation for other device types. + Default: ``None`` + """ + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + clip_value = float(clip_value) + + grads = [p.grad for p in parameters if p.grad is not None] + # pyrefly: ignore [bad-argument-type] + grouped_grads = _group_tensors_by_device_and_dtype([grads]) + + for (device, _), ([grads], _) in grouped_grads.items(): + if ( + foreach is None + and _has_foreach_support(cast(list[Tensor], grads), device=device) + ) or (foreach and _device_has_foreach_support(device)): + torch._foreach_clamp_min_(cast(list[Tensor], grads), -clip_value) + torch._foreach_clamp_max_(cast(list[Tensor], grads), clip_value) + elif foreach: + raise RuntimeError( + f"foreach=True was passed, but can't use the foreach API on {device.type} tensors" + ) + else: + for grad in grads: + cast(Tensor, grad).clamp_(min=-clip_value, max=clip_value) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/convert_parameters.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/convert_parameters.py new file mode 100644 index 0000000000000000000000000000000000000000..6a56da711ecda3c6e3d5770783f100a8890bbf55 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/convert_parameters.py @@ -0,0 +1,90 @@ +from collections.abc import Iterable + +import torch + + +def parameters_to_vector(parameters: Iterable[torch.Tensor]) -> torch.Tensor: + r"""Flatten an iterable of parameters into a single vector. + + Args: + parameters (Iterable[Tensor]): an iterable of Tensors that are the + parameters of a model. + + Returns: + The parameters represented by a single vector + """ + # Flag for the device where the parameter is located + param_device = None + + vec = [] + for param in parameters: + # Ensure the parameters are located in the same device + param_device = _check_param_device(param, param_device) + + vec.append(param.view(-1)) + return torch.cat(vec) + + +def vector_to_parameters(vec: torch.Tensor, parameters: Iterable[torch.Tensor]) -> None: + r"""Copy slices of a vector into an iterable of parameters. + + Args: + vec (Tensor): a single vector representing the parameters of a model. + parameters (Iterable[Tensor]): an iterable of Tensors that are the + parameters of a model. + """ + # Ensure vec of type Tensor + if not isinstance(vec, torch.Tensor): + raise TypeError(f"expected torch.Tensor, but got: {torch.typename(vec)}") + # Flag for the device where the parameter is located + param_device = None + + # Pointer for slicing the vector for each parameter + pointer = 0 + for param in parameters: + # Ensure the parameters are located in the same device + param_device = _check_param_device(param, param_device) + + # The length of the parameter + num_param = param.numel() + # Slice the vector, reshape it, and replace the old data of the parameter + param.data = vec[pointer : pointer + num_param].view_as(param).data + + # Increment the pointer + pointer += num_param + + +def _check_param_device(param: torch.Tensor, old_param_device: int | None) -> int: + r"""Check if the parameters are located on the same device. + + Currently, the conversion between model parameters and single vector form is not supported + for multiple allocations, e.g. parameters in different GPUs/PrivateUse1s, or mixture of CPU/GPU/PrivateUse1. + + Args: + param ([Tensor]): a Tensor of a parameter of a model + old_param_device (int): the device where the first parameter of a + model is allocated. + + Returns: + old_param_device (int): report device for the first time + """ + # Meet the first parameter + support_device_types = ["cuda", torch._C._get_privateuse1_backend_name()] + if old_param_device is None: + old_param_device = ( + param.get_device() if param.device.type in support_device_types else -1 + ) + else: + warn = False + if ( + param.device.type in support_device_types + ): # Check if in same GPU/PrivateUse1 + warn = param.get_device() != old_param_device + else: # Check if in CPU + warn = old_param_device != -1 + if warn: + raise TypeError( + "Found two parameters on different devices, " + "this is currently not supported." + ) + return old_param_device diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/fusion.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/fusion.py new file mode 100644 index 0000000000000000000000000000000000000000..0a9ac9a9668e1de589b1b2c330c4b8d2904854c3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/fusion.py @@ -0,0 +1,196 @@ +from __future__ import annotations + +import copy +from typing import TypeVar + +import torch + + +__all__ = [ + "fuse_conv_bn_eval", + "fuse_conv_bn_weights", + "fuse_linear_bn_eval", + "fuse_linear_bn_weights", +] + +ConvT = TypeVar("ConvT", bound="torch.nn.modules.conv._ConvNd") +LinearT = TypeVar("LinearT", bound="torch.nn.Linear") + + +def fuse_conv_bn_eval( + conv: ConvT, + bn: torch.nn.modules.batchnorm._BatchNorm, + transpose: bool = False, +) -> ConvT: + r"""Fuse a convolutional module and a BatchNorm module into a single, new convolutional module. + + Args: + conv (torch.nn.modules.conv._ConvNd): A convolutional module. + bn (torch.nn.modules.batchnorm._BatchNorm): A BatchNorm module. + transpose (bool, optional): If True, transpose the convolutional weight. Defaults to False. + + Returns: + torch.nn.modules.conv._ConvNd: The fused convolutional module. + + .. note:: + Both ``conv`` and ``bn`` must be in eval mode, and ``bn`` must have its running buffers computed. + """ + if conv.training or bn.training: + raise AssertionError("Fusion only for eval!") + fused_conv = copy.deepcopy(conv) + + if bn.running_mean is None or bn.running_var is None: + raise AssertionError("bn.running_mean and bn.running_var must not be None") + fused_conv.weight, fused_conv.bias = fuse_conv_bn_weights( + fused_conv.weight, + fused_conv.bias, + bn.running_mean, + bn.running_var, + bn.eps, + bn.weight, + bn.bias, + transpose, + ) + + return fused_conv + + +def fuse_conv_bn_weights( + conv_w: torch.Tensor, + conv_b: torch.Tensor | None, + bn_rm: torch.Tensor, + bn_rv: torch.Tensor, + bn_eps: float, + bn_w: torch.Tensor | None, + bn_b: torch.Tensor | None, + transpose: bool = False, +) -> tuple[torch.nn.Parameter, torch.nn.Parameter]: + r"""Fuse convolutional module parameters and BatchNorm module parameters into new convolutional module parameters. + + Args: + conv_w (torch.Tensor): Convolutional weight. + conv_b (Optional[torch.Tensor]): Convolutional bias. + bn_rm (torch.Tensor): BatchNorm running mean. + bn_rv (torch.Tensor): BatchNorm running variance. + bn_eps (float): BatchNorm epsilon. + bn_w (Optional[torch.Tensor]): BatchNorm weight. + bn_b (Optional[torch.Tensor]): BatchNorm bias. + transpose (bool, optional): If True, transpose the conv weight. Defaults to False. + + Returns: + Tuple[torch.nn.Parameter, torch.nn.Parameter]: Fused convolutional weight and bias. + """ + conv_weight_dtype = conv_w.dtype + conv_bias_dtype = conv_b.dtype if conv_b is not None else conv_weight_dtype + if conv_b is None: + conv_b = torch.zeros_like(bn_rm) + if bn_w is None: + bn_w = torch.ones_like(bn_rm) + if bn_b is None: + bn_b = torch.zeros_like(bn_rm) + bn_var_rsqrt = torch.rsqrt(bn_rv + bn_eps) + + if transpose: + shape = [1, -1] + [1] * (len(conv_w.shape) - 2) + else: + shape = [-1, 1] + [1] * (len(conv_w.shape) - 2) + + fused_conv_w = (conv_w * (bn_w * bn_var_rsqrt).reshape(shape)).to( + dtype=conv_weight_dtype + ) + fused_conv_b = ((conv_b - bn_rm) * bn_var_rsqrt * bn_w + bn_b).to( + dtype=conv_bias_dtype + ) + + return ( + torch.nn.Parameter(fused_conv_w, conv_w.requires_grad), + torch.nn.Parameter(fused_conv_b, conv_b.requires_grad), + ) + + +def fuse_linear_bn_eval( + linear: LinearT, + bn: torch.nn.modules.batchnorm._BatchNorm, +) -> LinearT: + r"""Fuse a linear module and a BatchNorm module into a single, new linear module. + + Args: + linear (torch.nn.Linear): A Linear module. + bn (torch.nn.modules.batchnorm._BatchNorm): A BatchNorm module. + + Returns: + torch.nn.Linear: The fused linear module. + + .. note:: + Both ``linear`` and ``bn`` must be in eval mode, and ``bn`` must have its running buffers computed. + """ + if linear.training or bn.training: + raise AssertionError("Fusion only for eval!") + fused_linear = copy.deepcopy(linear) + + """ + Linear-BN needs to be fused while preserving the shapes of linear weight/bias. + To preserve the shapes of linear weight/bias, the channel dim of bn needs to be broadcastable with the last dim of linear, + because bn operates over the channel dim, (N, C_in, H, W) while linear operates over the last dim, (*, H_in). + To be broadcastable, the number of features in bn and + the number of output features from linear must satisfy the following condition: + 1. they are equal, or + 2. the number of features in bn is 1 + Otherwise, skip the folding path + """ + if linear.out_features != bn.num_features and bn.num_features != 1: + raise AssertionError( + f"To fuse, linear.out_features == bn.num_features or bn.num_features == 1, " + f"got linear.out_features={linear.out_features} and bn.num_features={bn.num_features}" + ) + + if bn.running_mean is None or bn.running_var is None: + raise AssertionError("bn.running_mean and bn.running_var must not be None") + fused_linear.weight, fused_linear.bias = fuse_linear_bn_weights( + fused_linear.weight, + fused_linear.bias, + bn.running_mean, + bn.running_var, + bn.eps, + bn.weight, + bn.bias, + ) + + return fused_linear + + +def fuse_linear_bn_weights( + linear_w: torch.Tensor, + linear_b: torch.Tensor | None, + bn_rm: torch.Tensor, + bn_rv: torch.Tensor, + bn_eps: float, + bn_w: torch.Tensor, + bn_b: torch.Tensor, +) -> tuple[torch.nn.Parameter, torch.nn.Parameter]: + r"""Fuse linear module parameters and BatchNorm module parameters into new linear module parameters. + + Args: + linear_w (torch.Tensor): Linear weight. + linear_b (Optional[torch.Tensor]): Linear bias. + bn_rm (torch.Tensor): BatchNorm running mean. + bn_rv (torch.Tensor): BatchNorm running variance. + bn_eps (float): BatchNorm epsilon. + bn_w (torch.Tensor): BatchNorm weight. + bn_b (torch.Tensor): BatchNorm bias. + + Returns: + Tuple[torch.nn.Parameter, torch.nn.Parameter]: Fused linear weight and bias. + """ + linear_weight_dtype = linear_w.dtype + linear_bias_dtype = linear_b.dtype if linear_b is not None else linear_weight_dtype + if linear_b is None: + linear_b = torch.zeros_like(bn_rm) + bn_scale = bn_w * torch.rsqrt(bn_rv + bn_eps) + + fused_w = linear_w * bn_scale.unsqueeze(-1).to(dtype=linear_weight_dtype) + fused_b = ((linear_b - bn_rm) * bn_scale + bn_b).to(dtype=linear_bias_dtype) + + return torch.nn.Parameter(fused_w, linear_w.requires_grad), torch.nn.Parameter( + fused_b, linear_b.requires_grad + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/init.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/init.py new file mode 100644 index 0000000000000000000000000000000000000000..10fa03b7c01c2eac7e474ef55f433e4704e6c778 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/init.py @@ -0,0 +1,55 @@ +# mypy: allow-untyped-defs +import inspect + +import torch + + +def skip_init(module_cls, *args, **kwargs): + r""" + Given a module class object and args / kwargs, instantiate the module without initializing parameters / buffers. + + This can be useful if initialization is slow or if custom initialization will + be performed, making the default initialization unnecessary. There are some caveats to this, due to + the way this function is implemented: + + 1. The module must accept a `device` arg in its constructor that is passed to any parameters + or buffers created during construction. + + 2. The module must not perform any computation on parameters in its constructor except + initialization (i.e. functions from :mod:`torch.nn.init`). + + If these conditions are satisfied, the module can be instantiated with parameter / buffer values + uninitialized, as if having been created using :func:`torch.empty`. + + Args: + module_cls: Class object; should be a subclass of :class:`torch.nn.Module` + args: args to pass to the module's constructor + kwargs: kwargs to pass to the module's constructor + + Returns: + Instantiated module with uninitialized parameters / buffers + + Example:: + + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> import torch + >>> m = torch.nn.utils.skip_init(torch.nn.Linear, 5, 1) + >>> m.weight + Parameter containing: + tensor([[0.0000e+00, 1.5846e+29, 7.8307e+00, 2.5250e-29, 1.1210e-44]], + requires_grad=True) + >>> m2 = torch.nn.utils.skip_init(torch.nn.Linear, in_features=6, out_features=1) + >>> m2.weight + Parameter containing: + tensor([[-1.4677e+24, 4.5915e-41, 1.4013e-45, 0.0000e+00, -1.4677e+24, + 4.5915e-41]], requires_grad=True) + + """ + if not issubclass(module_cls, torch.nn.Module): + raise RuntimeError(f"Expected a Module; got {module_cls}") + if "device" not in inspect.signature(module_cls).parameters: + raise RuntimeError("Module must support a 'device' arg to skip initialization") + + final_device = kwargs.pop("device", "cpu") + kwargs["device"] = "meta" + return module_cls(*args, **kwargs).to_empty(device=final_device) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/memory_format.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/memory_format.py new file mode 100644 index 0000000000000000000000000000000000000000..d7ee4d40a05611c1951269bd6c7d15556b186e20 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/memory_format.py @@ -0,0 +1,174 @@ +from __future__ import annotations + +from typing import TypeVar + +import torch + + +_M = TypeVar("_M", bound="torch.nn.Module") + + +def convert_conv2d_weight_memory_format( + module: _M, memory_format: torch.memory_format +) -> _M: + r"""Convert ``memory_format`` of ``nn.Conv2d.weight`` to ``memory_format``. + + The conversion recursively applies to nested ``nn.Module``, including ``module``. + Note that it only changes the memory_format, but not the semantics of each dimensions. + This function is used to facilitate the computation to adopt NHWC kernels, which + provides considerable speed up for fp16 data on CUDA devices with compute capability >= 7.0 + + .. note:: + Calling ``model.to(memory_format=torch.channels_last)`` is more aggressive + than the utility function ``convert_conv2d_weight_memory_format``. Any + layer with 4d weight will be affected by ``model.to``, which does not + necessarily benefit from conversion to specified ``memory_format``. + One place we are confident in is that NHWC(channels_last) conversion for + convolution in cuDNN, as it is beneficial to run convolution in NHWC, + even in cases where we have to apply permutation to input tensors. + + Hence our strategy here is to convert only the weight of convolution to + channels_last. This ensures that; + 1. Fast convolution kernels will be used, the benefit of which could + outweigh overhead of permutation (if input is not in the same format). + 2. No unnecessary permutations are applied on layers that do not benefit + from memory_format conversion. + + The optimal case is that, layers between convolution layers are channels + last compatible. Input tensor would be permuted to channels last when it + encounters the first convolution layer and stay in that memory format. + Hence following convolutions will not need to permute its input tensor. + + In case where a channels last incompatible layer is between convolution + layers, we need to permute the input tensor back to contiguous format + for that layer. The input tensor will go through the remaining layers in + contiguous format and be permuted to channels last when it encounters + another convolution layer. There's no point in propagating that + permutation to an earlier layer, as most layers are quite agnostic to + ``memory_format``. + + This claim might change when PyTorch supports fusion of permutation, as + there might have been a better spot to fuse the permutation other than + immediately before a convolution. + + Args: + module (nn.Module): ``nn.Conv2d`` & ``nn.ConvTranspose2d`` or container + ``nn.Module`` + memory_format: user specified ``memory_format``, + e.g. ``torch.channels_last`` or ``torch.contiguous_format`` + + Returns: + The original module with updated ``nn.Conv2d`` + + Example: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> # xdoctest: +REQUIRES(env:CUBLAS_WORKSPACE_CONFIG) + >>> input = torch.randint( + ... 1, 10, (2, 8, 4, 4), dtype=torch.float16, device="cuda" + ... ) + >>> model = nn.Sequential( + >>> nn.Conv2d(8, 4, 3)).cuda().half() + >>> # This is identical to: + >>> # nn.utils.convert_conv2d_weight_memory_format(model, torch.channels_last) + >>> model = nn.utils.convert_conv2d_weight_memory_format( + ... model, torch.channels_last + ... ) + >>> out = model(input) + """ + # TODO: expand this to `_ConvNd` when channels_last support is extended + # beyond only 4d tensors. + if isinstance(module, (torch.nn.Conv2d, torch.nn.ConvTranspose2d)): + weight_data = module.weight.detach().clone(memory_format=memory_format) + module.weight.data = weight_data.resize_( + weight_data.size(), memory_format=memory_format + ) + for child in module.children(): + convert_conv2d_weight_memory_format(child, memory_format) + + return module + + +def convert_conv3d_weight_memory_format( + module: _M, memory_format: torch.memory_format +) -> _M: + r"""Convert ``memory_format`` of ``nn.Conv3d.weight`` to ``memory_format`` + The conversion recursively applies to nested ``nn.Module``, including ``module``. + Note that it only changes the memory_format, but not the semantics of each dimensions. + This function is used to facilitate the computation to adopt NHWC kernels, which + provides considerable speed up for fp16 data on CUDA devices with compute capability >= 7.0 + + .. note:: + Calling ``model.to(memory_format=torch.channels_last_3d)`` is more aggressive + than the utility function ``convert_conv3d_weight_memory_format``. Any + layer with 4d weight will be affected by ``model.to``, which does not + necessarily benefit from conversion to specified ``memory_format``. + One place we are confident in is that NDHWC(channels_last_3d) conversion for + convolution in cuDNN, as it is beneficial to run convolution in NDHWC, + even in cases where we have to apply permutation to input tensors. + + Hence our strategy here is to convert only the weight of convolution to + channels_last_3d. This ensures that; + 1. Fast convolution kernels will be used, the benefit of which could + outweigh overhead of permutation (if input is not in the same format). + 2. No unnecessary permutations are applied on layers that do not benefit + from memory_format conversion. + + The optimal case is that, layers between convolution layers are channels + last compatible. Input tensor would be permuted to channels last when it + encounters the first convolution layer and stay in that memory format. + Hence following convolutions will not need to permute its input tensor. + + In case where a channels last incompatible layer is between convolution + layers, we need to permute the input tensor back to contiguous format + for that layer. The input tensor will go through the remaining layers in + contiguous format and be permuted to channels last when it encounters + another convolution layer. There's no point in propagating that + permutation to an earlier layer, as most layers are quite agnostic to + ``memory_format``. + + This claim might change when PyTorch supports fusion of permutation, as + there might have been a better spot to fuse the permutation other than + immediately before a convolution. + + Args: + module (nn.Module): ``nn.Conv3d`` & ``nn.ConvTranspose3d`` or container + ``nn.Module`` + memory_format: user specified ``memory_format``, + e.g. ``torch.channels_last`` or ``torch.contiguous_format`` + + Returns: + The original module with updated ``nn.Conv3d`` + + Example: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> # xdoctest: +REQUIRES(env:CUBLAS_WORKSPACE_CONFIG) + >>> input = torch.randint( + ... 1, 10, (2, 8, 4, 4, 4), dtype=torch.float16, device="cuda" + ... ) + >>> model = nn.Sequential( + >>> nn.Conv3d(8, 4, 3)).cuda().half() + >>> # This is identical to: + >>> # nn.utils.convert_conv3d_weight_memory_format(model, torch.channels_last_3d) + >>> model = nn.utils.convert_conv3d_weight_memory_format( + ... model, torch.channels_last_3d + ... ) + >>> out = model(input) + """ + + # TODO: expand this to `_ConvNd` when channels_last support is extended + # beyond only 4d tensors. + if isinstance(module, (torch.nn.Conv3d, torch.nn.ConvTranspose3d)): + weight_data = module.weight.detach().clone(memory_format=memory_format) + module.weight.data = weight_data.resize_( + weight_data.size(), memory_format=memory_format + ) + for child in module.children(): + convert_conv3d_weight_memory_format(child, memory_format) + + return module + + +__all__ = [ + "convert_conv2d_weight_memory_format", + "convert_conv3d_weight_memory_format", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/parametrizations.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/parametrizations.py new file mode 100644 index 0000000000000000000000000000000000000000..b0eb0d77220f9b35b672d37aca6c3efdb9566929 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/parametrizations.py @@ -0,0 +1,636 @@ +# mypy: allow-untyped-defs +from enum import auto, Enum + +import torch +import torch.nn.functional as F +from torch import Tensor +from torch.nn.modules import Module +from torch.nn.utils import parametrize + + +__all__ = ["orthogonal", "spectral_norm", "weight_norm"] + + +def _is_orthogonal(Q, eps=None): + n, k = Q.size(-2), Q.size(-1) + Id = torch.eye(k, dtype=Q.dtype, device=Q.device) + # A reasonable eps, but not too large + eps = 10.0 * n * torch.finfo(Q.dtype).eps + return torch.allclose(Q.mH @ Q, Id, atol=eps) + + +def _make_orthogonal(A): + """Assume that A is a tall matrix. + + Compute the Q factor s.t. A = QR (A may be complex) and diag(R) is real and non-negative. + """ + X, tau = torch.geqrf(A) + Q = torch.linalg.householder_product(X, tau) + # The diagonal of X is the diagonal of R (which is always real) so we normalise by its signs + Q *= X.diagonal(dim1=-2, dim2=-1).sgn().unsqueeze(-2) + return Q + + +class _OrthMaps(Enum): + matrix_exp = auto() + cayley = auto() + householder = auto() + + +class _Orthogonal(Module): + base: Tensor + + def __init__( + self, weight, orthogonal_map: _OrthMaps, *, use_trivialization=True + ) -> None: + super().__init__() + + # Note [Householder complex] + # For complex tensors, it is not possible to compute the tensor `tau` necessary for + # linalg.householder_product from the reflectors. + # To see this, note that the reflectors have a shape like: + # 0 0 0 + # * 0 0 + # * * 0 + # which, for complex matrices, give n(n-1) (real) parameters. Now, you need n^2 parameters + # to parametrize the unitary matrices. Saving tau on its own does not work either, because + # not every combination of `(A, tau)` gives a unitary matrix, meaning that if we optimise + # them as independent tensors we would not maintain the constraint + # An equivalent reasoning holds for rectangular matrices + if weight.is_complex() and orthogonal_map == _OrthMaps.householder: + raise ValueError( + "The householder parametrization does not support complex tensors." + ) + + self.shape = weight.shape + self.orthogonal_map = orthogonal_map + if use_trivialization: + self.register_buffer("base", None) + + def forward(self, X: torch.Tensor) -> torch.Tensor: + n, k = X.size(-2), X.size(-1) + transposed = n < k + if transposed: + X = X.mT + n, k = k, n + # Here n > k and X is a tall matrix + if ( + self.orthogonal_map == _OrthMaps.matrix_exp + or self.orthogonal_map == _OrthMaps.cayley + ): + # We just need n x k - k(k-1)/2 parameters + X = X.tril() + if n != k: + # Embed into a square matrix + X = torch.cat( + [X, X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1 + ) + A = X - X.mH + # A is skew-symmetric (or skew-hermitian) + if self.orthogonal_map == _OrthMaps.matrix_exp: + Q = torch.matrix_exp(A) + elif self.orthogonal_map == _OrthMaps.cayley: + # Computes the Cayley retraction (I+A/2)(I-A/2)^{-1} + Id = torch.eye(n, dtype=A.dtype, device=A.device) + Q = torch.linalg.solve( + torch.add(Id, A, alpha=-0.5), torch.add(Id, A, alpha=0.5) + ) + # Q is now orthogonal (or unitary) of size (..., n, n) + if n != k: + # pyrefly: ignore [unbound-name] + Q = Q[..., :k] + # Q is now the size of the X (albeit perhaps transposed) + else: + # X is real here, as we do not support householder with complex numbers + A = X.tril(diagonal=-1) + tau = 2.0 / (1.0 + (A * A).sum(dim=-2)) + Q = torch.linalg.householder_product(A, tau) + # The diagonal of X is 1's and -1's + # We do not want to differentiate through this or update the diagonal of X hence the casting + Q = Q * X.diagonal(dim1=-2, dim2=-1).int().unsqueeze(-2) + + if hasattr(self, "base"): + # pyrefly: ignore [unbound-name] + Q = self.base @ Q + if transposed: + # pyrefly: ignore [unbound-name] + Q = Q.mT + return Q # type: ignore[possibly-undefined] + + @torch.autograd.no_grad() + def right_inverse(self, Q: torch.Tensor) -> torch.Tensor: + if Q.shape != self.shape: + raise ValueError( + f"Expected a matrix or batch of matrices of shape {self.shape}. " + f"Got a tensor of shape {Q.shape}." + ) + + Q_init = Q + n, k = Q.size(-2), Q.size(-1) + transpose = n < k + if transpose: + Q = Q.mT + n, k = k, n + + # We always make sure to always copy Q in every path + if not hasattr(self, "base"): + # Note [right_inverse expm cayley] + # If we do not have use_trivialization=True, we just implement the inverse of the forward + # map for the Householder. To see why, think that for the Cayley map, + # we would need to find the matrix X \in R^{n x k} such that: + # Y = torch.cat([X.tril(), X.new_zeros(n, n - k).expand(*X.shape[:-2], -1, -1)], dim=-1) + # A = Y - Y.mH + # cayley(A)[:, :k] + # gives the original tensor. It is not clear how to do this. + # Perhaps via some algebraic manipulation involving the QR like that of + # Corollary 2.2 in Edelman, Arias and Smith? + if ( + self.orthogonal_map == _OrthMaps.cayley + or self.orthogonal_map == _OrthMaps.matrix_exp + ): + raise NotImplementedError( + "It is not possible to assign to the matrix exponential " + "or the Cayley parametrizations when use_trivialization=False." + ) + + # If parametrization == _OrthMaps.householder, make Q orthogonal via the QR decomposition. + # Here Q is always real because we do not support householder and complex matrices. + # See note [Householder complex] + A, tau = torch.geqrf(Q) + # We want to have a decomposition X = QR with diag(R) > 0, as otherwise we could + # decompose an orthogonal matrix Q as Q = (-Q)@(-Id), which is a valid QR decomposition + # The diagonal of Q is the diagonal of R from the qr decomposition + A.diagonal(dim1=-2, dim2=-1).sign_() + # Equality with zero is ok because LAPACK returns exactly zero when it does not want + # to use a particular reflection + A.diagonal(dim1=-2, dim2=-1)[tau == 0.0] *= -1 + return A.mT if transpose else A + else: + if n == k: + # We check whether Q is orthogonal + if not _is_orthogonal(Q): + Q = _make_orthogonal(Q) + else: # Is orthogonal + Q = Q.clone() + else: + # Complete Q into a full n x n orthogonal matrix + N = torch.randn( + *(Q.size()[:-2] + (n, n - k)), dtype=Q.dtype, device=Q.device + ) + Q = torch.cat([Q, N], dim=-1) + Q = _make_orthogonal(Q) + self.base = Q + + # It is necessary to return the -Id, as we use the diagonal for the + # Householder parametrization. Using -Id makes: + # householder(torch.zeros(m,n)) == torch.eye(m,n) + # Poor man's version of eye_like + neg_Id = torch.zeros_like(Q_init) + neg_Id.diagonal(dim1=-2, dim2=-1).fill_(-1.0) + return neg_Id + + +def orthogonal( + module: Module, + name: str = "weight", + orthogonal_map: str | None = None, + *, + use_trivialization: bool = True, +) -> Module: + r"""Apply an orthogonal or unitary parametrization to a matrix or a batch of matrices. + + Letting :math:`\mathbb{K}` be :math:`\mathbb{R}` or :math:`\mathbb{C}`, the parametrized + matrix :math:`Q \in \mathbb{K}^{m \times n}` is **orthogonal** as + + .. math:: + + \begin{align*} + Q^{\text{H}}Q &= \mathrm{I}_n \mathrlap{\qquad \text{if }m \geq n}\\ + QQ^{\text{H}} &= \mathrm{I}_m \mathrlap{\qquad \text{if }m < n} + \end{align*} + + where :math:`Q^{\text{H}}` is the conjugate transpose when :math:`Q` is complex + and the transpose when :math:`Q` is real-valued, and + :math:`\mathrm{I}_n` is the `n`-dimensional identity matrix. + In plain words, :math:`Q` will have orthonormal columns whenever :math:`m \geq n` + and orthonormal rows otherwise. + + If the tensor has more than two dimensions, we consider it as a batch of matrices of shape `(..., m, n)`. + + The matrix :math:`Q` may be parametrized via three different ``orthogonal_map`` in terms of the original tensor: + + - ``"matrix_exp"``/``"cayley"``: + the :func:`~torch.matrix_exp` :math:`Q = \exp(A)` and the `Cayley map`_ + :math:`Q = (\mathrm{I}_n + A/2)(\mathrm{I}_n - A/2)^{-1}` are applied to a skew-symmetric + :math:`A` to give an orthogonal matrix. + - ``"householder"``: computes a product of Householder reflectors + (:func:`~torch.linalg.householder_product`). + + ``"matrix_exp"``/``"cayley"`` often make the parametrized weight converge faster than + ``"householder"``, but they are slower to compute for very thin or very wide matrices. + + If ``use_trivialization=True`` (default), the parametrization implements the "Dynamic Trivialization Framework", + where an extra matrix :math:`B \in \mathbb{K}^{n \times n}` is stored under + ``module.parametrizations.weight[0].base``. This helps the + convergence of the parametrized layer at the expense of some extra memory use. + See `Trivializations for Gradient-Based Optimization on Manifolds`_ . + + Initial value of :math:`Q`: + If the original tensor is not parametrized and ``use_trivialization=True`` (default), the initial value + of :math:`Q` is that of the original tensor if it is orthogonal (or unitary in the complex case) + and it is orthogonalized via the QR decomposition otherwise (see :func:`torch.linalg.qr`). + Same happens when it is not parametrized and ``orthogonal_map="householder"`` even when ``use_trivialization=False``. + Otherwise, the initial value is the result of the composition of all the registered + parametrizations applied to the original tensor. + + .. note:: + This function is implemented using the parametrization functionality + in :func:`~torch.nn.utils.parametrize.register_parametrization`. + + + .. _`Cayley map`: https://en.wikipedia.org/wiki/Cayley_transform#Matrix_map + .. _`Trivializations for Gradient-Based Optimization on Manifolds`: https://arxiv.org/abs/1909.09501 + + Args: + module (nn.Module): module on which to register the parametrization. + name (str, optional): name of the tensor to make orthogonal. Default: ``"weight"``. + orthogonal_map (str, optional): One of the following: ``"matrix_exp"``, ``"cayley"``, ``"householder"``. + Default: ``"matrix_exp"`` if the matrix is square or complex, ``"householder"`` otherwise. + use_trivialization (bool, optional): whether to use the dynamic trivialization framework. + Default: ``True``. + + Returns: + The original module with an orthogonal parametrization registered to the specified + weight + + Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK) + >>> orth_linear = orthogonal(nn.Linear(20, 40)) + >>> orth_linear + ParametrizedLinear( + in_features=20, out_features=40, bias=True + (parametrizations): ModuleDict( + (weight): ParametrizationList( + (0): _Orthogonal() + ) + ) + ) + >>> # xdoctest: +IGNORE_WANT + >>> Q = orth_linear.weight + >>> torch.dist(Q.T @ Q, torch.eye(20)) + tensor(4.9332e-07) + """ + weight = getattr(module, name, None) + if not isinstance(weight, Tensor): + raise ValueError( + f"Module '{module}' has no parameter or buffer with name '{name}'" + ) + + # We could implement this for 1-dim tensors as the maps on the sphere + # but I believe it'd bite more people than it'd help + if weight.ndim < 2: + raise ValueError( + "Expected a matrix or batch of matrices. " + f"Got a tensor of {weight.ndim} dimensions." + ) + + if orthogonal_map is None: + orthogonal_map = ( + "matrix_exp" + if weight.size(-2) == weight.size(-1) or weight.is_complex() + else "householder" + ) + + orth_enum = getattr(_OrthMaps, orthogonal_map, None) + if orth_enum is None: + raise ValueError( + 'orthogonal_map has to be one of "matrix_exp", "cayley", "householder". ' + f"Got: {orthogonal_map}" + ) + orth = _Orthogonal(weight, orth_enum, use_trivialization=use_trivialization) + parametrize.register_parametrization(module, name, orth, unsafe=True) + return module + + +class _WeightNorm(Module): + def __init__( + self, + dim: int | None = 0, + ) -> None: + super().__init__() + if dim is None: + dim = -1 + self.dim = dim + + def forward(self, weight_g, weight_v): + return torch._weight_norm(weight_v, weight_g, self.dim) + + def right_inverse(self, weight): + weight_g = torch.norm_except_dim(weight, 2, self.dim) + weight_v = weight + + return weight_g, weight_v + + +def weight_norm(module: Module, name: str = "weight", dim: int = 0): + r"""Apply weight normalization to a parameter in the given module. + + .. math:: + \mathbf{w} = g \dfrac{\mathbf{v}}{\|\mathbf{v}\|} + + Weight normalization is a reparameterization that decouples the magnitude + of a weight tensor from its direction. This replaces the parameter specified + by :attr:`name` with two parameters: one specifying the magnitude + and one specifying the direction. + + By default, with ``dim=0``, the norm is computed independently per output + channel/plane. To compute a norm over the entire weight tensor, use + ``dim=None``. + + See https://arxiv.org/abs/1602.07868 + + Args: + module (Module): containing module + name (str, optional): name of weight parameter + dim (int, optional): dimension over which to compute the norm + + Returns: + The original module with the weight norm hook + + Example:: + + >>> m = weight_norm(nn.Linear(20, 40), name='weight') + >>> m + ParametrizedLinear( + in_features=20, out_features=40, bias=True + (parametrizations): ModuleDict( + (weight): ParametrizationList( + (0): _WeightNorm() + ) + ) + ) + >>> m.parametrizations.weight.original0.size() + torch.Size([40, 1]) + >>> m.parametrizations.weight.original1.size() + torch.Size([40, 20]) + + """ + _weight_norm = _WeightNorm(dim) + parametrize.register_parametrization(module, name, _weight_norm, unsafe=True) + + def _weight_norm_compat_hook( + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) -> None: + g_key = f"{prefix}{name}_g" + v_key = f"{prefix}{name}_v" + if g_key in state_dict and v_key in state_dict: + original0 = state_dict.pop(g_key) + original1 = state_dict.pop(v_key) + state_dict[f"{prefix}parametrizations.{name}.original0"] = original0 + state_dict[f"{prefix}parametrizations.{name}.original1"] = original1 + + module._register_load_state_dict_pre_hook(_weight_norm_compat_hook) + return module + + +class _SpectralNorm(Module): + def __init__( + self, + weight: torch.Tensor, + n_power_iterations: int = 1, + dim: int = 0, + eps: float = 1e-12, + ) -> None: + super().__init__() + ndim = weight.ndim + if dim >= ndim or dim < -ndim: + raise IndexError( + "Dimension out of range (expected to be in range of " + f"[-{ndim}, {ndim - 1}] but got {dim})" + ) + + if n_power_iterations <= 0: + raise ValueError( + "Expected n_power_iterations to be positive, but " + f"got n_power_iterations={n_power_iterations}" + ) + self.dim = dim if dim >= 0 else dim + ndim + self.eps = eps + if ndim > 1: + # For ndim == 1 we do not need to approximate anything (see _SpectralNorm.forward) + self.n_power_iterations = n_power_iterations + weight_mat = self._reshape_weight_to_matrix(weight) + h, w = weight_mat.size() + + u = weight_mat.new_empty(h).normal_(0, 1) + v = weight_mat.new_empty(w).normal_(0, 1) + self.register_buffer("_u", F.normalize(u, dim=0, eps=self.eps)) + self.register_buffer("_v", F.normalize(v, dim=0, eps=self.eps)) + + # Start with u, v initialized to some reasonable values by performing a number + # of iterations of the power method + self._power_method(weight_mat, 15) + + def _reshape_weight_to_matrix(self, weight: torch.Tensor) -> torch.Tensor: + # Precondition + if weight.ndim <= 1: + raise AssertionError( + f"Expected weight to have more than 1 dimension, got {weight.ndim}" + ) + + if self.dim != 0: + # permute dim to front + weight = weight.permute( + self.dim, *(d for d in range(weight.dim()) if d != self.dim) + ) + + return weight.flatten(1) + + @torch.autograd.no_grad() + def _power_method(self, weight_mat: torch.Tensor, n_power_iterations: int) -> None: + # See original note at torch/nn/utils/spectral_norm.py + # NB: If `do_power_iteration` is set, the `u` and `v` vectors are + # updated in power iteration **in-place**. This is very important + # because in `DataParallel` forward, the vectors (being buffers) are + # broadcast from the parallelized module to each module replica, + # which is a new module object created on the fly. And each replica + # runs its own spectral norm power iteration. So simply assigning + # the updated vectors to the module this function runs on will cause + # the update to be lost forever. And the next time the parallelized + # module is replicated, the same randomly initialized vectors are + # broadcast and used! + # + # Therefore, to make the change propagate back, we rely on two + # important behaviors (also enforced via tests): + # 1. `DataParallel` doesn't clone storage if the broadcast tensor + # is already on correct device; and it makes sure that the + # parallelized module is already on `device[0]`. + # 2. If the out tensor in `out=` kwarg has correct shape, it will + # just fill in the values. + # Therefore, since the same power iteration is performed on all + # devices, simply updating the tensors in-place will make sure that + # the module replica on `device[0]` will update the _u vector on the + # parallelized module (by shared storage). + # + # However, after we update `u` and `v` in-place, we need to **clone** + # them before using them to normalize the weight. This is to support + # backproping through two forward passes, e.g., the common pattern in + # GAN training: loss = D(real) - D(fake). Otherwise, engine will + # complain that variables needed to do backward for the first forward + # (i.e., the `u` and `v` vectors) are changed in the second forward. + + # Precondition + if weight_mat.ndim <= 1: + raise AssertionError( + f"Expected weight_mat to have more than 1 dimension, got {weight_mat.ndim}" + ) + + for _ in range(n_power_iterations): + # Spectral norm of weight equals to `u^T W v`, where `u` and `v` + # are the first left and right singular vectors. + # This power iteration produces approximations of `u` and `v`. + self._u = F.normalize( + torch.mv(weight_mat, self._v), # type: ignore[has-type] + dim=0, + eps=self.eps, + out=self._u, # type: ignore[has-type] + ) + self._v = F.normalize( + torch.mv(weight_mat.H, self._u), # type: ignore[has-type] + dim=0, + eps=self.eps, + out=self._v, # type: ignore[has-type] + ) + + def forward(self, weight: torch.Tensor) -> torch.Tensor: + if weight.ndim == 1: + # Faster and more exact path, no need to approximate anything + return F.normalize(weight, dim=0, eps=self.eps) + else: + weight_mat = self._reshape_weight_to_matrix(weight) + if self.training: + self._power_method(weight_mat, self.n_power_iterations) + # See above on why we need to clone + u = self._u.clone(memory_format=torch.contiguous_format) + v = self._v.clone(memory_format=torch.contiguous_format) + # The proper way of computing this should be through F.bilinear, but + # it seems to have some efficiency issues: + # https://github.com/pytorch/pytorch/issues/58093 + sigma = torch.vdot(u, torch.mv(weight_mat, v)) + return weight / sigma + + def right_inverse(self, value: torch.Tensor) -> torch.Tensor: + # we may want to assert here that the passed value already + # satisfies constraints + return value + + +def spectral_norm( + module: Module, + name: str = "weight", + n_power_iterations: int = 1, + eps: float = 1e-12, + dim: int | None = None, +) -> Module: + r"""Apply spectral normalization to a parameter in the given module. + + .. math:: + \mathbf{W}_{SN} = \dfrac{\mathbf{W}}{\sigma(\mathbf{W})}, + \sigma(\mathbf{W}) = \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2} + + When applied on a vector, it simplifies to + + .. math:: + \mathbf{x}_{SN} = \dfrac{\mathbf{x}}{\|\mathbf{x}\|_2} + + Spectral normalization stabilizes the training of discriminators (critics) + in Generative Adversarial Networks (GANs) by reducing the Lipschitz constant + of the model. :math:`\sigma` is approximated performing one iteration of the + `power method`_ every time the weight is accessed. If the dimension of the + weight tensor is greater than 2, it is reshaped to 2D in power iteration + method to get spectral norm. + + + See `Spectral Normalization for Generative Adversarial Networks`_ . + + .. _`power method`: https://en.wikipedia.org/wiki/Power_iteration + .. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957 + + .. note:: + This function is implemented using the parametrization functionality + in :func:`~torch.nn.utils.parametrize.register_parametrization`. It is a + reimplementation of :func:`torch.nn.utils.spectral_norm`. + + .. note:: + When this constraint is registered, the singular vectors associated to the largest + singular value are estimated rather than sampled at random. These are then updated + performing :attr:`n_power_iterations` of the `power method`_ whenever the tensor + is accessed with the module on `training` mode. + + .. note:: + If the `_SpectralNorm` module, i.e., `module.parametrization.weight[idx]`, + is in training mode on removal, it will perform another power iteration. + If you'd like to avoid this iteration, set the module to eval mode + before its removal. + + Args: + module (nn.Module): containing module + name (str, optional): name of weight parameter. Default: ``"weight"``. + n_power_iterations (int, optional): number of power iterations to + calculate spectral norm. Default: ``1``. + eps (float, optional): epsilon for numerical stability in + calculating norms. Default: ``1e-12``. + dim (int, optional): dimension corresponding to number of outputs. + Default: ``0``, except for modules that are instances of + ConvTranspose{1,2,3}d, when it is ``1`` + + Returns: + The original module with a new parametrization registered to the specified + weight + + Example:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK) + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> snm = spectral_norm(nn.Linear(20, 40)) + >>> snm + ParametrizedLinear( + in_features=20, out_features=40, bias=True + (parametrizations): ModuleDict( + (weight): ParametrizationList( + (0): _SpectralNorm() + ) + ) + ) + >>> torch.linalg.matrix_norm(snm.weight, 2) + tensor(1.0081, grad_fn=) + """ + weight = getattr(module, name, None) + if not isinstance(weight, Tensor): + raise ValueError( + f"Module '{module}' has no parameter or buffer with name '{name}'" + ) + + if dim is None: + if isinstance( + module, + ( + torch.nn.ConvTranspose1d, + torch.nn.ConvTranspose2d, + torch.nn.ConvTranspose3d, + ), + ): + dim = 1 + else: + dim = 0 + parametrize.register_parametrization( + module, name, _SpectralNorm(weight, n_power_iterations, dim, eps) + ) + return module diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/parametrize.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/parametrize.py new file mode 100644 index 0000000000000000000000000000000000000000..81a3ad796abebdd50289aa6088d630b49341b496 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/parametrize.py @@ -0,0 +1,886 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +import collections +import copyreg +from collections.abc import Sequence +from contextlib import contextmanager +from copy import deepcopy + +import torch +from torch import Tensor +from torch.__future__ import get_swap_module_params_on_conversion +from torch._library.opaque_object import is_opaque_reference_type +from torch._opaque_base import OpaqueBase +from torch.nn.modules.container import Module, ModuleDict, ModuleList +from torch.nn.parameter import Parameter +from torch.utils._python_dispatch import is_traceable_wrapper_subclass + + +__all__ = [ + "cached", + "ParametrizationList", + "register_parametrization", + "is_parametrized", + "remove_parametrizations", + "type_before_parametrizations", + "transfer_parametrizations_and_params", +] + +_cache_enabled = 0 +_cache: dict[tuple[int, str], Tensor | None] = {} + + +@contextmanager +def cached(): + r"""Context manager that enables the caching system within parametrizations registered with :func:`register_parametrization`. + + The value of the parametrized objects is computed and cached the first time + they are required when this context manager is active. The cached values are + discarded when leaving the context manager. + + This is useful when using a parametrized parameter more than once in the forward pass. + An example of this is when parametrizing the recurrent kernel of an RNN or when + sharing weights. + + The simplest way to activate the cache is by wrapping the forward pass of the neural network + + .. code-block:: python + + import torch.nn.utils.parametrize as P + + ... + with P.cached(): + output = model(inputs) + + in training and evaluation. One may also wrap the parts of the modules that use + several times the parametrized tensors. For example, the loop of an RNN with a + parametrized recurrent kernel: + + .. code-block:: python + + with P.cached(): + for x in xs: + out_rnn = self.rnn_cell(x, out_rnn) + """ + global _cache + global _cache_enabled + _cache_enabled += 1 + try: + yield + finally: + _cache_enabled -= 1 + if not _cache_enabled: + _cache = {} + + +def _register_parameter_or_buffer(module, name, X) -> None: + if isinstance(X, Parameter): + module.register_parameter(name, X) + else: + module.register_buffer(name, X) + + +def _maybe_set(dest: Tensor, src: Tensor) -> None: + should_swap = ( + get_swap_module_params_on_conversion() or is_traceable_wrapper_subclass(dest) + ) + if should_swap: + if isinstance(dest, Parameter) and not isinstance(src, Parameter): + src = Parameter(src, requires_grad=dest.requires_grad) + torch.utils.swap_tensors(dest, src) + else: + dest.set_(src) # type: ignore[call-overload] + + +class ParametrizationList(ModuleList): + r"""A sequential container that holds and manages the original parameters or buffers of a parametrized :class:`torch.nn.Module`. + + It is the type of ``module.parametrizations[tensor_name]`` when ``module[tensor_name]`` + has been parametrized with :func:`register_parametrization`. + + If the first registered parametrization has a ``right_inverse`` that returns one tensor or + does not have a ``right_inverse`` (in which case we assume that ``right_inverse`` is the identity), + it will hold the tensor under the name ``original``. + If it has a ``right_inverse`` that returns more than one tensor, these will be registered as + ``original0``, ``original1``, ... + + .. warning:: + This class is used internally by :func:`register_parametrization`. It is documented + here for completeness. It shall not be instantiated by the user. + + Args: + modules (sequence): sequence of modules representing the parametrizations + original (Parameter or Tensor): parameter or buffer that is parametrized + unsafe (bool): a boolean flag that denotes whether the parametrization + may change the dtype and shape of the tensor. Default: `False` + Warning: the parametrization is not checked for consistency upon registration. + Enable this flag at your own risk. + """ + + original: Tensor + unsafe: bool + + def __init__( + self, + modules: Sequence[Module], + original: Tensor | Parameter, + unsafe: bool = False, + ) -> None: + # We require this because we need to treat differently the first parametrization + # This should never throw, unless this class is used from the outside + if len(modules) == 0: + raise ValueError("ParametrizationList requires one or more modules.") + + super().__init__(modules) + self.unsafe = unsafe + + # In plain words: + # module.weight must keep its dtype and shape. + # Furthermore, if there is no right_inverse or the right_inverse returns a tensor, + # this should be of the same dtype as the original tensor + # + # We check that the following invariants hold: + # X = module.weight + # Y = param.right_inverse(X) + # assert isinstance(Y, Tensor) or + # (isinstance(Y, collections.abc.Sequence) and all(isinstance(t, Tensor) for t in Y)) + # Z = param(Y) if isinstance(Y, Tensor) else param(*Y) + # # Consistency checks + # assert X.dtype == Z.dtype and X.shape == Z.shape + # # If it has one input, this allows to be able to use set_ to be able to + # # move data to/from the original tensor without changing its id (which is what the + # # optimizer uses to track parameters) + # if isinstance(Y, Tensor) + # assert X.dtype == Y.dtype + # Below we use original = X, new = Y + + original_shape = original.shape + original_dtype = original.dtype + + # Compute new + with torch.no_grad(): + new = original + for module in reversed(self): # type: ignore[call-overload] + if hasattr(module, "right_inverse"): + try: + new = module.right_inverse(new) # type: ignore[operator] + except NotImplementedError: + pass + # else, or if it throws, we assume that right_inverse is the identity + + if not isinstance(new, Tensor) and not isinstance(new, Sequence): + raise ValueError( + "'right_inverse' must return a Tensor or a Sequence of tensors (list, tuple...). " + f"Got {type(new).__name__}" + ) + + # Set the number of original tensors + self.is_tensor = isinstance(new, Tensor) + self.ntensors = 1 if self.is_tensor else len(new) + + # Register the tensor(s) + if self.is_tensor: + # pyrefly: ignore [missing-attribute] + if original.dtype != new.dtype: + raise ValueError( + "When `right_inverse` outputs one tensor, it may not change the dtype.\n" + f"original.dtype: {original.dtype}\n" + # pyrefly: ignore [missing-attribute] + f"right_inverse(original).dtype: {new.dtype}" + ) + + # pyrefly: ignore [missing-attribute] + if original.device != new.device: + raise ValueError( + "When `right_inverse` outputs one tensor, it may not change the device.\n" + f"original.device: {original.device}\n" + # pyrefly: ignore [missing-attribute] + f"right_inverse(original).device: {new.device}" + ) + + # Set the original to original so that the user does not need to re-register the parameter + # manually in the optimiser + with torch.no_grad(): + # pyrefly: ignore [bad-argument-type] + _maybe_set(original, new) + _register_parameter_or_buffer(self, "original", original) + else: + for i, originali in enumerate(new): + match originali: + case OpaqueBase(): + if not is_opaque_reference_type(type(originali)): + raise ValueError( + f"'right_inverse' must return a Tensor or a reference-type " + f"opaque. Got element {i} of the sequence with type " + f"{type(originali).__name__}." + ) + setattr(self, f"original{i}", originali) + case Tensor(): + # If the original tensor was a Parameter that required grad, we expect the user to + # add the new parameters to the optimizer after registering the parametrization + # (this is documented) + if isinstance(original, Parameter): + originali = Parameter(originali, original.requires_grad) + originali.requires_grad_(original.requires_grad) + _register_parameter_or_buffer(self, f"original{i}", originali) + case _: + raise ValueError( + "'right_inverse' must return a Tensor or a Sequence of tensors " + "(list, tuple...). " + f"Got element {i} of the sequence with type {type(originali).__name__}." + ) + + if not self.unsafe: + # Consistency checks: + # Since f : A -> B, right_inverse : B -> A, Z and original should live in B + # Z = forward(right_inverse(original)) + Z = self() + if not isinstance(Z, Tensor): + raise ValueError( + f"A parametrization must return a tensor. Got {type(Z).__name__}." + ) + if Z.dtype != original_dtype: + raise ValueError( + "Registering a parametrization may not change the dtype of the tensor, unless `unsafe` flag is enabled.\n" + f"unparametrized dtype: {original_dtype}\n" + f"parametrized dtype: {Z.dtype}" + ) + if Z.shape != original_shape: + raise ValueError( + "Registering a parametrization may not change the shape of the tensor, unless `unsafe` flag is enabled.\n" + f"unparametrized shape: {original_shape}\n" + f"parametrized shape: {Z.shape}" + ) + + def right_inverse(self, value: Tensor) -> None: + r"""Call the ``right_inverse`` methods of the parametrizations in the inverse registration order. + + Then, it stores the result in ``self.original`` if ``right_inverse`` outputs one tensor + or in ``self.original0``, ``self.original1``, ... if it outputs several. + + Args: + value (Tensor): Value to which initialize the module + """ + # All the exceptions in this function should almost never throw. + # They could throw if, for example, right_inverse function returns a different + # dtype when given a different input, which should most likely be caused by a + # bug in the user's code + + with torch.no_grad(): + # See https://github.com/pytorch/pytorch/issues/53103 + for module in reversed(self): # type: ignore[call-overload] + if hasattr(module, "right_inverse"): + value = module.right_inverse(value) # type: ignore[operator] + else: + raise RuntimeError( + f"parametrization {type(module).__name__} does not implement " + "right_inverse." + ) + if self.is_tensor: + # These exceptions should only throw when a right_inverse function does not + # return the same dtype for every input, which should most likely be caused by a bug + if not isinstance(value, Tensor): + raise ValueError( + f"`right_inverse` should return a tensor. Got {type(value).__name__}" + ) + if value.dtype != self.original.dtype: + raise ValueError( + f"The tensor returned by `right_inverse` has dtype {value.dtype} " + f"while `original` has dtype {self.original.dtype}" + ) + # We know that the result is going to have the same dtype + _maybe_set(self.original, value) + else: + if not isinstance(value, collections.abc.Sequence): + raise ValueError( + "'right_inverse' must return a sequence of tensors. " + f"Got {type(value).__name__}." + ) + if len(value) != self.ntensors: + raise ValueError( + "'right_inverse' must return a sequence of tensors of length " + f"{self.ntensors}. Got a sequence of length {len(value)}." + ) + for i, tensor in enumerate(value): + original_i = getattr(self, f"original{i}") + match tensor: + case OpaqueBase(): + if is_opaque_reference_type(type(tensor)): + setattr(self, f"original{i}", tensor) + continue + # Fall-through + case Tensor(): + if original_i.dtype != tensor.dtype: + raise ValueError( + f"Tensor {i} returned by `right_inverse` has dtype {tensor.dtype} " + f"while `original{i}` has dtype {original_i.dtype}" + ) + _maybe_set(original_i, tensor) + continue + raise ValueError( + f"'right_inverse' must return a sequence of tensors " + f"or reference-type opaques. Got element {i} of type " + f"{type(tensor).__name__}." + ) + + def forward(self) -> Tensor: + if torch.jit.is_scripting(): + raise RuntimeError("Parametrization is not working with scripting.") + # Unpack the originals for the first parametrization + if self.is_tensor: + x = self[0](self.original) + else: + originals = (getattr(self, f"original{i}") for i in range(self.ntensors)) + x = self[0](*originals) + # It's not possible to call self[1:] here, so we have to be a bit more cryptic + # Also we want to skip all non-integer keys + curr_idx = 1 + while hasattr(self, str(curr_idx)): + x = self[curr_idx](x) + curr_idx += 1 + return x + + +def _inject_new_class(module: Module) -> None: + r"""Set up a module to be parametrized. + + This works by substituting the class of the module by a class + that extends it to be able to inject a property + + Args: + module (nn.Module): module into which to inject the property + """ + cls = module.__class__ + + def default_deepcopy(self, memo): + # Just emulate a standard deepcopy procedure when __deepcopy__ doesn't exist in the current class. + obj = memo.get(id(self), None) + if obj is not None: + return obj + replica = self.__new__(self.__class__) + memo[id(self)] = replica + replica.__dict__ = deepcopy(self.__dict__, memo) + # Also save all slots if they exist. + slots_to_save = copyreg._slotnames(self.__class__) # type: ignore[attr-defined] + for slot in slots_to_save: + if hasattr(self, slot): + setattr(replica, slot, deepcopy(getattr(self, slot), memo)) + return replica + + def getstate(self): + raise RuntimeError( + "Serialization of parametrized modules is only " + "supported through state_dict(). See:\n" + "https://pytorch.org/tutorials/beginner/saving_loading_models.html" + "#saving-loading-a-general-checkpoint-for-inference-and-or-resuming-training" + ) + + dct = {"__getstate__": getstate} + # We don't allow serialization of parametrized modules but should still allow deepcopying. + # Default 'deepcopy' function invokes __deepcopy__ method instead of __getstate__ when it exists. + if not hasattr(cls, "__deepcopy__"): + dct["__deepcopy__"] = default_deepcopy # type: ignore[assignment] + + param_cls = type( + f"Parametrized{cls.__name__}", + (cls,), + dct, + ) + + module.__class__ = param_cls + + +def _inject_property(module: Module, tensor_name: str) -> None: + r"""Injects a property into module[tensor_name]. + + It assumes that the class in the module has already been modified from its + original one using _inject_new_class and that the tensor under :attr:`tensor_name` + has already been moved out + + Args: + module (nn.Module): module into which to inject the property + tensor_name (str): name of the name of the property to create + """ + # We check the precondition. + # This should never fire if register_parametrization is correctly implemented + if hasattr(module, tensor_name): + raise AssertionError(f"Module already has an attribute named '{tensor_name}'") + + @torch.jit.unused + def get_cached_parametrization(parametrization) -> Tensor: + global _cache + key = (id(module), tensor_name) + tensor = _cache.get(key) + if tensor is None: + tensor = parametrization() + _cache[key] = tensor + return tensor + + def get_parametrized(self) -> Tensor: + if torch.jit.is_scripting(): + raise RuntimeError("Parametrization is not working with scripting.") + parametrization = self.parametrizations[tensor_name] + # pyrefly: ignore [redundant-condition] + if _cache_enabled: + if torch.jit.is_scripting(): + # Scripting + raise RuntimeError( + "Caching is not implemented for scripting. " + "Either disable caching or avoid scripting." + ) + elif torch._C._get_tracing_state() is not None: + # Tracing + raise RuntimeError( + "Cannot trace a model while caching parametrizations." + ) + else: + return get_cached_parametrization(parametrization) + else: + # If caching is not active, this function just evaluates the parametrization + return parametrization() + + def set_original(self, value: Tensor) -> None: + if torch.jit.is_scripting(): + raise RuntimeError("Parametrization is not working with scripting.") + self.parametrizations[tensor_name].right_inverse(value) + + setattr(module.__class__, tensor_name, property(get_parametrized, set_original)) + + +def register_parametrization( + module: Module, + tensor_name: str, + parametrization: Module, + *, + unsafe: bool = False, +) -> Module: + r"""Register a parametrization to a tensor in a module. + + Assume that ``tensor_name="weight"`` for simplicity. When accessing ``module.weight``, + the module will return the parametrized version ``parametrization(module.weight)``. + If the original tensor requires a gradient, the backward pass will differentiate + through :attr:`parametrization`, and the optimizer will update the tensor accordingly. + + The first time that a module registers a parametrization, this function will add an attribute + ``parametrizations`` to the module of type :class:`~ParametrizationList`. + + The list of parametrizations on the tensor ``weight`` will be accessible under + ``module.parametrizations.weight``. + + The original tensor will be accessible under + ``module.parametrizations.weight.original``. + + Parametrizations may be concatenated by registering several parametrizations + on the same attribute. + + The training mode of a registered parametrization is updated on registration + to match the training mode of the host module + + Parametrized parameters and buffers have an inbuilt caching system that can be activated + using the context manager :func:`cached`. + + A :attr:`parametrization` may optionally implement a method with signature + + .. code-block:: python + + def right_inverse(self, X: Tensor) -> Union[Tensor, Sequence[Tensor]] + + This method is called on the unparametrized tensor when the first parametrization + is registered to compute the initial value of the original tensor. + If this method is not implemented, the original tensor will be just the unparametrized tensor. + + If all the parametrizations registered on a tensor implement `right_inverse` it is possible + to initialize a parametrized tensor by assigning to it, as shown in the example below. + + It is possible for the first parametrization to depend on several inputs. + This may be implemented returning a tuple of tensors from ``right_inverse`` + (see the example implementation of a ``RankOne`` parametrization below). + + In this case, the unconstrained tensors are also located under ``module.parametrizations.weight`` + with names ``original0``, ``original1``,... + + .. note:: + + If unsafe=False (default) both the forward and right_inverse methods will be called + once to perform a number of consistency checks. + If unsafe=True, then right_inverse will be called if the tensor is not parametrized, + and nothing will be called otherwise. + + .. note:: + + In most situations, ``right_inverse`` will be a function such that + ``forward(right_inverse(X)) == X`` (see + `right inverse `_). + Sometimes, when the parametrization is not surjective, it may be reasonable + to relax this. + + .. warning:: + + If a parametrization depends on several inputs, :func:`~register_parametrization` + will register a number of new parameters. If such parametrization is registered + after the optimizer is created, these new parameters will need to be added manually + to the optimizer. See :meth:`torch.Optimizer.add_param_group`. + + Args: + module (nn.Module): module on which to register the parametrization + tensor_name (str): name of the parameter or buffer on which to register + the parametrization + parametrization (nn.Module): the parametrization to register + Keyword args: + unsafe (bool): a boolean flag that denotes whether the parametrization + may change the dtype and shape of the tensor. Default: `False` + Warning: the parametrization is not checked for consistency upon registration. + Enable this flag at your own risk. + + Raises: + ValueError: if the module does not have a parameter or a buffer named :attr:`tensor_name` + + Examples: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_LAPACK) + >>> import torch + >>> import torch.nn as nn + >>> import torch.nn.utils.parametrize as P + >>> + >>> class Symmetric(nn.Module): + >>> def forward(self, X): + >>> return X.triu() + X.triu(1).T # Return a symmetric matrix + >>> + >>> def right_inverse(self, A): + >>> return A.triu() + >>> + >>> m = nn.Linear(5, 5) + >>> P.register_parametrization(m, "weight", Symmetric()) + >>> print(torch.allclose(m.weight, m.weight.T)) # m.weight is now symmetric + True + >>> A = torch.rand(5, 5) + >>> A = A + A.T # A is now symmetric + >>> m.weight = A # Initialize the weight to be the symmetric matrix A + >>> print(torch.allclose(m.weight, A)) + True + + >>> class RankOne(nn.Module): + >>> def forward(self, x, y): + >>> # Form a rank 1 matrix multiplying two vectors + >>> return x.unsqueeze(-1) @ y.unsqueeze(-2) + >>> + >>> def right_inverse(self, Z): + >>> # Project Z onto the rank 1 matrices + >>> U, S, Vh = torch.linalg.svd(Z, full_matrices=False) + >>> # Return rescaled singular vectors + >>> s0_sqrt = S[0].sqrt().unsqueeze(-1) + >>> return U[..., :, 0] * s0_sqrt, Vh[..., 0, :] * s0_sqrt + >>> + >>> linear_rank_one = P.register_parametrization( + ... nn.Linear(4, 4), "weight", RankOne() + ... ) + >>> print(torch.linalg.matrix_rank(linear_rank_one.weight).item()) + 1 + + """ + parametrization.train(module.training) + if is_parametrized(module, tensor_name): + # Correctness checks. + # If A is the space of tensors with shape and dtype equal to module.weight + # we check that parametrization.forward and parametrization.right_inverse are + # functions from A to A + if not unsafe: + Y = getattr(module, tensor_name) + X = parametrization(Y) + if not isinstance(X, Tensor): + raise ValueError( + f"A parametrization must return a tensor. Got {type(X).__name__}." + ) + if X.dtype != Y.dtype: + raise ValueError( + "Registering a parametrization may not change the dtype of the tensor, unless the `unsafe` flag is enabled.\n" + f"module.{tensor_name}.dtype: {Y.dtype}\n" + f"parametrization(module.{tensor_name}).dtype: {X.dtype}" + ) + if X.shape != Y.shape: + raise ValueError( + "Registering a parametrization may not change the shape of the tensor, unless the `unsafe` flag is enabled.\n" + f"module.{tensor_name}.shape: {Y.shape}\n" + f"parametrization(module.{tensor_name}).shape: {X.shape}" + ) + if hasattr(parametrization, "right_inverse"): + try: + Z = parametrization.right_inverse(X) # type: ignore[operator] + except NotImplementedError: + pass + else: + if not isinstance(Z, Tensor): + raise ValueError( + f"parametrization.right_inverse must return a tensor. Got: {type(Z).__name__}" + ) + if Z.dtype != Y.dtype: + raise ValueError( + "The tensor returned by parametrization.right_inverse must have the same dtype " + f"as module.{tensor_name}, unless the `unsafe` flag is enabled.\n" + f"module.{tensor_name}.dtype: {Y.dtype}\n" + f"returned dtype: {Z.dtype}" + ) + if Z.shape != Y.shape: + raise ValueError( + "The tensor returned by parametrization.right_inverse must have the same shape " + f"as module.{tensor_name}, unless the `unsafe` flag is enabled.\n" + f"module.{tensor_name}.shape: {Y.shape}\n" + f"returned shape: {Z.shape}" + ) + # else right_inverse is assumed to be the identity + + # add the new parametrization to the parametrization list + if not isinstance(module.parametrizations, ModuleDict): + raise AssertionError( + f"Expected module.parametrizations to be a ModuleDict, " + f"got {type(module.parametrizations).__name__}" + ) + module.parametrizations[tensor_name].append(parametrization) # type: ignore[operator] + # If unsafe was True in previous parametrization, keep it enabled + module.parametrizations[tensor_name].unsafe |= unsafe # type: ignore[index, union-attr, operator] + elif tensor_name in module._buffers or tensor_name in module._parameters: + # Set the parametrization mechanism + # Fetch the original buffer or parameter + original = getattr(module, tensor_name) + # We create this early to check for possible errors + parametrizations = ParametrizationList( + [parametrization], original, unsafe=unsafe + ) + # Delete the previous parameter or buffer + delattr(module, tensor_name) + # If this is the first parametrization registered on the module, + # we prepare the module to inject the property + if not is_parametrized(module): + # Change the class + _inject_new_class(module) + # Inject a ``ModuleDict`` into the instance under module.parametrizations + module.parametrizations = ModuleDict() + # Add a property into the class + _inject_property(module, tensor_name) + # Add a ParametrizationList + if not isinstance(module.parametrizations, ModuleDict): + raise AssertionError( + f"Expected module.parametrizations to be a ModuleDict, " + f"got {type(module.parametrizations).__name__}" + ) + module.parametrizations[tensor_name] = parametrizations + else: + raise ValueError( + f"Module '{module}' does not have a parameter, a buffer, or a " + f"parametrized element with name '{tensor_name}'" + ) + return module + + +def is_parametrized(module: Module, tensor_name: str | None = None) -> bool: + r"""Determine if a module has a parametrization. + + Args: + module (nn.Module): module to query + tensor_name (str, optional): name of the parameter in the module + Default: ``None`` + Returns: + ``True`` if :attr:`module` has a parametrization for the parameter named :attr:`tensor_name`, + or if it has any parametrization when :attr:`tensor_name` is ``None``; + otherwise ``False`` + """ + parametrizations = getattr(module, "parametrizations", None) + if parametrizations is None or not isinstance(parametrizations, ModuleDict): + return False + if tensor_name is None: + # Check that there is at least one parametrized buffer or Parameter + return len(parametrizations) > 0 + else: + return tensor_name in parametrizations + + +def remove_parametrizations( + module: Module, + tensor_name: str, + leave_parametrized: bool = True, +) -> Module: + r"""Remove the parametrizations on a tensor in a module. + + - If ``leave_parametrized=True``, ``module[tensor_name]`` will be set to + its current output. In this case, the parametrization shall not change the ``dtype`` + of the tensor. + - If ``leave_parametrized=False``, ``module[tensor_name]`` will be set to + the unparametrised tensor in ``module.parametrizations[tensor_name].original``. + This is only possible when the parametrization depends on just one tensor. + + Args: + module (nn.Module): module from which remove the parametrization + tensor_name (str): name of the parametrization to be removed + leave_parametrized (bool, optional): leave the attribute :attr:`tensor_name` parametrized. + Default: ``True`` + + Returns: + Module: module + + Raises: + ValueError: if ``module[tensor_name]`` is not parametrized + ValueError: if ``leave_parametrized=False`` and the parametrization depends on several tensors + """ + if not is_parametrized(module, tensor_name): + raise ValueError( + f"Module {module} does not have a parametrization on {tensor_name}" + ) + + # Fetch the original tensor + if not isinstance(module.parametrizations, ModuleDict): + raise AssertionError( + f"Expected module.parametrizations to be a ModuleDict, " + f"got {type(module.parametrizations).__name__}" + ) + parametrizations = module.parametrizations[tensor_name] + + if parametrizations.is_tensor: + original = parametrizations.original + if not isinstance(original, torch.Tensor): + raise AssertionError( + f"Expected original to be a Tensor (is_tensor promised us a Tensor), " + f"got {type(original).__name__}" + ) + if leave_parametrized: + with torch.no_grad(): + t = getattr(module, tensor_name) + # We know they have the same dtype because we have checked this when registering the + # parametrizations. As such, we can use set_ + # We do this so that the parameter does not to change the id() + # This way the user does not need to update the optimizer + with torch.no_grad(): + if type(original) is torch.Tensor: + _maybe_set(original, t) + else: + try: + _maybe_set(original, t) + except RuntimeError as e: + # TODO: Fix this for tensor subclasses that are parameters: + # RuntimeError: set_storage is not allowed on a Tensor created from .data or .detach(). + raise RuntimeError( + "Calling remove_parametrizations() with leave_parametrized=True " + "for a parameter that is an instance of a tensor subclass requires " + "set_() to be implemented correctly for the tensor subclass." + "Alternatively, one can opt into the swap_tensors path" + "Either set leave_parametrized=False or provide a working implementation" + "for set_() in the tensor subclass or set " + "torch.__future__.set_swap_module_params_on_conversion(True)." + ) from e + else: + if leave_parametrized: + # We cannot use no_grad because we need to know whether one or more + # original tensors required grad + t = getattr(module, tensor_name) + # We'll have to trust the user to add it to the optimizer + original = Parameter(t) if t.requires_grad else t + else: + raise ValueError( + "Cannot leave unparametrized (`leave_parametrized=False`) a tensor " + "that is parametrized in terms of a sequence of tensors." + ) + + # Delete the property that manages the parametrization + delattr(module.__class__, tensor_name) + # Delete the ParametrizationList + del module.parametrizations[tensor_name] + + # Restore the parameter / buffer into the main class + _register_parameter_or_buffer(module, tensor_name, original) + + # Roll back the parametrized class if no other buffer or parameter + # is currently parametrized in this class + if not is_parametrized(module): + delattr(module, "parametrizations") + # Restore class + orig_cls = module.__class__.__bases__[0] + module.__class__ = orig_cls + return module + + +def type_before_parametrizations(module: Module) -> type: + r"""Return the module type before parametrizations were applied and if not, then it returns the module type. + + Args: + module (nn.Module): module to get type of + """ + if is_parametrized(module): + return module.__class__.__bases__[0] + else: + return type(module) + + +def transfer_parametrizations_and_params( + from_module: Module, + to_module: Module, + tensor_name: str | None = None, +) -> Module: + r"""Transfer parametrizations and the parameters they parametrize from :attr:`from_module` to :attr:`to_module`. + + If :attr:`tensor_name` is specified, only transfers the specified parameter, otherwise + transfers all parametrized parameters. If those parameters do not exist in to_module, it will create them. + Does nothing if from_module is not parametrized. + + Args: + from_module (nn.Module): module to transfer from + to_module (nn.Module): module to transfer to + tensor_name (str, optional): parameter to transfer + + Returns: + Module: to_module + """ + if is_parametrized(from_module): + if not isinstance(from_module.parametrizations, ModuleDict): + raise AssertionError( + f"Expected from_module.parametrizations to be a ModuleDict, " + f"got {type(from_module.parametrizations).__name__}" + ) + + # get list of all params or the single param to transfer + parameters_to_transfer: list | ModuleDict = ( + from_module.parametrizations if tensor_name is None else [tensor_name] + ) + + if not hasattr(parameters_to_transfer, "__iter__"): + raise AssertionError( + f"Expected parameters_to_transfer to be iterable, " + f"got {type(parameters_to_transfer).__name__}" + ) + for parameter_name in parameters_to_transfer: + # initialize the to-be-transferred param in to_module if it doesn't exist already + if not hasattr(to_module, parameter_name): + setattr( + to_module, + parameter_name, + Parameter(getattr(from_module, parameter_name)), + ) + + # apply the params's parametrizations to to_module + for param_func in from_module.parametrizations[ # type: ignore[attr-defined] + parameter_name + ]: + register_parametrization(to_module, parameter_name, param_func) + if not isinstance(to_module.parametrizations, ModuleDict): + raise AssertionError( + f"Expected to_module.parametrizations to be a ModuleDict, " + f"got {type(to_module.parametrizations).__name__}" + ) + + # make values match, original values can be stored in either original or + # original0, original1..., need to check both cases + if hasattr(from_module.parametrizations[parameter_name], "original"): + to_module.parametrizations[ + parameter_name + ].original = from_module.parametrizations[parameter_name].original + else: + num = 0 + orig_num = "original" + str(num) + # loop through each original# until all values have been set + while hasattr(from_module.parametrizations[parameter_name], orig_num): + setattr( + to_module.parametrizations[parameter_name], + orig_num, + getattr(from_module.parametrizations[parameter_name], orig_num), + ) + num = num + 1 + orig_num = "original" + str(num) + + return to_module diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/prune.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/prune.py new file mode 100644 index 0000000000000000000000000000000000000000..240fce1c320698c5121d627a46d124e76f8e78d6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/prune.py @@ -0,0 +1,1395 @@ +# mypy: allow-untyped-defs +r"""Pruning methods.""" + +import numbers +from abc import ABC, abstractmethod +from collections.abc import Iterable + +import torch + + +class BasePruningMethod(ABC): + r"""Abstract base class for creation of new pruning techniques. + + Provides a skeleton for customization requiring the overriding of methods + such as :meth:`compute_mask` and :meth:`apply`. + """ + + _tensor_name: str + + def __call__(self, module, inputs): + r"""Multiply the mask into original tensor and store the result. + + Multiplies the mask (stored in ``module[name + '_mask']``) + into the original tensor (stored in ``module[name + '_orig']``) + and stores the result into ``module[name]`` by using :meth:`apply_mask`. + + Args: + module (nn.Module): module containing the tensor to prune + inputs: not used. + """ + setattr(module, self._tensor_name, self.apply_mask(module)) + + @abstractmethod + def compute_mask(self, t, default_mask): + r"""Compute and returns a mask for the input tensor ``t``. + + Starting from a base ``default_mask`` (which should be a mask of ones + if the tensor has not been pruned yet), generate a random mask to + apply on top of the ``default_mask`` according to the specific pruning + method recipe. + + Args: + t (torch.Tensor): tensor representing the importance scores of the + parameter to prune. + default_mask (torch.Tensor): Base mask from previous pruning + iterations, that need to be respected after the new mask is + applied. Same dims as ``t``. + + Returns: + mask (torch.Tensor): mask to apply to ``t``, of same dims as ``t`` + """ + + def apply_mask(self, module): + r"""Simply handles the multiplication between the parameter being pruned and the generated mask. + + Fetches the mask and the original tensor from the module + and returns the pruned version of the tensor. + + Args: + module (nn.Module): module containing the tensor to prune + + Returns: + pruned_tensor (torch.Tensor): pruned version of the input tensor + """ + # to carry out the multiplication, the mask needs to have been computed, + # so the pruning method must know what tensor it's operating on + if self._tensor_name is None: + raise AssertionError( + f"Module {module} has to be pruned" + ) # this gets set in apply() + mask = getattr(module, self._tensor_name + "_mask") + orig = getattr(module, self._tensor_name + "_orig") + pruned_tensor = mask.to(dtype=orig.dtype) * orig + return pruned_tensor + + @classmethod + def apply(cls, module, name, *args, importance_scores=None, **kwargs): + r"""Add pruning on the fly and reparameterization of a tensor. + + Adds the forward pre-hook that enables pruning on the fly and + the reparameterization of a tensor in terms of the original tensor + and the pruning mask. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + args: arguments passed on to a subclass of + :class:`BasePruningMethod` + importance_scores (torch.Tensor): tensor of importance scores (of + same shape as module parameter) used to compute mask for pruning. + The values in this tensor indicate the importance of the + corresponding elements in the parameter being pruned. + If unspecified or None, the parameter will be used in its place. + kwargs: keyword arguments passed on to a subclass of a + :class:`BasePruningMethod` + """ + + def _get_composite_method(cls, module, name, *args, **kwargs): + # Check if a pruning method has already been applied to + # `module[name]`. If so, store that in `old_method`. + old_method = None + found = 0 + # there should technically be only 1 hook with hook.name == name + # assert this using `found` + hooks_to_remove = [] + for k, hook in module._forward_pre_hooks.items(): + # if it exists, take existing thing, remove hook, then + # go through normal thing + if isinstance(hook, BasePruningMethod) and hook._tensor_name == name: + old_method = hook + hooks_to_remove.append(k) + found += 1 + if found > 1: + raise AssertionError( + f"Avoid adding multiple pruning hooks to the " + f"same tensor {name} of module {module}. Use a PruningContainer." + ) + + for k in hooks_to_remove: + del module._forward_pre_hooks[k] + + # Apply the new pruning method, either from scratch or on top of + # the previous one. + method = cls(*args, **kwargs) # new pruning + # Have the pruning method remember what tensor it's been applied to + method._tensor_name = name + + # combine `methods` with `old_method`, if `old_method` exists + if old_method is not None: # meaning that there was a hook + # if the hook is already a pruning container, just add the + # new pruning method to the container + if isinstance(old_method, PruningContainer): + old_method.add_pruning_method(method) + method = old_method # rename old_method --> method + + # if the hook is simply a single pruning method, create a + # container, add the old pruning method and the new one + elif isinstance(old_method, BasePruningMethod): + container = PruningContainer(old_method) + # Have the pruning method remember the name of its tensor + # setattr(container, '_tensor_name', name) + container.add_pruning_method(method) + method = container # rename container --> method + return method + + method = _get_composite_method(cls, module, name, *args, **kwargs) + # at this point we have no forward_pre_hooks but we could have an + # active reparameterization of the tensor if another pruning method + # had been applied (in which case `method` would be a PruningContainer + # and not a simple pruning method). + + # Pruning is to be applied to the module's tensor named `name`, + # starting from the state it is found in prior to this iteration of + # pruning. The pruning mask is calculated based on importances scores. + + orig = getattr(module, name) + if importance_scores is not None: + if importance_scores.shape != orig.shape: + raise AssertionError( + f"importance_scores should have the same shape as parameter " + f"{name} of {module}, got {importance_scores.shape} vs {orig.shape}" + ) + else: + importance_scores = orig + + # If this is the first time pruning is applied, take care of moving + # the original tensor to a new parameter called name + '_orig' and + # and deleting the original parameter + if not isinstance(method, PruningContainer): + # copy `module[name]` to `module[name + '_orig']` + module.register_parameter(name + "_orig", orig) + # temporarily delete `module[name]` + del module._parameters[name] + default_mask = torch.ones_like(orig) # temp + # If this is not the first time pruning is applied, all of the above + # has been done before in a previous pruning iteration, so we're good + # to go + else: + default_mask = ( + getattr(module, name + "_mask") + .detach() + .clone(memory_format=torch.contiguous_format) + ) + + # Use try/except because if anything goes wrong with the mask + # computation etc., you'd want to roll back. + try: + # get the final mask, computed according to the specific method + mask = method.compute_mask(importance_scores, default_mask=default_mask) + # reparameterize by saving mask to `module[name + '_mask']`... + module.register_buffer(name + "_mask", mask) + # ... and the new pruned tensor to `module[name]` + setattr(module, name, method.apply_mask(module)) + # associate the pruning method to the module via a hook to + # compute the function before every forward() (compile by run) + module.register_forward_pre_hook(method) + + except Exception as e: + if not isinstance(method, PruningContainer): + orig = getattr(module, name + "_orig") + module.register_parameter(name, orig) + del module._parameters[name + "_orig"] + raise e + + return method + + def prune(self, t, default_mask=None, importance_scores=None): + r"""Compute and returns a pruned version of input tensor ``t``. + + According to the pruning rule specified in :meth:`compute_mask`. + + Args: + t (torch.Tensor): tensor to prune (of same dimensions as + ``default_mask``). + importance_scores (torch.Tensor): tensor of importance scores (of + same shape as ``t``) used to compute mask for pruning ``t``. + The values in this tensor indicate the importance of the + corresponding elements in the ``t`` that is being pruned. + If unspecified or None, the tensor ``t`` will be used in its place. + default_mask (torch.Tensor, optional): mask from previous pruning + iteration, if any. To be considered when determining what + portion of the tensor that pruning should act on. If None, + default to a mask of ones. + + Returns: + pruned version of tensor ``t``. + """ + if importance_scores is not None: + if importance_scores.shape != t.shape: + raise AssertionError( + f"importance_scores should have the same shape as tensor t, " + f"got {importance_scores.shape} vs {t.shape}" + ) + else: + importance_scores = t + default_mask = default_mask if default_mask is not None else torch.ones_like(t) + return t * self.compute_mask(importance_scores, default_mask=default_mask) + + def remove(self, module) -> None: + r"""Remove the pruning reparameterization from a module. + + The pruned parameter named ``name`` remains permanently pruned, + and the parameter named ``name+'_orig'`` is removed from the parameter list. + Similarly, the buffer named ``name+'_mask'`` is removed from the buffers. + + Note: + Pruning itself is NOT undone or reversed! + """ + # before removing pruning from a tensor, it has to have been applied + if self._tensor_name is None: + raise AssertionError( + f"Module {module} has to be pruned before pruning can be removed" + ) # this gets set in apply() + + # to update module[name] to latest trained weights + weight = self.apply_mask(module) # masked weights + + # delete and reset + if hasattr(module, self._tensor_name): + delattr(module, self._tensor_name) + orig = module._parameters[self._tensor_name + "_orig"] + orig.data = weight.data + del module._parameters[self._tensor_name + "_orig"] + del module._buffers[self._tensor_name + "_mask"] + setattr(module, self._tensor_name, orig) + + +class PruningContainer(BasePruningMethod): + """Container holding a sequence of pruning methods for iterative pruning. + + Keeps track of the order in which pruning methods are applied and handles + combining successive pruning calls. + + Accepts as argument an instance of a BasePruningMethod or an iterable of + them. + """ + + def __init__(self, *args) -> None: + self._pruning_methods: tuple[BasePruningMethod, ...] = () + if not isinstance(args, Iterable): # only 1 item + self._tensor_name = args._tensor_name + self.add_pruning_method(args) + + elif len(args) == 1: # only 1 item in a tuple + self._tensor_name = args[0]._tensor_name + + self.add_pruning_method(args[0]) + else: # manual construction from list or other iterable (or no args) + for method in args: + self.add_pruning_method(method) + + def add_pruning_method(self, method) -> None: + r"""Add a child pruning ``method`` to the container. + + Args: + method (subclass of BasePruningMethod): child pruning method + to be added to the container. + """ + # check that we're adding a pruning method to the container + if not isinstance(method, BasePruningMethod) and method is not None: + raise TypeError(f"{type(method)} is not a BasePruningMethod subclass") + elif method is not None and self._tensor_name != method._tensor_name: + raise ValueError( + "Can only add pruning methods acting on " + f"the parameter named '{self._tensor_name}' to PruningContainer {self}." + + f" Found '{method._tensor_name}'" + ) + # if all checks passed, add to _pruning_methods tuple + self._pruning_methods += (method,) # type: ignore[operator] + + def __len__(self) -> int: + return len(self._pruning_methods) + + def __iter__(self): + return iter(self._pruning_methods) + + def __getitem__(self, idx): + return self._pruning_methods[idx] + + def compute_mask(self, t, default_mask): + r"""Apply the latest ``method`` by computing the new partial masks and returning its combination with the ``default_mask``. + + The new partial mask should be computed on the entries or channels + that were not zeroed out by the ``default_mask``. + Which portions of the tensor ``t`` the new mask will be calculated from + depends on the ``PRUNING_TYPE`` (handled by the type handler): + + * for 'unstructured', the mask will be computed from the raveled + list of nonmasked entries; + + * for 'structured', the mask will be computed from the nonmasked + channels in the tensor; + + * for 'global', the mask will be computed across all entries. + + Args: + t (torch.Tensor): tensor representing the parameter to prune + (of same dimensions as ``default_mask``). + default_mask (torch.Tensor): mask from previous pruning iteration. + + Returns: + mask (torch.Tensor): new mask that combines the effects + of the ``default_mask`` and the new mask from the current + pruning ``method`` (of same dimensions as ``default_mask`` and + ``t``). + """ + + def _combine_masks(method, t, mask): + r"""Combine the masks from all pruning methods and returns a new mask. + + Args: + method (a BasePruningMethod subclass): pruning method + currently being applied. + t (torch.Tensor): tensor representing the parameter to prune + (of same dimensions as mask). + mask (torch.Tensor): mask from previous pruning iteration + + Returns: + new_mask (torch.Tensor): new mask that combines the effects + of the old mask and the new mask from the current + pruning method (of same dimensions as mask and t). + """ + new_mask = mask # start off from existing mask + new_mask = new_mask.to(dtype=t.dtype) + + # compute a slice of t onto which the new pruning method will operate + if method.PRUNING_TYPE == "unstructured": + # prune entries of t where the mask is 1 + slc = mask == 1 + + # for struct pruning, exclude channels that have already been + # entirely pruned + elif method.PRUNING_TYPE == "structured": + if not hasattr(method, "dim"): + raise AttributeError( + "Pruning methods of PRUNING_TYPE " + '"structured" need to have the attribute `dim` defined.' + ) + + # find the channels to keep by removing the ones that have been + # zeroed out already (i.e. where sum(entries) == 0) + n_dims = t.dim() # "is this a 2D tensor? 3D? ..." + dim = method.dim + # convert negative indexing + if dim < 0: + dim = n_dims + dim + # if dim is still negative after subtracting it from n_dims + if dim < 0: + raise IndexError( + f"Index is out of bounds for tensor with dimensions {n_dims}" + ) + # find channels along dim = dim that aren't already tots 0ed out + keep_channel = mask.sum(dim=[d for d in range(n_dims) if d != dim]) != 0 + # create slice to identify what to prune + slc = [slice(None)] * n_dims + slc[dim] = keep_channel + + elif method.PRUNING_TYPE == "global": + n_dims = len(t.shape) # "is this a 2D tensor? 3D? ..." + slc = [slice(None)] * n_dims + + else: + raise ValueError(f"Unrecognized PRUNING_TYPE {method.PRUNING_TYPE}") + + # compute the new mask on the unpruned slice of the tensor t + if isinstance(slc, list): + slc = tuple(slc) + partial_mask = method.compute_mask(t[slc], default_mask=mask[slc]) + new_mask[slc] = partial_mask.to(dtype=new_mask.dtype) + + return new_mask + + method = self._pruning_methods[-1] + mask = _combine_masks(method, t, default_mask) + return mask + + +class Identity(BasePruningMethod): + r"""Utility pruning method that does not prune any units but generates the pruning parametrization with a mask of ones.""" + + PRUNING_TYPE = "unstructured" + + def compute_mask(self, t, default_mask): + mask = default_mask + return mask + + @classmethod + def apply(cls, module, name): # type: ignore[override] + r"""Add pruning on the fly and reparameterization of a tensor. + + Adds the forward pre-hook that enables pruning on the fly and + the reparameterization of a tensor in terms of the original tensor + and the pruning mask. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + """ + return super().apply(module, name) + + +class RandomUnstructured(BasePruningMethod): + r"""Prune (currently unpruned) units in a tensor at random. + + Args: + name (str): parameter name within ``module`` on which pruning + will act. + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + """ + + PRUNING_TYPE = "unstructured" + + def __init__(self, amount) -> None: + # Check range of validity of pruning amount + _validate_pruning_amount_init(amount) + self.amount = amount + + def compute_mask(self, t, default_mask): + # Check that the amount of units to prune is not > than the number of + # parameters in t + tensor_size = t.nelement() + # Compute number of units to prune: amount if int, + # else amount * tensor_size + nparams_toprune = _compute_nparams_toprune(self.amount, tensor_size) + # This should raise an error if the number of units to prune is larger + # than the number of units in the tensor + _validate_pruning_amount(nparams_toprune, tensor_size) + + mask = default_mask.clone(memory_format=torch.contiguous_format) + + if nparams_toprune != 0: # k=0 not supported by torch.kthvalue + prob = torch.rand_like(t) + topk = torch.topk(prob.view(-1), k=nparams_toprune) + mask.view(-1)[topk.indices] = 0 + + return mask + + @classmethod + def apply(cls, module, name, amount): # type: ignore[override] + r"""Add pruning on the fly and reparameterization of a tensor. + + Adds the forward pre-hook that enables pruning on the fly and + the reparameterization of a tensor in terms of the original tensor + and the pruning mask. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + """ + return super().apply(module, name, amount=amount) + + +class L1Unstructured(BasePruningMethod): + r"""Prune (currently unpruned) units in a tensor by zeroing out the ones with the lowest L1-norm. + + Args: + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + """ + + PRUNING_TYPE = "unstructured" + + def __init__(self, amount) -> None: + # Check range of validity of pruning amount + _validate_pruning_amount_init(amount) + self.amount = amount + + def compute_mask(self, t, default_mask): + # Check that the amount of units to prune is not > than the number of + # parameters in t + tensor_size = t.nelement() + # Compute number of units to prune: amount if int, + # else amount * tensor_size + nparams_toprune = _compute_nparams_toprune(self.amount, tensor_size) + # This should raise an error if the number of units to prune is larger + # than the number of units in the tensor + _validate_pruning_amount(nparams_toprune, tensor_size) + + mask = default_mask.clone(memory_format=torch.contiguous_format) + + if nparams_toprune != 0: # k=0 not supported by torch.kthvalue + # largest=True --> top k; largest=False --> bottom k + # Prune the smallest k + topk = torch.topk(torch.abs(t).view(-1), k=nparams_toprune, largest=False) + # topk will have .indices and .values + mask.view(-1)[topk.indices] = 0 + + return mask + + @classmethod + def apply(cls, module, name, amount, importance_scores=None): # type: ignore[override] + r"""Add pruning on the fly and reparameterization of a tensor. + + Adds the forward pre-hook that enables pruning on the fly and + the reparameterization of a tensor in terms of the original tensor + and the pruning mask. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + importance_scores (torch.Tensor): tensor of importance scores (of same + shape as module parameter) used to compute mask for pruning. + The values in this tensor indicate the importance of the corresponding + elements in the parameter being pruned. + If unspecified or None, the module parameter will be used in its place. + """ + return super().apply( + module, name, amount=amount, importance_scores=importance_scores + ) + + +class RandomStructured(BasePruningMethod): + r"""Prune entire (currently unpruned) channels in a tensor at random. + + Args: + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + dim (int, optional): index of the dim along which we define + channels to prune. Default: -1. + """ + + PRUNING_TYPE = "structured" + + def __init__(self, amount, dim=-1) -> None: + # Check range of validity of amount + _validate_pruning_amount_init(amount) + self.amount = amount + self.dim = dim + + def compute_mask(self, t, default_mask): + r"""Compute and returns a mask for the input tensor ``t``. + + Starting from a base ``default_mask`` (which should be a mask of ones + if the tensor has not been pruned yet), generate a random mask to + apply on top of the ``default_mask`` by randomly zeroing out channels + along the specified dim of the tensor. + + Args: + t (torch.Tensor): tensor representing the parameter to prune + default_mask (torch.Tensor): Base mask from previous pruning + iterations, that need to be respected after the new mask is + applied. Same dims as ``t``. + + Returns: + mask (torch.Tensor): mask to apply to ``t``, of same dims as ``t`` + + Raises: + IndexError: if ``self.dim >= len(t.shape)`` + """ + # Check that tensor has structure (i.e. more than 1 dimension) such + # that the concept of "channels" makes sense + _validate_structured_pruning(t) + + # Check that self.dim is a valid dim to index t, else raise IndexError + _validate_pruning_dim(t, self.dim) + + # Check that the amount of channels to prune is not > than the number of + # channels in t along the dim to prune + tensor_size = t.shape[self.dim] + # Compute number of units to prune: amount if int, + # else amount * tensor_size + nparams_toprune = _compute_nparams_toprune(self.amount, tensor_size) + # This should raise an error if the number of units to prune is larger + # than the number of units in the tensor + _validate_pruning_amount(nparams_toprune, tensor_size) + + # Compute binary mask by initializing it to all 0s and then filling in + # 1s wherever topk.indices indicates, along self.dim. + # mask has the same shape as tensor t + def make_mask(t, dim, nchannels, nchannels_toprune): + # generate a random number in [0, 1] to associate to each channel + prob = torch.rand(nchannels) + # generate mask for each channel by 0ing out the channels that + # got assigned the k = nchannels_toprune lowest values in prob + threshold = torch.kthvalue(prob, k=nchannels_toprune).values + channel_mask = prob > threshold + + mask = torch.zeros_like(t) + slc = [slice(None)] * len(t.shape) + slc[dim] = channel_mask + slc = tuple(slc) + mask[slc] = 1 + return mask + + if nparams_toprune == 0: # k=0 not supported by torch.kthvalue + mask = default_mask + else: + # apply the new structured mask on top of prior (potentially + # unstructured) mask + mask = make_mask(t, self.dim, tensor_size, nparams_toprune) + mask *= default_mask.to(dtype=mask.dtype) + return mask + + @classmethod + def apply(cls, module, name, amount, dim=-1): # type: ignore[override] + r"""Add pruning on the fly and reparameterization of a tensor. + + Adds the forward pre-hook that enables pruning on the fly and + the reparameterization of a tensor in terms of the original tensor + and the pruning mask. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + dim (int, optional): index of the dim along which we define + channels to prune. Default: -1. + """ + return super().apply(module, name, amount=amount, dim=dim) + + +class LnStructured(BasePruningMethod): + r"""Prune entire (currently unpruned) channels in a tensor based on their L\ ``n``-norm. + + Args: + amount (int or float): quantity of channels to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid + entries for argument ``p`` in :func:`torch.norm`. + dim (int, optional): index of the dim along which we define + channels to prune. Default: -1. + """ + + PRUNING_TYPE = "structured" + + def __init__(self, amount, n, dim=-1) -> None: + # Check range of validity of amount + _validate_pruning_amount_init(amount) + self.amount = amount + self.n = n + self.dim = dim + + def compute_mask(self, t, default_mask): + r"""Compute and returns a mask for the input tensor ``t``. + + Starting from a base ``default_mask`` (which should be a mask of ones + if the tensor has not been pruned yet), generate a mask to apply on + top of the ``default_mask`` by zeroing out the channels along the + specified dim with the lowest L\ ``n``-norm. + + Args: + t (torch.Tensor): tensor representing the parameter to prune + default_mask (torch.Tensor): Base mask from previous pruning + iterations, that need to be respected after the new mask is + applied. Same dims as ``t``. + + Returns: + mask (torch.Tensor): mask to apply to ``t``, of same dims as ``t`` + + Raises: + IndexError: if ``self.dim >= len(t.shape)`` + """ + # Check that tensor has structure (i.e. more than 1 dimension) such + # that the concept of "channels" makes sense + _validate_structured_pruning(t) + # Check that self.dim is a valid dim to index t, else raise IndexError + _validate_pruning_dim(t, self.dim) + + # Check that the amount of channels to prune is not > than the number of + # channels in t along the dim to prune + tensor_size = t.shape[self.dim] + # Compute number of units to prune: amount if int, + # else amount * tensor_size + nparams_toprune = _compute_nparams_toprune(self.amount, tensor_size) + nparams_tokeep = tensor_size - nparams_toprune + # This should raise an error if the number of units to prune is larger + # than the number of units in the tensor + _validate_pruning_amount(nparams_toprune, tensor_size) + + # Structured pruning prunes entire channels so we need to know the + # L_n norm along each channel to then find the topk based on this + # metric + norm = _compute_norm(t, self.n, self.dim) + # largest=True --> top k; largest=False --> bottom k + # Keep the largest k channels along dim=self.dim + topk = torch.topk(norm, k=nparams_tokeep, largest=True) + # topk will have .indices and .values + + # Compute binary mask by initializing it to all 0s and then filling in + # 1s wherever topk.indices indicates, along self.dim. + # mask has the same shape as tensor t + def make_mask(t, dim, indices): + # init mask to 0 + mask = torch.zeros_like(t) + # e.g.: slc = [None, None, None], if len(t.shape) = 3 + slc = [slice(None)] * len(t.shape) + # replace a None at position=dim with indices + # e.g.: slc = [None, None, [0, 2, 3]] if dim=2 & indices=[0,2,3] + slc[dim] = indices + slc = tuple(slc) + # use slc to slice mask and replace all its entries with 1s + # e.g.: mask[:, :, [0, 2, 3]] = 1 + mask[slc] = 1 + return mask + + if nparams_toprune == 0: # k=0 not supported by torch.kthvalue + mask = default_mask + else: + mask = make_mask(t, self.dim, topk.indices) + mask *= default_mask.to(dtype=mask.dtype) + + return mask + + @classmethod + def apply(cls, module, name, amount, n, dim, importance_scores=None): # type: ignore[override] + r"""Add pruning on the fly and reparameterization of a tensor. + + Adds the forward pre-hook that enables pruning on the fly and + the reparameterization of a tensor in terms of the original tensor + and the pruning mask. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid + entries for argument ``p`` in :func:`torch.norm`. + dim (int): index of the dim along which we define channels to + prune. + importance_scores (torch.Tensor): tensor of importance scores (of same + shape as module parameter) used to compute mask for pruning. + The values in this tensor indicate the importance of the corresponding + elements in the parameter being pruned. + If unspecified or None, the module parameter will be used in its place. + """ + return super().apply( + module, + name, + amount=amount, + n=n, + dim=dim, + importance_scores=importance_scores, + ) + + +class CustomFromMask(BasePruningMethod): + PRUNING_TYPE = "global" + + def __init__(self, mask) -> None: + self.mask = mask + + def compute_mask(self, t, default_mask): + if default_mask.shape != self.mask.shape: + raise AssertionError( + f"default_mask shape {default_mask.shape} must match " + f"self.mask shape {self.mask.shape}" + ) + mask = default_mask * self.mask.to(dtype=default_mask.dtype) + return mask + + @classmethod + def apply(cls, module, name, mask): # type: ignore[override] + r"""Add pruning on the fly and reparameterization of a tensor. + + Adds the forward pre-hook that enables pruning on the fly and + the reparameterization of a tensor in terms of the original tensor + and the pruning mask. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + """ + return super().apply(module, name, mask=mask) + + +def identity(module, name): + r"""Apply pruning reparameterization without pruning any units. + + Applies pruning reparameterization to the tensor corresponding to the + parameter called ``name`` in ``module`` without actually pruning any + units. Modifies module in place (and also return the modified module) + by: + + 1) adding a named buffer called ``name+'_mask'`` corresponding to the + binary mask applied to the parameter ``name`` by the pruning method. + 2) replacing the parameter ``name`` by its pruned version, while the + original (unpruned) parameter is stored in a new parameter named + ``name+'_orig'``. + + Note: + The mask is a tensor of ones. + + Args: + module (nn.Module): module containing the tensor to prune. + name (str): parameter name within ``module`` on which pruning + will act. + + Returns: + module (nn.Module): modified (i.e. pruned) version of the input module + + Examples: + >>> # xdoctest: +SKIP + >>> m = prune.identity(nn.Linear(2, 3), "bias") + >>> print(m.bias_mask) + tensor([1., 1., 1.]) + """ + Identity.apply(module, name) + return module + + +def random_unstructured(module, name, amount): + r"""Prune tensor by removing random (currently unpruned) units. + + Prunes tensor corresponding to parameter called ``name`` in ``module`` + by removing the specified ``amount`` of (currently unpruned) units + selected at random. + Modifies module in place (and also return the modified module) by: + + 1) adding a named buffer called ``name+'_mask'`` corresponding to the + binary mask applied to the parameter ``name`` by the pruning method. + 2) replacing the parameter ``name`` by its pruned version, while the + original (unpruned) parameter is stored in a new parameter named + ``name+'_orig'``. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + + Returns: + module (nn.Module): modified (i.e. pruned) version of the input module + + Examples: + >>> # xdoctest: +SKIP + >>> m = prune.random_unstructured(nn.Linear(2, 3), "weight", amount=1) + >>> torch.sum(m.weight_mask == 0) + tensor(1) + + """ + RandomUnstructured.apply(module, name, amount) + return module + + +def l1_unstructured(module, name, amount, importance_scores=None): + r"""Prune tensor by removing units with the lowest L1-norm. + + Prunes tensor corresponding to parameter called ``name`` in ``module`` + by removing the specified `amount` of (currently unpruned) units with the + lowest L1-norm. + Modifies module in place (and also return the modified module) + by: + + 1) adding a named buffer called ``name+'_mask'`` corresponding to the + binary mask applied to the parameter ``name`` by the pruning method. + 2) replacing the parameter ``name`` by its pruned version, while the + original (unpruned) parameter is stored in a new parameter named + ``name+'_orig'``. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + importance_scores (torch.Tensor): tensor of importance scores (of same + shape as module parameter) used to compute mask for pruning. + The values in this tensor indicate the importance of the corresponding + elements in the parameter being pruned. + If unspecified or None, the module parameter will be used in its place. + + Returns: + module (nn.Module): modified (i.e. pruned) version of the input module + + Examples: + >>> # xdoctest: +SKIP + >>> m = prune.l1_unstructured(nn.Linear(2, 3), "weight", amount=0.2) + >>> m.state_dict().keys() + odict_keys(['bias', 'weight_orig', 'weight_mask']) + """ + L1Unstructured.apply( + module, name, amount=amount, importance_scores=importance_scores + ) + return module + + +def random_structured(module, name, amount, dim): + r"""Prune tensor by removing random channels along the specified dimension. + + Prunes tensor corresponding to parameter called ``name`` in ``module`` + by removing the specified ``amount`` of (currently unpruned) channels + along the specified ``dim`` selected at random. + Modifies module in place (and also return the modified module) + by: + + 1) adding a named buffer called ``name+'_mask'`` corresponding to the + binary mask applied to the parameter ``name`` by the pruning method. + 2) replacing the parameter ``name`` by its pruned version, while the + original (unpruned) parameter is stored in a new parameter named + ``name+'_orig'``. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + dim (int): index of the dim along which we define channels to prune. + + Returns: + module (nn.Module): modified (i.e. pruned) version of the input module + + Examples: + >>> # xdoctest: +SKIP + >>> m = prune.random_structured(nn.Linear(5, 3), "weight", amount=3, dim=1) + >>> columns_pruned = int(sum(torch.sum(m.weight, dim=0) == 0)) + >>> print(columns_pruned) + 3 + """ + RandomStructured.apply(module, name, amount, dim) + return module + + +def ln_structured(module, name, amount, n, dim, importance_scores=None): + r"""Prune tensor by removing channels with the lowest L\ ``n``-norm along the specified dimension. + + Prunes tensor corresponding to parameter called ``name`` in ``module`` + by removing the specified ``amount`` of (currently unpruned) channels + along the specified ``dim`` with the lowest L\ ``n``-norm. + Modifies module in place (and also return the modified module) + by: + + 1) adding a named buffer called ``name+'_mask'`` corresponding to the + binary mask applied to the parameter ``name`` by the pruning method. + 2) replacing the parameter ``name`` by its pruned version, while the + original (unpruned) parameter is stored in a new parameter named + ``name+'_orig'``. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + amount (int or float): quantity of parameters to prune. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid + entries for argument ``p`` in :func:`torch.norm`. + dim (int): index of the dim along which we define channels to prune. + importance_scores (torch.Tensor): tensor of importance scores (of same + shape as module parameter) used to compute mask for pruning. + The values in this tensor indicate the importance of the corresponding + elements in the parameter being pruned. + If unspecified or None, the module parameter will be used in its place. + + Returns: + module (nn.Module): modified (i.e. pruned) version of the input module + + Examples: + >>> from torch.nn.utils import prune + >>> m = prune.ln_structured( + ... nn.Conv2d(5, 3, 2), "weight", amount=0.3, dim=1, n=float("-inf") + ... ) + """ + LnStructured.apply( + module, name, amount, n, dim, importance_scores=importance_scores + ) + return module + + +def global_unstructured( + parameters, pruning_method, importance_scores=None, **kwargs +) -> None: + r""" + Globally prunes tensors corresponding to all parameters in ``parameters`` by applying the specified ``pruning_method``. + + Modifies modules in place by: + + 1) adding a named buffer called ``name+'_mask'`` corresponding to the + binary mask applied to the parameter ``name`` by the pruning method. + 2) replacing the parameter ``name`` by its pruned version, while the + original (unpruned) parameter is stored in a new parameter named + ``name+'_orig'``. + + Args: + parameters (Iterable of (module, name) tuples): parameters of + the model to prune in a global fashion, i.e. by aggregating all + weights prior to deciding which ones to prune. module must be of + type :class:`nn.Module`, and name must be a string. + pruning_method (function): a valid pruning function from this module, + or a custom one implemented by the user that satisfies the + implementation guidelines and has ``PRUNING_TYPE='unstructured'``. + importance_scores (dict): a dictionary mapping (module, name) tuples to + the corresponding parameter's importance scores tensor. The tensor + should be the same shape as the parameter, and is used for computing + mask for pruning. + If unspecified or None, the parameter will be used in place of its + importance scores. + kwargs: other keyword arguments such as: + amount (int or float): quantity of parameters to prune across the + specified parameters. + If ``float``, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If ``int``, it represents the + absolute number of parameters to prune. + + Raises: + TypeError: if ``PRUNING_TYPE != 'unstructured'`` + + Note: + Since global structured pruning doesn't make much sense unless the + norm is normalized by the size of the parameter, we now limit the + scope of global pruning to unstructured methods. + + Examples: + >>> from torch.nn.utils import prune + >>> from collections import OrderedDict + >>> net = nn.Sequential( + ... OrderedDict( + ... [ + ... ("first", nn.Linear(10, 4)), + ... ("second", nn.Linear(4, 1)), + ... ] + ... ) + ... ) + >>> parameters_to_prune = ( + ... (net.first, "weight"), + ... (net.second, "weight"), + ... ) + >>> prune.global_unstructured( + ... parameters_to_prune, + ... pruning_method=prune.L1Unstructured, + ... amount=10, + ... ) + >>> print(sum(torch.nn.utils.parameters_to_vector(net.buffers()) == 0)) + tensor(10) + + """ + # ensure parameters is a list or generator of tuples + if not isinstance(parameters, Iterable): + raise TypeError("global_unstructured(): parameters is not an Iterable") + + importance_scores = importance_scores if importance_scores is not None else {} + if not isinstance(importance_scores, dict): + raise TypeError("global_unstructured(): importance_scores must be of type dict") + + # flatten importance scores to consider them all at once in global pruning + relevant_importance_scores = torch.nn.utils.parameters_to_vector( + # pyrefly: ignore [bad-argument-type] + [ + importance_scores.get((module, name), getattr(module, name)) + for (module, name) in parameters + ] + ) + # similarly, flatten the masks (if they exist), or use a flattened vector + # of 1s of the same dimensions as t + default_mask = torch.nn.utils.parameters_to_vector( + [ + getattr(module, name + "_mask", torch.ones_like(getattr(module, name))) + for (module, name) in parameters + ] + ) + + # use the canonical pruning methods to compute the new mask, even if the + # parameter is now a flattened out version of `parameters` + container = PruningContainer() + container._tensor_name = "temp" # to make it match that of `method` + method = pruning_method(**kwargs) + method._tensor_name = "temp" # to make it match that of `container` + if method.PRUNING_TYPE != "unstructured": + raise TypeError( + 'Only "unstructured" PRUNING_TYPE supported for ' + f"the `pruning_method`. Found method {pruning_method} of type {method.PRUNING_TYPE}" + ) + + container.add_pruning_method(method) + + # use the `compute_mask` method from `PruningContainer` to combine the + # mask computed by the new method with the pre-existing mask + final_mask = container.compute_mask(relevant_importance_scores, default_mask) + + # Pointer for slicing the mask to match the shape of each parameter + pointer = 0 + for module, name in parameters: + param = getattr(module, name) + # The length of the parameter + num_param = param.numel() + # Slice the mask, reshape it + param_mask = final_mask[pointer : pointer + num_param].view_as(param) + # Assign the correct pre-computed mask to each parameter and add it + # to the forward_pre_hooks like any other pruning method + custom_from_mask(module, name, mask=param_mask) + + # Increment the pointer to continue slicing the final_mask + pointer += num_param + + +def custom_from_mask(module, name, mask): + r"""Prune tensor corresponding to parameter called ``name`` in ``module`` by applying the pre-computed mask in ``mask``. + + Modifies module in place (and also return the modified module) by: + + 1) adding a named buffer called ``name+'_mask'`` corresponding to the + binary mask applied to the parameter ``name`` by the pruning method. + 2) replacing the parameter ``name`` by its pruned version, while the + original (unpruned) parameter is stored in a new parameter named + ``name+'_orig'``. + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + mask (Tensor): binary mask to be applied to the parameter. + + Returns: + module (nn.Module): modified (i.e. pruned) version of the input module + + Examples: + >>> from torch.nn.utils import prune + >>> m = prune.custom_from_mask( + ... nn.Linear(5, 3), name="bias", mask=torch.tensor([0, 1, 0]) + ... ) + >>> print(m.bias_mask) + tensor([0., 1., 0.]) + + """ + CustomFromMask.apply(module, name, mask) + return module + + +def remove(module, name): + r"""Remove the pruning reparameterization from a module and the pruning method from the forward hook. + + The pruned parameter named ``name`` remains permanently pruned, and the parameter + named ``name+'_orig'`` is removed from the parameter list. Similarly, + the buffer named ``name+'_mask'`` is removed from the buffers. + + Note: + Pruning itself is NOT undone or reversed! + + Args: + module (nn.Module): module containing the tensor to prune + name (str): parameter name within ``module`` on which pruning + will act. + + Examples: + >>> m = random_unstructured(nn.Linear(5, 7), name="weight", amount=0.2) + >>> m = remove(m, name="weight") + """ + for k, hook in module._forward_pre_hooks.items(): + if isinstance(hook, BasePruningMethod) and hook._tensor_name == name: + hook.remove(module) + del module._forward_pre_hooks[k] + return module + + raise ValueError( + f"Parameter '{name}' of module {module} has to be pruned before pruning can be removed" + ) + + +def is_pruned(module) -> bool: + r"""Check if a module is pruned by looking for pruning pre-hooks. + + Check whether ``module`` is pruned by looking for + ``forward_pre_hooks`` in its modules that inherit from the + :class:`BasePruningMethod`. + + Args: + module (nn.Module): object that is either pruned or unpruned + + Returns: + binary answer to whether ``module`` is pruned. + + Examples: + >>> from torch.nn.utils import prune + >>> m = nn.Linear(5, 7) + >>> print(prune.is_pruned(m)) + False + >>> prune.random_unstructured(m, name="weight", amount=0.2) + >>> print(prune.is_pruned(m)) + True + """ + for _, submodule in module.named_modules(): + for hook in submodule._forward_pre_hooks.values(): + if isinstance(hook, BasePruningMethod): + return True + return False + + +def _validate_pruning_amount_init(amount) -> None: + r"""Validate helper to check the range of amount at init. + + Args: + amount (int or float): quantity of parameters to prune. + If float, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If int, it represents the + absolute number of parameters to prune. + + Raises: + ValueError: if amount is a float not in [0, 1], or if it's a negative + integer. + TypeError: if amount is neither a float nor an integer. + + Note: + This does not take into account the number of parameters in the + tensor to be pruned, which is known only at prune. + """ + if not isinstance(amount, numbers.Real): + raise TypeError(f"Invalid type for amount: {amount}. Must be int or float.") + + if (isinstance(amount, numbers.Integral) and amount < 0) or ( + not isinstance(amount, numbers.Integral) # so it's a float + and (float(amount) > 1.0 or float(amount) < 0.0) + ): + raise ValueError( + f"amount={amount} should either be a float in the range [0, 1] or a non-negative integer" + ) + + +def _validate_pruning_amount(amount, tensor_size) -> None: + r"""Validate that the pruning amount is meaningful wrt to the size of the data. + + Validation helper to check that the amount of parameters to prune + is meaningful wrt to the size of the data (`tensor_size`). + + Args: + amount (int or float): quantity of parameters to prune. + If float, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If int, it represents the + absolute number of parameters to prune. + tensor_size (int): absolute number of parameters in the tensor + to prune. + """ + # TODO: consider removing this check and allowing users to specify + # a number of units to prune that is greater than the number of units + # left to prune. In this case, the tensor will just be fully pruned. + + if isinstance(amount, numbers.Integral) and amount > tensor_size: + raise ValueError( + f"amount={amount} should be smaller than the number of parameters to prune={tensor_size}" + ) + + +def _validate_structured_pruning(t) -> None: + r"""Validate that the tensor to be pruned is at least 2-Dimensional. + + Validation helper to check that the tensor to be pruned is multi- + dimensional, such that the concept of "channels" is well-defined. + + Args: + t (torch.Tensor): tensor representing the parameter to prune + + Raises: + ValueError: if the tensor `t` is not at least 2D. + """ + shape = t.shape + if len(shape) <= 1: + raise ValueError( + "Structured pruning can only be applied to " + "multidimensional tensors. Found tensor of shape " + f"{shape} with {len(shape)} dims" + ) + + +def _compute_nparams_toprune(amount, tensor_size): + r"""Convert the pruning amount from a percentage to absolute value. + + Since amount can be expressed either in absolute value or as a + percentage of the number of units/channels in a tensor, this utility + function converts the percentage to absolute value to standardize + the handling of pruning. + + Args: + amount (int or float): quantity of parameters to prune. + If float, should be between 0.0 and 1.0 and represent the + fraction of parameters to prune. If int, it represents the + absolute number of parameters to prune. + tensor_size (int): absolute number of parameters in the tensor + to prune. + + Returns: + int: the number of units to prune in the tensor + """ + # incorrect type already checked in _validate_pruning_amount_init + if isinstance(amount, numbers.Integral): + return amount + else: + return round(amount * tensor_size) + + +def _validate_pruning_dim(t, dim) -> None: + r"""Validate that the pruning dimension is within the bounds of the tensor dimension. + + Args: + t (torch.Tensor): tensor representing the parameter to prune + dim (int): index of the dim along which we define channels to prune + """ + if dim >= t.dim(): + raise IndexError(f"Invalid index {dim} for tensor of size {t.shape}") + + +def _compute_norm(t, n, dim): + r"""Compute the L_n-norm of a tensor along all dimensions except for the specified dimension. + + The L_n-norm will be computed across all entries in tensor `t` along all dimension + except for the one identified by dim. + Example: if `t` is of shape, say, 3x2x4 and dim=2 (the last dim), + then norm will have Size [4], and each entry will represent the + `L_n`-norm computed using the 3x2=6 entries for each of the 4 channels. + + Args: + t (torch.Tensor): tensor representing the parameter to prune + n (int, float, inf, -inf, 'fro', 'nuc'): See documentation of valid + entries for argument p in torch.norm + dim (int): dim identifying the channels to prune + + Returns: + norm (torch.Tensor): L_n norm computed across all dimensions except + for `dim`. By construction, `norm.shape = t.shape[-1]`. + """ + # dims = all axes, except for the one identified by `dim` + dims = list(range(t.dim())) + # convert negative indexing + if dim < 0: + dim = dims[dim] + dims.remove(dim) + + norm = torch.norm(t, p=n, dim=dims) + return norm diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/rnn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/rnn.py new file mode 100644 index 0000000000000000000000000000000000000000..84e208921baf64aa5c8f8c72c736ce253d7f6ee0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/rnn.py @@ -0,0 +1,589 @@ +import warnings +from collections.abc import Callable, Iterable +from typing import Any, NamedTuple, TypeVar +from typing_extensions import Self + +import torch +from torch import _VF, Tensor +from torch.utils._typing_utils import copy_method_params + + +__all__ = [ + "PackedSequence", + "invert_permutation", + "pack_padded_sequence", + "pad_packed_sequence", + "pad_sequence", + "unpad_sequence", + "pack_sequence", + "unpack_sequence", +] + +_T = TypeVar("_T") +_R = TypeVar("_R") + + +class PackedSequence_(NamedTuple): + data: torch.Tensor + batch_sizes: torch.Tensor + sorted_indices: torch.Tensor | None + unsorted_indices: torch.Tensor | None + + +def bind(optional: _T | None, fn: Callable[[_T], _R]) -> _R | None: + if optional is None: + return None + return fn(optional) + + +class PackedSequence(PackedSequence_): + r"""Holds the data and list of :attr:`batch_sizes` of a packed sequence. + + All RNN modules accept packed sequences as inputs. + + Note: + Instances of this class should never be created manually. They are meant + to be instantiated by functions like :func:`pack_padded_sequence`. + + Batch sizes represent the number elements at each sequence step in + the batch, not the varying sequence lengths passed to + :func:`pack_padded_sequence`. For instance, given data ``abc`` and ``x`` + the :class:`PackedSequence` would contain data ``axbc`` with + ``batch_sizes=[2,1,1]``. + + Attributes: + data (Tensor): Tensor containing packed sequence + batch_sizes (Tensor): Tensor of integers holding + information about the batch size at each sequence step + sorted_indices (Tensor, optional): Tensor of integers holding how this + :class:`PackedSequence` is constructed from sequences. + unsorted_indices (Tensor, optional): Tensor of integers holding how this + to recover the original sequences with correct order. + + .. note:: + :attr:`data` can be on arbitrary device and of arbitrary dtype. + :attr:`sorted_indices` and :attr:`unsorted_indices` must be ``torch.int64`` + tensors on the same device as :attr:`data`. + + However, :attr:`batch_sizes` should always be a CPU ``torch.int64`` tensor. + + This invariant is maintained throughout :class:`PackedSequence` class, + and all functions that construct a :class:`PackedSequence` in PyTorch + (i.e., they only pass in tensors conforming to this constraint). + """ + + def __new__( + cls, + data: Tensor, + batch_sizes: Tensor | None = None, + sorted_indices: Tensor | None = None, + unsorted_indices: Tensor | None = None, + ) -> Self: + return super().__new__( + cls, + *_packed_sequence_init_args( + data, batch_sizes, sorted_indices, unsorted_indices + ), + ) + + # NOTE [ device and dtype of a PackedSequence ] + # + # See the note above in doc string (starting with ":attr:`data` can be on + # arbitrary device..."). + def pin_memory(self) -> Self: + # Why not convert `batch_sizes`? + # See NOTE [ device and dtype of a PackedSequence ] + return type(self)( + self.data.pin_memory(), + self.batch_sizes, + bind(self.sorted_indices, lambda t: t.pin_memory()), + bind(self.unsorted_indices, lambda t: t.pin_memory()), + ) + + @copy_method_params(torch.Tensor.to) + def to(self, *args: Any, **kwargs: Any) -> Self: + r"""Perform dtype and/or device conversion on `self.data`. + + It has similar signature as :meth:`torch.Tensor.to` + + .. note:: + + If the ``self.data`` Tensor already has the correct :class:`torch.dtype` + and :class:`torch.device`, then ``self`` is returned. + Otherwise, returns a copy with the desired configuration. + """ + + # Why not convert `batch_sizes`? + # See NOTE [ device and dtype of a PackedSequence ] + data = self.data.to(*args, **kwargs) + if data is self.data: + return self + else: + _device, _dtype, non_blocking, convert_to_format = torch._C._nn._parse_to( + *args, **kwargs + ) + + # Does not forward device or dtype arg/kwargs, device is set from data.device + def call_to(t: torch.Tensor) -> torch.Tensor: + return t.to( + data.device, + non_blocking=non_blocking, + memory_format=convert_to_format, + ) + + sorted_indices = bind(self.sorted_indices, call_to) + unsorted_indices = bind(self.unsorted_indices, call_to) + return type(self)(data, self.batch_sizes, sorted_indices, unsorted_indices) + + @copy_method_params(torch.Tensor.cuda) + def cuda(self, *args: Any, **kwargs: Any) -> Self: + # Tests to see if 'cuda' should be added to kwargs + ex = torch.tensor((), dtype=self.data.dtype, device=self.data.device).to( + *args, **kwargs + ) + if ex.is_cuda: + return self.to(*args, **kwargs) + kwargs["device"] = "cuda" + return self.to(*args, **kwargs) + + @copy_method_params(torch.Tensor.cpu) + def cpu(self, *args: Any, **kwargs: Any) -> Self: + ex = torch.tensor((), dtype=self.data.dtype, device=self.data.device).to( + *args, **kwargs + ) + if ex.device.type == "cpu": + return self.to(*args, **kwargs) + kwargs["device"] = "cpu" + return self.to(*args, **kwargs) + + def double(self) -> Self: + return self.to(dtype=torch.double) + + def float(self) -> Self: + return self.to(dtype=torch.float) + + def half(self) -> Self: + return self.to(dtype=torch.half) + + def long(self) -> Self: + return self.to(dtype=torch.long) + + def int(self) -> Self: + return self.to(dtype=torch.int) + + def short(self) -> Self: + return self.to(dtype=torch.short) + + def char(self) -> Self: + return self.to(dtype=torch.int8) + + def byte(self) -> Self: + return self.to(dtype=torch.uint8) + + @property + def is_cuda(self) -> bool: + r"""Return true if `self.data` stored on a gpu.""" + return self.data.is_cuda + + def is_pinned(self) -> bool: + r"""Return true if `self.data` stored on in pinned memory.""" + return self.data.is_pinned() + + +# TorchScript doesn't support constructors on named tuples, so we use this helper +# method to construct PackedSequence +def _packed_sequence_init_args( + data: Tensor, + batch_sizes: Tensor | None = None, + sorted_indices: Tensor | None = None, + unsorted_indices: Tensor | None = None, +) -> tuple[Tensor, Tensor, Tensor | None, Tensor | None]: + # NB: if unsorted_indices is provided, it should be the inverse permutation + # to sorted_indices. Don't assert it here because the PackedSequence ctor + # should only be used internally. + + if unsorted_indices is None: + unsorted_indices = invert_permutation(sorted_indices) + + # support being called as `PackedSequence(data, batch_sizes, sorted_indices)` + if batch_sizes is not None: + # TODO: Re-enable this check (.type isn't supported in TorchScript) + if batch_sizes.device.type != "cpu": + raise ValueError( + "batch_sizes should always be on CPU. " + "Instances of PackedSequence should never be created manually. " + "They should be instantiated by functions like pack_sequence " + "and pack_padded_sequences in nn.utils.rnn. " + "https://pytorch.org/docs/stable/nn.html#torch.nn.utils.rnn.pack_sequence" + ) + return data, batch_sizes, sorted_indices, unsorted_indices + + # support being called as `PackedSequence((data, batch_sizes), *, sorted_indices)` + else: + if not (isinstance(data, (list, tuple)) and len(data) == 2): + raise AssertionError("Expected data to be a list or tuple of length 2") + return data[0], data[1], sorted_indices, unsorted_indices + + +def _packed_sequence_init( + data: Tensor, + batch_sizes: Tensor | None = None, + sorted_indices: Tensor | None = None, + unsorted_indices: Tensor | None = None, +) -> PackedSequence: + data, batch_sizes, sorted_indices, unsorted_indices = _packed_sequence_init_args( + data, batch_sizes, sorted_indices, unsorted_indices + ) + return PackedSequence(data, batch_sizes, sorted_indices, unsorted_indices) + + +def invert_permutation(permutation: Tensor | None) -> Tensor | None: + """Returns the inverse of ``permutation``. + + This is useful for converting between sorted and unsorted indices in + a :class:`~nn.utils.rnn.PackedSequence`. + + Args: + permutation (Tensor, optional): a 1-D tensor of indices to invert + """ + if permutation is None: + return None + output = torch.empty_like(permutation, memory_format=torch.legacy_contiguous_format) + output.scatter_( + 0, permutation, torch.arange(0, permutation.numel(), device=permutation.device) + ) + return output + + +def pack_padded_sequence( + input: Tensor, + lengths: Tensor | list[int], + batch_first: bool = False, + enforce_sorted: bool = True, +) -> PackedSequence: + r"""Packs a Tensor containing padded sequences of variable length. + + :attr:`input` can be of size ``T x B x *`` (if :attr:`batch_first` is ``False``) + or ``B x T x *`` (if :attr:`batch_first` is ``True``) where ``T`` is the length + of the longest sequence, ``B`` is the batch size, and ``*`` is any number of dimensions + (including 0). + + For unsorted sequences, use `enforce_sorted = False`. If :attr:`enforce_sorted` is + ``True``, the sequences should be sorted by length in a decreasing order, i.e. + ``input[:,0]`` should be the longest sequence, and ``input[:,B-1]`` the shortest + one. `enforce_sorted = True` is only necessary for ONNX export. + + It is an inverse operation to :func:`pad_packed_sequence`, and hence :func:`pad_packed_sequence` + can be used to recover the underlying tensor packed in :class:`PackedSequence`. + + Note: + This function accepts any input that has at least two dimensions. You + can apply it to pack the labels, and use the output of the RNN with + them to compute the loss directly. A Tensor can be retrieved from + a :class:`PackedSequence` object by accessing its ``.data`` attribute. + + Args: + input (Tensor): padded batch of variable length sequences. + lengths (Tensor or list(int)): list of sequence lengths of each batch + element (must be on the CPU if provided as a tensor). + batch_first (bool, optional): if ``True``, the input is expected in ``B x T x *`` + format, ``T x B x *`` otherwise. Default: ``False``. + enforce_sorted (bool, optional): if ``True``, the input is expected to + contain sequences sorted by length in a decreasing order. If + ``False``, the input will get sorted unconditionally. Default: ``True``. + + .. warning:: + The dim of ``input`` tensor will be truncated if its length larger than + correspond value in ``length``. + + Returns: + a :class:`PackedSequence` object + """ + if not isinstance(lengths, torch.Tensor): + if torch._C._get_tracing_state(): + warnings.warn( + "pack_padded_sequence has been called with a Python list of " + "sequence lengths. The tracer cannot track the data flow of Python " + "values, and it will treat them as constants, likely rendering " + "the trace incorrect for any other combination of lengths.", + stacklevel=2, + ) + lengths = torch.as_tensor(lengths, dtype=torch.int64, device="cpu") + else: + lengths = lengths.to(dtype=torch.int64) + + if enforce_sorted: + sorted_indices = None + else: + lengths, sorted_indices = torch.sort(lengths, descending=True) + sorted_indices = sorted_indices.to(input.device) + batch_dim = 0 if batch_first else 1 + input = input.index_select(batch_dim, sorted_indices) + + data, batch_sizes = _VF._pack_padded_sequence(input, lengths, batch_first) + return _packed_sequence_init(data, batch_sizes, sorted_indices, None) + + +def pad_packed_sequence( + sequence: PackedSequence, + batch_first: bool = False, + padding_value: float = 0.0, + total_length: int | None = None, +) -> tuple[Tensor, Tensor]: + r"""Pad a packed batch of variable length sequences. + + It is an inverse operation to :func:`pack_padded_sequence`. + + The returned Tensor's data will be of size ``T x B x *`` (if :attr:`batch_first` is ``False``) + or ``B x T x *`` (if :attr:`batch_first` is ``True``) , where ``T`` is the length of the longest + sequence and ``B`` is the batch size. + + Example: + >>> from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence + >>> seq = torch.tensor([[1, 2, 0], [3, 0, 0], [4, 5, 6]]) + >>> lens = [2, 1, 3] + >>> packed = pack_padded_sequence( + ... seq, lens, batch_first=True, enforce_sorted=False + ... ) + >>> packed + PackedSequence(data=tensor([4, 1, 3, 5, 2, 6]), batch_sizes=tensor([3, 2, 1]), + sorted_indices=tensor([2, 0, 1]), unsorted_indices=tensor([1, 2, 0])) + >>> seq_unpacked, lens_unpacked = pad_packed_sequence(packed, batch_first=True) + >>> seq_unpacked + tensor([[1, 2, 0], + [3, 0, 0], + [4, 5, 6]]) + >>> lens_unpacked + tensor([2, 1, 3]) + + .. note:: + :attr:`total_length` is useful to implement the + ``pack sequence -> recurrent network -> unpack sequence`` pattern in a + :class:`~torch.nn.Module` wrapped in :class:`~torch.nn.DataParallel`. + See :ref:`this FAQ section ` for + details. + + Args: + sequence (PackedSequence): batch to pad + batch_first (bool, optional): if ``True``, the output will be in ``B x T x *`` + format, ``T x B x *`` otherwise. + padding_value (float, optional): values for padded elements. + total_length (int, optional): if not ``None``, the output will be padded to + have length :attr:`total_length`. This method will throw :class:`ValueError` + if :attr:`total_length` is less than the max sequence length in + :attr:`sequence`. + + Returns: + Tuple of Tensor containing the padded sequence, and a Tensor + containing the list of lengths of each sequence in the batch. + Batch elements will be re-ordered as they were ordered originally when + the batch was passed to ``pack_padded_sequence`` or ``pack_sequence``. + """ + max_seq_length = sequence.batch_sizes.size(0) + if total_length is not None: + if total_length < max_seq_length: + raise ValueError( + "Expected total_length to be at least the length " + "of the longest sequence in input, but got " + f"total_length={total_length} and max sequence length being {max_seq_length}" + ) + max_seq_length = total_length + padded_output, lengths = _VF._pad_packed_sequence( + sequence.data, sequence.batch_sizes, batch_first, padding_value, max_seq_length + ) + unsorted_indices = sequence.unsorted_indices + if unsorted_indices is not None: + batch_dim = 0 if batch_first else 1 + return ( + padded_output.index_select(batch_dim, unsorted_indices), + lengths[unsorted_indices.cpu()], + ) + return padded_output, lengths + + +# NOTE: for JIT-compatibility, we need to be more restrictive here and use specific types instead of Iterable. +def pad_sequence( + sequences: Tensor | list[Tensor], + batch_first: bool = False, + padding_value: float = 0.0, + padding_side: str = "right", +) -> Tensor: + r"""Pad a list of variable length Tensors with :attr:`padding_value`. + + ``pad_sequence`` stacks a list of Tensors along a new dimension, and pads them + to equal length. :attr:`sequences` can be list of sequences with size ``L x *``, + where `L` is length of the sequence and ``*`` is any number of dimensions + (including ``0``). If :attr:`batch_first` is ``False``, the output is of size + ``T x B x *``, and ``B x T x *`` otherwise, where ``B`` is the batch size + (the number of elements in :attr:`sequences`), ``T`` is the length of the longest + sequence. + + Example: + >>> from torch.nn.utils.rnn import pad_sequence + >>> a = torch.ones(25, 300) + >>> b = torch.ones(22, 300) + >>> c = torch.ones(15, 300) + >>> pad_sequence([a, b, c]).size() + torch.Size([25, 3, 300]) + + Note: + This function returns a Tensor of size ``T x B x *`` or ``B x T x *`` + where `T` is the length of the longest sequence. This function assumes + trailing dimensions and type of all the Tensors in sequences are same. + + Args: + sequences (list[Tensor]): list of variable length sequences. + batch_first (bool, optional): if ``True``, the output will be in ``B x T x *`` + format, ``T x B x *`` otherwise. + padding_value (float, optional): value for padded elements. Default: ``0``. + padding_side (str, optional): the side to pad the sequences on. + Default: ``'right'``. + + Returns: + Tensor of size ``T x B x *`` if :attr:`batch_first` is ``False``. + Tensor of size ``B x T x *`` otherwise + """ + if not (torch.jit.is_tracing() or torch.jit.is_scripting()): + # JIT doesn't support `Iterable` + if not isinstance(sequences, Iterable): + msg = ( + "pad_sequence: Expected iterable for input sequences, but got arg of type: " + f"{type(sequences)}" + ) + raise RuntimeError(msg) + + # In JIT context this leads to, + # RuntimeError: cannot statically infer the expected size of a list in this context + sequences = tuple(sequences) # type: ignore[assignment] + else: + # For JIT, we only support Union[Tensor, Tuple[Tensor]] + if isinstance(sequences, torch.Tensor): + sequences = sequences.unbind(0) # type: ignore[assignment] + + # assuming trailing dimensions and type of all the Tensors + # in sequences are same and fetching those from sequences[0] + return torch._C._nn.pad_sequence( + sequences, # type: ignore[arg-type] + batch_first, + padding_value, + padding_side, # type: ignore[arg-type] + ) + + +def unpad_sequence( + padded_sequences: Tensor, + lengths: Tensor, + batch_first: bool = False, +) -> list[Tensor]: + r"""Unpad padded Tensor into a list of variable length Tensors. + + ``unpad_sequence`` unstacks padded Tensor into a list of variable length Tensors. + + Example: + >>> from torch.nn.utils.rnn import pad_sequence, unpad_sequence + >>> a = torch.ones(25, 300) + >>> b = torch.ones(22, 300) + >>> c = torch.ones(15, 300) + >>> sequences = [a, b, c] + >>> padded_sequences = pad_sequence(sequences) + >>> lengths = torch.as_tensor([v.size(0) for v in sequences]) + >>> unpadded_sequences = unpad_sequence(padded_sequences, lengths) + >>> torch.allclose(sequences[0], unpadded_sequences[0]) + True + >>> torch.allclose(sequences[1], unpadded_sequences[1]) + True + >>> torch.allclose(sequences[2], unpadded_sequences[2]) + True + + Args: + padded_sequences (Tensor): padded sequences. + lengths (Tensor): length of original (unpadded) sequences. + batch_first (bool, optional): whether batch dimension first or not. Default: ``False``. + + Returns: + a list of :class:`Tensor` objects + """ + unpadded_sequences = [] + + if not batch_first: + padded_sequences.transpose_(0, 1) + + max_length = padded_sequences.shape[1] + idx = torch.arange(max_length, device=lengths.device) + + for seq, length in zip(padded_sequences, lengths, strict=True): + mask = idx < length + unpacked_seq = seq[mask] + unpadded_sequences.append(unpacked_seq) + + return unpadded_sequences + + +def pack_sequence( + sequences: list[Tensor], + enforce_sorted: bool = True, +) -> PackedSequence: + r"""Packs a list of variable length Tensors. + + Consecutive call of the next functions: ``pad_sequence``, ``pack_padded_sequence``. + + ``sequences`` should be a list of Tensors of size ``L x *``, where `L` is + the length of a sequence and `*` is any number of trailing dimensions, + including ``0``. + + For unsorted sequences, use `enforce_sorted = False`. If ``enforce_sorted`` + is ``True``, the sequences should be sorted in the order of decreasing length. + ``enforce_sorted = True`` is only necessary for ONNX export. + + Example: + >>> from torch.nn.utils.rnn import pack_sequence + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5]) + >>> c = torch.tensor([6]) + >>> pack_sequence([a, b, c]) + PackedSequence(data=tensor([1, 4, 6, 2, 5, 3]), batch_sizes=tensor([3, 2, 1]), sorted_indices=None, unsorted_indices=None) + + Args: + sequences (list[Tensor]): A list of sequences of decreasing length. + enforce_sorted (bool, optional): if ``True``, checks that the input + contains sequences sorted by length in a decreasing order. If + ``False``, this condition is not checked. Default: ``True``. + + Returns: + a :class:`PackedSequence` object + """ + lengths = torch.as_tensor([v.size(0) for v in sequences]) + return pack_padded_sequence( + pad_sequence(sequences), lengths, enforce_sorted=enforce_sorted + ) + + +def unpack_sequence(packed_sequences: PackedSequence) -> list[Tensor]: + r"""Unpack PackedSequence into a list of variable length Tensors. + + ``packed_sequences`` should be a PackedSequence object. + + Example: + >>> from torch.nn.utils.rnn import pack_sequence, unpack_sequence + >>> a = torch.tensor([1, 2, 3]) + >>> b = torch.tensor([4, 5]) + >>> c = torch.tensor([6]) + >>> sequences = [a, b, c] + >>> print(sequences) + [tensor([1, 2, 3]), tensor([4, 5]), tensor([6])] + >>> packed_sequences = pack_sequence(sequences) + >>> print(packed_sequences) + PackedSequence(data=tensor([1, 4, 6, 2, 5, 3]), batch_sizes=tensor([3, 2, 1]), sorted_indices=None, unsorted_indices=None) + >>> unpacked_sequences = unpack_sequence(packed_sequences) + >>> print(unpacked_sequences) + [tensor([1, 2, 3]), tensor([4, 5]), tensor([6])] + + Args: + packed_sequences (PackedSequence): A PackedSequence object. + + Returns: + a list of :class:`Tensor` objects + """ + padded_sequences, lengths = pad_packed_sequence(packed_sequences, batch_first=True) + unpacked_sequences = unpad_sequence(padded_sequences, lengths, batch_first=True) + return unpacked_sequences diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/spectral_norm.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/spectral_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..79fea5a50b79334b11c51c6d74b6cba4a1cc1ff6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/spectral_norm.py @@ -0,0 +1,364 @@ +# mypy: allow-untyped-defs +"""Spectral Normalization from https://arxiv.org/abs/1802.05957.""" + +from typing import Any, TypeVar + +import torch +import torch.nn.functional as F +from torch.nn.modules import Module + + +__all__ = [ + "SpectralNorm", + "SpectralNormLoadStateDictPreHook", + "SpectralNormStateDictHook", + "spectral_norm", + "remove_spectral_norm", +] + + +class SpectralNorm: + # Invariant before and after each forward call: + # u = F.normalize(W @ v) + # NB: At initialization, this invariant is not enforced + + _version: int = 1 + # At version 1: + # made `W` not a buffer, + # added `v` as a buffer, and + # made eval mode use `W = u @ W_orig @ v` rather than the stored `W`. + name: str + dim: int + n_power_iterations: int + eps: float + + def __init__( + self, + name: str = "weight", + n_power_iterations: int = 1, + dim: int = 0, + eps: float = 1e-12, + ) -> None: + self.name = name + self.dim = dim + if n_power_iterations <= 0: + raise ValueError( + "Expected n_power_iterations to be positive, but " + f"got n_power_iterations={n_power_iterations}" + ) + self.n_power_iterations = n_power_iterations + self.eps = eps + + def reshape_weight_to_matrix(self, weight: torch.Tensor) -> torch.Tensor: + weight_mat = weight + if self.dim != 0: + # permute dim to front + weight_mat = weight_mat.permute( + self.dim, *[d for d in range(weight_mat.dim()) if d != self.dim] + ) + height = weight_mat.size(0) + return weight_mat.reshape(height, -1) + + def compute_weight(self, module: Module, do_power_iteration: bool) -> torch.Tensor: + # NB: If `do_power_iteration` is set, the `u` and `v` vectors are + # updated in power iteration **in-place**. This is very important + # because in `DataParallel` forward, the vectors (being buffers) are + # broadcast from the parallelized module to each module replica, + # which is a new module object created on the fly. And each replica + # runs its own spectral norm power iteration. So simply assigning + # the updated vectors to the module this function runs on will cause + # the update to be lost forever. And the next time the parallelized + # module is replicated, the same randomly initialized vectors are + # broadcast and used! + # + # Therefore, to make the change propagate back, we rely on two + # important behaviors (also enforced via tests): + # 1. `DataParallel` doesn't clone storage if the broadcast tensor + # is already on correct device; and it makes sure that the + # parallelized module is already on `device[0]`. + # 2. If the out tensor in `out=` kwarg has correct shape, it will + # just fill in the values. + # Therefore, since the same power iteration is performed on all + # devices, simply updating the tensors in-place will make sure that + # the module replica on `device[0]` will update the _u vector on the + # parallelized module (by shared storage). + # + # However, after we update `u` and `v` in-place, we need to **clone** + # them before using them to normalize the weight. This is to support + # backproping through two forward passes, e.g., the common pattern in + # GAN training: loss = D(real) - D(fake). Otherwise, engine will + # complain that variables needed to do backward for the first forward + # (i.e., the `u` and `v` vectors) are changed in the second forward. + weight = getattr(module, self.name + "_orig") + u = getattr(module, self.name + "_u") + v = getattr(module, self.name + "_v") + weight_mat = self.reshape_weight_to_matrix(weight) + + if do_power_iteration: + with torch.no_grad(): + for _ in range(self.n_power_iterations): + # Spectral norm of weight equals to `u^T W v`, where `u` and `v` + # are the first left and right singular vectors. + # This power iteration produces approximations of `u` and `v`. + v = F.normalize( + torch.mv(weight_mat.t(), u), dim=0, eps=self.eps, out=v + ) + u = F.normalize(torch.mv(weight_mat, v), dim=0, eps=self.eps, out=u) + if self.n_power_iterations > 0: + # See above on why we need to clone + u = u.clone(memory_format=torch.contiguous_format) + v = v.clone(memory_format=torch.contiguous_format) + + sigma = torch.dot(u, torch.mv(weight_mat, v)) + weight = weight / sigma + return weight + + def remove(self, module: Module) -> None: + with torch.no_grad(): + weight = self.compute_weight(module, do_power_iteration=False) + delattr(module, self.name) + delattr(module, self.name + "_u") + delattr(module, self.name + "_v") + delattr(module, self.name + "_orig") + module.register_parameter(self.name, torch.nn.Parameter(weight.detach())) + + def __call__(self, module: Module, inputs: Any) -> None: + setattr( + module, + self.name, + self.compute_weight(module, do_power_iteration=module.training), + ) + + def _solve_v_and_rescale(self, weight_mat, u, target_sigma): + # Tries to returns a vector `v` s.t. `u = F.normalize(W @ v)` + # (the invariant at top of this class) and `u @ W @ v = sigma`. + # This uses pinverse in case W^T W is not invertible. + v = torch.linalg.multi_dot( + [weight_mat.t().mm(weight_mat).pinverse(), weight_mat.t(), u.unsqueeze(1)] + ).squeeze(1) + return v.mul_(target_sigma / torch.dot(u, torch.mv(weight_mat, v))) + + @staticmethod + def apply( + module: Module, name: str, n_power_iterations: int, dim: int, eps: float + ) -> "SpectralNorm": + for hook in module._forward_pre_hooks.values(): + if isinstance(hook, SpectralNorm) and hook.name == name: + raise RuntimeError( + f"Cannot register two spectral_norm hooks on the same parameter {name}" + ) + + fn = SpectralNorm(name, n_power_iterations, dim, eps) + weight = module._parameters[name] + if weight is None: + raise ValueError( + f"`SpectralNorm` cannot be applied as parameter `{name}` is None" + ) + if isinstance(weight, torch.nn.parameter.UninitializedParameter): + raise ValueError( + "The module passed to `SpectralNorm` can't have uninitialized parameters. " + "Make sure to run the dummy forward before applying spectral normalization" + ) + + with torch.no_grad(): + weight_mat = fn.reshape_weight_to_matrix(weight) + + h, w = weight_mat.size() + # randomly initialize `u` and `v` + u = F.normalize(weight.new_empty(h).normal_(0, 1), dim=0, eps=fn.eps) + v = F.normalize(weight.new_empty(w).normal_(0, 1), dim=0, eps=fn.eps) + + delattr(module, fn.name) + module.register_parameter(fn.name + "_orig", weight) + # We still need to assign weight back as fn.name because all sorts of + # things may assume that it exists, e.g., when initializing weights. + # However, we can't directly assign as it could be an nn.Parameter and + # gets added as a parameter. Instead, we register weight.data as a plain + # attribute. + setattr(module, fn.name, weight.data) + module.register_buffer(fn.name + "_u", u) + module.register_buffer(fn.name + "_v", v) + + module.register_forward_pre_hook(fn) + module._register_state_dict_hook(SpectralNormStateDictHook(fn)) + module._register_load_state_dict_pre_hook(SpectralNormLoadStateDictPreHook(fn)) + return fn + + +class SpectralNormLoadStateDictPreHook: + # See docstring of SpectralNorm._version on the changes to spectral_norm. + def __init__(self, fn) -> None: + self.fn = fn + + # For state_dict with version None, (assuming that it has gone through at + # least one training forward), we have + # + # u = F.normalize(W_orig @ v) + # W = W_orig / sigma, where sigma = u @ W_orig @ v + # + # To compute `v`, we solve `W_orig @ x = u`, and let + # v = x / (u @ W_orig @ x) * (W / W_orig). + def __call__( + self, + state_dict, + prefix, + local_metadata, + strict, + missing_keys, + unexpected_keys, + error_msgs, + ) -> None: + fn = self.fn + version = local_metadata.get("spectral_norm", {}).get( + fn.name + ".version", None + ) + if version is None or version < 1: + weight_key = prefix + fn.name + if ( + version is None + and all(weight_key + s in state_dict for s in ("_orig", "_u", "_v")) + and weight_key not in state_dict + ): + # Detect if it is the updated state dict and just missing metadata. + # This could happen if the users are crafting a state dict themselves, + # so we just pretend that this is the newest. + return + has_missing_keys = False + for suffix in ("_orig", "", "_u"): + key = weight_key + suffix + if key not in state_dict: + has_missing_keys = True + if strict: + missing_keys.append(key) + if has_missing_keys: + return + with torch.no_grad(): + weight_orig = state_dict[weight_key + "_orig"] + weight = state_dict.pop(weight_key) + sigma = (weight_orig / weight).mean() + weight_mat = fn.reshape_weight_to_matrix(weight_orig) + u = state_dict[weight_key + "_u"] + v = fn._solve_v_and_rescale(weight_mat, u, sigma) + state_dict[weight_key + "_v"] = v + + +class SpectralNormStateDictHook: + # See docstring of SpectralNorm._version on the changes to spectral_norm. + def __init__(self, fn) -> None: + self.fn = fn + + def __call__(self, module, state_dict, prefix, local_metadata) -> None: + if "spectral_norm" not in local_metadata: + local_metadata["spectral_norm"] = {} + key = self.fn.name + ".version" + if key in local_metadata["spectral_norm"]: + raise RuntimeError(f"Unexpected key in metadata['spectral_norm']: {key}") + local_metadata["spectral_norm"][key] = self.fn._version + + +T_module = TypeVar("T_module", bound=Module) + + +def spectral_norm( + module: T_module, + name: str = "weight", + n_power_iterations: int = 1, + eps: float = 1e-12, + dim: int | None = None, +) -> T_module: + r"""Apply spectral normalization to a parameter in the given module. + + .. math:: + \mathbf{W}_{SN} = \dfrac{\mathbf{W}}{\sigma(\mathbf{W})}, + \sigma(\mathbf{W}) = \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2} + + Spectral normalization stabilizes the training of discriminators (critics) + in Generative Adversarial Networks (GANs) by rescaling the weight tensor + with spectral norm :math:`\sigma` of the weight matrix calculated using + power iteration method. If the dimension of the weight tensor is greater + than 2, it is reshaped to 2D in power iteration method to get spectral + norm. This is implemented via a hook that calculates spectral norm and + rescales weight before every :meth:`~Module.forward` call. + + See `Spectral Normalization for Generative Adversarial Networks`_ . + + .. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957 + + Args: + module (nn.Module): containing module + name (str, optional): name of weight parameter + n_power_iterations (int, optional): number of power iterations to + calculate spectral norm + eps (float, optional): epsilon for numerical stability in + calculating norms + dim (int, optional): dimension corresponding to number of outputs, + the default is ``0``, except for modules that are instances of + ConvTranspose{1,2,3}d, when it is ``1`` + + Returns: + The original module with the spectral norm hook + + .. note:: + This function has been reimplemented as + :func:`torch.nn.utils.parametrizations.spectral_norm` using the new + parametrization functionality in + :func:`torch.nn.utils.parametrize.register_parametrization`. Please use + the newer version. This function will be deprecated in a future version + of PyTorch. + + Example:: + + >>> m = spectral_norm(nn.Linear(20, 40)) + >>> m + Linear(in_features=20, out_features=40, bias=True) + >>> m.weight_u.size() + torch.Size([40]) + + """ + if dim is None: + if isinstance( + module, + ( + torch.nn.ConvTranspose1d, + torch.nn.ConvTranspose2d, + torch.nn.ConvTranspose3d, + ), + ): + dim = 1 + else: + dim = 0 + SpectralNorm.apply(module, name, n_power_iterations, dim, eps) + + return module + + +def remove_spectral_norm(module: T_module, name: str = "weight") -> T_module: + r"""Remove the spectral normalization reparameterization from a module. + + Args: + module (Module): containing module + name (str, optional): name of weight parameter + + Example: + >>> m = spectral_norm(nn.Linear(40, 10)) + >>> remove_spectral_norm(m) + """ + for k, hook in module._forward_pre_hooks.items(): + if isinstance(hook, SpectralNorm) and hook.name == name: + hook.remove(module) + del module._forward_pre_hooks[k] + break + else: + raise ValueError(f"spectral_norm of '{name}' not found in {module}") + + for k, hook in module._state_dict_hooks.items(): + if isinstance(hook, SpectralNormStateDictHook) and hook.fn.name == name: + del module._state_dict_hooks[k] + break + + for k, hook in module._load_state_dict_pre_hooks.items(): + if isinstance(hook, SpectralNormLoadStateDictPreHook) and hook.fn.name == name: + del module._load_state_dict_pre_hooks[k] + break + + return module diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/stateless.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/stateless.py new file mode 100644 index 0000000000000000000000000000000000000000..70f0afdeb52923a029a1843e1f2cfc702ab7473b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/stateless.py @@ -0,0 +1,279 @@ +# mypy: allow-untyped-defs +import contextlib +from typing import Any +from typing_extensions import deprecated + +import torch +from torch import Tensor +from torch.nn.utils._named_member_accessor import NamedMemberAccessor + + +__all__ = ["functional_call"] + + +def _untie_named_tensors_map( + module: "torch.nn.Module", + parameters_and_buffers: dict[str, Tensor], +) -> dict[str, Tensor]: + """ + Unties all tied tensors in the module to parameters_and_buffers. + + This function returns a new untied_parameters_and_buffers dictionary and leave the original + untied_parameters_and_buffers dictionary unchanged. It adds new (missing) keys for tied tensors + in the module to untied_parameters_and_buffers. The value of the new key is the user-given value + in the original parameters_and_buffers dictionary. + + If there are more than one user-given values for the same tied tensor, it will raise an error. + + For example, if the module has two tied weights self.foo and self.tied_foo and the user passes + {'foo': foo_value, ...}, this will return {'foo': foo_value, 'tied_foo': foo_value, ...}. If the + user passes {'foo': foo_value, 'tied_foo': tied_foo_value, ...}, it will raise an error. If the + user passes {'foo': foo_value, 'tied_foo': foo_value, ...}, it will not raise an error. + + Args: + module (torch.nn.Module): the module to determine which tensors are tied. + parameters_and_buffers (Dict[str, Tensor]): a map of {name: tensor} for reparamaterizing the module. + + Returns: + A new untied version of the parameters_and_buffers dictionary. + + Raises: + ValueError: if there are more than one user-given values for the same tied tensor. + """ + # A map of {name: tensor} for all tensors (including tied ones) in the module. + all_named_tensors: dict[str, Tensor] = {} + all_named_tensors.update(module.named_parameters(remove_duplicate=False)) + all_named_tensors.update(module.named_buffers(remove_duplicate=False)) + + # A map of {tensor: set(all_tied_names)} for all tensor names in the module. + tensor_to_tied_names_map: dict[Tensor, set[str]] = {} + for name, tensor in all_named_tensors.items(): + if tensor not in tensor_to_tied_names_map: + tensor_to_tied_names_map[tensor] = set() + tensor_to_tied_names_map[tensor].add(name) + + # A map of {tied_name: set(all_tied_names)} for all tensor names in the module. + # If a name is not tied, it will not be in this map. + tied_names_map: dict[str, set[str]] = {} + for tied_names in tensor_to_tied_names_map.values(): + if len(tied_names) > 1: + for tied_name in tied_names: + tied_names_map[tied_name] = tied_names + + # Make sure the user didn't pass multiple values for the same tied tensor. + given_names = set(parameters_and_buffers.keys()) + # same as given_names.intersection(tied_names_map.keys()) but dynamo can't + # handle that + given_names_for_tied_tensors: set[str] = set() + for name in given_names: + if name in tied_names_map: + given_names_for_tied_tensors.add(name) + + for given_name in given_names_for_tied_tensors: + tied_names = tied_names_map[given_name] + if ( + # Detect if there are multiple keys present for the same tied tensor. + len(tied_names.intersection(given_names_for_tied_tensors)) > 1 + # Only raise an error if the user passed multiple values for the same tied tensor. + # If all given values are the same, don't raise. + and len({parameters_and_buffers[tied_name] for tied_name in tied_names}) + != 1 + ): + raise ValueError( + f"functional_call got multiple values for keys {sorted(tied_names)}, " + f"which are tied. Consider using tie_weights=False" + ) + + # Untie the given named tensor map + # Make a copy for not modifying the original dict + untied_parameters_and_buffers = parameters_and_buffers.copy() + for given_name in given_names_for_tied_tensors: + for tied_name in tied_names_map[given_name]: + untied_parameters_and_buffers[tied_name] = parameters_and_buffers[ + given_name + ] + return untied_parameters_and_buffers + + +@contextlib.contextmanager +def _reparametrize_module( + module: "torch.nn.Module", + parameters_and_buffers: dict[str, Tensor], + tie_weights: bool = False, + strict: bool = False, + stack_weights: bool = False, +): + if tie_weights: + untied_parameters_and_buffers = _untie_named_tensors_map( + module, parameters_and_buffers + ) + else: + untied_parameters_and_buffers = parameters_and_buffers + + accessor = NamedMemberAccessor(module) + if strict: + missing_keys, unexpected_keys = accessor.check_keys( + untied_parameters_and_buffers + ) + error_msgs = [] + if len(unexpected_keys) > 0: + error_msgs.append( + f"Unexpected key(s): {', '.join(map(repr, unexpected_keys))}." + ) + if len(missing_keys) > 0: + error_msgs.append(f"Missing key(s): {', '.join(map(repr, missing_keys))}.") + if len(error_msgs) > 0: + raise RuntimeError( + "Error(s) in reparametrizing for {}:\n\t{}".format( + module._get_name(), "\n\t".join(error_msgs) + ) + ) + + orig_parameters_and_buffers: dict[str, Tensor] = {} + try: + orig_parameters_and_buffers, _ = accessor.swap_tensors_dict( + untied_parameters_and_buffers, allow_missing=True + ) + yield + finally: + if stack_weights: + # When stacking is enabled, we will restore the weights in LIFO order. + orig_parameters_and_buffers = dict( + reversed(orig_parameters_and_buffers.items()) + ) + new_parameters_and_buffers, _ = accessor.swap_tensors_dict( + orig_parameters_and_buffers, allow_missing=True + ) + # Sometimes the module is not completely stateless and has some in-place modifications on + # the _parameters and _buffers dictionaries. + # Write the changed parameters and buffers back to the original dict. + parameters_and_buffers.update( + { + k: new_parameters_and_buffers[k] + for k in parameters_and_buffers + if k in new_parameters_and_buffers + } + ) + + +@deprecated( + "`torch.nn.utils.stateless.functional_call` is deprecated as of PyTorch 2.0 " + "and will be removed in a future version of PyTorch. " + "Please use `torch.func.functional_call` instead which is a drop-in replacement.", + category=FutureWarning, +) +def functional_call( + module: "torch.nn.Module", + parameters_and_buffers: dict[str, Tensor], + args: Any | tuple | None = None, + kwargs: dict[str, Any] | None = None, + *, + tie_weights: bool = True, + strict: bool = False, +): + r"""Perform a functional call on the module by replacing the module parameters and buffers with the provided ones. + + .. warning:: + + This API is deprecated as of PyTorch 2.0 and will be removed in a future + version of PyTorch. Please use :func:`torch.func.functional_call` instead, + which is a drop-in replacement for this API. + + .. note:: If the module has active parametrizations, passing a value in the + :attr:`parameters_and_buffers` argument with the name set to the regular parameter + name will completely disable the parametrization. + If you want to apply the parametrization function to the value passed + please set the key as ``{submodule_name}.parametrizations.{parameter_name}.original``. + + .. note:: If the module performs in-place operations on parameters/buffers, these will be reflected + in the `parameters_and_buffers` input. + + Example:: + + >>> a = {'foo': torch.zeros(())} + >>> # xdoctest: +SKIP + >>> mod = Foo() # does self.foo = self.foo + 1 + >>> print(mod.foo) # tensor(0.) + >>> functional_call(mod, a, torch.ones(())) + >>> print(mod.foo) # tensor(0.) + >>> print(a['foo']) # tensor(1.) + + .. note:: If the module has tied weights, whether or not functional_call respects the tying is determined by the + tie_weights flag. + + Example:: + + >>> a = {'foo': torch.zeros(())} + >>> # xdoctest: +SKIP + >>> mod = Foo() # has both self.foo and self.foo_tied which are tied. Returns x + self.foo + self.foo_tied + >>> print(mod.foo) # tensor(1.) + >>> mod(torch.zeros(())) # tensor(2.) + >>> functional_call(mod, a, torch.zeros(())) # tensor(0.) since it will change self.foo_tied too + >>> functional_call(mod, a, torch.zeros(()), tie_weights=False) # tensor(1.)--self.foo_tied is not updated + >>> new_a = {'foo': torch.zeros(()), 'foo_tied': torch.zeros(())} + >>> functional_call(mod, new_a, torch.zeros()) # tensor(0.) + + Args: + module (torch.nn.Module): the module to call + parameters_and_buffers (dict of str and Tensor): the parameters that will be used in + the module call. + args (Any or tuple): arguments to be passed to the module call. If not a tuple, considered a single argument. + kwargs (dict): keyword arguments to be passed to the module call + tie_weights (bool, optional): If True, then parameters and buffers tied in the original model will be treated as + tied in the reparamaterized version. Therefore, if True and different values are passed for the tied + parameters and buffers, it will error. If False, it will not respect the originally tied parameters and + buffers unless the values passed for both weights are the same. Default: True. + strict (bool, optional): If True, then the parameters and buffers passed in must match the parameters and + buffers in the original module. Therefore, if True and there are any missing or unexpected keys, it will + error. Default: False. + + Returns: + Any: the result of calling ``module``. + """ + return _functional_call( + module, + parameters_and_buffers, + args, + kwargs, + tie_weights=tie_weights, + strict=strict, + ) + + +def _functional_call( + module: "torch.nn.Module", + parameters_and_buffers: dict[str, Tensor], + args: Any | tuple | None = None, + kwargs: dict[str, Any] | None = None, + *, + tie_weights: bool = True, + strict: bool = False, +): + # TODO allow kwargs such as unsafe and others for parametrization + if ( + torch.jit.is_tracing() + or torch.jit.is_scripting() + or isinstance( + module, + ( + torch.jit.RecursiveScriptModule, + torch.jit.ScriptModule, + torch.jit.ScriptFunction, + ), + ) + ): + raise RuntimeError("The stateless API can't be used with Jitted modules") + if isinstance(module, torch.nn.DataParallel): + raise RuntimeError( + "The stateless API can't be used with nn.DataParallel module" + ) + if kwargs is None: + kwargs = {} + if args is None: + args = () + elif not isinstance(args, tuple): + args = (args,) + with _reparametrize_module( + module, parameters_and_buffers, tie_weights=tie_weights, strict=strict + ): + return module(*args, **kwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/weight_norm.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/weight_norm.py new file mode 100644 index 0000000000000000000000000000000000000000..64b6b408b4de464be3f41e3e7e2b6e33bb0f1a54 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/nn/utils/weight_norm.py @@ -0,0 +1,165 @@ +# mypy: allow-untyped-defs +r"""Weight Normalization from https://arxiv.org/abs/1602.07868.""" + +from typing import Any, TypeVar +from typing_extensions import deprecated + +from torch import _weight_norm, norm_except_dim +from torch.nn.modules import Module +from torch.nn.parameter import Parameter, UninitializedParameter + + +__all__ = ["WeightNorm", "weight_norm", "remove_weight_norm"] + + +class WeightNorm: + name: str + dim: int + + def __init__(self, name: str, dim: int) -> None: + if dim is None: + dim = -1 + self.name = name + self.dim = dim + + # TODO Make return type more specific + def compute_weight(self, module: Module) -> Any: + g = getattr(module, self.name + "_g") + v = getattr(module, self.name + "_v") + return _weight_norm(v, g, self.dim) + + @staticmethod + @deprecated( + "`torch.nn.utils.weight_norm` is deprecated " + "in favor of `torch.nn.utils.parametrizations.weight_norm`.", + category=FutureWarning, + ) + def apply(module, name: str, dim: int) -> "WeightNorm": + for hook in module._forward_pre_hooks.values(): + if isinstance(hook, WeightNorm) and hook.name == name: + raise RuntimeError( + f"Cannot register two weight_norm hooks on the same parameter {name}" + ) + + if dim is None: + dim = -1 + + fn = WeightNorm(name, dim) + + weight = getattr(module, name) + if isinstance(weight, UninitializedParameter): + raise ValueError( + "The module passed to `WeightNorm` can't have uninitialized parameters. " + "Make sure to run the dummy forward before applying weight normalization" + ) + # remove w from parameter list + del module._parameters[name] + + # add g and v as new parameters and express w as g/||v|| * v + module.register_parameter( + name + "_g", Parameter(norm_except_dim(weight, 2, dim).data) + ) + module.register_parameter(name + "_v", Parameter(weight.data)) + setattr(module, name, fn.compute_weight(module)) + + # recompute weight before every forward() + module.register_forward_pre_hook(fn) + + return fn + + def remove(self, module: Module) -> None: + weight = self.compute_weight(module) + delattr(module, self.name) + del module._parameters[self.name + "_g"] + del module._parameters[self.name + "_v"] + setattr(module, self.name, Parameter(weight.data)) + + def __call__(self, module: Module, inputs: Any) -> None: + setattr(module, self.name, self.compute_weight(module)) + + +T_module = TypeVar("T_module", bound=Module) + + +def weight_norm(module: T_module, name: str = "weight", dim: int = 0) -> T_module: + r"""Apply weight normalization to a parameter in the given module. + + .. math:: + \mathbf{w} = g \dfrac{\mathbf{v}}{\|\mathbf{v}\|} + + Weight normalization is a reparameterization that decouples the magnitude + of a weight tensor from its direction. This replaces the parameter specified + by :attr:`name` (e.g. ``'weight'``) with two parameters: one specifying the magnitude + (e.g. ``'weight_g'``) and one specifying the direction (e.g. ``'weight_v'``). + Weight normalization is implemented via a hook that recomputes the weight + tensor from the magnitude and direction before every :meth:`~Module.forward` + call. + + By default, with ``dim=0``, the norm is computed independently per output + channel/plane. To compute a norm over the entire weight tensor, use + ``dim=None``. + + See https://arxiv.org/abs/1602.07868 + + .. warning:: + + This function is deprecated. Use :func:`torch.nn.utils.parametrizations.weight_norm` + which uses the modern parametrization API. The new ``weight_norm`` is compatible + with ``state_dict`` generated from old ``weight_norm``. + + Migration guide: + + * The magnitude (``weight_g``) and direction (``weight_v``) are now expressed + as ``parametrizations.weight.original0`` and ``parametrizations.weight.original1`` + respectively. If this is bothering you, please comment on + https://github.com/pytorch/pytorch/issues/102999 + + * To remove the weight normalization reparameterization, use + :func:`torch.nn.utils.parametrize.remove_parametrizations`. + + * The weight is no longer recomputed once at module forward; instead, it will + be recomputed on every access. To restore the old behavior, use + :func:`torch.nn.utils.parametrize.cached` before invoking the module + in question. + + Args: + module (Module): containing module + name (str, optional): name of weight parameter + dim (int, optional): dimension over which to compute the norm + + Returns: + The original module with the weight norm hook + + Example:: + + >>> m = weight_norm(nn.Linear(20, 40), name='weight') + >>> m + Linear(in_features=20, out_features=40, bias=True) + >>> m.weight_g.size() + torch.Size([40, 1]) + >>> m.weight_v.size() + torch.Size([40, 20]) + + """ + WeightNorm.apply(module, name, dim) + return module + + +def remove_weight_norm(module: T_module, name: str = "weight") -> T_module: + r"""Remove the weight normalization reparameterization from a module. + + Args: + module (Module): containing module + name (str, optional): name of weight parameter + + Example: + >>> m = weight_norm(nn.Linear(20, 40)) + >>> remove_weight_norm(m) + """ + for k, hook in module._forward_pre_hooks.items(): + if isinstance(hook, WeightNorm) and hook.name == name: + hook.remove(module) + del module._forward_pre_hooks[k] + return module + + raise ValueError(f"weight_norm of '{name}' not found in {module}") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/numa/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/numa/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/numa/binding.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/numa/binding.py new file mode 100644 index 0000000000000000000000000000000000000000..616dd848797de2ddf7f45ce7e59a1573294c7ea3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/numa/binding.py @@ -0,0 +1,727 @@ +""" +In NUMA (Non-Uniform Memory Access) systems, accessing memory on remote NUMA +nodes incurs additional latency. PyTorch provides NUMA binding utilities to +promote memory locality by binding worker processes to CPUs near their assigned GPUs. + +In practice, NUMA binding typically results in 1-10% overall performance improvements, +but some workloads may obtain much greater benefits or none at all. + +To enable NUMA binding, use the ``--numa-binding`` flag with :ref:`torchrun `, e.g.: + +.. code-block:: bash + + torchrun --numa-binding=node --nproc_per_node=8 train.py + +Alternatively, pass :class:`NumaOptions` to ``LaunchConfig`` +when using ``elastic_launch``. + +See :class:`AffinityMode` for available binding modes. +""" + +import os +import shutil +import traceback +from collections import defaultdict +from collections.abc import Callable, Iterable +from dataclasses import asdict, dataclass +from enum import Enum +from functools import wraps +from logging import getLogger +from typing import ParamSpec, TypeVar + +import torch +from torch._utils_internal import signpost_event + + +__all__ = [ + "AffinityMode", + "NumaOptions", +] + +logger = getLogger(__name__) + + +class AffinityMode(str, Enum): + NODE = "node" + """ + Each worker process and its threads will be bound to all the CPUs + on the NUMA node containing the GPU whose local index equals the worker's local rank. + If in doubt, use this option rather than the others. + + **Ex.:** If GPU 3 (i.e. ``torch.device("cuda:3")``) lives on NUMA node 1, then the worker + whose local rank is 3 will only be able to run on the CPUs of NUMA node 1. + """ + + SOCKET = "socket" + """ + Each worker process and its threads will be bound to all the CPUs on all the NUMA nodes of the + socket containing the GPU whose local index equals the worker's local rank. + + **Ex.:** If socket 0 contains GPU 3 and NUMA nodes 0-1, then the worker whose + local rank is 3 will be bound to the CPUs of NUMA nodes 0-1. + + For cases where there is only one NUMA node per socket anyway, this is equivalent to NODE. + """ + + EXCLUSIVE = "exclusive" + """ + Each worker process and its threads will be bound to an exclusive subset of CPUs + on the NUMA node containing the GPU whose local index equals the worker's local rank. + The CPUs on the NUMA node are divided evenly among all GPUs on that node, so no two + workers share the same CPU cores. + + **Ex.:** If NUMA node 1 has 16 physical cores and GPUs 2 and 3, then the worker whose + local rank is 2 will be bound to cores 0-7, and the worker whose local rank is 3 will + be bound to cores 8-15. + """ + + CORE_COMPLEX = "core-complex" + """ + Each worker process and its threads will be bound to a single core complex (a group of cores + sharing the same L3 cache) on the NUMA node containing the GPU whose local index equals + the worker's local rank. Each worker is bound to a different core complex when possible. + + **Ex.:** If NUMA node 1 has two core complexes (cores 0-7 sharing one L3 cache, cores 8-15 + sharing another) and GPUs 2 and 3, then the worker whose local rank is 2 will be bound to + cores 0-7, and the worker whose local rank is 3 will be bound to cores 8-15. + """ + + +@dataclass(frozen=True) +class NumaOptions: + affinity_mode: AffinityMode + + should_fall_back_if_binding_fails: bool = False + """ + If ``True``, we will silence any exceptions that occur during NUMA binding itself + rather than raising them. + + There are no expected exceptions, so avoid using this option. Its purpose is simply + to mitigate crash risk while conducting mass rollouts of NUMA binding. + """ + + +def _maybe_wrap_command_args_with_numa_binding( + command_args: tuple[str, ...], + *, + gpu_index: int, + numa_options: NumaOptions | None, +) -> tuple[str, ...]: + """ + Wraps command arguments with numactl to apply NUMA CPU binding. + + This function prepends numactl with appropriate CPU affinity flags to the + provided command arguments, binding the process to CPUs associated with + the specified GPU's NUMA node. + + Args: + command_args: The original command arguments to wrap. + gpu_index: The index of the GPU that will be used by the subprocess. + numa_options: Configuration for NUMA binding behavior. If None, returns + the original command_args unchanged. + + Returns: + Tuple of command arguments, potentially wrapped with numactl for NUMA binding. + Returns the original command_args if numa_options is None or if binding fails + and fallback is enabled. + """ + if numa_options is None: + return command_args + + kwargs = { + "command_args": command_args, + "gpu_index": gpu_index, + "numa_options": asdict(numa_options), + } + + try: + logical_cpu_indices = _get_validated_logical_cpus_to_bind_to( + gpu_index=gpu_index, + numa_options=numa_options, + ) + + wrapped_command_args = _assemble_numactl_command_args( + original_command_args=command_args, + logical_cpu_indices=logical_cpu_indices, + ) + signpost_event( + category="numa_binding", + name="apply_success", + parameters={ + **kwargs, + "wrapped_command": wrapped_command_args, + }, + ) + return wrapped_command_args + except Exception: + # pyrefly: ignore [bad-argument-type] + _handle_exception(numa_options=numa_options, logger_kwargs=kwargs) + return command_args + + +_TParams = ParamSpec("_TParams") +_TReturn = TypeVar("_TReturn") + + +def _maybe_wrap_with_numa_binding( + func: Callable[_TParams, _TReturn], + *, + gpu_index: int, + numa_options: NumaOptions | None, +) -> Callable[_TParams, _TReturn]: + """ + Wraps a function to apply NUMA CPU binding before execution. + + This decorator applies NUMA CPU affinity to all threads in the current process + before calling the wrapped function, binding them to CPUs associated with the + specified GPU's NUMA node. + + Args: + func: The function to wrap with NUMA binding. + gpu_index: The index of the GPU that will be used. + numa_options: Configuration for NUMA binding behavior. If None, returns + the original function unchanged. + + Returns: + A wrapped function that applies NUMA binding before execution, or the + original function if numa_options is None. + """ + if numa_options is None: + return func + + @wraps(func) + def wrapped(*args: _TParams.args, **kwargs: _TParams.kwargs) -> _TReturn: + _maybe_apply_numa_binding_to_current_process( + gpu_index=gpu_index, + # pyrefly: ignore [bad-argument-type] + numa_options=numa_options, + ) + return func(*args, **kwargs) + + return wrapped + + +def _maybe_apply_numa_binding_to_current_process( + *, gpu_index: int, numa_options: NumaOptions +) -> None: + kwargs = { + "gpu_index": gpu_index, + "numa_options": asdict(numa_options), + } + + try: + logical_cpu_indices = _get_validated_logical_cpus_to_bind_to( + gpu_index=gpu_index, + numa_options=numa_options, + ) + + _bind_all_threads_in_current_process_to_logical_cpus( + logical_cpu_indices=logical_cpu_indices + ) + + signpost_event( + category="numa_binding", + name="apply_success", + parameters={ + **kwargs, + "logical_cpu_indices": _get_ranges_str_from_ints(logical_cpu_indices), + }, + ) + except Exception: + # pyrefly: ignore [bad-argument-type] + _handle_exception(numa_options=numa_options, logger_kwargs=kwargs) + + +def _assemble_numactl_command_args( + *, original_command_args: tuple[str, ...], logical_cpu_indices: set[int] +) -> tuple[str, ...]: + return ( + "numactl", + f"--physcpubind={_get_ranges_str_from_ints(logical_cpu_indices)}", + *original_command_args, + ) + + +def _handle_exception( + *, numa_options: NumaOptions, logger_kwargs: dict[str, object] +) -> None: + signpost_event( + category="numa_binding", + name="apply_exception", + parameters={ + **logger_kwargs, + "traceback": traceback.format_exc(), + }, + ) + logger.exception("Failed to apply NUMA binding for input=%r", logger_kwargs) + if numa_options.should_fall_back_if_binding_fails: + logger.warning( + "Continuing executing without applying NUMA binding, despite exception %s", + traceback.format_exc(), + ) + return + # This function is called within an except block, so silence the warning + # about raise without an exception. + raise # noqa: PLE0704 + + +def _get_validated_logical_cpus_to_bind_to( + *, + gpu_index: int, + numa_options: NumaOptions, +) -> set[int]: + logical_cpu_indices = _get_logical_cpus_to_bind_to( + gpu_index=gpu_index, numa_options=numa_options + ) + _raise_if_binding_invalid(logical_cpu_indices=logical_cpu_indices) + + return logical_cpu_indices + + +def _raise_if_binding_invalid(*, logical_cpu_indices: set[int]) -> None: + # NOTE: numactl CLI is only actually necessary for the str entrypoint path, + # but for simplicity we will just check it no matter what. + if shutil.which("numactl") is None: + raise RuntimeError("numactl CLI is required for NUMA binding") + + if not logical_cpu_indices: + raise RuntimeError("Must bind to a non-empty set of CPU indices") + + +def _bind_all_threads_in_current_process_to_logical_cpus( + *, logical_cpu_indices: set[int] +) -> None: + # Save the original affinity of the main thread before changing it + # pyrefly: ignore [missing-attribute] + original_main_thread_affinity = os.sched_getaffinity(0) # type: ignore[attr-defined] + + # 0 represents the current thread. + # This is outside the try/except because the main thread should always bind successfully. + # pyrefly: ignore [missing-attribute] + os.sched_setaffinity(0, logical_cpu_indices) # type: ignore[attr-defined] + + for tid_str in os.listdir("/proc/self/task"): + try: + tid = int(tid_str) + # pyrefly: ignore [missing-attribute] + tid_affinity = os.sched_getaffinity(tid) # type: ignore[attr-defined] + + # Defensive check to ensure we do not overwrite affinity on any threads + # that have already had their affinity set elsewhere. + if tid_affinity == original_main_thread_affinity: + # pyrefly: ignore [missing-attribute] + os.sched_setaffinity(tid, logical_cpu_indices) # type: ignore[attr-defined] + except Exception: + # Thread may have exited or otherwise become invalid + pass + + +def _get_logical_cpus_to_bind_to( + *, + gpu_index: int, + numa_options: NumaOptions, +) -> set[int]: + """ + Args: + gpu_index: The index of the GPU that will be used by the subprocess. + Example: 0 + numa_options: See NumaOptions for details. + + Returns: + Set of logical CPU indices to bind to. + """ + if numa_options.affinity_mode == AffinityMode.NODE: + logical_cpus = _node_get_logical_cpus_to_bind_to(gpu_index=gpu_index) + elif numa_options.affinity_mode == AffinityMode.SOCKET: + logical_cpus = _socket_get_logical_cpus_to_bind_to(gpu_index=gpu_index) + elif numa_options.affinity_mode == AffinityMode.EXCLUSIVE: + logical_cpus = _exclusive_get_logical_cpus_to_bind_to(gpu_index=gpu_index) + elif numa_options.affinity_mode == AffinityMode.CORE_COMPLEX: + logical_cpus = _core_complex_get_logical_cpus_to_bind_to(gpu_index=gpu_index) + else: + raise ValueError(f"Affinity mode {numa_options.affinity_mode} not supported.") + + return logical_cpus + + +def _node_get_logical_cpus_to_bind_to(*, gpu_index: int) -> set[int]: + """ + Core logic of 'node' numa strategy. + """ + numa_node_index = _get_numa_node_index_for_gpu_index(gpu_index=gpu_index) + + return _get_allowed_logical_cpu_indices_for_numa_node( + numa_node_index=numa_node_index + ) + + +def _socket_get_logical_cpus_to_bind_to(*, gpu_index: int) -> set[int]: + """ + Core logic of 'socket' numa strategy. + """ + numa_node_index_of_gpu = _get_numa_node_index_for_gpu_index(gpu_index=gpu_index) + socket_index = _get_socket_index_for_numa_node( + numa_node_index=numa_node_index_of_gpu + ) + numa_node_indices = _get_numa_node_indices_for_socket_index( + socket_index=socket_index + ) + + logical_cpus = set() + for numa_node_index in numa_node_indices: + logical_cpus.update( + _get_allowed_logical_cpu_indices_for_numa_node( + numa_node_index=numa_node_index + ) + ) + + return logical_cpus + + +def _exclusive_get_logical_cpus_to_bind_to(*, gpu_index: int) -> set[int]: + """ + Core logic of 'exclusive' numa strategy. + """ + numa_node_index = _get_numa_node_index_for_gpu_index(gpu_index=gpu_index) + + gpu_indices = _get_gpu_indices_for_numa_node(numa_node_index=numa_node_index) + gpu_indices = sorted(gpu_indices) + original_gpu_relative_index = gpu_indices.index(gpu_index) + + allowed_logical_cpu_indices = _get_allowed_logical_cpu_indices_for_numa_node( + numa_node_index=numa_node_index + ) + + # Arbitrarily use the min logical cpu index on the physical core to + # represent the physical core. + physical_core_to_allowed_logical_cpu_indices = _group_by( + allowed_logical_cpu_indices, + lambda logical_cpu_index: min( + _get_logical_cpu_indices_sharing_same_physical_core_as( + logical_cpu_index=logical_cpu_index + ) + ), + ) + # Sort the dict for consistency (dicts maintain order in Python) + physical_core_to_allowed_logical_cpu_indices = dict( + sorted(physical_core_to_allowed_logical_cpu_indices.items()) + ) + + num_physical_cores_per_gpu = len( + physical_core_to_allowed_logical_cpu_indices + ) // len(gpu_indices) + # Often, the number of physical cores will not be perfectly divisible by the number + # of GPUs. In those cases, give the lowest GPU indices an extra core + num_gpus_to_give_one_extra_physical_core = len( + physical_core_to_allowed_logical_cpu_indices + ) % len(gpu_indices) + + if num_physical_cores_per_gpu < 1: + raise RuntimeError( + f"There are only {len(physical_core_to_allowed_logical_cpu_indices)} physical cores on {numa_node_index=}," + + f" but there are {len(gpu_indices)} GPUs associated with this NUMA node." + ) + + # Compute slice indices for this GPU + start = original_gpu_relative_index * num_physical_cores_per_gpu + min( + original_gpu_relative_index, num_gpus_to_give_one_extra_physical_core + ) + end = ( + start + + num_physical_cores_per_gpu + + ( + 1 + if original_gpu_relative_index < num_gpus_to_give_one_extra_physical_core + else 0 + ) + ) + + # Slice and flatten the logical CPUs from the selected physical cores + logical_cpu_indices_for_original_gpu = { + logical_cpu_index + for logical_cpu_indices in list( + physical_core_to_allowed_logical_cpu_indices.values() + )[start:end] + for logical_cpu_index in logical_cpu_indices + } + + return logical_cpu_indices_for_original_gpu + + +def _core_complex_get_logical_cpus_to_bind_to(*, gpu_index: int) -> set[int]: + """ + Core logic of 'core-complex' numa strategy. + + Each GPU is assigned a full core complex (group of cores sharing L3 cache) + within its affined NUMA node. + """ + numa_node_index = _get_numa_node_index_for_gpu_index(gpu_index=gpu_index) + + gpu_indices = _get_gpu_indices_for_numa_node(numa_node_index=numa_node_index) + gpu_indices = sorted(gpu_indices) + original_gpu_relative_index = gpu_indices.index(gpu_index) + + allowed_logical_cpu_indices = _get_allowed_logical_cpu_indices_for_numa_node( + numa_node_index=numa_node_index + ) + + # Arbitrarily use the min logical cpu index on the max level cache + # to represent the max level cache. + max_level_cache_to_allowed_logical_cpu_indices = _group_by( + allowed_logical_cpu_indices, + lambda logical_cpu_index: min( + _get_logical_cpus_sharing_same_max_level_cache_as( + logical_cpu_index=logical_cpu_index + ) + ), + ) + + max_level_cache_to_allowed_logical_cpu_indices = dict( + sorted( + max_level_cache_to_allowed_logical_cpu_indices.items(), + # First, prioritize caches with more available cpus + # Second, prioritize lower index cpus (just for clarity/consistency) + key=lambda item: (-len(item[1]), item[0]), + ) + ) + + cache_index_for_original_gpu = original_gpu_relative_index % len( + max_level_cache_to_allowed_logical_cpu_indices + ) + logical_cpu_indices_for_original_gpu = list( + max_level_cache_to_allowed_logical_cpu_indices.values() + )[cache_index_for_original_gpu] + + return logical_cpu_indices_for_original_gpu + + +K = TypeVar("K") +V = TypeVar("V") + + +def _group_by(values: Iterable[V], get_key: Callable[[V], K]) -> dict[K, set[V]]: + """ + Groups elements with same key into sets. + """ + key_to_values: defaultdict[K, set[V]] = defaultdict(set) + for value in values: + key = get_key(value) + key_to_values[key].add(value) + return key_to_values + + +def _get_logical_cpu_indices_sharing_same_physical_core_as( + *, logical_cpu_index: int +) -> set[int]: + thread_siblings_list_absolute_path = ( + f"/sys/devices/system/cpu/cpu{logical_cpu_index}/topology/thread_siblings_list" + ) + with open(thread_siblings_list_absolute_path) as f: + return _get_set_of_int_from_ranges_str(f.read()) + + +def _get_logical_cpus_sharing_same_max_level_cache_as( + *, logical_cpu_index: int +) -> set[int]: + cpu_cache_dir_absolute_path = ( + f"/sys/devices/system/cpu/cpu{logical_cpu_index}/cache" + ) + + max_level = -1 + logical_cpus_sharing_max_level_cache = set() + for entry in os.listdir(cpu_cache_dir_absolute_path): + if not entry.startswith("index") or not entry[5:].isdecimal(): + continue + cache_index_absolute_path = os.path.join(cpu_cache_dir_absolute_path, entry) + + # Filter out other cache types like Instruction + type_absolute_path = os.path.join(cache_index_absolute_path, "type") + with open(type_absolute_path) as type_file: + if type_file.read().strip() not in {"Unified", "Data"}: + continue + + level_absolute_path = os.path.join(cache_index_absolute_path, "level") + with open(level_absolute_path) as level_file: + level = int(level_file.read()) + if level <= max_level: + continue + + max_level = level + shared_cpu_list_absolute_path = os.path.join( + cache_index_absolute_path, "shared_cpu_list" + ) + with open(shared_cpu_list_absolute_path) as share_cpu_list_file: + logical_cpus_sharing_max_level_cache = _get_set_of_int_from_ranges_str( + share_cpu_list_file.read() + ) + + return logical_cpus_sharing_max_level_cache + + +def _get_allowed_logical_cpu_indices_for_numa_node(*, numa_node_index: int) -> set[int]: + all_cpu_indices = _get_cpu_indices_for_numa_node_MAYBE_NOT_ALLOWED( + numa_node_index=numa_node_index + ) + allowed_cpu_indices = _get_allowed_cpu_indices_for_current_thread() + return all_cpu_indices & allowed_cpu_indices + + +def _get_cpu_indices_for_numa_node_MAYBE_NOT_ALLOWED( + *, numa_node_index: int +) -> set[int]: + """ + Returns: + Indices of all CPUs associated with numa_node_index. However, the list + is not filtered based on whether the thread is allowed to use them. + """ + cpulist_absolute_path = f"/sys/devices/system/node/node{numa_node_index}/cpulist" + try: + with open(cpulist_absolute_path) as f: + cpu_range_str = f.read() + except FileNotFoundError as e: + raise RuntimeError( + f"Could not determine CPUs corresponding to {numa_node_index=}." + ) from e + return _get_set_of_int_from_ranges_str(cpu_range_str) + + +def _get_gpu_count() -> int: + return torch.cuda.device_count() + + +def _get_numa_node_index_for_gpu_index(*, gpu_index: int) -> int: + device_properties = torch.cuda.get_device_properties(gpu_index) + + domain = device_properties.pci_domain_id # type: ignore[attr-defined] + bus = device_properties.pci_bus_id # type: ignore[attr-defined] + device = device_properties.pci_device_id # type: ignore[attr-defined] + + # Format to sysfs PCI address: "0000:dc:00.0" + pci_addr = f"{domain:04x}:{bus:02x}:{device:02x}.0" + + pci_numa_node_absolute_path = f"/sys/bus/pci/devices/{pci_addr}/numa_node" + with open(pci_numa_node_absolute_path) as f: + # In systems with only one NUMA node, this will + # often be saved as -1. In those cases, there is obviously + # at least one numa node, 0, so we use that. + return max(int(f.read().strip()), 0) + + +def _get_gpu_indices_for_numa_node(*, numa_node_index: int) -> set[int]: + return { + gpu_index + for gpu_index in range(_get_gpu_count()) + if _get_numa_node_index_for_gpu_index(gpu_index=gpu_index) == numa_node_index + } + + +def _get_socket_index_for_numa_node(*, numa_node_index: int) -> int: + arbitrary_cpu_index = _get_arbitrary_allowed_cpu_index_for_numa_node( + numa_node_index=numa_node_index + ) + + return _get_socket_index_for_cpu(cpu_index=arbitrary_cpu_index) + + +def _get_socket_index_for_cpu(*, cpu_index: int) -> int: + package_id_absolute_path = ( + f"/sys/devices/system/cpu/cpu{cpu_index}/topology/physical_package_id" + ) + try: + with open(package_id_absolute_path) as f: + return int(f.read().strip()) + except FileNotFoundError as e: + raise RuntimeError(f"Could not determine socket for {cpu_index=}") from e + + +def _get_arbitrary_allowed_cpu_index_for_numa_node(*, numa_node_index: int) -> int: + return min( + _get_allowed_logical_cpu_indices_for_numa_node(numa_node_index=numa_node_index) + ) + + +def _get_set_of_int_from_ranges_str(ranges_str: str) -> set[int]: + """ + Util for parsing a string of int ranges, as in a sysfs file. + + Args: + ranges_str: E.g., "0-2,4,6-7" + + Returns: + E.g., {0, 1, 2, 4, 6, 7} + """ + ints: set[int] = set() + for range_str in ranges_str.split(","): + range_str = range_str.strip() + if not range_str: + continue + if "-" in range_str: + start_str, end_str = range_str.split("-") + start, end = int(start_str), int(end_str) + ints.update(range(start, end + 1)) + else: + ints.add(int(range_str)) + return ints + + +def _get_ranges_str_from_ints(ints: Iterable[int]) -> str: + """ + Convert a set of integers to a compact string with ranges. + + Args: + ints: E.g., {0, 1, 2, 4, 6, 7} + + Returns: + E.g., "0-2,4,6-7" + """ + if not ints: + return "" + + sorted_ints = sorted(ints) + ranges = [] + start = prev = sorted_ints[0] + + for num in sorted_ints[1:]: + if num == prev + 1: + prev = num + else: + if start == prev: + ranges.append(f"{start}") + else: + ranges.append(f"{start}-{prev}") + start = prev = num + + # Append the last range + if start == prev: + ranges.append(f"{start}") + else: + ranges.append(f"{start}-{prev}") + + return ",".join(ranges) + + +def _get_systemwide_numa_node_indices() -> set[int]: + with open("/sys/devices/system/node/possible") as f: + possible_nodes_str = f.read() + + return _get_set_of_int_from_ranges_str(possible_nodes_str) + + +def _get_numa_node_indices_for_socket_index(*, socket_index: int) -> set[int]: + systemwide_numa_node_indices = _get_systemwide_numa_node_indices() + + matching_numa_node_indices = set() + for numa_node_index in systemwide_numa_node_indices: + arbitrary_cpu_index = _get_arbitrary_allowed_cpu_index_for_numa_node( + numa_node_index=numa_node_index + ) + if socket_index == _get_socket_index_for_cpu(cpu_index=arbitrary_cpu_index): + matching_numa_node_indices.add(numa_node_index) + + return matching_numa_node_indices + + +def _get_allowed_cpu_indices_for_current_thread() -> set[int]: + # 0 denotes current thread + # pyrefly: ignore [missing-attribute] + return os.sched_getaffinity(0) # type:ignore[attr-defined] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..346155f88b44255b02f55c5cd88c5e2b8d0a3daa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/__init__.py @@ -0,0 +1,361 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + + +__all__ = [ + # Modules + "errors", + "ops", + # Public functions + "export", + "is_in_onnx_export", + # Base error + "OnnxExporterError", + "ONNXProgram", + "ExportableModule", + "InputObserver", +] + +from typing import Any, TYPE_CHECKING + +import torch +from torch._C import _onnx as _C_onnx +from torch._C._onnx import ( # Deprecated members that are excluded from __all__ + OperatorExportTypes as OperatorExportTypes, + TensorProtoDataType as TensorProtoDataType, + TrainingMode as TrainingMode, +) + +from . import errors, ops +from ._internal.exporter._exportable_module import ExportableModule +from ._internal.exporter._input_observer import InputObserver +from ._internal.exporter._onnx_program import ONNXProgram +from ._internal.torchscript_exporter import ( # Deprecated members that are excluded from __all__ + symbolic_helper, + symbolic_opset10, + symbolic_opset9, + utils, +) +from ._internal.torchscript_exporter._type_utils import ( + JitScalarType, # Deprecated members that are excluded from __all__ +) +from ._internal.torchscript_exporter.utils import ( # Deprecated members that are excluded from __all__ + register_custom_op_symbolic, + select_model_mode_for_export, # pyrefly: ignore # deprecated + unregister_custom_op_symbolic, +) +from .errors import OnnxExporterError + + +if TYPE_CHECKING: + import os + from collections.abc import Callable, Collection, Mapping, Sequence + +# Set namespace for exposed private names +ONNXProgram.__module__ = "torch.onnx" +ExportableModule.__module__ = "torch.onnx" +OnnxExporterError.__module__ = "torch.onnx" +InputObserver.__module__ = "torch.onnx" + +# TODO(justinchuby): Remove these two properties +producer_name = "pytorch" +producer_version = _C_onnx.PRODUCER_VERSION + + +def export( + model: torch.nn.Module + | torch.export.ExportedProgram + | torch.jit.ScriptModule + | torch.jit.ScriptFunction, + args: tuple[Any, ...] = (), + f: str | os.PathLike | None = None, + *, + kwargs: dict[str, Any] | None = None, + verbose: bool | None = None, + input_names: Sequence[str] | None = None, + output_names: Sequence[str] | None = None, + opset_version: int | None = None, + dynamo: bool = True, + # Dynamo only options + external_data: bool = True, + dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None = None, + custom_translation_table: dict[Callable, Callable] | None = None, + report: bool = False, + optimize: bool = True, + verify: bool = False, + profile: bool = False, + dump_exported_program: bool = False, + artifacts_dir: str | os.PathLike = ".", + # BC options + export_params: bool = True, + keep_initializers_as_inputs: bool = False, + dynamic_axes: Mapping[str, Mapping[int, str]] + | Mapping[str, Sequence[int]] + | None = None, + # Deprecated options + training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL, + operator_export_type: _C_onnx.OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX, + do_constant_folding: bool = True, + custom_opsets: Mapping[str, int] | None = None, + export_modules_as_functions: bool | Collection[type[torch.nn.Module]] = False, + autograd_inlining: bool = True, +) -> ONNXProgram | None: + r"""Exports a model into ONNX format. + + Setting ``dynamo=True`` enables the new ONNX export logic + which is based on :class:`torch.export.ExportedProgram` and a more modern + set of translation logic. This is the recommended and default way to export models + to ONNX. + + When ``dynamo=True``: + + The exporter tries the following strategies to get an ExportedProgram for conversion to ONNX. + + #. If the model is already an ExportedProgram, it will be used as-is. + #. Use :func:`torch.export.export` and set ``strict=False``. + #. Use :func:`torch.export.export` and set ``strict=True``. + + Args: + model: The model to be exported. + args: Example positional inputs. Any non-Tensor arguments will be hard-coded into the + exported model; any Tensor arguments will become inputs of the exported model, + in the order they occur in the tuple. + f: Path to the output ONNX model file. E.g. "model.onnx". This argument is kept for + backward compatibility. It is recommended to leave unspecified (None) + and use the returned :class:`torch.onnx.ONNXProgram` to serialize the model + to a file instead. + kwargs: Optional example keyword inputs. + verbose: Whether to enable verbose logging. + input_names: names to assign to the input nodes of the graph, in order. + output_names: names to assign to the output nodes of the graph, in order. + opset_version: The version of the + `default (ai.onnx) opset `_ + to target. You should set ``opset_version`` according to the supported opset versions + of the runtime backend or compiler you want to run the exported model with. + Leave as default (``None``) to use the recommended version, or refer to + the ONNX operators documentation for more information. + dynamo: Whether to export the model with ``torch.export`` ExportedProgram instead of TorchScript. + external_data: Whether to save the model weights as an external data file. + This is required for models with large weights that exceed the ONNX file size limit (2GB). + When False, the weights are saved in the ONNX file with the model architecture. + dynamic_shapes: A dictionary or a tuple of dynamic shapes for the model inputs. Refer to + :func:`torch.export.export` for more details. This is only used (and preferred) when dynamo is True. + Note that dynamic_shapes is designed to be used when the model is exported with dynamo=True, while + dynamic_axes is used when dynamo=False. + custom_translation_table: A dictionary of custom decompositions for operators in the model. + The dictionary should have the callable target in the fx Node as the key (e.g. ``torch.ops.aten.stft.default``), + and the value should be a function that builds that graph using ONNX Script. This option + is only valid when dynamo is True. + report: Whether to generate a markdown report for the export process. This option + is only valid when dynamo is True. + optimize: Whether to optimize the exported model. This option + is only valid when dynamo is True. Default is True. + verify: Whether to verify the exported model using ONNX Runtime. This option + is only valid when dynamo is True. + profile: Whether to profile the export process. This option + is only valid when dynamo is True. + dump_exported_program: Whether to dump the :class:`torch.export.ExportedProgram` to a file. + This is useful for debugging the exporter. This option is only valid when dynamo is True. + artifacts_dir: The directory to save the debugging artifacts like the report and the serialized + exported program. This option is only valid when dynamo is True. + export_params: **When ``f`` is specified**: If false, parameters (weights) will not be exported. + + You can also leave it unspecified and use the returned :class:`torch.onnx.ONNXProgram` + to control how initializers are treated when serializing the model. + keep_initializers_as_inputs: **When ``f`` is specified**: If True, all the + initializers (typically corresponding to model weights) in the + exported graph will also be added as inputs to the graph. If False, + then initializers are not added as inputs to the graph, and only + the user inputs are added as inputs. + + Set this to True if you intend to supply model weights at runtime. + Set it to False if the weights are static to allow for better optimizations + (e.g. constant folding) by backends/runtimes. + + You can also leave it unspecified and use the returned :class:`torch.onnx.ONNXProgram` + to control how initializers are treated when serializing the model. + dynamic_axes: + Deprecated: Prefer specifying ``dynamic_shapes`` when ``dynamo=True``. + + By default the exported model will have the shapes of all input and output tensors + set to exactly match those given in ``args``. To specify axes of tensors as + dynamic (i.e. known only at run-time), set ``dynamic_axes`` to a dict with schema: + + * KEY (str): an input or output name. Each name must also be provided in ``input_names`` or + ``output_names``. + * VALUE (dict or list): If a dict, keys are axis indices and values are axis names. If a + list, each element is an axis index. + + For example:: + + class SumModule(torch.nn.Module): + def forward(self, x): + return torch.sum(x, dim=1) + + + torch.onnx.export( + SumModule(), + (torch.ones(2, 2),), + "onnx.pb", + input_names=["x"], + output_names=["sum"], + ) + + Produces:: + + input { + name: "x" + ... + shape { + dim { + dim_value: 2 # axis 0 + } + dim { + dim_value: 2 # axis 1 + ... + output { + name: "sum" + ... + shape { + dim { + dim_value: 2 # axis 0 + ... + + While:: + + torch.onnx.export( + SumModule(), + (torch.ones(2, 2),), + "onnx.pb", + input_names=["x"], + output_names=["sum"], + dynamic_axes={ + # dict value: manually named axes + "x": {0: "my_custom_axis_name"}, + # list value: automatic names + "sum": [0], + }, + ) + + Produces:: + + input { + name: "x" + ... + shape { + dim { + dim_param: "my_custom_axis_name" # axis 0 + } + dim { + dim_value: 2 # axis 1 + ... + output { + name: "sum" + ... + shape { + dim { + dim_param: "sum_dynamic_axes_1" # axis 0 + ... + + training: Deprecated option. Instead, set the training mode of the model before exporting. + operator_export_type: Deprecated option. Only ONNX is supported. + do_constant_folding: Deprecated option. + custom_opsets: Deprecated option. + export_modules_as_functions: Deprecated option. + autograd_inlining: Deprecated option. + + Returns: + :class:`torch.onnx.ONNXProgram` if dynamo is True, otherwise None. + + .. versionchanged:: 2.6 + ``training`` is now deprecated. Instead, set the training mode of the model before exporting. + ``operator_export_type`` is now deprecated. Only ONNX is supported. + ``do_constant_folding`` is now deprecated. It is always enabled. + ``export_modules_as_functions`` is now deprecated. + ``autograd_inlining`` is now deprecated. + .. versionchanged:: 2.7 + ``optimize`` is now True by default. + .. versionchanged:: 2.9 + ``dynamo`` is now True by default. + .. versionchanged:: 2.11 + ``fallback`` option has been removed. + """ + if dynamo is True or isinstance( + model, (torch.export.ExportedProgram, ExportableModule) + ): + from torch.onnx._internal.exporter import _compat + + if isinstance(args, torch.Tensor): + args = (args,) + + return _compat.export_compat( + model, + args, + f, + kwargs=kwargs, + export_params=export_params, + verbose=verbose, + input_names=input_names, + output_names=output_names, + opset_version=opset_version, + custom_translation_table=custom_translation_table, + dynamic_axes=dynamic_axes, + keep_initializers_as_inputs=keep_initializers_as_inputs, + external_data=external_data, + dynamic_shapes=dynamic_shapes, + report=report, + optimize=optimize, + verify=verify, + profile=profile, + dump_exported_program=dump_exported_program, + artifacts_dir=artifacts_dir, + ) + else: + import warnings + + from ._internal.torchscript_exporter.utils import export + + warnings.warn( + "You are using the legacy TorchScript-based ONNX export. Starting in PyTorch 2.9, " + "the new torch.export-based ONNX exporter has become the default. " + "Learn more about the new export logic: https://docs.pytorch.org/docs/stable/onnx_export.html. " + "For exporting control flow: " + "https://pytorch.org/tutorials/beginner/onnx/export_control_flow_model_to_onnx_tutorial.html", + category=DeprecationWarning, + stacklevel=2, + ) + + if dynamic_shapes: + raise ValueError( + "The exporter only supports dynamic shapes " + "through parameter dynamic_axes when dynamo=False." + ) + + export( + model, + args, + f, # type: ignore[arg-type] + kwargs=kwargs, + export_params=export_params, + verbose=verbose is True, + input_names=input_names, + output_names=output_names, + opset_version=opset_version, + dynamic_axes=dynamic_axes, + keep_initializers_as_inputs=keep_initializers_as_inputs, + training=training, + operator_export_type=operator_export_type, + do_constant_folding=do_constant_folding, + custom_opsets=custom_opsets, + export_modules_as_functions=export_modules_as_functions, + autograd_inlining=autograd_inlining, + ) + return None + + +def is_in_onnx_export() -> bool: + """Returns whether it is in the middle of ONNX export.""" + from torch.onnx._internal.exporter import _flags + from torch.onnx._internal.torchscript_exporter._globals import GLOBALS + + return GLOBALS.in_onnx_export or _flags._is_onnx_exporting diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_constants.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..87ff04da8cd1e6528bb3f902f56a826f17fc5f87 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_constants.py @@ -0,0 +1,24 @@ +"""Constant values used in ONNX.""" + +ONNX_ARCHIVE_MODEL_PROTO_NAME = "__MODEL_PROTO" + +ONNX_BASE_OPSET = 9 +ONNX_MIN_OPSET = 7 +ONNX_MAX_OPSET = 23 +ONNX_TORCHSCRIPT_EXPORTER_MAX_OPSET = 20 +ONNX_DEFAULT_OPSET = 20 +ONNX_CONSTANT_FOLDING_MIN_OPSET = 9 + +PYTORCH_GITHUB_ISSUES_URL = "https://github.com/pytorch/pytorch/issues" + +INT64_MAX = 9223372036854775807 +INT32_MAX = 2147483647 +INT16_MAX = 32767 +INT8_MAX = 127 +UINT8_MAX = 255 + +INT64_MIN = -9223372036854775808 +INT32_MIN = -2147483648 +INT16_MIN = -32768 +INT8_MIN = -128 +UINT8_MIN = 0 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_flags.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_flags.py new file mode 100644 index 0000000000000000000000000000000000000000..bcc9f633e7584edf65273796f8baf1e70484e827 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_flags.py @@ -0,0 +1,55 @@ +"""Internal feature flags for torch.onnx. + +NOTE: These flags are experimental only. Any flag here can be removed at any +time without notice. +""" + +import logging +import os + + +logger = logging.getLogger(__name__) + + +def _load_boolean_flag( + name: str, + *, + this_will: str, + deprecated: bool = False, + default: bool = False, +) -> bool: + """Load a boolean flag from environment variable. + + Args: + name: The name of the environment variable. + this_will: A string that describes what this flag will do. + deprecated: Whether this flag is deprecated. + default: The default value if envvar not defined. + """ + undefined = os.getenv(name) is None + state = os.getenv(name) == "1" + if state: + if deprecated: + logger.error( + "Experimental flag %s is deprecated. Please remove it from your environment.", + name, + ) + else: + logger.warning( + "Experimental flag %s is enabled. This will %s.", name, this_will + ) + if undefined: + state = default + return state + + +ENABLE_DRAFT_EXPORT: bool = _load_boolean_flag( + "TORCH_ONNX_ENABLE_DRAFT_EXPORT", + this_will="enable torch.export.draft_export as a strategy for capturing models", + default=False, +) +PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS: bool = _load_boolean_flag( + "TORCH_ONNX_PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS", + this_will="set prefer_deferred_runtime_asserts_over_guards when calling torch.export", + default=True, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/_lazy_import.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/_lazy_import.py new file mode 100644 index 0000000000000000000000000000000000000000..3e538b44e1dc60a29b8dc322d054c388a379ffba --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/_lazy_import.py @@ -0,0 +1,39 @@ +"""Utility to lazily import modules.""" + +from __future__ import annotations + +import importlib +from typing import Any, TYPE_CHECKING + + +class _LazyModule: + """Lazily import a module.""" + + def __init__(self, module_name: str) -> None: + self._name = module_name + self._module: Any = None + + def __repr__(self) -> str: + return f"" + + def __getattr__(self, attr: str) -> object: + if self._module is None: + self._module = importlib.import_module(".", self._name) + return getattr(self._module, attr) + + +# Import the following modules during type checking to enable code intelligence features, +# such as auto-completion in tools like pylance, even when these modules are not explicitly +# imported in user code. +# NOTE: Add additional used imports here. +if TYPE_CHECKING: + import onnx + import onnx_ir # type: ignore[import-untyped, import-not-found] + import onnxscript + import onnxscript._framework_apis.torch_2_11 as onnxscript_apis + +else: + onnx = _LazyModule("onnx") + onnx_ir = _LazyModule("onnx_ir") + onnxscript = _LazyModule("onnxscript") + onnxscript_apis = _LazyModule("onnxscript._framework_apis.torch_2_11") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_analysis.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_analysis.py new file mode 100644 index 0000000000000000000000000000000000000000..53860413526eed8fb858eb70606169bd8e14aee3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_analysis.py @@ -0,0 +1,271 @@ +"""Compatibility analyzer for PyTorch models.""" + +# mypy: allow-untyped-defs +# flake8: noqa: B950 We do not need flake8 as it complains line length +from __future__ import annotations + +import dataclasses +import operator +import textwrap +import traceback +from collections import defaultdict +from typing import TYPE_CHECKING + +import torch +import torch._export.serde.schema +from torch.export import graph_signature +from torch.onnx._internal.exporter import _dispatching, _registration + + +if TYPE_CHECKING: + import torch.fx + + +@dataclasses.dataclass +class ModelInfo: + """Information about the model.""" + + parameter_count: defaultdict[torch.dtype, int] = dataclasses.field( + default_factory=lambda: defaultdict(int) + ) + buffer_count: defaultdict[torch.dtype, int] = dataclasses.field( + default_factory=lambda: defaultdict(int) + ) + fx_node_count: int = 0 + fx_node_op_count: defaultdict[str, int] = dataclasses.field( + default_factory=lambda: defaultdict(int) + ) + fx_node_target_count: defaultdict[str, int] = dataclasses.field( + default_factory=lambda: defaultdict(int) + ) + dispatch_failures: list[tuple[torch.fx.Node, str]] = dataclasses.field( + default_factory=list + ) + inputs: dict[str, torch._export.serde.schema.TensorMeta] = dataclasses.field( + default_factory=dict + ) + outputs: dict[str, torch._export.serde.schema.TensorMeta] = dataclasses.field( + default_factory=dict + ) + + +def _count_weights( + exported_program: torch.export.ExportedProgram, +) -> tuple[defaultdict[torch.dtype, int], defaultdict[torch.dtype, int]]: + """Count the size of the parameters in the exported program.""" + + parameter_count: defaultdict[torch.dtype, int] = defaultdict(int) + buffer_count: defaultdict[torch.dtype, int] = defaultdict(int) + for parameter in exported_program.parameters(): + dtype = parameter.dtype + parameter_count[dtype] += parameter.numel() + + for buffer in exported_program.buffers(): + dtype = buffer.dtype + buffer_count[dtype] += buffer.numel() + + return parameter_count, buffer_count + + +def _format_model_info(model_info: ModelInfo) -> str: + """Format the information about the model.""" + lines = [ + textwrap.dedent( + f"""\ + PyTorch ONNX Conversion Analysis + + ## Model Information + + The model has {sum(model_info.parameter_count.values())} parameters and {sum(model_info.buffer_count.values())} buffers (non-trainable parameters). + Number of parameters per dtype: + ```python + {model_info.parameter_count} + ``` + Number of buffers per dtype: + ```python + {model_info.buffer_count} + ``` + """ + ), + "Inputs:", + *[f"- `{name}`: `{meta}`" for name, meta in model_info.inputs.items()], + "", + "Outputs:", + *[f"- `{name}`: `{meta}`" for name, meta in model_info.outputs.items()], + "", + f"The FX graph has {model_info.fx_node_count} nodes in total. Number of FX nodes per op:", + ] + for op, count in model_info.fx_node_op_count.items(): + lines.append(f"- `{op}`: {count}") + lines.append("\n") + lines.append("Of the call_function nodes, the counts of operators used are:\n") + sorted_targets = sorted( + model_info.fx_node_target_count.items(), + key=operator.itemgetter(1), + reverse=True, + ) + for target, count in sorted_targets: + lines.append(f"- `{target}`: {count}") + + lines.append("") + lines.append("## ONNX Conversion Information") + lines.append("") + + if model_info.dispatch_failures: + lines.append( + "The model contains operators the dispatcher could not find registered ONNX decompositions for. " + "This may be due to missing implementations, decompositions not registered " + "correctly, or a bug in the dispatcher." + ) + lines.append("") + lines.append("Errors grouped by operator:\n") + + target_to_nodes = defaultdict(list) + for node, _ in model_info.dispatch_failures: + target_to_nodes[str(node.target)].append(node) + + target_to_messages = {} + for node, message in model_info.dispatch_failures: + if str(node.target) not in target_to_messages: + target_to_messages[str(node.target)] = message + + for target, nodes in sorted( + target_to_nodes.items(), key=operator.itemgetter(0), reverse=True + ): + message = textwrap.indent( + f"{target_to_messages[target]}. Example node: `{nodes[0].format_node()}`. All nodes: `{nodes}`", + " ", + ) + lines.append(f"- `{target}`: {message}") + else: + lines.append("All operators in the model have registered ONNX decompositions.") + + return "\n".join(lines) + + +def _get_io_specs(exported_program: torch.export.ExportedProgram) -> tuple[dict, dict]: + """Get the input and output specs of the exported program.""" + + nodes: dict[str, torch.fx.Node] = { + node.name: node for node in exported_program.graph.nodes + } + user_inputs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == graph_signature.InputKind.USER_INPUT + ] + user_outputs = [ + spec + for spec in exported_program.graph_signature.output_specs + if spec.kind == graph_signature.OutputKind.USER_OUTPUT + ] + inputs: dict[str, torch._export.serde.schema.TensorMeta | str] = {} + outputs: dict[str, torch._export.serde.schema.TensorMeta | str] = {} + for spec in user_inputs: + inputs = _log_spec_into_io_specs(spec, nodes, inputs) + for spec in user_outputs: + outputs = _log_spec_into_io_specs(spec, nodes, outputs) + return inputs, outputs + + +def _log_spec_into_io_specs( + spec: graph_signature.InputSpec, + nodes: dict[str, torch.fx.Node], + inputs_or_outputs: dict[str, torch._export.serde.schema.TensorMeta | str], +) -> dict[str, torch._export.serde.schema.TensorMeta | str]: + # If dynamic is set to a constant input, it becomes a + # symbolic argument, which is not a tensor. + if isinstance(spec.arg, graph_signature.ConstantArgument): + # Constant input does not have tensor_meta. + return inputs_or_outputs + # Symbolic arguments are not tensors, so it does not have tensor_meta, + # but we need to provide a string representation for them to inform users. + name = spec.arg.name + if isinstance( + spec.arg, + ( + graph_signature.SymIntArgument, + graph_signature.SymFloatArgument, + graph_signature.SymBoolArgument, + ), + ): + argument_to_str: dict[type[graph_signature.ArgumentSpec], str] = { + graph_signature.SymIntArgument: "SymInt", + graph_signature.SymFloatArgument: "SymFloat", + graph_signature.SymBoolArgument: "SymBool", + } + inputs_or_outputs[name] = argument_to_str[type(spec.arg)] + return inputs_or_outputs + # FIXME: tensor_meta is None sometimes when the exported program still knows the shape/type + inputs_or_outputs[name] = nodes[name].meta["tensor_meta"] + return inputs_or_outputs + + +def _count_fx_targets( + exported_program: torch.export.ExportedProgram, +) -> defaultdict[str, int]: + """Count the number of targets for each node in the exported program.""" + fx_node_target_count: defaultdict[str, int] = defaultdict(int) + for node in exported_program.graph.nodes: + if node.op == "call_function": + fx_node_target_count[str(node.target)] += 1 + return fx_node_target_count + + +def analyze( + exported_program: torch.export.ExportedProgram, + registry: _registration.ONNXRegistry | None = None, + file=None, +) -> None: + """Analyze the compatibility of the exported program.""" + # Get basic information about the model + model_info = ModelInfo() + model_info.parameter_count, model_info.buffer_count = _count_weights( + exported_program + ) + model_info.fx_node_count = len(exported_program.graph.nodes) + model_info.fx_node_target_count = _count_fx_targets(exported_program) + inputs, outputs = _get_io_specs(exported_program) + model_info.inputs = inputs + model_info.outputs = outputs + + if registry is None: + registry = _registration.ONNXRegistry.from_torchlib() + + # Try to find ops for every node in the graph + for node in exported_program.graph.nodes: + model_info.fx_node_op_count[node.op] += 1 + if node.op == "call_function": + try: + onnx_function, message = _dispatching.dispatch(node, registry) + except Exception as e: + message = "Critical Error in dispatcher:\n" + formatted_exception = "\n".join( + traceback.format_exception(type(e), e, e.__traceback__) + ) + message += f"```pytb\n{formatted_exception}\n```\n" + onnx_function = None + if onnx_function is None: + model_info.dispatch_failures.append((node, message)) + + # Print the results + report = _format_model_info(model_info) + print(report, file=file, flush=True) + + +def compare_ops( + program_a: torch.export.ExportedProgram, program_b: torch.export.ExportedProgram +) -> tuple[set[str], set[str]]: + """Compare and get unique ops in two exported programs. + + Args: + program_a: The first exported program. + program_b: The second exported program. + + Returns: + A tuple of two sets, where the first set contains the unique ops in the first program + and the second set contains the unique ops in the second program. + """ + program_a_ops = set(_count_fx_targets(program_a)) + program_b_ops = set(_count_fx_targets(program_b)) + return program_a_ops - program_b_ops, program_b_ops - program_a_ops diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_building.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_building.py new file mode 100644 index 0000000000000000000000000000000000000000..52e0e215fea2c6bd9819a5b5dfd21f3a3a449dc0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_building.py @@ -0,0 +1,734 @@ +"""NOTES: + +We need a typing module that will handling Python to ONNX type promotion for use. +For example, if we have torch.ops.aten.add(Tensor, 1.0), we need to promote 1.0 +to the same type as Tensor. The same thing needs to work for +torch.ops.aten.add(1.0, Tensor) as well, which means we need a mechanism to` +""" + +# mypy: allow-untyped-defs +# mypy: disable-error-code=union-attr +from __future__ import annotations + +import copy +import inspect +import logging +from collections.abc import Iterable, Mapping, Sequence +from typing import Any, TYPE_CHECKING + +from onnxscript import evaluator + +import torch +from torch.onnx._internal._lazy_import import onnx_ir as ir, onnxscript +from torch.onnx._internal.exporter import _errors, _schemas, _tensors + + +if TYPE_CHECKING: + import onnx + + +logger = logging.getLogger(__name__) + +ValidAttributeType = ( + ir.TensorProtocol + | int + | float + | bool + | str + | Sequence[int] + | Sequence[float] + | None +) + +AllowedArgType = ir.Value | Sequence[ir.Value | ValidAttributeType] | ValidAttributeType + + +# Logic for adapting inputs from general Python or PyTorch inputs to ONNX ir.Value +def _construct_named_inputs_and_attrs( + signature: ir.schemas.OpSignature, + args: Sequence[AllowedArgType], + kwargs: Mapping[str, AllowedArgType], +) -> tuple[dict[str, AllowedArgType], dict[str, ValidAttributeType]]: + """Construct two mappings: name to inputs and named to attributes based on the signature and args/kwargs. + + This function uses the OpSignature to determine which argument in args and kwargs corresponds to + which parameter in the signature. ONNX node inputs are stored in named_inputs, and attributes are + stored in named_attrs. If an _optional input_ is not provided, it is filled with None. + + Args: + signature: The OpSignature for the node. + args: The positional arguments for the node. + kwargs: The keyword arguments for the node. + + Returns: + A tuple of two mappings: named_inputs and named_attrs. + + Raises: + ValueError: If a required parameter is not provided. + """ + # 1. Construct the (named_inputs, named_attrs) mapping based on (args, kwargs) and the signature. + # a. Loop over all parameters in the signature and args together + # b. Depending on param.is_input, Record named_inputs[param.name] = arg or named_attrs[param.name] = arg + # c. Handle kwargs as well + # d. Fill in None if the input is not provided + named_inputs: dict[str, Any] = {} + named_attrs: dict[str, Any] = {} + reversed_args_stack = list(reversed(args)) + for param in signature.params: + if isinstance(param, ir.schemas.Parameter): + # Handle inputs + if reversed_args_stack: + # First exhaust the positional arguments + if param.variadic: + # Handle variadic arguments + named_inputs[param.name] = tuple(args) + reversed_args_stack.clear() + else: + named_inputs[param.name] = reversed_args_stack.pop() # type: ignore[assignment] + elif param.name in kwargs: + named_inputs[param.name] = kwargs[param.name] # type: ignore[assignment] + elif param.required: + raise ValueError( + f"Required parameter '{param.name}' is not provided. " + f"Signature: {signature}. Args: {args}. Kwargs: {kwargs}." + ) + else: + logger.debug( + "Optional parameter '%s' is not provided. Added as None. Signature: %s", + param.name, + signature, + ) + named_inputs[param.name] = None # type: ignore[assignment] + else: + # Handle attributes + attribute: ValidAttributeType | ir.Attr + if not isinstance(param, ir.schemas.AttributeParameter): + raise AssertionError(f"Expected AttributeParameter, got {type(param)}") + if reversed_args_stack: + # First exhaust the positional arguments + attribute = reversed_args_stack.pop() # type: ignore[assignment] + elif param.name in kwargs: + attribute = kwargs[param.name] # type: ignore[assignment] + elif param.default is not None: + attribute = param.default + else: + attribute = None + + if attribute is None: + if param.required: + raise ValueError( + f"Required attribute '{param.name}' is not provided. " + f"Signature: {signature}. Args: {args}. Kwargs: {kwargs}." + ) + else: + logger.debug( + "Optional attribute '%s' is None. Dropped. Signature: %s", + param.name, + signature, + ) + continue + + if isinstance(attribute, ir.Attr): + # Turn the attribute from an default value into an actual parameter for the node + attr_copied = copy.copy(attribute) + # Make sure the name is the same as the parameter name and not the name of the default parameter + attr_copied.name = param.name + attribute = attr_copied + + if isinstance(attribute, int) and param.type == ir.AttributeType.FLOAT: + # Convert the attribute to float if needed. This happens in PyTorch + # where an attribute marked as float can be passed as an int. + attribute = float(attribute) + named_attrs[param.name] = attribute + return named_inputs, named_attrs # type: ignore[return-value] + + +def _resolve_parameter_dtypes( + signature: ir.schemas.OpSignature, named_inputs: Mapping[str, AllowedArgType] +) -> Mapping[ir.schemas.TypeConstraintParam, ir.TypeProtocol]: + """Determine which parameter takes which type. + + Handle non-tensor input corner cases and type promotion. + + Requires: + All ir.Value in name_inputs should have type set. Their type should be + compatible with the type_constraint of the corresponding parameter in the signature. + + Args: + signature: The OpSignature for the node. + named_inputs: The mapping of parameter names to their arguments. + + Returns: + A mapping of Constraint names to ir.TypeProtocol. + """ + # a. Create type_binding: dict[str, ir.TypeProtocol] + # b. Iterate over all named_inputs + # b0. Find the corresponding parameter in the signature + # b1. If the argument is a Python constant, skip. + # b2. If the argument is a ir.Value, Bind {constraint: arg.type}. + type_binding = {} + for name, arg in named_inputs.items(): + param = signature.params_map[name] + if not isinstance(param, ir.schemas.Parameter): + raise AssertionError(f"Expected Parameter, got {type(param)}") + if isinstance(arg, (int, float, bool, str, Sequence, torch.Tensor)): + # Skip the Python constants because we do not know what dtype they should take yet + continue + elif isinstance(arg, ir.Value): + if arg.type is None: + # Skip the ir.Value if the type is not set + continue + # NOTE: We assume arg.type is compatible with the type_constraint + if arg.type is None: + raise AssertionError(f"Expected type to be set for {arg}") + # TODO(justinchuby): Implement type promotion logic here. + type_binding[param.type_constraint] = arg.type + return type_binding + + +def _determine_input_dtype( + param: ir.schemas.Parameter, + arg: AllowedArgType, + type_binding: Mapping[ir.schemas.TypeConstraintParam, ir.TypeProtocol], +) -> ir.DataType: + """Determine the dtype of the input that is a mix of Python constants and ir.Value.""" + if param.type_constraint in type_binding: + # A known dtype is available because it was resolved + return type_binding[param.type_constraint].dtype + if len(param.type_constraint.allowed_types) == 1: + # Only one type is allowed by the type constraint + return next(iter(param.type_constraint.allowed_types)).dtype + + # No dtype information available. Infer from the Python constant or (in the Sequence case) + # from a mix of Python constants and ir.Value + if isinstance(arg, bool): + return ir.DataType.BOOL + if isinstance(arg, float): + return ir.DataType.FLOAT + if isinstance(arg, int): + return ir.DataType.INT64 + if isinstance(arg, str): + return ir.DataType.STRING + if isinstance(arg, (ir.Tensor, ir.TensorProtocol)): + return arg.dtype + if isinstance(arg, complex): + return ir.DataType.FLOAT + if arg is None: + return ir.DataType.UNDEFINED + + # Handle sequences + if isinstance(arg, (tuple, list)): + if len(arg) == 0: + # Special case: Treat empty sequence as INT64 as they are typically used for shape + return ir.DataType.INT64 + + # Try to obtain the dtype from one of the values + for val in arg: + if isinstance(val, ir.Value) and val.dtype is not None: + return val.dtype + + if any(isinstance(val, float) for val in arg): + # If any float is present, the dtype is float + return ir.DataType.FLOAT + elif any(isinstance(val, int) for val in arg): + # Otherwise if any int is present, the dtype is int + return ir.DataType.INT64 + + raise ValueError( + f"Could not determine the dtype for the input '{param.name}'. " + f"param={param}, arg={arg}, param_type_constraint={param.type_constraint}, " + f"type_binding={type_binding}" + ) + + +def _allowed_types_are_sequence_types(allowed_types: Iterable[ir.TypeProtocol]) -> bool: + """Check if all allowed types are Sequence types.""" + return all(isinstance(t, ir.SequenceType) for t in allowed_types) + + +def _get_or_create_constant( + constant_farm: dict[ + tuple[ + bool + | int + | float + | str + | tuple[int, ...] + | tuple[float, ...] + | tuple[bool, ...], + ir.DataType, + ], + ir.Value, + ], + arg: bool + | int + | float + | str + | tuple[int, ...] + | tuple[float, ...] + | tuple[bool, ...] + | list[int] + | list[float] + | list[bool], + dtype: ir.DataType, + opset: onnxscript.values.Opset, +) -> ir.Value: + # float representation of complex numbers + if isinstance(arg, complex): + # Convert the complex number to a float + arg = (arg.real, arg.imag) + + if isinstance(arg, list): + # Make the arg hashable + # pyrefly: ignore [bad-argument-type] + arg = tuple(arg) + + constant_value = constant_farm.get((arg, dtype)) # type: ignore[arg-type] + if constant_value is None: + constant_tensor = ir.tensor(value=arg, dtype=dtype) + constant_value = opset.Constant(value=constant_tensor) + constant_farm[(arg, dtype)] = constant_value # type: ignore[arg-type,index] + return constant_value # type: ignore[return-value] + + +def _process_python_constants( + signature: ir.schemas.OpSignature, + named_inputs: dict[str, AllowedArgType], + type_binding: Mapping[ir.schemas.TypeConstraintParam, ir.TypeProtocol], + constant_farm: dict[ + tuple[ + bool | int | float | str | tuple[int, ...] | tuple[float, ...], + ir.DataType, + ], + ir.Value, + ], + opset: onnxscript.values.Opset, +) -> dict[str, ir.Value | None]: + """Convert Python constants to Constant nodes and list to Sequence nodes based on the dtype information. + + The added constants will be replacing values in named_inputs in place. + + Args: + signature: The OpSignature for the node. + named_inputs: The mapping of parameter names to their arguments. + type_binding: A mapping of Constraint names to ir.DataType. + constant_farm: A dictionary of {(py_value, ir.DataType): ir.Value} to store the deduplicated constants. + opset: The Opset to use for creating Constant nodes. + + Returns: + A mapping of parameter names to Python constants converted to constant Nodes. + """ + # 3. Convert Python constants to Constant nodes based on the dtype information; + # construct sequences + # a. Iterate over all parameters in the signature the second time + # b. If the parameter is in to_resolve_type: + # - If param.constraint in type_binding, + # Get the constant from constant_farm (deduplicated); + # otherwise set named_inputs[param.name] = Constant(value, dtype=type_binding[param.constraint]) + # - Otherwise, set named_inputs[param.name] = Constant(value) + for name, arg in named_inputs.items(): + param = signature.params_map[name] + if not isinstance(param, ir.schemas.Parameter): + raise AssertionError(f"Expected Parameter, got {type(param)}") + + if isinstance(arg, ir.Value): + # TODO(justinchuby): Cast the ir.Value here if needed + continue + + if ( + isinstance(arg, Sequence) + and len(arg) > 0 + and any(isinstance(val, ir.Value) for val in arg) + ): + # Skip the sequence of ir.Value. This is a variadic input or a Sequence input + # It will be handled by _process_python_sequences + continue + if param.variadic: + # Handled by _process_python_sequences + continue + if _allowed_types_are_sequence_types(param.type_constraint.allowed_types): + # Handled by _process_python_sequences + continue + + dtype = _determine_input_dtype(param, arg, type_binding) + + if arg is None: + constant_value = None + elif isinstance(arg, (ir.Tensor, ir.TensorProtocol)): + constant_value = opset.Constant(value=arg) + else: + # Deduplicate the constants + constant_value = _get_or_create_constant(constant_farm, arg, dtype, opset) # type: ignore[arg-type] + + named_inputs[param.name] = constant_value + return named_inputs # type: ignore[return-value] + + +def _reshape_to_1d_tensor(opset: onnxscript.values.Opset, arg: ir.Value) -> ir.Value: + """Reshape the input to a 1D tensor.""" + + return opset.Reshape( + arg, opset.Constant(value=ir.tensor([-1], dtype=ir.DataType.INT64)) + ) + + +def _process_python_sequences( + signature: ir.schemas.OpSignature, + named_inputs: dict[str, AllowedArgType], + type_binding: Mapping[ir.schemas.TypeConstraintParam, ir.TypeProtocol], + constant_farm: dict[ + tuple[ + bool + | int + | float + | str + | ir.TensorProtocol + | tuple[bool, ...] + | tuple[int, ...] + | tuple[float, ...], + ir.DataType, + ], + ir.Value, + ], + opset: onnxscript.values.Opset, +): + """Handle three types of sequences. + + 1. Variadic inputs + 2. Sequence input of ir.Value, + 3. Sequence of Python constants that contains ir.Value + """ + for name, arg in named_inputs.items(): + param = signature.params_map[name] + if not isinstance(param, ir.schemas.Parameter): + raise AssertionError(f"Expected Parameter, got {type(param)}") + + if not isinstance(arg, (tuple, list)): + continue + + if len(arg) == 0: + # Skip empty sequences + continue + + # 1. Sequence input of ir.Value + if _allowed_types_are_sequence_types(param.type_constraint.allowed_types): + # Turn the list into a Sequence node + # Constant op creation will be handled by the variadic case below when calling + # the SequenceConstruct op. + named_inputs[name] = opset.SequenceConstruct(*arg) + continue + + # 2. Variadic inputs + # NOTE: Variadic operators like Max can be called with mixed ir.Value and Python constants + # like `Max(0, ir.Value())` + # We need to convert the Python constants to Constant nodes + if param.variadic: + if all(isinstance(val, ir.Value) for val in arg): + # Skip the variadic input if all values are ir.Value + continue + + dtype = _determine_input_dtype(param, arg, type_binding) + new_args = [] + for val in arg: + if isinstance(val, ir.Value): + new_args.append(val) + else: + constant_tensor = ir.tensor(value=val, dtype=dtype) # type: ignore[arg-type] + constant_value = opset.Constant(value=constant_tensor) + new_args.append(constant_value) + named_inputs[name] = new_args + continue + else: + # 3. Concat the list as a single input + # E.g. [Value, 42] should be converted to op.Concat(Value, Constant(42)) + # when the expected input type is INT64 + # We assume this only happens for 0D cases + if all(isinstance(val, ir.Value) for val in arg): + expanded_args = [_reshape_to_1d_tensor(opset, val) for val in arg] + named_inputs[name] = opset.Concat(*expanded_args, axis=0) + continue + + dtype = _determine_input_dtype(param, arg, type_binding) + new_args = [] + for val in arg: + if isinstance(val, ir.Value): + new_args.append(_reshape_to_1d_tensor(opset, val)) + elif val is None: + # Skip None values + continue + elif isinstance(val, (ir.Tensor, ir.TensorProtocol)): + new_args.append( + _reshape_to_1d_tensor(opset, opset.Constant(value=val)) + ) + else: + # Turn the Python constant into 1D tensor for the constant + if not isinstance(val, (bool, int, float)): + raise AssertionError(f"Expected int or float, got {type(val)}") + new_args.append( + _get_or_create_constant(constant_farm, [val], dtype, opset) # type: ignore[arg-type] + ) + named_inputs[name] = opset.Concat(*new_args, axis=0) + continue + return named_inputs + + +def _determine_output_number( + signature: ir.schemas.OpSignature, named_attrs: Mapping[str, ValidAttributeType] +) -> int: + """Determine the number of outputs for the node with heuristics.""" + if signature.domain == "": + if signature.name == "BatchNormalization": + if not named_attrs.get("training_mode", 0): + return 1 + if signature.name == "Split": + num_outputs = named_attrs.get("num_outputs") + if num_outputs is not None and isinstance(num_outputs, int): + return num_outputs + else: + raise ValueError( + "Could not determine the number of outputs for Split. " + "num_outputs must be provided" + ) + return len(signature.outputs) + + +def _construct_node( + signature: ir.schemas.OpSignature, + named_inputs: Mapping[str, ir.Value | None], + named_attrs: Mapping[str, ValidAttributeType], + opset: onnxscript.values.Opset, + num_outputs: int, +) -> ir.Node: + """Construct the node with the inputs and attributes. + + Variadic inputs are flattened. + + Args: + signature: The OpSignature for the node. + named_inputs: The mapping of parameter names to their arguments. When we + do not have the schema of an operator, we do not know the names of + the inputs, in which case the names can be anything because they + are not used in this function. The data structure is passed in for + consistency with the other functions. + named_attrs: The mapping of attribute names to their values. + num_outputs: The number of outputs for the node. + """ + inputs: list[ir.Value | None] = [] + # Flatten variadic inputs + for value in named_inputs.values(): + if isinstance(value, Sequence): + inputs.extend(value) + else: + inputs.append(value) + + # If final inputs are None, strip them from the node inputs + for input in reversed(inputs): + if input is not None: + break + inputs.pop() + + # Construct and filter out None attributes + attributes = [ + attr + for attr in ir.convenience.convert_attributes(named_attrs) + if attr.value is not None + ] + outputs = [_tensors.SymbolicTensor(opset) for _ in range(num_outputs)] + return ir.Node( + signature.domain, + signature.name, + inputs=inputs, + attributes=attributes, + outputs=outputs, + version=signature.since_version, + ) + + +class OpRecorder(evaluator.Evaluator): + """An onnxscript Evaluator that captures the graph into ONNX IR.""" + + def __init__( + self, opset: onnxscript.values.Opset, constant_farm: dict[Any, ir.Value] + ) -> None: + self.nodes: list[ir.Node] = [] + self.opset = opset + self.functions: dict[ + ir.OperatorIdentifier, onnxscript.OnnxFunction | ir.Function + ] = {} + self.constant_farm = constant_farm + + def _call_op( + self, + op_signature: ir.schemas.OpSignature, + named_inputs: dict[str, AllowedArgType], + named_attrs: dict[str, ValidAttributeType], + num_outputs: int, + ) -> Sequence[_tensors.SymbolicTensor]: + """Record nodes for the given opschema and arguments. + + Args: + op_signature: The OpSchema containing the node signature. + named_inputs: The mapping of parameter names to their arguments. + named_attrs: The mapping of attribute names to their values. + """ + type_binding = _resolve_parameter_dtypes(op_signature, named_inputs) + try: + converted_named_inputs = _process_python_constants( + op_signature, named_inputs, type_binding, self.constant_farm, self.opset + ) + converted_named_inputs = _process_python_sequences( + op_signature, + converted_named_inputs, # type: ignore[arg-type] + type_binding, + self.constant_farm, + self.opset, + ) + + except Exception as e: + raise _errors.GraphConstructionError( + f"Error processing Python constants for operator '{op_signature.domain}::{op_signature.name}'. " + f"named_inputs={named_inputs}, named_attrs={named_attrs}, opset={self.opset}, op_signature={op_signature}." + ) from e + + try: + self.nodes.append( + node := _construct_node( + op_signature, + converted_named_inputs, + named_attrs, + self.opset, + num_outputs, + ) + ) + except Exception as e: + raise _errors.GraphConstructionError( + f"Error constructing node for operator '{op_signature.domain}::{op_signature.name}'. " + f"named_inputs={named_inputs}, converted_named_inputs={converted_named_inputs}, " + f"named_attrs={named_attrs}, opset={self.opset}, op_signature={op_signature}." + ) from e + return node.outputs # type: ignore[return-value] + + def eval( + self, + schema: onnx.defs.OpSchema, + args: Sequence[AllowedArgType], # type: ignore[override] + kwargs: Mapping[str, AllowedArgType], + ) -> _tensors.SymbolicTensor | Sequence[_tensors.SymbolicTensor]: + try: + op_signature = ir.schemas.OpSignature.from_op_schema(schema) + named_inputs, named_attrs = _construct_named_inputs_and_attrs( + op_signature, args, kwargs + ) + # TODO(justinchuby): Handle cast + if schema.name == "CastLike": + if len(named_inputs) != 2: + raise AssertionError(f"Expected 2 inputs, got {len(named_inputs)}") + # Skip CastLike if the input and output types are the same + src_input = named_inputs["input"] + target_type = named_inputs["target_type"] + + if ( + isinstance(src_input, ir.Value) + and isinstance(target_type, ir.Value) + and src_input.dtype is not None + and target_type.dtype is not None + ): + # dtypes are available + if src_input.dtype == target_type.dtype: + # Same type. No cast needed + return src_input # type: ignore[return-value] + else: + # Create a Cast node + return self.opset.Cast(src_input, to=target_type.dtype) # type: ignore[union-attr,return-value] + + num_outputs = _determine_output_number(op_signature, named_attrs) + outputs = self._call_op( + op_signature, named_inputs, named_attrs, num_outputs + ) + if len(outputs) == 1: + return outputs[0] + return outputs + except Exception as e: + raise _errors.GraphConstructionError( + f"Error calling operator '{schema.name}' with args {args} and kwargs {kwargs}." + ) from e + + def eval_function( # type: ignore[override] + self, + function: onnxscript.OnnxFunction, + args: Sequence[AllowedArgType], + kwargs: Mapping[str, AllowedArgType], + ) -> _tensors.SymbolicTensor | Sequence[_tensors.SymbolicTensor] | bool | int: + try: + # NOTE: signature should be written to function in the registration process + if hasattr(function, "_pt_onnx_signature"): + op_signature = function._pt_onnx_signature # type: ignore[attr-defined] + else: + op_signature = _schemas.op_signature_from_function( + function, + function.function_ir.domain, + function.name, + since_version=function.opset.version, + ) + function._pt_onnx_signature = op_signature # type: ignore[attr-defined] + + named_inputs, named_attrs = _construct_named_inputs_and_attrs( + op_signature, args, kwargs + ) + + # TODO(after torchlib migration): Remove traceable function handling + # NOTE: We need to call traceable functions after the _construct_named_inputs_and_attrs + # call because it will filter out the unexpected kwargs for us. + if function.traceable: + # Trace the function call instead of adding the function as a node + # Turn the ir.Attr objects into Python constants first + named_attrs = { + name: attr.value if isinstance(attr, ir.Attr) else attr + for name, attr in named_attrs.items() + } + + # Use the type binding to resolve the dtypes of the inputs, and + # convert Python constants to Constant nodes + type_binding = _resolve_parameter_dtypes(op_signature, named_inputs) + try: + # _process_python_sequences is not here because we want to preserve python list + # properties for the function call + converted_named_inputs = _process_python_constants( + op_signature, + named_inputs, + type_binding, + self.constant_farm, + self.opset, + ) + + except Exception as e: + raise _errors.GraphConstructionError( + f"Error processing Python constants for operator '{op_signature.domain}::{op_signature.name}'. " + f"named_inputs={named_inputs}, named_attrs={named_attrs}, opset={self.opset}, op_signature={op_signature}." + ) from e + + return function.function(**converted_named_inputs, **named_attrs) + + outputs = self._call_op( + op_signature, + named_inputs, + named_attrs, + len(op_signature.outputs), + ) + + self.functions[(function.function_ir.domain, function.name, "")] = function + if len(outputs) == 1: + return outputs[0] + return outputs + except Exception as e: + try: + source_file = inspect.getsourcefile(function.function) + _, lineno = inspect.getsourcelines(function.function) + except Exception: + source_file = lineno = None + raise _errors.GraphConstructionError( + f"Error calling function '{function.name}' with args {args} and kwargs {kwargs}." + + f" The function is defined at '{source_file}:{lineno}'." + if source_file + else "" + ) from e diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_capture_strategies.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_capture_strategies.py new file mode 100644 index 0000000000000000000000000000000000000000..0b262ff77ebb4a1ae4e5edc640c6ccaaa7204a4d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_capture_strategies.py @@ -0,0 +1,304 @@ +"""Strategies for capturing ExportedPrograms.""" + +# mypy: allow-untyped-defs +from __future__ import annotations + +import abc +import contextlib +import dataclasses +import datetime +import logging +import pathlib +from typing import Any, TYPE_CHECKING + +import torch +from torch.onnx import _flags + + +if TYPE_CHECKING: + import os + from collections.abc import Callable + + +logger = logging.getLogger(__name__) + + +def _verbose_printer(verbose: bool | None) -> Callable[..., None]: + """Prints messages based on `verbose`.""" + if verbose is False: + return lambda *_, **__: None + + return lambda *args, **kwargs: print("[torch.onnx]", *args, **kwargs) + + +def _take_first_line(text: str) -> str: + """Take the first line of a text.""" + lines = text.split("\n", maxsplit=1) + first_line = lines[0] + if len(lines) > 1: + first_line += "[...]" + return first_line + + +@contextlib.contextmanager +def _patch_dynamo_unsupported_functions(): + """Patch PyTorch to bypass some functions torch.export.export does not support.""" + # TODO: Remove the patches once dynamo supports these functions. + import torch.jit + + # Replace torch.jit.isinstance with isinstance + jit_isinstance = torch.jit.isinstance + # pyrefly: ignore [bad-assignment] + torch.jit.isinstance = isinstance + logger.info("Replaced torch.jit.isinstance with isinstance to allow dynamo tracing") + try: + yield + finally: + torch.jit.isinstance = jit_isinstance + + +@dataclasses.dataclass +class Result: + exported_program: torch.export.ExportedProgram | None + strategy: str + exception: Exception | None = None + + @property + def success(self) -> bool: + """Whether the capture was successful. + + An exception can still be recorded even if the capture was successful. In + this case the exception is informational only. For example, draft_export + can record an exception if there are warnings during the export. The exceptions + will go into the onnx export report when report=True. + """ + return self.exported_program is not None + + +class CaptureStrategy(abc.ABC): + """Strategy for capturing a module as ExportedProgram. + + To use a strategy, create an instance and call it with the model, args, kwargs, and dynamic_shapes. + Example:: + + strategy = TorchExportNonStrictStrategy(verbose=True) + result = strategy(model, args, kwargs, dynamic_shapes) + """ + + def __init__( + self, + *, + verbose: bool = False, + dump: bool = False, + artifacts_dir: str | os.PathLike = ".", + timestamp: str | None = None, + ) -> None: + """Initialize the strategy. + + Args: + verbose: Whether to print verbose messages. + dump: Whether to dump the intermediate artifacts to a file. + """ + self._verbose_print = _verbose_printer(verbose) + self._dump = dump + self._artifacts_dir = pathlib.Path(artifacts_dir) + self._timestamp = timestamp or datetime.datetime.now().strftime( + "%Y-%m-%d_%H-%M-%S-%f" + ) + self._exception: Exception | None = None + + def __call__( + self, + model: torch.nn.Module | torch.jit.ScriptFunction, + args: tuple[Any, ...], + kwargs: dict[str, Any] | None, + dynamic_shapes, + ) -> Result: + self._enter(model) + if kwargs is None: + kwargs = {} + try: + exported_program = self._capture(model, args, kwargs, dynamic_shapes) + except Exception as e: + self._failure(model, e) + return Result( + exported_program=None, + strategy=self.__class__.__name__, + exception=e, + ) + self._success(model) + return Result( + exported_program, + strategy=self.__class__.__name__, + exception=self._exception, + ) + + @abc.abstractmethod + def _capture( + self, model, args, kwargs, dynamic_shapes + ) -> torch.export.ExportedProgram: + raise NotImplementedError + + def _enter(self, model: torch.nn.Module | torch.jit.ScriptFunction) -> None: + return + + def _success(self, model: torch.nn.Module | torch.jit.ScriptFunction) -> None: + return + + def _failure( + self, model: torch.nn.Module | torch.jit.ScriptFunction, e: Exception + ) -> None: + return + + +class TorchExportStrictStrategy(CaptureStrategy): + def _capture( + self, model, args, kwargs, dynamic_shapes + ) -> torch.export.ExportedProgram: + with ( + _patch_dynamo_unsupported_functions(), + # Support the dynamism with 0/1 input dim + torch.fx.experimental._config.patch(backed_size_oblivious=True), # type: ignore[attr-defined] + ): + try: + return torch.export.export( + model, + args, + kwargs=kwargs, + dynamic_shapes=dynamic_shapes, + strict=True, + prefer_deferred_runtime_asserts_over_guards=_flags.PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS, + ) + except torch._dynamo.exc.UserError as exc: + # Refine the dynamic shapes based on the suggested fixes. + try: + new_shapes = torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes( + exc.msg, dynamic_shapes + ) + except Exception: + # If the dynamic shapes cannot be refined, re-raise the exception. + raise exc from None + return torch.export.export( + model, + args, + kwargs=kwargs, + dynamic_shapes=new_shapes, + strict=True, + prefer_deferred_runtime_asserts_over_guards=_flags.PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS, + ) + + def _enter(self, model) -> None: + model_repr = _take_first_line(repr(model)) + self._verbose_print( + f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=True)`..." + ) + + def _success(self, model) -> None: + model_repr = _take_first_line(repr(model)) + self._verbose_print( + f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=True)`... ✅" + ) + + def _failure(self, model, e) -> None: + del e # Unused + model_repr = _take_first_line(repr(model)) + self._verbose_print( + f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=True)`... ❌" + ) + + +class TorchExportNonStrictStrategy(CaptureStrategy): + def _capture( + self, model, args, kwargs, dynamic_shapes + ) -> torch.export.ExportedProgram: + with ( + # Support the dynamism with 0/1 input dim + torch.fx.experimental._config.patch(backed_size_oblivious=True), # type: ignore[attr-defined] + ): + try: + return torch.export.export( + model, + args, + kwargs=kwargs, + dynamic_shapes=dynamic_shapes, + strict=False, + prefer_deferred_runtime_asserts_over_guards=_flags.PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS, + ) + except torch._dynamo.exc.UserError as exc: + # Refine the dynamic shapes based on the suggested fixes. + try: + new_shapes = torch.export.dynamic_shapes.refine_dynamic_shapes_from_suggested_fixes( + exc.msg, dynamic_shapes + ) + except Exception: + # If the dynamic shapes cannot be refined, re-raise the exception. + raise exc from None + return torch.export.export( + model, + args, + kwargs=kwargs, + dynamic_shapes=new_shapes, + strict=False, + prefer_deferred_runtime_asserts_over_guards=_flags.PREFER_DEFERRED_RUNTIME_ASSERTS_OVER_GUARDS, + ) + + def _enter(self, model) -> None: + model_repr = _take_first_line(repr(model)) + self._verbose_print( + f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=False)`..." + ) + + def _success(self, model) -> None: + model_repr = _take_first_line(repr(model)) + self._verbose_print( + f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=False)`... ✅" + ) + + def _failure(self, model, e) -> None: + del e # Unused + model_repr = _take_first_line(repr(model)) + self._verbose_print( + f"Obtain model graph for `{model_repr}` with `torch.export.export(..., strict=False)`... ❌" + ) + + +class TorchExportDraftExportStrategy(CaptureStrategy): + def _capture( + self, model, args, kwargs, dynamic_shapes + ) -> torch.export.ExportedProgram: + ep = torch.export.draft_export( + model, args, kwargs=kwargs, dynamic_shapes=dynamic_shapes + ) + report = ep._report # type: ignore[attr-defined] + if not report.successful(): + self._exception = RuntimeError(str(report)) + self._verbose_print(f"Draft Export report:\n{report}") + return ep + + def _enter(self, model) -> None: + model_repr = _take_first_line(repr(model)) + self._verbose_print( + f"Obtain model graph for `{model_repr}` with `torch.export.draft_export`..." + ) + + def _success(self, model) -> None: + model_repr = _take_first_line(repr(model)) + self._verbose_print( + f"Obtain model graph for `{model_repr}` with `torch.export.draft_export`... ✅" + ) + + def _failure(self, model, e) -> None: + del e # Unused + model_repr = _take_first_line(repr(model)) + self._verbose_print( + f"Obtain model graph for `{model_repr}` with `torch.export.draft_export`... ❌" + ) + + +CAPTURE_STRATEGIES: tuple[type[CaptureStrategy], ...] = ( + TorchExportNonStrictStrategy, # strict=False is preferred over strict=True because it does not have dynamo issues + TorchExportStrictStrategy, +) + +if _flags.ENABLE_DRAFT_EXPORT: + CAPTURE_STRATEGIES = (*CAPTURE_STRATEGIES, TorchExportDraftExportStrategy) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_compat.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_compat.py new file mode 100644 index 0000000000000000000000000000000000000000..f6292e1e4d09c26f5eec3ce8e26372b24ebc4fb4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_compat.py @@ -0,0 +1,202 @@ +"""Compatibility functions for the torch.onnx.export API.""" + +# mypy: allow-untyped-defs +# mypy: disable-error-code=attr-defined +from __future__ import annotations + +import io +import logging +import warnings +from collections.abc import Callable, Mapping, Sequence +from typing import Any, TYPE_CHECKING + +import torch +from torch.onnx import _constants as onnx_constants +from torch.onnx._internal._lazy_import import onnx +from torch.onnx._internal.exporter import ( + _constants, + _core, + _dynamic_shapes, + _exportable_module, + _onnx_program, + _registration, +) + + +if TYPE_CHECKING: + import os + +logger = logging.getLogger(__name__) + + +def _get_torch_export_args( + args: tuple[Any, ...], + kwargs: dict[str, Any] | None, +) -> tuple[tuple[Any, ...], dict[str, Any] | None]: + """Obtain the arguments for torch.onnx.export from the model and the input arguments.""" + if not kwargs and args and isinstance(args[-1], dict): + kwargs = args[-1] + args = args[:-1] + return args, kwargs + + +def export_compat( + model: torch.nn.Module + | torch.export.ExportedProgram + | torch.jit.ScriptModule + | torch.jit.ScriptFunction, + args: tuple[Any, ...], + f: str | os.PathLike | None = None, + *, + kwargs: dict[str, Any] | None = None, + export_params: bool = True, + verbose: bool | None = None, + input_names: Sequence[str] | None = None, + output_names: Sequence[str] | None = None, + opset_version: int | None = onnx_constants.ONNX_DEFAULT_OPSET, + custom_translation_table: dict[Callable, Callable] | None = None, + dynamic_axes: Mapping[str, Mapping[int, str]] + | Mapping[str, Sequence[int]] + | None = None, + dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None = None, + keep_initializers_as_inputs: bool = False, + external_data: bool = True, + report: bool = False, + optimize: bool = True, + verify: bool = False, + profile: bool = False, + dump_exported_program: bool = False, + artifacts_dir: str | os.PathLike = ".", +) -> _onnx_program.ONNXProgram: + if opset_version is None: + opset_version = onnx_constants.ONNX_DEFAULT_OPSET + + if isinstance(model, torch.nn.Module): + if model.training: + warnings.warn( + "Exporting a model while it is in training mode. " + "Please ensure that this is intended, as it may lead to " + "different behavior during inference. " + "Calling model.eval() before export is recommended.", + UserWarning, + stacklevel=3, + ) + + if isinstance(model, _exportable_module.ExportableModule): + # Skip argument extraction if args or kwargs are provided + if not args and not kwargs: + args, kwargs = model.example_arguments() + if input_names is None: + input_names = model.input_names() + if output_names is None: + output_names = model.output_names() + if dynamic_shapes is None: + dynamic_shapes = model.dynamic_shapes() + + if isinstance(model, torch.export.ExportedProgram): + # We know the model is already exported program, so the args, kwargs, and dynamic_shapes + # are not used + dynamic_shapes = dynamic_shapes or {} + else: + args, kwargs = _get_torch_export_args(args, kwargs) + if dynamic_shapes is None and dynamic_axes is not None: + warnings.warn( + "# 'dynamic_axes' is not recommended when dynamo=True, " + "and may lead to 'torch._dynamo.exc.UserError: Constraints violated.' " + "Supply the 'dynamic_shapes' argument instead if export is unsuccessful.", + UserWarning, + stacklevel=3, + ) + try: + dynamic_shapes, args, kwargs = ( + _dynamic_shapes.from_dynamic_axes_to_dynamic_shapes( + model, + args, + kwargs, + dynamic_axes=dynamic_axes, + input_names=input_names, + output_names=set(output_names or ()), + ) + ) + except Exception as e: + raise RuntimeError( + "# Failed to convert 'dynamic_axes' to 'dynamic_shapes'. " + "Please provide 'dynamic_shapes' directly. " + "Refer to the documentation for 'torch.export.export' for more information on dynamic shapes." + ) from e + + dynamic_shapes_with_export_dim, need_axis_mapping = ( + _dynamic_shapes.convert_str_to_export_dim(dynamic_shapes) + ) + + if opset_version < _constants.TORCHLIB_OPSET: + logger.warning( + "Setting ONNX exporter to use operator set version %s because " + "the requested opset_version %s is a lower version than we have implementations for. " + "Automatic version conversion will be performed, which may not be successful " + "at converting to the requested version. If version conversion is unsuccessful, " + "the opset version of the exported model will be kept at %s. " + "Please consider setting opset_version >=%s to leverage latest ONNX features", + _constants.TORCHLIB_OPSET, + opset_version, + _constants.TORCHLIB_OPSET, + _constants.TORCHLIB_OPSET, + ) + registry_opset_version = _constants.TORCHLIB_OPSET + else: + registry_opset_version = opset_version + + registry = _registration.ONNXRegistry().from_torchlib( + opset_version=registry_opset_version + ) + if custom_translation_table is not None: + for torch_op, onnx_op in custom_translation_table.items(): + # TODO(justinchuby): Support complex inputs with annotations + if isinstance(onnx_op, Sequence): + raise TypeError( + "The value in custom_translation_table should be a single callable, not a sequence" + ) + registry.register_op(torch_op, onnx_op, is_complex=False) + + onnx_program = _core.export( + model, + args, + kwargs, + registry=registry, + dynamic_shapes=dynamic_shapes_with_export_dim, + input_names=input_names, + output_names=output_names, + profile=profile, + report=report, + verify=verify, + dump_exported_program=dump_exported_program, + artifacts_dir=artifacts_dir, + verbose=verbose, + optimize=optimize, + opset_version=opset_version, + ) + + if need_axis_mapping and dynamic_shapes is not None: + onnx_program._rename_dynamic_axes(dynamic_shapes) + + if f is not None: + if isinstance(f, io.BytesIO): + # For legacy export compatibility, we allow f to be a BytesIO object. + # This is not explicitly supported but we may need to maintain the + # behavior indefinitely. + warnings.warn( + "Saving ONNX model to a BytesIO object is deprecated. " + "Please use a file path instead.", + DeprecationWarning, + stacklevel=2, + ) + onnx.save(onnx_program.model_proto, f) + else: + onnx_program.save( + f, + include_initializers=export_params, + keep_initializers_as_inputs=keep_initializers_as_inputs, + external_data=external_data, + ) + + return onnx_program diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_constants.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..bb5206e6e0b7eaa80c4ac445ab9eb0616f598df1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_constants.py @@ -0,0 +1,7 @@ +# ir_version used for the ONNX file. See https://github.com/onnx/onnx/blob/main/docs/IR.md#onnx-versioning +ONNX_IR_VERSION = 10 +# The opset version torchlib is implemented with. Update this number when updating torchlib +TORCHLIB_OPSET = 18 +TORCHLIB_DOMAIN = "pkg.torch.onnx" +# Domain used for functions translated from subgraphs +LOCAL_FUNCTION_DOMAIN = "pkg.torch.__subgraph__" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_core.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_core.py new file mode 100644 index 0000000000000000000000000000000000000000..3b1e1dfc2162039c51eb7b8ae91dc455ee1eb019 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_core.py @@ -0,0 +1,1730 @@ +# mypy: allow-untyped-defs +# flake8: noqa: B950 We do not need flake8 as it complains line length +from __future__ import annotations + +import ctypes +import datetime +import inspect +import itertools +import logging +import operator +import pathlib +import sys +import textwrap +import traceback +import typing +from collections.abc import Callable, Mapping, Sequence +from typing import Any, Literal + +import onnxscript +import onnxscript.evaluator +from onnxscript import ir +from onnxscript.ir import convenience as ir_convenience + +import torch +import torch.fx +from torch.export import graph_signature +from torch.onnx._internal._lazy_import import onnxscript_apis +from torch.onnx._internal.exporter import ( + _analysis, + _building, + _capture_strategies, + _constants, + _dispatching, + _errors, + _flags, + _fx_passes, + _ir_passes, + _onnx_program, + _registration, + _reporting, + _tensors, + _type_casting, + _verification, +) + + +if typing.TYPE_CHECKING: + import os + + import numpy.typing as npt + + +# Define utilities to convert PyTorch data types so users do not need to specify manually +_TORCH_DTYPE_TO_ONNX: dict[torch.dtype, ir.DataType] = { + torch.bfloat16: ir.DataType.BFLOAT16, + torch.bool: ir.DataType.BOOL, + torch.complex128: ir.DataType.COMPLEX128, + torch.complex64: ir.DataType.COMPLEX64, + torch.float16: ir.DataType.FLOAT16, + torch.float32: ir.DataType.FLOAT, + torch.float64: ir.DataType.DOUBLE, + torch.float8_e4m3fn: ir.DataType.FLOAT8E4M3FN, + torch.float8_e4m3fnuz: ir.DataType.FLOAT8E4M3FNUZ, + torch.float8_e5m2: ir.DataType.FLOAT8E5M2, + torch.float8_e5m2fnuz: ir.DataType.FLOAT8E5M2FNUZ, + torch.float4_e2m1fn_x2: ir.DataType.FLOAT4E2M1, + torch.int16: ir.DataType.INT16, + torch.int32: ir.DataType.INT32, + torch.int64: ir.DataType.INT64, + torch.int8: ir.DataType.INT8, + torch.uint8: ir.DataType.UINT8, + torch.uint16: ir.DataType.UINT16, + torch.uint32: ir.DataType.UINT32, + torch.uint64: ir.DataType.UINT64, +} +_BLUE = "\033[96m" +_END = "\033[0m" + +_STEP_ONE_ERROR_MESSAGE = textwrap.dedent( + f"""\ + Failed to export the model with torch.export. {_BLUE}This is step 1/3{_END} of exporting the model to ONNX. Next steps: + - Modify the model code for `torch.export.export` to succeed. Refer to https://pytorch.org/docs/stable/generated/exportdb/index.html for more information. + - Debug `torch.export.export` and submit a PR to PyTorch. + - Create an issue in the PyTorch GitHub repository against the {_BLUE}*torch.export*{_END} component and attach the full error stack as well as reproduction scripts.""" +) + +_STEP_TWO_ERROR_MESSAGE = textwrap.dedent( + f"""\ + Failed to decompose the FX graph for ONNX compatibility. {_BLUE}This is step 2/3{_END} of exporting the model to ONNX. Next steps: + - Create an issue in the PyTorch GitHub repository against the {_BLUE}*torch.export*{_END} component and attach the full error stack as well as reproduction scripts. + - Create an error report with `torch.onnx.export(..., report=True)`, and save the ExportedProgram as a pt2 file. Create an issue in the PyTorch GitHub repository against the {_BLUE}*onnx*{_END} component. Attach the error report and the pt2 model.""" +) + +_STEP_THREE_ERROR_MESSAGE = textwrap.dedent( + f"""\ + Failed to convert the exported program to an ONNX model. {_BLUE}This is step 3/3{_END} of exporting the model to ONNX. Next steps: + - If there is a missing ONNX function, implement it and register it to the registry. + - If there is an internal error during ONNX conversion, debug the error and submit a PR to PyTorch. + - Create an error report with `torch.onnx.export(..., report=True)`, and save the ExportedProgram as a pt2 file. Create an issue in the PyTorch GitHub repository against the {_BLUE}*onnx*{_END} component. Attach the error report and the pt2 model.""" +) + +logger = logging.getLogger(__name__) +# The current tracer that is being used to trace the operators, +# used by torch/onnx/_internal/exporter/_torchlib/ops/hop.py +current_tracer: _building.OpRecorder | None = None + + +def torch_dtype_to_onnx_dtype(dtype: torch.dtype) -> ir.DataType: + return _TORCH_DTYPE_TO_ONNX[dtype] + + +class TorchTensor(ir.Tensor): + def __init__(self, tensor: torch.Tensor, name: str | None = None) -> None: + # Pass the tensor as the raw data to ir.Tensor's constructor + if tensor.dtype == torch.float4_e2m1fn_x2: + # Change the shape to the unpacked shape + shape = ir.Shape(_type_casting.get_float4_shape(tensor), frozen=True) + else: + # The base class will set the shape to the tensor's shape + shape = None + super().__init__( + tensor, + dtype=torch_dtype_to_onnx_dtype(tensor.dtype), + shape=shape, + name=name, + ) + + def numpy(self) -> npt.NDArray: + self.raw: torch.Tensor + + # Handle dtypes that are not natively supported by NumPy: + # We pick an uint dtype that has the same size as the original dtype, + # view the tensor as that dtype so that it is convertible to NumPy, + # and then view it back to the proper dtype (using ml_dtypes obtained by + # calling dtype.numpy()). + if self.dtype == ir.DataType.BFLOAT16: + return ( + self.raw.view(torch.uint16).numpy(force=True).view(self.dtype.numpy()) + ) + if self.dtype in { + ir.DataType.FLOAT8E4M3FN, + ir.DataType.FLOAT8E4M3FNUZ, + ir.DataType.FLOAT8E5M2, + ir.DataType.FLOAT8E5M2FNUZ, + }: + return self.raw.view(torch.uint8).numpy(force=True).view(self.dtype.numpy()) + if self.dtype == ir.DataType.FLOAT4E2M1: + return _type_casting.unpack_float4x2_as_uint8(self.raw).view( + self.dtype.numpy() + ) + + return self.raw.numpy(force=True) + + def __array__(self, dtype: Any = None, copy: bool | None = None) -> npt.NDArray: + del copy # Unused, but needed for the signature + if dtype is None: + return self.numpy() + return self.numpy().__array__(dtype) + + def _get_cbytes(self): + """Get a ctypes byte array pointing to the tensor data.""" + import torch._subclasses.fake_tensor + + with torch._subclasses.fake_tensor.unset_fake_temporarily(): + # Disable any fake mode so calling detach() etc. will return a real tensor + tensor = self.raw.detach().cpu().contiguous() + + if isinstance(tensor, torch._subclasses.fake_tensor.FakeTensor): + raise TypeError( + f"Cannot take content out from the FakeTensor ('{self.name}'). Please replace the tensor " + "with a tensor backed by real data using ONNXProgram.apply_weights() " + "or save the model without initializers by setting include_initializers=False." + ) + + # Return the tensor to ensure it is not garbage collected while the ctypes array is in use + return tensor, ( + ctypes.c_ubyte * tensor.element_size() * tensor.numel() + ).from_address(tensor.data_ptr()) + + def tobytes(self) -> bytes: + # On big-endian machines, call the super's tobytes() which returns a little-endian result. + if sys.byteorder == "big": + return super().tobytes() + # Implement tobytes to support native PyTorch types so we can use types like bloat16 + # Reading from memory directly is also more efficient because + # it avoids copying to a NumPy array + _, data = self._get_cbytes() + return bytes(data) + + def tofile(self, file) -> None: + # On big-endian machines, call the super's tofile() which returns a little-endian result. + if sys.byteorder == "big": + return super().tofile(file) + _, data = self._get_cbytes() + return file.write(data) + + +# https://github.com/pytorch/pytorch/blob/ee6cb6daa173896f8ea1876266a19775aaa4f610/torch/export/graph_signature.py#L56C1-L62C19 +# class InputKind(Enum): +# USER_INPUT = auto() +# PARAMETER = auto() +# BUFFER = auto() +# CONSTANT_TENSOR = auto() +# CUSTOM_OBJ = auto() +# TOKEN = auto() + +# https://github.com/pytorch/pytorch/blob/ee6cb6daa173896f8ea1876266a19775aaa4f610/torch/export/graph_signature.py#L89C1-L96C19 +# class OutputKind(Enum): +# USER_OUTPUT = auto() +# LOSS_OUTPUT = auto() +# BUFFER_MUTATION = auto() +# GRADIENT_TO_PARAMETER = auto() +# GRADIENT_TO_USER_INPUT = auto() +# USER_INPUT_MUTATION = auto() +# TOKEN = auto() + + +def _set_shape_types( + values: Sequence[ir.Value], + meta_vals: Sequence[torch.Tensor], + complex_to_float: bool = True, +) -> None: + if not isinstance(meta_vals, Sequence): + logger.warning( + "Expected meta_vals to be a sequence, but got %s. There may be an internal error.", + meta_vals, + ) + meta_vals = (meta_vals,) + for value, meta_val in zip(values, meta_vals): + _set_shape_type(value, meta_val, complex_to_float=complex_to_float) + + +def _set_shape_type( + value: ir.Value, + meta_val: torch.Tensor + | torch.SymBool + | torch.SymInt + | torch.SymFloat + | tuple[torch.Tensor], + complex_to_float: bool, +) -> None: + if isinstance(meta_val, tuple): + logger.warning("Setting shape and type of tensors is not supported yet") + if isinstance(meta_val, torch.Tensor): + dims = [] + shape: tuple[int, ...] + if meta_val.dtype == torch.float4_e2m1fn_x2: + # Change the shape to the unpacked shape + shape = _type_casting.get_float4_shape(meta_val) + else: + shape = meta_val.shape + for dim in shape: + if isinstance(dim, int): + dims.append(dim) + else: + # pyrefly: ignore [bad-argument-type] + dims.append(str(dim.node)) + + # If the dtype is set already (e.g. by the onnx_symbolic ops), + # we don't need to set it again. + # + # When a user specifies complex in onnx_symbolic, we consider that to + # be the intention even though non of the ONNX ops deals with complex values. + # In this case, we don't change the dtype or the shape of the tensor. + if value.dtype is None: + value.dtype = torch_dtype_to_onnx_dtype(meta_val.dtype) + if complex_to_float and meta_val.dtype.is_complex: + value.dtype = torch_dtype_to_onnx_dtype(meta_val.dtype.to_real()) + # Add 2 as the last dimension if the tensor is complex to hold the real/imag parts + dims.append(2) + + value.shape = ir.Shape(dims) + elif isinstance(meta_val, (int, torch.SymInt)): + # aten::sym_size output is a int, not a tensor, which stands + # for the size of one dim. We treat it as a scalar. + value.dtype = ir.DataType.INT64 + value.shape = ir.Shape([]) + elif isinstance(meta_val, (bool, torch.SymBool)): + value.dtype = ir.DataType.BOOL + value.shape = ir.Shape([]) + elif isinstance(meta_val, (float, torch.SymFloat)): + value.dtype = ir.DataType.FLOAT + value.shape = ir.Shape([]) + + +def _get_qualified_module_name(cls: Any) -> str: + if isinstance(cls, str): + return cls + module = cls.__module__ + if module is None or module == str.__class__.__module__: + return cls.__name__ + return module + "." + cls.__name__ + + +def _get_node_namespace(node: torch.fx.Node) -> tuple[str, list[str], list[str]]: + """Get the namespace and scope of the node. + + Example:: + + { + 'L__self__': ('', ), + 'L__self___avgpool': ('avgpool', ) + } + + Will yield + + namespace: ": torchvision.models.resnet.ResNet/avgpool: torch.nn.modules.pooling.AdaptiveAvgPool2d/node_name: node_target" + class_hierarchy: ["torchvision.models.resnet.ResNet", "torch.nn.modules.pooling.AdaptiveAvgPool2d", ] + name_scopes: ["", "avgpool", ] + + Args: + node: The node to get the namespace and scope of. + + Returns: + (namespace, class_hierarchy, name_scope) + """ + nn_module_stack = node.meta.get("nn_module_stack") + logger.debug("%s", nn_module_stack) + if nn_module_stack is None: + logger.warning( + "nn_module_stack not found for node '%s'. Skip adding metadata...", + node.name, + ) + return f"{node.name}: {node.target}", [str(node.target)], [node.name] + namespaces = [] + class_hierarchy = [] + name_scopes = [] + for name, nn_module in nn_module_stack.values(): + name_scopes.append(name) + nn_module_name = _get_qualified_module_name(nn_module) + class_hierarchy.append(nn_module_name) + namespaces.append(f"{name}: {_get_qualified_module_name(nn_module)}") + namespaces.append(f"{node.name}: {node.target}") + class_hierarchy.append(str(node.target)) + name_scopes.append(node.name) + + return "/".join(namespaces), class_hierarchy, name_scopes + + +def _set_node_metadata(fx_node: torch.fx.Node, ir_node: ir.Node) -> None: + """Adds namespace and other node metadata to the ONNX node.""" + namespace, class_hierarchy, name_scopes = _get_node_namespace(fx_node) + ir_node.metadata_props["namespace"] = namespace + ir_node.metadata_props["pkg.torch.onnx.class_hierarchy"] = repr(class_hierarchy) + ir_node.metadata_props["pkg.torch.onnx.name_scopes"] = repr(name_scopes) + ir_node.metadata_props["pkg.torch.onnx.fx_node"] = str(fx_node.format_node()) + ir_node.metadata_props["pkg.torch.onnx.stack_trace"] = fx_node.meta.get( + "stack_trace", "" + ) + + +def _handle_getitem_node( + node: torch.fx.Node, node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]] +) -> ir.Value: + """Handle a getitem node. + + Add the input value it is getting to the mapping, then return the value. + + There are two cases for this node: + 1. The output is a Sequence (traced), we can simply get the value from the sequence + 2. The output is produced by a SplitToSequence node, we need to get the value from the sequence value + This function only handles the first case + """ + if len(node.all_input_nodes) != 1: + raise AssertionError(f"Expected 1 input node, got {len(node.all_input_nodes)}") + source = node.all_input_nodes[0] + source_outputs = node_name_to_values[source.name] + if not isinstance(source_outputs, Sequence): + raise AssertionError( + f"Expected {source.name} to output sequence, got {node_name_to_values[source.name]}" + ) + index = typing.cast(int, node.args[1]) + value = source_outputs[index] + # Save the getitem value to the values mapping to in case + # it is one of the graph outputs + node_name_to_values[node.name] = value + # Rename the name of value with the getitem name. + value.name = node.name + return value + + +def _handle_call_function_node( + graph_like: ir.Graph | ir.Function, + node: torch.fx.Node, + node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]], +) -> None: + """Handle a call_function node. + + Args: + graph: The ONNX graph at construction. + node: The FX node to translate. + node_name_to_values: A mapping of FX node names to their produced ir.Value. + """ + if node.target is operator.getitem: + _handle_getitem_node(node, node_name_to_values) + # Add op to the graph + op = str(node.target) + fx_inputs, attributes, input_names, output_names = _get_inputs_and_attributes(node) + inputs: list[ir.Value | None] = [] + for i, input_ in enumerate(fx_inputs): + if input_ is None: + inputs.append(None) + elif hasattr(input_, "name"): + if isinstance(input_, torch.fx.Node) and input_.target is operator.getitem: + actual_input = _handle_getitem_node(input_, node_name_to_values) + inputs.append(actual_input) + else: + value = node_name_to_values[input_.name] + if isinstance(value, Sequence): + raise AssertionError(f"Unexpected sequence value for {input_.name}") + inputs.append(value) + else: + attributes[f"arg_{i}"] = input_ + + outputs = [ir.Value(name=name) for name in output_names] + if len(outputs) > 1: + _set_shape_types(outputs, node.meta["val"], complex_to_float=False) + node_name_to_values[node.name] = outputs + else: + _set_shape_type(outputs[0], node.meta["val"], complex_to_float=False) + node_name_to_values[node.name] = outputs[0] + ir_node = ir.Node( + "pkg.torch.ops", + op, + inputs, + attributes=ir_convenience.convert_attributes(attributes), + outputs=outputs, + name=node.name, + ) + ir_node.meta["node"] = node + ir_node.metadata_props["pkg.torch.onnx.input_names"] = repr(input_names) + # Record the nn.Module stack for the node + _set_node_metadata(node, ir_node) + + graph_like.append(ir_node) + + +def _convert_fx_arg_to_onnx_arg( + arg, + node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]], + node_name_to_local_functions: dict[str, ir.Function], +) -> Any: + """Convert an FX argument to an ONNX compatible argument. + + This function + - Converts a torch dtype to an integer + - Converts a torch device/memory_format/layout to a string + - Converts a torch.fx.Node to an ir.Value + - Converts a sequence of torch.fx.Node to a sequence of ir.Value + - Converts a get_attr node to an ir.Function + """ + if arg is None: + # None arguments are not modified because when the arg is an ONNX input + # we need to preserve the None value; when the arg is an ONNX attribute, + # we want to drop the value. + # The actual dropping of a None attribute value is done by OpRecorder + return None + if hasattr(arg, "name"): + if isinstance(arg, torch.fx.Node) and arg.target is operator.getitem: + source = arg.all_input_nodes[0] + source_outputs = node_name_to_values[source.name] + if isinstance(source_outputs, Sequence): + # If the node is getting an input from another node, get the actual value the node is retrieving + return _handle_getitem_node(arg, node_name_to_values) + else: + # `source_outputs` is a sequence(tensor()) value and we need to + # use SequenceAt to get the value. This is handled by torchlib + pass + if isinstance(arg, torch.fx.Node) and arg.op == "get_attr": + return node_name_to_local_functions[arg.name] + # If the input is a node, get the value from the mapping + return node_name_to_values[arg.name] + if isinstance(arg, (list, tuple)): + return [ + _convert_fx_arg_to_onnx_arg( + elem, node_name_to_values, node_name_to_local_functions + ) + for elem in arg + ] + if isinstance(arg, (torch.device, torch.memory_format, torch.layout)): + return str(arg) + if isinstance(arg, torch.dtype): + return torch_dtype_to_onnx_dtype(arg) + # Maybe a Python value + return arg + + +def _get_onnxscript_opset(opset_version: int) -> onnxscript.values.Opset: + return onnxscript.values.Opset("", opset_version) + + +def _is_onnx_op(op: Any) -> bool: + """Whether the op overload is an ONNX custom op implemented with PyTorch.""" + if not isinstance(op, torch._ops.OpOverload): + return False + return op.name().startswith("onnx::") + + +def _parse_onnx_op(op: torch._ops.OpOverload) -> tuple[str, int]: + """Parse the ONNX custom op overload name to get the op type and opset version.""" + name = op.name()[len("onnx::") :] + name, _, opset = name.partition(".opset") + return name, int(opset) + + +def _handle_call_function_node_with_lowering( + model: ir.Model, + node: torch.fx.Node, + node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]], + *, + graph_like: ir.Graph | ir.Function, + constant_farm: dict[Any, ir.Value], + registry: _registration.ONNXRegistry, + opset: onnxscript.values.Opset, + node_name_to_local_functions: dict[str, ir.Function], +) -> None: + """Translate a call_function node to an ONNX node. + + Args: + model: The ONNX model at construction. + node: The FX node to translate. + node_name_to_values: A mapping of FX node names to their produced ONNX ``Value``. + graph_like: The current ONNX graph at construction. + Must add nodes to this graph because it can be a subgraph that is currently being constructed. + constant_farm: A mapping of constant values to existing ONNX ``Value``s. + registry: The registry of all aten to ONNX decomposition functions. + opset: The ONNX Script opset object for constructing ONNX nodes. + node_name_to_local_functions: A mapping of subgraph names to the corresponding ONNX functions. + """ + if node.target is operator.getitem: + source = node.all_input_nodes[0] + source_outputs = node_name_to_values[source.name] + if isinstance(source_outputs, Sequence): + _handle_getitem_node(node, node_name_to_values) + return + else: + # `source_outputs` is a sequence(tensor()) value and we need to + # use SequenceAt to get the value. This is handled by torchlib + pass + + # Map FX inputs to ONNX inputs and fill optional inputs. + # torch_args and torch_kwargs are for op-level validation + fx_args = node.args + fx_kwargs = node.kwargs + + # Replace the input FX nodes with ONNX values + onnx_args = [ + _convert_fx_arg_to_onnx_arg( + input_, node_name_to_values, node_name_to_local_functions + ) + for input_ in fx_args + ] + + onnx_kwargs = {} + for key, value in fx_kwargs.items(): + onnx_kwargs[key] = _convert_fx_arg_to_onnx_arg( + value, node_name_to_values, node_name_to_local_functions + ) + if key == "dtype" and onnx_kwargs[key] is None: + # Set dtype to -1 if it is None + # TODO(justinchuby): Maybe keep it as None? + onnx_kwargs[key] = -1 + + if _is_onnx_op(node.target): + # Handle torch.ops.onnx.* ops. These ops can be directly added to the graph + op_type, opset_version = _parse_onnx_op(node.target) # type: ignore[arg-type] + # If final inputs are None, strip them from the node inputs + for input_ in reversed(onnx_args): + if input_ is not None: + break + onnx_args.pop() + onnx_node = ir.Node( + "", + op_type, + onnx_args, + ir.convenience.convert_attributes(onnx_kwargs), + name=node.name, + num_outputs=len(node.target._schema.returns), # type: ignore[union-attr] + version=opset_version, + ) + # Store the single node in a list to be consistent with the rest of the code for further processing + onnx_nodes = [onnx_node] + if len(onnx_node.outputs) == 1: + outputs = onnx_node.outputs[0] + else: + outputs = onnx_node.outputs # type: ignore[assignment] + else: + # Find the matching ONNX overload for the node + # TODO: Log the message here to expose false positives + onnx_function, message = _dispatching.dispatch(node, registry) + + if onnx_function is None: + raise _errors.DispatchError( + f"No ONNX function found for {node.target!r}. Failure message: {message}" + ) + + with onnxscript.evaluator.default_as( + tracer := _building.OpRecorder(opset, constant_farm) + ): + global current_tracer + current_tracer = tracer + try: + outputs = onnx_function(*onnx_args, **onnx_kwargs) + except Exception as e: + raise _errors.GraphConstructionError( + f"Error when calling function '{onnx_function}' with args '{onnx_args}' and kwargs '{onnx_kwargs}'" + ) from e + finally: + current_tracer = None + + # Add the defined functions to the model + for identifier, onnxscript_function in tracer.functions.items(): + if identifier in model.functions: + continue + if isinstance(onnxscript_function, ir.Function): + ir_function = onnxscript_function + else: + # TODO: Get IR function directly when onnxscript is updated + proto = onnxscript_function.to_function_proto() + ir_function = ir.serde.deserialize_function(proto) + model.functions[identifier] = ir_function + # Opset imports are added to the model in the final add_opset_imports pass + + onnx_nodes = tracer.nodes + del tracer # tracer is no longer needed + + # NOTE: Instead of using the output names from node.target._schema, + # we always use the index if there are more than one outputs so the + # names can be programmatically reconstructed. This is useful for + # comparing values from the ONNX graph with those from the FX graph. + # + # When there are multiple outputs, the output names will be + # node_name__0, node_name__1, etc. + if isinstance(outputs, Sequence): + _set_shape_types(outputs, node.meta["val"], complex_to_float=True) + node_name_to_values[node.name] = outputs + for i, output in enumerate(outputs): + output.name = f"{node.name}__{i}" + # Set the name of the producing node using the value name for correspondence + producer = output.producer() + if producer is not None: + producer.name = f"node_{output.name}" + else: + _set_shape_type(outputs, node.meta["val"], complex_to_float=True) + node_name_to_values[node.name] = outputs + outputs.name = node.name + producer = outputs.producer() + if producer is not None: + producer.name = f"node_{outputs.name}" + + for ir_node in onnx_nodes: + ir_node.meta["node"] = node + # Record the nn.Module stack for the node + _set_node_metadata(node, ir_node) + + # Add the traced nodes to the current graph + # Must add nodes to this graph, not model.graph, because it can be a subgraph that is currently being constructed + graph_like.extend(onnx_nodes) + + +def _handle_placeholder_node( + node: torch.fx.Node, + node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]], + *, + graph_like: ir.Graph | ir.Function, + lower: str, + opset: onnxscript.values.Opset, +) -> None: + # Placeholder nodes are user inputs + # We need to create a new tensor for each user input + # and add it to the graph's inputs + name = node.name + input_ = _tensors.SymbolicTensor(opset, name=name) + input_.meta["node"] = node + _set_shape_type(input_, node.meta["val"], complex_to_float=lower != "none") + node_name_to_values[name] = input_ + # The inputs should be add to the graph here + graph_like.inputs.append(input_) + + +def _handle_get_attr_node( + node: torch.fx.Node, + *, + owned_graphs: Mapping[str, ir.Function], + node_name_to_local_functions: dict[str, ir.Function], +) -> None: + """Handle a get_attr node by assigning the corresponding ONNX function to the node name. + + An example ExportedProgram that has uses get_attr nodes is: + + ExportedProgram: + class GraphModule(torch.nn.Module): + def forward(self, arg0_1: "f32[5]"): + true_graph_0 = self.true_graph_0 # get_attr + false_graph_0 = self.false_graph_0 # get_attr + conditional = torch.ops.higher_order.cond(False, true_graph_0, false_graph_0, [arg0_1]); true_graph_0 = false_graph_0 = arg0_1 = None + getitem: "f32[5]" = conditional[0]; conditional = None + return (getitem,) + + class (torch.nn.Module): + def forward(self, arg0_1: "f32[5]"): + cos: "f32[5]" = torch.ops.aten.cos.default(arg0_1); arg0_1 = None + return (cos,) + + class (torch.nn.Module): + def forward(self, arg0_1: "f32[5]"): + sin: "f32[5]" = torch.ops.aten.sin.default(arg0_1); arg0_1 = None + return (sin,) + + Args: + node: The FX node to translate. + owned_graphs: A mapping of subgraph names to the corresponding ONNX functions. + node_name_to_local_functions: A mapping of local function names to their corresponding ONNX functions. + """ + if not isinstance(node.target, str): + logger.warning( + "Expected node.target for the node %s to be a string, but got '%s'. There may be an internal error.", + node, + type(node.target), + ) + return + function = owned_graphs[node.target] + node_name_to_local_functions[node.name] = function + + +def _handle_output_node( + node: torch.fx.Node, + node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]], + graph_like: ir.Graph | ir.Function, +) -> None: + """Handle an output node by adding the output to the graph's outputs. + + Args: + node: The FX node to translate. + node_name_to_values: A mapping of FX node names to their produced ONNX ``Value``. + graph_like: The ONNX graph at construction. + """ + if not isinstance(node.args[0], Sequence): + output_nodes = (node.args[0],) + else: + # node.args[0] can be a tuple with more than one elements. This happens when, + # for example, a subgraph has multiple outputs. We flatten them all as ONNX graph outputs + output_nodes = node.args[0] + for output in output_nodes: + if output is None: + logger.warning( + "Output node %s has None output. The output is ignored in the exported graph. Please ensure the graph output order is expected", + node.name, + ) + continue + output_value_name = output.name # type: ignore[union-attr] + if not isinstance(output_value_name, str): + raise AssertionError(f"Bug: Expected {output_value_name!r} to be a string") + values = node_name_to_values[output_value_name] + if isinstance(values, Sequence): + graph_like.outputs.extend(values) + return + graph_like.outputs.append(values) + + +def _translate_fx_graph( + fx_graph: torch.fx.Graph, + model: ir.Model, + *, + graph_like: ir.Graph | ir.Function, + owned_graphs: Mapping[str, ir.Function], + lower: Literal["at_conversion", "none"], + registry: _registration.ONNXRegistry, +) -> dict[str, ir.Value | Sequence[ir.Value]]: + """Translate a submodule to an ONNX function. + + Any functions used by the traced functions will be added to the model. + + Args: + fx_graph: The FX graph module to translate. + model: The ONNX model at construction. + current_scope: The current name scope of the submodule, excluding the current module name. + E.g. "true_graph_0.false_graph_0". + graph_name: The name of the submodule. E.g. "true_graph_0". + graph: The ONNX graph at construction. + owned_graphs: The subgraphs owned by the current graph. + lower: The lowering strategy to use. + registry: The registry of all aten to ONNX decomposition functions. + + Returns: + A mapping of FX node names to their produced ONNX ``Value``. + """ + node_name_to_values: dict[str, ir.Value | Sequence[ir.Value]] = {} + # The reason we need node_name_to_local_functions in addition to owned_graphs + # is because the get_attr nodes may assign a different name than the GraphModule name + # to the subgraph. This is not typical but is valid Python. + node_name_to_local_functions: dict[str, ir.Function] = {} + constant_farm: dict[Any, ir.Value] = {} + opset = _get_onnxscript_opset(registry.opset_version) + + for node in fx_graph.nodes: + logger.debug( + "%s", (node.name, node.args, node.target, node.op, node.type, node.kwargs) + ) + try: + if node.op == "placeholder": + _handle_placeholder_node( + node, + node_name_to_values, + graph_like=graph_like, + lower=lower, + opset=opset, + ) + elif node.op == "call_function": + if lower == "at_conversion": + _handle_call_function_node_with_lowering( + model, + node, + node_name_to_values, + graph_like=graph_like, + constant_farm=constant_farm, + registry=registry, + opset=opset, + node_name_to_local_functions=node_name_to_local_functions, + ) + else: + # No lowering + _handle_call_function_node(graph_like, node, node_name_to_values) + elif node.op == "get_attr": + _handle_get_attr_node( + node, + owned_graphs=owned_graphs, + node_name_to_local_functions=node_name_to_local_functions, + ) + elif node.op == "output": + _handle_output_node( + node, + node_name_to_values, + graph_like=graph_like, + ) + except Exception as e: + raise _errors.ConversionError( + f"Error when translating node {node.format_node()}. See the stack trace for more information." + ) from e + return node_name_to_values + + +def _get_inputs_and_attributes( + node: torch.fx.Node, +) -> tuple[list[torch.fx.Node | None], dict[str, Any], list[str], list[str]]: + """Find and Fill in the not provided kwargs with default values. + + Returns: + (inputs, attributes, input_names, output_names) + """ + if inspect.isbuiltin(node.target) or isinstance(node.target, str): + inputs = list(node.args) + return inputs, {}, [], [node.name] # type: ignore[return-value] + + # The target should be an ATen operator now + if not hasattr(node.target, "_schema"): + raise AssertionError( + f"The target should be an ATen operator now, but node target {node.target} has no schema" + ) + node_schema: torch.FunctionSchema = node.target._schema + + # This function assumes the order of arguments in FX op is the + # same as the order of arguments in TorchScript op. + inputs: list[Any] = [] # type: ignore[no-redef] + input_names: list[str] = [] + attributes: dict[str, Any] = {} + + if inspect.isbuiltin(node.target): + inputs = list(node.args) + else: + for arg, schema_arg in zip(node.args, node_schema.arguments): + if arg is None or isinstance(arg, torch.fx.Node): + inputs.append(arg) + input_names.append(schema_arg.name) + elif isinstance(arg, Sequence) and all( + elem is None or isinstance(elem, torch.fx.Node) for elem in arg + ): + inputs.extend(arg) + input_names.extend([schema_arg.name] * len(arg)) + elif isinstance(arg, torch.device): + attributes[schema_arg.name] = str(arg) + elif isinstance(arg, torch.dtype): + attributes[schema_arg.name] = torch_dtype_to_onnx_dtype(arg) + else: + attributes[schema_arg.name] = arg + for schema_arg in node_schema.arguments: + if schema_arg.name not in node.kwargs: + continue + kwarg = node.kwargs[schema_arg.name] + if schema_arg.name in { + "layout", + "device", + "requires_grad", + "memory_format", + "implicit", + } or isinstance(kwarg, torch.device): + attr = str(kwarg) + elif isinstance(kwarg, torch.dtype): + attr = torch_dtype_to_onnx_dtype(kwarg) # type: ignore[assignment] + else: + attr = kwarg # type: ignore[assignment] + + attributes[schema_arg.name] = attr + + output_names = [f"{node.name}_{output.name}" for output in node_schema.returns] + + return inputs, attributes, input_names, output_names # type: ignore[return-value] + + +def _maybe_start_profiler(should_profile: bool) -> Any: + if should_profile: + import pyinstrument # type: ignore[import-not-found] + + profiler = pyinstrument.Profiler(async_mode="disabled") + profiler.start() + return profiler + return None + + +def _maybe_stop_profiler_and_get_result(profiler) -> str | None: + if profiler is None: + return None + profiler.stop() + return profiler.output_text(unicode=True) + + +def _format_exception(e: Exception) -> str: + """Format the full traceback as Python would show it.""" + return "\n".join(traceback.format_exception(type(e), e, e.__traceback__)) + + +def _summarize_exception_stack(e: BaseException) -> str: + """Format the exception stack by showing the text of each exception.""" + causes = [e] + while e.__cause__ is not None: + causes.append(e.__cause__) + e = e.__cause__ + return ( + "\n\n## Exception summary\n\n" + + "⬆️\n".join([f"{type(e)}: {e}\n" for e in reversed(causes)]) + + "\n(Refer to the full stack trace above for more information.)" + ) + + +def _format_exceptions_for_all_strategies( + results: list[_capture_strategies.Result], +) -> str: + """Format all the exceptions from the capture strategies.""" + return "\n".join( + [ + f"# ⚠️ Errors from strategy '{result.strategy}': -----------------------\n\n" + f"{_format_exception(result.exception)}\n" + for result in results + if result.exception is not None + ] + ) + + +def exported_program_to_ir( + exported_program: torch.export.ExportedProgram, + *, + registry: _registration.ONNXRegistry | None = None, + lower: Literal["at_conversion", "none"] = "at_conversion", +) -> ir.Model: + """Convert an exported program to an ONNX IR model. + + Reference: + - ExportedProgram spec: https://pytorch.org/docs/stable/export.ir_spec.html + + Args: + exported_program: The exported program to convert. + lower: Whether to lower the graph to core ONNX operators. + at_conversion: Lower when translating the FX graph to ONNX IR. + none: Do not lower the graph. + registry: The registry of all ONNX Script decomposition. + """ + if registry is None: + registry = _registration.ONNXRegistry.from_torchlib() + if lower != "none": + exported_program = _prepare_exported_program_for_export( + exported_program, registry=registry + ) + return _exported_program_to_onnx_program( + exported_program, registry=registry, lower=lower + ).model + + +def _prepare_exported_program_for_export( + exported_program: torch.export.ExportedProgram, + *, + registry: _registration.ONNXRegistry, +) -> torch.export.ExportedProgram: + """Decompose and apply pre-export transformations to the exported program.""" + + with ( + # Support the dynamism with 0/1 input dim + torch.fx.experimental._config.patch(backed_size_oblivious=True), # type: ignore[attr-defined] + ): + # Decompose the graph given the implemented torch ops in ONNX + exported_program = _fx_passes.decompose_with_registry( + exported_program, registry + ) + + graph_module = exported_program.graph_module + # Include explicit type promotion nodes + _fx_passes.insert_type_promotion_nodes(graph_module) + graph_module = _fx_passes.remove_assertion_nodes(graph_module) + # Reassign the graph module to save some runtime. + exported_program._graph_module = graph_module + return exported_program + + +def _get_scope_name(scoped_name: str) -> tuple[str, str]: + """Get the scope and name of a node. + + Examples:: + >>> _get_scope_name('') + ('', '') + >>> _get_scope_name('true_graph') + ('', 'true_graph') + >>> _get_scope_name('true_graph.false_graph') + ('true_graph', 'false_graph') + >>> _get_scope_name('true_graph.false_graph.some_graph') + ('true_graph.false_graph', 'some_graph') + + Args: + scoped_name: The scoped name of the node. + + Returns: + (scope, name) + """ + if "." in scoped_name: + scope, name = scoped_name.rsplit(".", 1) + else: + scope, name = "", scoped_name + return scope, name + + +def _exported_program_to_onnx_program( + exported_program: torch.export.ExportedProgram, + *, + registry: _registration.ONNXRegistry, + lower: Literal["at_conversion", "none"] = "at_conversion", +) -> _onnx_program.ONNXProgram: + """Convert an exported program to an ONNX Program. + + The exported_program field in the returned ONNXProgram is one that is after + decompositions have been applied. + + Reference: + - ExportedProgram spec: https://pytorch.org/docs/stable/export.ir_spec.html + + Args: + exported_program: The exported program to convert. The exported program + should be the one that is after decompositions have been applied. + lower: Whether to lower the graph to core ONNX operators. + at_conversion: Lower when translating the FX graph to ONNX IR. + none: Do not lower the graph. + registry: The registry of all ONNX Script decomposition. + """ + model = ir.Model( + graph=ir.Graph( + [], + [], + nodes=[], + # Opset imports are added to the model in the final add_opset_imports pass + name="main_graph", + metadata_props={ + "pkg.torch.export.ExportedProgram.graph_signature": str( + exported_program.graph_signature + ), + "pkg.torch.export.ExportedProgram.range_constraints": str( + exported_program.range_constraints + ), + }, + ), + ir_version=_constants.ONNX_IR_VERSION, + producer_name="pytorch", + producer_version=torch.__version__, + ) + + # A dictionary storing the translated subgraphs as ONNX functions made available to outer graphs + # {: {: }} + scoped_subgraphs: dict[str, dict[str, ir.Function]] = {} + values = None + + # 1. Translate all nodes in all subgraphs and the main graph + # Create a dictionary of values for the main graph for step 2-3 to add inputs and outputs + module: torch.fx.GraphModule + # Reverse the order of the modules so that the innermost module is processed first + # and made available to the outer module + for name, module in reversed( + tuple(exported_program.graph_module.named_modules(remove_duplicate=False)) + ): + # Obtain the graphs (previously built) owned by the current module + owned_graphs = scoped_subgraphs.setdefault(name, {}) + fx_graph = module.graph + + graph_like: ir.Graph | ir.Function + if name == "": + # Root graph + graph_like = model.graph + else: + function_name = name.replace(".", "__") + # Inputs and outputs will be created within _translate_fx_graph + func = ir.Function( + domain=_constants.LOCAL_FUNCTION_DOMAIN, + name=function_name, + graph=ir.Graph((), (), nodes=()), + attributes=(), + ) + # Make this function available to the outer graph + scope, subgraph_name = _get_scope_name(name) + scoped_subgraphs.setdefault(scope, {})[subgraph_name] = func + model.functions[func.identifier()] = func + graph_like = func + + values = _translate_fx_graph( + fx_graph, + model, + graph_like=graph_like, + owned_graphs=owned_graphs, + lower=lower, + registry=registry, + ) + + if name != "": + raise AssertionError("The last module processed should be the root module") + if values is None: + raise AssertionError("values must be non-None") + + # Clear the input/output of the main graph and add them back in step 2-3 + # using the more accurate graph signature + model.graph.inputs.clear() + model.graph.outputs.clear() + + # 2. Add user inputs and all parameters/buffers to the graph. + # Since the node names and the tensor names are different, we need to rename + # the nodes to match the tensor names later. For now we will just use the node names. + user_inputs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == graph_signature.InputKind.USER_INPUT + ] + non_user_inputs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind != graph_signature.InputKind.USER_INPUT + ] + + for spec in itertools.chain(user_inputs, non_user_inputs): + # Put the user inputs first and then the parameters/buffers + if isinstance(spec.arg, graph_signature.ConstantArgument): + logger.debug("Skipping constant argument %s", spec.arg) + continue + value_name = spec.arg.name + input_kind = spec.kind + persistent = spec.persistent + value = values[value_name] + + if isinstance(value, Sequence): + raise AssertionError( + f"Input '{value_name}' should not be a sequence. This is unexpected." + ) + + value.metadata_props["pkg.torch.export.graph_signature.InputSpec.kind"] = ( + input_kind.name + ) + value.metadata_props[ + "pkg.torch.export.graph_signature.InputSpec.persistent" + ] = str(persistent) + + if input_kind == graph_signature.InputKind.USER_INPUT: + # Add only user inputs to the graph + # Subsequent passes can decide if they want to add initializers as inputs + model.graph.inputs.append(value) + else: + model.graph.initializers[value_name] = value + + # 3. Add user outputs to the graph and assign metadata to all outputs + user_outputs = [ + spec + for spec in exported_program.graph_signature.output_specs + if spec.kind == graph_signature.OutputKind.USER_OUTPUT + ] + non_user_outputs = [ + spec + for spec in exported_program.graph_signature.output_specs + if spec.kind != graph_signature.OutputKind.USER_OUTPUT + ] + for spec in itertools.chain(user_outputs, non_user_outputs): + if isinstance(spec.arg, graph_signature.ConstantArgument): + logger.warning("Skipping constant argument %s", spec.arg) + continue + value_name = spec.arg.name + output_kind = spec.kind + value = values[value_name] + + if not isinstance(value, (ir.Value, Sequence)): + raise TypeError( + f"Output '{value_name}' should be an ir.Value. Actual type is '{type(value)}': {value!r}. " + "This may be due to an incorrect implementation of the ONNX function that produced this output." + ) + + # The output value may be a sequence, meaning the operator has multiple outputs + _values = (value,) if not isinstance(value, Sequence) else value + + if len(_values) > 1: + logger.warning( + "Model output '%s' has multiple values: %s (output spec: %s). Please make sure this is expected.", + value_name, + _values, + spec, + ) + + for value in _values: + value.metadata_props["pkg.torch.export.graph_signature.OutputSpec.kind"] = ( + output_kind.name + ) + if output_kind == graph_signature.OutputKind.USER_OUTPUT: + model.graph.outputs.append(value) + + # 4. Rename the initializers to match the tensor names + for name, param_name in itertools.chain( + exported_program.graph_signature.inputs_to_parameters.items(), + exported_program.graph_signature.inputs_to_buffers.items(), + exported_program.graph_signature.inputs_to_lifted_tensor_constants.items(), + ): + initializer = model.graph.initializers.pop(name) + initializer.name = param_name + # Record the original name so users can search the metadata and correspond + # with the FX graph + initializer.metadata_props["pkg.torch.onnx.original_node_name"] = name + model.graph.initializers[param_name] = initializer + + # 5. Add initializers to the graph + # ExportedProgram stores parameters and buffers in state_dict, + # but non_persistent_buffers and lifted_tensor_constants are not there + # so we need to get them from the name_* apis. + for name, torch_tensor in itertools.chain( + exported_program.named_parameters(), + # pyrefly: ignore [bad-argument-type] + exported_program.named_buffers(), + exported_program.constants.items(), + ): + initializer = model.graph.initializers.get(name) # type: ignore[assignment] + if initializer is None: + logger.warning("Tensor '%s' is not one of the initializers", name) + continue + if not isinstance(torch_tensor, torch.Tensor): + raise NotImplementedError( + f"Tensor '{name}' should be a torch.Tensor. Actual type is '{type(torch_tensor)}': {torch_tensor!r}. " + "This is unexpected and not yet supported." + ) + + # Turn complex tensors into float tensors when converting to ONNX + complex_to_float = lower != "none" + if complex_to_float: + if torch_tensor.dtype.is_complex: + torch_tensor = torch.view_as_real(torch_tensor) + + ir_tensor = TorchTensor(torch_tensor, name=name) + initializer.const_value = ir_tensor + _set_shape_type( + initializer, + torch_tensor, + complex_to_float=complex_to_float, + ) + + # TODO: Decide if we should keep mutated buffers as inputs/outputs + + # Collect and add opset imports to the model + _ir_passes.add_opset_imports(model) + + return _onnx_program.ONNXProgram(model, exported_program) + + +def _verbose_printer(verbose: bool | None) -> Callable[..., None]: + """Prints messages based on `verbose`.""" + if verbose is False: + return lambda *_, **__: None + + return lambda *args, **kwargs: print("[torch.onnx]", *args, **kwargs) + + +@_flags.set_onnx_exporting_flag +def export( + model: torch.nn.Module + | torch.export.ExportedProgram + | torch.fx.GraphModule + | torch.jit.ScriptModule + | torch.jit.ScriptFunction, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + *, + registry: _registration.ONNXRegistry | None = None, + dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None = None, + input_names: Sequence[str] | None = None, + output_names: Sequence[str] | None = None, + report: bool = False, + verify: bool = False, + profile: bool = False, + dump_exported_program: bool = False, + artifacts_dir: str | os.PathLike = ".", + verbose: bool | None = None, + optimize: bool = True, + opset_version: int | None = None, +) -> _onnx_program.ONNXProgram: + """Export a PyTorch model to ONNXProgram. + + Args: + model: The model to export. This can be a PyTorch nn.Module or an ExportedProgram. + args: The arguments to pass to the model. + kwargs: The keyword arguments to pass to the model. + registry: The registry of all ONNX decompositions. + dynamic_shapes: Dynamic shapes in the graph. + input_names: If provided, rename the inputs. + output_names: If provided, rename the outputs. + report: Whether to generate an error report if the export fails. + verify: Whether to verify the ONNX model after exporting. + profile: Whether to profile the export process. When report is True, + the profile result will be saved in the report. Otherwise, the profile + result will be printed. + dump_exported_program: Whether to save the exported program to a file. + artifacts_dir: The directory to save the exported program and error reports. + verbose: Whether to print verbose messages. If None (default), some messages will be printed. + optimize: Whether to optimize the exported ONNX graph. + opset_version: The ONNX opset version to use. If None, use the default opset version + from the registry. + + Returns: + The ONNXProgram with the exported IR graph. + + Raises: + TorchExportError: If the export process fails with torch.export. + ConversionError: If the ExportedProgram to ONNX translation fails. + """ + # Set up the error reporting facilities + timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S-%f") + profiler = _maybe_start_profiler(profile) + + # Create the artifacts directory if it does not exist + artifacts_dir = pathlib.Path(artifacts_dir) + if report or profile or dump_exported_program: + artifacts_dir.mkdir(parents=True, exist_ok=True) + + verbose_print = _verbose_printer(verbose) + export_status = _reporting.ExportStatus() + failed_results: list[_capture_strategies.Result] = [] + + program: torch.export.ExportedProgram | None = None + capture_strategy: str | None = None + # Step 1: Export the model with torch.export.export if the model is not already an ExportedProgram + if isinstance(model, torch.export.ExportedProgram): + # We know the model is already exported program, so the args, kwargs, and dynamic_shapes + # are not used. + program = model + # torch.export.export has strict default to False + export_status.torch_export_non_strict = True + else: + # Convert an nn.Module to an ExportedProgram + # Try everything 🐰 (all paths for getting an ExportedProgram) + result: _capture_strategies.Result | None = None + for strategy_class in _capture_strategies.CAPTURE_STRATEGIES: + strategy = strategy_class( # type: ignore[abstract] + verbose=verbose is not False, # Treat None as verbose + dump=dump_exported_program, + artifacts_dir=artifacts_dir, + timestamp=timestamp, + ) + result = strategy(model, args, kwargs, dynamic_shapes=dynamic_shapes) + + # Record the status + if strategy_class is _capture_strategies.TorchExportNonStrictStrategy: + export_status.torch_export_non_strict = result.success + elif strategy_class is _capture_strategies.TorchExportStrictStrategy: + export_status.torch_export_strict = result.success + elif strategy_class is _capture_strategies.TorchExportDraftExportStrategy: + export_status.torch_export_draft_export = result.success + + if result.exception is not None: + failed_results.append(result) + if result.success: + if result.exported_program is None: + raise AssertionError("exported_program must be non-None on success") + program = result.exported_program + break + + if result is None: + raise AssertionError("result must be non-None") + capture_strategy = result.strategy + if result.exported_program is None: + # If all strategies fail, produce an error report and raise the first error + profile_result = _maybe_stop_profiler_and_get_result(profiler) + + if report: + report_path = artifacts_dir / _reporting.construct_report_file_name( + timestamp, export_status + ) + + try: + _reporting.create_torch_export_error_report( + report_path, + _format_exceptions_for_all_strategies(failed_results), + export_status=export_status, + profile_result=profile_result, + ) + except Exception as e_report: + verbose_print( + f"Failed to save error report due to an error: {e_report}" + ) + else: + report_path = None + + first_error = failed_results[0].exception + if first_error is None: + raise AssertionError("first_error must be non-None") + + # NOTE: We only throw the torch.export (first) exception because we want to + # focus on the torch.export.export error. Errors from other strategies like + # torch.jit.trace is due to the fallback and can be confusing to users. + # We save all errors in the error report. + raise _errors.TorchExportError( + _STEP_ONE_ERROR_MESSAGE + + ( + f"\nError report has been saved to '{report_path}'." + if report + else "" + ) + + _summarize_exception_stack(first_error) + ) from first_error + + if program is None: + raise AssertionError("program must be non-None") + + if dump_exported_program: + verbose_print("Dumping ExportedProgram because `dump_exported_program=True`...") + program_path = artifacts_dir / f"onnx_export_{timestamp}.pt2" + try: + torch.export.save(program, program_path) + except Exception as e: + verbose_print(f"Failed to save ExportedProgram due to an error: {e}") + else: + verbose_print(f"ExportedProgram has been saved to '{program_path}'.") + + # Step 2: Decompose the exported program and insert type promotion nodes + verbose_print("Run decompositions...") + + try: + # Build the ONNX function registry + if registry is None: + registry = _registration.ONNXRegistry.from_torchlib() + + # Process the exported program to run decompositions and type promotions etc. + decomposed_program = _prepare_exported_program_for_export( + program, registry=registry + ) + except Exception as e: + export_status.decomposition = False + verbose_print("Run decompositions... ❌") + profile_result = _maybe_stop_profiler_and_get_result(profiler) + + if report: + report_path = artifacts_dir / _reporting.construct_report_file_name( + timestamp, export_status + ) + + # Run the analysis to get the error report + try: + _reporting.create_onnx_export_report( + report_path, + f"{_format_exceptions_for_all_strategies(failed_results)}\n\n{_format_exception(e)}", + program, + export_status=export_status, + profile_result=profile_result, + registry=registry, + ) + except Exception: + logger.exception("Failed to save report due to an error.") + else: + report_path = None + + raise _errors.ConversionError( + _STEP_TWO_ERROR_MESSAGE + + (f"\nError report has been saved to '{report_path}'." if report else "") + + _summarize_exception_stack(e) + ) from e + else: + export_status.decomposition = True + verbose_print("Run decompositions... ✅") + + # Step 3: Translate the decomposed program to ONNX and produce ONNXProgram + verbose_print("Translate the graph into ONNX...") + if report or profile: + pre_decomp_unique_ops, post_decomp_unique_ops = _analysis.compare_ops( + program, decomposed_program + ) + else: + pre_decomp_unique_ops = None + post_decomp_unique_ops = None + + try: + # Convert the exported program to an ONNX model + onnx_program = _exported_program_to_onnx_program( + decomposed_program, registry=registry + ) + # Record the strategy used for getting the exported program for unit test assertions + onnx_program._capture_strategy = capture_strategy + + export_status.onnx_translation = True + verbose_print("Translate the graph into ONNX... ✅") + except Exception as e: + export_status.onnx_translation = False + verbose_print("Translate the graph into ONNX... ❌") + profile_result = _maybe_stop_profiler_and_get_result(profiler) + + if report: + report_path = artifacts_dir / _reporting.construct_report_file_name( + timestamp, export_status + ) + + try: + if pre_decomp_unique_ops is None: + raise AssertionError("pre_decomp_unique_ops must be non-None") + if post_decomp_unique_ops is None: + raise AssertionError("post_decomp_unique_ops must be non-None") + + # Run the analysis to get the error report + _reporting.create_onnx_export_report( + report_path, + f"{_format_exceptions_for_all_strategies(failed_results)}\n\n{_format_exception(e)}", + decomposed_program, + decomp_comparison=_reporting.format_decomp_comparison( + pre_decomp_unique_ops, post_decomp_unique_ops + ), + export_status=export_status, + profile_result=profile_result, + registry=registry, + ) + verbose_print(f"Export report has been saved to '{report_path}'.") + except Exception: + logger.exception("Failed to save report due to an error.") + else: + report_path = None + + raise _errors.ConversionError( + _STEP_THREE_ERROR_MESSAGE + + (f"\nError report has been saved to '{report_path}'." if report else "") + + _summarize_exception_stack(e) + ) from e + + profile_result = _maybe_stop_profiler_and_get_result(profiler) + + if onnx_program.exported_program is None: + raise AssertionError("exported_program must be non-None") + + # Converter opset version and optimize + if opset_version is not None: + onnx_program.model = onnxscript_apis.convert_version( + onnx_program.model, opset_version + ) + + if optimize: + verbose_print("Optimize the ONNX graph...") + onnx_program.optimize() + verbose_print("Optimize the ONNX graph... ✅") + + # Run the ONNX passes + if input_names: + _ir_passes.rename_inputs(onnx_program.model, input_names) + if output_names: + _ir_passes.rename_outputs(onnx_program.model, output_names) + + if not verify: + # Return if verification is not requested + if report: + try: + if pre_decomp_unique_ops is None: + raise AssertionError("pre_decomp_unique_ops must be non-None") + if post_decomp_unique_ops is None: + raise AssertionError("post_decomp_unique_ops must be non-None") + report_path = artifacts_dir / _reporting.construct_report_file_name( + timestamp, export_status + ) + _reporting.create_onnx_export_report( + report_path, + "No errors" + if not failed_results + else _format_exceptions_for_all_strategies(failed_results), + onnx_program.exported_program, + decomp_comparison=_reporting.format_decomp_comparison( + pre_decomp_unique_ops, post_decomp_unique_ops + ), + export_status=export_status, + profile_result=profile_result, + model=onnx_program.model, + registry=registry, + ) + verbose_print(f"Export report has been saved to '{report_path}'.") + except Exception: + logger.exception("Failed to save report due to an error.") + elif profile and profile_result is not None: + verbose_print("Profile result:") + verbose_print(profile_result) + return onnx_program + + # Step 4: (verify=True) Check the ONNX model with ONNX checker + try: + verbose_print("Check the ONNX model...") + onnxscript_apis.check_model(onnx_program.model) + export_status.onnx_checker = True + verbose_print("Check the ONNX model... ✅") + except Exception as e: + export_status.onnx_checker = False + verbose_print("Check the ONNX model... ❌") + if report: + try: + if pre_decomp_unique_ops is None: + raise AssertionError("pre_decomp_unique_ops must be non-None") + if post_decomp_unique_ops is None: + raise AssertionError("post_decomp_unique_ops must be non-None") + report_path = artifacts_dir / _reporting.construct_report_file_name( + timestamp, export_status + ) + _reporting.create_onnx_export_report( + report_path, + f"{_format_exceptions_for_all_strategies(failed_results)}\n\n{_format_exception(e)}", + onnx_program.exported_program, + decomp_comparison=_reporting.format_decomp_comparison( + pre_decomp_unique_ops, post_decomp_unique_ops + ), + export_status=export_status, + profile_result=profile_result, + model=onnx_program.model, + registry=registry, + ) + verbose_print(f"Export report has been saved to '{report_path}'.") + except Exception: + logger.exception("Failed to save report due to an error.") + logger.warning( + "Conversion successful but the ONNX model fails ONNX checker. " # noqa: G004 + "Please create an issue " + f"in the PyTorch GitHub repository against the {_BLUE}*onnx*{_END} component and " + "attach the full error stack as well as reproduction scripts. ", + exc_info=e, + ) + return onnx_program + + # Step 5: (verify=True) Execute the model with ONNX Runtime + try: + verbose_print("Execute the model with ONNX Runtime...") + verification_results = _verification.verify_onnx_program(onnx_program) + verbose_print("Execute the model with ONNX Runtime... ✅") + export_status.onnx_runtime = True + onnx_runtime_error_message = None + except Exception as e: + verbose_print("Execute the model with ONNX Runtime... ❌") + export_status.onnx_runtime = False + onnx_runtime_error_message = _format_exception(e) + verification_message = None + + else: + # Step 6: (verify=True) Validate the output values + verbose_print("Verify output accuracy...") + export_status.output_accuracy = True + for verification_result in verification_results: + # TODO(justinchuby): The threshold is arbitrary right now + if verification_result.max_abs_diff >= 5e-3: + logger.warning( + "Output '%s' has a large absolute difference of %f. ", + verification_result.name, + verification_result.max_abs_diff, + ) + export_status.output_accuracy = False + if verification_result.max_rel_diff >= 1e-1: + logger.warning( + "Output '%s' has a large relative difference of %f. ", + verification_result.name, + verification_result.max_rel_diff, + ) + export_status.output_accuracy = False + if export_status.output_accuracy: + verbose_print("Verify output accuracy... ✅") + else: + verbose_print("Verify output accuracy... ❌") + verification_message = _reporting.format_verification_infos( + verification_results + ) + + if report: + try: + if pre_decomp_unique_ops is None: + raise AssertionError("pre_decomp_unique_ops must be non-None") + if post_decomp_unique_ops is None: + raise AssertionError("post_decomp_unique_ops must be non-None") + + traceback_lines = [] + if failed_results: + traceback_lines.append( + _format_exceptions_for_all_strategies(failed_results) + ) + if onnx_runtime_error_message: + traceback_lines.append("# ⚠️ ONNX Runtime error -----------------------") + traceback_lines.append(onnx_runtime_error_message) + if not traceback_lines: + traceback_lines.append("No errors") + + report_path = artifacts_dir / _reporting.construct_report_file_name( + timestamp, export_status + ) + _reporting.create_onnx_export_report( + report_path, + "\n\n".join(traceback_lines), + onnx_program.exported_program, + profile_result=profile_result, + export_status=export_status, + decomp_comparison=_reporting.format_decomp_comparison( + pre_decomp_unique_ops, post_decomp_unique_ops + ), + model=onnx_program.model, + registry=registry, + verification_result=verification_message, + ) + verbose_print(f"Export report has been saved to '{report_path}'.") + except Exception: + logger.exception("Failed to save report due to an error.") + + # Release the inference session created during verification + onnx_program.release() + return onnx_program diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_decomp.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_decomp.py new file mode 100644 index 0000000000000000000000000000000000000000..1db085a1fcbf13cdb7b4c0c5968b75b357ab866b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_decomp.py @@ -0,0 +1,77 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import torch +import torch._ops + + +if TYPE_CHECKING: + from collections.abc import Callable + + from torch.onnx._internal.exporter import _registration + + +def get_onnx_implemented_overloads( + registry: _registration.ONNXRegistry, +) -> list[_registration.TorchOp]: + """ + Creates a set of OperatorBase and Callable objects that represent ONNX-supported PyTorch operations. + + Args: + registry: The ONNX registry for PyTorch. + + Returns: + A collection of OperatorBase and Callable objects representing ONNX-supported PyTorch operations. + """ + registered_ops: list[_registration.TorchOp] = [] + for onnx_decomp_meta in registry.functions.values(): + if len(onnx_decomp_meta) == 0: + raise AssertionError("onnx_decomp_meta must not be empty") + # Different OnnxDecompMeta for the same TorchOp should + # have the same fx_target. + fx_target = onnx_decomp_meta[0].fx_target + registered_ops.append(fx_target) + return registered_ops + + +def create_onnx_friendly_decomposition_table( + onnx_registered_ops: set[_registration.TorchOp], +) -> dict[_registration.TorchOp, Callable]: + """ + This function creates a dictionary of op overloads and their decomposition functions + for ops that do not have ONNX symbolic functions. If an op already has an ONNX symbolic function, + its decomposition function is excluded from the table. The decomposition table is a subset of PyTorch's + built-in aten-to-aten decomposition. + + Args: + onnx_registered_ops: All ops that have an ONNX decomposition implemented. + + Returns: + Dict[torch._ops.OperatorBase, Callable]: A dictionary that maps op overloads to their corresponding + decomposition functions. + """ + decomposition_table: dict[_registration.TorchOp, Callable] = {} + + for op_overload, decomp_fn in itertools.chain( + torch.export.default_decompositions().items(), # type: ignore[attr-defined] + torch._decomp.decomposition_table.items(), # type: ignore[attr-defined] + ): + # Skip decomposition for op_overload as long as that op_overload has a corresponding ONNX + # symbolic function. + # NOTE: Do not skip torch._refs decomps. They are fine because otherwise the model is + # not exportable anyways. + if op_overload in onnx_registered_ops: + continue + # If it is HOP, we filter those out as well. + if not hasattr(op_overload, "_schema"): + continue + # NOTE: torch._decomp.decomposition_table covers more ops + # than torch.export.default_decompositions, but the latter is + # more critical to torch.onnx.export. + if op_overload in decomposition_table: + continue + decomposition_table[op_overload] = decomp_fn + return decomposition_table diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_dispatching.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_dispatching.py new file mode 100644 index 0000000000000000000000000000000000000000..7ce267ed0f265262abab785cc0d094fecfdf7194 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_dispatching.py @@ -0,0 +1,58 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.fx + + +if TYPE_CHECKING: + from collections.abc import Callable + + from torch.onnx._internal.exporter import _registration + + +def _arg_has_complex_dtype(arg) -> bool: + """Check if the node has complex dtype recursively.""" + if ( + isinstance(arg, torch.fx.Node) + and "val" in arg.meta + and isinstance(arg.meta["val"], torch.Tensor) + and torch.is_complex(arg.meta["val"]) + ): + return True + elif isinstance(arg, list): + return any(_arg_has_complex_dtype(item) for item in arg) + return False + + +def dispatch( + node: torch.fx.Node, registry: _registration.ONNXRegistry +) -> tuple[Callable | None, str]: + """Dispatch a node to an ONNX function based on the node's target and the ONNX registry. + + Args: + node: The node to dispatch. + registry: The ONNX registry to use for dispatching. + + Returns: + A tuple containing the matched ONNX function and a string describing the reason for failure or success. + """ + decomp_metas = registry.get_decomps(node.target) # type: ignore[arg-type] + # Determine if the node has complex inputs. + is_complex = any(_arg_has_complex_dtype(arg) for arg in node.args) or any( + _arg_has_complex_dtype(arg) for arg in node.kwargs.values() + ) + if is_complex: + decomp_metas = [decomp for decomp in decomp_metas if decomp.is_complex] + if not decomp_metas: + return None, "No decompositions registered for the complex-valued input" + else: + decomp_metas = [decomp for decomp in decomp_metas if not decomp.is_complex] + if not decomp_metas: + return None, "No decompositions registered for the real-valued input" + + # NOTE: Complex overload type matching logic has been removed to keep this simple + # There should no longer be overloads (for the same opset version) in torchlib anymore + return (decomp_metas[0].onnx_function, "The first implementation is used") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_dynamic_shapes.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_dynamic_shapes.py new file mode 100644 index 0000000000000000000000000000000000000000..66199ade5290db85c4bb2ad71b3807088d8473e5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_dynamic_shapes.py @@ -0,0 +1,343 @@ +"""Compatibility functions for the torch.onnx.export API.""" + +# mypy: allow-untyped-defs +from __future__ import annotations + +import inspect +import warnings +from typing import Any, TYPE_CHECKING + +import torch +from torch.export.dynamic_shapes import _DimHint, Dim +from torch.onnx._internal._lazy_import import onnx_ir as ir +from torch.utils import _pytree + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +def from_dynamic_axes_to_dynamic_shapes( + model, + args: tuple[Any, ...], + kwargs: dict[str, Any] | None, + *, + dynamic_axes=None, + output_names: set[str], + input_names: Sequence[str] | None = None, +) -> tuple[dict[str, Any | None] | None, tuple[Any, ...], dict[str, Any] | None]: + """ + Converts dynamic_axes into dynamic_shapes by wrapping the axis names with ``torch.export.Dim.DYNAMIC``. + + dynamic_axes examples: + (1) dynamic_axes = {"x": {0: "my_custom_axis_name_1"}, "y": {1: "my_custom_axis_name_2"}} + (2) dynamic_axes = {"x": [0], "y": [1]} + + these will be converted to dynamic_shapes respectively: + (1) dynamic_shapes = {"x": {0: Dim.DYNAMIC}, "y": {1: Dim.DYNAMIC}} + (2) dynamic_shapes = {"x": {0: Dim.DYNAMIC}, "y": {1: Dim.DYNAMIC}} + + Detail on Dim.DYNAMIC: `#133620 `_ + """ + + warnings.warn( + "from_dynamic_axes_to_dynamic_shapes is deprecated and will be removed in a future release. " + "This function converts 'dynamic_axes' format (including custom axis names) to 'dynamic_shapes' format. " + "Instead of relying on this conversion, provide 'dynamic_shapes' directly with custom names.", + DeprecationWarning, + stacklevel=2, + ) + + # https://github.com/pytorch/pytorch/pull/128371 + # 1. The function does not need to provide dynamic_shapes to torch.export.export + if dynamic_axes is None: + return None, args, kwargs + + if input_names is None: + input_names = [] + + if kwargs is None: + kwargs = {} + + dynamic_shapes: dict[str, Any | None] = {} + for input_name, axes in dynamic_axes.items(): + # NOTE: torch.export.Dim.DYNAMIC does its best to infer the min and max values + # from the model, but it's not guaranteed to be dynamic. + if input_name in output_names: + # output names are not needed for dynamic_shapes + continue + if isinstance(axes, dict): + if any(not isinstance(k, int) for k in axes): + raise ValueError( + "The axis in dynamic_axes must be in the form of: dict[int, str] or list[int]." + ) + # str will be converted to Dim.DYNAMIC in convert_str_to_export_dim + dynamic_shapes[input_name] = axes + elif isinstance(axes, list): + if any(not isinstance(k, int) for k in axes): + raise ValueError( + "The axis in dynamic_axes must be in the form of: dict[int, str] or list[int]." + ) + dynamic_shapes[input_name] = dict.fromkeys(axes, torch.export.Dim.DYNAMIC) + elif axes is None: + dynamic_shapes[input_name] = None + else: + raise ValueError( + "Unsupported dynamic_axes format. Please provide a dict or a list." + ) + + for input_name in input_names: + if input_name not in dynamic_shapes: + dynamic_shapes[input_name] = None + + # Order the inputs according to the signature of the model + sig = _signature(model) + inputs = [] + for idx, param_name in enumerate(sig.parameters): + if idx < len(args): + inputs.append(args[idx]) + elif param_name in kwargs: + inputs.append(kwargs[param_name]) + + # We need tree structure to represent dynamic_shapes + dynamic_shapes = _unflatten_dynamic_shapes_with_inputs_tree(inputs, dynamic_shapes) + + # Since the dynamic_shapes are now in the order of the model parameters, + # we need to convert args and kwargs to the order of the model parameters. + return dynamic_shapes, tuple(inputs), {} + + +def from_dynamic_shapes_to_dynamic_axes( + dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any], + input_names: Sequence[str], + exception: Exception, +) -> dict[str, Any] | None: + """ + Converts dynamic_shapes into dynamic_axes by removing torch.export.Dim wrapping + and converting to list or dict form based on whether dimension names are present. + + dynamic_shapes examples: + (1) dynamic_shapes = {"x": {0: Dim("my_custom_axis_name_1")}, "y": {1: Dim("my_custom_axis_name_2")}} + (2) dynamic_shapes = ({0: Dim("my_custom_axis_name_1"}, {1: Dim("my_custom_axis_name_2")}) + + these will be converted to dynamic_axes respectively: + (1) dynamic_axes = {"x": [0], "y": [1]} + (2) dynamic_axes = {"x": [0], "y": [1]} + + NOTE: If the model input is nested, so is the dynamic_shapes, we need to flatten the dynamic_shapes, + and then assign the axes to the input names in the order they are provided. + + NOTE: input_names are used to assign the axes to the correct input names. If the input names are not + provided, or less than the dynamic inputs/axes, it raises an error. + """ + + flat_dynamic_shapes, _ = _flatten_dynamic_shapes_to_axes(dynamic_shapes) + + if len(input_names) < len(flat_dynamic_shapes): + raise ValueError( + "To construct dynamic_axes from dynamic_shapes, " + f"number of input names ({len(input_names)}) should be greater than or equal to " + f"the number of graph inputs(flat) ({len(flat_dynamic_shapes)})" + ) from exception + + dynamic_axes: dict[str, list[int]] = {} + # input names are assigned in order + for input_name, axes in zip(input_names, flat_dynamic_shapes): + if axes is None: + continue + + converted_axes: list[int] = [] + if isinstance(axes, dict): + for axis, dim in axes.items(): + if dim is None: + continue + converted_axes.append(axis) + dynamic_axes[input_name] = converted_axes + elif isinstance(axes, (list, tuple)): + for idx, dim in enumerate(axes): + if dim is None: + continue + converted_axes.append(idx) + dynamic_axes[input_name] = converted_axes + return dynamic_axes + + +def _any_str_or_dim_in_dynamic_shapes( + dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any], +) -> bool: + """Check if there is any string or Dim in the dynamic_shapes.""" + flat_dynamic_shapes, _ = _flatten_dynamic_shapes_to_axes(dynamic_shapes) + # This indicates the dynamic_shapes includes something we don't support in axes, and it's flattened + # to itself. Otherwise, flat_dynamic_shapes should be a list of dict/list/tuple (or None). + if any( + not isinstance(axes, (dict, list, tuple)) and axes is not None + for axes in flat_dynamic_shapes + ): + return False + # both str and Dim can provide custom names + for axes in flat_dynamic_shapes: + if isinstance(axes, dict): + for dim in axes.values(): + if isinstance(dim, (str, Dim)): + return True + elif isinstance(axes, (list, tuple)): + for dim in axes: + if isinstance(dim, (str, Dim)): + return True + return False + + +def convert_str_to_export_dim( + dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] | None, +) -> tuple[dict[str, Any] | tuple[Any, ...] | list[Any] | None, bool]: + # 1. If there is no string in dynamic_shapes, we do not touch dynamic_shapes + if dynamic_shapes is None or not _any_str_or_dim_in_dynamic_shapes(dynamic_shapes): + return dynamic_shapes, False + # 2. Convert "name" to Dim.DYNAMIC with flattening and identify if there is any string + # to be replaced with Dim.DYNAMIC, and then unflatten it back to the original structure. + # for example: {"y": {0: "dim_0"}, "x": {1: "dim_1"}} + # to {"y": {0: Dim.DYNAMIC}, "x": {1: Dim.DYNAMIC}} + dynamic_shapes_with_export_dim: list[ + list[Dim | _DimHint | None] | dict[int, Dim | _DimHint | None] | None + ] = [] + flat_dynamic_shapes, tree_structure = _flatten_dynamic_shapes_to_axes( + dynamic_shapes + ) + for axes in flat_dynamic_shapes: + if axes is None: + dynamic_shapes_with_export_dim.append(None) + elif isinstance(axes, dict): + converted_axes_dict: dict[int, Dim | _DimHint | None] = {} + for axis, dim in axes.items(): + if isinstance(dim, str): + converted_axes_dict[axis] = torch.export.Dim.DYNAMIC + else: + converted_axes_dict[axis] = dim + dynamic_shapes_with_export_dim.append(converted_axes_dict) + elif isinstance(axes, (list, tuple)): + converted_axes_list: list[Dim | _DimHint | None] = [] + for dim in axes: + if isinstance(dim, str): + converted_axes_list.append(torch.export.Dim.DYNAMIC) + else: + converted_axes_list.append(dim) + dynamic_shapes_with_export_dim.append(converted_axes_list) + + dynamic_shapes_with_export_dim = _pytree.tree_unflatten( + dynamic_shapes_with_export_dim, tree_structure + ) + return ( + dynamic_shapes_with_export_dim, + True, + ) + + +def create_rename_mapping( + inputs, dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any] +) -> dict[str, str]: + """Create a mapping from old names to new names for dynamic axes.""" + + # NOTE: There's no need to handle cases where kwargs are out of order with the model signature, + # as torch.export.export supports dynamism only when kwargs and dynamic_shapes are provided in order. + # Reference: https://github.com/pytorch/pytorch/blob/49082f9dba3b79a344cb03652972ddbe7c3729cc/torch/export/_trace.py#L2034 + + flat_dynamic_shapes, _ = _flatten_dynamic_shapes_to_axes(dynamic_shapes) + if len(inputs) != len(flat_dynamic_shapes): + warnings.warn( + "# ONNX model has different number of inputs than the flatten dynamic_shapes. " + "The dynamic axes will not be renamed.", + UserWarning, + stacklevel=3, + ) + return {} + rename_mapping: dict[str, str] = {} + # NOTE: We assume that the flat_dynamic_shapes is in the same order as the inputs + # When the axis is static, or it connects to _DimHint in dynamic shapes, we skip renaming + for idx, axes in enumerate(flat_dynamic_shapes): + input = inputs[idx] + if isinstance(axes, dict): + for dim, axis in axes.items(): + if not isinstance(input.shape[dim], ir.SymbolicDim): + continue + old_name = input.shape[dim].value + if old_name is None: + continue + # _DimHint, int and None exists in dynamic shapes, we skip renaming + if isinstance(axis, (_DimHint, int)) or axis is None: + continue + # NOTE: ExportedProgram could give the axes the same name if they share + # the same shape constraints. + custom_name = _get_custom_axis_name(axis) + if input.shape[dim].value in rename_mapping: + warnings.warn( + f"# The axis name: {custom_name} will not be used, since it shares " + f"the same shape constraints with another axis: {rename_mapping[input.shape[dim].value]}.", + stacklevel=2, + ) + continue + rename_mapping[input.shape[dim].value] = custom_name + elif isinstance(axes, (list, tuple)): + for dim, axis in enumerate(axes): + if not isinstance(input.shape[dim], ir.SymbolicDim): + continue + old_name = input.shape[dim].value + if old_name is None: + continue + # _DimHint, int and None exists in dynamic shapes, we skip renaming + if isinstance(axis, (_DimHint, int)) or axis is None: + continue + # NOTE: ExportedProgram could give the axes the same name if they share + # the same shape constraints. + custom_name = _get_custom_axis_name(axis) + if input.shape[dim].value in rename_mapping: + warnings.warn( + f"# The axis name: {custom_name} will not be used, since it shares " + f"the same shape constraints with another axis: {rename_mapping[input.shape[dim].value]}.", + UserWarning, + stacklevel=3, + ) + continue + rename_mapping[input.shape[dim].value] = _get_custom_axis_name(axis) + return rename_mapping + + +def _get_custom_axis_name(axis: Dim | str) -> str: + """Get the custom axis name from a torch.export.Dim.""" + if isinstance(axis, Dim): + return axis.__name__ + return axis + + +def _unflatten_dynamic_shapes_with_inputs_tree( + inputs: list[Any], + dynamic_shapes: dict[str, Any], +) -> dict[str, Any | None]: + _, tree_structure = _pytree.tree_flatten(inputs) + return _pytree.tree_unflatten(dynamic_shapes.values(), tree_structure) + + +def _flatten_dynamic_shapes_to_axes( + dynamic_shapes: dict[str, Any | None] | tuple[Any, ...] | list[Any], +) -> tuple[list[Any], _pytree.TreeSpec]: + # If it's a dict/list/tuple with torch.export.Dim, we consider it's an axis to dim mapping + def is_axes(x) -> bool: + return ( + isinstance(x, dict) + and all( + isinstance(k, int) + and (v is None or isinstance(v, (Dim, _DimHint, str, int))) + for k, v in x.items() + ) + ) or ( + isinstance(x, (list, tuple)) + and all(v is None or isinstance(v, (Dim, _DimHint, str, int)) for v in x) + ) + + return _pytree.tree_flatten(dynamic_shapes, is_leaf=is_axes) + + +def _signature(model) -> inspect.Signature: + should_be_callable = getattr(model, "forward", model) + if callable(should_be_callable): + return inspect.signature(should_be_callable) + raise ValueError("model has no forward method and is not callable") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_errors.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_errors.py new file mode 100644 index 0000000000000000000000000000000000000000..ff41bbe695fe7d0ebe60c40014332abc36430d0f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_errors.py @@ -0,0 +1,21 @@ +"""Error classes for the ONNX exporter.""" + +from __future__ import annotations + +import torch.onnx.errors + + +class TorchExportError(torch.onnx.errors.OnnxExporterError): + """Error during graph capturing using torch.export.""" + + +class ConversionError(torch.onnx.errors.OnnxExporterError): + """Error during ExportedProgram to ONNX conversion.""" + + +class DispatchError(ConversionError): + """Error during ONNX Function dispatching.""" + + +class GraphConstructionError(ConversionError): + """Error during ONNX graph construction.""" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_exportable_module.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_exportable_module.py new file mode 100644 index 0000000000000000000000000000000000000000..ef16ffd8da6d131fb43e4394eed2a171bc834b50 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_exportable_module.py @@ -0,0 +1,196 @@ +"""Abstract interface for ONNX exportable modules.""" + +from __future__ import annotations + +import abc +from typing import Any, TYPE_CHECKING + +import torch + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +class ExportableModule(torch.nn.Module, abc.ABC): + """Abstract interface for ONNX exportable modules. + + Inherit from this class and implement the defined abstract methods + to create a module that can be exported to ONNX format. + + Example:: + + class Model(torch.nn.Module): + def forward(self, x): + return x * 2 + + + class MyExportableModule(torch.onnx.ExportableModule): + def __init__(self): + super().__init__() + self.model = Model() + + def forward(self, x): + return self.model(x) + + def example_arguments(self): + return (torch.randn(2, 3, 224, 224),), None + + def input_names(self): + return ("input",) + + def output_names(self): + return ("output",) + + def dynamic_shapes(self): + return ({0: "batch_size"},) + + + exportable_module = MyExportableModule() + onnx_program = exportable_module.to_onnx() + # The model can also be supplied directly to torch.onnx.export + onnx_program = torch.onnx.export(exportable_module) + """ + + @abc.abstractmethod + def example_arguments(self) -> tuple[tuple[Any], dict[str, Any] | None]: + """Return example arguments for the model's forward method. + + This method must be implemented by subclasses to provide sample inputs + that can be used for tracing, testing, and ONNX export. The returned + arguments should be representative of the expected input shapes and types + during inference. + + Example:: + + def example_arguments(self): + # For a model expecting a single tensor input + return (torch.randn(1, 3, 224, 224),), None + + + def example_arguments(self): + # For a model with multiple inputs and keyword arguments + return (torch.randn(1, 3, 224, 224), torch.randn(1, 512)), { + "temperature": 1.0 + } + + Returns: + A tuple containing: + + - A tuple of positional arguments to pass to the forward method + - A dictionary of keyword arguments (or None if no kwargs are needed) + """ + raise NotImplementedError + + def dynamic_shapes(self) -> Any: + """Return dynamic shape specifications for the model's inputs. + + Override this method to specify which dimensions of the input tensors + should be treated as dynamic during ONNX export. This allows the exported + model to accept inputs with varying sizes along the specified dimensions. + + Example:: + + def dynamic_shapes(self): + # Specify batch dimension as dynamic for input named 'x' + return {"x": {0: "batch_size"}} + + + def dynamic_shapes(self): + # Multiple dynamic dimensions + return { + "input": {0: "batch_size", 2: "height", 3: "width"}, + "mask": {0: "batch_size"}, + } + + Note: + The default implementation returns None, indicating all dimensions are static. + + Returns: + Dynamic shape specification compatible with ``torch.export.export``. + Return None if all input dimensions should be static. The format can be: + + - A dictionary mapping input names to dimension specifications + - A tuple/list of dimension specifications corresponding to inputs + - Any format accepted by the ``dynamic_shapes`` parameter of ``torch.export.export`` + """ + return None + + def input_names(self) -> Sequence[str] | None: + """Return names for the model's input tensors. + + Override this method to provide custom names for the input tensors in the + exported ONNX model. These names will be used as identifiers in the ONNX + graph and can be useful for debugging and model inspection. + + Example:: + + def input_names(self): + return ["image", "mask"] + + + def input_names(self): + # For a single input + return ["input_tensor"] + + Note: + The default implementation returns None, which results in auto-generated names. + + Returns: + A sequence of strings representing input names, or None to use default names. + The number of names should match the number of positional arguments in the + forward method. + """ + return None + + def output_names(self) -> Sequence[str] | None: + """Return names for the model's output tensors. + + Override this method to provide custom names for the output tensors in the + exported ONNX model. These names will be used as identifiers in the ONNX + graph and can be useful for debugging and model inspection. + + Example:: + + def output_names(self): + return ["logits", "probabilities"] + + + def output_names(self): + # For a single output + return ["prediction"] + + Note: + The default implementation returns None, which results in auto-generated names. + + Returns: + A sequence of strings representing output names, or None to use default names. + The number of names should match the number of outputs from the forward method. + For models returning multiple outputs, provide a name for each output. + """ + return None + + def to_onnx(self, **kwargs: Any) -> torch.onnx.ONNXProgram: + """Export the module to ONNX format. + + This method provides a convenient wrapper around ``torch.onnx.export`` that + automatically uses the example arguments, dynamic shapes, and input/output + names defined by the module. Additional export options can be specified via + keyword arguments. + + See Also: ``torch.onnx.export`` for complete documentation of export options. + + Args: + **kwargs: Additional keyword arguments to pass to ``torch.onnx.export``. + Common options include: + + - ``opset_version`` (int): The ONNX opset version to target + - ``optimize`` (bool): Whether to apply optimizations to the exported model + + Returns: + An ONNXProgram object containing the exported model and metadata. + """ + result = torch.onnx.export(self, **kwargs) + if result is None: + raise AssertionError("result must be non-None") + return result diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_flags.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_flags.py new file mode 100644 index 0000000000000000000000000000000000000000..8e9d8c9db6e444f802d7b50998ef2a5ee13028de --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_flags.py @@ -0,0 +1,32 @@ +"""Internal flags for ONNX export.""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING, TypeVar +from typing_extensions import ParamSpec + + +if TYPE_CHECKING: + from collections.abc import Callable + + +_is_onnx_exporting = False + +# Use ParamSpec to preserve parameter types instead of erasing to Any +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +def set_onnx_exporting_flag(func: Callable[_P, _R]) -> Callable[_P, _R]: + @functools.wraps(func) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + global _is_onnx_exporting + _is_onnx_exporting = True + try: + return func(*args, **kwargs) + finally: + # Ensure it resets even if an exception occurs + _is_onnx_exporting = False + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_fx_passes.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_fx_passes.py new file mode 100644 index 0000000000000000000000000000000000000000..ef599c27922a001fcdd04055e6a53759bf1fc719 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_fx_passes.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import torch +import torch.export +import torch.fx +from torch.onnx._internal.exporter import _decomp, _registration +from torch.onnx._internal.fx import passes + + +def decompose_with_registry( + exported_program: torch.export.ExportedProgram, registry: _registration.ONNXRegistry +) -> torch.export.ExportedProgram: + """Decompose the exported program with the given registry. + + This function is needed so it shows clearly on the profiler results. + """ + onnx_registered_ops = set(_decomp.get_onnx_implemented_overloads(registry)) + decomp_table = _decomp.create_onnx_friendly_decomposition_table(onnx_registered_ops) + return exported_program.run_decompositions(decomp_table) + + +def insert_type_promotion_nodes( + graph_module: torch.fx.GraphModule, +) -> None: + """Inplace pass to insert explicit type promotion nodes, recursively through nested modules.""" + for module in graph_module.modules(): + if not isinstance(module, torch.fx.GraphModule): + raise AssertionError(f"Expected GraphModule, got {type(module)}") + passes.InsertTypePromotion(module).run() + + +def remove_assertion_nodes(graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule: + """Remove all assertion and check nodes from the FX graph""" + aten_assertion_targets = { + torch.ops.aten.sym_constrain_range_for_size.default, + torch.ops.aten._assert_async.default, + torch.ops.aten._assert_async.msg, + torch.ops.aten._assert_scalar.default, + torch.ops.aten._assert_tensor_metadata.default, + } + for gm in graph_module.modules(): + for node in gm.graph.nodes: # type: ignore[union-attr] + if node.op == "call_function" and node.target in aten_assertion_targets: + gm.graph.erase_node(node) # type: ignore[operator, union-attr] + gm.recompile() # type: ignore[operator] + return graph_module diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_input_observer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_input_observer.py new file mode 100644 index 0000000000000000000000000000000000000000..820db724f288a24ba7760c6dce7e30f28b9603a4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_input_observer.py @@ -0,0 +1,1224 @@ +from __future__ import annotations + + +__all__ = ["InputObserver"] + +import contextlib +import inspect +import time +from typing import Any, TYPE_CHECKING + +import torch +from torch.onnx._internal.exporter import _onnx_program + + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + import onnxruntime as ort + + +def _flatten_unflatten_for_dynamic_shapes( + obj: Any, + change_function: Callable[[torch.Tensor], Any] | None = None, +) -> Any: + """Returns the object in a different structure similar to what + the definition of the dynamic shapes should use. + + Args: + obj: Object from a custom class. + change_function: If not None, this function is called to modify the tensors + in the structure itself, like replace them by a shape. + + Returns: + The flattened object. + """ + if isinstance(obj, torch.Tensor): + return change_function(obj) if change_function else obj + flat, spec = torch.utils._pytree.tree_flatten(obj) + start = 0 + end = 0 + subtrees = [] + for subspec in ( + spec.children() if hasattr(spec, "children") else spec.children_specs + ): + end += subspec.num_leaves + value = subspec.unflatten(flat[start:end]) + value = _flatten_unflatten_for_dynamic_shapes( + value, change_function=change_function + ) + subtrees.append(value) + start = end + if spec.type is dict: + # This is a dictionary. + return dict(zip(spec.context, subtrees)) + if spec.type is tuple: + return tuple(subtrees) + if spec.type is list: + return list(subtrees) + if spec.type is None and not subtrees: + return None + if spec.context: + # This is a custom class with attributes. + # It is returned as a list. + return list(subtrees) + raise ValueError( + f"Unable to interpret spec type {spec.type} " + f"(type is {type(spec.type)}, context is {spec.context}), " + f"spec={spec}, subtrees={subtrees}" + ) + + +def _infer_dynamic_dimensions( + shape_list: Sequence[tuple[int, ...]], set_batch_dimension: bool = False +) -> list[int]: + """Returns the list of dynamic dimensions given a list of shapes + corresponding to the same tensor. + + Args: + shape_list: + List of shapes, they must all have the same length. + set_batch_dimension: + Forces the first dimension to be treated as dynamic, + even if all shapes have the same value for that dimension. + + Returns: + list of dynamic dimensions + """ + unique_ranks = {len(shape) for shape in shape_list} + torch._check( + len(unique_ranks) == 1, + lambda: f"All shapes in shape_list must have the same rank but {shape_list=}.", + ) + rank = unique_ranks.pop() + dynamic = [] + for i in range(rank): + dims = [shape[i] for shape in shape_list] + if len(set(dims)) > 1 or (i == 0 and set_batch_dimension): + dynamic.append(i) + return dynamic + + +class InputCandidate: + """Retains one set of inputs given to the forward method or any + other method the class :class:`InputObserver` is stealing from. + Any class is allowed as long as it can be flattened. + + Args: + args: Positional arguments. + kwargs: Optional arguments. + clone: Clone the inputs before storing them. Some tensors + may be modified inplace, the original value must be retained. + cst_kwargs: Any optional arguments constant over multiple calls. + int, float, str, bool values must be stored here. + + The constructor flattens the received arguments. + Any necessary flattening function should have been registered first. + """ + + def __init__( + self, + args: tuple[Any, ...], + kwargs: dict[str, Any], + clone: bool, + cst_kwargs: dict[str, int | str | float | bool], + ): + self.args = args + self.kwargs = kwargs + self.flat_list, self.spec = torch.utils._pytree.tree_flatten((args, kwargs)) + self._position_to_args_kwargs: list[int | str] | None = None + self._n_tensors_for_args_kwargs: dict[int | str, int] | None = None + self.cst_kwargs = cst_kwargs.copy() + + if clone: + self.flat_list = [ + (None if not isinstance(t, torch.Tensor) else t.clone().detach()) + for t in self.flat_list + ] + self.args, self.kwargs = torch.utils._pytree.tree_unflatten( + self.flat_list, self.spec + ) + + self.aligned_spec: torch.utils._pytree.PyTreeSpec | None = None + self.aligned_flat_list: list[torch.Tensor | None] | None = None + + def __str__(self) -> str: + return ( + f"{self.__class__.__name__}({len(self.args)} args, " + f"{len(self.kwargs)} kwargs, {len(self.flat_list)} tensors, " + f"{len(self.aligned_flat_list or [])} aligned tensors)" + ) + + def __len__(self) -> int: + """Returns the number of flattened tensors, None tensors are included.""" + return len(self.flat_list) + + def str_obs(self) -> str: + """Prints out some information about the observations.""" + return ( + f"InputCandidate(args=#{len(self.args)}(...), " + f"kwargs=#{len(self.kwargs)}{{...}}, cst_kwargs={self.cst_kwargs})" + ) + + def build_mappings(self) -> list[int | str]: + if self._position_to_args_kwargs is not None: + return self._position_to_args_kwargs + self._n_tensors_for_args_kwargs = {} + + flat_index_to_args: list[int | str] = [] + for index_args, a in enumerate(self.args): + size = len(torch.utils._pytree.tree_flatten(a)[0]) + self._n_tensors_for_args_kwargs[index_args] = size + flat_index_to_args.extend([index_args] * size) + for k, v in self.kwargs.items(): + size = len(torch.utils._pytree.tree_flatten(v)[0]) + self._n_tensors_for_args_kwargs[k] = size + flat_index_to_args.extend([k] * size) + + self._position_to_args_kwargs = flat_index_to_args + return self._position_to_args_kwargs + + @property + def position_to_args_kwargs(self) -> list[int | str]: + """Returns the corresponding args or kwargs + for every tensor in the flattened inputs. + """ + if self._position_to_args_kwargs is None: + self.build_mappings() + # pyrefly: ignore [bad-return] + return self._position_to_args_kwargs + + @property + def n_tensors_for_args_kwargs(self) -> dict[int | str, int]: + """Returns the number of flat tensors in every args or kwargs.""" + if self._n_tensors_for_args_kwargs is None: + self.build_mappings() + # pyrefly: ignore [bad-return] + return self._n_tensors_for_args_kwargs + + def _set_aligned_flat_list( + self, + aligned_flat_list: list[torch.Tensor | None], + aligned_spec: torch.utils._pytree.PyTreeSpec, + ): + self.aligned_flat_list = aligned_flat_list + self.aligned_spec = aligned_spec + + def align_with( + self, + best_candidate: InputCandidate, + captured_inputs: dict[int | str, int], + signature_names: list[str], + ): + """Two candidates are considered as aligned if after being flattened + if they have the same number of tensors (None allowed).""" + if self.cst_kwargs != best_candidate.cst_kwargs: + raise RuntimeError( + f"Two calls were made with different constant values, " + f"{self.cst_kwargs} != {best_candidate.cst_kwargs}" + ) + + args = self.args + if len(self.args) > len(best_candidate.args): + # We need to move some args to kwargs as the best_candidate does. + new_kwargs = {} + for i in range(len(best_candidate.args), len(self.args)): + new_kwargs[signature_names[i]] = args[i] + args = args[: len(best_candidate.args)] + kwargs = {**new_kwargs, **self.kwargs} + else: + kwargs = self.kwargs + + flat = [] + for i in range(len(best_candidate.args)): + if i < len(args) and (isinstance(args[i], torch.Tensor) or args[i]): + ts = torch.utils._pytree.tree_flatten(self.args[i])[0] + if i in captured_inputs and captured_inputs[i] != len(ts): + raise RuntimeError( + f"Positional argument {i} has {len(ts)} tensors " + f"but previously got {captured_inputs[i]} tensors. " + f"Inference is impossible in that case." + ) + captured_inputs[i] = len(ts) + flat.extend(ts) + continue + # If the argument i is not specified or is None or an empty container. + flat.extend( + [None for _ in range(best_candidate.n_tensors_for_args_kwargs[i])] + ) + + for k in best_candidate.kwargs: + if k in kwargs and (isinstance(kwargs[k], torch.Tensor) or kwargs[k]): + ts = torch.utils._pytree.tree_flatten(kwargs[k])[0] + if k in captured_inputs and captured_inputs[k] != len(ts): + raise RuntimeError( + f"Named argument {k!r} has {len(ts)} tensors " + f"but previously got {captured_inputs[k]} tensors in " + f"kwargs={list(kwargs)}. " + f"Inference is impossible in that case." + ) + captured_inputs[k] = len(ts) + flat.extend(ts) + continue + # If the argument k is not specified or is None or an empty container. + flat.extend( + [None for _ in range(best_candidate.n_tensors_for_args_kwargs[k])] + ) + + self._set_aligned_flat_list(flat, best_candidate.spec) + + @property + def n_aligned_tensors(self) -> int: + if self.aligned_flat_list is None: + raise RuntimeError("This input was not aligned with the others.") + return len(self.aligned_flat_list) + + +class InputObserverInfo: + """Contains all the necessary information to infer dynamic shapes + and the arguments to send to :func:`torch.export.export`. + + Args: + signature_names: Names of the arguments of the method + the collector tensors come from. They are used if it becomes + necessary to move positional arguments to named ones. + They are used a second time because :func:`torch.export.export` + cares about the order in kwargs and dynamic shapes, it needs + to be the same in the ordered dictionaries `add_inputs` receive. + default_values: Default values defined by the signature of the function, + any value equal to that is ignored to simplify the export. + value_if_missing: If an argument is missing, + a default value will be taken in this dictionary, + this is used when after the prefill step, an argument + disappears (such as `pixel_values`) and another one + is added (such as `past_key_values`). + The values are only to infer dynamic shapes and arguments, + not to run the model. + args_name_and_position: Name of parameter `*args` + and its position if it exists. + kwargs_name: Name of the variable keyword parameter `**kwargs` if it exists. + + This is used by class :class:`InputObserver`. + """ + + def __init__( + self, + signature_names: list[str], + default_values: dict[str, int | bool | str | float], + value_if_missing: dict[str | int, Any], + args_name_and_position: tuple[str, int] | None, + kwargs_name: str | None, + ): + self.default_values = default_values + self.value_if_missing = value_if_missing + self.inputs: list[InputCandidate] = [] + self.outputs_specs: list[torch.utils._pytree.PyTreeSpec] = [] + self.flat_outputs: list[list[torch.Tensor | None]] = [] + self.latencies: list[float] = [] + self.args_name_and_position = args_name_and_position + self.kwargs_name = kwargs_name + self.signature_names = signature_names + self._best_candidate: InputCandidate | None = None + self._captured_inputs: dict[int | str, int] | None = None + + def __len__(self) -> int: + """Returns the number of collected set of inputs/outputs.""" + return len(self.inputs) + + def add_inputs(self, args: tuple[Any, ...], kwargs: dict[str, Any]): + """Stores one set of inputs. They are deepcopied. + + Args: + args: Positional arguments. + kwargs: Named arguments. + """ + cst_kwargs = { + k: v + for k, v in kwargs.items() + if k in self.signature_names + and isinstance(v, (int, float, bool, str)) + and v != self.default_values.get(k, None) + and self.default_values.get(k, None) is not None + } + kwargs = { + k: v + for k, v in kwargs.items() + if v is not None and not isinstance(v, (int, float, bool, str)) + } + + # adds value_if_missing attributes + for k, v in self.value_if_missing.items(): + if isinstance(k, str): + if k not in kwargs: + # Validate that `value_if_missing` keys are compatible + # with the observed signature. + # If the function does not accept **kwargs, + # all value_if_missing keys must be + # present in the observed signature names. + if k not in self.signature_names and not self.kwargs_name: + raise ValueError( + f"Unexpected keyword argument {k!r} " + f"provided as a value_if_missing input " + "for a function that does not accept it. " + f"All value_if_missing keys must " + f"be in the observed signature: {tuple(self.signature_names)}." + ) + kwargs[k] = v + elif isinstance(k, int): + if k >= len(self.signature_names): + raise ValueError( + f"Unexpected keyword argument {k=} " + f"provided as a value_if_missing input " + "for a function that does not accept it. " + f"All value_if_missing indices must " + f"be in the observed signature: {tuple(self.signature_names)}." + ) + if k >= len(args): + raise NotImplementedError( + f"Unexpected keyword argument {k=} " + f"provided as a value_if_missing input " + "for a function that does not accept it. " + f"All value_if_missing indices must " + f"be in the observed signature: {tuple(self.signature_names)}, " + f"only {len(args)} were given." + ) + list_args = list(args) + list_args[k] = v + args = tuple(list_args) + else: + raise TypeError( + f"Unexpected type {type(k)} for a missing value. The key is {k!r}." + ) + + # kwargs may come in a different order each time. + # dictionaries are ordered and torch.export.export expects + # dynamic shapes and kwargs to follow the same order. + + ordered_kwargs = {k: kwargs[k] for k in self.signature_names if k in kwargs} + for k, v in kwargs.items(): + if k not in ordered_kwargs: + ordered_kwargs[k] = v + + candidate = InputCandidate( + args, ordered_kwargs, clone=True, cst_kwargs=cst_kwargs + ) + self.inputs.append(candidate) + if self._best_candidate is None or len(self._best_candidate) < len(candidate): + self._best_candidate = candidate + + def add_outputs(self, res: torch.Tensor | tuple[torch.Tensor, ...], latency: float): + """Stores outputs. They are deepcopied.""" + flat_res, spec = torch.utils._pytree.tree_flatten(res) + self.outputs_specs.append(spec) + self.flat_outputs.append( + [(None if t is None else t.clone().detach()) for t in flat_res] + ) + self.latencies.append(latency) + + def align_inputs_none_values(self): + """Once the best candidate is chosen, this method aligns every set of inputs + on the best candidate, it inserts None at the right position when + optional inputs are not specified. We consider a set of inputs is aligned + if this method does not change the original flattened inputs. + """ + if not self.inputs or self._best_candidate is None: + raise RuntimeError("No inputs were captured.") + + if all(candidate.aligned_flat_list is not None for candidate in self.inputs): + # No new inputs, no alignment is necessary. + return + + # Let's reprocess everything. + self._captured_inputs = {} + for candidate in self.inputs: + if len(set(candidate.kwargs) | set(self._best_candidate.kwargs)) > len( + self._best_candidate.kwargs + ): + raise RuntimeError( + f"At least one call to the observed model " + f"must contain all the named arguments. " + f"candidate kwargs={list(candidate.kwargs)}, " + f"best candidate kwargs={list(self._best_candidate.kwargs)}." + ) + candidate.align_with( + self._best_candidate, self._captured_inputs, self.signature_names + ) + + def infer_dynamic_shapes( + self, + set_batch_dimension_for: set[int | str] | bool | None = None, + return_flat: bool = False, + ) -> tuple[dict[int, Any] | None, ...] | dict[str, dict[int, Any] | None]: + """Infers dynamic shapes based on the collected tensors. + Most of the time, models do support a batch dimension + but this batch dimension has the same value for every input sample. + Instead of running inference on new samples, argument `set_batch_dimension_for` + can be used to tell the first dimension is a dynamic dimension for a particular + set of inputs referenced by their name (str) or their position (int). + + Args: + set_batch_dimension_for (set[int | str] | bool | None): Set of input identifiers, + by name (``str``) or position (``int``), for which the first dimension + should be treated as a dynamic batch dimension. If ``None`` or empty, + no additional batch dimensions are marked as dynamic. + return_flat: Tells the function to return a flat tuple instead of + nested structured. This option is used internally to infer arguments. + """ + self.align_inputs_none_values() + assert self._best_candidate is not None # noqa: S101 + assert self._best_candidate.flat_list is not None # noqa: S101 + assert self._best_candidate.aligned_flat_list is not None # noqa: S101 + + def _set_batch_dimension(name_or_position) -> bool: + if not set_batch_dimension_for: + return False + if ( + isinstance(set_batch_dimension_for, bool) and set_batch_dimension_for + ) or name_or_position in set_batch_dimension_for: + return True + if isinstance(name_or_position, int): + torch._check( + name_or_position < len(self.signature_names), + lambda: f"argument at position {name_or_position} is out of boundary", + ) + if self.signature_names[name_or_position] in set_batch_dimension_for: + return True + return False + + def _set_batch_dimension_for_flat_index(index) -> bool: + return _set_batch_dimension( + # pyrefly: ignore[missing-attribute] + self._best_candidate.position_to_args_kwargs[index] + ) + + if len(self._best_candidate.flat_list) != len( + self._best_candidate.aligned_flat_list + ): + raise NotImplementedError( + "infer_dynamic_shapes is not implemented " + "when the best candidate is not 'aligned'. " + "This happens when there is no stored set of inputs where " + "all optional inputs showing in other sets are defined." + ) + + if len({inputs.n_aligned_tensors for inputs in self.inputs}) != 1: + raise NotImplementedError( + f"infer_dynamic_shapes is not implemented " + f"when the number of input tensors are not the same in " + f"every set of inputs " + f"{[inputs.n_aligned_tensors for inputs in self.inputs]}." + ) + shape_lists = [ + [(None if t is None else t.shape) for t in candidate.aligned_flat_list] + for candidate in self.inputs + if candidate.aligned_flat_list is not None + ] + n_tensors = len(shape_lists[0]) + dynamic_shapes = [ + _infer_dynamic_dimensions( + [s for s in [shapes[index] for shapes in shape_lists] if s is not None], + set_batch_dimension=_set_batch_dimension_for_flat_index(index), + ) + for index in range(n_tensors) + ] + cst = torch.export.Dim.DYNAMIC + flat_dynamic_shapes = [dict.fromkeys(dims, cst) for dims in dynamic_shapes] + if return_flat: + return tuple(flat_dynamic_shapes) + + # Let's regroup. + if len(flat_dynamic_shapes) == len(self._best_candidate.args) + len( + self._best_candidate.kwargs + ): + # It means forward method is called with tensors only. + if ( + not self._best_candidate.kwargs + and not self._best_candidate.cst_kwargs + and not self.args_name_and_position + ): + # only positional arguments + return tuple(flat_dynamic_shapes) + if not self._best_candidate.args: + # only named arguments + ds = dict(zip(list(self._best_candidate.kwargs), flat_dynamic_shapes)) + return self._post_process_for_kwargs( + {**ds, **dict.fromkeys(self._best_candidate.cst_kwargs, None)} + ) + if not self.args_name_and_position: + # positional arguments needs to be moved to the named arguments + n_args = len(self._best_candidate.args) + pos_names = self.signature_names[:n_args] + return self._post_process_for_kwargs( + { + **dict(zip(pos_names, flat_dynamic_shapes[:n_args])), + **dict( + zip( + list(self._best_candidate.kwargs), + flat_dynamic_shapes[n_args:], + ) + ), + **dict.fromkeys(self._best_candidate.cst_kwargs, None), + } + ) + # positional arguments needs to be moved to the named arguments + n_args = min(len(self._best_candidate.args), self.args_name_and_position[1]) + i_kwargs = max( + len(self._best_candidate.args), self.args_name_and_position[1] + ) + var_pos = self.args_name_and_position[0] + pos_names = self.signature_names[:n_args] + return self._post_process_for_kwargs( + { + **dict(zip(pos_names, flat_dynamic_shapes[:n_args])), + var_pos: tuple(flat_dynamic_shapes[n_args:i_kwargs]), + **dict( + zip( + list(self._best_candidate.kwargs), + flat_dynamic_shapes[i_kwargs:], + ) + ), + **dict.fromkeys(self._best_candidate.cst_kwargs, None), + } + ) + + # nested types, here comes the fun part because the shapes cannot be unflattened, + # custom classes must appear in their flattened shape. + # This does not work in all cases but every time every available argument is flattened + # with the same number of tensors. The function does not check + # if that assumption is true. + flat_inputs, _max_spec = torch.utils._pytree.tree_flatten( + (self._best_candidate.args, self._best_candidate.kwargs) + ) + torch._check( + len(flat_inputs) == len(flat_dynamic_shapes), + ( + f"Length mismatch len(flat_inputs)={len(flat_inputs)}, " + f"len(flat_dynamic_shapes)={len(flat_dynamic_shapes)}" + ), + ) + + index = 0 + + def change_function(t): + nonlocal index + if index >= len(flat_dynamic_shapes): + raise RuntimeError( + f"Flattened {index} tensors when there are only " + f"{len(flat_dynamic_shapes)}." + ) + res = flat_dynamic_shapes[index] + index += 1 + return res + + ds_args, ds_kwargs = _flatten_unflatten_for_dynamic_shapes( + (self._best_candidate.args, self._best_candidate.kwargs), + change_function=change_function, + ) + if self._best_candidate.cst_kwargs: + ds_kwargs = { + **ds_kwargs, + **dict.fromkeys(self._best_candidate.cst_kwargs, None), + } + if not ds_kwargs and not self.args_name_and_position: + return tuple(ds_args) + if not ds_args: + return self._post_process_for_kwargs(ds_kwargs) + + if not self.args_name_and_position: + pos_names = self.signature_names[: len(ds_args)] + return self._post_process_for_kwargs( + {**dict(zip(pos_names, ds_args)), **ds_kwargs} + ) + + n_args = min(len(ds_args), self.args_name_and_position[1]) + pos_names = self.signature_names[:n_args] + return self._post_process_for_kwargs( + { + **dict(zip(pos_names, ds_args[:n_args])), + self.args_name_and_position[0]: tuple(ds_args[n_args:]), + **ds_kwargs, + } + ) + + def infer_arguments( + self, + index_or_candidate: InputCandidate | int | None = None, + /, + flat: bool = False, + as_args_kwargs: bool = False, + ) -> ( + list[torch.Tensor | None] + | tuple[torch.Tensor, ...] + | dict[str, torch.Tensor] + | tuple[list[torch.Tensor] | tuple[torch.Tensor, ...], dict[str, torch.Tensor]] + ): + """Infers arguments based on the collected tensors. + + Args: + index_or_candidate: If missing, the method selects one set of inputs + among the available ones, usually the set of inputs containing + with the highest number of tensors. + It then replaces None values and missing tensors with empty tensors. + If not missing, it can be an integer to fetch one of the stored set + or some inputs. + flat: If True, it returns a flattened list of tensors, + if False, it returns a tuple or a dictionary preserving + the nested structures. The flat version is used internally. + It produces a single list of tensors easier to process or modify + rather than a nested structure holding the same tensors. + The original structure can be restored with + ``torch.utils._pytree.tree_unflatten(flat_list, self.aligned_spec)``. + This mechanism is used to replace None values by empty tensors. + as_args_kwargs: If True, the method always returns `(args, kwargs)`, + otherwise, it returns either a tuple (only args) or a dictionary + (only kwargs) or raises an exception if it cannot do so. + Returns: + Inferred arguments, every optional tensor is replaced by an empty tensor. + """ + # This is already checked by _build_inputs_completed_with_none_values + # but this is not always well captured by tools checking types. + self.align_inputs_none_values() + assert self._best_candidate is not None # noqa: S101 + candidate = None + if index_or_candidate is None: + for cand in self.inputs: + args, kwargs = cand.args, cand.kwargs + if len(args) == len(self._best_candidate.args or ()) and len( + kwargs + ) == len(self._best_candidate.kwargs or {}): + candidate = cand + break + elif isinstance(index_or_candidate, int): + torch._check( + index_or_candidate < len(self.inputs), + lambda: ( + f"No stored input set for index=" + f"{index_or_candidate}<{len(self.inputs)}." + ), + ) + candidate = self.inputs[index_or_candidate] + else: + candidate = index_or_candidate + + assert candidate is not None # noqa: S101 + if candidate.aligned_flat_list is None: + raise RuntimeError( + f"Candidate {candidate} has no aligned flat list of tensors, " + f"index_or_candidate={index_or_candidate}. You should call " + f"method 'align_with'." + ) + + aligned_flat_list = candidate.aligned_flat_list + assert aligned_flat_list is not None # noqa: S101 + if any(t is None for t in aligned_flat_list): + dynamic_shapes = self.infer_dynamic_shapes(return_flat=True) + assert isinstance(dynamic_shapes, tuple) # noqa: S101 + aligned_flat_list = list(aligned_flat_list) + for index in range(len(aligned_flat_list)): + if aligned_flat_list[index] is not None: + continue + shape = dynamic_shapes[index] + all_non_empty_tensors = [ + c.aligned_flat_list[index] + for c in self.inputs + if c.aligned_flat_list is not None + ] + all_non_empty_tensors_not_none = [ + t for t in all_non_empty_tensors if t is not None + ] + if not all_non_empty_tensors_not_none: + raise RuntimeError( + f"There is no tensor at position {index} in any flattened inputs." + ) + tensor = all_non_empty_tensors_not_none.pop() + if tensor.numel() == 0: + aligned_flat_list[index] = tensor + continue + if not shape: + aligned_flat_list[index] = torch.zeros( + tensor.shape, dtype=tensor.dtype, device=tensor.device + ) + continue + dim = max(shape) + torch._check( + dim < tensor.ndim, + lambda index=index, shape=shape, tshape=tensor.shape: ( + f"Tensor shape {tshape} does not match the " + f"dynamic shape {shape} at position {index}." + ), + ) + new_shape = list(tensor.shape) + new_shape[dim] = 0 + aligned_flat_list[index] = torch.empty( + tuple(new_shape), dtype=tensor.dtype, device=tensor.device + ) + if flat: + return aligned_flat_list + args, kwargs = torch.utils._pytree.tree_unflatten( + aligned_flat_list, + # pyrefly: ignore[bad-argument-type] + candidate.aligned_spec, + ) + if self._best_candidate.cst_kwargs: + kwargs = {**kwargs, **self._best_candidate.cst_kwargs} + + if not as_args_kwargs: + if not kwargs: + return args + if not args: + return kwargs + + # We need to move args to kwargs + if self.args_name_and_position: + raise RuntimeError( + "Cannot return arguments " + "as a single tuple or a single dictionary " + "because of '*args' in the function signature. " + "You need to set `as_args_kwargs=True`." + ) + n_args = len(args) + pos_names = self.signature_names[:n_args] + return {**dict(zip(pos_names, args[:n_args])), **kwargs} + + # Generic case. + return tuple(args), kwargs + + def _post_process_for_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]: + """:func:`torch.export.export` requires dynamic shapes and keyword arguments + that are not part of the explicit function signature but are collected via + ``**`` to be wrapped under the corresponding parameter name + (``self.kwargs_name``) as ``{: {'param': shape or tensor}}``. + This function ensures this wrapping is performed when ``self.kwargs_name`` is set. + """ + if not self.kwargs_name: + # Nothing to do here. + return kwargs + to_be_moved = {k for k in kwargs if k not in self.signature_names} + if not to_be_moved: + return kwargs + keywords = {k: v for k, v in kwargs.items() if k in to_be_moved} + new_kwargs = {k: v for k, v in kwargs.items() if k not in to_be_moved} + if self.kwargs_name in new_kwargs: + raise ValueError( + f"Keyword argument name collision: received a keyword argument " + f"'{self.kwargs_name}' which conflicts with the **{self.kwargs_name} " + "parameter used to collect extra keyword arguments. " + "Passing a keyword argument with this name is not supported." + ) + return {**new_kwargs, self.kwargs_name: keywords} + + +class InputObserver: + """Steals forward method to collect inputs and outputs. + This information is used to infer dynamic shapes and + export arguments. + + Args: + value_if_missing: If an argument is missing, + a default value will be taken in this dictionary, + this is used when after the prefill step, an argument + disappears (such as `pixel_values`) and another one + is added (such as `past_key_values`). + The values are only to infer dynamic shapes and arguments, + not to run the model. + + Examples + -------- + >>> input_observer = InputObserver() + >>> with input_observer(model): + >>> model(x1, y1) + >>> model(x2, y2) + >>> ep = torch.export.export( # or torch.onnx.export + >>> model, + >>> input_observer.infer_arguments(), + >>> dynamic_shapes.input_observer.infer_dynamic_shapes(), + >>> ) + + With LLM: + + >>> input_observer = InputObserver() + >>> with input_observer(model): + >>> model.generate(input_ids) + >>> ep = torch.export.export( # or torch.onnx.export + >>> model, + >>> (), + >>> kwargs=input_observer.infer_arguments(), + >>> dynamic_shapes.input_observer.infer_dynamic_shapes(), + >>> ) + + The last example considers an LLM taking images and text as inputs. + The first call to the forward method which we try to export has `pixel_values` + but no `past_key_values`. The next calls do not have `pixel_values` but + `past_key_values`. The observer understands `pixel_values` and `past_key_values` + are needed but they may not be both specified at the same time. + Since `pixel_values` only appears in the first call, the observer cannot + tell how to infer an empty tensor for this argument. That's what the argument + `value_if_missing` is for. The following example is more than a dummy example + but shows how to use it with ``transformers``. + + .. code-block:: python + + from transformers import pipeline + + model_id = "tiny-random/gemma-3" + pipe = pipeline( + "image-text-to-text", + model=model_id, + device="cpu", + trust_remote_code=True, + max_new_tokens=3, + dtype=torch.float16, + ) + messages = [ + { + "role": "system", + "content": [{"type": "text", "text": "You are a helpful assistant."}], + }, + { + "role": "user", + "content": [ + { + "type": "image", + "url": "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/p-blog/candy.JPG", + }, + {"type": "text", "text": "What animal is on the candy?"}, + ], + }, + ] + observer = InputObserver( + value_if_missing=dict( + pixel_values=torch.empty((0, 3, 896, 896), dtype=torch.float16) + ) + ) + with observer(pipe.model): + pipe(text=messages, max_new_tokens=4) + + .. versionadded:: 2.11.0 + """ + + def __init__(self, value_if_missing: dict[str | int, Any] | None = None): + self.info: InputObserverInfo | None = None + self.value_if_missing = value_if_missing or {} + + def _replaced_method( + self, + *args, + _captured_method: Callable | None = None, + _store_n_calls: int = 3, + **kwargs, + ): + if _captured_method is None: + raise AssertionError("_captured_forward cannot be None") + if self.info is None: + raise AssertionError("info cannot be None") + n_stored = len(self.info) + if n_stored < _store_n_calls: + self.info.add_inputs(args, kwargs) + begin = time.perf_counter() + res = _captured_method(*args, **kwargs) + duration = time.perf_counter() - begin + if n_stored < _store_n_calls: + self.info.add_outputs(res, latency=duration) + return res + + def num_obs(self) -> int: + """Returns the number of stored set of inputs.""" + return 0 if not self.info else len(self.info) + + @contextlib.contextmanager + def __call__( + self, + model: torch.nn.Module, + store_n_calls: int = 3, + method_name: str = "forward", + ): + """Starts collecting inputs and outputs of a specific method. + The model method is replaced by a new one collecting tensors + before and after the inner one is called. + The original method is restored after the collection. + + Args: + model: Model + store_n_calls: The collection stops after this many calls + to avoid taking too much memory. + method_name: Method name to spy on. + """ + if not hasattr(model, method_name): + raise ValueError( + f"Model type {model} does not have a method {method_name!r}." + ) + captured_method = getattr(model, method_name) + sig = inspect.signature(captured_method) + if self.info is None: + kwargs_names = [ + p + for p in sig.parameters + if sig.parameters[p].kind == inspect.Parameter.VAR_KEYWORD + ] + args_names = [ + (p, i) + for (i, p) in enumerate(sig.parameters) + if sig.parameters[p].kind == inspect.Parameter.VAR_POSITIONAL + ] + self.info = InputObserverInfo( + signature_names=list(sig.parameters), + default_values={ + p.name: p.default + for p in sig.parameters.values() + if p.default != inspect.Parameter.empty + and isinstance(p.default, (int, bool, str, float)) + }, + value_if_missing=self.value_if_missing, + args_name_and_position=args_names[0] if args_names else None, + kwargs_name=kwargs_names[0] if kwargs_names else None, + ) + n_already_stored = len(self.info) + lambda_method = lambda *args, _cm=captured_method, _snc=( # noqa: E731 + store_n_calls + n_already_stored + ), **kwargs: self._replaced_method( + *args, _captured_method=_cm, _store_n_calls=_snc, **kwargs + ) + + # It may happen that the signature of the forward is used to trigger a preprocessing. + # This is used in GenerationMixin (transformers): + # position_ids_key = "decoder_position_ids" if ... else "position_ids" + # if position_ids_key in set(inspect.signature(self.forward).parameters.keys()): + lambda_method.__signature__ = sig # type: ignore[attr-defined] + + setattr(model, method_name, lambda_method) + + try: + yield self + finally: + setattr(model, method_name, captured_method) + + def _check_captured(self): + if self.info is None: + raise RuntimeError("No inputs were captured.") + + def infer_dynamic_shapes( + self, set_batch_dimension_for: set[int | str] | bool | None = None + ) -> tuple[dict[int, Any] | None, ...] | dict[str, dict[int, Any] | None]: + """ + Infers dynamic shapes. Most of the time, models do support a batch dimension + but this batch dimension has the same value for every input sample. + Instead of running inference on new samples, argument `set_batch_dimension_for` + can be used to tell the first dimension is a dynamic dimension for a particular + set of inputs referenced by their name (str) or their position (int). + + Args: + set_batch_dimension_for (set[int | str] | bool | None): A set of input + identifiers (by position as ``int`` or by name as ``str``) for + which the first dimension should be treated as a dynamic batch + dimension. If ``None``, no dimensions are explicitly marked as + dynamic. + """ + self._check_captured() + if self.info is None: + raise AssertionError("info must be non-None") + return self.info.infer_dynamic_shapes( + set_batch_dimension_for=set_batch_dimension_for + ) + + def infer_arguments( + self, + index_or_args_or_kwargs: tuple[Any] | dict[str, Any] | int | None = None, + flat: bool = False, + as_args_kwargs: bool = False, + ) -> ( + list[torch.Tensor | None] + | tuple[torch.Tensor, ...] + | dict[str, torch.Tensor] + | tuple[list[torch.Tensor] | tuple[torch.Tensor, ...], dict[str, torch.Tensor]] + ): + """Infers arguments based on the collected tensors. + + Args: + index_or_args_or_kwargs: If missing, the method selects one set of inputs + among the available ones, usually the set of inputs containing + with the highest number of tensors. + It then replaces None values and missing tensors with empty tensors. + If not missing, it can be an integer to fetch one of the stored set + or some inputs. + flat: If True, it returns a flattened list of tensors, + if False, it returns a tuple or a dictionary preserving + the nested structures. The flat version is used internally. + It produces a single list of tensors easier to process or modify + rather than a nested structure holding the same tensors. + The original structure can be restored with + ``torch.utils._pytree.tree_unflatten(flat_list, self.aligned_spec)``. + This mechanism is used to replace None values by empty tensors. + as_args_kwargs: If True, the method always returns `(args, kwargs)`, + otherwise, it returns either a tuple (only args) or a dictionary + (only kwargs) or raises an exception if it cannot do so. + Returns: + Inferred arguments, every optional tensor is replaced by an empty tensor. + """ + self._check_captured() + assert self.info is not None # noqa: S101 + index_or_candidate: int | InputCandidate | None = None + if index_or_args_or_kwargs is None or isinstance(index_or_args_or_kwargs, int): + index_or_candidate = index_or_args_or_kwargs + else: + if isinstance(index_or_args_or_kwargs, tuple): + index_or_candidate = InputCandidate( + args=index_or_args_or_kwargs, kwargs={}, clone=False, cst_kwargs={} + ) + elif isinstance(index_or_args_or_kwargs, dict): + index_or_candidate = InputCandidate( + args=(), + kwargs={ + k: v + for k, v in index_or_args_or_kwargs.items() + if k not in self.info.default_values + }, + clone=False, + cst_kwargs={ + k: v + for k, v in index_or_args_or_kwargs.items() + if k in self.info.default_values + }, + ) + else: + raise ValueError( + f"Unexpected type {type(index_or_args_or_kwargs)} " + f"for index_or_args_or_kwargs." + ) + self.info.align_inputs_none_values() + index_or_candidate.align_with( + # pyrefly: ignore[bad-argument-type] + self.info._best_candidate, + # pyrefly: ignore[bad-argument-type] + self.info._captured_inputs, + self.info.signature_names, + ) + return self.info.infer_arguments( + index_or_candidate, + flat=flat, + as_args_kwargs=as_args_kwargs, + ) + + def check_discrepancies( + self, + onnx_program: torch.onnx.ONNXProgram, + atol: float = 1e-4, + rtol: float = 0.1, + progress_bar: bool = False, + initializer: Callable[ + [str | bytes], ort.InferenceSession + ] = _onnx_program._ort_session_initializer, + skip_none: bool = True, + ) -> list[dict[str, str | int | float | bool]]: + """Computes the discrepancies between the saved inputs and outputs + with the saved onnx model. + + Args: + onnx_program: Exported Model to verify. + atol: Absolute tolerance, recommended values, 1e-4 for float, 1e-2 for float16. + rtol: Relative tolerance. + progress_bar: Shows a progress bar (requires `tqdm`). + initializer: The function called to initialize the ONNX Runtime inference + session with the specified model. By default, it uses the + `_ort_session_initializer` function. + skip_none: Does not check discrepancies when an output is None. + + Returns: + A list of dictionaries, ready to be consumed by a dataframe. + + The function catches exceptions, it shows the error in the returned + summary. + """ + # For big models, we should consider taking a filename to avoid the users + # creating the model proto twice. + self._check_captured() + assert self.info is not None # noqa: S101 + + onnx_program.initialize_inference_session(initializer) + + input_names = [i.name for i in onnx_program.model.graph.inputs] + io_sets = list( + zip(self.info.inputs, self.info.flat_outputs, self.info.latencies) + ) + if progress_bar: + from tqdm import tqdm + + loop = tqdm(io_sets) + else: + loop = io_sets + data: list[dict[str, Any]] = [] + for inputs, outputs, latency in loop: + assert inputs.aligned_flat_list is not None # noqa: S101 + if len(input_names) != len(inputs.aligned_flat_list): + raise RuntimeError( + f"There are ({len(inputs.aligned_flat_list)}) " + f"tensors but the model expects {len(input_names)}." + ) + n_none = sum(t is None for t in inputs.aligned_flat_list) + n_empty = sum(t is None or t.numel() == 0 for t in inputs.aligned_flat_list) + + feeds = dict(zip(input_names, self.info.infer_arguments(inputs, flat=True))) + + begin = time.perf_counter() + try: + ort_outputs = onnx_program(**feeds) + error = None + except Exception as e: + error = str(e) + ort_outputs = None + + duration = time.perf_counter() - begin + if error: + diff: dict[str, str | int | float | bool] = dict( + error=error, SUCCESS=False + ) + elif ort_outputs is None or len(outputs) != len(ort_outputs): + diff = dict(SUCCESS=False, error="not the same number of outputs") + else: + success = True + err_abs = 0.0 + err_rel = 0.0 + error = "" + for torch_tensor, ort_tensor in zip(outputs, ort_outputs): + if torch_tensor is None or ort_tensor is None: + if type(torch_tensor) is not type(ort_tensor) and not skip_none: + success = False + error = "missing output" + break + continue + if torch_tensor.shape != ort_tensor.shape: + success = False + error = "not the same shape" + break + if torch_tensor.dtype != ort_tensor.dtype: + success = False + error = "not the same type" + break + err = (torch_tensor - ort_tensor).abs().max().item() + err_abs = max(err_abs, err) + if err_abs > atol: + success = False + err = ( + ( + (torch_tensor - ort_tensor).abs() + / (torch_tensor.abs() + rtol) + ) + .max() + .item() + ) + err_rel = max(err_rel, err) + if err_rel > rtol: + success = False + diff = dict(SUCCESS=success, abs=err_abs, rel=err_rel) + diff.update( + dict( + index=len(data), + duration_torch=latency, + ort_duration=duration, + n_inputs=len(input_names), + n_none=n_none, + n_empty=n_empty, + ) + ) + data.append(diff) + onnx_program.release() + return data diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_ir_passes.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_ir_passes.py new file mode 100644 index 0000000000000000000000000000000000000000..ee8fca22437e01fbc77d37151a6ac3c94baaa1f9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_ir_passes.py @@ -0,0 +1,166 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import logging +import re +from typing import TYPE_CHECKING + +from torch.onnx._internal._lazy_import import onnx_ir as ir +from torch.onnx._internal.exporter import _constants + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# The opset domain for ONNX operators +_ONNX_DOMAIN = "" + +logger = logging.getLogger(__name__) + + +def rename_inputs(model: ir.Model, new_names: Sequence[str]) -> None: + unique_names = frozenset(new_names) + if len(unique_names) != len(new_names): + seen = set() + duplicates = [] + for name in new_names: + if name in seen: + duplicates.append(name) + seen.add(name) + raise ValueError(f"Input names cannot be duplicated: {duplicates}") + + for input, new_name in zip(model.graph.inputs, new_names): + input.metadata_props["pkg.torch.onnx.original_node_name"] = str(input.name) + input.name = new_name + + +def rename_outputs(model: ir.Model, new_names: Sequence[str]) -> None: + unique_names = frozenset(new_names) + if len(unique_names) != len(new_names): + seen = set() + duplicates = [] + for name in new_names: + if name in seen: + duplicates.append(name) + seen.add(name) + raise ValueError(f"Output names cannot be duplicated: {duplicates}") + + for output, new_name in zip(model.graph.outputs, new_names): + output.metadata_props["pkg.torch.onnx.original_node_name"] = str(output.name) + output.name = new_name + + +def _all_values(model: ir.Model): + """Yield all values in a model.""" + # Yield all values in the model + yield from model.graph.inputs + yield from model.graph.initializers.values() + for node in ir.traversal.RecursiveGraphIterator(model.graph): + yield from node.outputs + # Yield all values in functions + for function in model.functions.values(): + yield from function.inputs + for node in ir.traversal.RecursiveGraphIterator(function): + yield from node.outputs + + +def _replace_names(shape_expr: str, rename_mapping: dict[str, str]) -> str: + """Replace all known names in a shape expression with new names.""" + for old_name, new_name in rename_mapping.items(): + shape_expr = re.sub( + rf"(? None: + """Rename dynamic axes in a model according to the specified dynamic_axes names.""" + + # Create a mapping from string to string for easier replacement + string_mapping: dict[str, str] = {} + for key, value in tuple(rename_mapping.items()): + if isinstance(key, ir.SymbolicDim): + if isinstance(key.value, str): + string_mapping[key.value] = value + else: + raise ValueError( + f"Invalid SymbolicDim value in rename_mapping: {key.value!r}. " + "Expected str." + ) + elif isinstance(key, str): + string_mapping[key] = value + else: + raise ValueError( + f"Invalid key type in rename_mapping: {type(key)}({key!r}). Expected " + "str or ir.SymbolicDim." + ) + + # NOTE: Mapping needs to be sorted by length because the shape expression + # could have multiple ways to be expressed, for example, + # {"s1": sequence_length, "s11": "past_sequence_length", "s1 + s11": "masked_sequence_length"} + # We prefer the replacement starts from the longest match. + sorted_rename_mapping = dict( + sorted(string_mapping.items(), key=lambda item: len(item[0]), reverse=True) + ) + for value in _all_values(model): + if value.shape is None: + continue + new_shape = [] + changed = False + for dim in value.shape: + if not isinstance(dim, ir.SymbolicDim): + new_shape.append(dim) + continue + dim_name = dim.value + if dim_name in sorted_rename_mapping: + new_shape.append(sorted_rename_mapping[dim_name]) + changed = True + elif dim_name is not None: + # For example: "2*s1", "s1+1", "s1-1", "s1*s2", "s1/s2" + new_name = _replace_names(dim_name, sorted_rename_mapping) + new_shape.append(new_name) + if new_name != dim_name: + changed = True + else: + new_shape.append(None) + if changed: + value.shape = ir.Shape(new_shape) + + +def _maybe_set_opset_version( + opset_imports: dict[str, int], domain: str, version: int | None +) -> None: + """Set the opset version for the domain.""" + if domain in opset_imports and opset_imports[domain] != 1: + # Already set + return + if domain == _ONNX_DOMAIN: + opset_imports[domain] = _constants.TORCHLIB_OPSET + return + if version is None: + # We don't know the opset version, so set it to 1 + # This is valid for the custom function domains like "pkg.torch.__subgraph__" + opset_imports[domain] = 1 + return + # Set the known opset version for the domain + opset_imports[domain] = version + + +def add_opset_imports(model: ir.Model) -> None: + """Collect all opsets used and add opset imports to the model and functions.""" + for node in ir.traversal.RecursiveGraphIterator(model.graph): + domain = node.domain + _maybe_set_opset_version(model.opset_imports, domain, node.version) + + for function in model.functions.values(): + for node in ir.traversal.RecursiveGraphIterator(function): + domain = node.domain + _maybe_set_opset_version(function.opset_imports, domain, node.version) + for domain, version in function.opset_imports.items(): + # Add all opsets used in the function to the model, because ONNX Runtime + # does not handle adding the opset imports to the model after inlining during inference. + # This should happen after all opsets are collected for the function from its nodes. + _maybe_set_opset_version(model.opset_imports, domain, version) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_isolated.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_isolated.py new file mode 100644 index 0000000000000000000000000000000000000000..543a1e20d027f28975ecd3b137ac2d1f421eabf4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_isolated.py @@ -0,0 +1,65 @@ +"""Isolated calls to methods that may segfault.""" + +from __future__ import annotations + +import multiprocessing +import os +import warnings +from typing import Any, TYPE_CHECKING, TypeVar +from typing_extensions import ParamSpec, TypeVarTuple, Unpack + + +if TYPE_CHECKING: + from collections.abc import Callable + + +_P = ParamSpec("_P") +_R = TypeVar("_R") +_Ts = TypeVarTuple("_Ts") + +_IS_WINDOWS = os.name == "nt" + + +def _call_function_and_return_exception( + func: Callable[[Unpack[_Ts]], _R], args: tuple[Unpack[_Ts]], kwargs: dict[str, Any] +) -> _R | Exception: + """Call function and return a exception if there is one.""" + + try: + return func(*args, **kwargs) + except Exception as e: + return e + + +def safe_call(func: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs) -> _R: + """Call a function in a separate process. + + Args: + func: The function to call. + args: The positional arguments to pass to the function. + kwargs: The keyword arguments to pass to the function. + + Returns: + The return value of the function. + + Raises: + Exception: If the function raised an exception. + """ + if _IS_WINDOWS: + # On Windows, we cannot create a new process with fork. + warnings.warn( + f"A new process is not created for {func} on Windows.", stacklevel=1 + ) + return func(*args, **kwargs) + + with multiprocessing.get_context("fork").Pool(1) as pool: + # It is important to fork a process here to prevent the main logic from + # running again when the user does not place it under a `if __name__ == "__main__":` + # block. + result = pool.apply_async( + _call_function_and_return_exception, (func, args, kwargs) + ) + result = result.get(timeout=5) + if isinstance(result, Exception): + raise result + return result diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_onnx_program.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_onnx_program.py new file mode 100644 index 0000000000000000000000000000000000000000..5d7484d85bed55c5f646cd5877f2b019ce63a60a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_onnx_program.py @@ -0,0 +1,521 @@ +# mypy: allow-untyped-defs +# mypy: disable-error-code="attr-defined,name-defined" +from __future__ import annotations + + +__all__ = ["ONNXProgram"] + +import contextlib +import copy +import gc +import logging +import os +import tempfile +import textwrap +import warnings +from collections.abc import Callable, Sequence +from typing import Any, TYPE_CHECKING + +import torch +from torch.onnx._internal._lazy_import import onnx, onnx_ir as ir, onnxscript_apis +from torch.onnx._internal.exporter import _dynamic_shapes, _ir_passes +from torch.utils import _pytree + + +# NOTE: DO NOT import module from torch.onnx._internal to this module in the global scope +# because ONNXProgram is exposed to the public API + +if TYPE_CHECKING: + import numpy as np + import onnxruntime as ort + +_LARGE_MODEL_THRESHOLD = 1536 * 1024 * 1024 # 1536MB +_NP_UNSUPPORTED_DTYPES_8BIT = frozenset( + { + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + torch.float8_e5m2, + torch.float8_e5m2fnuz, + } +) + +logger = logging.getLogger(__name__) + + +def _ort_session_initializer(model: str | bytes) -> ort.InferenceSession: + """Initialize an ONNX Runtime inference session with the specified model.""" + import onnxruntime as ort + + session_options = ort.SessionOptions() + session_options.log_severity_level = 3 # 3: Error + possible_providers = ( + "CUDAExecutionProvider", + "CPUExecutionProvider", + ) + available_providers = set(ort.get_available_providers()) + providers = [ + provider for provider in possible_providers if provider in available_providers + ] + return ort.InferenceSession( + model, providers=providers, sess_options=session_options + ) + + +def _count_initializer_size(graph: ir.Graph) -> int: + """Count the total size of the initializers in bytes.""" + return sum( + v.const_value.nbytes + for v in graph.initializers.values() + if v.const_value is not None + ) + + +@contextlib.contextmanager +def _set_graph_outputs( + graph: ir.Graph, + outputs: list[ir.Value], +): + """Temporarily set the outputs of the graph. + + Args: + graph: The graph to set the outputs for. + outputs: The outputs to set. + """ + original_outputs = list(graph.outputs) + graph.outputs.clear() + graph.outputs.extend(outputs) + try: + yield + finally: + graph.outputs.clear() + graph.outputs.extend(original_outputs) + + +def _create_value_mapping(graph: ir.Graph) -> dict[str, ir.Value]: + """Return a dictionary mapping names to values in the graph. + + The mapping does not include values from subgraphs. + + Args: + graph: The graph to extract the mapping from. + + Returns: + A dictionary mapping names to values. + """ + values: dict[str, ir.Value] = {} + values.update(graph.initializers) + # The names of the values can be None or "", which we need to exclude + for input in graph.inputs: + if not input.name: + continue + values[input.name] = input + for node in graph: + for value in node.outputs: + if not value.name: + continue + values[value.name] = value + return values + + +def _to_numpy_array(input: torch.Tensor | int | float | str | bool) -> np.ndarray: + if isinstance(input, (int, float, str, bool)): + return ir.tensor(input).numpy() + + from torch.onnx._internal.exporter import _core + + return _core.TorchTensor(input).numpy() + + +def _from_numpy_array(array: np.ndarray) -> torch.Tensor: + """Convert a NumPy array to a PyTorch tensor.""" + import ml_dtypes # type: ignore[import-not-found] + import numpy as np + + if array.dtype == ml_dtypes.bfloat16: + return torch.from_numpy(array.view(np.uint16)).view(torch.bfloat16) + if array.dtype == ml_dtypes.float8_e4m3fn: + return torch.from_numpy(array.view(np.uint8)).view(torch.float8_e4m3fn) + if array.dtype == ml_dtypes.float8_e4m3fnuz: + return torch.from_numpy(array.view(np.uint8)).view(torch.float8_e4m3fnuz) + if array.dtype == ml_dtypes.float8_e5m2: + return torch.from_numpy(array.view(np.uint8)).view(torch.float8_e5m2) + if array.dtype == ml_dtypes.float8_e5m2fnuz: + return torch.from_numpy(array.view(np.uint8)).view(torch.float8_e5m2fnuz) + return torch.from_numpy(array) + + +def _to_ort_value(input: torch.Tensor | int | float | str | bool) -> ort.OrtValue: + """Convert a PyTorch tensor to an ONNX Runtime OrtValue.""" + import numpy as np + import onnxruntime as ort + + from torch.onnx._internal.exporter import _core + + if isinstance(input, (int, float, str, bool)): + # Convert scalar values to OrtValue + dtype_mapping = { + int: np.int64, + float: np.float32, + } + # pyrefly: ignore [bad-argument-type, no-matching-overload] + dtype = dtype_mapping.get(type(input)) + return ort.OrtValue.ortvalue_from_numpy(np.array(input, dtype=dtype)) + + if input.dtype == torch.bfloat16 or input.dtype in _NP_UNSUPPORTED_DTYPES_8BIT: + if hasattr(ort.OrtValue, "ortvalue_from_numpy_with_onnx_type"): + # This requires ONNX Runtime 1.21 or newer + if input.dtype == torch.bfloat16: + uint_type = torch.uint16 + else: + uint_type = torch.uint8 + onnx_type = _core.torch_dtype_to_onnx_dtype(input.dtype) + # Make tensor contiguous to ensure view() works + input = input.contiguous() + return ort.OrtValue.ortvalue_from_numpy_with_onnx_type( + input.view(uint_type).numpy(force=True), onnx_element_type=onnx_type + ) + raise RuntimeError( + f"Failed to convert tensor of type '{input.dtype}' to OrtValue. " + "Please ensure that ONNX Runtime is built with DLPack support or is the latest version" + ) + # TODO(#151064): Use dlpack when ORT properly supports it + return ort.OrtValue.ortvalue_from_numpy(input.numpy(force=True)) + + +def _from_ort_value(value: ort.OrtValue) -> torch.Tensor: + if value.element_type() in ( + ir.DataType.BFLOAT16, + ir.DataType.FLOAT8E4M3FN, + ir.DataType.FLOAT8E4M3FNUZ, + ir.DataType.FLOAT8E5M2, + ir.DataType.FLOAT8E5M2FNUZ, + ): + # This requires ONNX Runtime 1.21 or newer + try: + return torch.from_dlpack(value._get_c_value()) + except Exception as e: + raise RuntimeError( + "Failed to convert OrtValue to torch.Tensor. " + "Please ensure that ONNX Runtime is built with DLPack support or is the latest version" + ) from e + return torch.from_numpy(value.numpy()) + + +class ONNXProgram: + """A class to represent an ONNX program that is callable with torch tensors. + + Attributes: + model: The ONNX model as an ONNX IR model object. + exported_program: The exported program that produced the ONNX model. + """ + + def __init__( + self, model: ir.Model, exported_program: torch.export.ExportedProgram | None + ) -> None: + """Initialize the ONNX program with the specified model and exported program. + Args: + model: The ONNX model. + exported_program: The exported program that produced the ONNX model. Optional. + """ + self.model: ir.Model = model + self.exported_program = exported_program + self._inference_session: ort.InferenceSession | None = None + self._tempdir: tempfile.TemporaryDirectory | None = None + # Strategy used to capture the exported program + self._capture_strategy: str | None = None + + def __repr__(self) -> str: + return f"""\ +ONNXProgram( + model= +{textwrap.indent(str(self.model), " " * 8)} + , + exported_program= +{textwrap.indent(str(self.exported_program), " " * 8)} +) +""" + + def __call__(self, *args, **kwargs) -> Sequence[torch.Tensor]: + """Run the ONNX model with the same arguments you would provide to the GraphModule.""" + import onnxruntime as ort + + flatten_args = _process_args(args, kwargs) + + if self._inference_session is None: + self.initialize_inference_session() + + if self._inference_session is None: + raise AssertionError("_inference_session must be non-None") + + ort_input = { + k.name: _to_ort_value(v) + for k, v in zip(self.model.graph.inputs, flatten_args) + } + run_options = ort.RunOptions() + run_options.log_severity_level = 3 # 3: Error + logger.debug("Running the inference session with %s arguments.", len(ort_input)) + outputs = self._inference_session.run_with_ort_values( + None, ort_input, run_options=run_options + ) + logger.debug("Inference session run completed.") + return tuple(_from_ort_value(output) for output in outputs) + + def call_reference(self, *args, **kwargs) -> Sequence[torch.Tensor]: + """Run the ONNX model using the reference backend.""" + import onnx.reference + + evaluator = onnx.reference.ReferenceEvaluator(self.model_proto) + + flatten_args = _process_args(args, kwargs) + ref_input = { + k.name: _to_numpy_array(v) + for k, v in zip(self.model.graph.inputs, flatten_args) + } + outputs = evaluator.run(None, ref_input) # type: ignore[arg-type] + if not isinstance(outputs, Sequence): + raise AssertionError(f"Expected Sequence, got {type(outputs)}") + return tuple(_from_numpy_array(output) for output in outputs) + + def compute_values( + self, value_names: Sequence[str], args=(), kwargs=None + ) -> Sequence[torch.Tensor]: + """Compute the values of the specified names in the ONNX model. + + This method is used to compute the values of the specified names in the ONNX model. + The values are returned as a dictionary mapping names to tensors. + + Args: + value_names: The names of the values to compute. + + Returns: + A dictionary mapping names to tensors. + """ + if kwargs is None: + kwargs = {} + self.release() + values = _create_value_mapping(self.model.graph) + for name in value_names: + if name not in values: + raise ValueError( + f"Value '{name}' not found in the model. " + "Please provide a valid value name." + ) + temporary_outputs = [values[name] for name in value_names] + with _set_graph_outputs(self.model.graph, temporary_outputs): + try: + result = self(*args, **kwargs) + finally: + self.release() + return result + + @property + def model_proto(self) -> onnx.ModelProto: + """Return the ONNX ``ModelProto`` object.""" + return ir.serde.serialize_model(self.model) + + def optimize(self) -> None: + """Optimize the ONNX model. + + This method optimizes the ONNX model by performing constant folding and + eliminating redundancies in the graph. The optimization is done in-place. + """ + self.model = onnxscript_apis.optimize(self.model) + + def save( + self, + destination: str | os.PathLike, + *, + include_initializers: bool = True, + keep_initializers_as_inputs: bool = False, + external_data: bool | None = None, + ) -> None: + """Save the ONNX model to the specified destination. + + When ``external_data`` is ``True`` or the model is larger than 2GB, + the weights are saved as external data in a separate file. + + Initializer (model weights) serialization behaviors: + + * ``include_initializers=True``, ``keep_initializers_as_inputs=False`` (default): + The initializers are included in the saved model. + * ``include_initializers=True``, ``keep_initializers_as_inputs=True``: + The initializers are included in the saved model and kept as model inputs. + Choose this option if you want the ability to override the model weights + during inference. + * ``include_initializers=False``, ``keep_initializers_as_inputs=False``: + The initializers are not included in the saved model and are not listed + as model inputs. Choose this option if you want to attach the initializers + to the ONNX model in a separate, post-processing, step. + * ``include_initializers=False``, ``keep_initializers_as_inputs=True``: + The initializers are not included in the saved model but are listed as model + inputs. Choose this option if you want to supply the initializers during + inference and want to minimize the size of the saved model. + + Args: + destination: The path to save the ONNX model to. + include_initializers: Whether to include the initializers in the saved model. + keep_initializers_as_inputs: Whether to keep the initializers as inputs in the saved model. + If `True`, the initializers are added as inputs to the model which means they can be overwritten. + by providing the initializers as model inputs. + external_data: Whether to save the weights as external data in a separate file. + + Raises: + TypeError: If ``external_data`` is ``True`` and ``destination`` is not a file path. + """ + original_initializers = copy.copy(self.model.graph.initializers) + original_inputs = copy.copy(self.model.graph.inputs) + + # Adjust the model based on options + if not include_initializers: + self.model.graph.initializers.clear() + if keep_initializers_as_inputs: + self.model.graph.inputs.extend(original_initializers.values()) # type: ignore[arg-type] + + try: + # Save the model to disk + if ( + external_data + or _count_initializer_size(self.model.graph) > _LARGE_MODEL_THRESHOLD + ): + onnxscript_apis.save_model_with_external_data(self.model, destination) + else: + ir.save(self.model, destination) + finally: + # Revert the changes to the model + if not include_initializers: + self.model.graph.initializers.update(original_initializers) + if keep_initializers_as_inputs: + self.model.graph.inputs.clear() + self.model.graph.inputs.extend(original_inputs) + + def apply_weights(self, state_dict: dict[str, torch.Tensor]) -> None: + """Apply the weights from the specified state dict to the ONNX model. + + Use this method to replace FakeTensors or other weights. + + Args: + state_dict: The state dict containing the weights to apply to the ONNX model. + """ + from torch.onnx._internal.exporter import _core + + for name, tensor in state_dict.items(): + if name in self.model.graph.initializers: + self.model.graph.initializers[name].const_value = _core.TorchTensor( + tensor, name + ) + else: + warnings.warn( + f"Weight '{name}' not found in the model. Skipped applying.", + category=torch.onnx.errors.OnnxExporterWarning, + stacklevel=1, + ) + + def initialize_inference_session( + self, + initializer: Callable[ + [str | bytes], ort.InferenceSession + ] = _ort_session_initializer, + ) -> None: + """Initialize the ONNX Runtime inference session. + + Args: + initializer: The function to initialize the ONNX Runtime inference + session with the specified model. By default, it uses the + :func:`_ort_session_initializer` function. + """ + # TODO(justinchuby): Allow different inference options + logger.debug("Initializing the inference session.") + if ( + byte_size := _count_initializer_size(self.model.graph) + ) > _LARGE_MODEL_THRESHOLD: + logger.debug("The model initializers is larger than 1.5GB (%s).", byte_size) + # Save the model to a temporary file if too large + self._tempdir = tempfile.TemporaryDirectory(ignore_cleanup_errors=True) + model_path = os.path.join(self._tempdir.name, "model.onnx") + self.save(model_path, external_data=True) + model = model_path + else: + model = self.model_proto.SerializeToString() # type: ignore[assignment] + + self._inference_session = initializer(model) + logger.debug("Inference session initialized.") + + def release(self) -> None: + """Release the inference session. + + You may call this method to release the resources used by the inference session. + """ + # Release the inference session first so that the model file can be deleted + if self._inference_session is not None: + self._inference_session = None + gc.collect() + if self._tempdir is not None: + self._tempdir.cleanup() + self._tempdir = None + + def rename_axes(self, rename_mapping: dict[str | ir.SymbolicDim, str]) -> None: + """Rename axes in a model according to the specified rename mapping. + + Example:: + + batch = onnx_program.model.graph.inputs[0].shape[0] + seq_len = onnx_program.model.graph.inputs[0].shape[2] + rename_mapping = { + batch: "batch", + seq_len: "seq_len", + } + onnx_program.rename_axes(rename_mapping) + + Args: + rename_mapping: A dictionary mapping old axes to new axis names. + Keys can be either: + + * String axis names (e.g., "s1", "s2") + * SymbolicDim objects obtained from the model + (e.g., onnx_program.model.graph.inputs[0].shape[0]) + + Values must be strings representing the new axis names. + """ + _ir_passes.rename_axis(self.model, rename_mapping) + + def _rename_dynamic_axes( + self, + dynamic_shapes: dict[str, Any] | tuple[Any, ...] | list[Any], + ) -> None: + """Rename dynamic axes in a model according to the specified dynamic_axes names.""" + rename_mapping = _dynamic_shapes.create_rename_mapping( + self.model.graph.inputs, dynamic_shapes + ) + _ir_passes.rename_axis(self.model, rename_mapping) + + +def _process_args(args, kwargs) -> tuple[torch.Tensor, ...]: + """Process input arguments for the ONNX model.""" + args = _flatten_inputs(args, kwargs) + args = _remove_none_from_inputs(args) + args = _convert_complex_to_real_representation(args) + return args + + +def _flatten_inputs(model_args, model_kwargs): + flattened_args, _ = _pytree.tree_flatten((model_args, model_kwargs)) + return flattened_args + + +def _remove_none_from_inputs(model_args): + return tuple(arg for arg in model_args if arg is not None) + + +def _convert_complex_to_real_representation(model_args): + """Convert complex dtype tensors to real representation tensors. + + ONNX does not support complex dtype tensors. Thus, we convert complex dtype tensors + to real representation tensors (i.e., float dtype tensors with an extra dimension + representing the real and imaginary parts of the complex number). + """ + return tuple( + torch.view_as_real(arg.resolve_conj()) + if isinstance(arg, torch.Tensor) and arg.is_complex() + else arg + for arg in model_args + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_registration.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_registration.py new file mode 100644 index 0000000000000000000000000000000000000000..d22186944e1fe7ac7fc19912405667c7301fe87c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_registration.py @@ -0,0 +1,303 @@ +"""Module for handling ATen to ONNX functions registration. + +https://github.com/pytorch/pytorch/blob/6aa5bb1a76dee8112f1a9e7c194c790b5cdc6462/torch/onnx/_internal/fx/registration.py +""" + +# NOTE: Why do we need a different registry than the one in torchlib? +# The registry in torchlib is used to register functions that are already implemented in +# torchlib, and is designed to be a static singleton. It does not take into account custom ops or different +# opsets etc. The registry implemented for the exporter is designed to be modifiable at +# export time by users, and is designed with dispatching in mind. + +# mypy: allow-untyped-defs +from __future__ import annotations + +import dataclasses +import importlib.util +import logging +import math +import operator +import types +from collections.abc import Callable +from typing import Literal, TypeAlias + +import torch +import torch._ops +from torch.onnx._internal._lazy_import import onnx_ir as ir, onnxscript, onnxscript_apis +from torch.onnx._internal.exporter import _constants, _schemas +from torch.onnx._internal.exporter._torchlib import _torchlib_registry + + +TorchOp: TypeAlias = torch._ops.OpOverload | types.BuiltinFunctionType | Callable + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class OnnxDecompMeta: + """A wrapper of onnx-script function with additional metadata. + + onnx_function: The onnx-script function from torchlib. + fx_target: The PyTorch node callable target. + signature: The ONNX signature of the function. When None, the signature is inferred. + is_custom: Whether the function is a custom function. + is_complex: Whether the function is a function that handles complex valued inputs. + opset_introduced: + The ONNX opset version in which the function was introduced. + Its specifies the minimum ONNX opset version required to use the function. + device: The device the function is registered to. If None, it is registered to all devices. + skip_signature_inference: Whether to skip signature inference for the function. + """ + + onnx_function: Callable + fx_target: TorchOp + signature: ir.schemas.OpSignature | None + is_custom: bool = False + is_complex: bool = False + opset_introduced: int = 18 + device: Literal["cuda", "cpu"] | str | None = None # noqa: PYI051 + skip_signature_inference: bool = False + + def __post_init__(self) -> None: + if self.signature is None and not self.skip_signature_inference: + try: + if isinstance(self.onnx_function, onnxscript.OnnxFunction): + signature = _schemas.op_signature_from_function( + self.onnx_function, + self.onnx_function.function_ir.domain, + self.onnx_function.name, + since_version=self.onnx_function.opset.version, + ) + else: + signature = _schemas.op_signature_from_function( + self.onnx_function, "__traced", self.onnx_function.__name__ + ) + except Exception as e: + # Log an warning if the op is custom. Raise exception for builtin ops. + if not self.is_custom: + raise + else: + # When the function is targeting an HOP, for example, it will accept + # functions as arguments and fail to generate an ONNX signature. + # In this case we set signature to None and dispatch to this function always. + logger.warning( # noqa: G200 + "Failed to infer the signature for function '%s' because '%s'" + "All nodes targeting `%s` will be dispatched to this function", + self.onnx_function, + e, + self.fx_target, + ) + else: + self.signature = signature + self.onnx_function._pt_onnx_signature = signature # type: ignore[attr-defined] + + +def _get_overload(qualified_name: str) -> torch._ops.OpOverload | None: + """Obtain the torch op from ::[.]""" + # TODO(justinchuby): Handle arbitrary custom ops + namespace, opname_overload = qualified_name.split("::") + op_name, *maybe_overload = opname_overload.split(".", 1) + if namespace == "_operator": + # Builtin functions + return getattr(operator, op_name) + if namespace == "math": + return getattr(math, op_name) + if namespace == "torchvision": + if importlib.util.find_spec("torchvision") is None: + logger.warning("torchvision is not installed. Skipping %s", qualified_name) + return None + try: + op_packet = getattr(getattr(torch.ops, namespace), op_name) + if maybe_overload: + overload = maybe_overload[0] + elif "default" in op_packet._overload_names or "" in op_packet._overload_names: + # Has a default overload + overload = "default" + else: + logger.warning( + "'%s' does not have a 'default' overload. This could be an error in specifying the op name. Ignoring.", + qualified_name, + stacklevel=1, + ) + return None + + return getattr(op_packet, overload) # type: ignore[call-overload] + except AttributeError: + if qualified_name.endswith("getitem"): + # This is a special case where we registered the function incorrectly, + # but for BC reasons (pt<=2.4) we need to keep it. + return None + logger.info("'%s' is not found in this version of PyTorch.", qualified_name) + return None + except Exception: + logger.exception("Failed to find torch op '%s'", qualified_name) + return None + + +class ONNXRegistry: + """Registry for ONNX functions. + + The registry maintains a mapping from qualified names to symbolic functions under a + fixed opset version. It supports registering custom onnx-script functions and for + dispatcher to dispatch calls to the appropriate function. + + """ + + def __init__(self) -> None: + """Initializes the registry""" + self._opset_version = _constants.TORCHLIB_OPSET + self.functions: dict[TorchOp | str, list[OnnxDecompMeta]] = {} + + @property + def opset_version(self) -> int: + """The ONNX opset version the exporter should target.""" + return self._opset_version + + @classmethod + def from_torchlib(cls, opset_version=_constants.TORCHLIB_OPSET) -> ONNXRegistry: + """Populates the registry with ATen functions from torchlib. + + Args: + torchlib_registry: The torchlib registry to use for populating the registry. + """ + registry = cls() + registry._opset_version = opset_version + for meta in _torchlib_registry.get_torchlib_ops(): + registry._register(meta.fx_target, meta) + + # TODO(justinchuby): Remove this once torchlib is migrated to PyTorch + torchlib_ops = onnxscript_apis.get_torchlib_ops() + + for torchlib_meta in torchlib_ops: + qualified_name = torchlib_meta.qualified_name + overload_func = torchlib_meta.function + try: + # NOTE: This is heavily guarded with try-except because we don't want + # to fail the entire registry population if one function fails. + target = _get_overload(qualified_name) + if target is None: + continue + + meta = OnnxDecompMeta( + onnx_function=overload_func, + fx_target=target, + signature=None, + is_custom=False, + is_complex=torchlib_meta.is_complex, + ) + registry._register(target, meta) + except Exception: + logger.exception("Failed to register '%s'. Skipped", qualified_name) + continue + + registry._cleanup_registry_based_on_opset_version() + return registry + + def _register( + self, + target: TorchOp, + onnx_decomposition: OnnxDecompMeta, + ) -> None: + """Registers a OnnxDecompMeta to an operator. + + Args: + target: The PyTorch node callable target. + onnx_decomposition: The OnnxDecompMeta to register. + """ + target_or_name: str | TorchOp + if isinstance(target, torch._ops.OpOverload): + # Get the qualified name of the aten op because torch._ops.OpOverload lookup in + # a dictionary is unreliable for some reason. + target_or_name = target.name() + else: + target_or_name = target + if onnx_decomposition.is_custom: + self.functions.setdefault(target_or_name, []).insert(0, onnx_decomposition) + else: + self.functions.setdefault(target_or_name, []).append(onnx_decomposition) + + def register_op( + self, + target: TorchOp, + function: Callable, + is_complex: bool = False, + ) -> None: + """Registers a custom operator: torch.ops.... + + Args: + target: The PyTorch node callable target. + function: The onnx-script function to register. + is_complex: Whether the function is a function that handles complex valued inputs. + """ + if isinstance(target, torch._ops.OpOverloadPacket): + raise TypeError( + f"Target '{target}' should be provided as an OpOverload instead of an " + "OpOverloadPacket. You can get the default overload with " + ".default" + ) + + self._register( + target, + OnnxDecompMeta( + onnx_function=function, + fx_target=target, + signature=None, + is_custom=True, + is_complex=is_complex, + ), + ) + + def get_decomps(self, target: TorchOp) -> list[OnnxDecompMeta]: + """Returns a list of OnnxDecompMeta for the given op: torch.ops.... + + The list is ordered by the time of registration. The custom operators should come + first in the list. + + Args: + target: The PyTorch node callable target. + Returns: + A list of OnnxDecompMeta corresponding to the given name, or None if + the name is not in the registry. + """ + target_or_name: str | TorchOp + if isinstance(target, torch._ops.OpOverload): + # Get the qualified name of the aten op because torch._ops.OpOverload lookup in + # a dictionary is unreliable for some reason. + target_or_name = target.name() + else: + target_or_name = target + decomps = self.functions.get(target_or_name, []) + return sorted(decomps, key=lambda x: x.is_custom, reverse=True) + + def is_registered(self, target: TorchOp) -> bool: + """Returns whether the given op is registered: torch.ops.... + + Args: + target: The PyTorch node callable target. + + Returns: + True if the given op is registered, otherwise False. + """ + return bool(self.get_decomps(target)) + + def _cleanup_registry_based_on_opset_version(self) -> None: + """Pick the implementation with the highest opset version valid until the current opset version.""" + cleaned_functions = {} + for target_or_name, decomps in self.functions.items(): + # Filter decompositions to only include those with opset_introduced <= opset_version + decomps = [d for d in decomps if d.opset_introduced <= self.opset_version] + + # Keep only the decomposition with the highest opset_introduced + if decomps: + # Find the maximum opset_introduced + max_opset = max(d.opset_introduced for d in decomps) + + # Keep all decompositions with the maximum opset_introduced + cleaned_functions[target_or_name] = [ + d for d in decomps if d.opset_introduced == max_opset + ] + + self.functions = cleaned_functions + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(functions={self.functions})" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_reporting.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_reporting.py new file mode 100644 index 0000000000000000000000000000000000000000..cef2fbcc8a436af55925f89800e57cdebc533747 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_reporting.py @@ -0,0 +1,207 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import dataclasses +import re +from typing import TYPE_CHECKING + +from torch.onnx._internal.exporter import _analysis, _registration, _verification + + +if TYPE_CHECKING: + import os + + import onnx_ir as ir # type: ignore[import-untyped] + + import torch + + +@dataclasses.dataclass +class ExportStatus: + # Whether torch.export.export(..., strict=True) succeeds + torch_export_strict: bool | None = None + # Whether torch.export.export(..., strict=False) succeeds + torch_export_non_strict: bool | None = None + # Whether torch.export.draft_export() succeeds + torch_export_draft_export: bool | None = None + # Whether decomposition succeeds + decomposition: bool | None = None + # Whether ONNX translation succeeds + onnx_translation: bool | None = None + # Whether ONNX model passes onnx.checker.check_model + onnx_checker: bool | None = None + # Whether ONNX model runs successfully with ONNX Runtime + onnx_runtime: bool | None = None + # Whether the output of the ONNX model is accurate + output_accuracy: bool | None = None + + +def _status_emoji(status: bool | None) -> str: + if status is None: + return "⚪" + return "✅" if status else "❌" + + +def _format_export_status(status: ExportStatus) -> str: + return ( + f"```\n" + f"{_status_emoji(status.torch_export_non_strict)} Obtain model graph with `torch.export.export(..., strict=False)`\n" + f"{_status_emoji(status.torch_export_strict)} Obtain model graph with `torch.export.export(..., strict=True)`\n" + f"{_status_emoji(status.torch_export_draft_export)} Obtain model graph with `torch.export.draft_export`\n" + f"{_status_emoji(status.decomposition)} Decompose operators for ONNX compatibility\n" + f"{_status_emoji(status.onnx_translation)} Translate the graph into ONNX\n" + f"{_status_emoji(status.onnx_checker)} Run `onnx.checker` on the ONNX model\n" + f"{_status_emoji(status.onnx_runtime)} Execute the model with ONNX Runtime\n" + f"{_status_emoji(status.output_accuracy)} Validate model output accuracy\n" + f"```\n\n" + ) + + +def _strip_color_from_string(text: str) -> str: + # This regular expression matches ANSI escape codes + # https://github.com/pytorch/pytorch/blob/9554a9af8788c57e1c5222c39076a5afcf0998ae/torch/_dynamo/utils.py#L2785-L2788 + ansi_escape = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]") + return ansi_escape.sub("", text) + + +def _format_exported_program(exported_program: torch.export.ExportedProgram) -> str: + # Adapted from https://github.com/pytorch/pytorch/pull/128476 + # to remove colors + # Even though we can call graph_module.print_readable directly, since the + # colored option was added only recently, we can't guarantee that the + # version of PyTorch used by the user has this option. Therefore, we + # still call str(ExportedProgram) + text = f"```python\n{_strip_color_from_string(str(exported_program))}\n```\n\n" + return text + + +def construct_report_file_name(timestamp: str, status: ExportStatus) -> str: + # Status could be None. So we need to check for False explicitly. + if not ( + status.torch_export_non_strict + or status.torch_export_strict + or status.torch_export_draft_export + ): + # All strategies failed + postfix = "pt_export" + elif status.decomposition is False: + postfix = "decomp" + elif status.onnx_translation is False: + postfix = "conversion" + elif status.onnx_checker is False: + postfix = "checker" + elif status.onnx_runtime is False: + postfix = "runtime" + elif status.output_accuracy is False: + postfix = "accuracy" + elif ( + status.torch_export_strict is False + or status.torch_export_non_strict is False + or status.torch_export_draft_export is False + ): + # Some strategies failed + postfix = "strategies" + else: + postfix = "success" + return f"onnx_export_{timestamp}_{postfix}.md" + + +def format_decomp_comparison( + pre_decomp_unique_ops: set[str], + post_decomp_unique_ops: set[str], +) -> str: + """Format the decomposition comparison result. + + Args: + unique_ops_in_a: The unique ops in the first program. + unique_ops_in_b: The unique ops in the second program. + + Returns: + The formatted comparison result. + """ + return ( + f"Ops exist only in the ExportedProgram before decomposition: `{sorted(pre_decomp_unique_ops)}`\n\n" + f"Ops exist only in the ExportedProgram after decomposition: `{sorted(post_decomp_unique_ops)}`\n" + ) + + +def format_verification_infos( + verification_infos: list[_verification.VerificationInfo], +) -> str: + """Format the verification result. + + Args: + verification_infos: The verification result. + + Returns: + The formatted verification result. + """ + return "\n".join( + f"`{info.name}`: `max_abs_diff={info.max_abs_diff:e}`, `max_rel_diff={info.max_rel_diff:e}`, " + f"`abs_diff_hist={info.abs_diff_hist}`, `rel_diff_hist={info.rel_diff_hist}`" + for info in verification_infos + ) + + +def create_torch_export_error_report( + filename: str | os.PathLike, + formatted_traceback: str, + *, + export_status: ExportStatus, + profile_result: str | None, +) -> None: + with open(filename, "w", encoding="utf-8") as f: + f.write("# PyTorch ONNX Conversion Error Report\n\n") + f.write(_format_export_status(export_status)) + f.write("Error message:\n\n") + f.write("```pytb\n") + f.write(formatted_traceback) + f.write("```\n\n") + if profile_result is not None: + f.write("## Profiling result\n\n") + f.write("```\n") + f.write(profile_result) + f.write("```\n") + + +def create_onnx_export_report( + filename: str | os.PathLike, + formatted_traceback: str, + program: torch.export.ExportedProgram, + *, + decomp_comparison: str | None = None, + export_status: ExportStatus, + profile_result: str | None, + model: ir.Model | None = None, + registry: _registration.ONNXRegistry | None = None, + verification_result: str | None = None, +) -> None: + with open(filename, "w", encoding="utf-8") as f: + f.write("# PyTorch ONNX Conversion Report\n\n") + f.write(_format_export_status(export_status)) + f.write("## Error messages\n\n") + f.write("```pytb\n") + f.write(formatted_traceback) + f.write("\n```\n\n") + f.write("## Exported program\n\n") + f.write(_format_exported_program(program)) + if model is not None: + f.write("## ONNX model\n\n") + f.write("```python\n") + f.write(str(model)) + f.write("\n```\n\n") + f.write("## Analysis\n\n") + _analysis.analyze(program, file=f, registry=registry) + if decomp_comparison is not None: + f.write("\n## Decomposition comparison\n\n") + f.write(decomp_comparison) + f.write("\n") + if verification_result is not None: + f.write("\n## Verification results\n\n") + f.write(verification_result) + f.write("\n") + if profile_result is not None: + f.write("\n## Profiling result\n\n") + f.write("```\n") + f.write(profile_result) + f.write("```\n") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_schemas.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_schemas.py new file mode 100644 index 0000000000000000000000000000000000000000..9b702baeb6351529b6059feb8373a805fc24914a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_schemas.py @@ -0,0 +1,332 @@ +# mypy: allow-untyped-defs +"""Helpers for constructing ONNX operator signatures from Python functions.""" + +from __future__ import annotations + +import collections.abc +import inspect +import logging +import types +import typing +from collections.abc import Sequence +from typing import Any, Optional, TypeVar, Union + +from torch.onnx._internal._lazy_import import onnx_ir as ir, onnxscript + + +logger = logging.getLogger(__name__) + + +# Map from python type to corresponding ONNX AttributeProto type +_PY_TYPE_TO_ATTR_TYPE = { + float: ir.AttributeType.FLOAT, + int: ir.AttributeType.INT, + str: ir.AttributeType.STRING, + bool: ir.AttributeType.INT, + ir.Tensor: ir.AttributeType.TENSOR, + ir.TensorProtocol: ir.AttributeType.TENSOR, + ir.Graph: ir.AttributeType.GRAPH, + ir.GraphProtocol: ir.AttributeType.GRAPH, +} + +# Map from python type to corresponding ONNX AttributeProto type, +# for repeated (i.e., list of) values +_LIST_TYPE_TO_ATTR_TYPE = { + float: ir.AttributeType.FLOATS, + int: ir.AttributeType.INTS, + str: ir.AttributeType.STRINGS, + bool: ir.AttributeType.INTS, + ir.Tensor: ir.AttributeType.TENSORS, + ir.TensorProtocol: ir.AttributeType.TENSORS, + ir.Graph: ir.AttributeType.GRAPHS, + ir.GraphProtocol: ir.AttributeType.GRAPHS, +} + +_ALL_VALUE_TYPES = ( + {ir.TensorType(dtype) for dtype in ir.DataType} + | {ir.SequenceType(ir.TensorType(dtype)) for dtype in ir.DataType} + | {ir.OptionalType(ir.TensorType(dtype)) for dtype in ir.DataType} +) + +# TypeAnnotationValue represents the (value of) valid type-annotations recognized +# by ONNX Script. Currently, it supports +# - float, int, str (primitive attribute types) +# - Sequence[float], Sequence[int], Sequence[str] (attribute types) +# - Tensor types +# - Sequence[Tensor] types +# - Union of above 2 +# - TypeVars with above bounds +# - Above types with annotation attached +TypeAnnotationValue = Any + + +def _is_optional(type_: type) -> bool: + """Returns whether a type_ is an Optional.""" + origin_type = typing.get_origin(type_) + if origin_type is Union and type(None) in typing.get_args(type_): + # Python < 3.10 + return True + if origin_type is Optional: + # Python >= 3.10 + return True + if ( + hasattr(types, "UnionType") + and origin_type is types.UnionType + and type(None) in typing.get_args(type_) + ): + # Python >= 3.10 + return True + return False + + +def _get_attr_type(type_: type) -> ir.AttributeType: + """Obtain the type of the attribute from a Python class.""" + try: + if type_ in _PY_TYPE_TO_ATTR_TYPE: + return _PY_TYPE_TO_ATTR_TYPE[type_] + origin_type = typing.get_origin(type_) + if origin_type is None: + return ir.AttributeType.UNDEFINED + if origin_type in ( + collections.abc.Sequence, + Sequence, + list, + list, + tuple, + tuple, + ): + inner_type = typing.get_args(type_)[0] + if inner_type in _LIST_TYPE_TO_ATTR_TYPE: + return _LIST_TYPE_TO_ATTR_TYPE[inner_type] + except TypeError: + logger.warning("TypeError when checking %s.", type_, exc_info=True) + return ir.AttributeType.UNDEFINED + + +def _get_type_constraint_name(type_: TypeAnnotationValue) -> str | None: + """Returns the name of the type constraint for a given type annotation. + + Args: + type_: A Python type. + + Returns: + The name of the type constraint if it is a TypeVar. + - Prefixes the name with "Sequence_" if the type annotation is a Sequence[]. + """ + if isinstance(type_, TypeVar): + return type_.__name__ + if _is_optional(type_): + subtypes = typing.get_args(type_) + for subtype in subtypes: + if subtype is type(None): + continue + type_param_name = _get_type_constraint_name(subtype) + return type_param_name if type_param_name else None + origin_type = typing.get_origin(type_) + if isinstance(origin_type, type) and issubclass(origin_type, Sequence): + subtypes = typing.get_args(type_) + type_param_name = _get_type_constraint_name(subtypes[0]) + return f"Sequence_{type_param_name}" if type_param_name else None + return None + + +def _get_allowed_types_from_type_annotation( + type_: TypeAnnotationValue, +) -> set[ir.TypeProtocol]: + """Obtain the allowed types from a type annotation.""" + if type_ is onnxscript.onnx_types.TensorType: + # Any tensor type + return {ir.TensorType(dtype) for dtype in ir.DataType} + + allowed_types: set[ir.TypeProtocol] + + if isinstance(type_, TypeVar): + allowed_types = set() + if constraints := type_.__constraints__: + for constraint in constraints: + allowed_types.update( + _get_allowed_types_from_type_annotation(constraint) + ) + else: + bound = type_.__bound__ + if bound is None: + allowed_types = _ALL_VALUE_TYPES # type: ignore[assignment] + else: + allowed_types.update(_get_allowed_types_from_type_annotation(bound)) + return allowed_types + if hasattr(type_, "dtype"): + # A single tensor type like INT64, FLOAT, etc. + return {ir.TensorType(ir.DataType(type_.dtype))} + if _is_optional(type_): + allowed_types = set() + subtypes = typing.get_args(type_) + for subtype in subtypes: + if subtype is type(None): + continue + allowed_types.update(_get_allowed_types_from_type_annotation(subtype)) + # NOTE: We do not consider dynamic optional types like optional(float) because they are not very useful. + return allowed_types + + origin_type = typing.get_origin(type_) + if origin_type is Union: + allowed_types = set() + subtypes = typing.get_args(type_) + for subtype in subtypes: + if subtype is type(None): + raise AssertionError( + "Union should not contain None type because it is handled by _is_optional." + ) + allowed_types.update(_get_allowed_types_from_type_annotation(subtype)) + return allowed_types + + if isinstance(origin_type, type) and issubclass(origin_type, Sequence): + subtypes = typing.get_args(type_) + return { + ir.SequenceType(t) + for t in _get_allowed_types_from_type_annotation(subtypes[0]) + } + + # Allow everything by default + return _ALL_VALUE_TYPES # type: ignore[return-value] + + +def op_signature_from_function( + func, + domain: str, + name: str | None = None, + overload: str = "", + *, + since_version: int = 1, +) -> ir.schemas.OpSignature: + """Produce an OpSignature from a function using type annotation.""" + + py_signature = inspect.signature(func) + # Not using inspect.get_annotations because typing.get_type_hints seems to handle more cases + # https://github.com/python/cpython/issues/102405 + type_hints = typing.get_type_hints(func) + + params: list[ir.schemas.Parameter | ir.schemas.AttributeParameter] = [] + # Create a mapping from type to a unique name + type_constraints: dict[str, ir.schemas.TypeConstraintParam] = {} + + for param in py_signature.parameters.values(): + if param.name not in type_hints: + logger.debug( + "Missing annotation for parameter '%s' from %s. Treating as an Input.", + param.name, + py_signature, + ) + type_constraint = ir.schemas.TypeConstraintParam.any_value( + f"T_{param.name}" + ) + type_constraints[param.name] = type_constraint + kwargs: dict[str, Any] = {} + if param.default is not inspect.Parameter.empty: + kwargs["default"] = param.default + params.append( + ir.schemas.Parameter( + name=param.name, + type_constraint=type_constraint, + required=param.default is inspect.Parameter.empty, + # TODO: Handle variadic + variadic=False, + **kwargs, + ) + ) + else: + type_ = type_hints[param.name] + if (attr_type := _get_attr_type(type_)) != ir.AttributeType.UNDEFINED: + # Construct the default attribute + if param.default is not inspect.Parameter.empty: + # TODO: Use ir_convenience instead to handle int as float + default = ir.Attr(param.name, attr_type, param.default) + else: + default = None + params.append( + ir.schemas.AttributeParameter( + name=param.name, + type=attr_type, + required=param.default is inspect.Parameter.empty, + default=default, + ) + ) + else: + # Obtain the type constraint from the type annotation + + # 1. Get a type constraint name from the type annotation + # If the type annotation is a TypeVar or Optional[TypeVar], get its name + # Otherwise, name it T_{param.name} + type_constraint_name = _get_type_constraint_name(type_) + if type_constraint_name is None: + type_constraint_name = f"T_{param.name}" + + # 2. If the type constraint param is already initialized, use it + if type_constraint_name in type_constraints: + type_constraint = type_constraints[type_constraint_name] + else: + # 3. Otherwise, create a new TypeConstraintParam + type_constraint = ir.schemas.TypeConstraintParam( + name=type_constraint_name, + allowed_types=_get_allowed_types_from_type_annotation(type_), + ) + type_constraints[type_constraint_name] = type_constraint + # 4. Create Parameter + kwargs: dict[str, Any] = {} + if param.default is not inspect.Parameter.empty: + kwargs["default"] = param.default + params.append( + ir.schemas.Parameter( + name=param.name, + type_constraint=type_constraint, + required=param.default is inspect.Parameter.empty, + # TODO: Handle variadic + variadic=False, + **kwargs, + ) + ) + + return_type = type_hints.get("return") + + outputs = [] + if return_type is None: + # No returns + pass + else: + if typing.get_origin(return_type) is tuple: + # Multiple returns + return_types = typing.get_args(return_type) + else: + return_types = [return_type] # type: ignore[assignment] + + for i, return_type_i in enumerate(return_types): + if ( + return_param_name := _get_type_constraint_name(return_type_i) + ) in type_constraints: + # pyrefly: ignore [bad-index] + type_constraint = type_constraints[return_param_name] + else: + return_param_name = f"TReturn{i}" + type_constraint = ir.schemas.TypeConstraintParam( + name=return_param_name, + allowed_types=_get_allowed_types_from_type_annotation( + return_type_i + ), + ) + type_constraints[return_param_name] = type_constraint + outputs.append( + ir.schemas.Parameter( + name=return_param_name, + type_constraint=type_constraint, + required=True, + variadic=False, + ) + ) + + return ir.schemas.OpSignature( + domain=domain, + name=name or func.__name__, + overload=overload, + params=params, + outputs=outputs, + since_version=since_version, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_tensors.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_tensors.py new file mode 100644 index 0000000000000000000000000000000000000000..981e2fef97de0004baafcc4998289204f4ae20b3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_tensors.py @@ -0,0 +1,100 @@ +"""Subclass of ir.Value that supports Python operators.""" + +# mypy: allow-untyped-defs +from __future__ import annotations + +from typing import TYPE_CHECKING + +from torch.onnx._internal._lazy_import import onnx_ir as ir + + +if TYPE_CHECKING: + import onnxscript + + +class SymbolicTensor(ir.Value): + """A subclass of ir.Value that supports Python operators.""" + + def __init__( + self, + opset: onnxscript.values.Opset, + name: str | None = None, + shape: ir.Shape | None = None, + type: ir.TypeProtocol | None = None, + doc_string: str | None = None, + const_value: ir.TensorProtocol | None = None, + ) -> None: + super().__init__( + name=name, + shape=shape, + type=type, + doc_string=doc_string, + const_value=const_value, + ) + self._opset = opset + + @property + def rank(self) -> int | None: + if self.shape is None: + return None + return len(self.shape) + + # TODO: Implement indexing + + def __mod__(self, other): + if self.dtype in { + ir.DataType.FLOAT, + ir.DataType.DOUBLE, + ir.DataType.FLOAT16, + ir.DataType.BFLOAT16, + }: + return self._opset.Mod(self, other, fmod=1) + return self._opset.Mod(self, other) + + def __ne__(self, other): + return self._opset.Not(self._opset.Equal(self, other)) + + def __neg__(self): + return self._opset.Neg(self) + + def __add__(self, other): + return self._opset.Add(self, other) + + def __radd__(self, other): + return self._opset.Add(other, self) + + def __rand__(self, other): + return self._opset.And(other, self) + + def __mul__(self, other): + return self._opset.Mul(self, other) + + def __rmul__(self, other): + return self._opset.Mul(other, self) + + def __matmul__(self, other): + return self._opset.MatMul(self, other) + + def __pow__(self, other): + return self._opset.Pow(self, other) + + def __sub__(self, other): + return self._opset.Sub(self, other) + + def __rsub__(self, other): + return self._opset.Sub(other, self) + + def __truediv__(self, other): + return self._opset.Div(self, other) + + def __lt__(self, other): + return self._opset.Less(self, other) + + def __le__(self, other): + return self._opset.LessOrEqual(self, other) + + def __ge__(self, other): + return self._opset.GreaterOrEqual(self, other) + + def __gt__(self, other): + return self._opset.Greater(self, other) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_testing.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_testing.py new file mode 100644 index 0000000000000000000000000000000000000000..c34c2f1a38c3d639d7d48818b1db1ac7ad2ba7ba --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_testing.py @@ -0,0 +1,102 @@ +"""Test utilities for ONNX export.""" + +from __future__ import annotations + + +__all__ = ["assert_onnx_program"] + +from typing import Any, Literal, TYPE_CHECKING + +import torch +from torch.utils import _pytree + + +if TYPE_CHECKING: + from torch.onnx._internal.exporter import _onnx_program + + +def assert_onnx_program( + program: _onnx_program.ONNXProgram, + *, + rtol: float | None = None, + atol: float | None = None, + args: tuple[Any, ...] | None = None, + kwargs: dict[str, Any] | None = None, + strategy: str | None = "TorchExportNonStrictStrategy", + backend: Literal["onnxruntime", "reference"] = "onnxruntime", +) -> None: + """Assert that the ONNX model produces the same output as the PyTorch ExportedProgram. + + Args: + program: The ``ONNXProgram`` to verify. + rtol: Relative tolerance. + atol: Absolute tolerance. + args: The positional arguments to pass to the program. + If None, the default example inputs in the ExportedProgram will be used. + kwargs: The keyword arguments to pass to the program. + If None, the default example inputs in the ExportedProgram will be used. + strategy: Assert the capture strategy used to export the program. Values can be + class names like "TorchExportNonStrictStrategy". + If None, the strategy is not asserted. + backend: The backend to use for evaluating the ONNX program. + Supported values are "onnxruntime" and "reference". + """ + if strategy is not None: + if program._capture_strategy != strategy: + raise ValueError( + f"Expected strategy '{strategy}' is used to capture the exported program, " + f"but got '{program._capture_strategy}'." + ) + exported_program = program.exported_program + if exported_program is None: + raise ValueError( + "The ONNXProgram does not contain an ExportedProgram. " + "To verify the ONNX program, initialize ONNXProgram with an ExportedProgram, " + "or assign the ExportedProgram to the ONNXProgram.exported_program attribute." + ) + if args is None and kwargs is None: + # User did not provide example inputs, use the default example inputs + if exported_program.example_inputs is None: + raise ValueError( + "No example inputs provided and the exported_program does not contain example inputs. " + "Please provide arguments to verify the ONNX program." + ) + args, kwargs = exported_program.example_inputs + if args is None: + args = () + if kwargs is None: + kwargs = {} + torch_module = exported_program.module() + torch_outputs, _ = _pytree.tree_flatten(torch_module(*args, **kwargs)) + # ONNX outputs are always real, so we need to convert torch complex outputs to real representations + torch_outputs_adapted = [] + for output in torch_outputs: + # ONNX graph does not support None outputs, so we skip them + if output is None: + continue + if not isinstance(output, torch.Tensor): + torch_outputs_adapted.append(torch.tensor(output)) + elif torch.is_complex(output): + torch_outputs_adapted.append(torch.view_as_real(output)) + else: + torch_outputs_adapted.append(output) + + # Obtain the ONNX outputs using the specified backend + if backend == "onnxruntime": + onnx_outputs = program(*args, **kwargs) + elif backend == "reference": + onnx_outputs = program.call_reference(*args, **kwargs) + else: + raise ValueError( + f"Unsupported backend '{backend}'. Supported backends are 'onnxruntime' and 'reference'." + ) + + # TODO(justinchuby): Include output names in the error message + torch.testing.assert_close( + tuple(onnx_outputs), + tuple(torch_outputs_adapted), + rtol=rtol, + atol=atol, + equal_nan=True, + check_device=False, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/_tensor_typing.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/_tensor_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..3212cf32468e1d8530439c56c9bd3b9e9a61e974 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/_tensor_typing.py @@ -0,0 +1,75 @@ +"""Typings for function definitions.""" + +from __future__ import annotations + +from typing import TypeVar, Union + +from onnxscript import ( + BFLOAT16, + BOOL, + COMPLEX128, + COMPLEX64, + DOUBLE, + FLOAT, + FLOAT16, + INT16, + INT32, + INT64, + INT8, + STRING, + UINT8, +) + + +# NOTE: We do not care about unsigned types beyond UINT8 because PyTorch does not us them. +# More detail can be found: https://pytorch.org/docs/stable/tensors.html + +TensorType = Union[ # noqa: UP007 + BFLOAT16, + BOOL, + COMPLEX64, + COMPLEX128, + DOUBLE, + FLOAT, + FLOAT16, + INT8, + INT16, + INT32, + INT64, + UINT8, +] +_FloatType = Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16] # noqa: UP007 +IntType = Union[INT8, INT16, INT32, INT64] # noqa: UP007 +RealType = Union[ # noqa: UP007 + BFLOAT16, + FLOAT16, + FLOAT, + DOUBLE, + INT8, + INT16, + INT32, + INT64, +] + +TTensor = TypeVar("TTensor", bound=TensorType) +# Duplicate TTensor for inputs/outputs that accept the same set of types as TTensor +# but do not constrain the type to be the same as the other inputs/outputs +TTensor2 = TypeVar("TTensor2", bound=TensorType) +TTensorOrString = TypeVar("TTensorOrString", bound=Union[TensorType, STRING]) # noqa: UP007 +TFloat = TypeVar("TFloat", bound=_FloatType) +TFloatOrUInt8 = TypeVar( + "TFloatOrUInt8", + bound=Union[FLOAT, FLOAT16, DOUBLE, INT8, UINT8], # noqa: UP007 +) +TInt = TypeVar("TInt", bound=IntType) +TReal = TypeVar("TReal", bound=RealType) +TRealUnlessInt16OrInt8 = TypeVar( + "TRealUnlessInt16OrInt8", + bound=Union[FLOAT16, FLOAT, DOUBLE, BFLOAT16, INT32, INT64], # noqa: UP007 +) +TRealUnlessFloat16OrInt8 = TypeVar( + "TRealUnlessFloat16OrInt8", + bound=Union[DOUBLE, FLOAT, INT16, INT32, INT64], # noqa: UP007 +) +TRealOrUInt8 = TypeVar("TRealOrUInt8", bound=Union[RealType, UINT8]) # noqa: UP007 +TFloatHighPrecision = TypeVar("TFloatHighPrecision", bound=Union[FLOAT, DOUBLE]) # noqa: UP007 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/_torchlib_registry.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/_torchlib_registry.py new file mode 100644 index 0000000000000000000000000000000000000000..e900c1e2daa443b50ba32ffbbbb8228e045ce393 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/_torchlib_registry.py @@ -0,0 +1,95 @@ +"""Registry for aten functions.""" + +from __future__ import annotations + + +__all__ = ["onnx_impl", "get_torchlib_ops"] + +import logging +from collections.abc import Callable, Sequence +from typing import Any, TypeVar +from typing_extensions import ParamSpec + +import onnxscript + +import torch +from torch.onnx._internal.exporter import _constants, _registration + + +# Use ParamSpec for better type preservation instead of bound Callable TypeVar +_P = ParamSpec("_P") +_R = TypeVar("_R") + +logger = logging.getLogger("__name__") + + +_registry: list[_registration.OnnxDecompMeta] = [] + + +def onnx_impl( + target: _registration.TorchOp | tuple[_registration.TorchOp, ...], + *, + trace_only: bool = False, + complex: bool = False, + opset_introduced: int = 18, + no_compile: bool = False, + private: bool = False, +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: + """Register an ONNX implementation of a torch op.""" + + if isinstance(target, torch._ops.OpOverloadPacket): + raise TypeError( + f"Target '{target}' should be provided as an OpOverload instead of an " + "OpOverloadPacket. You can get the default overload with " + ".default" + ) + + def wrapper( + func: Callable[_P, _R], + ) -> Callable[_P, _R]: + processed_func: Any + if no_compile: + processed_func = func + else: + torchlib_opset = onnxscript.values.Opset( + domain=_constants.TORCHLIB_DOMAIN, version=1 + ) + + if not trace_only: + # Compile the function + processed_func = onnxscript.script(opset=torchlib_opset)(func) + else: + processed_func = onnxscript.TracedOnnxFunction(torchlib_opset, func) + + if not private: + # TODO(justinchuby): Simplify the logic and remove the private attribute + # Skip registration if private + if not isinstance(target, Sequence): + targets = (target,) + else: + targets = target # type: ignore[assignment] + + for t in targets: + _registry.append( + _registration.OnnxDecompMeta( + onnx_function=processed_func, + fx_target=t, + signature=None, + is_complex=complex, + opset_introduced=opset_introduced, + skip_signature_inference=no_compile, + ) + ) + return processed_func # type: ignore[return-value] + + return wrapper + + +def get_torchlib_ops() -> tuple[_registration.OnnxDecompMeta, ...]: + # Trigger op registration + from torch.onnx._internal.exporter._torchlib import ops + + del ops + if len(_registry) == 0: + raise AssertionError("_registry must not be empty") + return tuple(_registry) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b7834043414f2191959cf760bb8c7e13b8d543ec --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/__init__.py @@ -0,0 +1,6 @@ +from __future__ import annotations + + +__all__ = ["core", "hop", "nn", "symbolic", "symops"] + +from torch.onnx._internal.exporter._torchlib.ops import core, hop, nn, symbolic, symops diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/core.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/core.py new file mode 100644 index 0000000000000000000000000000000000000000..36d53b113edc20bee27ef941cbbb740c34727a52 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/core.py @@ -0,0 +1,47 @@ +"""torch.ops.aten operators under the `core` module.""" +# mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value,type-var,operator,no-untyped-def,index" +# pyrefly: ignore-errors +# ruff: noqa: TCH001,TCH002 + +from __future__ import annotations + +import operator + +from onnxscript.onnx_opset import opset18 as op + +import torch +from torch.onnx._internal.exporter._torchlib._tensor_typing import TReal, TRealOrUInt8 +from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl + + +aten = torch.ops.aten + + +@onnx_impl((aten.abs.default, operator.abs), trace_only=True) +def aten_abs(self: TRealOrUInt8) -> TRealOrUInt8: + """abs(Tensor self) -> Tensor""" + + return op.Abs(self) + + +@onnx_impl(aten.abs.default, complex=True, trace_only=True) +def aten_abs_complex(self: TRealOrUInt8) -> TRealOrUInt8: + """abs(Tensor self) -> Tensor""" + + return op.ReduceL2(self, [-1], keepdims=False) + + +@onnx_impl((aten.add.Tensor, aten.add.Scalar, operator.add), trace_only=True) +def aten_add(self: TReal, other: TReal, alpha: float = 1.0) -> TReal: + """add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor""" + if alpha != 1.0: + alpha = op.CastLike(alpha, other) + other = op.Mul(other, alpha) + return op.Add(self, other) + + +@onnx_impl((aten.add.Tensor, aten.add.Scalar), trace_only=True, complex=True) +def aten_add_complex(self: TReal, other: TReal, alpha: float = 1.0) -> TReal: + """add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor""" + + return aten_add(self, other, alpha=alpha) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/hop.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/hop.py new file mode 100644 index 0000000000000000000000000000000000000000..3dece30a84aa51ed752f7bb418303ead0b5e7f2d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/hop.py @@ -0,0 +1,370 @@ +"""Implementation for higher-order operators.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +from torch.onnx._internal._lazy_import import onnx_ir as ir +from torch.onnx._internal.exporter import _core +from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +def call_op( + op_type: str, + *args: ir.Value | None, + _num_outputs: int = 1, + _domain: str = "", + **kwargs: int | float | str | bool | ir.Graph | ir.TensorProtocol | Sequence[int], +) -> Sequence[ir.Value]: + """Call an operator with the given arguments and keyword arguments. + + Arguments are always inputs, while keyword arguments are attributes. + """ + # This is a wrapper around the IR node creation that hooks into the _builder.OpRecorder + # tracer so that all nodes created are recorded the same way as if we were to use + # onnxscript ops directly. + from onnxscript.ir import convenience as ir_convenience + + if _core.current_tracer is None: + raise AssertionError("current_tracer must be non-None") + tracer = _core.current_tracer + + inputs = list(args) + + # If final inputs are None, strip them from the node inputs + for input in reversed(inputs): + if input is not None: + break + inputs.pop() + + # Construct and filter out None attributes + attributes = [ + attr + for attr in ir_convenience.convert_attributes(kwargs) + if attr.value is not None # type: ignore[union-attr] + ] + tracer.nodes.append( + node := ir.Node( + _domain, + op_type, + inputs=inputs, + attributes=attributes, + num_outputs=_num_outputs, + version=tracer.opset.version, + ) + ) + return node.outputs + + +@onnx_impl(torch.ops.higher_order.cond, no_compile=True) +def higher_order_cond( + cond: ir.Value, + true_func: ir.Function, + false_func: ir.Function, + inputs: Sequence[ir.Value], +) -> Sequence[ir.Value]: + then_node = ir.Node( + true_func.domain, true_func.name, inputs, num_outputs=len(true_func.outputs) + ) + else_node = ir.Node( + false_func.domain, false_func.name, inputs, num_outputs=len(false_func.outputs) + ) + + # ONNX Runtime complains about duplicate output names if we don't rename them. + # But the doesn't seem to be an actual violation of SSA form without renaming. + for func_out, out in zip(true_func.outputs, then_node.outputs): + out.name = f"{func_out.name}_{true_func.name}" + for func_out, out in zip(false_func.outputs, else_node.outputs): + out.name = f"{func_out.name}_{false_func.name}" + + return call_op( + "If", + cond, + _num_outputs=len(true_func.outputs), + then_branch=ir.Graph( + (), then_node.outputs, nodes=[then_node], name=true_func.name + ), + else_branch=ir.Graph( + (), else_node.outputs, nodes=[else_node], name=false_func.name + ), + ) + + +@onnx_impl(torch.ops.higher_order.scan, no_compile=True) +def higher_order_scan( + body_func: ir.Function, + scan_inits: Sequence[ir.Value], + scan_inputs: Sequence[ir.Value], + additional_inputs: Sequence[ir.Value] | None, + reverse: bool = False, +) -> Sequence[ir.Value]: + """https://github.com/pytorch/pytorch/blob/66ac724b56e6c37a534f3e066423ef2f41d7477f/torch/_higher_order_ops/scan.py#L109""" + subgraph_inputs = [ + *[ + ir.Value( + name=f"{inp.name}_{body_func.name}__subgraph_in", + shape=inp.shape, + type=ir.TensorType(inp.dtype), # type: ignore[arg-type] + ) + for inp in scan_inits + ], + *[ + ir.Value( + name=f"{inp.name}_{body_func.name}__subgraph_in", + # The iterated element passed to the body subgraph does not have a sequence axis. + # It will have a rank one less than the rank of the corresponding scan_input. + shape=ir.Shape(inp.shape[1:]), # type: ignore[index] + type=ir.TensorType(inp.dtype), # type: ignore[arg-type] + ) + for inp in scan_inputs + ], + ] + # The one and only node in the Scan subgraph that calls the body_func + body_node = ir.Node( + body_func.domain, + body_func.name, + [ + *subgraph_inputs, + *(additional_inputs or []), + ], + num_outputs=len(body_func.outputs), + ) + + # ONNX Runtime complains about duplicate output names if we don't rename them. + # But the doesn't seem to be an actual violation of SSA form without renaming. + for func_out, out in zip(body_func.outputs, body_node.outputs): + out.name = f"{func_out.name}_{body_func.name}" + + n_outputs = len(body_func.outputs) - len(scan_inits) + return call_op( + "Scan", + *scan_inits, + *scan_inputs, + _num_outputs=len(body_func.outputs), + body=ir.Graph( + subgraph_inputs, + body_node.outputs, + nodes=[body_node], + name=body_func.name, + ), + num_scan_inputs=len(scan_inputs), + scan_input_directions=[(1 if reverse else 0) for _ in scan_inputs], + scan_output_directions=[(1 if reverse else 0) for _ in range(n_outputs)], + ) + + +@onnx_impl(torch.ops.higher_order.while_loop, no_compile=True) +def higher_order_while_loop( + cond_func: ir.Function, + body_func: ir.Function, + carried_inputs: Sequence[ir.Value | int | float], + additional_inputs: Sequence[ir.Value], +) -> Sequence[ir.Value]: + """Implementation of while_loop using ONNX Loop operator. + + The ONNX Loop operator implements a generic looping construct with the signature: + Loop(M, cond, v_initial) -> (v_final_and_scan_outputs) + + For while_loop, we use: + - M: None (no trip count limit) + - cond: initial condition value + - v_initial: carried_inputs (loop-carried dependencies) + + The body subgraph takes: + - iteration_num (int): current iteration number + - condition_in (bool): loop continuation condition from previous iteration + - loop_carried_dependencies: the carried values + - additional_inputs: any additional inputs (constants/parameters) + + The body subgraph returns: + - condition_out (bool): whether to continue looping + - loop_carried_dependencies: updated carried values + """ + + # Create subgraph inputs for the Loop body + # ONNX Loop body signature: (iter_num, cond_in, loop_carried_deps..., additional_inputs...) + + # Start subgraph construction + subgraph_carried_inputs = [] + + for i, inp in enumerate(carried_inputs): + if isinstance(inp, ir.Value): + subgraph_carried_inputs.append( + ir.Value( + name=f"{inp.name}_{body_func.name}__subgraph_in", + shape=inp.shape, + type=ir.TensorType(inp.dtype), # type: ignore[arg-type] + ) + ) + elif isinstance(inp, int): + subgraph_carried_inputs.append( + ir.Value( + name=f"carried_input_{i}_{body_func.name}__subgraph_in", + shape=ir.Shape([]), + type=ir.TensorType(ir.DataType.INT64), + ) + ) + elif isinstance(inp, float): + subgraph_carried_inputs.append( + ir.Value( + name=f"carried_input_{i}_{body_func.name}__subgraph_in", + shape=ir.Shape([]), + type=ir.TensorType(ir.DataType.FLOAT), + ) + ) + else: + raise NotImplementedError( + f"Unsupported type for carried input: {type(inp)} ({inp}). " + "Expected ir.Value, int, or float." + ) + + subgraph_inputs = [ + # Iteration number (int scalar, unused) + ir.Value( + name=f"iter_num_{body_func.name}", + shape=ir.Shape([]), + type=ir.TensorType(ir.DataType.INT64), + ), + # Condition input (bool scalar, unused) + ir.Value( + name=f"cond_in_{body_func.name}", + shape=ir.Shape([]), + type=ir.TensorType(ir.DataType.BOOL), + ), + # Loop-carried dependencies + *subgraph_carried_inputs, + ] + + # Create the combined body function that handles both condition and body logic + # First, call the body function with the same inputs + body_node = ir.Node( + body_func.domain, + body_func.name, + [ + *subgraph_carried_inputs, # carried inputs + *additional_inputs, + ], + num_outputs=len(body_func.outputs), # carried inputs + ) + + # Then call the condition function with carried inputs + additional inputs + cond_node = ir.Node( + cond_func.domain, + cond_func.name, + [ + *body_node.outputs, # updated carried inputs from body + *additional_inputs, + ], + num_outputs=len(cond_func.outputs), + ) + + if len(cond_func.outputs) != 1: + raise AssertionError("Condition function must return a single boolean value.") + + # ONNX Runtime complains about duplicate output names if we don't rename them + for func_out, out in zip(body_func.outputs, body_node.outputs): + out.name = f"{func_out.name}_{body_func.name}" + for func_out, out in zip(cond_func.outputs, cond_node.outputs): + out.name = f"{func_out.name}_{cond_func.name}" + + # The Loop body must return: (cond_out, loop_carried_deps...) + # We use the condition output and the body outputs + loop_body_outputs = [ + cond_node.outputs[0], # condition output (bool) + *body_node.outputs, # updated carried inputs + ] + + body_graph = ir.Graph( + subgraph_inputs, + loop_body_outputs, + nodes=[body_node, cond_node], + name=f"{body_func.name}_loop_body", + ) + + # End subgraph construction + + carried_inputs_values: list[ir.Value] = [] + for inp in carried_inputs: + if isinstance(inp, ir.Value): + carried_inputs_values.append(inp) + elif isinstance(inp, int): + const = call_op("Constant", value=ir.tensor(inp))[0] + carried_inputs_values.append(const) + elif isinstance(inp, float): + const = call_op("Constant", value=ir.tensor(inp))[0] + carried_inputs_values.append(const) + else: + raise NotImplementedError( + f"Unsupported type for carried input: {type(inp)} ({inp}). " + "Expected ir.Value, int, or float." + ) + + # Get initial condition by calling cond_func with initial inputs + initial_outputs = call_op( + cond_func.name, + *carried_inputs_values, + *additional_inputs, + _num_outputs=len(cond_func.outputs), + _domain=cond_func.domain, + ) + + if len(initial_outputs) != 1: + raise AssertionError("Condition function must return a single boolean value.") + + # Create the Loop operator call + # Loop(M, cond, v_initial) where M is empty (no trip count limit) + loop_outputs = call_op( + "Loop", + # M (trip count) - empty string means no limit + None, + # cond - initial condition + initial_outputs[0], + # v_initial - carried inputs (loop-carried dependencies) + *carried_inputs_values, + _num_outputs=len(carried_inputs_values), + body=body_graph, + ) + + return loop_outputs + + +@onnx_impl(torch.ops.higher_order.invoke_subgraph, no_compile=True) +def higher_order_invoke_subgraph( + subgraph: ir.Function, + identifier: str | None, + *operands: ir.Value, +) -> Sequence[ir.Value]: + """Export invoke_subgraph HOP by creating a direct function call. + + This preserves the function as a separate entity in the ONNX graph + instead of inlining it, which is the purpose of invoke_subgraph. + + Note: The onnxscript optimizer should be configured to not inline functions + created by invoke_subgraph to preserve the intended structure. + + Args: + subgraph: The function to invoke + identifier: Optional identifier for the subgraph (used for caching in PyTorch, + not needed for ONNX export as the function reference provides all necessary information) + *operands: Input values to pass to the function + + Returns: + Sequence of output values from the function call + """ + # This key can be used by downstream to avoid inlining + subgraph.metadata_props["pkg.torch.ops.higher_order.invoke_subgraph.identifier"] = ( + str(identifier) + ) + + # Create the function call node + return call_op( + subgraph.name, + *operands, + _num_outputs=len(subgraph.outputs), + _domain=subgraph.domain, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/nn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/nn.py new file mode 100644 index 0000000000000000000000000000000000000000..14c3ad6d2a3ab7e11a0b4d5129f086dad23c9498 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/nn.py @@ -0,0 +1,377 @@ +"""torch.ops.aten operators under the `core` module.""" +# mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value,type-var,operator,no-untyped-def,index" +# pyrefly: ignore-errors +# ruff: noqa: TC001,TC002 +# flake8: noqa: B950 + +from __future__ import annotations + +from typing import Sequence, TYPE_CHECKING # noqa: UP035 + +from onnxscript.onnx_opset import ( # type: ignore[attr-defined] + opset20 as op20, + opset21 as op21, + opset23 as op23, +) + +import torch +from torch.onnx._internal._lazy_import import onnx_ir as ir +from torch.onnx._internal.exporter._torchlib._tensor_typing import TFloat, TReal +from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl + + +if TYPE_CHECKING: + from onnxscript.values import Opset + +aten = torch.ops.aten + + +@onnx_impl(aten.gelu.default, trace_only=True, opset_introduced=20) +def aten_gelu_opset20( + self: TReal, + approximate: str = "none", +) -> TReal: + """gelu(Tensor self, *, str approximate="none") -> Tensor""" + return op20.Gelu(self, approximate=approximate) + + +@onnx_impl(aten.group_norm.default, trace_only=True, opset_introduced=21) +def aten_group_norm( + input: TFloat, + num_groups: int, + weight: TFloat | None = None, + bias: TFloat | None = None, + eps: float = 1e-05, + cudnn_enabled: bool = True, +) -> TFloat: + """group_norm(Tensor input, int num_groups, Tensor? weight=None, Tensor? bias=None, float eps=1e-05, bool cudnn_enabled=True) -> Tensor""" + + c = op21.Shape(input, start=1, end=2) + if weight is None: + weight = op21.ConstantOfShape(c, value=ir.tensor([1.0], dtype=input.dtype)) + if bias is None: + bias = op21.ConstantOfShape(c, value=ir.tensor([0.0], dtype=input.dtype)) + return op21.GroupNormalization( + input, weight, bias, epsilon=eps, num_groups=num_groups + ) + + +@onnx_impl(aten.rms_norm.default, trace_only=True, opset_introduced=23) +def aten_rms_norm( + input: TFloat, + normalized_shape: Sequence[int], + weight: TFloat | None = None, + eps: float | None = None, +) -> TFloat: + """rms_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor""" + + # Default eps value if not provided + if eps is None: + eps = torch.finfo(torch.float).eps # Observed from decomp + + # Calculate axis: the first normalization dimension + # For normalized_shape with D dimensions, normalize over last D dimensions + # Since ONNX RMSNormalization supports negative axis values, we use -len(normalized_shape) + # which correctly maps to the first axis of the normalized dimensions + normalized_dims = len(normalized_shape) + axis = -normalized_dims + + # Create weight tensor if not provided + if weight is None: + weight = op23.ConstantOfShape( + op23.Shape(input), value=ir.tensor([1], dtype=input.dtype) + ) + + return op23.RMSNormalization(input, weight, axis=axis, epsilon=eps) + + +@onnx_impl( + aten.scaled_dot_product_attention.default, trace_only=True, opset_introduced=23 +) +def aten_scaled_dot_product_attention_23( + query: TFloat, + key: TFloat, + value: TFloat, + attn_mask: TFloat | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + enable_gqa: bool = False, +) -> TFloat: + """scaled_dot_product_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None, bool enable_gqa=False) -> Tensor + + Reference: + 1. https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html + 2. https://onnx.ai/onnx/operators/onnx__Attention.html + + Attempts to convert SDPA to Attention onnx op and fallbacks to an onnx graph equivalent to the following PyTorch code:: + scale_factor = 1 / math.sqrt(Q.size(-1)) if scale is None else scale + attn_mask = ( + torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) + if is_causal + else attn_mask + ) + attn_mask = ( + attn_mask.masked_fill(not attn_mask, -float("inf")) + if attn_mask.dtype == torch.bool + else attn_mask + ) + attn_weight = torch.softmax( + (Q @ K.transpose(-2, -1) * scale_factor) + attn_mask, dim=-1 + ) + attn_weight = torch.dropout(attn_weight, dropout_p) + return attn_weight @ V + + where Q, K, V are the query, key, and value tensors, respectively. + L is the target sequence length, S is the source sequence length, and E is the embedding size. + """ + if is_causal and attn_mask is not None: + raise AssertionError("is_causal and attn_mask cannot be set at the same time") + if not (len(query.shape) == 4 and len(key.shape) == 4 and len(value.shape) == 4): + raise AssertionError("only 4D query, key, and value are supported") + + # Attention onnx op can only handle non-training scenarios where dropout is disabled. + if dropout_p == 0: + if enable_gqa: + if not ( + query.shape[1] > key.shape[1] == value.shape[1] + and query.shape[1] % key.shape[1] == 0 + ): + raise AssertionError( + "SDPA (GQA or MQA) requires q_num_heads > kv_num_heads & " + "q_num_heads % kv_num_heads == 0" + ) + else: + if not (query.shape[1] == key.shape[1] == value.shape[1]): + raise AssertionError("SDPA (MHA) requires q_num_heads = kv_num_heads") + + # NOTE: num_heads attributes (q_num_heads/kv_num_heads) should not be specified for 4D. + # They are not populated with 4D inputs because this information directly comes from input shapes: + # `q_num_heads=query.shape[1]` and `kv_num_heads=key.shape[1]`. + # This dimension is usually static but it could not be dynamic if also given as an attribute. + # num_heads attributes are needed for 3D attention inputs: + # (shape: [B, S, N*H]), 4D shape is ([B, N, S, H]). + + Y, _, _, _ = op23.Attention( + query, + key, + value, + attn_mask=attn_mask, + scale=scale, + is_causal=is_causal, + ) + return Y + + if scale is None: + scale = _attention_scale(query, op23) + scale = op23.CastLike(scale, query) + + if is_causal: + attn_mask = _causal_attention_mask(query, key, op23) + + if enable_gqa: + key, value = _attention_repeat_kv_for_group_query(query, key, value, op23) + + if attn_mask is None: + return _aten_scaled_dot_product_attention_no_mask_onnx( + query, key, value, scale, dropout_p, op23 + ) + + return _aten_scaled_dot_product_attention_float_mask_onnx( + query, key, value, attn_mask, scale, dropout_p, op23 + ) + + +def _attention_repeat_kv_for_group_query( + query: TFloat, key: TFloat, value: TFloat, op: Opset +) -> tuple[TFloat, TFloat]: + """Expand key and value for group query attention. + + repeat_interleave is applied on key and value to match the number of heads in query. + + Args: + query: Tensor of shape [B, q_num_heads, q_S, E] + key: Tensor of shape [B, k_num_heads, kv_S, E] + value: Tensor of shape [B, v_num_heads, kv_S, E] + + Returns: + Tuple of (expanded_key, expanded_value) where: + - expanded_key: Tensor of shape [B, q_num_heads, kv_S, E] + - expanded_value: Tensor of shape [B, q_num_heads, kv_S, E] + """ + + if not ( + query.shape[1] > key.shape[1] == value.shape[1] + and query.shape[1] % key.shape[1] == 0 + ): + raise AssertionError( + "SDPA (GQA or MQA) requires q_num_heads > kv_num_heads & " + "q_num_heads % kv_num_heads == 0" + ) + + # NOTE: QKV are expected to be 4D tensors + + batch_size = op.Shape(query, start=0, end=1) # [B] + q_num_heads = op.Shape(query, start=1, end=2) # [Hq] + kv_num_heads = op.Shape(key, start=1, end=2) # [Hk] + qk_head_size = op.Shape(key, start=3, end=4) # [Dk] + v_head_size = op.Shape(value, start=3, end=4) # [Dv] + new_kv_seq_len = op.Shape(key, start=2, end=3) # [T] + + interleave_dim = op.Div(q_num_heads, kv_num_heads) # Hq / Hk + two = op.Constant(value_int=2) + k_unsqueezed = op.Unsqueeze(key, two) # [B, Hk, 1, T, Dk] + v_unsqueezed = op.Unsqueeze(value, two) # [B, Hv, 1, T, Dv] + + k_expand_shape = op.Concat( + batch_size, kv_num_heads, interleave_dim, new_kv_seq_len, qk_head_size, axis=0 + ) + k_expand = op.Expand(k_unsqueezed, k_expand_shape) + v_expand_shape = op.Concat( + batch_size, kv_num_heads, interleave_dim, new_kv_seq_len, v_head_size, axis=0 + ) + v_expand = op.Expand(v_unsqueezed, v_expand_shape) + + k_attention_shape = op.Concat( + batch_size, q_num_heads, new_kv_seq_len, qk_head_size, axis=0 + ) + v_attention_shape = op.Concat( + batch_size, q_num_heads, new_kv_seq_len, v_head_size, axis=0 + ) + + expanded_key = op.Reshape(k_expand, k_attention_shape) + expanded_value = op.Reshape(v_expand, v_attention_shape) + + return expanded_key, expanded_value + + +def _attention_scale(query: TFloat, op: Opset) -> TFloat: + """Calculate the scale factor for the attention result. + + Args: + query: Tensor of shape [..., L, E] + + Returns: + Scalar scale factor := 1 / math.sqrt(query.size(-1)) + """ + q_shape = op.Shape(query) + q_last_dim = op.Gather(q_shape, op.Constant(value_ints=[-1])) + embedding_size = op.CastLike(q_last_dim, query) + one = op.Constant(value_float=1.0) + cast_one = op.CastLike(one, query) + scale = op.Div(cast_one, op.Sqrt(embedding_size)) + return scale + + +def _causal_attention_mask(query: TFloat, key: TFloat, op: Opset) -> TFloat: + """Create a causal mask for the given query and key tensors. + + Equivalent to:: + mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) + attn_mask = torch.zeros(L, S, dtype=torch.float) + attn_mask = attn_mask.masked_fill(not mask, -float("inf")) + + Args: + query: Tensor of shape [..., L, E] + key: Tensor of shape [..., S, E] + + Returns: + Tensor of shape [L, S] + """ + q_shape = op.Shape(query) + k_shape = op.Shape(key) + + target_length = op.Slice( + q_shape, op.Constant(value_ints=[-2]), op.Constant(value_ints=[-1]) + ) + source_length = op.Slice( + k_shape, op.Constant(value_ints=[-2]), op.Constant(value_ints=[-1]) + ) + # attn_mask = torch.ones(L, S) := { + size = op.Concat(target_length, source_length, axis=0) + attn_mask = op.Expand(op.Constant(value_float=1.0), size) + # } + attn_mask = op.Trilu(attn_mask, upper=0) + # The causal mask has 0s in the lower triangle and -inf in the upper triangle. + attn_mask = op.Where( + op.Equal(attn_mask, op.Constant(value_float=0.0)), + op.Constant(value_float=-float("inf")), + op.Constant(value_float=0.0), + ) + attn_mask = op.CastLike(attn_mask, query) + return attn_mask + + +def _aten_scaled_dot_product_attention_no_mask_onnx( + query: TFloat, + key: TFloat, + value: TFloat, + scale: TFloat, + dropout_p: float, + op: Opset, +) -> TFloat: + # Swap the last two axes of key + key_last_dim = op.Shape(key, start=-1) + key_second_last_dim = op.Shape(key, start=-2, end=-1) + key_first_dims = op.Shape(key, end=-2) + # Contract the dimensions that are not the last two so we can transpose + # with a static permutation. + key_squeezed_shape = op.Concat( + op.Constant(value_ints=[-1]), key_second_last_dim, key_last_dim, axis=0 + ) + key_squeezed = op.Reshape(key, key_squeezed_shape) + key_squeezed_transposed = op.Transpose(key_squeezed, perm=[0, 2, 1]) + key_transposed_shape = op.Concat( + key_first_dims, key_last_dim, key_second_last_dim, axis=0 + ) + key_transposed = op.Reshape(key_squeezed_transposed, key_transposed_shape) + + # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653 + # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math + query_scaled = op.Mul(query, op.Sqrt(scale)) + key_transposed_scaled = op.Mul( + key_transposed, op.CastLike(op.Sqrt(scale), key_transposed) + ) + attn_weight = op.Softmax( + op.MatMul(query_scaled, key_transposed_scaled), + axis=-1, + ) + attn_weight, _ = op.Dropout(attn_weight, dropout_p) + return op.MatMul(attn_weight, value) + + +def _aten_scaled_dot_product_attention_float_mask_onnx( + query: TFloat, + key: TFloat, + value: TFloat, + attn_mask: TFloat, + scale: TFloat, + dropout_p: float, + op: Opset, +) -> TFloat: + # Swap the last two axes of key + key_last_dim = op.Shape(key, start=-1) + key_second_last_dim = op.Shape(key, start=-2, end=-1) + key_first_dims = op.Shape(key, end=-2) + # Contract the dimensions that are not the last two so we can transpose + # with a static permutation. + key_squeezed_shape = op.Concat( + op.Constant(value_ints=[-1]), key_second_last_dim, key_last_dim, axis=0 + ) + key_squeezed = op.Reshape(key, key_squeezed_shape) + key_squeezed_transposed = op.Transpose(key_squeezed, perm=[0, 2, 1]) + key_transposed_shape = op.Concat( + key_first_dims, key_last_dim, key_second_last_dim, axis=0 + ) + key_transposed = op.Reshape(key_squeezed_transposed, key_transposed_shape) + + # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653 + # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math + query_scaled = op.Mul(query, op.Sqrt(scale)) + key_transposed_scaled = op.Mul(key_transposed, op.Sqrt(scale)) + attn_weight = op.Softmax( + op.Add(op.MatMul(query_scaled, key_transposed_scaled), attn_mask), + axis=-1, + ) + attn_weight, _ = op.Dropout(attn_weight, dropout_p) + return op.MatMul(attn_weight, value) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/symbolic.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/symbolic.py new file mode 100644 index 0000000000000000000000000000000000000000..f1a237b68e57b864b502d003173938ea8ecdc255 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/symbolic.py @@ -0,0 +1,150 @@ +"""Implementation for higher-order operators.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from onnxscript.ir import convenience as ir_convenience + +import torch +from torch.onnx._internal._lazy_import import onnx_ir as ir +from torch.onnx._internal.exporter import _core +from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl +from torch.onnx.ops import _symbolic_impl + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +def _call_symbolic_op( + op_type: str, + domain: str, + args: Sequence[ir.Value | None], + kwargs: dict[str, int | float | str | bool | list[int] | list[float] | list[str]], + dtypes: Sequence[int], + version: int | None, + metadata_props: dict[str, str] | None, +) -> Sequence[ir.Value]: + """Call an operator with the given arguments and keyword arguments. + + Arguments are always inputs, while keyword arguments are attributes. + """ + # This is a wrapper around the IR node creation that hooks into the _builder.OpRecorder + # tracer so that all nodes created are recorded the same way as if we were to use + # onnxscript ops directly. + + if _core.current_tracer is None: + raise AssertionError("current_tracer must be non-None") + tracer = _core.current_tracer + + inputs = list(args) + + # If final inputs are None, strip them from the node inputs + for input in reversed(inputs): + if input is not None: + break + inputs.pop() + + # Construct and filter out None attributes + attributes = [ + attr + for attr in ir_convenience.convert_attributes(kwargs) # type: ignore[arg-type] + if attr.value is not None # type: ignore[union-attr] + ] + tracer.nodes.append( + node := ir.Node( + domain, + op_type, + inputs=inputs, + attributes=attributes, + num_outputs=len(dtypes), + version=version, + metadata_props=metadata_props, + ) + ) + # Set the dtypes for the outputs. We set them here because the graph builder + # Uses PyTorch types which are sometimes inaccurate when they are ONNX only + # types like float4e2m1. + for value, dtype in zip(node.outputs, dtypes): + value.dtype = ir.DataType(dtype) + # The shape is set by the graph builder. We don't need to set it here. + return node.outputs + + +@onnx_impl(torch.ops.onnx_symbolic._symbolic.default, no_compile=True) +def onnx_symbolic_symbolic( + inputs: Sequence[ir.Value | None], + op_type: str, + onnx_dtype: int, + *, + shape: Sequence[int | ir.Value], + attr_keys: Sequence[str], + attr_types: Sequence[str], + attr_pos: Sequence[tuple[int, int]], + attr_ints: Sequence[int], + attr_floats: Sequence[float], + attr_strs: Sequence[str], + metadata_props_keys: Sequence[str] = (), + metadata_props_values: Sequence[str] = (), + domain: str = "", + version: int | None = None, +) -> ir.Value: + del shape # Unused. The shapes are set by the graph builder + encoded = _symbolic_impl.EncodedAttrs( + attr_keys=list(attr_keys), + attr_types=list(attr_types), + attr_pos=list(attr_pos), + attr_ints=list(attr_ints), + attr_floats=list(attr_floats), + attr_strs=list(attr_strs), + ) + attrs = encoded.to_dict() + return _call_symbolic_op( + op_type, + domain, + inputs, + attrs, + dtypes=[onnx_dtype], + version=version, + metadata_props=dict(zip(metadata_props_keys, metadata_props_values)), + )[0] + + +@onnx_impl(torch.ops.onnx_symbolic._symbolic_multi_out.default, no_compile=True) +def onnx_symbolic_symbolic_multi_out( + inputs: Sequence[ir.Value | None], + op_type: str, + onnx_dtypes: Sequence[int], + *, + shapes: Sequence[Sequence[int | ir.Value]], + attr_keys: Sequence[str], + attr_types: Sequence[str], + attr_pos: Sequence[tuple[int, int]], + attr_ints: Sequence[int], + attr_floats: Sequence[float], + attr_strs: Sequence[str], + metadata_props_keys: Sequence[str] = (), + metadata_props_values: Sequence[str] = (), + domain: str = "", + version: int | None = None, +) -> Sequence[ir.Value]: + del shapes # Unused. The shapes are set by the graph builder + encoded = _symbolic_impl.EncodedAttrs( + attr_keys=list(attr_keys), + attr_types=list(attr_types), + attr_pos=list(attr_pos), + attr_ints=list(attr_ints), + attr_floats=list(attr_floats), + attr_strs=list(attr_strs), + ) + attrs = encoded.to_dict() + return _call_symbolic_op( + op_type, + domain, + inputs, + attrs, + dtypes=onnx_dtypes, + version=version, + metadata_props=dict(zip(metadata_props_keys, metadata_props_values)), + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/symops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/symops.py new file mode 100644 index 0000000000000000000000000000000000000000..cdaee46802768e93be405d533f2b14cd3f7a4cfd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_torchlib/ops/symops.py @@ -0,0 +1,64 @@ +"""Implementation for torch.sym* ops.""" + +# mypy: disable-error-code="misc,arg-type,type-arg,valid-type,assignment,return-value,type-var,operator,no-untyped-def,index" +# pyrefly: ignore-errors +# ruff: noqa: TCH001,TCH002,TC003 + +from __future__ import annotations + +from collections.abc import Sequence + +from onnxscript.onnx_opset import opset18 as op + +import torch +from torch.onnx._internal.exporter._torchlib._tensor_typing import ( + BOOL, + FLOAT, + IntType, + TensorType, + TTensor, +) +from torch.onnx._internal.exporter._torchlib._torchlib_registry import onnx_impl + + +@onnx_impl(torch.sym_float, trace_only=True) +def sym_float(self: TensorType) -> FLOAT: + """sym_float(SymInt self) -> SymFloat""" + return op.Cast(self, to=FLOAT.dtype) + + +@onnx_impl(torch.sym_max, trace_only=True) +def sym_max(x: IntType, y: IntType) -> IntType: + """sym_max(SymInt x, SymInt y) -> SymInt""" + return op.Max(x, y) + + +@onnx_impl(torch.sym_min, trace_only=True) +def sym_min(x: IntType, y: IntType) -> IntType: + """sym_min(SymInt x, SymInt y) -> SymInt""" + return op.Min(x, y) + + +@onnx_impl(torch.sym_not, trace_only=True) +def sym_not(self: BOOL) -> BOOL: + """sym_not(SymBool self) -> SymBool""" + return op.Not(self) + + +@onnx_impl(torch.sym_sum, trace_only=True) +def sym_sum(args: Sequence[IntType]) -> IntType: + """sym_sum(SymInt[] args) -> SymInt""" + if len(args) == 0: + return op.Constant(value_int=0) + if len(args) == 1: + return args[0] + result = op.Add(args[0], args[1]) + for i in range(2, len(args)): + result = op.Add(result, args[i]) + return result + + +@onnx_impl(torch.sym_ite, trace_only=True) +def sym_ite(b: BOOL, t: TTensor, f: TTensor) -> TTensor: + """sym_ite(SymBool b, Tensor t, Tensor f) -> Tensor""" + return op.Where(b, t, f) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_type_casting.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_type_casting.py new file mode 100644 index 0000000000000000000000000000000000000000..d108bccae2e054492ab47c6122a4e3a2b2706888 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_type_casting.py @@ -0,0 +1,34 @@ +import numpy as np + +import torch + + +def unpack_float4x2_as_uint8(tensor: torch.Tensor) -> np.ndarray: + """Convert a float4x2 tensor to unpacked uint8 np array.""" + if tensor.dtype != torch.float4_e2m1fn_x2: + raise AssertionError(f"Expected float4_e2m1fn_x2, got {tensor.dtype}") + data = tensor.view(torch.uint8).numpy(force=True).flatten() + result_size = tensor.numel() * 2 + result = np.empty([result_size], dtype=np.uint8) + array_low = data & np.uint8(0x0F) + array_high = data & np.uint8(0xF0) + array_high >>= np.uint8(4) + result[0::2] = array_low + result[1::2] = array_high + result.resize(get_float4_shape(tensor), refcheck=False) + return result + + +def get_float4_shape(tensor: torch.Tensor) -> tuple[int, ...]: + """Get the shape of an unpacked float4 tensor. + + The float4_e2m1fn_x2 type is a shell type described in + https://github.com/pytorch/pytorch/issues/146414. + + the shell dtype is takes up 1 byte per element and semantically represents + two fp4 values packed into 1 byte. Semantically it represents (*tensor.shape[:-1], tensor.shape[-1]*2) + fp4 elements. + """ + if tensor.dtype != torch.float4_e2m1fn_x2: + raise AssertionError(f"Expected float4_e2m1fn_x2, got {tensor.dtype}") + return (*tensor.shape[:-1], tensor.shape[-1] * 2) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_verification.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_verification.py new file mode 100644 index 0000000000000000000000000000000000000000..ac5e0418e078463c450cdcbd688ce4233a9823ef --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/exporter/_verification.py @@ -0,0 +1,344 @@ +from __future__ import annotations + + +__all__ = [ + "VerificationInfo", + "verify_onnx_program", +] + +import dataclasses +import logging +import math +from typing import Any, TYPE_CHECKING + +import torch +from torch.utils import _pytree + + +if TYPE_CHECKING: + from torch.onnx._internal._lazy_import import onnx_ir as ir + from torch.onnx._internal.exporter import _onnx_program + + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class VerificationInfo: + """Verification information for a value in the ONNX program. + + This class contains the maximum absolute difference, maximum relative difference, + and histograms of absolute and relative differences between the expected and actual + values. It also includes the expected and actual data types. + + The histograms are represented as tuples of tensors, where the first tensor is the + histogram counts and the second tensor is the bin edges. + + Attributes: + name: The name of the value (output or intermediate). + max_abs_diff: The maximum absolute difference between the expected and actual values. + max_rel_diff: The maximum relative difference between the expected and actual values. + abs_diff_hist: A tuple of tensors representing the histogram of absolute differences. + The first tensor is the histogram counts and the second tensor is the bin edges. + rel_diff_hist: A tuple of tensors representing the histogram of relative differences. + The first tensor is the histogram counts and the second tensor is the bin edges. + expected_dtype: The data type of the expected value. + actual_dtype: The data type of the actual value. + """ + + name: str + max_abs_diff: float + max_rel_diff: float + abs_diff_hist: tuple[torch.Tensor, torch.Tensor] + rel_diff_hist: tuple[torch.Tensor, torch.Tensor] + expected_dtype: torch.dtype + actual_dtype: torch.dtype + # NOTE: We don't need to include shape because the expected shape is already known + # and checked by the runtime + + @classmethod + def from_tensors( + cls, + name: str, + expected: torch.Tensor | float | int | bool, + actual: torch.Tensor | float | int | bool, + ) -> VerificationInfo: + """Create a VerificationInfo object from two tensors. + + Args: + name: The name of the value. + expected: The expected tensor. + actual: The actual tensor. + + Returns: + VerificationInfo: The VerificationInfo object. + """ + if not isinstance(expected, torch.Tensor): + expected = torch.tensor(expected) + if not isinstance(actual, torch.Tensor): + actual = torch.tensor(actual) + + max_abs_diff, max_rel_diff, abs_diff, rel_diff = _compare_tensors( + expected, actual + ) + bins = torch.tensor( + [0.0, 1e-6, 1e-5, 1e-4, 1e-3, 1e-2, 1e-1, 1.0, 10, 1000000], + dtype=torch.float, + ) + abs_diff_hist = torch.histogram(abs_diff.float(), bins=bins) + rel_diff_hist = torch.histogram(rel_diff.float(), bins=bins) + return cls( + name=name, + max_abs_diff=max_abs_diff, + max_rel_diff=max_rel_diff, + abs_diff_hist=abs_diff_hist, + rel_diff_hist=rel_diff_hist, + expected_dtype=expected.dtype, + actual_dtype=actual.dtype, + ) + + def asdict(self) -> dict[str, Any]: + """Convert the VerificationInfo object to a dictionary. + + Returns: + A dictionary representation of the VerificationInfo object. + """ + return { + "name": self.name, + "max_abs_diff": self.max_abs_diff, + "max_rel_diff": self.max_rel_diff, + "abs_diff_hist": [ + self.abs_diff_hist[0].tolist(), + self.abs_diff_hist[1].tolist(), + ], + "rel_diff_hist": [ + self.rel_diff_hist[0].tolist(), + self.rel_diff_hist[1].tolist(), + ], + "expected_dtype": str(self.expected_dtype), + "actual_dtype": str(self.actual_dtype), + } + + +def _compare_tensors( + expected: torch.Tensor, + actual: torch.Tensor, +) -> tuple[float, float, torch.Tensor, torch.Tensor]: + # Move tensors to the same device + expected = expected.detach().cpu() + actual = actual.detach().cpu() + if expected.numel() == 0 or actual.numel() == 0: + return math.inf, math.inf, torch.tensor(math.inf), torch.tensor(math.inf) + if expected.dtype == torch.bool: + expected = expected.to(torch.float32) + actual = actual.to(torch.float32) + if torch.is_complex(expected): + expected = torch.view_as_real(expected) + abs_diff = torch.abs(expected - actual) + eps = 1e-7 + normalizer = torch.abs(expected) + eps + rel_diff = abs_diff / normalizer + + max_absolute_difference = abs_diff.max().item() + max_relative_difference = rel_diff.max().item() + + return max_absolute_difference, max_relative_difference, abs_diff, rel_diff + + +def verify_onnx_program( + onnx_program: _onnx_program.ONNXProgram, + args: tuple[Any, ...] | None = None, + kwargs: dict[str, Any] | None = None, + compare_intermediates: bool = False, +) -> list[VerificationInfo]: + """Verify the ONNX model by comparing the values with the expected values from ExportedProgram. + + Args: + onnx_program: The ONNX program to verify. + args: The input arguments for the model. + kwargs: The keyword arguments for the model. + compare_intermediates: Whether to verify intermediate values. This is going + to take longer time, so it is disabled by default. + + Returns: + VerificationInfo objects containing the verification information for each value. + """ + exported_program = onnx_program.exported_program + if exported_program is None: + raise ValueError( + "The ONNX program does not contain an exported_program. " + "Please provide an exported_program to verify the ONNX program." + ) + if args is None and kwargs is None: + # User did not provide example inputs, use the default example inputs + if exported_program.example_inputs is None: + raise ValueError( + "No example inputs provided and the exported_program does not contain example inputs. " + "Please provide arguments to verify the ONNX program." + ) + args, kwargs = exported_program.example_inputs + if args is None: + args = () + if kwargs is None: + kwargs = {} + + # Flatten args for ONNX program and the VerificationInterpreter + flat_args, _ = exported_program._get_flat_args_with_check(args, kwargs) + + if not compare_intermediates: + # Compare the output values + torch_outputs, _ = _pytree.tree_flatten( + exported_program.module()(*args, **kwargs) + ) + onnx_outputs = onnx_program(*flat_args) + results = [] + for torch_output, onnx_output, output_val in zip( + torch_outputs, onnx_outputs, onnx_program.model.graph.outputs + ): + results.append( + VerificationInfo.from_tensors( + name=str(output_val.name), + expected=torch_output, + actual=onnx_output, + ) + ) + return results + + # Use the _VerificationInterpreter to get the intermediate values + # By design the output values are included too + interpreter = _VerificationInterpreter(onnx_program) + interpreter.run(*flat_args) + + return interpreter.verification_infos + + +def _create_value_mapping(graph: ir.Graph) -> dict[str, ir.Value]: + """Return a dictionary mapping names to values in the graph. + + The mapping does not include values from subgraphs. + + Args: + graph: The graph to extract the mapping from. + + Returns: + A dictionary mapping names to values. + """ + values: dict[str, ir.Value] = {} + values.update(graph.initializers) + # The names of the values can be None or "", which we need to exclude + for input in graph.inputs: + if not input.name: + continue + values[input.name] = input + for node in graph: + for value in node.outputs: + if not value.name: + continue + values[value.name] = value + return values + + +class _VerificationInterpreter(torch.fx.Interpreter): + """Interpreter for verifying converted ONNX model accuracy by comparing intermediate values. + + To compare models, first initialize the interpreter with an ONNX program. + Then, call the :meth:`run` method with the input arguments to execute the model. + The :meth:`run` method will execute the model and populate the + :attr:`verification_infos` attribute with the verification information for each value. + + :: + onnx_program = torch.onnx.export(model, args, dynamo=True) + interpreter = _VerificationInterpreter(onnx_program) + interpreter.run(*args) + verification_infos = interpreter.verification_infos + for info in verification_infos: + print("value name:", info.name, info) + + The verification information includes the maximum absolute difference, maximum relative + difference, and histograms of absolute and relative differences between the expected + and actual values. See :class:`VerificationInfo` for more details. + + Attributes: + verification_infos: A list of verification information for each value. + It is populated when the `run` method is called. + """ + + def __init__(self, onnx_program: torch.onnx.ONNXProgram) -> None: + """Initialize the _VerificationInterpreter with an ONNX program. + + Args: + onnx_program: The ONNX program to verify. + """ + if onnx_program.exported_program is None: + raise ValueError( + "The ONNX program does not contain an exported_program. " + "Please provide an exported_program to verify the ONNX program." + ) + super().__init__(onnx_program.exported_program.module()) + self._onnx_program = onnx_program + self._onnx_values = _create_value_mapping(onnx_program.model.graph) + self._args: tuple[Any, ...] = () + self.verification_infos: list[VerificationInfo] = [] + + def run( + self, + *args: Any, + initial_env: dict[torch.fx.Node, Any] | None = None, + enable_io_processing: bool = True, + ) -> Any: + """Run the interpreter with the given input arguments. + + This method executes the model and populates the :attr:`verification_infos` attribute + with the verification information for each value. + + Args: + args: The input arguments for the model. + initial_env: The initial environment for the interpreter. + enable_io_processing: Whether to enable IO processing. + + Returns: + Any: The result of executing the model. + """ + self.verification_infos = [] + self._args = args + return super().run( + *args, + initial_env=initial_env, + enable_io_processing=enable_io_processing, + ) + + def run_node(self, n: torch.fx.Node) -> Any: + result = super().run_node(n) + if n.op != "call_function": + return result + node_name = n.name + if node_name not in self._onnx_values: + return result + try: + (onnx_result,) = self._onnx_program.compute_values([node_name], self._args) + except Exception: + logger.warning( + "Failed to compute value for node %s", node_name, exc_info=True + ) + return result + info = VerificationInfo.from_tensors( + name=node_name, + expected=result, + actual=onnx_result, + ) + self.verification_infos.append(info) + if info.max_abs_diff > 0.01 or info.max_rel_diff > 0.1: + logger.warning( + "Verification info for node %s: max_abs_diff: %s, max_rel_diff: %s", + node_name, + info.max_abs_diff, + info.max_rel_diff, + ) + else: + logger.info( + "Verification info for node %s: max_abs_diff: %s, max_rel_diff: %s", + node_name, + info.max_abs_diff, + info.max_rel_diff, + ) + return result diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/_pass.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/_pass.py new file mode 100644 index 0000000000000000000000000000000000000000..95b7892fec4df0361165e86eb01fedfbb2f7e7a4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/_pass.py @@ -0,0 +1,237 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import abc +import contextlib +import dataclasses +import difflib +import io +import sys +from typing import Any, TYPE_CHECKING + +import torch +import torch.fx +from torch._subclasses.fake_tensor import unset_fake_temporarily + + +if TYPE_CHECKING: + from collections.abc import Callable + + from torch._subclasses import fake_tensor + + +@dataclasses.dataclass +class PackageInfo: + package_name: str + version: str | None + commit_hash: str | None + + def to_onnx_domain_string(self) -> str: + return ".".join( + filter(None, ("pkg", self.package_name, self.version, self.commit_hash)) + ) + + @classmethod + def from_python_class(cls, python_class_name: type | str) -> PackageInfo: + if isinstance(python_class_name, type): + python_class_name = python_class_name.__module__ + package_name = python_class_name.split(".")[0] + package = __import__(package_name) + version = getattr(package, "__version__", None) + # TODO: Figure out how to retrieve commit hash. + commit_hash = None + return cls(package_name, version, commit_hash) + + +@dataclasses.dataclass +class GraphModuleOnnxMeta: + package_info: PackageInfo + + +@contextlib.contextmanager +def _patch_difflib_sequence_matcher_init(): + """Context patching `difflib.SequenceMatcher` for fx readable graph. + + Under this context, the `autojunk` argument of `difflib.SequenceMatcher` will always + be considered as `False`. This is to prevent `difflib.SequenceMatcher` recognizing + stacktrace messages in fx readable graph as junk, as these messages tend to be long (>200) + and repeat multiple times, which falls under the junk filter criteria. + + `difflib.SequenceMatcher` is used underneath by all sorts of diffing functions + in `difflib`, including `difflib.unified_diff`, `difflib.ndiff`, `difflib.context_diff`. + Unfortunately, there is no way to pass `autojunk` argument to these functions, and + they all default to `True`. This context patching will affect all of them. + + `Reference: Automatic junk heuristic `_ + """ + original_init = difflib.SequenceMatcher.__init__ + + def patched_init(self, isjunk=None, a="", b="", autojunk=True) -> None: + original_init(self, isjunk, a, b, autojunk=False) + + difflib.SequenceMatcher.__init__ = patched_init # type: ignore[assignment] + try: + yield + finally: + difflib.SequenceMatcher.__init__ = original_init # type: ignore[assignment] + + +def _unified_diff(a: str, b: str) -> str: + """Return a string containing the unified diff of two strings. + + This function calls a patched version of `difflib.unified_diff` with `autojunk` set + to `False` for `difflib.SequenceMatcher` class. More details can be found in + `_patch_difflib_sequence_matcher_init` function. + + Args: + a: The first string. + b: The second string. + + Returns: + The unified diff of the two strings. If there is no diff, return "". + + Example:: + + >>> a = '''class GraphModule(torch.nn.Module): + ... def forward(self, input_ids : torch.Tensor, attention_mask : torch.Tensor): + ... # File: /modeling.py:770, code: input_ids = input_ids.view(-1, input_shape[-1]) + ... view = input_ids.view(-1, 3); input_ids = None + ... ''' + >>> b = '''class (torch.nn.Module): + ... def forward(self, input_ids: i64[1, 3], attention_mask: i64[1, 3]): + ... # File: /modeling.py:770, code: input_ids = input_ids.view(-1, input_shape[-1]) + ... view: i64[1, 3] = torch.ops.aten.view.default(input_ids, [-1, 3]); input_ids = None + ... ''' + >>> print(_unified_diff(a, b)) + --- + +++ + @@ -1,4 +1,4 @@ + -class GraphModule(torch.nn.Module): + - def forward(self, input_ids : torch.Tensor, attention_mask : torch.Tensor): + +class (torch.nn.Module): + + def forward(self, input_ids: i64[1, 3], attention_mask: i64[1, 3]): + # File: /modeling.py:770, code: input_ids = input_ids.view(-1, input_shape[-1]) + - view = input_ids.view(-1, 3); input_ids = None + + view: i64[1, 3] = torch.ops.aten.view.default(input_ids, [-1, 3]); input_ids = None + """ + + a_list = a.splitlines(keepends=True) + b_list = b.splitlines(keepends=True) + + with _patch_difflib_sequence_matcher_init(): + # Set `n` to `sys.maxsize` to show entire graph when there is a diff. + diff = "".join(difflib.unified_diff(a_list, b_list, n=sys.maxsize)) + + if not diff: + return "" + return diff + + +def _transform_diagnose_call_message_formatter( + run: Callable, + self: Transform, + *args: Any, + **kwargs: Any, +) -> str: + return f"Running {self.__class__.__name__} pass. " + + +def maybe_fx_graph_tabular(graph: torch.fx.Graph) -> str | None: + """Return the Graph nodes in tabular format. Equivalent to stdout of `graph.print_tabular()`. + If `tabulate` is not installed, return `None`. + + Args: + graph: The Graph to print. + + Returns: + The Graph printed in a tabular format. None if `tabulate` is not installed. + """ + f = io.StringIO() + with contextlib.redirect_stdout(f): + try: + graph.print_tabular() + except ImportError: + return None + return f.getvalue() + + +class Transform(abc.ABC): + """Base class for FX graph transformations to be used by FX-ONNX exporter. + + Similar to `FX Interpreter `_, + specializations of this class execute the FX graph Node-by-Node. + Methods in the `Transform` class can be overridden to customize the behavior of the model. + This pattern can be useful for many things, including writing code transformations as well as analysis passes. + + The following methods can be overridden:: + + _run() + +-- run_node() + +-- placeholder() + +-- get_attr() + +-- call_function() + +-- call_method() + +-- call_module() + +-- output() + + One important aspect to note is that if the transformation modifies the model input and/or output signature, + (e.g. additional inputs/outputs are added to the model), :class:`InputAdaptStep` and/or :class:`OutputAdaptStep` + are needed to reconcile :attr:`ONNXProgram.model_proto`. + That is, the model signature and the model representation must match. + + TODO(bowbao): Add more overridable methods in call hierarchy + TODO(bowbao): Create an example once more overridable methods are added. + """ + + module: torch.fx.GraphModule + """The module to be transformed.""" + + fake_mode: fake_tensor.FakeTensorMode | None + """The existing fake mode detected from `self.module`.""" + + def __init__( + self, + module: torch.fx.GraphModule, + ) -> None: + """Initialize the transform. + + Args: + module: The module to be transformed. + """ + self.module = module + self.fake_mode = self._detect_fake_mode() + + def _detect_fake_mode(self) -> fake_tensor.FakeTensorMode | None: + """Detect fake mode from the graph. + + Scan through all nodes in graph and their meta['val'] to detect fake mode. + """ + fake_tensors = [node.meta.get("val") for node in self.module.graph.nodes] + with unset_fake_temporarily(): + return torch._dynamo.utils.detect_fake_mode(fake_tensors) + + def _maybe_fakefy_args( + self, fake_mode: fake_tensor.FakeTensorMode | None, *args: Any + ) -> tuple[Any, ...]: + if fake_mode is None: + return args + # NB: This should hit the cache if tensors were fakefied before. + # E.g., when the fx graph is produced by Dynamo. + return tuple( + fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t for t in args + ) + + @abc.abstractmethod + def _run(self, *args, **kwargs) -> torch.fx.GraphModule: ... + + def run(self, *args, **kwargs) -> torch.fx.GraphModule: + """Run the transform on `self.module`. + + Note that this method may or may not mutate `self.module`, and the returned + `GraphModule` could be either `self.module` or a new `GraphModule`. + + Args: + *args: Positional arguments for `self.module` to run. + **kwargs: Keyword arguments for `self.module` to run. + """ + return self._run(*args, **kwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/passes/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/passes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eff83563a5a0852f68db92c5059d6eb2f17faa72 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/passes/__init__.py @@ -0,0 +1,6 @@ +from .type_promotion import InsertTypePromotion + + +__all__ = [ + "InsertTypePromotion", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/passes/type_promotion.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/passes/type_promotion.py new file mode 100644 index 0000000000000000000000000000000000000000..1f500e170cf749f67297786890a33c0fcb08455e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/passes/type_promotion.py @@ -0,0 +1,1682 @@ +# mypy: allow-untyped-defs +# Owner(s): ["module: onnx"] +from __future__ import annotations + +import abc +import dataclasses +import inspect +import logging +from typing import Any, TYPE_CHECKING + +import torch +import torch._dispatch.python +import torch._ops +import torch.fx +import torch.fx.traceback as fx_traceback +from torch import _prims_common, _refs +from torch._prims_common import ( + ELEMENTWISE_TYPE_PROMOTION_KIND, + wrappers as _prims_common_wrappers, +) +from torch._refs import linalg as _linalg_refs, nn as _nn_refs, special as _special_refs +from torch._refs.nn import functional as _functional_refs +from torch.fx.experimental import proxy_tensor +from torch.onnx._internal.fx import _pass, type_utils as fx_type_utils +from torch.utils import _python_dispatch, _pytree + + +if TYPE_CHECKING: + from collections.abc import Callable, Mapping, Sequence + from types import ModuleType + + from torch._subclasses import fake_tensor + + +logger = logging.getLogger(__name__) + + +def _try_getclosurevars(func): + try: + return inspect.getclosurevars(func) + except TypeError: + return None + + +@dataclasses.dataclass +class TypePromotionSnapshot: + """Type promotion snapshot for a fx node and its inputs. + + Contains the promoted dtype for args and kwargs that needs promoting. + Contains the expected node output dtype. + """ + + args_dtypes: Mapping[int, torch.dtype] + """Mapping from arg position to dtype to promote to.""" + + kwargs_dtypes: Mapping[str, torch.dtype] + """Mapping from kwarg name to dtype to promote to.""" + + out_dtype: torch.dtype + """Expected output dtype of the node.""" + + +class TypePromotionRule(abc.ABC): + """Base class for type promotion rule per 'torch.ops.{namespace}.{op_name}'.""" + + def __init__(self, namespace: str, op_name: str) -> None: + self.namespace = namespace + self.op_name = op_name + + # Make this abstract as well because subclass needs to override __eq__(). + # A class that overrides __eq__() and does not define __hash__() will have its __hash__() implicitly set to None. + # Ref: https://docs.python.org/3/reference/datamodel.html#object.__hash__ + @abc.abstractmethod + def __hash__(self) -> int: ... + + @abc.abstractmethod + def __repr__(self) -> str: ... + + @abc.abstractmethod + def __eq__(self, other: object) -> bool: ... + + def is_valid(self) -> bool: + """Check if the rule is valid.""" + # This always returns a module. If the module does not exist it will be created. + module = getattr(torch.ops, self.namespace) + py_op = getattr(module, self.op_name, None) + if py_op is None: + logger.warning( + "Cannot find op: %s in module: %s", self.op_name, self.namespace + ) + return False + if not isinstance(py_op, torch._ops.OpOverloadPacket): + logger.warning( + "Op: torch.ops.%s.%s is not an OpOverloadPacket, got: %s", + self.namespace, + self.op_name, + type(py_op), + ) + return False + + return True + + @abc.abstractmethod + def preview_type_promotion( + self, args: tuple, kwargs: dict + ) -> TypePromotionSnapshot: + """Preview type promotion results for provided set of args and kwargs. + + Returns a TypePromotionSnapshot object that contains the promoted dtypes for + the arguments and the expected output dtype. + """ + ... + + +class ElementwiseTypePromotionRule(TypePromotionRule): + """Defines how to perform elementwise type promotion for 'torch.ops.{namespace}.{op_name}'.""" + + _USE_OPMATH: bool = False + """Whether to use opmath to compute the promoted input dtype. + If used, upcasts will be inserted everywhere for lower precision models. + Set to False and have torchlib handle upcasts in op implementation internally. + """ + + def __init__( + self, + namespace: str, + op_name: str, + promote_args_positions: Sequence[int], + promote_kwargs_names: Sequence[str], + promotion_kind: _prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND, + ) -> None: + """Constructs a TypePromotionRule for elementwise operators. + + Args: + namespace: Namespace of the op. E.g. 'aten' in 'torch.ops.aten.add'. + op_name: Name of the op. E.g. 'add' in 'torch.ops.aten.add'. + promote_args_positions: Positions of args to promote. + promote_kwargs_names: Names of kwargs to promote. + promotion_kind: Type promotion kind. Refer to [_prims_common.elementwise_dtypes](https://github.com/pytorch/pytorch/blob/main/torch/_prims_common/__init__.py) for detail. # noqa: B950 + """ + super().__init__(namespace, op_name) + self.promote_args_positions = promote_args_positions + self.promote_kwargs_names = promote_kwargs_names + self.promotion_kind = promotion_kind + + def __repr__(self) -> str: + return ( + f"ElementwiseTypePromotionRule('{self.namespace}', '{self.op_name}', " + f"{self.promote_args_positions}, {self.promote_kwargs_names}, {self.promotion_kind})" + ) + + # pyrefly: ignore [bad-override] + def __eq__(self, other: object, /) -> bool: + if not isinstance(other, ElementwiseTypePromotionRule): + return False + return ( + self.namespace == other.namespace + and self.op_name == other.op_name + and self.promote_args_positions == other.promote_args_positions + and self.promote_kwargs_names == other.promote_kwargs_names + and self.promotion_kind == other.promotion_kind + ) + + def __hash__(self) -> int: + return f"{type(self)}:{self.namespace}.{self.op_name}".__hash__() + + def _consolidate_input_dtype( + self, computed_dtype: torch.dtype, result_dtype: torch.dtype + ) -> torch.dtype: + """ + Although opmath is the right thing to do to retain on-par precision, it inserts + upcasts everywhere in the graph. This is particularly hard for backend to optimize + since there is no way to differentiate between inserted upcasts and model code + casts. Hence we consolidate the input dtype to the result dtype to avoid this. + """ + if not self._USE_OPMATH and self.promotion_kind in ( + _prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + _prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, + ): + return result_dtype + return computed_dtype + + def preview_type_promotion( + self, args: tuple, kwargs: dict + ) -> TypePromotionSnapshot: + candidate_args = { + i: args[i] + for i in self.promote_args_positions + if i < len(args) and args[i] is not None + } + candidate_kwargs = { + name: kwargs[name] + for name in self.promote_kwargs_names + if name in kwargs and kwargs[name] is not None + } + + computed_dtype, result_dtype = _prims_common.elementwise_dtypes( + *_pytree.arg_tree_leaves(*candidate_args.values(), **candidate_kwargs), + type_promotion_kind=self.promotion_kind, + ) + + consolidated_input_dtype = self._consolidate_input_dtype( + computed_dtype, result_dtype + ) + + return TypePromotionSnapshot( + dict.fromkeys(candidate_args.keys(), consolidated_input_dtype), + dict.fromkeys(candidate_kwargs.keys(), consolidated_input_dtype), + result_dtype, + ) + + +class DivElementwiseTypePromotionRule(ElementwiseTypePromotionRule): + """Reference type promotion rule from torch._refs.div. + + Rule depends on the value of the `rounding_mode` argument. + """ + + def __init__(self) -> None: + super().__init__( + "aten", + "div", + promote_args_positions=(0, 1), + promote_kwargs_names=(), + promotion_kind=_prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + ) + + def preview_type_promotion( + self, args: tuple, kwargs: dict + ) -> TypePromotionSnapshot: + rounding_mode = kwargs.get("rounding_mode") + if rounding_mode is None: + # true_divide + self.promotion_kind = ( + _prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ) + return super().preview_type_promotion(args, kwargs) + if rounding_mode == "trunc": + # trunc_divide + self.promotion_kind = _prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + return super().preview_type_promotion(args, kwargs) + if rounding_mode == "floor": + # floor_divide + self.promotion_kind = _prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + return super().preview_type_promotion(args, kwargs) + raise ValueError(f"Unknown rounding_mode: {rounding_mode}") + + +class ReductionTypePromotionRule(TypePromotionRule): + def __init__( + self, + namespace: str, + op_name: str, + promotion_kind: _prims_common.REDUCTION_OUTPUT_TYPE_KIND, + ) -> None: + """Constructs a TypePromotionRule for reduction operators. + + Args: + namespace: Namespace of the op. E.g. 'aten' in 'torch.ops.aten.sum'. + op_name: Name of the op. E.g. 'sum' in 'torch.ops.aten.sum'. + promotion_kind: Type promotion kind. Refer to [_prims_common.reduction_dtypes]((https://github.com/pytorch/pytorch/blob/main/torch/_prims_common/__init__.py)) for detail. # noqa: B950 + """ + super().__init__(namespace, op_name) + self.promotion_kind = promotion_kind + + def __repr__(self) -> str: + return f"ReductionTypePromotionRule('{self.namespace}', '{self.op_name}', {self.promotion_kind})" + + # pyrefly: ignore [bad-override] + def __eq__(self, other: object, /) -> bool: + if not isinstance(other, ElementwiseTypePromotionRule): + return False + return ( + self.namespace == other.namespace + and self.op_name == other.op_name + and self.promotion_kind == other.promotion_kind + ) + + def __hash__(self) -> int: + return f"{type(self)}:{self.namespace}.{self.op_name}".__hash__() + + def preview_type_promotion( + self, args: tuple, kwargs: dict + ) -> TypePromotionSnapshot: + if len(args) < 1: + raise AssertionError( + f"Reduction op torch.ops.{self.namespace}.{self.op_name} expects at least one argument" + ) + arg = args[0] + if not isinstance(arg, torch.Tensor): + raise AssertionError(f"{type(arg)=} is not torch.Tensor") + dtype: torch.dtype | None = kwargs.get("dtype") + + computation_dtype, result_dtype = _prims_common.reduction_dtypes( + arg, self.promotion_kind, dtype + ) + if result_dtype is None: + # Inspecting code, this can only happen when `promotion_kind` is `KEEP_PROMOTED_TYPE`. + # Hence set same as computation_dtype. + result_dtype = computation_dtype + + return TypePromotionSnapshot( + {0: computation_dtype}, + {}, + result_dtype, + ) + + +class AllOrAnyReductionTypePromotionRule(ReductionTypePromotionRule): + """Reference type promotion rule from torch.ops.aten.all or torch.ops.aten.any. + + This is a special case where computation dtype is always torch.bool. + The result dtype is always uint8 if `dtype` kwarg is uint8, otherwise torch.bool. + """ + + def __init__(self, op_name: str) -> None: + super().__init__( + "aten", + op_name, + _prims_common.REDUCTION_OUTPUT_TYPE_KIND.ALWAYS_BOOL, + ) + + def preview_type_promotion( + self, args: tuple, kwargs: dict + ) -> TypePromotionSnapshot: + if len(args) < 1: + raise AssertionError( + f"Reduction op torch.ops.{self.namespace}.{self.op_name} expects at least one argument" + ) + arg = args[0] + if not isinstance(arg, torch.Tensor): + raise AssertionError(f"{type(arg)=} is not torch.Tensor") + computation_dtype = torch.bool + # Preserves uint8 -- probably a legacy mask thing + result_dtype = torch.uint8 if arg.dtype == torch.uint8 else torch.bool + return TypePromotionSnapshot( + {0: computation_dtype}, + {}, + result_dtype, + ) + + +class SumLikeReductionTypePromotionRule(ReductionTypePromotionRule): + """Reference type promotion rule from torch.ops.aten.sum. + + This is a special case where computation dtype is always torch.int64 for integral arg, + unless overridden by `dtype` kwarg. + """ + + def preview_type_promotion( + self, args: tuple, kwargs: dict + ) -> TypePromotionSnapshot: + if len(args) < 1: + raise AssertionError( + f"Reduction op torch.ops.{self.namespace}.{self.op_name} expects at least one argument" + ) + arg = args[0] + if not isinstance(arg, torch.Tensor): + raise AssertionError(f"{type(arg)=} is not torch.Tensor") + dtype: torch.dtype | None = kwargs.get("dtype") + # The below logic is copied from `torch/_refs/__init__.py` reduction ops impl. + if dtype is None: + if _prims_common.is_boolean_dtype( + arg.dtype + ) or _prims_common.is_integer_dtype(arg.dtype): + dtype = torch.int64 + else: + dtype = arg.dtype + return super().preview_type_promotion(args, {"dtype": dtype}) + + +# NOTE: [Update type promotion rule] +# BELOW TABLE IS GENERATED FROM `TypePromotionRuleSetGenerator.generate_from_torch_refs`. +# DO NOT EDIT MANUALLY !!! +# For missing rules or discrepancies, please +# 1. Run `pytest test/onnx/test_fx_type_promotion.py` to validate if the generated rule set is current. +# If it is not, update with new generated set. +# 2. If discrepancies still exist, consider debugging torch._refs or report a bug. +# 3. If rules are still missing, add them to `_EXTRA_TYPE_PROMOTION_RULE_SET` or report a bug. +# Check `TypePromotionRule` class for how each rule is defined and used. +_GENERATED_ATEN_TYPE_PROMOTION_RULE_SET = { + ElementwiseTypePromotionRule( + "aten", "abs", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "abs_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "acos", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "acos_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "acosh", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "acosh_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "add", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "add_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "addcdiv", [0, 1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "addcdiv_", [0, 1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "addcmul", [0, 1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "addcmul_", [0, 1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "addr", [0, 1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "asin", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "asin_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "asinh", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "asinh_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "atan", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "atan2", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "atan2_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "atan_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "atanh", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "atanh_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "bitwise_and", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "bitwise_and_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", + "bitwise_left_shift", + [0, 1], + [], + ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + ), + ElementwiseTypePromotionRule( + "aten", + "bitwise_left_shift_", + [0, 1], + [], + ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + ), + ElementwiseTypePromotionRule( + "aten", "bitwise_not", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "bitwise_not_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "bitwise_or", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "bitwise_or_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", + "bitwise_right_shift", + [0, 1], + [], + ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + ), + ElementwiseTypePromotionRule( + "aten", + "bitwise_right_shift_", + [0, 1], + [], + ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + ), + ElementwiseTypePromotionRule( + "aten", "bitwise_xor", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "bitwise_xor_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "cat", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.NO_OPMATH + ), + ElementwiseTypePromotionRule( + "aten", "cauchy", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "cauchy_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "ceil", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "ceil_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "celu", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "celu_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "clamp", [0, 1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "clamp_", [0, 1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "copysign", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "copysign_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "cos", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "cos_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "cosh", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "cosh_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "deg2rad", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "deg2rad_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "digamma", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "digamma_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "dot", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "elu", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "elu_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "eq", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "eq_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "erf", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "erf_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "erfc", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "erfc_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "erfinv", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "erfinv_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "exp", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "exp2", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "exp2_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "exp_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "expm1", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "expm1_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "exponential", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "exponential_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "fill", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.NO_OPMATH + ), + ElementwiseTypePromotionRule( + "aten", "floor", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "floor_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "floor_divide", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "floor_divide_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "fmax", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "fmin", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "fmod", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "fmod_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "frac", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "frac_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "gcd", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "gcd_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "ge", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "ge_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "gelu", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "geometric", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "geometric_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "glu", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "gt", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "gt_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "hardtanh", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "heaviside", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "heaviside_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "huber_loss", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "hypot", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "hypot_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "i0", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "i0_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "igamma", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "igamma_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "igammac", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "igammac_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "isfinite", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "isinf", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "isnan", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "isneginf", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "isposinf", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "isreal", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "l1_loss", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "lcm", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "lcm_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "le", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "le_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "leaky_relu", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "lerp", [0, 1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "lerp_", [0, 1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "lgamma", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "lgamma_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "log", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "log10", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "log10_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "log1p", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "log1p_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "log2", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "log2_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "log_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "log_normal", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "log_normal_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "logaddexp", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "logaddexp2", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "logical_and", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "logical_and_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "logical_not", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "logical_not_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "logical_or", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "logical_or_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "logical_xor", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "logical_xor_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "logit", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "logsumexp", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "lt", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "lt_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "maximum", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "minimum", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "mish", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "mish_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "mse_loss", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "mul", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "mul_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "ne", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "ne_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "neg", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "neg_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "nextafter", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.NO_OPMATH + ), + ElementwiseTypePromotionRule( + "aten", "nextafter_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.NO_OPMATH + ), + ElementwiseTypePromotionRule( + "aten", "nll_loss", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "normal", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "pdist", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", + "poisson_nll_loss", + [0, 1], + [], + ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT, + ), + ElementwiseTypePromotionRule( + "aten", "prelu", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "rad2deg", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "rad2deg_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "reciprocal", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "reciprocal_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "relu", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "remainder", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "remainder_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "round", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "rsqrt", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "rsqrt_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "selu", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "selu_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "sgn", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "sgn_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "sigmoid", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "sigmoid_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "sign", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "sign_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "signbit", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.ALWAYS_BOOL + ), + ElementwiseTypePromotionRule( + "aten", "sin", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "sin_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "sinc", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "sinc_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "sinh", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "sinh_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", + "smooth_l1_loss", + [0, 1], + [], + ELEMENTWISE_TYPE_PROMOTION_KIND.COMPLEX_TO_FLOAT, + ), + ElementwiseTypePromotionRule( + "aten", "softplus", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "sqrt", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "sqrt_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "square", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.BOOL_TO_LONG + ), + ElementwiseTypePromotionRule( + "aten", "square_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.BOOL_TO_LONG + ), + ElementwiseTypePromotionRule( + "aten", "sub", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "sub_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "tan", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "tan_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "tanh", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "tanh_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "threshold", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "threshold_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "true_divide", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "true_divide_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "trunc", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "trunc_", [0], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "vdot", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT + ), + ElementwiseTypePromotionRule( + "aten", "where", [1, 2], [], ELEMENTWISE_TYPE_PROMOTION_KIND.NO_OPMATH + ), + ElementwiseTypePromotionRule( + "aten", "xlogy", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), + ElementwiseTypePromotionRule( + "aten", "xlogy_", [0, 1], [], ELEMENTWISE_TYPE_PROMOTION_KIND.INT_TO_FLOAT + ), +} + +# Manually curated extra type promotion rules. Please see NOTE [Update type promotion rule] +# before adding new rules. +_EXTRA_TYPE_PROMOTION_RULE_SET = { + # torch._refs skips type promotion decoration for `clamp_min` and `clamp_max` since + # the call is routed to the decorated `aten.clamp` op. + ElementwiseTypePromotionRule( + "aten", + "clamp_max", + promote_args_positions=(0, 1), + promote_kwargs_names=(), + promotion_kind=_prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + ), + ElementwiseTypePromotionRule( + "aten", + "clamp_min", + promote_args_positions=(0, 1), + promote_kwargs_names=(), + promotion_kind=_prims_common.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT, + ), + # torch.ops.aten.div.Tensor_mode applies different type promotion rules + # depending on the value of the `mode` argument. + DivElementwiseTypePromotionRule(), + # Manually curating reduction ops since the logic is written inside the op reference + # implementation. + AllOrAnyReductionTypePromotionRule("all"), + AllOrAnyReductionTypePromotionRule("any"), + ReductionTypePromotionRule( + "aten", + "amax", + promotion_kind=_prims_common.REDUCTION_OUTPUT_TYPE_KIND.SAME, + ), + ReductionTypePromotionRule( + "aten", + "amin", + promotion_kind=_prims_common.REDUCTION_OUTPUT_TYPE_KIND.SAME, + ), + # torch.ops.aten.mean is a special case that does not need type promotion. + ReductionTypePromotionRule( + "aten", + "std", + promotion_kind=_prims_common.REDUCTION_OUTPUT_TYPE_KIND.COMPLEX_TO_FLOAT, + ), + ReductionTypePromotionRule( + "aten", + "std_mean", + promotion_kind=_prims_common.REDUCTION_OUTPUT_TYPE_KIND.COMPLEX_TO_FLOAT, + ), + ReductionTypePromotionRule( + "aten", + "var", + promotion_kind=_prims_common.REDUCTION_OUTPUT_TYPE_KIND.COMPLEX_TO_FLOAT, + ), + SumLikeReductionTypePromotionRule( + "aten", + "cumprod", + promotion_kind=_prims_common.REDUCTION_OUTPUT_TYPE_KIND.SAME, + ), + SumLikeReductionTypePromotionRule( + "aten", + "cumsum", + promotion_kind=_prims_common.REDUCTION_OUTPUT_TYPE_KIND.SAME, + ), + SumLikeReductionTypePromotionRule( + "aten", + "prod", + promotion_kind=_prims_common.REDUCTION_OUTPUT_TYPE_KIND.SAME, + ), + SumLikeReductionTypePromotionRule( + "aten", + "sum", + promotion_kind=_prims_common.REDUCTION_OUTPUT_TYPE_KIND.SAME, + ), +} + + +class ElementwiseTypePromotionRuleSetGenerator: + """Hackly distilling info from reference ops decorated with elementwise type promotion rule. + + The goal is to retrieve the decorator + + ```python + @elementwise_type_promotion_wrapper( + type_promoting_args=("a", "b"), + type_promotion_kind=type_promotion_kind, + ) + ``` + + from the reference ops. It provides info as for which arguments are promoted + and what kind of promotion is applied. + """ + + @classmethod + def generate_from_torch_refs(cls) -> set[ElementwiseTypePromotionRule]: + """Parse type promotion rules from reference ops under torch._C._refs.""" + rule_set = set() + rule_set.update(cls._parse_torch_refs(_refs)) + rule_set.update(cls._parse_torch_refs(_nn_refs)) + rule_set.update(cls._parse_torch_refs(_linalg_refs)) + rule_set.update(cls._parse_torch_refs(_special_refs)) + rule_set.update(cls._parse_torch_refs(_functional_refs)) + return rule_set + + @classmethod + def _parse_torch_refs( + cls, ref_module: ModuleType + ) -> set[ElementwiseTypePromotionRule]: + logger.info("Processing module: %s", ref_module.__name__) + rule_set = set() + for name in ref_module.__all__: + decorated_op = getattr(ref_module, name) + rule = cls._parse_type_promotion_rule_from_refs_op(decorated_op) + if rule is not None and rule.is_valid(): + rule_set.add(rule) + + return rule_set + + @classmethod + def _parse_type_promotion_rule_from_refs_op( + cls, + decorated_op: Callable, + ) -> ElementwiseTypePromotionRule | None: + """Retrieve and parse type promotion decorator from op under torch._refs.""" + fn = decorated_op + type_promo_wrapper = None + while fn_closure_vars := _try_getclosurevars(fn): + if "fn" not in fn_closure_vars.nonlocals: + break + if "self" in fn_closure_vars.nonlocals and isinstance( + fn_closure_vars.nonlocals["self"], + _prims_common_wrappers.elementwise_type_promotion_wrapper, + ): + type_promo_wrapper = fn_closure_vars.nonlocals["self"] + break + fn = fn_closure_vars.nonlocals["fn"] + + if type_promo_wrapper is not None: + signature = inspect.signature(decorated_op) + + pos = 0 + promote_args_positions = [] + promote_kwargs_names = [] + + if type_promo_wrapper.type_promoting_arg_names is not None: + for name, param in signature.parameters.items(): + if name in type_promo_wrapper.type_promoting_arg_names: + if param.kind in ( + param.POSITIONAL_OR_KEYWORD, + param.POSITIONAL_ONLY, + ): + promote_args_positions.append(pos) + elif param.kind == param.KEYWORD_ONLY: + promote_kwargs_names.append(name) + pos += 1 + + return ElementwiseTypePromotionRule( + "aten", + decorated_op.__name__, + promote_args_positions=promote_args_positions, + promote_kwargs_names=promote_kwargs_names, + promotion_kind=type_promo_wrapper.type_promotion_kind, + ) + + logger.warning( + "Cannot find type promotion rule for: %s.%s", + decorated_op.__module__, + decorated_op.__name__, + ) + return None + + +class TypePromotionTable: + """Type promotion table for torch.ops.""" + + def __init__(self) -> None: + self._rule_table = {} + for rule in _GENERATED_ATEN_TYPE_PROMOTION_RULE_SET: + self.add_rule(rule) + for rule in _EXTRA_TYPE_PROMOTION_RULE_SET: + self.add_rule(rule) + + def add_rule(self, rule: TypePromotionRule) -> None: + """Add a type promotion rule for a python op in a torch.ops module. + + Args: + rule: Type promotion rule. + module: Module containing the op. E.g. torch.ops.aten. + + Raises: + ValueError: If the rule is invalid. + """ + if not rule.is_valid(): + raise ValueError(f"Invalid type promotion rule: {rule}") + self._rule_table[f"{rule.namespace}.{rule.op_name}"] = rule + + def get_rule(self, py_op: torch._ops.OpOverloadPacket) -> TypePromotionRule | None: + """Get type promotion rule for a python op under 'torch.ops.'.""" + return self._rule_table.get(str(py_op), None) + + +def get_type_promotion_rule( + node: torch.fx.Node, + type_promotion_table: TypePromotionTable, +) -> TypePromotionRule | None: + """Get type promotion rule for a node. + + Args: + node: Node to get type promotion rule for. + type_promotion_table: Type promotion table. + + Returns: + Type promotion rule for the node. None if no rule is found or if the node is not + representing a torch operator. + """ + op = node.target + if not isinstance(op, torch._ops.OpOverload): + return None + if (rule := type_promotion_table.get_rule(op.overloadpacket)) is None: + return None + + return rule + + +class _OpTraceDispatchMode(_python_dispatch.TorchDispatchMode): + """Trace ops that were dispatched. + + Utilize the dispatch mechanism in [`__torch_dispatch__`](https://dev-discuss.pytorch.org/t/what-and-why-is-torch-dispatch/557) + to trace op overloads that were dispatched to. This is used to find the compatible + op overload for a given op overload packet for different set of args and kwargs. + """ + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.traced_ops = [] + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + self.traced_ops.append(func) + return func(*args, **kwargs) + + +def find_compatible_op_overload( + op: torch._ops.OpOverloadPacket, args: tuple, kwargs: dict +) -> torch._ops.OpOverload: + """Find compatible OpOverload for an OpOverloadPacket using provided args and kwargs. + + Each "call_function" fx.Node in the fx.GraphModule has a target that represents a torch._ops.OpOverload. + The OpOverload contains an OpOverloadPacket that holds all the available overloads for the operation. + + During the type promotion pass, there are cases where the types of the args and kwargs may change, + such as promoting Python numbers to tensors. Consequently, the original OpOverload might not be + compatible with the updated args and kwargs. This function is used to identify the compatible + OpOverload for the given args and kwargs. + + Args: + op: OpOverloadPacket to find compatible OpOverload for. + args: The positional arguments to consider for compatibility. + kwargs: The keyword arguments to consider for compatibility. + + Returns: + torch._ops.OpOverload: The compatible OpOverload found for the given args and kwargs. + + Raises: + RuntimeError: If no compatible op overload is found. + + Examples: + >>> import torch + >>> packet = torch.ops.aten.pow + >>> args = (torch.tensor([1.0, 2.0]), 2) + >>> find_compatible_op_overload(packet, args, {})._overloadname + 'Tensor_Scalar' + >>> args = (torch.tensor([1.0, 2.0]), torch.tensor(2.0)) + >>> find_compatible_op_overload(packet, args, {})._overloadname + 'Tensor_Tensor' + """ + # Utilize the dispatch mechanism to find the compatible op overload. + op_trace_dispatch_mode = _OpTraceDispatchMode() + with op_trace_dispatch_mode: + op(*args, **kwargs) + if len(op_trace_dispatch_mode.traced_ops) < 1: + raise AssertionError("Expected at least 1 traced op, got 0") + + new_op_overload = op_trace_dispatch_mode.traced_ops[0] + if not isinstance(new_op_overload, torch._ops.OpOverload): + raise AssertionError(f"Expected OpOverload, got {type(new_op_overload)}") + if new_op_overload.overloadpacket != op: + raise AssertionError( + f"Expected same OpOverload packet, got {new_op_overload.overloadpacket} != {op}" + ) + + return new_op_overload + + +class _TypePromotionInterpreter(torch.fx.Interpreter): + """Interpreter that inserts type promotion for each node.""" + + def __init__( + self, + module: torch.fx.GraphModule, + type_promotion_table: TypePromotionTable, + ) -> None: + super().__init__(module) + self.type_promotion_table = type_promotion_table + + def _run_node_and_set_meta(self, node) -> Any: + """Run node and set meta according to `fx_traceback.get_current_meta()`. + + This should be used on new nodes or nodes that have been modified. + By default `Interpreter.run_node` does not update `node.meta`. + Set `node.meta` to the current meta, except for `node.meta["val"]`, which is + recomputed. + """ + out = super().run_node(node) + # Update interpreter env state with new output value. + self.env[node] = out + node.meta.update( + (k, v) + for k, v in fx_traceback.get_current_meta().items() + if k not in node.meta + ) + node.meta["val"] = proxy_tensor.extract_val(out) + return out + + def _create_node( + self, + graph: torch.fx.Graph, + op_type: str, + target: torch.fx.node.Target, + args: tuple, + kwargs: dict, + ) -> torch.fx.Node: + """Create a node and set its metadata.""" + if op_type not in ( + "call_function", + "call_method", + "get_attr", + "call_module", + "placeholder", + "output", + ): + raise AssertionError(f"Unexpected op_type: {op_type}") + node = getattr(graph, op_type)(target, args, kwargs) + self._run_node_and_set_meta(node) + return node + + def _rerun_node_after_type_promotion( + self, + node: torch.fx.Node, + expected_out_dtype: torch.dtype, + ) -> None: + """Rerun a node after type promotion and update node.meta["val"] with the output value.""" + node_val = node.meta.get("val", None) + if node_val is None: + raise AssertionError(f"Node {node} node.meta['val'] is not set.") + args, kwargs = self.fetch_args_kwargs_from_env(node) + target = node.target + if not isinstance(target, torch._ops.OpOverload): + raise AssertionError(f"Expected OpOverload, got {type(target)}") + node.target = find_compatible_op_overload(target.overloadpacket, args, kwargs) + + new_node_val = self._run_node_and_set_meta(node) + if not isinstance(new_node_val, type(node_val)): + raise AssertionError( + f"run_node output type should not change between runs. " + f"Got {type(new_node_val)}, expect {type(node_val)}." + ) + + if isinstance(node_val, torch.Tensor): + prev_node_dtype = node_val.dtype + + if prev_node_dtype != expected_out_dtype: + raise AssertionError( + f"node.meta['val'].dtype({prev_node_dtype}) does not agree with " + f"type promotion rule({expected_out_dtype})." + ) + + if new_node_val.dtype != expected_out_dtype: + # With explicit type promotion, the expected result dtype may not be + # the same as the computation dtype. This is referred to as "op math". + # We need to explicitly cast the output back to the expected dtype. + # See more about "op math" topic at `_prims_common.elementwise_dtypes`. + graph = node.graph + with graph.inserting_after(node): + output_cast_node = self._create_node( + graph, + "call_function", + torch.ops.prims.convert_element_type.default, + (node,), + {"dtype": expected_out_dtype}, + ) + node.replace_all_uses_with(output_cast_node) + output_cast_node.args = (node,) + logger.info( + "Node '%s' output dtype becomes %s due to op math. " + "Cast back to %s.", + node, + new_node_val.dtype, + expected_out_dtype, + ) + + elif fx_type_utils.is_torch_symbolic_type(node_val): + raise NotImplementedError( + "Type promotion does not support node output of sym types." + ) + elif isinstance(node_val, (list, tuple)): + raise NotImplementedError( + "Type promotion does not support node output of list or tuple." + ) + else: + raise RuntimeError(f"Unexpected node output type: {type(node_val)}.") + + def _maybe_promote_arg( + self, + node: torch.fx.Node, + fx_arg: torch.fx.node.Argument, + dtype: torch.dtype | None, + ) -> torch.fx.node.Argument: + """Promote fx_arg to dtype if necessary.""" + if dtype is None: + logger.info( + "Argument %s is not promoted. Not mentioned by type promotion rule.", + fx_arg, + ) + return fx_arg + + if isinstance(fx_arg, torch.fx.Node): + arg_val = self.env[fx_arg] + if isinstance(arg_val, torch.Tensor): + if (old_dtype := arg_val.dtype) != dtype: + # Promote tensor to dtype. + graph = node.graph + with graph.inserting_before(node): + logger.info( + "Argument %s(%s) is promoted to %s.", + fx_arg, + old_dtype, + dtype, + ) + return self._create_node( + graph, + "call_function", + torch.ops.prims.convert_element_type.default, + (fx_arg,), + {"dtype": dtype}, + ) + logger.info("Argument %s is not promoted. Already %s.", fx_arg, dtype) + return fx_arg + elif fx_type_utils.is_torch_symbolic_type(arg_val): + arg_type = type(arg_val) + equivalent_dtype = fx_type_utils.from_scalar_type_to_torch_dtype( + arg_type + ) + if equivalent_dtype is None: + raise AssertionError(f"Unexpected arg_type: {arg_type}") + if equivalent_dtype != dtype: + # Promote Sym number to tensor of dtype. + graph = node.graph + with graph.inserting_before(node): + logger.info( + "Argument %s(Scalar of equivalent dtype: %s) " + "is promoted to %s.", + fx_arg, + equivalent_dtype, + dtype, + ) + return self._create_node( + graph, + "call_function", + torch.ops.aten.scalar_tensor.default, + (fx_arg,), + {"dtype": dtype}, + ) + logger.info("Argument %s is not promoted. Already %s.", fx_arg, dtype) + return fx_arg + elif ( + equivalent_dtype := fx_type_utils.from_scalar_type_to_torch_dtype( + type(fx_arg) + ) + ) is not None: + if equivalent_dtype != dtype: + # Promote number to tensor of dtype. + # The op should have overload that supports tensor for this arg, otherwise + # the type promotion rule should not suggest promoting this arg. + graph = node.graph + with graph.inserting_before(node): + logger.info( + "Argument %s(Scalar of equivalent dtype: %s) " + "is promoted to %s.", + fx_arg, + equivalent_dtype, + dtype, + ) + return self._create_node( + graph, + "call_function", + torch.ops.aten.scalar_tensor.default, + (fx_arg,), + {"dtype": dtype}, + ) + logger.info("Argument %s is not promoted. Already %s.", fx_arg, dtype) + return fx_arg + elif isinstance(fx_arg, (tuple, list)): + logger.info("Argument %s is a tuple/list. Promoting each element.", fx_arg) + return type(fx_arg)( + self._maybe_promote_arg(node, fx_arg_elem, dtype) + for fx_arg_elem in fx_arg + ) + + raise NotImplementedError(f"Unknown fx arg type: {type(fx_arg)}") + + def _maybe_promote_node( + self, + node: torch.fx.Node, + rule: TypePromotionRule, + ) -> torch.fx.Node: + """Promote node inputs and outputs according to type promotion rule.""" + args, kwargs = self.fetch_args_kwargs_from_env(node) + type_promotion_info = rule.preview_type_promotion(args, kwargs) + new_args = [] + new_kwargs = {} + for i, arg in enumerate(node.args): + new_args.append( + self._maybe_promote_arg( + node, arg, type_promotion_info.args_dtypes.get(i, None) + ) + ) + + for name, arg in node.kwargs.items(): + new_kwargs[name] = self._maybe_promote_arg( + node, arg, type_promotion_info.kwargs_dtypes.get(name, None) + ) + new_args = tuple(new_args) + + if node.args != new_args or node.kwargs != new_kwargs: + node.args = new_args + node.kwargs = new_kwargs + self._rerun_node_after_type_promotion(node, type_promotion_info.out_dtype) + + return node + + def run_node(self, n: torch.fx.Node) -> Any: + """This method is an override which inserts type promotion nodes as needed. + + For each `call_function` node, an initial check is conducted to determine if a type + promotion rule is applicable. If a relevant rule exists, type casting nodes are + introduced for the corresponding arguments. The OpOverload of the node is updated + to one that accommodates the promoted types. Should the output type be different, + type casting node is inserted for this output. + + The call `super().run_node(node)` is guaranteed to be invoked for each node. + In the case of new or modified nodes, the result of `super().run_node(node)` is + used to update its `node.meta["val"]` value. + """ + with self._set_current_node(n): + if rule := get_type_promotion_rule(n, self.type_promotion_table): + self._maybe_promote_node(n, rule) + + return super().run_node(n) + + +class InsertTypePromotion(_pass.Transform): + """Explicitly insert type promotion ops to the graph. + + Underneath, the main pass is driven by `_TypePromotionInterpreter`, which is a subclass + of `torch.fx.Interpreter` to interpret the fx.Graph and perform the insertion of type + promotion operations. + + By re-running the new and modified nodes using the interpreter, we can update the + metadata, specifically the fake tensor stored under node.meta["val"], and ensure it + reflects the latest changes. + """ + + def __init__( + self, + module: torch.fx.GraphModule, + type_promotion_table: TypePromotionTable | None = None, + ) -> None: + super().__init__(module) + self.interpreter = _TypePromotionInterpreter( + module, type_promotion_table or TypePromotionTable() + ) + + def _fetch_fake_args( + self, + ) -> Sequence[ + fake_tensor.FakeTensor + | float + | int + | bool + | torch.SymInt + | torch.SymFloat + | torch.SymBool + | None + ]: + """Fetch fake args from fx graph. + + For each argument, try to fetch fake tensor from the matching placeholder node. + """ + fake_args = [] + for node in self.module.graph.nodes: + if node.op == "placeholder": + try: + # Meta value can be torch.Tensor, int, float, bool, + # torch.SymInt, torch.SymFloat, torch.SymBool. + meta_value = _val = node.meta.get("val", None) + except RuntimeError as e: + if not node.users: + # If the placeholder is not used, we can safely ignore it and put + # None as placeholder. + meta_value = None + else: + raise RuntimeError( + "Cannot fetch symbolic fake args from fx graph. " + "InsertTypePromotion pass needs to run with pre-existing fake args, " + "Otherwise the pass will produce inaccurate dynamic shape. " + ) from e + + fake_args.append(meta_value) + return fake_args + + def _run(self, *args, **kwargs) -> torch.fx.GraphModule: + if args: + raise AssertionError( + "`InsertTypePromotion` deduces symbolic fake arguments from the graph. " + "It does not accept concrete arguments as input because this pass requires " + "re-running the graph. When executed with newly faked concrete arguments, " + "the pass loses the symbolic dynamic shape information." + ) + if kwargs: + raise AssertionError("`kwargs` is not supported") + + fake_args = self._fetch_fake_args() + fake_mode = self.fake_mode + if fake_mode is None: + raise AssertionError("Cannot detect fake_mode.") + + # Use the python dispatcher to run through some python kernels which + # can better handle symints. Without this, some SymInts can become static + # when there are dynamic shapes. + dispatcher_mode = torch._dispatch.python.enable_python_dispatcher() + with fake_mode, dispatcher_mode, fx_traceback.preserve_node_meta(): + self.interpreter.run(*fake_args) + + return self.module diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/type_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/type_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ca434fa213f79e3b8e4c7640c8a00ebed74d1c16 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/fx/type_utils.py @@ -0,0 +1,38 @@ +# mypy: allow-untyped-defs +"""Utilities for converting and operating on ONNX and torch types.""" + +from __future__ import annotations + +from typing import Any +from typing_extensions import TypeIs + +import torch + + +def is_torch_symbolic_type( + value: Any, +) -> TypeIs[torch.SymBool | torch.SymInt | torch.SymFloat]: + return isinstance(value, (torch.SymBool, torch.SymInt, torch.SymFloat)) + + +def from_scalar_type_to_torch_dtype(scalar_type: type) -> torch.dtype | None: + return _SCALAR_TYPE_TO_TORCH_DTYPE.get(scalar_type) + + +_PYTHON_TYPE_TO_TORCH_DTYPE = { + bool: torch.bool, + int: torch.int64, + float: torch.float32, + complex: torch.complex64, +} + +_SYM_TYPE_TO_TORCH_DTYPE = { + torch.SymInt: torch.int64, + torch.SymFloat: torch.float32, + torch.SymBool: torch.bool, +} + +_SCALAR_TYPE_TO_TORCH_DTYPE: dict[type, torch.dtype] = { + **_PYTHON_TYPE_TO_TORCH_DTYPE, + **_SYM_TYPE_TO_TORCH_DTYPE, # type: ignore[dict-item] +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/_experimental.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/_experimental.py new file mode 100644 index 0000000000000000000000000000000000000000..6f1fee027331a3a215446726b9050ff5f73892b4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/_experimental.py @@ -0,0 +1,27 @@ +"""Experimental classes and functions used by ONNX export.""" + +import dataclasses +from collections.abc import Mapping, Sequence + +import torch +import torch._C._onnx as _C_onnx + + +@dataclasses.dataclass +class ExportOptions: + """Arguments used by :func:`torch.onnx.export`.""" + + # TODO(justinchuby): Deprecate and remove this class. + + export_params: bool = True + verbose: bool = False + training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL + input_names: Sequence[str] | None = None + output_names: Sequence[str] | None = None + operator_export_type: _C_onnx.OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX + opset_version: int | None = None + do_constant_folding: bool = True + dynamic_axes: Mapping[str, Mapping[int, str] | Sequence[int]] | None = None + keep_initializers_as_inputs: bool | None = None + custom_opsets: Mapping[str, int] | None = None + export_modules_as_functions: bool | set[type[torch.nn.Module]] = False diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/_globals.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/_globals.py new file mode 100644 index 0000000000000000000000000000000000000000..9e27c1dbeb8ad11f36dc8f36e0474800b4ad2895 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/_globals.py @@ -0,0 +1,81 @@ +"""Globals used internally by the ONNX exporter. + +Do not use this module outside of `torch.onnx` and its tests. + +Be very judicious when adding any new global variables. Do not create new global +variables unless they are absolutely necessary. +""" + +import torch._C._onnx as _C_onnx + +# This module should only depend on _constants and nothing else in torch.onnx to keep +# dependency direction clean. +from torch.onnx import _constants + + +class _InternalGlobals: + """Globals used internally by ONNX exporter. + + NOTE: Be very judicious when adding any new variables. Do not create new + global variables unless they are absolutely necessary. + """ + + def __init__(self) -> None: + self._export_onnx_opset_version = _constants.ONNX_DEFAULT_OPSET + self._training_mode: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL + self._in_onnx_export: bool = False + # Whether the user's model is training during export + self.export_training: bool = False + self.operator_export_type: _C_onnx.OperatorExportTypes = ( + _C_onnx.OperatorExportTypes.ONNX + ) + self.onnx_shape_inference: bool = True + self._autograd_inlining: bool = True + + @property + def training_mode(self) -> _C_onnx.TrainingMode: + """The training mode for the exporter.""" + return self._training_mode + + @training_mode.setter + def training_mode(self, training_mode: _C_onnx.TrainingMode) -> None: + if not isinstance(training_mode, _C_onnx.TrainingMode): + raise TypeError( + "training_mode must be of type 'torch.onnx.TrainingMode'. This is " + "likely a bug in torch.onnx." + ) + self._training_mode = training_mode + + @property + def export_onnx_opset_version(self) -> int: + """Opset version used during export.""" + return self._export_onnx_opset_version + + @export_onnx_opset_version.setter + def export_onnx_opset_version(self, value: int) -> None: + self._export_onnx_opset_version = value + + @property + def in_onnx_export(self) -> bool: + """Whether it is in the middle of ONNX export.""" + return self._in_onnx_export + + @in_onnx_export.setter + def in_onnx_export(self, value: bool) -> None: + if type(value) is not bool: + raise TypeError("in_onnx_export must be a boolean") + self._in_onnx_export = value + + @property + def autograd_inlining(self) -> bool: + """Whether Autograd must be inlined.""" + return self._autograd_inlining + + @autograd_inlining.setter + def autograd_inlining(self, value: bool) -> None: + if type(value) is not bool: + raise TypeError("autograd_inlining must be a boolean") + self._autograd_inlining = value + + +GLOBALS = _InternalGlobals() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/_type_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/_type_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..c0ffa56a38e3b59efe3dd2667f77295ddf529593 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/_type_utils.py @@ -0,0 +1,393 @@ +# mypy: allow-untyped-defs +"""Utilities for converting and operating on ONNX, JIT and torch types.""" + +from __future__ import annotations + +import enum +import typing +from typing import Literal + +import torch +from torch._C import _onnx as _C_onnx +from torch.onnx import errors + + +if typing.TYPE_CHECKING: + # Hack to help mypy to recognize torch._C.Value + from torch import _C # noqa: F401 + +ScalarName = Literal[ + "Byte", + "Char", + "Double", + "Float", + "Half", + "Int", + "Long", + "Short", + "Bool", + "ComplexHalf", + "ComplexFloat", + "ComplexDouble", + "QInt8", + "QUInt8", + "QInt32", + "BFloat16", + "Float8E5M2", + "Float8E4M3FN", + "Float8E5M2FNUZ", + "Float8E4M3FNUZ", + "Undefined", +] + +TorchName = Literal[ + "bool", + "uint8_t", + "int8_t", + "double", + "float", + "half", + "int", + "int64_t", + "int16_t", + "complex32", + "complex64", + "complex128", + "qint8", + "quint8", + "qint32", + "bfloat16", + "float8_e5m2", + "float8_e4m3fn", + "float8_e5m2fnuz", + "float8_e4m3fnuz", +] + + +class JitScalarType(enum.IntEnum): + """Scalar types defined in torch. + + Use ``JitScalarType`` to convert from torch and JIT scalar types to ONNX scalar types. + + Examples: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_ONNX) + >>> # xdoctest: +IGNORE_WANT("win32 has different output") + >>> JitScalarType.from_value(torch.ones(1, 2)).onnx_type() + TensorProtoDataType.FLOAT + + >>> JitScalarType.from_value(torch_c_value_with_type_float).onnx_type() + TensorProtoDataType.FLOAT + + >>> JitScalarType.from_dtype(torch.get_default_dtype).onnx_type() + TensorProtoDataType.FLOAT + + """ + + # Order defined in https://github.com/pytorch/pytorch/blob/344defc9733a45fee8d0c4d3f5530f631e823196/c10/core/ScalarType.h + UINT8 = 0 + INT8 = enum.auto() # 1 + INT16 = enum.auto() # 2 + INT = enum.auto() # 3 + INT64 = enum.auto() # 4 + HALF = enum.auto() # 5 + FLOAT = enum.auto() # 6 + DOUBLE = enum.auto() # 7 + COMPLEX32 = enum.auto() # 8 + COMPLEX64 = enum.auto() # 9 + COMPLEX128 = enum.auto() # 10 + BOOL = enum.auto() # 11 + QINT8 = enum.auto() # 12 + QUINT8 = enum.auto() # 13 + QINT32 = enum.auto() # 14 + BFLOAT16 = enum.auto() # 15 + FLOAT8E5M2 = enum.auto() # 16 + FLOAT8E4M3FN = enum.auto() # 17 + FLOAT8E5M2FNUZ = enum.auto() # 18 + FLOAT8E4M3FNUZ = enum.auto() # 19 + UNDEFINED = enum.auto() # 20 + + @classmethod + def _from_name(cls, name: ScalarName | TorchName | str | None) -> JitScalarType: + """Convert a JIT scalar type or torch type name to ScalarType. + + Note: DO NOT USE this API when `name` comes from a `torch._C.Value.type()` calls. + A "RuntimeError: INTERNAL ASSERT FAILED at "../aten/src/ATen/core/jit_type_base.h" can + be raised in several scenarios where shape info is not present. + Instead use `from_value` API which is safer. + + Args: + name: JIT scalar type name (Byte) or torch type name (uint8_t). + + Returns: + JitScalarType + + Raises: + OnnxExporterError: if name is not a valid scalar type name or if it is None. + """ + if name is None: + raise errors.OnnxExporterError("Scalar type name cannot be None") + if valid_scalar_name(name): + return _SCALAR_NAME_TO_TYPE[name] # type: ignore[index] + if valid_torch_name(name): + return _TORCH_NAME_TO_SCALAR_TYPE[name] # type: ignore[index] + + raise errors.OnnxExporterError(f"Unknown torch or scalar type: '{name}'") + + @classmethod + def from_dtype(cls, dtype: torch.dtype | None) -> JitScalarType: + """Convert a torch dtype to JitScalarType. + + Note: DO NOT USE this API when `dtype` comes from a `torch._C.Value.type()` calls. + A "RuntimeError: INTERNAL ASSERT FAILED at "../aten/src/ATen/core/jit_type_base.h" can + be raised in several scenarios where shape info is not present. + Instead use `from_value` API which is safer. + + Args: + dtype: A torch.dtype to create a JitScalarType from + + Returns: + JitScalarType + + Raises: + OnnxExporterError: if dtype is not a valid torch.dtype or if it is None. + """ + if dtype not in _DTYPE_TO_SCALAR_TYPE: + raise errors.OnnxExporterError(f"Unknown dtype: {dtype}") + # pyrefly: ignore [bad-index] + return _DTYPE_TO_SCALAR_TYPE[dtype] + + @classmethod + def from_onnx_type( + cls, onnx_type: int | _C_onnx.TensorProtoDataType | None + ) -> JitScalarType: + """Convert a ONNX data type to JitScalarType. + + Args: + onnx_type: A torch._C._onnx.TensorProtoDataType to create a JitScalarType from + + Returns: + JitScalarType + + Raises: + OnnxExporterError: if dtype is not a valid torch.dtype or if it is None. + """ + if onnx_type not in _ONNX_TO_SCALAR_TYPE: + raise errors.OnnxExporterError(f"Unknown onnx_type: {onnx_type}") + # pyrefly: ignore [redundant-cast] + return _ONNX_TO_SCALAR_TYPE[typing.cast(_C_onnx.TensorProtoDataType, onnx_type)] + + @classmethod + def from_value( + cls, value: torch._C.Value | torch.Tensor | None, default=None + ) -> JitScalarType: + """Create a JitScalarType from an value's scalar type. + + Args: + value: An object to fetch scalar type from. + default: The JitScalarType to return if a valid scalar cannot be fetched from value + + Returns: + JitScalarType. + + Raises: + OnnxExporterError: if value does not have a valid scalar type and default is None. + SymbolicValueError: when value.type()'s info are empty and default is None + """ + + if not isinstance(value, (torch._C.Value, torch.Tensor)) or ( + isinstance(value, torch._C.Value) and value.node().mustBeNone() + ): + # default value of type JitScalarType is returned when value is not valid + if default is None: + raise errors.OnnxExporterError( + "value must be either torch._C.Value or torch.Tensor objects." + ) + elif not isinstance(default, JitScalarType): + raise errors.OnnxExporterError( + "default value must be a JitScalarType object." + ) + return default + + # Each value type has their own way of storing scalar type + if isinstance(value, torch.Tensor): + return cls.from_dtype(value.dtype) + if isinstance(value.type(), torch.ListType): + try: + return cls.from_dtype(value.type().getElementType().dtype()) + except RuntimeError: + return cls._from_name(str(value.type().getElementType())) + if isinstance(value.type(), torch._C.OptionalType): + if value.type().getElementType().dtype() is None: + if isinstance(default, JitScalarType): + return default + raise errors.OnnxExporterError( + "default value must be a JitScalarType object." + ) + return cls.from_dtype(value.type().getElementType().dtype()) + + scalar_type = None + if value.node().kind() != "prim::Constant" or not isinstance( + value.type(), torch._C.NoneType + ): + # value must be a non-list torch._C.Value scalar + scalar_type = value.type().scalarType() + + if scalar_type is not None: + return cls._from_name(scalar_type) + + # When everything fails... try to default + if default is not None: + return default + raise errors.SymbolicValueError( + f"Cannot determine scalar type for this '{type(value.type())}' instance and " + "a default value was not provided.", + value, + ) + + def scalar_name(self) -> ScalarName: + """Convert a JitScalarType to a JIT scalar type name.""" + return _SCALAR_TYPE_TO_NAME[self] + + def torch_name(self) -> TorchName: + """Convert a JitScalarType to a torch type name.""" + return _SCALAR_TYPE_TO_TORCH_NAME[self] + + def dtype(self) -> torch.dtype: + """Convert a JitScalarType to a torch dtype.""" + return _SCALAR_TYPE_TO_DTYPE[self] + + def onnx_type(self) -> _C_onnx.TensorProtoDataType: + """Convert a JitScalarType to an ONNX data type.""" + if self not in _SCALAR_TYPE_TO_ONNX: + raise errors.OnnxExporterError( + f"Scalar type {self} cannot be converted to ONNX" + ) + return _SCALAR_TYPE_TO_ONNX[self] + + def onnx_compatible(self) -> bool: + """Return whether this JitScalarType is compatible with ONNX.""" + return ( + self in _SCALAR_TYPE_TO_ONNX + and self != JitScalarType.UNDEFINED + and self != JitScalarType.COMPLEX32 + ) + + +def valid_scalar_name(scalar_name: ScalarName | str) -> bool: + """Return whether the given scalar name is a valid JIT scalar type name.""" + return scalar_name in _SCALAR_NAME_TO_TYPE + + +def valid_torch_name(torch_name: TorchName | str) -> bool: + """Return whether the given torch name is a valid torch type name.""" + return torch_name in _TORCH_NAME_TO_SCALAR_TYPE + + +# https://github.com/pytorch/pytorch/blob/344defc9733a45fee8d0c4d3f5530f631e823196/c10/core/ScalarType.h +_SCALAR_TYPE_TO_NAME: dict[JitScalarType, ScalarName] = { + JitScalarType.BOOL: "Bool", + JitScalarType.UINT8: "Byte", + JitScalarType.INT8: "Char", + JitScalarType.INT16: "Short", + JitScalarType.INT: "Int", + JitScalarType.INT64: "Long", + JitScalarType.HALF: "Half", + JitScalarType.FLOAT: "Float", + JitScalarType.DOUBLE: "Double", + JitScalarType.COMPLEX32: "ComplexHalf", + JitScalarType.COMPLEX64: "ComplexFloat", + JitScalarType.COMPLEX128: "ComplexDouble", + JitScalarType.QINT8: "QInt8", + JitScalarType.QUINT8: "QUInt8", + JitScalarType.QINT32: "QInt32", + JitScalarType.BFLOAT16: "BFloat16", + JitScalarType.FLOAT8E5M2: "Float8E5M2", + JitScalarType.FLOAT8E4M3FN: "Float8E4M3FN", + JitScalarType.FLOAT8E5M2FNUZ: "Float8E5M2FNUZ", + JitScalarType.FLOAT8E4M3FNUZ: "Float8E4M3FNUZ", + JitScalarType.UNDEFINED: "Undefined", +} + +_SCALAR_NAME_TO_TYPE: dict[ScalarName, JitScalarType] = { + v: k for k, v in _SCALAR_TYPE_TO_NAME.items() +} + +_SCALAR_TYPE_TO_TORCH_NAME: dict[JitScalarType, TorchName] = { + JitScalarType.BOOL: "bool", + JitScalarType.UINT8: "uint8_t", + JitScalarType.INT8: "int8_t", + JitScalarType.INT16: "int16_t", + JitScalarType.INT: "int", + JitScalarType.INT64: "int64_t", + JitScalarType.HALF: "half", + JitScalarType.FLOAT: "float", + JitScalarType.DOUBLE: "double", + JitScalarType.COMPLEX32: "complex32", + JitScalarType.COMPLEX64: "complex64", + JitScalarType.COMPLEX128: "complex128", + JitScalarType.QINT8: "qint8", + JitScalarType.QUINT8: "quint8", + JitScalarType.QINT32: "qint32", + JitScalarType.BFLOAT16: "bfloat16", + JitScalarType.FLOAT8E5M2: "float8_e5m2", + JitScalarType.FLOAT8E4M3FN: "float8_e4m3fn", + JitScalarType.FLOAT8E5M2FNUZ: "float8_e5m2fnuz", + JitScalarType.FLOAT8E4M3FNUZ: "float8_e4m3fnuz", +} + +_TORCH_NAME_TO_SCALAR_TYPE: dict[TorchName, JitScalarType] = { + v: k for k, v in _SCALAR_TYPE_TO_TORCH_NAME.items() +} + +_SCALAR_TYPE_TO_ONNX = { + JitScalarType.BOOL: _C_onnx.TensorProtoDataType.BOOL, + JitScalarType.UINT8: _C_onnx.TensorProtoDataType.UINT8, + JitScalarType.INT8: _C_onnx.TensorProtoDataType.INT8, + JitScalarType.INT16: _C_onnx.TensorProtoDataType.INT16, + JitScalarType.INT: _C_onnx.TensorProtoDataType.INT32, + JitScalarType.INT64: _C_onnx.TensorProtoDataType.INT64, + JitScalarType.HALF: _C_onnx.TensorProtoDataType.FLOAT16, + JitScalarType.FLOAT: _C_onnx.TensorProtoDataType.FLOAT, + JitScalarType.DOUBLE: _C_onnx.TensorProtoDataType.DOUBLE, + JitScalarType.COMPLEX64: _C_onnx.TensorProtoDataType.COMPLEX64, + JitScalarType.COMPLEX128: _C_onnx.TensorProtoDataType.COMPLEX128, + JitScalarType.BFLOAT16: _C_onnx.TensorProtoDataType.BFLOAT16, + JitScalarType.UNDEFINED: _C_onnx.TensorProtoDataType.UNDEFINED, + JitScalarType.COMPLEX32: _C_onnx.TensorProtoDataType.UNDEFINED, + JitScalarType.QINT8: _C_onnx.TensorProtoDataType.INT8, + JitScalarType.QUINT8: _C_onnx.TensorProtoDataType.UINT8, + JitScalarType.QINT32: _C_onnx.TensorProtoDataType.INT32, + JitScalarType.FLOAT8E5M2: _C_onnx.TensorProtoDataType.FLOAT8E5M2, + JitScalarType.FLOAT8E4M3FN: _C_onnx.TensorProtoDataType.FLOAT8E4M3FN, + JitScalarType.FLOAT8E5M2FNUZ: _C_onnx.TensorProtoDataType.FLOAT8E5M2FNUZ, + JitScalarType.FLOAT8E4M3FNUZ: _C_onnx.TensorProtoDataType.FLOAT8E4M3FNUZ, +} + +_ONNX_TO_SCALAR_TYPE = {v: k for k, v in _SCALAR_TYPE_TO_ONNX.items()} + +# source of truth is +# https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/tensor_dtypes.cpp +_SCALAR_TYPE_TO_DTYPE = { + JitScalarType.BOOL: torch.bool, + JitScalarType.UINT8: torch.uint8, + JitScalarType.INT8: torch.int8, + JitScalarType.INT16: torch.short, + JitScalarType.INT: torch.int, + JitScalarType.INT64: torch.int64, + JitScalarType.HALF: torch.half, + JitScalarType.FLOAT: torch.float, + JitScalarType.DOUBLE: torch.double, + JitScalarType.COMPLEX32: torch.complex32, + JitScalarType.COMPLEX64: torch.complex64, + JitScalarType.COMPLEX128: torch.complex128, + JitScalarType.QINT8: torch.qint8, + JitScalarType.QUINT8: torch.quint8, + JitScalarType.QINT32: torch.qint32, + JitScalarType.BFLOAT16: torch.bfloat16, + JitScalarType.FLOAT8E5M2: torch.float8_e5m2, + JitScalarType.FLOAT8E4M3FN: torch.float8_e4m3fn, + JitScalarType.FLOAT8E5M2FNUZ: torch.float8_e5m2fnuz, + JitScalarType.FLOAT8E4M3FNUZ: torch.float8_e4m3fnuz, +} + +_DTYPE_TO_SCALAR_TYPE = {v: k for k, v in _SCALAR_TYPE_TO_DTYPE.items()} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/jit_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/jit_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b15612c744c0d46edecbe7470bc47b3ed1d877d5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/jit_utils.py @@ -0,0 +1,374 @@ +# mypy: allow-untyped-defs +"""Utilities for manipulating the torch.Graph object and the torchscript.""" + +from __future__ import annotations + +import dataclasses +import re +import typing +from collections.abc import Iterable, Sequence +from typing import Any + +import torch +from torch import _C +from torch.onnx._internal.torchscript_exporter import registration +from torch.onnx._internal.torchscript_exporter._globals import GLOBALS + + +_ATTR_PATTERN = re.compile("^(.+)_(([ifstgz])|(ty))$") +_SKIP_NODE_ATTRIBUTES = {"inplace", "aten"} + + +@dataclasses.dataclass +class GraphContext: + """Extra context for symbolic functions with all methods from torch.Graph. + + NOTE: This class is not meant for external consumption. Please do not depend on + it outside of torch.onnx as the interface may evolve. + + Attributes: + graph: The _C.Graph being constructed. + block: The current _C.Block being constructed. + opset: The opset version. + original_node: Current node that is being converted from. + params_dict: Mapping from graph initializer name to IValue. + env: Mapping from Torch domain graph Value to ONNX domain graph Value. + values_in_env: Set of all values in env, for constant-time lookups. + new_nodes: List that tracks all new nodes that are added (used to make + sure metadata is propagated to all new nodes). + """ + + graph: _C.Graph + block: _C.Block + opset: int + original_node: _C.Node + params_dict: dict[str, _C.IValue] + env: dict[_C.Value, _C.Value] + values_in_env: set[_C.Value] + new_nodes: list[_C.Node] = dataclasses.field(default_factory=list) + + # Relay methods from _C.Graph for compatibility with symbolic functions that expect + # a _C.Graph + def __getattr__(self, name: str) -> Any: + return getattr(self.graph, name) + + def op( + self, + opname: str, + *raw_args: torch.Tensor | _C.Value, + outputs: int = 1, + **kwargs, + ): + """Creates an ONNX operator "opname", taking "raw_args" as inputs and "kwargs" as attributes. + + The set of operators and the inputs/attributes they take + is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md + + Args: + opname: The ONNX operator name, e.g., `Abs` or `Add`, or an operator qualified + with a namespace, e.g., `aten::add`. + raw_args: The inputs to the operator; usually provided + as arguments to the `symbolic` definition. + outputs: The number of outputs this operator returns. + By default an operator is assumed to return a single output. + If `outputs` is greater than one, this functions returns a tuple + of output `Value`, representing each output of the ONNX operator + in order. + kwargs: The attributes of the ONNX operator, whose keys are named + according to the following convention: `alpha_f` indicates + the `alpha` attribute with type `f`. The valid type specifiers are + `f` (float), `i` (int), `s` (string) or `t` (Tensor). An attribute + specified with type float accepts either a single float, or a + list of floats (e.g., you would say `dims_i` for a `dims` attribute + that takes a list of integers). + + Returns: + The value representing the single output of this operator (see the `outputs` + keyword argument for multi-return nodes). + """ + return _add_op(self, opname, *raw_args, outputs=outputs, **kwargs) + + def aten_op(self, operator: str, *args, overload_name: str = "", **kwargs): + """Generates an ONNX ATen op node. + + This function is for backward compatibility with the old symbolic functions. + """ + return self.op( + "aten::ATen", + *args, + operator_s=operator, + overload_name_s=overload_name, + **kwargs, + ) + + # NOTE: For backward compatibility with the old symbolic functions. + # We are probably going to remove this only after the fx exporter is established. + at = aten_op + + def onnxscript_op( + self, + onnx_fn, + *raw_args: torch.Tensor | _C.Value, + outputs: int = 1, + **kwargs, + ): + """Creates an ONNX operator from onnx-script function, taking "raw_args" as inputs and "kwargs" as attributes. + + onnx-script repository: https://github.com/microsoft/onnx-script + + Args: + onnx_fn: ONNXFunction from onnx-script; An example can be found at + https://github.com/microsoft/onnx-script#example + raw_args: The inputs to the operator; usually provided + as arguments to the `symbolic` definition. + outputs: The number of outputs this operator returns. + By default an operator is assumed to return a single output. + If `outputs` is greater than one, this functions returns a tuple + of output `Value`, representing each output of the ONNX operator + in order. + kwargs: The attributes of the ONNX operator, whose keys are named + according to the following convention: `alpha_f` indicates + the `alpha` attribute with type `f`. The valid type specifiers are + `f` (float), `i` (int), `s` (string) or `t` (Tensor). An attribute + specified with type float accepts either a single float, or a + list of floats (e.g., you would say `dims_i` for a `dims` attribute + that takes a list of integers). + + Returns: + The value representing the single output of this operator (see the `outputs` + keyword argument for multi-return nodes). + """ + # NOTE(titaiwang): This is using class attributes, and it needs to be updated + # if onnx-script makes any change on these. + symbolic_name = f"{onnx_fn.opset.domain}::{onnx_fn.name}" + opset_version = onnx_fn.opset.version + + registration.custom_onnx_symbolic(symbolic_name, opset_version)(onnx_fn) + + return _add_op(self, symbolic_name, *raw_args, outputs=outputs, **kwargs) + + +def add_op_with_blocks( + graph_context: GraphContext, + opname: str, + *inputs: _C.Value, + outputs: int = 1, + n_blocks: int = 1, + **attributes, +) -> tuple[Any, tuple[GraphContext, ...], _C.Node]: + """Creates an ONNX operator "opname", taking inputs and attributes. + + Args: + graph_context: The context for the current graph. + opname: The ONNX operator name, e.g., `Abs` or `Add`, or an operator qualified + with a namespace, e.g., `aten::add`. + inputs: The inputs to the operator. + outputs: The number of outputs this operator returns. + By default an operator is assumed to return a single output. + If `outputs` is greater than one, this functions returns a tuple + of output `Value`, representing each output of the ONNX operator + in order. + n_blocks: The number of sub-blocks to create in the node. + attributes: The attributes of the ONNX operator. + + Returns: + A tuple of (output_values, new_contexts, node) where: + output_values: One or more output value of this operator + (see the `outputs` keyword argument for multi-return nodes). + new_contexts: A tuple of new graph contexts for each sub-block. + node: The node representing the operator. + """ + + output_values = graph_context.op(opname, *inputs, outputs=outputs, **attributes) + if isinstance(output_values, Sequence): + node = output_values[0].node() + else: + node = output_values.node() + + new_contexts = [] + for _ in range(n_blocks): + new_block = node.addBlock() + # Create shallow copy of the graph context and update the block + new_context = dataclasses.replace(graph_context, block=new_block) + new_contexts.append(new_context) + + return output_values, tuple(new_contexts), node + + +def _add_op( + graph_context: GraphContext, + opname: str, + *args: torch.Tensor | _C.Value, + outputs: int = 1, + **kwargs, +): + """Creates an ONNX operator "opname", taking "args" as inputs and attributes "kwargs". + + The set of operators and the inputs/attributes they take + is documented at https://github.com/onnx/onnx/blob/master/docs/Operators.md + + Args: + graph_context: The Torch Graph or Block. + opname: The ONNX operator name, e.g., `Abs` or `Add`, or an operator qualified + with a namespace, e.g., `aten::add`. + args: The inputs to the operator; usually provided + as arguments to the `symbolic` definition. + outputs: The number of outputs this operator returns. + By default an operator is assumed to return a single output. + If `outputs` is greater than one, this functions returns a tuple + of output `Value`, representing each output of the ONNX operator + in order. + kwargs: The attributes of the ONNX operator, whose keys are named + according to the following convention: `alpha_f` indicates + the `alpha` attribute with type `f`. The valid type specifiers are + `f` (float), `i` (int), `s` (string) or `t` (Tensor). An attribute + specified with type float accepts either a single float, or a + list of floats (e.g., you would say `dims_i` for a `dims` attribute + that takes a list of integers). + + Returns: + (Union[_C.Value, Tuple[_C.Value, ...]]) + The value representing the single output of this operator (see the `outputs` + keyword argument for multi-return nodes). + """ + inputs = [_const_if_tensor(graph_context, arg) for arg in args] + # Filter out None attributes, this can be convenient client side because + # now they can pass through None attributes, and have them not show up + attributes = {k: v for k, v in kwargs.items() if v is not None} + + if "::" not in opname: + opname = "onnx::" + opname + + node = _create_node( + graph_context.block, + opname, + inputs, + attributes, + params_dict=graph_context.params_dict, + opset_version=graph_context.opset, + n_outputs=outputs, + shape_inference=GLOBALS.onnx_shape_inference, + ) + graph_context.new_nodes.append(node) + + if outputs == 1: + return node.output() + return tuple(node.outputs()) + + +def _const_if_tensor(graph_context: GraphContext, arg): + if arg is None: + return arg + if isinstance(arg, _C.Value): + return arg + + return _add_op(graph_context, "onnx::Constant", value_z=arg) + + +def _create_node( + graph_or_block: _C.Graph | _C.Block, + domain_op: str, + inputs: Sequence, + attributes: dict, + params_dict: dict, + opset_version: int, + n_outputs: int, + shape_inference: bool = True, +) -> _C.Node: + """Creates an node 'domain_op', taking inputs and attributes.""" + if isinstance(graph_or_block, _C.Graph): + graph = graph_or_block + node = graph.create(domain_op, inputs, n_outputs) + node = graph.insertNode(node) + elif isinstance(graph_or_block, _C.Block): + block = graph_or_block + node = block.addNode(domain_op, inputs) + + # Block does not have create defined, so we need to add outputs manually + if n_outputs > 1: + for _ in range(1, n_outputs): + node.addOutput() + + node_outputs = tuple(node.outputs()) # type: ignore[possibly-undefined] + if len(node_outputs) != n_outputs: + raise AssertionError( + f"len(node_outputs)={len(node_outputs)} != n_outputs={n_outputs}" + ) + + aten = domain_op.startswith("aten::") + + # Add all attributes + for key, value in sorted(attributes.items()): + if key in _SKIP_NODE_ATTRIBUTES: + continue + # pyrefly: ignore [unbound-name] + _add_attribute(node, key, value, aten=aten) + if shape_inference: + # pyrefly: ignore [unbound-name] + _C._jit_pass_onnx_node_shape_type_inference(node, params_dict, opset_version) + # pyrefly: ignore [unbound-name] + return node + + +def _is_onnx_list(value): + return isinstance(value, Iterable) and not isinstance( + value, (str, bytes, torch.Tensor) + ) + + +def _scalar(x: torch.Tensor): + """Convert a scalar tensor into a Python value.""" + if x.numel() != 1: + raise AssertionError(f"Expected numel() == 1, got {x.numel()}") + return x[0] + + +def _add_attribute(node: _C.Node, key: str, value: Any, aten: bool): + r"""Initializes the right attribute based on type of value.""" + m = _ATTR_PATTERN.match(key) + if m is None: + raise ValueError( + f"Invalid attribute specifier '{key}' names " + "must be suffixed with type, e.g. 'dim_i' or 'dims_i'" + ) + name, kind = m.group(1), m.group(2) + if _is_onnx_list(value): + kind += "s" + + return getattr(node, f"{kind}_")(name, value) + + +def _is_tensor(x: _C.Value) -> bool: + return x.type().isSubtypeOf(_C.TensorType.get()) + + +def get_device_from_value(value: _C.Value) -> torch.device | None: + if not _is_tensor(value): + return None + tensor_type = typing.cast(_C.TensorType, value.type()) + return tensor_type.device() + + +def parse_node_kind(kind: str) -> tuple[str, str]: + """Parse node kind into domain and Op name.""" + if "::" not in kind: + raise ValueError(f"Node kind: {kind} is invalid. '::' is not in node kind.") + domain, opname = kind.split("::", 1) + if "::" in opname: + raise ValueError(f"Node kind: {kind} is invalid. '::' should only appear once.") + return domain, opname + + +def is_aten(domain: str) -> bool: + """Check if the domain is official.""" + return domain == "aten" + + +def is_prim(domain: str) -> bool: + """Check if the domain is official.""" + return domain == "prim" + + +def is_onnx(domain: str) -> bool: + """Check if the domain is official.""" + return domain == "onnx" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/onnx_proto_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/onnx_proto_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..902963311b820bc95042fcb48a468986cf3ed341 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/onnx_proto_utils.py @@ -0,0 +1,251 @@ +# mypy: allow-untyped-defs +"""Utilities for manipulating the onnx and onnx-script dependencies and ONNX proto.""" + +from __future__ import annotations + +import glob +import os +import shutil +from typing import Any, TYPE_CHECKING + +import torch +import torch.serialization +from torch.onnx import errors +from torch.onnx._internal.torchscript_exporter import jit_utils, registration + + +if TYPE_CHECKING: + import io + from collections.abc import Mapping + + +def export_as_test_case( + model_bytes: bytes, inputs_data, outputs_data, name: str, dir: str +) -> str: + """Export an ONNX model as a self contained ONNX test case. + + The test case contains the model and the inputs/outputs data. The directory structure + is as follows: + + dir + \u251c\u2500\u2500 test_ + \u2502 \u251c\u2500\u2500 model.onnx + \u2502 \u2514\u2500\u2500 test_data_set_0 + \u2502 \u251c\u2500\u2500 input_0.pb + \u2502 \u251c\u2500\u2500 input_1.pb + \u2502 \u251c\u2500\u2500 output_0.pb + \u2502 \u2514\u2500\u2500 output_1.pb + + Args: + model_bytes: The ONNX model in bytes. + inputs_data: The inputs data, nested data structure of numpy.ndarray. + outputs_data: The outputs data, nested data structure of numpy.ndarray. + + Returns: + The path to the test case directory. + """ + try: + import onnx + except ImportError as exc: + raise ImportError( + "Export test case to ONNX format failed: Please install ONNX." + ) from exc + + test_case_dir = os.path.join(dir, "test_" + name) + os.makedirs(test_case_dir, exist_ok=True) + _export_file( + model_bytes, + os.path.join(test_case_dir, "model.onnx"), + {}, + ) + data_set_dir = os.path.join(test_case_dir, "test_data_set_0") + if os.path.exists(data_set_dir): + shutil.rmtree(data_set_dir) + os.makedirs(data_set_dir) + + proto = onnx.load_model_from_string(model_bytes) # type: ignore[attr-defined] + + for i, (input_proto, input) in enumerate(zip(proto.graph.input, inputs_data)): + export_data(input, input_proto, os.path.join(data_set_dir, f"input_{i}.pb")) + for i, (output_proto, output) in enumerate(zip(proto.graph.output, outputs_data)): + export_data(output, output_proto, os.path.join(data_set_dir, f"output_{i}.pb")) + + return test_case_dir + + +def load_test_case(dir: str) -> tuple[bytes, Any, Any]: + """Load a self contained ONNX test case from a directory. + + The test case must contain the model and the inputs/outputs data. The directory structure + should be as follows: + + dir + \u251c\u2500\u2500 test_ + \u2502 \u251c\u2500\u2500 model.onnx + \u2502 \u2514\u2500\u2500 test_data_set_0 + \u2502 \u251c\u2500\u2500 input_0.pb + \u2502 \u251c\u2500\u2500 input_1.pb + \u2502 \u251c\u2500\u2500 output_0.pb + \u2502 \u2514\u2500\u2500 output_1.pb + + Args: + dir: The directory containing the test case. + + Returns: + model_bytes: The ONNX model in bytes. + inputs: the inputs data, mapping from input name to numpy.ndarray. + outputs: the outputs data, mapping from output name to numpy.ndarray. + """ + try: + import onnx + from onnx import numpy_helper # type: ignore[attr-defined] + except ImportError as exc: + raise ImportError( + "Load test case from ONNX format failed: Please install ONNX." + ) from exc + + with open(os.path.join(dir, "model.onnx"), "rb") as f: + model_bytes = f.read() + + test_data_dir = os.path.join(dir, "test_data_set_0") + + inputs = {} + input_files = glob.glob(os.path.join(test_data_dir, "input_*.pb")) + for input_file in input_files: + tensor = onnx.load_tensor(input_file) # type: ignore[attr-defined] + inputs[tensor.name] = numpy_helper.to_array(tensor) + outputs = {} + output_files = glob.glob(os.path.join(test_data_dir, "output_*.pb")) + for output_file in output_files: + tensor = onnx.load_tensor(output_file) # type: ignore[attr-defined] + outputs[tensor.name] = numpy_helper.to_array(tensor) + + return model_bytes, inputs, outputs + + +def export_data(data, value_info_proto, f: str) -> None: + """Export data to ONNX protobuf format. + + Args: + data: The data to export, nested data structure of numpy.ndarray. + value_info_proto: The ValueInfoProto of the data. The type of the ValueInfoProto + determines how the data is stored. + f: The file to write the data to. + """ + try: + from onnx import numpy_helper # type: ignore[attr-defined] + except ImportError as exc: + raise ImportError( + "Export data to ONNX format failed: Please install ONNX." + ) from exc + + with open(f, "wb") as opened_file: + if value_info_proto.type.HasField("map_type"): + opened_file.write( + numpy_helper.from_dict(data, value_info_proto.name).SerializeToString() + ) + elif value_info_proto.type.HasField("sequence_type"): + opened_file.write( + numpy_helper.from_list(data, value_info_proto.name).SerializeToString() + ) + elif value_info_proto.type.HasField("optional_type"): + opened_file.write( + numpy_helper.from_optional( + data, value_info_proto.name + ).SerializeToString() + ) + else: + if not value_info_proto.type.HasField("tensor_type"): + raise AssertionError("Expected tensor_type field to be set") + opened_file.write( + numpy_helper.from_array(data, value_info_proto.name).SerializeToString() + ) + + +def _export_file( + model_bytes: bytes, + f: io.BytesIO | str, + export_map: Mapping[str, bytes], +) -> None: + """export/write model bytes into directory/protobuf/zip""" + if len(export_map) != 0: + raise AssertionError(f"export_map must be empty, got {len(export_map)} items") + with torch.serialization._open_file_like(f, "wb") as opened_file: + opened_file.write(model_bytes) + + +def _add_onnxscript_fn( + model_bytes: bytes, + custom_opsets: Mapping[str, int], +) -> bytes: + """Insert model-included custom onnx-script function into ModelProto""" + try: + import onnx + except ImportError as e: + raise errors.OnnxExporterError("Module onnx is not installed!") from e + + # For > 2GB model, onnx.load_fromstring would fail. However, because + # in _export_onnx, the tensors should be saved separately if the proto + # size > 2GB, and if it for some reason did not, the model would fail on + # serialization anyway in terms of the protobuf limitation. So we don't + # need to worry about > 2GB model getting here. + model_proto = onnx.load_model_from_string(model_bytes) # type: ignore[attr-defined] + + # Iterate graph nodes to insert only the included custom + # function_proto into model_proto + onnx_function_list = [] # type: ignore[var-annotated] + included_node_func: set[str] = set() + # onnx_function_list and included_node_func are expanded in-place + _find_onnxscript_op( + model_proto.graph, included_node_func, custom_opsets, onnx_function_list + ) + + if onnx_function_list: + model_proto.functions.extend(onnx_function_list) + model_bytes = model_proto.SerializeToString() + return model_bytes + + +def _find_onnxscript_op( + graph_proto, + included_node_func: set[str], + custom_opsets: Mapping[str, int], + onnx_function_list: list, +): + """Recursively iterate ModelProto to find ONNXFunction op as it may contain control flow Op.""" + for node in graph_proto.node: + node_kind = node.domain + "::" + node.op_type + # Recursive needed for control flow nodes: IF/Loop which has inner graph_proto + for attr in node.attribute: + if attr.g is not None: + _find_onnxscript_op( + attr.g, included_node_func, custom_opsets, onnx_function_list + ) + # Only custom Op with ONNX function and aten with symbolic_fn should be found in registry + onnx_function_group = registration.registry.get_function_group(node_kind) + # Ruled out corner cases: onnx/prim in registry + if ( + node.domain + and not jit_utils.is_aten(node.domain) + and not jit_utils.is_prim(node.domain) + and not jit_utils.is_onnx(node.domain) + and onnx_function_group is not None + and node_kind not in included_node_func + ): + specified_version = custom_opsets.get(node.domain, 1) + onnx_fn = onnx_function_group.get(specified_version) + if onnx_fn is not None: + if hasattr(onnx_fn, "to_function_proto"): + onnx_function_proto = onnx_fn.to_function_proto() # type: ignore[attr-defined] + onnx_function_list.append(onnx_function_proto) + included_node_func.add(node_kind) + continue + + raise errors.UnsupportedOperatorError( + node_kind, + specified_version, + onnx_function_group.get_min_supported() + if onnx_function_group + else None, + ) + return onnx_function_list, included_node_func diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/registration.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/registration.py new file mode 100644 index 0000000000000000000000000000000000000000..3ac1fdb048ea17d0db6708e7d1e320e8283bf288 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/registration.py @@ -0,0 +1,337 @@ +# mypy: allow-untyped-defs +"""Module for handling symbolic function registration.""" + +import warnings +from collections.abc import Callable, Collection, Sequence +from typing import Generic, TypeVar +from typing_extensions import ParamSpec + +from torch.onnx import _constants, errors + + +OpsetVersion = int + + +def _dispatch_opset_version( + target: OpsetVersion, registered_opsets: Collection[OpsetVersion] +) -> OpsetVersion | None: + """Finds the registered opset given a target opset version and the available opsets. + + Args: + target: The target opset version. + registered_opsets: The available opsets. + + Returns: + The registered opset version. + """ + if not registered_opsets: + return None + + descending_registered_versions = sorted(registered_opsets, reverse=True) + # Linear search for the opset version, which is fine since the number of opset + # versions is small. + + if target >= _constants.ONNX_BASE_OPSET: + # Always look down toward opset 1 when the target is >= ONNX_BASE_OPSET (opset 9). + # When a custom op is register at opset 1, we want to be able to discover it as a + # fallback for all opsets >= ONNX_BASE_OPSET. + for version in descending_registered_versions: + if version <= target: + return version + return None + + # target < opset 9. This is the legacy behavior to support opset 7 and opset 8. + # for caffe2 support. We search up toward opset 9. + for version in reversed(descending_registered_versions): + # Count back up until _constants.ONNX_BASE_OPSET + if target <= version <= _constants.ONNX_BASE_OPSET: + return version + + return None + + +_K = TypeVar("_K") +_V = TypeVar("_V") +_R = TypeVar("_R") +_P = ParamSpec("_P") + + +class OverrideDict(Collection[_K], Generic[_K, _V]): + """A dictionary that merges built-in and custom symbolic functions. + + It supports overriding and un-overriding built-in symbolic functions with custom + ones. + """ + + def __init__(self) -> None: + self._base: dict[_K, _V] = {} + self._overrides: dict[_K, _V] = {} + self._merged: dict[_K, _V] = {} + + def set_base(self, key: _K, value: _V) -> None: + self._base[key] = value + if key not in self._overrides: + self._merged[key] = value + + def in_base(self, key: _K) -> bool: + """Checks if a key is in the base dictionary.""" + return key in self._base + + def override(self, key: _K, value: _V) -> None: + """Overrides a base key-value with a new pair.""" + self._overrides[key] = value + self._merged[key] = value + + def remove_override(self, key: _K) -> None: + """Un-overrides a key-value pair.""" + self._overrides.pop(key, None) # type: ignore[arg-type] + self._merged.pop(key, None) # type: ignore[arg-type] + if key in self._base: + self._merged[key] = self._base[key] + + def overridden(self, key: _K) -> bool: + """Checks if a key-value pair is overridden.""" + return key in self._overrides + + def __getitem__(self, key: _K) -> _V: + return self._merged[key] + + def get(self, key: _K, default: _V | None = None): + return self._merged.get(key, default) + + def __contains__(self, key: object) -> bool: + return key in self._merged + + def __iter__(self): + return iter(self._merged) + + def __len__(self) -> int: + return len(self._merged) + + def __repr__(self) -> str: + return f"OverrideDict(base={self._base}, overrides={self._overrides})" + + def __bool__(self) -> bool: + return bool(self._merged) + + +class _SymbolicFunctionGroup: + """Different versions of symbolic functions registered to the same name. + + O(number of registered versions of an op) search is performed to find the most + recent version of the op. + + The registration is delayed until op is used to improve startup time. + + Function overloads with different arguments are not allowed. + Custom op overrides are supported. + """ + + def __init__(self, name: str) -> None: + self._name = name + # A dictionary of functions, keyed by the opset version. + self._functions: OverrideDict[OpsetVersion, Callable] = OverrideDict() + + def __repr__(self) -> str: + return f"_SymbolicFunctionGroup({self._name}, registered={self._functions})" + + def __getitem__(self, key: OpsetVersion) -> Callable: + result = self.get(key) + if result is None: + raise KeyError(key) + return result + + # TODO(justinchuby): Add @functools.lru_cache(maxsize=None) if lookup time becomes + # a problem. + def get(self, opset: OpsetVersion) -> Callable | None: + """Find the most recent version of the function.""" + version = _dispatch_opset_version(opset, self._functions) + if version is None: + return None + + return self._functions[version] + + def add(self, func: Callable, opset: OpsetVersion) -> None: + """Adds a symbolic function. + + Args: + func: The function to add. + opset: The opset version of the function to add. + """ + if self._functions.in_base(opset): + warnings.warn( + f"Symbolic function '{self._name}' already registered for opset {opset}. " + f"Replacing the existing function with new function. This is unexpected. " + f"Please report it on {_constants.PYTORCH_GITHUB_ISSUES_URL}.", + errors.OnnxExporterWarning, + stacklevel=2, + ) + self._functions.set_base(opset, func) + + def add_custom(self, func: Callable, opset: OpsetVersion) -> None: + """Adds a custom symbolic function. + + Args: + func: The symbolic function to register. + opset: The corresponding opset version. + """ + self._functions.override(opset, func) + + def remove_custom(self, opset: OpsetVersion) -> None: + """Removes a custom symbolic function. + + Args: + opset: The opset version of the custom function to remove. + """ + if not self._functions.overridden(opset): + warnings.warn( + f"No custom function registered for '{self._name}' opset {opset}", + stacklevel=2, + ) + return + self._functions.remove_override(opset) + + def get_min_supported(self) -> OpsetVersion: + """Returns the lowest built-in opset version supported by the function.""" + return min(self._functions) + + +class SymbolicRegistry: + """Registry for symbolic functions. + + The registry maintains a mapping from qualified names to symbolic functions. + It is used to register new symbolic functions and to dispatch calls to + the appropriate function. + """ + + def __init__(self) -> None: + self._registry: dict[str, _SymbolicFunctionGroup] = {} + + def register( + self, name: str, opset: OpsetVersion, func: Callable, custom: bool = False + ) -> None: + """Registers a symbolic function. + + Args: + name: The qualified name of the function to register. In the form of 'domain::op'. + E.g. 'aten::add'. + opset: The opset version of the function to register. + func: The symbolic function to register. + custom: Whether the function is a custom function that overrides existing ones. + + Raises: + ValueError: If the separator '::' is not in the name. + """ + if "::" not in name: + raise ValueError( + f"The name must be in the form of 'domain::op', not '{name}'" + ) + symbolic_functions = self._registry.setdefault( + name, _SymbolicFunctionGroup(name) + ) + if custom: + symbolic_functions.add_custom(func, opset) + else: + symbolic_functions.add(func, opset) + + def unregister(self, name: str, opset: OpsetVersion) -> None: + """Unregisters a symbolic function. + + Args: + name: The qualified name of the function to unregister. + opset: The opset version of the function to unregister. + """ + if name not in self._registry: + return + self._registry[name].remove_custom(opset) + + def get_function_group(self, name: str) -> _SymbolicFunctionGroup | None: + """Returns the function group for the given name.""" + return self._registry.get(name) + + def is_registered_op(self, name: str, version: int) -> bool: + """Returns whether the given op is registered for the given opset version.""" + functions = self.get_function_group(name) + if functions is None: + return False + return functions.get(version) is not None + + def all_functions(self) -> set[str]: + """Returns the set of all registered function names.""" + return set(self._registry) + + +def onnx_symbolic( + name: str, + opset: OpsetVersion | Sequence[OpsetVersion], + decorate: Sequence[Callable] | None = None, + custom: bool = False, +) -> Callable: + """Registers a symbolic function. + + Usage:: + + ``` + @onnx_symbolic( + "aten::symbolic_b", + opset=10, + decorate=[quantized_aten_handler(scale=1 / 128, zero_point=0)], + ) + @symbolic_helper.parse_args("v", "v", "b") + def symbolic_b(g: _C.Graph, x: _C.Value, y: _C.Value, arg1: bool) -> _C.Value: ... + ``` + + Args: + name: The qualified name of the function in the form of 'domain::op'. + E.g. 'aten::add'. + opset: The opset versions of the function to register at. + decorate: A sequence of decorators to apply to the function. + custom: Whether the function is a custom symbolic function. + + Raises: + ValueError: If the separator '::' is not in the name. + """ + + def wrapper(func: Callable[_P, _R]) -> Callable[_P, _R]: + decorated = func + if decorate is not None: + for decorate_func in decorate: + decorated = decorate_func(decorated) + + global registry + nonlocal opset + if isinstance(opset, OpsetVersion): + opset = (opset,) + for opset_version in opset: + registry.register(name, opset_version, decorated, custom=custom) + + # Return the original function because the decorators in "decorate" are only + # specific to the instance being registered. + return func + + return wrapper + + +def custom_onnx_symbolic( + name: str, + opset: OpsetVersion | Sequence[OpsetVersion], + decorate: Sequence[Callable] | None = None, +) -> Callable: + """Registers a custom symbolic function. + + Args: + name: the qualified name of the function. + opset: the opset version of the function. + decorate: a sequence of decorators to apply to the function. + + Returns: + The decorator. + + Raises: + ValueError: If the separator '::' is not in the name. + """ + return onnx_symbolic(name, opset, decorate, custom=True) + + +# The registry for all symbolic functions. +registry = SymbolicRegistry() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_helper.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..0e00a7a5cd263948216fd0382215e30b09d2ed13 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_helper.py @@ -0,0 +1,2407 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + + +__all__ = [ + "_apply_params", + "_arange_cast_helper", + "_arange_helper", + "_argmin_argmax_helper", + "_as_list_type", + "_avgpool_helper", + "_batchnorm_helper", + "_block_list_in_opset", + "_embedding_bag_helper", + "_flatten_helper", + "_generate_wrapped_number", + "_get_const", + "_get_dim_for_cross", + "_get_interpolate_attributes", + "_get_tensor_dim_size", + "_get_tensor_rank", + "_get_tensor_sizes", + "_handle_reduce_dim_none", + "_if_scalar_type_as", + "_index_fill_reshape_helper", + "_interpolate_get_scales_and_mode", + "_interpolate_get_scales_if_available", + "_interpolate_get_scales", + "_interpolate_helper", + "_interpolate_size_to_scales", + "_interpolate_warning", + "_is_bool", + "_is_constant", + "_is_fp", + "_is_list", + "_is_none", + "_is_onnx_constant", + "_is_packed_list", + "_is_scalar_list", + "_is_split_static", + "_is_tensor_list", + "_is_tensor", + "_is_tuple_construct", + "_is_value", + "_linalg_vector_norm_helper", + "_lt_helper", + "_max_helper", + "_maybe_cast_reduce_op_input", + "_maybe_cast_to_type", + "_maybe_get_const", + "_maybe_get_scalar", + "_min_helper", + "_node_get", + "_numel_helper", + "_onnx_opset_unsupported_detailed", + "_onnx_opset_unsupported", + "_onnx_unsupported", + "_op_with_optional_float_cast", + "_optional_input_placeholder_tensor", + "_overload_by_arg_count", + "_parse_arg", + "_reduce_op_symbolic_helper", + "_reduce_with_dtype_helper", + "_reducesum_helper", + "_repeat_interleave_single_value_repeat_helper", + "_repeat_interleave_split_helper", + "_reshape_helper", + "_scalar", + "_scatter_helper", + "_select_helper", + "_size_helper", + "_slice_helper", + "_sort_helper", + "_squeeze_helper", + "_topk_helper", + "_try_get_scalar_type", + "_type_promote_from_values", + "_unbind_helper", + "_unimplemented", + "_unpack_list", + "_unpack_quantized_tensor", + "_unpack_tuple", + "_unsqueeze_helper", + "_var_mean_helper", + "args_have_same_dtype", + "cast_pytorch_to_onnx", + "check_training_mode", + "dequantize_helper", + "is_complex_value", + "parse_args", + "pytorch_name_to_type", + "quantize_helper", + "quantized_args", + "requantize_bias_helper", + "scalar_name_to_pytorch", + "scalar_type_to_onnx", + "scalar_type_to_pytorch_type", +] + +import functools +import inspect +import math +import sys +import typing +import warnings +from typing import ( + Any, + Concatenate as _Concatenate, + Literal, + NoReturn, + TypeVar as _TypeVar, +) +from typing_extensions import ParamSpec as _ParamSpec, TypeIs as _TypeIs + +import torch +import torch._C._onnx as _C_onnx +from torch import _C +from torch.onnx import _constants, errors +from torch.onnx._internal.torchscript_exporter import _type_utils, jit_utils, utils +from torch.onnx._internal.torchscript_exporter._globals import GLOBALS + + +if typing.TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from torch.types import Number + +_T = _TypeVar("_T") +_U = _TypeVar("_U") +_P = _ParamSpec("_P") + +# --------------------------------------------------------------------------------- +# Helper functions +# --------------------------------------------------------------------------------- + +_ValueDescriptor = Literal[ + "v", + "i", + "is", + "f", + "fs", + "b", + "s", + "t", + "none", +] + + +def _parse_arg( + value, + desc: _ValueDescriptor, + arg_name: str | None = None, + node_name: str | None = None, +): + if desc == "none": + return value + if desc == "v" or not _is_value(value): + return value + + node = value.node() + if node.mustBeNone(): + return None + if node.kind() == "onnx::Constant": + node_val = _node_get(node, "value") + if desc == "i": + return int(node_val) + elif desc == "f": + return float(node_val) + elif desc == "b": + return bool(node_val) + elif desc == "s": + return str(node_val) + elif desc == "t": + return node_val + elif desc == "is": + return [int(v) for v in node_val] + elif desc == "fs": + return [float(v) for v in node_val] + else: + raise errors.SymbolicValueError( + f"ONNX symbolic does not understand the Constant node '{node}' " + f"specified with descriptor '{desc}'.", + value, + ) + elif node.kind() == "prim::ListConstruct": + if desc == "is": + for v in node.inputs(): + element_node = v.node() + if element_node.kind() != "onnx::Constant": + raise errors.SymbolicValueError( + f"Failed to export a node '{element_node}' " + f"(in list node {node}) " + f"because it is not constant. " + f"Please try to make things (e.g. kernel sizes) static if possible.", + value, + ) + return [int(_node_get(v.node(), "value")) for v in value.node().inputs()] + else: + raise errors.SymbolicValueError( + f"ONNX symbolic does not know how to unpack the ListConstruct node that " + f"is not a list of integers: '{node}'", + value, + ) + + if arg_name is None or node_name is None: + raise errors.SymbolicValueError( + f"Expected node type 'onnx::Constant', got '{node.kind()}'.", + value, + ) + + raise errors.SymbolicValueError( + "Expected node type 'onnx::Constant' " + f"for argument '{arg_name}' of node '{node_name}', got '{node.kind()}'.", + value, + ) + + +def _node_get(node: _C.Node, key: str): + """Gets attributes of a node which is polymorphic over return type.""" + if not isinstance(node, _C.Node): + raise AssertionError(f"Expected _C.Node, got {type(node)}") + sel = node.kindOf(key) + return getattr(node, sel)(key) + + +def _is_onnx_constant(value: _C.Value): + """Whether a Value is an ONNX constant.""" + return value.node().kind() == "onnx::Constant" + + +def _maybe_get_const( + value: _C.Value | torch.Tensor | Number | Sequence | None, + descriptor: _ValueDescriptor, +): + # NOTE: prim::Constant at this stage usually means something not compatible in ONNX, + # otherwise it'd be converted to onnx::Constant + # TODO(justinchuby): Replace insinstance with _is_value once we figure out mypy + if isinstance(value, _C.Value) and _is_onnx_constant(value): + return _parse_arg(value, descriptor) + return value + + +def _maybe_get_scalar(value): + value_t = _maybe_get_const(value, "t") + if isinstance(value_t, torch.Tensor) and value_t.shape == (): + return value_t + return value + + +def _get_const(value, desc, arg_name): + if not _is_constant(value): + raise errors.SymbolicValueError( + f"ONNX symbolic expected a constant value of the '{arg_name}' argument, " + f"got '{value}'", + value, + ) + return _parse_arg(value, desc) + + +def _unpack_list(list_value: _C.Value) -> list[_C.Value]: + list_node = list_value.node() + if list_node.kind() != "prim::ListConstruct": + raise errors.SymbolicValueError( + f"ONNX symbolic expected node type prim::ListConstruct, got '{list_node}'.", + list_value, + ) + return list(list_node.inputs()) + + +def _unpack_tuple(tuple_value: _C.Value) -> tuple[_C.Value, ...]: + tuple_node = tuple_value.node() + if not _is_tuple_construct(tuple_value): + raise errors.SymbolicValueError( + f"ONNX symbolic expected node type 'prim::TupleConstruct', " + f"got '{tuple_node.kind()}'.", + tuple_value, + ) + return tuple(tuple_node.inputs()) + + +def _unpack_quantized_tensor(tuple_value: _C.Value) -> tuple[_C.Value, ...]: + """Unpacks a quantized tensor into a tuple of tensor and scale/zero_point. + Args: + tuple_value: A tuple of tensor, scale, zero_point, and optionally axis. + Returns: + A tuple of tensor, scale, zero_point, and optionally axis. + """ + tuple_node = tuple_value.node() + # A quantized tensor is represented as tuple of the form (tensor, scale, zero_point, ) + if not _is_tuple_construct(tuple_value): + raise errors.SymbolicValueError( + f"ONNX symbolic expected the output of `{tuple_node}` to be a quantized " + f"tensor. Is this likely due to missing support for quantized " + f"`{tuple_node.kind()}`. Please create an issue on {_constants.PYTORCH_GITHUB_ISSUES_URL}", + tuple_value, + ) + unpacked = tuple(tuple_node.inputs()) + if not (len(unpacked) == 3 or len(unpacked) == 4): + raise AssertionError(f"Expected 3 or 4 elements, got {len(unpacked)}") + return unpacked + + +# Check if list_value is output from prim::ListConstruct +# This is usually called before _unpack_list to ensure the list can be unpacked. +def _is_packed_list(list_value: Any) -> bool: + return _is_value(list_value) and list_value.node().kind() == "prim::ListConstruct" + + +def parse_args( + *arg_descriptors: _ValueDescriptor, +) -> Callable[[Callable[_Concatenate[_U, _P], _T]], Callable[_Concatenate[_U, _P], _T]]: + """A decorator which converts args from torch._C.Value to built-in types. + + For example: + + ``` + @parse_args('v', 'i', 'fs') + foo(g, a, b, c): + # a is torch._C.Value + # b is int + # c is list of floats + ``` + + Args: + arg_descriptors: list of str, where each element is + a string that specifies the type to convert to. Valid descriptors: + "v": no conversion, keep torch._C.Value. + "i": int + "is": list of int + "f": float + "fs": list of float + "b": bool + "s": str + "t": torch.Tensor + "none": the variable is unused + """ + + def decorator( + fn: Callable[_Concatenate[_U, _P], _T], + ) -> Callable[_Concatenate[_U, _P], _T]: + fn._arg_descriptors = arg_descriptors # type: ignore[attr-defined] + + @functools.wraps(fn) + def wrapper(g: _U, *args: _P.args, **kwargs: _P.kwargs) -> _T: + # some args may be optional, so the length may be smaller + FILE_BUG_MSG = ( + "If you believe this is not due to custom symbolic implementation within your code or " + "an external library, please file an issue at " + "https://github.com/pytorch/pytorch/issues/new?template=bug-report.yml to report this bug." + ) + if len(arg_descriptors) < len(args): + raise AssertionError( + f"A mismatch between the number of arguments ({len(args)}) and " + f"their descriptors ({len(arg_descriptors)}) was found at " + f"symbolic function '{fn.__name__}'. {FILE_BUG_MSG}" + ) + + try: + sig = inspect.signature(fn) + arg_names = list(sig.parameters.keys())[1:] + fn_name = fn.__name__ + except Exception: + # FIXME(justinchuby): Avoid catching Exception. + # Catch a more specific exception instead. + arg_names = [None] * len(args) # type: ignore[list-item] + fn_name = None + # pyrefly: ignore [bad-assignment] + args = [ + _parse_arg(arg, arg_desc, arg_name, fn_name) # type: ignore[method-assign] + for arg, arg_desc, arg_name in zip(args, arg_descriptors, arg_names) + ] + # only support _outputs in kwargs + if len(kwargs) > 1: + raise AssertionError( + f"Symbolic function {fn.__name__}'s '**kwargs' can contain a single " + f"key/value entry. " + f"{FILE_BUG_MSG}" + ) + + if len(kwargs) == 1: + if "_outputs" not in kwargs: + raise AssertionError( + f"Symbolic function {fn.__name__}'s '**kwargs' can only contain " + f"'_outputs' key at '**kwargs'. " + f"{FILE_BUG_MSG}" + ) + return fn(g, *args, **kwargs) + + return wrapper + + return decorator + + +def quantized_args( + *arg_q_descriptors: bool, + scale: float | None = None, + zero_point: int | None = None, + quantize_output: bool = True, +) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: + """A decorator which extends support for quantized version of the base operator. + + Quantization is detected by examining the arguments that are annotated by + `arg_q_descriptors`. + + If quantization is detected, the base operator symbolic function will be wrapped with + argument de-quantization and output quantization. + + Otherwise, only the base symbolic function will be invoked. + + For example: + + ``` + @quantized_args(True, False) + def foo(g, x, y): + return x + y + ``` + + is equivalent to + + ``` + def q_foo(g, x, y): + if is_quantized_tensor(x): + x = dequantize(x) + out = foo(g, x, y) + return quantize(out) + else: + return foo(g, x, y) + ``` + + Args: + arg_q_descriptors: A sequence of bool, where each element represents if the + argument is QTensor for quantized version of this operator. It defaults + to False for unspecified (variable length) arguments. + scale: Quantized output scale. If None, derive from + the first quantized input scale. + zero_point: Quantized output zero point. If None, + derive from the first quantized input zero point. + quantize_output: If True, quantize the output of the base operator. Default is True + """ + + def decorator(fn): + @functools.wraps(fn) + def wrapper(g, *args, **kwargs): + nonlocal scale + nonlocal zero_point + if scale is not None: + _scale = g.op("Constant", value_t=torch.tensor(scale)) + else: + _scale = None + if zero_point is not None: + _zero_point = g.op("Constant", value_t=torch.tensor(zero_point)) + else: + _zero_point = None + + # Support variable length arguments by marking unspecified ones as non-quantized + arg_q_descriptors_extended = arg_q_descriptors + (False,) * ( + len(args) - len(arg_q_descriptors) + ) + descriptor_args = tuple(zip(arg_q_descriptors_extended, args)) + + def _is_arg_quantized(descriptor, arg): + return descriptor and _is_value(arg) and _is_tuple_construct(arg) + + # Run regular symbolic function if none of the argument is QTensor. + is_quantized: list[bool] = [] + for descriptor, arg in descriptor_args: + # ListConstruct + if _is_packed_list(arg): + is_quantized.extend( + _is_arg_quantized(descriptor, arg_input) + for arg_input in arg.node().inputs() + ) + else: + is_quantized.append(_is_arg_quantized(descriptor, arg)) + + if not any(is_quantized): + return fn(g, *args, **kwargs) + + # Dequantize arguments that are quantized + non_quantized_args = [] + for descriptor, arg in descriptor_args: + if _is_arg_quantized(descriptor, arg): + # Quantized arg is a tuple of (value, scale, zero_point) + dequantized_arg, arg_scale, arg_zero_point, _ = dequantize_helper( + g, arg + ) + non_quantized_args.append(dequantized_arg) + # Set scale and zero_point to the first quantized input if not already set + if _scale is None: + _scale = arg_scale + if _zero_point is None: + _zero_point = arg_zero_point + # ListConstruct + elif _is_packed_list(arg): + for arg_input in arg.node().inputs(): + if _is_arg_quantized(descriptor, arg_input): + # Quantized arg is a tuple of (value, scale, zero_point) + ( + dequantized_arg, + arg_scale, + arg_zero_point, + _, + ) = dequantize_helper(g, arg_input) + # Set scale and zero_point to the first quantized input if not already set + if _scale is None: + _scale = arg_scale + if _zero_point is None: + _zero_point = arg_zero_point + arg_input.replaceAllUsesWith(dequantized_arg) + non_quantized_args.append(arg) + else: + # Non-quantized arg + non_quantized_args.append(arg) + # TODO(justinchuby): Only single output is supported for now. We may want to + # support multiple outputs in the future. + output = fn(g, *non_quantized_args, **kwargs) + + if _scale is None: + raise AssertionError("Bug: Scale must be set for quantized operator") + if _zero_point is None: + raise AssertionError( + "Bug: Zero point must be set for quantized operator" + ) + + if quantize_output: + return quantize_helper(g, output, _scale, _zero_point) + return output + + return wrapper + + return decorator + + +def _scalar(x: Any) -> Number | None: + """Convert a scalar tensor into a Python value.""" + if isinstance(x, torch.Tensor) and x.shape == (): + return x.item() + return None + + +def _if_scalar_type_as(self, tensor): + """ + Convert self into the same type of tensor, as necessary. + We only support implicit casting for scalars, so we never + actually need to insert an ONNX cast operator here; just + fix up the scalar. + """ + if isinstance(self, _C.Value): + return self + + scalar_type = _type_utils.JitScalarType.from_value( + tensor, _type_utils.JitScalarType.UNDEFINED + ) + if scalar_type != _type_utils.JitScalarType.UNDEFINED: + ty = scalar_type.scalar_name().lower() + return getattr(self, ty)() + return self + + +def _is_none(x: Any) -> bool: + return x is None or (x.node().mustBeNone() if isinstance(x, _C.Value) else False) + + +def _is_value(x: Any) -> _TypeIs[_C.Value]: + return isinstance(x, _C.Value) + + +def _is_constant(value: Any) -> bool: + return not _is_value(value) or value.node().kind() in { + "onnx::Constant", + "prim::Constant", + } + + +def _is_tensor(x: _C.Value) -> bool: + return x.type().isSubtypeOf(_C.TensorType.get()) + + +# Note: _C.JitType is not exposed to Python and cannot be checked in runtime. +def _as_list_type(jit_type: _C.JitType) -> _C.ListType | None: + if isinstance(jit_type, _C.ListType): + return jit_type + return None + + +def _is_list(x: _C.Value) -> bool: + return _as_list_type(x.type()) is not None + + +def _is_tensor_list(x: _C.Value) -> bool: + x_type = _as_list_type(x.type()) + if x_type is None: + return False + return isinstance(x_type.getElementType(), _C.TensorType) + + +def _is_scalar_list(x: _C.Value) -> bool: + """Checks if x is a scalar list, for example: List[float], List[int]. + + Besides checking the type is ListType, we also check if the data type is + a valid ONNX data type. + """ + x_type = _as_list_type(x.type()) + if x_type is None: + return False + scalar_type = _type_utils.JitScalarType.from_value(x) + return scalar_type.onnx_compatible() + + +def _is_tuple_construct(x: _C.Value) -> bool: + return x.node().kind() == "prim::TupleConstruct" + + +def is_complex_value(x: _C.Value) -> bool: + if not _is_value(x): + raise AssertionError("Expected a Value") + return _type_utils.JitScalarType.from_value( + x, _type_utils.JitScalarType.UNDEFINED + ) in { + _type_utils.JitScalarType.COMPLEX32, + _type_utils.JitScalarType.COMPLEX64, + _type_utils.JitScalarType.COMPLEX128, + } + + +def _get_tensor_rank(x: _C.Value) -> int | None: + if not _is_tensor(x) or x.type() is None: + return None + x_type = x.type() + x_type = typing.cast(_C.TensorType, x_type) + return x_type.dim() + + +def _get_tensor_sizes(x: _C.Value, allow_nonstatic: bool = True): + if not _is_tensor(x) or x.type() is None: + return None + x_type = x.type() + x_type = typing.cast(_C.TensorType, x_type) + if allow_nonstatic: + # Each individual symbol is returned as None. + # e.g. [1, "a", "b"] -> [1, None, None] + return x_type.varyingSizes() + # returns None, if exists any symbol in sizes. + # e.g. [1, "a", "b"] -> None + return x_type.sizes() + + +def _get_tensor_dim_size(x: _C.Value, dim: int) -> int | None: + sizes = _get_tensor_sizes(x) + return sizes[dim] if sizes else None + + +def _get_dim_for_cross(x: _C.Value, dim: int | None): + if dim == -1: + tensor_rank = _get_tensor_rank(x) + if tensor_rank is None: + raise AssertionError("Expected tensor_rank to be non-None") + return dim + tensor_rank + # If dim is not given, it defaults to the first dimension found with the size 3 + if dim is None: + sizes = _get_tensor_sizes(x) + if sizes is None: + raise AssertionError("Expected sizes to be non-None") + for index, size in enumerate(sizes): + if size is not None and size == 3: + return index + return dim + + +def _unimplemented(op: str, msg: str, value: _C.Value | None = None) -> None: + # For BC reasons, the behavior for Caffe2 does not raise exception for unimplemented operators + if GLOBALS.operator_export_type == _C_onnx.OperatorExportTypes.ONNX: + _onnx_unsupported(f"{op}, {msg}", value) + + +def _onnx_unsupported(op_name: str, value: _C.Value | None = None) -> NoReturn: + message = ( + f"Unsupported: ONNX export of operator {op_name}. " + f"Please feel free to request support or submit a pull request " + f"on PyTorch GitHub: {_constants.PYTORCH_GITHUB_ISSUES_URL}" + ) + if isinstance(value, _C.Value): + raise errors.SymbolicValueError( + message, + value, + ) + raise errors.OnnxExporterError(message) + + +def _onnx_opset_unsupported( + op_name: str, + current_opset: int, + supported_opset: int, + value: _C.Value | None = None, +) -> NoReturn: + message = ( + f"Unsupported: ONNX export of {op_name} in opset {current_opset}. " + f"Please try opset version {supported_opset}." + ) + if isinstance(value, _C.Value): + raise errors.SymbolicValueError( + message, + value, + ) + raise errors.OnnxExporterError(message) + + +def _onnx_opset_unsupported_detailed( + op_name: str, + current_opset: int, + supported_opset: int, + reason: str, + value: _C.Value | None = None, +) -> NoReturn: + message = ( + f"Unsupported: ONNX export of {op_name} in " + f"opset {current_opset}. {reason}. Please try opset version {supported_opset}." + ) + if isinstance(value, _C.Value): + raise errors.SymbolicValueError( + message, + value, + ) + raise errors.OnnxExporterError(message) + + +def _block_list_in_opset(name: str): + def symbolic_fn(*args, **kwargs): + raise errors.OnnxExporterError( + f"ONNX export failed on {name}, which is not implemented for opset " + f"{GLOBALS.export_onnx_opset_version}. " + "Try exporting with other opset versions." + ) + + return symbolic_fn + + +def _try_get_scalar_type(*args) -> _type_utils.JitScalarType | None: + for arg in args: + scalar_type = _type_utils.JitScalarType.from_value( + arg, _type_utils.JitScalarType.UNDEFINED + ) + if scalar_type != _type_utils.JitScalarType.UNDEFINED: + return scalar_type + return None + + +def _type_promote_from_values(*args) -> _type_utils.JitScalarType: + undef = _type_utils.JitScalarType.UNDEFINED + jit_types = [_try_get_scalar_type(arg) for arg in args] + if len(jit_types) == 0: + return undef + if len(jit_types) == 1: + return jit_types[0] # type: ignore[return-value] + new_dtype = jit_types[0].dtype() # type: ignore[union-attr] + for t in jit_types: + new_dtype = torch.promote_types(new_dtype, t.dtype()) # type: ignore[union-attr] + return _type_utils.JitScalarType.from_dtype(new_dtype) + + +def _maybe_cast_to_type( + g: jit_utils.GraphContext, value, jit_type: _type_utils.JitScalarType +): + if ( + _type_utils.JitScalarType.from_value(value, _type_utils.JitScalarType.UNDEFINED) + != jit_type + ): + return g.op( + "Cast", + value, + to_i=jit_type.onnx_type(), + ) + return value + + +def _select_helper(g: jit_utils.GraphContext, self, dim, index, apply_reshape=True): + index_const = _maybe_get_scalar(index) + index_dim = _get_tensor_rank(index) + if not _is_value(index_const): + # Index is a constant scalar. Make it a size 1 constant tensor. + index = g.op("Constant", value_t=torch.LongTensor([index_const])) + elif index_dim is not None and apply_reshape: + if index_dim == 0: + # Index is a scalar. Reshape it to a size 1 tensor. + index = _reshape_helper( + g, index, g.op("Constant", value_t=torch.LongTensor([1])) + ) + + index_scalar_type = _type_utils.JitScalarType.from_value( + index, _type_utils.JitScalarType.UNDEFINED + ) + if index_scalar_type not in { + _type_utils.JitScalarType.INT64, + _type_utils.JitScalarType.INT, + }: + index = g.op("Cast", index, to_i=_C_onnx.TensorProtoDataType.INT64) + return g.op("Gather", self, index, axis_i=dim) + + +def _slice_helper( + g: jit_utils.GraphContext, + input, + axes, + starts, + ends, + steps=None, +): + if g.opset <= 9: + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import ( + _slice as _slice9, + ) + + return _slice9(g, input, axes, starts, ends) + else: + from torch.onnx._internal.torchscript_exporter.symbolic_opset10 import ( + _slice as _slice10, + ) + + return _slice10(g, input, axes, starts, ends, steps) + + +def _is_fp(value) -> bool: + return _type_utils.JitScalarType.from_value( + value, _type_utils.JitScalarType.UNDEFINED + ) in { + _type_utils.JitScalarType.FLOAT, + _type_utils.JitScalarType.DOUBLE, + _type_utils.JitScalarType.HALF, + _type_utils.JitScalarType.BFLOAT16, + } + + +def _is_bool(value) -> bool: + return ( + _type_utils.JitScalarType.from_value(value, _type_utils.JitScalarType.UNDEFINED) + == _type_utils.JitScalarType.BOOL + ) + + +def _generate_wrapped_number(g: jit_utils.GraphContext, scalar): + """Creates a wrapped number based on https://github.com/pytorch/pytorch/issues/9515. + + A Tensor is a considered a "wrapped number" if it is + auto-wrapped from a C++ or Python number type. Integer types are + wrapped as 0-dim int64 tensors and floating-point types are + wrapped as 0-dim double tensors. + + The input to this function is constant value. If the data type + is a floating point type, it is converted to a 0-dim double + tensor, else it is converted to a 0-dim tensor of its original type + """ + if isinstance(scalar, torch.Tensor): + raise AssertionError("Expected scalar, not torch.Tensor") + if isinstance(scalar, float): + return g.op("Constant", value_t=torch.tensor(scalar, dtype=torch.double)) + return g.op("Constant", value_t=torch.tensor(scalar)) + + +def _sort_helper(g: jit_utils.GraphContext, input, dim, descending=True, out=None): + if out is not None: + _unimplemented("Sort", "Out parameter is not supported") + shape_ = g.op("Shape", input) + dim_size_ = g.op( + "Gather", + shape_, + g.op("Constant", value_t=torch.tensor([dim], dtype=torch.int64)), + ) + if g.opset <= 10: + if not descending: + _unimplemented("Sort", "Ascending is not supported") + return g.op("TopK", input, dim_size_, axis_i=dim, outputs=2) + else: + return g.op( + "TopK", input, dim_size_, axis_i=dim, largest_i=descending, outputs=2 + ) + + +def _topk_helper( + g: jit_utils.GraphContext, input, k, dim, largest=True, sorted=False, out=None +): + if out is not None: + _unimplemented("TopK", "Out parameter is not supported") + if not _is_value(k): + k = g.op("Constant", value_t=torch.tensor([k], dtype=torch.int64)) + else: + k = _reshape_helper(g, k, g.op("Constant", value_t=torch.tensor([1]))) + if _try_get_scalar_type(k) != _type_utils.JitScalarType.INT64: + k = g.op("Cast", k, to_i=_C_onnx.TensorProtoDataType.INT64) + if g.opset <= 10: + if not largest: + _unimplemented("TopK", "Ascending is not supported") + return g.op("TopK", input, k, axis_i=dim, outputs=2) + else: + return g.op( + "TopK", input, k, axis_i=dim, largest_i=largest, sorted_i=sorted, outputs=2 + ) + + +def _lt_helper(g: jit_utils.GraphContext, input, other): + if g.opset <= 8: + from torch.onnx._internal.torchscript_exporter.symbolic_opset8 import lt as _lt8 + + return _lt8(g, input, other) + else: + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import lt as _lt9 + + return _lt9(g, input, other) + + +def _interpolate_warning(interpolate_mode): + onnx_op = ( + "onnx:Resize" if GLOBALS.export_onnx_opset_version >= 10 else "onnx:Upsample" + ) + warnings.warn( + "You are trying to export the model with " + + onnx_op + + " for ONNX opset version " + "" + str(GLOBALS.export_onnx_opset_version) + ". " + "This operator might cause results to not match the expected results by PyTorch.\n" + "ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11. " + "Attributes to determine how to transform the input were added in onnx:Resize in opset 11 " + "to support Pytorch's behavior (like coordinate_transformation_mode and nearest_mode).\n" + "We recommend using opset 11 and above for models using this operator.", + stacklevel=2, + ) + + +def _unsqueeze_helper(g: jit_utils.GraphContext, input, axes_i): + if len(axes_i) == 0: + # unnecessary unsqueeze if axes length==0 + return input + elif _is_constant(axes_i[0]): + if g.opset >= 13: + axes = g.op("Constant", value_t=torch.tensor(axes_i, dtype=torch.long)) + return g.op("Unsqueeze", input, axes) + return g.op("Unsqueeze", input, axes_i=axes_i) + # Tensor type + if g.opset < 13: + raise errors.SymbolicValueError( + "Opset version must be >= 13 for Unsqueeze with dynamic axes.", input + ) + return g.op("Unsqueeze", input, axes_i[0]) + + +def _squeeze_helper(g: jit_utils.GraphContext, input, axes_i): + if _is_constant(axes_i[0]): + if g.opset >= 13: + axes = g.op("Constant", value_t=torch.tensor(axes_i, dtype=torch.long)) + return g.op("Squeeze", input, axes) + return g.op("Squeeze", input, axes_i=axes_i) + # Tensor type + if g.opset < 13: + raise errors.SymbolicValueError( + "Opset version must be >= 13 for Squeeze with dynamic axes.", input + ) + axes_t = axes_i[0] + axes_rank = _get_tensor_rank(axes_t) + if axes_rank is None: + raise AssertionError("Expected axes_rank to be non-None") + if axes_rank > 1: + raise errors.SymbolicValueError( + "For Squeeze axses as input, the axes rank must be one in ONNX spec.", input + ) + elif axes_rank == 0: + # The axes is a scalar. Unsqueeze it to a rank 1 tensor. + axes_t = _unsqueeze_helper(g, axes_t, [0]) + return g.op("Squeeze", input, axes_t) + return g.op("Squeeze", input, axes_t) + + +def _reducesum_helper( + g: jit_utils.GraphContext, + input, + axes_i=None, + keepdims_i=1, + noop_with_empty_axes_i=0, +): + keepdims_i = _maybe_get_const(keepdims_i, "i") + if g.opset >= 13: + if axes_i: + if not _is_value(axes_i): + axes_i = g.op( + "Constant", value_t=torch.tensor(axes_i, dtype=torch.long) + ) + return g.op( + "ReduceSum", + input, + axes_i, + keepdims_i=keepdims_i, + noop_with_empty_axes_i=noop_with_empty_axes_i, + ) + return g.op( + "ReduceSum", + input, + keepdims_i=keepdims_i, + noop_with_empty_axes_i=noop_with_empty_axes_i, + ) + else: + return g.op("ReduceSum", input, axes_i=axes_i, keepdims_i=keepdims_i) + + +def _interpolate_size_to_scales(g: jit_utils.GraphContext, input, output_size, dim): + output_size = _maybe_get_const(output_size, "is") + if _is_value(output_size): + offset = 2 + offsets = g.op("Constant", value_t=torch.ones(offset, dtype=torch.float32)) + dividend = g.op("Cast", output_size, to_i=_C_onnx.TensorProtoDataType.FLOAT) + divisor = _slice_helper( + g, g.op("Shape", input), axes=[0], ends=[sys.maxsize], starts=[offset] + ) + divisor = g.op("Cast", divisor, to_i=_C_onnx.TensorProtoDataType.FLOAT) + scale_dims = g.op("Div", dividend, divisor) + scales = g.op("Concat", offsets, scale_dims, axis_i=0) + else: + scales_constant = [ + 1.0 + if i < 2 + else float(output_size[-(dim - i)]) + / float(input.type().sizes()[-(dim - i)]) + for i in range(dim) + ] + scales = g.op( + "Constant", value_t=torch.tensor(scales_constant, dtype=torch.float32) + ) + return scales + + +def _interpolate_get_scales_if_available(g: jit_utils.GraphContext, scales): + available_scales = _maybe_get_const(scales[0], "fs") != -1 and not _is_none( + scales[0] + ) + + if not available_scales: + return None + + offsets = g.op("Constant", value_t=torch.ones(2, dtype=torch.float32)) + scales_list = g.op( + "Constant", value_t=torch.tensor(_maybe_get_const(scales[0], "fs")) + ) + scales = g.op("Concat", offsets, scales_list, axis_i=0) + return scales + + +def _get_interpolate_attributes(g: jit_utils.GraphContext, mode, args): + if mode == "nearest": + align_corners = None + scales = args[0:] + else: + align_corners = args[0] + scales = args[1:] + scales = _interpolate_get_scales_if_available(g, scales) + return scales, align_corners + + +def _interpolate_get_scales(g: jit_utils.GraphContext, scale_factor, dim): + offsets = g.op("Constant", value_t=torch.ones(2, dtype=torch.float32)) + scale_factor_rank = _get_tensor_rank(scale_factor) + if isinstance(scale_factor.type(), _C.ListType) or ( + scale_factor_rank is not None and scale_factor_rank > 0 + ): + return g.op("Concat", offsets, scale_factor, axis_i=0) + else: + scale_factor = _unsqueeze_helper(g, scale_factor, [0]) + scale_factor = g.op( + "Cast", scale_factor, to_i=_C_onnx.TensorProtoDataType.FLOAT + ) + scales = [scale_factor for i in range(dim - 2)] + scale_factor = g.op("Concat", offsets, *scales, axis_i=0) + return scale_factor + + +def _interpolate_get_scales_and_mode( + g: jit_utils.GraphContext, input, size, scale_factor, mode, align_corners +): + mode = _maybe_get_const(mode, "s") + if "linear" in mode: + mode = "linear" + if "cubic" in mode: + mode = "cubic" + _interpolate_warning(mode) + + align_corners = _maybe_get_const(align_corners, "b") + if isinstance(align_corners, bool) and align_corners: + return _unimplemented("interpolate", "align_corners == True") + + if not input.type().dim(): + return _unimplemented("interpolate", "missing input shape") + dim = input.type().dim() + + if not _is_none(scale_factor): + scale_factor = _interpolate_get_scales(g, scale_factor, dim) + elif not _is_none(size): + if not _is_packed_list(size): + is_scalar = _maybe_get_const(size, "t").dim() == 0 + if is_scalar: + size = _unsqueeze_helper(g, size, [0]) + size = [size for i in range(dim - 2)] + size = g.op("Concat", *size, axis_i=0) + scale_factor = _interpolate_size_to_scales(g, input, size, dim) + else: + return _unimplemented( + "interpolate", "Both size and scales are None in __interpolate" + ) + return scale_factor, mode + + +def _argmin_argmax_helper( + g: jit_utils.GraphContext, + input: torch._C.Value, + dim: torch._C.Value, + keepdim: bool, + op_name: str, +): + def op_wrapper(input, axis_i, keepdims_i): + if g.opset >= 12: + return g.op( + op_name, + input, + axis_i=axis_i, + keepdims_i=keepdims_i, + select_last_index_i=False, + ) + return g.op(op_name, input, axis_i=axis_i, keepdims_i=keepdims_i) + + if _is_none(dim): + flattened = _reshape_helper( + g, input, g.op("Constant", value_t=torch.tensor([-1])) + ) + output = op_wrapper(flattened, axis_i=0, keepdims_i=False) + if keepdim: + input_shape = g.op("Shape", input) + input_shape_shape = g.op("Shape", input_shape) + new_shape = g.op( + "ConstantOfShape", + input_shape_shape, + value_t=torch.tensor([1], dtype=torch.int64), + ) + output = g.op("Reshape", output, new_shape) + return output + + dim = _parse_arg(dim, "i") + return op_wrapper(input, axis_i=dim, keepdims_i=keepdim) + + +def _interpolate_helper(name, dim, interpolate_mode): + @quantized_args(True, False, False) + def symbolic_fn(g, input, output_size, *args): + scales, align_corners = _get_interpolate_attributes(g, interpolate_mode, args) + align_corners = _maybe_get_scalar(align_corners) + coordinate_transformation_mode = ( + "asymmetric" + if interpolate_mode == "nearest" + else "align_corners" + if align_corners + else "half_pixel" + ) + + if scales is None: + input_size = g.op("Shape", input) + input_size_beg = _slice_helper( + g, input_size, axes=[0], ends=[2], starts=[0] + ) + output_size = g.op( + "Cast", output_size, to_i=_C_onnx.TensorProtoDataType.INT64 + ) + output_size = g.op("Concat", input_size_beg, output_size, axis_i=0) + + if g.opset >= 13: + empty_roi = _optional_input_placeholder_tensor(g) + empty_scales = _optional_input_placeholder_tensor(g) + else: + empty_roi = g.op( + "Constant", value_t=torch.tensor([], dtype=torch.float32) + ) + empty_scales = g.op( + "Constant", value_t=torch.tensor([], dtype=torch.float32) + ) + + return g.op( + "Resize", + input, + empty_roi, + empty_scales, + output_size, + coordinate_transformation_mode_s=coordinate_transformation_mode, + cubic_coeff_a_f=-0.75, # only valid when mode="cubic" + mode_s=interpolate_mode, # nearest, linear, or cubic + nearest_mode_s="floor", + ) # only valid when mode="nearest" + else: + if g.opset >= 13: + empty_roi = _optional_input_placeholder_tensor(g) + else: + empty_roi = g.op( + "Constant", value_t=torch.tensor([], dtype=torch.float32) + ) + + return g.op( + "Resize", + input, + empty_roi, + scales, + coordinate_transformation_mode_s=coordinate_transformation_mode, + cubic_coeff_a_f=-0.75, # only valid when mode="cubic" + mode_s=interpolate_mode, # nearest, linear, or cubic + nearest_mode_s="floor", + ) # only valid when mode="nearest" + + return symbolic_fn + + +def __interpolate_helper( + g: jit_utils.GraphContext, + input, + size, + scale_factor, + mode, + align_corners, + recompute_scale_factor, +): + mode = _maybe_get_const(mode, "s") + if "linear" in mode: + mode = "linear" + if "cubic" in mode: + mode = "cubic" + align_corners = _maybe_get_const(align_corners, "b") + align_corners = False if not isinstance(align_corners, bool) else align_corners + coordinate_transformation_mode = ( + "asymmetric" + if mode == "nearest" + else "align_corners" + if align_corners + else "half_pixel" + ) + + if not _is_none(size): + input_size = g.op("Shape", input) + input_size = _slice_helper(g, input_size, axes=[0], ends=[2], starts=[0]) + # in some cases size is not a packed list but size is a scalar + # We need to also verify that (_maybe_get_const(size, "t").dim() == 0) + # but this information is not always available. Try to get the dim, + # and if not assume that it is not a scalar. + try: + is_scalar = not _is_packed_list(size) and ( + _maybe_get_const(size, "t").dim() == 0 + ) + except AttributeError: + is_scalar = not _is_packed_list(size) + if not is_scalar: + warnings.warn( + "Cannot verify if the output_size is a scalar " + "while exporting interpolate. Assuming that it is not a scalar.", + stacklevel=2, + ) + + if is_scalar: + rank = _get_tensor_rank(input) + if rank is None: + return _unimplemented( + "interpolate (with a scalar output_size)", + "missing input shape (try giving an array of output_size values)", + ) + size = _unsqueeze_helper(g, size, [0]) + size = [size for i in range(rank - 2)] + size = g.op("Concat", *size, axis_i=0) + size = g.op("Cast", size, to_i=_C_onnx.TensorProtoDataType.INT64) + size = g.op("Concat", input_size, size, axis_i=0) + + if g.opset >= 13: + empty_roi = _optional_input_placeholder_tensor(g) + empty_scales = _optional_input_placeholder_tensor(g) + else: + empty_roi = g.op("Constant", value_t=torch.tensor([], dtype=torch.float32)) + empty_scales = g.op( + "Constant", value_t=torch.tensor([], dtype=torch.float32) + ) + + return g.op( + "Resize", + input, + empty_roi, + empty_scales, + size, + coordinate_transformation_mode_s=coordinate_transformation_mode, + cubic_coeff_a_f=-0.75, # only valid when mode="cubic" + mode_s=mode, # nearest, linear, or cubic + nearest_mode_s="floor", + ) + else: # if not _is_none(scales) + rank = _get_tensor_rank(input) + if rank is None: + return _unimplemented("interpolate (with scales)", "missing input shape") + + if g.opset >= 13: + empty_roi = _optional_input_placeholder_tensor(g) + else: + empty_roi = g.op("Constant", value_t=torch.tensor([], dtype=torch.float32)) + + scales = _interpolate_get_scales(g, scale_factor, rank) + return g.op( + "Resize", + input, + empty_roi, + scales, + coordinate_transformation_mode_s=coordinate_transformation_mode, + cubic_coeff_a_f=-0.75, # only valid when mode="cubic" + mode_s=mode, # nearest, linear, or cubic + nearest_mode_s="floor", + ) # only valid when mode="nearest" + + +def _unbind_helper(g: jit_utils.GraphContext, self, dim, _outputs): + if g.opset < 11: + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import unbind + elif g.opset <= 12: + from torch.onnx._internal.torchscript_exporter.symbolic_opset11 import ( + unbind, # type: ignore[no-redef] + ) + else: + from torch.onnx._internal.torchscript_exporter.symbolic_opset13 import ( + unbind, # type: ignore[no-redef] + ) + return unbind(g, self, dim, _outputs) + + +def _scatter_helper(g: jit_utils.GraphContext, self, dim, index, src): + if g.opset <= 10: + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import scatter + else: + # for mypy, scatter was imported two lines above + from torch.onnx._internal.torchscript_exporter.symbolic_opset11 import ( + scatter, # type: ignore[no-redef] + ) + return scatter(g, self, dim, index, src) + + +def _repeat_interleave_split_helper(g: jit_utils.GraphContext, self, reps, dim): + if g.opset <= 12: + split_out = g.op("Split", self, split_i=[1] * reps, axis_i=dim, outputs=reps) + else: + from torch.onnx._internal.torchscript_exporter.symbolic_opset13 import split + + repeats = g.op("Constant", value_t=torch.tensor([1] * reps)) + split_out = split(g, self, repeats, dim, _outputs=reps) + return split_out if reps > 1 else [split_out] + + +def _repeat_interleave_single_value_repeat_helper( + g: jit_utils.GraphContext, self, repeats, dim +): + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import ( + flatten, + unsqueeze, + ) + + if not _is_tensor(repeats): + repeats = g.op("Constant", value_t=torch.LongTensor(repeats)) + + const_repeats: bool = _is_constant(repeats) + reps = _maybe_get_const(repeats, "t") + + # Convert 'repeats' to 1-d if it is 0-d. + if _get_tensor_rank(repeats) == 0: + repeats = g.op("Reshape", repeats, g.op("Constant", value_t=torch.tensor([1]))) + + # Create a new dim of size 1, then expand it to be 'repeats' long, and finally collapse it. + unsqueezed = unsqueeze(g, self, dim + 1) + + # repeats_per_dim is 1 for all dims except for the new unsqueezed dim, where it has value 'repeats'. + if const_repeats: + # 'Repeats' is a constant, 'repeats_per_dim' can be a constant. + onehot = torch.ones(_get_tensor_rank(unsqueezed), dtype=torch.int64) # type: ignore[arg-type] + onehot[dim + 1] = reps + repeats_per_dim = g.op("Constant", value_t=onehot) + else: + # 'Repeats' is a variable, 'repeats_per_dim' cannot be a constant. + onehot = g.op( + "OneHot", + unsqueeze(g, dim + 1, 0), # indices, must be >= 1-dimensional + g.op( + "Constant", value_t=torch.tensor(_get_tensor_rank(unsqueezed)) + ), # depth + g.op( + "Concat", g.op("Constant", value_t=torch.tensor([1])), repeats, axis_i=0 + ), # on/off values + ) + repeats_per_dim = flatten(g, onehot, 0, 1) + + tiled = g.op("Tile", unsqueezed, repeats_per_dim) + return flatten(g, tiled, dim, dim + 1) + + +def _arange_cast_helper( + g: jit_utils.GraphContext, end, start=None, step=None, dtype=None +) -> tuple[ + _type_utils.JitScalarType, + _C.Value | None, + _C.Value | None, + _C.Value | None, +]: + def _is_all_integral(scalars): + for scalar in scalars: + scalar_type = _type_utils.JitScalarType.from_value( + scalar, _type_utils.JitScalarType.UNDEFINED + ) + if ( + scalar_type != _type_utils.JitScalarType.INT64 + and scalar_type != _type_utils.JitScalarType.UNDEFINED + ): + return False + return True + + # This logic is based on torch.arange docs. If "dtype" is provided, + # infer input types from dtype. If not, then check if any of start, stop, + # or step are floating point, and infer the type from get_default. + # Otherwise, the dtype is inferred to be torch.int64. + if dtype is None or (_is_value(dtype) and _is_none(dtype)): + if _is_all_integral([start, end, step]): + scalar_type = _type_utils.JitScalarType.INT64 + else: + scalar_type = _type_utils.JitScalarType.from_dtype( + torch.get_default_dtype() + ) + else: + if not isinstance(dtype, int): + raise AssertionError(f"Expected dtype to be int, got {type(dtype)}") + # TODO(justinchuby): Check if dtype is indeed a int. + scalar_type = _type_utils.JitScalarType(dtype) + + start = g.op("Cast", start, to_i=scalar_type.onnx_type()) if start else None + end = g.op("Cast", end, to_i=scalar_type.onnx_type()) if end else None + step = g.op("Cast", step, to_i=scalar_type.onnx_type()) if step else None + return scalar_type, end, start, step + + +def _arange_helper(g: jit_utils.GraphContext, *args): + if g.opset <= 10: + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import arange + else: + from torch.onnx._internal.torchscript_exporter.symbolic_opset11 import ( + arange, # type: ignore[no-redef] + ) + return arange(g, *args) + + +def _size_helper(g: jit_utils.GraphContext, self, dim): + full_shape = g.op("Shape", self) + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import select + + return select(g, full_shape, g.op("Constant", value_t=torch.tensor([0])), dim) + + +def _index_fill_reshape_helper(g: jit_utils.GraphContext, self, dim, index): + # 1. reshape index => [1, ..., 1, dim, 1, ..., 1] + # 2. expand index => [..., dim, ...], same shape as self except for dim. + # 3. expand value as well. + # 4. apply onnx::scatter. + + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import expand + + if g.opset <= 10: + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import scatter + else: + # for mypy, scatter was imported two lines above + from torch.onnx._internal.torchscript_exporter.symbolic_opset11 import ( + scatter, # type: ignore[no-redef] + ) + + if self.type().dim() is None: + return _unimplemented("index_fill", "input rank not accessible") + self_dim = self.type().dim() + dim_value = _parse_arg(dim, "i") + if dim_value < 0: + dim_value += self_dim + unsqueezed_index = _unsqueeze_helper( + g, index, [i for i in range(self_dim) if i != dim_value] + ) + expanded_index_shape = scatter( + g, g.op("Shape", self), 0, _unsqueeze_helper(g, dim, [0]), g.op("Shape", index) + ) + expanded_index = expand(g, unsqueezed_index, expanded_index_shape, None) + return expanded_index_shape, expanded_index + + +# By default, when any value in the 'shape' input is equal to zero +# the corresponding dimension value is copied from the input tensor dynamically. +# allowzero=1 indicates that if any value in the 'shape' input is set to zero, +# the zero value is honored, similar to NumPy. +# allowzero=1 is only supported for opset version >= 14. +def _reshape_helper(g: jit_utils.GraphContext, input, shape, allowzero=0): + shape = _maybe_get_const(shape, "is") + if not _is_value(shape): + shape = g.op("Constant", value_t=torch.LongTensor(shape)) + if g.opset <= 13: + if allowzero == 1: + _onnx_opset_unsupported( + "Reshape with allowzero=1", GLOBALS.export_onnx_opset_version, 14, input + ) + return g.op("Reshape", input, shape) + else: + return g.op("Reshape", input, shape, allowzero_i=allowzero) + + +def _batchnorm_helper( + g: jit_utils.GraphContext, input, weight, bias, running_mean, running_var +): + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import _var_mean + + batch_size = _get_tensor_dim_size(input, 0) + channel_size = _get_tensor_dim_size(input, 1) + + if weight is None or _is_none(weight): + if channel_size is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of batch_norm for unknown channel size.", + input, + ) + weight_value = torch.tensor( + [1.0] * channel_size, + dtype=_type_utils.JitScalarType.from_value(input).dtype(), + ) + weight = g.op("Constant", value_t=weight_value) + if bias is None or _is_none(bias): + if channel_size is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of batch_norm for unknown channel size.", + input, + ) + bias_value = torch.tensor( + [0.0] * channel_size, + dtype=_type_utils.JitScalarType.from_value(input).dtype(), + ) + bias = g.op("Constant", value_t=bias_value) + # If track_running_stats is set to False batch statistics are instead used during evaluation time + if ( + running_mean is None + or _is_none(running_mean) + or running_var is None + or _is_none(running_var) + ): + if batch_size is None or channel_size is None: + raise AssertionError("batch_size and channel_size must be non-None") + reshape_in = _reshape_helper( + g, + input, + g.op( + "Constant", + value_t=torch.tensor([batch_size, channel_size, -1], dtype=torch.int64), + ), + ) + trans_in = g.op("Transpose", reshape_in, perm_i=[0, 2, 1]) + running_var, running_mean = _var_mean( + g, + trans_in, + g.op("Constant", value_t=torch.tensor([0, 1], dtype=torch.int64)), + False, + False, + ) + return weight, bias, running_mean, running_var + + +def _avgpool_helper( + tuple_fn: Callable[[Any], Sequence[int]], + padding: int | Sequence[int], + kernel_size, + stride, + divisor_override, + name, +) -> tuple[int, ...]: + if divisor_override and divisor_override.node().kind() != "prim::Constant": + _unimplemented(name, "divisor_override") + return tuple(tuple_fn(padding)) + + +def check_training_mode(op_train_mode: int, op_name: str) -> None: + """Warns the user if the model's training mode and the export mode do not agree.""" + if GLOBALS.training_mode == _C_onnx.TrainingMode.PRESERVE: + return + + if op_train_mode: + op_mode_enum = _C_onnx.TrainingMode.TRAINING + else: + op_mode_enum = _C_onnx.TrainingMode.EVAL + if op_mode_enum == GLOBALS.training_mode: + # The modes agree. Do nothing + return + + op_mode_text = f"train={bool(op_train_mode)}" + # Setting the model mode could result in op_mode != GLOBALS.training_mode + # if the model is a FuncModule. In this case we warn the user of + # the state and export depending on op_mode + # This is to support use-cases of fixing certain layer weights + # in training. + warnings.warn( + f"ONNX export mode is set to {GLOBALS.training_mode}, but operator '{op_name}' " + f"is set to {op_mode_text}. Exporting with {op_mode_text}.", + stacklevel=2, + ) + + +def _flatten_helper(g: jit_utils.GraphContext, input, start_dim, end_dim, dim): + input_size = g.op("Shape", input) + slice1 = _slice_helper(g, input_size, axes=[0], starts=[0], ends=[start_dim]) + slices = [slice1, g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long))] + if end_dim < dim - 1: + slice3 = _slice_helper( + g, input_size, axes=[0], starts=[end_dim + 1], ends=[dim] + ) + slices = [ + slice1, + g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), + slice3, + ] + + final_shape = g.op("Concat", *slices, axis_i=0) + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import ( + _reshape_from_tensor, + ) + + return _reshape_from_tensor(g, input, final_shape) + + +def _is_split_static(split_size_or_sizes, _outputs): + if _outputs is None: + return False + if ( + _is_value(split_size_or_sizes) + and split_size_or_sizes.node().kind() != "onnx::Constant" + ): + return False + return True + + +def _optional_input_placeholder_tensor(g): + n = g.op("prim::Constant") + n.setType(_C.OptionalType.ofTensor()) + return n + + +def _handle_reduce_dim_none(g: jit_utils.GraphContext, self, op_name): + rank = _get_tensor_rank(self) + if rank is not None and any( + _get_tensor_dim_size(self, i) == 0 for i in range(rank) + ): + # If input tensor is empty, according to ONNX ReduceSum definition, + # set keepdims=1 so that the resulted tensor has the same rank as the input. + return g.op(op_name, self, keepdims_i=1) + return g.op(op_name, self, keepdims_i=0) + + +def dequantize_helper( + g: jit_utils.GraphContext, + qtensor: _C.Value, + qdtype: _C_onnx.TensorProtoDataType | None = None, +) -> tuple[_C.Value, _C.Value, _C.Value, _C.Value | None]: + """Appends to graph `g` ONNX nodes that dequantizes `qtensor` into `tensor`. + + Args: + g: Graph, the ONNX IR graph that is under construction. + qtensor: torch._C.Value, either a tuple of (quantized_tensor, scale, zero_point) + for per tensor quantization, or + (quantized_tensor, scale, zero_point, axis) for per channel quantization, + representing the quantized tensor. + qdtype: torch.onnx.TensorProtoDataType default None, if not None, represents the + data type of quantized tensor. It must be either + torch.onnx.TensorProtoDataType.UINT8 or torch.onnx.TensorProtoDataType.INT8. + """ + unpacked_qtensors = _unpack_quantized_tensor(qtensor) + tensor, scale, zero_point = unpacked_qtensors[:3] + axis = unpacked_qtensors[3] if len(unpacked_qtensors) >= 4 else None + axis_i = _get_const(axis, "i", "axis") + input_qdtype = _type_utils.JitScalarType.from_value(tensor) + if qdtype is None: + if input_qdtype is not None: + qdtype = input_qdtype.onnx_type() + else: + qdtype = _C_onnx.TensorProtoDataType.UINT8 + value = g.op("Cast", tensor, to_i=qdtype) + scale = g.op("Cast", scale, to_i=_C_onnx.TensorProtoDataType.FLOAT) + zero_point = g.op("Cast", zero_point, to_i=qdtype) + + if axis_i is not None and GLOBALS.export_onnx_opset_version < 13: + _onnx_opset_unsupported_detailed( + "DequantizeLinear", + GLOBALS.export_onnx_opset_version, + 13, + "Attribute axis is not supported.", + qtensor, + ) + + return ( + g.op("DequantizeLinear", value, scale, zero_point, axis_i=axis_i), + scale, + zero_point, + axis, + ) + + +def quantize_helper( + g: jit_utils.GraphContext, + tensor: _C.Value, + scale: _C.Value, + zero_point: _C.Value, + axis: _C.Value | None = None, +) -> _C.Value: + """Appends to graph `g` ONNX nodes that quantizes `tensor` based on `scale`, `zero_point` and `axis`. + + Args: + g: Graph, the ONNX IR graph that is under construction. + tensor: torch._C.Value, representing the tensor to be quantized. + scale: torch._C.Value, quantized scale. + zero_point: torch._C.Value, quantized zero point. + axis: Optional[torch._C.Value] default None, if None, represents per tensor quantization. + Otherwise, represents per channel quantization, along given axis. + + Returns: + A TupleConstruct storing information of the quantized tensor. + """ + if ( + axis is not None + and not _is_none(axis) + and GLOBALS.export_onnx_opset_version < 13 + ): + _onnx_opset_unsupported_detailed( + "QuantizeLinear", + GLOBALS.export_onnx_opset_version, + 13, + "Attribute axis is not supported.", + tensor, + ) + + if scale is None: + raise AssertionError("scale must be non-None") + if ( + _type_utils.JitScalarType.from_value(scale, _type_utils.JitScalarType.UNDEFINED) + != _type_utils.JitScalarType.FLOAT + ): + scale = g.op("Cast", scale, to_i=_C_onnx.TensorProtoDataType.FLOAT) + + if zero_point is None: + raise AssertionError("zero_point must be non-None") + if _type_utils.JitScalarType.from_value( + zero_point, _type_utils.JitScalarType.UNDEFINED + ) not in { + _type_utils.JitScalarType.UINT8, + _type_utils.JitScalarType.INT8, + }: + zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.UINT8) + output = g.op( + "QuantizeLinear", + tensor, + scale, + zero_point, + axis_i=_get_const(axis, "i", "axis"), + ) + args = [output, scale, zero_point] + if axis is not None and not _is_none(axis): + args.append(axis) + return g.op("prim::TupleConstruct", *args) + + +def requantize_bias_helper( + g: jit_utils.GraphContext, bias, input_scale, weight_scale, axis=None +): + """In PyTorch, bias is float and is quantized to int32 implicitly inside the quantized ATen op kernel. + In ONNX we need to make the quantization explicit because operators expect all of their inputs to be quantized. + Since int32 is not a supported output type by ONNX operator `QuantizeLinear`, quantization is exported using + regular operators. + """ + bias_scale = g.op("Mul", weight_scale, input_scale) + bias_scale_shape = g.op("Shape", bias_scale) + bias_zero_point = g.op( + "ConstantOfShape", bias_scale_shape, value_t=torch.tensor([0], dtype=torch.int) + ) + q_bias = g.op( + "Cast", g.op("Div", bias, bias_scale), to_i=_C_onnx.TensorProtoDataType.INT32 + ) + axis_args = [] + if axis is not None and not _is_none(axis): + axis_args.append(axis) + return g.op("prim::TupleConstruct", q_bias, bias_scale, bias_zero_point, *axis_args) + + +def args_have_same_dtype(args): + if not args: + raise AssertionError("args must be non-empty") + base_dtype = _type_utils.JitScalarType.from_value(args[0]) + has_same_dtype = all( + _type_utils.JitScalarType.from_value(elem) == base_dtype for elem in args + ) + return has_same_dtype + + +def _op_with_optional_float_cast(g: jit_utils.GraphContext, op_name, *args, **kwargs): + """Some PyTorch operators (e.g., Clip/Min/ReLU/Pad) are super set of ONNX in terms of data types. + This function maximizes the exportability of PyTorch-ONNX by allowing ONNX-unsupported PyTorch + operator data type. For example, `Cast(Clip(Cast(INPUT)))` can be used to mimic + `Clip(INPUT)` (opset version < 12). + + Args: + g (torch._C.Graph): graph to write the ONNX representation into. + op_name (str): operator name in ONNX. + *args (tuple): operands to the operator. + **kwargs (dict): attributes to the operator along with "opset_before" (optional, None by default) + indicating the smallest opset version to trigger such casting behavior and "target_float_t" + (optional, torch.onnx.JitScalarType.FLOAT by default) indicating the data type of internal operator. + + Returns: + Optional[torch._C.Value, Tuple[torch._C.Value, ...]]: output(s) of the operator. + """ + opset_before = kwargs.pop("opset_before", None) + target_float_t = kwargs.pop("target_float_t", _type_utils.JitScalarType.FLOAT) + + inputs = list(args) + dtype_0 = _type_utils.JitScalarType.from_value(inputs[0]) + + require_cast = not _is_fp(inputs[0]) and ( + opset_before is None or GLOBALS.export_onnx_opset_version < opset_before + ) + + if require_cast: + for input in inputs: + if input.isCompleteTensor(): + input_scalar_type = _type_utils.JitScalarType.from_value(input) + if input_scalar_type != dtype_0: + raise errors.SymbolicValueError( + f"Inputs of {op_name} must have same dtype." + f"Got {dtype_0.scalar_name()} and {input_scalar_type.scalar_name()}", + input, + ) + for i, input in enumerate(inputs): + if input.isCompleteTensor() and not _is_fp(input): + inputs[i] = g.op( + "Cast", + input, + to_i=target_float_t.onnx_type(), + ) + + self = g.op(op_name, *inputs, **kwargs) + + if require_cast: + self = g.op("Cast", self, to_i=dtype_0.onnx_type()) + + return self + + +def _maybe_cast_reduce_op_input(g: jit_utils.GraphContext, self): + scalar_type = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.UNDEFINED + ) + if scalar_type != _type_utils.JitScalarType.UNDEFINED: + # This check only covers traced modules where dtype is present + # pytorch reduce-ops cast all other integral types to int64 + if not _is_fp(self) and scalar_type != _type_utils.JitScalarType.INT64: + self = g.op("Cast", self, to_i=_C_onnx.TensorProtoDataType.INT64) + return self + + +def _apply_params(*args, **kwargs): + """Returns a decorator that calls the decorated (higher-order) function with the given parameters.""" + + def _apply(fn): + return fn(*args, **kwargs) + + return _apply + + +def _reduce_op_symbolic_helper(onnx_op_name, allow_multi_dim_support=True): + def symbolic(g, self, dim=None, keepdim=None): + self = _maybe_cast_reduce_op_input(g, self) + if dim is None or dim == (): + # Dim can be 0, which will cause (not dim) == True. So we don't want to do + # (not dim) + # all-reduce path + return _handle_reduce_dim_none(g, self, onnx_op_name) + else: + # dim-reduce path + keepdim = _get_const(keepdim, "i", "keepdim") + if g.opset < 18: + desc = "is" if allow_multi_dim_support else "i" + dim = _get_const(dim, desc, "dim") + dim_list = dim if allow_multi_dim_support else [dim] + return g.op(onnx_op_name, self, axes_i=dim_list, keepdims_i=keepdim) + else: + if _is_value(dim): + axes = dim + else: + if allow_multi_dim_support: + axes = g.op( + "Constant", value_t=torch.tensor(dim, dtype=torch.long) + ) + else: + axes = g.op( + "Constant", value_t=torch.tensor([dim], dtype=torch.long) + ) + return g.op(onnx_op_name, self, axes, keepdims_i=keepdim) + + return symbolic + + +def _overload_by_arg_count(fn): + @functools.wraps(fn) + def wrapper(g, *args): + overloads = fn(g, *args) + for overload in overloads: + arg_descriptors = overload._arg_descriptors + if len(arg_descriptors) == len(args): + return overload(g, *args) + return _unimplemented(f"aten::{fn.__name__}", f"with {len(args)} arguments") + + return wrapper + + +def _reduce_with_dtype_helper( + onnx_op: str, name: str, allow_multi_dim_support: bool = True +): + symbolic = _reduce_op_symbolic_helper( + onnx_op, allow_multi_dim_support=allow_multi_dim_support + ) + + @_overload_by_arg_count + def reduce(g, *args, **kwargs): + @quantized_args(True) + @parse_args("v", "none") + def reduce_nodim(g, self, dtype): + dtype_onnx = None + if dtype.node().kind() == "onnx::Constant": + dtype = _get_const(dtype, "i", "dtype") + dtype_onnx = _type_utils.JitScalarType(dtype).onnx_type() + self = g.op("Cast", self, to_i=dtype_onnx) + elif dtype.node().kind() != "prim::Constant": + return _unimplemented(name, "dtype", dtype) + result = symbolic(g, self) + if dtype_onnx is not None: + result_dtype_onnx = _type_utils.JitScalarType.from_value( + result + ).onnx_type() + if result_dtype_onnx != dtype_onnx: + result = g.op("Cast", result, to_i=dtype_onnx) + return result + + dim_desc = "is" if allow_multi_dim_support else "i" + + @quantized_args(True) + @parse_args("v", dim_desc, "i", "none") # type: ignore[arg-type] + def reduce_dim(g, self, dim, keepdim, dtype): + dtype_onnx = None + if dtype.node().kind() == "onnx::Constant": + dtype = _get_const(dtype, "i", "dtype") + dtype_onnx = _type_utils.JitScalarType(dtype).onnx_type() + self = g.op("Cast", self, to_i=dtype_onnx) + elif dtype.node().kind() != "prim::Constant": + return _unimplemented(name, "dtype", dtype) + result = symbolic(g, self, dim, keepdim) + if dtype_onnx is not None: + result_dtype_onnx = _type_utils.JitScalarType.from_value( + result + ).onnx_type() + if result_dtype_onnx != dtype_onnx: + result = g.op("Cast", result, to_i=dtype_onnx) + return result + + return reduce_nodim, reduce_dim + + return reduce + + +def _max_helper(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None): + # torch.max(input) + if dim_or_y is None and keepdim is None: + return g.op("ReduceMax", self, keepdims_i=0) + # torch.max(input, other) + if keepdim is None: + return _op_with_optional_float_cast(g, "Max", self, dim_or_y, opset_before=12) + # torch.max(input, dim, keepdim) + else: + keepdim = _get_const(keepdim, "i", "keepdim") + dim = _get_const(dim_or_y, "i", "dim") + if g.opset < 18: + max = g.op("ReduceMax", self, axes_i=[dim], keepdims_i=keepdim) + else: + axes = g.op("Constant", value_t=torch.tensor([dim], dtype=torch.long)) + max = g.op("ReduceMax", self, axes, keepdims_i=keepdim) + indices = g.op("ArgMax", self, axis_i=dim, keepdims_i=keepdim) + return max, indices + + +def _min_helper(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None): + # torch.min(input) + if dim_or_y is None and keepdim is None: + return g.op("ReduceMin", self, keepdims_i=0) + # torch.min(input, other) + if keepdim is None: + return _op_with_optional_float_cast(g, "Min", self, dim_or_y, opset_before=12) + # torch.min(input, dim, keepdim) + else: + keepdim = _get_const(keepdim, "i", "keepdim") + dim = _get_const(dim_or_y, "i", "dim") + if g.opset < 18: + min = g.op("ReduceMin", self, axes_i=[dim], keepdims_i=keepdim) + else: + axes = g.op("Constant", value_t=torch.tensor([dim], dtype=torch.long)) + min = g.op("ReduceMin", self, axes, keepdims_i=keepdim) + indices = g.op("ArgMin", self, axis_i=dim, keepdims_i=keepdim) + return min, indices + + +def _numel_helper(g: jit_utils.GraphContext, self): + shape = g.op("Shape", self) + return g.op("ReduceProd", shape, keepdims_i=0) + + +@parse_args("v", "is", "i", "i") +def _var_mean_helper(g: jit_utils.GraphContext, input, dim, correction, keepdim): + if g.opset < 18: + if dim is None: + mean = g.op("ReduceMean", input, keepdims_i=0) + t_mean = mean + num_elements = _numel_helper(g, input) + else: + mean = g.op("ReduceMean", input, axes_i=dim, keepdims_i=keepdim) + t_mean = g.op("ReduceMean", input, axes_i=dim, keepdims_i=1) + redudced_dims = g.op("Shape", input) + # dim could contain one or multiple dimensions + redudced_dims = g.op( + "Gather", + redudced_dims, + g.op("Constant", value_t=torch.tensor(dim)), + axis_i=0, + ) + num_elements = g.op("ReduceProd", redudced_dims, keepdims_i=0) + sub_v = g.op("Sub", input, t_mean) + sqr_sub = g.op("Mul", sub_v, sub_v) + keepdim_mean = 0 if dim is None else keepdim + var = g.op("ReduceMean", sqr_sub, axes_i=dim, keepdims_i=keepdim_mean) + # Correct bias in calculating variance, by dividing it over (N - correction) instead on N + if correction is None: + correction = 1 + if correction != 0: + num_elements = g.op( + "Cast", num_elements, to_i=_C_onnx.TensorProtoDataType.FLOAT + ) + one = g.op("Constant", value_t=torch.tensor(correction, dtype=torch.float)) + mul = g.op("Mul", var, num_elements) + var = g.op("Div", mul, g.op("Sub", num_elements, one)) + return var, mean + else: + axes = None + if dim is None: + mean = g.op("ReduceMean", input, keepdims_i=0) + t_mean = mean + num_elements = _numel_helper(g, input) + else: + axes = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long)) + mean = g.op("ReduceMean", input, axes, keepdims_i=keepdim) + t_mean = g.op("ReduceMean", input, axes, keepdims_i=1) + redudced_dims = g.op("Shape", input) + # dim could contain one or multiple dimensions + redudced_dims = g.op( + "Gather", + redudced_dims, + g.op("Constant", value_t=torch.tensor(dim)), + axis_i=0, + ) + num_elements = g.op("ReduceProd", redudced_dims, keepdims_i=0) + sub_v = g.op("Sub", input, t_mean) + sqr_sub = g.op("Mul", sub_v, sub_v) + keepdim_mean = 0 if dim is None else keepdim + if axes is None: + var = g.op("ReduceMean", sqr_sub, keepdims_i=keepdim_mean) + else: + var = g.op("ReduceMean", sqr_sub, axes, keepdims_i=keepdim_mean) + # Correct bias in calculating variance, by dividing it over (N - correction) instead on N + if correction is None: + correction = 1 + if correction != 0: + num_elements = g.op( + "Cast", num_elements, to_i=_C_onnx.TensorProtoDataType.FLOAT + ) + one = g.op("Constant", value_t=torch.tensor(correction, dtype=torch.float)) + mul = g.op("Mul", var, num_elements) + var = g.op("Div", mul, g.op("Sub", num_elements, one)) + return var, mean + + +def _embedding_bag_helper( + g: jit_utils.GraphContext, + embedding_matrix, + indices, + offsets, + scale_grad_by_freq, + mode, + sparse, + per_sample_weights, + include_last_offset, + padding_idx, +): + if scale_grad_by_freq and GLOBALS.export_training: + return _onnx_unsupported( + "embedding_bag with scale_grad_by_freq for training mode" + ) + if padding_idx is not None and padding_idx >= 0: + raise RuntimeError("embedding_bag with padding_idx") + + loop_condition = g.op("Constant", value_t=torch.tensor(1)) + loop_condition = g.op("Cast", loop_condition, to_i=_C_onnx.TensorProtoDataType.BOOL) + zero = g.op("Constant", value_t=torch.tensor([0])) + + indices_len = _unsqueeze_helper( + g, + _size_helper(g, indices, g.op("Constant", value_t=torch.tensor(0))), + [0], + ) + if not include_last_offset: + offsets = [offsets, indices_len] + offsets = g.op("Concat", *offsets, axis_i=0) + + # Offsets holds the starting index position of each bag. So we create a list of the indices slices (determined by + # offsets) and gather those indices in indices_row. Then we use this subset of indices to gather from embeddings. + # The embeddings output is a loop scan output, so we can avoid creating a sequence and inserting elements in. + offsets_starts = _slice_helper( + g, offsets, axes=[0], starts=[0], ends=[sys.maxsize], steps=[1] + ) + offsets_ends = _slice_helper( + g, offsets, axes=[0], starts=[1], ends=[sys.maxsize], steps=[1] + ) + + loop_len = _size_helper(g, offsets_ends, g.op("Constant", value_t=torch.tensor(0))) + + loop, (loop_context,), _ = jit_utils.add_op_with_blocks( + g, "Loop", loop_len, loop_condition, n_blocks=1 + ) + loop_block = loop_context.block + + # FIXME(justinchuby): We need to handle what happens when we call b.op on a node return + block_input_iter = utils._add_input_to_block(loop_block) + utils._add_input_to_block(loop_block) + + indices_start = loop_context.op( + "Gather", offsets_starts, block_input_iter, axis_i=0 + ) + indices_end = loop_context.op("Gather", offsets_ends, block_input_iter, axis_i=0) + indices_start = _unsqueeze_helper(loop_context, indices_start, [0]) + indices_end = _unsqueeze_helper(loop_context, indices_end, [0]) + + indices_row = loop_context.op("Slice", indices, indices_start, indices_end, zero) + embeddings = loop_context.op("Gather", embedding_matrix, indices_row, axis_i=0) + if not _is_none(per_sample_weights): + per_sample_weights_row = loop_context.op( + "Slice", per_sample_weights, indices_start, indices_end, zero + ) + per_sample_weights_row = _unsqueeze_helper( + loop_context, per_sample_weights_row, [1] + ) + embeddings = loop_context.op("Mul", embeddings, per_sample_weights_row) + if mode == 0: + embeddings = _reducesum_helper( + loop_context, embeddings, axes_i=[0], keepdims_i=0 + ) + elif mode == 1: + if loop_context.opset < 18: + embeddings = loop_context.op( + "ReduceMean", embeddings, axes_i=[0], keepdims_i=0 + ) + else: + axes = loop_context.op( + "Constant", value_t=torch.tensor([0], dtype=torch.long) + ) + embeddings = loop_context.op("ReduceMean", embeddings, axes, keepdims_i=0) + else: + if loop_context.opset < 18: + embeddings = loop_context.op( + "ReduceMax", embeddings, axes_i=[0], keepdims_i=0 + ) + else: + axes = loop_context.op( + "Constant", value_t=torch.tensor([0], dtype=torch.long) + ) + embeddings = loop_context.op("ReduceMax", embeddings, axes, keepdims_i=0) + + cond_out = loop_context.op( + "Cast", loop_condition, to_i=_C_onnx.TensorProtoDataType.BOOL + ) + utils._add_output_to_block(loop_block, cond_out) + utils._add_output_to_block(loop_block, embeddings) + + # aten::embedding_bag returns a tuple of 4 elements: output, offset2bag, bag_size, max_indices. + # But the last three outputs are not used in torch.nn.EmbeddingBag or torch.nn.functional.embedding_bag. + return loop.node().output(), None, None, None + + +def _linalg_vector_norm_helper( + g: jit_utils.GraphContext, + self: torch._C.Value, + ord: float, + dim: Sequence[int] | None, + keepdim: bool, + dtype: torch._C.Value, +): + axes = None + # Conditions based on https://pytorch.org/docs/stable/generated/torch.linalg.vector_norm.html + if _is_none(dim): + self = _reshape_helper(g, self, [-1]) + keepdim = False + elif g.opset >= 18: + axes = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long)) + + if ord == math.inf: + if g.opset < 18: + result = g.op( + "ReduceMax", g.op("Abs", self), axes_i=dim, keepdims_i=keepdim + ) + else: + if axes is None: + result = g.op("ReduceMax", g.op("Abs", self), keepdims_i=keepdim) + else: + result = g.op("ReduceMax", g.op("Abs", self), axes, keepdims_i=keepdim) + elif ord == -math.inf: + if g.opset < 18: + result = g.op( + "ReduceMin", g.op("Abs", self), axes_i=dim, keepdims_i=keepdim + ) + else: + if axes is None: + result = g.op("ReduceMin", g.op("Abs", self), keepdims_i=keepdim) + else: + result = g.op("ReduceMin", g.op("Abs", self), axes, keepdims_i=keepdim) + elif ord == 0: + if g.opset < 11: + return _onnx_opset_unsupported_detailed( + "linalg_vector_norm", 9, 11, "ord=0 not supported", self + ) + else: + if dim is None: + self = _reshape_helper( + g, + self, + g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)), + ) + keepdim = False + + cond_op = g.op( + "Not", + g.op("Equal", self, g.op("Constant", value_t=torch.LongTensor([0]))), + ) + cond_op = g.op( + "Cast", + cond_op, + to_i=_type_utils.JitScalarType.from_value(self).onnx_type(), + ) + return _reducesum_helper(g, cond_op, axes_i=dim, keepdims_i=keepdim) + elif ord == 1: + if g.opset < 18: + result = _reduce_op_symbolic_helper("ReduceL1")( + g, self, dim=dim, keepdim=keepdim + ) + else: + if axes is None: + result = _reduce_op_symbolic_helper("ReduceL1")( + g, self, keepdim=keepdim + ) + else: + result = _reduce_op_symbolic_helper("ReduceL1")( + g, self, axes, keepdim=keepdim + ) + elif ord == 2: + if g.opset < 18: + result = _reduce_op_symbolic_helper("ReduceL2")( + g, self, dim=dim, keepdim=keepdim + ) + else: + if axes is None: + result = _reduce_op_symbolic_helper("ReduceL2")( + g, self, keepdim=keepdim + ) + else: + result = _reduce_op_symbolic_helper("ReduceL2")( + g, self, axes, keepdim=keepdim + ) + else: + ord_op = g.op("Constant", value_t=torch.tensor(ord, dtype=torch.float32)) + result = _reducesum_helper( + g, g.op("Pow", g.op("Abs", self), ord_op), axes_i=dim, keepdims_i=keepdim + ) + result = g.op( + "Pow", + result, + g.op( + "Div", + g.op("Constant", value_t=torch.tensor(1, dtype=torch.float32)), + ord_op, + ), + ) + + if not _is_none(dtype): + dtype = _get_const(dtype, "i", "dtype") + result = g.op("Cast", result, to_i=_type_utils.JitScalarType(dtype).onnx_type()) # type: ignore[arg-type] + return result + + +# Deprecated. Internally use _type_utils.ScalarType +# TODO: remove these once we support Type's in the JIT IR and we can once again +# use the unified toType operator +cast_pytorch_to_onnx = { + "Byte": _C_onnx.TensorProtoDataType.UINT8, + "Char": _C_onnx.TensorProtoDataType.INT8, + "Double": _C_onnx.TensorProtoDataType.DOUBLE, + "Float": _C_onnx.TensorProtoDataType.FLOAT, + "Half": _C_onnx.TensorProtoDataType.FLOAT16, + "Int": _C_onnx.TensorProtoDataType.INT32, + "Long": _C_onnx.TensorProtoDataType.INT64, + "Short": _C_onnx.TensorProtoDataType.INT16, + "Bool": _C_onnx.TensorProtoDataType.BOOL, + "ComplexFloat": _C_onnx.TensorProtoDataType.COMPLEX64, + "ComplexDouble": _C_onnx.TensorProtoDataType.COMPLEX128, + "BFloat16": _C_onnx.TensorProtoDataType.BFLOAT16, + "Undefined": _C_onnx.TensorProtoDataType.UNDEFINED, +} + +# Deprecated. Internally use _type_utils.ScalarType +scalar_name_to_pytorch = { + "uint8_t": "Byte", + "int8_t": "Char", + "double": "Double", + "float": "Float", + "half": "Half", + "int": "Int", + "int64_t": "Long", + "int16_t": "Short", + "bool": "Bool", + "complex64": "ComplexFloat", + "complex128": "ComplexDouble", + "qint8": "QInt8", + "quint8": "QUInt8", + "qint32": "QInt32", + "bfloat16": "BFloat16", +} + + +# Deprecated. Internally use _type_utils.ScalarType +# This indicates each scalar type's corresponding +# torch type. Related source: +# https://github.com/pytorch/pytorch/blob/344defc9733a45fee8d0c4d3f5530f631e823196/c10/core/ScalarType.h +scalar_type_to_pytorch_type = [ + torch.uint8, # 0 + torch.int8, # 1 + torch.short, # 2 + torch.int, # 3 + torch.int64, # 4 + torch.half, # 5 + torch.float, # 6 + torch.double, # 7 + torch.complex32, # 8 + torch.complex64, # 9 + torch.complex128, # 10 + torch.bool, # 11 + torch.qint8, # 12 + torch.quint8, # 13 + torch.qint32, # 14 + torch.bfloat16, # 15 +] + +# Deprecated. Internally use _type_utils.ScalarType +# source of truth is +# https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/tensor_dtypes.cpp +pytorch_name_to_type = { + "Byte": torch.uint8, + "Char": torch.int8, + "Double": torch.double, + "Float": torch.float, + "Half": torch.half, + "Int": torch.int, + "Long": torch.int64, + "Short": torch.short, + "Bool": torch.bool, + "ComplexFloat": torch.complex64, + "ComplexDouble": torch.complex128, + "QInt8": torch.qint8, + "QUInt8": torch.quint8, + "QInt32": torch.qint32, + "BFloat16": torch.bfloat16, +} + + +# Deprecated. Internally use _type_utils.ScalarType +scalar_type_to_onnx = [ + cast_pytorch_to_onnx["Byte"], # 0 + cast_pytorch_to_onnx["Char"], # 1 + cast_pytorch_to_onnx["Short"], # 2 + cast_pytorch_to_onnx["Int"], # 3 + cast_pytorch_to_onnx["Long"], # 4 + cast_pytorch_to_onnx["Half"], # 5 + cast_pytorch_to_onnx["Float"], # 6 + cast_pytorch_to_onnx["Double"], # 7 + cast_pytorch_to_onnx["Undefined"], # 8 + cast_pytorch_to_onnx["ComplexFloat"], # 9 + cast_pytorch_to_onnx["ComplexDouble"], # 10 + cast_pytorch_to_onnx["Bool"], # 11 + cast_pytorch_to_onnx["Char"], # 12 + cast_pytorch_to_onnx["Byte"], # 13 + cast_pytorch_to_onnx["Int"], # 14 + cast_pytorch_to_onnx["BFloat16"], # 15 +] + +# Global set to store the list of quantized operators in the network. +# This is currently only used in the conversion of quantized ops from PT -> C2 via ONNX. +_quantized_ops: set[int] = set() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset10.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset10.py new file mode 100644 index 0000000000000000000000000000000000000000..d63968e4f278a8a9cdf2a1d8590e6b8a5b0140e5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset10.py @@ -0,0 +1,1191 @@ +# mypy: allow-untyped-defs +# mypy: disable-error-code=arg-type +from __future__ import annotations + +import functools +import sys +import warnings +from typing import TYPE_CHECKING + +import torch +import torch._C._onnx as _C_onnx +from torch import _C +from torch.onnx import _constants, errors +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, + symbolic_opset9 as opset9, +) +from torch.onnx._internal.torchscript_exporter._globals import GLOBALS + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in README.md + +# This file exports ONNX ops for opset 10 +# Opset 10 is supported by ONNX release 1.5.0 +# release on 04/24/19 + + +__all__ = [ + "dequantize", + "div", + "embedding_bag", + "fake_quantize_per_tensor_affine", + "flip", + "fmod", + "isfinite", + "isinf", + "nan_to_num", + "quantize_per_tensor", + "quantized_add_relu", + "quantized_add", + "quantized_cat", + "quantized_conv1d_relu", + "quantized_conv2d_relu", + "quantized_conv3d_relu", + "quantized_conv1d", + "quantized_conv2d", + "quantized_conv3d", + "quantized_conv_transpose1d", + "quantized_conv_transpose2d", + "quantized_conv_transpose3d", + "quantized_group_norm", + "quantized_hardswish", + "quantized_instance_norm", + "quantized_layer_norm", + "quantized_leaky_relu", + "quantized_linear", + "quantized_linear_relu", + "quantized_mul", + "quantized_sigmoid", + "slice", + "sort", + "topk", +] + + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=10) + + +@_onnx_symbolic("aten::div") +def div(g: jit_utils.GraphContext, self, other, *args): + if len(args) == 0: + return opset9.true_divide(g, self, other) + else: + return _div_rounding_mode(g, self, other, *args) + + +@symbolic_helper.parse_args("v", "v", "s") +def _div_rounding_mode(g: jit_utils.GraphContext, self, other, rounding_mode): + if rounding_mode == "floor": + return _floor_divide(g, self, other) + else: + return opset9._div_rounding_mode(g, self, other, rounding_mode) + + +@_onnx_symbolic("aten::_floor_divide") +def _floor_divide(g: jit_utils.GraphContext, self, other): + if symbolic_helper._is_fp(self) or symbolic_helper._is_fp(other): + out = opset9.true_divide(g, self, other) + return g.op("Floor", out) + else: + # Integer division does truncation rounding + div = g.op("Div", self, other) + # Division is negative if: self < 0 != other < 0 + zero = g.op("Constant", value_t=torch.tensor(0, dtype=torch.int64)) + negative = g.op("Xor", g.op("Less", self, zero), g.op("Less", other, zero)) + + # For negative numbers with self % other != 0, subtract 1 to round down instead of up + mod = g.op("Mod", self, other, fmod_i=0) + fixup_mask = g.op("And", negative, g.op("Not", g.op("Equal", mod, zero))) + + one = g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64)) + fixup = g.op("Sub", div, one) + return g.op("Where", fixup_mask, fixup, div) + + +@_onnx_symbolic("aten::sort") +@symbolic_helper.parse_args("v", "i", "i", "none") +def sort(g: jit_utils.GraphContext, self, dim, descending, out=None): + return symbolic_helper._sort_helper(g, self, dim, descending=descending, out=out) + + +@_onnx_symbolic("aten::topk") +@symbolic_helper.parse_args("v", "v", "i", "i", "i", "none") +def topk(g: jit_utils.GraphContext, self, k, dim, largest, sorted, out=None): + return symbolic_helper._topk_helper( + g, self, k, dim, largest=largest, sorted=sorted, out=out + ) + + +def _aten_max_pool_onnx( + g: jit_utils.GraphContext, + self: _C.Value, + kernel_shape: Sequence[int], + strides: Sequence[int], + pads: Sequence[int], + dilations: Sequence[int], + ceil_mode: bool, + unbatched_rank: int, +) -> _C.Value: + self_rank = g.op("Size", g.op("Shape", self)) + if self_rank == unbatched_rank: # C,H,W -> N,C,H,W and N=1 + self = g.op( + "Unsqueeze", + self, + g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64)), + ) + + pool_result, _ = g.op( + "MaxPool", + self, + outputs=2, + ceil_mode_i=ceil_mode, + dilations_i=dilations, + kernel_shape_i=kernel_shape, + pads_i=pads, + strides_i=strides, + ) + + if self_rank == unbatched_rank: + pool_result = g.op( + "Squeeze", + pool_result, + g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64)), + ) + + return pool_result + + +# For MaxPool +def _adjust_attributes_of_max_pool( + expand_size: int, + kernel_size: Sequence[int] | int, + stride: Sequence[int] | int, + padding: Sequence[int] | int, + dilation: Sequence[int] | int, +) -> tuple[Sequence[int], Sequence[int], Sequence[int], Sequence[int]]: + """Adjust attributes of avg_pool to match ONNX specification.""" + + if isinstance(dilation, int): + dilation = [dilation] * expand_size + + if isinstance(kernel_size, int): + kernel_shape = [kernel_size] * expand_size + else: + kernel_shape = kernel_size # type: ignore[assignment] + + if isinstance(padding, int): + pads = [padding] * expand_size * 2 # type: ignore[operator, assignment] + elif len(padding) == 1: + pads = padding * expand_size * 2 # type: ignore[operator, assignment] + elif len(padding) == 2: + # 2D padding + pads = padding * 2 # type: ignore[operator, assignment] + elif len(padding) == 3: + # 3D padding + pads = padding * 2 # type: ignore[operator, assignment] + else: + # When padding is already done for all dimensions, + # we don't need to double it + # eg: (1, 1, 1, 1, 1, 1) + pads = padding # type: ignore[assignment] + + if isinstance(stride, int): + strides = [stride] * expand_size + elif not stride: + strides = kernel_shape + else: + strides = stride # type: ignore[assignment] + + # pyrefly: ignore [bad-return] + return (kernel_shape, strides, pads, dilation) + + +def _aten_max_pool_with_indices_onnx( + g: jit_utils.GraphContext, + self: _C.Value, + kernel_shape: Sequence[int], + strides: Sequence[int], + pads: Sequence[int], + dilations: Sequence[int], + ceil_mode: bool, + unbatched_rank: int, + n_dims_one: Sequence[int], + n_dims_zero: Sequence[int], + n_dims_axes: Sequence[int], +) -> tuple[_C.Value, Sequence[int]]: + self_rank = g.op("Size", g.op("Shape", self)) + if self_rank == unbatched_rank: # C,H,W -> N,C,H,W and N=1 + self = g.op( + "Unsqueeze", + self, + g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64)), + ) + + pool_result, indices = g.op( + "MaxPool", + self, + outputs=2, + ceil_mode_i=ceil_mode, + dilations_i=dilations, + kernel_shape_i=kernel_shape, + pads_i=pads, + strides_i=strides, + ) + _, flatten_indices = g.op( + "MaxPool", + self, + outputs=2, + dilations_i=dilations, + kernel_shape_i=n_dims_one, + strides_i=n_dims_one, + ) + + ends = g.op("Constant", value_t=torch.tensor(n_dims_one)) + starts = g.op("Constant", value_t=torch.tensor(n_dims_zero)) + axes = g.op("Constant", value_t=torch.tensor(n_dims_axes)) + + delta = g.op("Slice", flatten_indices, starts, ends, axes) + indices = g.op("Sub", indices, delta) + + if self_rank == unbatched_rank: + pool_result = g.op( + "Squeeze", pool_result, value_t=torch.tensor([0], dtype=torch.int64) + ) + indices = g.op("Squeeze", indices, value_t=torch.tensor([0], dtype=torch.int64)) + + return (pool_result, indices) + + +@_onnx_symbolic( + "aten::max_pool1d", + decorate=[symbolic_helper._apply_params("max_pool1d", 1, return_indices=False)], +) +@_onnx_symbolic( + "aten::max_pool2d", + decorate=[symbolic_helper._apply_params("max_pool2d", 2, return_indices=False)], +) +@_onnx_symbolic( + "aten::max_pool3d", + decorate=[symbolic_helper._apply_params("max_pool3d", 3, return_indices=False)], +) +@_onnx_symbolic( + "aten::max_pool1d_with_indices", + decorate=[ + symbolic_helper._apply_params( + "max_pool1d_with_indices", + 1, + return_indices=True, + ) + ], +) +@_onnx_symbolic( + "aten::max_pool2d_with_indices", + decorate=[ + symbolic_helper._apply_params( + "max_pool2d_with_indices", + 2, + return_indices=True, + ) + ], +) +@_onnx_symbolic( + "aten::max_pool3d_with_indices", + decorate=[ + symbolic_helper._apply_params( + "max_pool3d_with_indices", + 3, + return_indices=True, + ) + ], +) +def _max_pool(name: str, expand_size: int, return_indices: bool): + @symbolic_helper.quantized_args(True, False, False, False, False, False) + @symbolic_helper.parse_args("v", "is", "is", "is", "is", "i") + def symbolic_fn( + g: jit_utils.GraphContext, + input: _C.Value, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: int | Sequence[int], + dilation: Sequence[int], + ceil_mode: bool, + ): + kernel_shape, strides, pads, dilations = _adjust_attributes_of_max_pool( + expand_size, kernel_size, stride, padding, dilation + ) + + if return_indices: + return _aten_max_pool_with_indices_onnx( + g, + input, + kernel_shape, + strides, + pads, + dilations, + ceil_mode, + expand_size + 1, + ([1] * expand_size), + ([0] * expand_size), + ([2 + i for i in range(expand_size)]), + ) + else: + return _aten_max_pool_onnx( + g, + input, + kernel_shape, + strides, + pads, + dilations, + ceil_mode, + expand_size + 1, + ) + + return symbolic_fn + + +# For AvgPool +def _adjust_attributes_of_avg_pool( + expand_size: int, + kernel_size: Sequence[int] | int, + stride: Sequence[int] | int, + padding: Sequence[int] | int, +) -> tuple[Sequence[int], Sequence[int], Sequence[int]]: + """Adjust attributes of avg_pool to match ONNX specification.""" + + if isinstance(kernel_size, int): + kernel_shape = [kernel_size] * expand_size + else: + kernel_shape = kernel_size # type: ignore[assignment] + + if isinstance(padding, int): + pads = [padding] * expand_size * 2 + elif len(padding) == 1: + pads = padding * expand_size * 2 # type: ignore[operator, assignment] + elif len(padding) == 2: + pads = padding * expand_size # type: ignore[operator, assignment] + else: + pads = padding * 2 # type: ignore[operator, assignment] + + if isinstance(stride, int): + strides = [stride] * expand_size + elif not stride: + strides = kernel_shape + else: + strides = stride # type: ignore[assignment] + + # pyrefly: ignore [bad-return] + return (kernel_shape, strides, pads) + + +@_onnx_symbolic( + "aten::avg_pool1d", + decorate=[symbolic_helper._apply_params("avg_pool1d", 1)], +) +@_onnx_symbolic( + "aten::avg_pool2d", + decorate=[symbolic_helper._apply_params("avg_pool2d", 2)], +) +@_onnx_symbolic( + "aten::avg_pool3d", + decorate=[symbolic_helper._apply_params("avg_pool3d", 3)], +) +def _avg_pool(name, expand_size): + @symbolic_helper.quantized_args(True, False, False, False, False, False, False) + @symbolic_helper.parse_args("v", "is", "is", "is", "i", "i", "none") + def symbolic_fn( + g, + input: _C.Value, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: int | Sequence[int], + ceil_mode: int, + count_include_pad: int, + divisor_override=None, + ): + kernel_shape, strides, pads = _adjust_attributes_of_avg_pool( + expand_size, kernel_size, stride, padding + ) + + result = g.op( + "AveragePool", + input, + ceil_mode_i=ceil_mode, + count_include_pad_i=count_include_pad, + kernel_shape_i=kernel_shape, + pads_i=pads, + strides_i=strides, + ) + + return result + + return symbolic_fn + + +@_onnx_symbolic( + "aten::upsample_nearest1d", + decorate=[symbolic_helper._apply_params("upsample_nearest1d", 3, "nearest")], +) +@_onnx_symbolic( + "aten::upsample_nearest2d", + decorate=[symbolic_helper._apply_params("upsample_nearest2d", 4, "nearest")], +) +@_onnx_symbolic( + "aten::upsample_nearest3d", + decorate=[symbolic_helper._apply_params("upsample_nearest3d", 5, "nearest")], +) +@_onnx_symbolic( + "aten::upsample_linear1d", + decorate=[symbolic_helper._apply_params("upsample_linear1d", 3, "linear")], +) +@_onnx_symbolic( + "aten::upsample_bilinear2d", + decorate=[symbolic_helper._apply_params("upsample_bilinear2d", 4, "linear")], +) +@_onnx_symbolic( + "aten::upsample_trilinear3d", + decorate=[symbolic_helper._apply_params("upsample_trilinear3d", 5, "linear")], +) +def _interpolate(name, dim, interpolate_mode): + @symbolic_helper.quantized_args(True, False, False) + def symbolic_fn(g, input, output_size, *args): + scales, align_corners = symbolic_helper._get_interpolate_attributes( + g, interpolate_mode, args + ) + symbolic_helper._interpolate_warning(interpolate_mode) + align_corners = symbolic_helper._maybe_get_scalar(align_corners) + if align_corners: + return symbolic_helper._unimplemented(name, "align_corners == True", input) + if scales is None: + scales = symbolic_helper._interpolate_size_to_scales( + g, input, output_size, dim + ) + return g.op("Resize", input, scales, mode_s=interpolate_mode) + + return symbolic_fn + + +@_onnx_symbolic("aten::__interpolate") +def __interpolate( + g: jit_utils.GraphContext, + input, + size, + scale_factor, + mode, + align_corners, + recompute_scale_factor, + antialias, +): + scales, mode = symbolic_helper._interpolate_get_scales_and_mode( + g, input, size, scale_factor, mode, align_corners + ) + return g.op("Resize", input, scales, mode_s=mode) + + +def _slice( + g: jit_utils.GraphContext, + input: torch._C.Value, + axes: list | torch.Tensor | torch._C.Value, + starts: list | torch.Tensor | torch._C.Value, + ends: list | torch.Tensor | torch._C.Value, + steps: list | torch.Tensor | torch._C.Value | None = None, +): + def is_none_value(value): + if value is None: + return True + return ( + isinstance(value, torch._C.Value) + and value.node().kind() == "prim::Constant" + and isinstance(value.type(), _C.NoneType) + ) + + def to_slice_input(list_or_value, default_value=None): + # Convert input param into a 1D torch.Value. + if is_none_value(list_or_value) and default_value is not None: + list_or_value = [default_value] + + if isinstance(list_or_value, torch.Tensor): + return g.op("Constant", value_t=list_or_value.clone().detach()) + elif isinstance(list_or_value, list): + return g.op("Constant", value_t=torch.tensor(list_or_value)) + + rank = symbolic_helper._get_tensor_rank(list_or_value) + if rank == 0: + return symbolic_helper._unsqueeze_helper(g, list_or_value, [0]) + if rank == 1: + return list_or_value + raise errors.SymbolicValueError( + f"Rank must be 0 or 1, not {rank}", list_or_value + ) + + def get_const_value(list_or_value): + if isinstance(list_or_value, (list, torch.Tensor)): + if len(list_or_value) == 1: + return list_or_value[0] + return None + return symbolic_helper._maybe_get_const(list_or_value, "i") + + # Check if slice is a no-op + if ( + get_const_value(starts) == 0 + and get_const_value(ends) == _constants.INT64_MAX + and (steps is None or get_const_value(steps) == 1) + ): + return input + + axes = to_slice_input(axes) + starts = to_slice_input(starts, default_value=0) + ends = to_slice_input(ends, default_value=_constants.INT64_MAX) + if steps is None: + return g.op("Slice", input, starts, ends, axes) + steps = to_slice_input(steps, default_value=1) + return g.op("Slice", input, starts, ends, axes, steps) + + +@_onnx_symbolic("aten::slice") +def slice(g: jit_utils.GraphContext, self, *args): + if len(args) == 4: + # aten::slice(Tensor self, int dim, int? start=None, int? end=None, int step=1) -> Tensor + dims, start, end, step = args + elif len(args) == 3: + # aten::slice(t[] l, int? start=None, int? end=None, int step=1) -> t[] + start, end, step = args + dims = [0] + else: + raise errors.SymbolicValueError("Unknown aten::slice signature", self) + + return symbolic_helper._slice_helper( + g, + self, + axes=dims, + starts=start, + ends=end, + steps=step, + ) + + +@_onnx_symbolic("aten::flip") +@symbolic_helper.parse_args("v", "is") +def flip(g: jit_utils.GraphContext, input, dims): + return symbolic_helper._slice_helper( + g, + input, + axes=dims, + starts=[-1] * len(dims), + ends=[-_constants.INT64_MAX] * len(dims), + steps=[-1] * len(dims), + ) + + +@_onnx_symbolic("aten::fmod") +def fmod(g: jit_utils.GraphContext, input, other): + return g.op("Mod", input, other, fmod_i=1) + + +@_onnx_symbolic("aten::embedding_bag") +@symbolic_helper.parse_args("v", "v", "v", "i", "i", "i", "v", "i", "i") +def embedding_bag( + g: jit_utils.GraphContext, + embedding_matrix, + indices, + offsets, + scale_grad_by_freq, + mode, + sparse, + per_sample_weights, + include_last_offset, + padding_idx, +): + if scale_grad_by_freq and GLOBALS.export_training: + return symbolic_helper._onnx_unsupported( + "embedding_bag with scale_grad_by_freq for training mode" + ) + if padding_idx is not None and padding_idx >= 0: + raise RuntimeError("embedding_bag with padding_idx") + + warnings.warn( + "Export of embedding_bag with dynamic input/offsets shape is not supported in opset 10. " + "Please use opset 11 or higher to export model for dynamic input shape.'", + stacklevel=2, + ) + offsets_dim_0 = symbolic_helper._get_tensor_dim_size(offsets, 0) + if offsets_dim_0 is not None: + if include_last_offset: + offset_len = offsets_dim_0 - 1 + offsets_extended = offsets + else: + offset_len = offsets_dim_0 + offsets_extended = [ + offsets, + g.op("Constant", value_t=torch.tensor([sys.maxsize])), + ] + offsets_extended = g.op("Concat", *offsets_extended, axis_i=0) + list_ = [] + for i in range(offset_len): + start_ = symbolic_helper._unsqueeze_helper( + g, + opset9.select(g, offsets_extended, torch.tensor(0), torch.tensor(i)), + [0], + ) + end_ = symbolic_helper._unsqueeze_helper( + g, + opset9.select( + g, offsets_extended, torch.tensor(0), torch.tensor(i + 1) + ), + [0], + ) + axes_ = g.op("Constant", value_t=torch.tensor([0])) + indices_row = g.op("Slice", indices, start_, end_, axes_) + + embeddings = g.op("Gather", embedding_matrix, indices_row) + if not symbolic_helper._is_none(per_sample_weights): + per_sample_weights_row = g.op( + "Slice", per_sample_weights, start_, end_, axes_ + ) + per_sample_weights_row = symbolic_helper._unsqueeze_helper( + g, per_sample_weights_row, [1] + ) + embeddings = g.op("Mul", embeddings, per_sample_weights_row) + if mode == 0: + embeddings = symbolic_helper._reducesum_helper( + g, embeddings, axes_i=[0], keepdims_i=0 + ) + elif mode == 1: + embeddings = g.op("ReduceMean", embeddings, axes_i=[0], keepdims_i=0) + else: + embeddings = g.op("ReduceMax", embeddings, axes_i=[0], keepdims_i=0) + + embeddings = symbolic_helper._unsqueeze_helper(g, embeddings, [0]) + list_.append(embeddings) + + output = g.op("Concat", *list_, axis_i=0) + # aten::embedding_bag returns a tuple of 4 elements: output, offset2bag, bag_size, max_indices. + # But the last three outputs are not used in torch.nn.EmbeddingBag or torch.nn.functional.embedding_bag. + return output, None, None, None + else: + return symbolic_helper._onnx_unsupported( + "embedding_bag with unknown shape of offsets for opset 10 is not supported. " + "please use opset 11 or higher." + ) + + +@_onnx_symbolic("aten::fake_quantize_per_tensor_affine") +@symbolic_helper.parse_args("v", "v", "v", "i", "i") +def fake_quantize_per_tensor_affine( + g: jit_utils.GraphContext, + inputs, + scale, + zero_point, + quant_min=-128, + quant_max=127, +): + # NOTE: (0, 127) is a special case. PyTorch restricts activations to be in the range (0, 127). + # https://github.com/pytorch/pytorch/blob/b34b192d6b97325c9f78e5995c48c8498ede34bd/torch/ao/quantization/observer.py#L1422 + if (quant_min, quant_max) == (0, 127): + symbolic_helper._onnx_opset_unsupported_detailed( + "fake_quantize_per_tensor_affine", + 10, + 13, + "Quantize range (0, 127) not supported, requires opset 13 Clip", + inputs, + ) + if (quant_min, quant_max) not in [(0, 255), (-128, 127)]: + raise errors.SymbolicValueError( + f"For (quant_min, quant_max), ONNX allows only (0, 255) and (-128, 127). " + f"Got ({quant_min}, {quant_max})", + inputs, + ) + scale = symbolic_helper._maybe_get_scalar(scale) + if scale is None: + symbolic_helper._onnx_opset_unsupported_detailed( + "fake_quantize_per_tensor_affine", + 10, + 13, + "Non-constant scale not supported", + inputs, + ) + + scale = scale.float().data # Avoid exporter generating double type + if quant_min == 0: + zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.UINT8) + else: + zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.INT8) + return g.op( + "DequantizeLinear", + g.op("QuantizeLinear", inputs, scale, zero_point), + scale, + zero_point, + ) + + +@_onnx_symbolic("aten::isinf") +def isinf(g: jit_utils.GraphContext, input): + return g.op("IsInf", g.op("Cast", input, to_i=_C_onnx.TensorProtoDataType.DOUBLE)) + + +@_onnx_symbolic("aten::isfinite") +def isfinite(g: jit_utils.GraphContext, input): + inf_node = isinf(g, input) + nan_node = opset9.isnan(g, input) + return opset9.__not_(g, opset9.__or_(g, inf_node, nan_node)) + + +@_onnx_symbolic("aten::quantize_per_tensor") +def quantize_per_tensor(g: jit_utils.GraphContext, input, scale, zero_point, dtype): + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + # TODO(justinchuby): Extract all the cast ops into a helper function. + zero_point = g.op( + "Cast", zero_point, to_i=_type_utils.JitScalarType(dtype).onnx_type() + ) + scale = g.op("Cast", scale, to_i=_C_onnx.TensorProtoDataType.FLOAT) + return symbolic_helper.quantize_helper(g, input, scale, zero_point) + + +@_onnx_symbolic("aten::dequantize") +def dequantize(g: jit_utils.GraphContext, input): + return symbolic_helper.dequantize_helper(g, input)[0] + + +@_onnx_symbolic("aten::nan_to_num") +@symbolic_helper.parse_args("v", "f", "f", "f") +def nan_to_num(g: jit_utils.GraphContext, input, nan, posinf, neginf): + # Cannot create a int type tensor with inf/nan values, so we simply + # return the original tensor + if not symbolic_helper._is_fp(input): + return input + input_dtype = _type_utils.JitScalarType.from_value(input).dtype() + if nan is None: + nan = 0.0 + nan_cond = opset9.isnan(g, input) + nan_result = g.op( + "Where", + nan_cond, + g.op("Constant", value_t=torch.tensor([nan], dtype=input_dtype)), + input, + ) + + # For None values of posinf, neginf we use the greatest/lowest finite + # value representable by input's dtype. + finfo = torch.finfo(input_dtype) + if posinf is None: + posinf = finfo.max + posinf_cond = opset9.logical_and( + g, + isinf(g, nan_result), + opset9.gt(g, nan_result, g.op("Constant", value_t=torch.LongTensor([0]))), + ) + nan_posinf_result = g.op( + "Where", + posinf_cond, + g.op("Constant", value_t=torch.tensor([posinf], dtype=input_dtype)), + nan_result, + ) + + if neginf is None: + neginf = finfo.min + neginf_cond = opset9.logical_and( + g, + isinf(g, nan_posinf_result), + opset9.lt( + g, nan_posinf_result, g.op("Constant", value_t=torch.LongTensor([0])) + ), + ) + return g.op( + "Where", + neginf_cond, + g.op("Constant", value_t=torch.tensor([neginf], dtype=input_dtype)), + nan_posinf_result, + ) + + +# Quantized symbolics --------------------------------------------------------- +# https://github.com/pytorch/pytorch/wiki/PyTorch-ONNX-exporter#quantized-model-export +# Support starts from opset 10 because `DequantizeLinear` and `QuantizeLinear` were +# introduced in opset version 10. +@_onnx_symbolic("quantized::linear") +def quantized_linear( + g: jit_utils.GraphContext, q_input, q_weight, bias, op_scale, op_zero_point +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.linear(g, input, weight, bias) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::linear_relu") +def quantized_linear_relu( + g: jit_utils.GraphContext, q_input, q_weight, bias, op_scale, op_zero_point +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.linear(g, input, weight, bias) + output = opset9.relu(g, output) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::add") +def quantized_add(g: jit_utils.GraphContext, x, y, op_scale, op_zero_point): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + y, _, _, _ = symbolic_helper.dequantize_helper(g, y) + + output = opset9.add(g, x, y) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::add_relu") +def quantized_add_relu(g: jit_utils.GraphContext, x, y, op_scale, op_zero_point): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + y, _, _, _ = symbolic_helper.dequantize_helper(g, y) + + output = opset9.add(g, x, y) + output = opset9.relu(g, output) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::mul") +def quantized_mul(g: jit_utils.GraphContext, x, y, op_scale, op_zero_point): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + y, _, _, _ = symbolic_helper.dequantize_helper(g, y) + + output = opset9.mul(g, x, y) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::hardswish") +def quantized_hardswish(g: jit_utils.GraphContext, x, op_scale, op_zero_point): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + + output = opset9.hardswish(g, x) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::sigmoid") +def quantized_sigmoid(g: jit_utils.GraphContext, x, op_scale, op_zero_point): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + + output = opset9.sigmoid(g, x) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::leaky_relu") +def quantized_leaky_relu( + g: jit_utils.GraphContext, x, negative_slope, inplace, op_scale, op_zero_point +): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + + output = opset9.leaky_relu(g, x, negative_slope, inplace) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::layer_norm") +def quantized_layer_norm( + g: jit_utils.GraphContext, + x, + normalized_shape, + weight, + bias, + eps, + op_scale, + op_zero_point, +): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + + output = opset9.layer_norm(g, x, normalized_shape, weight, bias, eps, False) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::group_norm") +def quantized_group_norm( + g: jit_utils.GraphContext, + x, + num_groups, + weight, + bias, + eps, + op_scale, + op_zero_point, +): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + + output = opset9.group_norm(g, x, num_groups, weight, bias, eps, False) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::instance_norm") +@symbolic_helper.parse_args("v", "v", "v", "f", "v", "v") +def quantized_instance_norm( + g: jit_utils.GraphContext, + q_input, + weight, + bias, + eps, + op_scale, + op_zero_point, +): + input, _, _, _ = symbolic_helper.dequantize_helper(g, q_input) + + output = opset9.instance_norm( + g, input, weight, bias, None, None, False, 0.0, eps, False + ) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv1d_relu") +def quantized_conv1d_relu( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv1d(g, input, weight, bias, stride, padding, dilation, groups) + output = opset9.relu(g, output) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv2d_relu") +def quantized_conv2d_relu( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv2d(g, input, weight, bias, stride, padding, dilation, groups) + output = opset9.relu(g, output) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv3d_relu") +def quantized_conv3d_relu( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv3d(g, input, weight, bias, stride, padding, dilation, groups) + output = opset9.relu(g, output) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv1d") +def quantized_conv1d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv1d(g, input, weight, bias, stride, padding, dilation, groups) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv2d") +def quantized_conv2d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv2d(g, input, weight, bias, stride, padding, dilation, groups) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv3d") +def quantized_conv3d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv3d(g, input, weight, bias, stride, padding, dilation, groups) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv_transpose1d") +def quantized_conv_transpose1d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + output_padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv_transpose2d( + g, input, weight, bias, stride, padding, output_padding, groups, dilation + ) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv_transpose2d") +def quantized_conv_transpose2d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + output_padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv_transpose2d( + g, input, weight, bias, stride, padding, output_padding, groups, dilation + ) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv_transpose3d") +def quantized_conv_transpose3d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + output_padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, _ = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper(g, bias, input_scale, weight_scale) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv_transpose3d( + g, input, weight, bias, stride, padding, output_padding, groups, dilation + ) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::cat") +@symbolic_helper.parse_args("v", "i", "v", "v") +def quantized_cat( + g: jit_utils.GraphContext, + q_inputs: _C.Value, + dim: int, + op_scale: _C.Value, + op_zero_point: _C.Value, +) -> _C.Value: + unpacked_inputs = symbolic_helper._unpack_list(q_inputs) + dequantized = [ + symbolic_helper.dequantize_helper(g, input)[0] for input in unpacked_inputs + ] + concatenated = g.op("Concat", *dequantized, axis_i=dim) + return symbolic_helper.quantize_helper(g, concatenated, op_scale, op_zero_point) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset11.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset11.py new file mode 100644 index 0000000000000000000000000000000000000000..07d0424c4e1984cd43cad2ed70566fcad01b9d5f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset11.py @@ -0,0 +1,1476 @@ +# mypy: allow-untyped-defs +# mypy: disable-error-code=arg-type +"""This file exports ONNX ops for opset 11.""" + +from __future__ import annotations + +import functools +import sys +import warnings +from typing import TYPE_CHECKING + +import torch +from torch import _C +from torch._C import _onnx as _C_onnx +from torch.onnx import errors +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, + symbolic_opset10 as opset10, + symbolic_opset9 as opset9, + utils, +) + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in README.md + +__all__ = [ + "add", + "append", + "arange", + "argsort", + "atleast_1d", + "atleast_2d", + "atleast_3d", + "cat", + "chunk", + "clamp_max", + "clamp_min", + "clamp", + "constant_pad_nd", + "cumsum", + "Delete", + "embedding_bag", + "embedding_renorm", + "flatten", + "gather", + "hardtanh", + "hstack", + "im2col", + "index_fill", + "index", + "index_copy", + "index_put", + "insert", + "linalg_det", + "linalg_vector_norm", + "logdet", + "masked_scatter", + "masked_select", + "mm", + "narrow", + "normal", + "pad", + "pixel_shuffle", + "pop", + "prim_constant_chunk", + "reflection_pad", + "relu6", + "remainder", + "replication_pad", + "round", + "scatter", + "select", + "size", + "sort", + "split_with_sizes", + "split", + "squeeze", + "stack", + "topk", + "unbind", + "unique_dim", + "unsqueeze", + "vstack", +] + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=11) + + +@_onnx_symbolic("aten::hardtanh") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "f", "f") +def hardtanh(g: jit_utils.GraphContext, self: _C.Value, min_val: float, max_val: float): + scalar_type = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.FLOAT + ) + min_val = g.op( + "Constant", + value_t=torch.tensor(min_val, dtype=scalar_type.dtype()), + ) + max_val = g.op( + "Constant", + value_t=torch.tensor(max_val, dtype=scalar_type.dtype()), + ) + return symbolic_helper._op_with_optional_float_cast( + g, "Clip", self, min_val, max_val, opset_before=12 + ) + + +@_onnx_symbolic("aten::clamp") +def clamp(g: jit_utils.GraphContext, self, min, max): + def _cast_if_not_none(tensor, dtype): + if tensor is not None and not symbolic_helper._is_none(tensor): + return g.op( + "Cast", + tensor, + to_i=dtype.onnx_type(), + ) + else: + return tensor + + scalar_type = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.UNDEFINED + ) + if scalar_type != _type_utils.JitScalarType.UNDEFINED: + min = _cast_if_not_none(min, scalar_type) + max = _cast_if_not_none(max, scalar_type) + + if symbolic_helper._is_none(min): + return clamp_max(g, self, max) + elif symbolic_helper._is_none(max): + return clamp_min(g, self, min) + else: + if ( + symbolic_helper._get_tensor_rank(min) == 0 + and symbolic_helper._get_tensor_rank(max) == 0 + ): + return symbolic_helper._op_with_optional_float_cast( + g, "Clip", self, min, max, opset_before=12 + ) + else: + return clamp_max(g, clamp_min(g, self, min), max) + + +@_onnx_symbolic("aten::clamp_min") +@symbolic_helper.parse_args("v", "v") +def clamp_min(g: jit_utils.GraphContext, self, min): + min = g.op("Cast", min, to_i=_type_utils.JitScalarType.from_value(self).onnx_type()) + if symbolic_helper._get_tensor_rank(min) == 0: + max = opset9.unused(g) + return symbolic_helper._op_with_optional_float_cast( + g, "Clip", self, min, max, opset_before=12 + ) + else: + return symbolic_helper._op_with_optional_float_cast( + g, "Max", self, min, opset_before=12 + ) + + +@_onnx_symbolic("aten::clamp_max") +@symbolic_helper.parse_args("v", "v") +def clamp_max(g: jit_utils.GraphContext, self, max): + max = g.op("Cast", max, to_i=_type_utils.JitScalarType.from_value(self).onnx_type()) + if symbolic_helper._get_tensor_rank(max) == 0: + min = opset9.unused(g) + return symbolic_helper._op_with_optional_float_cast( + g, "Clip", self, min, max, opset_before=12 + ) + else: + return symbolic_helper._op_with_optional_float_cast( + g, "Min", self, max, opset_before=12 + ) + + +@_onnx_symbolic("aten::relu6") +def relu6(g: jit_utils.GraphContext, input): + scalar_type = _type_utils.JitScalarType.from_value( + input, _type_utils.JitScalarType.FLOAT + ) + min_val = g.op( + "Constant", + value_t=torch.tensor(0, dtype=scalar_type.dtype()), + ) + max_val = g.op( + "Constant", + value_t=torch.tensor(6, dtype=scalar_type.dtype()), + ) + return clamp(g, input, min_val, max_val) + + +@_onnx_symbolic("aten::select") +# Opset 11 gather accepts negative indices +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "i", "v") +def select(g: jit_utils.GraphContext, self, dim, index): + return g.op("Gather", self, index, axis_i=dim) + + +@_onnx_symbolic("aten::index_put") +def index_put( + g: jit_utils.GraphContext, self, indices_list_value, values, accumulate=False +): + if symbolic_helper._is_packed_list(indices_list_value): + indices_list = symbolic_helper._unpack_list(indices_list_value) + else: + indices_list = [indices_list_value] + accumulate = symbolic_helper._parse_arg(accumulate, "b") + + if len(indices_list) == 0: + return values + + if len(indices_list) > 1: + for idx_ in range(len(indices_list)): + if symbolic_helper._is_bool(indices_list[idx_]): + indices_list[idx_] = g.op("NonZero", indices_list[idx_]) + index = indices_list[0] + + for ind in indices_list[1:]: + index = opset9.add(g, index, ind) + broadcast_index_shape = g.op("Shape", index) + indices_list = [ + symbolic_helper._unsqueeze_helper( + g, opset9.expand(g, ind, broadcast_index_shape, None), [-1] + ) + for ind in indices_list + ] + index = g.op("Concat", *indices_list, axis_i=-1) + else: + # Replace index_put node with masked_scatter or masked_fill + # when inputs to the index_put node contains a single boolean input. + # + # index_put -> masked_fill + # * input index contains single tensor of Bool type (e.g.: %24 <- %23). + # * input value contains single element (e.g.: %18). + # + # Torch IR + # %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6) + # %16 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = + # aten::to(%8, %26, %27, %11, %12, %28, %29, %15) + # %18 : Float(requires_grad=0, device=cpu) = prim::Constant[value={1}]() + # %23 : Bool(8, strides=[1], device=cpu) = aten::view(%16, %22) + # %24 : Tensor?[] = prim::ListConstruct(%23) + # %25 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = + # aten::index_put(%mask, %24, %18, %30) + # return (%25) + # + # + # index_put -> masked_scatter + # * input index contains single tensor of Bool type (e.g.: %32 <- %31). + # * input value contains multiple elements (e.g.: %28). + # + # Torch IR + # %mask : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) = aten::clone(%0, %6) + # %28 : Float(8, strides=[1], requires_grad=0, device=cpu) + # = prim::Constant[value= 1 1 1 1 1 1 1 1 [ CPUFloatType{8} ]]() + # %15 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) + # = aten::ne(%mask, %some_const) + # %23 : Bool(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) + # = aten::to(%15, %34, %35, %18, %19, %36, %37, %22) + # %38 : Long(requires_grad=0, device=cpu) = prim::Constant[value={0}]() + # %30 : int[] = prim::Constant[value=[-1]]() + # %31 : Bool(8, strides=[1], device=cpu) = aten::view(%23, %30) + # %32 : Tensor?[] = prim::ListConstruct(%31) + # %33 : Float(2, 2, 2, strides=[4, 2, 1], requires_grad=0, device=cpu) + # = aten::index_put(%mask, %32, %28, %38) + # return (%33) + index = indices_list[0] + bool_inp = index + if symbolic_helper._is_bool(bool_inp): + rank = symbolic_helper._get_tensor_rank(values) + if rank is not None and rank == 0: + return opset9.masked_fill(g, self, bool_inp, values) + mask_rank = symbolic_helper._get_tensor_rank(bool_inp) + self_rank = symbolic_helper._get_tensor_rank(self) + if ( + mask_rank is not None + and self_rank is not None + and self_rank > mask_rank + ): + # Unsqueeze 'bool_inp' to be broadcastable to shape of 'self'. + bool_inp = symbolic_helper._unsqueeze_helper( + g, bool_inp, list(range(mask_rank, self_rank)) + ) + return masked_scatter(g, self, bool_inp, values) + broadcast_index_shape = g.op("Shape", index) + index = symbolic_helper._unsqueeze_helper(g, index, [-1]) + sub_data_shape = symbolic_helper._slice_helper( + g, g.op("Shape", self), axes=[0], starts=[len(indices_list)], ends=[sys.maxsize] + ) + values_shape = g.op("Concat", broadcast_index_shape, sub_data_shape, axis_i=0) + # Check if values is a singular value and expand accordingly + rank = symbolic_helper._get_tensor_rank(values) + if rank is not None and rank == 0: + values = opset9.expand(g, values, values_shape, None) + values = symbolic_helper._reshape_helper(g, values, values_shape) + + self_scalar_type = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.UNDEFINED + ) + if self_scalar_type != _type_utils.JitScalarType.UNDEFINED: + values_scalar_type = _type_utils.JitScalarType.from_value( + values, _type_utils.JitScalarType.UNDEFINED + ) + if self_scalar_type != values_scalar_type: + values = g.op("Cast", values, to_i=self_scalar_type.onnx_type()) + elif accumulate: + raise errors.SymbolicValueError("self does not have a valid scalar type.", self) + + if accumulate: + zeros = g.op( + "ConstantOfShape", + g.op("Shape", self), + value_t=torch.tensor([0], dtype=self_scalar_type.dtype()), + ) + result = g.op("ScatterND", zeros, index, values) + result = add(g, self, result) + else: + result = g.op("ScatterND", self, index, values) + + return result + + +@_onnx_symbolic("aten::pixel_shuffle") +@symbolic_helper.parse_args("v", "i") +def pixel_shuffle(g: jit_utils.GraphContext, self, upscale_factor): + rank = symbolic_helper._get_tensor_rank(self) + if rank is not None and rank != 4: + return symbolic_helper._unimplemented("pixel_shuffle", "only support 4d input") + return g.op("DepthToSpace", self, blocksize_i=upscale_factor, mode_s="CRD") + + +@_onnx_symbolic( + "aten::upsample_nearest1d", + decorate=[symbolic_helper._apply_params("upsample_nearest1d", 3, "nearest")], +) +@_onnx_symbolic( + "aten::upsample_nearest2d", + decorate=[symbolic_helper._apply_params("upsample_nearest2d", 4, "nearest")], +) +@_onnx_symbolic( + "aten::upsample_nearest3d", + decorate=[symbolic_helper._apply_params("upsample_nearest3d", 5, "nearest")], +) +@_onnx_symbolic( + "aten::upsample_linear1d", + decorate=[symbolic_helper._apply_params("upsample_linear1d", 3, "linear")], +) +@_onnx_symbolic( + "aten::upsample_bilinear2d", + decorate=[symbolic_helper._apply_params("upsample_bilinear2d", 4, "linear")], +) +@_onnx_symbolic( + "aten::upsample_trilinear3d", + decorate=[symbolic_helper._apply_params("upsample_trilinear3d", 5, "linear")], +) +@_onnx_symbolic( + "aten::upsample_bicubic2d", + decorate=[symbolic_helper._apply_params("upsample_bicubic2d", 4, "cubic")], +) +def _interpolate(name: str, dim: int, interpolate_mode: str): + return symbolic_helper._interpolate_helper(name, dim, interpolate_mode) + + +@_onnx_symbolic("aten::__interpolate") +@symbolic_helper.quantized_args(True, False, False, False, False, False, False) +def __interpolate( + g: jit_utils.GraphContext, + input, + size, + scale_factor, + mode, + align_corners, + recompute_scale_factor, + antialias, +): + return symbolic_helper.__interpolate_helper( + g, input, size, scale_factor, mode, align_corners, recompute_scale_factor + ) + + +@_onnx_symbolic("aten::gather") +@symbolic_helper.parse_args("v", "i", "v", "v") +def gather(g: jit_utils.GraphContext, self, dim, index, sparse_grad=False): + if symbolic_helper._maybe_get_const(sparse_grad, "i"): + return symbolic_helper._unimplemented("gather", "sparse_grad == True") + return g.op("GatherElements", self, index, axis_i=dim) + + +@_onnx_symbolic("aten::scatter") +@symbolic_helper.parse_args("v", "i", "v", "v") +def scatter(g: jit_utils.GraphContext, self, dim, index, src): + src_type = _type_utils.JitScalarType.from_value(src) + src = symbolic_helper._maybe_get_scalar(src) + if symbolic_helper._is_value(src): + return g.op("ScatterElements", self, index, src, axis_i=dim) + else: + # Check if scalar "src" has same type as self (PyTorch allows different + # type for scalar src (but not when src is tensor)). If not, insert Cast node. + if _type_utils.JitScalarType.from_value(self) != src_type: + src = g.op( + "Cast", + src, + to_i=_type_utils.JitScalarType.from_value(self).onnx_type(), + ) + return g.op( + "ScatterElements", self, index, opset9.expand_as(g, src, index), axis_i=dim + ) + + +@_onnx_symbolic("aten::cumsum") +@symbolic_helper.parse_args("v", "i", "none") +def cumsum(g: jit_utils.GraphContext, self, dim, dtype=None): + dim_tensor = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.int)) + if dtype and dtype.node().kind() != "prim::Constant": + parsed_dtype = symbolic_helper._get_const(dtype, "i", "dtype") + cast = g.op( + "Cast", self, to_i=_type_utils.JitScalarType(parsed_dtype).onnx_type() + ) + else: + cast = self + csum = g.op("CumSum", cast, dim_tensor) + return csum + + +@_onnx_symbolic("aten::masked_select") +def masked_select(g: jit_utils.GraphContext, self, mask): + index = opset9.nonzero(g, opset9.expand_as(g, mask, self)) + return g.op("GatherND", self, index) + + +@_onnx_symbolic("aten::masked_scatter") +def masked_scatter(g: jit_utils.GraphContext, self, mask, source): + index = opset9.nonzero(g, opset9.expand_as(g, mask, self)) + # NOTE: source can have more elements than needed. + # It could also have arbitrary shape. + # This is not supported by ONNX::ScatterND, so we need to flatten and slice source tensor. + source = symbolic_helper._reshape_helper(g, source, torch.LongTensor([-1])) + source = symbolic_helper._slice_helper( + g, + source, + axes=torch.LongTensor([0]), + starts=torch.LongTensor([0]), + ends=opset9.size(g, index, torch.LongTensor([0])), + ) + return g.op("ScatterND", self, index, source) + + +@_onnx_symbolic("aten::len") +def _len(g: jit_utils.GraphContext, self): + if ( + symbolic_helper._is_tensor_list(self) + or self.node().kind() == "onnx::SplitToSequence" + ): + return g.op("SequenceLength", self) + sz_0 = size(g, self, g.op("Constant", value_t=torch.LongTensor([0]))) + return symbolic_helper._squeeze_helper(g, sz_0, [0]) + + +@_onnx_symbolic("aten::__getitem_") +def __getitem_(g: jit_utils.GraphContext, self, i): + if symbolic_helper._is_tensor_list(self): + # SequenceAt requires that the input be a List of Tensors + return g.op("SequenceAt", self, i) + else: + from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import ( + __getitem_ as getitem, + ) + + return getitem(g, self, i) + + +@_onnx_symbolic("aten::_set_item") +def _set_item(g: jit_utils.GraphContext, tensor_list, i, v): + tensor_list = g.op("SequenceErase", tensor_list, i) + return g.op("SequenceInsert", tensor_list, v, i) + + +@_onnx_symbolic("aten::append") +def append(g: jit_utils.GraphContext, self, tensor): + return g.op("SequenceInsert", self, tensor) + + +@_onnx_symbolic("aten::add") +def add(g: jit_utils.GraphContext, self, other, alpha=None): + if symbolic_helper._is_value(self) and symbolic_helper._is_tensor_list(self): + tensor_list_node = other.node() + if tensor_list_node.kind() != "prim::ListConstruct": + return symbolic_helper._unimplemented( + "add", "does not support adding dynamic tensor list to another" + ) + tensors = symbolic_helper._unpack_list(other) + l = self + for t in tensors: + l = g.op("SequenceInsert", l, t) + return l + + return opset9.add(g, self, other, alpha) + + +@_onnx_symbolic("aten::insert") +def insert(g: jit_utils.GraphContext, self, pos, tensor): + return g.op("SequenceInsert", self, tensor, pos) + + +@_onnx_symbolic("aten::pop") +def pop(g: jit_utils.GraphContext, tensor_list, dim): + return g.op("SequenceErase", tensor_list, dim) + + +@_onnx_symbolic("aten::Delete") +def Delete(g: jit_utils.GraphContext, tensor_list, dim): + return g.op("SequenceErase", tensor_list, dim) + + +@_onnx_symbolic("aten::cat") +@symbolic_helper.quantized_args(True) +def cat(g: jit_utils.GraphContext, tensor_list, dim): + if symbolic_helper._is_packed_list(tensor_list): + return opset9.cat(g, tensor_list, dim) + else: + dim = symbolic_helper._get_const(dim, "i", "dim") + return g.op("ConcatFromSequence", tensor_list, axis_i=dim) + + +@_onnx_symbolic("aten::stack") +def stack(g: jit_utils.GraphContext, tensor_list, dim): + if symbolic_helper._is_packed_list(tensor_list): + return opset9.stack(g, tensor_list, dim) + else: + dim = symbolic_helper._get_const(dim, "i", "dim") + return g.op("ConcatFromSequence", tensor_list, axis_i=dim, new_axis_i=1) + + +@_onnx_symbolic("aten::_unique2") +@symbolic_helper.parse_args("v", "i", "i", "i") +def _unique2(g: jit_utils.GraphContext, self, sorted, return_inverse, return_counts): + u, _indices, inverse_indices, counts = g.op( + "Unique", self, sorted_i=sorted, outputs=4 + ) + return u, inverse_indices, counts + + +@_onnx_symbolic("aten::unique_dim") +@symbolic_helper.parse_args("v", "i", "i", "i", "i") +def unique_dim( + g: jit_utils.GraphContext, self, dim, sorted, return_inverse, return_counts +): + u, _indices, inverse_indices, counts = g.op( + "Unique", self, axis_i=dim, sorted_i=sorted, outputs=4 + ) + return u, inverse_indices, counts + + +@_onnx_symbolic("aten::topk") +@symbolic_helper.parse_args("v", "v", "i", "i", "i", "none") +def topk(g: jit_utils.GraphContext, self, k, dim, largest, sorted, out=None): + return symbolic_helper._topk_helper( + g, self, k, dim, largest=largest, sorted=sorted, out=out + ) + + +@_onnx_symbolic("aten::sort") +@symbolic_helper.parse_args("v", "i", "i", "none") +def sort(g: jit_utils.GraphContext, self, dim, descending, out=None): + return symbolic_helper._sort_helper(g, self, dim, descending=descending, out=out) + + +@_onnx_symbolic("aten::argsort") +@symbolic_helper.parse_args("v", "i", "i", "none") +def argsort(g: jit_utils.GraphContext, self, dim, descending, out=None): + _, indices = symbolic_helper._sort_helper( + g, self, dim, descending=descending, out=out + ) + return indices + + +@_onnx_symbolic("aten::round") +@symbolic_helper.parse_args("v", "i") +def round(g: jit_utils.GraphContext, self, decimals=0): + if not symbolic_helper._is_fp(self): + return self + if decimals == 0: + return g.op("Round", self) + mul = g.op("Mul", self, g.op("Constant", value_t=torch.tensor(pow(10, decimals)))) + round = g.op("Round", mul) + return g.op( + "Mul", round, g.op("Constant", value_t=torch.tensor(pow(10, -1 * decimals))) + ) + + +@_onnx_symbolic("aten::remainder") +def remainder(g: jit_utils.GraphContext, input, other): + if symbolic_helper._is_fp(input) or symbolic_helper._is_fp(other): + return opset9.remainder(g, input, other) + return g.op("Mod", input, other, fmod_i=0) + + +@_onnx_symbolic("aten::split") +@symbolic_helper.parse_args("v", "v", "i", "i") +def split(g: jit_utils.GraphContext, self, split_size_or_sizes, dim, _outputs=None): + if not symbolic_helper._is_split_static(split_size_or_sizes, _outputs): + split_out = g.op("SplitToSequence", self, split_size_or_sizes, axis_i=dim) + if _outputs is None: + return split_out + # Convert to multiple slice nodes iff number of splits and number of outputs are statically known. + if ( + symbolic_helper._is_packed_list(split_size_or_sizes) + and len(symbolic_helper._unpack_list(split_size_or_sizes)) == _outputs + ): + split_sizes = [ + symbolic_helper._unsqueeze_helper(g, v, [0]) + for v in symbolic_helper._unpack_list(split_size_or_sizes) + ] + start = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long)) + axis = g.op("Constant", value_t=torch.tensor([dim], dtype=torch.long)) + res = [] + for i in range(_outputs): + end = g.op( + "Add", start, split_sizes[i] + ) # split_sizes is a list of same length as _outputs + res.append(g.op("Slice", self, start, end, axis)) + start = end + return res + return [ + g.op( + "SequenceAt", + split_out, + g.op("Constant", value_t=torch.tensor([i], dtype=torch.long)), + ) + for i in range(_outputs) + ] + else: + return opset9.split(g, self, split_size_or_sizes, dim, _outputs) + + +@_onnx_symbolic("aten::split_with_sizes") +@symbolic_helper.parse_args("v", "v", "i", "i") +def split_with_sizes(g: jit_utils.GraphContext, self, split_sizes, dim, _outputs=None): + return split(g, self, split_sizes, dim, _outputs) + + +@_onnx_symbolic("aten::unbind") +@symbolic_helper.parse_args("v", "i", "i") +def unbind(g: jit_utils.GraphContext, self, dim=0, _outputs=None): + if _outputs is None: + return g.op( + "SplitToSequence", + self, + g.op("Constant", value_t=torch.tensor(1, dtype=torch.long)), + axis_i=dim, + keepdims_i=0, + ) + else: + return opset9.unbind(g, self, dim, _outputs) + + +def _prepare_onnx_paddings(g: jit_utils.GraphContext, input, pad): + """Generate paddings in ONNX order based on pad in pytorch. + + Args: + input: the input tensor. + pad: the paddings in pytorch. + The order is dim_n_begin, dim_n_end, dim_n-1_begin, dim_n-1_end, ..., dim_m_begin, dim_m_end, + where m is in range [0, n]. + """ + if ( + not symbolic_helper._is_packed_list(pad) + and symbolic_helper._is_list(pad) + and symbolic_helper._is_scalar_list(pad) + ): + pad = g.op("ConcatFromSequence", pad, axis_i=0, new_axis_i=1) + # The desired order of paddings is + # dim_0_begin, dim_1_begin, ... , dim_0_end, ..., dim_n_end. + # n is the dimension of input. + # Assume zero-dimensions in the beginning, pad the "pad" sequence with zeros in the beginning + pad_len = opset9.size(g, pad, g.op("Constant", value_t=torch.tensor([0]))) + # Set extension = [0] * (dim * 2 - len(pad)) + rank = symbolic_helper._get_tensor_rank(input) + if rank is None: + rank = g.op("Size", g.op("Shape", input)) + else: + rank = g.op("Constant", value_t=torch.tensor(rank, dtype=torch.int64)) + extension = g.op( + "Sub", + g.op("Mul", rank, g.op("Constant", value_t=torch.tensor(2, dtype=torch.int64))), + pad_len, + ) + # Concat pad with extension: paddings = [dim_n_begin, dim_n_end, dim_n-1_begin, dim_n-1_end, 0, 0, ... ] + # Currently ONNX only supports int64 type for Pad + pad = g.op("Cast", pad, to_i=_C_onnx.TensorProtoDataType.INT64) + paddings = g.op( + "Concat", + pad, + g.op( + "ConstantOfShape", extension, value_t=torch.tensor([0], dtype=torch.int64) + ), + axis_i=0, + ) + # Reshape and reverse order and collate first beginnings and then ends + # paddings = [[..., 0, dim_n-1_begin, dim_n_begin], + # [..., 0, dim_n-1_end, dim_n_end]] + # Reshape back to 1-D paddings = [..., 0, dim_n - 1_begin, dim_n_begin, ..., 0, dim_n - 1_end, dim_n_end] + paddings = symbolic_helper._reshape_helper( + g, paddings, g.op("Constant", value_t=torch.tensor([-1, 2])) + ) + paddings = g.op("Transpose", opset10.flip(g, paddings, [0]), perm_i=[1, 0]) + paddings = symbolic_helper._reshape_helper( + g, paddings, g.op("Constant", value_t=torch.tensor([-1])) + ) + padding_c = g.op("Cast", paddings, to_i=_C_onnx.TensorProtoDataType.INT64) + return padding_c + + +@_onnx_symbolic("aten::constant_pad_nd") +def constant_pad_nd(g: jit_utils.GraphContext, input, padding, value=None): + mode = "constant" + value = symbolic_helper._maybe_get_scalar(value) + value = symbolic_helper._if_scalar_type_as(value, input) + pad = _prepare_onnx_paddings(g, input, padding) + return g.op("Pad", input, pad, value, mode_s=mode) + + +@_onnx_symbolic("aten::reflection_pad1d") +@_onnx_symbolic("aten::reflection_pad2d") +@_onnx_symbolic("aten::reflection_pad3d") +def reflection_pad(g: jit_utils.GraphContext, input, padding): + mode = "reflect" + paddings = _prepare_onnx_paddings(g, input, padding) + return g.op("Pad", input, paddings, mode_s=mode) + + +@_onnx_symbolic("aten::replication_pad1d") +@_onnx_symbolic("aten::replication_pad2d") +@_onnx_symbolic("aten::replication_pad3d") +def replication_pad(g: jit_utils.GraphContext, input, padding): + mode = "edge" + paddings = _prepare_onnx_paddings(g, input, padding) + return g.op("Pad", input, paddings, mode_s=mode) + + +@_onnx_symbolic("aten::pad") +def pad( + g: jit_utils.GraphContext, + input: _C.Value, + pad: _C.Value, + mode: _C.Value, + value: _C.Value, +): + mode = symbolic_helper._parse_arg(mode, "s") + if mode == "replicate": + return replication_pad(g, input, pad) + elif mode == "reflect": + return reflection_pad(g, input, pad) + elif mode == "constant": + return constant_pad_nd(g, input, pad, value) + elif mode == "circular": + return opset9._pad_circular(g, input, pad) + else: + raise errors.SymbolicValueError(f"Unrecognized padding mode {mode}", input) + + +@_onnx_symbolic("aten::linalg_det") +def linalg_det(g: jit_utils.GraphContext, self): + return g.op("Det", self) + + +@_onnx_symbolic("aten::logdet") +def logdet(g: jit_utils.GraphContext, input): + return opset9.log(g, linalg_det(g, input)) + + +@_onnx_symbolic("aten::arange") +def arange(g: jit_utils.GraphContext, *args): + def _get_arange_dtype(dtype): + dtype = symbolic_helper._maybe_get_const(dtype, "i") + return dtype + + if len(args) == 2 and all(isinstance(val, int) for val in args): + # aten::arange(Scalar start, Scalar end) + dtype = torch.int64 + # Start index. + start = g.op( + "Constant", + value_t=torch.tensor(args[0], dtype=dtype), + ) + # End (exclusive) index. + end = g.op( + "Constant", + value_t=torch.tensor(args[1], dtype=dtype), + ) + # Step size from start to end indexes. + delta_default = g.op( + "Constant", + value_t=torch.tensor(1, dtype=dtype), + ) + return g.op("Range", start, end, delta_default) + elif len(args) == 2 or len(args) == 5: + if len(args) == 2: + # aten::arange(Scalar end, Tensor out) + dtype = None + else: + # aten::arange(Scalar end, ScalarType dtype, Layout, Device, bool pin_memory) + dtype = _get_arange_dtype(args[1]) + type_, end, start, step = symbolic_helper._arange_cast_helper( + g, end=args[0], dtype=dtype + ) + start_default = g.op( + "Constant", + value_t=torch.tensor(0, dtype=type_.dtype()), + ) + delta_default = g.op( + "Constant", + value_t=torch.tensor(1, dtype=type_.dtype()), + ) + # pyrefly: ignore [bad-argument-type] + return g.op("Range", start_default, end, delta_default) + elif len(args) == 4 or len(args) == 7: + if len(args) == 4: + # aten::arange(Scalar start, Scalar end, Scalar step, Tensor out) + dtype = None + else: + # aten::arange(Scalar start, Scalar end, Scalar step, ScalarType dtype, Layout, Device, bool pin_memory) + dtype = _get_arange_dtype(args[3]) + _, end, start, step = symbolic_helper._arange_cast_helper( + g, start=args[0], end=args[1], step=args[2], dtype=dtype + ) + # pyrefly: ignore [bad-argument-type] + return g.op("Range", start, end, step) + elif len(args) == 6: + # aten::arange(Scalar start, Scalar end, ScalarType dtype, Layout, Device, bool pin_memory) + dtype = _get_arange_dtype(args[2]) + type_, end, start, step = symbolic_helper._arange_cast_helper( + g, start=args[0], end=args[1], dtype=dtype + ) + delta_default = g.op( + "Constant", + value_t=torch.tensor(1, dtype=type_.dtype()), + ) + # pyrefly: ignore [bad-argument-type] + return g.op("Range", start, end, delta_default) + else: + return symbolic_helper._unimplemented( + "aten::arange", f"with {len(args)} arguments" + ) + + +@_onnx_symbolic("aten::_dim_arange") +@symbolic_helper.parse_args("v", "i") +def _dim_arange(g: jit_utils.GraphContext, like, dim): + like_shape = g.op("Shape", like) + stop = g.op( + "Gather", like_shape, g.op("Constant", value_t=torch.tensor(dim)), axis_i=0 + ) + return arange(g, stop, 4, None, None, None) + + +@_onnx_symbolic("aten::size") +@symbolic_helper.quantized_args(True, quantize_output=False) +def size(g: jit_utils.GraphContext, self, dim=None): + if dim is None: + return g.op("Shape", self) + return symbolic_helper._size_helper(g, self, dim) + + +@_onnx_symbolic("aten::squeeze") +def squeeze(g: jit_utils.GraphContext, self, dim=None): + if dim is None: + return g.op("Squeeze", self) + + # dim as a tensor + if not symbolic_helper._is_constant(dim): + return symbolic_helper._squeeze_helper(g, self, [dim]) + + dim = symbolic_helper._get_const(dim, "i", "dim") + + input_rank = symbolic_helper._get_tensor_rank(self) + adjusted_dim = dim + if input_rank is not None and dim < 0: + adjusted_dim += input_rank + dim_size = symbolic_helper._get_tensor_dim_size(self, adjusted_dim) + if (dim < 0 and input_rank is None) or dim_size is None: + # If onnx shape inference is not on, export always as dynamic. + # Because we cannot tell if observed static shape is also static at runtime. + # create "cond" node (condition is shape[i]==1) + dim_constant = g.op("Constant", value_t=torch.tensor([dim])) + size = symbolic_helper._size_helper(g, self, dim_constant) + const_one = g.op("Constant", value_t=torch.ones(1, dtype=torch.int64)) + cond = g.op("Equal", size, const_one) + # create the "If" node and add the "then" and "else" blocks to it. + if_op, (if_context, else_context), _ = jit_utils.add_op_with_blocks( + g, "If", cond, n_blocks=2 + ) + squeeze_ = symbolic_helper._squeeze_helper(if_context, self, [dim]) + utils._add_output_to_block(if_context.block, squeeze_) + identity_ = else_context.op("Identity", self) + utils._add_output_to_block(else_context.block, identity_) + return if_op + + # For static input shape + dim = adjusted_dim + if dim_size > 1: + warnings.warn( + "This model contains a squeeze operation on dimension " + + str(dim) + + ". The size of " + + "this dimension in the given input is " + + str(dim_size) + + ". The model will " + + "be exported without the squeeze node. If the model is intended to be used with dynamic " + + "input shapes, please export with dynamic_axes argument.", + stacklevel=2, + ) + return self + return symbolic_helper._squeeze_helper(g, self, [dim]) + + +@_onnx_symbolic("aten::unsqueeze") +def unsqueeze(g: jit_utils.GraphContext, self, dim): + if symbolic_helper._is_constant(dim): + dim = symbolic_helper._get_const(dim, "i", "dim") + + return symbolic_helper._unsqueeze_helper(g, self, [dim]) + + +@_onnx_symbolic("aten::mm") +def mm(g: jit_utils.GraphContext, self, other): + return g.op("Gemm", self, other, beta_f=0.0, alpha_f=1.0) + + +@_onnx_symbolic("aten::index") +def index(g: jit_utils.GraphContext, self, index): + if symbolic_helper._is_packed_list(index): + indices = symbolic_helper._unpack_list(index) + else: + indices = [index] + + # Handle single mask index. + if len(indices) == 1: + index = indices[0] + if not symbolic_helper._is_none(index) and ( + symbolic_helper._is_bool(index) + or _type_utils.JitScalarType.from_value(index) + == _type_utils.JitScalarType.UINT8 + ): + index = opset9.nonzero(g, index) + return g.op("GatherND", self, index) + return opset9.index(g, self, index) + + +@_onnx_symbolic("aten::index_fill") +def index_fill(g: jit_utils.GraphContext, self, dim, index, value): + expanded_index_shape, expanded_index = symbolic_helper._index_fill_reshape_helper( + g, self, dim, index + ) + value = symbolic_helper._maybe_get_scalar(value) + value = symbolic_helper._if_scalar_type_as(value, self) + expanded_value = opset9.expand(g, value, expanded_index_shape, None) + return scatter(g, self, dim, expanded_index, expanded_value) + + +@_onnx_symbolic("aten::index_copy") +def index_copy(g: jit_utils.GraphContext, self, dim, index, source): + _expanded_index_shape, expanded_index = symbolic_helper._index_fill_reshape_helper( + g, self, dim, index + ) + return scatter(g, self, dim, expanded_index, source) + + +@_onnx_symbolic("aten::bitwise_right_shift") +@_onnx_symbolic("aten::__rshift_") +def __rshift_(g: jit_utils.GraphContext, self, other): + # make sure to cast other to self's type + # (when self is long, make sure that other is not float) + if _type_utils.JitScalarType.from_value( + other, _type_utils.JitScalarType.UNDEFINED + ) != _type_utils.JitScalarType.from_value(self): + other = g.op( + "Cast", + other, + to_i=_type_utils.JitScalarType.from_value(self).onnx_type(), + ) + + if ( + _type_utils.JitScalarType.from_value(self, _type_utils.JitScalarType.UNDEFINED) + == _type_utils.JitScalarType.UINT8 + ): + return g.op("BitShift", self, other, direction_s="RIGHT") + + two = g.op("Constant", value_t=torch.tensor(2, dtype=torch.float32)) + # exponent (same type as self) has to be float or double in onnx::Pow + if not symbolic_helper._is_fp(self): + other = g.op("Cast", other, to_i=_C_onnx.TensorProtoDataType.FLOAT) + two_pow = g.op("Pow", two, other) + two_pow = g.op( + "Cast", + two_pow, + to_i=_type_utils.JitScalarType.from_value(self).onnx_type(), + ) + rshift = g.op("Div", self, two_pow) + return rshift + + +@_onnx_symbolic("aten::bitwise_left_shift") +@_onnx_symbolic("aten::__lshift_") +def __lshift_(g: jit_utils.GraphContext, self, other): + # make sure to cast other to self's type + # (when self is long, make sure that other is not float) + if _type_utils.JitScalarType.from_value( + other, _type_utils.JitScalarType.UNDEFINED + ) != _type_utils.JitScalarType.from_value(self): + other = g.op( + "Cast", + other, + to_i=_type_utils.JitScalarType.from_value(self).onnx_type(), + ) + + if ( + _type_utils.JitScalarType.from_value(self, _type_utils.JitScalarType.UNDEFINED) + == _type_utils.JitScalarType.UINT8 + ): + return g.op("BitShift", self, other, direction_s="LEFT") + + two = g.op("Constant", value_t=torch.tensor(2, dtype=torch.float32)) + # exponent (same type as self) has to be float or double in onnx::Pow + if not symbolic_helper._is_fp(self): + other = g.op("Cast", other, to_i=_C_onnx.TensorProtoDataType.FLOAT) + two_pow = g.op("Pow", two, other) + two_pow = g.op( + "Cast", + two_pow, + to_i=_type_utils.JitScalarType.from_value(self).onnx_type(), + ) + lshift = g.op("Mul", self, two_pow) + return lshift + + +def _get_im2col_indices_along_dim( + g: jit_utils.GraphContext, input_d, kernel_size_d, dilation_d, padding_d, stride_d +): + # Input is always 4-D (N, C, H, W) + # Calculate indices of sliding blocks along spatial dimension + # Slide kernel over input each dim d: + # each dimension d ranges from 0 to input[d]+2xpadding[d]-dilation[d]x(kernel_size[d]-1) + # with steps = stride + + blocks_d = g.op( + "Add", input_d, g.op("Constant", value_t=torch.tensor(padding_d * 2)) + ) + blocks_d = g.op( + "Sub", + blocks_d, + g.op("Constant", value_t=torch.tensor(dilation_d * (kernel_size_d - 1))), + ) + + # Stride kernel over input and find starting indices along dim d + blocks_d_indices = g.op( + "Range", + g.op("Constant", value_t=torch.tensor(0)), + blocks_d, + g.op("Constant", value_t=torch.tensor(stride_d)), + ) + + # Apply dilation on kernel and find its indices along dim d + kernel_grid = torch.arange(0, kernel_size_d * dilation_d, dilation_d) + kernel_grid = g.op("Constant", value_t=kernel_grid.unsqueeze(0)) + + # Broadcast and add kernel staring positions (indices) with + # kernel_grid along dim d, to get block indices along dim d + blocks_d_indices = symbolic_helper._unsqueeze_helper( + g, blocks_d_indices, [0] + ) # Reshape to [1, -1] + kernel_mask = symbolic_helper._reshape_helper( + g, kernel_grid, g.op("Constant", value_t=torch.tensor([-1, 1])) + ) + block_mask = g.op("Add", blocks_d_indices, kernel_mask) + + return block_mask + + +def _get_im2col_padded_input(g: jit_utils.GraphContext, input, padding_h, padding_w): + # Input is always 4-D tensor (N, C, H, W) + # Padding tensor has the following format: (padding_h, padding_w) + # Reshape the padding to follow ONNX format: (dim1_begin, dim2_begin,...,dim1_end, dim2_end,...) + pad = g.op("Constant", value_t=torch.LongTensor([0, 0, padding_h, padding_w] * 2)) + return g.op("Pad", input, pad) + + +def _get_im2col_output_shape(g: jit_utils.GraphContext, input, kernel_h, kernel_w): + batch_dim = size(g, input, g.op("Constant", value_t=torch.tensor(0))) + channel_dim = size(g, input, g.op("Constant", value_t=torch.tensor(1))) + channel_unfolded = g.op( + "Mul", channel_dim, g.op("Constant", value_t=torch.tensor(kernel_h * kernel_w)) + ) + + return g.op( + "Concat", + symbolic_helper._unsqueeze_helper(g, batch_dim, [0]), + symbolic_helper._unsqueeze_helper(g, channel_unfolded, [0]), + g.op("Constant", value_t=torch.tensor([-1])), + axis_i=0, + ) + + +@_onnx_symbolic("aten::im2col") +@symbolic_helper.parse_args("v", "is", "is", "is", "is") +def im2col(g: jit_utils.GraphContext, input, kernel_size, dilation, padding, stride): + # Input is always 4-D tensor (N, C, H, W) + # All other args are int[2] + + input_h = size(g, input, g.op("Constant", value_t=torch.tensor(2))) + input_w = size(g, input, g.op("Constant", value_t=torch.tensor(3))) + + stride_h, stride_w = stride[0], stride[1] + padding_h, padding_w = padding[0], padding[1] + dilation_h, dilation_w = dilation[0], dilation[1] + kernel_h, kernel_w = kernel_size[0], kernel_size[1] + + blocks_row_indices = _get_im2col_indices_along_dim( + g, input_h, kernel_h, dilation_h, padding_h, stride_h + ) + blocks_col_indices = _get_im2col_indices_along_dim( + g, input_w, kernel_w, dilation_w, padding_w, stride_w + ) + + output_shape = _get_im2col_output_shape(g, input, kernel_h, kernel_w) + padded_input = _get_im2col_padded_input(g, input, padding_h, padding_w) + + # For a 4D matrix of size (1, 1, 3, 3) as below with kernel_size=2, stride=1, and dilation=1 + # [[[[1., 2., 3.,], + # [4., 5., 6.,], + # [7., 8., 9.,]]]] + # First gather indices along rows (dim=2) with blocks_row_indices = [[0,1], [1,2]] to get: + # [[[[[1., 2., 3.], + # [4., 5., 6.]], + # [[4., 5., 6.], + # [7., 8., 9.]]]]] + # And then gather along cols (dim=4) with blocks_row_indices = [[0,1], [1,2]] to get: + # [[[[[[1., 2.], + # [4., 5.]], + # [[2., 3.], + # [5., 6]]], + # [[[4., 5.], + # [7., 8.]], + # [[5., 6.], + # [8., 9.]]]]]] + # Transpose dims 3 (depth) and 4 (rows), and then reshape to output shape (1, 1, 4, 4) to get: + # [[[1., 2., 4., 5.], + # [2., 3., 5., 6.], + # [4., 5., 7., 8.], + # [5., 6., 8., 9.]]] + output = g.op("Gather", padded_input, blocks_row_indices, axis_i=2) + output = g.op("Gather", output, blocks_col_indices, axis_i=4) + output = g.op("Transpose", output, perm_i=[0, 1, 2, 4, 3, 5]) + return symbolic_helper._reshape_helper(g, output, output_shape) + + +@_onnx_symbolic("aten::narrow") +def narrow(g: jit_utils.GraphContext, input, dim, start, length): + end = g.op("Add", start, length) + return symbolic_helper._slice_helper(g, input, axes=dim, starts=start, ends=end) + + +@_onnx_symbolic("aten::flatten") +@symbolic_helper.quantized_args(True, False, False) +@symbolic_helper.parse_args("v", "i", "i") +def flatten(g: jit_utils.GraphContext, input, start_dim, end_dim): + dim = symbolic_helper._get_tensor_rank(input) + if dim == 1: + return input + # use ONNX's Flatten operator for cases where the output shape is 2D + if start_dim == 1: + if end_dim == -1 or (dim is not None and end_dim == dim - 1): + return g.op("Flatten", input, axis_i=start_dim) + elif start_dim == 0: + if end_dim == -2 or (dim is not None and end_dim == dim - 2): + return g.op("Flatten", input, axis_i=end_dim + 1) + if dim is None: + return symbolic_helper._unimplemented( + "dim", + "ONNX and PyTorch use different strategies to split the input. " + "Input rank must be known at export time.", + ) + # if end_dim is negative add dim + if end_dim < 0: + end_dim = dim + end_dim + + return symbolic_helper._flatten_helper(g, input, start_dim, end_dim, dim) + + +@_onnx_symbolic("aten::linalg_vector_norm") +@symbolic_helper.parse_args("v", "f", "is", "b", "v") +def linalg_vector_norm( + g: jit_utils.GraphContext, + self, + ord, + dim: Sequence[int] | None, + keepdim: bool, + dtype, +): + return symbolic_helper._linalg_vector_norm_helper(g, self, ord, dim, keepdim, dtype) + + +@_onnx_symbolic("aten::embedding_bag") +@symbolic_helper.parse_args("v", "v", "v", "i", "i", "i", "v", "i", "i") +def embedding_bag( + g: jit_utils.GraphContext, + embedding_matrix, + indices, + offsets, + scale_grad_by_freq, + mode, + sparse, + per_sample_weights, + include_last_offset, + padding_idx, +): + return symbolic_helper._embedding_bag_helper( + g, + embedding_matrix, + indices, + offsets, + scale_grad_by_freq, + mode, + sparse, + per_sample_weights, + include_last_offset, + padding_idx, + ) + + +@_onnx_symbolic("aten::embedding_renorm") +@symbolic_helper.parse_args("v", "v", "f", "f") +def embedding_renorm(g: jit_utils.GraphContext, weight, indices, max_norm, norm_type): + unique_indices = g.op("Unique", indices) + partial_weight = g.op("Gather", weight, unique_indices) + norm_i = int(norm_type) + if norm_i == 1: + norm_type = "ReduceL1" + elif norm_i == 2: + norm_type = "ReduceL2" + else: + raise errors.SymbolicValueError( + f"Unsupported: ONNX export of embedding_renorm with norm: {norm_i}. " + "Only 1. and 2. are supported.", + weight, + ) + partial_weight_norm = g.op(norm_type, partial_weight, axes_i=[1], keepdims_i=1) + # https://github.com/pytorch/pytorch/blob/0a07488ed2c47765e337e290bd138c0e6e459cbd/aten/src/ATen/native/Embedding.cpp#L177 + # Add 1e-7 to prevent division by zero. + partial_weight_norm_ = g.op( + "Add", partial_weight_norm, g.op("Constant", value_t=torch.tensor(1e-7)) + ) + max_norm = torch.tensor(max_norm) + scales = g.op("Div", max_norm, partial_weight_norm_) + partial_weight_renorm = g.op("Mul", partial_weight, scales) + partial_weight_renorm = g.op( + "Where", + g.op("Greater", partial_weight_norm, max_norm), + partial_weight_renorm, + partial_weight, + ) + return g.op( + "ScatterND", + weight, + symbolic_helper._unsqueeze_helper(g, unique_indices, [1]), + partial_weight_renorm, + ) + + +@_onnx_symbolic("aten::chunk") +def chunk(g: jit_utils.GraphContext, self, chunks, dim): + # Calculate chunk size for dynamic chunk + dim_size = g.op("Gather", g.op("Shape", self), dim, axis_i=0) + chunk_size_s = g.op( + "Sub", chunks, g.op("Constant", value_t=torch.tensor([1], dtype=torch.long)) + ) + chunk_size = g.op("Div", g.op("Add", dim_size, chunk_size_s), chunks) + # Create splits vector + chunk_vec = [ + opset9.expand(g, chunk_size, chunk_size_s, None), + g.op("Sub", dim_size, g.op("Mul", chunk_size, chunk_size_s)), + ] + chunk_vec = g.op("Concat", *chunk_vec, axis_i=0) + return split(g, self, chunk_vec, dim) + + +@_onnx_symbolic("aten::normal") +def normal( + g: jit_utils.GraphContext, + mean, + std, + sizes=None, + generator=None, + dtype=None, + layout=None, + device=None, + pin_memory=None, +): + # If you can sample from a given distribution with mean 0 and variance 1, then you can easily sample from a + # scale-location transformation of that distribution, which has mean mu and variance sigma's square. If x is a sample + # from a mean 0 and variance 1 distribution then + # sigma x+mu + # is a sample with mean mu and variance sigma's square. + if sizes is not None and not symbolic_helper._is_none(sizes): + mean = opset9.expand(g, mean, sizes, None) + result = opset9.mul(g, std, g.op("RandomNormalLike", mean)) + return add(g, result, mean) + + +@_onnx_symbolic("aten::atleast_1d") +def atleast_1d(g: jit_utils.GraphContext, self: torch._C.Value): + # NOTE: If it's 0D, reshape to 1D + + # NOTE: self could be a packed list or a tensor + if symbolic_helper._is_value(self) and symbolic_helper._is_packed_list(self): + tensor_list = symbolic_helper._unpack_list(self) + new_tensor_list = [] + for tensor in tensor_list: + new_tensor = tensor + tensor_rank = symbolic_helper._get_tensor_rank(tensor) + if tensor_rank == 0: + new_tensor = symbolic_helper._reshape_helper( + g, new_tensor, g.op("Constant", value_t=torch.tensor([1])) + ) + new_tensor_list.append(new_tensor) + return g.op("SequenceConstruct", *new_tensor_list) + + tensor_rank = symbolic_helper._get_tensor_rank(self) + if tensor_rank == 0: + self = symbolic_helper._reshape_helper( + g, self, g.op("Constant", value_t=torch.tensor([1])) + ) + return self + + +@_onnx_symbolic("aten::atleast_2d") +def atleast_2d(g: jit_utils.GraphContext, self: torch._C.Value): + # NOTE: If it's 0D, reshape to 2D + # If it's 1D, unsqueeze to 2D + + # NOTE: self could be a packed list or a tensor + if symbolic_helper._is_value(self) and symbolic_helper._is_packed_list(self): + tensor_list = symbolic_helper._unpack_list(self) + new_tensor_list = [] + for tensor in tensor_list: + new_tensor = tensor + tensor_rank = symbolic_helper._get_tensor_rank(tensor) + if tensor_rank == 0: + new_tensor = symbolic_helper._reshape_helper( + g, new_tensor, g.op("Constant", value_t=torch.tensor([1, 1])) + ) + elif tensor_rank == 1: + new_tensor = symbolic_helper._unsqueeze_helper( + g, new_tensor, axes_i=[0] + ) + new_tensor_list.append(new_tensor) + return g.op("SequenceConstruct", *new_tensor_list) + + tensor_rank = symbolic_helper._get_tensor_rank(self) + if tensor_rank == 0: + self = symbolic_helper._reshape_helper( + g, self, g.op("Constant", value_t=torch.tensor([1, 1])) + ) + elif tensor_rank == 1: + self = symbolic_helper._unsqueeze_helper(g, self, axes_i=[0]) + return self + + +@_onnx_symbolic("aten::atleast_3d") +def atleast_3d(g: jit_utils.GraphContext, self: torch._C.Value): + # NOTE: If it's 0D, reshape to 3D + # If it's 1D, unsqueeze to 3D + # If it's 2D, unsqueeze to 3D + + # NOTE: self could be a packed list or a tensor + if symbolic_helper._is_value(self) and symbolic_helper._is_packed_list(self): + tensor_list = symbolic_helper._unpack_list(self) + new_tensor_list = [] + for tensor in tensor_list: + new_tensor = tensor + tensor_rank = symbolic_helper._get_tensor_rank(tensor) + if tensor_rank == 0: + new_tensor = symbolic_helper._reshape_helper( + g, new_tensor, g.op("Constant", value_t=torch.tensor([1, 1, 1])) + ) + elif tensor_rank == 1: + new_tensor = symbolic_helper._unsqueeze_helper( + g, new_tensor, axes_i=[0] + ) + new_tensor = symbolic_helper._unsqueeze_helper( + g, new_tensor, axes_i=[-1] + ) + elif tensor_rank == 2: + new_tensor = symbolic_helper._unsqueeze_helper( + g, new_tensor, axes_i=[-1] + ) + new_tensor_list.append(new_tensor) + return g.op("SequenceConstruct", *new_tensor_list) + + tensor_rank = symbolic_helper._get_tensor_rank(self) + if tensor_rank == 0: + self = symbolic_helper._reshape_helper( + g, self, g.op("Constant", value_t=torch.tensor([1, 1, 1])) + ) + elif tensor_rank == 1: + self = symbolic_helper._unsqueeze_helper(g, self, axes_i=[0]) + self = symbolic_helper._unsqueeze_helper(g, self, axes_i=[-1]) + elif tensor_rank == 2: + self = symbolic_helper._unsqueeze_helper(g, self, axes_i=[-1]) + return self + + +@_onnx_symbolic("prim::ConstantChunk") +def prim_constant_chunk(g: jit_utils.GraphContext, self, chunks, dim): + input_shape = g.op("Shape", self) + axis = g.op("Constant", value_t=torch.tensor([dim], dtype=torch.long)) + input_shape_dim = g.op("Gather", input_shape, axis, axis_i=0) + start = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long)) + chunk_size = g.op("Constant", value_t=torch.tensor([chunks], dtype=torch.long)) + chunk_size_minus_1 = g.op( + "Constant", value_t=torch.tensor([chunks - 1], dtype=torch.long) + ) + input_shape_dim_shift = g.op("Add", input_shape_dim, chunk_size_minus_1) + chunk_dim = g.op("Div", input_shape_dim_shift, chunk_size) + res = [] + for i in range(chunks): + index = g.op("Constant", value_t=torch.tensor([i + 1], dtype=torch.long)) + end = g.op("Mul", chunk_dim, index) + res.append(g.op("Slice", self, start, end, axis)) + start = end + return res + + +@_onnx_symbolic("aten::hstack") +def hstack(g: jit_utils.GraphContext, tensor_list: _C.Value): + tensor_list = atleast_1d(g, tensor_list) + first_tensor = g.op( + "SequenceAt", + tensor_list, + g.op("Constant", value_t=torch.tensor(0, dtype=torch.long)), + ) + first_tensor_shape = g.op("Shape", first_tensor) + first_tensor_dim = g.op("Size", first_tensor_shape) + + const_one = g.op("Constant", value_t=torch.tensor(1, dtype=torch.long)) + equal_to_one = g.op("Equal", first_tensor_dim, const_one) + + ( + if_op_greater, + (if_context_equal, else_context_equal), + _, + ) = jit_utils.add_op_with_blocks(g, "If", equal_to_one, n_blocks=2, outputs=1) + result_if = if_context_equal.op( + "ConcatFromSequence", tensor_list, axis_i=0, new_axis_i=0 + ) + utils._add_output_to_block(if_context_equal.block, result_if) + result_else = else_context_equal.op( + "ConcatFromSequence", tensor_list, axis_i=1, new_axis_i=0 + ) + utils._add_output_to_block(else_context_equal.block, result_else) + result = if_op_greater.node().output() + + return result + + +@_onnx_symbolic("aten::vstack") +def vstack(g: jit_utils.GraphContext, tensor_list: _C.Value): + tensor_list = atleast_2d(g, tensor_list) + return g.op("ConcatFromSequence", tensor_list, axis_i=0, new_axis_i=0) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset12.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset12.py new file mode 100644 index 0000000000000000000000000000000000000000..61136b76b33cb4eb61850c8fa7e67883da9f1848 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset12.py @@ -0,0 +1,469 @@ +# mypy: allow-untyped-defs +# mypy: disable-error-code=arg-type +from __future__ import annotations + +import functools +import sys + +import torch +from torch._C import _onnx as _C_onnx +from torch.onnx import errors +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, + symbolic_opset9 as opset9, + utils, +) + + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in README.md + +# This file exports ONNX ops for opset 12 + +__all__ = [ + "argmax", + "argmin", + "binary_cross_entropy_with_logits", + "celu", + "cross_entropy_loss", + "dropout", + "einsum", + "ge", + "le", + "native_dropout", + "nll_loss", + "nll_loss2d", + "nll_loss_nd", + "outer", + "pow", + "tensordot", + "unfold", +] + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=12) + + +def _einsum_helper(g: jit_utils.GraphContext, equation, tensors): + if not tensors: + raise RuntimeError("Einsum inputs are empty.") + # ONNX does not support bool for Einsum inputs. + if symbolic_helper._is_bool(tensors[0]): + tensors = [ + g.op("Cast", tensor, to_i=_C_onnx.TensorProtoDataType.INT64) + for tensor in tensors + ] + return g.op( + "Cast", + g.op("Einsum", *tensors, equation_s=equation), + to_i=_C_onnx.TensorProtoDataType.BOOL, + ) + else: + return g.op("Einsum", *tensors, equation_s=equation) + + +@_onnx_symbolic("aten::einsum") +@symbolic_helper.parse_args("s", "v", "is") +def einsum(g: jit_utils.GraphContext, equation, tensor_list, path=None): + tensors = symbolic_helper._unpack_list(tensor_list) + return _einsum_helper(g, equation, tensors) + + +@_onnx_symbolic("aten::outer") +@symbolic_helper.parse_args("v", "v") +def outer(g: jit_utils.GraphContext, input, other): + # make sure to cast other to self's type + if _type_utils.JitScalarType.from_value( + other, _type_utils.JitScalarType.UNDEFINED + ) != _type_utils.JitScalarType.from_value(input): + other = g.op( + "Cast", + other, + to_i=_type_utils.JitScalarType.from_value(input).onnx_type(), + ) + return _einsum_helper(g, "i,j->ij", [input, other]) + + +def _dropout_returns_masked_input_and_mask( + g: jit_utils.GraphContext, input: torch._C.Value, p: float, train: bool +) -> tuple[torch._C.Value, torch._C.Value | None]: + symbolic_helper.check_training_mode(train, "dropout") + # In eval mode, dropout is non-op. That is, if the node's + # train param is set to False, dropout just returns its inputs. + if not train: + return input, None + p = g.op("Constant", value_t=torch.tensor(p)) + t = g.op("Constant", value_t=torch.tensor(train, dtype=torch.bool)) + r, mask = g.op("Dropout", input, p, t, outputs=2) + return r, mask + + +@_onnx_symbolic("aten::dropout") +@symbolic_helper.parse_args("v", "f", "b") +def dropout(g: jit_utils.GraphContext, input, p, train): + masked, _ = _dropout_returns_masked_input_and_mask(g, input, p, train) + return masked + + +@_onnx_symbolic("aten::native_dropout") +@symbolic_helper.parse_args("v", "f", "b") +def native_dropout(g: jit_utils.GraphContext, input, p, train): + return _dropout_returns_masked_input_and_mask(g, input, p, train) + + +@_onnx_symbolic("aten::nll_loss") +def nll_loss(g: jit_utils.GraphContext, self, target, weight, reduction, ignore_index): + # none reduction : onnx::Constant[value={0}] + # mean reduction : onnx::Constant[value={1}] + # sum reduction : onnx::Constant[value={2}] + reduction = symbolic_helper._maybe_get_const(reduction, "i") + reduction_vals = ["none", "mean", "sum"] + reduction = reduction_vals[reduction] + + # in onnx NegativeLogLikelihoodLoss specification, ignore_index is optional without default value. + # therefore we need to set ignore_index attribute even if it is not specified (e.g. ignore_index=-100). + ignore_index = symbolic_helper._maybe_get_const(ignore_index, "i") + if weight.node().mustBeNone(): + nllloss = g.op( + "NegativeLogLikelihoodLoss", + self, + target, + reduction_s=reduction, + ignore_index_i=ignore_index, + ) + else: + nllloss = g.op( + "NegativeLogLikelihoodLoss", + self, + target, + weight, + reduction_s=reduction, + ignore_index_i=ignore_index, + ) + + return nllloss + + +@_onnx_symbolic("aten::nll_loss2d") +def nll_loss2d( + g: jit_utils.GraphContext, self, target, weight, reduction, ignore_index +): + return nll_loss(g, self, target, weight, reduction, ignore_index) + + +@_onnx_symbolic("aten::nll_loss_nd") +def nll_loss_nd( + g: jit_utils.GraphContext, self, target, weight, reduction, ignore_index +): + return nll_loss(g, self, target, weight, reduction, ignore_index) + + +@_onnx_symbolic("aten::cross_entropy_loss") +def cross_entropy_loss( + g: jit_utils.GraphContext, + self, + target, + weight, + reduction, + ignore_index, + label_smoothing, +): + # none reduction : onnx::Constant[value={0}] + # mean reduction : onnx::Constant[value={1}] + # sum reduction : onnx::Constant[value={2}] + reduction = symbolic_helper._maybe_get_const(reduction, "i") + reduction_vals = ["none", "mean", "sum"] + reduction = reduction_vals[reduction] + + label_smoothing = symbolic_helper._maybe_get_const(label_smoothing, "f") + if label_smoothing is not None and label_smoothing > 0.0: + raise errors.SymbolicValueError( + "Unsupported: ONNX does not support label_smoothing", self + ) + + # in onnx SoftmaxCrossEntropyLoss specification, ignore_index is optional without default value. + # therefore we need to set ignore_index attribute even if it is not specified (e.g. ignore_index=-100). + ignore_index = symbolic_helper._maybe_get_const(ignore_index, "i") + if weight.node().mustBeNone(): + celoss = g.op( + "SoftmaxCrossEntropyLoss", + self, + target, + reduction_s=reduction, + ignore_index_i=ignore_index, + ) + else: + celoss = g.op( + "SoftmaxCrossEntropyLoss", + self, + target, + weight, + reduction_s=reduction, + ignore_index_i=ignore_index, + ) + + return celoss + + +@_onnx_symbolic("aten::binary_cross_entropy_with_logits") +@symbolic_helper.parse_args("v", "v", "v", "v", "i") +def binary_cross_entropy_with_logits( + g: jit_utils.GraphContext, input, target, weight, pos_weight, reduction +): + p = g.op("Constant", value_t=torch.tensor([1])) + sig_x = opset9.sigmoid(g, input) + log_sig_x = opset9.log(g, sig_x) + sub_1_x = opset9.sub(g, p, sig_x) + sub_1_y = opset9.sub(g, p, target) + log_1_x = opset9.log(g, sub_1_x) + if pos_weight is None or symbolic_helper._is_none(pos_weight): + output = opset9.neg( + g, + opset9.add( + g, opset9.mul(g, target, log_sig_x), opset9.mul(g, sub_1_y, log_1_x) + ), + ) + else: + output = opset9.neg( + g, + opset9.add( + g, + opset9.mul(g, opset9.mul(g, target, log_sig_x), pos_weight), + opset9.mul(g, sub_1_y, log_1_x), + ), + ) + + if weight is not None and not symbolic_helper._is_none(weight): + output = opset9.mul(g, weight, output) + + reduction = symbolic_helper._maybe_get_const(reduction, "i") + if reduction == 0: + return output + elif reduction == 1: + return g.op("ReduceMean", output, keepdims_i=0) + elif reduction == 2: + return g.op("ReduceSum", output, keepdims_i=0) + else: + return symbolic_helper._onnx_unsupported( + "binary_cross_entropy_with_logits with reduction other than none, mean, or sum", + input, + ) + + +@_onnx_symbolic("aten::celu") +def celu(g: jit_utils.GraphContext, self, alpha): + alpha = symbolic_helper._maybe_get_const(alpha, "f") + # if the input is of type double cast it to float + if ( + _type_utils.JitScalarType.from_value(self, _type_utils.JitScalarType.UNDEFINED) + == _type_utils.JitScalarType.DOUBLE + ): + self = g.op("Cast", self, to_i=_C_onnx.TensorProtoDataType.FLOAT) + out = g.op("Celu", self, alpha_f=alpha) + return g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.DOUBLE) + + return g.op("Celu", self, alpha_f=alpha) + + +@_onnx_symbolic("aten::argmax") +@symbolic_helper.parse_args("v", "v", "b") +def argmax( + g: jit_utils.GraphContext, + input: torch._C.Value, + dim: torch._C.Value, + keepdim: bool, +): + return symbolic_helper._argmin_argmax_helper(g, input, dim, keepdim, "ArgMax") + + +@_onnx_symbolic("aten::argmin") +@symbolic_helper.parse_args("v", "v", "b") +def argmin( + g: jit_utils.GraphContext, + input: torch._C.Value, + dim: torch._C.Value, + keepdim: bool, +): + return symbolic_helper._argmin_argmax_helper(g, input, dim, keepdim, "ArgMin") + + +@_onnx_symbolic("aten::pow") +def pow(g: jit_utils.GraphContext, self, exponent): + return g.op("Pow", self, exponent) + + +@_onnx_symbolic("aten::ge") +def ge(g: jit_utils.GraphContext, input, other): + return g.op("GreaterOrEqual", input, other) + + +@_onnx_symbolic("aten::le") +def le(g: jit_utils.GraphContext, input, other): + return g.op("LessOrEqual", input, other) + + +@_onnx_symbolic("aten::unfold") +@symbolic_helper.parse_args("v", "i", "v", "v") +def unfold(g: jit_utils.GraphContext, input, dimension, size, step): + const_size = symbolic_helper._maybe_get_const(size, "i") + const_step = symbolic_helper._maybe_get_const(step, "i") + if not symbolic_helper._is_value(const_size) and not symbolic_helper._is_value( + const_step + ): + return opset9.unfold(g, input, dimension, const_size, const_step) + + sizedim = symbolic_helper._get_tensor_dim_size(input, dimension) + if sizedim is not None: + low_start = g.op("Constant", value_t=torch.tensor(0)) + low_end = g.op("Constant", value_t=torch.tensor(sizedim)) + hi_end = g.op("Constant", value_t=torch.tensor(sizedim + 1)) + low_indices = g.op("Range", low_start, low_end, step) + hi_indices = g.op("Range", size, hi_end, step) + + low_size = symbolic_helper._size_helper( + g, low_indices, g.op("Constant", value_t=torch.tensor(0)) + ) + hi_size = symbolic_helper._size_helper( + g, hi_indices, g.op("Constant", value_t=torch.tensor(0)) + ) + + ndim = symbolic_helper._get_tensor_rank(input) + if ndim is None: + raise AssertionError("ndim must be non-None") + perm = list(range(ndim)) + perm.append(perm.pop(dimension)) + + unsqueeze_list = [] + loop_condition = g.op("Constant", value_t=torch.tensor(1)) + loop_condition = g.op( + "Cast", loop_condition, to_i=_C_onnx.TensorProtoDataType.BOOL + ) + loop_len = g.op("Min", low_size, hi_size) + + loop, (loop_context,), _ = jit_utils.add_op_with_blocks( + g, "Loop", loop_len, loop_condition, n_blocks=1 + ) + + loop_block = loop_context.block + block_input_iter = utils._add_input_to_block(loop_block) + cond = utils._add_input_to_block(loop_block) # noqa: F841 + + starts = loop_context.op("Gather", low_indices, block_input_iter) + ends = loop_context.op("Gather", hi_indices, block_input_iter) + axes = loop_context.op("Constant", value_t=torch.tensor([2])) + starts = symbolic_helper._unsqueeze_helper(loop_context, starts, [0]) + ends = symbolic_helper._unsqueeze_helper(loop_context, ends, [0]) + stack = loop_context.op("Slice", input, starts, ends, axes) + + unsqueeze = symbolic_helper._unsqueeze_helper( + loop_context, loop_context.op("Transpose", stack, perm_i=perm), [dimension] + ) + unsqueeze_list.append(unsqueeze) + concat = loop_context.op("Concat", *unsqueeze_list, axis_i=0) + + cond_out = loop_context.op( + "Cast", + loop_condition, + # pyrefly: ignore [bad-argument-type] + _C_onnx.TensorProtoDataType.BOOL, + ) + utils._add_output_to_block(loop_block, cond_out) + utils._add_output_to_block(loop_block, concat) + + loop_output = loop.node().output() + perm = [0, 1, 2, 3, 4] + perm[0], perm[dimension + 1] = perm[dimension + 1], perm[0] + transpose = g.op("Transpose", loop_output, perm_i=perm) + squeeze = symbolic_helper._squeeze_helper(g, transpose, [0]) + + return squeeze + + return symbolic_helper._unimplemented("Unfold", "input size not accessible") + + +@_onnx_symbolic("aten::tensordot") +@symbolic_helper.parse_args("v", "v", "is", "is", "v") +def tensordot(g: jit_utils.GraphContext, input_a, input_b, dims_a, dims_b, out=None): + if out is not None: + symbolic_helper._unimplemented( + "Tensordot", "Out parameter is not supported for tensordot." + ) + + dim_count_a = symbolic_helper._get_tensor_rank(input_a) + if dim_count_a is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of tensordot for tensor(input_a) of unknown rank.", + input_a, + ) + + dim_count_b = symbolic_helper._get_tensor_rank(input_b) + if dim_count_b is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of tensordot for tensor(input_b) of unknown rank.", + input_b, + ) + + dims_a = [ + (dims_a[i] + dim_count_a) if (dims_a[i] < 0) else dims_a[i] + for i in range(len(dims_a)) + ] + dims_b = [ + (dims_b[i] + dim_count_b) if (dims_b[i] < 0) else dims_b[i] + for i in range(len(dims_b)) + ] + + left_dims_a = [i for i in range(dim_count_a) if (i not in dims_a)] + left_dims_b = [i for i in range(dim_count_b) if (i not in dims_b)] + + new_input_a = opset9.permute(g, input_a, left_dims_a + dims_a) + new_input_b = opset9.permute(g, input_b, dims_b + left_dims_b) + + input_shape = g.op("Shape", new_input_a) + left_sizes_a = symbolic_helper._slice_helper( + g, input_shape, axes=[0], starts=[0], ends=[len(left_dims_a)] + ) + shape_sizes = [ + left_sizes_a, + g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), + ] + output_a = opset9._reshape_from_tensor(g, new_input_a, shape_sizes) + + input_shape = g.op("Shape", output_a) + slices = symbolic_helper._slice_helper( + g, input_shape, axes=[0], starts=[-1], ends=[sys.maxsize] + ) + shape_sizes = [ + g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), + slices, + ] + output_a = opset9._reshape_from_tensor(g, new_input_a, shape_sizes) + + input_shape = g.op("Shape", new_input_b) + left_sizes_b = symbolic_helper._slice_helper( + g, input_shape, axes=[0], starts=[len(dims_b)], ends=[sys.maxsize] + ) + slices = symbolic_helper._slice_helper( + g, input_shape, axes=[0], starts=[0], ends=[len(dims_b)] + ) + shape_sizes = [ + slices, + g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), + ] + output_b = opset9._reshape_from_tensor(g, new_input_b, shape_sizes) + + input_shape = g.op("Shape", output_b) + slices = symbolic_helper._slice_helper( + g, input_shape, axes=[0], starts=[-1], ends=[sys.maxsize] + ) + shape_sizes = [ + g.op("Constant", value_t=torch.tensor([-1], dtype=torch.long)), + slices, + ] + output_b = opset9._reshape_from_tensor(g, new_input_b, shape_sizes) + + output = einsum(g, "ij,jk->ik", g.op("prim::ListConstruct", *[output_a, output_b])) + + shape_sizes = [left_sizes_a, left_sizes_b] + return opset9._reshape_from_tensor(g, output, shape_sizes) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset13.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset13.py new file mode 100644 index 0000000000000000000000000000000000000000..a03c41081221b90ff2e67af790ad6e7c4f3a5b35 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset13.py @@ -0,0 +1,1118 @@ +# mypy: allow-untyped-defs +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in README.md + +# This file exports ONNX ops for opset 13 +import functools + +import torch +import torch._C._onnx as _C_onnx +from torch.onnx import _constants, errors +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, + symbolic_opset11 as opset11, + symbolic_opset9 as opset9, + utils, +) + + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=13) + + +@_onnx_symbolic("aten::softmax") +@symbolic_helper.parse_args("v", "i", "none") +def softmax(g: jit_utils.GraphContext, input, dim, dtype=None): + softmax = g.op("Softmax", input, axis_i=dim) + if dtype and dtype.node().kind() != "prim::Constant": + parsed_dtype = symbolic_helper._get_const(dtype, "i", "dtype") + softmax = g.op( + "Cast", softmax, to_i=_type_utils.JitScalarType(parsed_dtype).onnx_type() + ) + + return softmax + + +@_onnx_symbolic("aten::log_softmax") +@symbolic_helper.parse_args("v", "i", "none") +def log_softmax(g: jit_utils.GraphContext, input, dim, dtype=None): + return_op = g.op("LogSoftmax", input, axis_i=dim) + if dtype and dtype.node().kind() != "prim::Constant": + parsed_dtype = symbolic_helper._get_const(dtype, "i", "dtype") + return_op = g.op( + "Cast", return_op, to_i=_type_utils.JitScalarType(parsed_dtype).onnx_type() + ) + return return_op + + +@_onnx_symbolic("aten::frobenius_norm") +@symbolic_helper.parse_args("v", "v", "i") +def frobenius_norm(g: jit_utils.GraphContext, self, dim=None, keepdim=False): + dim_val = symbolic_helper._maybe_get_const(dim, "is") + if not symbolic_helper._is_value(dim_val) and len(dim_val) == 0: + return g.op("ReduceL2", self, keepdims_i=0) + sqr = g.op("Mul", self, self) + sumsqr = symbolic_helper._reducesum_helper(g, sqr, dim, keepdims_i=keepdim) + return g.op("Sqrt", sumsqr) + + +@_onnx_symbolic("aten::split") +@symbolic_helper.parse_args("v", "v", "i", "i") +def split(g: jit_utils.GraphContext, self, split_size_or_sizes, dim, _outputs=None): + if not symbolic_helper._is_split_static(split_size_or_sizes, _outputs): + split_out = g.op("SplitToSequence", self, split_size_or_sizes, axis_i=dim) + if _outputs is None: + return split_out + # Convert to multiple slice nodes iff number of splits and number of outputs are statically known. + if ( + symbolic_helper._is_packed_list(split_size_or_sizes) + and len(symbolic_helper._unpack_list(split_size_or_sizes)) == _outputs + ): + split_sizes = [ + symbolic_helper._unsqueeze_helper(g, v, [0]) + for v in symbolic_helper._unpack_list(split_size_or_sizes) + ] + + start = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long)) + axis = g.op("Constant", value_t=torch.tensor([dim], dtype=torch.long)) + res = [] + for i in range(_outputs): + end = g.op( + "Add", start, split_sizes[i] + ) # split_sizes is a list of same length as _outputs + res.append(g.op("Slice", self, start, end, axis)) + start = end + return res + return [ + g.op( + "SequenceAt", + split_out, + g.op("Constant", value_t=torch.tensor([i], dtype=torch.long)), + ) + for i in range(_outputs) + ] + + split_val = symbolic_helper._node_get(split_size_or_sizes.node(), "value") + if split_val.dim() > 0: + # pyrefly: ignore [bad-argument-type] + return g.op("Split", self, split_size_or_sizes, axis_i=dim, outputs=_outputs) + split_size = symbolic_helper._get_const(split_size_or_sizes, "i", "split_size") + + size = symbolic_helper._get_tensor_dim_size(self, dim) + if size is None: + if _outputs is not None: + size = split_size * _outputs + else: + raise errors.SymbolicValueError( + "Unknown dimension size not supported", self + ) + splits = [split_size] * (size // split_size) + leftover = size % split_size + if leftover: + splits.append(leftover) + splits = g.op("Constant", value_t=torch.tensor(splits)) + # pyrefly: ignore [bad-argument-type] + return g.op("Split", self, splits, axis_i=dim, outputs=_outputs) + + +@_onnx_symbolic("aten::split_with_sizes") +def split_with_sizes(g: jit_utils.GraphContext, self, split_sizes, dim, _outputs=None): + return split(g, self, split_sizes, dim, _outputs) + + +@_onnx_symbolic("aten::unsafe_split") +def unsafe_split( + g: jit_utils.GraphContext, self, split_size_or_sizes, dim, _outputs=None +): + return split(g, self, split_size_or_sizes, dim, _outputs) + + +@_onnx_symbolic("aten::unsafe_split_with_sizes") +def unsafe_split_with_sizes( + g: jit_utils.GraphContext, self, split_sizes, dim, _outputs=None +): + return split_with_sizes(g, self, split_sizes, dim, _outputs) + + +@_onnx_symbolic("aten::tensor_split") +@symbolic_helper.parse_args("v", "v", "i", "i") +def tensor_split( + g: jit_utils.GraphContext, self, indices_or_sections, dim, _outputs=None +): + axis = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long)) + axis = opset11.unsqueeze(g, axis, 0) + const_1 = g.op("Constant", value_t=torch.tensor(1, dtype=torch.long)) + + if symbolic_helper._is_split_static(indices_or_sections, _outputs): + split_val = symbolic_helper._node_get(indices_or_sections.node(), "value") + + if split_val.dim() > 0: + start = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long)) + res = [] + if _outputs is None: + raise AssertionError("_outputs must be non-None") + for i in range(_outputs - 1): + end = g.op( + "Gather", + indices_or_sections, + g.op("Constant", value_t=torch.tensor([i], dtype=torch.long)), + axis_i=0, + ) + res.append(g.op("Slice", self, start, end, axis)) + start = end + + end = symbolic_helper._size_helper(g, self, axis) + res.append(g.op("Slice", self, start, end, axis)) + return res + + split_size = symbolic_helper._get_const( + indices_or_sections, "i", "indices_or_sections" + ) + + size = symbolic_helper._get_tensor_dim_size(self, dim) + if size is None: + if _outputs is not None: + size = split_size * _outputs + else: + raise errors.SymbolicValueError( + "Unknown dimension size not supported", self + ) + + min_split_size = size // split_size + num_splits_one_extra = size % split_size + + splits = num_splits_one_extra * [min_split_size + 1] + leftover = (split_size - num_splits_one_extra) * [min_split_size] + + splits = g.op( + "Constant", value_t=torch.tensor(splits + leftover, dtype=torch.long) + ) + # pyrefly: ignore [bad-argument-type] + return g.op("Split", self, splits, axis_i=dim, outputs=_outputs) + + if ( + symbolic_helper._is_tensor(indices_or_sections) + and symbolic_helper._get_tensor_rank(indices_or_sections) == 1 + ): + loop_len = symbolic_helper._size_helper( + g, indices_or_sections, g.op("Constant", value_t=torch.tensor(0)) + ) + loop_len = opset11.unsqueeze(g, loop_len, 0) + loop_condition = g.op("Cast", const_1, to_i=_C_onnx.TensorProtoDataType.BOOL) + + # To make the first slice in the below loop work, + # we pad a zero to the first position so that it will be the initial start of slice. + padding_0 = g.op("Constant", value_t=torch.tensor([0], dtype=torch.long)) + indices_or_sections = g.op("Concat", padding_0, indices_or_sections, axis_i=0) + + final_splits = g.op("SequenceEmpty") + # Loop inputs + loop, (loop_context,), _ = jit_utils.add_op_with_blocks( + g, "Loop", loop_len, loop_condition, final_splits, outputs=1, n_blocks=1 + ) + + loop_block = loop_context.block + block_input_iter = utils._add_input_to_block(loop_block) + cond = utils._add_input_to_block(loop_block) # noqa: F841 + final_splits = utils._add_input_to_block(loop_block) + + start = loop_context.op( + "Gather", indices_or_sections, block_input_iter, axis_i=0 + ) + end = loop_context.op( + "Gather", + indices_or_sections, + loop_context.op("Add", block_input_iter, const_1), + axis_i=0, + ) + + slice = loop_context.op("Slice", self, start, end, axis) + final_splits = loop_context.op("SequenceInsert", final_splits, slice) + + # Loop outputs + cond_out = loop_context.op("Identity", loop_condition) + utils._add_output_to_block(loop_block, cond_out) + utils._add_output_to_block(loop_block, final_splits) + + loop_out = loop.node().output() + start = g.op( + "Gather", + indices_or_sections, + g.op("Constant", value_t=torch.tensor(-1, dtype=torch.long)), + axis_i=0, + ) + start = opset11.unsqueeze(g, start, 0) + end = symbolic_helper._size_helper(g, self, axis) + + last_slice = g.op("Slice", self, start, end, axis) + + return g.op("SequenceInsert", loop_out, last_slice) + + else: # scalar tensor + dim_size = symbolic_helper._size_helper(g, self, axis) + min_split_size = g.op("Div", dim_size, indices_or_sections) + min_split_size_plus_1 = g.op( + "Add", + min_split_size, + const_1, + ) + num_splits_one_extra = g.op("Mod", dim_size, indices_or_sections) + splits = g.op("Tile", min_split_size_plus_1, num_splits_one_extra) + leftover = g.op( + "Tile", + min_split_size, + g.op( + "Sub", + opset11.unsqueeze(g, indices_or_sections, 0), + num_splits_one_extra, + ), + ) + + splits = g.op("Concat", splits, leftover, axis_i=0) + if _outputs is None: + return g.op("SplitToSequence", self, splits, axis_i=dim) + return g.op("Split", self, splits, axis_i=dim, outputs=_outputs) + + +@_onnx_symbolic("aten::unbind") +@symbolic_helper.parse_args("v", "i", "i") +def unbind(g: jit_utils.GraphContext, self, dim=0, _outputs=None): + if _outputs is None: + return g.op( + "SplitToSequence", + self, + g.op("Constant", value_t=torch.tensor(1, dtype=torch.long)), + axis_i=dim, + keepdims_i=0, + ) + + splits = g.op("Constant", value_t=torch.tensor([1] * _outputs)) + outputs = g.op("Split", self, splits, axis_i=dim, outputs=_outputs) + outputs = [outputs] if _outputs == 1 else outputs + squeezed_outputs = [ + g.op("Squeeze", out, g.op("Constant", value_t=torch.tensor([dim]))) + for out in outputs + ] + return squeezed_outputs + + +@_onnx_symbolic("aten::nonzero_numpy") +# Emitted from `torch.nonzero(x, as_tuple=True)` +def nonzero_numpy(g: jit_utils.GraphContext, input, _outputs=None): + return unbind(g, opset9.nonzero(g, input), 1, _outputs=_outputs) + + +@_onnx_symbolic("aten::where") +@symbolic_helper.parse_args("v", "v", "v", "i") +def where(g: jit_utils.GraphContext, condition, self=None, other=None, _outputs=None): + # Assumes that torch.where's first argument takes only Bool and Byte tensors. + if not symbolic_helper._is_bool(condition): + condition = g.op("Cast", condition, to_i=_C_onnx.TensorProtoDataType.BOOL) + if self is None: + condition = opset9.nonzero(g, condition) + return symbolic_helper._unbind_helper( + g, condition, g.op("Constant", value_t=torch.tensor(1)), _outputs + ) + # pyrefly: ignore [bad-argument-type] + return g.op("Where", condition, self, other) + + +@_onnx_symbolic("aten::fake_quantize_per_channel_affine") +@symbolic_helper.parse_args("v", "v", "v", "i", "i", "i") +def fake_quantize_per_channel_affine( + g: jit_utils.GraphContext, + inputs, + scale, + zero_point, + axis, + quant_min=-128, + quant_max=127, +): + # NOTE: (0, 127) is allowed as special case. PyTorch restricts activations to be in the range (0, 127). + # https://github.com/pytorch/pytorch/blob/b34b192d6b97325c9f78e5995c48c8498ede34bd/torch/ao/quantization/observer.py#L1422 + if (quant_min, quant_max) not in [(0, 255), (-128, 127), (0, 127)]: + raise errors.SymbolicValueError( + "For (quant_min, quant_max), ONNX allows only (0, 127), (0, 255) and (-128, 127). " + f"Got ({quant_min}, {quant_max})", + inputs, + ) + # ONNX defines zero_point to be int8 or uint8 + if quant_min == 0: + zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.UINT8) + else: + zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.INT8) + quantized = g.op("QuantizeLinear", inputs, scale, zero_point, axis_i=axis) + if (quant_min, quant_max) == (0, 127): + quantized = g.op( + "Clip", + quantized, + opset9.unused(g), + g.op("Constant", value_t=torch.tensor(127, dtype=torch.uint8)), + ) + return g.op("DequantizeLinear", quantized, scale, zero_point, axis_i=axis) + + +@_onnx_symbolic("aten::fake_quantize_per_tensor_affine") +@symbolic_helper.parse_args("v", "v", "v", "i", "i") +def fake_quantize_per_tensor_affine( + g: jit_utils.GraphContext, + inputs, + scale, + zero_point, + quant_min=-128, + quant_max=127, +): + # NOTE: (0, 127) is allowed as special case. PyTorch restricts activations to be in the range (0, 127). + # https://github.com/pytorch/pytorch/blob/b34b192d6b97325c9f78e5995c48c8498ede34bd/torch/ao/quantization/observer.py#L1422 + if (quant_min, quant_max) not in [(0, 255), (-128, 127), (0, 127)]: + raise errors.SymbolicValueError( + "For (quant_min, quant_max), ONNX allows only (0, 127), (0, 255) and (-128, 127). " + f"Got ({quant_min}, {quant_max})", + inputs, + ) + if quant_min == 0: + zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.UINT8) + else: + zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.INT8) + if ( + _type_utils.JitScalarType.from_value(scale, _type_utils.JitScalarType.UNDEFINED) + != _type_utils.JitScalarType.FLOAT + ): + scale = g.op("Cast", scale, to_i=_C_onnx.TensorProtoDataType.FLOAT) + quantized = g.op("QuantizeLinear", inputs, scale, zero_point) + if (quant_min, quant_max) == (0, 127): + quantized = g.op( + "Clip", + quantized, + opset9.unused(g), + g.op("Constant", value_t=torch.tensor(127, dtype=torch.uint8)), + ) + return g.op("DequantizeLinear", quantized, scale, zero_point) + + +def _reduce_op_symbolic(onnx_op_name): + def symbolic(g, self, dim=None, keepdim=None): + self = symbolic_helper._maybe_cast_reduce_op_input(g, self) + if dim is None: + # all-reduce path + return symbolic_helper._handle_reduce_dim_none(g, self, onnx_op_name) + else: + keepdim = symbolic_helper._get_const(keepdim, "i", "keepdim") + return g.op(onnx_op_name, self, dim, keepdims_i=keepdim) + + return symbolic + + +@_onnx_symbolic( + "aten::sum", + decorate=[symbolic_helper._apply_params("ReduceSum", "sum")], +) +def _reduce_with_dtype(onnx_op, name): + symbolic = _reduce_op_symbolic(onnx_op) + + @symbolic_helper._overload_by_arg_count + def reduce(g, *args, **kwargs): + @symbolic_helper.parse_args("v", "none") + def reduce_nodim(g, self, dtype): + dtype_onnx = None + if dtype.node().kind() == "onnx::Constant": + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + dtype_onnx = _type_utils.JitScalarType(dtype).onnx_type() + self = g.op("Cast", self, to_i=dtype_onnx) + elif dtype.node().kind() != "prim::Constant": + return symbolic_helper._unimplemented(name, "dtype", dtype) + result = symbolic(g, self) + if dtype_onnx is not None: + result_dtype_onnx = _type_utils.JitScalarType.from_value( + result + ).onnx_type() + if result_dtype_onnx != dtype_onnx: + result = g.op("Cast", result, to_i=dtype_onnx) + return result + + @symbolic_helper.parse_args("v", "v", "i", "none") + def reduce_dim(g, self, dim, keepdim, dtype): + dtype_onnx = None + if dtype.node().kind() == "onnx::Constant": + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + dtype_onnx = _type_utils.JitScalarType(dtype).onnx_type() + self = g.op("Cast", self, to_i=dtype_onnx) + elif dtype.node().kind() != "prim::Constant": + return symbolic_helper._unimplemented(name, "dtype", dtype) + result = symbolic(g, self, dim, keepdim) + if dtype_onnx is not None: + result_dtype_onnx = _type_utils.JitScalarType.from_value( + result + ).onnx_type() + if result_dtype_onnx != dtype_onnx: + result = g.op("Cast", result, to_i=dtype_onnx) + return result + + return reduce_nodim, reduce_dim + + return reduce + + +# Ported from +# https://github.com/microsoft/onnxscript/blob/6b1b81700b4523f31d8c6d3321e5d8ef5d42b764/onnxscript/function_libs/torch_aten/ops/core.py#L6097 +# NOTE: Supporting aten::unflatten before opset13 needs helper function to adjust ONNX op changes in Concat, Slice, ... +@_onnx_symbolic("aten::unflatten") +def unflatten(g: jit_utils.GraphContext, input, dim, unflattened_size): + input_dim = symbolic_helper._get_tensor_rank(input) + if input_dim is None: + return symbolic_helper._unimplemented( + "dim", + "ONNX and PyTorch use different strategies to split the input. " + "Input rank must be known at export time.", + ) + + # dim could be negative + input_dim = g.op("Constant", value_t=torch.tensor([input_dim], dtype=torch.int64)) + dim = g.op("Add", input_dim, dim) + dim = g.op("Mod", dim, input_dim) + + input_size = g.op("Shape", input) + + head_start_idx = g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64)) + head_end_idx = g.op( + "Reshape", dim, g.op("Constant", value_t=torch.tensor([1], dtype=torch.int64)) + ) + head_part_rank = g.op("Slice", input_size, head_start_idx, head_end_idx) + + dim_plus_one = g.op( + "Add", dim, g.op("Constant", value_t=torch.tensor([1], dtype=torch.int64)) + ) + tail_start_idx = g.op( + "Reshape", + dim_plus_one, + g.op("Constant", value_t=torch.tensor([1], dtype=torch.int64)), + ) + tail_end_idx = g.op( + "Constant", value_t=torch.tensor([_constants.INT64_MAX], dtype=torch.int64) + ) + tail_part_rank = g.op("Slice", input_size, tail_start_idx, tail_end_idx) + + final_shape = g.op( + "Concat", head_part_rank, unflattened_size, tail_part_rank, axis_i=0 + ) + + return symbolic_helper._reshape_helper(g, input, final_shape) + + +@_onnx_symbolic("aten::unsafe_chunk") +@symbolic_helper.parse_args("v", "i", "i", "i") +def unsafe_chunk(g: jit_utils.GraphContext, self, chunks, dim, _outputs=None): + if _outputs is None: + return g.op( + "SplitToSequence", + self, + g.op("Constant", value_t=torch.tensor(1, dtype=torch.long)), + axis_i=dim, + keepdims_i=0, + ) + + size = symbolic_helper._get_tensor_dim_size(self, dim) + if size is None: + return symbolic_helper._unimplemented("unsafe_chunk", "unknown dimension size") + split_size = (size + chunks - 1) // chunks + splits = [split_size] * (size // split_size) + leftover = size % split_size + if leftover: + splits.append(leftover) + + # TODO: So far we don"t have a module using this method. We"ll keep + # this as a constant unless we see a request of dynamics in any + # user's modules. + splits = g.op("Constant", value_t=torch.tensor(splits, dtype=torch.long)) + return g.op("Split", self, splits, axis_i=dim, outputs=_outputs) + + +@_onnx_symbolic("aten::tile") +def tile(g: jit_utils.GraphContext, self, dims): + self_shape = g.op("Shape", self) + self_rank = g.op("Size", self_shape) + dims_rank = g.op("Size", dims) + diff = g.op("Sub", self_rank, dims_rank) + const_zero = g.op("Constant", value_t=torch.tensor([0])) + + # 1. If dims is shorter than self.shape pad dims with 1 + dims_shorter_than_self_shape = g.op("Greater", diff, const_zero) + ( + if_op_greater, + (if_context_greater, else_context_greater), + _, + ) = jit_utils.add_op_with_blocks( + g, "If", dims_shorter_than_self_shape, n_blocks=2, outputs=1 + ) + const_one = if_context_greater.op("Constant", value_t=torch.LongTensor([1])) + diff_1d_greater = if_context_greater.op("Reshape", diff, const_one) + exapnd_ones_greater = if_context_greater.op("Expand", const_one, diff_1d_greater) + dims_ = if_context_greater.op("Concat", exapnd_ones_greater, dims, axis_i=0) + utils._add_output_to_block(if_context_greater.block, dims_) + identity_dim = else_context_greater.op("Identity", dims) + utils._add_output_to_block(else_context_greater.block, identity_dim) + dims_final = if_op_greater.node().output() + + # 2. If dims is longer than self.shape pad self.shape with 1 + dims_longer_than_self_shape = g.op("Less", diff, const_zero) + ( + if_op_less, + (if_context_less, else_context_less), + _, + ) = jit_utils.add_op_with_blocks( + g, "If", dims_longer_than_self_shape, n_blocks=2, outputs=1 + ) + const_one = if_context_less.op("Constant", value_t=torch.LongTensor([1])) + diff_1d_less = if_context_less.op( + "Reshape", + if_context_less.op("Abs", diff), + const_one, + ) + exapnd_ones_less = if_context_less.op("Expand", const_one, diff_1d_less) + self_final_shape = if_context_less.op( + "Concat", exapnd_ones_less, self_shape, axis_i=0 + ) + self_ = if_context_less.op("Reshape", self, self_final_shape) + utils._add_output_to_block(if_context_less.block, self_) + identity_self = else_context_less.op("Identity", self) + utils._add_output_to_block(else_context_less.block, identity_self) + self_final = if_op_less.node().output() + + dims_final = g.op("Cast", dims_final, to_i=_C_onnx.TensorProtoDataType.INT64) + return g.op("Tile", self_final, dims_final) + + +@_onnx_symbolic("aten::repeat_interleave") +def repeat_interleave( + g: jit_utils.GraphContext, self, repeats, dim=None, output_size=None +): + repeats_dim = symbolic_helper._get_tensor_rank(repeats) + repeats_sizes = symbolic_helper._get_tensor_sizes(repeats) + input_sizes = symbolic_helper._get_tensor_sizes(self) + if repeats_dim is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of repeat_interleave for unknown repeats rank.", + self, + ) + if repeats_sizes is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of repeat_interleave for unknown repeats size.", + self, + ) + if input_sizes is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of repeat_interleave for unknown input size.", + self, + ) + + final_dim = dim + # if dim is None flatten + # By default, use the flattened input array, and return a flat output array + if symbolic_helper._is_none(dim): + self = symbolic_helper._reshape_helper( + g, self, g.op("Constant", value_t=torch.tensor([-1])) + ) + dim = torch.tensor(0, dtype=torch.int64) + else: + dim = symbolic_helper._maybe_get_scalar(dim) + + # Handle cases where dim is negative + if dim < 0: + dim += len(input_sizes) + + output_sizes = input_sizes.copy() + for idx, input_size in enumerate(input_sizes): + if input_size is None: + output_sizes[idx], input_sizes[idx] = 0, -1 + + # Check if all indices should be repeated the same number of times. + if repeats_dim == 0 or (repeats_dim == 1 and repeats_sizes[0] == 1): + return symbolic_helper._repeat_interleave_single_value_repeat_helper( + g, self, repeats, dim + ) + + cond_dynamic_repeats = repeats_dim == 1 and repeats_sizes[0] is None + # If input size is dynamic or repeats vector is dynamic + if output_sizes[dim] == 0 or cond_dynamic_repeats: + reps = symbolic_helper._size_helper(g, self, dim) + reps = opset11.unsqueeze(g, reps, 0) + + # Check if repeats is dynamic + # As repeats is dynamic, we use a where node as a substitute for the if statement + # If repests_dim = 1, expand repeats otherwise use original tensor + if cond_dynamic_repeats: + repeat_dim = symbolic_helper._size_helper( + g, repeats, g.op("Constant", value_t=torch.LongTensor([0])) + ) + repeat_cond = g.op( + "Equal", repeat_dim, g.op("Constant", value_t=torch.LongTensor([1])) + ) + repeats = where(g, repeat_cond, g.op("Expand", repeats, reps), repeats) + # There are cases when the repeats are 1-d tensor with multiple repeats, but dim + # provided along one of the dynamic axes provided. A simple example would be + # input.shape -> [1, 1, *] where * represents the dynamic axes, and dim = 2 + # Now, repeat interleaving can be performed in pytorch when the value of * matches + # with the number of elements in repeat, for example if * -> 2, number of repeats + # should be 2 as well. + else: + return opset9.repeat_interleave(g, self, repeats, final_dim) + + reps_like = g.op( + "ConstantOfShape", + g.op("Shape", repeats), + value_t=torch.tensor([1], dtype=torch.long), + ) + r_splits = split(g, repeats, reps_like, 0) + i_splits = split(g, self, reps_like, dim) + + output_sizes[dim], input_sizes[dim] = -1, 1 + + # Create a loop to iterate over each value along the dimension + # and perform individual interleaving using the repeats tensor + # Loop is of the following pattern + # input (trip_count, cond) + # int trip_count = ...; + # bool cond = ...; + # for (int i=0; i < trip_count && cond; ++i) { + # cond = ...; + # } + + # Loop conditions + loop_condition = g.op("Constant", value_t=torch.tensor(1)) + loop_condition = g.op("Cast", loop_condition, to_i=_C_onnx.TensorProtoDataType.BOOL) + loop_len = reps + + # Create an empty sequence to store final expansions + final_splits = g.op("SequenceEmpty") + + # Loop inputs + loop, (loop_context,), _ = jit_utils.add_op_with_blocks( + g, "Loop", loop_len, loop_condition, final_splits, n_blocks=1 + ) + + loop_block = loop_context.block + block_input_iter = utils._add_input_to_block(loop_block) + cond = utils._add_input_to_block(loop_block) # noqa: F841 + final_splits = utils._add_input_to_block(loop_block) + + r_split = loop_context.op("SequenceAt", r_splits, block_input_iter) + i_split = loop_context.op("SequenceAt", i_splits, block_input_iter) + + i_split = opset11.unsqueeze(loop_context, i_split, dim + 1) + r_concat = [ + loop_context.op("Constant", value_t=torch.LongTensor(input_sizes[: dim + 1])), + r_split, + loop_context.op("Constant", value_t=torch.LongTensor(input_sizes[dim + 1 :])), + ] + r_concat = loop_context.op("Concat", *r_concat, axis_i=0) + i_split = opset9.expand(loop_context, i_split, r_concat, None) + i_split = symbolic_helper._reshape_helper( + loop_context, i_split, g.op("Constant", value_t=torch.LongTensor(output_sizes)) + ) + final_splits = loop_context.op("SequenceInsert", final_splits, i_split) + + # Loop outputs + cond_out = loop_context.op( + "Cast", loop_condition, to_i=_C_onnx.TensorProtoDataType.BOOL + ) + utils._add_output_to_block(loop_block, cond_out) + utils._add_output_to_block(loop_block, final_splits) + + loop_out = loop.node().output() + loop_out = g.op("ConcatFromSequence", loop_out, axis_i=dim) + return loop_out + + +@_onnx_symbolic("aten::diagonal") +@symbolic_helper.parse_args("v", "i", "i", "i") +def diagonal(g: jit_utils.GraphContext, self, offset, dim1, dim2): + rank = symbolic_helper._get_tensor_rank(self) + # Replace negative indexing when rank is known + if rank is not None: + dim1 = dim1 if dim1 >= 0 else dim1 + rank + dim2 = dim2 if dim2 >= 0 else dim2 + rank + + dim1_size = opset9.size( + g, self, dim=g.op("Constant", value_t=torch.LongTensor([dim1])) + ) + dim2_size = opset9.size( + g, self, dim=g.op("Constant", value_t=torch.LongTensor([dim2])) + ) + # Create appropriate mask + mask_shape = g.op("Concat", dim1_size, dim2_size, axis_i=0) + mask = opset9.zeros(g, mask_shape, None, None, None) + mask = g.op("EyeLike", mask, k_i=offset) + # dim1 and dim2 appended as a dimension at the end of the shape + + if rank is not None: + axes = list(range(rank)) + axes.remove(dim1) + axes.remove(dim2) + self = g.op("Transpose", self, perm_i=axes + [dim1, dim2]) + else: + return symbolic_helper._unimplemented("diagonal", "unknown input rank") + + # Multiply input and mask to calculate values along diagonal + # The mask consists of one values where diagonal values are to be calculated + # For example: + # [[1.1, 1.2, 1.3], * [[1, 0, 0] = [[1.1, 0, 0], + # [2.1, 2.2, 2.3], [0, 1, 0] [0, 2.2, 0], + # [3.1, 3.2, 3.3]] [0, 0, 1]] [0, 0, 3.3]] + result = g.op("Mul", self, mask) + result = symbolic_helper._reducesum_helper(g, result, axes_i=[-1], keepdims_i=0) + + # Calculate gather indices based on offset and dims + # If offset is greater than zero, set offset to zero as this aids in + # calculation of selection window + offset_op = g.op("Constant", value_t=torch.LongTensor([offset])) + if offset >= 0: + diag_size = g.op( + "Max", + g.op("Min", dim1_size, g.op("Sub", dim2_size, offset_op)), + g.op("Constant", value_t=torch.LongTensor([0])), + ) + offset = 0 + else: + diag_size = g.op( + "Max", + g.op("Min", g.op("Add", dim1_size, offset_op), dim2_size), + g.op("Constant", value_t=torch.LongTensor([0])), + ) + diag_size = g.op("Concat", diag_size, axis_i=0) + + # Calculate which diagonal values to select + # For example, in cases with offsets: + # [[0, 1.1, 0] + # [0, 0, 2.2]] + # we need to select the last two columns, so we create a tensor + # with all columns that are to be selected + # So in this example, it is [1, 2] + select_window_ones_fill = opset9.ones(g, diag_size, 4, None, None) + select_window = g.op( + "CumSum", + select_window_ones_fill, + g.op("Constant", value_t=torch.LongTensor([0])), + ) + select_window = g.op( + "Add", + select_window, + g.op("Constant", value_t=torch.LongTensor([abs(offset) - 1])), + ) + + gather_shape = [ + opset9.size(g, result, dim=g.op("Constant", value_t=torch.LongTensor([axis]))) + for axis in list(range(rank))[:-2] + ] + gather_shape.append(diag_size) + gather_shape = g.op("Concat", *gather_shape, axis_i=0) + gather_indices = opset9.zeros(g, gather_shape, 4, None, None) + + # There might be cases where offset value is greater than number of rows/columns + # and might cause the diagonal to overrun and as a result of this, diag_size would be zero. + # For example, if + # offset = 9, dim1_size = 2 (columns), dim2_size = 4 (rows) + # diag_size = max(min(2, (4-9)), 0) = 0, based on calculation above + # Cases with diagonal overrun always result in diag_size = max(0, -ve value) = 0 + # In cases without diagonal overrun, we select the appropriate rows/columns along which we + # are calculating diagonal values. In cases with diagonal overrun, we return a tensor which has + # the dimension of the row/column where overrun occurred as 0-dim, as we are essentially + # returning an empty tensor + overrun_cond = g.op( + "Not", + g.op( + "Equal", + diag_size, + g.op("Constant", value_t=torch.tensor(0, dtype=torch.int64)), + ), + ) + + if_op, (if_context, else_context), _ = jit_utils.add_op_with_blocks( + g, "If", overrun_cond, n_blocks=2 + ) + + gather_indices_if_block = if_context.op("Add", gather_indices, select_window) + gather_indices_if_block = symbolic_helper._unsqueeze_helper( + if_context, gather_indices_if_block, [rank - 1] + ) + final_non_overrun = if_context.op( + "GatherND", result, gather_indices_if_block, batch_dims_i=rank - 2 + ) + final_overrun = opset9.zeros(else_context, gather_shape, 6, None, None) + utils._add_output_to_block(if_context.block, final_non_overrun) + utils._add_output_to_block(else_context.block, final_overrun) + return if_op + + +# Quantized ops + + +@_onnx_symbolic("quantized::linear") +def quantized_linear( + g: jit_utils.GraphContext, q_input, q_weight, bias, op_scale, op_zero_point +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.linear(g, input, weight, bias) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::linear_relu") +def quantized_linear_relu( + g: jit_utils.GraphContext, q_input, q_weight, bias, op_scale, op_zero_point +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.linear(g, input, weight, bias) + output = opset9.relu(g, output) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv1d_relu") +def quantized_conv1d_relu( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv1d(g, input, weight, bias, stride, padding, dilation, groups) + output = opset9.relu(g, output) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv2d_relu") +def quantized_conv2d_relu( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv2d(g, input, weight, bias, stride, padding, dilation, groups) + output = opset9.relu(g, output) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv3d_relu") +def quantized_conv3d_relu( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv3d(g, input, weight, bias, stride, padding, dilation, groups) + output = opset9.relu(g, output) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv1d") +def quantized_conv1d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv1d(g, input, weight, bias, stride, padding, dilation, groups) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv2d") +def quantized_conv2d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv2d(g, input, weight, bias, stride, padding, dilation, groups) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv3d") +def quantized_conv3d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv3d(g, input, weight, bias, stride, padding, dilation, groups) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv_transpose1d") +def quantized_conv_transpose1d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + output_padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv_transpose2d( + g, input, weight, bias, stride, padding, output_padding, groups, dilation + ) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv_transpose2d") +def quantized_conv_transpose2d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + output_padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv_transpose2d( + g, input, weight, bias, stride, padding, output_padding, groups, dilation + ) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +@_onnx_symbolic("quantized::conv_transpose3d") +def quantized_conv_transpose3d( + g: jit_utils.GraphContext, + q_input, + q_weight, + bias, + stride, + padding, + output_padding, + dilation, + groups, + op_scale, + op_zero_point, +): + input, input_scale, _, _ = symbolic_helper.dequantize_helper(g, q_input) + weight, weight_scale, _, axis = symbolic_helper.dequantize_helper(g, q_weight) + q_bias = symbolic_helper.requantize_bias_helper( + g, bias, input_scale, weight_scale, axis + ) + bias, _, _, _ = symbolic_helper.dequantize_helper(g, q_bias) + + output = opset9.conv_transpose3d( + g, input, weight, bias, stride, padding, output_padding, groups, dilation + ) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py new file mode 100644 index 0000000000000000000000000000000000000000..7beb3de13f27260a5dd5d139f4ccbd8474d5cefd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset14.py @@ -0,0 +1,301 @@ +# mypy: allow-untyped-defs +# mypy: disable-error-code=arg-type +"""This file exports ONNX ops for opset 14. + +Note [ONNX operators that are added/updated in opset 14] +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +New operators: + HardSwish, Trilu + +Updated operators: + Reshape + Add, Sub, Mul, Div + GRU, LSTM, RNN + BatchNorm, Cumsum, Relu +""" + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in README.md +from __future__ import annotations + +import functools + +import torch +from torch.onnx import _constants +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, +) +from torch.onnx._internal.torchscript_exporter._globals import GLOBALS + + +__all__ = [ + "hardswish", + "tril", + "triu", + "reshape", + "batch_norm", + "quantized_hardswish", + "scaled_dot_product_attention", +] + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=14) + + +@_onnx_symbolic("aten::hardswish") +@symbolic_helper.parse_args("v") +def hardswish(g: jit_utils.GraphContext, self): + return g.op("HardSwish", self) + + +@_onnx_symbolic("aten::tril") +def tril(g: jit_utils.GraphContext, self, diagonal, out=None): + return g.op("Trilu", self, diagonal, upper_i=0) + + +@_onnx_symbolic("aten::triu") +def triu(g: jit_utils.GraphContext, self, diagonal, out=None): + return g.op("Trilu", self, diagonal, upper_i=1) + + +@_onnx_symbolic("aten::reshape") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "v") +def reshape(g: jit_utils.GraphContext, self, shape): + # NOTE: Due to bug in ORT https://github.com/microsoft/onnxruntime/issues/10664 + # Reshape export cannot utilize the new allowzero attribute introduced in opset 14. + return symbolic_helper._reshape_helper(g, self, shape, allowzero=0) + + +@_onnx_symbolic("aten::batch_norm") +@symbolic_helper.parse_args("v", "v", "v", "v", "v", "i", "f", "f", "i") +def batch_norm( + g: jit_utils.GraphContext, + input, + weight, + bias, + running_mean, + running_var, + training, + momentum, + eps, + cudnn_enabled, +): + if ( + torch.is_autocast_enabled() + and not symbolic_helper.args_have_same_dtype( + [input, weight, bias, running_mean, running_var] + ) + and GLOBALS.export_onnx_opset_version < 15 + ): + return symbolic_helper._onnx_opset_unsupported_detailed( + "BatchNormalization", + 14, + 15, + "All input tensors must have the same `dtype`." + " Turn off Autocast or export using opset version 15.", + input, + ) + + symbolic_helper.check_training_mode(training, "batch_norm") + weight, bias, running_mean, running_var = symbolic_helper._batchnorm_helper( + g, input, weight, bias, running_mean, running_var + ) + out = g.op( + "BatchNormalization", + input, + weight, + bias, + running_mean, + running_var, + epsilon_f=eps, + momentum_f=1 - momentum, + training_mode_i=0 if not training else 1, + outputs=1 if not training else 3, + ) + if not training: + return out + else: + res, new_running_mean, new_running_var = out + new_running_mean.setType(running_mean.type()) + new_running_var.setType(running_var.type()) + return res + + +@_onnx_symbolic("quantized::hardswish") +def quantized_hardswish(g: jit_utils.GraphContext, x, op_scale, op_zero_point): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + + output = hardswish(g, x) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +# Ported from +# https://github.com/microsoft/onnxscript/blob/6b1b81700b4523f31d8c6d3321e5d8ef5d42b764/onnxscript/function_libs/torch_aten/ops/nn.py#L1504 +# aten_scaled_dot_product_attention +# NOTE: Need op.Trilu +@_onnx_symbolic("aten::scaled_dot_product_attention") +@symbolic_helper.parse_args("v", "v", "v", "v", "f", "b", "v", "b") +def scaled_dot_product_attention( + g: jit_utils.GraphContext, + query: torch._C.Value, + key: torch._C.Value, + value: torch._C.Value, + attn_mask: torch._C.Value | None = None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: torch._C.Value | None = None, + enable_gqa: bool = False, +): + if is_causal and not symbolic_helper._is_none(attn_mask): + raise AssertionError("is_causal and attn_mask cannot be set at the same time") + if enable_gqa: + raise AssertionError( + "conversion of scaled_dot_product_attention not implemented if enable_gqa is True" + ) + + if symbolic_helper._is_none(scale): + scale = _attention_scale(g, query) + + if is_causal: + attn_mask = _causal_attention_mask(g, query, key) + + # Swap the last two axes of key + # NOTE: onnx-script has different logic here, because the attribute perms in + # transpose needs list of ints + key_shape_builtin = symbolic_helper._get_tensor_rank(key) + # pyrefly: ignore [bad-argument-type, no-matching-overload] + key_transposed_axes = list(range(key_shape_builtin)) + key_transposed_axes[-1], key_transposed_axes[-2] = ( + key_transposed_axes[-2], + key_transposed_axes[-1], + ) + key_transposed = g.op("Transpose", key, perm_i=key_transposed_axes) + + # https://github.com/pytorch/pytorch/blob/12da0c70378b5be9135c6fda62a9863bce4a4818/aten/src/ATen/native/transformers/attention.cpp#L653 + # Scale q, k before matmul for stability see https://tinyurl.com/sudb9s96 for math + # pyrefly: ignore [bad-argument-type] + query_scaled = g.op("Mul", query, g.op("Sqrt", scale)) + # pyrefly: ignore [bad-argument-type] + key_transposed_scaled = g.op("Mul", key_transposed, g.op("Sqrt", scale)) + mul_qk = g.op("MatMul", query_scaled, key_transposed_scaled) + + if symbolic_helper._is_none(attn_mask): + mul_qk_add = mul_qk + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + elif ( + _type_utils.JitScalarType.from_value(attn_mask) + == _type_utils.JitScalarType.BOOL + ): + # Turn the Boolean mask to float: attn_mask.masked_fill(not attn_mask, -float('inf')) + const_zero = g.op("Constant", value_t=torch.tensor([0.0])) + const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")])) + # pyrefly: ignore [bad-argument-type] + attn_mask = g.op("Where", attn_mask, const_zero, const_neg_inf) + mul_qk_add = g.op("Add", mul_qk, attn_mask) + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + # When using scaled dot product attention with a boolean mask, the softmax operation might return NaN values + # due to the presence of -inf in an entire row (padding tokens), resulting in 0/0 (NaN) in the softmax output. + # This is because there's no safe softmax imp in ONNX, so we need to handle NaN values explicitly to match + # the behavior of PyTorch with boolean masks. + attn_weight = g.op("Where", g.op("IsNaN", attn_weight), const_zero, attn_weight) + elif _type_utils.JitScalarType.from_value(attn_mask) in ( + _type_utils.JitScalarType.FLOAT, + _type_utils.JitScalarType.HALF, + _type_utils.JitScalarType.BFLOAT16, + ): + # pyrefly: ignore [bad-argument-type] + mul_qk_add = g.op("Add", mul_qk, attn_mask) + attn_weight = g.op("Softmax", mul_qk_add, axis_i=-1) + else: + raise ValueError( + f"Unsupported type for attn_mask: {_type_utils.JitScalarType.from_value(attn_mask)}" + ) + + if dropout_p != 0: + attn_weight = g.op( + "Dropout", + attn_weight, + g.op("Constant", value_t=torch.tensor(dropout_p, dtype=torch.float)), + ) + + return g.op("MatMul", attn_weight, value) + + +def _attention_scale( + g: jit_utils.GraphContext, query: torch._C.Value +) -> torch._C.Value: + """Calculate the scale factor for the attention result. + + Args: + query: Tensor of shape [..., L, E] + + Returns: + Scalar scale factor := 1 / math.sqrt(query.size(-1)) + """ + query_shape = g.op("Shape", query) + query_shape_last = g.op( + "Slice", + query_shape, + g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)), + g.op( + "Constant", value_t=torch.tensor([_constants.INT64_MAX], dtype=torch.int64) + ), + ) + embedding_size = g.op( + "Cast", + query_shape_last, + to_i=_type_utils.JitScalarType.from_value(query).onnx_type(), + ) + const_one = g.op("Constant", value_t=torch.tensor([1.0], dtype=torch.float)) + scale = g.op("Div", const_one, g.op("Sqrt", embedding_size)) + # Add a Cast to convert the scale back to original type + scale = g.op( + "Cast", + scale, + to_i=_type_utils.JitScalarType.from_value(query).onnx_type(), + ) + return scale + + +def _causal_attention_mask( + g: jit_utils.GraphContext, query: torch._C.Value, key: torch._C.Value +) -> torch._C.Value: + """Create a causal mask for the given query and key tensors. + + Equivalent to:: + mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0) + attn_mask = torch.zeros(L, S, dtype=torch.float) + attn_mask = attn_mask.masked_fill(not mask, -float("inf")) + + Args: + query: Tensor of shape [..., L, E] + key: Tensor of shape [..., S, E] + + Returns: + Tensor of shape [L, S] + """ + + query_shape = g.op("Shape", query) + key_shape = g.op("Shape", key) + + last_idx = g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)) + second_last_idx = g.op("Constant", value_t=torch.tensor([-2], dtype=torch.int64)) + target_length = g.op("Slice", query_shape, second_last_idx, last_idx) + source_length = g.op("Slice", key_shape, second_last_idx, last_idx) + # attn_mask = torch.ones(L, S) := { + size = g.op("Concat", target_length, source_length, axis_i=0) + const_one = g.op("Constant", value_t=torch.tensor([1.0])) + attn_mask = g.op("Expand", const_one, size) + # } + attn_mask = g.op("Trilu", attn_mask, upper_i=0) + # The causal mask has 0s in the lower triangle and -inf in the upper triangle. + const_zero = g.op("Constant", value_t=torch.tensor([0.0])) + const_neg_inf = g.op("Constant", value_t=torch.tensor([-float("inf")])) + attn_mask = g.op( + "Where", g.op("Equal", attn_mask, const_zero), const_neg_inf, const_zero + ) + return attn_mask diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset15.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset15.py new file mode 100644 index 0000000000000000000000000000000000000000..4f86a7f2f86254109bac74df1a54014dd93c1787 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset15.py @@ -0,0 +1,84 @@ +# mypy: allow-untyped-defs +"""This file exports ONNX ops for opset 15. + +Note [ONNX operators that are added/updated in opset 15] +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +https://github.com/onnx/onnx/blob/master/docs/Changelog.md#version-15-of-the-default-onnx-operator-set +New operators: + Bernoulli + CastLike + Optional + OptionalGetElement + OptionalHasElement + +Updated operators: + BatchNormalization https://github.com/onnx/onnx/pull/3545 + Backwards compatible + TODO: test coverage for mixed types inputs. + Pow https://github.com/onnx/onnx/pull/3412 + Backwards compatible + TODO: bfloat16 support. + Shape https://github.com/onnx/onnx/pull/3580 + Backwards compatible + TODO: optional start/end attribute. +""" + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in README.md + +import functools + +import torch +from torch import _C +from torch.onnx._internal.torchscript_exporter import ( + jit_utils, + registration, + symbolic_helper, + symbolic_opset9 as opset9, +) + + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=15) + + +@_onnx_symbolic("aten::__is_") +def aten__is_(g: jit_utils.GraphContext, self, other): + if symbolic_helper._is_none(other): + if isinstance(self.type(), _C.OptionalType): + none = g.op("OptionalHasElement", self) + return g.op("Not", none) + else: + return g.op("Constant", value_t=torch.BoolTensor([0])) + return opset9.eq(g, self, other) + + +@_onnx_symbolic("aten::__isnot_") +@opset9.wrap_logical_op_with_negation # type: ignore[has-type] +def aten__isnot_(g: jit_utils.GraphContext, self, other): + return aten__is_(g, self, other) + + +@_onnx_symbolic("aten::bernoulli") +def bernoulli(g: jit_utils.GraphContext, input, p=None, generator=None, out=None): + if out is not None and not symbolic_helper._is_none(out): + symbolic_helper._unimplemented( + "Bernoulli", "out parameter is not supported for bernoulli", input + ) + if generator is not None and not symbolic_helper._is_none(generator): + symbolic_helper._unimplemented( + "Bernoulli", "generator is not supported for bernoulli", input + ) + if p is None or symbolic_helper._is_none(p): + return g.op("Bernoulli", input) + return opset9.bernoulli(g, input, p, generator, out) + + +@_onnx_symbolic("prim::unchecked_cast") +def prim_unchecked_cast(g: jit_utils.GraphContext, self): + # exists to refine the type of the Value + # if x is Optional[Tensor], unchecked_cast will cast + # x to Tensor, so the rest of the graph knows that x is a Tensor. + if isinstance(self.type(), _C.OptionalType): + return g.op("OptionalGetElement", self) + + return self diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset16.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset16.py new file mode 100644 index 0000000000000000000000000000000000000000..a617270a2a7c69f3faf6da45566dec9d2096b0f6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset16.py @@ -0,0 +1,191 @@ +# mypy: allow-untyped-defs +"""This file exports ONNX ops for opset 16. + +Note [ONNX Operators that are added/updated in opset 16] + +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-16-of-the-default-onnx-operator-set +New operators: + GridSample https://github.com/onnx/onnx/pull/3557 + +Updated operators: + Identity + If + LeakyRelu + Loop + PRelu + RoiAlign + Scan + ScatterElements + ScatterND + Where + GreaterOrEqual + LessOrEqual +""" + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in README.md + +import functools + +import torch +from torch.nn.functional import ( + GRID_SAMPLE_INTERPOLATION_MODES, + GRID_SAMPLE_PADDING_MODES, +) +from torch.onnx import errors +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, + utils, +) + + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=16) + + +# note (mkozuki): Why `grid_sampler` instead of `grid_sample`? +# Because `torch.nn.functional.grid_sample` calls `torch.grid_sampler`. +@_onnx_symbolic("aten::grid_sampler") +@symbolic_helper.parse_args("v", "v", "i", "i", "b") +def grid_sampler( + g: jit_utils.GraphContext, + input, + grid, + mode_enum, + padding_mode_enum, + align_corners, +): + # Check the input and grid tensor rank beforehand. + if symbolic_helper._get_tensor_rank(input) == 5: + return symbolic_helper._onnx_unsupported("GridSample with 5D volumetric input") + mode_s = {v: k for k, v in GRID_SAMPLE_INTERPOLATION_MODES.items()}[mode_enum] # type: ignore[call-arg] + padding_mode_s = {v: k for k, v in GRID_SAMPLE_PADDING_MODES.items()}[ # type: ignore[call-arg] + padding_mode_enum + ] + return g.op( + "GridSample", + input, + grid, + align_corners_i=int(align_corners), + mode_s=mode_s, + padding_mode_s=padding_mode_s, + ) + + +@_onnx_symbolic("aten::scatter_add") +@symbolic_helper.parse_args("v", "i", "v", "v") +def scatter_add(g: jit_utils.GraphContext, self, dim, index, src): + src_type = _type_utils.JitScalarType.from_value( + src, _type_utils.JitScalarType.UNDEFINED + ) + src_sizes = symbolic_helper._get_tensor_sizes(src) + index_sizes = symbolic_helper._get_tensor_sizes(index) + + if len(src_sizes) != len(index_sizes): + return symbolic_helper._unimplemented( + "scatter_add", + f"`index` ({index_sizes}) should have the same dimensionality as `src` ({src_sizes})", + ) + + # PyTorch only allows index shape <= src shape, so we can only consider + # taking index as subset size to src, like PyTorch does. When sizes for src + # and index are not matched or there are dynamic axes, we take index shape to + # slice src to accommodate. + if src_sizes != index_sizes or None in index_sizes: + adjusted_shape = g.op("Shape", index) + starts = g.op("Constant", value_t=torch.tensor([0] * len(index_sizes))) + src = g.op("Slice", src, starts, adjusted_shape) + + src = symbolic_helper._maybe_get_scalar(src) + if symbolic_helper._is_value(src): + return g.op("ScatterElements", self, index, src, axis_i=dim, reduction_s="add") + else: + # Check if scalar "src" has same type as self (PyTorch allows different + # type for scalar src (but not when src is tensor)). If not, insert Cast node. + if _type_utils.JitScalarType.from_value(self) != src_type: + src = g.op( + "Cast", + src, + to_i=_type_utils.JitScalarType.from_value(self).onnx_type(), + ) + + return g.op( + "ScatterElements", + self, + index, + src, + axis_i=dim, + reduction_s="add", + ) + + +@_onnx_symbolic("aten::scatter_reduce") +@symbolic_helper.parse_args("v", "i", "v", "v", "s", "b") +def scatter_reduce( + g: jit_utils.GraphContext, + self: torch._C.Value, + dim: int, + index: torch._C.Value, + src: torch._C.Value, + reduce: str, + include_self: bool, +): + if reduce == "mean": + raise errors.OnnxExporterError( + "ONNX does not support mean reduction for scatter_reduce" + ) + if not include_self: + raise errors.OnnxExporterError( + "ONNX does not support include_self=False for scatter_reduce" + ) + + reduce_mode = { # convert torch string name to onnx string name + "mean": "none", # 'mean' doesn't support in ONNX 1.14 definition + "sum": "add", + "prod": "mul", + "amin": "min", + "amax": "max", + } + onnx_reduce = reduce_mode[reduce] + + self_rank = g.op("Size", g.op("Shape", self)) + + # if self_rank == 0: # assert (index_rank == 0 and rank_src == 0) + self_rank_is_zero = g.op( + "Equal", self_rank, g.op("Constant", value_t=torch.tensor(0, dtype=torch.int64)) + ) + if_op, (if_context, else_context), _ = jit_utils.add_op_with_blocks( + g, "If", self_rank_is_zero, n_blocks=2, outputs=3 + ) + neg_1 = if_context.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)) + + self_reshape = if_context.op("Reshape", self, neg_1) + utils._add_output_to_block(if_context.block, self_reshape) + index_reshape = if_context.op("Reshape", index, neg_1) + utils._add_output_to_block(if_context.block, index_reshape) + src_reshape = if_context.op("Reshape", src, neg_1) + utils._add_output_to_block(if_context.block, src_reshape) + + self_identity = else_context.op("Identity", self) + utils._add_output_to_block(else_context.block, self_identity) + index_identitye = else_context.op("Identity", index) + utils._add_output_to_block(else_context.block, index_identitye) + src_identity = else_context.op("Identity", src) + utils._add_output_to_block(else_context.block, src_identity) + + result = g.op("ScatterElements", *if_op, axis_i=dim, reduction_s=onnx_reduce) + + # if self_rank == 0: + if_op, (if_context, else_context), _ = jit_utils.add_op_with_blocks( + g, "If", self_rank_is_zero, n_blocks=2, outputs=1 + ) + result_squeezed = if_context.op("Squeeze", result) + utils._add_output_to_block(if_context.block, result_squeezed) + result_identity = else_context.op("Identity", result) + utils._add_output_to_block(else_context.block, result_identity) + result_final = if_op.node().output() + + return result_final diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset17.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset17.py new file mode 100644 index 0000000000000000000000000000000000000000..bb17d8f72942b11e58b45c718693d2bb7baafe96 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset17.py @@ -0,0 +1,252 @@ +# mypy: allow-untyped-defs +# mypy: disable-error-code=arg-type +"""This file exports ONNX ops for opset 17. + +Note [ONNX Operators that are added/updated in opset 17] + +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-17-of-the-default-onnx-operator-set +New operators: + BlackmanWindow + DFT + HammingWindow + HannWindow + LayerNormalization + MelWeightMatrix + STFT + SequenceMap +""" + +import functools +from collections.abc import Sequence + +import torch +from torch import _C +from torch.onnx import errors +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, +) + + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in README.md + +__all__ = ["layer_norm", "stft", "quantized_layer_norm"] + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=17) + + +@_onnx_symbolic("aten::layer_norm") +@symbolic_helper.parse_args("v", "is", "v", "v", "f", "none") +def layer_norm( + g: jit_utils.GraphContext, + input: _C.Value, + normalized_shape: Sequence[int], + weight: _C.Value, + bias: _C.Value, + eps: float, + cudnn_enable: bool, +): + # normalized_shape: input shape from an expected input of size + # axis: The first normalization dimension. + # layer_norm normalizes on the last D dimensions, + # where D is the size of normalized_shape + axis = -len(normalized_shape) + scalar_type = _type_utils.JitScalarType.from_value( + input, _type_utils.JitScalarType.FLOAT + ) + dtype = scalar_type.dtype() + if symbolic_helper._is_none(weight): + weight_value = torch.ones(normalized_shape, dtype=dtype) + weight = g.op("Constant", value_t=weight_value) + if symbolic_helper._is_none(bias): + bias_value = torch.zeros(normalized_shape, dtype=dtype) + bias = g.op("Constant", value_t=bias_value) + return g.op( + "LayerNormalization", + input, + weight, + bias, + epsilon_f=eps, + axis_i=axis, + ) + + +@_onnx_symbolic("quantized::layer_norm") +def quantized_layer_norm( + g: jit_utils.GraphContext, + x, + normalized_shape, + weight, + bias, + eps, + op_scale, + op_zero_point, +): + x, _, _, _ = symbolic_helper.dequantize_helper(g, x) + + output = layer_norm(g, x, normalized_shape, weight, bias, eps, False) + + return symbolic_helper.quantize_helper(g, output, op_scale, op_zero_point) + + +def _compute_edge_sizes(n_fft, window_size): + """Helper function to compute the sizes of the edges (left and right) + of a given window centered within an FFT size.""" + left = (n_fft - window_size) // 2 + right = n_fft - left - window_size + return left, right + + +@_onnx_symbolic("aten::stft") +@symbolic_helper.parse_args("v", "i", "i", "i", "v", "b", "b", "b", "b") +def stft( + g: jit_utils.GraphContext, + input: _C.Value, + n_fft: int, + hop_length: int | None = None, + win_length: int | None = None, + window: _C.Value | None = None, + normalized: bool = False, + onesided: bool | None = True, + return_complex: bool | None = False, + align_to_window: bool | None = None, +) -> _C.Value: + """Associates `torch.stft` with the `STFT` ONNX operator. + Note that torch.stft calls _VF.stft, without centering or padding options. + Hence, this function does not contain these two arguments. + See torch.stft source code for more info. + + Args: + g: Graph to write the ONNX representation into + input: Input tensor for the transformation + n_fft: FFT size + hop_length: Size of the hop. Defaults to `floot(n_fft // 4)` + win_length: Size of the analysis window. Defaults to `n_fft` + window: Analysis window. Defaults to a window of all ones + normalized: Whether to return a normalized STFT + onesided: Whether to return only half (+1) of the results, given the + symmetry of the STFT + return_complex: Whether to return the complex value (Note: Must be + `False` or `None`) + + Returns: + op: Operator for torch.stft associated with STFT (ONNX) + """ + # Checks + if return_complex: + raise errors.SymbolicValueError( + msg="STFT does not currently support complex types", value=input + ) + + if align_to_window is not None: + raise errors.SymbolicValueError( + msg="STFT does not currently support the align_to_window option", + value=input, + ) # TODO(#145944): add compatibility with align_to_window option. + + # Get STFT sizes + frame_step_value = hop_length if hop_length is not None else n_fft // 4 + frame_step_const = g.op( + "Constant", value_t=torch.tensor(frame_step_value, dtype=torch.int64) + ) + frame_length_const = g.op( + "Constant", value_t=torch.tensor(n_fft, dtype=torch.int64) + ) + + # Pre-process input if needed + signal = input + signal_rank = symbolic_helper._get_tensor_rank(signal) + if signal_rank == 1: + # Add batch dimension + signal = g.op( + "Unsqueeze", + signal, + g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64)), + ) + elif signal_rank is None or signal_rank > 2: + raise errors.SymbolicValueError( + msg="STFT can only take inputs of 1 [signal] or 2 [batch, signal] dimensions. " + f"Current rank of signal is {signal_rank}, please reduce it.", + value=input, + ) + + # Get window and make sure it's the same size as `win_length` or `n_fft` + # pyrefly: ignore [bad-argument-type] + n_win = symbolic_helper._get_tensor_dim_size(window, dim=0) + if n_win is not None: + win_length_default = win_length if win_length else n_fft + if n_win != win_length_default: + raise AssertionError( + "Analysis window size must equal `win_length` or `n_fft`. " + f"Please, set `win_length` or `n_fft` to match `window` size ({n_win})" + ) + + # Center window around zeros if needed (required by ONNX's STFT) + if n_win < n_fft: + left, right = _compute_edge_sizes(n_fft, n_win) + left_win = g.op("Constant", value_t=torch.zeros(left)) + right_win = g.op("Constant", value_t=torch.zeros(right)) + # pyrefly: ignore [bad-argument-type] + window = g.op("Concat", left_win, window, right_win, axis_i=0) + + # Create window, if needed + if symbolic_helper._is_none(window): + if win_length: + if win_length > n_fft: + raise errors.SymbolicValueError( + msg="The analysis window can't be longer than the size of the FFT. " + f"Please set `win_length` ({win_length}) to `n_fft` ({n_fft}) or less.", + value=input, + ) + + # Center window, if needed + left, right = _compute_edge_sizes(n_fft, win_length) + torch_window = torch.hstack( + (torch.zeros(left), torch.ones(win_length), torch.zeros(right)) + ) + else: + # Rectangle window + torch_window = torch.ones(n_fft) + if torch_window.shape[0] != n_fft: + raise AssertionError( + f"torch_window.shape[0]={torch_window.shape[0]} != n_fft={n_fft}" + ) + window = g.op("Constant", value_t=torch_window) + window = g.op( + "Cast", + # pyrefly: ignore [bad-argument-type] + window, + to_i=_type_utils.JitScalarType.from_value(signal).onnx_type(), + ) + + # Run STFT + result = g.op( + "STFT", + signal, + frame_step_const, + window, + frame_length_const, + onesided_i=1 if onesided is None or onesided else 0, + ) + + # Transpose to mimic torch.stft's behavior + result = g.op("Transpose", result, perm_i=[0, 2, 1, 3]) + + # Remove batch dimension, if needed + if signal_rank == 1: + result = g.op( + "Squeeze", + result, + g.op("Constant", value_t=torch.tensor([0], dtype=torch.int64)), + ) + + # Normalize, if needed + if normalized: + sqrt_nfft = torch.sqrt(torch.tensor(n_fft, dtype=signal.type().dtype())) + result = g.op("Div", result, g.op("Constant", value_t=sqrt_nfft)) + + return result diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset18.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset18.py new file mode 100644 index 0000000000000000000000000000000000000000..1d622bdd141e305c71b412e9bc45f234527a9f0c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset18.py @@ -0,0 +1,272 @@ +# mypy: allow-untyped-defs +"""This file exports ONNX ops for opset 18. + +Note [ONNX Operators that are added/updated in opset 18] + +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-18-of-the-default-onnx-operator-set +New operators: + BitwiseAnd + CenterCropPad + Col2Im + Mish + OptionalGetElement + OptionalHasElement + Pad + Resize + ScatterElements + ScatterND + Split +""" + +import functools +from collections.abc import Sequence + +import torch +from torch import _C +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, + symbolic_opset9 as opset9, +) + + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in symbolic_helper.py + +__all__ = [ + "col2im", +] + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=18) + + +@_onnx_symbolic("aten::__and_") +@_onnx_symbolic("aten::bitwise_and") +def __and_(g: jit_utils.GraphContext, self, other): + # do type promotion (scalars don't seem to apply) + args = [self, other] + # type promotion doesn't happen with torch.bitwise_and(tensor, scalar) + prom_args = [arg for arg in args if symbolic_helper._get_tensor_rank(arg)] + if len(prom_args) == 0: + prom_args = args + promotion_jit_type = symbolic_helper._type_promote_from_values(*prom_args) + self = symbolic_helper._maybe_cast_to_type(g, self, promotion_jit_type) + other = symbolic_helper._maybe_cast_to_type(g, other, promotion_jit_type) + if promotion_jit_type == _type_utils.JitScalarType.BOOL: + return g.op("And", self, other) + return g.op("BitwiseAnd", self, other) + + +@_onnx_symbolic("aten::col2im") +@symbolic_helper.parse_args("v", "v", "v", "is", "is", "is") +def col2im( + g, + input: _C.Value, + output_size: _C.Value, + kernel_size: _C.Value, + dilation: Sequence[int], + padding: Sequence[int], + stride: Sequence[int], +): + # convert [i0, i1, ..., in] into [i0, i0, i1, i1, ..., in, in] + adjusted_padding: list[int] = [] + for pad in padding: + adjusted_padding.extend(pad for _ in range(2)) + + num_dimensional_axis = symbolic_helper._get_tensor_sizes(output_size)[0] + if not adjusted_padding: + adjusted_padding = [0, 0] * num_dimensional_axis + + if not dilation: + dilation = [1] * num_dimensional_axis + + if not stride: + stride = [1] * num_dimensional_axis + + return g.op( + "Col2Im", + input, + output_size, + kernel_size, + dilations_i=dilation, + pads_i=adjusted_padding, + strides_i=stride, + ) + + +@_onnx_symbolic( + "aten::mean", decorate=[symbolic_helper._apply_params("ReduceMean", "mean")] +) +@_onnx_symbolic( + "aten::prod", + decorate=[ + symbolic_helper._apply_params( + "ReduceProd", "prod", allow_multi_dim_support=False + ) + ], +) +def _reduce_with_dtype(onnx_op: str, name: str, allow_multi_dim_support: bool = True): + return symbolic_helper._reduce_with_dtype_helper( + onnx_op, name, allow_multi_dim_support + ) + + +@_onnx_symbolic("aten::native_layer_norm") +@symbolic_helper.quantized_args(True, False, False, False) +@symbolic_helper.parse_args("v", "is", "v", "v", "f") +def _native_layer_norm( + g: jit_utils.GraphContext, + input: _C.Value, + normalized_shape: Sequence[int], + weight: _C.Value, + bias: _C.Value, + eps: float, +) -> tuple[_C.Value, _C.Value, _C.Value]: + return opset9.native_layer_norm(g, input, normalized_shape, weight, bias, eps) + + +@_onnx_symbolic("aten::glu") +@symbolic_helper.parse_args("v", "i") +def _glu(g: jit_utils.GraphContext, input, dim): + dim_size = symbolic_helper._get_tensor_dim_size(input, dim) + if dim_size is not None: + if dim_size % 2 != 0: + raise AssertionError(f"dim_size must be even, got {dim_size}") + + first, second = g.op("Split", input, axis_i=dim, num_outputs_i=2, outputs=2) + return g.op("Mul", first, g.op("Sigmoid", second)) + + +@_onnx_symbolic("aten::max") +# torch.max (same for torch.min) actually has two interfaces smashed together: +# torch.max(x, dim, keepdim) and torch.max(x, y) +# TODO(justinchuby): Support multiple quantized args in output +def max(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None): + return symbolic_helper._max_helper(g, self, dim_or_y, keepdim) + + +@_onnx_symbolic("aten::maximum") +@symbolic_helper.quantized_args(True, True) +def maximum(g: jit_utils.GraphContext, input, other): + # pyrefly: ignore [no-matching-overload] + return max(g, input, dim_or_y=other) + + +@_onnx_symbolic("aten::min") +# TODO(justinchuby): Support multiple quantized args in output +def min(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None): + return symbolic_helper._min_helper(g, self, dim_or_y, keepdim) + + +@_onnx_symbolic("aten::minimum") +@symbolic_helper.quantized_args(True, True) +def minimum(g: jit_utils.GraphContext, input, other): + # pyrefly: ignore [no-matching-overload] + return min(g, input, dim_or_y=other) + + +@_onnx_symbolic("aten::amax") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "is", "i") +def amax(g: jit_utils.GraphContext, self, dim, keepdim): + axes = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long)) + return g.op("ReduceMax", self, axes, keepdims_i=keepdim) + + +@_onnx_symbolic("aten::amin") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "is", "i") +def amin(g: jit_utils.GraphContext, self, dim, keepdim): + axes = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long)) + return g.op("ReduceMin", self, axes, keepdims_i=keepdim) + + +@_onnx_symbolic("aten::aminmax") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "v", "i") +def aminmax(g: jit_utils.GraphContext, self, dim, keepdim): + if not symbolic_helper._is_none(dim): + dim = symbolic_helper._get_const(dim, "i", "dim") + axes = g.op("Constant", value_t=torch.tensor([dim], dtype=torch.long)) + return g.op("ReduceMin", self, axes, keepdims_i=keepdim), g.op( + "ReduceMax", self, axes, keepdims_i=keepdim + ) + else: + return g.op("ReduceMin", self, keepdims_i=keepdim), g.op( + "ReduceMax", self, keepdims_i=keepdim + ) + + +@_onnx_symbolic("aten::var_mean") +def _var_mean(g: jit_utils.GraphContext, input, *args): + if len(args) == 1: + return symbolic_helper._var_mean_helper(g, input, None, args[0], None) + else: + return symbolic_helper._var_mean_helper(g, input, *args) + + +@_onnx_symbolic("aten::logsumexp") +@symbolic_helper.parse_args("v", "is", "i") +def _logsumexp(g: jit_utils.GraphContext, input, dim, keepdim): + if dim is None: + return g.op("ReduceLogSumExp", input, keepdims_i=0) + else: + axes = g.op("Constant", value_t=torch.tensor(dim, dtype=torch.long)) + return g.op("ReduceLogSumExp", input, axes, keepdims_i=keepdim) + + +@_onnx_symbolic("aten::linalg_matrix_norm") +@symbolic_helper.parse_args("v", "v", "is", "b", "v") +def _linalg_matrix_norm( + g: jit_utils.GraphContext, + self: torch._C.Value, + ord: torch._C.Value, + dim: list[int], + keepdim: bool, + dtype: torch._C.Value, +): + return opset9.linalg_matrix_norm(g, self, ord, dim, keepdim, dtype) + + +@_onnx_symbolic("aten::embedding_bag") +@symbolic_helper.parse_args("v", "v", "v", "i", "i", "i", "v", "i", "i") +def embedding_bag( + g: jit_utils.GraphContext, + embedding_matrix, + indices, + offsets, + scale_grad_by_freq, + mode, + sparse, + per_sample_weights, + include_last_offset, + padding_idx, +): + return symbolic_helper._embedding_bag_helper( + g, + embedding_matrix, + indices, + offsets, + scale_grad_by_freq, + mode, + sparse, + per_sample_weights, + include_last_offset, + padding_idx, + ) + + +@_onnx_symbolic("aten::linalg_vector_norm") +@symbolic_helper.parse_args("v", "f", "is", "b", "v") +def linalg_vector_norm( + g: jit_utils.GraphContext, + self: torch._C.Value, + ord: float, + dim: Sequence[int] | None, + keepdim: bool, + dtype: torch._C.Value, +): + return symbolic_helper._linalg_vector_norm_helper(g, self, ord, dim, keepdim, dtype) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset19.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset19.py new file mode 100644 index 0000000000000000000000000000000000000000..781bc2d200c7e7b6c60dcad726574acc7f54b2fd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset19.py @@ -0,0 +1,31 @@ +"""This file exports ONNX ops for opset 19. + +Note [ONNX Operators that are added/updated in opset 19] + +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-19-of-the-default-onnx-operator-set +New operators: +AveragePool +Cast +CastLike +Constant +DeformConv +DequantizeLinear +Equal +Identity +If +Loop +Pad +QuantizeLinear +Reshape +Resize +Scan +Shape +Size +""" + + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in symbolic_helper.py + +__all__: list[str] = [] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset20.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset20.py new file mode 100644 index 0000000000000000000000000000000000000000..8e8ca44a26a4e8d2829f800abab7a1c14d63658f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset20.py @@ -0,0 +1,95 @@ +# mypy: allow-untyped-defs +"""This file exports ONNX ops for opset 20. + +Note [ONNX Operators that are added/updated in opset 20] + +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +https://github.com/onnx/onnx/blob/main/docs/Changelog.md#version-20-of-the-default-onnx-operator-set +New operators: + AffineGrid + ConstantOfShape + DFT + Gelu + GridSample + ImageDecoder + IsInf + IsNaN + ReduceMax + ReduceMin + RegexFullMatch + StringConcat + StringSplit +""" + +import functools + +import torch.nn.functional as F +from torch import _C +from torch.onnx._internal.torchscript_exporter import ( + jit_utils, + registration, + symbolic_helper, +) + + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in symbolic_helper.py + +__all__ = ["_grid_sampler", "_affine_grid_generator", "gelu"] + + +def convert_grid_sample_mode(mode_s): + return ( + "linear" if mode_s == "bilinear" else "cubic" if mode_s == "bicubic" else mode_s + ) + + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=20) + + +@_onnx_symbolic("aten::grid_sampler") +@symbolic_helper.parse_args("v", "v", "i", "i", "b") +def _grid_sampler( + g: jit_utils.GraphContext, + input: _C.Value, + grid: _C.Value, + mode_enum: int, + padding_mode_enum: int, + align_corners: bool, +): + mode_s = {v: k for k, v in F.GRID_SAMPLE_INTERPOLATION_MODES.items()}[mode_enum] # type: ignore[call-arg, index] + # mode string changes at https://onnx.ai/onnx/operators/text_diff_GridSample_16_20.html + mode_s = convert_grid_sample_mode(mode_s) + padding_mode_s = {v: k for k, v in F.GRID_SAMPLE_PADDING_MODES.items()}[ # type: ignore[call-arg, index] + padding_mode_enum # type: ignore[index] + ] + return g.op( + "GridSample", + input, + grid, + align_corners_i=int(align_corners), + mode_s=mode_s, + padding_mode_s=padding_mode_s, + ) + + +@_onnx_symbolic("aten::affine_grid_generator") +@symbolic_helper.parse_args("v", "v", "b") +def _affine_grid_generator( + g: jit_utils.GraphContext, + theta: _C.Value, + size: _C.Value, + align_corners: bool, +): + return g.op( + "AffineGrid", + theta, + size, + align_corners_i=int(align_corners), + ) + + +@_onnx_symbolic("aten::gelu") +@symbolic_helper.parse_args("v", "s") +def gelu(g: jit_utils.GraphContext, self: _C.Value, approximate: str = "none"): + return g.op("Gelu", self, approximate_s=approximate) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset7.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset7.py new file mode 100644 index 0000000000000000000000000000000000000000..ae9a5039d397e6cc0f59987dac531ef9a66856fc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset7.py @@ -0,0 +1,73 @@ +# mypy: allow-untyped-defs +""" +Note [ONNX operators that are added/updated from opset 7 to opset 8] +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +New operators: + Expand + +Updated operators: + Min, Max, Sum, Mean: supports multidirectional broadcasting. + MaxPool: added optional indices output. + Scan +""" + +import functools +import warnings + +from torch.onnx._internal.torchscript_exporter import ( + jit_utils, + registration, + symbolic_helper, + symbolic_opset9 as opset9, +) + + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=7) + +block_listed_operators = ( + "scan", + "expand", + "expand_as", + "meshgrid", + "adaptive_max_pool1d", + "adaptive_max_pool2d", + "adaptive_max_pool3d", + "max_pool1d_with_indices", + "max_pool2d_with_indices", + "max_pool3d_with_indices", +) + + +# NOTE: max, min, sum, mean: broadcasting is not supported in opset 7. +# torch.max (same for torch.min) actually has two interfaces smashed together: +# torch.max(x, dim, keepdim) and torch.max(x, y) +@_onnx_symbolic("aten::max") +def max(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None): + # torch.max(input, other) + if keepdim is None and dim_or_y is not None: + warnings.warn( + "Multidirectional broadcasting is not supported in opset 7. " + "This might cause the onnx model to be incorrect, if inputs to max operators " + "have different shapes", + stacklevel=2, + ) + return opset9.max(g, self, dim_or_y, keepdim) + + +@_onnx_symbolic("aten::min") +def min(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None): + # torch.min(input, other) + if keepdim is None and dim_or_y is not None: + warnings.warn( + "Multidirectional broadcasting is not supported in opset 7. " + "This might cause the onnx model to be incorrect, if inputs to min operators " + "have different shapes", + stacklevel=2, + ) + return opset9.min(g, self, dim_or_y, keepdim) + + +for block_listed_op in block_listed_operators: + _onnx_symbolic(f"aten::{block_listed_op}")( + symbolic_helper._block_list_in_opset(block_listed_op) + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset8.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset8.py new file mode 100644 index 0000000000000000000000000000000000000000..921dd5d119f17681ce336f5913d6e2180a65b6c4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset8.py @@ -0,0 +1,470 @@ +# mypy: allow-untyped-defs +""" +Note [ONNX operators that are added/updated from opset 8 to opset 9] +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +New operators: + Compress + ConstantOfShape + EyeLike + MaxUnpool + OneHot + Sinh + Cosh + Asinh + Acosh + Atanh + Shrink + IsNaN + Sign + Erf + Scatter + Where + NonZero + TfIdfVectorizer + MeanVarianceNormalization + +Updated operators: + BatchNormalization: removed spatial attribute. + Greater, Less, Constant, MatMul, PRelu, Gemm, Flatten: more data types{integers} supported. + Cast: more data types{string} supported. + Upsample: moved scales from attribute to input. + Scan +""" + +import functools +import warnings + +import torch +from torch._C import _onnx as _C_onnx +from torch.onnx import errors +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, + symbolic_opset9 as opset9, +) + + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=8) + +block_listed_operators = ( + "nonzero", + "where", + "scatter", + "scatter_add", + "erf", + "sign", + "isnan", + "gather", + "arange", + "masked_fill", + "index_fill", + "index_copy", + "repeat_interleave", + "any", + "all", +) + +for block_listed_op in block_listed_operators: + _onnx_symbolic(f"aten::{block_listed_op}")( + symbolic_helper._block_list_in_opset(block_listed_op) + ) + + +@_onnx_symbolic( + "aten::upsample_nearest1d", + decorate=[symbolic_helper._apply_params("upsample_nearest1d", 3, "nearest")], +) +@_onnx_symbolic( + "aten::upsample_nearest2d", + decorate=[symbolic_helper._apply_params("upsample_nearest2d", 4, "nearest")], +) +@_onnx_symbolic( + "aten::upsample_nearest3d", + decorate=[symbolic_helper._apply_params("upsample_nearest3d", 5, "nearest")], +) +@_onnx_symbolic( + "aten::upsample_linear1d", + decorate=[symbolic_helper._apply_params("upsample_linear1d", 3, "linear")], +) +@_onnx_symbolic( + "aten::upsample_bilinear2d", + decorate=[symbolic_helper._apply_params("upsample_bilinear2d", 4, "linear")], +) +@_onnx_symbolic( + "aten::upsample_trilinear3d", + decorate=[symbolic_helper._apply_params("upsample_trilinear3d", 5, "linear")], +) +def _interpolate(name, dim, interpolate_mode): + def symbolic_fn(g, input, output_size, *args): + scales, align_corners = symbolic_helper._get_interpolate_attributes( + g, interpolate_mode, args + ) + symbolic_helper._interpolate_warning(interpolate_mode) + align_corners = symbolic_helper._maybe_get_scalar(align_corners) + if align_corners: + return symbolic_helper._unimplemented(name, "align_corners == True", input) + output_size = symbolic_helper._maybe_get_const(output_size, "is") + if symbolic_helper._is_value(output_size): + return symbolic_helper._unimplemented( + name, "torch._C.Value (output_size) indexing" + ) + if scales is None: + scales = [ + 1.0 + if i < 2 + else float(output_size[-(dim - i)]) + / float(input.type().sizes()[-(dim - i)]) + for i in range(dim) + ] + return g.op("Upsample", input, mode_s=interpolate_mode, scales_f=scales) + + return symbolic_fn + + +@_onnx_symbolic("aten::__interpolate") +def __interpolate( + g: jit_utils.GraphContext, + input, + size, + scale_factor, + mode, + align_corners, + recompute_scale_factor, + antialias, +): + align_corners = symbolic_helper._maybe_get_const(align_corners, "b") + if not symbolic_helper._is_none(align_corners) and align_corners: + return symbolic_helper._unimplemented("interpolate", "align_corners == True") + + if not symbolic_helper._is_none(scale_factor) and symbolic_helper._is_value( + scale_factor + ): + return symbolic_helper._unimplemented( + "interpolate", "dynamic scales in opset 8" + ) + + if not symbolic_helper._is_none(size) and symbolic_helper._is_value(size): + return symbolic_helper._unimplemented("interpolate", "dynamic size in opset 8") + + scales, mode = symbolic_helper._interpolate_get_scales_and_mode( + g, input, size, scale_factor, mode, align_corners + ) + return g.op("Upsample", input, mode_s=mode, scales_f=scales) + + +# NOTE: We should create a wrapper for this kind of operation, after resolving the shape/type propagation +# issue for "cast" operators. Some symbolic functions depend on shape information of input tensor, which +# is lost after casting. +def _try_cast_integer_to_float(g: jit_utils.GraphContext, *args): + floating_scalar_types = { + _type_utils.JitScalarType.HALF, + _type_utils.JitScalarType.FLOAT, + _type_utils.JitScalarType.DOUBLE, + } + old_type = None + # Cast the input tensor to Float if its scalarType is known and is not floating number. + # If casting is performed, return the old scalarType, otherwise return None. + arg0_type = _type_utils.JitScalarType.from_value( + args[0], _type_utils.JitScalarType.UNDEFINED + ) + if arg0_type != _type_utils.JitScalarType.UNDEFINED: + old_type = arg0_type + if old_type not in floating_scalar_types: + old_type = old_type.scalar_name() # type: ignore[assignment] + args = tuple( + g.op("Cast", arg, to_i=_C_onnx.TensorProtoDataType.FLOAT) + for arg in args + ) + else: + return (None,) + args + else: + warnings.warn( + "Only floating datatype is supported for these operators: " + "{Greater, Less, MatMul, PRelu, Gemm, Flatten}. This might cause " + "the onnx model to be incorrect, if inputs have integer datatypes.", + stacklevel=2, + ) + return (old_type,) + args + + +def _cast_to_type(g: jit_utils.GraphContext, input, to_type): + if to_type is None: + return input + return g.op("Cast", input, to_i=symbolic_helper.cast_pytorch_to_onnx[to_type]) + + +def _comparison_operator(g: jit_utils.GraphContext, input, other, op_name): + other = symbolic_helper._maybe_get_scalar(other) + other = symbolic_helper._if_scalar_type_as(other, input) + _, input, other = _try_cast_integer_to_float(g, input, other) + return g.op(op_name, input, other) + + +# NOTE: For symbolics {gt, lt, bmm, matmul, prelu, mm, addmm, view, flatten}, +# integer input type not supported in opset8. Cast to float if possible. +@_onnx_symbolic("aten::gt") +def gt(g: jit_utils.GraphContext, input, other): + return _comparison_operator(g, input, other, "Greater") + + +@_onnx_symbolic("aten::lt") +def lt(g: jit_utils.GraphContext, input, other): + return _comparison_operator(g, input, other, "Less") + + +@_onnx_symbolic("aten::bmm") +def bmm(g: jit_utils.GraphContext, self, other): + if symbolic_helper._try_get_scalar_type(self): + old_type, self, other = _try_cast_integer_to_float(g, self, other) + return _cast_to_type(g, g.op("MatMul", self, other), old_type) + else: + return g.op("MatMul", self, other) + + +@_onnx_symbolic("aten::matmul") +def matmul(g: jit_utils.GraphContext, self, other): + return bmm(g, self, other) + + +@_onnx_symbolic("aten::prelu") +def prelu(g: jit_utils.GraphContext, self, weight): + self_rank = symbolic_helper._get_tensor_rank(self) + weight_sizes = symbolic_helper._get_tensor_sizes(weight) + if self_rank is not None and self_rank > 2: + weight = g.op("Unsqueeze", weight, axes_i=list(range(1, self_rank - 1))) + elif self_rank == 0 and weight_sizes == [1]: + # self and weight are both scalar but weight has rank == 1, squeeze weight. + weight = symbolic_helper._squeeze_helper(g, weight, [0]) + if symbolic_helper._try_get_scalar_type(self): + old_type, self, weight = _try_cast_integer_to_float(g, self, weight) + return _cast_to_type(g, g.op("PRelu", self, weight), old_type) + else: + return g.op("PRelu", self, weight) + + +@_onnx_symbolic("aten::mm") +def mm(g: jit_utils.GraphContext, self, other): + # Create a dummy C tensor. Only needed for API purposes, the value is + # since beta = 0 + scalar_type = symbolic_helper._try_get_scalar_type(self, other) + if scalar_type is None: + raise errors.SymbolicValueError( + "mm can only operate on tensors with known types", self + ) + zero_constant = g.op( + "Constant", + value_t=torch.tensor([0], dtype=scalar_type.dtype()), + ) + + if symbolic_helper._try_get_scalar_type(self): + old_type, self, other, zero_constant = _try_cast_integer_to_float( + g, self, other, zero_constant + ) + return _cast_to_type( + g, + g.op("Gemm", self, other, zero_constant, beta_f=0.0, alpha_f=1.0), + old_type, + ) + return g.op("Gemm", self, other, zero_constant, beta_f=0.0, alpha_f=1.0) + + +@_onnx_symbolic("aten::addmm") +@symbolic_helper.parse_args("v", "v", "v", "t", "t") +def addmm(g: jit_utils.GraphContext, self, mat1, mat2, beta, alpha): + if symbolic_helper._try_get_scalar_type(self): + old_type, self, mat1, mat2 = _try_cast_integer_to_float(g, self, mat1, mat2) + return _cast_to_type( + g, + g.op( + "Gemm", + mat1, + mat2, + self, + beta_f=symbolic_helper._scalar(beta), + alpha_f=symbolic_helper._scalar(alpha), + ), + old_type, + ) + else: + return g.op( + "Gemm", + mat1, + mat2, + self, + beta_f=symbolic_helper._scalar(beta), + alpha_f=symbolic_helper._scalar(alpha), + ) + + +@_onnx_symbolic("aten::flatten") +def flatten(g: jit_utils.GraphContext, input, start_dim, end_dim): + start_dim_i = symbolic_helper._get_const(start_dim, "i", "start_dim") + end_dim_i = symbolic_helper._get_const(end_dim, "i", "end_dim") + + dim = input.type().dim() + if end_dim_i < 0: + end_dim_i = dim + end_dim_i + # use ONNX's Flatten operator for cases where the output shape is 2D + if start_dim_i == 1 and end_dim_i == dim - 1: + if symbolic_helper._try_get_scalar_type(input): + old_type, input = _try_cast_integer_to_float(g, input) + return _cast_to_type( + g, g.op("Flatten", input, axis_i=start_dim_i), old_type + ) + else: + return g.op("Flatten", input, axis_i=start_dim_i) + if start_dim_i == 0 and end_dim_i == dim - 2: + if symbolic_helper._try_get_scalar_type(input): + old_type, input = _try_cast_integer_to_float(g, input) + return _cast_to_type( + g, g.op("Flatten", input, axis_i=end_dim_i + 1), old_type + ) + else: + return g.op("Flatten", input, axis_i=end_dim_i + 1) + + return opset9.flatten(g, input, start_dim, end_dim) + + +def _constant_fill(g: jit_utils.GraphContext, sizes, dtype: int, const_value): + if dtype is None: + scalar_type = _type_utils.JitScalarType.FLOAT + else: + scalar_type = _type_utils.JitScalarType(dtype) + if not scalar_type.dtype().is_floating_point: + result = g.op( + "ConstantFill", + sizes, + dtype_i=_type_utils.JitScalarType.FLOAT.onnx_type(), + input_as_shape_i=1, + value_f=const_value, + ) + return g.op("Cast", result, to_i=scalar_type.onnx_type()) + else: + return g.op( + "ConstantFill", + sizes, + dtype_i=scalar_type.onnx_type(), + input_as_shape_i=1, + value_f=const_value, + ) + + +@_onnx_symbolic("aten::empty") +@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v") +def empty( + g: jit_utils.GraphContext, + sizes, + dtype, + layout, + device, + pin_memory=False, + memory_format=None, +): + return zeros(g, sizes, dtype, layout, device, pin_memory) + + +@_onnx_symbolic("aten::empty_like") +@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v") +def empty_like( + g: jit_utils.GraphContext, + input, + dtype, + layout, + device, + pin_memory=False, + memory_format=None, +): + return zeros_like(g, input, dtype, layout, device, pin_memory) + + +@_onnx_symbolic("aten::zeros") +@symbolic_helper.parse_args("v", "i", "v", "v", "v") +def zeros(g: jit_utils.GraphContext, sizes, dtype, layout, device, pin_memory=False): + # NOTE: no way to set device and layout in ONNX, so we ignore it + return _constant_fill(g, sizes, dtype, 0) + + +@_onnx_symbolic("aten::zeros_like") +@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v") +def zeros_like( + g: jit_utils.GraphContext, + input, + dtype, + layout, + device, + pin_memory=False, + memory_format=None, +): + shape = g.op("Shape", input) + return _constant_fill(g, shape, dtype, 0) + + +@_onnx_symbolic("aten::ones") +@symbolic_helper.parse_args("v", "i", "v", "v", "v") +def ones(g: jit_utils.GraphContext, sizes, dtype, layout, device, pin_memory=False): + return _constant_fill(g, sizes, dtype, 1) + + +@_onnx_symbolic("aten::ones_like") +@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v") +def ones_like( + g: jit_utils.GraphContext, + input, + dtype, + layout, + device, + pin_memory=False, + memory_format=None, +): + shape = g.op("Shape", input) + return _constant_fill(g, shape, dtype, 1) + + +@_onnx_symbolic("aten::full") +def full( + g: jit_utils.GraphContext, sizes, value, dtype, layout, device, pin_memory=False +): + const_value = symbolic_helper._maybe_get_const(value, "t") + if symbolic_helper._is_value(const_value): + tmp = zeros(g, sizes, dtype, layout, device) + return opset9.add(g, tmp, value, g.op("Constant", value_t=torch.tensor(1))) + else: + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + return _constant_fill(g, sizes, dtype, const_value) + + +@_onnx_symbolic("aten::full_like") +@symbolic_helper.parse_args("v", "f", "i", "v", "v", "v", "v") +def full_like( + g: jit_utils.GraphContext, + input, + fill_value, + dtype, + layout, + device, + pin_memory=False, + memory_format=None, +): + shape = g.op("Shape", input) + return _constant_fill(g, shape, dtype, fill_value) + + +@_onnx_symbolic("aten::repeat") +def repeat(g: jit_utils.GraphContext, self, repeats): + if not symbolic_helper._is_value(repeats): + repeats = g.op("Constant", value_t=torch.LongTensor(repeats)) + if symbolic_helper._is_packed_list(repeats): + repeat_size_len = len(symbolic_helper._unpack_list(repeats)) + else: + const_repeats = symbolic_helper._maybe_get_const(repeats, "is") + repeat_size_len = len(const_repeats) + if self.isCompleteTensor(): + sizes = self.type().sizes() + diff_dims = repeat_size_len - len(sizes) + if diff_dims > 0: + self = opset9.view( + g, self, g.op("Constant", value_t=torch.tensor([1] * diff_dims + sizes)) + ) + return g.op("Tile", self, repeats) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset9.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset9.py new file mode 100644 index 0000000000000000000000000000000000000000..d781c0dfbdf63d4882339c75ebbaee893aabb275 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/symbolic_opset9.py @@ -0,0 +1,6668 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +# mypy: disable-error-code=arg-type +"""This file exports ONNX ops for opset 9. + +Opset 9 is supported by ONNX release 1.4.1 +release on 01/23/19 +""" + +from __future__ import annotations + +import builtins +import functools +import math +import sys +import warnings +from typing import TYPE_CHECKING + +import torch +import torch._C._onnx as _C_onnx +import torch.nn.modules.utils +import torch.onnx +from torch import _C +from torch.onnx import _constants, errors +from torch.onnx._internal.torchscript_exporter import ( + _type_utils, + jit_utils, + registration, + symbolic_helper, +) +from torch.onnx._internal.torchscript_exporter._globals import GLOBALS + + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + from torch.types import Number + +# EDITING THIS FILE? READ THIS FIRST! +# see Note [Edit Symbolic Files] in README.md + +__all__ = [ + "abs", + "acos", + "add", + "addcmul", + "addmm", + "alias", + "amax", + "amin", + "aminmax", + "arange", + "argmax", + "argmin", + "as_strided", + "as_tensor", + "asin", + "atan", + "atan2", + "baddbmm", + "batch_norm", + "bernoulli", + "bitwise_not", + "bitwise_or", + "bmm", + "broadcast_tensors", + "broadcast_to", + "bucketize", + "cat", + "cdist", + "ceil", + "clamp_max", + "clamp_min", + "clamp", + "clone", + "constant_pad_nd", + "contiguous", + "conv_tbc", + "conv_transpose1d", + "conv_transpose2d", + "conv_transpose3d", + "conv1d", + "conv2d", + "conv3d", + "convert_element_type", + "convolution", + "cos", + "cosine_similarity", + "cross", + "cumsum", + "detach", + "dim", + "div", + "dot", + "dropout", + "elu", + "embedding_bag", + "embedding", + "empty_like", + "empty", + "eq", + "erf", + "exp", + "expand_as", + "expand", + "eye", + "fill", + "flatten", + "floor_divide", + "floor", + "floordiv", + "frobenius_norm", + "full_like", + "full", + "gather", + "ge", + "gelu", + "get_pool_ceil_padding", + "glu", + "group_norm", + "gt", + "hann_window", + "hardshrink", + "hardsigmoid", + "hardswish", + "hardtanh", + "index_add", + "index_copy", + "index_fill", + "index_put", + "index_select", + "index", + "instance_norm", + "is_floating_point", + "is_pinned", + "isnan", + "item", + "kl_div", + "layer_norm", + "le", + "leaky_relu", + "lerp", + "lift", + "linalg_cross", + "linalg_matrix_norm", + "linalg_norm", + "linalg_vector_norm", + "linear", + "linspace", + "log_sigmoid", + "log_softmax", + "log", + "log10", + "log1p", + "log2", + "logical_and", + "logical_not", + "logical_or", + "logical_xor", + "logit", + "logsumexp", + "lstm_cell", + "lstm", + "lt", + "masked_fill", + "masked_fill_", + "matmul", + "max_pool1d_with_indices", + "max_pool2d_with_indices", + "max_pool3d_with_indices", + "max", + "maximum", + "meshgrid", + "min", + "minimum", + "mish", + "mm", + "movedim", + "mse_loss", + "mul", + "multinomial", + "mv", + "narrow", + "native_layer_norm", + "ne", + "neg", + "new_empty", + "new_full", + "new_ones", + "new_zeros", + "nonzero_numpy", + "nonzero", + "norm", + "numel", + "numpy_T", + "one_hot", + "ones_like", + "ones", + "onnx_placeholder", + "pad", + "pairwise_distance", + "permute", + "pixel_shuffle", + "pixel_unshuffle", + "pow", + "prelu", + "prim_constant_chunk", + "prim_constant_split", + "prim_constant", + "prim_data", + "prim_device", + "prim_dtype", + "prim_if", + "prim_layout", + "prim_list_construct", + "prim_list_unpack", + "prim_loop", + "prim_max", + "prim_min", + "prim_shape", + "prim_tolist", + "prim_tuple_construct", + "prim_type", + "prim_unchecked_cast", + "prim_uninitialized", + "rand_like", + "rand", + "randint_like", + "randint", + "randn_like", + "randn", + "reciprocal", + "reflection_pad", + "relu", + "relu6", + "remainder", + "repeat_interleave", + "repeat", + "replication_pad", + "reshape_as", + "reshape", + "roll", + "rrelu", + "rsqrt", + "rsub", + "scalar_tensor", + "scatter_add", + "scatter", + "select", + "selu", + "sigmoid", + "sign", + "silu", + "sin", + "size", + "slice", + "softmax", + "softplus", + "softshrink", + "sort", + "split_with_sizes", + "split", + "sqrt", + "square", + "squeeze", + "stack", + "std_mean", + "std", + "sub", + "t", + "take", + "tan", + "tanh", + "tanhshrink", + "tensor", + "threshold", + "to", + "topk", + "transpose", + "true_divide", + "type_as", + "unbind", + "unfold", + "unsafe_chunk", + "unsafe_split_with_sizes", + "unsafe_split", + "unsqueeze", + "unsupported_complex_operators", + "noop_complex_operators", + "unused", + "var_mean", + "var", + "view_as", + "view", + "where", + "wrap_logical_op_with_cast_to", + "wrap_logical_op_with_negation", + "zeros_like", + "zeros", + "zero", +] + + +_onnx_symbolic = functools.partial(registration.onnx_symbolic, opset=9) + + +def _export(name: str): + """Exports the function in the current global namespace.""" + + def wrapper(func): + globals()[name] = func + __all__.append(name) + return func + + return wrapper + + +def unused(g): + """Represents "missing" optional inputs.""" + n = g.op("prim::Constant") + n.setType(_C.OptionalType.ofTensor()) + return n + + +@_onnx_symbolic("aten::_shape_as_tensor") +def _shape_as_tensor(g: jit_utils.GraphContext, input): + return g.op("Shape", input) + + +@_onnx_symbolic("aten::_reshape_from_tensor") +def _reshape_from_tensor(g: jit_utils.GraphContext, input, shape): + if isinstance(shape, list): + shape = g.op("Concat", *shape, axis_i=0) + return reshape(g, input, shape) + + +@_onnx_symbolic("aten::reshape") +@symbolic_helper.quantized_args(True) +def reshape(g: jit_utils.GraphContext, self, shape): + return symbolic_helper._reshape_helper(g, self, shape) + + +@_onnx_symbolic("aten::reshape_as") +@symbolic_helper.quantized_args(True) +def reshape_as(g: jit_utils.GraphContext, self, other): + shape = g.op("Shape", other) + return reshape(g, self, shape) + + +@_onnx_symbolic("aten::add") +def add(g: jit_utils.GraphContext, self, other, alpha=None): + """ + This function takes the add function and returns the corresponding ONNX operator. + + This function is not meant to be called directly by the user. + + Args: + g (GraphContext): The graph context. + self (Tensor): The first operand. + other (Tensor): The second operand. + alpha (float, optional): The scaling factor for the second operand. Defaults to None. + + Returns: + ONNX operator. + """ + if symbolic_helper._is_value(self) and symbolic_helper._is_tensor_list(self): + return symbolic_helper._onnx_opset_unsupported_detailed( + "Add", 9, 11, "Add between list of tensors not supported", self + ) + if alpha and symbolic_helper._scalar(symbolic_helper._maybe_get_scalar(alpha)) != 1: + other = g.op("Mul", other, alpha) + return g.op("Add", self, other) + + +@_onnx_symbolic("aten::sub") +def sub(g: jit_utils.GraphContext, self, other, alpha=None): + """ + Consumes sub function and returns the corresponding ONNX operator. + + This function is not meant to be called directly by the user. + + Args: + g (GraphContext): The graph context. + self (Tensor): The first operand. + other (Tensor): The second operand. + alpha (Optional[Tensor]): A scaling factor to apply to the second operand. + If `alpha` is not provided, it defaults to 1. + + Returns: + ONNX operator + """ + if alpha and symbolic_helper._scalar(symbolic_helper._maybe_get_scalar(alpha)) != 1: + other = g.op("Mul", other, alpha) + return g.op("Sub", self, other) + + +@_onnx_symbolic("aten::rsub") +def rsub(g: jit_utils.GraphContext, self, other, alpha=None): + return sub(g, other, self, alpha=alpha) + + +@_onnx_symbolic("aten::mul") +def mul(g: jit_utils.GraphContext, self, other): + if symbolic_helper._is_bool(self) and symbolic_helper._is_bool(other): + # ONNX Mul doesn't support Boolean, so use And as an equivalent operator. + return g.op("And", self, other) + else: + return g.op("Mul", self, other) + + +@_onnx_symbolic("aten::div") +def div(g: jit_utils.GraphContext, self, other, *args): + if len(args) == 0: + return true_divide(g, self, other) + else: + return _div_rounding_mode(g, self, other, *args) + + +@_onnx_symbolic("aten::addcmul") +@symbolic_helper.parse_args("v", "v", "v", "f") +def addcmul(g: jit_utils.GraphContext, self, tensor1, tensor2, value=1.0): + value_tens = g.op("Constant", value_t=torch.tensor([value])) + return add(g, self, mul(g, mul(g, tensor1, tensor2), value_tens)) + + +@symbolic_helper.parse_args("v", "v", "s") +def _div_rounding_mode(g: jit_utils.GraphContext, self, other, rounding_mode): + if rounding_mode is None: + return true_divide(g, self, other) + elif rounding_mode == "floor": + return _floor_divide(g, self, other) + elif rounding_mode == "trunc": + return _trunc_divide(g, self, other) + else: + raise errors.SymbolicValueError( + f'Unsupported rounding mode: "{rounding_mode}". Expected None, "floor" or "trunc"', + self, + ) + + +def _trunc_divide(g: jit_utils.GraphContext, self, other): + out = g.op("Div", self, other) + # the correct operation is truncate, which is not supported in ONNX, + # we cannot call floor since it will behave differently for negative numbers + # (eg. -0.1 should become -0 ) + # - if scalar_type information are not available, assume that + # we need to call floor (treat as float) + out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.INT64) + + # Matching PyTorch's behavior: + # - if self is fp the output's type is self's type + # - if self is not fp and other is fp, the output is of type JitScalarType.FLOAT + # - self is not fp and other is not fp, the output's type is self's output type + # - the output type defaults to Float + scalar_type = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.UNDEFINED + ) + if scalar_type != _type_utils.JitScalarType.UNDEFINED: + if not symbolic_helper._is_fp(self) and symbolic_helper._is_fp(other): + out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.FLOAT) + else: + out = g.op( + "Cast", + out, + to_i=scalar_type.onnx_type(), + ) + else: + out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.FLOAT) + return out + + +def _floor_divide(g: jit_utils.GraphContext, self, other): + if symbolic_helper._is_fp(self) or symbolic_helper._is_fp(other): + out = true_divide(g, self, other) + return g.op("Floor", out) + else: + # Integer division does truncation rounding + div = g.op("Div", self, other) + # Division is negative if: self < 0 != other < 0 + zero = g.op("Constant", value_t=torch.tensor(0, dtype=torch.int64)) + negative = g.op( + "Xor", + symbolic_helper._lt_helper(g, self, zero), + symbolic_helper._lt_helper(g, other, zero), + ) + + # For negative numbers with self % other != 0, subtract 1 to round down instead of up + mod = g.op("Sub", self, g.op("Mul", div, other)) + fixup_mask = g.op("And", negative, g.op("Not", g.op("Equal", mod, zero))) + + one = g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64)) + fixup = g.op("Mul", fixup_mask, one) + return g.op("Sub", div, fixup) + + +@_onnx_symbolic("aten::floor_divide") +def floor_divide(g: jit_utils.GraphContext, self, other): + # Deprecated behavior, floor_divide actually truncates + return _trunc_divide(g, self, other) + + +@_onnx_symbolic("aten::floordiv") +def floordiv(g: jit_utils.GraphContext, self, other): + return floor_divide(g, self, other) + + +@_onnx_symbolic("aten::true_divide") +def true_divide(g: jit_utils.GraphContext, self, other): + """Division where both inputs are cast to floating types + + If both inputs are floating, performs div as usual + If only one input is a floating type, the other input is cast to its type + If neither input is a floating type, both inputs are cast to the default scalar type + """ + + # Case 1: either values are floating + # Performs div as usual. + # Implicit casting will be handled in scalar type analysis pass. + if symbolic_helper._is_fp(self) or symbolic_helper._is_fp(other): + return g.op("Div", self, other) + + # Case 2: neither is floating + # Casts both inputs to the default scalar type + scalar_type = torch.get_default_dtype() + onnx_scalar_type = _C_onnx.TensorProtoDataType.FLOAT + if scalar_type is not torch.float and scalar_type is not torch.double: + raise AssertionError(f"Expected float or double, got {scalar_type}") + if torch.get_default_dtype() is torch.double: + onnx_scalar_type = _C_onnx.TensorProtoDataType.DOUBLE + + self = g.op("Cast", self, to_i=onnx_scalar_type) + other = g.op("Cast", other, to_i=onnx_scalar_type) + return g.op("Div", self, other) + + +@_onnx_symbolic("aten::reciprocal") +def reciprocal(g: jit_utils.GraphContext, self): + # torch.reciprocal implicitly casts to float, so we do the same. + if not symbolic_helper._is_fp(self): + self = g.op("Cast", self, to_i=_C_onnx.TensorProtoDataType.FLOAT) + return g.op("Reciprocal", self) + + +@_onnx_symbolic("aten::cat") +@symbolic_helper.parse_args("v", "i") +def cat(g: jit_utils.GraphContext, tensor_list, dim): + """Implement concatenation of pytorch tensors in ONNX along the specified `dim` dimension. + + Parameters: + g (jit_utils.GraphContext): Graph context. + tensor_list (List[torch.Tensor]): List of tensors to concatenate. + dim (int): Dimension along which to concatenate the tensors. + + Returns: + ONNX graph node representing the concatenated tensor. + """ + tensors = symbolic_helper._unpack_list(tensor_list) + # torch.cat ignores empty tensors such as `torch.Tensor([])` + # These needs to be removed as input from ONNX's concat too, otherwise shape inference + # will likely fail due to inputs with different ranks (0 for empty tensor, > 0 for anything else) + nonempty_tensors = [] + for t in tensors: + if symbolic_helper._is_constant(t) and not symbolic_helper._get_tensor_dim_size( + t, 0 + ): + continue + nonempty_tensors.append(t) + if len(nonempty_tensors) == 0: + raise AssertionError("nonempty_tensors must not be empty") + if not all( + symbolic_helper._get_tensor_rank(nonempty_tensors[0]) is None + or symbolic_helper._get_tensor_rank(t) is None + or symbolic_helper._get_tensor_rank(t) + == symbolic_helper._get_tensor_rank(nonempty_tensors[0]) + for t in nonempty_tensors + ): + raise AssertionError("All tensors must have the same rank") + tensor_list.node().removeAllInputs() + for t in nonempty_tensors: + tensor_list.node().addInput(t) + + tensors = symbolic_helper._unpack_list(tensor_list) + return g.op("Concat", *tensors, axis_i=dim) + + +@_onnx_symbolic("aten::stack") +@symbolic_helper.parse_args("v", "i") +def stack(g: jit_utils.GraphContext, tensor_list, dim): + unsqueezed = [ + symbolic_helper._unsqueeze_helper(g, t, [dim]) + for t in symbolic_helper._unpack_list(tensor_list) + ] + return g.op("Concat", *unsqueezed, axis_i=dim) + + +@_onnx_symbolic("aten::list") +def _list(g: jit_utils.GraphContext, self): + return self + + +@_onnx_symbolic("aten::mm") +def mm(g: jit_utils.GraphContext, self, other): + # Create a dummy C tensor. Only needed for API purposes, the value is + # since beta = 0 + C = g.op("Constant", value_t=torch.tensor([1])) + return g.op("Gemm", self, other, C, beta_f=0.0, alpha_f=1.0) + + +@_onnx_symbolic("aten::bmm") +def bmm(g: jit_utils.GraphContext, self, other): + return g.op("MatMul", self, other) + + +@_onnx_symbolic("aten::matmul") +def matmul(g: jit_utils.GraphContext, self, other): + return g.op("MatMul", self, other) + + +@_onnx_symbolic("aten::addmm") +@symbolic_helper.parse_args("v", "v", "v", "t", "t") +def addmm(g: jit_utils.GraphContext, self, mat1, mat2, beta, alpha): + scalar_type = None + self_scalar_type = symbolic_helper._try_get_scalar_type(self) + mat1_scalar_type = symbolic_helper._try_get_scalar_type(mat1) + mat2_scalar_type = symbolic_helper._try_get_scalar_type(mat2) + if self_scalar_type is not None: + scalar_type = self_scalar_type + elif mat1_scalar_type is not None: + scalar_type = mat1_scalar_type + elif mat2_scalar_type is not None: + scalar_type = mat2_scalar_type + + mat1_rank = symbolic_helper._get_tensor_rank(mat1) + mat2_rank = symbolic_helper._get_tensor_rank(mat2) + + def is_not_none_nor(v, u): + return v is not None and v != u + + if scalar_type is not None and ( + is_not_none_nor(mat1_rank, 2) or is_not_none_nor(mat2_rank, 2) + ): + res1 = g.op("MatMul", mat1, mat2) + res2 = self + + alpha = symbolic_helper._scalar(alpha) + beta = symbolic_helper._scalar(beta) + + if alpha != 1: + alpha = g.op( + "Constant", value_t=torch.tensor(alpha, dtype=scalar_type.dtype()) + ) + res1 = g.op("Mul", res1, alpha) + if beta != 1: + beta = g.op( + "Constant", + value_t=torch.tensor( + symbolic_helper._scalar(beta), dtype=scalar_type.dtype() + ), + ) + res2 = g.op("Mul", res2, beta) + + return g.op("Add", res1, res2) + + return g.op( + "Gemm", + mat1, + mat2, + self, + beta_f=symbolic_helper._scalar(beta), + alpha_f=symbolic_helper._scalar(alpha), + ) + + +@_onnx_symbolic("aten::neg") +def neg(g: jit_utils.GraphContext, self): + return g.op("Neg", self) + + +@_onnx_symbolic("aten::sqrt") +def sqrt(g: jit_utils.GraphContext, self): + if _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.UNDEFINED + ) in { + _type_utils.JitScalarType.UINT8, + _type_utils.JitScalarType.INT8, + _type_utils.JitScalarType.INT16, + _type_utils.JitScalarType.INT, + _type_utils.JitScalarType.INT64, + }: + # torch converts all int inputs to sqrt to float + self = g.op("Cast", self, to_i=_C_onnx.TensorProtoDataType.FLOAT) + + return g.op("Sqrt", self) + + +@_onnx_symbolic("aten::rsqrt") +def rsqrt(g: jit_utils.GraphContext, self): + return g.op( + "Div", symbolic_helper._if_scalar_type_as(torch.ones(1), self), sqrt(g, self) + ) + + +@_onnx_symbolic("aten::tanh") +# Fixed scale and zero_point, discovered from aten/src/ATen/native/quantized/cpu/qtanh.cpp +@symbolic_helper.quantized_args(True, scale=2.0 / 256.0, zero_point=128) +def tanh(g: jit_utils.GraphContext, self): + return g.op("Tanh", self) + + +@_onnx_symbolic("aten::sin") +def sin(g: jit_utils.GraphContext, self): + return g.op("Sin", self) + + +@_onnx_symbolic("aten::cos") +def cos(g: jit_utils.GraphContext, self): + return g.op("Cos", self) + + +@_onnx_symbolic("aten::tan") +def tan(g: jit_utils.GraphContext, self): + return g.op("Tan", self) + + +@_onnx_symbolic("aten::asin") +def asin(g: jit_utils.GraphContext, self): + return g.op("Asin", self) + + +@_onnx_symbolic("aten::acos") +def acos(g: jit_utils.GraphContext, self): + return g.op("Acos", self) + + +@_onnx_symbolic("aten::atan") +def atan(g: jit_utils.GraphContext, self): + return g.op("Atan", self) + + +@_onnx_symbolic("aten::atan2") +def atan2(g: jit_utils.GraphContext, self, other): + # self is y, and other is x on coordinate + slope = g.op("Div", self, other) + atan = g.op("Atan", slope) + const_zero = g.op("Constant", value_t=torch.tensor(0)) + const_pi = g.op("Constant", value_t=torch.tensor(math.pi)) + + condition_second_or_third_quadrant = g.op("Greater", self, const_zero) + second_third_quadrant = g.op( + "Where", + condition_second_or_third_quadrant, + g.op("Add", atan, const_pi), + g.op("Sub", atan, const_pi), + ) + + condition_14_or_23_quadrant = g.op("Less", other, const_zero) + result = g.op("Where", condition_14_or_23_quadrant, second_third_quadrant, atan) + + return result + + +@_onnx_symbolic("aten::sigmoid") +# Fixed scale and zero_point, discovered from aten/src/ATen/native/quantized/cpu/qsigmoid.cpp +@symbolic_helper.quantized_args(True, scale=1.0 / 256.0, zero_point=0) +def sigmoid(g: jit_utils.GraphContext, self): + """Converts the corresponding PyTorch function into ONNX operators. + + It is not meant to be called directly by a user. + + Args: + g (jit_utils.GraphContext): Graph context. + self (Tensor): the input tensor. + Returns: + ONNX operator + """ + return g.op("Sigmoid", self) + + +@_onnx_symbolic("aten::sign") +def sign(g: jit_utils.GraphContext, self): + return g.op("Sign", self) + + +@symbolic_helper.quantized_args(True) +def _slice(g: jit_utils.GraphContext, input, axes, starts, ends): + if len(starts) != len(ends): + raise AssertionError(f"len(starts)={len(starts)} != len(ends)={len(ends)}") + if len(starts) == 1 and starts[0] == 0 and ends[0] == _constants.INT64_MAX: + return input + return g.op("Slice", input, axes_i=axes, starts_i=starts, ends_i=ends) + + +@_onnx_symbolic( + "aten::sum", decorate=[symbolic_helper._apply_params("ReduceSum", "sum")] +) +@_onnx_symbolic( + "aten::mean", decorate=[symbolic_helper._apply_params("ReduceMean", "mean")] +) +# torch.prod does not support multidimensional "dim" +@_onnx_symbolic( + "aten::prod", + decorate=[ + symbolic_helper._apply_params( + "ReduceProd", "prod", allow_multi_dim_support=False + ) + ], +) +def _reduce_with_dtype(onnx_op: str, name: str, allow_multi_dim_support: bool = True): + return symbolic_helper._reduce_with_dtype_helper( + onnx_op, name, allow_multi_dim_support + ) + + +@_onnx_symbolic("aten::cumsum") +@symbolic_helper.parse_args("v", "i", "none") +def cumsum(g: jit_utils.GraphContext, input, dim, dtype) -> None: + symbolic_helper._onnx_opset_unsupported("cumsum", 9, 11, input) + + +@_onnx_symbolic("aten::_sample_dirichlet") +def _sample_dirichlet(g: jit_utils.GraphContext, self, generator): + return symbolic_helper._onnx_unsupported("_sample_dirichlet", self) + + +@_onnx_symbolic("aten::_standard_gamma") +def _standard_gamma(g: jit_utils.GraphContext, self, generator): + return symbolic_helper._onnx_unsupported("_standard_gamma", self) + + +@_onnx_symbolic("aten::t") +def t(g: jit_utils.GraphContext, self): + rank = symbolic_helper._get_tensor_rank(self) + if rank is None or rank < 2: + # The transpose of a 1d or 0d tensor is itself. ONNX does not define the behavior + # clearly and onnxruntime fails on these cases. So we add an Identity node to + # mirror the behavior of eager mode. + return g.op("Identity", self) + return g.op("Transpose", self, perm_i=(1, 0)) + + +@_onnx_symbolic("aten::numpy_T") +@symbolic_helper.quantized_args(True) +def numpy_T(g: jit_utils.GraphContext, input): + ndim = symbolic_helper._get_tensor_rank(input) + if ndim is None: + raise AssertionError("ndim must be non-None") + perm = list(reversed(range(ndim))) + return g.op("Transpose", input, perm_i=perm) + + +@_onnx_symbolic("aten::expand") +@symbolic_helper.quantized_args(True) +def expand(g: jit_utils.GraphContext, self, size, implicit): + """Implement the expand function for a pytorch tensor in ONNX according to specified `size`""" + size = symbolic_helper._maybe_get_const(size, "is") + if not symbolic_helper._is_value(size): + size = g.op("Constant", value_t=torch.LongTensor(size)) + elif symbolic_helper._is_packed_list(size): + # Expand with -1 dim value means dim is unchanged. + # Since onnx::expand supports two-way broadcasting, + # -1 dim value can be exported to onnx as 1 + size = symbolic_helper._reshape_helper( + g, stack(g, size, 0), g.op("Constant", value_t=torch.tensor([-1])) + ) + dtype = _type_utils.JitScalarType.INT64 + ones = ones_like(g, size, dtype) + neg_ones = mul(g, ones, g.op("Constant", value_t=torch.tensor(-1))) + size = where(g, g.op("Equal", size, neg_ones), ones, size) + return g.op("Expand", self, size) + + +@_onnx_symbolic("aten::broadcast_to") +@symbolic_helper.quantized_args(True) +def broadcast_to(g: jit_utils.GraphContext, self, size): + size = symbolic_helper._maybe_get_const(size, "is") + if not symbolic_helper._is_value(size): + size = g.op("Constant", value_t=torch.LongTensor(size)) + elif symbolic_helper._is_packed_list(size): + # Expand with -1 dim value means dim is unchanged. + # Since onnx::expand supports two-way broadcasting, + # -1 dim value can be exported to onnx as 1 + size = symbolic_helper._reshape_helper( + g, stack(g, size, 0), g.op("Constant", value_t=torch.tensor([-1])) + ) + dtype = _type_utils.JitScalarType.INT64 + ones = ones_like(g, size, dtype) + neg_ones = mul(g, ones, g.op("Constant", value_t=torch.tensor(-1))) + size = where(g, g.op("Equal", size, neg_ones), ones, size) + return g.op("Expand", self, size) + + +@_onnx_symbolic("aten::expand_as") +@symbolic_helper.quantized_args(True, True) +def expand_as(g: jit_utils.GraphContext, self, other): + self_t = symbolic_helper._maybe_get_const(self, "t") + if isinstance(self_t, torch.Tensor): + orig_type = self_t.dtype + self_t = self_t.to(torch.double) + dims = [] + for d in range(self_t.dim()): + if torch.equal(self_t.mean(d).unsqueeze(d).expand_as(self_t), self_t): + dims.append(d) + self = g.op( + "Constant", value_t=self_t.mean(dims, keepdim=True).to(orig_type) + ) + + shape = g.op("Shape", other) + return g.op("Expand", self, shape) + + +@_onnx_symbolic("aten::embedding") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "v", "i", "b", "v") +def embedding( + g: jit_utils.GraphContext, + weight, + indices, + padding_idx, + scale_grad_by_freq, + sparse, +): + if scale_grad_by_freq and GLOBALS.export_training: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of embedding with scale_grad_by_freq=True " + "for training mode. ONNX does not support scaling the gradients.", + weight, + ) + if padding_idx >= 0 and GLOBALS.export_training: + warnings.warn( + "Warning: ONNX export of embedding with padding_idx >= 0 " + "for training mode. " + "ONNX does not support not updating the embedding vector at padding_idx during training.", + stacklevel=2, + ) + + return g.op("Gather", weight, indices) + + +@_onnx_symbolic("aten::embedding_bag") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "v", "v", "i", "i", "i", "v", "i", "i") +def embedding_bag( + g: jit_utils.GraphContext, + embedding_matrix, + indices, + offsets, + scale_grad_by_freq, + mode, + sparse, + per_sample_weights, + include_last_offset, + padding_idx, +): + if not symbolic_helper._is_none(per_sample_weights): + return symbolic_helper._onnx_unsupported( + "embedding_bag with per_sample_weights" + ) + + return symbolic_helper._onnx_unsupported("embedding_bag", embedding_matrix) + + +@_onnx_symbolic("aten::size") +@symbolic_helper.quantized_args(True, quantize_output=False) +def size(g: jit_utils.GraphContext, self, dim=None): + if dim is None: + return g.op("Shape", self) + if symbolic_helper._maybe_get_const(dim, "i") < 0: + rank = symbolic_helper._get_tensor_rank(self) + if rank is not None: + dim = symbolic_helper._maybe_get_const(dim, "i") + rank + dim = g.op("Constant", value_t=torch.tensor(dim)) + return symbolic_helper._size_helper(g, self, dim) + + +@_onnx_symbolic("aten::transpose") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "i", "i") +def transpose(g: jit_utils.GraphContext, self, dim0, dim1): + if dim0 == dim1: # micro-optimization + return self + + # NB: Transpose in ONNX is actually a Permute + rank = symbolic_helper._get_tensor_rank(self) + if rank is not None: + axes = list(range(rank)) + axes[dim0], axes[dim1] = axes[dim1], axes[dim0] + return g.op("Transpose", self, perm_i=axes) + else: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of transpose for tensor of unknown rank.", + self, + ) + + +@_onnx_symbolic("aten::permute") +@symbolic_helper.parse_args("v", "is") +def permute(g: jit_utils.GraphContext, self, dims): + if dims == list(range(len(dims))): + return self + return g.op("Transpose", self, perm_i=dims) + + +@_onnx_symbolic("aten::view") +@symbolic_helper.quantized_args(True) +def view(g: jit_utils.GraphContext, self, size): + return reshape(g, self, size) + + +@_onnx_symbolic("aten::view_as") +def view_as(g: jit_utils.GraphContext, self, other): + shape = g.op("Shape", other) + return reshape(g, self, shape) + + +@_onnx_symbolic("aten::unsafe_chunk") +@symbolic_helper.parse_args("v", "i", "i", "i") +def unsafe_chunk(g: jit_utils.GraphContext, self, chunks, dim, _outputs=None): + if _outputs is None: + return symbolic_helper._onnx_opset_unsupported_detailed( + "unsafe_chunk", 9, 11, "Dynamic number of outputs not supported", self + ) + size = symbolic_helper._get_tensor_dim_size(self, dim) + if size is None: + return symbolic_helper._unimplemented( + "unsafe_chunk", "unknown dimension size", self + ) + split_size = (size + chunks - 1) // chunks + splits = [split_size] * (size // split_size) + leftover = size % split_size + if leftover: + splits.append(leftover) + return g.op("Split", self, split_i=splits, axis_i=dim, outputs=_outputs) + + +@_onnx_symbolic("aten::split") +@symbolic_helper.parse_args("v", "v", "i", "i") +def split(g: jit_utils.GraphContext, self, split_size_or_sizes, dim, _outputs=None): + if not symbolic_helper._is_split_static(split_size_or_sizes, _outputs): + return symbolic_helper._onnx_opset_unsupported_detailed( + "split", 9, 11, "Dynamic number of outputs not supported", self + ) + split_val = symbolic_helper._node_get(split_size_or_sizes.node(), "value") + if split_val.dim() > 0: + return split_with_sizes(g, self, split_size_or_sizes, dim, _outputs) + split_size = symbolic_helper._get_const(split_size_or_sizes, "i", "split_size") + + size = symbolic_helper._get_tensor_dim_size(self, dim) + if size is None: + if _outputs is not None: + size = split_size * _outputs + else: + return symbolic_helper._onnx_opset_unsupported_detailed( + "split", 9, 11, "Unknown dimension size not supported", self + ) + splits = [split_size] * (size // split_size) + leftover = size % split_size + if leftover: + splits.append(leftover) + # pyrefly: ignore [bad-argument-type] + return g.op("Split", self, split_i=splits, axis_i=dim, outputs=_outputs) + + +@_onnx_symbolic("aten::unsafe_split") +def unsafe_split( + g: jit_utils.GraphContext, self, split_size_or_sizes, dim, _outputs=None +): + return split(g, self, split_size_or_sizes, dim, _outputs) + + +@_onnx_symbolic("aten::split_with_sizes") +@symbolic_helper.parse_args("v", "is", "i", "i") +def split_with_sizes(g: jit_utils.GraphContext, self, split_sizes, dim, _outputs=None): + if not symbolic_helper._is_split_static(split_sizes, _outputs): + return symbolic_helper._onnx_opset_unsupported_detailed( + "split_with_sizes", 9, 11, "Dynamic number of outputs not supported", self + ) + # pyrefly: ignore [bad-argument-type] + return g.op("Split", self, split_i=split_sizes, axis_i=dim, outputs=_outputs) + + +@_onnx_symbolic("aten::unsafe_split_with_sizes") +def unsafe_split_with_sizes( + g: jit_utils.GraphContext, self, split_sizes, dim, _outputs=None +): + return split_with_sizes(g, self, split_sizes, dim, _outputs) + + +@_onnx_symbolic("aten::unbind") +@symbolic_helper.parse_args("v", "i", "i") +def unbind(g: jit_utils.GraphContext, self, dim=0, _outputs=None): + if _outputs is None: + return symbolic_helper._onnx_opset_unsupported_detailed( + "unbind", 9, 11, "Dynamic number of outputs not supported", self + ) + + outputs = g.op("Split", self, split_i=[1] * _outputs, axis_i=dim, outputs=_outputs) + outputs = [outputs] if _outputs == 1 else outputs + squeezed_outputs = [ + symbolic_helper._squeeze_helper(g, out, [dim]) for out in outputs + ] + return squeezed_outputs + + +@_onnx_symbolic("aten::select") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "i", "v") +def select(g: jit_utils.GraphContext, self, dim, index): + """Implement the select functionality for a pytorch tensor in ONNX. + + Selects elements from the input tensor along the specified `dim` dimension based on the `index` tensor. + """ + index = symbolic_helper._maybe_get_scalar(index) + if (not symbolic_helper._is_value(index)) and (index < 0): + if index == -1: + end_index = _constants.INT64_MAX + else: + end_index = index + 1 + slice_node = symbolic_helper._slice_helper( + g, self, axes=[dim], starts=[index], ends=[end_index] + ) + return symbolic_helper._squeeze_helper(g, slice_node, [dim]) + else: + # FIXME(justinchuby): can index be an int and not a value? + return g.op("Gather", self, index, axis_i=dim) + + +@_onnx_symbolic("aten::square") +def square(g: jit_utils.GraphContext, self): + return g.op("Mul", self, self) + + +@_onnx_symbolic("aten::squeeze") +def squeeze(g: jit_utils.GraphContext, self, dim=None): + if dim is None: + return g.op("Squeeze", self) + + squeeze_dim = symbolic_helper._get_const(dim, "i", "dim") + # Handle negative dims + if squeeze_dim < 0: + rank = symbolic_helper._get_tensor_rank(self) + if rank is not None: + warnings.warn( + "ONNX export squeeze with negative axis " + + str(squeeze_dim) + + " might cause the onnx model to be incorrect. " + + "Negative axis is not supported in ONNX. " + + "Axis is converted to " + + str(squeeze_dim + rank) + + " based on input shape at export time. " + + "Passing an tensor of different rank in execution will be incorrect.", + stacklevel=2, + ) + squeeze_dim += rank + else: + return symbolic_helper._unimplemented( + "squeeze", "negative axis with unknown input rank", self + ) + + dim_size = symbolic_helper._get_tensor_dim_size(self, squeeze_dim) + if dim_size is None: + warnings.warn( + "This model contains a squeeze operation on dimension " + + str(squeeze_dim) + + " on an input " + + "with unknown shape. Note that if the size of dimension " + + str(squeeze_dim) + + " of the input " + + "is not 1, the ONNX model will return an error. Opset version 11 supports squeezing on " + + "non-singleton dimensions, it is recommended to export this model using opset " + + "version 11 or higher.", + stacklevel=2, + ) + return symbolic_helper._squeeze_helper(g, self, axes_i=[squeeze_dim]) + if dim_size > 1: + warnings.warn( + "This model contains a squeeze operation on dimension " + + str(squeeze_dim) + + ". The size of " + + "this dimension in the given input is " + + str(dim_size) + + ". The model will " + + "be exported without the squeeze node. If the model is intended to be used with dynamic " + + "input shapes, please use opset version 11 to " + + "export the model.", + stacklevel=2, + ) + return self + + warnings.warn( + "This model contains a squeeze operation on dimension " + + str(squeeze_dim) + + ". If the model is " + + "intended to be used with dynamic input shapes, please use opset version 11 to export the model.", + stacklevel=2, + ) + return symbolic_helper._squeeze_helper(g, self, axes_i=[squeeze_dim]) + + +@_onnx_symbolic("aten::prelu") +def prelu(g: jit_utils.GraphContext, self, weight): + self_rank = symbolic_helper._get_tensor_rank(self) + weight_sizes = symbolic_helper._get_tensor_sizes(weight) + weight_rank = len(weight_sizes) + if self_rank is not None: + if self_rank > 2: + # make weight unidirectional broadcastable + weight = symbolic_helper._unsqueeze_helper( + g, weight, list(range(1, self_rank - 1)) + ) + elif self_rank == 0 and weight_sizes == [1]: + # self and weight are both scalar but weight has rank == 1, squeeze weight. + weight = symbolic_helper._squeeze_helper(g, weight, [0]) + weight_rank = 0 + + if self_rank is not None and weight_rank is not None: + if self_rank < weight_rank: + raise AssertionError( + f"rank(x) should be >= rank(slope) but got {self_rank} < {weight_rank}" + ) + return g.op("PRelu", self, weight) + + +@_onnx_symbolic("aten::silu") +def silu(g: jit_utils.GraphContext, input): + return g.op("Mul", input, g.op("Sigmoid", input)) + + +@_onnx_symbolic("aten::mish") +def mish(g: jit_utils.GraphContext, input): + return g.op("Mul", input, g.op("Tanh", g.op("Softplus", input))) + + +@_onnx_symbolic("aten::relu") +@symbolic_helper.quantized_args(True) +def relu(g: jit_utils.GraphContext, input): + return symbolic_helper._op_with_optional_float_cast( + g, "Relu", input, opset_before=14 + ) + + +@_onnx_symbolic("aten::relu6") +@symbolic_helper.quantized_args(True) +def relu6(g: jit_utils.GraphContext, input): + return clamp(g, input, 0, 6) + + +@_onnx_symbolic("aten::ceil") +def ceil(g: jit_utils.GraphContext, input): + return g.op("Ceil", input) + + +@_onnx_symbolic("aten::floor") +def floor(g: jit_utils.GraphContext, input): + return g.op("Floor", input) + + +@_onnx_symbolic("aten::len") +def _len(g: jit_utils.GraphContext, self): + sz_0 = size(g, self, g.op("Constant", value_t=torch.LongTensor([0]))) + return symbolic_helper._squeeze_helper(g, sz_0, [0]) + + +@_onnx_symbolic("aten::threshold") +@symbolic_helper.parse_args("v", "t", "t") +def threshold(g: jit_utils.GraphContext, self, threshold, value): + # See Note [Export inplace] + if symbolic_helper._scalar(threshold) != 0: + return symbolic_helper._unimplemented("threshold", "non-zero threshold", self) + if symbolic_helper._scalar(value) != 0: + return symbolic_helper._unimplemented("threshold", "non-zero value", self) + return g.op("Relu", self) + + +@_onnx_symbolic("aten::leaky_relu") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "f", "b") +def leaky_relu( + g: jit_utils.GraphContext, + input: _C.Value, + negative_slope: float, + inplace: bool = False, +): + # See Note [Export inplace] + return g.op("LeakyRelu", input, alpha_f=negative_slope) + + +@_onnx_symbolic("aten::glu") +@symbolic_helper.parse_args("v", "i") +def glu(g: jit_utils.GraphContext, input, dim): + dim_size = symbolic_helper._get_tensor_dim_size(input, dim) + if dim_size is not None: + if dim_size % 2 != 0: + raise AssertionError(f"dim_size must be even, got {dim_size}") + + first, second = g.op("Split", input, axis_i=dim, outputs=2) + return g.op("Mul", first, g.op("Sigmoid", second)) + + +@_onnx_symbolic("aten::softmax") +@symbolic_helper.parse_args("v", "i", "none") +def softmax(g: jit_utils.GraphContext, input, dim, dtype=None): + # Softmax does normalization at vector level. + # PyTorch and ONNX use different strategies to split the input tensor into vectors. + # Thus dim and axis have different meanings. + # PyTorch slices the input tensor into vectors along the `dim`-th dimension. + # ONNX reshapes the input into a 2-D tensor, and `axis` indicates where the input is coerced. + # If input is a 2 x 3 tensor: + # input = [[1.0, 1.0, 1.0], + # [1.0, 1,0, 1,0]] + # with dim = 0, the result is: + # result = [[0.5, 0.5, 0.5], + # [0.5, 0.5, 0.5]] + # with axis = 0, the result is: + # result = [[0.167, 0.167, 0.167], + # [0.167, 0.167, 0.167]] + # So only when dim and axis both equal to ndim - 1 (the last dimension), + # their semantics are equivalent. + # So use softmax when dim and axis both equal to ndim - 1, + # otherwise transpose the input to put the vectors to be normalized to the last dimension. + # When input rank is not known at export time we compute softmax using a subgraph + # with other operators + input_dim = symbolic_helper._get_tensor_rank(input) + if input_dim is not None: + # TODO: remove this as onnx opset 11 spec allows negative axes + if dim < 0: + dim = input_dim + dim + + is_transpose_required = input_dim != dim + 1 + + if is_transpose_required: + axes = list(range(input_dim)) + axes[dim], axes[-1] = axes[-1], axes[dim] + input = g.op("Transpose", input, perm_i=axes) + dim = input_dim - 1 + + softmax = g.op("Softmax", input, axis_i=dim) + if dtype and dtype.node().kind() != "prim::Constant": + parsed_dtype = symbolic_helper._get_const(dtype, "i", "dtype") + softmax = g.op( + "Cast", + softmax, + to_i=_type_utils.JitScalarType(parsed_dtype).onnx_type(), + ) + + if is_transpose_required: + softmax = g.op("Transpose", softmax, perm_i=axes) # type: ignore[possibly-undefined] + return softmax + + # Apply max normalization. + input = g.op("Sub", input, g.op("ReduceMax", input, axes_i=[dim], keepdims_i=1)) + + exp = g.op("Exp", input) + sum = symbolic_helper._reducesum_helper(g, exp, axes_i=[dim]) + softmax = g.op("Div", exp, sum) + if dtype and dtype.node().kind() != "prim::Constant": + parsed_dtype = symbolic_helper._get_const(dtype, "i", "dtype") + softmax = g.op( + "Cast", softmax, to_i=_type_utils.JitScalarType(parsed_dtype).onnx_type() + ) + return softmax + + +@_onnx_symbolic("aten::softplus") +def softplus(g: jit_utils.GraphContext, self, beta, threshold): + beta_const = symbolic_helper._maybe_get_const(beta, "f") + if beta_const != 1: + return g.op("Div", g.op("Softplus", g.op("Mul", self, beta)), beta) + return g.op("Softplus", self) + + +@_onnx_symbolic("aten::get_pool_ceil_padding") +def get_pool_ceil_padding(input, kernel_size, stride, padding): + # TODO(justinchuby): Looks like this op is deprecated in torch + sizes = symbolic_helper._get_tensor_sizes(input) + dim = sizes[-len(padding) :] if sizes is not None else None + if dim is None or any(i is None for i in dim): + return symbolic_helper._unimplemented( + "get_pool_ceil_padding", "input size not accessible", input + ) + ceiled_output_dim = [ + math.ceil((dim[i] + 2 * padding[i] - kernel_size[i]) / float(stride[i])) + 1 + for i in range(len(padding)) + ] + # ensure last pooling starts inside + ceiled_output_dim = [ + ( + ceiled_output_dim[i] - 1 + if (((ceiled_output_dim[i] - 1) * stride[i]) >= (dim[i] + padding[i])) + else ceiled_output_dim[i] + ) + for i in range(len(ceiled_output_dim)) + ] + padding_ceil = [ + ( + 0 + if (stride[i] == 1) + else ( + kernel_size[i] + - ( + dim[i] + + 2 * padding[i] + - ((ceiled_output_dim[i] - 1) * stride[i] + 1) + ) + ) + ) + for i in range(len(padding)) + ] + # ensure padding is not > kernel_size + padding_ceil = [ + ( + ( + int(padding_ceil[i]) + if padding_ceil[i] < kernel_size[i] - 1 + else int(kernel_size[i] - 1) + ) + if ((padding_ceil[i] + 2 * padding[i]) >= (kernel_size[i])) + else int(padding_ceil[i]) + ) + for i in range(len(padding_ceil)) + ] + return padding_ceil + + +@_onnx_symbolic( + "aten::max_pool1d", + decorate=[ + symbolic_helper._apply_params( + "max_pool1d", torch.nn.modules.utils._single, 1, return_indices=False + ), + _export("max_pool1d"), + ], +) +@_onnx_symbolic( + "aten::max_pool2d", + decorate=[ + symbolic_helper._apply_params( + "max_pool2d", torch.nn.modules.utils._pair, 2, return_indices=False + ), + _export("max_pool2d"), + ], +) +@_onnx_symbolic( + "aten::max_pool3d", + decorate=[ + symbolic_helper._apply_params( + "max_pool3d", torch.nn.modules.utils._triple, 3, return_indices=False + ), + _export("max_pool3d"), + ], +) +def _max_pool(name, tuple_fn, ndims, return_indices): + @symbolic_helper.quantized_args(True, False, False, False, False, False) + @symbolic_helper.parse_args("v", "is", "is", "is", "is", "i") + def symbolic_fn(g, input, kernel_size, stride, padding, dilation, ceil_mode): + if set(tuple_fn(dilation)) != {1}: + return symbolic_helper._unimplemented(name, "dilation", input) + if not stride: + stride = kernel_size + padding = tuple(tuple_fn(padding)) + if ceil_mode: + padding_ceil = get_pool_ceil_padding(input, kernel_size, stride, padding) + padding = padding + tuple(a + b for (a, b) in zip(padding_ceil, padding)) + else: + padding = padding * 2 + kwargs = { + "kernel_shape_i": tuple_fn(kernel_size), + "pads_i": padding, + "strides_i": tuple_fn(stride), + } + # easy but hacky way to get flattened indices values + # to be used to convert the indices values to non-flattened. + # In ONNX the indices are computed as a flatten 1-D tensor, + # so the values in indices are in [0, N x C x D1 x ... x Dn). + # To convert the indices to the same format used by Pytorch, + # we first execute a maxpool with a kernel and stride of 1 on the same input. + # This will result in a tensor of indices in which each index will have it's own value. + # Using this tensor as a reference, we extract the first index of each axis and subtract + # it from each index of this axis in the indices to convert. + # This step will result in a tensor were each dimension has values of indices within + # the dimension it is in. + # For more information : + # https://github.com/pytorch/pytorch/pull/16455#issuecomment-460776407 + if return_indices: + r, indices = g.op("MaxPool", input, outputs=2, **kwargs) + _, flattened_indices = g.op( + "MaxPool", + input, + outputs=2, + kernel_shape_i=[1 for _ in range(ndims)], + strides_i=[1 for _ in range(ndims)], + ) + # convert indices to have non-flattened indices values + s = symbolic_helper._slice_helper( + g, + flattened_indices, + axes=[2 + i for i in range(ndims)], + starts=list(tuple_fn(0)), + ends=list(tuple_fn(1)), + ) + indices = sub(g, indices, s) + return r, indices + else: + r = g.op("MaxPool", input, outputs=1, **kwargs) + return r + + return symbolic_fn + + +max_pool1d_with_indices = _onnx_symbolic("aten::max_pool1d_with_indices")( + _max_pool( + "max_pool1d_with_indices", + torch.nn.modules.utils._single, + 1, + return_indices=True, + ) +) +max_pool2d_with_indices = _onnx_symbolic("aten::max_pool2d_with_indices")( + _max_pool( + "max_pool2d_with_indices", + torch.nn.modules.utils._pair, + 2, + return_indices=True, + ) +) +max_pool3d_with_indices = _onnx_symbolic("aten::max_pool3d_with_indices")( + _max_pool( + "max_pool3d_with_indices", + torch.nn.modules.utils._triple, + 3, + return_indices=True, + ) +) + + +@_onnx_symbolic( + "aten::avg_pool1d", + decorate=[ + symbolic_helper._apply_params("avg_pool1d", torch.nn.modules.utils._single), + _export("avg_pool1d"), + ], +) +@_onnx_symbolic( + "aten::avg_pool2d", + decorate=[ + symbolic_helper._apply_params("avg_pool2d", torch.nn.modules.utils._pair), + _export("avg_pool2d"), + ], +) +@_onnx_symbolic( + "aten::avg_pool3d", + decorate=[ + symbolic_helper._apply_params("avg_pool3d", torch.nn.modules.utils._triple), + _export("avg_pool3d"), + ], +) +def _avg_pool(name, tuple_fn): + @symbolic_helper.quantized_args(True) + @symbolic_helper.parse_args("v", "is", "is", "is", "i", "i", "none") + def symbolic_fn( + g, + input: _C.Value, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: int | Sequence[int], + ceil_mode: int, + count_include_pad: int, + divisor_override=None, + ): + if not stride: + stride = kernel_size + padding = symbolic_helper._avgpool_helper( + tuple_fn, padding, kernel_size, stride, divisor_override, name + ) + if not isinstance(padding, tuple): + raise AssertionError(f"Expected padding to be tuple, got {type(padding)}") + adjusted_padding = padding + # Although onnx::AvgPool provides count_include_pad, + # The corner case of Average Pooling with ceil_mode on + # PyTorch allows sliding window go off bound, which leads to + # this accommodation. + # More detail on https://github.com/pytorch/pytorch/issues/57178 + if count_include_pad: + input = symbolic_helper._op_with_optional_float_cast( + g, + "Pad", + input, + pads_i=((0,) * 2 + padding) * 2, + mode_s="constant", + value_f=0.0, + opset_before=11, + ) + adjusted_padding = (0,) * len(padding) + if ceil_mode: + padding_ceil = get_pool_ceil_padding(input, kernel_size, stride, padding) + adjusted_padding = adjusted_padding + tuple( + a + b for (a, b) in zip(padding_ceil, adjusted_padding) + ) + else: + adjusted_padding = adjusted_padding * 2 + output = g.op( + "AveragePool", + input, + kernel_shape_i=tuple_fn(kernel_size), + strides_i=tuple_fn(stride), + pads_i=adjusted_padding, + ) + return output + + return symbolic_fn + + +@_onnx_symbolic( + "aten::adaptive_avg_pool1d", + decorate=[ + symbolic_helper._apply_params( + "adaptive_avg_pool1d", "AveragePool", torch.nn.modules.utils._single + ), + _export("adaptive_avg_pool1d"), + ], +) +@_onnx_symbolic( + "aten::adaptive_avg_pool2d", + decorate=[ + symbolic_helper._apply_params( + "adaptive_avg_pool2d", "AveragePool", torch.nn.modules.utils._pair + ), + _export("adaptive_avg_pool2d"), + ], +) +@_onnx_symbolic( + "aten::adaptive_avg_pool3d", + decorate=[ + symbolic_helper._apply_params( + "adaptive_avg_pool3d", "AveragePool", torch.nn.modules.utils._triple + ), + _export("adaptive_avg_pool3d"), + ], +) +@_onnx_symbolic( + "aten::adaptive_max_pool1d", + decorate=[ + symbolic_helper._apply_params( + "adaptive_max_pool1d", + "MaxPool", + torch.nn.modules.utils._single, + max_pool1d_with_indices, + ), + _export("adaptive_max_pool1d"), + ], +) +@_onnx_symbolic( + "aten::adaptive_max_pool2d", + decorate=[ + symbolic_helper._apply_params( + "adaptive_max_pool2d", + "MaxPool", + torch.nn.modules.utils._pair, + max_pool2d_with_indices, + ), + _export("adaptive_max_pool2d"), + ], +) +@_onnx_symbolic( + "aten::adaptive_max_pool3d", + decorate=[ + symbolic_helper._apply_params( + "adaptive_max_pool3d", + "MaxPool", + torch.nn.modules.utils._triple, + max_pool3d_with_indices, + ), + _export("adaptive_max_pool3d"), + ], +) +def _adaptive_pool(name, type, tuple_fn, fn=None): + @symbolic_helper.quantized_args(True, False) + def symbolic_fn(g, input, output_size): + # _adaptive_pool is supported for cases where output_size is 1 for all dimensions, + # by executing a GlobalPool. + # It is also supported for cases where the output size is a factor of the input size. + # For these cases the stride and kernel size are uniform along all the indices of + # the same dimension, which makes it possible to export it to ONNX. + # for MaxPool, GlobalMaxPool does not return indices, + # so we try using max_poolxd_with_indices, and if it is not possible + # (input is not a complete tensor or output size not factor of input size) + # then we call GlobalAveragePool and return None for the indices + output_size_value = output_size + try: + output_size = symbolic_helper._parse_arg(output_size, "is") + except Exception: + # FIXME(justinchuby): Avoid catching Exception. + # Catch a more specific exception instead. + return symbolic_helper._onnx_unsupported( + "adaptive pooling, since output_size is not constant.", input + ) + if output_size == [1] * len(output_size) and type == "AveragePool": + return g.op("GlobalAveragePool", input) + sizes = symbolic_helper._get_tensor_sizes(input) + try: + dim = sizes[2:] + except Exception: + # FIXME(justinchuby): Avoid catching Exception. + # Catch a more specific exception instead. + dim = None + if dim is None or any(i is None for i in dim): + if output_size == [1] * len(output_size): + return g.op("GlobalMaxPool", input), None + return symbolic_helper._unimplemented( + name, "input size not accessible", input + ) + # verify if output size % input size = 0 for all dim + mod = [dim[i] % output_size[i] for i in range(len(dim))] + if mod != [0] * len(mod): + if output_size == [1] * len(output_size): + return g.op("GlobalMaxPool", input), None + return symbolic_helper._unimplemented( + name, "output size that are not factor of input size", output_size_value + ) + k = [int(dim[i] / output_size[i]) for i in range(len(dim))] + # call max_poolxd_with_indices to get indices in the output + if type == "MaxPool": + # pyrefly: ignore [not-callable] + return fn(g, input, k, k, (0,) * len(dim), (1,) * len(dim), False) + output = g.op(type, input, kernel_shape_i=tuple_fn(k), strides_i=tuple_fn(k)) + return output + + return symbolic_fn + + +def _prepare_onnx_paddings(dim: int, pad): + """Generate paddings in ONNX order based on pad in pytorch. + Args: + dim: the dimension of the tensor. + pad: the paddings in pytorch. + The order is dim_n_begin, dim_n_end, dim_n-1_begin, dim_n-1_end, ... + """ + # The desired order of paddings is + # dim_0_begin, dim_1_begin, ... , dim_0_end, ..., dim_n_end. + # n is the dimension of input. + # assume zero-dimensions in the beginning + paddings = list(pad[:]) + [0] * (dim * 2 - len(pad)) + # reverse order and collate first beginnings and then ends + paddings = paddings[-2::-2] + paddings[-1::-2] + return paddings + + +def _convert_padding_node(input): + padding = symbolic_helper._maybe_get_const(input, "is") + if symbolic_helper._is_value(padding) and symbolic_helper._is_packed_list(padding): + input_list = symbolic_helper._unpack_list(padding) + try: + padding = [ + symbolic_helper._get_const(v, "i", "padding") for v in input_list + ] + except Exception: + # FIXME(justinchuby): Avoid catching Exception. + # Catch a more specific exception instead. + return symbolic_helper._onnx_opset_unsupported_detailed( + "Pad", 9, 11, "The sizes of the padding must be constant", input + ) + return padding + + +@_onnx_symbolic("aten::constant_pad_nd") +def constant_pad_nd(g: jit_utils.GraphContext, input, padding, value): + mode = "constant" + try: + value = symbolic_helper._get_const(value, "f", "value") + except Exception: + # FIXME(justinchuby): Avoid catching Exception. + # Catch a more specific exception instead. + return symbolic_helper._onnx_opset_unsupported_detailed( + "Pad", 9, 11, "The value for the padding must be constant", value + ) + + padding = _convert_padding_node(padding) + # pyrefly: ignore [bad-argument-type] + paddings = _prepare_onnx_paddings(symbolic_helper._get_tensor_rank(input), padding) + return symbolic_helper._op_with_optional_float_cast( + g, "Pad", input, pads_i=paddings, mode_s=mode, value_f=value, opset_before=11 + ) + + +def _pad_circular(g: jit_utils.GraphContext, input: _C.Value, pad: _C.Value): + padding = _convert_padding_node(pad) + if len(padding) % 2 != 0: + raise AssertionError(f"padding length must be even, got {len(padding)}") + ndim = len(padding) // 2 + + cur = input + for idx in range(ndim): + pad_r = padding[-(2 * idx + 1)] + pad_l = padding[-(2 * idx + 2)] + tensors = [] + if pad_l > 0: + left = symbolic_helper._slice_helper( + g, cur, axes=[2 + idx], starts=[-(pad_l)], ends=[_constants.INT64_MAX] + ) + tensors.append(left) + + if pad_l < 0 or pad_r < 0: + start = builtins.max(0, -pad_l) + end = -(builtins.max(0, -pad_r)) + middle = symbolic_helper._slice_helper( + g, + cur, + axes=[2 + idx], + starts=[start], + ends=[end], + ) + tensors.append(middle) + else: + tensors.append(cur) + + if pad_r > 0: + right = symbolic_helper._slice_helper( + g, cur, axes=[2 + idx], starts=[0], ends=[pad_r] + ) + tensors.append(right) + + cur = g.op("Concat", *tensors, axis_i=(2 + idx)) + + return cur + + +@_onnx_symbolic("aten::reflection_pad1d") +@_onnx_symbolic("aten::reflection_pad2d") +@_onnx_symbolic("aten::reflection_pad3d") +def reflection_pad(g: jit_utils.GraphContext, input, padding): + mode = "reflect" + padding = _convert_padding_node(padding) + # pyrefly: ignore [bad-argument-type] + paddings = _prepare_onnx_paddings(symbolic_helper._get_tensor_rank(input), padding) + return symbolic_helper._op_with_optional_float_cast( + g, "Pad", input, pads_i=paddings, mode_s=mode, opset_before=11 + ) + + +@_onnx_symbolic("aten::replication_pad1d") +@_onnx_symbolic("aten::replication_pad2d") +@_onnx_symbolic("aten::replication_pad3d") +def replication_pad(g: jit_utils.GraphContext, input, padding): + mode = "edge" + padding = _convert_padding_node(padding) + # pyrefly: ignore [bad-argument-type] + paddings = _prepare_onnx_paddings(symbolic_helper._get_tensor_rank(input), padding) + return symbolic_helper._op_with_optional_float_cast( + g, "Pad", input, pads_i=paddings, mode_s=mode, opset_before=11 + ) + + +@_onnx_symbolic("aten::pad") +def pad( + g: jit_utils.GraphContext, + input: _C.Value, + pad: _C.Value, + mode: _C.Value, + value: _C.Value, +): + mode = symbolic_helper._parse_arg(mode, "s") + if mode == "replicate": + return replication_pad(g, input, pad) + elif mode == "reflect": + return reflection_pad(g, input, pad) + elif mode == "constant": + return constant_pad_nd(g, input, pad, value) + elif mode == "circular": + return _pad_circular(g, input, pad) + else: + raise errors.SymbolicValueError(f"Unrecognized padding mode {mode}", input) + + +@_onnx_symbolic( + "aten::upsample_nearest1d", + decorate=[ + symbolic_helper._apply_params("upsample_nearest1d", 3, "nearest"), + _export("upsample_nearest1d"), + ], +) +@_onnx_symbolic( + "aten::upsample_nearest2d", + decorate=[ + symbolic_helper._apply_params("upsample_nearest2d", 4, "nearest"), + _export("upsample_nearest2d"), + ], +) +@_onnx_symbolic( + "aten::upsample_nearest3d", + decorate=[ + symbolic_helper._apply_params("upsample_nearest3d", 5, "nearest"), + _export("upsample_nearest3d"), + ], +) +@_onnx_symbolic( + "aten::upsample_linear1d", + decorate=[ + symbolic_helper._apply_params("upsample_linear1d", 3, "linear"), + _export("upsample_linear1d"), + ], +) +@_onnx_symbolic( + "aten::upsample_bilinear2d", + decorate=[ + symbolic_helper._apply_params("upsample_bilinear2d", 4, "linear"), + _export("upsample_bilinear2d"), + ], +) +@_onnx_symbolic( + "aten::upsample_trilinear3d", + decorate=[ + symbolic_helper._apply_params("upsample_trilinear3d", 5, "linear"), + _export("upsample_trilinear3d"), + ], +) +def _interpolate(name: str, dim: int, interpolate_mode: str): + def symbolic_fn(g, input, output_size, *args): + scales, align_corners = symbolic_helper._get_interpolate_attributes( + g, interpolate_mode, args + ) + symbolic_helper._interpolate_warning(interpolate_mode) + align_corners = symbolic_helper._maybe_get_scalar(align_corners) + if align_corners: + return symbolic_helper._unimplemented(name, "align_corners == True", input) + if scales is None: + scales = symbolic_helper._interpolate_size_to_scales( + g, input, output_size, dim + ) + return g.op("Upsample", input, scales, mode_s=interpolate_mode) + + return symbolic_fn + + +@_onnx_symbolic("aten::__interpolate") +def __interpolate( + g: jit_utils.GraphContext, + input, + size, + scale_factor, + mode, + align_corners, + recompute_scale_factor, + antialias, +): + scales, mode = symbolic_helper._interpolate_get_scales_and_mode( + g, input, size, scale_factor, mode, align_corners + ) + return g.op("Upsample", input, scales, mode_s=mode) + + +@_onnx_symbolic("aten::bitwise_not") +def bitwise_not(g: jit_utils.GraphContext, input): + if not symbolic_helper._is_bool(input): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise Not " + "for non-boolean input values", + input, + ) + return g.op("Not", input) + + +@_onnx_symbolic("aten::bitwise_or") +def bitwise_or(g, self, other): + if not symbolic_helper._is_bool(self): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise OR " + "for non-boolean input values. self: ", + self, + ) + if not symbolic_helper._is_bool(other): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise OR " + "for non-boolean input values. other: ", + other, + ) + return g.op("Or", self, other) + + +def wrap_logical_op_with_cast_to(to_type): + def decorator(fn): + @functools.wraps(fn) + def wrap_with_cast(g, input, other): + to_i = symbolic_helper.cast_pytorch_to_onnx[to_type] + return fn(g, g.op("Cast", input, to_i=to_i), g.op("Cast", other, to_i=to_i)) + + return wrap_with_cast + + return decorator + + +def wrap_logical_op_with_negation(func: Callable) -> Callable: + @functools.wraps(func) + def wrap_with_not(g, input, other): + return g.op("Not", func(g, input, other)) + + return wrap_with_not + + +@_onnx_symbolic("aten::__not_") +def __not_(g: jit_utils.GraphContext, self): + if not symbolic_helper._is_bool(self): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise Not " + "for non-boolean input values", + self, + ) + return g.op("Not", self) + + +@_onnx_symbolic("aten::eq") +@symbolic_helper.quantized_args(True, True) +def eq(g: jit_utils.GraphContext, self, other): + if isinstance(self.type(), _C.DeviceObjType) and isinstance( + other.type(), _C.DeviceObjType + ): + # ONNX doesn't have devices, so consider them all to be equal. + # The no-op check for equality will get constant-folded. + return g.op("Constant", value_t=torch.tensor(True, dtype=torch.bool)) + self_node = self.node() + other_node = other.node() + if self_node.kind() == other_node.kind() == "onnx::Constant": + if self_node.kindOf("value") == other_node.kindOf("value") == "s": + # Exporting strings to ONNX is not supported. + # If both strings are constant, we can compare them directly. + # The no-op check for equality will get constant-folded. + return g.op( + "Constant", + value_t=torch.tensor( + self_node.s("value") == other_node.s("value"), + dtype=torch.bool, + ), + ) + + return g.op("Equal", self, other) + + +@_onnx_symbolic("aten::ne") +@symbolic_helper.quantized_args(True, True) +@wrap_logical_op_with_negation +def ne(g: jit_utils.GraphContext, self, other): + return eq(g, self, other) + + +@_onnx_symbolic("aten::gt") +@symbolic_helper.quantized_args(True, True) +def gt(g: jit_utils.GraphContext, input, other): + return _gt_impl(g, input, other) + + +def _gt_impl(g: jit_utils.GraphContext, input, other): + if symbolic_helper._is_bool(input) and symbolic_helper._is_bool(other): + input = g.op("Cast", input, to_i=_C_onnx.TensorProtoDataType.INT32) + other = g.op("Cast", other, to_i=_C_onnx.TensorProtoDataType.INT32) + return g.op("Greater", input, other) + + +@_onnx_symbolic("aten::lt") +@symbolic_helper.quantized_args(True, True) +def lt(g: jit_utils.GraphContext, input, other): + return _lt_impl(g, input, other) + + +def _lt_impl(g: jit_utils.GraphContext, input, other): + if symbolic_helper._is_bool(input) and symbolic_helper._is_bool(other): + input = g.op("Cast", input, to_i=_C_onnx.TensorProtoDataType.INT32) + other = g.op("Cast", other, to_i=_C_onnx.TensorProtoDataType.INT32) + return g.op("Less", input, other) + + +@_onnx_symbolic("aten::ge") +@symbolic_helper.quantized_args(True, True) +@wrap_logical_op_with_negation +def ge(g: jit_utils.GraphContext, input, other): + return _lt_impl(g, input, other) + + +@_onnx_symbolic("aten::le") +@symbolic_helper.quantized_args(True, True) +@wrap_logical_op_with_negation +def le(g: jit_utils.GraphContext, input, other): + return _gt_impl(g, input, other) + + +@_onnx_symbolic("aten::__and_") +def __and_(g: jit_utils.GraphContext, input, other): + if not symbolic_helper._is_bool(input): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise AND " + "for non-boolean input values", + input, + ) + if not symbolic_helper._is_bool(other): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise AND " + "for non-boolean input values", + other, + ) + return g.op("And", input, other) + + +@_onnx_symbolic("aten::__or_") +def __or_(g: jit_utils.GraphContext, input, other): + if not symbolic_helper._is_bool(input): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise OR " + "for non-boolean input values", + input, + ) + if not symbolic_helper._is_bool(other): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise OR " + "for non-boolean input values", + other, + ) + return g.op("Or", input, other) + + +@_onnx_symbolic("aten::__xor_") +def __xor_(g: jit_utils.GraphContext, input, other): + if not symbolic_helper._is_bool(input): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise XOR " + "for non-boolean input values", + input, + ) + if not symbolic_helper._is_bool(other): + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting bitwise XOR " + "for non-boolean input values", + other, + ) + return g.op("Xor", input, other) + + +@_onnx_symbolic("aten::logical_and") +@wrap_logical_op_with_cast_to("Bool") +def logical_and(g: jit_utils.GraphContext, input, other): + return g.op("And", input, other) + + +@_onnx_symbolic("aten::logical_or") +@wrap_logical_op_with_cast_to("Bool") +def logical_or(g: jit_utils.GraphContext, input, other): + return g.op("Or", input, other) + + +@_onnx_symbolic("aten::logical_xor") +@wrap_logical_op_with_cast_to("Bool") +def logical_xor(g: jit_utils.GraphContext, input, other): + return g.op("Xor", input, other) + + +@_onnx_symbolic("aten::logical_not") +def logical_not(g: jit_utils.GraphContext, input): + return g.op("Not", g.op("Cast", input, to_i=_C_onnx.TensorProtoDataType.BOOL)) + + +@_onnx_symbolic("aten::__rshift_") +def __rshift_(g: jit_utils.GraphContext, self, other): + # make sure to cast other to self's type + # (when self is long, make sure that other is not float) + self_scalar_type = _type_utils.JitScalarType.from_value(self) + if ( + _type_utils.JitScalarType.from_value(other, _type_utils.JitScalarType.UNDEFINED) + != self_scalar_type + ): + other = g.op( + "Cast", + other, + to_i=self_scalar_type.onnx_type(), + ) + + two = g.op("Constant", value_t=torch.tensor(2, dtype=torch.float32)) + # exponent (same type as self) has to be float or double in onnx::Pow + if not symbolic_helper._is_fp(self): + other = g.op("Cast", other, to_i=_C_onnx.TensorProtoDataType.FLOAT) + two_pow = g.op("Pow", two, other) + two_pow = g.op( + "Cast", + two_pow, + to_i=self_scalar_type.onnx_type(), + ) + rshift = g.op("Div", self, two_pow) + return rshift + + +@_onnx_symbolic("aten::__lshift_") +def __lshift_(g: jit_utils.GraphContext, self, other): + # make sure to cast other to self's type + # (when self is long, make sure that other is not float) + self_scalar_type = _type_utils.JitScalarType.from_value(self) + if ( + _type_utils.JitScalarType.from_value(other, _type_utils.JitScalarType.UNDEFINED) + != self_scalar_type + ): + other = g.op( + "Cast", + other, + to_i=self_scalar_type.onnx_type(), + ) + + two = g.op("Constant", value_t=torch.tensor(2, dtype=torch.float32)) + # exponent (same type as self) has to be float or double in onnx::Pow + if not symbolic_helper._is_fp(self): + other = g.op("Cast", other, to_i=_C_onnx.TensorProtoDataType.FLOAT) + two_pow = g.op("Pow", two, other) + two_pow = g.op( + "Cast", + two_pow, + to_i=self_scalar_type.onnx_type(), + ) + lshift = g.op("Mul", self, two_pow) + return lshift + + +@_onnx_symbolic("aten::where") +@symbolic_helper.parse_args("v", "v", "v", "i") +def where(g: jit_utils.GraphContext, condition, self=None, other=None, _outputs=None): + # Assumes that torch.where's first argument takes only Bool and Byte tensors. + if not symbolic_helper._is_bool(condition): + condition = g.op("Cast", condition, to_i=_C_onnx.TensorProtoDataType.BOOL) + if self is None: + condition = nonzero(g, condition) + return symbolic_helper._unbind_helper( + g, condition, g.op("Constant", value_t=torch.tensor(1)), _outputs + ) + # pyrefly: ignore [bad-argument-type] + return g.op("Where", condition, self, other) + + +@_onnx_symbolic("aten::log_softmax") +@symbolic_helper.parse_args("v", "i", "none") +def log_softmax(g: jit_utils.GraphContext, input, dim, dtype=None): + # PyTorch dim and ONNX axis have different meanings. + # See Softmax comment for details. + # TODO: remove this as onnx opset 11 spec allows negative axes + input_dim = symbolic_helper._get_tensor_rank(input) + if input_dim is None: + return symbolic_helper._unimplemented( + "dim", + "ONNX and PyTorch use different strategies to split the input. " + "Input rank must be known at export time.", + ) + if dim < 0: + dim = input_dim + dim + is_transpose_required = input_dim != dim + 1 + # ONNX only supports log_softmax with dim = -1. Transpose must be added before and after log_softmax to support other cases. + if is_transpose_required: + axes = list(range(input_dim)) + axes[dim], axes[-1] = axes[-1], axes[dim] + input = g.op("Transpose", input, perm_i=axes) + dim = input_dim - 1 + return_op = g.op("LogSoftmax", input, axis_i=dim) + if dtype and dtype.node().kind() != "prim::Constant": + parsed_dtype = symbolic_helper._get_const(dtype, "i", "dtype") + return_op = g.op( + "Cast", return_op, to_i=_type_utils.JitScalarType(parsed_dtype).onnx_type() + ) + if is_transpose_required: + return_op = g.op("Transpose", return_op, perm_i=axes) # type: ignore[possibly-undefined] + return return_op + + +@_onnx_symbolic("aten::_log_softmax") +@symbolic_helper.parse_args("v", "i", "i") +def _log_softmax(g: jit_utils.GraphContext, input, dim, half_to_float): + if ( + half_to_float + and _type_utils.JitScalarType.from_value( + input, _type_utils.JitScalarType.UNDEFINED + ) + == _type_utils.JitScalarType.HALF + ): + input = g.op("Cast", input, to_i=_C_onnx.TensorProtoDataType.FLOAT) + return log_softmax(g, input, dim) + + +@_onnx_symbolic("aten::_convolution") +@symbolic_helper.parse_args( + "v", "v", "v", "is", "is", "is", "i", "is", "i", "i", "i", "i", "i" +) +def _convolution( + g: jit_utils.GraphContext, + input, + weight, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + benchmark, + deterministic, + cudnn_enabled, + allow_tf32=None, +): + weight_size = symbolic_helper._get_tensor_sizes(weight) + try: + kernel_shape = weight_size[2:] + except Exception: + # FIXME(justinchuby): Avoid catching Exception. + # Catch a more specific exception instead. + kernel_shape = None + + if kernel_shape is None or any(i is None for i in kernel_shape): + raise errors.SymbolicValueError( + "Unsupported: ONNX export of convolution for kernel of unknown shape.", + input, + ) + + args = [input, weight] + # ONNX only supports 1D bias + if ( + not symbolic_helper._is_none(bias) + and symbolic_helper._get_tensor_rank(bias) == 1 + ): + args.append(bias) + + kwargs = { + "kernel_shape_i": weight_size[2:], + "strides_i": stride, + # NB: ONNX supports asymmetric padding, whereas PyTorch supports only + # symmetric padding + "pads_i": padding + padding, + "dilations_i": dilation, + "group_i": groups, + } + + if any(o != 0 for o in output_padding): + # ONNX supports both output_shape and output_padding. they are equivalent expressive. + # output_padding is more straightforward, so we use it here. + # output_shape = stride * (input_shape - 1) + output_padding + kernel_shape - padding * 2 + if not transposed: + raise AssertionError("output_padding requires transposed=True") + if len(stride) != len(output_padding): + raise AssertionError( + f"len(stride)={len(stride)} != len(output_padding)={len(output_padding)}" + ) + kwargs["output_padding_i"] = output_padding + + n = g.op("ConvTranspose" if transposed else "Conv", *args, **kwargs) + + if ( + not symbolic_helper._is_none(bias) + and symbolic_helper._get_tensor_rank(bias) != 1 + ): + return g.op("Add", n, bias) + else: + return n + + +@_onnx_symbolic("aten::_convolution_mode") +@symbolic_helper.parse_args( + "v", + "v", + "v", + "is", + "s", + "is", + "i", +) +def _convolution_mode( + g: jit_utils.GraphContext, + input, + weight, + bias, + stride, + padding, + dilation, + groups, +): + weight_size = symbolic_helper._get_tensor_sizes(weight) + try: + kernel_shape = weight_size[2:] + except Exception: + # FIXME(justinchuby): Avoid catching Exception. + # Catch a more specific exception instead. + kernel_shape = None + + if kernel_shape is None or any(i is None for i in kernel_shape): + raise errors.SymbolicValueError( + "Unsupported: ONNX export of convolution for kernel of unknown shape.", + input, + ) + + args = [input, weight] + # ONNX only supports 1D bias + if ( + not symbolic_helper._is_none(bias) + and symbolic_helper._get_tensor_rank(bias) == 1 + ): + args.append(bias) + + if padding == "valid": + padding = "VALID" + elif padding == "same": + padding = "SAME_UPPER" + kwargs = { + "kernel_shape_i": weight_size[2:], + "strides_i": stride, + "auto_pad_s": padding, + "dilations_i": dilation, + "group_i": groups, + } + + n = g.op("Conv", *args, **kwargs) + + if ( + not symbolic_helper._is_none(bias) + and symbolic_helper._get_tensor_rank(bias) != 1 + ): + return g.op("Add", n, bias) + else: + return n + + +@_onnx_symbolic("aten::convolution") +@symbolic_helper.parse_args("v", "v", "v", "is", "is", "is", "i", "is", "i") +def convolution( + g: jit_utils.GraphContext, + input, + weight, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, +): + return _convolution( + g, + input, + weight, + bias, + stride, + padding, + dilation, + transposed, + output_padding, + groups, + None, + None, + None, + None, + ) + + +@_onnx_symbolic("aten::conv1d") +@symbolic_helper.parse_args("v", "v", "v", "is", "v", "is", "i") +def conv1d( + g: jit_utils.GraphContext, input, weight, bias, stride, padding, dilation, groups +): + str_padding = symbolic_helper._parse_arg(padding, "s") + if str_padding in ["valid", "same"]: + return _convolution_mode( + g, + input, + weight, + bias, + stride, + str_padding, + dilation, + groups, + ) + else: + padding = symbolic_helper._parse_arg(padding, "is") + return _convolution( + g, + input, + weight, + bias, + stride, + padding, + dilation, + False, + (), + groups, + None, + None, + None, + None, + ) + + +@_onnx_symbolic("aten::conv2d") +@symbolic_helper.parse_args("v", "v", "v", "is", "v", "is", "i") +def conv2d( + g: jit_utils.GraphContext, input, weight, bias, stride, padding, dilation, groups +): + str_padding = symbolic_helper._parse_arg(padding, "s") + if str_padding in ["valid", "same"]: + return _convolution_mode( + g, + input, + weight, + bias, + stride, + str_padding, + dilation, + groups, + ) + else: + padding = symbolic_helper._parse_arg(padding, "is") + return _convolution( + g, + input, + weight, + bias, + stride, + padding, + dilation, + False, + (), + groups, + None, + None, + None, + None, + ) + + +@_onnx_symbolic("aten::conv3d") +@symbolic_helper.parse_args("v", "v", "v", "is", "v", "is", "i") +def conv3d( + g: jit_utils.GraphContext, input, weight, bias, stride, padding, dilation, groups +): + str_padding = symbolic_helper._parse_arg(padding, "s") + if str_padding in ["valid", "same"]: + return _convolution_mode( + g, + input, + weight, + bias, + stride, + str_padding, + dilation, + groups, + ) + else: + padding = symbolic_helper._parse_arg(padding, "is") + return _convolution( + g, + input, + weight, + bias, + stride, + padding, + dilation, + False, + (), + groups, + None, + None, + None, + None, + ) + + +@_onnx_symbolic("aten::conv_transpose1d") +@symbolic_helper.parse_args("v", "v", "v", "is", "is", "is", "i", "is") +def conv_transpose1d( + g: jit_utils.GraphContext, + input, + weight, + bias, + stride, + padding, + output_padding, + groups, + dilation, +): + return _convolution( + g, + input, + weight, + bias, + stride, + padding, + dilation, + True, + output_padding, + groups, + None, + None, + None, + None, + ) + + +@_onnx_symbolic("aten::conv_transpose2d") +@symbolic_helper.parse_args("v", "v", "v", "is", "is", "is", "i", "is") +def conv_transpose2d( + g: jit_utils.GraphContext, + input, + weight, + bias, + stride, + padding, + output_padding, + groups, + dilation, +): + return _convolution( + g, + input, + weight, + bias, + stride, + padding, + dilation, + True, + output_padding, + groups, + None, + None, + None, + None, + ) + + +@_onnx_symbolic("aten::conv_transpose3d") +@symbolic_helper.parse_args("v", "v", "v", "is", "is", "is", "i", "is") +def conv_transpose3d( + g: jit_utils.GraphContext, + input, + weight, + bias, + stride, + padding, + output_padding, + groups, + dilation, +): + return _convolution( + g, + input, + weight, + bias, + stride, + padding, + dilation, + True, + output_padding, + groups, + None, + None, + None, + None, + ) + + +@_onnx_symbolic("aten::batch_norm") +@symbolic_helper.parse_args("v", "v", "v", "v", "v", "i", "f", "f", "i") +def batch_norm( + g: jit_utils.GraphContext, + input, + weight, + bias, + running_mean, + running_var, + training, + momentum, + eps, + cudnn_enabled, +): + symbolic_helper.check_training_mode(training, "batch_norm") + + if ( + torch.is_autocast_enabled() + and not symbolic_helper.args_have_same_dtype( + [input, weight, bias, running_mean, running_var] + ) + and GLOBALS.export_onnx_opset_version < 15 + ): + return symbolic_helper._onnx_opset_unsupported_detailed( + "BatchNormalization", + 9, + 15, + "All input tensors must have the same `dtype`." + " Turn off Autocast or export using opset version 15.", + input, + ) + + weight, bias, running_mean, running_var = symbolic_helper._batchnorm_helper( + g, input, weight, bias, running_mean, running_var + ) + out = g.op( + "BatchNormalization", + input, + weight, + bias, + running_mean, + running_var, + epsilon_f=eps, + momentum_f=1 - momentum, + outputs=1 if not training else 5, + ) + if not training: + return out + else: + res, new_running_mean, new_running_var, saved_mean, saved_var = out + new_running_mean.setType(running_mean.type()) + new_running_var.setType(running_var.type()) + saved_mean.setDebugName("batch_norm_dead_output-" + saved_mean.debugName()) + saved_var.setDebugName("batch_norm_dead_output-" + saved_var.debugName()) + return res + + +@_onnx_symbolic("aten::native_layer_norm") +@symbolic_helper.quantized_args(True, False, False, False) +@symbolic_helper.parse_args("v", "is", "v", "v", "f") +def native_layer_norm( + g: jit_utils.GraphContext, + input: _C.Value, + normalized_shape: Sequence[int], + weight: _C.Value, + bias: _C.Value, + eps: float, +) -> tuple[_C.Value, _C.Value, _C.Value]: + axes = [-i for i in range(len(normalized_shape), 0, -1)] + + two_cst = symbolic_helper._generate_wrapped_number(g, 2.0) + eps_cst = symbolic_helper._generate_wrapped_number(g, eps) + + if g.opset < 18: + mean = g.op("ReduceMean", input, axes_i=axes) + else: + mean = g.op( + "ReduceMean", + input, + g.op("Constant", value_t=torch.tensor(axes, dtype=torch.long)), + ) + + numerator = sub(g, input, mean) + + # Cast it to eps dtype to avoid precision loss + is_type_half = ( + _type_utils.JitScalarType.from_value(numerator) + == _type_utils.JitScalarType.HALF + ) + if is_type_half: + eps_dtype = _type_utils.JitScalarType.from_value(eps_cst) + numerator = g.op( + "Cast", numerator, to_i=_type_utils.JitScalarType(eps_dtype).onnx_type() + ) + + # variance = e((x - e(x))^2), and (x - e(x)) is the numerator in the layer_norm formula + if g.opset < 18: + # pyrefly: ignore [no-matching-overload] + variance = g.op("ReduceMean", pow(g, numerator, two_cst), axes_i=axes) + else: + variance = g.op( + "ReduceMean", + # pyrefly: ignore [no-matching-overload] + pow(g, numerator, two_cst), + g.op("Constant", value_t=torch.tensor(axes, dtype=torch.long)), + ) + + denominator = sqrt(g, g.op("Add", variance, eps_cst)) + normalized = g.op("Div", numerator, denominator) + + # Cast back to input type as eps related ops are all done + if is_type_half: + input_dtype = _type_utils.JitScalarType.from_value(input) + normalized = g.op( + "Cast", normalized, to_i=_type_utils.JitScalarType(input_dtype).onnx_type() + ) + + if not (weight is None or symbolic_helper._is_none(weight)): + normalized = mul(g, normalized, weight) + if not (bias is None or symbolic_helper._is_none(bias)): + normalized = add(g, normalized, bias) + + # rdenominator := 1 / sqrt(variance + eps) + # According to aten::native_layer_norm, rdenominator should have the same dtype as input, + # mean and normalized, so we need to Cast it back + if is_type_half: + denominator = g.op( + "Cast", + denominator, + to_i=_type_utils.JitScalarType(input_dtype).onnx_type(), # type: ignore[possibly-undefined] + ) + rdenominator = g.op("Reciprocal", denominator) + else: + rdenominator = reciprocal(g, denominator) + + return normalized, mean, rdenominator + + +@_onnx_symbolic("aten::layer_norm") +@symbolic_helper.quantized_args(True, False, False, False) +@symbolic_helper.parse_args("v", "is", "v", "v", "f", "b") +def layer_norm( + g: jit_utils.GraphContext, + input: _C.Value, + normalized_shape: Sequence[int], + weight: _C.Value, + bias: _C.Value, + eps: float, + cudnn_enable: bool, +) -> _C.Value: + normalized, _, _ = native_layer_norm(g, input, normalized_shape, weight, bias, eps) + return normalized + + +@_onnx_symbolic("aten::instance_norm") +@symbolic_helper.parse_args("v", "v", "v", "v", "v", "b", "f", "f", "b") +def instance_norm( + g: jit_utils.GraphContext, + input, + weight, + bias, + running_mean, + running_var, + use_input_stats: bool, + momentum: Number, + eps: Number, + cudnn_enabled: bool, +): + symbolic_helper.check_training_mode(use_input_stats, "instance_norm") + channel_size = symbolic_helper._get_tensor_dim_size(input, 1) + if weight is None or symbolic_helper._is_none(weight): + if channel_size is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of instance_norm for unknown channel size.", + input, + ) + weight_value = torch.tensor( + [1.0] * channel_size, + dtype=_type_utils.JitScalarType.from_value(input).dtype(), + ) + weight = g.op("Constant", value_t=weight_value) + if bias is None or symbolic_helper._is_none(bias): + if channel_size is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of instance_norm for unknown channel size.", + input, + ) + bias_value = torch.tensor( + [0.0] * channel_size, + dtype=_type_utils.JitScalarType.from_value(input).dtype(), + ) + bias = g.op("Constant", value_t=bias_value) + if ( + running_mean is None + or symbolic_helper._is_none(running_mean) + or running_var is None + or symbolic_helper._is_none(running_var) + ): + return g.op("InstanceNormalization", input, weight, bias, epsilon_f=eps) + else: + input_size = symbolic_helper._get_tensor_sizes(input) + # If input shape is [N, C, H, W], reshape to [1, N * C, H, W] and call batch_norm. + # For more information instance_norm(): + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/Normalization.cpp#L542 + input_size_reshape = input_size.copy() + n = input_size[0] + if n is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of instance_norm training for unknown " + "batch size.", + input, + ) + c = input_size[1] + input_size_reshape[0] = 1 + input_size_reshape[1] = n * c + weight_ = repeat( + g, weight, g.op("Constant", value_t=torch.tensor([n], dtype=torch.int64)) + ) + bias_ = repeat( + g, bias, g.op("Constant", value_t=torch.tensor([n], dtype=torch.int64)) + ) + running_mean_ = repeat( + g, + running_mean, + g.op("Constant", value_t=torch.tensor([n], dtype=torch.int64)), + ) + running_var_ = repeat( + g, + running_var, + g.op("Constant", value_t=torch.tensor([n], dtype=torch.int64)), + ) + input_reshaped = g.op( + "Reshape", + input, + g.op("Constant", value_t=torch.LongTensor(input_size_reshape)), + ) + out = batch_norm( + g, + input_reshaped, + weight_, + bias_, + running_mean_, + running_var_, + use_input_stats, + momentum, + eps, + cudnn_enabled, + ) + return view(g, out, g.op("Constant", value_t=torch.tensor(input_size))) + + +@_onnx_symbolic("aten::unfold") +@symbolic_helper.parse_args("v", "i", "i", "i") +def unfold(g: jit_utils.GraphContext, input, dimension, size, step): + sizes = symbolic_helper._get_tensor_sizes(input) + # FIXME(justinchuby): Get rid of the try catch here to improve readability + try: + sizedim = sizes[dimension] + except Exception: + # FIXME(justinchuby): Avoid catching Exception. + # Catch a more specific exception instead. + sizedim = None + if sizedim is not None: + low_indices = range(0, sizedim, step) + hi_indices = range(size, sizedim + 1, step) + stack = [ + symbolic_helper._slice_helper( + g, input, axes=[dimension], starts=[low], ends=[hi] + ) + for low, hi in zip(low_indices, hi_indices) + ] + ndim = len(sizes) + perm = list(range(ndim)) + perm.append(perm.pop(dimension)) + unsqueeze = [ + symbolic_helper._unsqueeze_helper( + g, g.op("Transpose", t, perm_i=perm), [dimension] + ) + for t in stack + ] + return g.op("Concat", *unsqueeze, axis_i=dimension) + else: + return symbolic_helper._unimplemented( + "Unfold", "input size not accessible", input + ) + + +@_onnx_symbolic("aten::elu") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "t", "t", "t") +def elu(g: jit_utils.GraphContext, input, alpha, scale, input_scale): + if scale and scale != 1.0: + return symbolic_helper._unimplemented( + "scale", "does not support scale in Elu", scale + ) + if input_scale and input_scale != 1.0: + return symbolic_helper._unimplemented( + "input_scale", "does not support input_scale in Elu", input_scale + ) + # See Note [Export inplace] + return g.op("Elu", input, alpha_f=symbolic_helper._scalar(alpha)) + + +@_onnx_symbolic("aten::selu") +@symbolic_helper.quantized_args(True) +def selu(g: jit_utils.GraphContext, input): + return g.op("Selu", input) + + +@_onnx_symbolic("aten::index_select") +@symbolic_helper.parse_args("v", "i", "v") +def index_select(g: jit_utils.GraphContext, self, dim, index): + # In case of a scalar index, index_select returns a tensor with the same rank as the input. + # To match this behavior in ONNX, we make index a 1D tensor so that the following gather + # also produces a tensor with the same rank as the input. + return symbolic_helper._select_helper(g, self, dim, index) + + +@_onnx_symbolic("aten::index_put") +def index_put(g: jit_utils.GraphContext, self, indices_list_value, values, accumulate): + if symbolic_helper._is_packed_list(indices_list_value): + indices_list = symbolic_helper._unpack_list(indices_list_value) + else: + indices_list = [indices_list_value] + + accumulate = symbolic_helper._parse_arg(accumulate, "b") + + if len(indices_list) == 0: + if accumulate: + return add(g, self, values) + return values + symbolic_helper._onnx_opset_unsupported("index_put", 9, 11, self) + + +@_onnx_symbolic("aten::index_fill") +def index_fill(g: jit_utils.GraphContext, self, dim, index, value): + expanded_index_shape, expanded_index = symbolic_helper._index_fill_reshape_helper( + g, self, dim, index + ) + value = symbolic_helper._maybe_get_scalar(value) + value = symbolic_helper._if_scalar_type_as(value, self) + expanded_value = expand(g, value, expanded_index_shape, None) + + return scatter(g, self, dim, expanded_index, expanded_value) + + +@_onnx_symbolic("aten::index_copy") +def index_copy(g: jit_utils.GraphContext, self, dim, index, source): + _expanded_index_shape, expanded_index = symbolic_helper._index_fill_reshape_helper( + g, self, dim, index + ) + return scatter(g, self, dim, expanded_index, source) + + +@_onnx_symbolic("aten::bucketize") +@symbolic_helper.parse_args("v", "v", "b", "b") +def bucketize( + g: jit_utils.GraphContext, self, boundaries, out_int32=False, right=False +): + out_type = _C_onnx.TensorProtoDataType.INT64 + if out_int32: + out_type = _C_onnx.TensorProtoDataType.INT32 + # A tensor expanded_boundaries is created such that it + # contains a copy of boundaries for each element of self. + new_shape = g.op("Concat", g.op("Shape", boundaries), g.op("Shape", self), axis_i=0) + # Unsqueeze step is performed to respect ONNX's numpy style broadcasting for comparison ops + # https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md + tensor_rank = symbolic_helper._get_tensor_rank(self) + if tensor_rank is None: + raise AssertionError("tensor_rank must be non-None") + unsqueeze_axes = list(range(1, tensor_rank + 1)) + expanded_boundaries = expand( + g, + symbolic_helper._unsqueeze_helper(g, boundaries, unsqueeze_axes), + new_shape, + None, + ) + # Compare each element of self to boundaries to get a tensor + # with leading 1s and trailing 0s. + # e.g., 4 > [1, 3, 4] = [1, 1, 0] + # The index of the last 1 is the bucket where the element should go. + if right: + cond = ge(g, self, expanded_boundaries) + else: + cond = gt(g, self, expanded_boundaries) + cond_out = g.op("Cast", cond, to_i=out_type) + # Sum to get the number of 1s corresponding to each element, + # which is the same as the bucket index. + # e.g., sum(4 > [1, 3, 4]) = sum([1, 1, 0]) = 2 + return symbolic_helper._reducesum_helper(g, cond_out, axes_i=[0], keepdims_i=0) + + +@_onnx_symbolic("aten::type_as") +def type_as(g: jit_utils.GraphContext, self, other): + self_dtype = symbolic_helper._try_get_scalar_type(self) + other_dtype = symbolic_helper._try_get_scalar_type(other) + if self_dtype == other_dtype and self_dtype is not None: + return self + if other_dtype is not None: + return g.op( + "Cast", + self, + to_i=other_dtype.onnx_type(), + ) + + raise errors.SymbolicValueError( + "Unsupported: ONNX export of type_as for tensor " + "of unknown dtype. Please check if the dtype of the " + "parameter passed to the type_as function is correct.", + other, + ) + + +@_onnx_symbolic("aten::cosine_similarity") +@symbolic_helper.parse_args("v", "v", "i", "f") +def cosine_similarity(g: jit_utils.GraphContext, x1, x2, dim, eps): + cross = symbolic_helper._reducesum_helper( + g, mul(g, x1, x2), axes_i=[dim], keepdims_i=0 + ) + x1_l2 = symbolic_helper._reducesum_helper( + g, mul(g, x1, x1), axes_i=[dim], keepdims_i=0 + ) + x2_l2 = symbolic_helper._reducesum_helper( + g, mul(g, x2, x2), axes_i=[dim], keepdims_i=0 + ) + # pyrefly: ignore [no-matching-overload] + div_tens = max( + g, sqrt(g, mul(g, x1_l2, x2_l2)), g.op("Constant", value_t=torch.tensor([eps])) + ) + return div(g, cross, div_tens) + + +@_onnx_symbolic("aten::pairwise_distance") +def pairwise_distance(g: jit_utils.GraphContext, input1, input2, p, eps, keepdim): + if not symbolic_helper._is_value(eps): + eps = g.op("Constant", value_t=torch.tensor([eps])) + inv_p = div( + g, + g.op("Constant", value_t=torch.tensor([1], dtype=torch.float)), + add(g, p, eps), + ) + summation = symbolic_helper._reducesum_helper( + g, + # pyrefly: ignore [no-matching-overload] + pow(g, sub(g, input1, input2), p), + axes_i=[-1], + keepdims_i=symbolic_helper._parse_arg(keepdim, "i"), + ) + # pyrefly: ignore [no-matching-overload] + return pow(g, summation, inv_p) + + +@_onnx_symbolic("aten::clone") +# ignore clone operators that are inserted by PyTorch autograd +def clone(g: jit_utils.GraphContext, input, unused_memory_format): + return input + + +@_onnx_symbolic("aten::abs") +def abs(g: jit_utils.GraphContext, self): + return g.op("Abs", self) + + +@_onnx_symbolic("aten::log") +def log(g: jit_utils.GraphContext, self): + return g.op("Log", self) + + +@_onnx_symbolic("aten::log1p") +def log1p(g: jit_utils.GraphContext, self): + return log(g, add(g, symbolic_helper._if_scalar_type_as(torch.ones(1), self), self)) + + +@_onnx_symbolic("aten::log10") +def log10(g: jit_utils.GraphContext, self): + _ln10 = 2.30258509299404568401 + return g.op("Div", log(g, self), g.op("Constant", value_t=torch.tensor([_ln10]))) + + +@_onnx_symbolic("aten::pow") +def pow(g: jit_utils.GraphContext, self, exponent): + f_dtype = _type_utils.JitScalarType.from_value(self) + if not symbolic_helper._is_fp(self): + f_dtype = _type_utils.JitScalarType.FLOAT + self = g.op("Cast", self, to_i=f_dtype.onnx_type()) + if not symbolic_helper._is_fp(exponent): + exponent = g.op( + "Cast", + exponent, + to_i=f_dtype.onnx_type(), + ) + pow = g.op("Pow", self, exponent) + return pow + + +@_onnx_symbolic("aten::clamp") +def clamp(g: jit_utils.GraphContext, self, min, max): + # min or max may be None that we need to dispatch to + # Clip separately, as ONNX does not have None syntax + if symbolic_helper._is_none(min): + return clamp_max(g, self, max) + elif symbolic_helper._is_none(max): + return clamp_min(g, self, min) + else: + if symbolic_helper._is_constant(min) and symbolic_helper._is_constant(max): + return symbolic_helper._op_with_optional_float_cast( + g, + "Clip", + self, + min_f=symbolic_helper._parse_arg(min, "f"), + max_f=symbolic_helper._parse_arg(max, "f"), + opset_before=12, + ) + else: + return clamp_max(g, clamp_min(g, self, min), max) + + +@_onnx_symbolic("aten::clamp_min") +@symbolic_helper.parse_args("v", "v") +def clamp_min(g: jit_utils.GraphContext, self, min): + if symbolic_helper._is_constant(min): + return symbolic_helper._op_with_optional_float_cast( + g, "Clip", self, min_f=symbolic_helper._parse_arg(min, "f"), opset_before=12 + ) + else: + dtype = _type_utils.JitScalarType.from_value(self) + min = g.op("Cast", min, to_i=dtype.onnx_type()) + return symbolic_helper._op_with_optional_float_cast( + g, "Max", self, min, opset_before=12 + ) + + +@_onnx_symbolic("aten::clamp_max") +@symbolic_helper.parse_args("v", "v") +def clamp_max(g: jit_utils.GraphContext, self, max): + if symbolic_helper._is_constant(max): + return symbolic_helper._op_with_optional_float_cast( + g, "Clip", self, max_f=symbolic_helper._parse_arg(max, "f"), opset_before=12 + ) + else: + dtype = _type_utils.JitScalarType.from_value(self) + max = g.op("Cast", max, to_i=dtype.onnx_type()) + return symbolic_helper._op_with_optional_float_cast( + g, "Min", self, max, opset_before=12 + ) + + +@_onnx_symbolic("aten::max") +# torch.max (same for torch.min) actually has two interfaces smashed together: +# torch.max(x, dim, keepdim) and torch.max(x, y) +# TODO(justinchuby): Support multiple quantized args in output +def max(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None): + return symbolic_helper._max_helper(g, self, dim_or_y, keepdim) + + +@_onnx_symbolic("aten::maximum") +@symbolic_helper.quantized_args(True, True) +def maximum(g: jit_utils.GraphContext, input, other): + # pyrefly: ignore [no-matching-overload] + return max(g, input, dim_or_y=other) + + +@_onnx_symbolic("aten::min") +# TODO(justinchuby): Support multiple quantized args in output +def min(g: jit_utils.GraphContext, self, dim_or_y=None, keepdim=None): + return symbolic_helper._min_helper(g, self, dim_or_y, keepdim) + + +@_onnx_symbolic("aten::minimum") +@symbolic_helper.quantized_args(True, True) +def minimum(g: jit_utils.GraphContext, input, other): + # pyrefly: ignore [no-matching-overload] + return min(g, input, dim_or_y=other) + + +@_onnx_symbolic("aten::amax") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "is", "i") +def amax(g: jit_utils.GraphContext, self, dim, keepdim): + return g.op("ReduceMax", self, axes_i=dim, keepdims_i=keepdim) + + +@_onnx_symbolic("aten::amin") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "is", "i") +def amin(g: jit_utils.GraphContext, self, dim, keepdim): + return g.op("ReduceMin", self, axes_i=dim, keepdims_i=keepdim) + + +@_onnx_symbolic("aten::aminmax") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "v", "i") +def aminmax(g: jit_utils.GraphContext, self, dim, keepdim): + reduce_kwargs = {"keepdims_i": keepdim} + if not symbolic_helper._is_none(dim): + dim = symbolic_helper._get_const(dim, "i", "dim") + reduce_kwargs["axes_i"] = [dim] + + return g.op("ReduceMin", self, **reduce_kwargs), g.op( + "ReduceMax", self, **reduce_kwargs + ) + + +@_onnx_symbolic("aten::exp") +def exp(g: jit_utils.GraphContext, self): + return g.op("Exp", self) + + +@_onnx_symbolic("aten::dropout_") +@_onnx_symbolic("aten::dropout") +@symbolic_helper.parse_args("v", "f", "i") +def dropout(g: jit_utils.GraphContext, input, p, train): + symbolic_helper.check_training_mode(train, "dropout") + # if train is False, dropout is no-op + if not train: + return input + r, _ = g.op("Dropout", input, ratio_f=p, outputs=2) + return r + + +@_onnx_symbolic( + "aten::alpha_dropout_", + decorate=[symbolic_helper._apply_params("aten::alpha_dropout_")], +) # See Note [Export inplace] +@_onnx_symbolic( + "aten::feature_alpha_dropout_", + decorate=[symbolic_helper._apply_params("aten::feature_alpha_dropout_")], +) +@_onnx_symbolic( + "aten::feature_dropout_", + decorate=[symbolic_helper._apply_params("aten::feature_dropout_")], +) +@_onnx_symbolic( + "aten::feature_alpha_dropout", + decorate=[symbolic_helper._apply_params("aten::feature_alpha_dropout")], +) +@_onnx_symbolic( + "aten::alpha_dropout", + decorate=[symbolic_helper._apply_params("aten::alpha_dropout")], +) +@_onnx_symbolic( + "aten::feature_dropout", + decorate=[symbolic_helper._apply_params("aten::feature_dropout")], +) +def _unsupported_dropout(name: str): + @symbolic_helper.parse_args("v", "none", "b") + def feature_dropout(g, input, p, train): + # NB: In inference mode, FeatureDropout is exported as an identity op. + if train: + return symbolic_helper._unimplemented(name, "training mode", input) + return input + + return feature_dropout + + +@_onnx_symbolic("aten::norm") +@symbolic_helper.parse_args("v", "t", "is", "i", "v") +def norm(g: jit_utils.GraphContext, self, p, dim, keepdim, dtype=None): + if p == 1: + f = symbolic_helper._reduce_op_symbolic_helper("ReduceL1") + elif p == 2: + f = symbolic_helper._reduce_op_symbolic_helper("ReduceL2") + else: + raise errors.SymbolicValueError( + "ONNX export only p-norms with p of 1 or 2", self + ) + result = f(g, self, dim=dim, keepdim=keepdim) + if dtype is not None: + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + result = g.op("Cast", result, to_i=_type_utils.JitScalarType(dtype).onnx_type()) + return result + + +@_onnx_symbolic("aten::conv_tbc") +@symbolic_helper.parse_args("v", "v", "v", "i") +def conv_tbc(g: jit_utils.GraphContext, input, weight, bias, pad): + # input must have 3 dimensions, see: + # https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/ConvolutionTBC.cpp#L8-L10 + # input = (time, batch, in_channels) + # weight = (kernel_width, in_channels, out_channels) + # bias = (out_channels,) + input = g.op("Transpose", input, perm_i=[1, 2, 0]) + weight = g.op("Transpose", weight, perm_i=[2, 1, 0]) + conv = conv1d(g, input, weight, bias, [1], [pad], [1], 1) + return g.op("Transpose", conv, perm_i=[2, 0, 1]) + + +@_onnx_symbolic("aten::_unique") +@symbolic_helper.parse_args("v", "i", "i") +def _unique(g: jit_utils.GraphContext, input, sorted, return_inverse): + return symbolic_helper._onnx_unsupported("_unique", input) + + +@_onnx_symbolic("aten::_unique2") +@symbolic_helper.parse_args("v", "i", "i", "i") +def _unique2( + g: jit_utils.GraphContext, input, sorted, return_inverse, return_counts +) -> None: + symbolic_helper._onnx_opset_unsupported("_unique2", 9, 11, input) + + +@_onnx_symbolic("aten::empty") +@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v") +def empty( + g: jit_utils.GraphContext, + sizes, + dtype, + layout, + device, + pin_memory=False, + memory_format=None, +): + return zeros(g, sizes, dtype, layout, device, pin_memory) + + +@_onnx_symbolic("aten::empty_like") +@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v") +def empty_like( + g: jit_utils.GraphContext, + input, + dtype=None, + layout=None, + device=None, + pin_memory=False, + memory_format=None, +): + return zeros_like(g, input, dtype, layout, device, pin_memory) + + +@_onnx_symbolic("aten::new_empty") +def new_empty( + g: jit_utils.GraphContext, self, sizes, dtype, layout, device, pin_memory=False +): + self_dtype = symbolic_helper._try_get_scalar_type(self) + if symbolic_helper._is_none(dtype) and self_dtype is not None: + dtype = self_dtype + return empty(g, sizes, dtype, layout, device, pin_memory) + + +@_onnx_symbolic("aten::scalar_tensor") +def scalar_tensor(g: jit_utils.GraphContext, scalar, dtype, *options): + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + if dtype is None: + dtype = _type_utils.JitScalarType.FLOAT + scalar = g.op("Cast", scalar, to_i=_type_utils.JitScalarType(dtype).onnx_type()) + return scalar + + +@_onnx_symbolic("aten::tensor") +def tensor( + g: jit_utils.GraphContext, data, dtype=None, device=None, requires_grad=False +): + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + if symbolic_helper._is_packed_list(data): + if dtype is None: + dtype = _type_utils.JitScalarType.from_value( + symbolic_helper._unpack_list(data)[0] + ) + input_list = [] + for t in symbolic_helper._unpack_list(data): + shape_reference = g.op("Constant", value_t=torch.LongTensor([1])) + t = symbolic_helper._reshape_helper(g, t, shape_reference) + t = g.op("Cast", t, to_i=_type_utils.JitScalarType(dtype).onnx_type()) + input_list.append(t) + return g.op("Concat", *input_list, axis_i=0) + else: + if dtype is None: + dtype = _type_utils.JitScalarType.from_value(data) + if symbolic_helper._is_list(data) and ( + symbolic_helper._is_tensor_list(data) + or symbolic_helper._is_scalar_list(data) + ): + data = g.op("ConcatFromSequence", data, axis_i=0, new_axis_i=1) + return g.op("Cast", data, to_i=_type_utils.JitScalarType(dtype).onnx_type()) + + +@_onnx_symbolic("aten::as_tensor") +def as_tensor(g: jit_utils.GraphContext, data, dtype=None, device=None): + return tensor(g, data, dtype, device) + + +@_onnx_symbolic("aten::zeros") +@symbolic_helper.parse_args("v", "i", "v", "v", "v") +def zeros(g: jit_utils.GraphContext, sizes, dtype, layout, device, pin_memory=False): + # NOTE: no way to set device, layout and pin_memory in ONNX, so we ignore it + if dtype is None: + scalar_type = _type_utils.JitScalarType.FLOAT + else: + scalar_type = _type_utils.JitScalarType(dtype) + sizes_ = symbolic_helper._maybe_get_const(sizes, "is") + if isinstance(sizes_, list) and len(sizes_) == 0: + sizes = g.op("Constant", value_t=torch.tensor([]).to(torch.int64)) + return g.op( + "ConstantOfShape", + sizes, + value_t=torch.tensor([0], dtype=scalar_type.dtype()), + ) + + +@_onnx_symbolic("aten::zeros_like") +@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v") +def zeros_like( + g: jit_utils.GraphContext, + input, + dtype=None, + layout=None, + device=None, + pin_memory=False, + memory_format=None, +): + shape = g.op("Shape", input) + if symbolic_helper._is_none(dtype): + scalar_type = _type_utils.JitScalarType.from_value( + input, _type_utils.JitScalarType.FLOAT + ) + else: + # pyrefly: ignore [bad-argument-type] + scalar_type = _type_utils.JitScalarType(dtype) + return g.op( + "ConstantOfShape", + shape, + value_t=torch.tensor([0], dtype=scalar_type.dtype()), + ) + + +@_onnx_symbolic("aten::new_zeros") +def new_zeros( + g: jit_utils.GraphContext, self, sizes, dtype, layout, device, pin_memory=False +): + self_dtype = symbolic_helper._try_get_scalar_type(self) + + if symbolic_helper._is_none(dtype) and self_dtype is not None: + dtype = self_dtype + return zeros(g, sizes, dtype, layout, device, pin_memory) + + +@_onnx_symbolic("aten::zero") +def zero(g: jit_utils.GraphContext, self): + self_dtype = symbolic_helper._try_get_scalar_type(self) + return zeros_like(g, self, self_dtype) + + +@_onnx_symbolic("aten::ones") +@symbolic_helper.parse_args("v", "i", "v", "v", "v") +def ones(g: jit_utils.GraphContext, sizes, dtype, layout, device, pin_memory=False): + if dtype is None: + scalar_type = _type_utils.JitScalarType.FLOAT + else: + scalar_type = _type_utils.JitScalarType(dtype) + sizes_ = symbolic_helper._maybe_get_const(sizes, "is") + if isinstance(sizes_, list) and len(sizes_) == 0: + sizes = g.op("Constant", value_t=torch.tensor([]).to(torch.int64)) + return g.op( + "ConstantOfShape", + sizes, + value_t=torch.tensor([1], dtype=scalar_type.dtype()), + ) + + +@_onnx_symbolic("aten::ones_like") +@symbolic_helper.parse_args("v", "i", "v", "v", "v", "v") +def ones_like( + g: jit_utils.GraphContext, + input, + dtype=None, + layout=None, + device=None, + pin_memory=False, + memory_format=None, +): + shape = g.op("Shape", input) + if symbolic_helper._is_none(dtype): + scalar_type = _type_utils.JitScalarType.from_value( + input, _type_utils.JitScalarType.FLOAT + ) + else: + # pyrefly: ignore [bad-argument-type] + scalar_type = _type_utils.JitScalarType(dtype) + return g.op( + "ConstantOfShape", + shape, + value_t=torch.tensor([1], dtype=scalar_type.dtype()), + ) + + +@_onnx_symbolic("aten::new_ones") +def new_ones( + g: jit_utils.GraphContext, self, sizes, dtype, layout, device, pin_memory=False +): + self_dtype = symbolic_helper._try_get_scalar_type(self) + if symbolic_helper._is_none(dtype) and self_dtype is not None: + dtype = self_dtype + return ones(g, sizes, dtype, layout, device, pin_memory) + + +@_onnx_symbolic("aten::full") +def full( + g: jit_utils.GraphContext, sizes, value, dtype, layout, device, pin_memory=False +): + const_value = symbolic_helper._maybe_get_const(value, "t") + if symbolic_helper._is_value(const_value): + dtype = _type_utils.JitScalarType.FLOAT if dtype is None else dtype + tmp = zeros(g, sizes, dtype, layout, device) + return add(g, tmp, value, g.op("Constant", value_t=torch.tensor(1))) + else: + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + if dtype is None: + scalar_type = _type_utils.JitScalarType.FLOAT + else: + scalar_type = _type_utils.JitScalarType(dtype) + sizes_ = symbolic_helper._maybe_get_const(sizes, "is") + if isinstance(sizes_, list) and len(sizes_) == 0: + sizes = g.op("Constant", value_t=torch.tensor([]).to(torch.int64)) + return g.op( + "ConstantOfShape", + sizes, + value_t=const_value.view(1).to(scalar_type.dtype()), + ) + + +@_onnx_symbolic("aten::full_like") +def full_like( + g: jit_utils.GraphContext, + input, + fill_value, + dtype=None, + layout=None, + device=None, + pin_memory=False, + memory_format=None, +): + fill_value = symbolic_helper._maybe_get_const(fill_value, "f") + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + if dtype is None: + scalar_type = _type_utils.JitScalarType.from_value( + input, _type_utils.JitScalarType.FLOAT + ) + else: + scalar_type = _type_utils.JitScalarType(dtype) + if symbolic_helper._is_value(fill_value): + tmp = zeros_like(g, input, dtype, layout, device) + fill_value = g.op("Cast", fill_value, to_i=scalar_type.onnx_type()) + return add(g, tmp, fill_value, g.op("Constant", value_t=torch.tensor(1))) + else: + shape = g.op("Shape", input) + return g.op( + "ConstantOfShape", + shape, + value_t=torch.tensor([fill_value], dtype=scalar_type.dtype()), + ) + + +@_onnx_symbolic("aten::new_full") +def new_full( + g: jit_utils.GraphContext, + self, + size, + fill_value, + dtype, + layout, + device, + pin_memory=False, +): + self_dtype = symbolic_helper._try_get_scalar_type(self) + if symbolic_helper._is_none(dtype) and self_dtype is not None: + dtype = self_dtype + return full(g, size, fill_value, dtype, layout, device, pin_memory) + + +@_onnx_symbolic("aten::eye") +def eye(g: jit_utils.GraphContext, *args): + if len(args) == 5: + # aten::eye(n, dtype, layout, device, pin_memory) + n, dtype, layout, device, _pin_memory = args + dim_size = symbolic_helper._unsqueeze_helper(g, n, [0]) + shape = g.op("Concat", dim_size, dim_size, axis_i=0) + tensor = zeros(g, shape, dtype, layout, device) + return g.op("EyeLike", tensor) + if len(args) == 6: + # aten::eye(n, m, dtype, layout, device, pin_memory) + n, m, dtype, layout, device, _pin_memory = args + shape = g.op( + "Concat", + symbolic_helper._unsqueeze_helper(g, n, [0]), + symbolic_helper._unsqueeze_helper(g, m, [0]), + axis_i=0, + ) + tensor = zeros(g, shape, dtype, layout, device) + return g.op("EyeLike", tensor) + + return symbolic_helper._unimplemented("aten::eye", f"with {len(args)} arguments") + + +@_onnx_symbolic("aten::slice") +def slice(g: jit_utils.GraphContext, self, *args): + if len(args) == 4: + # aten::slice(Tensor self, int dim, int start, int end, int step) -> Tensor + dim, start, end, step = args + step = symbolic_helper._parse_arg(step, "i") + if step != 1: + raise errors.SymbolicValueError("step!=1 is currently not supported", self) + is_start_none = start.node().kind() == "prim::Constant" and isinstance( + start.type(), _C.NoneType + ) + is_end_none = end.node().kind() == "prim::Constant" and isinstance( + end.type(), _C.NoneType + ) + is_start_onnx_const = start.node().kind() == "onnx::Constant" + is_end_onnx_const = end.node().kind() == "onnx::Constant" + if ( + ((not is_start_none) and (not is_start_onnx_const)) + or ((not is_end_none) and (not is_end_onnx_const)) + or dim.node().kind() != "onnx::Constant" + ): + if GLOBALS.operator_export_type == _C_onnx.OperatorExportTypes.ONNX: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of Slice with dynamic inputs. DynamicSlice " + "is a deprecated experimental op. Please use statically allocated " + "variables or export to a higher opset version.", + self, + ) + else: + start_unsqueezed = symbolic_helper._unsqueeze_helper(g, start, [0]) + end_unsqueezed = symbolic_helper._unsqueeze_helper(g, end, [0]) + dim_unsqueezed = symbolic_helper._unsqueeze_helper(g, dim, [0]) + return g.op( + "DynamicSlice", + self, + start_unsqueezed, + end_unsqueezed, + dim_unsqueezed, + ) + else: + start = 0 if is_start_none else symbolic_helper._parse_arg(start, "i") + end = ( + _constants.INT64_MAX + if is_end_none + else symbolic_helper._parse_arg(end, "i") + ) + dim = symbolic_helper._parse_arg(dim, "i") + return symbolic_helper._slice_helper( + g, self, axes=[dim], starts=[start], ends=[end] + ) + elif len(args) == 3: + # aten::slice(t[] l, int start, int end, int step) -> t[] + start, end, step = args + dim = 0 + is_start_none = start.node().kind() == "prim::Constant" and isinstance( + start.type(), _C.NoneType + ) + is_end_none = end.node().kind() == "prim::Constant" and isinstance( + end.type(), _C.NoneType + ) + start = 0 if is_start_none else symbolic_helper._parse_arg(start, "i") + end = ( + _constants.INT64_MAX + if is_end_none + else symbolic_helper._parse_arg(end, "i") + ) + return symbolic_helper._slice_helper( + g, self, axes=[dim], starts=[start], ends=[end] + ) + + return symbolic_helper._unimplemented("aten::slice", f"with {len(args)} arguments") + + +@_onnx_symbolic("aten::hardtanh") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "f", "f") +def hardtanh(g: jit_utils.GraphContext, self: _C.Value, min_val: float, max_val: float): + return symbolic_helper._op_with_optional_float_cast( + g, "Clip", self, min_f=min_val, max_f=max_val, opset_before=12 + ) + + +@_onnx_symbolic("aten::hardswish") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v") +def hardswish(g: jit_utils.GraphContext, self): + hs = hardsigmoid(g, self) + return g.op("Mul", self, hs) + + +@_onnx_symbolic("aten::hardsigmoid") +# Fixed scale and zero_point, discovered from aten/src/ATen/native/quantized/cpu/qhardsigmoid.cpp +@symbolic_helper.quantized_args(True, scale=1.0 / 256.0, zero_point=0) +@symbolic_helper.parse_args("v") +def hardsigmoid(g: jit_utils.GraphContext, self): + # Set alpha_f to 1 / 6 to make op equivalent to PyTorch's definition of Hardsigmoid. + # See https://pytorch.org/docs/stable/generated/torch.nn.Hardsigmoid.html + return g.op("HardSigmoid", self, alpha_f=1 / 6) + + +@_onnx_symbolic("aten::tanhshrink") +@symbolic_helper.parse_args("v") +def tanhshrink(g: jit_utils.GraphContext, self): + return g.op("Sub", self, tanh(g, self)) + + +@_onnx_symbolic("aten::hardshrink") +@symbolic_helper.parse_args("v", "f") +def hardshrink(g: jit_utils.GraphContext, self, lambd): + scalar_type = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.FLOAT + ) + lambd_op = g.op( + "Constant", + value_t=torch.tensor(lambd, dtype=scalar_type.dtype()), + ) + cond = logical_or(g, gt(g, self, lambd_op), lt(g, self, neg(g, lambd_op))) + return g.op( + "Where", + cond, + self, + g.op( + "Constant", + value_t=torch.tensor(0, dtype=scalar_type.dtype()), + ), + ) + + +@_onnx_symbolic("aten::softshrink") +@symbolic_helper.parse_args("v", "f") +def softshrink(g: jit_utils.GraphContext, self, lambd): + scalar_type = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.FLOAT + ) + lambd_op = g.op( + "Constant", + value_t=torch.tensor(lambd, dtype=scalar_type.dtype()), + ) + gt_cond = gt(g, self, lambd_op) + gt_out = g.op( + "Where", + gt_cond, + sub(g, self, lambd_op), + g.op( + "Constant", + value_t=torch.tensor(0, dtype=scalar_type.dtype()), + ), + ) + lt_cond = lt(g, self, neg(g, lambd_op)) + lt_out = g.op( + "Where", + lt_cond, + add(g, self, lambd_op), + g.op( + "Constant", + value_t=torch.tensor(0, dtype=scalar_type.dtype()), + ), + ) + return add(g, gt_out, lt_out) + + +@_onnx_symbolic("aten::alias") +def alias(g: jit_utils.GraphContext, self): + return self + + +@_onnx_symbolic("aten::unsqueeze") +@symbolic_helper.parse_args("v", "i") +def unsqueeze(g: jit_utils.GraphContext, self, dim): + """Implement unsqueezing a pytorch tensor in ONNX by inserting a new dimension at the specified `dim`""" + # Handle negative dim + if dim < 0: + rank = symbolic_helper._get_tensor_rank(self) + if rank is not None: + warnings.warn( + "ONNX export unsqueeze with negative axis " + + str(dim) + + " might cause the onnx model to be incorrect. " + + "Negative axis is not supported in ONNX. " + + "Axis is converted to " + + str(dim + rank + 1) + + " based on input shape at export time. " + + "Passing an tensor of different rank in execution will be incorrect.", + stacklevel=2, + ) + dim = dim + rank + 1 + else: + return symbolic_helper._unimplemented( + "unsqueeze", "negative axis with unknown input rank", self + ) + + return symbolic_helper._unsqueeze_helper(g, self, axes_i=[dim]) + + +@_onnx_symbolic("aten::sort") +# TODO(justinchuby): Support multiple quantized args in output +@symbolic_helper.parse_args("v", "i", "i", "none") +def sort(g: jit_utils.GraphContext, self, dim, descending, out=None): + if out is not None: + symbolic_helper._unimplemented( + "Sort", "Out parameter is not supported for sort", self + ) + self_sizes = symbolic_helper._get_tensor_sizes(self) + try: + dim_size = self_sizes[dim] + except Exception: + # FIXME(justinchuby): Avoid catching Exception. + # Catch a more specific exception instead. + dim_size = None + + if dim_size is None: + return symbolic_helper._unimplemented("Sort", "input size not accessible", self) + + return g.op("TopK", self, k_i=dim_size, axis_i=dim, outputs=2) + + +@_onnx_symbolic("aten::numel") +def numel(g: jit_utils.GraphContext, self): + return symbolic_helper._numel_helper(g, self) + + +@_onnx_symbolic("aten::topk") +# TODO(justinchuby): Support multiple quantized args in output +@symbolic_helper.parse_args("v", "i", "i", "i", "i", "none") +def topk(g: jit_utils.GraphContext, self, k, dim, largest, sorted, out=None): + if out is not None: + symbolic_helper._unimplemented( + "TopK", "Out parameter is not supported for topk", self + ) + if not largest: + symbolic_helper._unimplemented("TopK", "Ascending TopK is not supported", self) + + return g.op("TopK", self, k_i=k, axis_i=dim, outputs=2) + + +@_onnx_symbolic("prim::convert_element_type") +def convert_element_type(g: jit_utils.GraphContext, self, *args): + dtype = symbolic_helper._get_const(args[0], "i", "dtype") + return g.op("Cast", self, to_i=_type_utils.JitScalarType(dtype).onnx_type()) + + +@_onnx_symbolic("aten::to") +def to(g: jit_utils.GraphContext, self, *args): + def is_aten_to_device_only(args): + if len(args) == 4: + # aten::to(Tensor, Device, bool, bool, memory_format) + return ( + args[0].node().kind() == "prim::device" + or args[0].type().isSubtypeOf(_C.ListType.ofInts()) + or isinstance(args[0].type(), _C.DeviceObjType) + ) + elif len(args) == 5: + # aten::to(Tensor, Device, ScalarType, bool, bool, memory_format) + # When dtype is None, this is a aten::to(device) call + dtype = symbolic_helper._get_const(args[1], "i", "dtype") + return dtype is None + elif len(args) in (6, 7): + # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, memory_format) -> Tensor + # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, bool, memory_format) -> Tensor + # When dtype is None, this is a aten::to(device) call + dtype = symbolic_helper._get_const(args[0], "i", "dtype") + return dtype is None + return False + + # ONNX doesn't have a concept of a device, so we ignore device-only casts + if is_aten_to_device_only(args): + return self + + if len(args) == 4: + # TestONNXRuntime::test_ones_bool shows args[0] of aten::to() can be onnx::Constant[value=]() + # In this case, the constant value is a tensor not int, + # so symbolic_helper._maybe_get_const(args[0], 'i') would not work. + dtype = args[0] + if ( + symbolic_helper._is_value(args[0]) + and args[0].node().kind() == "onnx::Constant" + ): + tval = symbolic_helper._node_get(args[0].node(), "value") + if isinstance(tval, torch.Tensor): + if len(tval.shape) == 0: + tval = tval.item() + dtype = int(tval) + else: + dtype = tval + + if symbolic_helper._is_value(dtype) or isinstance(dtype, torch.Tensor): + # aten::to(Tensor, Tensor, bool, bool, memory_format) + dtype = _type_utils.JitScalarType.from_value(args[0]) + return g.op( + "Cast", + self, + to_i=dtype.onnx_type(), + ) + else: + # aten::to(Tensor, ScalarType, bool, bool, memory_format) + # memory_format is ignored + return g.op("Cast", self, to_i=_type_utils.JitScalarType(dtype).onnx_type()) + elif len(args) == 5: + # aten::to(Tensor, Device, ScalarType, bool, bool, memory_format) + dtype = symbolic_helper._get_const(args[1], "i", "dtype") + # memory_format is ignored + return g.op("Cast", self, to_i=_type_utils.JitScalarType(dtype).onnx_type()) + elif len(args) == 6: + # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, memory_format) -> Tensor + dtype = symbolic_helper._get_const(args[0], "i", "dtype") + # Layout, device and memory_format are ignored + return g.op("Cast", self, to_i=_type_utils.JitScalarType(dtype).onnx_type()) + elif len(args) == 7: + # aten::to(Tensor, ScalarType, Layout, Device, bool, bool, bool, memory_format) -> Tensor + dtype = symbolic_helper._get_const(args[0], "i", "dtype") + # Layout, device and memory_format are ignored + return g.op("Cast", self, to_i=_type_utils.JitScalarType(dtype).onnx_type()) + + return symbolic_helper._onnx_unsupported("Unknown aten::to signature", self) + + +@_onnx_symbolic("aten::repeat") +def repeat(g: jit_utils.GraphContext, self, repeats): + dtype = _type_utils.JitScalarType.INT64 + shape_ = ones_like(g, repeats, dtype) + self = g.op("Expand", self, shape_) + return g.op("Tile", self, repeats) + + +@_onnx_symbolic("aten::repeat_interleave") +def repeat_interleave( + g: jit_utils.GraphContext, self, repeats, dim=None, output_size=None +): + repeats_dim = symbolic_helper._get_tensor_rank(repeats) + repeats_sizes = symbolic_helper._get_tensor_sizes(repeats) + input_sizes = symbolic_helper._get_tensor_sizes(self) + if repeats_dim is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of repeat_interleave for unknown repeats rank.", + self, + ) + if repeats_sizes is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of repeat_interleave for unknown repeats size.", + self, + ) + if input_sizes is None: + raise errors.SymbolicValueError( + "Unsupported: ONNX export of repeat_interleave for unknown input size.", + self, + ) + + # if dim is None flatten + # By default, use the flattened input array, and return a flat output array + if symbolic_helper._is_none(dim): + self = symbolic_helper._reshape_helper( + g, self, g.op("Constant", value_t=torch.tensor([-1])) + ) + dim = torch.tensor(0, dtype=torch.int64) + else: + dim = symbolic_helper._maybe_get_scalar(dim) + + # Handle cases where dim is negative + if dim < 0: + dim += len(input_sizes) + + input_sizes_temp = input_sizes.copy() + for idx, input_size in enumerate(input_sizes): + if input_size is None: + input_sizes[idx], input_sizes_temp[idx] = 0, -1 + + # Cases where repeats is an int or single value tensor + if repeats_dim == 0 or (repeats_dim == 1 and repeats_sizes[0] == 1): + if input_sizes[dim] == 0: + return symbolic_helper._onnx_opset_unsupported_detailed( + "repeat_interleave", + 9, + 13, + "Unsupported along dimension with unknown input size", + self, + ) + return symbolic_helper._repeat_interleave_single_value_repeat_helper( + g, self, repeats, dim + ) + + # Cases where repeats is a 1 dim Tensor + elif repeats_dim == 1: + if input_sizes[dim] == 0: + return symbolic_helper._onnx_opset_unsupported_detailed( + "repeat_interleave", + 9, + 13, + "Unsupported along dimension with unknown input size", + self, + ) + if repeats_sizes[0] is None: + return symbolic_helper._onnx_opset_unsupported_detailed( + "repeat_interleave", + 9, + 13, + "Unsupported for cases with dynamic repeats", + self, + ) + if repeats_sizes[0] != input_sizes[dim]: + raise AssertionError("repeats must have the same size as input along dim") + reps = repeats_sizes[0] + else: + raise errors.SymbolicValueError("repeats must be 0-dim or 1-dim tensor", self) + + final_splits = [] + r_splits = symbolic_helper._repeat_interleave_split_helper(g, repeats, reps, 0) + i_splits = symbolic_helper._repeat_interleave_split_helper(g, self, reps, dim) + input_sizes[dim], input_sizes_temp[dim] = -1, 1 + for idx, r_split in enumerate(r_splits): + i_split = unsqueeze(g, i_splits[idx], dim + 1) + r_concat = [ + g.op("Constant", value_t=torch.LongTensor(input_sizes_temp[: dim + 1])), + r_split, + g.op("Constant", value_t=torch.LongTensor(input_sizes_temp[dim + 1 :])), + ] + r_concat = g.op("Concat", *r_concat, axis_i=0) + i_split = expand(g, i_split, r_concat, None) + i_split = symbolic_helper._reshape_helper( + g, + i_split, + g.op("Constant", value_t=torch.LongTensor(input_sizes)), + allowzero=0, + ) + final_splits.append(i_split) + return g.op("Concat", *final_splits, axis_i=dim) + + +@_onnx_symbolic("aten::pixel_shuffle") +@symbolic_helper.parse_args("v", "i") +def pixel_shuffle(g: jit_utils.GraphContext, self, upscale_factor): + dims = symbolic_helper._get_tensor_sizes(self) + if len(dims) != 4: + return symbolic_helper._unimplemented( + "pixel_shuffle", "only support 4d input", self + ) + if any(i is None for i in dims[1:]): + after_view = symbolic_helper._reshape_helper( + g, + symbolic_helper._unsqueeze_helper(g, self, [2, 3]), + g.op( + "Constant", + value_t=torch.tensor([0, -1, upscale_factor, upscale_factor, 0, 0]), + ), + allowzero=0, + ) + after_transpose = g.op("Transpose", after_view, perm_i=[0, 1, 4, 2, 5, 3]) + # For dynamic input shapes, two reshapes are performed + reshape_h = symbolic_helper._reshape_helper( + g, + after_transpose, + g.op("Constant", value_t=torch.tensor([0, 0, -1, 1, 0, 0])), + allowzero=0, + ) + reshape_w = symbolic_helper._reshape_helper( + g, + reshape_h, + g.op("Constant", value_t=torch.tensor([0, 0, 0, 0, -1, 1])), + allowzero=0, + ) + return symbolic_helper._squeeze_helper(g, reshape_w, [3, 5]) + else: + output_channel = dims[1] // upscale_factor // upscale_factor + after_view = symbolic_helper._reshape_helper( + g, + self, + g.op( + "Constant", + value_t=torch.tensor( + [ + -1, + output_channel, + upscale_factor, + upscale_factor, + dims[2], + dims[3], + ] + ), + ), + allowzero=0, + ) + after_transpose = g.op("Transpose", after_view, perm_i=[0, 1, 4, 2, 5, 3]) + return symbolic_helper._reshape_helper( + g, + after_transpose, + g.op( + "Constant", + value_t=torch.tensor( + [ + -1, + output_channel, + dims[2] * upscale_factor, + dims[3] * upscale_factor, + ] + ), + ), + allowzero=0, + ) + + +@_onnx_symbolic("aten::pixel_unshuffle") +@symbolic_helper.parse_args("v", "i") +def pixel_unshuffle(g: jit_utils.GraphContext, self, downscale_factor): + dims = symbolic_helper._get_tensor_sizes(self) + if len(dims) != 4: + return symbolic_helper._unimplemented( + "pixel_shuffle", "only support 4d input", self + ) + if any(i is None for i in dims[1:]): + # For dynamic input shapes, two reshapes are performed + reshape_h = symbolic_helper._reshape_helper( + g, + symbolic_helper._unsqueeze_helper(g, self, [3]), + g.op("Constant", value_t=torch.tensor([0, 0, -1, downscale_factor, 0])), + allowzero=0, + ) + reshape_w = symbolic_helper._reshape_helper( + g, + reshape_h, + g.op("Constant", value_t=torch.tensor([0, 0, 0, 0, -1, downscale_factor])), + allowzero=0, + ) + after_transpose = g.op("Transpose", reshape_w, perm_i=[0, 1, 3, 5, 2, 4]) + final_reshape = symbolic_helper._reshape_helper( + g, + after_transpose, + g.op("Constant", value_t=torch.tensor([0, -1, 1, 1, 0, 0])), + allowzero=0, + ) + return symbolic_helper._squeeze_helper(g, final_reshape, [2, 3]) + else: + output_channel = dims[1] * downscale_factor * downscale_factor + after_view = symbolic_helper._reshape_helper( + g, + self, + g.op( + "Constant", + value_t=torch.tensor( + [ + -1, + dims[1], + dims[2] // downscale_factor, + downscale_factor, + dims[3] // downscale_factor, + downscale_factor, + ] + ), + ), + allowzero=0, + ) + after_transpose = g.op("Transpose", after_view, perm_i=[0, 1, 3, 5, 2, 4]) + return symbolic_helper._reshape_helper( + g, + after_transpose, + g.op( + "Constant", + value_t=torch.tensor( + [ + -1, + output_channel, + dims[2] // downscale_factor, + dims[3] // downscale_factor, + ] + ), + ), + allowzero=0, + ) + + +def _generic_rnn( + g: jit_utils.GraphContext, + variant, + input, + initial_states, + all_weights, + has_biases, + num_layers, + dropout, + train, + bidirectional, + batch_first=None, + batch_sizes=None, +): + warnings.warn( + "Exporting a model to ONNX with a batch_size other than 1, " + + "with a variable length with " + + variant + + " can cause an error " + + "when running the ONNX model with a different batch size. " + + "Make sure to save the model with a batch size of 1, " + + "or define the initial states (h0/c0) as inputs of the model. ", + stacklevel=2, + ) + + onnxActivations = [ + "Relu", + "Tanh", + "Sigmoid", + "Affine", + "LeakyRelu", + "ThresholdedRelu", + "ScaledTanh", + "HardSigmoid", + "Elu", + "Softsign", + "Softplus", + ] + variantToOnnxActivationMap = dict( + zip([act_fun.lower() for act_fun in onnxActivations], onnxActivations) + ) + weights_per_layer = 4 if has_biases else 2 + # this means that projections are used inside LSTM, so need to tell user that it's not supported + if variant == "LSTM" and len(all_weights) != num_layers * weights_per_layer * ( + 1 + bidirectional + ): + return symbolic_helper._unimplemented("LSTM", "LSTMs with projections", input) + expected_weights = num_layers * weights_per_layer * (1 + bidirectional) + if len(all_weights) != expected_weights: + raise AssertionError( + f"Expected {expected_weights} weights, got {len(all_weights)}" + ) + layer_weights = [ + all_weights[i : i + weights_per_layer] + for i in range(0, len(all_weights), weights_per_layer) + ] + if batch_first: + # batch, seq, feat -> seq, batch, feat + input = g.op("Transpose", input, perm_i=[1, 0, 2]) + if dropout and train: + return symbolic_helper._unimplemented( + "RNN/GRU/LSTM", "dropout in training mode", input + ) + + if variant.startswith("RNN"): + nonlinearity = variantToOnnxActivationMap[variant[4:].lower()] + variant = "RNN" + + w_hh = all_weights[1] + hidden_size = symbolic_helper._get_tensor_dim_size(w_hh, 1) + if hidden_size is None: + return symbolic_helper._unimplemented( + "RNN/GRU/LSTM", "unknown hidden size", input + ) + + unidirectional = not bidirectional + + prev_output = input + + h_outs = [] + if variant == "RNN" or variant == "GRU": + h0 = initial_states + elif variant == "LSTM": + h0, c0 = initial_states + c_outs = [] + + sequence_lens = unused(g) if batch_sizes is None else batch_sizes + + if variant == "GRU": + # pytorch is reset, input, hidden + # onnx is input, reset, hidden + reform_permutation = [(1, 2), (0, 1), (2, 3)] + elif variant == "LSTM": + # pytorch is input, forget, cell, output. + # onnx is input, output, forget, cell. + reform_permutation = [(0, 1), (3, 4), (1, 3)] + + def reform_weights(g, w, n, intervals): + slices = [ + symbolic_helper._slice_helper(g, w, axes=[0], starts=[x * n], ends=[y * n]) + for x, y in intervals + ] + return g.op("Concat", *slices, axis_i=0) + + def transform_weights_no_bias(layer_index): + weights = layer_weights[layer_index] + if variant == "RNN": + weight_ih, weight_hh = weights + elif variant == "GRU" or variant == "LSTM": + weight_ih, weight_hh = ( + reform_weights(g, w, hidden_size, reform_permutation) for w in weights + ) + return tuple( + symbolic_helper._unsqueeze_helper(g, x, [0]) + for x in (weight_ih, weight_hh) # type: ignore[possibly-undefined] + ) + + def transform_weights(layer_index): + weights = layer_weights[layer_index] + if variant == "RNN": + weight_ih, weight_hh, bias_ih, bias_hh = weights + elif variant == "GRU" or variant == "LSTM": + weight_ih, weight_hh, bias_ih, bias_hh = ( + reform_weights(g, w, hidden_size, reform_permutation) for w in weights + ) + bias_concat = g.op("Concat", bias_ih, bias_hh, axis_i=0) # type: ignore[possibly-undefined] + return tuple( + symbolic_helper._unsqueeze_helper(g, x, [0]) + for x in (weight_ih, weight_hh, bias_concat) # type: ignore[possibly-undefined] + ) + + def retrieve_state(x, start, end): + return ( + x + if num_layers == 1 + else symbolic_helper._slice_helper( + g, x, axes=[0], starts=[start], ends=[end] + ) + ) + + for i in range(num_layers): + if unidirectional: + if weights_per_layer == 4: + weight_ih, weight_hh, bias_concat = transform_weights(i) + else: + weight_ih, weight_hh = transform_weights_no_bias(i) + bias_concat = unused(g) + + state_indices = i, i + 1 + else: + if weights_per_layer == 4: + weight_ih_f, weight_hh_f, bias_f = transform_weights(2 * i) + weight_ih_b, weight_hh_b, bias_b = transform_weights(2 * i + 1) + bias_concat = g.op("Concat", bias_f, bias_b, axis_i=0) + else: + weight_ih_f, weight_hh_f = transform_weights_no_bias(2 * i) + weight_ih_b, weight_hh_b = transform_weights_no_bias(2 * i + 1) + bias_concat = unused(g) + + weight_ih = g.op("Concat", weight_ih_f, weight_ih_b, axis_i=0) + weight_hh = g.op("Concat", weight_hh_f, weight_hh_b, axis_i=0) + + state_indices = 2 * i, 2 * i + 2 + + inputs = [prev_output, weight_ih, weight_hh, bias_concat, sequence_lens] + + inputs.append(retrieve_state(h0, *state_indices)) # type: ignore[possibly-undefined] + if variant == "LSTM": + inputs.append(retrieve_state(c0, *state_indices)) # type: ignore[possibly-undefined] + + extra_kwargs = {} if unidirectional else {"direction_s": "bidirectional"} + if variant == "RNN": + if bidirectional: + activation = [nonlinearity, nonlinearity] # type: ignore[possibly-undefined] + else: + activation = [nonlinearity] # type: ignore[possibly-undefined] + + prev_output, h_out = g.op( + "RNN", + *inputs, + outputs=2, + hidden_size_i=hidden_size, + activations_s=activation, + **extra_kwargs, + ) + elif variant == "GRU": + prev_output, h_out = g.op( + "GRU", + *inputs, + outputs=2, + hidden_size_i=hidden_size, + linear_before_reset_i=1, + **extra_kwargs, + ) + elif variant == "LSTM": + prev_output, h_out, c_out = g.op( + "LSTM", *inputs, outputs=3, hidden_size_i=hidden_size, **extra_kwargs + ) + + if bidirectional: + # The ONNX RNN/GRU/LSTM produce an output of dimensions + # seq_len, num_directions, batch, hidden_size + # We have to convert to match pytorch's expected + # seq_len, batch, num_directions * hidden_size + # by first moving num_directions before hidden_size with + # Transpose, and then combining it with hidden_size + # with Reshape. + prev_output = g.op("Transpose", prev_output, perm_i=[0, 2, 1, 3]) + prev_output = symbolic_helper._reshape_helper( + g, + prev_output, + g.op("Constant", value_t=torch.LongTensor([0, 0, -1])), + allowzero=0, + ) + else: + prev_output = symbolic_helper._squeeze_helper(g, prev_output, [1]) + + h_outs.append(h_out) # type: ignore[possibly-undefined] + if variant == "LSTM": + c_outs.append(c_out) # type: ignore[possibly-undefined] + if batch_first: + # seq, batch, num_directions * hidden_size -> batch, seq, num_directions * hidden_size + prev_output = g.op("Transpose", prev_output, perm_i=[1, 0, 2]) + h_outs = h_out if num_layers == 1 else g.op("Concat", *h_outs, axis_i=0) # type: ignore[possibly-undefined] + if variant == "RNN" or variant == "GRU": + return prev_output, h_outs + elif variant == "LSTM": + c_outs = c_out if num_layers == 1 else g.op("Concat", *c_outs, axis_i=0) # type: ignore[possibly-undefined] + return prev_output, h_outs, c_outs + + +@symbolic_helper.parse_args("v", "v", "v", "i", "i", "f", "i", "i", "i") +def _lstm_full( + g: jit_utils.GraphContext, + input, + hidden_v, + weight_v, + has_biases, + num_layers, + dropout, + train, + bidirectional, + batch_first, +): + hidden, weight = ( + symbolic_helper._unpack_list(hidden_v), + symbolic_helper._unpack_list(weight_v), + ) + return _generic_rnn( + g, + "LSTM", + input, + hidden, + weight, + has_biases, + num_layers, + dropout, + train, + bidirectional, + batch_first, + ) + + +@symbolic_helper.parse_args("v", "v", "v", "v", "i", "i", "f", "i", "i") +def _lstm_packed( + g: jit_utils.GraphContext, + input, + batch_sizes, + hidden_v, + weight_v, + has_biases, + num_layers, + dropout, + train, + bidirectional, +): + hidden, weight = ( + symbolic_helper._unpack_list(hidden_v), + symbolic_helper._unpack_list(weight_v), + ) + return _generic_rnn( + g, + "LSTM", + input, + hidden, + weight, + has_biases, + num_layers, + dropout, + train, + bidirectional, + batch_sizes=batch_sizes, + ) + + +@_onnx_symbolic("aten::lstm") +def lstm(g: jit_utils.GraphContext, *args): + if symbolic_helper._is_tensor_list(args[3]): + return _lstm_packed(g, *args) + else: + return _lstm_full(g, *args) + + +@_onnx_symbolic("aten::lstm_cell") +def lstm_cell(g: jit_utils.GraphContext, self, hidden, w_ih, w_hh, b_ih, b_hh): + input = symbolic_helper._unsqueeze_helper(g, self, [0]) + hidden = symbolic_helper._unpack_list(hidden) + hidden = [symbolic_helper._unsqueeze_helper(g, x, [0]) for x in hidden] + weight = ( + (w_ih, w_hh, b_ih, b_hh) if symbolic_helper._is_tensor(b_ih) else (w_ih, w_hh) + ) + has_biases = bool(symbolic_helper._is_tensor(b_ih)) + _, h_outs, c_outs = _generic_rnn( + g, + "LSTM", + input, + hidden, + weight, + has_biases, + num_layers=1, + dropout=0, + train=0, + bidirectional=False, + batch_first=False, + ) + return symbolic_helper._squeeze_helper( + g, h_outs, [0] + ), symbolic_helper._squeeze_helper(g, c_outs, [0]) + + +@_onnx_symbolic( + "aten::gru", decorate=[symbolic_helper._apply_params("GRU"), _export("gru")] +) +@_onnx_symbolic( + "aten::rnn_tanh", + decorate=[symbolic_helper._apply_params("RNN_TANH"), _export("rnn_tanh")], +) +@_onnx_symbolic( + "aten::rnn_relu", + decorate=[symbolic_helper._apply_params("RNN_RELU"), _export("rnn_relu")], +) +def _one_hidden_rnn(kind: str): + @symbolic_helper.parse_args("v", "v", "v", "i", "i", "f", "i", "i", "i") + def _rnn_full( + g, + input, + hidden, + weight_v, + has_biases, + num_layers, + dropout, + train, + bidirectional, + batch_first, + ): + weight = symbolic_helper._unpack_list(weight_v) + return _generic_rnn( + g, + kind, + input, + hidden, + weight, + has_biases, + num_layers, + dropout, + train, + bidirectional, + batch_first, + ) + + @symbolic_helper.parse_args("v", "v", "v", "v", "i", "i", "f", "i", "i") + def _rnn_packed( + g, + input, + batch_sizes, + hidden, + weight_v, + has_biases, + num_layers, + dropout, + train, + bidirectional, + ): + weight = symbolic_helper._unpack_list(weight_v) + return _generic_rnn( + g, + kind, + input, + hidden, + weight, + has_biases, + num_layers, + dropout, + train, + bidirectional, + batch_sizes=batch_sizes, + ) + + def symbolic(g, *args): + if symbolic_helper._is_tensor_list(args[3]): + return _rnn_packed(g, *args) + else: + return _rnn_full(g, *args) + + return symbolic + + +@_onnx_symbolic("aten::_dim_arange") +@symbolic_helper.parse_args("v", "i") +def _dim_arange(g: jit_utils.GraphContext, like, dim): + like_shape = g.op("Shape", like) + stop = g.op( + "Gather", like_shape, g.op("Constant", value_t=torch.tensor(dim)), axis_i=0 + ) + # aten::arange(Scalar end, ScalarType dtype, Layout, Device, bool pin_memory) + return arange(g, stop, 4, None, None, None) + + +@_onnx_symbolic("aten::detach") +def detach(g: jit_utils.GraphContext, input): + # Erase aten::detach nodes because ONNX is inference only + return input + + +@_onnx_symbolic("aten::contiguous") +@symbolic_helper.parse_args("v", "i") +def contiguous(g: jit_utils.GraphContext, input, memory_format): + if memory_format > 2: # allower values are any, preserve and contiguous_format + raise errors.SymbolicValueError( + "onnx memory_format support is not implemented", input + ) + return input + + +@_onnx_symbolic("aten::_pack_padded_sequence") +@symbolic_helper.parse_args("v", "v", "i") +def _pack_padded_sequence(g: jit_utils.GraphContext, input, lengths, batch_first): + # Currently there is no PackPadded operator in ONNX. We rely on an + # optimization pass to remove this later. It is an error if all + # PackPadded operators cannot be optimized out. + if batch_first: + input = g.op("Transpose", input, perm_i=[1, 0, 2]) + if not lengths.type().isSubtypeOf(torch._C.TensorType.get()): + raise errors.SymbolicValueError( + "'lengths' must be a Tensor for ONNX export", input + ) + # We know it's a TensorType so this check is now safe. + # It's really only necessary because those operators expand to something that + # only works with int32 types in Caffe2... + if ( + _type_utils.JitScalarType.from_value( + lengths, _type_utils.JitScalarType.UNDEFINED + ) + != _type_utils.JitScalarType.INT + ): + lengths = g.op("Cast", lengths, to_i=_C_onnx.TensorProtoDataType.INT32) + return g.op("prim::PackPadded", input, lengths, outputs=2) + + +@_onnx_symbolic("aten::_pad_packed_sequence") +@symbolic_helper.parse_args("v", "v", "i", "t", "v") +def _pad_packed_sequence( + g: jit_utils.GraphContext, + data, + batch_sizes, + batch_first, + padding_value, + total_length, +): + # Ignore total_length as it is not supported in _symbolic_pad_packed_sequence + # It is only useful/used when training using data_parallel model, so + # It shouldn't be relevant for ONNX anyway + data, lengths = g.op("prim::PadPacked", data, batch_sizes, outputs=2) + if batch_first: + data = g.op("Transpose", data, perm_i=[1, 0, 2]) + return data, lengths + + +@_onnx_symbolic("aten::randint") +def randint(g: jit_utils.GraphContext, low, high, shapes, dtype, *options): + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + low_i = symbolic_helper._get_const(low, "i", "low") + high_i = symbolic_helper._get_const(high, "i", "high") + if dtype is None: + scalar_type = _type_utils.JitScalarType.INT64 + else: + scalar_type = _type_utils.JitScalarType(dtype) + if low_i is None: + raise symbolic_helper._onnx_unsupported("randint", low) + if high_i is None: + raise symbolic_helper._onnx_unsupported("randint", high) + + shape = symbolic_helper._maybe_get_const(shapes, "is") + if symbolic_helper._is_value(shape): + shape_const = g.op( + "ConstantOfShape", + shapes, + value_t=torch.tensor([0], dtype=torch.float), + ) + randn = g.op( + "RandomUniformLike", + shape_const, + low_f=low_i, + high_f=high_i, + ) + else: + randn = g.op( + "RandomUniform", + shape_i=shape, + low_f=low_i, + high_f=high_i, + ) + + # cast to integer type + int_dtype = _type_utils.JitScalarType.INT64 + randint = g.op("Cast", randn, to_i=int_dtype.onnx_type()) + if int_dtype != scalar_type: + randint = g.op("Cast", randint, to_i=scalar_type.onnx_type()) + return randint + + +@_onnx_symbolic("aten::randint_like") +def randint_like(g: jit_utils.GraphContext, self, low, high, dtype, *options): + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + low_i = symbolic_helper._get_const(low, "i", "low") + high_i = symbolic_helper._get_const(high, "i", "high") + if dtype is None: + scalar_type = _type_utils.JitScalarType.INT64 + else: + scalar_type = _type_utils.JitScalarType(dtype) + if low_i is None: + raise symbolic_helper._onnx_unsupported("randint", low) + if high_i is None: + raise symbolic_helper._onnx_unsupported("randint", high) + + randn = g.op( + "RandomUniformLike", + self, + low_f=low_i, + high_f=high_i, + ) + + # cast to integer type + int_dtype = _type_utils.JitScalarType.INT64 + randint = g.op("Cast", randn, to_i=int_dtype.onnx_type()) + if int_dtype != scalar_type: + randint = g.op("Cast", randint, to_i=scalar_type.onnx_type()) + return randint + + +@_onnx_symbolic("aten::randn") +def randn(g: jit_utils.GraphContext, shapes, dtype, *options): + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + if dtype is None: + scalar_type = _type_utils.JitScalarType.FLOAT + else: + scalar_type = _type_utils.JitScalarType(dtype) + shape = symbolic_helper._maybe_get_const(shapes, "is") + if symbolic_helper._is_value(shape): + shape_const = g.op( + "ConstantOfShape", + shapes, + value_t=torch.tensor([0], dtype=torch.float), + ) + return g.op( + "RandomNormalLike", + shape_const, + dtype_i=scalar_type.onnx_type(), + ) + return g.op( + "RandomNormal", + shape_i=shape, + dtype_i=scalar_type.onnx_type(), + ) + + +@_onnx_symbolic("aten::rand") +def rand(g: jit_utils.GraphContext, shapes, dtype, *options): + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + if dtype is None: + scalar_type = _type_utils.JitScalarType.FLOAT + else: + scalar_type = _type_utils.JitScalarType(dtype) + shape = symbolic_helper._maybe_get_const(shapes, "is") + if symbolic_helper._is_value(shape): + shape_const = g.op( + "ConstantOfShape", + shapes, + value_t=torch.tensor([0], dtype=torch.float), + ) + return g.op( + "RandomUniformLike", + shape_const, + dtype_i=scalar_type.onnx_type(), + ) + return g.op( + "RandomUniform", + shape_i=shape, + dtype_i=scalar_type.onnx_type(), + ) + + +@_onnx_symbolic("aten::randn_like") +def randn_like( + g: jit_utils.GraphContext, + self, + dtype, + layout=None, + device=None, + pin_memory=False, + memory_format=None, +): + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + if dtype is None: + scalar_type = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.FLOAT + ) + else: + scalar_type = _type_utils.JitScalarType(dtype) + return g.op("RandomNormalLike", self, dtype_i=scalar_type.onnx_type()) + + +@_onnx_symbolic("aten::rand_like") +def rand_like( + g: jit_utils.GraphContext, + self, + dtype, + layout=None, + device=None, + pin_memory=False, + memory_format=None, +): + dtype = symbolic_helper._get_const(dtype, "i", "dtype") + if dtype is None: + dtype = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.FLOAT + ) + return g.op( + "RandomUniformLike", self, dtype_i=_type_utils.JitScalarType(dtype).onnx_type() + ) + + +@_onnx_symbolic("aten::rrelu") +@symbolic_helper.parse_args("v", "f", "f", "i", "none") +def rrelu(g: jit_utils.GraphContext, input, lower, upper, training, generator): + if not training: + slope = (upper + lower) / 2.0 + return g.op("LeakyRelu", input, alpha_f=slope) + p = g.op("RandomUniformLike", input, high_f=upper, low_f=lower) + return g.op("PRelu", input, p) + + +@_onnx_symbolic("aten::bernoulli") +def bernoulli(g: jit_utils.GraphContext, input, p=None, generator=None, out=None): + if out is not None and not symbolic_helper._is_none(out): + symbolic_helper._unimplemented( + "Bernoulli", "out parameter is not supported for bernoulli", input + ) + if generator is not None and not symbolic_helper._is_none(generator): + symbolic_helper._unimplemented( + "Bernoulli", "generator is not supported for bernoulli", input + ) + + dtype = _type_utils.JitScalarType.from_value( + input, _type_utils.JitScalarType.UNDEFINED + ) + if dtype == _type_utils.JitScalarType.UNDEFINED: + return symbolic_helper._unimplemented( + "Bernoulli", "input dtype not accessible", input + ) + + rands = g.op( + "RandomUniformLike", + input, + high_f=1.0, + low_f=0.0, + dtype_i=dtype.onnx_type(), + ) + prob = p if p is not None and not symbolic_helper._is_none(p) else input + output = g.op("Less", rands, prob) + return g.op("Cast", output, to_i=dtype.onnx_type()) + + +@_onnx_symbolic("aten::log_sigmoid") +@symbolic_helper.parse_args("v") +def log_sigmoid(g: jit_utils.GraphContext, input): + p = g.op("Sigmoid", input) + return g.op("Log", p) + + +@_onnx_symbolic("aten::erf") +@symbolic_helper.parse_args("v") +def erf(g: jit_utils.GraphContext, input): + return g.op("Erf", input) + + +@_onnx_symbolic("aten::flatten") +@symbolic_helper.quantized_args(True, False, False) +@symbolic_helper.parse_args("v", "i", "i") +def flatten(g: jit_utils.GraphContext, input, start_dim, end_dim): + dim = symbolic_helper._get_tensor_rank(input) + if dim is None: + return symbolic_helper._unimplemented( + "dim", + "ONNX and PyTorch use different strategies to split the input. " + "Input rank must be known at export time.", + input, + ) + + if dim == 0: + return symbolic_helper._reshape_helper(g, input, [1]) + if dim == 1: + return g.op("Identity", input) + # TODO: remove this as onnx opset 11 spec allows negative axes + if end_dim < 0: + end_dim = dim + end_dim + # use ONNX's Flatten operator for cases where the output shape is 2D + if start_dim == 1 and end_dim == dim - 1: + return g.op("Flatten", input, axis_i=start_dim) + if start_dim == 0 and end_dim == dim - 2: + return g.op("Flatten", input, axis_i=end_dim + 1) + + return symbolic_helper._flatten_helper(g, input, start_dim, end_dim, dim) + + +@_onnx_symbolic("aten::nonzero") +@symbolic_helper.parse_args("v") +def nonzero(g: jit_utils.GraphContext, input): + """Emitted from `torch.nonzero(x, as_tuple=False)`""" + return t(g, g.op("NonZero", input)) + + +@_onnx_symbolic("aten::nonzero_numpy") +# Emitted from `torch.nonzero(x, as_tuple=True)` +def nonzero_numpy(g: jit_utils.GraphContext, input, _outputs=None): + return unbind(g, nonzero(g, input), 1, _outputs=_outputs) + + +@_onnx_symbolic("aten::isnan") +@symbolic_helper.parse_args("v") +def isnan(g: jit_utils.GraphContext, input): + output = g.op("IsNaN", input) + return output + + +@_onnx_symbolic("aten::any") +def _any(g: jit_utils.GraphContext, *args): + # aten::any(Tensor self) + if len(args) == 1: + input = args[0] + dim, keepdim = None, 0 + # aten::any(Tensor self, int[]? dim, bool keepdim) + else: + input, dim, keepdim = args + # Can be int list or single int + dim = symbolic_helper._parse_arg(dim, "t") + dim = [int(d) for d in dim.view(-1)] + keepdim = symbolic_helper._parse_arg(keepdim, "i") + input = g.op("Cast", input, to_i=_C_onnx.TensorProtoDataType.INT64) + input_sum = symbolic_helper._reducesum_helper( + g, input, axes_i=dim, keepdims_i=keepdim + ) + return gt(g, input_sum, g.op("Constant", value_t=torch.tensor(0, dtype=torch.long))) + + +@_onnx_symbolic("aten::all") +def _all(g: jit_utils.GraphContext, *args): + input = g.op("Not", args[0]) + # aten::all(Tensor self) + if len(args) == 1: + return g.op("Not", _any(g, input)) + # aten::all(Tensor self, int[]? dim, bool keepdim) + else: + return g.op("Not", _any(g, input, args[1], args[2])) + + +@_onnx_symbolic("aten::narrow") +@symbolic_helper.parse_args("v", "i", "i", "i") +def narrow(g: jit_utils.GraphContext, input, dim, start, length): + return symbolic_helper._slice_helper( + g, input, axes=[dim], starts=[start], ends=[start + length] + ) + + +@_onnx_symbolic("aten::argmax") +@symbolic_helper.parse_args("v", "v", "b") +def argmax( + g: jit_utils.GraphContext, + input: torch._C.Value, + dim: torch._C.Value, + keepdim: bool, +): + return symbolic_helper._argmin_argmax_helper(g, input, dim, keepdim, "ArgMax") + + +@_onnx_symbolic("aten::argmin") +@symbolic_helper.parse_args("v", "v", "b") +def argmin( + g: jit_utils.GraphContext, + input: torch._C.Value, + dim: torch._C.Value, + keepdim: bool, +): + return symbolic_helper._argmin_argmax_helper(g, input, dim, keepdim, "ArgMin") + + +@_onnx_symbolic("aten::scatter") +@symbolic_helper.parse_args("v", "i", "v", "v") +def scatter(g: jit_utils.GraphContext, self, dim, index, src): + src_type = _type_utils.JitScalarType.from_value( + src, _type_utils.JitScalarType.UNDEFINED + ) + src = symbolic_helper._maybe_get_scalar(src) + if symbolic_helper._is_value(src): + return g.op("Scatter", self, index, src, axis_i=dim) + else: + # Check if scalar "src" has same type as self (PyTorch allows different + # type for scalar src (but not when src is tensor)). If not, insert Cast node. + self_scalar_type = _type_utils.JitScalarType.from_value(self) + if self_scalar_type != src_type: + src = g.op("Cast", src, to_i=self_scalar_type.onnx_type()) + return g.op("Scatter", self, index, expand_as(g, src, index), axis_i=dim) + + +@_onnx_symbolic("aten::scatter_add") +@symbolic_helper.parse_args("v", "i", "v", "v") +def scatter_add(g: jit_utils.GraphContext, self, dim, index, src): + scalar_type = symbolic_helper._try_get_scalar_type(self) + if scalar_type is None: + return symbolic_helper._unimplemented( + "scatter_add", "input dtype not accessible", self + ) + sizes = symbolic_helper._get_tensor_sizes(self, allow_nonstatic=False) + if sizes: + to_add = g.op("Constant", value_t=torch.zeros(sizes, dtype=scalar_type.dtype())) + else: + to_add = zeros_like(g, self, scalar_type) + to_add = symbolic_helper._scatter_helper(g, to_add, dim, index, src) + return add(g, self, to_add) + + +@_onnx_symbolic("aten::log2") +def log2(g: jit_utils.GraphContext, self): + _ln2 = 0.693147180559945309 + return g.op("Div", log(g, self), g.op("Constant", value_t=torch.tensor(_ln2))) + + +@_onnx_symbolic("aten::is_floating_point") +def is_floating_point(g: jit_utils.GraphContext, self): + if symbolic_helper._is_fp(self): + return g.op("Constant", value_t=torch.BoolTensor([1])) + return g.op("Constant", value_t=torch.BoolTensor([0])) + + +@_onnx_symbolic("aten::__is_") +def __is_(g: jit_utils.GraphContext, self, other): + if symbolic_helper._is_none(other): + if symbolic_helper._is_none(self): + return g.op("Constant", value_t=torch.BoolTensor([1])) + return g.op("Constant", value_t=torch.BoolTensor([0])) + return eq(g, self, other) + + +@_onnx_symbolic("aten::__isnot_") +@wrap_logical_op_with_negation +def __isnot_(g: jit_utils.GraphContext, self, other): + return __is_(g, self, other) + + +@_onnx_symbolic("aten::one_hot") +def one_hot(g: jit_utils.GraphContext, self, num_classes): + values = g.op("Constant", value_t=torch.LongTensor([0, 1])) + # onnxruntime supports limited type combinations for OneHot. + if _type_utils.JitScalarType.from_value( + num_classes, _type_utils.JitScalarType.UNDEFINED + ) in { + _type_utils.JitScalarType.UINT8, + _type_utils.JitScalarType.INT8, + _type_utils.JitScalarType.INT, + _type_utils.JitScalarType.INT16, + }: + num_classes = g.op("Cast", num_classes, to_i=_C_onnx.TensorProtoDataType.INT64) + return g.op("OneHot", self, num_classes, values, axis_i=-1) + + +@_onnx_symbolic("aten::gather") +@symbolic_helper.parse_args("v", "i", "v", "v") +def gather(g: jit_utils.GraphContext, self, dim, index, sparse_grad=False): + if symbolic_helper._maybe_get_const(sparse_grad, "i"): + return symbolic_helper._unimplemented("gather", "sparse_grad == True", self) + # NOTE: This workaround is needed since GatherElement is only supported + # since opset 11, and Gather in ONNX is not the same as torch.gather. + scalar_type = _type_utils.JitScalarType.from_value(self) + values = g.op("Constant", value_t=torch.LongTensor([0, 1])) + depth = size(g, self, g.op("Constant", value_t=torch.LongTensor([dim]))) + index = g.op( + "Cast", + g.op("OneHot", index, depth, values, axis_i=dim), + to_i=scalar_type.onnx_type(), + ) + mul = g.op("Mul", symbolic_helper._unsqueeze_helper(g, self, [dim + 1]), index) + return symbolic_helper._reducesum_helper(g, mul, axes_i=[dim], keepdims_i=0) + + +@symbolic_helper.parse_args("v", "is", "i", "i") +def _var_mean(g: jit_utils.GraphContext, input, dim, correction, keepdim): + return symbolic_helper._var_mean_helper(g, input, dim, correction, keepdim) + + +@_onnx_symbolic("aten::std") +def std(g: jit_utils.GraphContext, input, *args): + var, _ = var_mean(g, input, *args) + return g.op("Sqrt", var) + + +@_onnx_symbolic("aten::var") +def var(g: jit_utils.GraphContext, input, *args): + var, _ = var_mean(g, input, *args) + return var + + +@_onnx_symbolic("aten::var_mean") +def var_mean(g: jit_utils.GraphContext, input, *args): + if len(args) == 1: + return _var_mean(g, input, None, args[0], None) + else: + return _var_mean(g, input, *args) + + +@_onnx_symbolic("aten::std_mean") +def std_mean(g: jit_utils.GraphContext, input, *args): + var, mean = var_mean(g, input, *args) + return g.op("Sqrt", var), mean + + +@_onnx_symbolic("aten::logsumexp") +@symbolic_helper.parse_args("v", "is", "i") +def logsumexp(g: jit_utils.GraphContext, input, dim, keepdim): + return g.op("ReduceLogSumExp", input, axes_i=dim, keepdims_i=keepdim) + + +@_onnx_symbolic("aten::arange") +def arange(g: jit_utils.GraphContext, *args): + def _get_arange_dtype(dtype): + dtype = symbolic_helper._maybe_get_const(dtype, "i") + return dtype + + def _float_step_convert(range_tensor): + if symbolic_helper._is_fp(range_tensor): + range_tensor = g.op( + "Cast", + g.op("Ceil", range_tensor), + to_i=_type_utils.JitScalarType.INT64.onnx_type(), + ) + return range_tensor + + if len(args) == 2 or len(args) == 5: + if len(args) == 2: + # aten::arange(Scalar end, Tensor out) + dtype = None + else: + # aten::arange(Scalar end, ScalarType dtype, Layout, Device, bool pin_memory) + dtype = _get_arange_dtype(args[1]) + dtype, end, start, step = symbolic_helper._arange_cast_helper( + g, end=args[0], dtype=dtype + ) + end = symbolic_helper._unsqueeze_helper(g, end, [0]) + range_tensor = _float_step_convert(end) + arange_tensor = symbolic_helper._squeeze_helper( + g, nonzero(g, ones(g, range_tensor, dtype, None, None)), [1] + ) + return g.op( + "Cast", arange_tensor, to_i=_type_utils.JitScalarType(dtype).onnx_type() + ) + elif len(args) == 4 or len(args) == 7: + if len(args) == 4: + # aten::arange(Scalar start, Scalar end, Scalar step, Tensor out) + dtype = None + else: + # aten::arange(Scalar start, Scalar end, Scalar step, ScalarType dtype, Layout, Device, bool pin_memory) + dtype = _get_arange_dtype(args[3]) + dtype, end, start, step = symbolic_helper._arange_cast_helper( + g, start=args[0], end=args[1], step=args[2], dtype=dtype + ) + step = symbolic_helper._unsqueeze_helper(g, step, [0]) + end = symbolic_helper._unsqueeze_helper(g, end, [0]) + start = symbolic_helper._unsqueeze_helper(g, start, [0]) + range_tensor = _float_step_convert(g.op("Div", g.op("Sub", end, start), step)) + arange_tensor = symbolic_helper._squeeze_helper( + g, nonzero(g, ones(g, range_tensor, None, None, None)), [1] + ) + arange_tensor = g.op("Add", g.op("Mul", arange_tensor, step), start) + return g.op( + "Cast", arange_tensor, to_i=_type_utils.JitScalarType(dtype).onnx_type() + ) + elif len(args) == 6: + # aten::arange(Scalar start, Scalar end, ScalarType dtype, Layout, Device, bool pin_memory) + dtype = _get_arange_dtype(args[2]) + dtype, end, start, step = symbolic_helper._arange_cast_helper( + g, start=args[0], end=args[1], dtype=dtype + ) + end = symbolic_helper._unsqueeze_helper(g, end, [0]) + start = symbolic_helper._unsqueeze_helper(g, start, [0]) + range_tensor = _float_step_convert(g.op("Sub", end, start)) + arange_tensor = g.op( + "Add", + symbolic_helper._squeeze_helper( + g, nonzero(g, ones(g, range_tensor, dtype, *(args[3:]))), [1] + ), + start, + ) + return g.op( + "Cast", arange_tensor, to_i=_type_utils.JitScalarType(dtype).onnx_type() + ) + + return symbolic_helper._unimplemented("aten::arange", f"with {len(args)} arguments") + + +@_onnx_symbolic("aten::linspace") +def linspace( + g: jit_utils.GraphContext, start, end, steps, dtype, layout, device, pin_memory +): + range_tensor = symbolic_helper._arange_helper(g, steps, None) + step = div( + g, + sub(g, end, start), + sub(g, steps, g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64))), + ) + return add(g, mul(g, range_tensor, step), start) + + +@_onnx_symbolic("aten::lift") +def lift(g: jit_utils.GraphContext, self): + # at::lift() is a no-op from the perspective of tracing for onnx + return self + + +@_onnx_symbolic("aten::masked_fill") +def masked_fill(g: jit_utils.GraphContext, self, mask, value): + """Implement the masked_fill functionality available for a pytorch tensor in ONNX. + + Fills elements of the input tensor with `value` where `mask` is True. + """ + mask = g.op("Cast", mask, to_i=_C_onnx.TensorProtoDataType.BOOL) + value = symbolic_helper._maybe_get_scalar(value) + return g.op("Where", mask, symbolic_helper._if_scalar_type_as(value, self), self) + + +@_onnx_symbolic("aten::masked_fill_") +def masked_fill_(g: jit_utils.GraphContext, self, mask, value): + return masked_fill(g, self, mask, value) + + +@_onnx_symbolic("aten::index") +def index(g: jit_utils.GraphContext, self, index): + if symbolic_helper._is_packed_list(index): + indices = symbolic_helper._unpack_list(index) + else: + indices = [index] + + def try_mask_to_index(index): + if not symbolic_helper._is_none(index) and ( + _type_utils.JitScalarType.from_value( + index, _type_utils.JitScalarType.UNDEFINED + ) + == _type_utils.JitScalarType.UINT8 + or symbolic_helper._is_bool(index) + ): + if g.opset < 9: + raise errors.SymbolicValueError( + "Exporting masked indices are only supported after ONNX opset 9.", + self, + ) + warnings.warn( + "Exporting aten::index operator with indices of type Byte. " + "Only 1-D indices are supported. In any other case, " + "this will produce an incorrect ONNX graph.", + stacklevel=2, + ) + index = symbolic_helper._squeeze_helper(g, nonzero(g, index), [1]) + return index + + indices = [try_mask_to_index(idx) for idx in indices] + if len(indices) == 1: + return symbolic_helper._select_helper( + g, self, 0, indices[0], apply_reshape=False + ) + else: + # Multiple tensors as indices. Each tensor could either be + # 1. prim::Constant() + # representing ":" in python indexing. E.g. tensor[:, :] + # 2. prim::Constant[value=...] or tensor output + # representing advanced indexing. E.g. tensor[[0, 1], [2, 0]]. + # For more info on advanced indexing, + # check https://numpy.org/doc/stable/user/basics.indexing.html#advanced-indexing + + # Consider a general case of + # t: [x_1, y_1, y_2, ..., x_m, ..., y_n] + # where t is a tensor of rank m+n, {x_i} are axes where tensor index is provided, and {y_i} are axes for ":". + # Same results can be achieved through transposing t into + # t: [x_1, x_2, ..., x_m, y_1, y_2, ..., y_n] + # and use gatherND. However ONNX does not have gatherND, to use 1d gather we'll need to flatten t + # and process the tensor indices. + # t: [x_1 * x_2 * ... * x_m, y_1 * y_2 * ... * y_n] + # tensor index = \sum_{i=1}^m (ind_i * \prod_{j=i+1}^m (x_j)) + # After gather, reshape and transpose back. + adv_idx_indices = [ + i for i, idx in enumerate(indices) if not symbolic_helper._is_none(idx) + ] + + if len(adv_idx_indices) == 0: + return self + elif len(adv_idx_indices) == 1: + return index_select( + g, self, adv_idx_indices[0], indices[adv_idx_indices[0]] + ) + else: + rank = symbolic_helper._get_tensor_rank(self) + if rank is None: + return symbolic_helper._unimplemented( + "aten::index", + "operator of advanced indexing on tensor of unknown rank. ", + self, + ) + # TODO: If indexing is supported natively in ONNX in future opsets, + # update the warning to recommend exporting with higher opset version. + warnings.warn( + "Exporting aten::index operator of advanced indexing in opset " + f"{GLOBALS.export_onnx_opset_version}" + " is achieved by combination of multiple ONNX operators, " + "including Reshape, Transpose, Concat, and Gather. " + "If indices include negative values, the exported graph will produce incorrect results.", + stacklevel=2, + ) + adv_idx_count = len(adv_idx_indices) + shape_tensor = _shape_as_tensor(g, self) + dim_tensor_list = [ + g.op( + "Gather", + shape_tensor, + g.op("Constant", value_t=torch.LongTensor([dim])), + axis_i=0, + ) + for dim in range(rank) + ] + + self = g.op( + "Transpose", + self, + perm_i=adv_idx_indices + + [i for i in range(rank) if i not in adv_idx_indices], + ) + self = g.op("Flatten", self, axis_i=adv_idx_count) + + # Note that tensor indices will be broadcasted while accumulating. Thus we get the final subarray shape as well. + cum_adv_index = indices[adv_idx_indices[-1]] + multiplier = dim_tensor_list[adv_idx_indices[-1]] + for i in range(adv_idx_count - 2, -1, -1): + adv_index = g.op("Mul", indices[adv_idx_indices[i]], multiplier) + cum_adv_index = g.op("Add", cum_adv_index, adv_index) + multiplier = g.op( + "Mul", multiplier, dim_tensor_list[adv_idx_indices[i]] + ) + + # perform gather + self = index_select(g, self, 0, cum_adv_index) + + cum_adv_index_shape_tensor = _shape_as_tensor(g, cum_adv_index) + # check if all advanced indices are consecutive. + # Refer to https://numpy.org/doc/stable/user/basics.indexing.html#combining-advanced-and-basic-indexing + # to understand how the subarray position is decided. + if adv_idx_indices == list( + range(adv_idx_indices[0], adv_idx_indices[-1] + 1) + ): + # unfold regular index axes + folded_adv_idx_shape_list = [ + g.op("Constant", value_t=torch.LongTensor([-1])) + ] + [ + dim_tensor_list[i] for i in range(rank) if i not in adv_idx_indices + ] + folded_adv_idx_shape = g.op( + "Concat", *folded_adv_idx_shape_list, axis_i=0 + ) + self = symbolic_helper._reshape_helper(g, self, folded_adv_idx_shape) + + # Transpose folded advanced indexed axis to its original location. + adv_idx_permute = ( + list(range(1, adv_idx_indices[0] + 1)) + + [0] + + list(range(adv_idx_indices[0] + 1, rank - adv_idx_count + 1)) + ) + self = g.op("Transpose", self, perm_i=adv_idx_permute) + + # unfold advanced index axes + final_shape_list = ( + [dim_tensor_list[i] for i in range(adv_idx_indices[0])] + + [cum_adv_index_shape_tensor] + + [ + dim_tensor_list[i] + for i in range(adv_idx_indices[0], rank) + if i not in adv_idx_indices + ] + ) + final_shape = g.op("Concat", *final_shape_list, axis_i=0) + else: + final_shape = g.op( + "Concat", + cum_adv_index_shape_tensor, + *[ + dim_tensor_list[i] + for i in range(rank) + if i not in adv_idx_indices + ], + axis_i=0, + ) + + return symbolic_helper._reshape_helper(g, self, final_shape) + + +@_onnx_symbolic("aten::linalg_norm") +@symbolic_helper.parse_args("v", "v", "is", "b", "v") +def linalg_norm( + g: jit_utils.GraphContext, + self: torch._C.Value, + ord: torch._C.Value, + dim: Sequence[int] | None, + keepdim: bool, + dtype: torch._C.Value, +): + # Conditions based on https://pytorch.org/docs/stable/generated/torch.linalg.norm.html + ord_value = None + if dim is None: + if symbolic_helper._is_none(ord): + self = symbolic_helper._reshape_helper(g, self, [-1]) + ord = g.op("Constant", value_t=torch.LongTensor([2])) + self_dim = symbolic_helper._get_tensor_rank(self) + if self_dim is None: + return symbolic_helper._unimplemented( + "dim", "Input rank must be known at export time.", self + ) + if self_dim == 1: + ord_value = symbolic_helper._parse_arg(ord, "f") + else: + dim = [0, 1] + else: + if len(dim) == 1: + if symbolic_helper._is_none(ord): + ord = g.op("Constant", value_t=torch.LongTensor([2])) + ord_value = symbolic_helper._parse_arg(ord, "f") + if ord_value: + return linalg_vector_norm(g, self, ord_value, dim, keepdim, dtype) + return linalg_matrix_norm(g, self, ord, dim, keepdim, dtype) + + +@_onnx_symbolic("aten::linalg_vector_norm") +@symbolic_helper.parse_args("v", "f", "is", "b", "v") +def linalg_vector_norm( + g: jit_utils.GraphContext, + self: torch._C.Value, + ord: float, + dim: Sequence[int] | None, + keepdim: bool, + dtype: torch._C.Value, +): + return symbolic_helper._linalg_vector_norm_helper(g, self, ord, dim, keepdim, dtype) + + +@_onnx_symbolic("aten::linalg_matrix_norm") +@symbolic_helper.parse_args("v", "v", "is", "b", "v") +def linalg_matrix_norm( + g: jit_utils.GraphContext, + self: torch._C.Value, + ord: torch._C.Value, + dim: list[int], + keepdim: bool, + dtype: torch._C.Value, +): + # Conditions based on https://pytorch.org/docs/stable/generated/torch.linalg.matrix_norm.html + ord_value = symbolic_helper._parse_arg(ord, "s") + if ord_value == "fro": + return frobenius_norm(g, self, dim, keepdim) + elif ord_value == "nuc": + return symbolic_helper._unimplemented("linalg.matrix_norm", "ord==nuc", self) + else: + ord_value = symbolic_helper._parse_arg(ord, "f") + if ord_value is None: + return frobenius_norm(g, self, dim, keepdim) + if ord_value == 2 or ord_value == -2: + # ord = 2/-2 unimplemented due to lack of operators + # used to calculate singular values + return symbolic_helper._unimplemented("linalg.matrix_norm", "ord==2", self) + # Wrap the dim vector to handle negative dim values + self_dim = symbolic_helper._get_tensor_rank(self) + if self_dim is None: + return symbolic_helper._unimplemented( + "linalg.matrix_norm", "Input rank must be known at export time.", self + ) + # Common implementation for cases with + # ord = 1/-1 and ord = inf/-inf + if dim[0] < 0: + dim[0] += self_dim + if dim[1] < 0: + dim[1] += self_dim + + if ord_value == math.inf or ord_value == -math.inf: + dim[0], dim[1] = dim[1], dim[0] + if dim[1] > dim[0] and not keepdim: + dim[1] -= 1 + sum = symbolic_helper._reducesum_helper( + g, g.op("Abs", self), axes_i=[dim[0]], keepdims_i=keepdim + ) + if ord_value > 0: + # pyrefly: ignore [no-matching-overload] + result, _indices = max( + g, + sum, + dim_or_y=g.op("Constant", value_t=torch.LongTensor([dim[1]])), + keepdim=keepdim, + ) + else: + # pyrefly: ignore [no-matching-overload] + result, _indices = min( + g, + sum, + dim_or_y=g.op("Constant", value_t=torch.LongTensor([dim[1]])), + keepdim=keepdim, + ) + return result + + +@_onnx_symbolic("aten::linalg_cross") +@symbolic_helper.parse_args("v", "v", "i") +def linalg_cross(g: jit_utils.GraphContext, input, other, dim=-1): + return cross(g, input, other, dim) + + +@_onnx_symbolic("aten::frobenius_norm") +@symbolic_helper.parse_args("v", "is", "b") +def frobenius_norm(g: jit_utils.GraphContext, self, dim=None, keepdim=False): + sqr = g.op("Mul", self, self) + sumsqr = symbolic_helper._reducesum_helper(g, sqr, axes_i=dim, keepdims_i=keepdim) + return g.op("Sqrt", sumsqr) + + +@_onnx_symbolic("aten::multinomial") +@symbolic_helper.parse_args("v", "i", "b", "v") +def multinomial( + g: jit_utils.GraphContext, input, num_samples, replacement=False, generator=None +): + if generator is not None and not symbolic_helper._is_none(generator): + symbolic_helper._unimplemented( + "Multinomial", "generator is not supported for multinomial", input + ) + if not replacement and num_samples > 1: + symbolic_helper._unimplemented( + "Multinomial", + "replacement=False when num_samples > 1 is not supported for multinomial", + input, + ) + + log_input = log(g, input) + return g.op( + "Multinomial", + log_input, + dtype_i=_C_onnx.TensorProtoDataType.INT64, + sample_size_i=num_samples, + ) + + +@_onnx_symbolic("aten::baddbmm") +def baddbmm(g: jit_utils.GraphContext, self, batch1, batch2, beta, alpha): + scalar_type = _type_utils.JitScalarType.from_value(self) + batch_mul = matmul(g, batch1, batch2) + mul_a = mul( + g, + batch_mul, + g.op("Cast", alpha, to_i=scalar_type.onnx_type()), + ) + mul_b = mul( + g, + self, + g.op("Cast", beta, to_i=scalar_type.onnx_type()), + ) + return add(g, mul_a, mul_b) + + +@_onnx_symbolic("aten::meshgrid") +@symbolic_helper.parse_args("v", "s") +def meshgrid(g: jit_utils.GraphContext, tensor_list, indexing: str | None = None): + if indexing is None: + indexing = "ij" + elif indexing not in {"ij", "xy"}: + raise errors.SymbolicValueError( + f"Unsupported indexing: {indexing}", tensor_list + ) + unpacked_tensor_list = symbolic_helper._unpack_list(tensor_list) + if indexing == "xy": + unpacked_tensor_list[:2] = unpacked_tensor_list[1::-1] + tensors = [ + symbolic_helper._reshape_helper( + g, t, g.op("Constant", value_t=torch.LongTensor([-1])) + ) + for t in unpacked_tensor_list + ] + tensors_shape = [g.op("Shape", t) for t in tensors] + out_shape = g.op("Concat", *tensors_shape, axis_i=0) + out = [] + for i, t in enumerate(tensors): + shape_i = [g.op("Constant", value_t=torch.ones(1, dtype=torch.int64))] * len( + tensors + ) + shape_i[i] = tensors_shape[i] + t_reshaped = _reshape_from_tensor(g, t, g.op("Concat", *shape_i, axis_i=0)) + out.append(g.op("Expand", t_reshaped, out_shape)) + if indexing == "xy": + out[0], out[1] = out[1], out[0] + return g.op("prim::ListConstruct", *out) + + +@_onnx_symbolic("aten::remainder") +def remainder(g: jit_utils.GraphContext, input, other): + div = _floor_divide(g, input, other) + quo = g.op("Mul", div, other) + return g.op("Sub", input, quo) + + +@_onnx_symbolic("aten::gelu") +@symbolic_helper.parse_args("v", "s") +def gelu(g: jit_utils.GraphContext, self: torch._C.Value, approximate: str = "none"): + if approximate == "tanh": + kBeta = math.sqrt(2 / math.pi) + kKappa = 0.044715 + + beta = torch.tensor(kBeta, dtype=torch.double) + kappa = torch.tensor(kKappa, dtype=torch.double) + one = torch.tensor(1.0, dtype=torch.double) + half = torch.tensor(0.5, dtype=torch.double) + + self_cube = mul(g, self, mul(g, self, self)) + inner = mul(g, beta, add(g, self, mul(g, kappa, self_cube))) + return mul(g, half, mul(g, self, add(g, one, g.op("Tanh", inner)))) + else: + _sqrt2 = 1.4142135623730951 + erf = g.op("Erf", g.op("Div", self, torch.tensor(_sqrt2, dtype=torch.double))) + erf_plusone = add( + g, erf, g.op("Constant", value_t=torch.tensor(1, dtype=torch.double)) + ) + return mul( + g, + mul(g, self, erf_plusone), + g.op("Constant", value_t=torch.tensor(0.5, dtype=torch.double)), + ) + + +@_onnx_symbolic("aten::group_norm") +@symbolic_helper.quantized_args(True, False, False, False) +@symbolic_helper.parse_args("v", "i", "v", "v", "f", "i") +def group_norm( + g: jit_utils.GraphContext, input, num_groups, weight, bias, eps, cudnn_enabled +): + channel_size = symbolic_helper._get_tensor_dim_size(input, 1) + if channel_size is not None: + if channel_size % num_groups != 0: + raise AssertionError( + f"channel_size ({channel_size}) must be divisible by num_groups ({num_groups})" + ) + input_rank = symbolic_helper._get_tensor_rank(input) + if input_rank is None: + return symbolic_helper._unimplemented("group_norm", "unknown input rank", input) + # 0 in the shape list keeps dimension value unchanged. + shape = [0, num_groups, -1] + input_reshaped = symbolic_helper._reshape_helper( + g, input, g.op("Constant", value_t=torch.LongTensor(shape)) + ) + + # C is always divisible by num_groups + # Due to shape difference. we need to apply weight and bias after + # instance norm computation and reshape + weight_ = g.op( + "Constant", + value_t=torch.tensor( + [1.0] * num_groups, + dtype=_type_utils.JitScalarType.from_value(input).dtype(), + ), + ) + bias_ = g.op( + "Constant", + value_t=torch.tensor( + [0.0] * num_groups, + dtype=_type_utils.JitScalarType.from_value(input).dtype(), + ), + ) + + norm_reshaped = g.op( + "InstanceNormalization", input_reshaped, weight_, bias_, epsilon_f=eps + ) + norm = symbolic_helper._reshape_helper(g, norm_reshaped, g.op("Shape", input)) + + if weight is None or weight.node().mustBeNone(): + weight_value = torch.tensor( + [1.0], dtype=_type_utils.JitScalarType.from_value(input).dtype() + ) + weight = g.op("Constant", value_t=weight_value) + if bias is None or bias.node().mustBeNone(): + bias_value = torch.tensor( + [0.0], dtype=_type_utils.JitScalarType.from_value(input).dtype() + ) + bias = g.op("Constant", value_t=bias_value) + + # Norm has shape [N, C, *] so we reshape weight and bias to [C, *] + axes = list(range(1, input_rank - 1)) + return add( + g, + mul(g, norm, symbolic_helper._unsqueeze_helper(g, weight, axes)), + symbolic_helper._unsqueeze_helper(g, bias, axes), + ) + + +@_onnx_symbolic("aten::_weight_norm") +@symbolic_helper.parse_args("v", "v", "i") +def _weight_norm(g: jit_utils.GraphContext, weight_v, weight_g, dim): + rank = symbolic_helper._get_tensor_rank(weight_v) + if rank is not None: + # W = g * ((v) / ||v||) + # Compute norm_except_dim for l2 norm. dim = None means over all dims + # torch's weight_norm module sets dim = -1 if it's None. + # This conflicts the logic for negative axes to access dims backwards + # TODO: Might need a fix in torch group_norm module + axes = list(range(rank)) + if dim is not None: + if dim < -1: + dim += rank + if dim != -1: + axes.remove(dim) + norm_v = norm(g, weight_v, 2, axes, 1) + div = g.op("Div", weight_v, norm_v) + return g.op("Mul", div, weight_g) + raise errors.SymbolicValueError( + "Unsupported: ONNX export of _weight_norm for tensor of unknown rank.", + weight_v, + ) + + +@_onnx_symbolic("aten::dim") +def dim(g: jit_utils.GraphContext, self): + """Implement the dim functionality available for a pytorch tensor in ONNX""" + # ONNX does not support dim directly in this opset so we can use 2 ops to get the info + shape = g.op("Shape", self) + return g.op("Size", shape) + + +@_onnx_symbolic("aten::__contains_") +def __contains_(g: jit_utils.GraphContext, self, element): + unpacked_list = symbolic_helper._unpack_list(self) + if all( + symbolic_helper._is_constant(x) for x in unpacked_list + ) and symbolic_helper._is_constant(element): + return g.op( + "Constant", + value_t=torch.tensor( + symbolic_helper._node_get(element.node(), "value") + in (symbolic_helper._node_get(x.node(), "value") for x in unpacked_list) + ), + ) + + raise errors.SymbolicValueError( + "Unsupported: ONNX export of __contains__ for non-constant list or element.", + self, + ) + + +@_onnx_symbolic("aten::__getitem_") +def __getitem_(g: jit_utils.GraphContext, self, i): + return select(g, self, g.op("Constant", value_t=torch.tensor([0])), i) + + +@_onnx_symbolic("aten::item") +def item(g: jit_utils.GraphContext, self): + return self + + +@_onnx_symbolic("aten::take") +def take(g: jit_utils.GraphContext, self, index): + self_flattened = symbolic_helper._reshape_helper( + g, self, g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)) + ) + out = index_select(g, self_flattened, 0, index) + out = reshape_as(g, out, index) + return out + + +def _kl_div_log_target_impl(g: jit_utils.GraphContext, input, target): + diff_ = sub(g, target, input) + exp_ = exp(g, target) + output = mul(g, exp_, diff_) + return output + + +def _kl_div_non_log_target_impl(g: jit_utils.GraphContext, input, target): + log_ = log(g, target) + diff_ = sub(g, log_, input) + output_pos = mul(g, target, diff_) + zeros_ = zeros_like(g, output_pos) + mask_ = gt(g, target, g.op("Constant", value_t=torch.tensor(0))) + output = where(g, mask_, output_pos, zeros_) + return output + + +@_onnx_symbolic("aten::kl_div") +@symbolic_helper.parse_args("v", "v", "i", "b") +def kl_div(g: jit_utils.GraphContext, input, target, reduction, log_target): + if log_target: + output = _kl_div_log_target_impl(g, input, target) + else: + output = _kl_div_non_log_target_impl(g, input, target) + + if reduction == 0: + return output + elif reduction == 1: + return g.op("ReduceMean", output, keepdims_i=0) + elif reduction == 2: + return symbolic_helper._reducesum_helper(g, output, keepdims_i=0) + else: + return symbolic_helper._onnx_unsupported( + "kl_div with reduction other than none, mean, or sum.", input + ) + + +@_onnx_symbolic("aten::mse_loss") +@symbolic_helper.parse_args("v", "v", "i") +def mse_loss(g: jit_utils.GraphContext, input, target, reduction): + output = mul(g, sub(g, input, target), sub(g, input, target)) + if reduction == 0: + return output + elif reduction == 1: + return g.op("ReduceMean", output, keepdims_i=0) + elif reduction == 2: + return symbolic_helper._reducesum_helper(g, output, keepdims_i=0) + else: + return symbolic_helper._onnx_unsupported( + "mse_loss with reduction other than none, mean, or sum.", input + ) + + +@_onnx_symbolic("aten::as_strided") +@symbolic_helper.quantized_args(True) +@symbolic_helper.parse_args("v", "v", "is", "i") +def as_strided(g: jit_utils.GraphContext, self, sizes, strides, offset=None): + sizes = symbolic_helper._maybe_get_const(sizes, "is") + rank = len(strides) + self_1d = symbolic_helper._reshape_helper( + g, self, g.op("Constant", value_t=torch.tensor([-1], dtype=torch.int64)) + ) + ind: torch.Tensor | None + if not symbolic_helper._is_value(sizes): + ind = torch.tensor([0], dtype=torch.long) + for i, (size, stride) in enumerate(zip(sizes, strides)): + r_size = [1] * rank + r_size[i] = -1 + ind = ind + torch.arange(size).view(r_size) * stride + if offset: + ind = ind + offset + return g.op("Gather", self_1d, g.op("Constant", value_t=ind)) + else: + ind = None + for i, stride in enumerate(strides): + r_size = [1] * rank + r_size[i] = -1 + size = select( + g, + sizes, + g.op("Constant", value_t=torch.tensor([0])), + g.op("Constant", value_t=torch.tensor(i)), + ) + tmp_ind = symbolic_helper._reshape_helper( + g, + arange(g, size, 4, None, None, None), + g.op("Constant", value_t=torch.tensor(r_size)), + ) + tmp_ind = g.op( + "Mul", tmp_ind, g.op("Constant", value_t=torch.tensor([stride])) + ) + if ind is None: + ind = tmp_ind + else: + ind = g.op("Add", ind, tmp_ind) + if offset: + # pyrefly: ignore [bad-argument-type] + ind = g.op("Add", ind, g.op("Constant", torch.tensor([offset]))) + # pyrefly: ignore [bad-argument-type] + return g.op("Gather", self_1d, ind) + + +@_onnx_symbolic("aten::__derive_index") +def __derive_index(g: jit_utils.GraphContext, index, start, step): + return g.op("Add", start, g.op("Mul", index, step)) + + +@_onnx_symbolic("aten::__range_length") +# Source code for aten op can be found here: pytorch/torch/csrc/jit/runtime/register_prim_ops.cpp +# if (step > 0 && lo < hi) { +# push(stack, 1 + (hi - 1 - lo) / step); +# } else if (step < 0 && lo > hi) { +# push(stack, 1 + (lo - 1 - hi) / (0 - step)); +# } else { +# push(stack, 0); +# } +def __range_length(g: jit_utils.GraphContext, lo, hi, step): + sub = g.op("Sub", hi, lo) + div = g.op("Ceil", true_divide(g, sub, step)) + return g.op("Cast", div, to_i=_C_onnx.TensorProtoDataType.INT64) + + +@_onnx_symbolic("aten::linear") +def linear(g: jit_utils.GraphContext, input, weight, bias): + rank = symbolic_helper._get_tensor_rank(input) + weight = t(g, weight) + if rank == 2 and not bias.node().mustBeNone(): + alpha = g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64)) + beta = g.op("Constant", value_t=torch.tensor(1, dtype=torch.int64)) + output = addmm(g, bias, input, weight, alpha, beta) + else: + output = matmul(g, input, weight) + if not bias.node().mustBeNone(): + output = add(g, bias, output) + + return output + + +@_onnx_symbolic("aten::hann_window") +@symbolic_helper.parse_args("v", "b", "i", "v", "v", "v", "v") +def hann_window( + g: jit_utils.GraphContext, + window_length, + periodic=True, + dtype: int | None = None, + layout=None, + device=None, + pin_memory=None, + requires_grad=False, +): + if dtype is None: + dtype_ = torch.get_default_dtype() + if not dtype_ or not dtype_.is_floating_point: + dtype_ = torch.float + scalar_type = _type_utils.JitScalarType.from_dtype(dtype_) + else: + scalar_type = _type_utils.JitScalarType(dtype) + + n_array = arange(g, window_length, 4, None, None, None) + output = g.op("Cast", n_array, to_i=_C_onnx.TensorProtoDataType.FLOAT) + output = mul( + g, g.op("Constant", value_t=torch.tensor(math.pi, dtype=torch.float)), output + ) + + if periodic is False: + window_length = sub( + g, window_length, g.op("Constant", value_t=torch.tensor(1, dtype=torch.int)) + ) + output = div(g, output, window_length) + output = g.op( + "Cast", + square(g, sin(g, output)), + to_i=scalar_type.onnx_type(), + ) + + return output + + +@_onnx_symbolic("aten::mv") +def mv(g: jit_utils.GraphContext, self, vec): + return matmul(g, self, vec) + + +@_onnx_symbolic("aten::dot") +def dot(g: jit_utils.GraphContext, self, other): + return matmul(g, self, other) + + +@_onnx_symbolic("aten::movedim") +@symbolic_helper.parse_args("v", "t", "t") +def movedim(g: jit_utils.GraphContext, self, source, destination): + # This is a pythonic implementation mostly taken from aten/src/ATen/native/TensorShape.cpp::movedim + source = source.view(-1) + destination = destination.view(-1) + + if source.size() != destination.size(): + raise AssertionError( + f"source.size()={source.size()} != destination.size()={destination.size()}" + ) + + if (source == destination).all(): + return self + + self_rank = symbolic_helper._get_tensor_rank(self) + if self_rank is None: + raise AssertionError("self_rank must be non-None") + + perm = list(range(self_rank)) + + src_dims = perm.copy() + dst_dims = perm.copy() + + for src, dst in zip(source.tolist(), destination.tolist()): + perm[dst] = src + src_dims[src] = -1 + dst_dims[dst] = -1 + + src_dims = [dim for dim in src_dims if dim != -1] + dst_dims = [dim for dim in dst_dims if dim != -1] + + for src, dst in zip(src_dims, dst_dims): + perm[dst] = src + + return g.op("Transpose", self, perm_i=perm) + + +@_onnx_symbolic("aten::fill") +@symbolic_helper.parse_args("v", "v") +def fill(g: jit_utils.GraphContext, self, value): + scalar_type = _type_utils.JitScalarType.from_value( + self, _type_utils.JitScalarType.FLOAT + ) + return full_like(g, self, value, scalar_type) + + +@_onnx_symbolic("aten::index_add") +def index_add(g: jit_utils.GraphContext, self, dim, index, other, alpha=None): + warnings.warn( + "Warning: ONNX export does not support duplicated values in 'index' field, " + + "this will cause the ONNX model to be incorrect.", + stacklevel=2, + ) + + # ONNX does not support "alpha" argument, unlike aten index_add + # See: https://github.com/pytorch/pytorch/pull/65993#issuecomment-953151102 for more context + if alpha and symbolic_helper._scalar(symbolic_helper._maybe_get_scalar(alpha)) != 1: + return symbolic_helper._unimplemented("index_add", "alpha != 1", self) + + dim = symbolic_helper._maybe_get_const(dim, "i") + if dim is None: + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting 'index_add_()' function with " + "unknown 'dim' value.", + self, + ) + + self_dim_rank = symbolic_helper._get_tensor_rank(self) + other_dim_rank = symbolic_helper._get_tensor_rank(other) + + if self_dim_rank is None or other_dim_rank is None: + raise errors.SymbolicValueError( + "ONNX export does NOT support exporting 'index_add_()' function while " + "the rank of self tensor or tensor to be added is unknown.", + self, + ) + + if other_dim_rank != self_dim_rank: + delta = self_dim_rank - other_dim_rank + for _ in range(delta): + other = symbolic_helper._unsqueeze_helper( + g, other, [symbolic_helper._get_tensor_rank(other)] + ) + + other_dim_size = symbolic_helper._get_tensor_dim_size(other, dim) + self_dim_size = symbolic_helper._get_tensor_dim_size(self, dim) + + if (other_dim_size is not None) and (self_dim_size is not None): + if other_dim_size > self_dim_size: + raise errors.SymbolicValueError( + "ONNX export does not support exporting 'index_add_()' function with " + "duplicated values in 'index' parameter yet.", + self, + ) + + # Construct a new shape. It's almost as same as self except the size of the 'dim' + # dimension is 1, so that we can expand other dimensions as expected. + new_shape_axes = list(range(self_dim_rank)) + new_shape_starts = [0 for i in range(self_dim_rank)] + new_shape_ends = [sys.maxsize if (i != dim) else 1 for i in range(self_dim_rank)] + + new_shape = symbolic_helper._slice_helper( + g, self, axes=new_shape_axes, starts=new_shape_starts, ends=new_shape_ends + ) + other = expand_as(g, other, new_shape) + + for _ in range(dim): + index = symbolic_helper._unsqueeze_helper(g, index, [0]) + + for _ in range(self_dim_rank - dim - 1): + index = symbolic_helper._unsqueeze_helper( + g, index, [symbolic_helper._get_tensor_rank(index)] + ) + + return scatter_add(g, self, dim, expand_as(g, index, other), other) + + +@_onnx_symbolic("aten::roll") +@symbolic_helper.parse_args("v", "is", "is") +def roll(g: jit_utils.GraphContext, self, shifts, dims): + if len(shifts) != len(dims): + raise AssertionError(f"len(shifts)={len(shifts)} != len(dims)={len(dims)}") + + result = self + for i in range(len(shifts)): + shapes = [] + shape = symbolic_helper._slice_helper( + g, result, axes=[dims[i]], starts=[-shifts[i]], ends=[sys.maxsize] + ) + shapes.append(shape) + shape = symbolic_helper._slice_helper( + g, result, axes=[dims[i]], starts=[0], ends=[-shifts[i]] + ) + shapes.append(shape) + result = g.op("Concat", *shapes, axis_i=dims[i]) + + return result + + +@_onnx_symbolic("aten::cross") +@symbolic_helper.parse_args("v", "v", "i") +def cross(g: jit_utils.GraphContext, input, other, dim=None): + dim = symbolic_helper._get_dim_for_cross(input, dim) + # If we have two tensors such that + # A = [a, b, c], B = [d, e, f], we permute the tensor such that we have + # After first roll, + # A' = [b, c, a], B' = [f, d, e], so that we calculate (b*f, c*d, a*e) + roll_x_1 = roll(g, input, [2], [dim]) + roll_y_1 = roll(g, other, [1], [dim]) + # After second roll, + # A' = [c, a, b], B' = [e, f, d], so that we calculate (c*e, a*f, b*d) + roll_x_2 = roll(g, input, [1], [dim]) + roll_y_2 = roll(g, other, [2], [dim]) + # cross product is calculated as + # result = [(b*f - c*e), (c*d - a*f), (a*e - b*d)] + return sub(g, mul(g, roll_x_1, roll_y_1), mul(g, roll_x_2, roll_y_2)) + + +@_onnx_symbolic("aten::cdist") +def cdist( + g: jit_utils.GraphContext, + x1, + x2, + p=2.0, + compute_mode="use_mm_for_euclid_dist_if_necessary", +): + # X1.shape = (B * P * D), X2.shape = (B * R * D) + # In order to respect numpy style broadcasting as demonstrated in + # https://github.com/onnx/onnx/blob/main/docs/Broadcasting.md + # we unsqueeze both input tensors + row_size_x1 = symbolic_helper._get_tensor_dim_size(x1, -2) + row_size_x2 = symbolic_helper._get_tensor_dim_size(x2, -2) + p_float = symbolic_helper._parse_arg(p, "f") + compute_mode = symbolic_helper._parse_arg(compute_mode, "i") + if p_float == 2.0 and ( + compute_mode == 1 + or ( + compute_mode is None + and ( + row_size_x1 is None + or row_size_x2 is None + or (row_size_x1 >= 25 and row_size_x2 >= 25) + ) + ) + ): + return _euclidean_dist(g, x1, x2) + rank = symbolic_helper._get_tensor_rank(x1) + if rank is None: + raise AssertionError("rank must be non-None") + broadcasted_x1 = symbolic_helper._unsqueeze_helper(g, x1, [rank - 1]) + broadcasted_x2 = symbolic_helper._unsqueeze_helper(g, x2, [rank - 2]) + return pairwise_distance( + g, broadcasted_x1, broadcasted_x2, p, eps=1e-06, keepdim=False + ) + + +def _euclidean_dist(g: jit_utils.GraphContext, x1, x2): + # X1.shape = (B * P * D), X2.shape = (B * R * D) + # using matrix multiplication to accelerate the calculation of + # the euclidean distance + rank = symbolic_helper._get_tensor_rank(x1) + if rank is None: + raise AssertionError("rank must be non-None") + x1_norm = symbolic_helper._reducesum_helper( + g, + # pyrefly: ignore [no-matching-overload] + pow(g, x1, symbolic_helper._generate_wrapped_number(g, 2.0)), + axes_i=[-1], + keepdims_i=True, + ) + x1_pad = ones_like(g, x1_norm) + x2_norm = symbolic_helper._reducesum_helper( + g, + # pyrefly: ignore [no-matching-overload] + pow(g, x2, symbolic_helper._generate_wrapped_number(g, 2.0)), + axes_i=[-1], + keepdims_i=True, + ) + x2_pad = ones_like(g, x2_norm) + x1_ = g.op( + "Concat", + *[ + mul(g, symbolic_helper._generate_wrapped_number(g, -2.0), x1), + x1_norm, + x1_pad, + ], + axis_i=-1, + ) + x2_ = g.op("Concat", *[x2, x2_pad, x2_norm], axis_i=-1) + result = matmul(g, x1_, transpose(g, x2_, -2, -1)) + dtype = _type_utils.JitScalarType.from_value(result) + min = g.op( + "Cast", symbolic_helper._generate_wrapped_number(g, 0.0), to_i=dtype.onnx_type() + ) + result = symbolic_helper._op_with_optional_float_cast( + g, "Max", result, min, opset_before=12 + ) + result = sqrt(g, result) + return result + + +@_onnx_symbolic("aten::lerp") +def lerp(g: jit_utils.GraphContext, self, end, weight): + # Conditional for better numeric. This has been discussed in + # https://github.com/pytorch/pytorch/pull/18871 + diff = g.op("Sub", end, self) + return where( + g, + g.op("Less", weight, g.op("Constant", value_t=torch.tensor(0.5))), + g.op("Add", self, g.op("Mul", weight, diff)), + g.op( + "Sub", + end, + g.op( + "Mul", + diff, + g.op("Sub", g.op("Constant", value_t=torch.tensor(1.0)), weight), + ), + ), + ) + + +@_onnx_symbolic("aten::broadcast_tensors") +def broadcast_tensors(g: jit_utils.GraphContext, self): + all_tensors = symbolic_helper._unpack_list(self) + t_with_final_shape = zeros_like(g, all_tensors[0]) + + # Add operator supports multidirectional broadcasting. So we leverage this function + # to infer the final shape generated by the broadcast. + for t in all_tensors: + t_with_final_shape = add(g, t_with_final_shape, t) + + t_list = [expand_as(g, t, t_with_final_shape) for t in all_tensors] + return g.op("prim::ListConstruct", *t_list) + + +@_onnx_symbolic("aten::is_pinned") +def is_pinned(g: jit_utils.GraphContext, self, device=None) -> None: + # Unused by ONNX. + return None + + +@_onnx_symbolic("prim::ConstantSplit") +def prim_constant_split(g: jit_utils.GraphContext, self, split_size, dim): + size = symbolic_helper._get_tensor_dim_size(self, dim) + if size is None: + return symbolic_helper._unimplemented( + "prim::ConstantSplit", "unknown dimension size", self + ) + splits = [split_size] * (size // split_size) + leftover = size % split_size + if leftover: + splits.append(leftover) + return g.op("Split", self, split_i=splits, axis_i=dim, outputs=len(splits)) + + +# TODO: It would be better to export this as a chunk directly, as this is +# less sensitive to changes in input size. +# TODO: Once we have proper scoping, stop reimplementing chunk, delete this +# method, and use the desugared version +@_onnx_symbolic("prim::ConstantChunk") +def prim_constant_chunk(g: jit_utils.GraphContext, self, chunks, dim): + dim_size = symbolic_helper._get_tensor_dim_size(self, dim) + if dim_size is None: + return symbolic_helper._unimplemented( + "prim::ConstantChunk", "unknown dimension size", self + ) + split_size = (dim_size + chunks - 1) // chunks + return prim_constant_split(g, self, split_size, dim) + + +@_onnx_symbolic("prim::shape") +def prim_shape(g: jit_utils.GraphContext, self): + return g.op("Shape", self) + + +@_onnx_symbolic("prim::max") +def prim_max(g: jit_utils.GraphContext, self, other): + return symbolic_helper._op_with_optional_float_cast( + g, "Max", self, other, opset_before=12 + ) + + +@_onnx_symbolic("prim::min") +def prim_min(g: jit_utils.GraphContext, self, other=None): + if not other: + if symbolic_helper._is_packed_list(self): + self = stack(g, self, g.op("Constant", value_t=torch.tensor([0]))) + # pyrefly: ignore [no-matching-overload] + return min(g, self) + # pyrefly: ignore [no-matching-overload] + return min(g, self, other) + + +@_onnx_symbolic("prim::data") +def prim_data(g: jit_utils.GraphContext, self): + return self + + +@_onnx_symbolic("prim::layout") +def prim_layout(g: jit_utils.GraphContext, self): + # Always return 'torch.strided'. Other layout types are not supported by JIT 'TensorType'. + # Layout class defined in 'c10/core/Layout.h'. + return g.op("Constant", value_t=torch.tensor(0)) + + +@_onnx_symbolic("prim::ListConstruct") +def prim_list_construct(g: jit_utils.GraphContext, *inputs, **kwargs) -> None: + return None + + +@_onnx_symbolic("prim::ListUnpack") +def prim_list_unpack( + g: jit_utils.GraphContext, *inputs, **kwargs +) -> list[_C.Value] | None: + if len(inputs) == 1 and inputs[0].node().kind() == "prim::ListConstruct": + # Cancel the previous node if it is ListConstruct by returning its inputs + # TODO(justinchuby): Use a public method in the helper module + return symbolic_helper._unpack_list(inputs[0]) + + return None + + +@_onnx_symbolic("prim::TupleConstruct") +def prim_tuple_construct(g: jit_utils.GraphContext, *inputs, **kwargs) -> None: + return None + + +@_onnx_symbolic("prim::Uninitialized") +def prim_uninitialized(g: jit_utils.GraphContext, *inputs, **kwargs) -> None: + return None + + +# exists to refine the type of the Value +# if x is an optional Tensor, unchecked_cast will cast +# x to Tensor, so the rest of the graph knows that x is a Tensor +# this doesn't do anything in runtime and is a noop in ONNX +@_onnx_symbolic("prim::unchecked_cast") +def prim_unchecked_cast(g: jit_utils.GraphContext, self): + return self + + +@_onnx_symbolic("prim::dtype") +def prim_dtype(g: jit_utils.GraphContext, self): + scalar_type = symbolic_helper._try_get_scalar_type(self) + if scalar_type is None: + scalar_type = _type_utils.JitScalarType.FLOAT + # This node records a torch dtype as int + return g.op("Constant", value_t=torch.tensor(scalar_type)) + + +@_onnx_symbolic("prim::tolist") +def prim_tolist(g: jit_utils.GraphContext, input, dim_val, elem_ty_val): + """tolist is currently supported only for 1D input tensors. + + dim_val and elem_ty_val represent dimension and type annotations + that need to match dimension and type of the input tensor. + """ + dim = symbolic_helper._maybe_get_const(dim_val, "i") + if dim > 1: + return symbolic_helper._unimplemented("prim::tolist", "dim_val > 1", input) + return input + + +# ----------------------------------------------------------------------------- +# Symbolic functions that need extra context +# ----------------------------------------------------------------------------- +@_onnx_symbolic("prim::device") +def prim_device(g: jit_utils.GraphContext, *inputs, **kwargs) -> None: + output_type = g.original_node.output().type() + if isinstance(output_type, _C.DeviceObjType): + return None + + return symbolic_helper._unimplemented( + "prim::device", + f"output type should be 'DeviceObjType', not '{output_type.kind()}'", + g.original_node.output(), + ) + + +@_onnx_symbolic("prim::Loop") +def prim_loop(g: jit_utils.GraphContext, *inputs, **attrs) -> list[_C.Value]: + node = g.original_node + env = g.env + values_in_env = g.values_in_env + params_dict = g.params_dict + + operator_export_type = GLOBALS.operator_export_type + opset_version = GLOBALS.export_onnx_opset_version + + old_blocks = tuple(node.blocks()) + _new_op_outputs, new_block_contexts, new_node = jit_utils.add_op_with_blocks( + g, "Loop", *inputs, outputs=node.outputsSize(), n_blocks=len(old_blocks) + ) + + for old_block, new_block_context in zip(old_blocks, new_block_contexts): + # Copy input metadata to subblock + # + # prim::Loop(iter, cond, input_1, ..., input_n) + # block0(iter, input_1, ..., input_n) + # + # For `Loop` node, copy metadata for `iter`, `input_1`, ..., `input_n`. + for i, b_in in enumerate(old_block.inputs()): + if i == 0 and i < len(inputs): + b_in.setType(inputs[i].type()) + # For optional block inputs, they may switch between None not-None inside + # the loop body, so if the loop input is not optional, the block input may + # still need to be optional. + if ( + i > 0 + and (i + 1) < len(inputs) + and not isinstance(b_in.type(), _C.OptionalType) + ): + b_in.setType(inputs[i + 1].type()) + torch._C._jit_pass_onnx_block( + old_block, + new_block_context.block, + operator_export_type, + env, + values_in_env, + False, + ) + fixed_outputs = torch._C._jit_pass_fixup_onnx_controlflow_node( + new_node, opset_version + ) + # Run shape type inference for Loop after subblock is converted. + if GLOBALS.onnx_shape_inference: + torch._C._jit_pass_onnx_node_shape_type_inference( + new_node, params_dict, opset_version + ) + return fixed_outputs + + +@_onnx_symbolic("prim::If") +def prim_if(g: jit_utils.GraphContext, *inputs, **attrs) -> list[_C.Value]: + n = g.original_node + block = g.block + env = g.env + values_in_env = g.values_in_env + params_dict = g.params_dict + + operator_export_type = GLOBALS.operator_export_type + opset_version = GLOBALS.export_onnx_opset_version + + static_if = inputs[0].node().kind() == "onnx::Constant" + if static_if: + # Fold static if + # + # The torch IR + # graph(%embedding_matrix.1 : Float(10, 15, strides=[15, 1], requires_grad=0, device=cpu), + # %input.1 : Long(6, strides=[1], requires_grad=0, device=cpu), ... + # %65 : Bool(requires_grad=0, device=cpu) = prim::Constant[value={0}]() + # %21 : Long(device=cpu) = aten::eq(%20, %64) + # %22 : Long(device=cpu) = prim::If(%21) + # block0(): + # %23 : Long(device=cpu) = aten::is_floating_point(%input.1) + # -> (%23) + # block1(): + # -> (%65) + # %input.53 : Tensor, %weight : Tensor = prim::If(%22) + # block0(): + # -> (%embedding_matrix.1, %input.1) + # block1(): + # -> (%input.1, %embedding_matrix.1) + # %26 : int[] = aten::size(%input.53) + # + # The converted ONNX graph + # %10 : Bool(device=cpu) = onnx::Constant[value={0}]() + # %14 : Bool(device=cpu) = onnx::Equal(%13, %8) + # %15 : Bool(requires_grad=0, device=cpu) = onnx::Constant[value={0}]() + # %16 : Long(1, strides=[1], device=cpu) = onnx::Shape(%input.1) + input_flag = symbolic_helper._node_get(inputs[0].node(), "value").tolist() + const_value = ( + all(input_flag) if isinstance(input_flag, list) else bool(input_flag) + ) + block_idx = 0 if const_value else 1 + current_b = list(n.blocks())[block_idx] + env = torch._C._jit_pass_onnx_block( + current_b, + block, + operator_export_type, + env, + values_in_env, + True, + ) + if_output_list = list(n.outputs()) + current_b_list = list(current_b.outputs()) + + final_b_list = [] + for idx in range(len(if_output_list)): + if current_b_list[idx] not in env: + raise errors.SymbolicValueError( + f"The sub block ATen output {current_b_list[idx]} is not in env.", + current_b_list[idx], + ) # type:ignore[operator] + onnx_b = env[current_b_list[idx]] + final_b_list.append(onnx_b) + return final_b_list + else: + old_blocks = tuple(n.blocks()) + _new_op_outputs, new_block_contexts, new_node = jit_utils.add_op_with_blocks( + g, "If", *inputs, outputs=n.outputsSize(), n_blocks=len(old_blocks) + ) + + for old_block, new_block_context in zip(old_blocks, new_block_contexts): + torch._C._jit_pass_onnx_block( + old_block, + new_block_context.block, + operator_export_type, + env, + values_in_env, + False, + ) + fixed_outputs = torch._C._jit_pass_fixup_onnx_controlflow_node( + new_node, opset_version + ) + # Run shape type inference for If after subblock is converted. + if GLOBALS.onnx_shape_inference: + torch._C._jit_pass_onnx_node_shape_type_inference( + new_node, params_dict, opset_version + ) + return fixed_outputs + + +@_onnx_symbolic("prim::Constant") +def prim_constant(g: jit_utils.GraphContext, *inputs, **attrs): + node = g.original_node + + if node.mustBeNone(): + return None + # This must go before checking for string values, because some device constants + # have string values, but we want to keep them as unconverted Device types so + # that eq() can work on them. + if isinstance(node.output().type(), _C.DeviceObjType): + return None + if node.kindOf("value") == "t": + return g.op("Constant", value_t=symbolic_helper._node_get(node, "value")) + if node.kindOf("value") == "s": + return g.op("Constant", value_s=symbolic_helper._node_get(node, "value")) + if node.output().type().isSubtypeOf( + _C.ListType.ofInts() + ) or node.output().type().isSubtypeOf(_C.ListType.ofFloats()): + return g.op( + "Constant", value_t=torch.tensor(symbolic_helper._node_get(node, "value")) + ) + if node.output().type().isSubtypeOf(_C.ListType.ofStrings()): + str_constants = [ + g.op("Constant", value_s=s) + for s in symbolic_helper._node_get(node, "value") + ] + return g.op("prim::ListConstruct", *str_constants) + + raise errors.SymbolicValueError( + f"Unsupported prim::Constant kind: '{node.kindOf('value')}'. " + f"Please send a bug report at {_constants.PYTORCH_GITHUB_ISSUES_URL}.", + node.output(), + ) + + +@_onnx_symbolic("prim::type") +def prim_type(g: jit_utils.GraphContext, device_value: _C.Value, *args, **kwargs): + if device_value.node().kind() == "prim::device": + device = jit_utils.get_device_from_value(device_value.node().input()) + if device is not None: + return g.op("Constant", value_s=str(device)) + + return symbolic_helper._unimplemented( + "prim::type", + "Device type cannot be statically determined.", + device_value, + ) + + +@_onnx_symbolic("onnx::Placeholder") +def onnx_placeholder(g: jit_utils.GraphContext, *inputs, **attrs): + node = g.original_node + block = g.block + env = g.env + values_in_env = g.values_in_env + + return torch._C._jit_onnx_convert_pattern_from_subblock( + block, node, env, values_in_env + ) + + +@_onnx_symbolic("aten::resolve_conj") +@_onnx_symbolic("aten::resolve_neg") +def noop_complex_operators(g: jit_utils.GraphContext, input: _C.Value): + # ONNX does not have operators to *directly* manipulate real/imaginary components + # However, a few torch APIs (e.g. .tolist()) use complex operations when input is real, + # which results in failures due to missing operators for complex numbers + + # `aten::resolve_conj` and `aten::resolve_neg` can safely be implemented as no-op + return input + + +@_onnx_symbolic("aten::_conj") +@_onnx_symbolic("aten::conj_physical") +def unsupported_complex_operators(g: jit_utils.GraphContext, input: _C.Value): + # ONNX does not have operators to *directly* manipulate real/imaginary components + # However, a few torch APIs (e.g. .tolist()) use complex operations when input is real, + # which results in failures due to missing operators for complex numbers + + # While `aten::_conj` and `aten::conj_physical` raise exception when input is complex + if symbolic_helper.is_complex_value(input): + # FIXME(justinchuby): report correct name for symbolic being executed + return symbolic_helper._onnx_unsupported( + "aten::_conj, aten::conj_physical", + input, + ) + + # they can safely be implemented as no-op for real numbers only + return noop_complex_operators(g, input) + + +@_onnx_symbolic("aten::logit") +def logit(g: jit_utils.GraphContext, self: torch._C.Value, eps: torch._C.Value): + one = g.op("Constant", value_t=torch.tensor(1.0)) + + if not symbolic_helper._is_none(eps): + eps = g.op( + "Cast", eps, to_i=_type_utils.JitScalarType.from_value(self).onnx_type() + ) + one_sub_eps = g.op("Sub", one, eps) + self_less_equal_one_sub_eps = g.op("Greater", one_sub_eps, self) + temporary_self = g.op("Where", self_less_equal_one_sub_eps, self, one_sub_eps) + + temporary_self_less_eps = g.op("Less", temporary_self, eps) + z = g.op("Where", temporary_self_less_eps, eps, temporary_self) + else: + z = self + + sub = g.op("Sub", one, z) + div = g.op("Div", z, sub) + return g.op("Log", div) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..024966aa6f4242978669ef4d55abc770c9c0c78f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/utils.py @@ -0,0 +1,1942 @@ +# mypy: allow-untyped-defs +"""Functions to export models into the ONNX IR format. + +These models can be loaded with the ONNX library and then +converted to models which run on other deep learning frameworks. +""" + +from __future__ import annotations + + +__all__ = [ + "select_model_mode_for_export", + "disable_apex_o2_state_dict_hook", + "setup_onnx_logging", + "exporter_context", + "export", + "model_signature", + "warn_on_static_input_change", + "unpack_quantized_tensor", + "unconvertible_ops", + "register_custom_op_symbolic", + "unregister_custom_op_symbolic", + "_add_block", + "_add_input_to_block", + "_add_output_to_block", + "_apply_friendly_debug_names", + "_check_flatten_did_not_remove", + "_create_jit_graph", + "_decide_add_node_names", + "_decide_constant_folding", + "_decide_input_format", + "_decide_keep_init_as_input", + "_export", + "_get_aten_op_overload_name", + "_get_example_outputs", + "_get_module_attributes", + "_get_named_param_dict", + "_get_param_count_list", + "_is_constant_tensor_list", + "_model_to_graph", + "_optimize_graph", + "_pre_trace_quant_model", + "_reset_trace_module_map", + "_resolve_args_by_export_type", + "_run_symbolic_function", + "_run_symbolic_method", + "_set_input_and_output_names", + "_setup_trace_module_map", + "_should_aten_fallback", + "_signature", + "_split_tensor_list_constants", + "_trace_and_get_graph_from_model", + "_trace", + "_trigger_symbolic_function_registration", + "_validate_dynamic_axes", + "_verify_custom_op_name", +] + +import contextlib +import copy +import inspect +import re +import typing +import warnings +from typing import Any, cast +from typing_extensions import deprecated + +import torch +import torch._C._onnx as _C_onnx +import torch.jit._trace +from torch import _C +from torch.onnx import _constants, errors +from torch.onnx._internal.torchscript_exporter import ( + jit_utils, + onnx_proto_utils, + registration, + symbolic_helper, +) +from torch.onnx._internal.torchscript_exporter._globals import GLOBALS + + +if typing.TYPE_CHECKING: + from collections.abc import Callable, Collection, Mapping, Sequence + + +# TODO(justinchuby): Remove dependency to this global variable from constant_fold.cpp +# Skip check due to cannot import IValue from torch._C +_params_dict = {} # type: ignore[var-annotated] + + +@deprecated("Please set training mode before exporting the model", category=None) +@contextlib.contextmanager +def select_model_mode_for_export(model, mode: _C_onnx.TrainingMode): + """A context manager to temporarily set the training mode of ``model`` + to ``mode``, resetting it when we exit the with-block. + + .. deprecated:: 2.7 + Please set training mode before exporting the model. + + Args: + model: Same type and meaning as ``model`` arg to :func:`export`. + mode: Same type and meaning as ``training`` arg to :func:`export`. + """ + if not isinstance(mode, _C_onnx.TrainingMode): + raise TypeError( + f"'mode' should be a torch.onnx.TrainingMode enum, but got '{type(mode)}'." + ) + originally_training: bool = False + + if hasattr(model, "training"): + originally_training = model.training + + # ONNX opset 12 has better support for training amenable models, with updated + # versions of the dropout and batch_norm operators + if mode == _C_onnx.TrainingMode.TRAINING or ( + mode == _C_onnx.TrainingMode.PRESERVE and originally_training + ): + GLOBALS.export_training = True + if GLOBALS.export_onnx_opset_version < 12: + warnings.warn( + "You are exporting the model in training mode with onnx opset " + f"version {GLOBALS.export_onnx_opset_version}. " + "Opset versions lower than opset 12 will not be able to export " + "nodes such as Dropout and BatchNorm correctly.", + stacklevel=2, + ) + else: + GLOBALS.export_training = False + + GLOBALS.training_mode = mode + if mode == _C_onnx.TrainingMode.TRAINING: + model.train(True) + elif mode == _C_onnx.TrainingMode.EVAL: + model.train(False) + # else mode == _C_onnx.TrainingMode.PRESERVE, do nothing + + try: + yield + finally: + if hasattr(model, "training") and mode != _C_onnx.TrainingMode.PRESERVE: + model.train(originally_training) + + +@deprecated( + "Please remove usage of this function. Copy its logic if it is required in user code", + category=None, +) +@contextlib.contextmanager +def disable_apex_o2_state_dict_hook(model: torch.nn.Module | torch.jit.ScriptFunction): + """A context manager to temporarily disable the Apex O2 hook that returns. + + .. deprecated:: 2.7 + Please remove usage of this function. + """ + # Apex O2 hook state_dict to return fp16 weights as fp32. + # Exporter cannot identify them as same tensors. + # Since this hook is only used by optimizer, it is safe to + # remove this hook while exporting. + if not isinstance(model, torch.jit.ScriptFunction): + model_hooks = {} # type: ignore[var-annotated] + for module in model.modules(): + for key, hook in module._state_dict_hooks.items(): + if type(hook).__name__ == "O2StateDictHook": + if module not in model_hooks: + model_hooks[module] = {} + model_hooks[module][key] = hook + if module in model_hooks: + for key in model_hooks[module]: + module._state_dict_hooks.pop(key) + try: + yield + finally: + # Add the hooks back + for module, m_map in model_hooks.items(): + for key, hook in m_map.items(): + module._state_dict_hooks[key] = hook + else: + try: + yield + finally: + pass + + +@deprecated("The feature will be removed. Please remove usage of this function") +@contextlib.contextmanager +def setup_onnx_logging(verbose: bool): + """A context manager to temporarily set the ONNX logging verbosity. + + .. deprecated:: 2.7 + Please remove usage of this function. + """ + is_originally_enabled = _C._jit_is_onnx_log_enabled + if is_originally_enabled or verbose: # type: ignore[truthy-function] + _C._jit_set_onnx_log_enabled(True) + try: + yield + finally: + if not is_originally_enabled: # type: ignore[truthy-function] + _C._jit_set_onnx_log_enabled(False) + + +@deprecated( + "The feature will be removed. Please remove usage of this function " + "and implement equivalent logic if needed", + category=None, +) +@contextlib.contextmanager +def exporter_context(model, mode: _C_onnx.TrainingMode, verbose: bool): + """A context manager to temporarily set the training mode of ``model`` + to ``mode``, disable the Apex O2 hook, and set the ONNX logging verbosity. + + .. deprecated:: 2.7 + Please set training mode before exporting the model. + """ + with ( + select_model_mode_for_export(model, mode) as mode_ctx, + disable_apex_o2_state_dict_hook(model) as apex_ctx, + setup_onnx_logging(verbose) as log_ctx, + ): + yield (mode_ctx, apex_ctx, log_ctx) + + +def export( + model: torch.nn.Module | torch.jit.ScriptModule | torch.jit.ScriptFunction, + args: tuple[Any, ...] | torch.Tensor, + f: str, + *, + kwargs: dict[str, Any] | None = None, + export_params: bool = True, + verbose: bool = False, + training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL, + input_names: Sequence[str] | None = None, + output_names: Sequence[str] | None = None, + operator_export_type: _C_onnx.OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX, + opset_version: int | None = None, + do_constant_folding: bool = True, + dynamic_axes: Mapping[str, Mapping[int, str]] + | Mapping[str, Sequence[int]] + | None = None, + keep_initializers_as_inputs: bool | None = None, + custom_opsets: Mapping[str, int] | None = None, + export_modules_as_functions: bool | Collection[type[torch.nn.Module]] = False, + autograd_inlining: bool = True, +) -> None: + r"""Exports a model into ONNX format. + + If ``model`` is not a :class:`torch.jit.ScriptModule` nor a + :class:`torch.jit.ScriptFunction`, this runs + ``model`` once in order to convert it to a TorchScript graph to be exported + (the equivalent of :func:`torch.jit.trace`). Thus this has the same limited support + for dynamic control flow as :func:`torch.jit.trace`. + + Args: + model: The model to be exported. + args: + + args can be structured either as: + + 1. ONLY A TUPLE OF ARGUMENTS:: + + args = (x, y, z) + + The tuple should contain model inputs such that ``model(*args)`` is a valid + invocation of the model. Any non-Tensor arguments will be hard-coded into the + exported model; any Tensor arguments will become inputs of the exported model, + in the order they occur in the tuple. + + 2. A TENSOR:: + + args = torch.Tensor([1]) + + This is equivalent to a 1-ary tuple of that Tensor. + + 3. A TUPLE OF ARGUMENTS ENDING WITH A DICTIONARY OF NAMED ARGUMENTS:: + + args = (x, {"y": input_y, "z": input_z}) + + All but the last element of the tuple will be passed as non-keyword arguments, + and named arguments will be set from the last element. If a named argument is + not present in the dictionary, it is assigned the default value, or None if a + default value is not provided. + + .. warning:: + This behavior will be deprecated in a future release. Please use the + kwargs argument instead. + + .. note:: + If a dictionary is the last element of the args tuple, it will be + interpreted as containing named arguments. In order to pass a dict as the + last non-keyword arg, provide an empty dict as the last element of the args + tuple. For example, instead of:: + + torch.onnx.export( + model, + ( + x, + # WRONG: will be interpreted as named arguments + {y: z}, + ), + "test.onnx.pb", + ) + + Write:: + + torch.onnx.export(model, (x, {y: z}, {}), "test.onnx.pb") + + f: Path to the output ONNX model file. E.g. "model.onnx". + kwargs: Named arguments to the model. + export_params: If True, all parameters will + be exported. Set this to False if you want to export an untrained model. + In this case, the exported model will first take all of its parameters + as arguments, with the ordering as specified by ``model.state_dict().values()`` + verbose: if True, prints a description of the + model being exported to stdout. In addition, the final ONNX graph will include the + field ``doc_string``` from the exported model which mentions the source code locations + for ``model``. If True, ONNX exporter logging will be turned on. + training: + * ``TrainingMode.EVAL``: export the model in inference mode. + * ``TrainingMode.PRESERVE``: export the model in inference mode if model.training is + False and in training mode if model.training is True. + * ``TrainingMode.TRAINING``: export the model in training mode. Disables optimizations + which might interfere with training. + input_names (list of str, default empty list): names to assign to the + input nodes of the graph, in order. + output_names (list of str, default empty list): names to assign to the + output nodes of the graph, in order. + operator_export_type (enum, default OperatorExportTypes.ONNX): + + .. warning:: + This option will be deprecated in a future release. Future exported + graphs will always use the default opset domain. + + * ``OperatorExportTypes.ONNX``: Export all ops as regular ONNX ops + (in the default opset domain). + * ``OperatorExportTypes.ONNX_FALLTHROUGH``: Try to convert all ops + to standard ONNX ops in the default opset domain. If unable to do so + (e.g. because support has not been added to convert a particular torch op to ONNX), + fall back to exporting the op into a custom opset domain without conversion. Applies + to `custom ops `_ + as well as ATen ops. For the exported model to be usable, the runtime must support + these non-standard ops. + * ``OperatorExportTypes.ONNX_ATEN``: All ATen ops (in the TorchScript namespace "aten") + are exported as ATen ops (in opset domain "org.pytorch.aten"). + `ATen `_ is PyTorch's built-in tensor library, so + this instructs the runtime to use PyTorch's implementation of these ops. + + .. warning:: + + Models exported this way are probably runnable only by Caffe2. + + This may be useful if the numeric differences in implementations of operators are + causing large differences in behavior between PyTorch and Caffe2 (which is more + common on untrained models). + + * ``OperatorExportTypes.ONNX_ATEN_FALLBACK``: Try to export each ATen op + (in the TorchScript namespace "aten") as a regular ONNX op. If we are unable to do so + (e.g. because support has not been added to convert a particular torch op to ONNX), + fall back to exporting an ATen op. See documentation on OperatorExportTypes.ONNX_ATEN for + context. + For example:: + + graph(%0 : Float): + %3 : int = prim::Constant[value=0]() + # conversion unsupported + %4 : Float = aten::triu(%0, %3) + # conversion supported + %5 : Float = aten::mul(%4, %0) + return (%5) + + Assuming ``aten::triu`` is not supported in ONNX, this will be exported as:: + + graph(%0 : Float): + %1 : Long() = onnx::Constant[value={0}]() + # not converted + %2 : Float = aten::ATen[operator="triu"](%0, %1) + # converted + %3 : Float = onnx::Mul(%2, %0) + return (%3) + + .. warning:: + + Models exported this way are probably runnable only by Caffe2. + + opset_version (int, default 18): The version of the + `default (ai.onnx) opset `_ + to target. Must be >= 7. + do_constant_folding: Apply the constant-folding optimization. + Constant-folding will replace some of the ops that have all constant inputs + with pre-computed constant nodes. + dynamic_axes: + + By default the exported model will have the shapes of all input and output tensors + set to exactly match those given in ``args``. To specify axes of tensors as + dynamic (i.e. known only at run-time), set ``dynamic_axes`` to a dict with schema: + + * KEY (str): an input or output name. Each name must also be provided in ``input_names`` or + ``output_names``. + * VALUE (dict or list): If a dict, keys are axis indices and values are axis names. If a + list, each element is an axis index. + + For example:: + + class SumModule(torch.nn.Module): + def forward(self, x): + return torch.sum(x, dim=1) + + + torch.onnx.export( + SumModule(), + (torch.ones(2, 2),), + "onnx.pb", + input_names=["x"], + output_names=["sum"], + ) + + Produces:: + + input { + name: "x" + ... + shape { + dim { + dim_value: 2 # axis 0 + } + dim { + dim_value: 2 # axis 1 + ... + output { + name: "sum" + ... + shape { + dim { + dim_value: 2 # axis 0 + ... + + While:: + + torch.onnx.export( + SumModule(), + (torch.ones(2, 2),), + "onnx.pb", + input_names=["x"], + output_names=["sum"], + dynamic_axes={ + # dict value: manually named axes + "x": {0: "my_custom_axis_name"}, + # list value: automatic names + "sum": [0], + }, + ) + + Produces:: + + input { + name: "x" + ... + shape { + dim { + dim_param: "my_custom_axis_name" # axis 0 + } + dim { + dim_value: 2 # axis 1 + ... + output { + name: "sum" + ... + shape { + dim { + dim_param: "sum_dynamic_axes_1" # axis 0 + ... + + keep_initializers_as_inputs: If True, all the + initializers (typically corresponding to parameters) in the + exported graph will also be added as inputs to the graph. If False, + then initializers are not added as inputs to the graph, and only + the non-parameter inputs are added as inputs. + This may allow for better optimizations (e.g. constant folding) by + backends/runtimes. + + If True, `deduplicate_initializers` pass will not be executed. This means + initializers with duplicated values will not be deduplicated and + will be treated as distinct inputs to the graph. This allows different + input initializers to be supplied at the runtime following export. + + If ``opset_version < 9``, initializers MUST be part of graph + inputs and this argument will be ignored and the behavior will be + equivalent to setting this argument to True. + + custom_opsets (dict[str, int], default empty dict): A dict with schema: + + * KEY (str): opset domain name + * VALUE (int): opset version + + If a custom opset is referenced by ``model`` but not mentioned in this dictionary, + the opset version is set to 1. Only custom opset domain name and version should be + indicated through this argument. + + export_modules_as_functions: Flag to enable + exporting all ``nn.Module`` forward calls as local functions in ONNX. Or a set to indicate the + particular types of modules to export as local functions in ONNX. + This feature requires ``opset_version`` >= 15, otherwise the export will fail. This is because + ``opset_version`` < 15 implies IR version < 8, which means no local function support. + Module variables will be exported as function attributes. There are two categories of function + attributes. + + 1. Annotated attributes: class variables that have type annotations via + `PEP 526-style `_ + will be exported as attributes. + Annotated attributes are not used inside the subgraph of ONNX local function because + they are not created by PyTorch JIT tracing, but they may be used by consumers + to determine whether or not to replace the function with a particular fused kernel. + + 2. Inferred attributes: variables that are used by operators inside the module. Attribute names + will have prefix "inferred::". This is to differentiate from predefined attributes retrieved from + python module annotations. Inferred attributes are used inside the subgraph of ONNX local function. + + * ``False`` (default): export ``nn.Module`` forward calls as fine grained nodes. + * ``True``: export all ``nn.Module`` forward calls as local function nodes. + * Set of type of nn.Module: export ``nn.Module`` forward calls as local function nodes, + only if the type of the ``nn.Module`` is found in the set. + + autograd_inlining: Flag used to control whether to inline autograd functions. + Refer to https://github.com/pytorch/pytorch/pull/74765 for more details. + + Raises: + :class:`torch.onnx.errors.CheckerError`: If the ONNX checker detects an invalid ONNX graph. + :class:`torch.onnx.errors.UnsupportedOperatorError`: If the ONNX graph cannot be exported because it + uses an operator that is not supported by the exporter. + :class:`torch.onnx.errors.OnnxExporterError`: Other errors that can occur during export. + All errors are subclasses of :class:`errors.OnnxExporterError`. + """ + if operator_export_type != _C_onnx.OperatorExportTypes.ONNX: + warnings.warn( + "Setting `operator_export_type` to something other than default is deprecated. " + "The option will be removed in a future release.", + stacklevel=2, + category=DeprecationWarning, + ) + if training == _C_onnx.TrainingMode.TRAINING: + warnings.warn( + "Setting `training` to something other than default is deprecated. " + "The option will be removed in a future release. Please set the training mode " + "before exporting the model.", + stacklevel=2, + category=DeprecationWarning, + ) + + args = (args,) if isinstance(args, torch.Tensor) else args + if kwargs is not None: + args = args + (kwargs,) + + _export( + model, + args, + f, + export_params, + verbose, + training, + input_names, + output_names, + operator_export_type=operator_export_type, + opset_version=opset_version, + do_constant_folding=do_constant_folding, + dynamic_axes=dynamic_axes, + keep_initializers_as_inputs=keep_initializers_as_inputs, + custom_opsets=custom_opsets, + export_modules_as_functions=export_modules_as_functions, + autograd_inlining=autograd_inlining, + ) + + return None + + +def _is_constant_tensor_list(node) -> bool | None: + if node.kind() != "prim::Constant": + return False + output_type = node.output().type() + if output_type.isSubtypeOf(_C.ListType.ofTensors()): + return True + if output_type.isSubtypeOf(_C.ListType(_C.OptionalType.ofTensor())): + return True + + +# ONNX can't handle constants that are lists of tensors, which can +# get generated in constant prop. So we split them back into prim::ListConstructs + + +def _split_tensor_list_constants(g, block) -> None: + for node in block.nodes(): + for subblock in node.blocks(): + _split_tensor_list_constants(g, subblock) + if _is_constant_tensor_list(node): + inputs = [] + for val in node.output().toIValue(): + input = g.insertConstant(val) + input.node().moveBefore(node) + input.node().copyMetadata(node) + inputs.append(input) + + lc = ( + g.create("prim::ListConstruct", inputs) + .insertBefore(node) + .output() + .setType(_C.ListType.ofTensors()) + ) + lc.node().copyMetadata(node) + node.output().replaceAllUsesWith(lc) + + +def _optimize_graph( + graph: _C.Graph, + operator_export_type: _C_onnx.OperatorExportTypes, + _disable_torch_constant_prop: bool = False, + fixed_batch_size: bool = False, + params_dict=None, + dynamic_axes=None, + input_names=None, + module=None, +): + if params_dict is None: + params_dict = {} + + # Inline everything + _C._jit_pass_inline(graph) + + # Remove fork/wait nodes + _C._jit_pass_inline_fork_wait(graph) + _C._jit_pass_lint(graph) + if GLOBALS.autograd_inlining: + _C._jit_pass_onnx_autograd_function_process(graph) + _C._jit_pass_lower_all_tuples(graph) + + # we now record some ops like ones/zeros + # into a trace where we previously recorded constants. + # use constant prop to maintain our current level of onnx support + # without implementing symbolics for all of them + if _disable_torch_constant_prop is False: + _C._jit_pass_constant_propagation(graph) + + _split_tensor_list_constants(graph, graph) + # run dce to eliminate dead parts of the graph that might have been + # left behind by things like symbolic_override + _C._jit_pass_dce(graph) + _C._jit_pass_lint(graph) + + # CSE should improve perf when Autocast is used with disabled cache + # Autocast is disabled due to a limitation on tracer as described at https://github.com/pytorch/pytorch/issues/84092 + # Must run before _C._jit_pass_erase_number_types to prevent type substitution + if _C._jit_pass_cse(graph): + _C._jit_pass_onnx_lint(graph) + + _C._jit_pass_canonicalize_graph_fuser_ops(graph) + _C._jit_pass_lint(graph) + _C._jit_pass_peephole(graph, True) + _C._jit_pass_fuse_addmm(graph) + _C._jit_pass_lint(graph) + + _C._jit_pass_peephole(graph, True) + _C._jit_pass_lower_all_tuples(graph) + # in _jit_pass_onnx, symbolic functions are called for each node for conversion. + # However, there are nodes that cannot be converted without additional context. + # For example, the number of outputs from split (and whether it is static or dynamic) is unknown + # until the point where it is unpacked by listUnpack node. + # This pass does a preprocess, and prepares the nodes such that enough context can be received + # by the symbolic function. + _C._jit_pass_onnx_remove_inplace_ops_for_onnx(graph, module) + _C._jit_pass_onnx_preprocess(graph) + + # onnx does not support tuples, so try to remove them + _C._jit_pass_lint(graph) + + # onnx only supports tensors, but 1 / 2 = 0.5 and tensor(1) / tensor(2) = 0 + _C._jit_pass_prepare_division_for_onnx(graph) + + _C._jit_pass_onnx_remove_print(graph) + _C._jit_pass_onnx_preprocess_caffe2(graph) + + symbolic_helper._quantized_ops.clear() + # Unpack quantized weights for conv and linear ops and insert into graph. + _C._jit_pass_onnx_unpack_quantized_weights(graph, params_dict) + # onnx only supports tensors, so we turn all out number types into tensors + _C._jit_pass_erase_number_types(graph) + if GLOBALS.onnx_shape_inference: + input_names = [] if input_names is None else input_names + dynamic_axes = {} if dynamic_axes is None else dynamic_axes + _C._jit_pass_onnx_set_dynamic_input_shape(graph, dynamic_axes, input_names) + _C._jit_pass_onnx_lint(graph) + + graph = _C._jit_pass_onnx(graph, operator_export_type) + _C._jit_pass_onnx_lint(graph) + _C._jit_pass_lint(graph) + + _C._jit_pass_onnx_scalar_type_analysis( + graph, True, GLOBALS.export_onnx_opset_version + ) + _C._jit_pass_lint(graph) + + _C._jit_pass_onnx_peephole( + graph, GLOBALS.export_onnx_opset_version, fixed_batch_size + ) + _C._jit_pass_lint(graph) + + # graph is not a valid jit graph anymore because types have been replaced + # (e.g. int with Tensor), so it now contains operators that don't actually + # exist. We can't run normal dead code elimination because it'd fail trying + # to look up if an operator has side effects, but we can run a dead code + # elimination variant that doesn't need to look up if an op has side effects. + _C._jit_pass_dce_allow_deleting_nodes_with_side_effects(graph) + _C._jit_pass_lint(graph) + graph = _C._jit_pass_canonicalize(graph) + _C._jit_pass_lint(graph) + if GLOBALS.onnx_shape_inference: + try: + _C._jit_pass_onnx_graph_shape_type_inference( + graph, params_dict, GLOBALS.export_onnx_opset_version + ) + except RuntimeError: + # NOTE: shape type inference error should not stop the export process + # https://github.com/pytorch/pytorch/issues/132205 + pass + + return graph + + +def warn_on_static_input_change(input_states) -> None: + """Warns that changes to input dictionaries and strings won't take effect in the traced ONNX graph. + + We accept dictionaries and strings as ONNX inputs, but they should be only for + configuration use. we detect here if these inputs are modified, and if so we warn + the user that the changes won't take effect in the traced ONNX graph. + """ + for input, traced_input in zip(input_states[0], input_states[1]): + if isinstance(input, dict): + if list(input.keys()) != list(traced_input.keys()): + warning = ( + "We detected that you are modifying a dictionary that is an input to your " + "model. " + "Note that dictionaries are allowed as inputs in ONNX but they should be " + "handled with care. " + "Usages of dictionaries is not recommended, and should not be used except " + "for configuration use. " + "Also note that the order and values of the keys must remain the same. " + ) + warnings.warn(warning, stacklevel=2) + elif isinstance(input, str): + if input != traced_input: + warning = ( + "The model seems to have string inputs/outputs. " + "Note that strings will not appear as inputs/outputs of the ONNX graph. " + ) + warnings.warn(warning, stacklevel=2) + + +def _resolve_args_by_export_type(arg_name, arg_value, operator_export_type): + """Resolves the arguments that are ignored when export_type != operator_export_type.ONNX.""" + return arg_value + + +def _decide_keep_init_as_input( + keep_initializers_as_inputs: bool | None, + operator_export_type: _C_onnx.OperatorExportTypes, + opset_version: int, +): + """Decides whether the initializers in the graph should be listed as ONNX graph inputs. + + This method encapsulates the logic to decide whether the initializers in the graph + should be listed as ONNX graph inputs (i.e., whether to choose ONNX IR v3 or v4). + If keep_initializers_as_inputs is not specified (None), then we decide whether to keep + initializers as graph inputs (val_keep_init_as_ip) based on export type. If export type + is ONNX, then do not keep initializers as input (val_keep_init_as_ip=False). For all other + export types keep initializers as input (val_keep_init_as_ip=True). + If keep_initializers_as_inputs is specified, then respect it. Unless opset version <= 8, + in which case it must be ignored because for opset version <= 8, all initializers MUST be + part of graph input (only ONNX IR v3 is allowed), i.e. val_keep_init_as_ip=True. + + Special handling is needed for opset version 8 or lower, because irrespective + of user input for keep_initializers_as_inputs, the graph must follow ONNX IR v3 + semantics, i.e. all initializers must be listed as ONNX graph input. + """ + + if opset_version < 9: + if keep_initializers_as_inputs is False: + warnings.warn( + "Setting 'keep_initializers_as_inputs=False' for opset version" + "8 or lower would lead to an invalid ONNX graph. Therefore, " + "'keep_initializers_as_inputs=False' is ignored during export." + "Exported model will have initializers as graph inputs (compliant " + " to ONNX IR v3).", + stacklevel=2, + ) + return True # i.e. True == initializers are part of graph input (ONNX IR v3) + val_keep_init_as_ip = ( + True if keep_initializers_as_inputs is None else keep_initializers_as_inputs + ) + if ( + keep_initializers_as_inputs is None + and operator_export_type is _C_onnx.OperatorExportTypes.ONNX + ): + val_keep_init_as_ip = False + return val_keep_init_as_ip + + +def _decide_add_node_names(add_node_names, operator_export_type): + return _resolve_args_by_export_type( + "add_node_names", add_node_names, operator_export_type + ) + + +def _decide_constant_folding(do_constant_folding, operator_export_type, training): + do_constant_folding = _resolve_args_by_export_type( + "do_constant_folding", do_constant_folding, operator_export_type + ) + if do_constant_folding and ( + training is not None and training is not _C_onnx.TrainingMode.EVAL + ): + warnings.warn( + "It is recommended that constant folding be turned off ('do_constant_folding=False') " + "when exporting the model in training-amenable mode, i.e. with 'training=TrainingMode.TRAIN' " + "or 'training=TrainingMode.PRESERVE' (when model is in training mode). Otherwise, some " + "learnable model parameters may not translate correctly in the exported ONNX model " + "because constant folding mutates model parameters. Please consider " + "turning off constant folding or setting the training=TrainingMode.EVAL.", + stacklevel=2, + ) + return do_constant_folding + + +def _signature(model) -> inspect.Signature: + should_be_callable = getattr(model, "forward", model) + if callable(should_be_callable): + return inspect.signature(should_be_callable) + raise ValueError("model has no forward method and is not callable") + + +def _decide_input_format(model, args): + try: + sig = _signature(model) + except ValueError as e: + warnings.warn(f"{e}, skipping _decide_input_format", stacklevel=2) + return args + try: + ordered_list_keys = list(sig.parameters.keys()) + if ordered_list_keys[0] == "self": + ordered_list_keys = ordered_list_keys[1:] + args_dict: dict = {} + if isinstance(args, list): + args_list = args + elif isinstance(args, tuple): + args_list = list(args) + else: + args_list = [args] + if isinstance(args_list[-1], dict): + args_dict = args_list[-1] + args_list = args_list[:-1] + n_nonkeyword = len(args_list) + for optional_arg in ordered_list_keys[n_nonkeyword:]: + if optional_arg in args_dict: + args_list.append(args_dict[optional_arg]) + # Check if this arg has a default value + else: + param = sig.parameters[optional_arg] + if param.default != param.empty: + args_list.append(param.default) + args = args_list if isinstance(args, list) else tuple(args_list) + # Cases of models with no input args + except IndexError: + warnings.warn("No input args, skipping _decide_input_format", stacklevel=2) + except Exception as e: + warnings.warn(f"Skipping _decide_input_format\n {e.args[0]}", stacklevel=2) + return args + + +def _trace(func, args, operator_export_type, return_outs=False): + # Special case for common case of passing a single Tensor + if isinstance(args, torch.Tensor): + args = (args,) + + trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph( + func, + args, + strict=False, + _force_outplace=False, + _return_inputs_states=True, + ) + warn_on_static_input_change(inputs_states) + + trace_graph = _optimize_graph(trace_graph, operator_export_type, params_dict={}) + if return_outs: + return trace_graph, torch_out + return trace_graph + + +def _trace_and_get_graph_from_model(model, args): + # A basic sanity check: make sure the state_dict keys are the same + # before and after running the model. Fail fast! + orig_state_dict_keys = torch.jit._unique_state_dict(model).keys() + + # Disable Autocast cache because it replaces kernel's weight and bias + # by (undesired) constants. + # No perf impact for when there are reused weights since https://github.com/pytorch/pytorch/pull/85665 + prev_autocast_cache_enabled = torch.is_autocast_cache_enabled() + torch.set_autocast_cache_enabled(False) + trace_graph, torch_out, inputs_states = torch.jit._get_trace_graph( + model, + args, + strict=False, + _force_outplace=False, + _return_inputs_states=True, + ) + torch.set_autocast_cache_enabled(prev_autocast_cache_enabled) + + warn_on_static_input_change(inputs_states) + + if orig_state_dict_keys != torch.jit._unique_state_dict(model).keys(): + raise RuntimeError( + "state_dict changed after running the tracer; " + "something weird is happening in your model!" + ) + + return trace_graph, torch_out + + +def _get_param_count_list(method_graph, args_params): + param_count_list = [] + for input_, arg_params_ in zip(method_graph.inputs(), args_params): + if "PackedParams" in str(input_.type()): + in_vars, _ = torch.jit._flatten(arg_params_) + param_count_list.append(len(in_vars)) + else: + param_count_list.append(arg_params_ is not None) + + return param_count_list + + +def _check_flatten_did_not_remove(original, jit_flattened) -> None: + """torch.jit._flatten removes None. Check if it did so in this case.""" + + def flatten(x): + if isinstance(x, (list, tuple)): + for inner in x: + yield from flatten(inner) + elif isinstance(x, dict): + for inner in x.values(): + yield from flatten(inner) + else: + yield x + + flattened_with_none = list(flatten(original)) + num_none = len(flattened_with_none) - len(jit_flattened) + if num_none < 0: + raise AssertionError(f"num_none must be >= 0, got {num_none}") + if num_none: + raise ValueError( + f"args contained {num_none} None's after flattening. " + "When exporting a ScriptModule or ScriptFunction, no args may " + "be None because that breaks type propagation." + ) + + +def _create_jit_graph( + model: torch.nn.Module | torch.jit.ScriptFunction, args: Sequence[Any] +) -> tuple[_C.Graph, list[_C.IValue], Any | None, _C.ScriptModule | None]: + if isinstance(model, (torch.jit.ScriptFunction, torch.jit.ScriptModule)): + flattened_args = tuple(torch.jit._flatten(tuple(args))[0]) + _check_flatten_did_not_remove(args, flattened_args) + torch_out = None + + if isinstance(model, torch.jit.ScriptModule): + try: + graph = model.forward.graph # type: ignore[attr-defined] + except AttributeError as e: + raise RuntimeError("'forward' method must be a script method") from e + _C._jit_pass_onnx_function_substitution(graph) + freezed_module = _C._freeze_module( + cast(_C.ScriptModule, model._c), preserveParameters=True + ) + module, params = _C._jit_onnx_list_model_parameters(freezed_module) + method_graph = module._get_method("forward").graph + args_params = tuple(args) + tuple(params) + param_count_list = _get_param_count_list(method_graph, args_params) + in_vars, _ = torch.jit._flatten(args_params) + graph = _C._propagate_and_assign_input_shapes( + method_graph, tuple(in_vars), param_count_list, False, False + ) + return graph, params, torch_out, module + + # torch.jit.ScriptFunction + params = [] + graph = model.graph + _C._jit_pass_onnx_function_substitution(graph) + param_count_list = _get_param_count_list(graph, args) + graph = _C._propagate_and_assign_input_shapes( + graph, flattened_args, param_count_list, False, False + ) + return graph, params, torch_out, None + + graph, torch_out = _trace_and_get_graph_from_model(model, args) + _C._jit_pass_onnx_lint(graph) + state_dict = torch.jit._unique_state_dict(model) + params = list(state_dict.values()) + graph_inputs = list(graph.inputs()) + user_input_num = len(graph_inputs) - len(state_dict) + param_names = list(state_dict.keys()) + for i, inp in enumerate(graph_inputs): + if i >= user_input_num: + inp.setDebugName(param_names[i - user_input_num]) + _C._jit_pass_onnx_function_substitution(graph) + return graph, params, torch_out, None + + +def _get_named_param_dict(graph, params): + input_and_param_names = [val.debugName() for val in graph.inputs()] + param_names = input_and_param_names[len(input_and_param_names) - len(params) :] + _params_dict = dict(zip(param_names, params)) + return _params_dict + + +def _get_example_outputs(model, args): + input_args = copy.deepcopy(args) + input_kwargs = {} + if input_args and isinstance(input_args[-1], dict): + input_kwargs = input_args[-1] + input_args = input_args[:-1] + + example_outputs = model(*input_args, **input_kwargs) + if isinstance(example_outputs, list): + example_outputs = [example_outputs] + elif not isinstance(example_outputs, tuple): + example_outputs = (example_outputs,) + + return example_outputs + + +_qtype_vtype_map = { + torch.quint8: torch.uint8, + torch.qint8: torch.int8, + torch.qint32: torch.int32, + torch.quint4x2: torch.int8, +} + + +def unpack_quantized_tensor(value, cast_onnx_accepted=True): + if isinstance(value, torch.Tensor) and value.dtype in _qtype_vtype_map: + q_value_dequantize = value.dequantize() + q_scale = ( + torch.tensor(value.q_scale(), dtype=torch.double) + if cast_onnx_accepted + else torch.tensor(value.q_scale(), dtype=torch.float32) + ) + q_zero_point = ( + torch.tensor(value.q_zero_point(), dtype=torch.int64) + if cast_onnx_accepted + else torch.tensor(value.q_zero_point(), dtype=_qtype_vtype_map[value.dtype]) + ) + q_value = q_value_dequantize / q_scale + q_zero_point + q_value = q_value.to(dtype=_qtype_vtype_map[value.dtype]) + return q_value, q_scale, q_zero_point + else: + return (value,) + + +def _pre_trace_quant_model(model, args): + r"""Returns `torch.jit.trace(model, args)` if model is quantized. Otherwise do nothing and return + original model. + + This is due to https://github.com/pytorch/pytorch/issues/75761. + """ + if any( + hasattr(m, "_packed_params") for m in getattr(model, "modules", list)() + ) or any(getattr(arg, "is_quantized", False) for arg in args): + return torch.jit.trace(model, args) + return model + + +def _model_to_graph( + model, + args, + verbose=False, + input_names=None, + output_names=None, + operator_export_type=_C_onnx.OperatorExportTypes.ONNX, + do_constant_folding=True, + _disable_torch_constant_prop=False, + fixed_batch_size=False, + training=_C_onnx.TrainingMode.EVAL, + dynamic_axes=None, +) -> tuple[ + _C.Graph, + dict[str, torch.Tensor], + torch.Tensor + | tuple[torch.Tensor, ...] + | list[torch.Tensor] + | dict[str, torch.Tensor] + | Any + | None, +]: + """Converts model into an ONNX graph. + + Returns: + graph: A TorchScript IR Graph with ONNX nodes. + params_dict: Dict from input param name to param value. + torch_out: The output tensors resulting from the trace of ``model``. + If ``model`` is a :class:`torch.jit.ScriptModule` or :class:`torch.jit.ScriptFunction`, + this will be None, since we are not doing any tracing. + """ + # TODO: can we simplify this to always return a tuple of Tensor or None? + + # Special case for common case of passing a single Tensor + if isinstance(args, (torch.Tensor, int, float, bool)): + args = (args,) + + model = _pre_trace_quant_model(model, args) + graph, params, torch_out, module = _create_jit_graph(model, args) + params_dict = _get_named_param_dict(graph, params) + + try: + graph = _optimize_graph( + graph, + operator_export_type, + _disable_torch_constant_prop=_disable_torch_constant_prop, + fixed_batch_size=fixed_batch_size, + params_dict=params_dict, + dynamic_axes=dynamic_axes, + input_names=input_names, + module=module, + ) + except Exception: + _C._jit_onnx_log("Torch IR graph at exception: ", graph) + raise + + is_script = isinstance(model, (torch.jit.ScriptFunction, torch.jit.ScriptModule)) + if is_script: + example_outputs = _get_example_outputs(model, args) + example_outputs_final = () + for example_output in example_outputs: + example_outputs_final += unpack_quantized_tensor(example_output) + out_vars, desc = torch.jit._flatten(example_outputs_final) + _C._jit_pass_onnx_assign_output_shape( + graph, + out_vars, + desc, + GLOBALS.onnx_shape_inference, + is_script, + GLOBALS.export_onnx_opset_version, + ) + + # NB: ONNX requires complete information about output types, which might be + # erased by some optimizations, so we need to set it explicitly again. + else: + if not isinstance(torch_out, (list, tuple)): + output_wrapped = [torch_out] + else: + output_wrapped = torch_out # type: ignore[assignment] + + output_tensors, out_desc = torch.jit._flatten(tuple(output_wrapped)) + # assign_output_shape pass is not compatible with quantized outputs. + # Quantized outputs are flattened to 3 values in ONNX, while packed as + # single value in PyTorch. + if not any(getattr(out, "is_quantized", False) for out in output_tensors): + _C._jit_pass_onnx_assign_output_shape( + graph, + output_tensors, + out_desc, + GLOBALS.onnx_shape_inference, + is_script, + GLOBALS.export_onnx_opset_version, + ) + + _set_input_and_output_names(graph, input_names, output_names) + params_dict = _get_named_param_dict(graph, params) + + if ( + do_constant_folding + and GLOBALS.export_onnx_opset_version + >= _constants.ONNX_CONSTANT_FOLDING_MIN_OPSET + ): + if training is None or training == _C_onnx.TrainingMode.EVAL: + params_dict = _C._jit_pass_onnx_eval_peephole(graph, params_dict) + + params_dict = _C._jit_pass_onnx_constant_fold( + graph, params_dict, GLOBALS.export_onnx_opset_version + ) + _C._jit_pass_dce_allow_deleting_nodes_with_side_effects(graph) + + if GLOBALS.onnx_shape_inference: + try: + _C._jit_pass_onnx_graph_shape_type_inference( + graph, params_dict, GLOBALS.export_onnx_opset_version + ) + except RuntimeError: + # NOTE: shape type inference error should not stop the export process + # https://github.com/pytorch/pytorch/issues/132205 + pass + + params_dict = _C._jit_pass_onnx_eliminate_unused_items(graph, params_dict) + + # For ONNX opset < 9, constants only have three data types: float16, float, double. + # In this pass transform constants of other data types to float/double + cast operator. + if GLOBALS.export_onnx_opset_version < 9: + _C._jit_pass_onnx_cast_all_constant_to_floating(graph) + + params_dict = _C._jit_pass_filter_non_tensor_arguments(params_dict) + _C._jit_decay_packed_param_input_types(graph) + + # If output names lack a proper name and are identified only by their unique + # give them a legible name for debugging purposes + _apply_friendly_debug_names(graph, params_dict) + + return graph, params_dict, torch_out + + +@deprecated( + "Unconvertible ops are not definitive. Please remove usage of this function" +) +def unconvertible_ops( + model, + args, + training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL, + opset_version: int | None = None, +) -> tuple[_C.Graph, list[str]]: + """Returns an approximated list of all ops that are yet supported by :mod:`torch.onnx`. + + .. deprecated:: 2.5 + Unconvertible ops are not definitive. Please remove usage of this function. + + The list is approximated because some ops may be removed during the conversion + process and don't need to be converted. Some other ops may have partial support + that will fail conversion with particular inputs. Please open a Github Issue + for op support requests. + + Args: + model: Same as the `model` parameter in :func:`torch.onnx.export`. + args: Same as the `args` parameter in :func:`torch.onnx.export`. + training: Same as the `training` parameter in :func:`torch.onnx.export`. + opset_version: Same as the `opset_version` parameter in :func:`torch.onnx.export`. + + Returns: + The JIT graph and a list of unconvertible ops in the format of "domain::op". + """ + + opset_version = opset_version or _constants.ONNX_DEFAULT_OPSET + GLOBALS.export_onnx_opset_version = opset_version + + try: + with exporter_context(model, training, verbose=False): + # Create a mostly clean JIT graph that contains the plain aten and + # other ops we can check with the symbolic registry. + # NOTE: We don't want to actually convert any ops to ONNX or run any + # symbolic functions because there is a higher chance that a pass + # fails or an unconvertible op messes up the graph during ONNX conversion. + # This way we can always generate a list just by looking at the names + # of the ops in the graph. + args = _decide_input_format(model, args) + model = _pre_trace_quant_model(model, args) + graph, _, _, module = _create_jit_graph(model, args) + _C._jit_pass_inline(graph) + _C._jit_pass_onnx_remove_inplace_ops_for_onnx(graph, module) + _C._jit_pass_erase_number_types(graph) + _C._jit_pass_dce_allow_deleting_nodes_with_side_effects(graph) + except Exception as e: + raise errors.OnnxExporterError( + "Failed to discover unconvertible ops because of errors during the JIT graph " + "generation process." + ) from e + + unsupported_ops = [] + for node in graph.nodes(): + domain_op = node.kind() + if domain_op.startswith(("onnx::", "prim::")): + # We consider onnx and prim ops as supported ops, even though some "prim" + # ops are not implemented as symbolic functions, because they may be + # eliminated in the conversion passes. Users may still see errors caused + # by prim ops even though they don't show up in the list. + continue + if not registration.registry.is_registered_op( + domain_op.rstrip("_"), opset_version + ): + # We consider all registered ops supported, even though some of them are + # only partially supported, because there is not yet a good way to check + # if an op is fully supported. + # TODO(justinchuby): Create a way to check if an op is fully supported. + unsupported_ops.append(domain_op) + return graph, unsupported_ops + + +def _setup_trace_module_map( + model: torch.nn.Module | torch.jit.ScriptModule, + export_modules_as_functions: bool | Collection[type[torch.nn.Module]], +) -> set[str]: + def __register_attribute_hook() -> None: + attr_name = "_onnx_attrs" + + def _track_module_attributes_forward_pre_hook(module, input) -> None: + setattr(module, attr_name, _get_module_attributes(module)) + + def _track_module_attributes_forward_hook(module, input, output) -> None: + tracing_state = _C._get_tracing_state() + if not tracing_state: + return + + graph = tracing_state.graph() + onnx_attrs = {} + if hasattr(module, attr_name): + onnx_attrs = getattr(module, attr_name) + delattr(module, attr_name) + + _C._jit_pass_onnx_track_scope_attributes(graph, onnx_attrs) + + for m in model.modules(): + m.register_forward_hook(_track_module_attributes_forward_hook) + m.register_forward_pre_hook(_track_module_attributes_forward_pre_hook) + + def _unqualified_variable_name(qualified_name: str) -> str: + """ + Parse qualified variable name and return the unqualified version. + + Pure numeric atoms are considered inadequate, so this function will look past them, + and start from the first non-numeric atom. + + Example: + >>> _unqualified_variable_name("__main__.Foo.bar") + 'bar' + >>> _unqualified_variable_name("__main__.Foo.bar.0") + 'bar.0' + """ + name_atoms = qualified_name.split(".") + for i, atom in reversed(list(enumerate(name_atoms))): + if not atom.isnumeric(): + return ".".join(name_atoms[i:]) + return qualified_name + + trace_module_map = { + _m: torch._C._jit_onnx_create_full_scope_name( + torch.typename(type(_m)), _unqualified_variable_name(_n) + ) + for _n, _m in model.named_modules() + } + torch.jit._trace._trace_module_map = trace_module_map + if isinstance(export_modules_as_functions, bool) and export_modules_as_functions: + module_typenames = {torch.typename(type(module)) for module in trace_module_map} + elif isinstance(export_modules_as_functions, set) and export_modules_as_functions: + + def _find_typename(v): + if isinstance(v, type): + return torch.typename(v) + else: + raise RuntimeError( + "Only type of the `nn.Module` should be " + "passed in the set for argument `export_modules_as_functions`. " + f"Got `{type(v).__name__}`." + ) + + module_typenames = {_find_typename(v) for v in export_modules_as_functions} + else: + module_typenames = set() + + if module_typenames: + __register_attribute_hook() + + return module_typenames + + +def _reset_trace_module_map() -> None: + torch.jit._trace._trace_module_map = None + _C._jit_pass_onnx_clear_scope_records() + + +def _get_module_attributes(module): + annotations = typing.get_type_hints(type(module)) + base_m_annotations = typing.get_type_hints(torch.nn.Module) + [annotations.pop(k, None) for k in base_m_annotations] + # Check whether module attributes can be accessed. Some classes + # define attributes but don't provide access to them in their + # constructor. + # + # For example, torch.nn.Embedding has the `freeze` variable and its + # type specified in the class but the attribute is not created in the + # constructor. In other words, there is no `self.freeze = ` + # in the constructor. + # + # Reference: https://github.com/pytorch/pytorch/blob/92de1d322223fb5584e384971b32c46b93bc2f4b/torch/nn/modules/sparse.py#L120 + attrs = {} + for k in annotations: + try: + attrs[k] = getattr(module, k) + except AttributeError: + _C._jit_onnx_log(f"Skipping module attribute '{k}'") + continue + return attrs + + +def _trigger_symbolic_function_registration() -> None: + """Trigger the registration of symbolic functions for all supported opsets.""" + + from torch.onnx._internal.torchscript_exporter import ( # noqa: F401 + symbolic_opset10, + symbolic_opset11, + symbolic_opset12, + symbolic_opset13, + symbolic_opset14, + symbolic_opset15, + symbolic_opset16, + symbolic_opset17, + symbolic_opset18, + symbolic_opset19, + symbolic_opset20, + symbolic_opset7, + symbolic_opset8, + symbolic_opset9, + ) + + +def _export( + model, + args, + f, + export_params=True, + verbose=False, + training=_C_onnx.TrainingMode.EVAL, + input_names=None, + output_names=None, + operator_export_type=_C_onnx.OperatorExportTypes.ONNX, + export_type=None, + opset_version=None, + do_constant_folding=True, + dynamic_axes=None, + keep_initializers_as_inputs=None, + fixed_batch_size=False, + custom_opsets=None, + add_node_names=True, + onnx_shape_inference=True, + export_modules_as_functions: Any = False, + autograd_inlining=True, +): + if GLOBALS.in_onnx_export is not False: + raise AssertionError("GLOBALS.in_onnx_export must be False") + + _trigger_symbolic_function_registration() + + if isinstance(model, torch.nn.DataParallel): + raise ValueError( + "torch.nn.DataParallel is not supported by ONNX " + "exporter, please use 'attribute' module to " + "unwrap model from torch.nn.DataParallel. Try " + "torch.onnx.export(model.module, ...)" + ) + + GLOBALS.onnx_shape_inference = onnx_shape_inference + + if opset_version is None: + opset_version = _constants.ONNX_DEFAULT_OPSET + + if opset_version > _constants.ONNX_TORCHSCRIPT_EXPORTER_MAX_OPSET: + warnings.warn( + f"Exporting to ONNX opset version {opset_version} is not supported. " + f"by 'torch.onnx.export()'. " + f"The highest opset version supported is {_constants.ONNX_TORCHSCRIPT_EXPORTER_MAX_OPSET}. " + f"To use a newer opset version, consider 'torch.onnx.export(..., dynamo=True)'. ", + stacklevel=2, + category=errors.OnnxExporterWarning, + ) + + if export_modules_as_functions and opset_version < 15: + raise ValueError( + "`export_modules_as_functions` is not supported for `opset_version` < 15." + "This is because `opset_version` < 15 implies IR version < 8, which means " + "no local function support. " + ) + if not operator_export_type: + operator_export_type = _C_onnx.OperatorExportTypes.ONNX + + # By default, training=TrainingMode.EVAL, + # which is good because running a model in training mode could result in + # internal buffers getting updated, dropout getting applied, etc. + # If you really know what you're doing, you can turn + # training=TrainingMode.TRAINING or training=TrainingMode.PRESERVE, + # (to preserve whatever the original training mode was.) + GLOBALS.export_onnx_opset_version = opset_version + GLOBALS.operator_export_type = operator_export_type + + try: + GLOBALS.in_onnx_export = True + _autograd_inlining_previous = GLOBALS.autograd_inlining + GLOBALS.autograd_inlining = autograd_inlining + + module_typenames_to_export_as_functions: set[str] = set() + if isinstance(model, (torch.nn.Module, torch.jit.ScriptModule)): + module_typenames_to_export_as_functions = _setup_trace_module_map( + model, export_modules_as_functions + ) + + with exporter_context(model, training, verbose): + val_keep_init_as_ip = _decide_keep_init_as_input( + keep_initializers_as_inputs, + operator_export_type, + opset_version, + ) + val_add_node_names = _decide_add_node_names( + add_node_names, operator_export_type + ) + val_do_constant_folding = _decide_constant_folding( + do_constant_folding, operator_export_type, training + ) + # Normally f can be a file-like object, but for large models, the external data format requires a + # valid `model_file_location`. Code in export.cpp will enforce this. + if isinstance(f, str): + model_file_location = f + else: + model_file_location = "" + args = _decide_input_format(model, args) + if dynamic_axes is None: + dynamic_axes = {} + _validate_dynamic_axes(dynamic_axes, model, input_names, output_names) + + graph, params_dict, torch_out = _model_to_graph( + model, + args, + verbose, + input_names, + output_names, + operator_export_type, + val_do_constant_folding, + fixed_batch_size=fixed_batch_size, + training=training, + dynamic_axes=dynamic_axes, + ) + + if custom_opsets is None: + custom_opsets = {} + + _C._jit_pass_dce_allow_deleting_nodes_with_side_effects(graph) + node_attr_to_name = {} # type: ignore[var-annotated] + if module_typenames_to_export_as_functions: + # NOTE: cannot call DCE after this pass. DCE will remove function definition nodes. + node_attr_to_name = _C._jit_pass_onnx_function_extraction( + graph, + module_typenames_to_export_as_functions, + list(params_dict.keys()), + ) + + if keep_initializers_as_inputs is not True: + params_dict = _C._jit_pass_onnx_deduplicate_initializers( # type: ignore[assignment] + graph, + params_dict, # type: ignore[arg-type] + getattr(model, "training", False), # type: ignore[arg-type] + ) + _C._jit_pass_onnx_assign_scoped_names_for_node_and_value(graph) + defer_weight_export = False + if export_params: + ( + proto, + export_map, + _val_use_external_data_format, + _node_names, + ) = graph._export_onnx( # type: ignore[attr-defined] + params_dict, + opset_version, + dynamic_axes, + defer_weight_export, + operator_export_type, + not verbose, + val_keep_init_as_ip, + custom_opsets, + val_add_node_names, + model_file_location, + node_attr_to_name, + ) + else: + ( + proto, + export_map, + _, + _, + ) = graph._export_onnx( # type: ignore[attr-defined] + {}, + opset_version, + dynamic_axes, + defer_weight_export, + operator_export_type, + not verbose, + val_keep_init_as_ip, + custom_opsets, + val_add_node_names, + model_file_location, + node_attr_to_name, + ) + # insert function_proto into model_proto. + proto = onnx_proto_utils._add_onnxscript_fn( + proto, + custom_opsets, + ) + if verbose: + _C._jit_onnx_log("Exported graph: ", graph) + onnx_proto_utils._export_file(proto, f, export_map) + finally: + if not GLOBALS.in_onnx_export: + raise AssertionError("GLOBALS.in_onnx_export must be True") + GLOBALS.in_onnx_export = False + GLOBALS.autograd_inlining = _autograd_inlining_previous + _reset_trace_module_map() + + return torch_out + + +def _apply_friendly_debug_names(graph, params) -> None: + for n in graph.nodes(): + for v in n.inputs(): + old_name = v.debugName() + if old_name != str(v.unique()): + continue + new_name = f"{n.kind()}_{v.unique()}" + v.setDebugName(new_name) + if old_name in params: + params[new_name] = params.pop(old_name) + + +def _set_input_and_output_names(graph, input_names, output_names) -> None: + def set_names(node_list, name_list, descriptor) -> None: + if name_list is None: + return + if len(name_list) > len(node_list): + raise RuntimeError( + f"number of {descriptor} names provided ({len(name_list)}) " + f"exceeded number of {descriptor}s ({len(node_list)})" + ) + + # Mark if the output node DebugName is set before. + output_node_set = set() + for i, (name, node) in enumerate(zip(name_list, node_list)): + # Duplicated output node, insert onnx::Identity to avoid setting the same DebugName after setDebugName(). + if descriptor == "output": + if node in output_node_set: + identity_node = graph.create("onnx::Identity") + identity_node.insertAfter(node.node()) + identity_node.addInput(node) + identity_node.output().setType(node.type()) + graph.return_node().replaceInput(i, identity_node.output()) + node = identity_node.output() + output_node_set.add(node) + + if node.debugName() != name: + node.setDebugName(name) + + set_names(list(graph.inputs()), input_names, "input") + set_names(list(graph.outputs()), output_names, "output") + + +def _run_symbolic_method(g, op_name, symbolic_fn, args): + r""" + This trampoline function gets invoked for every symbolic method + call from C++. + """ + try: + graph_context = jit_utils.GraphContext( + graph=g, + block=g.block(), + opset=GLOBALS.export_onnx_opset_version, + original_node=None, # type: ignore[arg-type] + params_dict=_params_dict, + env={}, + values_in_env=set(), + new_nodes=[], + ) + return symbolic_fn(graph_context, *args) + except TypeError as e: + # Handle the specific case where we didn't successfully dispatch + # to symbolic_fn. Otherwise, the backtrace will have the clues + # you need. + e.args = (f"{e.args[0]} (occurred when translating {op_name})",) + raise + + +def _add_block(node: _C.Node) -> _C.Block: + return node.addBlock() + + +def _add_input_to_block(block: _C.Block): + return block.addInputToBlock() # type: ignore[attr-defined] + + +def _add_output_to_block(block: _C.Block, value: _C.Value) -> int: + return block.registerOutput(value) + + +def _should_aten_fallback( + name: str, opset_version: int, operator_export_type: _C_onnx.OperatorExportTypes +) -> bool: + # For all builds, if domain=="aten" and operator_export_type==ONNX_ATEN, + # an aten::ATen operator is created regardless of symbolics existence + + is_exportable_aten_op = registration.registry.is_registered_op(name, opset_version) + is_onnx_aten_export = operator_export_type == _C_onnx.OperatorExportTypes.ONNX_ATEN + is_aten_fallback_export = ( + operator_export_type == _C_onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK + ) + + if not name.startswith("aten::"): + return False + + if is_onnx_aten_export or (is_aten_fallback_export and not is_exportable_aten_op): + return True + + return False + + +def _get_aten_op_overload_name(n: _C.Node) -> str: + # Returns `overload_name` attribute to ATen ops on non-Caffe2 builds + schema = n.schema() + if not schema.startswith("aten::"): + return "" + return _C.parse_schema(schema).overload_name + + +def _run_symbolic_function( + graph: _C.Graph, + block: _C.Block, + node: _C.Node, + inputs: Any, + env: dict[_C.Value, _C.Value], + values_in_env: set[_C.Value], + new_nodes: list[_C.Node], + operator_export_type=_C_onnx.OperatorExportTypes.ONNX, +) -> _C.Value | Sequence[_C.Value | None] | None: + """Runs a symbolic function. + + The function is used in C++ to export the node to ONNX. + + Returns: + A single or a tuple of Values. + None when the node gets cloned as is into the new graph. + """ + + opset_version = GLOBALS.export_onnx_opset_version + + # See Note [Export inplace] + node_kind = node.kind() + if node_kind.endswith("_"): + # Treat relu_ -> relu; add_ -> add etc. + ns_op_name = node_kind[:-1] + else: + ns_op_name = node_kind + + namespace, op_name = jit_utils.parse_node_kind(ns_op_name) + + graph_context = jit_utils.GraphContext( + graph=graph, + block=block, + opset=opset_version, + original_node=node, + params_dict=_params_dict, + env=env, + values_in_env=values_in_env, + new_nodes=new_nodes, + ) + + # Direct ATen export requested + if _should_aten_fallback(ns_op_name, opset_version, operator_export_type): + attrs = { + k + "_" + node.kindOf(k)[0]: symbolic_helper._node_get(node, k) + for k in node.attributeNames() + } + outputs = node.outputsSize() + attrs["outputs"] = outputs + return graph_context.aten_op( + op_name, + *inputs, + overload_name=_get_aten_op_overload_name(node), + **attrs, + ) + + try: + domain = namespace + symbolic_function_name = f"{domain}::{op_name}" + + symbolic_function_group = registration.registry.get_function_group( + symbolic_function_name + ) + if symbolic_function_group is not None: + symbolic_fn = symbolic_function_group.get(opset_version) + if symbolic_fn is not None: + # TODO Wrap almost identical attrs assignment or comment the difference. + attrs = { + k: symbolic_helper._node_get(node, k) for k in node.attributeNames() + } + return symbolic_fn(graph_context, *inputs, **attrs) + + attrs = { + k + "_" + node.kindOf(k)[0]: symbolic_helper._node_get(node, k) + for k in node.attributeNames() + } + if namespace == "onnx": + # Clone node to trigger ONNX shape inference + return graph_context.op( + op_name, *inputs, **attrs, outputs=node.outputsSize() + ) # type: ignore[attr-defined] + + raise errors.UnsupportedOperatorError( + symbolic_function_name, + opset_version, + symbolic_function_group.get_min_supported() + if symbolic_function_group + else None, + ) + + except RuntimeError: + if operator_export_type == _C_onnx.OperatorExportTypes.ONNX_FALLTHROUGH: + return None + elif operator_export_type == _C_onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK: + # Emit ATen op for non-Caffe2 builds when `operator_export_type==ONNX_ATEN_FALLBACK` + attrs = { + k + "_" + node.kindOf(k)[0]: symbolic_helper._node_get(node, k) + for k in node.attributeNames() + } + return graph_context.aten_op( + op_name, + *inputs, + overload_name=_get_aten_op_overload_name(node), + **attrs, + ) + raise + except TypeError as e: + # Handle the specific case where we didn't successfully dispatch. + # Otherwise, the backtrace will have the clues you need. + e.args = (f"{e.args[0]} \n(Occurred when translating {op_name}).",) + raise + + +def _verify_custom_op_name(symbolic_name: str) -> None: + if not re.match(r"^[a-zA-Z0-9-_]+::[a-zA-Z-_]+[a-zA-Z0-9-_]*$", symbolic_name): + raise errors.OnnxExporterError( + f"Failed to register operator {symbolic_name}. " + "The symbolic name must match the format domain::name, " + "and should start with a letter and contain only " + "alphanumerical characters" + ) + + ns, _ = jit_utils.parse_node_kind(symbolic_name) + if ns == "onnx": + raise ValueError( + f"Failed to register operator {symbolic_name}. {ns} domain cannot be modified." + ) + + +def register_custom_op_symbolic( + symbolic_name: str, + symbolic_fn: Callable, + opset_version: int, +) -> None: + """Registers a symbolic function for a custom operator. + + When the user registers symbolic for custom/contrib ops, + it is highly recommended to add shape inference for that operator via setType API, + otherwise the exported graph may have incorrect shape inference in some extreme cases. + An example of setType is `test_aten_embedding_2` in `test_operators.py`. + + See "Custom Operators" in the module documentation for an example usage. + + Args: + symbolic_name (str): The name of the custom operator in "::" + format. + symbolic_fn (Callable): A function that takes in the ONNX graph and + the input arguments to the current operator, and returns new + operator nodes to add to the graph. + opset_version (int): The ONNX opset version in which to register. + """ + if symbolic_name.startswith("::"): + symbolic_name = f"aten{symbolic_name}" + + _verify_custom_op_name(symbolic_name) + + registration.custom_onnx_symbolic(symbolic_name, opset_version)(symbolic_fn) + + +def unregister_custom_op_symbolic(symbolic_name: str, opset_version: int) -> None: + """Unregisters ``symbolic_name``. + + See "Custom Operators" in the module documentation for an example usage. + + Args: + symbolic_name (str): The name of the custom operator in "::" + format. + opset_version (int): The ONNX opset version in which to unregister. + """ + if symbolic_name.startswith("::"): + symbolic_name = f"aten{symbolic_name}" + + _verify_custom_op_name(symbolic_name) + + registration.registry.unregister(symbolic_name, opset_version) + + +def _validate_dynamic_axes(dynamic_axes, model, input_names, output_names) -> None: + """Ensures dynamic axes argument is follows the expected format.""" + if len(dynamic_axes) == 0: + return + + if hasattr(model, "graph"): + # Extracting set of valid input/output names that shall be used for dynamic_axes + if (input_names is None) or len(input_names) == 0: + input_names = [x.debugName() for x in model.graph.inputs()] + if (output_names is None) or len(output_names) == 0: + output_names = [y.debugName() for y in model.graph.outputs()] + + valid_names = set((input_names or []) + (output_names or [])) + + # If dynamic axes are provided as a list rather than dictionary, they should + # first get converted to a dictionary in expected format. If desired axes names + # are not provided for dynamic axes, automatic names shall be generated for + # provided dynamic axes of specified input/output + for key, value in dynamic_axes.items(): + if key not in valid_names: + warnings.warn( + f"Provided key {key} for dynamic axes is not a valid input/output name", + stacklevel=2, + ) + if isinstance(value, list): + warnings.warn( + "No names were found for specified dynamic axes of provided input." + f"Automatically generated names will be applied to each dynamic axes of input {key}", + stacklevel=2, + ) + + value_dict = {} + for i, x in enumerate(value): + if not isinstance(x, int): + raise ValueError( + "The type of axis index is expected to be an integer" + ) + if x in value_dict: + warnings.warn( + f"Duplicate dynamic axis index {x} was provided for input {key}.", + stacklevel=2, + ) + else: + value_dict[x] = str(key) + "_dynamic_axes_" + str(i + 1) + dynamic_axes[key] = value_dict + + +def model_signature(model: torch.nn.Module | Callable) -> inspect.Signature: + return inspect.signature( + model.forward if isinstance(model, torch.nn.Module) else model + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/verification.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/verification.py new file mode 100644 index 0000000000000000000000000000000000000000..57f764e106ba33a607bc86f2f76feaa298c85796 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/_internal/torchscript_exporter/verification.py @@ -0,0 +1,516 @@ +# mypy: allow-untyped-defs +"""The ONNX verification module provides a set of tools to verify the correctness of ONNX models.""" + +from __future__ import annotations + + +__all__ = [ + "OnnxBackend", + "VerificationOptions", + "verify", +] + +import contextlib +import copy +import dataclasses +import enum +import io +import os +import tempfile +import warnings +from collections.abc import Mapping, Sequence +from typing import Any + +import numpy as np +import numpy.typing as npt + +import torch +import torch._C._onnx as _C_onnx +from torch.onnx._internal.torchscript_exporter import utils +from torch.types import Number + + +# Everything below are deprecated ############################################## + +_ORT_PROVIDERS = ("CPUExecutionProvider",) + +_NumericType = Number | torch.Tensor | np.ndarray +_ModelType = torch.nn.Module | torch.jit.ScriptModule +_InputArgsType = torch.Tensor | tuple[Any, ...] +_InputKwargsType = Mapping[str, Any] +_OutputsType = Sequence[_NumericType] | Sequence + + +class OnnxBackend(enum.Enum): + """Enum class for ONNX backend used for export verification. + + .. deprecated:: 2.7 + Consider using ``torch.onnx.export(..., dynamo=True)`` and use the returned + ``ONNXProgram`` to test the ONNX model. + """ + + REFERENCE = "ONNXReferenceEvaluator" + ONNX_RUNTIME_CPU = "CPUExecutionProvider" + ONNX_RUNTIME_CUDA = "CUDAExecutionProvider" + + +@dataclasses.dataclass +class VerificationOptions: + """Options for ONNX export verification. + + .. deprecated:: 2.7 + Consider using ``torch.onnx.export(..., dynamo=True)`` and use the returned + ``ONNXProgram`` to test the ONNX model. + + Attributes: + flatten: If True, unpack nested list/tuple/dict inputs into a flattened list of + Tensors for ONNX. Set this to False if nested structures are to be preserved + for ONNX, which is usually the case with exporting ScriptModules. Default True. + ignore_none: Whether to ignore None type in torch output, which is usually the + case with tracing. Set this to False, if torch output should keep None type, + which is usually the case with exporting ScriptModules. Default to True. + check_shape: Whether to check the shapes between PyTorch and ONNX Runtime outputs + are exactly the same. Set this to False to allow output shape broadcasting. + Default to True. + check_dtype: Whether to check the dtypes between PyTorch and ONNX Runtime outputs + are consistent. Default to True. + backend: ONNX backend for verification. Default to OnnxBackend.ONNX_RUNTIME_CPU. + rtol: relative tolerance in comparison between ONNX and PyTorch outputs. + atol: absolute tolerance in comparison between ONNX and PyTorch outputs. + remained_onnx_input_idx: If provided, only the specified inputs will be passed + to the ONNX model. Supply a list when there are unused inputs in the model. + Since unused inputs will be removed in the exported ONNX model, supplying + all inputs will cause an error on unexpected inputs. This parameter tells + the verifier which inputs to pass into the ONNX model. + acceptable_error_percentage: acceptable percentage of element mismatches in comparison. + It should be a float of value between 0.0 and 1.0. + """ + + flatten: bool = True + ignore_none: bool = True + check_shape: bool = True + check_dtype: bool = True + backend: OnnxBackend = OnnxBackend.ONNX_RUNTIME_CPU + rtol: float = 1e-3 + atol: float = 1e-7 + remained_onnx_input_idx: Sequence[int] | None = None + acceptable_error_percentage: float | None = None + + +def _flatten_tuples(elem): + flattened = [] + for t in elem: + if isinstance(t, tuple): + flattened.extend(_flatten_tuples(t)) + else: + flattened.append(t) + return flattened + + +def _to_numpy(elem) -> list | npt.NDArray: + if isinstance(elem, torch.Tensor): + if elem.requires_grad: + return elem.detach().cpu().numpy() + else: + return elem.cpu().numpy() + elif isinstance(elem, (list, tuple)): + return [_to_numpy(inp) for inp in elem] + elif isinstance(elem, (bool, int, float)): + return np.array(elem) + elif isinstance(elem, dict): + flattened = [] + for k in elem: + flattened.extend([_to_numpy(k), _to_numpy(elem[k])]) + return flattened + return elem + + +def _inline_flatten_list(inputs, res_list) -> list: + for i in inputs: + res_list.append(i) if not isinstance( + i, (list, tuple) + ) else _inline_flatten_list(i, res_list) + return res_list + + +def _unpack_to_numpy(values, cast_onnx_accepted=True) -> list: + value_unpacked = [] + for value in values: + value_unpacked.extend( + utils.unpack_quantized_tensor(value, cast_onnx_accepted=cast_onnx_accepted) + ) + return [_to_numpy(v) for v in value_unpacked] + + +def _run_onnx(onnx_session, inputs) -> _OutputsType: + kw_inputs = {} + if inputs and isinstance(inputs[-1], dict): + kw_inputs = inputs[-1] + inputs = inputs[:-1] + inputs = _unpack_to_numpy(_flatten_tuples(inputs)) + ort_inputs = {} + for input_name, input in kw_inputs.items(): + ort_inputs[input_name] = _to_numpy(input) + inputs = _to_numpy(inputs) + if hasattr(onnx_session, "get_inputs"): + # onnxruntime.InferenceSession + input_names = [i.name for i in onnx_session.get_inputs()] + elif hasattr(onnx_session, "input_names"): + # onnx.reference.ReferenceEvaluator + input_names = onnx_session.input_names + else: + raise ValueError(f"Unknown ONNX backend type: {type(onnx_session)}.") + + for i, input in enumerate(inputs): + if i == len(input_names) or input_names[i] in ort_inputs: + raise ValueError( + f"got too many positional inputs. inputs: {inputs}. kw_inputs: {kw_inputs}. " + f"input names: {input_names}." + ) + ort_inputs[input_names[i]] = input + onnx_outs = onnx_session.run(None, ort_inputs) + return onnx_outs + + +def _ort_session( + model: str | io.BytesIO, ort_providers: Sequence[str] = _ORT_PROVIDERS +): + try: + import onnxruntime # type: ignore[import] + except ImportError as e: + raise ImportError("onnxruntime is required for export verification.") from e + + if ort_providers is None: + ort_providers = _ORT_PROVIDERS + + session_options = onnxruntime.SessionOptions() + # suppress ort warnings. + # 0:Verbose, 1:Info, 2:Warning. 3:Error, 4:Fatal. Default is 2. + session_options.log_severity_level = 3 + ort_session = onnxruntime.InferenceSession( + model if isinstance(model, str) else model.getvalue(), + session_options, + providers=ort_providers, + ) + return ort_session + + +def _onnx_backend_session(model: str | io.BytesIO, backend: OnnxBackend): + if backend == OnnxBackend.REFERENCE: + raise NotImplementedError + elif backend in {OnnxBackend.ONNX_RUNTIME_CPU, OnnxBackend.ONNX_RUNTIME_CUDA}: + onnx_session = _ort_session(model, (backend.value,)) + else: + raise ValueError(f"Unsupported backend: {backend}") + return onnx_session + + +def _compare_onnx_pytorch_outputs_in_np( + onnx_outs: _OutputsType, + pt_outs: _OutputsType, + options: VerificationOptions, +) -> None: + if len(onnx_outs) != len(pt_outs): + raise AssertionError( + f"Number of outputs differ ONNX runtime: ({len(onnx_outs)}) PyTorch: ({len(pt_outs)})" + ) + acceptable_error_percentage = options.acceptable_error_percentage + if acceptable_error_percentage and ( + acceptable_error_percentage > 1.0 or acceptable_error_percentage < 0.0 + ): + raise ValueError( + "If set, acceptable_error_percentage should be between 0.0 and 1.0" + ) + + for ort_out, pt_out in zip(onnx_outs, pt_outs): + try: + # TODO: Remove `check_shape` option once every shape inconsistent issue is addressed. + if not options.check_shape: + # Allow different but broadcastable output shapes. + ort_out, pt_out = np.broadcast_arrays(ort_out, pt_out) + torch.testing.assert_close( + ort_out, + pt_out, + rtol=options.rtol, + atol=options.atol, + check_dtype=options.check_dtype, + equal_nan=True, + ) + except AssertionError as e: + if acceptable_error_percentage: + error_percentage = 1 - np.sum( + np.isclose(ort_out, pt_out, rtol=options.rtol, atol=options.atol) + ) / np.prod(ort_out.shape) # pyrefly: ignore [missing-attribute] + if error_percentage <= acceptable_error_percentage: + warnings.warn( + f"Suppressed AssertionError:\n{e}.\n" + f"Error percentage {error_percentage} " + f"within acceptable range {acceptable_error_percentage}.", + stacklevel=2, + ) + continue + # pyrefly: ignore [missing-attribute] + if ort_out.dtype == np.uint8 or ort_out.dtype == np.int8: + warnings.warn("ONNX output is quantized", stacklevel=2) + # pyrefly: ignore [missing-attribute] + if pt_out.dtype == np.uint8 or pt_out.dtype == np.int8: + warnings.warn("PyTorch output is quantized", stacklevel=2) + raise + + +def _compare_onnx_pytorch_outputs( + onnx_outs: _OutputsType, + pt_outs: Any, + options: VerificationOptions, +) -> None: + """ + Compare ONNX and PyTorch outputs. + + Args: + onnx_outs: outputs from ONNX backend. + pt_outs: outputs from PyTorch. + options: options for verification. + + Raises: + AssertionError: if outputs from ONNX model and PyTorch model are not + equal up to specified precision. + ValueError: if arguments provided are invalid. + """ + if options.ignore_none: + # torch.jit._flatten filters None type + pt_outs, _ = torch.jit._flatten(pt_outs) + else: + pt_outs = _inline_flatten_list([pt_outs], []) + pt_outs_np = _unpack_to_numpy(pt_outs, cast_onnx_accepted=False) + onnx_outs = _inline_flatten_list(onnx_outs, []) + _compare_onnx_pytorch_outputs_in_np(onnx_outs, pt_outs_np, options) + + +def _prepare_input_for_pytorch(args, kwargs): + """Prepare input for PyTorch model execution. + + Any future changes/formatting to the input before dispatching to the PyTorch + model should be made in this function. + + Args: + args: positional arguments for PyTorch model forward method. + kwargs: keyword arguments for PyTorch model forward method. + + Returns: + args: positional arguments for PyTorch model forward method. + kwargs: keyword arguments for PyTorch model forward method. + """ + if isinstance(args, (torch.Tensor, dict)): + args = (args,) + # In-place operators will update input tensor data as well. + # Thus inputs are replicated before every forward call. + args = copy.deepcopy(args) + if kwargs: + kwargs = copy.deepcopy(kwargs) + else: + kwargs = {} + return args, kwargs + + +def _prepare_input_for_export(args, kwargs): + """Prepare input for ONNX model export. + + Any future changes/formatting to the input before dispatching to the + :func:`torch.onnx.export` api should be made in this function. + + Args: + args: positional arguments for PyTorch model forward method. + kwargs: keyword arguments for PyTorch model forward method. + + Returns: + onnx_inputs: positional arguments for ONNX model export, as `args` in + :func:`torch.onnx.export`. + """ + args, kwargs = _prepare_input_for_pytorch(args, kwargs) + if not kwargs and len(args) > 0 and isinstance(args[-1], dict): + onnx_inputs = args + ({},) + elif kwargs: + onnx_inputs = args + (kwargs,) + else: + onnx_inputs = args + return onnx_inputs + + +def _prepare_input_for_onnx( + args, kwargs, remained_onnx_input_idx: Sequence[int] | None, flatten: bool +): + """Prepare input for ONNX model execution in ONNX backend. + + Any future changes/formatting to the input before dispatching to the ONNX backend + run should be made in this function. + + Args: + args: positional arguments for PyTorch model forward method. + kwargs: keyword arguments for PyTorch model forward method. + remained_onnx_input_idx: indices of inputs to be used for ONNX model execution. + flatten: whether to flatten the input before dispatching to the ONNX model execution. + + Returns: + onnx_inputs: positional arguments for ONNX model execution in ONNX backend. + """ + onnx_inputs = _prepare_input_for_export(args, kwargs) + if flatten: + onnx_inputs, _ = torch.jit._flatten(onnx_inputs) + elif onnx_inputs and onnx_inputs[-1] == {}: + # Handle empty kwargs (normally removed by flatten). + onnx_inputs = onnx_inputs[:-1] + if remained_onnx_input_idx is not None: + return [onnx_inputs[i] for i in remained_onnx_input_idx] + else: + return onnx_inputs + + +def _try_clone_model(model): + """Used for preserving original model in case forward mutates model states.""" + try: + return copy.deepcopy(model) + except Exception: + warnings.warn( + "Failed to clone model. Model state might be mutated during verification.", + stacklevel=2, + ) + return model + + +def _compare_onnx_pytorch_model( + pt_model: _ModelType, + onnx_model_f: str | io.BytesIO, + input_args: _InputArgsType, + input_kwargs: _InputKwargsType | None, + additional_test_inputs: Sequence[_InputArgsType] | None, + options: VerificationOptions, +) -> None: + """Compare outputs from ONNX model runs with outputs from PyTorch model runs. + + Args: + pt_model: PyTorch model. + onnx_model_f: ONNX model file path or file-like object. + input_args: positional arguments for PyTorch model forward method. + input_kwargs: keyword arguments for PyTorch model forward method. + additional_test_inputs: additional positional arguments for PyTorch model + forward method. + options: options for verification. + + Raises: + AssertionError: if outputs from ONNX model and PyTorch model are not + equal up to specified precision. + """ + onnx_session = _onnx_backend_session(onnx_model_f, options.backend) + + def compare_onnx_pytorch_model_with_input(input_args, input_kwargs) -> None: + pt_args, pt_kwargs = _prepare_input_for_pytorch(input_args, input_kwargs) + # TODO: remove this and treat mutating model separately. See #77679 + pt_model_copy = _try_clone_model(pt_model) + pt_outs = pt_model_copy(*pt_args, **pt_kwargs) + + onnx_inputs = _prepare_input_for_onnx( + input_args, input_kwargs, options.remained_onnx_input_idx, options.flatten + ) + + onnx_outs = _run_onnx(onnx_session, onnx_inputs) + + _compare_onnx_pytorch_outputs( + onnx_outs=onnx_outs, + pt_outs=pt_outs, + options=options, + ) + + compare_onnx_pytorch_model_with_input(input_args, input_kwargs) + + if additional_test_inputs: + for test_input_args in additional_test_inputs: + compare_onnx_pytorch_model_with_input(test_input_args, {}) + + +def verify( + model: _ModelType, + input_args: _InputArgsType, + input_kwargs: _InputKwargsType | None = None, + do_constant_folding: bool = True, + dynamic_axes: Mapping[str, Mapping[int, str] | Mapping[str, Sequence[int]]] + | None = None, + input_names: Sequence[str] | None = None, + output_names: Sequence[str] | None = None, + training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL, + opset_version: int | None = None, + keep_initializers_as_inputs: bool = True, + verbose: bool = False, + fixed_batch_size: bool = False, + use_external_data: bool = False, + additional_test_inputs: Sequence[_InputArgsType] | None = None, + options: VerificationOptions | None = None, +) -> None: + """Verify model export to ONNX against original PyTorch model. + + .. deprecated:: 2.7 + Consider using ``torch.onnx.export(..., dynamo=True)`` and use the returned + ``ONNXProgram`` to test the ONNX model. + + Args: + model: See :func:`torch.onnx.export`. + input_args: See :func:`torch.onnx.export`. + input_kwargs: See :func:`torch.onnx.export`. + do_constant_folding: See :func:`torch.onnx.export`. + dynamic_axes: See :func:`torch.onnx.export`. + input_names: See :func:`torch.onnx.export`. + output_names: See :func:`torch.onnx.export`. + training: See :func:`torch.onnx.export`. + opset_version: See :func:`torch.onnx.export`. + keep_initializers_as_inputs: See :func:`torch.onnx.export`. + verbose: See :func:`torch.onnx.export`. + fixed_batch_size: Legacy argument, used only by rnn test cases. + use_external_data: Explicitly specify whether to export the model with external data. + additional_test_inputs: List of tuples. Each tuple is a group of + input arguments to test. Currently only ``*args`` are supported. + options: A VerificationOptions object that controls the verification behavior. + + Raises: + AssertionError: if outputs from ONNX model and PyTorch model are not + equal up to specified precision. + ValueError: if arguments provided are invalid. + """ + if options is None: + options = VerificationOptions() + + if training == torch.onnx.TrainingMode.TRAINING: + model.train() + elif training == torch.onnx.TrainingMode.EVAL: + model.eval() + with torch.no_grad(), contextlib.ExitStack() as stack: + model_f: str | io.BytesIO = io.BytesIO() + if use_external_data: + tmpdir_path = stack.enter_context(tempfile.TemporaryDirectory()) + model_f = os.path.join(tmpdir_path, "model.onnx") + + inputs_for_export = _prepare_input_for_export(input_args, input_kwargs) + + # TODO(#77679): remove this and treat mutating model separately. + model_copy = _try_clone_model(model) + utils._export( + model, + inputs_for_export, + model_f, + opset_version=opset_version, + do_constant_folding=do_constant_folding, + keep_initializers_as_inputs=keep_initializers_as_inputs, + dynamic_axes=dynamic_axes, + input_names=input_names, + output_names=output_names, + fixed_batch_size=fixed_batch_size, + training=training, + verbose=verbose, + ) + + _compare_onnx_pytorch_model( + pt_model=model_copy, + onnx_model_f=model_f, + input_args=input_args, + input_kwargs=input_kwargs, + additional_test_inputs=additional_test_inputs, + options=options, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/errors.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/errors.py new file mode 100644 index 0000000000000000000000000000000000000000..3645e01d7a7a2c5ddd8f44a1f2a3883a9b437fd1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/errors.py @@ -0,0 +1,101 @@ +"""ONNX exporter exceptions.""" + +from __future__ import annotations + + +__all__ = [ + "OnnxExporterWarning", + "SymbolicValueError", + "UnsupportedOperatorError", +] + +import textwrap +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from torch import _C + + +class OnnxExporterWarning(UserWarning): + """Warnings in the ONNX exporter.""" + + +class OnnxExporterError(RuntimeError): + """Errors raised by the ONNX exporter. This is the base class for all exporter errors.""" + + +class UnsupportedOperatorError(OnnxExporterError): + """Raised when an operator is unsupported by the exporter.""" + + # NOTE: This is legacy and is only used by the torchscript exporter + # Clean up when the torchscript exporter is removed + def __init__(self, name: str, version: int, supported_version: int | None) -> None: + if supported_version is not None: + msg = ( + f"Exporting the operator '{name}' to ONNX opset version {version} " + "is not supported. Support for this operator was added in version " + f"{supported_version}, try exporting with this version" + ) + elif name.startswith(("aten::", "prim::", "quantized::")): + msg = ( + f"Exporting the operator '{name}' to ONNX opset version {version} " + "is not supported" + ) + else: + msg = ( + f"ONNX export failed on an operator with unrecognized namespace {name}. " + "If you are trying to export a custom operator, make sure you registered it with " + "the right domain and version." + ) + + super().__init__(msg) + + +class SymbolicValueError(OnnxExporterError): + """Errors around TorchScript values and nodes.""" + + # NOTE: This is legacy and is only used by the torchscript exporter + # Clean up when the torchscript exporter is removed + def __init__(self, msg: str, value: _C.Value) -> None: + message = ( + f"{msg} [Caused by the value '{value}' (type '{value.type()}') in the " + f"TorchScript graph. The containing node has kind '{value.node().kind()}'.] " + ) + + code_location = value.node().sourceRange() + if code_location: + message += f"\n (node defined in {code_location})" + + try: + # Add its input and output to the message. + message += "\n\n" + message += textwrap.indent( + ( + "Inputs:\n" + + ( + "\n".join( + f" #{i}: {input_} (type '{input_.type()}')" + for i, input_ in enumerate(value.node().inputs()) + ) + or " Empty" + ) + + "\n" + + "Outputs:\n" + + ( + "\n".join( + f" #{i}: {output} (type '{output.type()}')" + for i, output in enumerate(value.node().outputs()) + ) + or " Empty" + ) + ), + " ", + ) + except AttributeError: + message += ( + " Failed to obtain its input and output for debugging. " + "Please refer to the TorchScript graph for debugging information." + ) + + super().__init__(message) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/operators.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/operators.py new file mode 100644 index 0000000000000000000000000000000000000000..dd31ba8c304043d6b6a6bad92c60d8f729a034f3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/operators.py @@ -0,0 +1,47 @@ +"""This file provides a location for operators that help exporting models via onnx. + +E.g. `shape_as_tensor` and `reshape_from_tensor_shape` +are to make all dynamic sizes operations traceable. + +NOTE: at one point these functions were implemented differently. +Since then we have implemented these directly in ATen, so this +file is kept purely for backward-compatibility. +""" + +from __future__ import annotations + + +__all__: list[str] = [] + +import torch + + +"""Get the shape of a tensor as a tensor. + +Args: + x (Tensor): The input tensor. + +Returns: + Tensor: A tensor of shape [len(x.shape)] containing the size of each dimension of x. + +Example: + >>> x = torch.randn(2, 3) + >>> shape_as_tensor(x) + tensor([2, 3]) + +""" +shape_as_tensor = torch._shape_as_tensor + +"""Reshape a tensor to the given shape. + +This function is used to make dynamic size operations traceable when exporting models via ONNX. +This function is kept for backward-compatibility. It is implemented directly in ATen. + +Parameters: + x (Tensor): the tensor to be reshaped. + shape (Tensor): the target shape. + +Returns: + Tensor: the reshaped tensor. +""" +reshape_from_tensor_shape = torch._reshape_from_tensor diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8da3fc8e587236cb6793ee1123129d3ccf2d49dd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/__init__.py @@ -0,0 +1,467 @@ +"""ONNX operators as native torch.fx operators. + +This module provides a set of functions to create ONNX operators in the FX graph +which are exportable to ONNX. +""" + +# flake8: noqa: B950 +from __future__ import annotations + + +__all__ = [ + "aten_decompositions", + "symbolic", + "symbolic_multi_out", + "rotary_embedding", + "attention", +] + + +from typing import TYPE_CHECKING + +import torch +from torch.onnx.ops import _impl, _symbolic_impl + + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + +# https://github.com/onnx/onnx/blob/f542e1f06699ea7e1db5f62af53355b64338c723/onnx/onnx.proto#L597 +_TORCH_DTYPE_TO_ONNX_DTYPE = { + torch.float32: 1, # FLOAT + torch.uint8: 2, # UINT8 + torch.int8: 3, # INT8 + torch.uint16: 4, # UINT16 + torch.int16: 5, # INT16 + torch.int32: 6, # INT32 + torch.int64: 7, # INT64 + str: 8, # STRING + torch.bool: 9, # BOOL + torch.float16: 10, # FLOAT16 + torch.double: 11, # DOUBLE + torch.uint32: 12, # UINT32 + torch.uint64: 13, # UINT64 + torch.complex64: 14, # COMPLEX64 + torch.complex128: 15, # COMPLEX128 + torch.bfloat16: 16, # BFLOAT16 + torch.float8_e4m3fn: 17, # FLOAT8E4M3FN + torch.float8_e4m3fnuz: 18, # FLOAT8E4M3FNUZ + torch.float8_e5m2: 19, # FLOAT8E5M2 + torch.float8_e5m2fnuz: 20, # FLOAT8E5M2FNUZ + # 21 = UINT4 + # 22 = INT4 + torch.float4_e2m1fn_x2: 23, # FLOAT4E2M1 +} + + +def aten_decompositions() -> dict[torch._ops.OpOverload, Callable]: + """Return the ONNX to ATen decomp table.""" + return _impl.ONNX_ATEN_DECOMP_TABLE + + +def _parse_domain_op_type(domain_op: str) -> tuple[str, str]: + split = domain_op.split("::", 1) + if len(split) == 1: + domain = "" + op_type = split[0] + else: + domain = split[0] + op_type = split[1] + return domain, op_type + + +def symbolic( + domain_op: str, + /, + inputs: Sequence[torch.Tensor | None], + attrs: dict[ + str, + int + | float + | str + | bool + | Sequence[int] + | Sequence[float] + | Sequence[str] + | Sequence[bool], + ] + | None = None, + *, + dtype: torch.dtype | int, + shape: Sequence[int | torch.SymInt], + version: int | None = None, + metadata_props: dict[str, str] | None = None, +) -> torch.Tensor: + """Create a symbolic FX operator to represent an arbitrary ONNX operator. + + This function is used to create a symbolic operator with a single output. + To create an operator with multiple outputs, use :func:`symbolic_multi_out`. + + You may use ``if torch.onnx.is_in_onnx_export()`` to conditionally enable the + symbolic logic only during ``torch.onnx.export()``. + + Example:: + + class CustomOp(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Normal torch operators can interleave with the symbolic ops during ONNX export + x = x + 1 + + # Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain. + # The output tensor will have the specified dtype and shape + val = torch.onnx.ops.symbolic( + "custom_domain::CustomOp", + (x,), + dict(attr_key="attr_value"), + dtype=x.dtype, + shape=x.shape, + version=1, + ) + + # The result of the symbolic op can be used in normal torch operations during ONNX export + return torch.nn.functional.relu(val) + + + # You may then export this model to ONNX using torch.onnx.export(..., dynamo=True). + + Args: + domain_op: The domain and operator name, separated by "::". For example, + "custom_domain::CustomOp". + inputs: The input tensors to the operator. + attrs: The attributes of the operator. The keys are attribute names and + the values are attribute values. Valid attribute types are int, float, + str, bool, and lists of int, float, str, and bool. Tensor attributes + are unsupported. + dtype: The data type of the output tensor.This can be either a torch.dtype + or an integer representing the ONNX data type. + shape: The shape of the output tensor. This can be a list of integers or + SymInt values. + version: The version of the opset used for the operator. + metadata_props: Metadata properties for the ONNX node. + This is a dictionary of str-str pairs. + + Returns: + The output tensor of the operator. + """ + if not isinstance(dtype, int): + torch._check( + dtype in _TORCH_DTYPE_TO_ONNX_DTYPE, lambda: f"Unsupported dtype: {dtype}" + ) + dtype = _TORCH_DTYPE_TO_ONNX_DTYPE[dtype] + domain, op_type = _parse_domain_op_type(domain_op) + if attrs is None: + attrs = {} + encoded_attrs = _symbolic_impl.EncodedAttrs.from_dict(attrs) + # TODO: Parse domain + return _symbolic_impl._symbolic( + inputs, + op_type, + dtype, + shape=shape, + attr_keys=encoded_attrs.attr_keys, + attr_types=encoded_attrs.attr_types, + attr_pos=encoded_attrs.attr_pos, + attr_ints=encoded_attrs.attr_ints, + attr_floats=encoded_attrs.attr_floats, + attr_strs=encoded_attrs.attr_strs, + metadata_props_keys=metadata_props.keys() if metadata_props else [], + metadata_props_values=metadata_props.values() if metadata_props else [], + domain=domain, + version=version, + ) + + +def symbolic_multi_out( + domain_op: str, + /, + inputs: Sequence[torch.Tensor | None], + attrs: dict[ + str, + int + | float + | str + | bool + | Sequence[int] + | Sequence[float] + | Sequence[str] + | Sequence[bool], + ] + | None = None, + *, + dtypes: Sequence[torch.dtype | int], + shapes: Sequence[Sequence[int | torch.SymInt]], + version: int | None = None, + metadata_props: dict[str, str] | None = None, +) -> Sequence[torch.Tensor]: + """Create a symbolic FX operator to represent an arbitrary ONNX operator with multiple outputs. + + You may use ``if torch.onnx.is_in_onnx_export()`` to conditionally enable the + symbolic logic only during ``torch.onnx.export()``. + + Example:: + + class CustomOp(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Normal torch operators can interleave with the symbolic ops during ONNX export + x = x + 1 + + # Create a symbolic ONNX operator with the name "CustomOp" in the "custom_domain" domain. + # The output tensors will have the specified dtypes and shapes + (out1, out2) = torch.onnx.ops.symbolic_multi_out( + "custom_domain::CustomOp", + (x,), + dict(attr_key="attr_value"), + dtypes=(x.dtype, torch.float32), + shapes=(x.shape, [1, 2, 3]), + version=1, + ) + + # The result of the symbolic op can be used in normal torch operations during ONNX export + return torch.nn.functional.relu(out1 + out2) + + + # You may then export this model to ONNX using torch.onnx.export(..., dynamo=True). + + Args: + domain_op: The domain and operator name, separated by "::". For example, + "custom_domain::CustomOp". + inputs: The input tensors to the operator. + attrs: The attributes of the operator. The keys are attribute names and + the values are attribute values. Valid attribute types are int, float, + str, bool, and lists of int, float, str, and bool. Tensor attributes + are unsupported. + dtypes: The data types of the output tensors. This can be a list of + torch.dtype or integers representing the ONNX data types. The length + of this list must be the number of outputs. + shapes: The shapes of the output tensors. This can be a list of lists of + integers or SymInt values. The length of this list must be the number of outputs. + version: The version of the opset used for the operator. + metadata_props: Metadata properties for the ONNX node. + This is a dictionary of str-str pairs. + + Returns: + A list of output tensors of the operator. + """ + torch._check( + len(shapes) == len(dtypes), + lambda: f"Number of shapes ({len(shapes)}) must match number of dtypes ({len(dtypes)})", + ) + onnx_dtypes = [] + for dtype in dtypes: + if not isinstance(dtype, int): + torch._check( + dtype in _TORCH_DTYPE_TO_ONNX_DTYPE, + lambda: f"Unsupported dtype: {dtype}", + ) + onnx_dtypes.append(_TORCH_DTYPE_TO_ONNX_DTYPE[dtype]) + else: + onnx_dtypes.append(dtype) + domain, op_type = _parse_domain_op_type(domain_op) + if attrs is None: + attrs = {} + encoded_attrs = _symbolic_impl.EncodedAttrs.from_dict(attrs) + # Use the size of dtypes to determine the number of outputs + return _symbolic_impl._symbolic_multi_out( + inputs, + op_type, + onnx_dtypes, + shapes=shapes, + attr_keys=encoded_attrs.attr_keys, + attr_types=encoded_attrs.attr_types, + attr_pos=encoded_attrs.attr_pos, + attr_ints=encoded_attrs.attr_ints, + attr_floats=encoded_attrs.attr_floats, + attr_strs=encoded_attrs.attr_strs, + metadata_props_keys=metadata_props.keys() if metadata_props else [], + metadata_props_values=metadata_props.values() if metadata_props else [], + domain=domain, + version=version, + ) + + +def rotary_embedding( + X: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + position_ids: torch.Tensor | None = None, + *, + interleaved: bool = False, + num_heads: int = 0, + rotary_embedding_dim: int = 0, +) -> torch.Tensor: + """RotaryEmbedding op in ONNX. + + https://onnx.ai/onnx/operators/onnx__RotaryEmbedding.html + + RotaryEmbedding is the implementation of rotary positional embeddings (RoPE) based on the paper https://arxiv.org/pdf/2104.09864. + The key advantage of RoPE is that it allows the model to understand both the absolute position of a token and the relative distances + between tokens. This is achieved through a rotational mechanism where the extent of rotation is computed based on the token's absolute position (position_ids). + + The rotational mechanism is defined by sine and cosine functions that are used to represent the rotation angles. + For each token in the sequence, its positional embedding is computed by rotating its embedding vector. This is done by splitting the + embedding vector either into two halves or interleaving every alternate token and applying the rotation matrix to each half of the embedding vector. + The rotation matrix is parameterized by the token's position in the sequence. The rotated halves of the embedding vector are concatenated + to form the final positional embedding for each token. The rotated positional embeddings are used in the self-attention mechanism. + The rotation ensures that the model captures both absolute and relative positional information. + + Args: + X: The input tensor representing the token embeddings. 4D tensor with + shape `(batch_size, num_heads, sequence_length, head_size)` or 3D tensor + with shape `(batch_size, sequence_length, hidden_size)`. For cases with + a 4D input tensor, `head_size` has to be even. For cases with a 3D input + tensor, `num_heads` attribute must be provided and `hidden_size` must + be an even multiple of `num_heads` where `hidden_size = num_heads * head_size` + cos_cache: The cosine values for the rotation. 2D tensor with shape `(max_position_id_plus_1, head_size / 2)` + for full rotation or `(max_position_id_plus_1, rotary_embedding_dim / 2)` + for partial rotation when `position_ids` are provided. 3D tensor with shape + `(batch_size, sequence_length, head_size / 2)` for full rotation or + `(batch_size, sequence_length, rotary_embedding_dim / 2)` for partial + rotation when `position_ids` are not provided. `max_position_id_plus_1` + is a parameter to the model. + sin_cache: The sine values for the rotation. 2D tensor with shape `(max_position_id_plus_1, head_size / 2)` + for full rotation or `(max_position_id_plus_1, rotary_embedding_dim / 2)` + for partial rotation when `position_ids` are provided. 3D tensor with shape + `(batch_size, sequence_length, head_size / 2)` for full rotation or + `(batch_size, sequence_length, rotary_embedding_dim / 2)` for partial rotation + when `position_ids` are not provided. `max_position_id_plus_1` is a parameter + to the model. + position_ids: The position indices for the tokens. 2D tensor with shape + `(batch_size, sequence_length)`. + interleaved: Rotate using interleaved pattern. Default value is 0 (False). + num_heads: Number of attention heads. Must be provided when input is a 3D tensor. + rotary_embedding_dim: Rotary embedding dimension used to apply partial rotary embeddings. + + Returns: + Tensor with same shape as input. + """ + return _impl.rotary_embedding_23( + X, + cos_cache, + sin_cache, + position_ids=position_ids, + interleaved=interleaved, + num_heads=num_heads, + rotary_embedding_dim=rotary_embedding_dim, + ) + + +def attention( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + attn_mask: torch.Tensor | None = None, + past_key: torch.Tensor | None = None, + past_value: torch.Tensor | None = None, + *, + is_causal: bool = False, + kv_num_heads: int = 0, + q_num_heads: int = 0, + qk_matmul_output_mode: int = 0, + scale: float | None = None, + softcap: float = 0.0, + softmax_precision: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Attention op in ONNX. + + https://onnx.ai/onnx/operators/onnx__Attention.html + + Computes scaled dot product attention on query, key and value tensors, using an optional attention mask if passed. + + This operator covers self and cross variants of the attention operation based on sequence lengths of K, Q and V. + + For self attention, ``kv_sequence_length`` equals to ``q_sequence_length``. + + For cross attention, query and key might have different lengths. + + This operator also covers the 3 following variants based on the number of heads: + + 1. Multi-headed Attention (MHA): Described in the paper https://arxiv.org/pdf/1706.03762, `q_num_heads = kv_num_heads`. + 2. Group-query Attention (GQA): Described in the paper https://arxiv.org/pdf/2305.13245, `q_num_heads > kv_num_heads`, `q_num_heads % kv_num_heads == 0`. + 3. Multi-query Attention (MQA): Described in the paper https://arxiv.org/pdf/1911.02150, `q_num_heads > kv_num_heads`, `kv_num_heads=1`. + + Attention bias to be added is calculated based on ``attn_mask`` input and ``is_causal` `attribute``, only one of which can be provided. + + 1. If ``is_causal`` is set to `1`, the attention masking is a lower triangular matrix when the mask is a square matrix. The attention masking has the form of the upper left causal bias due to the alignment. + 2. `attn_mask`: A boolean mask where a value of `True` indicates that the element should take part in attention or a float mask of the same type as query, key, value that is added to the attention score. + + Both past and present state key/values are optional. They shall be used together, and not allowed to use only one of them. + The following pattern is applied to the Q, K and V inputs after appropriate reshaping of K and V inputs based on sequence lengths and num heads provided:: + + The following pattern is applied by this operator: + Q K V + | | | + Q*sqrt(scale) K*sqrt(scale) | + | | | + | Transpose | + | | | + ---MatMul--- | + | | + at_mask---Add | + | | + softcap (if provided) | + | | + Softmax | + | | + -----MatMul------ + | + Y + + Args: + Q: Query tensor. 4D tensor with shape `(batch_size, q_num_heads, q_sequence_length, head_size)` or 3D tensor + with shape `(batch_size, q_sequence_length, q_hidden_size)`. For cases with a 3D input tensor, + `q_hidden_size = q_num_heads * head_size` + K: Key tensor. 4D tensor with shape `(batch_size, kv_num_heads, kv_sequence_length, head_size)` or 3D tensor + with shape `(batch_size, kv_sequence_length, k_hidden_size)`. For cases with a 3D input tensor, + `k_hidden_size = kv_num_heads * head_size` + V: Value tensor. 4D tensor with shape `(batch_size, kv_num_heads, kv_sequence_length, v_head_size)` or 3D tensor + with shape `(batch_size, kv_sequence_length, v_hidden_size)`. For cases with a 3D input tensor, + `v_hidden_size = kv_num_heads * v_head_size` + attn_mask: Attention mask. Shape must be broadcastable to 4D tensor with shape + `(batch_size, q_num_heads, q_sequence_length, total_sequence_length)` where + `total_sequence_length = past_sequence_length + kv_sequence_length`. Two types of masks are supported. + A boolean mask where a value of True indicates that the element should take part in attention. + Also supports a float mask of the same type as query, key, value that is added to the attention score. + past_key: Past state cache for key with shape `(batch_size, kv_num_heads, past_sequence_length, head_size)` + past_value: Past state cache for value with shape `(batch_size, kv_num_heads, past_sequence_length, v_head_size)` + is_causal: If set to True, the attention masking is a lower triangular matrix when the mask is a square matrix. + The attention masking has the form of the upper left causal bias due to the alignment. + kv_num_heads: Number of heads of key and value. Must be used with 3D inputs of Q, K and V. + q_num_heads: Number of heads of query. Must be used with 3D inputs of Q, K and V. + qk_matmul_output_mode: If set to 0, qk_matmul_output is the output of qk matmul. If set to 1, + qk_matmul_output includes the addition of the attention mask to the output of qk matmul. + If set to 2, qk_matmul_output is the output after the softcap operation. If set to 3, + qk_matmul_output is the output after the softmax operation. Default value is 0. + scale: Scaling factor applied to Q*K^T. Default value is 1/sqrt(head_size). To prevent numerical overflow, + scale Q, K by sqrt(scale) before matmul. + softcap: Softcap value for attention weights. Default value is 0. + softmax_precision: The floating-point precision used in softmax computation. If softmax precision is not provided, + the same precision as the input of softmax (Q and K) is used. + + Returns: + A tuple containing: + - The output tensor. 4D tensor with shape `(batch_size, q_num_heads, q_sequence_length, v_head_size)` or 3D tensor + with shape `(batch_size, q_sequence_length, hidden_size)`. For cases with a 3D input tensor, + `hidden_size = q_num_heads * v_head_size` + - Updated key cache with shape `(batch_size, kv_num_heads, total_sequence_length, head_size)` where + `total_sequence_length = past_sequence_length + kv_sequence_length`. + - Updated value cache with shape `(batch_size, kv_num_heads, total_sequence_length, v_head_size)` where + `total_sequence_length = past_sequence_length + kv_sequence_length`. + - The output of QK matmul. 4D tensor with shape `(batch_size, q_num_heads, q_sequence_length, total_sequence_length)` + where `total_sequence_length = past_sequence_length + kv_sequence_length`. + """ + return _impl.attention_23( + Q, + K, + V, + attn_mask=attn_mask, + past_key=past_key, + past_value=past_value, + is_causal=is_causal, + kv_num_heads=kv_num_heads, + q_num_heads=q_num_heads, + qk_matmul_output_mode=qk_matmul_output_mode, + scale=scale, + softcap=softcap, + softmax_precision=softmax_precision, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/_dtype_mappings.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/_dtype_mappings.py new file mode 100644 index 0000000000000000000000000000000000000000..0023e356d89f1e27659c198877869c030293a660 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/_dtype_mappings.py @@ -0,0 +1,27 @@ +import torch + + +ONNX_DTYPE_TO_TORCH_DTYPE: dict[int, torch.dtype] = { + 1: torch.float32, # FLOAT + 2: torch.uint8, # UINT8 + 3: torch.int8, # INT8 + 4: torch.uint16, # UINT16 + 5: torch.int16, # INT16 + 6: torch.int32, # INT32 + 7: torch.int64, # INT64 + 9: torch.bool, # BOOL + 10: torch.float16, # FLOAT16 + 11: torch.double, # DOUBLE + 12: torch.uint32, # UINT32 + 13: torch.uint64, # UINT64 + 14: torch.complex64, # COMPLEX64 + 15: torch.complex128, # COMPLEX128 + 16: torch.bfloat16, # BFLOAT16 + 17: torch.float8_e4m3fn, # FLOAT8E4M3FN + 18: torch.float8_e4m3fnuz, # FLOAT8E4M3FNUZ + 19: torch.float8_e5m2, # FLOAT8E5M2 + 20: torch.float8_e5m2fnuz, # FLOAT8E5M2FNUZ + 21: torch.uint8, # UINT4 + 22: torch.uint8, # INT4 + 23: torch.float4_e2m1fn_x2, # FLOAT4E2M1 +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/_impl.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..be50d953b619b4c65ac798208d981dd05742ced1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/_impl.py @@ -0,0 +1,547 @@ +"""Implementations of ONNX operators as native Torch ops. + +NOTE: Fake implementations: + Refer to https://docs.pytorch.org/docs/stable/library.html#torch.library.register_fake + for more details on how to create fake kernels. +""" + +# flake8: noqa: B950 +import math +from collections.abc import Callable +from typing import TypeVar +from typing_extensions import ParamSpec + +import torch +from torch.onnx.ops import _dtype_mappings + + +# Use ParamSpec for better type preservation instead of bound Callable TypeVar +_P = ParamSpec("_P") +_R = TypeVar("_R") + +# ONNX to ATen decomp table +ONNX_ATEN_DECOMP_TABLE: dict[torch._ops.OpOverload, Callable] = {} +_ATTENTION_23_ALLOWED_INTERMEDIATE_PRECISIONS = frozenset( + { + 1, # FLOAT + 10, # FLOAT16 + 11, # DOUBLE + 16, # BFLOAT16 + } +) + + +def _onnx_op( + op_type: str, opset_version: int, fake_impl: Callable[_P, _R] +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: + """Decorator to register an ONNX operator with a custom implementation.""" + + def decorator(func: Callable[_P, _R]) -> Callable[_P, _R]: + overload = f"opset{opset_version}" + torch_op = torch.library.custom_op( + f"onnx::{op_type}.{overload}", mutates_args=() + )(func) + ONNX_ATEN_DECOMP_TABLE[getattr(getattr(torch.ops.onnx, op_type), overload)] = ( + func # type: ignore[assignment] + ) + torch_op.register_fake(fake_impl) + return torch_op # type: ignore[return-value] + + return decorator + + +def _rotary_embedding_23_fake_impl( + x: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + position_ids: torch.Tensor | None = None, + *, + interleaved: bool = False, + num_heads: int = 0, + rotary_embedding_dim: int = 0, +) -> torch.Tensor: + """Fake implementation for RotaryEmbedding-23 for torch.compile purposes.""" + return x.clone() + + +@_onnx_op("RotaryEmbedding", 23, _rotary_embedding_23_fake_impl) +def rotary_embedding_23( + x: torch.Tensor, + cos_cache: torch.Tensor, + sin_cache: torch.Tensor, + position_ids: torch.Tensor | None = None, + *, + interleaved: bool = False, + num_heads: int = 0, + rotary_embedding_dim: int = 0, +) -> torch.Tensor: + """RotaryEmbedding-23 https://onnx.ai/onnx/operators/onnx__RotaryEmbedding.html#rotaryembedding-23""" + # x has shape (batch_size, num_heads, sequence_length, head_size) + # or (batch_size, sequence_length, hidden_size) + input_shape = x.shape + input_rank = len(input_shape) + batch_size = input_shape[0] + sequence_length = input_shape[-2] + + # Validate position_ids and caches match x + if position_ids is not None: + torch._check( + position_ids.dim() == 2, + lambda: f"position_ids must be 2D when provided. Received shape {position_ids.shape}", + ) + torch._check( + position_ids.shape[0] == batch_size, + lambda: f"position_ids first dim (batch) must match x.shape[0] ({batch_size}). Received {position_ids.shape[0]}", + ) + torch._check( + position_ids.shape[1] == sequence_length, + lambda: f"position_ids second dim (sequence) must match x.shape[-2] ({sequence_length}). Received {position_ids.shape[1]}", + ) + torch._check( + cos_cache.dim() == 2 and sin_cache.dim() == 2, + lambda: "cos_cache/sin_cache must be 2D when position_ids is provided. " + f"Received cos_cache shape {cos_cache.shape}, sin_cache shape {sin_cache.shape}", + ) + else: + torch._check( + cos_cache.dim() == 3 and sin_cache.dim() == 3, + lambda: "cos_cache/sin_cache must be 3D when position_ids is not provided. " + f"Received cos_cache shape {cos_cache.shape}, sin_cache shape {sin_cache.shape}", + ) + + # First ensure x has shape [batch_size, num_heads, seq_len, head_size] + # So that the rotation logic can be shared with reshaped 3D inputs + if input_rank == 4: + # Reshape from (batch_size, num_heads, seq_len, head_size) + # to [batch_size, seq_len, num_heads, head_size] + x = torch.permute(x, (0, 2, 1, 3)) + elif input_rank == 3: + torch._check( + num_heads != 0, + lambda: f"num_heads must be provided for 3D inputs. Received input tensor with shape {input_shape}", + ) + hidden_size = input_shape[2] + head_size = hidden_size // num_heads + new_shape = [batch_size, sequence_length, num_heads, head_size] + x = torch.reshape(x, new_shape) + + torch._check(len(x.shape) == 4, lambda: "x should be a 4D tensor by now") + head_size = x.shape[3] + + # Fully or partially perform rotation on x based on rotary_embedding_dim attribute + if rotary_embedding_dim == 0: + # If rotary_embedding_dim not provided, perform full rotation by using head_size + rotary_embedding_dim = head_size + x_rotate = x[:, :, :, :rotary_embedding_dim] + x_not_rotate = x[:, :, :, rotary_embedding_dim:] + rotary_embedding_dim_half = rotary_embedding_dim // 2 + + # Retrieve sin and cos caches using position ids + if position_ids is not None: + cos = cos_cache[ + position_ids + ] # Shape: [batch_size, sequence_length, head_size/2] + sin = sin_cache[ + position_ids + ] # Shape: [batch_size, sequence_length, head_size/2] + else: + cos = cos_cache # Shape: [batch_size, sequence_length, rotary_embedding_dim/2] + sin = sin_cache # Shape: [batch_size, sequence_length, rotary_embedding_dim/2] + + torch._check( + cos.shape[0] == batch_size and cos.shape[1] == sequence_length, + lambda: f"cos has shape {cos.shape} but expected (batch={batch_size}, seq={sequence_length}, ...)", + ) + torch._check( + sin.shape[0] == batch_size and sin.shape[1] == sequence_length, + lambda: f"sin has shape {sin.shape} but expected (batch={batch_size}, seq={sequence_length}, ...)", + ) + torch._check( + cos.shape[-1] == rotary_embedding_dim_half, + lambda: f"Last dimension of cos cache ({cos.shape[-1]}) should match rotary_embedding_dim/2 ({rotary_embedding_dim_half}).", + ) + torch._check( + sin.shape[-1] == rotary_embedding_dim_half, + lambda: f"Last dimension of sin cache ({sin.shape[-1]}) should match rotary_embedding_dim/2 ({rotary_embedding_dim_half}).", + ) + cos = torch.unsqueeze( + cos, 2 + ) # Shape: [batch_size, sequence_length, 1, rotary_embedding_dim/2] + sin = torch.unsqueeze( + sin, 2 + ) # Shape: [batch_size, sequence_length, 1, rotary_embedding_dim/2] + + # Either divide the x in halves or interleave (based on interleaved attribute) + if interleaved: + x1 = x_rotate[:, :, :, 0::2] + x2 = x_rotate[:, :, :, 1::2] + else: + x1, x2 = torch.chunk(x_rotate, 2, dim=-1) + + # Calculate real and imaginary values + real = cos * x1 - sin * x2 + imag = sin * x1 + cos * x2 + + # Inserted rotated embeddings back to the original x + if interleaved: + # x_rotate[:, :, :, 0::2] = real + # x_rotate[:, :, :, 1::2] = imag + real = torch.unsqueeze(real, -1) + imag = torch.unsqueeze(imag, -1) + x_rotate_concat = torch.cat((real, imag), dim=-1) + x_rotate = torch.reshape(x_rotate_concat, x_rotate.shape) + else: + x_rotate = torch.cat((real, imag), dim=-1) + output = torch.cat((x_rotate, x_not_rotate), dim=-1) + if input_rank == 3: + return torch.reshape(output, input_shape) + + # Return the dimensions to the original order + return torch.permute(output, (0, 2, 1, 3)) + + +def _get_scale_factor(scale: float | None, head_size: int) -> float: + """Get the scale factor for attention computation.""" + return scale if scale is not None else (1.0 / math.sqrt(head_size)) + + +def _reshape_3d_to_4d( + tensor: torch.Tensor, batch_size: int, num_heads: int +) -> torch.Tensor: + """Reshape 3D tensor to 4D for multi-head attention.""" + sequence_length, hidden_size = tensor.shape[1], tensor.shape[2] + head_size = hidden_size // num_heads + return ( + tensor.view(batch_size, sequence_length, num_heads, head_size) + .transpose(1, 2) + .contiguous() + ) + + +def _get_qk_output_for_aten_spda( + Q: torch.Tensor, + K: torch.Tensor, + current_q_num_heads: int, + current_kv_num_heads: int, + scale: float | None, + qk_matmul_output_mode: int, +) -> torch.Tensor: + """Get QK output tensor based on the specified mode.""" + if qk_matmul_output_mode == 0: + return _compute_qk_output_for_mode_0( + Q, K, current_q_num_heads, current_kv_num_heads, scale + ) + else: + # For other modes, return a zero tensor with correct shape + return torch.zeros_like(torch.matmul(Q, K.transpose(-2, -1))) + + +def _validate_gqa_configuration( + current_q_num_heads: int, current_kv_num_heads: int +) -> None: + """Validate Group Query Attention configuration.""" + torch._check( + current_q_num_heads % current_kv_num_heads == 0, + lambda: f"q_num_heads ({current_q_num_heads}) must be divisible by kv_num_heads ({current_kv_num_heads}) for GQA", + ) + + +def _compute_qk_output_for_mode_0( + Q: torch.Tensor, + K: torch.Tensor, + current_q_num_heads: int, + current_kv_num_heads: int, + scale: float | None, +) -> torch.Tensor: + """Helper function to compute QK output for qk_matmul_output_mode == 0.""" + # Handle GQA manually for QK output + K_for_qk = K + if current_q_num_heads != current_kv_num_heads: + repeat_factor = current_q_num_heads // current_kv_num_heads + K_for_qk = K.repeat_interleave(repeat_factor, dim=1) + + scale_factor = _get_scale_factor(scale, Q.shape[3]) + # Scale both Q and K by sqrt(scale_factor) for numerical stability + sqrt_scale = math.sqrt(scale_factor) + Q_scaled = Q * sqrt_scale + K_scaled = K_for_qk * sqrt_scale + return torch.matmul(Q_scaled, K_scaled.transpose(-2, -1)) + + +def _attention_23_fake_impl( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + attn_mask: torch.Tensor | None = None, + past_key: torch.Tensor | None = None, + past_value: torch.Tensor | None = None, + *, + is_causal: bool = False, + kv_num_heads: int = 0, + q_num_heads: int = 0, + qk_matmul_output_mode: int = 0, + scale: float | None = None, + softcap: float = 0.0, + softmax_precision: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Fake implementation for Attention-23 for torch.compile purposes.""" + batch_size = Q.shape[0] + + # Handle 3D vs 4D input shapes + if len(Q.shape) == 3: + # 3D input: (batch_size, sequence_length, hidden_size) + q_sequence_length = Q.shape[1] + output_shape = Q.shape # Same shape as Q for 3D output + + # For present_key and present_value, we need 4D shapes + if past_key is not None: + present_key_shape = ( + batch_size, + kv_num_heads, + past_key.shape[2] + K.shape[1], # Combined sequence length + K.shape[2] // kv_num_heads, # head_size + ) + else: + present_key_shape = ( + batch_size, + kv_num_heads, + K.shape[1], # sequence_length + K.shape[2] // kv_num_heads, # head_size + ) + present_value_shape = present_key_shape # Same shape as present_key + + # QK output shape for 3D input (reshaped to 4D internally) + qk_output_shape = ( + batch_size, + q_num_heads, + q_sequence_length, + present_key_shape[2], # kv_sequence_length + ) + else: + # 4D input: (batch_size, num_heads, sequence_length, head_size) + q_sequence_length = Q.shape[2] + # Same shape as Q for 4D output + output_shape = Q.shape # type: ignore[assignment] + + # Handle past key/value concatenation + if past_key is not None: + present_key_shape = ( + K.shape[0], # batch_size + K.shape[1], # num_heads + past_key.shape[2] + K.shape[2], # Combined sequence length + K.shape[3], # head_size + ) + else: + present_key_shape = K.shape # type: ignore[assignment] + present_value_shape = present_key_shape # Same shape as present_key + + # QK output shape + qk_output_shape = ( + Q.shape[0], # batch_size + Q.shape[1], # q_num_heads + Q.shape[2], # q_sequence_length + present_key_shape[2], # kv_sequence_length + ) + + # Create fake tensors with correct shapes and dtypes + output = torch.empty(output_shape, dtype=Q.dtype, device=Q.device) + present_key = torch.empty(present_key_shape, dtype=K.dtype, device=K.device) + present_value = torch.empty(present_value_shape, dtype=V.dtype, device=V.device) + qk_output = torch.empty(qk_output_shape, dtype=Q.dtype, device=Q.device) + + return output, present_key, present_value, qk_output + + +@_onnx_op("Attention", 23, _attention_23_fake_impl) +def attention_23( + Q: torch.Tensor, + K: torch.Tensor, + V: torch.Tensor, + attn_mask: torch.Tensor | None = None, + past_key: torch.Tensor | None = None, + past_value: torch.Tensor | None = None, + *, + is_causal: bool = False, + kv_num_heads: int = 0, + q_num_heads: int = 0, + qk_matmul_output_mode: int = 0, + scale: float | None = None, + softcap: float = 0.0, + softmax_precision: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Attention-23 https://onnx.ai/onnx/operators/onnx__Attention.html#attention-23""" + + num_head_dim, sequence_dim, head_dim = 1, 2, 3 + + # Store original input shape to determine output shape + input_shape_len = len(Q.shape) + batch_size = Q.shape[0] + + # Reshape 3D inputs to 4D format + if len(Q.shape) == 3: + torch._check( + q_num_heads != 0 and kv_num_heads != 0, + lambda: "q_num_heads and kv_num_heads must be provided for 3D inputs", + ) + q_sequence_length = Q.shape[1] + Q = _reshape_3d_to_4d(Q, batch_size, q_num_heads) + K = _reshape_3d_to_4d(K, batch_size, kv_num_heads) + V = _reshape_3d_to_4d(V, batch_size, kv_num_heads) + + torch._check( + len(Q.shape) == 4 and len(K.shape) == 4 and len(V.shape) == 4, + lambda: "Q, K, and V should be 4D tensors by now", + ) + + # Calculate scale factor if not provided + q_head_size = Q.shape[head_dim] + scale = _get_scale_factor(scale, q_head_size) + + # Handle past key/value caches + present_key = ( + torch.cat([past_key, K], dim=sequence_dim) + if past_key is not None + else K.clone() + ) + present_value = ( + torch.cat([past_value, V], dim=sequence_dim) + if past_value is not None + else V.clone() + ) + + # Update K and V to include past states + K, V = present_key, present_value + + # Get current dimensions + current_q_num_heads = Q.shape[num_head_dim] + current_kv_num_heads = K.shape[num_head_dim] + q_sequence_length = Q.shape[sequence_dim] + kv_sequence_length = K.shape[sequence_dim] + + # Check if we can use the optimized scaled_dot_product_attention (most optimized) + can_use_sdpa = ( + softcap == 0.0 # No softcap + and qk_matmul_output_mode == 0 # Default QK output mode + and softmax_precision is None # No custom softmax precision + and (attn_mask is None or attn_mask.dtype == torch.bool) + ) + + _validate_gqa_configuration(current_q_num_heads, current_kv_num_heads) + + if can_use_sdpa: + # Use PyTorch's optimized scaled_dot_product_attention + output = torch.nn.functional.scaled_dot_product_attention( + Q, + K, + V, + attn_mask=attn_mask, + dropout_p=0.0, + is_causal=is_causal, + scale=scale, + enable_gqa=bool( + current_q_num_heads != current_kv_num_heads + ), # Ensure enable_gqa is not SymBool + ) + + qk_output = _get_qk_output_for_aten_spda( + Q, + K, + current_q_num_heads, + current_kv_num_heads, + scale, + qk_matmul_output_mode, + ) + else: + # Fallback to manual implementation for complex cases + + # Handle Group Query Attention (GQA) and Multi-Query Attention (MQA) + if current_q_num_heads != current_kv_num_heads: + repeat_factor = current_q_num_heads // current_kv_num_heads + K = K.repeat_interleave(repeat_factor, dim=num_head_dim) + V = V.repeat_interleave(repeat_factor, dim=num_head_dim) + + # Create attention bias + attn_bias = torch.zeros( + q_sequence_length, kv_sequence_length, dtype=Q.dtype, device=Q.device + ) + + # Apply causal masking + if is_causal: + torch._check( + attn_mask is None, lambda: "Cannot use both is_causal and attn_mask" + ) + causal_mask = torch.tril( + torch.ones( + q_sequence_length, + kv_sequence_length, + dtype=torch.bool, + device=Q.device, + ) + ) + attn_bias = attn_bias.masked_fill(~causal_mask, float("-inf")) + + # Apply attention mask + if attn_mask is not None: + if attn_mask.dtype == torch.bool: + # Boolean mask: True means participate in attention + attn_bias = attn_bias.masked_fill(~attn_mask, float("-inf")) + else: + # Float mask: added to attention scores + attn_bias = attn_bias + attn_mask + + # Apply scaling factor + scale_factor = _get_scale_factor(scale, Q.shape[3]) + + # Scale both Q and K by sqrt(scale_factor) for numerical stability + sqrt_scale = math.sqrt(scale_factor) + Q_scaled = Q * sqrt_scale + K_scaled = K * sqrt_scale + + # Compute Q @ K^T + qk_matmul_output = torch.matmul(Q_scaled, K_scaled.transpose(-2, -1)) + + # Initialize QK output based on mode + qk_output = qk_matmul_output # Default case for mode 0 + + # Add attention bias + qk_with_bias = qk_matmul_output + attn_bias + + if qk_matmul_output_mode == 1: + qk_output = qk_with_bias + + # Apply softcap if provided + if softcap > 0.0: + qk_with_bias = softcap * torch.tanh(qk_with_bias / softcap) + + if qk_matmul_output_mode == 2: + qk_output = qk_with_bias + + # Apply softmax with optional precision casting + if softmax_precision is not None: + # Map ONNX data type to torch dtype + if softmax_precision in _ATTENTION_23_ALLOWED_INTERMEDIATE_PRECISIONS: + original_dtype = qk_with_bias.dtype + qk_with_bias = qk_with_bias.to( + _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[softmax_precision] + ) + qk_softmax = torch.softmax(qk_with_bias, dim=-1) + qk_softmax = qk_softmax.to(original_dtype) + else: + qk_softmax = torch.softmax(qk_with_bias, dim=-1) + else: + qk_softmax = torch.softmax(qk_with_bias, dim=-1) + + if qk_matmul_output_mode == 3: + qk_output = qk_softmax + + # Compute attention output + output = torch.matmul(qk_softmax, V) + + # Reshape output back to 3D if input was 3D + if input_shape_len == 3: + # output: (batch_size, q_num_heads, q_sequence_length, v_head_size) -> (batch_size, q_sequence_length, hidden_size) + output = ( + output.transpose(1, 2).contiguous().view(batch_size, q_sequence_length, -1) + ) + + return output, present_key, present_value, qk_output diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/_symbolic_impl.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/_symbolic_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..85963dd85b1d2127b97b28155abaf4e9fb75a401 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/ops/_symbolic_impl.py @@ -0,0 +1,304 @@ +"""Implementation of symbolic FX ops to represent arbitrary ONNX ops. + +This module provides a way to create symbolic FX operators that can represent +arbitrary ONNX operators. + +The operators are called "symbolic" because they don't do any actual computation +but instead serve as placeholders in the computation graph. + +Each implementation contains two parts: A "real" implementation that produce all +zeros based on the input shape and dtype, and a "fake" implementation that does more +or less the same thing but is required by the `torch.library.custom_op` interface. +""" + +# flake8: noqa: B950 +import dataclasses +from collections.abc import Sequence + +import torch +from torch.onnx.ops import _dtype_mappings + + +_INT_TYPE = "i" +_FLOAT_TYPE = "f" +_STRING_TYPE = "s" +_INT_SEQ_TYPE = "is" +_FLOAT_SEQ_TYPE = "fs" +_STRING_SEQ_TYPE = "ss" + + +@dataclasses.dataclass +class EncodedAttrs: + """Class to encode attributes from dictionary into lists of FX compatible attributes. + + Since FX does not support dictionaries, we need to encode the attributes into + lists. This class provides a way to encode and decode the attributes. + + Attributes: + attr_keys: List of attribute keys. + attr_types: List of attribute types. Values can be "i" (int), "f" (float), + "s" (string), "is" (int sequence), "fs" (float sequence), or "ss" (string sequence). + attr_pos: List of tuples representing the start and end positions of each + attribute in the corresponding list. + attr_ints: List of integer attributes. + attr_floats: List of float attributes. + attr_strs: List of string attributes. + """ + + attr_keys: list[str] + attr_types: list[str] + attr_pos: list[tuple[int, int]] + attr_ints: list[int] + attr_floats: list[float] + attr_strs: list[str] + + @classmethod + def from_dict( + cls, + attrs: dict[ + str, + int + | float + | str + | bool + | Sequence[int] + | Sequence[float] + | Sequence[str] + | Sequence[bool], + ], + ) -> "EncodedAttrs": + encoded = cls( + attr_keys=[], + attr_types=[], + attr_pos=[], + attr_ints=[], + attr_floats=[], + attr_strs=[], + ) + for k, v in attrs.items(): + encoded.attr_keys.append(k) + if isinstance(v, int): + start_pos = len(encoded.attr_ints) + encoded.attr_ints.append(v) + encoded.attr_pos.append((start_pos, start_pos + 1)) + encoded.attr_types.append(_INT_TYPE) + elif isinstance(v, float): + start_pos = len(encoded.attr_floats) + encoded.attr_floats.append(v) + encoded.attr_pos.append((start_pos, start_pos + 1)) + encoded.attr_types.append(_FLOAT_TYPE) + elif isinstance(v, str): + start_pos = len(encoded.attr_strs) + encoded.attr_strs.append(v) + encoded.attr_pos.append((start_pos, start_pos + 1)) + encoded.attr_types.append(_STRING_TYPE) + elif isinstance(v, Sequence): + if len(v) == 0: + raise ValueError(f"Empty sequence for attribute {k}") + if any(isinstance(elem, float) for elem in v): + start_pos = len(encoded.attr_floats) + encoded.attr_floats.extend([float(elem) for elem in v]) + encoded.attr_pos.append((start_pos, start_pos + len(v))) + encoded.attr_types.append(_FLOAT_SEQ_TYPE) + elif isinstance(v[0], int): + start_pos = len(encoded.attr_ints) + encoded.attr_ints.extend([int(elem) for elem in v]) + encoded.attr_pos.append((start_pos, start_pos + len(v))) + encoded.attr_types.append(_INT_SEQ_TYPE) + elif isinstance(v[0], str): + start_pos = len(encoded.attr_strs) + encoded.attr_strs.extend([str(elem) for elem in v]) + encoded.attr_pos.append((start_pos, start_pos + len(v))) + encoded.attr_types.append(_STRING_SEQ_TYPE) + else: + raise ValueError(f"Unsupported sequence type for attribute {k}") + else: + raise ValueError(f"Unsupported attribute type for {k}: {type(v)}") + if len(encoded.attr_keys) != len(encoded.attr_types): + raise AssertionError( + f"Mismatch between number of attribute keys and types: {len(encoded.attr_keys)} != {len(encoded.attr_types)}" + ) + if len(encoded.attr_keys) != len(encoded.attr_pos): + raise AssertionError( + f"Mismatch between number of attribute keys and positions: {len(encoded.attr_keys)} != {len(encoded.attr_pos)}" + ) + return encoded + + def to_dict( + self, + ) -> dict[ + str, + int | float | str | list[int] | list[float] | list[str], + ]: + """Convert the encoded attributes back to a dictionary for creating an ONNX node.""" + attrs: dict[ + str, + int | float | str | list[int] | list[float] | list[str], + ] = {} + for i, key in enumerate(self.attr_keys): + attr_type = self.attr_types[i] + if attr_type == _INT_TYPE: + attrs[key] = self.attr_ints[self.attr_pos[i][0]] + elif attr_type == _FLOAT_TYPE: + attrs[key] = self.attr_floats[self.attr_pos[i][0]] + elif attr_type == _STRING_TYPE: + attrs[key] = self.attr_strs[self.attr_pos[i][0]] + elif attr_type == _FLOAT_SEQ_TYPE: + attrs[key] = self.attr_floats[self.attr_pos[i][0] : self.attr_pos[i][1]] + elif attr_type == _INT_SEQ_TYPE: + attrs[key] = self.attr_ints[self.attr_pos[i][0] : self.attr_pos[i][1]] + elif attr_type == _STRING_SEQ_TYPE: + attrs[key] = self.attr_strs[self.attr_pos[i][0] : self.attr_pos[i][1]] + else: + raise ValueError(f"Unsupported attribute type: {attr_type}") + return attrs + + +@torch.library.custom_op( + "onnx_symbolic::_symbolic", + mutates_args=(), + schema=( + "(Tensor?[] inputs, str op_type, int onnx_dtype, *," + " SymInt[] shape, str[] attr_keys, str[] attr_types, int[][] attr_pos," + " int[] attr_ints, float[] attr_floats, str[] attr_strs, str[] metadata_props_keys," + " str[] metadata_props_values, str domain='', int? version=None" + ") -> Tensor" + ), +) +def _symbolic( + inputs: Sequence[torch.Tensor | None], + op_type: str, + onnx_dtype: int, + *, + shape: Sequence[int | torch.SymInt], + attr_keys: Sequence[str], + attr_types: Sequence[str], + attr_pos: Sequence[tuple[int, int]], + attr_ints: Sequence[int], + attr_floats: Sequence[float], + attr_strs: Sequence[str], + metadata_props_keys: Sequence[str] = (), + metadata_props_values: Sequence[str] = (), + domain: str = "", + version: int | None = None, +) -> torch.Tensor: + torch._check( + onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE, + lambda: f"{onnx_dtype} is invalid as an ONNX data type. Valid values are {list(_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE.keys())}", + ) + return torch.zeros( + shape, dtype=_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[onnx_dtype] + ) + + +@_symbolic.register_fake +def _( + inputs: Sequence[torch.Tensor], + op_type: str, + onnx_dtype: int, + *, + shape: Sequence[int | torch.SymInt], + attr_keys: Sequence[str], + attr_types: Sequence[str], + attr_pos: Sequence[tuple[int, int]], + attr_ints: Sequence[int], + attr_floats: Sequence[float], + attr_strs: Sequence[str], + metadata_props_keys: Sequence[str] = (), + metadata_props_values: Sequence[str] = (), + domain: str = "", + version: int | None = None, +) -> torch.Tensor: + torch._check( + onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE, + lambda: f"{onnx_dtype} is invalid as an ONNX data type. Valid values are {list(_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE.keys())}", + ) + # NOTE(justinchuby): Use zeros instead of torch.empty because I haven't figured + # out how it can handle empty shapes + return torch.zeros( + shape, dtype=_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[onnx_dtype] + ) + + +@torch.library.custom_op( + "onnx_symbolic::_symbolic_multi_out", + mutates_args=(), + schema=( + "(Tensor?[] inputs, str op_type, int[] onnx_dtypes, *," + " SymInt[][] shapes, str[] attr_keys, str[] attr_types, int[][] attr_pos," + " int[] attr_ints, float[] attr_floats, str[] attr_strs, str[] metadata_props_keys," + " str[] metadata_props_values, str domain='', int? version=None" + ") -> Tensor[]" + ), +) +def _symbolic_multi_out( + inputs: Sequence[torch.Tensor | None], + op_type: str, + onnx_dtypes: Sequence[int], + *, + shapes: Sequence[Sequence[int | torch.SymInt]], + attr_keys: Sequence[str], + attr_types: Sequence[str], + attr_pos: Sequence[tuple[int, int]], + attr_ints: Sequence[int], + attr_floats: Sequence[float], + attr_strs: Sequence[str], + metadata_props_keys: Sequence[str] = (), + metadata_props_values: Sequence[str] = (), + domain: str = "", + version: int | None = None, +) -> list[torch.Tensor]: + outputs = [] + torch._check( + len(shapes) == len(onnx_dtypes), + lambda: f"Number of shapes ({len(shapes)}) must match number of ONNX dtypes ({len(onnx_dtypes)})", + ) + for shape, onnx_dtype in zip(shapes, onnx_dtypes): + torch._check( + onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE, + lambda: f"{onnx_dtype} is invalid as an ONNX data type. Valid values are {list(_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE.keys())}", + ) + outputs.append( + torch.zeros( + shape, dtype=_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[onnx_dtype] + ) + ) + return outputs + + +@_symbolic_multi_out.register_fake +def _( + inputs: Sequence[torch.Tensor], + op_type: str, + onnx_dtypes: Sequence[int], + *, + shapes: Sequence[Sequence[int | torch.SymInt]], + attr_keys: Sequence[str], + attr_types: Sequence[str], + attr_pos: Sequence[tuple[int, int]], + attr_ints: Sequence[int], + attr_floats: Sequence[float], + attr_strs: Sequence[str], + metadata_props_keys: Sequence[str] = (), + metadata_props_values: Sequence[str] = (), + domain: str = "", + version: int | None = None, +) -> list[torch.Tensor]: + outputs = [] + torch._check( + len(shapes) == len(onnx_dtypes), + lambda: f"Number of shapes ({len(shapes)}) must match number of ONNX dtypes ({len(onnx_dtypes)})", + ) + for shape, onnx_dtype in zip(shapes, onnx_dtypes): + torch._check( + onnx_dtype in _dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE, + lambda: f"{onnx_dtype} is invalid as an ONNX data type. Valid values are {list(_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE.keys())}", + ) + # NOTE(justinchuby): Use zeros instead of torch.empty because I haven't figured + # out how it can handle empty shapes + outputs.append( + torch.zeros( + shape, dtype=_dtype_mappings.ONNX_DTYPE_TO_TORCH_DTYPE[onnx_dtype] + ) + ) + return outputs diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_helper.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_helper.py new file mode 100644 index 0000000000000000000000000000000000000000..76b50a8eb3f77a8afa155e924d1734b601711dba --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_helper.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_helper.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_helper import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset10.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset10.py new file mode 100644 index 0000000000000000000000000000000000000000..9bda69b81ab603fdb4932fa69edefa8525b2173f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset10.py @@ -0,0 +1,11 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset10.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset10 import * # noqa: F401,F403 +from torch.onnx._internal.torchscript_exporter.symbolic_opset10 import ( # noqa: F401 + _slice, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset11.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset11.py new file mode 100644 index 0000000000000000000000000000000000000000..276ef7209bf69df1697838ac65354c73ecc29ea4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset11.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset11.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset11 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset12.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset12.py new file mode 100644 index 0000000000000000000000000000000000000000..63e137734e8a74c8644fb273590db9396c4d9511 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset12.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset12.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset12 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset13.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset13.py new file mode 100644 index 0000000000000000000000000000000000000000..18aff9295be8cdabfca7d6395ad562bcc7e5834d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset13.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset13.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset13 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset14.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset14.py new file mode 100644 index 0000000000000000000000000000000000000000..367aa9eb0832afb7a1b9f0b08936ef349656a218 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset14.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset14.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset14 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset15.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset15.py new file mode 100644 index 0000000000000000000000000000000000000000..e04e3b04521275021659a1cbbd6cae90b3440192 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset15.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset15.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset15 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset16.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset16.py new file mode 100644 index 0000000000000000000000000000000000000000..9a248bb0f26c5469278e096ad1c7df486748e62d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset16.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset16.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset16 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset17.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset17.py new file mode 100644 index 0000000000000000000000000000000000000000..800acd446b5dc3739d351f4796603ce23747cdd6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset17.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset17.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset17 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset18.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset18.py new file mode 100644 index 0000000000000000000000000000000000000000..cc07a60f018d8c7c21ea01e1972252d8b849c1a6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset18.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset18.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset18 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset19.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset19.py new file mode 100644 index 0000000000000000000000000000000000000000..4f7a54fc1dd38a62b137430dae4a5ccce93d2b73 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset19.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset19.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset19 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset20.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset20.py new file mode 100644 index 0000000000000000000000000000000000000000..56635a781161106c02b05f2d5b82d229e8f9360f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset20.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset20.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset20 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset7.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset7.py new file mode 100644 index 0000000000000000000000000000000000000000..c11e769677ec424a8bfc231fc1d21b8c859ef314 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset7.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset7.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset7 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset8.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset8.py new file mode 100644 index 0000000000000000000000000000000000000000..0e4411649f3e09b2267e76f0307a833d6d15a7b0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset8.py @@ -0,0 +1,8 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset8.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset8 import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset9.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset9.py new file mode 100644 index 0000000000000000000000000000000000000000..bd0f4795340ae8d66f426d30629ab689ee4417a3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/symbolic_opset9.py @@ -0,0 +1,14 @@ +"""Backward compatibility module for torch.onnx.symbolic_opset9.""" + +from __future__ import annotations + + +__all__: list[str] = [] + +from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import * # noqa: F401,F403 +from torch.onnx._internal.torchscript_exporter.symbolic_opset9 import ( # noqa: F401 + _prepare_onnx_paddings, + _reshape_from_tensor, + _slice, + _var_mean, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/testing.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/testing.py new file mode 100644 index 0000000000000000000000000000000000000000..aa168b32746f5f0ee67dda5c1035d8c975f2ac02 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/testing.py @@ -0,0 +1,8 @@ +"""Utilities to aid in testing exported ONNX models.""" + +__all__ = ["assert_onnx_program"] + +from torch.onnx._internal.exporter._testing import assert_onnx_program + + +assert_onnx_program.__module__ = "torch.onnx.testing" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..d387069b0b0c96242f3b1f7debe1ed909573020e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/utils.py @@ -0,0 +1,9 @@ +"""Backward compatibility module for torch.onnx.utils.""" + +from __future__ import annotations + + +__all__: list[str] = [] + + +from torch.onnx._internal.torchscript_exporter.utils import * # noqa: F401,F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/verification.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/verification.py new file mode 100644 index 0000000000000000000000000000000000000000..70d901acb47a941bf03caa25f6e8dfaadebe170b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/onnx/verification.py @@ -0,0 +1,12 @@ +"""A set of tools to verify the correctness of ONNX models.""" + +__all__ = ["VerificationInfo", "verify_onnx_program"] + +from torch.onnx._internal.exporter._verification import ( + VerificationInfo, + verify_onnx_program, +) + + +VerificationInfo.__module__ = "torch.onnx.verification" +verify_onnx_program.__module__ = "torch.onnx.verification" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1060a6287a8e6e4b59aa1a46527cf0001de1ccfe --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/__init__.py @@ -0,0 +1,66 @@ +""" +:mod:`torch.optim` is a package implementing various optimization algorithms. + +Most commonly used methods are already supported, and the interface is general +enough, so that more sophisticated ones can also be easily integrated in the +future. +""" + +from torch.optim import lr_scheduler as lr_scheduler, swa_utils as swa_utils +from torch.optim._adafactor import Adafactor as Adafactor +from torch.optim._muon import Muon as Muon +from torch.optim.adadelta import Adadelta as Adadelta +from torch.optim.adagrad import Adagrad as Adagrad +from torch.optim.adam import Adam as Adam +from torch.optim.adamax import Adamax as Adamax +from torch.optim.adamw import AdamW as AdamW +from torch.optim.asgd import ASGD as ASGD +from torch.optim.lbfgs import LBFGS as LBFGS +from torch.optim.nadam import NAdam as NAdam +from torch.optim.optimizer import Optimizer as Optimizer +from torch.optim.radam import RAdam as RAdam +from torch.optim.rmsprop import RMSprop as RMSprop +from torch.optim.rprop import Rprop as Rprop +from torch.optim.sgd import SGD as SGD +from torch.optim.sparse_adam import SparseAdam as SparseAdam + + +Adafactor.__module__ = "torch.optim" +Muon.__module__ = "torch.optim" + + +del adadelta # type: ignore[name-defined] # noqa: F821 +del adagrad # type: ignore[name-defined] # noqa: F821 +del adam # type: ignore[name-defined] # noqa: F821 +del adamw # type: ignore[name-defined] # noqa: F821 +del sparse_adam # type: ignore[name-defined] # noqa: F821 +del adamax # type: ignore[name-defined] # noqa: F821 +del asgd # type: ignore[name-defined] # noqa: F821 +del sgd # type: ignore[name-defined] # noqa: F821 +del radam # type: ignore[name-defined] # noqa: F821 +del rprop # type: ignore[name-defined] # noqa: F821 +del rmsprop # type: ignore[name-defined] # noqa: F821 +del optimizer # type: ignore[name-defined] # noqa: F821 +del nadam # type: ignore[name-defined] # noqa: F821 +del lbfgs # type: ignore[name-defined] # noqa: F821 + +__all__ = [ + "Adafactor", + "Adadelta", + "Adagrad", + "Adam", + "Adamax", + "AdamW", + "ASGD", + "LBFGS", + "lr_scheduler", + "Muon", + "NAdam", + "Optimizer", + "RAdam", + "RMSprop", + "Rprop", + "SGD", + "SparseAdam", + "swa_utils", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_adafactor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_adafactor.py new file mode 100644 index 0000000000000000000000000000000000000000..99efd40cbfcc46bc8f3aef69c132f1a4e013228e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_adafactor.py @@ -0,0 +1,661 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +from typing import cast, TYPE_CHECKING + +import torch +from torch import Tensor + +from .optimizer import ( + _disable_dynamo_if_unsupported, + _get_scalar_dtype, + _maximize_doc, + _params_doc, + _to_scalar, + Optimizer, + ParamsT, + TensorListList, +) + + +__all__ = ["Adafactor", "adafactor"] + + +class Adafactor(Optimizer): + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-2, + beta2_decay: float = -0.8, + eps: tuple[float | None, float] = (None, 1e-3), + d: float = 1.0, + weight_decay: float = 0.0, + *, + foreach: bool | None = None, + maximize: bool = False, + ) -> None: + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Learning rate should be >= 0 but is: {lr}") + if not 0.0 >= beta2_decay: + raise ValueError(f"beta2_decay should be <= 0 but is: {beta2_decay}") + if eps[0] is not None and not 0.0 <= eps[0]: + raise ValueError(f"epsilon1 should be >= 0 but is: {eps[0]}") + if not 0.0 <= eps[1]: + raise ValueError(f"epsilon2 should be >= 0 but is: {eps[1]}") + if not 1.0 <= d: + raise ValueError(f"Clipping threshold d should be >= 1 but is: {d}") + if not 0.0 <= weight_decay: + raise ValueError(f"weight_decay should be >= 0 but is: {weight_decay}") + defaults = { + "lr": lr, + "beta2_decay": beta2_decay, + "eps": eps, + "d": d, + "weight_decay": weight_decay, + "foreach": foreach, + "maximize": maximize, + } + super().__init__(params, defaults) + + def __setstate__(self, state): + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("foreach", None) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = torch.tensor(step_val, dtype=_get_scalar_dtype()) + + def _init_group( + self, + group, + params_with_grad, + grads, + row_vars, + col_vars, + variances, + state_steps, + ) -> bool: + for p in group["params"]: + if p.grad is None: + continue + if torch.is_complex(p): + raise RuntimeError("Adafactor does not support complex parameters") + if p.grad.is_sparse: + raise RuntimeError("Adafactor does not support sparse gradients") + + params_with_grad.append(p) + grads.append(p.grad) + + state = self.state[p] + + # State initialization + if len(state) == 0: + # note(crcrpar): Deliberately host `step` on CPU if both capturable and fused are off. + # This is because kernel launches are costly on CUDA and XLA. + state["step"] = torch.tensor(0.0, dtype=_get_scalar_dtype()) + + if p.grad.dim() > 1: + row_shape = list(p.grad.shape) + row_shape[-1] = 1 + # Row factor of variance, NOT the same shape as grads (will be reduced along last dim) + state["row_var"] = p.grad.new_zeros(row_shape) + + col_shape = list(p.grad.shape) + col_shape[-2] = 1 + # Col factor of variance, NOT the same shape as grads (will be reduced along penultimate dim) + state["col_var"] = p.grad.new_zeros(col_shape) + else: + state["variance"] = torch.zeros_like( + p.grad, memory_format=torch.preserve_format + ) + + row_vars.append(state.get("row_var", None)) + col_vars.append(state.get("col_var", None)) + variances.append(state.get("variance", None)) + state_steps.append(state["step"]) + return False # has_complex + + @torch.no_grad() + def step(self, closure=None): + r"""Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._accelerator_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + row_vars: list[Tensor | None] = [] + col_vars: list[Tensor | None] = [] + variances: list[Tensor | None] = [] + state_steps: list[Tensor] = [] + eps1, eps2 = group["eps"] + + has_complex = self._init_group( + group, + params_with_grad, + grads, + row_vars, + col_vars, + variances, + state_steps, + ) + + adafactor( + params_with_grad, + grads, + row_vars, + col_vars, + variances, + state_steps, + d=group["d"], + lr=group["lr"], + beta2_decay=group["beta2_decay"], + weight_decay=group["weight_decay"], + eps1=eps1, + eps2=eps2, + foreach=group["foreach"], + maximize=group["maximize"], + grad_scale=getattr(self, "grad_scale", None), + found_inf=getattr(self, "found_inf", None), + has_complex=has_complex, + ) + + return loss + + +Adafactor.__doc__ = ( + r"""Implements Adafactor algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{(lr)}, \: \tau + \text{(}\beta_2\text{ decay)}, \: \theta_0 \text{(params)}, \: f(\theta) \text{(objective)}, \\ + &\hspace{15mm} \: \epsilon_1, \epsilon_2 \text{ (epsilons)}, \: d \text{(clipping threshold)}, \\ + &\hspace{15mm} \: \lambda \text{(weight decay)}, + \: \textit{maximize} \\ + &\textbf{initialize} : \: R_0 \leftarrow 0 \text{ (second moment row factor)}, \\ + &\hspace{23mm} \: C_0 \leftarrow 0 \text{ (second moment col factor)}, \\ + &\hspace{23mm} \: \widehat{V}_0 \leftarrow 0 \text{ (second moment for vectors)} \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + + &\hspace{5mm}\textbf{if} \: \textit{maximize}: \\ + &\hspace{10mm}G_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}G_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\widehat{\beta}_{2_t} \leftarrow 1 - t^{\tau} \\ + &\hspace{5mm}\rho_t \leftarrow min(lr, \frac{1}{\sqrt{t}}) \\ + &\hspace{5mm}\alpha_t \leftarrow max(\epsilon_2, + \text{RMS}(\theta_{t-1}))\rho_t \\ + &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \gamma \lambda \theta_{t-1} \\ + &\hspace{5mm}\textbf{if} \: \text{dim}(G_t) > 1: \\ + &\hspace{10mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ + (1-\widehat{\beta}_{2_t})(G_t \odot G_t) \cdot 1_m \\ + &\hspace{10mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ + (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t) \\ + &\hspace{10mm}\widehat{V}_t \leftarrow + \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}\widehat{V}_t \leftarrow \widehat{\beta}_{2_t}\widehat{V}_{t-1}+ + (1-\widehat{\beta}_{2_t}) \cdot (G_t \odot G_t) \\ + &\hspace{5mm}U_t \leftarrow + \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ + &\hspace{5mm}\widehat{U}_t \leftarrow \frac{U_t}{max(1, \frac{\text{RMS}(U_t)}{d})} \\ + &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \alpha_t \widehat{U}_t \\ + + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to `Adafactor: Adaptive Learning Rates with Sublinear Memory Cost`_. + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): unlike other optimizers, Adafactor does not require a + learning rate, and Noam Shazeer and Mitchell Stern do not use lr at all. + Deviating from the paper, this implementation uses lr for applying weight + decay and as the maximum value for relative step size rho_t. Note that in + the paper, a constant of 0.01 is used as the maximum value for relative + step size, and so we set 0.01 as the default value. (default: 1e-2) + beta2_decay (float, optional): the decay rate of beta2. beta2 standardly refers + to the coefficient used for computing the running average of the gradient + squared. (default: -0.8) + eps (Tuple[float, float], optional): epsilon1 is the term added to the denominator + of the update calculation to improve numerical stability. This use of epsilon1 + deviates from the algorithm written in the paper! See note below for more details. + epsilon2 is the term used to avoid having too small a weight update when applying + parameter scaling. (default: (None, 1e-3)) + d (float, optional): the clipping threshold, used to avoid larger-than-desired + updates. + weight_decay (float, optional): weight decay coefficient (default: 1e-2) + foreach (bool, optional): whether foreach implementation of optimizer is used. Note + that the foreach implementation uses ~ sizeof(params) more peak memory than the + for-loop version due to the intermediates being a tensorlist vs just one tensor. + As Adafactor is commonly used when memory is prohibitive, Adafactor will default + to the slower single tensor for-loop implementation unless this flag is explicitly + True. This behavior is contrary to other optimizers, which will attempt defaulting + to foreach on CUDA for faster runtime. (default: None) + {_maximize_doc}""" + + r""" + .. Note:: + The implementation of Adafactor subtly differs from Noam Shazeer and Mitchell Stern + and implementations in some other frameworks with its use of learning rate and + :math:`\epsilon_1`. + + Regarding the learning rate hyperparameter: Noam Shazeer and Mitchell Stern do not + use lr at all, as the stated algorithm uses :math:`\rho_t` and update clipping to + affect the step size. + + This implementation allows `lr` to influence the maximum value for :math:`\rho_t`: + + .. math:: + \begin{aligned} + &\hspace{5mm}\rho_t \leftarrow min(lr, \frac{1}{\sqrt{t}}) + \end{aligned} + + This differs from Noam Shazeer and Mitchell Stern, who use a constant of 0.01 as + the maximum value of :math:`\rho_t` + + .. math:: + \begin{aligned} + &\hspace{5mm}\rho_t \leftarrow min(0.01, \frac{1}{\sqrt{t}}) + \end{aligned} + + Noam Shazeer and Mitchell Stern do not enforce an opinion on how weight decay should + be computed, and so we use the learning rate as a coefficient for decoupled weight + decay, similar to what is suggested in `Decoupled Weight Decay Regularization`_. + + Regarding the use of :math:`\epsilon_1`: The implementation attempts to replicate the + presumed intention of Noam Shazeer and Mitchell Stern to use :math:`\epsilon_1` as + a stabilizing term when the squared gradient becomes small. + + This stabilization can be written as + + .. math:: + \begin{aligned} + &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ + (1-\widehat{\beta}_{2_t})(G_t \odot G_t + 1_n \cdot 1^\top_m) \cdot 1_m \\ + &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ + (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + 1_n \cdot 1^\top_m) \\ + &\hspace{5mm}\widehat{V}_t \leftarrow + \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ + &\hspace{5mm}U_t \leftarrow \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ + \end{aligned} + + where the row and column factors of gradient squared :math:`R_t` and :math:`C_t` + are left alone, and we apply :math:`\epsilon_1` at the final calculation of + the variance estimate :math:`\widehat{V}_t` and for the update :math:`U_t`. + + This is in contrast to Noam Shazeer and Mitchell Stern and other frameworks which + apply :math:`\epsilon_1` to both row and column factors of the squared gradient, but + not in the calculations after: + + .. math:: + \begin{aligned} + &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ + (1-\widehat{\beta}_{2_t})(G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \cdot 1_m \\ + &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ + (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \\ + &\hspace{5mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{1^\top_n \cdot R_t} \\ + &\hspace{5mm}U_t \leftarrow \frac{G_t}{\sqrt{\widehat{V}_t}} \\ + \end{aligned} + + You may note that Noam Shazeer and Mitchell Stern describe using the sum of squared gradients, + while this implementation uses the mean instead. This choice is mathematically equivalent and + allows for greater numerical stability for large sums. + + .. _Adafactor\: Adaptive Learning Rates with Sublinear Memory Cost: + https://arxiv.org/pdf/1804.04235 + .. _Decoupled Weight Decay Regularization: + https://arxiv.org/abs/1711.05101 + """ +) + + +def _single_tensor_adafactor( + params: list[Tensor], + grads: list[Tensor], + # If grad is 1-dimensional (aka a vector), there is no factorization necessary + # so row_var and col_var will be None while variance will be filled. + # Contrarily, for a grad with multiple dimensions, we will factor along the last + # 2 dimensions, and so row_var and col_var will be filled and variance will be None. + row_vars: list[Tensor | None], + col_vars: list[Tensor | None], + variances: list[Tensor | None], + state_steps: list[Tensor], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + d: float, + lr: Tensor | float, + beta2_decay: float, + weight_decay: float, + eps1: float | None, + eps2: float, + maximize: bool, + has_complex: bool, +) -> None: + if grad_scale is not None or found_inf is not None: + raise AssertionError("Grad scaling should occur outside of optimizer.step()") + + if torch.jit.is_scripting(): + # this assert is due to JIT being dumb and not realizing that the ops below + # have overloads to handle both float and Tensor lrs, so we just assert it's + # a float since most people using JIT are using floats + if not isinstance(lr, float): + raise AssertionError(f"Expected lr to be a float, but got {type(lr)}") + + else: + lr = _to_scalar(lr) + + for i, param in enumerate(params): + grad = grads[i] if not maximize else -grads[i] + step_t = state_steps[i] + row_var = row_vars[i] + col_var = col_vars[i] + variance = variances[i] + if eps1 is None: + eps1 = torch.finfo(param.dtype).eps + + # update step + step_t += 1 + step_float = step_t.item() + + one_minus_beta2_t = step_float**beta2_decay + rho_t = min(lr, 1 / (step_float**0.5)) + alpha = max(eps2, param.norm(2).item() / (param.numel() ** 0.5)) * rho_t + + # Perform stepweight decay + if weight_decay != 0: + param.mul_(1 - lr * weight_decay) + + if grad.dim() > 1: + if row_var is None or col_var is None: + raise AssertionError( + "row_var and col_var should be defined when grad is multidimensional" + ) + # same as (g * g).mean(dim=-1) w/o materializing an intermediate size g + row_mean = ( + torch.norm(grad, dim=-1, keepdim=True).square_().div_(grad.size(-1)) + ) + row_var.lerp_(row_mean, one_minus_beta2_t) + # same as (g * g).mean(dim=-2) w/o materializing an intermediate size g + col_mean = ( + torch.norm(grad, dim=-2, keepdim=True).square_().div_(grad.size(-2)) + ) + col_var.lerp_(col_mean, one_minus_beta2_t) + var_estimate = row_var @ col_var + var_estimate.div_(row_var.mean(dim=-2, keepdim=True).clamp_(min=eps1)) + else: + if variance is None: + raise AssertionError("variance should be defined when grad is a vector") + grad_squared = grad * grad + variance.lerp_(grad_squared, one_minus_beta2_t) + # avoid writing into variance during update + var_estimate = variance.clone() + + # square the eps1 as we sqrt after to keep eps1's magnitude + update = var_estimate.clamp_(min=eps1 * eps1).rsqrt_() + update.mul_(grad) + denom = max(1.0, update.norm(2).item() / ((update.numel() ** 0.5) * d)) + param.add_(update, alpha=-alpha / denom) + + +def _group_tensors_by_device_dtype_and_is_multidim( + tensorlists: TensorListList, +) -> dict[ + tuple[torch.device | None, torch.dtype | None, bool], + list[list[Tensor | None]], +]: + """Groups tensors by device, dtype, AND multidimensionality -- whether the tensor + has multiple dims or just one dim (is a vector). This allows the foreach impl of + Adafactor to assume that every group of params will either be factored or not.""" + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype(tensorlists) + ultra_grouped_tensors: dict[ + tuple[torch.device | None, torch.dtype | None, bool], + list[list[Tensor | None]], + ] = {} + for (device, dtype), (tensorlists, _) in grouped_tensors.items(): + matrix_key = (device, dtype, True) + vector_key = (device, dtype, False) + + # assumes grad is the second tensorlist + for j, tensor in enumerate(tensorlists[1]): + if tensor is None: + raise AssertionError("grad should not be None") + if tensor.dim() > 1: + if matrix_key not in ultra_grouped_tensors: + ultra_grouped_tensors[matrix_key] = [[] for _ in tensorlists] + for i in range(len(tensorlists)): + ultra_grouped_tensors[matrix_key][i].append(tensorlists[i][j]) + else: + if vector_key not in ultra_grouped_tensors: + ultra_grouped_tensors[vector_key] = [[] for _ in tensorlists] + for i in range(len(tensorlists)): + ultra_grouped_tensors[vector_key][i].append(tensorlists[i][j]) + return ultra_grouped_tensors + + +def _multi_tensor_adafactor( + params: list[Tensor], + grads: list[Tensor], + # If grad is 1-dimensional (aka a vector), there is no factorization necessary + # so row_var and col_var will be None while variance will be filled. + # Contrarily, for a grad with multiple dimensions, we will factor along the last + # 2 dimensions, and so row_var and col_var will be filled and variance will be None. + row_vars: list[Tensor | None], + col_vars: list[Tensor | None], + variances: list[Tensor | None], + state_steps: list[Tensor], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + d: float, + lr: Tensor | float, + beta2_decay: float, + weight_decay: float, + eps1: float | None, + eps2: float, + maximize: bool, + has_complex: bool, +) -> None: + if len(params) == 0: + return + + if grad_scale is not None or found_inf is not None: + raise AssertionError("Grad scaling should occur outside of optimizer.step()") + + lr = _to_scalar(lr) + + grouped_tensors = _group_tensors_by_device_dtype_and_is_multidim( + [params, grads, row_vars, col_vars, variances, state_steps] # type: ignore[list-item] + ) + for (_, dtype, is_multidim), ( + ( + device_params_, + device_grads_, + device_row_vars_, + device_col_vars_, + device_variances_, + device_state_steps_, + ) + ) in grouped_tensors.items(): + device_params = cast(list[Tensor], device_params_) + device_grads = cast(list[Tensor], device_grads_) + device_state_steps = cast(list[Tensor], device_state_steps_) + if eps1 is None: + if dtype is None: + raise AssertionError( + "dtype is needed to compute eps1 when eps1 is unset" + ) + eps1 = torch.finfo(dtype).eps + + if TYPE_CHECKING: + assert device_state_steps[0] is not None + + if maximize: + device_grads = torch._foreach_neg(device_grads) # type: ignore[assignment] + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and device_state_steps[0].is_cpu: + torch._foreach_add_( + device_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(device_state_steps, 1.0) + + one_minus_beta2_ts = [] + beta2_ts = [] + rho_ts = [] + for s in device_state_steps: + one_minus_beta2_ts.append(s.item() ** beta2_decay) + beta2_ts.append(1 - s.item() ** beta2_decay) + rho_ts.append(min(lr, 1 / (s.item() ** 0.5))) + + alphas = [ + max(eps2, p.norm(2).item() / (p.numel() ** 0.5)) * r + for p, r in zip(device_params, rho_ts, strict=True) + ] + + # Perform stepweight decay + if weight_decay != 0: + torch._foreach_mul_(device_params, 1 - lr * weight_decay) + + if is_multidim: + device_row_vars = cast(list[Tensor], device_row_vars_) + device_col_vars = cast(list[Tensor], device_col_vars_) + if device_row_vars[0] is None or device_col_vars[0] is None: + raise AssertionError( + "row_var and col_var should be defined when grad is multidimensional" + ) + # same as (g * g).mean(dim=-1) w/o materializing an intermediate size g + row_means = [ + torch.norm(grad, dim=-1, keepdim=True) for grad in device_grads + ] + torch._foreach_mul_(row_means, row_means) + torch._foreach_div_(row_means, [grad.size(-1) for grad in device_grads]) + torch._foreach_lerp_(device_row_vars, row_means, one_minus_beta2_ts) + del row_means + + # same as (g * g).mean(dim=-2) w/o materializing an intermediate size g + col_means = [ + torch.norm(grad, dim=-2, keepdim=True) for grad in device_grads + ] + torch._foreach_mul_(col_means, col_means) + torch._foreach_div_(col_means, [grad.size(-2) for grad in device_grads]) + torch._foreach_lerp_(device_col_vars, col_means, one_minus_beta2_ts) + del col_means + + var_estimates = [ + row_var @ col_var + for row_var, col_var in zip( + device_row_vars, device_col_vars, strict=True + ) + ] + row_var_means = [ + row_var.mean(dim=-2, keepdim=True) for row_var in device_row_vars + ] + torch._foreach_clamp_min_(row_var_means, eps1) + torch._foreach_div_(var_estimates, row_var_means) + del row_var_means + else: + device_variances = cast(list[Tensor], device_variances_) + if device_variances[0] is None: + raise AssertionError("variance should be defined when grad is a vector") + + grads_squared = torch._foreach_mul(device_grads, device_grads) + torch._foreach_lerp_(device_variances, grads_squared, one_minus_beta2_ts) + del grads_squared + + # avoid writing into variance during update + var_estimates = [v.clone() for v in device_variances] + + # square the eps1 as we sqrt after to keep eps1's magnitude + torch._foreach_clamp_min_(var_estimates, eps1 * eps1) + torch._foreach_rsqrt_(var_estimates) + torch._foreach_mul_(var_estimates, device_grads) + updates = var_estimates + + alphas = [ + -a / (max(1.0, update.norm(2).item() / ((update.numel() ** 0.5) * d))) + for a, update in zip(alphas, updates, strict=True) + ] + torch._foreach_mul_(updates, alphas) + torch._foreach_add_(device_params, updates) + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_adafactor) +def adafactor( + params: list[Tensor], + grads: list[Tensor], + row_vars: list[Tensor | None], + col_vars: list[Tensor | None], + variances: list[Tensor | None], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + foreach: bool | None = None, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, + has_complex: bool = False, + *, + d: float, + lr: float | Tensor, + beta2_decay: float, + weight_decay: float, + eps1: float, + eps2: float, + maximize: bool, +) -> None: + r"""Functional API that performs Adafactor algorithm computation. + + See :class:`~torch.optim.Adafactor` for details. + """ + if not torch.compiler.is_compiling() and not all( + isinstance(t, torch.Tensor) for t in state_steps + ): + raise RuntimeError( + "`state_steps` argument must contain a list of singleton tensors" + ) + + if foreach: + func = _multi_tensor_adafactor + else: + func = _single_tensor_adafactor + + func( + params, + grads, + row_vars, + col_vars, + variances, + state_steps, + d=d, + lr=lr, + beta2_decay=beta2_decay, + weight_decay=weight_decay, + eps1=eps1, + eps2=eps2, + maximize=maximize, + grad_scale=grad_scale, + found_inf=found_inf, + has_complex=has_complex, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_functional.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_functional.py new file mode 100644 index 0000000000000000000000000000000000000000..ba97bc997937889b93ca30b8b783f38dd895a227 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_functional.py @@ -0,0 +1,84 @@ +# mypy: allow-untyped-defs +r"""Functional interface.""" + +import math + +from torch import Tensor + +from .adadelta import adadelta # type: ignore[attr-defined] # noqa: F401 +from .adagrad import _make_sparse, adagrad # type: ignore[attr-defined] # noqa: F401 +from .adam import adam # type: ignore[attr-defined] # noqa: F401 +from .adamax import adamax # type: ignore[attr-defined] # noqa: F401 +from .adamw import adamw # type: ignore[attr-defined] # noqa: F401 +from .asgd import asgd # type: ignore[attr-defined] # noqa: F401 +from .nadam import nadam # type: ignore[attr-defined] # noqa: F401 +from .radam import radam # type: ignore[attr-defined] # noqa: F401 +from .rmsprop import rmsprop # type: ignore[attr-defined] # noqa: F401 +from .rprop import rprop # type: ignore[attr-defined] # noqa: F401 +from .sgd import sgd # type: ignore[attr-defined] # noqa: F401 + + +# TODO: use foreach API in optim._functional to do all the computation + + +def sparse_adam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + state_steps: list[int], + *, + eps: float, + beta1: float, + beta2: float, + lr: float, + maximize: bool, +) -> None: + r"""Functional API that performs Sparse Adam algorithm computation. + + See :class:`~torch.optim.SparseAdam` for details. + """ + for i, param in enumerate(params): + grad = grads[i] + grad = grad if not maximize else -grad + grad = grad.coalesce() # the update is non-linear so indices must be unique + grad_indices = grad._indices() + grad_values = grad._values() + if grad_values.numel() == 0: + # Skip update for empty grad + continue + size = grad.size() + + exp_avg = exp_avgs[i] + exp_avg_sq = exp_avg_sqs[i] + step = state_steps[i] + + def make_sparse(values): + constructor = grad.new + if grad_indices.dim() == 0 or values.dim() == 0: + return constructor().resize_as_(grad) + return constructor(grad_indices, values, size) + + # Decay the first and second moment running average coefficient + # old <- b * old + (1 - b) * new + # <==> old += (1 - b) * (new - old) + old_exp_avg_values = exp_avg.sparse_mask(grad)._values() + exp_avg_update_values = grad_values.sub(old_exp_avg_values).mul_(1 - beta1) + exp_avg.add_(make_sparse(exp_avg_update_values)) + old_exp_avg_sq_values = exp_avg_sq.sparse_mask(grad)._values() + exp_avg_sq_update_values = ( + grad_values.pow(2).sub_(old_exp_avg_sq_values).mul_(1 - beta2) + ) + exp_avg_sq.add_(make_sparse(exp_avg_sq_update_values)) + + # Dense addition again is intended, avoiding another sparse_mask + numer = exp_avg_update_values.add_(old_exp_avg_values) + exp_avg_sq_update_values.add_(old_exp_avg_sq_values) + denom = exp_avg_sq_update_values.sqrt_().add_(eps) + del exp_avg_update_values, exp_avg_sq_update_values + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + step_size = lr * math.sqrt(bias_correction2) / bias_correction1 + + param.add_(make_sparse(-step_size * numer.div_(denom))) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_multi_tensor/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_multi_tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..b6818e5a50f3b6792e249dc6e45b64bd29d0a067 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_multi_tensor/__init__.py @@ -0,0 +1,31 @@ +""" +:mod:`torch.optim._multi_tensor` is a package implementing various optimization algorithms. + +Most commonly used methods are already supported, and the interface is general +enough, so that more sophisticated ones can be also easily integrated in the +future. +""" + +from functools import partialmethod + +from torch import optim + + +def partialclass(cls, *args, **kwargs): # noqa: D103 + class NewCls(cls): + __init__ = partialmethod(cls.__init__, *args, **kwargs) + + return NewCls + + +Adam = partialclass(optim.Adam, foreach=True) +AdamW = partialclass(optim.AdamW, foreach=True) +NAdam = partialclass(optim.NAdam, foreach=True) +SGD = partialclass(optim.SGD, foreach=True) +RAdam = partialclass(optim.RAdam, foreach=True) +RMSprop = partialclass(optim.RMSprop, foreach=True) +Rprop = partialclass(optim.Rprop, foreach=True) +ASGD = partialclass(optim.ASGD, foreach=True) +Adamax = partialclass(optim.Adamax, foreach=True) +Adadelta = partialclass(optim.Adadelta, foreach=True) +Adagrad = partialclass(optim.Adagrad, foreach=True) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_multi_tensor/__init__.pyi b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_multi_tensor/__init__.pyi new file mode 100644 index 0000000000000000000000000000000000000000..97c3e2df989303c0f4a1cf76977cc47e25dfaaf8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_multi_tensor/__init__.pyi @@ -0,0 +1,15 @@ +from functools import partial + +from torch import optim + +Adam = partial(optim.Adam, foreach=True) +AdamW = partial(optim.AdamW, foreach=True) +NAdam = partial(optim.NAdam, foreach=True) +SGD = partial(optim.SGD, foreach=True) +RAdam = partial(optim.RAdam, foreach=True) +RMSprop = partial(optim.RMSprop, foreach=True) +Rprop = partial(optim.Rprop, foreach=True) +ASGD = partial(optim.ASGD, foreach=True) +Adamax = partial(optim.Adamax, foreach=True) +Adadelta = partial(optim.Adadelta, foreach=True) +Adagrad = partial(optim.Adagrad, foreach=True) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_muon.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_muon.py new file mode 100644 index 0000000000000000000000000000000000000000..2e45e07c4a596fb93f435130c344bb634ee0541c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/_muon.py @@ -0,0 +1,383 @@ +# mypy: allow-untyped-defs +# mypy: disable-error-code=arg-type +"""Implementation of the Muon optimizer.""" + +import math +from collections.abc import MutableMapping + +import torch +from torch import Tensor + +from .optimizer import ( + _disable_dynamo_if_unsupported, + _params_doc, + _to_scalar, + Optimizer, + ParamsT, +) + + +__all__ = ["Muon"] + +# Constants from Keller Jordan's Muon post: https://kellerjordan.github.io/posts/muon/ +# github permlink: https://github.com/KellerJordan/Muon/blob/f90a42b28e00b8d9d2d05865fe90d9f39abcbcbd/muon.py#L16 +EPS = 1e-7 +DEFAULT_A = 3.4445 +DEFAULT_B = -4.7750 +DEFAULT_C = 2.0315 +DEFAULT_NS_STEPS = 5 + + +def _zeropower_via_newtonschulz( + grad: Tensor, ns_coefficients: tuple[float, float, float], ns_steps: int, eps: float +) -> Tensor: + """ + Newton-Schulz iteration to compute the zeroth power / orthogonalization of G. We opt to use a + quintic iteration whose coefficients are selected to maximize the slope at zero. For the purpose + of minimizing steps, it turns out to be empirically effective to keep increasing the slope at + zero even beyond the point where the iteration no longer converges all the way to one everywhere + on the interval. This iteration therefore does not produce UV^T but rather something like US'V^T + where S' is diagonal with S_{ii}' ~ Uniform(0.5, 1.5), which turns out not to hurt model + performance at all relative to UV^T, where USV^T = G is the SVD. + + Implementation reference: https://github.com/KellerJordan/Muon/blob/master/muon.py + with suggestions by @jxbz, @leloykun, and @YouJiacheng. + """ + if ns_steps >= 100: + raise ValueError( + "Number of steps must be less than 100 for computational efficiency" + ) + if len(grad.shape) != 2: + raise ValueError("Input tensor gradient must be a 2D matrix") + if len(ns_coefficients) != 3: + raise ValueError("Coefficients must be a tuple of exactly 3 values") + a, b, c = ns_coefficients + ortho_grad = grad.bfloat16() + if grad.size(0) > grad.size(1): + ortho_grad = ortho_grad.T + # Ensure spectral norm is at most 1 + ortho_grad.div_(ortho_grad.norm().clamp(min=eps)) + # Perform the NS iterations + for _ in range(ns_steps): + gram_matrix = ortho_grad @ ortho_grad.T + gram_update = torch.addmm( + gram_matrix, gram_matrix, gram_matrix, beta=b, alpha=c + ) + ortho_grad = torch.addmm(ortho_grad, gram_update, ortho_grad, beta=a) + + if grad.size(0) > grad.size(1): + ortho_grad = ortho_grad.T + return ortho_grad + + +def _adjust_lr(lr: float, adjust_lr_fn: str | None, param_shape: torch.Size) -> float: + """Default learning rate adjustment used by Muon.""" + A, B = param_shape[:2] + + if adjust_lr_fn is None or adjust_lr_fn == "original": + # pyrefly: ignore [no-matching-overload] + adjusted_ratio = math.sqrt(max(1, A / B)) + elif adjust_lr_fn == "match_rms_adamw": + adjusted_ratio = 0.2 * math.sqrt(max(A, B)) + else: + adjusted_ratio = 1.0 + return lr * adjusted_ratio + + +class Muon(Optimizer): + def __init__( + self, + params: ParamsT, + lr: float = 1e-3, + weight_decay: float = 0.1, + momentum: float = 0.95, + nesterov: bool = True, + ns_coefficients: tuple[float, float, float] = (DEFAULT_A, DEFAULT_B, DEFAULT_C), + eps: float = EPS, + ns_steps: int = DEFAULT_NS_STEPS, + adjust_lr_fn: str | None = None, + ) -> None: + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Learning rate should be >= 0 but is: {lr}") + if not 0.0 <= momentum: + raise ValueError(f"momentum should be >= 0 but is: {momentum}") + if not 0.0 <= weight_decay: + raise ValueError(f"weight decay should be >= 0 but is: {weight_decay}") + if adjust_lr_fn is not None and adjust_lr_fn not in [ + "original", + "match_rms_adamw", + ]: + raise ValueError( + f"Adjust learning rate function {adjust_lr_fn} is not supported" + ) + + defaults = { + "lr": lr, + "weight_decay": weight_decay, + "momentum": momentum, + "nesterov": nesterov, + "ns_coefficients": ns_coefficients, + "eps": eps, + "ns_steps": ns_steps, + "adjust_lr_fn": adjust_lr_fn, + } + super().__init__(params, defaults) + + for group in self.param_groups: + for p in group["params"]: + if p.ndim != 2: + raise ValueError( + f"Muon only supports 2D parameters whereas we found a parameter with size: {p.size()}" + ) + + def _init_group( + self, + group: MutableMapping, + params_with_grad: list[Tensor], + grads: list[Tensor], + muon_momentum_bufs: list[Tensor], + ) -> bool: + for p in group["params"]: + if p.grad is None: + continue + + if torch.is_complex(p): + raise RuntimeError("Muon does not support complex parameters") + if p.grad.is_sparse: + raise RuntimeError("Muon does not support sparse gradients") + + params_with_grad.append(p) + grads.append(p.grad) + + state = self.state[p] + + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like( + p.grad, memory_format=torch.preserve_format + ) + muon_momentum_bufs.append(state["momentum_buffer"]) + + return False # has_complex + + @torch.no_grad() + def step(self, closure=None): + """Performs a single optimization step.""" + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + lr = group["lr"] + weight_decay = group["weight_decay"] + momentum = group["momentum"] + + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + muon_momentum_bufs: list[Tensor] = [] + + has_complex = self._init_group( + group, + params_with_grad, + grads, + muon_momentum_bufs, + ) + + muon( + params_with_grad, + grads, + muon_momentum_bufs, + lr=lr, + weight_decay=weight_decay, + momentum=momentum, + nesterov=group["nesterov"], + ns_coefficients=group["ns_coefficients"], + eps=group["eps"], + ns_steps=group["ns_steps"], + adjust_lr_fn=group["adjust_lr_fn"], + has_complex=has_complex, + ) + return loss + + +Muon.__doc__ = ( + r"""Implements Muon algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{ (lr)},\ \lambda \text{ (weight decay)},\ + \mu \text{ (momentum)},\ \textit{nesterov}\in\{True,False\},\\ + &\hspace{13mm}(a,b,c)\ \text{ (NS coefficients)},\ + \varepsilon \text{ (epsilon)},\ k \text{ (NS steps)},\ + \theta_0 \text{ (params)},\ f(\theta) \text{ (objective)} \\ + &\textbf{initialize} : B_0 \leftarrow 0 \text{ (momentum buffer)} \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for}\ t=1\ \textbf{to}\ \ldots\ \textbf{do} \\[0.25ex] + &\hspace{5mm} g_t \leftarrow \nabla_{\theta} f_t(\theta_{t-1}) \\[0.25ex] + &\hspace{5mm} B_t \leftarrow \mu B_{t-1} + g_t \\[0.25ex] + &\hspace{5mm} \widetilde{B}_t \leftarrow + \begin{cases} + g_t + \mu B_t, & \text{if nesterov}=True \\ + B_t, & \text{if nesterov}=False + \end{cases} \\[1.0ex] + &\hspace{5mm} O_t \leftarrow \mathrm{NS}^{(a,b,c)}_{k}\!\big(\widetilde{B}_t;\ \varepsilon\big) \\[0.5ex] + &\hspace{5mm} \theta_t \leftarrow \theta_{t-1} - \gamma\,\lambda\,\theta_{t-1} + \quad\text{(decoupled weight decay)} \\[0.25ex] + + &\hspace{5mm} \gamma \leftarrow \mathrm{AdjustLR}\!\big(\gamma;\ \mathrm{shape}\!\big(\theta_t \big) \big) \\[0.25ex] + &\hspace{5mm} \theta_t \leftarrow \theta_t - \gamma\, O_t \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\mathbf{return}\ \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt}s + \end{aligned} + + Here, :math:`\mathrm{NS}^{(a,b,c)}_{k}(\cdot;\varepsilon)` denotes :math:`k` iterations of the + Newton–Schulz orthogonalization operator parameterized by coefficients :math:`(a,b,c)` + with numerical stabilization :math:`\varepsilon`. + + The purpose for :math:`\mathrm{AdjustLR}\!\big(\gamma;\ \mathrm{shape}\!\big(\theta_t \big) \big)` + is to make the orthogonalized update have a consistent :math:`RMS` across rectangular matrices. + + Keller's original implementation scales the update by :math:`\sqrt{\max\!\left(1, \frac{A}{B}\right)}`, + where :math:`A` and :math:`B` are dimension of the matrix being optimized. + + Moonshot's implementation also focuses on matching :math:`RMS` of AdamW. The adjustment is computed as: + :math:`\gamma \leftarrow {0.2}\gamma\,\sqrt{\max\!\left({A}, {B}\right)}` + The method is adopted from `Muon is Scalable for LLM Training`_. Research + results show that with this adjustment Muon can directly reuse the learning rate + and weight decay tuned for AdamW. + + We provide two options for the learning rate adjustment: "original", which follows Keller's + implementation, and "match_rms_adamw", which refers to Moonshot's implementation. This gives users the + flexibility to choose between the two. If `adjust_lr_fn` is not specified, the default is "original". + + For further details regarding the algorithm we refer to `Muon: An optimizer for hidden layers in neural networks`_ + and `Muon is Scalable for LLM Training`_. + """ + + rf""" + Args: + {_params_doc}. Note that Muon is an optimizer for 2D parameters of neural network hidden layers. Other + parameters, such as bias, and embedding, should be optimized by a standard method such as AdamW. + lr (float, Tensor, optional): learning rate (default: 1e-3). + weight_decay (float, optional): weight decay (L2 penalty). (default: 0.1) + momentum (float, optional): momentum factor (default: 0.95) + nesterov (bool, optional): enables Nesterov momentum. Only applicable + when momentum is non-zero + ns_coefficients (tuple of three floats, optional): coefficients \(a,b,c\) for the + Newton–Schulz orthogonalization polynomial (default: ({DEFAULT_A}, {DEFAULT_B}, {DEFAULT_C})) + eps (float, optional): term added to the denominator for numerical stability. (default: {EPS}) + ns_steps (int, optional): number of Newton–Schulz iteration steps. (default: {DEFAULT_NS_STEPS}) + adjust_lr_fn (str, optional): function to adjust learning rate. One of "original" and "match_rms_adamw". + If not specified, we will default to use "original". (default: None) + + Example: + >>> # xdoctest: +SKIP + >>> # Muon only supports 2D params; use a standard optimizer + >>> # such as AdamW for biases, embeddings, and other non-2D + >>> # parameters. + >>> muon_params = [ + ... p for p in model.parameters() if p.ndim == 2 + ... ] + >>> other_params = [ + ... p for p in model.parameters() if p.ndim != 2 + ... ] + >>> optim_muon = torch.optim.Muon( + ... muon_params, lr=0.02, momentum=0.95 + ... ) + >>> optim_adamw = torch.optim.AdamW( + ... other_params, lr=3e-4, weight_decay=0.01 + ... ) + >>> optim_muon.zero_grad() + >>> optim_adamw.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optim_muon.step() + >>> optim_adamw.step() + + .. _Muon\: An optimizer for hidden layers in neural networks: + https://kellerjordan.github.io/posts/muon/ + .. _Muon is Scalable for LLM Training: + https://arxiv.org/pdf/2502.16982 + + """ +) + + +def _single_tensor_muon( + params: list[Tensor], + grads: list[Tensor], + muon_momentum_bufs: list[Tensor], + *, + lr: float, + weight_decay: float, + momentum: float, + nesterov: bool, + ns_coefficients: tuple[float, float, float], + ns_steps: int, + eps: float, + adjust_lr_fn: str | None, + has_complex: bool, +) -> None: + lr = _to_scalar(lr) + if has_complex: + raise ValueError("Complex parameters are not supported") + + for i, param in enumerate(params): + grad = grads[i] + if grad.ndim != 2: + raise ValueError("Param gradient must be a 2D matrix") + + buf = muon_momentum_bufs[i] + buf.lerp_(grad, 1 - momentum) + update = grad.lerp(buf, momentum) if nesterov else buf + + update = _zeropower_via_newtonschulz(update, ns_coefficients, ns_steps, eps) + + adjusted_lr = _adjust_lr(lr, adjust_lr_fn, param.shape) + + param.mul_(1 - lr * weight_decay) + param.add_(update, alpha=-adjusted_lr) + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_muon) +def muon( + params: list[Tensor], + grads: list[Tensor], + muon_momentum_bufs: list[Tensor], + *, + foreach: bool | None = None, + lr: float, + weight_decay: float, + momentum: float, + nesterov: bool, + ns_coefficients: tuple[float, float, float], + ns_steps: int, + eps: float, + adjust_lr_fn: str | None, + has_complex: bool, +) -> None: + r"""Functional API that performs Muon algorithm computation. + + See :class:`~torch.optim.Muon` for details. + """ + if foreach is not None and foreach: + raise RuntimeError("Foreach is not supported for Muon yet") + + func = _single_tensor_muon + + func( + params, + grads, + muon_momentum_bufs, + lr=lr, + weight_decay=weight_decay, + momentum=momentum, + nesterov=nesterov, + ns_coefficients=ns_coefficients, + ns_steps=ns_steps, + eps=eps, + adjust_lr_fn=adjust_lr_fn, + has_complex=has_complex, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adadelta.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adadelta.py new file mode 100644 index 0000000000000000000000000000000000000000..63b54583d3029aa6e0ad08d19289bb66f3384452 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adadelta.py @@ -0,0 +1,473 @@ +# mypy: allow-untyped-defs +from typing import Any, cast + +import torch +from torch import Tensor + +from .optimizer import ( + _capturable_doc, + _default_to_fused_or_foreach, + _differentiable_doc, + _disable_dynamo_if_unsupported, + _foreach_doc, + _get_capturable_supported_devices, + _get_scalar_dtype, + _maximize_doc, + _params_doc, + _to_scalar, + _use_grad_for_differentiable, + _view_as_real, + Optimizer, + ParamsT, +) + + +__all__ = ["Adadelta", "adadelta"] + + +class Adadelta(Optimizer): + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1.0, + rho: float = 0.9, + eps: float = 1e-6, + weight_decay: float = 0, + foreach: bool | None = None, + *, + capturable: bool = False, + maximize: bool = False, + differentiable: bool = False, + ) -> None: + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= rho <= 1.0: + raise ValueError(f"Invalid rho value: {rho}") + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + + defaults = { + "lr": lr, + "rho": rho, + "eps": eps, + "weight_decay": weight_decay, + "maximize": maximize, + "capturable": capturable, + "foreach": foreach, + "differentiable": differentiable, + } + super().__init__(params, defaults) + + def __setstate__(self, state): + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("foreach", None) + group.setdefault("maximize", False) + group.setdefault("differentiable", False) + group.setdefault("capturable", False) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = ( + torch.tensor( + step_val, dtype=_get_scalar_dtype(), device=p.device + ) + if group["capturable"] + else torch.tensor(step_val, dtype=_get_scalar_dtype()) + ) + + def _init_group( + self, + group: dict[str, Any], + params_with_grad: list[Tensor], + grads: list[Tensor], + square_avgs: list[Tensor], + acc_deltas: list[Tensor], + state_steps: list[Tensor], + ): + has_complex = False + p: Tensor + for p in group["params"]: + if p.grad is None: + continue + has_complex |= torch.is_complex(p) + params_with_grad.append(p) + if p.grad.is_sparse: + raise RuntimeError("Adadelta does not support sparse gradients") + grads.append(p.grad) + + state = self.state[p] + + # Lazy state initialization + if len(state) == 0: + state["step"] = ( + torch.zeros((), dtype=_get_scalar_dtype(), device=p.device) + if group["capturable"] + else torch.zeros((), dtype=_get_scalar_dtype()) + ) + + state["square_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + state["acc_delta"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + square_avgs.append(state["square_avg"]) + acc_deltas.append(state["acc_delta"]) + state_steps.append(state["step"]) + + return has_complex + + @_use_grad_for_differentiable + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._accelerator_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + square_avgs: list[Tensor] = [] + acc_deltas: list[Tensor] = [] + state_steps: list[Tensor] = [] + ( + lr, + rho, + eps, + weight_decay, + foreach, + maximize, + differentiable, + capturable, + ) = ( + group["lr"], + group["rho"], + group["eps"], + group["weight_decay"], + group["foreach"], + group["maximize"], + group["differentiable"], + group["capturable"], + ) + + has_complex = self._init_group( + group, params_with_grad, grads, square_avgs, acc_deltas, state_steps + ) + + adadelta( + params_with_grad, + grads, + square_avgs, + acc_deltas, + state_steps, + lr=lr, + rho=rho, + eps=eps, + weight_decay=weight_decay, + foreach=foreach, + maximize=maximize, + differentiable=differentiable, + capturable=capturable, + has_complex=has_complex, + ) + + return loss + + +Adadelta.__doc__ = ( + r"""Implements Adadelta algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{ (lr)}, \: \theta_0 \text{ (params)}, + \: f(\theta) \text{ (objective)}, \: \rho \text{ (decay)}, + \: \lambda \text{ (weight decay)} \\ + &\textbf{initialize} : v_0 \leftarrow 0 \: \text{ (square avg)}, + \: u_0 \leftarrow 0 \: \text{ (accumulate variables)} \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + &\hspace{5mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}if \: \lambda \neq 0 \\ + &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ + &\hspace{5mm} v_t \leftarrow v_{t-1} \rho + g^2_t (1 - \rho) \\ + &\hspace{5mm}\Delta x_t \leftarrow \frac{\sqrt{u_{t-1} + + \epsilon }}{ \sqrt{v_t + \epsilon} }g_t \hspace{21mm} \\ + &\hspace{5mm} u_t \leftarrow u_{t-1} \rho + + \Delta x^2_t (1 - \rho) \\ + &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \gamma \Delta x_t \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to `ADADELTA: An Adaptive Learning Rate Method`_. + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): coefficient that scale delta before it is applied + to the parameters (default: 1.0) + rho (float, optional): coefficient used for computing a running average + of squared gradients (default: 0.9). A higher value of `rho` will + result in a slower average, which can be helpful for preventing + oscillations in the learning process. + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-6). + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + {_foreach_doc} + {_capturable_doc} + {_maximize_doc} + {_differentiable_doc} + + .. _ADADELTA\: An Adaptive Learning Rate Method: + https://arxiv.org/abs/1212.5701 + + """ +) + + +def _single_tensor_adadelta( + params: list[Tensor], + grads: list[Tensor], + square_avgs: list[Tensor], + acc_deltas: list[Tensor], + state_steps: list[Tensor], + *, + lr: float, + rho: float, + eps: float, + weight_decay: float, + maximize: bool, + differentiable: bool, + capturable: bool, + has_complex: bool, +) -> None: + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices( + supports_xla=False + ) + if not all( + p.device.type == step.device.type + and p.device.type in capturable_supported_devices + for p, step in zip(params, state_steps, strict=True) + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + if not torch.jit.is_scripting(): + lr = _to_scalar(lr) + + for param, grad, square_avg, acc_delta, step in zip( + params, grads, square_avgs, acc_deltas, state_steps, strict=True + ): + step += 1 + grad = grad if not maximize else -grad + + if weight_decay != 0: + grad = grad.add(param, alpha=weight_decay) + + if torch.is_complex(param): + square_avg = torch.view_as_real(square_avg) + acc_delta = torch.view_as_real(acc_delta) + grad = torch.view_as_real(grad) + + square_avg.mul_(rho).addcmul_(grad, grad, value=1 - rho) + std = square_avg.add(eps).sqrt_() + delta = acc_delta.add(eps).sqrt_() + if differentiable: + delta = delta.clone() + delta.div_(std).mul_(grad) + acc_delta.mul_(rho).addcmul_(delta, delta, value=1 - rho) + + if torch.is_complex(param): + delta = torch.view_as_complex(delta) + param.add_(delta, alpha=-lr) + + +def _multi_tensor_adadelta( + params: list[Tensor], + grads: list[Tensor], + square_avgs: list[Tensor], + acc_deltas: list[Tensor], + state_steps: list[Tensor], + *, + lr: float, + rho: float, + eps: float, + weight_decay: float, + maximize: bool, + differentiable: bool, + capturable: bool, + has_complex: bool, +) -> None: + if differentiable: + raise AssertionError("_foreach ops don't support autograd") + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices( + supports_xla=False + ) + if not all( + p.device.type == step.device.type + and p.device.type in capturable_supported_devices + for p, step in zip(params, state_steps, strict=True) + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + if len(params) == 0: + return + + lr = _to_scalar(lr) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, square_avgs, acc_deltas, state_steps] # type: ignore[list-item] + ) + for ( + device_params_, + device_grads_, + device_square_avgs_, + device_acc_deltas_, + device_state_steps_, + ), _ in grouped_tensors.values(): + device_params = cast(list[Tensor], device_params_) + device_grads = cast(list[Tensor], device_grads_) + device_square_avgs = cast(list[Tensor], device_square_avgs_) + device_acc_deltas = cast(list[Tensor], device_acc_deltas_) + device_state_steps = cast(list[Tensor], device_state_steps_) + if has_complex: + _view_as_real( + device_params, device_grads, device_square_avgs, device_acc_deltas + ) + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and device_state_steps[0].is_cpu: + torch._foreach_add_( + device_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(device_state_steps, 1) + + if maximize: + device_grads = torch._foreach_neg(device_grads) # type: ignore[assignment] + + if weight_decay != 0: + # Reuse the intermediate memory (device_grads) already allocated for maximize + if maximize: + torch._foreach_add_(device_grads, device_params, alpha=weight_decay) + else: + device_grads = torch._foreach_add( # type: ignore[assignment] + device_grads, device_params, alpha=weight_decay + ) + + torch._foreach_mul_(device_square_avgs, rho) + torch._foreach_addcmul_( + device_square_avgs, device_grads, device_grads, value=1 - rho + ) + + std = torch._foreach_add(device_square_avgs, eps) + torch._foreach_sqrt_(std) + + deltas = torch._foreach_add(device_acc_deltas, eps) + torch._foreach_sqrt_(deltas) + torch._foreach_div_(deltas, std) + torch._foreach_mul_(deltas, device_grads) + + torch._foreach_mul_(device_acc_deltas, rho) + torch._foreach_addcmul_(device_acc_deltas, deltas, deltas, value=1 - rho) + + # If LR is a tensor, the else branch will internally call item() + # which will cause silent incorrectness if we are capturing + if capturable and isinstance(lr, torch.Tensor): + torch._foreach_mul_(deltas, -lr) + torch._foreach_add_(device_params, deltas) + else: + torch._foreach_add_(device_params, deltas, alpha=-lr) + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_adadelta) +def adadelta( + params: list[Tensor], + grads: list[Tensor], + square_avgs: list[Tensor], + acc_deltas: list[Tensor], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + capturable: bool = False, + foreach: bool | None = None, + differentiable: bool = False, + has_complex: bool = False, + *, + lr: float, + rho: float, + eps: float, + weight_decay: float, + maximize: bool, +) -> None: + r"""Functional API that performs Adadelta algorithm computation. + + See :class:`~torch.optim.Adadelta` for details. + """ + + # this check is slow during compilation, so we skip it + # if it's strictly needed we can add this check back in dynamo + if not torch.compiler.is_compiling() and not all( + isinstance(t, torch.Tensor) for t in state_steps + ): + raise RuntimeError( + "API has changed, `state_steps` argument must contain a list of singleton tensors" + ) + + # We still respect when the user inputs False for foreach. + if foreach is None: + _, foreach = _default_to_fused_or_foreach( + params, differentiable, use_fused=False + ) + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + + if foreach and not torch.jit.is_scripting(): + func = _multi_tensor_adadelta + else: + func = _single_tensor_adadelta + + func( + params, + grads, + square_avgs, + acc_deltas, + state_steps, + lr=lr, + rho=rho, + eps=eps, + weight_decay=weight_decay, + maximize=maximize, + differentiable=differentiable, + capturable=capturable, + has_complex=has_complex, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adagrad.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adagrad.py new file mode 100644 index 0000000000000000000000000000000000000000..d18e32fe17a9861088f4c598b3cb9ff0d7e6be36 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adagrad.py @@ -0,0 +1,626 @@ +# mypy: allow-untyped-defs +from typing import cast + +import torch +from torch import Tensor + +from .optimizer import ( + _default_to_fused_or_foreach, + _device_dtype_check_for_fused, + _differentiable_doc, + _foreach_doc, + _get_scalar_dtype, + _get_value, + _maximize_doc, + _params_doc, + _to_scalar, + _use_grad_for_differentiable, + _view_as_real, + DeviceDict, + Optimizer, + ParamsT, +) + + +__all__ = ["Adagrad", "adagrad"] + + +class Adagrad(Optimizer): + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-2, + lr_decay: float = 0, + weight_decay: float = 0, + initial_accumulator_value: float = 0, + eps: float = 1e-10, + foreach: bool | None = None, + *, + maximize: bool = False, + differentiable: bool = False, + fused: bool | None = None, + ) -> None: + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= lr_decay: + raise ValueError(f"Invalid lr_decay value: {lr_decay}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + if not 0.0 <= initial_accumulator_value: + raise ValueError( + f"Invalid initial_accumulator_value value: {initial_accumulator_value}" + ) + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + + defaults = { + "lr": lr, + "lr_decay": lr_decay, + "eps": eps, + "weight_decay": weight_decay, + "initial_accumulator_value": initial_accumulator_value, + "foreach": foreach, + "maximize": maximize, + "differentiable": differentiable, + "fused": fused, + } + super().__init__(params, defaults) + + if fused: + if differentiable: + raise RuntimeError("`fused` does not support `differentiable`") + if foreach: + raise RuntimeError("`fused` and `foreach` cannot be `True` together.") + self._need_device_dtype_check_for_fused = True + self._step_supports_amp_scaling = True + + for group in self.param_groups: + for p in group["params"]: + state = self.state[p] + state["step"] = ( + torch.zeros( + (), + dtype=_get_scalar_dtype(is_fused=group["fused"]), + device=p.device, + ) + if group["fused"] + else torch.tensor(0.0, dtype=_get_scalar_dtype()) + ) + init_value = ( + complex(initial_accumulator_value, initial_accumulator_value) + if torch.is_complex(p) + else initial_accumulator_value + ) + state["sum"] = torch.full_like( + p, init_value, memory_format=torch.preserve_format + ) + + def __setstate__(self, state): + super().__setstate__(state) + # define "fused" for + # MYPY error: Name "fused" may be undefined + fused = None + for group in self.param_groups: + group.setdefault("foreach", None) + group.setdefault("maximize", False) + group.setdefault("differentiable", False) + fused = group.setdefault("fused", None) + + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = ( + torch.tensor( + step_val, + dtype=_get_scalar_dtype(is_fused=fused), + device=p.device, + ) + if group["fused"] + else torch.tensor(step_val, dtype=_get_scalar_dtype()) + ) + + state_values = list(self.state.values()) + step_is_tensor = (len(state_values) != 0) and torch.is_tensor( + state_values[0]["step"] + ) + if not step_is_tensor: + for s in state_values: + s["step"] = torch.tensor( + float(s["step"]), dtype=_get_scalar_dtype(is_fused=fused) + ) + + def share_memory(self) -> None: + """Calls tensor.share_memory_() on the state sum tensors.""" + for group in self.param_groups: + for p in group["params"]: + state = self.state[p] + state["sum"].share_memory_() + + def _init_group(self, group, params_with_grad, grads, state_sums, state_steps): + has_sparse_grad, has_complex = False, False + for p in group["params"]: + if p.grad is not None: + if group["fused"] and getattr( + self, + "_need_device_dtype_check_for_fused", + True, + ): + _device_dtype_check_for_fused(p) + self._need_device_dtype_check_for_fused = False + has_sparse_grad |= p.grad.is_sparse + has_complex |= torch.is_complex(p) + params_with_grad.append(p) + grads.append(p.grad) + state = self.state[p] + if len(state) == 0: + if group["fused"]: + _device_dtype_check_for_fused(p) + + state["step"] = ( + torch.zeros( + (), + dtype=_get_scalar_dtype(is_fused=group["fused"]), + device=p.device, + ) + if group["fused"] + else torch.tensor(0.0, dtype=_get_scalar_dtype()) + ) + + initial_accumulator_value = self.defaults[ + "initial_accumulator_value" + ] + init_value = ( + complex(initial_accumulator_value, initial_accumulator_value) + if torch.is_complex(p) + else initial_accumulator_value + ) + state["sum"] = torch.full_like( + p, init_value, memory_format=torch.preserve_format + ) + state_sums.append(state["sum"]) + state_steps.append(state["step"]) + + return has_sparse_grad, has_complex + + @_use_grad_for_differentiable + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + state_sums: list[Tensor] = [] + state_steps: list[Tensor] = [] + + has_sparse_grad, has_complex = self._init_group( + group, params_with_grad, grads, state_sums, state_steps + ) + + adagrad( + params_with_grad, + grads, + state_sums, + state_steps, + lr=group["lr"], + weight_decay=group["weight_decay"], + lr_decay=group["lr_decay"], + eps=group["eps"], + has_sparse_grad=has_sparse_grad, + foreach=group["foreach"], + maximize=group["maximize"], + differentiable=group["differentiable"], + has_complex=has_complex, + fused=group["fused"], + grad_scale=getattr(self, "grad_scale", None), + found_inf=getattr(self, "found_inf", None), + ) + + return loss + + +Adagrad.__doc__ = ( + r"""Implements Adagrad algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{ (lr)}, \: \theta_0 \text{ (params)}, \: f(\theta) + \text{ (objective)}, \: \lambda \text{ (weight decay)}, \\ + &\hspace{12mm} \tau \text{ (initial accumulator value)}, \: \eta\text{ (lr decay)}\\ + &\textbf{initialize} : state\_sum_0 \leftarrow \tau \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + &\hspace{5mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm} \tilde{\gamma} \leftarrow \gamma / (1 +(t-1) \eta) \\ + &\hspace{5mm} \textbf{if} \: \lambda \neq 0 \\ + &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ + &\hspace{5mm}state\_sum_t \leftarrow state\_sum_{t-1} + g^2_t \\ + &\hspace{5mm}\theta_t \leftarrow + \theta_{t-1}- \tilde{\gamma} \frac{g_t}{\sqrt{state\_sum_t}+\epsilon} \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to `Adaptive Subgradient Methods for Online Learning + and Stochastic Optimization`_. + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 1e-2) + lr_decay (float, optional): learning rate decay (default: 0) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + initial_accumulator_value (float, optional): initial value of the + sum of squares of gradients (default: 0) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-10) + {_foreach_doc} + {_maximize_doc} + {_differentiable_doc} + fused (bool, optional): whether the fused implementation (CPU and CUDA only) is used. + Currently, `torch.float64`, `torch.float32`, `torch.float16`, and `torch.bfloat16` + are supported. (default: None). Please note that the fused implementations does not + support sparse or complex gradients. + .. _Adaptive Subgradient Methods for Online Learning and Stochastic + Optimization: http://jmlr.org/papers/v12/duchi11a.html + + """ +) + + +def adagrad( + params: list[Tensor], + grads: list[Tensor], + state_sums: list[Tensor], + state_steps: list[Tensor], + fused: bool | None = None, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting these as kwargs for now as functional API is compiled by torch/distributed/optim + has_sparse_grad: bool = False, + foreach: bool | None = None, + differentiable: bool = False, + has_complex: bool = False, + *, + lr: float, + weight_decay: float, + lr_decay: float, + eps: float, + maximize: bool, +) -> None: + r"""Functional API that performs Adagrad algorithm computation. + + See :class:`~torch.optim.Adagrad` for details. + """ + if not all(isinstance(t, torch.Tensor) for t in state_steps): + raise RuntimeError( + "API has changed, `state_steps` argument must contain a list of singleton tensors" + ) + + # Respect when the user inputs False/True for foreach or fused. We only want to change + # the default when neither have been user-specified. Note that we default to foreach + # and pass False to use_fused. This is not a mistake--we want to give the fused impl + # bake-in time before making it the default, even if it is typically faster. + if fused is None and foreach is None: + _, foreach = _default_to_fused_or_foreach( + params, differentiable, use_fused=False + ) + + if fused is None: + fused = False + if foreach is None: + foreach = False + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + if fused and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with fused optimizers") + + if fused and not torch.jit.is_scripting(): + func = _fused_adagrad + elif foreach and not torch.jit.is_scripting(): + func = _multi_tensor_adagrad + else: + func = _single_tensor_adagrad + + func( + params, + grads, + state_sums, + state_steps, + lr=lr, + weight_decay=weight_decay, + lr_decay=lr_decay, + eps=eps, + has_sparse_grad=has_sparse_grad, + maximize=maximize, + differentiable=differentiable, + has_complex=has_complex, + grad_scale=grad_scale, + found_inf=found_inf, + ) + + +def _make_sparse(grad, grad_indices, values): + size = grad.size() + return torch.sparse_coo_tensor(grad_indices, values, size) + + +def _single_tensor_adagrad( + params: list[Tensor], + grads: list[Tensor], + state_sums: list[Tensor], + state_steps: list[Tensor], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + lr: float, + weight_decay: float, + lr_decay: float, + eps: float, + has_sparse_grad: bool, + maximize: bool, + differentiable: bool, + has_complex: bool, +) -> None: + if grad_scale is not None or found_inf is not None: + raise AssertionError("Expected grad_scale and found_inf to be None") + + if not torch.jit.is_scripting(): + lr = _to_scalar(lr) + + for param, grad, state_sum, step_t in zip( + params, grads, state_sums, state_steps, strict=True + ): + # update step + step_t += 1 + step = _get_value(step_t) + grad = grad if not maximize else -grad + + if weight_decay != 0: + if grad.is_sparse: + raise RuntimeError( + "weight_decay option is not compatible with sparse gradients" + ) + grad = grad.add(param, alpha=weight_decay) + + clr = lr / (1 + (step - 1) * lr_decay) + + if grad.is_sparse: + grad = grad.coalesce() # the update is non-linear so indices must be unique + grad_indices = grad._indices() + grad_values = grad._values() + + state_sum.add_(_make_sparse(grad, grad_indices, grad_values.pow(2))) + std = state_sum.sparse_mask(grad) + std_values = std._values().sqrt_().add_(eps) + param.add_( + _make_sparse(grad, grad_indices, grad_values / std_values), alpha=-clr + ) + else: + is_complex = torch.is_complex(param) + if is_complex: + grad = torch.view_as_real(grad) + state_sum = torch.view_as_real(state_sum) + param = torch.view_as_real(param) + state_sum.addcmul_(grad, grad, value=1) + if differentiable: + std = state_sum.sqrt() + eps + else: + std = state_sum.sqrt().add_(eps) + param.addcdiv_(grad, std, value=-clr) + if is_complex: + param = torch.view_as_complex(param) + state_sum = torch.view_as_complex(state_sum) + + +def _multi_tensor_adagrad( + params: list[Tensor], + grads: list[Tensor], + state_sums: list[Tensor], + state_steps: list[Tensor], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + lr: float, + weight_decay: float, + lr_decay: float, + eps: float, + has_sparse_grad: bool, + maximize: bool, + differentiable: bool, + has_complex: bool, +) -> None: + if differentiable: + raise AssertionError("_foreach ops don't support autograd") + if grad_scale is not None or found_inf is not None: + raise AssertionError("Expected grad_scale and found_inf to be None") + + # Foreach functions will throw errors if given empty lists + if len(params) == 0: + return + + lr = _to_scalar(lr) + + grouped_tensorlists = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, state_sums, state_steps] # type: ignore[list-item] + ) + for ( + device_params_, + device_grads_, + device_state_sums_, + device_state_steps_, + ), _ in grouped_tensorlists.values(): + device_params = cast(list[Tensor], device_params_) + device_grads = cast(list[Tensor], device_grads_) + device_state_sums = cast(list[Tensor], device_state_sums_) + device_state_steps = cast(list[Tensor], device_state_steps_) + + device_has_sparse_grad = has_sparse_grad and any( + grad.is_sparse for grad in device_grads + ) + + if device_has_sparse_grad: + _single_tensor_adagrad( + device_params, + device_grads, + device_state_sums, + device_state_steps, + lr=lr, + weight_decay=weight_decay, + lr_decay=lr_decay, + eps=eps, + has_sparse_grad=True, + maximize=maximize, + differentiable=differentiable, + has_complex=has_complex, + grad_scale=grad_scale, + found_inf=found_inf, + ) + continue + + # Handle complex parameters + if has_complex: + _view_as_real(device_params, device_grads, device_state_sums) + + if maximize: + device_grads = torch._foreach_neg(device_grads) # type: ignore[assignment] + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and device_state_steps[0].is_cpu: + torch._foreach_add_( + device_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(device_state_steps, 1) + + if weight_decay != 0: + # Reuse the intermediate memory (device_grads) already allocated for maximize + if maximize: + torch._foreach_add_(device_grads, device_params, alpha=weight_decay) + else: + device_grads = torch._foreach_add( # type: ignore[assignment] + device_grads, device_params, alpha=weight_decay + ) + + minus_clr = [ + -lr / (1 + (_get_value(step) - 1) * lr_decay) for step in device_state_steps + ] + + torch._foreach_addcmul_(device_state_sums, device_grads, device_grads, value=1) + + std = torch._foreach_sqrt(device_state_sums) + torch._foreach_add_(std, eps) + + if weight_decay != 0 or maximize: + # Again, reuse the intermediate memory (device_grads) already allocated + torch._foreach_mul_(device_grads, minus_clr) + numerator = device_grads + else: + numerator = torch._foreach_mul(device_grads, minus_clr) # type: ignore[assignment] + + torch._foreach_addcdiv_(device_params, numerator, std) + + +def _fused_adagrad( + params: list[Tensor], + grads: list[Tensor], + state_sums: list[Tensor], + state_steps: list[Tensor], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + lr: float | Tensor, + weight_decay: float, + lr_decay: float, + eps: float, + has_sparse_grad: bool, + maximize: bool, + differentiable: bool, + has_complex: bool, +) -> None: + if not params: + return + if has_sparse_grad or has_complex: + raise RuntimeError("`fused` does not support sparse grad or complex param") + + if differentiable: + raise RuntimeError( + "adagrad with fused=True does not support differentiable=True" + ) + + grad_scale_dict: DeviceDict = ( + {grad_scale.device: grad_scale} if grad_scale is not None else {} + ) + found_inf_dict: DeviceDict = ( + {found_inf.device: found_inf} if found_inf is not None else {} + ) + lr_dict: DeviceDict | None = ( + {lr.device: lr} if isinstance(lr, Tensor) and str(lr.device) != "cpu" else None + ) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, state_sums, state_steps] # type: ignore[list-item] + ) + for (device, _), ( + ( + device_params_, + device_grads_, + device_state_sums_, + device_state_steps_, + ), + _, + ) in grouped_tensors.items(): + device_params = cast(list[Tensor], device_params_) + device_grads = cast(list[Tensor], device_grads_) + device_state_sums = cast(list[Tensor], device_state_sums_) + device_state_steps = cast(list[Tensor], device_state_steps_) + + device_grad_scale, device_found_inf = None, None + if grad_scale is not None: + device_grad_scale = grad_scale_dict.setdefault( + device, grad_scale.to(device, non_blocking=True) + ) + if found_inf is not None: + device_found_inf = found_inf_dict.setdefault( + device, found_inf.to(device, non_blocking=True) + ) + if lr_dict is not None and device not in lr_dict: + lr_dict[device] = lr.to(device=device, non_blocking=True) # type: ignore[union-attr] + lr = lr_dict[device] + torch._foreach_add_(device_state_steps, 1) + torch._fused_adagrad_( + device_params, + device_grads, + device_state_sums, + device_state_steps, + lr=lr, + lr_decay=lr_decay, + weight_decay=weight_decay, + eps=eps, + maximize=maximize, + grad_scale=device_grad_scale, + found_inf=device_found_inf, + ) + if device_found_inf is not None: + torch._foreach_sub_( + device_state_steps, [device_found_inf] * len(device_state_steps) + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adam.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adam.py new file mode 100644 index 0000000000000000000000000000000000000000..ca3750484c464d9ed905a12a49b75c9182c0643c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adam.py @@ -0,0 +1,990 @@ +# mypy: allow-untyped-defs +from typing import cast + +import torch +from torch import Tensor + +from .optimizer import ( + _capturable_doc, + _default_to_fused_or_foreach, + _device_dtype_check_for_fused, + _differentiable_doc, + _disable_dynamo_if_unsupported, + _foreach_doc, + _fused_doc, + _get_capturable_supported_devices, + _get_scalar_dtype, + _get_value, + _maximize_doc, + _params_doc, + _stack_if_compiling, + _to_scalar, + _use_grad_for_differentiable, + _view_as_real, + DeviceDict, + DeviceDtypeDict, + Optimizer, + ParamsT, +) + + +__all__ = ["Adam", "adam"] + + +class Adam(Optimizer): + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-3, + betas: tuple[float | Tensor, float | Tensor] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0, + amsgrad: bool = False, + *, + foreach: bool | None = None, + maximize: bool = False, + capturable: bool = False, + differentiable: bool = False, + fused: bool | None = None, + decoupled_weight_decay: bool = False, + ) -> None: + if isinstance(lr, Tensor): + if foreach and not capturable: + raise ValueError( + "lr as a Tensor is not supported for capturable=False and foreach=True" + ) + if lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + if not ( + (isinstance(betas[0], float) and isinstance(betas[1], float)) + or (isinstance(betas[0], Tensor) and isinstance(betas[1], Tensor)) + ): + raise ValueError("betas must be either both floats or both Tensors") + if isinstance(betas[0], Tensor): + if not capturable and foreach: + raise ValueError( + "betas[0] as a Tensor is not supported for capturable=False and foreach=True" + ) + if betas[0].numel() != 1: + raise ValueError("Tensor betas[0] must be 1-element") + if isinstance(betas[1], Tensor): + if not capturable and foreach: + raise ValueError( + "betas[1] as a Tensor is not supported for capturable=False and foreach=True" + ) + if betas[1].numel() != 1: + raise ValueError("Tensor betas[1] must be 1-element") + betas = tuple(map(_to_scalar, betas)) + + defaults = { + "lr": lr, + "betas": betas, + "eps": eps, + "weight_decay": weight_decay, + "amsgrad": amsgrad, + "maximize": maximize, + "foreach": foreach, + "capturable": capturable, + "differentiable": differentiable, + "fused": fused, + "decoupled_weight_decay": decoupled_weight_decay, + } + super().__init__(params, defaults) + + if fused: + if differentiable: + raise RuntimeError("`fused` does not support `differentiable`") + self._step_supports_amp_scaling = True + # TODO(crcrpar): [low prec params & their higher prec copy] + # Support AMP with FP16/BF16 model params which would need + # higher prec copy of params to do update math in higher prec to + # alleviate the loss of information. + if foreach: + raise RuntimeError("`fused` and `foreach` cannot be `True` together.") + + def __setstate__(self, state): + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("amsgrad", False) + group.setdefault("maximize", False) + group.setdefault("foreach", None) + group.setdefault("capturable", False) + group.setdefault("differentiable", False) + group.setdefault("decoupled_weight_decay", False) + fused = group.setdefault("fused", None) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = ( + torch.tensor( + step_val, + dtype=_get_scalar_dtype(is_fused=fused), + device=p.device, + ) + if group["capturable"] or group["fused"] + else torch.tensor(step_val, dtype=_get_scalar_dtype()) + ) + + def _init_group( + self, + group, + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + ): + has_complex = False + for p in group["params"]: + if p.grad is not None: + has_complex |= torch.is_complex(p) + params_with_grad.append(p) + if p.grad.is_sparse: + raise RuntimeError( + "Adam does not support sparse gradients, please consider SparseAdam instead" + ) + grads.append(p.grad) + + state = self.state[p] + # Lazy state initialization + if len(state) == 0: + if group["fused"]: + _device_dtype_check_for_fused(p) + # note(crcrpar): [special device hosting for step] + # Deliberately host `step` on CPU if both capturable and fused are off. + # This is because kernel launches are costly on CUDA and XLA. + state["step"] = ( + torch.zeros( + (), + dtype=_get_scalar_dtype(is_fused=group["fused"]), + device=p.device, + ) + if group["capturable"] or group["fused"] + else torch.tensor(0.0, dtype=_get_scalar_dtype()) + ) + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + if group["amsgrad"]: + # Maintains max of all exp. moving avg. of sq. grad. values + state["max_exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + exp_avgs.append(state["exp_avg"]) + exp_avg_sqs.append(state["exp_avg_sq"]) + + if group["amsgrad"]: + max_exp_avg_sqs.append(state["max_exp_avg_sq"]) + if group["differentiable"] and state["step"].requires_grad: + raise RuntimeError( + "`requires_grad` is not supported for `step` in differentiable mode" + ) + + # Foreach without capturable does not support a tensor lr + if ( + group["foreach"] + and torch.is_tensor(group["lr"]) + and not group["capturable"] + ): + raise RuntimeError( + "lr as a Tensor is not supported for capturable=False and foreach=True" + ) + + state_steps.append(state["step"]) + return has_complex + + @_use_grad_for_differentiable + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._accelerator_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + exp_avgs: list[Tensor] = [] + exp_avg_sqs: list[Tensor] = [] + max_exp_avg_sqs: list[Tensor] = [] + state_steps: list[Tensor] = [] + beta1, beta2 = group["betas"] + + has_complex = self._init_group( + group, + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + ) + + adam( + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + amsgrad=group["amsgrad"], + has_complex=has_complex, + beta1=beta1, + beta2=beta2, + lr=group["lr"], + weight_decay=group["weight_decay"], + eps=group["eps"], + maximize=group["maximize"], + foreach=group["foreach"], + capturable=group["capturable"], + differentiable=group["differentiable"], + fused=group["fused"], + grad_scale=getattr(self, "grad_scale", None), + found_inf=getattr(self, "found_inf", None), + decoupled_weight_decay=group["decoupled_weight_decay"], + ) + + return loss + + +Adam.__doc__ = ( + r"""Implements Adam algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{ (lr)}, \beta_1, \beta_2 + \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)} \\ + &\hspace{13mm} \lambda \text{ (weight decay)}, \: \textit{amsgrad}, + \:\textit{maximize}, \: \epsilon \text{ (epsilon)} \\ + &\textbf{initialize} : m_0 \leftarrow 0 \text{ ( first moment)}, + v_0\leftarrow 0 \text{ (second moment)},\: v_0^{max}\leftarrow 0 \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + + &\hspace{5mm}\textbf{if} \: \textit{maximize}: \\ + &\hspace{10mm}g_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\textbf{if} \: \lambda \neq 0 \\ + &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ + &\hspace{5mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ + &\hspace{5mm}v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) g^2_t \\ + &\hspace{5mm}\widehat{m_t} \leftarrow m_t/\big(1-\beta_1^t \big) \\ + &\hspace{5mm}\textbf{if} \: amsgrad \\ + &\hspace{10mm} v_t^{max} \leftarrow \mathrm{max}(v_{t-1}^{max},v_t) \\ + &\hspace{10mm}\widehat{v_t} \leftarrow v_t^{max}/\big(1-\beta_2^t \big) \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}\widehat{v_t} \leftarrow v_t/\big(1-\beta_2^t \big) \\ + &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/ + \big(\sqrt{\widehat{v_t}} + \epsilon \big) \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to `Adam: A Method for Stochastic Optimization`_. + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 1e-3). A tensor LR + is not yet supported for all our implementations. Please use a float + LR if you are not also specifying fused=True or capturable=True. + betas (tuple[float | Tensor, float | Tensor], optional): + coefficients used for computing running averages of gradient and + its square. If a tensor is provided, must be 1-element. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + decoupled_weight_decay (bool, optional): if True, this optimizer is + equivalent to AdamW and the algorithm will not accumulate weight + decay in the momentum nor variance. (default: False) + amsgrad (bool, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) + {_foreach_doc} + {_maximize_doc} + {_capturable_doc} + {_differentiable_doc} + {_fused_doc} + .. Note:: + A prototype implementation of Adam and AdamW for MPS supports `torch.float32` and `torch.float16`. + .. _Adam\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + + """ +) + + +def _single_tensor_adam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + max_exp_avg_sqs: list[Tensor], + state_steps: list[Tensor], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + amsgrad: bool, + has_complex: bool, + beta1: float | Tensor, + beta2: float | Tensor, + lr: float | Tensor, + weight_decay: float, + eps: float, + maximize: bool, + capturable: bool, + differentiable: bool, + decoupled_weight_decay: bool, +) -> None: + if grad_scale is not None or found_inf is not None: + raise AssertionError("Expected grad_scale and found_inf to be None") + + if torch.jit.is_scripting(): + # this assert is due to JIT being dumb and not realizing that the ops below + # have overloads to handle both float and Tensor lrs, so we just assert it's + # a float since most people using JIT are using floats + if not isinstance(lr, float): + raise AssertionError(f"Expected lr to be a float, but got {type(lr)}") + if not isinstance(beta1, float): + raise AssertionError(f"Expected beta1 to be a float, but got {type(beta1)}") + if not isinstance(beta2, float): + raise AssertionError(f"Expected beta2 to be a float, but got {type(beta2)}") + else: + lr = _to_scalar(lr) + beta1 = _to_scalar(beta1) + beta2 = _to_scalar(beta2) + + # We only shuffle around the beta when it is a Tensor, otherwise, we prefer + # treating it as a scalar. + # Note: ensure type declaration is under conditional check for isinstance + # or else torchscript will get cranky about the DeviceDict type. + if isinstance(beta1, Tensor): + beta1_dict: DeviceDtypeDict | None = {(beta1.device, beta1.dtype): beta1} + else: + beta1_dict = None + + for i, param in enumerate(params): + grad = grads[i] if not maximize else -grads[i] + exp_avg = exp_avgs[i] + exp_avg_sq = exp_avg_sqs[i] + step_t = state_steps[i] + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices() + if not ( + param.device.type == step_t.device.type + and param.device.type in capturable_supported_devices + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + # update step + step_t += 1 + + if weight_decay != 0: + if decoupled_weight_decay: + # Perform stepweight decay + param.mul_(1 - lr * weight_decay) + else: + # Nested if is necessary to bypass jitscript rules + if differentiable and isinstance(weight_decay, Tensor): + if weight_decay.requires_grad: + grad = grad.addcmul_(param.clone(), weight_decay) + else: + grad = grad.add(param, alpha=weight_decay) + else: + grad = grad.add(param, alpha=weight_decay) + + if torch.is_complex(param): + grad = torch.view_as_real(grad) + exp_avg = torch.view_as_real(exp_avg) + exp_avg_sq = torch.view_as_real(exp_avg_sq) + if amsgrad: + max_exp_avg_sqs[i] = torch.view_as_real(max_exp_avg_sqs[i]) + param = torch.view_as_real(param) + + device = param.device + + if beta1_dict is not None: + dtype = param.dtype # type: ignore[union-attr] + + # cast to workaround https://github.com/pytorch/pytorch/issues/140601 + key = (device, dtype) + if key not in beta1_dict: + beta1_dict[key] = beta1.to( # type: ignore[union-attr] + device=device, dtype=dtype, non_blocking=True + ) + + device_beta1: float | Tensor = beta1_dict[key] + else: + device_beta1 = beta1 + + # Decay the first and second moment running average coefficient + + exp_avg.lerp_(grad, 1 - device_beta1) + + # Nested if is necessary to bypass jitscript rules + if differentiable and isinstance(beta2, Tensor): + if beta2.requires_grad: + # Using lerp to only use 2 operations bc addcmul's value cannot be a tensor + # Showing equivalence of differentiable path and nondifferentiable path + # expavg * b2 + grad^2 * (1-b2) + # add expavg * (1-b2) - expavg * (1-b2) = 0 + # expavg * b2 + expavg * (1-b2) - expavg * (1-b2) + grad^2 * (1-b2) + # expavg - expavg * (1-b2) + grad^2 * (1-b2) + # expavg + (grad^2 - expavg) * (1-b2) + # expavg.lerp(grad^2, 1-beta2) + exp_avg_sq.lerp_(torch.square(grad), weight=1 - beta2) + else: + exp_avg_sq.mul_(beta2).addcmul_( + grad, grad, value=cast(float, 1 - beta2) + ) + else: + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) # type: ignore[arg-type] + + if capturable or differentiable: + step = step_t + + # Nested if is necessary to bypass jitscript rules + if differentiable and isinstance(beta1, Tensor): + if beta1.requires_grad: + bias_correction1 = 1 - beta1 ** step.clone() + else: + bias_correction1 = 1 - beta1**step + else: + bias_correction1 = 1 - beta1**step + + # Nested if is necessary to bypass jitscript rules + if differentiable and isinstance(beta2, Tensor): + if beta2.requires_grad: + bias_correction2 = 1 - beta2 ** step.clone() + else: + bias_correction2 = 1 - beta2**step + else: + bias_correction2 = 1 - beta2**step + + step_size = lr / bias_correction1 + step_size_neg = step_size.neg() + + bias_correction2_sqrt = bias_correction2.sqrt() + + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. till now + if differentiable: + max_exp_avg_sq = max_exp_avg_sqs[i].clone() + else: + max_exp_avg_sq = max_exp_avg_sqs[i] + + max_exp_avg_sqs[i].copy_(torch.maximum(max_exp_avg_sq, exp_avg_sq)) + + # Uses the max. for normalizing running avg. of gradient + # Folds in (admittedly ugly) 1-elem step_size math here to avoid extra param-set-sized read+write + # (can't fold it into addcdiv_ below because addcdiv_ requires value is a Number, not a Tensor) + denom = ( + max_exp_avg_sqs[i].sqrt() / (bias_correction2_sqrt * step_size_neg) + ).add_(eps / step_size_neg) + else: + denom = ( + exp_avg_sq.sqrt() / (bias_correction2_sqrt * step_size_neg) + ).add_(eps / step_size_neg) + + if differentiable: + param.addcdiv_(exp_avg.clone(), denom) + else: + param.addcdiv_(exp_avg, denom) + else: + step = _get_value(step_t) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + + step_size = lr / bias_correction1 + + bias_correction2_sqrt = bias_correction2**0.5 + + if amsgrad: + # Maintains the maximum of all 2nd moment running avg. till now + torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i]) + + # Use the max. for normalizing running avg. of gradient + denom = (max_exp_avg_sqs[i].sqrt() / bias_correction2_sqrt).add_(eps) + else: + denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps) + + param.addcdiv_(exp_avg, denom, value=-step_size) # type: ignore[arg-type] + + # Lastly, switch back to complex view + if amsgrad and torch.is_complex(params[i]): + max_exp_avg_sqs[i] = torch.view_as_complex(max_exp_avg_sqs[i]) + + +def _multi_tensor_adam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + max_exp_avg_sqs: list[Tensor], + state_steps: list[Tensor], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + amsgrad: bool, + has_complex: bool, + beta1: float | Tensor, + beta2: float | Tensor, + lr: float | Tensor, + weight_decay: float, + eps: float, + maximize: bool, + capturable: bool, + differentiable: bool, + decoupled_weight_decay: bool, +) -> None: + if len(params) == 0: + return + + if isinstance(lr, Tensor): + if not capturable: + raise RuntimeError( + "lr as a Tensor is not supported for capturable=False and foreach=True" + ) + if lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + + if isinstance(beta1, Tensor): + if not capturable: + raise ValueError( + "beta1 as a Tensor is not supported for capturable=False and foreach=True" + ) + if beta1.numel() != 1: + raise ValueError("Tensor beta1 must be 1-element") + + if isinstance(beta2, Tensor): + if not capturable: + raise ValueError( + "beta2 as a Tensor is not supported for capturable=False and foreach=True" + ) + if beta2.numel() != 1: + raise ValueError("Tensor beta2 must be 1-element") + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices( + supports_xla=False + ) + if not all( + p.device.type == step.device.type + and p.device.type in capturable_supported_devices + for p, step in zip(params, state_steps, strict=True) + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + if grad_scale is not None or found_inf is not None: + raise AssertionError("Expected grad_scale and found_inf to be None") + + if differentiable: + raise AssertionError("_foreach ops don't support autograd") + + lr = _to_scalar(lr) + beta1 = _to_scalar(beta1) + beta2 = _to_scalar(beta2) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps] # type: ignore[list-item] + ) + + # We only shuffle around the beta when it is a Tensor and on CUDA, otherwise, we prefer + # treating it as a scalar. + beta1_dict: DeviceDict | None = ( # type: ignore[attr-defined] + {beta1.device: beta1} + if isinstance(beta1, Tensor) and str(beta1.device) != "cpu" + else None + ) + + for ( + device_params_, + device_grads_, + device_exp_avgs_, + device_exp_avg_sqs_, + device_max_exp_avg_sqs_, + device_state_steps_, + ), _ in grouped_tensors.values(): + device_params = cast(list[Tensor], device_params_) + device_grads = cast(list[Tensor], device_grads_) + device_exp_avgs = cast(list[Tensor], device_exp_avgs_) + device_exp_avg_sqs = cast(list[Tensor], device_exp_avg_sqs_) + device_state_steps = cast(list[Tensor], device_state_steps_) + + device = device_params[0].device + if beta1_dict is not None and device not in beta1_dict: + beta1_dict[device] = beta1.to(device=device, non_blocking=True) # type: ignore[union-attr, attr-defined] + + device_beta1 = beta1_dict[device] if beta1_dict else beta1 + + # Handle complex parameters + if has_complex: + if amsgrad: + device_max_exp_avg_sqs = cast(list[Tensor], device_max_exp_avg_sqs_) + _view_as_real( + device_params, + device_grads, + device_exp_avgs, + device_exp_avg_sqs, + device_max_exp_avg_sqs, + ) + else: + _view_as_real( + device_params, device_grads, device_exp_avgs, device_exp_avg_sqs + ) + + if maximize: + device_grads = torch._foreach_neg(device_grads) # type: ignore[assignment] + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and device_state_steps[0].is_cpu: + torch._foreach_add_( + device_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(device_state_steps, 1) + + if weight_decay != 0: + if decoupled_weight_decay: + # Perform stepweight decay + torch._foreach_mul_(device_params, 1 - lr * weight_decay) + else: + # Reuse the intermediate memory (device_grads) already allocated for maximize + if maximize: + torch._foreach_add_(device_grads, device_params, alpha=weight_decay) + else: + device_grads = torch._foreach_add( # type: ignore[assignment] + device_grads, device_params, alpha=weight_decay + ) + + # Decay the first and second moment running average coefficient + # Use device beta1 if beta1 is a tensor to ensure all + # tensors are on the same device + torch._foreach_lerp_( + device_exp_avgs, device_grads, cast(float, 1 - device_beta1) + ) + + torch._foreach_mul_(device_exp_avg_sqs, beta2) + + # Due to the strictness of the _foreach_addcmul API, we can't have a single + # tensor scalar as the scalar arg (only python number is supported there) + # as a result, separate out the value mul + # Filed https://github.com/pytorch/pytorch/issues/139795 + if isinstance(beta2, torch.Tensor): + scaled_device_grads = torch._foreach_mul(device_grads, 1 - beta2) # type: ignore[assignment] + value = 1.0 + else: + scaled_device_grads = device_grads # type: ignore[assignment] + value = 1 - beta2 + + torch._foreach_addcmul_( + device_exp_avg_sqs, scaled_device_grads, device_grads, value + ) + + # Delete the local intermediate(s) since they won't be used anymore to save on peak memory + del device_grads + del scaled_device_grads + + bias_correction1: tuple[Tensor, ...] | list[Tensor] + bias_correction2: tuple[Tensor, ...] | list[Tensor] + bias_correction2_sqrt: tuple[Tensor, ...] | list[Tensor] + + if capturable: + bias_correction1 = torch._foreach_pow(beta1, device_state_steps) # type: ignore[arg-type] + bias_correction2 = torch._foreach_pow(beta2, device_state_steps) # type: ignore[arg-type] + # foreach_sub doesn't allow a scalar as the first arg + torch._foreach_sub_(bias_correction1, 1) + torch._foreach_sub_(bias_correction2, 1) + # we do not negate bias_correction1 as it'll need to be negated later anyway + torch._foreach_neg_(bias_correction2) + + # foreach_div doesn't allow a scalar as the first arg + torch._foreach_div_(bias_correction1, lr) + torch._foreach_reciprocal_(bias_correction1) + + torch._foreach_sqrt_(bias_correction2) + + # Re-assign for clarity as we maintain minimal intermediates: we'll have + # step_size = - lr / (1 - beta1 ^ t) where t = num_steps + # bias_correction2_sqrt = sqrt(1 - beta2 ^ t) + step_size = bias_correction1 + bias_correction2_sqrt = bias_correction2 + + if amsgrad: + device_max_exp_avg_sqs = cast(list[Tensor], device_max_exp_avg_sqs_) + # Maintains the maximum of all 2nd moment running avg. till now + torch._foreach_maximum_(device_max_exp_avg_sqs, device_exp_avg_sqs) # type: ignore[assignment] + + # Set intermediate to the max. for normalizing running avg. of gradient when amsgrad + exp_avg_sq_sqrt = torch._foreach_sqrt(device_max_exp_avg_sqs) + else: + exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs) + + torch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt) + torch._foreach_add_(exp_avg_sq_sqrt, eps) + torch._foreach_div_(exp_avg_sq_sqrt, step_size) + + # at this point, exp_avg_sq_sqrt = - (1 - beta^t) * [sqrt(exp_avg_sq / (1 - beta2^t)) + eps] / lr + torch._foreach_addcdiv_(device_params, device_exp_avgs, exp_avg_sq_sqrt) + else: + bias_correction1 = [ + 1 - beta1 ** _get_value(step) for step in device_state_steps + ] + bias_correction2 = [ + 1 - beta2 ** _get_value(step) for step in device_state_steps + ] + + step_size = _stack_if_compiling([(lr / bc) * -1 for bc in bias_correction1]) + + bias_correction2_sqrt = [bc**0.5 for bc in bias_correction2] # type: ignore[arg-type] + + if amsgrad: + device_max_exp_avg_sqs = cast(list[Tensor], device_max_exp_avg_sqs_) + # Maintains the maximum of all 2nd moment running avg. till now + torch._foreach_maximum_(device_max_exp_avg_sqs, device_exp_avg_sqs) + + # Use the max. for normalizing running avg. of gradient + exp_avg_sq_sqrt = torch._foreach_sqrt(device_max_exp_avg_sqs) + else: + exp_avg_sq_sqrt = torch._foreach_sqrt(device_exp_avg_sqs) + + torch._foreach_div_(exp_avg_sq_sqrt, bias_correction2_sqrt) + torch._foreach_add_(exp_avg_sq_sqrt, eps) + torch._foreach_addcdiv_( + device_params, + device_exp_avgs, + exp_avg_sq_sqrt, + step_size, # type: ignore[arg-type] + ) + + +def _fused_adam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + max_exp_avg_sqs: list[Tensor], + state_steps: list[Tensor], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + amsgrad: bool, + has_complex: bool, # Needed for consistency. + beta1: float | Tensor, + beta2: float | Tensor, + lr: float | Tensor, + weight_decay: float, + eps: float, + maximize: bool, + capturable: bool, # Needed for consistency. + differentiable: bool, + decoupled_weight_decay: bool, +) -> None: + if not params: + return + if differentiable: + raise RuntimeError("Adam with fused=True does not support differentiable=True") + + beta1 = _to_scalar(beta1) + beta2 = _to_scalar(beta2) + + grad_scale_dict: DeviceDict = ( + {grad_scale.device: grad_scale} if grad_scale is not None else {} + ) + found_inf_dict: DeviceDict = ( + {found_inf.device: found_inf} if found_inf is not None else {} + ) + + # We only shuffle around the lr when it is a Tensor and on CUDA, otherwise, we prefer + # treating it as a scalar. + lr_dict: DeviceDict | None = ( + {lr.device: lr} if isinstance(lr, Tensor) and str(lr.device) != "cpu" else None + ) + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, exp_avgs, exp_avg_sqs, max_exp_avg_sqs, state_steps] # type: ignore[list-item] + ) + for (device, _), ( + ( + device_params_, + device_grads_, + device_exp_avgs_, + device_exp_avg_sqs_, + device_max_exp_avg_sqs, + device_state_steps_, + ), + _, + ) in grouped_tensors.items(): + device_params = cast(list[Tensor], device_params_) + device_grads = cast(list[Tensor], device_grads_) + device_exp_avgs = cast(list[Tensor], device_exp_avgs_) + device_exp_avg_sqs = cast(list[Tensor], device_exp_avg_sqs_) + device_state_steps = cast(list[Tensor], device_state_steps_) + + device_grad_scale, device_found_inf = None, None + if grad_scale is not None: + device_grad_scale = grad_scale_dict.setdefault( + device, grad_scale.to(device, non_blocking=True) + ) + if found_inf is not None: + device_found_inf = found_inf_dict.setdefault( + device, found_inf.to(device, non_blocking=True) + ) + if lr_dict is not None and device not in lr_dict: + lr_dict[device] = lr.to(device=device, non_blocking=True) # type: ignore[union-attr] + lr = lr_dict[device] + torch._foreach_add_(device_state_steps, 1) + func = torch._fused_adam_ if not decoupled_weight_decay else torch._fused_adamw_ + # pyrefly: ignore [no-matching-overload] + func( + device_params, + device_grads, + device_exp_avgs, + device_exp_avg_sqs, + device_max_exp_avg_sqs, # type: ignore[arg-type] + device_state_steps, + amsgrad=amsgrad, + lr=lr, # type: ignore[arg-type] + beta1=beta1, + beta2=beta2, + weight_decay=weight_decay, + eps=eps, + maximize=maximize, + grad_scale=device_grad_scale, + found_inf=device_found_inf, + ) + if device_found_inf is not None: + torch._foreach_sub_( + device_state_steps, [device_found_inf] * len(device_state_steps) + ) + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_adam) +def adam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + max_exp_avg_sqs: list[Tensor], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + foreach: bool | None = None, + capturable: bool = False, + differentiable: bool = False, + fused: bool | None = None, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, + has_complex: bool = False, + decoupled_weight_decay: bool = False, + *, + amsgrad: bool, + beta1: float | Tensor, + beta2: float | Tensor, + lr: float | Tensor, + weight_decay: float, + eps: float, + maximize: bool, +) -> None: + r"""Functional API that performs Adam algorithm computation. + + See :class:`~torch.optim.Adam` for details. + """ + # Respect when the user inputs False/True for foreach or fused. We only want to change + # the default when neither have been user-specified. Note that we default to foreach + # and pass False to use_fused. This is not a mistake--we want to give the fused impl + # bake-in time before making it the default, even if it is typically faster. + if fused is None and foreach is None: + _, foreach = _default_to_fused_or_foreach( + params, differentiable, use_fused=False + ) + # Do not flip on foreach for the unsupported case where lr is a Tensor and capturable=False. + if foreach and isinstance(lr, Tensor) and not capturable: + foreach = False + if fused is None: + fused = False + if foreach is None: + foreach = False + + # this check is slow during compilation, so we skip it + # if it's strictly needed we can add this check back in dynamo + if not torch.compiler.is_compiling() and not all( + isinstance(t, torch.Tensor) for t in state_steps + ): + raise RuntimeError( + "API has changed, `state_steps` argument must contain a list of singleton tensors" + ) + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + if fused and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with fused optimizers") + + if fused and not torch.jit.is_scripting(): + func = _fused_adam + elif foreach and not torch.jit.is_scripting(): + func = _multi_tensor_adam + else: + func = _single_tensor_adam + + func( + params, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + amsgrad=amsgrad, + has_complex=has_complex, + beta1=beta1, + beta2=beta2, + lr=lr, + weight_decay=weight_decay, + eps=eps, + maximize=maximize, + capturable=capturable, + differentiable=differentiable, + grad_scale=grad_scale, + found_inf=found_inf, + decoupled_weight_decay=decoupled_weight_decay, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adamax.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adamax.py new file mode 100644 index 0000000000000000000000000000000000000000..fc354285a432ebf73f399d7cbe85ec57c34ac6b5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adamax.py @@ -0,0 +1,485 @@ +# mypy: allow-untyped-defs +from typing import cast + +import torch +from torch import Tensor + +from .optimizer import ( + _capturable_doc, + _default_to_fused_or_foreach, + _differentiable_doc, + _disable_dynamo_if_unsupported, + _foreach_doc, + _get_capturable_supported_devices, + _get_scalar_dtype, + _get_value, + _maximize_doc, + _params_doc, + _to_scalar, + _use_grad_for_differentiable, + _view_as_real, + Optimizer, + ParamsT, +) + + +__all__ = ["Adamax", "adamax"] + + +class Adamax(Optimizer): + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 2e-3, + betas: tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0, + foreach: bool | None = None, + *, + maximize: bool = False, + differentiable: bool = False, + capturable: bool = False, + ) -> None: + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + + defaults = { + "lr": lr, + "betas": betas, + "eps": eps, + "weight_decay": weight_decay, + "foreach": foreach, + "maximize": maximize, + "differentiable": differentiable, + "capturable": capturable, + } + super().__init__(params, defaults) + + def __setstate__(self, state): + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("foreach", None) + group.setdefault("maximize", False) + group.setdefault("differentiable", False) + group.setdefault("capturable", False) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = ( + torch.tensor( + step_val, dtype=_get_scalar_dtype(), device=p.device + ) + if group["capturable"] + else torch.tensor(step_val, dtype=_get_scalar_dtype()) + ) + + def _init_group( + self, group, params_with_grad, grads, exp_avgs, exp_infs, state_steps + ): + has_complex = False + for p in group["params"]: + if p.grad is None: + continue + has_complex |= torch.is_complex(p) + params_with_grad.append(p) + if p.grad.is_sparse: + raise RuntimeError("Adamax does not support sparse gradients") + grads.append(p.grad) + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = ( + torch.zeros((), dtype=_get_scalar_dtype(), device=p.device) + if group["capturable"] + else torch.tensor(0.0, dtype=_get_scalar_dtype()) + ) + state["exp_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + state["exp_inf"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + exp_avgs.append(state["exp_avg"]) + exp_infs.append(state["exp_inf"]) + state_steps.append(state["step"]) + + return has_complex + + @_use_grad_for_differentiable + def step(self, closure=None): + """Performs a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._accelerator_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + exp_avgs: list[Tensor] = [] + exp_infs: list[Tensor] = [] + state_steps: list[Tensor] = [] + + beta1, beta2 = group["betas"] + eps = group["eps"] + lr = group["lr"] + weight_decay = group["weight_decay"] + foreach = group["foreach"] + maximize = group["maximize"] + differentiable = group["differentiable"] + capturable = group["capturable"] + + has_complex = self._init_group( + group, params_with_grad, grads, exp_avgs, exp_infs, state_steps + ) + + adamax( + params_with_grad, + grads, + exp_avgs, + exp_infs, + state_steps, + eps=eps, + beta1=beta1, + beta2=beta2, + lr=lr, + weight_decay=weight_decay, + foreach=foreach, + maximize=maximize, + differentiable=differentiable, + capturable=capturable, + has_complex=has_complex, + ) + + return loss + + +Adamax.__doc__ = ( + r"""Implements Adamax algorithm (a variant of Adam based on infinity norm). + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{ (lr)}, \beta_1, \beta_2 + \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)}, + \: \lambda \text{ (weight decay)}, \\ + &\hspace{13mm} \epsilon \text{ (epsilon)} \\ + &\textbf{initialize} : m_0 \leftarrow 0 \text{ ( first moment)}, + u_0 \leftarrow 0 \text{ ( infinity norm)} \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + &\hspace{5mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}if \: \lambda \neq 0 \\ + &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ + &\hspace{5mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ + &\hspace{5mm}u_t \leftarrow \mathrm{max}(\beta_2 u_{t-1}, |g_{t}|+\epsilon) \\ + &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \frac{\gamma m_t}{(1-\beta^t_1) u_t} \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to `Adam: A Method for Stochastic Optimization`_. + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 2e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + {_foreach_doc} + {_maximize_doc} + {_differentiable_doc} + {_capturable_doc} + + .. _Adam\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + + """ +) + + +def _single_tensor_adamax( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_infs: list[Tensor], + state_steps: list[Tensor], + *, + eps: float, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + maximize: bool, + differentiable: bool, + capturable: bool, + has_complex: bool, +) -> None: + if not torch.jit.is_scripting(): + lr = _to_scalar(lr) + + for i, param in enumerate(params): + grad = grads[i] + grad = grad if not maximize else -grad + exp_avg = exp_avgs[i] + exp_inf = exp_infs[i] + step_t = state_steps[i] + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices() + if not ( + param.device.type == step_t.device.type + and param.device.type in capturable_supported_devices + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + # update step + step_t += 1 + + if weight_decay != 0: + grad = grad.add(param, alpha=weight_decay) + + if torch.is_complex(param): + param = torch.view_as_real(param) + grad = torch.view_as_real(grad) + exp_avg = torch.view_as_real(exp_avg) + exp_inf = torch.view_as_real(exp_inf) + + # Update biased first moment estimate. + exp_avg.lerp_(grad, 1 - beta1) + # Update the exponentially weighted infinity norm. + if not differentiable: + torch.maximum( + exp_inf.mul_(beta2), + grad.abs().add_(eps), + out=exp_inf, + ) + else: + norm_buf = torch.cat( + [exp_inf.mul_(beta2).unsqueeze(0), grad.abs().add_(eps).unsqueeze_(0)], + 0, + ) + exp_inf.copy_(torch.amax(norm_buf, 0, keepdim=False)) + + if capturable: + # why jump through extra hoops and negate bias_correction? check out #121238 + # once fixed, we should use bias_correction with addcdiv value=-1 for readability + neg_bias_correction = beta1**step_t - 1 + neg_bias_correction.div_(lr) + denom = exp_inf * neg_bias_correction + param.addcdiv_(exp_avg, denom) + else: + bias_correction = 1 - beta1 ** _get_value(step_t) + clr = lr / bias_correction + + param.addcdiv_(exp_avg, exp_inf, value=-clr) + + +def _multi_tensor_adamax( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_infs: list[Tensor], + state_steps: list[Tensor], + *, + eps: float, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + maximize: bool, + differentiable: bool, + capturable: bool, + has_complex: bool, +) -> None: + if differentiable: + raise AssertionError("_foreach ops don't support autograd") + + if len(params) == 0: + return + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices( + supports_xla=False + ) + if not all( + p.device.type == step.device.type + and p.device.type in capturable_supported_devices + for p, step in zip(params, state_steps, strict=True) + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + lr = _to_scalar(lr) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, exp_avgs, exp_infs, state_steps] # type: ignore[list-item] + ) + for ( + grouped_params_, + grouped_grads_, + grouped_exp_avgs_, + grouped_exp_infs_, + grouped_state_steps_, + ), _ in grouped_tensors.values(): + grouped_params = cast(list[Tensor], grouped_params_) + grouped_grads = cast(list[Tensor], grouped_grads_) + grouped_exp_avgs = cast(list[Tensor], grouped_exp_avgs_) + grouped_exp_infs = cast(list[Tensor], grouped_exp_infs_) + grouped_state_steps = cast(list[Tensor], grouped_state_steps_) + + if has_complex: + _view_as_real( + grouped_params, grouped_grads, grouped_exp_avgs, grouped_exp_infs + ) + + if maximize: + grouped_grads = torch._foreach_neg(grouped_grads) # type: ignore[assignment] + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and grouped_state_steps[0].is_cpu: + torch._foreach_add_( + grouped_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(grouped_state_steps, 1) + + if weight_decay != 0: + if maximize: + # Reuse the intermediate memory (grouped_grads) already allocated for maximize + torch._foreach_add_(grouped_grads, grouped_params, alpha=weight_decay) + else: + grouped_grads = torch._foreach_add( # type: ignore[assignment] + grouped_grads, grouped_params, alpha=weight_decay + ) + + # Update biased first moment estimate. + torch._foreach_lerp_(grouped_exp_avgs, grouped_grads, 1 - beta1) + + # Update the exponentially weighted infinity norm. + torch._foreach_mul_(grouped_exp_infs, beta2) + + # in this case, we need to introduce a copy of the grads + # since one has not been introduced previously + if not maximize and weight_decay == 0: + grouped_grads = torch._foreach_abs(grouped_grads) # type: ignore[assignment] + else: + torch._foreach_abs_(grouped_grads) + + torch._foreach_add_(grouped_grads, eps) + torch._foreach_maximum_(grouped_exp_infs, grouped_grads) + + bias_corrections: tuple[Tensor, ...] | list[Tensor] + if capturable: + bias_corrections = torch._foreach_pow(beta1, grouped_state_steps) + # foreach_sub doesn't allow a scalar as the first arg + torch._foreach_sub_(bias_corrections, 1) + torch._foreach_div_(bias_corrections, lr) + + denom = torch._foreach_mul(grouped_exp_infs, bias_corrections) + torch._foreach_addcdiv_(grouped_params, grouped_exp_avgs, denom) + else: + bias_corrections = [ + 1 - beta1 ** _get_value(step) for step in grouped_state_steps + ] + step_size = [(_get_value(lr) / bc) * -1 for bc in bias_corrections] + torch._foreach_addcdiv_( + grouped_params, grouped_exp_avgs, grouped_exp_infs, step_size + ) + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_adamax) +def adamax( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_infs: list[Tensor], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + foreach: bool | None = None, + maximize: bool = False, + differentiable: bool = False, + capturable: bool = False, + has_complex: bool = False, + *, + eps: float, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, +) -> None: + r"""Functional API that performs adamax algorithm computation. + + See :class:`~torch.optim.Adamax` for details. + """ + + if not torch.compiler.is_compiling() and not all( + isinstance(t, torch.Tensor) for t in state_steps + ): + raise RuntimeError( + "API has changed, `state_steps` argument must contain a list of singleton tensors" + ) + + if foreach is None: + _, foreach = _default_to_fused_or_foreach( + params, differentiable, use_fused=False + ) + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + + if foreach and not torch.jit.is_scripting(): + func = _multi_tensor_adamax + else: + func = _single_tensor_adamax + + func( + params, + grads, + exp_avgs, + exp_infs, + state_steps, + eps=eps, + beta1=beta1, + beta2=beta2, + lr=lr, + weight_decay=weight_decay, + maximize=maximize, + differentiable=differentiable, + has_complex=has_complex, + capturable=capturable, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adamw.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adamw.py new file mode 100644 index 0000000000000000000000000000000000000000..fbb3e453274694cc786fec1690137c0f15fde5f1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/adamw.py @@ -0,0 +1,181 @@ +# mypy: allow-untyped-defs + +from torch import Tensor + +from .adam import Adam, adam +from .optimizer import ( + _capturable_doc, + _differentiable_doc, + _foreach_doc, + _fused_doc, + _maximize_doc, + _params_doc, + ParamsT, +) + + +__all__ = ["AdamW", "adamw"] + + +class AdamW(Adam): + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-3, + betas: tuple[float | Tensor, float | Tensor] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 1e-2, + amsgrad: bool = False, + *, + maximize: bool = False, + foreach: bool | None = None, + capturable: bool = False, + differentiable: bool = False, + fused: bool | None = None, + ) -> None: + super().__init__( + params, + lr, + betas, + eps, + weight_decay, + amsgrad, + foreach=foreach, + maximize=maximize, + capturable=capturable, + differentiable=differentiable, + fused=fused, + decoupled_weight_decay=True, + ) + + # Preserve decoupled_weight_decay from AdamW for backwards compatibility. The following + # guarantees that decoupled_weight_decay will always be True for loading any state into + # AdamW + def __setstate__(self, state): + super().__setstate__(state) + for group in self.param_groups: + group["decoupled_weight_decay"] = True + + +AdamW.__doc__ = ( + r"""Implements AdamW algorithm, where weight decay does not accumulate in the momentum nor variance. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{(lr)}, \: \beta_1, \beta_2 + \text{(betas)}, \: \theta_0 \text{(params)}, \: f(\theta) \text{(objective)}, + \: \epsilon \text{ (epsilon)} \\ + &\hspace{13mm} \lambda \text{(weight decay)}, \: \textit{amsgrad}, + \: \textit{maximize} \\ + &\textbf{initialize} : m_0 \leftarrow 0 \text{ (first moment)}, v_0 \leftarrow 0 + \text{ ( second moment)}, \: v_0^{max}\leftarrow 0 \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + + &\hspace{5mm}\textbf{if} \: \textit{maximize}: \\ + &\hspace{10mm}g_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm} \theta_t \leftarrow \theta_{t-1} - \gamma \lambda \theta_{t-1} \\ + &\hspace{5mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ + &\hspace{5mm}v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) g^2_t \\ + &\hspace{5mm}\widehat{m_t} \leftarrow m_t/\big(1-\beta_1^t \big) \\ + &\hspace{5mm}\textbf{if} \: amsgrad \\ + &\hspace{10mm} v_t^{max} \leftarrow \mathrm{max}(v_{t-1}^{max},v_t) \\ + &\hspace{10mm}\widehat{v_t} \leftarrow v_t^{max}/\big(1-\beta_2^t \big) \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}\widehat{v_t} \leftarrow v_t/\big(1-\beta_2^t \big) \\ + &\hspace{5mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t}/ + \big(\sqrt{\widehat{v_t}} + \epsilon \big) \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to `Decoupled Weight Decay Regularization`_. + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 1e-3). A tensor LR + is not yet supported for all our implementations. Please use a float + LR if you are not also specifying fused=True or capturable=True. + betas (tuple[float | Tensor, float | Tensor], optional): + coefficients used for computing running averages of gradient and + its square. If a tensor is provided, must be 1-element. (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay coefficient (default: 1e-2) + amsgrad (bool, optional): whether to use the AMSGrad variant of this + algorithm from the paper `On the Convergence of Adam and Beyond`_ + (default: False) + {_maximize_doc} + {_foreach_doc} + {_capturable_doc} + {_differentiable_doc} + {_fused_doc} + .. Note:: + A prototype implementation of Adam and AdamW for MPS supports `torch.float32` and `torch.float16`. + .. _Decoupled Weight Decay Regularization: + https://arxiv.org/abs/1711.05101 + .. _On the Convergence of Adam and Beyond: + https://openreview.net/forum?id=ryQu7f-RZ + + """ +) + + +# @_disable_dynamo_if_unsupported logic occurs in the decorator that's applied to F.adam +def adamw( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + max_exp_avg_sqs: list[Tensor], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + foreach: bool | None = None, + capturable: bool = False, + differentiable: bool = False, + fused: bool | None = None, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, + has_complex: bool = False, + *, + amsgrad: bool, + beta1: float | Tensor, + beta2: float | Tensor, + lr: float | Tensor, + weight_decay: float, + eps: float, + maximize: bool, +) -> None: + r"""Functional API that performs AdamW algorithm computation. + + See :class:`~torch.optim.AdamW` for details. + """ + adam( + params, + grads, + exp_avgs, + exp_avg_sqs, + max_exp_avg_sqs, + state_steps, + foreach=foreach, + capturable=capturable, + differentiable=differentiable, + fused=fused, + grad_scale=grad_scale, + found_inf=found_inf, + has_complex=has_complex, + amsgrad=amsgrad, + beta1=beta1, + beta2=beta2, + lr=lr, + weight_decay=weight_decay, + eps=eps, + maximize=maximize, + decoupled_weight_decay=True, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/asgd.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/asgd.py new file mode 100644 index 0000000000000000000000000000000000000000..d36e6e0ec8c70def5340a033e46886f7ddb9be46 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/asgd.py @@ -0,0 +1,478 @@ +# mypy: allow-untyped-defs +from typing import cast + +import torch +from torch import Tensor + +from .optimizer import ( + _capturable_doc, + _default_to_fused_or_foreach, + _differentiable_doc, + _disable_dynamo_if_unsupported, + _foreach_doc, + _get_capturable_supported_devices, + _get_scalar_dtype, + _get_value, + _maximize_doc, + _params_doc, + _to_scalar, + _use_grad_for_differentiable, + _view_as_real, + Optimizer, + ParamsT, +) + + +__all__ = ["ASGD", "asgd"] + + +class ASGD(Optimizer): + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-2, + lambd: float = 1e-4, + alpha: float = 0.75, + t0: float = 1e6, + weight_decay: float = 0, + foreach: bool | None = None, + maximize: bool = False, + differentiable: bool = False, + capturable: bool = False, + ) -> None: + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + + defaults = { + "lr": lr, + "lambd": lambd, + "alpha": alpha, + "t0": t0, + "weight_decay": weight_decay, + "foreach": foreach, + "maximize": maximize, + "differentiable": differentiable, + "capturable": capturable, + } + super().__init__(params, defaults) + + def __setstate__(self, state): + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("foreach", None) + group.setdefault("maximize", False) + group.setdefault("differentiable", False) + group.setdefault("capturable", False) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0: + if not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = torch.tensor( + step_val, dtype=_get_scalar_dtype(), device=p.device + ) + if not torch.is_tensor(p_state["eta"]): + p_state["eta"] = torch.tensor( + p_state["eta"], dtype=_get_scalar_dtype(), device=p.device + ) + if not torch.is_tensor(p_state["mu"]): + p_state["mu"] = torch.tensor( + p_state["mu"], dtype=_get_scalar_dtype(), device=p.device + ) + + def _init_group(self, group, params_with_grad, grads, mus, axs, etas, state_steps): + has_complex = False + for p in group["params"]: + if p.grad is not None: + has_complex |= torch.is_complex(p) + params_with_grad.append(p) + if p.grad.is_sparse: + raise RuntimeError("ASGD does not support sparse gradients") + grads.append(p.grad) + + state = self.state[p] + # State initialization + if len(state) == 0: + state["step"] = torch.zeros( + (), device=p.device, dtype=_get_scalar_dtype() + ) + state["eta"] = ( + torch.as_tensor( + _to_scalar(group["lr"]), + device=p.device, + dtype=_get_scalar_dtype(), + ) + .clone() + .detach() + ) + state["mu"] = torch.ones( + (), device=p.device, dtype=_get_scalar_dtype() + ) + state["ax"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + mus.append(state["mu"]) + axs.append(state["ax"]) + etas.append(state["eta"]) + state_steps.append(state["step"]) + return has_complex + + @_use_grad_for_differentiable + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._accelerator_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + mus: list[Tensor] = [] + axs: list[Tensor] = [] + etas: list[Tensor] = [] + state_steps: list[Tensor] = [] + + has_complex = self._init_group( + group, params_with_grad, grads, mus, axs, etas, state_steps + ) + + asgd( + params_with_grad, + grads, + axs, + mus, + etas, + state_steps, + lambd=group["lambd"], + lr=group["lr"], + t0=group["t0"], + alpha=group["alpha"], + weight_decay=group["weight_decay"], + foreach=group["foreach"], + maximize=group["maximize"], + differentiable=group["differentiable"], + capturable=group["capturable"], + has_complex=has_complex, + ) + + return loss + + +ASGD.__doc__ = rf"""Implements Averaged Stochastic Gradient Descent. + + It has been proposed in `Acceleration of stochastic approximation by + averaging`_. + + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 1e-2) + lambd (float, optional): decay term (default: 1e-4) + alpha (float, optional): power for eta update (default: 0.75) + t0 (float, optional): point at which to start averaging (default: 1e6) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + {_foreach_doc} + {_maximize_doc} + {_differentiable_doc} + {_capturable_doc} + + .. _Acceleration of stochastic approximation by averaging: + https://meyn.ece.ufl.edu/wp-content/uploads/sites/77/archive/spm_files/Courses/ECE555-2011/555media/poljud92.pdf + + """ + + +def _single_tensor_asgd( + params: list[Tensor], + grads: list[Tensor], + axs: list[Tensor], + mus: list[Tensor], + etas: list[Tensor], + state_steps: list[Tensor], + *, + lambd: float, + lr: float, + t0: float, + alpha: float, + weight_decay: float, + maximize: bool, + differentiable: bool, + capturable: bool, + has_complex: bool, +) -> None: + if not torch.jit.is_scripting(): + lr = _to_scalar(lr) + + for i, param in enumerate(params): + grad = grads[i] + grad = grad if not maximize else -grad + mu = mus[i] + ax = axs[i] + eta = etas[i] + step_t = state_steps[i] + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices() + if not ( + param.device.type + == mu.device.type + == eta.device.type + == step_t.device.type + and param.device.type in capturable_supported_devices + ): + raise AssertionError( + f"If capturable=True, params, mus, etas, and state_steps must be " + f"on supported devices: {capturable_supported_devices}." + ) + + if torch.is_complex(param): + grad = torch.view_as_real(grad) + param = torch.view_as_real(param) + ax = torch.view_as_real(ax) + + # update step + step_t += 1 + + if weight_decay != 0: + grad = grad.add(param, alpha=weight_decay) + + if capturable: + param.mul_(1 - lambd * eta) + param.addcmul_(grad, eta, value=-1) # update parameter + else: + eta_value = _get_value(eta) + param.mul_(1 - lambd * eta_value) # decay term + param.add_(grad, alpha=-eta_value) # update parameter + + # averaging + if capturable or mu.item() != 1: + ax.add_(param.sub(ax).mul_(mu)) + else: + ax.copy_(param) + + if capturable: + eta.copy_(lr / ((1 + lambd * lr * step_t) ** alpha)) + mu.copy_(1 / torch.maximum(step_t - t0, torch.ones_like(step_t))) + else: + step = _get_value(step_t) + new_eta = torch.as_tensor(lr / ((1 + lambd * lr * step) ** alpha)) + eta.copy_(new_eta) + new_mu = torch.as_tensor(1 / max(1, step - t0)) + mu.copy_(new_mu) + + +def _multi_tensor_asgd( + params: list[Tensor], + grads: list[Tensor], + axs: list[Tensor], + mus: list[Tensor], + etas: list[Tensor], + state_steps: list[Tensor], + *, + lambd: float, + lr: float, + t0: float, + alpha: float, + weight_decay: float, + maximize: bool, + differentiable: bool, + capturable: bool, + has_complex: bool, +) -> None: + if len(params) == 0: + return + + if differentiable: + raise AssertionError("_foreach ops don't support autograd") + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices( + supports_xla=False + ) + if not all( + p.device.type == mu.device.type == eta.device.type == step.device.type + and p.device.type in capturable_supported_devices + for p, mu, eta, step in zip(params, mus, etas, state_steps, strict=True) + ): + raise AssertionError( + f"If capturable=True, params, mus, etas, and state_steps must be on " + f"supported devices: {capturable_supported_devices}." + ) + + lr = _to_scalar(lr) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, axs, mus, etas, state_steps] # type: ignore[list-item] + ) + for (device, _), ( + ( + grouped_params_, + grouped_grads_, + grouped_axs_, + grouped_mus_, + grouped_etas_, + grouped_state_steps_, + ), + _, + ) in grouped_tensors.items(): + grouped_params = cast(list[Tensor], grouped_params_) + grouped_grads = cast(list[Tensor], grouped_grads_) + grouped_axs = cast(list[Tensor], grouped_axs_) + grouped_mus = cast(list[Tensor], grouped_mus_) + grouped_etas = cast(list[Tensor], grouped_etas_) + grouped_state_steps = cast(list[Tensor], grouped_state_steps_) + + if has_complex: + _view_as_real(grouped_params, grouped_grads, grouped_axs) + + if maximize: + grouped_grads = torch._foreach_neg(grouped_grads) # type: ignore[assignment] + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and grouped_state_steps[0].is_cpu: + torch._foreach_add_( + grouped_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(grouped_state_steps, 1) + + # intermediate = grad + param * lambd + intermediate: tuple[Tensor, ...] | list[Tensor] + if weight_decay != 0: + if maximize: + torch._foreach_add_(grouped_grads, grouped_params, alpha=weight_decay) + intermediate = grouped_grads + else: + intermediate = torch._foreach_add( + grouped_grads, grouped_params, alpha=weight_decay + ) + + torch._foreach_add_(intermediate, grouped_params, alpha=lambd) + else: + intermediate = torch._foreach_add( + grouped_grads, grouped_params, alpha=lambd + ) + + # update param + # param * (1 - lambd * eta) - eta * grad + # => param - param * lambd * eta - eta * grad + # => param - eta * intermediate + torch._foreach_addcmul_(grouped_params, intermediate, grouped_etas, value=-1) + del intermediate + + # update grouped_axs + # averaging: ax = ax + mu * (param - ax) + # Note (mlazos): We can't use lerp here since it requires weight to be float64 + # and our grouping code requires dtypes to match for all tensors in a group (and it should, since + # we use the mus in other places) + # all dtypes need to match, so we could introduce a cast in a loop + # but since this only adds one additional kernel launch, this looks like the cleaner + # and faster solution + intermediate = torch._foreach_sub(grouped_params, grouped_axs) + torch._foreach_addcmul_(grouped_axs, intermediate, grouped_mus) + del intermediate + + new_etas: tuple[Tensor, ...] | list[Tensor] + new_mus: tuple[Tensor, ...] | list[Tensor] + if capturable: + # update grouped_mus + new_mus = torch._foreach_sub(grouped_state_steps, t0) + torch._foreach_maximum_(new_mus, 1.0) + torch._foreach_reciprocal_(new_mus) + torch._foreach_copy_(grouped_mus, new_mus) + del new_mus + + # update eta = lr / ((1 + lambd * lr * step)^alpha) + new_etas = torch._foreach_mul(grouped_state_steps, lambd) + torch._foreach_mul_(new_etas, lr) + torch._foreach_add_(new_etas, 1) + torch._foreach_pow_(new_etas, alpha) + torch._foreach_reciprocal_(new_etas) + torch._foreach_mul_(new_etas, lr) + torch._foreach_copy_(grouped_etas, new_etas) + else: + new_etas = [ + torch.as_tensor(lr / ((1 + lambd * lr * step) ** alpha), device=device) + for step in grouped_state_steps + ] + new_mus = [ + torch.as_tensor(1 / max(1, _get_value(step) - t0), device=device) + for step in grouped_state_steps + ] + torch._foreach_copy_(grouped_etas, new_etas) + torch._foreach_copy_(grouped_mus, new_mus) + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_asgd) +def asgd( + params: list[Tensor], + grads: list[Tensor], + axs: list[Tensor], + mus: list[Tensor], + etas: list[Tensor], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + foreach: bool | None = None, + maximize: bool = False, + differentiable: bool = False, + capturable: bool = False, + has_complex: bool = False, + *, + lambd: float, + lr: float, + t0: float, + alpha: float, + weight_decay: float, +) -> None: + r"""Functional API that performs asgd algorithm computation. + + See :class:`~torch.optim.ASGD` for details. + """ + if foreach is None: + _, foreach = _default_to_fused_or_foreach( + params, differentiable, use_fused=False + ) + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + + if foreach and not torch.jit.is_scripting(): + func = _multi_tensor_asgd + else: + func = _single_tensor_asgd + + func( + params, + grads, + axs, + mus, + etas, + state_steps, + lambd=lambd, + lr=lr, + t0=t0, + alpha=alpha, + weight_decay=weight_decay, + maximize=maximize, + differentiable=differentiable, + capturable=capturable, + has_complex=has_complex, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/lbfgs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/lbfgs.py new file mode 100644 index 0000000000000000000000000000000000000000..af6b6e8b02bb0b4350afc6208a4c7042798c77bc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/lbfgs.py @@ -0,0 +1,530 @@ +# mypy: allow-untyped-defs + +import torch +from torch import Tensor + +from .optimizer import _to_scalar, Optimizer, ParamsT + + +__all__ = ["LBFGS"] + + +def _cubic_interpolate(x1, f1, g1, x2, f2, g2, bounds=None): + # ported from https://github.com/torch/optim/blob/master/polyinterp.lua + # Compute bounds of interpolation area + if bounds is not None: + xmin_bound, xmax_bound = bounds + else: + xmin_bound, xmax_bound = (x1, x2) if x1 <= x2 else (x2, x1) + + # Code for most common case: cubic interpolation of 2 points + # w/ function and derivative values for both + # Solution in this case (where x2 is the farthest point): + # d1 = g1 + g2 - 3*(f1-f2)/(x1-x2); + # d2 = sqrt(d1^2 - g1*g2); + # min_pos = x2 - (x2 - x1)*((g2 + d2 - d1)/(g2 - g1 + 2*d2)); + # t_new = min(max(min_pos,xmin_bound),xmax_bound); + d1 = g1 + g2 - 3 * (f1 - f2) / (x1 - x2) + d2_square = d1**2 - g1 * g2 + if d2_square >= 0: + d2 = d2_square.sqrt() + if x1 <= x2: + min_pos = x2 - (x2 - x1) * ((g2 + d2 - d1) / (g2 - g1 + 2 * d2)) + else: + min_pos = x1 - (x1 - x2) * ((g1 + d2 - d1) / (g1 - g2 + 2 * d2)) + return min(max(min_pos, xmin_bound), xmax_bound) + else: + return (xmin_bound + xmax_bound) / 2.0 + + +def _strong_wolfe( + obj_func, x, t, d, f, g, gtd, c1=1e-4, c2=0.9, tolerance_change=1e-9, max_ls=25 +): + # ported from https://github.com/torch/optim/blob/master/lswolfe.lua + d_norm = d.abs().max() + g = g.clone(memory_format=torch.contiguous_format) + # evaluate objective and gradient using initial step + f_new, g_new = obj_func(x, t, d) + ls_func_evals = 1 + gtd_new = g_new.dot(d) + + # bracket an interval containing a point satisfying the Wolfe criteria + t_prev, f_prev, g_prev, gtd_prev = 0, f, g, gtd + done = False + ls_iter = 0 + while ls_iter < max_ls: + # check conditions + if f_new > (f + c1 * t * gtd) or (ls_iter > 1 and f_new >= f_prev): + bracket = [t_prev, t] + bracket_f = [f_prev, f_new] + bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)] + bracket_gtd = [gtd_prev, gtd_new] + break + + if abs(gtd_new) <= -c2 * gtd: + bracket = [t] + bracket_f = [f_new] + bracket_g = [g_new] + done = True + break + + if gtd_new >= 0: + bracket = [t_prev, t] + bracket_f = [f_prev, f_new] + bracket_g = [g_prev, g_new.clone(memory_format=torch.contiguous_format)] + bracket_gtd = [gtd_prev, gtd_new] + break + + # interpolate + min_step = t + 0.01 * (t - t_prev) + max_step = t * 10 + tmp = t + t = _cubic_interpolate( + t_prev, f_prev, gtd_prev, t, f_new, gtd_new, bounds=(min_step, max_step) + ) + + # next step + t_prev = tmp + f_prev = f_new + g_prev = g_new.clone(memory_format=torch.contiguous_format) + gtd_prev = gtd_new + f_new, g_new = obj_func(x, t, d) + ls_func_evals += 1 + gtd_new = g_new.dot(d) + ls_iter += 1 + + # reached max number of iterations? + if ls_iter == max_ls: + bracket = [0, t] + bracket_f = [f, f_new] + bracket_g = [g, g_new] + + # zoom phase: we now have a point satisfying the criteria, or + # a bracket around it. We refine the bracket until we find the + # exact point satisfying the criteria + insuf_progress = False + # find high and low points in bracket + low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[-1] else (1, 0) # type: ignore[possibly-undefined] + while not done and ls_iter < max_ls: + # line-search bracket is so small + if abs(bracket[1] - bracket[0]) * d_norm < tolerance_change: # type: ignore[possibly-undefined] + break + + # compute new trial value + t = _cubic_interpolate( + # pyrefly: ignore [unbound-name] + bracket[0], + # pyrefly: ignore [unbound-name] + bracket_f[0], + bracket_gtd[0], # type: ignore[possibly-undefined] + # pyrefly: ignore [unbound-name] + bracket[1], + # pyrefly: ignore [unbound-name] + bracket_f[1], + # pyrefly: ignore [unbound-name] + bracket_gtd[1], + ) + + # test that we are making sufficient progress: + # in case `t` is so close to boundary, we mark that we are making + # insufficient progress, and if + # + we have made insufficient progress in the last step, or + # + `t` is at one of the boundary, + # we will move `t` to a position which is `0.1 * len(bracket)` + # away from the nearest boundary point. + # pyrefly: ignore [unbound-name] + eps = 0.1 * (max(bracket) - min(bracket)) + # pyrefly: ignore [unbound-name] + if min(max(bracket) - t, t - min(bracket)) < eps: + # interpolation close to boundary + # pyrefly: ignore [unbound-name] + if insuf_progress or t >= max(bracket) or t <= min(bracket): + # evaluate at 0.1 away from boundary + # pyrefly: ignore [unbound-name] + if abs(t - max(bracket)) < abs(t - min(bracket)): + # pyrefly: ignore [unbound-name] + t = max(bracket) - eps + else: + # pyrefly: ignore [unbound-name] + t = min(bracket) + eps + insuf_progress = False + else: + insuf_progress = True + else: + insuf_progress = False + + # Evaluate new point + f_new, g_new = obj_func(x, t, d) + ls_func_evals += 1 + gtd_new = g_new.dot(d) + ls_iter += 1 + + # pyrefly: ignore [unbound-name] + if f_new > (f + c1 * t * gtd) or f_new >= bracket_f[low_pos]: + # Armijo condition not satisfied or not lower than lowest point + # pyrefly: ignore [unbound-name] + bracket[high_pos] = t + # pyrefly: ignore [unbound-name] + bracket_f[high_pos] = f_new + bracket_g[high_pos] = g_new.clone(memory_format=torch.contiguous_format) # type: ignore[possibly-undefined] + # pyrefly: ignore [unbound-name] + bracket_gtd[high_pos] = gtd_new + # pyrefly: ignore [unbound-name] + low_pos, high_pos = (0, 1) if bracket_f[0] <= bracket_f[1] else (1, 0) + else: + if abs(gtd_new) <= -c2 * gtd: + # Wolfe conditions satisfied + done = True + # pyrefly: ignore [unbound-name] + elif gtd_new * (bracket[high_pos] - bracket[low_pos]) >= 0: + # old high becomes new low + # pyrefly: ignore [unbound-name] + bracket[high_pos] = bracket[low_pos] + # pyrefly: ignore [unbound-name] + bracket_f[high_pos] = bracket_f[low_pos] + bracket_g[high_pos] = bracket_g[low_pos] # type: ignore[possibly-undefined] + # pyrefly: ignore [unbound-name] + bracket_gtd[high_pos] = bracket_gtd[low_pos] + + # new point becomes new low + # pyrefly: ignore [unbound-name] + bracket[low_pos] = t + # pyrefly: ignore [unbound-name] + bracket_f[low_pos] = f_new + bracket_g[low_pos] = g_new.clone(memory_format=torch.contiguous_format) # type: ignore[possibly-undefined] + # pyrefly: ignore [unbound-name] + bracket_gtd[low_pos] = gtd_new + + # return stuff + t = bracket[low_pos] # type: ignore[possibly-undefined] + # pyrefly: ignore [unbound-name] + f_new = bracket_f[low_pos] + g_new = bracket_g[low_pos] # type: ignore[possibly-undefined] + return f_new, g_new, t, ls_func_evals + + +class LBFGS(Optimizer): + """Implements L-BFGS algorithm. + + Heavily inspired by `minFunc + `_. + + .. warning:: + This optimizer doesn't support per-parameter options and parameter + groups (there can be only one). + + .. warning:: + Right now all parameters have to be on a single device. This will be + improved in the future. + + .. note:: + This is a very memory intensive optimizer (it requires additional + ``param_bytes * (history_size + 1)`` bytes). If it doesn't fit in memory + try reducing the history size, or use a different algorithm. + + Args: + params (iterable): iterable of parameters to optimize. Parameters must be real. + lr (float, optional): learning rate (default: 1) + max_iter (int, optional): maximal number of iterations per optimization step + (default: 20) + max_eval (int, optional): maximal number of function evaluations per optimization + step (default: max_iter * 1.25). + tolerance_grad (float, optional): termination tolerance on first order optimality + (default: 1e-7). + tolerance_change (float, optional): termination tolerance on function + value/parameter changes (default: 1e-9). + history_size (int, optional): update history size (default: 100). + line_search_fn (str, optional): either 'strong_wolfe' or None (default: None). + """ + + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1, + max_iter: int = 20, + max_eval: int | None = None, + tolerance_grad: float = 1e-7, + tolerance_change: float = 1e-9, + history_size: int = 100, + line_search_fn: str | None = None, + ) -> None: + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if max_eval is None: + max_eval = max_iter * 5 // 4 + defaults = { + "lr": lr, + "max_iter": max_iter, + "max_eval": max_eval, + "tolerance_grad": tolerance_grad, + "tolerance_change": tolerance_change, + "history_size": history_size, + "line_search_fn": line_search_fn, + } + super().__init__(params, defaults) + + if len(self.param_groups) != 1: + raise ValueError( + "LBFGS doesn't support per-parameter options (parameter groups)" + ) + + self._params = self.param_groups[0]["params"] + self._numel_cache = None + + def _numel(self): + if self._numel_cache is None: + self._numel_cache = sum( + 2 * p.numel() if torch.is_complex(p) else p.numel() + for p in self._params + ) + + return self._numel_cache + + def _gather_flat_grad(self): + views = [] + for p in self._params: + if p.grad is None: + view = p.new(p.numel()).zero_() + elif p.grad.is_sparse: + view = p.grad.to_dense().view(-1) + else: + view = p.grad.view(-1) + if torch.is_complex(view): + view = torch.view_as_real(view).view(-1) + views.append(view) + return torch.cat(views, 0) + + def _add_grad(self, step_size, update) -> None: + offset = 0 + for p in self._params: + if torch.is_complex(p): + p = torch.view_as_real(p) + numel = p.numel() + # view as to avoid deprecated pointwise semantics + p.add_(update[offset : offset + numel].view_as(p), alpha=step_size) + offset += numel + if offset != self._numel(): + raise AssertionError(f"Expected offset {offset} to equal {self._numel()}") + + def _clone_param(self): + return [p.clone(memory_format=torch.contiguous_format) for p in self._params] + + def _set_param(self, params_data) -> None: + for p, pdata in zip(self._params, params_data, strict=True): + p.copy_(pdata) + + def _directional_evaluate(self, closure, x, t, d): + self._add_grad(t, d) + loss = float(closure()) + flat_grad = self._gather_flat_grad() + self._set_param(x) + return loss, flat_grad + + @torch.no_grad() + def step(self, closure): # type: ignore[override] + """Perform a single optimization step. + + Args: + closure (Callable): A closure that reevaluates the model + and returns the loss. + """ + if len(self.param_groups) != 1: + raise AssertionError( + f"Expected exactly one param_group, but got {len(self.param_groups)}" + ) + + # Make sure the closure is always called with grad enabled + closure = torch.enable_grad()(closure) + + group = self.param_groups[0] + lr = _to_scalar(group["lr"]) + max_iter = group["max_iter"] + max_eval = group["max_eval"] + tolerance_grad = group["tolerance_grad"] + tolerance_change = group["tolerance_change"] + line_search_fn = group["line_search_fn"] + history_size = group["history_size"] + + # NOTE: LBFGS has only global state, but we register it as state for + # the first param, because this helps with casting in load_state_dict + state = self.state[self._params[0]] + state.setdefault("func_evals", 0) + state.setdefault("n_iter", 0) + + # evaluate initial f(x) and df/dx + orig_loss = closure() + loss = float(orig_loss) + current_evals = 1 + state["func_evals"] += 1 + + flat_grad = self._gather_flat_grad() + opt_cond = flat_grad.abs().max() <= tolerance_grad + + # optimal condition + if opt_cond: + return orig_loss + + # tensors cached in state (for tracing) + d = state.get("d") + t = state.get("t") + old_dirs = state.get("old_dirs") + old_stps = state.get("old_stps") + ro = state.get("ro") + H_diag = state.get("H_diag") + prev_flat_grad = state.get("prev_flat_grad") + prev_loss = state.get("prev_loss") + + n_iter = 0 + # optimize for a max of max_iter iterations + while n_iter < max_iter: + # keep track of nb of iterations + n_iter += 1 + state["n_iter"] += 1 + + ############################################################ + # compute gradient descent direction + ############################################################ + if state["n_iter"] == 1: + d = flat_grad.neg() + old_dirs = [] + old_stps = [] + ro = [] + H_diag = 1 + else: + # do lbfgs update (update memory) + y = flat_grad.sub(prev_flat_grad) + s = d.mul(t) + ys = y.dot(s) # y*s + if ys > 1e-10: + # updating memory + if len(old_dirs) == history_size: + # shift history by one (limited-memory) + old_dirs.pop(0) + old_stps.pop(0) + ro.pop(0) + + # store new direction/step + old_dirs.append(y) + old_stps.append(s) + ro.append(1.0 / ys) + + # update scale of initial Hessian approximation + H_diag = ys / y.dot(y) # (y*y) + + # compute the approximate (L-BFGS) inverse Hessian + # multiplied by the gradient + num_old = len(old_dirs) + + if "al" not in state: + state["al"] = [None] * history_size + al = state["al"] + + # iteration in L-BFGS loop collapsed to use just one buffer + q = flat_grad.neg() + for i in range(num_old - 1, -1, -1): + al[i] = old_stps[i].dot(q) * ro[i] + q.add_(old_dirs[i], alpha=-al[i]) + + # multiply by initial Hessian + # r/d is the final direction + d = r = torch.mul(q, H_diag) + for i in range(num_old): + be_i = old_dirs[i].dot(r) * ro[i] + r.add_(old_stps[i], alpha=al[i] - be_i) + + if prev_flat_grad is None: + prev_flat_grad = flat_grad.clone(memory_format=torch.contiguous_format) + else: + prev_flat_grad.copy_(flat_grad) + prev_loss = loss + + ############################################################ + # compute step length + ############################################################ + # reset initial guess for step size + if state["n_iter"] == 1: + t = min(1.0, 1.0 / flat_grad.abs().sum()) * lr + else: + t = lr + + # directional derivative + gtd = flat_grad.dot(d) # g * d + + # directional derivative is below tolerance + if gtd > -tolerance_change: + break + + # optional line search: user function + ls_func_evals = 0 + if line_search_fn is not None: + # perform line search, using user function + if line_search_fn != "strong_wolfe": + raise RuntimeError("only 'strong_wolfe' is supported") + else: + x_init = self._clone_param() + + def obj_func(x, t, d): + return self._directional_evaluate(closure, x, t, d) + + loss, flat_grad, t, ls_func_evals = _strong_wolfe( + obj_func, + x_init, + t, + d, + loss, + flat_grad, + gtd, + max_ls=max_eval - current_evals, + ) + self._add_grad(t, d) + opt_cond = flat_grad.abs().max() <= tolerance_grad + else: + # no line search, simply move with fixed-step + self._add_grad(t, d) + if n_iter != max_iter: + # re-evaluate function only if not in last iteration + # the reason we do this: in a stochastic setting, + # no use to re-evaluate that function here + with torch.enable_grad(): + loss = closure() + loss = float(loss) + flat_grad = self._gather_flat_grad() + opt_cond = flat_grad.abs().max() <= tolerance_grad + ls_func_evals = 1 + + # update func eval + current_evals += ls_func_evals + state["func_evals"] += ls_func_evals + + ############################################################ + # check conditions + ############################################################ + if n_iter == max_iter: + break + + if current_evals >= max_eval: + break + + # optimal condition + if opt_cond: + break + + # lack of progress + if d.mul(t).abs().max() <= tolerance_change: + break + + if abs(loss - prev_loss) < tolerance_change: + break + + state["d"] = d + state["t"] = t + state["old_dirs"] = old_dirs + state["old_stps"] = old_stps + state["ro"] = ro + state["H_diag"] = H_diag + state["prev_flat_grad"] = prev_flat_grad + state["prev_loss"] = prev_loss + + return orig_loss diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/lr_scheduler.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/lr_scheduler.py new file mode 100644 index 0000000000000000000000000000000000000000..b2049c73fc472beff2a7ddbab5186832da8ec6e8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/lr_scheduler.py @@ -0,0 +1,2598 @@ +# mypy: allow-untyped-defs +r"""Learning Rate Scheduler.""" + +from __future__ import annotations + +import math +import types +import warnings +from bisect import bisect_right +from collections import Counter +from functools import partial, wraps +from typing import Any, cast, Literal, SupportsFloat, TYPE_CHECKING, TypedDict +from typing_extensions import override, Self +from weakref import ref + +from torch import inf, Tensor + +from .optimizer import _to_scalar, Optimizer + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable, Sequence + + +__all__ = [ + "LambdaLR", + "MultiplicativeLR", + "StepLR", + "MultiStepLR", + "ConstantLR", + "LinearLR", + "ExponentialLR", + "SequentialLR", + "CosineAnnealingLR", + "ChainedScheduler", + "ReduceLROnPlateau", + "CyclicLR", + "CosineAnnealingWarmRestarts", + "OneCycleLR", + "PolynomialLR", + "LRScheduler", +] + +EPOCH_DEPRECATION_WARNING = ( + "The epoch parameter in `scheduler.step()` was not necessary and is being " + "deprecated where possible. Please use `scheduler.step()` to step the " + "scheduler. During the deprecation, if epoch is different from None, the " + "closed form is used instead of the new chainable form, where available. " + "Please open an issue if you are unable to replicate your use case: " + "https://github.com/pytorch/pytorch/issues/new/choose." +) + + +def _format_param(name: str, optimizer: Optimizer, param): + """Return correctly formatted lr/momentum for each param group.""" + + def _copy(_param): + return _param.clone() if isinstance(_param, Tensor) else _param + + if isinstance(param, (list, tuple)): + if len(param) != len(optimizer.param_groups): + raise ValueError( + f"{name} must have the same length as optimizer.param_groups. " + f"{name} has {len(param)} values, param_groups has {len(optimizer.param_groups)}." + ) + else: + param = [param] * len(optimizer.param_groups) + + return list(map(_copy, param)) + + +def _param_groups_val_list(optimizer: Optimizer, key: str) -> list[Any]: + """Create a list containing group[key] for each optimizer param_group. + Prevents aliasing when group[key] could be a Tensor. + Raises a KeyError when group[key] does not exist. + """ + return [ + group[key].clone() if isinstance(group[key], Tensor) else group[key] + for group in optimizer.param_groups + ] + + +def _update_param_group_val( + param_group: dict[str, Any], key: str, val: float | Tensor +) -> None: + """Set param_group[key] to val without aliasing or assignment when they're + both tensors. Raises a KeyError if param_group[key] does not exist. + """ + if isinstance(param_group[key], Tensor): + param_group[key].fill_(_to_scalar(val)) + else: + param_group[key] = val + + +class LRScheduler: + r"""Base class for all learning rate schedulers. + + Subclasses implement :meth:`get_lr` and optionally override :meth:`step` to + define scheduling behavior. + + Args: + optimizer (Optimizer): The optimizer this scheduler will adjust the + learning rates of. + last_epoch (int): Index of the last epoch seen by the scheduler. Use + ``-1`` (default) to initialize the scheduler. Only use a non-default + value when restoring this scheduler from a saved checkpoint. + + .. warning:: + Initializing a scheduler overwrites its optimizer's + ``param_group["lr"]``\s. When restoring a checkpoint, initialize the + scheduler **before** calling your optimizer's + :meth:`~torch.optim.Optimizer.load_state_dict` to avoid overwriting the + loaded learning rates. + """ + + _get_lr_called_within_step: bool = False + _is_initial: bool = False + + def __init__( + self, + optimizer: Optimizer, + last_epoch: int = -1, + ) -> None: # noqa: D107 + # Attach optimizer + if not isinstance(optimizer, Optimizer): + raise TypeError(f"{type(optimizer).__name__} is not an Optimizer") + self.optimizer = optimizer + + # Initialize epoch and base learning rates + if last_epoch == -1: + for group in optimizer.param_groups: + initial_lr = group["lr"] + if isinstance(initial_lr, Tensor): + initial_lr = initial_lr.clone() + group.setdefault("initial_lr", initial_lr) + else: + for i, group in enumerate(optimizer.param_groups): + if "initial_lr" not in group: + raise KeyError( + f"param 'initial_lr' is not specified in param_groups[{i}] when resuming scheduler with last_epoch >= 0.\n" + "This typically happens when:\n" + "1. You're trying to resume training from a checkpoint but haven't properly loaded the optimizer state\n" + "2. You're using last_epoch >= 0 for a fresh training run (not recommended)" + ) + self.base_lrs: list[float | Tensor] = _param_groups_val_list( + optimizer, "initial_lr" + ) + self.last_epoch = last_epoch + + # Following https://github.com/pytorch/pytorch/issues/20124 + # We would like to ensure that `lr_scheduler.step()` is called after + # `optimizer.step()` + def patch_track_step_called(opt: Optimizer): + if hasattr(opt.step, "_wrapped_by_lr_sched"): + # we've already patched + return opt.step + + def wrap_step(step_fn): + opt_ref = ref(self.optimizer) + func = step_fn.__func__ + + @wraps(func) + def wrapper(*args, **kwargs): + opt = opt_ref() + opt._opt_called = True # type: ignore[union-attr] + return func.__get__(opt, opt.__class__)(*args, **kwargs) + + wrapper._wrapped_by_lr_sched = True # type: ignore[attr-defined] + return wrapper + + opt.step = wrap_step(opt.step) # type: ignore[method-assign] + + patch_track_step_called(self.optimizer) + self._initial_step() + + def _initial_step(self) -> None: + """Initialize step counts and perform a step.""" + self._step_count = 0 + with _initial_mode(self): + self.step() + + def state_dict(self) -> dict[str, Any]: + """Return the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in ``self.__dict__`` which + is not the optimizer. + """ + return { + key: value for key, value in self.__dict__.items() if key != "optimizer" + } + + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load the scheduler's state. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + self.__dict__.update(state_dict) + + def get_last_lr(self) -> list[float | Tensor]: + r"""Get the most recent learning rates computed by this scheduler. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates with entries + for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`, with the same types as + their ``group["lr"]``\s. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + # We always update self._last_lr with _param_groups_val_list, so it's a + # .clone() of the group["lr"]s. If we didn't do this, the user could + # corrupt their learning rates by modifying the outputs in place. + return self._last_lr + + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + raise NotImplementedError + + def step(self, epoch: int | None = None) -> None: + """Step the scheduler. + + Args: + epoch (int, optional): + .. deprecated:: 1.4 + If provided, sets :attr:`last_epoch` to ``epoch`` and uses + :meth:`_get_closed_form_lr` if it is available. This is not + universally supported. Use :meth:`step` without arguments + instead. + + .. note:: + Call this method after calling the optimizer's + :meth:`~torch.optim.Optimizer.step`. + """ + # Raise a warning if old pattern is detected + # https://github.com/pytorch/pytorch/issues/20124 + if self._step_count == 1: + if not hasattr(self.optimizer.step, "_wrapped_by_lr_sched"): + warnings.warn( + "Seems like `optimizer.step()` has been overridden after learning rate scheduler " + "initialization. Please, make sure to call `optimizer.step()` before " + "`lr_scheduler.step()`. See more details at " + "https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate", + UserWarning, + stacklevel=2, + ) + + # Just check if there were two first lr_scheduler.step() calls before optimizer.step() + elif not getattr(self.optimizer, "_opt_called", False): + warnings.warn( + "Detected call of `lr_scheduler.step()` before `optimizer.step()`. " + "In PyTorch 1.1.0 and later, you should call them in the opposite order: " + "`optimizer.step()` before `lr_scheduler.step()`. Failure to do this " + "will result in PyTorch skipping the first value of the learning rate schedule. " + "See more details at " + "https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate", + UserWarning, + stacklevel=2, + ) + + self._step_count += 1 + if epoch is not None: + warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning, stacklevel=2) + self._update_lr(epoch) + + def _update_lr(self, epoch: int | None = None) -> None: + with _enable_get_lr_call(self): + if epoch is None: + self.last_epoch += 1 + values = self.get_lr() + else: + self.last_epoch = epoch + if hasattr(self, "_get_closed_form_lr"): + values = cast(list[float | Tensor], self._get_closed_form_lr()) + else: + values = self.get_lr() + + for param_group, lr in zip(self.optimizer.param_groups, values, strict=True): + _update_param_group_val(param_group, "lr", lr) + + self._last_lr: list[float | Tensor] = _param_groups_val_list( + self.optimizer, "lr" + ) + + +def _warn_get_lr_called_within_step(lr_scheduler: LRScheduler) -> None: + if not lr_scheduler._get_lr_called_within_step: + warnings.warn( + "To get the last learning rate computed by the scheduler, " + "please use `get_last_lr()`.", + UserWarning, + stacklevel=2, + ) + + +# Including _LRScheduler for backwards compatibility +# Subclass instead of assign because we want __name__ of _LRScheduler to be _LRScheduler (assigning would make it LRScheduler). +class _LRScheduler(LRScheduler): + pass + + +class _enable_get_lr_call: + def __init__(self, o: LRScheduler) -> None: + self.o = o + + def __enter__(self) -> Self: + self.o._get_lr_called_within_step = True + return self + + def __exit__(self, type, value, traceback) -> None: + self.o._get_lr_called_within_step = False + + +class _initial_mode: + def __init__(self, o: LRScheduler) -> None: + self.o = o + + def __enter__(self): + self.o._is_initial = True + + def __exit__(self, type, value, traceback): + self.o._is_initial = False + + +class LambdaLR(LRScheduler): + """Sets the initial learning rate. + + The learning rate of each parameter group is set to the initial lr + times a given function. When last_epoch=-1, sets initial lr as lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + lr_lambda (function or list): A function which computes a multiplicative + factor given an integer parameter epoch, or a list of such + functions, one for each group in optimizer.param_groups. + last_epoch (int): The index of last epoch. Default: -1. + + Example: + >>> # xdoctest: +SKIP + >>> # Assuming optimizer has two groups. + >>> num_epochs = 100 + >>> lambda1 = lambda epoch: epoch // 30 + >>> lambda2 = lambda epoch: 0.95**epoch + >>> scheduler = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2]) + >>> for epoch in range(num_epochs): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + >>> + >>> # Alternatively, you can use a single lambda function for all groups. + >>> scheduler = LambdaLR(opt, lr_lambda=lambda epoch: epoch // 30) + >>> for epoch in range(num_epochs): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/LambdaLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + lr_lambda: Callable[[int], float] | list[Callable[[int], float]], + last_epoch: int = -1, + ) -> None: # noqa: D107 + self.optimizer = optimizer + + self.lr_lambdas: list[Callable[[int], float]] + if not isinstance(lr_lambda, list) and not isinstance(lr_lambda, tuple): + self.lr_lambdas = [lr_lambda] * len(optimizer.param_groups) + else: + if len(lr_lambda) != len(optimizer.param_groups): + raise ValueError( + f"Expected {len(optimizer.param_groups)} lr_lambdas, but got {len(lr_lambda)}" + ) + self.lr_lambdas = list(lr_lambda) + super().__init__(optimizer, last_epoch) + + @override + def state_dict(self) -> dict[str, Any]: + """Return the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in ``self.__dict__`` which is not the optimizer. + The learning rate lambda functions will only be saved if they are callable objects + and not if they are functions or lambdas. + + When saving or loading the scheduler, please make sure to also save or load the state of the optimizer. + """ + state_dict = { + key: value + for key, value in self.__dict__.items() + if key not in ("optimizer", "lr_lambdas") + } + state_dict["lr_lambdas"] = [None] * len(self.lr_lambdas) + + for idx, fn in enumerate(self.lr_lambdas): + if not isinstance(fn, types.FunctionType): + state_dict["lr_lambdas"][idx] = fn.__dict__.copy() + + return state_dict + + @override + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load the scheduler's state. + + When saving or loading the scheduler, please make sure to also save or load the state of the optimizer. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + lr_lambdas = state_dict.pop("lr_lambdas") + self.__dict__.update(state_dict) + # Restore state_dict keys in order to prevent side effects + # https://github.com/pytorch/pytorch/issues/32756 + state_dict["lr_lambdas"] = lr_lambdas + + for idx, fn in enumerate(lr_lambdas): + if fn is not None: + self.lr_lambdas[idx].__dict__.update(fn) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Scales the :attr:`base_lrs` by the outputs of the :attr:`lr_lambdas` at + :attr:`last_epoch`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + _warn_get_lr_called_within_step(self) + + return [ + base_lr * lmbda(self.last_epoch) + for lmbda, base_lr in zip(self.lr_lambdas, self.base_lrs, strict=True) + ] + + +class MultiplicativeLR(LRScheduler): + """Multiply the learning rate of each parameter group by the factor given in the specified function. + + When last_epoch=-1, set initial lr as lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + lr_lambda (function or list): A function which computes a multiplicative + factor given an integer parameter epoch, or a list of such + functions, one for each group in optimizer.param_groups. + last_epoch (int): The index of last epoch. Default: -1. + + Example: + >>> # xdoctest: +SKIP + >>> lmbda = lambda epoch: 0.95 + >>> scheduler = MultiplicativeLR(optimizer, lr_lambda=lmbda) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/MultiplicativeLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + lr_lambda: Callable[[int], float] | list[Callable[[int], float]], + last_epoch: int = -1, + ) -> None: # noqa: D107 + self.optimizer = optimizer + + self.lr_lambdas: list[Callable[[int], float]] + if not isinstance(lr_lambda, list) and not isinstance(lr_lambda, tuple): + self.lr_lambdas = [lr_lambda] * len(optimizer.param_groups) + else: + if len(lr_lambda) != len(optimizer.param_groups): + raise ValueError( + f"Expected {len(optimizer.param_groups)} lr_lambdas, but got {len(lr_lambda)}" + ) + self.lr_lambdas = list(lr_lambda) + for lr_lambda in self.lr_lambdas: + if not callable(lr_lambda): + raise TypeError( + f"lr_lambda should be a function, but got {type(lr_lambda).__name__}" + ) + super().__init__(optimizer, last_epoch) + + @override + def state_dict(self) -> dict[str, Any]: + """Return the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in ``self.__dict__`` which + is not the optimizer. + The learning rate lambda functions will only be saved if they are callable objects + and not if they are functions or lambdas. + """ + state_dict = { + key: value + for key, value in self.__dict__.items() + if key not in ("optimizer", "lr_lambdas") + } + state_dict["lr_lambdas"] = [None] * len(self.lr_lambdas) + + for idx, fn in enumerate(self.lr_lambdas): + if not isinstance(fn, types.FunctionType): + state_dict["lr_lambdas"][idx] = fn.__dict__.copy() + + return state_dict + + @override + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load the scheduler's state. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + lr_lambdas = state_dict.pop("lr_lambdas") + self.__dict__.update(state_dict) + # Restore state_dict keys in order to prevent side effects + # https://github.com/pytorch/pytorch/issues/32756 + state_dict["lr_lambdas"] = lr_lambdas + + for idx, fn in enumerate(lr_lambdas): + if fn is not None: + self.lr_lambdas[idx].__dict__.update(fn) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Scales the current ``group["lr"]``\s in each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` by the outputs of the + :attr:`lr_lambdas` at :attr:`last_epoch`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + _warn_get_lr_called_within_step(self) + + if not self._is_initial: + return [ + group["lr"] * lmbda(self.last_epoch) + for lmbda, group in zip( + self.lr_lambdas, self.optimizer.param_groups, strict=True + ) + ] + else: + return _param_groups_val_list(self.optimizer, "lr") + + +class StepLR(LRScheduler): + """Decays the learning rate of each parameter group by gamma every step_size epochs. + + Notice that such decay can happen simultaneously with other changes to the learning rate + from outside this scheduler. When last_epoch=-1, sets initial lr as lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + step_size (int): Period of learning rate decay. + gamma (float): Multiplicative factor of learning rate decay. + Default: 0.1. + last_epoch (int): The index of last epoch. Default: -1. + + Example: + >>> # xdoctest: +SKIP + >>> # Assuming optimizer uses lr = 0.05 for all groups + >>> # lr = 0.05 if epoch < 30 + >>> # lr = 0.005 if 30 <= epoch < 60 + >>> # lr = 0.0005 if 60 <= epoch < 90 + >>> # ... + >>> scheduler = StepLR(optimizer, step_size=30, gamma=0.1) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/StepLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + step_size: int, + gamma: float = 0.1, + last_epoch: int = -1, + ) -> None: # noqa: D107 + self.step_size = step_size + self.gamma = gamma + super().__init__(optimizer, last_epoch) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + If the current epoch is a non-zero multiple of :attr:`step_size`, we + scale the current ``group["lr"]``\s in the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` by :attr:`gamma`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + _warn_get_lr_called_within_step(self) + + if (self.last_epoch == 0) or (self.last_epoch % self.step_size != 0): + return _param_groups_val_list(self.optimizer, "lr") + return [group["lr"] * self.gamma for group in self.optimizer.param_groups] + + def _get_closed_form_lr(self) -> list[float | Tensor]: + r"""Compute learning rates for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` at :attr:`last_epoch` using + a closed-form formula. + + Uses :attr:`base_lrs` to compute learning rates. This method is called + when an epoch is passed to :meth:`step`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + """ + return [ + base_lr * self.gamma ** (self.last_epoch // self.step_size) + for base_lr in self.base_lrs + ] + + +class MultiStepLR(LRScheduler): + """Decays the learning rate of each parameter group by gamma once the number of epoch reaches one of the milestones. + + Notice that such decay can happen simultaneously with other changes to the learning rate + from outside this scheduler. When last_epoch=-1, sets initial lr as lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + milestones (list): List of epoch indices. Must be increasing. + gamma (float): Multiplicative factor of learning rate decay. + Default: 0.1. + last_epoch (int): The index of last epoch. Default: -1. + + Example: + >>> # xdoctest: +SKIP + >>> # Assuming optimizer uses lr = 0.05 for all groups + >>> # lr = 0.05 if epoch < 30 + >>> # lr = 0.005 if 30 <= epoch < 80 + >>> # lr = 0.0005 if epoch >= 80 + >>> scheduler = MultiStepLR(optimizer, milestones=[30, 80], gamma=0.1) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/MultiStepLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + milestones: Iterable[int], + gamma: float = 0.1, + last_epoch: int = -1, + ) -> None: # noqa: D107 + self.milestones = Counter(milestones) + self.gamma = gamma + super().__init__(optimizer, last_epoch) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + If the current epoch is in :attr:`milestones`, decays the + ``group["lr"]``\s in the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` by :attr:`gamma`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + + .. note:: + If the current epoch appears in :attr:`milestones` ``n`` times, we + scale by :attr:`gamma` to the power of ``n`` + """ + _warn_get_lr_called_within_step(self) + + if self.last_epoch not in self.milestones: + return _param_groups_val_list(self.optimizer, "lr") + return [ + group["lr"] * self.gamma ** self.milestones[self.last_epoch] + for group in self.optimizer.param_groups + ] + + def _get_closed_form_lr(self): + r"""Compute learning rates for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` at :attr:`last_epoch` using + a closed-form formula. + + Uses :attr:`base_lrs` to compute learning rates. This method is called + when an epoch is passed to :meth:`step`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + """ + milestones = sorted(self.milestones.elements()) + return [ + base_lr * self.gamma ** bisect_right(milestones, self.last_epoch) + for base_lr in self.base_lrs + ] + + +class ConstantLR(LRScheduler): + """Multiply the learning rate of each parameter group by a small constant factor. + + The multiplication is done until the number of epoch reaches a pre-defined milestone: total_iters. + Notice that such multiplication of the small constant factor can + happen simultaneously with other changes to the learning rate from outside this scheduler. + When last_epoch=-1, sets initial lr as lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + factor (float): The number we multiply learning rate until the milestone. Default: 1./3. + total_iters (int): The number of steps that the scheduler multiplies the learning rate by the factor. + Default: 5. + last_epoch (int): The index of the last epoch. Default: -1. + + Example: + >>> # xdoctest: +SKIP + >>> # Assuming optimizer uses lr = 0.05 for all groups + >>> # lr = 0.025 if epoch == 0 + >>> # lr = 0.025 if epoch == 1 + >>> # lr = 0.025 if epoch == 2 + >>> # lr = 0.025 if epoch == 3 + >>> # ... + >>> # lr = 0.05 if epoch >= 40 + >>> scheduler = ConstantLR(optimizer, factor=0.5, total_iters=40) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/ConstantLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + factor: float = 1.0 / 3, + total_iters: int = 5, + last_epoch: int = -1, + ) -> None: # noqa: D107 + if factor > 1.0 or factor < 0: + raise ValueError( + "Constant multiplicative factor expected to be between 0 and 1." + ) + + self.factor = factor + self.total_iters = total_iters + super().__init__(optimizer, last_epoch) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + When :attr:`last_epoch` is 0, this method scales the ``group["lr"]``\s + in each of the optimizer's :attr:`~torch.optim.Optimizer.param_groups` + by :attr:`factor`. Once :attr:`total_iters` is reached, it undoes this, + scaling by ``1 / factor``. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + _warn_get_lr_called_within_step(self) + + if self.last_epoch == 0: + return [group["lr"] * self.factor for group in self.optimizer.param_groups] + + if self.last_epoch != self.total_iters: + return _param_groups_val_list(self.optimizer, "lr") + + return [ + group["lr"] * (1.0 / self.factor) for group in self.optimizer.param_groups + ] + + def _get_closed_form_lr(self): + r"""Compute learning rates for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` at :attr:`last_epoch` using + a closed-form formula. + + Uses :attr:`base_lrs` to compute learning rates. This method is called + when an epoch is passed to :meth:`step`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + """ + return [ + base_lr + * (self.factor + (self.last_epoch >= self.total_iters) * (1 - self.factor)) + for base_lr in self.base_lrs + ] + + +class LinearLR(LRScheduler): + """Decays the learning rate of each parameter group by linearly changing small multiplicative factor. + + The multiplication is done until the number of epoch reaches a pre-defined milestone: total_iters. + Notice that such decay can happen simultaneously with other changes to the learning rate + from outside this scheduler. When last_epoch=-1, sets initial lr as lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + start_factor (float): The number we multiply learning rate in the first epoch. + The multiplication factor changes towards end_factor in the following epochs. + Default: 1./3. + end_factor (float): The number we multiply learning rate at the end of linear changing + process. Default: 1.0. + total_iters (int): The number of iterations that multiplicative factor reaches to 1. + Default: 5. + last_epoch (int): The index of the last epoch. Default: -1. + + Example: + >>> # xdoctest: +SKIP + >>> # Assuming optimizer uses lr = 0.05 for all groups + >>> # lr = 0.003687 if epoch == 0 + >>> # lr = 0.004875 if epoch == 1 + >>> # lr = 0.006062 if epoch == 2 + >>> # lr = 0.00725 if epoch == 3 + >>> # ... + >>> # lr = 0.05 if epoch >= 40 + >>> scheduler = LinearLR(optimizer, start_factor=0.05, total_iters=40) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/LinearLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + start_factor: float = 1.0 / 3, + end_factor: float = 1.0, + total_iters: int = 5, + last_epoch: int = -1, + ) -> None: # noqa: D107 + if start_factor > 1.0 or start_factor <= 0: + raise ValueError( + "Starting multiplicative factor expected to be greater than 0 and less or equal to 1." + ) + + if end_factor > 1.0 or end_factor < 0: + raise ValueError( + "Ending multiplicative factor expected to be between 0 and 1." + ) + + self.start_factor = start_factor + self.end_factor = end_factor + self.total_iters = total_iters + super().__init__(optimizer, last_epoch) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Scales the ``group["lr"]``\s in the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` such that successive steps + interpolate linearly from :attr:`start_factor` up to :attr:`end_factor` + across :attr:`total_iters` steps. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + _warn_get_lr_called_within_step(self) + + if self.last_epoch == 0: + return [ + group["lr"] * self.start_factor for group in self.optimizer.param_groups + ] + + if self._is_initial or self.last_epoch > self.total_iters: + return _param_groups_val_list(self.optimizer, "lr") + + return [ + group["lr"] + * ( + 1.0 + + (self.end_factor - self.start_factor) + / ( + self.total_iters * self.start_factor + + (self.last_epoch - 1) * (self.end_factor - self.start_factor) + ) + ) + for group in self.optimizer.param_groups + ] + + def _get_closed_form_lr(self): + r"""Compute learning rates for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` at :attr:`last_epoch` using + a closed-form formula. + + Uses :attr:`base_lrs` to compute learning rates. This method is called + when an epoch is passed to :meth:`step`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + """ + return [ + base_lr + * ( + self.start_factor + + (self.end_factor - self.start_factor) + * min(self.total_iters, self.last_epoch) + / self.total_iters + ) + for base_lr in self.base_lrs + ] + + +class ExponentialLR(LRScheduler): + """Decays the learning rate of each parameter group by gamma every epoch. + + When last_epoch=-1, sets initial lr as lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + gamma (float): Multiplicative factor of learning rate decay. + last_epoch (int): The index of last epoch. Default: -1. + + Example: + >>> # xdoctest: +SKIP + >>> scheduler = ExponentialLR(optimizer, gamma=0.95) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/ExponentialLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + gamma: float, + last_epoch: int = -1, + ) -> None: # noqa: D107 + self.gamma = gamma + super().__init__(optimizer, last_epoch) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Multiplies the current ``group["lr"]``\s in the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` by :attr:`gamma`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + _warn_get_lr_called_within_step(self) + + # when loading from a checkpoint, we don't want _initial_step (called from the constructor) + # to update the lr one more step ahead of itself. + if self._is_initial: + return _param_groups_val_list(self.optimizer, "lr") + return [group["lr"] * self.gamma for group in self.optimizer.param_groups] + + def _get_closed_form_lr(self): + r"""Compute learning rates for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` at :attr:`last_epoch` using + a closed-form formula. + + Uses :attr:`base_lrs` to compute learning rates. This method is called + when an epoch is passed to :meth:`step`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + """ + return [base_lr * self.gamma**self.last_epoch for base_lr in self.base_lrs] + + +class SequentialLR(LRScheduler): + """Contains a list of schedulers expected to be called sequentially during the optimization process. + + Specifically, the schedulers will be called according to the milestone points, which should provide exact + intervals by which each scheduler should be called at a given epoch. + + Args: + optimizer (Optimizer): Wrapped optimizer. + schedulers (list): List of chained schedulers. + milestones (list): List of integers that reflects milestone points. + last_epoch (int): The index of last epoch. Default: -1. + + Example: + >>> # xdoctest: +SKIP + >>> # Assuming optimizer uses lr = 0.05 for all groups + >>> # lr = 0.005 if epoch == 0 + >>> # lr = 0.005 if epoch == 1 + >>> # lr = 0.005 if epoch == 2 + >>> # ... + >>> # lr = 0.05 if epoch == 20 + >>> # lr = 0.045 if epoch == 21 + >>> # lr = 0.0405 if epoch == 22 + >>> scheduler1 = ConstantLR(optimizer, factor=0.1, total_iters=20) + >>> scheduler2 = ExponentialLR(optimizer, gamma=0.9) + >>> scheduler = SequentialLR( + ... optimizer, + ... schedulers=[scheduler1, scheduler2], + ... milestones=[20], + ... ) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/SequentialLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + schedulers: list[LRScheduler], + milestones: list[int], + last_epoch: int = -1, + ) -> None: # noqa: D107 + if len(schedulers) < 1: + raise ValueError( + f"{self.__class__.__name__} expects at least one scheduler, but got no scheduler." + ) + + for scheduler_idx, scheduler in enumerate(schedulers): + if not hasattr(scheduler, "optimizer"): + raise TypeError( + f"{self.__class__.__name__} at index {scheduler_idx} should have `optimizer` as its attribute." + ) + if isinstance(scheduler, ReduceLROnPlateau): + raise ValueError( + f"{self.__class__.__name__} does not support `ReduceLROnPlateau` scheduler as it " + "requires additional kwargs to be specified when calling `step`, " + f"but got one at index {scheduler_idx} in the given schedulers sequence." + ) + if optimizer != scheduler.optimizer: + raise ValueError( + f"{self.__class__.__name__} expects all schedulers to belong to the same optimizer, but " + f"got scheduler {scheduler.__class__.__name__} at index {scheduler_idx} has {scheduler.optimizer}, " + f"which is different from {optimizer.__class__.__name__}." + ) + + if len(milestones) != len(schedulers) - 1: + raise ValueError( + "Sequential Schedulers expects number of schedulers provided to be one more " + f"than the number of milestone points, but got number of schedulers {len(schedulers)} and the " + f"number of milestones to be equal to {len(milestones)}" + ) + self._schedulers = schedulers + self._milestones = milestones + self.last_epoch = last_epoch + 1 + self.optimizer = optimizer + + # Reset learning rates back to initial values + for group in self.optimizer.param_groups: + _update_param_group_val(group, "lr", group["initial_lr"]) + + # "Undo" the step performed by other schedulers + self.recursive_undo() + + # Perform the initial step for only the first scheduler + self._schedulers[0]._initial_step() + + self._last_lr = schedulers[0].get_last_lr() + + def recursive_undo(self, sched=None) -> None: + """ + Recursively undo any step performed by the initialisation of + schedulers. + """ + scheds = self if sched is None else sched + + if hasattr(scheds, "_schedulers"): + for s in scheds._schedulers: + self.recursive_undo(s) + elif hasattr(scheds, "last_epoch"): + scheds.last_epoch -= 1 + + def step(self) -> None: # type: ignore[override] + """Perform a step.""" + self.last_epoch += 1 + idx = bisect_right(self._milestones, self.last_epoch) + scheduler = self._schedulers[idx] + if idx > 0 and self._milestones[idx - 1] == self.last_epoch: + scheduler._update_lr(0) + else: + scheduler.step() + + self._last_lr = scheduler.get_last_lr() + + @override + def state_dict(self) -> dict[str, Any]: + """Return the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in ``self.__dict__`` which + is not the optimizer. + The wrapped scheduler states will also be saved. + """ + state_dict = { + key: value + for key, value in self.__dict__.items() + if key not in ("optimizer", "_schedulers") + } + state_dict["_schedulers"] = [None] * len(self._schedulers) + + for idx, s in enumerate(self._schedulers): + # pyrefly: ignore [unsupported-operation] + state_dict["_schedulers"][idx] = s.state_dict() + + return state_dict + + @override + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load the scheduler's state. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + _schedulers = state_dict.pop("_schedulers") + self.__dict__.update(state_dict) + # Restore state_dict keys in order to prevent side effects + # https://github.com/pytorch/pytorch/issues/32756 + state_dict["_schedulers"] = _schedulers + + for idx, s in enumerate(_schedulers): + self._schedulers[idx].load_state_dict(s) + + +class PolynomialLR(LRScheduler): + """Decays the learning rate of each parameter group using a polynomial function in the given total_iters. + + When last_epoch=-1, sets initial lr as lr. + + Args: + optimizer (Optimizer): Wrapped optimizer. + total_iters (int): The number of steps that the scheduler decays the learning rate. Default: 5. + power (float): The power of the polynomial. Default: 1.0. + + Example: + >>> # xdoctest: +SKIP("undefined vars") + >>> # Assuming optimizer uses lr = 0.05 for all groups + >>> # lr = 0.0490 if epoch == 0 + >>> # lr = 0.0481 if epoch == 1 + >>> # lr = 0.0472 if epoch == 2 + >>> # ... + >>> # lr = 0.0 if epoch >= 50 + >>> scheduler = PolynomialLR(optimizer, total_iters=50, power=0.9) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/PolynomialLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + total_iters: int = 5, + power: float = 1.0, + last_epoch: int = -1, + ) -> None: # noqa: D107 + self.total_iters = total_iters + self.power = power + super().__init__(optimizer, last_epoch) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Scales the ``group["lr"]``\s in the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` such that the learning rates + follow + + .. math:: + \texttt{base\_lr} \cdot \left(1 - \frac{\texttt{last\_epoch}} + {\texttt{total\_iters}} \right)^\texttt{power} + + Returns the current learning rates unchanged after :attr:`total_iters` + is reached. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + _warn_get_lr_called_within_step(self) + + if self._is_initial or self.last_epoch > self.total_iters: + return _param_groups_val_list(self.optimizer, "lr") + + decay_factor = ( + (1.0 - self.last_epoch / self.total_iters) + / (1.0 - (self.last_epoch - 1) / self.total_iters) + ) ** self.power + return [group["lr"] * decay_factor for group in self.optimizer.param_groups] + + def _get_closed_form_lr(self) -> list[float | Tensor]: + r"""Compute learning rates for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` at :attr:`last_epoch` using + a closed-form formula. + + Uses :attr:`base_lrs` to compute learning rates. This method is called + when an epoch is passed to :meth:`step`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + """ + return [ + ( + base_lr + * (1.0 - min(self.total_iters, self.last_epoch) / self.total_iters) + ** self.power + ) + for base_lr in self.base_lrs + ] + + +class CosineAnnealingLR(LRScheduler): + r""" + Set the learning rate of each parameter group using a cosine annealing schedule. + + The learning rate is updated recursively using: + + .. math:: + \eta_{t+1} = \eta_{\min} + (\eta_t - \eta_{\min}) \cdot + \frac{1 + \cos\left(\frac{(T_{cur}+1) \pi}{T_{max}}\right)} + {1 + \cos\left(\frac{T_{cur} \pi}{T_{max}}\right)} + + This implements a recursive approximation of the closed-form schedule proposed in + `SGDR: Stochastic Gradient Descent with Warm Restarts`_: + + .. math:: + \eta_t = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min}) \left( + 1 + \cos\left(\frac{T_{cur} \pi}{T_{max}}\right) \right) + + where: + + - :math:`\eta_t` is the learning rate at step :math:`t` + - :math:`T_{cur}` is the number of epochs since the last restart + - :math:`T_{max}` is the maximum number of epochs in a cycle + + Note: + Although SGDR includes periodic restarts, this implementation performs cosine annealing + **without restarts**, so :math:`T_{cur} = t` and increases monotonically with each call + to :meth:`step`. + + Args: + optimizer (Optimizer): Wrapped optimizer. + T_max (int): Maximum number of iterations. + eta_min (float): Minimum learning rate. Default: 0. + last_epoch (int): The index of the last epoch. Default: -1. + + .. _SGDR\: Stochastic Gradient Descent with Warm Restarts: + https://arxiv.org/abs/1608.03983 + + Example: + >>> # xdoctest: +SKIP + >>> num_epochs = 100 + >>> scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs) + >>> for epoch in range(num_epochs): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/CosineAnnealingLR.png + """ + + def __init__( + self, + optimizer: Optimizer, + T_max: int, + eta_min: float = 0.0, + last_epoch: int = -1, + ) -> None: # noqa: D107 + self.T_max = T_max + self.eta_min = eta_min + super().__init__(optimizer, last_epoch) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Scales the ``group["lr"]``\s in the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` such that their learning + rates approximate + + .. math:: + \texttt{eta\_min} + \frac{1}{2} (\texttt{base\_lr} - + \texttt{eta\_min}) \left(1 + \cos\left(\pi \cdot + \frac{\texttt{last\_epoch}}{\texttt{T\_max}}\right) \right) + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + _warn_get_lr_called_within_step(self) + + if self._is_initial: + return _param_groups_val_list(self.optimizer, "lr") + elif self._step_count == 1 and self.last_epoch > 0: + return [ + self.eta_min + + (base_lr - self.eta_min) + * (1 + math.cos((self.last_epoch) * math.pi / self.T_max)) + / 2 + for base_lr, group in zip( + self.base_lrs, self.optimizer.param_groups, strict=True + ) + ] + elif (self.last_epoch - 1 - self.T_max) % (2 * self.T_max) == 0: + return [ + group["lr"] + + (base_lr - self.eta_min) * (1 - math.cos(math.pi / self.T_max)) / 2 + for base_lr, group in zip( + self.base_lrs, self.optimizer.param_groups, strict=True + ) + ] + return [ + (1 + math.cos(math.pi * self.last_epoch / self.T_max)) + / (1 + math.cos(math.pi * (self.last_epoch - 1) / self.T_max)) + * (group["lr"] - self.eta_min) + + self.eta_min + for group in self.optimizer.param_groups + ] + + def _get_closed_form_lr(self) -> list[float | Tensor]: + r"""Compute learning rates for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` at :attr:`last_epoch` using + a closed-form formula. + + Uses :attr:`base_lrs` to compute learning rates. This method is called + when an epoch is passed to :meth:`step`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + """ + return [ + self.eta_min + + (base_lr - self.eta_min) + * (1 + math.cos(math.pi * self.last_epoch / self.T_max)) + / 2 + for base_lr in self.base_lrs + ] + + +class ChainedScheduler(LRScheduler): + """Chains a list of learning rate schedulers. + + Takes in a sequence of chainable learning rate schedulers and calls their + step() functions consecutively in just one call to step(). + + Args: + schedulers (sequence): sequence of chained schedulers. + optimizer (Optimizer, optional): Wrapped optimizer. Default: None. + + Example: + >>> # xdoctest: +SKIP + >>> # Assuming optimizer uses lr = 0.05 for all groups + >>> # lr = 0.05 if epoch == 0 + >>> # lr = 0.0450 if epoch == 1 + >>> # lr = 0.0405 if epoch == 2 + >>> # ... + >>> # lr = 0.00675 if epoch == 19 + >>> # lr = 0.06078 if epoch == 20 + >>> # lr = 0.05470 if epoch == 21 + >>> scheduler1 = ConstantLR(optimizer, factor=0.1, total_iters=20) + >>> scheduler2 = ExponentialLR(optimizer, gamma=0.9) + >>> scheduler = ChainedScheduler([scheduler1, scheduler2], optimizer=optimizer) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/ChainedScheduler.png + """ + + def __init__( + self, schedulers: Sequence[LRScheduler], optimizer: Optimizer | None = None + ) -> None: # noqa: D107 + if len(schedulers) < 1: + raise ValueError( + f"{self.__class__.__name__} expects at least one scheduler to be chained, but got no scheduler." + ) + + optimizer = optimizer or schedulers[0].optimizer + for scheduler_idx, scheduler in enumerate(schedulers): + if not hasattr(scheduler, "optimizer"): + raise TypeError( + f"{self.__class__.__name__} at index {scheduler_idx} should have `optimizer` as its attribute." + ) + if isinstance(scheduler, ReduceLROnPlateau): + raise ValueError( + f"{self.__class__.__name__} does not support `ReduceLROnPlateau` scheduler as it " + "requires additional kwargs to be specified when calling `step`, " + f"but got one at index {scheduler_idx} in the given schedulers sequence." + ) + if optimizer != scheduler.optimizer: + raise ValueError( + f"{self.__class__.__name__} expects all schedulers to belong to the same optimizer, but " + f"got scheduler {scheduler.__class__.__name__} at index {scheduler_idx} has {scheduler.optimizer}, " + f"which is different from {optimizer.__class__.__name__}." + ) + self._schedulers = schedulers + self.optimizer = optimizer + self._last_lr = _param_groups_val_list(self._schedulers[-1].optimizer, "lr") + + def step(self) -> None: # type: ignore[override] + """Perform a step.""" + for scheduler in self._schedulers: + scheduler.step() + self._last_lr = _param_groups_val_list(self._schedulers[-1].optimizer, "lr") + + @override + def state_dict(self) -> dict[str, Any]: + """Return the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in ``self.__dict__`` which + is not the optimizer. + The wrapped scheduler states will also be saved. + """ + state_dict = { + key: value + for key, value in self.__dict__.items() + if key not in ("optimizer", "_schedulers") + } + state_dict["_schedulers"] = [None] * len(self._schedulers) + + for idx, s in enumerate(self._schedulers): + # pyrefly: ignore [unsupported-operation] + state_dict["_schedulers"][idx] = s.state_dict() + + return state_dict + + @override + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load the scheduler's state. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + _schedulers = state_dict.pop("_schedulers") + self.__dict__.update(state_dict) + # Restore state_dict keys in order to prevent side effects + # https://github.com/pytorch/pytorch/issues/32756 + state_dict["_schedulers"] = _schedulers + + for idx, s in enumerate(_schedulers): + self._schedulers[idx].load_state_dict(s) + + +class ReduceLROnPlateau(LRScheduler): + """Reduce learning rate when a metric has stopped improving. + + Models often benefit from reducing the learning rate by a factor + of 2-10 once learning stagnates. This scheduler reads a metrics + quantity and if no improvement is seen for a 'patience' number + of epochs, the learning rate is reduced. + + Args: + optimizer (Optimizer): Wrapped optimizer. + mode (str): One of `min`, `max`. In `min` mode, lr will + be reduced when the quantity monitored has stopped + decreasing; in `max` mode it will be reduced when the + quantity monitored has stopped increasing. Default: 'min'. + factor (float): Factor by which the learning rate will be + reduced. new_lr = lr * factor. Default: 0.1. + patience (int): The number of allowed epochs with no improvement after + which the learning rate will be reduced. + For example, consider the case of having no patience (`patience = 0`). + In the first epoch, a baseline is established and is always considered good as there's no previous baseline. + In the second epoch, if the performance is worse than the baseline, + we have what is considered an intolerable epoch. + Since the count of intolerable epochs (1) is greater than the patience level (0), + the learning rate is reduced at the end of this epoch. + From the third epoch onwards, the learning rate continues to be reduced at the end of each epoch + if the performance is worse than the baseline. If the performance improves or remains the same, + the learning rate is not adjusted. + Default: 10. + threshold (float): Threshold for measuring the new optimum, + to only focus on significant changes. Default: 1e-4. + threshold_mode (str): One of `rel`, `abs`. In `rel` mode, + dynamic_threshold = best * ( 1 + threshold ) in 'max' + mode or best * ( 1 - threshold ) in `min` mode. + In `abs` mode, dynamic_threshold = best + threshold in + `max` mode or best - threshold in `min` mode. Default: 'rel'. + cooldown (int): Number of epochs to wait before resuming + normal operation after lr has been reduced. Default: 0. + min_lr (float or list): A scalar or a list of scalars. A + lower bound on the learning rate of all param groups + or each group respectively. Default: 0. + eps (float): Minimal decay applied to lr. If the difference + between new and old lr is smaller than eps, the update is + ignored. Default: 1e-8. + + Example: + >>> # xdoctest: +SKIP + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = ReduceLROnPlateau(optimizer, "min") + >>> for epoch in range(10): + >>> train(...) + >>> val_loss = validate(...) + >>> # Note that step should be called after validate() + >>> scheduler.step(val_loss) + + .. image:: ../scripts/lr_scheduler_images/ReduceLROnPlateau.png + """ + + def __init__( + self, + optimizer: Optimizer, + mode: Literal["min", "max"] = "min", + factor: float = 0.1, + patience: int = 10, + threshold: float = 1e-4, + threshold_mode: Literal["rel", "abs"] = "rel", + cooldown: int = 0, + min_lr: list[float] | float = 0, + eps: float = 1e-8, + ) -> None: # noqa: D107 + if factor >= 1.0: + raise ValueError("Factor should be < 1.0.") + self.factor = factor + + # Attach optimizer + if not isinstance(optimizer, Optimizer): + raise TypeError(f"{type(optimizer).__name__} is not an Optimizer") + self.optimizer = optimizer + + if isinstance(min_lr, (list, tuple)): + if len(min_lr) != len(optimizer.param_groups): + raise ValueError( + f"expected {len(optimizer.param_groups)} min_lrs, got {len(min_lr)}" + ) + self.default_min_lr = None + self.min_lrs = list(min_lr) + else: + self.default_min_lr = min_lr + self.min_lrs = [min_lr] * len(optimizer.param_groups) + + self.patience = patience + self.cooldown = cooldown + self.eps = eps + self.last_epoch = 0 + self._last_lr = _param_groups_val_list(self.optimizer, "lr") + self._init_is_better( + mode=mode, threshold=threshold, threshold_mode=threshold_mode + ) + self._reset() + + def _reset(self) -> None: + """Reset num_bad_epochs counter and cooldown counter.""" + self.best = self.mode_worse + self.cooldown_counter = 0 + self.num_bad_epochs = 0 + + def step(self, metrics: SupportsFloat, epoch=None) -> None: # type: ignore[override] + """Perform a step.""" + # convert `metrics` to float, in case it's a zero-dim Tensor + current = float(metrics) + if epoch is None: + epoch = self.last_epoch + 1 + else: + warnings.warn(EPOCH_DEPRECATION_WARNING, UserWarning, stacklevel=2) + self.last_epoch = epoch + + if self._is_better(current, self.best): + self.best = current + self.num_bad_epochs = 0 + else: + self.num_bad_epochs += 1 + + if self.in_cooldown: + self.cooldown_counter -= 1 + self.num_bad_epochs = 0 # ignore any bad epochs in cooldown + + if self.num_bad_epochs > self.patience: + self._reduce_lr(epoch) + self.cooldown_counter = self.cooldown + self.num_bad_epochs = 0 + + self._last_lr = _param_groups_val_list(self.optimizer, "lr") + + def _reduce_lr(self, epoch) -> None: + if len(self.optimizer.param_groups) != len(self.min_lrs): + if self.default_min_lr is None: + raise RuntimeError( + "The number of param groups in the `optimizer` " + f"({len(self.optimizer.param_groups)}) differs " + f"from when `ReduceLROnPlateau` was initialized " + f"({len(self.min_lrs)}), usually due to a new " + "param group being added to the optimizer. Please " + "modify the `min_lrs` field to match the length " + "of the `optimizer` param groups." + ) + else: + self.min_lrs = [self.default_min_lr] * len(self.optimizer.param_groups) + + for i, param_group in enumerate(self.optimizer.param_groups): + old_lr = float(param_group["lr"]) + new_lr = max(old_lr * self.factor, self.min_lrs[i]) + if old_lr - new_lr > self.eps: + _update_param_group_val(param_group, "lr", new_lr) + + @property + def in_cooldown(self): # noqa: D102 + return self.cooldown_counter > 0 + + def _is_better(self, a, best): # noqa: D102 + if self.mode == "min" and self.threshold_mode == "rel": + rel_epsilon = 1.0 - self.threshold + return a < best * rel_epsilon + + elif self.mode == "min" and self.threshold_mode == "abs": + return a < best - self.threshold + + elif self.mode == "max" and self.threshold_mode == "rel": + rel_epsilon = self.threshold + 1.0 + return a > best * rel_epsilon + + else: # mode == 'max' and epsilon_mode == 'abs': + return a > best + self.threshold + + def _init_is_better(self, mode, threshold, threshold_mode) -> None: + if mode not in {"min", "max"}: + raise ValueError("mode " + mode + " is unknown!") + if threshold_mode not in {"rel", "abs"}: + raise ValueError("threshold mode " + threshold_mode + " is unknown!") + + # the worse value for the chosen mode + if mode == "min": + self.mode_worse = inf + else: # mode == 'max': + self.mode_worse = -inf + + self.mode = mode + self.threshold = threshold + self.threshold_mode = threshold_mode + + @override + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load the scheduler's state.""" + self.__dict__.update(state_dict) + self._init_is_better( + mode=self.mode, threshold=self.threshold, threshold_mode=self.threshold_mode + ) + + +class CyclicLR(LRScheduler): + r"""Sets the learning rate of each parameter group according to cyclical learning rate policy (CLR). + + The policy cycles the learning rate between two boundaries with a constant frequency, + as detailed in the paper `Cyclical Learning Rates for Training Neural Networks`_. + The distance between the two boundaries can be scaled on a per-iteration + or per-cycle basis. + + Cyclical learning rate policy changes the learning rate after every batch. + `step` should be called after a batch has been used for training. + + This class has three built-in policies, as put forth in the paper: + + * "triangular": A basic triangular cycle without amplitude scaling. + * "triangular2": A basic triangular cycle that scales initial amplitude by half each cycle. + * "exp_range": A cycle that scales initial amplitude by :math:`\text{gamma}^{\text{cycle iterations}}` + at each cycle iteration. + + This implementation was adapted from the github repo: `bckenstler/CLR`_ + + Args: + optimizer (Optimizer): Wrapped optimizer. + base_lr (float or list): Initial learning rate which is the + lower boundary in the cycle for each parameter group. + max_lr (float or list): Upper learning rate boundaries in the cycle + for each parameter group. Functionally, + it defines the cycle amplitude (max_lr - base_lr). + The lr at any cycle is the sum of base_lr + and some scaling of the amplitude; therefore + max_lr may not actually be reached depending on + scaling function. + step_size_up (int): Number of training iterations in the + increasing half of a cycle. Default: 2000 + step_size_down (int): Number of training iterations in the + decreasing half of a cycle. If step_size_down is None, + it is set to step_size_up. Default: None + mode (str): One of {triangular, triangular2, exp_range}. + Values correspond to policies detailed above. + If scale_fn is not None, this argument is ignored. + Default: 'triangular' + gamma (float): Constant in 'exp_range' scaling function: + gamma**(cycle iterations) + Default: 1.0 + scale_fn (function): Custom scaling policy defined by a single + argument lambda function, where + 0 <= scale_fn(x) <= 1 for all x >= 0. + If specified, then 'mode' is ignored. + Default: None + scale_mode (str): {'cycle', 'iterations'}. + Defines whether scale_fn is evaluated on + cycle number or cycle iterations (training + iterations since start of cycle). + Default: 'cycle' + cycle_momentum (bool): If ``True``, momentum is cycled inversely + to learning rate between 'base_momentum' and 'max_momentum'. + Default: True + base_momentum (float or list): Lower momentum boundaries in the cycle + for each parameter group. Note that momentum is cycled inversely + to learning rate; at the peak of a cycle, momentum is + 'base_momentum' and learning rate is 'max_lr'. + Default: 0.8 + max_momentum (float or list): Upper momentum boundaries in the cycle + for each parameter group. Functionally, + it defines the cycle amplitude (max_momentum - base_momentum). + The momentum at any cycle is the difference of max_momentum + and some scaling of the amplitude; therefore + base_momentum may not actually be reached depending on + scaling function. Note that momentum is cycled inversely + to learning rate; at the start of a cycle, momentum is 'max_momentum' + and learning rate is 'base_lr' + Default: 0.9 + last_epoch (int): The index of the last batch. This parameter is used when + resuming a training job. Since `step()` should be invoked after each + batch instead of after each epoch, this number represents the total + number of *batches* computed, not the total number of epochs computed. + When last_epoch=-1, the schedule is started from the beginning. + Default: -1 + + Example: + >>> # xdoctest: +SKIP + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> scheduler = torch.optim.lr_scheduler.CyclicLR( + ... optimizer, + ... base_lr=0.01, + ... max_lr=0.1, + ... step_size_up=10, + ... ) + >>> data_loader = torch.utils.data.DataLoader(...) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/CyclicLR.png + + .. _Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186 + .. _bckenstler/CLR: https://github.com/bckenstler/CLR + """ + + def __init__( + self, + optimizer: Optimizer, + base_lr: float | list[float], + max_lr: float | list[float], + step_size_up: int = 2000, + step_size_down: int | None = None, + mode: Literal["triangular", "triangular2", "exp_range"] = "triangular", + gamma: float = 1.0, + scale_fn: Callable[[float], float] | None = None, + scale_mode: Literal["cycle", "iterations"] = "cycle", + cycle_momentum: bool = True, + base_momentum: float = 0.8, + max_momentum: float = 0.9, + last_epoch: int = -1, + ) -> None: # noqa: D107 + # Attach optimizer + if not isinstance(optimizer, Optimizer): + raise TypeError(f"{type(optimizer).__name__} is not an Optimizer") + self.optimizer = optimizer + + base_lrs = _format_param("base_lr", optimizer, base_lr) + if last_epoch == -1: + for lr, group in zip(base_lrs, optimizer.param_groups, strict=True): + _update_param_group_val(group, "lr", lr) + + self.max_lrs = _format_param("max_lr", optimizer, max_lr) + + # pyrefly: ignore [bad-assignment] + step_size_up = float(step_size_up) + step_size_down = ( + # pyrefly: ignore [bad-assignment] + float(step_size_down) if step_size_down is not None else step_size_up + ) + # pyrefly: ignore [unsupported-operation] + self.total_size = step_size_up + step_size_down + self.step_ratio = step_size_up / self.total_size + + if mode not in ["triangular", "triangular2", "exp_range"] and scale_fn is None: + raise ValueError("mode is invalid and scale_fn is None") + + self.mode = mode + self.gamma = gamma + + self._scale_fn_ref: Callable[[float], float] + self._scale_fn_custom = scale_fn + self.scale_mode = scale_mode + self._init_scale_fn() + + self.cycle_momentum = cycle_momentum + if cycle_momentum: + if ( + "momentum" not in optimizer.defaults + and "betas" not in optimizer.defaults + ): + raise ValueError( + "optimizer must support momentum or beta1 with `cycle_momentum` option enabled" + ) + + self.use_beta1 = "betas" in self.optimizer.defaults + self.base_momentums = _format_param( + "base_momentum", optimizer, base_momentum + ) + self.max_momentums = _format_param("max_momentum", optimizer, max_momentum) + if last_epoch == -1: + for m_momentum, b_momentum, group in zip( + self.max_momentums, + self.base_momentums, + optimizer.param_groups, + strict=True, + ): + if self.use_beta1: + group["betas"] = (m_momentum, *group["betas"][1:]) + else: + group["momentum"] = m_momentum + group["max_momentum"] = m_momentum + group["base_momentum"] = b_momentum + + super().__init__(optimizer, last_epoch) + self.base_lrs = base_lrs + + def _init_scale_fn(self) -> None: + if self._scale_fn_custom is not None: + return + if self.mode == "triangular": + self._scale_fn_ref = self._triangular_scale_fn + self.scale_mode = "cycle" + elif self.mode == "triangular2": + self._scale_fn_ref = self._triangular2_scale_fn + self.scale_mode = "cycle" + elif self.mode == "exp_range": + self._scale_fn_ref = partial(self._exp_range_scale_fn, self.gamma) + self.scale_mode = "iterations" + + def scale_fn(self, x) -> float: + """Get the scaling policy.""" + if self._scale_fn_custom is not None: + return self._scale_fn_custom(x) + else: + return self._scale_fn_ref(x) # static method + + @staticmethod + def _triangular_scale_fn(x: float) -> float: + return 1.0 + + @staticmethod + def _triangular2_scale_fn(x: float) -> float: + return 1 / (2.0 ** (x - 1)) + + @staticmethod + def _exp_range_scale_fn(gamma: float, x: float) -> float: + return gamma**x + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Advances each ``group["lr"]`` in the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` along a cycle between the + group's ``base_lr`` and ``max_lr`` using :meth:`scale_fn`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + + .. note:: + This method treats :attr:`last_epoch` as the index of the previous + batch. + + .. note:: + When :attr:`cycle_momentum` is ``True``, this method has a side + effect of updating the optimizer's momentum. + """ + _warn_get_lr_called_within_step(self) + + cycle = math.floor(1 + self.last_epoch / self.total_size) + x = 1.0 + self.last_epoch / self.total_size - cycle + if x <= self.step_ratio: + scale_factor = x / self.step_ratio + else: + scale_factor = (x - 1) / (self.step_ratio - 1) + + lrs = [] + for base_lr, max_lr in zip(self.base_lrs, self.max_lrs, strict=True): + base_height = (max_lr - base_lr) * scale_factor + if self.scale_mode == "cycle": + lr = base_lr + base_height * self.scale_fn(cycle) + else: + lr = base_lr + base_height * self.scale_fn(self.last_epoch) + lrs.append(lr) + + if self.cycle_momentum: + momentums = [] + for base_momentum, max_momentum in zip( + self.base_momentums, self.max_momentums, strict=True + ): + base_height = (max_momentum - base_momentum) * scale_factor + if self.scale_mode == "cycle": + momentum = max_momentum - base_height * self.scale_fn(cycle) + else: + momentum = max_momentum - base_height * self.scale_fn( + self.last_epoch + ) + momentums.append(momentum) + for param_group, momentum in zip( + self.optimizer.param_groups, momentums, strict=True + ): + if self.use_beta1: + param_group["betas"] = (momentum, *param_group["betas"][1:]) + else: + param_group["momentum"] = momentum + + return lrs + + @override + def state_dict(self) -> dict[str, Any]: # noqa: D102 + """Return the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in ``self.__dict__`` which + is not the optimizer. + The learning rate lambda functions will only be saved if they are callable objects + and not if they are functions or lambdas. + + When saving or loading the scheduler, please make sure to also save or load the state of the optimizer. + """ + state = super().state_dict() + # We are dropping the `_scale_fn_ref` attribute because it is a + # `weakref.WeakMethod` and can't be pickled. + state.pop("_scale_fn_ref", None) + fn = state.pop("_scale_fn_custom") + state["_scale_fn_custom"] = None + if fn is not None and not isinstance(fn, types.FunctionType): + # The _scale_fn_custom will only be saved if it is a callable object + # and not if it is a function or lambda. + state["_scale_fn_custom"] = fn.__dict__.copy() + + return state + + @override + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load the scheduler's state.""" + fn = state_dict.pop("_scale_fn_custom") + super().load_state_dict(state_dict) + if fn is not None: + self._scale_fn_custom.__dict__.update(fn) + self._init_scale_fn() + + +class CosineAnnealingWarmRestarts(LRScheduler): + r"""Set the learning rate of each parameter group using a cosine annealing schedule. + + The :math:`\eta_{max}` is set to the initial lr, :math:`T_{cur}` + is the number of epochs since the last restart and :math:`T_{i}` is the number + of epochs between two warm restarts in SGDR: + + .. math:: + \eta_t = \eta_{min} + \frac{1}{2}(\eta_{max} - \eta_{min})\left(1 + + \cos\left(\frac{T_{cur}}{T_{i}}\pi\right)\right) + + When :math:`T_{cur}=T_{i}`, set :math:`\eta_t = \eta_{min}`. + When :math:`T_{cur}=0` after restart, set :math:`\eta_t=\eta_{max}`. + + It has been proposed in + `SGDR: Stochastic Gradient Descent with Warm Restarts`_. + + Args: + optimizer (Optimizer): Wrapped optimizer. + T_0 (int): Number of iterations until the first restart. + T_mult (int, optional): A factor by which :math:`T_{i}` increases after a restart. Default: 1. + eta_min (float, optional): Minimum learning rate. Default: 0. + last_epoch (int, optional): The index of the last epoch. Default: -1. + + .. _SGDR\: Stochastic Gradient Descent with Warm Restarts: + https://arxiv.org/abs/1608.03983 + + Example: + >>> # xdoctest: +SKIP + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.05) + >>> scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts( + ... optimizer, T_0=20 + ... ) + >>> for epoch in range(100): + >>> train(...) + >>> validate(...) + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/CosineAnnealingWarmRestarts.png + """ + + def __init__( + self, + optimizer: Optimizer, + T_0: int, + T_mult: int = 1, + eta_min: float = 0.0, + last_epoch: int = -1, + ) -> None: # noqa: D107 + if T_0 <= 0 or not isinstance(T_0, int): + raise ValueError(f"Expected positive integer T_0, but got {T_0}") + if T_mult < 1 or not isinstance(T_mult, int): + raise ValueError(f"Expected integer T_mult >= 1, but got {T_mult}") + if not isinstance(eta_min, (float, int)): + raise ValueError( + f"Expected float or int eta_min, but got {eta_min} of type {type(eta_min)}" + ) + self.T_0 = T_0 + self.T_i = T_0 + self.T_mult = T_mult + self.eta_min = eta_min + self.T_cur = last_epoch + super().__init__(optimizer, last_epoch) + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Computes learning rates for the optimizer's + :attr:`~torch.optim.Optimizer.param_groups` following: + + .. math:: + \texttt{eta\_min} + \frac{1}{2}(\texttt{base\_lr} - + \texttt{eta\_min})\left(1 + \cos\left(\pi \cdot + \frac{\texttt{T\_cur}}{\texttt{T\_i}}\right)\right) + + Where :attr:`T_cur` is the number of epochs since the last restart and + :attr:`T_i` is the number of epochs between two restarts. Both + :attr:`T_cur` and :attr:`T_i` are updated in :meth:`step`, and + :attr:`T_i` becomes :attr:`T_mult` times larger after each restart. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + _warn_get_lr_called_within_step(self) + + return [ + self.eta_min + + (base_lr - self.eta_min) + * (1 + math.cos(math.pi * self.T_cur / self.T_i)) + / 2 + for base_lr in self.base_lrs + ] + + @override + def step(self, epoch=None) -> None: + """Step could be called after every batch update. + + Example: + >>> # xdoctest: +SKIP("Undefined vars") + >>> scheduler = CosineAnnealingWarmRestarts(optimizer, T_0, T_mult) + >>> iters = len(dataloader) + >>> for epoch in range(20): + >>> for i, sample in enumerate(dataloader): + >>> inputs, labels = sample['inputs'], sample['labels'] + >>> optimizer.zero_grad() + >>> outputs = net(inputs) + >>> loss = criterion(outputs, labels) + >>> loss.backward() + >>> optimizer.step() + >>> scheduler.step(epoch + i / iters) + + This function can be called in an interleaved way. + + Example: + >>> # xdoctest: +SKIP("Undefined vars") + >>> scheduler = CosineAnnealingWarmRestarts(optimizer, T_0, T_mult) + >>> for epoch in range(20): + >>> scheduler.step() + >>> scheduler.step(26) + >>> scheduler.step() # scheduler.step(27), instead of scheduler(20) + """ + if epoch is None and self.last_epoch < 0: + epoch = 0 + + if epoch is None: + epoch = self.last_epoch + 1 + self.T_cur = self.T_cur + 1 + if self.T_cur >= self.T_i: + self.T_cur = self.T_cur % self.T_i + self.T_i = self.T_i * self.T_mult + else: + if epoch < 0: + raise ValueError(f"Expected non-negative epoch, but got {epoch}") + if epoch >= self.T_0: + if self.T_mult == 1: + self.T_cur = epoch % self.T_0 + else: + n = int( + math.log( + (epoch / self.T_0 * (self.T_mult - 1) + 1), self.T_mult + ) + ) + self.T_cur = epoch - self.T_0 * (self.T_mult**n - 1) / ( + self.T_mult - 1 + ) + self.T_i = self.T_0 * self.T_mult ** (n) + else: + self.T_i = self.T_0 + self.T_cur = epoch + self.last_epoch = math.floor(epoch) + + with _enable_get_lr_call(self): + for param_group, lr in zip( + self.optimizer.param_groups, self.get_lr(), strict=True + ): + _update_param_group_val(param_group, "lr", lr) + + self._last_lr = _param_groups_val_list(self.optimizer, "lr") + + +class _SchedulePhase(TypedDict): + end_step: float + start_lr: str + end_lr: str + start_momentum: str + end_momentum: str + + +class OneCycleLR(LRScheduler): + r"""Sets the learning rate of each parameter group according to the 1cycle learning rate policy. + + The 1cycle policy anneals the learning rate from an initial learning rate to some maximum + learning rate and then from that maximum learning rate to some minimum learning rate much + lower than the initial learning rate. + This policy was initially described in the paper `Super-Convergence: + Very Fast Training of Neural Networks Using Large Learning Rates`_. + + The 1cycle learning rate policy changes the learning rate after every batch. + `step` should be called after a batch has been used for training. + + This scheduler is not chainable. + + Note also that the total number of steps in the cycle can be determined in one + of two ways (listed in order of precedence): + + #. A value for total_steps is explicitly provided. + #. A number of epochs (epochs) and a number of steps per epoch + (steps_per_epoch) are provided. + In this case, the number of total steps is inferred by + total_steps = epochs * steps_per_epoch + + You must either provide a value for total_steps or provide a value for both + epochs and steps_per_epoch. + + The default behaviour of this scheduler follows the fastai implementation of 1cycle, which + claims that "unpublished work has shown even better results by using only two phases". To + mimic the behaviour of the original paper instead, set ``three_phase=True``. + + Args: + optimizer (Optimizer): Wrapped optimizer. + max_lr (float or list): Upper learning rate boundaries in the cycle + for each parameter group. + total_steps (int): The total number of steps in the cycle. Note that + if a value is not provided here, then it must be inferred by providing + a value for epochs and steps_per_epoch. + Default: None + epochs (int): The number of epochs to train for. This is used along + with steps_per_epoch in order to infer the total number of steps in the cycle + if a value for total_steps is not provided. + Default: None + steps_per_epoch (int): The number of steps per epoch to train for. This is + used along with epochs in order to infer the total number of steps in the + cycle if a value for total_steps is not provided. + Default: None + pct_start (float): The percentage of the cycle (in number of steps) spent + increasing the learning rate. + Default: 0.3 + anneal_strategy (str): {'cos', 'linear'} + Specifies the annealing strategy: "cos" for cosine annealing, "linear" for + linear annealing. + Default: 'cos' + cycle_momentum (bool): If ``True``, momentum is cycled inversely + to learning rate between 'base_momentum' and 'max_momentum'. + Default: True + base_momentum (float or list): Lower momentum boundaries in the cycle + for each parameter group. Note that momentum is cycled inversely + to learning rate; at the peak of a cycle, momentum is + 'base_momentum' and learning rate is 'max_lr'. + Default: 0.85 + max_momentum (float or list): Upper momentum boundaries in the cycle + for each parameter group. Functionally, + it defines the cycle amplitude (max_momentum - base_momentum). + Note that momentum is cycled inversely + to learning rate; at the start of a cycle, momentum is 'max_momentum' + and learning rate is 'base_lr' + Default: 0.95 + div_factor (float): Determines the initial learning rate via + initial_lr = max_lr/div_factor + Default: 25 + final_div_factor (float): Determines the minimum learning rate via + min_lr = initial_lr/final_div_factor + Default: 1e4 + three_phase (bool): If ``True``, use a third phase of the schedule to annihilate the + learning rate according to 'final_div_factor' instead of modifying the second + phase (the first two phases will be symmetrical about the step indicated by + 'pct_start'). + last_epoch (int): The index of the last batch. This parameter is used when + resuming a training job. Since `step()` should be invoked after each + batch instead of after each epoch, this number represents the total + number of *batches* computed, not the total number of epochs computed. + When last_epoch=-1, the schedule is started from the beginning. + Default: -1 + + Example: + >>> # xdoctest: +SKIP + >>> data_loader = torch.utils.data.DataLoader(...) + >>> optimizer = torch.optim.SGD(model.parameters(), lr=1e-4, momentum=0.9) + >>> scheduler = torch.optim.lr_scheduler.OneCycleLR( + ... optimizer, max_lr=0.01, steps_per_epoch=len(data_loader), epochs=10 + ... ) + >>> for epoch in range(10): + >>> for batch in data_loader: + >>> train_batch(...) + >>> optimizer.step() + >>> scheduler.step() + + .. image:: ../scripts/lr_scheduler_images/OneCycleLR.png + + .. _Super-Convergence\: Very Fast Training of Neural Networks Using Large Learning Rates: + https://arxiv.org/abs/1708.07120 + """ + + def __init__( + self, + optimizer: Optimizer, + max_lr: float | list[float], + total_steps: int | None = None, + epochs: int | None = None, + steps_per_epoch: int | None = None, + pct_start: float = 0.3, + anneal_strategy: Literal["cos", "linear"] = "cos", + cycle_momentum: bool = True, + base_momentum: float | list[float] = 0.85, + max_momentum: float | list[float] = 0.95, + div_factor: float = 25.0, + final_div_factor: float = 1e4, + three_phase: bool = False, + last_epoch: int = -1, + ) -> None: # noqa: D107 + # Validate optimizer + if not isinstance(optimizer, Optimizer): + raise TypeError(f"{type(optimizer).__name__} is not an Optimizer") + self.optimizer = optimizer + + # Validate total_steps + if total_steps is not None: + if total_steps <= 0 or not isinstance(total_steps, int): + raise ValueError( + f"Expected positive integer total_steps, but got {total_steps}" + ) + self.total_steps = total_steps + elif epochs is not None and steps_per_epoch is not None: + if not isinstance(epochs, int) or epochs <= 0: + raise ValueError(f"Expected positive integer epochs, but got {epochs}") + if not isinstance(steps_per_epoch, int) or steps_per_epoch <= 0: + raise ValueError( + f"Expected positive integer steps_per_epoch, but got {steps_per_epoch}" + ) + self.total_steps = epochs * steps_per_epoch + else: + raise ValueError( + "You must define either total_steps OR (epochs AND steps_per_epoch)" + ) + + self._schedule_phases: list[_SchedulePhase] + if three_phase: + self._schedule_phases = [ + { + "end_step": float(pct_start * self.total_steps) - 1, + "start_lr": "initial_lr", + "end_lr": "max_lr", + "start_momentum": "max_momentum", + "end_momentum": "base_momentum", + }, + { + "end_step": float(2 * pct_start * self.total_steps) - 2, + "start_lr": "max_lr", + "end_lr": "initial_lr", + "start_momentum": "base_momentum", + "end_momentum": "max_momentum", + }, + { + "end_step": self.total_steps - 1, + "start_lr": "initial_lr", + "end_lr": "min_lr", + "start_momentum": "max_momentum", + "end_momentum": "max_momentum", + }, + ] + else: + self._schedule_phases = [ + { + "end_step": float(pct_start * self.total_steps) - 1, + "start_lr": "initial_lr", + "end_lr": "max_lr", + "start_momentum": "max_momentum", + "end_momentum": "base_momentum", + }, + { + "end_step": self.total_steps - 1, + "start_lr": "max_lr", + "end_lr": "min_lr", + "start_momentum": "base_momentum", + "end_momentum": "max_momentum", + }, + ] + + # Validate pct_start + if pct_start < 0 or pct_start > 1 or not isinstance(pct_start, float): + raise ValueError( + f"Expected float between 0 and 1 pct_start, but got {pct_start}" + ) + + # Validate anneal_strategy + if anneal_strategy not in ["cos", "linear"]: + raise ValueError( + f"anneal_strategy must be one of 'cos' or 'linear', instead got {anneal_strategy}" + ) + else: + self._anneal_func_type = anneal_strategy + + # Initialize learning rate variables + max_lrs = _format_param("max_lr", self.optimizer, max_lr) + if last_epoch == -1: + for idx, group in enumerate(self.optimizer.param_groups): + group["initial_lr"] = max_lrs[idx] / div_factor + group["max_lr"] = max_lrs[idx] + group["min_lr"] = group["initial_lr"] / final_div_factor + + # Initialize momentum variables + self.cycle_momentum = cycle_momentum + if self.cycle_momentum: + if ( + "momentum" not in self.optimizer.defaults + and "betas" not in self.optimizer.defaults + ): + raise ValueError( + "optimizer must support momentum or beta1 with `cycle_momentum` option enabled" + ) + self.use_beta1 = "betas" in self.optimizer.defaults + max_momentums = _format_param("max_momentum", optimizer, max_momentum) + base_momentums = _format_param("base_momentum", optimizer, base_momentum) + if last_epoch == -1: + for m_momentum, b_momentum, group in zip( + max_momentums, base_momentums, optimizer.param_groups, strict=True + ): + if self.use_beta1: + group["betas"] = (m_momentum, *group["betas"][1:]) + else: + group["momentum"] = m_momentum + group["max_momentum"] = m_momentum + group["base_momentum"] = b_momentum + + super().__init__(optimizer, last_epoch) + + def _anneal_func(self, *args, **kwargs): + if hasattr(self, "_anneal_func_type"): + if self._anneal_func_type == "cos": + return self._annealing_cos(*args, **kwargs) + elif self._anneal_func_type == "linear": + return self._annealing_linear(*args, **kwargs) + else: + raise ValueError(f"Unknown _anneal_func_type: {self._anneal_func_type}") + else: + # For BC + return self.anneal_func(*args, **kwargs) # type: ignore[attr-defined] + + @staticmethod + def _annealing_cos(start, end, pct): + """Cosine anneal from `start` to `end` as pct goes from 0.0 to 1.0.""" + cos_out = math.cos(math.pi * pct) + 1 + return end + (start - end) / 2.0 * cos_out + + @staticmethod + def _annealing_linear(start, end, pct): + """Linearly anneal from `start` to `end` as pct goes from 0.0 to 1.0.""" + return (end - start) * pct + start + + @override + def get_lr(self) -> list[float | Tensor]: + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Finds the appropriate :attr:`_schedule_phases` entry for the current + step and interpolates between its ``start_lr`` and ``end_lr`` using + :meth:`_anneal_func`. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + + .. note:: + When :attr:`cycle_momentum` is ``True``, this method has a side + effect of updating the optimizer's momentum. + """ + _warn_get_lr_called_within_step(self) + + lrs = [] + step_num = self.last_epoch + + if step_num > self.total_steps: + raise ValueError( + f"Tried to step {step_num} times. The specified number of total steps is {self.total_steps}" + ) + + for group in self.optimizer.param_groups: + start_step = 0.0 + for i, phase in enumerate(self._schedule_phases): + end_step = phase["end_step"] + if step_num <= end_step or i == len(self._schedule_phases) - 1: + pct = (step_num - start_step) / (end_step - start_step) + computed_lr = self._anneal_func( + group[phase["start_lr"]], group[phase["end_lr"]], pct + ) + if self.cycle_momentum: + computed_momentum = self._anneal_func( + group[phase["start_momentum"]], + group[phase["end_momentum"]], + pct, + ) + break + start_step = phase["end_step"] + + lrs.append(computed_lr) # type: ignore[possibly-undefined] + if self.cycle_momentum: + if self.use_beta1: + group["betas"] = (computed_momentum, *group["betas"][1:]) # type: ignore[possibly-undefined] + else: + group["momentum"] = computed_momentum # type: ignore[possibly-undefined] + + return lrs diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/nadam.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/nadam.py new file mode 100644 index 0000000000000000000000000000000000000000..d338333d60934a207d84b5c9d6d8497ab31b42ae --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/nadam.py @@ -0,0 +1,673 @@ +# mypy: allow-untyped-defs +r"""Implementation for the NAdam algorithm.""" + +from typing import cast + +import torch +from torch import Tensor + +from .optimizer import ( + _capturable_doc, + _default_to_fused_or_foreach, + _differentiable_doc, + _disable_dynamo_if_unsupported, + _foreach_doc, + _get_capturable_supported_devices, + _get_scalar_dtype, + _get_value, + _maximize_doc, + _params_doc, + _stack_if_compiling, + _to_scalar, + _use_grad_for_differentiable, + _view_as_real, + Optimizer, + ParamsT, +) + + +__all__ = ["NAdam", "nadam"] + + +class NAdam(Optimizer): # noqa: D101 + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 2e-3, + betas: tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0, + momentum_decay: float = 4e-3, + decoupled_weight_decay: bool = False, + *, + foreach: bool | None = None, + maximize: bool = False, + capturable: bool = False, + differentiable: bool = False, + ) -> None: # noqa: D107 + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + if not 0.0 <= momentum_decay: + raise ValueError(f"Invalid momentum_decay value: {momentum_decay}") + defaults = { + "lr": lr, + "betas": betas, + "eps": eps, + "weight_decay": weight_decay, + "momentum_decay": momentum_decay, + "decoupled_weight_decay": decoupled_weight_decay, + "maximize": maximize, + "foreach": foreach, + "capturable": capturable, + "differentiable": differentiable, + } + super().__init__(params, defaults) + + def __setstate__(self, state): # noqa: D105 + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("maximize", False) + group.setdefault("foreach", None) + group.setdefault("capturable", False) + group.setdefault("differentiable", False) + group.setdefault("decoupled_weight_decay", False) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0: + if not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = ( + torch.tensor( + step_val, dtype=_get_scalar_dtype(), device=p.device + ) + if group["capturable"] + else torch.tensor(step_val, dtype=_get_scalar_dtype()) + ) + if not torch.is_tensor(p_state["mu_product"]): + mu_prod_val = p_state["mu_product"] + p_state["mu_product"] = ( + torch.tensor( + mu_prod_val, dtype=_get_scalar_dtype(), device=p.device + ) + if group["capturable"] + else torch.tensor(mu_prod_val, dtype=_get_scalar_dtype()) + ) + + def _init_group( + self, + group, + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + mu_products, + state_steps, + ): + has_complex = False + for p in group["params"]: + if p.grad is not None: + has_complex |= torch.is_complex(p) + params_with_grad.append(p) + if p.grad.is_sparse: + raise RuntimeError("NAdam does not support sparse gradients") + grads.append(p.grad) + + state = self.state[p] + # Lazy state initialization + if len(state) == 0: + # note(crcrpar): [special device hosting for step] + # Deliberately host `step` and `mu_product` on CPU if capturable is False. + # This is because kernel launches are costly on CUDA and XLA. + state["step"] = ( + torch.zeros((), dtype=_get_scalar_dtype(), device=p.device) + if group["capturable"] + else torch.tensor(0.0, dtype=_get_scalar_dtype()) + ) + state["mu_product"] = ( + torch.ones((), dtype=_get_scalar_dtype(), device=p.device) + if group["capturable"] + else torch.tensor(1.0, dtype=_get_scalar_dtype()) + ) + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + exp_avgs.append(state["exp_avg"]) + exp_avg_sqs.append(state["exp_avg_sq"]) + mu_products.append(state["mu_product"]) + state_steps.append(state["step"]) + return has_complex + + @_use_grad_for_differentiable + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._accelerator_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + exp_avgs: list[Tensor] = [] + exp_avg_sqs: list[Tensor] = [] + mu_products: list[Tensor] = [] + state_steps: list[Tensor] = [] + beta1, beta2 = cast(tuple[float, float], group["betas"]) + + has_complex = self._init_group( + group, + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + mu_products, + state_steps, + ) + + nadam( + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + mu_products, + state_steps, + beta1=beta1, + beta2=beta2, + lr=group["lr"], + weight_decay=group["weight_decay"], + momentum_decay=group["momentum_decay"], + eps=group["eps"], + maximize=group["maximize"], + decoupled_weight_decay=group["decoupled_weight_decay"], + foreach=group["foreach"], + capturable=group["capturable"], + differentiable=group["differentiable"], + has_complex=has_complex, + ) + + return loss + + +NAdam.__doc__ = ( + r"""Implements NAdam algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma_t \text{ (lr)}, \: \beta_1,\beta_2 \text{ (betas)}, + \: \theta_0 \text{ (params)}, \: f(\theta) \text{ (objective)} \\ + &\hspace{13mm} \: \lambda \text{ (weight decay)}, \:\psi \text{ (momentum decay)} \\ + &\hspace{13mm} \: \textit{decoupled\_weight\_decay}, \:\textit{maximize} \\ + &\textbf{initialize} : m_0 \leftarrow 0 \text{ ( first moment)}, + v_0 \leftarrow 0 \text{ ( second moment)} \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + &\hspace{5mm}\textbf{if} \: \textit{maximize}: \\ + &\hspace{10mm}g_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm} \theta_t \leftarrow \theta_{t-1} \\ + &\hspace{5mm} \textbf{if} \: \lambda \neq 0 \\ + &\hspace{10mm}\textbf{if} \: \textit{decoupled\_weight\_decay} \\ + &\hspace{15mm} \theta_t \leftarrow \theta_{t-1} - \gamma \lambda \theta_{t-1} \\ + &\hspace{10mm}\textbf{else} \\ + &\hspace{15mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ + &\hspace{5mm} \mu_t \leftarrow \beta_1 \big(1 - \frac{1}{2} 0.96^{t \psi} \big) \\ + &\hspace{5mm} \mu_{t+1} \leftarrow \beta_1 \big(1 - \frac{1}{2} 0.96^{(t+1)\psi}\big)\\ + &\hspace{5mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ + &\hspace{5mm}v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) g^2_t \\ + &\hspace{5mm}\widehat{m_t} \leftarrow \mu_{t+1} m_t/(1-\prod_{i=1}^{t+1}\mu_i)\\[-1.ex] + & \hspace{11mm} + (1-\mu_t) g_t /(1-\prod_{i=1}^{t} \mu_{i}) \\ + &\hspace{5mm}\widehat{v_t} \leftarrow v_t/\big(1-\beta_2^t \big) \\ + &\hspace{5mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t}/ + \big(\sqrt{\widehat{v_t}} + \epsilon \big) \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to `Incorporating Nesterov Momentum into Adam`_. + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 2e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + momentum_decay (float, optional): momentum momentum_decay (default: 4e-3) + decoupled_weight_decay (bool, optional): whether to decouple the weight + decay as in AdamW to obtain NAdamW. If True, the algorithm does not + accumulate weight decay in the momentum nor variance. (default: False) + {_foreach_doc} + {_maximize_doc} + {_capturable_doc} + {_differentiable_doc} + + .. _Incorporating Nesterov Momentum into Adam: + https://openreview.net/forum?id=OM0jvwB8jIp57ZJjtNEZ + .. _Decoupled Weight Decay Regularization: + https://arxiv.org/abs/1711.05101 + + """ +) + + +def _single_tensor_nadam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + mu_products: list[Tensor], + state_steps: list[Tensor], + *, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + momentum_decay: float, + eps: float, + decoupled_weight_decay: bool, + maximize: bool, + capturable: bool, + differentiable: bool, + has_complex: bool, +) -> None: + if not torch.jit.is_scripting(): + lr = _to_scalar(lr) + + for i, param in enumerate(params): + grad = grads[i] if not maximize else -grads[i] + exp_avg = exp_avgs[i] + exp_avg_sq = exp_avg_sqs[i] + mu_product = mu_products[i] + step_t = state_steps[i] + + if torch.is_complex(param): + param = torch.view_as_real(param) + grad = torch.view_as_real(grad) + exp_avg = torch.view_as_real(exp_avg) + exp_avg_sq = torch.view_as_real(exp_avg_sq) + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices() + if not ( + param.device.type == mu_product.device.type == step_t.device.type + and param.device.type in capturable_supported_devices + ): + raise AssertionError( + f"If capturable=True, params, mu_products and state_steps must be " + f"on supported devices: {capturable_supported_devices}." + ) + + # update step + step_t += 1 + + if capturable: + step = step_t + else: + step = _get_value(step_t) + + bias_correction2 = 1 - beta2**step + + if weight_decay != 0: + if decoupled_weight_decay: + # Perform stepweight decay + param.mul_(1 - lr * weight_decay) + else: + grad = grad.add(param, alpha=weight_decay) + + # calculate the momentum cache \mu^{t} and \mu^{t+1} + mu = beta1 * (1.0 - 0.5 * (0.96 ** (step * momentum_decay))) + mu_next = beta1 * (1.0 - 0.5 * (0.96 ** ((step + 1) * momentum_decay))) + + # update mu_product + mu_product *= mu + + # decay the first and second moment running average coefficient + exp_avg.lerp_(grad, 1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + denom = exp_avg_sq.div(bias_correction2).sqrt() + + if differentiable or capturable: + denom = denom.add(eps) + # Make autograd track the operations + # by updating the grad and exp_avg directly and not using the + # scalar "value" argument of addcdiv. + mu_product_next = mu_product * mu_next + grad = grad * (-lr * (1.0 - mu) / (1.0 - mu_product)) + exp_avg = exp_avg * (-lr * mu_next / (1.0 - mu_product_next)) + param.addcdiv_(grad, denom) + param.addcdiv_(exp_avg, denom) + else: + mu_product_next = _get_value(mu_product) * mu_next + denom.add_(eps) + param.addcdiv_( + grad, denom, value=(-lr * (1.0 - mu) / (1.0 - _get_value(mu_product))) + ) + param.addcdiv_( + exp_avg, + denom, + value=cast(float, (-lr * mu_next) / (1.0 - mu_product_next)), + ) + + +def _multi_tensor_nadam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + mu_products: list[Tensor], + state_steps: list[Tensor], + *, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + momentum_decay: float, + eps: float, + decoupled_weight_decay: bool, + maximize: bool, + capturable: bool, + differentiable: bool, + has_complex: bool, +) -> None: + if len(params) == 0: + return + + if differentiable: + raise AssertionError("_foreach ops don't support autograd") + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices( + supports_xla=False + ) + if not all( + p.device.type == mp.device.type == step.device.type + and p.device.type in capturable_supported_devices + for p, mp, step in zip(params, mu_products, state_steps, strict=True) + ): + raise AssertionError( + "If capturable=True, " + "params, mu_products, and state_steps must be on supported devices: " + f"{capturable_supported_devices}." + ) + + lr = _to_scalar(lr) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, exp_avgs, exp_avg_sqs, mu_products, state_steps] # type: ignore[list-item] + ) + for ( + grouped_params_, + grouped_grads_, + grouped_exp_avgs_, + grouped_exp_avg_sqs_, + grouped_mu_products_, + grouped_state_steps_, + ), _ in grouped_tensors.values(): + grouped_params = cast(list[Tensor], grouped_params_) + grouped_grads = cast(list[Tensor], grouped_grads_) + grouped_exp_avgs = cast(list[Tensor], grouped_exp_avgs_) + grouped_exp_avg_sqs = cast(list[Tensor], grouped_exp_avg_sqs_) + grouped_mu_products = cast(list[Tensor], grouped_mu_products_) + grouped_state_steps = cast(list[Tensor], grouped_state_steps_) + + # handle complex + if has_complex: + _view_as_real( + grouped_params, grouped_grads, grouped_exp_avgs, grouped_exp_avg_sqs + ) + + if maximize: + grouped_grads = torch._foreach_neg(grouped_grads) # type: ignore[assignment] + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and grouped_state_steps[0].is_cpu: + torch._foreach_add_( + grouped_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(grouped_state_steps, 1) + + if weight_decay != 0: + if decoupled_weight_decay: + # Perform stepweight decay + torch._foreach_mul_(grouped_params, 1 - lr * weight_decay) + else: + # Reuse the intermediate memory (grouped_grads) already allocated for maximize + if maximize: + torch._foreach_add_( + grouped_grads, grouped_params, alpha=weight_decay + ) + else: + grouped_grads = torch._foreach_add( # type: ignore[assignment] + grouped_grads, grouped_params, alpha=weight_decay + ) + + # Decay the first and second moment running average coefficient + torch._foreach_lerp_(grouped_exp_avgs, grouped_grads, 1 - beta1) + + torch._foreach_mul_(grouped_exp_avg_sqs, beta2) + torch._foreach_addcmul_( + grouped_exp_avg_sqs, grouped_grads, grouped_grads, 1 - beta2 + ) + + exp_avg_sq_sqrt = torch._foreach_sqrt(grouped_exp_avg_sqs) + + bias_correction_sqrt: tuple[Tensor, ...] | list[Tensor] + mus: tuple[Tensor, ...] | list[Tensor] + mu_nexts: tuple[Tensor, ...] | list[Tensor] + if capturable: + # mus will be beta1 * (1 - 0.5 * 0.96 ** (step * momentum_decay)) + exponent = torch._foreach_mul(grouped_state_steps, momentum_decay) + mus = torch._foreach_pow(0.96, exponent) + torch._foreach_mul_(mus, -0.5) + torch._foreach_add_(mus, 1.0) + torch._foreach_mul_(mus, beta1) + + # mu_nexts will be beta1 * (1 - 0.5 * 0.96 ** ((step + 1) * momentum_decay)) + torch._foreach_add_(exponent, momentum_decay) + mu_nexts = torch._foreach_pow(0.96, exponent) + torch._foreach_mul_(mu_nexts, -0.5) + torch._foreach_add_(mu_nexts, 1.0) + torch._foreach_mul_(mu_nexts, beta1) + + # save peak memory as we don't need exponent anymore + del exponent + + bias_correction_sqrt = torch._foreach_pow(beta2, grouped_state_steps) + # foreach_sub doesn't allow a scalar as the first arg + torch._foreach_sub_(bias_correction_sqrt, 1.0) + torch._foreach_neg_(bias_correction_sqrt) + torch._foreach_sqrt_(bias_correction_sqrt) + else: + bias_correction_sqrt = [ + (1 - beta2 ** _get_value(step)) ** 0.5 for step in grouped_state_steps + ] + mus = [ + beta1 * (1.0 - 0.5 * (0.96 ** (_get_value(step) * momentum_decay))) + for step in grouped_state_steps + ] + mu_nexts = [ + beta1 + * (1.0 - 0.5 * (0.96 ** ((_get_value(step) + 1) * momentum_decay))) + for step in grouped_state_steps + ] + + # update mu_products + torch._foreach_mul_(grouped_mu_products, mus) + + torch._foreach_div_(exp_avg_sq_sqrt, bias_correction_sqrt) + torch._foreach_add_(exp_avg_sq_sqrt, eps) + + # explicitly delete bias_correction refs to save memory + del bias_correction_sqrt + + if capturable: + # Build up the step_size multiplier for grad, reusing mus' memory + torch._foreach_sub_(mus, 1.0) + torch._foreach_mul_(mus, lr) + # foreach_sub doesn't allow a scalar as the first arg + denom = torch._foreach_sub(grouped_mu_products, 1.0) + torch._foreach_neg_(denom) + torch._foreach_div_(mus, denom) + # - lr * (1 - mu) / (1 - mu_product) + step_size_grads = mus + # explicitly delete denom to save memory + del denom + + # Build up the step_size multiplier for exp_avg, reusing mu_nexts' memory + denom = torch._foreach_mul(grouped_mu_products, mu_nexts) + torch._foreach_mul_(mu_nexts, lr) + # foreach_sub doesn't allow a scalar as the first arg, but it's okay because + # we need a negative here anyway + torch._foreach_sub_(denom, 1.0) + torch._foreach_div_(mu_nexts, denom) + # - lr * mu_next / (1 - mu_product * mu_next) + step_size_expavg = mu_nexts + # explicitly delete denom to save memory + del denom + + # we cannot inplace into step_size_grads cuz it is a list of ScalarTensors + # and mul'ing with grouped_grads will result in a list of bigger Tensors + numerator = torch._foreach_mul(step_size_grads, grouped_grads) + torch._foreach_addcmul_(numerator, step_size_expavg, grouped_exp_avgs) + + # finally, update params + torch._foreach_addcdiv_(grouped_params, numerator, exp_avg_sq_sqrt) + else: + step_size_grads = _stack_if_compiling( + [ + (_get_value(lr) * (1.0 - mu) / (1.0 - _get_value(mu_product))) * -1 + for mu_product, mu in zip(grouped_mu_products, mus, strict=True) + ] + ) + step_size_expavg = _stack_if_compiling( + [ + ( + _get_value(lr) + * mu_next + / (1.0 - _get_value(mu_product) * mu_next) + ) + * -1 + for mu_product, mu_next in zip( + grouped_mu_products, mu_nexts, strict=True + ) + ] + ) + + torch._foreach_addcdiv_( + grouped_params, + grouped_grads, + exp_avg_sq_sqrt, + step_size_grads, # type: ignore[arg-type] + ) + torch._foreach_addcdiv_( + grouped_params, + grouped_exp_avgs, + exp_avg_sq_sqrt, + step_size_expavg, # type: ignore[arg-type] + ) + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_nadam) +def nadam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + mu_products: list[Tensor], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + decoupled_weight_decay: bool = False, + foreach: bool | None = None, + capturable: bool = False, + differentiable: bool = False, + has_complex: bool = False, + maximize: bool = False, + *, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + momentum_decay: float, + eps: float, +) -> None: + r"""Functional API that performs NAdam algorithm computation. + + See :class:`~torch.optim.NAdam` for details. + """ + if not all(isinstance(t, torch.Tensor) for t in state_steps): + raise RuntimeError( + "API has changed, `state_steps` argument must contain a list of singleton tensors" + ) + + if not all(isinstance(t, torch.Tensor) for t in mu_products): + raise RuntimeError( + "API has changed, `mu_products` argument must contain a list of singleton tensors" + ) + + if foreach is None: + _, foreach = _default_to_fused_or_foreach( + params, differentiable, use_fused=False + ) + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + + if foreach and not torch.jit.is_scripting(): + func = _multi_tensor_nadam + else: + func = _single_tensor_nadam + + func( + params, + grads, + exp_avgs, + exp_avg_sqs, + mu_products, + state_steps, + beta1=beta1, + beta2=beta2, + lr=lr, + weight_decay=weight_decay, + momentum_decay=momentum_decay, + maximize=maximize, + decoupled_weight_decay=decoupled_weight_decay, + eps=eps, + capturable=capturable, + differentiable=differentiable, + has_complex=has_complex, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/optimizer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..cf9a3c53bf8e8a8952a06df9864fb0226e88867a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/optimizer.py @@ -0,0 +1,1195 @@ +# mypy: allow-untyped-defs +"""Base optimizer.""" + +import functools +import warnings +from collections import defaultdict, OrderedDict +from collections.abc import Callable, Hashable, Iterable, Sequence +from copy import deepcopy +from itertools import chain +from typing import Any, cast, overload, TypeAlias, TypeVar +from typing_extensions import deprecated, ParamSpec, Self + +import torch +import torch.utils.hooks as hooks +from torch.utils._foreach_utils import ( + _get_foreach_kernels_supported_devices, + _get_fused_kernels_supported_devices, + _group_tensors_by_device_and_dtype, + Indices, + TensorListList, +) +from torch.utils.hooks import RemovableHandle + + +_T = TypeVar("_T") +_P = ParamSpec("_P") + +Args: TypeAlias = tuple[Any, ...] +Kwargs: TypeAlias = dict[str, Any] +StateDict: TypeAlias = dict[str, Any] +DeviceDict: TypeAlias = dict[torch.device | None, torch.Tensor] +DeviceDtypeDict: TypeAlias = dict[tuple[torch.device, torch.dtype] | None, torch.Tensor] + +GlobalOptimizerPreHook: TypeAlias = Callable[ + ["Optimizer", Args, Kwargs], tuple[Args, Kwargs] | None +] +GlobalOptimizerPostHook: TypeAlias = Callable[["Optimizer", Args, Kwargs], None] + +__all__ = [ + "Optimizer", + "register_optimizer_step_pre_hook", + "register_optimizer_step_post_hook", +] +_global_optimizer_pre_hooks: dict[int, GlobalOptimizerPreHook] = OrderedDict() +_global_optimizer_post_hooks: dict[int, GlobalOptimizerPostHook] = OrderedDict() +_foreach_supported_types = [torch.Tensor, torch.nn.parameter.Parameter] + + +class _RequiredParameter: + """Singleton class representing a required parameter for an Optimizer.""" + + def __repr__(self) -> str: + return "" + + +required = _RequiredParameter() + + +def _use_grad_for_differentiable(func: Callable[_P, _T]) -> Callable[_P, _T]: + def _use_grad(*args: _P.args, **kwargs: _P.kwargs) -> _T: + import torch._dynamo + + self = cast(Optimizer, args[0]) # assume first positional arg is `self` + prev_grad = torch.is_grad_enabled() + try: + # Note on graph break below: + # we need to graph break to ensure that aot respects the no_grad annotation. + # This is important for perf because without this, functionalization will generate an epilogue + # which updates the mutated parameters of the optimizer which is *not* visible to inductor, as a result, + # inductor will allocate for every parameter in the model, which is horrible. + # With this, aot correctly sees that this is an inference graph, and functionalization will generate + # an epilogue which is appended to the graph, which *is* visible to inductor, as a result, inductor sees that + # step is in place and is able to avoid the extra allocation. + # In the future, we will either 1) continue to graph break on backward, so this graph break does not matter + # or 2) have a fully fused forward and backward graph, which will have no_grad by default, and we can remove this + # graph break to allow the fully fused fwd-bwd-optimizer graph to be compiled. + # see https://github.com/pytorch/pytorch/issues/104053 + torch.set_grad_enabled(self.defaults["differentiable"]) + torch._dynamo.graph_break() + ret = func(*args, **kwargs) + finally: + torch._dynamo.graph_break() + torch.set_grad_enabled(prev_grad) + return ret + + functools.update_wrapper(_use_grad, func) + return _use_grad + + +def _get_value(x): + # item is significantly faster than a cpu tensor in eager mode + if not torch.jit.is_scripting() and torch.compiler.is_compiling(): + return x + else: + return x.item() if isinstance(x, torch.Tensor) else x + + +def _stack_if_compiling(x): + if not torch.jit.is_scripting() and torch.compiler.is_compiling(): + return torch.stack(x) + else: + return x + + +def _disable_dynamo_if_unsupported( + single_tensor_fn: Callable[..., object] | None = None, +) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: + # workaround for torchscript BC + # it requires all called functions to be in the + # global environment at the site at which the + # maybe_fallback closure is created + if single_tensor_fn: + globals()[single_tensor_fn.__name__] = single_tensor_fn + + def wrapper(func: Callable[_P, _T]) -> Callable[_P, _T]: + import inspect + + disabled_func = torch._disable_dynamo(func) + ps = inspect.signature(func).parameters + has_state_steps = True + try: + state_steps_ind = list(ps.keys()).index("state_steps") + except ValueError: + has_state_steps = False + + # Today, there are cases where we stack state steps + # and pass them as the value arg of foreach ops. + # Having state steps on cuda as the value arg is not supported in eager, + # but this only occurs in the rare case that the user explicitly deletes + # the capturable flag. If capturable=True, this is not a problem. + @functools.wraps(func) + def maybe_fallback(*args: _P.args, **kwargs: _P.kwargs): + if torch.compiler.is_compiling() and ( + not kwargs.get("capturable", False) + and has_state_steps + and (arg := args[state_steps_ind]) + and isinstance(arg, Sequence) + and arg[0].is_cuda + or ( + "state_steps" in kwargs + and (kwarg := kwargs["state_steps"]) + and isinstance(kwarg, Sequence) + and kwarg[0].is_cuda + ) + ): + return disabled_func(*args, **kwargs) + else: + return func(*args, **kwargs) + + return maybe_fallback + + return wrapper + + +# For any optimizer with a faster implementation, we attempt to default to the +# fastest + stablest whenever possible. For foreach, the requirements are to have +# native params all on CUDA. For fused, there's currently the additional requirement +# that the tensors' dtypes must be floating point. Neither alternative supports +# torch.jit.script nor differentiable, so we fall back to the single tensor +# implementation in those cases. +def _default_to_fused_or_foreach( + params: list[torch.Tensor], differentiable: bool, use_fused: bool = False +) -> tuple[bool, bool]: + if torch.jit.is_scripting() or differentiable: + return False, False + + fused_supported_devices = _get_fused_kernels_supported_devices() + foreach_supported_devices = _get_foreach_kernels_supported_devices() + fused = use_fused and all( + p is None + or ( + type(p) in _foreach_supported_types + and p.device.type in fused_supported_devices + and torch.is_floating_point(p) + ) + for p in params + ) + foreach = not fused and all( + p is None + or ( + type(p) in _foreach_supported_types + and p.device.type in foreach_supported_devices + ) + for p in params + ) + return fused, foreach + + +def _device_dtype_check_for_fused( + p: torch.Tensor, cuda_unsupported: bool = False +) -> None: + fused_supported_devices = _get_fused_kernels_supported_devices() + if cuda_unsupported: + fused_supported_devices.remove("cuda") + if not (p.device.type in fused_supported_devices and torch.is_floating_point(p)): + raise RuntimeError( + "`fused=True` requires all the params to be floating point Tensors of " + f"supported devices: {fused_supported_devices} but {p.dtype} and {p.device.type}" + ) + + +def _view_as_real(params, *state_and_grads) -> None: + for i, p in enumerate(params): + if torch.is_complex(p): + params[i] = torch.view_as_real(params[i]) + for s in state_and_grads: + s[i] = torch.view_as_real(s[i]) + + +def _get_scalar_dtype(is_fused=None): + if is_fused: + return torch.float32 + return ( + torch.float64 if torch.get_default_dtype() == torch.float64 else torch.float32 + ) + + +def _get_capturable_supported_devices(supports_xla: bool = True) -> list[str]: + r"""Return the device type list that supports capturable optimizer.""" + capturable_supported_devices = ["cuda", "xpu", "hpu"] + if not torch.jit.is_scripting(): + capturable_supported_devices.append(torch._C._get_privateuse1_backend_name()) + if supports_xla: + capturable_supported_devices.append("xla") + return capturable_supported_devices + + +def _to_scalar(x: float | torch.Tensor): + r"""This function converts a hyperparameter to a 0-dimension (scalar) tensor + if it is a nonzero-dimensions 1-element tensor. If it is not a tensor, it is + kept as is. + + Args: + x (float or Tensor): A hyperparameter of the optimizer. + If it is Tensor, it is needed to be 1-element. + + Returns: + float or Tensor: + a scalar tensor if x is Tensor otherwise Python scalar (float) value. + """ + if isinstance(x, torch.Tensor) and x.dim() != 0: + return x.squeeze() + else: + return x + + +# Common doc strings among optimizers +_params_doc = r"""params (iterable): iterable of parameters or named_parameters to optimize + or iterable of dicts defining parameter groups. When using named_parameters, + all parameters in all groups should be named""" + +_foreach_doc = r"""foreach (bool, optional): whether foreach implementation of optimizer + is used. If unspecified by the user (so foreach is None), we will try to use + foreach over the for-loop implementation on CUDA, since it is usually + significantly more performant. Note that the foreach implementation uses + ~ sizeof(params) more peak memory than the for-loop version due to the intermediates + being a tensorlist vs just one tensor. If memory is prohibitive, batch fewer + parameters through the optimizer at a time or switch this flag to False (default: None)""" + +_fused_doc = r"""fused (bool, optional): whether the fused implementation is used. + Currently, `torch.float64`, `torch.float32`, `torch.float16`, and `torch.bfloat16` + are supported. (default: None) + + .. note:: The foreach and fused implementations are typically faster than the for-loop, + single-tensor implementation, with fused being theoretically fastest with both + vertical and horizontal fusion. As such, if the user has not specified either + flag (i.e., when foreach = fused = None), we will attempt defaulting to the foreach + implementation when the tensors are all on CUDA. Why not fused? Since the fused + implementation is relatively new, we want to give it sufficient bake-in time. + To specify fused, pass True for fused. To force running the for-loop + implementation, pass False for either foreach or fused. """ + +_capturable_doc = r"""capturable (bool, optional): whether this instance is safe to + capture in a graph, whether for CUDA graphs or for torch.compile support. + Tensors are only capturable when on supported :ref:`accelerators`. + Passing True can impair ungraphed performance, so if you don't intend to graph + capture this instance, leave it False (default: False)""" + +_differentiable_doc = r"""differentiable (bool, optional): whether autograd should + occur through the optimizer step in training. Otherwise, the step() + function runs in a torch.no_grad() context. Setting to True can impair + performance, so leave it False if you don't intend to run autograd + through this instance (default: False)""" + +_maximize_doc = r"""maximize (bool, optional): maximize the objective with respect to the + params, instead of minimizing (default: False)""" + + +def register_optimizer_step_pre_hook(hook: GlobalOptimizerPreHook) -> RemovableHandle: + r"""Register a pre hook common to all optimizers. + + The hook should have the following signature:: + + hook(optimizer, args, kwargs) -> None or modified args and kwargs + + Args: + hook (Callable): A user defined hook which is registered on all optimizers. + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + a handle that can be used to remove the added hook by calling + ``handle.remove()`` + """ + handle = hooks.RemovableHandle(_global_optimizer_pre_hooks) + _global_optimizer_pre_hooks[handle.id] = hook + return handle + + +def register_optimizer_step_post_hook(hook: GlobalOptimizerPostHook) -> RemovableHandle: + r"""Register a post hook common to all optimizers. + + The hook should have the following signature:: + + hook(optimizer, args, kwargs) -> None + + Args: + hook (Callable): A user defined hook which is registered on all optimizers. + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + a handle that can be used to remove the added hook by calling + ``handle.remove()`` + """ + handle = hooks.RemovableHandle(_global_optimizer_post_hooks) + _global_optimizer_post_hooks[handle.id] = hook + return handle + + +ParamsT: TypeAlias = ( + Iterable[torch.Tensor] + | Iterable[dict[str, Any]] + | Iterable[tuple[str, torch.Tensor]] +) + +R = TypeVar("R") +T = TypeVar("T") + + +class Optimizer: + r"""Base class for all optimizers. + + .. warning:: + Parameters need to be specified as collections that have a deterministic + ordering that is consistent between runs. Examples of objects that don't + satisfy those properties are sets and iterators over values of dictionaries. + + Args: + params (iterable): an iterable of :class:`torch.Tensor` s or + :class:`dict` s. Specifies what Tensors should be optimized. + defaults: (dict): a dict containing default values of optimization + options (used when a parameter group doesn't specify them). + """ + + OptimizerPreHook: TypeAlias = Callable[ + [Self, Args, Kwargs], # type: ignore[misc] + tuple[Args, Kwargs] | None, + ] + OptimizerPostHook: TypeAlias = Callable[[Self, Args, Kwargs], None] # type: ignore[misc] + + _optimizer_step_pre_hooks: dict[int, OptimizerPreHook] + _optimizer_step_post_hooks: dict[int, OptimizerPostHook] + # pyrefly: ignore [not-a-type] + _optimizer_state_dict_pre_hooks: 'OrderedDict[int, Callable[["Optimizer"], None]]' + _optimizer_state_dict_post_hooks: ( + # pyrefly: ignore [not-a-type] + 'OrderedDict[int, Callable[["Optimizer", StateDict], StateDict | None]]' + ) + _optimizer_load_state_dict_pre_hooks: ( + # pyrefly: ignore [not-a-type] + 'OrderedDict[int, Callable[["Optimizer", StateDict], StateDict | None]]' + ) + _optimizer_load_state_dict_post_hooks: ( + # pyrefly: ignore [not-a-type] + 'OrderedDict[int, Callable[["Optimizer"], None]]' + ) + + def __init__(self, params: ParamsT, defaults: dict[str, Any]) -> None: # noqa: D107 + torch._C._log_api_usage_once("python.optimizer") + self.defaults = defaults + self._optimizer_step_pre_hooks = OrderedDict() + self._optimizer_step_post_hooks = OrderedDict() + self._optimizer_state_dict_pre_hooks = OrderedDict() + self._optimizer_state_dict_post_hooks = OrderedDict() + self._optimizer_load_state_dict_pre_hooks = OrderedDict() + self._optimizer_load_state_dict_post_hooks = OrderedDict() + + self._patch_step_function() + + if isinstance(params, torch.Tensor): + raise TypeError( + "params argument given to the optimizer should be " + "an iterable of Tensors or dicts, but got " + torch.typename(params) + ) + + self.state: defaultdict[torch.Tensor, Any] = defaultdict(dict) + self.param_groups: list[dict[str, Any]] = [] + + param_groups = list(params) + if len(param_groups) == 0: + raise ValueError("optimizer got an empty parameter list") + if not isinstance(param_groups[0], dict): + param_groups = [{"params": param_groups}] + + for param_group in param_groups: + self.add_param_group(cast(dict, param_group)) + + # Allows _accelerator_graph_capture_health_check to rig a poor man's TORCH_WARN_ONCE in python, + # which I don't think exists + # https://github.com/pytorch/pytorch/issues/72948 + self._warned_capturable_if_run_uncaptured = True + + def __getstate__(self) -> dict[str, Any]: # noqa: D105 + return { + "defaults": self.defaults, + "state": self.state, + "param_groups": self.param_groups, + } + + def __setstate__(self, state: dict[str, Any]) -> None: # noqa: D105 + self.__dict__.update(state) + if "_optimizer_step_pre_hooks" not in self.__dict__: + self._optimizer_step_pre_hooks = OrderedDict() + if "_optimizer_step_post_hooks" not in self.__dict__: + self._optimizer_step_post_hooks = OrderedDict() + if "_optimizer_state_dict_pre_hooks" not in self.__dict__: + self._optimizer_state_dict_pre_hooks = OrderedDict() + if "_optimizer_state_dict_post_hooks" not in self.__dict__: + self._optimizer_state_dict_post_hooks = OrderedDict() + if "_optimizer_load_state_dict_pre_hooks" not in self.__dict__: + self._optimizer_load_state_dict_pre_hooks = OrderedDict() + if "_optimizer_load_state_dict_post_hooks" not in self.__dict__: + self._optimizer_load_state_dict_post_hooks = OrderedDict() + self._patch_step_function() # To support multiprocessing pickle/unpickle + self.defaults.setdefault("differentiable", False) + + def __repr__(self) -> str: # noqa: D105 + format_string = self.__class__.__name__ + " (" + for i, group in enumerate(self.param_groups): + format_string += "\n" + format_string += f"Parameter Group {i}\n" + for key in sorted(group.keys()): + if key != "params": + format_string += f" {key}: {group[key]}\n" + format_string += ")" + return format_string + + # Currently needed by Adam and AdamW + def _accelerator_graph_capture_health_check(self) -> None: + # Note [torch.compile x capturable] + # If we are compiling, we try to take the capturable path automatically by + # setting the flag to True during tracing. Due to this, we skip all the checks + # normally required for determining whether we can use CUDA/XPU graphs and + # shunt the responsibility to torch.inductor. This saves time during tracing + # since the checks are slow without sacrificing UX since inductor will warn + # later if CUDA/XPU graphs cannot be enabled, e.g., + # https://github.com/pytorch/pytorch/blob/d3ba8901d8640eb16f88b2bfef9df7fa383d4b47/torch/_inductor/compile_fx.py#L390. + # Thus, when compiling, inductor will determine if cudagraphs + # can be enabled based on whether there is input mutation or CPU tensors. + if torch.compiler.is_compiling(): + return + + # Determine available accelerator device + accelerator = None + if torch.cuda.is_available(): + accelerator = (torch.cuda, "CUDA") + elif torch.xpu.is_available(): + accelerator = (torch.xpu, "XPU") + + if accelerator: + device_module, device_name = accelerator + capturing = device_module.is_current_stream_capturing() + + if capturing and not all( + group["capturable"] for group in self.param_groups + ): + raise RuntimeError( + f"Attempting {device_name} graph capture of step() for an instance of " + + self.__class__.__name__ + + " but param_groups' capturable is False." + ) + + if ( + (not getattr(self, "_warned_capturable_if_run_uncaptured", False)) + and all(group["capturable"] for group in self.param_groups) + and (not capturing) + ): + warnings.warn( + "This instance was constructed with capturable=True or some of all the param_groups came with capturable=True, " + f"but step() is running without {device_name} graph capture. If you never intend to graph-capture this " + "instance, capturable=True can impair performance, and you should set capturable=False.", + stacklevel=2, + ) + self._warned_capturable_if_run_uncaptured = True + + # Backward compatibility alias for internal callers still using the old name. + _cuda_graph_capture_health_check = deprecated( + "Use _accelerator_graph_capture_health_check instead", + )(_accelerator_graph_capture_health_check) + + def _optimizer_step_code(self) -> None: + """Entry point for `torch.profile.profiler`. + + When python tracing is enabled the profiler will hook into this + function at the CPython level to inspect the optimizer's parameters and + param groups. It is called it after `step()` since many optimizers + lazily initialize state. + + This is a workaround due to lack of a proper step hook on the optimizer, + and will be removed if it exists. + """ + + @staticmethod + def profile_hook_step(func: Callable[_P, R]) -> Callable[_P, R]: # noqa: D102 + @functools.wraps(func) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> R: + self, *_ = args + self = cast(Optimizer, self) + profile_name = f"Optimizer.step#{self.__class__.__name__}.step" + with torch.autograd.profiler.record_function(profile_name): + # call optimizer step pre hooks + for pre_hook in chain( + _global_optimizer_pre_hooks.values(), + self._optimizer_step_pre_hooks.values(), + ): + result = pre_hook(self, args, kwargs) + if result is not None: + if isinstance(result, tuple) and len(result) == 2: + args, kwargs = result # type: ignore[assignment] + else: + raise RuntimeError( + f"{func} must return None or a tuple of (new_args, new_kwargs), but got {result}." + ) + + # pyrefly: ignore [invalid-param-spec] + out = func(*args, **kwargs) + self._optimizer_step_code() + + # call optimizer step post hooks + for post_hook in chain( + self._optimizer_step_post_hooks.values(), + _global_optimizer_post_hooks.values(), + ): + post_hook(self, args, kwargs) + + return out + + return wrapper + + @staticmethod + def _group_tensors_by_device_and_dtype( + tensorlistlist: TensorListList, + with_indices: bool = False, + ) -> ( + dict[tuple[None, None], tuple[TensorListList, Indices]] + | dict[tuple[torch.device, torch.dtype], tuple[TensorListList, Indices]] + ): + """Group a list of lists of tensors by device and dtype. + + Skips this step if we are compiling since this will occur during inductor lowering. + """ + if torch.compiler.is_compiling(): + return {(None, None): (tensorlistlist, list(range(len(tensorlistlist[0]))))} + else: + return _group_tensors_by_device_and_dtype(tensorlistlist, with_indices) # type: ignore[return-value, arg-type] + + def _patch_step_function(self) -> None: + self._zero_grad_profile_name = ( + f"Optimizer.zero_grad#{self.__class__.__name__}.zero_grad" + ) + hooked = getattr(self.__class__.step, "hooked", None) + if not hooked: + self.__class__.step = self.profile_hook_step(self.__class__.step) # type: ignore[assignment] + self.__class__.step.hooked = True # type: ignore[attr-defined] + + def register_step_pre_hook(self, hook: OptimizerPreHook) -> RemovableHandle: + r"""Register an optimizer step pre hook which will be called before optimizer step. + + It should have the following signature:: + + hook(optimizer, args, kwargs) -> None or modified args and kwargs + + The ``optimizer`` argument is the optimizer instance being used. If + args and kwargs are modified by the pre-hook, then the transformed + values are returned as a tuple containing the new_args and new_kwargs. + + Args: + hook (Callable): The user defined hook to be registered. + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + a handle that can be used to remove the added hook by calling + ``handle.remove()`` + """ + handle = hooks.RemovableHandle(self._optimizer_step_pre_hooks) + self._optimizer_step_pre_hooks[handle.id] = hook + return handle + + def register_step_post_hook(self, hook: OptimizerPostHook) -> RemovableHandle: + r"""Register an optimizer step post hook which will be called after optimizer step. + + It should have the following signature:: + + hook(optimizer, args, kwargs) -> None + + The ``optimizer`` argument is the optimizer instance being used. + + Args: + hook (Callable): The user defined hook to be registered. + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + a handle that can be used to remove the added hook by calling + ``handle.remove()`` + """ + handle = hooks.RemovableHandle(self._optimizer_step_post_hooks) + self._optimizer_step_post_hooks[handle.id] = hook + return handle + + def register_state_dict_pre_hook( + self, hook: Callable[["Optimizer"], None], prepend: bool = False + ) -> RemovableHandle: # noqa: D101 + r"""Register a state dict pre-hook which will be called before :meth:`~torch.optim.Optimizer.state_dict` is called. + + It should have the following signature:: + + hook(optimizer) -> None + + The ``optimizer`` argument is the optimizer instance being used. + The hook will be called with argument ``self`` before calling ``state_dict`` on ``self``. + The registered hook can be used to perform pre-processing before the ``state_dict`` + call is made. + + Args: + hook (Callable): The user defined hook to be registered. + prepend (bool): If True, the provided pre ``hook`` will be fired before + all the already registered pre-hooks on ``state_dict``. Otherwise, + the provided ``hook`` will be fired after all the already registered + pre-hooks. (default: False) + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + a handle that can be used to remove the added hook by calling + ``handle.remove()`` + """ + handle = hooks.RemovableHandle(self._optimizer_state_dict_pre_hooks) + self._optimizer_state_dict_pre_hooks[handle.id] = hook + if prepend: + self._optimizer_state_dict_pre_hooks.move_to_end(handle.id, last=False) + return handle + + def register_state_dict_post_hook( + self, + hook: Callable[["Optimizer", StateDict], StateDict | None], + prepend: bool = False, + ) -> RemovableHandle: + r"""Register a state dict post-hook which will be called after :meth:`~torch.optim.Optimizer.state_dict` is called. + + It should have the following signature:: + + hook(optimizer, state_dict) -> state_dict or None + + The hook will be called with arguments ``self`` and ``state_dict`` after generating + a ``state_dict`` on ``self``. The hook may modify the state_dict inplace or optionally + return a new one. The registered hook can be used to perform post-processing + on the ``state_dict`` before it is returned. + + Args: + hook (Callable): The user defined hook to be registered. + prepend (bool): If True, the provided post ``hook`` will be fired before + all the already registered post-hooks on ``state_dict``. Otherwise, + the provided ``hook`` will be fired after all the already registered + post-hooks. (default: False) + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + a handle that can be used to remove the added hook by calling + ``handle.remove()`` + """ + handle = hooks.RemovableHandle(self._optimizer_state_dict_post_hooks) + self._optimizer_state_dict_post_hooks[handle.id] = hook + if prepend: + self._optimizer_state_dict_post_hooks.move_to_end(handle.id, last=False) + return handle + + @torch._disable_dynamo + def state_dict(self) -> StateDict: + r"""Return the state of the optimizer as a :class:`dict`. + + It contains two entries: + + * ``state``: a Dict holding current optimization state. Its content + differs between optimizer classes, but some common characteristics + hold. For example, state is saved per parameter, and the parameter + itself is NOT saved. ``state`` is a Dictionary mapping parameter ids + to a Dict with state corresponding to each parameter. + * ``param_groups``: a List containing all parameter groups where each + parameter group is a Dict. Each parameter group contains metadata + specific to the optimizer, such as learning rate and weight decay, + as well as a List of parameter IDs of the parameters in the group. + If a param group was initialized with ``named_parameters()`` the names + content will also be saved in the state dict. + + NOTE: The parameter IDs may look like indices but they are just IDs + associating state with param_group. When loading from a state_dict, + the optimizer will zip the param_group ``params`` (int IDs) and the + optimizer ``param_groups`` (actual ``nn.Parameter`` s) in order to + match state WITHOUT additional verification. + + A returned state dict might look something like: + + .. code-block:: text + + { + 'state': { + 0: {'momentum_buffer': tensor(...), ...}, + 1: {'momentum_buffer': tensor(...), ...}, + 2: {'momentum_buffer': tensor(...), ...}, + 3: {'momentum_buffer': tensor(...), ...} + }, + 'param_groups': [ + { + 'lr': 0.01, + 'weight_decay': 0, + ... + 'params': [0] + 'param_names' ['param0'] (optional) + }, + { + 'lr': 0.001, + 'weight_decay': 0.5, + ... + 'params': [1, 2, 3] + 'param_names': ['param1', 'layer.weight', 'layer.bias'] (optional) + } + ] + } + + """ + for pre_hook in self._optimizer_state_dict_pre_hooks.values(): + pre_hook(self) + + # Save order indices instead of Tensors + param_mappings: dict[int, int] = {} + start_index = 0 + + def pack_group(group: dict[str, Any]) -> dict[str, Any]: + nonlocal start_index + packed = {k: v for k, v in group.items() if k != "params"} + param_mappings.update( + { + id(p): i + for i, p in enumerate(group["params"], start_index) + if id(p) not in param_mappings + } + ) + packed["params"] = [param_mappings[id(p)] for p in group["params"]] + start_index += len(packed["params"]) + return packed + + param_groups = [pack_group(g) for g in self.param_groups] + # Remap state to use order indices as keys + packed_state = { + (param_mappings[id(k)] if isinstance(k, torch.Tensor) else k): v + for k, v in self.state.items() + } + + state_dict = { + "state": packed_state, + "param_groups": param_groups, + } + + for post_hook in self._optimizer_state_dict_post_hooks.values(): + hook_result = post_hook(self, state_dict) + if hook_result is not None: + state_dict = hook_result + return state_dict + + @staticmethod + def _process_value_according_to_param_policy( + param: torch.Tensor, + value: torch.Tensor, + param_id: int, + param_groups: list[dict[Any, Any]], + key: Hashable = None, + ) -> torch.Tensor: + # Floating-point types are a bit special here. They are the only ones + # that are assumed to always match the type of params. + # Make sure state['step'] is not casted https://github.com/pytorch/pytorch/issues/74424 + # UNLESS fused or capturable, see note [special device hosting for step] + fused = False + capturable = False + if param_groups is None: + raise AssertionError("Expected param_groups to be set") + for pg in param_groups: + if param_id in pg["params"]: + fused = pg.get("fused", False) + capturable = pg.get("capturable", False) + break + if key == "step": + if capturable or fused: + return value.to(dtype=torch.float32, device=param.device) + else: + return value + else: + if param.is_floating_point(): + return value.to(dtype=param.dtype, device=param.device) + else: + return value.to(device=param.device) + + def register_load_state_dict_pre_hook( + self, + hook: Callable[["Optimizer", StateDict], StateDict | None], + prepend: bool = False, + ) -> RemovableHandle: # noqa: D205 D400 + r"""Register a load_state_dict pre-hook which will be called before + :meth:`~torch.optim.Optimizer.load_state_dict` is called. It should have the + following signature:: + + hook(optimizer, state_dict) -> state_dict or None + + The ``optimizer`` argument is the optimizer instance being used and the + ``state_dict`` argument is a shallow copy of the ``state_dict`` the user + passed in to ``load_state_dict``. The hook may modify the state_dict inplace + or optionally return a new one. If a state_dict is returned, it will be used + to be loaded into the optimizer. + + The hook will be called with argument ``self`` and ``state_dict`` before + calling ``load_state_dict`` on ``self``. The registered hook can be used to + perform pre-processing before the ``load_state_dict`` call is made. + + Args: + hook (Callable): The user defined hook to be registered. + prepend (bool): If True, the provided pre ``hook`` will be fired before + all the already registered pre-hooks on ``load_state_dict``. Otherwise, + the provided ``hook`` will be fired after all the already registered + pre-hooks. (default: False) + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + a handle that can be used to remove the added hook by calling + ``handle.remove()`` + """ + handle = hooks.RemovableHandle(self._optimizer_load_state_dict_pre_hooks) + self._optimizer_load_state_dict_pre_hooks[handle.id] = hook + if prepend: + self._optimizer_load_state_dict_pre_hooks.move_to_end(handle.id, last=False) + return handle + + def register_load_state_dict_post_hook( + self, hook: Callable[["Optimizer"], None], prepend: bool = False + ) -> RemovableHandle: # noqa: D205 D400 + r"""Register a load_state_dict post-hook which will be called after + :meth:`~torch.optim.Optimizer.load_state_dict` is called. It should have the + following signature:: + + hook(optimizer) -> None + + The ``optimizer`` argument is the optimizer instance being used. + + The hook will be called with argument ``self`` after calling + ``load_state_dict`` on ``self``. The registered hook can be used to + perform post-processing after ``load_state_dict`` has loaded the + ``state_dict``. + + Args: + hook (Callable): The user defined hook to be registered. + prepend (bool): If True, the provided post ``hook`` will be fired before + all the already registered post-hooks on ``load_state_dict``. Otherwise, + the provided ``hook`` will be fired after all the already registered + post-hooks. (default: False) + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + a handle that can be used to remove the added hook by calling + ``handle.remove()`` + """ + handle = hooks.RemovableHandle(self._optimizer_load_state_dict_post_hooks) + self._optimizer_load_state_dict_post_hooks[handle.id] = hook + if prepend: + self._optimizer_load_state_dict_post_hooks.move_to_end( + handle.id, last=False + ) # type: ignore[attr-defined] + return handle + + @torch._disable_dynamo + def load_state_dict(self, state_dict: StateDict) -> None: + r"""Load the optimizer state. + + Args: + state_dict (dict): optimizer state. Should be an object returned + from a call to :meth:`state_dict`. + + .. warning:: + Make sure this method is called after initializing :class:`torch.optim.lr_scheduler.LRScheduler`, + as calling it beforehand will overwrite the loaded learning rates. + + .. note:: + The names of the parameters (if they exist under the "param_names" key of each param group + in :meth:`state_dict`) will not affect the loading process. + To use the parameters' names for custom cases (such as when the parameters in the loaded state dict + differ from those initialized in the optimizer), + a custom ``register_load_state_dict_pre_hook`` should be implemented to adapt the loaded dict + accordingly. + If ``param_names`` exist in loaded state dict ``param_groups`` they will be saved and override + the current names, if present, in the optimizer state. If they do not exist in loaded state dict, + the optimizer ``param_names`` will remain unchanged. + + Example: + >>> # xdoctest: +SKIP + >>> model = torch.nn.Linear(10, 10) + >>> optim = torch.optim.SGD(model.parameters(), lr=3e-4) + >>> scheduler1 = torch.optim.lr_scheduler.LinearLR( + ... optim, + ... start_factor=0.1, + ... end_factor=1, + ... total_iters=20, + ... ) + >>> scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR( + ... optim, + ... T_max=80, + ... eta_min=3e-5, + ... ) + >>> lr = torch.optim.lr_scheduler.SequentialLR( + ... optim, + ... schedulers=[scheduler1, scheduler2], + ... milestones=[20], + ... ) + >>> lr.load_state_dict(torch.load("./save_seq.pt")) + >>> # now load the optimizer checkpoint after loading the LRScheduler + >>> optim.load_state_dict(torch.load("./save_optim.pt")) + + """ + # shallow copy, to be consistent with module API + state_dict = state_dict.copy() + + for pre_hook in self._optimizer_load_state_dict_pre_hooks.values(): + hook_result = pre_hook(self, state_dict) + if hook_result is not None: + state_dict = hook_result + + # Validate the state_dict + groups = self.param_groups + + # Deepcopy as we write into saved_groups later to update state + saved_groups = deepcopy(state_dict["param_groups"]) + + if len(groups) != len(saved_groups): + raise ValueError( + "loaded state dict has a different number of parameter groups" + ) + param_lens = (len(g["params"]) for g in groups) + saved_lens = (len(g["params"]) for g in saved_groups) + if any( + p_len != s_len for p_len, s_len in zip(param_lens, saved_lens, strict=True) + ): + raise ValueError( + "loaded state dict contains a parameter group " + "that doesn't match the size of optimizer's group" + ) + + # Update the state + id_map = dict( + zip( + chain.from_iterable(g["params"] for g in saved_groups), + chain.from_iterable(g["params"] for g in groups), + strict=True, + ) + ) + + def _cast(param, value, param_id=None, param_groups=None, key=None): + r"""Make a deep copy of value, casting all tensors to device of param.""" + if isinstance(value, torch.Tensor): + return Optimizer._process_value_according_to_param_policy( + param, + value, + # pyrefly: ignore [bad-argument-type] + param_id, + # pyrefly: ignore [bad-argument-type] + param_groups, + key, + ) + elif isinstance(value, dict): + return { + k: _cast( + param, v, param_id=param_id, param_groups=param_groups, key=k + ) + for k, v in value.items() + } + elif isinstance(value, Iterable): + # pyrefly: ignore [bad-instantiation] + return type(value)( + # pyrefly: ignore [bad-argument-count] + _cast(param, v, param_id=param_id, param_groups=param_groups) + for v in value + ) # type: ignore[call-arg] + else: + return value + + # Copy state assigned to params (and cast tensors to appropriate types). + # State that is not assigned to params is copied as is (needed for + # backward compatibility). + state: defaultdict[torch.Tensor, dict[Any, Any]] = defaultdict(dict) + for k, v in state_dict["state"].items(): + if k in id_map: + param = id_map[k] + state[param] = _cast( + param, v, param_id=k, param_groups=state_dict["param_groups"] + ) + else: + state[k] = v + + # Update parameter groups, setting their 'params' value + def update_group( + group: dict[str, Any], new_group: dict[str, Any] + ) -> dict[str, Any]: + new_group["params"] = group["params"] + if "param_names" in group and "param_names" not in new_group: + new_group["param_names"] = group["param_names"] + return new_group + + param_groups = [ + update_group(g, ng) for g, ng in zip(groups, saved_groups, strict=True) + ] + self.__setstate__({"state": state, "param_groups": param_groups}) + + for post_hook in self._optimizer_load_state_dict_post_hooks.values(): + post_hook(self) + + @torch._disable_dynamo + def zero_grad(self, set_to_none: bool = True) -> None: + r"""Reset the gradients of all optimized :class:`torch.Tensor` s. + + Args: + set_to_none (bool, optional): Instead of setting to zero, set the grads to None. Default: ``True`` + + This will in general have lower memory footprint, and can modestly improve performance. + However, it changes certain behaviors. For example: + + 1. When the user tries to access a gradient and perform manual ops on it, + a None attribute or a Tensor full of 0s will behave differently. + 2. If the user requests ``zero_grad(set_to_none=True)`` followed by a backward pass, ``.grad``\ s + are guaranteed to be None for params that did not receive a gradient. + 3. ``torch.optim`` optimizers have a different behavior if the gradient is 0 or None + (in one case it does the step with a gradient of 0 and in the other it skips + the step altogether). + """ + foreach = self.defaults.get("foreach", False) or self.defaults.get( + "fused", False + ) + + if not hasattr(self, "_zero_grad_profile_name"): + self._patch_step_function() + + per_device_and_dtype_grads: ( + defaultdict[torch.device, defaultdict[torch.dtype, list[torch.Tensor]]] + | None + ) + if foreach: + per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list)) + else: + per_device_and_dtype_grads = None + + with torch.autograd.profiler.record_function(self._zero_grad_profile_name): + for group in self.param_groups: + for p in group["params"]: + if p.grad is not None: + if set_to_none: + p.grad = None + else: + if p.grad.grad_fn is not None: + p.grad.detach_() + else: + p.grad.requires_grad_(False) + if not foreach or p.grad.is_sparse: + p.grad.zero_() + else: + if per_device_and_dtype_grads is None: + raise AssertionError( + "Expected per_device_and_dtype_grads to be set" + ) + per_device_and_dtype_grads[p.grad.device][ + p.grad.dtype + ].append(p.grad) + if foreach: + if per_device_and_dtype_grads is None: + raise AssertionError( + "Expected per_device_and_dtype_grads to be set" + ) + for per_dtype_grads in per_device_and_dtype_grads.values(): + for grads in per_dtype_grads.values(): + torch._foreach_zero_(grads) + + @overload + def step(self, closure: None = None) -> None: ... + + @overload + def step(self, closure: Callable[[], float]) -> float: ... + + def step(self, closure: Callable[[], float] | None = None) -> float | None: + r"""Perform a single optimization step to update parameter. + + Args: + closure (Callable): A closure that reevaluates the model and + returns the loss. Optional for most optimizers. + """ + raise NotImplementedError + + @torch._disable_dynamo + def add_param_group(self, param_group: dict[str, Any]) -> None: + r"""Add a param group to the :class:`Optimizer` s `param_groups`. + + This can be useful when fine tuning a pre-trained network as frozen layers can be made + trainable and added to the :class:`Optimizer` as training progresses. + + Args: + param_group (dict): Specifies what Tensors should be optimized along with group + specific optimization options. + """ + if not isinstance(param_group, dict): + raise TypeError(f"param_group must be a dict, but got {type(param_group)}") + + params = param_group["params"] + if isinstance(params, torch.Tensor): + param_group["params"] = [params] + elif isinstance(params, set): + raise TypeError( + "optimizer parameters need to be organized in ordered collections, but " + "the ordering of tensors in sets will change between runs. Please use a list instead." + ) + else: + param_group["params"] = list(params) + + extracted_param_tensors = [] + extracted_param_names = [] + for param in param_group["params"]: + if isinstance(param, tuple): + param_name = param[0] + extracted_param_names.append(param_name) + extracted_param_tensors.append(param[1]) + else: + extracted_param_tensors.append(param) + + param_group["params"] = extracted_param_tensors + if len(extracted_param_names) != 0: + if len(extracted_param_names) == len(extracted_param_tensors): + param_group["param_names"] = extracted_param_names + else: + raise ValueError( + "all optimizer params should be with/without names. Some param names are missing" + ) + + for param in param_group["params"]: + if not isinstance(param, torch.Tensor): + raise TypeError( + "optimizer can only optimize Tensors, " + "but one of the params is " + torch.typename(param) + ) + if not self.defaults.get("differentiable", None) and not ( + param.is_leaf or param.retains_grad + ): + raise ValueError("can't optimize a non-leaf Tensor") + + for name, default in self.defaults.items(): + if default is required and name not in param_group: + raise ValueError( + f"parameter group didn't specify a value of required optimization parameter {name}" + ) + else: + param_group.setdefault(name, default) + + params = param_group["params"] + if len(params) != len(set(params)): + warnings.warn( + "optimizer contains a parameter group with duplicate parameters; " + "in future, this will cause an error; " + "see github.com/pytorch/pytorch/issues/40967 for more information", + stacklevel=3, + ) + + param_set: set[torch.Tensor] = set() + for group in self.param_groups: + param_set.update(set(group["params"])) + if ("param_names" in param_group) != ("param_names" in group): + current_group_txt = ( + "with names" if "param_names" in param_group else "without names" + ) + raise ValueError( + "all optimizer param groups should be with/without names. " + f"cannot add param group {current_group_txt} to the optimizer" + ) + + if not param_set.isdisjoint(set(param_group["params"])): + raise ValueError("some parameters appear in more than one parameter group") + + self.param_groups.append(param_group) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/radam.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/radam.py new file mode 100644 index 0000000000000000000000000000000000000000..d414b9c7edbeeeb5c0cab2e54a93415ecd3f97b3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/radam.py @@ -0,0 +1,626 @@ +# mypy: allow-untyped-defs +r"""Implementation for the RAdam algorithm.""" + +from typing import cast + +import torch +from torch import Tensor + +from .optimizer import ( + _capturable_doc, + _default_to_fused_or_foreach, + _differentiable_doc, + _disable_dynamo_if_unsupported, + _foreach_doc, + _get_capturable_supported_devices, + _get_scalar_dtype, + _get_value, + _maximize_doc, + _params_doc, + _to_scalar, + _use_grad_for_differentiable, + _view_as_real, + Optimizer, + ParamsT, +) + + +__all__ = ["RAdam", "radam"] + + +class RAdam(Optimizer): # noqa: D101 + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-3, + betas: tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + weight_decay: float = 0, + decoupled_weight_decay: bool = False, + *, + foreach: bool | None = None, + maximize: bool = False, + capturable: bool = False, + differentiable: bool = False, + ) -> None: # noqa: D107 + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + + defaults = { + "lr": lr, + "betas": betas, + "eps": eps, + "weight_decay": weight_decay, + "maximize": maximize, + "foreach": foreach, + "capturable": capturable, + "decoupled_weight_decay": decoupled_weight_decay, + "differentiable": differentiable, + } + super().__init__(params, defaults) + + def __setstate__(self, state): # noqa: D105 + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("foreach", None) + group.setdefault("maximize", False) + group.setdefault("differentiable", False) + group.setdefault("decoupled_weight_decay", False) + group.setdefault("capturable", False) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = ( + torch.tensor( + step_val, dtype=_get_scalar_dtype(), device=p.device + ) + if group["capturable"] + else torch.tensor(step_val, dtype=_get_scalar_dtype()) + ) + + def _init_group( + self, group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps + ): + has_complex = False + for p in group["params"]: + if p.grad is not None: + has_complex |= torch.is_complex(p) + params_with_grad.append(p) + if p.grad.is_sparse: + raise RuntimeError("RAdam does not support sparse gradients") + grads.append(p.grad) + + state = self.state[p] + # Lazy state initialization + if len(state) == 0: + state["step"] = ( + torch.zeros((), dtype=_get_scalar_dtype(), device=p.device) + if group["capturable"] + else torch.tensor(0.0, dtype=_get_scalar_dtype()) + ) + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + exp_avgs.append(state["exp_avg"]) + exp_avg_sqs.append(state["exp_avg_sq"]) + state_steps.append(state["step"]) + + return has_complex + + @_use_grad_for_differentiable + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._accelerator_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + exp_avgs: list[Tensor] = [] + exp_avg_sqs: list[Tensor] = [] + state_steps: list[Tensor] = [] + beta1, beta2 = cast(tuple[float, float], group["betas"]) + + has_complex = self._init_group( + group, params_with_grad, grads, exp_avgs, exp_avg_sqs, state_steps + ) + + radam( + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + state_steps, + beta1=beta1, + beta2=beta2, + lr=group["lr"], + weight_decay=group["weight_decay"], + eps=group["eps"], + maximize=group["maximize"], + foreach=group["foreach"], + capturable=group["capturable"], + differentiable=group["differentiable"], + decoupled_weight_decay=group["decoupled_weight_decay"], + has_complex=has_complex, + ) + + return loss + + +RAdam.__doc__ = ( + r"""Implements RAdam algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{ (lr)}, \: \beta_1, \beta_2 + \text{ (betas)}, \: \theta_0 \text{ (params)}, \:f(\theta) \text{ (objective)}, \: + \lambda \text{ (weightdecay)}, \:\textit{maximize} \\ + &\hspace{13mm} \epsilon \text{ (epsilon)}, \textit{decoupled\_weight\_decay} \\ + &\textbf{initialize} : m_0 \leftarrow 0 \text{ ( first moment)}, + v_0 \leftarrow 0 \text{ ( second moment)}, \\ + &\hspace{18mm} \rho_{\infty} \leftarrow 2/(1-\beta_2) -1 \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + &\hspace{6mm}\textbf{if} \: \textit{maximize}: \\ + &\hspace{12mm}g_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{6mm}\textbf{else} \\ + &\hspace{12mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{6mm} \theta_t \leftarrow \theta_{t-1} \\ + &\hspace{6mm} \textbf{if} \: \lambda \neq 0 \\ + &\hspace{12mm}\textbf{if} \: \textit{decoupled\_weight\_decay} \\ + &\hspace{18mm} \theta_t \leftarrow \theta_{t} - \gamma \lambda \theta_{t} \\ + &\hspace{12mm}\textbf{else} \\ + &\hspace{18mm} g_t \leftarrow g_t + \lambda \theta_{t} \\ + &\hspace{6mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ + &\hspace{6mm}v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) g^2_t \\ + &\hspace{6mm}\widehat{m_t} \leftarrow m_t/\big(1-\beta_1^t \big) \\ + &\hspace{6mm}\rho_t \leftarrow \rho_{\infty} - + 2 t \beta^t_2 /\big(1-\beta_2^t \big) \\[0.1.ex] + &\hspace{6mm}\textbf{if} \: \rho_t > 5 \\ + &\hspace{12mm} l_t \leftarrow \frac{\sqrt{ (1-\beta^t_2) }}{ \sqrt{v_t} +\epsilon } \\ + &\hspace{12mm} r_t \leftarrow + \sqrt{\frac{(\rho_t-4)(\rho_t-2)\rho_{\infty}}{(\rho_{\infty}-4)(\rho_{\infty}-2) \rho_t}} \\ + &\hspace{12mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t} r_t l_t \\ + &\hspace{6mm}\textbf{else} \\ + &\hspace{12mm}\theta_t \leftarrow \theta_t - \gamma \widehat{m_t} \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to `On the variance of the adaptive learning rate and beyond`_. + + This implementation provides an option to use either the original weight_decay implementation as in Adam + (where the weight_decay is applied to the gradient) or the one from AdamW (where weight_decay is applied + to the weight) through the decoupled_weight_decay option. When decoupled_weight_decay is set to False + (default), it uses the original Adam style weight decay, otherwise, it uses the AdamW style which + corresponds more closely to the `author's implementation`_ in the RAdam paper. Further information + about decoupled weight decay can be found in `Decoupled Weight Decay Regularization`_. + + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + decoupled_weight_decay (bool, optional): whether to decouple the weight + decay as in AdamW to obtain RAdamW. If True, the algorithm does not + accumulate weight decay in the momentum nor variance. (default: False) + {_foreach_doc} + {_maximize_doc} + {_capturable_doc} + {_differentiable_doc} + + .. _On the variance of the adaptive learning rate and beyond: + https://arxiv.org/abs/1908.03265 + .. _author's implementation: + https://github.com/LiyuanLucasLiu/RAdam + .. _Decoupled Weight Decay Regularization: + https://arxiv.org/abs/1711.05101 + + """ +) + + +def _single_tensor_radam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + state_steps: list[Tensor], + *, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + eps: float, + decoupled_weight_decay: bool, + differentiable: bool, + maximize: bool, + capturable: bool, + has_complex: bool, +) -> None: + if not torch.jit.is_scripting(): + lr = _to_scalar(lr) + + for i, param in enumerate(params): + grad = grads[i] if not maximize else -grads[i] + exp_avg = exp_avgs[i] + exp_avg_sq = exp_avg_sqs[i] + step_t = state_steps[i] + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices() + if not ( + param.device.type == step_t.device.type + and param.device.type in capturable_supported_devices + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + if torch.is_complex(param): + param = torch.view_as_real(param) + grad = torch.view_as_real(grad) + exp_avg = torch.view_as_real(exp_avg) + exp_avg_sq = torch.view_as_real(exp_avg_sq) + + # update step + step_t += 1 + step = step_t if capturable else _get_value(step_t) + + if weight_decay != 0: + if decoupled_weight_decay: + param.mul_(1 - lr * weight_decay) + else: + grad = grad.add(param, alpha=weight_decay) + + # Decay the first and second moment running average coefficient + exp_avg.lerp_(grad, 1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) + + bias_correction1 = 1 - beta1**step + bias_correction2 = 1 - beta2**step + + # correcting bias for the first moving moment + bias_corrected_exp_avg = exp_avg / bias_correction1 + + # maximum length of the approximated SMA + rho_inf = 2 / (1 - beta2) - 1 + # compute the length of the approximated SMA + rho_t = rho_inf - 2 * step * (beta2**step) / bias_correction2 + + def _compute_rect(): + return ( + (rho_t - 4) + * (rho_t - 2) + * rho_inf + / ((rho_inf - 4) * (rho_inf - 2) * rho_t) + ) ** 0.5 + + def _compute_adaptive_lr(): + exp_avg_sq_sqrt = exp_avg_sq.sqrt() + if differentiable: + exp_avg_sq_sqrt = exp_avg_sq_sqrt.add(eps) + else: + exp_avg_sq_sqrt = exp_avg_sq_sqrt.add_(eps) + + return (bias_correction2**0.5) / exp_avg_sq_sqrt + + # Compute the variance rectification term and update parameters accordingly + if capturable: + update = torch.where( + rho_t > 5.0, _compute_rect() * _compute_adaptive_lr(), 1.0 + ) + param.add_(bias_corrected_exp_avg * lr * update, alpha=-1.0) + else: + if rho_t > 5.0: + param.add_( + bias_corrected_exp_avg + * lr + * _compute_adaptive_lr() + * _compute_rect(), + alpha=-1.0, + ) + else: + param.add_(bias_corrected_exp_avg * lr, alpha=-1.0) + + +def _multi_tensor_radam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + state_steps: list[Tensor], + *, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + eps: float, + decoupled_weight_decay: bool, + differentiable: bool, + maximize: bool, + capturable: bool, + has_complex: bool, +) -> None: + if len(params) == 0: + return + + if differentiable: + raise AssertionError("_foreach ops don't support autograd") + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices( + supports_xla=False + ) + if not all( + p.device.type == step.device.type + and p.device.type in capturable_supported_devices + for p, step in zip(params, state_steps, strict=True) + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + lr = _to_scalar(lr) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, exp_avgs, exp_avg_sqs, state_steps] # type: ignore[list-item] + ) + for ( + grouped_params_, + grouped_grads_, + grouped_exp_avgs_, + grouped_exp_avg_sqs_, + grouped_state_steps_, + ), _ in grouped_tensors.values(): + grouped_params = cast(list[Tensor], grouped_params_) + grouped_grads = cast(list[Tensor], grouped_grads_) + grouped_exp_avgs = cast(list[Tensor], grouped_exp_avgs_) + grouped_exp_avg_sqs = cast(list[Tensor], grouped_exp_avg_sqs_) + grouped_state_steps = cast(list[Tensor], grouped_state_steps_) + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and grouped_state_steps[0].is_cpu: + torch._foreach_add_( + grouped_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(grouped_state_steps, 1) + + if has_complex: + _view_as_real( + grouped_params, grouped_grads, grouped_exp_avgs, grouped_exp_avg_sqs + ) + + if maximize: + grouped_grads = torch._foreach_neg(grouped_grads) # type: ignore[assignment] + + # maximum length of the approximated SMA + rho_inf = 2 / (1 - beta2) - 1 + # compute the length of the approximated SMA + bias_correction1: tuple[Tensor, ...] | list[Tensor] + bias_correction2: tuple[Tensor, ...] | list[Tensor] + rho_t_list: tuple[Tensor, ...] | list[Tensor] + if capturable: + bias_correction1 = torch._foreach_pow(beta2, grouped_state_steps) + torch._foreach_neg_(bias_correction1) + torch._foreach_add_(bias_correction1, 1) + bias_correction2 = torch._foreach_pow(beta2, grouped_state_steps) + torch._foreach_mul_(bias_correction2, grouped_state_steps) + torch._foreach_mul_(bias_correction2, 2) + torch._foreach_div_(bias_correction2, bias_correction1) + torch._foreach_neg_(bias_correction2) + torch._foreach_add_(bias_correction2, rho_inf) + rho_t_list = bias_correction2 + else: + rho_t_list = [ + rho_inf + - 2 + * _get_value(step) + * (beta2 ** _get_value(step)) + / (1 - beta2 ** _get_value(step)) + for step in grouped_state_steps + ] + + if weight_decay != 0: + if decoupled_weight_decay: + torch._foreach_mul_(grouped_params, 1 - lr * weight_decay) + else: + # Reuse the intermediate memory (grouped_grads) already allocated for maximize + if maximize: + torch._foreach_add_( + grouped_grads, grouped_params, alpha=weight_decay + ) + else: + grouped_grads = torch._foreach_add( # type: ignore[assignment] + grouped_grads, grouped_params, alpha=weight_decay + ) + + # Decay the first and second moment running average coefficient + torch._foreach_lerp_(grouped_exp_avgs, grouped_grads, 1 - beta1) + + torch._foreach_mul_(grouped_exp_avg_sqs, beta2) + torch._foreach_addcmul_( + grouped_exp_avg_sqs, grouped_grads, grouped_grads, 1 - beta2 + ) + + # Delete the local intermediate since it won't be used anymore to save on peak memory + del grouped_grads + + if capturable: + num = torch._foreach_sub(rho_t_list, 4) + sub2 = torch._foreach_sub(rho_t_list, 2) + torch._foreach_mul_(num, sub2) + del sub2 + torch._foreach_mul_(num, rho_inf) + rho_inf = (rho_inf - 4) * (rho_inf - 2) + denom = torch._foreach_mul(rho_t_list, rho_inf) + torch._foreach_div_(num, denom) + del denom + torch._foreach_sqrt_(num) + + # TODO(mlazos): we should try and get a foreach_where op https://github.com/pytorch/pytorch/issues/117884 + rect = [ + torch.where(rho_t > 5.0, n, 0.0) + for n, rho_t in zip(num, rho_t_list, strict=True) + ] + del num + del rho_t_list + unrect_step_size = [torch.where(rect > 0, 0.0, 1.0) for rect in rect] + torch._foreach_mul_(unrect_step_size, lr) + + bias_correction1 = torch._foreach_pow(beta1, grouped_state_steps) + torch._foreach_neg_(bias_correction1) + torch._foreach_add_(bias_correction1, 1) + + torch._foreach_div_(unrect_step_size, bias_correction1) + torch._foreach_neg_(unrect_step_size) + + bias_correction2 = torch._foreach_pow(beta2, grouped_state_steps) + torch._foreach_neg_(bias_correction2) + torch._foreach_add_(bias_correction2, 1) + torch._foreach_sqrt_(bias_correction2) + torch._foreach_mul_(bias_correction2, lr) + torch._foreach_mul_(bias_correction2, rect) + del rect + torch._foreach_neg_(bias_correction2) + torch._foreach_div_(bias_correction2, bias_correction1) + del bias_correction1 + else: + rect = [ + ( # type: ignore[misc] + (rho_t - 4) # type: ignore[arg-type] + * (rho_t - 2) + * rho_inf + / ((rho_inf - 4) * (rho_inf - 2) * rho_t) + ) + ** 0.5 + if rho_t > 5 + else 0 + for rho_t in rho_t_list + ] + unrectified = [0 if rect > 0 else 1.0 for rect in rect] + + bias_correction1 = [ + 1 - beta1 ** _get_value(step) for step in grouped_state_steps + ] + unrect_step_size = [ + (lr * rect / bc) * -1 + for rect, bc in zip(unrectified, bias_correction1, strict=True) + ] + bias_correction2 = [ + ((1 - beta2 ** _get_value(step)) ** 0.5) * (lr * rect / bc) * -1 + for step, rect, bc in zip( + grouped_state_steps, rect, bias_correction1, strict=True + ) + ] + + buffer = torch._foreach_sqrt(grouped_exp_avg_sqs) + torch._foreach_add_(buffer, eps) + torch._foreach_div_(buffer, bias_correction2) + torch._foreach_reciprocal_(buffer) + torch._foreach_add_(buffer, unrect_step_size) + + # Here, buffer = sqrt(1 - beta2^t) * rect_step_size / (sqrt(v) + eps) + unrect_step_size + torch._foreach_addcmul_(grouped_params, grouped_exp_avgs, buffer) + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_radam) +def radam( + params: list[Tensor], + grads: list[Tensor], + exp_avgs: list[Tensor], + exp_avg_sqs: list[Tensor], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + decoupled_weight_decay: bool = False, + foreach: bool | None = None, + differentiable: bool = False, + capturable: bool = False, + has_complex: bool = False, + maximize: bool = False, + *, + beta1: float, + beta2: float, + lr: float, + weight_decay: float, + eps: float, +) -> None: + r"""Functional API that performs RAdam algorithm computation. + + See :class:`~torch.optim.RAdam` for details. + """ + if not all(isinstance(t, torch.Tensor) for t in state_steps): + raise RuntimeError( + "API has changed, `state_steps` argument must contain a list of singleton tensors" + ) + + if foreach is None: + _, foreach = _default_to_fused_or_foreach( + params, differentiable, use_fused=False + ) + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + + if foreach and not torch.jit.is_scripting(): + func = _multi_tensor_radam + else: + func = _single_tensor_radam + + func( + params, + grads, + exp_avgs, + exp_avg_sqs, + state_steps, + beta1=beta1, + beta2=beta2, + lr=lr, + weight_decay=weight_decay, + eps=eps, + maximize=maximize, + decoupled_weight_decay=decoupled_weight_decay, + differentiable=differentiable, + capturable=capturable, + has_complex=has_complex, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/rmsprop.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/rmsprop.py new file mode 100644 index 0000000000000000000000000000000000000000..eefce02fac6fe375d2f9b85c57e28834fd46edd7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/rmsprop.py @@ -0,0 +1,542 @@ +# mypy: allow-untyped-defs +r"""Implementation for the RMSprop algorithm.""" + +from typing import cast + +import torch +from torch import Tensor + +from .optimizer import ( + _capturable_doc, + _default_to_fused_or_foreach, + _differentiable_doc, + _disable_dynamo_if_unsupported, + _foreach_doc, + _get_capturable_supported_devices, + _get_scalar_dtype, + _maximize_doc, + _params_doc, + _to_scalar, + _use_grad_for_differentiable, + _view_as_real, + Optimizer, + ParamsT, +) + + +__all__ = ["RMSprop", "rmsprop"] + + +class RMSprop(Optimizer): # noqa: D101 + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-2, + alpha: float = 0.99, + eps: float = 1e-8, + weight_decay: float = 0, + momentum: float = 0, + centered: bool = False, + capturable: bool = False, + foreach: bool | None = None, + maximize: bool = False, + differentiable: bool = False, + ) -> None: # noqa: D107 + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 <= eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= momentum: + raise ValueError(f"Invalid momentum value: {momentum}") + if not 0.0 <= weight_decay: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + if not 0.0 <= alpha: + raise ValueError(f"Invalid alpha value: {alpha}") + + defaults = { + "lr": lr, + "momentum": momentum, + "alpha": alpha, + "eps": eps, + "centered": centered, + "weight_decay": weight_decay, + "capturable": capturable, + "foreach": foreach, + "maximize": maximize, + "differentiable": differentiable, + } + super().__init__(params, defaults) + + def __setstate__(self, state): # noqa: D105 + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("momentum", 0) + group.setdefault("centered", False) + group.setdefault("foreach", None) + group.setdefault("maximize", False) + group.setdefault("differentiable", False) + group.setdefault("capturable", False) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = ( + torch.tensor( + step_val, dtype=_get_scalar_dtype(), device=p.device + ) + if group["capturable"] + else torch.tensor(step_val, dtype=_get_scalar_dtype()) + ) + + def _init_group( + self, + group, + params_with_grad, + grads, + square_avgs, + momentum_buffer_list, + grad_avgs, + state_steps, + ): + has_complex = False + for p in group["params"]: + if p.grad is None: + continue + has_complex |= torch.is_complex(p) + params_with_grad.append(p) + + if p.grad.is_sparse: + raise RuntimeError("RMSprop does not support sparse gradients") + grads.append(p.grad) + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = ( + torch.zeros((), dtype=_get_scalar_dtype(), device=p.device) + if group["capturable"] + else torch.zeros((), dtype=_get_scalar_dtype()) + ) + state["square_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + if group["momentum"] > 0: + state["momentum_buffer"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + if group["centered"]: + state["grad_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + square_avgs.append(state["square_avg"]) + state_steps.append(state["step"]) + + if group["momentum"] > 0: + momentum_buffer_list.append(state["momentum_buffer"]) + if group["centered"]: + grad_avgs.append(state["grad_avg"]) + + return has_complex + + @_use_grad_for_differentiable + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._accelerator_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + square_avgs: list[Tensor] = [] + grad_avgs: list[Tensor] = [] + momentum_buffer_list: list[Tensor] = [] + state_steps: list[Tensor] = [] + + has_complex = self._init_group( + group, + params_with_grad, + grads, + square_avgs, + momentum_buffer_list, + grad_avgs, + state_steps, + ) + + rmsprop( + params_with_grad, + grads, + square_avgs, + grad_avgs, + momentum_buffer_list, + state_steps, + lr=group["lr"], + alpha=group["alpha"], + eps=group["eps"], + weight_decay=group["weight_decay"], + momentum=group["momentum"], + centered=group["centered"], + foreach=group["foreach"], + maximize=group["maximize"], + differentiable=group["differentiable"], + capturable=group["capturable"], + has_complex=has_complex, + ) + + return loss + + +RMSprop.__doc__ = ( + r"""Implements RMSprop algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \alpha \text{ (alpha)}, \: \gamma \text{ (lr)}, + \: \theta_0 \text{ (params)}, \: f(\theta) \text{ (objective)} \\ + &\hspace{13mm} \lambda \text{ (weight decay)},\: \mu \text{ (momentum)}, + \: centered, \: \epsilon \text{ (epsilon)} \\ + &\textbf{initialize} : v_0 \leftarrow 0 \text{ (square average)}, \: + \textbf{b}_0 \leftarrow 0 \text{ (buffer)}, \: g^{ave}_0 \leftarrow 0 \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + &\hspace{5mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}if \: \lambda \neq 0 \\ + &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ + &\hspace{5mm}v_t \leftarrow \alpha v_{t-1} + (1 - \alpha) g^2_t + \hspace{8mm} \\ + &\hspace{5mm} \tilde{v_t} \leftarrow v_t \\ + &\hspace{5mm}if \: centered \\ + &\hspace{10mm} g^{ave}_t \leftarrow g^{ave}_{t-1} \alpha + (1-\alpha) g_t \\ + &\hspace{10mm} \tilde{v_t} \leftarrow \tilde{v_t} - \big(g^{ave}_{t} \big)^2 \\ + &\hspace{5mm}if \: \mu > 0 \\ + &\hspace{10mm} \textbf{b}_t\leftarrow \mu \textbf{b}_{t-1} + + g_t/ \big(\sqrt{\tilde{v_t}} + \epsilon \big) \\ + &\hspace{10mm} \theta_t \leftarrow \theta_{t-1} - \gamma \textbf{b}_t \\ + &\hspace{5mm} else \\ + &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - + \gamma g_t/ \big(\sqrt{\tilde{v_t}} + \epsilon \big) \hspace{3mm} \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to + `lecture notes `_ by G. Hinton. + and centered version `Generating Sequences + With Recurrent Neural Networks `_. + The implementation here takes the square root of the gradient average before + adding epsilon (note that TensorFlow interchanges these two operations). The effective + learning rate is thus :math:`\gamma/(\sqrt{v} + \epsilon)` where :math:`\gamma` + is the scheduled learning rate and :math:`v` is the weighted moving average + of the squared gradient. + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 1e-2) + alpha (float, optional): smoothing constant (default: 0.99) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + momentum (float, optional): momentum factor (default: 0) + centered (bool, optional) : if ``True``, compute the centered RMSProp, + the gradient is normalized by an estimation of its variance + {_capturable_doc} + {_foreach_doc} + {_maximize_doc} + {_differentiable_doc} + + """ +) + + +def _single_tensor_rmsprop( + params: list[Tensor], + grads: list[Tensor], + square_avgs: list[Tensor], + grad_avgs: list[Tensor], + momentum_buffer_list: list[Tensor], + state_steps: list[Tensor], + *, + lr: float, + alpha: float, + eps: float, + weight_decay: float, + momentum: float, + centered: bool, + maximize: bool, + differentiable: bool, + capturable: bool, + has_complex: bool, +) -> None: + if not torch.jit.is_scripting(): + lr = _to_scalar(lr) + + for i, param in enumerate(params): + step = state_steps[i] + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices() + if not ( + param.device.type == step.device.type + and param.device.type in capturable_supported_devices + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + grad = grads[i] + grad = grad if not maximize else -grad + square_avg = square_avgs[i] + + step += 1 + + if weight_decay != 0: + grad = grad.add(param, alpha=weight_decay) + + is_complex_param = torch.is_complex(param) + if is_complex_param: + param = torch.view_as_real(param) + grad = torch.view_as_real(grad) + square_avg = torch.view_as_real(square_avg) + + square_avg.mul_(alpha).addcmul_(grad, grad, value=1 - alpha) + + if centered: + grad_avg = grad_avgs[i] + if is_complex_param: + grad_avg = torch.view_as_real(grad_avg) + grad_avg.lerp_(grad, 1 - alpha) + avg = square_avg.addcmul(grad_avg, grad_avg, value=-1).sqrt_() + else: + avg = square_avg.sqrt() + + if differentiable: + avg = avg.add(eps) + else: + avg = avg.add_(eps) + + if momentum > 0: + buf = momentum_buffer_list[i] + if is_complex_param: + buf = torch.view_as_real(buf) + buf.mul_(momentum).addcdiv_(grad, avg) + param.add_(buf, alpha=-lr) + else: + param.addcdiv_(grad, avg, value=-lr) + + +def _multi_tensor_rmsprop( + params: list[Tensor], + grads: list[Tensor], + square_avgs: list[Tensor], + grad_avgs: list[Tensor], + momentum_buffer_list: list[Tensor], + state_steps: list[Tensor], + *, + lr: float, + alpha: float, + eps: float, + weight_decay: float, + momentum: float, + centered: bool, + maximize: bool, + differentiable: bool, + capturable: bool, + has_complex: bool, +) -> None: + if len(params) == 0: + return + + if differentiable: + raise AssertionError("_foreach ops don't support autograd") + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices() + if not all( + p.device.type == step.device.type + and p.device.type in capturable_supported_devices + for p, step in zip(params, state_steps, strict=True) + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + lr = _to_scalar(lr) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, square_avgs, grad_avgs, momentum_buffer_list, state_steps] # type: ignore[list-item] + ) + for ( + ( + grouped_params_, + grouped_grads_, + grouped_square_avgs_, + grouped_grad_avgs_, + grouped_momentum_buffer_list_, + grouped_state_steps_, + ) + ), _ in grouped_tensors.values(): + grouped_params = cast(list[Tensor], grouped_params_) + grouped_grads = cast(list[Tensor], grouped_grads_) + grouped_square_avgs = cast(list[Tensor], grouped_square_avgs_) + grouped_state_steps = cast(list[Tensor], grouped_state_steps_) + + if has_complex: + state_and_grads = [grouped_grads, grouped_square_avgs] + if momentum > 0: + grouped_momentum_buffer_list = cast( + list[Tensor], grouped_momentum_buffer_list_ + ) + state_and_grads.append(grouped_momentum_buffer_list) + if centered: + grouped_grad_avgs = cast(list[Tensor], grouped_grad_avgs_) + state_and_grads.append(grouped_grad_avgs) + _view_as_real(grouped_params, *state_and_grads) + + if maximize: + grouped_grads = torch._foreach_neg(grouped_grads) # type: ignore[assignment] + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and grouped_state_steps[0].is_cpu: + torch._foreach_add_( + grouped_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(grouped_state_steps, 1) + + if weight_decay != 0: + # Reuse the intermediate memory (grouped_grads) already allocated for maximize + if maximize: + torch._foreach_add_(grouped_grads, grouped_params, alpha=weight_decay) + else: + grouped_grads = torch._foreach_add( # type: ignore[assignment] + grouped_grads, grouped_params, alpha=weight_decay + ) + + torch._foreach_mul_(grouped_square_avgs, alpha) + torch._foreach_addcmul_( + grouped_square_avgs, grouped_grads, grouped_grads, value=1 - alpha + ) + + if centered: + grouped_grad_avgs = cast(list[Tensor], grouped_grad_avgs_) + torch._foreach_lerp_(grouped_grad_avgs, grouped_grads, 1 - alpha) + avg = torch._foreach_addcmul( + grouped_square_avgs, grouped_grad_avgs, grouped_grad_avgs, value=-1 + ) + torch._foreach_sqrt_(avg) + torch._foreach_add_(avg, eps) + else: + avg = torch._foreach_sqrt(grouped_square_avgs) + torch._foreach_add_(avg, eps) + + if momentum > 0: + grouped_momentum_buffer_list = cast( + list[Tensor], grouped_momentum_buffer_list_ + ) + torch._foreach_mul_(grouped_momentum_buffer_list, momentum) + torch._foreach_addcdiv_(grouped_momentum_buffer_list, grouped_grads, avg) + # If LR is a tensor, the else branch will internally call item() + # which will cause silent incorrectness if we are capturing + if capturable and isinstance(lr, torch.Tensor): + momentum_lr = torch._foreach_mul(grouped_momentum_buffer_list, -lr) + torch._foreach_add_(grouped_params, momentum_lr) + else: + torch._foreach_add_( + grouped_params, grouped_momentum_buffer_list, alpha=-lr + ) + else: + # If LR is a tensor, the else branch will internally call item() + # which will cause silent incorrectness if we are capturing + if capturable and isinstance(lr, torch.Tensor): + torch._foreach_div_(avg, -lr) + torch._foreach_addcdiv_(grouped_params, grouped_grads, avg) + else: + torch._foreach_addcdiv_(grouped_params, grouped_grads, avg, value=-lr) + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_rmsprop) +def rmsprop( + params: list[Tensor], + grads: list[Tensor], + square_avgs: list[Tensor], + grad_avgs: list[Tensor], + momentum_buffer_list: list[Tensor], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + foreach: bool | None = None, + maximize: bool = False, + differentiable: bool = False, + capturable: bool = False, + has_complex: bool = False, + *, + lr: float, + alpha: float, + eps: float, + weight_decay: float, + momentum: float, + centered: bool, +) -> None: + r"""Functional API that performs rmsprop algorithm computation. + + See :class:`~torch.optim.RMSProp` for details. + """ + # this check is slow during compilation, so we skip it + # if it's strictly needed we can add this check back in dynamo + if not torch.compiler.is_compiling() and not all( + isinstance(t, torch.Tensor) for t in state_steps + ): + raise RuntimeError( + "API has changed, `state_steps` argument must contain a list of singleton tensors" + ) + + if foreach is None: + _, foreach = _default_to_fused_or_foreach( + params, differentiable, use_fused=False + ) + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + + if foreach and not torch.jit.is_scripting(): + func = _multi_tensor_rmsprop + else: + func = _single_tensor_rmsprop + + func( + params, + grads, + square_avgs, + grad_avgs, + momentum_buffer_list, + state_steps, + lr=lr, + alpha=alpha, + eps=eps, + weight_decay=weight_decay, + momentum=momentum, + centered=centered, + maximize=maximize, + capturable=capturable, + differentiable=differentiable, + has_complex=has_complex, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/rprop.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/rprop.py new file mode 100644 index 0000000000000000000000000000000000000000..ea9da851412864decad0568b8be1346b83806261 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/rprop.py @@ -0,0 +1,472 @@ +# mypy: allow-untyped-defs +r"""Implementation for the Resilient backpropagation.""" + +from typing import cast + +import torch +from torch import Tensor + +from .optimizer import ( + _capturable_doc, + _default_to_fused_or_foreach, + _differentiable_doc, + _disable_dynamo_if_unsupported, + _foreach_doc, + _get_capturable_supported_devices, + _get_scalar_dtype, + _maximize_doc, + _params_doc, + _to_scalar, + _use_grad_for_differentiable, + _view_as_real, + Optimizer, + ParamsT, +) + + +__all__ = ["Rprop", "rprop"] + + +class Rprop(Optimizer): # noqa: D101 + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-2, + etas: tuple[float, float] = (0.5, 1.2), + step_sizes: tuple[float, float] = (1e-6, 50), + *, + capturable: bool = False, + foreach: bool | None = None, + maximize: bool = False, + differentiable: bool = False, + ) -> None: # noqa: D107 + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 <= lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 < etas[0] < 1.0 < etas[1]: + raise ValueError(f"Invalid eta values: {etas[0]}, {etas[1]}") + + defaults = { + "lr": lr, + "etas": etas, + "step_sizes": step_sizes, + "foreach": foreach, + "maximize": maximize, + "differentiable": differentiable, + "capturable": capturable, + } + super().__init__(params, defaults) + + def __setstate__(self, state): # noqa: D105 + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("foreach", None) + group.setdefault("maximize", False) + group.setdefault("differentiable", False) + group.setdefault("capturable", False) + for p in group["params"]: + p_state = self.state.get(p, []) + if len(p_state) != 0 and not torch.is_tensor(p_state["step"]): + step_val = float(p_state["step"]) + p_state["step"] = ( + torch.tensor( + step_val, dtype=_get_scalar_dtype(), device=p.device + ) + if group["capturable"] + else torch.tensor(step_val, dtype=_get_scalar_dtype()) + ) + + def _init_group(self, group, params, grads, prevs, step_sizes, state_steps): + has_complex = False + for p in group["params"]: + if p.grad is None: + continue + has_complex |= torch.is_complex(p) + params.append(p) + grad = p.grad + if grad.is_sparse: + raise RuntimeError("Rprop does not support sparse gradients") + + grads.append(grad) + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = ( + torch.zeros((), dtype=_get_scalar_dtype(), device=p.device) + if group["capturable"] + else torch.zeros((), dtype=_get_scalar_dtype()) + ) + + state["prev"] = torch.zeros_like(p, memory_format=torch.preserve_format) + if p.dtype.is_complex: + # Complex Number should be as if they are two independent real numbers. + # Hence the step_size shouldn't be zero for imaginary part. + state["step_size"] = torch.full_like( + grad, complex(group["lr"], group["lr"]) + ) + else: + state["step_size"] = torch.full_like(grad, _to_scalar(group["lr"])) + + prevs.append(state["prev"]) + step_sizes.append(state["step_size"]) + state_steps.append(state["step"]) + + return has_complex + + @_use_grad_for_differentiable + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + self._accelerator_graph_capture_health_check() + + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params: list[Tensor] = [] + grads: list[Tensor] = [] + prevs: list[Tensor] = [] + step_sizes: list[Tensor] = [] + state_steps: list[Tensor] = [] + + etaminus, etaplus = group["etas"] + step_size_min, step_size_max = group["step_sizes"] + foreach = group["foreach"] + maximize = group["maximize"] + + has_complex = self._init_group( + group, params, grads, prevs, step_sizes, state_steps + ) + + rprop( + params, + grads, + prevs, + step_sizes, + state_steps, + step_size_min=step_size_min, + step_size_max=step_size_max, + etaminus=etaminus, + etaplus=etaplus, + foreach=foreach, + maximize=maximize, + differentiable=group["differentiable"], + capturable=group["capturable"], + has_complex=has_complex, + ) + + return loss + + +Rprop.__doc__ = ( + r"""Implements the resilient backpropagation algorithm. + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \theta_0 \in \mathbf{R}^d \text{ (params)},f(\theta) + \text{ (objective)}, \\ + &\hspace{13mm} \eta_{+/-} \text{ (etaplus, etaminus)}, \Gamma_{max/min} + \text{ (step sizes)} \\ + &\textbf{initialize} : g^0_{prev} \leftarrow 0, + \: \eta_0 \leftarrow \text{lr (learning rate)} \\ + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + &\hspace{5mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm} \textbf{for} \text{ } i = 0, 1, \ldots, d-1 \: \mathbf{do} \\ + &\hspace{10mm} \textbf{if} \: g^i_{prev} g^i_t > 0 \\ + &\hspace{15mm} \eta^i_t \leftarrow \mathrm{min}(\eta^i_{t-1} \eta_{+}, + \Gamma_{max}) \\ + &\hspace{10mm} \textbf{else if} \: g^i_{prev} g^i_t < 0 \\ + &\hspace{15mm} \eta^i_t \leftarrow \mathrm{max}(\eta^i_{t-1} \eta_{-}, + \Gamma_{min}) \\ + &\hspace{15mm} g^i_t \leftarrow 0 \\ + &\hspace{10mm} \textbf{else} \: \\ + &\hspace{15mm} \eta^i_t \leftarrow \eta^i_{t-1} \\ + &\hspace{5mm}\theta_t \leftarrow \theta_{t-1}- \eta_t \mathrm{sign}(g_t) \\ + &\hspace{5mm}g_{prev} \leftarrow g_t \\ + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + For further details regarding the algorithm we refer to the paper + `A Direct Adaptive Method for Faster Backpropagation Learning: The RPROP Algorithm + `_.""" # codespell:ignore + + rf""" + + Args: + {_params_doc} + lr (float, optional): learning rate (default: 1e-2) + etas (Tuple[float, float], optional): pair of (etaminus, etaplus), that + are multiplicative increase and decrease factors + (default: (0.5, 1.2)) + step_sizes (Tuple[float, float], optional): a pair of minimal and + maximal allowed step sizes (default: (1e-6, 50)) + {_capturable_doc} + {_foreach_doc} + {_maximize_doc} + {_differentiable_doc} + + """ +) + + +def _single_tensor_rprop( + params: list[Tensor], + grads: list[Tensor], + prevs: list[Tensor], + step_sizes: list[Tensor], + state_steps: list[Tensor], + *, + step_size_min: float, + step_size_max: float, + etaminus: float, + etaplus: float, + maximize: bool, + capturable: bool, + differentiable: bool, + has_complex: bool, +) -> None: + for i, param in enumerate(params): + grad = grads[i] + grad = grad if not maximize else -grad + prev = prevs[i] + step_size = step_sizes[i] + step = state_steps[i] + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices() + if not ( + param.device.type == step.device.type + and param.device.type in capturable_supported_devices + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + step += 1 + + if torch.is_complex(param): + grad = torch.view_as_real(grad) + prev = torch.view_as_real(prev) + param = torch.view_as_real(param) + step_size = torch.view_as_real(step_size) + if differentiable: + sign = grad.mul(prev.clone()).sign() + else: + sign = grad.mul(prev).sign() + + if capturable: + sign.copy_(torch.where(sign.gt(0), etaplus, sign)) + sign.copy_(torch.where(sign.lt(0), etaminus, sign)) + sign.copy_(torch.where(sign.eq(0), 1, sign)) + else: + sign[sign.gt(0)] = etaplus + sign[sign.lt(0)] = etaminus + sign[sign.eq(0)] = 1 + + # update stepsizes with step size updates + step_size.mul_(sign).clamp_(step_size_min, step_size_max) + + # for dir<0, dfdx=0 + # for dir>=0 dfdx=dfdx + grad = grad.clone(memory_format=torch.preserve_format) + if capturable: + grad.copy_(torch.where(sign.eq(etaminus), 0, grad)) + else: + grad[sign.eq(etaminus)] = 0 + + # update parameters + param.addcmul_(grad.sign(), step_size, value=-1) + prev.copy_(grad) + + +def _multi_tensor_rprop( + params: list[Tensor], + grads: list[Tensor], + prevs: list[Tensor], + step_sizes: list[Tensor], + state_steps: list[Tensor], + *, + step_size_min: float, + step_size_max: float, + etaminus: float, + etaplus: float, + maximize: bool, + capturable: bool, + differentiable: bool, + has_complex: bool, +) -> None: + if len(params) == 0: + return + + if differentiable: + raise AssertionError("_foreach ops don't support autograd") + + # If compiling, the compiler will handle cudagraph checks, see note [torch.compile x capturable] + if not torch.compiler.is_compiling() and capturable: + capturable_supported_devices = _get_capturable_supported_devices() + if not all( + p.device.type == step.device.type + and p.device.type in capturable_supported_devices + for p, step in zip(params, state_steps, strict=True) + ): + raise AssertionError( + f"If capturable=True, params and state_steps must be on supported devices: {capturable_supported_devices}." + ) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, prevs, step_sizes, state_steps] # type: ignore[list-item] + ) + for ( + grouped_params_, + grouped_grads_, + grouped_prevs_, + grouped_step_sizes_, + grouped_state_steps_, + ), _ in grouped_tensors.values(): + grouped_params = cast(list[Tensor], grouped_params_) + grouped_grads = cast(list[Tensor], grouped_grads_) + grouped_prevs = cast(list[Tensor], grouped_prevs_) + grouped_step_sizes = cast(list[Tensor], grouped_step_sizes_) + grouped_state_steps = cast(list[Tensor], grouped_state_steps_) + + # Update steps + # If steps are on CPU, foreach will fall back to the slow path, which is a for-loop calling t.add(1) over + # and over. 1 will then be wrapped into a Tensor over and over again, which is slower than if we just + # wrapped it once now. The alpha is required to assure we go to the right overload. + if not torch.compiler.is_compiling() and grouped_state_steps[0].is_cpu: + torch._foreach_add_( + grouped_state_steps, torch.tensor(1.0, device="cpu"), alpha=1.0 + ) + else: + torch._foreach_add_(grouped_state_steps, 1) + + # Handle complex params + if has_complex: + _view_as_real( + grouped_params, grouped_grads, grouped_prevs, grouped_step_sizes + ) + + signs = torch._foreach_mul(grouped_grads, grouped_prevs) + if maximize: + torch._foreach_neg_(signs) + + # At the end of the step, grouped_prevs will contain the current grads, so we reuse + # grouped_prevs memory instead of creating a new buffer, but, for clarity, we reassign + # to keep referring to the buffer as grouped_grads. + torch._foreach_copy_(grouped_prevs, grouped_grads) + if maximize: + torch._foreach_neg_(grouped_prevs) + grouped_grads = grouped_prevs + + torch._foreach_sign_(signs) + if capturable: + for sign in signs: + sign.copy_(torch.where(sign.gt(0), etaplus, sign)) + sign.copy_(torch.where(sign.lt(0), etaminus, sign)) + sign.copy_(torch.where(sign.eq(0), 1, sign)) + else: + for sign in signs: + sign[sign.gt(0)] = etaplus + sign[sign.lt(0)] = etaminus + sign[sign.eq(0)] = 1 + + # update stepsizes with step size updates + torch._foreach_mul_(grouped_step_sizes, signs) + for step_size in grouped_step_sizes: + step_size.clamp_(step_size_min, step_size_max) + + # for dir<0, dfdx=0 + # for dir>=0 dfdx=dfdx + grouped_grads = list(grouped_grads) + for i in range(len(grouped_grads)): + grouped_grads[i].copy_( + torch.where(signs[i].eq(etaminus), 0, grouped_grads[i]) + ) + + # explicitly del signs as it's not used after here to save memory + del signs + + # update parameters + grad_signs = [grad.sign() for grad in grouped_grads] + torch._foreach_addcmul_( + grouped_params, grad_signs, grouped_step_sizes, value=-1 + ) + + # Logically, you may expect grouped_prevs to get updated to grouped_grads, but that's + # basically already happened since we've been using grouped_prevs' memory to store + # updated grouped_grads! + + +@_disable_dynamo_if_unsupported(single_tensor_fn=_single_tensor_rprop) +def rprop( + params: list[Tensor], + grads: list[Tensor], + prevs: list[Tensor], + step_sizes: list[Tensor], + state_steps: list[Tensor], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + foreach: bool | None = None, + capturable: bool = False, + maximize: bool = False, + differentiable: bool = False, + has_complex: bool = False, + *, + step_size_min: float, + step_size_max: float, + etaminus: float, + etaplus: float, +) -> None: + r"""Functional API that performs rprop algorithm computation. + + See :class:`~torch.optim.Rprop` for details. + """ + # this check is slow during compilation, so we skip it + # if it's strictly needed we can add this check back in dynamo + if not torch.compiler.is_compiling() and not all( + isinstance(t, torch.Tensor) for t in state_steps + ): + raise RuntimeError( + "API has changed, `state_steps` argument must contain a list of singleton tensors" + ) + + if foreach is None: + _, foreach = _default_to_fused_or_foreach( + params, differentiable, use_fused=False + ) + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + + if foreach and not torch.jit.is_scripting(): + func = _multi_tensor_rprop + else: + func = _single_tensor_rprop + + func( + params, + grads, + prevs, + step_sizes, + state_steps, + step_size_min=step_size_min, + step_size_max=step_size_max, + etaminus=etaminus, + etaplus=etaplus, + capturable=capturable, + maximize=maximize, + differentiable=differentiable, + has_complex=has_complex, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/sgd.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/sgd.py new file mode 100644 index 0000000000000000000000000000000000000000..7c8783fd231827436cee31f13b706143a6267829 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/sgd.py @@ -0,0 +1,544 @@ +# mypy: allow-untyped-defs +r"""Implementation for Stochastic Gradient Descent optimizer.""" + +from typing import cast + +import torch +from torch import Tensor + +from .optimizer import ( + _default_to_fused_or_foreach, + _device_dtype_check_for_fused, + _differentiable_doc, + _foreach_doc, + _fused_doc, + _maximize_doc, + _params_doc, + _to_scalar, + _use_grad_for_differentiable, + DeviceDict, + Optimizer, + ParamsT, +) + + +__all__ = ["SGD", "sgd"] + + +class SGD(Optimizer): # noqa: D101 + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-3, + momentum: float = 0, + dampening: float = 0, + weight_decay: float | Tensor = 0, + nesterov: bool = False, + *, + maximize: bool = False, + foreach: bool | None = None, + differentiable: bool = False, + fused: bool | None = None, + ) -> None: # noqa: D107 + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if lr < 0.0: + raise ValueError(f"Invalid learning rate: {lr}") + if momentum < 0.0: + raise ValueError(f"Invalid momentum value: {momentum}") + if weight_decay < 0.0: + raise ValueError(f"Invalid weight_decay value: {weight_decay}") + + defaults = { + "lr": lr, + "momentum": momentum, + "dampening": dampening, + "weight_decay": weight_decay, + "nesterov": nesterov, + "maximize": maximize, + "foreach": foreach, + "differentiable": differentiable, + "fused": fused, + } + if nesterov and (momentum <= 0 or dampening != 0): + raise ValueError("Nesterov momentum requires a momentum and zero dampening") + super().__init__(params, defaults) + + if fused: + self._step_supports_amp_scaling = True + self._need_device_dtype_check_for_fused = True + if differentiable: + raise RuntimeError("`fused` does not support `differentiable`") + if foreach: + raise RuntimeError("`fused` and `foreach` cannot be `True` together.") + + def __setstate__(self, state): # noqa: D105 + super().__setstate__(state) + for group in self.param_groups: + group.setdefault("nesterov", False) + group.setdefault("maximize", False) + group.setdefault("foreach", None) + group.setdefault("differentiable", False) + group.setdefault("fused", False) + + def _init_group(self, group, params, grads, momentum_buffer_list): + has_sparse_grad = False + + for p in group["params"]: + if p.grad is not None: + if group["fused"] and getattr( + self, "_need_device_dtype_check_for_fused", True + ): + _device_dtype_check_for_fused(p) + self._need_device_dtype_check_for_fused = False + params.append(p) + grads.append(p.grad) + if p.grad.is_sparse: + has_sparse_grad = True + + if group["momentum"] != 0: + state = self.state[p] + momentum_buffer_list.append(state.get("momentum_buffer")) + + return has_sparse_grad + + @_use_grad_for_differentiable + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params: list[Tensor] = [] + grads: list[Tensor] = [] + momentum_buffer_list: list[Tensor | None] = [] + + has_sparse_grad = self._init_group( + group, params, grads, momentum_buffer_list + ) + + sgd( + params, + grads, + momentum_buffer_list, + weight_decay=group["weight_decay"], + momentum=group["momentum"], + lr=group["lr"], + dampening=group["dampening"], + nesterov=group["nesterov"], + maximize=group["maximize"], + has_sparse_grad=has_sparse_grad, + foreach=group["foreach"], + fused=group["fused"], + grad_scale=getattr(self, "grad_scale", None), + found_inf=getattr(self, "found_inf", None), + ) + + if group["momentum"] != 0: + # update momentum_buffers in state + for p, momentum_buffer in zip( + params, momentum_buffer_list, strict=True + ): + state = self.state[p] + state["momentum_buffer"] = momentum_buffer + + return loss + + +SGD.__doc__ = ( + r"""Implements stochastic gradient descent (optionally with momentum). + + .. math:: + \begin{aligned} + &\rule{110mm}{0.4pt} \\ + &\textbf{input} : \gamma \text{ (lr)}, \: \theta_0 \text{ (params)}, \: f(\theta) + \text{ (objective)}, \: \lambda \text{ (weight decay)}, \\ + &\hspace{13mm} \:\mu \text{ (momentum)}, \:\tau \text{ (dampening)}, + \:\textit{ nesterov,}\:\textit{ maximize} \\[-1.ex] + &\rule{110mm}{0.4pt} \\ + &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ + &\hspace{5mm}\textbf{if} \: \textit{maximize}: \\ + &\hspace{10mm}g_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\textbf{else} \\ + &\hspace{10mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ + &\hspace{5mm}\textbf{if} \: \lambda \neq 0 \\ + &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ + &\hspace{5mm}\textbf{if} \: \mu \neq 0 \\ + &\hspace{10mm}\textbf{if} \: t > 1 \\ + &\hspace{15mm} \textbf{b}_t \leftarrow \mu \textbf{b}_{t-1} + (1-\tau) g_t \\ + &\hspace{10mm}\textbf{else} \\ + &\hspace{15mm} \textbf{b}_t \leftarrow g_t \\ + &\hspace{10mm}\textbf{if} \: \textit{nesterov} \\ + &\hspace{15mm} g_t \leftarrow g_{t} + \mu \textbf{b}_t \\ + &\hspace{10mm}\textbf{else} \\[-1.ex] + &\hspace{15mm} g_t \leftarrow \textbf{b}_t \\ + &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \gamma g_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + &\bf{return} \: \theta_t \\[-1.ex] + &\rule{110mm}{0.4pt} \\[-1.ex] + \end{aligned} + + Nesterov momentum is based on the formula from + `On the importance of initialization and momentum in deep learning`__. + """ + + rf""" + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 1e-3) + momentum (float, optional): momentum factor (default: 0) + dampening (float, optional): dampening for momentum (default: 0) + weight_decay (float, optional): weight decay (L2 penalty) (default: 0) + nesterov (bool, optional): enables Nesterov momentum. Only applicable + when momentum is non-zero. (default: False) + {_maximize_doc} + {_foreach_doc} + {_differentiable_doc} + {_fused_doc} + """ + + r""" + + Example: + >>> # xdoctest: +SKIP + >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9) + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + + __ http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf + + .. note:: + The implementation of SGD with Momentum/Nesterov subtly differs from + Sutskever et al. and implementations in some other frameworks. + + Considering the specific case of Momentum, the update can be written as + + .. math:: + \begin{aligned} + v_{t+1} & = \mu * v_{t} + g_{t+1}, \\ + p_{t+1} & = p_{t} - \text{lr} * v_{t+1}, + \end{aligned} + + where :math:`p`, :math:`g`, :math:`v` and :math:`\mu` denote the + parameters, gradient, velocity, and momentum respectively. + + This is in contrast to Sutskever et al. and + other frameworks which employ an update of the form + + .. math:: + \begin{aligned} + v_{t+1} & = \mu * v_{t} + \text{lr} * g_{t+1}, \\ + p_{t+1} & = p_{t} - v_{t+1}. + \end{aligned} + + The Nesterov version is analogously modified. + + Moreover, the initial value of the momentum buffer is set to the + gradient value at the first step. This is in contrast to some other + frameworks that initialize it to all zeros. One notable side effect + of this decision is that the first momentum value will not be scaled + by dampening. Dampening will be applied starting at the second step. + + """ +) + + +def sgd( + params: list[Tensor], + d_p_list: list[Tensor], + momentum_buffer_list: list[Tensor | None], + # kwonly args with defaults are not supported by functions compiled with torchscript issue #70627 + # setting this as kwarg for now as functional API is compiled by torch/distributed/optim + has_sparse_grad: bool = False, + foreach: bool | None = None, + fused: bool | None = None, + grad_scale: Tensor | None = None, + found_inf: Tensor | None = None, + *, + weight_decay: float, + momentum: float, + lr: float, + dampening: float, + nesterov: bool, + maximize: bool, +) -> None: + r"""Functional API that performs SGD algorithm computation. + + See :class:`~torch.optim.SGD` for details. + """ + # Respect when the user inputs False/True for foreach or fused. We only want to change + # the default when neither have been user-specified. Note that we default to foreach + # and pass False to use_fused. This is not a mistake--we want to give the fused impl + # bake-in time before making it the default, even if it is typically faster. + if foreach is None and fused is None: + # why must we be explicit about an if statement for torch.jit.is_scripting here? + # because JIT can't handle Optionals nor fancy conditionals when scripting + if not torch.jit.is_scripting(): + fused, foreach = _default_to_fused_or_foreach( + params, differentiable=False, use_fused=False + ) + else: + foreach = False + fused = False + if foreach is None: + foreach = False + if fused is None: + fused = False + + if foreach and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with foreach optimizers") + if fused and torch.jit.is_scripting(): + raise RuntimeError("torch.jit.script not supported with fused optimizers") + + if foreach and not torch.jit.is_scripting(): + func = _multi_tensor_sgd + elif fused and not torch.jit.is_scripting(): + func = _fused_sgd + else: + func = _single_tensor_sgd + + func( + params, + d_p_list, + momentum_buffer_list, + weight_decay=weight_decay, + momentum=momentum, + lr=lr, + dampening=dampening, + nesterov=nesterov, + has_sparse_grad=has_sparse_grad, + maximize=maximize, + grad_scale=grad_scale, + found_inf=found_inf, + ) + + +def _single_tensor_sgd( + params: list[Tensor], + grads: list[Tensor], + momentum_buffer_list: list[Tensor | None], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + weight_decay: float, + momentum: float, + lr: float, + dampening: float, + nesterov: bool, + maximize: bool, + has_sparse_grad: bool, +) -> None: + if grad_scale is not None or found_inf is not None: + raise AssertionError("Expected grad_scale and found_inf to be None") + + if not torch.jit.is_scripting(): + lr = _to_scalar(lr) + + for i, param in enumerate(params): + grad = grads[i] if not maximize else -grads[i] + + if weight_decay != 0: + # Nested if is necessary to bypass jitscript rules + if isinstance(weight_decay, Tensor): + if weight_decay.requires_grad: + # usually this is the differentiable path, which is why the param.clone() is needed + grad = grad.addcmul_(param.clone(), weight_decay) + else: + grad = grad.add(param, alpha=weight_decay) + else: + grad = grad.add(param, alpha=weight_decay) + + if momentum != 0: + buf = momentum_buffer_list[i] + + if buf is None: + buf = grad.detach().clone() + momentum_buffer_list[i] = buf + else: + buf.mul_(momentum).add_(grad, alpha=1 - dampening) + + if nesterov: + grad = grad.add(buf, alpha=momentum) + else: + grad = buf + + # Nested if is necessary to bypass jitscript rules + if isinstance(lr, Tensor): + if lr.requires_grad: + param.addcmul_(grad, lr, value=-1) + else: + # pyrefly: ignore [bad-argument-type] + param.add_(grad, alpha=-lr) + else: + param.add_(grad, alpha=-lr) + + +def _multi_tensor_sgd( + params: list[Tensor], + grads: list[Tensor], + momentum_buffer_list: list[Tensor | None], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + weight_decay: float, + momentum: float, + lr: float, + dampening: float, + nesterov: bool, + maximize: bool, + has_sparse_grad: bool, +) -> None: + if grad_scale is not None or found_inf is not None: + raise AssertionError("Expected grad_scale and found_inf to be None") + + if len(params) == 0: + return + + lr = _to_scalar(lr) + + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, momentum_buffer_list], # type: ignore[list-item] + with_indices=True, + ) + for ( + device_params_, + device_grads_, + device_momentum_buffer_list, + ), indices in grouped_tensors.values(): + device_params: list[Tensor] = cast(list[Tensor], device_params_) + device_grads: list[Tensor] = cast(list[Tensor], device_grads_) + + device_has_sparse_grad = has_sparse_grad and any( + grad.is_sparse for grad in device_grads + ) + + if maximize: + device_grads = torch._foreach_neg(device_grads) # type: ignore[assignment] + + if weight_decay != 0: + # Reuse the intermediate memory (device_grads) already allocated for maximize + if maximize: + torch._foreach_add_(device_grads, device_params, alpha=weight_decay) + else: + device_grads = torch._foreach_add( # type: ignore[assignment] + device_grads, device_params, alpha=weight_decay + ) + + if momentum != 0: + bufs: list[Tensor] = [] + + all_states_with_momentum_buffer = True + for i in range(len(device_momentum_buffer_list)): + if device_momentum_buffer_list[i] is None: + all_states_with_momentum_buffer = False + break + else: + bufs.append(cast(Tensor, device_momentum_buffer_list[i])) + + if all_states_with_momentum_buffer: + torch._foreach_mul_(bufs, momentum) + torch._foreach_add_(bufs, device_grads, alpha=1 - dampening) + else: + bufs = [] + + for i in range(len(device_momentum_buffer_list)): + if device_momentum_buffer_list[i] is None: + buf = device_momentum_buffer_list[i] = momentum_buffer_list[ + indices[i] + ] = device_grads[i].detach().clone() + else: + buf = cast(Tensor, device_momentum_buffer_list[i]) + buf.mul_(momentum).add_(device_grads[i], alpha=1 - dampening) + + bufs.append(buf) + + if nesterov: + torch._foreach_add_(device_grads, bufs, alpha=momentum) + else: + device_grads = bufs + + if not device_has_sparse_grad: + # handle internal item() call if lr is a tensor + if isinstance(lr, torch.Tensor) and torch.compiler.is_compiling(): + grads_x_lr = torch._foreach_mul(device_grads, -lr) + torch._foreach_add_(device_params, grads_x_lr) + else: + torch._foreach_add_(device_params, device_grads, alpha=-lr) + else: + # foreach APIs don't support sparse + for i in range(len(device_params)): + device_params[i].add_(device_grads[i], alpha=-lr) + + +def _fused_sgd( + params: list[Tensor], + grads: list[Tensor], + momentum_buffer_list: list[Tensor | None], + grad_scale: Tensor | None, + found_inf: Tensor | None, + *, + weight_decay: float, + momentum: float, + lr: float, + dampening: float, + nesterov: bool, + maximize: bool, + has_sparse_grad: bool, +) -> None: + if not params: + return + if has_sparse_grad: + raise RuntimeError("`_fused_sgd` does not support sparse gradients") + grad_scale_dict: DeviceDict = ( + {grad_scale.device: grad_scale} if grad_scale is not None else {} + ) + found_inf_dict: DeviceDict = ( + {found_inf.device: found_inf} if found_inf is not None else {} + ) + + no_momentum_buffer = momentum == 0 + is_first_step = ( + all(t is None for t in momentum_buffer_list) and not no_momentum_buffer + ) + if is_first_step: + for i, g in enumerate(grads): + momentum_buffer_list[i] = torch.empty_like(g) + grouped_tensors = Optimizer._group_tensors_by_device_and_dtype( + [params, grads, momentum_buffer_list], # type: ignore[list-item] + with_indices=False, + ) + for (device, _), ( + (device_params_, device_grads_, device_momentum_buffer_list), + _, + ) in grouped_tensors.items(): + device_params: list[Tensor] = cast(list[Tensor], device_params_) + device_grads: list[Tensor] = cast(list[Tensor], device_grads_) + device_grad_scale, device_found_inf = None, None + if grad_scale is not None: + device_grad_scale = grad_scale_dict.setdefault( + device, grad_scale.to(device) + ) + if found_inf_dict is not None and found_inf is not None: + device_found_inf = found_inf_dict.setdefault(device, found_inf.to(device)) + torch._fused_sgd_( + device_params, + device_grads, + [] + if no_momentum_buffer + else cast(list[Tensor], device_momentum_buffer_list), + weight_decay=weight_decay, + momentum=momentum, + lr=lr, + dampening=dampening, + nesterov=nesterov, + maximize=maximize, + is_first_step=is_first_step, + grad_scale=device_grad_scale, + found_inf=device_found_inf, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/sparse_adam.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/sparse_adam.py new file mode 100644 index 0000000000000000000000000000000000000000..d6196cb20cd4e3c5ab1af7a16272d17a961bcad8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/sparse_adam.py @@ -0,0 +1,189 @@ +# mypy: allow-untyped-defs + +import torch +from torch import Tensor + +from . import _functional as F +from .optimizer import _maximize_doc, _params_doc, _to_scalar, Optimizer, ParamsT + + +__all__ = ["SparseAdam"] + + +class SparseAdam(Optimizer): + def __init__( + self, + params: ParamsT, + lr: float | Tensor = 1e-3, + betas: tuple[float, float] = (0.9, 0.999), + eps: float = 1e-8, + maximize: bool = False, + ) -> None: + if isinstance(lr, Tensor) and lr.numel() != 1: + raise ValueError("Tensor lr must be 1-element") + if not 0.0 < lr: + raise ValueError(f"Invalid learning rate: {lr}") + if not 0.0 < eps: + raise ValueError(f"Invalid epsilon value: {eps}") + if not 0.0 <= betas[0] < 1.0: + raise ValueError(f"Invalid beta parameter at index 0: {betas[0]}") + if not 0.0 <= betas[1] < 1.0: + raise ValueError(f"Invalid beta parameter at index 1: {betas[1]}") + + defaults = { + "lr": lr, + "betas": betas, + "eps": eps, + "maximize": maximize, + } + super().__init__(params, defaults) + + sparse_params = [] + complex_params = [] + for index, param_group in enumerate(self.param_groups): + if not isinstance(param_group, dict): + raise AssertionError( + f"param_groups must be a list of dicts, but got {type(param_group)}" + ) + # given param group, convert given params to a list first before iterating + for d_index, d_param in enumerate(param_group["params"]): + if d_param.is_sparse: + sparse_params.append([index, d_index]) + if d_param.is_complex(): + complex_params.append([index, d_index]) + if sparse_params: + raise ValueError( + f"Sparse params at indices {sparse_params}: SparseAdam requires dense parameter tensors" + ) + if complex_params: + raise ValueError( + f"Complex params at indices {complex_params}: SparseAdam does not support complex parameters" + ) + + @torch.no_grad() + def step(self, closure=None): + """Perform a single optimization step. + + Args: + closure (Callable, optional): A closure that reevaluates the model + and returns the loss. + """ + loss = None + if closure is not None: + with torch.enable_grad(): + loss = closure() + + for group in self.param_groups: + params_with_grad: list[Tensor] = [] + grads: list[Tensor] = [] + exp_avgs: list[Tensor] = [] + exp_avg_sqs: list[Tensor] = [] + state_steps: list[int] = [] + beta1, beta2 = group["betas"] + maximize = group.get("maximize", False) + + for p in group["params"]: + if p.grad is not None: + params_with_grad.append(p) + if not p.grad.is_sparse: + raise RuntimeError( + "SparseAdam does not support dense gradients, please consider Adam instead" + ) + grads.append(p.grad) + + state = self.state[p] + + # State initialization + if len(state) == 0: + state["step"] = 0 + # Exponential moving average of gradient values + state["exp_avg"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + # Exponential moving average of squared gradient values + state["exp_avg_sq"] = torch.zeros_like( + p, memory_format=torch.preserve_format + ) + + exp_avgs.append(state["exp_avg"]) + exp_avg_sqs.append(state["exp_avg_sq"]) + + # update the steps for each param group update + state["step"] += 1 + # record the step after step update + state_steps.append(state["step"]) + + F.sparse_adam( + params_with_grad, + grads, + exp_avgs, + exp_avg_sqs, + state_steps, + eps=group["eps"], + beta1=beta1, + beta2=beta2, + lr=_to_scalar(group["lr"]), + maximize=maximize, + ) + + return loss + + +SparseAdam.__doc__ = rf"""SparseAdam implements a masked version of the Adam algorithm + suitable for sparse gradients. Currently, due to implementation constraints (explained + below), SparseAdam is only intended for a narrow subset of use cases, specifically + parameters of a dense layout with gradients of a sparse layout. This occurs in a + special case where the module backwards produces grads already in a sparse layout. + One example NN module that behaves as such is ``nn.Embedding(sparse=True)``. + + SparseAdam approximates the Adam algorithm by masking out the parameter and moment + updates corresponding to the zero values in the gradients. Whereas the Adam algorithm + will update the first moment, the second moment, and the parameters based on all values + of the gradients, SparseAdam only updates the moments and parameters corresponding + to the non-zero values of the gradients. + + A simplified way of thinking about the `intended` implementation is as such: + + 1. Create a mask of the non-zero values in the sparse gradients. For example, + if your gradient looks like [0, 5, 0, 0, 9], the mask would be [0, 1, 0, 0, 1]. + 2. Apply this mask over the running moments and do computation on only the + non-zero values. + 3. Apply this mask over the parameters and only apply an update on non-zero values. + + In actuality, we use sparse layout Tensors to optimize this approximation, which means the + more gradients that are masked by not being materialized, the more performant the optimization. + Since we rely on using sparse layout tensors, we infer that any materialized value in the + sparse layout is non-zero and we do NOT actually verify that all values are not zero! + It is important to not conflate a semantically sparse tensor (a tensor where many + of its values are zeros) with a sparse layout tensor (a tensor where ``.is_sparse`` + returns ``True``). The SparseAdam approximation is intended for `semantically` sparse + tensors and the sparse layout is only a implementation detail. A clearer implementation + would be to use MaskedTensors, but those are experimental. + + + .. note:: + + If you suspect your gradients are semantically sparse (but do not have sparse + layout), this variant may not be the best for you. Ideally, you want to avoid + materializing anything that is suspected to be sparse in the first place, since + needing to convert all your grads from dense layout to sparse layout may outweigh + the performance gain. Here, using Adam may be the best alternative, unless you + can easily rig up your module to output sparse grads similar to + ``nn.Embedding(sparse=True)``. If you insist on converting your grads, you can do + so by manually overriding your parameters' ``.grad`` fields with their sparse + equivalents before calling ``.step()``. + + + Args: + {_params_doc} + lr (float, Tensor, optional): learning rate (default: 1e-3) + betas (Tuple[float, float], optional): coefficients used for computing + running averages of gradient and its square (default: (0.9, 0.999)) + eps (float, optional): term added to the denominator to improve + numerical stability (default: 1e-8) + {_maximize_doc} + + .. _Adam\: A Method for Stochastic Optimization: + https://arxiv.org/abs/1412.6980 + + """ diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/swa_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/swa_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ec60e81aff9b06afa28c1631360280db8e9478f3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/optim/swa_utils.py @@ -0,0 +1,592 @@ +# mypy: allow-untyped-defs +r"""Implementation for Stochastic Weight Averaging implementation.""" + +from __future__ import annotations + +import itertools +import math +import warnings +from copy import deepcopy +from typing import Any, cast, Literal, TYPE_CHECKING +from typing_extensions import override + +import torch +from torch import Tensor +from torch.nn import Module +from torch.optim.lr_scheduler import _format_param, LRScheduler +from torch.utils._foreach_utils import _get_foreach_kernels_supported_devices + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + + from .optimizer import Optimizer + + +__all__ = [ + "AveragedModel", + "update_bn", + "SWALR", + "get_ema_multi_avg_fn", + "get_swa_multi_avg_fn", + "get_ema_avg_fn", + "get_swa_avg_fn", +] + +from torch.utils._foreach_utils import _group_tensors_by_device_and_dtype + + +PARAM_LIST = tuple[Tensor, ...] | list[Tensor] + + +def get_ema_multi_avg_fn(decay=0.999): + """Get the function applying exponential moving average (EMA) across multiple params. + + The EMA is computed as: + + .. math:: + W_0^{\\text{EMA}} = W_0^{\\text{model}} + + .. math:: + W_{t+1}^{\\text{EMA}} = \\text{decay} \\times W_t^{\\text{EMA}} + (1 - \\text{decay}) \\times W_{t+1}^{\\text{model}} + + where :math:`W_t^{\\text{EMA}}` is the EMA parameter at step :math:`t`, + :math:`W_t^{\\text{model}}` is the model parameter at step :math:`t`, + and :math:`\\text{decay}` is the decay rate (default: 0.999). + + Args: + decay (float): Decay rate for EMA. Must be in the range [0, 1]. Default: 0.999 + + Returns: + Callable: A function that updates EMA parameters given current model parameters + """ + + if decay < 0.0 or decay > 1.0: + raise ValueError( + f"Invalid decay value {decay} provided. Please provide a value in [0,1] range." + ) + + @torch.no_grad() + def ema_update( + ema_param_list: PARAM_LIST, current_param_list: PARAM_LIST, _ + ) -> None: + # foreach lerp only handles float and complex + if torch.is_floating_point(ema_param_list[0]) or torch.is_complex( + ema_param_list[0] + ): + torch._foreach_lerp_(ema_param_list, current_param_list, 1 - decay) + else: + for p_ema, p_model in zip(ema_param_list, current_param_list, strict=True): + p_ema.copy_(p_ema * decay + p_model * (1 - decay)) + + return ema_update + + +def get_swa_multi_avg_fn(): + """Get the function applying stochastic weight average (SWA) across multiple params.""" + + @torch.no_grad() + def swa_update( + averaged_param_list: PARAM_LIST, + current_param_list: PARAM_LIST, + num_averaged: Tensor | int, + ) -> None: + # foreach lerp only handles float and complex + if torch.is_floating_point(averaged_param_list[0]) or torch.is_complex( + averaged_param_list[0] + ): + torch._foreach_lerp_( + averaged_param_list, + current_param_list, + cast(float, 1 / (num_averaged + 1)), + ) + else: + diffs = torch._foreach_sub(current_param_list, averaged_param_list) + if isinstance(num_averaged, Tensor): + torch._foreach_addcdiv_( + averaged_param_list, + diffs, + [num_averaged + 1] * len(averaged_param_list), + ) + else: + torch._foreach_add_( + averaged_param_list, diffs, alpha=1.0 / (num_averaged + 1) + ) + + return swa_update + + +def get_ema_avg_fn(decay=0.999): + """Get the function applying exponential moving average (EMA) across multiple params. + + The EMA is computed as: + + .. math:: + W_0^{\\text{EMA}} = W_0^{\\text{model}} + + .. math:: + W_{t+1}^{\\text{EMA}} = \\text{decay} \\times W_t^{\\text{EMA}} + (1 - \\text{decay}) \\times W_{t+1}^{\\text{model}} + + where :math:`W_t^{\\text{EMA}}` is the EMA parameter at step :math:`t`, + :math:`W_t^{\\text{model}}` is the model parameter at step :math:`t`, + and :math:`\\text{decay}` is the decay rate (default: 0.999). + + Args: + decay (float): Decay rate for EMA. Must be in the range [0, 1]. Default: 0.999 + + Returns: + Callable: A function that updates EMA parameters given current model parameters + """ + + if decay < 0.0 or decay > 1.0: + raise ValueError( + f"Invalid decay value {decay} provided. Please provide a value in [0,1] range." + ) + + @torch.no_grad() + def ema_update(ema_param: Tensor, current_param: Tensor, num_averaged): + return decay * ema_param + (1 - decay) * current_param + + return ema_update + + +def get_swa_avg_fn(): + """Get the function applying stochastic weight average (SWA) across a single param.""" + + @torch.no_grad() + def swa_update( + averaged_param: Tensor, current_param: Tensor, num_averaged: Tensor | int + ): + return averaged_param + (current_param - averaged_param) / (num_averaged + 1) + + return swa_update + + +class AveragedModel(Module): + r"""Implements averaged model for Stochastic Weight Averaging (SWA) and Exponential Moving Average (EMA). + + Stochastic Weight Averaging was proposed in `Averaging Weights Leads to + Wider Optima and Better Generalization`_ by Pavel Izmailov, Dmitrii + Podoprikhin, Timur Garipov, Dmitry Vetrov and Andrew Gordon Wilson + (UAI 2018). + + Exponential Moving Average is a variation of `Polyak averaging`_, + but using exponential weights instead of equal weights across iterations. + + AveragedModel class creates a copy of the provided module :attr:`model` + on the device :attr:`device` and allows to compute running averages of the + parameters of the :attr:`model`. + + Args: + model (torch.nn.Module): model to use with SWA/EMA + device (torch.device, optional): if provided, the averaged model will be + stored on the :attr:`device` + avg_fn (function, optional): the averaging function used to update + parameters; the function must take in the current value of the + :class:`AveragedModel` parameter, the current value of :attr:`model` + parameter, and the number of models already averaged; if None, + an equally weighted average is used (default: None) + multi_avg_fn (function, optional): the averaging function used to update + parameters inplace; the function must take in the current values of the + :class:`AveragedModel` parameters as a list, the current values of :attr:`model` + parameters as a list, and the number of models already averaged; if None, + an equally weighted average is used (default: None) + use_buffers (bool): if ``True``, it will compute running averages for + both the parameters and the buffers of the model. (default: ``False``) + + Example: + >>> # xdoctest: +SKIP("undefined variables") + >>> loader, optimizer, model, loss_fn = ... + >>> swa_model = torch.optim.swa_utils.AveragedModel(model) + >>> scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, + >>> T_max=300) + >>> swa_start = 160 + >>> swa_scheduler = SWALR(optimizer, swa_lr=0.05) + >>> for i in range(300): + >>> for input, target in loader: + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + >>> if i > swa_start: + >>> swa_model.update_parameters(model) + >>> swa_scheduler.step() + >>> else: + >>> scheduler.step() + >>> + >>> # Update bn statistics for the swa_model at the end + >>> torch.optim.swa_utils.update_bn(loader, swa_model) + + You can also use custom averaging functions with the `avg_fn` or `multi_avg_fn` parameters. + If no averaging function is provided, the default is to compute + equally-weighted average of the weights (SWA). + + Example: + >>> # xdoctest: +SKIP("undefined variables") + >>> # Compute exponential moving averages of the weights and buffers + >>> ema_model = torch.optim.swa_utils.AveragedModel(model, + >>> torch.optim.swa_utils.get_ema_multi_avg_fn(0.9), use_buffers=True) + + .. note:: + When using SWA/EMA with models containing Batch Normalization you may + need to update the activation statistics for Batch Normalization. + This can be done either by using the :meth:`torch.optim.swa_utils.update_bn` + or by setting :attr:`use_buffers` to `True`. The first approach updates the + statistics in a post-training step by passing data through the model. The + second does it during the parameter update phase by averaging all buffers. + Empirical evidence has shown that updating the statistics in normalization + layers increases accuracy, but you may wish to empirically test which + approach yields the best results in your problem. + + .. note:: + :attr:`avg_fn` and `multi_avg_fn` are not saved in the :meth:`state_dict` of the model. + + .. note:: + When :meth:`update_parameters` is called for the first time (i.e. + :attr:`n_averaged` is `0`) the parameters of `model` are copied + to the parameters of :class:`AveragedModel`. For every subsequent + call of :meth:`update_parameters` the function `avg_fn` is used + to update the parameters. + + .. _Averaging Weights Leads to Wider Optima and Better Generalization: + https://arxiv.org/abs/1803.05407 + .. _There Are Many Consistent Explanations of Unlabeled Data: Why You Should + Average: + https://arxiv.org/abs/1806.05594 + .. _SWALP: Stochastic Weight Averaging in Low-Precision Training: + https://arxiv.org/abs/1904.11943 + .. _Stochastic Weight Averaging in Parallel: Large-Batch Training That + Generalizes Well: + https://arxiv.org/abs/2001.02312 + .. _Polyak averaging: + https://paperswithcode.com/method/polyak-averaging + """ + + n_averaged: Tensor + + def __init__( + self, + model: Module, + device: int | torch.device | None = None, + avg_fn: Callable[[Tensor, Tensor, Tensor | int], Tensor] | None = None, + multi_avg_fn: Callable[[PARAM_LIST, PARAM_LIST, Tensor | int], None] + | None = None, + use_buffers=False, + ) -> None: # noqa: D107 + super().__init__() + if avg_fn is not None and multi_avg_fn is not None: + raise AssertionError( + "Only one of avg_fn and multi_avg_fn should be provided" + ) + self.module = deepcopy(model) + if device is not None: + self.module = self.module.to(device) + self.register_buffer( + "n_averaged", torch.tensor(0, dtype=torch.long, device=device) + ) + self.avg_fn = avg_fn + self.multi_avg_fn = multi_avg_fn + self.use_buffers = use_buffers + + def forward(self, *args, **kwargs): + """Forward pass.""" + return self.module(*args, **kwargs) + + def update_parameters(self, model: Module) -> None: + """Update model parameters.""" + self_param = ( + # pyrefly: ignore [bad-argument-type] + itertools.chain(self.module.parameters(), self.module.buffers()) + if self.use_buffers + else self.parameters() + ) + model_param = ( + # pyrefly: ignore [bad-argument-type] + itertools.chain(model.parameters(), model.buffers()) + if self.use_buffers + else model.parameters() + ) + self_param_detached: list[Tensor | None] = [] + model_param_detached: list[Tensor | None] = [] + copy_param = bool(self.n_averaged == 0) + for p_averaged, p_model in zip(self_param, model_param, strict=False): + p_model_ = p_model.detach().to(p_averaged.device) + self_param_detached.append(p_averaged.detach()) + model_param_detached.append(p_model_) + if copy_param: + p_averaged.detach().copy_(p_model_) + + if self.n_averaged > 0: + if self.multi_avg_fn is not None or self.avg_fn is None: + grouped_tensors = _group_tensors_by_device_and_dtype( + [self_param_detached, model_param_detached] + ) + for (device, _), ( + [self_params, model_params], + _, + ) in grouped_tensors.items(): + if self.multi_avg_fn: + self.multi_avg_fn( + self_params, # type: ignore[arg-type] + model_params, # type: ignore[arg-type] + self.n_averaged.to(device), + ) + elif ( + device is not None + and device.type in _get_foreach_kernels_supported_devices() + ): + multi_avg_fn = get_swa_multi_avg_fn() + multi_avg_fn( + self_params, model_params, self.n_averaged.to(device) + ) + else: + avg_fn = get_swa_avg_fn() + n_averaged = self.n_averaged.to(device) + for p_averaged, p_model in zip( # type: ignore[assignment] + self_params, model_params, strict=True + ): + # pyrefly: ignore [missing-attribute] + p_averaged.copy_(avg_fn(p_averaged, p_model, n_averaged)) + else: + for p_averaged, p_model in zip( # type: ignore[assignment] + self_param_detached, model_param_detached, strict=True + ): + # pyrefly: ignore [missing-attribute] + n_averaged = self.n_averaged.to(p_averaged.device) + # pyrefly: ignore [missing-attribute] + p_averaged.detach().copy_( + # pyrefly: ignore [missing-attribute, bad-argument-type] + self.avg_fn(p_averaged.detach(), p_model, n_averaged) + ) + + if not self.use_buffers: + # If not apply running averages to the buffers, + # keep the buffers in sync with the source model. + for b_swa, b_model in zip( + self.module.buffers(), model.buffers(), strict=True + ): + b_swa.detach().copy_(b_model.detach().to(b_swa.device)) + self.n_averaged += 1 + + +@torch.no_grad() +def update_bn( + loader: Iterable[Any], + model: Module, + device: int | torch.device | None = None, +) -> None: + r"""Update BatchNorm running_mean, running_var buffers in the model. + + It performs one pass over data in `loader` to estimate the activation + statistics for BatchNorm layers in the model. + + Args: + loader (torch.utils.data.DataLoader): dataset loader to compute the + activation statistics on. Each data batch should be either a + tensor, or a list/tuple whose first element is a tensor + containing data. + model (torch.nn.Module): model for which we seek to update BatchNorm + statistics. + device (torch.device, optional): If set, data will be transferred to + :attr:`device` before being passed into :attr:`model`. + + Example: + >>> # xdoctest: +SKIP("Undefined variables") + >>> loader, model = ... + >>> torch.optim.swa_utils.update_bn(loader, model) + + .. note:: + The `update_bn` utility assumes that each data batch in :attr:`loader` + is either a tensor or a list or tuple of tensors; in the latter case it + is assumed that :meth:`model.forward()` should be called on the first + element of the list or tuple corresponding to the data batch. + """ + momenta = {} + for module in model.modules(): + if isinstance(module, torch.nn.modules.batchnorm._BatchNorm): + module.reset_running_stats() + momenta[module] = module.momentum + + if not momenta: + return + + was_training = model.training + model.train() + for module in momenta: + module.momentum = None + + for input in loader: + if isinstance(input, (list, tuple)): + input = input[0] + if device is not None: + input = input.to(device) + + model(input) + + for bn_module in momenta: + bn_module.momentum = momenta[bn_module] + model.train(was_training) + + +class SWALR(LRScheduler): + r"""Anneals the learning rate in each parameter group to a fixed value. + + This learning rate scheduler is meant to be used with Stochastic Weight + Averaging (SWA) method (see `torch.optim.swa_utils.AveragedModel`). + + Args: + optimizer (torch.optim.Optimizer): wrapped optimizer + swa_lrs (float or list): the learning rate value for all param groups + together or separately for each group. + annealing_epochs (int): number of epochs in the annealing phase + (default: 10) + annealing_strategy (str): "cos" or "linear"; specifies the annealing + strategy: "cos" for cosine annealing, "linear" for linear annealing + (default: "cos") + last_epoch (int): the index of the last epoch (default: -1) + + The :class:`SWALR` scheduler can be used together with other + schedulers to switch to a constant learning rate late in the training + as in the example below. + + Example: + >>> # xdoctest: +SKIP("Undefined variables") + >>> loader, optimizer, model = ... + >>> lr_lambda = lambda epoch: 0.9 + >>> scheduler = torch.optim.lr_scheduler.MultiplicativeLR(optimizer, + >>> lr_lambda=lr_lambda) + >>> swa_scheduler = torch.optim.swa_utils.SWALR(optimizer, + >>> anneal_strategy="linear", anneal_epochs=20, swa_lr=0.05) + >>> swa_start = 160 + >>> for i in range(300): + >>> for input, target in loader: + >>> optimizer.zero_grad() + >>> loss_fn(model(input), target).backward() + >>> optimizer.step() + >>> if i > swa_start: + >>> swa_scheduler.step() + >>> else: + >>> scheduler.step() + + .. _Averaging Weights Leads to Wider Optima and Better Generalization: + https://arxiv.org/abs/1803.05407 + """ + + def __init__( + self, + optimizer: Optimizer, + swa_lr: float, + anneal_epochs=10, + anneal_strategy: Literal["cos", "linear"] = "cos", + last_epoch=-1, + ) -> None: # noqa: D107 + swa_lrs = _format_param("swa_lr", optimizer, swa_lr) + for swa_lr, group in zip(swa_lrs, optimizer.param_groups, strict=True): + group["swa_lr"] = swa_lr + if anneal_strategy not in ["cos", "linear"]: + raise ValueError( + "anneal_strategy must by one of 'cos' or 'linear', " + f"instead got {anneal_strategy}" + ) + self._set_anneal_func(anneal_strategy) + if not isinstance(anneal_epochs, int) or anneal_epochs < 0: + raise ValueError( + f"anneal_epochs must be equal or greater than 0, got {anneal_epochs}" + ) + self.anneal_epochs = anneal_epochs + super().__init__(optimizer, last_epoch) + + @staticmethod + def _linear_anneal(t): + return t + + @staticmethod + def _cosine_anneal(t): + return (1 - math.cos(math.pi * t)) / 2 + + @staticmethod + def _get_initial_lr(lr, swa_lr, alpha): + if alpha == 1: + return swa_lr + return (lr - alpha * swa_lr) / (1 - alpha) + + @override + def get_lr(self): + r"""Compute the next learning rate for each of the optimizer's + :attr:`~torch.optim.Optimizer.param_groups`. + + Uses :attr:`anneal_func` to interpolate between each group's + ``group["lr"]`` and ``group["swa_lr"]`` over :attr:`anneal_epochs` + epochs. Once :attr:`anneal_epochs` is reached, keeps the learning rate + fixed at ``group["swa_lr"]``. + + Returns: + list[float | Tensor]: A :class:`list` of learning rates for each of + the optimizer's :attr:`~torch.optim.Optimizer.param_groups` with the + same types as their current ``group["lr"]``\s. + + .. note:: + If you're trying to inspect the most recent learning rate, use + :meth:`get_last_lr()` instead. + + .. note:: + The returned :class:`~torch.Tensor`\s are copies, and never alias + the optimizer's ``group["lr"]``\s. + """ + # `_get_lr_called_within_step` is only available `_enable_get_lr_call`, + # so we ignore the type error here. See `LRScheduler.step()` for more details. + if not self._get_lr_called_within_step: + warnings.warn( + "To get the last learning rate computed by the scheduler, " + "please use `get_last_lr()`.", + UserWarning, + stacklevel=2, + ) + # Set in `LRScheduler._initial_step()` + step = self._step_count - 1 + if self.anneal_epochs == 0: + step = max(1, step) + # pyrefly: ignore [no-matching-overload] + prev_t = max(0, min(1, (step - 1) / max(1, self.anneal_epochs))) + prev_alpha = self.anneal_func(prev_t) + prev_lrs = [ + self._get_initial_lr(group["lr"], group["swa_lr"], prev_alpha) + for group in self.optimizer.param_groups + ] + # pyrefly: ignore [no-matching-overload] + t = max(0, min(1, step / max(1, self.anneal_epochs))) + alpha = self.anneal_func(t) + return [ + group["swa_lr"] * alpha + lr * (1 - alpha) + for group, lr in zip(self.optimizer.param_groups, prev_lrs, strict=True) + ] + + def _set_anneal_func(self, anneal_strategy: Literal["cos", "linear"]) -> None: + self._anneal_strategy = anneal_strategy + if anneal_strategy == "cos": + self.anneal_func = self._cosine_anneal + else: + self.anneal_func = self._linear_anneal + + @override + def state_dict(self) -> dict[str, Any]: + """Return the state of the scheduler as a :class:`dict`. + + It contains an entry for every variable in self.__dict__ which + is not the optimizer or anneal_func. + """ + return { + key: value + for key, value in self.__dict__.items() + if key not in ("optimizer", "anneal_func") + } + + @override + def load_state_dict(self, state_dict: dict[str, Any]) -> None: + """Load the scheduler's state. + + Args: + state_dict (dict): scheduler state. Should be an object returned + from a call to :meth:`state_dict`. + """ + self.__dict__.update(state_dict) + self._set_anneal_func(self._anneal_strategy) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/overrides.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/overrides.py new file mode 100644 index 0000000000000000000000000000000000000000..67920aef7d853a12d55f5ae3316f2087419e1ae2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/overrides.py @@ -0,0 +1,2166 @@ +""" +Python implementation of ``__torch_function__`` + +While most of the torch API and handling for ``__torch_function__`` happens +at the C++ level, some of the torch API is written in Python so we need +python-level handling for ``__torch_function__`` overrides as well. The main +developer-facing functionality in this file are handle_torch_function and +has_torch_function. See torch/functional.py and test/test_overrides.py +for usage examples. + +Note +---- +heavily inspired by NumPy's ``__array_function__`` (see: +https://github.com/pytorch/pytorch/issues/24015 and +https://www.numpy.org/neps/nep-0018-array-function-protocol.html +) + +If changing this file in a way that can affect ``__torch_function__`` overhead, +please report the benchmarks in ``benchmarks/overrides_benchmark``. See the +instructions in the ``README.md`` in that directory. +""" + +import __future__ # noqa: F404 + +import collections +import contextlib +import functools +import sys +import types +import warnings +from collections.abc import Callable, Iterable +from functools import wraps +from typing import Any, cast, TypeVar +from typing_extensions import ParamSpec + +import torch +from torch._C import ( + _add_docstr, + _get_function_stack_at, + _has_torch_function, + _has_torch_function_unary, + _has_torch_function_variadic, + _is_torch_function_mode_enabled, + _len_torch_function_stack, + _pop_torch_function_stack, + _push_on_torch_function_stack, +) + + +__all__ = [ + "get_ignored_functions", + "get_overridable_functions", + "get_testing_overrides", + "handle_torch_function", + "has_torch_function", + "resolve_name", + "is_tensor_like", + "is_tensor_method_or_property", + "wrap_torch_function", + "enable_reentrant_dispatch", +] + +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +def _disable_user_warnings( + func: Callable[_P, _R], + regex: str = ".*is deprecated, please use.*", + module: str = "torch", +) -> Callable[_P, _R]: + """ + Decorator that temporarily disables ``UserWarning``s for the given ``module`` if the warning message matches the + given ``regex`` pattern. + + Arguments + --------- + func : function + Function to disable the warnings for. + regex : str + A regex pattern compilable by ``re.compile``. This is used to match the ``UserWarning`` message. + module : str + The python module to which the filtering should be restricted. + + Returns + ------- + function + The wrapped function. + """ + + @wraps(func) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", category=UserWarning, message=regex, module=module + ) + return func(*args, **kwargs) + + return wrapper + + +@functools.cache +@_disable_user_warnings +def get_ignored_functions() -> set[Callable]: + """ + Return public functions that cannot be overridden by ``__torch_function__``. + + Returns + ------- + set[Callable] + A tuple of functions that are publicly available in the torch API but cannot + be overridden with ``__torch_function__``. Mostly this is because none of the + arguments of these functions are tensors or tensor-likes. + + Examples + -------- + >>> torch.Tensor.as_subclass in torch.overrides.get_ignored_functions() + True + >>> torch.add in torch.overrides.get_ignored_functions() + False + """ + Tensor = torch.Tensor + functions = { + torch.typename, + torch.is_tensor, + torch.is_storage, + torch.set_default_tensor_type, + torch.set_default_device, + torch.get_default_device, + torch.set_rng_state, + torch.get_rng_state, + torch.manual_seed, + torch.initial_seed, + torch.seed, + torch.thread_safe_generator, + torch.save, + torch.load, + torch.set_printoptions, + torch.fork, + torch.get_default_dtype, + torch.get_num_interop_threads, + torch.get_num_threads, + torch.init_num_threads, + torch.import_ir_module, + torch.import_ir_module_from_buffer, + torch.is_anomaly_enabled, + torch.is_anomaly_check_nan_enabled, + torch.is_grad_enabled, + torch.merge_type_from_type_comment, + torch.parse_ir, + torch.parse_schema, + torch.parse_type_comment, + torch.set_anomaly_enabled, + torch.set_flush_denormal, + torch.set_num_interop_threads, + torch.set_num_threads, + torch.wait, + torch.as_tensor, + torch.from_numpy, + torch.tensor, + torch.default_generator, + torch.has_cuda, + torch.has_cudnn, + torch.has_lapack, + torch.device, + torch.dtype, + torch.finfo, + torch.has_mkl, + torch.has_mps, + torch.has_mkldnn, + torch.has_openmp, + torch.iinfo, + torch.memory_format, + torch.qscheme, + torch.set_grad_enabled, + torch.no_grad, + torch.enable_grad, + torch.inference_mode, + torch.is_inference_mode_enabled, + torch.layout, + torch.align_tensors, + torch.arange, + torch.as_strided, + torch.bartlett_window, + torch.blackman_window, + torch.broadcast_shapes, + torch.can_cast, + torch.compile, + torch.cudnn_affine_grid_generator, + torch.cudnn_batch_norm, + torch.cudnn_convolution, + torch.cudnn_convolution_transpose, + torch.cudnn_convolution_relu, + torch.cudnn_convolution_add_relu, + torch.cudnn_grid_sampler, + torch.cudnn_is_acceptable, + torch.miopen_ctc_loss, + torch.empty, + torch.empty_permuted, + torch.empty_strided, + torch.empty_quantized, + torch.export.export, + torch.export.load, + torch.export.register_dataclass, + torch.export.save, + torch.eye, + torch.fft.fftfreq, + torch.fft.rfftfreq, + torch.from_file, + torch.full, + torch.fill, + torch.hamming_window, + torch.hann_window, + torch.kaiser_window, + torch.linspace, + torch.logspace, + torch.mkldnn_adaptive_avg_pool2d, + torch.mkldnn_convolution, + torch.mkldnn_max_pool2d, + torch.mkldnn_max_pool3d, + torch.mkldnn_linear_backward_weights, + torch.mkldnn_rnn_layer, + torch.normal, + torch.ones, + torch.promote_types, + torch.rand, + torch.rand_like, + torch.randn, + torch.randn_like, + torch.randint, + torch.randint_like, + torch.randperm, + torch.range, + torch.result_type, + torch.scalar_tensor, + torch.sparse_coo_tensor, + torch.sparse_compressed_tensor, + torch.sparse_csr_tensor, + torch.sparse_csc_tensor, + torch.sparse_bsr_tensor, + torch.sparse_bsc_tensor, + torch.sym_constrain_range, + torch.sym_constrain_range_for_size, + torch.sym_fresh_size, + torch.tril_indices, + torch.triu_indices, + torch.vander, + torch.zeros, + torch._jit_internal.boolean_dispatch, + torch.nn.functional.assert_int_or_pair, + torch.nn.functional.upsample, + torch.nn.functional.upsample_bilinear, + torch.nn.functional.upsample_nearest, + torch.nn.functional.has_torch_function, + torch.nn.functional.has_torch_function_unary, + torch.nn.functional.has_torch_function_variadic, + torch.nn.functional.handle_torch_function, + torch.nn.functional.grouped_mm, + torch.nn.functional.scaled_grouped_mm, + torch.nn.functional.scaled_mm, + torch.nn.functional.sigmoid, + torch.nn.functional.hardsigmoid, + torch.nn.functional.tanh, + torch.nn.functional._canonical_mask, + torch.nn.functional._none_or_dtype, + # Doesn't actually take or return tensor arguments + torch.nn.init.calculate_gain, + # These are deprecated; don't test them + torch.nn.init.uniform, + torch.nn.init.normal, + torch.nn.init.constant, + torch.nn.init.eye, + torch.nn.init.dirac, + torch.nn.init.xavier_uniform, + torch.nn.init.xavier_normal, + torch.nn.init.kaiming_uniform, + torch.nn.init.kaiming_normal, + torch.nn.init.orthogonal, + torch.nn.init.sparse, + torch.nested.to_padded_tensor, + has_torch_function, + handle_torch_function, + torch.set_autocast_enabled, + torch.is_autocast_enabled, + torch.set_autocast_dtype, + torch.get_autocast_dtype, + torch.clear_autocast_cache, + torch.set_autocast_cpu_enabled, + torch.is_autocast_cpu_enabled, + torch.set_autocast_xla_enabled, + torch.is_autocast_xla_enabled, + torch.set_autocast_ipu_enabled, + torch.is_autocast_ipu_enabled, + torch.set_autocast_cpu_dtype, + torch.get_autocast_cpu_dtype, + torch.set_autocast_ipu_dtype, + torch.get_autocast_ipu_dtype, + torch.get_autocast_gpu_dtype, + torch.set_autocast_gpu_dtype, + torch.get_autocast_xla_dtype, + torch.set_autocast_xla_dtype, + torch.autocast_increment_nesting, + torch.autocast_decrement_nesting, + torch.is_autocast_cache_enabled, + torch.set_autocast_cache_enabled, + torch.nn.functional.hardswish, + torch.is_vulkan_available, + torch.are_deterministic_algorithms_enabled, + torch.use_deterministic_algorithms, + torch.is_deterministic_algorithms_warn_only_enabled, + torch.set_deterministic_debug_mode, + torch.get_device_module, + torch.get_deterministic_debug_mode, + torch.set_float32_matmul_precision, + torch.get_float32_matmul_precision, + torch.unify_type_list, + torch.is_warn_always_enabled, + torch.set_warn_always, + torch.vmap, + torch.cond, + torch.frombuffer, + torch.asarray, + torch._functional_sym_constrain_range, + torch._make_dep_token, + Tensor.__delitem__, + Tensor.__dir__, + Tensor.__getattribute__, + Tensor.__init__, + Tensor.__iter__, + Tensor.__init_subclass__, + Tensor.__delattr__, + Tensor.__setattr__, + Tensor.__torch_function__, + Tensor.__torch_dispatch__, + Tensor.__new__, + Tensor.__class__, + Tensor.__subclasshook__, + Tensor.__hash__, + Tensor.as_subclass, + Tensor.eig, + Tensor.lstsq, + Tensor.reinforce, + Tensor.new, + Tensor.new_tensor, + Tensor.new_empty, + Tensor.new_empty_strided, + Tensor.new_zeros, + Tensor.new_ones, + Tensor.new_full, + Tensor._make_subclass, + Tensor.solve, + Tensor.symeig, + Tensor.stride, + Tensor.unflatten, + Tensor.to_sparse_coo, + Tensor.to_sparse_csr, + Tensor.to_sparse_csc, + Tensor.to_sparse_bsr, + Tensor.to_sparse_bsc, + Tensor._to_sparse, + Tensor._to_sparse_csr, + Tensor._to_sparse_csc, + Tensor._to_sparse_bsr, + Tensor._to_sparse_bsc, + Tensor._typed_storage, + Tensor._reduce_ex_internal, + Tensor._fix_weakref, + Tensor._view_func, + Tensor._view_func_unsafe, + Tensor._rev_view_func_unsafe, + Tensor._dtensor__new__, + Tensor._make_wrapper_subclass, + Tensor._python_dispatch.__get__, + Tensor._has_symbolic_sizes_strides.__get__, + Tensor._conj, + Tensor._conj_physical, + Tensor._lazy_clone, + Tensor._neg_view, + Tensor._is_zerotensor, + Tensor._is_all_true, + Tensor._is_any_true, + Tensor._addmm_activation, + Tensor.to_padded_tensor, + Tensor._use_count, + Tensor._philox_normal_, + Tensor._philox_uniform_, + } + + if sys.version_info >= (3, 14): + functions.add(Tensor.__annotate__) + + return functions + + +@functools.cache +def get_default_nowrap_functions() -> set[Callable]: + """ + Return public functions that do not wrap in a subclass when invoked by + the default ``Tensor.__torch_function__`` that preserves subclasses. Typically, + these functions represent field accesses (i.e., retrieving a Tensor that + is stored somewhere on the Tensor) as opposed to computation. Users of + these functions expect object identity to be preserved over multiple accesses + (e.g., ``a.grad is a.grad``) which cannot be upheld if we're wrapping on + the fly every time (furthermore, the tensor stored here might already be + the subclass, in which case wrapping really ought not to happen). + + Not ALL property accessors have this property; for example ``Tensor.T`` actually + just creates a new transposed tensor on the fly, and so we SHOULD interpose on + these calls (you need to check the implementation of the function to see if + this is the case or not). Additionally, if a property accessor doesn't return a Tensor, + it doesn't have to be on this list (though it is harmless if it is). + """ + Tensor = torch.Tensor + return { + Tensor._base.__get__, + Tensor.grad.__get__, + Tensor._grad.__get__, + } + + +@functools.cache +@_disable_user_warnings +def get_testing_overrides() -> dict[Callable, Callable]: + """Return a dict containing dummy overrides for all overridable functions + + Returns + ------- + Dict[Callable, Callable] + A dictionary that maps overridable functions in the PyTorch API to + lambda functions that have the same signature as the real function + and unconditionally return -1. These lambda functions are useful + for testing API coverage for a type that defines ``__torch_function__``. + + Examples + -------- + >>> import inspect + >>> my_add = torch.overrides.get_testing_overrides()[torch.add] + >>> inspect.signature(my_add) + + """ + # Every function in the PyTorchAPI that can be overridden needs an entry + # in this dict. + # + # Optimally we would use inspect to get the function signature and define + # the lambda function procedurally but that is blocked by generating + # function signatures for native kernels that can be consumed by inspect. + # See Issue #28233. + Tensor = torch.Tensor + ret: dict[Callable, Callable] = { + torch.abs: lambda input, out=None: -1, + torch.absolute: lambda input, out=None: -1, + torch.adaptive_avg_pool1d: lambda input, output_size: -1, + torch.adaptive_max_pool1d: lambda inputs, output_size: -1, + torch.acos: lambda input, out=None: -1, + torch.adjoint: lambda input: -1, + torch.arccos: lambda input, out=None: -1, + torch.acosh: lambda input, out=None: -1, + torch.arccosh: lambda input, out=None: -1, + torch.add: lambda input, other, out=None: -1, + torch.addbmm: lambda input, batch1, batch2, alpha=1, beta=1, out=None: -1, + torch.addcdiv: lambda input, tensor1, tensor2, value=1, out=None: -1, + torch.addcmul: lambda input, tensor1, tensor2, value=1, out=None: -1, + torch.addmm: lambda input, mat1, mat2, beta=1, alpha=1, out=None: -1, + torch.addmv: lambda input, mat, vec, beta=1, alpha=1, out=None: -1, + torch.addr: lambda input, vec1, vec2, beta=1, alpha=1, out=None: -1, + torch.affine_grid_generator: lambda theta, size, align_corners: -1, + torch.all: lambda input, dim=None: -1, + torch.allclose: lambda input, other, rtol=1e-05, atol=1e-08, equal_nan=False: -1, + torch.alpha_dropout: lambda input, p, train, inplace=False: -1, + torch.amax: lambda input, dim=None: -1, + torch.amin: lambda input, dim=None: -1, + torch.aminmax: lambda input, dim=None, keepdim=False, out=None: -1, + torch.angle: lambda input, out=None: -1, + torch.any: lambda input, dim=None, keepdim=False, out=None: -1, + torch.argmax: lambda input: -1, + torch.argmin: lambda input: -1, + torch.argsort: lambda input, dim=None: -1, + torch.asin: lambda input, out=None: -1, + torch._assert_async: lambda input, msg: -1, + torch.arcsin: lambda input, out=None: -1, + torch.asinh: lambda input, out=None: -1, + torch.arcsinh: lambda input, out=None: -1, + torch.atan: lambda input, out=None: -1, + torch.arctan: lambda input, out=None: -1, + torch.atan2: lambda input, other, out=None: -1, + torch.arctan2: lambda input, other, out=None: -1, + torch.atanh: lambda input, out=None: -1, + torch.arctanh: lambda input, out=None: -1, + torch.atleast_1d: lambda *tensors: -1, + torch.atleast_2d: lambda *tensors: -1, + torch.atleast_3d: lambda *tensors: -1, + torch.avg_pool1d: lambda input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True: -1, + torch.baddbmm: lambda input, batch1, batch2, alpha=1, beta=1, out=None: -1, + torch.batch_norm: lambda input, weight, bias, running_mean, running_var, training, momentum, eps, cudnn_enabled: -1, + torch.batch_norm_backward_elemt: lambda grad_out, input, mean, invstd, weight, sum_dy, sum_dy_xmu, count_tensor: -1, + torch.batch_norm_backward_reduce: lambda grad_out, input, mean, invstd, weight, input_g, weight_g, bias_g: -1, + torch.batch_norm_elemt: lambda input, weight, bias, mean, invstd, eps: -1, + torch.batch_norm_gather_stats: lambda input, mean, invstd, running_mean, running_var, momentum, eps, count: -1, + torch.batch_norm_gather_stats_with_counts: lambda input, mean, invstd, running_mean, running_var, momentum, eps, count: -1, + torch.batch_norm_stats: lambda input, eps: -1, + torch.batch_norm_update_stats: lambda input, running_mean, running_var, momentum: -1, + torch.bernoulli: lambda input, generator=None, out=None: -1, + torch.bilinear: lambda input1, input2, weight, bias: -1, + torch.binary_cross_entropy_with_logits: ( + lambda input, target, weight=None, size_average=None, reduce=None, reduction="mean", pos_weight=None: -1 + ), + torch.bincount: lambda input, weights=None, minlength=0: -1, + torch.binomial: lambda count, prob, generator=None: -1, + torch.bitwise_and: lambda input, other, out=None: -1, + torch.bitwise_not: lambda input, out=None: -1, + torch.bitwise_or: lambda input, other, out=None: -1, + torch.bitwise_xor: lambda input, other, out=None: -1, + torch.bitwise_left_shift: lambda input, other, out=None: -1, + torch.bitwise_right_shift: lambda input, other, out=None: -1, + torch.block_diag: lambda *tensors: -1, + torch.bmm: lambda input, mat2, out_dtype=None, out=None: -1, + torch.broadcast_tensors: lambda *tensors: -1, + torch.broadcast_to: lambda self, size: -1, + torch.bucketize: lambda input, boundaries, out_int32=False, right=False, out=None: -1, + torch.cartesian_prod: lambda *tensors: -1, + torch.cat: lambda tensors, dim=0, out=None: -1, + torch.concat: lambda tensors, dim=0, out=None: -1, # alias for torch.cat + torch.concatenate: lambda tensors, dim=0, out=None: -1, # alias for torch.concatenate + torch.cdist: lambda x1, x2, p=2.0, compute_mode="use_mm_for_euclid_dist_if_necessary": -1, + torch.ceil: lambda input, out=None: -1, + torch.celu: lambda input, alpha=1.0, inplace=False: -1, + torch.chain_matmul: lambda *matrices, out=None: -1, + torch.channel_shuffle: lambda input, groups: -1, + torch.cholesky: lambda input, upper=False, out=None: -1, + torch.linalg.cholesky: lambda input, out=None: -1, + torch.linalg.cholesky_ex: lambda input, check_errors=False, out=None: -1, + torch.cholesky_inverse: lambda input, upper=False, out=None: -1, + torch.cholesky_solve: lambda input1, input2, upper=False, out=None: -1, + torch.choose_qparams_optimized: lambda input, numel, n_bins, ratio, bit_width: -1, + torch.chunk: lambda input, chunks, dim=0: -1, + torch.clamp: lambda input, min=None, max=None, out=None: -1, + torch.clip: lambda input, min=None, max=None, out=None: -1, + torch.clamp_min: lambda input, min, out=None: -1, + torch.clamp_max: lambda input, max, out=None: -1, + torch.column_stack: lambda tensors, out=None: -1, + torch.cov: lambda input, correction=1, fweights=None, aweights=None: -1, + torch.clone: lambda input: -1, + torch.combinations: lambda input, r=2, with_replacement=False: -1, + torch.complex: lambda real, imag: -1, + torch.copysign: lambda input, other, out=None: -1, + torch.polar: lambda abs, ang: -1, + torch.linalg.cond: lambda input, ord=None: -1, + torch.conj: lambda input, out=None: -1, + torch.conj_physical: lambda input, out=None: -1, + torch.resolve_conj: lambda input, out=None: -1, + torch.resolve_neg: lambda input, out=None: -1, + torch.constant_pad_nd: lambda input, pad, value=0: -1, + torch.conv1d: lambda input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1: -1, + torch.conv2d: lambda input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1: -1, + torch.conv3d: lambda input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1: -1, + torch.convolution: lambda input, weight, bias, stride, padding, dilation, transposed, output_adding, groups: -1, + torch.conv_tbc: lambda input, weight, bias, pad=0: -1, + torch.conv_transpose1d: lambda input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1: -1, + torch.conv_transpose2d: lambda input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1: -1, + torch.conv_transpose3d: lambda input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1: -1, + torch.corrcoef: lambda input: -1, + torch.cos: lambda input, out=None: -1, + torch.cosine_embedding_loss: lambda input1, input2, target, margin=0, size_average=None, reduce=None, reduction="mean": -1, + torch.cosh: lambda input, out=None: -1, + torch.cosine_similarity: lambda x1, x2, dim=1, eps=1e-8: -1, + torch.count_nonzero: lambda input: -1, + torch.cross: lambda input, other, dim=None, out=None: -1, + torch.linalg.cross: lambda input, other, dim=-1, out=None: -1, + torch.ctc_loss: ( + lambda log_probs, targets, input_lengths, target_lengths, blank=0, reduction="mean", zero_infinity=False: -1 + ), + torch.cummax: lambda input, dim, out=None: -1, + torch.cummin: lambda input, dim, out=None: -1, + torch.cumprod: lambda input, dim, out=None, dtype=None: -1, + torch.cumsum: lambda input, dim, out=None, dtype=None: -1, + torch.cumulative_trapezoid: lambda y, x=None, dim=-1: -1, + torch.logcumsumexp: lambda input, dim, out=None: -1, + torch.deg2rad: lambda input, out=None: -1, + torch.dequantize: lambda input: -1, + torch.det: lambda input: -1, + torch.linalg.det: lambda input: -1, # alias for torch.det # type: ignore[attr-defined] + torch.detach: lambda input: -1, + torch.diag: lambda input, diagonal=0, out=None: -1, + torch.diag_embed: lambda input, diagonal=0, out=None: -1, + torch.diagflat: lambda input, offset=0: -1, + torch.diff: lambda input, n=1, dim=-1, prepend=None, append=None, out=None: -1, + torch.diagonal: lambda input, offset=0, dim1=0, dim2=1: -1, + torch.linalg.diagonal: lambda input, offset=0, dim1=-2, dim2=-1: -1, + torch.diagonal_scatter: lambda input, src, offset=0, dim1=0, dim2=1: -1, + torch.as_strided_scatter: lambda self, src, size, stride, storage_offset=None: -1, + torch.digamma: lambda input, out=None: -1, + torch.dist: lambda input, other, p=2: -1, + torch.div: lambda input, other, rounding_mode=None, out=None: -1, + torch.divide: lambda input, other, rounding_mode=None, out=None: -1, + torch.dot: lambda input, other, out=None: -1, + torch.dropout: lambda input, p, train, inplace=False: -1, + torch.dsmm: lambda input, mat2, out_dtype=None: -1, + torch.hsmm: lambda mat1, mat2: -1, + torch.dsplit: lambda input, indices_or_sections: -1, + torch.dstack: lambda tensors, out=None: -1, + torch.linalg.eig: lambda input, out=None: -1, + torch.linalg.eigvals: lambda input, out=None: -1, + torch.linalg.eigh: lambda input, UPLO="L", out=None: -1, + torch.linalg.eigvalsh: lambda input, UPLO="L", out=None: -1, + torch.einsum: lambda equation, *operands: -1, + torch.embedding: ( + lambda input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False: -1 # noqa: B950 + ), + torch.embedding_bag: ( + lambda input, weight, offsets, max_norm=None, norm_type=2, scale_grad_by_freq=False, mode="mean", sparse=False, per_sample_weights=None, padding_idx=None: -1 # noqa: B950 + ), + torch.empty_like: lambda input, dtype=None, layout=None, device=None, requires_grad=False: -1, + torch.eq: lambda input, other, out=None: -1, + torch.equal: lambda input, other: -1, + torch.erf: lambda input, out=None: -1, + torch.erfc: lambda input, out=None: -1, + torch.erfinv: lambda input, out=None: -1, + torch.exp: lambda input, out=None: -1, + torch.exp2: lambda input, out=None: -1, + torch.expm1: lambda input, out=None: -1, + torch.fake_quantize_per_channel_affine: lambda input, scale, zero_point, axis, quant_min, quant_max: -1, + torch.fake_quantize_per_tensor_affine: lambda input, scale, zero_point, quant_min, quant_max: -1, + torch.fused_moving_avg_obs_fake_quant: ( + lambda x, observer_on, fake_quant_on, averaging_const, running_min, running_max, scale, zero_point, quant_min, quant_max, ch_axis, per_row_fake_quant=False, symmetric_quant=False: -1 # noqa: B950 + ), + torch.fbgemm_linear_fp16_weight: lambda input, packed_weight, bias, output: -1, + torch.fbgemm_linear_fp16_weight_fp32_activation: lambda input, packed_weight, bias, output: -1, + torch.fbgemm_linear_int8_weight: lambda input, weight, packed, col_offsets, weight_scale, weight_zero_point, bias: -1, # noqa: B950 + torch.fbgemm_linear_int8_weight_fp32_activation: ( + lambda input, weight, packed, col_offsets, weight_scale, weight_zero_point, bias: -1 + ), + torch.fbgemm_linear_quantize_weight: lambda input: -1, + torch.fbgemm_pack_gemm_matrix_fp16: lambda input: -1, + torch.fbgemm_pack_quantized_matrix: lambda input, a, b: -1, + torch.feature_alpha_dropout: lambda input, p, train: -1, + torch.feature_dropout: lambda input, p, train: -1, + torch.fft.ifft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.rfft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.irfft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.hfft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.ihfft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fft.hfft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.ihfft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.hfftn: lambda input, s=None, dim=-1, norm=None: -1, + torch.fft.ihfftn: lambda input, s=None, dim=-1, norm=None: -1, + torch.fft.fftn: lambda input, s=None, dim=None, norm=None: -1, + torch.fft.ifftn: lambda input, s=None, dim=None, norm=None: -1, + torch.fft.rfftn: lambda input, s=None, dim=None, norm=None: -1, + torch.fft.irfftn: lambda input, s=None, dim=None, norm=None: -1, + torch.fft.fft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.ifft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.rfft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.irfft2: lambda input, s=None, dim=(-2, -1), norm=None: -1, + torch.fft.fftshift: lambda input, dim=None: -1, + torch.fft.ifftshift: lambda input, dim=None: -1, + torch.fft.fft: lambda input, n=None, dim=-1, norm=None: -1, + torch.fix: lambda input, out=None: -1, + torch.flatten: lambda input, start_dim=0, end_dim=-1: -1, + torch.flip: lambda input, dims: -1, + torch.fliplr: lambda input: -1, + torch.flipud: lambda input: -1, + torch.frobenius_norm: lambda input, dim=None, keepdim=False, out=None: -1, + torch.floor: lambda input, out=None: -1, + torch.floor_divide: lambda input, other: -1, + torch.float_power: lambda input, exponent, out=None: -1, + torch.fmod: lambda input, other, out=None: -1, + torch.frac: lambda input, out=None: -1, + torch.frexp: lambda input, out=None: -1, + torch.full_like: lambda input, fill_value, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False: -1, # noqa: B950 + torch._functional_assert_async: lambda input, msg, dep_token: -1, + torch.lu_unpack: lambda LU_data, LU_pivots, unpack_data=True, unpack_pivots=True: -1, + torch.gather: lambda input, dim, index, out=None, sparse_grad=False: -1, + torch.gcd: lambda input, other, out=None: -1, + torch.ge: lambda input, other, out=None: -1, + torch.get_device: lambda input: -1, + torch.greater_equal: lambda input, other, out=None: -1, + torch.geqrf: lambda input, out=None: -1, + torch.i0: lambda input, out=None: -1, + torch.inner: lambda input, other, out=None: -1, + torch.outer: lambda input, vec2, out=None: -1, + torch.ger: lambda input, vec2, out=None: -1, # alias for torch.outer + torch.gradient: lambda input, spacing=None, dim=None, edge_order=1: -1, + torch.grid_sampler: lambda input, grid, interpolation_mode, padding_mode, align_corners: -1, + torch.grid_sampler_2d: lambda input, grid, interpolation_mode, padding_mode, align_corners: -1, + torch.grid_sampler_3d: lambda input, grid, interpolation_mode, padding_mode, align_corners: -1, + torch.group_norm: lambda input, num_groups, weight=None, bias=None, eps=1e-05, cudnn_enabled=True: -1, + torch.gru: lambda input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first: -1, + torch.gru_cell: lambda input, hx, w_ih, w_hh, b_ih=None, b_hh=None: -1, + torch.gt: lambda input, other, out=None: -1, + torch.greater: lambda input, other, out=None: -1, + torch.hardshrink: lambda input, lambd=0.5: -1, + torch.hash_tensor: lambda input, dim=None, keepdim=False, mode=0, out=None: -1, + torch.heaviside: lambda input, values, out=None: -1, + torch.hinge_embedding_loss: lambda input, target, margin=1.0, size_average=None, reduce=None, reduction="mean": -1, # noqa: B950 + torch.histc: lambda input, bins=100, min=0, max=0, out=None: -1, + torch.histogram: lambda input, bins=100, min=None, max=None, weight=None, density=False, out=None: -1, + torch.histogramdd: lambda input, bins, range=None, weight=None, density=False: -1, + torch.linalg.householder_product: lambda input, tau: -1, + torch.hspmm: lambda mat1, mat2, out=None: -1, + torch.hsplit: lambda input, indices_or_sections: -1, + torch.hstack: lambda tensors, out=None: -1, + torch.hypot: lambda input, other, out=None: -1, + torch.igamma: lambda input, other, out=None: -1, + torch.igammac: lambda input, other, out=None: -1, + torch.imag: lambda input, out=None: -1, + torch.index_add: lambda input, dim, index, source: -1, + torch.index_copy: lambda input, dim, index, source: -1, + torch.index_put: lambda input, indices, values, accumulate=False: -1, + torch.index_select: lambda input, dim, index, out=None: -1, + torch.index_fill: lambda input, dim, index, value: -1, + torch.index_reduce: lambda input, dim, index, source, reduce, include_input=True: -1, + torch.isfinite: lambda tensor: -1, + torch.isin: lambda e, te, assume_unique=False, invert=False: -1, + torch.isinf: lambda tensor: -1, + torch.isreal: lambda tensor: -1, + torch.isposinf: lambda input, out=None: -1, + torch.isneginf: lambda input, out=None: -1, + torch.instance_norm: ( + lambda input, running_mean, running_var, weight, bias, use_input_stats, momentum, eps, cudnn_enabled: -1 + ), + torch.int_repr: lambda input: -1, + torch.inverse: lambda input, out=None: -1, + torch.linalg.inv: lambda input, out=None: -1, + torch.linalg.inv_ex: lambda input, check_errors=False, out=None: -1, + torch.is_complex: lambda input: -1, + torch.is_conj: lambda input: -1, + torch.is_neg: lambda input: -1, + torch.is_distributed: lambda input: -1, + torch.is_inference: lambda input: -1, + torch.is_floating_point: lambda input: -1, + torch.is_nonzero: lambda input: -1, + torch.is_same_size: lambda input, other: -1, + torch.is_signed: lambda input: -1, + torch.isclose: lambda input, other, rtol=1e-05, atol=1e-08, equal_nan=False: -1, + torch.isnan: lambda input: -1, + torch.istft: ( + lambda input, n_fft, hop_length=None, win_length=None, window=None, center=True, normalized=False, onesided=None, length=None, return_complex=False: -1 # noqa: B950 + ), + torch.kl_div: lambda input, target, size_average=None, reduce=None, reduction="mean", log_target=False: -1, + torch.kron: lambda input, other: -1, + torch.kthvalue: lambda input, k, dim=None, keepdim=False, out=None: -1, + torch.linalg.ldl_factor_ex: lambda input, hermitian=False, check_errors=False, out=None: -1, + torch.linalg.ldl_factor: lambda input, hermitian=False, out=None: -1, + torch.linalg.ldl_solve: lambda LD, pivots, B, hermitian=False, out=None: -1, + torch.layer_norm: lambda input, normalized_shape, weight=None, bias=None, eps=1e-05, cudnn_enabled=True: -1, + torch.lcm: lambda input, other, out=None: -1, + torch.ldexp: lambda input, other, out=None: -1, + torch.le: lambda input, other, out=None: -1, + torch.less_equal: lambda input, other, out=None: -1, + torch.lerp: lambda input, end, weight, out=None: -1, + torch.lgamma: lambda input, out=None: -1, + torch.lobpcg: lambda input, k=None, B=None, X=None, n=None, iK=None, niter=None, tol=None, largest=None, method=None, tracker=None, ortho_iparams=None, ortho_fparams=None, ortho_bparams=None: -1, # noqa: B950 + torch.log: lambda input, out=None: -1, + torch.log_softmax: lambda input, dim, dtype=None: -1, + torch.log10: lambda input, out=None: -1, + torch.log1p: lambda input, out=None: -1, + torch.log2: lambda input, out=None: -1, + torch.logaddexp: lambda input, other, out=None: -1, + torch.logaddexp2: lambda input, other, out=None: -1, + torch.logdet: lambda input: -1, + torch.xlogy: lambda x, y, out=None: -1, + torch.logical_and: lambda input, other, out=None: -1, + torch.logical_not: lambda input, out=None: -1, + torch.logical_or: lambda input, other, out=None: -1, + torch.logical_xor: lambda input, other, out=None: -1, + torch.logit: lambda input, eps=None: -1, + torch.logsumexp: lambda input, names, keepdim=False, out=None: -1, + torch.lstm: lambda data, batch_sizes, hx, params, has_biases, num_layers, dropout, train, bidirectional: -1, + torch.lstm_cell: lambda input, hx, w_ih, w_hh, b_ih=None, b_hh=None: -1, + torch.lt: lambda input, other, out=None: -1, + torch.less: lambda input, other, out=None: -1, + torch.lu: lambda A, pivot=True, get_infos=False, out=None: -1, + torch.lu_solve: lambda b, LU_data, LU_pivots, out=None: -1, + torch.margin_ranking_loss: lambda input1, input2, target, margin=0, size_average=None, reduce=None, reduction="mean": -1, # type: ignore[attr-defined] # noqa: B950 + torch.masked_fill: lambda input, mask, value: -1, + torch.masked_scatter: lambda input, mask, source: -1, + torch.masked_select: lambda input, mask, out=None: -1, + torch.matmul: lambda input, other, out=None: -1, + torch.linalg.lu: lambda input, pivot=True, out=None: -1, + torch.linalg.lu_factor: lambda input, pivot=True, out=None: -1, + torch.linalg.lu_factor_ex: lambda input, pivot=True, check_errors=False, out=None: -1, + torch.linalg.lu_solve: lambda LU, pivots, B, left=True, adjoint=False, out=None: -1, + torch.linalg.matmul: lambda input, other, out=None: -1, # alias for torch.matmul + torch.matrix_power: lambda input, n: -1, + torch.linalg.matrix_power: lambda input, n, out=None: -1, + torch.linalg.matrix_rank: lambda input, tol=None, hermitian=False: -1, + torch.linalg.multi_dot: lambda tensors, out=None: -1, + torch.matrix_exp: lambda input: -1, + torch.linalg.matrix_exp: lambda input: -1, + torch.max: lambda input, out=None: -1, + torch.maximum: lambda input, other, out=None: -1, + torch.fmax: lambda input, other, out=None: -1, + torch.max_pool1d: lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False: -1, + torch.max_pool2d: lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False: -1, + torch.max_pool3d: lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False: -1, + torch.max_pool1d_with_indices: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.mean: lambda input, dim=None: -1, + torch.nanmean: lambda input, dim=None, keepdim=False, dtype=None, out=None: -1, + torch.median: lambda input, dim=None: -1, + torch.nanmedian: lambda input, dim=None: -1, + torch.meshgrid: lambda *tensors, **kwargs: -1, + torch.min: lambda input, out=None: -1, + torch.minimum: lambda input, other, out=None: -1, + torch.fmin: lambda input, other, out=None: -1, + torch.miopen_batch_norm: ( + lambda input, weight, bias, running_mean, running_var, training, exponential_average_factor, epsilon: -1 + ), + torch.miopen_convolution: lambda input, weight, bias, padding, stride, dilation, groups, benchmark, deterministic: -1, # noqa: B950 + torch.miopen_convolution_add_relu: lambda input, weight, z, alpha, bias, stride, padding, dilation, groups: -1, + torch.miopen_convolution_relu: lambda input, weight, bias, stride, padding, dilation, groups: -1, + torch.miopen_convolution_transpose: ( + lambda input, weight, bias, padding, output_padding, stride, dilation, groups, benchmark, deterministic: -1 + ), + torch.miopen_depthwise_convolution: ( + lambda input, weight, bias, padding, stride, dilation, groups, benchmark, deterministic: -1 + ), + torch.miopen_rnn: ( + lambda input, weight, weight_stride0, hx, cx, mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state: -1 # noqa: B950 + ), + torch.mm: lambda input, mat2, out_dtype=None, out=None: -1, + torch.mode: lambda input, dim=-1, keepdim=False, out=None: -1, + torch.movedim: lambda input, source, destination: -1, + torch.moveaxis: lambda input, source, destination: -1, + torch.msort: lambda input, descending=False, out=None: -1, + torch.mul: lambda input, other, out=None: -1, + torch.multiply: lambda input, other, out=None: -1, + torch.multinomial: lambda input, num_samples, replacement=False, out=None: -1, + torch.mv: lambda input, vec, out=None: -1, + torch.mvlgamma: lambda input, p: -1, + torch.narrow: lambda input, dim, start, length: -1, + torch.nan_to_num: lambda input, nan=0.0, posinf=None, neginf=None, out=None: -1, + torch.native_batch_norm: lambda input, weight, bias, running_mean, running_var, training, momentum, eps: -1, + torch._native_batch_norm_legit: lambda input, weight, bias, training, momentum, eps: -1, + torch.native_dropout: lambda input, p, train: -1, + torch.native_layer_norm: lambda input, normalized_shape, weight=None, bias=None, eps=1e-05: -1, + torch._fused_rms_norm: lambda input, normalized_shape, weight=None, eps=1e-05: -1, + torch.native_group_norm: lambda input, weight, bias, N, C, HxW, group, eps: -1, + torch.native_norm: lambda input, p=2, dim=None, keepdim=False, dtype=None: -1, + torch.native_channel_shuffle: lambda input, groups: -1, + torch.ne: lambda input, other, out=None: -1, + torch.not_equal: lambda input, other, out=None: -1, + torch.neg: lambda input, out=None: -1, + torch.negative: lambda input, out=None: -1, + torch.nextafter: lambda input, other, out=None: -1, + torch.nn.functional.adaptive_avg_pool2d: lambda input, output_size: -1, + torch.nn.functional.adaptive_avg_pool3d: lambda input, output_size: -1, + torch.nn.functional.adaptive_max_pool1d: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool1d_with_indices: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool2d: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool2d_with_indices: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool3d: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.adaptive_max_pool3d_with_indices: lambda input, output_size, return_indices=False: -1, + torch.nn.functional.affine_grid: lambda theta, size, align_corners=None: -1, + torch.nn.functional.alpha_dropout: lambda input, p=0.5, training=False, inplace=False: -1, + torch.nn.functional.avg_pool2d: ( + lambda input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None: -1 # noqa: B950 + ), + torch.nn.functional.avg_pool3d: ( + lambda input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None: -1 # noqa: B950 + ), + torch.nn.functional.batch_norm: ( + lambda input, running_mean, running_var, weight=None, bias=None, training=False, momentum=0.1, eps=1e-05: -1 + ), + torch.nn.functional.bilinear: lambda input1, input2, weight, bias=None: -1, + torch.nn.functional.binary_cross_entropy: ( + lambda input, target, weight=None, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.binary_cross_entropy_with_logits: ( + lambda input, target, weight=None, size_average=None, reduce=None, reduction="mean", pos_weight=None: -1 + ), + torch.nn.functional.celu: lambda input, alpha=1.0, inplace=False: -1, + torch.nn.functional.cosine_embedding_loss: ( + lambda input1, input2, target, margin=0, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.cross_entropy: ( + lambda input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="mean", label_smoothing=0.0: -1 # noqa: B950 + ), + torch.nn.functional.ctc_loss: ( + lambda log_probs, targets, input_lengths, target_lengths, blank=0, reduction="mean", zero_infinity=False: -1 + ), + torch.nn.functional.dropout: lambda input, p=0.5, training=True, inplace=False: -1, + torch.nn.functional.dropout1d: lambda input, p=0.5, training=True, inplace=False: -1, + torch.nn.functional.dropout2d: lambda input, p=0.5, training=True, inplace=False: -1, + torch.nn.functional.dropout3d: lambda input, p=0.5, training=True, inplace=False: -1, + torch.nn.functional.elu: lambda input, alpha=1.0, inplace=False: -1, + torch.nn.functional.embedding: ( + lambda input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False: -1 # noqa: B950 + ), + torch.nn.functional.embedding_bag: ( + lambda input, weight, offsets=None, max_norm=None, norm_type=2, scale_grad_by_freq=False, mode="mean", sparse=False, per_sample_weights=None, include_last_offset=False, padding_idx=None: -1 # noqa: B950 + ), + torch.nn.functional.feature_alpha_dropout: lambda input, p=0.5, training=False, inplace=False: -1, + torch.nn.functional.fold: lambda input, output_size, kernel_size, dilation=1, padding=0, stride=1: -1, + torch.nn.functional.fractional_max_pool2d: ( + lambda input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None: -1 # noqa: B950 + ), + torch.nn.functional.fractional_max_pool2d_with_indices: ( + lambda input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None: -1 # noqa: B950 + ), + torch.nn.functional.fractional_max_pool3d: ( + lambda input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None: -1 # noqa: B950 + ), + torch.nn.functional.fractional_max_pool3d_with_indices: ( + lambda input, kernel_size, output_size=None, output_ratio=None, return_indices=False, _random_samples=None: -1 # noqa: B950 + ), + torch.nn.functional.gaussian_nll_loss: lambda input, target, var, full=False, eps=1e-06, reduction="mean": -1, + torch.nn.functional.gelu: lambda input, approximate="none": -1, + torch.nn.functional.glu: lambda input, dim=-1: -1, + torch.nn.functional.grid_sample: lambda input, grid, mode="bilinear", padding_mode="zeros", align_corners=None: -1, # noqa: B950 + torch.nn.functional.group_norm: lambda input, num_groups, weight=None, bias=None, eps=1e-05: -1, + torch.nn.functional.gumbel_softmax: lambda logits, tau=1, hard=False, eps=1e-10, dim=-1: -1, + torch.nn.functional.hardshrink: lambda input, lambd=0.5: -1, + torch.nn.functional.hardtanh: lambda input, min_val=-1.0, max_val=1.0, inplace=False: -1, + torch.nn.functional.hinge_embedding_loss: ( + lambda input, target, margin=1.0, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.instance_norm: ( + lambda input, running_mean=None, running_var=None, weight=None, bias=None, use_input_stats=True, momentum=0.1, eps=1e-05: -1 # noqa: B950 + ), + torch.nn.functional.interpolate: ( + lambda input, size=None, scale_factor=None, mode="nearest", align_corners=None, recompute_scale_factor=None, antialias=False: -1 # noqa: B950 + ), + torch.nn.functional.kl_div: lambda input, target, size_average=None, reduce=None, reduction="mean", log_target=False: -1, # noqa: B950 + torch.nn.functional.l1_loss: lambda input, target, size_average=None, reduce=None, reduction="mean", weight=None: -1, + torch.nn.functional.layer_norm: lambda input, normalized_shape, weight=None, bias=None, eps=1e-05: -1, + torch.nn.functional.leaky_relu: lambda input, negative_slope=0.01, inplace=False: -1, + torch.nn.functional.linear: lambda input, weight, bias=None: -1, + torch.nn.functional.local_response_norm: lambda input, size, alpha=0.0001, beta=0.75, k=1.0: -1, + torch.nn.functional.log_softmax: lambda input, dim=None, _stacklevel=3, dtype=None: -1, + torch.nn.functional.logsigmoid: lambda input: -1, + torch.nn.functional.lp_pool1d: lambda input, norm_type, kernel_size, stride=None, ceil_mode=False: -1, + torch.nn.functional.lp_pool2d: lambda input, norm_type, kernel_size, stride=None, ceil_mode=False: -1, + torch.nn.functional.lp_pool3d: lambda input, norm_type, kernel_size, stride=None, ceil_mode=False: -1, + torch.nn.functional.margin_ranking_loss: ( + lambda input1, input2, target, margin=0, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.max_pool1d: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False: -1 + ), + torch.nn.functional.max_pool1d_with_indices: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.nn.functional.max_pool2d: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False: -1 + ), + torch.nn.functional.max_pool2d_with_indices: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.nn.functional.max_pool3d: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.nn.functional.max_pool3d_with_indices: ( + lambda input, kernel_size, stride=None, padding=0, dilation=1, return_indices=False, ceil_mode=False: -1 + ), + torch.nn.functional.max_unpool1d: lambda input, indices, kernel_size, stride=None, padding=0, output_size=None: -1, # noqa: B950 + torch.nn.functional.max_unpool2d: lambda input, indices, kernel_size, stride=None, padding=0, output_size=None: -1, # noqa: B950 + torch.nn.functional.max_unpool3d: lambda input, indices, kernel_size, stride=None, padding=0, output_size=None: -1, # noqa: B950 + torch.nn.functional.mse_loss: lambda input, target, size_average=None, reduce=None, reduction="mean", weight=None: -1, + torch.nn.functional.multi_head_attention_forward: ( + lambda query, key, value, embed_dim_to_check, num_heads, in_proj_weight, in_proj_bias, bias_k, bias_v, add_zero_attn, dropout_p, out_proj_weight, out_proj_bias, training=True, key_padding_mask=None, need_weights=True, attn_mask=None, use_separate_proj_weight=False, q_proj_weight=None, k_proj_weight=None, v_proj_weight=None, static_k=None, static_v=None, average_attn_weights=None, is_causal=False: -1 # noqa: B950 + ), + torch.nn.functional.multi_margin_loss: ( + lambda input, target, p=1, margin=1.0, weight=None, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.multilabel_margin_loss: ( + lambda input, target, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.multilabel_soft_margin_loss: ( + lambda input, target, weight=None, size_average=None, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.nll_loss: ( + lambda input, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction="mean": -1 + ), + torch.nn.functional.normalize: lambda input, p=2, dim=1, eps=1e-12, out=None: -1, + torch.nn.functional.one_hot: lambda tensor, num_classes=-1: -1, + torch.nn.functional.pad: lambda input, pad, mode="constant", value=0: -1, + torch.nn.functional.pairwise_distance: lambda x1, x2, p=2.0, eps=1e-06, keepdim=False: -1, + torch.nn.functional.poisson_nll_loss: ( + lambda input, target, log_input=True, full=False, size_average=None, eps=1e-08, reduce=None, reduction="mean": -1 # noqa: B950 + ), + torch.nn.functional.prelu: lambda input, weight: -1, + torch.nn.functional.relu: lambda input, inplace=False: -1, + torch.nn.functional.relu6: lambda input, inplace=False: -1, + torch.nn.functional.rms_norm: lambda input, normalized_shape, weight=None, eps=1e-6: -1, + torch.nn.functional.rrelu: lambda input, lower=0.125, upper=0.3333333333333333, training=False, inplace=False: -1, # noqa: B950 + torch.nn.functional.selu: lambda input, inplace=False: -1, + torch.nn.functional.silu: lambda input, inplace=False: -1, + torch.nn.functional.mish: lambda input, inplace=False: -1, + torch.nn.functional.scaled_dot_product_attention: lambda query, key, value, attn_mask=None, dropout_p=0.0: -1, + torch.nn.functional.smooth_l1_loss: lambda input, target, size_average=None, reduce=None, reduction="mean", beta=1.0: -1, # noqa: B950 + torch.nn.functional.huber_loss: lambda input, target, reduction="mean", delta=1.0, weight=None: -1, + torch.nn.functional.soft_margin_loss: lambda input, target, size_average=None, reduce=None, reduction="mean": -1, # noqa: B950 + torch.nn.functional.softmax: lambda input, dim=None, _stacklevel=3, dtype=None: -1, + torch.nn.functional.softmin: lambda input, dim=None, _stacklevel=3, dtype=None: -1, + torch.nn.functional.softplus: lambda input, beta=1, threshold=20: -1, + torch.nn.functional.softshrink: lambda input, lambd=0.5: -1, + torch.nn.functional.softsign: lambda input: -1, + torch.nn.functional.tanhshrink: lambda input: -1, + torch.nn.functional.threshold: lambda input, threshold, value, inplace=False: -1, + torch.nn.functional.triplet_margin_loss: ( + lambda anchor, positive, negative, margin=1.0, p=2, eps=1e-06, swap=False, size_average=None, reduce=None, reduction="mean": -1 # noqa: B950 + ), + torch.nn.functional.triplet_margin_with_distance_loss: ( + lambda anchor, positive, negative, *, distance_function=None, margin=1.0, swap=False, reduction="mean": -1 + ), + torch.nn.functional.unfold: lambda input, kernel_size, dilation=1, padding=0, stride=1: -1, + torch.nn.init.uniform_: lambda tensor, a=0.0, b=1.0, generator=None: -1, + torch.nn.init.normal_: lambda tensor, mean=0.0, std=1.0, generator=None: -1, + torch.nn.init.constant_: lambda tensor, val: -1, + torch.nn.init.kaiming_uniform_: lambda tensor, a=0, mode="fan_in", nonlinearity="leaky_relu", generator=None: -1, # noqa: B950 + torch.nonzero: lambda input, as_tuple=False: -1, + torch.nonzero_static: lambda input, *, size, fill_value=-1: -1, + torch.argwhere: lambda input: -1, + torch.norm: lambda input, p="fro", dim=None, keepdim=False, out=None, dtype=None: -1, + torch.linalg.norm: lambda input, ord=None, dim=None, keepdim=False, out=None, dtype=None: -1, + torch.linalg.vector_norm: lambda input, ord=2, dim=None, keepdim=False, out=None, dtype=None: -1, + torch.linalg.matrix_norm: lambda input, ord="fro", dim=( + -2, + -1, + ), keepdim=False, out=None, dtype=None: -1, + torch.norm_except_dim: lambda v, pow=2, dim=0: -1, + torch.nuclear_norm: lambda input, p="fro", dim=None, keepdim=False, out=None, dtype=None: -1, + torch.numel: lambda input: -1, + torch.orgqr: lambda input, tau: -1, + torch.ormqr: lambda input, input2, input3, left=True, transpose=False: -1, + torch.pairwise_distance: lambda x1, x2, p=2.0, eps=1e-06, keepdim=False: -1, + torch.permute: lambda self, dim: -1, + torch.pca_lowrank: lambda input, q=None, center=True, niter=2: -1, + torch.pdist: lambda input, p=2: -1, + torch.pinverse: lambda input, rcond=1e-15: -1, + torch.linalg.pinv: lambda input, rcond=1e-15, hermitian=False: -1, + torch.pixel_shuffle: lambda input, upscale_factor: -1, + torch.pixel_unshuffle: lambda input, downscale_factor: -1, + torch.poisson: lambda input, generator=None: -1, + torch.poisson_nll_loss: lambda input, target, log_input, full, eps, reduction: -1, + torch.polygamma: lambda input, n, out=None: -1, + torch.positive: lambda input, out=None: -1, + torch.prelu: lambda input, weight: -1, + torch.ones_like: lambda input, dtype=None, layout=None, device=None, requires_grad=False: -1, + torch.pow: lambda input, exponent, out=None: -1, + torch.prod: lambda input, dtype=None: -1, + torch.put: lambda input, index, source, accumulate=False: -1, + torch.q_per_channel_axis: lambda input: -1, + torch.q_per_channel_scales: lambda input: -1, + torch.q_per_channel_zero_points: lambda input: -1, + torch.q_scale: lambda input: -1, + torch.q_zero_point: lambda input: -1, + torch.qr: lambda input, some=True, out=None: -1, + torch.linalg.qr: lambda input, mode="reduced", out=None: -1, + torch.quantile: lambda input, q, dim=None, keepdim=False, interpolation="linear", out=None: -1, + torch.nanquantile: lambda input, q, dim=None, keepdim=False, interpolation="linear", out=None: -1, + torch.quantize_per_channel: lambda input, scales, zero_points, axis, dtype: -1, + torch.quantize_per_tensor: lambda input, scale, zero_point, dtype: -1, + torch.quantize_per_tensor_dynamic: lambda input, dtype, reduce_range: -1, + torch.quantized_batch_norm: lambda input, weight, bias, mean, var, eps, output_scale, output_zero_point: -1, + torch.quantized_gru_cell: ( + lambda input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh: -1 # noqa: B950 + ), + torch.quantized_lstm_cell: ( + lambda input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh: -1 # noqa: B950 + ), + torch.quantized_max_pool1d: ( + lambda input, kernel_size, stride=(), padding=(0,), dilation=( + 1, + ), ceil_mode=False: -1 + ), + torch.quantized_max_pool2d: ( + lambda input, kernel_size, stride=(), padding=(0, 0), dilation=( + 1, + 1, + ), ceil_mode=False: -1 + ), + torch.quantized_max_pool3d: ( + lambda input, kernel_size, stride=(), padding=(0, 0, 0), dilation=( + 1, + 1, + 1, + ), ceil_mode=False: -1 + ), + torch.quantized_rnn_relu_cell: ( + lambda input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh: -1 # noqa: B950 + ), + torch.quantized_rnn_tanh_cell: ( + lambda input, hx, w_ih, w_hh, b_ih, b_hh, packed_ih, packed_hh, col_offsets_ih, col_offsets_hh, scale_ih, scale_hh, zero_point_ih, zero_point_hh: -1 # noqa: B950 + ), + torch.rad2deg: lambda input, out=None: -1, + torch.ravel: lambda input: -1, + torch.real: lambda input, out=None: -1, + torch.vdot: lambda input, other, out=None: -1, + torch.linalg.vecdot: lambda input, other, dim=-1, out=None: -1, + torch.view_as_real: lambda input: -1, + torch.view_as_complex: lambda input: -1, + torch.reciprocal: lambda input, out=None: -1, + torch.relu: lambda input, inplace=False: -1, + torch.remainder: lambda input, other, out=None: -1, + torch.renorm: lambda input, p, dim, maxnorm, out=None: -1, + torch.repeat_interleave: lambda input, dim=None: -1, + torch.reshape: lambda input, shape: -1, + torch.rms_norm: lambda input, normalized_shape, weight=None, eps=1e-6: -1, + torch.rnn_relu: lambda input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first: -1, # noqa: B950 + torch.rnn_relu_cell: lambda input, hx, w_ih, w_hh, b_ih=None, b_hh=None: -1, + torch.rnn_tanh: lambda input, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first: -1, # noqa: B950 + torch.rnn_tanh_cell: lambda input, hx, w_ih, w_hh, b_ih=None, b_hh=None: -1, + torch.roll: lambda input, shifts, dims=None: -1, + torch.rot90: lambda input, k=1, dims=(0, 1): -1, + torch.round: lambda input, out=None: -1, + torch.row_stack: lambda tensors, out=None: -1, # alias for torch.vstack + torch._rowwise_prune: (lambda weight, mask, compressed_indices_dtype: -1), + torch.rrelu: lambda input, lower=1.0 / 8, upper=1.0 / 3, training=False, inplace=False: -1, + torch.rsqrt: lambda input, out=None: -1, + torch.rsub: lambda input, other, alpha=1: -1, + torch.saddmm: lambda input, mat1, mat2, beta=1, alpha=1, out=None: -1, + torch.scatter: lambda input, dim, index, src: -1, + torch.scatter_add: lambda input, dim, index, src: -1, + torch.scatter_reduce: lambda input, dim, index, src, reduce, include_self=True: -1, + torch.searchsorted: lambda sorted_sequence, input, out_int32=False, right=False, out=None: -1, + torch._segment_reduce: lambda data, reduce="max", lengths=None, indices=None, offsets=None, axis=0, unsafe=False: -1, # noqa: B950 + torch.select: lambda input, dim, index: -1, + torch.select_scatter: lambda input, src, dim, index: -1, + torch.slice_inverse: lambda input, src, dim=0, start=None, end=None, step=1: -1, + torch.slice_scatter: lambda input, src, dim=0, start=None, end=None, step=1: -1, + torch.selu: lambda input, inplace=False: -1, + torch.sigmoid: lambda input, out=None: -1, + torch.sign: lambda input, out=None: -1, + torch.signbit: lambda input, out=None: -1, + torch.sgn: lambda input, out=None: -1, + torch.sin: lambda input, out=None: -1, + torch.sinc: lambda input, out=None: -1, + torch.sinh: lambda input, out=None: -1, + torch.slogdet: lambda input: -1, + torch.linalg.slogdet: lambda input: -1, + torch.smm: lambda input, mat2, out_dtype=None: -1, + torch.spmm: lambda input, mat2, out_dtype=None: -1, + torch.softmax: lambda input, dim, dtype=None: -1, + torch.linalg.solve: lambda A, B, left=True, out=None: -1, + torch.linalg.solve_ex: lambda A, B, left=True, check_errors=False, out=None: -1, + torch.sort: lambda input, dim=-1, descending=False, *, stable=False, out=None: -1, + torch.split: lambda tensor, split_size_or_sections, dim=0: -1, + torch.split_with_sizes: lambda tensor, split_size_or_sections, dim=0: -1, + torch.sqrt: lambda input, out=None: -1, + torch.square: lambda input, out=None: -1, + torch.squeeze: lambda input, dim=None, out=None: -1, + torch.sspaddmm: lambda input, mat1, mat2, beta=1, alpha=1, out=None: -1, + torch.stack: lambda tensors, dim=0, out=None: -1, + torch.std: lambda input, dim=None: -1, + torch.std_mean: lambda input, dim=None: -1, + torch.stft: ( + lambda input, n_fft, hop_length=None, win_length=None, window=None, center=True, pad_mode="reflect", normalized=False, onesided=True, return_complex=None, align_to_window=None: -1 # noqa: B950 + ), + torch.sub: lambda input, other, out=None: -1, + torch.subtract: lambda input, other, out=None: -1, + torch.sum: lambda input, dim=None: -1, + torch.sym_float: lambda input: -1, + torch.sym_int: lambda input: -1, + torch.sym_max: lambda a, b: -1, + torch.sym_min: lambda a, b: -1, + torch.sym_not: lambda input: -1, + torch.sym_ite: lambda a, b, c: -1, + torch.sym_sum: lambda args: -1, + torch._sym_sqrt: lambda input: -1, + torch._sym_cos: lambda input: -1, + torch._sym_cosh: lambda input: -1, + torch._sym_sin: lambda input: -1, + torch._sym_sinh: lambda input: -1, + torch._sym_tan: lambda input: -1, + torch._sym_tanh: lambda input: -1, + torch._sym_asin: lambda input: -1, + torch._sym_acos: lambda input: -1, + torch._sym_atan: lambda input: -1, + torch.nansum: lambda input, dim=None: -1, + torch.svd: lambda input, some=True, compute_uv=True, out=None: -1, + torch.svd_lowrank: lambda input, q=6, niter=2, M=None: -1, + torch.linalg.svd: lambda input, full_matrices=True, out=None: -1, + torch.linalg.svdvals: lambda input, out=None: -1, + torch.swapaxes: lambda input, dim0, dim1: -1, + torch.swapdims: lambda input, axis0, axis1: -1, + torch.special.airy_ai: lambda input: -1, + torch.special.bessel_j0: lambda input: -1, + torch.special.bessel_j1: lambda input: -1, + torch.special.bessel_y0: lambda input: -1, + torch.special.bessel_y1: lambda input: -1, + torch.special.chebyshev_polynomial_t: lambda input, n, out=None: -1, + torch.special.chebyshev_polynomial_u: lambda input, n, out=None: -1, + torch.special.chebyshev_polynomial_v: lambda input, n, out=None: -1, + torch.special.chebyshev_polynomial_w: lambda input, n, out=None: -1, + torch.special.digamma: lambda input: -1, + torch.special.entr: lambda input: -1, + torch.special.erf: lambda input: -1, + torch.special.erfc: lambda input: -1, + torch.special.erfcx: lambda input: -1, + torch.special.erfinv: lambda input: -1, + torch.special.exp2: lambda input: -1, + torch.special.expit: lambda input: -1, + torch.special.expm1: lambda input: -1, + torch.special.gammainc: lambda input, other, out=None: -1, + torch.special.gammaincc: lambda input, other, out=None: -1, + torch.special.gammaln: lambda input: -1, + torch.special.hermite_polynomial_h: lambda input, n, out=None: -1, + torch.special.hermite_polynomial_he: lambda input, n, out=None: -1, + torch.special.i0: lambda input: -1, + torch.special.i0e: lambda input: -1, + torch.special.i1: lambda input: -1, + torch.special.i1e: lambda input: -1, + torch.special.laguerre_polynomial_l: lambda input, n, out=None: -1, + torch.special.legendre_polynomial_p: lambda input, n, out=None: -1, + torch.special.log1p: lambda input: -1, + torch.special.log_ndtr: lambda input: -1, + torch.special.log_softmax: lambda input, dim, dtype=None: -1, + torch.special.logit: lambda input: -1, + torch.special.logsumexp: lambda input, dim, keepdim=False, out=None: -1, + torch.special.modified_bessel_i0: lambda input: -1, + torch.special.modified_bessel_i1: lambda input: -1, + torch.special.modified_bessel_k0: lambda input: -1, + torch.special.modified_bessel_k1: lambda input: -1, + torch.special.multigammaln: lambda input, p: -1, + torch.special.ndtr: lambda input: -1, + torch.special.ndtri: lambda input: -1, + torch.special.polygamma: lambda input, n, out=None: -1, + torch.special.psi: lambda input: -1, + torch.special.round: lambda input: -1, + torch.special.scaled_modified_bessel_k0: lambda input: -1, + torch.special.scaled_modified_bessel_k1: lambda input: -1, + torch.special.shifted_chebyshev_polynomial_t: lambda input, n, out=None: -1, + torch.special.shifted_chebyshev_polynomial_u: lambda input, n, out=None: -1, + torch.special.shifted_chebyshev_polynomial_v: lambda input, n, out=None: -1, + torch.special.shifted_chebyshev_polynomial_w: lambda input, n, out=None: -1, + torch.special.sinc: lambda input: -1, + torch.special.softmax: lambda input, dim, dtype=None: -1, + torch.special.spherical_bessel_j0: lambda input: -1, + torch.special.xlog1py: lambda input, other, out=None: -1, + torch.special.xlogy: lambda input, other, out=None: -1, + torch.special.zeta: lambda self, other, out=None: -1, + torch.t: lambda input: -1, + torch.take: lambda input, index: -1, + torch.take_along_dim: lambda input, indices, dim=None, out=None: -1, + torch.tan: lambda input, out=None: -1, + torch.tanh: lambda input, out=None: -1, + torch.linalg.tensorinv: lambda a, ind=2: -1, + torch.linalg.tensorsolve: lambda a, b, dims=None: -1, + torch.tensordot: lambda a, b, dims=2, out=None: -1, + torch.tensor_split: lambda input, indices_or_sections, dim=0: -1, + torch.threshold: lambda input, threshold, value, inplace=False: -1, + torch.tile: lambda input, dims: -1, + torch.topk: lambda input, k, dim=-1, descending=False, out=None: -1, + torch.trace: lambda input: -1, + torch.transpose: lambda input, dim0, dim1: -1, + torch.trapz: lambda y, x=None, dim=-1: -1, + torch.trapezoid: lambda y, x=None, dim=-1: -1, + torch.triangular_solve: lambda input, A, upper=True, transpose=False, unitriangular=False: -1, + torch.linalg.solve_triangular: lambda input, B, upper, left=True, unitriangular=False: -1, + torch.tril: lambda input, diagonal=0, out=None: -1, + torch.triplet_margin_loss: ( + lambda anchor, positive, negative, margin=1.0, p=2, eps=1e-06, swap=False, size_average=None, reduce=None, reduction="mean": -1 # noqa: B950 + ), + torch.triu: lambda input, diagonal=0, out=None: -1, + torch.true_divide: lambda input, other: -1, + torch.trunc: lambda input, out=None: -1, + torch.unbind: lambda input, dim=0: -1, + torch.unflatten: lambda input, dim, sizes, names: -1, + torch.unique: lambda input, sorted=True, return_inverse=False, return_counts=False, dim=None: -1, + torch.unique_consecutive: lambda input, return_inverse=False, return_counts=False, dim=None: -1, + torch.unravel_index: lambda indices, shape: -1, + torch.unsafe_chunk: lambda input, chunks, dim=0: -1, + torch.unsafe_split: lambda tensor, split_size_or_sections, dim=0: -1, + torch.unsafe_split_with_sizes: lambda tensor, split_size_or_sections, dim=0: -1, + torch.unsqueeze: lambda input, dim, out=None: -1, + torch.linalg.vander: lambda x, N=None: -1, + torch.var: lambda input, dim=None: -1, + torch.var_mean: lambda input, dim=None: -1, + torch.vsplit: lambda input, indices_or_sections: -1, + torch.vstack: lambda tensors, out=None: -1, + torch.where: lambda condition, x=None, y=None: -1, + torch._wrapped_linear_prepack: lambda weight, weight_scale, weight_zero_point, bias : -1, + torch._wrapped_quantized_linear_prepacked: ( + lambda input, input_scale, input_zero_point, prepacked, out_scale, out_zero_point, out_channel : -1 # noqa: B950 + ), + torch.zeros_like: lambda input, dtype=None, layout=None, device=None, requires_grad=False: -1, + torch._fw_primal_copy: lambda self, level: -1, + torch._make_dual_copy: lambda primal, tangent, level: -1, + torch.view_as_real_copy: lambda self: -1, + torch.view_as_complex_copy: lambda self: -1, + torch._conj_copy: lambda self: -1, + torch._neg_view_copy: lambda self: -1, + torch.as_strided_copy: lambda self, size, stride, storage_offset=None: -1, + torch._sparse_broadcast_to_copy: lambda self, size: -1, + torch.diagonal_copy: lambda self, offset=0, dim1=0, dim2=1: -1, + torch.expand_copy: lambda self, size, *, implicit=False: -1, + torch.narrow_copy: lambda self, dim, start, length: -1, + torch.permute_copy: lambda self, dims: -1, + torch._reshape_alias_copy: lambda self, size, stride: -1, + torch.select_copy: lambda self, dim, index: -1, + torch.detach_copy: lambda self: -1, + torch.slice_copy: lambda self, dim=0, start=None, end=None, step=1: -1, + torch.split_copy: lambda self, split_size, dim=0: -1, + torch.split_with_sizes_copy: lambda self, split_sizes, dim=0: -1, + torch.squeeze_copy: lambda self, dim: -1, + torch.t_copy: lambda self: -1, + torch.transpose_copy: lambda self, dim0, dim1: -1, + torch.unsqueeze_copy: lambda self, dim: -1, + torch._indices_copy: lambda self: -1, + torch._values_copy: lambda self: -1, + torch.indices_copy: lambda self: -1, + torch.values_copy: lambda self: -1, + torch.crow_indices_copy: lambda self: -1, + torch.col_indices_copy: lambda self: -1, + torch.ccol_indices_copy: lambda self: -1, + torch.row_indices_copy: lambda self: -1, + torch.unbind_copy: lambda self, dim=0: -1, + torch.view_copy: lambda self, dtype: -1, + torch.unfold_copy: lambda self, dimension, size, step: -1, + torch.alias_copy: lambda self: -1, + Tensor.__floordiv__: lambda self, other: -1, + Tensor.__rfloordiv__: lambda self, other: -1, + Tensor.__ifloordiv__: lambda self, other: -1, + Tensor.__truediv__: lambda self, other: -1, + Tensor.__rtruediv__: lambda self, other: -1, + Tensor.__itruediv__: lambda self, other: -1, + Tensor.__lshift__: lambda self, other: -1, + Tensor.__rlshift__: lambda self, other: -1, + Tensor.__ilshift__: lambda self, other: -1, + Tensor.__rshift__: lambda self, other: -1, + Tensor.__rrshift__: lambda self, other: -1, + Tensor.__irshift__: lambda self, other: -1, + Tensor.__and__: lambda self, other: -1, + Tensor.__or__: lambda self, other: -1, + Tensor.__xor__: lambda self, other: -1, + Tensor.__float__: lambda self: -1, + Tensor.__complex__: lambda self: -1, + Tensor.__array__: lambda self, dtype: -1, + Tensor.__bool__: lambda self: -1, + Tensor.__contains__: lambda self, other: -1, + Tensor.__neg__: lambda self: -1, + Tensor.__invert__: lambda self: -1, + Tensor.__mod__: lambda self, other: -1, + Tensor.__rmod__: lambda self, other: -1, + Tensor.__imod__: lambda self, other: -1, + Tensor.__array_wrap__: lambda self, array: -1, + Tensor.__getitem__: lambda self, idx: -1, + Tensor.__deepcopy__: lambda self, memo: -1, + Tensor.__int__: lambda self: -1, + Tensor.__long__: lambda self: -1, + Tensor.__index__: lambda self: -1, + Tensor.__len__: lambda self: -1, + Tensor.__format__: lambda self, format_spec: -1, + Tensor.__reduce_ex__: lambda self, proto: -1, + Tensor.__reversed__: lambda self: -1, + Tensor.__repr__: lambda self, *, tensor_contents=None: -1, + Tensor.__setitem__: lambda self, k, v: -1, + Tensor.__setstate__: lambda self, d: -1, + Tensor.T.__get__: lambda self: -1, + Tensor.H.__get__: lambda self: -1, + Tensor.mT.__get__: lambda self: -1, + Tensor.mH.__get__: lambda self: -1, + Tensor._backward_hooks.__get__: lambda self: -1, + Tensor._post_accumulate_grad_hooks.__get__: lambda self: -1, + Tensor._base.__get__: lambda self: -1, + Tensor._cdata.__get__: lambda self: -1, + Tensor.grad.__get__: lambda self: -1, + Tensor._grad.__get__: lambda self: -1, + Tensor._grad_fn.__get__: lambda self: -1, + Tensor.grad_fn.__get__: lambda self: -1, + Tensor.grad_dtype.__get__: lambda self: -1, + Tensor._version.__get__: lambda self: -1, + Tensor._autocast_to_reduced_precision: lambda self, cuda_enabled, cpu_enabled, cuda_dtype, cpu_dtype: -1, + Tensor._autocast_to_full_precision: lambda self, cuda_enabled, cpu_enabled: -1, + Tensor._clear_non_serializable_cached_data: lambda self: -1, + Tensor.data.__get__: lambda self: -1, + Tensor.device.__get__: lambda self: -1, + Tensor.dtype.__get__: lambda self: -1, + Tensor.is_cuda.__get__: lambda self: -1, + Tensor.is_cpu.__get__: lambda self: -1, + Tensor.is_xla.__get__: lambda self: -1, + Tensor.is_xpu.__get__: lambda self: -1, + Tensor.is_ipu.__get__: lambda self: -1, + Tensor.is_leaf.__get__: lambda self: -1, + Tensor.retains_grad.__get__: lambda self: -1, + Tensor.is_meta.__get__: lambda self: -1, + Tensor.is_mps.__get__: lambda self: -1, + Tensor.is_mtia.__get__: lambda self: -1, + Tensor.is_nested.__get__: lambda self: -1, + Tensor.is_maia.__get__: lambda self: -1, + Tensor.is_mkldnn.__get__: lambda self: -1, + Tensor.is_quantized.__get__: lambda self: -1, + Tensor.is_sparse.__get__: lambda self: -1, + Tensor.is_sparse_csr.__get__: lambda self: -1, + Tensor.is_vulkan.__get__: lambda self: -1, + Tensor.itemsize.__get__: lambda self: -1, + Tensor.layout.__get__: lambda self: -1, + Tensor.name.__get__: lambda self: -1, + Tensor.names.__get__: lambda self: -1, + Tensor.nbytes.__get__: lambda self: -1, + Tensor.ndim.__get__: lambda self: -1, + Tensor.output_nr.__get__: lambda self: -1, + Tensor.requires_grad.__get__: lambda self: -1, + Tensor.shape.__get__: lambda self: -1, + Tensor.volatile.__get__: lambda self: -1, + Tensor.real.__get__: lambda self: -1, + Tensor.imag.__get__: lambda self: -1, + Tensor.__cuda_array_interface__.__get__: lambda self: -1, + Tensor.type: lambda self, dtype=None, non_blocking=False, **kwargs: -1, + Tensor._dimI: lambda self: -1, + Tensor._dimV: lambda self: -1, + Tensor._indices: lambda self: -1, + Tensor._is_view: lambda self: -1, + Tensor._nnz: lambda self: -1, + Tensor.crow_indices: lambda self: -1, + Tensor.col_indices: lambda self: -1, + Tensor.ccol_indices: lambda self: -1, + Tensor.row_indices: lambda self: -1, + Tensor._update_names: lambda self, names, inplace: -1, + Tensor._values: lambda self: -1, + Tensor.adjoint: lambda self: -1, + Tensor.align_as: lambda self, other: -1, + Tensor.align_to: lambda self, order, ellipsis_idx: -1, + Tensor.apply_: lambda self, callable: -1, + Tensor.as_strided: lambda self, size, stride: -1, + Tensor.as_strided_: lambda self, size, stride: -1, + Tensor.backward: lambda self, gradient=None, retain_graph=None, create_graph=False, inputs=None: -1, + Tensor.bfloat16: lambda self, memory_format=torch.preserve_format: -1, + Tensor.bool: lambda self, memory_format=torch.preserve_format: -1, + Tensor.byte: lambda self, memory_format=torch.preserve_format: -1, + Tensor.char: lambda self, memory_format=torch.preserve_format: -1, + Tensor.cauchy_: lambda self, median=0, sigma=1, *, generator=None: -1, + Tensor.coalesce: lambda self: -1, + Tensor._coalesced_: lambda self, coalesced: -1, + Tensor.contiguous: lambda self, memory_format=torch.contiguous_format: -1, + Tensor.copy_: lambda self, src, non_blocking=False: -1, + Tensor.cpu: lambda self, memory_format=torch.preserve_format: -1, + Tensor.cuda: lambda self, memory_format=torch.preserve_format: -1, + Tensor.mtia: lambda self, memory_format=torch.preserve_format: -1, + Tensor.xpu: lambda self, memory_format=torch.preserve_format: -1, + Tensor.ipu: lambda self, memory_format=torch.preserve_format: -1, + Tensor.data_ptr: lambda self: -1, + Tensor.dense_dim: lambda self: -1, + Tensor.diagonal_scatter: lambda self, src, offset=0, dim1=0, dim2=1: -1, + Tensor.dim: lambda self: -1, + Tensor.dim_order: lambda self, ambiguity_check=False: -1, + Tensor.double: lambda self, memory_format=torch.preserve_format: -1, + Tensor.cdouble: lambda self, memory_format=torch.preserve_format: -1, + Tensor.element_size: lambda self: -1, + Tensor.expand: lambda self, size: -1, + Tensor.expand_as: lambda self, other: -1, + Tensor.exponential_: lambda self, lambd=1, *, generator=None: -1, + Tensor.fill_: lambda self, value: -1, + Tensor.fill_diagonal_: lambda self, value: -1, + Tensor.float: lambda self, memory_format=torch.preserve_format: -1, + Tensor.cfloat: lambda self, memory_format=torch.preserve_format: -1, + Tensor.geometric_: lambda self, p, *, generator=None: -1, + Tensor.get_device: lambda self: -1, + Tensor.half: lambda self, memory_format=torch.preserve_format: -1, + Tensor.chalf: lambda self, memory_format=torch.preserve_format: -1, + Tensor.has_names: lambda self: -1, + Tensor.indices: lambda self: -1, + Tensor.int: lambda self, memory_format=torch.preserve_format: -1, + Tensor.is_coalesced: lambda self: -1, + Tensor.is_contiguous: lambda self: -1, + Tensor.is_inference: lambda self: -1, + Tensor.is_pinned: lambda self: -1, + Tensor.is_set_to: lambda self, tensor: -1, + Tensor.is_shared: lambda self: -1, + Tensor.item: lambda self: -1, + Tensor.log_normal_: lambda self, mean=1, std=2, *, generator=None: -1, + Tensor.log_softmax: lambda self, dim: -1, + Tensor.long: lambda self, memory_format=torch.preserve_format: -1, + Tensor.map_: lambda self, tensor, callable: -1, + Tensor.map2_: lambda self, x, y, callable: -1, + Tensor.mm: lambda self, mat2, out_dtype=None: -1, + Tensor.module_load: lambda self, other, assign=False: -1, + Tensor.narrow_copy: lambda self, dimension, start, length: -1, + Tensor.ndimension: lambda self: -1, + Tensor.nelement: lambda self: -1, + Tensor._nested_tensor_size: lambda self: -1, + Tensor._nested_tensor_storage_offsets: lambda self: -1, + Tensor._nested_tensor_strides: lambda self: -1, + Tensor.normal_: lambda self: -1, + Tensor.numpy: lambda self: -1, + Tensor.permute: lambda self, dim: -1, + Tensor.pin_memory: lambda self: -1, + Tensor.put_: lambda self, indices, tensor, accumulate=False: -1, + Tensor.qscheme: lambda self: -1, + Tensor.random_: lambda self, from_=0, to=None, *, generator=None: -1, + Tensor.record_stream: lambda self, stream: -1, + Tensor.refine_names: lambda self, names: -1, + Tensor.register_hook: lambda self, hook: -1, + Tensor.register_post_accumulate_grad_hook: lambda self, hook: -1, + Tensor.rename: lambda self, name: -1, + Tensor.repeat: lambda self, *size: -1, + Tensor.requires_grad_: lambda self, requires_grad=True: -1, + Tensor.reshape_as: lambda self, other: -1, + Tensor.resize: lambda self, *size: -1, + Tensor.resize_: lambda self, size: -1, + Tensor.resize_as: lambda self, other: -1, + Tensor.resize_as_sparse_: lambda self, other: -1, + Tensor.retain_grad: lambda self: -1, + Tensor.set_: lambda self, source=None, storage_offset=0, size=None, stride=None: -1, + Tensor.select_scatter: lambda self, src, dim, index: -1, + Tensor.share_memory_: lambda self: -1, + Tensor.short: lambda self, memory_format=torch.preserve_format: -1, + Tensor.size: lambda self: -1, + Tensor.slice_scatter: lambda self, src, dim=0, start=None, end=None, step=1: -1, + Tensor.sparse_dim: lambda self: -1, + Tensor.sparse_mask: lambda self, mask: -1, + Tensor._sparse_mask_projection: lambda self, mask, accumulate_matches=False: -1, + Tensor.sparse_resize_: lambda self, size1, size2, dense_dim: -1, + Tensor.sparse_resize_and_clear_: lambda self, size1, size2, dense_dim: -1, + Tensor.sspaddmm: lambda self, mat1, mat2, beta=1, alpha=1, out=None: -1, + Tensor.storage: lambda self: -1, + Tensor.untyped_storage: lambda self: -1, + Tensor.storage_offset: lambda self: -1, + Tensor.storage_type: lambda self: -1, + Tensor.sum_to_size: lambda self, size: -1, + Tensor.tile: lambda self, *reps: -1, + Tensor.to: lambda self, dtype, non_blocking=False, copy=False, memory_format=torch.preserve_format: -1, + Tensor.to_dense: lambda self, dtype=None, *, masked_grad=None: -1, + Tensor._to_dense: lambda self, dtype=None, masked_grad=None: -1, + Tensor.to_sparse: lambda self: -1, + Tensor.tolist: lambda self: -1, + Tensor.to_mkldnn: lambda self: -1, + Tensor.type_as: lambda self, other: -1, + Tensor.unfold: lambda self, dimension, size, step: -1, + Tensor.uniform_: lambda self, from_=0, to=1: -1, + Tensor.values: lambda self: -1, + Tensor.view: lambda self, shape: -1, + Tensor.view_as: lambda self, other: -1, + Tensor.zero_: lambda self: -1, + Tensor.__dlpack__: lambda self, stream=None, max_version=None, dl_device=None, copy=None: -1, + Tensor.__dlpack_device__: lambda self: -1, + Tensor.index: lambda self, a, b: -1, + torch.linalg.lstsq: lambda self, b, cond=None, driver=None: -1, + } # fmt: skip + + privateuse1_backend_name = ( + torch.utils.backend_registration._privateuse1_backend_name + ) + if hasattr(Tensor, privateuse1_backend_name): + ret[getattr(Tensor, privateuse1_backend_name)] = ( + lambda self, device=None, non_blocking=False, **kwargs: -1 + ) + ret[getattr(Tensor, f"is_{privateuse1_backend_name}").__get__] = lambda self: -1 + + ret2 = {} + ignored = get_ignored_functions() + + for k, v in ret.items(): + # Generate methods like __add__ and add_ by default from add + names = [ + k.__name__, # Default method + k.__name__ + "_", # Inplace variant + "__" + k.__name__ + "__", # Dunder method + "__i" + k.__name__ + "__", # Inplace dunder method + "__r" + k.__name__ + "__", # Reverse dunder method + ] + + if k.__name__.startswith("bitwise_"): + # bitwise_ have dunder methods of the form ____ + # And so on. + subname = k.__name__[len("bitwise_") :] + names.extend( + ["__" + subname + "__", "__i" + subname + "__", "__r" + subname + "__"] + ) + + for name in names: + func = getattr(Tensor, name, None) + if callable(func) and func not in ret and func not in ignored: + ret2[func] = v + + ret.update(ret2) + + # Distributed functions are added after the auto-generation loop above + # to avoid generating spurious Tensor method entries (e.g., dist.reduce + # would otherwise generate __reduce__ on Tensor). + if torch.distributed.is_available(): + import torch.distributed as dist + + ret.update( + { + dist.broadcast: lambda tensor, src=None, group=None, async_op=False, group_src=None: -1, + dist.all_reduce: lambda tensor, op=None, group=None, async_op=False: -1, + dist.reduce: lambda tensor, dst=None, op=None, group=None, async_op=False, group_dst=None: -1, + dist.all_reduce_coalesced: lambda tensors, op=None, group=None, async_op=False: -1, + dist.all_gather: lambda tensor_list, tensor, group=None, async_op=False: -1, + dist.all_gather_into_tensor: lambda output_tensor, input_tensor, group=None, async_op=False: -1, + dist.all_gather_coalesced: lambda output_tensor_lists, input_tensor_list, group=None, async_op=False: -1, + dist.gather: lambda tensor, gather_list=None, dst=None, group=None, async_op=False, group_dst=None: -1, + dist.scatter: lambda tensor, scatter_list=None, src=None, group=None, async_op=False, group_src=None: -1, + dist.reduce_scatter: lambda output, input_list, op=None, group=None, async_op=False: -1, + dist.reduce_scatter_tensor: lambda output, input, op=None, group=None, async_op=False: -1, + dist.all_to_all_single: lambda output, input, output_split_sizes=None, input_split_sizes=None, group=None, async_op=False: -1, # noqa: B950 + dist.all_to_all: lambda output_tensor_list, input_tensor_list, group=None, async_op=False: -1, + dist.isend: lambda tensor, dst=None, group=None, tag=0, group_dst=None: -1, + dist.irecv: lambda tensor, src=None, group=None, tag=0, group_src=None: -1, + dist.send: lambda tensor, dst=None, group=None, tag=0, group_dst=None: -1, + dist.recv: lambda tensor, src=None, group=None, tag=0, group_src=None: -1, + } + ) # fmt: skip + + return ret + + +def wrap_torch_function( + dispatcher: Callable[_P, Iterable[Any]], +) -> Callable[[Callable[_P, _R]], Callable[_P, _R]]: + """Wraps a given function with ``__torch_function__`` -related functionality. + + Parameters + ---------- + dispatcher: Callable + A callable that returns an iterable of Tensor-likes passed into the function. + + Note + ---- + This decorator may reduce the performance of your code. Generally, it's enough to express + your code as a series of functions that, themselves, support __torch_function__. If you + find yourself in the rare situation where this is not the case, e.g. if you're wrapping a + low-level library and you also need it to work for Tensor-likes, then this function is available. + + Examples + -------- + >>> def dispatcher(a): # Must have the same signature as func + ... return (a,) + >>> @torch.overrides.wrap_torch_function(dispatcher) + >>> def func(a): # This will make func dispatchable by __torch_function__ + ... return a + 0 + """ + + def inner(func: Callable[_P, _R]) -> Callable[_P, _R]: + @functools.wraps(func) + def wrapped(*args: _P.args, **kwargs: _P.kwargs) -> _R: + relevant_args = dispatcher(*args, **kwargs) + if has_torch_function(relevant_args): + return handle_torch_function( + cast(Callable[_P, _R], wrapped), relevant_args, *args, **kwargs + ) + + return func(*args, **kwargs) + + return cast(Callable[_P, _R], wrapped) + + return inner + + +def _get_overloaded_args( + relevant_args: Iterable[Any], + get_type_fn: Callable[[Any], type] | None = None, +) -> list[Any]: + """Returns a list of arguments on which to call __torch_function__. + + Checks arguments in relevant_args for __torch_function__ implementations, + storing references to the arguments and their types in overloaded_args and + overloaded_types in order of calling precedence. Only distinct types are + considered. If a type is a subclass of another type it will have higher + precedence, otherwise the precedence order is the same as the order of + arguments in relevant_args, that is, from left-to-right in the argument list. + + The precedence-determining algorithm implemented in this function is + described in `NEP-0018`_. + + See torch::append_overloaded_arg for the equivalent function in the C++ + implementation. + + Parameters + ---------- + relevant_args : iterable of array-like + Iterable of array-like arguments to check for __torch_function__ + methods. + + get_type_fn : callable, optional + Function to call on each argument in relevant_args to get its type. + + Returns + ------- + overloaded_args : list + Arguments from relevant_args on which to call __torch_function__ + methods, in the order in which they should be called. + + .. _NEP-0018: + https://numpy.org/neps/nep-0018-array-function-protocol.html + """ + if get_type_fn is None: + get_type_fn = type + + # If torch function is not enabled, there are no overloaded types + if not torch._C._is_torch_function_enabled(): + return [] + # Runtime is O(num_arguments * num_unique_types) + overloaded_types: set[type] = set() + overloaded_args: list[Any] = [] + for arg in relevant_args: + arg_type = get_type_fn(arg) + # We only collect arguments if they have a unique type, which ensures + # reasonable performance even with a long list of possibly overloaded + # arguments. + # + # NB: Important to exclude _disabled_torch_function_impl, otherwise + # https://github.com/pytorch/pytorch/issues/64687 + if ( + arg_type not in overloaded_types + and hasattr(arg_type, "__torch_function__") + and arg_type.__torch_function__ + is not torch._C._disabled_torch_function_impl + ): + # Create lists explicitly for the first type (usually the only one + # done) to avoid setting up the iterator for overloaded_args. + if overloaded_types: + overloaded_types.add(arg_type) + # By default, insert argument at the end, but if it is + # subclass of another argument, insert it before that argument. + # This ensures "subclasses before superclasses". + index = len(overloaded_args) + for i, old_arg in enumerate(overloaded_args): + if issubclass(arg_type, get_type_fn(old_arg)): + index = i + break + overloaded_args.insert(index, arg) + else: + overloaded_types = {arg_type} + overloaded_args = [arg] + return overloaded_args + + +def handle_torch_function( + public_api: Callable[_P, _R], + relevant_args: Iterable[Any], + *args: _P.args, + **kwargs: _P.kwargs, +) -> _R: + """Implement a function with checks for ``__torch_function__`` overrides. + + See torch::autograd::handle_torch_function for the equivalent of this + function in the C++ implementation. + + Arguments + --------- + public_api : function + Function exposed by the public torch API originally called like + ``public_api(*args, **kwargs)`` on which arguments are now being + checked. + relevant_args : iterable + Iterable of arguments to check for __torch_function__ methods. + args : tuple + Arbitrary positional arguments originally passed into ``public_api``. + kwargs : tuple + Arbitrary keyword arguments originally passed into ``public_api``. + + Returns + ------- + object + Result from calling ``implementation`` or an ``__torch_function__`` + method, as appropriate. + + Raises + ------ + TypeError : if no implementation is found. + + Example + ------- + >>> def func(a): + ... if has_torch_function_unary(a): + ... return handle_torch_function(func, (a,), a) + ... return a + 0 + """ + # Check for __torch_function__ methods. + overloaded_args = _get_overloaded_args(relevant_args) + # overloaded_args already have unique types. + types = tuple(map(type, overloaded_args)) + + # Check for __torch_function__ mode. + if _is_torch_function_mode_enabled(): + # if we're here, the mode must be set to a TorchFunctionStackMode + # this unsets it and calls directly into TorchFunctionStackMode's torch function + with _pop_mode_temporarily() as mode: + result = mode.__torch_function__(public_api, types, args, kwargs) + if result is not NotImplemented: + return result + + # Call overrides + for overloaded_arg in overloaded_args: + # This call needs to become a classmethod call in the future. + # See https://github.com/pytorch/pytorch/issues/63767 + torch_func_method = overloaded_arg.__torch_function__ + if ( + hasattr(torch_func_method, "__self__") + and torch_func_method.__self__ is overloaded_arg + and torch_func_method is not torch._C._disabled_torch_function_impl + ): + warnings.warn( + "Defining your `__torch_function__ as a plain method is deprecated and " + "will be an error in future, please define it as a classmethod.", + DeprecationWarning, + stacklevel=2, + ) + + # Use `public_api` instead of `implementation` so __torch_function__ + # implementations can do equality/identity comparisons. + result = torch_func_method(public_api, types, args, kwargs) + + if result is not NotImplemented: + return result + + func_name = f"{public_api.__module__}.{public_api.__name__}" + msg = ( + f"no implementation found for '{func_name}' on types that implement " + f"__torch_function__: {[type(arg) for arg in overloaded_args]}" + ) + if _is_torch_function_mode_enabled(): + msg += f" nor in mode {_get_current_function_mode()}" + raise TypeError(msg) + + +has_torch_function = _add_docstr( + _has_torch_function, + r"""Check for __torch_function__ implementations in the elements of an iterable + or if a __torch_function__ mode is enabled. Considers exact ``Tensor`` s + and ``Parameter`` s non-dispatchable. Use this to guard a call to + :func:`handle_torch_function`; don't use it to test if something + is Tensor-like, use :func:`is_tensor_like` instead. + Arguments + --------- + relevant_args : iterable + Iterable or arguments to check for __torch_function__ methods. + Returns + ------- + bool + True if any of the elements of relevant_args have __torch_function__ + implementations, False otherwise. + See Also + ________ + torch.is_tensor_like + Checks if something is a Tensor-like, including an exact ``Tensor``. + """, +) + +has_torch_function_unary = _add_docstr( + _has_torch_function_unary, + r"""Special case of `has_torch_function` for single inputs. + Instead of: + `has_torch_function((t,))` + call: + `has_torch_function_unary(t)` + which skips unnecessary packing and unpacking work. + """, +) + +has_torch_function_variadic = _add_docstr( + _has_torch_function_variadic, + r"""Special case of `has_torch_function` that skips tuple creation. + + This uses the METH_FASTCALL protocol introduced in Python 3.7 + + Instead of: + `has_torch_function((a, b))` + call: + `has_torch_function_variadic(a, b)` + which skips unnecessary packing and unpacking work. + """, +) + + +@functools.cache +def _get_overridable_functions() -> tuple[ + dict[Any, list[Callable]], dict[Callable, str] +]: + overridable_funcs = collections.defaultdict(list) + index = {} + tested_namespaces = [ + ("torch", torch, torch.__all__), + ("torch.functional", torch.functional, torch.functional.__all__), + ("torch.nn.functional", torch.nn.functional, dir(torch.nn.functional)), + ("torch.nn.init", torch.nn.init, dir(torch.nn.init)), + ("torch.Tensor", torch.Tensor, dir(torch.Tensor)), + ("torch.linalg", torch.linalg, dir(torch.linalg)), + ("torch.fft", torch.fft, dir(torch.fft)), + ("torch.special", torch.special, dir(torch.special)), + ] + for namespace_str, namespace, ns_funcs in tested_namespaces: + for func_name in ns_funcs: + ignore = False + # ignore private functions or functions that are deleted in torch.__init__ + if namespace is not torch.Tensor: + if func_name.startswith("__"): + continue + elif func_name.startswith("_"): + ignore = True + elif func_name.endswith("_"): + ignore = True + elif not func_name[0].islower(): + ignore = True + elif func_name == "unique_dim": + continue + else: + func = getattr(namespace, func_name) + if getattr(object, func_name, None) == func: + continue + if func_name == "__weakref__": + continue + func = getattr(namespace, func_name) + if namespace is torch.Tensor and getattr(object, func_name, None) == func: + continue + # ignore re-exported modules + if isinstance(func, types.ModuleType): + continue + # ignore __future__ imports + if isinstance(func, __future__._Feature): + continue + + if not callable(func) and hasattr(func, "__get__"): + index[func.__get__] = f"{namespace_str}.{func_name}.__get__" + index[func.__set__] = f"{namespace_str}.{func_name}.__set__" + if ignore: + continue + if func.__get__ in get_ignored_functions(): + msg = ( + "{}.{} is in the tuple returned by torch._overrides.get_ignored_functions " + "but still has an explicit override" + ) + if func.__get__ in get_testing_overrides(): + raise AssertionError(msg.format(namespace, func.__name__)) + continue + else: + overridable_funcs[func].append(func.__get__) + continue + + if not callable(func): + continue + + index[func] = f"{namespace_str}.{func_name}" + + if ignore: + continue + + # cannot be overridden by __torch_function__ + if func in get_ignored_functions(): + msg = ( + "{}.{} is in the tuple returned by torch._overrides.get_ignored_functions " + "but still has an explicit override" + ) + if func in get_testing_overrides(): + raise AssertionError(msg.format(namespace, func.__name__)) + continue + overridable_funcs[namespace].append(func) + return overridable_funcs, index + + +@_disable_user_warnings +def get_overridable_functions() -> dict[Any, list[Callable]]: + """List functions that are overridable via __torch_function__ + + Returns + ------- + Dict[Any, List[Callable]] + A dictionary that maps namespaces that contain overridable functions + to functions in that namespace that can be overridden. + """ + return _get_overridable_functions()[0] + + +@_disable_user_warnings +def resolve_name(f): + """Get a human readable string name for a function passed to + __torch_function__ + + Arguments + --------- + f : Callable + Function to resolve the name of. + + Returns + ------- + str + Name of the function; if eval'ed it should give back the input + function. + """ + if isinstance(f, (torch._ops.OpOverload, torch._ops.OpOverloadPacket)): + return str(f) + return _get_overridable_functions()[1].get(f) + + +@functools.cache +def _get_tensor_methods() -> set[Callable]: + """Returns a set of the overridable methods on ``torch.Tensor``""" + overridable_funcs = get_overridable_functions() + methods = set(overridable_funcs[torch.Tensor]) + return methods + + +@_disable_user_warnings +def is_tensor_method_or_property(func: Callable) -> bool: + """ + Returns True if the function passed in is a handler for a + method or property belonging to ``torch.Tensor``, as passed + into ``__torch_function__``. + + .. note:: + For properties, their ``__get__`` method must be passed in. + + This may be needed, in particular, for the following reasons: + + 1. Methods/properties sometimes don't contain a `__module__` slot. + 2. They require that the first passed-in argument is an instance + of ``torch.Tensor``. + + Examples + -------- + >>> is_tensor_method_or_property(torch.Tensor.add) + True + >>> is_tensor_method_or_property(torch.add) + False + """ + return func in _get_tensor_methods() or func.__name__ == "__get__" + + +def is_tensor_like(inp): + """ + Returns ``True`` if the passed-in input is a Tensor-like. + + Currently, this occurs whenever there's a ``__torch_function__`` + attribute on the type of the input. + + Examples + -------- + A subclass of tensor is generally a Tensor-like. + + >>> class SubTensor(torch.Tensor): ... + >>> is_tensor_like(SubTensor([0])) + True + + Built-in or user types aren't usually Tensor-like. + + >>> is_tensor_like(6) + False + >>> is_tensor_like(None) + False + >>> class NotATensor: ... + >>> is_tensor_like(NotATensor()) + False + + But, they can be made Tensor-like by implementing __torch_function__. + + >>> class TensorLike: + ... @classmethod + ... def __torch_function__(cls, func, types, args, kwargs): + ... return -1 + >>> is_tensor_like(TensorLike()) + True + """ + return type(inp) is torch.Tensor or hasattr(inp, "__torch_function__") + + +class TorchFunctionMode: + """ + A ``TorchFunctionMode`` allows you to override the meaning of all + ``__torch_function__`` overridable functions within a dynamic scope, + without having to actually create a tensor subclass or manually + monkey-patch functions in the PyTorch API. Some common situations + where you should use a mode: + + * You want to override the meaning of factory functions, or other + functions that do not otherwise take a tensor as an argument + (these cannot be overridden with tensor subclasses). + + * You want to override the behavior of all functions without needing + to wrap your inputs in tensor subclasses; e.g., if you are just + interested in logging intermediate computations. + + * You want to control the order of execution of various tensor + subclasses explicitly, rather than implicitly via the return of + ``NotImplemented``. + + Independent subclasses of :class:`TorchFunctionMode` are compositional: + modes can be pushed onto a stack using ``with MyMode():``. + When you call functions in the PyTorch API inside your + ``__torch_function__`` implementation, by default, they will forward on to + the next mode on the mode stack. If you want recursively call back into + your current ``__torch_function__`` implementation, either explicitly + invoke ``self.__torch_function__(...)``, or use the context manager + ``enable_torch_function_mode(self, replace=self.inner)`` to make PyTorch + API self-referential (beware of infinite loops, in this case!) + """ + + inner: "TorchFunctionMode" + + # Force metaclass to generate constructor at the base of the hierarchy + def __init__(self) -> None: + pass + + def __torch_function__(self, func, types, args=(), kwargs=None): + raise NotImplementedError + + def __enter__(self): + _push_mode(self) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + _pop_mode() + + @classmethod + def push(cls, *args, **kwargs): + warnings.warn( + "`Mode.push()` is no longer necessary and can be replaced with just `with Mode()`", + stacklevel=2, + ) + instance = cls(*args, **kwargs) + return instance + + +def _get_current_function_mode(): + stack_len = _len_torch_function_stack() + return _get_function_stack_at(stack_len - 1) if stack_len > 0 else None + + +def _get_current_function_mode_stack(): + stack_len = _len_torch_function_stack() + return [_get_function_stack_at(i) for i in range(stack_len)] + + +def _push_mode(mode): + _push_on_torch_function_stack(mode) + + +def _pop_mode(): + old = _pop_torch_function_stack() + return old + + +@contextlib.contextmanager +def _pop_mode_temporarily(): + old = _pop_mode() + try: + yield old + finally: + _push_mode(old) + + +class BaseTorchFunctionMode(TorchFunctionMode): + def __torch_function__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + return func(*args, **kwargs) + + +@contextlib.contextmanager +def _enable_torch_function(): + old_state = torch._C._get_torch_function_state() + try: + torch._C._set_torch_function_state(torch._C._TorchFunctionState.ENABLED) + yield + finally: + torch._C._set_torch_function_state(old_state) + + +@contextlib.contextmanager +def enable_reentrant_dispatch(): + # NB: this can't simply be + # `enable_reentrant_dispatch = torch._C._RestorePythonTLSSnapshot` + # because: + # 1. torch._C._RestorePythonTLSSnapshot is unavailable when this file + # initially gets imported. Probably an import order thing. + # 2. enable_reentrant_dispatch is technically public API; assigning + # it the object would change the __module__ to look private. + with torch._C._RestorePythonTLSSnapshot(): + try: + yield + finally: + pass diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..66cace5931ac17c548becfddbb0e56dbbdac3d38 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/__init__.py @@ -0,0 +1,12 @@ +from .analyze.is_from_package import is_from_package +from .file_structure_representation import Directory +from .glob_group import GlobGroup +from .importer import ( + Importer, + ObjMismatchError, + ObjNotFoundError, + OrderedImporter, + sys_importer, +) +from .package_exporter import EmptyMatchError, PackageExporter, PackagingError +from .package_importer import PackageImporter diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_digraph.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_digraph.py new file mode 100644 index 0000000000000000000000000000000000000000..b98b49b507a3777d9e65877e9f9a1be7d299557e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_digraph.py @@ -0,0 +1,173 @@ +# mypy: allow-untyped-defs +from collections import deque + + +class DiGraph: + """Really simple unweighted directed graph data structure to track dependencies. + + The API is pretty much the same as networkx so if you add something just + copy their API. + """ + + def __init__(self): + # Dict of node -> dict of arbitrary attributes + self._node = {} + # Nested dict of node -> successor node -> nothing. + # (didn't implement edge data) + self._succ = {} + # Nested dict of node -> predecessor node -> nothing. + self._pred = {} + + # Keep track of the order in which nodes are added to + # the graph. + self._node_order = {} + self._insertion_idx = 0 + + def add_node(self, n, **kwargs): + """Add a node to the graph. + + Args: + n: the node. Can we any object that is a valid dict key. + **kwargs: any attributes you want to attach to the node. + """ + if n not in self._node: + self._node[n] = kwargs + self._succ[n] = {} + self._pred[n] = {} + self._node_order[n] = self._insertion_idx + self._insertion_idx += 1 + else: + self._node[n].update(kwargs) + + def add_edge(self, u, v): + """Add an edge to graph between nodes ``u`` and ``v`` + + ``u`` and ``v`` will be created if they do not already exist. + """ + # add nodes + self.add_node(u) + self.add_node(v) + + # add the edge + self._succ[u][v] = True + self._pred[v][u] = True + + def successors(self, n): + """Returns an iterator over successor nodes of n.""" + try: + return iter(self._succ[n]) + except KeyError as e: + raise ValueError(f"The node {n} is not in the digraph.") from e + + def predecessors(self, n): + """Returns an iterator over predecessors nodes of n.""" + try: + return iter(self._pred[n]) + except KeyError as e: + raise ValueError(f"The node {n} is not in the digraph.") from e + + @property + def edges(self): + """Returns an iterator over all edges (u, v) in the graph""" + for n, successors in self._succ.items(): + for succ in successors: + yield n, succ + + @property + def nodes(self): + """Returns a dictionary of all nodes to their attributes.""" + return self._node + + def __iter__(self): + """Iterate over the nodes.""" + return iter(self._node) + + def __contains__(self, n): + """Returns True if ``n`` is a node in the graph, False otherwise.""" + try: + return n in self._node + except TypeError: + return False + + def forward_transitive_closure(self, src: str) -> set[str]: + """Returns a set of nodes that are reachable from src""" + + result = set(src) + working_set = deque(src) + while len(working_set) > 0: + cur = working_set.popleft() + for n in self.successors(cur): + if n not in result: + result.add(n) + working_set.append(n) + return result + + def backward_transitive_closure(self, src: str) -> set[str]: + """Returns a set of nodes that are reachable from src in reverse direction""" + + result = set(src) + working_set = deque(src) + while len(working_set) > 0: + cur = working_set.popleft() + for n in self.predecessors(cur): + if n not in result: + result.add(n) + working_set.append(n) + return result + + def all_paths(self, src: str, dst: str): + """Returns a subgraph rooted at src that shows all the paths to dst.""" + + result_graph = DiGraph() + # First compute forward transitive closure of src (all things reachable from src). + forward_reachable_from_src = self.forward_transitive_closure(src) + + if dst not in forward_reachable_from_src: + return result_graph + + # Second walk the reverse dependencies of dst, adding each node to + # the output graph iff it is also present in forward_reachable_from_src. + # we don't use backward_transitive_closures for optimization purposes + working_set = deque(dst) + while len(working_set) > 0: + cur = working_set.popleft() + for n in self.predecessors(cur): + if n in forward_reachable_from_src: + result_graph.add_edge(n, cur) + # only explore further if its reachable from src + working_set.append(n) + + return result_graph.to_dot() + + def first_path(self, dst: str) -> list[str]: + """Returns a list of nodes that show the first path that resulted in dst being added to the graph.""" + path = [] + + while dst: + path.append(dst) + candidates = self._pred[dst].keys() + dst, min_idx = "", None + for candidate in candidates: + idx = self._node_order.get(candidate, None) + if idx is None: + break + if min_idx is None or idx < min_idx: + min_idx = idx + dst = candidate + + return list(reversed(path)) + + def to_dot(self) -> str: + """Returns the dot representation of the graph. + + Returns: + A dot representation of the graph. + """ + edges = "\n".join(f'"{f}" -> "{t}";' for f, t in self.edges) + return f"""\ +digraph G {{ +rankdir = LR; +node [shape=box]; +{edges} +}} +""" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_directory_reader.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_directory_reader.py new file mode 100644 index 0000000000000000000000000000000000000000..52197fb1c84a99c79b4b587e5cdeba6053809b70 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_directory_reader.py @@ -0,0 +1,66 @@ +# mypy: allow-untyped-defs +import os.path +from glob import glob +from typing import cast + +import torch +from torch.types import Storage + + +__serialization_id_record_name__ = ".data/serialization_id" + + +# because get_storage_from_record returns a tensor!? +class _HasStorage: + def __init__(self, storage): + self._storage = storage + + def storage(self): + return self._storage + + +class DirectoryReader: + """ + Class to allow PackageImporter to operate on unzipped packages. Methods + copy the behavior of the internal PyTorchFileReader class (which is used for + accessing packages in all other cases). + + N.B.: ScriptObjects are not depickleable or accessible via this DirectoryReader + class due to ScriptObjects requiring an actual PyTorchFileReader instance. + """ + + def __init__(self, directory): + self.directory = directory + + def get_record(self, name): + filename = f"{self.directory}/{name}" + with open(filename, "rb") as f: + return f.read() + + def get_storage_from_record(self, name, numel, dtype): + filename = f"{self.directory}/{name}" + nbytes = torch._utils._element_size(dtype) * numel + storage = cast(Storage, torch.UntypedStorage) + return _HasStorage(storage.from_file(filename=filename, nbytes=nbytes)) + + def has_record(self, path): + full_path = os.path.join(self.directory, path) + return os.path.isfile(full_path) + + def get_all_records( + self, + ): + files = [ + filename[len(self.directory) + 1 :] + for filename in glob(f"{self.directory}/**", recursive=True) + if not os.path.isdir(filename) + ] + return files + + def serialization_id( + self, + ): + if self.has_record(__serialization_id_record_name__): + return self.get_record(__serialization_id_record_name__) + else: + return "" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_importlib.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_importlib.py new file mode 100644 index 0000000000000000000000000000000000000000..609efd294c4c9650d890fd36aafc9f521068ce8b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_importlib.py @@ -0,0 +1,95 @@ +# mypy: allow-untyped-defs +import _warnings +import os.path + + +# note: implementations +# copied from cpython's import code + + +# _zip_searchorder defines how we search for a module in the Zip +# archive: we first search for a package __init__, then for +# non-package .pyc, and .py entries. The .pyc entries +# are swapped by initzipimport() if we run in optimized mode. Also, +# '/' is replaced by path_sep there. + +_zip_searchorder = ( + ("/__init__.py", True), + (".py", False), +) + + +# Replace any occurrences of '\r\n?' in the input string with '\n'. +# This converts DOS and Mac line endings to Unix line endings. +def _normalize_line_endings(source): + source = source.replace(b"\r\n", b"\n") + source = source.replace(b"\r", b"\n") + return source + + +def _resolve_name(name, package, level): + """Resolve a relative module name to an absolute one.""" + bits = package.rsplit(".", level - 1) + if len(bits) < level: + raise ValueError("attempted relative import beyond top-level package") + base = bits[0] + return f"{base}.{name}" if name else base + + +def _sanity_check(name, package, level): + """Verify arguments are "sane".""" + if not isinstance(name, str): + raise TypeError(f"module name must be str, not {type(name)}") + if level < 0: + raise ValueError("level must be >= 0") + if level > 0: + if not isinstance(package, str): + raise TypeError("__package__ not set to a string") + elif not package: + raise ImportError("attempted relative import with no known parent package") + if not name and level == 0: + raise ValueError("Empty module name") + + +def _calc___package__(globals): + """Calculate what __package__ should be. + + __package__ is not guaranteed to be defined or could be set to None + to represent that its proper value is unknown. + + """ + package = globals.get("__package__") + spec = globals.get("__spec__") + if package is not None: + if spec is not None and package != spec.parent: + _warnings.warn( # noqa: G010 + f"__package__ != __spec__.parent ({package!r} != {spec.parent!r})", # noqa: G004 + ImportWarning, + stacklevel=3, + ) + return package + elif spec is not None: + return spec.parent + else: + _warnings.warn( # noqa: G010 + "can't resolve package from __spec__ or __package__, " + "falling back on __name__ and __path__", + ImportWarning, + stacklevel=3, + ) + package = globals["__name__"] + if "__path__" not in globals: + package = package.rpartition(".")[0] + return package + + +def _normalize_path(path): + """Normalize a path by ensuring it is a string. + + If the resulting string contains path separators, an exception is raised. + """ + parent, file_name = os.path.split(path) + if parent: + raise ValueError(f"{path!r} must be only a file name") + else: + return file_name diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_mangling.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_mangling.py new file mode 100644 index 0000000000000000000000000000000000000000..7b26109c19e8c2cb4a67c55ad4ce7e663d970533 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_mangling.py @@ -0,0 +1,66 @@ +# mypy: allow-untyped-defs +"""Import mangling. +See mangling.md for details. +""" + +import re + + +_mangle_index = 0 + + +class PackageMangler: + """ + Used on import, to ensure that all modules imported have a shared mangle parent. + """ + + def __init__(self) -> None: + global _mangle_index + self._mangle_index = _mangle_index + # Increment the global index + _mangle_index += 1 + # Angle brackets are used so that there is almost no chance of + # confusing this module for a real module. Plus, it is Python's + # preferred way of denoting special modules. + self._mangle_parent = f"" + + def mangle(self, name) -> str: + if len(name) == 0: + raise AssertionError("name must not be empty") + return self._mangle_parent + "." + name + + def demangle(self, mangled: str) -> str: + """ + Note: This only demangles names that were mangled by this specific + PackageMangler. It will pass through names created by a different + PackageMangler instance. + """ + if mangled.startswith(self._mangle_parent + "."): + return mangled.partition(".")[2] + + # wasn't a mangled name + return mangled + + def parent_name(self): + return self._mangle_parent + + +def is_mangled(name: str) -> bool: + return bool(re.match(r"", name)) + + +def demangle(name: str) -> str: + """ + Note: Unlike PackageMangler.demangle, this version works on any + mangled name, irrespective of which PackageMangler created it. + """ + if is_mangled(name): + _first, sep, last = name.partition(".") + # If there is only a base mangle prefix, e.g. '', + # then return an empty string. + return last if len(sep) != 0 else "" + return name + + +def get_mangle_prefix(name: str) -> str: + return name.partition(".")[0] if is_mangled(name) else name diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_mock.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_mock.py new file mode 100644 index 0000000000000000000000000000000000000000..44876b1a1d3fb3ef4a485eaf16f26755d5bb00f2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_mock.py @@ -0,0 +1,123 @@ +# mypy: allow-untyped-defs +_magic_methods = [ + "__subclasscheck__", + "__hex__", + "__rmul__", + "__float__", + "__idiv__", + "__setattr__", + "__div__", + "__invert__", + "__nonzero__", + "__rshift__", + "__eq__", + "__pos__", + "__round__", + "__rand__", + "__or__", + "__complex__", + "__divmod__", + "__len__", + "__reversed__", + "__copy__", + "__reduce__", + "__deepcopy__", + "__rdivmod__", + "__rrshift__", + "__ifloordiv__", + "__hash__", + "__iand__", + "__xor__", + "__isub__", + "__oct__", + "__ceil__", + "__imod__", + "__add__", + "__truediv__", + "__unicode__", + "__le__", + "__delitem__", + "__sizeof__", + "__sub__", + "__ne__", + "__pow__", + "__bytes__", + "__mul__", + "__itruediv__", + "__bool__", + "__iter__", + "__abs__", + "__gt__", + "__iadd__", + "__enter__", + "__floordiv__", + "__call__", + "__neg__", + "__and__", + "__ixor__", + "__getitem__", + "__exit__", + "__cmp__", + "__getstate__", + "__index__", + "__contains__", + "__floor__", + "__lt__", + "__getattr__", + "__mod__", + "__trunc__", + "__delattr__", + "__instancecheck__", + "__setitem__", + "__ipow__", + "__ilshift__", + "__long__", + "__irshift__", + "__imul__", + "__lshift__", + "__dir__", + "__ge__", + "__int__", + "__ior__", +] + + +class MockedObject: + _name: str + + def __new__(cls, *args, **kwargs): + # _suppress_err is set by us in the mocked module impl, so that we can + # construct instances of MockedObject to hand out to people looking up + # module attributes. + + # Any other attempt to construct a MockedObject instance (say, in the + # unpickling process) should give an error. + if not kwargs.get("_suppress_err"): + raise NotImplementedError( + f"Object '{cls._name}' was mocked out during packaging " + f"but it is being used in '__new__'. If this error is " + "happening during 'load_pickle', please ensure that your " + "pickled object doesn't contain any mocked objects." + ) + # Otherwise, this is just a regular object creation + # (e.g. `x = MockedObject("foo")`), so pass it through normally. + return super().__new__(cls) + + def __init__(self, name: str, _suppress_err: bool): + self.__dict__["_name"] = name + + def __repr__(self): + return f"MockedObject({self._name})" + + +def install_method(method_name): + def _not_implemented(self, *args, **kwargs): + raise NotImplementedError( + f"Object '{self._name}' was mocked out during packaging but it is being used in {method_name}" + ) + + setattr(MockedObject, method_name, _not_implemented) + + +for method_name in _magic_methods: + install_method(method_name) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_package_pickler.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_package_pickler.py new file mode 100644 index 0000000000000000000000000000000000000000..a1ecd994dbb2811ad719c5226eeadc3bc4aeadf5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_package_pickler.py @@ -0,0 +1,144 @@ +import sys +from pickle import ( + _compat_pickle, # pyrefly: ignore [missing-module-attribute] + _extension_registry, # pyrefly: ignore [missing-module-attribute] + _getattribute, # pyrefly: ignore [missing-module-attribute] + _Pickler, + EXT1, + EXT2, + EXT4, + GLOBAL, + PicklingError, + STACK_GLOBAL, +) +from struct import pack +from types import FunctionType + +from .importer import Importer, ObjMismatchError, ObjNotFoundError, sys_importer + + +class _PyTorchLegacyPickler(_Pickler): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._persistent_id = None + + def persistent_id(self, obj): + if self._persistent_id is None: + return super().persistent_id(obj) + return self._persistent_id(obj) + + +class PackagePickler(_PyTorchLegacyPickler): + """Package-aware pickler. + + This behaves the same as a normal pickler, except it uses an `Importer` + to find objects and modules to save. + """ + + def __init__(self, importer: Importer, *args, **kwargs): + self.importer = importer + super().__init__(*args, **kwargs) + + # Make sure the dispatch table copied from _Pickler is up-to-date. + # Previous issues have been encountered where a library (e.g. dill) + # mutate _Pickler.dispatch, PackagePickler makes a copy when this lib + # is imported, then the offending library removes its dispatch entries, + # leaving PackagePickler with a stale dispatch table that may cause + # unwanted behavior. + self.dispatch = _Pickler.dispatch.copy() # type: ignore[misc] + self.dispatch[FunctionType] = PackagePickler.save_global # type: ignore[assignment] + + def save_global(self, obj, name=None): + # ruff: noqa: F841 + # unfortunately the pickler code is factored in a way that + # forces us to copy/paste this function. The only change is marked + # CHANGED below. + write = self.write # type: ignore[attr-defined] + memo = self.memo # type: ignore[attr-defined] + + # CHANGED: import module from module environment instead of __import__ + try: + module_name, name = self.importer.get_name(obj, name) + except (ObjNotFoundError, ObjMismatchError) as err: + raise PicklingError(f"Can't pickle {obj}: {str(err)}") from err + + module = self.importer.import_module(module_name) + if sys.version_info >= (3, 14): + # pickle._getattribute signature changes in 3.14 + # to take iterable and return just the object (not tuple) + # We need to get the parent object that contains the attribute + name_parts = name.split(".") + if "" in name_parts: + raise PicklingError(f"Can't pickle local object {obj!r}") + if len(name_parts) == 1: + parent = module + else: + parent = _getattribute(module, name_parts[:-1]) + else: + _, parent = _getattribute(module, name) + # END CHANGED + + if self.proto >= 2: # type: ignore[attr-defined] + code = _extension_registry.get((module_name, name)) + if code: + if code <= 0: + raise AssertionError( + f"expected positive extension code, got {code}" + ) + if code <= 0xFF: + write(EXT1 + pack("= 3. + if self.proto >= 4: # type: ignore[attr-defined] + self.save(module_name) # type: ignore[attr-defined] + self.save(name) # type: ignore[attr-defined] + write(STACK_GLOBAL) + elif parent is not module: + self.save_reduce(getattr, (parent, lastname)) # type: ignore[attr-defined] + elif self.proto >= 3: # type: ignore[attr-defined] + write( + GLOBAL + + bytes(module_name, "utf-8") + + b"\n" + + bytes(name, "utf-8") + + b"\n" + ) + else: + if self.fix_imports: # type: ignore[attr-defined] + r_name_mapping = _compat_pickle.REVERSE_NAME_MAPPING + r_import_mapping = _compat_pickle.REVERSE_IMPORT_MAPPING + if (module_name, name) in r_name_mapping: + module_name, name = r_name_mapping[(module_name, name)] + elif module_name in r_import_mapping: + module_name = r_import_mapping[module_name] + try: + write( + GLOBAL + + bytes(module_name, "ascii") + + b"\n" + + bytes(name, "ascii") + + b"\n" + ) + except UnicodeEncodeError as exc: + raise PicklingError( + f"can't pickle global identifier '{module}.{name}' using " + f"pickle protocol {self.proto:d}" # type: ignore[attr-defined] + ) from exc + + self.memoize(obj) # type: ignore[attr-defined] + + +def create_pickler(data_buf, importer, protocol=4): + if importer is sys_importer: + # if we are using the normal import library system, then + # we can use the C implementation of pickle which is faster + return _PyTorchLegacyPickler(data_buf, protocol=protocol) + else: + return PackagePickler(importer, data_buf, protocol=protocol) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_package_unpickler.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_package_unpickler.py new file mode 100644 index 0000000000000000000000000000000000000000..890e6b4e03ba076e30512712d57c4bf715c4c8bb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_package_unpickler.py @@ -0,0 +1,27 @@ +# mypy: allow-untyped-defs +import _compat_pickle +import pickle + +from .importer import Importer + + +class PackageUnpickler(pickle._Unpickler): # type: ignore[name-defined] + """Package-aware unpickler. + + This behaves the same as a normal unpickler, except it uses `importer` to + find any global names that it encounters while unpickling. + """ + + def __init__(self, importer: Importer, *args, **kwargs): + super().__init__(*args, **kwargs) + self._importer = importer + + def find_class(self, module, name): + # Subclasses may override this. + if self.proto < 3 and self.fix_imports: # type: ignore[attr-defined] + if (module, name) in _compat_pickle.NAME_MAPPING: + module, name = _compat_pickle.NAME_MAPPING[(module, name)] + elif module in _compat_pickle.IMPORT_MAPPING: + module = _compat_pickle.IMPORT_MAPPING[module] + mod = self._importer.import_module(module) + return getattr(mod, name) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_stdlib.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_stdlib.py new file mode 100644 index 0000000000000000000000000000000000000000..4277ce7f3b1a528305dd6b8b8158f9cc6c716ea1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/_stdlib.py @@ -0,0 +1,20 @@ +# mypy: allow-untyped-defs +"""List of Python standard library modules. + +Sadly, there is no reliable way to tell whether a module is part of the +standard library except by comparing to a canonical list. + +This is taken from https://github.com/PyCQA/isort/tree/develop/isort/stdlibs, +which itself is sourced from the Python documentation. +""" + +import sys + + +def is_stdlib_module(module: str) -> bool: + base_module = module.partition(".")[0] + return base_module in _get_stdlib_modules() + + +def _get_stdlib_modules(): + return sys.stdlib_module_names diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8ef7a1716af241e21f97f593abde2a2b75960814 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/__init__.py @@ -0,0 +1,2 @@ +from .find_first_use_of_broken_modules import find_first_use_of_broken_modules +from .trace_dependencies import trace_dependencies diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/find_first_use_of_broken_modules.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/find_first_use_of_broken_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..fd911a0ce804ea9f59074aaf4c967fea7bd410ca --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/find_first_use_of_broken_modules.py @@ -0,0 +1,33 @@ +from torch.package.package_exporter import PackagingError + + +__all__ = ["find_first_use_of_broken_modules"] + + +def find_first_use_of_broken_modules(exc: PackagingError) -> dict[str, list[str]]: + """ + Find all broken modules in a PackagingError, and for each one, return the + dependency path in which the module was first encountered. + + E.g. broken module m.n.o was added to a dependency graph while processing a.b.c, + then re-encountered while processing d.e.f. This method would return + {'m.n.o': ['a', 'b', 'c']} + + Args: + exc: a PackagingError + + Returns: A dict from broken module names to lists of module names in the path. + """ + + if not isinstance(exc, PackagingError): + raise AssertionError( + f"exception must be a PackagingError, got {type(exc).__name__}" + ) + uses = {} + broken_module_names = [ + m for m, attr in exc.dependency_graph.nodes.items() if attr.get("error", False) + ] + for module_name in broken_module_names: + path = exc.dependency_graph.first_path(module_name) + uses[module_name] = path + return uses diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/is_from_package.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/is_from_package.py new file mode 100644 index 0000000000000000000000000000000000000000..800f87eb48672cc6f2e086ff9ef31d93d6bf1805 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/is_from_package.py @@ -0,0 +1,16 @@ +from types import ModuleType +from typing import Any + +from .._mangling import is_mangled + + +def is_from_package(obj: Any) -> bool: + """ + Return whether an object was loaded from a package. + + Note: packaged objects from externed modules will return ``False``. + """ + if type(obj) is ModuleType: + return is_mangled(obj.__name__) + else: + return is_mangled(type(obj).__module__) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/trace_dependencies.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/trace_dependencies.py new file mode 100644 index 0000000000000000000000000000000000000000..839c2da8cabc4b258a5fa3e619ef09d1b5f433dc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/analyze/trace_dependencies.py @@ -0,0 +1,65 @@ +# mypy: allow-untyped-defs +import sys +from collections.abc import Callable, Iterable +from typing import Any + + +__all__ = ["trace_dependencies"] + + +def trace_dependencies( + callable: Callable[[Any], Any], inputs: Iterable[tuple[Any, ...]] +) -> list[str]: + """Trace the execution of a callable in order to determine which modules it uses. + + Args: + callable: The callable to execute and trace. + inputs: The input to use during tracing. The modules used by 'callable' when invoked by each set of inputs + are union-ed to determine all modules used by the callable for the purpooses of packaging. + + Returns: A list of the names of all modules used during callable execution. + """ + modules_used = set() + + def record_used_modules(frame, event, arg): + # If the event being profiled is not a Python function + # call, there is nothing to do. + if event != "call": + return + + # This is the name of the function that was called. + name = frame.f_code.co_name + module = None + + # Try to determine the name of the module that the function + # is in: + # 1) Check the global namespace of the frame. + # 2) Check the local namespace of the frame. + # 3) To handle class instance method calls, check + # the attribute named 'name' of the object + # in the local namespace corresponding to "self". + if name in frame.f_globals: + module = frame.f_globals[name].__module__ + elif name in frame.f_locals: + module = frame.f_locals[name].__module__ + elif "self" in frame.f_locals: + method = getattr(frame.f_locals["self"], name, None) + module = method.__module__ if method else None + + # If a module was found, add it to the set of used modules. + if module: + modules_used.add(module) + + try: + # Attach record_used_modules as the profiler function. + sys.setprofile(record_used_modules) + + # Execute the callable with all inputs. + for inp in inputs: + callable(*inp) + + finally: + # Detach the profiler function. + sys.setprofile(None) + + return list(modules_used) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/file_structure_representation.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/file_structure_representation.py new file mode 100644 index 0000000000000000000000000000000000000000..b926645e8e84bfd5cd4e9be8c666f8415d4dbfc7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/file_structure_representation.py @@ -0,0 +1,137 @@ +# mypy: allow-untyped-defs + +from .glob_group import GlobGroup, GlobPattern + + +__all__ = ["Directory"] + + +class Directory: + """A file structure representation. Organized as Directory nodes that have lists of + their Directory children. Directories for a package are created by calling + :meth:`PackageImporter.file_structure`.""" + + def __init__(self, name: str, is_dir: bool): + self.name = name + self.is_dir = is_dir + self.children: dict[str, Directory] = {} + + def _get_dir(self, dirs: list[str]) -> "Directory": + """Builds path of Directories if not yet built and returns last directory + in list. + + Args: + dirs (List[str]): List of directory names that are treated like a path. + + Returns: + :class:`Directory`: The last Directory specified in the dirs list. + """ + if len(dirs) == 0: + return self + dir_name = dirs[0] + if dir_name not in self.children: + self.children[dir_name] = Directory(dir_name, True) + return self.children[dir_name]._get_dir(dirs[1:]) + + def _add_file(self, file_path: str): + """Adds a file to a Directory. + + Args: + file_path (str): Path of file to add. Last element is added as a file while + other paths items are added as directories. + """ + *dirs, file = file_path.split("/") + dir = self._get_dir(dirs) + dir.children[file] = Directory(file, False) + + def has_file(self, filename: str) -> bool: + """Checks if a file is present in a :class:`Directory`. + + Args: + filename (str): Path of file to search for. + Returns: + bool: If a :class:`Directory` contains the specified file. + """ + lineage = filename.split("/", maxsplit=1) + child = lineage[0] + grandchildren = lineage[1] if len(lineage) > 1 else None + if child in self.children: + if grandchildren is None: + return True + else: + return self.children[child].has_file(grandchildren) + return False + + def __str__(self): + str_list: list[str] = [] + self._stringify_tree(str_list) + return "".join(str_list) + + def _stringify_tree( + self, + str_list: list[str], + preamble: str = "", + dir_ptr: str = "\u2500\u2500\u2500 ", + ): + """Recursive method to generate print-friendly version of a Directory.""" + space = " " + branch = "\u2502 " + tee = "\u251c\u2500\u2500 " + last = "\u2514\u2500\u2500 " + + # add this directory's representation + str_list.append(f"{preamble}{dir_ptr}{self.name}\n") + + # add directory's children representations + if dir_ptr == tee: + preamble = preamble + branch + else: + preamble = preamble + space + + file_keys: list[str] = [] + dir_keys: list[str] = [] + for key, val in self.children.items(): + if val.is_dir: + dir_keys.append(key) + else: + file_keys.append(key) + + for index, key in enumerate(sorted(dir_keys)): + if (index == len(dir_keys) - 1) and len(file_keys) == 0: + self.children[key]._stringify_tree(str_list, preamble, last) + else: + self.children[key]._stringify_tree(str_list, preamble, tee) + for index, file in enumerate(sorted(file_keys)): + pointer = last if (index == len(file_keys) - 1) else tee + str_list.append(f"{preamble}{pointer}{file}\n") + + +def _create_directory_from_file_list( + filename: str, + file_list: list[str], + include: "GlobPattern" = "**", + exclude: "GlobPattern" = (), +) -> Directory: + """Return a :class:`Directory` file structure representation created from a list of files. + + Args: + filename (str): The name given to the top-level directory that will be the + relative root for all file paths found in the file_list. + + file_list (List[str]): List of files to add to the top-level directory. + + include (list[str] | str): An optional pattern that limits what is included from the file_list to + files whose name matches the pattern. + + exclude (list[str] | str): An optional pattern that excludes files whose name match the pattern. + + Returns: + :class:`Directory`: a :class:`Directory` file structure representation created from a list of files. + """ + glob_pattern = GlobGroup(include, exclude=exclude, separator="/") + + top_dir = Directory(filename, True) + for file in file_list: + if glob_pattern.matches(file): + top_dir._add_file(file) + return top_dir diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/find_file_dependencies.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/find_file_dependencies.py new file mode 100644 index 0000000000000000000000000000000000000000..f8579db07811d31a58bb9ec81515685a9d97f78d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/find_file_dependencies.py @@ -0,0 +1,95 @@ +# mypy: allow-untyped-defs +import ast + +from ._importlib import _resolve_name + + +class _ExtractModuleReferences(ast.NodeVisitor): + """ + Extract the list of global variables a block of code will read and write + """ + + @classmethod + def run(cls, src: str, package: str) -> list[tuple[str, str | None]]: + visitor = cls(package) + tree = ast.parse(src) + visitor.visit(tree) + return list(visitor.references.keys()) + + def __init__(self, package): + super().__init__() + self.package = package + self.references = {} + + def _absmodule(self, module_name: str, level: int) -> str: + if level > 0: + return _resolve_name(module_name, self.package, level) + return module_name + + def visit_Import(self, node): + for alias in node.names: + self.references[(alias.name, None)] = True + + def visit_ImportFrom(self, node): + name = self._absmodule(node.module, 0 if node.level is None else node.level) + for alias in node.names: + # from my_package import foo + # foo may be a module, so we have to add it to the list of + # potential references, if import of it fails, we will ignore it + if alias.name != "*": + self.references[(name, alias.name)] = True + else: + self.references[(name, None)] = True + + def _grab_node_int(self, node): + return node.value + + def _grab_node_str(self, node): + return node.value + + def visit_Call(self, node): + # __import__ calls aren't routed to the visit_Import/From nodes + if hasattr(node.func, "id") and node.func.id == "__import__": + try: + name = self._grab_node_str(node.args[0]) + fromlist: list[str] = [] + level = 0 + if len(node.args) > 3: + fromlist.extend(self._grab_node_str(v) for v in node.args[3].elts) + elif hasattr(node, "keywords"): + for keyword in node.keywords: + if keyword.arg == "fromlist": + fromlist.extend( + self._grab_node_str(v) for v in keyword.value.elts + ) + if len(node.args) > 4: + level = self._grab_node_int(node.args[4]) + elif hasattr(node, "keywords"): + for keyword in node.keywords: + if keyword.arg == "level": + level = self._grab_node_int(keyword.value) + if fromlist == []: + # the top-level package (the name up till the first dot) is returned + # when the fromlist argument is empty in normal import system, + # we need to include top level package to match this behavior and last + # level package to capture the intended dependency of user + self.references[(name, None)] = True + top_name = name.rsplit(".", maxsplit=1)[0] + if top_name != name: + top_name = self._absmodule(top_name, level) + self.references[(top_name, None)] = True + else: + name = self._absmodule(name, level) + for alias in fromlist: + # fromlist args may be submodules, so we have to add the fromlist args + # to the list of potential references. If import of an arg fails we + # will ignore it, similar to visit_ImportFrom + if alias != "*": + self.references[(name, alias)] = True + else: + self.references[(name, None)] = True + except Exception: + return + + +find_files_source_depends_on = _ExtractModuleReferences.run diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/glob_group.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/glob_group.py new file mode 100644 index 0000000000000000000000000000000000000000..3f20d7cbc6542401986fa178ec818eaf07e6a516 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/glob_group.py @@ -0,0 +1,86 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import re +from collections.abc import Iterable + + +GlobPattern = str | Iterable[str] + + +class GlobGroup: + """A set of patterns that candidate strings will be matched against. + + A candidate is composed of a list of segments separated by ``separator``, e.g. "foo.bar.baz". + + A pattern contains one or more segments. Segments can be: + - A literal string (e.g. "foo"), which matches exactly. + - A string containing a wildcard (e.g. "torch*", or "foo*baz*"). The wildcard matches + any string, including the empty string. + - A double wildcard ("**"). This matches against zero or more complete segments. + + Examples: + ``torch.**``: matches ``torch`` and all its submodules, e.g. ``torch.nn`` and ``torch.nn.functional``. + ``torch.*``: matches ``torch.nn`` or ``torch.functional``, but not ``torch.nn.functional``. + ``torch*.**``: matches ``torch``, ``torchvision``, and all their submodules. + + A candidates will match the ``GlobGroup`` if it matches any of the ``include`` patterns and + none of the ``exclude`` patterns. + + Args: + include (str | Iterable[str]): A string or list of strings, + each representing a pattern to be matched against. A candidate + will match if it matches *any* include pattern + exclude (str | Iterable[str]): A string or list of strings, + each representing a pattern to be matched against. A candidate + will be excluded from matching if it matches *any* exclude pattern. + separator (str): A string that delimits segments in candidates and + patterns. By default this is "." which corresponds to how modules are + named in Python. Another common value for this is "/", which is + the Unix path separator. + """ + + def __init__( + self, include: GlobPattern, *, exclude: GlobPattern = (), separator: str = "." + ): + self._dbg = f"GlobGroup(include={include}, exclude={exclude})" + self.include = GlobGroup._glob_list(include, separator) + self.exclude = GlobGroup._glob_list(exclude, separator) + self.separator = separator + + def __str__(self): + return self._dbg + + def __repr__(self): + return self._dbg + + def matches(self, candidate: str) -> bool: + candidate = self.separator + candidate + return any(p.fullmatch(candidate) for p in self.include) and all( + not p.fullmatch(candidate) for p in self.exclude + ) + + @staticmethod + def _glob_list(elems: GlobPattern, separator: str = "."): + if isinstance(elems, str): + return [GlobGroup._glob_to_re(elems, separator)] + else: + return [GlobGroup._glob_to_re(e, separator) for e in elems] + + @staticmethod + def _glob_to_re(pattern: str, separator: str = "."): + # to avoid corner cases for the first component, we prefix the candidate string + # with '.' so `import torch` will regex against `.torch`, assuming '.' is the separator + def component_to_re(component): + if "**" in component: + if component == "**": + return "(" + re.escape(separator) + "[^" + separator + "]+)*" + else: + raise ValueError("** can only appear as an entire path segment") + else: + return re.escape(separator) + ("[^" + separator + "]*").join( + re.escape(x) for x in component.split("*") + ) + + result = "".join(component_to_re(c) for c in pattern.split(separator)) + return re.compile(result) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/importer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/importer.py new file mode 100644 index 0000000000000000000000000000000000000000..66a996067b73bb071cfeb6559e9966c23d409c42 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/importer.py @@ -0,0 +1,262 @@ +import importlib +import logging +import sys +from abc import ABC, abstractmethod +from pickle import ( + _getattribute, # pyrefly: ignore [missing-module-attribute] + _Pickler, + whichmodule as _pickle_whichmodule, # pyrefly: ignore [missing-module-attribute] +) +from types import ModuleType +from typing import Any + +from ._mangling import demangle, get_mangle_prefix, is_mangled + + +__all__ = ["ObjNotFoundError", "ObjMismatchError", "Importer", "OrderedImporter"] +log = logging.getLogger(__name__) + + +class ObjNotFoundError(Exception): + """Raised when an importer cannot find an object by searching for its name.""" + + +class ObjMismatchError(Exception): + """Raised when an importer found a different object with the same name as the user-provided one.""" + + +class Importer(ABC): + """Represents an environment to import modules from. + + By default, you can figure out what module an object belongs by checking + __module__ and importing the result using __import__ or importlib.import_module. + + torch.package introduces module importers other than the default one. + Each PackageImporter introduces a new namespace. Potentially a single + name (e.g. 'foo.bar') is present in multiple namespaces. + + It supports two main operations: + import_module: module_name -> module object + get_name: object -> (parent module name, name of obj within module) + + The guarantee is that following round-trip will succeed or throw an ObjNotFoundError/ObjMisMatchError. + module_name, obj_name = env.get_name(obj) + module = env.import_module(module_name) + obj2 = getattr(module, obj_name) + assert obj1 is obj2 + """ + + modules: dict[str, ModuleType] + + @abstractmethod + def import_module(self, module_name: str) -> ModuleType: + """Import `module_name` from this environment. + + The contract is the same as for importlib.import_module. + """ + + def get_name(self, obj: Any, name: str | None = None) -> tuple[str, str]: + """Given an object, return a name that can be used to retrieve the + object from this environment. + + Args: + obj: An object to get the module-environment-relative name for. + name: If set, use this name instead of looking up __name__ or __qualname__ on `obj`. + This is only here to match how Pickler handles __reduce__ functions that return a string, + don't use otherwise. + Returns: + A tuple (parent_module_name, attr_name) that can be used to retrieve `obj` from this environment. + Use it like: + mod = importer.import_module(parent_module_name) + obj = getattr(mod, attr_name) + + Raises: + ObjNotFoundError: we couldn't retrieve `obj by name. + ObjMisMatchError: we found a different object with the same name as `obj`. + """ + if name is None and obj and _Pickler.dispatch.get(type(obj)) is None: + # Honor the string return variant of __reduce__, which will give us + # a global name to search for in this environment. + # TODO: I guess we should do copyreg too? + reduce = getattr(obj, "__reduce__", None) + if reduce is not None: + try: + rv = reduce() + if isinstance(rv, str): + name = rv + except Exception: + pass + if name is None: + name = getattr(obj, "__qualname__", None) + if name is None: + name = obj.__name__ + + orig_module_name = self.whichmodule(obj, name) + # Demangle the module name before importing. If this obj came out of a + # PackageImporter, `__module__` will be mangled. See mangling.md for + # details. + module_name = demangle(orig_module_name) + + # Check that this name will indeed return the correct object + try: + module = self.import_module(module_name) + if sys.version_info >= (3, 14): + # pickle._getatribute signature changes in 3.14 + # to take iterable and return just one object + obj2 = _getattribute(module, name.split(".")) + else: + obj2, _ = _getattribute(module, name) + except (ImportError, KeyError, AttributeError): + raise ObjNotFoundError( + f"{obj} was not found as {module_name}.{name}" + ) from None + + if obj is obj2: + return module_name, name + + def get_obj_info(obj): + if name is None: + raise AssertionError("name must not be None") + module_name = self.whichmodule(obj, name) + is_mangled_ = is_mangled(module_name) + location = ( + get_mangle_prefix(module_name) + if is_mangled_ + else "the current Python environment" + ) + importer_name = ( + f"the importer for {get_mangle_prefix(module_name)}" + if is_mangled_ + else "'sys_importer'" + ) + return module_name, location, importer_name + + obj_module_name, obj_location, obj_importer_name = get_obj_info(obj) + obj2_module_name, obj2_location, obj2_importer_name = get_obj_info(obj2) + msg = ( + f"\n\nThe object provided is from '{obj_module_name}', " + f"which is coming from {obj_location}." + f"\nHowever, when we import '{obj2_module_name}', it's coming from {obj2_location}." + "\nTo fix this, make sure this 'PackageExporter's importer lists " + f"{obj_importer_name} before {obj2_importer_name}." + ) + raise ObjMismatchError(msg) + + def whichmodule(self, obj: Any, name: str) -> str: + """Find the module name an object belongs to. + + This should be considered internal for end-users, but developers of + an importer can override it to customize the behavior. + + Taken from pickle.py, but modified to exclude the search into sys.modules + """ + module_name = getattr(obj, "__module__", None) + if module_name is not None: + return module_name + + # Protect the iteration by using a list copy of self.modules against dynamic + # modules that trigger imports of other modules upon calls to getattr. + for module_name, module in self.modules.copy().items(): + if ( + module_name == "__main__" + or module_name == "__mp_main__" # bpo-42406 + or module is None + ): + continue + try: + if _getattribute(module, name)[0] is obj: + return module_name + except AttributeError: + pass + + return "__main__" + + +class _SysImporter(Importer): + """An importer that implements the default behavior of Python.""" + + def import_module(self, module_name: str): + return importlib.import_module(module_name) + + def whichmodule(self, obj: Any, name: str) -> str: + # In Python 3.14+, pickle.whichmodule tries to import the module, + # which fails for mangled package names like ''. + # Check __module__ first before calling pickle.whichmodule. + module_name = getattr(obj, "__module__", None) + if module_name is not None: + return module_name + return _pickle_whichmodule(obj, name) + + +sys_importer = _SysImporter() + + +class OrderedImporter(Importer): + """A compound importer that takes a list of importers and tries them one at a time. + + The first importer in the list that returns a result "wins". + """ + + def __init__(self, *args): + self._importers: list[Importer] = list(args) + + def _is_torchpackage_dummy(self, module): + """Returns true iff this module is an empty PackageNode in a torch.package. + + If you intern `a.b` but never use `a` in your code, then `a` will be an + empty module with no source. This can break cases where we are trying to + re-package an object after adding a real dependency on `a`, since + OrderedImportere will resolve `a` to the dummy package and stop there. + + See: https://github.com/pytorch/pytorch/pull/71520#issuecomment-1029603769 + """ + if not getattr(module, "__torch_package__", False): + return False + if not hasattr(module, "__path__"): + return False + if not hasattr(module, "__file__"): + return True + return module.__file__ is None + + def get_name(self, obj: Any, name: str | None = None) -> tuple[str, str]: + for importer in self._importers: + try: + return importer.get_name(obj, name) + except (ObjNotFoundError, ObjMismatchError) as e: + warning_message = ( + f"Tried to call get_name with obj {obj}, " + f"and name {name} on {importer} and got {e}" + ) + log.warning(warning_message) + raise ObjNotFoundError( + f"Could not find obj {obj} and name {name} in any of the importers {self._importers}" + ) + + def import_module(self, module_name: str) -> ModuleType: + last_err = None + for importer in self._importers: + if not isinstance(importer, Importer): + raise TypeError( + f"{importer} is not a Importer. " + "All importers in OrderedImporter must inherit from Importer." + ) + try: + module = importer.import_module(module_name) + if self._is_torchpackage_dummy(module): + continue + return module + except ModuleNotFoundError as err: + last_err = err + + if last_err is not None: + raise last_err + else: + raise ModuleNotFoundError(module_name) + + def whichmodule(self, obj: Any, name: str) -> str: + for importer in self._importers: + module_name = importer.whichmodule(obj, name) + if module_name != "__main__": + return module_name + + return "__main__" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/package_exporter.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/package_exporter.py new file mode 100644 index 0000000000000000000000000000000000000000..fffbaa6648434ec45663aafa92d1cd94125c3317 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/package_exporter.py @@ -0,0 +1,1211 @@ +# mypy: allow-untyped-defs +import collections +import importlib.machinery +import io +import linecache +import os +import pickletools +import platform +import types +from collections import defaultdict, OrderedDict +from collections.abc import Callable, Sequence +from dataclasses import dataclass +from enum import Enum +from importlib.machinery import SourceFileLoader +from pathlib import Path +from typing import Any, cast, IO + +import torch +from torch.serialization import location_tag, normalize_storage_type +from torch.types import FileLike, Storage +from torch.utils.hooks import RemovableHandle + +from ._digraph import DiGraph +from ._importlib import _normalize_path +from ._mangling import demangle, is_mangled +from ._package_pickler import create_pickler +from ._stdlib import is_stdlib_module +from .find_file_dependencies import find_files_source_depends_on +from .glob_group import GlobGroup, GlobPattern +from .importer import Importer, OrderedImporter, sys_importer + + +__all__ = [ + "PackagingErrorReason", + "EmptyMatchError", + "PackagingError", + "PackageExporter", +] + +_gate_torchscript_serialization = True + +ActionHook = Callable[["PackageExporter", str], None] + + +class _ModuleProviderAction(Enum): + """Represents one of the actions that :class:`PackageExporter` can take on a module. + + See :meth:`PackageExporter.extern` and friends for a description of what the actions do. + """ + + INTERN = 1 + EXTERN = 2 + MOCK = 3 + DENY = 4 + # Special case: when a module is mocked, PackageExporter writes out a + # `_mock` module that implements our mocking stubs. If we re-package code, + # we may encounter a `_mock` module from the original package. If we do, + # just ignore it and write a `_mock` module once. + REPACKAGED_MOCK_MODULE = 5 + # Special case: PackageImporter adds a fake module + # (`torch_package_importer`) that allows packaged code to access it. Don't + # re-export this. + SKIP = 6 + + +class PackagingErrorReason(Enum): + """Listing of different reasons a dependency may fail to package. + + This enum is used to provide good error messages when + :class:`PackagingError` is raised. + """ + + def __repr__(self): + return f"<{self.__class__.__name__}.{self.name}>" + + IS_EXTENSION_MODULE = ( + "Module is a C extension module. torch.package supports Python modules only." + ) + NO_DUNDER_FILE = "Module had no __file__ defined." + SOURCE_FILE_NOT_FOUND = ( + "Module had a __file__, but we could not find it in your filesystem." + ) + DEPENDENCY_RESOLUTION_FAILED = "Dependency resolution failed." + NO_ACTION = ( + "Module did not match against any action pattern. Extern, mock, or intern it." + ) + DENIED = "Module was denied by a pattern." + MOCKED_BUT_STILL_USED = ( + "Module was mocked out, but is still being used in the package. " + "Please intern or extern the mocked modules if objects are supposed to be in " + "the package." + ) + + +@dataclass +class _PatternInfo: + """Holds :class:`PackageExporter`-specific info about how to execute matches against""" + + # What action to take on a module that matches this pattern. + action: _ModuleProviderAction + # The value of `allow_empty` the user gave when specifying the pattern. + allow_empty: bool + # Whether this pattern has been matched during packaging. + was_matched: bool + + def __init__(self, action, allow_empty): + self.action = action + self.allow_empty = allow_empty + self.was_matched = False + + +class EmptyMatchError(Exception): + """This is an exception that is thrown when a mock or extern is marked as + ``allow_empty=False``, and is not matched with any module during packaging. + """ + + +class PackagingError(Exception): + """This exception is raised when there is an issue with exporting a package. + ``PackageExporter`` will attempt to gather up all the errors and present + them to you at once. + """ + + def __init__(self, dependency_graph: DiGraph, debug=False): + # Group errors by reason. + broken: dict[PackagingErrorReason, list[str]] = defaultdict(list) + for module_name, attrs in dependency_graph.nodes.items(): + error = attrs.get("error") + if error is None: + continue + if error == PackagingErrorReason.NO_ACTION: + if "action" in attrs: + raise AssertionError( + f"module {module_name} has NO_ACTION error but action is set" + ) + broken[error].append(module_name) + + message = io.StringIO() + message.write("\n") + + for reason, module_names in broken.items(): + message.write(f"* {reason.value}\n") + for module_name in module_names: + message.write(f" {module_name}\n") + + # Print additional context if it's provided. + error_context = dependency_graph.nodes[module_name].get("error_context") + if error_context is not None: + message.write(f" Context: {error_context}\n") + if module_name in _DISALLOWED_MODULES: + message.write( + " Note: While we usually use modules in the python standard library " + f"from the local environment, `{module_name}` has a lot of system " + "level access and therefore can pose a security risk. We heavily " + f"recommend removing `{module_name}` from your packaged code. However, if that " + "is not possible, add it to the extern list by calling " + f'PackageExporter.extern("`{module_name}`")\n' + ) + if debug: + module_path = dependency_graph.first_path(module_name) + message.write( + f" A path to {module_name}: {' -> '.join(module_path)}\n" + ) + if not debug: + message.write("\n") + message.write( + "Set debug=True when invoking PackageExporter for a visualization of where " + "broken modules are coming from!\n" + ) + # Save the dependency graph so that tooling can get at it. + self.dependency_graph = dependency_graph + super().__init__(message.getvalue()) + + +class PackageExporter: + """Exporters allow you to write packages of code, pickled Python data, and + arbitrary binary and text resources into a self-contained package. + + Imports can load this code in a hermetic way, such that code is loaded + from the package rather than the normal Python import system. This allows + for the packaging of PyTorch model code and data so that it can be run + on a server or used in the future for transfer learning. + + The code contained in packages is copied file-by-file from the original + source when it is created, and the file format is a specially organized + zip file. Future users of the package can unzip the package, and edit the code + in order to perform custom modifications to it. + + The importer for packages ensures that code in the module can only be loaded from + within the package, except for modules explicitly listed as external using :meth:`extern`. + The file ``extern_modules`` in the zip archive lists all the modules that a package externally depends on. + This prevents "implicit" dependencies where the package runs locally because it is importing + a locally-installed package, but then fails when the package is copied to another machine. + + When source code is added to the package, the exporter can optionally scan it + for further code dependencies (``dependencies=True``). It looks for import statements, + resolves relative references to qualified module names, and performs an action specified by the user + (See: :meth:`extern`, :meth:`mock`, and :meth:`intern`). + """ + + """A importer that will be searched in order to find the modules referenced by other modules or by + pickled objects. The default module environment just uses sys_importer, which searches the Python environment. + """ + importer: Importer + + def __init__( + self, + f: FileLike, + importer: Importer | Sequence[Importer] = sys_importer, + debug: bool = False, + ) -> None: + """ + Create an exporter. + + Args: + f: The location to export to. Can be a ``string``/``Path`` object containing a filename + or a binary I/O object. + importer: If a single Importer is passed, use that to search for modules. + If a sequence of importers are passed, an ``OrderedImporter`` will be constructed out of them. + debug: If set to True, add path of broken modules to PackagingErrors. + """ + torch._C._log_api_usage_once("torch.package.PackageExporter") + self.debug = debug + if isinstance(f, (str, os.PathLike)): + f = os.fspath(f) + self.buffer: IO[bytes] | None = None + else: # is a byte buffer + self.buffer = f + + self.zip_file = torch._C.PyTorchFileWriter(f) + self.zip_file.set_min_version(6) + self._written_files: set[str] = set() + + self.serialized_reduces: dict[int, Any] = {} + + # A graph tracking all the modules and pickle objects added to this + # package and the dependencies between them. + # - Each node is a module name (or a pickle name that looks like '') + # - Each directed edge (u, v) means u depends on v. + # - Nodes may contain metadata that describe how to write the thing to the zipfile. + self.dependency_graph = DiGraph() + self.script_module_serializer = torch._C.ScriptModuleSerializer(self.zip_file) + self.storage_context = self.script_module_serializer.storage_context() + + # These are OrderedDicts for compatibility with RemovableHandle. + # Generic OrderedDict type annotations are not present until 3.7. + # The real type signature is OrderedDict[int, Callable[[PackageExporter, str], None]] + self._extern_hooks: OrderedDict = OrderedDict() + self._mock_hooks: OrderedDict = OrderedDict() + self._intern_hooks: OrderedDict = OrderedDict() + + if isinstance(importer, Importer): + self.importer = importer + else: + if not isinstance(importer, collections.abc.Sequence): + raise TypeError( + "importer arg should be an Importer or a sequence of Importers, " + f"got {type(importer)} instead." + ) + self.importer = OrderedImporter(*importer) + + self.patterns: dict[GlobGroup, _PatternInfo] = {} + self._unique_id = 0 + + def save_source_file( + self, module_name: str, file_or_directory: str, dependencies=True + ): + """Adds the local file system ``file_or_directory`` to the source package to provide the code + for ``module_name``. + + Args: + module_name (str): e.g. ``"my_package.my_subpackage"``, code will be saved to provide code for this package. + file_or_directory (str): the path to a file or directory of code. When a directory, all python files in the directory + are recursively copied using :meth:`save_source_file`. If a file is named ``"/__init__.py"`` the code is treated + as a package. + dependencies (bool, optional): If ``True``, we scan the source for dependencies. + """ + path = Path(file_or_directory) + if path.is_dir(): + to_save = [] # list of tuples with arguments to save_source_string + module_path = module_name.replace(".", "/") + for filename in path.glob("**/*.py"): + relative_path = filename.relative_to(path).as_posix() + archivename = module_path + "/" + relative_path + submodule_name = None + if filename.name == "__init__.py": + submodule_name = archivename[: -len("/__init__.py")].replace( + "/", "." + ) + is_package = True + else: + submodule_name = archivename[: -len(".py")].replace("/", ".") + is_package = False + + # we delay the call to save_source_string so that we record all the source files + # being provided by this directory structure _before_ attempting to resolve the dependencies + # on the source. This makes sure we don't try to copy over modules that will just get + # overwritten by this directory blob + to_save.append( + ( + submodule_name, + _read_file(str(filename)), + is_package, + dependencies, + ) + ) + + for item in to_save: + self.save_source_string(*item) + else: + is_package = path.name == "__init__.py" + self.save_source_string( + module_name, + _read_file(file_or_directory), + is_package, + dependencies, + ) + + def get_unique_id(self) -> str: + """Get an id. This id is guaranteed to only be handed out once for this package.""" + ret = str(self._unique_id) + self._unique_id += 1 + return ret + + def _get_dependencies( + self, src: str, module_name: str, is_package: bool + ) -> list[str]: + """Return all modules that this source code depends on. + + Dependencies are found by scanning the source code for import-like statements. + + Arguments: + src: The Python source code to analyze for dependencies. + module_name: The name of the module that ``src`` corresponds to. + is_package: Whether this module should be treated as a package. + See :py:meth:`save_source_string` for more info. + + Returns: + A list containing modules detected as direct dependencies in + ``src``. The items in the list are guaranteed to be unique. + """ + package_name = ( + module_name if is_package else module_name.rsplit(".", maxsplit=1)[0] + ) + try: + dep_pairs = find_files_source_depends_on(src, package_name) + except Exception as e: + self.dependency_graph.add_node( + module_name, + error=PackagingErrorReason.DEPENDENCY_RESOLUTION_FAILED, + error_context=str(e), + ) + return [] + + # Use a dict to get uniquing but also deterministic order + dependencies = {} + for dep_module_name, dep_module_obj in dep_pairs: + # handle the case where someone did something like `from pack import sub` + # where `sub` is a submodule. In this case we don't have to save pack, just sub. + # this ensures we don't pick up additional dependencies on pack. + # However, in the case where `sub` is not a submodule but an object, then we do have + # to save pack. + if dep_module_obj is not None: + possible_submodule = f"{dep_module_name}.{dep_module_obj}" + if self._module_exists(possible_submodule): + dependencies[possible_submodule] = True + # we don't need to save `pack` + continue + if self._module_exists(dep_module_name): + dependencies[dep_module_name] = True + + return list(dependencies.keys()) + + def save_source_string( + self, + module_name: str, + src: str, + is_package: bool = False, + dependencies: bool = True, + ): + """Adds ``src`` as the source code for ``module_name`` in the exported package. + + Args: + module_name (str): e.g. ``my_package.my_subpackage``, code will be saved to provide code for this package. + src (str): The Python source code to save for this package. + is_package (bool, optional): If ``True``, this module is treated as a package. Packages are allowed to have submodules + (e.g. ``my_package.my_subpackage.my_subsubpackage``), and resources can be saved inside them. Defaults to ``False``. + dependencies (bool, optional): If ``True``, we scan the source for dependencies. + """ + self.dependency_graph.add_node( + module_name, + source=src, + is_package=is_package, + provided=True, + action=_ModuleProviderAction.INTERN, + ) + + if dependencies: + deps = self._get_dependencies(src, module_name, is_package) + + for dep in deps: + self.dependency_graph.add_edge(module_name, dep) + self.add_dependency(dep) + + def _write_source_string( + self, + module_name: str, + src: str, + is_package: bool = False, + ): + """Write ``src`` as the source code for ``module_name`` in the zip archive. + + Arguments are otherwise the same as for :meth:`save_source_string`. + """ + extension = "/__init__.py" if is_package else ".py" + filename = module_name.replace(".", "/") + extension + + self._write(filename, src) + + def _import_module(self, module_name: str): + try: + return self.importer.import_module(module_name) + except ModuleNotFoundError: + if not is_mangled(module_name): + raise + msg = ( + f"Module not found: '{module_name}'. Make sure the PackageImporter that " + "created this module is present in `self.importer`" + ) + raise ModuleNotFoundError(msg) from None + + def _module_exists(self, module_name: str) -> bool: + try: + self._import_module(module_name) + return True + except Exception: + return False + + def _get_source_of_module(self, module: types.ModuleType) -> str | None: + filename = None + spec = getattr(module, "__spec__", None) + if spec is not None: + loader = getattr(spec, "loader", None) + if loader is not None and isinstance(loader, SourceFileLoader): + try: + filename = loader.get_filename(module.__name__) + except ImportError: + pass + if filename is None: + filename = getattr(module, "__file__", None) + if isinstance(filename, str) and filename.endswith(".py"): + return "".join(linecache.getlines(filename, module.__dict__)) + return None + + def add_dependency(self, module_name: str, dependencies=True): + """Given a module, add it to the dependency graph according to patterns + specified by the user. + """ + if ( + module_name in self.dependency_graph + and self.dependency_graph.nodes[module_name].get("provided") is True + ): + return + + # Special case: PackageImporter provides a special module called + # `torch_package_importer` that allows packaged modules to reference + # their PackageImporter. We don't want to re-export this. + if module_name == "torch_package_importer": + self.dependency_graph.add_node( + module_name, + action=_ModuleProviderAction.SKIP, + provided=True, + ) + return + + if module_name == "_mock": + self.dependency_graph.add_node( + module_name, + action=_ModuleProviderAction.REPACKAGED_MOCK_MODULE, + provided=True, + ) + return + + if self._can_implicitly_extern(module_name): + self.dependency_graph.add_node( + module_name, action=_ModuleProviderAction.EXTERN, provided=True + ) + return + + for pattern, pattern_info in self.patterns.items(): + if pattern.matches(module_name): + pattern_info.was_matched = True + self.dependency_graph.add_node( + module_name, action=pattern_info.action, provided=True + ) + + if pattern_info.action == _ModuleProviderAction.DENY: + # Requiring a denied module just adds an error to the graph. + self.dependency_graph.add_node( + module_name, error=PackagingErrorReason.DENIED + ) + + # If we are interning this module, we need to retrieve its + # dependencies and package those as well. + if pattern_info.action == _ModuleProviderAction.INTERN: + self._intern_module(module_name, dependencies) + return + + # No patterns have matched. Explicitly add this as an error. + self.dependency_graph.add_node( + module_name, error=PackagingErrorReason.NO_ACTION + ) + + def save_module(self, module_name: str, dependencies=True): + """Save the code for ``module`` into the package. Code for the module is resolved using the ``importers`` path to find the + module object, and then using its ``__file__`` attribute to find the source code. + + Args: + module_name (str): e.g. ``my_package.my_subpackage``, code will be saved to provide code + for this package. + dependencies (bool, optional): If ``True``, we scan the source for dependencies. + """ + if not isinstance(module_name, str): + raise TypeError( + "save_module() expects a string input, did you perhaps mean to pass `__name__`?" + ) + + self._intern_module(module_name, dependencies) + + def _intern_module( + self, + module_name: str, + dependencies: bool, + ): + """Adds the module to the dependency graph as an interned module, + along with any metadata needed to write it out to the zipfile at serialization time. + """ + module_obj = self._import_module(module_name) + # Subtle: if the import above succeeded, either: + # 1. The module name is not mangled, and this was just a regular import, or + # 2. The module name is mangled, but one of the importers was able to + # recognize the mangling and import it. + # Either way, it is now safe to demangle this name so that we don't + # serialize the mangled version to the package. + module_name = demangle(module_name) + + # Find dependencies of this module and require them as well. + is_package = hasattr(module_obj, "__path__") + source = self._get_source_of_module(module_obj) + if source is None: + # Couldn't find a source! Add it to our dependency graph as broken + # and continue. + filename = getattr(module_obj, "__file__", None) + error_context = None + if filename is None: + packaging_error = PackagingErrorReason.NO_DUNDER_FILE + elif filename.endswith(tuple(importlib.machinery.EXTENSION_SUFFIXES)): + packaging_error = PackagingErrorReason.IS_EXTENSION_MODULE + else: + packaging_error = PackagingErrorReason.SOURCE_FILE_NOT_FOUND + error_context = f"filename: {filename}" + self.dependency_graph.add_node( + module_name, + action=_ModuleProviderAction.INTERN, + is_package=is_package, + error=packaging_error, + error_context=error_context, + provided=True, + ) + return + + self.dependency_graph.add_node( + module_name, + action=_ModuleProviderAction.INTERN, + is_package=is_package, + source=source, + provided=True, + ) + + if dependencies: + deps = self._get_dependencies(source, module_name, is_package) + for dep in deps: + self.dependency_graph.add_edge(module_name, dep) + self.add_dependency(dep) + + def save_pickle( + self, + package: str, + resource: str, + obj: Any, + dependencies: bool = True, + pickle_protocol: int = 3, + ): + """Save a python object to the archive using pickle. Equivalent to :func:`torch.save` but saving into + the archive rather than a stand-alone file. Standard pickle does not save the code, only the objects. + If ``dependencies`` is true, this method will also scan the pickled objects for which modules are required + to reconstruct them and save the relevant code. + + To be able to save an object where ``type(obj).__name__`` is ``my_module.MyObject``, + ``my_module.MyObject`` must resolve to the class of the object according to the ``importer`` order. When saving objects that + have previously been packaged, the importer's ``import_module`` method will need to be present in the ``importer`` list + for this to work. + + Args: + package (str): The name of module package this resource should go in (e.g. ``"my_package.my_subpackage"``). + resource (str): A unique name for the resource, used to identify it to load. + obj (Any): The object to save, must be picklable. + dependencies (bool, optional): If ``True``, we scan the source for dependencies. + """ + + if pickle_protocol not in (3, 4): + raise AssertionError( + f"torch.package only supports pickle protocols 3 and 4, got {pickle_protocol}" + ) + + filename = self._filename(package, resource) + # Write the pickle data for `obj` + data_buf = io.BytesIO() + pickler = create_pickler(data_buf, self.importer, protocol=pickle_protocol) + pickler.persistent_id = self._persistent_id + pickler.dump(obj) + data_value = data_buf.getvalue() + mocked_modules = defaultdict(list) + name_in_dependency_graph = f"<{package}.{resource}>" + self.dependency_graph.add_node( + name_in_dependency_graph, + action=_ModuleProviderAction.INTERN, + provided=True, + is_pickle=True, + ) + + def _check_mocked_error(module: str | None, field: str | None): + """ + checks if an object (field) comes from a mocked module and then adds + the pair to mocked_modules which contains mocked modules paired with their + list of mocked objects present in the pickle. + + We also hold the invariant that the first user defined rule that applies + to the module is the one we use. + """ + + if not isinstance(module, str): + raise AssertionError(f"module must be str, got {type(module).__name__}") + if not isinstance(field, str): + raise AssertionError(f"field must be str, got {type(field).__name__}") + if self._can_implicitly_extern(module): + return + for pattern, pattern_info in self.patterns.items(): + if pattern.matches(module): + if pattern_info.action == _ModuleProviderAction.MOCK: + mocked_modules[module].append(field) + return + + if dependencies: + all_dependencies = [] + module = None + field = None + memo: defaultdict[int, str] = defaultdict(None) + memo_count = 0 + # pickletools.dis(data_value) + # pyrefly: ignore [bad-assignment] + for opcode, arg, _pos in pickletools.genops(data_value): + if pickle_protocol == 4: + if ( + opcode.name == "SHORT_BINUNICODE" + or opcode.name == "BINUNICODE" + or opcode.name == "BINUNICODE8" + ): + if not isinstance(arg, str): + raise AssertionError( + f"expected str arg for {opcode.name}, got {type(arg).__name__}" + ) + module = field + field = arg + memo[memo_count] = arg + elif ( + opcode.name == "LONG_BINGET" + or opcode.name == "BINGET" + or opcode.name == "GET" + ): + if not isinstance(arg, int): + raise AssertionError( + f"expected int arg for {opcode.name}, got {type(arg).__name__}" + ) + module = field + field = memo.get(arg, None) + elif opcode.name == "MEMOIZE": + memo_count += 1 + elif opcode.name == "STACK_GLOBAL": + if module is None: + # If not module was passed on in the entries preceding this one, continue. + continue + if not isinstance(module, str): + raise AssertionError( + f"module must be str, got {type(module).__name__}" + ) + if module not in all_dependencies: + all_dependencies.append(module) + _check_mocked_error(module, field) + elif ( + pickle_protocol == 3 and opcode.name == "GLOBAL" + ): # a global reference + if not isinstance(arg, str): + raise AssertionError( + f"expected str arg for GLOBAL, got {type(arg).__name__}" + ) + module, field = arg.split(" ") + if module not in all_dependencies: + all_dependencies.append(module) + _check_mocked_error(module, field) + for module_name in all_dependencies: + self.dependency_graph.add_edge(name_in_dependency_graph, module_name) + + """ If an object happens to come from a mocked module, then we collect these errors and spit them + out with the other errors found by package exporter. + """ + if module_name in mocked_modules: + if not isinstance(module_name, str): + raise AssertionError( + f"module_name must be str, got {type(module_name).__name__}" + ) + fields = mocked_modules[module_name] + self.dependency_graph.add_node( + module_name, + action=_ModuleProviderAction.MOCK, + error=PackagingErrorReason.MOCKED_BUT_STILL_USED, + error_context=f"Object(s) '{fields}' from module `{module_name}` was mocked out during packaging " + f"but is being used in resource - `{resource}` in package `{package}`. ", + provided=True, + ) + else: + self.add_dependency(module_name) + + self._write(filename, data_value) + + def save_text(self, package: str, resource: str, text: str): + """Save text data to the package. + + Args: + package (str): The name of module package this resource should go it (e.g. ``"my_package.my_subpackage"``). + resource (str): A unique name for the resource, used to identify it to load. + text (str): The contents to save. + """ + return self.save_binary(package, resource, text.encode("utf-8")) + + def save_binary(self, package, resource, binary: bytes): + """Save raw bytes to the package. + + Args: + package (str): The name of module package this resource should go it (e.g. ``"my_package.my_subpackage"``). + resource (str): A unique name for the resource, used to identify it to load. + binary (str): The data to save. + """ + filename = self._filename(package, resource) + self._write(filename, binary) + + def register_extern_hook(self, hook: ActionHook) -> RemovableHandle: + """Registers an extern hook on the exporter. + + The hook will be called each time a module matches against an :meth:`extern` pattern. + It should have the following signature:: + + hook(exporter: PackageExporter, module_name: str) -> None + + Hooks will be called in order of registration. + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + A handle that can be used to remove the added hook by calling + ``handle.remove()``. + """ + handle = RemovableHandle(self._extern_hooks) + self._extern_hooks[handle.id] = hook + return handle + + def register_mock_hook(self, hook: ActionHook) -> RemovableHandle: + """Registers a mock hook on the exporter. + + The hook will be called each time a module matches against a :meth:`mock` pattern. + It should have the following signature:: + + hook(exporter: PackageExporter, module_name: str) -> None + + Hooks will be called in order of registration. + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + A handle that can be used to remove the added hook by calling + ``handle.remove()``. + """ + handle = RemovableHandle(self._mock_hooks) + self._mock_hooks[handle.id] = hook + return handle + + def register_intern_hook(self, hook: ActionHook) -> RemovableHandle: + """Registers an intern hook on the exporter. + + The hook will be called each time a module matches against an :meth:`intern` pattern. + It should have the following signature:: + + hook(exporter: PackageExporter, module_name: str) -> None + + Hooks will be called in order of registration. + + Returns: + :class:`torch.utils.hooks.RemovableHandle`: + A handle that can be used to remove the added hook by calling + ``handle.remove()``. + """ + handle = RemovableHandle(self._intern_hooks) + self._intern_hooks[handle.id] = hook + return handle + + def intern( + self, + include: "GlobPattern", + *, + exclude: "GlobPattern" = (), + allow_empty: bool = True, + ): + """Specify modules that should be packaged. A module must match some ``intern`` pattern in order to be + included in the package and have its dependencies processed recursively. + + Args: + include (list[str] | str): A string e.g. "my_package.my_subpackage", or list of strings + for the names of the modules to be externed. This can also be a glob-style pattern, as described in :meth:`mock`. + + exclude (list[str] | str): An optional pattern that excludes some patterns that match the include string. + + allow_empty (bool): An optional flag that specifies whether the intern modules specified by this call + to the ``intern`` method must be matched to some module during packaging. If an ``intern`` module glob + pattern is added with ``allow_empty=False``, and :meth:`close` is called (either explicitly or via ``__exit__``) + before any modules match that pattern, an exception is thrown. If ``allow_empty=True``, no such exception is thrown. + + """ + self.patterns[GlobGroup(include, exclude=exclude)] = _PatternInfo( + _ModuleProviderAction.INTERN, allow_empty + ) + + def mock( + self, + include: "GlobPattern", + *, + exclude: "GlobPattern" = (), + allow_empty: bool = True, + ): + """Replace some required modules with a mock implementation. Mocked modules will return a fake + object for any attribute accessed from it. Because we copy file-by-file, the dependency resolution will sometimes + find files that are imported by model files but whose functionality is never used + (e.g. custom serialization code or training helpers). + Use this function to mock this functionality out without having to modify the original code. + + Args: + include (list[str] | str): A string e.g. ``"my_package.my_subpackage"``, or list of strings + for the names of the modules to be mocked out. Strings can also be a glob-style pattern + string that may match multiple modules. Any required dependencies that match this pattern + string will be mocked out automatically. + + Examples : + ``'torch.**'`` -- matches ``torch`` and all submodules of torch, e.g. ``'torch.nn'`` + and ``'torch.nn.functional'`` + + ``'torch.*'`` -- matches ``'torch.nn'`` or ``'torch.functional'``, but not + ``'torch.nn.functional'`` + + exclude (list[str] | str): An optional pattern that excludes some patterns that match the include string. + e.g. ``include='torch.**', exclude='torch.foo'`` will mock all torch packages except ``'torch.foo'``, + Default: is ``[]``. + + allow_empty (bool): An optional flag that specifies whether the mock implementation(s) specified by this call + to the :meth:`mock` method must be matched to some module during packaging. If a mock is added with + ``allow_empty=False``, and :meth:`close` is called (either explicitly or via ``__exit__``) and the mock has + not been matched to a module used by the package being exported, an exception is thrown. + If ``allow_empty=True``, no such exception is thrown. + + """ + self.patterns[GlobGroup(include, exclude=exclude)] = _PatternInfo( + _ModuleProviderAction.MOCK, allow_empty + ) + + def extern( + self, + include: "GlobPattern", + *, + exclude: "GlobPattern" = (), + allow_empty: bool = True, + ): + """Include ``module`` in the list of external modules the package can import. + This will prevent dependency discovery from saving + it in the package. The importer will load an external module directly from the standard import system. + Code for extern modules must also exist in the process loading the package. + + Args: + include (list[str] | str): A string e.g. ``"my_package.my_subpackage"``, or list of strings + for the names of the modules to be externed. This can also be a glob-style pattern, as + described in :meth:`mock`. + + exclude (list[str] | str): An optional pattern that excludes some patterns that match the + include string. + + allow_empty (bool): An optional flag that specifies whether the extern modules specified by this call + to the ``extern`` method must be matched to some module during packaging. If an extern module glob + pattern is added with ``allow_empty=False``, and :meth:`close` is called (either explicitly or via + ``__exit__``) before any modules match that pattern, an exception is thrown. If ``allow_empty=True``, + no such exception is thrown. + + """ + self.patterns[GlobGroup(include, exclude=exclude)] = _PatternInfo( + _ModuleProviderAction.EXTERN, allow_empty + ) + + def deny(self, include: "GlobPattern", *, exclude: "GlobPattern" = ()): + """Blocklist modules who names match the given glob patterns from the list of modules the package can import. + If a dependency on any matching packages is found, a :class:`PackagingError` is raised. + + Args: + include (list[str] | str): A string e.g. ``"my_package.my_subpackage"``, or list of strings + for the names of the modules to be externed. This can also be a glob-style pattern, as described in :meth:`mock`. + + exclude (list[str] | str): An optional pattern that excludes some patterns that match the include string. + """ + self.patterns[GlobGroup(include, exclude=exclude)] = _PatternInfo( + _ModuleProviderAction.DENY, allow_empty=True + ) + + def _persistent_id(self, obj): + if torch.is_storage(obj) or isinstance(obj, torch.storage.TypedStorage): + storage: Storage + if isinstance(obj, torch.storage.TypedStorage): + # TODO: Once we decide to break serialization FC, we can + # remove this case + untyped_storage = obj._untyped_storage + storage_type_str = obj.pickle_storage_type() + storage_type = getattr(torch, storage_type_str) + storage = cast(Storage, untyped_storage) + storage_numel = obj.size() + + elif isinstance(obj, torch.UntypedStorage): + untyped_storage = obj + storage = cast(Storage, untyped_storage) + storage_type = normalize_storage_type(type(storage)) + storage_numel = storage.nbytes() + else: + raise RuntimeError(f"storage type not recognized: {type(obj)}") + + location = location_tag(storage) + + # serialize storage if not already written + storage_present = self.storage_context.has_storage(storage) + storage_id = self.storage_context.get_or_add_storage(storage) + if not storage_present: + if storage.device.type != "cpu": + storage = storage.cpu() + num_bytes = storage.nbytes() + self.zip_file.write_record( + f".data/{storage_id}.storage", storage, num_bytes + ) + return ("storage", storage_type, storage_id, location, storage_numel) + + if hasattr(obj, "__reduce_package__"): + if _gate_torchscript_serialization and isinstance( + obj, torch.jit.RecursiveScriptModule + ): + raise Exception( # noqa: TRY002 + "Serializing ScriptModules directly into a package is a beta feature. " + "To use, set global " + "`torch.package.package_exporter._gate_torchscript_serialization` to `False`." + ) + if self.serialized_reduces.get(id(obj)) is None: + self.serialized_reduces[id(obj)] = ( + "reduce_package", + id(obj), + *obj.__reduce_package__(self), + ) + + return self.serialized_reduces[id(obj)] + + return None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + # If __exit__ was called because an exception was raised, we do not + # attempt to finalize the package. Instead, control is returned to the + # caller to continue raising the exception. + if exc_type is not None: + # Do the bare minimum to leave the open buffer in a valid state. + self._finalize_zip() + return + + self.close() + + def _write(self, filename, str_or_bytes): + if filename in self._written_files: + raise AssertionError( + f"Tried to write file '{filename}', but it already exists in this archive. " + "Please file a bug." + ) + self._written_files.add(filename) + + if is_mangled(filename): + raise AssertionError( + f"Tried to save a torch.package'd module as '{filename}'. " + "Directly saving torch.package'd modules is not allowed." + ) + if isinstance(str_or_bytes, str): + str_or_bytes = str_or_bytes.encode("utf-8") + self.zip_file.write_record(filename, str_or_bytes, len(str_or_bytes)) + + def _validate_dependency_graph(self): + # 1. Check the graph for any errors inserted during dependency analysis. + for attrs in self.dependency_graph.nodes.values(): + if "error" in attrs: + raise PackagingError(self.dependency_graph, debug=self.debug) + + # 2. Check that all patterns for which allow_empty=False have been matched at least once. + for pattern, pattern_info in self.patterns.items(): + if not pattern_info.allow_empty and not pattern_info.was_matched: + raise EmptyMatchError( + f"Exporter did not match any modules to {pattern}, which was marked as allow_empty=False" + ) + + def _write_mock_file(self): + if "_mock.py" not in self._written_files: + mock_file = str(Path(__file__).parent / "_mock.py") + self._write_source_string("_mock", _read_file(mock_file), is_package=False) + + def _execute_dependency_graph(self): + """Takes a finalized dependency graph describing how to package all + modules and executes it, writing to the ZIP archive. + """ + self._validate_dependency_graph() + + extern_modules = [] + for module_name, attrs in self.dependency_graph.nodes.items(): + action = attrs["action"] + + if action == _ModuleProviderAction.EXTERN: + for hook in self._extern_hooks.values(): + hook(self, module_name) + + extern_modules.append(module_name) + + elif action == _ModuleProviderAction.MOCK: + for hook in self._mock_hooks.values(): + hook(self, module_name) + + self._write_mock_file() + + is_package = hasattr(self._import_module(module_name), "__path__") + self._write_source_string(module_name, _MOCK_IMPL, is_package) + + elif action == _ModuleProviderAction.INTERN: + for hook in self._intern_hooks.values(): + hook(self, module_name) + + # The node in the dependency graph contains metadata that tells us + # how to intern the module. + if "provided" not in attrs: + raise AssertionError( + f"Module was marked `intern` but not provided: {module_name}" + ) + + if attrs.get("is_pickle") is True: + # This node came from save_pickle, we don't need to write any source for it. + continue + + is_package = attrs["is_package"] + source = attrs["source"] + self._write_source_string(module_name, source, is_package) + + elif action == _ModuleProviderAction.REPACKAGED_MOCK_MODULE: + self._write_mock_file() + elif action == _ModuleProviderAction.SKIP: + continue + else: + raise AssertionError( + f"Invalid action: {module_name}, {action}. Please report a bug to PyTorch." + ) + + extern_file_contents = "\n".join(extern_modules) + "\n" + self._write(".data/extern_modules", extern_file_contents) + + def _write_python_version(self): + """Writes the python version that the package was created with to .data/python_version""" + self._write(".data/python_version", platform.python_version()) + + def close(self): + """Write the package to the filesystem. Any calls after :meth:`close` are now invalid. + It is preferable to use resource guard syntax instead:: + + with PackageExporter("file.zip") as e: + ... + """ + self._execute_dependency_graph() + self._write_python_version() + + self.script_module_serializer.write_files() + self._finalize_zip() + + def _finalize_zip(self): + """Called at the very end of packaging to leave the zipfile in a closed but valid state.""" + del self.zip_file + if self.buffer: + self.buffer.flush() + + def _filename(self, package, resource): + package_path = package.replace(".", "/") + resource = _normalize_path(resource) + return f"{package_path}/{resource}" + + def _can_implicitly_extern(self, module_name: str): + top_level_package_name = module_name.partition(".")[0] + return top_level_package_name == "torch" or ( + top_level_package_name not in _DISALLOWED_MODULES + and is_stdlib_module(top_level_package_name) + ) + + def dependency_graph_string(self) -> str: + """Returns digraph string representation of dependencies in package. + + Returns: + A string representation of dependencies in package. + """ + return self.dependency_graph.to_dot() + + def _nodes_with_action_type( + self, action: _ModuleProviderAction | None + ) -> list[str]: + result = [] + for name, node_dict in self.dependency_graph.nodes.items(): + node_action = node_dict.get("action", None) + if node_action == action and "is_pickle" not in node_dict: + result.append(name) + result.sort() + return result + + def externed_modules(self) -> list[str]: + """Return all modules that are currently externed. + + Returns: + A list containing the names of modules which will be + externed in this package. + """ + return self._nodes_with_action_type(_ModuleProviderAction.EXTERN) + + def interned_modules(self) -> list[str]: + """Return all modules that are currently interned. + + Returns: + A list containing the names of modules which will be + interned in this package. + """ + return self._nodes_with_action_type(_ModuleProviderAction.INTERN) + + def mocked_modules(self) -> list[str]: + """Return all modules that are currently mocked. + + Returns: + A list containing the names of modules which will be + mocked in this package. + """ + return self._nodes_with_action_type(_ModuleProviderAction.MOCK) + + def denied_modules(self) -> list[str]: + """Return all modules that are currently denied. + + Returns: + A list containing the names of modules which will be + denied in this package. + """ + return self._nodes_with_action_type(_ModuleProviderAction.DENY) + + def get_rdeps(self, module_name: str) -> list[str]: + """Return a list of all modules which depend on the module ``module_name``. + + Returns: + A list containing the names of modules which depend on ``module_name``. + """ + if module_name in self.dependency_graph._pred: + return list(self.dependency_graph._pred[module_name].keys()) + else: + return [] + + def all_paths(self, src: str, dst: str) -> str: + """Return a dot representation of the subgraph + that has all paths from src to dst. + + Returns: + A dot representation containing all paths from src to dst. + (https://graphviz.org/doc/info/lang.html) + """ + return self.dependency_graph.all_paths(src, dst) + + +# even though these are in the standard library, we do not allow them to be +# automatically externed since they offer a lot of system level access +_DISALLOWED_MODULES = ["sys", "io"] + +_MOCK_IMPL = """\ +from _mock import MockedObject +def __getattr__(attr: str): + return MockedObject(__name__ + '.' + attr, _suppress_err=True) +""" + + +def _read_file(filename: str) -> str: + with open(filename, "rb") as f: + b = f.read() + return b.decode("utf-8") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/package_importer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/package_importer.py new file mode 100644 index 0000000000000000000000000000000000000000..a42af10794fa56f204836d84fed460f38facb665 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/package/package_importer.py @@ -0,0 +1,806 @@ +# mypy: allow-untyped-defs +import builtins +import importlib +import importlib.machinery +import inspect +import io +import linecache +import os +import sys +import types +from collections.abc import Callable, Iterable +from contextlib import contextmanager +from typing import Any, cast, TYPE_CHECKING +from weakref import WeakValueDictionary + +import torch +from torch.serialization import _get_restore_location, _maybe_decode_ascii +from torch.types import FileLike + +from ._directory_reader import DirectoryReader +from ._importlib import ( + _calc___package__, + _normalize_line_endings, + _normalize_path, + _resolve_name, + _sanity_check, +) +from ._mangling import demangle, PackageMangler +from ._package_unpickler import PackageUnpickler +from .file_structure_representation import _create_directory_from_file_list, Directory +from .importer import Importer + + +if TYPE_CHECKING: + from .glob_group import GlobPattern + +__all__ = ["PackageImporter"] + + +# This is a list of imports that are implicitly allowed even if they haven't +# been marked as extern. This is to work around the fact that Torch implicitly +# depends on numpy and package can't track it. +# https://github.com/pytorch/multipy/issues/46 # codespell:ignore multipy +IMPLICIT_IMPORT_ALLOWLIST: Iterable[str] = [ + "numpy", + "numpy.core", + "numpy.core._multiarray_umath", + # FX GraphModule might depend on builtins module and users usually + # don't extern builtins. Here we import it here by default. + "builtins", +] + + +# Compatibility name mapping to facilitate upgrade of external modules. +# The primary motivation is to enable Numpy upgrade that many modules +# depend on. The latest release of Numpy removed `numpy.str` and +# `numpy.bool` breaking unpickling for many modules. +EXTERN_IMPORT_COMPAT_NAME_MAPPING: dict[str, dict[str, Any]] = { + "numpy": { + "str": str, + "bool": bool, + }, +} + + +class PackageImporter(Importer): + """Importers allow you to load code written to packages by :class:`PackageExporter`. + Code is loaded in a hermetic way, using files from the package + rather than the normal python import system. This allows + for the packaging of PyTorch model code and data so that it can be run + on a server or used in the future for transfer learning. + + The importer for packages ensures that code in the module can only be loaded from + within the package, except for modules explicitly listed as external during export. + The file ``extern_modules`` in the zip archive lists all the modules that a package externally depends on. + This prevents "implicit" dependencies where the package runs locally because it is importing + a locally-installed package, but then fails when the package is copied to another machine. + """ + + """The dictionary of already loaded modules from this package, equivalent to ``sys.modules`` but + local to this importer. + """ + + modules: dict[str, types.ModuleType] + + def __init__( + self, + file_or_buffer: FileLike | torch._C.PyTorchFileReader, + module_allowed: Callable[[str], bool] = lambda module_name: True, + ): + """Open ``file_or_buffer`` for importing. This checks that the imported package only requires modules + allowed by ``module_allowed`` + + Args: + file_or_buffer: a file-like object (has to implement :meth:`read`, :meth:`readline`, :meth:`tell`, and :meth:`seek`), + a string, or an ``os.PathLike`` object containing a filename. + module_allowed (Callable[[str], bool], optional): A method to determine if a externally provided module + should be allowed. Can be used to ensure packages loaded do not depend on modules that the server + does not support. Defaults to allowing anything. + + Raises: + ImportError: If the package will use a disallowed module. + """ + torch._C._log_api_usage_once("torch.package.PackageImporter") + + self.zip_reader: Any + if isinstance(file_or_buffer, torch._C.PyTorchFileReader): + self.filename = "" + self.zip_reader = file_or_buffer + elif isinstance(file_or_buffer, (os.PathLike, str)): + self.filename = os.fspath(file_or_buffer) + if not os.path.isdir(self.filename): + self.zip_reader = torch._C.PyTorchFileReader(self.filename) + else: + self.zip_reader = DirectoryReader(self.filename) + else: + self.filename = "" + self.zip_reader = torch._C.PyTorchFileReader(file_or_buffer) + + torch._C._log_api_usage_metadata( + "torch.package.PackageImporter.metadata", + { + "serialization_id": self.zip_reader.serialization_id(), + "file_name": self.filename, + }, + ) + + self.root = _PackageNode(None) + self.modules = {} + self.extern_modules = self._read_extern() + + for extern_module in self.extern_modules: + if not module_allowed(extern_module): + raise ImportError( + f"package '{file_or_buffer}' needs the external module '{extern_module}' " + f"but that module has been disallowed" + ) + self._add_extern(extern_module) + + for fname in self.zip_reader.get_all_records(): + self._add_file(fname) + + self.patched_builtins = builtins.__dict__.copy() + self.patched_builtins["__import__"] = self.__import__ + # Allow packaged modules to reference their PackageImporter + self.modules["torch_package_importer"] = self # type: ignore[assignment] + + self._mangler = PackageMangler() + + # used for reduce deserializaiton + self.storage_context: Any = None + self.last_map_location = None + + # used for torch.serialization._load + self.Unpickler = lambda *args, **kwargs: PackageUnpickler(self, *args, **kwargs) + + # pyrefly: ignore [bad-override] + def import_module(self, name: str, package=None): + """Load a module from the package if it hasn't already been loaded, and then return + the module. Modules are loaded locally + to the importer and will appear in ``self.modules`` rather than ``sys.modules``. + + Args: + name (str): Fully qualified name of the module to load. + package ([type], optional): Unused, but present to match the signature of importlib.import_module. Defaults to ``None``. + + Returns: + types.ModuleType: The (possibly already) loaded module. + """ + # We should always be able to support importing modules from this package. + # This is to support something like: + # obj = importer.load_pickle(...) + # importer.import_module(obj.__module__) <- this string will be mangled + # + # Note that _mangler.demangle will not demangle any module names + # produced by a different PackageImporter instance. + name = self._mangler.demangle(name) + + return self._gcd_import(name) + + def load_binary(self, package: str, resource: str) -> bytes: + """Load raw bytes. + + Args: + package (str): The name of module package (e.g. ``"my_package.my_subpackage"``). + resource (str): The unique name for the resource. + + Returns: + bytes: The loaded data. + """ + + path = self._zipfile_path(package, resource) + return self.zip_reader.get_record(path) + + def load_text( + self, + package: str, + resource: str, + encoding: str = "utf-8", + errors: str = "strict", + ) -> str: + """Load a string. + + Args: + package (str): The name of module package (e.g. ``"my_package.my_subpackage"``). + resource (str): The unique name for the resource. + encoding (str, optional): Passed to ``decode``. Defaults to ``'utf-8'``. + errors (str, optional): Passed to ``decode``. Defaults to ``'strict'``. + + Returns: + str: The loaded text. + """ + data = self.load_binary(package, resource) + return data.decode(encoding, errors) + + def load_pickle(self, package: str, resource: str, map_location=None) -> Any: + """Unpickles the resource from the package, loading any modules that are needed to construct the objects + using :meth:`import_module`. + + Args: + package (str): The name of module package (e.g. ``"my_package.my_subpackage"``). + resource (str): The unique name for the resource. + map_location: Passed to `torch.load` to determine how tensors are mapped to devices. Defaults to ``None``. + + Returns: + Any: The unpickled object. + """ + pickle_file = self._zipfile_path(package, resource) + restore_location = _get_restore_location(map_location) + loaded_storages = {} + loaded_reduces = {} + storage_context = torch._C.DeserializationStorageContext() + + def load_tensor(dtype, size, key, location, restore_location): + name = f"{key}.storage" + + if storage_context.has_storage(name): + storage = storage_context.get_storage(name, dtype)._typed_storage() + else: + tensor = self.zip_reader.get_storage_from_record( + ".data/" + name, size, dtype + ) + if isinstance(self.zip_reader, torch._C.PyTorchFileReader): + storage_context.add_storage(name, tensor) + storage = tensor._typed_storage() + loaded_storages[key] = restore_location(storage, location) + + def persistent_load(saved_id): + if not isinstance(saved_id, tuple): + raise AssertionError( + f"saved_id must be a tuple, got {type(saved_id).__name__}" + ) + typename = _maybe_decode_ascii(saved_id[0]) + data = saved_id[1:] + + if typename == "storage": + storage_type, key, location, size = data + if storage_type is torch.UntypedStorage: + dtype = torch.uint8 + else: + dtype = storage_type.dtype + + if key not in loaded_storages: + load_tensor( + dtype, + size, + key, + _maybe_decode_ascii(location), + restore_location, + ) + storage = loaded_storages[key] + # TODO: Once we decide to break serialization FC, we can + # stop wrapping with TypedStorage + return torch.storage.TypedStorage( + wrap_storage=storage._untyped_storage, dtype=dtype, _internal=True + ) + elif typename == "reduce_package": + # to fix BC breaking change, objects on this load path + # will be loaded multiple times erroneously + if len(data) == 2: + func, args = data + return func(self, *args) + reduce_id, func, args = data + if reduce_id not in loaded_reduces: + loaded_reduces[reduce_id] = func(self, *args) + return loaded_reduces[reduce_id] + else: + f"Unknown typename for persistent_load, expected 'storage' or 'reduce_package' but got '{typename}'" + + # Load the data (which may in turn use `persistent_load` to load tensors) + data_file = io.BytesIO(self.zip_reader.get_record(pickle_file)) + unpickler = self.Unpickler(data_file) + unpickler.persistent_load = persistent_load # type: ignore[assignment] + + @contextmanager + def set_deserialization_context(): + # to let reduce_package access deserializaiton context + self.storage_context = storage_context + self.last_map_location = map_location + try: + yield + finally: + self.storage_context = None + self.last_map_location = None + + with set_deserialization_context(): + result = unpickler.load() + + # TODO from zdevito: + # This stateful weird function will need to be removed in our efforts + # to unify the format. It has a race condition if multiple python + # threads try to read independent files + torch._utils._validate_loaded_sparse_tensors() + + return result + + def id(self): + """ + Returns internal identifier that torch.package uses to distinguish :class:`PackageImporter` instances. + Looks like:: + + + """ + return self._mangler.parent_name() + + def file_structure( + self, *, include: "GlobPattern" = "**", exclude: "GlobPattern" = () + ) -> Directory: + """Returns a file structure representation of package's zipfile. + + Args: + include (list[str] | str): An optional string e.g. ``"my_package.my_subpackage"``, or optional list of strings + for the names of the files to be included in the zipfile representation. This can also be + a glob-style pattern, as described in :meth:`PackageExporter.mock` + + exclude (list[str] | str): An optional pattern that excludes files whose name match the pattern. + + Returns: + :class:`Directory` + """ + return _create_directory_from_file_list( + self.filename, self.zip_reader.get_all_records(), include, exclude + ) + + def python_version(self): + """Returns the version of python that was used to create this package. + + Note: this function is experimental and not Forward Compatible. The plan is to move this into a lock + file later on. + + Returns: + :class:`str | None` a python version e.g. 3.8.9 or None if no version was stored with this package + """ + python_version_path = ".data/python_version" + return ( + self.zip_reader.get_record(python_version_path).decode("utf-8").strip() + if self.zip_reader.has_record(python_version_path) + else None + ) + + def _read_extern(self): + return ( + self.zip_reader.get_record(".data/extern_modules") + .decode("utf-8") + .splitlines(keepends=False) + ) + + def _make_module( + self, name: str, filename: str | None, is_package: bool, parent: str + ): + mangled_filename = self._mangler.mangle(filename) if filename else None + spec = importlib.machinery.ModuleSpec( + name, + self, # type: ignore[arg-type] + origin="", + is_package=is_package, + ) + module = importlib.util.module_from_spec(spec) + self.modules[name] = module + module.__name__ = self._mangler.mangle(name) + ns = module.__dict__ + ns["__spec__"] = spec + ns["__loader__"] = self + ns["__file__"] = mangled_filename + ns["__cached__"] = None + ns["__builtins__"] = self.patched_builtins + ns["__torch_package__"] = True + + # Add this module to our private global registry. It should be unique due to mangling. + if module.__name__ in _package_imported_modules: + raise AssertionError( + f"module {module.__name__} already exists in _package_imported_modules" + ) + _package_imported_modules[module.__name__] = module + + # preemptively install on the parent to prevent IMPORT_FROM from trying to + # access sys.modules + self._install_on_parent(parent, name, module) + + if filename is not None: + if mangled_filename is None: + raise AssertionError( + "mangled_filename must not be None when filename is set" + ) + # preemptively install the source in `linecache` so that stack traces, + # `inspect`, etc. work. + if filename in linecache.cache: # type: ignore[attr-defined] + raise AssertionError(f"filename {filename} already in linecache.cache") + linecache.lazycache(mangled_filename, ns) + + code = self._compile_source(filename, mangled_filename) + exec(code, ns) + + return module + + def _load_module(self, name: str, parent: str): + cur: _PathNode = self.root + for atom in name.split("."): + if not isinstance(cur, _PackageNode) or atom not in cur.children: + if name in IMPLICIT_IMPORT_ALLOWLIST: + module = self.modules[name] = importlib.import_module(name) + return module + raise ModuleNotFoundError( + f'No module named "{name}" in self-contained archive "{self.filename}"' + f" and the module is also not in the list of allowed external modules: {self.extern_modules}", + name=name, + ) + cur = cur.children[atom] + if isinstance(cur, _ExternNode): + module = self.modules[name] = importlib.import_module(name) + + if compat_mapping := EXTERN_IMPORT_COMPAT_NAME_MAPPING.get(name): + for old_name, new_name in compat_mapping.items(): + module.__dict__.setdefault(old_name, new_name) + + return module + return self._make_module( + name, + cur.source_file, # type: ignore[attr-defined] + isinstance(cur, _PackageNode), + parent, + ) + + def _compile_source(self, fullpath: str, mangled_filename: str): + source = self.zip_reader.get_record(fullpath) + source = _normalize_line_endings(source) + return compile(source, mangled_filename, "exec", dont_inherit=True) + + # note: named `get_source` so that linecache can find the source + # when this is the __loader__ of a module. + def get_source(self, module_name) -> str: + # linecache calls `get_source` with the `module.__name__` as the argument, so we must demangle it here. + module = self.import_module(demangle(module_name)) + return self.zip_reader.get_record(demangle(module.__file__)).decode("utf-8") + + # note: named `get_resource_reader` so that importlib.resources can find it. + # This is otherwise considered an internal method. + def get_resource_reader(self, fullname): + try: + package = self._get_package(fullname) + except ImportError: + return None + if package.__loader__ is not self: + return None + return _PackageResourceReader(self, fullname) + + def _install_on_parent(self, parent: str, name: str, module: types.ModuleType): + if not parent: + return + # Set the module as an attribute on its parent. + parent_module = self.modules[parent] + if parent_module.__loader__ is self: + setattr(parent_module, name.rpartition(".")[2], module) + + # note: copied from cpython's import code, with call to create module replaced with _make_module + def _do_find_and_load(self, name): + parent = name.rpartition(".")[0] + module_name_no_parent = name.rpartition(".")[-1] + if parent: + if parent not in self.modules: + self._gcd_import(parent) + # Crazy side-effects! + if name in self.modules: + return self.modules[name] + parent_module = self.modules[parent] + + try: + parent_module.__path__ # type: ignore[attr-defined] + + except AttributeError: + # when we attempt to import a package only containing pybinded files, + # the parent directory isn't always a package as defined by python, + # so we search if the package is actually there or not before calling the error. + if isinstance( + parent_module.__loader__, + importlib.machinery.ExtensionFileLoader, + ): + if name not in self.extern_modules: + msg = ( + _ERR_MSG + + "; {!r} is a c extension module which was not externed. C extension modules \ + need to be externed by the PackageExporter in order to be used as we do not support interning them.}." + ).format(name, name) + raise ModuleNotFoundError(msg, name=name) from None + if not isinstance( + parent_module.__dict__.get(module_name_no_parent), + types.ModuleType, + ): + msg = ( + _ERR_MSG + + "; {!r} is a c extension package which does not contain {!r}." + ).format(name, parent, name) + raise ModuleNotFoundError(msg, name=name) from None + else: + msg = (_ERR_MSG + "; {!r} is not a package").format(name, parent) + raise ModuleNotFoundError(msg, name=name) from None + + module = self._load_module(name, parent) + + self._install_on_parent(parent, name, module) + + return module + + # note: copied from cpython's import code + def _find_and_load(self, name): + module = self.modules.get(name, _NEEDS_LOADING) + if module is _NEEDS_LOADING: + return self._do_find_and_load(name) + + if module is None: + message = f"import of {name} halted; None in sys.modules" + raise ModuleNotFoundError(message, name=name) + + # To handle https://github.com/pytorch/pytorch/issues/57490, where std's + # creation of fake submodules via the hacking of sys.modules is not import + # friendly + if name == "os": + self.modules["os.path"] = cast(Any, module).path + elif name == "typing": + if sys.version_info < (3, 13): + self.modules["typing.io"] = cast(Any, module).io + self.modules["typing.re"] = cast(Any, module).re + + return module + + def _gcd_import(self, name, package=None, level=0): + """Import and return the module based on its name, the package the call is + being made from, and the level adjustment. + + This function represents the greatest common denominator of functionality + between import_module and __import__. This includes setting __package__ if + the loader did not. + + """ + _sanity_check(name, package, level) + if level > 0: + name = _resolve_name(name, package, level) + + return self._find_and_load(name) + + # note: copied from cpython's import code + def _handle_fromlist(self, module, fromlist, *, recursive=False): + """Figure out what __import__ should return. + + The import_ parameter is a callable which takes the name of module to + import. It is required to decouple the function from assuming importlib's + import implementation is desired. + + """ + module_name = demangle(module.__name__) + # The hell that is fromlist ... + # If a package was imported, try to import stuff from fromlist. + if hasattr(module, "__path__"): + for x in fromlist: + if not isinstance(x, str): + if recursive: + where = module_name + ".__all__" + else: + where = "``from list''" + raise TypeError( + f"Item in {where} must be str, not {type(x).__name__}" + ) + elif x == "*": + if not recursive and hasattr(module, "__all__"): + self._handle_fromlist(module, module.__all__, recursive=True) + elif not hasattr(module, x): + from_name = f"{module_name}.{x}" + try: + self._gcd_import(from_name) + except ModuleNotFoundError as exc: + # Backwards-compatibility dictates we ignore failed + # imports triggered by fromlist for modules that don't + # exist. + if ( + exc.name == from_name + and self.modules.get(from_name, _NEEDS_LOADING) is not None + ): + continue + raise + return module + + def __import__(self, name, globals=None, locals=None, fromlist=(), level=0): + if level == 0: + module = self._gcd_import(name) + else: + globals_ = globals if globals is not None else {} + package = _calc___package__(globals_) + module = self._gcd_import(name, package, level) + if not fromlist: + # Return up to the first dot in 'name'. This is complicated by the fact + # that 'name' may be relative. + if level == 0: + return self._gcd_import(name.partition(".")[0]) + elif not name: + return module + else: + # Figure out where to slice the module's name up to the first dot + # in 'name'. + cut_off = len(name) - len(name.partition(".")[0]) + # Slice end needs to be positive to alleviate need to special-case + # when ``'.' not in name``. + module_name = demangle(module.__name__) + return self.modules[module_name[: len(module_name) - cut_off]] + else: + return self._handle_fromlist(module, fromlist) + + def _get_package(self, package): + """Take a package name or module object and return the module. + + If a name, the module is imported. If the passed or imported module + object is not a package, raise an exception. + """ + if hasattr(package, "__spec__"): + if package.__spec__.submodule_search_locations is None: + raise TypeError(f"{package.__spec__.name!r} is not a package") + else: + return package + else: + module = self.import_module(package) + if module.__spec__.submodule_search_locations is None: + raise TypeError(f"{package!r} is not a package") + else: + return module + + def _zipfile_path(self, package, resource=None): + package = self._get_package(package) + if package.__loader__ is not self: + raise AssertionError( + f"package.__loader__ must be self, got {package.__loader__}" + ) + name = demangle(package.__name__) + if resource is not None: + resource = _normalize_path(resource) + return f"{name.replace('.', '/')}/{resource}" + else: + return f"{name.replace('.', '/')}" + + def _get_or_create_package(self, atoms: list[str]) -> "_PackageNode | _ExternNode": + cur = self.root + for i, atom in enumerate(atoms): + node = cur.children.get(atom, None) + if node is None: + node = cur.children[atom] = _PackageNode(None) + if isinstance(node, _ExternNode): + return node + if isinstance(node, _ModuleNode): + name = ".".join(atoms[:i]) + raise ImportError( + f"inconsistent module structure. module {name} is not a package, but has submodules" + ) + if not isinstance(node, _PackageNode): + raise AssertionError( + f"expected _PackageNode, got {type(node).__name__}" + ) + cur = node + return cur + + def _add_file(self, filename: str): + """Assembles a Python module out of the given file. Will ignore files in the .data directory. + + Args: + filename (str): the name of the file inside of the package archive to be added + """ + *prefix, last = filename.split("/") + if len(prefix) > 1 and prefix[0] == ".data": + return + package = self._get_or_create_package(prefix) + if isinstance(package, _ExternNode): + raise ImportError( + f"inconsistent module structure. package contains a module file {filename}" + f" that is a subpackage of a module marked external." + ) + if last == "__init__.py": + package.source_file = filename + elif last.endswith(".py"): + package_name = last[: -len(".py")] + package.children[package_name] = _ModuleNode(filename) + + def _add_extern(self, extern_name: str): + *prefix, last = extern_name.split(".") + package = self._get_or_create_package(prefix) + if isinstance(package, _ExternNode): + return # the shorter extern covers this extern case + package.children[last] = _ExternNode() + + +_NEEDS_LOADING = object() +_ERR_MSG_PREFIX = "No module named " +_ERR_MSG = _ERR_MSG_PREFIX + "{!r}" + + +class _PathNode: + __slots__ = [] + + +class _PackageNode(_PathNode): + def __init__(self, source_file: str | None): + self.source_file = source_file + self.children: dict[str, _PathNode] = {} + + +class _ModuleNode(_PathNode): + __slots__ = ["source_file"] + + def __init__(self, source_file: str): + self.source_file = source_file + + +class _ExternNode(_PathNode): + pass + + +# A private global registry of all modules that have been package-imported. +_package_imported_modules: WeakValueDictionary = WeakValueDictionary() + +# `inspect` by default only looks in `sys.modules` to find source files for classes. +# Patch it to check our private registry of package-imported modules as well. +_orig_getfile = inspect.getfile + + +def _patched_getfile(object): + if inspect.isclass(object): + if object.__module__ in _package_imported_modules: + return _package_imported_modules[object.__module__].__file__ + return _orig_getfile(object) + + +inspect.getfile = _patched_getfile + + +class _PackageResourceReader: + """Private class used to support PackageImporter.get_resource_reader(). + + Confirms to the importlib.abc.ResourceReader interface. Allowed to access + the innards of PackageImporter. + """ + + def __init__(self, importer, fullname): + self.importer = importer + self.fullname = fullname + + def open_resource(self, resource): + from io import BytesIO + + return BytesIO(self.importer.load_binary(self.fullname, resource)) + + def resource_path(self, resource): + # The contract for resource_path is that it either returns a concrete + # file system path or raises FileNotFoundError. + if isinstance( + self.importer.zip_reader, DirectoryReader + ) and self.importer.zip_reader.has_record( + os.path.join(self.fullname, resource) + ): + return os.path.join( + self.importer.zip_reader.directory, self.fullname, resource + ) + raise FileNotFoundError + + def is_resource(self, name): + path = self.importer._zipfile_path(self.fullname, name) + return self.importer.zip_reader.has_record(path) + + def contents(self): + from pathlib import Path + + filename = self.fullname.replace(".", "/") + + fullname_path = Path(self.importer._zipfile_path(self.fullname)) + files = self.importer.zip_reader.get_all_records() + subdirs_seen = set() + for filename in files: + try: + relative = Path(filename).relative_to(fullname_path) + except ValueError: + continue + # If the path of the file (which is relative to the top of the zip + # namespace), relative to the package given when the resource + # reader was created, has a parent, then it's a name in a + # subdirectory and thus we skip it. + parent_name = relative.parent.name + if len(parent_name) == 0: + yield relative.name + elif parent_name not in subdirs_seen: + subdirs_seen.add(parent_name) + yield parent_name diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..153d4560e264198ddb8f04649aef22456afbb452 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/__init__.py @@ -0,0 +1,60 @@ +r""" +PyTorch Profiler is a tool that allows the collection of performance metrics during training and inference. +Profiler's context manager API can be used to better understand what model operators are the most expensive, +examine their input shapes and stack traces, study device kernel activity and visualize the execution trace. + +.. note:: + An earlier version of the API in :mod:`torch.autograd` module is considered legacy and will be deprecated. + +""" + +import os +from typing import Any +from typing_extensions import TypeVarTuple, Unpack + +from torch._C._autograd import _supported_activities, DeviceType, kineto_available +from torch._C._profiler import _ExperimentalConfig, ProfilerActivity, RecordScope +from torch._environment import is_fbcode +from torch.autograd.profiler import KinetoStepTracker, record_function +from torch.optim.optimizer import Optimizer, register_optimizer_step_post_hook + +from .profiler import ( + _KinetoProfile, + ExecutionTraceObserver, + profile, + ProfilerAction, + schedule, + supported_activities, + tensorboard_trace_handler, +) + + +__all__ = [ + "profile", + "schedule", + "supported_activities", + "tensorboard_trace_handler", + "ProfilerAction", + "ProfilerActivity", + "kineto_available", + "DeviceType", + "record_function", + "ExecutionTraceObserver", +] + +from . import itt + + +_Ts = TypeVarTuple("_Ts") + + +def _optimizer_post_hook( + optimizer: Optimizer, args: tuple[Unpack[_Ts]], kwargs: dict[str, Any] +) -> None: + KinetoStepTracker.increment_step("Optimizer") + + +if os.environ.get("KINETO_USE_DAEMON", "") or ( + is_fbcode() and os.environ.get("KINETO_FORCE_OPTIMIZER_HOOK", "") +): + _ = register_optimizer_step_post_hook(_optimizer_post_hook) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/_memory_profiler.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/_memory_profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..39e9452360657bdaba7ee5299be5653f9b326e37 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/_memory_profiler.py @@ -0,0 +1,1234 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import collections +import dataclasses +import enum +import itertools as it +import logging +from typing import Any, cast, Literal, TYPE_CHECKING + +import torch +from torch._C._profiler import ( + _EventType, + _ExtraFields_Allocation, + _ExtraFields_TorchOp, + _ProfilerEvent, + _TensorMetadata, + RecordScope, +) +from torch._utils import _element_size +from torch.profiler import _utils + + +if TYPE_CHECKING: + from collections.abc import Iterator + + from torch._C import FunctionSchema + from torch._C._autograd import _ProfilerResult + + +KeyAndID = tuple["Key", int] +TensorAndID = tuple["TensorKey", int] + +log = logging.getLogger(__name__) + + +class Category(enum.Enum): + INPUT = enum.auto() + TEMPORARY = enum.auto() + ACTIVATION = enum.auto() + GRADIENT = enum.auto() + AUTOGRAD_DETAIL = enum.auto() + PARAMETER = enum.auto() + OPTIMIZER_STATE = enum.auto() + + +_CATEGORY_TO_COLORS = { + Category.PARAMETER: "darkgreen", + Category.OPTIMIZER_STATE: "goldenrod", + Category.INPUT: "black", + Category.TEMPORARY: "mediumpurple", + Category.ACTIVATION: "red", + Category.GRADIENT: "mediumblue", + Category.AUTOGRAD_DETAIL: "royalblue", + None: "grey", +} + +_CATEGORY_TO_INDEX = {c: i for i, c in enumerate(_CATEGORY_TO_COLORS)} + + +class Action(enum.Enum): + PREEXISTING = enum.auto() + CREATE = enum.auto() + INCREMENT_VERSION = enum.auto() + DESTROY = enum.auto() + + +_ACTION_TO_INDEX = {i: i.value for i in Action} + + +@dataclasses.dataclass(eq=True, unsafe_hash=False, frozen=True) +class Key: + device: torch.device + + +@dataclasses.dataclass +class _Storage: + """Bundle storage pointer and id. + + All profiling logic should use `allocation_id`, however it is useful to + print storage pointers for debugging and unit tests sometimes look up + values using the storage data pointer of a live Tensor.""" + + ptr: int + allocation_id: int + + def __repr__(self) -> str: + return f"{hex(self.ptr):>18} ({self.allocation_id})" + + def __eq__(self, other: object) -> bool: + return isinstance(other, _Storage) and self.allocation_id == other.allocation_id + + def __hash__(self) -> int: + return hash(self.allocation_id) + + +@dataclasses.dataclass(eq=True, unsafe_hash=True, frozen=True) +class TensorKey(Key): + """Hashable identifier for a storage which has been assigned an ID. + + A detailed description of Tensor IDs and why they are needed is given in + `torch/csrc/profiler/collection.h` when `TensorID` is declared. To + summarize, multiple Storage buffers can map to the same logical Tensor. + This dataclass is used to refer to a concrete in-memory StorageImpl of + a Tensor. + """ + + id: int + storage: _Storage + + def __repr__(self) -> str: + return f"id={self.id}: {repr(self.storage):<24} ({self.device})" + + def __lt__(self, other: TensorKey) -> bool: + return self._as_sortable < other._as_sortable + + @staticmethod + def _make( + tensor_id: int | None, + storage_ptr: int | None, + allocation_id: int | None, + device: torch.device, + ) -> TensorKey | None: + if ( + tensor_id is not None + and storage_ptr is not None + and allocation_id is not None + ): + return TensorKey(device, tensor_id, _Storage(storage_ptr, allocation_id)) + return None + + @classmethod + def from_allocation(cls, alloc: _ExtraFields_Allocation) -> TensorKey | None: + return cls._make(alloc.id, alloc.ptr, alloc.allocation_id, alloc.device) + + @classmethod + def from_tensor(cls, t: _TensorMetadata | None) -> TensorKey | None: + if t is not None: + return cls._make(t.id, t.storage_data_ptr, t.allocation_id, t.device) + return None + + @property + def _as_sortable(self) -> tuple[int, int, str, int]: + return self.id, self.storage.allocation_id, self.device.type, self.device.index + + +def _extract_parameters_and_gradients( + node: _ProfilerEvent, +) -> Iterator[tuple[TensorKey | None, TensorKey | None]]: + children = node.children + + # AccumulateGrad is used in the Autograd engine to handle gradient updates. + # There are two possible cases: + # 1) This is a newly created gradient Tensor. In that case there is nothing + # to accumulate, so autograd simply detaches the Tensor. + # + # 2) There is a preexisting gradient Tensor and we need to add the newly + # computed update. This is done with an in-place add (aten::add_) op. + # (The underscore suffix denotes "in-place".) + if ( + node.typed[0] == _EventType.TorchOp + and node.typed[1].scope == RecordScope.BACKWARD_FUNCTION + # TODO(robieta): Move away from load bearing names + and node.name == "torch::autograd::AccumulateGrad" + and children + and children[0].typed[0] == _EventType.TorchOp + and children[0].name in ("aten::detach", "aten::add_") + and children[0].typed[1].inputs + and isinstance(children[0].typed[1].inputs[0], _TensorMetadata) + ): + yield None, TensorKey.from_tensor(children[0].typed[1].inputs[0]) + + # We directly instrument `torch.nn.Module` and `torch.optim.Optimizer` + # NOTE: The values captured by the python tracer are cached; they can be + # used to build up labels but do not imply that a Tensor was live at + # a particular time. + elif node.typed[0] == _EventType.PyCall: + typed_fields = node.typed[1] + if typed_fields.module is not None and typed_fields.optimizer is not None: + raise AssertionError("module and optimizer cannot both be set") + if typed_fields.module is not None: + for _, p, p_grad in typed_fields.module.parameters: + yield TensorKey.from_tensor(p), TensorKey.from_tensor(p_grad) + + if typed_fields.optimizer is not None: + for p, p_grad, _ in typed_fields.optimizer.parameters: + yield TensorKey.from_tensor(p), TensorKey.from_tensor(p_grad) + + +def extract_parameters(node: _ProfilerEvent) -> Iterator[TensorKey]: + for p, _p_grad in _extract_parameters_and_gradients(node): + if p is not None: + yield p + + +def extract_gradients( + node: _ProfilerEvent, +) -> Iterator[tuple[TensorKey | None, TensorKey]]: + for p, p_grad in _extract_parameters_and_gradients(node): + if p_grad is not None: + yield p, p_grad + + +def get_scopes(event: _ProfilerEvent | None) -> tuple[RecordScope, ...]: + scopes = [] + while event: + if event.typed[0] == _EventType.TorchOp: + scopes.append(event.typed[1].scope) + event = event.parent + return tuple(scopes) + + +class SchemaMatcher: + """Lookup operator schema based on profiled name. + + When profiling we record the operator's name but not the schema. However + some analysis requires that information. Fortunately we can look up + registered schema from the recorded name. We do not, however, record the + overload and so we must compare the profiled arguments with all overloads + to determine viable matches. + + Note: Once https://github.com/pytorch/pytorch/issues/78871 is completed + this code will be obsolete. + """ + + @classmethod + def inputs_are_mutable(cls, t: _ExtraFields_TorchOp) -> tuple[bool | None, ...]: + """Determine which inputs may have mutated based on function schema. + + Note that we don't need to resolve down to a single schema to perform + this analysis. An input is mutable if it is mutable in any overload. In + practice, however, it is overwhelmingly common to match a single + overload. If we cannot find any valid schema then we must be + conservative and assume all inputs are mutable. + """ + mutable: list[bool] | None = None + for schema in cls.match_schemas(t): + mutable = mutable or [False for _ in schema.arguments] + for i, arg in enumerate(schema.arguments): + # pyrefly: ignore [unsupported-operation] + mutable[i] |= getattr(arg.alias_info, "is_write", False) + + return tuple(mutable or (None for _ in t.inputs)) + + @classmethod + def match_schemas(cls, t: _ExtraFields_TorchOp) -> tuple[FunctionSchema, ...]: + signature = tuple( + # Tensor + TensorKey.from_tensor(i) + if isinstance(i, _TensorMetadata) + # + # TensorList + else [TensorKey.from_tensor(j) for j in i] + if isinstance(i, list) + # + # Scalar and uncaptured inputs. + else i + for i in t.inputs + ) + + def matches(schema) -> bool: + return len(schema.arguments) == len(signature) and all( + cls._types_match(observed, schema_arg.type) + for observed, schema_arg in zip( + signature, schema.arguments, strict=True + ) + ) + + return tuple(s for s in cls.lookup_schemas(t.name) or () if matches(s)) + + @classmethod + def _types_match(cls, observed, schema_type) -> bool: + if isinstance(schema_type, torch._C.OptionalType): + schema_type = schema_type.getElementType() + return observed is None or cls._types_match(observed, schema_type) + + if isinstance(schema_type, torch._C.AnyType): + return True + + if schema_type.isSubtypeOf(torch._C.ListType.ofTensors()): + return isinstance(observed, list) and all( + isinstance(i, TensorKey) for i in observed + ) + + type_map: tuple[tuple[Any, type | tuple[type, ...]], ...] = ( + (torch._C.TensorType, TensorKey), + (torch._C.NoneType, type(None)), + (torch._C.BoolType, bool), + (torch._C.IntType, int), + (torch._C.FloatType, float), + (torch._C.ComplexType, complex), + (torch._C.NumberType, (bool, int, float, complex)), + ) + + for jit_type, py_types in type_map: + if isinstance(schema_type, jit_type): + return isinstance(observed, py_types) + + # Profiler only records a subset of possible argument types. If we + # reach this point then the schema must call for a type that profiler + # does not record. Thus, the schema can only be a match if `observed` + # is also None. + return observed is None + + @staticmethod + def lookup_schemas(name: str) -> tuple[FunctionSchema, ...] | None: + # TODO(robieta): + # _jit_get_schemas_for_operator is quite expensive. (~100us / call) + # Consider adding `functools.lru_cache` if that becomes an issue. + + try: + # Schema lookup will throw if `name` is malformed. (For example, + # schemas must be namespaced and schema lookup will fail if name + # does not include "::".) We simply catch the exception and return + # `None` to denote that `name` cannot be an operator name. + # + # Note that record_function annotations also go through this path, + # so it is expected that some names will not correspond to PyTorch + # operators. + if "::" not in name: + return None + return tuple(torch._C._jit_get_schemas_for_operator(name)) + except RuntimeError: + return None + + +class OpTree: + def __init__(self, result: _ProfilerResult) -> None: + self._root_nodes = result.experimental_event_tree() + self._sorted_nodes = tuple(sorted(self.dfs(), key=lambda x: x.start_time_ns)) + + def dfs(self, *args, **kwargs) -> Iterator[_ProfilerEvent]: + yield from _utils.traverse_dfs(self._root_nodes, *args, **kwargs) + + @property + def sorted_nodes(self) -> tuple[_ProfilerEvent, ...]: + return self._sorted_nodes + + +class SizeMap: + def __init__(self, op_tree: OpTree) -> None: + self._values: dict[TensorKey, int] = {} + + for node in op_tree.sorted_nodes: + if node.typed[0] == _EventType.TorchOp: + for t in self._flat_tensor_inputs(node.typed[1]): + self._update_values(t) + + elif node.typed[0] == _EventType.PyCall: + typed_fields = node.typed[1] + if ( + typed_fields.module is not None + and typed_fields.optimizer is not None + ): + raise AssertionError("module and optimizer cannot both be set") + if typed_fields.module is not None: + for _, p, p_grad in typed_fields.module.parameters: + self._update_values(p) + self._update_values(p_grad) + + if typed_fields.optimizer is not None: + for p, p_grad, state in typed_fields.optimizer.parameters: + self._update_values(p) + self._update_values(p_grad) + for _, t in state: + self._update_values(t) + + allocations: dict[TensorKey, int] = {} + for node in op_tree.sorted_nodes: + if node.typed[0] == _EventType.Allocation: + alloc_fields = node.typed[1] + key = TensorKey.from_allocation(alloc_fields) + if key: + new_size = abs(alloc_fields.alloc_size) + prior_size = allocations.setdefault(key, new_size) + + # It is possible to resize Storage in PyTorch, however we + # key on data pointer so most resizes will be treated as a + # change in storage. The one corner case that cannot be + # handled is `realloc` which successfully resizes the + # storage. At time of writing this is not done anywhere in + # the core PyTorch codebase. + if prior_size != new_size: + delta = f"{prior_size} vs. {new_size}" + log.warning("Mismatch between allocation and free: %s", delta) + + self._values.update(allocations) + + def _update_values(self, t: _TensorMetadata | None) -> None: + key = TensorKey.from_tensor(t) + if key is not None and t is not None and t.layout == torch.strided: + # Scalars are represented as zero dim Tensors + n = max( + i[0] * i[1] for i in zip(t.sizes or [1], t.strides or [1], strict=True) + ) + + num_bytes = n * _element_size(t.dtype) + if num_bytes < 0: + raise AssertionError(f"num_bytes must be non-negative, got {num_bytes}") + self._values[key] = max(self._values.get(key, 0), num_bytes) + + @staticmethod + def _flat_tensor_inputs(op: _ExtraFields_TorchOp) -> Iterator[_TensorMetadata]: + for i in op.inputs: + if isinstance(i, _TensorMetadata): + yield i + elif isinstance(i, list): + yield from i + + def __getitem__(self, key: TensorKey): + return self._values[key] + + +@dataclasses.dataclass() +class DataFlowEdge: + input_version: int | None = None + mutated: bool | None = False + + @property + def is_allocation(self) -> bool: + return self.input_version is None + + @property + def is_deletion(self) -> bool: + return self.mutated is None + + +class DataFlowNode: + def __init__(self, event: _ProfilerEvent, graph: DataFlowGraph) -> None: + self._event = event + self._graph = graph + self._edges: dict[TensorKey, DataFlowEdge] = self._determine_edges() + + for key, edge in self._edges.items(): + if edge.mutated and not edge.is_allocation: + self._graph.bump(key) + + # Make sure the version bumping behavior matches what we expect. + versions = {k: (v, self._graph.lookup(k)) for k, v in self.outputs.items()} + if not all(i == j for i, j in versions.values()): + raise AssertionError(f"version mismatch: {versions}, {self._edges}") + + def _determine_edges(self) -> dict[TensorKey, DataFlowEdge]: + subtree = tuple(_utils.traverse_dfs([self._event])) + + # Start by populating edges from op inputs and outputs. + mutable_by_key: dict[TensorKey | None, set[bool | None]] = {} + for op in (i.typed[1] for i in subtree if i.typed[0] == _EventType.TorchOp): + for op_input, mutable in zip( + op.inputs, SchemaMatcher.inputs_are_mutable(op), strict=True + ): + # Tensor + if isinstance(op_input, _TensorMetadata): + key = TensorKey.from_tensor(op_input) + mutable_by_key.setdefault(key, set()).add(mutable) + + # TensorList + elif isinstance(op_input, list): + for op_input_i in op_input: + key = TensorKey.from_tensor(op_input_i) + mutable_by_key.setdefault(key, set()).add(mutable) + + edges: collections.defaultdict[TensorKey | None, DataFlowEdge] + edges = collections.defaultdict(DataFlowEdge) + for key, mutable_set in mutable_by_key.items(): + if key is not None: + edges[key].input_version = self._graph.lookup(key) if key else -1 + + # We consider an op to be mutated if we encounter a schema where it + # is a mutable argument OR if it is ambiguous. (We never explicitly + # see it in any schema.) + mutated = (True in mutable_set) or (tuple(mutable_set) == (None,)) + edges[key].mutated = mutated + + # Then handle deletions. Note that deleting a Tensor implicitly adds + # it as an input edge. + for i in subtree: + if i.typed[0] == _EventType.Allocation and i.typed[1].alloc_size < 0: + key = TensorKey.from_allocation(i.typed[1]) + edge = edges[key] + if key is not None and edge.mutated is None: + raise AssertionError(f"Double delete: {key}") + edge.mutated = None + edge.input_version = self._graph.lookup(key) if key else -1 + + # And finally handle allocations. This step must be last, because the + # previous two steps optimistically add input edges. + for i in subtree: + if i.typed[0] == _EventType.Allocation and i.typed[1].alloc_size > 0: + edges[TensorKey.from_allocation(i.typed[1])].input_version = None + + # We don't need to sort the inputs, but it makes debugging and unit tests nicer. + return dict(sorted((k, v) for k, v in edges.items() if k is not None)) + + @property + def inputs(self) -> dict[TensorKey, tuple[bool, int]]: + return { + # MyPy can't see through `is_allocation` to know that + # `v.input_version` is not None. + k: (bool(v.mutated), cast(int, v.input_version)) + for k, v in self._edges.items() + if not v.is_allocation + } + + @property + def outputs(self) -> dict[TensorKey, int]: + return { + k: 0 if v.input_version is None else v.input_version + 1 + for k, v in self._edges.items() + if (v.is_allocation and not v.is_deletion) or v.mutated + } + + @property + def intermediates(self) -> tuple[TensorKey, ...]: + return tuple( + k for k, v in self._edges.items() if v.is_allocation and v.is_deletion + ) + + @property + def start_time(self) -> int: + return self._event.start_time_ns + + +class DataFlowGraph: + def __init__(self, op_tree: OpTree) -> None: + self._op_tree = op_tree + self._leaf_events = self._extract_leaf_events(op_tree) + self._active_version: dict[TensorKey, int | None] = {} + self._flow_nodes = [DataFlowNode(e, self) for e in self.leaf_events] + self._flow_nodes.sort(key=lambda x: x.start_time) + self.validate() + + @property + def flow_nodes(self) -> tuple[DataFlowNode, ...]: + return tuple(self._flow_nodes) + + def validate(self) -> None: + # Check that each (Tensor, version) pair has a unique creation node + outputs: set[tuple[TensorKey, int]] = set() + for node in self.flow_nodes: + node_outputs = set(node.outputs.items()) + duplicates = outputs & node_outputs + if duplicates: + raise AssertionError( + f"duplicate outputs: {node._event.name} {node._edges} {duplicates}" + ) + outputs |= node_outputs + + # And check that `self._nodes` forms a valid topologically sorted DAG. + tensor_versions: dict[TensorKey, int] = {} + for node in self.flow_nodes: + for key, (_, version) in node.inputs.items(): + expected = tensor_versions.get(key, 0) + if expected != version: + raise AssertionError( + f"version mismatch for input: expected {expected}, got {version}" + ) + + for key, version in node.outputs.items(): + prior_version = tensor_versions.get(key, version) + if version < prior_version: + raise AssertionError( + f"version regression: {version} < {prior_version}" + ) + tensor_versions[key] = version + + @property + def leaf_events(self) -> tuple[_ProfilerEvent, ...]: + return self._leaf_events + + @staticmethod + def _extract_leaf_events(op_tree: OpTree) -> tuple[_ProfilerEvent, ...]: + """Partially traverse the op tree and extract top level ops. + + Consider the following code: + ``` + with record_function("My annotation"): + x.zero_() + y.zero_() + ``` + + The op tree (assuming no Autograd) will look like: + + TorchOp: "My annotation" + TorchOp: zero_ + TorchOp: fill_ + TorchOp: zero_ + TorchOp: fill_ + + The recursive structure of operator calls makes data flow unwieldy. + In order to simplify analysis we would like to select the highest level + ops to represent in the graph. In this case those are the `zero_` ops; + the fact that `fill_` is called is an implementation detail. We also + do not want to group everything under "My annotation" as this could + create overly coarse bundles and lose critical semantics. + + To address this issue we walk over the graph and select the topmost + torch ops ** which match at least one operator schema **. These form + the leaves of the first pass through the op tree. (As well as any + allocations or frees which do are not part of a kernel.) These events + form the logical nodes in our data flow graph. + """ + + leaf_events: list[_ProfilerEvent] = [] + + def leaf_op(e: _ProfilerEvent) -> bool: + return e.typed[0] == _EventType.TorchOp and ( + e.typed[1].scope == RecordScope.BACKWARD_FUNCTION + or bool(SchemaMatcher.match_schemas(e.typed[1])) + ) + + def children_fn(e: _ProfilerEvent): + if leaf_op(e) or e.tag == _EventType.Allocation: + leaf_events.append(e) + return [] + + return e.children + + for _ in op_tree.dfs(children_fn=children_fn): + pass + + return tuple(sorted(leaf_events, key=lambda x: x.start_time_ns)) + + def lookup(self, key: TensorKey) -> int: + version = self._active_version.setdefault(key, 0) + if version is None: + raise AssertionError(f"version for key {key} is None") + return version + + def bump(self, key: TensorKey) -> None: + prior_version = self._active_version.get(key, None) + if prior_version is None: + raise AssertionError(f"prior_version for key {key} is None") + self._active_version[key] = prior_version + 1 + + def delete(self, key: TensorKey) -> None: + if self._active_version.setdefault(key, 0) is None: + raise AssertionError(f"cannot delete key {key}, already deleted") + self._active_version[key] = None + + +@dataclasses.dataclass +class CategoryElement: + by_id: Category | None = None + by_key: dict[TensorKey, Category] = dataclasses.field(default_factory=dict) + by_version: dict[TensorAndID, Category] = dataclasses.field(default_factory=dict) + + # Used by unit tests to check internals. (And consequently by + # MemoryProfile.lookup) This should not be used in any other capacity. + _by_id_keyset: set[TensorKey] = dataclasses.field(default_factory=set) + + +@dataclasses.dataclass +class CategoryDict: + _values: collections.defaultdict[int, CategoryElement] = dataclasses.field( + default_factory=lambda: collections.defaultdict(CategoryElement) + ) + + def set_by_id(self, key: TensorKey, category: Category) -> None: + self._values[key.id].by_id = category + self._values[key.id]._by_id_keyset.add(key) + + def set_by_key(self, key: TensorKey, category: Category) -> None: + self._values[key.id].by_key[key] = category + + def set_by_version(self, key: TensorKey, version: int, category: Category) -> None: + self._values[key.id].by_version[(key, version)] = category + + def setdefault_by_version( + self, key: TensorKey, version: int, category: Category + ) -> None: + self._values[key.id].by_version.setdefault((key, version), category) + + def get(self, key: Key, version: int) -> Category | None: + if isinstance(key, Key) and not isinstance(key, TensorKey): + return None + element = self._values[key.id] + return ( + element.by_id + or element.by_key.get(key, None) + or element.by_version.get((key, version), None) + ) + + +class MemoryProfile: + def __init__(self, result: _ProfilerResult) -> None: + self._op_tree = OpTree(result) + self._data_flow_graph = DataFlowGraph(self._op_tree) + self._size_map = SizeMap(self._op_tree) + self._categories = CategoryDict() + + self._set_gradients_and_temporaries() + self._set_parameters_using_python_tracer() + self._set_inputs() + self._set_parameters_using_data_flow() + self._set_activations() + self._set_optimizer_state() + self._set_autograd_detail() + + @property + def timeline(self) -> tuple[tuple[int, Action, KeyAndID, int], ...]: + output: list[tuple[int, Action, KeyAndID, int]] = [] + allocation_times: dict[tuple[TensorKey, bool], int] = {} + live_unknown: dict[tuple[int, torch.device], Literal[True]] = {} + + for event in self._op_tree.dfs(): + if event.typed[0] == _EventType.Allocation: + alloc_fields = event.typed[1] + alloc_size = alloc_fields.alloc_size + is_allocation = alloc_size > 0 + t = event.start_time_ns + + tkey = TensorKey.from_allocation(alloc_fields) + if tkey is not None: + allocation_times[(tkey, is_allocation)] = t + + else: + key = Key(alloc_fields.device) + ptr_and_device = (alloc_fields.ptr, key.device) + if is_allocation: + if ptr_and_device in live_unknown: + output.append( + (t, Action.INCREMENT_VERSION, (key, 0), alloc_size) + ) + else: + live_unknown[ptr_and_device] = True + output.append((t, Action.CREATE, (key, 0), alloc_size)) + else: + output.append((t, Action.DESTROY, (key, 0), -alloc_size)) + if not live_unknown.pop(ptr_and_device, False): + output.append( + (-1, Action.PREEXISTING, (key, 0), -alloc_size) + ) + + snapshot = self._category_snapshot() + last_version = dict(sorted(snapshot.keys())) + + events: list[tuple[int, Action, TensorAndID]] = [ + (-1, Action.PREEXISTING, (key, version)) + for key, version in snapshot + if (key, True) not in allocation_times and version == 0 + ] + + for node in self._data_flow_graph.flow_nodes: + for key, edge in node._edges.items(): + if edge.is_allocation: + t = allocation_times[(key, True)] + events.append((t, Action.CREATE, (key, 0))) + + elif edge.mutated: + t = node._event.start_time_ns + version = edge.input_version + if version is None: + raise AssertionError(f"input_version is None for key {key}") + events.append((t, Action.INCREMENT_VERSION, (key, version))) + + if edge.is_deletion: + t = allocation_times[(key, False)] + events.append((t, Action.DESTROY, (key, last_version[key]))) + + output.extend( + (time, action, (key, version), self._size_map[key]) + for time, action, (key, version) in events + ) + + output.sort(key=lambda x: (x[0], x[1].value)) + return tuple(output) + + def _is_gradient(self, *args, **kwargs) -> bool: + return self._categories.get(*args, **kwargs) == Category.GRADIENT + + def _category_snapshot(self) -> dict[TensorAndID, Category | None]: + all_tensor_versions: set[TensorAndID] = set() + + for node in self._data_flow_graph.flow_nodes: + all_tensor_versions.update(((k, v) for k, (_, v) in node.inputs.items())) + all_tensor_versions.update((key, 0) for key in node.intermediates) + all_tensor_versions.update(node.outputs.items()) + + for i in self._categories._values.values(): + all_tensor_versions.update((key, 0) for key in i._by_id_keyset) + + return { + (key, version): self._categories.get(key, version) + for key, version in sorted(all_tensor_versions) + } + + def _any_version_depends_on_gradient(self) -> set[int]: + """Extract IDs of Tensors which depend or will depend on a gradient. + + Note that this weakened definition of "depends" requires us to loop + over the data flow graph multiple times because it allows dependency + information to flow backward through edges and removes the guarantee + that nodes are topologically sorted. (Or indeed, even that a valid + topological order exists.) Put another way, we have converted an + acyclic data flow graph into a cyclic graph and we are attempting to + partition cycles involving a gradient from the rest of the graph. + """ + depends_on_gradient: set[int] = set() + while True: + start_size = len(depends_on_gradient) + for node in self._data_flow_graph.flow_nodes: + ids = tuple( + key.id + for key, (_, version) in node.inputs.items() + if self._categories.get(key, version) + in (Category.GRADIENT, Category.PARAMETER) + or key.id in depends_on_gradient + ) + + if ids: + depends_on_gradient.update(ids) + + depends_on_gradient.update(key.id for key in node.outputs) + + # We are guaranteed to exit because there is a finite set of + # TensorAndID pairs. In practice we do not expect to loop more than + # three times: once to identify the core parameter update loop, + # once to fold the first step into that loop, and a third time + # where no new elements are added. + if len(depends_on_gradient) == start_size: + return depends_on_gradient + + def _set_gradients_and_temporaries(self) -> None: + """Mark Tensors which are unambiguous and simple to reason about.""" + + # Gradients are straightforward to detect. We directly check the + # `.grad` property in the Python tracer, and we can detect any new + # gradient Tensors from `AccumulateGrad` ops. + for event in self._op_tree.dfs(): + for _, p_grad in extract_gradients(event): + self._categories.set_by_id(p_grad, Category.GRADIENT) + + # Similarly, temporary Tensors are easy to identify and are useful to + # flag since they can make memory use "spikier" than one would + # otherwise expect. + for node in self._data_flow_graph.flow_nodes: + for i in node.intermediates: + self._categories.set_by_key(i, Category.TEMPORARY) + + def _set_parameters_using_python_tracer(self) -> None: + for event in self._op_tree.dfs(): + for p in extract_parameters(event): + if p is not None: + self._categories.set_by_id(p, Category.PARAMETER) + + def _set_inputs(self) -> None: + """Mark inputs based on which Tensors are updated using gradients. + + The process for differentiating between inputs and activations is more + involved. Most Tensors in a training loop depend on at least one + gradient: parameters depend on them through updates, and activations + and optimizer state depend on them transitively through parameters. + Critically, we do not need to know which Tensors are parameters to + apply this method; we can simply walk the data flow graph to build the + set of all values which depend on a gradient and then obtain the set + of inputs from the conjugate set. + + There is, however, one hiccup. The first time we see a parameter is + generally on the forward pass of the first step. We know from + inspection of the data flow graph that v1 of that Tensor depends on + a gradient (provided we profile an optimizer step), but not v0. To + address this problem we weaken the definition of "depends on a + gradient" to "any version of this Tensor depends on a gradient", + which in turn strengthens the criteria for the input set enough to + filter the activations in the forward pass of the first step.""" + + # All of this analysis is predicated on using at least one training + # step (or parameters from the python tracer) to partition the graph. + # Absent that we cannot determine which Tensors are inputs and which + # ones are part of the model. + depends_on_gradient = self._any_version_depends_on_gradient() + + # We only want to annotate Tensors which actually contribute to the + # model calculation. + produces_gradient: set[TensorAndID] = set() + for node in reversed(self._data_flow_graph.flow_nodes): + tensors = {(key, version) for key, (_, version) in node.inputs.items()} + tensors |= node.outputs.items() + if any( + self._categories.get(*i) in (Category.GRADIENT, Category.PARAMETER) + or i in produces_gradient + for i in tensors + ): + produces_gradient |= tensors + + # Don't include Tensors created in the backward pass, as these are + # generally Autograd implementation details rather than proper inputs. + input_candidates = produces_gradient.copy() + for node in self._data_flow_graph.flow_nodes: + if RecordScope.BACKWARD_FUNCTION in get_scopes(node._event): + input_candidates -= set(node.outputs.items()) + + for key, version in input_candidates: + if key.id not in depends_on_gradient: + self._categories.setdefault_by_version(key, version, Category.INPUT) + + def _set_parameters_using_data_flow(self) -> None: + """Deduce which Tensors are parameters. + + Consider the following code for the step of SGD with momentum + (nesterov=False), where `d_p` is the gradient of `param` and `buf` is + the momentum buffer. + ``` + buf.mul_(momentum).add_(d_p, alpha=1 - dampening) + d_p = buf + param.add_(d_p, alpha=-lr) + ``` + Both `param` and `buf` take a gradient and perform an in-place update. + + The python tracer will inspect calls to `nn.Module.forward` and + `optim.Optimizer.step` to extract parameter and optimizer state + respectively (including parameters), so this is generally a non-issue. + + However as a fallback we can also exploit several properties of + parameters to distinguish them from other model state. + + First, they are directly used in the forward pass. (At this point we + haven't established which parts of the graph correspond to the forward + pass but we can deduce enough to suffice.) Some mutable state such as + batch norm moving averages also contribute to the forward pass, but + optimizer state does not. + + Second, a parameter is by definition used to compute at least one + gradient and depends on at least one gradient. + """ + snapshot = self._category_snapshot() + + # Determine which Tensors might be parameters based on forward pass + # data flow. Note this these are only candidates; we filter nodes that + # we know are part of the backward pass but that doesn't guarantee that + # they are part of the forward pass. + candidate_parameters: set[TensorAndID] = set() + candidate_fwd_tensors: set[TensorAndID] = { + i for i, category in snapshot.items() if category == Category.INPUT + } + + for node in self._data_flow_graph.flow_nodes: + inputs = {(key, value) for key, (_, value) in node.inputs.items()} + if ( + # Don't check nodes in the backward pass. + RecordScope.BACKWARD_FUNCTION not in get_scopes(node._event) + and not any(self._is_gradient(*i) for i in inputs) + and not any(self._is_gradient(*i) for i in node.outputs.items()) + # + # and only check nodes which depend on an input. + and candidate_fwd_tensors.intersection(inputs) + ): + candidate_fwd_tensors |= node.outputs.items() + candidate_parameters |= inputs.difference(candidate_fwd_tensors) + + # Require that each parameter eventually contributes to the value of a gradient + used_for_gradient: set[TensorAndID] = set() + for node in reversed(self._data_flow_graph.flow_nodes): + if any( + self._is_gradient(*i) or i in used_for_gradient + for i in node.outputs.items() + ): + used_for_gradient.update( + (key, version) for key, (_, version) in node.inputs.items() + ) + candidate_parameters.intersection_update(used_for_gradient) + + # and depends on a gradient. + parameter_keys = {key.id for key, _ in candidate_parameters} + parameter_keys &= self._any_version_depends_on_gradient() + + for key, _ in snapshot: + if key.id in parameter_keys: + self._categories.set_by_id(key, Category.PARAMETER) + + def _set_activations(self) -> None: + """Flood the graph to identify activations.""" + + required = {Category.INPUT, Category.ACTIVATION} + also_allowed = {Category.PARAMETER, Category.TEMPORARY} + for node in self._data_flow_graph.flow_nodes: + inputs = {(key, value) for key, (_, value) in node.inputs.items()} + input_categories = {self._categories.get(*i) for i in inputs} + + if ( + (input_categories & required) + and not (input_categories - (required | also_allowed)) + # + # Stop filling when we reach the backward pass. + and RecordScope.BACKWARD_FUNCTION not in get_scopes(node._event) + ): + for i in node.outputs.items(): + self._categories.setdefault_by_version(*i, Category.ACTIVATION) + + def _set_optimizer_state(self) -> None: + for event in self._op_tree.dfs(): + if event.typed[0] == _EventType.PyCall and event.typed[1].optimizer: + parameters = event.typed[1].optimizer.parameters + for _, t in it.chain.from_iterable( + (state for _, _, state in parameters) + ): + key = TensorKey.from_tensor(t) + if key is not None: + self._categories.set_by_id(key, Category.OPTIMIZER_STATE) + + def _set_autograd_detail(self) -> None: + prior = {None, Category.AUTOGRAD_DETAIL} + for node in self._data_flow_graph.flow_nodes: + if RecordScope.BACKWARD_FUNCTION in get_scopes(node._event): + for key, version in node.outputs.items(): + if version == 0 or self._categories.get(key, version - 1) in prior: + self._categories.setdefault_by_version( + key, version, Category.AUTOGRAD_DETAIL + ) + + +class MemoryProfileTimeline: + def __init__(self, memory_profile) -> None: + """The minimum representation of the memory profile timeline + includes the memory timeline and categories. The timeline + consists of [timestamp, action, (TensorKey, version), numbytes] + elements, to denote any actions (pre-existing, create, destroy, + or increment_version) that occurred to a specific Tensor for a + chunk of memory. The categories help map each (TensorKey, + version) pair into a category.""" + self.timeline = memory_profile.timeline + self.categories = memory_profile._categories + + def _coalesce_timeline(self, device_str): + """Convert the memory timeline and categories into a memory plot + consisting of timestamps and their respective sizes by category + for a given device. + + Input: device + Output: [timestamps, sizes by category] + """ + device = torch.device(device_str) + times: list[int] = [] + sizes: list[list[int]] = [] + + def update(key, version, delta) -> None: + category = ( + self.categories.get(key, version) + if isinstance(key, TensorKey) + else None + ) + index = _CATEGORY_TO_INDEX[category] + 1 + sizes[-1][index] += int(delta) + + t_min = -1 + for t, action, (key, version), numbytes in self.timeline: + if key.device != device: + continue + + # Convert timestamps from ns to us, to match trace events. + if t != -1: + t = int(t / 1000) + + # Save the smallest timestamp to populate pre-existing allocs. + if t_min == -1 or (t < t_min and t > 0): + t_min = t + + # Handle timestep + if len(times) == 0: + times.append(t) + sizes.append([0] + [0 for _ in _CATEGORY_TO_INDEX]) + + elif t != times[-1]: + times.append(t) + sizes.append(sizes[-1].copy()) + + # Handle memory and categories + if action in (Action.PREEXISTING, Action.CREATE): + update(key, version, numbytes) + + elif action == Action.INCREMENT_VERSION: + update(key, version, -numbytes) + update(key, version + 1, numbytes) + + elif action == Action.DESTROY: + update(key, version, -numbytes) + + else: + raise ValueError(f"Unknown action: {action}") + + times = [t_min if t < 0 else t for t in times] + return times, sizes + + def export_memory_timeline(self, path, device_str) -> None: + """Saves the memory timeline as [times, sizes by category] + as a JSON formatted file to the given path for the given + device.""" + times, sizes = self._coalesce_timeline(device_str) + # TODO: Write a faster serialize (orjson not available in CI) + import json + + with open(path, "w") as f: + json.dump([times, sizes], f) + + def export_memory_timeline_raw(self, path, device_str) -> None: + """Saves the memory timeline as raw memory event tuples in the + form of (timestamp, action, numbytes, category) + as a JSON formatted file to the given path for the given + device.""" + device = torch.device(device_str) + raw_events: list[tuple[int, int, int, int]] = [] + + def get_category_index(key, version): + category = ( + self.categories.get(key, version) + if isinstance(key, TensorKey) + else None + ) + return _CATEGORY_TO_INDEX[category] + + for t, action, (key, version), numbytes in self.timeline: + if key.device != device: + continue + + if action in (Action.PREEXISTING, Action.CREATE): + raw_events.append( + # pyrefly: ignore [bad-argument-type] + ( + t, + _ACTION_TO_INDEX[action], + numbytes, + get_category_index(key, version), + ) + ) + + elif action == Action.INCREMENT_VERSION: + raw_events.append( + # pyrefly: ignore [bad-argument-type] + ( + t, + _ACTION_TO_INDEX[action], + -numbytes, + get_category_index(key, version), + ) + ) + raw_events.append( + # pyrefly: ignore [bad-argument-type] + ( + t, + _ACTION_TO_INDEX[action], + numbytes, + get_category_index(key, version + 1), + ) + ) + + elif action == Action.DESTROY: + raw_events.append( + # pyrefly: ignore [bad-argument-type] + ( + t, + _ACTION_TO_INDEX[action], + -numbytes, + get_category_index(key, version), + ) + ) + + else: + raise ValueError(f"Unknown action: {action}") + + import json + + with open(path, "w") as f: + json.dump(raw_events, f) + + def export_memory_timeline_html( + self, path, device_str, figsize=(20, 12), title=None + ) -> None: + """Exports the memory timeline as an HTML file which contains + the memory timeline plot embedded as a PNG file.""" + # Check if user has matplotlib installed, return gracefully if not. + import importlib.util + + matplotlib_spec = importlib.util.find_spec("matplotlib") + if matplotlib_spec is None: + print( + "export_memory_timeline_html failed because matplotlib was not found." + ) + return + + from base64 import b64encode + from tempfile import NamedTemporaryFile + + import matplotlib.pyplot as plt + import numpy as np + + mt = self._coalesce_timeline(device_str) + times, sizes = np.array(mt[0]), np.array(mt[1]) + # For this timeline, start at 0 to match Chrome traces. + t_min = min(times) + times -= t_min + stacked = np.cumsum(sizes, axis=1) / 1024**3 + device = torch.device(device_str) + max_memory_allocated = torch.cuda.max_memory_allocated(device) + max_memory_reserved = torch.cuda.max_memory_reserved(device) + + # Plot memory timeline as stacked data + fig = plt.figure(figsize=figsize, dpi=80) + axes = fig.gca() + for category, color in _CATEGORY_TO_COLORS.items(): + i = _CATEGORY_TO_INDEX[category] + axes.fill_between( + times / 1e3, stacked[:, i], stacked[:, i + 1], color=color, alpha=0.7 + ) + fig.legend(["Unknown" if i is None else i.name for i in _CATEGORY_TO_COLORS]) + # Usually training steps are in magnitude of ms. + axes.set_xlabel("Time (ms)") + axes.set_ylabel("Memory (GB)") + title = "\n\n".join( + ([title] if title else []) + + [ + f"Max memory allocated: {max_memory_allocated / (1024**3):.2f} GiB \n" + f"Max memory reserved: {max_memory_reserved / (1024**3):.2f} GiB" + ] + ) + axes.set_title(title) + + # Embed the memory timeline image into the HTML file + with NamedTemporaryFile("w+b", suffix=".png") as tmpfile: + fig.savefig(tmpfile, format="png") + + tmpfile.seek(0, 0) + encoded = b64encode(tmpfile.read()).decode("utf-8") + if not encoded: + raise AssertionError("failed to encode image as base64") + html = f""" +GPU Memory Timeline HTML + + + +""" + + with open(path, "w", encoding="utf-8") as f: + f.write(html) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/_pattern_matcher.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/_pattern_matcher.py new file mode 100644 index 0000000000000000000000000000000000000000..988df7ce8d900a3bfc4651f633c252627f171dc6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/_pattern_matcher.py @@ -0,0 +1,681 @@ +# mypy: allow-untyped-defs +import json +import math +import os +import re + +import torch +import torch.utils.benchmark as benchmark +from torch._C._profiler import ( + _EventType, + _ExtraFields_PyCall, + _ExtraFields_PyCCall, + _ExtraFields_TorchOp, + _ProfilerEvent, +) +from torch.profiler import profile +from torch.profiler._utils import index_of_first_match, traverse_bfs, traverse_dfs + + +class Pattern: + """ + Base class for all patterns, subclass this class and implement match() + to define custom patterns. + + In subclass, define description and skip property. + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + self.prof = prof + self.should_benchmark = should_benchmark + self.name = "Please specify a name for pattern" + self.description = "Please specify a description for pattern" + self.url = "" + if prof.profiler is None or prof.profiler.kineto_results is None: + raise AssertionError("profiler and kineto_results must not be None") + self.event_tree = prof.profiler.kineto_results.experimental_event_tree() + self.tid_root: dict[int, list[_ProfilerEvent]] = {} + for event in self.event_tree: + self.tid_root.setdefault(event.start_tid, []).append(event) + + @property + def skip(self) -> bool: + return False + + def report(self, event: _ProfilerEvent): + msg = ( + f"{self.description}\n[Source Code Location] {source_code_location(event)}" + ) + return msg + + def eventTreeTraversal(self): + """ + Traverse the event tree and yield all events. + Override this method in subclass to customize the traversal. + """ + yield from traverse_dfs(self.event_tree) + + def summary(self, events: list[_ProfilerEvent]): + default_summary = f"{self.name}: {len(events)} events matched." + if self.should_benchmark: + # If benchmark summary is not empty, use it. + return ( + self.benchmark_summary(events) + if hasattr(self, "benchmark") # type: ignore[attr-defined] + else default_summary + ) + return default_summary + + def benchmark_summary(self, events: list[_ProfilerEvent]) -> str: + def format_time(time_ns: int) -> str: + unit_lst = ["ns", "us", "ms"] + for unit in unit_lst: + if time_ns < 1000: + return f"{time_ns:.2f} {unit}" + time_ns //= 1000 + return f"{time_ns:.2f} s" + + if not hasattr(self, "benchmark"): + raise AssertionError("Please implement benchmark()") + shapes_factor_map = self.benchmark(events) # type: ignore[attr-defined] + original_time = sum(event.duration_time_ns for event in events) + new_time = sum( + shapes_factor_map[input_shapes(event)] * event.duration_time_ns + for event in events + ) + return ( + f"{self.name}: {len(events)} events matched. " + f"Total Estimated Speedup: {format_time(original_time - new_time)} ({round(original_time / new_time, 2)}X)" + ) + + def match(self, event: _ProfilerEvent): + """ + Return True if the event matches the pattern. + This method should be overridden in subclass. + """ + raise NotImplementedError + + def matched_events(self): + if self.skip: + return [] + matched_events = [ + event for event in self.eventTreeTraversal() if self.match(event) + ] + return matched_events + + def root_of(self, event: _ProfilerEvent): + while event.parent: + event = event.parent + return event + + def siblings_of(self, event: _ProfilerEvent): + if event.parent: + children = event.parent.children + else: + children = self.tid_root[event.start_tid] + index = children.index(event) + return children[:index], children[index + 1 :] + + def next_of(self, event: _ProfilerEvent): + _, next_events = self.siblings_of(event) + return next_events[0] if next_events else None + + def prev_of(self, event: _ProfilerEvent): + prev_events, _ = self.siblings_of(event) + return prev_events[-1] if prev_events else None + + def go_up_until(self, event: _ProfilerEvent, predicate): + if not event: + return None + while event.parent and not predicate(event): + event = event.parent + return event + + +# Patterns + + +class NamePattern(Pattern): + def __init__( + self, prof: profile, name: str, should_benchmark: bool = False + ) -> None: + super().__init__(prof, should_benchmark) + self.description = f"Matched Name Event: {name}" + self.name = name + + def match(self, event: _ProfilerEvent): + return re.search(self.name, event.name) is not None + + +class ExtraCUDACopyPattern(Pattern): + """ + This pattern identifies if we creates a constant tensor on CPU and immediately moves it to GPU. + example: torch.zeros((100, 100)).to("cuda") + + Pattern: + built-in method |built-in method + ... | aten::to + aten::fill_/aten::zero_ | aten::_to_copy + + Algorithm: + We start at node aten::to, go parent events' previous events, + and check if we have a aten::fill_/aten::zero_ as we keep going down the tree. + We always select the last child in the children list when we go down the tree. + If at any step we failed, it is not a match. + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Extra CUDA Copy Pattern" + self.description = "Filled a CPU tensor and immediately moved it to GPU. Please initialize it on GPU." + self.url = "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#create-tensors-directly-on-the-target-device" + self.init_ops = { + "aten::fill_", + "aten::zero_", + "aten::normal_", + "aten::uniform_", + } + + @property + def skip(self) -> bool: + return not self.prof.with_stack or not self.prof.record_shapes + + def match(self, event): + # TODO: We should also check tensor identities + if event.name != "aten::to": + return False + to_event = event + if not event.children: + return False + event = event.children[-1] + if event.name != "aten::_to_copy": + return False + if not event.children: + return False + event = event.children[-1] + if event.name != "aten::copy_": + return False + # aten::copy_ should have the first 2 args dtype the same + dtypes = input_dtypes(event) + if len(dtypes) < 2: + return False + if dtypes[0] is None or dtypes[0] != dtypes[1]: + return False + event = to_event + # Up one level + event = event.parent + if event is None: + return False + # Check if we have a aten::fill_ in previous leaf + event = self.prev_of(event) + if event is None: + return False + while event.children: + event = event.children[-1] + # aten::zero_ is a special optimization case where fill_ is not called + if event.name in self.init_ops: + return True + return event.name in self.init_ops + # TODO: Check if tensor is reused + + def benchmark(self, events: list[_ProfilerEvent]): + shapes_factor_map = {input_shapes(event): 0.0 for event in events} + for shape in shapes_factor_map: + size = shape[0] + to_timer = benchmark.Timer( + stmt='torch.ones(size).to("cuda")', globals={"size": size} + ) + de_timer = benchmark.Timer( + stmt='torch.ones(size, device="cuda")', globals={"size": size} + ) + to_time = to_timer.timeit(10).mean + de_time = de_timer.timeit(10).mean + shapes_factor_map[shape] = de_time / to_time + return shapes_factor_map + + +class ForLoopIndexingPattern(Pattern): + """ + This pattern identifies if we use a for loop to index a tensor that + can be vectorized. + example: + tensor = torch.empty((100, 100)) + for i in range(100): + tensor[i] = i + + Pattern: + aten::select | ... | aten::select | ... (Repeat) + + Algorithm: + We start at node aten::select, and we check if we can find this alternating patterns. + We also keep a dictionary to avoid duplicate match in the for loop. + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "For Loop Indexing Pattern" + self.description = "For loop indexing detected. Vectorization recommended." + self.visited: set[int] = set() + + def eventTreeTraversal(self): + """ + We need to use BFS traversal order to avoid duplicate match. + """ + yield from traverse_bfs(self.event_tree) + + def match(self, event: _ProfilerEvent): + if event.name != "aten::select": + return False + if event.id in self.visited: + return False + repeat_count = 1 + _, next = self.siblings_of(event) + if len(next) <= 1: + return False + + # Custom event list matching + def same_ops(list1, list2) -> bool: + if len(list1) != len(list2): + return False + for op1, op2 in zip(list1, list2, strict=True): + if op1.name != op2.name: + return False + return True + + # Record the ops between two aten::select + next_select_idx = index_of_first_match(next, lambda e: e.name == "aten::select") + if next_select_idx is None: + return False + indexing_ops = [event] + next[:next_select_idx] + next = next[len(indexing_ops) - 1 :] + for i in range(0, len(next), len(indexing_ops)): + if same_ops(indexing_ops, next[i : i + len(indexing_ops)]): + repeat_count += 1 + self.visited.add(next[i].id) + else: + break + return repeat_count >= 10 + + +class FP32MatMulPattern(Pattern): + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "FP32 MatMul Pattern" + self.description = ( + "You are currently using GPU that supports TF32. " + "Please enable TF32 by setting 'torch.backends.cuda.matmul.allow_tf32 = True'" + ) + self.url = "https://pytorch.org/docs/stable/notes/cuda.html#tensorfloat-32-tf32-on-ampere-devices" + + @property + def skip(self): + if torch.version.hip is not None: + has_tf32 = False + else: + # Anything less than sm_80 is not Ampere which doesn't support TF32 + has_tf32 = all( + int(re.sub("sm_|compute_", "", arch)) >= 80 + for arch in torch.cuda.get_arch_list() + ) + return has_tf32 is False or super().skip or not self.prof.record_shapes + + def match(self, event: _ProfilerEvent) -> bool: + # If we saw this pattern once, we don't need to match it again + if event.tag != _EventType.TorchOp: + return False + if not isinstance(event.extra_fields, _ExtraFields_TorchOp): + raise AssertionError( + f"expected _ExtraFields_TorchOp, got {type(event.extra_fields).__name__}" + ) + if event.name == "aten::mm": + if event.extra_fields.allow_tf32_cublas is False: + return True + return False + + def report(self, event: _ProfilerEvent): + return self.description + + def benchmark(self, events: list[_ProfilerEvent]): + shapes_factor_map = {input_shapes(event): 0.0 for event in events} + for shape in shapes_factor_map: + matrixA = torch.randn(shape[0], device="cuda", dtype=torch.float32) + matrixB = torch.randn(shape[1], device="cuda", dtype=torch.float32) + fp32_timer = benchmark.Timer( + stmt="torch.mm(matrixA, matrixB)", + globals={"matrixA": matrixA, "matrixB": matrixB}, + ) + tf32_timer = benchmark.Timer( + stmt="torch.mm(matrixA, matrixB)", + setup="torch.backends.cuda.matmul.allow_tf32 = True", + globals={"matrixA": matrixA, "matrixB": matrixB}, + ) + torch.backends.cuda.matmul.allow_tf32 = False + fp32_time = fp32_timer.timeit(10).mean + tf32_time = tf32_timer.timeit(10).mean + shapes_factor_map[shape] = tf32_time / fp32_time + return shapes_factor_map + + +class OptimizerSingleTensorPattern(Pattern): + """ + This pattern identifies if we are using the single-tensor version of an optimizer. + example: + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + By adding foreach=True to enable multi-tensor optimizer, we can gain speedup when + the kernels are relatively small. + + Pattern: + XXXXX: _single_tenser_ + + Algorithm: + String match + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Optimizer Single Tensor Pattern" + self.optimizers_with_foreach = ["adam", "sgd", "adamw"] + self.description = ( + "Detected optimizer running with single tensor implementation. " + "Please enable multi tensor implementation by passing 'foreach=True' into optimizer." + ) + self.url = "" + + def match(self, event: _ProfilerEvent) -> bool: + for optimizer in self.optimizers_with_foreach: + if event.name.endswith(f"_single_tensor_{optimizer}"): + return True + return False + + +class SynchronizedDataLoaderPattern(Pattern): + """ + This pattern identifies if we are using num_workers=0 in DataLoader. + example: + torch.utils.data.DataLoader(dataset, batch_size=batch_size) + Add num_workers=N to the arguments. N depends on system configuration. + + Pattern: + dataloader.py(...): __iter__ + dataloader.py(...): _get_iterator + NOT dataloader.py(...): check_worker_number_rationality + + Algorithm: + If we don't see check_worker_number_rationality call in the dataloader __iter__, + It is not an asynchronous dataloader. + + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Synchronized DataLoader Pattern" + self.description = ( + "Detected DataLoader running with synchronized implementation. " + "Please enable asynchronous dataloading by setting num_workers > 0 when initializing DataLoader." + ) + self.url = ( + "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html" + "#enable-async-data-loading-and-augmentation" + ) + + def match(self, event: _ProfilerEvent) -> bool: + def is_dataloader_function(name: str, function_name: str): + return name.startswith( + os.path.join("torch", "utils", "data", "dataloader.py") + ) and name.endswith(function_name) + + # TODO: fixme! Due to lifetime issues of the function name, this field might + # actually point to an already freed string when the even is a PyCall. + # Just silently skip this to unblock testing. + try: + event.name + except UnicodeDecodeError: + return False + + if not is_dataloader_function(event.name, "__iter__"): + return False + if not event.children: + return False + event = event.children[0] + if not is_dataloader_function(event.name, "_get_iterator"): + return False + if not event.children: + return False + event = event.children[0] + return not is_dataloader_function(event.name, "check_worker_number_rationality") + # TODO: We should also check if the loader is bottleneck. + + +class GradNotSetToNonePattern(Pattern): + """ + This pattern identifies if we are not setting grad to None in zero_grad. + example: + optimizer.zero_grad() + By setting set_to_none=True, we can gain speedup + + Pattern: + XXXXX: _zero_grad + NOT aten::zeros + aten::zero_ + + aten::zero_ is called on each parameter in the model. + We also want to make sure it is not called by aten::zeros. + + Algorithm: + String match + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Gradient Set To Zero Instead of None Pattern" + self.description = ( + "Detected gradient set to zero instead of None. " + "Please add 'set_to_none=True' when calling zero_grad()." + ) + self.url = ( + "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html" + "#disable-gradient-calculation-for-validation-or-inference" + ) + + def match(self, event: _ProfilerEvent) -> bool: + if not event.name.endswith(": zero_grad"): + return False + if not event.children: + return False + + for sub_event in traverse_dfs(event.children): + if ( + sub_event.name == "aten::zero_" + and sub_event.parent.name != "aten::zeros" + ): + return True + # TODO: We should also check if the optimizer's numerical behavior will change. + return False + + +class Conv2dBiasFollowedByBatchNorm2dPattern(Pattern): + """ + This pattern identifies if we are enabling bias in Conv2d which is followed by BatchNorm2d. + Bias doesn't do anything when followed by batchnorm. + Pattern: + nn.Module: Conv2d | nn.Module: BatchNorm2d + ... + aten::conv2d AND dtype of third argument is not null + The third argument is the bias + Algorithm: + String match + """ + + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Enabling Bias in Conv2d Followed By BatchNorm Pattern" + self.description = "Detected bias enabled in Conv2d that is followed by BatchNorm2d. Please set 'bias=False' in Conv2d." + self.url = ( + "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html" + "#disable-bias-for-convolutions-directly-followed-by-a-batch-norm" + ) + + @property + def skip(self): + return self.prof.record_shapes is False or super().skip + + def match(self, event: _ProfilerEvent): + if event.name != "aten::conv2d": + return False + if len(input_dtypes(event)) < 3 or input_dtypes(event)[2] is None: + return False + # This means bias=True + event = self.go_up_until( + event, lambda e: e.name.startswith("nn.Module: Conv2d") + ) + if not event: + return False + event = self.next_of(event) + if not event: + return False + return event.name.startswith("nn.Module: BatchNorm2d") + + +class MatMulDimInFP16Pattern(Pattern): + def __init__(self, prof: profile, should_benchmark: bool = False) -> None: + super().__init__(prof, should_benchmark) + self.name = "Matrix Multiplication Dimension Not Aligned Pattern" + self.description = "Detected matmul with dimension not aligned. Please use matmul with aligned dimension." + self.url = "https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html#use-mixed-precision-and-amp" + + @property + def skip(self) -> bool: + return not self.prof.with_stack or not self.prof.record_shapes + + def match(self, event: _ProfilerEvent) -> bool: + def mutiple_of(shapes, multiple): + return all(dim % multiple == 0 for shape in shapes for dim in shape[-2:]) + + if event.name not in ("aten::mm", "aten::bmm", "aten::addmm"): + return False + if not input_dtypes(event): + return False + arg_dtype = input_dtypes(event)[0] + if arg_dtype in (torch.bfloat16, torch.half) and not mutiple_of( + input_shapes(event), 8 + ): + return True + return False + + def benchmark(self, events: list[_ProfilerEvent]): + def closest_multiple(shapes, multiple): + return [multiple * math.ceil(shape / multiple) for shape in shapes] + + shapes_factor_map = {input_shapes(event): 0.0 for event in events} + for shape in shapes_factor_map: + matrixA = torch.randn(shape[0], device="cuda", dtype=torch.float16) + matrixB = torch.randn(shape[1], device="cuda", dtype=torch.float16) + not_aligned_dim_timer = benchmark.Timer( + stmt="torch.mm(matrixA, matrixB)", + globals={"matrixA": matrixA, "matrixB": matrixB}, + ) + matrixA = torch.randn( + closest_multiple(shape[0], 8), device="cuda", dtype=torch.float16 + ) + matrixB = torch.randn( + closest_multiple(shape[1], 8), device="cuda", dtype=torch.float16 + ) + aligned_dim_timer = benchmark.Timer( + stmt="torch.mm(matrixA, matrixB)", + globals={"matrixA": matrixA, "matrixB": matrixB}, + ) + not_aligned_dim_time = not_aligned_dim_timer.timeit(10).mean + aligned_dim_time = aligned_dim_timer.timeit(10).mean + shapes_factor_map[shape] = aligned_dim_time / not_aligned_dim_time + return shapes_factor_map + + +def source_code_location(event: _ProfilerEvent | None) -> str: + while event: + if event.tag == _EventType.PyCall or event.tag == _EventType.PyCCall: + if not isinstance( + event.extra_fields, (_ExtraFields_PyCall, _ExtraFields_PyCCall) + ): + raise AssertionError( + f"expected _ExtraFields_PyCall or _ExtraFields_PyCCall, " + f"got {type(event.extra_fields).__name__}" + ) + if not event.extra_fields.caller.file_name.startswith("torch" + os.sep): + return f"{event.extra_fields.caller.file_name}:{event.extra_fields.caller.line_number}" + event = event.parent + return "No source code location found" + + +def input_shapes(event: _ProfilerEvent): + if not isinstance(event.extra_fields, _ExtraFields_TorchOp): + raise AssertionError( + f"expected _ExtraFields_TorchOp, got {type(event.extra_fields).__name__}" + ) + return tuple(tuple(getattr(i, "sizes", ())) for i in event.extra_fields.inputs) + + +def input_dtypes(event: _ProfilerEvent): + if not isinstance(event.extra_fields, _ExtraFields_TorchOp): + raise AssertionError( + f"expected _ExtraFields_TorchOp, got {type(event.extra_fields).__name__}" + ) + return tuple(getattr(i, "dtype", None) for i in event.extra_fields.inputs) + + +def report_all_anti_patterns( + prof, + should_benchmark: bool = False, + print_enable: bool = True, + json_report_dir: str | None = None, +) -> None: + report_dict: dict = {} + anti_patterns = [ + ExtraCUDACopyPattern(prof, should_benchmark), + # ForLoopIndexingPattern(prof, should_benchmark), + FP32MatMulPattern(prof, should_benchmark), + OptimizerSingleTensorPattern(prof, should_benchmark), + SynchronizedDataLoaderPattern(prof, should_benchmark), + GradNotSetToNonePattern(prof, should_benchmark), + Conv2dBiasFollowedByBatchNorm2dPattern(prof, should_benchmark), + MatMulDimInFP16Pattern(prof, should_benchmark), + ] + reported = set() + summaries = [] + message_list = [f"{'-' * 40}TorchTidy Report{'-' * 40}"] + message_list.append("Matched Events:") + + for anti_pattern in anti_patterns: + matched_events = anti_pattern.matched_events() + if not matched_events: + continue + summaries.append(anti_pattern.summary(matched_events)) + for event in matched_events: + report_msg = anti_pattern.report(event) + if report_msg not in reported: + message_list.append(report_msg) + reported.add(report_msg) + src_location, line_no = source_code_location(event).split(":") + report_dict.setdefault(src_location, []).append( + { + "line_number": int(line_no), + "name": anti_pattern.name, + "url": anti_pattern.url, + "message": anti_pattern.description, + } + ) + + if json_report_dir is not None: + json_report_path = os.path.join(json_report_dir, "torchtidy_report.json") + if os.path.exists(json_report_path): + with open(json_report_path) as f: + exisiting_report = json.load(f) + exisiting_report.update(report_dict) + report_dict = exisiting_report + with open(json_report_path, "w") as f: + json.dump(report_dict, f, indent=4) + + message_list.append("Summary:") + message_list += summaries + message_list.append(f"{'-' * 40}TorchTidy Report{'-' * 40}") + if print_enable: + print("\n".join(message_list)) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a8d053dc20e74e1979f36c8e4a8b644f9978d55a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/_utils.py @@ -0,0 +1,577 @@ +# mypy: allow-untyped-defs +import functools +import operator +import re +from collections import deque +from dataclasses import dataclass +from typing import Any, Literal, TYPE_CHECKING + +from torch.autograd.profiler import profile +from torch.profiler import DeviceType + + +if TYPE_CHECKING: + from torch.autograd import _KinetoEvent + + +def _traverse(tree, next_fn, children_fn=lambda x: x.children, reverse: bool = False): + order = reversed if reverse else lambda x: x + remaining = deque(order(tree)) + while remaining: + curr_event = next_fn(remaining) + yield curr_event + for child_event in order(children_fn(curr_event)): + remaining.append(child_event) + + +traverse_dfs = functools.partial(_traverse, next_fn=lambda x: x.pop(), reverse=True) +traverse_bfs = functools.partial( + _traverse, next_fn=lambda x: x.popleft(), reverse=False +) + + +@dataclass +class EventMetrics: + duration_time_ns: int = 0 + self_time_ns: int = 0 + idle_time_ns: int = 0 + queue_depth: int = 0 + + @property + def fraction_idle_time(self): + if self.duration_time_ns == 0: + return 0.0 + return self.idle_time_ns / self.duration_time_ns + + +@dataclass +class Interval: + start: int + end: int + queue_depth: int = 0 + + +class EventKey: + def __init__(self, event) -> None: + self.event = event + + def __hash__(self): + return hash(self.event.id) + + def __eq__(self, other): + return self.event.id == other.event.id + + def __repr__(self) -> str: + return f"{self.event.name}" + + def intervals_overlap(self, intervals: list[Interval]): + overlap_time = 0 + intervals = sorted(intervals, key=lambda x: x.start) + + if intervals: + overlap_start = max(self.event.start_time_ns, intervals[0].start) + overlap_end = min(self.event.end_time_ns, intervals[0].end) + + if overlap_start < overlap_end: + overlap_time += overlap_end - overlap_start + + i, j = 0, 1 + while j < len(intervals): + prev_interval = intervals[i] + curr_interval = intervals[j] + j += 1 + if prev_interval.end > curr_interval.start: + # Completely subsumed by previous interval + if prev_interval.end > curr_interval.end: + j += 1 + continue + else: + curr_interval.start = prev_interval.end + i = j + + overlap_start = max(self.event.start_time_ns, curr_interval.start) + overlap_end = min(self.event.end_time_ns, curr_interval.end) + if overlap_start < overlap_end: + overlap_time += overlap_end - overlap_start + + return overlap_time + + +class BasicEvaluation: + def __init__(self, prof: profile) -> None: + self.profile = prof + self.metrics: dict[EventKey, EventMetrics] = {} + self.compute_self_time() + self.event_keys = sorted( + self.metrics.keys(), key=lambda x: x.event.start_time_ns + ) + self.events = [e.event for e in self.event_keys] + self.cuda_events: list[_KinetoEvent] = [] + self.queue_depth_list = self.compute_queue_depth() + self.compute_idle_time() + + def compute_self_time(self) -> None: + """ + Computes event's self time(total time - time in child ops). + """ + if self.profile.kineto_results is None: + raise AssertionError("kineto_results must not be None") + stack = deque(self.profile.kineto_results.experimental_event_tree()) + + # standard iterating dfs + while stack: + curr_event = stack.pop() + self_time = curr_event.duration_time_ns + for child_event in curr_event.children: + self_time -= child_event.duration_time_ns + stack.append(child_event) + if EventKey(curr_event) in self.metrics: + raise AssertionError( + f"Duplicate id: {curr_event.id}, {curr_event.name}" + ) + self.metrics[EventKey(curr_event)] = EventMetrics(self_time_ns=self_time) + self.metrics[ + EventKey(curr_event) + ].duration_time_ns = curr_event.duration_time_ns + + def compute_queue_depth(self): + """ + Computes queue_depth at each event. This will calculate the queue depth data for + All the events in the tree. + This will return a list of Interval of queue depth data of cuda launch and kernels. + """ + if self.profile.kineto_results is None: + raise AssertionError("kineto_results must not be None") + cuda_event_list = self.profile.kineto_results.events() + + def is_cuda_launch_kernel(e): + """Check if the event is a CUDA launch kernel.""" + launch_patterns = { + "cudaLaunchKernel", # Standard CUDA + "cudaLaunchKernelExC", # Extended C + "__cudaLaunchKernel", # Internal + "cudaLaunchCooperativeKernel", # Collaborative (single-device) + "cudaLaunchCooperativeKernelMultiDevice", # Collaborative (multi-devices) + } + name = str(getattr(e, "name", e)) + return any(name.startswith(pattern) for pattern in launch_patterns) + + def is_cuda_kernel(e): + """Check if the event is a CUDA runtime kernel.""" + # Check if the kernel is CUDA + if e.device_type() != DeviceType.CUDA: + return False + + name = str(getattr(e, "name", e)).lower() + + # Exclude memory operations + exclude_patterns = {"mem", "cpy", "alloc", "free"} + + return not any(pattern in name for pattern in exclude_patterns) + + cuda_launch_events = sorted( + (e for e in cuda_event_list if is_cuda_launch_kernel(e)), + key=lambda x: x.start_ns(), + ) + cuda_kernel_events = sorted( + (e for e in cuda_event_list if is_cuda_kernel(e)), + key=lambda x: x.start_ns(), + ) + + self.cuda_events = sorted( + cuda_launch_events + cuda_kernel_events, key=lambda x: x.start_ns() + ) + + kernel_mapping: dict[_KinetoEvent, int] = {} + last_mapped_kernel = 0 + for cuda_launch_event in cuda_launch_events: + index = index_of_first_match( + cuda_kernel_events, + lambda x: x.linked_correlation_id() + == cuda_launch_event.linked_correlation_id(), + start=last_mapped_kernel, + ) + kernel_mapping[cuda_launch_event] = index + last_mapped_kernel = index if index is not None else last_mapped_kernel + + current_kernel_index = 0 + spawned_kernel_index = -1 + + all_events = cuda_launch_events + cuda_kernel_events + self.events + + def new_old_event_comparator(event): + if hasattr(event, "start_us"): + return event.start_us() * 1000 + if hasattr(event, "start_ns"): + return event.start_ns() + if hasattr(event, "start_time_ns"): + return event.start_time_ns + raise Exception("Unknown Event Type") # noqa: TRY002 + + queue_depth_list: list[Interval] = [] + all_events.sort(key=new_old_event_comparator) + for event in all_events: + # Find latest cuda kernel event + if hasattr(event, "start_us"): + start_time = event.start_us() * 1000 + # pyrefly: ignore [missing-attribute] + end_time = (event.start_us() + event.duration_us()) * 1000 + # Find current spawned cuda kernel event + if event in kernel_mapping and kernel_mapping[event] is not None: + spawned_kernel_index = kernel_mapping[event] + if hasattr(event, "start_ns"): + start_time = event.start_ns() + end_time = event.start_ns() + event.duration_ns() + # Find current spawned cuda kernel event + if event in kernel_mapping and kernel_mapping[event] is not None: + spawned_kernel_index = kernel_mapping[event] + elif hasattr(event, "start_time_ns"): + start_time = event.start_time_ns # type: ignore[attr-defined] + end_time = event.end_time_ns # type: ignore[attr-defined] + + while ( + current_kernel_index < len(cuda_kernel_events) + and (cuda_kernel_events[current_kernel_index].start_ns()) <= start_time # type: ignore[possibly-undefined] + ): + current_kernel_index += 1 + current_queue_depth = spawned_kernel_index - current_kernel_index + 1 + current_queue_depth = max(current_queue_depth, 0) + + if hasattr(event, "start_us") or hasattr(event, "start_ns"): + queue_depth_list.append( + Interval(start_time, end_time, current_queue_depth) # type: ignore[possibly-undefined] + ) + elif hasattr(event, "start_time_ns"): + self.metrics[EventKey(event)].queue_depth = current_queue_depth + + return queue_depth_list + + def compute_idle_time(self) -> None: + """ + Computes idle time of the profile. + """ + # Based on queue_depth_list, we can calculate idle time for all the events + idle = False + idle_start = 0 + idle_intervals: list[Interval] = [] + if self.queue_depth_list and self.events: + idle_intervals += [ + Interval(self.events[0].start_time_ns, self.queue_depth_list[0].start), + Interval(self.queue_depth_list[-1].end, self.events[-1].end_time_ns), + ] + + for data_point in self.queue_depth_list: + if data_point.queue_depth == 0 and not idle: + idle_start = data_point.end + idle = True + if data_point.queue_depth > 0 and idle: + idle_intervals.append(Interval(idle_start, data_point.start)) + idle = False + + event_list = [e.event for e in self.metrics] + for event in event_list: + self.metrics[EventKey(event)].idle_time_ns = EventKey( + event + ).intervals_overlap(idle_intervals) + + def rank_events(self, length): + """ + Filter and Rank the events based on some heuristics: + 1) Events that are in the falling phase of the queue depth. + 2) Events that have a high idle_time, self_time difference. + + Parameters: + length: The number of events to return. + """ + + # Find the interval when qd is falling to 0 + import torch + + queue_depth_list = list(reversed(self.queue_depth_list)) + qd_values = [e.queue_depth for e in queue_depth_list] + + bottom_threashold = 0 + top_threashold = 4 + decrease_interval = [] + i = 0 + while i < len(qd_values): + if qd_values[i] > bottom_threashold: + i += 1 + continue + for j in range(i + 1, len(qd_values)): + # Find next zero and if the max value between them exceeds + # the threshold, then we have a falling interval + next_minimum_idx = index_of_first_match( + qd_values, lambda x: x <= bottom_threashold, start=j + ) + peak_idx = argmax(qd_values, start=j, end=next_minimum_idx) + + # if is a valid peak, we add to list and continue + if peak_idx is not None and qd_values[peak_idx] >= top_threashold: + decrease_interval.append( + Interval( + queue_depth_list[peak_idx].start, queue_depth_list[i].start + ) + ) + i = next_minimum_idx if next_minimum_idx is not None else i + break + i += 1 + # Filter out events that are not in the decrease interval + event_list = [ + event + for event in self.metrics + if event.intervals_overlap(decrease_interval) + ] + if event_list: + self_time = torch.tensor( + [self.metrics[event].self_time_ns for event in event_list], + dtype=torch.float32, + ) + idle_time = torch.tensor( + [self.metrics[event].fraction_idle_time for event in event_list], + dtype=torch.float32, + ) + normalized_gain = (idle_time - torch.mean(idle_time)) / torch.std(idle_time) + normalized_self = (self_time - torch.mean(self_time)) / torch.std(self_time) + heuristic_score_list = normalized_gain + 0.6 * normalized_self + + # Sort events by heuristic + event_list = [ + event + for _, event in sorted( + zip(heuristic_score_list, event_list, strict=True), + key=operator.itemgetter(0), + reverse=True, + ) + ] + event_list = event_list[:length] + return event_list + + def get_optimizable_events(self, length: int = 1, print_enable: bool = True): + event_list = self.rank_events(length) + if not print_enable: + return event_list + output = "Optimizable events:\n" if event_list else "No events to optimize\n" + + output += "\n".join( + [ + f"""{"-" * 80} +Event: {event} +Source code location: {source_code_location(event.event)} +Percentage idle time: {self.metrics[event].fraction_idle_time * 100:.2f}% +{"-" * 80}""" + for event in event_list + ] + ) + if print_enable: + print(output) + return event_list + + +def index_of_first_match(seq, predicate, start=0, end=None): + if end is None or end >= len(seq): + end = len(seq) + for i in range(start, end): + if predicate(seq[i]): + return i + return None + + +def argmax(seq, key=lambda x: x, start=0, end=None): + seq = seq[start:end] + if len(seq) == 0: + return None + return seq.index(max(seq, key=key)) + start + + +def source_code_location(event): + while event is not None: + match = re.search(r"\.py\(.*\)", event.name) + if match is None: + event = event.parent + continue + return event.name + return "No source code location found" + + +# Provide an OSS workaround for cudagraphs + CUPTI issue +# https://github.com/pytorch/pytorch/issues/75504 +# TODO(dberard) - deprecate / remove workaround for CUDA >= 12, when +# we stop supporting older CUDA versions. +def _init_for_cuda_graphs() -> None: + from torch.autograd.profiler import profile + + with profile(): + pass + + +@dataclass +class TimelineEvent: + """Represents an event in the profiler timeline.""" + + timestamp: int + event_type: Literal["start", "end", "regular"] + marker_type: Literal["filename", "node"] | None + identifier: str | int | None + event: dict[str, Any] + + +@dataclass +class ContextStackEntry: + """Represents a context (filename or node) in the stack.""" + + context_type: Literal["filename", "node"] + identifier: str | int + metadata: dict | None + tid: int | None = None # Thread ID associated with this context + + +def map_recorded_events_to_aten_ops_with_stack_trace(traced_data): + """ + Maps recorded profiler events to their corresponding fx nodes and adds stack traces. + + Builds a timeline of all events (regular ops and FX markers for filenames/nodes), + sorts by timestamp, then processes chronologically while maintaining a context stack of active + filename/node scopes. Regular events are augmented with stack traces and node names from the + innermost active context. Runtime is O(n log n) for n events. + + Args: + traced_data: Json of profiler events from Chrome trace + + Returns: + Dict mapping recorded event names to their aten operations with added stack traces + """ + from torch.fx.traceback import _FX_METADATA_REGISTRY + + trace_events = traced_data.get("traceEvents", []) + + # Create event timeline + event_timeline: list[TimelineEvent] = [] + + def is_fx_marker_event(event): + return ( + event.get("cat") == "cpu_op" + and event.get("name", "").startswith("## ") + and event.get("name", "").endswith(" ##") + ) + + def append_fx_marker_event(event_type, identifier, event): + start_ts = event["ts"] + end_ts = start_ts + event["dur"] + event_timeline.append( + TimelineEvent(start_ts, "start", event_type, identifier, event) + ) + event_timeline.append( + TimelineEvent(end_ts, "end", event_type, identifier, event) + ) + + for event in trace_events: + if "ts" not in event or "dur" not in event: + continue + + if is_fx_marker_event(event): + content = event["name"][3:-3] + + if content.endswith(".py"): + append_fx_marker_event("filename", content, event) + else: + try: + node_index = int(content) + except ValueError: + pass + append_fx_marker_event("node", node_index, event) # type: ignore[possibly-undefined] + + else: + # Regular event that needs augmentation + start_ts = event["ts"] + event_timeline.append(TimelineEvent(start_ts, "regular", None, None, event)) + + # Sort by timestamp + event_timeline.sort(key=lambda x: x.timestamp) + + # Process events in chronological order with a stack + context_stack: list[ContextStackEntry] = [] + + # Invariant: all start event has a corresponding end event + for timeline_event in event_timeline: + match timeline_event.event_type: + case "start": + if timeline_event.identifier is None: + raise AssertionError("identifier must not be None for start event") + + if timeline_event.marker_type == "filename": + if not isinstance(timeline_event.identifier, str): + raise AssertionError( + f"identifier must be str for filename marker, " + f"got {type(timeline_event.identifier).__name__}" + ) + # Push filename context - query metadata registry on-demand + metadata = _FX_METADATA_REGISTRY.get(timeline_event.identifier) + tid = timeline_event.event.get("tid") + context_stack.append( + ContextStackEntry( + "filename", timeline_event.identifier, metadata, tid + ) + ) + elif timeline_event.marker_type == "node": + # Find the current filename from stack + current_file_metadata = None + tid = timeline_event.event.get("tid") + for ctx_entry in reversed(context_stack): + if ( + ctx_entry.context_type == "filename" + and ctx_entry.tid == tid + ): + current_file_metadata = ctx_entry.metadata + break + + if current_file_metadata: + node_metadata = current_file_metadata.get("node_metadata", {}) + if timeline_event.identifier in node_metadata: + node_meta: dict | None = node_metadata[ + timeline_event.identifier + ] + context_stack.append( + ContextStackEntry( + "node", timeline_event.identifier, node_meta, tid + ) + ) + + case "end": + # Pop from stack - search backwards to find matching context + for i in range(len(context_stack) - 1, -1, -1): + ctx_entry = context_stack[i] + if ( + timeline_event.marker_type == ctx_entry.context_type + and timeline_event.identifier == ctx_entry.identifier + ): + context_stack.pop(i) + break + + case "regular": + # Apply metadata from current context stack + # Find the most specific context (node takes precedence over filename) + # Only augment events with the same tid as the file/node event matched + current_stack_trace = None + current_node_name = None + event_tid = timeline_event.event.get("tid") + + for ctx_entry in reversed(context_stack): + # Only apply metadata from contexts with matching tid + if ctx_entry.tid == event_tid: + if ctx_entry.context_type == "node" and ctx_entry.metadata: + current_stack_trace = ctx_entry.metadata.get( + "stack_trace", "No model stack trace available" + ) + current_node_name = ctx_entry.metadata.get("name", "") + # Do we want to only attach the stack trace of the lowest node or stack trace of all nodes + # if nodes are nested, e.g. in nested graph modules + break + + # Augment the event + if current_stack_trace or current_node_name: + args = timeline_event.event.setdefault("args", {}) + if current_stack_trace: + args["stack_trace"] = current_stack_trace + if current_node_name: + args["node_name"] = current_node_name diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/itt.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/itt.py new file mode 100644 index 0000000000000000000000000000000000000000..7b1a6eac0f0bc8d69988fe59f7acae26302572c6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/itt.py @@ -0,0 +1,81 @@ +# mypy: allow-untyped-defs +from contextlib import contextmanager +from typing import NoReturn + + +try: + from torch._C import _itt +except ImportError: + + class _ITTStub: + @staticmethod + def _fail(*args, **kwargs) -> NoReturn: + raise RuntimeError( + "ITT functions not installed. Are you sure you have a ITT build?" + ) + + @staticmethod + def is_available() -> bool: + return False + + rangePush = _fail + rangePop = _fail + mark = _fail + + _itt = _ITTStub() # type: ignore[assignment] + + +__all__ = ["is_available", "range_push", "range_pop", "mark", "range"] + + +def is_available(): + """ + Check if ITT feature is available or not + """ + return _itt.is_available() + + +def range_push(msg): + """ + Pushes a range onto a stack of nested range span. Returns zero-based + depth of the range that is started. + + Arguments: + msg (str): ASCII message to associate with range + """ + return _itt.rangePush(msg) + + +def range_pop(): + """ + Pops a range off of a stack of nested range spans. Returns the + zero-based depth of the range that is ended. + """ + return _itt.rangePop() + + +def mark(msg): + """ + Describe an instantaneous event that occurred at some point. + + Arguments: + msg (str): ASCII message to associate with the event. + """ + return _itt.mark(msg) + + +@contextmanager +def range(msg, *args, **kwargs): + """ + Context manager / decorator that pushes an ITT range at the beginning + of its scope, and pops it at the end. If extra arguments are given, + they are passed as arguments to msg.format(). + + Args: + msg (str): message to associate with the range + """ + range_push(msg.format(*args, **kwargs)) + try: + yield + finally: + range_pop() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/profiler.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..2dc5a63cb33f47ede3b3f241beb16c1f29094a60 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/profiler.py @@ -0,0 +1,1255 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import gzip +import json +import os +import shutil +import tempfile +from abc import ABC, abstractmethod +from enum import Enum +from functools import partial +from typing import Any, TYPE_CHECKING +from typing_extensions import deprecated, Self +from warnings import warn + +import torch +import torch.autograd.profiler as prof +from torch._C import _get_privateuse1_backend_name +from torch._C._profiler import ( + _add_execution_trace_observer, + _disable_execution_trace_observer, + _enable_execution_trace_observer, + _ExperimentalConfig, + _remove_execution_trace_observer, +) +from torch._environment import is_fbcode +from torch._utils_internal import profiler_allow_cudagraph_cupti_lazy_reinit_cuda12 +from torch.autograd import kineto_available, ProfilerActivity +from torch.profiler._memory_profiler import MemoryProfile, MemoryProfileTimeline + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable + + +__all__ = [ + "supported_activities", + "ProfilerAction", + "schedule", + "tensorboard_trace_handler", + "profile", + "ExecutionTraceObserver", +] +PROFILER_STEP_NAME = "ProfilerStep" + +_WARNINGS_SHOWN = set() + + +def _warn_once(msg, category=UserWarning, stacklevel=2): + if msg not in _WARNINGS_SHOWN: + _WARNINGS_SHOWN.add(msg) + warn(msg, category=category, stacklevel=stacklevel) + + +class _NumpyEncoder(json.JSONEncoder): + """ + Json encoder for numpy types (np.int, np.float, np.array etc.) + Returns default encoder if numpy is not available + """ + + def default(self, obj): + """Encode NumPy types to JSON""" + try: + import numpy as np + except ImportError: + return json.JSONEncoder.default(self, obj) + if isinstance(obj, np.integer): + return int(obj) + elif isinstance(obj, np.floating): + return float(obj) + elif isinstance(obj, np.ndarray): + return obj.tolist() + else: + return json.JSONEncoder.default(self, obj) + + +def supported_activities(): + """ + Returns a set of supported profiler tracing activities. + + Note: profiler uses CUPTI library to trace on-device CUDA kernels. + In case when CUDA is enabled but CUPTI is not available, passing + ``ProfilerActivity.CUDA`` to profiler results in using the legacy CUDA + profiling code (same as in the legacy ``torch.autograd.profiler``). + This, in turn, results in including CUDA time in the profiler table output, + but not in the JSON trace. + """ + return torch.autograd._supported_activities() + + +class _ITraceObserver(ABC): + """Abstract interface for a Trace observer. + This satisfies 3 methods: start, stop and cleanup""" + + @abstractmethod + def start(self): + pass + + @abstractmethod + def stop(self): + pass + + @abstractmethod + def cleanup(self): + pass + + +def _parse_activities( + activities: Iterable[ProfilerActivity | dict[ProfilerActivity, list[str]]], +) -> tuple[set[ProfilerActivity], dict[ProfilerActivity, set[str]]]: + """Parse a mixed activities list into a set of activities and a filter dict. + + Each item is either a bare ``ProfilerActivity`` (collect all defaults) or a + ``dict[ProfilerActivity, list[str]]`` (collect only the named subset). + An empty list value (e.g. ``{CUDA: []}``) means collect nothing for that group. + """ + parsed_activities: set[ProfilerActivity] = set() + activity_filters: dict[ProfilerActivity, set[str]] = {} + for item in activities: + if isinstance(item, ProfilerActivity): + if item in parsed_activities: + raise ValueError(f"Activity {item} specified more than once") + parsed_activities.add(item) + elif isinstance(item, dict): + for key, val in item.items(): + if key in parsed_activities: + raise ValueError(f"Activity {key} specified more than once") + parsed_activities.add(key) + activity_filters[key] = set(val) + else: + raise TypeError(f"Expected ProfilerActivity or dict, got {type(item)}") + return parsed_activities, activity_filters + + +class _KinetoProfile: + """Low-level profiler wrap the autograd profile + + Args: + activities (iterable): list of activity groups (CPU, CUDA) to use in profiling, supported values: + ``torch.profiler.ProfilerActivity.CPU``, ``torch.profiler.ProfilerActivity.CUDA``, + ``torch.profiler.ProfilerActivity.XPU``. + Default value: ProfilerActivity.CPU and (when available) ProfilerActivity.CUDA + or (when available) ProfilerActivity.XPU. + + Each item can be a ``ProfilerActivity`` enum (collects all default + activity types for that group) or a ``dict`` mapping a ``ProfilerActivity`` + to a list of individual activity type names to collect, e.g. + ``{ProfilerActivity.CUDA: ["GPU_MEMCPY", "CUDA_RUNTIME"]}``. + An empty list (e.g. ``{ProfilerActivity.CUDA: []}``) means collect + nothing for that group. + The same activity group must not appear more than once. + record_shapes (bool): save information about operator's input shapes. + profile_memory (bool): track tensor memory allocation/deallocation (see ``export_memory_timeline`` + for more details). + with_stack (bool): record source information (file and line number) for the ops. + with_flops (bool): use formula to estimate the FLOPS of specific operators + (matrix multiplication and 2D convolution). + with_modules (bool): record module hierarchy (including function names) + corresponding to the callstack of the op. e.g. If module A's forward call's + module B's forward which contains an aten::add op, + then aten::add's module hierarchy is A.B + Note that this support exist, at the moment, only for TorchScript models + and not eager mode models. + experimental_config (_ExperimentalConfig) : A set of experimental options + used by profiler libraries like Kineto. Note, backward compatibility is not guaranteed. + execution_trace_observer (ExecutionTraceObserver) : A PyTorch Execution Trace Observer object. + `PyTorch Execution Traces `__ offer a graph based + representation of AI/ML workloads and enable replay benchmarks, simulators, and emulators. + When this argument is included the observer start() and stop() will be called for the + same time window as PyTorch profiler. + acc_events (bool): Enable the accumulation of FunctionEvents across multiple profiling cycles + post_processing_timeout_s (float): Optional timeout in seconds for post-processing profiler + results. In this context, post-processing happens after the profiling itself has finished. + If specified, event parsing will stop after this duration and return partial results. Useful + for handling large traces that may take too long to process. + + + .. note:: + This API is experimental and subject to change in the future. + + Enabling shape and stack tracing results in additional overhead. + When record_shapes=True is specified, profiler will temporarily hold references to the tensors; + that may further prevent certain optimizations that depend on the reference count and introduce + extra tensor copies. + """ + + def __init__( + self, + *, + activities: Iterable[ProfilerActivity | dict[ProfilerActivity, list[str]]] + | None = None, + record_shapes: bool = False, + profile_memory: bool = False, + with_stack: bool = False, + with_flops: bool = False, + with_modules: bool = False, + experimental_config: _ExperimentalConfig | None = None, + execution_trace_observer: _ITraceObserver | None = None, + acc_events: bool = False, + custom_trace_id_callback: Callable[[], str] | None = None, + post_processing_timeout_s: float | None = None, + ) -> None: + if activities is not None: + self.activities, self.activity_filters = _parse_activities(activities) + else: + self.activities = supported_activities() + self.activity_filters: dict[ProfilerActivity, set[str]] = {} + self.record_shapes = record_shapes + self.with_flops = with_flops + self.profile_memory = profile_memory + self.with_stack = with_stack + self.with_modules = with_modules + self.experimental_config = experimental_config + self.execution_trace_observer = execution_trace_observer + self.acc_events = acc_events + self.custom_trace_id_callback = custom_trace_id_callback + self.post_processing_timeout_s = post_processing_timeout_s + self.profiler: prof.profile | None = None + self.has_cudagraphs = False + self.mem_tl: MemoryProfileTimeline | None = None + self.use_device = None + if ProfilerActivity.CUDA in self.activities: + # pyrefly: ignore [bad-assignment] + self.use_device = "cuda" + elif ProfilerActivity.XPU in self.activities: + # pyrefly: ignore [bad-assignment] + self.use_device = "xpu" + elif ProfilerActivity.MTIA in self.activities: + # pyrefly: ignore [bad-assignment] + self.use_device = "mtia" + elif ProfilerActivity.HPU in self.activities: + # pyrefly: ignore [bad-assignment] + self.use_device = "hpu" + elif ProfilerActivity.PrivateUse1 in self.activities: + # pyrefly: ignore [bad-assignment] + self.use_device = _get_privateuse1_backend_name() + + # user-defined metadata to be amended to the trace + self.preset_metadata: dict[str, str] = {} + + def start(self) -> None: + self.prepare_trace() + self.start_trace() + + def stop(self) -> None: + self.stop_trace() + + def prepare_trace(self) -> None: + if hasattr(torch, "_inductor"): + import torch._inductor.config as inductor_config + + self.has_cudagraphs = inductor_config.triton.cudagraphs + if (self.profiler is None) or (not self.acc_events): + self.profiler = prof.profile( + use_cpu=(ProfilerActivity.CPU in self.activities), + use_device=self.use_device, + record_shapes=self.record_shapes, + with_flops=self.with_flops, + profile_memory=self.profile_memory, + with_stack=self.with_stack, + with_modules=self.with_modules, + use_kineto=True, + experimental_config=self.experimental_config, + acc_events=self.acc_events, + custom_trace_id_callback=self.custom_trace_id_callback, + post_processing_timeout_s=self.post_processing_timeout_s, + activity_filters=self.activity_filters + if self.activity_filters + else None, + ) + if (self.profiler is not None) and (not self.acc_events): + _warn_once( + "Warning: Profiler clears events at the end of each cycle." + "Only events from the current cycle will be reported." + "To keep events across cycles, set acc_events=True." + ) + self.profiler._prepare_trace() + + def start_trace(self) -> None: + if self.execution_trace_observer: + self.execution_trace_observer.start() + if self.profiler is None: + raise AssertionError("Profiler must be initialized before starting trace") + self.profiler._start_trace() + + if self.profile_memory: + self.add_metadata_json("profile_memory", "1") + if self.with_stack: + self.add_metadata_json("with_stack", "1") + if self.record_shapes: + self.add_metadata_json("record_shapes", "1") + if self.with_modules: + self.add_metadata_json("with_modules", "1") + if self.with_flops: + self.add_metadata_json("with_flops", "1") + + if kineto_available(): + dist_info = self._get_distributed_info() + if dist_info: + self.add_metadata_json( + "distributedInfo", json.dumps(dist_info, cls=_NumpyEncoder) + ) + + cuda_version = None + if hasattr(torch, "version"): + from torch.torch_version import TorchVersion + + cuda_version = TorchVersion(getattr(torch.version, "cuda", "0.0")) + + if self.has_cudagraphs and ( + (cuda_version and cuda_version < "12.6") + or not profiler_allow_cudagraph_cupti_lazy_reinit_cuda12() + ): + os.environ["DISABLE_CUPTI_LAZY_REINIT"] = "1" + self.add_metadata_json("DISABLE_CUPTI_LAZY_REINIT", "1") + # FIXME: CUDA Graph does not work well with CUPTI teardown. + # 1) crashes on 1st lazy CUPTI re-init after teardown (CUDA 11) + # 2) crashes on 2nd non-lazy CUPTI re-init after teardown (CUDA 12) + # Workaround: turn off CUPTI teardown when using CUDA Graphs. + os.environ["TEARDOWN_CUPTI"] = "0" + + # Insert the preset user metadata to the trace + for k, v in self.preset_metadata.items(): + self.add_metadata_json(k, v) + + def stop_trace(self) -> None: + if self.execution_trace_observer: + self.execution_trace_observer.stop() + if self.profiler is None: + raise AssertionError("Profiler must be initialized before stopping trace") + self.profiler.__exit__(None, None, None) + + def export_chrome_trace(self, path: str): + """ + Exports the collected trace in Chrome JSON format. If kineto is enabled, only + last cycle in schedule is exported. + """ + if self.profiler is None: + raise AssertionError( + "Profiler must be initialized before exporting chrome trace" + ) + if path.endswith(".gz"): + with tempfile.NamedTemporaryFile("w+b", suffix=".json") as fp: + retvalue = self.profiler.export_chrome_trace(fp.name) + with open(fp.name, "rb") as fin, gzip.open(path, "wb") as fout: + fout.writelines(fin) + return retvalue + else: + return self.profiler.export_chrome_trace(path) + + def export_stacks(self, path: str, metric: str = "self_cpu_time_total"): + """Save stack traces to a file + + Args: + path (str): save stacks file to this location; + metric (str): metric to use: "self_cpu_time_total" or "self_cuda_time_total" + """ + if self.profiler is None: + raise AssertionError("Profiler must be initialized before exporting stacks") + return self.profiler.export_stacks(path, metric) + + def toggle_collection_dynamic( + self, enable: bool, activities: Iterable[ProfilerActivity] + ) -> None: + """Toggle collection of activities on/off at any point of collection. Currently supports toggling Torch Ops + (CPU) and CUDA activity supported in Kineto + + Args: + activities (iterable): list of activity groups to use in profiling, supported values: + ``torch.profiler.ProfilerActivity.CPU``, ``torch.profiler.ProfilerActivity.CUDA`` + Examples: + + .. code-block:: python + + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) as p: + code_to_profile_0() + // turn off collection of all CUDA activity + p.toggle_collection_dynamic(False, [torch.profiler.ProfilerActivity.CUDA]) + code_to_profile_1() + // turn on collection of all CUDA activity + p.toggle_collection_dynamic(True, [torch.profiler.ProfilerActivity.CUDA]) + code_to_profile_2() + print(p.key_averages().table( + sort_by="self_cuda_time_total", row_limit=-1)) + """ + if self.profiler is None: + return + self.profiler.toggle_collection_dynamic(enable, activities) + + def key_averages( + self, + group_by_input_shape: bool = False, + group_by_stack_n: int = 0, + group_by_overload_name: bool = False, + ): + """Averages events, grouping them by operator name and (optionally) input shapes, stack + and overload name. + + .. note:: + To use shape/stack functionality make sure to set record_shapes/with_stack + when creating profiler context manager. + """ + if self.profiler is None: + raise AssertionError( + "Profiler must be initialized before getting key averages" + ) + return self.profiler.key_averages( + group_by_input_shape, group_by_stack_n, group_by_overload_name + ) + + def events(self): + """ + Returns the list of unaggregated profiler events, + to be used in the trace callback or after the profiling is finished + """ + if self.profiler is None: + raise AssertionError("Profiler must be initialized before accessing events") + return self.profiler.function_events + + def add_metadata(self, key: str, value: str) -> None: + """ + Adds a user defined metadata with a string key and a string value + into the trace file + """ + wrapped_value = '"' + value.replace('"', '\\"') + '"' + torch.autograd._add_metadata_json(key, wrapped_value) + + def add_metadata_json(self, key: str, value: str) -> None: + """ + Adds a user defined metadata with a string key and a valid json value + into the trace file + """ + torch.autograd._add_metadata_json(key, value) + + def preset_metadata_json(self, key: str, value: str) -> None: + """ + Preset a user defined metadata when the profiler is not started + and added into the trace file later. + Metadata is in the format of a string key and a valid json value + """ + self.preset_metadata[key] = value + + def _get_distributed_info(self): + import torch.distributed as dist + + if not dist.is_available() or not dist.is_initialized(): + return None + + backend = dist.get_backend() + dist_info = { + "backend": backend, + "rank": dist.get_rank(), + "world_size": dist.get_world_size(), + "pg_count": dist.get_pg_count(), + "pg_config": dist.distributed_c10d._get_all_pg_configs(), + } + if backend == "nccl": + nccl_version = torch.cuda.nccl.version() + # pyrefly: ignore [bad-typed-dict-key, unsupported-operation] + dist_info["nccl_version"] = ".".join(str(v) for v in nccl_version) + return dist_info + + def _memory_profile(self) -> MemoryProfile: + required = ("record_shapes", "profile_memory", "with_stack") + missing = [f"{i}=True" for i in required if not getattr(self, i)] + if missing: + raise ValueError(f"{', '.join(missing)} required for memory profiling.") + + if self.profiler is None or self.profiler.kineto_results is None: + raise AssertionError( + "Profiler and kineto_results must be initialized for memory profiling" + ) + return MemoryProfile(self.profiler.kineto_results) + + @deprecated( + "`export_memory_timeline` is deprecated and will be removed in a future version. " + "Please use `torch.cuda.memory._record_memory_history` and `torch.cuda.memory._export_memory_snapshot` instead.", + category=FutureWarning, + ) + def export_memory_timeline(self, path: str, device: str | None = None) -> None: + """Export memory event information from the profiler collected + tree for a given device, and export a timeline plot. There are 3 + exportable files using ``export_memory_timeline``, each controlled by the + ``path``'s suffix. + + - For an HTML compatible plot, use the suffix ``.html``, and a memory timeline + plot will be embedded as a PNG file in the HTML file. + + - For plot points consisting of ``[times, [sizes by category]]``, where + ``times`` are timestamps and ``sizes`` are memory usage for each category. + The memory timeline plot will be saved a JSON (``.json``) or gzipped JSON + (``.json.gz``) depending on the suffix. + + - For raw memory points, use the suffix ``.raw.json.gz``. Each raw memory + event will consist of ``(timestamp, action, numbytes, category)``, where + ``action`` is one of ``[PREEXISTING, CREATE, INCREMENT_VERSION, DESTROY]``, + and ``category`` is one of the enums from + ``torch.profiler._memory_profiler.Category``. + + Output: Memory timeline written as gzipped JSON, JSON, or HTML. + + .. deprecated:: + ``export_memory_timeline`` is deprecated and will be removed in a future version. + Please use ``torch.cuda.memory._record_memory_history`` and + ``torch.cuda.memory._export_memory_snapshot`` instead. + """ + # Default to device 0, if unset. Fallback on cpu. + if device is None: + if self.use_device and self.use_device != "cuda": + device = self.use_device + ":0" + else: + device = "cuda:0" if torch.cuda.is_available() else "cpu" + + # Construct the memory timeline plot data + self.mem_tl = MemoryProfileTimeline(self._memory_profile()) + + # Depending on the file suffix, save the data as json.gz or json. + # For html, we can embed the image into an HTML file. + if path.endswith(".html"): + self.mem_tl.export_memory_timeline_html(path, device) + elif path.endswith(".gz"): + with tempfile.NamedTemporaryFile("w+t", suffix=".json") as fp: + if path.endswith("raw.json.gz"): + self.mem_tl.export_memory_timeline_raw(fp.name, device) + else: + self.mem_tl.export_memory_timeline(fp.name, device) + with open(fp.name) as fin, gzip.open(path, "wt") as fout: + fout.writelines(fin) + else: + self.mem_tl.export_memory_timeline(path, device) + + +class ProfilerAction(Enum): + """ + Profiler actions that can be taken at the specified intervals + """ + + NONE = 0 + WARMUP = 1 + RECORD = 2 + RECORD_AND_SAVE = 3 + + +def schedule( + *, + wait: int, + warmup: int, + active: int, + repeat: int = 0, + skip_first: int = 0, + skip_first_wait: int = 0, +) -> Callable: + """ + Returns a callable that can be used as profiler ``schedule`` argument. The profiler will skip + the first ``skip_first`` steps, then wait for ``wait`` steps, then do the warmup for the next ``warmup`` steps, + then do the active recording for the next ``active`` steps and then repeat the cycle starting with ``wait`` steps. + The optional number of cycles is specified with the ``repeat`` parameter, the zero value means that + the cycles will continue until the profiling is finished. + + The ``skip_first_wait`` parameter controls whether the first ``wait`` stage should be skipped. + This can be useful if a user wants to wait longer than ``skip_first`` between cycles, but not + for the first profile. For example, if ``skip_first`` is 10 and ``wait`` is 20, the first cycle will + wait 10 + 20 = 30 steps before warmup if ``skip_first_wait`` is zero, but will wait only 10 + steps if ``skip_first_wait`` is non-zero. All subsequent cycles will then wait 20 steps between the + last active and warmup. + """ + + def schedule_fn(step: int) -> ProfilerAction: + if step < 0: + raise AssertionError(f"Step must be non-negative. Got {step}.") + if step < skip_first: + return ProfilerAction.NONE + else: + step -= skip_first + # If wait >> skip_first and we want to grab profiling early, shift left by wait if skip_first_wait is True + if skip_first_wait != 0: + step += wait + num_steps = wait + warmup + active + if repeat > 0 and step / num_steps >= repeat: + return ProfilerAction.NONE + mod_step = step % num_steps + if mod_step < wait: + return ProfilerAction.NONE + elif mod_step < wait + warmup: + return ProfilerAction.WARMUP + else: + return ( + ProfilerAction.RECORD + if mod_step < num_steps - 1 + else ProfilerAction.RECORD_AND_SAVE + ) + + if wait < 0 or warmup < 0 or active <= 0 or repeat < 0 or skip_first < 0: + raise AssertionError( + f"Invalid profiler schedule arguments. Got wait={wait} (need >= 0), warmup={warmup} (need >= 0), " + f"active={active} (need > 0), repeat={repeat} (need >= 0), skip_first={skip_first} (need >= 0)." + ) + if warmup == 0: + warn( + "Profiler won't be using warmup, this can skew profiler results", + stacklevel=2, + ) + return schedule_fn + + +def _default_schedule_fn(_: int) -> ProfilerAction: + """ + Default profiler behavior - immediately starts recording the events, + keeps doing it on every profiler step. + """ + return ProfilerAction.RECORD + + +def tensorboard_trace_handler( + dir_name: str, worker_name: str | None = None, use_gzip: bool = False +): + """ + Outputs tracing files to directory of ``dir_name``, then that directory can be + directly delivered to tensorboard as logdir. + ``worker_name`` should be unique for each worker in distributed scenario, + it will be set to '[hostname]_[pid]' by default. + """ + import socket + import time + + def handler_fn(prof) -> None: + nonlocal worker_name + if not os.path.isdir(dir_name): + try: + os.makedirs(dir_name, exist_ok=True) + except Exception as e: + raise RuntimeError("Can't create directory: " + dir_name) from e + if not worker_name: + worker_name = f"{socket.gethostname()}_{os.getpid()}" + # Use nanosecond here to avoid naming clash when exporting the trace + file_name = f"{worker_name}.{time.time_ns()}.pt.trace.json" + if use_gzip: + file_name = file_name + ".gz" + prof.export_chrome_trace(os.path.join(dir_name, file_name)) + + return handler_fn + + +class profile(_KinetoProfile): + """Profiler context manager. + + Args: + activities (iterable): list of activity groups (CPU, CUDA) to use in profiling, supported values: + ``torch.profiler.ProfilerActivity.CPU``, ``torch.profiler.ProfilerActivity.CUDA``, + ``torch.profiler.ProfilerActivity.XPU``. + Default value: ProfilerActivity.CPU and (when available) ProfilerActivity.CUDA + or (when available) ProfilerActivity.XPU. + + Each item can be a ``ProfilerActivity`` enum (collects all default + activity types for that group) or a ``dict`` mapping a ``ProfilerActivity`` + to a list of individual activity type names to collect, e.g. + ``{ProfilerActivity.CUDA: ["GPU_MEMCPY", "CUDA_RUNTIME"]}``. + An empty list (e.g. ``{ProfilerActivity.CUDA: []}``) means collect + nothing for that group. + The same activity group must not appear more than once. + schedule (Callable): callable that takes step (int) as a single parameter and returns + ``ProfilerAction`` value that specifies the profiler action to perform at each step. + on_trace_ready (Callable): callable that is called at each step when ``schedule`` + returns ``ProfilerAction.RECORD_AND_SAVE`` during the profiling. + record_shapes (bool): save information about operator's input shapes. + profile_memory (bool): track tensor memory allocation/deallocation. + with_stack (bool): record source information (file and line number) for the ops. + with_flops (bool): use formula to estimate the FLOPs (floating point operations) of specific operators + (matrix multiplication and 2D convolution). + with_modules (bool): record module hierarchy (including function names) + corresponding to the callstack of the op. e.g. If module A's forward call's + module B's forward which contains an aten::add op, + then aten::add's module hierarchy is A.B + Note that this support exist, at the moment, only for TorchScript models + and not eager mode models. + experimental_config (_ExperimentalConfig) : A set of experimental options + used for Kineto library features. Note, backward compatibility is not guaranteed. + execution_trace_observer (ExecutionTraceObserver) : A PyTorch Execution Trace Observer object. + `PyTorch Execution Traces `__ offer a graph based + representation of AI/ML workloads and enable replay benchmarks, simulators, and emulators. + When this argument is included the observer start() and stop() will be called for the + same time window as PyTorch profiler. See the examples section below for a code sample. + acc_events (bool): Enable the accumulation of FunctionEvents across multiple profiling cycles + post_processing_timeout_s (float): Optional timeout in seconds for post-processing profiler + results. If specified, event parsing will stop after this duration and return partial + results. Useful for handling large traces that may take too long to process. + use_cuda (bool): + .. deprecated:: 1.8.1 + use ``activities`` instead. + + .. note:: + Use :func:`~torch.profiler.schedule` to generate the callable schedule. + Non-default schedules are useful when profiling long training jobs + and allow the user to obtain multiple traces at the different iterations + of the training process. + The default schedule simply records all the events continuously for the + duration of the context manager. + + .. note:: + Use :func:`~torch.profiler.tensorboard_trace_handler` to generate result files for TensorBoard: + + ``on_trace_ready=torch.profiler.tensorboard_trace_handler(dir_name)`` + + After profiling, result files can be found in the specified directory. Use the command: + + ``tensorboard --logdir dir_name`` + + to see the results in TensorBoard. + For more information, see + `PyTorch Profiler TensorBoard Plugin `__ + + .. note:: + Enabling shape and stack tracing results in additional overhead. + When record_shapes=True is specified, profiler will temporarily hold references to the tensors; + that may further prevent certain optimizations that depend on the reference count and introduce + extra tensor copies. + + + Examples: + + .. code-block:: python + + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ] + ) as p: + code_to_profile() + print(p.key_averages().table(sort_by="self_cuda_time_total", row_limit=-1)) + + Using the profiler's ``schedule``, ``on_trace_ready`` and ``step`` functions: + + .. code-block:: python + + # Non-default profiler schedule allows user to turn profiler on and off + # on different iterations of the training loop; + # trace_handler is called every time a new trace becomes available + def trace_handler(prof): + print( + prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=-1) + ) + # prof.export_chrome_trace("/tmp/test_trace_" + str(prof.step_num) + ".json") + + + with torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + # In this example with wait=1, warmup=1, active=2, repeat=1, + # profiler will skip the first step/iteration, + # start warming up on the second, record + # the third and the forth iterations, + # after which the trace will become available + # and on_trace_ready (when set) is called; + # the cycle repeats starting with the next step + schedule=torch.profiler.schedule(wait=1, warmup=1, active=2, repeat=1), + on_trace_ready=trace_handler, + # on_trace_ready=torch.profiler.tensorboard_trace_handler('./log') + # used when outputting for tensorboard + ) as p: + for iter in range(N): + code_iteration_to_profile(iter) + # send a signal to the profiler that the next iteration has started + p.step() + + The following sample shows how to setup up an Execution Trace Observer (`execution_trace_observer`) + + .. code-block:: python + + with torch.profiler.profile( + ... + execution_trace_observer=( + ExecutionTraceObserver().register_callback("./execution_trace.json") + ), + ) as p: + for iter in range(N): + code_iteration_to_profile(iter) + p.step() + + You can also refer to test_execution_trace_with_kineto() in tests/profiler/test_profiler.py. + Note: One can also pass any object satisfying the _ITraceObserver interface. + """ + + def __init__( + self, + *, + activities: Iterable[ProfilerActivity | dict[ProfilerActivity, list[str]]] + | None = None, + schedule: Callable[[int], ProfilerAction] | None = None, + on_trace_ready: Callable[..., Any] | None = None, + record_shapes: bool = False, + profile_memory: bool = False, + with_stack: bool = False, + with_flops: bool = False, + with_modules: bool = False, + experimental_config: _ExperimentalConfig | None = None, + execution_trace_observer: _ITraceObserver | None = None, + acc_events: bool = False, + # deprecated: + use_cuda: bool | None = None, + custom_trace_id_callback: Callable[[], str] | None = None, + post_processing_timeout_s: float | None = None, + ) -> None: + # Extract activities for the use_cuda deprecation check. + if activities is not None: + activities_set: set[ProfilerActivity] = set() + for item in activities: + if isinstance(item, ProfilerActivity): + activities_set.add(item) + elif isinstance(item, dict): + activities_set.update(item.keys()) + else: + activities_set = supported_activities() + if use_cuda is not None: + warn( + "`use_cuda` is deprecated, use `activities` argument instead", + FutureWarning, + stacklevel=2, + ) + if use_cuda: + activities_set.add(ProfilerActivity.CUDA) + elif ProfilerActivity.CUDA in activities_set: + activities_set.remove(ProfilerActivity.CUDA) + if len(activities_set) == 0: + raise AssertionError("No valid profiler activities found") + + super().__init__( + activities=activities, + record_shapes=record_shapes, + profile_memory=profile_memory, + with_stack=with_stack, + with_flops=with_flops, + with_modules=with_modules, + experimental_config=experimental_config, + execution_trace_observer=execution_trace_observer + or ExecutionTraceObserver.build_execution_trace_obs_from_env(), + acc_events=acc_events, + custom_trace_id_callback=custom_trace_id_callback, + post_processing_timeout_s=post_processing_timeout_s, + ) + + if schedule: + self.schedule = schedule + # add step markers into the trace and table view + self.record_steps = True + else: + self.schedule = _default_schedule_fn + self.record_steps = False + self.on_trace_ready = on_trace_ready + self.step_num = 0 + self.current_action = self.schedule(self.step_num) + self.step_rec_fn: prof.record_function | None = None + + self.action_map: dict[ + tuple[ProfilerAction, ProfilerAction | None], list[Any] + ] = { + # key is (prev_action, current_action), value is action list corresponding to the state pair. + (ProfilerAction.NONE, ProfilerAction.NONE): [], + (ProfilerAction.NONE, ProfilerAction.WARMUP): [self.prepare_trace], + (ProfilerAction.NONE, ProfilerAction.RECORD): [ + self.prepare_trace, + self.start_trace, + ], + (ProfilerAction.NONE, ProfilerAction.RECORD_AND_SAVE): [ + self.prepare_trace, + self.start_trace, + ], + (ProfilerAction.WARMUP, ProfilerAction.NONE): [ + partial(warn, "Incorrect schedule: WARMUP followed by NONE"), + self.start_trace, + self.stop_trace, + ], + (ProfilerAction.WARMUP, ProfilerAction.WARMUP): [], + (ProfilerAction.WARMUP, ProfilerAction.RECORD): [self.start_trace], + (ProfilerAction.WARMUP, ProfilerAction.RECORD_AND_SAVE): [self.start_trace], + (ProfilerAction.RECORD, ProfilerAction.NONE): [ + partial(warn, "Incorrect schedule: RECORD followed by NONE"), + self.stop_trace, + ], + (ProfilerAction.RECORD, ProfilerAction.WARMUP): [ + partial(warn, "Incorrect schedule: RECORD followed by WARMUP"), + self.stop_trace, + ], + (ProfilerAction.RECORD, ProfilerAction.RECORD): [], + (ProfilerAction.RECORD, ProfilerAction.RECORD_AND_SAVE): [], + (ProfilerAction.RECORD_AND_SAVE, ProfilerAction.NONE): [ + self.stop_trace, + self._trace_ready, + ], + (ProfilerAction.RECORD_AND_SAVE, ProfilerAction.WARMUP): [ + self.stop_trace, + self._trace_ready, + self.prepare_trace, + ], + (ProfilerAction.RECORD_AND_SAVE, ProfilerAction.RECORD): [ + self.stop_trace, + self._trace_ready, + self.prepare_trace, + self.start_trace, + ], + (ProfilerAction.RECORD_AND_SAVE, ProfilerAction.RECORD_AND_SAVE): [ + self.stop_trace, + self._trace_ready, + self.prepare_trace, + self.start_trace, + ], + # used for exit action + (ProfilerAction.WARMUP, None): [self.start_trace, self.stop_trace], + (ProfilerAction.RECORD, None): [self.stop_trace, self._trace_ready], + (ProfilerAction.RECORD_AND_SAVE, None): [ + self.stop_trace, + self._trace_ready, + ], + } + # Start tracking increments to profiler step, this will be used + # by Kineto + prof.KinetoStepTracker.init_step_count(PROFILER_STEP_NAME) + + def __enter__(self): + self.start() + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.stop() + prof.KinetoStepTracker.erase_step_count(PROFILER_STEP_NAME) + if self.execution_trace_observer: + self.execution_trace_observer.cleanup() + + def start(self) -> None: + self._transit_action(ProfilerAction.NONE, self.current_action) + if self.record_steps: + self.step_rec_fn = prof.record_function( + "ProfilerStep#" + str(self.step_num) + ) + self.step_rec_fn.__enter__() + + def stop(self) -> None: + if self.record_steps and self.step_rec_fn: + self.step_rec_fn.__exit__(None, None, None) + self._transit_action(self.current_action, None) + + def step(self) -> None: + """ + Signals the profiler that the next profiling step has started. + """ + if self.record_steps and self.step_rec_fn: + self.step_rec_fn.__exit__(None, None, None) + prev_action = self.current_action + self.step_num += 1 + self.current_action = self.schedule(self.step_num) + + self._transit_action(prev_action, self.current_action) + if os.environ.get("KINETO_USE_DAEMON", "") or ( + is_fbcode() and os.environ.get("KINETO_FORCE_STEP_HOOK", "") + ): + prof.KinetoStepTracker.increment_step(PROFILER_STEP_NAME) + + if self.record_steps: + self.step_rec_fn = prof.record_function( + "ProfilerStep#" + str(self.step_num) + ) + self.step_rec_fn.__enter__() + + def set_custom_trace_id_callback(self, callback) -> None: + """ + Sets a callback to be called when a new trace ID is generated. + """ + self.custom_trace_id_callback = callback + + def get_trace_id(self): + """ + Returns the current trace ID. + """ + if self.profiler is None: + return None + return self.profiler.trace_id + + def _trace_ready(self) -> None: + if self.on_trace_ready: + self.on_trace_ready(self) + + def _transit_action(self, prev_action, current_action) -> None: + action_list = self.action_map.get((prev_action, current_action)) + if action_list: + for action in action_list: + action() + + def _stats(self) -> prof._ProfilerStats | None: + if self.profiler is None: + return None + return self.profiler._stats + + +class ExecutionTraceObserver(_ITraceObserver): + """Execution Trace Observer + + Each process can have a single ExecutionTraceObserver instance. The observer + can be added to record function callbacks via calling register_callback() + explicitly. Without calling unregister_callback(), repeated calls to + register_callback() will not add additional observers to record function + callbacks. Once an ExecutionTraceObserver is created, the start() and stop() + methods control when the event data is recorded. + + Deleting or calling unregister_callback() will remove the observer from the + record function callbacks, finalize the output file, and will stop + incurring any overheads. + """ + + def __init__(self) -> None: + """ + Initializes the default states. + """ + self._registered = False + self._execution_trace_running = False + self.extra_resources_collection = False + self.resources_dir: str = "" + self.output_file_path: str = "" + self.output_file_path_observer: str = "" + + def __del__(self) -> None: + """ + Calls unregister_callback() to make sure to finalize outputs. + """ + self.unregister_callback() + + @staticmethod + def build_execution_trace_obs_from_env() -> ExecutionTraceObserver | None: + """ + Returns an ExecutionTraceObserver instance if the environment variable + ENABLE_PYTORCH_EXECUTION_TRACE is set to 1, otherwise returns None. + + Configures the observer to also collect extra resources if the environment variable + ``ENABLE_PYTORCH_EXECUTION_TRACE_EXTRAS=1``. These are resources such as generated kernels, + index tensor data etc. that are required to make the Execution Trace replayable. + """ + if os.environ.get("ENABLE_PYTORCH_EXECUTION_TRACE", "0") == "1": + try: + with tempfile.NamedTemporaryFile( + "w+t", suffix=".et.json", delete=False + ) as fp: + filename = fp.name + except Exception as e: + warn( + f"Execution trace will not be recorded. Exception on creating default temporary file: {e}", + stacklevel=2, + ) + return None + et = ExecutionTraceObserver() + et.register_callback(filename) + # additionally, check if the env requires us to collect extra resources + if os.environ.get("ENABLE_PYTORCH_EXECUTION_TRACE_EXTRAS", "0") == "1": + et.set_extra_resource_collection(True) + else: + et.set_extra_resource_collection(False) + return et + return None + + def set_extra_resource_collection(self, val) -> None: + """ + Collects extra resources such as generated kernels, index tensor data, and any other + metadata that is required to complete the Execution Trace content. + + The caller should call this method with val=True after calling register_callback() if they want + to collect the extra resources. + """ + self.extra_resources_collection = val + if self.extra_resources_collection: + self.get_resources_dir(can_create=True) + return + + def register_callback(self, output_file_path: str) -> Self: + """ + Adds ET observer to record function callbacks. The data will be + written to output_file_path. + """ + + def get_temp_uncompressed_file() -> str: + with tempfile.NamedTemporaryFile("w+b", suffix=".json", delete=False) as fp: + return fp.name + + if not self._registered: + self.output_file_path = output_file_path + if output_file_path.endswith(".gz"): + output_file_path = get_temp_uncompressed_file() + self.output_file_path_observer = output_file_path + self._registered = _add_execution_trace_observer(output_file_path) + return self + + def get_resources_dir(self, can_create=False) -> str | None: + """ + Generates the resources directory for the generated kernels, + or index tensor data or any other metadata that is required + to complete the Execution Trace content. + + The directory is created right where the ET file is being output. + + Only works if the observer has called set_extra_resource_collection(val=True). + + Returns None if the observer is not configured with extra resource collection. + """ + if not self.extra_resources_collection: + return None + if self.resources_dir: + # already created + return self.resources_dir + generated_path = ExecutionTraceObserver.get_resources_dir_for_et_path( + self.output_file_path, create_dir=can_create + ) + if not generated_path: + # could not find of create the resources dir + return None + self.resources_dir = generated_path + return self.resources_dir + + @staticmethod + def get_resources_dir_for_et_path( + trace_path, create_dir: bool = False + ) -> str | None: + work_dir, file_name = os.path.split(trace_path) + resource_dir = os.path.join( + work_dir, os.path.splitext(file_name)[0] + "_resources" + ) + if not os.path.exists(resource_dir): + if create_dir: + try: + os.mkdir(resource_dir) + except Exception: + warn( + f"Execution trace exception when creating {resource_dir}", + stacklevel=2, + ) + return None + else: + return None + return resource_dir + + def unregister_callback(self) -> None: + """ + Removes ET observer from record function callbacks. + """ + + def _save_triton_kernels() -> None: + try: + resource_dir = self.get_resources_dir() + except Exception as e: + warn( + f"Execution trace exception when generating resource directory: {e}", + stacklevel=2, + ) + return + if not resource_dir: + return + + # Save the kernel paths for the generated kernels + from torch._inductor.codecache import PyCodeCache + + kernel_files = [ + v.__file__ + for v in PyCodeCache.modules + if getattr(v, "__file__", None) is not None + ] + + for kernel_file in kernel_files: + if kernel_file is None: + continue + name = os.path.basename(kernel_file) + dst = os.path.join(resource_dir, name) + shutil.copyfile(kernel_file, dst) + + def _save_gz_file(uncompressed_file: str, output_file: str) -> None: + print(f"Execution Trace: compressing {uncompressed_file} to {output_file}") + with open(uncompressed_file, "rb") as fin: + with gzip.open(output_file, "wb") as fout: + fout.writelines(fin) + os.remove(uncompressed_file) + + if self._registered: + self.stop() + + try: + _save_triton_kernels() + except Exception as e: + warn(f"Execution trace failed to save kernels: {e}", stacklevel=2) + + _remove_execution_trace_observer() + if self.output_file_path.endswith("gz"): + _save_gz_file(self.output_file_path_observer, self.output_file_path) + + self._registered = False + + @property + def is_registered(self): + """ + Returns True if the execution trace observer is registered, otherwise False. + """ + return self._registered + + def is_running(self): + """ + Returns True if the observer is running, otherwise False. + """ + return self._execution_trace_running + + def start(self) -> None: + """ + Starts to capture. + """ + if self._registered and not self._execution_trace_running: + _enable_execution_trace_observer() + self._execution_trace_running = True + self._record_pg_config() + + def stop(self) -> None: + """ + Stops to capture. + """ + if self._execution_trace_running: + _disable_execution_trace_observer() + self._execution_trace_running = False + + def cleanup(self) -> None: + """ + Calls unregister_callback() to make sure to finalize outputs. + """ + self.unregister_callback() + + def get_output_file_path(self) -> str | None: + """ + Returns the output file name or None. + """ + if self.output_file_path: + return self.output_file_path + else: + return None + + def _record_pg_config(self) -> None: + # Records the PG config info to the trace as node: + # ## process_group:init ## + if ( + self.is_registered + and torch.distributed.is_available() + and torch.distributed.is_initialized() + ): + pg_config_info = torch.distributed.distributed_c10d._world.pg_config_info + torch.autograd._record_function_with_args_enter( + "## process_group:init ##", + json.dumps(pg_config_info, cls=_NumpyEncoder), + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/python_tracer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/python_tracer.py new file mode 100644 index 0000000000000000000000000000000000000000..315a370a61bbad98f05f133812e1e8630d9ad8c5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/profiler/python_tracer.py @@ -0,0 +1,20 @@ +import os +import site +import sys + +import torch + + +def _prefix_regex() -> list[str]: + raw_paths = ( + site.getsitepackages() + + sys.path + + [site.getuserbase()] + + [site.getusersitepackages()] + + [os.path.dirname(os.path.dirname(torch.__file__))] + ) + + path_prefixes = sorted({os.path.abspath(i) for i in raw_paths}, reverse=True) + if not all(isinstance(i, str) for i in path_prefixes): + raise AssertionError("all path_prefixes must be strings") + return [i + os.sep for i in path_prefixes] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/py.typed b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..eb25784041d14e40ae392b6ead2d79acba1894ed --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/__init__.py @@ -0,0 +1,90 @@ +# mypy: allow-untyped-defs +from .fake_quantize import * # noqa: F403 +from .fuse_modules import fuse_modules +from .fuser_method_mappings import * # noqa: F403 +from .observer import * # noqa: F403 +from .qconfig import * # noqa: F403 +from .quant_type import * # noqa: F403 +from .quantization_mappings import * # noqa: F403 +from .quantize import * # noqa: F403 +from .quantize_jit import * # noqa: F403 +from .stubs import * # noqa: F403 + + +def default_eval_fn(model, calib_data): + r""" + Default evaluation function takes a torch.utils.data.Dataset or a list of + input Tensors and run the model on the dataset + """ + for data, _target in calib_data: + model(data) + + +__all__ = [ + "QuantWrapper", + "QuantStub", + "DeQuantStub", + # Top level API for eager mode quantization + "quantize", + "quantize_dynamic", + "quantize_qat", + "prepare", + "convert", + "prepare_qat", + # Top level API for graph mode quantization on TorchScript + "quantize_jit", + "quantize_dynamic_jit", + # pyrefly: ignore [bad-dunder-all] + "_prepare_ondevice_dynamic_jit", + # pyrefly: ignore [bad-dunder-all] + "_convert_ondevice_dynamic_jit", + # pyrefly: ignore [bad-dunder-all] + "_quantize_ondevice_dynamic_jit", + # Top level API for graph mode quantization on GraphModule(torch.fx) + # 'fuse_fx', 'quantize_fx', # TODO: add quantize_dynamic_fx + # 'prepare_fx', 'prepare_dynamic_fx', 'convert_fx', + "QuantType", # quantization type + # custom module APIs + "get_default_static_quant_module_mappings", + "get_static_quant_module_class", + "get_default_dynamic_quant_module_mappings", + "get_default_qat_module_mappings", + "get_default_qconfig_propagation_list", + "get_default_compare_output_module_list", + "get_quantized_operator", + "get_fuser_method", + # Sub functions for `prepare` and `swap_module` + "propagate_qconfig_", + "add_quant_dequant", + "swap_module", + "default_eval_fn", + # Observers + "ObserverBase", + # pyrefly: ignore [bad-dunder-all] + "WeightObserver", + "HistogramObserver", + "observer", + "default_observer", + "default_weight_observer", + "default_placeholder_observer", + "default_per_channel_weight_observer", + # FakeQuantize (for qat) + "default_fake_quant", + "default_weight_fake_quant", + "default_fixed_qparams_range_neg1to1_fake_quant", + "default_fixed_qparams_range_0to1_fake_quant", + "default_per_channel_weight_fake_quant", + "default_histogram_fake_quant", + # QConfig + "QConfig", + "default_qconfig", + "default_dynamic_qconfig", + "float16_dynamic_qconfig", + "float_qparams_weight_only_qconfig", + # QAT utilities + "default_qat_qconfig", + "prepare_qat", + "quantize_qat", + # module transformations + "fuse_modules", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/_numeric_suite.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/_numeric_suite.py new file mode 100644 index 0000000000000000000000000000000000000000..49ccc8e69523f7dbee2335b788a2cb3a7db618a2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/_numeric_suite.py @@ -0,0 +1,28 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/ns/_numeric_suite.py`, while adding an import statement +here. +""" + +from torch.ao.ns._numeric_suite import ( + _convert_tuple_to_list, + _dequantize_tensor_list, + _find_match, + _get_logger_dict_helper, + _is_identical_module_type, + compare_model_outputs, + compare_model_stub, + compare_weights, + get_logger_dict, + get_matching_activations, + Logger, + NON_LEAF_MODULE_TO_ADD_OBSERVER_ALLOW_LIST, + OutputLogger, + prepare_model_outputs, + prepare_model_with_stubs, + Shadow, + ShadowLogger, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/_numeric_suite_fx.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/_numeric_suite_fx.py new file mode 100644 index 0000000000000000000000000000000000000000..55cd7085740d0ce8de79491acbfc4888ebba21f8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/_numeric_suite_fx.py @@ -0,0 +1,26 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/ns/_numeric_suite_fx.py`, while adding an import statement +here. +""" + +from torch.ao.ns._numeric_suite_fx import ( + _add_loggers_impl, + _add_loggers_one_model, + _add_shadow_loggers_impl, + _extract_logger_info_one_model, + _extract_weights_impl, + _extract_weights_one_model, + add_loggers, + add_shadow_loggers, + extend_logger_results_with_comparison, + extract_logger_info, + extract_shadow_logger_info, + extract_weights, + NSTracer, + OutputLogger, + RNNReturnType, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/_quantized_conversions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/_quantized_conversions.py new file mode 100644 index 0000000000000000000000000000000000000000..2045d4dd741fd7eeac75737f09a9659f23d94dfd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/_quantized_conversions.py @@ -0,0 +1,146 @@ +# mypy: allow-untyped-defs +import torch + + +# Pack pairs of int4 values into int8, in row major order; first int4 +# value goes into lower order bits, and second int4 value into higher +# order bits of resulting int8 value. +def pack_int4_to_int8(weight): + if weight.dim() != 2: + raise AssertionError(f"weight must be 2D, got {weight.dim()}D") + if weight.shape[1] % 2 != 0: + raise AssertionError(f"weight.shape[1] must be even, got {weight.shape[1]}") + if weight.dtype != torch.int8: + raise AssertionError(f"weight.dtype must be int8, got {weight.dtype}") + return ((weight[:, 1::2] & 0xF) << 4) | (weight[:, 0::2] & 0xF) + + +# Unpack quandruples of bits in int8 values into int4 values, in row +# major order; lower 4 bits go into first int4 value goes, and upper 4 +# bits go into second int4 value. +def unpack_int8_to_int4(weight): + if weight.dim() != 2: + raise AssertionError(f"weight must be 2D, got {weight.dim()}D") + if weight.dtype != torch.int8: + raise AssertionError(f"weight.dtype must be int8, got {weight.dtype}") + return torch.stack((weight & 0xF, (weight >> 4) & 0xF), dim=2).view( + weight.shape[0], 2 * weight.shape[1] + ) + + +# Transpose the weight matrix, and then reorder its elements according +# to underlying requirements of CUTLASS library, so that it could be +# used for CUTLASS-based mixed datatypes linear operation. +def quantized_weight_reorder_for_mixed_dtypes_linear_cutlass( + weight, dtypeq, transpose=False +): + if weight.dim() != 2: + raise AssertionError(f"weight must be 2D, got {weight.dim()}D") + if weight.dtype != torch.int8: + raise AssertionError(f"weight.dtype must be int8, got {weight.dtype}") + if dtypeq != torch.int8 and dtypeq != torch.quint4x2: + raise AssertionError(f"dtypeq must be int8 or quint4x2, got {dtypeq}") + if weight.device.type != "cuda": + raise AssertionError(f"weight must be on CUDA, got {weight.device.type}") + + device = weight.device + + # subbyte_transpose + if not transpose: + if dtypeq == torch.int8: + outp = weight.T + elif dtypeq == torch.quint4x2: + outp = pack_int4_to_int8(unpack_int8_to_int4(weight.view(torch.int8)).T) + else: + outp = weight + + ncols, nrows = outp.shape # type: ignore[possibly-undefined] + divisor = 32 if dtypeq == torch.quint4x2 else 64 + if nrows % divisor != 0: + raise AssertionError(f"nrows must be divisible by {divisor}, got {nrows}") + if ncols % 64 != 0: + raise AssertionError(f"ncols must be divisible by 64, got {ncols}") + + # permute_B_rows_for_mixed_gemm + # (permute cols actually, as transpose is applied first here) + if dtypeq == torch.quint4x2: + cols_permuted = ( + torch.tensor( + [0, 4, 8, 12, 1, 5, 9, 13, 2, 6, 10, 14, 3, 7, 11, 15], + device=device, + ) + + (torch.arange(0, nrows // 16, device=device).reshape(-1, 1) * 16).expand( + nrows // 16, 16 + ) + ).view(-1) + else: + cols_permuted = ( + torch.tensor( + [0, 1, 4, 5, 8, 9, 12, 13, 2, 3, 6, 7, 10, 11, 14, 15], + device=device, + ) + + (torch.arange(0, nrows // 16, device=device).reshape(-1, 1) * 16).expand( + nrows // 16, 16 + ) + ).view(-1) + # pyrefly: ignore [unbound-name] + outp = outp.index_copy(1, cols_permuted, outp) + + # interleave_column_major_tensor + magic0 = 4 if dtypeq == torch.quint4x2 else 2 + magic1 = 32 // magic0 + + tmp0 = ( + (torch.arange(0, ncols // magic0, device=device) * (nrows // 4 * magic0)) + .view(-1, 1) + .repeat(1, nrows // 4 * magic0) + .view(-1) + ) + tmp1 = ( + (torch.arange(0, nrows // 4 // magic1, device=device) * (magic0 * magic1)) + .view(-1, 1) + .repeat(1, magic1) + .view(-1) + .repeat(ncols) + ) + tmp2 = ( + (torch.arange(0, magic0, device=device) * magic1) + .view(-1, 1) + .repeat(1, nrows // 4) + .view(-1) + .repeat(ncols // magic0) + ) + tmp3 = torch.arange(0, magic1, device=device).repeat(nrows // 4 * ncols // magic1) + + outp_offsets = tmp0 + tmp1 + tmp2 + tmp3 + + tmp = outp.view(-1).view(torch.int32) + outp = torch.zeros_like(tmp) + outp.scatter_(0, outp_offsets, tmp) + outp = outp.view(weight.dtype) + + # add_bias_and_interleave_quantized_tensor_inplace + tmp = outp.view(-1) + + outp = torch.empty_like(tmp) + if dtypeq == torch.int8: + tmp = (tmp.to(torch.int) + 128).to(tmp.dtype) + outp[0::4] = tmp[0::4] + outp[1::4] = tmp[2::4] + outp[2::4] = tmp[1::4] + outp[3::4] = tmp[3::4] + elif dtypeq == torch.quint4x2: + tmp0 = ((tmp & 0xF) + 8) & 0xF + tmp0 = (tmp0[1::2] << 4) | tmp0[0::2] + tmp1 = (((tmp >> 4) & 0xF) + 8) & 0xF + tmp1 = (tmp1[1::2] << 4) | tmp1[0::2] + outp[0::4] = tmp0[0::2] + outp[1::4] = tmp0[1::2] + outp[2::4] = tmp1[0::2] + outp[3::4] = tmp1[1::2] + + if dtypeq == torch.quint4x2: + nrows *= 2 + ncols //= 2 + + return outp.view(nrows, ncols).view(torch.uint8) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fake_quantize.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fake_quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..69a5d730bfb68e89e24beb04ad13fd3fa5881ae9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fake_quantize.py @@ -0,0 +1,32 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/fake_quantize.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.fake_quantize import ( + _is_fake_quant_script_module, + _is_per_channel, + _is_per_tensor, + _is_symmetric_quant, + default_fake_quant, + default_fixed_qparams_range_0to1_fake_quant, + default_fixed_qparams_range_neg1to1_fake_quant, + default_fused_act_fake_quant, + default_fused_per_channel_wt_fake_quant, + default_fused_wt_fake_quant, + default_histogram_fake_quant, + default_per_channel_weight_fake_quant, + default_weight_fake_quant, + disable_fake_quant, + disable_observer, + enable_fake_quant, + enable_observer, + FakeQuantize, + FakeQuantizeBase, + FixedQParamsFakeQuantize, + FusedMovingAvgObsFakeQuantize, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fuse_modules.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fuse_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..bce403549d68584ec22089c22b14f17010d6252d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fuse_modules.py @@ -0,0 +1,22 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/fuse_modules.py`, while adding an import statement +here. +""" + +# TODO: These functions are not used outside the `fuse_modules.py` +# Keeping here for now, need to remove them later. +from torch.ao.quantization.fuse_modules import ( + _fuse_modules, + _get_module, + _set_module, + fuse_known_modules, + fuse_modules, + get_fuser_method, +) + +# for backward compatibility +from torch.ao.quantization.fuser_method_mappings import fuse_conv_bn, fuse_conv_bn_relu diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fuser_method_mappings.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fuser_method_mappings.py new file mode 100644 index 0000000000000000000000000000000000000000..5a68fbf02015ff162ebbd4e26bf85a94328322c8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fuser_method_mappings.py @@ -0,0 +1,16 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/fuser_method_mappings.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.fuser_method_mappings import ( + _DEFAULT_OP_LIST_TO_FUSER_METHOD, + fuse_conv_bn, + fuse_conv_bn_relu, + fuse_linear_bn, + get_fuser_method, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c01cbd457374c27e40b07daca5ae1644a701767d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/__init__.py @@ -0,0 +1,15 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.convert import convert +from torch.ao.quantization.fx.fuse import fuse + +# omitting files that's unlikely to be used right now, for example +# the newly added lower_to_fbgemm etc. +from torch.ao.quantization.fx.prepare import prepare diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/_equalize.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/_equalize.py new file mode 100644 index 0000000000000000000000000000000000000000..d6b8611d4a769a9c1e93682180becc5117020d55 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/_equalize.py @@ -0,0 +1,39 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx._equalize import ( + _convert_equalization_ref, + _InputEqualizationObserver, + _WeightEqualizationObserver, + calculate_equalization_scale, + clear_weight_quant_obs_node, + convert_eq_obs, + CUSTOM_MODULE_SUPP_LIST, + custom_module_supports_equalization, + default_equalization_qconfig, + EqualizationQConfig, + fused_module_supports_equalization, + get_equalization_qconfig_dict, + get_layer_sqnr_dict, + get_op_node_and_weight_eq_obs, + input_equalization_observer, + is_equalization_observer, + maybe_get_next_equalization_scale, + maybe_get_next_input_eq_obs, + maybe_get_weight_eq_obs_node, + nn_module_supports_equalization, + node_supports_equalization, + remove_node, + reshape_scale, + scale_input_observer, + scale_weight_functional, + scale_weight_node, + update_obs_for_equalization, + weight_equalization_observer, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/convert.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/convert.py new file mode 100644 index 0000000000000000000000000000000000000000..30a661da41e5e2bb417a0e0aa6c7088a1b8ea7e4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/convert.py @@ -0,0 +1,10 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.convert import convert diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/fuse.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/fuse.py new file mode 100644 index 0000000000000000000000000000000000000000..22ad750e9f8784376cecee4f5d10cfcd1488a7ac --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/fuse.py @@ -0,0 +1,10 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.fuse import fuse diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/fusion_patterns.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/fusion_patterns.py new file mode 100644 index 0000000000000000000000000000000000000000..982d919655f36320c87e066fa04e8ab10e70a719 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/fusion_patterns.py @@ -0,0 +1,10 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.fuse_handler import DefaultFuseHandler, FuseHandler diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/graph_module.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/graph_module.py new file mode 100644 index 0000000000000000000000000000000000000000..74b63903d7400c037ca15ac7b9cf200d70d07ab9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/graph_module.py @@ -0,0 +1,18 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.graph_module import ( + _is_observed_module, + _is_observed_standalone_module, + FusedGraphModule, + GraphModule, + ObservedGraphModule, + ObservedStandaloneGraphModule, + QuantizedGraphModule, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/match_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/match_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8585a21ad445dd20338d24267d8a0f05f96d0f92 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/match_utils.py @@ -0,0 +1,15 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.match_utils import ( + _find_matches, + _is_match, + _MatchResult, + MatchAllNode, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/pattern_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/pattern_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..fa601d1eb619c14a37f95177b9850942ab361974 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/pattern_utils.py @@ -0,0 +1,36 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.pattern_utils import ( + _register_fusion_pattern, + _register_quant_pattern, + get_default_fusion_patterns, + get_default_output_activation_post_process_map, + get_default_quant_patterns, + QuantizeHandler, +) + + +# QuantizeHandler.__module__ = _NAMESPACE +_register_fusion_pattern.__module__ = "torch.ao.quantization.fx.pattern_utils" +get_default_fusion_patterns.__module__ = "torch.ao.quantization.fx.pattern_utils" +_register_quant_pattern.__module__ = "torch.ao.quantization.fx.pattern_utils" +get_default_quant_patterns.__module__ = "torch.ao.quantization.fx.pattern_utils" +get_default_output_activation_post_process_map.__module__ = ( + "torch.ao.quantization.fx.pattern_utils" +) + +# __all__ = [ +# "QuantizeHandler", +# "_register_fusion_pattern", +# "get_default_fusion_patterns", +# "_register_quant_pattern", +# "get_default_quant_patterns", +# "get_default_output_activation_post_process_map", +# ] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/prepare.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/prepare.py new file mode 100644 index 0000000000000000000000000000000000000000..a6007ef242af5d33566065a0b9d570399deccf94 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/prepare.py @@ -0,0 +1,10 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.prepare import prepare diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/quantization_patterns.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/quantization_patterns.py new file mode 100644 index 0000000000000000000000000000000000000000..89f8d4406e9126525d6c1518c6743a5c84c7b760 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/quantization_patterns.py @@ -0,0 +1,49 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.quantize_handler import ( + BatchNormQuantizeHandler, + BinaryOpQuantizeHandler, + CatQuantizeHandler, + ConvReluQuantizeHandler, + CopyNodeQuantizeHandler, + CustomModuleQuantizeHandler, + DefaultNodeQuantizeHandler, + EmbeddingQuantizeHandler, + FixedQParamsOpQuantizeHandler, + GeneralTensorShapeOpQuantizeHandler, + LinearReLUQuantizeHandler, + QuantizeHandler, + RNNDynamicQuantizeHandler, + StandaloneModuleQuantizeHandler, +) + + +QuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +BinaryOpQuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +CatQuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +ConvReluQuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +LinearReLUQuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +BatchNormQuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +EmbeddingQuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +RNNDynamicQuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +DefaultNodeQuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +FixedQParamsOpQuantizeHandler.__module__ = ( + "torch.ao.quantization.fx.quantization_patterns" +) +CopyNodeQuantizeHandler.__module__ = "torch.ao.quantization.fx.quantization_patterns" +CustomModuleQuantizeHandler.__module__ = ( + "torch.ao.quantization.fx.quantization_patterns" +) +GeneralTensorShapeOpQuantizeHandler.__module__ = ( + "torch.ao.quantization.fx.quantization_patterns" +) +StandaloneModuleQuantizeHandler.__module__ = ( + "torch.ao.quantization.fx.quantization_patterns" +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/quantization_types.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/quantization_types.py new file mode 100644 index 0000000000000000000000000000000000000000..0820ea057078ea89da763b1c5864089b8682a9f3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/quantization_types.py @@ -0,0 +1,10 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.utils import Pattern, QuantizerCls diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e45c82b8fb6f2379a5805442666f5551c2680683 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/fx/utils.py @@ -0,0 +1,21 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +appropriate files under `torch/ao/quantization/fx/`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.utils import ( + all_node_args_have_no_tensors, + assert_and_get_unique_device, + create_getattr_from_value, + get_custom_module_class_keys, + get_linear_prepack_op_for_dtype, + get_new_attr_name_with_prefix, + get_non_observable_arg_indexes_and_types, + get_qconv_prepack_op, + graph_module_from_producer_nodes, + maybe_get_next_module, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/observer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/observer.py new file mode 100644 index 0000000000000000000000000000000000000000..2163e2717b0697d34fe23e05dbb69c3a555da4b3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/observer.py @@ -0,0 +1,37 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/observer.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.observer import ( + _is_activation_post_process, + _is_per_channel_script_obs_instance, + _ObserverBase, + _PartialWrapper, + _with_args, + _with_callable_args, + ABC, + default_debug_observer, + default_dynamic_quant_observer, + default_float_qparams_observer, + default_histogram_observer, + default_observer, + default_per_channel_weight_observer, + default_placeholder_observer, + default_weight_observer, + get_observer_state_dict, + HistogramObserver, + load_observer_state_dict, + MinMaxObserver, + MovingAverageMinMaxObserver, + MovingAveragePerChannelMinMaxObserver, + NoopObserver, + ObserverBase, + PerChannelMinMaxObserver, + PlaceholderObserver, + RecordingObserver, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/qconfig.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/qconfig.py new file mode 100644 index 0000000000000000000000000000000000000000..75398d3343f93a813346e04d3709885eaaeae5cb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/qconfig.py @@ -0,0 +1,31 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/qconfig.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.qconfig import ( + _add_module_to_qconfig_obs_ctr, + _assert_valid_qconfig, + default_activation_only_qconfig, + default_debug_qconfig, + default_dynamic_qconfig, + default_per_channel_qconfig, + default_qat_qconfig, + default_qat_qconfig_v2, + default_qconfig, + default_weight_only_qconfig, + float16_dynamic_qconfig, + float16_static_qconfig, + float_qparams_weight_only_qconfig, + get_default_qat_qconfig, + get_default_qconfig, + per_channel_dynamic_qconfig, + QConfig, + qconfig_equals, + QConfigAny, + QConfigDynamic, # pyrefly: ignore # deprecated +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quant_type.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quant_type.py new file mode 100644 index 0000000000000000000000000000000000000000..8555f03792661f39c85c8facf3f911786cc25d0f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quant_type.py @@ -0,0 +1,10 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/quant_type.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.quant_type import _get_quant_type_to_str, QuantType diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantization_mappings.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantization_mappings.py new file mode 100644 index 0000000000000000000000000000000000000000..faa24d391d31ad65cb54d580a7dc6e8f1ff36f83 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantization_mappings.py @@ -0,0 +1,30 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/quantization_mappings.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.quantization_mappings import ( + _get_special_act_post_process, + _has_special_act_post_process, + _INCLUDE_QCONFIG_PROPAGATE_LIST, + DEFAULT_DYNAMIC_QUANT_MODULE_MAPPINGS, + DEFAULT_FLOAT_TO_QUANTIZED_OPERATOR_MAPPINGS, + DEFAULT_MODULE_TO_ACT_POST_PROCESS, + DEFAULT_QAT_MODULE_MAPPINGS, + DEFAULT_REFERENCE_STATIC_QUANT_MODULE_MAPPINGS, + DEFAULT_STATIC_QUANT_MODULE_MAPPINGS, + get_default_compare_output_module_list, + get_default_dynamic_quant_module_mappings, + get_default_float_to_quantized_operator_mappings, + get_default_qat_module_mappings, + get_default_qconfig_propagation_list, + get_default_static_quant_module_mappings, + get_dynamic_quant_module_class, + get_quantized_operator, + get_static_quant_module_class, + no_observer_set, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantize.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantize.py new file mode 100644 index 0000000000000000000000000000000000000000..600d3a46fed0346e3ae8909872cd5bf3c733860c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantize.py @@ -0,0 +1,30 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/quantize.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.quantize import ( + _add_observer_, + _convert, + _get_observer_dict, + _get_unique_devices_, + _is_activation_post_process, + _observer_forward_hook, + _propagate_qconfig_helper, + _register_activation_post_process_hook, + _remove_activation_post_process, + _remove_qconfig, + add_quant_dequant, + convert, + prepare, + prepare_qat, + propagate_qconfig_, + quantize, + quantize_dynamic, + quantize_qat, + swap_module, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantize_fx.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantize_fx.py new file mode 100644 index 0000000000000000000000000000000000000000..649142c7a7eee9885d96b37f70e582f3ea9a9f8d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantize_fx.py @@ -0,0 +1,26 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/quantize_fx.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.fx.graph_module import ObservedGraphModule +from torch.ao.quantization.quantize_fx import ( + _check_is_graph_module, + _convert_fx, + _convert_standalone_module_fx, + _fuse_fx, + _prepare_fx, + _prepare_standalone_module_fx, + _swap_ff_with_fxff, + convert_fx, + fuse_fx, + prepare_fx, + prepare_qat_fx, + QuantizationTracer, + Scope, + ScopeContextManager, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantize_jit.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantize_jit.py new file mode 100644 index 0000000000000000000000000000000000000000..aa627dc7bb51ef7ea1fde7e2e5da283c9f6c8900 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/quantize_jit.py @@ -0,0 +1,26 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/quantize_jit.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.quantize_jit import ( + _check_forward_method, + _check_is_script_module, + _convert_jit, + _prepare_jit, + _prepare_ondevice_dynamic_jit, + _quantize_jit, + convert_dynamic_jit, + convert_jit, + fuse_conv_bn_jit, + prepare_dynamic_jit, + prepare_jit, + quantize_dynamic_jit, + quantize_jit, + script_qconfig, + script_qconfig_dict, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/stubs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/stubs.py new file mode 100644 index 0000000000000000000000000000000000000000..d3fd5c63683dc572c35cabc202ee4ddb2b0053c6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/stubs.py @@ -0,0 +1,10 @@ +# flake8: noqa: F401 +r""" +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/stubs.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.stubs import DeQuantStub, QuantStub, QuantWrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7d51d58f38d7462713f84ab62427852c1dd8e52c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quantization/utils.py @@ -0,0 +1,29 @@ +# flake8: noqa: F401 +r""" +Utils shared by different modes of quantization (eager/graph) + +This file is in the process of migration to `torch/ao/quantization`, and +is kept here for compatibility while the migration process is ongoing. +If you are adding a new entry/functionality, please, add it to the +`torch/ao/quantization/utils.py`, while adding an import statement +here. +""" + +from torch.ao.quantization.utils import ( + activation_dtype, + activation_is_int8_quantized, + activation_is_statically_quantized, + calculate_qmin_qmax, + check_min_max_valid, + get_combined_dict, + get_qconfig_dtypes, + get_qparam_dict, + get_quant_type, + get_swapped_custom_module_class, + getattr_from_fqn, + is_per_channel, + is_per_tensor, + weight_dtype, + weight_is_quantized, + weight_is_statically_quantized, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quasirandom.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quasirandom.py new file mode 100644 index 0000000000000000000000000000000000000000..f9e6619cab180da41e0d7ec5968cdfd20da9097c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/quasirandom.py @@ -0,0 +1,216 @@ +# mypy: allow-untyped-defs + +import torch + + +class SobolEngine: + r""" + The :class:`torch.quasirandom.SobolEngine` is an engine for generating + (scrambled) Sobol sequences. Sobol sequences are an example of low + discrepancy quasi-random sequences. + + This implementation of an engine for Sobol sequences is capable of + sampling sequences up to a maximum dimension of 21201. It uses direction + numbers from https://web.maths.unsw.edu.au/~fkuo/sobol/ obtained using the + search criterion D(6) up to the dimension 21201. This is the recommended + choice by the authors. + + References: + - Art B. Owen. Scrambling Sobol and Niederreiter-Xing points. + Journal of Complexity, 14(4):466-489, December 1998. + + - I. M. Sobol. The distribution of points in a cube and the accurate + evaluation of integrals. + Zh. Vychisl. Mat. i Mat. Phys., 7:784-802, 1967. + + Args: + dimension (Int): The dimensionality of the sequence to be drawn + scramble (bool, optional): Setting this to ``True`` will produce + scrambled Sobol sequences. Scrambling is + capable of producing better Sobol + sequences. Default: ``False``. + seed (Int, optional): This is the seed for the scrambling. The seed + of the random number generator is set to this, + if specified. Otherwise, it uses a random seed. + Default: ``None`` + + Examples:: + + >>> # xdoctest: +SKIP("unseeded random state") + >>> soboleng = torch.quasirandom.SobolEngine(dimension=5) + >>> soboleng.draw(3) + tensor([[0.0000, 0.0000, 0.0000, 0.0000, 0.0000], + [0.5000, 0.5000, 0.5000, 0.5000, 0.5000], + [0.7500, 0.2500, 0.2500, 0.2500, 0.7500]]) + """ + + MAXBIT = 30 + MAXDIM = 21201 + + def __init__(self, dimension, scramble=False, seed=None): + if dimension > self.MAXDIM or dimension < 1: + raise ValueError( + "Supported range of dimensionality " + f"for SobolEngine is [1, {self.MAXDIM}]" + ) + + self.seed = seed + self.scramble = scramble + self.dimension = dimension + + cpu = torch.device("cpu") + + self.sobolstate = torch.zeros( + dimension, self.MAXBIT, device=cpu, dtype=torch.long + ) + torch._sobol_engine_initialize_state_(self.sobolstate, self.dimension) + + if not self.scramble: + self.shift = torch.zeros(self.dimension, device=cpu, dtype=torch.long) + else: + self._scramble() + + self.quasi = self.shift.clone(memory_format=torch.contiguous_format) + self._first_point = (self.quasi / 2**self.MAXBIT).reshape(1, -1) + self.num_generated = 0 + + def draw( + self, + n: int = 1, + out: torch.Tensor | None = None, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + r""" + Function to draw a sequence of :attr:`n` points from a Sobol sequence. + Note that the samples are dependent on the previous samples. The size + of the result is :math:`(n, dimension)`. + + Args: + n (Int, optional): The length of sequence of points to draw. + Default: 1 + out (Tensor, optional): The output tensor + dtype (:class:`torch.dtype`, optional): the desired data type of the + returned tensor. + Default: ``None`` + """ + if dtype is None: + dtype = torch.get_default_dtype() + + if self.num_generated == 0: + if n == 1: + result = self._first_point.to(dtype) + else: + result, self.quasi = torch._sobol_engine_draw( + self.quasi, + n - 1, + self.sobolstate, + self.dimension, + self.num_generated, + dtype=dtype, + ) + result = torch.cat((self._first_point.to(dtype), result), dim=-2) + else: + result, self.quasi = torch._sobol_engine_draw( + self.quasi, + n, + self.sobolstate, + self.dimension, + self.num_generated - 1, + dtype=dtype, + ) + + self.num_generated += n + + if out is not None: + out.resize_as_(result).copy_(result) + return out + + return result + + def draw_base2( + self, + m: int, + out: torch.Tensor | None = None, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + r""" + Function to draw a sequence of :attr:`2**m` points from a Sobol sequence. + Note that the samples are dependent on the previous samples. The size + of the result is :math:`(2**m, dimension)`. + + Args: + m (Int): The (base2) exponent of the number of points to draw. + out (Tensor, optional): The output tensor + dtype (:class:`torch.dtype`, optional): the desired data type of the + returned tensor. + Default: ``None`` + """ + n = 2**m + total_n = self.num_generated + n + if not (total_n & (total_n - 1) == 0): + raise ValueError( + "The balance properties of Sobol' points require " + f"n to be a power of 2. {self.num_generated} points have been " + f"previously generated, then: n={self.num_generated}+2**{m}={total_n}. " + "If you still want to do this, please use " + "'SobolEngine.draw()' instead." + ) + return self.draw(n=n, out=out, dtype=dtype) + + def reset(self): + r""" + Function to reset the ``SobolEngine`` to base state. + """ + self.quasi.copy_(self.shift) + self.num_generated = 0 + return self + + def fast_forward(self, n): + r""" + Function to fast-forward the state of the ``SobolEngine`` by + :attr:`n` steps. This is equivalent to drawing :attr:`n` samples + without using the samples. + + Args: + n (Int): The number of steps to fast-forward by. + """ + if self.num_generated == 0: + torch._sobol_engine_ff_( + self.quasi, n - 1, self.sobolstate, self.dimension, self.num_generated + ) + else: + torch._sobol_engine_ff_( + self.quasi, n, self.sobolstate, self.dimension, self.num_generated - 1 + ) + self.num_generated += n + return self + + def _scramble(self): + g: torch.Generator | None = None + if self.seed is not None: + g = torch.Generator() + g.manual_seed(self.seed) + + cpu = torch.device("cpu") + + # Generate shift vector + shift_ints = torch.randint( + 2, (self.dimension, self.MAXBIT), device=cpu, generator=g + ) + self.shift = torch.mv( + shift_ints, torch.pow(2, torch.arange(0, self.MAXBIT, device=cpu)) + ) + + # Generate lower triangular matrices (stacked across dimensions) + ltm_dims = (self.dimension, self.MAXBIT, self.MAXBIT) + ltm = torch.randint(2, ltm_dims, device=cpu, generator=g).tril() + + torch._sobol_engine_scramble_(self.sobolstate, ltm, self.dimension) + + def __repr__(self): + fmt_string = [f"dimension={self.dimension}"] + if self.scramble: + fmt_string += ["scramble=True"] + if self.seed is not None: + fmt_string += [f"seed={self.seed}"] + return self.__class__.__name__ + "(" + ", ".join(fmt_string) + ")" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/random.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/random.py new file mode 100644 index 0000000000000000000000000000000000000000..dd137d3a2d4fb688f063b2b2eac7fb4afdf5bc96 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/random.py @@ -0,0 +1,270 @@ +# mypy: allow-untyped-defs +import contextlib +import warnings +from collections.abc import Generator +from typing import TYPE_CHECKING + +import torch + + +__all__ = [ + "set_rng_state", + "get_rng_state", + "manual_seed", + "seed", + "initial_seed", + "fork_rng", + "thread_safe_generator", +] + + +if TYPE_CHECKING: + from torch.utils.data._utils.worker import WorkerInfo + +from torch._C import default_generator + + +def set_rng_state(new_state: torch.Tensor) -> None: + r"""Sets the random number generator state. + + .. note:: This function only works for CPU. For CUDA, please use + :func:`torch.manual_seed`, which works for both CPU and CUDA. + + Args: + new_state (torch.ByteTensor): The desired state + """ + default_generator.set_state(new_state) + + +def get_rng_state() -> torch.Tensor: + r"""Returns the random number generator state as a `torch.ByteTensor`. + + .. note:: The returned state is for the default generator on CPU only. + + See also: :func:`torch.random.fork_rng`. + """ + return default_generator.get_state() + + +def manual_seed(seed) -> torch._C.Generator: + r"""Sets the seed for generating random numbers on all devices. Returns a + `torch.Generator` object. + + Args: + seed (int): The desired seed. Value must be within the inclusive range + `[-0x8000_0000_0000_0000, 0xffff_ffff_ffff_ffff]`. Otherwise, a RuntimeError + is raised. Negative inputs are remapped to positive values with the formula + `0xffff_ffff_ffff_ffff + seed`. + """ + return _manual_seed_impl(seed) + + +def _manual_seed_impl(seed) -> torch._C.Generator: + seed = int(seed) + import torch.cuda + + if not torch.cuda._is_in_bad_fork(): + torch.cuda.manual_seed_all(seed) + + import torch.mps + + if not torch.mps._is_in_bad_fork(): + torch.mps.manual_seed(seed) + + import torch.xpu + + if not torch.xpu._is_in_bad_fork(): + torch.xpu.manual_seed_all(seed) + + import torch.mtia + + if not torch.mtia._is_in_bad_fork(): + torch.mtia.manual_seed_all(seed) + + _seed_custom_device(seed) + + return default_generator.manual_seed(seed) + + +def seed() -> int: + r"""Sets the seed for generating random numbers to a non-deterministic + random number on all devices. Returns a 64 bit number used to seed the RNG. + """ + seed = default_generator.seed() + import torch.cuda + + if not torch.cuda._is_in_bad_fork(): + torch.cuda.manual_seed_all(seed) + + import torch.mps + + if not torch.mps._is_in_bad_fork(): + torch.mps.manual_seed(seed) + + import torch.xpu + + if not torch.xpu._is_in_bad_fork(): + torch.xpu.manual_seed_all(seed) + + import torch.mtia + + if not torch.mtia._is_in_bad_fork(): + torch.mtia.manual_seed_all(seed) + + _seed_custom_device(seed) + + return seed + + +def _seed_custom_device(seed) -> None: + r"""Sets the seed to generate random numbers for custom device. + + Args: + seed (int): The desired seed. + + See [Note: support the custom device with privateuse1] + """ + seed = int(seed) + custom_backend_name = torch._C._get_privateuse1_backend_name() + if hasattr(torch, custom_backend_name): + custom_device_mod = getattr(torch, custom_backend_name) + _bad_fork_name = "_is_in_bad_fork" + _seed_all_name = "manual_seed_all" + if hasattr(custom_device_mod, _bad_fork_name) and hasattr( + custom_device_mod, _seed_all_name + ): + if not getattr(custom_device_mod, _bad_fork_name)(): + getattr(custom_device_mod, _seed_all_name)(seed) + else: + message = f"Set seed for `{custom_backend_name}` device does not take effect, please add API's " + message += f"`{_bad_fork_name}` and `{_seed_all_name}` to `{custom_backend_name}` device module." + warnings.warn(message, UserWarning, stacklevel=3) + + +def initial_seed() -> int: + r"""Returns the initial seed for generating random numbers as a + Python `long`. + + .. note:: The returned seed is for the default generator on CPU only. + """ + return default_generator.initial_seed() + + +_fork_rng_warned_already = False + + +@contextlib.contextmanager +def fork_rng( + devices=None, + enabled=True, + _caller="fork_rng", + _devices_kw="devices", + device_type="cuda", +) -> Generator: + """ + Forks the RNG, so that when you return, the RNG is reset + to the state that it was previously in. + + Args: + devices (iterable of Device IDs): devices for which to fork + the RNG. CPU RNG state is always forked. By default, :meth:`fork_rng` operates + on all devices, but will emit a warning if your machine has a lot + of devices, since this function will run very slowly in that case. + If you explicitly specify devices, this warning will be suppressed + enabled (bool): if ``False``, the RNG is not forked. This is a convenience + argument for easily disabling the context manager without having + to delete it and unindent your Python code under it. + device_type (str): device type str, default is `cuda`. As for supported device, + see details in :ref:`accelerator` + """ + + if device_type == "meta": + yield + return + + device_type = torch.device(device_type).type + device_mod = getattr(torch, device_type, None) + if device_mod is None: + raise RuntimeError( + f"torch has no module of `{device_type}`, you should register " + + "a module by `torch._register_device_module`." + ) + global _fork_rng_warned_already + + # Internal arguments: + # _caller: the function which called fork_rng, which the user used + # _devices_kw: the devices keyword of _caller + + if not enabled: + yield + return + + if devices is None: + num_devices = device_mod.device_count() + if num_devices > 1 and not _fork_rng_warned_already: + message = ( + f"{device_type.upper()} reports that you have {num_devices} available devices, and " + f"you have used {_caller} without explicitly specifying which devices are being used. " + f"For safety, we initialize *every* {device_type.upper()} device by default, which can " + f"be quite slow if you have a lot of {device_type.upper()}s. If you know that you are only" + f" making use of a few {device_type.upper()} devices, set the environment variable " + f"{device_type.upper()}_VISIBLE_DEVICES or the '{_devices_kw}' keyword argument of {_caller} " + "with the set of devices you are actually using. For example, if you are using CPU only, " + "set device.upper()_VISIBLE_DEVICES= or devices=[]; if you are using device 0 only, " + f"set {device_type.upper()}_VISIBLE_DEVICES=0 or devices=[0]. To initialize all devices " + f"and suppress this warning, set the '{_devices_kw}' keyword argument to " + f"`range(torch.{device_type}.device_count())`." + ) + warnings.warn(message, stacklevel=2) + _fork_rng_warned_already = True + devices = list(range(num_devices)) + else: + # Protect against user passing us a generator; we need to traverse this + # multiple times but a generator will be exhausted upon first traversal + devices = list(devices) + + cpu_rng_state = torch.get_rng_state() + device_rng_states = [device_mod.get_rng_state(device) for device in devices] + + try: + yield + finally: + torch.set_rng_state(cpu_rng_state) + for device, device_rng_state in zip(devices, device_rng_states): + device_mod.set_rng_state(device_rng_state, device) + + +def thread_safe_generator() -> torch.Generator | None: + """Returns a thread-safe random number generator for use in DataLoader workers. + This function provides a convenient way for transforms and user code to use + thread-safe random number generation without manually checking worker context. + When called in a DataLoader thread worker, returns the worker's thread-local + :class:`torch.Generator`. When called in the main process or process workers, + returns ``None`` (which causes PyTorch functions to use the default global RNG). + Returns: + Optional[torch.Generator]: Thread-local generator in thread workers, None otherwise. + Example:: + >>> from torch.random import thread_safe_generator + >>> generator = thread_safe_generator() + >>> torch.randint(0, 10, (5,), generator=generator) + Example with transforms:: + >>> from torch.random import thread_safe_generator + >>> class MyRandomTransform: + ... def __call__(self, img): + ... generator = thread_safe_generator() + ... offset = torch.randint(0, 10, (2,), generator=generator) + ... return img[..., offset[0]:, offset[1]:] + """ + # Lazy import to avoid circular dependency during torch module initialization + # torch.__init__ loads torch.random early, but torch.utils.data triggers + # torch.distributed which needs torch to be fully initialized + from torch.utils.data import get_worker_info + + worker_info: WorkerInfo | None = get_worker_info() + if ( + worker_info is not None + and worker_info.worker_method == "thread" + and worker_info.rng is not None + ): + return worker_info.rng.torch_generator + return None diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/return_types.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/return_types.py new file mode 100644 index 0000000000000000000000000000000000000000..4c74fe01537e7762243ae0c79657524b125d10b4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/return_types.py @@ -0,0 +1,52 @@ +import inspect + +import torch +from torch.utils._pytree import register_pytree_node, SequenceKey + + +__all__ = ["pytree_register_structseq", "all_return_types"] + +all_return_types = [] + +# error: Module has no attribute "_return_types" +return_types = torch._C._return_types # type: ignore[attr-defined] + + +def pytree_register_structseq(cls): + def structseq_flatten(structseq): + return list(structseq), None + + def structseq_flatten_with_keys(structseq): + values, context = structseq_flatten(structseq) + return [(SequenceKey(i), v) for i, v in enumerate(values)], context + + def structseq_unflatten(values, context): + return cls(values) + + register_pytree_node( + cls, + structseq_flatten, + structseq_unflatten, + flatten_with_keys_fn=structseq_flatten_with_keys, + ) + + +for name in dir(return_types): + if name.startswith("__"): + continue + + _attr = getattr(return_types, name) + globals()[name] = _attr + + if not name.startswith("_"): + # pyrefly: ignore [unresolvable-dunder-all] + __all__.append(name) + all_return_types.append(_attr) + + # Today everything in torch.return_types is a structseq, aka a "namedtuple"-like + # thing defined by the Python C-API. We're going to need to modify this when that + # is no longer the case. + # NB: I don't know how to check that something is a "structseq" so we do a fuzzy + # check for tuple + if inspect.isclass(_attr) and issubclass(_attr, tuple): + pytree_register_structseq(_attr) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/return_types.pyi b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/return_types.pyi new file mode 100644 index 0000000000000000000000000000000000000000..e8dce0869e225ea2a50200240469d220ed296e34 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/return_types.pyi @@ -0,0 +1,605 @@ +# @generated by tools/pyi/gen_pyi.py from torch/_C/return_types.pyi.in +# mypy: allow-untyped-defs + +from typing import Final, NoReturn +from typing_extensions import Self + +from torch import SymInt, Tensor +from torch.types import ( # noqa: F401 + _bool, + _device, + _dtype, + _float, + _int, + _layout, + _qscheme, + _size, + Number, +) + +__all__ = [ + "pytree_register_structseq", + "all_return_types", + "_fake_quantize_per_tensor_affine_cachemask_tensor_qparams", + "_fused_moving_avg_obs_fq_helper", + "_linalg_det", + "_linalg_eigh", + "_linalg_slogdet", + "_linalg_solve_ex", + "_linalg_svd", + "_lu_with_info", + "_scaled_dot_product_cudnn_attention", + "_scaled_dot_product_efficient_attention", + "_scaled_dot_product_flash_attention", + "_scaled_dot_product_flash_attention_for_cpu", + "_unpack_dual", + "aminmax", + "cummax", + "cummin", + "frexp", + "geqrf", + "histogram", + "histogramdd", + "kthvalue", + "lu_unpack", + "max", + "median", + "min", + "mode", + "nanmedian", + "qr", + "slogdet", + "sort", + "svd", + "topk", + "triangular_solve", +] + +def pytree_register_structseq(cls: type) -> None: ... + +class _fake_quantize_per_tensor_affine_cachemask_tensor_qparams(tuple[Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def mask(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _fused_moving_avg_obs_fq_helper(tuple[Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def mask(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_det(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def result(self) -> Tensor: ... + @property + def LU(self) -> Tensor: ... + @property + def pivots(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_eigh(tuple[Tensor, Tensor]): # fmt: skip + @property + def eigenvalues(self) -> Tensor: ... + @property + def eigenvectors(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_slogdet(tuple[Tensor, Tensor, Tensor, Tensor]): # fmt: skip + @property + def sign(self) -> Tensor: ... + @property + def logabsdet(self) -> Tensor: ... + @property + def LU(self) -> Tensor: ... + @property + def pivots(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 4 + n_sequence_fields: Final[_int] = 4 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_solve_ex(tuple[Tensor, Tensor, Tensor, Tensor]): # fmt: skip + @property + def result(self) -> Tensor: ... + @property + def LU(self) -> Tensor: ... + @property + def pivots(self) -> Tensor: ... + @property + def info(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 4 + n_sequence_fields: Final[_int] = 4 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _linalg_svd(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def U(self) -> Tensor: ... + @property + def S(self) -> Tensor: ... + @property + def Vh(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _lu_with_info(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def LU(self) -> Tensor: ... + @property + def pivots(self) -> Tensor: ... + @property + def info(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _scaled_dot_product_cudnn_attention(tuple[Tensor, Tensor, Tensor, Tensor, _int | SymInt, _int | SymInt, Tensor, Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def logsumexp(self) -> Tensor: ... + @property + def cum_seq_q(self) -> Tensor: ... + @property + def cum_seq_k(self) -> Tensor: ... + @property + def max_q(self) -> _int | SymInt: ... + @property + def max_k(self) -> _int | SymInt: ... + @property + def philox_seed(self) -> Tensor: ... + @property + def philox_offset(self) -> Tensor: ... + @property + def debug_attn_mask(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor, _int | SymInt, _int | SymInt, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 9 + n_sequence_fields: Final[_int] = 9 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _scaled_dot_product_efficient_attention(tuple[Tensor, Tensor, Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def log_sumexp(self) -> Tensor: ... + @property + def philox_seed(self) -> Tensor: ... + @property + def philox_offset(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 4 + n_sequence_fields: Final[_int] = 4 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _scaled_dot_product_flash_attention(tuple[Tensor, Tensor, Tensor, Tensor, _int | SymInt, _int | SymInt, Tensor, Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def logsumexp(self) -> Tensor: ... + @property + def cum_seq_q(self) -> Tensor: ... + @property + def cum_seq_k(self) -> Tensor: ... + @property + def max_q(self) -> _int | SymInt: ... + @property + def max_k(self) -> _int | SymInt: ... + @property + def rng_state(self) -> Tensor: ... + @property + def unused(self) -> Tensor: ... + @property + def debug_attn_mask(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor, Tensor, _int | SymInt, _int | SymInt, Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 9 + n_sequence_fields: Final[_int] = 9 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _scaled_dot_product_flash_attention_for_cpu(tuple[Tensor, Tensor]): # fmt: skip + @property + def output(self) -> Tensor: ... + @property + def logsumexp(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class _unpack_dual(tuple[Tensor, Tensor]): # fmt: skip + @property + def primal(self) -> Tensor: ... + @property + def tangent(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class aminmax(tuple[Tensor, Tensor]): # fmt: skip + @property + def min(self) -> Tensor: ... + @property + def max(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class cummax(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class cummin(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class frexp(tuple[Tensor, Tensor]): # fmt: skip + @property + def mantissa(self) -> Tensor: ... + @property + def exponent(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class geqrf(tuple[Tensor, Tensor]): # fmt: skip + @property + def a(self) -> Tensor: ... + @property + def tau(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class histogram(tuple[Tensor, Tensor]): # fmt: skip + @property + def hist(self) -> Tensor: ... + @property + def bin_edges(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class histogramdd(tuple[Tensor, tuple[Tensor, ...]]): # fmt: skip + @property + def hist(self) -> Tensor: ... + @property + def bin_edges(self) -> tuple[Tensor, ...]: ... + def __new__( + cls, + sequence: tuple[Tensor, tuple[Tensor, ...]], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class kthvalue(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class lu_unpack(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def P(self) -> Tensor: ... + @property + def L(self) -> Tensor: ... + @property + def U(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class max(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class median(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class min(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class mode(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class nanmedian(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class qr(tuple[Tensor, Tensor]): # fmt: skip + @property + def Q(self) -> Tensor: ... + @property + def R(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class slogdet(tuple[Tensor, Tensor]): # fmt: skip + @property + def sign(self) -> Tensor: ... + @property + def logabsdet(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class sort(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class svd(tuple[Tensor, Tensor, Tensor]): # fmt: skip + @property + def U(self) -> Tensor: ... + @property + def S(self) -> Tensor: ... + @property + def V(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 3 + n_sequence_fields: Final[_int] = 3 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class topk(tuple[Tensor, Tensor]): # fmt: skip + @property + def values(self) -> Tensor: ... + @property + def indices(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +class triangular_solve(tuple[Tensor, Tensor]): # fmt: skip + @property + def solution(self) -> Tensor: ... + @property + def cloned_coefficient(self) -> Tensor: ... + def __new__( + cls, + sequence: tuple[Tensor, Tensor], + ) -> Self: # fmt: skip + ... + n_fields: Final[_int] = 2 + n_sequence_fields: Final[_int] = 2 + n_unnamed_fields: Final[_int] = 0 + def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing + +all_return_types: list[type] = ... diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/serialization.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/serialization.py new file mode 100644 index 0000000000000000000000000000000000000000..32314ae0723e391652c4ec555257152d4c08e24b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/serialization.py @@ -0,0 +1,2201 @@ +# mypy: allow-untyped-defs +import copyreg +import difflib +import functools +import io +import os +import pickle +import re +import shutil +import struct +import sys +import tarfile +import tempfile +import threading +import warnings +from collections.abc import Callable +from contextlib import closing, contextmanager +from enum import Enum +from typing import Any, cast, Generic, IO, TypeAlias, TypeVar +from typing_extensions import TypeIs + +import torch +import torch._weights_only_unpickler as _weights_only_unpickler +from torch._sources import get_source_lines_and_file +from torch._utils import _import_dotted_name +from torch.storage import _get_dtype_from_pickle_storage_type +from torch.types import FileLike, Storage + + +__all__ = [ + "SourceChangeWarning", + "mkdtemp", + "register_package", + "check_module_version_greater_or_equal", + "validate_cuda_device", + "validate_hpu_device", + "location_tag", + "default_restore_location", + "normalize_storage_type", + "storage_to_tensor_type", + "save", + "load", + "StorageType", + "LoadEndianness", + "get_crc32_options", + "set_crc32_options", + "get_default_load_endianness", + "set_default_load_endianness", + "get_default_mmap_options", + "set_default_mmap_options", + "clear_safe_globals", + "get_safe_globals", + "add_safe_globals", + "safe_globals", + "get_unsafe_globals_in_checkpoint", + "skip_data", +] + +DEFAULT_PROTOCOL = 2 + +LONG_SIZE = struct.Struct("=l").size +INT_SIZE = struct.Struct("=i").size +SHORT_SIZE = struct.Struct("=h").size + +MAGIC_NUMBER = 0x1950A86A20F9469CFC6C +PROTOCOL_VERSION = 1001 +STORAGE_KEY_SEPARATOR = "," + +MAP_LOCATION: TypeAlias = ( + Callable[[Storage, str], Storage] | torch.device | str | dict[str, str] | None +) +STORAGE: TypeAlias = Storage | torch.storage.TypedStorage | torch.UntypedStorage + +IS_WINDOWS = sys.platform == "win32" + +UNSAFE_MESSAGE = ( + "In PyTorch 2.6, we changed the default value of the `weights_only` argument in `torch.load` " + "from `False` to `True`. Re-running `torch.load` with `weights_only` set to `False` will likely succeed, " + "but it can result in arbitrary code execution. Do it only if you got the file from a " + "trusted source." +) + +if not IS_WINDOWS: + from mmap import MAP_PRIVATE, MAP_SHARED +else: + MAP_SHARED, MAP_PRIVATE = None, None # type: ignore[assignment] + + +def _default_to_weights_only(pickle_module): + is_fbcode = not hasattr(torch.version, "git_version") + return pickle_module is None and not is_fbcode + + +# _serialization_tls is used to store thread local state specific to serialization +# that needs to be propagated to other files, in particular we use this for +# (1) map_location (needed for wrapper subclasses/third party devices to torch._utils) +# (2) skip_data (needed for torch.Tensor.__reduce_ex__ for skip_data ctx) +# (3) materialize_fake_tensors (needed for torch.Tensor.__reduce_ex__ for skip_data ctx) +class _SerializationLocal(threading.local): + def __init__(self): + super().__init__() + self.map_location: MAP_LOCATION | None = None + self.skip_data: bool = False + self.materialize_fake_tensors: bool = False + + +_serialization_tls = _SerializationLocal() + + +class SourceChangeWarning(Warning): + pass + + +@contextmanager +def mkdtemp(): + path = tempfile.mkdtemp() + try: + yield path + finally: + shutil.rmtree(path) + + +_package_registry: list[ + tuple[ + int, + Callable[[STORAGE], str | None], + Callable[[STORAGE, str], STORAGE | None], + ] +] = [] + + +class LoadEndianness(Enum): + NATIVE = 1 + LITTLE = 2 + BIG = 3 + + +def get_default_load_endianness() -> LoadEndianness | None: + """ + Get fallback byte order for loading files + + If byteorder mark is not present in saved checkpoint, + this byte order is used as fallback. + By default, it's "native" byte order. + + Returns: + default_load_endian: Optional[LoadEndianness] + """ + from torch.utils.serialization import config + + return config.load.endianness + + +def set_default_load_endianness(endianness): + """ + Set fallback byte order for loading files + + If byteorder mark is not present in saved checkpoint, + this byte order is used as fallback. + By default, it's "native" byte order. + + Args: + endianness: the new fallback byte order + """ + if not isinstance(endianness, LoadEndianness) and endianness is not None: + raise TypeError("Invalid argument type in function set_default_load_endianness") + from torch.utils.serialization import config + + config.load.endianness = endianness + + +def get_crc32_options() -> bool: + """ + Get whether :func:`torch.save` computes and writes crc32 for each record. + + Defaults to ``True``. + """ + from torch.utils.serialization import config + + return config.save.compute_crc32 + + +def set_crc32_options(compute_crc32: bool): + """ + Set whether :func:`torch.save` computes and writes crc32 for each record. + + .. note:: + Setting this to ``False`` may make unzipping of the ``torch.save`` output + fail or warn due to corrupted CRC32. However ``torch.load`` will be + able to load the file. + + Args: + compute_crc32 (bool): set crc32 computation flag + """ + from torch.utils.serialization import config + + config.save.compute_crc32 = compute_crc32 + + +def get_default_mmap_options() -> int | None: + """ + Get default mmap options for :func:`torch.load` with ``mmap=True``. + + Defaults to ``mmap.MAP_PRIVATE``. + + + Returns: + default_mmap_options: int + """ + from torch.utils.serialization import config + + return config.load.mmap_flags + + +def _get_storage_alignment() -> int: + """ + Gets alignment for storages in torch.save files/ + + Defaults to 64. + + Returns: + storage_alginment: int + """ + from torch.utils.serialization import config + + return config.save.storage_alignment + + +class set_default_mmap_options: + """ + Context manager or function to set default mmap options for :func:`torch.load` with ``mmap=True`` to flags. + + For now, only either ``mmap.MAP_PRIVATE`` or ``mmap.MAP_SHARED`` are supported. + Please open an issue if you need any other option to be added here. + + .. note:: + This feature is currently not supported for Windows. + + Args: + flags: ``mmap.MAP_PRIVATE`` or ``mmap.MAP_SHARED`` + """ + + def __init__(self, flags: int) -> None: + if IS_WINDOWS: + raise RuntimeError( + "Changing the default mmap options is currently not supported for Windows" + ) + if flags != MAP_PRIVATE and flags != MAP_SHARED: + raise ValueError( + "Invalid argument in function set_default_mmap_options, " + f"expected mmap.MAP_PRIVATE or mmap.MAP_SHARED, but got {flags}" + ) + # global config + from torch.utils.serialization import config + + self.prev = config.load.mmap_flags + config.load.mmap_flags = flags + + def __enter__(self) -> None: + pass + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + from torch.utils.serialization import config + + config.load.mmap_flags = self.prev + + +def clear_safe_globals() -> None: + """ + Clears the list of globals that are safe for ``weights_only`` load. + """ + _weights_only_unpickler._clear_safe_globals() + + +def get_safe_globals() -> list[Callable | tuple[Callable, str]]: + """ + Returns the list of user-added globals that are safe for ``weights_only`` load. + """ + return _weights_only_unpickler._get_safe_globals() + + +def add_safe_globals(safe_globals: list[Callable | tuple[Callable, str]]) -> None: + """ + Marks the given globals as safe for ``weights_only`` load. For example, functions + added to this list can be called during unpickling, classes could be instantiated + and have state set. + + Each item in the list can either be a function/class or a tuple of the form + (function/class, string) where string is the full path of the function/class. + + Within the serialized format, each function is identified with its full + path as ``{__module__}.{__qualname__}``. When calling this API, you can provide this + full path that should match the one in the checkpoint otherwise the default + ``{fn.__module__}.{fn.__qualname__}`` will be used. + + Args: + safe_globals (List[Union[Callable, Tuple[Callable, str]]]): list of globals to mark as safe + + Example: + >>> # xdoctest: +SKIP("Can't torch.save(t, ...) as doctest thinks MyTensor is defined on torch.serialization") + >>> import tempfile + >>> class MyTensor(torch.Tensor): + ... pass + >>> t = MyTensor(torch.randn(2, 3)) + >>> with tempfile.NamedTemporaryFile() as f: + ... torch.save(t, f.name) + # Running `torch.load(f.name, weights_only=True)` will fail with + # Unsupported global: GLOBAL __main__.MyTensor was not an allowed global by default. + # Check the code and make sure MyTensor is safe to be used when loaded from an arbitrary checkpoint. + ... torch.serialization.add_safe_globals([MyTensor]) + ... torch.load(f.name, weights_only=True) + # MyTensor([[-0.5024, -1.8152, -0.5455], + # [-0.8234, 2.0500, -0.3657]]) + """ + _weights_only_unpickler._add_safe_globals(safe_globals) + + +class safe_globals(_weights_only_unpickler._safe_globals): + r"""Context-manager that adds certain globals as safe for ``weights_only`` load. + + Args: + safe_globals: List of globals for weights_only load. + + Example: + >>> # xdoctest: +SKIP("Can't torch.save(t, ...) as doctest thinks MyTensor is defined on torch.serialization") + >>> import tempfile + >>> class MyTensor(torch.Tensor): + ... pass + >>> t = MyTensor(torch.randn(2, 3)) + >>> with tempfile.NamedTemporaryFile() as f: + ... torch.save(t, f.name) + # Running `torch.load(f.name, weights_only=True)` will fail with + # Unsupported global: GLOBAL __main__.MyTensor was not an allowed global by default. + # Check the code and make sure MyTensor is safe to be used when loaded from an arbitrary checkpoint. + ... with torch.serialization.safe_globals([MyTensor]): + ... torch.load(f.name, weights_only=True) + # MyTensor([[-0.5024, -1.8152, -0.5455], + # [-0.8234, 2.0500, -0.3657]]) + >>> assert torch.serialization.get_safe_globals() == [] + """ + + +def get_unsafe_globals_in_checkpoint(f: FileLike) -> list[str]: + """Returns a list of strings of functions/classes in a ``torch.save`` object that are not safe for ``weights_only``. + + For a given function or class ``f``, the corresponding string will be of the form + ``{f.__module__}.{f.__name__}``. + + This function will return any GLOBALs in the checkpoint that are not in the set marked safe + for ``weights_only`` (either via :func:`add_safe_globals` or :class:`safe_globals` context or + allowlisted by ``torch`` by default). + + .. note:: + This function will statically disassemble the pickle file in the checkpoint. + The implication is any classes dynamically pushed onto the stack during unpickling + will not be included in the output. + + Args: + f: File-like object or string containing the checkpoint object saved via ``torch.save`` + + Returns: + A list of strings of pickle GLOBALs in the checkpoint that are not allowlisted for ``weights_only``. + """ + default_safe_globals_strings = set( + _weights_only_unpickler._get_allowed_globals().keys() + ) + user_safe_global_strings = set( + _weights_only_unpickler._get_user_allowed_globals().keys() + ) + safe_global_strings = default_safe_globals_strings.union(user_safe_global_strings) + + with _open_file_like(f, "rb") as opened_file: + if not _is_zipfile(opened_file): + raise ValueError("Expected input to be a checkpoint returned by torch.save") + with _open_zipfile_reader(opened_file) as zip_file: + if _is_torchscript_zip(zip_file): + raise ValueError( + "Expected input to be a checkpoint returned by torch.save but got a torchscript checkpoint" + ) + data_file = io.BytesIO(zip_file.get_record("data.pkl")) + all_globals = _weights_only_unpickler.get_globals_in_pkl(data_file) + return list(all_globals.difference(safe_global_strings)) + + +class skip_data: + """ + Context-manager that skips writing/reading storage bytes for ``torch.save`` / ``torch.load`` calls. + + For the save path, storages will still be saved, but the space that their bytes would usually be written to + will be empty space. The storage bytes can then be populated in a separate pass. + + For the load path, tensors will be loaded per the checkpoint but their storages will not be populated with data. + + .. warning:: + The ``skip_data`` context manager is an early prototype and is subject to change. + + Args: + materialize_fake_tensors: Whether to materialize FakeTensors during save. This is a no-op for the load path. + + Example: + >>> # xdoctest: +SKIP("NamedTemporaryFile on Windows") + >>> import tempfile + >>> t = torch.randn(2, 3) + >>> with tempfile.NamedTemporaryFile() as f: + ... with torch.serialization.skip_data(): + ... torch.save(t, f.name) + ... torch.load(f.name, weights_only=True) + tensor([[0., 0., 0.], + [0., 0., 0.]]) + """ + + def __init__(self, materialize_fake_tensors: bool = False): + self.materialize_fake_tensors = materialize_fake_tensors + + def __enter__(self): + global _serialization_tls + self._old_skip_data = _serialization_tls.skip_data + self._old_materialize_fake_tensors = _serialization_tls.materialize_fake_tensors + _serialization_tls.skip_data = True + _serialization_tls.materialize_fake_tensors = self.materialize_fake_tensors + + def __exit__(self, type, value, tb): + global _serialization_tls + _serialization_tls.skip_data = self._old_skip_data + _serialization_tls.materialize_fake_tensors = self._old_materialize_fake_tensors + + +def _is_zipfile(f) -> bool: + # This is a stricter implementation than zipfile.is_zipfile(). + # zipfile.is_zipfile() is True if the magic number appears anywhere in the + # binary. Since we expect the files here to be generated by torch.save or + # torch.jit.save, it's safe to only check the start bytes and avoid + # collisions and assume the zip has only 1 file. + # See bugs.python.org/issue28494. + + start = f.tell() + # Read the first few bytes and match against the ZIP file signature + local_header_magic_number = b"PK\x03\x04" + read_bytes = f.read(len(local_header_magic_number)) + f.seek(start) + return read_bytes == local_header_magic_number + + +def register_package( + priority: int, + tagger: Callable[[STORAGE], str | None], + deserializer: Callable[[STORAGE, str], STORAGE | None], +): + """ + Registers callables for tagging and deserializing storage objects with an associated priority. + Tagging associates a device with a storage object at save time while deserializing moves a + storage object to an appropriate device at load time. :attr:`tagger` and :attr:`deserializer` + are run in the order given by their :attr:`priority` until a tagger/deserializer returns a + value that is not `None`. + + To override the deserialization behavior for a device in the global registry, one can register a + tagger with a higher priority than the existing tagger. + + This function can also be used to register a tagger and deserializer for new devices. + + Args: + priority: Indicates the priority associated with the tagger and deserializer, where a lower + value indicates higher priority. + tagger: Callable that takes in a storage object and returns its tagged device as a string + or None. + deserializer: Callable that takes in storage object and a device string and returns a storage + object on the appropriate device or None. + + Returns: + `None` + + Example: + >>> def ipu_tag(obj): + >>> if obj.device.type == 'ipu': + >>> return 'ipu' + >>> def ipu_deserialize(obj, location): + >>> if location.startswith('ipu'): + >>> ipu = getattr(torch, "ipu", None) + >>> assert ipu is not None, "IPU device module is not loaded" + >>> assert torch.ipu.is_available(), "ipu is not available" + >>> return obj.ipu(location) + >>> torch.serialization.register_package(11, ipu_tag, ipu_deserialize) + """ + queue_elem = (priority, tagger, deserializer) + _package_registry.append(queue_elem) + _package_registry.sort() + + +def check_module_version_greater_or_equal( + module, + req_version_tuple, + error_if_malformed=True, +): + """ + Check if a module's version satisfies requirements + + Usually, a module's version string will be like 'x.y.z', which would be represented + as a tuple (x, y, z), but sometimes it could be an unexpected format. If the version + string does not match the given tuple's format up to the length of the tuple, then + error and exit or emit a warning. + + Args: + module: the module to check the version of + req_version_tuple: tuple (usually of ints) representing the required version + error_if_malformed: whether we should exit if module version string is malformed + + Returns: + requirement_is_met: bool + """ + try: + version_strs = module.__version__.split(".") + # Cast module version fields to match the types of the required version + module_version = tuple( + type(req_field)(version_strs[idx]) + for idx, req_field in enumerate(req_version_tuple) + ) + requirement_is_met = module_version >= req_version_tuple + + except Exception as e: + message = ( + f"'{module.__name__}' module version string is malformed '{module.__version__}' and cannot be compared" + f" with tuple {str(req_version_tuple)}" + ) + if error_if_malformed: + raise RuntimeError(message) from e + else: + warnings.warn( + message + ", but continuing assuming that requirement is met", + stacklevel=2, + ) + requirement_is_met = True + + return requirement_is_met + + +def _cpu_tag(obj): + if obj.device.type == "cpu": + return "cpu" + + +def _mps_tag(obj): + if obj.device.type == "mps": + return "mps" + + +def _meta_tag(obj): + if obj.device.type == "meta": + return "meta" + + +def _backend_tag(backend_name, obj): + if backend_name == "privateuse1": + backend_name = torch._C._get_privateuse1_backend_name() + if obj.device.type == backend_name: + if obj.device.index is None: + return backend_name + else: + return backend_name + ":" + str(obj.device.index) + + +def _cpu_deserialize(obj, location): + if location == "cpu": + return obj + + +def _mps_deserialize(obj, location): + if location.startswith("mps"): + return obj.mps() + + +def _meta_deserialize(obj, location): + if location == "meta": + return torch.UntypedStorage(obj.nbytes(), device="meta") + + +def _is_meta_location(map_location): + """ + Check if map_location specifies the meta device. + + This is used to skip reading storage data from disk when loading + to the meta device, since meta tensors don't hold actual data. + + Args: + map_location: The map_location argument passed to torch.load + + Returns: + True if map_location is definitively the meta device, False otherwise. + For dict or callable map_location, returns False since we can't + easily determine the target device without evaluating it. + """ + if map_location is None: + return False + if isinstance(map_location, str): + return map_location == "meta" + if isinstance(map_location, torch.device): + return map_location.type == "meta" + # dict or callable - can't easily determine + return False + + +def _validate_device(location, backend_name): + """ + Check whether the device index of specified backend is valid + + In case of privateuse1 backend, your must first register a device_module for + privateuse1 using torch._register_device_module. Implement the following + methods in device_module like cuda: device_module._utils._get_device_index(location, True), + device_module.device_count(). + + Args: + location: string of device + backend_name: the backend name or the name of privateuse1, which can be renamed + + Returns: + device_index: int + """ + if not hasattr(torch, backend_name): + raise RuntimeError( + f"The {backend_name.upper()} device module is not registered. " + "If you are running on a CPU-only machine, " + "please use torch.load with map_location=torch.device('cpu') " + "to map your storages to the CPU." + ) + device_module = getattr(torch, backend_name) + if hasattr(device_module, "_utils") and hasattr( + device_module._utils, "_get_device_index" + ): + device_index = device_module._utils._get_device_index(location, True) + device = torch.device(backend_name, device_index) + else: + device = torch.device(location) + device_index = device.index if device.index else 0 + if hasattr(device_module, "is_available") and not device_module.is_available(): + raise RuntimeError( + f"Attempting to deserialize object on a {backend_name.upper()} " + f"device but torch.{backend_name}.is_available() is False. " + "If you are running on a CPU-only machine, " + "please use torch.load with map_location=torch.device('cpu') " + "to map your storages to the CPU." + ) + if hasattr(device_module, "device_count"): + device_count = device_module.device_count() + if device_index >= device_count: + raise RuntimeError( + f"Attempting to deserialize object on {backend_name.upper()} device " + f"{device_index} but torch.{backend_name}.device_count() is {device_count}. " + "Please use torch.load with map_location to map your storages " + "to an existing device." + ) + return device + + +def validate_cuda_device(location): + return _validate_device(location, "cuda").index + + +def validate_hpu_device(location): + return _validate_device(location, "hpu").index + + +def _deserialize(backend_name, obj, location): + if backend_name == "privateuse1": + backend_name = torch._C._get_privateuse1_backend_name() + if location == backend_name or bool( + re.match(f"{backend_name}(:|[0-9]+)", location) + ): + device = _validate_device(location, backend_name) + return obj.to(device=device) + + +register_package(10, _cpu_tag, _cpu_deserialize) +register_package( + 20, + functools.partial(_backend_tag, "cuda"), + functools.partial(_deserialize, "cuda"), +) +register_package(21, _mps_tag, _mps_deserialize) +register_package(22, _meta_tag, _meta_deserialize) +register_package( + 23, + functools.partial(_backend_tag, "privateuse1"), + functools.partial(_deserialize, "privateuse1"), +) +register_package( + 24, + functools.partial(_backend_tag, "hpu"), + functools.partial(_deserialize, "hpu"), +) +register_package( + 25, + functools.partial(_backend_tag, "xpu"), + functools.partial(_deserialize, "xpu"), +) +register_package( + 26, + functools.partial(_backend_tag, "mtia"), + functools.partial(_deserialize, "mtia"), +) + + +def location_tag( + storage: Storage | torch.storage.TypedStorage | torch.UntypedStorage, +): + for _, tagger, _ in _package_registry: + location = tagger(storage) + if location: + return location + raise RuntimeError( + "don't know how to determine data location of " + torch.typename(storage) + ) + + +def default_restore_location(storage, location): + """ + Restores `storage` using a deserializer function registered for the `location`. + + This function looks in the registry for deserializer functions that match the `location`. + If found, it attempts to use them, in priority order, to restore `storage` until one + returns a not `None` result. If no deserializer can be found in the registry, or all found fail + to bear a result, it raises a `RuntimeError`. + + Args: + storage (STORAGE): the storage object to restore + location (str): the location tag associated with the storage object + + Returns: + storage: Optional[STORAGE] + + Raises: + RuntimeError: If no deserializer matching `location` is found in the registry or if + all matching ones return `None`. + """ + for _, _, fn in _package_registry: + result = fn(storage, location) + if result is not None: + return result + raise RuntimeError( + "don't know how to restore data location of " + + torch.typename(storage) + + " (tagged with " + + location + + ")" + ) + + +def normalize_storage_type(storage_type): + return getattr(torch, storage_type.__name__) + + +def storage_to_tensor_type(storage): + storage_type = type(storage) + module = _import_dotted_name(storage_type.__module__) + return getattr(module, storage_type.__name__.replace("Storage", "Tensor")) + + +def _is_path(name_or_buffer: object) -> TypeIs[str | os.PathLike]: + return isinstance(name_or_buffer, (str, os.PathLike)) + + +T = TypeVar("T") + + +class _opener(Generic[T]): + def __init__(self, file_like: T) -> None: + self.file_like: T = file_like + + def __enter__(self): + return self.file_like + + def __exit__(self, *args): + pass + + +class _open_file(_opener[IO[bytes]]): + def __init__(self, name: str | os.PathLike[str], mode: str) -> None: + super().__init__(open(name, mode)) # noqa: SIM115 + + def __exit__(self, *args): + self.file_like.close() + + +class _open_buffer_reader(_opener[IO[bytes]]): + def __init__(self, buffer: IO[bytes]) -> None: + super().__init__(buffer) + _check_seekable(buffer) + + +class _open_buffer_writer(_opener[IO[bytes]]): + def __exit__(self, *args): + self.file_like.flush() + + +def _open_file_like(name_or_buffer: FileLike, mode: str) -> _opener[IO[bytes]]: + if _is_path(name_or_buffer): + return _open_file(name_or_buffer, mode) + else: + if "w" in mode: + return _open_buffer_writer(name_or_buffer) + elif "r" in mode: + return _open_buffer_reader(name_or_buffer) + else: + raise RuntimeError(f"Expected 'r' or 'w' in mode but got {mode}") + + +class _open_zipfile_reader(_opener[torch._C.PyTorchFileReader]): + def __init__(self, name_or_buffer: str | IO[bytes]) -> None: + super().__init__(torch._C.PyTorchFileReader(name_or_buffer)) + + +class _open_zipfile_writer_file(_opener[torch._C.PyTorchFileWriter]): + def __init__(self, name: str) -> None: + self.file_stream = None + self.name = name + try: + self.name.encode("ascii") + except UnicodeEncodeError: + # PyTorchFileWriter only supports ascii filename. + # For filenames with non-ascii characters, we rely on Python + # for writing out the file. + # pyrefly: ignore [bad-assignment] + self.file_stream = io.FileIO(self.name, mode="w") + super().__init__( + torch._C.PyTorchFileWriter( # pyrefly: ignore # no-matching-overload + self.file_stream, get_crc32_options(), _get_storage_alignment() + ) + ) + else: + super().__init__( + torch._C.PyTorchFileWriter( + self.name, get_crc32_options(), _get_storage_alignment() + ) + ) + + def __exit__(self, *args) -> None: + self.file_like.write_end_of_file() + if self.file_stream is not None: + self.file_stream.close() + + +class _open_zipfile_writer_buffer(_opener[torch._C.PyTorchFileWriter]): + def __init__(self, buffer: IO[bytes]) -> None: + if not callable(getattr(buffer, "write", None)): + msg = f"Buffer of {str(type(buffer)).strip('<>')} has no callable attribute 'write'" + if not hasattr(buffer, "write"): + raise AttributeError(msg) + raise TypeError(msg) + self.buffer = buffer + super().__init__( + torch._C.PyTorchFileWriter( + buffer, get_crc32_options(), _get_storage_alignment() + ) + ) + + def __exit__(self, *args) -> None: + self.file_like.write_end_of_file() + self.buffer.flush() + + +def _open_zipfile_writer(name_or_buffer: str | IO[bytes]) -> _opener: + container: type[_opener] + if _is_path(name_or_buffer): + container = _open_zipfile_writer_file + else: + container = _open_zipfile_writer_buffer + return container(name_or_buffer) # type: ignore[arg-type] + + +def _is_compressed_file(f) -> bool: + compress_modules = ["gzip"] + try: + return f.__module__ in compress_modules + except AttributeError: + return False + + +def _should_read_directly(f): + """ + Checks if f is a file that should be read directly. It should be read + directly if it is backed by a real file (has a fileno) and is not a + a compressed file (e.g. gzip) + """ + if _is_compressed_file(f): + return False + try: + return f.fileno() >= 0 + except io.UnsupportedOperation: + return False + except AttributeError: + return False + + +def _check_seekable(f) -> bool: + def raise_err_msg(patterns, e): + for p in patterns: + if p in str(e): + msg = ( + str(e) + + ". You can only torch.load from a file that is seekable." + + " Please pre-load the data into a buffer like io.BytesIO and" + + " try to load from it instead." + ) + raise type(e)(msg) + raise e + + try: + f.seek(f.tell()) + return True + except (io.UnsupportedOperation, AttributeError) as e: + raise_err_msg(["seek", "tell"], e) + return False + + +def _check_dill_version(pickle_module) -> None: + """Checks if using dill as the pickle module, and if so, checks if it is the correct version. + If dill version is lower than 0.3.1, a ValueError is raised. + + Args: + pickle_module: module used for pickling metadata and objects + + """ + if pickle_module is not None and pickle_module.__name__ == "dill": + required_dill_version = (0, 3, 1) + if not check_module_version_greater_or_equal( + pickle_module, required_dill_version, False + ): + raise ValueError( + ( + "'torch' supports dill >= {}, but you have dill {}." + " Please upgrade dill or switch to 'pickle'" + ).format( + ".".join([str(num) for num in required_dill_version]), + pickle_module.__version__, + ) + ) + + +def _check_save_filelike(f): + if not _is_path(f) and not hasattr(f, "write"): + raise AttributeError( + "expected 'f' to be string, path, or a file-like object with " + "a 'write' attribute" + ) + + +def save( + obj: object, + f: FileLike, + pickle_module: Any = pickle, + pickle_protocol: int = DEFAULT_PROTOCOL, + _use_new_zipfile_serialization: bool = True, + _disable_byteorder_record: bool = False, +) -> None: + # Reference: https://github.com/pytorch/pytorch/issues/54354 + # The first line of this docstring overrides the one Sphinx generates for the + # documentation. We need it so that Sphinx doesn't leak `pickle`s path from + # the build environment (e.g. `>> # xdoctest: +SKIP("makes cwd dirty") + >>> # Save to file + >>> x = torch.tensor([0, 1, 2, 3, 4]) + >>> torch.save(x, "tensor.pt") + >>> # Save to io.BytesIO buffer + >>> buffer = io.BytesIO() + >>> torch.save(x, buffer) + """ + torch._C._log_api_usage_once("torch.save") + _check_dill_version(pickle_module) + _check_save_filelike(f) + + if isinstance(f, (str, os.PathLike)): + f = os.fspath(f) + + if _use_new_zipfile_serialization: + with _open_zipfile_writer(f) as opened_zipfile: + _save( + obj, + opened_zipfile, + pickle_module, + pickle_protocol, + _disable_byteorder_record, + ) + return + else: + global _serialization_tls + if _serialization_tls.skip_data: + raise RuntimeError( + "Cannot use skip_data=True with _use_new_zipfile_serialization=False" + ) + with _open_file_like(f, "wb") as opened_file: + _legacy_save(obj, opened_file, pickle_module, pickle_protocol) + + +def _legacy_save(obj, f, pickle_module, pickle_protocol) -> None: + import torch.nn as nn + + serialized_container_types = {} + serialized_storages: dict[str, tuple[torch.UntypedStorage, torch.dtype]] = {} + + # Since loading storages that view the same data with different dtypes is + # not supported, we need to keep track of the dtype associated with each + # storage data_ptr and throw an error if the dtype is ever different. + # TODO: This feature could be added in the future + storage_dtypes: dict[int, torch.dtype] = {} + + def persistent_id(obj: Any) -> tuple | None: + # FIXME: the docs say that persistent_id should only return a string + # but torch store returns tuples. This works only in the binary protocol + # see + # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects + # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537 + if isinstance(obj, type) and issubclass(obj, nn.Module): + if obj in serialized_container_types: + return None + serialized_container_types[obj] = True + source_file = source = None + try: + source_lines, _, source_file = get_source_lines_and_file(obj) + source = "".join(source_lines) + except ( + Exception + ): # saving the source is optional, so we can ignore any errors + warnings.warn( + "Couldn't retrieve source code for container of " + "type " + obj.__name__ + ". It won't be checked " + "for correctness upon loading.", + stacklevel=2, + ) + return ("module", obj, source_file, source) + + if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj): + storage: torch.UntypedStorage + + if isinstance(obj, torch.storage.TypedStorage): + # TODO: Once we decide to break serialization FC, this case + # can be deleted + storage = obj._untyped_storage + storage_dtype = obj.dtype + storage_type_str = obj._pickle_storage_type() + storage_type = getattr(torch, storage_type_str) + dtype = obj.dtype + storage_numel = obj._size() + + elif isinstance(obj, torch.UntypedStorage): + storage = obj + storage_dtype = torch.uint8 + storage_type = normalize_storage_type(type(obj)) + dtype = torch.uint8 + storage_numel = storage.nbytes() + else: + raise TypeError(f"type not recognized: {type(obj)}") + + # If storage is allocated, ensure that any other saved storages + # pointing to the same data all have the same dtype. If storage is + # not allocated, don't perform this check + if storage.data_ptr() != 0: + if storage.data_ptr() in storage_dtypes: + if storage_dtype != storage_dtypes[storage.data_ptr()]: + raise RuntimeError( + "Cannot save multiple tensors or storages that " + "view the same data as different types" + ) + else: + storage_dtypes[storage.data_ptr()] = storage_dtype + + view_metadata: tuple[str, int, int] | None + + # Offset is always 0, but we keep it for backwards compatibility + # with the old serialization format (which supported storage views) + offset = 0 + storage_key = str(storage._cdata) + location = location_tag(storage) + + # TODO: There's an issue here with FC. It might be impossible to + # solve, but it's worth noting. Imagine we save a list `[storage, + # tensor]`, where `tensor.storage()` is the same as `storage`, and + # `tensor.element_size() > 1`. Let's say that `tensor.dtype == + # torch.float`. The storage will be serialized with element size + # of 1, since we're choosing to serialize the first occurrence of + # a duplicate storage. Since this legacy serialization format saves + # the numel of the storage, rather than nbytes directly, we'll be + # effectively saving nbytes in this case. We'll be able to load it + # and the tensor back up with no problems in _this_ and future + # versions of pytorch, but in older versions, here's the problem: + # the storage will be loaded up as a UntypedStorage, and then the + # FloatTensor will loaded and the UntypedStorage will be assigned to + # it. Since the storage dtype does not match the tensor dtype, this + # will cause an error. If we reverse the list, like `[tensor, + # storage]`, then we will save the `tensor.storage()` as a faked + # `FloatStorage`, and the saved size will be the correct + # dtype-specific numel count that old versions expect. `tensor` + # will be able to load up properly in old versions, pointing to + # a FloatStorage. However, `storage` is still being translated to + # a UntypedStorage, and it will try to resolve to the same + # FloatStorage that `tensor` contains. This will also cause an + # error. It doesn't seem like there's any way around this. + # Probably, we just cannot maintain FC for the legacy format if the + # saved list contains both a tensor and a storage that point to the + # same data. We should still be able to maintain FC for lists of + # just tensors, as long as all views share the same dtype as the + # tensor they are viewing. + + if storage_key not in serialized_storages: + serialized_storages[storage_key] = (storage, dtype) + is_view = storage._cdata != storage._cdata + if is_view: + view_metadata = (str(storage._cdata), offset, storage.nbytes()) + else: + view_metadata = None + + res = ( + "storage", + storage_type, + storage_key, + location, + storage_numel, + view_metadata, + ) + return res + return None + + sys_info = { + "protocol_version": PROTOCOL_VERSION, + "little_endian": sys.byteorder == "little", + "type_sizes": { + "short": SHORT_SIZE, + "int": INT_SIZE, + "long": LONG_SIZE, + }, + } + + pickle_module.dump(MAGIC_NUMBER, f, protocol=pickle_protocol) + pickle_module.dump(PROTOCOL_VERSION, f, protocol=pickle_protocol) + pickle_module.dump(sys_info, f, protocol=pickle_protocol) + + class PyTorchLegacyPickler(pickle_module.Pickler): + def persistent_id(self, obj): + return persistent_id(obj) # noqa: F821 + + pickler = PyTorchLegacyPickler(f, protocol=pickle_protocol) + pickler.dump(obj) + + # The class def keeps the persistent_id closure alive, leaking memory. + del persistent_id + + serialized_storage_keys = sorted(serialized_storages.keys()) + pickle_module.dump(serialized_storage_keys, f, protocol=pickle_protocol) + f.flush() + for key in serialized_storage_keys: + storage, dtype = serialized_storages[key] + storage._write_file( + f, _should_read_directly(f), True, torch._utils._element_size(dtype) + ) + + +def _save( + obj, + zip_file, + pickle_module, + pickle_protocol, + _disable_byteorder_record, +): + serialized_storages: dict[str, torch.storage.UntypedStorage] = {} + id_map: dict[int, str] = {} + + # Since loading storages that view the same data with different dtypes is + # not supported, we need to keep track of the dtype associated with each + # storage data_ptr and throw an error if the dtype is ever different. + # TODO: This feature could be added in the future + storage_dtypes: dict[int, torch.dtype] = {} + + def persistent_id(obj): + # FIXME: the docs say that persistent_id should only return a string + # but torch store returns tuples. This works only in the binary protocol + # see + # https://docs.python.org/2/library/pickle.html#pickling-and-unpickling-external-objects + # https://github.com/python/cpython/blob/master/Lib/pickle.py#L527-L537 + if isinstance(obj, torch.storage.TypedStorage) or torch.is_storage(obj): + if isinstance(obj, torch.storage.TypedStorage): + # TODO: Once we decide to break serialization FC, this case + # can be deleted + storage = obj._untyped_storage + storage_dtype = obj.dtype + storage_type_str = obj._pickle_storage_type() + storage_type = getattr(torch, storage_type_str) + storage_numel = obj._size() + + else: + storage = obj + storage_dtype = torch.uint8 + storage_type = normalize_storage_type(type(obj)) + storage_numel = storage.nbytes() + + # If storage is allocated, ensure that any other saved storages + # pointing to the same data all have the same dtype. If storage is + # not allocated, don't perform this check + if str(storage.device) != "meta" and storage.data_ptr() != 0: + if storage.data_ptr() in storage_dtypes: + if storage_dtype != storage_dtypes[storage.data_ptr()]: + raise RuntimeError( + "Cannot save multiple tensors or storages that " + "view the same data as different types" + ) + else: + storage_dtypes[storage.data_ptr()] = storage_dtype + + storage_key = id_map.setdefault(storage._cdata, str(len(id_map))) + if hasattr(obj, "_fake_device") and obj._fake_device is not None: + location = str(obj._fake_device) + else: + location = location_tag(storage) + serialized_storages[storage_key] = storage + + return ("storage", storage_type, storage_key, location, storage_numel) + + return None + + # Write the pickle data for `obj` + data_buf = io.BytesIO() + + class PyTorchPickler(pickle_module.Pickler): # type: ignore[name-defined] + def persistent_id(self, obj): + return persistent_id(obj) # noqa: F821 + + pickler = PyTorchPickler(data_buf, protocol=pickle_protocol) + pickler.dump(obj) + + # The class def keeps the persistent_id closure alive, leaking memory. + del persistent_id + + data_value = data_buf.getvalue() + zip_file.write_record("data.pkl", data_value, len(data_value)) + # .format_version is used to track + # 1. version 1 represents the order of storages being changed from + # lexicographical based on keys to numerically ordered based on keys + # 2. version 2 represents including storage_alignment as a record + # within the zipfile + zip_file.write_record(".format_version", "1", len("1")) + storage_alignment = str(_get_storage_alignment()) + zip_file.write_record( + ".storage_alignment", storage_alignment, len(storage_alignment) + ) + + # Write byte order marker + if not _disable_byteorder_record: + if sys.byteorder not in ["little", "big"]: + raise ValueError("Unknown endianness type: " + sys.byteorder) + + zip_file.write_record("byteorder", sys.byteorder, len(sys.byteorder)) + + # Write each tensor to a file named tensor/the_tensor_key in the zip archive + for key in serialized_storages: + name = f"data/{key}" + storage = serialized_storages[key] + num_bytes = storage.nbytes() + global _serialization_tls + if _serialization_tls.skip_data: + zip_file.write_record_metadata(name, num_bytes) + else: + # given that we copy things around anyway, we might use storage.cpu() + # this means to that to get tensors serialized, you need to implement + # .cpu() on the underlying Storage + if storage.device.type != "cpu": + from torch.utils.serialization import config + + if ( + config.save.use_pinned_memory_for_d2h + and ( + acc := torch.accelerator.current_accelerator( + check_available=True + ) + ) + is not None + and acc.type == storage.device.type + ): + new_storage = torch.empty( + num_bytes, dtype=torch.uint8, device="cpu", pin_memory=True + ).untyped_storage() + new_storage.copy_(storage) + torch.accelerator.current_stream(storage.device.index).synchronize() + storage = new_storage + else: + storage = storage.cpu() + # Now that it is on the CPU we can directly copy it into the zip file + zip_file.write_record(name, storage, num_bytes) + + +def load( + f: FileLike, + map_location: MAP_LOCATION = None, + pickle_module: Any = None, + *, + weights_only: bool | None = None, + mmap: bool | None = None, + **pickle_load_args: Any, +) -> Any: + # Reference: https://github.com/pytorch/pytorch/issues/54354 + # The first line of this docstring overrides the one Sphinx generates for the + # documentation. We need it so that Sphinx doesn't leak `pickle`s path from + # the build environment (e.g. `>> # xdoctest: +SKIP("undefined filepaths") + >>> torch.load("tensors.pt", weights_only=True) + # Load all tensors onto the CPU + >>> torch.load( + ... "tensors.pt", + ... map_location=torch.device("cpu"), + ... weights_only=True, + ... ) + # Load all tensors onto the CPU, using a function + >>> torch.load( + ... "tensors.pt", + ... map_location=lambda storage, loc: storage, + ... weights_only=True, + ... ) + # Load all tensors onto GPU 1 + >>> torch.load( + ... "tensors.pt", + ... map_location=lambda storage, loc: storage.cuda(1), # type: ignore[attr-defined] + ... weights_only=True, + ... ) # type: ignore[attr-defined] + # Map tensors from GPU 1 to GPU 0 + >>> torch.load( + ... "tensors.pt", + ... map_location={"cuda:1": "cuda:0"}, + ... weights_only=True, + ... ) + # Load tensor from io.BytesIO object + # Loading from a buffer setting weights_only=False, warning this can be unsafe + >>> with open("tensor.pt", "rb") as f: + ... buffer = io.BytesIO(f.read()) + >>> torch.load(buffer, weights_only=False) + # Load a module with 'ascii' encoding for unpickling + # Loading from a module setting weights_only=False, warning this can be unsafe + >>> torch.load("module.pt", encoding="ascii", weights_only=False) + """ + torch._C._log_api_usage_once("torch.load") + DOCS_MESSAGE = ( + "\n\nCheck the documentation of torch.load to learn more about types accepted by default with " + "weights_only https://pytorch.org/docs/stable/generated/torch.load.html." + ) + + def _get_wo_message(message: str) -> str: + unsafe_global_pattern = r"GLOBAL (\S+) was not an allowed global by default." + has_unsafe_global = re.search(unsafe_global_pattern, message) is not None + blocklist_pattern = r"whose module (\S+) is blocked" + has_blocklist = re.search(blocklist_pattern, message) is not None + import_pattern = r"(\S+) must be (\S+) to load" + has_import = re.search(import_pattern, message) is not None + if has_unsafe_global: + updated_message = ( + "Weights only load failed. This file can still be loaded, to do so you have two options, " + "\033[1mdo those steps only if you trust the source of the checkpoint\033[0m. " + f"\n\t(1) {UNSAFE_MESSAGE}\n\t(2) Alternatively, to load with `weights_only=True` please check " + "the recommended steps in the following error message.\n\tWeightsUnpickler error: " + + message + ) + else: + if has_import: + return f"Weights only load failed. {message}\n {UNSAFE_MESSAGE}\n" + else: + updated_message = f"Weights only load failed. {UNSAFE_MESSAGE}\n" + if not has_blocklist: + updated_message += ( + "Please file an issue with the following so that we can make " + "`weights_only=True` compatible with your use case: WeightsUnpickler error: " + ) + updated_message += "\n\n" + message + return updated_message + DOCS_MESSAGE + + weights_only_not_set = weights_only is None + + if weights_only_not_set: + weights_only = _default_to_weights_only(pickle_module) + + true_values = ["1", "y", "yes", "true"] + # Add ability to force safe only or non-safe weight loads via environment variables + force_weights_only_load = ( + os.getenv("TORCH_FORCE_WEIGHTS_ONLY_LOAD", "0") in true_values + ) + force_no_weights_only_load = ( + os.getenv("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "0") in true_values + ) + + if force_weights_only_load and force_no_weights_only_load: + raise RuntimeError( + "Only one of `TORCH_FORCE_WEIGHTS_ONLY_LOAD` or `TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD` " + "should be set, but both were set." + ) + elif force_weights_only_load: + weights_only = True + elif force_no_weights_only_load: + # TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD can only override if callsite did not explicitly set weights_only + if weights_only_not_set: + warnings.warn( + "Environment variable TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD detected, since the" + "`weights_only` argument was not explicitly passed to `torch.load`, forcing weights_only=False.", + UserWarning, + stacklevel=2, + ) + weights_only = False + + if weights_only: + if pickle_module is not None: + raise RuntimeError( + "Can not safely load weights when explicit pickle_module is specified" + ) + else: + if pickle_module is None: + pickle_module = pickle + + if pickle_load_args != {} and weights_only: + warnings.warn("pickle_load_args only works if `weights_only=False`.") + + # make flipping default BC-compatible + if mmap is None: + from torch.utils.serialization import config + + mmap = config.load.mmap + + _check_dill_version(pickle_module) + + if "encoding" not in pickle_load_args: + pickle_load_args["encoding"] = "utf-8" + + with _open_file_like(f, "rb") as opened_file: + if _is_zipfile(opened_file): + # The zipfile reader is going to advance the current file position. + # If we want to actually tail call to torch.jit.load, we need to + # reset back to the original position. + orig_position = opened_file.tell() + overall_storage = None + with _open_zipfile_reader(opened_file) as opened_zipfile: + if _is_torchscript_zip(opened_zipfile): + warnings.warn( + "'torch.load' received a zip file that looks like a TorchScript archive" + " dispatching to 'torch.jit.load' (call 'torch.jit.load' directly to" + " silence this warning)", + UserWarning, + stacklevel=2, + ) + if weights_only: + raise RuntimeError( + "Cannot use ``weights_only=True`` with TorchScript archives passed to " + "``torch.load``. " + UNSAFE_MESSAGE + ) + opened_file.seek(orig_position) + return torch.jit.load(opened_file, map_location=map_location) + if mmap: + if not _is_path(f): + raise ValueError( + "f must be a file path in order to use the mmap argument" + ) + size = os.path.getsize(f) + if not IS_WINDOWS: + shared = get_default_mmap_options() == MAP_SHARED + else: + shared = False + overall_storage = torch.UntypedStorage.from_file( + os.fspath(f), + shared, + size, + ) + if weights_only: + try: + return _load( + opened_zipfile, + map_location, + _weights_only_unpickler, + overall_storage=overall_storage, + **pickle_load_args, + ) + except pickle.UnpicklingError as e: + raise pickle.UnpicklingError(_get_wo_message(str(e))) from None + return _load( + opened_zipfile, + map_location, + pickle_module, + overall_storage=overall_storage, + **pickle_load_args, + ) + if mmap: + f_name = "" if not isinstance(f, str) else f"{f}, " + raise RuntimeError( + "mmap can only be used with files saved with " + f"`torch.save({f_name}_use_new_zipfile_serialization=True), " + "please torch.save your checkpoint with this option in order to use mmap." + ) + if weights_only: + try: + return _legacy_load( + opened_file, + map_location, + _weights_only_unpickler, + **pickle_load_args, + ) + except pickle.UnpicklingError as e: + raise pickle.UnpicklingError(_get_wo_message(str(e))) from None + return _legacy_load( + opened_file, map_location, pickle_module, **pickle_load_args + ) + + +# Register pickling support for layout instances such as +# torch.sparse_coo, etc +def _get_layout(name): + """Get layout extension object from its string representation.""" + cache = _get_layout.cache # type: ignore[attr-defined] + if not cache: + for v in torch.__dict__.values(): + if isinstance(v, torch.layout): + cache[str(v)] = v + return cache[name] + + +# There are yet not good way to type annotate function attributes https://github.com/python/mypy/issues/2087 +_get_layout.cache = {} # type: ignore[attr-defined] +copyreg.pickle(torch.layout, lambda obj: (_get_layout, (str(obj),))) + + +def _legacy_load(f, map_location, pickle_module, **pickle_load_args): + deserialized_objects: dict[int, Any] = {} + + restore_location = _get_restore_location(map_location) + + class UnpicklerWrapper(pickle_module.Unpickler): # type: ignore[name-defined] + def find_class(self, mod_name, name): + if type(name) is str and "Storage" in name: + try: + return StorageType(name) + except KeyError: + pass + return super().find_class(mod_name, name) + + def _check_container_source(container_type, source_file, original_source): + try: + current_source = "".join(get_source_lines_and_file(container_type)[0]) + except Exception: # saving the source is optional, so we can ignore any errors + warnings.warn( + "Couldn't retrieve source code for container of " + "type " + container_type.__name__ + ". It won't be checked " + "for correctness upon loading.", + stacklevel=2, + ) + return + if original_source != current_source: + if container_type.dump_patches: + file_name = container_type.__name__ + ".patch" + diff = difflib.unified_diff( + current_source.split("\n"), + original_source.split("\n"), + source_file, + source_file, + lineterm="", + ) + lines = "\n".join(diff) + try: + with open(file_name, "a+") as f: + file_size = f.seek(0, 2) + f.seek(0) + if file_size == 0: + f.write(lines) + elif file_size != len(lines) or f.read() != lines: + raise OSError + msg = ( + "Saved a reverse patch to " + file_name + ". " + "Run `patch -p0 < " + file_name + "` to revert your " + "changes." + ) + except OSError: + msg = ( + "Tried to save a patch, but couldn't create a " + "writable file " + file_name + ". Make sure it " + "doesn't exist and your working directory is " + "writable." + ) + else: + msg = ( + "you can retrieve the original source code by " + "accessing the object's source attribute or set " + "`torch.nn.Module.dump_patches = True` and use the " + "patch tool to revert the changes." + ) + msg = f"source code of class '{torch.typename(container_type)}' has changed. {msg}" + warnings.warn(msg, SourceChangeWarning, stacklevel=2) + + def legacy_load(f): + deserialized_objects: dict[int, Any] = {} + + def persistent_load(saved_id): + if isinstance(saved_id, tuple): + # Ignore containers that don't have any sources saved + if all(saved_id[1:]): + _check_container_source(*saved_id) + return saved_id[0] + return deserialized_objects[int(saved_id)] + + with ( + closing( + tarfile.open(fileobj=f, mode="r:", format=tarfile.PAX_FORMAT) + ) as tar, + mkdtemp() as tmpdir, + ): + if pickle_module is _weights_only_unpickler: + raise RuntimeError( + "Cannot use ``weights_only=True`` with files saved in the " + "legacy .tar format. " + UNSAFE_MESSAGE + ) + tar.extract("storages", path=tmpdir) + with open(os.path.join(tmpdir, "storages"), "rb", 0) as f: + num_storages = pickle_module.load(f, **pickle_load_args) + for _ in range(num_storages): + args = pickle_module.load(f, **pickle_load_args) + key, location, storage_type = args + dtype = storage_type._dtype + obj = cast(Storage, torch.UntypedStorage)._new_with_file( + f, torch._utils._element_size(dtype) + ) + obj = restore_location(obj, location) + # TODO: Once we decide to break serialization FC, we can + # stop wrapping with TypedStorage + deserialized_objects[key] = torch.storage.TypedStorage( + wrap_storage=obj, dtype=dtype, _internal=True + ) + + storage_views = pickle_module.load(f, **pickle_load_args) + for target_cdata, root_cdata, offset, numel in storage_views: + root = deserialized_objects[root_cdata] + element_size = torch._utils._element_size(root.dtype) + offset_bytes = offset * element_size + # TODO: Once we decide to break serialization FC, we can + # stop wrapping with TypedStorage + deserialized_objects[target_cdata] = torch.storage.TypedStorage( + wrap_storage=root._untyped_storage[ + offset_bytes : offset_bytes + numel * element_size + ], + dtype=root.dtype, + _internal=True, + ) + + tar.extract("tensors", path=tmpdir) + with open(os.path.join(tmpdir, "tensors"), "rb", 0) as f: + num_tensors = pickle_module.load(f, **pickle_load_args) + for _ in range(num_tensors): + args = pickle_module.load(f, **pickle_load_args) + key, storage_id, _original_tensor_type = args + storage = deserialized_objects[storage_id] + (ndim,) = struct.unpack(" str: + # When using encoding='bytes' in Py3, some **internal** keys stored as + # strings in Py2 are loaded as bytes. This function decodes them with + # ascii encoding, one that Py3 uses by default. + # + # NOTE: This should only be used on internal keys (e.g., `typename` and + # `location` in `persistent_load` below! + if isinstance(bytes_str, bytes): + return bytes_str.decode("ascii") + return bytes_str + + +def _get_restore_location(map_location): + if map_location is None: + restore_location = default_restore_location + elif isinstance(map_location, dict): + + def restore_location(storage, location): + location = map_location.get(location, location) + return default_restore_location(storage, location) + + elif isinstance(map_location, (str, bytes)): + + def restore_location(storage, location): + return default_restore_location(storage, map_location) + + elif isinstance(map_location, torch.device): + + def restore_location(storage, location): + return default_restore_location(storage, str(map_location)) + + else: + + def restore_location(storage, location): + result = map_location(storage, location) + if result is None: + result = default_restore_location(storage, location) + return result + + return restore_location + + +class StorageType: + def __init__(self, name): + self._dtype = _get_dtype_from_pickle_storage_type(name) + + @property + def dtype(self): + return self._dtype + + def __str__(self): + return f"StorageType(dtype={self.dtype})" + + +def _load( + zip_file, + map_location, + pickle_module, + pickle_file="data.pkl", + overall_storage=None, + **pickle_load_args, +): + restore_location = _get_restore_location(map_location) + + loaded_storages = {} + + is_meta_map_location = _is_meta_location(map_location) + + can_calculate_storage_offsets = False + if zip_file.has_record(".format_version"): + version = zip_file.get_record(".format_version") + can_calculate_storage_offsets = version >= b"1" + + # check if byteswapping is needed + byteordername = "byteorder" + byteorderdata = None + if zip_file.has_record(byteordername): + byteorderdata = zip_file.get_record(byteordername) + if byteorderdata not in [b"little", b"big"]: + raise ValueError("Unknown endianness type: " + byteorderdata.decode()) + elif ( + get_default_load_endianness() == LoadEndianness.LITTLE + or get_default_load_endianness() is None + ): + byteorderdata = b"little" + elif get_default_load_endianness() == LoadEndianness.BIG: + byteorderdata = b"big" + elif get_default_load_endianness() == LoadEndianness.NATIVE: + pass + else: + raise ValueError("Invalid load endianness type") + + storage_alignment = 64 + if zip_file.has_record(".storage_alignment"): + storage_alignment = int(zip_file.get_record(".storage_alignment")) + + if ( + not zip_file.has_record(byteordername) + and get_default_load_endianness() is None + and sys.byteorder == "big" + ): + # Default behaviour was changed + # See https://github.com/pytorch/pytorch/issues/101688 + warnings.warn( + "The default load endianness for checkpoints without a byteorder mark " + "on big endian machines was changed from 'native' to 'little' endian, " + "to avoid this behavior please use " + "torch.serialization.set_default_load_endianness to set " + "the desired default load endianness", + UserWarning, + stacklevel=2, + ) + + from torch.utils.serialization import config + + calculate_storage_offsets = config.load.calculate_storage_offsets + run_debug_asserts = os.environ.get("TORCH_SERIALIZATION_DEBUG", "0") == "1" + current_offset = None + # constants from miniz.h/miniz.c + data_descripter_size64 = 24 + data_descripter_size32 = 16 + mz_uint32_max = 0xFFFFFFFF + offsets: dict[str, int] = dict() + + def _get_offset(key, name, numel): + """ + Return the offset of the storage associated with key with record name `name` and size numel. + It is expected that the zipfile header of this storage starts at current_offset. + + WARNING: This function relies on the behavior of the zipwriter in miniz.c. In particular, + the behavior of `mz_zip_writer_add_mem_ex_v2`. The behavior of this function must be kept + in sync with that of miniz! + + After reading a storage of size numel that starts at storage_offset + if it is the first time that storage was read, update nonlocal variable + current_offset to the start of the next zipfile header by incrementing + it by numel and the data descriptor size. + """ + nonlocal current_offset, offsets + if name in offsets: + storage_offset = offsets[name] + return storage_offset + + if current_offset is None: + if key != "0": + raise AssertionError(f"expected key '0', got {key!r}") + current_offset = zip_file.get_record_offset(name) + local_header_offset = zip_file.get_record_header_offset(name) + storage_offset = current_offset + else: + storage_offset = zip_file.get_record_offset_no_read( + current_offset, name, numel, storage_alignment + ) + local_header_offset = current_offset + + # This is only actually needed for storages that have typed_storage._data_ptr() == 0 + # after being read. Otherwise persistent_load would never "re-call" load_tensor + # for a given key. + offsets[name] = storage_offset + + # Increment current_offset to offset where next zipfile header starts + current_offset = storage_offset + numel + # add size of data descriptor after payload + if numel > 0: + if local_header_offset >= mz_uint32_max or numel >= mz_uint32_max: + current_offset += data_descripter_size64 + else: + current_offset += data_descripter_size32 + + return storage_offset + + def load_tensor(dtype, nbytes, key, location): + name = f"data/{key}" + if torch._guards.detect_fake_mode(None) is not None or is_meta_map_location: + storage = torch.UntypedStorage(nbytes, device="meta") + if can_calculate_storage_offsets: + storage._checkpoint_offset = _get_offset(key, name, nbytes) + else: + storage._checkpoint_offset = zip_file.get_record_offset(name) + elif _serialization_tls.skip_data: + storage = torch.UntypedStorage(nbytes) + elif overall_storage is not None: + if can_calculate_storage_offsets and calculate_storage_offsets: + storage_offset = _get_offset(key, name, nbytes) + if run_debug_asserts: + if storage_offset != zip_file.get_record_offset(name): + raise RuntimeError( + "This is a debug assert that was run as the `TORCH_SERIALIZATION_DEBUG` environment " + f"variable was set: Incorrect offset for {name}, got {storage_offset} expected " + f"{zip_file.get_record_offset(name)}" + ) + else: + storage_offset = zip_file.get_record_offset(name) + storage = overall_storage[storage_offset : storage_offset + nbytes] + else: + if can_calculate_storage_offsets and run_debug_asserts: + # This is debug code that we use to test the validity of + # torch.utils.serialization.config.load.calculate_storage_offsets throughout CI + storage_offset = _get_offset(key, name, nbytes) + if storage_offset != zip_file.get_record_offset(name): + raise RuntimeError( + "This is a debug assert that was run as the `TORCH_SERIALIZATION_DEBUG` environment " + f"variable was set: Incorrect offset for {name}, got {storage_offset} expected " + f"{zip_file.get_record_offset(name)}" + ) + storage = ( + zip_file.get_storage_from_record(name, nbytes, torch.UntypedStorage) + ._typed_storage() + ._untyped_storage + ) + # swap here if byteswapping is needed + if byteorderdata is not None: + if byteorderdata.decode() != sys.byteorder: + storage.byteswap(dtype) + + # TODO: Once we decide to break serialization FC, we can + # stop wrapping with TypedStorage + + if is_meta_map_location: + # Skip restore_location for meta map_location. Since we already created + # a meta storage above, calling restore_location would just redundantly + # call _meta_deserialize which creates another meta storage with the same + # size. + wrap_storage = storage + elif torch._guards.detect_fake_mode(None) is None: + wrap_storage = restore_location(storage, location) + else: + storage._fake_device = location + wrap_storage = storage + + typed_storage = torch.storage.TypedStorage( + wrap_storage=wrap_storage, + dtype=dtype, + _internal=True, + ) + + if typed_storage._data_ptr() != 0: + loaded_storages[key] = typed_storage + + return typed_storage + + def persistent_load(saved_id): + if not isinstance(saved_id, tuple): + raise AssertionError( + f"saved_id must be a tuple, got {type(saved_id).__name__}" + ) + typename = _maybe_decode_ascii(saved_id[0]) + data = saved_id[1:] + + if typename != "storage": + raise AssertionError( + f"Unknown typename for persistent_load, expected 'storage' but got '{typename}'" + ) + storage_type, key, location, numel = data + if storage_type is torch.UntypedStorage: + dtype = torch.uint8 + else: + dtype = storage_type.dtype + + if key in loaded_storages: + typed_storage = loaded_storages[key] + else: + nbytes = numel * torch._utils._element_size(dtype) + typed_storage = load_tensor( + dtype, nbytes, key, _maybe_decode_ascii(location) + ) + + return typed_storage + + load_module_mapping: dict[str, str] = { + # See https://github.com/pytorch/pytorch/pull/51633 + "torch.tensor": "torch._tensor" + } + + # Need to subclass Unpickler instead of directly monkey-patching the find_class method + # because it's marked readonly in pickle. + # The type: ignore is because mypy can't statically determine the type of this class. + class UnpicklerWrapper(pickle_module.Unpickler): # type: ignore[name-defined] + # from https://stackoverflow.com/questions/13398462/unpickling-python-objects-with-a-changed-module-path/13405732 + # Lets us override the imports that pickle uses when unpickling an object. + # This is useful for maintaining BC if we change a module path that tensor instantiation relies on. + def find_class(self, mod_name, name): + if type(name) is str and "Storage" in name: + try: + return StorageType(name) + except KeyError: + pass + mod_name = load_module_mapping.get(mod_name, mod_name) + return super().find_class(mod_name, name) + + # Load the data (which may in turn use `persistent_load` to load tensors) + data_file = io.BytesIO(zip_file.get_record(pickle_file)) + + unpickler = UnpicklerWrapper(data_file, **pickle_load_args) + unpickler.persistent_load = persistent_load + # Needed for tensors where storage device and rebuild tensor device are + # not connected (wrapper subclasses and tensors rebuilt using numpy) + global _serialization_tls + _serialization_tls.map_location = map_location + result = unpickler.load() + _serialization_tls.map_location = None + + torch._utils._validate_loaded_sparse_tensors() + torch._C._log_api_usage_metadata( + "torch.load.metadata", {"serialization_id": zip_file.serialization_id()} + ) + return result + + +def _is_torchscript_zip(zip_file): + return "constants.pkl" in zip_file.get_all_records() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/ATen/ATenConfig.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/ATen/ATenConfig.cmake new file mode 100644 index 0000000000000000000000000000000000000000..a5d6e54343cb894053a78d02458df06ac39e7ff2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/ATen/ATenConfig.cmake @@ -0,0 +1,9 @@ +# Find the TH includes and library +# +# ATEN_INCLUDE_DIR -- where to find the includes +# ATEN_LIBRARIES -- list of libraries to link against +# ATEN_FOUND -- set to 1 if found + +set(ATEN_FOUND 1) +set(ATEN_INCLUDE_DIR "/tmp/build_wheels_tmp.3712/python-3.11/torch/torch/include") +set(ATEN_LIBRARIES "") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Caffe2Config.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Caffe2Config.cmake new file mode 100644 index 0000000000000000000000000000000000000000..a6bf8473c02fdd160f4814a83697d92cb74730bb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Caffe2Config.cmake @@ -0,0 +1,140 @@ +# - Config file for the Caffe2 package +# It defines the following variable(s) +# CAFFE2_INCLUDE_DIRS - include directories for FooBar +# as well as Caffe2 targets for other cmake libraries to use. + +# library version information + +# Utils functions. +include("${CMAKE_CURRENT_LIST_DIR}/public/utils.cmake") + +# Depending on whether Caffe2 uses gflags during compile time or +# not, invoke gflags. +if(OFF) + include("${CMAKE_CURRENT_LIST_DIR}/public/gflags.cmake") + if(NOT TARGET gflags) + message(FATAL_ERROR + "Your installed Caffe2 version uses gflags but the gflags library " + "cannot be found. Did you accidentally remove it, or have you set " + "the right CMAKE_PREFIX_PATH and/or GFLAGS_ROOT_DIR? If you do not " + "have gflags, you will need to install gflags and set the library " + "path accordingly.") + endif() +endif() + +# Depending on whether Caffe2 uses glog during compile time or +# not, invoke glog. +if(OFF) + include("${CMAKE_CURRENT_LIST_DIR}/public/glog.cmake") + if(NOT TARGET glog::glog) + message(FATAL_ERROR + "Your installed Caffe2 version uses glog but the glog library " + "cannot be found. Did you accidentally remove it, or have you set " + "the right CMAKE_PREFIX_PATH and/or GFLAGS_ROOT_DIR? If you do not " + "have glog, you will need to install glog and set the library " + "path accordingly.") + endif() +endif() + +# Protobuf +if(OFF) + if(NOT TARGET protobuf::libprotobuf) + # Define protobuf::libprotobuf as a dummy target to resolve references to + # protobuf::libprotobuf in Caffe2Targets.cmake. + add_library(dummy INTERFACE) + add_library(protobuf::libprotobuf ALIAS dummy) + endif() +else() + include("${CMAKE_CURRENT_LIST_DIR}/public/protobuf.cmake") + if(NOT TARGET protobuf::libprotobuf) + message(FATAL_ERROR + "Your installed Caffe2 version uses protobuf but the protobuf library " + "cannot be found. Did you accidentally remove it, or have you set " + "the right CMAKE_PREFIX_PATH? If you do not have protobuf, you will " + "need to install protobuf and set the library path accordingly.") + endif() + message(STATUS "Caffe2: Protobuf version " ${Protobuf_VERSION}) + # If during build time we know the protobuf version, we will also do a sanity + # check to ensure that the protobuf library that Caffe2 found is consistent + # with the compiled version. + if(FALSE) + if(NOT (${Protobuf_VERSION} VERSION_EQUAL Protobuf_VERSION_NOTFOUND)) + message(FATAL_ERROR + "Your installed Caffe2 is built with protobuf " + "Protobuf_VERSION_NOTFOUND" + ", while your current cmake setting discovers protobuf version " + ${Protobuf_VERSION} + ". Please specify a protobuf version that is the same as the built " + "version.") + endif() + endif() +endif() + +if (OFF) + include("${CMAKE_CURRENT_LIST_DIR}/public/LoadHIP.cmake") +endif() + +if(ON) + # The file public/cuda.cmake exclusively uses CAFFE2_USE_*. + # If Caffe2 was compiled with the libraries below, they must + # be found again when including the Caffe2 target. + set(CAFFE2_USE_CUDA ON) + + # Add current directory to module path so we pick up FindCUDAToolkit.cmake + set(old_CMAKE_MODULE_PATH "${CMAKE_MODULE_PATH}") + list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}") + include("${CMAKE_CURRENT_LIST_DIR}/public/cuda.cmake") + set(CMAKE_MODULE_PATH "${old_CMAKE_MODULE_PATH}") + + if(ON AND NOT CAFFE2_USE_CUDA) + message(FATAL_ERROR + "Your installed Caffe2 version uses CUDA but I cannot find the CUDA " + "libraries. Please set the proper CUDA prefixes and / or install " + "CUDA.") + endif() +endif() + +if(OFF) + # Add current directory to module path so we pick up FindSYCLToolkit.cmake + set(old_CMAKE_MODULE_PATH "${CMAKE_MODULE_PATH}") + list(APPEND CMAKE_MODULE_PATH "${CMAKE_CURRENT_LIST_DIR}") + include("${CMAKE_CURRENT_LIST_DIR}/public/xpu.cmake") + set(CMAKE_MODULE_PATH "${old_CMAKE_MODULE_PATH}") + + if(OFF AND NOT PYTORCH_FOUND_XPU) + message(FATAL_ERROR + "Your installed Caffe2 version uses XPU but I cannot find the XPU runtime" + "libraries. Please set the proper oneAPI paths and / or install " + "oneAPI.") + endif() +endif() + +if() + include("${CMAKE_CURRENT_LIST_DIR}/public/mkl.cmake") +endif() + +if(ON) + include("${CMAKE_CURRENT_LIST_DIR}/public/mkldnn.cmake") +endif() + +# import targets +include ("${CMAKE_CURRENT_LIST_DIR}/Caffe2Targets.cmake") + +# Interface libraries, that allows one to build proper link flags. +# We will also define a helper variable, Caffe2_MAIN_LIBS, that resolves to +# the main caffe2 libraries in cases of cuda presence / absence. +set(Caffe2_MAIN_LIBS torch_library) + +# include directory. +# +# Newer versions of CMake set the INTERFACE_INCLUDE_DIRECTORIES property +# of the imported targets. It is hence not necessary to add this path +# manually to the include search path for targets which link to gflags. +# The following lines are here for backward compatibility, in case one +# would like to use the old-style include path. +get_filename_component( + CMAKE_CURRENT_LIST_DIR "${CMAKE_CURRENT_LIST_FILE}" PATH) +# Note: the current list dir is _INSTALL_PREFIX/share/cmake/Gloo. +get_filename_component( + _INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/../../../" ABSOLUTE) +set(CAFFE2_INCLUDE_DIRS "${_INSTALL_PREFIX}/include") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Caffe2Targets-release.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Caffe2Targets-release.cmake new file mode 100644 index 0000000000000000000000000000000000000000..95b59eba00319cc797d167be055b0a488438afc1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Caffe2Targets-release.cmake @@ -0,0 +1,60 @@ +#---------------------------------------------------------------- +# Generated CMake target import file for configuration "Release". +#---------------------------------------------------------------- + +# Commands may need to know the format version. +set(CMAKE_IMPORT_FILE_VERSION 1) + +# Import target "c10_cuda" for configuration "Release" +set_property(TARGET c10_cuda APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE) +set_target_properties(c10_cuda PROPERTIES + IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib/libc10_cuda.so" + IMPORTED_SONAME_RELEASE "libc10_cuda.so" + ) + +list(APPEND _cmake_import_check_targets c10_cuda ) +list(APPEND _cmake_import_check_files_for_c10_cuda "${_IMPORT_PREFIX}/lib/libc10_cuda.so" ) + +# Import target "c10" for configuration "Release" +set_property(TARGET c10 APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE) +set_target_properties(c10 PROPERTIES + IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib/libc10.so" + IMPORTED_SONAME_RELEASE "libc10.so" + ) + +list(APPEND _cmake_import_check_targets c10 ) +list(APPEND _cmake_import_check_files_for_c10 "${_IMPORT_PREFIX}/lib/libc10.so" ) + +# Import target "torch_cpu" for configuration "Release" +set_property(TARGET torch_cpu APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE) +set_target_properties(torch_cpu PROPERTIES + IMPORTED_LINK_DEPENDENT_LIBRARIES_RELEASE "XNNPACK;microkernels-prod" + IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib/libtorch_cpu.so" + IMPORTED_SONAME_RELEASE "libtorch_cpu.so" + ) + +list(APPEND _cmake_import_check_targets torch_cpu ) +list(APPEND _cmake_import_check_files_for_torch_cpu "${_IMPORT_PREFIX}/lib/libtorch_cpu.so" ) + +# Import target "torch_cuda" for configuration "Release" +set_property(TARGET torch_cuda APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE) +set_target_properties(torch_cuda PROPERTIES + IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib/libtorch_cuda.so" + IMPORTED_SONAME_RELEASE "libtorch_cuda.so" + ) + +list(APPEND _cmake_import_check_targets torch_cuda ) +list(APPEND _cmake_import_check_files_for_torch_cuda "${_IMPORT_PREFIX}/lib/libtorch_cuda.so" ) + +# Import target "torch" for configuration "Release" +set_property(TARGET torch APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE) +set_target_properties(torch PROPERTIES + IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib/libtorch.so" + IMPORTED_SONAME_RELEASE "libtorch.so" + ) + +list(APPEND _cmake_import_check_targets torch ) +list(APPEND _cmake_import_check_files_for_torch "${_IMPORT_PREFIX}/lib/libtorch.so" ) + +# Commands beyond this point should not need to know the version. +set(CMAKE_IMPORT_FILE_VERSION) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Caffe2Targets.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Caffe2Targets.cmake new file mode 100644 index 0000000000000000000000000000000000000000..60e9fc0a1fcb4a0e9db80dcfc247d701e3ba14ca --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Caffe2Targets.cmake @@ -0,0 +1,177 @@ +# Generated by CMake + +if("${CMAKE_MAJOR_VERSION}.${CMAKE_MINOR_VERSION}" LESS 2.8) + message(FATAL_ERROR "CMake >= 3.0.0 required") +endif() +if(CMAKE_VERSION VERSION_LESS "3.0.0") + message(FATAL_ERROR "CMake >= 3.0.0 required") +endif() +cmake_policy(PUSH) +cmake_policy(VERSION 3.0.0...3.29) +#---------------------------------------------------------------- +# Generated CMake target import file. +#---------------------------------------------------------------- + +# Commands may need to know the format version. +set(CMAKE_IMPORT_FILE_VERSION 1) + +# Protect against multiple inclusion, which would fail when already imported targets are added once more. +set(_cmake_targets_defined "") +set(_cmake_targets_not_defined "") +set(_cmake_expected_targets "") +foreach(_cmake_expected_target IN ITEMS headeronly c10_cuda c10 torch_cpu torch_cpu_library torch_cuda torch_cuda_library torch torch_library) + list(APPEND _cmake_expected_targets "${_cmake_expected_target}") + if(TARGET "${_cmake_expected_target}") + list(APPEND _cmake_targets_defined "${_cmake_expected_target}") + else() + list(APPEND _cmake_targets_not_defined "${_cmake_expected_target}") + endif() +endforeach() +unset(_cmake_expected_target) +if(_cmake_targets_defined STREQUAL _cmake_expected_targets) + unset(_cmake_targets_defined) + unset(_cmake_targets_not_defined) + unset(_cmake_expected_targets) + unset(CMAKE_IMPORT_FILE_VERSION) + cmake_policy(POP) + return() +endif() +if(NOT _cmake_targets_defined STREQUAL "") + string(REPLACE ";" ", " _cmake_targets_defined_text "${_cmake_targets_defined}") + string(REPLACE ";" ", " _cmake_targets_not_defined_text "${_cmake_targets_not_defined}") + message(FATAL_ERROR "Some (but not all) targets in this export set were already defined.\nTargets Defined: ${_cmake_targets_defined_text}\nTargets not yet defined: ${_cmake_targets_not_defined_text}\n") +endif() +unset(_cmake_targets_defined) +unset(_cmake_targets_not_defined) +unset(_cmake_expected_targets) + + +# Compute the installation prefix relative to this file. +get_filename_component(_IMPORT_PREFIX "${CMAKE_CURRENT_LIST_FILE}" PATH) +get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH) +get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH) +get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH) +if(_IMPORT_PREFIX STREQUAL "/") + set(_IMPORT_PREFIX "") +endif() + +# Create imported target headeronly +add_library(headeronly INTERFACE IMPORTED) + +# Create imported target c10_cuda +add_library(c10_cuda SHARED IMPORTED) + +set_target_properties(c10_cuda PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${_IMPORT_PREFIX}/include" + INTERFACE_LINK_LIBRARIES "c10;torch::cudart" +) + +# Create imported target c10 +add_library(c10 SHARED IMPORTED) + +set_target_properties(c10 PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${_IMPORT_PREFIX}/include" + INTERFACE_LINK_LIBRARIES "headeronly" +) + +# Create imported target torch_cpu +add_library(torch_cpu SHARED IMPORTED) + +set_target_properties(torch_cpu PROPERTIES + INTERFACE_COMPILE_DEFINITIONS "USE_DISTRIBUTED;USE_C10D_GLOO;USE_C10D_MPI;USE_RPC;USE_TENSORPIPE" + INTERFACE_INCLUDE_DIRECTORIES "${_IMPORT_PREFIX}/include" + INTERFACE_LINK_LIBRARIES "protobuf::libprotobuf;c10" +) + +# Create imported target torch_cpu_library +add_library(torch_cpu_library INTERFACE IMPORTED) + +set_target_properties(torch_cpu_library PROPERTIES + INTERFACE_COMPILE_DEFINITIONS "\$" + INTERFACE_COMPILE_OPTIONS "\$" + INTERFACE_INCLUDE_DIRECTORIES "\$" + INTERFACE_LINK_LIBRARIES "-Wl,--no-as-needed,\"\$\" -Wl,--as-needed;\$" + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "\$" +) + +# Create imported target torch_cuda +add_library(torch_cuda SHARED IMPORTED) + +set_target_properties(torch_cuda PROPERTIES + INTERFACE_COMPILE_DEFINITIONS "USE_C10D_NCCL" + INTERFACE_INCLUDE_DIRECTORIES "${_IMPORT_PREFIX}/include;${_IMPORT_PREFIX}/include" + INTERFACE_LINK_LIBRARIES "torch::cudart;c10_cuda;torch_cpu_library" + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "include" +) + +# Create imported target torch_cuda_library +add_library(torch_cuda_library INTERFACE IMPORTED) + +set_target_properties(torch_cuda_library PROPERTIES + INTERFACE_COMPILE_DEFINITIONS "\$" + INTERFACE_COMPILE_OPTIONS "\$" + INTERFACE_INCLUDE_DIRECTORIES "\$" + INTERFACE_LINK_LIBRARIES "-Wl,--no-as-needed,\"\$\" -Wl,--as-needed;\$" + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "\$" +) + +# Create imported target torch +add_library(torch SHARED IMPORTED) + +set_target_properties(torch PROPERTIES + INTERFACE_LINK_LIBRARIES "torch_cpu_library;torch_cuda_library" +) + +# Create imported target torch_library +add_library(torch_library INTERFACE IMPORTED) + +set_target_properties(torch_library PROPERTIES + INTERFACE_COMPILE_DEFINITIONS "\$" + INTERFACE_COMPILE_OPTIONS "\$" + INTERFACE_INCLUDE_DIRECTORIES "\$" + INTERFACE_LINK_LIBRARIES "-Wl,--no-as-needed,\"\$\" -Wl,--as-needed;\$" + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "\$" +) + +# Load information for each installed configuration. +file(GLOB _cmake_config_files "${CMAKE_CURRENT_LIST_DIR}/Caffe2Targets-*.cmake") +foreach(_cmake_config_file IN LISTS _cmake_config_files) + include("${_cmake_config_file}") +endforeach() +unset(_cmake_config_file) +unset(_cmake_config_files) + +# Cleanup temporary variables. +set(_IMPORT_PREFIX) + +# Loop over all imported files and verify that they actually exist +foreach(_cmake_target IN LISTS _cmake_import_check_targets) + if(CMAKE_VERSION VERSION_LESS "3.28" + OR NOT DEFINED _cmake_import_check_xcframework_for_${_cmake_target} + OR NOT IS_DIRECTORY "${_cmake_import_check_xcframework_for_${_cmake_target}}") + foreach(_cmake_file IN LISTS "_cmake_import_check_files_for_${_cmake_target}") + if(NOT EXISTS "${_cmake_file}") + message(FATAL_ERROR "The imported target \"${_cmake_target}\" references the file + \"${_cmake_file}\" +but this file does not exist. Possible reasons include: +* The file was deleted, renamed, or moved to another location. +* An install or uninstall procedure did not complete successfully. +* The installation package was faulty and contained + \"${CMAKE_CURRENT_LIST_FILE}\" +but not all the files it references. +") + endif() + endforeach() + endif() + unset(_cmake_file) + unset("_cmake_import_check_files_for_${_cmake_target}") +endforeach() +unset(_cmake_target) +unset(_cmake_import_check_targets) + +# This file does not depend on other imported targets which have +# been exported from the same project but in a separate export set. + +# Commands beyond this point should not need to know the version. +set(CMAKE_IMPORT_FILE_VERSION) +cmake_policy(POP) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindCUDAToolkit.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindCUDAToolkit.cmake new file mode 100644 index 0000000000000000000000000000000000000000..ec9ae530aa6b2bdceb87f966e706fb5c2a36349a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindCUDAToolkit.cmake @@ -0,0 +1,1081 @@ + +# This module is back-ported from CMake 3.17 and above to work with CMake 3.10 + +# Distributed under the OSI-approved BSD 3-Clause License. See accompanying +# file Copyright.txt or https://cmake.org/licensing for details. + +#[=======================================================================[.rst: +FindCUDAToolkit +--------------- + +.. versionadded:: 3.17 + +This script locates the NVIDIA CUDA toolkit and the associated libraries, but +does not require the ``CUDA`` language be enabled for a given project. This +module does not search for the NVIDIA CUDA Samples. + +.. versionadded:: 3.19 + QNX support. + +Search Behavior +^^^^^^^^^^^^^^^ + +The CUDA Toolkit search behavior uses the following order: + +1. If the ``CUDA`` language has been enabled we will use the directory + containing the compiler as the first search location for ``nvcc``. + +2. If the ``CUDAToolkit_ROOT`` cmake configuration variable (e.g., + ``-DCUDAToolkit_ROOT=/some/path``) *or* environment variable is defined, it + will be searched. If both an environment variable **and** a + configuration variable are specified, the *configuration* variable takes + precedence. + + The directory specified here must be such that the executable ``nvcc`` or + the appropriate ``version.txt`` file can be found underneath the specified + directory. + +3. If the CUDA_PATH environment variable is defined, it will be searched + for ``nvcc``. + +4. The user's path is searched for ``nvcc`` using :command:`find_program`. If + this is found, no subsequent search attempts are performed. Users are + responsible for ensuring that the first ``nvcc`` to show up in the path is + the desired path in the event that multiple CUDA Toolkits are installed. + +5. On Unix systems, if the symbolic link ``/usr/local/cuda`` exists, this is + used. No subsequent search attempts are performed. No default symbolic link + location exists for the Windows platform. + +6. The platform specific default install locations are searched. If exactly one + candidate is found, this is used. The default CUDA Toolkit install locations + searched are: + + +-------------+-------------------------------------------------------------+ + | Platform | Search Pattern | + +=============+=============================================================+ + | macOS | ``/Developer/NVIDIA/CUDA-X.Y`` | + +-------------+-------------------------------------------------------------+ + | Other Unix | ``/usr/local/cuda-X.Y`` | + +-------------+-------------------------------------------------------------+ + | Windows | ``C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vX.Y`` | + +-------------+-------------------------------------------------------------+ + + Where ``X.Y`` would be a specific version of the CUDA Toolkit, such as + ``/usr/local/cuda-9.0`` or + ``C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v9.0`` + + .. note:: + + When multiple CUDA Toolkits are installed in the default location of a + system(e.g., both ``/usr/local/cuda-9.0`` and ``/usr/local/cuda-10.0`` + exist but the ``/usr/local/cuda`` symbolic link does **not** exist), this + package is marked as **not** found. + + There are too many factors involved in making an automatic decision in + the presence of multiple CUDA Toolkits being installed. In this + situation, users are encouraged to either (1) set ``CUDAToolkit_ROOT`` or + (2) ensure that the correct ``nvcc`` executable shows up in ``$PATH`` for + :command:`find_program` to find. + +Arguments +^^^^^^^^^ + +``[]`` + The ``[]`` argument requests a version with which the package found + should be compatible. See :ref:`find_package version format ` + for more details. + +Options +^^^^^^^ + +``REQUIRED`` + If specified, configuration will error if a suitable CUDA Toolkit is not + found. + +``QUIET`` + If specified, the search for a suitable CUDA Toolkit will not produce any + messages. + +``EXACT`` + If specified, the CUDA Toolkit is considered found only if the exact + ``VERSION`` specified is recovered. + +Imported targets +^^^^^^^^^^^^^^^^ + +An :ref:`imported target ` named ``CUDA::toolkit`` is provided. + +This module defines :prop_tgt:`IMPORTED` targets for each +of the following libraries that are part of the CUDAToolkit: + +- :ref:`CUDA Runtime Library` +- :ref:`CUDA Driver Library` +- :ref:`cuBLAS` +- :ref:`cuFFT` +- :ref:`cuRAND` +- :ref:`cuSOLVER` +- :ref:`cuSPARSE` +- :ref:`cuPTI` +- :ref:`NPP` +- :ref:`nvBLAS` +- :ref:`nvGRAPH` +- :ref:`nvJPEG` +- :ref:`nvidia-ML` +- :ref:`nvRTC` +- :ref:`nvToolsExt` +- :ref:`OpenCL` +- :ref:`cuLIBOS` + +.. _`cuda_toolkit_rt_lib`: + +CUDA Runtime Library +"""""""""""""""""""" + +The CUDA Runtime library (cudart) are what most applications will typically +need to link against to make any calls such as `cudaMalloc`, and `cudaFree`. + +Targets Created: + +- ``CUDA::cudart`` +- ``CUDA::cudart_static`` + +.. _`cuda_toolkit_driver_lib`: + +CUDA Driver Library +"""""""""""""""""""" + +The CUDA Driver library (cuda) are used by applications that use calls +such as `cuMemAlloc`, and `cuMemFree`. + +Targets Created: + +- ``CUDA::cuda_driver`` + +.. _`cuda_toolkit_cuBLAS`: + +cuBLAS +"""""" + +The `cuBLAS `_ library. + +Targets Created: + +- ``CUDA::cublas`` +- ``CUDA::cublas_static`` +- ``CUDA::cublasLt`` starting in CUDA 10.1 +- ``CUDA::cublasLt_static`` starting in CUDA 10.1 + +.. _`cuda_toolkit_cuFFT`: + +cuFFT +""""" + +The `cuFFT `_ library. + +Targets Created: + +- ``CUDA::cufft`` +- ``CUDA::cufftw`` +- ``CUDA::cufft_static`` +- ``CUDA::cufft_static_nocallback`` starting in CUDA 9.2, requires CMake 3.23+ +- ``CUDA::cufftw_static`` + +cuRAND +"""""" + +The `cuRAND `_ library. + +Targets Created: + +- ``CUDA::curand`` +- ``CUDA::curand_static`` + +.. _`cuda_toolkit_cuSOLVER`: + +cuSOLVER +"""""""" + +The `cuSOLVER `_ library. + +Targets Created: + +- ``CUDA::cusolver`` +- ``CUDA::cusolver_static`` + +.. _`cuda_toolkit_cuSPARSE`: + +cuSPARSE +"""""""" + +The `cuSPARSE `_ library. + +Targets Created: + +- ``CUDA::cusparse`` +- ``CUDA::cusparse_static`` + +.. _`cuda_toolkit_cupti`: + +cupti +""""" + +The `NVIDIA CUDA Profiling Tools Interface `_. + +Targets Created: + +- ``CUDA::cupti`` +- ``CUDA::cupti_static`` + +.. _`cuda_toolkit_NPP`: + +NPP +""" + +The `NPP `_ libraries. + +Targets Created: + +- `nppc`: + + - ``CUDA::nppc`` + - ``CUDA::nppc_static`` + +- `nppial`: Arithmetic and logical operation functions in `nppi_arithmetic_and_logical_operations.h` + + - ``CUDA::nppial`` + - ``CUDA::nppial_static`` + +- `nppicc`: Color conversion and sampling functions in `nppi_color_conversion.h` + + - ``CUDA::nppicc`` + - ``CUDA::nppicc_static`` + +- `nppicom`: JPEG compression and decompression functions in `nppi_compression_functions.h` + Removed starting in CUDA 11.0, use :ref:`nvJPEG` instead. + + - ``CUDA::nppicom`` + - ``CUDA::nppicom_static`` + +- `nppidei`: Data exchange and initialization functions in `nppi_data_exchange_and_initialization.h` + + - ``CUDA::nppidei`` + - ``CUDA::nppidei_static`` + +- `nppif`: Filtering and computer vision functions in `nppi_filter_functions.h` + + - ``CUDA::nppif`` + - ``CUDA::nppif_static`` + +- `nppig`: Geometry transformation functions found in `nppi_geometry_transforms.h` + + - ``CUDA::nppig`` + - ``CUDA::nppig_static`` + +- `nppim`: Morphological operation functions found in `nppi_morphological_operations.h` + + - ``CUDA::nppim`` + - ``CUDA::nppim_static`` + +- `nppist`: Statistics and linear transform in `nppi_statistics_functions.h` and `nppi_linear_transforms.h` + + - ``CUDA::nppist`` + - ``CUDA::nppist_static`` + +- `nppisu`: Memory support functions in `nppi_support_functions.h` + + - ``CUDA::nppisu`` + - ``CUDA::nppisu_static`` + +- `nppitc`: Threshold and compare operation functions in `nppi_threshold_and_compare_operations.h` + + - ``CUDA::nppitc`` + - ``CUDA::nppitc_static`` + +- `npps`: + + - ``CUDA::npps`` + - ``CUDA::npps_static`` + +.. _`cuda_toolkit_nvBLAS`: + +nvBLAS +"""""" + +The `nvBLAS `_ libraries. +This is a shared library only. + +Targets Created: + +- ``CUDA::nvblas`` + +.. _`cuda_toolkit_nvGRAPH`: + +nvGRAPH +""""""" + +The `nvGRAPH `_ library. +Removed starting in CUDA 11.0 + +Targets Created: + +- ``CUDA::nvgraph`` +- ``CUDA::nvgraph_static`` + + +.. _`cuda_toolkit_nvJPEG`: + +nvJPEG +"""""" + +The `nvJPEG `_ library. +Introduced in CUDA 10. + +Targets Created: + +- ``CUDA::nvjpeg`` +- ``CUDA::nvjpeg_static`` + +.. _`cuda_toolkit_nvRTC`: + +nvRTC +""""" + +The `nvRTC `_ (Runtime Compilation) library. +This is a shared library only. + +Targets Created: + +- ``CUDA::nvrtc`` + +.. _`cuda_toolkit_nvml`: + +nvidia-ML +""""""""" + +The `NVIDIA Management Library `_. +This is a shared library only. + +Targets Created: + +- ``CUDA::nvml`` + +.. _`cuda_toolkit_nvToolsExt`: + +nvToolsExt +"""""""""" + +The `NVIDIA Tools Extension `_. +This is a shared library only. + +Targets Created: + +- ``CUDA::nvToolsExt`` + +.. _`cuda_toolkit_opencl`: + +OpenCL +"""""" + +The `NVIDIA OpenCL Library `_. +This is a shared library only. + +Targets Created: + +- ``CUDA::OpenCL`` + +.. _`cuda_toolkit_cuLIBOS`: + +cuLIBOS +""""""" + +The cuLIBOS library is a backend thread abstraction layer library which is +static only. The ``CUDA::cublas_static``, ``CUDA::cusparse_static``, +``CUDA::cufft_static``, ``CUDA::curand_static``, and (when implemented) NPP +libraries all automatically have this dependency linked. + +Target Created: + +- ``CUDA::culibos`` + +**Note**: direct usage of this target by consumers should not be necessary. + +.. _`cuda_toolkit_cuRAND`: + + + +Result variables +^^^^^^^^^^^^^^^^ + +``CUDAToolkit_FOUND`` + A boolean specifying whether or not the CUDA Toolkit was found. + +``CUDAToolkit_VERSION`` + The exact version of the CUDA Toolkit found (as reported by + ``nvcc --version`` or ``version.txt``). + +``CUDAToolkit_VERSION_MAJOR`` + The major version of the CUDA Toolkit. + +``CUDAToolkit_VERSION_MINOR`` + The minor version of the CUDA Toolkit. + +``CUDAToolkit_VERSION_PATCH`` + The patch version of the CUDA Toolkit. + +``CUDAToolkit_BIN_DIR`` + The path to the CUDA Toolkit library directory that contains the CUDA + executable ``nvcc``. + +``CUDAToolkit_INCLUDE_DIRS`` + The path to the CUDA Toolkit ``include`` folder containing the header files + required to compile a project linking against CUDA. + +``CUDAToolkit_LIBRARY_DIR`` + The path to the CUDA Toolkit library directory that contains the CUDA + Runtime library ``cudart``. + +``CUDAToolkit_LIBRARY_ROOT`` + .. versionadded:: 3.18 + + The path to the CUDA Toolkit directory containing the nvvm directory and + version.txt. + +``CUDAToolkit_TARGET_DIR`` + The path to the CUDA Toolkit directory including the target architecture + when cross-compiling. When not cross-compiling this will be equivalent to + the parent directory of ``CUDAToolkit_BIN_DIR``. + +``CUDAToolkit_NVCC_EXECUTABLE`` + The path to the NVIDIA CUDA compiler ``nvcc``. Note that this path may + **not** be the same as + :variable:`CMAKE_CUDA_COMPILER _COMPILER>`. ``nvcc`` must be + found to determine the CUDA Toolkit version as well as determining other + features of the Toolkit. This variable is set for the convenience of + modules that depend on this one. + + +#]=======================================================================] + +# NOTE: much of this was simply extracted from FindCUDA.cmake. + +# James Bigler, NVIDIA Corp (nvidia.com - jbigler) +# Abe Stephens, SCI Institute -- http://www.sci.utah.edu/~abe/FindCuda.html +# +# Copyright (c) 2008 - 2009 NVIDIA Corporation. All rights reserved. +# +# Copyright (c) 2007-2009 +# Scientific Computing and Imaging Institute, University of Utah +# +# This code is licensed under the MIT License. See the FindCUDA.cmake script +# for the text of the license. + +# The MIT License +# +# License for the specific language governing rights and limitations under +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included +# in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# +############################################################################### + +# The toolkit is located during compiler detection for CUDA and stored in CMakeCUDACompiler.cmake as +# CMAKE_CUDA_COMPILER_TOOLKIT_ROOT and CMAKE_CUDA_COMPILER_LIBRARY_ROOT. +# We compute the rest based on those here to avoid re-searching and to avoid finding a possibly +# different installation. +if(CMAKE_CUDA_COMPILER_TOOLKIT_ROOT) + set(CUDAToolkit_ROOT_DIR "${CMAKE_CUDA_COMPILER_TOOLKIT_ROOT}") + set(CUDAToolkit_LIBRARY_ROOT "${CMAKE_CUDA_COMPILER_LIBRARY_ROOT}") + set(CUDAToolkit_VERSION "${CMAKE_CUDA_COMPILER_TOOLKIT_VERSION}") + + if(CUDAToolkit_VERSION MATCHES [=[([0-9]+)\.([0-9]+)\.([0-9]+)]=]) + set(CUDAToolkit_VERSION_MAJOR "${CMAKE_MATCH_1}") + set(CUDAToolkit_VERSION_MINOR "${CMAKE_MATCH_2}") + set(CUDAToolkit_VERSION_PATCH "${CMAKE_MATCH_3}") + endif() +else() + function(_CUDAToolkit_find_root_dir ) + cmake_parse_arguments(arg "" "" "SEARCH_PATHS;FIND_FLAGS" ${ARGN}) + + if(NOT CUDAToolkit_BIN_DIR) + if(NOT CUDAToolkit_SENTINEL_FILE) + find_program(CUDAToolkit_NVCC_EXECUTABLE + NAMES nvcc nvcc.exe + PATHS ${arg_SEARCH_PATHS} + ${arg_FIND_FLAGS} + ) + endif() + + if(NOT CUDAToolkit_NVCC_EXECUTABLE) + find_file(CUDAToolkit_SENTINEL_FILE + NAMES version.txt + PATHS ${arg_SEARCH_PATHS} + NO_DEFAULT_PATH + ) + endif() + + if(EXISTS "${CUDAToolkit_NVCC_EXECUTABLE}") + # If NVCC exists then invoke it to find the toolkit location. + # This allows us to support wrapper scripts (e.g. ccache or colornvcc), CUDA Toolkit, + # NVIDIA HPC SDK, and distro's splayed layouts + execute_process(COMMAND ${CUDAToolkit_NVCC_EXECUTABLE} "-v" "__cmake_determine_cuda" + OUTPUT_VARIABLE _CUDA_NVCC_OUT ERROR_VARIABLE _CUDA_NVCC_OUT) + if(_CUDA_NVCC_OUT MATCHES "\\#\\$ TOP=([^\r\n]*)") + get_filename_component(CUDAToolkit_BIN_DIR "${CMAKE_MATCH_1}/bin" ABSOLUTE) + else() + get_filename_component(CUDAToolkit_BIN_DIR "${CUDAToolkit_NVCC_EXECUTABLE}" DIRECTORY) + endif() + unset(_CUDA_NVCC_OUT) + + mark_as_advanced(CUDAToolkit_BIN_DIR) + set(CUDAToolkit_BIN_DIR "${CUDAToolkit_BIN_DIR}" CACHE PATH "" FORCE) + endif() + + if(CUDAToolkit_SENTINEL_FILE) + get_filename_component(CUDAToolkit_BIN_DIR ${CUDAToolkit_SENTINEL_FILE} DIRECTORY ABSOLUTE) + set(CUDAToolkit_BIN_DIR "${CUDAToolkit_BIN_DIR}/bin") + + set(CUDAToolkit_BIN_DIR "${CUDAToolkit_BIN_DIR}" CACHE PATH "" FORCE) + mark_as_advanced(CUDAToolkit_BIN_DIR) + endif() + endif() + + if(CUDAToolkit_BIN_DIR) + get_filename_component(CUDAToolkit_ROOT_DIR ${CUDAToolkit_BIN_DIR} DIRECTORY ABSOLUTE) + set(CUDAToolkit_ROOT_DIR "${CUDAToolkit_ROOT_DIR}" PARENT_SCOPE) + endif() + + endfunction() + + # For NVCC we can easily deduce the SDK binary directory from the compiler path. + if(CMAKE_CUDA_COMPILER_LOADED AND NOT CUDAToolkit_BIN_DIR AND CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA") + get_filename_component(CUDAToolkit_BIN_DIR "${CMAKE_CUDA_COMPILER}" DIRECTORY) + set(CUDAToolkit_BIN_DIR "${CUDAToolkit_BIN_DIR}" CACHE PATH "") + # Try language provided path first. + _CUDAToolkit_find_root_dir(SEARCH_PATHS "${CUDAToolkit_BIN_DIR}" FIND_FLAGS NO_DEFAULT_PATH) + mark_as_advanced(CUDAToolkit_BIN_DIR) + endif() + + # Try user provided path + if(NOT CUDAToolkit_ROOT_DIR AND CUDAToolkit_ROOT) + _CUDAToolkit_find_root_dir(SEARCH_PATHS "${CUDAToolkit_ROOT}" FIND_FLAGS PATH_SUFFIXES bin NO_DEFAULT_PATH) + endif() + if(NOT CUDAToolkit_ROOT_DIR) + _CUDAToolkit_find_root_dir(FIND_FLAGS PATHS ENV CUDA_PATH PATH_SUFFIXES bin) + endif() + + # If the user specified CUDAToolkit_ROOT but the toolkit could not be found, this is an error. + if(NOT CUDAToolkit_ROOT_DIR AND (DEFINED CUDAToolkit_ROOT OR DEFINED ENV{CUDAToolkit_ROOT})) + # Declare error messages now, print later depending on find_package args. + set(fail_base "Could not find nvcc executable in path specified by") + set(cuda_root_fail "${fail_base} CUDAToolkit_ROOT=${CUDAToolkit_ROOT}") + set(env_cuda_root_fail "${fail_base} environment variable CUDAToolkit_ROOT=$ENV{CUDAToolkit_ROOT}") + + if(CUDAToolkit_FIND_REQUIRED) + if(DEFINED CUDAToolkit_ROOT) + message(FATAL_ERROR ${cuda_root_fail}) + elseif(DEFINED ENV{CUDAToolkit_ROOT}) + message(FATAL_ERROR ${env_cuda_root_fail}) + endif() + else() + if(NOT CUDAToolkit_FIND_QUIETLY) + if(DEFINED CUDAToolkit_ROOT) + message(STATUS ${cuda_root_fail}) + elseif(DEFINED ENV{CUDAToolkit_ROOT}) + message(STATUS ${env_cuda_root_fail}) + endif() + endif() + set(CUDAToolkit_FOUND FALSE) + unset(fail_base) + unset(cuda_root_fail) + unset(env_cuda_root_fail) + return() + endif() + endif() + + # CUDAToolkit_ROOT cmake / env variable not specified, try platform defaults. + # + # - Linux: /usr/local/cuda-X.Y + # - macOS: /Developer/NVIDIA/CUDA-X.Y + # - Windows: C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\vX.Y + # + # We will also search the default symlink location /usr/local/cuda first since + # if CUDAToolkit_ROOT is not specified, it is assumed that the symlinked + # directory is the desired location. + if(NOT CUDAToolkit_ROOT_DIR) + if(UNIX) + if(NOT APPLE) + set(platform_base "/usr/local/cuda-") + else() + set(platform_base "/Developer/NVIDIA/CUDA-") + endif() + else() + set(platform_base "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v") + endif() + + # Build out a descending list of possible cuda installations, e.g. + file(GLOB possible_paths "${platform_base}*") + # Iterate the glob results and create a descending list. + set(versions) + foreach(p ${possible_paths}) + # Extract version number from end of string + string(REGEX MATCH "[0-9][0-9]?\\.[0-9]$" p_version ${p}) + if(IS_DIRECTORY ${p} AND p_version) + list(APPEND versions ${p_version}) + endif() + endforeach() + + # Sort numerically in descending order, so we try the newest versions first. + if(CMAKE_VERSION VERSION_GREATER_EQUAL 3.18) + list(SORT versions COMPARE NATURAL ORDER DESCENDING) + elseif(versions) + # Alphabetical sort here is not ideal but better than nothing + list(SORT versions) + list(REVERSE versions) + endif() + + # With a descending list of versions, populate possible paths to search. + set(search_paths) + foreach(v ${versions}) + list(APPEND search_paths "${platform_base}${v}") + endforeach() + + # Force the global default /usr/local/cuda to the front on Unix. + if(UNIX) + list(INSERT search_paths 0 "/usr/local/cuda") + endif() + + # Now search for the toolkit again using the platform default search paths. + _CUDAToolkit_find_root_dir(SEARCH_PATHS "${search_paths}" FIND_FLAGS PATH_SUFFIXES bin) + + # We are done with these variables now, cleanup for caller. + unset(platform_base) + unset(possible_paths) + unset(versions) + unset(search_paths) + + if(NOT CUDAToolkit_ROOT_DIR) + if(CUDAToolkit_FIND_REQUIRED) + message(FATAL_ERROR "Could not find nvcc, please set CUDAToolkit_ROOT.") + elseif(NOT CUDAToolkit_FIND_QUIETLY) + message(STATUS "Could not find nvcc, please set CUDAToolkit_ROOT.") + endif() + + set(CUDAToolkit_FOUND FALSE) + return() + endif() + endif() +endif() + +if(NOT CUDAToolkit_BIN_DIR) + set(CUDAToolkit_BIN_DIR "${CUDAToolkit_ROOT_DIR}/bin") +endif() + +if(NOT CUDAToolkit_NVCC_EXECUTABLE) + set(CUDAToolkit_NVCC_EXECUTABLE "${CUDAToolkit_BIN_DIR}/nvcc${CMAKE_EXECUTABLE_SUFFIX}") +endif() + +if(CMAKE_CUDA_COMPILER_TOOLKIT_VERSION) + set(CUDAToolkit_VERSION "${CMAKE_CUDA_COMPILER_TOOLKIT_VERSION}") +else() + function(_CUDAToolkit_find_version_file result_variable) + # We first check for a non-scattered installation to prefer it over a scattered installation. + if(CUDAToolkit_ROOT AND EXISTS "${CUDAToolkit_ROOT}/version.txt") + set(${result_variable} "${CUDAToolkit_ROOT}/version.txt" PARENT_SCOPE) + elseif(CUDAToolkit_ROOT_DIR AND EXISTS "${CUDAToolkit_ROOT_DIR}/version.txt") + set(${result_variable} "${CUDAToolkit_ROOT_DIR}/version.txt" PARENT_SCOPE) + elseif(CMAKE_SYSROOT_LINK AND EXISTS "${CMAKE_SYSROOT_LINK}/usr/lib/cuda/version.txt") + set(${result_variable} "${CMAKE_SYSROOT_LINK}/usr/lib/cuda/version.txt" PARENT_SCOPE) + elseif(EXISTS "${CMAKE_SYSROOT}/usr/lib/cuda/version.txt") + set(${result_variable} "${CMAKE_SYSROOT}/usr/lib/cuda/version.txt" PARENT_SCOPE) + endif() + endfunction() + + _CUDAToolkit_find_version_file( _CUDAToolkit_version_file ) + if(_CUDAToolkit_version_file) + # CUDAToolkit_LIBRARY_ROOT contains the device library and version file. + get_filename_component(CUDAToolkit_LIBRARY_ROOT "${_CUDAToolkit_version_file}" DIRECTORY ABSOLUTE) + endif() + unset(_CUDAToolkit_version_file) + + if(CUDAToolkit_NVCC_EXECUTABLE AND + CMAKE_CUDA_COMPILER_VERSION AND + CUDAToolkit_NVCC_EXECUTABLE STREQUAL CMAKE_CUDA_COMPILER) + # Need to set these based off the already computed CMAKE_CUDA_COMPILER_VERSION value + # This if statement will always match, but is used to provide variables for MATCH 1,2,3... + if(CMAKE_CUDA_COMPILER_VERSION MATCHES [=[([0-9]+)\.([0-9]+)\.([0-9]+)]=]) + set(CUDAToolkit_VERSION_MAJOR "${CMAKE_MATCH_1}") + set(CUDAToolkit_VERSION_MINOR "${CMAKE_MATCH_2}") + set(CUDAToolkit_VERSION_PATCH "${CMAKE_MATCH_3}") + set(CUDAToolkit_VERSION "${CMAKE_CUDA_COMPILER_VERSION}") + endif() + elseif(CUDAToolkit_NVCC_EXECUTABLE) + # Compute the version by invoking nvcc + execute_process(COMMAND ${CUDAToolkit_NVCC_EXECUTABLE} "--version" OUTPUT_VARIABLE NVCC_OUT) + if(NVCC_OUT MATCHES [=[ V([0-9]+)\.([0-9]+)\.([0-9]+)]=]) + set(CUDAToolkit_VERSION_MAJOR "${CMAKE_MATCH_1}") + set(CUDAToolkit_VERSION_MINOR "${CMAKE_MATCH_2}") + set(CUDAToolkit_VERSION_PATCH "${CMAKE_MATCH_3}") + set(CUDAToolkit_VERSION "${CMAKE_MATCH_1}.${CMAKE_MATCH_2}.${CMAKE_MATCH_3}") + endif() + unset(NVCC_OUT) + else() + _CUDAToolkit_find_version_file(version_file) + if(version_file) + file(READ "${version_file}" VERSION_INFO) + if(VERSION_INFO MATCHES [=[CUDA Version ([0-9]+)\.([0-9]+)\.([0-9]+)]=]) + set(CUDAToolkit_VERSION_MAJOR "${CMAKE_MATCH_1}") + set(CUDAToolkit_VERSION_MINOR "${CMAKE_MATCH_2}") + set(CUDAToolkit_VERSION_PATCH "${CMAKE_MATCH_3}") + set(CUDAToolkit_VERSION "${CMAKE_MATCH_1}.${CMAKE_MATCH_2}.${CMAKE_MATCH_3}") + endif() + endif() + endif() +endif() + +# Find target directory when crosscompiling. +if(CMAKE_CROSSCOMPILING) + if(CMAKE_SYSTEM_PROCESSOR STREQUAL "armv7-a") + # Support for NVPACK + set(CUDAToolkit_TARGET_NAME "armv7-linux-androideabi") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "arm") + set(CUDAToolkit_TARGET_NAME "armv7-linux-gnueabihf") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") + if(ANDROID_ARCH_NAME STREQUAL "arm64") + set(CUDAToolkit_TARGET_NAME "aarch64-linux-androideabi") + elseif(CMAKE_SYSTEM_NAME STREQUAL "QNX") + set(CUDAToolkit_TARGET_NAME "aarch64-qnx") + else() + set(CUDAToolkit_TARGET_NAME "aarch64-linux") + endif(ANDROID_ARCH_NAME STREQUAL "arm64") + elseif(CMAKE_SYSTEM_PROCESSOR STREQUAL "x86_64") + set(CUDAToolkit_TARGET_NAME "x86_64-linux") + endif() + + if(EXISTS "${CUDAToolkit_ROOT_DIR}/targets/${CUDAToolkit_TARGET_NAME}") + set(CUDAToolkit_TARGET_DIR "${CUDAToolkit_ROOT_DIR}/targets/${CUDAToolkit_TARGET_NAME}") + # add known CUDA target root path to the set of directories we search for programs, libraries and headers + list(PREPEND CMAKE_FIND_ROOT_PATH "${CUDAToolkit_TARGET_DIR}") + + # Mark that we need to pop the root search path changes after we have + # found all cuda libraries so that searches for our cross-compilation + # libraries work when another cuda sdk is in CMAKE_PREFIX_PATH or + # PATh + set(_CUDAToolkit_Pop_ROOT_PATH True) + endif() +endif() + +# If not already set we can simply use the toolkit root or it's a scattered installation. +if(NOT CUDAToolkit_TARGET_DIR) + # Not cross compiling + set(CUDAToolkit_TARGET_DIR "${CUDAToolkit_ROOT_DIR}") + # Now that we have the real ROOT_DIR, find components inside it. + list(APPEND CMAKE_PREFIX_PATH ${CUDAToolkit_ROOT_DIR}) + + # Mark that we need to pop the prefix path changes after we have + # found the cudart library. + set(_CUDAToolkit_Pop_Prefix True) +endif() + +# CUDAToolkit_TARGET_DIR always points to the directory containing the include directory. +# On a scattered installation /usr, on a non-scattered something like /usr/local/cuda or /usr/local/cuda-10.2/targets/aarch64-linux. +if(EXISTS "${CUDAToolkit_TARGET_DIR}/include/cuda_runtime.h") + set(CUDAToolkit_INCLUDE_DIR "${CUDAToolkit_TARGET_DIR}/include") +elseif(NOT CUDAToolkit_FIND_QUIETLY) + message(STATUS "Unable to find cuda_runtime.h in \"${CUDAToolkit_TARGET_DIR}/include\" for CUDAToolkit_INCLUDE_DIR.") +endif() + +# The NVHPC layout moves math library headers and libraries to a sibling directory. +# Create a separate variable so this directory can be selectively added to math targets. +if(NOT EXISTS "${CUDAToolkit_INCLUDE_DIR}/cublas_v2.h") + set(CUDAToolkit_MATH_INCLUDE_DIR "${CUDAToolkit_TARGET_DIR}/../../math_libs/include") + get_filename_component(CUDAToolkit_MATH_INCLUDE_DIR "${CUDAToolkit_MATH_INCLUDE_DIR}" ABSOLUTE) + if(NOT EXISTS "${CUDAToolkit_MATH_INCLUDE_DIR}/cublas_v2.h") + if(NOT CUDAToolkit_FIND_QUIETLY) + message(STATUS "Unable to find cublas_v2.h in either \"${CUDAToolkit_INCLUDE_DIR}\" or \"${CUDAToolkit_MATH_INCLUDE_DIR}\"") + endif() + unset(CUDAToolkit_MATH_INCLUDE_DIR) + endif() +endif() + +# Find the CUDA Runtime Library libcudart +find_library(CUDA_CUDART + NAMES cudart + PATH_SUFFIXES lib64 lib/x64 +) +find_library(CUDA_CUDART + NAMES cudart + PATH_SUFFIXES lib64/stubs lib/x64/stubs +) + +if(NOT CUDA_CUDART AND NOT CUDAToolkit_FIND_QUIETLY) + message(STATUS "Unable to find cudart library.") +endif() + +if(_CUDAToolkit_Pop_Prefix) + list(REMOVE_AT CMAKE_PREFIX_PATH -1) + unset(_CUDAToolkit_Pop_Prefix) +endif() + +#----------------------------------------------------------------------------- +# Perform version comparison and validate all required variables are set. +include(FindPackageHandleStandardArgs) +find_package_handle_standard_args(CUDAToolkit + REQUIRED_VARS + CUDAToolkit_INCLUDE_DIR + CUDAToolkit_VERSION + CUDA_CUDART + CUDAToolkit_BIN_DIR + VERSION_VAR + CUDAToolkit_VERSION +) + +mark_as_advanced(CUDA_CUDART + CUDAToolkit_INCLUDE_DIR + CUDAToolkit_NVCC_EXECUTABLE + CUDAToolkit_SENTINEL_FILE + ) + +#----------------------------------------------------------------------------- +# Construct result variables +if(CUDAToolkit_FOUND) + set(CUDAToolkit_INCLUDE_DIRS ${CUDAToolkit_INCLUDE_DIR}) + get_filename_component(CUDAToolkit_LIBRARY_DIR ${CUDA_CUDART} DIRECTORY ABSOLUTE) +endif() + +#----------------------------------------------------------------------------- +# Construct import targets +if(CUDAToolkit_FOUND) + + function(_CUDAToolkit_find_and_add_import_lib lib_name) + cmake_parse_arguments(arg "" "" "ALT;DEPS;EXTRA_HINTS;EXTRA_PATH_SUFFIXES;EXTRA_INCLUDE_DIRS" ${ARGN}) + + set(search_names ${lib_name} ${arg_ALT}) + + find_library(CUDA_${lib_name}_LIBRARY + NAMES ${search_names} + HINTS ${CUDAToolkit_LIBRARY_DIR} + ENV CUDA_PATH + ${arg_EXTRA_HINTS} + PATH_SUFFIXES nvidia/current lib64 lib/x64 lib + ${arg_EXTRA_PATH_SUFFIXES} + ) + # Don't try any stub directories until we have exhausted all other + # search locations. + find_library(CUDA_${lib_name}_LIBRARY + NAMES ${search_names} + HINTS ${CUDAToolkit_LIBRARY_DIR} + ENV CUDA_PATH + ${arg_EXTRA_HINTS} + PATH_SUFFIXES lib64/stubs lib/x64/stubs lib/stubs stubs + # Support NVHPC splayed math library layout + ../../math_libs/${CUDAToolkit_VERSION_MAJOR}.${CUDAToolkit_VERSION_MINOR}/lib64 + ../../math_libs/lib64 + ) + + mark_as_advanced(CUDA_${lib_name}_LIBRARY) + + if(NOT TARGET CUDA::${lib_name} AND CUDA_${lib_name}_LIBRARY) + add_library(CUDA::${lib_name} UNKNOWN IMPORTED) + set_property(TARGET CUDA::${lib_name} APPEND PROPERTY + INTERFACE_INCLUDE_DIRECTORIES "${CUDAToolkit_INCLUDE_DIRS}") + set_property(TARGET CUDA::${lib_name} APPEND PROPERTY + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "${CUDAToolkit_INCLUDE_DIRS}") + if(DEFINED CUDAToolkit_MATH_INCLUDE_DIR) + string(FIND ${CUDA_${lib_name}_LIBRARY} "math_libs" math_libs) + if(NOT ${math_libs} EQUAL -1) + set_property(TARGET CUDA::${lib_name} APPEND PROPERTY + INTERFACE_INCLUDE_DIRECTORIES "${CUDAToolkit_MATH_INCLUDE_DIRS}") + set_property(TARGET CUDA::${lib_name} APPEND PROPERTY + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "${CUDAToolkit_MATH_INCLUDE_DIRS}") + endif() + endif() + set_property(TARGET CUDA::${lib_name} PROPERTY IMPORTED_LOCATION "${CUDA_${lib_name}_LIBRARY}") + foreach(dep ${arg_DEPS}) + if(TARGET CUDA::${dep}) + set_property(TARGET CUDA::${lib_name} APPEND PROPERTY + INTERFACE_LINK_LIBRARIES CUDA::${dep}) + endif() + endforeach() + if(arg_EXTRA_INCLUDE_DIRS) + set_property(TARGET CUDA::${lib_name} APPEND PROPERTY + INTERFACE_INCLUDE_DIRECTORIES "${arg_EXTRA_INCLUDE_DIRS}") + set_property(TARGET CUDA::${lib_name} APPEND PROPERTY + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "${arg_EXTRA_INCLUDE_DIRS}") + endif() + endif() + endfunction() + + if(NOT TARGET CUDA::toolkit) + add_library(CUDA::toolkit IMPORTED INTERFACE) + set_property(TARGET CUDA::toolkit APPEND PROPERTY + INTERFACE_INCLUDE_DIRECTORIES "${CUDAToolkit_INCLUDE_DIRS}") + set_property(TARGET CUDA::toolkit APPEND PROPERTY + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES "${CUDAToolkit_INCLUDE_DIRS}") + endif() + + _CUDAToolkit_find_and_add_import_lib(cuda_driver ALT cuda) + + _CUDAToolkit_find_and_add_import_lib(cudart) + _CUDAToolkit_find_and_add_import_lib(cudart_static) + + # setup dependencies that are required for cudart_static when building + # on linux. These are generally only required when using the CUDA toolkit + # when CUDA language is disabled + if(NOT TARGET CUDA::cudart_static_deps + AND TARGET CUDA::cudart_static) + + add_library(CUDA::cudart_static_deps IMPORTED INTERFACE) + set_property(TARGET CUDA::cudart_static APPEND PROPERTY + INTERFACE_LINK_LIBRARIES CUDA::cudart_static_deps) + + if(UNIX AND (CMAKE_C_COMPILER OR CMAKE_CXX_COMPILER)) + find_package(Threads REQUIRED) + set_property(TARGET CUDA::cudart_static_deps APPEND PROPERTY + INTERFACE_LINK_LIBRARIES Threads::Threads ${CMAKE_DL_LIBS}) + endif() + + if(UNIX AND NOT APPLE AND NOT (CMAKE_SYSTEM_NAME STREQUAL "QNX")) + # On Linux, you must link against librt when using the static cuda runtime. + find_library(CUDAToolkit_rt_LIBRARY rt) + mark_as_advanced(CUDAToolkit_rt_LIBRARY) + if(NOT CUDAToolkit_rt_LIBRARY) + message(WARNING "Could not find librt library, needed by CUDA::cudart_static") + else() + set_property(TARGET CUDA::cudart_static_deps APPEND PROPERTY + INTERFACE_LINK_LIBRARIES ${CUDAToolkit_rt_LIBRARY}) + endif() + endif() + endif() + + _CUDAToolkit_find_and_add_import_lib(culibos) # it's a static library + foreach(cuda_lib cublasLt cufft curand cusparse nppc nvjpeg) + _CUDAToolkit_find_and_add_import_lib(${cuda_lib}) + _CUDAToolkit_find_and_add_import_lib(${cuda_lib}_static DEPS culibos) + endforeach() + + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 11.0.0) + # cublas depends on cublasLt + # https://docs.nvidia.com/cuda/archive/11.0/cublas/index.html#static-library + _CUDAToolkit_find_and_add_import_lib(cublas DEPS cublasLt) + _CUDAToolkit_find_and_add_import_lib(cublas_static DEPS cublasLt_static) + else() + _CUDAToolkit_find_and_add_import_lib(cublas) + _CUDAToolkit_find_and_add_import_lib(cublas_static DEPS culibos) + endif() + + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 11.4) + _CUDAToolkit_find_and_add_import_lib(cuFile ALT cufile DEPS culibos) + _CUDAToolkit_find_and_add_import_lib(cuFile_static ALT cufile_static DEPS culibos) + + _CUDAToolkit_find_and_add_import_lib(cuFile_rdma ALT cufile_rdma DEPS cuFile culibos) + _CUDAToolkit_find_and_add_import_lib(cuFile_rdma_static ALT cufile_rdma_static DEPS cuFile_static culibos) + endif() + + # cuFFTW depends on cuFFT + _CUDAToolkit_find_and_add_import_lib(cufftw DEPS cufft) + _CUDAToolkit_find_and_add_import_lib(cufftw_static DEPS cufft_static) + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 9.2) + _CUDAToolkit_find_and_add_import_lib(cufft_static_nocallback DEPS culibos) + endif() + + # cuSOLVER depends on cuBLAS, and cuSPARSE + _CUDAToolkit_find_and_add_import_lib(cusolver DEPS cublas cusparse) + _CUDAToolkit_find_and_add_import_lib(cusolver_static DEPS cublas_static cusparse_static culibos) + + + if(CUDAToolkit_VERSION VERSION_GREATER_EQUAL 10.1.2) + # cusolver depends on liblapack_static.a starting with CUDA 10.1 update 2, + # https://docs.nvidia.com/cuda/archive/11.5.0/cusolver/index.html#static-link-lapack + _CUDAToolkit_find_and_add_import_lib(cusolver_lapack_static ALT lapack_static) # implementation detail static lib + _CUDAToolkit_find_and_add_import_lib(cusolver_static DEPS cusolver_lapack_static) + endif() + + if(CUDAToolkit_VERSION VERSION_GREATER 11.2.1) + # cusolver depends on libcusolver_metis and cublasLt + # https://docs.nvidia.com/cuda/archive/11.2.2/cusolver/index.html#link-dependency + _CUDAToolkit_find_and_add_import_lib(cusolver DEPS cublasLt) + + _CUDAToolkit_find_and_add_import_lib(cusolver_metis_static ALT metis_static) # implementation detail static lib + _CUDAToolkit_find_and_add_import_lib(cusolver_static DEPS cusolver_metis_static cublasLt_static) + endif() + + # nvGRAPH depends on cuRAND, and cuSOLVER. + _CUDAToolkit_find_and_add_import_lib(nvgraph DEPS curand cusolver) + _CUDAToolkit_find_and_add_import_lib(nvgraph_static DEPS curand_static cusolver_static) + + # Process the majority of the NPP libraries. + foreach(cuda_lib nppial nppicc nppidei nppif nppig nppim nppist nppitc npps nppicom nppisu) + _CUDAToolkit_find_and_add_import_lib(${cuda_lib} DEPS nppc) + _CUDAToolkit_find_and_add_import_lib(${cuda_lib}_static DEPS nppc_static) + endforeach() + + find_path(CUDAToolkit_CUPTI_INCLUDE_DIR cupti.h PATHS + "${CUDAToolkit_ROOT_DIR}/extras/CUPTI/include" + "${CUDAToolkit_INCLUDE_DIR}/../extras/CUPTI/include" + "${CUDAToolkit_INCLUDE_DIR}" + NO_DEFAULT_PATH) + mark_as_advanced(CUDAToolkit_CUPTI_INCLUDE_DIR) + + if(CUDAToolkit_CUPTI_INCLUDE_DIR) + _CUDAToolkit_find_and_add_import_lib(cupti + EXTRA_PATH_SUFFIXES ../extras/CUPTI/lib64/ + ../extras/CUPTI/lib/ + EXTRA_INCLUDE_DIRS "${CUDAToolkit_CUPTI_INCLUDE_DIR}") + _CUDAToolkit_find_and_add_import_lib(cupti_static + EXTRA_PATH_SUFFIXES ../extras/CUPTI/lib64/ + ../extras/CUPTI/lib/ + EXTRA_INCLUDE_DIRS "${CUDAToolkit_CUPTI_INCLUDE_DIR}") + endif() + + _CUDAToolkit_find_and_add_import_lib(nvrtc DEPS cuda_driver) + + _CUDAToolkit_find_and_add_import_lib(nvml ALT nvidia-ml nvml) + + # nvtools can be installed outside the CUDA toolkit directory, + # so search the NVTOOLSEXT_PATH windows only environment variable + set(nvToolsExt_EXTRA_PATH) + if(WIN32) + set(nvToolsExt_EXTRA_PATH "C:\\Program Files\\NVIDIA Corporation\\NvToolsExt") + endif() + + find_path(CUDAToolkit_nvToolsExt_INCLUDE_DIR nvToolsExt.h + PATHS "${CUDAToolkit_INCLUDE_DIR}" + "${CUDAToolkit_ROOT_DIR}" + ENV NVTOOLSEXT_PATH + "${nvToolsExt_EXTRA_PATH}" + PATH_SUFFIXES include + NO_DEFAULT_PATH) + mark_as_advanced(CUDAToolkit_nvToolsExt_INCLUDE_DIR) + + if(CUDAToolkit_nvToolsExt_INCLUDE_DIR) + _CUDAToolkit_find_and_add_import_lib(nvToolsExt + ALT nvToolsExt64 nvToolsExt64_1 + EXTRA_HINTS ENV NVTOOLSEXT_PATH + "${nvToolsExt_EXTRA_PATH}" + EXTRA_INCLUDE_DIRS "${CUDAToolkit_nvToolsExt_INCLUDE_DIR}") + endif() + + _CUDAToolkit_find_and_add_import_lib(OpenCL) +endif() + +unset(CUDAToolkit_ROOT_DIR) + +if(_CUDAToolkit_Pop_ROOT_PATH) + list(REMOVE_AT CMAKE_FIND_ROOT_PATH 0) + unset(_CUDAToolkit_Pop_ROOT_PATH) +endif() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindCUDSS.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindCUDSS.cmake new file mode 100644 index 0000000000000000000000000000000000000000..b614e1c492b99f7b3adf456b0b88bdf5cd26fd0b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindCUDSS.cmake @@ -0,0 +1,67 @@ +# Find the CUDSS library +# +# The following variables are optionally searched for defaults +# CUDSS_ROOT: Base directory where CUDSS is found +# CUDSS_INCLUDE_DIR: Directory where CUDSS header is searched for +# CUDSS_LIBRARY: Directory where CUDSS library is searched for +# +# The following are set after configuration is done: +# CUDSS_FOUND +# CUDSS_INCLUDE_PATH +# CUDSS_LIBRARY_PATH + +include(FindPackageHandleStandardArgs) + +set(CUDSS_ROOT $ENV{CUDSS_ROOT_DIR} CACHE PATH "Folder containing NVIDIA CUDSS") +if (DEFINED $ENV{CUDSS_ROOT_DIR}) + message(WARNING "CUDSS_ROOT_DIR is deprecated. Please set CUDSS_ROOT instead.") +endif() +list(APPEND CUDSS_ROOT $ENV{CUDSS_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}) + +# Compatible layer for CMake <3.12. CUDSS_ROOT will be accounted in for searching paths and libraries for CMake >=3.12. +list(APPEND CMAKE_PREFIX_PATH ${CUDSS_ROOT}) + +set(CUDSS_INCLUDE_DIR $ENV{CUDSS_INCLUDE_DIR} CACHE PATH "Folder containing NVIDIA CUDSS header files") + +find_path(CUDSS_INCLUDE_PATH cudss.h + HINTS ${CUDSS_INCLUDE_DIR} + PATH_SUFFIXES cuda/include cuda include) + +set(CUDSS_LIBRARY $ENV{CUDSS_LIBRARY} CACHE PATH "Path to the CUDSS library file (e.g., libcudss.so)") + +set(CUDSS_LIBRARY_NAME "libcudss.so") +if(MSVC) + set(CUDSS_LIBRARY_NAME "cudss.lib") +endif() + +find_library(CUDSS_LIBRARY_PATH ${CUDSS_LIBRARY_NAME} + PATHS ${CUDSS_LIBRARY} + PATH_SUFFIXES lib lib64 cuda/lib cuda/lib64 lib/x64) + +find_package_handle_standard_args(CUDSS DEFAULT_MSG CUDSS_LIBRARY_PATH CUDSS_INCLUDE_PATH) + +if(CUDSS_FOUND) + # Get CUDSS version + file(READ ${CUDSS_INCLUDE_PATH}/cudss.h CUDSS_HEADER_CONTENTS) + string(REGEX MATCH "define CUDSS_VER_MAJOR * +([0-9]+)" + CUDSS_VERSION_MAJOR "${CUDSS_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUDSS_VER_MAJOR * +([0-9]+)" "\\1" + CUDSS_VERSION_MAJOR "${CUDSS_VERSION_MAJOR}") + string(REGEX MATCH "define CUDSS_VER_MINOR * +([0-9]+)" + CUDSS_VERSION_MINOR "${CUDSS_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUDSS_VER_MINOR * +([0-9]+)" "\\1" + CUDSS_VERSION_MINOR "${CUDSS_VERSION_MINOR}") + string(REGEX MATCH "define CUDSS_VER_PATCH * +([0-9]+)" + CUDSS_VERSION_PATCH "${CUDSS_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUDSS_VER_PATCH * +([0-9]+)" "\\1" + CUDSS_VERSION_PATCH "${CUDSS_VERSION_PATCH}") + # Assemble CUDSS version. Use minor version since current major version is 0. + if(NOT CUDSS_VERSION_MINOR) + set(CUDSS_VERSION "?") + else() + set(CUDSS_VERSION + "${CUDSS_VERSION_MAJOR}.${CUDSS_VERSION_MINOR}.${CUDSS_VERSION_PATCH}") + endif() +endif() + +mark_as_advanced(CUDSS_ROOT CUDSS_INCLUDE_DIR CUDSS_LIBRARY CUDSS_VERSION) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindCUSPARSELT.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindCUSPARSELT.cmake new file mode 100644 index 0000000000000000000000000000000000000000..6c15bde147469ddc84980dca0c756e8f26e1ddb1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindCUSPARSELT.cmake @@ -0,0 +1,67 @@ +# Find the CUSPARSELT library +# +# The following variables are optionally searched for defaults +# CUSPARSELT_ROOT: Base directory where CUSPARSELT is found +# CUSPARSELT_INCLUDE_DIR: Directory where CUSPARSELT header is searched for +# CUSPARSELT_LIBRARY: Directory where CUSPARSELT library is searched for +# +# The following are set after configuration is done: +# CUSPARSELT_FOUND +# CUSPARSELT_INCLUDE_PATH +# CUSPARSELT_LIBRARY_PATH + +include(FindPackageHandleStandardArgs) + +set(CUSPARSELT_ROOT $ENV{CUSPARSELT_ROOT_DIR} CACHE PATH "Folder containing NVIDIA cuSPARSELt") +if (DEFINED $ENV{CUSPARSELT_ROOT_DIR}) + message(WARNING "CUSPARSELT_ROOT_DIR is deprecated. Please set CUSPARSELT_ROOT instead.") +endif() +list(APPEND CUSPARSELT_ROOT $ENV{CUSPARSELT_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}) + +# Compatible layer for CMake <3.12. CUSPARSELT_ROOT will be accounted in for searching paths and libraries for CMake >=3.12. +list(APPEND CMAKE_PREFIX_PATH ${CUSPARSELT_ROOT}) + +set(CUSPARSELT_INCLUDE_DIR $ENV{CUSPARSELT_INCLUDE_DIR} CACHE PATH "Folder containing NVIDIA cuSPARSELt header files") + +find_path(CUSPARSELT_INCLUDE_PATH cusparseLt.h + HINTS ${CUSPARSELT_INCLUDE_DIR} + PATH_SUFFIXES cuda/include cuda include) + +set(CUSPARSELT_LIBRARY $ENV{CUSPARSELT_LIBRARY} CACHE PATH "Path to the cusparselt library file (e.g., libcusparseLt.so)") + +set(CUSPARSELT_LIBRARY_NAME "libcusparseLt.so") +if(MSVC) + set(CUSPARSELT_LIBRARY_NAME "cusparseLt.lib") +endif() + +find_library(CUSPARSELT_LIBRARY_PATH ${CUSPARSELT_LIBRARY_NAME} + PATHS ${CUSPARSELT_LIBRARY} + PATH_SUFFIXES lib lib64 cuda/lib cuda/lib64 lib/x64) + +find_package_handle_standard_args(CUSPARSELT DEFAULT_MSG CUSPARSELT_LIBRARY_PATH CUSPARSELT_INCLUDE_PATH) + +if(CUSPARSELT_FOUND) + # Get cuSPARSELt version + file(READ ${CUSPARSELT_INCLUDE_PATH}/cusparseLt.h CUSPARSELT_HEADER_CONTENTS) + string(REGEX MATCH "define CUSPARSELT_VER_MAJOR * +([0-9]+)" + CUSPARSELT_VERSION_MAJOR "${CUSPARSELT_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUSPARSELT_VER_MAJOR * +([0-9]+)" "\\1" + CUSPARSELT_VERSION_MAJOR "${CUSPARSELT_VERSION_MAJOR}") + string(REGEX MATCH "define CUSPARSELT_VER_MINOR * +([0-9]+)" + CUSPARSELT_VERSION_MINOR "${CUSPARSELT_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUSPARSELT_VER_MINOR * +([0-9]+)" "\\1" + CUSPARSELT_VERSION_MINOR "${CUSPARSELT_VERSION_MINOR}") + string(REGEX MATCH "define CUSPARSELT_VER_PATCH * +([0-9]+)" + CUSPARSELT_VERSION_PATCH "${CUSPARSELT_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUSPARSELT_VER_PATCH * +([0-9]+)" "\\1" + CUSPARSELT_VERSION_PATCH "${CUSPARSELT_VERSION_PATCH}") + # Assemble cuSPARSELt version. Use minor version since current major version is 0. + if(NOT CUSPARSELT_VERSION_MINOR) + set(CUSPARSELT_VERSION "?") + else() + set(CUSPARSELT_VERSION + "${CUSPARSELT_VERSION_MAJOR}.${CUSPARSELT_VERSION_MINOR}.${CUSPARSELT_VERSION_PATCH}") + endif() +endif() + +mark_as_advanced(CUSPARSELT_ROOT CUSPARSELT_INCLUDE_DIR CUSPARSELT_LIBRARY CUSPARSELT_VERSION) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindSYCLToolkit.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindSYCLToolkit.cmake new file mode 100644 index 0000000000000000000000000000000000000000..337afa1bfe4178d1af041c6504c1124b8c31d482 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/FindSYCLToolkit.cmake @@ -0,0 +1,141 @@ +# This will define the following variables: +# SYCL_FOUND : True if the system has the SYCL library. +# SYCL_INCLUDE_DIR : Include directories needed to use SYCL. +# SYCL_LIBRARY_DIR : The path to the SYCL library. +# SYCL_LIBRARY : SYCL library fullname. +# SYCL_COMPILER_VERSION : SYCL compiler version. + +include(FindPackageHandleStandardArgs) + +set(SYCL_ROOT "") +if(DEFINED ENV{SYCL_ROOT}) + set(SYCL_ROOT $ENV{SYCL_ROOT}) +elseif(DEFINED ENV{CMPLR_ROOT}) + set(SYCL_ROOT $ENV{CMPLR_ROOT}) +else() + # Use the default path to ensure proper linking with torch::xpurt when the user is working with libtorch. + if(CMAKE_SYSTEM_NAME MATCHES "Linux") + set(SYCL_ROOT "/opt/intel/oneapi/compiler/latest") + elseif(CMAKE_SYSTEM_NAME MATCHES "Windows") + set(SYCL_ROOT "C:/Program Files (x86)/Intel/oneAPI/compiler/latest") + endif() + if(NOT EXISTS ${SYCL_ROOT}) + set(SYCL_ROOT "") + endif() +endif() + +string(COMPARE EQUAL "${SYCL_ROOT}" "" nosyclfound) +if(nosyclfound) + set(SYCL_FOUND False) + set(SYCL_REASON_FAILURE "SYCL library not set!!") + set(SYCL_NOT_FOUND_MESSAGE "${SYCL_REASON_FAILURE}") + return() +endif() + +# Find SYCL compiler executable. +find_program( + SYCL_COMPILER + NAMES icx + PATHS "${SYCL_ROOT}" + PATH_SUFFIXES bin bin64 + NO_DEFAULT_PATH + ) + +function(parse_sycl_compiler_version version_number) + # Execute the SYCL compiler with the --version flag to match the version string. + execute_process(COMMAND ${SYCL_COMPILER} --version OUTPUT_VARIABLE SYCL_VERSION_STRING) + string(REGEX REPLACE "Intel\\(R\\) (.*) Compiler ([0-9]+\\.[0-9]+\\.[0-9]+) (.*)" "\\2" + SYCL_VERSION_STRING_MATCH ${SYCL_VERSION_STRING}) + string(REPLACE "." ";" SYCL_VERSION_LIST ${SYCL_VERSION_STRING_MATCH}) + # Split the version number list into major, minor, and patch components. + list(GET SYCL_VERSION_LIST 0 VERSION_MAJOR) + list(GET SYCL_VERSION_LIST 1 VERSION_MINOR) + list(GET SYCL_VERSION_LIST 2 VERSION_PATCH) + # Calculate the version number in the format XXXXYYZZ, using the formula (major * 10000 + minor * 100 + patch). + math(EXPR VERSION_NUMBER_MATCH "${VERSION_MAJOR} * 10000 + ${VERSION_MINOR} * 100 + ${VERSION_PATCH}") + set(${version_number} "${VERSION_NUMBER_MATCH}" PARENT_SCOPE) +endfunction() + +if(SYCL_COMPILER) + parse_sycl_compiler_version(SYCL_COMPILER_VERSION) +endif() + +if(NOT SYCL_COMPILER_VERSION) + set(SYCL_FOUND False) + set(SYCL_REASON_FAILURE "Cannot parse sycl compiler version to get SYCL_COMPILER_VERSION!") + set(SYCL_NOT_FOUND_MESSAGE "${SYCL_REASON_FAILURE}") + return() +endif() + +# Find include path from binary. +find_file( + SYCL_INCLUDE_DIR + NAMES include + HINTS ${SYCL_ROOT} + NO_DEFAULT_PATH + ) + +# Find include/sycl path from include path. +find_file( + SYCL_INCLUDE_SYCL_DIR + NAMES sycl + HINTS ${SYCL_ROOT}/include/ + NO_DEFAULT_PATH + ) + +# Due to the unrecognized compilation option `-fsycl` in other compiler. +list(APPEND SYCL_INCLUDE_DIR ${SYCL_INCLUDE_SYCL_DIR}) + +# Find library directory from binary. +find_file( + SYCL_LIBRARY_DIR + NAMES lib lib64 + HINTS ${SYCL_ROOT} + NO_DEFAULT_PATH + ) + +# Define the old version of SYCL toolkit that is compatible with the current version of PyTorch. +set(PYTORCH_2_5_SYCL_TOOLKIT_VERSION 20249999) + +# By default, we use libsycl.so on Linux and sycl.lib on Windows as the SYCL library name. +if (SYCL_COMPILER_VERSION VERSION_LESS_EQUAL PYTORCH_2_5_SYCL_TOOLKIT_VERSION) + # Don't use if(WIN32) here since this requires cmake>=3.25 and file is installed + # and used by other projects. + # See: https://cmake.org/cmake/help/v3.25/variable/LINUX.html + if(CMAKE_SYSTEM_NAME MATCHES "Windows") + # On Windows, the SYCL library is named sycl7.lib until PYTORCH_2_5_SYCL_TOOLKIT_VERSION. + # sycl.lib is supported in the later version. + set(sycl_lib_suffix "7") + endif() +endif() + +# Find SYCL library fullname. +find_library( + SYCL_LIBRARY + NAMES "sycl${sycl_lib_suffix}" + HINTS ${SYCL_LIBRARY_DIR} + NO_DEFAULT_PATH +) + +# Find OpenCL library fullname, which is a dependency of oneDNN. +find_library( + OCL_LIBRARY + NAMES OpenCL + HINTS ${SYCL_LIBRARY_DIR} + NO_DEFAULT_PATH +) + +if((NOT SYCL_LIBRARY) OR (NOT OCL_LIBRARY)) + set(SYCL_FOUND False) + set(SYCL_REASON_FAILURE "SYCL library is incomplete!!") + set(SYCL_NOT_FOUND_MESSAGE "${SYCL_REASON_FAILURE}") + return() +endif() + +find_package_handle_standard_args( + SYCL + FOUND_VAR SYCL_FOUND + REQUIRED_VARS SYCL_INCLUDE_DIR SYCL_LIBRARY_DIR SYCL_LIBRARY + REASON_FAILURE_MESSAGE "${SYCL_REASON_FAILURE}" + VERSION_VAR SYCL_COMPILER_VERSION + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/FindCUDA.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/FindCUDA.cmake new file mode 100644 index 0000000000000000000000000000000000000000..55c4e83012d820995f59b717ecb676452f9ccbec --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/FindCUDA.cmake @@ -0,0 +1,10 @@ +# This is a wrapper of the upstream `./upstream/FindCUDA.cmake` that +# automatically includes `./upstream/CMakeInitializeConfigs.cmake` before +# `./upstream/FindCUDA.cmake`. The `CMakeInitializeConfigs.cmake`, which is +# absent in old CMake versions, creates some necessary variables for the later +# to run. +# See ./README.md for details. + +set(UPSTREAM_FIND_CUDA_DIR "${CMAKE_CURRENT_LIST_DIR}/upstream/") + +include("${UPSTREAM_FIND_CUDA_DIR}/FindCUDA.cmake") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/FindCUDNN.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/FindCUDNN.cmake new file mode 100644 index 0000000000000000000000000000000000000000..82134328c803dc87a89564638540a6cbcfa2d906 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/FindCUDNN.cmake @@ -0,0 +1,78 @@ +# Find the CUDNN libraries +# +# The following variables are optionally searched for defaults +# CUDNN_ROOT: Base directory where CUDNN is found +# CUDNN_INCLUDE_DIR: Directory where CUDNN header is searched for +# CUDNN_LIBRARY: Directory where CUDNN library is searched for +# CUDNN_STATIC: Are we looking for a static library? (default: no) +# +# The following are set after configuration is done: +# CUDNN_FOUND +# CUDNN_INCLUDE_PATH +# CUDNN_LIBRARY_PATH +# + +include(FindPackageHandleStandardArgs) + +set(CUDNN_ROOT $ENV{CUDNN_ROOT_DIR} CACHE PATH "Folder containing NVIDIA cuDNN") +if (DEFINED $ENV{CUDNN_ROOT_DIR}) + message(WARNING "CUDNN_ROOT_DIR is deprecated. Please set CUDNN_ROOT instead.") +endif() +list(APPEND CUDNN_ROOT $ENV{CUDNN_ROOT_DIR} ${CUDA_TOOLKIT_ROOT_DIR}) + +# Compatible layer for CMake <3.12. CUDNN_ROOT will be accounted in for searching paths and libraries for CMake >=3.12. +list(APPEND CMAKE_PREFIX_PATH ${CUDNN_ROOT}) + +set(CUDNN_INCLUDE_DIR $ENV{CUDNN_INCLUDE_DIR} CACHE PATH "Folder containing NVIDIA cuDNN header files") + +find_path(CUDNN_INCLUDE_PATH cudnn.h + HINTS ${CUDNN_INCLUDE_DIR} + PATH_SUFFIXES cuda/include cuda include) + +option(CUDNN_STATIC "Look for static CUDNN" OFF) +if (CUDNN_STATIC) + set(CUDNN_LIBNAME "libcudnn_static.a") +else() + set(CUDNN_LIBNAME "cudnn") +endif() + +set(CUDNN_LIBRARY $ENV{CUDNN_LIBRARY} CACHE PATH "Path to the cudnn library file (e.g., libcudnn.so)") +if (CUDNN_LIBRARY MATCHES ".*cudnn_static.a" AND NOT CUDNN_STATIC) + message(WARNING "CUDNN_LIBRARY points to a static library (${CUDNN_LIBRARY}) but CUDNN_STATIC is OFF.") +endif() + +find_library(CUDNN_LIBRARY_PATH ${CUDNN_LIBNAME} + PATHS ${CUDNN_LIBRARY} + PATH_SUFFIXES lib lib64 cuda/lib cuda/lib64 lib/x64) + +find_package_handle_standard_args(CUDNN DEFAULT_MSG CUDNN_LIBRARY_PATH CUDNN_INCLUDE_PATH) + +if(CUDNN_FOUND) + # Get cuDNN version + if(EXISTS ${CUDNN_INCLUDE_PATH}/cudnn_version.h) + file(READ ${CUDNN_INCLUDE_PATH}/cudnn_version.h CUDNN_HEADER_CONTENTS) + else() + file(READ ${CUDNN_INCLUDE_PATH}/cudnn.h CUDNN_HEADER_CONTENTS) + endif() + string(REGEX MATCH "define CUDNN_MAJOR * +([0-9]+)" + CUDNN_VERSION_MAJOR "${CUDNN_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUDNN_MAJOR * +([0-9]+)" "\\1" + CUDNN_VERSION_MAJOR "${CUDNN_VERSION_MAJOR}") + string(REGEX MATCH "define CUDNN_MINOR * +([0-9]+)" + CUDNN_VERSION_MINOR "${CUDNN_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUDNN_MINOR * +([0-9]+)" "\\1" + CUDNN_VERSION_MINOR "${CUDNN_VERSION_MINOR}") + string(REGEX MATCH "define CUDNN_PATCHLEVEL * +([0-9]+)" + CUDNN_VERSION_PATCH "${CUDNN_HEADER_CONTENTS}") + string(REGEX REPLACE "define CUDNN_PATCHLEVEL * +([0-9]+)" "\\1" + CUDNN_VERSION_PATCH "${CUDNN_VERSION_PATCH}") + # Assemble cuDNN version + if(NOT CUDNN_VERSION_MAJOR) + set(CUDNN_VERSION "?") + else() + set(CUDNN_VERSION + "${CUDNN_VERSION_MAJOR}.${CUDNN_VERSION_MINOR}.${CUDNN_VERSION_PATCH}") + endif() +endif() + +mark_as_advanced(CUDNN_ROOT CUDNN_INCLUDE_DIR CUDNN_LIBRARY CUDNN_VERSION) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/CMakeInitializeConfigs.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/CMakeInitializeConfigs.cmake new file mode 100644 index 0000000000000000000000000000000000000000..5517e8f0624b1e5538b761e1f4891227007d0045 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/CMakeInitializeConfigs.cmake @@ -0,0 +1,40 @@ +# Distributed under the OSI-approved BSD 3-Clause License. See accompanying +# file Copyright.txt or https://cmake.org/licensing for details. + +# Present in upstream, but not supported on versions of cmake we need to support +# include_guard(GLOBAL) + +# Initializes `<_PREFIX>_` variables from the corresponding +# `<_PREFIX>__INIT`, for the configurations currently used. +function(cmake_initialize_per_config_variable _PREFIX _DOCSTRING) + string(STRIP "${${_PREFIX}_INIT}" _INIT) + set("${_PREFIX}" "${_INIT}" + CACHE STRING "${_DOCSTRING} during all build types.") + mark_as_advanced("${_PREFIX}") + + if (NOT CMAKE_NOT_USING_CONFIG_FLAGS) + set(_CONFIGS Debug Release MinSizeRel RelWithDebInfo) + + get_property(_GENERATOR_IS_MULTI_CONFIG GLOBAL PROPERTY GENERATOR_IS_MULTI_CONFIG) + if (_GENERATOR_IS_MULTI_CONFIG) + list(APPEND _CONFIGS ${CMAKE_CONFIGURATION_TYPES}) + else() + if (NOT CMAKE_NO_BUILD_TYPE) + set(CMAKE_BUILD_TYPE "${CMAKE_BUILD_TYPE_INIT}" CACHE STRING + "Choose the type of build, options are: None Debug Release RelWithDebInfo MinSizeRel ...") + endif() + list(APPEND _CONFIGS ${CMAKE_BUILD_TYPE}) + endif() + + list(REMOVE_DUPLICATES _CONFIGS) + foreach(_BUILD_TYPE IN LISTS _CONFIGS) + if (NOT "${_BUILD_TYPE}" STREQUAL "") + string(TOUPPER "${_BUILD_TYPE}" _BUILD_TYPE) + string(STRIP "${${_PREFIX}_${_BUILD_TYPE}_INIT}" _INIT) + set("${_PREFIX}_${_BUILD_TYPE}" "${_INIT}" + CACHE STRING "${_DOCSTRING} during ${_BUILD_TYPE} builds.") + mark_as_advanced("${_PREFIX}_${_BUILD_TYPE}") + endif() + endforeach() + endif() +endfunction() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA.cmake new file mode 100644 index 0000000000000000000000000000000000000000..411a246656b3bdaba6abc238fd35caf959c9cca0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA.cmake @@ -0,0 +1,1981 @@ +#.rst: +# FindCUDA +# -------- +# +# .. note:: +# +# The FindCUDA module has been superseded by first-class support +# for the CUDA language in CMake. It is no longer necessary to +# use this module or call ``find_package(CUDA)``. This module +# now exists only for compatibility with projects that have not +# been ported. +# +# Instead, list ``CUDA`` among the languages named in the top-level +# call to the :command:`project` command, or call the +# :command:`enable_language` command with ``CUDA``. +# Then one can add CUDA (``.cu``) sources to programs directly +# in calls to :command:`add_library` and :command:`add_executable`. +# +# Tools for building CUDA C files: libraries and build dependencies. +# +# This script locates the NVIDIA CUDA C tools. It should work on Linux, +# Windows, and macOS and should be reasonably up to date with CUDA C +# releases. +# +# This script makes use of the standard :command:`find_package` arguments of +# ````, ``REQUIRED`` and ``QUIET``. ``CUDA_FOUND`` will report if an +# acceptable version of CUDA was found. +# +# The script will prompt the user to specify ``CUDA_TOOLKIT_ROOT_DIR`` if +# the prefix cannot be determined by the location of nvcc in the system +# path and ``REQUIRED`` is specified to :command:`find_package`. To use +# a different installed version of the toolkit set the environment variable +# ``CUDA_BIN_PATH`` before running cmake (e.g. +# ``CUDA_BIN_PATH=/usr/local/cuda1.0`` instead of the default +# ``/usr/local/cuda``) or set ``CUDA_TOOLKIT_ROOT_DIR`` after configuring. If +# you change the value of ``CUDA_TOOLKIT_ROOT_DIR``, various components that +# depend on the path will be relocated. +# +# It might be necessary to set ``CUDA_TOOLKIT_ROOT_DIR`` manually on certain +# platforms, or to use a CUDA runtime not installed in the default +# location. In newer versions of the toolkit the CUDA library is +# included with the graphics driver -- be sure that the driver version +# matches what is needed by the CUDA runtime version. +# +# The following variables affect the behavior of the macros in the +# script (in alphebetical order). Note that any of these flags can be +# changed multiple times in the same directory before calling +# ``CUDA_ADD_EXECUTABLE``, ``CUDA_ADD_LIBRARY``, ``CUDA_COMPILE``, +# ``CUDA_COMPILE_PTX``, ``CUDA_COMPILE_FATBIN``, ``CUDA_COMPILE_CUBIN`` +# or ``CUDA_WRAP_SRCS``:: +# +# CUDA_64_BIT_DEVICE_CODE (Default matches host bit size) +# -- Set to ON to compile for 64 bit device code, OFF for 32 bit device code. +# Note that making this different from the host code when generating object +# or C files from CUDA code just won't work, because size_t gets defined by +# nvcc in the generated source. If you compile to PTX and then load the +# file yourself, you can mix bit sizes between device and host. +# +# CUDA_ATTACH_VS_BUILD_RULE_TO_CUDA_FILE (Default ON) +# -- Set to ON if you want the custom build rule to be attached to the source +# file in Visual Studio. Turn OFF if you add the same cuda file to multiple +# targets. +# +# This allows the user to build the target from the CUDA file; however, bad +# things can happen if the CUDA source file is added to multiple targets. +# When performing parallel builds it is possible for the custom build +# command to be run more than once and in parallel causing cryptic build +# errors. VS runs the rules for every source file in the target, and a +# source can have only one rule no matter how many projects it is added to. +# When the rule is run from multiple targets race conditions can occur on +# the generated file. Eventually everything will get built, but if the user +# is unaware of this behavior, there may be confusion. It would be nice if +# this script could detect the reuse of source files across multiple targets +# and turn the option off for the user, but no good solution could be found. +# +# CUDA_BUILD_CUBIN (Default OFF) +# -- Set to ON to enable and extra compilation pass with the -cubin option in +# Device mode. The output is parsed and register, shared memory usage is +# printed during build. +# +# CUDA_BUILD_EMULATION (Default OFF for device mode) +# -- Set to ON for Emulation mode. -D_DEVICEEMU is defined for CUDA C files +# when CUDA_BUILD_EMULATION is TRUE. +# +# CUDA_LINK_LIBRARIES_KEYWORD (Default "") +# -- The keyword to use for internal +# target_link_libraries calls. The default is to use no keyword which +# uses the old "plain" form of target_link_libraries. Note that is matters +# because whatever is used inside the FindCUDA module must also be used +# outside - the two forms of target_link_libraries cannot be mixed. +# +# CUDA_GENERATED_OUTPUT_DIR (Default CMAKE_CURRENT_BINARY_DIR) +# -- Set to the path you wish to have the generated files placed. If it is +# blank output files will be placed in CMAKE_CURRENT_BINARY_DIR. +# Intermediate files will always be placed in +# CMAKE_CURRENT_BINARY_DIR/CMakeFiles. +# +# CUDA_HOST_COMPILATION_CPP (Default ON) +# -- Set to OFF for C compilation of host code. +# +# CUDA_HOST_COMPILER (Default CMAKE_C_COMPILER) +# -- Set the host compiler to be used by nvcc. Ignored if -ccbin or +# --compiler-bindir is already present in the CUDA_NVCC_FLAGS or +# CUDA_NVCC_FLAGS_ variables. For Visual Studio targets, +# the host compiler is constructed with one or more visual studio macros +# such as $(VCInstallDir), that expands out to the path when +# the command is run from within VS. +# If the CUDAHOSTCXX environment variable is set it will +# be used as the default. +# +# CUDA_NVCC_FLAGS +# CUDA_NVCC_FLAGS_ +# -- Additional NVCC command line arguments. NOTE: multiple arguments must be +# semi-colon delimited (e.g. --compiler-options;-Wall) +# +# CUDA_PROPAGATE_HOST_FLAGS (Default ON) +# -- Set to ON to propagate CMAKE_{C,CXX}_FLAGS and their configuration +# dependent counterparts (e.g. CMAKE_C_FLAGS_DEBUG) automatically to the +# host compiler through nvcc's -Xcompiler flag. This helps make the +# generated host code match the rest of the system better. Sometimes +# certain flags give nvcc problems, and this will help you turn the flag +# propagation off. This does not affect the flags supplied directly to nvcc +# via CUDA_NVCC_FLAGS or through the OPTION flags specified through +# CUDA_ADD_LIBRARY, CUDA_ADD_EXECUTABLE, or CUDA_WRAP_SRCS. Flags used for +# shared library compilation are not affected by this flag. +# +# CUDA_PROPAGATE_HOST_FLAGS_BLACKLIST (Default "") +# -- A list containing the host flags that should not be propagated when +# CUDA_PROPAGATE_HOST_FLAGS is ON. +# +# CUDA_SEPARABLE_COMPILATION (Default OFF) +# -- If set this will enable separable compilation for all CUDA runtime object +# files. If used outside of CUDA_ADD_EXECUTABLE and CUDA_ADD_LIBRARY +# (e.g. calling CUDA_WRAP_SRCS directly), +# CUDA_COMPUTE_SEPARABLE_COMPILATION_OBJECT_FILE_NAME and +# CUDA_LINK_SEPARABLE_COMPILATION_OBJECTS should be called. +# +# CUDA_SOURCE_PROPERTY_FORMAT +# -- If this source file property is set, it can override the format specified +# to CUDA_WRAP_SRCS (OBJ, PTX, CUBIN, or FATBIN). If an input source file +# is not a .cu file, setting this file will cause it to be treated as a .cu +# file. See documentation for set_source_files_properties on how to set +# this property. +# +# CUDA_USE_STATIC_CUDA_RUNTIME (Default ON) +# -- When enabled the static version of the CUDA runtime library will be used +# in CUDA_LIBRARIES. If the version of CUDA configured doesn't support +# this option, then it will be silently disabled. +# +# CUDA_VERBOSE_BUILD (Default OFF) +# -- Set to ON to see all the commands used when building the CUDA file. When +# using a Makefile generator the value defaults to VERBOSE (run make +# VERBOSE=1 to see output), although setting CUDA_VERBOSE_BUILD to ON will +# always print the output. +# +# The script creates the following macros (in alphebetical order):: +# +# CUDA_ADD_CUFFT_TO_TARGET( cuda_target ) +# -- Adds the cufft library to the target (can be any target). Handles whether +# you are in emulation mode or not. +# +# CUDA_ADD_CUBLAS_TO_TARGET( cuda_target ) +# -- Adds the cublas library to the target (can be any target). Handles +# whether you are in emulation mode or not. +# +# CUDA_ADD_EXECUTABLE( cuda_target file0 file1 ... +# [WIN32] [MACOSX_BUNDLE] [EXCLUDE_FROM_ALL] [OPTIONS ...] ) +# -- Creates an executable "cuda_target" which is made up of the files +# specified. All of the non CUDA C files are compiled using the standard +# build rules specified by CMAKE and the cuda files are compiled to object +# files using nvcc and the host compiler. In addition CUDA_INCLUDE_DIRS is +# added automatically to include_directories(). Some standard CMake target +# calls can be used on the target after calling this macro +# (e.g. set_target_properties and target_link_libraries), but setting +# properties that adjust compilation flags will not affect code compiled by +# nvcc. Such flags should be modified before calling CUDA_ADD_EXECUTABLE, +# CUDA_ADD_LIBRARY or CUDA_WRAP_SRCS. +# +# CUDA_ADD_LIBRARY( cuda_target file0 file1 ... +# [STATIC | SHARED | MODULE] [EXCLUDE_FROM_ALL] [OPTIONS ...] ) +# -- Same as CUDA_ADD_EXECUTABLE except that a library is created. +# +# CUDA_BUILD_CLEAN_TARGET() +# -- Creates a convenience target that deletes all the dependency files +# generated. You should make clean after running this target to ensure the +# dependency files get regenerated. +# +# CUDA_COMPILE( generated_files file0 file1 ... [STATIC | SHARED | MODULE] +# [OPTIONS ...] ) +# -- Returns a list of generated files from the input source files to be used +# with ADD_LIBRARY or ADD_EXECUTABLE. +# +# CUDA_COMPILE_PTX( generated_files file0 file1 ... [OPTIONS ...] ) +# -- Returns a list of PTX files generated from the input source files. +# +# CUDA_COMPILE_FATBIN( generated_files file0 file1 ... [OPTIONS ...] ) +# -- Returns a list of FATBIN files generated from the input source files. +# +# CUDA_COMPILE_CUBIN( generated_files file0 file1 ... [OPTIONS ...] ) +# -- Returns a list of CUBIN files generated from the input source files. +# +# CUDA_COMPUTE_SEPARABLE_COMPILATION_OBJECT_FILE_NAME( output_file_var +# cuda_target +# object_files ) +# -- Compute the name of the intermediate link file used for separable +# compilation. This file name is typically passed into +# CUDA_LINK_SEPARABLE_COMPILATION_OBJECTS. output_file_var is produced +# based on cuda_target the list of objects files that need separable +# compilation as specified by object_files. If the object_files list is +# empty, then output_file_var will be empty. This function is called +# automatically for CUDA_ADD_LIBRARY and CUDA_ADD_EXECUTABLE. Note that +# this is a function and not a macro. +# +# CUDA_INCLUDE_DIRECTORIES( path0 path1 ... ) +# -- Sets the directories that should be passed to nvcc +# (e.g. nvcc -Ipath0 -Ipath1 ... ). These paths usually contain other .cu +# files. +# +# +# CUDA_LINK_SEPARABLE_COMPILATION_OBJECTS( output_file_var cuda_target +# nvcc_flags object_files) +# -- Generates the link object required by separable compilation from the given +# object files. This is called automatically for CUDA_ADD_EXECUTABLE and +# CUDA_ADD_LIBRARY, but can be called manually when using CUDA_WRAP_SRCS +# directly. When called from CUDA_ADD_LIBRARY or CUDA_ADD_EXECUTABLE the +# nvcc_flags passed in are the same as the flags passed in via the OPTIONS +# argument. The only nvcc flag added automatically is the bitness flag as +# specified by CUDA_64_BIT_DEVICE_CODE. Note that this is a function +# instead of a macro. +# +# CUDA_SELECT_NVCC_ARCH_FLAGS(out_variable [target_CUDA_architectures]) +# -- Selects GPU arch flags for nvcc based on target_CUDA_architectures +# target_CUDA_architectures : Auto | Common | All | LIST(ARCH_AND_PTX ...) +# - "Auto" detects local machine GPU compute arch at runtime. +# - "Common" and "All" cover common and entire subsets of architectures +# ARCH_AND_PTX : NAME | NUM.NUM | NUM.NUM(NUM.NUM) | NUM.NUM+PTX +# NAME: Kepler Maxwell Kepler+Tesla Maxwell+Tegra Pascal Volta Turing +# NUM: Any number. Only those pairs are currently accepted by NVCC though: +# 3.5 3.7 5.0 5.2 5.3 6.0 6.1 6.2 7.0 7.2 7.5 +# Returns LIST of flags to be added to CUDA_NVCC_FLAGS in ${out_variable} +# Additionally, sets ${out_variable}_readable to the resulting numeric list +# Example: +# CUDA_SELECT_NVCC_ARCH_FLAGS(ARCH_FLAGS 3.0 3.5+PTX 5.2(5.0) Maxwell) +# LIST(APPEND CUDA_NVCC_FLAGS ${ARCH_FLAGS}) +# +# More info on CUDA architectures: https://en.wikipedia.org/wiki/CUDA +# Note that this is a function instead of a macro. +# +# CUDA_WRAP_SRCS ( cuda_target format generated_files file0 file1 ... +# [STATIC | SHARED | MODULE] [OPTIONS ...] ) +# -- This is where all the magic happens. CUDA_ADD_EXECUTABLE, +# CUDA_ADD_LIBRARY, CUDA_COMPILE, and CUDA_COMPILE_PTX all call this +# function under the hood. +# +# Given the list of files (file0 file1 ... fileN) this macro generates +# custom commands that generate either PTX or linkable objects (use "PTX" or +# "OBJ" for the format argument to switch). Files that don't end with .cu +# or have the HEADER_FILE_ONLY property are ignored. +# +# The arguments passed in after OPTIONS are extra command line options to +# give to nvcc. You can also specify per configuration options by +# specifying the name of the configuration followed by the options. General +# options must precede configuration specific options. Not all +# configurations need to be specified, only the ones provided will be used. +# +# OPTIONS -DFLAG=2 "-DFLAG_OTHER=space in flag" +# DEBUG -g +# RELEASE --use_fast_math +# RELWITHDEBINFO --use_fast_math;-g +# MINSIZEREL --use_fast_math +# +# For certain configurations (namely VS generating object files with +# CUDA_ATTACH_VS_BUILD_RULE_TO_CUDA_FILE set to ON), no generated file will +# be produced for the given cuda file. This is because when you add the +# cuda file to Visual Studio it knows that this file produces an object file +# and will link in the resulting object file automatically. +# +# This script will also generate a separate cmake script that is used at +# build time to invoke nvcc. This is for several reasons. +# +# 1. nvcc can return negative numbers as return values which confuses +# Visual Studio into thinking that the command succeeded. The script now +# checks the error codes and produces errors when there was a problem. +# +# 2. nvcc has been known to not delete incomplete results when it +# encounters problems. This confuses build systems into thinking the +# target was generated when in fact an unusable file exists. The script +# now deletes the output files if there was an error. +# +# 3. By putting all the options that affect the build into a file and then +# make the build rule dependent on the file, the output files will be +# regenerated when the options change. +# +# This script also looks at optional arguments STATIC, SHARED, or MODULE to +# determine when to target the object compilation for a shared library. +# BUILD_SHARED_LIBS is ignored in CUDA_WRAP_SRCS, but it is respected in +# CUDA_ADD_LIBRARY. On some systems special flags are added for building +# objects intended for shared libraries. A preprocessor macro, +# _EXPORTS is defined when a shared library compilation is +# detected. +# +# Flags passed into add_definitions with -D or /D are passed along to nvcc. +# +# +# +# The script defines the following variables:: +# +# CUDA_VERSION_MAJOR -- The major version of cuda as reported by nvcc. +# CUDA_VERSION_MINOR -- The minor version. +# CUDA_VERSION +# CUDA_VERSION_STRING -- CUDA_VERSION_MAJOR.CUDA_VERSION_MINOR +# CUDA_HAS_FP16 -- Whether a short float (float16,fp16) is supported. +# +# CUDA_TOOLKIT_ROOT_DIR -- Path to the CUDA Toolkit (defined if not set). +# CUDA_SDK_ROOT_DIR -- Path to the CUDA SDK. Use this to find files in the +# SDK. This script will not directly support finding +# specific libraries or headers, as that isn't +# supported by NVIDIA. If you want to change +# libraries when the path changes see the +# FindCUDA.cmake script for an example of how to clear +# these variables. There are also examples of how to +# use the CUDA_SDK_ROOT_DIR to locate headers or +# libraries, if you so choose (at your own risk). +# CUDA_INCLUDE_DIRS -- Include directory for cuda headers. Added automatically +# for CUDA_ADD_EXECUTABLE and CUDA_ADD_LIBRARY. +# CUDA_LIBRARIES -- Cuda RT library. +# CUDA_CUFFT_LIBRARIES -- Device or emulation library for the Cuda FFT +# implementation (alternative to: +# CUDA_ADD_CUFFT_TO_TARGET macro) +# CUDA_CUBLAS_LIBRARIES -- Device or emulation library for the Cuda BLAS +# implementation (alternative to: +# CUDA_ADD_CUBLAS_TO_TARGET macro). +# CUDA_cudart_static_LIBRARY -- Statically linkable cuda runtime library. +# Only available for CUDA version 5.5+ +# CUDA_cudadevrt_LIBRARY -- Device runtime library. +# Required for separable compilation. +# CUDA_cupti_LIBRARY -- CUDA Profiling Tools Interface library. +# Only available for CUDA version 4.0+. +# CUDA_curand_LIBRARY -- CUDA Random Number Generation library. +# Only available for CUDA version 3.2+. +# CUDA_cusolver_LIBRARY -- CUDA Direct Solver library. +# Only available for CUDA version 7.0+. +# CUDA_cusparse_LIBRARY -- CUDA Sparse Matrix library. +# Only available for CUDA version 3.2+. +# CUDA_npp_LIBRARY -- NVIDIA Performance Primitives lib. +# Only available for CUDA version 4.0+. +# CUDA_nppc_LIBRARY -- NVIDIA Performance Primitives lib (core). +# Only available for CUDA version 5.5+. +# CUDA_nppi_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 5.5 - 8.0. +# CUDA_nppial_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_nppicc_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_nppicom_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_nppidei_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_nppif_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_nppig_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_nppim_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_nppist_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_nppisu_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_nppitc_LIBRARY -- NVIDIA Performance Primitives lib (image processing). +# Only available for CUDA version 9.0. +# CUDA_npps_LIBRARY -- NVIDIA Performance Primitives lib (signal processing). +# Only available for CUDA version 5.5+. +# CUDA_nvcuvenc_LIBRARY -- CUDA Video Encoder library. +# Only available for CUDA version 3.2+. +# Windows only. +# CUDA_nvcuvid_LIBRARY -- CUDA Video Decoder library. +# Only available for CUDA version 3.2+. +# Windows only. +# + +# James Bigler, NVIDIA Corp (nvidia.com - jbigler) +# Abe Stephens, SCI Institute -- http://www.sci.utah.edu/~abe/FindCuda.html +# +# Copyright (c) 2008 - 2009 NVIDIA Corporation. All rights reserved. +# +# Copyright (c) 2007-2009 +# Scientific Computing and Imaging Institute, University of Utah +# +# This code is licensed under the MIT License. See the FindCUDA.cmake script +# for the text of the license. + +# The MIT License +# +# License for the specific language governing rights and limitations under +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included +# in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# +############################################################################### + +# FindCUDA.cmake + +include(FindPackageHandleStandardArgs) +# This macro helps us find the location of helper files we will need the full path to +macro(CUDA_FIND_HELPER_FILE _name _extension) + set(_full_name "${_name}.${_extension}") + # CMAKE_CURRENT_LIST_FILE contains the full path to the file currently being + # processed. Using this variable, we can pull out the current path, and + # provide a way to get access to the other files we need local to here. + get_filename_component(CMAKE_CURRENT_LIST_DIR "${CMAKE_CURRENT_LIST_FILE}" PATH) + set(CUDA_${_name} "${CMAKE_CURRENT_LIST_DIR}/FindCUDA/${_full_name}") + if(NOT EXISTS "${CUDA_${_name}}") + set(error_message "${_full_name} not found in ${CMAKE_CURRENT_LIST_DIR}/FindCUDA") + if(CUDA_FIND_REQUIRED) + message(FATAL_ERROR "${error_message}") + else() + if(NOT CUDA_FIND_QUIETLY) + message(STATUS "${error_message}") + endif() + endif() + endif() + # Set this variable as internal, so the user isn't bugged with it. + set(CUDA_${_name} ${CUDA_${_name}} CACHE INTERNAL "Location of ${_full_name}" FORCE) +endmacro() + +##################################################################### +## CUDA_INCLUDE_NVCC_DEPENDENCIES +## + +# So we want to try and include the dependency file if it exists. If +# it doesn't exist then we need to create an empty one, so we can +# include it. + +# If it does exist, then we need to check to see if all the files it +# depends on exist. If they don't then we should clear the dependency +# file and regenerate it later. This covers the case where a header +# file has disappeared or moved. + +macro(CUDA_INCLUDE_NVCC_DEPENDENCIES dependency_file) + set(CUDA_NVCC_DEPEND) + set(CUDA_NVCC_DEPEND_REGENERATE FALSE) + + + # Include the dependency file. Create it first if it doesn't exist . The + # INCLUDE puts a dependency that will force CMake to rerun and bring in the + # new info when it changes. DO NOT REMOVE THIS (as I did and spent a few + # hours figuring out why it didn't work. + if(NOT EXISTS ${dependency_file}) + file(WRITE ${dependency_file} "#FindCUDA.cmake generated file. Do not edit.\n") + endif() + # Always include this file to force CMake to run again next + # invocation and rebuild the dependencies. + #message("including dependency_file = ${dependency_file}") + include(${dependency_file}) + + # Now we need to verify the existence of all the included files + # here. If they aren't there we need to just blank this variable and + # make the file regenerate again. +# if(DEFINED CUDA_NVCC_DEPEND) +# message("CUDA_NVCC_DEPEND set") +# else() +# message("CUDA_NVCC_DEPEND NOT set") +# endif() + if(CUDA_NVCC_DEPEND) + #message("CUDA_NVCC_DEPEND found") + foreach(f ${CUDA_NVCC_DEPEND}) + # message("searching for ${f}") + if(NOT EXISTS ${f}) + #message("file ${f} not found") + set(CUDA_NVCC_DEPEND_REGENERATE TRUE) + endif() + endforeach() + else() + #message("CUDA_NVCC_DEPEND false") + # No dependencies, so regenerate the file. + set(CUDA_NVCC_DEPEND_REGENERATE TRUE) + endif() + + #message("CUDA_NVCC_DEPEND_REGENERATE = ${CUDA_NVCC_DEPEND_REGENERATE}") + # No incoming dependencies, so we need to generate them. Make the + # output depend on the dependency file itself, which should cause the + # rule to re-run. + if(CUDA_NVCC_DEPEND_REGENERATE) + set(CUDA_NVCC_DEPEND ${dependency_file}) + #message("Generating an empty dependency_file: ${dependency_file}") + file(WRITE ${dependency_file} "#FindCUDA.cmake generated file. Do not edit.\n") + endif() + +endmacro() + +############################################################################### +############################################################################### +# Setup variables' defaults +############################################################################### +############################################################################### + +# Allow the user to specify if the device code is supposed to be 32 or 64 bit. +if(CMAKE_SIZEOF_VOID_P EQUAL 8) + set(CUDA_64_BIT_DEVICE_CODE_DEFAULT ON) +else() + set(CUDA_64_BIT_DEVICE_CODE_DEFAULT OFF) +endif() +option(CUDA_64_BIT_DEVICE_CODE "Compile device code in 64 bit mode" ${CUDA_64_BIT_DEVICE_CODE_DEFAULT}) + +# Attach the build rule to the source file in VS. This option +option(CUDA_ATTACH_VS_BUILD_RULE_TO_CUDA_FILE "Attach the build rule to the CUDA source file. Enable only when the CUDA source file is added to at most one target." ON) + +# Prints out extra information about the cuda file during compilation +option(CUDA_BUILD_CUBIN "Generate and parse .cubin files in Device mode." OFF) + +# Set whether we are using emulation or device mode. +option(CUDA_BUILD_EMULATION "Build in Emulation mode" OFF) + +# Where to put the generated output. +set(CUDA_GENERATED_OUTPUT_DIR "" CACHE PATH "Directory to put all the output files. If blank it will default to the CMAKE_CURRENT_BINARY_DIR") + +# Parse HOST_COMPILATION mode. +option(CUDA_HOST_COMPILATION_CPP "Generated file extension" ON) + +# Extra user settable flags +cmake_initialize_per_config_variable(CUDA_NVCC_FLAGS "Semi-colon delimit multiple arguments.") + +if(DEFINED ENV{CUDAHOSTCXX}) + set(CUDA_HOST_COMPILER "$ENV{CUDAHOSTCXX}" CACHE FILEPATH "Host side compiler used by NVCC") +elseif(CMAKE_GENERATOR MATCHES "Visual Studio") + set(_CUDA_MSVC_HOST_COMPILER "$(VCInstallDir)Tools/MSVC/$(VCToolsVersion)/bin/Host$(Platform)/$(PlatformTarget)") + if(MSVC_VERSION LESS 1910) + set(_CUDA_MSVC_HOST_COMPILER "$(VCInstallDir)bin") + endif() + + set(CUDA_HOST_COMPILER "${_CUDA_MSVC_HOST_COMPILER}" CACHE FILEPATH "Host side compiler used by NVCC") + +else() + if(APPLE + AND "${CMAKE_C_COMPILER_ID}" MATCHES "Clang" + AND "${CMAKE_C_COMPILER}" MATCHES "/cc$") + # Using cc which is symlink to clang may let NVCC think it is GCC and issue + # unhandled -dumpspecs option to clang. Also in case neither + # CMAKE_C_COMPILER is defined (project does not use C language) nor + # CUDA_HOST_COMPILER is specified manually we should skip -ccbin and let + # nvcc use its own default C compiler. + # Only care about this on APPLE with clang to avoid + # following symlinks to things like ccache + if(DEFINED CMAKE_C_COMPILER AND NOT DEFINED CUDA_HOST_COMPILER) + get_filename_component(c_compiler_realpath "${CMAKE_C_COMPILER}" REALPATH) + # if the real path does not end up being clang then + # go back to using CMAKE_C_COMPILER + if(NOT "${c_compiler_realpath}" MATCHES "/clang$") + set(c_compiler_realpath "${CMAKE_C_COMPILER}") + endif() + else() + set(c_compiler_realpath "") + endif() + set(CUDA_HOST_COMPILER "${c_compiler_realpath}" CACHE FILEPATH "Host side compiler used by NVCC") + elseif(MSVC AND "${CMAKE_C_COMPILER}" MATCHES "clcache|sccache") + # NVCC does not think it will work if it is passed clcache.exe or sccache.exe + # as the host compiler, which means that builds with CC=cl.exe won't work. + # Best to just feed it whatever the actual cl.exe is as the host compiler. + set(CUDA_HOST_COMPILER "cl.exe" CACHE FILEPATH "Host side compiler used by NVCC") + else() + set(CUDA_HOST_COMPILER "${CMAKE_C_COMPILER}" + CACHE FILEPATH "Host side compiler used by NVCC") + endif() +endif() + +# Propagate the host flags to the host compiler via -Xcompiler +option(CUDA_PROPAGATE_HOST_FLAGS "Propagate C/CXX_FLAGS and friends to the host compiler via -Xcompile" ON) + +# Blacklisted flags to prevent propagation +set(CUDA_PROPAGATE_HOST_FLAGS_BLACKLIST "" CACHE STRING "Blacklisted flags to prevent propagation") + +# Enable CUDA_SEPARABLE_COMPILATION +option(CUDA_SEPARABLE_COMPILATION "Compile CUDA objects with separable compilation enabled. Requires CUDA 5.0+" OFF) + +# Specifies whether the commands used when compiling the .cu file will be printed out. +option(CUDA_VERBOSE_BUILD "Print out the commands run while compiling the CUDA source file. With the Makefile generator this defaults to VERBOSE variable specified on the command line, but can be forced on with this option." OFF) + +mark_as_advanced( + CUDA_64_BIT_DEVICE_CODE + CUDA_ATTACH_VS_BUILD_RULE_TO_CUDA_FILE + CUDA_GENERATED_OUTPUT_DIR + CUDA_HOST_COMPILATION_CPP + CUDA_NVCC_FLAGS + CUDA_PROPAGATE_HOST_FLAGS + CUDA_PROPAGATE_HOST_FLAGS_BLACKLIST + CUDA_BUILD_CUBIN + CUDA_BUILD_EMULATION + CUDA_VERBOSE_BUILD + CUDA_SEPARABLE_COMPILATION + ) + +# Single config generators like Makefiles or Ninja don't usually have +# CMAKE_CONFIGURATION_TYPES defined (but note that it can be defined if set by +# projects or developers). Even CMAKE_BUILD_TYPE might not be defined for +# single config generators (and should not be defined for multi-config +# generators). To ensure we get a complete superset of all possible +# configurations, we combine CMAKE_CONFIGURATION_TYPES, CMAKE_BUILD_TYPE and +# all of the standard configurations, then weed out duplicates with +# list(REMOVE_DUPLICATES). Looping over the unique set then ensures we have +# each configuration-specific set of nvcc flags defined and marked as advanced. +set(CUDA_configuration_types ${CMAKE_CONFIGURATION_TYPES} ${CMAKE_BUILD_TYPE} Debug MinSizeRel Release RelWithDebInfo) +list(REMOVE_DUPLICATES CUDA_configuration_types) + +############################################################################### +############################################################################### +# Locate CUDA, Set Build Type, etc. +############################################################################### +############################################################################### + +macro(cuda_unset_include_and_libraries) + unset(CUDA_TOOLKIT_INCLUDE CACHE) + unset(CUDA_CUDART_LIBRARY CACHE) + unset(CUDA_CUDA_LIBRARY CACHE) + # Make sure you run this before you unset CUDA_VERSION. + unset(CUDA_cudart_static_LIBRARY CACHE) + unset(CUDA_cudadevrt_LIBRARY CACHE) + unset(CUDA_cublas_LIBRARY CACHE) + unset(CUDA_cublas_device_LIBRARY CACHE) + unset(CUDA_cublasemu_LIBRARY CACHE) + unset(CUDA_cublasLt_LIBRARY CACHE) + unset(CUDA_cufft_LIBRARY CACHE) + unset(CUDA_cufftemu_LIBRARY CACHE) + unset(CUDA_cupti_LIBRARY CACHE) + unset(CUDA_curand_LIBRARY CACHE) + unset(CUDA_cusolver_LIBRARY CACHE) + unset(CUDA_cusparse_LIBRARY CACHE) + unset(CUDA_npp_LIBRARY CACHE) + unset(CUDA_nppc_LIBRARY CACHE) + unset(CUDA_nppi_LIBRARY CACHE) + unset(CUDA_npps_LIBRARY CACHE) + unset(CUDA_nvcuvenc_LIBRARY CACHE) + unset(CUDA_nvcuvid_LIBRARY CACHE) + unset(CUDA_GPU_DETECT_OUTPUT CACHE) +endmacro() + +# Check to see if the CUDA_TOOLKIT_ROOT_DIR and CUDA_SDK_ROOT_DIR have changed, +# if they have then clear the cache variables, so that will be detected again. +if(NOT "${CUDA_TOOLKIT_ROOT_DIR}" STREQUAL "${CUDA_TOOLKIT_ROOT_DIR_INTERNAL}") + unset(CUDA_TOOLKIT_TARGET_DIR CACHE) + unset(CUDA_NVCC_EXECUTABLE CACHE) + cuda_unset_include_and_libraries() + unset(CUDA_VERSION CACHE) +endif() + +if(NOT "${CUDA_TOOLKIT_TARGET_DIR}" STREQUAL "${CUDA_TOOLKIT_TARGET_DIR_INTERNAL}") + cuda_unset_include_and_libraries() +endif() + +# +# End of unset() +# + +# +# Start looking for things +# + +# Search for the cuda distribution. +if(NOT CUDA_TOOLKIT_ROOT_DIR AND NOT CMAKE_CROSSCOMPILING) + # Search in the CUDA_BIN_PATH first. + find_program(CUDA_TOOLKIT_ROOT_DIR_NVCC + NAMES nvcc nvcc.exe + PATHS + ENV CUDA_TOOLKIT_ROOT + ENV CUDA_PATH + ENV CUDA_BIN_PATH + PATH_SUFFIXES bin bin64 + DOC "Toolkit location." + NO_DEFAULT_PATH + ) + + # Now search default paths + find_program(CUDA_TOOLKIT_ROOT_DIR_NVCC + NAMES nvcc nvcc.exe + PATHS /opt/cuda/bin + PATH_SUFFIXES cuda/bin + DOC "Toolkit location." + ) + + if (CUDA_TOOLKIT_ROOT_DIR_NVCC) + get_filename_component(CUDA_TOOLKIT_ROOT_DIR_NVCC_PAR "${CUDA_TOOLKIT_ROOT_DIR_NVCC}" DIRECTORY) + get_filename_component(CUDA_TOOLKIT_ROOT_DIR "${CUDA_TOOLKIT_ROOT_DIR_NVCC_PAR}" DIRECTORY CACHE) + string(REGEX REPLACE "[/\\\\]?bin[64]*[/\\\\]?$" "" CUDA_TOOLKIT_ROOT_DIR ${CUDA_TOOLKIT_ROOT_DIR}) + # We need to force this back into the cache. + set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_TOOLKIT_ROOT_DIR} CACHE PATH "Toolkit location." FORCE) + set(CUDA_TOOLKIT_TARGET_DIR ${CUDA_TOOLKIT_ROOT_DIR}) + endif() + unset(CUDA_TOOLKIT_ROOT_DIR_NVCC CACHE) + + if (NOT EXISTS ${CUDA_TOOLKIT_ROOT_DIR}) + if(CUDA_FIND_REQUIRED) + message(FATAL_ERROR "Specify CUDA_TOOLKIT_ROOT_DIR") + elseif(NOT CUDA_FIND_QUIETLY) + message("CUDA_TOOLKIT_ROOT_DIR not found or specified") + endif() + endif () +endif () + +if(CMAKE_CROSSCOMPILING) + SET (CUDA_TOOLKIT_ROOT $ENV{CUDA_TOOLKIT_ROOT}) + if(CMAKE_SYSTEM_PROCESSOR STREQUAL "armv7-a") + # Support for NVPACK + set (CUDA_TOOLKIT_TARGET_NAMES "armv7-linux-androideabi") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "arm") + # Support for arm cross compilation + set(CUDA_TOOLKIT_TARGET_NAMES "armv7-linux-gnueabihf") + elseif(CMAKE_SYSTEM_PROCESSOR MATCHES "aarch64") + # Support for aarch64 cross compilation + if (ANDROID_ARCH_NAME STREQUAL "arm64") + set(CUDA_TOOLKIT_TARGET_NAMES "aarch64-linux-androideabi") + else() + set(CUDA_TOOLKIT_TARGET_NAMES "aarch64-linux" "sbsa-linux") + endif (ANDROID_ARCH_NAME STREQUAL "arm64") + endif() + + foreach(CUDA_TOOLKIT_TARGET_NAME IN LISTS CUDA_TOOLKIT_TARGET_NAMES) + if (EXISTS "${CUDA_TOOLKIT_ROOT}/targets/${CUDA_TOOLKIT_TARGET_NAME}") + set(CUDA_TOOLKIT_TARGET_DIR "${CUDA_TOOLKIT_ROOT}/targets/${CUDA_TOOLKIT_TARGET_NAME}" CACHE PATH "CUDA Toolkit target location.") + SET (CUDA_TOOLKIT_ROOT_DIR ${CUDA_TOOLKIT_ROOT} CACHE PATH "Toolkit location." FORCE) + mark_as_advanced(CUDA_TOOLKIT_TARGET_DIR) + break() + endif() + endforeach() + + # add known CUDA targetr root path to the set of directories we search for programs, libraries and headers + set( CMAKE_FIND_ROOT_PATH "${CUDA_TOOLKIT_TARGET_DIR};${CMAKE_FIND_ROOT_PATH}") + macro( cuda_find_host_program ) + if (COMMAND find_host_program) + find_host_program( ${ARGN} ) + else() + find_program( ${ARGN} ) + endif() + endmacro() +else() + # for non-cross-compile, find_host_program == find_program and CUDA_TOOLKIT_TARGET_DIR == CUDA_TOOLKIT_ROOT_DIR + macro( cuda_find_host_program ) + find_program( ${ARGN} ) + endmacro() + SET (CUDA_TOOLKIT_TARGET_DIR ${CUDA_TOOLKIT_ROOT_DIR}) +endif() + + +# CUDA_NVCC_EXECUTABLE +if(DEFINED ENV{CUDA_NVCC_EXECUTABLE}) + set(CUDA_NVCC_EXECUTABLE "$ENV{CUDA_NVCC_EXECUTABLE}" CACHE FILEPATH "The CUDA compiler") +else() + cuda_find_host_program(CUDA_NVCC_EXECUTABLE + NAMES nvcc + PATHS "${CUDA_TOOLKIT_ROOT_DIR}" + ENV CUDA_PATH + ENV CUDA_BIN_PATH + PATH_SUFFIXES bin bin64 + NO_DEFAULT_PATH + ) + # Search default search paths, after we search our own set of paths. + cuda_find_host_program(CUDA_NVCC_EXECUTABLE nvcc) +endif() + +if(CUDA_NVCC_EXECUTABLE AND NOT CUDA_VERSION) + # Compute the version. + execute_process(COMMAND ${CUDA_NVCC_EXECUTABLE} "--version" + OUTPUT_VARIABLE NVCC_OUT + RESULT_VARIABLE NVCC_RC) + if(NOT (${NVCC_RC} EQUAL 0)) + message(WARNING "Failed to execute '${CUDA_NVCC_EXECUTABLE} --version'") + set(CUDA_FOUND FALSE) + return() + endif() + string(REGEX REPLACE ".*release ([0-9]+)\\.([0-9]+).*" "\\1" CUDA_VERSION_MAJOR ${NVCC_OUT}) + string(REGEX REPLACE ".*release ([0-9]+)\\.([0-9]+).*" "\\2" CUDA_VERSION_MINOR ${NVCC_OUT}) + set(CUDA_VERSION "${CUDA_VERSION_MAJOR}.${CUDA_VERSION_MINOR}" CACHE STRING "Version of CUDA as computed from nvcc.") + mark_as_advanced(CUDA_VERSION) +else() + # Need to set these based off of the cached value + string(REGEX REPLACE "([0-9]+)\\.([0-9]+).*" "\\1" CUDA_VERSION_MAJOR "${CUDA_VERSION}") + string(REGEX REPLACE "([0-9]+)\\.([0-9]+).*" "\\2" CUDA_VERSION_MINOR "${CUDA_VERSION}") +endif() + +# Always set this convenience variable +set(CUDA_VERSION_STRING "${CUDA_VERSION}") + +# CUDA_TOOLKIT_INCLUDE +find_path(CUDA_TOOLKIT_INCLUDE + device_functions.h # Header included in toolkit + PATHS ${CUDA_TOOLKIT_TARGET_DIR} + ENV CUDA_PATH + ENV CUDA_INC_PATH + PATH_SUFFIXES include + NO_DEFAULT_PATH + ) +# Search default search paths, after we search our own set of paths. +find_path(CUDA_TOOLKIT_INCLUDE device_functions.h) +mark_as_advanced(CUDA_TOOLKIT_INCLUDE) + +set(CUDA_HAS_FP16 TRUE) + +# Set the user list of include dir to nothing to initialize it. +set (CUDA_NVCC_INCLUDE_DIRS_USER "") +set (CUDA_INCLUDE_DIRS ${CUDA_TOOLKIT_INCLUDE}) + +macro(cuda_find_library_local_first_with_path_ext _var _names _doc _path_ext ) + if(CMAKE_SIZEOF_VOID_P EQUAL 8) + # CUDA 3.2+ on Windows moved the library directories, so we need the new + # and old paths. + set(_cuda_64bit_lib_dir "${_path_ext}lib/x64" "${_path_ext}lib64" "${_path_ext}libx64" ) + endif() + # CUDA 3.2+ on Windows moved the library directories, so we need to new + # (lib/Win32) and the old path (lib). + find_library(${_var} + NAMES ${_names} + PATHS "${CUDA_TOOLKIT_TARGET_DIR}" + ENV CUDA_PATH + ENV CUDA_LIB_PATH + PATH_SUFFIXES ${_cuda_64bit_lib_dir} "${_path_ext}lib/Win32" "${_path_ext}lib" "${_path_ext}libWin32" + DOC ${_doc} + NO_DEFAULT_PATH + ) + if (NOT CMAKE_CROSSCOMPILING) + # Search default search paths, after we search our own set of paths. + find_library(${_var} + NAMES ${_names} + PATHS "/usr/lib/nvidia-current" + DOC ${_doc} + ) + endif() +endmacro() + +macro(cuda_find_library_local_first _var _names _doc) + cuda_find_library_local_first_with_path_ext( "${_var}" "${_names}" "${_doc}" "" ) +endmacro() + +macro(find_library_local_first _var _names _doc ) + cuda_find_library_local_first( "${_var}" "${_names}" "${_doc}" "" ) +endmacro() + + +# CUDA_LIBRARIES +cuda_find_library_local_first(CUDA_CUDART_LIBRARY cudart "\"cudart\" library") + +cuda_find_library_local_first(CUDA_cudart_static_LIBRARY cudart_static "static CUDA runtime library") +mark_as_advanced(CUDA_cudart_static_LIBRARY) + + +if(CUDA_cudart_static_LIBRARY) + # If static cudart available, use it by default, but provide a user-visible option to disable it. + option(CUDA_USE_STATIC_CUDA_RUNTIME "Use the static version of the CUDA runtime library if available" ON) +else() + # If not available, silently disable the option. + set(CUDA_USE_STATIC_CUDA_RUNTIME OFF CACHE INTERNAL "") +endif() + +if(CUDA_USE_STATIC_CUDA_RUNTIME) + set(CUDA_CUDART_LIBRARY_VAR CUDA_cudart_static_LIBRARY) +else() + set(CUDA_CUDART_LIBRARY_VAR CUDA_CUDART_LIBRARY) +endif() + +cuda_find_library_local_first(CUDA_cudadevrt_LIBRARY cudadevrt "\"cudadevrt\" library") +mark_as_advanced(CUDA_cudadevrt_LIBRARY) + +if(CUDA_USE_STATIC_CUDA_RUNTIME) + if(UNIX) + # Check for the dependent libraries. Here we look for pthreads. + if (DEFINED CMAKE_THREAD_PREFER_PTHREAD) + set(_cuda_cmake_thread_prefer_pthread ${CMAKE_THREAD_PREFER_PTHREAD}) + endif() + set(CMAKE_THREAD_PREFER_PTHREAD 1) + + # Many of the FindXYZ CMake comes with makes use of try_compile with int main(){return 0;} + # as the source file. Unfortunately this causes a warning with -Wstrict-prototypes and + # -Werror causes the try_compile to fail. We will just temporarily disable other flags + # when doing the find_package command here. + set(_cuda_cmake_c_flags ${CMAKE_C_FLAGS}) + set(CMAKE_C_FLAGS "-fPIC") + find_package(Threads REQUIRED) + set(CMAKE_C_FLAGS ${_cuda_cmake_c_flags}) + + if (DEFINED _cuda_cmake_thread_prefer_pthread) + set(CMAKE_THREAD_PREFER_PTHREAD ${_cuda_cmake_thread_prefer_pthread}) + unset(_cuda_cmake_thread_prefer_pthread) + else() + unset(CMAKE_THREAD_PREFER_PTHREAD) + endif() + + if(NOT APPLE) + #On Linux, you must link against librt when using the static cuda runtime. + find_library(CUDA_rt_LIBRARY rt) + if (NOT CUDA_rt_LIBRARY) + message(WARNING "Expecting to find librt for libcudart_static, but didn't find it.") + endif() + endif() + endif() +endif() + +cuda_find_library_local_first_with_path_ext(CUDA_cupti_LIBRARY cupti "\"cupti\" library" "extras/CUPTI/") +mark_as_advanced(CUDA_cupti_LIBRARY) + +# Set the CUDA_LIBRARIES variable. This is the set of stuff to link against if you are +# using the CUDA runtime. For the dynamic version of the runtime, most of the +# dependencies are brought in, but for the static version there are additional libraries +# and linker commands needed. +# Initialize to empty +set(CUDA_LIBRARIES) + +# If we are using emulation mode and we found the cudartemu library then use +# that one instead of cudart. +if(CUDA_BUILD_EMULATION AND CUDA_CUDARTEMU_LIBRARY) + list(APPEND CUDA_LIBRARIES ${CUDA_CUDARTEMU_LIBRARY}) +elseif(CUDA_USE_STATIC_CUDA_RUNTIME AND CUDA_cudart_static_LIBRARY) + list(APPEND CUDA_LIBRARIES ${CUDA_cudart_static_LIBRARY} ${CMAKE_THREAD_LIBS_INIT} ${CMAKE_DL_LIBS}) + if (CUDA_rt_LIBRARY) + list(APPEND CUDA_LIBRARIES ${CUDA_rt_LIBRARY}) + endif() + if(APPLE) + # We need to add the default path to the driver (libcuda.dylib) as an rpath, so that + # the static cuda runtime can find it at runtime. + list(APPEND CUDA_LIBRARIES -Wl,-rpath,/usr/local/cuda/lib) + endif() +else() + list(APPEND CUDA_LIBRARIES ${CUDA_CUDART_LIBRARY}) +endif() + +# 1.1 toolkit on linux doesn't appear to have a separate library on +# some platforms. +cuda_find_library_local_first(CUDA_CUDA_LIBRARY cuda "\"cuda\" library (older versions only).") + +mark_as_advanced( + CUDA_CUDA_LIBRARY + CUDA_CUDART_LIBRARY + ) + +####################### +# Look for some of the toolkit helper libraries +macro(FIND_CUDA_HELPER_LIBS _name) + cuda_find_library_local_first(CUDA_${_name}_LIBRARY ${_name} "\"${_name}\" library") + mark_as_advanced(CUDA_${_name}_LIBRARY) +endmacro() + +if(CUDA_BUILD_EMULATION) + message(FATAL_ERROR "CUDA_BUILD_EMULATION is not supported in version 3.1 and onwards. You must disable it to proceed. You have version ${CUDA_VERSION}.") +endif() + +find_cuda_helper_libs(cufft) +find_cuda_helper_libs(cublas) +find_cuda_helper_libs(cublasLt) +# cusparse showed up in version 3.2 +find_cuda_helper_libs(cusparse) +find_cuda_helper_libs(curand) +if (WIN32) + find_cuda_helper_libs(nvcuvenc) + find_cuda_helper_libs(nvcuvid) +endif() + +# In CUDA 9.0 NPP was nppi was removed +find_cuda_helper_libs(nppc) +find_cuda_helper_libs(nppial) +find_cuda_helper_libs(nppicc) +find_cuda_helper_libs(nppicom) +find_cuda_helper_libs(nppidei) +find_cuda_helper_libs(nppif) +find_cuda_helper_libs(nppig) +find_cuda_helper_libs(nppim) +find_cuda_helper_libs(nppist) +find_cuda_helper_libs(nppisu) +find_cuda_helper_libs(nppitc) +find_cuda_helper_libs(npps) +set(CUDA_npp_LIBRARY "${CUDA_nppc_LIBRARY};${CUDA_nppial_LIBRARY};${CUDA_nppicc_LIBRARY};${CUDA_nppicom_LIBRARY};${CUDA_nppidei_LIBRARY};${CUDA_nppif_LIBRARY};${CUDA_nppig_LIBRARY};${CUDA_nppim_LIBRARY};${CUDA_nppist_LIBRARY};${CUDA_nppisu_LIBRARY};${CUDA_nppitc_LIBRARY};${CUDA_npps_LIBRARY}") +# cusolver showed up in version 7.0 +find_cuda_helper_libs(cusolver) + +if (CUDA_BUILD_EMULATION) + set(CUDA_CUFFT_LIBRARIES ${CUDA_cufftemu_LIBRARY}) + set(CUDA_CUBLAS_LIBRARIES ${CUDA_cublasemu_LIBRARY}) +else() + set(CUDA_CUFFT_LIBRARIES ${CUDA_cufft_LIBRARY}) + set(CUDA_CUBLAS_LIBRARIES ${CUDA_cublas_LIBRARY} ${CUDA_cublas_device_LIBRARY} ${CUDA_cublasLt_LIBRARY}) +endif() + +######################## +# Look for the SDK stuff. As of CUDA 3.0 NVSDKCUDA_ROOT has been replaced with +# NVSDKCOMPUTE_ROOT with the old CUDA C contents moved into the C subdirectory +find_path(CUDA_SDK_ROOT_DIR common/inc/cutil.h + HINTS + "$ENV{NVSDKCOMPUTE_ROOT}/C" + ENV NVSDKCUDA_ROOT + "[HKEY_LOCAL_MACHINE\\SOFTWARE\\NVIDIA Corporation\\Installed Products\\NVIDIA SDK 10\\Compute;InstallDir]" + PATHS + "/Developer/GPU\ Computing/C" + ) + +# Keep the CUDA_SDK_ROOT_DIR first in order to be able to override the +# environment variables. +set(CUDA_SDK_SEARCH_PATH + "${CUDA_SDK_ROOT_DIR}" + "${CUDA_TOOLKIT_ROOT_DIR}/local/NVSDK0.2" + "${CUDA_TOOLKIT_ROOT_DIR}/NVSDK0.2" + "${CUDA_TOOLKIT_ROOT_DIR}/NV_CUDA_SDK" + "$ENV{HOME}/NVIDIA_CUDA_SDK" + "$ENV{HOME}/NVIDIA_CUDA_SDK_MACOSX" + "/Developer/CUDA" + ) + +# Example of how to find an include file from the CUDA_SDK_ROOT_DIR + +# find_path(CUDA_CUT_INCLUDE_DIR +# cutil.h +# PATHS ${CUDA_SDK_SEARCH_PATH} +# PATH_SUFFIXES "common/inc" +# DOC "Location of cutil.h" +# NO_DEFAULT_PATH +# ) +# # Now search system paths +# find_path(CUDA_CUT_INCLUDE_DIR cutil.h DOC "Location of cutil.h") + +# mark_as_advanced(CUDA_CUT_INCLUDE_DIR) + + +# Example of how to find a library in the CUDA_SDK_ROOT_DIR + +# # cutil library is called cutil64 for 64 bit builds on windows. We don't want +# # to get these confused, so we are setting the name based on the word size of +# # the build. + +# if(CMAKE_SIZEOF_VOID_P EQUAL 8) +# set(cuda_cutil_name cutil64) +# else() +# set(cuda_cutil_name cutil32) +# endif() + +# find_library(CUDA_CUT_LIBRARY +# NAMES cutil ${cuda_cutil_name} +# PATHS ${CUDA_SDK_SEARCH_PATH} +# # The new version of the sdk shows up in common/lib, but the old one is in lib +# PATH_SUFFIXES "common/lib" "lib" +# DOC "Location of cutil library" +# NO_DEFAULT_PATH +# ) +# # Now search system paths +# find_library(CUDA_CUT_LIBRARY NAMES cutil ${cuda_cutil_name} DOC "Location of cutil library") +# mark_as_advanced(CUDA_CUT_LIBRARY) +# set(CUDA_CUT_LIBRARIES ${CUDA_CUT_LIBRARY}) + + + +############################# +# Check for required components +set(CUDA_FOUND TRUE) + +set(CUDA_TOOLKIT_ROOT_DIR_INTERNAL "${CUDA_TOOLKIT_ROOT_DIR}" CACHE INTERNAL + "This is the value of the last time CUDA_TOOLKIT_ROOT_DIR was set successfully." FORCE) +set(CUDA_TOOLKIT_TARGET_DIR_INTERNAL "${CUDA_TOOLKIT_TARGET_DIR}" CACHE INTERNAL + "This is the value of the last time CUDA_TOOLKIT_TARGET_DIR was set successfully." FORCE) +set(CUDA_SDK_ROOT_DIR_INTERNAL "${CUDA_SDK_ROOT_DIR}" CACHE INTERNAL + "This is the value of the last time CUDA_SDK_ROOT_DIR was set successfully." FORCE) + +find_package_handle_standard_args(CUDA + REQUIRED_VARS + CUDA_TOOLKIT_ROOT_DIR + CUDA_NVCC_EXECUTABLE + CUDA_INCLUDE_DIRS + ${CUDA_CUDART_LIBRARY_VAR} + VERSION_VAR + CUDA_VERSION + ) + + + +############################################################################### +############################################################################### +# Macros +############################################################################### +############################################################################### + +############################################################################### +# Add include directories to pass to the nvcc command. +macro(CUDA_INCLUDE_DIRECTORIES) + foreach(dir ${ARGN}) + list(APPEND CUDA_NVCC_INCLUDE_DIRS_USER ${dir}) + endforeach() +endmacro() + + +############################################################################## +cuda_find_helper_file(parse_cubin cmake) +cuda_find_helper_file(make2cmake cmake) +cuda_find_helper_file(run_nvcc cmake) +include("${CMAKE_CURRENT_LIST_DIR}/FindCUDA/select_compute_arch.cmake") + +############################################################################## +# Separate the OPTIONS out from the sources +# +macro(CUDA_GET_SOURCES_AND_OPTIONS _sources _cmake_options _options) + set( ${_sources} ) + set( ${_cmake_options} ) + set( ${_options} ) + set( _found_options FALSE ) + foreach(arg ${ARGN}) + if("x${arg}" STREQUAL "xOPTIONS") + set( _found_options TRUE ) + elseif( + "x${arg}" STREQUAL "xWIN32" OR + "x${arg}" STREQUAL "xMACOSX_BUNDLE" OR + "x${arg}" STREQUAL "xEXCLUDE_FROM_ALL" OR + "x${arg}" STREQUAL "xSTATIC" OR + "x${arg}" STREQUAL "xSHARED" OR + "x${arg}" STREQUAL "xMODULE" + ) + list(APPEND ${_cmake_options} ${arg}) + else() + if ( _found_options ) + list(APPEND ${_options} ${arg}) + else() + # Assume this is a file + list(APPEND ${_sources} ${arg}) + endif() + endif() + endforeach() +endmacro() + +############################################################################## +# Parse the OPTIONS from ARGN and set the variables prefixed by _option_prefix +# +macro(CUDA_PARSE_NVCC_OPTIONS _option_prefix) + set( _found_config ) + foreach(arg ${ARGN}) + # Determine if we are dealing with a perconfiguration flag + foreach(config ${CUDA_configuration_types}) + string(TOUPPER ${config} config_upper) + if (arg STREQUAL "${config_upper}") + set( _found_config _${arg}) + # Set arg to nothing to keep it from being processed further + set( arg ) + endif() + endforeach() + + if ( arg ) + list(APPEND ${_option_prefix}${_found_config} "${arg}") + endif() + endforeach() +endmacro() + +############################################################################## +# Helper to add the include directory for CUDA only once +function(CUDA_ADD_CUDA_INCLUDE_ONCE) + get_directory_property(_include_directories INCLUDE_DIRECTORIES) + set(_add TRUE) + if(_include_directories) + foreach(dir ${_include_directories}) + if("${dir}" STREQUAL "${CUDA_INCLUDE_DIRS}") + set(_add FALSE) + endif() + endforeach() + endif() + if(_add) + include_directories(${CUDA_INCLUDE_DIRS}) + endif() +endfunction() + +function(CUDA_BUILD_SHARED_LIBRARY shared_flag) + set(cmake_args ${ARGN}) + # If SHARED, MODULE, or STATIC aren't already in the list of arguments, then + # add SHARED or STATIC based on the value of BUILD_SHARED_LIBS. + list(FIND cmake_args SHARED _cuda_found_SHARED) + list(FIND cmake_args MODULE _cuda_found_MODULE) + list(FIND cmake_args STATIC _cuda_found_STATIC) + if( _cuda_found_SHARED GREATER -1 OR + _cuda_found_MODULE GREATER -1 OR + _cuda_found_STATIC GREATER -1) + set(_cuda_build_shared_libs) + else() + if (BUILD_SHARED_LIBS) + set(_cuda_build_shared_libs SHARED) + else() + set(_cuda_build_shared_libs STATIC) + endif() + endif() + set(${shared_flag} ${_cuda_build_shared_libs} PARENT_SCOPE) +endfunction() + +############################################################################## +# Helper to avoid clashes of files with the same basename but different paths. +# This doesn't attempt to do exactly what CMake internals do, which is to only +# add this path when there is a conflict, since by the time a second collision +# in names is detected it's already too late to fix the first one. For +# consistency sake the relative path will be added to all files. +function(CUDA_COMPUTE_BUILD_PATH path build_path) + #message("CUDA_COMPUTE_BUILD_PATH([${path}] ${build_path})") + # Only deal with CMake style paths from here on out + file(TO_CMAKE_PATH "${path}" bpath) + if (IS_ABSOLUTE "${bpath}") + # Absolute paths are generally unnecessary, especially if something like + # file(GLOB_RECURSE) is used to pick up the files. + + string(FIND "${bpath}" "${CMAKE_CURRENT_BINARY_DIR}" _binary_dir_pos) + if (_binary_dir_pos EQUAL 0) + file(RELATIVE_PATH bpath "${CMAKE_CURRENT_BINARY_DIR}" "${bpath}") + else() + file(RELATIVE_PATH bpath "${CMAKE_CURRENT_SOURCE_DIR}" "${bpath}") + endif() + endif() + + # This recipe is from cmLocalGenerator::CreateSafeUniqueObjectFileName in the + # CMake source. + + # Remove leading / + string(REGEX REPLACE "^[/]+" "" bpath "${bpath}") + # Avoid absolute paths by removing ':' + string(REPLACE ":" "_" bpath "${bpath}") + # Avoid relative paths that go up the tree + string(REPLACE "../" "__/" bpath "${bpath}") + # Avoid spaces + string(REPLACE " " "_" bpath "${bpath}") + + # Strip off the filename. I wait until here to do it, since removing the + # basename can make a path that looked like path/../basename turn into + # path/.. (notice the trailing slash). + get_filename_component(bpath "${bpath}" PATH) + + set(${build_path} "${bpath}" PARENT_SCOPE) + #message("${build_path} = ${bpath}") +endfunction() + +############################################################################## +# This helper macro populates the following variables and setups up custom +# commands and targets to invoke the nvcc compiler to generate C or PTX source +# dependent upon the format parameter. The compiler is invoked once with -M +# to generate a dependency file and a second time with -cuda or -ptx to generate +# a .cpp or .ptx file. +# INPUT: +# cuda_target - Target name +# format - PTX, CUBIN, FATBIN or OBJ +# FILE1 .. FILEN - The remaining arguments are the sources to be wrapped. +# OPTIONS - Extra options to NVCC +# OUTPUT: +# generated_files - List of generated files +############################################################################## +############################################################################## + +macro(CUDA_WRAP_SRCS cuda_target format generated_files) + + # Put optional arguments in list. + set(_argn_list "${ARGN}") + # If one of the given optional arguments is "PHONY", make a note of it, then + # remove it from the list. + list(FIND _argn_list "PHONY" _phony_idx) + if("${_phony_idx}" GREATER "-1") + set(_target_is_phony true) + list(REMOVE_AT _argn_list ${_phony_idx}) + else() + set(_target_is_phony false) + endif() + + # If CMake doesn't support separable compilation, complain + if(CUDA_SEPARABLE_COMPILATION AND CMAKE_VERSION VERSION_LESS "2.8.10.1") + message(SEND_ERROR "CUDA_SEPARABLE_COMPILATION isn't supported for CMake versions less than 2.8.10.1") + endif() + + # Set up all the command line flags here, so that they can be overridden on a per target basis. + + set(nvcc_flags "") + + # Emulation if the card isn't present. + if (CUDA_BUILD_EMULATION) + # Emulation. + set(nvcc_flags ${nvcc_flags} --device-emulation -D_DEVICEEMU -g) + else() + # Device mode. No flags necessary. + endif() + + if(CUDA_HOST_COMPILATION_CPP) + set(CUDA_C_OR_CXX CXX) + else() + message(WARNING "--host-compilation flag is deprecated in CUDA version >= 3.0. Removing --host-compilation C flag" ) + set(CUDA_C_OR_CXX C) + endif() + + set(generated_extension ${CMAKE_${CUDA_C_OR_CXX}_OUTPUT_EXTENSION}) + + if(CUDA_64_BIT_DEVICE_CODE) + set(nvcc_flags ${nvcc_flags} -m64) + else() + set(nvcc_flags ${nvcc_flags} -m32) + endif() + + if(CUDA_TARGET_CPU_ARCH) + set(nvcc_flags ${nvcc_flags} "--target-cpu-architecture=${CUDA_TARGET_CPU_ARCH}") + endif() + + # This needs to be passed in at this stage, because VS needs to fill out the + # various macros from within VS. Note that CCBIN is only used if + # -ccbin or --compiler-bindir isn't used and CUDA_HOST_COMPILER matches + # _CUDA_MSVC_HOST_COMPILER + if(CMAKE_GENERATOR MATCHES "Visual Studio") + set(ccbin_flags -D "\"CCBIN:PATH=${_CUDA_MSVC_HOST_COMPILER}\"" ) + else() + set(ccbin_flags) + endif() + + # Figure out which configure we will use and pass that in as an argument to + # the script. We need to defer the decision until compilation time, because + # for VS projects we won't know if we are making a debug or release build + # until build time. + if(CMAKE_GENERATOR MATCHES "Visual Studio") + set( CUDA_build_configuration "$(ConfigurationName)" ) + else() + set( CUDA_build_configuration "${CMAKE_BUILD_TYPE}") + endif() + + # Initialize our list of includes with the user ones followed by the CUDA system ones. + set(CUDA_NVCC_INCLUDE_DIRS ${CUDA_NVCC_INCLUDE_DIRS_USER} "${CUDA_INCLUDE_DIRS}") + if(_target_is_phony) + # If the passed in target name isn't a real target (i.e., this is from a call to one of the + # cuda_compile_* functions), need to query directory properties to get include directories + # and compile definitions. + get_directory_property(_dir_include_dirs INCLUDE_DIRECTORIES) + get_directory_property(_dir_compile_defs COMPILE_DEFINITIONS) + + list(APPEND CUDA_NVCC_INCLUDE_DIRS "${_dir_include_dirs}") + set(CUDA_NVCC_COMPILE_DEFINITIONS "${_dir_compile_defs}") + else() + # Append the include directories for this target via generator expression, which is + # expanded by the FILE(GENERATE) call below. This generator expression captures all + # include dirs set by the user, whether via directory properties or target properties + list(APPEND CUDA_NVCC_INCLUDE_DIRS "$") + + # Do the same thing with compile definitions + set(CUDA_NVCC_COMPILE_DEFINITIONS "$") + endif() + + + # Reset these variables + set(CUDA_WRAP_OPTION_NVCC_FLAGS) + foreach(config ${CUDA_configuration_types}) + string(TOUPPER ${config} config_upper) + set(CUDA_WRAP_OPTION_NVCC_FLAGS_${config_upper}) + endforeach() + + CUDA_GET_SOURCES_AND_OPTIONS(_cuda_wrap_sources _cuda_wrap_cmake_options _cuda_wrap_options ${_argn_list}) + CUDA_PARSE_NVCC_OPTIONS(CUDA_WRAP_OPTION_NVCC_FLAGS ${_cuda_wrap_options}) + + # Figure out if we are building a shared library. BUILD_SHARED_LIBS is + # respected in CUDA_ADD_LIBRARY. + set(_cuda_build_shared_libs FALSE) + # SHARED, MODULE + list(FIND _cuda_wrap_cmake_options SHARED _cuda_found_SHARED) + list(FIND _cuda_wrap_cmake_options MODULE _cuda_found_MODULE) + if(_cuda_found_SHARED GREATER -1 OR _cuda_found_MODULE GREATER -1) + set(_cuda_build_shared_libs TRUE) + endif() + # STATIC + list(FIND _cuda_wrap_cmake_options STATIC _cuda_found_STATIC) + if(_cuda_found_STATIC GREATER -1) + set(_cuda_build_shared_libs FALSE) + endif() + + # CUDA_HOST_FLAGS + if(_cuda_build_shared_libs) + # If we are setting up code for a shared library, then we need to add extra flags for + # compiling objects for shared libraries. + set(CUDA_HOST_SHARED_FLAGS ${CMAKE_SHARED_LIBRARY_${CUDA_C_OR_CXX}_FLAGS}) + else() + set(CUDA_HOST_SHARED_FLAGS) + endif() + + macro(_filter_blocklisted_host_flags CUDA_FLAGS) + string(REGEX REPLACE "[ \t]+" ";" ${CUDA_FLAGS} "${${CUDA_FLAGS}}") + foreach(_blacklisted ${CUDA_PROPAGATE_HOST_FLAGS_BLACKLIST}) + list(REMOVE_ITEM ${CUDA_FLAGS} "${_blacklisted}") + endforeach() + string(REPLACE ";" " " ${CUDA_FLAGS} "${${CUDA_FLAGS}}") + endmacro() + + # Only add the CMAKE_{C,CXX}_FLAGS if we are propagating host flags. We + # always need to set the SHARED_FLAGS, though. + if(CUDA_PROPAGATE_HOST_FLAGS) + set(_cuda_C_FLAGS "${CMAKE_${CUDA_C_OR_CXX}_FLAGS}") + _filter_blocklisted_host_flags(_cuda_C_FLAGS) + set(_cuda_host_flags "set(CMAKE_HOST_FLAGS ${_cuda_C_FLAGS} ${CUDA_HOST_SHARED_FLAGS})") + else() + set(_cuda_host_flags "set(CMAKE_HOST_FLAGS ${CUDA_HOST_SHARED_FLAGS})") + endif() + + set(_cuda_nvcc_flags_config "# Build specific configuration flags") + # Loop over all the configuration types to generate appropriate flags for run_nvcc.cmake + foreach(config ${CUDA_configuration_types}) + string(TOUPPER ${config} config_upper) + # CMAKE_FLAGS are strings and not lists. By not putting quotes around CMAKE_FLAGS + # we convert the strings to lists (like we want). + + if(CUDA_PROPAGATE_HOST_FLAGS) + # nvcc chokes on -g3 in versions previous to 3.0, so replace it with -g + set(_cuda_fix_g3 FALSE) + + set(_cuda_C_FLAGS "${CMAKE_${CUDA_C_OR_CXX}_FLAGS_${config_upper}}") + _filter_blocklisted_host_flags(_cuda_C_FLAGS) + if(_cuda_fix_g3) + string(REPLACE "-g3" "-g" _cuda_C_FLAGS "${_cuda_C_FLAGS}") + endif() + + string(APPEND _cuda_host_flags "\nset(CMAKE_HOST_FLAGS_${config_upper} ${_cuda_C_FLAGS})") + endif() + + # Note that if we ever want CUDA_NVCC_FLAGS_ to be string (instead of a list + # like it is currently), we can remove the quotes around the + # ${CUDA_NVCC_FLAGS_${config_upper}} variable like the CMAKE_HOST_FLAGS_ variable. + string(APPEND _cuda_nvcc_flags_config "\nset(CUDA_NVCC_FLAGS_${config_upper} ${CUDA_NVCC_FLAGS_${config_upper}} ;; ${CUDA_WRAP_OPTION_NVCC_FLAGS_${config_upper}})") + endforeach() + + # Process the C++14 flag. If the host sets the flag, we need to add it to nvcc and + # remove it from the host. This is because -Xcompile -std=c++ will choke nvcc (it uses + # the C preprocessor). In order to get this to work correctly, we need to use nvcc's + # specific c++14 flag. + if( "${_cuda_host_flags}" MATCHES "-std=c\\+\\+11") + # Add the c++14 flag to nvcc if it isn't already present. Note that we only look at + # the main flag instead of the configuration specific flags. + if( NOT "${CUDA_NVCC_FLAGS}" MATCHES "-std=c\\+\\+14" ) + list(APPEND nvcc_flags --std c++14) + endif() + string(REGEX REPLACE "[-]+std=c\\+\\+14" "" _cuda_host_flags "${_cuda_host_flags}") + endif() + + if(_cuda_build_shared_libs) + list(APPEND nvcc_flags "-D${cuda_target}_EXPORTS") + endif() + + # Reset the output variable + set(_cuda_wrap_generated_files "") + + # Iterate over the macro arguments and create custom + # commands for all the .cu files. + foreach(file ${_argn_list}) + # Ignore any file marked as a HEADER_FILE_ONLY + get_source_file_property(_is_header ${file} HEADER_FILE_ONLY) + # Allow per source file overrides of the format. Also allows compiling non-.cu files. + get_source_file_property(_cuda_source_format ${file} CUDA_SOURCE_PROPERTY_FORMAT) + if((${file} MATCHES "\\.cu$" OR _cuda_source_format) AND NOT _is_header) + + if(NOT _cuda_source_format) + set(_cuda_source_format ${format}) + endif() + # If file isn't a .cu file, we need to tell nvcc to treat it as such. + if(NOT file MATCHES "\\.cu$") + set(cuda_language_flag -x=cu) + else() + set(cuda_language_flag) + endif() + + if( ${_cuda_source_format} MATCHES "OBJ") + set( cuda_compile_to_external_module OFF ) + else() + set( cuda_compile_to_external_module ON ) + if( ${_cuda_source_format} MATCHES "PTX" ) + set( cuda_compile_to_external_module_type "ptx" ) + elseif( ${_cuda_source_format} MATCHES "CUBIN") + set( cuda_compile_to_external_module_type "cubin" ) + elseif( ${_cuda_source_format} MATCHES "FATBIN") + set( cuda_compile_to_external_module_type "fatbin" ) + else() + message( FATAL_ERROR "Invalid format flag passed to CUDA_WRAP_SRCS or set with CUDA_SOURCE_PROPERTY_FORMAT file property for file '${file}': '${_cuda_source_format}'. Use OBJ, PTX, CUBIN or FATBIN.") + endif() + endif() + + if(cuda_compile_to_external_module) + # Don't use any of the host compilation flags for PTX targets. + set(CUDA_HOST_FLAGS) + set(CUDA_NVCC_FLAGS_CONFIG) + else() + set(CUDA_HOST_FLAGS ${_cuda_host_flags}) + set(CUDA_NVCC_FLAGS_CONFIG ${_cuda_nvcc_flags_config}) + endif() + + # Determine output directory + cuda_compute_build_path("${file}" cuda_build_path) + set(cuda_compile_intermediate_directory "${CMAKE_CURRENT_BINARY_DIR}/CMakeFiles/${cuda_target}.dir/${cuda_build_path}") + if(CUDA_GENERATED_OUTPUT_DIR) + set(cuda_compile_output_dir "${CUDA_GENERATED_OUTPUT_DIR}") + else() + if ( cuda_compile_to_external_module ) + set(cuda_compile_output_dir "${CMAKE_CURRENT_BINARY_DIR}") + else() + set(cuda_compile_output_dir "${cuda_compile_intermediate_directory}") + endif() + endif() + + # Add a custom target to generate a c or ptx file. ###################### + + get_filename_component( basename ${file} NAME ) + if( cuda_compile_to_external_module ) + set(generated_file_path "${cuda_compile_output_dir}") + set(generated_file_basename "${cuda_target}_generated_${basename}.${cuda_compile_to_external_module_type}") + set(format_flag "-${cuda_compile_to_external_module_type}") + file(MAKE_DIRECTORY "${cuda_compile_output_dir}") + else() + set(generated_file_path "${cuda_compile_output_dir}/${CMAKE_CFG_INTDIR}") + set(generated_file_basename "${cuda_target}_generated_${basename}${generated_extension}") + if(CUDA_SEPARABLE_COMPILATION) + set(format_flag "-dc") + else() + set(format_flag "-c") + endif() + endif() + + # Set all of our file names. Make sure that whatever filenames that have + # generated_file_path in them get passed in through as a command line + # argument, so that the ${CMAKE_CFG_INTDIR} gets expanded at run time + # instead of configure time. + set(generated_file "${generated_file_path}/${generated_file_basename}") + set(cmake_dependency_file "${cuda_compile_intermediate_directory}/${generated_file_basename}.depend") + set(NVCC_generated_dependency_file "${cuda_compile_intermediate_directory}/${generated_file_basename}.NVCC-depend") + set(generated_cubin_file "${generated_file_path}/${generated_file_basename}.cubin.txt") + set(custom_target_script_pregen "${cuda_compile_intermediate_directory}/${generated_file_basename}.cmake.pre-gen") + set(custom_target_script "${cuda_compile_intermediate_directory}/${generated_file_basename}$<$>:.$>.cmake") + + # Setup properties for obj files: + if( NOT cuda_compile_to_external_module ) + set_source_files_properties("${generated_file}" + PROPERTIES + EXTERNAL_OBJECT true # This is an object file not to be compiled, but only be linked. + ) + endif() + + # Don't add CMAKE_CURRENT_SOURCE_DIR if the path is already an absolute path. + get_filename_component(file_path "${file}" PATH) + if(IS_ABSOLUTE "${file_path}") + set(source_file "${file}") + else() + set(source_file "${CMAKE_CURRENT_SOURCE_DIR}/${file}") + endif() + + if( NOT cuda_compile_to_external_module AND CUDA_SEPARABLE_COMPILATION) + list(APPEND ${cuda_target}_SEPARABLE_COMPILATION_OBJECTS "${generated_file}") + endif() + + # Bring in the dependencies. Creates a variable CUDA_NVCC_DEPEND ####### + cuda_include_nvcc_dependencies(${cmake_dependency_file}) + + # Convenience string for output ######################################### + if(CUDA_BUILD_EMULATION) + set(cuda_build_type "Emulation") + else() + set(cuda_build_type "Device") + endif() + + # Build the NVCC made dependency file ################################### + set(build_cubin OFF) + if ( NOT CUDA_BUILD_EMULATION AND CUDA_BUILD_CUBIN ) + if ( NOT cuda_compile_to_external_module ) + set ( build_cubin ON ) + endif() + endif() + + # Configure the build script + configure_file("${CUDA_run_nvcc}" "${custom_target_script_pregen}" @ONLY) + file(GENERATE + OUTPUT "${custom_target_script}" + INPUT "${custom_target_script_pregen}" + ) + + # So if a user specifies the same cuda file as input more than once, you + # can have bad things happen with dependencies. Here we check an option + # to see if this is the behavior they want. + if(CUDA_ATTACH_VS_BUILD_RULE_TO_CUDA_FILE) + set(main_dep MAIN_DEPENDENCY ${source_file}) + else() + set(main_dep DEPENDS ${source_file}) + endif() + + if(CUDA_VERBOSE_BUILD) + set(verbose_output ON) + elseif(CMAKE_GENERATOR MATCHES "Makefiles") + set(verbose_output "$(VERBOSE)") + # This condition lets us also turn on verbose output when someone + # specifies CMAKE_VERBOSE_MAKEFILE, even if the generator isn't + # the Makefiles generator (this is important for us, Ninja users.) + elseif(CMAKE_VERBOSE_MAKEFILE) + set(verbose_output ON) + else() + set(verbose_output OFF) + endif() + + # Create up the comment string + file(RELATIVE_PATH generated_file_relative_path "${CMAKE_BINARY_DIR}" "${generated_file}") + if(cuda_compile_to_external_module) + set(cuda_build_comment_string "Building NVCC ${cuda_compile_to_external_module_type} file ${generated_file_relative_path}") + else() + set(cuda_build_comment_string "Building NVCC (${cuda_build_type}) object ${generated_file_relative_path}") + endif() + + set(_verbatim VERBATIM) + if(ccbin_flags MATCHES "\\$\\(VCInstallDir\\)") + set(_verbatim "") + endif() + + # Build the generated file and dependency file ########################## + add_custom_command( + OUTPUT ${generated_file} + # These output files depend on the source_file and the contents of cmake_dependency_file + ${main_dep} + DEPENDS ${CUDA_NVCC_DEPEND} + DEPENDS ${custom_target_script} + # Make sure the output directory exists before trying to write to it. + COMMAND ${CMAKE_COMMAND} -E make_directory "${generated_file_path}" + COMMAND ${CMAKE_COMMAND} ARGS + -D verbose:BOOL=${verbose_output} + ${ccbin_flags} + -D build_configuration:STRING=${CUDA_build_configuration} + -D "generated_file:STRING=${generated_file}" + -D "generated_cubin_file:STRING=${generated_cubin_file}" + -P "${custom_target_script}" + WORKING_DIRECTORY "${cuda_compile_intermediate_directory}" + COMMENT "${cuda_build_comment_string}" + ${_verbatim} + ) + + # Make sure the build system knows the file is generated. + set_source_files_properties(${generated_file} PROPERTIES GENERATED TRUE) + + list(APPEND _cuda_wrap_generated_files ${generated_file}) + + # Add the other files that we want cmake to clean on a cleanup ########## + list(APPEND CUDA_ADDITIONAL_CLEAN_FILES "${cmake_dependency_file}") + list(REMOVE_DUPLICATES CUDA_ADDITIONAL_CLEAN_FILES) + set(CUDA_ADDITIONAL_CLEAN_FILES ${CUDA_ADDITIONAL_CLEAN_FILES} CACHE INTERNAL "List of intermediate files that are part of the cuda dependency scanning.") + + endif() + endforeach() + + # Set the return parameter + set(${generated_files} ${_cuda_wrap_generated_files}) +endmacro() + +function(_cuda_get_important_host_flags important_flags flag_string) + if(CMAKE_GENERATOR MATCHES "Visual Studio") + string(REGEX MATCHALL "/M[DT][d]?" flags "${flag_string}") + list(APPEND ${important_flags} ${flags}) + else() + string(REGEX MATCHALL "-fPIC" flags "${flag_string}") + list(APPEND ${important_flags} ${flags}) + endif() + set(${important_flags} ${${important_flags}} PARENT_SCOPE) +endfunction() + +############################################################################### +############################################################################### +# Separable Compilation Link +############################################################################### +############################################################################### + +# Compute the filename to be used by CUDA_LINK_SEPARABLE_COMPILATION_OBJECTS +function(CUDA_COMPUTE_SEPARABLE_COMPILATION_OBJECT_FILE_NAME output_file_var cuda_target object_files) + if (object_files) + set(generated_extension ${CMAKE_${CUDA_C_OR_CXX}_OUTPUT_EXTENSION}) + set(output_file "${CMAKE_CURRENT_BINARY_DIR}/CMakeFiles/${cuda_target}.dir/${CMAKE_CFG_INTDIR}/${cuda_target}_intermediate_link${generated_extension}") + else() + set(output_file) + endif() + + set(${output_file_var} "${output_file}" PARENT_SCOPE) +endfunction() + +# Setup the build rule for the separable compilation intermediate link file. +function(CUDA_LINK_SEPARABLE_COMPILATION_OBJECTS output_file cuda_target options object_files) + if (object_files) + + set_source_files_properties("${output_file}" + PROPERTIES + EXTERNAL_OBJECT TRUE # This is an object file not to be compiled, but only + # be linked. + GENERATED TRUE # This file is generated during the build + ) + + # For now we are ignoring all the configuration specific flags. + set(nvcc_flags) + CUDA_PARSE_NVCC_OPTIONS(nvcc_flags ${options}) + if(CUDA_64_BIT_DEVICE_CODE) + list(APPEND nvcc_flags -m64) + else() + list(APPEND nvcc_flags -m32) + endif() + # If -ccbin, --compiler-bindir has been specified, don't do anything. Otherwise add it here. + list( FIND nvcc_flags "-ccbin" ccbin_found0 ) + list( FIND nvcc_flags "--compiler-bindir" ccbin_found1 ) + if( ccbin_found0 LESS 0 AND ccbin_found1 LESS 0 AND CUDA_HOST_COMPILER ) + # Match VERBATIM check below. + if(CUDA_HOST_COMPILER MATCHES "\\$\\(VCInstallDir\\)") + list(APPEND nvcc_flags -ccbin "\"${CUDA_HOST_COMPILER}\"") + else() + list(APPEND nvcc_flags -ccbin "${CUDA_HOST_COMPILER}") + endif() + endif() + + # Create a list of flags specified by CUDA_NVCC_FLAGS_${CONFIG} and CMAKE_${CUDA_C_OR_CXX}_FLAGS* + set(config_specific_flags) + set(flags) + foreach(config ${CUDA_configuration_types}) + string(TOUPPER ${config} config_upper) + # Add config specific flags + foreach(f ${CUDA_NVCC_FLAGS_${config_upper}}) + list(APPEND config_specific_flags $<$:${f}>) + endforeach() + set(important_host_flags) + _cuda_get_important_host_flags(important_host_flags "${CMAKE_${CUDA_C_OR_CXX}_FLAGS_${config_upper}}") + foreach(f ${important_host_flags}) + list(APPEND flags $<$:-Xcompiler> $<$:${f}>) + endforeach() + endforeach() + # Add CMAKE_${CUDA_C_OR_CXX}_FLAGS + set(important_host_flags) + _cuda_get_important_host_flags(important_host_flags "${CMAKE_${CUDA_C_OR_CXX}_FLAGS}") + foreach(f ${important_host_flags}) + list(APPEND flags -Xcompiler ${f}) + endforeach() + + # Add our general CUDA_NVCC_FLAGS with the configuration specific flags + set(nvcc_flags ${CUDA_NVCC_FLAGS} ${config_specific_flags} ${nvcc_flags}) + + file(RELATIVE_PATH output_file_relative_path "${CMAKE_BINARY_DIR}" "${output_file}") + + # Some generators don't handle the multiple levels of custom command + # dependencies correctly (obj1 depends on file1, obj2 depends on obj1), so + # we work around that issue by compiling the intermediate link object as a + # pre-link custom command in that situation. + set(do_obj_build_rule TRUE) + if (MSVC_VERSION GREATER 1599 AND MSVC_VERSION LESS 1800) + # VS 2010 and 2012 have this problem. + set(do_obj_build_rule FALSE) + endif() + + set(_verbatim VERBATIM) + if(nvcc_flags MATCHES "\\$\\(VCInstallDir\\)") + set(_verbatim "") + endif() + + if (do_obj_build_rule) + add_custom_command( + OUTPUT ${output_file} + DEPENDS ${object_files} + COMMAND ${CUDA_NVCC_EXECUTABLE} ${nvcc_flags} -dlink ${object_files} -o ${output_file} + ${flags} + COMMENT "Building NVCC intermediate link file ${output_file_relative_path}" + COMMAND_EXPAND_LISTS + ${_verbatim} + ) + else() + get_filename_component(output_file_dir "${output_file}" DIRECTORY) + add_custom_command( + TARGET ${cuda_target} + PRE_LINK + COMMAND ${CMAKE_COMMAND} -E echo "Building NVCC intermediate link file ${output_file_relative_path}" + COMMAND ${CMAKE_COMMAND} -E make_directory "${output_file_dir}" + COMMAND ${CUDA_NVCC_EXECUTABLE} ${nvcc_flags} ${flags} -dlink ${object_files} -o "${output_file}" + COMMAND_EXPAND_LISTS + ${_verbatim} + ) + endif() + endif() +endfunction() + +############################################################################### +############################################################################### +# ADD LIBRARY +############################################################################### +############################################################################### +macro(CUDA_ADD_LIBRARY cuda_target) + + CUDA_ADD_CUDA_INCLUDE_ONCE() + + # Separate the sources from the options + CUDA_GET_SOURCES_AND_OPTIONS(_sources _cmake_options _options ${ARGN}) + CUDA_BUILD_SHARED_LIBRARY(_cuda_shared_flag ${ARGN}) + # Create custom commands and targets for each file. + CUDA_WRAP_SRCS( ${cuda_target} OBJ _generated_files ${_sources} + ${_cmake_options} ${_cuda_shared_flag} + OPTIONS ${_options} ) + + # Compute the file name of the intermedate link file used for separable + # compilation. + CUDA_COMPUTE_SEPARABLE_COMPILATION_OBJECT_FILE_NAME(link_file ${cuda_target} "${${cuda_target}_SEPARABLE_COMPILATION_OBJECTS}") + + # Add the library. + add_library(${cuda_target} ${_cmake_options} + ${_generated_files} + ${_sources} + ${link_file} + ) + + # Add a link phase for the separable compilation if it has been enabled. If + # it has been enabled then the ${cuda_target}_SEPARABLE_COMPILATION_OBJECTS + # variable will have been defined. + CUDA_LINK_SEPARABLE_COMPILATION_OBJECTS("${link_file}" ${cuda_target} "${_options}" "${${cuda_target}_SEPARABLE_COMPILATION_OBJECTS}") + + target_link_libraries(${cuda_target} ${CUDA_LINK_LIBRARIES_KEYWORD} + ${CUDA_LIBRARIES} + ) + + if(CUDA_SEPARABLE_COMPILATION) + target_link_libraries(${cuda_target} ${CUDA_LINK_LIBRARIES_KEYWORD} + ${CUDA_cudadevrt_LIBRARY} + ) + endif() + + # We need to set the linker language based on what the expected generated file + # would be. CUDA_C_OR_CXX is computed based on CUDA_HOST_COMPILATION_CPP. + set_target_properties(${cuda_target} + PROPERTIES + LINKER_LANGUAGE ${CUDA_C_OR_CXX} + ) + +endmacro() + + +############################################################################### +############################################################################### +# ADD EXECUTABLE +############################################################################### +############################################################################### +macro(CUDA_ADD_EXECUTABLE cuda_target) + + CUDA_ADD_CUDA_INCLUDE_ONCE() + + # Separate the sources from the options + CUDA_GET_SOURCES_AND_OPTIONS(_sources _cmake_options _options ${ARGN}) + # Create custom commands and targets for each file. + CUDA_WRAP_SRCS( ${cuda_target} OBJ _generated_files ${_sources} OPTIONS ${_options} ) + + # Compute the file name of the intermedate link file used for separable + # compilation. + CUDA_COMPUTE_SEPARABLE_COMPILATION_OBJECT_FILE_NAME(link_file ${cuda_target} "${${cuda_target}_SEPARABLE_COMPILATION_OBJECTS}") + + # Add the library. + add_executable(${cuda_target} ${_cmake_options} + ${_generated_files} + ${_sources} + ${link_file} + ) + + # Add a link phase for the separable compilation if it has been enabled. If + # it has been enabled then the ${cuda_target}_SEPARABLE_COMPILATION_OBJECTS + # variable will have been defined. + CUDA_LINK_SEPARABLE_COMPILATION_OBJECTS("${link_file}" ${cuda_target} "${_options}" "${${cuda_target}_SEPARABLE_COMPILATION_OBJECTS}") + + target_link_libraries(${cuda_target} ${CUDA_LINK_LIBRARIES_KEYWORD} + ${CUDA_LIBRARIES} + ) + + # We need to set the linker language based on what the expected generated file + # would be. CUDA_C_OR_CXX is computed based on CUDA_HOST_COMPILATION_CPP. + set_target_properties(${cuda_target} + PROPERTIES + LINKER_LANGUAGE ${CUDA_C_OR_CXX} + ) + +endmacro() + + +############################################################################### +############################################################################### +# (Internal) helper for manually added cuda source files with specific targets +############################################################################### +############################################################################### +macro(cuda_compile_base cuda_target format generated_files) + # Update a counter in this directory, to keep phony target names unique. + set(_cuda_target "${cuda_target}") + get_property(_counter DIRECTORY PROPERTY _cuda_internal_phony_counter) + if(_counter) + math(EXPR _counter "${_counter} + 1") + else() + set(_counter 1) + endif() + string(APPEND _cuda_target "_${_counter}") + set_property(DIRECTORY PROPERTY _cuda_internal_phony_counter ${_counter}) + + # Separate the sources from the options + CUDA_GET_SOURCES_AND_OPTIONS(_sources _cmake_options _options ${ARGN}) + + # Create custom commands and targets for each file. + CUDA_WRAP_SRCS( ${_cuda_target} ${format} _generated_files ${_sources} + ${_cmake_options} OPTIONS ${_options} PHONY) + + set( ${generated_files} ${_generated_files}) + +endmacro() + +############################################################################### +############################################################################### +# CUDA COMPILE +############################################################################### +############################################################################### +macro(CUDA_COMPILE generated_files) + cuda_compile_base(cuda_compile OBJ ${generated_files} ${ARGN}) +endmacro() + +############################################################################### +############################################################################### +# CUDA COMPILE PTX +############################################################################### +############################################################################### +macro(CUDA_COMPILE_PTX generated_files) + cuda_compile_base(cuda_compile_ptx PTX ${generated_files} ${ARGN}) +endmacro() + +############################################################################### +############################################################################### +# CUDA COMPILE FATBIN +############################################################################### +############################################################################### +macro(CUDA_COMPILE_FATBIN generated_files) + cuda_compile_base(cuda_compile_fatbin FATBIN ${generated_files} ${ARGN}) +endmacro() + +############################################################################### +############################################################################### +# CUDA COMPILE CUBIN +############################################################################### +############################################################################### +macro(CUDA_COMPILE_CUBIN generated_files) + cuda_compile_base(cuda_compile_cubin CUBIN ${generated_files} ${ARGN}) +endmacro() + + +############################################################################### +############################################################################### +# CUDA ADD CUFFT TO TARGET +############################################################################### +############################################################################### +macro(CUDA_ADD_CUFFT_TO_TARGET target) + if (CUDA_BUILD_EMULATION) + target_link_libraries(${target} ${CUDA_LINK_LIBRARIES_KEYWORD} ${CUDA_cufftemu_LIBRARY}) + else() + target_link_libraries(${target} ${CUDA_LINK_LIBRARIES_KEYWORD} ${CUDA_cufft_LIBRARY}) + endif() +endmacro() + +############################################################################### +############################################################################### +# CUDA ADD CUBLAS TO TARGET +############################################################################### +############################################################################### +macro(CUDA_ADD_CUBLAS_TO_TARGET target) + if (CUDA_BUILD_EMULATION) + target_link_libraries(${target} ${CUDA_LINK_LIBRARIES_KEYWORD} ${CUDA_cublasemu_LIBRARY}) + else() + target_link_libraries(${target} ${CUDA_LINK_LIBRARIES_KEYWORD} ${CUDA_cublas_LIBRARY} ${CUDA_cublas_device_LIBRARY} ${CUDA_cublasLt_LIBRARY}) + endif() +endmacro() + +############################################################################### +############################################################################### +# CUDA BUILD CLEAN TARGET +############################################################################### +############################################################################### +macro(CUDA_BUILD_CLEAN_TARGET) + # Call this after you add all your CUDA targets, and you will get a + # convenience target. You should also make clean after running this target + # to get the build system to generate all the code again. + + set(cuda_clean_target_name clean_cuda_depends) + if (CMAKE_GENERATOR MATCHES "Visual Studio") + string(TOUPPER ${cuda_clean_target_name} cuda_clean_target_name) + endif() + add_custom_target(${cuda_clean_target_name} + COMMAND ${CMAKE_COMMAND} -E remove ${CUDA_ADDITIONAL_CLEAN_FILES}) + + # Clear out the variable, so the next time we configure it will be empty. + # This is useful so that the files won't persist in the list after targets + # have been removed. + set(CUDA_ADDITIONAL_CLEAN_FILES "" CACHE INTERNAL "List of intermediate files that are part of the cuda dependency scanning.") +endmacro() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/make2cmake.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/make2cmake.cmake new file mode 100644 index 0000000000000000000000000000000000000000..580f24a400d8c5662ec572c4631db9e3e47645d9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/make2cmake.cmake @@ -0,0 +1,106 @@ +# James Bigler, NVIDIA Corp (nvidia.com - jbigler) +# Abe Stephens, SCI Institute -- http://www.sci.utah.edu/~abe/FindCuda.html +# +# Copyright (c) 2008 - 2009 NVIDIA Corporation. All rights reserved. +# +# Copyright (c) 2007-2009 +# Scientific Computing and Imaging Institute, University of Utah +# +# This code is licensed under the MIT License. See the FindCUDA.cmake script +# for the text of the license. + +# The MIT License +# +# License for the specific language governing rights and limitations under +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included +# in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# + +####################################################################### +# This converts a file written in makefile syntax into one that can be included +# by CMake. + +# Input variables +# +# verbose:BOOL=<> OFF: Be as quiet as possible (default) +# ON : Extra output +# +# input_file:FILEPATH=<> Path to dependency file in makefile format +# +# output_file:FILEPATH=<> Path to file with dependencies in CMake readable variable +# + +file(READ ${input_file} depend_text) + +if (NOT "${depend_text}" STREQUAL "") + + # message("FOUND DEPENDS") + + string(REPLACE "\\ " " " depend_text ${depend_text}) + + # This works for the nvcc -M generated dependency files. + string(REGEX REPLACE "^.* : " "" depend_text ${depend_text}) + string(REGEX REPLACE "[ \\\\]*\n" ";" depend_text ${depend_text}) + + set(dependency_list "") + + foreach(file ${depend_text}) + + string(REGEX REPLACE "^ +" "" file ${file}) + + # OK, now if we had a UNC path, nvcc has a tendency to only output the first '/' + # instead of '//'. Here we will test to see if the file exists, if it doesn't then + # try to prepend another '/' to the path and test again. If it still fails remove the + # path. + + if(NOT EXISTS "${file}") + if (EXISTS "/${file}") + set(file "/${file}") + else() + if(verbose) + message(WARNING " Removing non-existent dependency file: ${file}") + endif() + set(file "") + endif() + endif() + + # Make sure we check to see if we have a file, before asking if it is not a directory. + # if(NOT IS_DIRECTORY "") will return TRUE. + if(file AND NOT IS_DIRECTORY "${file}") + # If softlinks start to matter, we should change this to REALPATH. For now we need + # to flatten paths, because nvcc can generate stuff like /bin/../include instead of + # just /include. + get_filename_component(file_absolute "${file}" ABSOLUTE) + list(APPEND dependency_list "${file_absolute}") + endif() + + endforeach() + +else() + # message("FOUND NO DEPENDS") +endif() + +# Remove the duplicate entries and sort them. +list(REMOVE_DUPLICATES dependency_list) +list(SORT dependency_list) + +foreach(file ${dependency_list}) + string(APPEND cuda_nvcc_depend " \"${file}\"\n") +endforeach() + +file(WRITE ${output_file} "# Generated by: make2cmake.cmake\nSET(CUDA_NVCC_DEPEND\n ${cuda_nvcc_depend})\n\n") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/parse_cubin.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/parse_cubin.cmake new file mode 100644 index 0000000000000000000000000000000000000000..25ceb49f3dd8e684e35cac49834c4db0aa5c338a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/parse_cubin.cmake @@ -0,0 +1,109 @@ +# James Bigler, NVIDIA Corp (nvidia.com - jbigler) +# Abe Stephens, SCI Institute -- http://www.sci.utah.edu/~abe/FindCuda.html +# +# Copyright (c) 2008 - 2009 NVIDIA Corporation. All rights reserved. +# +# Copyright (c) 2007-2009 +# Scientific Computing and Imaging Institute, University of Utah +# +# This code is licensed under the MIT License. See the FindCUDA.cmake script +# for the text of the license. + +# The MIT License +# +# License for the specific language governing rights and limitations under +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included +# in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. +# + +####################################################################### +# Parses a .cubin file produced by nvcc and reports statistics about the file. + + +file(READ ${input_file} file_text) + +if (NOT "${file_text}" STREQUAL "") + + string(REPLACE ";" "\\;" file_text ${file_text}) + string(REPLACE "\ncode" ";code" file_text ${file_text}) + + list(LENGTH file_text len) + + foreach(line ${file_text}) + + # Only look at "code { }" blocks. + if(line MATCHES "^code") + + # Break into individual lines. + string(REGEX REPLACE "\n" ";" line ${line}) + + foreach(entry ${line}) + + # Extract kernel names. + if (${entry} MATCHES "[^g]name = ([^ ]+)") + set(entry "${CMAKE_MATCH_1}") + + # Check to see if the kernel name starts with "_" + set(skip FALSE) + # if (${entry} MATCHES "^_") + # Skip the rest of this block. + # message("Skipping ${entry}") + # set(skip TRUE) + # else () + message("Kernel: ${entry}") + # endif () + + endif() + + # Skip the rest of the block if necessary + if(NOT skip) + + # Registers + if (${entry} MATCHES "reg([ ]+)=([ ]+)([^ ]+)") + set(entry "${CMAKE_MATCH_3}") + message("Registers: ${entry}") + endif() + + # Local memory + if (${entry} MATCHES "lmem([ ]+)=([ ]+)([^ ]+)") + set(entry "${CMAKE_MATCH_3}") + message("Local: ${entry}") + endif() + + # Shared memory + if (${entry} MATCHES "smem([ ]+)=([ ]+)([^ ]+)") + set(entry "${CMAKE_MATCH_3}") + message("Shared: ${entry}") + endif() + + if (${entry} MATCHES "^}") + message("") + endif() + + endif() + + + endforeach() + + endif() + + endforeach() + +else() + # message("FOUND NO DEPENDS") +endif() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/run_nvcc.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/run_nvcc.cmake new file mode 100644 index 0000000000000000000000000000000000000000..59c5c11a1091f34df89b681a926db602a1c75caa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/run_nvcc.cmake @@ -0,0 +1,303 @@ +# James Bigler, NVIDIA Corp (nvidia.com - jbigler) +# +# Copyright (c) 2008 - 2009 NVIDIA Corporation. All rights reserved. +# +# This code is licensed under the MIT License. See the FindCUDA.cmake script +# for the text of the license. + +# The MIT License +# +# License for the specific language governing rights and limitations under +# Permission is hereby granted, free of charge, to any person obtaining a +# copy of this software and associated documentation files (the "Software"), +# to deal in the Software without restriction, including without limitation +# the rights to use, copy, modify, merge, publish, distribute, sublicense, +# and/or sell copies of the Software, and to permit persons to whom the +# Software is furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included +# in all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL +# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING +# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER +# DEALINGS IN THE SOFTWARE. + + +########################################################################## +# This file runs the nvcc commands to produce the desired output file along with +# the dependency file needed by CMake to compute dependencies. In addition the +# file checks the output of each command and if the command fails it deletes the +# output files. + +# Input variables +# +# verbose:BOOL=<> OFF: Be as quiet as possible (default) +# ON : Describe each step +# +# build_configuration:STRING=<> Typically one of Debug, MinSizeRel, Release, or +# RelWithDebInfo, but it should match one of the +# entries in CUDA_HOST_FLAGS. This is the build +# configuration used when compiling the code. If +# blank or unspecified Debug is assumed as this is +# what CMake does. +# +# generated_file:STRING=<> File to generate. This argument must be passed in. +# +# generated_cubin_file:STRING=<> File to generate. This argument must be passed +# in if build_cubin is true. + +cmake_policy(PUSH) +cmake_policy(SET CMP0007 NEW) +cmake_policy(SET CMP0010 NEW) +if(NOT generated_file) + message(FATAL_ERROR "You must specify generated_file on the command line") +endif() + +# Set these up as variables to make reading the generated file easier +set(CMAKE_COMMAND "@CMAKE_COMMAND@") # path +set(source_file "@source_file@") # path +set(NVCC_generated_dependency_file "@NVCC_generated_dependency_file@") # path +set(cmake_dependency_file "@cmake_dependency_file@") # path +set(CUDA_make2cmake "@CUDA_make2cmake@") # path +set(CUDA_parse_cubin "@CUDA_parse_cubin@") # path +set(build_cubin @build_cubin@) # bool +set(CUDA_HOST_COMPILER "@CUDA_HOST_COMPILER@") # path +# We won't actually use these variables for now, but we need to set this, in +# order to force this file to be run again if it changes. +set(generated_file_path "@generated_file_path@") # path +set(generated_file_internal "@generated_file@") # path +set(generated_cubin_file_internal "@generated_cubin_file@") # path + +set(CUDA_NVCC_EXECUTABLE "@CUDA_NVCC_EXECUTABLE@") # path +set(CUDA_NVCC_FLAGS @CUDA_NVCC_FLAGS@ ;; @CUDA_WRAP_OPTION_NVCC_FLAGS@) # list +@CUDA_NVCC_FLAGS_CONFIG@ +set(nvcc_flags @nvcc_flags@) # list +set(CUDA_NVCC_INCLUDE_DIRS [==[@CUDA_NVCC_INCLUDE_DIRS@]==]) # list (needs to be in lua quotes to address backslashes) +string(REPLACE "\\" "/" CUDA_NVCC_INCLUDE_DIRS "${CUDA_NVCC_INCLUDE_DIRS}") +set(CUDA_NVCC_COMPILE_DEFINITIONS [==[@CUDA_NVCC_COMPILE_DEFINITIONS@]==]) # list (needs to be in lua quotes see #16510 ). +set(format_flag "@format_flag@") # string +set(cuda_language_flag @cuda_language_flag@) # list + +# Clean up list of include directories and add -I flags +list(REMOVE_DUPLICATES CUDA_NVCC_INCLUDE_DIRS) +set(CUDA_NVCC_INCLUDE_ARGS) +foreach(dir ${CUDA_NVCC_INCLUDE_DIRS}) + # Extra quotes are added around each flag to help nvcc parse out flags with spaces. + list(APPEND CUDA_NVCC_INCLUDE_ARGS "-I${dir}") +endforeach() + +# Clean up list of compile definitions, add -D flags, and append to nvcc_flags +list(REMOVE_DUPLICATES CUDA_NVCC_COMPILE_DEFINITIONS) +foreach(def ${CUDA_NVCC_COMPILE_DEFINITIONS}) + list(APPEND nvcc_flags "-D${def}") +endforeach() + +if(build_cubin AND NOT generated_cubin_file) + message(FATAL_ERROR "You must specify generated_cubin_file on the command line") +endif() + +# This is the list of host compilation flags. It C or CXX should already have +# been chosen by FindCUDA.cmake. +@CUDA_HOST_FLAGS@ + +# Take the compiler flags and package them up to be sent to the compiler via -Xcompiler +set(nvcc_host_compiler_flags "") +# If we weren't given a build_configuration, use Debug. +if(NOT build_configuration) + set(build_configuration Debug) +endif() +string(TOUPPER "${build_configuration}" build_configuration) +#message("CUDA_NVCC_HOST_COMPILER_FLAGS = ${CUDA_NVCC_HOST_COMPILER_FLAGS}") +foreach(flag ${CMAKE_HOST_FLAGS} ${CMAKE_HOST_FLAGS_${build_configuration}}) + # Extra quotes are added around each flag to help nvcc parse out flags with spaces. + string(APPEND nvcc_host_compiler_flags ",\"${flag}\"") +endforeach() +if (nvcc_host_compiler_flags) + set(nvcc_host_compiler_flags "-Xcompiler" ${nvcc_host_compiler_flags}) +endif() +#message("nvcc_host_compiler_flags = \"${nvcc_host_compiler_flags}\"") +# Add the build specific configuration flags +list(APPEND CUDA_NVCC_FLAGS ${CUDA_NVCC_FLAGS_${build_configuration}}) + +# Any -ccbin existing in CUDA_NVCC_FLAGS gets highest priority +list( FIND CUDA_NVCC_FLAGS "-ccbin" ccbin_found0 ) +list( FIND CUDA_NVCC_FLAGS "--compiler-bindir" ccbin_found1 ) +if( ccbin_found0 LESS 0 AND ccbin_found1 LESS 0 AND CUDA_HOST_COMPILER ) + if (CUDA_HOST_COMPILER STREQUAL "@_CUDA_MSVC_HOST_COMPILER@" AND DEFINED CCBIN) + set(CCBIN -ccbin "${CCBIN}") + else() + set(CCBIN -ccbin "${CUDA_HOST_COMPILER}") + endif() +endif() + +# cuda_execute_process - Executes a command with optional command echo and status message. +# +# status - Status message to print if verbose is true +# command - COMMAND argument from the usual execute_process argument structure +# ARGN - Remaining arguments are the command with arguments +# +# CUDA_result - return value from running the command +# +# Make this a macro instead of a function, so that things like RESULT_VARIABLE +# and other return variables are present after executing the process. +macro(cuda_execute_process status command) + set(_command ${command}) + if(NOT "x${_command}" STREQUAL "xCOMMAND") + message(FATAL_ERROR "Malformed call to cuda_execute_process. Missing COMMAND as second argument. (command = ${command})") + endif() + if(verbose) + execute_process(COMMAND "${CMAKE_COMMAND}" -E echo -- ${status}) + # Now we need to build up our command string. We are accounting for quotes + # and spaces, anything else is left up to the user to fix if they want to + # copy and paste a runnable command line. + set(cuda_execute_process_string) + foreach(arg ${ARGN}) + # If there are quotes, escape them, so they come through. + string(REPLACE "\"" "\\\"" arg ${arg}) + # Args with spaces need quotes around them to get them to be parsed as a single argument. + if(arg MATCHES " ") + list(APPEND cuda_execute_process_string "\"${arg}\"") + else() + list(APPEND cuda_execute_process_string ${arg}) + endif() + endforeach() + # Echo the command + execute_process(COMMAND ${CMAKE_COMMAND} -E echo ${cuda_execute_process_string}) + endif() + # Run the command + execute_process(COMMAND ${ARGN} RESULT_VARIABLE CUDA_result ) +endmacro() + +# Delete the target file +cuda_execute_process( + "Removing ${generated_file}" + COMMAND "${CMAKE_COMMAND}" -E remove "${generated_file}" + ) + +# For CUDA 2.3 and below, -G -M doesn't work, so remove the -G flag +# for dependency generation and hope for the best. +set(depends_CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS}") +set(CUDA_VERSION @CUDA_VERSION@) + +# nvcc doesn't define __CUDACC__ for some reason when generating dependency files. This +# can cause incorrect dependencies when #including files based on this macro which is +# defined in the generating passes of nvcc invocation. We will go ahead and manually +# define this for now until a future version fixes this bug. +set(CUDACC_DEFINE -D__CUDACC__) + +# Generate the dependency file +cuda_execute_process( + "Generating dependency file: ${NVCC_generated_dependency_file}" + COMMAND "${CUDA_NVCC_EXECUTABLE}" + -M + ${CUDACC_DEFINE} + "${source_file}" + -o "${NVCC_generated_dependency_file}" + ${CCBIN} + ${nvcc_flags} + ${nvcc_host_compiler_flags} + ${depends_CUDA_NVCC_FLAGS} + -DNVCC + ${CUDA_NVCC_INCLUDE_ARGS} + ) + +if(CUDA_result) + message(FATAL_ERROR "Error generating ${generated_file}") +endif() + +# Generate the cmake readable dependency file to a temp file. Don't put the +# quotes just around the filenames for the input_file and output_file variables. +# CMake will pass the quotes through and not be able to find the file. +cuda_execute_process( + "Generating temporary cmake readable file: ${cmake_dependency_file}.tmp" + COMMAND "${CMAKE_COMMAND}" + -D "input_file:FILEPATH=${NVCC_generated_dependency_file}" + -D "output_file:FILEPATH=${cmake_dependency_file}.tmp" + -D "verbose=${verbose}" + -P "${CUDA_make2cmake}" + ) + +if(CUDA_result) + message(FATAL_ERROR "Error generating ${generated_file}") +endif() + +# Copy the file if it is different +cuda_execute_process( + "Copy if different ${cmake_dependency_file}.tmp to ${cmake_dependency_file}" + COMMAND "${CMAKE_COMMAND}" -E copy_if_different "${cmake_dependency_file}.tmp" "${cmake_dependency_file}" + ) + +if(CUDA_result) + message(FATAL_ERROR "Error generating ${generated_file}") +endif() + +# Delete the temporary file +cuda_execute_process( + "Removing ${cmake_dependency_file}.tmp and ${NVCC_generated_dependency_file}" + COMMAND "${CMAKE_COMMAND}" -E remove "${cmake_dependency_file}.tmp" "${NVCC_generated_dependency_file}" + ) + +if(CUDA_result) + message(FATAL_ERROR "Error generating ${generated_file}") +endif() + +# Generate the code +cuda_execute_process( + "Generating ${generated_file}" + COMMAND "${CUDA_NVCC_EXECUTABLE}" + "${source_file}" + ${cuda_language_flag} + ${format_flag} -o "${generated_file}" + ${CCBIN} + ${nvcc_flags} + ${nvcc_host_compiler_flags} + ${CUDA_NVCC_FLAGS} + -DNVCC + ${CUDA_NVCC_INCLUDE_ARGS} + ) + +if(CUDA_result) + # Since nvcc can sometimes leave half done files make sure that we delete the output file. + cuda_execute_process( + "Removing ${generated_file}" + COMMAND "${CMAKE_COMMAND}" -E remove "${generated_file}" + ) + message(FATAL_ERROR "Error generating file ${generated_file}") +else() + if(verbose) + message("Generated ${generated_file} successfully.") + endif() +endif() + +# Cubin resource report commands. +if( build_cubin ) + # Run with -cubin to produce resource usage report. + cuda_execute_process( + "Generating ${generated_cubin_file}" + COMMAND "${CUDA_NVCC_EXECUTABLE}" + "${source_file}" + ${CUDA_NVCC_FLAGS} + ${nvcc_flags} + ${CCBIN} + ${nvcc_host_compiler_flags} + -DNVCC + -cubin + -o "${generated_cubin_file}" + ${CUDA_NVCC_INCLUDE_ARGS} + ) + + # Execute the parser script. + cuda_execute_process( + "Executing the parser script" + COMMAND "${CMAKE_COMMAND}" + -D "input_file:STRING=${generated_cubin_file}" + -P "${CUDA_parse_cubin}" + ) + +endif() + +cmake_policy(POP) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake new file mode 100644 index 0000000000000000000000000000000000000000..642b96ba119e85ccedf2f286e1dd2da7ecd0da20 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake @@ -0,0 +1,297 @@ +# Synopsis: +# CUDA_SELECT_NVCC_ARCH_FLAGS(out_variable [target_CUDA_architectures]) +# -- Selects GPU arch flags for nvcc based on target_CUDA_architectures +# target_CUDA_architectures : Auto | Common | All | LIST(ARCH_AND_PTX ...) +# - "Auto" detects local machine GPU compute arch at runtime. +# - "Common" and "All" cover common and entire subsets of architectures +# ARCH_AND_PTX : NAME | NUM.NUM | NUM.NUM(NUM.NUM) | NUM.NUM+PTX +# NAME: Kepler Maxwell Kepler+Tegra Kepler+Tesla Maxwell+Tegra Pascal Volta Turing Ampere +# NUM: Any number. Only those pairs are currently accepted by NVCC though: +# 3.5 3.7 5.0 5.2 5.3 6.0 6.2 7.0 7.2 7.5 8.0 +# Returns LIST of flags to be added to CUDA_NVCC_FLAGS in ${out_variable} +# Additionally, sets ${out_variable}_readable to the resulting numeric list +# Example: +# CUDA_SELECT_NVCC_ARCH_FLAGS(ARCH_FLAGS 3.0 3.5+PTX 5.2(5.0) Maxwell) +# LIST(APPEND CUDA_NVCC_FLAGS ${ARCH_FLAGS}) +# +# More info on CUDA architectures: https://en.wikipedia.org/wiki/CUDA +# + +if(CMAKE_CUDA_COMPILER_LOADED) # CUDA as a language + if(CMAKE_CUDA_COMPILER_ID STREQUAL "NVIDIA" + AND CMAKE_CUDA_COMPILER_VERSION MATCHES "^([0-9]+\\.[0-9]+)") + set(CUDA_VERSION "${CMAKE_MATCH_1}") + endif() +endif() + +# See: https://docs.nvidia.com/cuda/cuda-compiler-driver-nvcc/index.html#gpu-feature-list + +# This list will be used for CUDA_ARCH_NAME = All option +set(CUDA_KNOWN_GPU_ARCHITECTURES "Kepler" "Maxwell") + +# This list will be used for CUDA_ARCH_NAME = Common option (enabled by default) +set(CUDA_COMMON_GPU_ARCHITECTURES "5.0") + +# This list is used to filter CUDA archs when autodetecting +set(CUDA_ALL_GPU_ARCHITECTURES "5.0") + +if(CUDA_VERSION VERSION_GREATER "10.5") + list(APPEND CUDA_KNOWN_GPU_ARCHITECTURES "Ampere") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "8.0") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "8.0") + + if(CUDA_VERSION VERSION_LESS "11.1") + set(CUDA_LIMIT_GPU_ARCHITECTURE "8.0") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "8.0+PTX") + endif() +endif() + +if(NOT CUDA_VERSION VERSION_LESS "11.1") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "8.6") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "8.6") + set(CUDA_LIMIT_GPU_ARCHITECUTRE "8.6") + + if(CUDA_VERSION VERSION_LESS "11.8") + set(CUDA_LIMIT_GPU_ARCHITECTURE "8.9") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "8.6+PTX") + endif() +endif() + +if(NOT CUDA_VERSION VERSION_LESS "11.8") + list(APPEND CUDA_KNOWN_GPU_ARCHITECTURES "Ada") + list(APPEND CUDA_KNOWN_GPU_ARCHITECTURES "Hopper") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "8.9") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "9.0") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "8.9") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "9.0") + +endif() + +list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "9.0a") +list(APPEND CUDA_ALL_GPU_ARCHITECTURES "9.0a") + +if(CUDA_VERSION VERSION_GREATER "12.6") + list(APPEND CUDA_KNOWN_GPU_ARCHITECTURES "Blackwell") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "10.0") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "10.0a") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "10.1a") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "12.0") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "12.0a") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "10.0") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "10.0a") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "10.1a") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "12.0") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "12.0a") + if(NOT CUDA_VERSION VERSION_LESS "13.0") + list(REMOVE_ITEM CUDA_COMMON_GPU_ARCHITECTURES "10.1a") + list(REMOVE_ITEM CUDA_ALL_GPU_ARCHITECTURES "10.1a") + list(APPEND CUDA_COMMON_GPU_ARCHITECTURES "11.0a") + list(APPEND CUDA_ALL_GPU_ARCHITECTURES "11.0a") + endif() +endif() + + +################################################################################################ +# A function for automatic detection of GPUs installed (if autodetection is enabled) +# Usage: +# CUDA_DETECT_INSTALLED_GPUS(OUT_VARIABLE) +# +function(CUDA_DETECT_INSTALLED_GPUS OUT_VARIABLE) + if(NOT CUDA_GPU_DETECT_OUTPUT) + if(CMAKE_CUDA_COMPILER_LOADED) # CUDA as a language + set(file "${PROJECT_BINARY_DIR}/detect_cuda_compute_capabilities.cu") + else() + set(file "${PROJECT_BINARY_DIR}/detect_cuda_compute_capabilities.cpp") + endif() + + file(WRITE ${file} "" + "#include \n" + "#include \n" + "int main()\n" + "{\n" + " int count = 0;\n" + " if (cudaSuccess != cudaGetDeviceCount(&count)) return -1;\n" + " if (count == 0) return -1;\n" + " for (int device = 0; device < count; ++device)\n" + " {\n" + " cudaDeviceProp prop;\n" + " if (cudaSuccess == cudaGetDeviceProperties(&prop, device))\n" + " std::printf(\"%d.%d \", prop.major, prop.minor);\n" + " }\n" + " return 0;\n" + "}\n") + + if(CMAKE_CUDA_COMPILER_LOADED) # CUDA as a language + try_run(run_result compile_result ${PROJECT_BINARY_DIR} ${file} + RUN_OUTPUT_VARIABLE compute_capabilities) + else() + try_run(run_result compile_result ${PROJECT_BINARY_DIR} ${file} + CMAKE_FLAGS "-DINCLUDE_DIRECTORIES=${CUDA_INCLUDE_DIRS}" + LINK_LIBRARIES ${CUDA_LIBRARIES} + RUN_OUTPUT_VARIABLE compute_capabilities) + endif() + + # Filter unrelated content out of the output. + string(REGEX MATCHALL "[0-9]+\\.[0-9]+" compute_capabilities "${compute_capabilities}") + + if(run_result EQUAL 0) + string(REPLACE "2.1" "2.1(2.0)" compute_capabilities "${compute_capabilities}") + set(CUDA_GPU_DETECT_OUTPUT ${compute_capabilities} + CACHE INTERNAL "Returned GPU architectures from detect_gpus tool" FORCE) + endif() + endif() + + if(NOT CUDA_GPU_DETECT_OUTPUT) + message(STATUS "Automatic GPU detection failed. Building for common architectures.") + set(${OUT_VARIABLE} ${CUDA_COMMON_GPU_ARCHITECTURES} PARENT_SCOPE) + else() + # Filter based on CUDA version supported archs + set(CUDA_GPU_DETECT_OUTPUT_FILTERED "") + separate_arguments(CUDA_GPU_DETECT_OUTPUT) + foreach(ITEM IN ITEMS ${CUDA_GPU_DETECT_OUTPUT}) + if(CUDA_LIMIT_GPU_ARCHITECTURE AND (ITEM VERSION_GREATER CUDA_LIMIT_GPU_ARCHITECTURE OR + ITEM VERSION_EQUAL CUDA_LIMIT_GPU_ARCHITECTURE)) + list(GET CUDA_COMMON_GPU_ARCHITECTURES -1 NEWITEM) + string(APPEND CUDA_GPU_DETECT_OUTPUT_FILTERED " ${NEWITEM}") + else() + string(APPEND CUDA_GPU_DETECT_OUTPUT_FILTERED " ${ITEM}") + endif() + endforeach() + + set(${OUT_VARIABLE} ${CUDA_GPU_DETECT_OUTPUT_FILTERED} PARENT_SCOPE) + endif() +endfunction() + + +################################################################################################ +# Function for selecting GPU arch flags for nvcc based on CUDA architectures from parameter list +# Usage: +# SELECT_NVCC_ARCH_FLAGS(out_variable [list of CUDA compute archs]) +function(CUDA_SELECT_NVCC_ARCH_FLAGS out_variable) + set(CUDA_ARCH_LIST "${ARGN}") + + if("X${CUDA_ARCH_LIST}" STREQUAL "X" ) + set(CUDA_ARCH_LIST "Auto") + endif() + + set(cuda_arch_bin) + set(cuda_arch_ptx) + + if("${CUDA_ARCH_LIST}" STREQUAL "All") + set(CUDA_ARCH_LIST ${CUDA_KNOWN_GPU_ARCHITECTURES}) + elseif("${CUDA_ARCH_LIST}" STREQUAL "Common") + set(CUDA_ARCH_LIST ${CUDA_COMMON_GPU_ARCHITECTURES}) + elseif("${CUDA_ARCH_LIST}" STREQUAL "Auto") + CUDA_DETECT_INSTALLED_GPUS(CUDA_ARCH_LIST) + message(STATUS "Autodetected CUDA architecture(s): ${CUDA_ARCH_LIST}") + endif() + + # Now process the list and look for names + string(REGEX REPLACE "[ \t]+" ";" CUDA_ARCH_LIST "${CUDA_ARCH_LIST}") + list(REMOVE_DUPLICATES CUDA_ARCH_LIST) + foreach(arch_name ${CUDA_ARCH_LIST}) + set(arch_bin) + set(arch_ptx) + set(add_ptx FALSE) + # Check to see if we are compiling PTX + if(arch_name MATCHES "(.*)\\+PTX$") + set(add_ptx TRUE) + set(arch_name ${CMAKE_MATCH_1}) + endif() + if(arch_name MATCHES "^([0-9]+\\.[0-9][af]?(\\([0-9]+\\.[0-9]\\))?)$") + set(arch_bin ${CMAKE_MATCH_1}) + set(arch_ptx ${arch_bin}) + else() + # Look for it in our list of known architectures + if(${arch_name} STREQUAL "Kepler+Tesla") + set(arch_bin 3.7) + elseif(${arch_name} STREQUAL "Kepler") + set(arch_bin 3.5) + set(arch_ptx 3.5) + elseif(${arch_name} STREQUAL "Maxwell+Tegra") + set(arch_bin 5.3) + elseif(${arch_name} STREQUAL "Maxwell") + set(arch_bin 5.0 5.2) + set(arch_ptx 5.2) + elseif(${arch_name} STREQUAL "Pascal") + set(arch_bin 6.0 6.1) + set(arch_ptx 6.1) + elseif(${arch_name} STREQUAL "Volta+Tegra") + set(arch_bin 7.2) + elseif(${arch_name} STREQUAL "Volta") + set(arch_bin 7.0 7.0) + set(arch_ptx 7.0) + elseif(${arch_name} STREQUAL "Turing") + set(arch_bin 7.5) + set(arch_ptx 7.5) + elseif(${arch_name} STREQUAL "Ampere+Tegra") + set(arch_bin 8.7) + elseif(${arch_name} STREQUAL "Ampere") + set(arch_bin 8.0 8.6) + set(arch_ptx 8.0 8.6) + elseif(${arch_name} STREQUAL "Ada") + set(arch_bin 8.9) + set(arch_ptx 8.9) + elseif(${arch_name} STREQUAL "Hopper") + set(arch_bin 9.0) + set(arch_ptx 9.0) + elseif(${arch_name} STREQUAL "Blackwell+Tegra") + set(arch_bin 10.1) + elseif(${arch_name} STREQUAL "Blackwell") + set(arch_bin 10.0 12.0) + set(arch_ptx 10.0 12.0) + else() + message(SEND_ERROR "Found Unknown CUDA Architecture Name in CUDA_SELECT_NVCC_ARCH_FLAGS: ${arch_name} ") + endif() + endif() + if(NOT arch_bin) + message(SEND_ERROR "arch_bin wasn't set for some reason") + endif() + list(APPEND cuda_arch_bin ${arch_bin}) + if(add_ptx) + if (NOT arch_ptx) + set(arch_ptx ${arch_bin}) + endif() + list(APPEND cuda_arch_ptx ${arch_ptx}) + endif() + endforeach() + + # remove dots and convert to lists + string(REGEX REPLACE "\\." "" cuda_arch_bin "${cuda_arch_bin}") + string(REGEX REPLACE "\\." "" cuda_arch_ptx "${cuda_arch_ptx}") + string(REGEX MATCHALL "[0-9()]+[af]?" cuda_arch_bin "${cuda_arch_bin}") + string(REGEX MATCHALL "[0-9]+[af]?" cuda_arch_ptx "${cuda_arch_ptx}") + + if(cuda_arch_bin) + list(REMOVE_DUPLICATES cuda_arch_bin) + endif() + if(cuda_arch_ptx) + list(REMOVE_DUPLICATES cuda_arch_ptx) + endif() + + set(nvcc_flags "") + set(nvcc_archs_readable "") + + # Tell NVCC to add binaries for the specified GPUs + foreach(arch ${cuda_arch_bin}) + if(arch MATCHES "([0-9]+)\\(([0-9]+)\\)") + # User explicitly specified ARCH for the concrete CODE + list(APPEND nvcc_flags -gencode arch=compute_${CMAKE_MATCH_2},code=sm_${CMAKE_MATCH_1}) + list(APPEND nvcc_archs_readable sm_${CMAKE_MATCH_1}) + else() + # User didn't explicitly specify ARCH for the concrete CODE, we assume ARCH=CODE + list(APPEND nvcc_flags -gencode arch=compute_${arch},code=sm_${arch}) + list(APPEND nvcc_archs_readable sm_${arch}) + endif() + endforeach() + + # Tell NVCC to add PTX intermediate code for the specified architectures + foreach(arch ${cuda_arch_ptx}) + list(APPEND nvcc_flags -gencode arch=compute_${arch},code=compute_${arch}) + list(APPEND nvcc_archs_readable compute_${arch}) + endforeach() + + string(REPLACE ";" " " nvcc_archs_readable "${nvcc_archs_readable}") + set(${out_variable} ${nvcc_flags} PARENT_SCOPE) + set(${out_variable}_readable ${nvcc_archs_readable} PARENT_SCOPE) +endfunction() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindPackageMessage.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindPackageMessage.cmake new file mode 100644 index 0000000000000000000000000000000000000000..6821cee4f77a9d84c74f2c140870a2163ae5a5f0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/Modules_CUDA_fix/upstream/FindPackageMessage.cmake @@ -0,0 +1,47 @@ +# Distributed under the OSI-approved BSD 3-Clause License. See accompanying +# file Copyright.txt or https://cmake.org/licensing for details. + +#.rst: +# FindPackageMessage +# ------------------ +# +# +# +# FIND_PACKAGE_MESSAGE( "message for user" "find result details") +# +# This macro is intended to be used in FindXXX.cmake modules files. It +# will print a message once for each unique find result. This is useful +# for telling the user where a package was found. The first argument +# specifies the name (XXX) of the package. The second argument +# specifies the message to display. The third argument lists details +# about the find result so that if they change the message will be +# displayed again. The macro also obeys the QUIET argument to the +# find_package command. +# +# Example: +# +# :: +# +# if(X11_FOUND) +# FIND_PACKAGE_MESSAGE(X11 "Found X11: ${X11_X11_LIB}" +# "[${X11_X11_LIB}][${X11_INCLUDE_DIR}]") +# else() +# ... +# endif() + +function(FIND_PACKAGE_MESSAGE pkg msg details) + # Avoid printing a message repeatedly for the same find result. + if(NOT ${pkg}_FIND_QUIETLY) + string(REPLACE "\n" "" details "${details}") + set(DETAILS_VAR FIND_PACKAGE_MESSAGE_DETAILS_${pkg}) + if(NOT "${details}" STREQUAL "${${DETAILS_VAR}}") + # The message has not yet been printed. + message(STATUS "${msg}") + + # Save the find details in the cache to avoid printing the same + # message again. + set("${DETAILS_VAR}" "${details}" + CACHE INTERNAL "Details about finding ${pkg}") + endif() + endif() +endfunction() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/LoadHIP.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/LoadHIP.cmake new file mode 100644 index 0000000000000000000000000000000000000000..a87b16f5ba889277ad72e8f735cef71fe103a006 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/LoadHIP.cmake @@ -0,0 +1,287 @@ +set(PYTORCH_FOUND_HIP FALSE) + +# If ROCM_PATH is set, assume intention is to compile with +# ROCm support and error out if the ROCM_PATH does not exist. +# Else ROCM_PATH does not exist, try to get it from rocm-sdk, +# or assume a default of /opt/rocm +# In the latter case, if /opt/rocm does not exist emit status +# message and return. +if(DEFINED ENV{ROCM_PATH}) + file(TO_CMAKE_PATH "$ENV{ROCM_PATH}" ROCM_PATH) + if(NOT EXISTS ${ROCM_PATH}) + message(FATAL_ERROR + "ROCM_PATH environment variable is set to ${ROCM_PATH} but does not exist.\n" + "Set a valid ROCM_PATH or unset ROCM_PATH environment variable to fix.") + endif() +else() + # Try to get ROCM_PATH from rocm-sdk if available + find_program(ROCM_SDK_EXECUTABLE rocm-sdk) + if(ROCM_SDK_EXECUTABLE) + execute_process( + COMMAND ${ROCM_SDK_EXECUTABLE} path --root + OUTPUT_VARIABLE ROCM_SDK_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + RESULT_VARIABLE ROCM_SDK_RESULT + ERROR_QUIET + ) + if(ROCM_SDK_RESULT EQUAL 0 AND EXISTS "${ROCM_SDK_PATH}") + set(ROCM_PATH "${ROCM_SDK_PATH}") + message(STATUS "Found ROCm installation via rocm-sdk at: ${ROCM_PATH}") + endif() + endif() + + # Fall back to default paths if rocm-sdk did not work + if(NOT DEFINED ROCM_PATH OR NOT EXISTS ${ROCM_PATH}) + if(UNIX) + set(ROCM_PATH /opt/rocm) + else() # Win32 + set(ROCM_PATH C:/opt/rocm) + endif() + endif() + + if(NOT EXISTS ${ROCM_PATH}) + message(STATUS + "ROCM_PATH environment variable is not set and ${ROCM_PATH} does not exist.\n" + "Building without ROCm support.") + return() + endif() +endif() + +# MAGMA_HOME +if(NOT DEFINED ENV{MAGMA_HOME}) + set(MAGMA_HOME ${ROCM_PATH}/magma) + set(ENV{MAGMA_HOME} ${ROCM_PATH}/magma) +else() + file(TO_CMAKE_PATH "$ENV{MAGMA_HOME}" MAGMA_HOME) +endif() + +# MIOpen isn't a part of HIP-SDK for Windows and hence, may have a different +# installation directory. +if(WIN32) + if(NOT DEFINED ENV{MIOPEN_PATH}) + set(miopen_DIR C:/opt/miopen/lib/cmake/miopen) + else() + set(miopen_DIR $ENV{MIOPEN_PATH}/lib/cmake/miopen) + endif() +endif() + +torch_hip_get_arch_list(PYTORCH_ROCM_ARCH) +if(PYTORCH_ROCM_ARCH STREQUAL "") + message(FATAL_ERROR "No GPU arch specified for ROCm build. Please use PYTORCH_ROCM_ARCH environment variable to specify GPU archs to build for.") +endif() +message("Building PyTorch for GPU arch: ${PYTORCH_ROCM_ARCH}") + +# Add HIP to the CMAKE Module Path +# needed because the find_package call to this module uses the Module mode search +# https://cmake.org/cmake/help/latest/command/find_package.html#search-modes +if(UNIX) + set(CMAKE_MODULE_PATH ${ROCM_PATH}/lib/cmake/hip;${ROCM_PATH}/lib/${CMAKE_LIBRARY_ARCHITECTURE}/cmake/hip ${CMAKE_MODULE_PATH}) +else() # Win32 + set(CMAKE_MODULE_PATH ${ROCM_PATH}/cmake/ ${CMAKE_MODULE_PATH}) +endif() + +# Add ROCM_PATH to CMAKE_PREFIX_PATH, needed because the find_package +# call to individual ROCM components uses the Config mode search +list(APPEND CMAKE_PREFIX_PATH ${ROCM_PATH}) + +macro(find_package_and_print_version PACKAGE_NAME) + find_package("${PACKAGE_NAME}" ${ARGN}) + if(NOT ${PACKAGE_NAME}_FOUND) + message("Optional package ${PACKAGE_NAME} not found") + else() + message("${PACKAGE_NAME} VERSION: ${${PACKAGE_NAME}_VERSION}") + if(${PACKAGE_NAME}_INCLUDE_DIR) + list(APPEND ROCM_INCLUDE_DIRS ${${PACKAGE_NAME}_INCLUDE_DIR}) + endif() + endif() +endmacro() + +# Find the HIP Package +# MODULE argument is added for clarity that CMake is searching +# for FindHIP.cmake in Module mode +find_package_and_print_version(HIP 1.0 MODULE) +if(CMAKE_SYSTEM_NAME STREQUAL "Linux") + enable_language(HIP) +endif() + +if(HIP_FOUND) + set(PYTORCH_FOUND_HIP TRUE) + find_package_and_print_version(hip REQUIRED CONFIG) + if(HIP_VERSION) + # Check if HIP_VERSION contains a dash (e.g., "7.1.25421-32f9fa6ca5") + # and strip everything after it to get clean numeric version + string(FIND "${HIP_VERSION}" "-" DASH_POS) + if(NOT DASH_POS EQUAL -1) + string(SUBSTRING "${HIP_VERSION}" 0 ${DASH_POS} HIP_VERSION_CLEAN) + set(HIP_VERSION "${HIP_VERSION_CLEAN}") + else() + set(HIP_VERSION_CLEAN "${HIP_VERSION}") + endif() + message("HIP version: ${HIP_VERSION}") + else() + set(HIP_VERSION_CLEAN "") + endif() + +# The rocm-core package was only introduced in ROCm 6.4, so we make it optional. + find_package(rocm-core CONFIG) + + # Some old consumer HIP SDKs do not distribute rocm_version.h, so we allow + # falling back to the hip version, which everyone should have. + # rocm_version.h lives in the rocm-core package and hip_version.h lives in the + # hip (lower-case) package. Both are probed above and will be in + # ROCM_INCLUDE_DIRS if available. + find_file(ROCM_VERSION_HEADER_PATH + NAMES rocm-core/rocm_version.h hip/hip_version.h + NO_DEFAULT_PATH + PATHS ${ROCM_INCLUDE_DIRS} + ) + if(ROCM_VERSION_HEADER_PATH MATCHES "rocm-core/rocm_version.h$") + set(ROCM_LIB_NAME "ROCM") + else() + set(ROCM_LIB_NAME "HIP") + endif() + + if(NOT ROCM_VERSION_HEADER_PATH) + message(FATAL_ERROR "Could not find hip/hip_version.h or rocm-core/rocm_version.h in ${ROCM_INCLUDE_DIRS}") + endif() + get_filename_component(ROCM_HEADER_NAME ${ROCM_VERSION_HEADER_PATH} NAME) + + if(EXISTS ${ROCM_VERSION_HEADER_PATH}) + set(ROCM_HEADER_FILE ${ROCM_VERSION_HEADER_PATH}) + else() + message(FATAL_ERROR "********************* ${ROCM_HEADER_NAME} could not be found ******************\n") + endif() + + # Read the ROCM headerfile into a variable + message(STATUS "Reading ROCM version from: ${ROCM_HEADER_FILE}") + message(STATUS "Content: ${ROCM_HEADER_CONTENT}") + file(READ "${ROCM_HEADER_FILE}" ROCM_HEADER_CONTENT) + + # Below we use a RegEx to find ROCM version numbers. + # Note that CMake does not support \s for blank space. That is + # why in the regular expressions below we have a blank space in + # the square brackets. + # There are three steps: + # 1. Match regular expression + # 2. Strip the non-numerical part of the string + # 3. Strip leading and trailing spaces + + string(REGEX MATCH "${ROCM_LIB_NAME}_VERSION_MAJOR[ ]+[0-9]+" TEMP1 ${ROCM_HEADER_CONTENT}) + string(REPLACE "${ROCM_LIB_NAME}_VERSION_MAJOR" "" TEMP2 ${TEMP1}) + string(STRIP ${TEMP2} ROCM_VERSION_DEV_MAJOR) + string(REGEX MATCH "${ROCM_LIB_NAME}_VERSION_MINOR[ ]+[0-9]+" TEMP1 ${ROCM_HEADER_CONTENT}) + string(REPLACE "${ROCM_LIB_NAME}_VERSION_MINOR" "" TEMP2 ${TEMP1}) + string(STRIP ${TEMP2} ROCM_VERSION_DEV_MINOR) + string(REGEX MATCH "${ROCM_LIB_NAME}_VERSION_PATCH[ ]+[0-9]+" TEMP1 ${ROCM_HEADER_CONTENT}) + string(REPLACE "${ROCM_LIB_NAME}_VERSION_PATCH" "" TEMP2 ${TEMP1}) + string(STRIP ${TEMP2} ROCM_VERSION_DEV_PATCH) + + # Create ROCM_VERSION_DEV_INT which is later used as a preprocessor macros + set(ROCM_VERSION_DEV "${ROCM_VERSION_DEV_MAJOR}.${ROCM_VERSION_DEV_MINOR}.${ROCM_VERSION_DEV_PATCH}") + math(EXPR ROCM_VERSION_DEV_INT "(${ROCM_VERSION_DEV_MAJOR}*10000) + (${ROCM_VERSION_DEV_MINOR}*100) + ${ROCM_VERSION_DEV_PATCH}") + + message("\n***** ROCm version from ${ROCM_HEADER_NAME} ****\n") + message("ROCM_VERSION_DEV: ${ROCM_VERSION_DEV}") + message("ROCM_VERSION_DEV_MAJOR: ${ROCM_VERSION_DEV_MAJOR}") + message("ROCM_VERSION_DEV_MINOR: ${ROCM_VERSION_DEV_MINOR}") + message("ROCM_VERSION_DEV_PATCH: ${ROCM_VERSION_DEV_PATCH}") + message("ROCM_VERSION_DEV_INT: ${ROCM_VERSION_DEV_INT}") + + math(EXPR TORCH_HIP_VERSION "(${HIP_VERSION_MAJOR} * 100) + ${HIP_VERSION_MINOR}") + message("HIP_VERSION_MAJOR: ${HIP_VERSION_MAJOR}") + message("HIP_VERSION_MINOR: ${HIP_VERSION_MINOR}") + message("TORCH_HIP_VERSION: ${TORCH_HIP_VERSION}") + + # Find ROCM components using Config mode + # These components will be searced for recursively in ${ROCM_PATH} + message("\n***** Library versions from cmake find_package *****\n") + find_package_and_print_version(amd_comgr REQUIRED) + find_package_and_print_version(rocrand REQUIRED) + find_package_and_print_version(hiprand REQUIRED) + find_package_and_print_version(rocblas REQUIRED) + find_package_and_print_version(hipblas REQUIRED) + find_package_and_print_version(miopen REQUIRED) + find_package_and_print_version(hipfft REQUIRED) + find_package_and_print_version(hipsparse REQUIRED) + find_package_and_print_version(rocprim REQUIRED) + find_package_and_print_version(hipcub REQUIRED) + find_package_and_print_version(rocthrust REQUIRED) + find_package_and_print_version(hipsolver REQUIRED) + find_package_and_print_version(rocsolver REQUIRED) + find_package_and_print_version(rocshmem) + # workaround cmake 4 build issue + if(CMAKE_VERSION VERSION_GREATER_EQUAL "4.0.0") + message(WARNING "Work around hiprtc cmake failure for cmake >= 4") + set(CMAKE_POLICY_VERSION_MINIMUM 3.5) + find_package_and_print_version(hiprtc REQUIRED) + unset(CMAKE_POLICY_VERSION_MINIMUM) + else() + find_package_and_print_version(hiprtc REQUIRED) + endif() + find_package_and_print_version(hipblaslt REQUIRED) + + if(UNIX) + find_package_and_print_version(rccl) + find_package_and_print_version(hsa-runtime64 REQUIRED) + find_package_and_print_version(rocm_smi REQUIRED) + endif() + + # Optional components. + find_package_and_print_version(hipsparselt) # Will be required when ready. + + list(REMOVE_DUPLICATES ROCM_INCLUDE_DIRS) + + if(UNIX) + # roctx is part of roctracer + find_library(ROCM_ROCTX_LIB roctx64 HINTS ${ROCM_PATH}/lib) + + set(PROJECT_RANDOM_BINARY_DIR "${PROJECT_BINARY_DIR}") + + if(ROCM_VERSION_DEV VERSION_GREATER_EQUAL "5.7.0") + # check whether hipblaslt provides HIPBLASLT_MATMUL_MATRIX_SCALE_OUTER_VEC_32F + set(file "${PROJECT_BINARY_DIR}/hipblaslt_test_outer_vec.cc") + file(WRITE ${file} "" + "#define LEGACY_HIPBLAS_DIRECT\n" + "#include \n" + "int main() {\n" + " hipblasLtMatmulMatrixScale_t attr = HIPBLASLT_MATMUL_MATRIX_SCALE_OUTER_VEC_32F;\n" + " return 0;\n" + "}\n" + ) + try_compile(hipblaslt_compile_result_outer_vec ${PROJECT_RANDOM_BINARY_DIR} ${file} + CMAKE_FLAGS "-DINCLUDE_DIRECTORIES=${ROCM_INCLUDE_DIRS}" + COMPILE_DEFINITIONS -D__HIP_PLATFORM_AMD__ -D__HIP_PLATFORM_HCC__ + OUTPUT_VARIABLE hipblaslt_compile_output_outer_vec) + + # check whether hipblaslt provides HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER_VEC_EXT + set(file "${PROJECT_BINARY_DIR}/hipblaslt_test_vec_ext.cc") + file(WRITE ${file} "" + "#define LEGACY_HIPBLAS_DIRECT\n" + "#include \n" + "int main() {\n" + " hipblasLtMatmulDescAttributes_t attr = HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER_VEC_EXT;\n" + " return 0;\n" + "}\n" + ) + try_compile(hipblaslt_compile_result_vec_ext ${PROJECT_RANDOM_BINARY_DIR} ${file} + CMAKE_FLAGS "-DINCLUDE_DIRECTORIES=${ROCM_INCLUDE_DIRS}" + COMPILE_DEFINITIONS -D__HIP_PLATFORM_AMD__ -D__HIP_PLATFORM_HCC__ + OUTPUT_VARIABLE hipblaslt_compile_output_vec_ext) + + if(hipblaslt_compile_result_outer_vec) + set(HIPBLASLT_OUTER_VEC ON) + set(HIPBLASLT_VEC_EXT OFF) + message("hipblaslt is using scale pointer outer vec") + elseif(hipblaslt_compile_result_vec_ext) + set(HIPBLASLT_OUTER_VEC OFF) + set(HIPBLASLT_VEC_EXT ON) + message("hipblaslt is using scale pointer vec ext") + else() + set(HIPBLASLT_OUTER_VEC OFF) + set(HIPBLASLT_VEC_EXT OFF) + message("hipblaslt is NOT using scale pointer outer vec: ${hipblaslt_compile_output_outer_vec}") + message("hipblaslt is NOT using scale pointer vec ext: ${hipblaslt_compile_output_vec_ext}") + endif() + endif() + endif() +endif() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/cuda.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/cuda.cmake new file mode 100644 index 0000000000000000000000000000000000000000..239641d352c41894a236982880cfac77b4a0020e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/cuda.cmake @@ -0,0 +1,404 @@ +# ---[ cuda + +# Poor man's include guard +if(TARGET torch::cudart) + return() +endif() + +# sccache is only supported in CMake master and not in the newest official +# release (3.11.3) yet. Hence we need our own Modules_CUDA_fix to enable sccache. +list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_LIST_DIR}/../Modules_CUDA_fix) + +# We don't want to statically link cudart, because we rely on it's dynamic linkage in +# python (follow along torch/cuda/__init__.py and usage of cudaGetErrorName). +# Technically, we can link cudart here statically, and link libtorch_python.so +# to a dynamic libcudart.so, but that's just wasteful. +# However, on Windows, if this one gets switched off, the error "cuda: unknown error" +# will be raised when running the following code: +# >>> import torch +# >>> torch.cuda.is_available() +# >>> torch.cuda.current_device() +# More details can be found in the following links. +# https://github.com/pytorch/pytorch/issues/20635 +# https://github.com/pytorch/pytorch/issues/17108 +if(NOT MSVC) + set(CUDA_USE_STATIC_CUDA_RUNTIME OFF CACHE INTERNAL "") +endif() + +# Find CUDA. +find_package(CUDA) +if(NOT CUDA_FOUND) + # If user explicitly set USE_CUDA=1, error out instead of falling back + if(_USE_CUDA_EXPLICITLY_SET AND USE_CUDA) + message(FATAL_ERROR + "PyTorch: CUDA was explicitly requested (USE_CUDA=1) but cannot be found. " + "Please check your CUDA installation, ensure CUDA toolkit is installed, " + "and that CUDA_HOME or CMAKE_CUDA_COMPILER is set correctly. " + "If you want to build without CUDA, please set USE_CUDA=0.") + endif() + + message(WARNING + "PyTorch: CUDA cannot be found. Depending on whether you are building " + "PyTorch or a PyTorch dependent library, the next warning / error will " + "give you more info.") + set(CAFFE2_USE_CUDA OFF) + return() +endif() + +# Enable CUDA language support +set(CUDAToolkit_ROOT "${CUDA_TOOLKIT_ROOT_DIR}") +# Pass clang as host compiler, which according to the docs +# Must be done before CUDA language is enabled, see +# https://cmake.org/cmake/help/v3.15/variable/CMAKE_CUDA_HOST_COMPILER.html +if("${CMAKE_CXX_COMPILER_ID}" MATCHES "Clang") + set(CMAKE_CUDA_HOST_COMPILER "${CMAKE_CXX_COMPILER}") +endif() +enable_language(CUDA) +if("X${CMAKE_CUDA_STANDARD}" STREQUAL "X" ) + set(CMAKE_CUDA_STANDARD ${CMAKE_CXX_STANDARD}) +endif() +set(CMAKE_CUDA_STANDARD_REQUIRED ON) + +# CMP0074 - find_package will respect _ROOT variables +cmake_policy(PUSH) +if(CMAKE_VERSION VERSION_GREATER_EQUAL 3.12.0) + cmake_policy(SET CMP0074 NEW) +endif() + +find_package(CUDAToolkit REQUIRED) + +cmake_policy(POP) + +if(NOT CMAKE_CUDA_COMPILER_VERSION VERSION_EQUAL CUDAToolkit_VERSION) + message(FATAL_ERROR "Found two conflicting CUDA versions:\n" + "V${CMAKE_CUDA_COMPILER_VERSION} in '${CUDA_INCLUDE_DIRS}' and\n" + "V${CUDAToolkit_VERSION} in '${CUDAToolkit_INCLUDE_DIRS}'") +endif() + +message(STATUS "PyTorch: CUDA detected: " ${CUDA_VERSION}) +message(STATUS "PyTorch: CUDA nvcc is: " ${CUDA_NVCC_EXECUTABLE}) +message(STATUS "PyTorch: CUDA toolkit directory: " ${CUDA_TOOLKIT_ROOT_DIR}) +if(CUDA_VERSION VERSION_LESS 12.1) + message(FATAL_ERROR "PyTorch requires CUDA 12.1 or above.") +endif() + +if(CUDA_FOUND) + # Sometimes, we may mismatch nvcc with the CUDA headers we are + # compiling with, e.g., if a ccache nvcc is fed to us by CUDA_NVCC_EXECUTABLE + # but the PATH is not consistent with CUDA_HOME. It's better safe + # than sorry: make sure everything is consistent. + if(MSVC AND CMAKE_GENERATOR MATCHES "Visual Studio") + # When using Visual Studio, it attempts to lock the whole binary dir when + # `try_run` is called, which will cause the build to fail. + string(RANDOM BUILD_SUFFIX) + set(PROJECT_RANDOM_BINARY_DIR "${PROJECT_BINARY_DIR}/${BUILD_SUFFIX}") + else() + set(PROJECT_RANDOM_BINARY_DIR "${PROJECT_BINARY_DIR}") + endif() + set(file "${PROJECT_BINARY_DIR}/detect_cuda_version.cc") + file(WRITE ${file} "" + "#include \n" + "#include \n" + "int main() {\n" + " printf(\"%d.%d\", CUDA_VERSION / 1000, (CUDA_VERSION / 10) % 100);\n" + " return 0;\n" + "}\n" + ) + if(NOT CMAKE_CROSSCOMPILING) + try_run(run_result compile_result ${PROJECT_RANDOM_BINARY_DIR} ${file} + CMAKE_FLAGS "-DINCLUDE_DIRECTORIES=${CUDA_INCLUDE_DIRS}" + LINK_LIBRARIES ${CUDA_LIBRARIES} + RUN_OUTPUT_VARIABLE cuda_version_from_header + COMPILE_OUTPUT_VARIABLE output_var + ) + if(NOT compile_result) + message(FATAL_ERROR "PyTorch: Couldn't determine version from header: " ${output_var}) + endif() + message(STATUS "PyTorch: Header version is: " ${cuda_version_from_header}) + if(NOT cuda_version_from_header STREQUAL ${CUDA_VERSION_STRING}) + # Force CUDA to be processed for again next time + # TODO: I'm not sure if this counts as an implementation detail of + # FindCUDA + set(cuda_version_from_findcuda ${CUDA_VERSION_STRING}) + unset(CUDA_TOOLKIT_ROOT_DIR_INTERNAL CACHE) + # Not strictly necessary, but for good luck. + unset(CUDA_VERSION CACHE) + # Error out + message(FATAL_ERROR "FindCUDA says CUDA version is ${cuda_version_from_findcuda} (usually determined by nvcc), " + "but the CUDA headers say the version is ${cuda_version_from_header}. This often occurs " + "when you set both CUDA_HOME and CUDA_NVCC_EXECUTABLE to " + "non-standard locations, without also setting PATH to point to the correct nvcc. " + "Perhaps, try re-running this command again with PATH=${CUDA_TOOLKIT_ROOT_DIR}/bin:$PATH. " + "See above log messages for more diagnostics, and see https://github.com/pytorch/pytorch/issues/8092 for more details.") + endif() + endif() +endif() + +# ---[ CUDA libraries wrapper + +# find lbnvrtc.so +set(CUDA_NVRTC_LIB "${CUDA_nvrtc_LIBRARY}" CACHE FILEPATH "") +if(CUDA_NVRTC_LIB AND NOT CUDA_NVRTC_SHORTHASH) + find_package(Python COMPONENTS Interpreter) + execute_process( + COMMAND "${Python_EXECUTABLE}" -c + "import hashlib;hash=hashlib.sha256();hash.update(open('${CUDA_NVRTC_LIB}','rb').read());print(hash.hexdigest()[:8])" + RESULT_VARIABLE _retval + OUTPUT_VARIABLE CUDA_NVRTC_SHORTHASH) + if(NOT _retval EQUAL 0) + message(WARNING "Failed to compute shorthash for libnvrtc.so") + set(CUDA_NVRTC_SHORTHASH "XXXXXXXX") + else() + string(STRIP "${CUDA_NVRTC_SHORTHASH}" CUDA_NVRTC_SHORTHASH) + message(STATUS "${CUDA_NVRTC_LIB} shorthash is ${CUDA_NVRTC_SHORTHASH}") + endif() +endif() + +# Create new style imported libraries. +# Several of these libraries have a hardcoded path if CAFFE2_STATIC_LINK_CUDA +# is set. This path is where sane CUDA installations have their static +# libraries installed. This flag should only be used for binary builds, so +# end-users should never have this flag set. + +# cuda +add_library(caffe2::cuda INTERFACE IMPORTED) +set_property( + TARGET caffe2::cuda PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cuda_driver) + +# cudart +add_library(torch::cudart INTERFACE IMPORTED) +if(CAFFE2_STATIC_LINK_CUDA) + set_property( + TARGET torch::cudart PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cudart_static) +else() + set_property( + TARGET torch::cudart PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cudart) +endif() + + +# cublas +add_library(caffe2::cublas INTERFACE IMPORTED) +if(CAFFE2_STATIC_LINK_CUDA AND NOT WIN32) + set_property( + TARGET caffe2::cublas PROPERTY INTERFACE_LINK_LIBRARIES + # NOTE: cublas is always linked dynamically + CUDA::cublas CUDA::cublasLt) + set_property( + TARGET caffe2::cublas APPEND PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cudart_static rt) +else() + set_property( + TARGET caffe2::cublas PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cublas CUDA::cublasLt) +endif() + +# cudnn interface +# static linking is handled by USE_STATIC_CUDNN environment variable +if(CAFFE2_USE_CUDNN) + if(USE_STATIC_CUDNN) + set(CUDNN_STATIC ON CACHE BOOL "") + else() + set(CUDNN_STATIC OFF CACHE BOOL "") + endif() + + find_package(CUDNN) + + if(NOT CUDNN_FOUND) + message(WARNING + "Cannot find cuDNN library. Turning the option off") + set(CAFFE2_USE_CUDNN OFF) + else() + if(CUDNN_VERSION VERSION_LESS "8.1.0") + message(FATAL_ERROR "PyTorch requires cuDNN 8.1 and above.") + endif() + endif() + + add_library(torch::cudnn INTERFACE IMPORTED) + target_include_directories(torch::cudnn INTERFACE ${CUDNN_INCLUDE_PATH}) + if(CUDNN_STATIC AND NOT WIN32) + target_link_options(torch::cudnn INTERFACE + "-Wl,--exclude-libs,libcudnn_static.a") + else() + target_link_libraries(torch::cudnn INTERFACE ${CUDNN_LIBRARY_PATH}) + endif() +else() + message(STATUS "USE_CUDNN is set to 0. Compiling without cuDNN support") +endif() + +if(CAFFE2_USE_CUSPARSELT) + find_package(CUSPARSELT) + + if(NOT CUSPARSELT_FOUND) + message(WARNING + "Cannot find cuSPARSELt library. Turning the option off") + set(CAFFE2_USE_CUSPARSELT OFF) + else() + add_library(torch::cusparselt INTERFACE IMPORTED) + target_include_directories(torch::cusparselt INTERFACE ${CUSPARSELT_INCLUDE_PATH}) + target_link_libraries(torch::cusparselt INTERFACE ${CUSPARSELT_LIBRARY_PATH}) + endif() +else() + message(STATUS "USE_CUSPARSELT is set to 0. Compiling without cuSPARSELt support") +endif() + +if(USE_CUDSS) + find_package(CUDSS) + + if(NOT CUDSS_FOUND) + message(WARNING + "Cannot find CUDSS library. Turning the option off") + set(USE_CUDSS OFF) + else() + add_library(torch::cudss INTERFACE IMPORTED) + target_include_directories(torch::cudss INTERFACE ${CUDSS_INCLUDE_PATH}) + target_link_libraries(torch::cudss INTERFACE ${CUDSS_LIBRARY_PATH}) + endif() +else() + message(STATUS "USE_CUDSS is set to 0. Compiling without cuDSS support") +endif() + +# cufile +if(CAFFE2_USE_CUFILE) + add_library(torch::cufile INTERFACE IMPORTED) + if(CAFFE2_STATIC_LINK_CUDA AND NOT WIN32) + set_property( + TARGET torch::cufile PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cuFile_static) + else() + set_property( + TARGET torch::cufile PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cuFile) + endif() +else() + message(STATUS "USE_CUFILE is set to 0. Compiling without cuFile support") +endif() + +# curand +add_library(caffe2::curand INTERFACE IMPORTED) +if(CAFFE2_STATIC_LINK_CUDA AND NOT WIN32) + set_property( + TARGET caffe2::curand PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::curand_static) +else() + set_property( + TARGET caffe2::curand PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::curand) +endif() + +# cufft +add_library(caffe2::cufft INTERFACE IMPORTED) +if(CAFFE2_STATIC_LINK_CUDA AND NOT WIN32) + if(CUDA_VERSION VERSION_LESS_EQUAL 12.9) + set_property( + TARGET caffe2::cufft PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cufft_static_nocallback) + else() + set_property( + TARGET caffe2::cufft PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cufft_static) + endif() +else() + set_property( + TARGET caffe2::cufft PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::cufft) +endif() + +# nvrtc +# cuDNN frontend needs libnvrtc symbols, but linking through CUDA::nvrtc pulls +# CUDA::cuda_driver transitively. Keep a driver-free target for cuDNN users and +# reserve caffe2::nvrtc for the stub library that actually needs the driver API. +add_library(caffe2::nvrtc_runtime INTERFACE IMPORTED) +set_property( + TARGET caffe2::nvrtc_runtime PROPERTY INTERFACE_LINK_LIBRARIES + "${CUDA_NVRTC_LIB}") + +add_library(caffe2::nvrtc INTERFACE IMPORTED) +set_property( + TARGET caffe2::nvrtc PROPERTY INTERFACE_LINK_LIBRARIES + CUDA::nvrtc caffe2::cuda) + +# Add onnx namespace definition to nvcc +if(ONNX_NAMESPACE) + list(APPEND CUDA_NVCC_FLAGS "-DONNX_NAMESPACE=${ONNX_NAMESPACE}") +else() + list(APPEND CUDA_NVCC_FLAGS "-DONNX_NAMESPACE=onnx_c2") +endif() + +# Don't activate VC env again for Ninja generators with MSVC on Windows if CUDAHOSTCXX is not defined +# by adding --use-local-env. +if(MSVC AND CMAKE_GENERATOR STREQUAL "Ninja" AND NOT DEFINED ENV{CUDAHOSTCXX}) + list(APPEND CUDA_NVCC_FLAGS "--use-local-env") +endif() + +# setting nvcc arch flags +torch_cuda_get_nvcc_gencode_flag(NVCC_FLAGS_EXTRA) +# CMake 3.18 adds integrated support for architecture selection, but we can't rely on it +if(DEFINED CMAKE_CUDA_ARCHITECTURES) + message(WARNING + "pytorch is not compatible with `CMAKE_CUDA_ARCHITECTURES` and will ignore its value. " + "Please configure `TORCH_CUDA_ARCH_LIST` instead.") + set(CMAKE_CUDA_ARCHITECTURES OFF) +endif() + +list(APPEND CUDA_NVCC_FLAGS ${NVCC_FLAGS_EXTRA}) +message(STATUS "Added CUDA NVCC flags for: ${NVCC_FLAGS_EXTRA}") + +# disable some nvcc diagnostic that appears in boost, glog, glags, opencv, etc. +foreach(diag cc_clobber_ignored + field_without_dll_interface + base_class_has_different_dll_interface + dll_interface_conflict_none_assumed + dll_interface_conflict_dllexport_assumed + bad_friend_decl) + list(APPEND SUPPRESS_WARNING_FLAGS --diag_suppress=${diag}) +endforeach() +string(REPLACE ";" "," SUPPRESS_WARNING_FLAGS "${SUPPRESS_WARNING_FLAGS}") +list(APPEND CUDA_NVCC_FLAGS -Xcudafe ${SUPPRESS_WARNING_FLAGS}) + +set(CUDA_PROPAGATE_HOST_FLAGS_BLOCKLIST "-Werror") +if(MSVC) + list(APPEND CUDA_NVCC_FLAGS "--Werror" "cross-execution-space-call") + list(APPEND CUDA_NVCC_FLAGS "--no-host-device-move-forward") +endif() + +# Debug and Release symbol support +if(MSVC) + if(${CAFFE2_USE_MSVC_STATIC_RUNTIME}) + string(APPEND CMAKE_CUDA_FLAGS_DEBUG " -Xcompiler /MTd") + string(APPEND CMAKE_CUDA_FLAGS_MINSIZEREL " -Xcompiler /MT") + string(APPEND CMAKE_CUDA_FLAGS_RELEASE " -Xcompiler /MT") + string(APPEND CMAKE_CUDA_FLAGS_RELWITHDEBINFO " -Xcompiler /MT") + else() + string(APPEND CMAKE_CUDA_FLAGS_DEBUG " -Xcompiler /MDd") + string(APPEND CMAKE_CUDA_FLAGS_MINSIZEREL " -Xcompiler /MD") + string(APPEND CMAKE_CUDA_FLAGS_RELEASE " -Xcompiler /MD") + string(APPEND CMAKE_CUDA_FLAGS_RELWITHDEBINFO " -Xcompiler /MD") + endif() + if(CUDA_NVCC_FLAGS MATCHES "Zi") + list(APPEND CUDA_NVCC_FLAGS "-Xcompiler" "-FS") + endif() +elseif(CUDA_DEVICE_DEBUG) + list(APPEND CUDA_NVCC_FLAGS "-g" "-G") # -G enables device code debugging symbols +endif() + +# needed for compat with newer versions of clang that use C++20 mangling rules +if(CMAKE_CXX_COMPILER_ID MATCHES "Clang" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 18) + list(APPEND CUDA_NVCC_FLAGS "-Xcompiler=-fclang-abi-compat=17") +endif() + +# Set expt-relaxed-constexpr to suppress Eigen warnings +list(APPEND CUDA_NVCC_FLAGS "--expt-relaxed-constexpr") + +# Set expt-extended-lambda to support lambda on device +list(APPEND CUDA_NVCC_FLAGS "--expt-extended-lambda") + +foreach(FLAG ${CUDA_NVCC_FLAGS}) + string(FIND "${FLAG}" " " flag_space_position) + if(NOT flag_space_position EQUAL -1) + message(FATAL_ERROR "Found spaces in CUDA_NVCC_FLAGS entry '${FLAG}'") + endif() + string(APPEND CMAKE_CUDA_FLAGS " ${FLAG}") +endforeach() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/gflags.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/gflags.cmake new file mode 100644 index 0000000000000000000000000000000000000000..186cda1a909ab79431114d1c61de895069255389 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/gflags.cmake @@ -0,0 +1,83 @@ +# ---[ gflags + +# We will try to use the config mode first, and then manual find. +find_package(gflags CONFIG QUIET) +if(NOT TARGET gflags) + find_package(gflags MODULE QUIET) +endif() + +if(TARGET gflags) + message(STATUS "Caffe2: Found gflags with new-style gflags target.") +elseif(GFLAGS_FOUND) + message(STATUS "Caffe2: Found gflags with old-style gflag starget.") + add_library(gflags UNKNOWN IMPORTED) + set_property( + TARGET gflags PROPERTY IMPORTED_LOCATION ${GFLAGS_LIBRARY}) + set_property( + TARGET gflags PROPERTY INTERFACE_INCLUDE_DIRECTORIES + ${GFLAGS_INCLUDE_DIR}) +else() + message(STATUS + "Caffe2: Cannot find gflags automatically. Using legacy find.") + + # - Try to find GFLAGS in the legacy way. + # + # The following variables are optionally searched for defaults + # GFLAGS_ROOT_DIR: Base directory where all GFLAGS components are found + # + # The following are set after configuration is done: + # GFLAGS_FOUND + # GFLAGS_INCLUDE_DIRS + # GFLAGS_LIBRARIES + # GFLAGS_LIBRARYRARY_DIRS + include(FindPackageHandleStandardArgs) + set(GFLAGS_ROOT_DIR "" CACHE PATH "Folder contains Gflags") + + # We are testing only a couple of files in the include directories + if(WIN32) + find_path(GFLAGS_INCLUDE_DIR gflags/gflags.h + PATHS ${GFLAGS_ROOT_DIR}/src/windows) + else() + find_path(GFLAGS_INCLUDE_DIR gflags/gflags.h + PATHS ${GFLAGS_ROOT_DIR}) + endif() + + if(WIN32) + find_library(GFLAGS_LIBRARY_RELEASE + NAMES libgflags + PATHS ${GFLAGS_ROOT_DIR} + PATH_SUFFIXES Release) + + find_library(GFLAGS_LIBRARY_DEBUG + NAMES libgflags-debug + PATHS ${GFLAGS_ROOT_DIR} + PATH_SUFFIXES Debug) + set(GFLAGS_LIBRARY optimized ${GFLAGS_LIBRARY_RELEASE} debug ${GFLAGS_LIBRARY_DEBUG}) + else() + find_library(GFLAGS_LIBRARY gflags) + endif() + + find_package_handle_standard_args( + gflags DEFAULT_MSG GFLAGS_INCLUDE_DIR GFLAGS_LIBRARY) + + if(GFLAGS_FOUND) + message( + STATUS + "Caffe2: Found gflags (include: ${GFLAGS_INCLUDE_DIR}, " + "library: ${GFLAGS_LIBRARY})") + add_library(gflags UNKNOWN IMPORTED) + set_property( + TARGET gflags PROPERTY IMPORTED_LOCATION ${GFLAGS_LIBRARY}) + set_property( + TARGET gflags PROPERTY INTERFACE_INCLUDE_DIRECTORIES + ${GFLAGS_INCLUDE_DIR}) + endif() +endif() + +# After above, we should have the gflags target now. +if(NOT TARGET gflags) + message(WARNING + "Caffe2: gflags cannot be found. Depending on whether you are building " + "Caffe2 or a Caffe2 dependent library, the next warning / error will " + "give you more info.") +endif() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/glog.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/glog.cmake new file mode 100644 index 0000000000000000000000000000000000000000..bb03e81f29e3afed43ba95260cc5c298be881f72 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/glog.cmake @@ -0,0 +1,70 @@ +# ---[ glog + +# We will try to use the config mode first, and then manual find. +find_package(glog CONFIG QUIET) +if(NOT TARGET glog::glog) + find_package(glog MODULE QUIET) +endif() + +if(TARGET glog::glog) + message(STATUS "Caffe2: Found glog with new-style glog target.") +elseif(GLOG_FOUND) + message( + STATUS + "Caffe2: Found glog with old-style glog starget. Glog never shipped " + "old style glog targets, so somewhere in your cmake path there might " + "be a custom Findglog.cmake file that got triggered. We will make a " + "best effort to create the new style glog target for you.") + add_library(glog::glog UNKNOWN IMPORTED) + set_property( + TARGET glog::glog PROPERTY IMPORTED_LOCATION ${GLOG_LIBRARY}) + set_property( + TARGET glog::glog PROPERTY INTERFACE_INCLUDE_DIRECTORIES + ${GLOG_INCLUDE_DIR}) +else() + message(STATUS "Caffe2: Cannot find glog automatically. Using legacy find.") + + # - Try to find Glog + # + # The following variables are optionally searched for defaults + # GLOG_ROOT_DIR: Base directory where all GLOG components are found + # + # The following are set after configuration is done: + # GLOG_FOUND + # GLOG_INCLUDE_DIRS + # GLOG_LIBRARIES + # GLOG_LIBRARYRARY_DIRS + + include(FindPackageHandleStandardArgs) + set(GLOG_ROOT_DIR "" CACHE PATH "Folder contains Google glog") + if(NOT WIN32) + find_path(GLOG_INCLUDE_DIR glog/logging.h + PATHS ${GLOG_ROOT_DIR}) + endif() + + find_library(GLOG_LIBRARY glog + PATHS ${GLOG_ROOT_DIR} + PATH_SUFFIXES lib lib64) + + find_package_handle_standard_args(glog DEFAULT_MSG GLOG_INCLUDE_DIR GLOG_LIBRARY) + + if(GLOG_FOUND) + message(STATUS + "Caffe2: Found glog (include: ${GLOG_INCLUDE_DIR}, " + "library: ${GLOG_LIBRARY})") + add_library(glog::glog UNKNOWN IMPORTED) + set_property( + TARGET glog::glog PROPERTY IMPORTED_LOCATION ${GLOG_LIBRARY}) + set_property( + TARGET glog::glog PROPERTY INTERFACE_INCLUDE_DIRECTORIES + ${GLOG_INCLUDE_DIR}) + endif() +endif() + +# After above, we should have the glog::glog target now. +if(NOT TARGET glog::glog) + message(WARNING + "Caffe2: glog cannot be found. Depending on whether you are building " + "Caffe2 or a Caffe2 dependent library, the next warning / error will " + "give you more info.") +endif() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/mkl.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/mkl.cmake new file mode 100644 index 0000000000000000000000000000000000000000..2f6d1fd905aa303cc240b058318acdfb2483e9ad --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/mkl.cmake @@ -0,0 +1,40 @@ +find_package(MKL QUIET) + +if(TARGET caffe2::mkl) + return() +endif() + +add_library(caffe2::mkl INTERFACE IMPORTED) +target_include_directories(caffe2::mkl INTERFACE ${MKL_INCLUDE_DIR}) +target_link_libraries(caffe2::mkl INTERFACE ${MKL_LIBRARIES}) +foreach(MKL_LIB IN LISTS MKL_LIBRARIES) + if(EXISTS "${MKL_LIB}") + get_filename_component(MKL_LINK_DIR "${MKL_LIB}" DIRECTORY) + if(IS_DIRECTORY "${MKL_LINK_DIR}") + target_link_directories(caffe2::mkl INTERFACE "${MKL_LINK_DIR}") + endif() + endif() +endforeach() + +# TODO: This is a hack, it will not pick up architecture dependent +# MKL libraries correctly; see https://github.com/pytorch/pytorch/issues/73008 +set_property( + TARGET caffe2::mkl PROPERTY INTERFACE_LINK_DIRECTORIES + ${MKL_ROOT}/lib ${MKL_ROOT}/lib/intel64 ${MKL_ROOT}/lib/intel64_win ${MKL_ROOT}/lib/win-x64) + +if(UNIX) + if(USE_STATIC_MKL) + foreach(MKL_LIB_PATH IN LISTS MKL_LIBRARIES) + if(NOT EXISTS "${MKL_LIB_PATH}") + continue() + endif() + + get_filename_component(MKL_LIB_NAME "${MKL_LIB_PATH}" NAME) + + # Match archive libraries starting with "libmkl_" + if(MKL_LIB_NAME MATCHES "^libmkl_" AND MKL_LIB_NAME MATCHES ".a$") + target_link_options(caffe2::mkl INTERFACE "-Wl,--exclude-libs,${MKL_LIB_NAME}") + endif() + endforeach() + endif() +endif() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/mkldnn.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/mkldnn.cmake new file mode 100644 index 0000000000000000000000000000000000000000..87935625f9bfb543d1cdc7f2b59f11e8d4a709e7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/mkldnn.cmake @@ -0,0 +1,18 @@ +set(MKLDNN_USE_NATIVE_ARCH ${USE_NATIVE_ARCH}) + +if(CPU_AARCH64) + include(${CMAKE_CURRENT_LIST_DIR}/ComputeLibrary.cmake) +endif() + +find_package(MKLDNN QUIET) + +if(NOT TARGET caffe2::mkldnn) + add_library(caffe2::mkldnn INTERFACE IMPORTED) +endif() + +set_property( + TARGET caffe2::mkldnn PROPERTY INTERFACE_INCLUDE_DIRECTORIES + ${MKLDNN_INCLUDE_DIR}) +set_property( + TARGET caffe2::mkldnn PROPERTY INTERFACE_LINK_LIBRARIES + ${MKLDNN_LIBRARIES}) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/protobuf.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/protobuf.cmake new file mode 100644 index 0000000000000000000000000000000000000000..77ec3622b132dc7a7817716dd24ef986e6ac030d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/protobuf.cmake @@ -0,0 +1,92 @@ +# ---[ Protobuf + +# We will try to use the config mode first, and then manual find. +find_package(Protobuf CONFIG QUIET) +if(NOT Protobuf_FOUND) + find_package(Protobuf MODULE QUIET) +endif() + +if((TARGET protobuf::libprotobuf OR TARGET protobuf::libprotobuf-lite) AND TARGET protobuf::protoc) + # Hooray. This is the most ideal situation, meaning that you either have a + # Protobuf config file installed (like on Windows), or you are using a + # modern CMake that ships with a FindProtobuf.cmake file that produces + # modern targets. + message(STATUS "Caffe2: Found protobuf with new-style protobuf targets.") +elseif(Protobuf_FOUND OR PROTOBUF_FOUND) + # If the modern targets are not present, we will generate them for you for + # backward compatibility. This is backported from CMake's new FindProtobuf.cmake + # content. + if((NOT PROTOBUF_LIBRARY) AND (NOT PROTOBUF_LITE_LIBRARY)) + message(FATAL_ERROR + "Caffe2: Found protobuf with old style targets, but could not find targets." + " PROTOBUF_LIBRARY: " ${PROTOBUF_LIBRARY} + " PROTOBUF_LITE_LIBRARY: " ${PROTOBUF_LITE_LIBRARY} + " Protobuf_LIBRARY: " ${Protobuf_LIBRARY} + " Protobuf_LITE_LIBRARY: " ${Protobuf_LITE_LIBRARY}) + endif() + message(STATUS "Caffe2: Found protobuf with old-style protobuf targets.") + + if(PROTOBUF_LIBRARY) + if(NOT TARGET protobuf::libprotobuf) + add_library(protobuf::libprotobuf UNKNOWN IMPORTED) + set_target_properties(protobuf::libprotobuf PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${PROTOBUF_INCLUDE_DIRS}") + endif() + if(EXISTS "${PROTOBUF_LIBRARY}") + set_target_properties(protobuf::libprotobuf PROPERTIES + IMPORTED_LOCATION "${PROTOBUF_LIBRARY}") + endif() + if(EXISTS "${PROTOBUF_LIBRARY_RELEASE}") + set_property(TARGET protobuf::libprotobuf APPEND PROPERTY + IMPORTED_CONFIGURATIONS RELEASE) + set_target_properties(protobuf::libprotobuf PROPERTIES + IMPORTED_LOCATION_RELEASE "${PROTOBUF_LIBRARY_RELEASE}") + endif() + if(EXISTS "${PROTOBUF_LIBRARY_DEBUG}") + set_property(TARGET protobuf::libprotobuf APPEND PROPERTY + IMPORTED_CONFIGURATIONS DEBUG) + set_target_properties(protobuf::libprotobuf PROPERTIES + IMPORTED_LOCATION_DEBUG "${PROTOBUF_LIBRARY_DEBUG}") + endif() + endif() + + if(PROTOBUF_LITE_LIBRARY) + if(NOT TARGET protobuf::libprotobuf-lite) + add_library(protobuf::libprotobuf-lite UNKNOWN IMPORTED) + set_target_properties(protobuf::libprotobuf-lite PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${PROTOBUF_INCLUDE_DIRS}") + endif() + if(EXISTS "${PROTOBUF_LITE_LIBRARY}") + set_target_properties(protobuf::libprotobuf-lite PROPERTIES + IMPORTED_LOCATION "${PROTOBUF_LITE_LIBRARY}") + endif() + if(EXISTS "${PROTOBUF_LITE_LIBRARY_RELEASE}") + set_property(TARGET protobuf::libprotobuf-lite APPEND PROPERTY + IMPORTED_CONFIGURATIONS RELEASE) + set_target_properties(protobuf::libprotobuf-lite PROPERTIES + IMPORTED_LOCATION_RELEASE "${PROTOBUF_LITE_LIBRARY_RELEASE}") + endif() + if(EXISTS "${PROTOBUF_LITE_LIBRARY_DEBUG}") + set_property(TARGET protobuf::libprotobuf-lite APPEND PROPERTY + IMPORTED_CONFIGURATIONS DEBUG) + set_target_properties(protobuf::libprotobuf-lite PROPERTIES + IMPORTED_LOCATION_DEBUG "${PROTOBUF_LITE_LIBRARY_DEBUG}") + endif() + endif() + + if(PROTOBUF_PROTOC_EXECUTABLE) + if(NOT TARGET protobuf::protoc) + add_executable(protobuf::protoc IMPORTED) + endif() + set_property(TARGET protobuf::protoc PROPERTY + IMPORTED_LOCATION ${PROTOBUF_PROTOC_EXECUTABLE}) + endif() +endif() + +# After above, we should have the protobuf related target now. +if((NOT TARGET protobuf::libprotobuf) AND (NOT TARGET protobuf::libprotobuf-lite)) + message(WARNING + "Protobuf cannot be found. Depending on whether you are building Caffe2 " + "or a Caffe2 dependent library, the next warning / error will give you " + "more info.") +endif() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/utils.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/utils.cmake new file mode 100644 index 0000000000000000000000000000000000000000..a667ad55771b456bc63df59ed20b8ef4222448b5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/utils.cmake @@ -0,0 +1,568 @@ +################################################################################################ +# Exclude and prepend functionalities +function(exclude OUTPUT INPUT) +set(EXCLUDES ${ARGN}) +foreach(EXCLUDE ${EXCLUDES}) + list(REMOVE_ITEM INPUT "${EXCLUDE}") +endforeach() +set(${OUTPUT} ${INPUT} PARENT_SCOPE) +endfunction(exclude) + +function(prepend OUTPUT PREPEND) +set(OUT "") +foreach(ITEM ${ARGN}) + list(APPEND OUT "${PREPEND}${ITEM}") +endforeach() +set(${OUTPUT} ${OUT} PARENT_SCOPE) +endfunction(prepend) + +################################################################################################ +# Parses a version string that might have values beyond major, minor, and patch +# and set version variables for the library. +# Usage: +# caffe2_parse_version_str( ) +function(caffe2_parse_version_str LIBNAME VERSIONSTR) + string(REGEX REPLACE "^([0-9]+).*$" "\\1" ${LIBNAME}_VERSION_MAJOR "${VERSIONSTR}") + string(REGEX REPLACE "^[0-9]+\\.([0-9]+).*$" "\\1" ${LIBNAME}_VERSION_MINOR "${VERSIONSTR}") + string(REGEX REPLACE "[0-9]+\\.[0-9]+\\.([0-9]+).*$" "\\1" ${LIBNAME}_VERSION_PATCH "${VERSIONSTR}") + set(${LIBNAME}_VERSION_MAJOR ${${LIBNAME}_VERSION_MAJOR} ${ARGN} PARENT_SCOPE) + set(${LIBNAME}_VERSION_MINOR ${${LIBNAME}_VERSION_MINOR} ${ARGN} PARENT_SCOPE) + set(${LIBNAME}_VERSION_PATCH ${${LIBNAME}_VERSION_PATCH} ${ARGN} PARENT_SCOPE) + set(${LIBNAME}_VERSION "${${LIBNAME}_VERSION_MAJOR}.${${LIBNAME}_VERSION_MINOR}.${${LIBNAME}_VERSION_PATCH}" PARENT_SCOPE) +endfunction() + +### +# Removes common indentation from a block of text to produce code suitable for +# setting to `python -c`, or using with pycmd. This allows multiline code to be +# nested nicely in the surrounding code structure. +# +# This function respsects Python_EXECUTABLE if it defined, otherwise it uses +# `python` and hopes for the best. An error will be thrown if it is not found. +# +# Args: +# outvar : variable that will hold the stdout of the python command +# text : text to remove indentation from +# +function(dedent outvar text) + # Use Python_EXECUTABLE if it is defined, otherwise default to python + if("${Python_EXECUTABLE}" STREQUAL "") + set(_python_exe "python3") + else() + set(_python_exe "${Python_EXECUTABLE}") + endif() + set(_fixup_cmd "import sys; from textwrap import dedent; print(dedent(sys.stdin.read()))") + file(WRITE "${CMAKE_BINARY_DIR}/indented.txt" "${text}") + execute_process( + COMMAND "${_python_exe}" -c "${_fixup_cmd}" + INPUT_FILE "${CMAKE_BINARY_DIR}/indented.txt" + RESULT_VARIABLE _dedent_exitcode + OUTPUT_VARIABLE _dedent_text) + if(NOT _dedent_exitcode EQUAL 0) + message(ERROR " Failed to remove indentation from: \n\"\"\"\n${text}\n\"\"\" + Python dedent failed with error code: ${_dedent_exitcode}") + message(FATAL_ERROR " Python dedent failed with error code: ${_dedent_exitcode}") + endif() + # Remove supurflous newlines (artifacts of print) + string(STRIP "${_dedent_text}" _dedent_text) + set(${outvar} "${_dedent_text}" PARENT_SCOPE) +endfunction() + + +function(pycmd_no_exit outvar exitcode cmd) + # Use Python_EXECUTABLE if it is defined, otherwise default to python + if("${Python_EXECUTABLE}" STREQUAL "") + set(_python_exe "python") + else() + set(_python_exe "${Python_EXECUTABLE}") + endif() + # run the actual command + execute_process( + COMMAND "${_python_exe}" -c "${cmd}" + RESULT_VARIABLE _exitcode + OUTPUT_VARIABLE _output) + # Remove supurflous newlines (artifacts of print) + string(STRIP "${_output}" _output) + set(${outvar} "${_output}" PARENT_SCOPE) + set(${exitcode} "${_exitcode}" PARENT_SCOPE) +endfunction() + + +### +# Helper function to run `python -c ""` and capture the results of stdout +# +# Runs a python command and populates an outvar with the result of stdout. +# Common indentation in the text of `cmd` is removed before the command is +# executed, so the caller does not need to worry about indentation issues. +# +# This function respsects Python_EXECUTABLE if it defined, otherwise it uses +# `python` and hopes for the best. An error will be thrown if it is not found. +# +# Args: +# outvar : variable that will hold the stdout of the python command +# cmd : text representing a (possibly multiline) block of python code +# +function(pycmd outvar cmd) + dedent(_dedent_cmd "${cmd}") + pycmd_no_exit(_output _exitcode "${_dedent_cmd}") + + if(NOT _exitcode EQUAL 0) + message(ERROR " Failed when running python code: \"\"\"\n${_dedent_cmd}\n\"\"\"") + message(FATAL_ERROR " Python command failed with error code: ${_exitcode}") + endif() + # Remove supurflous newlines (artifacts of print) + string(STRIP "${_output}" _output) + set(${outvar} "${_output}" PARENT_SCOPE) +endfunction() + + +############################################################################## +# Macro to update cached options. +macro(caffe2_update_option variable value) + if(CAFFE2_CMAKE_BUILDING_WITH_MAIN_REPO) + get_property(__help_string CACHE ${variable} PROPERTY HELPSTRING) + set(${variable} ${value} CACHE BOOL ${__help_string} FORCE) + else() + set(${variable} ${value}) + endif() +endmacro() + + +############################################################################## +# Add an interface library definition that is dependent on the source. +# +# It's probably easiest to explain why this macro exists, by describing +# what things would look like if we didn't have this macro. +# +# Let's suppose we want to statically link against torch. We've defined +# a library in cmake called torch, and we might think that we just +# target_link_libraries(my-app PUBLIC torch). This will result in a +# linker argument 'libtorch.a' getting passed to the linker. +# +# Unfortunately, this link command is wrong! We have static +# initializers in libtorch.a that would get improperly pruned by +# the default link settings. What we actually need is for you +# to do -Wl,--whole-archive,libtorch.a -Wl,--no-whole-archive to ensure +# that we keep all symbols, even if they are (seemingly) not used. +# +# What caffe2_interface_library does is create an interface library +# that indirectly depends on the real library, but sets up the link +# arguments so that you get all of the extra link settings you need. +# The result is not a "real" library, and so we have to manually +# copy over necessary properties from the original target. +# +# (The discussion above is about static libraries, but a similar +# situation occurs for dynamic libraries: if no symbols are used from +# a dynamic library, it will be pruned unless you are --no-as-needed) +macro(caffe2_interface_library SRC DST) + add_library(${DST} INTERFACE) + add_dependencies(${DST} ${SRC}) + # Depending on the nature of the source library as well as the compiler, + # determine the needed compilation flags. + get_target_property(__src_target_type ${SRC} TYPE) + # Depending on the type of the source library, we will set up the + # link command for the specific SRC library. + if(${__src_target_type} STREQUAL "STATIC_LIBRARY") + # In the case of static library, we will need to add whole-static flags. + target_link_libraries(${DST} INTERFACE $) + # Link all interface link libraries of the src target as well. + # For static library, we need to explicitly depend on all the libraries + # that are the dependent library of the source library. Note that we cannot + # use the populated INTERFACE_LINK_LIBRARIES property, because if one of the + # dependent library is not a target, cmake creates a $ wrapper + # and then one is not able to find target "src". For more discussions, check + # https://cmake.org/Bug/print_bug_page.php?bug_id=15415 + # https://cmake.org/pipermail/cmake-developers/2013-May/019019.html + # Specifically the following quote + # + # """ + # For STATIC libraries we can define that the PUBLIC/PRIVATE/INTERFACE keys + # are ignored for linking and that it always populates both LINK_LIBRARIES + # LINK_INTERFACE_LIBRARIES. Note that for STATIC libraries the + # LINK_LIBRARIES property will not be used for anything except build-order + # dependencies. + # """ + target_link_libraries(${DST} INTERFACE + $) + elseif(${__src_target_type} STREQUAL "SHARED_LIBRARY") + if("${CMAKE_CXX_COMPILER_ID}" MATCHES "GNU") + target_link_libraries(${DST} INTERFACE + "-Wl,--no-as-needed,\"$\" -Wl,--as-needed") + else() + target_link_libraries(${DST} INTERFACE ${SRC}) + endif() + # Link all interface link libraries of the src target as well. + # For shared libraries, we can simply depend on the INTERFACE_LINK_LIBRARIES + # property of the target. + target_link_libraries(${DST} INTERFACE + $) + else() + message(FATAL_ERROR + "You made a CMake build file error: target " ${SRC} + " must be of type either STATIC_LIBRARY or SHARED_LIBRARY. However, " + "I got " ${__src_target_type} ".") + endif() + # For all other interface properties, manually inherit from the source target. + set_target_properties(${DST} PROPERTIES + INTERFACE_COMPILE_DEFINITIONS + $ + INTERFACE_COMPILE_OPTIONS + $ + INTERFACE_INCLUDE_DIRECTORIES + $ + INTERFACE_SYSTEM_INCLUDE_DIRECTORIES + $) +endmacro() + + +############################################################################## +# Creating a Caffe2 binary target with sources specified with relative path. +# Usage: +# caffe2_binary_target(target_name_or_src [] [] ...) +# If only target_name_or_src is specified, this target is build with one single +# source file and the target name is autogen from the filename. Otherwise, the +# target name is given by the first argument and the rest are the source files +# to build the target. +function(caffe2_binary_target target_name_or_src) + # https://cmake.org/cmake/help/latest/command/function.html + # Checking that ARGC is greater than # is the only way to ensure + # that ARGV# was passed to the function as an extra argument. + if(ARGC GREATER 1) + set(__target ${target_name_or_src}) + prepend(__srcs "${CMAKE_CURRENT_SOURCE_DIR}/" "${ARGN}") + else() + get_filename_component(__target ${target_name_or_src} NAME_WE) + prepend(__srcs "${CMAKE_CURRENT_SOURCE_DIR}/" "${target_name_or_src}") + endif() + add_executable(${__target} ${__srcs}) + target_link_libraries(${__target} torch_library) + # If we have Caffe2_MODULES defined, we will also link with the modules. + if(DEFINED Caffe2_MODULES) + target_link_libraries(${__target} ${Caffe2_MODULES}) + endif() + install(TARGETS ${__target} DESTINATION bin) +endfunction() + +function(caffe2_hip_binary_target target_name_or_src) + if(ARGC GREATER 1) + set(__target ${target_name_or_src}) + prepend(__srcs "${CMAKE_CURRENT_SOURCE_DIR}/" "${ARGN}") + else() + get_filename_component(__target ${target_name_or_src} NAME_WE) + prepend(__srcs "${CMAKE_CURRENT_SOURCE_DIR}/" "${target_name_or_src}") + endif() + + caffe2_binary_target(${target_name_or_src}) + + target_compile_options(${__target} PRIVATE ${HIP_CXX_FLAGS}) + target_include_directories(${__target} PRIVATE ${Caffe2_HIP_INCLUDE}) +endfunction() + + +############################################################################## +# Multiplex between adding libraries for CUDA versus HIP (AMD Software Stack). +# Usage: +# torch_cuda_based_add_library(cuda_target) +# +macro(torch_cuda_based_add_library cuda_target) + if(USE_ROCM) + hip_add_library(${cuda_target} ${ARGN}) + elseif(USE_CUDA) + add_library(${cuda_target} ${ARGN}) + else() + endif() +endmacro() + +############################################################################## +# Get the HIP arch flags specified by PYTORCH_ROCM_ARCH. +# Usage: +# torch_hip_get_arch_list(variable_to_store_flags) +# +macro(torch_hip_get_arch_list store_var) + if(DEFINED ENV{PYTORCH_ROCM_ARCH}) + set(_TMP $ENV{PYTORCH_ROCM_ARCH}) + else() + # Use arch of installed GPUs as default + execute_process(COMMAND "rocm_agent_enumerator" COMMAND bash "-c" "grep -v gfx000 | sort -u | xargs | tr -d '\n'" + RESULT_VARIABLE ROCM_AGENT_ENUMERATOR_RESULT + OUTPUT_VARIABLE ROCM_ARCH_INSTALLED) + if(NOT ROCM_AGENT_ENUMERATOR_RESULT EQUAL 0) + message(FATAL_ERROR " Could not detect ROCm arch for GPUs on machine. Result: '${ROCM_AGENT_ENUMERATOR_RESULT}'") + endif() + set(_TMP ${ROCM_ARCH_INSTALLED}) + endif() + string(REPLACE " " ";" ${store_var} "${_TMP}") +endmacro() + +############################################################################## +# Get the XPU arch flags specified by TORCH_XPU_ARCH_LIST. +# Usage: +# torch_xpu_get_arch_list(variable_to_store_flags) +# +macro(torch_xpu_get_arch_list store_var) + if(DEFINED ENV{TORCH_XPU_ARCH_LIST}) + set(${store_var} $ENV{TORCH_XPU_ARCH_LIST}) + endif() +endmacro() + +############################################################################## +# Get the NVCC arch flags specified by TORCH_CUDA_ARCH_LIST and CUDA_ARCH_NAME. +# Usage: +# torch_cuda_get_nvcc_gencode_flag(variable_to_store_flags) +# +macro(torch_cuda_get_nvcc_gencode_flag store_var) + # setting nvcc arch flags + # We need to support the explicitly and conveniently defined TORCH_CUDA_ARCH_LIST + if((NOT DEFINED TORCH_CUDA_ARCH_LIST) AND (DEFINED ENV{TORCH_CUDA_ARCH_LIST})) + set(TORCH_CUDA_ARCH_LIST $ENV{TORCH_CUDA_ARCH_LIST}) + endif() + if(DEFINED CUDA_ARCH_NAME) + message(WARNING + "CUDA_ARCH_NAME is no longer used. Use TORCH_CUDA_ARCH_LIST instead. " + "Right now, CUDA_ARCH_NAME is ${CUDA_ARCH_NAME} and " + "TORCH_CUDA_ARCH_LIST is ${TORCH_CUDA_ARCH_LIST}.") + if(NOT TORCH_CUDA_ARCH_LIST) + set(TORCH_CUDA_ARCH_LIST ${CUDA_ARCH_NAME}) + else() + list(APPEND TORCH_CUDA_ARCH_LIST ${CUDA_ARCH_NAME}) + endif() + endif() + + # Invoke cuda_select_nvcc_arch_flags from proper cmake FindCUDA. + cuda_select_nvcc_arch_flags(${store_var} ${TORCH_CUDA_ARCH_LIST}) +endmacro() + + +############################################################################## +# Add standard compile options. +# Usage: +# torch_compile_options(lib_name) +function(torch_compile_options libname) + set_property(TARGET ${libname} PROPERTY CXX_STANDARD 20) + + # until they can be unified, keep these lists synced with setup.py + if(MSVC) + + if(MSVC_Z7_OVERRIDE) + set(MSVC_DEBINFO_OPTION "/Z7") + else() + set(MSVC_DEBINFO_OPTION "/Zi") + endif() + + if(${MSVC_TOOLSET_VERSION} GREATER_EQUAL 142) + # Add /permissive- flag for conformance mode to the compiler. + # This will force more strict check to the code standard. + # 1. From MS official doc: https://learn.microsoft.com/en-us/cpp/build/reference/permissive-standards-conformance?view=msvc-170#remarks + # By default, the /permissive- option is set in new projects created by Visual Studio 2017 version 15.5 and later versions. + # We set the /permissive- flag from VS 2019 (MSVC_TOOLSET_VERSION 142) to avoid compiling issues for old toolkit. + # 2. For MSVC VERSION: https://cmake.org/cmake/help/latest/variable/MSVC_TOOLSET_VERSION.html + target_compile_options(${libname} PUBLIC $<$:/permissive->) + endif() + # This option enables a token-based preprocessor that conforms to C99 and C++11 and later standards. + # This option is available since VS 2017. + # For MS official doc: https://learn.microsoft.com/en-us/cpp/build/reference/zc-preprocessor + set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} /Zc:preprocessor" PARENT_SCOPE) + + target_compile_options(${libname} PUBLIC + $<$: + ${MSVC_RUNTIME_LIBRARY_OPTION} + $<$,$>:${MSVC_DEBINFO_OPTION}> + /EHsc + /bigobj> + ) + else() + set(private_compile_options + -Wall + -Wextra + -Wdeprecated + -Wunused + -Wno-unused-parameter + -Wno-missing-field-initializers + -Wno-array-bounds + -Wno-unknown-pragmas + -Wno-strict-overflow + -Wno-strict-aliasing + ) + if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + list(APPEND private_compile_options -Wredundant-move) + # -Wno-interference-size only exists in GCC 12+ + if(CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12) + list(APPEND private_compile_options -Wno-interference-size) + endif() + endif() + if(CMAKE_CXX_COMPILER_ID MATCHES "Clang") + if(NOT USE_CUDA) + # NS: One can not compile CUDA code with extra-semi flag as nvcc generates code like + # namespace MemoryOps_cu_d8602b38_109889 __attribute__((visibility("hidden"))) { }; + list(APPEND private_compile_options -Wextra-semi) + else() + # NVCC + clang15 reports deprecated copies from GPU lambda instantiations + list(APPEND private_compile_options -Wno-deprecated-copy) + # NVCC + clang18 reports spurious deprecated deprecated literal operator declaration when there were none + # I.e. failures look like torch/headeronly/util/complex.h:334:40: error: identifier '_if' preceded by whitespace in a literal operator declaration is deprecated + # but if one to look at the source code, there are no space there + list(APPEND private_compile_options -Wno-deprecated-literal-operator) + + endif() + list(APPEND private_compile_options -Wmove) + else() + list(APPEND private_compile_options + # Considered to be flaky. See the discussion at + # https://github.com/pytorch/pytorch/pull/9608 + -Wno-maybe-uninitialized) + endif() + + if(WERROR) + list(APPEND private_compile_options + -Werror + -Werror=ignored-attributes + -Werror=inconsistent-missing-override + -Werror=inconsistent-missing-destructor-override + -Werror=pedantic + -Werror=unused + -Wno-error=unused-parameter + ) + if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + list(APPEND private_compile_options -Werror=unused-but-set-variable -Werror=cpp) + endif() + if(CMAKE_CXX_COMPILER_ID MATCHES "Clang") + list(APPEND private_compile_options -Werror=macro-redefined -Werror=deprecated-copy-with-dtor) + endif() + endif() + endif() + + + target_compile_options(${libname} PRIVATE + $<$:${private_compile_options}>) + if(USE_CUDA) + foreach(option IN LISTS private_compile_options) + if(CMAKE_CUDA_HOST_COMPILER_ID STREQUAL "GNU") + if("${option}" STREQUAL "-Wextra-semi") + continue() + endif() + if("${option}" STREQUAL "-Wunused-private-field") + continue() + endif() + endif() + target_compile_options(${libname} PRIVATE $<$:-Xcompiler ${option}>) + endforeach() + endif() + + if(NOT WIN32 AND NOT USE_ASAN) + # Enable hidden visibility by default to make it easier to debug issues with + # TORCH_API annotations. Hidden visibility with selective default visibility + # behaves close enough to Windows' dllimport/dllexport. + # + # Unfortunately, hidden visibility messes up some ubsan warnings because + # templated classes crossing library boundary get duplicated (but identical) + # definitions. It's easier to just disable it. + target_compile_options(${libname} PRIVATE + $<$: -fvisibility=hidden>) + endif() + +endfunction() + +############################################################################## +# Set old-style FindCuda.cmake compile flags from modern CMake cuda flags. +# Usage: +# torch_update_find_cuda_flags() +function(torch_update_find_cuda_flags) + # Convert -O2 -Xcompiler="-O2 -Wall" to "-O2;-Xcompiler=-O2,-Wall" + if(USE_CUDA) + separate_arguments(FLAGS UNIX_COMMAND "${CMAKE_CUDA_FLAGS}") + string(REPLACE " " "," FLAGS "${FLAGS}") + set(CUDA_NVCC_FLAGS ${FLAGS} PARENT_SCOPE) + + separate_arguments(FLAGS_DEBUG UNIX_COMMAND "${CMAKE_CUDA_FLAGS_DEBUG}") + string(REPLACE " " "," FLAGS_DEBUG "${FLAGS_DEBUG}") + set(CUDA_NVCC_FLAGS_DEBUG "${FLAGS_DEBUG}" PARENT_SCOPE) + + separate_arguments(FLAGS_RELEASE UNIX_COMMAND "${CMAKE_CUDA_FLAGS_RELEASE}") + string(REPLACE " " "," FLAGS_RELEASE "${FLAGS_RELEASE}") + set(CUDA_NVCC_FLAGS_RELEASE "${FLAGS_RELEASE}" PARENT_SCOPE) + + separate_arguments(FLAGS_MINSIZEREL UNIX_COMMAND "${CMAKE_CUDA_FLAGS_MINSIZEREL}") + string(REPLACE " " "," FLAGS_MINSIZEREL "${FLAGS_MINSIZEREL}") + set(CUDA_NVCC_FLAGS_MINSIZEREL "${FLAGS_MINSIZEREL}" PARENT_SCOPE) + + separate_arguments(FLAGS_RELWITHDEBINFO UNIX_COMMAND "${CMAKE_CUDA_FLAGS_RELWITHDEBINFO}") + string(REPLACE " " "," FLAGS_RELWITHDEBINFO "${FLAGS_RELWITHDEBINFO}") + set(CUDA_NVCC_FLAGS_RELWITHDEBINFO "${FLAGS_RELWITHDEBINFO}" PARENT_SCOPE) + + message(STATUS "Converting CMAKE_CUDA_FLAGS to CUDA_NVCC_FLAGS:\n" + " CUDA_NVCC_FLAGS = ${FLAGS}\n" + " CUDA_NVCC_FLAGS_DEBUG = ${FLAGS_DEBUG}\n" + " CUDA_NVCC_FLAGS_RELEASE = ${FLAGS_RELEASE}\n" + " CUDA_NVCC_FLAGS_RELWITHDEBINFO = ${FLAGS_RELWITHDEBINFO}\n" + " CUDA_NVCC_FLAGS_MINSIZEREL = ${FLAGS_MINSIZEREL}") + endif() +endfunction() + +include(CheckCXXCompilerFlag) +include(CheckCCompilerFlag) +include(CheckLinkerFlag) + +############################################################################## +# CHeck if given flag is supported and append it to provided outputvar +# Also define HAS_UPPER_CASE_FLAG_NAME variable +# Usage: +# append_cxx_flag_if_supported("-Werror" CMAKE_CXX_FLAGS) +function(append_cxx_flag_if_supported flag outputvar) + string(TOUPPER "HAS${flag}" _FLAG_NAME) + string(REGEX REPLACE "[=-]" "_" _FLAG_NAME "${_FLAG_NAME}") + # GCC silents unknown -Wno-XXX flags, so we detect the corresponding -WXXX. + if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") + string(REGEX REPLACE "Wno-" "W" new_flag "${flag}") + else() + set(new_flag ${flag}) + endif() + check_cxx_compiler_flag("${new_flag}" ${_FLAG_NAME}) + if(${_FLAG_NAME}) + string(APPEND ${outputvar} " ${flag}") + set(${outputvar} "${${outputvar}}" PARENT_SCOPE) + endif() +endfunction() + +function(append_c_flag_if_supported flag outputvar) + string(TOUPPER "HAS${flag}" _FLAG_NAME) + string(REGEX REPLACE "[=-]" "_" _FLAG_NAME "${_FLAG_NAME}") + + # GCC silences unknown -Wno-XXX flags, so test the corresponding -WXXX. + if(CMAKE_C_COMPILER_ID STREQUAL "GNU") + string(REGEX REPLACE "^Wno-" "W" new_flag "${flag}") + else() + set(new_flag "${flag}") + endif() + + check_c_compiler_flag("${new_flag}" ${_FLAG_NAME}) + if(${_FLAG_NAME}) + string(APPEND ${outputvar} " ${flag}") + set(${outputvar} "${${outputvar}}" PARENT_SCOPE) + endif() +endfunction() + +function(target_compile_options_if_supported target flag) + set(_compile_options "") + append_cxx_flag_if_supported("${flag}" _compile_options) + if(NOT "${_compile_options}" STREQUAL "") + target_compile_options(${target} PRIVATE ${flag}) + endif() +endfunction() + +# Check if a global link option is supported +function(add_link_options_if_supported flag) + check_linker_flag(C "LINKER:${flag}" _supported) + if("${_supported}") + add_link_options("LINKER:${flag}") + else() + message(WARNING "Attempted to use unsupported link option : ${flag}.") + endif() +endfunction() + +function(target_link_options_if_supported tgt flag) + check_linker_flag(C "LINKER:${flag}" _supported) + if("${_supported}") + target_link_options("${tgt}" PRIVATE "LINKER:${flag}") + else() + message(WARNING "Attempted to use unsupported link option : ${flag}.") + endif() +endfunction() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/xpu.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/xpu.cmake new file mode 100644 index 0000000000000000000000000000000000000000..2731c2842c7f46bb029116e472f605fab08926e4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Caffe2/public/xpu.cmake @@ -0,0 +1,61 @@ +# ---[ xpu + +# Poor man's include guard +if(TARGET torch::xpurt) + return() +endif() + +set(XPU_HOST_CXX_FLAGS) + +# Find SYCL library. +find_package(SYCLToolkit REQUIRED) +if(NOT SYCL_FOUND) + set(PYTORCH_FOUND_XPU FALSE) + # Exit early to avoid populating XPU_HOST_CXX_FLAGS. + return() +endif() +set(PYTORCH_FOUND_XPU TRUE) + +# SYCL library interface +add_library(torch::sycl INTERFACE IMPORTED) + +set_property( + TARGET torch::sycl PROPERTY INTERFACE_INCLUDE_DIRECTORIES + ${SYCL_INCLUDE_DIR}) +set_property( + TARGET torch::sycl PROPERTY INTERFACE_LINK_LIBRARIES + ${SYCL_LIBRARY}) + +# xpurt +add_library(torch::xpurt INTERFACE IMPORTED) +set_property( + TARGET torch::xpurt PROPERTY INTERFACE_LINK_LIBRARIES + torch::sycl) + +# setting xpu arch flags +torch_xpu_get_arch_list(XPU_ARCH_FLAGS) +# propagate to torch-xpu-ops +set(TORCH_XPU_ARCH_LIST ${XPU_ARCH_FLAGS}) + +# Ensure SYCL device code compiles with C++20 (matching CMAKE_CXX_STANDARD). +# SYCL_FLAGS flows into SYCL_COMPILE_FLAGS in torch-xpu-ops' BuildFlags.cmake +# and is passed directly to icpx on the device compilation command line. +list(APPEND SYCL_FLAGS -std=c++20) + +# Ensure USE_XPU is enabled. +string(APPEND XPU_HOST_CXX_FLAGS " -DUSE_XPU") +string(APPEND XPU_HOST_CXX_FLAGS " -DSYCL_COMPILER_VERSION=${SYCL_COMPILER_VERSION}") + +if(DEFINED ENV{XPU_ENABLE_KINETO}) + set(XPU_ENABLE_KINETO TRUE) +else() + set(XPU_ENABLE_KINETO FALSE) +endif() + +if(WIN32) + if(${SYCL_COMPILER_VERSION} GREATER_EQUAL 20250101) + set(XPU_ENABLE_KINETO TRUE) + endif() +else() + set(XPU_ENABLE_KINETO TRUE) +endif() \ No newline at end of file diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Tensorpipe/TensorpipeTargets-release.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Tensorpipe/TensorpipeTargets-release.cmake new file mode 100644 index 0000000000000000000000000000000000000000..b59f8ceca10f56aaad16d71c32979919ea0537c1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Tensorpipe/TensorpipeTargets-release.cmake @@ -0,0 +1,39 @@ +#---------------------------------------------------------------- +# Generated CMake target import file for configuration "Release". +#---------------------------------------------------------------- + +# Commands may need to know the format version. +set(CMAKE_IMPORT_FILE_VERSION 1) + +# Import target "tensorpipe_uv" for configuration "Release" +set_property(TARGET tensorpipe_uv APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE) +set_target_properties(tensorpipe_uv PROPERTIES + IMPORTED_LINK_INTERFACE_LANGUAGES_RELEASE "C" + IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib64/libtensorpipe_uv.a" + ) + +list(APPEND _cmake_import_check_targets tensorpipe_uv ) +list(APPEND _cmake_import_check_files_for_tensorpipe_uv "${_IMPORT_PREFIX}/lib64/libtensorpipe_uv.a" ) + +# Import target "tensorpipe" for configuration "Release" +set_property(TARGET tensorpipe APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE) +set_target_properties(tensorpipe PROPERTIES + IMPORTED_LINK_INTERFACE_LANGUAGES_RELEASE "CXX" + IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib64/libtensorpipe.a" + ) + +list(APPEND _cmake_import_check_targets tensorpipe ) +list(APPEND _cmake_import_check_files_for_tensorpipe "${_IMPORT_PREFIX}/lib64/libtensorpipe.a" ) + +# Import target "tensorpipe_cuda" for configuration "Release" +set_property(TARGET tensorpipe_cuda APPEND PROPERTY IMPORTED_CONFIGURATIONS RELEASE) +set_target_properties(tensorpipe_cuda PROPERTIES + IMPORTED_LINK_INTERFACE_LANGUAGES_RELEASE "CXX" + IMPORTED_LOCATION_RELEASE "${_IMPORT_PREFIX}/lib64/libtensorpipe_cuda.a" + ) + +list(APPEND _cmake_import_check_targets tensorpipe_cuda ) +list(APPEND _cmake_import_check_files_for_tensorpipe_cuda "${_IMPORT_PREFIX}/lib64/libtensorpipe_cuda.a" ) + +# Commands beyond this point should not need to know the version. +set(CMAKE_IMPORT_FILE_VERSION) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Tensorpipe/TensorpipeTargets.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Tensorpipe/TensorpipeTargets.cmake new file mode 100644 index 0000000000000000000000000000000000000000..db80d5297ce171e898e7ca0ccc6043b577baed2b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Tensorpipe/TensorpipeTargets.cmake @@ -0,0 +1,122 @@ +# Generated by CMake + +if("${CMAKE_MAJOR_VERSION}.${CMAKE_MINOR_VERSION}" LESS 2.8) + message(FATAL_ERROR "CMake >= 2.8.12 required") +endif() +if(CMAKE_VERSION VERSION_LESS "2.8.12") + message(FATAL_ERROR "CMake >= 2.8.12 required") +endif() +cmake_policy(PUSH) +cmake_policy(VERSION 2.8.12...3.29) +#---------------------------------------------------------------- +# Generated CMake target import file. +#---------------------------------------------------------------- + +# Commands may need to know the format version. +set(CMAKE_IMPORT_FILE_VERSION 1) + +# Protect against multiple inclusion, which would fail when already imported targets are added once more. +set(_cmake_targets_defined "") +set(_cmake_targets_not_defined "") +set(_cmake_expected_targets "") +foreach(_cmake_expected_target IN ITEMS tensorpipe_uv tensorpipe tensorpipe_cuda) + list(APPEND _cmake_expected_targets "${_cmake_expected_target}") + if(TARGET "${_cmake_expected_target}") + list(APPEND _cmake_targets_defined "${_cmake_expected_target}") + else() + list(APPEND _cmake_targets_not_defined "${_cmake_expected_target}") + endif() +endforeach() +unset(_cmake_expected_target) +if(_cmake_targets_defined STREQUAL _cmake_expected_targets) + unset(_cmake_targets_defined) + unset(_cmake_targets_not_defined) + unset(_cmake_expected_targets) + unset(CMAKE_IMPORT_FILE_VERSION) + cmake_policy(POP) + return() +endif() +if(NOT _cmake_targets_defined STREQUAL "") + string(REPLACE ";" ", " _cmake_targets_defined_text "${_cmake_targets_defined}") + string(REPLACE ";" ", " _cmake_targets_not_defined_text "${_cmake_targets_not_defined}") + message(FATAL_ERROR "Some (but not all) targets in this export set were already defined.\nTargets Defined: ${_cmake_targets_defined_text}\nTargets not yet defined: ${_cmake_targets_not_defined_text}\n") +endif() +unset(_cmake_targets_defined) +unset(_cmake_targets_not_defined) +unset(_cmake_expected_targets) + + +# Compute the installation prefix relative to this file. +get_filename_component(_IMPORT_PREFIX "${CMAKE_CURRENT_LIST_FILE}" PATH) +get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH) +get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH) +get_filename_component(_IMPORT_PREFIX "${_IMPORT_PREFIX}" PATH) +if(_IMPORT_PREFIX STREQUAL "/") + set(_IMPORT_PREFIX "") +endif() + +# Create imported target tensorpipe_uv +add_library(tensorpipe_uv STATIC IMPORTED) + +set_target_properties(tensorpipe_uv PROPERTIES + INTERFACE_LINK_LIBRARIES "\$;\$;\$;\$" +) + +# Create imported target tensorpipe +add_library(tensorpipe STATIC IMPORTED) + +set_target_properties(tensorpipe PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${_IMPORT_PREFIX}/include" + INTERFACE_LINK_LIBRARIES "\$" +) + +# Create imported target tensorpipe_cuda +add_library(tensorpipe_cuda STATIC IMPORTED) + +set_target_properties(tensorpipe_cuda PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "/cvmfs/soft.computecanada.ca/easybuild/software/2023/x86-64-v3/Core/cudacore/13.2.0/include" + INTERFACE_LINK_LIBRARIES "tensorpipe;/cvmfs/soft.computecanada.ca/easybuild/software/2023/x86-64-v3/Core/cudacore/13.2.0/lib64/libcudart.so" +) + +# Load information for each installed configuration. +file(GLOB _cmake_config_files "${CMAKE_CURRENT_LIST_DIR}/TensorpipeTargets-*.cmake") +foreach(_cmake_config_file IN LISTS _cmake_config_files) + include("${_cmake_config_file}") +endforeach() +unset(_cmake_config_file) +unset(_cmake_config_files) + +# Cleanup temporary variables. +set(_IMPORT_PREFIX) + +# Loop over all imported files and verify that they actually exist +foreach(_cmake_target IN LISTS _cmake_import_check_targets) + if(CMAKE_VERSION VERSION_LESS "3.28" + OR NOT DEFINED _cmake_import_check_xcframework_for_${_cmake_target} + OR NOT IS_DIRECTORY "${_cmake_import_check_xcframework_for_${_cmake_target}}") + foreach(_cmake_file IN LISTS "_cmake_import_check_files_for_${_cmake_target}") + if(NOT EXISTS "${_cmake_file}") + message(FATAL_ERROR "The imported target \"${_cmake_target}\" references the file + \"${_cmake_file}\" +but this file does not exist. Possible reasons include: +* The file was deleted, renamed, or moved to another location. +* An install or uninstall procedure did not complete successfully. +* The installation package was faulty and contained + \"${CMAKE_CURRENT_LIST_FILE}\" +but not all the files it references. +") + endif() + endforeach() + endif() + unset(_cmake_file) + unset("_cmake_import_check_files_for_${_cmake_target}") +endforeach() +unset(_cmake_target) +unset(_cmake_import_check_targets) + +# This file does not depend on other imported targets which have +# been exported from the same project but in a separate export set. + +# Commands beyond this point should not need to know the version. +set(CMAKE_IMPORT_FILE_VERSION) +cmake_policy(POP) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfig.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfig.cmake new file mode 100644 index 0000000000000000000000000000000000000000..522ae095bb9199c21a4f91a19ba505cfa0dc2a4b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfig.cmake @@ -0,0 +1,164 @@ +# FindTorch +# ------- +# +# Finds the Torch library +# +# This will define the following variables: +# +# TORCH_FOUND -- True if the system has the Torch library +# TORCH_INCLUDE_DIRS -- The include directories for torch +# TORCH_LIBRARIES -- Libraries to link against +# TORCH_CXX_FLAGS -- Additional (required) compiler flags +# +# and the following imported targets: +# +# torch +macro(append_torchlib_if_found) + foreach (_arg ${ARGN}) + find_library(${_arg}_LIBRARY ${_arg} PATHS "${TORCH_INSTALL_PREFIX}/lib") + if(${_arg}_LIBRARY) + list(APPEND TORCH_LIBRARIES ${${_arg}_LIBRARY}) + else() + message(WARNING "library ${_arg} not found.") + endif() + endforeach() +endmacro() + +macro(append_wholearchive_lib_if_found) + foreach (_arg ${ARGN}) + find_library(${_arg}_LIBRARY ${_arg} PATHS "${TORCH_INSTALL_PREFIX}/lib") + if(${_arg}_LIBRARY) + if(APPLE) + list(APPEND TORCH_LIBRARIES "-Wl,-force_load,${${_arg}_LIBRARY}") + elseif(MSVC) + list(APPEND TORCH_LIBRARIES "-WHOLEARCHIVE:${${_arg}_LIBRARY}") + else() + # Linux + list(APPEND TORCH_LIBRARIES "-Wl,--whole-archive ${${_arg}_LIBRARY} -Wl,--no-whole-archive") + endif() + else() + message(WARNING "library ${_arg} not found.") + endif() + endforeach() +endmacro() + +include(FindPackageHandleStandardArgs) + +if(DEFINED ENV{TORCH_INSTALL_PREFIX}) + set(TORCH_INSTALL_PREFIX $ENV{TORCH_INSTALL_PREFIX}) +else() + # Assume we are in /share/cmake/Torch/TorchConfig.cmake + get_filename_component(CMAKE_CURRENT_LIST_DIR "${CMAKE_CURRENT_LIST_FILE}" PATH) + get_filename_component(TORCH_INSTALL_PREFIX "${CMAKE_CURRENT_LIST_DIR}/../../../" ABSOLUTE) +endif() + +# Include directories. +set(TORCH_INCLUDE_DIRS + ${TORCH_INSTALL_PREFIX}/include + ${TORCH_INSTALL_PREFIX}/include/torch/csrc/api/include) + +# Library dependencies. +if(ON) + find_package(Caffe2 REQUIRED PATHS ${CMAKE_CURRENT_LIST_DIR}/../Caffe2) + set(TORCH_LIBRARIES torch ${Caffe2_MAIN_LIBS}) + append_torchlib_if_found(c10) +else() + add_library(torch STATIC IMPORTED) # set imported_location at the bottom + #library need whole archive + append_wholearchive_lib_if_found(torch torch_cpu) + if(ON) + append_wholearchive_lib_if_found(torch_cuda c10_cuda) + endif() + if(OFF) + append_wholearchive_lib_if_found(torch_xpu c10_xpu) + endif() + + # We need manually add dependent libraries when they are not linked into the + # shared library. + # TODO: this list might be incomplete. + append_torchlib_if_found(c10) + + if(ON) + append_torchlib_if_found(nnpack) + endif() + + if(ON) + append_torchlib_if_found(pytorch_qnnpack) + endif() + + if(ON) + append_torchlib_if_found(XNNPACK) + append_torchlib_if_found(microkernels-prod) + endif() + + if(OFF) + append_torchlib_if_found(kleidiai) + endif() + + append_torchlib_if_found(caffe2_protos protobuf-lite protobuf protoc) + append_torchlib_if_found(onnx onnx_proto) + + append_torchlib_if_found(fmt) + append_torchlib_if_found(cpuinfo clog) + + append_torchlib_if_found(eigen_blas) + append_torchlib_if_found(pthreadpool) + + if(ON) + append_torchlib_if_found(fbgemm) + endif() + + if(ON) + append_torchlib_if_found(dnnl mkldnn) + endif() + + append_torchlib_if_found(sleef asmjit) +endif() + +if(ON) + append_torchlib_if_found(kineto) +endif() + +if(ON) + if(MSVC) + find_library(CAFFE2_NVRTC_LIBRARY caffe2_nvrtc PATHS "${TORCH_INSTALL_PREFIX}/lib") + list(APPEND TORCH_CUDA_LIBRARIES ${CAFFE2_NVRTC_LIBRARY}) + else() + set(TORCH_CUDA_LIBRARIES ${CUDA_NVRTC_LIB}) + endif() + if(TARGET torch::nvtoolsext) + list(APPEND TORCH_CUDA_LIBRARIES torch::nvtoolsext) + endif() + + if(ON) + find_library(C10_CUDA_LIBRARY c10_cuda PATHS "${TORCH_INSTALL_PREFIX}/lib") + list(APPEND TORCH_CUDA_LIBRARIES ${C10_CUDA_LIBRARY} ${Caffe2_PUBLIC_CUDA_DEPENDENCY_LIBS}) + endif() + list(APPEND TORCH_LIBRARIES ${TORCH_CUDA_LIBRARIES}) +endif() + +if(OFF AND ON) + append_torchlib_if_found(c10_xpu torch_xpu) +endif() + +find_library(TORCH_LIBRARY torch PATHS "${TORCH_INSTALL_PREFIX}/lib") +# the statements below changes target properties on +# - the imported target from Caffe2Targets.cmake in shared library mode (see the find_package above) +# - this is untested whether it is the correct (or desired) methodology in CMake +# - the imported target created in this file in static library mode +if(NOT ON) + # do not set this property on the shared library target, as it will cause confusion in some builds + # as the configuration specific property is set in the Caffe2Targets.cmake file + set_target_properties(torch PROPERTIES + IMPORTED_LOCATION "${TORCH_LIBRARY}" + ) +endif() +set_target_properties(torch PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${TORCH_INCLUDE_DIRS}" + CXX_STANDARD 20 +) +if(TORCH_CXX_FLAGS) + set_property(TARGET torch PROPERTY INTERFACE_COMPILE_OPTIONS "${TORCH_CXX_FLAGS}") +endif() + +find_package_handle_standard_args(Torch DEFAULT_MSG TORCH_LIBRARY TORCH_INCLUDE_DIRS) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfigVersion.cmake b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfigVersion.cmake new file mode 100644 index 0000000000000000000000000000000000000000..210c98e809cbecba676bda5cb45918392bbfb41a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/share/cmake/Torch/TorchConfigVersion.cmake @@ -0,0 +1,11 @@ +set(PACKAGE_VERSION "2.12.1") + +# Check whether the requested PACKAGE_FIND_VERSION is compatible +if("${PACKAGE_VERSION}" VERSION_LESS "${PACKAGE_FIND_VERSION}") + set(PACKAGE_VERSION_COMPATIBLE FALSE) +else() + set(PACKAGE_VERSION_COMPATIBLE TRUE) + if("${PACKAGE_VERSION}" VERSION_EQUAL "${PACKAGE_FIND_VERSION}") + set(PACKAGE_VERSION_EXACT TRUE) + endif() +endif() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/signal/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/signal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..74a23ba6fae924b8af8a2b4f0f798ca78b543545 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/signal/__init__.py @@ -0,0 +1,4 @@ +from . import windows + + +__all__ = ["windows"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/signal/windows/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/signal/windows/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f6749a92c6fc1525ea95c7d4d1e398229ab10b7a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/signal/windows/__init__.py @@ -0,0 +1,28 @@ +from .windows import ( + bartlett, + blackman, + cosine, + exponential, + gaussian, + general_cosine, + general_hamming, + hamming, + hann, + kaiser, + nuttall, +) + + +__all__ = [ + "bartlett", + "blackman", + "cosine", + "exponential", + "gaussian", + "general_cosine", + "general_hamming", + "hamming", + "hann", + "kaiser", + "nuttall", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/signal/windows/windows.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/signal/windows/windows.py new file mode 100644 index 0000000000000000000000000000000000000000..cda60aadfe1d6208354b045a86700e858cc946f0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/signal/windows/windows.py @@ -0,0 +1,883 @@ +# mypy: allow-untyped-defs +from collections.abc import Callable, Iterable +from math import sqrt +from typing import TypeVar + +import torch +from torch import Tensor +from torch._torch_docs import factory_common_args, merge_dicts, parse_kwargs + + +__all__ = [ + "bartlett", + "blackman", + "cosine", + "exponential", + "gaussian", + "general_cosine", + "general_hamming", + "hamming", + "hann", + "kaiser", + "nuttall", +] + +_T = TypeVar("_T") + +window_common_args = merge_dicts( + parse_kwargs( + """ + M (int): the length of the window. + In other words, the number of points of the returned window. + sym (bool, optional): If `False`, returns a periodic window suitable for use in spectral analysis. + If `True`, returns a symmetric window suitable for use in filter design. Default: `True`. +""" + ), + factory_common_args, + { + "normalization": "The window is normalized to 1 (maximum value is 1). However, the 1 doesn't appear if " + ":attr:`M` is even and :attr:`sym` is `True`.", + }, +) + + +def _add_docstr(*args: str) -> Callable[[_T], _T]: + r"""Adds docstrings to a given decorated function. + + Specially useful when then docstrings needs string interpolation, e.g., with + str.format(). + REMARK: Do not use this function if the docstring doesn't need string + interpolation, just write a conventional docstring. + + Args: + args (str): + """ + + def decorator(o: _T) -> _T: + o.__doc__ = "".join(args) + return o + + return decorator + + +def _window_function_checks( + function_name: str, M: int, dtype: torch.dtype, layout: torch.layout +) -> None: + r"""Performs common checks for all the defined windows. + This function should be called before computing any window. + + Args: + function_name (str): name of the window function. + M (int): length of the window. + dtype (:class:`torch.dtype`): the desired data type of returned tensor. + layout (:class:`torch.layout`): the desired layout of returned tensor. + """ + if M < 0: + raise ValueError( + f"{function_name} requires non-negative window length, got M={M}" + ) + if layout is not torch.strided: + raise ValueError( + f"{function_name} is implemented for strided tensors only, got: {layout}" + ) + if dtype not in [torch.float32, torch.float64]: + raise ValueError( + f"{function_name} expects float32 or float64 dtypes, got: {dtype}" + ) + + +@_add_docstr( + r""" +Computes a window with an exponential waveform. +Also known as Poisson window. + +The exponential window is defined as follows: + +.. math:: + w_n = \exp{\left(-\frac{|n - c|}{\tau}\right)} + +where `c` is the ``center`` of the window. + """, + r""" + +{normalization} + +Args: + {M} + +Keyword args: + center (float, optional): where the center of the window will be located. + Default: `M / 2` if `sym` is `False`, else `(M - 1) / 2`. + tau (float, optional): the decay value. + Tau is generally associated with a percentage, that means, that the value should + vary within the interval (0, 100]. If tau is 100, it is considered the uniform window. + Default: 1.0. + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric exponential window of size 10 and with a decay value of 1.0. + >>> # The center will be at (M - 1) / 2, where M is 10. + >>> torch.signal.windows.exponential(10) + tensor([0.0111, 0.0302, 0.0821, 0.2231, 0.6065, 0.6065, 0.2231, 0.0821, 0.0302, 0.0111]) + + >>> # Generates a periodic exponential window and decay factor equal to .5 + >>> torch.signal.windows.exponential(10, sym=False,tau=.5) + tensor([4.5400e-05, 3.3546e-04, 2.4788e-03, 1.8316e-02, 1.3534e-01, 1.0000e+00, 1.3534e-01, 1.8316e-02, 2.4788e-03, 3.3546e-04]) + """.format(**window_common_args), +) +def exponential( + M: int, + *, + center: float | None = None, + tau: float = 1.0, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + if dtype is None: + dtype = torch.get_default_dtype() + + _window_function_checks("exponential", M, dtype, layout) + + if tau <= 0: + raise ValueError(f"Tau must be positive, got: {tau} instead.") + + if sym and center is not None: + raise ValueError("Center must be None for symmetric windows") + + if M == 0: + return torch.empty( + (0,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad + ) + + if center is None: + center = (M if not sym and M > 1 else M - 1) / 2.0 + + constant = 1 / tau + + k = torch.linspace( + start=-center * constant, + end=(-center + (M - 1)) * constant, + steps=M, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + return torch.exp(-torch.abs(k)) + + +@_add_docstr( + r""" +Computes a window with a simple cosine waveform, following the same implementation as SciPy. +This window is also known as the sine window. + +The cosine window is defined as follows: + +.. math:: + w_n = \sin\left(\frac{\pi (n + 0.5)}{M}\right) + +This formula differs from the typical cosine window formula by incorporating a 0.5 term in the numerator, +which shifts the sample positions. This adjustment results in a window that starts and ends with non-zero values. + +""", + r""" + +{normalization} + +Args: + {M} + +Keyword args: + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric cosine window. + >>> torch.signal.windows.cosine(10) + tensor([0.1564, 0.4540, 0.7071, 0.8910, 0.9877, 0.9877, 0.8910, 0.7071, 0.4540, 0.1564]) + + >>> # Generates a periodic cosine window. + >>> torch.signal.windows.cosine(10, sym=False) + tensor([0.1423, 0.4154, 0.6549, 0.8413, 0.9595, 1.0000, 0.9595, 0.8413, 0.6549, 0.4154]) +""".format( + **window_common_args, + ), +) +def cosine( + M: int, + *, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + if dtype is None: + dtype = torch.get_default_dtype() + + _window_function_checks("cosine", M, dtype, layout) + + if M == 0: + return torch.empty( + (0,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad + ) + + start = 0.5 + constant = torch.pi / (M + 1 if not sym and M > 1 else M) + + k = torch.linspace( + start=start * constant, + end=(start + (M - 1)) * constant, + steps=M, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + return torch.sin(k) + + +@_add_docstr( + r""" +Computes a window with a gaussian waveform. + +The gaussian window is defined as follows: + +.. math:: + w_n = \exp{\left(-\left(\frac{n}{2\sigma}\right)^2\right)} + """, + r""" + +{normalization} + +Args: + {M} + +Keyword args: + std (float, optional): the standard deviation of the gaussian. It controls how narrow or wide the window is. + Default: 1.0. + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric gaussian window with a standard deviation of 1.0. + >>> torch.signal.windows.gaussian(10) + tensor([4.0065e-05, 2.1875e-03, 4.3937e-02, 3.2465e-01, 8.8250e-01, 8.8250e-01, 3.2465e-01, 4.3937e-02, 2.1875e-03, 4.0065e-05]) + + >>> # Generates a periodic gaussian window and standard deviation equal to 0.9. + >>> torch.signal.windows.gaussian(10, sym=False,std=0.9) + tensor([1.9858e-07, 5.1365e-05, 3.8659e-03, 8.4658e-02, 5.3941e-01, 1.0000e+00, 5.3941e-01, 8.4658e-02, 3.8659e-03, 5.1365e-05]) +""".format( + **window_common_args, + ), +) +def gaussian( + M: int, + *, + std: float = 1.0, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + if dtype is None: + dtype = torch.get_default_dtype() + + _window_function_checks("gaussian", M, dtype, layout) + + if std <= 0: + raise ValueError(f"Standard deviation must be positive, got: {std} instead.") + + if M == 0: + return torch.empty( + (0,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad + ) + + start = -(M if not sym and M > 1 else M - 1) / 2.0 + + constant = 1 / (std * sqrt(2)) + + k = torch.linspace( + start=start * constant, + end=(start + (M - 1)) * constant, + steps=M, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + return torch.exp(-(k**2)) # pyrefly: ignore [unsupported-operation] + + +@_add_docstr( + r""" +Computes the Kaiser window. + +The Kaiser window is defined as follows: + +.. math:: + w_n = I_0 \left( \beta \sqrt{1 - \left( {\frac{n - N/2}{N/2}} \right) ^2 } \right) / I_0( \beta ) + +where ``I_0`` is the zeroth order modified Bessel function of the first kind (see :func:`torch.special.i0`), and +``N = M - 1 if sym else M``. + """, + r""" + +{normalization} + +Args: + {M} + +Keyword args: + beta (float, optional): shape parameter for the window. Must be non-negative. Default: 12.0 + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric gaussian window with a standard deviation of 1.0. + >>> torch.signal.windows.kaiser(5) + tensor([4.0065e-05, 2.1875e-03, 4.3937e-02, 3.2465e-01, 8.8250e-01, 8.8250e-01, 3.2465e-01, 4.3937e-02, 2.1875e-03, 4.0065e-05]) + >>> # Generates a periodic gaussian window and standard deviation equal to 0.9. + >>> torch.signal.windows.kaiser(5, sym=False,std=0.9) + tensor([1.9858e-07, 5.1365e-05, 3.8659e-03, 8.4658e-02, 5.3941e-01, 1.0000e+00, 5.3941e-01, 8.4658e-02, 3.8659e-03, 5.1365e-05]) +""".format( + **window_common_args, + ), +) +def kaiser( + M: int, + *, + beta: float = 12.0, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + if dtype is None: + dtype = torch.get_default_dtype() + + _window_function_checks("kaiser", M, dtype, layout) + + if beta < 0: + raise ValueError(f"beta must be non-negative, got: {beta} instead.") + + if M == 0: + return torch.empty( + (0,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad + ) + + if M == 1: + return torch.ones( + (1,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad + ) + + # Avoid NaNs by casting `beta` to the appropriate dtype. + # pyrefly: ignore [bad-assignment] + beta = torch.tensor(beta, dtype=dtype, device=device) + + start = -beta + constant = 2.0 * beta / (M if not sym else M - 1) + end = torch.minimum( + # pyrefly: ignore [bad-argument-type] + beta, + # pyrefly: ignore [bad-argument-type] + start + (M - 1) * constant, + ) + + k = torch.linspace( + start=start, + end=end, + steps=M, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + return torch.i0(torch.sqrt(beta * beta - torch.pow(k, 2))) / torch.i0( + # pyrefly: ignore [bad-argument-type] + beta + ) + + +@_add_docstr( + r""" +Computes the Hamming window. + +The Hamming window is defined as follows: + +.. math:: + w_n = \alpha - \beta\ \cos \left( \frac{2 \pi n}{M - 1} \right) + """, + r""" + +{normalization} + +Arguments: + {M} + +Keyword args: + {sym} + alpha (float, optional): The coefficient :math:`\alpha` in the equation above. + beta (float, optional): The coefficient :math:`\beta` in the equation above. + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric Hamming window. + >>> torch.signal.windows.hamming(10) + tensor([0.0800, 0.1876, 0.4601, 0.7700, 0.9723, 0.9723, 0.7700, 0.4601, 0.1876, 0.0800]) + + >>> # Generates a periodic Hamming window. + >>> torch.signal.windows.hamming(10, sym=False) + tensor([0.0800, 0.1679, 0.3979, 0.6821, 0.9121, 1.0000, 0.9121, 0.6821, 0.3979, 0.1679]) +""".format(**window_common_args), +) +def hamming( + M: int, + *, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + return general_hamming( + M, + sym=sym, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + +@_add_docstr( + r""" +Computes the Hann window. + +The Hann window is defined as follows: + +.. math:: + w_n = \frac{1}{2}\ \left[1 - \cos \left( \frac{2 \pi n}{M - 1} \right)\right] = + \sin^2 \left( \frac{\pi n}{M - 1} \right) + """, + r""" + +{normalization} + +Arguments: + {M} + +Keyword args: + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric Hann window. + >>> torch.signal.windows.hann(10) + tensor([0.0000, 0.1170, 0.4132, 0.7500, 0.9698, 0.9698, 0.7500, 0.4132, 0.1170, 0.0000]) + + >>> # Generates a periodic Hann window. + >>> torch.signal.windows.hann(10, sym=False) + tensor([0.0000, 0.0955, 0.3455, 0.6545, 0.9045, 1.0000, 0.9045, 0.6545, 0.3455, 0.0955]) +""".format(**window_common_args), +) +def hann( + M: int, + *, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + return general_hamming( + M, + alpha=0.5, + sym=sym, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + +@_add_docstr( + r""" +Computes the Blackman window. + +The Blackman window is defined as follows: + +.. math:: + w_n = 0.42 - 0.5 \cos \left( \frac{2 \pi n}{M - 1} \right) + 0.08 \cos \left( \frac{4 \pi n}{M - 1} \right) + """, + r""" + +{normalization} + +Arguments: + {M} + +Keyword args: + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric Blackman window. + >>> torch.signal.windows.blackman(5) + tensor([-1.4901e-08, 3.4000e-01, 1.0000e+00, 3.4000e-01, -1.4901e-08]) + + >>> # Generates a periodic Blackman window. + >>> torch.signal.windows.blackman(5, sym=False) + tensor([-1.4901e-08, 2.0077e-01, 8.4923e-01, 8.4923e-01, 2.0077e-01]) +""".format(**window_common_args), +) +def blackman( + M: int, + *, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + if dtype is None: + dtype = torch.get_default_dtype() + + _window_function_checks("blackman", M, dtype, layout) + + return general_cosine( + M, + a=[0.42, 0.5, 0.08], + sym=sym, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + +@_add_docstr( + r""" +Computes the Bartlett window. + +The Bartlett window is defined as follows: + +.. math:: + w_n = 1 - \left| \frac{2n}{M - 1} - 1 \right| = \begin{cases} + \frac{2n}{M - 1} & \text{if } 0 \leq n \leq \frac{M - 1}{2} \\ + 2 - \frac{2n}{M - 1} & \text{if } \frac{M - 1}{2} < n < M \\ \end{cases} + """, + r""" + +{normalization} + +Arguments: + {M} + +Keyword args: + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric Bartlett window. + >>> torch.signal.windows.bartlett(10) + tensor([0.0000, 0.2222, 0.4444, 0.6667, 0.8889, 0.8889, 0.6667, 0.4444, 0.2222, 0.0000]) + + >>> # Generates a periodic Bartlett window. + >>> torch.signal.windows.bartlett(10, sym=False) + tensor([0.0000, 0.2000, 0.4000, 0.6000, 0.8000, 1.0000, 0.8000, 0.6000, 0.4000, 0.2000]) +""".format(**window_common_args), +) +def bartlett( + M: int, + *, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + if dtype is None: + dtype = torch.get_default_dtype() + + _window_function_checks("bartlett", M, dtype, layout) + + if M == 0: + return torch.empty( + (0,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad + ) + + if M == 1: + return torch.ones( + (1,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad + ) + + start = -1 + constant = 2 / (M if not sym else M - 1) + + k = torch.linspace( + start=start, + end=start + (M - 1) * constant, + steps=M, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + return 1 - torch.abs(k) + + +@_add_docstr( + r""" +Computes the general cosine window. + +The general cosine window is defined as follows: + +.. math:: + w_n = \sum^{M-1}_{i=0} (-1)^i a_i \cos{ \left( \frac{2 \pi i n}{M - 1}\right)} + """, + r""" + +{normalization} + +Arguments: + {M} + +Keyword args: + a (Iterable): the coefficients associated to each of the cosine functions. + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric general cosine window with 3 coefficients. + >>> torch.signal.windows.general_cosine(10, a=[0.46, 0.23, 0.31], sym=True) + tensor([0.5400, 0.3376, 0.1288, 0.4200, 0.9136, 0.9136, 0.4200, 0.1288, 0.3376, 0.5400]) + + >>> # Generates a periodic general cosine window with 2 coefficients. + >>> torch.signal.windows.general_cosine(10, a=[0.5, 1 - 0.5], sym=False) + tensor([0.0000, 0.0955, 0.3455, 0.6545, 0.9045, 1.0000, 0.9045, 0.6545, 0.3455, 0.0955]) +""".format(**window_common_args), +) +def general_cosine( + M, + *, + a: Iterable, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + if dtype is None: + dtype = torch.get_default_dtype() + + _window_function_checks("general_cosine", M, dtype, layout) + + if M == 0: + return torch.empty( + (0,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad + ) + + if M == 1: + return torch.ones( + (1,), dtype=dtype, layout=layout, device=device, requires_grad=requires_grad + ) + + if not isinstance(a, Iterable): + raise TypeError("Coefficients must be a list/tuple") + + if not a: + raise ValueError("Coefficients cannot be empty") + + constant = 2 * torch.pi / (M if not sym else M - 1) + + k = torch.linspace( + start=0, + end=(M - 1) * constant, + steps=M, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + a_i = torch.tensor( + [(-1) ** i * w for i, w in enumerate(a)], + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + i = torch.arange( + a_i.shape[0], + dtype=a_i.dtype, + device=a_i.device, + requires_grad=a_i.requires_grad, + ) + return (a_i.unsqueeze(-1) * torch.cos(i.unsqueeze(-1) * k)).sum(0) + + +@_add_docstr( + r""" +Computes the general Hamming window. + +The general Hamming window is defined as follows: + +.. math:: + w_n = \alpha - (1 - \alpha) \cos{ \left( \frac{2 \pi n}{M-1} \right)} + """, + r""" + +{normalization} + +Arguments: + {M} + +Keyword args: + alpha (float, optional): the window coefficient. Default: 0.54. + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +Examples:: + + >>> # Generates a symmetric Hamming window with the general Hamming window. + >>> torch.signal.windows.general_hamming(10, sym=True) + tensor([0.0800, 0.1876, 0.4601, 0.7700, 0.9723, 0.9723, 0.7700, 0.4601, 0.1876, 0.0800]) + + >>> # Generates a periodic Hann window with the general Hamming window. + >>> torch.signal.windows.general_hamming(10, alpha=0.5, sym=False) + tensor([0.0000, 0.0955, 0.3455, 0.6545, 0.9045, 1.0000, 0.9045, 0.6545, 0.3455, 0.0955]) +""".format(**window_common_args), +) +def general_hamming( + M, + *, + alpha: float = 0.54, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + return general_cosine( + M, + a=[alpha, 1.0 - alpha], + sym=sym, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) + + +@_add_docstr( + r""" +Computes the minimum 4-term Blackman-Harris window according to Nuttall. + +.. math:: + w_n = 1 - 0.36358 \cos{(z_n)} + 0.48917 \cos{(2z_n)} - 0.13659 \cos{(3z_n)} + 0.01064 \cos{(4z_n)} + +where :math:`z_n = \frac{2 \pi n}{M}`. + """, + """ + +{normalization} + +Arguments: + {M} + +Keyword args: + {sym} + {dtype} + {layout} + {device} + {requires_grad} + +References:: + + - A. Nuttall, "Some windows with very good sidelobe behavior," + IEEE Transactions on Acoustics, Speech, and Signal Processing, vol. 29, no. 1, pp. 84-91, + Feb 1981. https://doi.org/10.1109/TASSP.1981.1163506 + + - Heinzel G. et al., "Spectrum and spectral density estimation by the Discrete Fourier transform (DFT), + including a comprehensive list of window functions and some new flat-top windows", + February 15, 2002 https://holometer.fnal.gov/GH_FFT.pdf + +Examples:: + + >>> # Generates a symmetric Nutall window. + >>> torch.signal.windows.general_hamming(5, sym=True) + tensor([3.6280e-04, 2.2698e-01, 1.0000e+00, 2.2698e-01, 3.6280e-04]) + + >>> # Generates a periodic Nuttall window. + >>> torch.signal.windows.general_hamming(5, sym=False) + tensor([3.6280e-04, 1.1052e-01, 7.9826e-01, 7.9826e-01, 1.1052e-01]) +""".format(**window_common_args), +) +def nuttall( + M: int, + *, + sym: bool = True, + dtype: torch.dtype | None = None, + layout: torch.layout = torch.strided, + device: torch.device | None = None, + requires_grad: bool = False, +) -> Tensor: + return general_cosine( + M, + a=[0.3635819, 0.4891775, 0.1365995, 0.0106411], + sym=sym, + dtype=dtype, + layout=layout, + device=device, + requires_grad=requires_grad, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..78846265e9669fcea5c65dbdb208e3a119a4f931 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/__init__.py @@ -0,0 +1,737 @@ +# mypy: allow-untyped-defs +# The Tensor classes are added to this module by python_tensor.cpp +# A workaround to support both TorchScript and MyPy: +from typing import Any, TYPE_CHECKING + +import torch +from torch import Tensor +from torch._C import _add_docstr, _sparse # type: ignore[attr-defined] + +# Semi structured sparsity support +from .semi_structured import ( + SparseSemiStructuredTensor, + SparseSemiStructuredTensorCUSPARSELT, + SparseSemiStructuredTensorCUTLASS, + to_sparse_semi_structured, +) + + +if TYPE_CHECKING: + from torch.types import _dtype as DType + + DimOrDims = int | tuple[int, ...] | list[int] | None +else: + # The JIT doesn't understand Union, nor torch.dtype here + DType = int + DimOrDims = tuple[int] | None + + +__all__ = [ + "addmm", + "check_sparse_tensor_invariants", + "mm", + "sum", + "softmax", + # pyrefly: ignore [bad-dunder-all] + "solve", + "log_softmax", + "SparseSemiStructuredTensor", + "SparseSemiStructuredTensorCUTLASS", + "SparseSemiStructuredTensorCUSPARSELT", + "to_sparse_semi_structured", + "as_sparse_gradcheck", +] + +addmm = _add_docstr( + _sparse._sparse_addmm, + r""" +sparse.addmm(mat, mat1, mat2, *, beta=1., alpha=1.) -> Tensor + +This function does exact same thing as :func:`torch.addmm` in the forward, +except that it supports backward for sparse COO and CSR matrix :attr:`mat1`. +When :attr:`mat1` is a COO tensor it must have `sparse_dim = 2`. + +Supports both CSR and COO storage formats. + +.. note:: + **Gradient support:** + + - **COO @ Dense**: Backward is supported for both inputs. The gradient for the + sparse input is returned as a sparse COO tensor. + - **CSR @ Dense**: Backward is supported for both inputs. The gradient for the + sparse input is returned as a sparse CSR tensor. + - **CSC/BSR/BSC @ Dense**: Not supported. + - **Sparse @ Sparse** (COO @ COO, CSR @ CSR): Forward works, but backward is + not supported. + +Args: + mat (Tensor): a dense matrix to be added + mat1 (Tensor): a sparse matrix to be multiplied + mat2 (Tensor): a dense matrix to be multiplied + beta (Number, optional): multiplier for :attr:`mat` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) +""", +) + + +mm = _add_docstr( + _sparse._sparse_mm, + r""" + Performs a matrix multiplication of the sparse matrix :attr:`mat1` + and the (sparse or strided) matrix :attr:`mat2`. Similar to :func:`torch.mm`, if :attr:`mat1` is a + :math:`(n \times m)` tensor, :attr:`mat2` is a :math:`(m \times p)` tensor, out will be a + :math:`(n \times p)` tensor. + When :attr:`mat1` is a COO tensor it must have `sparse_dim = 2`. + + Supports both CSR and COO storage formats. + +.. note:: + **Gradient support:** + + - **COO @ Dense**: Backward is supported for both inputs. The gradient for the + sparse input is returned as a sparse COO tensor. + - **CSR @ Dense**: Backward is supported for both inputs. The gradient for the + sparse input is returned as a sparse CSR tensor. + - **CSC/BSR/BSC @ Dense**: Not supported. + - **Sparse @ Sparse** (COO @ COO, CSR @ CSR): Forward works, but backward is + not supported. + - **Mixed formats** (COO @ CSR, CSR @ COO): Not supported. + + This function also additionally accepts an optional :attr:`reduce` argument that allows + specification of an optional reduction operation, mathematically performs the following operation: + +.. math:: + + z_{ij} = \bigoplus_{k = 0}^{K - 1} x_{ik} y_{kj} + +where :math:`\bigoplus` defines the reduce operator. :attr:`reduce` is implemented only for +CSR storage format on CPU device. + +Args: + mat1 (Tensor): the first sparse matrix to be multiplied + mat2 (Tensor): the second matrix to be multiplied, which could be sparse or dense + reduce (str, optional): the reduction operation to apply for non-unique indices + (:obj:`"sum"`, :obj:`"mean"`, :obj:`"amax"`, :obj:`"amin"`). Default :obj:`"sum"`. + +Shape: + The format of the output tensor of this function follows: + - sparse x sparse -> sparse + - sparse x dense -> dense + +Example:: + + >>> a = torch.tensor([[1., 0, 2], [0, 3, 0]]).to_sparse().requires_grad_() + >>> a + tensor(indices=tensor([[0, 0, 1], + [0, 2, 1]]), + values=tensor([1., 2., 3.]), + size=(2, 3), nnz=3, layout=torch.sparse_coo, requires_grad=True) + >>> b = torch.tensor([[0, 1.], [2, 0], [0, 0]], requires_grad=True) + >>> b + tensor([[0., 1.], + [2., 0.], + [0., 0.]], requires_grad=True) + >>> y = torch.sparse.mm(a, b) + >>> y + tensor([[0., 1.], + [6., 0.]], grad_fn=) + >>> y.sum().backward() + >>> a.grad + tensor(indices=tensor([[0, 0, 1], + [0, 2, 1]]), + values=tensor([1., 0., 2.]), + size=(2, 3), nnz=3, layout=torch.sparse_coo) + >>> c = a.detach().to_sparse_csr() + >>> c + tensor(crow_indices=tensor([0, 2, 3]), + col_indices=tensor([0, 2, 1]), + values=tensor([1., 2., 3.]), size=(2, 3), nnz=3, + layout=torch.sparse_csr) + >>> y1 = torch.sparse.mm(c, b, 'sum') + >>> y1 + tensor([[0., 1.], + [6., 0.]], grad_fn=) + >>> y2 = torch.sparse.mm(c, b, 'max') + >>> y2 + tensor([[0., 1.], + [6., 0.]], grad_fn=) +""", +) + + +sampled_addmm = _add_docstr( + _sparse.sparse_sampled_addmm, + r""" +sparse.sampled_addmm(input, mat1, mat2, *, beta=1., alpha=1., out=None) -> Tensor + +Performs a matrix multiplication of the dense matrices :attr:`mat1` and :attr:`mat2` at the locations +specified by the sparsity pattern of :attr:`input`. The matrix :attr:`input` is added to the final result. + +Mathematically this performs the following operation: + +.. math:: + + \text{out} = \alpha\ (\text{mat1} \mathbin{@} \text{mat2})*\text{spy}(\text{input}) + \beta\ \text{input} + +where :math:`\text{spy}(\text{input})` is the sparsity pattern matrix of :attr:`input`, :attr:`alpha` +and :attr:`beta` are the scaling factors. +:math:`\text{spy}(\text{input})` has value 1 at the positions where :attr:`input` has non-zero values, and 0 elsewhere. + +.. note:: + :attr:`input` must be a sparse CSR tensor. :attr:`mat1` and :attr:`mat2` must be dense tensors. + +Args: + input (Tensor): a sparse CSR matrix of shape `(m, n)` to be added and used to compute + the sampled matrix multiplication + mat1 (Tensor): a dense matrix of shape `(m, k)` to be multiplied + mat2 (Tensor): a dense matrix of shape `(k, n)` to be multiplied + +Keyword args: + beta (Number, optional): multiplier for :attr:`input` (:math:`\beta`) + alpha (Number, optional): multiplier for :math:`mat1 @ mat2` (:math:`\alpha`) + out (Tensor, optional): output tensor. Ignored if `None`. Default: `None`. + +Examples:: + + >>> input = torch.eye(3, device='cuda').to_sparse_csr() + >>> mat1 = torch.randn(3, 5, device='cuda') + >>> mat2 = torch.randn(5, 3, device='cuda') + >>> torch.sparse.sampled_addmm(input, mat1, mat2) + tensor(crow_indices=tensor([0, 1, 2, 3]), + col_indices=tensor([0, 1, 2]), + values=tensor([ 0.2847, -0.7805, -0.1900]), device='cuda:0', + size=(3, 3), nnz=3, layout=torch.sparse_csr) + >>> torch.sparse.sampled_addmm(input, mat1, mat2).to_dense() + tensor([[ 0.2847, 0.0000, 0.0000], + [ 0.0000, -0.7805, 0.0000], + [ 0.0000, 0.0000, -0.1900]], device='cuda:0') + >>> torch.sparse.sampled_addmm(input, mat1, mat2, beta=0.5, alpha=0.5) + tensor(crow_indices=tensor([0, 1, 2, 3]), + col_indices=tensor([0, 1, 2]), + values=tensor([ 0.1423, -0.3903, -0.0950]), device='cuda:0', + size=(3, 3), nnz=3, layout=torch.sparse_csr) +""", +) + + +def sum(input: Tensor, dim: DimOrDims = None, dtype: DType | None = None) -> Tensor: + r"""Return the sum of each row of the given sparse tensor. + + Returns the sum of each row of the sparse tensor :attr:`input` in the given + dimensions :attr:`dim`. If :attr:`dim` is a list of dimensions, + reduce over all of them. When sum over all ``sparse_dim``, this method + returns a dense tensor instead of a sparse tensor. + + All summed :attr:`dim` are squeezed (see :func:`torch.squeeze`), resulting an output + tensor having :attr:`dim` fewer dimensions than :attr:`input`. + + During backward, only gradients at ``nnz`` locations of :attr:`input` + will propagate back. Note that the gradients of :attr:`input` is coalesced. + + Args: + input (Tensor): the input sparse tensor + dim (int or tuple of ints): a dimension or a list of dimensions to reduce. Default: reduce + over all dims. + dtype (:class:`torch.dtype`, optional): the desired data type of returned Tensor. + Default: dtype of :attr:`input`. + + Example:: + + >>> nnz = 3 + >>> dims = [5, 5, 2, 3] + >>> I = torch.cat([torch.randint(0, dims[0], size=(nnz,)), + torch.randint(0, dims[1], size=(nnz,))], 0).reshape(2, nnz) + >>> V = torch.randn(nnz, dims[2], dims[3]) + >>> size = torch.Size(dims) + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> S = torch.sparse_coo_tensor(I, V, size) + >>> S + tensor(indices=tensor([[2, 0, 3], + [2, 4, 1]]), + values=tensor([[[-0.6438, -1.6467, 1.4004], + [ 0.3411, 0.0918, -0.2312]], + + [[ 0.5348, 0.0634, -2.0494], + [-0.7125, -1.0646, 2.1844]], + + [[ 0.1276, 0.1874, -0.6334], + [-1.9682, -0.5340, 0.7483]]]), + size=(5, 5, 2, 3), nnz=3, layout=torch.sparse_coo) + + # when sum over only part of sparse_dims, return a sparse tensor + >>> torch.sparse.sum(S, [1, 3]) + tensor(indices=tensor([[0, 2, 3]]), + values=tensor([[-1.4512, 0.4073], + [-0.8901, 0.2017], + [-0.3183, -1.7539]]), + size=(5, 2), nnz=3, layout=torch.sparse_coo) + + # when sum over all sparse dim, return a dense tensor + # with summed dims squeezed + >>> torch.sparse.sum(S, [0, 1, 3]) + tensor([-2.6596, -1.1450]) + """ + if dtype is None: + if dim is not None: + return torch._sparse_sum(input, dim) + else: + return torch._sparse_sum(input) + else: + if dim is not None: + return torch._sparse_sum(input, dim, dtype=dtype) + else: + return torch._sparse_sum(input, dtype=dtype) + + +softmax = _add_docstr( + _sparse._sparse_softmax, + r""" +sparse.softmax(input, dim, *, dtype=None) -> Tensor + +Applies a softmax function. + +Softmax is defined as: + +:math:`\text{Softmax}(x_{i}) = \frac{exp(x_i)}{\sum_j exp(x_j)}` + +where :math:`i, j` run over sparse tensor indices and unspecified +entries are ignores. This is equivalent to defining unspecified +entries as negative infinity so that :math:`exp(x_k) = 0` when the +entry with index :math:`k` has not specified. + +It is applied to all slices along `dim`, and will re-scale them so +that the elements lie in the range `[0, 1]` and sum to 1. + +Args: + input (Tensor): input + dim (int): A dimension along which softmax will be computed. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. This is useful for preventing data type + overflows. Default: None +""", +) + + +spsolve = _add_docstr( + _sparse._spsolve, + r""" +sparse.spsolve(input, other, *, left=True) -> Tensor + +Computes the solution of a square system of linear equations with +a unique solution. Its purpose is similar to :func:`torch.linalg.solve`, +except that the system is defined by a sparse CSR matrix with layout +`sparse_csr`. + +Args: + input (Tensor): a sparse CSR matrix of shape `(n, n)` representing the + coefficients of the linear system. + other (Tensor): a dense matrix of shape `(n, )` representing the right-hand + side of the linear system. + left (bool, optional): whether to solve the system for `input @ out = other` + (default) or `out @ input = other`. Only `left=True` is supported. +""", +) + +log_softmax = _add_docstr( + _sparse._sparse_log_softmax, + r""" +sparse.log_softmax(input, dim, *, dtype=None) -> Tensor + +Applies a softmax function followed by logarithm. + +See :class:`~torch.sparse.softmax` for more details. + +Args: + input (Tensor): input + dim (int): A dimension along which softmax will be computed. + dtype (:class:`torch.dtype`, optional): the desired data type + of returned tensor. If specified, the input tensor is + casted to :attr:`dtype` before the operation is + performed. This is useful for preventing data type + overflows. Default: None +""", +) + + +spdiags = _add_docstr( + _sparse._spdiags, + r""" +sparse.spdiags(diagonals, offsets, shape, layout=None) -> Tensor + +Creates a sparse 2D tensor by placing the values from rows of +:attr:`diagonals` along specified diagonals of the output + +The :attr:`offsets` tensor controls which diagonals are set. + +- If :attr:`offsets[i]` = 0, it is the main diagonal +- If :attr:`offsets[i]` < 0, it is below the main diagonal +- If :attr:`offsets[i]` > 0, it is above the main diagonal + +The number of rows in :attr:`diagonals` must match the length of :attr:`offsets`, +and an offset may not be repeated. + +Args: + diagonals (Tensor): Matrix storing diagonals row-wise + offsets (Tensor): The diagonals to be set, stored as a vector + shape (2-tuple of ints): The desired shape of the result +Keyword args: + layout (:class:`torch.layout`, optional): The desired layout of the + returned tensor. ``torch.sparse_coo``, ``torch.sparse_csc`` and ``torch.sparse_csr`` + are supported. Default: ``torch.sparse_coo`` + +Examples: + +Set the main and first two lower diagonals of a matrix:: + + >>> diags = torch.arange(9).reshape(3, 3) + >>> diags + tensor([[0, 1, 2], + [3, 4, 5], + [6, 7, 8]]) + >>> s = torch.sparse.spdiags(diags, torch.tensor([0, -1, -2]), (3, 3)) + >>> s + tensor(indices=tensor([[0, 1, 2, 1, 2, 2], + [0, 1, 2, 0, 1, 0]]), + values=tensor([0, 1, 2, 3, 4, 6]), + size=(3, 3), nnz=6, layout=torch.sparse_coo) + >>> s.to_dense() + tensor([[0, 0, 0], + [3, 1, 0], + [6, 4, 2]]) + + +Change the output layout:: + + >>> diags = torch.arange(9).reshape(3, 3) + >>> diags + tensor([[0, 1, 2],[3, 4, 5], [6, 7, 8]) + >>> s = torch.sparse.spdiags(diags, torch.tensor([0, -1, -2]), (3, 3), layout=torch.sparse_csr) + >>> s + tensor(crow_indices=tensor([0, 1, 3, 6]), + col_indices=tensor([0, 0, 1, 0, 1, 2]), + values=tensor([0, 3, 1, 6, 4, 2]), size=(3, 3), nnz=6, + layout=torch.sparse_csr) + >>> s.to_dense() + tensor([[0, 0, 0], + [3, 1, 0], + [6, 4, 2]]) + +Set partial diagonals of a large output:: + + >>> diags = torch.tensor([[1, 2], [3, 4]]) + >>> offsets = torch.tensor([0, -1]) + >>> torch.sparse.spdiags(diags, offsets, (5, 5)).to_dense() + tensor([[1, 0, 0, 0, 0], + [3, 2, 0, 0, 0], + [0, 4, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0]]) + +.. note:: + + When setting the values along a given diagonal the index into the diagonal + and the index into the row of :attr:`diagonals` is taken as the + column index in the output. This has the effect that when setting a diagonal + with a positive offset `k` the first value along that diagonal will be + the value in position `k` of the row of :attr:`diagonals` + +Specifying a positive offset:: + + >>> diags = torch.tensor([[1, 2, 3], [1, 2, 3], [1, 2, 3]]) + >>> torch.sparse.spdiags(diags, torch.tensor([0, 1, 2]), (5, 5)).to_dense() + tensor([[1, 2, 3, 0, 0], + [0, 2, 3, 0, 0], + [0, 0, 3, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0]]) +""", +) + + +class check_sparse_tensor_invariants: + """A tool to control checking sparse tensor invariants. + + The following options exists to manage sparsr tensor invariants + checking in sparse tensor construction: + + 1. Using a context manager: + + .. code:: python + + with torch.sparse.check_sparse_tensor_invariants(): + run_my_model() + + 2. Using a procedural approach: + + .. code:: python + + prev_checks_enabled = torch.sparse.check_sparse_tensor_invariants.is_enabled() + torch.sparse.check_sparse_tensor_invariants.enable() + + run_my_model() + + if not prev_checks_enabled: + torch.sparse.check_sparse_tensor_invariants.disable() + + 3. Using function decoration: + + .. code:: python + + @torch.sparse.check_sparse_tensor_invariants() + def run_my_model(): + ... + + run_my_model() + + 4. Using ``check_invariants`` keyword argument in sparse tensor constructor call. + For example: + + >>> torch.sparse_csr_tensor([0, 1, 3], [0, 1], [1, 2], check_invariants=True) + Traceback (most recent call last): + File "", line 1, in + RuntimeError: `crow_indices[..., -1] == nnz` is not satisfied. + """ + + @staticmethod + def is_enabled(): + r"""Return True if the sparse tensor invariants checking is enabled. + + .. note:: + + Use :func:`torch.sparse.check_sparse_tensor_invariants.enable` or + :func:`torch.sparse.check_sparse_tensor_invariants.disable` to + manage the state of the sparse tensor invariants checks. + """ + return torch._C._check_sparse_tensor_invariants() + + @staticmethod + def enable(): + r"""Enable sparse tensor invariants checking in sparse tensor constructors. + + .. note:: + + By default, the sparse tensor invariants checks are disabled. Use + :func:`torch.sparse.check_sparse_tensor_invariants.is_enabled` to + retrieve the current state of sparse tensor invariants checking. + + .. note:: + + The sparse tensor invariants check flag is effective to all sparse + tensor constructors, both in Python and ATen. + + The flag can be locally overridden by the ``check_invariants`` + optional argument of the sparse tensor constructor functions. + """ + torch._C._set_check_sparse_tensor_invariants(True) + + @staticmethod + def disable(): + r"""Disable sparse tensor invariants checking in sparse tensor constructors. + + See :func:`torch.sparse.check_sparse_tensor_invariants.enable` for more information. + """ + torch._C._set_check_sparse_tensor_invariants(False) + + # context manager support + def __init__(self, enable=True): + self.state = enable + self.saved_state: bool | None = None + + def __enter__(self): + if self.saved_state is not None: + raise RuntimeError( + "This context manager instance is already activated." + " Use a different context manager instance for context nesting." + ) + self.saved_state = self.is_enabled() + torch._C._set_check_sparse_tensor_invariants(self.state) + + def __exit__(self, type, value, traceback): + if self.saved_state is None: + raise AssertionError("saved_state should not be None on exit") + torch._C._set_check_sparse_tensor_invariants(self.saved_state) + self.saved_state = None + + # decorator support + def __call__(self, mth): + def test_mth(*args, **kwargs): + with type(self)(self.state): + return mth(*args, **kwargs) + + return test_mth + + +def as_sparse_gradcheck(gradcheck): + """Decorate function, to extend gradcheck for sparse tensors. + + Decorator for torch.autograd.gradcheck or its functools.partial + variants that extends the gradcheck function with support to input + functions that operate on or/and return sparse tensors. + + The specified gradcheck function itself is guaranteed to operate + on strided tensors only. + + For example: + + >>> gradcheck = torch.sparse.as_sparse_gradcheck(torch.autograd.gradcheck) + >>> x = ( + ... torch.tensor([[0, 1], [2, 3]], dtype=torch.float64) + ... .to_sparse_coo() + ... .requires_grad_(True) + ... ) + >>> gradcheck(lambda x: x.to_sparse_csr(), x) + True + """ + + def gradcheck_with_sparse_support(func, inputs, **kwargs): + """ + Create gradcheck with support for sparse tensors. + + Same as :func:`torch.autograd.gradcheck` but with sparse tensors inputs and outputs support. + """ + masked = kwargs.pop("masked", False) + sparse_layouts = { + torch.sparse_coo, + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + } + sparse_compressed_layouts = { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + } + sparse_block_layouts = {torch.sparse_bsr, torch.sparse_bsc} + STRIDED_REPRESENTATION = "__STRIDED_REPRESENTATION__" + + def convert_to_strided_representation(args): + """Convert differentiable non-strided tensors to a representation containing differentiable strided tensors.""" + if not isinstance(args, (list, tuple)): + args = (args,) + new_args: list[Any] = [] + for obj in args: + if ( + isinstance(obj, torch.Tensor) + and obj.requires_grad + and obj.layout in sparse_layouts + ): + d = { + "layout": obj.layout, + "shape": obj.shape, + } + if not masked: + # Materialize unspecified elements with zero values + batch_dim = obj.ndim - obj.dense_dim() - obj.sparse_dim() + blocksize = ( + obj.values().shape[batch_dim + 1 : batch_dim + 3] + if obj.layout in sparse_block_layouts + else None + ) + full_mask = torch.ones( + obj.shape, device=obj.device, dtype=torch.bool + ).to_sparse( + layout=obj.layout, + blocksize=blocksize, + dense_dim=obj.dense_dim(), + ) + obj = obj.to_dense().sparse_mask(full_mask) + if obj.layout is torch.sparse_coo: + # pyrefly: ignore [no-matching-overload] + d.update( + # pyrefly: ignore [bad-argument-type] + indices=obj._indices(), + # pyrefly: ignore [bad-argument-type] + is_coalesced=obj.is_coalesced(), + ) + values = obj._values() + elif obj.layout in {torch.sparse_csr, torch.sparse_bsr}: + # pyrefly: ignore [no-matching-overload] + d.update( + # pyrefly: ignore [bad-argument-type] + compressed_indices=obj.crow_indices(), + # pyrefly: ignore [bad-argument-type] + plain_indices=obj.col_indices(), + ) + values = obj.values() + else: + # pyrefly: ignore [no-matching-overload] + d.update( + # pyrefly: ignore [bad-argument-type] + compressed_indices=obj.ccol_indices(), + # pyrefly: ignore [bad-argument-type] + plain_indices=obj.row_indices(), + ) + values = obj.values() + new_args.extend( + (STRIDED_REPRESENTATION, d, values.requires_grad_(True)) + ) + else: + new_args.append(obj) + return tuple(new_args) + + def restore_from_strided_representation(args): + """Restore non-strided differentiable tensosr from their strided representations.""" + new_args = [] + args = list(args) + while args: + a = args.pop(0) + if a == STRIDED_REPRESENTATION: + d, values = args.pop(0), args.pop(0) + if d["layout"] is torch.sparse_coo: + a = torch.sparse_coo_tensor( + d["indices"], + values, + size=d["shape"], + is_coalesced=d["is_coalesced"], + ) + elif d["layout"] in sparse_compressed_layouts: + a = torch.sparse_compressed_tensor( + d["compressed_indices"], + d["plain_indices"], + values, + size=d["shape"], + layout=d["layout"], + ) + else: + raise NotImplementedError( + f"conversion of {d['layout']} strided representation to tensor" + ) + new_args.append(a) + return tuple(new_args) + + def func_wrapper(*args, **kwargs): + restored_args = restore_from_strided_representation(args) + + # convert differentiable output sparse tensors to strided + # tensors: + outputs = func(*restored_args, **kwargs) + + strided_outputs = ( + tuple(outputs) if isinstance(outputs, (list, tuple)) else (outputs,) + ) + strided_outputs = tuple( + ( + o.to_dense(masked_grad=masked) + if isinstance(o, torch.Tensor) + and o.requires_grad + and o.layout in sparse_layouts + else o + ) + for o in strided_outputs + ) + + return ( + strided_outputs + if isinstance(outputs, (list, tuple)) + else strided_outputs[0] + ) + + args = (func_wrapper, convert_to_strided_representation(inputs)) + + return gradcheck(*args, **kwargs) + + return gradcheck_with_sparse_support diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_semi_structured_conversions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_semi_structured_conversions.py new file mode 100644 index 0000000000000000000000000000000000000000..354acdee16a26eeef851e344f9956ac11523cb92 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_semi_structured_conversions.py @@ -0,0 +1,358 @@ +# mypy: allow-untyped-defs +import torch + + +def _calculate_meta_reordering_scatter_offsets(m, meta_ncols, meta_dtype, device): + """ + This is PyTorch implementation of main part of reorder_meta() + function, from tools/util/include/cutlass/util/host_reorder.h file + of CUTLASS source tree. Furthermore, CUTLASS template for sparse + GEMM decides upon layout of this matrix, and at the moment for the + sparse GEMM executed on tensor cores, this is layout described by + ColumnMajorInterleaved<2> data structure, in + include/cutlass/layout/matrix.h of CUTLASS source tree. The + reordering of meta matrix into meta_reordered matrix calculated + according to these segments of CUTLASS code is re-implemented here. + Note that this calculation produces offsets for scattering metadata + matrix elements into reordered metadata matrix elements (or, + equivalently, for gathering reordered metadata matrix element back + into metadata matrix elements). + """ + dst_rows = torch.arange(0, m, device=device)[:, None].repeat(1, meta_ncols) + dst_cols = torch.arange(0, meta_ncols, device=device).repeat(m, 1) + + # Reorder the rows, then swizzle the 2x2 blocks. + group = 32 if meta_dtype.itemsize == 2 else 16 + interweave = 4 if meta_dtype.itemsize == 2 else 2 + dst_rows = ( + dst_rows // group * group + + (dst_rows % 8) * interweave + + (dst_rows % group) // 8 + ) + + topright = ((dst_rows % 2 == 0) & (dst_cols % 2 == 1)).to(torch.int8) + bottomleft = ((dst_rows % 2 == 1) & (dst_cols % 2 == 0)).to(torch.int8) + dst_rows += topright - bottomleft + dst_cols -= topright - bottomleft + + # Assumed that meta tensor is to be stored in CUTLASS + # InterleavedColumnMajor layout, and reverse engineered + # corresponding code to store values into this tensor. + interleave = 2 + cols_maj = dst_cols // interleave + cols_min = dst_cols % interleave + return (cols_maj * m * interleave + dst_rows * interleave + cols_min).view(-1) + + +def sparse_semi_structured_from_dense_cutlass(dense): + """ + This function converts dense matrix into sparse semi-structured + representation, producing "compressed" matrix, in the layout used by + CUTLASS backend, and corresponding metadata matrix. + """ + if dense.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional dense tensor, got {dense.dim()}-dimensional tensor" + ) + + m, k = dense.shape + device = dense.device + + meta_dtype = torch.int8 + if dense.dtype == torch.int8: + meta_dtype = torch.int32 + elif dense.dtype in [torch.half, torch.bfloat16, torch.float]: + meta_dtype = torch.int16 + else: + raise RuntimeError(f"Invalid datatype {dense.dtype} of dense matrix") + quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 + if quadbits_per_meta_elem not in (4, 8): + raise RuntimeError("Invalid number of elements per meta element calculated") + + if meta_dtype == torch.int32: + if m % 16 != 0: + raise RuntimeError( + f"Number of rows of dense matrix {m} must be divisible by 16" + ) + else: + if m % 32 != 0: + raise RuntimeError( + f"Number of rows of dense matrix {m} must be divisible by 32" + ) + if k % (4 * quadbits_per_meta_elem) != 0: + raise RuntimeError( + f"Number of columns of dense matrix {k} must be divisible by {4 * quadbits_per_meta_elem}" + ) + + if dense.dtype != torch.float: + ksparse = 4 + dense_4 = dense.view(-1, k // ksparse, ksparse) + m0, m1, _m2, m3 = (dense_4 != 0).unbind(-1) + else: + ksparse = 2 + dense_2 = dense.view(-1, k // ksparse, ksparse) + m0, _m2 = m1, m3 = (dense_2 != 0).unbind(-1) + meta_ncols = k // (ksparse * quadbits_per_meta_elem) + + # Encoding quadruples of True/False values as follows: + # [True, True, False, False] -> 0b0100 + # [True, False, True, False] -> 0b1000 + # [False, True, True, False] -> 0b1001 + # [True, False, False, True ] -> 0b1100 + # [False, True, False, True ] -> 0b1101 + # [False, False, True, True ] -> 0b1110 + # Thus, lower two bits in the encoding are index of the True value + # at the lowest index in the quadruple, and the higher two bits in + # the encoding are index of the other True value in the quadruple. + # In case there are less than two True values, than False value or + # values at some index or indices are considered True for the + # encoding. In case there are more than two True values, then the + # excess True value(s) at some indices are considered False for + # the encoding. The exact encodings used for these cases are as + # follows: + # [False, False, False, False] -> 0b1110 + # [False, False, False, True ] -> 0b1110 + # [False, False, True, False] -> 0b1110 + # [False, True, False, False] -> 0b1001 + # [False, True, True, True ] -> 0b1101 + # [True, False, False, False] -> 0b1000 + # [True, False, True, True ] -> 0b1100 + # [True, True, False, True ] -> 0b0100 + # [True, True, True, False] -> 0b0100 + # [True, True, True, True ] -> 0b0100 + # These particular encodings are chosen, with the help of Espresso + # logic minimizer software, for the purpose of minimization of + # corresponding Boolean functions, that translate non-zero flags + # into encoding bits. Note also possible choices for the first + # and last of these encodings were limited only to (0b0100, + # 0b1110), in order to produce valid encodings for 1:2 sparsity + # case. + + expr0 = m0 & m1 + expr1 = ~m0 & m1 + expr2 = ~m0 & ~m1 + bit0 = expr1 + bit1 = expr2 + bit2 = expr0 | expr2 | m3 + bit3 = expr1 | ~m1 + idxs0 = bit0 | (bit1.to(torch.int64) << 1) + idxs1 = bit2 | (bit3.to(torch.int64) << 1) + + if dense.dtype != torch.float: + sparse0 = dense_4.gather(-1, idxs0.unsqueeze(-1)) # type: ignore[possibly-undefined] + # pyrefly: ignore [unbound-name] + sparse1 = dense_4.gather(-1, idxs1.unsqueeze(-1)) + sparse = torch.stack((sparse0, sparse1), dim=-1).view(m, k // 2) + else: + sparse = dense_2.gather(-1, idxs0.unsqueeze(-1) // 2).view(m, k // 2) # type: ignore[possibly-undefined] + + meta_4 = idxs0 | (idxs1 << 2) + meta_n = meta_4.view((-1, meta_ncols, quadbits_per_meta_elem)).to(meta_dtype) + + if quadbits_per_meta_elem == 4: + meta = ( + meta_n[:, :, 0] + | (meta_n[:, :, 1] << 4) + | (meta_n[:, :, 2] << 8) + | (meta_n[:, :, 3] << 12) + ) + elif quadbits_per_meta_elem == 8: + meta = ( + meta_n[:, :, 0] + | (meta_n[:, :, 1] << 4) + | (meta_n[:, :, 2] << 8) + | (meta_n[:, :, 3] << 12) + | (meta_n[:, :, 4] << 16) + | (meta_n[:, :, 5] << 20) + | (meta_n[:, :, 6] << 24) + | (meta_n[:, :, 7] << 28) + ) + + # Reorder meta tensor elements. + meta_reordered = meta.new_empty((m * meta_ncols,)) # type: ignore[possibly-undefined] + meta_offsets = _calculate_meta_reordering_scatter_offsets( + m, meta_ncols, meta_dtype, device + ) + # pyrefly: ignore [unbound-name] + meta_reordered.scatter_(0, meta_offsets, meta.view(-1)) + + return (sparse, meta_reordered.view(m, meta_ncols)) + + +def sparse_semi_structured_to_dense_cutlass(sparse, meta_reordered): + """ + This function performs reverse of the function above - it + reconstructs dense matrix from a pair of "compressed" matrix, given + in the layout used by CUTLASS backend, and accompanying metadata + matrix. + """ + if sparse.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional sparse tensor, got {sparse.dim()}-dimensional tensor" + ) + + m, k = sparse.shape + device = sparse.device + + if meta_reordered.dim() != 2: + raise RuntimeError( + f"Expected 2-dimensional meta tensor, got {meta_reordered.dim()}-dimensional tensor" + ) + if meta_reordered.device != device: + raise RuntimeError( + f"Expected meta matrix to be on {device} device, got matrix on {meta_reordered.device} device" + ) + + meta_dtype = meta_reordered.dtype + if meta_dtype not in (torch.int16, torch.int32): + raise RuntimeError(f"Invalid datatype {meta_dtype} of meta matrix") + quadbits_per_meta_elem = meta_dtype.itemsize * 8 // 4 + + if sparse.dtype != torch.float: + ksparse = 4 + else: + ksparse = 2 + + meta_nrows, meta_ncols = meta_reordered.shape + if meta_nrows != m: + raise RuntimeError( + f"Number of rows of meta matrix {meta_nrows} must be equal to number of columns of spase matrix {m}" + ) + if meta_ncols * ksparse * quadbits_per_meta_elem != 2 * k: + raise RuntimeError( + f"Number of columns of sparse matrix {k} different from the {meta_ncols * ksparse * quadbits_per_meta_elem // 2}, " + "expected according to the number of columns of meta matrix" + ) + + # Undo meta tensor elements reordering. + meta_offsets = _calculate_meta_reordering_scatter_offsets( + m, meta_ncols, meta_dtype, device + ) + meta = torch.gather(meta_reordered.view(-1), 0, meta_offsets).view(m, meta_ncols) + + # Unpack sparse tensor back to original dense tensor, using + # information provided by meta tensor. Note that torch.float + # datatype is handled pretty much the same as + # torch.half/torch.bfloat16, as metadata for a pair of torch.float + # value is encoded as if underlying 8 bytes contain four + # torch.half/torch.bfloat16 values, where either first two or last + # two are zeros. + meta_2 = torch.empty( + (m, meta_ncols, 2 * quadbits_per_meta_elem), + dtype=meta_dtype, + device=device, + ) + if quadbits_per_meta_elem == 4: + meta_2[:, :, 0] = meta & 0b11 + meta_2[:, :, 1] = (meta >> 2) & 0b11 + meta_2[:, :, 2] = (meta >> 4) & 0b11 + meta_2[:, :, 3] = (meta >> 6) & 0b11 + meta_2[:, :, 4] = (meta >> 8) & 0b11 + meta_2[:, :, 5] = (meta >> 10) & 0b11 + meta_2[:, :, 6] = (meta >> 12) & 0b11 + meta_2[:, :, 7] = (meta >> 14) & 0b11 + elif quadbits_per_meta_elem == 8: + meta_2[:, :, 0] = meta & 0b11 + meta_2[:, :, 1] = (meta >> 2) & 0b11 + meta_2[:, :, 2] = (meta >> 4) & 0b11 + meta_2[:, :, 3] = (meta >> 6) & 0b11 + meta_2[:, :, 4] = (meta >> 8) & 0b11 + meta_2[:, :, 5] = (meta >> 10) & 0b11 + meta_2[:, :, 6] = (meta >> 12) & 0b11 + meta_2[:, :, 7] = (meta >> 14) & 0b11 + meta_2[:, :, 8] = (meta >> 16) & 0b11 + meta_2[:, :, 9] = (meta >> 18) & 0b11 + meta_2[:, :, 10] = (meta >> 20) & 0b11 + meta_2[:, :, 11] = (meta >> 22) & 0b11 + meta_2[:, :, 12] = (meta >> 24) & 0b11 + meta_2[:, :, 13] = (meta >> 26) & 0b11 + meta_2[:, :, 14] = (meta >> 28) & 0b11 + meta_2[:, :, 15] = (meta >> 30) & 0b11 + + dense_offsets = meta_2.view(-1) + ( + torch.arange(0, 2 * m * k // ksparse, device=device) * 4 + ).view(-1, 1).repeat(1, 2).view(-1) + + dense = torch.zeros((m * 2 * k,), dtype=sparse.dtype, device=device) + if sparse.dtype != torch.float: + dense.scatter_(0, dense_offsets, sparse.view(-1)) + else: + dense.view(torch.half).scatter_( + 0, dense_offsets, sparse.view(torch.half).view(-1) + ) + + return dense.view(m, 2 * k) + + +def _sparse_semi_structured_tile(dense): + """ + This function computes a 2:4 sparse tile by greedily taking the largest values. + + Since we take the largest values greedily, how the sorting algorithm handles duplicates affects + the ultimate sparsity pattern. + + Note that this function does not have the same sorting semantics as our CUDA backend, + which is exposed via `torch._sparse_semi_structured_tile` and thus returns a different pattern. + """ + + def greedy_prune_tile(tile): + num_kept_row = [0, 0, 0, 0] + num_kept_col = [0, 0, 0, 0] + + for x in tile.flatten().sort(descending=True, stable=True).indices: + r, c = x // 4, x % 4 + if num_kept_row[r] < 2 and num_kept_col[c] < 2: + num_kept_row[r] += 1 + num_kept_col[c] += 1 + else: + tile[r, c] = 0 + + for batch in dense.unfold(0, 4, 4).unfold(1, 4, 4): + for tile in batch: + greedy_prune_tile(tile) + + return dense + + +def _compute_compressed_swizzled_bitmask(dense): + """ + Calculates the compressed swizzled bitmask from a dense tensor + """ + + # first we need to convert the dense tensor to a bitmask + int_bitmask = dense.bool().to(torch.uint8) + + # Each thread is responsible for an 8x8 tile, which contains 4 4x4 tiles: + # A, B, C and D, as displayed in the following schema: + # +---+---+ + # | A | B | + # +---+---+ + # | C | D | + # +---+---+ + + # we first need to split into the 8x8 tiles + bitmask_8x8_chunks = int_bitmask.unfold(0, 8, 8).unfold(1, 8, 8) + + # then we unfold again to get our individual 4x4 tiles + bitmask_4x4_chunks = bitmask_8x8_chunks.unfold(2, 4, 4).unfold(3, 4, 4) + + # Each 4x4 bitmask defines two 8-bit integers, which encode the sparsity pattern + # of that tile. Note that the least significant bit is stored first. + # [1 1 0 0] + # [1 1 0 0] -> 0011 0011 -> 51 + # [0 0 1 1] 1100 1100 204 + # [0 0 1 1] + + # reshape tensor to expand tiles into 8-bit vectors + bitmask_binary_representation = bitmask_4x4_chunks.reshape( + *bitmask_4x4_chunks.shape[:2], 4, 2, 8 + ) + + # to convert from binary representation, we can do a matmul with powers of two + powers_of_two = 2 ** torch.arange(8, dtype=torch.float, device="cuda") + # To run on GPU: cast to float to do matmul and then cast back + compressed_swizzled_bitmask = ( + bitmask_binary_representation.to(torch.float) @ powers_of_two + ).to(torch.uint8) + + return compressed_swizzled_bitmask diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_semi_structured_ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_semi_structured_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..0254a33a4bc62b76484d9cd319e6f10b8dbbbdfa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_semi_structured_ops.py @@ -0,0 +1,300 @@ +# mypy: allow-untyped-defs +import contextlib + +import torch + + +__all__ = [ + "fallback_dispatcher", + "semi_sparse_values", + "semi_sparse_indices", + "semi_sparse_t", + "semi_sparse_view", + "semi_sparse_detach", + "semi_sparse_mm", + "semi_sparse_addmm", + "semi_sparse_linear", + "semi_sparse_scaled_mm", + "semi_sparse_clone", + "semi_sparse_to", + "semi_sparse_to_copy", +] + + +@contextlib.contextmanager +def no_dispatch(): + guard = torch._C._DisableTorchDispatch() + try: + yield + finally: + del guard + + +def fallback_dispatcher(func, types, args, kwargs): + with no_dispatch(): + return func(*args) + + +def semi_sparse_values(func, types, args=(), kwargs=None) -> torch.Tensor: + if len(args) != 1: + raise AssertionError(f"expected 1 arg, got {len(args)}") + A = args[0] + if not isinstance(A, torch.sparse.SparseSemiStructuredTensor): + raise AssertionError( + f"expected SparseSemiStructuredTensor, got {type(A).__name__}" + ) + if A.packed is None: + raise AssertionError("A.packed must not be None") + if A.meta is None: + m, k = A.shape + num_kept_elements = m * k // 2 + return A.packed.ravel()[:num_kept_elements:].view(m, -1) + else: + return A.packed.detach() + + +def semi_sparse_indices(func, types, args=(), kwargs=None) -> torch.Tensor: + if len(args) != 1: + raise AssertionError(f"expected 1 arg, got {len(args)}") + A = args[0] + if not isinstance(A, torch.sparse.SparseSemiStructuredTensor): + raise AssertionError( + f"expected SparseSemiStructuredTensor, got {type(A).__name__}" + ) + if A.packed is None: + raise AssertionError("A.packed must not be None") + if A.meta is None: + m, k = A.shape + num_kept_elements = m * k // 2 + metadata = A.packed.ravel()[num_kept_elements:].view(m, -1) + return metadata.view(torch.int32 if A.dtype == torch.int32 else torch.int16) + else: + return A.meta + + +def semi_sparse_t(func, types, args=(), kwargs=None) -> torch.Tensor: + if len(args) != 1: + raise AssertionError(f"expected 1 arg, got {len(args)}") + self = args[0] + if not isinstance(self, torch.sparse.SparseSemiStructuredTensor): + raise AssertionError( + f"expected SparseSemiStructuredTensor, got {type(self).__name__}" + ) + if len(self.shape) != 2: + raise AssertionError(f"expected 2D tensor, got {len(self.shape)}D") + # Because we cannot go from the compressed representation back to the dense representation currently, + # we just keep track of how many times we have been transposed. Depending on whether the sparse matrix + # is the first or second argument, we expect an even / odd number of calls to transpose respectively. + # pyrefly: ignore [no-matching-overload] + return self.__class__( + torch.Size([self.shape[-1], self.shape[0]]), + packed=self.packed_t, + meta=self.meta_t, + packed_t=self.packed, + meta_t=self.meta, + compressed_swizzled_bitmask=( + self.compressed_swizzled_bitmask.transpose(0, 1) + if self.compressed_swizzled_bitmask is not None + else None + ), + fuse_transpose_cusparselt=args[0].fuse_transpose_cusparselt, + alg_id_cusparselt=args[0].alg_id_cusparselt, + ) + + +def semi_sparse_view(func, types, args=(), kwargs=None) -> torch.Tensor: + if len(args) != 2: + raise AssertionError(f"expected 2 args, got {len(args)}") + self, shape = args + if tuple(shape) != self.shape: + raise NotImplementedError( + f"`view` is not implemented for SparseSemiStructuredTensor, except for the dummy case (shape={shape})" + ) + return self + + +def semi_sparse_detach(func, types, args, kwargs) -> torch.Tensor: + if len(args) != 1: + raise AssertionError(f"expected 1 arg, got {len(args)}") + self = args[0] + return self.__class__( + shape=self.shape, + packed=self.packed, + meta=self.meta, + packed_t=self.packed_t, + meta_t=self.meta_t, + compressed_swizzled_bitmask=self.compressed_swizzled_bitmask, + fuse_transpose_cusparselt=self.fuse_transpose_cusparselt, + alg_id_cusparselt=self.alg_id_cusparselt, + requires_grad=False, + ) + + +def semi_sparse_mm(func, types, args=(), kwargs=None) -> torch.Tensor: + if len(args) != 2: + raise AssertionError(f"expected 2 args, got {len(args)}") + A, B = args + if A.ndim != 2 or B.ndim != 2: + raise NotImplementedError( + "`SparseSemiStructuredTensor` matmul: Broadcasting is not implemented" + ) + if isinstance(A, torch.sparse.SparseSemiStructuredTensor): + return A._mm(B) + else: + B_t = B.t() + if not isinstance(B_t, torch.sparse.SparseSemiStructuredTensor): + raise AssertionError( + f"expected SparseSemiStructuredTensor, got {type(B_t).__name__}" + ) + return B_t._mm(A, should_transpose_dense=True).t() + + +def semi_sparse_addmm(func, types, args=(), kwargs=None) -> torch.Tensor: + if len(args) != 3: + raise AssertionError(f"expected 3 args, got {len(args)}") + bias, A, B = args + if A.ndim != 2 or B.ndim != 2: + raise NotImplementedError( + "`SparseSemiStructuredTensor` matmul: Broadcasting is not implemented" + ) + if bias.ndim != 1: + raise NotImplementedError( + f"`SparseSemiStructuredTensor` matmul: only bias dim=1 supported. Shape={bias.shape}" + ) + if isinstance(A, torch.sparse.SparseSemiStructuredTensor): + raise NotImplementedError( + "`SparseSemiStructuredTensor` matmul: only operand B of `addmm` can be sparse" + ) + B_t = B.t() + if not isinstance(B_t, torch.sparse.SparseSemiStructuredTensor): + raise AssertionError( + f"expected SparseSemiStructuredTensor, got {type(B_t).__name__}" + ) + row, _col = A.shape + return B_t._mm(A, bias=bias, should_transpose_dense=True).t() + + +def semi_sparse_linear(func, types, args=(), kwargs=None) -> torch.Tensor: + if len(args) not in [2, 3]: + raise AssertionError(f"expected 2 or 3 args, got {len(args)}") + A, B = args[:2] + bias = args[2] if len(args) == 3 else None + + shape = A.shape + A_2d = A.view(-1, shape[-1]) + if bias is None: + res = A_2d @ B.t() + else: + res = semi_sparse_addmm( + func=None, + types=None, + args=[bias, A_2d, B.t()], + ) + return res.view(*shape[:-1], -1) + + +def semi_sparse_scaled_mm(func, types, args=(), kwargs=None) -> torch.Tensor: + # pull all args, excluding use_fast_accum flag if set. + A, B, A_scale, B_scale, bias, scale_result, out_dtype = args[:7] + + if A.dtype != torch.float8_e4m3fn: + raise AssertionError(f"expected A.dtype float8_e4m3fn, got {A.dtype}") + if B.dtype != torch.float8_e4m3fn: + raise AssertionError(f"expected B.dtype float8_e4m3fn, got {B.dtype}") + # only cuSPARSELt supports float8_e4m3fn currently + if not isinstance(A, torch.sparse.SparseSemiStructuredTensorCUSPARSELT): + raise AssertionError( + f"expected SparseSemiStructuredTensorCUSPARSELT, got {type(A).__name__}" + ) + if A.packed is None: + raise AssertionError("A.packed must not be None") + # Currently we only support per-tensor scaling, with float32 scales + if A_scale.numel() != 1 or B_scale.numel() != 1: + raise AssertionError( + f"expected A_scale and B_scale to have numel 1, got {A_scale.numel()} and {B_scale.numel()}" + ) + if A_scale.dtype != torch.float32 or B_scale.dtype != torch.float32: + raise AssertionError( + f"expected A_scale and B_scale dtype float32, got {A_scale.dtype} and {B_scale.dtype}" + ) + + # cuSPARSELt lacks the A and B operand scaling support, so instead we use alpha to scale the result. + # Note that this limits us to per-tensor scalig only. + sparse_result = torch._cslt_sparse_mm( + A.packed, + B, + alpha=A_scale * B_scale, + out_dtype=out_dtype, + ) + return sparse_result + + +def semi_sparse_clone(func, types, args=(), kwargs=None) -> torch.Tensor: + if len(args) != 1: + raise AssertionError(f"expected 1 arg, got {len(args)}") + + self = args[0] + if not isinstance(self, torch.sparse.SparseSemiStructuredTensor): + raise AssertionError( + f"expected SparseSemiStructuredTensor, got {type(self).__name__}" + ) + + # pyrefly: ignore [no-matching-overload] + return self.__class__( + shape=self.shape, + packed=None if self.packed is None else self.packed.clone(), + meta=None if self.meta is None else self.meta.clone(), + packed_t=None if self.packed_t is None else self.packed_t.clone(), + meta_t=None if self.meta_t is None else self.meta_t.clone(), + compressed_swizzled_bitmask=( + None + if self.compressed_swizzled_bitmask is None + else self.compressed_swizzled_bitmask.clone() + ), + fuse_transpose_cusparselt=self.fuse_transpose_cusparselt, + alg_id_cusparselt=self.alg_id_cusparselt, + requires_grad=self.requires_grad, + ) + + +def semi_sparse_to_copy(func, types, args, kwargs=None) -> torch.Tensor: + self = args[0] + kwargs = kwargs or {} + + device = kwargs.get("device", None) + + if device is not None and torch.device(device).type == "cpu": + dense = self.to_dense() + return func(dense, **kwargs) + + raise NotImplementedError( + f"`_to_copy()` with kwargs={kwargs} is not implemented " + "for SparseSemiStructuredTensor. Only converting to CPU is supported currently." + ) + + +def semi_sparse_to(func, types, args, kwargs=None) -> torch.Tensor: + self = args[0] + remaining_args = args[1:] + kwargs = kwargs or {} + + # Determine the target device from args/kwargs + device = None + if remaining_args: + first_arg = remaining_args[0] + if isinstance(first_arg, (torch.device, str)): + try: + device = torch.device(first_arg) + except RuntimeError: + pass + if "device" in kwargs: + device = torch.device(kwargs["device"]) + + if device is not None and device.type == "cpu": + dense = self.to_dense() + return func(dense, *remaining_args, **kwargs) + + raise NotImplementedError( + f"`to()` with args={remaining_args}, kwargs={kwargs} is not implemented " + "for SparseSemiStructuredTensor. Only `to('cpu')` is supported currently." + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_triton_ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_triton_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..106306c6e4cdee3c4d8ef7d2b786a06b10e0a02a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_triton_ops.py @@ -0,0 +1,2602 @@ +# mypy: allow-untyped-decorators +# mypy: allow-untyped-defs +import math +import os +import weakref +from functools import lru_cache + +import torch +from torch._dynamo.utils import warn_once +from torch.utils._triton import has_triton + +from ._triton_ops_meta import get_meta + + +TORCH_SPARSE_BSR_SCATTER_MM_LRU_CACHE_SIZE = int( + os.getenv("TORCH_SPARSE_BSR_SCATTER_MM_LRU_CACHE_SIZE", 2) +) + + +def check(cond, msg): + if not cond: + raise ValueError(msg) + + +def check_bsr_layout(f_name, t): + check( + t.layout == torch.sparse_bsr, + f"{f_name}(): only BSR sparse format is supported for the sparse argument.", + ) + + +def check_device(f_name, t, device): + check( + t.device == device and t.device.type == "cuda", + f"{f_name}(): all inputs are expected to be on the same GPU device.", + ) + + +def check_mm_compatible_shapes(f_name, lhs, rhs): + check( + lhs.dim() >= 2 and rhs.dim() >= 2, + f"{f_name}(): all inputs involved in the matrix product are expected to be at least 2D, " + f"but got lhs.dim() == {lhs.dim()} and rhs.dim() == {rhs.dim()}.", + ) + + _m, kl = lhs.shape[-2:] + kr, _n = rhs.shape[-2:] + + check( + kl == kr, + f"{f_name}(): arguments' sizes involved in the matrix product are not compatible for matrix multiplication, " + f"got lhs.shape[-1] == {kl} which is not equal to rhs.shape[-2] == {kr}.", + ) + + +def check_dtype(f_name, t, dtype, *additional_dtypes): + check( + t.dtype == dtype + and t.dtype + in ((torch.half, torch.bfloat16, torch.float) + tuple(*additional_dtypes)), + f"{f_name}(): all inputs are expected to be of the same dtype " + f"and one of (half, bfloat16, float32) or {additional_dtypes}, " + f"but got dtype == {t.dtype}.", + ) + + +def check_blocksize(f_name, blocksize): + if len(blocksize) != 2: + raise AssertionError(f"blocksize must have length 2, got {len(blocksize)}") + + def is_power_of_two(v): + return not (v & (v - 1)) + + def is_compatible_blocksize(b): + res = True + for blocksize in b: + # Triton loads only blocks which are at least 16 and powers of 2. + res = (blocksize >= 16 and is_power_of_two(blocksize)) and res + return res + + check( + is_compatible_blocksize(blocksize), + f"{f_name}(): sparse inputs' blocksize ({blocksize[0]}, {blocksize[1]}) " + "should be at least 16 and a power of 2 in each dimension.", + ) + + +def make_triton_contiguous(t): + """Return input as a triton-contiguous tensor. + + A triton-contiguous tensor is defined as a tensor that has strides + with minimal value smaller than or equal to 1. + + While triton kernels support triton-non-contiguous tensors (all + strides being greater than 1) arguments, a considerable slow-down + occurs because tensor data is copied element-wise rather than + chunk-wise. Zero strides is assumed to not have this defect. + """ + if min(t.stride()) > 1: + # TODO: investigate if contiguity along other axes than the + # last one can be beneficial for performance + return t.contiguous() + else: + return t + + +def broadcast_batch_dims(f_name, *tensors): + try: + return torch.broadcast_shapes(*(t.shape[:-2] for t in tensors)) + except Exception: + check(False, f"{f_name}(): inputs' batch dimensions are not broadcastable!") + + +def slicer(dim, slice_range, *tensors): + for t in tensors: + slices = [slice(None)] * t.dim() + slices[dim] = slice_range + yield t[slices] + + +def multidim_slicer(dims, slices, *tensors): + for t in tensors: + s = [slice(None)] * t.dim() + for d, d_slice in zip(dims, slices, strict=False): + if d is not None: + s[d] = d_slice + yield t[tuple(s)] + + +def ptr_stride_extractor(*tensors): + for t in tensors: + yield t + yield from t.stride() + + +def grid_partitioner(full_grid, grid_blocks, tensor_dims_map): + if len(full_grid) < 0 or len(full_grid) > 3: + raise AssertionError(f"full_grid length must be 0-3, got {len(full_grid)}") + if len(grid_blocks) < 0 or len(grid_blocks) > 3: + raise AssertionError(f"grid_blocks length must be 0-3, got {len(grid_blocks)}") + + import itertools + + def generate_grid_points(): + for fg, mg in zip(full_grid, grid_blocks, strict=False): + yield range(0, fg, mg) + + def generate_sliced_tensors(slices): + for t, t_dims in tensor_dims_map.items(): + yield next(multidim_slicer(t_dims, slices, t)) + + for grid_point in itertools.product(*generate_grid_points()): + grid = [ + min(fg - gp, mg) + for fg, gp, mg in zip(full_grid, grid_point, grid_blocks, strict=False) + ] + slices = [slice(gp, gp + g) for gp, g in zip(grid_point, grid, strict=False)] + # grid_points are iterated in a "contiguous" order, i.e. + # left dimensions traversed slower than right dimensions. + # This order is reversed for CUDA grids. + yield grid[::-1], *generate_sliced_tensors(slices) + + +def launch_kernel(kernel, tensor_dims_map, full_grid, grid_blocks=None): + # cuda_max_grid = (2 ** 31 - 1, 2 ** 16 - 1, 2 ** 16 - 1) + cuda_max_grid = (2147483647, 65535, 65535)[::-1] + if grid_blocks is None: + grid_blocks = cuda_max_grid + else: + + def valid_grid_dim(g, mg): + if g is None: + return mg + else: + # grid must be at least 1 and no greater than mg + return max(1, min(g, mg)) + + grid_blocks = tuple( + valid_grid_dim(g, mg) + for g, mg in zip(grid_blocks, cuda_max_grid, strict=False) + ) # type: ignore[assignment] + + for grid, *sliced_tensors in grid_partitioner( + full_grid, grid_blocks, tensor_dims_map + ): + kernel(grid, *sliced_tensors) + + +def prepare_inputs(bsr, *dense_tensors): + # Introduce fake batch dimension if not present for convenience. + crow_indices = bsr.crow_indices().unsqueeze(0) + col_indices = bsr.col_indices().unsqueeze(0) + values = make_triton_contiguous(bsr.values().unsqueeze(0)) + tensors = [make_triton_contiguous(t.unsqueeze(0)) for t in dense_tensors] + + # Compute broadcasted batch dimension + batch_dims_broadcasted = torch.broadcast_shapes( + values.shape[:-3], *(t.shape[:-2] for t in tensors) + ) + + # Broadcast batch dimensions and squash. + # The result can be either a view or a copy. + def batch_broadcast_and_squash(t, batch_dims, invariant_dims): + return t.broadcast_to(batch_dims + invariant_dims).flatten( + 0, len(batch_dims) - 1 + ) + + crow_indices = batch_broadcast_and_squash( + crow_indices, batch_dims_broadcasted, (-1,) + ) + + col_indices = batch_broadcast_and_squash(col_indices, batch_dims_broadcasted, (-1,)) + values = batch_broadcast_and_squash( + values, batch_dims_broadcasted, values.shape[-3:] + ) + tensors = [ + batch_broadcast_and_squash(t, batch_dims_broadcasted, t.shape[-2:]) + for t in tensors + ] + + return crow_indices, col_indices, values, *tensors + + +def broadcast_batch_dims_bsr(f_name, bsr, *tensors): + batch_shape = broadcast_batch_dims(f_name, bsr, *tensors) + + crow_indices = bsr.crow_indices().broadcast_to(batch_shape + (-1,)) + col_indices = bsr.col_indices().broadcast_to(batch_shape + (-1,)) + values = bsr.values().broadcast_to(batch_shape + bsr.values().shape[-3:]) + size = batch_shape + bsr.shape[-2:] + return torch.sparse_compressed_tensor( + crow_indices, col_indices, values, size=size, layout=bsr.layout + ) + + +# NOTE: this function will ALWAYS create a view +def tile_to_blocksize(t, blocksize): + *rest, m, n = t.shape + new_shape = rest + [ + m // blocksize[0], + blocksize[0], + n // blocksize[1], + blocksize[1], + ] + # using .view instead of .reshape to ensure that the result is + # indeed a view: + return t.view(new_shape).transpose(-3, -2) + + +def as1Dbatch(tensor): + """Return tensor as 3D tensor by either prepending new dimensions to + the tensor shape (when ``tensor.ndim < 3``), or by collapsing + starting dimensions into the first dimension (when ``tensor.ndim > + 3``). + """ + while tensor.ndim < 3: + tensor = tensor.unsqueeze(0) + if tensor.ndim > 3: + tensor = tensor.flatten(0, tensor.ndim - 3) + if tensor.ndim != 3: + raise AssertionError( + f"tensor should have 3 dimensions after reshape, got {tensor.shape}" + ) + return tensor + + +def scatter_mm(blocks, others, indices_data, *, accumulators=None): + """Scattered matrix multiplication of tensors. + + A scattered matrix multiplication is defined as a series of matrix + multiplications applied to input tensors according to the input + and output mappings specified by indices data. + + The following indices data formats are supported for defining a + scattered matrix multiplication operation (:attr:`indices_data[0]` + holds the name of the indices data format as specified below): + + - ``"scatter_mm"`` - matrix multiplications scattered in batches + of tensors. + + If :attr:`blocks` is a :math:`(* \times M \times K) tensor, + :attr:`others` is a :math:`(* \times K \times N)` tensor, + :attr:`accumulators` is a :math:`(* \times M \times N)` tensor, + and :attr:`indices = indices_data['indices']` is a :math:`(* + \times 3)` tensor, then the operation is equivalent to the + following code:: + + c_offsets, pq = indices_data[1:] + for r in range(len(c_offsets) - 1): + for g in range(c_offsets[r], c_offsets[r + 1]): + p, q = pq[g] + accumulators[r] += blocks[p] @ others[q] + + - ``"bsr_strided_mm"`` - matrix multiplications scattered in + batches of tensors and a tensor. + + If :attr:`blocks` is a :math:`(Ms \times Ks) tensor, + :attr:`others` is a :math:`(* \times K \times N)` tensor, + :attr:`accumulators` is a :math:`(* \times M \times N)` tensor, then + the operation is equivalent to the following code:: + + c_indices, r_offsets, p_offsets, q_offsets, meta = indices_data[1:] + for b in range(nbatches): + for i, r in enumerate(r_offsets): + r0, r1 = divmod(r, N) + acc = accumulators[b, r0 : r0 + Ms, r1 : r1 + Ns] + for g in range(c_indices[i], c_indices[i + 1]): + p = p_offsets[g] + q0, q1 = divmod(q_offsets[g], N) + acc += blocks[p] @ others[b, q0 : q0 + Ks, q1 : q1 + Ns] + + where ``Ns = N // meta['SPLIT_N']``, and ``M`` and ``K`` are + integer multiples of ``Ms`` and ``Ks``, respectively. + + - ``"bsr_strided_mm_compressed"`` - matrix multiplications + scattered in batches of tensors and a tensor. A memory and + processor efficient version of ``"bsr_strided_mm"`` format. If + :attr:`blocks` is a :math:`(Ms \times Ks) tensor, :attr:`others` + is a :math:`(* \times K \times N)` tensor, :attr:`accumulators` + is a :math:`(* \times M \times N)` tensor, then the operation is + equivalent to the following code:: + + c_indices, r_offsets, q_offsets, meta = indices_data[1:] + for b in range(nbatches): + for r in r_offsets: + m = (r // N) // Ms + n = (r % N) // Ns + r0, r1 = divmod(r, N) + c0, c1 = c_indices[m], c_indices[m + 1] + acc = accumulators[b, r0 : r0 + Ms, r1 : r1 + Ns] + for i, p in enumerate(range(c0, c1)): + q = q_offsets[n * c1 + (SPLIT_N - n) * c0 + i] + q0, q1 = divmod(q, N) + acc += blocks[p] @ others[b, q0 : q0 + Ks, q1 : q1 + Ns] + + where ``Ns = N // meta['SPLIT_N']``, and ``M`` and ``K`` are + integer multiples of ``Ms`` and ``Ks``, respectively. + + Notice that the order of ``r_offsets`` items can be arbitrary; + this property enables defining swizzle operators via + rearrangements of ``r_offsets`` items.. + + Auxiliary functions are provided for pre-computing + :attr:`indices_data`. For example, + :func:`bsr_scatter_mm_indices_data` is used to define indices data + for matrix multiplication of BSR and strided tensors. + + Parameters + ---------- + blocks (Tensor): a 3-D tensor of first matrices to be multiplied + + others (Tensor): a tensor of second matrices to be multiplied. If + ``indices_data[0]=="scatter_mm"``, the tensor is a 1-D batch + tensor of second input matrices to be multiplied. Otherwise, the + second input matrices are slices of the :attr:`others` tensor. + indices_data (tuple): a format data that defines the inputs and + outputs of scattered matrix multiplications. + + Keyword arguments + ----------------- + + accumulators (Tensor, optional): a tensor of matrix product + accumulators. If ``indices_data[0]=="scatter_mm"``, the tensor + is a 1-D batch tensor of output matrices. Otherwise, output + matrices are slices of the :attr:`accumulators` tensor. + """ + indices_format = indices_data[0] + + if blocks.ndim != 3: + raise AssertionError(f"blocks must be 3D, got {blocks.ndim}D") + _P, Ms, Ks = blocks.shape + + if indices_format == "scatter_mm": + c_offsets, pq = indices_data[1:] + + if others.ndim != 3: + raise AssertionError(f"others must be 3D, got {others.ndim}D") + _Q, Ks_, Ns = others.shape + if Ks != Ks_: + raise AssertionError(f"blocks K ({Ks}) != others K ({Ks_})") + + if accumulators is None: + R = c_offsets.shape[0] - 1 + accumulators = torch.zeros( + (R, Ms, Ns), dtype=blocks.dtype, device=blocks.device + ) + else: + R, Ms_, Ns_ = accumulators.shape + if Ms_ != Ms: + raise AssertionError(f"accumulators Ms ({Ms_}) != blocks Ms ({Ms})") + if Ns_ != Ns: + raise AssertionError(f"accumulators Ns ({Ns_}) != others Ns ({Ns})") + + if Ms % 16 or Ks % 16 or Ns % 16 or _scatter_mm2 is None: + for r in range(c_offsets.shape[0] - 1): + g0 = c_offsets[r] + g1 = c_offsets[r + 1] + for g in range(g0, g1): + p, q = pq[g] + + accumulators[r] += blocks[p] @ others[q] + else: + _scatter_mm2(blocks, others, c_offsets, pq, accumulators) + return accumulators + + elif indices_format == "bsr_strided_mm": + others_shape = others.shape + others = as1Dbatch(others) + + B, K, N = others.shape + if K % Ks != 0: + raise AssertionError(f"K ({K}) must be divisible by Ks ({Ks})") + + c_indices, r_offsets, p_offsets, q_offsets, meta = indices_data[1:] + SPLIT_N = meta["SPLIT_N"] + + if accumulators is None: + M = Ms + (r_offsets.max().item() + 1) // N + accumulators = torch.zeros( + (*others_shape[:-2], M, N), dtype=blocks.dtype, device=blocks.device + ) + else: + M, N_ = accumulators.shape[-2:] + if N_ != N: + raise AssertionError(f"accumulators N ({N_}) != others N ({N})") + + accumulators_shape = accumulators.shape + accumulators = as1Dbatch(accumulators) + + Ns = N // SPLIT_N + + if Ms % 16 or Ks % 16 or Ns % 16 or _scatter_mm6 is None: + accumulators.zero_() + for b in range(B): + for r in range(r_offsets.shape[0]): + r_ = r_offsets[r].item() + g0 = c_indices[r].item() + g1 = c_indices[r + 1].item() + r0, r1 = divmod(r_, N) + acc = accumulators[b, r0 : r0 + Ms, r1 : r1 + Ns] + for g in range(g0, g1): + p, q = p_offsets[g], q_offsets[g] + q0, q1 = divmod(q.item(), N) + acc += blocks[p] @ others[b, q0 : q0 + Ks, q1 : q1 + Ns] + else: + _scatter_mm6( + blocks, + others, + c_indices, + r_offsets, + p_offsets, + q_offsets, + meta, + accumulators, + ) + return accumulators.view(accumulators_shape) + + elif indices_format == "bsr_strided_mm_compressed": + others_shape = others.shape + others = as1Dbatch(others) + + B, K, N = others.shape + if K % Ks != 0: + raise AssertionError(f"K ({K}) must be divisible by Ks ({Ks})") + + c_indices, r_offsets, q_offsets, meta = indices_data[1:] + SPLIT_N = meta["SPLIT_N"] + + if accumulators is None: + M = Ms + (r_offsets.max().item() + 1) // N + accumulators = torch.zeros( + (*others_shape[:-2], M, N), dtype=blocks.dtype, device=blocks.device + ) + else: + M, N_ = accumulators.shape[-2:] + if N_ != N: + raise AssertionError(f"accumulators N ({N_}) != others N ({N})") + + accumulators_shape = accumulators.shape + accumulators = as1Dbatch(accumulators) + + Ns = N // SPLIT_N + + if Ms % 16 or Ks % 16 or Ns % 16 or _scatter_mm6 is None: + for b in range(B): + for j in range(len(r_offsets)): + r0, r1 = divmod(r_offsets[j].item(), N) + m = r0 // Ms + n = r1 // Ns + c0 = c_indices[m].item() + c1 = c_indices[m + 1].item() + acc = accumulators[b, r0 : r0 + Ms, r1 : r1 + Ns] + for i, p in enumerate(range(c0, c1)): + q = q_offsets[n * c1 + (SPLIT_N - n) * c0 + i].item() + q0, q1 = divmod(q, N) + acc += blocks[p] @ others[b, q0 : q0 + Ks, q1 : q1 + Ns] + else: + p_offsets = torch.empty( + (0,), dtype=q_offsets.dtype, device=q_offsets.device + ) + _scatter_mm6( + blocks, + others, + c_indices, + r_offsets, + p_offsets, + q_offsets, + meta, + accumulators, + ) + return accumulators.view(accumulators_shape) + + else: + raise NotImplementedError(indices_format) + + +def scatter_mm_meta( + M, + K, + N, + Ms, + Ks, + GROUP_SIZE=None, + TILE_M=None, + TILE_N=None, + SPLIT_N=None, + num_warps=None, + num_stages=None, + **extra, +): + if {TILE_M, TILE_N, SPLIT_N, num_warps, num_stages, GROUP_SIZE} == {None}: + device_name = torch.cuda.get_device_name() + meta = get_meta( + "scatter_mm", + (M, K, N, Ms, Ks), + device_name, + version=(0, torch.float16, 0.5), + ) + if meta is not None: + meta.update(**extra) + return meta + # The following parameters are optimized for the performance + # equilibrium points of bsr-dense and dense-dense matrix + # multiplications when using GPU card NVIDIA GeForce RTX 2060 + # SUPER. For points far from the performance equilibrium + # points as well as for other GPU cards, the optimal + # parameters are likely different from what specified below. + if (M, K, N) == (256,) * 3: + if (Ms, Ks) == (16, 16): + SPLIT_N = 1 + TILE_M = 16 + TILE_N = 16 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (Ms, Ks) == (32, 32): + SPLIT_N = 2 + TILE_M = 32 + TILE_N = 16 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (Ms, Ks) == (64, 64): + SPLIT_N = 1 + TILE_M = 32 + TILE_N = 32 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (Ms, Ks) == (128, 128): + SPLIT_N = 1 + TILE_M = 32 + TILE_N = 32 + GROUP_SIZE = 2 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (M, K, N) == (512,) * 3: + if (Ms, Ks) == (16, 16): + SPLIT_N = 8 + TILE_M = 16 + TILE_N = 64 + GROUP_SIZE = 2 + num_stages = 1 + num_warps = 2 # noqa: E225,E231,E702 + elif (Ms, Ks) == (32, 32): + SPLIT_N = 8 + TILE_M = 32 + TILE_N = 64 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 2 # noqa: E225,E231,E702 + elif (Ms, Ks) == (64, 64): + SPLIT_N = 4 + TILE_M = 32 + TILE_N = 128 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (Ms, Ks) == (128, 128): + SPLIT_N = 8 + TILE_M = 64 + TILE_N = 64 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (M, K, N) == (1024,) * 3: + if (Ms, Ks) == (16, 16): + SPLIT_N = 4 + TILE_M = 16 + TILE_N = 128 + GROUP_SIZE = 2 + num_stages = 1 + num_warps = 1 # noqa: E225,E231,E702 + elif (Ms, Ks) == (32, 32): + SPLIT_N = 8 + TILE_M = 32 + TILE_N = 64 + GROUP_SIZE = 2 + num_stages = 1 + num_warps = 1 # noqa: E225,E231,E702 + elif (Ms, Ks) == (64, 64): + SPLIT_N = 16 + TILE_M = 64 + TILE_N = 64 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 2 # noqa: E225,E231,E702 + elif (Ms, Ks) == (128, 128): + SPLIT_N = 16 + TILE_M = 64 + TILE_N = 64 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (Ms, Ks) == (256, 256): + SPLIT_N = 16 + TILE_M = 64 + TILE_N = 64 + GROUP_SIZE = 2 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (M, K, N) == (2048,) * 3: + if (Ms, Ks) == (16, 16): + SPLIT_N = 4 + TILE_M = 16 + TILE_N = 128 + GROUP_SIZE = 8 + num_stages = 1 + num_warps = 1 # noqa: E225,E231,E702 + elif (Ms, Ks) == (32, 32): + SPLIT_N = 4 + TILE_M = 32 + TILE_N = 64 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 1 # noqa: E225,E231,E702 + elif (Ms, Ks) == (64, 64): + SPLIT_N = 4 + TILE_M = 64 + TILE_N = 128 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (Ms, Ks) == (128, 128): + SPLIT_N = 8 + TILE_M = 64 + TILE_N = 64 + GROUP_SIZE = 4 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (Ms, Ks) == (256, 256): + SPLIT_N = 4 + TILE_M = 64 + TILE_N = 64 + GROUP_SIZE = 2 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + elif (M, K, N) == (4096,) * 3: + if (Ms, Ks) == (16, 16): + SPLIT_N = 2 + TILE_M = 16 + TILE_N = 256 + GROUP_SIZE = 2 + num_stages = 1 + num_warps = 2 # noqa: E225,E231,E702 + elif (Ms, Ks) == (32, 32): + SPLIT_N = 2 + TILE_M = 32 + TILE_N = 64 + GROUP_SIZE = 2 + num_stages = 1 + num_warps = 1 # noqa: E225,E231,E702 + elif (Ms, Ks) == (64, 64): + SPLIT_N = 2 + TILE_M = 64 + TILE_N = 128 + GROUP_SIZE = 2 + num_stages = 1 + num_warps = 4 # noqa: E225,E231,E702 + + if SPLIT_N is None: + # Assume NVIDIA GeForce RTX 2060 SUPER: + # With the probality of 92% (99.9% when N > 512), the + # performance will not be worse more than 2% from the + # performance when using an optimal value. Otherwise, when N + # <= 512, using the following heuristics may give upto 15% + # lower performance. + SPLIT_N = { + 16: 1, + 32: 2, + 64: 4, + 128: 8, + 256: 16, + 512: 8, + 1024: 16, + 4096: 32, + 8192: 64, + }.get(N, 16) + if Ms >= 512 and N >= 2048: + SPLIT_N = 1 + Ns = N // SPLIT_N + if TILE_M is None: + TILE_M = min(64 if Ns < 512 else 32, Ms) + if TILE_N is None: + TILE_N = min(64 if Ns < 512 else 32, Ns) + num_stages = num_stages or 1 + if num_warps is None: + if min(M, N) > 1024: + num_warps = {16: 1, 32: 1, 64: 2}.get(Ms, 4) + elif min(M, N) == 1024: + num_warps = {16: 1, 32: 1, 64: 2}.get(Ms, 4) + elif min(M, N) == 256: + num_warps = {16: 1, 32: 4}.get(Ms, 4) + else: + num_warps = {16: 1, 32: 2}.get(Ms, 4) + GROUP_SIZE = GROUP_SIZE or 4 + + if TILE_M > Ms: + raise AssertionError(f"TILE_M ({TILE_M}) must be <= Ms ({Ms})") + if TILE_N > Ns: + raise AssertionError(f"TILE_N ({TILE_N}) must be <= Ns ({Ns})") + if Ms > M: + raise AssertionError(f"Ms ({Ms}) must be <= M ({M})") + if Ns > N: + raise AssertionError(f"Ns ({Ns}) must be <= N ({N})") + if Ks > K: + raise AssertionError(f"Ks ({Ks}) must be <= K ({K})") + + return dict( + TILE_M=TILE_M, + TILE_N=TILE_N, + GROUP_SIZE=GROUP_SIZE, + num_stages=num_stages, + num_warps=num_warps, + SPLIT_N=SPLIT_N, + **extra, + ) + + +def bsr_dense_addmm_meta( + M, + K, + N, + Ms, + Ks, + beta, + alpha, + SPLIT_N=None, + GROUP_SIZE_ROW=None, + num_warps=None, + num_stages=None, + sparsity=None, + dtype=None, + out_dtype=None, + _version=0, + **extra, +): + # Specifying _version is useful for situations when one wants to + # discard existing triton kernel tuning results, say, in testing + # bsr_dense_addmm_meta functionality. + if dtype is None: + dtype = torch.float16 + if out_dtype is None: + out_dtype = dtype + if sparsity is None: + sparsity = 0.5 + if {SPLIT_N, num_warps, num_stages, GROUP_SIZE_ROW} == {None}: + device_name = torch.cuda.get_device_name() + key = (M, K, N, Ms, Ks, beta == 0, beta == 1, alpha == 1) + if dtype is out_dtype: + version_dtype = dtype + else: + version_dtype = dtype, out_dtype + meta = get_meta( + "bsr_dense_addmm", + key, + device_name, + version=(_version, version_dtype, sparsity), + ) + if meta is None and sparsity != 0.5: + meta = get_meta( + "bsr_dense_addmm", + key, + device_name, + version=(_version, version_dtype, 0.5), + ) + if meta is None and dtype is not out_dtype: + meta = get_meta( + "bsr_dense_addmm", key, device_name, version=(_version, dtype, 0.5) + ) + if meta is None: + # find approximate meta such that N % SPLIT_N == 0. + matching_meta = get_meta( + "bsr_dense_addmm", + (*key[:2], "*", *key[3:]), + device_name, + version=(_version, version_dtype, 0.5), + ) + if matching_meta is None and dtype is not out_dtype: + matching_meta = get_meta( + "bsr_dense_addmm", + (*key[:2], "*", *key[3:]), + device_name, + version=(_version, dtype, 0.5), + ) + for mkey in sorted(matching_meta or {}): + meta_ = matching_meta[mkey] + n = mkey[2] + split_n = meta_["SPLIT_N"] + c = n // split_n + if N % c == 0 and n <= N: + meta = dict(meta_) + meta["SPLIT_N"] = N // c + if meta is not None: + meta.update(**extra) + return meta + else: + # see [Computing optimal kernel parameters] in + # _triton_ops_meta.py for ways to avoid this warning + # message + warn_once( + "bsr_dense_addmm uses non-optimal triton kernel parameters" + f" for {M=} {K=} {N=} {Ms=}, {Ks=} {beta=} {alpha=} {dtype=} {out_dtype=}" + ) + + SPLIT_N = SPLIT_N or max(N // Ms, 1) + GROUP_SIZE_ROW = GROUP_SIZE_ROW or 4 + num_stages = num_stages or 1 + num_warps = num_warps or 4 + return dict( + SPLIT_N=SPLIT_N, + GROUP_SIZE_ROW=GROUP_SIZE_ROW, + num_stages=num_stages, + num_warps=num_warps, + **extra, + ) + + +class TensorAsKey: + """A light-weight wrapper of a tensor that enables storing tensors as + keys with efficient memory reference based comparison as an + approximation to data equality based keys. + + Motivation: the hash value of a torch tensor is tensor instance + based that does not use data equality and makes the usage of + tensors as keys less useful. For instance, the result of + ``len({a.crow_indices(), a.crow_indices()})`` is `2`, although, + the tensor results from `crow_indices` method call are equal, in + fact, these share the same data storage. + On the other hand, for efficient caching of tensors we want to + avoid calling torch.equal that compares tensors item-wise. + + TensorAsKey offers a compromise in that it guarantees key equality + of tensors that references data in the same storage in the same + manner and without accessing underlying data. However, this + approach does not always guarantee correctness. For instance, for + a complex tensor ``x``, we have ``TensorAsKey(x) == + TensorAsKey(x.conj())`` while ``torch.equal(x, x.conj())`` would + return False. + """ + + def __init__(self, obj): + def get_tensor_key(obj): + # Warning: TensorAsKey does not track negative nor + # conjugate bits of its input object because in the use + # case of wrapping compressed/plain indices of compressed + # sparse tensors (that are always integer tensors with + # non-negative items) these bits are never set. However, + # when extending the use of TensorAsKey to float or + # complex tensors, the values of these bits (see is_neg + # and is_conj methods) must be included in the key as + # well. + if obj.dtype.is_floating_point or obj.dtype.is_complex: + raise AssertionError( + f"TensorAsKey does not support floating point or complex dtype: {obj.dtype}" + ) + return ( + obj.data_ptr(), + obj.storage_offset(), + obj.shape, + obj.stride(), + obj.dtype, + ) + + self._obj_ref = weakref.ref(obj) + if obj.layout is torch.strided: + self.key = get_tensor_key(obj) + elif obj.layout in {torch.sparse_csr, torch.sparse_bsr}: + self.key = ( + get_tensor_key(obj.crow_indices()), + get_tensor_key(obj.col_indices()), + ) + elif obj.layout in {torch.sparse_csc, torch.sparse_bsc}: + self.key = ( + get_tensor_key(obj.ccol_indices()), + get_tensor_key(obj.row_indices()), + ) + else: + raise NotImplementedError(obj.layout) + self._hash = hash(self.key) + + def __hash__(self): + return self._hash + + def __eq__(self, other): + if not isinstance(other, TensorAsKey): + return False + if self.obj is None or other.obj is None: + # dead objects always compare unequal unless these are + # same objects + return self is other + return self.key == other.key + + @property + def obj(self): + """Return object if alive, otherwise None.""" + return self._obj_ref() + + +@lru_cache(maxsize=TORCH_SPARSE_BSR_SCATTER_MM_LRU_CACHE_SIZE) +def _bsr_scatter_mm_indices_data( + indices_format, M, K, N, Ms, Ks, nbatches, SPLIT_N, compressed_sparse_tensor_as_key +): + bsr = compressed_sparse_tensor_as_key.obj + if bsr is None: + raise AssertionError("compressed_sparse_tensor_as_key.obj is None") + crow_indices, col_indices = bsr.crow_indices(), bsr.col_indices() + device = crow_indices.device + indices_dtype = torch.int32 + + if indices_format == "bsr_strided_mm_compressed": + Ns = N // SPLIT_N + q_offsets_lst = [] + b = torch.arange(SPLIT_N, dtype=indices_dtype, device=device) * Ns + for m in range(M // Ms): + r0 = crow_indices[m].item() + r1 = crow_indices[m + 1].item() + if r1 == r0: + continue + q_offsets_lst.append( + (col_indices[r0:r1] * (Ks * N)).repeat(SPLIT_N) + + b.repeat_interleave(r1 - r0) + ) + q_offsets = torch.cat(q_offsets_lst) + crow_indices_diff = crow_indices.diff() + non_zero_row_indices = crow_indices_diff.nonzero() + a = non_zero_row_indices * (Ms * N) + r_offsets = (a + b).view(-1) + c_indices = crow_indices + # swizzle operation: mm elements with longer sums are computed first: + nnz_per_row = crow_indices_diff[non_zero_row_indices].repeat_interleave(SPLIT_N) + nnz_per_row, indices = nnz_per_row.sort(descending=True, stable=True) + r_offsets = r_offsets[indices] + return (indices_format, c_indices, r_offsets, q_offsets) + + elif indices_format == "bsr_strided_mm": + Ns = N // SPLIT_N + p_offsets_lst = [] + q_offsets_lst = [] + b = torch.arange(SPLIT_N, dtype=indices_dtype, device=device) * Ns + for m in range(M // Ms): + r0 = crow_indices[m].item() + r1 = crow_indices[m + 1].item() + if r1 == r0: + continue + p_offsets_lst.append( + torch.arange(r0, r1, dtype=indices_dtype, device=device).repeat(SPLIT_N) + ) + q_offsets_lst.append( + (col_indices[r0:r1] * (Ks * N)).repeat(SPLIT_N) + + b.repeat_interleave(r1 - r0) + ) + q_offsets = torch.cat(q_offsets_lst) + crow_indices_diff = crow_indices.diff() + non_zero_row_indices = crow_indices_diff.nonzero() + a = non_zero_row_indices * (Ms * N) + r_offsets = (a + b).view(-1) + c_indices = torch.cat( + ( + crow_indices[:1], + torch.cumsum( + crow_indices_diff[non_zero_row_indices].repeat_interleave(SPLIT_N), + 0, + ), + ) + ) + p_offsets = torch.cat(p_offsets_lst) + return (indices_format, c_indices, r_offsets, p_offsets, q_offsets) + + elif indices_format == "scatter_mm": + Ns = Ms + c_indices = [0] + pq_offsets = [] + # todo: eliminate inner for-loops for efficiency + for b in range(nbatches): + for m in range(M // Ms): + r0 = crow_indices[m].item() + r1 = crow_indices[m + 1].item() + for n in range(N // Ns): + c_indices.append(c_indices[-1] + r1 - r0) + for t in range(r1 - r0): + p = r0 + t + q = (col_indices[p].item() + b * (K // Ks)) * (N // Ns) + n + pq_offsets.append([p, q]) + + return ( + indices_format, + torch.tensor(c_indices, dtype=indices_dtype, device=device), + torch.tensor(pq_offsets, dtype=indices_dtype, device=device), + ) + + else: + raise ValueError( + f"Invalid {indices_format=}. Expected bsr_strided_mm_compressed|bsr_strided_mm|scatter_mm" + ) + + +def bsr_scatter_mm_indices_data( + bsr, other, indices_format="bsr_strided_mm_compressed", **meta_input +): + """Computes indices data for :func:`scatter_mm` used in BSR and + strided tensor matrix multiplication. + """ + if bsr.dense_dim() != 0: + raise AssertionError(f"bsr.dense_dim() must be 0, got {bsr.dense_dim()}") + if bsr.ndim != 2: + raise AssertionError(f"bsr must be 2D (no batch dims), got {bsr.ndim}D") + blocksize = bsr.values().shape[-2:] + M, K = bsr.shape + Ms, Ks = blocksize + K_, N = other.shape[-2:] + if K_ != K: + raise AssertionError(f"other K ({K_}) != bsr K ({K})") + nbatches = other.shape[:-2].numel() + + meta = scatter_mm_meta(M, K, N, Ms, Ks, **meta_input) + if "allow_tf32" not in meta_input: + meta.update(allow_tf32=bsr.dtype in {torch.float16, torch.bfloat16}) + SPLIT_N = meta["SPLIT_N"] + indices_data = _bsr_scatter_mm_indices_data( + indices_format, M, K, N, Ms, Ks, nbatches, SPLIT_N, TensorAsKey(bsr) + ) + + if indices_format == "bsr_strided_mm_compressed": + meta.update(is_compressed=True) + return indices_data + (meta,) + elif indices_format == "bsr_strided_mm": + meta.update(is_compressed=False) + return indices_data + (meta,) + else: + return indices_data + + +def bsr_scatter_mm(bsr, other, indices_data=None, out=None): + """BSR @ strided -> strided""" + + if bsr.ndim != 2: + raise AssertionError(f"bsr must be 2D, got {bsr.ndim}D") + if other.ndim < 2: + raise AssertionError( + f"other must have at least 2 dimensions, got {other.ndim}D" + ) + + Ms, Ks, Ns = bsr.shape[-2], bsr.shape[-1], other.shape[-1] + blocksize = bsr.values().shape[-2:] + + if indices_data is None: + indices_data = bsr_scatter_mm_indices_data( + bsr, other, indices_format="bsr_strided_mm_compressed" + ) + + indices_format = indices_data[0] + + if out is None: + out = torch.empty( + (*other.shape[:-2], Ms, Ns), dtype=bsr.dtype, device=bsr.device + ) + out_shape = out.shape + out = as1Dbatch(out) + + if bsr._nnz() == 0: + out.zero_() + elif indices_format in {"bsr_strided_mm_compressed", "bsr_strided_mm"}: + out.zero_() + scatter_mm(bsr.values(), other, indices_data, accumulators=out) + elif indices_format == "scatter_mm": + nbatches = other.shape[:-2].numel() + accumulators = torch.zeros( + ( + nbatches * Ms // blocksize[0] * Ns // blocksize[0], + blocksize[0], + blocksize[0], + ), + dtype=bsr.dtype, + device=bsr.device, + ) + others = ( + as1Dbatch(other) + .transpose(-2, -1) + .view( + nbatches, + Ns // blocksize[0], + blocksize[0], + Ks // blocksize[1], + blocksize[1], + ) + .movedim( + (3, 1, 4, 2), (1, 2, 3, 4) + ) # equivalent to .transpose(-3, -2).transpose(-2, -1).transpose(-4, -3) + .flatten(0, 2) + ) + scatter_mm(bsr.values(), others, indices_data, accumulators=accumulators) + out.copy_( + accumulators.unflatten( + 0, (nbatches, Ms // blocksize[0], Ns // blocksize[0]) + ) + .movedim( + (1, 2, 3, 4), (3, 1, 4, 2) + ) # equivalent to .transpose(-4, -3).transpose(-2, -1).transpose(-3, -2) + .reshape(nbatches, Ns, Ms) + .transpose(-2, -1) + ) + else: + raise NotImplementedError(indices_format) + + return out.view(out_shape) + + +def _int_bsr_dense_addmm( + input: torch.Tensor, + bsr: torch.Tensor, + dense: torch.Tensor, + *, + beta=1, + alpha=1, + left_alpha: torch.Tensor | None = None, + right_alpha: torch.Tensor | None = None, + out: torch.Tensor | None = None, + skip_checks: bool = False, + max_grid: tuple[int | None, int | None, int | None] | None = None, + meta: dict | None = None, +): + if out is None and dense.dtype is torch.int8: + f_name = "_int_bsr_dense_addmm" + crow_indices = bsr.crow_indices() + batch_ndim = crow_indices.dim() - 1 + M = bsr.shape[batch_ndim] + N = dense.shape[-1] + original_batch_dims_broadcasted = broadcast_batch_dims(f_name, bsr, dense) + out = torch.empty( + original_batch_dims_broadcasted + (M, N), + dtype=torch.int32, + device=dense.device, + ) + return bsr_dense_addmm( + input, + bsr, + dense, + beta=beta, + alpha=alpha, + left_alpha=left_alpha, + right_alpha=right_alpha, + out=out, + skip_checks=skip_checks, + max_grid=max_grid, + meta=meta, + ) + + +def bsr_dense_addmm( + input: torch.Tensor, + bsr: torch.Tensor, + dense: torch.Tensor, + *, + beta=1, + alpha=1, + left_alpha: torch.Tensor | None = None, + right_alpha: torch.Tensor | None = None, + out: torch.Tensor | None = None, + skip_checks: bool = False, + max_grid: tuple[int | None, int | None, int | None] | None = None, + meta: dict | None = None, +): + """Compute + + out = beta * input + left_alpha.reshape(-1, 1) * (alpha * (bsr @ dense)) * right_alpha.reshape(1, -1) + + where left_alpha, right_alpha are (* + 1)-D tensors when + specified, otherwise, these are treated as tensors filled with + ones. + """ + f_name = "bsr_dense_addmm" + values = bsr.values() + crow_indices = bsr.crow_indices() + col_indices = bsr.col_indices() + batch_ndim = crow_indices.dim() - 1 + M, K = bsr.shape[batch_ndim : batch_ndim + 2] + blocksize = values.shape[batch_ndim + 1 : batch_ndim + 3] + N = dense.shape[-1] + + # todo: implement checks + + original_batch_dims_broadcasted = broadcast_batch_dims(f_name, bsr, dense) + if out is None: + out = dense.new_empty(original_batch_dims_broadcasted + (M, N)) + + if bsr._nnz() == 0 or alpha == 0 or N == 0 or M == 0 or K == 0: + if beta == 0: + out.zero_() + else: + out.copy_(input) + if beta != 1: + out.mul_(beta) + return out + + left_alpha_is_one = False + right_alpha_is_one = False + if left_alpha is None: + left_alpha_is_one = True + left_alpha = dense.new_empty(()).expand( + *original_batch_dims_broadcasted, M, N + ) # not referenced + else: + left_alpha = left_alpha.view(*original_batch_dims_broadcasted, M, 1).expand( + *original_batch_dims_broadcasted, M, N + ) + + if right_alpha is None: + right_alpha_is_one = True + right_alpha = dense.new_empty(()).expand( + *original_batch_dims_broadcasted, M, N + ) # not referenced + else: + right_alpha = right_alpha.view(*original_batch_dims_broadcasted, 1, N).expand( + *original_batch_dims_broadcasted, M, N + ) + if left_alpha.stride()[-1] != 0: + raise AssertionError( + f"left_alpha.stride()[-1] must be 0, got {left_alpha.stride()[-1]}" + ) + if right_alpha.stride()[-2] != 0: + raise AssertionError( + f"right_alpha.stride()[-2] must be 0, got {right_alpha.stride()[-2]}" + ) + + if meta is None: + sparsity = round(1 - bsr._nnz() * blocksize[0] * blocksize[1] / (M * K), 2) + meta = bsr_dense_addmm_meta( + M, + K, + N, + blocksize[0], + blocksize[1], + beta, + alpha, + sparsity=sparsity, + dtype=dense.dtype, + out_dtype=out.dtype, + ) + out_backup = out + + ( + crow_indices, + col_indices, + values, + input, + dense, + left_alpha, + right_alpha, + out, + ) = prepare_inputs(bsr, input, dense, left_alpha, right_alpha, out) + + BM, BK = blocksize + SPLIT_N = meta.get("SPLIT_N", N // BM) + BN = N // SPLIT_N + + out_untiled = out + out = tile_to_blocksize(out, (BM, BN)) + dense = tile_to_blocksize(dense, (BK, BN)) + input = tile_to_blocksize(input, (BM, BN)) + left_alpha = tile_to_blocksize(left_alpha, (BM, BN)) + right_alpha = tile_to_blocksize(right_alpha, (BM, BN)) + + # tl.dot supports float16, float32, int32 as accumulator types. + dot_out_dtype = { + torch.float16: tl.float32, + torch.bfloat16: tl.float32, + torch.float32: tl.float64, + torch.float64: tl.float64, + torch.int8: tl.int32, + torch.int32: tl.int32, + }[out.dtype] + + n_batches = dense.size(0) + n_block_rows = crow_indices.size(-1) - 1 + n_block_cols = dense.size(-3) + + full_grid = (n_batches, n_block_cols, n_block_rows) + if max_grid is not None: + grid_blocks = tuple(max_grid[:3][::-1]) + (None,) * (3 - len(max_grid[:3])) + else: + grid_blocks = None + + tensor_dims_map = { + values: (0, None, None), + crow_indices: (0, None, -1), + col_indices: (0, None, None), + input: (0, -3, -4), + dense: (0, -3, None), + left_alpha: (0, -3, -4), + right_alpha: (0, -3, -4), + out: (0, -3, -4), + } + + if alpha == 0: + raise AssertionError("alpha must not be 0") + + def kernel(grid, *sliced_tensors): + # pyrefly: ignore [unsupported-operation] + _bsr_strided_addmm_kernel[grid]( + *ptr_stride_extractor(*sliced_tensors), + beta, + alpha, + beta_is_one=beta == 1, + beta_is_nonzero=beta != 0, + alpha_is_one=alpha == 1, + left_alpha_is_one=left_alpha_is_one, + right_alpha_is_one=right_alpha_is_one, + BLOCKSIZE_ROW=BM, + BLOCKSIZE_INNER=BK, + BLOCKSIZE_COL=BN, + allow_tf32=dot_out_dtype == tl.float32, + acc_dtype=dot_out_dtype, + **meta, + ) + + launch_kernel(kernel, tensor_dims_map, full_grid, grid_blocks) + + if out.data_ptr() != out_backup.data_ptr(): + # prepare_inputs has made a copy of out, copy its content back + # to out_backup: + out_backup.copy_(out_untiled.view(out_backup.shape)) + + return out_backup + + +if has_triton(): + import triton + import triton.language as tl + + @triton.jit + def _sampled_addmm_kernel( + alpha, + beta, + IS_BETA_ZERO: tl.constexpr, + BLOCKSIZE_ROW: tl.constexpr, + BLOCKSIZE_COL: tl.constexpr, + k, + TILE_K: tl.constexpr, + values_ptr, + values_batch_stride, + values_nnz_stride, + values_row_block_stride, + values_col_block_stride, + crow_indices_ptr, + crow_indices_batch_stride, + crow_indices_stride, + col_indices_ptr, + col_indices_batch_stride, + col_indices_stride, + mat1_ptr, + mat1_batch_stride, + mat1_tiled_row_stride, + mat1_tiled_col_stride, + mat1_row_block_stride, + mat1_col_block_stride, + mat2_ptr, + mat2_batch_stride, + mat2_tiled_row_stride, + mat2_tiled_col_stride, + mat2_row_block_stride, + mat2_col_block_stride, + acc_dtype: tl.constexpr, + allow_tf32: tl.constexpr, + ): + batch_pid = tl.program_id(axis=1) + row_block_pid = tl.program_id(axis=0) + + crow_indices_offset_ptr = ( + crow_indices_ptr + + crow_indices_batch_stride * batch_pid + + crow_indices_stride * row_block_pid + ) + nnz_offset = tl.load(crow_indices_offset_ptr) + nnz_offset_next = tl.load(crow_indices_offset_ptr + crow_indices_stride) + + # Compute nnz for the row with number row_block_pid. + # If it is zero, skip the row. + row_nnz = nnz_offset_next - nnz_offset + if row_nnz == 0: + return + + row_block_arange = tl.arange(0, BLOCKSIZE_ROW) + col_block_arange = tl.arange(0, BLOCKSIZE_COL) + + # Pointers are set to the first block of the current row. + values_block_ptrs = ( + values_ptr + + values_batch_stride * batch_pid + + values_nnz_stride * nnz_offset + + values_row_block_stride * row_block_arange[:, None] + + values_col_block_stride * col_block_arange[None, :] + ) + + col_index_nnz_ptr = ( + col_indices_ptr + + col_indices_batch_stride * batch_pid + + col_indices_stride * nnz_offset + ) + + # Advance mat1 to the current tiled row, ignore columns. + mat1_block_ptrs = ( + mat1_ptr + + mat1_batch_stride * batch_pid + + mat1_tiled_row_stride * row_block_pid + + mat1_row_block_stride * row_block_arange[:, None] + ) + + # Advance mat2 in batch and block col dimension. + mat2_block_ptrs = ( + mat2_ptr + + mat2_batch_stride * batch_pid + + mat2_col_block_stride * col_block_arange[None, :] + ) + + k_tile_arange = tl.arange(0, TILE_K) + for _ in range(row_nnz): + acc_block = tl.zeros((BLOCKSIZE_ROW, BLOCKSIZE_COL), dtype=acc_dtype) + + # find column block index + col_block = tl.load(col_index_nnz_ptr) + + for k_tile in range(0, k, TILE_K): + k_offsets = k_tile + k_tile_arange + mask_k = k_offsets < k + + mat1_block = tl.load( + mat1_block_ptrs + mat1_col_block_stride * k_offsets[None, :], + # pyrefly: ignore [bad-index] + mask=mask_k[None, :], + other=0.0, + ) + + mat2_block = tl.load( + mat2_block_ptrs + + mat2_tiled_col_stride * col_block + + mat2_row_block_stride * k_offsets[:, None], + # pyrefly: ignore [bad-index] + mask=mask_k[:, None], + other=0.0, + ) + + acc_block += tl.dot( + mat1_block, mat2_block, allow_tf32=allow_tf32, out_dtype=acc_dtype + ) + + if IS_BETA_ZERO: + acc_block *= alpha + else: + acc_block = alpha * acc_block + beta * tl.load(values_block_ptrs) + + # write result + tl.store(values_block_ptrs, acc_block.to(values_ptr.dtype.element_ty)) + + # advance val/col_index ptrs to the next block in the row. + values_block_ptrs += values_nnz_stride + col_index_nnz_ptr += col_indices_stride + + @triton.jit + def _bsr_strided_dense_rowspace_kernel( + # values prologue + values_ptr, + values_batch_stride, + values_nnz_stride, + values_row_block_stride, + values_col_block_stride, + # values epilogue + # crow_indices prologue + crow_indices_ptr, + crow_indices_batch_stride, + crow_indices_stride, + # crow_indices epilogue + # col_indices prologue + col_indices_ptr, + col_indices_batch_stride, + col_indices_stride, + # col_indices epilogue + # dense prologue + dense_ptr, + dense_batch_stride, + dense_tiled_row_stride, + dense_tiled_col_stride, + dense_row_block_stride, + dense_col_block_stride, + # dense epilogue + # output prologue + output_ptr, + output_batch_stride, + output_tiled_row_stride, + output_tiled_col_stride, + output_row_block_stride, + output_col_block_stride, + # output epilogue + # + # gh-113754: Always keep all constexpr arguments at the end of + # triton kernel arguments list because with triton 2.1 or + # earlier non-contiguous outputs will corrupt CUDA state due + # to a triton bug (fixed in openai/triton#2262). + BLOCKSIZE_ROW: tl.constexpr, + BLOCKSIZE_COL: tl.constexpr, + acc_dtype: tl.constexpr, + allow_tf32: tl.constexpr, + GROUP_SIZE_ROW: tl.constexpr, + ): + batch_pid = tl.program_id(axis=2) + row_block_pid = tl.program_id(axis=0) + col_block_pid = tl.program_id(axis=1) + n_block_rows = tl.num_programs(axis=0) + n_block_cols = tl.num_programs(axis=1) + + row_block_pid, col_block_pid = tl.swizzle2d( + row_block_pid, col_block_pid, n_block_rows, n_block_cols, GROUP_SIZE_ROW + ) + + crow_indices_offset_ptr = ( + crow_indices_ptr + + crow_indices_batch_stride * batch_pid + + crow_indices_stride * row_block_pid + ) + nnz_offset = tl.load(crow_indices_offset_ptr) + nnz_offset_next = tl.load(crow_indices_offset_ptr + crow_indices_stride) + + # Compute nnz for the row with number row_block_pid. + # If it is zero, skip the row. + row_nnz = nnz_offset_next - nnz_offset + if row_nnz == 0: + return + + row_block_arange = tl.arange(0, BLOCKSIZE_ROW) + col_block_arange = tl.arange(0, BLOCKSIZE_COL) + + # Pointers are set to the first block of the current row. + values_block_ptrs = ( + values_ptr + + values_batch_stride * batch_pid + + values_nnz_stride * nnz_offset + + values_row_block_stride * row_block_arange[:, None] + + values_col_block_stride * col_block_arange[None, :] + ) + + # NOTE: dense is advanced into all dimensions but the tiled row one. + # That will be advanced in the loop according to values in col_indices. + dense_block_ptrs = ( + dense_ptr + + dense_batch_stride * batch_pid + + dense_tiled_col_stride * col_block_pid + + dense_row_block_stride * col_block_arange[:, None] + + dense_col_block_stride * row_block_arange[None, :] + ) + + # Pointers are set to exact write-to locations + output_ptrs = ( + output_ptr + + output_batch_stride * batch_pid + + output_tiled_row_stride * row_block_pid + + output_tiled_col_stride * col_block_pid + + output_row_block_stride * row_block_arange[:, None] + + output_col_block_stride * row_block_arange[None, :] + ) + + # Set pointer to the first nonzero element in the current row + col_index_nnz_ptr = ( + col_indices_ptr + + col_indices_batch_stride * batch_pid + + col_indices_stride * nnz_offset + ) + + output_acc_block = tl.zeros((BLOCKSIZE_ROW, BLOCKSIZE_COL), dtype=acc_dtype) + for _ in range(row_nnz): + values_block = tl.load(values_block_ptrs) + + # find which row of dense needs to get loaded + # for multiplication with values_block. + dense_row_idx = tl.load(col_index_nnz_ptr) + dense_block = tl.load( + dense_block_ptrs + dense_tiled_row_stride * dense_row_idx + ) + + # do block mm + output_acc_block += tl.dot( + values_block, dense_block, allow_tf32=allow_tf32, out_dtype=acc_dtype + ) + + # move val/col_index ptrs to the next block in the row + values_block_ptrs += values_nnz_stride + col_index_nnz_ptr += col_indices_stride + + # write back the result + tl.store(output_ptrs, output_acc_block.to(output_ptr.dtype.element_ty)) + + def _run_sampled_addmm_kernel( + alpha, + beta, + is_beta_zero, + blocksize, + k, + tile_k, + values, + crow_indices, + col_indices, + mat1, + mat2, + max_grid, + ): + n_batches = values.size(0) + n_block_rows = crow_indices.size(-1) - 1 + + full_grid = (n_batches, n_block_rows) + if max_grid is not None: + grid_blocks = tuple(max_grid[:2][::-1]) + (None,) * (2 - len(max_grid[:2])) + else: + grid_blocks = None + tensor_dims_map = { + values: (0, None), + crow_indices: (0, -1), + col_indices: (0, None), + mat1: (0, -4), + mat2: (0, None), + } + if values.dtype in (torch.half, torch.bfloat16): + acc_dtype = tl.float32 + allow_tf32 = True + else: + acc_dtype = tl.float64 + allow_tf32 = False + + def kernel(grid, *sliced_tensors): + _sampled_addmm_kernel[grid]( + alpha, + beta, + is_beta_zero, + *blocksize, + k, + tile_k, + *ptr_stride_extractor(*sliced_tensors), + acc_dtype=acc_dtype, + allow_tf32=allow_tf32, + num_stages=1, + num_warps=4, + ) + + launch_kernel(kernel, tensor_dims_map, full_grid, grid_blocks) + + def sampled_addmm( + input: torch.Tensor, + mat1: torch.Tensor, + mat2: torch.Tensor, + *, + beta=1.0, + alpha=1.0, + out: torch.Tensor | None = None, + skip_checks: bool = False, + max_grid: tuple[int | None, int | None, int | None] | None = None, + ): + f_name = "sampled_addmm" + + check_bsr_layout(f_name, input) + input_broadcasted = broadcast_batch_dims_bsr(f_name, input, mat1, mat2) + + if not skip_checks: + check_device(f_name, mat1, input.device) + check_device(f_name, mat2, input.device) + if beta != 0.0 and input.dtype is torch.bool: + check( + False, + f"{f_name}(): having beta == {beta} not equal to 0.0 with boolean mask is not allowed.", + ) + if input.dtype is not torch.bool: + check_dtype(f_name, mat1, input.dtype) + check_dtype(f_name, mat2, input.dtype) + else: + check_dtype(f_name, mat1, mat2.dtype) + check_mm_compatible_shapes(f_name, mat1, mat2) + if out is not None: + check_bsr_layout(f_name, out) + check_device(f_name, out, mat1.device) + check_dtype(f_name, out, input.dtype) + check( + out.shape == input_broadcasted.shape and out._nnz() == input._nnz(), + f"{f_name}(): Expects `out` to be of shape {input_broadcasted.shape} " + f"and with nnz equal to {input_broadcasted._nnz()} " + f"but got out.shape = {out.shape} and out.nnz = {out._nnz()}", + ) + + if out is None: + out = input_broadcasted.to(mat1.dtype, copy=True) + else: + out.copy_(input_broadcasted) + + if out.numel() == 0 or out._nnz() == 0: + return out + + blocksize = out.values().shape[-2:] + k = mat1.size(-1) + + # NOTE: (m, 0) @ (0, n) == zeros(m, n) + if alpha == 0.0 or k == 0: + out.values().mul_(beta) + return out + + # prepare inputs by reshaping them to be kernel-compatible + out_backup = out + crow_indices, col_indices, values, mat1, mat2 = prepare_inputs(out, mat1, mat2) + + mat1 = tile_to_blocksize(mat1, (blocksize[0], k)) + mat2 = tile_to_blocksize(mat2, (k, blocksize[1])) + tile_k = max(*blocksize) + + _run_sampled_addmm_kernel( + alpha, + beta, + beta == 0.0, + blocksize, + k, + tile_k, + values, + crow_indices, + col_indices, + mat1, + mat2, + max_grid, + ) + + # If nnz x block strides are not the same in out_backup.values and values, + # it means that out_backup.values and values are not the views of each other, + # so we have to copy. + if out_backup.values().stride()[-3:] != values.stride()[-3:]: + out_backup.values().copy_(values.reshape(out_backup.values().shape)) + return out_backup + + def bsr_dense_mm( + bsr: torch.Tensor, + dense: torch.Tensor, + *, + out: torch.Tensor | None = None, + skip_checks: bool = False, + max_grid: tuple[int | None, int | None, int | None] | None = None, + meta: dict | None = None, + ): + f_name = "bsr_dense_mm" + m, _kl = bsr.shape[-2:] + if not skip_checks: + check_bsr_layout(f_name, bsr) + check_device(f_name, bsr, dense.device) + check_dtype(f_name, bsr, dense.dtype, (torch.int8,)) + check_mm_compatible_shapes(f_name, bsr, dense) + + n = dense.size(-1) + row_block, col_block = bsr.values().shape[-2:] + check_blocksize(f_name, (row_block, col_block)) + check( + not n % 16, + f"{f_name}(): dense.size(-1) == {n} should be divisible by 16", + ) + else: + _kr, n = dense.shape[-2:] + + original_batch_dims_broadcasted = broadcast_batch_dims(f_name, bsr, dense) + + if out is not None and not skip_checks: + expected_out_shape = original_batch_dims_broadcasted + (m, n) + check( + out.shape == expected_out_shape, + "bsr_dense_mm(): `out` argument has wrong shape, " + f"expected {expected_out_shape}, but got {out.shape}.", + ) + check( + out.is_contiguous() or out.transpose(-2, -1).is_contiguous(), + "bsr_dense_mm(): only row-major/col-major `out` arguments are supported, " + "i.e. (out.is_contiguous() or out.transpose(-2, -1).is_contiguous()) " + "should be True.", + ) + + # Allocate out + if out is None: + out = dense.new_empty(original_batch_dims_broadcasted + (m, n)) + + # Short circuit if lhs is zero + if bsr._nnz() == 0: + return out.zero_() + + # with beta==0, addmm ignores input content, so we can use out + # as a placeholder for input because their shapes match: + return bsr_dense_addmm(out, bsr, dense, alpha=1, beta=0, out=out) + + @triton.jit + def _bsr_softmax_kernel( + crow_indices_ptr, + crow_indices_batch_stride, + crow_indices_stride, + values_ptr, + values_batch_stride, + values_row_block_stride, + values_nnz_col_block_stride, + row_block, + col_block, + MAX_ROW_NNZ: tl.constexpr, + TILE: tl.constexpr, + ): + batch_pid = tl.program_id(axis=2) + row_block_offset_pid = tl.program_id(axis=1) + row_block_pid = tl.program_id(axis=0) + + crow_indices_offset_ptr = ( + crow_indices_ptr + + crow_indices_batch_stride * batch_pid + + crow_indices_stride * row_block_pid + ) + nnz_offset = tl.load(crow_indices_offset_ptr) + nnz_offset_next = tl.load(crow_indices_offset_ptr + crow_indices_stride) + + # Compute nnz for the row with number row_block_pid. + # If it is zero, skip the row. + row_nnz = nnz_offset_next - nnz_offset + if row_nnz == 0: + return + + row_arange = tl.arange(0, TILE) + mask = row_arange < row_nnz * col_block + + curr_row_values_ptrs = ( + values_ptr + + values_batch_stride * batch_pid + + values_row_block_stride * row_block_offset_pid + + nnz_offset * col_block + ) + + # find max in the row + row_tile = tl.load( + curr_row_values_ptrs + row_arange, mask=mask, other=-float("inf") + ).to(tl.float32) + max_row_value = tl.max(row_tile, axis=0) + for _ in range(TILE, MAX_ROW_NNZ, TILE): + row_arange += TILE + mask = row_arange < row_nnz * col_block + row_tile = tl.load( + curr_row_values_ptrs + row_arange, mask=mask, other=-float("inf") + ).to(tl.float32) + curr_max_row_value = tl.max(row_tile, axis=0) + max_row_value = tl.where( + max_row_value > curr_max_row_value, max_row_value, curr_max_row_value + ) + + # find denominator for stable softmax + num = tl.exp(row_tile - max_row_value) + denom = tl.sum(num, axis=0) + for _ in range(TILE, MAX_ROW_NNZ, TILE): + row_arange -= TILE + mask = row_arange < row_nnz * col_block + row_tile = tl.load( + curr_row_values_ptrs + row_arange, mask=mask, other=-float("inf") + ).to(tl.float32) + num = tl.exp(row_tile - max_row_value) + denom += tl.sum(num, axis=0) + + # populate output + tl.store( + curr_row_values_ptrs + row_arange, + (num / denom).to(values_ptr.dtype.element_ty), + mask=mask, + ) + for _ in range(TILE, MAX_ROW_NNZ, TILE): + row_arange += TILE + mask = row_arange < row_nnz * col_block + row_tile = tl.load( + curr_row_values_ptrs + row_arange, mask=mask, other=-float("inf") + ).to(tl.float32) + num = tl.exp(row_tile - max_row_value) + tl.store( + curr_row_values_ptrs + row_arange, + (num / denom).to(values_ptr.dtype.element_ty), + mask=mask, + ) + + def bsr_softmax(input, max_row_nnz=None): + f_name = "bsr_softmax" + + check_bsr_layout(f_name, input) + check_dtype(f_name, input, input.dtype) + + if input._nnz() == 0 or input.numel() == 0: + return input.clone() + + m, n = input.shape[-2:] + nnz = input._nnz() + row_block, col_block = input.values().shape[-2:] + + if max_row_nnz is None: + max_row_nnz = triton.next_power_of_2(n) + else: + max_row_nnz = triton.next_power_of_2(max_row_nnz) + + crow_indices = input.crow_indices().unsqueeze(0).flatten(0, -2) + # reshape values from + # (b1, ..., bn, nnz, row_block, col_block) to + # (b1 * ... * bn, row_block, nnz * col_block). + # This simplifies batch dim manipulation and unlocks + # the possibility to access all nnzs in any given row. + if input.values().transpose(-3, -2).is_contiguous(): + # Need to clone to avoid `contiguous` returning a view. + values = input.values().clone() + else: + values = input.values() + values = ( + values.transpose(-3, -2) + .contiguous() + .unsqueeze(0) + .flatten(0, -4) + .reshape(-1, row_block, nnz * col_block) + ) + full_grid = (values.shape[0], row_block, m // row_block) + grid_blocks = None + tensor_dims_map = { + # We span nnz number of blocks, not nnz + 1, + # hence crow_indices[..., :-1] + crow_indices[..., :-1]: (0, None, -1), + values: (0, None, None), + } + + def kernel(grid, *sliced_tensors): + _bsr_softmax_kernel[grid]( + *ptr_stride_extractor(*sliced_tensors), + row_block, + col_block, + max_row_nnz, + # Triton's max numel is bounded by 2 ** 17. + min(2**17, max_row_nnz), + ) + + launch_kernel(kernel, tensor_dims_map, full_grid, grid_blocks) + + values = ( + values.reshape(-1, row_block, nnz, col_block) + .transpose(-3, -2) + .reshape(*input.values().shape) + ) + + return torch.sparse_compressed_tensor( + input.crow_indices().clone(), + input.col_indices().clone(), + values, + size=input.shape, + layout=input.layout, + ) + + def _scaled_dot_product_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor | None, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float | None = None, + ): + f_name = "_scaled_dot_product_attention" + check(not is_causal, f"{f_name}(): is_causal == True is not supported.") + check(attn_mask is not None, f"{f_name}(): attn_mask == None is not supported.") + if attn_mask is None: + raise AssertionError("attn_mask must not be None") + + check( + attn_mask.layout == torch.sparse_bsr, + f"{f_name}(): " + f"attn_mask.layout must be {torch.sparse_bsr}, but got " + f"attn_mask.layout == {attn_mask.layout}.", + ) + + check_device(f_name, key, query.device) + check_device(f_name, value, query.device) + check_device(f_name, attn_mask, query.device) + + check_dtype(f_name, key, query.dtype) + check_dtype(f_name, value, query.dtype) + if attn_mask.dtype is not torch.bool: + check_dtype(f_name, attn_mask, query.dtype) + + # pyrefly: ignore [not-callable] + sdpa = sampled_addmm( + attn_mask, query, key.transpose(-2, -1), beta=0.0, skip_checks=False + ) + if scale is None and query.size(-1) == 0 or scale == 0.0: + check( + False, + f"{f_name}(): current value of scale == {scale} " + "results in division by zero.", + ) + scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale + sdpa.values().mul_(scale_factor) + # pyrefly: ignore [not-callable] + sdpa = bsr_softmax(sdpa) + torch.nn.functional.dropout(sdpa.values(), p=dropout_p, inplace=True) + # pyrefly: ignore [not-callable] + sdpa = bsr_dense_mm(sdpa, value) + return sdpa + + @triton.jit + def _scatter_mm2_kernel( + M: tl.constexpr, + K: tl.constexpr, + N: tl.constexpr, + blocks_ptr, + blocks_stride_P, + blocks_stride_M, + blocks_stride_K, + others_ptr, + others_stride_Q, + others_stride_K, + others_stride_N, + accumulators_ptr, + accumulators_stride_R, + accumulators_stride_M, + accumulators_stride_N, + pq_offsets_ptr, + pq_offsets_stride, + pq_ptr, + pq_stride_T, + pq_stride_1, + dot_out_dtype: tl.constexpr, + TILE_M: tl.constexpr, + TILE_N: tl.constexpr, + allow_tf32: tl.constexpr, + ): + Ms = M // TILE_M + + pid_t = tl.program_id(axis=0) + + pid = tl.program_id(axis=1) + pid_m = pid // Ms + pid_n = pid % Ms + + rm = pid_m * TILE_M + tl.arange(0, TILE_M) + rn = pid_n * TILE_N + tl.arange(0, TILE_N) + rk = tl.arange(0, K) + + A_ptr = blocks_ptr + ( + rm[:, None] * blocks_stride_M + rk[None, :] * blocks_stride_K + ) + B_ptr = others_ptr + ( + rk[:, None] * others_stride_K + rn[None, :] * others_stride_N + ) + + g0 = tl.load(pq_offsets_ptr + pid_t * pq_offsets_stride) + g1 = tl.load(pq_offsets_ptr + (pid_t + 1) * pq_offsets_stride) + + if g0 == g1: + return + + acc_block = tl.zeros((TILE_M, TILE_N), dtype=dot_out_dtype) + + for i in range(g0, g1): + p = tl.load(pq_ptr + i * pq_stride_T) + q = tl.load(pq_ptr + i * pq_stride_T + pq_stride_1) + A = tl.load(A_ptr + p * blocks_stride_P) + B = tl.load(B_ptr + q * others_stride_Q) + acc_block += tl.dot(A, B, out_dtype=dot_out_dtype, allow_tf32=allow_tf32) + + C_ptr = ( + accumulators_ptr + + pid_t * accumulators_stride_R + + ( + rm[:, None] * accumulators_stride_M + + rn[None, :] * accumulators_stride_N + ) + ) + tl.store(C_ptr, acc_block.to(accumulators_ptr.dtype.element_ty)) + + def _scatter_mm2( + blocks: torch.Tensor, + others: torch.Tensor, + pq_offsets: torch.Tensor, + pq_indices: torch.Tensor, + accumulators: torch.Tensor, + ): + _P, M, K = blocks.shape + _Q, _, N = others.shape + + meta = dict( + TILE_M=max(16, M // 4), TILE_N=max(16, N // 4), num_stages=1, num_warps=2 + ) + + def grid(META): + return ( + pq_offsets.shape[0] - 1, + triton.cdiv(M, META["TILE_M"]) * triton.cdiv(N, META["TILE_N"]), + 1, + ) + + dot_out_dtype = { + torch.float16: tl.float32, + torch.bfloat16: tl.float32, + torch.float32: tl.float64, + torch.float64: tl.float64, + }[accumulators.dtype] + if "allow_tf32" not in meta: + meta.update(allow_tf32=dot_out_dtype == tl.float32) + _scatter_mm2_kernel[grid]( + M, + K, + N, + blocks, + blocks.stride(0), + blocks.stride(1), + blocks.stride(2), + others, + others.stride(0), + others.stride(1), + others.stride(2), + accumulators, + accumulators.stride(0), + accumulators.stride(1), + accumulators.stride(2), + pq_offsets, + pq_offsets.stride(0), + pq_indices, + pq_indices.stride(0), + pq_indices.stride(1), + dot_out_dtype=dot_out_dtype, + **meta, + ) + + @triton.jit + def _scatter_mm6_kernel( + nbatches, + Ms, + Ks: tl.constexpr, + N, + blocks_ptr, + blocks_stride_P, + blocks_stride_M, + blocks_stride_K, + others_ptr, + others_stride_B, + others_stride_K, + others_stride_N, + accumulators_ptr, + accumulators_stride_B, + accumulators_stride_M, + accumulators_stride_N, + c_indices_ptr, + r_offsets_ptr, + p_offsets_ptr, + q_offsets_ptr, + is_compressed: tl.constexpr, + dot_out_dtype: tl.constexpr, + SPLIT_N: tl.constexpr, + TILE_M: tl.constexpr, + TILE_N: tl.constexpr, + GROUP_SIZE: tl.constexpr, + allow_tf32: tl.constexpr, + ): + Ns = N // SPLIT_N + BLOCKS_M = Ms // TILE_M + BLOCKS_N = Ns // TILE_N + + pid_t_ = tl.program_id(axis=0) + pid = tl.program_id(axis=1) + pid_b = pid_t_ % nbatches + pid_t = pid_t_ // nbatches + + num_pid_in_group = GROUP_SIZE * BLOCKS_N + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE + group_size_m = min(BLOCKS_M - first_pid_m, GROUP_SIZE) + pid_m = first_pid_m + (pid % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + rm = pid_m * TILE_M + tl.arange(0, TILE_M) + rn = pid_n * TILE_N + tl.arange(0, TILE_N) + rk = tl.arange(0, Ks) + A_ptr = blocks_ptr + ( + rm[:, None] * blocks_stride_M + rk[None, :] * blocks_stride_K + ) + B_ptr = ( + others_ptr + + pid_b * others_stride_B + + (rk[:, None] * others_stride_K + rn[None, :] * others_stride_N) + ) + + # When is_compressed is True, r is the only variable that + # depends on pid_t. This property allows sorting r values + # before calling the kernel. The sorting of r is equivalent to + # defining swizzle operator outside of the kernel. + r = tl.load(r_offsets_ptr + pid_t) + + if is_compressed: + m = (r // N) // Ms + n = (r % N) // Ns + r0 = tl.load(c_indices_ptr + m) + r1 = tl.load(c_indices_ptr + m + 1) + g0 = n * r1 + (SPLIT_N - n) * r0 + nnz = r1 - r0 + else: + g0 = tl.load(c_indices_ptr + pid_t) + g1 = tl.load(c_indices_ptr + pid_t + 1) + nnz = g1 - g0 + + q_ptr = q_offsets_ptr + g0 + acc_block = tl.zeros((TILE_M, TILE_N), dtype=dot_out_dtype) + + if is_compressed: + A_ptr += r0 * blocks_stride_P # type: ignore[possibly-undefined] + for _ in range(nnz): + q = tl.load(q_ptr) + B = tl.load(B_ptr + q) + A = tl.load(A_ptr) + acc_block += tl.dot( + A, B, out_dtype=dot_out_dtype, allow_tf32=allow_tf32 + ) + A_ptr += blocks_stride_P + q_ptr += 1 + else: + p_ptr = p_offsets_ptr + g0 + for _ in range(nnz): + q = tl.load(q_ptr) + B = tl.load(B_ptr + q) + p = tl.load(p_ptr) + A = tl.load(A_ptr + p * blocks_stride_P) + p_ptr += 1 + q_ptr += 1 + acc_block += tl.dot( + A, B, out_dtype=dot_out_dtype, allow_tf32=allow_tf32 + ) + + C_ptr = ( + accumulators_ptr + + r + + pid_b * accumulators_stride_B + + ( + rm[:, None] * accumulators_stride_M + + rn[None, :] * accumulators_stride_N + ) + ) + tl.store(C_ptr, acc_block.to(accumulators_ptr.dtype.element_ty)) + + def _scatter_mm6( + blocks: torch.Tensor, + others: torch.Tensor, + c_indices: torch.Tensor, + r_offsets: torch.Tensor, + p_offsets: torch.Tensor, + q_offsets: torch.Tensor, + meta: dict, + accumulators: torch.Tensor, + force_contiguous: bool = True, + ): + SPLIT_N = meta["SPLIT_N"] + _P, Ms, Ks = blocks.shape + B, _K, N = others.shape + B_, _M, N_ = accumulators.shape + if N_ != N: + raise AssertionError(f"accumulators N ({N_}) != others N ({N})") + Ns = N // SPLIT_N + if B_ != B: + raise AssertionError(f"accumulators B ({B_}) != others B ({B})") + + def grid(META): + return ( + r_offsets.shape[0] * B, + triton.cdiv(Ms, META["TILE_M"]) * triton.cdiv(Ns, META["TILE_N"]), + ) + + dot_out_dtype = { + torch.float16: tl.float32, + torch.bfloat16: tl.float32, + torch.float32: tl.float64, + torch.float64: tl.float64, + }[accumulators.dtype] + if "allow_tf32" not in meta: + meta.update(allow_tf32=dot_out_dtype == tl.float32) + + if c_indices.stride(0) != 1: + raise AssertionError( + f"c_indices.stride(0) must be 1, got {c_indices.stride(0)}" + ) + if r_offsets.stride(0) != 1: + raise AssertionError( + f"r_offsets.stride(0) must be 1, got {r_offsets.stride(0)}" + ) + if p_offsets.stride(0) != 1: + raise AssertionError( + f"p_offsets.stride(0) must be 1, got {p_offsets.stride(0)}" + ) + if q_offsets.stride(0) != 1: + raise AssertionError( + f"q_offsets.stride(0) must be 1, got {q_offsets.stride(0)}" + ) + + # Re non-contiguous tensor arguments. Sometimes triton kernel + # launches may fail with + # + # RuntimeError: Triton Error [CUDA]: an illegal memory access was encountered + # + # that appears to be case when the size of a non-contiguous + # tensor argument is larger than a certain threshold. Could + # this be related to shared memory or L1 cache size of a GPU + # card? In anycase, ensuring that tensor arguments are + # contiguous seems to avoid the above exception. So, in the + # following we'll always convert tensor arguments to + # C-contiguous tensors. + + if force_contiguous: + blocks = blocks.contiguous() + others = others.contiguous() + if not accumulators.is_contiguous(): + accumulators_ = accumulators.contiguous() + else: + accumulators_ = accumulators + else: + accumulators_ = accumulators + + _scatter_mm6_kernel[grid]( + B, + Ms, + Ks, + N, + blocks, + blocks.stride(0), + blocks.stride(1), + blocks.stride(2), + others, + others.stride(0), + others.stride(1), + others.stride(2), + accumulators_, + accumulators_.stride(0), + accumulators_.stride(1), + accumulators_.stride(2), + c_indices, + r_offsets, + p_offsets, + q_offsets, + dot_out_dtype=dot_out_dtype, + **meta, + ) + + if force_contiguous and not accumulators.is_contiguous(): + accumulators.copy_(accumulators_) + + @triton.jit + def _bsr_strided_addmm_kernel( + # values prologue + values_ptr, + values_batch_stride, + values_nnz_stride, + values_row_block_stride, + values_col_block_stride, + # values epilogue + # crow_indices prologue + crow_indices_ptr, + crow_indices_batch_stride, + crow_indices_stride, + # crow_indices epilogue + # col_indices prologue + col_indices_ptr, + col_indices_batch_stride, + col_indices_stride, + # col_indices epilogue + # input prologue + input_ptr, + input_batch_stride, + input_tiled_row_stride, + input_tiled_col_stride, + input_row_block_stride, + input_col_block_stride, + # input epilogue + # dense prologue + dense_ptr, + dense_batch_stride, + dense_tiled_row_stride, + dense_tiled_col_stride, + dense_row_block_stride, + dense_col_block_stride, + # dense epilogue + # left_alpha prologue + left_alpha_ptr, + left_alpha_batch_stride, + left_alpha_tiled_row_stride, + left_alpha_tiled_col_stride: tl.constexpr, + left_alpha_row_block_stride, + left_alpha_col_block_stride: tl.constexpr, + # left_alpha epilogue + # right_alpha prologue + right_alpha_ptr, + right_alpha_batch_stride, + right_alpha_tiled_row_stride: tl.constexpr, + right_alpha_tiled_col_stride, + right_alpha_row_block_stride: tl.constexpr, + right_alpha_col_block_stride, + # right_alpha epilogue + # output prologue + output_ptr, + output_batch_stride, + output_tiled_row_stride, + output_tiled_col_stride, + output_row_block_stride, + output_col_block_stride, + # output epilogue + beta, + alpha, + beta_is_one: tl.constexpr, + beta_is_nonzero: tl.constexpr, + alpha_is_one: tl.constexpr, + left_alpha_is_one: tl.constexpr, + right_alpha_is_one: tl.constexpr, + BLOCKSIZE_ROW: tl.constexpr, + BLOCKSIZE_COL: tl.constexpr, + BLOCKSIZE_INNER: tl.constexpr, + acc_dtype: tl.constexpr, + allow_tf32: tl.constexpr, + GROUP_SIZE_ROW: tl.constexpr, + SPLIT_N: tl.constexpr, + ): + # left/right_alpha tensors are originally (* + 1)-dimensional + if left_alpha_tiled_col_stride != 0: + raise AssertionError( + f"left_alpha_tiled_col_stride must be 0, got {left_alpha_tiled_col_stride}" + ) + if left_alpha_col_block_stride != 0: + raise AssertionError( + f"left_alpha_col_block_stride must be 0, got {left_alpha_col_block_stride}" + ) + if right_alpha_tiled_row_stride != 0: + raise AssertionError( + f"right_alpha_tiled_row_stride must be 0, got {right_alpha_tiled_row_stride}" + ) + if right_alpha_row_block_stride != 0: + raise AssertionError( + f"right_alpha_row_block_stride must be 0, got {right_alpha_row_block_stride}" + ) + + batch_pid = tl.program_id(axis=2) + row_block_pid = tl.program_id(axis=0) + col_block_pid = tl.program_id(axis=1) + n_block_rows = tl.num_programs(axis=0) + n_block_cols = tl.num_programs(axis=1) + + row_block_pid, col_block_pid = tl.swizzle2d( + row_block_pid, col_block_pid, n_block_rows, n_block_cols, GROUP_SIZE_ROW + ) + + crow_indices_offset_ptr = ( + crow_indices_ptr + + crow_indices_batch_stride * batch_pid + + crow_indices_stride * row_block_pid + ) + nnz_offset = tl.load(crow_indices_offset_ptr) + nnz_offset_next = tl.load(crow_indices_offset_ptr + crow_indices_stride) + + # Compute nnz for the row with number row_block_pid. + row_nnz = nnz_offset_next - nnz_offset + + row_block_arange = tl.arange(0, BLOCKSIZE_ROW) + inner_block_arange = tl.arange(0, BLOCKSIZE_INNER) + col_block_arange = tl.arange(0, BLOCKSIZE_COL) + + # Pointers are set to the first block of the current row. + values_block_ptrs = ( + values_ptr + + values_batch_stride * batch_pid + + values_nnz_stride * nnz_offset + + values_row_block_stride * row_block_arange[:, None] + + values_col_block_stride * inner_block_arange[None, :] + ) + + # NOTE: dense is advanced into all dimensions but the tiled row one. + # That will be advanced in the loop according to values in col_indices. + dense_block_ptrs = ( + dense_ptr + + dense_batch_stride * batch_pid + + dense_tiled_col_stride * col_block_pid + + dense_row_block_stride * inner_block_arange[:, None] + + dense_col_block_stride * col_block_arange[None, :] + ) + + # Pointers are set to exact write-to locations + output_ptrs = ( + output_ptr + + output_batch_stride * batch_pid + + output_tiled_row_stride * row_block_pid + + output_tiled_col_stride * col_block_pid + + output_row_block_stride * row_block_arange[:, None] + + output_col_block_stride * col_block_arange[None, :] + ) + + # Set pointer to the first nonzero element in the current row + col_index_nnz_ptr = ( + col_indices_ptr + + col_indices_batch_stride * batch_pid + + col_indices_stride * nnz_offset + ) + + output_acc_block = tl.zeros((BLOCKSIZE_ROW, BLOCKSIZE_COL), dtype=acc_dtype) + + for _ in range(row_nnz): + values_block = tl.load(values_block_ptrs) + + # find which row of dense needs to get loaded + # for multiplication with values_block. + dense_row_idx = tl.load(col_index_nnz_ptr) + dense_block = tl.load( + dense_block_ptrs + dense_tiled_row_stride * dense_row_idx + ) + + # do block mm + output_acc_block += tl.dot( + values_block, dense_block, allow_tf32=allow_tf32, out_dtype=acc_dtype + ) + + # move val/col_index ptrs to the next block in the row + values_block_ptrs += values_nnz_stride + col_index_nnz_ptr += col_indices_stride + + if not alpha_is_one: + output_acc_block *= alpha + + if not left_alpha_is_one: + left_alpha_ptrs = ( + left_alpha_ptr + + left_alpha_batch_stride * batch_pid + + left_alpha_tiled_row_stride * row_block_pid + + left_alpha_tiled_col_stride * col_block_pid + + left_alpha_row_block_stride * row_block_arange[:, None] + + left_alpha_col_block_stride * col_block_arange[None, :] + ) + output_acc_block *= tl.load(left_alpha_ptrs) + + if not right_alpha_is_one: + right_alpha_ptrs = ( + right_alpha_ptr + + right_alpha_batch_stride * batch_pid + + right_alpha_tiled_row_stride * row_block_pid + + right_alpha_tiled_col_stride * col_block_pid + + right_alpha_row_block_stride * row_block_arange[:, None] + + right_alpha_col_block_stride * col_block_arange[None, :] + ) + output_acc_block *= tl.load(right_alpha_ptrs) + + if beta_is_nonzero: + input_ptrs = ( + input_ptr + + input_batch_stride * batch_pid + + input_tiled_row_stride * row_block_pid + + input_tiled_col_stride * col_block_pid + + input_row_block_stride * row_block_arange[:, None] + + input_col_block_stride * col_block_arange[None, :] + ) + if beta_is_one: + output_acc_block += tl.load(input_ptrs) + else: + output_acc_block += beta * tl.load(input_ptrs) + + # write back the result + tl.store(output_ptrs, output_acc_block.to(output_ptr.dtype.element_ty)) + +else: + bsr_softmax = None # type: ignore[assignment] + bsr_dense_mm = None # type: ignore[assignment] + sampled_addmm = None # type: ignore[assignment] + _scaled_dot_product_attention = None # type: ignore[assignment] + _scatter_mm2 = None # type: ignore[assignment] + _scatter_mm6 = None # type: ignore[assignment] + _bsr_strided_addmm_kernel = None # type: ignore[assignment] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_triton_ops_meta.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_triton_ops_meta.py new file mode 100644 index 0000000000000000000000000000000000000000..64f9ca15ca8dad16db6bd6df6a32969c2fdff3ba --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/_triton_ops_meta.py @@ -0,0 +1,7771 @@ +# mypy: allow-untyped-defs +"""Provides optimal triton kernel parameters. + +Aim +--- + +The usage of optimal triton kernel parameters may increase the +performance of operations several times. For example, for large tensor +shapes, the usage of a bsr tensor as mat1 argument in addmm-based +operations typically outperforms the corresponding operation with +strided-only inputs when the blocked representation of a tensor +provides a better alignment with memory access than what the strided +representation would provide. + +Pre-computed kernel parameters +------------------------------ + +This script finds and stores the optimal triton kernel parameters for +a specific set of shape configurations. For instance, the set of shape +configurations of the bsr_dense_addmm kernel is defined as + + input, out: M x N strided tensor + mat1: M x K bsr tensor with blocksize (BM, BK) and given sparsity + mat2: M x N strided tensor + dtype = float16, bfloat16, float32 + sparsity = 0.5 + M = 256, 512, ..., 16384 + K = M + N = 256, 512, ..., 131072 + BM = 16, 32, ..., 128 + BK = BM + alpha = 1 + beta = 0, 1 + GPUs: NVIDIA A100-SXM4-80GB + +Approximations +-------------- + +It is practically infeasible to pre-compute optimal kernel parameter +for all possible shape configurations as well as for all existing +GPUs. Therefore, we'll assume that the pre-computed optimal parameters +are good enough approximations when +1) the used GPU is any of NVIDIA A100 Tensor Core GPUs, +2) the actual sparsity of mat1 is different from sparsity value 0.5. + +If a particular shape configuration does not fall in the set of +pre-computed kernel parameters, or it does not match with the listed +approximations above, or the used GPU device is not a NVIDIA A100 GPU, +then a reference set of triton kernel parameters will be used when +executing operations. The reference kernel parameters are defined in +torch/sparse/_triton_ops.py, see bsr_dense_addmm_meta function, for +instance. + +Computing optimal kernel parameters +----------------------------------- + +If the approximations listed above are unacceptable, e.g. when one +seeks a maximal performance possible, the optimal kernel parameters +for a particular GPU can be computed by simply running this script in +the pytorch development tree:: + + cd /path/to/pytorch + python -m pip install --no-build-isolation -v -e . + python torch/sparse/_triton_ops_meta.py + +This will compute the optimal kernel parameters for the GPU device +available in the host system for all shape configurations listed in +"Pre-computed kernel parameters" above. The results will be stored in +the database of kernel parameters. Currently, this database is defined +as this module (see "BEGIN GENERATED DATA" comment below) that will be +modified when the script is run. Create a pytorch PR with the +corresponding modifications in this file to make the computed optimal +kernel parameters available for other users as pre-computed kernel +parameters. + +Moreover, one can compute the optimal kernel parameters for a specific +set of shape configurations and specific sparsity patterns. For that, +use tuning functions provided by this module: + + tune_bsr_dense_addmm(input, mat1, mat2, beta=1, alpha=1, out=None, verbose=False, store=False) -> meta + +The tuning functions return a dictionary of optimal kernel parameters +that can be passed to the corresponding operation, e.g. + + bsr_dense_addmm(..., meta=meta) + +Or, when store==True, the optimal kernel parameters will be stored in +the database of pre-computed kernel parameters in runtime so that all +addmm-based operations such as torch.addmm, torch.mm, +torch.nn.functional.linear will benefit from using the computed +optimal set of kernel parameters. + +Note that running tune_bsr_dense_addmm can take several minutes. So, +use it wisely, e.g. by implementing persistent storage of optimized +kernel parameters. See the source code of get_meta and +tune_bsr_dense_addmm to learn how to register a custom set of optimal +kernel parameters for addmm-based operations. + +""" + +__all__ = ["get_meta", "tune_bsr_dense_addmm", "tune__int_bsr_dense_addmm"] + +import inspect +import itertools +import re +import warnings +from typing import Any + +import torch +from torch.hub import tqdm +from torch.testing import make_tensor + + +def get_meta(op, key, device_name=None, version=(0, torch.float16, 0.5), exact=False): + """Return triton kernel meta parameters of the specified op and its inputs key. + + Parameters + ---------- + op (str): The name of an operation that implementation uses meta parameters. + key (tuple): A tuple of op input parameters, e.g. shapes, etc. + device_name (optional, str): The name of a device for which op + parameters are provided. + version (optional, hashable): Specifies the version of parameters. + exact (optional, bool): When True, the returned data (if + available) corresponds exactly to the specified device_name and + version information. Otherwise, if the corresponding data is not + available but there exists a data set that is computed for a + similar GPU device, then this data set will be returned. + + Returns + ------- + result (dict): The requested mapping of parameter names and + values, or None when no data is available. If the input `key` + contains `"*"`, the result will be a dictionary of keys and + mappings that match with the given `key`. + """ + if device_name is None: + device_name = torch.cuda.get_device_name() + + op_data = _operation_device_version_data.get((op, device_name, version)) + if op_data is None and not exact: + # A lack of op data could be due to using a (slightly) + # different GPU model compared to a model for which optimal + # meta parameters have been computed. In the following we'll + # assume that there is a set of GPU models that all have + # a similar set of optimal meta parameters. + if re.match(r"NVIDIA A100[^\d]", device_name) is not None: + device_name = "NVIDIA A100-SXM4-80GB" + else: + return + op_data = _operation_device_version_data.get((op, device_name, version)) + if op_data is None: + return + + matching_data = {} + if "*" in key: + for op_key in op_data: + if [ + None + for k1, k2 in zip(op_key, key, strict=True) + if k2 != "*" and k1 != k2 + ]: + continue + matching_data[op_key] = op_data[op_key] + else: + values = op_data.get(key) + if values is not None: + matching_data[key] = values + matching_meta = {} + for op_key, values in matching_data.items(): + if op == "scatter_mm": + names = ( + "GROUP_SIZE", + "SPLIT_N", + "TILE_M", + "TILE_N", + "num_stages", + "num_warps", + ) + meta = dict(zip(names, values, strict=True)) + elif op in {"bsr_dense_addmm", "_int_bsr_dense_addmm"}: + meta = dict( + zip( + ("GROUP_SIZE_ROW", "SPLIT_N", "num_stages", "num_warps"), + values, + strict=True, + ) + ) + else: + raise NotImplementedError(f"names for {op=}") + if "*" not in key: + return meta + + matching_meta[op_key] = meta + + if "*" in key: + return matching_meta + + +def update(op, device_name, version, key, value): + """Update the db of op parameters.""" + # skip storing possible optimization failures: + if not value: + warnings.warn( + f"skipping empty value for {op}: {device_name=} {version=} {key=}", + stacklevel=2, + ) + return + if (op, device_name, version) in _operation_device_version_data: + if _operation_device_version_data[op, device_name, version].get(key) == value: + return + _operation_device_version_data[op, device_name, version][key] = value + else: + _operation_device_version_data[op, device_name, version] = {key: value} + + +def dump(): + """Store the current runtime db state to the module file.""" + current_file = inspect.getfile(dump) + with open(current_file) as f: + current_content = f.read() + begin_data_str = "# BEGIN GENERATED DATA\n" + begin_data_index = current_content.find(begin_data_str) + end_data_index = current_content.find(" # END GENERATED DATA\n") + if begin_data_index == -1 or end_data_index == -1: + warnings.warn( + f"{current_file} cannot be updated:" + " BEGIN/END GENERATED DATA comment blocks appear to be corrupted", + stacklevel=2, + ) + return + + def sort_key(key): + op, device_name, version = key + version = tuple( + (str(item) if isinstance(item, torch.dtype) else item) for item in version + ) + return (op, device_name, version) + + part1 = current_content[: begin_data_index + len(begin_data_str)] + part2 = current_content[end_data_index:] + data_part = [] + for op_key in sorted(_operation_device_version_data, key=sort_key): + data_part.append(" " + repr(op_key).replace("'", '"') + ": {") + op_data = _operation_device_version_data[op_key] + data_part.extend(f" {key}: {op_data[key]}," for key in sorted(op_data)) + data_part.append(" },") + new_content = part1 + "\n".join(data_part) + "\n" + part2 + if current_content != new_content: + with open(current_file, "w") as f: + f.write(new_content) + + +def minimize( + target_func, + initial_parameters, + reference_parameters, + step_func, + max_step=2, + verbose=False, + all_values=None, +): + """Find a dict of parameters that minimizes the target function using + the initial dict of parameters and a step function that progresses + a specified parameter in a dict of parameters. + + Parameters + ---------- + target_func (callable): a functional with the signature + ``target_func(parameters: dict) -> float`` + initial_parameters (dict): a set of parameters used as an initial + value to the minimization process. + reference_parameters (dict): a set of parameters used as an + reference value with respect to which the speed up is computed. + step_func (callable): a functional with the signature + ``step_func(parameter_name:str, parameter_value:int, direction:int, parameters:dict) -> int`` + that increments or decrements (when ``direction`` is positive or + negative, respectively) the parameter with given name and value. + When return value is equal to ``parameter_value``, it means that + no step along the given direction can be made. + + Returns + ------- + parameters (dict): a set of parameters that minimizes the target + function. + speedup_incr (float): a speedup change given in percentage. + timing (float): the value of the target function at the parameters. + sensitivity_message (str): a message containing sensitivity. + information of parameters around the target function minimizer. + """ + + def to_key(parameters): + return tuple(parameters[k] for k in sorted(parameters)) + + def from_key(key, parameters): + return dict(zip(sorted(parameters), key, strict=True)) + + if all_values is None: + all_values = {} + + directions = list(range(-max_step, max_step + 1)) + names = sorted(initial_parameters) + all_directions = [] + for d_tuple in itertools.product(*((directions,) * len(names))): + dist = sum(map(abs, d_tuple)) + if dist > 0 and dist <= max_step: + all_directions.append((dist, d_tuple)) + all_directions.sort() + + try: + reference_target = target_func(reference_parameters) + except Exception as msg: + if verbose and "out of resource" not in str(msg): + print(f"{reference_parameters=} lead to failure: {msg}.") + reference_target = None + + if reference_target is not None: + all_values[to_key(reference_parameters)] = reference_target + + parameters = initial_parameters + try: + initial_target = target_func(parameters) + except Exception as msg: + if reference_target is None: + if verbose: + print( + f"{initial_parameters=} lead to failure: {msg}. Optimization failed!" + ) + return {}, -1, -1, f"{msg}" + if verbose and "out of resource" not in str(msg): + print( + f"{initial_parameters=} lead to failure: {msg}. Using reference parameters instead of initial parameters." + ) + parameters = reference_parameters + initial_target = reference_target + + if reference_target is None: + if verbose: + print("Using initial parameters instead of reference parameters.") + reference_target = initial_target + + initial_key = to_key(parameters) + minimal_target = all_values[initial_key] = initial_target + pbar = tqdm( + total=len(all_directions), + desc="Tuning...", + disable=not verbose, + ncols=75, + ) + while True: + for i, (_, d_tuple) in enumerate(all_directions): + pbar.update(1) + next_parameters = parameters.copy() + for name, direction in zip(names, d_tuple, strict=True): + value = next_parameters[name] + if direction == 0: + continue + next_value = step_func(name, value, direction, parameters) + if next_value == value: + break + next_parameters[name] = next_value + else: + next_key = to_key(next_parameters) + if next_key in all_values: + continue + try: + next_target = target_func(next_parameters) + except Exception as msg: + all_values[next_key] = str(msg) + if verbose and "out of resource" not in str(msg): + print(f"{next_parameters=} lead to failure: {msg}. Skipping.") + continue + all_values[next_key] = next_target + + if next_target < minimal_target: + minimal_target = next_target + parameters = next_parameters + # pyrefly: ignore [unsupported-operation] + pbar.total += i + 1 + break + else: + # ensure stable minimizer: + minimizer_keys = { + k + for k, v in all_values.items() + if isinstance(v, float) and abs(1 - v / minimal_target) < 0.001 + } + minimizer_key = ( + initial_key if initial_key in minimizer_keys else min(minimizer_keys) + ) + parameters = from_key(minimizer_key, parameters) + speedup_incr = (1 - minimal_target / reference_target) * 100 + if speedup_incr < 0: + if verbose: + print( + f"{speedup_incr=} is negative. Rerunning minimize with reference parameters as initial parameters." + ) + return minimize( + target_func, + reference_parameters, + reference_parameters, + step_func, + max_step=max_step, + verbose=verbose, + all_values=all_values, + ) + sensitivity = [] + for name in parameters: + value = parameters[name] + rel_diffs = [] + for direction in range(-max_step, max_step + 1): + if direction == 0: + continue + next_value = step_func(name, value, direction, parameters) + if next_value == value: + rel_diffs.append(0) + continue + next_parameters = parameters.copy() + next_parameters[name] = next_value + next_key = to_key(next_parameters) + next_target = all_values.get(next_key) + if next_target is None or isinstance(next_target, str): + rel_diffs.append(0) + continue + rel_diff = (next_target / minimal_target - 1) * 100 + rel_diffs.append(rel_diff) + sensitivity.append((max(rel_diffs), rel_diffs, name)) + + sensitivity_message = [f"timing0={initial_target:.3f}"] + for _, rel_diffs, name in sorted(sensitivity, reverse=True): + left_diffs = "|".join( + [f"{rel_diff:.1f}" for rel_diff in rel_diffs[:max_step]] + ) + right_diffs = "|".join( + [f"{rel_diff:.1f}" for rel_diff in rel_diffs[max_step:]] + ) + sensitivity_message.append( + f"{name}={parameters[name]} ({left_diffs}...{right_diffs} %)" + ) + sensitivity_message = ", ".join(sensitivity_message) + return parameters, speedup_incr, minimal_target, sensitivity_message + + +def create_blocked_tensor(B, M, N, blocksize, sparsity, dtype, device): + if sparsity < 0.0 or sparsity > 1.0: + raise AssertionError(f"sparsity should be between 0 and 1, got {sparsity}") + if M % blocksize[0] != 0: + raise AssertionError( + f"M ({M}) must be divisible by blocksize[0] ({blocksize[0]})" + ) + if N % blocksize[1] != 0: + raise AssertionError( + f"N ({N}) must be divisible by blocksize[1] ({blocksize[1]})" + ) + shape = (B, M // blocksize[0], N // blocksize[1])[int(B == 0) :] + A = torch.bernoulli( + torch.full(shape, 1 - sparsity, dtype=torch.float32, device=device) + ).to(dtype) + expected_nnz = int((1 - sparsity) * M * N / (blocksize[0] * blocksize[1])) + nonzero_indices = A.flatten().nonzero() + actual_nnz = nonzero_indices.shape[0] + if actual_nnz > expected_nnz: + selected_nonzeros = torch.randperm(actual_nnz)[: actual_nnz - expected_nnz] + A.flatten()[nonzero_indices[selected_nonzeros]] = 0 + elif actual_nnz < expected_nnz: + zero_indices = (A == 0).flatten().nonzero() + selected_zeros = torch.randperm(zero_indices.shape[0])[ + : expected_nnz - actual_nnz + ] + A.flatten()[zero_indices[selected_zeros]] = 1 + A = torch.repeat_interleave(A, blocksize[0], dim=-2) + A = torch.repeat_interleave(A, blocksize[1], dim=-1) + return A + + +def optimize_scatter_mm( + m, k, n, bm, bk, dtype=torch.float16, device="cuda", sparsity=0.5, force=False +): + import triton + + from torch.sparse._triton_ops import bsr_scatter_mm, bsr_scatter_mm_indices_data + + key = (m, k, n, bm, bk) + + version = (0, dtype, sparsity) + device_name = torch.cuda.get_device_name() + + reference_meta = dict( + GROUP_SIZE=1, + TILE_M=16, + TILE_N=16, + SPLIT_N=n // 16, + num_stages=1, + num_warps=1, + ) + + initial_meta = get_meta( + "scatter_mm", key, device_name=device_name, version=version, exact=True + ) + if initial_meta is None: + initial_meta = get_meta( + "bsr_dense_addmm", + key, + device_name=device_name, + version=(0, dtype, 0.5), + exact=True, + ) + if initial_meta is None: + initial_meta = reference_meta + elif not force: + return + + torch.manual_seed(0) + bsr = create_blocked_tensor( + 0, m, k, (bm, bk), sparsity, dtype, device + ).to_sparse_bsr((bm, bk)) + dense = make_tensor(k, n, dtype=dtype, device=device) + + def bench(meta, bsr=bsr, dense=dense): + indices_data = bsr_scatter_mm_indices_data( + bsr, dense, indices_format="bsr_strided_mm_compressed", **meta + ) + + def test_func(): + return bsr_scatter_mm(bsr, dense, indices_data=indices_data) + + ms_min = triton.testing.do_bench(test_func, warmup=500, rep=100) + + return ms_min + + def step_meta_parameter(name, value, direction, meta, m=m, n=n, k=k, bm=bm, bk=bk): + # return next value in positive or negative direction, or + # input value if the step will result an invalid + # value. The input value is assumed to be valid. + + is_log = name in {"SPLIT_N", "TILE_M", "TILE_N", "num_warps"} + min_value = dict( + SPLIT_N=1, TILE_M=16, TILE_N=16, num_warps=1, num_stages=1, GROUP_SIZE=1 + )[name] + max_value = dict( + SPLIT_N=n // meta["TILE_N"], TILE_M=bm, TILE_N=n // meta["SPLIT_N"] + ).get(name) + value_step = dict( + SPLIT_N=2, TILE_M=2, TILE_N=2, num_warps=2, num_stages=1, GROUP_SIZE=1 + )[name] + if is_log: + next_value = ( + value * value_step**direction + if direction > 0 + else value // (value_step ** abs(direction)) + ) + else: + next_value = value + value_step * direction + if min_value is not None: + next_value = max(next_value, min_value) + if max_value is not None: + next_value = min(next_value, max_value) + if name == "SPLIT_N" and n % next_value != 0: + return value + # Hard-skip parameter combinations that break CUDA state for pytorch: + if (dtype, name, next_value, m, n, k, bm, bk) in { + (torch.float32, "num_warps", 32, 256, 256, 256, 16, 16), + (torch.float32, "num_warps", 16, 256, 256, 256, 32, 32), + (torch.float32, "num_warps", 16, 256, 256, 256, 64, 64), + (torch.float32, "num_warps", 16, 256, 256, 256, 128, 128), + (torch.float32, "num_warps", 16, 512, 512, 256, 128, 128), + } and re.match(r"NVIDIA A100[^\d]", device_name) is not None: + return value + return next_value + + meta, speedup, timing, _sensitivity_message = minimize( + bench, initial_meta, reference_meta, step_meta_parameter + ) + if initial_meta is not reference_meta and initial_meta == meta and not force: + return + print(f"{meta=} {speedup=:.1f} % {timing=:.3f} ms") + if speedup < 0: + return + device_name = torch.cuda.get_device_name() + + update( + "scatter_mm", device_name, version, key, tuple(meta[k] for k in sorted(meta)) + ) + + +def tune__int_bsr_dense_addmm( + input, + bsr, + dense, + *, + beta=1, + alpha=1, + out=None, + store=False, + verbose=False, + force=False, +): + return tune_bsr_dense_addmm( + input, + bsr, + dense, + beta=beta, + alpha=alpha, + out=out, + store=store, + verbose=verbose, + force=force, + opname="_int_bsr_dense_addmm", + ) + + +def tune_bsr_dense_addmm( + input, + bsr, + dense, + *, + beta=1, + alpha=1, + left_alpha=None, + right_alpha=None, + out=None, + store=False, + verbose=False, + force=False, + opname=None, +): + """Tune bsr_dense_addmm kernel parameters against the given inputs. + + When store is True, the tuning results will be stored in the + database of kernel parameters. + """ + import triton + + if opname is None: + opname = "bsr_dense_addmm" + + if opname == "_int_bsr_dense_addmm": + from torch.sparse._triton_ops import _int_bsr_dense_addmm as bsr_dense_addmm + else: + from torch.sparse._triton_ops import bsr_dense_addmm + + N = dense.shape[-1] + values = bsr.values() + crow_indices = bsr.crow_indices() + batch_ndim = crow_indices.dim() - 1 + M, K = bsr.shape[batch_ndim : batch_ndim + 2] + BM, BK = values.shape[batch_ndim + 1 : batch_ndim + 3] + + # Reference parameters is a set of parameters that leads to a + # successful kernel call and the corresponding timing is used as a + # reference for computing speedups. Avoid changing the reference + # parameters when possible. + reference_meta = dict( + GROUP_SIZE_ROW=1, num_stages=1, num_warps=4, SPLIT_N=max(N // BM, 1) + ) + + # Compute the key of parameters: + sparsity = round(1 - bsr._nnz() * BM * BK / (M * K), 2) + dtype = bsr.dtype + if out is None: + out_dtype = dtype + else: + out_dtype = out.dtype + if out_dtype is dtype: + version_dtype = dtype + else: + version_dtype = (dtype, out_dtype) + version = (0, version_dtype, sparsity) + key = (M, K, N, BM, BK, beta == 0, beta == 1, alpha == 1) + + # For tuning, for an initial state, use parameters from the + # database if available, otherwise, use the reference parameters. + initial_meta = get_meta(opname, key, version=version, exact=True) + if initial_meta is None: + may_skip_update = False + initial_meta = get_meta(opname, key, version=(0, dtype, 0.5), exact=True) + if initial_meta is None: + initial_meta = reference_meta + elif not force: + return initial_meta + else: + may_skip_update = True + + # The target function that is minimized in the tuning process: + def bench(meta, input=input, bsr=bsr, dense=dense, alpha=alpha, out=out): + def test_func(): + return bsr_dense_addmm( + input, + bsr, + dense, + beta=beta, + alpha=alpha, + left_alpha=left_alpha, + right_alpha=right_alpha, + meta=meta, + out=out, + ) + + return triton.testing.do_bench(test_func, warmup=500, rep=100) + + # The step function that increments a specified meta parameter: + def step_meta_parameter(name, value, direction, meta, M=M, N=N, K=K, BM=BM, BK=BK): + # return next value in positive or negative direction, or + # input value if the step will result an invalid + # value. The input value is assumed to be valid. + is_log = name in {"SPLIT_N", "num_warps"} + min_value = dict(SPLIT_N=1, num_warps=1, num_stages=1, GROUP_SIZE_ROW=1)[name] + max_value = dict(SPLIT_N=max(N // BM, 1)).get(name) + value_step = dict(SPLIT_N=2, num_warps=2, num_stages=1, GROUP_SIZE_ROW=1)[name] + if is_log: + next_value = ( + value * value_step**direction + if direction > 0 + else value // (value_step ** abs(direction)) + ) + else: + next_value = value + value_step * direction + if min_value is not None: + next_value = max(next_value, min_value) + if max_value is not None: + next_value = min(next_value, max_value) + if name == "SPLIT_N" and N % next_value != 0: + return value + return next_value + + # Tune: + meta, speedup, timing, sensitivity_message = minimize( + bench, + initial_meta, + reference_meta, + step_meta_parameter, + max_step=2, + verbose=verbose, + ) + if verbose: + print(f"-> {sensitivity_message}, {speedup=:.1f} %, {timing=:.3f} ms") + + if store and not ( + may_skip_update and meta == initial_meta and initial_meta is not reference_meta + ): + device_name = torch.cuda.get_device_name() + update( + opname, + device_name, + version, + key, + tuple(meta[k] for k in sorted(meta)), + ) + + return meta + + +def optimize_bsr_dense_addmm( + m, + k, + n, + bm, + bk, + beta=1, + alpha=1, + use_left_alpha=False, + use_right_alpha=False, + dtype=torch.float16, + out_dtype=None, + device="cuda", + sparsity=0.5, + force=False, + verbose=False, + opname=None, +): + torch.manual_seed(0) + bsr = create_blocked_tensor( + 0, m, k, (bm, bk), sparsity, dtype, device + ).to_sparse_bsr((bm, bk)) + dense = make_tensor(k, n, dtype=dtype, device=device) + input = make_tensor(m, n, dtype=dtype, device=device) + left_alpha = make_tensor(m, dtype=dtype, device=device) if use_left_alpha else None + right_alpha = ( + make_tensor(n, dtype=dtype, device=device) if use_right_alpha else None + ) + if out_dtype is not None: + out = dense.new_empty((m, n), dtype=out_dtype) + else: + out = None + tune_bsr_dense_addmm( + input, + bsr, + dense, + beta=beta, + alpha=alpha, + left_alpha=left_alpha, + right_alpha=right_alpha, + out=out, + store=True, + force=force, + verbose=verbose, + opname=opname, + ) + + +def main(op="scatter_mm", force=False, dtype=torch.float16, verbose=True): + import itertools + + sizes_lst = [ + 256, + 512, + 1024, + 2048, + 4096, + 8192, + 16384, + 32768, + 65536, + 131072, + 50432, + 65792, + ] + sizes3_lst = [3 * sz for sz in [64, 128] + sizes_lst if sz <= 2048] + shapes_lst = [(sz, sz) for sz in sizes_lst[:-4] + sizes3_lst] + shapes_lst.extend([(3072, 768), (768, 3072)]) + shapes_lst.extend([(5120, 1280), (1280, 5120)]) + if dtype is torch.int8: + # triton does not support smaller blocks than 32 + blocksize_lst = [(32, 32), (64, 64), (128, 128), (256, 256)] + else: + blocksize_lst = [(16, 16), (32, 32), (64, 64), (128, 128)] + sparsity_lst = [0.5, 0.7, 0.3][:1] + for sparsity in sparsity_lst: + print(f"{op, dtype, sparsity=}") + try: + for (M, K), N, (BM, BK) in itertools.product( + shapes_lst, sizes_lst, blocksize_lst + ): + if not (BM <= M and BK <= K and M % BM == 0 and K % BK == 0): + continue + if op == "scatter_mm": + optimize_scatter_mm( + M, K, N, BM, BK, force=force, sparsity=sparsity, dtype=dtype + ) + elif op in {"bsr_dense_addmm", "_int_bsr_dense_addmm"}: + if M == K and N == 50432: + continue + print(f"{M, K, N, (BM, BK)=}") + for alpha, beta in [(1, 1), (1, 0)]: + optimize_bsr_dense_addmm( + M, + K, + N, + BM, + BK, + beta=beta, + alpha=alpha, + force=force, + sparsity=sparsity, + dtype=dtype, + verbose=verbose, + opname=op, + ) + else: + raise NotImplementedError(op) + except KeyboardInterrupt: + break + except Exception: + dump() + raise + dump() + + if 0: + # Check performance dependence on sparsity and apply + # adjustments when differences are noticeable (more than 10%). + # + # When using NVIDIA A100 GPU, the performance dependence on + # sparsity is insignificant (0 % ... 10 %) for majority of + # shapes/blocksizes combinations. However, for a very few + # specific size combinations, the effect of sparsity on + # performance can be up to 20 %. + for (M, K), N, (BM, BK) in itertools.product( + shapes_lst, sizes_lst, blocksize_lst + ): + meta_lst: list = [] + key = (M, K, N, BM, BK) + for sparsity1 in sparsity_lst: + torch.manual_seed(0) + bsr = create_blocked_tensor( + 0, M, K, (BM, BK), sparsity1, dtype, device="cuda" + ).to_sparse_bsr((BM, BK)) + dense = make_tensor(K, N, dtype=dtype, device="cuda") + meta_lst = [] + for sparsity in sparsity_lst: + meta = get_meta(op, key, version=(0, dtype, sparsity), exact=True) + if meta is None: + continue + + def bench(meta, bsr=bsr, dense=dense): + import triton + + if op == "scatter_mm": + from torch.sparse._triton_ops import ( + bsr_scatter_mm, + bsr_scatter_mm_indices_data, + ) + + indices_data = bsr_scatter_mm_indices_data( + bsr, + dense, + indices_format="bsr_strided_mm_compressed", + **meta, + ) + + def test_func(): + return bsr_scatter_mm( + bsr, dense, indices_data=indices_data + ) + + else: + raise NotImplementedError(op) + + ms_min = triton.testing.do_bench(test_func, warmup=500, rep=100) + + return ms_min + + meta_lst.append( + (bench(meta), sparsity, tuple(meta[k] for k in sorted(meta))) + ) + if not meta_lst: + continue + meta_lst = sorted(meta_lst) + index = next( + i for i, item in enumerate(meta_lst) if item[1] == sparsity1 + ) + if meta_lst[0][2] == meta_lst[index][2]: + continue + speeddiff = (1 - meta_lst[index][0] / meta_lst[0][0]) * 100 + if abs(speeddiff) < 10: + continue + + print(sparsity1, index, key, meta_lst, speeddiff) + + if index > 0: + device_name = torch.cuda.get_device_name() + meta = get_meta( + op, key, version=(0, dtype, meta_lst[0][1]), exact=True + ) + update( + op, + device_name, + (0, dtype, sparsity1), + key, + tuple(meta[k] for k in sorted(meta)), + ) + print("update") + dump() + + +_operation_device_version_data: dict[Any, dict] = { + # Warning: the data in between the BEGIN/END DATA comment lines + # below is generated. It can be updated either manually or via + # calling dump function defined above. + # + # Legend [op: key -> data]: + # scatter_mm : M, K, N, Ms, Ks -> GROUP_SIZE, SPLIT_N, TILE_M, TILE_N, num_stages, num_warps + # bsr_dense_addmm : M, K, N, Ms, Ks, beta==0, beta==1, alpha==1 -> GROUP_SIZE_ROW, SPLIT_N, num_stages, num_warps + # + # BEGIN GENERATED DATA + ("_int_bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.int8, 0.5)): { + (192, 192, 256, 32, 32, False, True, True): (2, 8, 1, 4), + (192, 192, 256, 32, 32, True, False, True): (2, 8, 5, 4), + (192, 192, 512, 32, 32, False, True, True): (1, 16, 1, 4), + (192, 192, 512, 32, 32, True, False, True): (1, 16, 5, 4), + (192, 192, 1024, 32, 32, False, True, True): (1, 32, 1, 4), + (192, 192, 1024, 32, 32, True, False, True): (4, 32, 4, 4), + (192, 192, 2048, 32, 32, False, True, True): (2, 64, 1, 4), + (192, 192, 2048, 32, 32, True, False, True): (3, 16, 5, 4), + (192, 192, 4096, 32, 32, False, True, True): (1, 128, 1, 4), + (192, 192, 4096, 32, 32, True, False, True): (1, 128, 1, 4), + (192, 192, 8192, 32, 32, False, True, True): (1, 256, 1, 4), + (192, 192, 8192, 32, 32, True, False, True): (1, 64, 3, 4), + (192, 192, 16384, 32, 32, False, True, True): (2, 512, 1, 4), + (192, 192, 16384, 32, 32, True, False, True): (5, 128, 1, 4), + (192, 192, 32768, 32, 32, False, True, True): (1, 1024, 1, 4), + (192, 192, 32768, 32, 32, True, False, True): (1, 256, 1, 4), + (192, 192, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (192, 192, 65536, 32, 32, True, False, True): (1, 512, 1, 4), + (192, 192, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (192, 192, 131072, 32, 32, True, False, True): (2, 512, 1, 4), + (256, 256, 256, 32, 32, False, True, True): (4, 8, 1, 4), + (256, 256, 256, 32, 32, True, False, True): (1, 8, 6, 4), + (256, 256, 256, 64, 64, False, True, True): (1, 4, 1, 16), + (256, 256, 256, 64, 64, True, False, True): (1, 4, 4, 4), + (256, 256, 256, 128, 128, False, True, True): (3, 2, 1, 16), + (256, 256, 256, 128, 128, True, False, True): (1, 2, 1, 4), + (256, 256, 512, 32, 32, False, True, True): (2, 16, 1, 4), + (256, 256, 512, 32, 32, True, False, True): (2, 16, 4, 4), + (256, 256, 512, 64, 64, False, True, True): (7, 8, 1, 16), + (256, 256, 512, 64, 64, True, False, True): (3, 8, 3, 4), + (256, 256, 512, 128, 128, False, True, True): (1, 4, 1, 32), + (256, 256, 512, 128, 128, True, False, True): (1, 4, 1, 4), + (256, 256, 1024, 32, 32, False, True, True): (1, 32, 1, 4), + (256, 256, 1024, 32, 32, True, False, True): (1, 8, 6, 4), + (256, 256, 1024, 64, 64, False, True, True): (2, 16, 1, 16), + (256, 256, 1024, 64, 64, True, False, True): (1, 16, 5, 4), + (256, 256, 1024, 128, 128, False, True, True): (4, 8, 1, 32), + (256, 256, 1024, 128, 128, True, False, True): (1, 8, 2, 4), + (256, 256, 2048, 32, 32, False, True, True): (1, 64, 1, 4), + (256, 256, 2048, 32, 32, True, False, True): (2, 32, 3, 2), + (256, 256, 2048, 64, 64, False, True, True): (2, 32, 1, 16), + (256, 256, 2048, 64, 64, True, False, True): (1, 16, 3, 4), + (256, 256, 2048, 128, 128, False, True, True): (1, 16, 1, 32), + (256, 256, 2048, 128, 128, True, False, True): (1, 16, 2, 4), + (256, 256, 4096, 32, 32, False, True, True): (2, 128, 1, 4), + (256, 256, 4096, 32, 32, True, False, True): (1, 32, 3, 2), + (256, 256, 4096, 64, 64, False, True, True): (2, 64, 1, 8), + (256, 256, 4096, 64, 64, True, False, True): (1, 64, 3, 2), + (256, 256, 4096, 128, 128, False, True, True): (2, 32, 1, 32), + (256, 256, 4096, 128, 128, True, False, True): (3, 32, 2, 8), + (256, 256, 8192, 32, 32, False, True, True): (1, 256, 1, 4), + (256, 256, 8192, 32, 32, True, False, True): (1, 64, 1, 4), + (256, 256, 8192, 64, 64, False, True, True): (1, 128, 1, 8), + (256, 256, 8192, 64, 64, True, False, True): (2, 128, 1, 4), + (256, 256, 8192, 128, 128, False, True, True): (4, 64, 1, 32), + (256, 256, 8192, 128, 128, True, False, True): (3, 64, 1, 4), + (256, 256, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (256, 256, 16384, 32, 32, True, False, True): (3, 128, 1, 4), + (256, 256, 16384, 64, 64, False, True, True): (2, 256, 1, 8), + (256, 256, 16384, 64, 64, True, False, True): (2, 256, 1, 4), + (256, 256, 16384, 128, 128, False, True, True): (2, 128, 1, 32), + (256, 256, 16384, 128, 128, True, False, True): (4, 128, 2, 4), + (256, 256, 32768, 32, 32, False, True, True): (2, 512, 1, 8), + (256, 256, 32768, 32, 32, True, False, True): (1, 256, 1, 4), + (256, 256, 32768, 64, 64, False, True, True): (1, 512, 1, 8), + (256, 256, 32768, 64, 64, True, False, True): (1, 512, 1, 4), + (256, 256, 32768, 128, 128, False, True, True): (2, 256, 1, 32), + (256, 256, 32768, 128, 128, True, False, True): (1, 256, 2, 4), + (256, 256, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (256, 256, 65536, 32, 32, True, False, True): (1, 512, 1, 4), + (256, 256, 65536, 64, 64, False, True, True): (1, 1024, 1, 8), + (256, 256, 65536, 64, 64, True, False, True): (1, 512, 1, 4), + (256, 256, 65536, 128, 128, False, True, True): (2, 512, 1, 16), + (256, 256, 65536, 128, 128, True, False, True): (1, 512, 1, 4), + (256, 256, 65792, 32, 32, False, True, True): (1, 1028, 1, 8), + (256, 256, 65792, 32, 32, True, False, True): (1, 514, 1, 4), + (256, 256, 65792, 64, 64, False, True, True): (1, 1028, 1, 8), + (256, 256, 65792, 64, 64, True, False, True): (4, 257, 1, 4), + (256, 256, 65792, 128, 128, False, True, True): (2, 514, 1, 16), + (256, 256, 65792, 128, 128, True, False, True): (3, 514, 1, 4), + (256, 256, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (256, 256, 131072, 32, 32, True, False, True): (2, 1024, 1, 4), + (256, 256, 131072, 64, 64, False, True, True): (1, 2048, 1, 8), + (256, 256, 131072, 64, 64, True, False, True): (2, 512, 1, 4), + (256, 256, 131072, 128, 128, False, True, True): (2, 1024, 1, 16), + (256, 256, 131072, 128, 128, True, False, True): (4, 1024, 1, 4), + (384, 384, 256, 32, 32, False, True, True): (1, 8, 1, 4), + (384, 384, 256, 32, 32, True, False, True): (5, 8, 5, 4), + (384, 384, 256, 64, 64, False, True, True): (2, 4, 1, 16), + (384, 384, 256, 64, 64, True, False, True): (1, 4, 5, 4), + (384, 384, 512, 32, 32, False, True, True): (2, 16, 1, 4), + (384, 384, 512, 32, 32, True, False, True): (1, 16, 4, 4), + (384, 384, 512, 64, 64, False, True, True): (3, 8, 1, 16), + (384, 384, 512, 64, 64, True, False, True): (3, 8, 3, 4), + (384, 384, 1024, 32, 32, False, True, True): (2, 32, 1, 4), + (384, 384, 1024, 32, 32, True, False, True): (1, 8, 6, 4), + (384, 384, 1024, 64, 64, False, True, True): (2, 16, 1, 16), + (384, 384, 1024, 64, 64, True, False, True): (1, 16, 5, 4), + (384, 384, 2048, 32, 32, False, True, True): (1, 64, 1, 4), + (384, 384, 2048, 32, 32, True, False, True): (3, 16, 4, 4), + (384, 384, 2048, 64, 64, False, True, True): (2, 32, 1, 16), + (384, 384, 2048, 64, 64, True, False, True): (1, 16, 4, 4), + (384, 384, 4096, 32, 32, False, True, True): (4, 64, 1, 8), + (384, 384, 4096, 32, 32, True, False, True): (4, 32, 1, 4), + (384, 384, 4096, 64, 64, False, True, True): (1, 64, 1, 8), + (384, 384, 4096, 64, 64, True, False, True): (1, 64, 1, 4), + (384, 384, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (384, 384, 8192, 32, 32, True, False, True): (3, 64, 1, 1), + (384, 384, 8192, 64, 64, False, True, True): (2, 128, 1, 8), + (384, 384, 8192, 64, 64, True, False, True): (1, 64, 2, 2), + (384, 384, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (384, 384, 16384, 32, 32, True, False, True): (1, 128, 1, 4), + (384, 384, 16384, 64, 64, False, True, True): (2, 256, 1, 8), + (384, 384, 16384, 64, 64, True, False, True): (2, 128, 1, 4), + (384, 384, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (384, 384, 32768, 32, 32, True, False, True): (1, 256, 1, 4), + (384, 384, 32768, 64, 64, False, True, True): (1, 512, 1, 8), + (384, 384, 32768, 64, 64, True, False, True): (1, 256, 3, 2), + (384, 384, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (384, 384, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (384, 384, 65536, 64, 64, False, True, True): (2, 1024, 1, 8), + (384, 384, 65536, 64, 64, True, False, True): (3, 256, 3, 4), + (384, 384, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (384, 384, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (384, 384, 131072, 64, 64, False, True, True): (1, 2048, 1, 8), + (384, 384, 131072, 64, 64, True, False, True): (2, 512, 3, 4), + (512, 512, 256, 32, 32, False, True, True): (1, 8, 1, 4), + (512, 512, 256, 32, 32, True, False, True): (4, 8, 4, 4), + (512, 512, 256, 64, 64, False, True, True): (3, 4, 1, 16), + (512, 512, 256, 64, 64, True, False, True): (2, 4, 5, 4), + (512, 512, 256, 128, 128, False, True, True): (4, 2, 1, 16), + (512, 512, 256, 128, 128, True, False, True): (1, 2, 3, 4), + (512, 512, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (512, 512, 256, 256, 256, True, False, True): (2, 1, 1, 32), + (512, 512, 512, 32, 32, False, True, True): (3, 16, 1, 4), + (512, 512, 512, 32, 32, True, False, True): (1, 8, 4, 2), + (512, 512, 512, 64, 64, False, True, True): (2, 8, 1, 16), + (512, 512, 512, 64, 64, True, False, True): (2, 8, 5, 4), + (512, 512, 512, 128, 128, False, True, True): (3, 4, 1, 16), + (512, 512, 512, 128, 128, True, False, True): (1, 4, 3, 4), + (512, 512, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (512, 512, 512, 256, 256, True, False, True): (2, 2, 1, 32), + (512, 512, 1024, 32, 32, False, True, True): (2, 32, 1, 4), + (512, 512, 1024, 32, 32, True, False, True): (4, 16, 3, 2), + (512, 512, 1024, 64, 64, False, True, True): (4, 16, 1, 16), + (512, 512, 1024, 64, 64, True, False, True): (1, 8, 4, 4), + (512, 512, 1024, 128, 128, False, True, True): (1, 8, 1, 32), + (512, 512, 1024, 128, 128, True, False, True): (1, 8, 3, 4), + (512, 512, 1024, 256, 256, False, True, True): (4, 4, 1, 32), + (512, 512, 1024, 256, 256, True, False, True): (2, 4, 1, 32), + (512, 512, 2048, 32, 32, False, True, True): (3, 32, 1, 8), + (512, 512, 2048, 32, 32, True, False, True): (1, 16, 3, 4), + (512, 512, 2048, 64, 64, False, True, True): (1, 32, 1, 8), + (512, 512, 2048, 64, 64, True, False, True): (1, 32, 3, 2), + (512, 512, 2048, 128, 128, False, True, True): (4, 16, 1, 32), + (512, 512, 2048, 128, 128, True, False, True): (1, 16, 3, 4), + (512, 512, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (512, 512, 2048, 256, 256, True, False, True): (3, 8, 1, 32), + (512, 512, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (512, 512, 4096, 32, 32, True, False, True): (5, 32, 1, 4), + (512, 512, 4096, 64, 64, False, True, True): (1, 64, 1, 8), + (512, 512, 4096, 64, 64, True, False, True): (1, 64, 1, 4), + (512, 512, 4096, 128, 128, False, True, True): (5, 32, 1, 32), + (512, 512, 4096, 128, 128, True, False, True): (2, 32, 3, 4), + (512, 512, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (512, 512, 4096, 256, 256, True, False, True): (3, 16, 1, 32), + (512, 512, 8192, 32, 32, False, True, True): (3, 128, 1, 8), + (512, 512, 8192, 32, 32, True, False, True): (3, 64, 1, 4), + (512, 512, 8192, 64, 64, False, True, True): (4, 128, 1, 8), + (512, 512, 8192, 64, 64, True, False, True): (1, 64, 3, 2), + (512, 512, 8192, 128, 128, False, True, True): (5, 64, 1, 32), + (512, 512, 8192, 128, 128, True, False, True): (1, 64, 2, 4), + (512, 512, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (512, 512, 8192, 256, 256, True, False, True): (1, 32, 1, 32), + (512, 512, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (512, 512, 16384, 32, 32, True, False, True): (2, 128, 1, 4), + (512, 512, 16384, 64, 64, False, True, True): (2, 256, 1, 8), + (512, 512, 16384, 64, 64, True, False, True): (1, 128, 3, 2), + (512, 512, 16384, 128, 128, False, True, True): (4, 128, 1, 16), + (512, 512, 16384, 128, 128, True, False, True): (2, 128, 1, 4), + (512, 512, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (512, 512, 16384, 256, 256, True, False, True): (2, 64, 1, 32), + (512, 512, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (512, 512, 32768, 32, 32, True, False, True): (2, 256, 1, 4), + (512, 512, 32768, 64, 64, False, True, True): (1, 512, 1, 8), + (512, 512, 32768, 64, 64, True, False, True): (1, 256, 3, 2), + (512, 512, 32768, 128, 128, False, True, True): (4, 256, 1, 16), + (512, 512, 32768, 128, 128, True, False, True): (2, 256, 1, 4), + (512, 512, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (512, 512, 32768, 256, 256, True, False, True): (2, 128, 1, 32), + (512, 512, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (512, 512, 65536, 32, 32, True, False, True): (2, 512, 1, 2), + (512, 512, 65536, 64, 64, False, True, True): (1, 1024, 1, 8), + (512, 512, 65536, 64, 64, True, False, True): (1, 512, 3, 2), + (512, 512, 65536, 128, 128, False, True, True): (4, 512, 1, 16), + (512, 512, 65536, 128, 128, True, False, True): (1, 512, 1, 4), + (512, 512, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (512, 512, 65536, 256, 256, True, False, True): (1, 256, 1, 32), + (512, 512, 65792, 32, 32, False, True, True): (1, 1028, 1, 8), + (512, 512, 65792, 32, 32, True, False, True): (1, 514, 3, 2), + (512, 512, 65792, 64, 64, False, True, True): (1, 1028, 1, 8), + (512, 512, 65792, 64, 64, True, False, True): (2, 257, 3, 4), + (512, 512, 65792, 128, 128, False, True, True): (4, 514, 1, 16), + (512, 512, 65792, 128, 128, True, False, True): (1, 514, 1, 4), + (512, 512, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (512, 512, 65792, 256, 256, True, False, True): (2, 257, 1, 32), + (512, 512, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (512, 512, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (512, 512, 131072, 64, 64, False, True, True): (1, 2048, 1, 8), + (512, 512, 131072, 64, 64, True, False, True): (1, 1024, 3, 2), + (512, 512, 131072, 128, 128, False, True, True): (4, 1024, 1, 16), + (512, 512, 131072, 128, 128, True, False, True): (1, 1024, 1, 4), + (512, 512, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (512, 512, 131072, 256, 256, True, False, True): (2, 512, 1, 32), + (768, 768, 256, 32, 32, False, True, True): (1, 8, 1, 4), + (768, 768, 256, 32, 32, True, False, True): (2, 8, 4, 4), + (768, 768, 256, 64, 64, False, True, True): (3, 4, 1, 16), + (768, 768, 256, 64, 64, True, False, True): (2, 4, 4, 4), + (768, 768, 256, 128, 128, False, True, True): (1, 2, 1, 8), + (768, 768, 256, 128, 128, True, False, True): (1, 2, 3, 4), + (768, 768, 512, 32, 32, False, True, True): (1, 16, 1, 4), + (768, 768, 512, 32, 32, True, False, True): (1, 4, 5, 4), + (768, 768, 512, 64, 64, False, True, True): (1, 8, 3, 32), + (768, 768, 512, 64, 64, True, False, True): (4, 8, 4, 4), + (768, 768, 512, 128, 128, False, True, True): (4, 4, 1, 16), + (768, 768, 512, 128, 128, True, False, True): (4, 4, 3, 4), + (768, 768, 1024, 32, 32, False, True, True): (1, 16, 1, 8), + (768, 768, 1024, 32, 32, True, False, True): (1, 8, 3, 4), + (768, 768, 1024, 64, 64, False, True, True): (3, 16, 1, 16), + (768, 768, 1024, 64, 64, True, False, True): (1, 8, 4, 4), + (768, 768, 1024, 128, 128, False, True, True): (3, 8, 1, 32), + (768, 768, 1024, 128, 128, True, False, True): (1, 8, 3, 4), + (768, 768, 2048, 32, 32, False, True, True): (2, 32, 1, 8), + (768, 768, 2048, 32, 32, True, False, True): (3, 16, 1, 4), + (768, 768, 2048, 64, 64, False, True, True): (1, 32, 1, 8), + (768, 768, 2048, 64, 64, True, False, True): (4, 8, 3, 4), + (768, 768, 2048, 128, 128, False, True, True): (1, 16, 1, 32), + (768, 768, 2048, 128, 128, True, False, True): (1, 16, 3, 4), + (768, 768, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (768, 768, 4096, 32, 32, True, False, True): (1, 32, 1, 1), + (768, 768, 4096, 64, 64, False, True, True): (2, 64, 1, 8), + (768, 768, 4096, 64, 64, True, False, True): (1, 32, 2, 2), + (768, 768, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (768, 768, 4096, 128, 128, True, False, True): (6, 32, 1, 4), + (768, 768, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (768, 768, 8192, 32, 32, True, False, True): (1, 64, 1, 4), + (768, 768, 8192, 64, 64, False, True, True): (1, 128, 1, 8), + (768, 768, 8192, 64, 64, True, False, True): (4, 32, 3, 4), + (768, 768, 8192, 128, 128, False, True, True): (2, 64, 1, 16), + (768, 768, 8192, 128, 128, True, False, True): (2, 64, 3, 4), + (768, 768, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (768, 768, 16384, 32, 32, True, False, True): (1, 128, 1, 4), + (768, 768, 16384, 64, 64, False, True, True): (1, 256, 1, 8), + (768, 768, 16384, 64, 64, True, False, True): (1, 128, 3, 2), + (768, 768, 16384, 128, 128, False, True, True): (2, 128, 1, 16), + (768, 768, 16384, 128, 128, True, False, True): (2, 128, 1, 4), + (768, 768, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (768, 768, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (768, 768, 32768, 64, 64, False, True, True): (2, 512, 1, 8), + (768, 768, 32768, 64, 64, True, False, True): (1, 256, 3, 2), + (768, 768, 32768, 128, 128, False, True, True): (2, 256, 1, 16), + (768, 768, 32768, 128, 128, True, False, True): (3, 256, 1, 4), + (768, 768, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (768, 768, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (768, 768, 65536, 64, 64, False, True, True): (2, 512, 1, 4), + (768, 768, 65536, 64, 64, True, False, True): (1, 512, 3, 2), + (768, 768, 65536, 128, 128, False, True, True): (2, 512, 1, 16), + (768, 768, 65536, 128, 128, True, False, True): (2, 512, 1, 4), + (768, 768, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (768, 768, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (768, 768, 131072, 64, 64, False, True, True): (2, 1024, 1, 4), + (768, 768, 131072, 64, 64, True, False, True): (2, 1024, 3, 2), + (768, 768, 131072, 128, 128, False, True, True): (2, 1024, 1, 16), + (768, 768, 131072, 128, 128, True, False, True): (2, 1024, 1, 4), + (768, 3072, 256, 32, 32, False, True, True): (3, 8, 4, 8), + (768, 3072, 256, 32, 32, True, False, True): (3, 8, 5, 4), + (768, 3072, 256, 64, 64, False, True, True): (1, 4, 4, 16), + (768, 3072, 256, 64, 64, True, False, True): (1, 4, 4, 4), + (768, 3072, 256, 128, 128, False, True, True): (2, 2, 1, 8), + (768, 3072, 256, 128, 128, True, False, True): (2, 2, 4, 4), + (768, 3072, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (768, 3072, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (768, 3072, 512, 32, 32, False, True, True): (1, 16, 1, 4), + (768, 3072, 512, 32, 32, True, False, True): (2, 4, 4, 4), + (768, 3072, 512, 64, 64, False, True, True): (3, 8, 4, 16), + (768, 3072, 512, 64, 64, True, False, True): (1, 8, 4, 4), + (768, 3072, 512, 128, 128, False, True, True): (2, 4, 1, 8), + (768, 3072, 512, 128, 128, True, False, True): (4, 4, 3, 4), + (768, 3072, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (768, 3072, 512, 256, 256, True, False, True): (1, 2, 1, 32), + (768, 3072, 1024, 32, 32, False, True, True): (1, 16, 1, 8), + (768, 3072, 1024, 32, 32, True, False, True): (3, 8, 3, 4), + (768, 3072, 1024, 64, 64, False, True, True): (2, 16, 1, 16), + (768, 3072, 1024, 64, 64, True, False, True): (1, 8, 3, 4), + (768, 3072, 1024, 128, 128, False, True, True): (1, 8, 1, 8), + (768, 3072, 1024, 128, 128, True, False, True): (3, 8, 4, 4), + (768, 3072, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (768, 3072, 1024, 256, 256, True, False, True): (4, 4, 1, 32), + (768, 3072, 2048, 32, 32, False, True, True): (3, 32, 1, 8), + (768, 3072, 2048, 32, 32, True, False, True): (4, 8, 3, 4), + (768, 3072, 2048, 64, 64, False, True, True): (5, 16, 1, 16), + (768, 3072, 2048, 64, 64, True, False, True): (6, 8, 3, 4), + (768, 3072, 2048, 128, 128, False, True, True): (2, 16, 1, 16), + (768, 3072, 2048, 128, 128, True, False, True): (1, 16, 4, 4), + (768, 3072, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (768, 3072, 2048, 256, 256, True, False, True): (1, 8, 1, 32), + (768, 3072, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (768, 3072, 4096, 32, 32, True, False, True): (1, 32, 3, 4), + (768, 3072, 4096, 64, 64, False, True, True): (1, 64, 1, 8), + (768, 3072, 4096, 64, 64, True, False, True): (2, 16, 3, 4), + (768, 3072, 4096, 128, 128, False, True, True): (1, 32, 1, 8), + (768, 3072, 4096, 128, 128, True, False, True): (2, 32, 2, 4), + (768, 3072, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (768, 3072, 4096, 256, 256, True, False, True): (1, 16, 1, 32), + (768, 3072, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (768, 3072, 8192, 32, 32, True, False, True): (1, 64, 1, 4), + (768, 3072, 8192, 64, 64, False, True, True): (1, 128, 1, 8), + (768, 3072, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (768, 3072, 8192, 128, 128, False, True, True): (2, 64, 1, 16), + (768, 3072, 8192, 128, 128, True, False, True): (2, 64, 3, 4), + (768, 3072, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (768, 3072, 8192, 256, 256, True, False, True): (1, 32, 1, 32), + (768, 3072, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (768, 3072, 16384, 32, 32, True, False, True): (1, 128, 1, 4), + (768, 3072, 16384, 64, 64, False, True, True): (1, 256, 1, 8), + (768, 3072, 16384, 64, 64, True, False, True): (2, 64, 3, 4), + (768, 3072, 16384, 128, 128, False, True, True): (2, 128, 1, 16), + (768, 3072, 16384, 128, 128, True, False, True): (2, 128, 3, 4), + (768, 3072, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (768, 3072, 16384, 256, 256, True, False, True): (1, 64, 1, 32), + (768, 3072, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (768, 3072, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (768, 3072, 32768, 64, 64, False, True, True): (1, 512, 1, 8), + (768, 3072, 32768, 64, 64, True, False, True): (3, 128, 3, 4), + (768, 3072, 32768, 128, 128, False, True, True): (2, 256, 1, 16), + (768, 3072, 32768, 128, 128, True, False, True): (2, 256, 3, 4), + (768, 3072, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (768, 3072, 32768, 256, 256, True, False, True): (1, 128, 1, 32), + (768, 3072, 50432, 32, 32, False, True, True): (1, 788, 1, 8), + (768, 3072, 50432, 32, 32, True, False, True): (1, 394, 3, 2), + (768, 3072, 50432, 64, 64, False, True, True): (1, 788, 1, 8), + (768, 3072, 50432, 64, 64, True, False, True): (2, 197, 3, 4), + (768, 3072, 50432, 128, 128, False, True, True): (2, 394, 1, 16), + (768, 3072, 50432, 128, 128, True, False, True): (2, 394, 3, 4), + (768, 3072, 50432, 256, 256, False, True, True): (1, 197, 1, 32), + (768, 3072, 50432, 256, 256, True, False, True): (1, 197, 1, 32), + (768, 3072, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (768, 3072, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (768, 3072, 65536, 64, 64, False, True, True): (1, 1024, 1, 8), + (768, 3072, 65536, 64, 64, True, False, True): (2, 256, 3, 4), + (768, 3072, 65536, 128, 128, False, True, True): (2, 512, 1, 16), + (768, 3072, 65536, 128, 128, True, False, True): (2, 512, 3, 4), + (768, 3072, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (768, 3072, 65536, 256, 256, True, False, True): (1, 256, 1, 32), + (768, 3072, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (768, 3072, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (768, 3072, 131072, 64, 64, False, True, True): (1, 2048, 1, 8), + (768, 3072, 131072, 64, 64, True, False, True): (2, 512, 3, 4), + (768, 3072, 131072, 128, 128, False, True, True): (2, 1024, 1, 16), + (768, 3072, 131072, 128, 128, True, False, True): (1, 1024, 3, 4), + (768, 3072, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (768, 3072, 131072, 256, 256, True, False, True): (1, 512, 1, 32), + (1024, 1024, 256, 32, 32, False, True, True): (1, 8, 1, 4), + (1024, 1024, 256, 32, 32, True, False, True): (1, 8, 5, 4), + (1024, 1024, 256, 64, 64, False, True, True): (1, 4, 1, 16), + (1024, 1024, 256, 64, 64, True, False, True): (4, 4, 4, 4), + (1024, 1024, 256, 128, 128, False, True, True): (1, 2, 1, 8), + (1024, 1024, 256, 128, 128, True, False, True): (1, 2, 3, 8), + (1024, 1024, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (1024, 1024, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (1024, 1024, 512, 32, 32, False, True, True): (5, 16, 1, 4), + (1024, 1024, 512, 32, 32, True, False, True): (2, 8, 4, 2), + (1024, 1024, 512, 64, 64, False, True, True): (4, 8, 1, 16), + (1024, 1024, 512, 64, 64, True, False, True): (1, 4, 3, 4), + (1024, 1024, 512, 128, 128, False, True, True): (3, 4, 1, 16), + (1024, 1024, 512, 128, 128, True, False, True): (1, 4, 2, 4), + (1024, 1024, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (1024, 1024, 512, 256, 256, True, False, True): (1, 2, 1, 32), + (1024, 1024, 1024, 32, 32, False, True, True): (1, 16, 1, 8), + (1024, 1024, 1024, 32, 32, True, False, True): (1, 8, 3, 4), + (1024, 1024, 1024, 64, 64, False, True, True): (3, 16, 1, 8), + (1024, 1024, 1024, 64, 64, True, False, True): (1, 16, 3, 2), + (1024, 1024, 1024, 128, 128, False, True, True): (1, 8, 1, 16), + (1024, 1024, 1024, 128, 128, True, False, True): (2, 8, 3, 8), + (1024, 1024, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (1024, 1024, 1024, 256, 256, True, False, True): (2, 4, 1, 32), + (1024, 1024, 2048, 32, 32, False, True, True): (2, 32, 1, 8), + (1024, 1024, 2048, 32, 32, True, False, True): (3, 16, 1, 4), + (1024, 1024, 2048, 64, 64, False, True, True): (1, 32, 1, 8), + (1024, 1024, 2048, 64, 64, True, False, True): (3, 32, 1, 4), + (1024, 1024, 2048, 128, 128, False, True, True): (4, 16, 1, 16), + (1024, 1024, 2048, 128, 128, True, False, True): (1, 16, 3, 4), + (1024, 1024, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (1024, 1024, 2048, 256, 256, True, False, True): (1, 8, 1, 32), + (1024, 1024, 4096, 32, 32, False, True, True): (4, 64, 1, 8), + (1024, 1024, 4096, 32, 32, True, False, True): (3, 32, 1, 4), + (1024, 1024, 4096, 64, 64, False, True, True): (3, 64, 1, 8), + (1024, 1024, 4096, 64, 64, True, False, True): (1, 32, 3, 2), + (1024, 1024, 4096, 128, 128, False, True, True): (4, 32, 1, 16), + (1024, 1024, 4096, 128, 128, True, False, True): (2, 32, 2, 4), + (1024, 1024, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (1024, 1024, 4096, 256, 256, True, False, True): (7, 16, 1, 32), + (1024, 1024, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (1024, 1024, 8192, 32, 32, True, False, True): (4, 64, 1, 4), + (1024, 1024, 8192, 64, 64, False, True, True): (2, 128, 1, 8), + (1024, 1024, 8192, 64, 64, True, False, True): (3, 32, 3, 4), + (1024, 1024, 8192, 128, 128, False, True, True): (4, 64, 1, 16), + (1024, 1024, 8192, 128, 128, True, False, True): (2, 64, 2, 4), + (1024, 1024, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (1024, 1024, 8192, 256, 256, True, False, True): (1, 32, 1, 32), + (1024, 1024, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (1024, 1024, 16384, 32, 32, True, False, True): (1, 128, 1, 4), + (1024, 1024, 16384, 64, 64, False, True, True): (1, 256, 1, 8), + (1024, 1024, 16384, 64, 64, True, False, True): (4, 64, 3, 4), + (1024, 1024, 16384, 128, 128, False, True, True): (4, 128, 1, 16), + (1024, 1024, 16384, 128, 128, True, False, True): (1, 128, 3, 4), + (1024, 1024, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (1024, 1024, 16384, 256, 256, True, False, True): (1, 64, 1, 32), + (1024, 1024, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (1024, 1024, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (1024, 1024, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (1024, 1024, 32768, 64, 64, True, False, True): (4, 128, 3, 4), + (1024, 1024, 32768, 128, 128, False, True, True): (4, 256, 1, 16), + (1024, 1024, 32768, 128, 128, True, False, True): (2, 256, 3, 4), + (1024, 1024, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (1024, 1024, 32768, 256, 256, True, False, True): (2, 128, 1, 32), + (1024, 1024, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (1024, 1024, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (1024, 1024, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (1024, 1024, 65536, 64, 64, True, False, True): (2, 256, 3, 4), + (1024, 1024, 65536, 128, 128, False, True, True): (4, 512, 1, 16), + (1024, 1024, 65536, 128, 128, True, False, True): (4, 512, 3, 4), + (1024, 1024, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (1024, 1024, 65536, 256, 256, True, False, True): (1, 256, 1, 32), + (1024, 1024, 65792, 32, 32, False, True, True): (1, 1028, 1, 8), + (1024, 1024, 65792, 32, 32, True, False, True): (1, 514, 3, 2), + (1024, 1024, 65792, 64, 64, False, True, True): (2, 514, 1, 4), + (1024, 1024, 65792, 64, 64, True, False, True): (4, 257, 3, 4), + (1024, 1024, 65792, 128, 128, False, True, True): (2, 514, 1, 16), + (1024, 1024, 65792, 128, 128, True, False, True): (2, 514, 2, 4), + (1024, 1024, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (1024, 1024, 65792, 256, 256, True, False, True): (1, 257, 1, 32), + (1024, 1024, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (1024, 1024, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (1024, 1024, 131072, 64, 64, False, True, True): (2, 1024, 1, 4), + (1024, 1024, 131072, 64, 64, True, False, True): (2, 512, 3, 4), + (1024, 1024, 131072, 128, 128, False, True, True): (4, 1024, 1, 16), + (1024, 1024, 131072, 128, 128, True, False, True): (1, 1024, 3, 4), + (1024, 1024, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (1024, 1024, 131072, 256, 256, True, False, True): (1, 512, 1, 32), + (1280, 5120, 65792, 32, 32, False, True, True): (1, 1028, 1, 8), + (1280, 5120, 65792, 32, 32, True, False, True): (1, 514, 3, 2), + (1280, 5120, 65792, 64, 64, False, True, True): (1, 1028, 1, 8), + (1280, 5120, 65792, 64, 64, True, False, True): (2, 257, 3, 4), + (1280, 5120, 65792, 128, 128, False, True, True): (2, 514, 1, 16), + (1280, 5120, 65792, 128, 128, True, False, True): (1, 514, 3, 4), + (1280, 5120, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (1280, 5120, 65792, 256, 256, True, False, True): (1, 257, 1, 32), + (1536, 1536, 256, 32, 32, False, True, True): (1, 8, 1, 4), + (1536, 1536, 256, 32, 32, True, False, True): (2, 8, 1, 8), + (1536, 1536, 256, 64, 64, False, True, True): (4, 4, 1, 16), + (1536, 1536, 256, 64, 64, True, False, True): (1, 4, 4, 4), + (1536, 1536, 256, 128, 128, False, True, True): (2, 2, 1, 16), + (1536, 1536, 256, 128, 128, True, False, True): (2, 2, 3, 4), + (1536, 1536, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (1536, 1536, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (1536, 1536, 512, 32, 32, False, True, True): (1, 8, 1, 8), + (1536, 1536, 512, 32, 32, True, False, True): (3, 4, 4, 4), + (1536, 1536, 512, 64, 64, False, True, True): (3, 8, 1, 16), + (1536, 1536, 512, 64, 64, True, False, True): (1, 4, 3, 4), + (1536, 1536, 512, 128, 128, False, True, True): (1, 4, 1, 16), + (1536, 1536, 512, 128, 128, True, False, True): (2, 4, 4, 4), + (1536, 1536, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (1536, 1536, 512, 256, 256, True, False, True): (1, 2, 1, 32), + (1536, 1536, 1024, 32, 32, False, True, True): (4, 16, 1, 8), + (1536, 1536, 1024, 32, 32, True, False, True): (2, 8, 1, 4), + (1536, 1536, 1024, 64, 64, False, True, True): (2, 16, 1, 16), + (1536, 1536, 1024, 64, 64, True, False, True): (2, 4, 3, 4), + (1536, 1536, 1024, 128, 128, False, True, True): (3, 8, 1, 32), + (1536, 1536, 1024, 128, 128, True, False, True): (4, 8, 3, 4), + (1536, 1536, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (1536, 1536, 1024, 256, 256, True, False, True): (1, 4, 1, 32), + (1536, 1536, 2048, 32, 32, False, True, True): (1, 32, 1, 8), + (1536, 1536, 2048, 32, 32, True, False, True): (1, 16, 1, 4), + (1536, 1536, 2048, 64, 64, False, True, True): (1, 32, 1, 8), + (1536, 1536, 2048, 64, 64, True, False, True): (1, 16, 2, 2), + (1536, 1536, 2048, 128, 128, False, True, True): (2, 16, 1, 16), + (1536, 1536, 2048, 128, 128, True, False, True): (4, 16, 2, 4), + (1536, 1536, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (1536, 1536, 2048, 256, 256, True, False, True): (1, 8, 1, 32), + (1536, 1536, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (1536, 1536, 4096, 32, 32, True, False, True): (1, 32, 1, 4), + (1536, 1536, 4096, 64, 64, False, True, True): (3, 64, 1, 8), + (1536, 1536, 4096, 64, 64, True, False, True): (1, 32, 3, 2), + (1536, 1536, 4096, 128, 128, False, True, True): (1, 32, 1, 8), + (1536, 1536, 4096, 128, 128, True, False, True): (2, 32, 2, 4), + (1536, 1536, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (1536, 1536, 4096, 256, 256, True, False, True): (2, 16, 1, 32), + (1536, 1536, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (1536, 1536, 8192, 32, 32, True, False, True): (1, 64, 1, 4), + (1536, 1536, 8192, 64, 64, False, True, True): (3, 128, 1, 8), + (1536, 1536, 8192, 64, 64, True, False, True): (1, 64, 3, 2), + (1536, 1536, 8192, 128, 128, False, True, True): (1, 64, 1, 8), + (1536, 1536, 8192, 128, 128, True, False, True): (1, 64, 2, 4), + (1536, 1536, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (1536, 1536, 8192, 256, 256, True, False, True): (2, 32, 1, 32), + (1536, 1536, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (1536, 1536, 16384, 32, 32, True, False, True): (1, 128, 3, 2), + (1536, 1536, 16384, 64, 64, False, True, True): (2, 128, 1, 4), + (1536, 1536, 16384, 64, 64, True, False, True): (2, 64, 3, 4), + (1536, 1536, 16384, 128, 128, False, True, True): (1, 128, 1, 8), + (1536, 1536, 16384, 128, 128, True, False, True): (2, 128, 2, 4), + (1536, 1536, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (1536, 1536, 16384, 256, 256, True, False, True): (2, 64, 1, 32), + (1536, 1536, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (1536, 1536, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (1536, 1536, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (1536, 1536, 32768, 64, 64, True, False, True): (3, 128, 3, 4), + (1536, 1536, 32768, 128, 128, False, True, True): (1, 256, 1, 8), + (1536, 1536, 32768, 128, 128, True, False, True): (1, 256, 2, 4), + (1536, 1536, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (1536, 1536, 32768, 256, 256, True, False, True): (2, 128, 1, 32), + (1536, 1536, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (1536, 1536, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (1536, 1536, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (1536, 1536, 65536, 64, 64, True, False, True): (1, 512, 3, 2), + (1536, 1536, 65536, 128, 128, False, True, True): (1, 512, 1, 8), + (1536, 1536, 65536, 128, 128, True, False, True): (1, 512, 3, 4), + (1536, 1536, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (1536, 1536, 65536, 256, 256, True, False, True): (2, 256, 1, 32), + (1536, 1536, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (1536, 1536, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (1536, 1536, 131072, 64, 64, False, True, True): (3, 1024, 1, 4), + (1536, 1536, 131072, 64, 64, True, False, True): (3, 512, 3, 4), + (1536, 1536, 131072, 128, 128, False, True, True): (1, 1024, 1, 8), + (1536, 1536, 131072, 128, 128, True, False, True): (1, 1024, 3, 4), + (1536, 1536, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (1536, 1536, 131072, 256, 256, True, False, True): (2, 512, 1, 32), + (2048, 2048, 256, 32, 32, False, True, True): (3, 8, 1, 4), + (2048, 2048, 256, 32, 32, True, False, True): (1, 4, 4, 2), + (2048, 2048, 256, 64, 64, False, True, True): (2, 4, 1, 16), + (2048, 2048, 256, 64, 64, True, False, True): (1, 2, 3, 4), + (2048, 2048, 256, 128, 128, False, True, True): (1, 2, 1, 8), + (2048, 2048, 256, 128, 128, True, False, True): (1, 2, 4, 4), + (2048, 2048, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (2048, 2048, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (2048, 2048, 512, 32, 32, False, True, True): (3, 8, 1, 8), + (2048, 2048, 512, 32, 32, True, False, True): (4, 4, 3, 2), + (2048, 2048, 512, 64, 64, False, True, True): (1, 8, 1, 8), + (2048, 2048, 512, 64, 64, True, False, True): (1, 8, 3, 4), + (2048, 2048, 512, 128, 128, False, True, True): (1, 4, 1, 8), + (2048, 2048, 512, 128, 128, True, False, True): (1, 4, 4, 4), + (2048, 2048, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (2048, 2048, 512, 256, 256, True, False, True): (2, 2, 1, 32), + (2048, 2048, 1024, 32, 32, False, True, True): (1, 16, 1, 8), + (2048, 2048, 1024, 32, 32, True, False, True): (3, 8, 1, 4), + (2048, 2048, 1024, 64, 64, False, True, True): (4, 16, 1, 8), + (2048, 2048, 1024, 64, 64, True, False, True): (1, 8, 3, 2), + (2048, 2048, 1024, 128, 128, False, True, True): (4, 8, 1, 16), + (2048, 2048, 1024, 128, 128, True, False, True): (2, 8, 2, 4), + (2048, 2048, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (2048, 2048, 1024, 256, 256, True, False, True): (3, 4, 1, 32), + (2048, 2048, 2048, 32, 32, False, True, True): (1, 32, 1, 8), + (2048, 2048, 2048, 32, 32, True, False, True): (1, 16, 1, 4), + (2048, 2048, 2048, 64, 64, False, True, True): (1, 32, 1, 8), + (2048, 2048, 2048, 64, 64, True, False, True): (1, 16, 3, 2), + (2048, 2048, 2048, 128, 128, False, True, True): (4, 16, 1, 16), + (2048, 2048, 2048, 128, 128, True, False, True): (2, 16, 2, 4), + (2048, 2048, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (2048, 2048, 2048, 256, 256, True, False, True): (1, 8, 1, 32), + (2048, 2048, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (2048, 2048, 4096, 32, 32, True, False, True): (1, 32, 1, 4), + (2048, 2048, 4096, 64, 64, False, True, True): (4, 64, 1, 8), + (2048, 2048, 4096, 64, 64, True, False, True): (2, 16, 3, 4), + (2048, 2048, 4096, 128, 128, False, True, True): (4, 32, 1, 8), + (2048, 2048, 4096, 128, 128, True, False, True): (1, 32, 2, 4), + (2048, 2048, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (2048, 2048, 4096, 256, 256, True, False, True): (4, 16, 1, 32), + (2048, 2048, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (2048, 2048, 8192, 32, 32, True, False, True): (1, 64, 1, 4), + (2048, 2048, 8192, 64, 64, False, True, True): (2, 64, 1, 4), + (2048, 2048, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (2048, 2048, 8192, 128, 128, False, True, True): (4, 64, 1, 8), + (2048, 2048, 8192, 128, 128, True, False, True): (2, 64, 2, 4), + (2048, 2048, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (2048, 2048, 8192, 256, 256, True, False, True): (4, 32, 1, 32), + (2048, 2048, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (2048, 2048, 16384, 32, 32, True, False, True): (1, 128, 3, 2), + (2048, 2048, 16384, 64, 64, False, True, True): (2, 128, 1, 4), + (2048, 2048, 16384, 64, 64, True, False, True): (2, 64, 3, 4), + (2048, 2048, 16384, 128, 128, False, True, True): (1, 128, 1, 8), + (2048, 2048, 16384, 128, 128, True, False, True): (2, 128, 2, 4), + (2048, 2048, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (2048, 2048, 16384, 256, 256, True, False, True): (4, 64, 1, 32), + (2048, 2048, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (2048, 2048, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (2048, 2048, 32768, 64, 64, False, True, True): (2, 256, 1, 4), + (2048, 2048, 32768, 64, 64, True, False, True): (2, 128, 3, 4), + (2048, 2048, 32768, 128, 128, False, True, True): (1, 256, 1, 8), + (2048, 2048, 32768, 128, 128, True, False, True): (2, 256, 2, 4), + (2048, 2048, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (2048, 2048, 32768, 256, 256, True, False, True): (4, 128, 1, 32), + (2048, 2048, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (2048, 2048, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (2048, 2048, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (2048, 2048, 65536, 64, 64, True, False, True): (2, 256, 3, 4), + (2048, 2048, 65536, 128, 128, False, True, True): (1, 512, 1, 8), + (2048, 2048, 65536, 128, 128, True, False, True): (1, 512, 2, 4), + (2048, 2048, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (2048, 2048, 65536, 256, 256, True, False, True): (4, 256, 1, 32), + (2048, 2048, 65792, 32, 32, False, True, True): (1, 1028, 1, 8), + (2048, 2048, 65792, 32, 32, True, False, True): (1, 514, 3, 2), + (2048, 2048, 65792, 64, 64, False, True, True): (1, 514, 1, 4), + (2048, 2048, 65792, 64, 64, True, False, True): (2, 257, 3, 4), + (2048, 2048, 65792, 128, 128, False, True, True): (1, 514, 1, 8), + (2048, 2048, 65792, 128, 128, True, False, True): (1, 514, 2, 4), + (2048, 2048, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (2048, 2048, 65792, 256, 256, True, False, True): (1, 257, 1, 32), + (2048, 2048, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (2048, 2048, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (2048, 2048, 131072, 64, 64, False, True, True): (1, 1024, 1, 4), + (2048, 2048, 131072, 64, 64, True, False, True): (2, 512, 3, 4), + (2048, 2048, 131072, 128, 128, False, True, True): (1, 1024, 1, 8), + (2048, 2048, 131072, 128, 128, True, False, True): (1, 1024, 3, 4), + (2048, 2048, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (2048, 2048, 131072, 256, 256, True, False, True): (4, 512, 1, 32), + (3072, 768, 256, 32, 32, False, True, True): (5, 4, 1, 8), + (3072, 768, 256, 32, 32, True, False, True): (2, 2, 4, 4), + (3072, 768, 256, 64, 64, False, True, True): (1, 4, 1, 16), + (3072, 768, 256, 64, 64, True, False, True): (2, 2, 3, 4), + (3072, 768, 256, 128, 128, False, True, True): (5, 2, 1, 16), + (3072, 768, 256, 128, 128, True, False, True): (1, 2, 5, 4), + (3072, 768, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (3072, 768, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (3072, 768, 512, 32, 32, False, True, True): (1, 8, 1, 8), + (3072, 768, 512, 32, 32, True, False, True): (5, 4, 1, 4), + (3072, 768, 512, 64, 64, False, True, True): (1, 8, 1, 8), + (3072, 768, 512, 64, 64, True, False, True): (3, 2, 3, 4), + (3072, 768, 512, 128, 128, False, True, True): (3, 4, 1, 32), + (3072, 768, 512, 128, 128, True, False, True): (2, 4, 3, 4), + (3072, 768, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (3072, 768, 512, 256, 256, True, False, True): (2, 2, 1, 32), + (3072, 768, 1024, 32, 32, False, True, True): (2, 16, 1, 8), + (3072, 768, 1024, 32, 32, True, False, True): (3, 8, 1, 4), + (3072, 768, 1024, 64, 64, False, True, True): (4, 16, 1, 8), + (3072, 768, 1024, 64, 64, True, False, True): (1, 8, 3, 2), + (3072, 768, 1024, 128, 128, False, True, True): (2, 8, 1, 32), + (3072, 768, 1024, 128, 128, True, False, True): (3, 8, 2, 4), + (3072, 768, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (3072, 768, 1024, 256, 256, True, False, True): (4, 4, 1, 32), + (3072, 768, 2048, 32, 32, False, True, True): (1, 32, 1, 8), + (3072, 768, 2048, 32, 32, True, False, True): (1, 16, 1, 4), + (3072, 768, 2048, 64, 64, False, True, True): (2, 32, 1, 8), + (3072, 768, 2048, 64, 64, True, False, True): (2, 8, 3, 4), + (3072, 768, 2048, 128, 128, False, True, True): (2, 16, 1, 16), + (3072, 768, 2048, 128, 128, True, False, True): (2, 16, 1, 4), + (3072, 768, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (3072, 768, 2048, 256, 256, True, False, True): (2, 8, 1, 32), + (3072, 768, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (3072, 768, 4096, 32, 32, True, False, True): (1, 32, 1, 2), + (3072, 768, 4096, 64, 64, False, True, True): (2, 64, 1, 8), + (3072, 768, 4096, 64, 64, True, False, True): (2, 32, 2, 2), + (3072, 768, 4096, 128, 128, False, True, True): (1, 32, 1, 8), + (3072, 768, 4096, 128, 128, True, False, True): (2, 32, 2, 4), + (3072, 768, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (3072, 768, 4096, 256, 256, True, False, True): (4, 16, 1, 32), + (3072, 768, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (3072, 768, 8192, 32, 32, True, False, True): (3, 64, 1, 2), + (3072, 768, 8192, 64, 64, False, True, True): (1, 128, 1, 8), + (3072, 768, 8192, 64, 64, True, False, True): (2, 64, 2, 2), + (3072, 768, 8192, 128, 128, False, True, True): (1, 64, 1, 8), + (3072, 768, 8192, 128, 128, True, False, True): (2, 64, 2, 4), + (3072, 768, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (3072, 768, 8192, 256, 256, True, False, True): (4, 32, 1, 32), + (3072, 768, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (3072, 768, 16384, 32, 32, True, False, True): (1, 128, 1, 2), + (3072, 768, 16384, 64, 64, False, True, True): (2, 128, 1, 4), + (3072, 768, 16384, 64, 64, True, False, True): (1, 128, 2, 2), + (3072, 768, 16384, 128, 128, False, True, True): (1, 128, 1, 8), + (3072, 768, 16384, 128, 128, True, False, True): (1, 128, 1, 4), + (3072, 768, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (3072, 768, 16384, 256, 256, True, False, True): (4, 64, 1, 32), + (3072, 768, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (3072, 768, 32768, 32, 32, True, False, True): (1, 256, 1, 2), + (3072, 768, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (3072, 768, 32768, 64, 64, True, False, True): (2, 256, 2, 2), + (3072, 768, 32768, 128, 128, False, True, True): (1, 256, 1, 8), + (3072, 768, 32768, 128, 128, True, False, True): (2, 256, 1, 4), + (3072, 768, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (3072, 768, 32768, 256, 256, True, False, True): (4, 128, 1, 32), + (3072, 768, 50432, 32, 32, False, True, True): (1, 788, 1, 8), + (3072, 768, 50432, 32, 32, True, False, True): (1, 394, 1, 2), + (3072, 768, 50432, 64, 64, False, True, True): (2, 394, 1, 4), + (3072, 768, 50432, 64, 64, True, False, True): (2, 394, 2, 2), + (3072, 768, 50432, 128, 128, False, True, True): (1, 394, 1, 8), + (3072, 768, 50432, 128, 128, True, False, True): (2, 394, 1, 4), + (3072, 768, 50432, 256, 256, False, True, True): (1, 197, 1, 32), + (3072, 768, 50432, 256, 256, True, False, True): (1, 197, 1, 32), + (3072, 768, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (3072, 768, 65536, 32, 32, True, False, True): (1, 512, 1, 2), + (3072, 768, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (3072, 768, 65536, 64, 64, True, False, True): (2, 512, 2, 2), + (3072, 768, 65536, 128, 128, False, True, True): (1, 512, 1, 8), + (3072, 768, 65536, 128, 128, True, False, True): (2, 512, 1, 4), + (3072, 768, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (3072, 768, 65536, 256, 256, True, False, True): (4, 256, 1, 32), + (3072, 768, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (3072, 768, 131072, 32, 32, True, False, True): (1, 1024, 1, 2), + (3072, 768, 131072, 64, 64, False, True, True): (2, 1024, 1, 4), + (3072, 768, 131072, 64, 64, True, False, True): (2, 1024, 2, 2), + (3072, 768, 131072, 128, 128, False, True, True): (1, 1024, 1, 8), + (3072, 768, 131072, 128, 128, True, False, True): (2, 1024, 1, 4), + (3072, 768, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (3072, 768, 131072, 256, 256, True, False, True): (4, 512, 1, 32), + (3072, 3072, 256, 32, 32, False, True, True): (1, 4, 1, 8), + (3072, 3072, 256, 32, 32, True, False, True): (2, 2, 5, 4), + (3072, 3072, 256, 64, 64, False, True, True): (2, 4, 1, 16), + (3072, 3072, 256, 64, 64, True, False, True): (3, 2, 3, 4), + (3072, 3072, 256, 128, 128, False, True, True): (1, 2, 1, 8), + (3072, 3072, 256, 128, 128, True, False, True): (1, 2, 5, 4), + (3072, 3072, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (3072, 3072, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (3072, 3072, 512, 32, 32, False, True, True): (1, 8, 1, 8), + (3072, 3072, 512, 32, 32, True, False, True): (3, 2, 3, 4), + (3072, 3072, 512, 64, 64, False, True, True): (1, 8, 1, 8), + (3072, 3072, 512, 64, 64, True, False, True): (3, 2, 3, 4), + (3072, 3072, 512, 128, 128, False, True, True): (2, 4, 1, 8), + (3072, 3072, 512, 128, 128, True, False, True): (2, 4, 4, 4), + (3072, 3072, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (3072, 3072, 512, 256, 256, True, False, True): (1, 2, 1, 32), + (3072, 3072, 1024, 32, 32, False, True, True): (1, 16, 1, 8), + (3072, 3072, 1024, 32, 32, True, False, True): (3, 8, 3, 4), + (3072, 3072, 1024, 64, 64, False, True, True): (2, 16, 1, 8), + (3072, 3072, 1024, 64, 64, True, False, True): (2, 4, 3, 4), + (3072, 3072, 1024, 128, 128, False, True, True): (1, 8, 1, 8), + (3072, 3072, 1024, 128, 128, True, False, True): (3, 8, 2, 4), + (3072, 3072, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (3072, 3072, 1024, 256, 256, True, False, True): (3, 4, 1, 32), + (3072, 3072, 2048, 32, 32, False, True, True): (1, 32, 1, 8), + (3072, 3072, 2048, 32, 32, True, False, True): (1, 16, 1, 4), + (3072, 3072, 2048, 64, 64, False, True, True): (1, 32, 1, 8), + (3072, 3072, 2048, 64, 64, True, False, True): (1, 16, 3, 2), + (3072, 3072, 2048, 128, 128, False, True, True): (1, 16, 1, 8), + (3072, 3072, 2048, 128, 128, True, False, True): (2, 16, 2, 4), + (3072, 3072, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (3072, 3072, 2048, 256, 256, True, False, True): (3, 8, 1, 32), + (3072, 3072, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (3072, 3072, 4096, 32, 32, True, False, True): (1, 32, 1, 4), + (3072, 3072, 4096, 64, 64, False, True, True): (1, 64, 1, 8), + (3072, 3072, 4096, 64, 64, True, False, True): (3, 16, 3, 4), + (3072, 3072, 4096, 128, 128, False, True, True): (1, 32, 1, 8), + (3072, 3072, 4096, 128, 128, True, False, True): (2, 32, 2, 4), + (3072, 3072, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (3072, 3072, 4096, 256, 256, True, False, True): (2, 16, 1, 32), + (3072, 3072, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (3072, 3072, 8192, 32, 32, True, False, True): (1, 64, 1, 2), + (3072, 3072, 8192, 64, 64, False, True, True): (1, 64, 1, 4), + (3072, 3072, 8192, 64, 64, True, False, True): (1, 64, 3, 2), + (3072, 3072, 8192, 128, 128, False, True, True): (1, 64, 1, 8), + (3072, 3072, 8192, 128, 128, True, False, True): (2, 64, 2, 4), + (3072, 3072, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (3072, 3072, 8192, 256, 256, True, False, True): (4, 32, 1, 32), + (3072, 3072, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (3072, 3072, 16384, 32, 32, True, False, True): (1, 128, 3, 2), + (3072, 3072, 16384, 64, 64, False, True, True): (1, 128, 1, 4), + (3072, 3072, 16384, 64, 64, True, False, True): (2, 64, 3, 4), + (3072, 3072, 16384, 128, 128, False, True, True): (1, 128, 1, 8), + (3072, 3072, 16384, 128, 128, True, False, True): (1, 128, 2, 4), + (3072, 3072, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (3072, 3072, 16384, 256, 256, True, False, True): (4, 64, 1, 32), + (3072, 3072, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (3072, 3072, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (3072, 3072, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (3072, 3072, 32768, 64, 64, True, False, True): (1, 256, 3, 2), + (3072, 3072, 32768, 128, 128, False, True, True): (1, 256, 1, 8), + (3072, 3072, 32768, 128, 128, True, False, True): (1, 256, 2, 4), + (3072, 3072, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (3072, 3072, 32768, 256, 256, True, False, True): (4, 128, 1, 32), + (3072, 3072, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (3072, 3072, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (3072, 3072, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (3072, 3072, 65536, 64, 64, True, False, True): (2, 256, 3, 4), + (3072, 3072, 65536, 128, 128, False, True, True): (1, 512, 1, 8), + (3072, 3072, 65536, 128, 128, True, False, True): (1, 512, 3, 4), + (3072, 3072, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (3072, 3072, 65536, 256, 256, True, False, True): (4, 256, 1, 32), + (3072, 3072, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (3072, 3072, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (3072, 3072, 131072, 64, 64, False, True, True): (1, 1024, 1, 4), + (3072, 3072, 131072, 64, 64, True, False, True): (1, 1024, 3, 2), + (3072, 3072, 131072, 128, 128, False, True, True): (1, 1024, 1, 8), + (3072, 3072, 131072, 128, 128, True, False, True): (1, 1024, 3, 4), + (3072, 3072, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (3072, 3072, 131072, 256, 256, True, False, True): (4, 512, 1, 32), + (4096, 4096, 256, 32, 32, False, True, True): (1, 4, 1, 8), + (4096, 4096, 256, 32, 32, True, False, True): (5, 2, 3, 4), + (4096, 4096, 256, 64, 64, False, True, True): (3, 4, 1, 8), + (4096, 4096, 256, 64, 64, True, False, True): (3, 4, 3, 2), + (4096, 4096, 256, 128, 128, False, True, True): (1, 2, 1, 8), + (4096, 4096, 256, 128, 128, True, False, True): (2, 2, 4, 4), + (4096, 4096, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (4096, 4096, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (4096, 4096, 512, 32, 32, False, True, True): (1, 8, 1, 8), + (4096, 4096, 512, 32, 32, True, False, True): (1, 4, 1, 4), + (4096, 4096, 512, 64, 64, False, True, True): (1, 8, 1, 8), + (4096, 4096, 512, 64, 64, True, False, True): (3, 4, 2, 2), + (4096, 4096, 512, 128, 128, False, True, True): (2, 4, 1, 8), + (4096, 4096, 512, 128, 128, True, False, True): (2, 4, 2, 4), + (4096, 4096, 512, 256, 256, False, True, True): (2, 2, 1, 32), + (4096, 4096, 512, 256, 256, True, False, True): (2, 2, 1, 32), + (4096, 4096, 1024, 32, 32, False, True, True): (4, 16, 1, 8), + (4096, 4096, 1024, 32, 32, True, False, True): (1, 8, 1, 4), + (4096, 4096, 1024, 64, 64, False, True, True): (1, 16, 1, 8), + (4096, 4096, 1024, 64, 64, True, False, True): (4, 4, 3, 4), + (4096, 4096, 1024, 128, 128, False, True, True): (2, 8, 1, 8), + (4096, 4096, 1024, 128, 128, True, False, True): (1, 8, 3, 4), + (4096, 4096, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (4096, 4096, 1024, 256, 256, True, False, True): (6, 4, 1, 32), + (4096, 4096, 2048, 32, 32, False, True, True): (1, 32, 1, 8), + (4096, 4096, 2048, 32, 32, True, False, True): (1, 16, 1, 4), + (4096, 4096, 2048, 64, 64, False, True, True): (4, 32, 1, 8), + (4096, 4096, 2048, 64, 64, True, False, True): (4, 8, 3, 4), + (4096, 4096, 2048, 128, 128, False, True, True): (2, 16, 1, 8), + (4096, 4096, 2048, 128, 128, True, False, True): (1, 16, 3, 4), + (4096, 4096, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (4096, 4096, 2048, 256, 256, True, False, True): (4, 8, 1, 32), + (4096, 4096, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (4096, 4096, 4096, 32, 32, True, False, True): (1, 32, 1, 4), + (4096, 4096, 4096, 64, 64, False, True, True): (1, 64, 1, 8), + (4096, 4096, 4096, 64, 64, True, False, True): (1, 32, 3, 2), + (4096, 4096, 4096, 128, 128, False, True, True): (1, 32, 1, 8), + (4096, 4096, 4096, 128, 128, True, False, True): (2, 32, 3, 4), + (4096, 4096, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (4096, 4096, 4096, 256, 256, True, False, True): (4, 16, 1, 32), + (4096, 4096, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (4096, 4096, 8192, 32, 32, True, False, True): (1, 64, 1, 4), + (4096, 4096, 8192, 64, 64, False, True, True): (1, 128, 1, 8), + (4096, 4096, 8192, 64, 64, True, False, True): (1, 64, 3, 2), + (4096, 4096, 8192, 128, 128, False, True, True): (1, 64, 1, 8), + (4096, 4096, 8192, 128, 128, True, False, True): (1, 64, 3, 4), + (4096, 4096, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (4096, 4096, 8192, 256, 256, True, False, True): (4, 32, 1, 32), + (4096, 4096, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (4096, 4096, 16384, 32, 32, True, False, True): (1, 128, 3, 2), + (4096, 4096, 16384, 64, 64, False, True, True): (1, 128, 1, 4), + (4096, 4096, 16384, 64, 64, True, False, True): (4, 64, 3, 4), + (4096, 4096, 16384, 128, 128, False, True, True): (1, 128, 1, 8), + (4096, 4096, 16384, 128, 128, True, False, True): (1, 128, 3, 4), + (4096, 4096, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (4096, 4096, 16384, 256, 256, True, False, True): (4, 64, 1, 32), + (4096, 4096, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (4096, 4096, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (4096, 4096, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (4096, 4096, 32768, 64, 64, True, False, True): (1, 256, 3, 2), + (4096, 4096, 32768, 128, 128, False, True, True): (1, 256, 1, 8), + (4096, 4096, 32768, 128, 128, True, False, True): (1, 256, 3, 4), + (4096, 4096, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (4096, 4096, 32768, 256, 256, True, False, True): (4, 128, 1, 32), + (4096, 4096, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (4096, 4096, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (4096, 4096, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (4096, 4096, 65536, 64, 64, True, False, True): (4, 256, 3, 4), + (4096, 4096, 65536, 128, 128, False, True, True): (1, 512, 1, 8), + (4096, 4096, 65536, 128, 128, True, False, True): (1, 512, 3, 4), + (4096, 4096, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (4096, 4096, 65536, 256, 256, True, False, True): (4, 256, 1, 32), + (4096, 4096, 65792, 32, 32, False, True, True): (1, 1028, 1, 8), + (4096, 4096, 65792, 32, 32, True, False, True): (1, 514, 3, 2), + (4096, 4096, 65792, 64, 64, False, True, True): (1, 1028, 1, 8), + (4096, 4096, 65792, 64, 64, True, False, True): (1, 514, 3, 2), + (4096, 4096, 65792, 128, 128, False, True, True): (1, 514, 1, 8), + (4096, 4096, 65792, 128, 128, True, False, True): (1, 514, 2, 4), + (4096, 4096, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (4096, 4096, 65792, 256, 256, True, False, True): (1, 257, 1, 32), + (4096, 4096, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (4096, 4096, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (4096, 4096, 131072, 64, 64, False, True, True): (1, 2048, 1, 8), + (4096, 4096, 131072, 64, 64, True, False, True): (1, 1024, 3, 2), + (4096, 4096, 131072, 128, 128, False, True, True): (1, 1024, 1, 8), + (4096, 4096, 131072, 128, 128, True, False, True): (1, 1024, 3, 4), + (4096, 4096, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (4096, 4096, 131072, 256, 256, True, False, True): (4, 512, 1, 32), + (5120, 1280, 65792, 32, 32, False, True, True): (1, 1028, 1, 8), + (5120, 1280, 65792, 32, 32, True, False, True): (1, 514, 1, 2), + (5120, 1280, 65792, 64, 64, False, True, True): (1, 514, 1, 4), + (5120, 1280, 65792, 64, 64, True, False, True): (1, 514, 2, 2), + (5120, 1280, 65792, 128, 128, False, True, True): (1, 514, 1, 8), + (5120, 1280, 65792, 128, 128, True, False, True): (1, 514, 2, 4), + (5120, 1280, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (5120, 1280, 65792, 256, 256, True, False, True): (1, 257, 1, 32), + (6144, 6144, 256, 32, 32, False, True, True): (2, 4, 1, 8), + (6144, 6144, 256, 32, 32, True, False, True): (2, 1, 4, 4), + (6144, 6144, 256, 64, 64, False, True, True): (1, 4, 1, 8), + (6144, 6144, 256, 64, 64, True, False, True): (5, 1, 3, 4), + (6144, 6144, 256, 128, 128, False, True, True): (1, 2, 1, 8), + (6144, 6144, 256, 128, 128, True, False, True): (1, 2, 3, 4), + (6144, 6144, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (6144, 6144, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (6144, 6144, 512, 32, 32, False, True, True): (1, 8, 1, 8), + (6144, 6144, 512, 32, 32, True, False, True): (1, 4, 4, 2), + (6144, 6144, 512, 64, 64, False, True, True): (2, 8, 1, 8), + (6144, 6144, 512, 64, 64, True, False, True): (2, 2, 3, 4), + (6144, 6144, 512, 128, 128, False, True, True): (3, 4, 1, 8), + (6144, 6144, 512, 128, 128, True, False, True): (2, 4, 3, 4), + (6144, 6144, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (6144, 6144, 512, 256, 256, True, False, True): (2, 2, 1, 32), + (6144, 6144, 1024, 32, 32, False, True, True): (1, 16, 1, 8), + (6144, 6144, 1024, 32, 32, True, False, True): (1, 8, 1, 4), + (6144, 6144, 1024, 64, 64, False, True, True): (1, 16, 1, 8), + (6144, 6144, 1024, 64, 64, True, False, True): (4, 4, 3, 4), + (6144, 6144, 1024, 128, 128, False, True, True): (1, 8, 1, 8), + (6144, 6144, 1024, 128, 128, True, False, True): (3, 8, 3, 4), + (6144, 6144, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (6144, 6144, 1024, 256, 256, True, False, True): (1, 4, 1, 32), + (6144, 6144, 2048, 32, 32, False, True, True): (1, 32, 1, 8), + (6144, 6144, 2048, 32, 32, True, False, True): (1, 16, 1, 4), + (6144, 6144, 2048, 64, 64, False, True, True): (1, 32, 1, 8), + (6144, 6144, 2048, 64, 64, True, False, True): (4, 8, 3, 4), + (6144, 6144, 2048, 128, 128, False, True, True): (1, 16, 1, 8), + (6144, 6144, 2048, 128, 128, True, False, True): (3, 16, 3, 4), + (6144, 6144, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (6144, 6144, 2048, 256, 256, True, False, True): (4, 8, 1, 32), + (6144, 6144, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (6144, 6144, 4096, 32, 32, True, False, True): (1, 32, 1, 4), + (6144, 6144, 4096, 64, 64, False, True, True): (1, 64, 1, 8), + (6144, 6144, 4096, 64, 64, True, False, True): (4, 16, 3, 4), + (6144, 6144, 4096, 128, 128, False, True, True): (1, 32, 1, 8), + (6144, 6144, 4096, 128, 128, True, False, True): (4, 32, 3, 4), + (6144, 6144, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (6144, 6144, 4096, 256, 256, True, False, True): (4, 16, 1, 32), + (6144, 6144, 8192, 32, 32, False, True, True): (1, 128, 1, 8), + (6144, 6144, 8192, 32, 32, True, False, True): (1, 64, 1, 4), + (6144, 6144, 8192, 64, 64, False, True, True): (1, 128, 1, 8), + (6144, 6144, 8192, 64, 64, True, False, True): (4, 32, 3, 4), + (6144, 6144, 8192, 128, 128, False, True, True): (1, 64, 1, 8), + (6144, 6144, 8192, 128, 128, True, False, True): (1, 64, 3, 4), + (6144, 6144, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (6144, 6144, 8192, 256, 256, True, False, True): (4, 32, 1, 32), + (6144, 6144, 16384, 32, 32, False, True, True): (1, 256, 1, 8), + (6144, 6144, 16384, 32, 32, True, False, True): (1, 128, 1, 4), + (6144, 6144, 16384, 64, 64, False, True, True): (1, 256, 1, 8), + (6144, 6144, 16384, 64, 64, True, False, True): (4, 64, 3, 4), + (6144, 6144, 16384, 128, 128, False, True, True): (1, 128, 1, 8), + (6144, 6144, 16384, 128, 128, True, False, True): (4, 128, 3, 4), + (6144, 6144, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (6144, 6144, 16384, 256, 256, True, False, True): (4, 64, 1, 32), + (6144, 6144, 32768, 32, 32, False, True, True): (1, 512, 1, 8), + (6144, 6144, 32768, 32, 32, True, False, True): (1, 256, 1, 4), + (6144, 6144, 32768, 64, 64, False, True, True): (1, 512, 1, 8), + (6144, 6144, 32768, 64, 64, True, False, True): (4, 128, 3, 4), + (6144, 6144, 32768, 128, 128, False, True, True): (1, 256, 1, 8), + (6144, 6144, 32768, 128, 128, True, False, True): (1, 256, 3, 4), + (6144, 6144, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (6144, 6144, 32768, 256, 256, True, False, True): (4, 128, 1, 32), + (6144, 6144, 65536, 32, 32, False, True, True): (1, 1024, 1, 8), + (6144, 6144, 65536, 32, 32, True, False, True): (1, 512, 1, 4), + (6144, 6144, 65536, 64, 64, False, True, True): (1, 1024, 1, 8), + (6144, 6144, 65536, 64, 64, True, False, True): (4, 256, 3, 4), + (6144, 6144, 65536, 128, 128, False, True, True): (1, 512, 1, 8), + (6144, 6144, 65536, 128, 128, True, False, True): (1, 512, 3, 4), + (6144, 6144, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (6144, 6144, 65536, 256, 256, True, False, True): (4, 256, 1, 32), + (6144, 6144, 131072, 32, 32, False, True, True): (1, 2048, 1, 8), + (6144, 6144, 131072, 32, 32, True, False, True): (1, 1024, 1, 4), + (6144, 6144, 131072, 64, 64, False, True, True): (1, 2048, 1, 8), + (6144, 6144, 131072, 64, 64, True, False, True): (4, 512, 3, 4), + (6144, 6144, 131072, 128, 128, False, True, True): (1, 1024, 1, 8), + (6144, 6144, 131072, 128, 128, True, False, True): (1, 1024, 3, 4), + (6144, 6144, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (6144, 6144, 131072, 256, 256, True, False, True): (4, 512, 1, 32), + (8192, 8192, 256, 32, 32, False, True, True): (1, 4, 1, 8), + (8192, 8192, 256, 32, 32, True, False, True): (3, 2, 3, 4), + (8192, 8192, 256, 64, 64, False, True, True): (1, 4, 1, 4), + (8192, 8192, 256, 64, 64, True, False, True): (1, 4, 1, 4), + (8192, 8192, 256, 128, 128, False, True, True): (1, 2, 1, 8), + (8192, 8192, 256, 128, 128, True, False, True): (2, 2, 3, 4), + (8192, 8192, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (8192, 8192, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (8192, 8192, 512, 32, 32, False, True, True): (4, 8, 1, 8), + (8192, 8192, 512, 32, 32, True, False, True): (2, 4, 4, 2), + (8192, 8192, 512, 64, 64, False, True, True): (4, 4, 1, 4), + (8192, 8192, 512, 64, 64, True, False, True): (3, 2, 3, 4), + (8192, 8192, 512, 128, 128, False, True, True): (1, 4, 1, 8), + (8192, 8192, 512, 128, 128, True, False, True): (1, 4, 3, 4), + (8192, 8192, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (8192, 8192, 512, 256, 256, True, False, True): (1, 2, 1, 32), + (8192, 8192, 1024, 32, 32, False, True, True): (4, 16, 1, 8), + (8192, 8192, 1024, 32, 32, True, False, True): (1, 8, 3, 2), + (8192, 8192, 1024, 64, 64, False, True, True): (4, 8, 1, 4), + (8192, 8192, 1024, 64, 64, True, False, True): (4, 4, 3, 4), + (8192, 8192, 1024, 128, 128, False, True, True): (1, 8, 1, 8), + (8192, 8192, 1024, 128, 128, True, False, True): (1, 8, 3, 4), + (8192, 8192, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (8192, 8192, 1024, 256, 256, True, False, True): (4, 4, 1, 32), + (8192, 8192, 2048, 32, 32, False, True, True): (4, 32, 1, 8), + (8192, 8192, 2048, 32, 32, True, False, True): (1, 16, 3, 2), + (8192, 8192, 2048, 64, 64, False, True, True): (4, 32, 1, 8), + (8192, 8192, 2048, 64, 64, True, False, True): (4, 8, 3, 4), + (8192, 8192, 2048, 128, 128, False, True, True): (4, 16, 1, 8), + (8192, 8192, 2048, 128, 128, True, False, True): (4, 16, 3, 4), + (8192, 8192, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (8192, 8192, 2048, 256, 256, True, False, True): (4, 8, 1, 32), + (8192, 8192, 4096, 32, 32, False, True, True): (4, 64, 1, 8), + (8192, 8192, 4096, 32, 32, True, False, True): (2, 32, 3, 2), + (8192, 8192, 4096, 64, 64, False, True, True): (4, 64, 1, 8), + (8192, 8192, 4096, 64, 64, True, False, True): (4, 16, 3, 4), + (8192, 8192, 4096, 128, 128, False, True, True): (4, 32, 1, 8), + (8192, 8192, 4096, 128, 128, True, False, True): (4, 32, 3, 4), + (8192, 8192, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (8192, 8192, 4096, 256, 256, True, False, True): (2, 16, 1, 32), + (8192, 8192, 8192, 32, 32, False, True, True): (4, 128, 1, 8), + (8192, 8192, 8192, 32, 32, True, False, True): (1, 64, 3, 2), + (8192, 8192, 8192, 64, 64, False, True, True): (4, 64, 1, 4), + (8192, 8192, 8192, 64, 64, True, False, True): (4, 32, 3, 4), + (8192, 8192, 8192, 128, 128, False, True, True): (4, 64, 1, 16), + (8192, 8192, 8192, 128, 128, True, False, True): (4, 64, 3, 4), + (8192, 8192, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (8192, 8192, 8192, 256, 256, True, False, True): (4, 32, 1, 32), + (8192, 8192, 16384, 32, 32, False, True, True): (4, 256, 1, 8), + (8192, 8192, 16384, 32, 32, True, False, True): (4, 128, 3, 2), + (8192, 8192, 16384, 64, 64, False, True, True): (4, 128, 1, 4), + (8192, 8192, 16384, 64, 64, True, False, True): (4, 64, 3, 4), + (8192, 8192, 16384, 128, 128, False, True, True): (4, 128, 1, 16), + (8192, 8192, 16384, 128, 128, True, False, True): (4, 128, 3, 4), + (8192, 8192, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (8192, 8192, 16384, 256, 256, True, False, True): (4, 64, 1, 32), + (8192, 8192, 32768, 32, 32, False, True, True): (4, 512, 1, 8), + (8192, 8192, 32768, 32, 32, True, False, True): (2, 256, 3, 2), + (8192, 8192, 32768, 64, 64, False, True, True): (4, 256, 1, 4), + (8192, 8192, 32768, 64, 64, True, False, True): (4, 128, 3, 4), + (8192, 8192, 32768, 128, 128, False, True, True): (4, 256, 1, 16), + (8192, 8192, 32768, 128, 128, True, False, True): (4, 256, 3, 4), + (8192, 8192, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (8192, 8192, 32768, 256, 256, True, False, True): (4, 128, 1, 32), + (8192, 8192, 65536, 32, 32, False, True, True): (4, 1024, 1, 8), + (8192, 8192, 65536, 32, 32, True, False, True): (4, 512, 3, 2), + (8192, 8192, 65536, 64, 64, False, True, True): (4, 512, 1, 4), + (8192, 8192, 65536, 64, 64, True, False, True): (4, 256, 3, 4), + (8192, 8192, 65536, 128, 128, False, True, True): (4, 512, 1, 16), + (8192, 8192, 65536, 128, 128, True, False, True): (4, 512, 3, 4), + (8192, 8192, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (8192, 8192, 65536, 256, 256, True, False, True): (4, 256, 1, 32), + (8192, 8192, 65792, 32, 32, False, True, True): (4, 1028, 1, 8), + (8192, 8192, 65792, 32, 32, True, False, True): (1, 514, 3, 2), + (8192, 8192, 65792, 64, 64, False, True, True): (4, 1028, 1, 8), + (8192, 8192, 65792, 64, 64, True, False, True): (2, 257, 3, 4), + (8192, 8192, 65792, 128, 128, False, True, True): (4, 514, 1, 16), + (8192, 8192, 65792, 128, 128, True, False, True): (2, 514, 3, 4), + (8192, 8192, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (8192, 8192, 65792, 256, 256, True, False, True): (1, 257, 1, 32), + (8192, 8192, 131072, 32, 32, False, True, True): (4, 2048, 1, 8), + (8192, 8192, 131072, 32, 32, True, False, True): (4, 1024, 3, 2), + (8192, 8192, 131072, 64, 64, False, True, True): (4, 1024, 1, 4), + (8192, 8192, 131072, 64, 64, True, False, True): (4, 512, 3, 4), + (8192, 8192, 131072, 128, 128, False, True, True): (4, 1024, 1, 16), + (8192, 8192, 131072, 128, 128, True, False, True): (4, 1024, 3, 4), + (8192, 8192, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (8192, 8192, 131072, 256, 256, True, False, True): (4, 512, 1, 32), + (16384, 16384, 256, 32, 32, False, True, True): (4, 4, 1, 8), + (16384, 16384, 256, 32, 32, True, False, True): (2, 2, 4, 2), + (16384, 16384, 256, 64, 64, False, True, True): (2, 2, 1, 4), + (16384, 16384, 256, 64, 64, True, False, True): (5, 1, 3, 4), + (16384, 16384, 256, 128, 128, False, True, True): (6, 2, 1, 8), + (16384, 16384, 256, 128, 128, True, False, True): (6, 2, 3, 4), + (16384, 16384, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (16384, 16384, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (16384, 16384, 512, 32, 32, False, True, True): (4, 8, 1, 8), + (16384, 16384, 512, 32, 32, True, False, True): (1, 4, 4, 2), + (16384, 16384, 512, 64, 64, False, True, True): (4, 4, 1, 4), + (16384, 16384, 512, 64, 64, True, False, True): (2, 2, 3, 4), + (16384, 16384, 512, 128, 128, False, True, True): (4, 4, 1, 8), + (16384, 16384, 512, 128, 128, True, False, True): (4, 4, 3, 4), + (16384, 16384, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (16384, 16384, 512, 256, 256, True, False, True): (2, 2, 1, 32), + (16384, 16384, 1024, 32, 32, False, True, True): (4, 16, 1, 8), + (16384, 16384, 1024, 32, 32, True, False, True): (1, 8, 3, 2), + (16384, 16384, 1024, 64, 64, False, True, True): (4, 8, 1, 4), + (16384, 16384, 1024, 64, 64, True, False, True): (4, 4, 3, 4), + (16384, 16384, 1024, 128, 128, False, True, True): (4, 4, 1, 8), + (16384, 16384, 1024, 128, 128, True, False, True): (4, 8, 3, 4), + (16384, 16384, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (16384, 16384, 1024, 256, 256, True, False, True): (4, 4, 1, 32), + (16384, 16384, 2048, 32, 32, False, True, True): (4, 32, 1, 8), + (16384, 16384, 2048, 32, 32, True, False, True): (2, 16, 3, 2), + (16384, 16384, 2048, 64, 64, False, True, True): (4, 16, 1, 4), + (16384, 16384, 2048, 64, 64, True, False, True): (4, 8, 3, 4), + (16384, 16384, 2048, 128, 128, False, True, True): (4, 16, 1, 8), + (16384, 16384, 2048, 128, 128, True, False, True): (4, 16, 3, 4), + (16384, 16384, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (16384, 16384, 2048, 256, 256, True, False, True): (4, 8, 1, 32), + (16384, 16384, 4096, 32, 32, False, True, True): (4, 64, 1, 8), + (16384, 16384, 4096, 32, 32, True, False, True): (2, 32, 3, 2), + (16384, 16384, 4096, 64, 64, False, True, True): (2, 32, 1, 4), + (16384, 16384, 4096, 64, 64, True, False, True): (4, 16, 3, 4), + (16384, 16384, 4096, 128, 128, False, True, True): (4, 32, 1, 8), + (16384, 16384, 4096, 128, 128, True, False, True): (4, 32, 3, 4), + (16384, 16384, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (16384, 16384, 4096, 256, 256, True, False, True): (4, 16, 1, 32), + (16384, 16384, 8192, 32, 32, False, True, True): (4, 128, 1, 8), + (16384, 16384, 8192, 32, 32, True, False, True): (2, 64, 3, 2), + (16384, 16384, 8192, 64, 64, False, True, True): (4, 64, 1, 4), + (16384, 16384, 8192, 64, 64, True, False, True): (4, 32, 3, 4), + (16384, 16384, 8192, 128, 128, False, True, True): (4, 64, 1, 16), + (16384, 16384, 8192, 128, 128, True, False, True): (4, 64, 3, 4), + (16384, 16384, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (16384, 16384, 8192, 256, 256, True, False, True): (4, 32, 1, 32), + (16384, 16384, 16384, 32, 32, False, True, True): (4, 256, 1, 8), + (16384, 16384, 16384, 32, 32, True, False, True): (2, 128, 3, 2), + (16384, 16384, 16384, 64, 64, False, True, True): (4, 128, 1, 4), + (16384, 16384, 16384, 64, 64, True, False, True): (4, 64, 3, 4), + (16384, 16384, 16384, 128, 128, False, True, True): (1, 64, 1, 8), + (16384, 16384, 16384, 128, 128, True, False, True): (4, 128, 3, 4), + (16384, 16384, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (16384, 16384, 16384, 256, 256, True, False, True): (4, 64, 1, 32), + (16384, 16384, 32768, 32, 32, False, True, True): (4, 512, 1, 8), + (16384, 16384, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (16384, 16384, 32768, 64, 64, False, True, True): (4, 256, 1, 4), + (16384, 16384, 32768, 64, 64, True, False, True): (4, 128, 3, 4), + (16384, 16384, 32768, 128, 128, False, True, True): (4, 256, 1, 16), + (16384, 16384, 32768, 128, 128, True, False, True): (4, 256, 3, 4), + (16384, 16384, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (16384, 16384, 32768, 256, 256, True, False, True): (4, 128, 1, 32), + (16384, 16384, 65536, 32, 32, False, True, True): (4, 1024, 1, 8), + (16384, 16384, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (16384, 16384, 65536, 64, 64, False, True, True): (2, 512, 1, 4), + (16384, 16384, 65536, 64, 64, True, False, True): (4, 256, 3, 4), + (16384, 16384, 65536, 128, 128, False, True, True): (4, 512, 1, 16), + (16384, 16384, 65536, 128, 128, True, False, True): (4, 512, 3, 4), + (16384, 16384, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (16384, 16384, 65536, 256, 256, True, False, True): (4, 256, 1, 32), + (16384, 16384, 65792, 32, 32, False, True, True): (4, 1028, 1, 8), + (16384, 16384, 65792, 32, 32, True, False, True): (1, 514, 3, 2), + (16384, 16384, 65792, 64, 64, False, True, True): (2, 514, 1, 4), + (16384, 16384, 65792, 64, 64, True, False, True): (2, 257, 3, 4), + (16384, 16384, 65792, 128, 128, False, True, True): (2, 514, 1, 16), + (16384, 16384, 65792, 128, 128, True, False, True): (2, 514, 3, 4), + (16384, 16384, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (16384, 16384, 65792, 256, 256, True, False, True): (1, 257, 1, 32), + (16384, 16384, 131072, 32, 32, False, True, True): (4, 1024, 1, 8), + (16384, 16384, 131072, 32, 32, True, False, True): (4, 512, 3, 4), + (16384, 16384, 131072, 64, 64, False, True, True): (4, 1024, 1, 4), + (16384, 16384, 131072, 64, 64, True, False, True): (4, 1024, 3, 2), + (16384, 16384, 131072, 128, 128, False, True, True): (2, 1024, 3, 8), + (16384, 16384, 131072, 128, 128, True, False, True): (4, 1024, 3, 4), + (16384, 16384, 131072, 256, 256, False, True, True): (4, 512, 1, 32), + (16384, 16384, 131072, 256, 256, True, False, True): (4, 512, 1, 32), + (32768, 32768, 256, 32, 32, False, True, True): (4, 4, 1, 8), + (32768, 32768, 256, 32, 32, True, False, True): (1, 2, 4, 2), + (32768, 32768, 256, 64, 64, False, True, True): (2, 2, 1, 4), + (32768, 32768, 256, 64, 64, True, False, True): (2, 1, 3, 4), + (32768, 32768, 256, 128, 128, False, True, True): (4, 2, 1, 8), + (32768, 32768, 256, 128, 128, True, False, True): (4, 2, 3, 4), + (32768, 32768, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (32768, 32768, 256, 256, 256, True, False, True): (1, 1, 1, 32), + (32768, 32768, 512, 32, 32, False, True, True): (4, 8, 1, 8), + (32768, 32768, 512, 32, 32, True, False, True): (1, 4, 3, 2), + (32768, 32768, 512, 64, 64, False, True, True): (4, 4, 1, 4), + (32768, 32768, 512, 64, 64, True, False, True): (4, 2, 3, 4), + (32768, 32768, 512, 128, 128, False, True, True): (1, 2, 1, 8), + (32768, 32768, 512, 128, 128, True, False, True): (4, 4, 3, 4), + (32768, 32768, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (32768, 32768, 512, 256, 256, True, False, True): (2, 2, 1, 32), + (32768, 32768, 1024, 32, 32, False, True, True): (4, 16, 1, 8), + (32768, 32768, 1024, 32, 32, True, False, True): (1, 8, 4, 2), + (32768, 32768, 1024, 64, 64, False, True, True): (4, 8, 1, 4), + (32768, 32768, 1024, 64, 64, True, False, True): (4, 4, 3, 4), + (32768, 32768, 1024, 128, 128, False, True, True): (1, 4, 1, 8), + (32768, 32768, 1024, 128, 128, True, False, True): (4, 8, 3, 4), + (32768, 32768, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (32768, 32768, 1024, 256, 256, True, False, True): (1, 4, 1, 32), + (32768, 32768, 2048, 32, 32, False, True, True): (2, 32, 1, 8), + (32768, 32768, 2048, 32, 32, True, False, True): (1, 16, 4, 2), + (32768, 32768, 2048, 64, 64, False, True, True): (2, 16, 1, 4), + (32768, 32768, 2048, 64, 64, True, False, True): (4, 8, 3, 4), + (32768, 32768, 2048, 128, 128, False, True, True): (1, 8, 1, 8), + (32768, 32768, 2048, 128, 128, True, False, True): (4, 16, 3, 4), + (32768, 32768, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (32768, 32768, 2048, 256, 256, True, False, True): (4, 8, 1, 32), + (32768, 32768, 4096, 32, 32, False, True, True): (2, 64, 1, 8), + (32768, 32768, 4096, 32, 32, True, False, True): (2, 32, 3, 2), + (32768, 32768, 4096, 64, 64, False, True, True): (2, 32, 1, 4), + (32768, 32768, 4096, 64, 64, True, False, True): (2, 16, 3, 4), + (32768, 32768, 4096, 128, 128, False, True, True): (1, 16, 1, 8), + (32768, 32768, 4096, 128, 128, True, False, True): (2, 32, 3, 4), + (32768, 32768, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (32768, 32768, 4096, 256, 256, True, False, True): (4, 16, 1, 32), + (32768, 32768, 8192, 32, 32, False, True, True): (2, 128, 1, 8), + (32768, 32768, 8192, 32, 32, True, False, True): (2, 64, 3, 2), + (32768, 32768, 8192, 64, 64, False, True, True): (2, 64, 1, 4), + (32768, 32768, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (32768, 32768, 8192, 128, 128, False, True, True): (1, 32, 1, 8), + (32768, 32768, 8192, 128, 128, True, False, True): (4, 64, 3, 4), + (32768, 32768, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (32768, 32768, 8192, 256, 256, True, False, True): (4, 32, 1, 32), + (32768, 32768, 16384, 32, 32, False, True, True): (2, 256, 1, 8), + (32768, 32768, 16384, 32, 32, True, False, True): (2, 128, 4, 2), + (32768, 32768, 16384, 64, 64, False, True, True): (2, 128, 1, 4), + (32768, 32768, 16384, 64, 64, True, False, True): (4, 64, 3, 4), + (32768, 32768, 16384, 128, 128, False, True, True): (1, 64, 1, 8), + (32768, 32768, 16384, 128, 128, True, False, True): (4, 128, 3, 4), + (32768, 32768, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (32768, 32768, 16384, 256, 256, True, False, True): (2, 64, 1, 32), + (32768, 32768, 32768, 32, 32, False, True, True): (2, 512, 1, 8), + (32768, 32768, 32768, 32, 32, True, False, True): (4, 256, 3, 2), + (32768, 32768, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (32768, 32768, 32768, 64, 64, True, False, True): (2, 128, 3, 4), + (32768, 32768, 32768, 128, 128, False, True, True): (1, 128, 1, 8), + (32768, 32768, 32768, 128, 128, True, False, True): (2, 256, 3, 4), + (32768, 32768, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (32768, 32768, 32768, 256, 256, True, False, True): (1, 128, 1, 32), + (32768, 32768, 65536, 32, 32, False, True, True): (2, 512, 1, 8), + (32768, 32768, 65536, 32, 32, True, False, True): (3, 512, 4, 2), + (32768, 32768, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (32768, 32768, 65536, 64, 64, True, False, True): (2, 512, 3, 2), + (32768, 32768, 65536, 128, 128, False, True, True): (1, 256, 1, 8), + (32768, 32768, 65536, 128, 128, True, False, True): (2, 512, 3, 4), + (32768, 32768, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (32768, 32768, 65536, 256, 256, True, False, True): (1, 256, 1, 32), + }, + ("_int_bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.int8, 0.56)): { + (192, 192, 256, 64, 64, False, True, True): (3, 4, 3, 32), + (192, 192, 256, 64, 64, True, False, True): (1, 4, 3, 4), + (192, 192, 512, 64, 64, False, True, True): (1, 8, 1, 16), + (192, 192, 512, 64, 64, True, False, True): (1, 8, 5, 4), + (192, 192, 1024, 64, 64, False, True, True): (4, 16, 1, 16), + (192, 192, 1024, 64, 64, True, False, True): (3, 16, 3, 4), + (192, 192, 2048, 64, 64, False, True, True): (5, 32, 1, 8), + (192, 192, 2048, 64, 64, True, False, True): (2, 32, 4, 4), + (192, 192, 4096, 64, 64, False, True, True): (4, 64, 1, 16), + (192, 192, 4096, 64, 64, True, False, True): (1, 32, 4, 4), + (192, 192, 8192, 64, 64, False, True, True): (2, 128, 1, 8), + (192, 192, 8192, 64, 64, True, False, True): (3, 64, 1, 4), + (192, 192, 16384, 64, 64, False, True, True): (2, 256, 1, 8), + (192, 192, 16384, 64, 64, True, False, True): (1, 128, 3, 2), + (192, 192, 32768, 64, 64, False, True, True): (2, 512, 1, 8), + (192, 192, 32768, 64, 64, True, False, True): (3, 128, 1, 4), + (192, 192, 65536, 64, 64, False, True, True): (3, 1024, 1, 8), + (192, 192, 65536, 64, 64, True, False, True): (1, 512, 3, 4), + (192, 192, 131072, 64, 64, False, True, True): (1, 2048, 1, 8), + (192, 192, 131072, 64, 64, True, False, True): (1, 512, 1, 4), + (384, 384, 256, 128, 128, False, True, True): (4, 2, 1, 16), + (384, 384, 256, 128, 128, True, False, True): (1, 2, 3, 4), + (384, 384, 512, 128, 128, False, True, True): (2, 4, 1, 16), + (384, 384, 512, 128, 128, True, False, True): (2, 4, 3, 4), + (384, 384, 1024, 128, 128, False, True, True): (3, 8, 1, 32), + (384, 384, 1024, 128, 128, True, False, True): (3, 8, 3, 4), + (384, 384, 2048, 128, 128, False, True, True): (3, 16, 1, 32), + (384, 384, 2048, 128, 128, True, False, True): (2, 16, 3, 4), + (384, 384, 4096, 128, 128, False, True, True): (3, 32, 1, 32), + (384, 384, 4096, 128, 128, True, False, True): (3, 32, 3, 4), + (384, 384, 8192, 128, 128, False, True, True): (2, 64, 1, 32), + (384, 384, 8192, 128, 128, True, False, True): (4, 64, 1, 4), + (384, 384, 16384, 128, 128, False, True, True): (2, 128, 1, 32), + (384, 384, 16384, 128, 128, True, False, True): (2, 128, 1, 4), + (384, 384, 32768, 128, 128, False, True, True): (3, 256, 1, 16), + (384, 384, 32768, 128, 128, True, False, True): (1, 256, 1, 4), + (384, 384, 65536, 128, 128, False, True, True): (4, 512, 1, 16), + (384, 384, 65536, 128, 128, True, False, True): (1, 512, 1, 4), + (384, 384, 131072, 128, 128, False, True, True): (4, 1024, 1, 16), + (384, 384, 131072, 128, 128, True, False, True): (1, 1024, 1, 4), + (768, 768, 256, 256, 256, False, True, True): (1, 1, 1, 32), + (768, 768, 256, 256, 256, True, False, True): (3, 1, 1, 32), + (768, 768, 512, 256, 256, False, True, True): (1, 2, 1, 32), + (768, 768, 512, 256, 256, True, False, True): (1, 2, 1, 32), + (768, 768, 1024, 256, 256, False, True, True): (1, 4, 1, 32), + (768, 768, 1024, 256, 256, True, False, True): (2, 4, 1, 32), + (768, 768, 2048, 256, 256, False, True, True): (1, 8, 1, 32), + (768, 768, 2048, 256, 256, True, False, True): (2, 8, 1, 32), + (768, 768, 4096, 256, 256, False, True, True): (1, 16, 1, 32), + (768, 768, 4096, 256, 256, True, False, True): (1, 16, 1, 32), + (768, 768, 8192, 256, 256, False, True, True): (1, 32, 1, 32), + (768, 768, 8192, 256, 256, True, False, True): (2, 32, 1, 32), + (768, 768, 16384, 256, 256, False, True, True): (1, 64, 1, 32), + (768, 768, 16384, 256, 256, True, False, True): (7, 64, 1, 32), + (768, 768, 32768, 256, 256, False, True, True): (1, 128, 1, 32), + (768, 768, 32768, 256, 256, True, False, True): (1, 128, 1, 32), + (768, 768, 65536, 256, 256, False, True, True): (1, 256, 1, 32), + (768, 768, 65536, 256, 256, True, False, True): (1, 256, 1, 32), + (768, 768, 131072, 256, 256, False, True, True): (1, 512, 1, 32), + (768, 768, 131072, 256, 256, True, False, True): (1, 512, 1, 32), + }, + ("_int_bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.int8, 1.0)): { + (256, 256, 256, 256, 256, False, True, True): (2, 1, 1, 4), + (256, 256, 256, 256, 256, True, False, True): (2, 1, 2, 1), + (256, 256, 512, 256, 256, False, True, True): (2, 1, 1, 2), + (256, 256, 512, 256, 256, True, False, True): (2, 2, 2, 8), + (256, 256, 1024, 256, 256, False, True, True): (1, 4, 1, 4), + (256, 256, 1024, 256, 256, True, False, True): (1, 2, 2, 4), + (256, 256, 2048, 256, 256, False, True, True): (1, 4, 1, 2), + (256, 256, 2048, 256, 256, True, False, True): (1, 8, 1, 2), + (256, 256, 4096, 256, 256, False, True, True): (1, 16, 1, 4), + (256, 256, 4096, 256, 256, True, False, True): (1, 16, 1, 2), + (256, 256, 8192, 256, 256, False, True, True): (1, 16, 3, 4), + (256, 256, 8192, 256, 256, True, False, True): (1, 8, 1, 4), + (256, 256, 16384, 256, 256, False, True, True): (2, 16, 1, 8), + (256, 256, 16384, 256, 256, True, False, True): (1, 32, 1, 2), + (256, 256, 32768, 256, 256, False, True, True): (1, 128, 1, 8), + (256, 256, 32768, 256, 256, True, False, True): (1, 128, 1, 4), + (256, 256, 65536, 256, 256, False, True, True): (1, 4, 1, 1), + (256, 256, 65536, 256, 256, True, False, True): (1, 128, 1, 4), + (256, 256, 65792, 256, 256, False, True, True): (1, 128, 2, 16), + (256, 256, 65792, 256, 256, True, False, True): (1, 16, 3, 4), + (256, 256, 131072, 256, 256, False, True, True): (1, 512, 1, 4), + (256, 256, 131072, 256, 256, True, False, True): (1, 512, 1, 2), + }, + ("bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.bfloat16, 0.5)): { + (16, 16, 16, 16, 16, False, False, False): (2, 1, 1, 2), + (16, 16, 16, 16, 16, False, False, True): (1, 1, 1, 4), + (16, 16, 16, 16, 16, False, True, False): (1, 1, 3, 16), + (16, 16, 16, 16, 16, False, True, True): (1, 1, 1, 8), + (16, 16, 16, 16, 16, True, False, False): (2, 1, 1, 8), + (16, 16, 16, 16, 16, True, False, True): (1, 1, 1, 8), + (16, 16, 32, 16, 16, False, False, False): (1, 2, 1, 8), + (16, 16, 32, 16, 16, False, False, True): (1, 2, 2, 4), + (16, 16, 32, 16, 16, False, True, False): (1, 1, 2, 4), + (16, 16, 32, 16, 16, False, True, True): (1, 1, 2, 4), + (16, 16, 32, 16, 16, True, False, False): (1, 1, 2, 4), + (16, 16, 32, 16, 16, True, False, True): (2, 2, 1, 2), + (16, 16, 64, 16, 16, False, False, False): (1, 4, 2, 4), + (16, 16, 64, 16, 16, False, False, True): (1, 2, 1, 2), + (16, 16, 64, 16, 16, False, True, False): (2, 1, 1, 2), + (16, 16, 64, 16, 16, False, True, True): (1, 4, 1, 8), + (16, 16, 64, 16, 16, True, False, False): (1, 4, 1, 1), + (16, 16, 64, 16, 16, True, False, True): (1, 4, 2, 4), + (16, 32, 16, 16, 16, False, False, False): (1, 1, 2, 2), + (16, 32, 16, 16, 16, False, False, True): (1, 1, 1, 4), + (16, 32, 16, 16, 16, False, True, False): (1, 1, 1, 2), + (16, 32, 16, 16, 16, False, True, True): (1, 1, 1, 1), + (16, 32, 16, 16, 16, True, False, False): (1, 1, 1, 2), + (16, 32, 16, 16, 16, True, False, True): (2, 1, 1, 2), + (16, 32, 16, 16, 32, False, False, False): (1, 1, 1, 4), + (16, 32, 16, 16, 32, False, False, True): (1, 1, 1, 8), + (16, 32, 16, 16, 32, False, True, False): (1, 1, 1, 8), + (16, 32, 16, 16, 32, False, True, True): (1, 1, 2, 4), + (16, 32, 16, 16, 32, True, False, False): (1, 1, 1, 2), + (16, 32, 16, 16, 32, True, False, True): (1, 1, 1, 1), + (16, 32, 32, 16, 16, False, False, False): (2, 2, 1, 4), + (16, 32, 32, 16, 16, False, False, True): (2, 2, 1, 2), + (16, 32, 32, 16, 16, False, True, False): (1, 1, 2, 8), + (16, 32, 32, 16, 16, False, True, True): (1, 2, 1, 1), + (16, 32, 32, 16, 16, True, False, False): (1, 1, 1, 8), + (16, 32, 32, 16, 16, True, False, True): (1, 2, 1, 4), + (16, 32, 32, 16, 32, False, False, False): (1, 1, 2, 8), + (16, 32, 32, 16, 32, False, False, True): (2, 1, 1, 8), + (16, 32, 32, 16, 32, False, True, False): (1, 1, 1, 4), + (16, 32, 32, 16, 32, False, True, True): (1, 1, 1, 4), + (16, 32, 32, 16, 32, True, False, False): (1, 2, 1, 8), + (16, 32, 32, 16, 32, True, False, True): (1, 1, 1, 4), + (16, 32, 64, 16, 16, False, False, False): (1, 4, 3, 8), + (16, 32, 64, 16, 16, False, False, True): (1, 4, 1, 4), + (16, 32, 64, 16, 16, False, True, False): (1, 4, 1, 4), + (16, 32, 64, 16, 16, False, True, True): (2, 4, 1, 4), + (16, 32, 64, 16, 16, True, False, False): (1, 2, 1, 4), + (16, 32, 64, 16, 16, True, False, True): (1, 2, 1, 4), + (16, 32, 64, 16, 32, False, False, False): (1, 4, 1, 8), + (16, 32, 64, 16, 32, False, False, True): (1, 4, 1, 4), + (16, 32, 64, 16, 32, False, True, False): (1, 4, 1, 2), + (16, 32, 64, 16, 32, False, True, True): (1, 2, 1, 4), + (16, 32, 64, 16, 32, True, False, False): (1, 2, 1, 4), + (16, 32, 64, 16, 32, True, False, True): (1, 2, 1, 2), + (16, 64, 16, 16, 32, False, False, False): (1, 1, 1, 2), + (16, 64, 16, 16, 32, False, False, True): (1, 1, 2, 2), + (16, 64, 16, 16, 32, False, True, False): (1, 1, 2, 8), + (16, 64, 16, 16, 32, False, True, True): (1, 1, 1, 4), + (16, 64, 16, 16, 32, True, False, False): (1, 1, 1, 8), + (16, 64, 16, 16, 32, True, False, True): (1, 1, 1, 4), + (16, 64, 32, 16, 32, False, False, False): (1, 2, 1, 2), + (16, 64, 32, 16, 32, False, False, True): (1, 2, 1, 4), + (16, 64, 32, 16, 32, False, True, False): (1, 2, 1, 4), + (16, 64, 32, 16, 32, False, True, True): (2, 2, 1, 4), + (16, 64, 32, 16, 32, True, False, False): (1, 2, 1, 4), + (16, 64, 32, 16, 32, True, False, True): (1, 2, 1, 8), + (16, 64, 64, 16, 32, False, False, False): (1, 2, 1, 4), + (16, 64, 64, 16, 32, False, False, True): (1, 4, 2, 2), + (16, 64, 64, 16, 32, False, True, False): (1, 1, 1, 4), + (16, 64, 64, 16, 32, False, True, True): (1, 4, 1, 2), + (16, 64, 64, 16, 32, True, False, False): (1, 2, 1, 4), + (16, 64, 64, 16, 32, True, False, True): (1, 4, 1, 4), + (32, 16, 16, 16, 16, False, False, False): (1, 1, 1, 8), + (32, 16, 16, 16, 16, False, False, True): (1, 1, 2, 4), + (32, 16, 16, 16, 16, False, True, False): (1, 1, 1, 4), + (32, 16, 16, 16, 16, False, True, True): (1, 1, 2, 4), + (32, 16, 16, 16, 16, True, False, False): (1, 1, 1, 2), + (32, 16, 16, 16, 16, True, False, True): (1, 1, 1, 4), + (32, 16, 32, 16, 16, False, False, False): (1, 1, 1, 4), + (32, 16, 32, 16, 16, False, False, True): (2, 2, 1, 4), + (32, 16, 32, 16, 16, False, True, False): (1, 2, 2, 2), + (32, 16, 32, 16, 16, False, True, True): (2, 2, 1, 4), + (32, 16, 32, 16, 16, True, False, False): (1, 2, 2, 8), + (32, 16, 32, 16, 16, True, False, True): (1, 2, 1, 2), + (32, 16, 64, 16, 16, False, False, False): (1, 4, 1, 4), + (32, 16, 64, 16, 16, False, False, True): (1, 4, 2, 4), + (32, 16, 64, 16, 16, False, True, False): (1, 2, 2, 2), + (32, 16, 64, 16, 16, False, True, True): (3, 4, 1, 4), + (32, 16, 64, 16, 16, True, False, False): (1, 2, 1, 2), + (32, 16, 64, 16, 16, True, False, True): (1, 2, 1, 4), + (32, 32, 16, 16, 16, False, False, False): (1, 1, 3, 4), + (32, 32, 16, 16, 16, False, False, True): (1, 1, 1, 4), + (32, 32, 16, 16, 16, False, True, False): (1, 1, 1, 2), + (32, 32, 16, 16, 16, False, True, True): (1, 1, 1, 4), + (32, 32, 16, 16, 16, True, False, False): (1, 1, 1, 4), + (32, 32, 16, 16, 16, True, False, True): (1, 1, 2, 2), + (32, 32, 16, 16, 32, False, False, False): (2, 1, 1, 4), + (32, 32, 16, 16, 32, False, False, True): (1, 1, 1, 4), + (32, 32, 16, 16, 32, False, True, False): (1, 1, 1, 4), + (32, 32, 16, 16, 32, False, True, True): (3, 1, 2, 4), + (32, 32, 16, 16, 32, True, False, False): (1, 1, 1, 4), + (32, 32, 16, 16, 32, True, False, True): (1, 1, 1, 4), + (32, 32, 16, 32, 32, False, False, False): (1, 1, 1, 8), + (32, 32, 16, 32, 32, False, False, True): (1, 1, 1, 4), + (32, 32, 16, 32, 32, False, True, False): (1, 1, 2, 1), + (32, 32, 16, 32, 32, False, True, True): (2, 1, 2, 2), + (32, 32, 16, 32, 32, True, False, False): (1, 1, 1, 8), + (32, 32, 16, 32, 32, True, False, True): (2, 1, 3, 4), + (32, 32, 32, 16, 16, False, False, False): (1, 2, 1, 4), + (32, 32, 32, 16, 16, False, False, True): (2, 2, 1, 4), + (32, 32, 32, 16, 16, False, True, False): (1, 1, 1, 8), + (32, 32, 32, 16, 16, False, True, True): (2, 2, 1, 4), + (32, 32, 32, 16, 16, True, False, False): (1, 1, 1, 4), + (32, 32, 32, 16, 16, True, False, True): (2, 2, 2, 4), + (32, 32, 32, 16, 32, False, False, False): (2, 2, 1, 8), + (32, 32, 32, 16, 32, False, False, True): (1, 2, 1, 2), + (32, 32, 32, 16, 32, False, True, False): (1, 2, 1, 4), + (32, 32, 32, 16, 32, False, True, True): (1, 2, 1, 4), + (32, 32, 32, 16, 32, True, False, False): (1, 2, 1, 4), + (32, 32, 32, 16, 32, True, False, True): (1, 2, 1, 2), + (32, 32, 32, 32, 32, False, False, False): (1, 1, 3, 8), + (32, 32, 32, 32, 32, False, False, True): (1, 1, 1, 8), + (32, 32, 32, 32, 32, False, True, False): (2, 1, 3, 4), + (32, 32, 32, 32, 32, False, True, True): (2, 1, 1, 2), + (32, 32, 32, 32, 32, True, False, False): (1, 1, 1, 2), + (32, 32, 32, 32, 32, True, False, True): (4, 1, 1, 1), + (32, 32, 64, 16, 16, False, False, False): (1, 4, 1, 4), + (32, 32, 64, 16, 16, False, False, True): (1, 4, 1, 4), + (32, 32, 64, 16, 16, False, True, False): (1, 2, 1, 8), + (32, 32, 64, 16, 16, False, True, True): (1, 4, 1, 2), + (32, 32, 64, 16, 16, True, False, False): (2, 4, 1, 2), + (32, 32, 64, 16, 16, True, False, True): (1, 4, 1, 2), + (32, 32, 64, 16, 32, False, False, False): (1, 2, 1, 8), + (32, 32, 64, 16, 32, False, False, True): (1, 4, 2, 2), + (32, 32, 64, 16, 32, False, True, False): (1, 2, 1, 4), + (32, 32, 64, 16, 32, False, True, True): (1, 4, 1, 4), + (32, 32, 64, 16, 32, True, False, False): (1, 4, 2, 2), + (32, 32, 64, 16, 32, True, False, True): (3, 4, 2, 2), + (32, 32, 64, 32, 32, False, False, False): (2, 2, 1, 4), + (32, 32, 64, 32, 32, False, False, True): (1, 2, 1, 4), + (32, 32, 64, 32, 32, False, True, False): (1, 1, 1, 8), + (32, 32, 64, 32, 32, False, True, True): (1, 1, 1, 4), + (32, 32, 64, 32, 32, True, False, False): (1, 2, 1, 2), + (32, 32, 64, 32, 32, True, False, True): (3, 2, 1, 8), + (32, 64, 16, 16, 32, False, False, False): (1, 1, 2, 2), + (32, 64, 16, 16, 32, False, False, True): (1, 1, 1, 4), + (32, 64, 16, 16, 32, False, True, False): (1, 1, 2, 4), + (32, 64, 16, 16, 32, False, True, True): (1, 1, 1, 4), + (32, 64, 16, 16, 32, True, False, False): (1, 1, 1, 2), + (32, 64, 16, 16, 32, True, False, True): (2, 1, 2, 2), + (32, 64, 16, 32, 32, False, False, False): (1, 1, 1, 1), + (32, 64, 16, 32, 32, False, False, True): (2, 1, 1, 4), + (32, 64, 16, 32, 32, False, True, False): (1, 1, 1, 1), + (32, 64, 16, 32, 32, False, True, True): (1, 1, 2, 2), + (32, 64, 16, 32, 32, True, False, False): (1, 1, 2, 4), + (32, 64, 16, 32, 32, True, False, True): (1, 1, 1, 4), + (32, 64, 32, 16, 32, False, False, False): (2, 2, 1, 4), + (32, 64, 32, 16, 32, False, False, True): (1, 2, 1, 4), + (32, 64, 32, 16, 32, False, True, False): (1, 1, 1, 4), + (32, 64, 32, 16, 32, False, True, True): (2, 2, 3, 4), + (32, 64, 32, 16, 32, True, False, False): (1, 1, 1, 2), + (32, 64, 32, 16, 32, True, False, True): (1, 2, 1, 2), + (32, 64, 32, 32, 32, False, False, False): (1, 1, 1, 2), + (32, 64, 32, 32, 32, False, False, True): (2, 1, 1, 4), + (32, 64, 32, 32, 32, False, True, False): (1, 1, 1, 8), + (32, 64, 32, 32, 32, False, True, True): (1, 1, 2, 4), + (32, 64, 32, 32, 32, True, False, False): (2, 1, 1, 4), + (32, 64, 32, 32, 32, True, False, True): (1, 1, 2, 4), + (32, 64, 64, 16, 32, False, False, False): (1, 4, 1, 4), + (32, 64, 64, 16, 32, False, False, True): (1, 4, 2, 4), + (32, 64, 64, 16, 32, False, True, False): (1, 4, 2, 2), + (32, 64, 64, 16, 32, False, True, True): (1, 4, 1, 4), + (32, 64, 64, 16, 32, True, False, False): (1, 4, 1, 8), + (32, 64, 64, 16, 32, True, False, True): (1, 4, 2, 1), + (32, 64, 64, 32, 32, False, False, False): (1, 1, 1, 4), + (32, 64, 64, 32, 32, False, False, True): (2, 2, 1, 4), + (32, 64, 64, 32, 32, False, True, False): (1, 1, 1, 4), + (32, 64, 64, 32, 32, False, True, True): (2, 2, 1, 4), + (32, 64, 64, 32, 32, True, False, False): (1, 2, 2, 4), + (32, 64, 64, 32, 32, True, False, True): (2, 2, 3, 4), + (64, 32, 16, 32, 32, False, False, False): (1, 1, 1, 4), + (64, 32, 16, 32, 32, False, False, True): (1, 1, 1, 4), + (64, 32, 16, 32, 32, False, True, False): (1, 1, 1, 8), + (64, 32, 16, 32, 32, False, True, True): (1, 1, 1, 4), + (64, 32, 16, 32, 32, True, False, False): (1, 1, 1, 16), + (64, 32, 16, 32, 32, True, False, True): (2, 1, 1, 4), + (64, 32, 32, 32, 32, False, False, False): (1, 1, 3, 4), + (64, 32, 32, 32, 32, False, False, True): (2, 1, 1, 4), + (64, 32, 32, 32, 32, False, True, False): (1, 1, 2, 4), + (64, 32, 32, 32, 32, False, True, True): (2, 1, 1, 4), + (64, 32, 32, 32, 32, True, False, False): (2, 1, 1, 16), + (64, 32, 32, 32, 32, True, False, True): (2, 1, 1, 4), + (64, 32, 64, 32, 32, False, False, False): (1, 2, 1, 4), + (64, 32, 64, 32, 32, False, False, True): (2, 2, 1, 4), + (64, 32, 64, 32, 32, False, True, False): (1, 1, 1, 4), + (64, 32, 64, 32, 32, False, True, True): (2, 2, 1, 4), + (64, 32, 64, 32, 32, True, False, False): (1, 2, 1, 8), + (64, 32, 64, 32, 32, True, False, True): (2, 2, 3, 4), + (64, 64, 16, 32, 32, False, False, False): (1, 1, 2, 16), + (64, 64, 16, 32, 32, False, False, True): (1, 1, 3, 4), + (64, 64, 16, 32, 32, False, True, False): (1, 1, 1, 2), + (64, 64, 16, 32, 32, False, True, True): (2, 1, 1, 4), + (64, 64, 16, 32, 32, True, False, False): (2, 1, 3, 2), + (64, 64, 16, 32, 32, True, False, True): (1, 1, 2, 4), + (64, 64, 32, 32, 32, False, False, False): (1, 1, 1, 8), + (64, 64, 32, 32, 32, False, False, True): (2, 1, 2, 4), + (64, 64, 32, 32, 32, False, True, False): (2, 1, 1, 4), + (64, 64, 32, 32, 32, False, True, True): (1, 1, 2, 4), + (64, 64, 32, 32, 32, True, False, False): (2, 1, 1, 4), + (64, 64, 32, 32, 32, True, False, True): (1, 1, 2, 4), + (64, 64, 64, 32, 32, False, False, False): (1, 2, 2, 4), + (64, 64, 64, 32, 32, False, False, True): (1, 2, 2, 2), + (64, 64, 64, 32, 32, False, True, False): (1, 2, 1, 2), + (64, 64, 64, 32, 32, False, True, True): (1, 2, 1, 4), + (64, 64, 64, 32, 32, True, False, False): (1, 2, 1, 4), + (64, 64, 64, 32, 32, True, False, True): (1, 2, 1, 4), + (192, 192, 256, 16, 16, False, True, True): (1, 8, 5, 4), + (192, 192, 256, 16, 16, True, False, True): (2, 8, 5, 2), + (192, 192, 256, 32, 32, False, True, True): (1, 8, 6, 4), + (192, 192, 256, 32, 32, True, False, True): (3, 8, 5, 2), + (192, 192, 512, 16, 16, False, True, True): (1, 16, 5, 2), + (192, 192, 512, 16, 16, True, False, True): (1, 8, 4, 2), + (192, 192, 512, 32, 32, False, True, True): (2, 16, 5, 4), + (192, 192, 512, 32, 32, True, False, True): (2, 8, 5, 2), + (192, 192, 1024, 16, 16, False, True, True): (1, 16, 3, 4), + (192, 192, 1024, 16, 16, True, False, True): (1, 16, 6, 2), + (192, 192, 1024, 32, 32, False, True, True): (1, 32, 3, 4), + (192, 192, 1024, 32, 32, True, False, True): (1, 16, 4, 2), + (192, 192, 2048, 16, 16, False, True, True): (1, 32, 1, 4), + (192, 192, 2048, 16, 16, True, False, True): (4, 32, 4, 2), + (192, 192, 2048, 32, 32, False, True, True): (1, 16, 3, 8), + (192, 192, 2048, 32, 32, True, False, True): (2, 32, 4, 2), + (192, 192, 4096, 16, 16, False, True, True): (2, 64, 1, 4), + (192, 192, 4096, 16, 16, True, False, True): (1, 32, 3, 2), + (192, 192, 4096, 32, 32, False, True, True): (1, 64, 1, 8), + (192, 192, 4096, 32, 32, True, False, True): (2, 32, 4, 4), + (192, 192, 8192, 16, 16, False, True, True): (1, 64, 1, 4), + (192, 192, 8192, 16, 16, True, False, True): (2, 32, 3, 1), + (192, 192, 8192, 32, 32, False, True, True): (3, 128, 1, 4), + (192, 192, 8192, 32, 32, True, False, True): (1, 64, 3, 4), + (192, 192, 16384, 16, 16, False, True, True): (1, 128, 1, 4), + (192, 192, 16384, 16, 16, True, False, True): (4, 64, 3, 1), + (192, 192, 16384, 32, 32, False, True, True): (1, 128, 1, 4), + (192, 192, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (192, 192, 32768, 16, 16, False, True, True): (2, 256, 1, 2), + (192, 192, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (192, 192, 32768, 32, 32, False, True, True): (2, 256, 1, 4), + (192, 192, 32768, 32, 32, True, False, True): (4, 128, 3, 4), + (192, 192, 65536, 16, 16, False, True, True): (2, 512, 1, 2), + (192, 192, 65536, 16, 16, True, False, True): (2, 256, 3, 2), + (192, 192, 65536, 32, 32, False, True, True): (2, 512, 1, 4), + (192, 192, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (192, 192, 131072, 16, 16, False, True, True): (4, 1024, 1, 2), + (192, 192, 131072, 16, 16, True, False, True): (3, 512, 3, 2), + (192, 192, 131072, 32, 32, False, True, True): (1, 1024, 1, 2), + (192, 192, 131072, 32, 32, True, False, True): (1, 512, 3, 4), + (256, 256, 256, 16, 16, False, True, True): (4, 8, 5, 1), + (256, 256, 256, 16, 16, True, False, True): (2, 8, 4, 2), + (256, 256, 256, 32, 32, False, True, True): (2, 8, 5, 2), + (256, 256, 256, 32, 32, True, False, True): (1, 8, 5, 4), + (256, 256, 256, 64, 64, False, True, True): (2, 4, 4, 4), + (256, 256, 256, 64, 64, True, False, True): (1, 4, 3, 4), + (256, 256, 256, 128, 128, False, True, True): (4, 2, 2, 8), + (256, 256, 256, 128, 128, True, False, True): (1, 2, 2, 8), + (256, 256, 512, 16, 16, False, True, True): (1, 16, 5, 1), + (256, 256, 512, 16, 16, True, False, True): (3, 16, 3, 2), + (256, 256, 512, 32, 32, False, True, True): (2, 8, 5, 2), + (256, 256, 512, 32, 32, True, False, True): (1, 16, 4, 4), + (256, 256, 512, 64, 64, False, True, True): (1, 8, 4, 4), + (256, 256, 512, 64, 64, True, False, True): (3, 8, 3, 4), + (256, 256, 512, 128, 128, False, True, True): (1, 4, 2, 8), + (256, 256, 512, 128, 128, True, False, True): (1, 4, 2, 8), + (256, 256, 1024, 16, 16, False, True, True): (1, 16, 5, 4), + (256, 256, 1024, 16, 16, True, False, True): (5, 16, 4, 2), + (256, 256, 1024, 32, 32, False, True, True): (1, 32, 5, 2), + (256, 256, 1024, 32, 32, True, False, True): (2, 16, 5, 2), + (256, 256, 1024, 64, 64, False, True, True): (1, 16, 4, 4), + (256, 256, 1024, 64, 64, True, False, True): (1, 16, 4, 4), + (256, 256, 1024, 128, 128, False, True, True): (1, 8, 2, 8), + (256, 256, 1024, 128, 128, True, False, True): (1, 8, 2, 8), + (256, 256, 2048, 16, 16, False, True, True): (1, 16, 4, 4), + (256, 256, 2048, 16, 16, True, False, True): (2, 32, 5, 1), + (256, 256, 2048, 32, 32, False, True, True): (1, 64, 4, 1), + (256, 256, 2048, 32, 32, True, False, True): (2, 32, 4, 2), + (256, 256, 2048, 64, 64, False, True, True): (8, 16, 5, 4), + (256, 256, 2048, 64, 64, True, False, True): (1, 16, 4, 4), + (256, 256, 2048, 128, 128, False, True, True): (2, 16, 2, 8), + (256, 256, 2048, 128, 128, True, False, True): (1, 16, 2, 8), + (256, 256, 4096, 16, 16, False, True, True): (1, 64, 1, 4), + (256, 256, 4096, 16, 16, True, False, True): (1, 16, 3, 2), + (256, 256, 4096, 32, 32, False, True, True): (6, 32, 3, 2), + (256, 256, 4096, 32, 32, True, False, True): (4, 32, 4, 2), + (256, 256, 4096, 64, 64, False, True, True): (6, 64, 3, 4), + (256, 256, 4096, 64, 64, True, False, True): (2, 64, 3, 4), + (256, 256, 4096, 128, 128, False, True, True): (1, 32, 2, 8), + (256, 256, 4096, 128, 128, True, False, True): (1, 32, 2, 8), + (256, 256, 8192, 16, 16, False, True, True): (2, 32, 3, 4), + (256, 256, 8192, 16, 16, True, False, True): (4, 64, 3, 2), + (256, 256, 8192, 32, 32, False, True, True): (1, 64, 3, 4), + (256, 256, 8192, 32, 32, True, False, True): (3, 128, 1, 2), + (256, 256, 8192, 64, 64, False, True, True): (9, 128, 1, 4), + (256, 256, 8192, 64, 64, True, False, True): (8, 128, 1, 4), + (256, 256, 8192, 128, 128, False, True, True): (7, 64, 1, 4), + (256, 256, 8192, 128, 128, True, False, True): (1, 32, 1, 16), + (256, 256, 16384, 16, 16, False, True, True): (3, 128, 3, 2), + (256, 256, 16384, 16, 16, True, False, True): (5, 64, 3, 2), + (256, 256, 16384, 32, 32, False, True, True): (3, 128, 3, 2), + (256, 256, 16384, 32, 32, True, False, True): (1, 128, 3, 2), + (256, 256, 16384, 64, 64, False, True, True): (3, 128, 1, 4), + (256, 256, 16384, 64, 64, True, False, True): (2, 128, 1, 4), + (256, 256, 16384, 128, 128, False, True, True): (7, 128, 1, 4), + (256, 256, 16384, 128, 128, True, False, True): (1, 128, 2, 8), + (256, 256, 32768, 16, 16, False, True, True): (2, 128, 3, 2), + (256, 256, 32768, 16, 16, True, False, True): (1, 128, 3, 2), + (256, 256, 32768, 32, 32, False, True, True): (1, 256, 3, 4), + (256, 256, 32768, 32, 32, True, False, True): (3, 256, 3, 2), + (256, 256, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (256, 256, 32768, 64, 64, True, False, True): (3, 256, 1, 4), + (256, 256, 32768, 128, 128, False, True, True): (9, 256, 1, 4), + (256, 256, 32768, 128, 128, True, False, True): (2, 256, 1, 4), + (256, 256, 65536, 16, 16, False, True, True): (1, 256, 3, 2), + (256, 256, 65536, 16, 16, True, False, True): (1, 256, 3, 2), + (256, 256, 65536, 32, 32, False, True, True): (2, 512, 3, 2), + (256, 256, 65536, 32, 32, True, False, True): (2, 512, 3, 2), + (256, 256, 65536, 64, 64, False, True, True): (2, 512, 1, 4), + (256, 256, 65536, 64, 64, True, False, True): (1, 512, 1, 4), + (256, 256, 65536, 128, 128, False, True, True): (7, 512, 1, 4), + (256, 256, 65536, 128, 128, True, False, True): (2, 512, 1, 4), + (256, 256, 131072, 16, 16, False, True, True): (1, 512, 3, 2), + (256, 256, 131072, 16, 16, True, False, True): (1, 512, 3, 2), + (256, 256, 131072, 32, 32, False, True, True): (1, 1024, 3, 2), + (256, 256, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (256, 256, 131072, 64, 64, False, True, True): (1, 1024, 1, 4), + (256, 256, 131072, 64, 64, True, False, True): (1, 1024, 1, 4), + (256, 256, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (256, 256, 131072, 128, 128, True, False, True): (1, 1024, 1, 4), + (384, 384, 256, 16, 16, False, True, True): (1, 8, 5, 2), + (384, 384, 256, 16, 16, True, False, True): (3, 4, 5, 2), + (384, 384, 256, 32, 32, False, True, True): (2, 8, 4, 4), + (384, 384, 256, 32, 32, True, False, True): (1, 4, 6, 2), + (384, 384, 256, 64, 64, False, True, True): (2, 4, 4, 4), + (384, 384, 256, 64, 64, True, False, True): (2, 4, 4, 4), + (384, 384, 512, 16, 16, False, True, True): (1, 8, 4, 2), + (384, 384, 512, 16, 16, True, False, True): (1, 4, 5, 4), + (384, 384, 512, 32, 32, False, True, True): (1, 8, 4, 4), + (384, 384, 512, 32, 32, True, False, True): (3, 8, 5, 2), + (384, 384, 512, 64, 64, False, True, True): (3, 8, 3, 4), + (384, 384, 512, 64, 64, True, False, True): (5, 8, 5, 4), + (384, 384, 1024, 16, 16, False, True, True): (3, 16, 4, 2), + (384, 384, 1024, 16, 16, True, False, True): (1, 8, 4, 4), + (384, 384, 1024, 32, 32, False, True, True): (6, 32, 3, 2), + (384, 384, 1024, 32, 32, True, False, True): (3, 8, 4, 4), + (384, 384, 1024, 64, 64, False, True, True): (3, 16, 3, 4), + (384, 384, 1024, 64, 64, True, False, True): (2, 16, 4, 4), + (384, 384, 2048, 16, 16, False, True, True): (1, 32, 1, 4), + (384, 384, 2048, 16, 16, True, False, True): (1, 16, 5, 2), + (384, 384, 2048, 32, 32, False, True, True): (1, 32, 1, 8), + (384, 384, 2048, 32, 32, True, False, True): (1, 8, 4, 4), + (384, 384, 2048, 64, 64, False, True, True): (4, 16, 3, 4), + (384, 384, 2048, 64, 64, True, False, True): (1, 16, 3, 8), + (384, 384, 4096, 16, 16, False, True, True): (5, 32, 1, 4), + (384, 384, 4096, 16, 16, True, False, True): (6, 32, 3, 2), + (384, 384, 4096, 32, 32, False, True, True): (1, 32, 1, 8), + (384, 384, 4096, 32, 32, True, False, True): (1, 16, 3, 4), + (384, 384, 4096, 64, 64, False, True, True): (1, 64, 1, 4), + (384, 384, 4096, 64, 64, True, False, True): (2, 32, 3, 4), + (384, 384, 8192, 16, 16, False, True, True): (2, 64, 1, 4), + (384, 384, 8192, 16, 16, True, False, True): (3, 32, 3, 2), + (384, 384, 8192, 32, 32, False, True, True): (5, 64, 1, 8), + (384, 384, 8192, 32, 32, True, False, True): (1, 32, 3, 2), + (384, 384, 8192, 64, 64, False, True, True): (1, 128, 1, 4), + (384, 384, 8192, 64, 64, True, False, True): (3, 64, 3, 4), + (384, 384, 16384, 16, 16, False, True, True): (1, 128, 1, 2), + (384, 384, 16384, 16, 16, True, False, True): (4, 128, 3, 2), + (384, 384, 16384, 32, 32, False, True, True): (3, 128, 1, 4), + (384, 384, 16384, 32, 32, True, False, True): (1, 128, 3, 2), + (384, 384, 16384, 64, 64, False, True, True): (3, 256, 1, 4), + (384, 384, 16384, 64, 64, True, False, True): (2, 128, 3, 4), + (384, 384, 32768, 16, 16, False, True, True): (1, 256, 1, 2), + (384, 384, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (384, 384, 32768, 32, 32, False, True, True): (1, 256, 1, 2), + (384, 384, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (384, 384, 32768, 64, 64, False, True, True): (2, 256, 1, 4), + (384, 384, 32768, 64, 64, True, False, True): (1, 256, 3, 4), + (384, 384, 65536, 16, 16, False, True, True): (4, 512, 1, 2), + (384, 384, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (384, 384, 65536, 32, 32, False, True, True): (1, 512, 1, 2), + (384, 384, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (384, 384, 65536, 64, 64, False, True, True): (3, 512, 1, 4), + (384, 384, 65536, 64, 64, True, False, True): (3, 256, 3, 4), + (384, 384, 131072, 16, 16, False, True, True): (1, 512, 1, 1), + (384, 384, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (384, 384, 131072, 32, 32, False, True, True): (1, 512, 1, 4), + (384, 384, 131072, 32, 32, True, False, True): (1, 512, 3, 4), + (384, 384, 131072, 64, 64, False, True, True): (3, 1024, 1, 4), + (384, 384, 131072, 64, 64, True, False, True): (3, 512, 3, 4), + (512, 512, 256, 16, 16, False, True, True): (2, 4, 5, 4), + (512, 512, 256, 16, 16, True, False, True): (3, 4, 5, 4), + (512, 512, 256, 32, 32, False, True, True): (1, 4, 5, 2), + (512, 512, 256, 32, 32, True, False, True): (4, 8, 5, 1), + (512, 512, 256, 64, 64, False, True, True): (4, 4, 5, 4), + (512, 512, 256, 64, 64, True, False, True): (5, 4, 5, 4), + (512, 512, 256, 128, 128, False, True, True): (3, 2, 2, 8), + (512, 512, 256, 128, 128, True, False, True): (2, 2, 2, 8), + (512, 512, 512, 16, 16, False, True, True): (1, 8, 5, 4), + (512, 512, 512, 16, 16, True, False, True): (4, 8, 5, 2), + (512, 512, 512, 32, 32, False, True, True): (1, 16, 4, 1), + (512, 512, 512, 32, 32, True, False, True): (1, 8, 5, 2), + (512, 512, 512, 64, 64, False, True, True): (4, 8, 5, 4), + (512, 512, 512, 64, 64, True, False, True): (2, 8, 5, 4), + (512, 512, 512, 128, 128, False, True, True): (2, 4, 2, 8), + (512, 512, 512, 128, 128, True, False, True): (1, 4, 2, 8), + (512, 512, 1024, 16, 16, False, True, True): (2, 8, 4, 4), + (512, 512, 1024, 16, 16, True, False, True): (1, 8, 4, 4), + (512, 512, 1024, 32, 32, False, True, True): (3, 16, 4, 2), + (512, 512, 1024, 32, 32, True, False, True): (1, 16, 5, 2), + (512, 512, 1024, 64, 64, False, True, True): (2, 8, 3, 4), + (512, 512, 1024, 64, 64, True, False, True): (2, 16, 3, 4), + (512, 512, 1024, 128, 128, False, True, True): (2, 8, 2, 8), + (512, 512, 1024, 128, 128, True, False, True): (3, 8, 2, 8), + (512, 512, 2048, 16, 16, False, True, True): (4, 16, 3, 2), + (512, 512, 2048, 16, 16, True, False, True): (1, 16, 4, 2), + (512, 512, 2048, 32, 32, False, True, True): (3, 32, 3, 2), + (512, 512, 2048, 32, 32, True, False, True): (2, 32, 3, 2), + (512, 512, 2048, 64, 64, False, True, True): (6, 32, 3, 2), + (512, 512, 2048, 64, 64, True, False, True): (1, 32, 3, 2), + (512, 512, 2048, 128, 128, False, True, True): (4, 16, 2, 8), + (512, 512, 2048, 128, 128, True, False, True): (1, 16, 2, 8), + (512, 512, 4096, 16, 16, False, True, True): (1, 16, 3, 2), + (512, 512, 4096, 16, 16, True, False, True): (4, 32, 3, 2), + (512, 512, 4096, 32, 32, False, True, True): (3, 32, 3, 2), + (512, 512, 4096, 32, 32, True, False, True): (2, 32, 3, 2), + (512, 512, 4096, 64, 64, False, True, True): (1, 32, 3, 4), + (512, 512, 4096, 64, 64, True, False, True): (1, 64, 3, 4), + (512, 512, 4096, 128, 128, False, True, True): (4, 32, 1, 4), + (512, 512, 4096, 128, 128, True, False, True): (4, 32, 2, 8), + (512, 512, 8192, 16, 16, False, True, True): (8, 64, 3, 2), + (512, 512, 8192, 16, 16, True, False, True): (4, 64, 3, 2), + (512, 512, 8192, 32, 32, False, True, True): (3, 64, 3, 2), + (512, 512, 8192, 32, 32, True, False, True): (3, 64, 3, 2), + (512, 512, 8192, 64, 64, False, True, True): (1, 64, 3, 4), + (512, 512, 8192, 64, 64, True, False, True): (7, 64, 3, 4), + (512, 512, 8192, 128, 128, False, True, True): (1, 64, 1, 4), + (512, 512, 8192, 128, 128, True, False, True): (4, 64, 2, 8), + (512, 512, 16384, 16, 16, False, True, True): (1, 64, 3, 2), + (512, 512, 16384, 16, 16, True, False, True): (1, 128, 3, 2), + (512, 512, 16384, 32, 32, False, True, True): (3, 128, 3, 2), + (512, 512, 16384, 32, 32, True, False, True): (1, 128, 3, 2), + (512, 512, 16384, 64, 64, False, True, True): (4, 64, 2, 4), + (512, 512, 16384, 64, 64, True, False, True): (2, 64, 2, 4), + (512, 512, 16384, 128, 128, False, True, True): (4, 128, 1, 4), + (512, 512, 16384, 128, 128, True, False, True): (2, 128, 1, 4), + (512, 512, 32768, 16, 16, False, True, True): (1, 128, 3, 2), + (512, 512, 32768, 16, 16, True, False, True): (1, 128, 3, 2), + (512, 512, 32768, 32, 32, False, True, True): (1, 256, 3, 2), + (512, 512, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (512, 512, 32768, 64, 64, False, True, True): (1, 256, 3, 4), + (512, 512, 32768, 64, 64, True, False, True): (2, 256, 3, 4), + (512, 512, 32768, 128, 128, False, True, True): (5, 256, 1, 4), + (512, 512, 32768, 128, 128, True, False, True): (4, 256, 1, 4), + (512, 512, 65536, 16, 16, False, True, True): (1, 256, 3, 2), + (512, 512, 65536, 16, 16, True, False, True): (1, 256, 3, 1), + (512, 512, 65536, 32, 32, False, True, True): (1, 512, 3, 2), + (512, 512, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (512, 512, 65536, 64, 64, False, True, True): (4, 256, 2, 4), + (512, 512, 65536, 64, 64, True, False, True): (2, 512, 3, 4), + (512, 512, 65536, 128, 128, False, True, True): (6, 512, 1, 4), + (512, 512, 65536, 128, 128, True, False, True): (4, 512, 1, 4), + (512, 512, 131072, 16, 16, False, True, True): (1, 512, 3, 2), + (512, 512, 131072, 16, 16, True, False, True): (1, 512, 3, 1), + (512, 512, 131072, 32, 32, False, True, True): (1, 1024, 3, 2), + (512, 512, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (512, 512, 131072, 64, 64, False, True, True): (4, 512, 2, 4), + (512, 512, 131072, 64, 64, True, False, True): (4, 1024, 3, 4), + (512, 512, 131072, 128, 128, False, True, True): (6, 1024, 1, 4), + (512, 512, 131072, 128, 128, True, False, True): (4, 1024, 1, 4), + (768, 768, 256, 16, 16, False, True, True): (1, 8, 4, 1), + (768, 768, 256, 16, 16, True, False, True): (3, 2, 6, 4), + (768, 768, 256, 32, 32, False, True, True): (3, 8, 3, 4), + (768, 768, 256, 32, 32, True, False, True): (1, 4, 4, 2), + (768, 768, 256, 64, 64, False, True, True): (2, 4, 3, 4), + (768, 768, 256, 64, 64, True, False, True): (1, 4, 4, 4), + (768, 768, 256, 128, 128, False, True, True): (2, 2, 3, 8), + (768, 768, 256, 128, 128, True, False, True): (4, 2, 3, 8), + (768, 768, 512, 16, 16, False, True, True): (4, 8, 4, 2), + (768, 768, 512, 16, 16, True, False, True): (4, 8, 6, 2), + (768, 768, 512, 32, 32, False, True, True): (1, 8, 4, 4), + (768, 768, 512, 32, 32, True, False, True): (3, 8, 4, 2), + (768, 768, 512, 64, 64, False, True, True): (1, 8, 3, 4), + (768, 768, 512, 64, 64, True, False, True): (1, 8, 4, 4), + (768, 768, 512, 128, 128, False, True, True): (1, 4, 3, 8), + (768, 768, 512, 128, 128, True, False, True): (4, 4, 3, 8), + (768, 768, 1024, 16, 16, False, True, True): (3, 16, 1, 4), + (768, 768, 1024, 16, 16, True, False, True): (1, 8, 5, 2), + (768, 768, 1024, 32, 32, False, True, True): (3, 16, 1, 8), + (768, 768, 1024, 32, 32, True, False, True): (1, 16, 3, 2), + (768, 768, 1024, 64, 64, False, True, True): (1, 8, 3, 4), + (768, 768, 1024, 64, 64, True, False, True): (2, 8, 3, 8), + (768, 768, 1024, 128, 128, False, True, True): (1, 8, 3, 8), + (768, 768, 1024, 128, 128, True, False, True): (1, 8, 3, 8), + (768, 768, 2048, 16, 16, False, True, True): (2, 16, 1, 2), + (768, 768, 2048, 16, 16, True, False, True): (1, 16, 3, 2), + (768, 768, 2048, 32, 32, False, True, True): (5, 32, 1, 4), + (768, 768, 2048, 32, 32, True, False, True): (3, 8, 3, 4), + (768, 768, 2048, 64, 64, False, True, True): (1, 16, 1, 8), + (768, 768, 2048, 64, 64, True, False, True): (3, 16, 3, 4), + (768, 768, 2048, 128, 128, False, True, True): (2, 16, 3, 8), + (768, 768, 2048, 128, 128, True, False, True): (1, 16, 3, 8), + (768, 768, 4096, 16, 16, False, True, True): (3, 32, 1, 4), + (768, 768, 4096, 16, 16, True, False, True): (2, 32, 3, 1), + (768, 768, 4096, 32, 32, False, True, True): (2, 64, 1, 4), + (768, 768, 4096, 32, 32, True, False, True): (1, 16, 4, 4), + (768, 768, 4096, 64, 64, False, True, True): (3, 64, 3, 4), + (768, 768, 4096, 64, 64, True, False, True): (2, 16, 3, 4), + (768, 768, 4096, 128, 128, False, True, True): (1, 32, 3, 8), + (768, 768, 4096, 128, 128, True, False, True): (4, 32, 3, 8), + (768, 768, 8192, 16, 16, False, True, True): (1, 64, 1, 2), + (768, 768, 8192, 16, 16, True, False, True): (4, 64, 3, 2), + (768, 768, 8192, 32, 32, False, True, True): (1, 64, 1, 8), + (768, 768, 8192, 32, 32, True, False, True): (2, 32, 3, 4), + (768, 768, 8192, 64, 64, False, True, True): (4, 64, 3, 4), + (768, 768, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (768, 768, 8192, 128, 128, False, True, True): (2, 64, 3, 8), + (768, 768, 8192, 128, 128, True, False, True): (1, 64, 3, 8), + (768, 768, 16384, 16, 16, False, True, True): (1, 128, 1, 2), + (768, 768, 16384, 16, 16, True, False, True): (1, 64, 4, 4), + (768, 768, 16384, 32, 32, False, True, True): (1, 128, 1, 8), + (768, 768, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (768, 768, 16384, 64, 64, False, True, True): (4, 128, 3, 4), + (768, 768, 16384, 64, 64, True, False, True): (1, 64, 3, 4), + (768, 768, 16384, 128, 128, False, True, True): (3, 128, 1, 4), + (768, 768, 16384, 128, 128, True, False, True): (3, 128, 2, 4), + (768, 768, 32768, 16, 16, False, True, True): (1, 256, 1, 2), + (768, 768, 32768, 16, 16, True, False, True): (1, 128, 4, 4), + (768, 768, 32768, 32, 32, False, True, True): (1, 128, 1, 2), + (768, 768, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (768, 768, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (768, 768, 32768, 64, 64, True, False, True): (2, 128, 3, 4), + (768, 768, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (768, 768, 32768, 128, 128, True, False, True): (2, 256, 2, 4), + (768, 768, 65536, 16, 16, False, True, True): (4, 512, 1, 2), + (768, 768, 65536, 16, 16, True, False, True): (1, 256, 4, 4), + (768, 768, 65536, 32, 32, False, True, True): (1, 256, 1, 2), + (768, 768, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (768, 768, 65536, 64, 64, False, True, True): (3, 512, 1, 4), + (768, 768, 65536, 64, 64, True, False, True): (2, 256, 3, 4), + (768, 768, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (768, 768, 65536, 128, 128, True, False, True): (2, 512, 2, 4), + (768, 768, 131072, 16, 16, False, True, True): (4, 1024, 1, 2), + (768, 768, 131072, 16, 16, True, False, True): (1, 512, 4, 1), + (768, 768, 131072, 32, 32, False, True, True): (1, 512, 1, 2), + (768, 768, 131072, 32, 32, True, False, True): (1, 512, 3, 4), + (768, 768, 131072, 64, 64, False, True, True): (1, 1024, 1, 4), + (768, 768, 131072, 64, 64, True, False, True): (2, 512, 3, 4), + (768, 768, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (768, 768, 131072, 128, 128, True, False, True): (1, 1024, 2, 4), + (768, 3072, 256, 16, 16, False, True, True): (3, 8, 6, 1), + (768, 3072, 256, 16, 16, True, False, True): (1, 4, 6, 2), + (768, 3072, 256, 32, 32, False, True, True): (1, 8, 4, 4), + (768, 3072, 256, 32, 32, True, False, True): (3, 4, 6, 4), + (768, 3072, 256, 64, 64, False, True, True): (2, 4, 3, 4), + (768, 3072, 256, 64, 64, True, False, True): (1, 4, 4, 4), + (768, 3072, 256, 128, 128, False, True, True): (2, 2, 3, 8), + (768, 3072, 256, 128, 128, True, False, True): (1, 2, 3, 8), + (768, 3072, 512, 16, 16, False, True, True): (1, 8, 4, 2), + (768, 3072, 512, 16, 16, True, False, True): (1, 8, 5, 2), + (768, 3072, 512, 32, 32, False, True, True): (1, 16, 3, 2), + (768, 3072, 512, 32, 32, True, False, True): (1, 8, 5, 2), + (768, 3072, 512, 64, 64, False, True, True): (1, 8, 3, 4), + (768, 3072, 512, 64, 64, True, False, True): (3, 8, 4, 4), + (768, 3072, 512, 128, 128, False, True, True): (1, 4, 3, 8), + (768, 3072, 512, 128, 128, True, False, True): (2, 4, 3, 8), + (768, 3072, 1024, 16, 16, False, True, True): (1, 16, 1, 4), + (768, 3072, 1024, 16, 16, True, False, True): (5, 4, 4, 4), + (768, 3072, 1024, 32, 32, False, True, True): (3, 8, 3, 4), + (768, 3072, 1024, 32, 32, True, False, True): (1, 8, 4, 4), + (768, 3072, 1024, 64, 64, False, True, True): (2, 16, 3, 4), + (768, 3072, 1024, 64, 64, True, False, True): (2, 16, 4, 4), + (768, 3072, 1024, 128, 128, False, True, True): (1, 8, 3, 8), + (768, 3072, 1024, 128, 128, True, False, True): (5, 8, 3, 8), + (768, 3072, 2048, 16, 16, False, True, True): (3, 16, 1, 2), + (768, 3072, 2048, 16, 16, True, False, True): (1, 8, 3, 4), + (768, 3072, 2048, 32, 32, False, True, True): (4, 16, 1, 8), + (768, 3072, 2048, 32, 32, True, False, True): (3, 8, 3, 4), + (768, 3072, 2048, 64, 64, False, True, True): (2, 16, 3, 4), + (768, 3072, 2048, 64, 64, True, False, True): (2, 16, 3, 4), + (768, 3072, 2048, 128, 128, False, True, True): (3, 16, 3, 8), + (768, 3072, 2048, 128, 128, True, False, True): (4, 16, 3, 8), + (768, 3072, 4096, 16, 16, False, True, True): (1, 32, 1, 4), + (768, 3072, 4096, 16, 16, True, False, True): (1, 16, 3, 1), + (768, 3072, 4096, 32, 32, False, True, True): (3, 32, 1, 8), + (768, 3072, 4096, 32, 32, True, False, True): (3, 16, 4, 4), + (768, 3072, 4096, 64, 64, False, True, True): (2, 32, 3, 4), + (768, 3072, 4096, 64, 64, True, False, True): (2, 16, 3, 4), + (768, 3072, 4096, 128, 128, False, True, True): (5, 32, 1, 4), + (768, 3072, 4096, 128, 128, True, False, True): (9, 32, 3, 8), + (768, 3072, 8192, 16, 16, False, True, True): (1, 32, 1, 4), + (768, 3072, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (768, 3072, 8192, 32, 32, False, True, True): (1, 64, 1, 8), + (768, 3072, 8192, 32, 32, True, False, True): (2, 64, 4, 2), + (768, 3072, 8192, 64, 64, False, True, True): (1, 64, 3, 4), + (768, 3072, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (768, 3072, 8192, 128, 128, False, True, True): (2, 64, 3, 8), + (768, 3072, 8192, 128, 128, True, False, True): (2, 64, 3, 8), + (768, 3072, 16384, 16, 16, False, True, True): (1, 64, 1, 4), + (768, 3072, 16384, 16, 16, True, False, True): (1, 64, 4, 1), + (768, 3072, 16384, 32, 32, False, True, True): (1, 128, 1, 8), + (768, 3072, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (768, 3072, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (768, 3072, 16384, 64, 64, True, False, True): (4, 64, 3, 4), + (768, 3072, 16384, 128, 128, False, True, True): (2, 128, 3, 8), + (768, 3072, 16384, 128, 128, True, False, True): (2, 128, 3, 8), + (768, 3072, 32768, 16, 16, False, True, True): (1, 128, 1, 4), + (768, 3072, 32768, 16, 16, True, False, True): (1, 128, 4, 1), + (768, 3072, 32768, 32, 32, False, True, True): (1, 256, 1, 8), + (768, 3072, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (768, 3072, 32768, 64, 64, False, True, True): (1, 256, 3, 4), + (768, 3072, 32768, 64, 64, True, False, True): (1, 128, 3, 4), + (768, 3072, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (768, 3072, 32768, 128, 128, True, False, True): (2, 256, 3, 8), + (768, 3072, 50432, 16, 16, False, True, True): (1, 197, 1, 4), + (768, 3072, 50432, 16, 16, True, False, True): (4, 197, 4, 4), + (768, 3072, 50432, 32, 32, False, True, True): (1, 197, 1, 4), + (768, 3072, 50432, 32, 32, True, False, True): (4, 197, 3, 4), + (768, 3072, 50432, 64, 64, False, True, True): (1, 394, 3, 4), + (768, 3072, 50432, 64, 64, True, False, True): (3, 197, 3, 4), + (768, 3072, 50432, 128, 128, False, True, True): (3, 394, 1, 4), + (768, 3072, 50432, 128, 128, True, False, True): (1, 394, 3, 8), + (768, 3072, 65536, 16, 16, False, True, True): (1, 256, 1, 4), + (768, 3072, 65536, 16, 16, True, False, True): (5, 256, 4, 1), + (768, 3072, 65536, 32, 32, False, True, True): (1, 256, 1, 4), + (768, 3072, 65536, 32, 32, True, False, True): (3, 256, 3, 4), + (768, 3072, 65536, 64, 64, False, True, True): (2, 512, 3, 4), + (768, 3072, 65536, 64, 64, True, False, True): (3, 256, 3, 4), + (768, 3072, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (768, 3072, 65536, 128, 128, True, False, True): (2, 512, 3, 8), + (768, 3072, 131072, 16, 16, False, True, True): (1, 512, 1, 4), + (768, 3072, 131072, 16, 16, True, False, True): (5, 512, 4, 1), + (768, 3072, 131072, 32, 32, False, True, True): (1, 512, 1, 4), + (768, 3072, 131072, 32, 32, True, False, True): (4, 512, 3, 4), + (768, 3072, 131072, 64, 64, False, True, True): (1, 1024, 3, 4), + (768, 3072, 131072, 64, 64, True, False, True): (1, 512, 3, 4), + (768, 3072, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (768, 3072, 131072, 128, 128, True, False, True): (1, 1024, 3, 8), + (1024, 1024, 256, 16, 16, False, True, True): (1, 4, 5, 4), + (1024, 1024, 256, 16, 16, True, False, True): (3, 4, 4, 4), + (1024, 1024, 256, 32, 32, False, True, True): (4, 4, 5, 2), + (1024, 1024, 256, 32, 32, True, False, True): (3, 4, 5, 2), + (1024, 1024, 256, 64, 64, False, True, True): (1, 4, 5, 4), + (1024, 1024, 256, 64, 64, True, False, True): (1, 4, 5, 4), + (1024, 1024, 256, 128, 128, False, True, True): (1, 2, 2, 8), + (1024, 1024, 256, 128, 128, True, False, True): (2, 2, 2, 8), + (1024, 1024, 512, 16, 16, False, True, True): (3, 4, 4, 4), + (1024, 1024, 512, 16, 16, True, False, True): (4, 8, 5, 2), + (1024, 1024, 512, 32, 32, False, True, True): (1, 8, 4, 2), + (1024, 1024, 512, 32, 32, True, False, True): (1, 8, 4, 2), + (1024, 1024, 512, 64, 64, False, True, True): (4, 8, 4, 4), + (1024, 1024, 512, 64, 64, True, False, True): (2, 8, 3, 4), + (1024, 1024, 512, 128, 128, False, True, True): (2, 4, 2, 8), + (1024, 1024, 512, 128, 128, True, False, True): (1, 4, 2, 8), + (1024, 1024, 1024, 16, 16, False, True, True): (3, 8, 4, 4), + (1024, 1024, 1024, 16, 16, True, False, True): (4, 8, 4, 2), + (1024, 1024, 1024, 32, 32, False, True, True): (1, 16, 3, 2), + (1024, 1024, 1024, 32, 32, True, False, True): (1, 16, 3, 2), + (1024, 1024, 1024, 64, 64, False, True, True): (1, 16, 3, 4), + (1024, 1024, 1024, 64, 64, True, False, True): (3, 16, 3, 2), + (1024, 1024, 1024, 128, 128, False, True, True): (1, 8, 2, 8), + (1024, 1024, 1024, 128, 128, True, False, True): (2, 8, 2, 8), + (1024, 1024, 2048, 16, 16, False, True, True): (3, 8, 3, 4), + (1024, 1024, 2048, 16, 16, True, False, True): (3, 8, 3, 2), + (1024, 1024, 2048, 32, 32, False, True, True): (5, 16, 3, 4), + (1024, 1024, 2048, 32, 32, True, False, True): (1, 16, 3, 2), + (1024, 1024, 2048, 64, 64, False, True, True): (6, 16, 4, 4), + (1024, 1024, 2048, 64, 64, True, False, True): (5, 16, 3, 4), + (1024, 1024, 2048, 128, 128, False, True, True): (4, 16, 2, 8), + (1024, 1024, 2048, 128, 128, True, False, True): (4, 16, 2, 8), + (1024, 1024, 4096, 16, 16, False, True, True): (8, 32, 3, 2), + (1024, 1024, 4096, 16, 16, True, False, True): (4, 32, 3, 2), + (1024, 1024, 4096, 32, 32, False, True, True): (2, 32, 3, 4), + (1024, 1024, 4096, 32, 32, True, False, True): (3, 32, 3, 2), + (1024, 1024, 4096, 64, 64, False, True, True): (3, 32, 3, 4), + (1024, 1024, 4096, 64, 64, True, False, True): (1, 32, 3, 4), + (1024, 1024, 4096, 128, 128, False, True, True): (4, 32, 2, 8), + (1024, 1024, 4096, 128, 128, True, False, True): (1, 32, 2, 8), + (1024, 1024, 8192, 16, 16, False, True, True): (4, 64, 3, 2), + (1024, 1024, 8192, 16, 16, True, False, True): (4, 64, 3, 2), + (1024, 1024, 8192, 32, 32, False, True, True): (8, 64, 3, 4), + (1024, 1024, 8192, 32, 32, True, False, True): (4, 32, 3, 4), + (1024, 1024, 8192, 64, 64, False, True, True): (4, 64, 3, 4), + (1024, 1024, 8192, 64, 64, True, False, True): (2, 64, 3, 4), + (1024, 1024, 8192, 128, 128, False, True, True): (4, 64, 2, 8), + (1024, 1024, 8192, 128, 128, True, False, True): (4, 64, 1, 4), + (1024, 1024, 16384, 16, 16, False, True, True): (1, 64, 3, 2), + (1024, 1024, 16384, 16, 16, True, False, True): (1, 64, 3, 2), + (1024, 1024, 16384, 32, 32, False, True, True): (1, 128, 3, 2), + (1024, 1024, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (1024, 1024, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (1024, 1024, 16384, 64, 64, True, False, True): (1, 128, 3, 4), + (1024, 1024, 16384, 128, 128, False, True, True): (2, 128, 1, 4), + (1024, 1024, 16384, 128, 128, True, False, True): (4, 128, 1, 4), + (1024, 1024, 32768, 16, 16, False, True, True): (1, 128, 3, 2), + (1024, 1024, 32768, 16, 16, True, False, True): (1, 128, 3, 2), + (1024, 1024, 32768, 32, 32, False, True, True): (1, 256, 3, 2), + (1024, 1024, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (1024, 1024, 32768, 64, 64, False, True, True): (2, 128, 2, 4), + (1024, 1024, 32768, 64, 64, True, False, True): (1, 256, 3, 4), + (1024, 1024, 32768, 128, 128, False, True, True): (2, 256, 1, 4), + (1024, 1024, 32768, 128, 128, True, False, True): (4, 256, 1, 4), + (1024, 1024, 65536, 16, 16, False, True, True): (1, 256, 3, 4), + (1024, 1024, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (1024, 1024, 65536, 32, 32, False, True, True): (9, 256, 3, 4), + (1024, 1024, 65536, 32, 32, True, False, True): (7, 256, 3, 4), + (1024, 1024, 65536, 64, 64, False, True, True): (2, 256, 2, 4), + (1024, 1024, 65536, 64, 64, True, False, True): (2, 512, 3, 4), + (1024, 1024, 65536, 128, 128, False, True, True): (2, 512, 1, 4), + (1024, 1024, 65536, 128, 128, True, False, True): (4, 512, 1, 4), + (1024, 1024, 131072, 16, 16, False, True, True): (11, 512, 3, 2), + (1024, 1024, 131072, 16, 16, True, False, True): (11, 512, 3, 2), + (1024, 1024, 131072, 32, 32, False, True, True): (4, 512, 3, 4), + (1024, 1024, 131072, 32, 32, True, False, True): (6, 512, 3, 4), + (1024, 1024, 131072, 64, 64, False, True, True): (2, 512, 2, 4), + (1024, 1024, 131072, 64, 64, True, False, True): (2, 1024, 3, 4), + (1024, 1024, 131072, 128, 128, False, True, True): (4, 1024, 1, 4), + (1024, 1024, 131072, 128, 128, True, False, True): (4, 1024, 1, 4), + (1280, 5120, 65792, 16, 16, False, True, True): (1, 257, 1, 4), + (1280, 5120, 65792, 16, 16, True, False, True): (5, 257, 4, 1), + (1280, 5120, 65792, 32, 32, False, True, True): (1, 514, 1, 8), + (1280, 5120, 65792, 32, 32, True, False, True): (2, 257, 3, 4), + (1280, 5120, 65792, 64, 64, False, True, True): (1, 514, 3, 4), + (1280, 5120, 65792, 64, 64, True, False, True): (1, 257, 3, 4), + (1280, 5120, 65792, 128, 128, False, True, True): (1, 514, 3, 8), + (1280, 5120, 65792, 128, 128, True, False, True): (2, 514, 3, 8), + (1536, 1536, 256, 16, 16, False, True, True): (1, 4, 6, 2), + (1536, 1536, 256, 16, 16, True, False, True): (3, 4, 5, 2), + (1536, 1536, 256, 32, 32, False, True, True): (2, 4, 3, 4), + (1536, 1536, 256, 32, 32, True, False, True): (1, 4, 5, 2), + (1536, 1536, 256, 64, 64, False, True, True): (2, 4, 3, 4), + (1536, 1536, 256, 64, 64, True, False, True): (1, 4, 4, 4), + (1536, 1536, 256, 128, 128, False, True, True): (3, 2, 3, 8), + (1536, 1536, 256, 128, 128, True, False, True): (6, 2, 3, 8), + (1536, 1536, 512, 16, 16, False, True, True): (1, 8, 1, 4), + (1536, 1536, 512, 16, 16, True, False, True): (3, 4, 5, 2), + (1536, 1536, 512, 32, 32, False, True, True): (1, 8, 1, 8), + (1536, 1536, 512, 32, 32, True, False, True): (1, 4, 4, 4), + (1536, 1536, 512, 64, 64, False, True, True): (3, 8, 5, 4), + (1536, 1536, 512, 64, 64, True, False, True): (3, 8, 3, 4), + (1536, 1536, 512, 128, 128, False, True, True): (2, 4, 3, 8), + (1536, 1536, 512, 128, 128, True, False, True): (3, 4, 3, 8), + (1536, 1536, 1024, 16, 16, False, True, True): (1, 8, 1, 2), + (1536, 1536, 1024, 16, 16, True, False, True): (2, 8, 4, 2), + (1536, 1536, 1024, 32, 32, False, True, True): (8, 16, 1, 4), + (1536, 1536, 1024, 32, 32, True, False, True): (3, 8, 4, 2), + (1536, 1536, 1024, 64, 64, False, True, True): (1, 16, 3, 4), + (1536, 1536, 1024, 64, 64, True, False, True): (3, 8, 3, 4), + (1536, 1536, 1024, 128, 128, False, True, True): (3, 8, 3, 8), + (1536, 1536, 1024, 128, 128, True, False, True): (3, 8, 3, 8), + (1536, 1536, 2048, 16, 16, False, True, True): (1, 16, 1, 4), + (1536, 1536, 2048, 16, 16, True, False, True): (1, 8, 3, 1), + (1536, 1536, 2048, 32, 32, False, True, True): (3, 16, 1, 8), + (1536, 1536, 2048, 32, 32, True, False, True): (3, 8, 4, 4), + (1536, 1536, 2048, 64, 64, False, True, True): (1, 16, 3, 4), + (1536, 1536, 2048, 64, 64, True, False, True): (3, 8, 3, 4), + (1536, 1536, 2048, 128, 128, False, True, True): (4, 16, 1, 4), + (1536, 1536, 2048, 128, 128, True, False, True): (6, 16, 3, 8), + (1536, 1536, 4096, 16, 16, False, True, True): (1, 32, 1, 2), + (1536, 1536, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (1536, 1536, 4096, 32, 32, False, True, True): (1, 32, 1, 8), + (1536, 1536, 4096, 32, 32, True, False, True): (5, 32, 4, 2), + (1536, 1536, 4096, 64, 64, False, True, True): (2, 32, 3, 4), + (1536, 1536, 4096, 64, 64, True, False, True): (2, 16, 3, 4), + (1536, 1536, 4096, 128, 128, False, True, True): (4, 32, 3, 8), + (1536, 1536, 4096, 128, 128, True, False, True): (4, 32, 3, 8), + (1536, 1536, 8192, 16, 16, False, True, True): (1, 64, 1, 2), + (1536, 1536, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (1536, 1536, 8192, 32, 32, False, True, True): (2, 64, 1, 8), + (1536, 1536, 8192, 32, 32, True, False, True): (2, 32, 3, 4), + (1536, 1536, 8192, 64, 64, False, True, True): (1, 64, 3, 4), + (1536, 1536, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (1536, 1536, 8192, 128, 128, False, True, True): (4, 64, 3, 8), + (1536, 1536, 8192, 128, 128, True, False, True): (1, 64, 3, 8), + (1536, 1536, 16384, 16, 16, False, True, True): (1, 128, 1, 2), + (1536, 1536, 16384, 16, 16, True, False, True): (1, 64, 4, 4), + (1536, 1536, 16384, 32, 32, False, True, True): (1, 64, 1, 2), + (1536, 1536, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (1536, 1536, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (1536, 1536, 16384, 64, 64, True, False, True): (1, 64, 3, 4), + (1536, 1536, 16384, 128, 128, False, True, True): (1, 128, 1, 4), + (1536, 1536, 16384, 128, 128, True, False, True): (1, 128, 2, 4), + (1536, 1536, 32768, 16, 16, False, True, True): (1, 256, 1, 2), + (1536, 1536, 32768, 16, 16, True, False, True): (1, 128, 3, 2), + (1536, 1536, 32768, 32, 32, False, True, True): (1, 128, 1, 2), + (1536, 1536, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (1536, 1536, 32768, 64, 64, False, True, True): (1, 256, 3, 4), + (1536, 1536, 32768, 64, 64, True, False, True): (1, 128, 3, 4), + (1536, 1536, 32768, 128, 128, False, True, True): (1, 256, 1, 4), + (1536, 1536, 32768, 128, 128, True, False, True): (2, 256, 2, 4), + (1536, 1536, 65536, 16, 16, False, True, True): (2, 512, 1, 2), + (1536, 1536, 65536, 16, 16, True, False, True): (1, 256, 4, 4), + (1536, 1536, 65536, 32, 32, False, True, True): (1, 256, 1, 2), + (1536, 1536, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (1536, 1536, 65536, 64, 64, False, True, True): (1, 512, 3, 4), + (1536, 1536, 65536, 64, 64, True, False, True): (3, 256, 3, 4), + (1536, 1536, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (1536, 1536, 65536, 128, 128, True, False, True): (4, 512, 2, 4), + (1536, 1536, 131072, 16, 16, False, True, True): (2, 1024, 1, 2), + (1536, 1536, 131072, 16, 16, True, False, True): (9, 512, 4, 4), + (1536, 1536, 131072, 32, 32, False, True, True): (1, 512, 1, 2), + (1536, 1536, 131072, 32, 32, True, False, True): (5, 512, 3, 4), + (1536, 1536, 131072, 64, 64, False, True, True): (1, 1024, 3, 4), + (1536, 1536, 131072, 64, 64, True, False, True): (2, 512, 3, 4), + (1536, 1536, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (1536, 1536, 131072, 128, 128, True, False, True): (1, 1024, 2, 4), + (2048, 2048, 256, 16, 16, False, True, True): (1, 4, 5, 2), + (2048, 2048, 256, 16, 16, True, False, True): (4, 4, 5, 2), + (2048, 2048, 256, 32, 32, False, True, True): (3, 4, 6, 2), + (2048, 2048, 256, 32, 32, True, False, True): (2, 4, 5, 2), + (2048, 2048, 256, 64, 64, False, True, True): (2, 4, 4, 4), + (2048, 2048, 256, 64, 64, True, False, True): (2, 4, 3, 4), + (2048, 2048, 256, 128, 128, False, True, True): (3, 2, 2, 8), + (2048, 2048, 256, 128, 128, True, False, True): (3, 2, 2, 8), + (2048, 2048, 512, 16, 16, False, True, True): (3, 4, 4, 4), + (2048, 2048, 512, 16, 16, True, False, True): (1, 4, 4, 4), + (2048, 2048, 512, 32, 32, False, True, True): (1, 4, 3, 4), + (2048, 2048, 512, 32, 32, True, False, True): (1, 4, 4, 2), + (2048, 2048, 512, 64, 64, False, True, True): (1, 8, 3, 4), + (2048, 2048, 512, 64, 64, True, False, True): (1, 8, 3, 4), + (2048, 2048, 512, 128, 128, False, True, True): (3, 4, 2, 8), + (2048, 2048, 512, 128, 128, True, False, True): (2, 4, 2, 8), + (2048, 2048, 1024, 16, 16, False, True, True): (3, 4, 3, 4), + (2048, 2048, 1024, 16, 16, True, False, True): (4, 8, 3, 2), + (2048, 2048, 1024, 32, 32, False, True, True): (3, 8, 3, 4), + (2048, 2048, 1024, 32, 32, True, False, True): (1, 8, 3, 2), + (2048, 2048, 1024, 64, 64, False, True, True): (1, 8, 3, 4), + (2048, 2048, 1024, 64, 64, True, False, True): (1, 8, 3, 4), + (2048, 2048, 1024, 128, 128, False, True, True): (4, 8, 1, 4), + (2048, 2048, 1024, 128, 128, True, False, True): (2, 8, 1, 4), + (2048, 2048, 2048, 16, 16, False, True, True): (4, 16, 3, 2), + (2048, 2048, 2048, 16, 16, True, False, True): (4, 16, 3, 2), + (2048, 2048, 2048, 32, 32, False, True, True): (1, 16, 3, 2), + (2048, 2048, 2048, 32, 32, True, False, True): (1, 16, 3, 2), + (2048, 2048, 2048, 64, 64, False, True, True): (4, 16, 3, 4), + (2048, 2048, 2048, 64, 64, True, False, True): (4, 16, 3, 4), + (2048, 2048, 2048, 128, 128, False, True, True): (6, 16, 2, 8), + (2048, 2048, 2048, 128, 128, True, False, True): (3, 16, 1, 4), + (2048, 2048, 4096, 16, 16, False, True, True): (4, 32, 4, 2), + (2048, 2048, 4096, 16, 16, True, False, True): (4, 32, 3, 2), + (2048, 2048, 4096, 32, 32, False, True, True): (4, 16, 3, 8), + (2048, 2048, 4096, 32, 32, True, False, True): (4, 16, 3, 8), + (2048, 2048, 4096, 64, 64, False, True, True): (1, 32, 3, 4), + (2048, 2048, 4096, 64, 64, True, False, True): (3, 32, 3, 4), + (2048, 2048, 4096, 128, 128, False, True, True): (2, 32, 1, 4), + (2048, 2048, 4096, 128, 128, True, False, True): (2, 32, 1, 4), + (2048, 2048, 8192, 16, 16, False, True, True): (4, 64, 4, 2), + (2048, 2048, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (2048, 2048, 8192, 32, 32, False, True, True): (4, 32, 4, 8), + (2048, 2048, 8192, 32, 32, True, False, True): (4, 32, 3, 8), + (2048, 2048, 8192, 64, 64, False, True, True): (4, 64, 3, 4), + (2048, 2048, 8192, 64, 64, True, False, True): (4, 64, 3, 4), + (2048, 2048, 8192, 128, 128, False, True, True): (2, 64, 1, 4), + (2048, 2048, 8192, 128, 128, True, False, True): (2, 64, 1, 4), + (2048, 2048, 16384, 16, 16, False, True, True): (4, 64, 3, 2), + (2048, 2048, 16384, 16, 16, True, False, True): (1, 64, 3, 2), + (2048, 2048, 16384, 32, 32, False, True, True): (4, 64, 3, 4), + (2048, 2048, 16384, 32, 32, True, False, True): (4, 64, 3, 4), + (2048, 2048, 16384, 64, 64, False, True, True): (4, 128, 3, 4), + (2048, 2048, 16384, 64, 64, True, False, True): (4, 128, 3, 4), + (2048, 2048, 16384, 128, 128, False, True, True): (2, 128, 1, 4), + (2048, 2048, 16384, 128, 128, True, False, True): (2, 128, 1, 4), + (2048, 2048, 32768, 16, 16, False, True, True): (8, 128, 3, 2), + (2048, 2048, 32768, 16, 16, True, False, True): (8, 128, 3, 4), + (2048, 2048, 32768, 32, 32, False, True, True): (8, 128, 3, 4), + (2048, 2048, 32768, 32, 32, True, False, True): (8, 128, 3, 4), + (2048, 2048, 32768, 64, 64, False, True, True): (1, 128, 2, 4), + (2048, 2048, 32768, 64, 64, True, False, True): (8, 256, 3, 4), + (2048, 2048, 32768, 128, 128, False, True, True): (2, 256, 1, 4), + (2048, 2048, 32768, 128, 128, True, False, True): (2, 256, 1, 4), + (2048, 2048, 65536, 16, 16, False, True, True): (9, 256, 4, 4), + (2048, 2048, 65536, 16, 16, True, False, True): (7, 256, 4, 4), + (2048, 2048, 65536, 32, 32, False, True, True): (7, 256, 3, 4), + (2048, 2048, 65536, 32, 32, True, False, True): (3, 256, 3, 4), + (2048, 2048, 65536, 64, 64, False, True, True): (2, 256, 2, 4), + (2048, 2048, 65536, 64, 64, True, False, True): (6, 512, 3, 4), + (2048, 2048, 65536, 128, 128, False, True, True): (2, 512, 1, 4), + (2048, 2048, 65536, 128, 128, True, False, True): (2, 512, 1, 4), + (2048, 2048, 131072, 16, 16, False, True, True): (9, 512, 4, 4), + (2048, 2048, 131072, 16, 16, True, False, True): (9, 512, 4, 4), + (2048, 2048, 131072, 32, 32, False, True, True): (7, 512, 4, 4), + (2048, 2048, 131072, 32, 32, True, False, True): (3, 512, 3, 4), + (2048, 2048, 131072, 64, 64, False, True, True): (2, 512, 2, 4), + (2048, 2048, 131072, 64, 64, True, False, True): (4, 1024, 3, 4), + (2048, 2048, 131072, 128, 128, False, True, True): (1, 1024, 1, 4), + (2048, 2048, 131072, 128, 128, True, False, True): (2, 1024, 1, 4), + (3072, 768, 256, 16, 16, False, True, True): (6, 4, 1, 4), + (3072, 768, 256, 16, 16, True, False, True): (3, 1, 4, 4), + (3072, 768, 256, 32, 32, False, True, True): (6, 8, 1, 2), + (3072, 768, 256, 32, 32, True, False, True): (1, 2, 4, 4), + (3072, 768, 256, 64, 64, False, True, True): (1, 4, 4, 4), + (3072, 768, 256, 64, 64, True, False, True): (4, 2, 4, 4), + (3072, 768, 256, 128, 128, False, True, True): (1, 2, 3, 8), + (3072, 768, 256, 128, 128, True, False, True): (1, 2, 3, 8), + (3072, 768, 512, 16, 16, False, True, True): (2, 4, 1, 4), + (3072, 768, 512, 16, 16, True, False, True): (1, 4, 4, 1), + (3072, 768, 512, 32, 32, False, True, True): (3, 8, 1, 4), + (3072, 768, 512, 32, 32, True, False, True): (1, 2, 3, 4), + (3072, 768, 512, 64, 64, False, True, True): (1, 8, 1, 4), + (3072, 768, 512, 64, 64, True, False, True): (4, 4, 3, 4), + (3072, 768, 512, 128, 128, False, True, True): (1, 4, 3, 8), + (3072, 768, 512, 128, 128, True, False, True): (1, 4, 3, 8), + (3072, 768, 1024, 16, 16, False, True, True): (1, 8, 1, 4), + (3072, 768, 1024, 16, 16, True, False, True): (3, 4, 3, 1), + (3072, 768, 1024, 32, 32, False, True, True): (1, 8, 1, 8), + (3072, 768, 1024, 32, 32, True, False, True): (1, 4, 4, 4), + (3072, 768, 1024, 64, 64, False, True, True): (1, 16, 3, 4), + (3072, 768, 1024, 64, 64, True, False, True): (1, 4, 3, 4), + (3072, 768, 1024, 128, 128, False, True, True): (1, 8, 3, 8), + (3072, 768, 1024, 128, 128, True, False, True): (2, 8, 3, 8), + (3072, 768, 2048, 16, 16, False, True, True): (3, 8, 1, 4), + (3072, 768, 2048, 16, 16, True, False, True): (2, 8, 3, 4), + (3072, 768, 2048, 32, 32, False, True, True): (3, 16, 1, 8), + (3072, 768, 2048, 32, 32, True, False, True): (3, 8, 3, 4), + (3072, 768, 2048, 64, 64, False, True, True): (1, 16, 1, 4), + (3072, 768, 2048, 64, 64, True, False, True): (1, 16, 3, 4), + (3072, 768, 2048, 128, 128, False, True, True): (1, 16, 3, 8), + (3072, 768, 2048, 128, 128, True, False, True): (2, 16, 2, 4), + (3072, 768, 4096, 16, 16, False, True, True): (1, 16, 1, 4), + (3072, 768, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (3072, 768, 4096, 32, 32, False, True, True): (2, 32, 1, 8), + (3072, 768, 4096, 32, 32, True, False, True): (7, 16, 3, 4), + (3072, 768, 4096, 64, 64, False, True, True): (2, 32, 1, 4), + (3072, 768, 4096, 64, 64, True, False, True): (2, 16, 2, 4), + (3072, 768, 4096, 128, 128, False, True, True): (1, 32, 3, 8), + (3072, 768, 4096, 128, 128, True, False, True): (3, 32, 2, 4), + (3072, 768, 8192, 16, 16, False, True, True): (2, 32, 1, 4), + (3072, 768, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (3072, 768, 8192, 32, 32, False, True, True): (4, 32, 1, 4), + (3072, 768, 8192, 32, 32, True, False, True): (4, 32, 3, 4), + (3072, 768, 8192, 64, 64, False, True, True): (2, 64, 1, 4), + (3072, 768, 8192, 64, 64, True, False, True): (4, 32, 2, 4), + (3072, 768, 8192, 128, 128, False, True, True): (3, 64, 1, 4), + (3072, 768, 8192, 128, 128, True, False, True): (6, 64, 2, 4), + (3072, 768, 16384, 16, 16, False, True, True): (1, 64, 1, 4), + (3072, 768, 16384, 16, 16, True, False, True): (1, 64, 1, 1), + (3072, 768, 16384, 32, 32, False, True, True): (1, 64, 1, 4), + (3072, 768, 16384, 32, 32, True, False, True): (4, 64, 3, 4), + (3072, 768, 16384, 64, 64, False, True, True): (4, 128, 1, 4), + (3072, 768, 16384, 64, 64, True, False, True): (4, 64, 2, 4), + (3072, 768, 16384, 128, 128, False, True, True): (3, 128, 1, 4), + (3072, 768, 16384, 128, 128, True, False, True): (4, 128, 2, 4), + (3072, 768, 32768, 16, 16, False, True, True): (1, 128, 1, 4), + (3072, 768, 32768, 16, 16, True, False, True): (8, 128, 4, 1), + (3072, 768, 32768, 32, 32, False, True, True): (1, 128, 1, 4), + (3072, 768, 32768, 32, 32, True, False, True): (8, 128, 3, 4), + (3072, 768, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (3072, 768, 32768, 64, 64, True, False, True): (1, 128, 2, 4), + (3072, 768, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (3072, 768, 32768, 128, 128, True, False, True): (8, 256, 2, 4), + (3072, 768, 50432, 16, 16, False, True, True): (1, 197, 1, 4), + (3072, 768, 50432, 16, 16, True, False, True): (7, 197, 4, 1), + (3072, 768, 50432, 32, 32, False, True, True): (1, 197, 1, 4), + (3072, 768, 50432, 32, 32, True, False, True): (4, 197, 3, 4), + (3072, 768, 50432, 64, 64, False, True, True): (1, 394, 1, 4), + (3072, 768, 50432, 64, 64, True, False, True): (3, 197, 2, 4), + (3072, 768, 50432, 128, 128, False, True, True): (3, 394, 1, 4), + (3072, 768, 50432, 128, 128, True, False, True): (8, 394, 2, 4), + (3072, 768, 65536, 16, 16, False, True, True): (1, 256, 1, 4), + (3072, 768, 65536, 16, 16, True, False, True): (15, 256, 4, 1), + (3072, 768, 65536, 32, 32, False, True, True): (1, 256, 1, 4), + (3072, 768, 65536, 32, 32, True, False, True): (15, 256, 3, 4), + (3072, 768, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (3072, 768, 65536, 64, 64, True, False, True): (2, 256, 2, 4), + (3072, 768, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (3072, 768, 65536, 128, 128, True, False, True): (3, 512, 2, 4), + (3072, 768, 131072, 16, 16, False, True, True): (1, 512, 1, 4), + (3072, 768, 131072, 16, 16, True, False, True): (15, 512, 4, 1), + (3072, 768, 131072, 32, 32, False, True, True): (1, 512, 1, 4), + (3072, 768, 131072, 32, 32, True, False, True): (9, 512, 3, 4), + (3072, 768, 131072, 64, 64, False, True, True): (1, 1024, 1, 4), + (3072, 768, 131072, 64, 64, True, False, True): (3, 512, 2, 4), + (3072, 768, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (3072, 768, 131072, 128, 128, True, False, True): (1, 1024, 2, 4), + (3072, 3072, 256, 16, 16, False, True, True): (5, 4, 1, 4), + (3072, 3072, 256, 16, 16, True, False, True): (1, 2, 5, 2), + (3072, 3072, 256, 32, 32, False, True, True): (5, 4, 1, 8), + (3072, 3072, 256, 32, 32, True, False, True): (1, 4, 4, 2), + (3072, 3072, 256, 64, 64, False, True, True): (2, 4, 4, 4), + (3072, 3072, 256, 64, 64, True, False, True): (2, 4, 4, 4), + (3072, 3072, 256, 128, 128, False, True, True): (1, 2, 3, 8), + (3072, 3072, 256, 128, 128, True, False, True): (1, 2, 3, 8), + (3072, 3072, 512, 16, 16, False, True, True): (5, 4, 1, 2), + (3072, 3072, 512, 16, 16, True, False, True): (1, 2, 3, 4), + (3072, 3072, 512, 32, 32, False, True, True): (3, 8, 1, 4), + (3072, 3072, 512, 32, 32, True, False, True): (1, 4, 4, 2), + (3072, 3072, 512, 64, 64, False, True, True): (1, 8, 2, 2), + (3072, 3072, 512, 64, 64, True, False, True): (2, 4, 3, 4), + (3072, 3072, 512, 128, 128, False, True, True): (2, 4, 3, 8), + (3072, 3072, 512, 128, 128, True, False, True): (1, 4, 3, 8), + (3072, 3072, 1024, 16, 16, False, True, True): (1, 8, 1, 4), + (3072, 3072, 1024, 16, 16, True, False, True): (2, 8, 3, 1), + (3072, 3072, 1024, 32, 32, False, True, True): (1, 16, 1, 4), + (3072, 3072, 1024, 32, 32, True, False, True): (1, 4, 4, 4), + (3072, 3072, 1024, 64, 64, False, True, True): (1, 8, 3, 4), + (3072, 3072, 1024, 64, 64, True, False, True): (2, 4, 3, 4), + (3072, 3072, 1024, 128, 128, False, True, True): (1, 8, 1, 4), + (3072, 3072, 1024, 128, 128, True, False, True): (2, 8, 3, 8), + (3072, 3072, 2048, 16, 16, False, True, True): (1, 16, 1, 2), + (3072, 3072, 2048, 16, 16, True, False, True): (2, 16, 4, 2), + (3072, 3072, 2048, 32, 32, False, True, True): (1, 16, 1, 8), + (3072, 3072, 2048, 32, 32, True, False, True): (3, 8, 4, 4), + (3072, 3072, 2048, 64, 64, False, True, True): (3, 16, 3, 4), + (3072, 3072, 2048, 64, 64, True, False, True): (3, 8, 3, 4), + (3072, 3072, 2048, 128, 128, False, True, True): (1, 16, 3, 8), + (3072, 3072, 2048, 128, 128, True, False, True): (5, 16, 3, 8), + (3072, 3072, 4096, 16, 16, False, True, True): (1, 32, 1, 2), + (3072, 3072, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (3072, 3072, 4096, 32, 32, False, True, True): (1, 32, 1, 8), + (3072, 3072, 4096, 32, 32, True, False, True): (3, 16, 3, 4), + (3072, 3072, 4096, 64, 64, False, True, True): (1, 32, 3, 4), + (3072, 3072, 4096, 64, 64, True, False, True): (3, 16, 3, 4), + (3072, 3072, 4096, 128, 128, False, True, True): (3, 32, 3, 8), + (3072, 3072, 4096, 128, 128, True, False, True): (3, 32, 3, 8), + (3072, 3072, 8192, 16, 16, False, True, True): (1, 64, 1, 2), + (3072, 3072, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (3072, 3072, 8192, 32, 32, False, True, True): (1, 64, 1, 8), + (3072, 3072, 8192, 32, 32, True, False, True): (6, 32, 3, 4), + (3072, 3072, 8192, 64, 64, False, True, True): (1, 64, 3, 4), + (3072, 3072, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (3072, 3072, 8192, 128, 128, False, True, True): (2, 64, 3, 8), + (3072, 3072, 8192, 128, 128, True, False, True): (1, 64, 3, 8), + (3072, 3072, 16384, 16, 16, False, True, True): (1, 128, 1, 2), + (3072, 3072, 16384, 16, 16, True, False, True): (4, 128, 4, 2), + (3072, 3072, 16384, 32, 32, False, True, True): (1, 64, 1, 2), + (3072, 3072, 16384, 32, 32, True, False, True): (4, 64, 3, 4), + (3072, 3072, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (3072, 3072, 16384, 64, 64, True, False, True): (4, 64, 3, 4), + (3072, 3072, 16384, 128, 128, False, True, True): (1, 128, 1, 4), + (3072, 3072, 16384, 128, 128, True, False, True): (1, 128, 3, 8), + (3072, 3072, 32768, 16, 16, False, True, True): (1, 256, 1, 2), + (3072, 3072, 32768, 16, 16, True, False, True): (8, 128, 4, 4), + (3072, 3072, 32768, 32, 32, False, True, True): (1, 256, 1, 8), + (3072, 3072, 32768, 32, 32, True, False, True): (5, 128, 3, 4), + (3072, 3072, 32768, 64, 64, False, True, True): (1, 256, 3, 4), + (3072, 3072, 32768, 64, 64, True, False, True): (3, 128, 3, 4), + (3072, 3072, 32768, 128, 128, False, True, True): (1, 256, 1, 4), + (3072, 3072, 32768, 128, 128, True, False, True): (3, 256, 2, 4), + (3072, 3072, 65536, 16, 16, False, True, True): (1, 512, 1, 2), + (3072, 3072, 65536, 16, 16, True, False, True): (7, 256, 4, 4), + (3072, 3072, 65536, 32, 32, False, True, True): (1, 256, 1, 2), + (3072, 3072, 65536, 32, 32, True, False, True): (5, 256, 3, 4), + (3072, 3072, 65536, 64, 64, False, True, True): (1, 512, 3, 4), + (3072, 3072, 65536, 64, 64, True, False, True): (3, 256, 3, 4), + (3072, 3072, 65536, 128, 128, False, True, True): (1, 512, 1, 4), + (3072, 3072, 65536, 128, 128, True, False, True): (3, 512, 2, 4), + (3072, 3072, 131072, 16, 16, False, True, True): (1, 1024, 1, 2), + (3072, 3072, 131072, 16, 16, True, False, True): (5, 512, 4, 4), + (3072, 3072, 131072, 32, 32, False, True, True): (1, 512, 1, 2), + (3072, 3072, 131072, 32, 32, True, False, True): (5, 512, 3, 4), + (3072, 3072, 131072, 64, 64, False, True, True): (1, 1024, 3, 4), + (3072, 3072, 131072, 64, 64, True, False, True): (3, 512, 3, 4), + (3072, 3072, 131072, 128, 128, False, True, True): (1, 1024, 1, 4), + (3072, 3072, 131072, 128, 128, True, False, True): (6, 1024, 2, 4), + (4096, 4096, 256, 16, 16, False, True, True): (2, 2, 5, 4), + (4096, 4096, 256, 16, 16, True, False, True): (2, 2, 4, 2), + (4096, 4096, 256, 32, 32, False, True, True): (1, 2, 4, 4), + (4096, 4096, 256, 32, 32, True, False, True): (3, 2, 4, 2), + (4096, 4096, 256, 64, 64, False, True, True): (3, 4, 3, 4), + (4096, 4096, 256, 64, 64, True, False, True): (1, 4, 3, 2), + (4096, 4096, 256, 128, 128, False, True, True): (1, 2, 2, 8), + (4096, 4096, 256, 128, 128, True, False, True): (1, 2, 2, 8), + (4096, 4096, 512, 16, 16, False, True, True): (4, 2, 3, 4), + (4096, 4096, 512, 16, 16, True, False, True): (1, 2, 3, 4), + (4096, 4096, 512, 32, 32, False, True, True): (1, 4, 3, 4), + (4096, 4096, 512, 32, 32, True, False, True): (3, 4, 3, 2), + (4096, 4096, 512, 64, 64, False, True, True): (4, 4, 4, 4), + (4096, 4096, 512, 64, 64, True, False, True): (3, 4, 3, 4), + (4096, 4096, 512, 128, 128, False, True, True): (2, 4, 2, 8), + (4096, 4096, 512, 128, 128, True, False, True): (2, 4, 1, 4), + (4096, 4096, 1024, 16, 16, False, True, True): (2, 8, 3, 2), + (4096, 4096, 1024, 16, 16, True, False, True): (2, 8, 3, 2), + (4096, 4096, 1024, 32, 32, False, True, True): (1, 8, 3, 4), + (4096, 4096, 1024, 32, 32, True, False, True): (1, 8, 3, 2), + (4096, 4096, 1024, 64, 64, False, True, True): (1, 8, 3, 4), + (4096, 4096, 1024, 64, 64, True, False, True): (1, 8, 3, 4), + (4096, 4096, 1024, 128, 128, False, True, True): (4, 8, 1, 4), + (4096, 4096, 1024, 128, 128, True, False, True): (2, 8, 2, 8), + (4096, 4096, 2048, 16, 16, False, True, True): (2, 8, 4, 4), + (4096, 4096, 2048, 16, 16, True, False, True): (2, 8, 4, 4), + (4096, 4096, 2048, 32, 32, False, True, True): (4, 8, 3, 8), + (4096, 4096, 2048, 32, 32, True, False, True): (4, 8, 4, 8), + (4096, 4096, 2048, 64, 64, False, True, True): (4, 16, 3, 4), + (4096, 4096, 2048, 64, 64, True, False, True): (4, 16, 3, 4), + (4096, 4096, 2048, 128, 128, False, True, True): (1, 16, 1, 4), + (4096, 4096, 2048, 128, 128, True, False, True): (4, 16, 1, 4), + (4096, 4096, 4096, 16, 16, False, True, True): (4, 32, 4, 4), + (4096, 4096, 4096, 16, 16, True, False, True): (2, 32, 4, 4), + (4096, 4096, 4096, 32, 32, False, True, True): (4, 16, 4, 8), + (4096, 4096, 4096, 32, 32, True, False, True): (4, 16, 4, 8), + (4096, 4096, 4096, 64, 64, False, True, True): (4, 32, 3, 4), + (4096, 4096, 4096, 64, 64, True, False, True): (2, 32, 3, 4), + (4096, 4096, 4096, 128, 128, False, True, True): (2, 32, 1, 4), + (4096, 4096, 4096, 128, 128, True, False, True): (2, 32, 1, 4), + (4096, 4096, 8192, 16, 16, False, True, True): (4, 64, 4, 2), + (4096, 4096, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (4096, 4096, 8192, 32, 32, False, True, True): (4, 32, 4, 8), + (4096, 4096, 8192, 32, 32, True, False, True): (4, 32, 4, 8), + (4096, 4096, 8192, 64, 64, False, True, True): (4, 64, 3, 4), + (4096, 4096, 8192, 64, 64, True, False, True): (4, 64, 3, 4), + (4096, 4096, 8192, 128, 128, False, True, True): (1, 64, 1, 4), + (4096, 4096, 8192, 128, 128, True, False, True): (1, 64, 1, 4), + (4096, 4096, 16384, 16, 16, False, True, True): (4, 64, 4, 4), + (4096, 4096, 16384, 16, 16, True, False, True): (4, 64, 4, 4), + (4096, 4096, 16384, 32, 32, False, True, True): (4, 64, 4, 8), + (4096, 4096, 16384, 32, 32, True, False, True): (4, 64, 4, 8), + (4096, 4096, 16384, 64, 64, False, True, True): (4, 128, 3, 4), + (4096, 4096, 16384, 64, 64, True, False, True): (4, 128, 3, 4), + (4096, 4096, 16384, 128, 128, False, True, True): (1, 128, 1, 4), + (4096, 4096, 16384, 128, 128, True, False, True): (1, 128, 1, 4), + (4096, 4096, 32768, 16, 16, False, True, True): (8, 128, 4, 4), + (4096, 4096, 32768, 16, 16, True, False, True): (5, 128, 4, 4), + (4096, 4096, 32768, 32, 32, False, True, True): (5, 128, 4, 4), + (4096, 4096, 32768, 32, 32, True, False, True): (3, 128, 4, 8), + (4096, 4096, 32768, 64, 64, False, True, True): (3, 256, 3, 4), + (4096, 4096, 32768, 64, 64, True, False, True): (2, 256, 3, 4), + (4096, 4096, 32768, 128, 128, False, True, True): (1, 256, 1, 4), + (4096, 4096, 32768, 128, 128, True, False, True): (1, 256, 1, 4), + (4096, 4096, 65536, 16, 16, False, True, True): (5, 256, 4, 4), + (4096, 4096, 65536, 16, 16, True, False, True): (5, 256, 4, 4), + (4096, 4096, 65536, 32, 32, False, True, True): (4, 256, 4, 8), + (4096, 4096, 65536, 32, 32, True, False, True): (4, 256, 4, 8), + (4096, 4096, 65536, 64, 64, False, True, True): (1, 512, 3, 4), + (4096, 4096, 65536, 64, 64, True, False, True): (3, 512, 3, 4), + (4096, 4096, 65536, 128, 128, False, True, True): (1, 512, 1, 4), + (4096, 4096, 65536, 128, 128, True, False, True): (1, 512, 1, 4), + (4096, 4096, 131072, 16, 16, False, True, True): (5, 512, 4, 4), + (4096, 4096, 131072, 16, 16, True, False, True): (5, 512, 4, 4), + (4096, 4096, 131072, 32, 32, False, True, True): (4, 512, 4, 4), + (4096, 4096, 131072, 32, 32, True, False, True): (2, 512, 3, 4), + (4096, 4096, 131072, 64, 64, False, True, True): (1, 1024, 3, 4), + (4096, 4096, 131072, 64, 64, True, False, True): (3, 1024, 3, 4), + (4096, 4096, 131072, 128, 128, False, True, True): (1, 1024, 1, 4), + (4096, 4096, 131072, 128, 128, True, False, True): (1, 1024, 1, 4), + (5120, 1280, 65792, 16, 16, False, True, True): (1, 257, 1, 4), + (5120, 1280, 65792, 16, 16, True, False, True): (11, 257, 4, 1), + (5120, 1280, 65792, 32, 32, False, True, True): (1, 257, 1, 4), + (5120, 1280, 65792, 32, 32, True, False, True): (5, 257, 3, 4), + (5120, 1280, 65792, 64, 64, False, True, True): (1, 514, 1, 4), + (5120, 1280, 65792, 64, 64, True, False, True): (5, 257, 2, 4), + (5120, 1280, 65792, 128, 128, False, True, True): (3, 514, 1, 4), + (5120, 1280, 65792, 128, 128, True, False, True): (7, 514, 2, 4), + (6144, 6144, 256, 16, 16, False, True, True): (1, 2, 1, 4), + (6144, 6144, 256, 16, 16, True, False, True): (3, 1, 4, 4), + (6144, 6144, 256, 32, 32, False, True, True): (3, 2, 1, 8), + (6144, 6144, 256, 32, 32, True, False, True): (1, 1, 4, 4), + (6144, 6144, 256, 64, 64, False, True, True): (4, 2, 3, 4), + (6144, 6144, 256, 64, 64, True, False, True): (3, 2, 4, 4), + (6144, 6144, 256, 128, 128, False, True, True): (2, 2, 3, 8), + (6144, 6144, 256, 128, 128, True, False, True): (1, 2, 3, 8), + (6144, 6144, 512, 16, 16, False, True, True): (4, 4, 1, 4), + (6144, 6144, 512, 16, 16, True, False, True): (3, 2, 3, 1), + (6144, 6144, 512, 32, 32, False, True, True): (1, 8, 1, 4), + (6144, 6144, 512, 32, 32, True, False, True): (1, 2, 3, 2), + (6144, 6144, 512, 64, 64, False, True, True): (2, 4, 3, 4), + (6144, 6144, 512, 64, 64, True, False, True): (2, 2, 3, 4), + (6144, 6144, 512, 128, 128, False, True, True): (1, 4, 3, 8), + (6144, 6144, 512, 128, 128, True, False, True): (1, 4, 3, 8), + (6144, 6144, 1024, 16, 16, False, True, True): (1, 8, 1, 2), + (6144, 6144, 1024, 16, 16, True, False, True): (4, 8, 4, 4), + (6144, 6144, 1024, 32, 32, False, True, True): (1, 8, 4, 2), + (6144, 6144, 1024, 32, 32, True, False, True): (1, 8, 4, 2), + (6144, 6144, 1024, 64, 64, False, True, True): (4, 8, 3, 4), + (6144, 6144, 1024, 64, 64, True, False, True): (1, 4, 3, 4), + (6144, 6144, 1024, 128, 128, False, True, True): (2, 8, 3, 8), + (6144, 6144, 1024, 128, 128, True, False, True): (1, 8, 3, 8), + (6144, 6144, 2048, 16, 16, False, True, True): (4, 4, 1, 4), + (6144, 6144, 2048, 16, 16, True, False, True): (2, 8, 4, 4), + (6144, 6144, 2048, 32, 32, False, True, True): (1, 16, 4, 2), + (6144, 6144, 2048, 32, 32, True, False, True): (4, 8, 4, 8), + (6144, 6144, 2048, 64, 64, False, True, True): (4, 16, 3, 4), + (6144, 6144, 2048, 64, 64, True, False, True): (2, 8, 3, 4), + (6144, 6144, 2048, 128, 128, False, True, True): (1, 16, 3, 8), + (6144, 6144, 2048, 128, 128, True, False, True): (4, 16, 3, 8), + (6144, 6144, 4096, 16, 16, False, True, True): (4, 8, 1, 4), + (6144, 6144, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (6144, 6144, 4096, 32, 32, False, True, True): (4, 16, 1, 2), + (6144, 6144, 4096, 32, 32, True, False, True): (2, 8, 3, 8), + (6144, 6144, 4096, 64, 64, False, True, True): (4, 32, 3, 4), + (6144, 6144, 4096, 64, 64, True, False, True): (4, 16, 3, 4), + (6144, 6144, 4096, 128, 128, False, True, True): (4, 32, 3, 8), + (6144, 6144, 4096, 128, 128, True, False, True): (4, 32, 3, 8), + (6144, 6144, 8192, 16, 16, False, True, True): (2, 16, 1, 2), + (6144, 6144, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (6144, 6144, 8192, 32, 32, False, True, True): (4, 32, 1, 2), + (6144, 6144, 8192, 32, 32, True, False, True): (4, 32, 4, 8), + (6144, 6144, 8192, 64, 64, False, True, True): (4, 64, 3, 4), + (6144, 6144, 8192, 64, 64, True, False, True): (4, 32, 3, 4), + (6144, 6144, 8192, 128, 128, False, True, True): (4, 64, 3, 8), + (6144, 6144, 8192, 128, 128, True, False, True): (4, 64, 3, 8), + (6144, 6144, 16384, 16, 16, False, True, True): (2, 32, 1, 2), + (6144, 6144, 16384, 16, 16, True, False, True): (4, 64, 4, 4), + (6144, 6144, 16384, 32, 32, False, True, True): (4, 64, 1, 2), + (6144, 6144, 16384, 32, 32, True, False, True): (4, 64, 3, 2), + (6144, 6144, 16384, 64, 64, False, True, True): (4, 128, 3, 4), + (6144, 6144, 16384, 64, 64, True, False, True): (2, 32, 3, 8), + (6144, 6144, 16384, 128, 128, False, True, True): (4, 128, 3, 8), + (6144, 6144, 16384, 128, 128, True, False, True): (4, 128, 3, 8), + (6144, 6144, 32768, 16, 16, False, True, True): (2, 64, 1, 2), + (6144, 6144, 32768, 16, 16, True, False, True): (3, 128, 4, 4), + (6144, 6144, 32768, 32, 32, False, True, True): (4, 128, 1, 2), + (6144, 6144, 32768, 32, 32, True, False, True): (3, 128, 3, 4), + (6144, 6144, 32768, 64, 64, False, True, True): (4, 256, 3, 4), + (6144, 6144, 32768, 64, 64, True, False, True): (2, 64, 3, 8), + (6144, 6144, 32768, 128, 128, False, True, True): (4, 256, 3, 8), + (6144, 6144, 32768, 128, 128, True, False, True): (4, 256, 3, 8), + (6144, 6144, 65536, 16, 16, False, True, True): (2, 128, 1, 2), + (6144, 6144, 65536, 16, 16, True, False, True): (4, 256, 4, 4), + (6144, 6144, 65536, 32, 32, False, True, True): (4, 256, 1, 2), + (6144, 6144, 65536, 32, 32, True, False, True): (4, 256, 3, 4), + (6144, 6144, 65536, 64, 64, False, True, True): (4, 512, 3, 4), + (6144, 6144, 65536, 64, 64, True, False, True): (2, 128, 3, 8), + (6144, 6144, 65536, 128, 128, False, True, True): (4, 512, 3, 8), + (6144, 6144, 65536, 128, 128, True, False, True): (4, 512, 3, 8), + (6144, 6144, 131072, 16, 16, False, True, True): (2, 256, 1, 2), + (6144, 6144, 131072, 16, 16, True, False, True): (5, 512, 4, 1), + (6144, 6144, 131072, 32, 32, False, True, True): (4, 512, 1, 2), + (6144, 6144, 131072, 32, 32, True, False, True): (4, 512, 3, 2), + (6144, 6144, 131072, 64, 64, False, True, True): (4, 1024, 3, 4), + (6144, 6144, 131072, 64, 64, True, False, True): (2, 256, 3, 8), + (6144, 6144, 131072, 128, 128, False, True, True): (4, 1024, 3, 8), + (6144, 6144, 131072, 128, 128, True, False, True): (4, 1024, 3, 8), + (8192, 8192, 256, 16, 16, False, True, True): (1, 1, 3, 4), + (8192, 8192, 256, 16, 16, True, False, True): (4, 1, 3, 4), + (8192, 8192, 256, 32, 32, False, True, True): (1, 2, 3, 4), + (8192, 8192, 256, 32, 32, True, False, True): (1, 2, 3, 4), + (8192, 8192, 256, 64, 64, False, True, True): (6, 2, 3, 8), + (8192, 8192, 256, 64, 64, True, False, True): (4, 2, 3, 8), + (8192, 8192, 256, 128, 128, False, True, True): (1, 2, 1, 4), + (8192, 8192, 256, 128, 128, True, False, True): (1, 2, 1, 4), + (8192, 8192, 512, 16, 16, False, True, True): (4, 4, 3, 2), + (8192, 8192, 512, 16, 16, True, False, True): (4, 4, 3, 4), + (8192, 8192, 512, 32, 32, False, True, True): (1, 4, 3, 4), + (8192, 8192, 512, 32, 32, True, False, True): (3, 4, 3, 2), + (8192, 8192, 512, 64, 64, False, True, True): (1, 4, 3, 4), + (8192, 8192, 512, 64, 64, True, False, True): (1, 4, 3, 4), + (8192, 8192, 512, 128, 128, False, True, True): (4, 4, 2, 8), + (8192, 8192, 512, 128, 128, True, False, True): (4, 4, 2, 8), + (8192, 8192, 1024, 16, 16, False, True, True): (4, 8, 4, 4), + (8192, 8192, 1024, 16, 16, True, False, True): (2, 8, 4, 4), + (8192, 8192, 1024, 32, 32, False, True, True): (2, 4, 4, 8), + (8192, 8192, 1024, 32, 32, True, False, True): (1, 4, 3, 4), + (8192, 8192, 1024, 64, 64, False, True, True): (4, 8, 3, 4), + (8192, 8192, 1024, 64, 64, True, False, True): (2, 8, 3, 4), + (8192, 8192, 1024, 128, 128, False, True, True): (4, 8, 1, 4), + (8192, 8192, 1024, 128, 128, True, False, True): (4, 8, 1, 4), + (8192, 8192, 2048, 16, 16, False, True, True): (2, 8, 4, 4), + (8192, 8192, 2048, 16, 16, True, False, True): (2, 8, 4, 4), + (8192, 8192, 2048, 32, 32, False, True, True): (2, 8, 4, 8), + (8192, 8192, 2048, 32, 32, True, False, True): (2, 8, 4, 8), + (8192, 8192, 2048, 64, 64, False, True, True): (4, 8, 2, 4), + (8192, 8192, 2048, 64, 64, True, False, True): (4, 16, 3, 4), + (8192, 8192, 2048, 128, 128, False, True, True): (4, 16, 1, 4), + (8192, 8192, 2048, 128, 128, True, False, True): (4, 16, 1, 4), + (8192, 8192, 4096, 16, 16, False, True, True): (4, 16, 4, 4), + (8192, 8192, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (8192, 8192, 4096, 32, 32, False, True, True): (2, 16, 4, 8), + (8192, 8192, 4096, 32, 32, True, False, True): (2, 16, 4, 8), + (8192, 8192, 4096, 64, 64, False, True, True): (4, 32, 3, 4), + (8192, 8192, 4096, 64, 64, True, False, True): (4, 16, 2, 4), + (8192, 8192, 4096, 128, 128, False, True, True): (4, 32, 1, 4), + (8192, 8192, 4096, 128, 128, True, False, True): (4, 32, 1, 4), + (8192, 8192, 8192, 16, 16, False, True, True): (4, 64, 4, 2), + (8192, 8192, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (8192, 8192, 8192, 32, 32, False, True, True): (2, 32, 4, 8), + (8192, 8192, 8192, 32, 32, True, False, True): (2, 32, 4, 8), + (8192, 8192, 8192, 64, 64, False, True, True): (4, 32, 3, 8), + (8192, 8192, 8192, 64, 64, True, False, True): (4, 32, 2, 4), + (8192, 8192, 8192, 128, 128, False, True, True): (4, 64, 1, 4), + (8192, 8192, 8192, 128, 128, True, False, True): (4, 64, 1, 4), + (8192, 8192, 16384, 16, 16, False, True, True): (4, 64, 4, 4), + (8192, 8192, 16384, 16, 16, True, False, True): (4, 64, 4, 4), + (8192, 8192, 16384, 32, 32, False, True, True): (4, 64, 3, 4), + (8192, 8192, 16384, 32, 32, True, False, True): (4, 64, 4, 8), + (8192, 8192, 16384, 64, 64, False, True, True): (4, 64, 2, 4), + (8192, 8192, 16384, 64, 64, True, False, True): (4, 64, 2, 4), + (8192, 8192, 16384, 128, 128, False, True, True): (4, 128, 1, 4), + (8192, 8192, 16384, 128, 128, True, False, True): (4, 128, 1, 4), + (8192, 8192, 32768, 16, 16, False, True, True): (3, 128, 4, 4), + (8192, 8192, 32768, 16, 16, True, False, True): (3, 128, 4, 4), + (8192, 8192, 32768, 32, 32, False, True, True): (2, 128, 4, 8), + (8192, 8192, 32768, 32, 32, True, False, True): (2, 128, 4, 8), + (8192, 8192, 32768, 64, 64, False, True, True): (2, 128, 2, 4), + (8192, 8192, 32768, 64, 64, True, False, True): (2, 128, 2, 4), + (8192, 8192, 32768, 128, 128, False, True, True): (4, 256, 1, 4), + (8192, 8192, 32768, 128, 128, True, False, True): (4, 256, 1, 4), + (8192, 8192, 65536, 16, 16, False, True, True): (3, 256, 4, 4), + (8192, 8192, 65536, 16, 16, True, False, True): (3, 256, 4, 4), + (8192, 8192, 65536, 32, 32, False, True, True): (2, 256, 3, 4), + (8192, 8192, 65536, 32, 32, True, False, True): (2, 256, 3, 4), + (8192, 8192, 65536, 64, 64, False, True, True): (2, 256, 2, 4), + (8192, 8192, 65536, 64, 64, True, False, True): (2, 256, 3, 8), + (8192, 8192, 65536, 128, 128, False, True, True): (4, 512, 1, 4), + (8192, 8192, 65536, 128, 128, True, False, True): (4, 512, 1, 4), + (8192, 8192, 131072, 16, 16, False, True, True): (3, 512, 4, 4), + (8192, 8192, 131072, 16, 16, True, False, True): (3, 512, 4, 4), + (8192, 8192, 131072, 32, 32, False, True, True): (2, 512, 4, 4), + (8192, 8192, 131072, 32, 32, True, False, True): (2, 512, 3, 4), + (8192, 8192, 131072, 64, 64, False, True, True): (4, 512, 2, 4), + (8192, 8192, 131072, 64, 64, True, False, True): (2, 512, 2, 4), + (8192, 8192, 131072, 128, 128, False, True, True): (4, 1024, 1, 4), + (8192, 8192, 131072, 128, 128, True, False, True): (4, 1024, 1, 4), + (16384, 16384, 256, 16, 16, False, True, True): (2, 2, 6, 4), + (16384, 16384, 256, 16, 16, True, False, True): (2, 2, 6, 4), + (16384, 16384, 256, 32, 32, False, True, True): (4, 2, 3, 2), + (16384, 16384, 256, 32, 32, True, False, True): (4, 2, 3, 2), + (16384, 16384, 256, 64, 64, False, True, True): (2, 2, 4, 4), + (16384, 16384, 256, 64, 64, True, False, True): (4, 2, 3, 8), + (16384, 16384, 256, 128, 128, False, True, True): (4, 2, 2, 8), + (16384, 16384, 256, 128, 128, True, False, True): (4, 2, 2, 8), + (16384, 16384, 512, 16, 16, False, True, True): (1, 2, 4, 4), + (16384, 16384, 512, 16, 16, True, False, True): (1, 2, 4, 4), + (16384, 16384, 512, 32, 32, False, True, True): (2, 2, 4, 8), + (16384, 16384, 512, 32, 32, True, False, True): (2, 2, 4, 8), + (16384, 16384, 512, 64, 64, False, True, True): (4, 4, 3, 4), + (16384, 16384, 512, 64, 64, True, False, True): (4, 4, 3, 4), + (16384, 16384, 512, 128, 128, False, True, True): (4, 4, 2, 8), + (16384, 16384, 512, 128, 128, True, False, True): (4, 4, 2, 8), + (16384, 16384, 1024, 16, 16, False, True, True): (3, 4, 4, 4), + (16384, 16384, 1024, 16, 16, True, False, True): (2, 8, 4, 4), + (16384, 16384, 1024, 32, 32, False, True, True): (2, 4, 4, 8), + (16384, 16384, 1024, 32, 32, True, False, True): (1, 4, 4, 8), + (16384, 16384, 1024, 64, 64, False, True, True): (2, 8, 3, 4), + (16384, 16384, 1024, 64, 64, True, False, True): (2, 8, 3, 4), + (16384, 16384, 1024, 128, 128, False, True, True): (4, 8, 1, 4), + (16384, 16384, 1024, 128, 128, True, False, True): (4, 8, 1, 4), + (16384, 16384, 2048, 16, 16, False, True, True): (2, 8, 4, 4), + (16384, 16384, 2048, 16, 16, True, False, True): (2, 8, 4, 4), + (16384, 16384, 2048, 32, 32, False, True, True): (1, 8, 4, 8), + (16384, 16384, 2048, 32, 32, True, False, True): (2, 8, 4, 8), + (16384, 16384, 2048, 64, 64, False, True, True): (2, 8, 2, 4), + (16384, 16384, 2048, 64, 64, True, False, True): (2, 8, 2, 4), + (16384, 16384, 2048, 128, 128, False, True, True): (4, 16, 1, 4), + (16384, 16384, 2048, 128, 128, True, False, True): (4, 16, 1, 4), + (16384, 16384, 4096, 16, 16, False, True, True): (2, 16, 4, 4), + (16384, 16384, 4096, 16, 16, True, False, True): (2, 16, 4, 4), + (16384, 16384, 4096, 32, 32, False, True, True): (1, 8, 3, 8), + (16384, 16384, 4096, 32, 32, True, False, True): (2, 16, 3, 4), + (16384, 16384, 4096, 64, 64, False, True, True): (2, 16, 2, 4), + (16384, 16384, 4096, 64, 64, True, False, True): (2, 16, 2, 4), + (16384, 16384, 4096, 128, 128, False, True, True): (4, 32, 1, 4), + (16384, 16384, 4096, 128, 128, True, False, True): (4, 32, 1, 4), + (16384, 16384, 8192, 16, 16, False, True, True): (4, 64, 4, 2), + (16384, 16384, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (16384, 16384, 8192, 32, 32, False, True, True): (2, 32, 4, 8), + (16384, 16384, 8192, 32, 32, True, False, True): (2, 32, 3, 4), + (16384, 16384, 8192, 64, 64, False, True, True): (2, 32, 4, 8), + (16384, 16384, 8192, 64, 64, True, False, True): (2, 32, 3, 8), + (16384, 16384, 8192, 128, 128, False, True, True): (4, 64, 1, 4), + (16384, 16384, 8192, 128, 128, True, False, True): (4, 64, 1, 4), + (16384, 16384, 16384, 16, 16, False, True, True): (1, 64, 4, 4), + (16384, 16384, 16384, 16, 16, True, False, True): (1, 64, 4, 4), + (16384, 16384, 16384, 32, 32, False, True, True): (1, 64, 3, 8), + (16384, 16384, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (16384, 16384, 16384, 64, 64, False, True, True): (1, 64, 2, 4), + (16384, 16384, 16384, 64, 64, True, False, True): (1, 64, 4, 8), + (16384, 16384, 16384, 128, 128, False, True, True): (4, 128, 1, 4), + (16384, 16384, 16384, 128, 128, True, False, True): (4, 128, 1, 4), + (16384, 16384, 32768, 16, 16, False, True, True): (1, 128, 4, 4), + (16384, 16384, 32768, 16, 16, True, False, True): (1, 128, 4, 4), + (16384, 16384, 32768, 32, 32, False, True, True): (1, 128, 4, 2), + (16384, 16384, 32768, 32, 32, True, False, True): (1, 128, 3, 8), + (16384, 16384, 32768, 64, 64, False, True, True): (2, 128, 2, 4), + (16384, 16384, 32768, 64, 64, True, False, True): (1, 128, 3, 8), + (16384, 16384, 32768, 128, 128, False, True, True): (4, 256, 1, 4), + (16384, 16384, 32768, 128, 128, True, False, True): (4, 256, 1, 4), + (16384, 16384, 65536, 16, 16, False, True, True): (1, 256, 4, 4), + (16384, 16384, 65536, 16, 16, True, False, True): (1, 256, 4, 4), + (16384, 16384, 65536, 32, 32, False, True, True): (1, 256, 3, 4), + (16384, 16384, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (16384, 16384, 65536, 64, 64, False, True, True): (1, 256, 2, 4), + (16384, 16384, 65536, 64, 64, True, False, True): (2, 256, 2, 4), + (16384, 16384, 65536, 128, 128, False, True, True): (4, 512, 1, 4), + (16384, 16384, 65536, 128, 128, True, False, True): (4, 512, 1, 4), + (16384, 16384, 131072, 16, 16, False, True, True): (2, 512, 4, 4), + (16384, 16384, 131072, 16, 16, True, False, True): (1, 512, 4, 4), + (16384, 16384, 131072, 32, 32, False, True, True): (1, 512, 4, 8), + (16384, 16384, 131072, 32, 32, True, False, True): (1, 512, 3, 4), + (16384, 16384, 131072, 64, 64, False, True, True): (2, 512, 2, 4), + (16384, 16384, 131072, 64, 64, True, False, True): (1, 512, 2, 4), + (16384, 16384, 131072, 128, 128, False, True, True): (4, 1024, 1, 4), + (16384, 16384, 131072, 128, 128, True, False, True): (4, 1024, 1, 4), + }, + ("bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.bfloat16, 0.56)): { + (192, 192, 256, 64, 64, False, True, True): (3, 4, 3, 4), + (192, 192, 256, 64, 64, True, False, True): (1, 4, 4, 4), + (192, 192, 512, 64, 64, False, True, True): (2, 8, 3, 4), + (192, 192, 512, 64, 64, True, False, True): (2, 8, 3, 4), + (192, 192, 1024, 64, 64, False, True, True): (1, 16, 3, 4), + (192, 192, 1024, 64, 64, True, False, True): (1, 16, 5, 4), + (192, 192, 2048, 64, 64, False, True, True): (3, 32, 3, 4), + (192, 192, 2048, 64, 64, True, False, True): (5, 32, 3, 4), + (192, 192, 4096, 64, 64, False, True, True): (1, 64, 4, 4), + (192, 192, 4096, 64, 64, True, False, True): (2, 32, 3, 4), + (192, 192, 8192, 64, 64, False, True, True): (1, 128, 2, 4), + (192, 192, 8192, 64, 64, True, False, True): (1, 64, 3, 4), + (192, 192, 16384, 64, 64, False, True, True): (1, 256, 1, 4), + (192, 192, 16384, 64, 64, True, False, True): (1, 64, 3, 4), + (192, 192, 32768, 64, 64, False, True, True): (2, 512, 1, 2), + (192, 192, 32768, 64, 64, True, False, True): (2, 256, 2, 4), + (192, 192, 65536, 64, 64, False, True, True): (3, 512, 1, 4), + (192, 192, 65536, 64, 64, True, False, True): (1, 512, 2, 4), + (192, 192, 131072, 64, 64, False, True, True): (5, 1024, 1, 4), + (192, 192, 131072, 64, 64, True, False, True): (4, 512, 2, 4), + (384, 384, 256, 128, 128, False, True, True): (3, 2, 3, 8), + (384, 384, 256, 128, 128, True, False, True): (1, 2, 3, 8), + (384, 384, 512, 128, 128, False, True, True): (4, 4, 3, 8), + (384, 384, 512, 128, 128, True, False, True): (3, 4, 3, 8), + (384, 384, 1024, 128, 128, False, True, True): (1, 8, 3, 8), + (384, 384, 1024, 128, 128, True, False, True): (2, 8, 3, 8), + (384, 384, 2048, 128, 128, False, True, True): (5, 16, 3, 8), + (384, 384, 2048, 128, 128, True, False, True): (5, 16, 3, 8), + (384, 384, 4096, 128, 128, False, True, True): (3, 32, 3, 8), + (384, 384, 4096, 128, 128, True, False, True): (6, 32, 3, 8), + (384, 384, 8192, 128, 128, False, True, True): (2, 64, 3, 8), + (384, 384, 8192, 128, 128, True, False, True): (4, 32, 2, 8), + (384, 384, 16384, 128, 128, False, True, True): (2, 128, 3, 8), + (384, 384, 16384, 128, 128, True, False, True): (5, 128, 2, 4), + (384, 384, 32768, 128, 128, False, True, True): (2, 256, 3, 8), + (384, 384, 32768, 128, 128, True, False, True): (3, 256, 2, 4), + (384, 384, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (384, 384, 65536, 128, 128, True, False, True): (1, 512, 2, 4), + (384, 384, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (384, 384, 131072, 128, 128, True, False, True): (1, 1024, 2, 4), + }, + ("bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.float16, 0.5)): { + (16, 16, 16, 16, 16, False, False, False): (1, 1, 1, 1), + (16, 16, 16, 16, 16, False, False, True): (1, 1, 2, 2), + (16, 16, 16, 16, 16, False, True, False): (1, 1, 1, 1), + (16, 16, 16, 16, 16, False, True, True): (1, 1, 1, 8), + (16, 16, 16, 16, 16, True, False, False): (3, 1, 3, 4), + (16, 16, 16, 16, 16, True, False, True): (1, 1, 2, 1), + (16, 16, 32, 16, 16, False, False, False): (1, 2, 1, 8), + (16, 16, 32, 16, 16, False, False, True): (1, 2, 1, 2), + (16, 16, 32, 16, 16, False, True, False): (2, 1, 1, 4), + (16, 16, 32, 16, 16, False, True, True): (1, 2, 1, 4), + (16, 16, 32, 16, 16, True, False, False): (1, 1, 1, 4), + (16, 16, 32, 16, 16, True, False, True): (1, 2, 1, 2), + (16, 16, 64, 16, 16, False, False, False): (1, 4, 1, 1), + (16, 16, 64, 16, 16, False, False, True): (1, 2, 2, 4), + (16, 16, 64, 16, 16, False, True, False): (1, 4, 1, 4), + (16, 16, 64, 16, 16, False, True, True): (1, 2, 1, 4), + (16, 16, 64, 16, 16, True, False, False): (1, 4, 1, 2), + (16, 16, 64, 16, 16, True, False, True): (1, 1, 1, 2), + (16, 32, 16, 16, 16, False, False, False): (1, 1, 2, 4), + (16, 32, 16, 16, 16, False, False, True): (1, 1, 1, 4), + (16, 32, 16, 16, 16, False, True, False): (1, 1, 1, 2), + (16, 32, 16, 16, 16, False, True, True): (1, 1, 1, 2), + (16, 32, 16, 16, 16, True, False, False): (1, 1, 2, 16), + (16, 32, 16, 16, 16, True, False, True): (1, 1, 1, 4), + (16, 32, 16, 16, 32, False, False, False): (2, 1, 1, 8), + (16, 32, 16, 16, 32, False, False, True): (2, 1, 1, 8), + (16, 32, 16, 16, 32, False, True, False): (1, 1, 2, 1), + (16, 32, 16, 16, 32, False, True, True): (1, 1, 1, 4), + (16, 32, 16, 16, 32, True, False, False): (2, 1, 1, 8), + (16, 32, 16, 16, 32, True, False, True): (1, 1, 2, 4), + (16, 32, 32, 16, 16, False, False, False): (1, 1, 1, 16), + (16, 32, 32, 16, 16, False, False, True): (1, 2, 1, 2), + (16, 32, 32, 16, 16, False, True, False): (1, 2, 1, 8), + (16, 32, 32, 16, 16, False, True, True): (3, 2, 1, 4), + (16, 32, 32, 16, 16, True, False, False): (1, 2, 1, 4), + (16, 32, 32, 16, 16, True, False, True): (1, 2, 1, 2), + (16, 32, 32, 16, 32, False, False, False): (1, 2, 1, 2), + (16, 32, 32, 16, 32, False, False, True): (1, 1, 1, 4), + (16, 32, 32, 16, 32, False, True, False): (1, 1, 2, 4), + (16, 32, 32, 16, 32, False, True, True): (1, 2, 1, 2), + (16, 32, 32, 16, 32, True, False, False): (1, 2, 1, 2), + (16, 32, 32, 16, 32, True, False, True): (1, 2, 1, 16), + (16, 32, 64, 16, 16, False, False, False): (1, 4, 1, 4), + (16, 32, 64, 16, 16, False, False, True): (2, 4, 1, 4), + (16, 32, 64, 16, 16, False, True, False): (1, 4, 1, 4), + (16, 32, 64, 16, 16, False, True, True): (1, 4, 1, 4), + (16, 32, 64, 16, 16, True, False, False): (3, 4, 1, 2), + (16, 32, 64, 16, 16, True, False, True): (1, 4, 1, 1), + (16, 32, 64, 16, 32, False, False, False): (1, 4, 1, 16), + (16, 32, 64, 16, 32, False, False, True): (1, 2, 1, 2), + (16, 32, 64, 16, 32, False, True, False): (1, 4, 2, 2), + (16, 32, 64, 16, 32, False, True, True): (1, 4, 1, 8), + (16, 32, 64, 16, 32, True, False, False): (1, 4, 1, 8), + (16, 32, 64, 16, 32, True, False, True): (1, 2, 1, 4), + (16, 64, 16, 16, 32, False, False, False): (1, 1, 1, 2), + (16, 64, 16, 16, 32, False, False, True): (1, 1, 1, 4), + (16, 64, 16, 16, 32, False, True, False): (2, 1, 2, 4), + (16, 64, 16, 16, 32, False, True, True): (1, 1, 1, 4), + (16, 64, 16, 16, 32, True, False, False): (1, 1, 1, 4), + (16, 64, 16, 16, 32, True, False, True): (1, 1, 1, 4), + (16, 64, 32, 16, 32, False, False, False): (1, 2, 1, 2), + (16, 64, 32, 16, 32, False, False, True): (1, 1, 1, 4), + (16, 64, 32, 16, 32, False, True, False): (1, 1, 1, 4), + (16, 64, 32, 16, 32, False, True, True): (1, 2, 3, 2), + (16, 64, 32, 16, 32, True, False, False): (1, 1, 1, 4), + (16, 64, 32, 16, 32, True, False, True): (1, 1, 2, 4), + (16, 64, 64, 16, 32, False, False, False): (1, 4, 1, 8), + (16, 64, 64, 16, 32, False, False, True): (1, 4, 1, 4), + (16, 64, 64, 16, 32, False, True, False): (1, 4, 1, 1), + (16, 64, 64, 16, 32, False, True, True): (2, 4, 1, 4), + (16, 64, 64, 16, 32, True, False, False): (1, 4, 1, 4), + (16, 64, 64, 16, 32, True, False, True): (1, 4, 1, 4), + (32, 16, 16, 16, 16, False, False, False): (2, 1, 2, 4), + (32, 16, 16, 16, 16, False, False, True): (2, 1, 1, 2), + (32, 16, 16, 16, 16, False, True, False): (1, 1, 2, 4), + (32, 16, 16, 16, 16, False, True, True): (1, 1, 1, 2), + (32, 16, 16, 16, 16, True, False, False): (1, 1, 1, 4), + (32, 16, 16, 16, 16, True, False, True): (2, 1, 1, 2), + (32, 16, 32, 16, 16, False, False, False): (1, 1, 1, 4), + (32, 16, 32, 16, 16, False, False, True): (1, 1, 1, 4), + (32, 16, 32, 16, 16, False, True, False): (1, 2, 1, 4), + (32, 16, 32, 16, 16, False, True, True): (2, 2, 1, 4), + (32, 16, 32, 16, 16, True, False, False): (2, 1, 1, 4), + (32, 16, 32, 16, 16, True, False, True): (2, 2, 1, 2), + (32, 16, 64, 16, 16, False, False, False): (1, 4, 1, 2), + (32, 16, 64, 16, 16, False, False, True): (1, 4, 1, 4), + (32, 16, 64, 16, 16, False, True, False): (1, 2, 1, 4), + (32, 16, 64, 16, 16, False, True, True): (1, 4, 1, 2), + (32, 16, 64, 16, 16, True, False, False): (1, 4, 2, 8), + (32, 16, 64, 16, 16, True, False, True): (1, 4, 1, 1), + (32, 32, 16, 16, 16, False, False, False): (1, 1, 1, 4), + (32, 32, 16, 16, 16, False, False, True): (2, 1, 1, 4), + (32, 32, 16, 16, 16, False, True, False): (1, 1, 2, 4), + (32, 32, 16, 16, 16, False, True, True): (1, 1, 2, 2), + (32, 32, 16, 16, 16, True, False, False): (1, 1, 1, 8), + (32, 32, 16, 16, 16, True, False, True): (1, 1, 1, 4), + (32, 32, 16, 16, 32, False, False, False): (1, 1, 3, 2), + (32, 32, 16, 16, 32, False, False, True): (2, 1, 1, 4), + (32, 32, 16, 16, 32, False, True, False): (3, 1, 1, 4), + (32, 32, 16, 16, 32, False, True, True): (1, 1, 1, 4), + (32, 32, 16, 16, 32, True, False, False): (2, 1, 1, 8), + (32, 32, 16, 16, 32, True, False, True): (1, 1, 3, 2), + (32, 32, 16, 32, 32, False, False, False): (1, 1, 1, 2), + (32, 32, 16, 32, 32, False, False, True): (2, 1, 1, 8), + (32, 32, 16, 32, 32, False, True, False): (1, 1, 1, 2), + (32, 32, 16, 32, 32, False, True, True): (1, 1, 1, 8), + (32, 32, 16, 32, 32, True, False, False): (1, 1, 2, 4), + (32, 32, 16, 32, 32, True, False, True): (1, 1, 1, 2), + (32, 32, 32, 16, 16, False, False, False): (1, 1, 1, 4), + (32, 32, 32, 16, 16, False, False, True): (1, 2, 1, 4), + (32, 32, 32, 16, 16, False, True, False): (1, 2, 1, 4), + (32, 32, 32, 16, 16, False, True, True): (1, 2, 1, 2), + (32, 32, 32, 16, 16, True, False, False): (1, 2, 1, 4), + (32, 32, 32, 16, 16, True, False, True): (1, 2, 1, 4), + (32, 32, 32, 16, 32, False, False, False): (1, 2, 1, 4), + (32, 32, 32, 16, 32, False, False, True): (1, 2, 1, 2), + (32, 32, 32, 16, 32, False, True, False): (1, 2, 1, 4), + (32, 32, 32, 16, 32, False, True, True): (1, 2, 1, 2), + (32, 32, 32, 16, 32, True, False, False): (1, 2, 1, 1), + (32, 32, 32, 16, 32, True, False, True): (1, 2, 1, 2), + (32, 32, 32, 32, 32, False, False, False): (1, 1, 1, 4), + (32, 32, 32, 32, 32, False, False, True): (2, 1, 1, 4), + (32, 32, 32, 32, 32, False, True, False): (1, 1, 1, 8), + (32, 32, 32, 32, 32, False, True, True): (1, 1, 1, 8), + (32, 32, 32, 32, 32, True, False, False): (1, 1, 3, 4), + (32, 32, 32, 32, 32, True, False, True): (1, 1, 1, 8), + (32, 32, 64, 16, 16, False, False, False): (1, 4, 1, 4), + (32, 32, 64, 16, 16, False, False, True): (1, 4, 1, 2), + (32, 32, 64, 16, 16, False, True, False): (1, 1, 1, 4), + (32, 32, 64, 16, 16, False, True, True): (1, 4, 1, 4), + (32, 32, 64, 16, 16, True, False, False): (1, 4, 1, 8), + (32, 32, 64, 16, 16, True, False, True): (1, 4, 1, 2), + (32, 32, 64, 16, 32, False, False, False): (1, 1, 1, 4), + (32, 32, 64, 16, 32, False, False, True): (1, 4, 1, 4), + (32, 32, 64, 16, 32, False, True, False): (1, 1, 1, 4), + (32, 32, 64, 16, 32, False, True, True): (1, 4, 1, 4), + (32, 32, 64, 16, 32, True, False, False): (2, 2, 1, 8), + (32, 32, 64, 16, 32, True, False, True): (1, 2, 1, 2), + (32, 32, 64, 32, 32, False, False, False): (1, 2, 1, 4), + (32, 32, 64, 32, 32, False, False, True): (1, 2, 1, 1), + (32, 32, 64, 32, 32, False, True, False): (1, 2, 2, 8), + (32, 32, 64, 32, 32, False, True, True): (1, 1, 1, 4), + (32, 32, 64, 32, 32, True, False, False): (1, 2, 1, 4), + (32, 32, 64, 32, 32, True, False, True): (2, 2, 1, 4), + (32, 64, 16, 16, 32, False, False, False): (1, 1, 1, 8), + (32, 64, 16, 16, 32, False, False, True): (1, 1, 1, 4), + (32, 64, 16, 16, 32, False, True, False): (2, 1, 1, 4), + (32, 64, 16, 16, 32, False, True, True): (1, 1, 1, 4), + (32, 64, 16, 16, 32, True, False, False): (1, 1, 2, 4), + (32, 64, 16, 16, 32, True, False, True): (1, 1, 2, 2), + (32, 64, 16, 32, 32, False, False, False): (1, 1, 1, 8), + (32, 64, 16, 32, 32, False, False, True): (2, 1, 1, 4), + (32, 64, 16, 32, 32, False, True, False): (1, 1, 1, 4), + (32, 64, 16, 32, 32, False, True, True): (1, 1, 2, 2), + (32, 64, 16, 32, 32, True, False, False): (1, 1, 1, 2), + (32, 64, 16, 32, 32, True, False, True): (2, 1, 2, 4), + (32, 64, 32, 16, 32, False, False, False): (1, 1, 1, 4), + (32, 64, 32, 16, 32, False, False, True): (1, 2, 1, 2), + (32, 64, 32, 16, 32, False, True, False): (1, 2, 3, 4), + (32, 64, 32, 16, 32, False, True, True): (2, 2, 1, 4), + (32, 64, 32, 16, 32, True, False, False): (1, 1, 1, 4), + (32, 64, 32, 16, 32, True, False, True): (1, 2, 2, 1), + (32, 64, 32, 32, 32, False, False, False): (1, 1, 1, 8), + (32, 64, 32, 32, 32, False, False, True): (1, 1, 1, 4), + (32, 64, 32, 32, 32, False, True, False): (1, 1, 2, 4), + (32, 64, 32, 32, 32, False, True, True): (1, 1, 1, 4), + (32, 64, 32, 32, 32, True, False, False): (2, 1, 1, 2), + (32, 64, 32, 32, 32, True, False, True): (1, 1, 1, 4), + (32, 64, 64, 16, 32, False, False, False): (1, 4, 2, 1), + (32, 64, 64, 16, 32, False, False, True): (3, 4, 1, 4), + (32, 64, 64, 16, 32, False, True, False): (1, 1, 1, 8), + (32, 64, 64, 16, 32, False, True, True): (1, 4, 1, 4), + (32, 64, 64, 16, 32, True, False, False): (1, 4, 1, 4), + (32, 64, 64, 16, 32, True, False, True): (2, 2, 3, 4), + (32, 64, 64, 32, 32, False, False, False): (1, 2, 1, 4), + (32, 64, 64, 32, 32, False, False, True): (1, 2, 1, 4), + (32, 64, 64, 32, 32, False, True, False): (1, 2, 2, 8), + (32, 64, 64, 32, 32, False, True, True): (1, 2, 1, 4), + (32, 64, 64, 32, 32, True, False, False): (1, 2, 2, 4), + (32, 64, 64, 32, 32, True, False, True): (1, 2, 1, 4), + (64, 32, 16, 32, 32, False, False, False): (1, 1, 1, 1), + (64, 32, 16, 32, 32, False, False, True): (1, 1, 2, 4), + (64, 32, 16, 32, 32, False, True, False): (2, 1, 1, 8), + (64, 32, 16, 32, 32, False, True, True): (1, 1, 1, 4), + (64, 32, 16, 32, 32, True, False, False): (2, 1, 1, 2), + (64, 32, 16, 32, 32, True, False, True): (1, 1, 1, 4), + (64, 32, 32, 32, 32, False, False, False): (3, 1, 1, 4), + (64, 32, 32, 32, 32, False, False, True): (1, 1, 1, 4), + (64, 32, 32, 32, 32, False, True, False): (1, 1, 1, 8), + (64, 32, 32, 32, 32, False, True, True): (1, 1, 1, 2), + (64, 32, 32, 32, 32, True, False, False): (1, 1, 1, 2), + (64, 32, 32, 32, 32, True, False, True): (1, 1, 1, 4), + (64, 32, 64, 32, 32, False, False, False): (1, 2, 1, 2), + (64, 32, 64, 32, 32, False, False, True): (3, 2, 1, 4), + (64, 32, 64, 32, 32, False, True, False): (1, 1, 1, 1), + (64, 32, 64, 32, 32, False, True, True): (1, 2, 1, 4), + (64, 32, 64, 32, 32, True, False, False): (1, 1, 3, 4), + (64, 32, 64, 32, 32, True, False, True): (1, 2, 2, 4), + (64, 64, 16, 32, 32, False, False, False): (1, 1, 2, 2), + (64, 64, 16, 32, 32, False, False, True): (1, 1, 3, 2), + (64, 64, 16, 32, 32, False, True, False): (1, 1, 1, 8), + (64, 64, 16, 32, 32, False, True, True): (1, 1, 2, 4), + (64, 64, 16, 32, 32, True, False, False): (1, 1, 2, 4), + (64, 64, 16, 32, 32, True, False, True): (2, 1, 2, 4), + (64, 64, 32, 32, 32, False, False, False): (1, 1, 2, 8), + (64, 64, 32, 32, 32, False, False, True): (1, 1, 2, 4), + (64, 64, 32, 32, 32, False, True, False): (1, 1, 1, 4), + (64, 64, 32, 32, 32, False, True, True): (1, 1, 1, 4), + (64, 64, 32, 32, 32, True, False, False): (1, 1, 1, 4), + (64, 64, 32, 32, 32, True, False, True): (2, 1, 2, 4), + (64, 64, 64, 32, 32, False, False, False): (1, 2, 1, 4), + (64, 64, 64, 32, 32, False, False, True): (1, 2, 1, 4), + (64, 64, 64, 32, 32, False, True, False): (1, 2, 1, 4), + (64, 64, 64, 32, 32, False, True, True): (3, 2, 1, 4), + (64, 64, 64, 32, 32, True, False, False): (1, 2, 1, 8), + (64, 64, 64, 32, 32, True, False, True): (1, 2, 3, 4), + (192, 192, 256, 16, 16, False, True, True): (1, 8, 4, 2), + (192, 192, 256, 16, 16, True, False, True): (1, 4, 4, 4), + (192, 192, 256, 32, 32, False, True, True): (2, 8, 5, 4), + (192, 192, 256, 32, 32, True, False, True): (2, 8, 5, 1), + (192, 192, 512, 16, 16, False, True, True): (3, 8, 4, 4), + (192, 192, 512, 16, 16, True, False, True): (5, 8, 5, 4), + (192, 192, 512, 32, 32, False, True, True): (1, 16, 5, 4), + (192, 192, 512, 32, 32, True, False, True): (1, 8, 6, 2), + (192, 192, 1024, 16, 16, False, True, True): (1, 16, 4, 4), + (192, 192, 1024, 16, 16, True, False, True): (3, 16, 5, 2), + (192, 192, 1024, 32, 32, False, True, True): (3, 16, 4, 4), + (192, 192, 1024, 32, 32, True, False, True): (1, 16, 5, 4), + (192, 192, 2048, 16, 16, False, True, True): (2, 16, 3, 4), + (192, 192, 2048, 16, 16, True, False, True): (1, 16, 4, 4), + (192, 192, 2048, 32, 32, False, True, True): (1, 32, 3, 4), + (192, 192, 2048, 32, 32, True, False, True): (3, 16, 4, 4), + (192, 192, 4096, 16, 16, False, True, True): (1, 64, 1, 4), + (192, 192, 4096, 16, 16, True, False, True): (1, 16, 3, 4), + (192, 192, 4096, 32, 32, False, True, True): (1, 128, 1, 4), + (192, 192, 4096, 32, 32, True, False, True): (2, 32, 4, 2), + (192, 192, 8192, 16, 16, False, True, True): (1, 64, 1, 4), + (192, 192, 8192, 16, 16, True, False, True): (2, 64, 3, 2), + (192, 192, 8192, 32, 32, False, True, True): (1, 128, 1, 4), + (192, 192, 8192, 32, 32, True, False, True): (4, 32, 3, 4), + (192, 192, 16384, 16, 16, False, True, True): (1, 128, 1, 4), + (192, 192, 16384, 16, 16, True, False, True): (1, 64, 3, 2), + (192, 192, 16384, 32, 32, False, True, True): (1, 128, 1, 4), + (192, 192, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (192, 192, 32768, 16, 16, False, True, True): (2, 256, 1, 2), + (192, 192, 32768, 16, 16, True, False, True): (1, 128, 3, 2), + (192, 192, 32768, 32, 32, False, True, True): (2, 256, 1, 4), + (192, 192, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (192, 192, 65536, 16, 16, False, True, True): (2, 512, 1, 2), + (192, 192, 65536, 16, 16, True, False, True): (1, 256, 3, 2), + (192, 192, 65536, 32, 32, False, True, True): (2, 512, 1, 4), + (192, 192, 65536, 32, 32, True, False, True): (2, 256, 3, 4), + (192, 192, 131072, 16, 16, False, True, True): (4, 1024, 1, 2), + (192, 192, 131072, 16, 16, True, False, True): (3, 512, 3, 2), + (192, 192, 131072, 32, 32, False, True, True): (1, 1024, 1, 4), + (192, 192, 131072, 32, 32, True, False, True): (3, 512, 3, 4), + (256, 256, 256, 16, 16, False, True, True): (4, 8, 6, 2), + (256, 256, 256, 16, 16, True, False, True): (5, 16, 5, 1), + (256, 256, 256, 32, 32, False, True, True): (1, 8, 7, 4), + (256, 256, 256, 32, 32, True, False, True): (1, 8, 5, 4), + (256, 256, 256, 64, 64, False, True, True): (1, 4, 5, 4), + (256, 256, 256, 64, 64, True, False, True): (2, 4, 3, 4), + (256, 256, 256, 128, 128, False, True, True): (1, 2, 2, 8), + (256, 256, 256, 128, 128, True, False, True): (1, 2, 2, 8), + (256, 256, 512, 16, 16, False, True, True): (4, 8, 4, 4), + (256, 256, 512, 16, 16, True, False, True): (4, 8, 6, 2), + (256, 256, 512, 32, 32, False, True, True): (3, 8, 5, 4), + (256, 256, 512, 32, 32, True, False, True): (2, 8, 5, 4), + (256, 256, 512, 64, 64, False, True, True): (2, 8, 4, 4), + (256, 256, 512, 64, 64, True, False, True): (1, 8, 7, 4), + (256, 256, 512, 128, 128, False, True, True): (2, 4, 2, 8), + (256, 256, 512, 128, 128, True, False, True): (5, 4, 2, 8), + (256, 256, 1024, 16, 16, False, True, True): (1, 8, 4, 4), + (256, 256, 1024, 16, 16, True, False, True): (1, 16, 4, 2), + (256, 256, 1024, 32, 32, False, True, True): (5, 32, 5, 1), + (256, 256, 1024, 32, 32, True, False, True): (1, 16, 4, 2), + (256, 256, 1024, 64, 64, False, True, True): (1, 16, 4, 4), + (256, 256, 1024, 64, 64, True, False, True): (2, 16, 3, 4), + (256, 256, 1024, 128, 128, False, True, True): (9, 8, 2, 8), + (256, 256, 1024, 128, 128, True, False, True): (1, 8, 2, 8), + (256, 256, 2048, 16, 16, False, True, True): (6, 32, 5, 2), + (256, 256, 2048, 16, 16, True, False, True): (2, 32, 4, 2), + (256, 256, 2048, 32, 32, False, True, True): (1, 32, 3, 2), + (256, 256, 2048, 32, 32, True, False, True): (1, 32, 3, 2), + (256, 256, 2048, 64, 64, False, True, True): (2, 32, 4, 4), + (256, 256, 2048, 64, 64, True, False, True): (2, 16, 4, 4), + (256, 256, 2048, 128, 128, False, True, True): (3, 16, 2, 8), + (256, 256, 2048, 128, 128, True, False, True): (4, 16, 2, 8), + (256, 256, 4096, 16, 16, False, True, True): (1, 32, 3, 4), + (256, 256, 4096, 16, 16, True, False, True): (3, 16, 3, 2), + (256, 256, 4096, 32, 32, False, True, True): (3, 32, 3, 2), + (256, 256, 4096, 32, 32, True, False, True): (1, 32, 3, 2), + (256, 256, 4096, 64, 64, False, True, True): (2, 32, 3, 4), + (256, 256, 4096, 64, 64, True, False, True): (2, 32, 3, 4), + (256, 256, 4096, 128, 128, False, True, True): (5, 32, 2, 8), + (256, 256, 4096, 128, 128, True, False, True): (1, 32, 2, 8), + (256, 256, 8192, 16, 16, False, True, True): (8, 32, 3, 4), + (256, 256, 8192, 16, 16, True, False, True): (1, 32, 3, 2), + (256, 256, 8192, 32, 32, False, True, True): (3, 64, 3, 4), + (256, 256, 8192, 32, 32, True, False, True): (2, 128, 1, 2), + (256, 256, 8192, 64, 64, False, True, True): (7, 128, 1, 4), + (256, 256, 8192, 64, 64, True, False, True): (4, 128, 1, 4), + (256, 256, 8192, 128, 128, False, True, True): (2, 64, 1, 4), + (256, 256, 8192, 128, 128, True, False, True): (4, 64, 1, 4), + (256, 256, 16384, 16, 16, False, True, True): (4, 128, 3, 2), + (256, 256, 16384, 16, 16, True, False, True): (5, 64, 3, 2), + (256, 256, 16384, 32, 32, False, True, True): (5, 128, 3, 2), + (256, 256, 16384, 32, 32, True, False, True): (5, 128, 3, 2), + (256, 256, 16384, 64, 64, False, True, True): (1, 256, 1, 4), + (256, 256, 16384, 64, 64, True, False, True): (5, 128, 3, 4), + (256, 256, 16384, 128, 128, False, True, True): (11, 128, 2, 8), + (256, 256, 16384, 128, 128, True, False, True): (3, 128, 1, 4), + (256, 256, 32768, 16, 16, False, True, True): (1, 128, 3, 4), + (256, 256, 32768, 16, 16, True, False, True): (2, 128, 3, 2), + (256, 256, 32768, 32, 32, False, True, True): (4, 256, 3, 2), + (256, 256, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (256, 256, 32768, 64, 64, False, True, True): (2, 256, 1, 4), + (256, 256, 32768, 64, 64, True, False, True): (2, 256, 1, 4), + (256, 256, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (256, 256, 32768, 128, 128, True, False, True): (2, 256, 1, 4), + (256, 256, 50432, 16, 16, False, True, True): (4, 197, 1, 4), + (256, 256, 50432, 16, 16, True, False, True): (4, 197, 3, 2), + (256, 256, 50432, 32, 32, False, True, True): (1, 394, 1, 2), + (256, 256, 50432, 32, 32, True, False, True): (4, 197, 3, 4), + (256, 256, 50432, 64, 64, False, True, True): (6, 394, 1, 4), + (256, 256, 50432, 64, 64, True, False, True): (4, 394, 2, 4), + (256, 256, 50432, 128, 128, False, True, True): (3, 394, 1, 4), + (256, 256, 50432, 128, 128, True, False, True): (1, 394, 2, 4), + (256, 256, 65536, 16, 16, False, True, True): (1, 256, 3, 2), + (256, 256, 65536, 16, 16, True, False, True): (1, 256, 3, 2), + (256, 256, 65536, 32, 32, False, True, True): (1, 512, 3, 2), + (256, 256, 65536, 32, 32, True, False, True): (4, 512, 3, 2), + (256, 256, 65536, 64, 64, False, True, True): (2, 512, 1, 4), + (256, 256, 65536, 64, 64, True, False, True): (5, 512, 1, 4), + (256, 256, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (256, 256, 65536, 128, 128, True, False, True): (1, 512, 1, 4), + (256, 256, 65792, 16, 16, False, True, True): (2, 257, 1, 4), + (256, 256, 65792, 16, 16, True, False, True): (1, 257, 3, 2), + (256, 256, 65792, 32, 32, False, True, True): (2, 257, 1, 4), + (256, 256, 65792, 32, 32, True, False, True): (1, 257, 3, 4), + (256, 256, 65792, 64, 64, False, True, True): (2, 514, 1, 4), + (256, 256, 65792, 64, 64, True, False, True): (2, 514, 2, 4), + (256, 256, 65792, 128, 128, False, True, True): (3, 514, 1, 4), + (256, 256, 65792, 128, 128, True, False, True): (1, 514, 2, 4), + (256, 256, 131072, 16, 16, False, True, True): (1, 512, 3, 1), + (256, 256, 131072, 16, 16, True, False, True): (1, 512, 3, 2), + (256, 256, 131072, 32, 32, False, True, True): (2, 1024, 3, 2), + (256, 256, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (256, 256, 131072, 64, 64, False, True, True): (1, 1024, 1, 4), + (256, 256, 131072, 64, 64, True, False, True): (1, 1024, 1, 4), + (256, 256, 131072, 128, 128, False, True, True): (7, 1024, 1, 4), + (256, 256, 131072, 128, 128, True, False, True): (1, 1024, 1, 4), + (384, 384, 256, 16, 16, False, True, True): (3, 16, 4, 1), + (384, 384, 256, 16, 16, True, False, True): (2, 4, 6, 2), + (384, 384, 256, 32, 32, False, True, True): (1, 8, 4, 4), + (384, 384, 256, 32, 32, True, False, True): (1, 4, 5, 2), + (384, 384, 256, 64, 64, False, True, True): (3, 4, 3, 4), + (384, 384, 256, 64, 64, True, False, True): (4, 4, 5, 4), + (384, 384, 512, 16, 16, False, True, True): (1, 16, 4, 1), + (384, 384, 512, 16, 16, True, False, True): (1, 8, 5, 2), + (384, 384, 512, 32, 32, False, True, True): (4, 16, 4, 2), + (384, 384, 512, 32, 32, True, False, True): (1, 8, 5, 2), + (384, 384, 512, 64, 64, False, True, True): (2, 8, 3, 4), + (384, 384, 512, 64, 64, True, False, True): (1, 8, 4, 4), + (384, 384, 1024, 16, 16, False, True, True): (1, 16, 4, 2), + (384, 384, 1024, 16, 16, True, False, True): (7, 8, 5, 2), + (384, 384, 1024, 32, 32, False, True, True): (2, 16, 3, 4), + (384, 384, 1024, 32, 32, True, False, True): (1, 16, 4, 2), + (384, 384, 1024, 64, 64, False, True, True): (6, 16, 3, 4), + (384, 384, 1024, 64, 64, True, False, True): (4, 16, 4, 4), + (384, 384, 2048, 16, 16, False, True, True): (1, 32, 1, 4), + (384, 384, 2048, 16, 16, True, False, True): (1, 16, 3, 2), + (384, 384, 2048, 32, 32, False, True, True): (1, 32, 1, 8), + (384, 384, 2048, 32, 32, True, False, True): (1, 8, 4, 4), + (384, 384, 2048, 64, 64, False, True, True): (2, 32, 1, 8), + (384, 384, 2048, 64, 64, True, False, True): (3, 16, 3, 4), + (384, 384, 4096, 16, 16, False, True, True): (5, 32, 1, 4), + (384, 384, 4096, 16, 16, True, False, True): (1, 32, 3, 2), + (384, 384, 4096, 32, 32, False, True, True): (1, 32, 1, 8), + (384, 384, 4096, 32, 32, True, False, True): (2, 16, 4, 4), + (384, 384, 4096, 64, 64, False, True, True): (1, 64, 1, 4), + (384, 384, 4096, 64, 64, True, False, True): (2, 32, 3, 4), + (384, 384, 8192, 16, 16, False, True, True): (2, 64, 1, 4), + (384, 384, 8192, 16, 16, True, False, True): (3, 32, 3, 2), + (384, 384, 8192, 32, 32, False, True, True): (4, 128, 1, 4), + (384, 384, 8192, 32, 32, True, False, True): (1, 32, 3, 2), + (384, 384, 8192, 64, 64, False, True, True): (1, 128, 1, 4), + (384, 384, 8192, 64, 64, True, False, True): (1, 64, 3, 4), + (384, 384, 16384, 16, 16, False, True, True): (1, 128, 1, 2), + (384, 384, 16384, 16, 16, True, False, True): (1, 64, 3, 2), + (384, 384, 16384, 32, 32, False, True, True): (1, 128, 1, 4), + (384, 384, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (384, 384, 16384, 64, 64, False, True, True): (5, 128, 3, 4), + (384, 384, 16384, 64, 64, True, False, True): (1, 128, 3, 4), + (384, 384, 32768, 16, 16, False, True, True): (2, 256, 1, 2), + (384, 384, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (384, 384, 32768, 32, 32, False, True, True): (1, 256, 1, 2), + (384, 384, 32768, 32, 32, True, False, True): (2, 128, 3, 4), + (384, 384, 32768, 64, 64, False, True, True): (3, 256, 1, 4), + (384, 384, 32768, 64, 64, True, False, True): (2, 256, 3, 4), + (384, 384, 65536, 16, 16, False, True, True): (2, 128, 1, 4), + (384, 384, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (384, 384, 65536, 32, 32, False, True, True): (1, 512, 1, 2), + (384, 384, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (384, 384, 65536, 64, 64, False, True, True): (3, 512, 1, 4), + (384, 384, 65536, 64, 64, True, False, True): (3, 256, 3, 4), + (384, 384, 131072, 16, 16, False, True, True): (2, 256, 1, 2), + (384, 384, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (384, 384, 131072, 32, 32, False, True, True): (1, 512, 1, 2), + (384, 384, 131072, 32, 32, True, False, True): (1, 512, 3, 4), + (384, 384, 131072, 64, 64, False, True, True): (3, 1024, 1, 4), + (384, 384, 131072, 64, 64, True, False, True): (3, 512, 3, 4), + (512, 512, 256, 16, 16, False, True, True): (1, 8, 5, 1), + (512, 512, 256, 16, 16, True, False, True): (2, 16, 5, 1), + (512, 512, 256, 32, 32, False, True, True): (2, 8, 5, 2), + (512, 512, 256, 32, 32, True, False, True): (4, 4, 5, 2), + (512, 512, 256, 64, 64, False, True, True): (1, 4, 5, 4), + (512, 512, 256, 64, 64, True, False, True): (3, 4, 5, 4), + (512, 512, 256, 128, 128, False, True, True): (1, 2, 2, 8), + (512, 512, 256, 128, 128, True, False, True): (1, 2, 2, 8), + (512, 512, 512, 16, 16, False, True, True): (1, 8, 4, 4), + (512, 512, 512, 16, 16, True, False, True): (4, 16, 5, 1), + (512, 512, 512, 32, 32, False, True, True): (4, 8, 5, 2), + (512, 512, 512, 32, 32, True, False, True): (7, 16, 4, 1), + (512, 512, 512, 64, 64, False, True, True): (3, 8, 5, 4), + (512, 512, 512, 64, 64, True, False, True): (1, 8, 4, 4), + (512, 512, 512, 128, 128, False, True, True): (4, 4, 2, 8), + (512, 512, 512, 128, 128, True, False, True): (4, 4, 2, 8), + (512, 512, 1024, 16, 16, False, True, True): (2, 8, 4, 4), + (512, 512, 1024, 16, 16, True, False, True): (2, 16, 4, 2), + (512, 512, 1024, 32, 32, False, True, True): (3, 16, 4, 2), + (512, 512, 1024, 32, 32, True, False, True): (3, 16, 3, 2), + (512, 512, 1024, 64, 64, False, True, True): (5, 8, 5, 4), + (512, 512, 1024, 64, 64, True, False, True): (4, 16, 3, 4), + (512, 512, 1024, 128, 128, False, True, True): (6, 8, 2, 8), + (512, 512, 1024, 128, 128, True, False, True): (4, 8, 2, 8), + (512, 512, 2048, 16, 16, False, True, True): (2, 16, 3, 4), + (512, 512, 2048, 16, 16, True, False, True): (1, 16, 4, 2), + (512, 512, 2048, 32, 32, False, True, True): (2, 32, 3, 2), + (512, 512, 2048, 32, 32, True, False, True): (2, 32, 3, 2), + (512, 512, 2048, 64, 64, False, True, True): (1, 32, 3, 4), + (512, 512, 2048, 64, 64, True, False, True): (1, 32, 3, 2), + (512, 512, 2048, 128, 128, False, True, True): (3, 16, 2, 8), + (512, 512, 2048, 128, 128, True, False, True): (1, 16, 2, 8), + (512, 512, 4096, 16, 16, False, True, True): (4, 32, 3, 2), + (512, 512, 4096, 16, 16, True, False, True): (1, 32, 3, 2), + (512, 512, 4096, 32, 32, False, True, True): (3, 32, 3, 2), + (512, 512, 4096, 32, 32, True, False, True): (3, 32, 3, 2), + (512, 512, 4096, 64, 64, False, True, True): (1, 32, 3, 4), + (512, 512, 4096, 64, 64, True, False, True): (1, 64, 1, 4), + (512, 512, 4096, 128, 128, False, True, True): (7, 32, 2, 8), + (512, 512, 4096, 128, 128, True, False, True): (1, 32, 2, 8), + (512, 512, 8192, 16, 16, False, True, True): (4, 64, 3, 2), + (512, 512, 8192, 16, 16, True, False, True): (1, 64, 3, 2), + (512, 512, 8192, 32, 32, False, True, True): (3, 64, 3, 2), + (512, 512, 8192, 32, 32, True, False, True): (1, 64, 3, 2), + (512, 512, 8192, 64, 64, False, True, True): (1, 64, 3, 4), + (512, 512, 8192, 64, 64, True, False, True): (1, 64, 3, 4), + (512, 512, 8192, 128, 128, False, True, True): (7, 64, 2, 8), + (512, 512, 8192, 128, 128, True, False, True): (1, 64, 1, 4), + (512, 512, 16384, 16, 16, False, True, True): (1, 128, 3, 2), + (512, 512, 16384, 16, 16, True, False, True): (1, 64, 3, 2), + (512, 512, 16384, 32, 32, False, True, True): (1, 128, 3, 2), + (512, 512, 16384, 32, 32, True, False, True): (1, 128, 3, 2), + (512, 512, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (512, 512, 16384, 64, 64, True, False, True): (4, 128, 3, 4), + (512, 512, 16384, 128, 128, False, True, True): (5, 128, 2, 8), + (512, 512, 16384, 128, 128, True, False, True): (2, 128, 1, 4), + (512, 512, 32768, 16, 16, False, True, True): (1, 128, 3, 4), + (512, 512, 32768, 16, 16, True, False, True): (1, 128, 3, 2), + (512, 512, 32768, 32, 32, False, True, True): (1, 256, 3, 2), + (512, 512, 32768, 32, 32, True, False, True): (1, 256, 3, 2), + (512, 512, 32768, 64, 64, False, True, True): (1, 256, 3, 4), + (512, 512, 32768, 64, 64, True, False, True): (1, 256, 3, 4), + (512, 512, 32768, 128, 128, False, True, True): (5, 256, 1, 4), + (512, 512, 32768, 128, 128, True, False, True): (1, 256, 1, 4), + (512, 512, 50432, 16, 16, False, True, True): (4, 197, 1, 4), + (512, 512, 50432, 16, 16, True, False, True): (4, 197, 3, 2), + (512, 512, 50432, 32, 32, False, True, True): (2, 197, 1, 4), + (512, 512, 50432, 32, 32, True, False, True): (4, 197, 3, 4), + (512, 512, 50432, 64, 64, False, True, True): (2, 394, 1, 4), + (512, 512, 50432, 64, 64, True, False, True): (4, 197, 2, 4), + (512, 512, 50432, 128, 128, False, True, True): (5, 394, 1, 4), + (512, 512, 50432, 128, 128, True, False, True): (6, 394, 2, 4), + (512, 512, 65536, 16, 16, False, True, True): (1, 256, 3, 2), + (512, 512, 65536, 16, 16, True, False, True): (1, 256, 3, 1), + (512, 512, 65536, 32, 32, False, True, True): (1, 512, 3, 2), + (512, 512, 65536, 32, 32, True, False, True): (1, 512, 3, 2), + (512, 512, 65536, 64, 64, False, True, True): (2, 256, 2, 4), + (512, 512, 65536, 64, 64, True, False, True): (1, 512, 3, 4), + (512, 512, 65536, 128, 128, False, True, True): (7, 512, 1, 4), + (512, 512, 65536, 128, 128, True, False, True): (5, 512, 1, 4), + (512, 512, 65792, 16, 16, False, True, True): (2, 257, 1, 4), + (512, 512, 65792, 16, 16, True, False, True): (1, 257, 3, 4), + (512, 512, 65792, 32, 32, False, True, True): (2, 257, 1, 4), + (512, 512, 65792, 32, 32, True, False, True): (1, 257, 3, 4), + (512, 512, 65792, 64, 64, False, True, True): (4, 514, 1, 4), + (512, 512, 65792, 64, 64, True, False, True): (4, 257, 2, 4), + (512, 512, 65792, 128, 128, False, True, True): (5, 514, 1, 4), + (512, 512, 65792, 128, 128, True, False, True): (4, 514, 2, 4), + (512, 512, 131072, 16, 16, False, True, True): (1, 512, 3, 1), + (512, 512, 131072, 16, 16, True, False, True): (1, 512, 3, 1), + (512, 512, 131072, 32, 32, False, True, True): (1, 1024, 3, 2), + (512, 512, 131072, 32, 32, True, False, True): (1, 1024, 3, 2), + (512, 512, 131072, 64, 64, False, True, True): (4, 512, 2, 4), + (512, 512, 131072, 64, 64, True, False, True): (2, 512, 2, 4), + (512, 512, 131072, 128, 128, False, True, True): (5, 1024, 1, 4), + (512, 512, 131072, 128, 128, True, False, True): (4, 1024, 1, 4), + (768, 768, 256, 16, 16, False, True, True): (1, 8, 4, 1), + (768, 768, 256, 16, 16, True, False, True): (3, 2, 5, 2), + (768, 768, 256, 32, 32, False, True, True): (1, 8, 4, 2), + (768, 768, 256, 32, 32, True, False, True): (2, 4, 6, 2), + (768, 768, 256, 64, 64, False, True, True): (3, 4, 3, 4), + (768, 768, 256, 64, 64, True, False, True): (2, 4, 4, 4), + (768, 768, 256, 128, 128, False, True, True): (1, 2, 3, 8), + (768, 768, 256, 128, 128, True, False, True): (2, 2, 3, 8), + (768, 768, 512, 16, 16, False, True, True): (1, 8, 4, 2), + (768, 768, 512, 16, 16, True, False, True): (2, 8, 5, 2), + (768, 768, 512, 32, 32, False, True, True): (1, 16, 1, 4), + (768, 768, 512, 32, 32, True, False, True): (3, 8, 5, 2), + (768, 768, 512, 64, 64, False, True, True): (4, 8, 3, 4), + (768, 768, 512, 64, 64, True, False, True): (2, 8, 4, 4), + (768, 768, 512, 128, 128, False, True, True): (1, 4, 3, 8), + (768, 768, 512, 128, 128, True, False, True): (3, 4, 3, 8), + (768, 768, 1024, 16, 16, False, True, True): (1, 16, 1, 4), + (768, 768, 1024, 16, 16, True, False, True): (1, 8, 5, 2), + (768, 768, 1024, 32, 32, False, True, True): (1, 16, 1, 8), + (768, 768, 1024, 32, 32, True, False, True): (1, 4, 4, 4), + (768, 768, 1024, 64, 64, False, True, True): (2, 16, 1, 8), + (768, 768, 1024, 64, 64, True, False, True): (1, 8, 3, 8), + (768, 768, 1024, 128, 128, False, True, True): (1, 8, 3, 8), + (768, 768, 1024, 128, 128, True, False, True): (3, 8, 3, 8), + (768, 768, 2048, 16, 16, False, True, True): (6, 16, 1, 2), + (768, 768, 2048, 16, 16, True, False, True): (2, 16, 4, 2), + (768, 768, 2048, 32, 32, False, True, True): (3, 32, 1, 4), + (768, 768, 2048, 32, 32, True, False, True): (6, 8, 3, 4), + (768, 768, 2048, 64, 64, False, True, True): (2, 32, 2, 2), + (768, 768, 2048, 64, 64, True, False, True): (1, 16, 4, 4), + (768, 768, 2048, 128, 128, False, True, True): (2, 16, 3, 8), + (768, 768, 2048, 128, 128, True, False, True): (4, 16, 3, 8), + (768, 768, 4096, 16, 16, False, True, True): (1, 32, 1, 4), + (768, 768, 4096, 16, 16, True, False, True): (2, 16, 3, 2), + (768, 768, 4096, 32, 32, False, True, True): (3, 32, 1, 8), + (768, 768, 4096, 32, 32, True, False, True): (1, 16, 4, 4), + (768, 768, 4096, 64, 64, False, True, True): (1, 64, 2, 4), + (768, 768, 4096, 64, 64, True, False, True): (1, 8, 3, 8), + (768, 768, 4096, 128, 128, False, True, True): (1, 32, 3, 8), + (768, 768, 4096, 128, 128, True, False, True): (2, 32, 3, 8), + (768, 768, 8192, 16, 16, False, True, True): (1, 64, 1, 2), + (768, 768, 8192, 16, 16, True, False, True): (2, 64, 3, 2), + (768, 768, 8192, 32, 32, False, True, True): (2, 64, 1, 8), + (768, 768, 8192, 32, 32, True, False, True): (2, 32, 3, 4), + (768, 768, 8192, 64, 64, False, True, True): (4, 64, 3, 4), + (768, 768, 8192, 64, 64, True, False, True): (1, 64, 3, 4), + (768, 768, 8192, 128, 128, False, True, True): (4, 64, 3, 8), + (768, 768, 8192, 128, 128, True, False, True): (2, 64, 3, 8), + (768, 768, 16384, 16, 16, False, True, True): (4, 128, 1, 2), + (768, 768, 16384, 16, 16, True, False, True): (1, 64, 3, 4), + (768, 768, 16384, 32, 32, False, True, True): (1, 128, 1, 8), + (768, 768, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (768, 768, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (768, 768, 16384, 64, 64, True, False, True): (1, 128, 3, 4), + (768, 768, 16384, 128, 128, False, True, True): (3, 128, 1, 4), + (768, 768, 16384, 128, 128, True, False, True): (1, 128, 2, 4), + (768, 768, 32768, 16, 16, False, True, True): (2, 256, 1, 2), + (768, 768, 32768, 16, 16, True, False, True): (1, 128, 4, 4), + (768, 768, 32768, 32, 32, False, True, True): (1, 128, 1, 2), + (768, 768, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (768, 768, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (768, 768, 32768, 64, 64, True, False, True): (1, 128, 3, 4), + (768, 768, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (768, 768, 32768, 128, 128, True, False, True): (3, 256, 2, 4), + (768, 768, 65536, 16, 16, False, True, True): (4, 512, 1, 2), + (768, 768, 65536, 16, 16, True, False, True): (1, 256, 4, 4), + (768, 768, 65536, 32, 32, False, True, True): (1, 256, 1, 2), + (768, 768, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (768, 768, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (768, 768, 65536, 64, 64, True, False, True): (1, 256, 3, 4), + (768, 768, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (768, 768, 65536, 128, 128, True, False, True): (2, 512, 2, 4), + (768, 768, 131072, 16, 16, False, True, True): (1, 512, 1, 1), + (768, 768, 131072, 16, 16, True, False, True): (1, 512, 4, 4), + (768, 768, 131072, 32, 32, False, True, True): (1, 512, 1, 2), + (768, 768, 131072, 32, 32, True, False, True): (1, 512, 3, 4), + (768, 768, 131072, 64, 64, False, True, True): (1, 1024, 1, 4), + (768, 768, 131072, 64, 64, True, False, True): (3, 512, 3, 4), + (768, 768, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (768, 768, 131072, 128, 128, True, False, True): (1, 1024, 2, 4), + (768, 3072, 256, 16, 16, False, True, True): (1, 8, 5, 2), + (768, 3072, 256, 16, 16, True, False, True): (3, 4, 7, 2), + (768, 3072, 256, 32, 32, False, True, True): (1, 8, 4, 2), + (768, 3072, 256, 32, 32, True, False, True): (1, 4, 5, 4), + (768, 3072, 256, 64, 64, False, True, True): (1, 4, 3, 4), + (768, 3072, 256, 64, 64, True, False, True): (1, 4, 5, 4), + (768, 3072, 256, 128, 128, False, True, True): (2, 2, 3, 8), + (768, 3072, 256, 128, 128, True, False, True): (2, 2, 3, 8), + (768, 3072, 512, 16, 16, False, True, True): (1, 8, 5, 2), + (768, 3072, 512, 16, 16, True, False, True): (1, 8, 5, 2), + (768, 3072, 512, 32, 32, False, True, True): (3, 8, 3, 4), + (768, 3072, 512, 32, 32, True, False, True): (1, 8, 7, 4), + (768, 3072, 512, 64, 64, False, True, True): (3, 8, 3, 4), + (768, 3072, 512, 64, 64, True, False, True): (3, 8, 5, 4), + (768, 3072, 512, 128, 128, False, True, True): (1, 4, 3, 8), + (768, 3072, 512, 128, 128, True, False, True): (1, 4, 3, 8), + (768, 3072, 1024, 16, 16, False, True, True): (4, 16, 1, 4), + (768, 3072, 1024, 16, 16, True, False, True): (2, 8, 5, 2), + (768, 3072, 1024, 32, 32, False, True, True): (1, 16, 6, 2), + (768, 3072, 1024, 32, 32, True, False, True): (1, 8, 4, 4), + (768, 3072, 1024, 64, 64, False, True, True): (2, 16, 4, 4), + (768, 3072, 1024, 64, 64, True, False, True): (2, 16, 4, 4), + (768, 3072, 1024, 128, 128, False, True, True): (1, 8, 3, 8), + (768, 3072, 1024, 128, 128, True, False, True): (3, 8, 3, 8), + (768, 3072, 2048, 16, 16, False, True, True): (1, 16, 1, 2), + (768, 3072, 2048, 16, 16, True, False, True): (1, 16, 5, 2), + (768, 3072, 2048, 32, 32, False, True, True): (4, 16, 1, 8), + (768, 3072, 2048, 32, 32, True, False, True): (2, 8, 3, 4), + (768, 3072, 2048, 64, 64, False, True, True): (2, 16, 3, 4), + (768, 3072, 2048, 64, 64, True, False, True): (2, 16, 3, 4), + (768, 3072, 2048, 128, 128, False, True, True): (3, 16, 3, 8), + (768, 3072, 2048, 128, 128, True, False, True): (1, 16, 3, 8), + (768, 3072, 4096, 16, 16, False, True, True): (1, 32, 1, 4), + (768, 3072, 4096, 16, 16, True, False, True): (1, 16, 3, 1), + (768, 3072, 4096, 32, 32, False, True, True): (3, 32, 1, 8), + (768, 3072, 4096, 32, 32, True, False, True): (2, 16, 3, 8), + (768, 3072, 4096, 64, 64, False, True, True): (2, 32, 3, 4), + (768, 3072, 4096, 64, 64, True, False, True): (2, 16, 3, 4), + (768, 3072, 4096, 128, 128, False, True, True): (5, 32, 1, 4), + (768, 3072, 4096, 128, 128, True, False, True): (4, 32, 3, 8), + (768, 3072, 8192, 16, 16, False, True, True): (1, 32, 1, 4), + (768, 3072, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (768, 3072, 8192, 32, 32, False, True, True): (1, 64, 1, 8), + (768, 3072, 8192, 32, 32, True, False, True): (2, 32, 3, 8), + (768, 3072, 8192, 64, 64, False, True, True): (2, 64, 3, 4), + (768, 3072, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (768, 3072, 8192, 128, 128, False, True, True): (1, 64, 3, 8), + (768, 3072, 8192, 128, 128, True, False, True): (2, 64, 3, 8), + (768, 3072, 16384, 16, 16, False, True, True): (1, 64, 1, 4), + (768, 3072, 16384, 16, 16, True, False, True): (1, 64, 4, 1), + (768, 3072, 16384, 32, 32, False, True, True): (1, 128, 1, 8), + (768, 3072, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (768, 3072, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (768, 3072, 16384, 64, 64, True, False, True): (1, 64, 3, 4), + (768, 3072, 16384, 128, 128, False, True, True): (2, 128, 3, 8), + (768, 3072, 16384, 128, 128, True, False, True): (1, 128, 3, 8), + (768, 3072, 32768, 16, 16, False, True, True): (1, 128, 1, 4), + (768, 3072, 32768, 16, 16, True, False, True): (1, 128, 4, 1), + (768, 3072, 32768, 32, 32, False, True, True): (1, 256, 1, 8), + (768, 3072, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (768, 3072, 32768, 64, 64, False, True, True): (1, 256, 3, 4), + (768, 3072, 32768, 64, 64, True, False, True): (1, 128, 3, 4), + (768, 3072, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (768, 3072, 32768, 128, 128, True, False, True): (5, 256, 3, 8), + (768, 3072, 50432, 16, 16, False, True, True): (1, 197, 1, 4), + (768, 3072, 50432, 16, 16, True, False, True): (4, 197, 4, 1), + (768, 3072, 50432, 32, 32, False, True, True): (2, 197, 1, 4), + (768, 3072, 50432, 32, 32, True, False, True): (4, 197, 3, 4), + (768, 3072, 50432, 64, 64, False, True, True): (1, 394, 3, 4), + (768, 3072, 50432, 64, 64, True, False, True): (1, 197, 3, 4), + (768, 3072, 50432, 128, 128, False, True, True): (3, 394, 1, 4), + (768, 3072, 50432, 128, 128, True, False, True): (3, 394, 2, 4), + (768, 3072, 65536, 16, 16, False, True, True): (1, 256, 1, 4), + (768, 3072, 65536, 16, 16, True, False, True): (5, 256, 4, 1), + (768, 3072, 65536, 32, 32, False, True, True): (2, 256, 1, 4), + (768, 3072, 65536, 32, 32, True, False, True): (3, 256, 3, 4), + (768, 3072, 65536, 64, 64, False, True, True): (1, 512, 3, 4), + (768, 3072, 65536, 64, 64, True, False, True): (1, 256, 3, 4), + (768, 3072, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (768, 3072, 65536, 128, 128, True, False, True): (2, 512, 3, 8), + (768, 3072, 131072, 16, 16, False, True, True): (1, 512, 1, 4), + (768, 3072, 131072, 16, 16, True, False, True): (5, 512, 4, 1), + (768, 3072, 131072, 32, 32, False, True, True): (2, 512, 1, 4), + (768, 3072, 131072, 32, 32, True, False, True): (2, 512, 3, 4), + (768, 3072, 131072, 64, 64, False, True, True): (1, 1024, 3, 4), + (768, 3072, 131072, 64, 64, True, False, True): (2, 512, 3, 4), + (768, 3072, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (768, 3072, 131072, 128, 128, True, False, True): (2, 1024, 3, 8), + (1024, 1024, 256, 16, 16, False, True, True): (3, 4, 5, 4), + (1024, 1024, 256, 16, 16, True, False, True): (3, 4, 5, 4), + (1024, 1024, 256, 32, 32, False, True, True): (2, 4, 6, 2), + (1024, 1024, 256, 32, 32, True, False, True): (2, 4, 6, 2), + (1024, 1024, 256, 64, 64, False, True, True): (1, 4, 4, 4), + (1024, 1024, 256, 64, 64, True, False, True): (2, 4, 6, 4), + (1024, 1024, 256, 128, 128, False, True, True): (1, 2, 2, 8), + (1024, 1024, 256, 128, 128, True, False, True): (1, 2, 2, 8), + (1024, 1024, 512, 16, 16, False, True, True): (3, 4, 5, 4), + (1024, 1024, 512, 16, 16, True, False, True): (3, 8, 4, 2), + (1024, 1024, 512, 32, 32, False, True, True): (1, 8, 4, 2), + (1024, 1024, 512, 32, 32, True, False, True): (1, 8, 4, 2), + (1024, 1024, 512, 64, 64, False, True, True): (2, 8, 3, 4), + (1024, 1024, 512, 64, 64, True, False, True): (1, 4, 4, 4), + (1024, 1024, 512, 128, 128, False, True, True): (7, 4, 2, 8), + (1024, 1024, 512, 128, 128, True, False, True): (1, 4, 2, 8), + (1024, 1024, 1024, 16, 16, False, True, True): (4, 8, 4, 2), + (1024, 1024, 1024, 16, 16, True, False, True): (3, 8, 5, 2), + (1024, 1024, 1024, 32, 32, False, True, True): (1, 8, 4, 4), + (1024, 1024, 1024, 32, 32, True, False, True): (1, 8, 4, 2), + (1024, 1024, 1024, 64, 64, False, True, True): (1, 16, 3, 4), + (1024, 1024, 1024, 64, 64, True, False, True): (3, 16, 3, 4), + (1024, 1024, 1024, 128, 128, False, True, True): (6, 8, 2, 8), + (1024, 1024, 1024, 128, 128, True, False, True): (4, 8, 2, 8), + (1024, 1024, 2048, 16, 16, False, True, True): (3, 8, 3, 4), + (1024, 1024, 2048, 16, 16, True, False, True): (3, 8, 3, 4), + (1024, 1024, 2048, 32, 32, False, True, True): (1, 16, 3, 4), + (1024, 1024, 2048, 32, 32, True, False, True): (1, 16, 3, 2), + (1024, 1024, 2048, 64, 64, False, True, True): (5, 16, 3, 4), + (1024, 1024, 2048, 64, 64, True, False, True): (5, 16, 3, 4), + (1024, 1024, 2048, 128, 128, False, True, True): (3, 16, 2, 8), + (1024, 1024, 2048, 128, 128, True, False, True): (4, 16, 2, 16), + (1024, 1024, 4096, 16, 16, False, True, True): (4, 32, 3, 2), + (1024, 1024, 4096, 16, 16, True, False, True): (8, 32, 3, 2), + (1024, 1024, 4096, 32, 32, False, True, True): (9, 32, 3, 2), + (1024, 1024, 4096, 32, 32, True, False, True): (1, 32, 3, 2), + (1024, 1024, 4096, 64, 64, False, True, True): (6, 32, 3, 4), + (1024, 1024, 4096, 64, 64, True, False, True): (1, 32, 3, 4), + (1024, 1024, 4096, 128, 128, False, True, True): (4, 32, 2, 8), + (1024, 1024, 4096, 128, 128, True, False, True): (4, 32, 1, 4), + (1024, 1024, 8192, 16, 16, False, True, True): (4, 64, 3, 2), + (1024, 1024, 8192, 16, 16, True, False, True): (4, 64, 3, 2), + (1024, 1024, 8192, 32, 32, False, True, True): (8, 64, 3, 2), + (1024, 1024, 8192, 32, 32, True, False, True): (6, 64, 3, 2), + (1024, 1024, 8192, 64, 64, False, True, True): (2, 64, 3, 4), + (1024, 1024, 8192, 64, 64, True, False, True): (2, 64, 3, 4), + (1024, 1024, 8192, 128, 128, False, True, True): (3, 64, 1, 4), + (1024, 1024, 8192, 128, 128, True, False, True): (2, 64, 1, 4), + (1024, 1024, 16384, 16, 16, False, True, True): (1, 64, 3, 4), + (1024, 1024, 16384, 16, 16, True, False, True): (1, 64, 3, 2), + (1024, 1024, 16384, 32, 32, False, True, True): (1, 128, 3, 4), + (1024, 1024, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (1024, 1024, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (1024, 1024, 16384, 64, 64, True, False, True): (1, 128, 3, 4), + (1024, 1024, 16384, 128, 128, False, True, True): (11, 128, 1, 4), + (1024, 1024, 16384, 128, 128, True, False, True): (4, 128, 1, 4), + (1024, 1024, 32768, 16, 16, False, True, True): (1, 128, 3, 4), + (1024, 1024, 32768, 16, 16, True, False, True): (1, 128, 3, 1), + (1024, 1024, 32768, 32, 32, False, True, True): (1, 256, 3, 2), + (1024, 1024, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (1024, 1024, 32768, 64, 64, False, True, True): (2, 128, 2, 4), + (1024, 1024, 32768, 64, 64, True, False, True): (1, 256, 3, 4), + (1024, 1024, 32768, 128, 128, False, True, True): (7, 256, 1, 4), + (1024, 1024, 32768, 128, 128, True, False, True): (4, 256, 1, 4), + (1024, 1024, 50432, 16, 16, False, True, True): (1, 197, 1, 4), + (1024, 1024, 50432, 16, 16, True, False, True): (4, 197, 3, 4), + (1024, 1024, 50432, 32, 32, False, True, True): (2, 197, 1, 4), + (1024, 1024, 50432, 32, 32, True, False, True): (1, 197, 3, 4), + (1024, 1024, 50432, 64, 64, False, True, True): (2, 394, 1, 4), + (1024, 1024, 50432, 64, 64, True, False, True): (1, 197, 2, 4), + (1024, 1024, 50432, 128, 128, False, True, True): (3, 394, 1, 4), + (1024, 1024, 50432, 128, 128, True, False, True): (2, 394, 2, 4), + (1024, 1024, 65536, 16, 16, False, True, True): (1, 256, 3, 4), + (1024, 1024, 65536, 16, 16, True, False, True): (1, 256, 3, 1), + (1024, 1024, 65536, 32, 32, False, True, True): (1, 512, 3, 2), + (1024, 1024, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (1024, 1024, 65536, 64, 64, False, True, True): (2, 256, 2, 4), + (1024, 1024, 65536, 64, 64, True, False, True): (1, 512, 3, 4), + (1024, 1024, 65536, 128, 128, False, True, True): (10, 512, 1, 4), + (1024, 1024, 65536, 128, 128, True, False, True): (4, 512, 1, 4), + (1024, 1024, 65792, 16, 16, False, True, True): (1, 257, 1, 4), + (1024, 1024, 65792, 16, 16, True, False, True): (10, 257, 4, 1), + (1024, 1024, 65792, 32, 32, False, True, True): (2, 257, 1, 4), + (1024, 1024, 65792, 32, 32, True, False, True): (1, 257, 3, 4), + (1024, 1024, 65792, 64, 64, False, True, True): (2, 514, 1, 4), + (1024, 1024, 65792, 64, 64, True, False, True): (2, 257, 2, 4), + (1024, 1024, 65792, 128, 128, False, True, True): (6, 514, 1, 4), + (1024, 1024, 65792, 128, 128, True, False, True): (2, 514, 2, 4), + (1024, 1024, 131072, 16, 16, False, True, True): (11, 512, 3, 2), + (1024, 1024, 131072, 16, 16, True, False, True): (11, 512, 3, 2), + (1024, 1024, 131072, 32, 32, False, True, True): (7, 1024, 3, 2), + (1024, 1024, 131072, 32, 32, True, False, True): (6, 512, 3, 4), + (1024, 1024, 131072, 64, 64, False, True, True): (1, 512, 2, 4), + (1024, 1024, 131072, 64, 64, True, False, True): (4, 1024, 3, 4), + (1024, 1024, 131072, 128, 128, False, True, True): (12, 1024, 1, 4), + (1024, 1024, 131072, 128, 128, True, False, True): (4, 1024, 1, 4), + (1280, 5120, 65792, 16, 16, False, True, True): (1, 257, 1, 4), + (1280, 5120, 65792, 16, 16, True, False, True): (5, 257, 4, 1), + (1280, 5120, 65792, 32, 32, False, True, True): (2, 257, 1, 4), + (1280, 5120, 65792, 32, 32, True, False, True): (2, 257, 3, 4), + (1280, 5120, 65792, 64, 64, False, True, True): (1, 514, 3, 4), + (1280, 5120, 65792, 64, 64, True, False, True): (2, 257, 3, 4), + (1280, 5120, 65792, 128, 128, False, True, True): (1, 514, 3, 8), + (1280, 5120, 65792, 128, 128, True, False, True): (1, 514, 3, 8), + (1536, 1536, 256, 16, 16, False, True, True): (5, 4, 4, 2), + (1536, 1536, 256, 16, 16, True, False, True): (3, 4, 5, 2), + (1536, 1536, 256, 32, 32, False, True, True): (2, 4, 4, 4), + (1536, 1536, 256, 32, 32, True, False, True): (1, 4, 6, 2), + (1536, 1536, 256, 64, 64, False, True, True): (5, 4, 4, 4), + (1536, 1536, 256, 64, 64, True, False, True): (2, 4, 4, 4), + (1536, 1536, 256, 128, 128, False, True, True): (1, 2, 3, 8), + (1536, 1536, 256, 128, 128, True, False, True): (2, 2, 3, 8), + (1536, 1536, 512, 16, 16, False, True, True): (1, 8, 1, 4), + (1536, 1536, 512, 16, 16, True, False, True): (3, 4, 4, 2), + (1536, 1536, 512, 32, 32, False, True, True): (1, 8, 1, 8), + (1536, 1536, 512, 32, 32, True, False, True): (1, 4, 4, 4), + (1536, 1536, 512, 64, 64, False, True, True): (3, 8, 3, 4), + (1536, 1536, 512, 64, 64, True, False, True): (5, 8, 3, 4), + (1536, 1536, 512, 128, 128, False, True, True): (3, 4, 3, 8), + (1536, 1536, 512, 128, 128, True, False, True): (1, 4, 3, 8), + (1536, 1536, 1024, 16, 16, False, True, True): (6, 8, 1, 2), + (1536, 1536, 1024, 16, 16, True, False, True): (2, 8, 5, 2), + (1536, 1536, 1024, 32, 32, False, True, True): (6, 8, 1, 8), + (1536, 1536, 1024, 32, 32, True, False, True): (2, 4, 3, 4), + (1536, 1536, 1024, 64, 64, False, True, True): (1, 16, 3, 4), + (1536, 1536, 1024, 64, 64, True, False, True): (3, 8, 3, 4), + (1536, 1536, 1024, 128, 128, False, True, True): (3, 8, 3, 8), + (1536, 1536, 1024, 128, 128, True, False, True): (3, 8, 3, 8), + (1536, 1536, 2048, 16, 16, False, True, True): (1, 16, 1, 4), + (1536, 1536, 2048, 16, 16, True, False, True): (1, 8, 3, 1), + (1536, 1536, 2048, 32, 32, False, True, True): (1, 16, 1, 8), + (1536, 1536, 2048, 32, 32, True, False, True): (4, 8, 3, 2), + (1536, 1536, 2048, 64, 64, False, True, True): (1, 16, 3, 4), + (1536, 1536, 2048, 64, 64, True, False, True): (3, 8, 3, 4), + (1536, 1536, 2048, 128, 128, False, True, True): (6, 16, 1, 4), + (1536, 1536, 2048, 128, 128, True, False, True): (4, 16, 3, 8), + (1536, 1536, 4096, 16, 16, False, True, True): (1, 32, 1, 2), + (1536, 1536, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (1536, 1536, 4096, 32, 32, False, True, True): (1, 32, 1, 8), + (1536, 1536, 4096, 32, 32, True, False, True): (3, 16, 3, 4), + (1536, 1536, 4096, 64, 64, False, True, True): (1, 32, 3, 4), + (1536, 1536, 4096, 64, 64, True, False, True): (1, 16, 3, 4), + (1536, 1536, 4096, 128, 128, False, True, True): (4, 32, 3, 8), + (1536, 1536, 4096, 128, 128, True, False, True): (2, 32, 3, 8), + (1536, 1536, 8192, 16, 16, False, True, True): (2, 64, 1, 2), + (1536, 1536, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (1536, 1536, 8192, 32, 32, False, True, True): (1, 64, 1, 8), + (1536, 1536, 8192, 32, 32, True, False, True): (12, 32, 3, 4), + (1536, 1536, 8192, 64, 64, False, True, True): (2, 64, 3, 4), + (1536, 1536, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (1536, 1536, 8192, 128, 128, False, True, True): (3, 64, 1, 4), + (1536, 1536, 8192, 128, 128, True, False, True): (4, 64, 3, 8), + (1536, 1536, 16384, 16, 16, False, True, True): (1, 128, 1, 2), + (1536, 1536, 16384, 16, 16, True, False, True): (1, 64, 4, 4), + (1536, 1536, 16384, 32, 32, False, True, True): (1, 64, 1, 2), + (1536, 1536, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (1536, 1536, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (1536, 1536, 16384, 64, 64, True, False, True): (1, 64, 3, 4), + (1536, 1536, 16384, 128, 128, False, True, True): (3, 128, 1, 4), + (1536, 1536, 16384, 128, 128, True, False, True): (1, 128, 2, 4), + (1536, 1536, 32768, 16, 16, False, True, True): (1, 256, 1, 2), + (1536, 1536, 32768, 16, 16, True, False, True): (1, 128, 3, 2), + (1536, 1536, 32768, 32, 32, False, True, True): (1, 128, 1, 2), + (1536, 1536, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (1536, 1536, 32768, 64, 64, False, True, True): (3, 256, 3, 4), + (1536, 1536, 32768, 64, 64, True, False, True): (1, 128, 3, 4), + (1536, 1536, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (1536, 1536, 32768, 128, 128, True, False, True): (1, 256, 2, 4), + (1536, 1536, 65536, 16, 16, False, True, True): (4, 512, 1, 2), + (1536, 1536, 65536, 16, 16, True, False, True): (1, 256, 4, 4), + (1536, 1536, 65536, 32, 32, False, True, True): (1, 256, 1, 2), + (1536, 1536, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (1536, 1536, 65536, 64, 64, False, True, True): (2, 512, 3, 4), + (1536, 1536, 65536, 64, 64, True, False, True): (1, 256, 3, 4), + (1536, 1536, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (1536, 1536, 65536, 128, 128, True, False, True): (2, 512, 2, 4), + (1536, 1536, 131072, 16, 16, False, True, True): (2, 1024, 1, 2), + (1536, 1536, 131072, 16, 16, True, False, True): (9, 512, 4, 4), + (1536, 1536, 131072, 32, 32, False, True, True): (1, 512, 1, 2), + (1536, 1536, 131072, 32, 32, True, False, True): (9, 512, 3, 4), + (1536, 1536, 131072, 64, 64, False, True, True): (1, 1024, 3, 4), + (1536, 1536, 131072, 64, 64, True, False, True): (1, 512, 3, 4), + (1536, 1536, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (1536, 1536, 131072, 128, 128, True, False, True): (1, 1024, 2, 4), + (2048, 2048, 256, 16, 16, False, True, True): (4, 4, 6, 2), + (2048, 2048, 256, 16, 16, True, False, True): (2, 8, 4, 1), + (2048, 2048, 256, 32, 32, False, True, True): (3, 4, 4, 2), + (2048, 2048, 256, 32, 32, True, False, True): (1, 4, 5, 2), + (2048, 2048, 256, 64, 64, False, True, True): (2, 4, 4, 4), + (2048, 2048, 256, 64, 64, True, False, True): (2, 4, 4, 4), + (2048, 2048, 256, 128, 128, False, True, True): (3, 2, 2, 8), + (2048, 2048, 256, 128, 128, True, False, True): (5, 2, 2, 8), + (2048, 2048, 512, 16, 16, False, True, True): (5, 4, 4, 4), + (2048, 2048, 512, 16, 16, True, False, True): (2, 4, 4, 2), + (2048, 2048, 512, 32, 32, False, True, True): (1, 4, 3, 4), + (2048, 2048, 512, 32, 32, True, False, True): (3, 4, 4, 2), + (2048, 2048, 512, 64, 64, False, True, True): (1, 8, 3, 4), + (2048, 2048, 512, 64, 64, True, False, True): (1, 8, 3, 2), + (2048, 2048, 512, 128, 128, False, True, True): (3, 4, 2, 8), + (2048, 2048, 512, 128, 128, True, False, True): (2, 4, 2, 8), + (2048, 2048, 1024, 16, 16, False, True, True): (3, 4, 3, 4), + (2048, 2048, 1024, 16, 16, True, False, True): (2, 8, 3, 2), + (2048, 2048, 1024, 32, 32, False, True, True): (3, 8, 3, 4), + (2048, 2048, 1024, 32, 32, True, False, True): (1, 8, 3, 2), + (2048, 2048, 1024, 64, 64, False, True, True): (1, 8, 3, 4), + (2048, 2048, 1024, 64, 64, True, False, True): (1, 8, 3, 4), + (2048, 2048, 1024, 128, 128, False, True, True): (4, 8, 2, 8), + (2048, 2048, 1024, 128, 128, True, False, True): (4, 8, 1, 4), + (2048, 2048, 2048, 16, 16, False, True, True): (4, 16, 3, 2), + (2048, 2048, 2048, 16, 16, True, False, True): (2, 16, 3, 2), + (2048, 2048, 2048, 32, 32, False, True, True): (1, 16, 3, 4), + (2048, 2048, 2048, 32, 32, True, False, True): (1, 16, 3, 2), + (2048, 2048, 2048, 64, 64, False, True, True): (1, 16, 3, 4), + (2048, 2048, 2048, 64, 64, True, False, True): (1, 16, 3, 4), + (2048, 2048, 2048, 128, 128, False, True, True): (6, 16, 2, 8), + (2048, 2048, 2048, 128, 128, True, False, True): (5, 16, 1, 4), + (2048, 2048, 4096, 16, 16, False, True, True): (4, 32, 4, 2), + (2048, 2048, 4096, 16, 16, True, False, True): (4, 32, 3, 2), + (2048, 2048, 4096, 32, 32, False, True, True): (4, 16, 3, 8), + (2048, 2048, 4096, 32, 32, True, False, True): (4, 16, 3, 4), + (2048, 2048, 4096, 64, 64, False, True, True): (4, 32, 3, 4), + (2048, 2048, 4096, 64, 64, True, False, True): (4, 32, 3, 4), + (2048, 2048, 4096, 128, 128, False, True, True): (4, 32, 2, 8), + (2048, 2048, 4096, 128, 128, True, False, True): (2, 32, 1, 4), + (2048, 2048, 8192, 16, 16, False, True, True): (4, 64, 4, 2), + (2048, 2048, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (2048, 2048, 8192, 32, 32, False, True, True): (4, 32, 3, 8), + (2048, 2048, 8192, 32, 32, True, False, True): (4, 32, 4, 8), + (2048, 2048, 8192, 64, 64, False, True, True): (2, 64, 3, 4), + (2048, 2048, 8192, 64, 64, True, False, True): (4, 64, 3, 4), + (2048, 2048, 8192, 128, 128, False, True, True): (3, 64, 1, 4), + (2048, 2048, 8192, 128, 128, True, False, True): (2, 64, 1, 4), + (2048, 2048, 16384, 16, 16, False, True, True): (4, 64, 3, 4), + (2048, 2048, 16384, 16, 16, True, False, True): (1, 64, 3, 4), + (2048, 2048, 16384, 32, 32, False, True, True): (4, 64, 3, 4), + (2048, 2048, 16384, 32, 32, True, False, True): (4, 64, 3, 4), + (2048, 2048, 16384, 64, 64, False, True, True): (4, 128, 3, 4), + (2048, 2048, 16384, 64, 64, True, False, True): (4, 128, 3, 4), + (2048, 2048, 16384, 128, 128, False, True, True): (3, 128, 1, 4), + (2048, 2048, 16384, 128, 128, True, False, True): (2, 128, 1, 4), + (2048, 2048, 32768, 16, 16, False, True, True): (8, 128, 3, 2), + (2048, 2048, 32768, 16, 16, True, False, True): (8, 128, 3, 4), + (2048, 2048, 32768, 32, 32, False, True, True): (8, 128, 3, 4), + (2048, 2048, 32768, 32, 32, True, False, True): (8, 128, 3, 4), + (2048, 2048, 32768, 64, 64, False, True, True): (8, 256, 3, 4), + (2048, 2048, 32768, 64, 64, True, False, True): (8, 256, 3, 4), + (2048, 2048, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (2048, 2048, 32768, 128, 128, True, False, True): (1, 256, 1, 4), + (2048, 2048, 50432, 16, 16, False, True, True): (1, 197, 1, 4), + (2048, 2048, 50432, 16, 16, True, False, True): (4, 197, 4, 1), + (2048, 2048, 50432, 32, 32, False, True, True): (2, 197, 1, 4), + (2048, 2048, 50432, 32, 32, True, False, True): (4, 197, 3, 4), + (2048, 2048, 50432, 64, 64, False, True, True): (2, 394, 3, 4), + (2048, 2048, 50432, 64, 64, True, False, True): (4, 197, 2, 4), + (2048, 2048, 50432, 128, 128, False, True, True): (3, 394, 1, 4), + (2048, 2048, 50432, 128, 128, True, False, True): (4, 394, 2, 4), + (2048, 2048, 65536, 16, 16, False, True, True): (9, 256, 3, 2), + (2048, 2048, 65536, 16, 16, True, False, True): (9, 256, 4, 4), + (2048, 2048, 65536, 32, 32, False, True, True): (7, 256, 3, 4), + (2048, 2048, 65536, 32, 32, True, False, True): (7, 256, 3, 4), + (2048, 2048, 65536, 64, 64, False, True, True): (2, 256, 2, 4), + (2048, 2048, 65536, 64, 64, True, False, True): (9, 512, 3, 4), + (2048, 2048, 65536, 128, 128, False, True, True): (5, 512, 1, 4), + (2048, 2048, 65536, 128, 128, True, False, True): (1, 512, 1, 4), + (2048, 2048, 65792, 16, 16, False, True, True): (1, 257, 1, 4), + (2048, 2048, 65792, 16, 16, True, False, True): (7, 257, 4, 1), + (2048, 2048, 65792, 32, 32, False, True, True): (2, 257, 1, 4), + (2048, 2048, 65792, 32, 32, True, False, True): (7, 257, 3, 4), + (2048, 2048, 65792, 64, 64, False, True, True): (1, 514, 3, 4), + (2048, 2048, 65792, 64, 64, True, False, True): (1, 257, 2, 4), + (2048, 2048, 65792, 128, 128, False, True, True): (3, 514, 1, 4), + (2048, 2048, 65792, 128, 128, True, False, True): (1, 514, 2, 4), + (2048, 2048, 131072, 16, 16, False, True, True): (9, 512, 3, 2), + (2048, 2048, 131072, 16, 16, True, False, True): (9, 512, 4, 4), + (2048, 2048, 131072, 32, 32, False, True, True): (7, 512, 3, 4), + (2048, 2048, 131072, 32, 32, True, False, True): (3, 512, 3, 4), + (2048, 2048, 131072, 64, 64, False, True, True): (1, 512, 2, 4), + (2048, 2048, 131072, 64, 64, True, False, True): (2, 1024, 3, 4), + (2048, 2048, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (2048, 2048, 131072, 128, 128, True, False, True): (1, 1024, 1, 4), + (3072, 768, 256, 16, 16, False, True, True): (6, 4, 1, 4), + (3072, 768, 256, 16, 16, True, False, True): (2, 1, 5, 2), + (3072, 768, 256, 32, 32, False, True, True): (1, 4, 1, 8), + (3072, 768, 256, 32, 32, True, False, True): (4, 2, 4, 4), + (3072, 768, 256, 64, 64, False, True, True): (1, 2, 3, 4), + (3072, 768, 256, 64, 64, True, False, True): (3, 4, 3, 4), + (3072, 768, 256, 128, 128, False, True, True): (1, 2, 3, 8), + (3072, 768, 256, 128, 128, True, False, True): (3, 2, 3, 8), + (3072, 768, 512, 16, 16, False, True, True): (1, 4, 1, 4), + (3072, 768, 512, 16, 16, True, False, True): (3, 4, 4, 1), + (3072, 768, 512, 32, 32, False, True, True): (5, 8, 1, 4), + (3072, 768, 512, 32, 32, True, False, True): (3, 4, 4, 2), + (3072, 768, 512, 64, 64, False, True, True): (1, 8, 1, 4), + (3072, 768, 512, 64, 64, True, False, True): (1, 4, 3, 4), + (3072, 768, 512, 128, 128, False, True, True): (3, 4, 3, 8), + (3072, 768, 512, 128, 128, True, False, True): (1, 4, 3, 8), + (3072, 768, 1024, 16, 16, False, True, True): (1, 8, 1, 4), + (3072, 768, 1024, 16, 16, True, False, True): (3, 4, 3, 1), + (3072, 768, 1024, 32, 32, False, True, True): (1, 16, 1, 4), + (3072, 768, 1024, 32, 32, True, False, True): (1, 4, 3, 8), + (3072, 768, 1024, 64, 64, False, True, True): (8, 16, 3, 2), + (3072, 768, 1024, 64, 64, True, False, True): (1, 4, 3, 4), + (3072, 768, 1024, 128, 128, False, True, True): (2, 8, 3, 8), + (3072, 768, 1024, 128, 128, True, False, True): (3, 8, 2, 4), + (3072, 768, 2048, 16, 16, False, True, True): (1, 8, 1, 4), + (3072, 768, 2048, 16, 16, True, False, True): (6, 8, 4, 4), + (3072, 768, 2048, 32, 32, False, True, True): (1, 16, 1, 8), + (3072, 768, 2048, 32, 32, True, False, True): (6, 8, 3, 4), + (3072, 768, 2048, 64, 64, False, True, True): (8, 16, 3, 4), + (3072, 768, 2048, 64, 64, True, False, True): (3, 16, 3, 4), + (3072, 768, 2048, 128, 128, False, True, True): (1, 16, 3, 8), + (3072, 768, 2048, 128, 128, True, False, True): (2, 16, 2, 4), + (3072, 768, 4096, 16, 16, False, True, True): (1, 16, 1, 4), + (3072, 768, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (3072, 768, 4096, 32, 32, False, True, True): (1, 32, 1, 8), + (3072, 768, 4096, 32, 32, True, False, True): (4, 16, 3, 4), + (3072, 768, 4096, 64, 64, False, True, True): (2, 32, 1, 4), + (3072, 768, 4096, 64, 64, True, False, True): (2, 16, 2, 4), + (3072, 768, 4096, 128, 128, False, True, True): (2, 32, 1, 16), + (3072, 768, 4096, 128, 128, True, False, True): (3, 32, 2, 4), + (3072, 768, 8192, 16, 16, False, True, True): (2, 32, 1, 4), + (3072, 768, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (3072, 768, 8192, 32, 32, False, True, True): (2, 32, 1, 4), + (3072, 768, 8192, 32, 32, True, False, True): (6, 32, 3, 4), + (3072, 768, 8192, 64, 64, False, True, True): (2, 64, 1, 4), + (3072, 768, 8192, 64, 64, True, False, True): (2, 32, 2, 4), + (3072, 768, 8192, 128, 128, False, True, True): (3, 64, 1, 4), + (3072, 768, 8192, 128, 128, True, False, True): (2, 64, 2, 4), + (3072, 768, 16384, 16, 16, False, True, True): (1, 64, 1, 4), + (3072, 768, 16384, 16, 16, True, False, True): (1, 64, 1, 1), + (3072, 768, 16384, 32, 32, False, True, True): (2, 64, 1, 4), + (3072, 768, 16384, 32, 32, True, False, True): (4, 64, 3, 4), + (3072, 768, 16384, 64, 64, False, True, True): (2, 128, 1, 4), + (3072, 768, 16384, 64, 64, True, False, True): (4, 64, 2, 4), + (3072, 768, 16384, 128, 128, False, True, True): (3, 128, 1, 4), + (3072, 768, 16384, 128, 128, True, False, True): (1, 128, 2, 4), + (3072, 768, 32768, 16, 16, False, True, True): (1, 128, 1, 4), + (3072, 768, 32768, 16, 16, True, False, True): (8, 256, 3, 2), + (3072, 768, 32768, 32, 32, False, True, True): (2, 128, 1, 4), + (3072, 768, 32768, 32, 32, True, False, True): (8, 128, 3, 4), + (3072, 768, 32768, 64, 64, False, True, True): (1, 256, 1, 4), + (3072, 768, 32768, 64, 64, True, False, True): (8, 128, 2, 4), + (3072, 768, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (3072, 768, 32768, 128, 128, True, False, True): (3, 256, 2, 4), + (3072, 768, 50432, 16, 16, False, True, True): (1, 197, 1, 4), + (3072, 768, 50432, 16, 16, True, False, True): (7, 197, 4, 1), + (3072, 768, 50432, 32, 32, False, True, True): (2, 197, 1, 4), + (3072, 768, 50432, 32, 32, True, False, True): (10, 197, 3, 4), + (3072, 768, 50432, 64, 64, False, True, True): (1, 394, 1, 4), + (3072, 768, 50432, 64, 64, True, False, True): (3, 197, 2, 4), + (3072, 768, 50432, 128, 128, False, True, True): (3, 394, 1, 4), + (3072, 768, 50432, 128, 128, True, False, True): (2, 394, 2, 4), + (3072, 768, 65536, 16, 16, False, True, True): (1, 256, 1, 4), + (3072, 768, 65536, 16, 16, True, False, True): (15, 256, 4, 1), + (3072, 768, 65536, 32, 32, False, True, True): (2, 256, 1, 4), + (3072, 768, 65536, 32, 32, True, False, True): (10, 256, 3, 4), + (3072, 768, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (3072, 768, 65536, 64, 64, True, False, True): (3, 256, 2, 4), + (3072, 768, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (3072, 768, 65536, 128, 128, True, False, True): (3, 512, 2, 4), + (3072, 768, 131072, 16, 16, False, True, True): (1, 512, 1, 4), + (3072, 768, 131072, 16, 16, True, False, True): (15, 512, 4, 1), + (3072, 768, 131072, 32, 32, False, True, True): (2, 512, 1, 4), + (3072, 768, 131072, 32, 32, True, False, True): (9, 512, 3, 4), + (3072, 768, 131072, 64, 64, False, True, True): (1, 1024, 1, 4), + (3072, 768, 131072, 64, 64, True, False, True): (3, 512, 2, 4), + (3072, 768, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (3072, 768, 131072, 128, 128, True, False, True): (3, 1024, 2, 4), + (3072, 3072, 256, 16, 16, False, True, True): (5, 4, 1, 4), + (3072, 3072, 256, 16, 16, True, False, True): (1, 2, 5, 2), + (3072, 3072, 256, 32, 32, False, True, True): (1, 4, 1, 8), + (3072, 3072, 256, 32, 32, True, False, True): (3, 4, 4, 2), + (3072, 3072, 256, 64, 64, False, True, True): (2, 4, 3, 4), + (3072, 3072, 256, 64, 64, True, False, True): (3, 4, 4, 4), + (3072, 3072, 256, 128, 128, False, True, True): (1, 2, 3, 8), + (3072, 3072, 256, 128, 128, True, False, True): (1, 2, 3, 8), + (3072, 3072, 512, 16, 16, False, True, True): (5, 4, 1, 2), + (3072, 3072, 512, 16, 16, True, False, True): (1, 2, 4, 4), + (3072, 3072, 512, 32, 32, False, True, True): (3, 8, 1, 4), + (3072, 3072, 512, 32, 32, True, False, True): (4, 2, 3, 4), + (3072, 3072, 512, 64, 64, False, True, True): (1, 8, 2, 2), + (3072, 3072, 512, 64, 64, True, False, True): (2, 4, 3, 4), + (3072, 3072, 512, 128, 128, False, True, True): (1, 4, 3, 8), + (3072, 3072, 512, 128, 128, True, False, True): (4, 4, 3, 8), + (3072, 3072, 1024, 16, 16, False, True, True): (1, 8, 1, 4), + (3072, 3072, 1024, 16, 16, True, False, True): (4, 8, 5, 2), + (3072, 3072, 1024, 32, 32, False, True, True): (1, 8, 1, 8), + (3072, 3072, 1024, 32, 32, True, False, True): (1, 4, 4, 4), + (3072, 3072, 1024, 64, 64, False, True, True): (3, 8, 3, 4), + (3072, 3072, 1024, 64, 64, True, False, True): (2, 4, 3, 4), + (3072, 3072, 1024, 128, 128, False, True, True): (3, 8, 1, 4), + (3072, 3072, 1024, 128, 128, True, False, True): (1, 8, 3, 8), + (3072, 3072, 2048, 16, 16, False, True, True): (1, 16, 1, 2), + (3072, 3072, 2048, 16, 16, True, False, True): (4, 16, 4, 2), + (3072, 3072, 2048, 32, 32, False, True, True): (1, 16, 1, 8), + (3072, 3072, 2048, 32, 32, True, False, True): (3, 8, 4, 4), + (3072, 3072, 2048, 64, 64, False, True, True): (3, 16, 3, 4), + (3072, 3072, 2048, 64, 64, True, False, True): (3, 8, 3, 4), + (3072, 3072, 2048, 128, 128, False, True, True): (4, 16, 3, 8), + (3072, 3072, 2048, 128, 128, True, False, True): (3, 16, 3, 8), + (3072, 3072, 4096, 16, 16, False, True, True): (1, 32, 1, 2), + (3072, 3072, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (3072, 3072, 4096, 32, 32, False, True, True): (1, 32, 1, 8), + (3072, 3072, 4096, 32, 32, True, False, True): (3, 16, 3, 4), + (3072, 3072, 4096, 64, 64, False, True, True): (1, 32, 3, 4), + (3072, 3072, 4096, 64, 64, True, False, True): (3, 16, 3, 4), + (3072, 3072, 4096, 128, 128, False, True, True): (1, 32, 3, 8), + (3072, 3072, 4096, 128, 128, True, False, True): (3, 32, 3, 8), + (3072, 3072, 8192, 16, 16, False, True, True): (1, 64, 1, 2), + (3072, 3072, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (3072, 3072, 8192, 32, 32, False, True, True): (1, 64, 1, 8), + (3072, 3072, 8192, 32, 32, True, False, True): (8, 32, 3, 4), + (3072, 3072, 8192, 64, 64, False, True, True): (3, 64, 3, 4), + (3072, 3072, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (3072, 3072, 8192, 128, 128, False, True, True): (2, 64, 3, 8), + (3072, 3072, 8192, 128, 128, True, False, True): (1, 64, 3, 8), + (3072, 3072, 16384, 16, 16, False, True, True): (1, 128, 1, 2), + (3072, 3072, 16384, 16, 16, True, False, True): (4, 128, 4, 2), + (3072, 3072, 16384, 32, 32, False, True, True): (1, 64, 1, 2), + (3072, 3072, 16384, 32, 32, True, False, True): (4, 64, 3, 4), + (3072, 3072, 16384, 64, 64, False, True, True): (1, 128, 3, 4), + (3072, 3072, 16384, 64, 64, True, False, True): (4, 64, 3, 4), + (3072, 3072, 16384, 128, 128, False, True, True): (3, 128, 1, 4), + (3072, 3072, 16384, 128, 128, True, False, True): (1, 128, 3, 8), + (3072, 3072, 32768, 16, 16, False, True, True): (1, 256, 1, 2), + (3072, 3072, 32768, 16, 16, True, False, True): (8, 128, 4, 4), + (3072, 3072, 32768, 32, 32, False, True, True): (1, 256, 1, 8), + (3072, 3072, 32768, 32, 32, True, False, True): (5, 128, 3, 4), + (3072, 3072, 32768, 64, 64, False, True, True): (1, 256, 3, 4), + (3072, 3072, 32768, 64, 64, True, False, True): (1, 128, 3, 4), + (3072, 3072, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (3072, 3072, 32768, 128, 128, True, False, True): (3, 256, 2, 4), + (3072, 3072, 65536, 16, 16, False, True, True): (1, 512, 1, 2), + (3072, 3072, 65536, 16, 16, True, False, True): (7, 256, 4, 4), + (3072, 3072, 65536, 32, 32, False, True, True): (1, 256, 1, 2), + (3072, 3072, 65536, 32, 32, True, False, True): (5, 256, 3, 4), + (3072, 3072, 65536, 64, 64, False, True, True): (1, 512, 3, 4), + (3072, 3072, 65536, 64, 64, True, False, True): (3, 256, 3, 4), + (3072, 3072, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (3072, 3072, 65536, 128, 128, True, False, True): (3, 512, 2, 4), + (3072, 3072, 131072, 16, 16, False, True, True): (1, 1024, 1, 2), + (3072, 3072, 131072, 16, 16, True, False, True): (5, 512, 4, 4), + (3072, 3072, 131072, 32, 32, False, True, True): (1, 512, 1, 2), + (3072, 3072, 131072, 32, 32, True, False, True): (3, 512, 3, 4), + (3072, 3072, 131072, 64, 64, False, True, True): (1, 1024, 3, 4), + (3072, 3072, 131072, 64, 64, True, False, True): (3, 512, 3, 4), + (3072, 3072, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (3072, 3072, 131072, 128, 128, True, False, True): (1, 1024, 2, 4), + (4096, 4096, 256, 16, 16, False, True, True): (2, 2, 6, 4), + (4096, 4096, 256, 16, 16, True, False, True): (2, 2, 5, 4), + (4096, 4096, 256, 32, 32, False, True, True): (7, 2, 4, 4), + (4096, 4096, 256, 32, 32, True, False, True): (1, 2, 4, 4), + (4096, 4096, 256, 64, 64, False, True, True): (3, 4, 3, 4), + (4096, 4096, 256, 64, 64, True, False, True): (3, 4, 3, 4), + (4096, 4096, 256, 128, 128, False, True, True): (1, 2, 2, 8), + (4096, 4096, 256, 128, 128, True, False, True): (1, 2, 2, 8), + (4096, 4096, 512, 16, 16, False, True, True): (4, 2, 3, 4), + (4096, 4096, 512, 16, 16, True, False, True): (2, 4, 3, 2), + (4096, 4096, 512, 32, 32, False, True, True): (3, 4, 3, 4), + (4096, 4096, 512, 32, 32, True, False, True): (3, 4, 3, 2), + (4096, 4096, 512, 64, 64, False, True, True): (3, 4, 3, 4), + (4096, 4096, 512, 64, 64, True, False, True): (3, 4, 3, 4), + (4096, 4096, 512, 128, 128, False, True, True): (2, 4, 2, 8), + (4096, 4096, 512, 128, 128, True, False, True): (2, 4, 1, 4), + (4096, 4096, 1024, 16, 16, False, True, True): (2, 8, 3, 2), + (4096, 4096, 1024, 16, 16, True, False, True): (2, 8, 3, 2), + (4096, 4096, 1024, 32, 32, False, True, True): (3, 8, 3, 4), + (4096, 4096, 1024, 32, 32, True, False, True): (1, 8, 3, 2), + (4096, 4096, 1024, 64, 64, False, True, True): (1, 8, 3, 4), + (4096, 4096, 1024, 64, 64, True, False, True): (1, 8, 3, 4), + (4096, 4096, 1024, 128, 128, False, True, True): (2, 8, 2, 8), + (4096, 4096, 1024, 128, 128, True, False, True): (2, 8, 2, 8), + (4096, 4096, 2048, 16, 16, False, True, True): (2, 8, 4, 4), + (4096, 4096, 2048, 16, 16, True, False, True): (2, 8, 4, 4), + (4096, 4096, 2048, 32, 32, False, True, True): (4, 8, 4, 8), + (4096, 4096, 2048, 32, 32, True, False, True): (4, 8, 4, 8), + (4096, 4096, 2048, 64, 64, False, True, True): (1, 16, 3, 4), + (4096, 4096, 2048, 64, 64, True, False, True): (4, 16, 3, 4), + (4096, 4096, 2048, 128, 128, False, True, True): (2, 16, 2, 8), + (4096, 4096, 2048, 128, 128, True, False, True): (4, 16, 1, 4), + (4096, 4096, 4096, 16, 16, False, True, True): (4, 32, 4, 4), + (4096, 4096, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (4096, 4096, 4096, 32, 32, False, True, True): (4, 16, 4, 8), + (4096, 4096, 4096, 32, 32, True, False, True): (4, 16, 3, 8), + (4096, 4096, 4096, 64, 64, False, True, True): (1, 32, 3, 4), + (4096, 4096, 4096, 64, 64, True, False, True): (1, 32, 3, 4), + (4096, 4096, 4096, 128, 128, False, True, True): (3, 32, 1, 4), + (4096, 4096, 4096, 128, 128, True, False, True): (2, 32, 1, 4), + (4096, 4096, 8192, 16, 16, False, True, True): (4, 64, 4, 2), + (4096, 4096, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (4096, 4096, 8192, 32, 32, False, True, True): (4, 32, 4, 8), + (4096, 4096, 8192, 32, 32, True, False, True): (4, 32, 4, 8), + (4096, 4096, 8192, 64, 64, False, True, True): (2, 64, 3, 4), + (4096, 4096, 8192, 64, 64, True, False, True): (2, 64, 3, 4), + (4096, 4096, 8192, 128, 128, False, True, True): (3, 64, 1, 4), + (4096, 4096, 8192, 128, 128, True, False, True): (1, 64, 1, 4), + (4096, 4096, 16384, 16, 16, False, True, True): (4, 64, 3, 4), + (4096, 4096, 16384, 16, 16, True, False, True): (4, 64, 4, 4), + (4096, 4096, 16384, 32, 32, False, True, True): (4, 64, 4, 8), + (4096, 4096, 16384, 32, 32, True, False, True): (4, 64, 4, 8), + (4096, 4096, 16384, 64, 64, False, True, True): (1, 64, 2, 4), + (4096, 4096, 16384, 64, 64, True, False, True): (1, 64, 3, 8), + (4096, 4096, 16384, 128, 128, False, True, True): (3, 128, 1, 4), + (4096, 4096, 16384, 128, 128, True, False, True): (1, 128, 1, 4), + (4096, 4096, 32768, 16, 16, False, True, True): (8, 128, 3, 2), + (4096, 4096, 32768, 16, 16, True, False, True): (5, 128, 4, 4), + (4096, 4096, 32768, 32, 32, False, True, True): (3, 128, 4, 4), + (4096, 4096, 32768, 32, 32, True, False, True): (3, 128, 4, 8), + (4096, 4096, 32768, 64, 64, False, True, True): (1, 128, 2, 4), + (4096, 4096, 32768, 64, 64, True, False, True): (3, 256, 3, 4), + (4096, 4096, 32768, 128, 128, False, True, True): (3, 256, 1, 4), + (4096, 4096, 32768, 128, 128, True, False, True): (1, 256, 1, 4), + (4096, 4096, 50432, 16, 16, False, True, True): (1, 197, 1, 4), + (4096, 4096, 50432, 16, 16, True, False, True): (4, 197, 4, 1), + (4096, 4096, 50432, 32, 32, False, True, True): (1, 197, 1, 4), + (4096, 4096, 50432, 32, 32, True, False, True): (2, 197, 3, 4), + (4096, 4096, 50432, 64, 64, False, True, True): (1, 394, 3, 4), + (4096, 4096, 50432, 64, 64, True, False, True): (1, 197, 2, 4), + (4096, 4096, 50432, 128, 128, False, True, True): (3, 394, 1, 4), + (4096, 4096, 50432, 128, 128, True, False, True): (1, 394, 2, 4), + (4096, 4096, 65536, 16, 16, False, True, True): (5, 256, 4, 4), + (4096, 4096, 65536, 16, 16, True, False, True): (5, 256, 4, 4), + (4096, 4096, 65536, 32, 32, False, True, True): (4, 256, 4, 8), + (4096, 4096, 65536, 32, 32, True, False, True): (4, 256, 3, 8), + (4096, 4096, 65536, 64, 64, False, True, True): (1, 256, 2, 4), + (4096, 4096, 65536, 64, 64, True, False, True): (1, 512, 3, 4), + (4096, 4096, 65536, 128, 128, False, True, True): (3, 512, 1, 4), + (4096, 4096, 65536, 128, 128, True, False, True): (1, 512, 1, 4), + (4096, 4096, 65792, 16, 16, False, True, True): (1, 257, 1, 4), + (4096, 4096, 65792, 16, 16, True, False, True): (5, 257, 4, 1), + (4096, 4096, 65792, 32, 32, False, True, True): (1, 257, 1, 4), + (4096, 4096, 65792, 32, 32, True, False, True): (1, 257, 3, 4), + (4096, 4096, 65792, 64, 64, False, True, True): (1, 514, 3, 4), + (4096, 4096, 65792, 64, 64, True, False, True): (1, 257, 2, 4), + (4096, 4096, 65792, 128, 128, False, True, True): (3, 514, 1, 4), + (4096, 4096, 65792, 128, 128, True, False, True): (1, 514, 2, 4), + (4096, 4096, 131072, 16, 16, False, True, True): (4, 512, 3, 4), + (4096, 4096, 131072, 16, 16, True, False, True): (5, 512, 4, 4), + (4096, 4096, 131072, 32, 32, False, True, True): (1, 512, 4, 8), + (4096, 4096, 131072, 32, 32, True, False, True): (4, 512, 4, 8), + (4096, 4096, 131072, 64, 64, False, True, True): (1, 512, 2, 4), + (4096, 4096, 131072, 64, 64, True, False, True): (1, 512, 2, 4), + (4096, 4096, 131072, 128, 128, False, True, True): (3, 1024, 1, 4), + (4096, 4096, 131072, 128, 128, True, False, True): (1, 1024, 1, 4), + (5120, 1280, 65792, 16, 16, False, True, True): (1, 257, 1, 4), + (5120, 1280, 65792, 16, 16, True, False, True): (7, 257, 4, 1), + (5120, 1280, 65792, 32, 32, False, True, True): (2, 257, 1, 4), + (5120, 1280, 65792, 32, 32, True, False, True): (5, 257, 3, 4), + (5120, 1280, 65792, 64, 64, False, True, True): (1, 514, 1, 4), + (5120, 1280, 65792, 64, 64, True, False, True): (5, 257, 2, 4), + (5120, 1280, 65792, 128, 128, False, True, True): (3, 514, 1, 4), + (5120, 1280, 65792, 128, 128, True, False, True): (4, 514, 2, 4), + (6144, 6144, 256, 16, 16, False, True, True): (1, 2, 1, 4), + (6144, 6144, 256, 16, 16, True, False, True): (1, 1, 4, 4), + (6144, 6144, 256, 32, 32, False, True, True): (3, 2, 1, 8), + (6144, 6144, 256, 32, 32, True, False, True): (2, 1, 3, 4), + (6144, 6144, 256, 64, 64, False, True, True): (2, 2, 3, 4), + (6144, 6144, 256, 64, 64, True, False, True): (6, 2, 4, 4), + (6144, 6144, 256, 128, 128, False, True, True): (2, 2, 3, 8), + (6144, 6144, 256, 128, 128, True, False, True): (1, 2, 3, 8), + (6144, 6144, 512, 16, 16, False, True, True): (4, 4, 1, 4), + (6144, 6144, 512, 16, 16, True, False, True): (3, 2, 3, 1), + (6144, 6144, 512, 32, 32, False, True, True): (1, 8, 1, 4), + (6144, 6144, 512, 32, 32, True, False, True): (2, 2, 3, 8), + (6144, 6144, 512, 64, 64, False, True, True): (4, 4, 3, 4), + (6144, 6144, 512, 64, 64, True, False, True): (6, 2, 3, 4), + (6144, 6144, 512, 128, 128, False, True, True): (3, 4, 1, 4), + (6144, 6144, 512, 128, 128, True, False, True): (4, 4, 3, 8), + (6144, 6144, 1024, 16, 16, False, True, True): (1, 8, 1, 2), + (6144, 6144, 1024, 16, 16, True, False, True): (4, 8, 4, 2), + (6144, 6144, 1024, 32, 32, False, True, True): (1, 8, 4, 2), + (6144, 6144, 1024, 32, 32, True, False, True): (1, 8, 4, 2), + (6144, 6144, 1024, 64, 64, False, True, True): (4, 8, 3, 4), + (6144, 6144, 1024, 64, 64, True, False, True): (1, 4, 3, 4), + (6144, 6144, 1024, 128, 128, False, True, True): (3, 8, 1, 4), + (6144, 6144, 1024, 128, 128, True, False, True): (1, 8, 3, 8), + (6144, 6144, 2048, 16, 16, False, True, True): (4, 4, 1, 4), + (6144, 6144, 2048, 16, 16, True, False, True): (2, 8, 4, 4), + (6144, 6144, 2048, 32, 32, False, True, True): (4, 8, 3, 4), + (6144, 6144, 2048, 32, 32, True, False, True): (2, 8, 3, 4), + (6144, 6144, 2048, 64, 64, False, True, True): (4, 16, 3, 4), + (6144, 6144, 2048, 64, 64, True, False, True): (2, 8, 3, 4), + (6144, 6144, 2048, 128, 128, False, True, True): (3, 16, 1, 4), + (6144, 6144, 2048, 128, 128, True, False, True): (4, 16, 3, 8), + (6144, 6144, 4096, 16, 16, False, True, True): (4, 8, 1, 4), + (6144, 6144, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (6144, 6144, 4096, 32, 32, False, True, True): (4, 16, 1, 2), + (6144, 6144, 4096, 32, 32, True, False, True): (2, 8, 3, 8), + (6144, 6144, 4096, 64, 64, False, True, True): (4, 32, 3, 4), + (6144, 6144, 4096, 64, 64, True, False, True): (4, 16, 3, 4), + (6144, 6144, 4096, 128, 128, False, True, True): (6, 32, 1, 4), + (6144, 6144, 4096, 128, 128, True, False, True): (4, 32, 3, 8), + (6144, 6144, 8192, 16, 16, False, True, True): (2, 16, 1, 2), + (6144, 6144, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (6144, 6144, 8192, 32, 32, False, True, True): (4, 32, 1, 2), + (6144, 6144, 8192, 32, 32, True, False, True): (4, 32, 3, 4), + (6144, 6144, 8192, 64, 64, False, True, True): (4, 64, 3, 4), + (6144, 6144, 8192, 64, 64, True, False, True): (4, 32, 3, 4), + (6144, 6144, 8192, 128, 128, False, True, True): (6, 64, 1, 4), + (6144, 6144, 8192, 128, 128, True, False, True): (4, 64, 3, 8), + (6144, 6144, 16384, 16, 16, False, True, True): (2, 32, 1, 2), + (6144, 6144, 16384, 16, 16, True, False, True): (4, 64, 4, 4), + (6144, 6144, 16384, 32, 32, False, True, True): (4, 64, 1, 2), + (6144, 6144, 16384, 32, 32, True, False, True): (4, 64, 3, 4), + (6144, 6144, 16384, 64, 64, False, True, True): (4, 128, 3, 4), + (6144, 6144, 16384, 64, 64, True, False, True): (1, 32, 3, 8), + (6144, 6144, 16384, 128, 128, False, True, True): (4, 128, 1, 4), + (6144, 6144, 16384, 128, 128, True, False, True): (4, 128, 3, 8), + (6144, 6144, 32768, 16, 16, False, True, True): (2, 64, 1, 2), + (6144, 6144, 32768, 16, 16, True, False, True): (5, 128, 4, 1), + (6144, 6144, 32768, 32, 32, False, True, True): (4, 128, 1, 2), + (6144, 6144, 32768, 32, 32, True, False, True): (3, 128, 3, 4), + (6144, 6144, 32768, 64, 64, False, True, True): (4, 256, 3, 4), + (6144, 6144, 32768, 64, 64, True, False, True): (2, 64, 3, 8), + (6144, 6144, 32768, 128, 128, False, True, True): (8, 256, 1, 4), + (6144, 6144, 32768, 128, 128, True, False, True): (4, 256, 3, 8), + (6144, 6144, 65536, 16, 16, False, True, True): (2, 128, 1, 2), + (6144, 6144, 65536, 16, 16, True, False, True): (5, 256, 4, 1), + (6144, 6144, 65536, 32, 32, False, True, True): (4, 256, 1, 2), + (6144, 6144, 65536, 32, 32, True, False, True): (2, 256, 3, 4), + (6144, 6144, 65536, 64, 64, False, True, True): (4, 512, 3, 4), + (6144, 6144, 65536, 64, 64, True, False, True): (1, 128, 3, 8), + (6144, 6144, 65536, 128, 128, False, True, True): (4, 512, 1, 4), + (6144, 6144, 65536, 128, 128, True, False, True): (4, 512, 3, 8), + (6144, 6144, 131072, 16, 16, False, True, True): (2, 256, 1, 2), + (6144, 6144, 131072, 16, 16, True, False, True): (3, 512, 4, 4), + (6144, 6144, 131072, 32, 32, False, True, True): (4, 512, 1, 2), + (6144, 6144, 131072, 32, 32, True, False, True): (4, 512, 3, 4), + (6144, 6144, 131072, 64, 64, False, True, True): (4, 1024, 3, 4), + (6144, 6144, 131072, 64, 64, True, False, True): (2, 256, 3, 8), + (6144, 6144, 131072, 128, 128, False, True, True): (4, 1024, 1, 4), + (6144, 6144, 131072, 128, 128, True, False, True): (4, 1024, 3, 8), + (8192, 8192, 256, 16, 16, False, True, True): (2, 2, 6, 4), + (8192, 8192, 256, 16, 16, True, False, True): (2, 4, 2, 2), + (8192, 8192, 256, 32, 32, False, True, True): (4, 2, 3, 4), + (8192, 8192, 256, 32, 32, True, False, True): (4, 2, 3, 4), + (8192, 8192, 256, 64, 64, False, True, True): (2, 2, 3, 8), + (8192, 8192, 256, 64, 64, True, False, True): (6, 2, 3, 8), + (8192, 8192, 256, 128, 128, False, True, True): (3, 2, 1, 4), + (8192, 8192, 256, 128, 128, True, False, True): (1, 2, 1, 4), + (8192, 8192, 512, 16, 16, False, True, True): (4, 4, 3, 2), + (8192, 8192, 512, 16, 16, True, False, True): (4, 4, 3, 4), + (8192, 8192, 512, 32, 32, False, True, True): (1, 4, 3, 4), + (8192, 8192, 512, 32, 32, True, False, True): (5, 4, 3, 2), + (8192, 8192, 512, 64, 64, False, True, True): (1, 4, 3, 4), + (8192, 8192, 512, 64, 64, True, False, True): (2, 2, 3, 8), + (8192, 8192, 512, 128, 128, False, True, True): (4, 4, 2, 8), + (8192, 8192, 512, 128, 128, True, False, True): (4, 4, 2, 8), + (8192, 8192, 1024, 16, 16, False, True, True): (4, 8, 4, 4), + (8192, 8192, 1024, 16, 16, True, False, True): (4, 8, 4, 4), + (8192, 8192, 1024, 32, 32, False, True, True): (2, 4, 4, 8), + (8192, 8192, 1024, 32, 32, True, False, True): (1, 4, 3, 4), + (8192, 8192, 1024, 64, 64, False, True, True): (4, 8, 3, 4), + (8192, 8192, 1024, 64, 64, True, False, True): (2, 8, 3, 4), + (8192, 8192, 1024, 128, 128, False, True, True): (4, 8, 2, 8), + (8192, 8192, 1024, 128, 128, True, False, True): (4, 8, 1, 4), + (8192, 8192, 2048, 16, 16, False, True, True): (2, 8, 4, 4), + (8192, 8192, 2048, 16, 16, True, False, True): (2, 8, 4, 4), + (8192, 8192, 2048, 32, 32, False, True, True): (2, 8, 4, 8), + (8192, 8192, 2048, 32, 32, True, False, True): (2, 8, 4, 8), + (8192, 8192, 2048, 64, 64, False, True, True): (4, 8, 2, 4), + (8192, 8192, 2048, 64, 64, True, False, True): (4, 16, 3, 4), + (8192, 8192, 2048, 128, 128, False, True, True): (6, 16, 1, 4), + (8192, 8192, 2048, 128, 128, True, False, True): (4, 16, 1, 4), + (8192, 8192, 4096, 16, 16, False, True, True): (4, 32, 4, 2), + (8192, 8192, 4096, 16, 16, True, False, True): (4, 32, 4, 2), + (8192, 8192, 4096, 32, 32, False, True, True): (2, 16, 4, 8), + (8192, 8192, 4096, 32, 32, True, False, True): (4, 16, 4, 8), + (8192, 8192, 4096, 64, 64, False, True, True): (4, 16, 2, 4), + (8192, 8192, 4096, 64, 64, True, False, True): (4, 16, 2, 4), + (8192, 8192, 4096, 128, 128, False, True, True): (6, 32, 1, 4), + (8192, 8192, 4096, 128, 128, True, False, True): (4, 32, 1, 4), + (8192, 8192, 8192, 16, 16, False, True, True): (4, 64, 4, 2), + (8192, 8192, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (8192, 8192, 8192, 32, 32, False, True, True): (2, 32, 4, 8), + (8192, 8192, 8192, 32, 32, True, False, True): (2, 32, 4, 8), + (8192, 8192, 8192, 64, 64, False, True, True): (2, 32, 2, 4), + (8192, 8192, 8192, 64, 64, True, False, True): (4, 32, 2, 4), + (8192, 8192, 8192, 128, 128, False, True, True): (6, 64, 1, 4), + (8192, 8192, 8192, 128, 128, True, False, True): (4, 64, 1, 4), + (8192, 8192, 16384, 16, 16, False, True, True): (4, 64, 3, 4), + (8192, 8192, 16384, 16, 16, True, False, True): (4, 64, 4, 4), + (8192, 8192, 16384, 32, 32, False, True, True): (4, 64, 4, 8), + (8192, 8192, 16384, 32, 32, True, False, True): (4, 64, 4, 8), + (8192, 8192, 16384, 64, 64, False, True, True): (4, 64, 2, 4), + (8192, 8192, 16384, 64, 64, True, False, True): (4, 64, 3, 8), + (8192, 8192, 16384, 128, 128, False, True, True): (6, 128, 1, 4), + (8192, 8192, 16384, 128, 128, True, False, True): (4, 128, 1, 4), + (8192, 8192, 32768, 16, 16, False, True, True): (3, 128, 4, 4), + (8192, 8192, 32768, 16, 16, True, False, True): (3, 128, 4, 4), + (8192, 8192, 32768, 32, 32, False, True, True): (2, 128, 4, 8), + (8192, 8192, 32768, 32, 32, True, False, True): (2, 128, 4, 8), + (8192, 8192, 32768, 64, 64, False, True, True): (2, 128, 2, 4), + (8192, 8192, 32768, 64, 64, True, False, True): (2, 128, 3, 8), + (8192, 8192, 32768, 128, 128, False, True, True): (6, 256, 1, 4), + (8192, 8192, 32768, 128, 128, True, False, True): (4, 256, 1, 4), + (8192, 8192, 50432, 16, 16, False, True, True): (1, 197, 1, 1), + (8192, 8192, 50432, 16, 16, True, False, True): (3, 197, 4, 1), + (8192, 8192, 50432, 32, 32, False, True, True): (2, 197, 1, 4), + (8192, 8192, 50432, 32, 32, True, False, True): (2, 197, 3, 4), + (8192, 8192, 50432, 64, 64, False, True, True): (2, 394, 3, 4), + (8192, 8192, 65536, 16, 16, False, True, True): (3, 256, 4, 4), + (8192, 8192, 65536, 16, 16, True, False, True): (4, 256, 4, 4), + (8192, 8192, 65536, 32, 32, False, True, True): (2, 256, 4, 8), + (8192, 8192, 65536, 32, 32, True, False, True): (2, 256, 3, 8), + (8192, 8192, 65536, 64, 64, False, True, True): (2, 256, 2, 4), + (8192, 8192, 65536, 64, 64, True, False, True): (4, 256, 3, 8), + (8192, 8192, 65536, 128, 128, False, True, True): (6, 512, 1, 4), + (8192, 8192, 65536, 128, 128, True, False, True): (4, 512, 1, 4), + (8192, 8192, 65792, 16, 16, False, True, True): (1, 257, 1, 1), + (8192, 8192, 65792, 16, 16, True, False, True): (3, 257, 4, 1), + (8192, 8192, 65792, 32, 32, False, True, True): (2, 257, 1, 4), + (8192, 8192, 65792, 32, 32, True, False, True): (1, 257, 3, 4), + (8192, 8192, 65792, 64, 64, False, True, True): (2, 514, 3, 4), + (8192, 8192, 65792, 64, 64, True, False, True): (1, 257, 3, 4), + (8192, 8192, 65792, 128, 128, False, True, True): (2, 514, 1, 4), + (8192, 8192, 65792, 128, 128, True, False, True): (2, 514, 3, 8), + (8192, 8192, 131072, 16, 16, False, True, True): (4, 512, 4, 4), + (8192, 8192, 131072, 16, 16, True, False, True): (3, 512, 4, 4), + (8192, 8192, 131072, 32, 32, False, True, True): (2, 512, 4, 8), + (8192, 8192, 131072, 32, 32, True, False, True): (2, 512, 4, 8), + (8192, 8192, 131072, 64, 64, False, True, True): (2, 512, 2, 4), + (8192, 8192, 131072, 64, 64, True, False, True): (2, 512, 2, 4), + (8192, 8192, 131072, 128, 128, False, True, True): (4, 1024, 1, 4), + (8192, 8192, 131072, 128, 128, True, False, True): (4, 1024, 1, 4), + (12288, 12288, 256, 16, 16, False, True, True): (4, 2, 1, 4), + (12288, 12288, 256, 16, 16, True, False, True): (1, 1, 3, 1), + (12288, 12288, 256, 32, 32, False, True, True): (4, 4, 1, 4), + (12288, 12288, 256, 32, 32, True, False, True): (2, 1, 3, 2), + (12288, 12288, 256, 64, 64, False, True, True): (4, 2, 3, 4), + (12288, 12288, 256, 64, 64, True, False, True): (3, 1, 3, 4), + (12288, 12288, 256, 128, 128, False, True, True): (6, 2, 1, 4), + (12288, 12288, 256, 128, 128, True, False, True): (4, 2, 3, 8), + (12288, 12288, 512, 16, 16, False, True, True): (4, 4, 1, 2), + (12288, 12288, 512, 16, 16, True, False, True): (4, 4, 4, 2), + (12288, 12288, 512, 32, 32, False, True, True): (4, 4, 4, 2), + (12288, 12288, 512, 32, 32, True, False, True): (2, 2, 3, 8), + (12288, 12288, 512, 64, 64, False, True, True): (4, 4, 3, 4), + (12288, 12288, 512, 64, 64, True, False, True): (8, 2, 3, 4), + (12288, 12288, 512, 128, 128, False, True, True): (4, 4, 3, 8), + (12288, 12288, 512, 128, 128, True, False, True): (4, 4, 3, 8), + (12288, 12288, 1024, 16, 16, False, True, True): (4, 8, 1, 2), + (12288, 12288, 1024, 16, 16, True, False, True): (2, 4, 4, 4), + (12288, 12288, 1024, 32, 32, False, True, True): (4, 4, 3, 4), + (12288, 12288, 1024, 32, 32, True, False, True): (1, 4, 3, 4), + (12288, 12288, 1024, 64, 64, False, True, True): (4, 8, 3, 4), + (12288, 12288, 1024, 64, 64, True, False, True): (2, 4, 3, 4), + (12288, 12288, 1024, 128, 128, False, True, True): (4, 8, 3, 8), + (12288, 12288, 1024, 128, 128, True, False, True): (4, 8, 3, 8), + (12288, 12288, 2048, 16, 16, False, True, True): (2, 4, 1, 4), + (12288, 12288, 2048, 16, 16, True, False, True): (2, 8, 4, 4), + (12288, 12288, 2048, 32, 32, False, True, True): (4, 8, 1, 2), + (12288, 12288, 2048, 32, 32, True, False, True): (2, 8, 4, 8), + (12288, 12288, 2048, 64, 64, False, True, True): (4, 16, 3, 4), + (12288, 12288, 2048, 64, 64, True, False, True): (2, 8, 3, 4), + (12288, 12288, 2048, 128, 128, False, True, True): (4, 16, 3, 8), + (12288, 12288, 2048, 128, 128, True, False, True): (4, 16, 3, 8), + (12288, 12288, 4096, 16, 16, False, True, True): (2, 8, 1, 4), + (12288, 12288, 4096, 16, 16, True, False, True): (2, 16, 4, 4), + (12288, 12288, 4096, 32, 32, False, True, True): (2, 16, 1, 2), + (12288, 12288, 4096, 32, 32, True, False, True): (2, 16, 3, 4), + (12288, 12288, 4096, 64, 64, False, True, True): (4, 32, 3, 4), + (12288, 12288, 4096, 64, 64, True, False, True): (2, 16, 3, 4), + (12288, 12288, 4096, 128, 128, False, True, True): (4, 32, 1, 4), + (12288, 12288, 4096, 128, 128, True, False, True): (4, 32, 3, 8), + (12288, 12288, 8192, 16, 16, False, True, True): (2, 32, 1, 1), + (12288, 12288, 8192, 16, 16, True, False, True): (4, 64, 4, 2), + (12288, 12288, 8192, 32, 32, False, True, True): (2, 32, 1, 2), + (12288, 12288, 8192, 32, 32, True, False, True): (2, 32, 3, 2), + (12288, 12288, 8192, 64, 64, False, True, True): (4, 64, 3, 4), + (12288, 12288, 8192, 64, 64, True, False, True): (2, 32, 3, 4), + (12288, 12288, 8192, 128, 128, False, True, True): (4, 64, 3, 8), + (12288, 12288, 8192, 128, 128, True, False, True): (2, 64, 3, 8), + (12288, 12288, 16384, 16, 16, False, True, True): (4, 128, 1, 2), + (12288, 12288, 16384, 16, 16, True, False, True): (4, 128, 4, 2), + (12288, 12288, 16384, 32, 32, False, True, True): (2, 64, 1, 2), + (12288, 12288, 16384, 32, 32, True, False, True): (2, 64, 3, 4), + (12288, 12288, 16384, 64, 64, False, True, True): (4, 128, 3, 4), + (12288, 12288, 16384, 64, 64, True, False, True): (2, 64, 3, 4), + (12288, 12288, 16384, 128, 128, False, True, True): (4, 128, 1, 4), + (12288, 12288, 16384, 128, 128, True, False, True): (4, 128, 3, 8), + (12288, 12288, 32768, 16, 16, False, True, True): (2, 128, 1, 1), + (12288, 12288, 32768, 16, 16, True, False, True): (3, 128, 4, 1), + (12288, 12288, 32768, 32, 32, False, True, True): (2, 128, 1, 2), + (12288, 12288, 32768, 32, 32, True, False, True): (2, 128, 3, 2), + (12288, 12288, 32768, 64, 64, False, True, True): (4, 256, 3, 4), + (12288, 12288, 32768, 64, 64, True, False, True): (1, 64, 3, 8), + (12288, 12288, 32768, 128, 128, False, True, True): (4, 256, 3, 8), + (12288, 12288, 32768, 128, 128, True, False, True): (4, 256, 3, 8), + (12288, 12288, 65536, 16, 16, False, True, True): (4, 512, 1, 2), + (12288, 12288, 65536, 16, 16, True, False, True): (3, 256, 4, 1), + (12288, 12288, 65536, 32, 32, False, True, True): (2, 256, 1, 2), + (12288, 12288, 65536, 32, 32, True, False, True): (2, 256, 3, 2), + (12288, 12288, 65536, 64, 64, False, True, True): (4, 512, 3, 4), + (12288, 12288, 65536, 64, 64, True, False, True): (2, 256, 3, 4), + (12288, 12288, 65536, 128, 128, False, True, True): (4, 512, 1, 4), + (12288, 12288, 65536, 128, 128, True, False, True): (4, 512, 3, 8), + (12288, 12288, 131072, 16, 16, False, True, True): (2, 512, 1, 1), + (12288, 12288, 131072, 16, 16, True, False, True): (2, 512, 4, 4), + (12288, 12288, 131072, 32, 32, False, True, True): (2, 512, 1, 2), + (12288, 12288, 131072, 32, 32, True, False, True): (2, 512, 3, 4), + (12288, 12288, 131072, 64, 64, False, True, True): (4, 1024, 3, 4), + (12288, 12288, 131072, 64, 64, True, False, True): (2, 512, 3, 4), + (12288, 12288, 131072, 128, 128, False, True, True): (4, 1024, 3, 8), + (12288, 12288, 131072, 128, 128, True, False, True): (4, 1024, 3, 8), + (16384, 16384, 256, 16, 16, False, True, True): (2, 2, 3, 2), + (16384, 16384, 256, 16, 16, True, False, True): (2, 2, 6, 4), + (16384, 16384, 256, 32, 32, False, True, True): (4, 2, 3, 4), + (16384, 16384, 256, 32, 32, True, False, True): (4, 2, 3, 2), + (16384, 16384, 256, 64, 64, False, True, True): (2, 2, 5, 4), + (16384, 16384, 256, 64, 64, True, False, True): (2, 2, 3, 8), + (16384, 16384, 256, 128, 128, False, True, True): (4, 2, 2, 8), + (16384, 16384, 256, 128, 128, True, False, True): (2, 2, 1, 4), + (16384, 16384, 512, 16, 16, False, True, True): (1, 2, 4, 4), + (16384, 16384, 512, 16, 16, True, False, True): (1, 2, 4, 4), + (16384, 16384, 512, 32, 32, False, True, True): (2, 2, 3, 8), + (16384, 16384, 512, 32, 32, True, False, True): (2, 2, 4, 8), + (16384, 16384, 512, 64, 64, False, True, True): (4, 4, 3, 4), + (16384, 16384, 512, 64, 64, True, False, True): (2, 4, 3, 4), + (16384, 16384, 512, 128, 128, False, True, True): (4, 4, 2, 8), + (16384, 16384, 512, 128, 128, True, False, True): (4, 4, 2, 8), + (16384, 16384, 1024, 16, 16, False, True, True): (4, 8, 4, 4), + (16384, 16384, 1024, 16, 16, True, False, True): (2, 4, 4, 4), + (16384, 16384, 1024, 32, 32, False, True, True): (2, 4, 4, 8), + (16384, 16384, 1024, 32, 32, True, False, True): (2, 4, 4, 8), + (16384, 16384, 1024, 64, 64, False, True, True): (4, 4, 2, 4), + (16384, 16384, 1024, 64, 64, True, False, True): (2, 4, 2, 4), + (16384, 16384, 1024, 128, 128, False, True, True): (6, 8, 1, 4), + (16384, 16384, 1024, 128, 128, True, False, True): (4, 8, 1, 4), + (16384, 16384, 2048, 16, 16, False, True, True): (2, 8, 4, 4), + (16384, 16384, 2048, 16, 16, True, False, True): (2, 8, 4, 4), + (16384, 16384, 2048, 32, 32, False, True, True): (2, 8, 4, 8), + (16384, 16384, 2048, 32, 32, True, False, True): (2, 8, 4, 8), + (16384, 16384, 2048, 64, 64, False, True, True): (2, 8, 2, 4), + (16384, 16384, 2048, 64, 64, True, False, True): (2, 8, 2, 4), + (16384, 16384, 2048, 128, 128, False, True, True): (4, 16, 2, 8), + (16384, 16384, 2048, 128, 128, True, False, True): (4, 16, 1, 4), + (16384, 16384, 4096, 16, 16, False, True, True): (2, 16, 4, 4), + (16384, 16384, 4096, 16, 16, True, False, True): (2, 16, 4, 4), + (16384, 16384, 4096, 32, 32, False, True, True): (1, 16, 4, 8), + (16384, 16384, 4096, 32, 32, True, False, True): (2, 16, 3, 4), + (16384, 16384, 4096, 64, 64, False, True, True): (1, 16, 2, 4), + (16384, 16384, 4096, 64, 64, True, False, True): (2, 16, 2, 4), + (16384, 16384, 4096, 128, 128, False, True, True): (4, 32, 2, 8), + (16384, 16384, 4096, 128, 128, True, False, True): (4, 32, 1, 4), + (16384, 16384, 8192, 16, 16, False, True, True): (2, 64, 4, 2), + (16384, 16384, 8192, 16, 16, True, False, True): (2, 64, 4, 2), + (16384, 16384, 8192, 32, 32, False, True, True): (2, 32, 4, 8), + (16384, 16384, 8192, 32, 32, True, False, True): (2, 32, 4, 8), + (16384, 16384, 8192, 64, 64, False, True, True): (2, 32, 2, 4), + (16384, 16384, 8192, 64, 64, True, False, True): (2, 32, 4, 8), + (16384, 16384, 8192, 128, 128, False, True, True): (4, 64, 2, 8), + (16384, 16384, 8192, 128, 128, True, False, True): (4, 64, 1, 4), + (16384, 16384, 16384, 16, 16, False, True, True): (1, 64, 4, 4), + (16384, 16384, 16384, 16, 16, True, False, True): (1, 64, 4, 4), + (16384, 16384, 16384, 32, 32, False, True, True): (1, 64, 4, 8), + (16384, 16384, 16384, 32, 32, True, False, True): (1, 64, 4, 8), + (16384, 16384, 16384, 64, 64, False, True, True): (1, 64, 2, 4), + (16384, 16384, 16384, 64, 64, True, False, True): (1, 64, 3, 8), + (16384, 16384, 16384, 128, 128, False, True, True): (4, 128, 1, 4), + (16384, 16384, 16384, 128, 128, True, False, True): (4, 128, 1, 4), + (16384, 16384, 32768, 16, 16, False, True, True): (1, 128, 4, 4), + (16384, 16384, 32768, 16, 16, True, False, True): (1, 128, 4, 4), + (16384, 16384, 32768, 32, 32, False, True, True): (1, 128, 3, 4), + (16384, 16384, 32768, 32, 32, True, False, True): (1, 128, 3, 8), + (16384, 16384, 32768, 64, 64, False, True, True): (2, 128, 2, 4), + (16384, 16384, 32768, 64, 64, True, False, True): (1, 128, 4, 8), + (16384, 16384, 32768, 128, 128, False, True, True): (4, 256, 2, 8), + (16384, 16384, 32768, 128, 128, True, False, True): (4, 256, 1, 4), + (16384, 16384, 65536, 16, 16, False, True, True): (1, 256, 3, 4), + (16384, 16384, 65536, 16, 16, True, False, True): (1, 256, 4, 4), + (16384, 16384, 65536, 32, 32, False, True, True): (1, 256, 4, 8), + (16384, 16384, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (16384, 16384, 65536, 64, 64, False, True, True): (2, 256, 2, 4), + (16384, 16384, 65536, 64, 64, True, False, True): (1, 256, 3, 8), + (16384, 16384, 65536, 128, 128, False, True, True): (4, 512, 2, 8), + (16384, 16384, 65536, 128, 128, True, False, True): (4, 512, 1, 4), + (16384, 16384, 65792, 16, 16, False, True, True): (1, 257, 1, 1), + (16384, 16384, 65792, 16, 16, True, False, True): (1, 257, 4, 1), + (16384, 16384, 65792, 32, 32, False, True, True): (1, 257, 1, 4), + (16384, 16384, 65792, 32, 32, True, False, True): (1, 257, 3, 4), + (16384, 16384, 65792, 64, 64, False, True, True): (2, 514, 3, 4), + (16384, 16384, 65792, 64, 64, True, False, True): (1, 257, 3, 4), + (16384, 16384, 65792, 128, 128, False, True, True): (2, 514, 3, 8), + (16384, 16384, 65792, 128, 128, True, False, True): (2, 514, 3, 8), + (16384, 16384, 131072, 16, 16, False, True, True): (1, 512, 4, 4), + (16384, 16384, 131072, 16, 16, True, False, True): (1, 512, 3, 2), + (16384, 16384, 131072, 32, 32, False, True, True): (1, 512, 4, 8), + (16384, 16384, 131072, 32, 32, True, False, True): (1, 512, 3, 2), + (16384, 16384, 131072, 64, 64, False, True, True): (1, 512, 2, 4), + (16384, 16384, 131072, 64, 64, True, False, True): (1, 512, 2, 4), + (16384, 16384, 131072, 128, 128, False, True, True): (4, 1024, 1, 4), + (16384, 16384, 131072, 128, 128, True, False, True): (4, 1024, 1, 4), + (24576, 24576, 256, 16, 16, False, True, True): (6, 2, 1, 2), + (24576, 24576, 256, 16, 16, True, False, True): (2, 2, 5, 4), + (24576, 24576, 256, 32, 32, False, True, True): (4, 4, 1, 4), + (24576, 24576, 256, 32, 32, True, False, True): (2, 2, 4, 2), + (24576, 24576, 256, 64, 64, False, True, True): (2, 2, 3, 4), + (24576, 24576, 256, 64, 64, True, False, True): (1, 1, 3, 4), + (24576, 24576, 256, 128, 128, False, True, True): (6, 2, 1, 4), + (24576, 24576, 256, 128, 128, True, False, True): (2, 2, 3, 8), + (24576, 24576, 512, 16, 16, False, True, True): (4, 4, 1, 2), + (24576, 24576, 512, 16, 16, True, False, True): (2, 2, 4, 4), + (24576, 24576, 512, 32, 32, False, True, True): (1, 2, 3, 4), + (24576, 24576, 512, 32, 32, True, False, True): (1, 2, 3, 4), + (24576, 24576, 512, 64, 64, False, True, True): (4, 4, 3, 4), + (24576, 24576, 512, 64, 64, True, False, True): (1, 2, 3, 4), + (24576, 24576, 512, 128, 128, False, True, True): (4, 4, 3, 8), + (24576, 24576, 512, 128, 128, True, False, True): (4, 4, 3, 8), + (24576, 24576, 1024, 16, 16, False, True, True): (2, 8, 1, 2), + (24576, 24576, 1024, 16, 16, True, False, True): (2, 4, 4, 4), + (24576, 24576, 1024, 32, 32, False, True, True): (2, 4, 1, 2), + (24576, 24576, 1024, 32, 32, True, False, True): (1, 4, 3, 4), + (24576, 24576, 1024, 64, 64, False, True, True): (4, 8, 3, 4), + (24576, 24576, 1024, 64, 64, True, False, True): (1, 4, 3, 4), + (24576, 24576, 1024, 128, 128, False, True, True): (4, 8, 3, 8), + (24576, 24576, 1024, 128, 128, True, False, True): (4, 8, 3, 8), + (24576, 24576, 2048, 16, 16, False, True, True): (1, 4, 1, 4), + (24576, 24576, 2048, 16, 16, True, False, True): (1, 8, 4, 4), + (24576, 24576, 2048, 32, 32, False, True, True): (2, 8, 1, 2), + (24576, 24576, 2048, 32, 32, True, False, True): (1, 8, 3, 4), + (24576, 24576, 2048, 64, 64, False, True, True): (4, 16, 3, 4), + (24576, 24576, 2048, 64, 64, True, False, True): (1, 4, 3, 8), + (24576, 24576, 2048, 128, 128, False, True, True): (4, 16, 3, 8), + (24576, 24576, 2048, 128, 128, True, False, True): (2, 16, 3, 8), + (24576, 24576, 4096, 16, 16, False, True, True): (2, 32, 1, 2), + (24576, 24576, 4096, 16, 16, True, False, True): (1, 16, 4, 4), + (24576, 24576, 4096, 32, 32, False, True, True): (1, 16, 1, 2), + (24576, 24576, 4096, 32, 32, True, False, True): (1, 16, 3, 4), + (24576, 24576, 4096, 64, 64, False, True, True): (4, 32, 3, 4), + (24576, 24576, 4096, 64, 64, True, False, True): (1, 8, 3, 8), + (24576, 24576, 4096, 128, 128, False, True, True): (4, 32, 3, 8), + (24576, 24576, 4096, 128, 128, True, False, True): (2, 32, 3, 8), + (24576, 24576, 8192, 16, 16, False, True, True): (1, 32, 1, 1), + (24576, 24576, 8192, 16, 16, True, False, True): (2, 64, 4, 2), + (24576, 24576, 8192, 32, 32, False, True, True): (1, 32, 1, 2), + (24576, 24576, 8192, 32, 32, True, False, True): (1, 32, 3, 4), + (24576, 24576, 8192, 64, 64, False, True, True): (4, 64, 3, 4), + (24576, 24576, 8192, 64, 64, True, False, True): (1, 32, 3, 4), + (24576, 24576, 8192, 128, 128, False, True, True): (4, 64, 3, 8), + (24576, 24576, 8192, 128, 128, True, False, True): (4, 64, 3, 8), + (24576, 24576, 16384, 16, 16, False, True, True): (2, 128, 1, 2), + (24576, 24576, 16384, 16, 16, True, False, True): (1, 64, 4, 4), + (24576, 24576, 16384, 32, 32, False, True, True): (1, 64, 1, 2), + (24576, 24576, 16384, 32, 32, True, False, True): (1, 64, 3, 2), + (24576, 24576, 16384, 64, 64, False, True, True): (2, 128, 3, 4), + (24576, 24576, 16384, 64, 64, True, False, True): (1, 32, 3, 8), + (24576, 24576, 16384, 128, 128, False, True, True): (4, 128, 3, 8), + (24576, 24576, 16384, 128, 128, True, False, True): (4, 128, 3, 8), + (24576, 24576, 32768, 16, 16, False, True, True): (1, 128, 1, 1), + (24576, 24576, 32768, 16, 16, True, False, True): (1, 128, 4, 4), + (24576, 24576, 32768, 32, 32, False, True, True): (1, 128, 1, 2), + (24576, 24576, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (24576, 24576, 32768, 64, 64, False, True, True): (2, 256, 3, 4), + (24576, 24576, 32768, 64, 64, True, False, True): (1, 128, 3, 4), + (24576, 24576, 32768, 128, 128, False, True, True): (4, 256, 3, 8), + (24576, 24576, 32768, 128, 128, True, False, True): (2, 256, 3, 8), + (24576, 24576, 65536, 16, 16, False, True, True): (2, 512, 1, 2), + (24576, 24576, 65536, 16, 16, True, False, True): (1, 256, 4, 4), + (32768, 32768, 256, 16, 16, False, True, True): (4, 2, 1, 2), + (32768, 32768, 256, 16, 16, True, False, True): (2, 2, 5, 4), + (32768, 32768, 256, 32, 32, False, True, True): (4, 2, 4, 2), + (32768, 32768, 256, 32, 32, True, False, True): (1, 1, 4, 8), + (32768, 32768, 256, 64, 64, False, True, True): (2, 2, 3, 4), + (32768, 32768, 256, 64, 64, True, False, True): (1, 1, 3, 8), + (32768, 32768, 256, 128, 128, False, True, True): (2, 2, 3, 8), + (32768, 32768, 256, 128, 128, True, False, True): (2, 2, 3, 8), + (32768, 32768, 512, 16, 16, False, True, True): (2, 2, 1, 4), + (32768, 32768, 512, 16, 16, True, False, True): (2, 2, 4, 2), + (32768, 32768, 512, 32, 32, False, True, True): (1, 2, 3, 4), + (32768, 32768, 512, 32, 32, True, False, True): (1, 2, 4, 8), + (32768, 32768, 512, 64, 64, False, True, True): (4, 4, 3, 4), + (32768, 32768, 512, 64, 64, True, False, True): (1, 2, 3, 4), + (32768, 32768, 512, 128, 128, False, True, True): (4, 4, 3, 8), + (32768, 32768, 512, 128, 128, True, False, True): (4, 4, 3, 8), + (32768, 32768, 1024, 16, 16, False, True, True): (2, 4, 1, 1), + (32768, 32768, 1024, 16, 16, True, False, True): (1, 4, 4, 2), + (32768, 32768, 1024, 32, 32, False, True, True): (2, 4, 1, 4), + (32768, 32768, 1024, 32, 32, True, False, True): (1, 4, 3, 4), + (32768, 32768, 1024, 64, 64, False, True, True): (4, 8, 3, 4), + (32768, 32768, 1024, 64, 64, True, False, True): (1, 4, 3, 4), + (32768, 32768, 1024, 128, 128, False, True, True): (4, 8, 3, 8), + (32768, 32768, 1024, 128, 128, True, False, True): (4, 8, 3, 8), + (32768, 32768, 2048, 16, 16, False, True, True): (1, 8, 1, 4), + (32768, 32768, 2048, 16, 16, True, False, True): (1, 8, 4, 4), + (32768, 32768, 2048, 32, 32, False, True, True): (2, 8, 1, 4), + (32768, 32768, 2048, 32, 32, True, False, True): (1, 8, 3, 4), + (32768, 32768, 2048, 64, 64, False, True, True): (4, 16, 3, 4), + (32768, 32768, 2048, 64, 64, True, False, True): (1, 8, 3, 4), + (32768, 32768, 2048, 128, 128, False, True, True): (4, 16, 3, 8), + (32768, 32768, 2048, 128, 128, True, False, True): (2, 16, 3, 8), + (32768, 32768, 4096, 16, 16, False, True, True): (1, 16, 1, 4), + (32768, 32768, 4096, 16, 16, True, False, True): (1, 16, 4, 4), + (32768, 32768, 4096, 32, 32, False, True, True): (2, 16, 1, 4), + (32768, 32768, 4096, 32, 32, True, False, True): (1, 16, 3, 4), + (32768, 32768, 4096, 64, 64, False, True, True): (2, 32, 3, 4), + (32768, 32768, 4096, 64, 64, True, False, True): (1, 16, 3, 4), + (32768, 32768, 4096, 128, 128, False, True, True): (4, 32, 3, 8), + (32768, 32768, 4096, 128, 128, True, False, True): (4, 32, 3, 8), + (32768, 32768, 8192, 16, 16, False, True, True): (1, 32, 1, 4), + (32768, 32768, 8192, 16, 16, True, False, True): (2, 64, 4, 1), + (32768, 32768, 8192, 32, 32, False, True, True): (2, 32, 1, 4), + (32768, 32768, 8192, 32, 32, True, False, True): (1, 32, 3, 4), + (32768, 32768, 8192, 64, 64, False, True, True): (2, 64, 3, 4), + (32768, 32768, 8192, 64, 64, True, False, True): (1, 32, 3, 4), + (32768, 32768, 8192, 128, 128, False, True, True): (4, 64, 3, 8), + (32768, 32768, 8192, 128, 128, True, False, True): (2, 64, 3, 8), + (32768, 32768, 16384, 16, 16, False, True, True): (1, 64, 1, 4), + (32768, 32768, 16384, 16, 16, True, False, True): (1, 64, 4, 1), + (32768, 32768, 16384, 32, 32, False, True, True): (2, 64, 1, 4), + (32768, 32768, 16384, 32, 32, True, False, True): (1, 64, 3, 4), + (32768, 32768, 16384, 64, 64, False, True, True): (2, 128, 3, 4), + (32768, 32768, 16384, 64, 64, True, False, True): (1, 64, 3, 4), + (32768, 32768, 16384, 128, 128, False, True, True): (4, 128, 3, 8), + (32768, 32768, 16384, 128, 128, True, False, True): (2, 128, 3, 8), + (32768, 32768, 32768, 16, 16, False, True, True): (1, 128, 1, 4), + (32768, 32768, 32768, 16, 16, True, False, True): (1, 128, 4, 1), + (32768, 32768, 32768, 32, 32, False, True, True): (2, 128, 1, 4), + (32768, 32768, 32768, 32, 32, True, False, True): (1, 128, 3, 4), + (32768, 32768, 32768, 64, 64, False, True, True): (2, 256, 3, 4), + (32768, 32768, 32768, 64, 64, True, False, True): (1, 128, 3, 4), + (32768, 32768, 32768, 128, 128, False, True, True): (2, 256, 3, 8), + (32768, 32768, 32768, 128, 128, True, False, True): (4, 256, 3, 8), + (32768, 32768, 65536, 16, 16, False, True, True): (1, 256, 1, 4), + (32768, 32768, 65536, 16, 16, True, False, True): (1, 256, 4, 1), + (32768, 32768, 65536, 32, 32, False, True, True): (1, 256, 3, 4), + (32768, 32768, 65536, 32, 32, True, False, True): (1, 256, 3, 4), + (32768, 32768, 65536, 64, 64, False, True, True): (1, 512, 3, 4), + (32768, 32768, 65536, 64, 64, True, False, True): (1, 256, 3, 4), + (32768, 32768, 65536, 128, 128, False, True, True): (4, 512, 1, 4), + (32768, 32768, 65536, 128, 128, True, False, True): (2, 512, 3, 8), + }, + ("bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.float16, 0.56)): { + (192, 192, 256, 64, 64, False, True, True): (1, 4, 3, 4), + (192, 192, 256, 64, 64, True, False, True): (1, 4, 3, 4), + (192, 192, 512, 64, 64, False, True, True): (1, 8, 5, 4), + (192, 192, 512, 64, 64, True, False, True): (1, 8, 3, 4), + (192, 192, 1024, 64, 64, False, True, True): (1, 16, 3, 2), + (192, 192, 1024, 64, 64, True, False, True): (1, 16, 3, 4), + (192, 192, 2048, 64, 64, False, True, True): (1, 32, 5, 4), + (192, 192, 2048, 64, 64, True, False, True): (4, 32, 5, 4), + (192, 192, 4096, 64, 64, False, True, True): (1, 64, 1, 8), + (192, 192, 4096, 64, 64, True, False, True): (1, 32, 3, 4), + (192, 192, 8192, 64, 64, False, True, True): (4, 128, 1, 4), + (192, 192, 8192, 64, 64, True, False, True): (3, 64, 3, 4), + (192, 192, 16384, 64, 64, False, True, True): (1, 256, 1, 4), + (192, 192, 16384, 64, 64, True, False, True): (3, 64, 2, 4), + (192, 192, 32768, 64, 64, False, True, True): (1, 512, 1, 2), + (192, 192, 32768, 64, 64, True, False, True): (2, 256, 2, 4), + (192, 192, 65536, 64, 64, False, True, True): (1, 512, 1, 4), + (192, 192, 65536, 64, 64, True, False, True): (2, 512, 2, 4), + (192, 192, 131072, 64, 64, False, True, True): (1, 1024, 1, 4), + (192, 192, 131072, 64, 64, True, False, True): (1, 512, 3, 4), + (384, 384, 256, 128, 128, False, True, True): (3, 2, 3, 8), + (384, 384, 256, 128, 128, True, False, True): (5, 2, 3, 8), + (384, 384, 512, 128, 128, False, True, True): (4, 4, 3, 8), + (384, 384, 512, 128, 128, True, False, True): (1, 4, 3, 8), + (384, 384, 1024, 128, 128, False, True, True): (1, 8, 3, 8), + (384, 384, 1024, 128, 128, True, False, True): (1, 8, 2, 8), + (384, 384, 2048, 128, 128, False, True, True): (3, 16, 3, 8), + (384, 384, 2048, 128, 128, True, False, True): (1, 16, 3, 8), + (384, 384, 4096, 128, 128, False, True, True): (3, 32, 3, 8), + (384, 384, 4096, 128, 128, True, False, True): (3, 32, 3, 8), + (384, 384, 8192, 128, 128, False, True, True): (2, 64, 3, 8), + (384, 384, 8192, 128, 128, True, False, True): (2, 64, 2, 4), + (384, 384, 16384, 128, 128, False, True, True): (1, 128, 2, 8), + (384, 384, 16384, 128, 128, True, False, True): (3, 128, 2, 4), + (384, 384, 32768, 128, 128, False, True, True): (2, 256, 3, 8), + (384, 384, 32768, 128, 128, True, False, True): (1, 256, 2, 4), + (384, 384, 65536, 128, 128, False, True, True): (7, 512, 1, 4), + (384, 384, 65536, 128, 128, True, False, True): (3, 512, 2, 4), + (384, 384, 131072, 128, 128, False, True, True): (5, 1024, 1, 4), + (384, 384, 131072, 128, 128, True, False, True): (1, 1024, 2, 4), + }, + ("bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.float32, 0.5)): { + (16, 16, 16, 16, 16, False, False, False): (2, 1, 1, 16), + (16, 16, 16, 16, 16, False, False, True): (1, 1, 2, 4), + (16, 16, 16, 16, 16, False, True, False): (1, 1, 2, 16), + (16, 16, 16, 16, 16, False, True, True): (2, 1, 2, 8), + (16, 16, 16, 16, 16, True, False, False): (1, 1, 1, 2), + (16, 16, 16, 16, 16, True, False, True): (2, 1, 1, 4), + (16, 16, 32, 16, 16, False, False, False): (1, 1, 1, 2), + (16, 16, 32, 16, 16, False, False, True): (1, 1, 2, 8), + (16, 16, 32, 16, 16, False, True, False): (1, 2, 1, 4), + (16, 16, 32, 16, 16, False, True, True): (1, 2, 2, 4), + (16, 16, 32, 16, 16, True, False, False): (1, 1, 2, 4), + (16, 16, 32, 16, 16, True, False, True): (1, 2, 2, 4), + (16, 16, 64, 16, 16, False, False, False): (1, 4, 1, 4), + (16, 16, 64, 16, 16, False, False, True): (2, 2, 1, 4), + (16, 16, 64, 16, 16, False, True, False): (1, 4, 1, 4), + (16, 16, 64, 16, 16, False, True, True): (1, 4, 1, 8), + (16, 16, 64, 16, 16, True, False, False): (1, 2, 1, 4), + (16, 16, 64, 16, 16, True, False, True): (1, 4, 2, 8), + (16, 32, 16, 16, 16, False, False, False): (1, 1, 2, 8), + (16, 32, 16, 16, 16, False, False, True): (2, 1, 1, 4), + (16, 32, 16, 16, 16, False, True, False): (1, 1, 1, 4), + (16, 32, 16, 16, 16, False, True, True): (1, 1, 1, 4), + (16, 32, 16, 16, 16, True, False, False): (1, 1, 1, 4), + (16, 32, 16, 16, 16, True, False, True): (1, 1, 2, 8), + (16, 32, 16, 16, 32, False, False, False): (1, 1, 2, 4), + (16, 32, 16, 16, 32, False, False, True): (2, 1, 2, 2), + (16, 32, 16, 16, 32, False, True, False): (1, 1, 1, 8), + (16, 32, 16, 16, 32, False, True, True): (1, 1, 1, 2), + (16, 32, 16, 16, 32, True, False, False): (3, 1, 1, 4), + (16, 32, 16, 16, 32, True, False, True): (1, 1, 1, 4), + (16, 32, 32, 16, 16, False, False, False): (1, 2, 1, 4), + (16, 32, 32, 16, 16, False, False, True): (2, 2, 1, 4), + (16, 32, 32, 16, 16, False, True, False): (1, 2, 1, 2), + (16, 32, 32, 16, 16, False, True, True): (1, 2, 1, 4), + (16, 32, 32, 16, 16, True, False, False): (1, 2, 1, 4), + (16, 32, 32, 16, 16, True, False, True): (1, 2, 1, 4), + (16, 32, 32, 16, 32, False, False, False): (1, 1, 2, 4), + (16, 32, 32, 16, 32, False, False, True): (1, 2, 1, 4), + (16, 32, 32, 16, 32, False, True, False): (1, 2, 2, 8), + (16, 32, 32, 16, 32, False, True, True): (1, 2, 1, 1), + (16, 32, 32, 16, 32, True, False, False): (1, 2, 1, 2), + (16, 32, 32, 16, 32, True, False, True): (1, 2, 1, 4), + (16, 32, 64, 16, 16, False, False, False): (1, 2, 1, 4), + (16, 32, 64, 16, 16, False, False, True): (2, 4, 1, 4), + (16, 32, 64, 16, 16, False, True, False): (1, 4, 2, 4), + (16, 32, 64, 16, 16, False, True, True): (1, 4, 1, 4), + (16, 32, 64, 16, 16, True, False, False): (1, 2, 2, 8), + (16, 32, 64, 16, 16, True, False, True): (1, 4, 1, 2), + (16, 32, 64, 16, 32, False, False, False): (1, 4, 1, 4), + (16, 32, 64, 16, 32, False, False, True): (1, 4, 3, 4), + (16, 32, 64, 16, 32, False, True, False): (1, 2, 1, 4), + (16, 32, 64, 16, 32, False, True, True): (1, 4, 1, 4), + (16, 32, 64, 16, 32, True, False, False): (1, 2, 1, 8), + (16, 32, 64, 16, 32, True, False, True): (1, 2, 1, 4), + (16, 64, 16, 16, 32, False, False, False): (1, 1, 1, 2), + (16, 64, 16, 16, 32, False, False, True): (1, 1, 1, 8), + (16, 64, 16, 16, 32, False, True, False): (1, 1, 1, 8), + (16, 64, 16, 16, 32, False, True, True): (1, 1, 1, 4), + (16, 64, 16, 16, 32, True, False, False): (1, 1, 1, 8), + (16, 64, 16, 16, 32, True, False, True): (1, 1, 1, 4), + (16, 64, 32, 16, 32, False, False, False): (1, 2, 1, 4), + (16, 64, 32, 16, 32, False, False, True): (1, 1, 1, 4), + (16, 64, 32, 16, 32, False, True, False): (1, 2, 1, 1), + (16, 64, 32, 16, 32, False, True, True): (1, 2, 1, 8), + (16, 64, 32, 16, 32, True, False, False): (2, 2, 1, 4), + (16, 64, 32, 16, 32, True, False, True): (2, 2, 1, 4), + (16, 64, 64, 16, 32, False, False, False): (1, 2, 1, 4), + (16, 64, 64, 16, 32, False, False, True): (1, 4, 1, 4), + (16, 64, 64, 16, 32, False, True, False): (1, 4, 1, 4), + (16, 64, 64, 16, 32, False, True, True): (1, 4, 1, 4), + (16, 64, 64, 16, 32, True, False, False): (1, 4, 1, 2), + (16, 64, 64, 16, 32, True, False, True): (3, 4, 1, 4), + (32, 16, 16, 16, 16, False, False, False): (1, 1, 2, 4), + (32, 16, 16, 16, 16, False, False, True): (1, 1, 1, 2), + (32, 16, 16, 16, 16, False, True, False): (1, 1, 2, 4), + (32, 16, 16, 16, 16, False, True, True): (1, 1, 2, 4), + (32, 16, 16, 16, 16, True, False, False): (1, 1, 3, 8), + (32, 16, 16, 16, 16, True, False, True): (1, 1, 2, 4), + (32, 16, 32, 16, 16, False, False, False): (1, 2, 1, 4), + (32, 16, 32, 16, 16, False, False, True): (1, 2, 3, 4), + (32, 16, 32, 16, 16, False, True, False): (1, 1, 1, 8), + (32, 16, 32, 16, 16, False, True, True): (1, 2, 1, 4), + (32, 16, 32, 16, 16, True, False, False): (1, 1, 1, 2), + (32, 16, 32, 16, 16, True, False, True): (1, 1, 1, 4), + (32, 16, 64, 16, 16, False, False, False): (1, 4, 1, 4), + (32, 16, 64, 16, 16, False, False, True): (3, 4, 1, 4), + (32, 16, 64, 16, 16, False, True, False): (1, 4, 1, 1), + (32, 16, 64, 16, 16, False, True, True): (1, 4, 1, 4), + (32, 16, 64, 16, 16, True, False, False): (1, 4, 1, 4), + (32, 16, 64, 16, 16, True, False, True): (1, 4, 1, 4), + (32, 32, 16, 16, 16, False, False, False): (1, 1, 1, 2), + (32, 32, 16, 16, 16, False, False, True): (2, 1, 1, 4), + (32, 32, 16, 16, 16, False, True, False): (1, 1, 1, 2), + (32, 32, 16, 16, 16, False, True, True): (2, 1, 1, 4), + (32, 32, 16, 16, 16, True, False, False): (3, 1, 2, 4), + (32, 32, 16, 16, 16, True, False, True): (1, 1, 2, 4), + (32, 32, 16, 16, 32, False, False, False): (2, 1, 1, 2), + (32, 32, 16, 16, 32, False, False, True): (1, 1, 1, 4), + (32, 32, 16, 16, 32, False, True, False): (1, 1, 1, 4), + (32, 32, 16, 16, 32, False, True, True): (1, 1, 1, 8), + (32, 32, 16, 16, 32, True, False, False): (1, 1, 1, 8), + (32, 32, 16, 16, 32, True, False, True): (1, 1, 1, 4), + (32, 32, 16, 32, 32, False, False, False): (2, 1, 1, 4), + (32, 32, 16, 32, 32, False, False, True): (1, 1, 2, 4), + (32, 32, 16, 32, 32, False, True, False): (2, 1, 1, 1), + (32, 32, 16, 32, 32, False, True, True): (2, 1, 2, 4), + (32, 32, 16, 32, 32, True, False, False): (1, 1, 1, 8), + (32, 32, 16, 32, 32, True, False, True): (1, 1, 1, 4), + (32, 32, 32, 16, 16, False, False, False): (1, 1, 1, 4), + (32, 32, 32, 16, 16, False, False, True): (1, 2, 1, 2), + (32, 32, 32, 16, 16, False, True, False): (2, 2, 1, 4), + (32, 32, 32, 16, 16, False, True, True): (1, 2, 2, 4), + (32, 32, 32, 16, 16, True, False, False): (1, 2, 1, 4), + (32, 32, 32, 16, 16, True, False, True): (2, 2, 1, 4), + (32, 32, 32, 16, 32, False, False, False): (1, 2, 1, 4), + (32, 32, 32, 16, 32, False, False, True): (1, 2, 1, 4), + (32, 32, 32, 16, 32, False, True, False): (1, 2, 1, 4), + (32, 32, 32, 16, 32, False, True, True): (1, 2, 1, 4), + (32, 32, 32, 16, 32, True, False, False): (2, 1, 1, 2), + (32, 32, 32, 16, 32, True, False, True): (2, 2, 2, 4), + (32, 32, 32, 32, 32, False, False, False): (1, 1, 1, 4), + (32, 32, 32, 32, 32, False, False, True): (1, 1, 1, 2), + (32, 32, 32, 32, 32, False, True, False): (1, 1, 1, 4), + (32, 32, 32, 32, 32, False, True, True): (1, 1, 2, 2), + (32, 32, 32, 32, 32, True, False, False): (1, 1, 1, 2), + (32, 32, 32, 32, 32, True, False, True): (1, 1, 2, 1), + (32, 32, 64, 16, 16, False, False, False): (2, 4, 1, 4), + (32, 32, 64, 16, 16, False, False, True): (1, 4, 2, 4), + (32, 32, 64, 16, 16, False, True, False): (1, 4, 1, 4), + (32, 32, 64, 16, 16, False, True, True): (1, 4, 1, 4), + (32, 32, 64, 16, 16, True, False, False): (1, 2, 1, 4), + (32, 32, 64, 16, 16, True, False, True): (2, 4, 1, 4), + (32, 32, 64, 16, 32, False, False, False): (1, 4, 1, 8), + (32, 32, 64, 16, 32, False, False, True): (1, 4, 1, 4), + (32, 32, 64, 16, 32, False, True, False): (1, 4, 1, 4), + (32, 32, 64, 16, 32, False, True, True): (2, 4, 1, 4), + (32, 32, 64, 16, 32, True, False, False): (1, 2, 2, 4), + (32, 32, 64, 16, 32, True, False, True): (2, 4, 1, 4), + (32, 32, 64, 32, 32, False, False, False): (2, 2, 1, 4), + (32, 32, 64, 32, 32, False, False, True): (1, 1, 1, 4), + (32, 32, 64, 32, 32, False, True, False): (1, 1, 1, 8), + (32, 32, 64, 32, 32, False, True, True): (2, 1, 1, 4), + (32, 32, 64, 32, 32, True, False, False): (1, 1, 1, 4), + (32, 32, 64, 32, 32, True, False, True): (1, 2, 1, 1), + (32, 64, 16, 16, 32, False, False, False): (1, 1, 2, 2), + (32, 64, 16, 16, 32, False, False, True): (2, 1, 1, 4), + (32, 64, 16, 16, 32, False, True, False): (1, 1, 1, 8), + (32, 64, 16, 16, 32, False, True, True): (1, 1, 3, 4), + (32, 64, 16, 16, 32, True, False, False): (1, 1, 1, 2), + (32, 64, 16, 16, 32, True, False, True): (1, 1, 2, 4), + (32, 64, 16, 32, 32, False, False, False): (1, 1, 1, 2), + (32, 64, 16, 32, 32, False, False, True): (1, 1, 3, 4), + (32, 64, 16, 32, 32, False, True, False): (1, 1, 2, 4), + (32, 64, 16, 32, 32, False, True, True): (1, 1, 1, 8), + (32, 64, 16, 32, 32, True, False, False): (1, 1, 2, 4), + (32, 64, 16, 32, 32, True, False, True): (1, 1, 1, 8), + (32, 64, 32, 16, 32, False, False, False): (1, 2, 1, 4), + (32, 64, 32, 16, 32, False, False, True): (1, 2, 3, 4), + (32, 64, 32, 16, 32, False, True, False): (1, 2, 1, 8), + (32, 64, 32, 16, 32, False, True, True): (3, 2, 1, 4), + (32, 64, 32, 16, 32, True, False, False): (1, 1, 1, 8), + (32, 64, 32, 16, 32, True, False, True): (1, 2, 1, 4), + (32, 64, 32, 32, 32, False, False, False): (1, 1, 1, 1), + (32, 64, 32, 32, 32, False, False, True): (1, 1, 1, 4), + (32, 64, 32, 32, 32, False, True, False): (1, 1, 1, 4), + (32, 64, 32, 32, 32, False, True, True): (1, 1, 1, 4), + (32, 64, 32, 32, 32, True, False, False): (1, 1, 1, 4), + (32, 64, 32, 32, 32, True, False, True): (1, 1, 2, 8), + (32, 64, 64, 16, 32, False, False, False): (2, 4, 1, 4), + (32, 64, 64, 16, 32, False, False, True): (1, 4, 1, 4), + (32, 64, 64, 16, 32, False, True, False): (1, 4, 1, 4), + (32, 64, 64, 16, 32, False, True, True): (2, 4, 1, 4), + (32, 64, 64, 16, 32, True, False, False): (1, 4, 1, 4), + (32, 64, 64, 16, 32, True, False, True): (1, 4, 1, 4), + (32, 64, 64, 32, 32, False, False, False): (2, 2, 1, 4), + (32, 64, 64, 32, 32, False, False, True): (1, 2, 1, 8), + (32, 64, 64, 32, 32, False, True, False): (1, 2, 1, 4), + (32, 64, 64, 32, 32, False, True, True): (1, 2, 1, 4), + (32, 64, 64, 32, 32, True, False, False): (2, 2, 1, 4), + (32, 64, 64, 32, 32, True, False, True): (1, 2, 3, 8), + (64, 32, 16, 32, 32, False, False, False): (1, 1, 1, 4), + (64, 32, 16, 32, 32, False, False, True): (3, 1, 2, 4), + (64, 32, 16, 32, 32, False, True, False): (2, 1, 1, 2), + (64, 32, 16, 32, 32, False, True, True): (1, 1, 1, 8), + (64, 32, 16, 32, 32, True, False, False): (1, 1, 1, 2), + (64, 32, 16, 32, 32, True, False, True): (1, 1, 1, 4), + (64, 32, 32, 32, 32, False, False, False): (1, 1, 1, 4), + (64, 32, 32, 32, 32, False, False, True): (1, 1, 2, 8), + (64, 32, 32, 32, 32, False, True, False): (1, 1, 1, 8), + (64, 32, 32, 32, 32, False, True, True): (1, 1, 1, 4), + (64, 32, 32, 32, 32, True, False, False): (1, 1, 2, 4), + (64, 32, 32, 32, 32, True, False, True): (1, 1, 3, 8), + (64, 32, 64, 32, 32, False, False, False): (1, 2, 1, 4), + (64, 32, 64, 32, 32, False, False, True): (2, 2, 1, 4), + (64, 32, 64, 32, 32, False, True, False): (1, 1, 1, 4), + (64, 32, 64, 32, 32, False, True, True): (1, 2, 1, 8), + (64, 32, 64, 32, 32, True, False, False): (2, 2, 1, 4), + (64, 32, 64, 32, 32, True, False, True): (1, 2, 1, 8), + (64, 64, 16, 32, 32, False, False, False): (1, 1, 2, 8), + (64, 64, 16, 32, 32, False, False, True): (2, 1, 2, 4), + (64, 64, 16, 32, 32, False, True, False): (1, 1, 1, 2), + (64, 64, 16, 32, 32, False, True, True): (1, 1, 2, 4), + (64, 64, 16, 32, 32, True, False, False): (1, 1, 1, 2), + (64, 64, 16, 32, 32, True, False, True): (1, 1, 2, 4), + (64, 64, 32, 32, 32, False, False, False): (1, 1, 1, 4), + (64, 64, 32, 32, 32, False, False, True): (2, 1, 1, 4), + (64, 64, 32, 32, 32, False, True, False): (1, 1, 1, 8), + (64, 64, 32, 32, 32, False, True, True): (2, 1, 1, 4), + (64, 64, 32, 32, 32, True, False, False): (1, 1, 1, 4), + (64, 64, 32, 32, 32, True, False, True): (1, 1, 1, 8), + (64, 64, 64, 32, 32, False, False, False): (2, 2, 1, 4), + (64, 64, 64, 32, 32, False, False, True): (1, 2, 1, 4), + (64, 64, 64, 32, 32, False, True, False): (1, 2, 1, 4), + (64, 64, 64, 32, 32, False, True, True): (2, 2, 1, 4), + (64, 64, 64, 32, 32, True, False, False): (1, 1, 1, 8), + (64, 64, 64, 32, 32, True, False, True): (1, 2, 2, 4), + (192, 192, 256, 16, 16, False, True, True): (1, 16, 3, 2), + (192, 192, 256, 16, 16, True, False, True): (1, 8, 5, 4), + (192, 192, 256, 32, 32, False, True, True): (2, 8, 4, 4), + (192, 192, 256, 32, 32, True, False, True): (1, 8, 5, 4), + (192, 192, 512, 16, 16, False, True, True): (2, 16, 3, 4), + (192, 192, 512, 16, 16, True, False, True): (1, 16, 5, 4), + (192, 192, 512, 32, 32, False, True, True): (1, 16, 3, 4), + (192, 192, 512, 32, 32, True, False, True): (2, 16, 3, 4), + (192, 192, 1024, 16, 16, False, True, True): (3, 16, 3, 4), + (192, 192, 1024, 16, 16, True, False, True): (2, 8, 3, 4), + (192, 192, 1024, 32, 32, False, True, True): (3, 32, 1, 4), + (192, 192, 1024, 32, 32, True, False, True): (3, 16, 3, 4), + (192, 192, 2048, 16, 16, False, True, True): (1, 32, 3, 4), + (192, 192, 2048, 16, 16, True, False, True): (2, 16, 3, 4), + (192, 192, 2048, 32, 32, False, True, True): (1, 64, 1, 4), + (192, 192, 2048, 32, 32, True, False, True): (1, 64, 2, 4), + (192, 192, 4096, 16, 16, False, True, True): (1, 64, 2, 4), + (192, 192, 4096, 16, 16, True, False, True): (1, 32, 3, 4), + (192, 192, 4096, 32, 32, False, True, True): (3, 128, 2, 4), + (192, 192, 4096, 32, 32, True, False, True): (1, 128, 2, 4), + (192, 192, 8192, 16, 16, False, True, True): (2, 64, 3, 4), + (192, 192, 8192, 16, 16, True, False, True): (1, 64, 3, 4), + (192, 192, 8192, 32, 32, False, True, True): (3, 128, 3, 4), + (192, 192, 8192, 32, 32, True, False, True): (1, 128, 2, 4), + (192, 192, 16384, 16, 16, False, True, True): (1, 256, 3, 2), + (192, 192, 16384, 16, 16, True, False, True): (1, 256, 3, 2), + (192, 192, 16384, 32, 32, False, True, True): (2, 256, 3, 4), + (192, 192, 16384, 32, 32, True, False, True): (2, 256, 3, 4), + (192, 192, 32768, 16, 16, False, True, True): (2, 512, 3, 2), + (192, 192, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (192, 192, 32768, 32, 32, False, True, True): (2, 512, 3, 4), + (192, 192, 32768, 32, 32, True, False, True): (2, 512, 3, 4), + (192, 192, 65536, 16, 16, False, True, True): (2, 1024, 3, 2), + (192, 192, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (192, 192, 65536, 32, 32, False, True, True): (2, 1024, 3, 4), + (192, 192, 65536, 32, 32, True, False, True): (2, 1024, 3, 4), + (192, 192, 131072, 16, 16, False, True, True): (2, 512, 3, 4), + (192, 192, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (192, 192, 131072, 32, 32, False, True, True): (2, 1024, 3, 4), + (192, 192, 131072, 32, 32, True, False, True): (2, 1024, 3, 4), + (256, 256, 256, 16, 16, False, True, True): (1, 16, 3, 4), + (256, 256, 256, 16, 16, True, False, True): (2, 16, 1, 4), + (256, 256, 256, 32, 32, False, True, True): (1, 8, 4, 8), + (256, 256, 256, 32, 32, True, False, True): (4, 8, 4, 4), + (256, 256, 256, 64, 64, False, True, True): (1, 4, 4, 8), + (256, 256, 256, 64, 64, True, False, True): (1, 4, 3, 8), + (256, 256, 256, 128, 128, False, True, True): (7, 2, 1, 32), + (256, 256, 256, 128, 128, True, False, True): (3, 2, 1, 32), + (256, 256, 512, 16, 16, False, True, True): (1, 16, 5, 4), + (256, 256, 512, 16, 16, True, False, True): (1, 16, 3, 2), + (256, 256, 512, 32, 32, False, True, True): (4, 16, 4, 4), + (256, 256, 512, 32, 32, True, False, True): (4, 16, 3, 4), + (256, 256, 512, 64, 64, False, True, True): (1, 8, 3, 8), + (256, 256, 512, 64, 64, True, False, True): (1, 8, 3, 8), + (256, 256, 512, 128, 128, False, True, True): (1, 4, 1, 32), + (256, 256, 512, 128, 128, True, False, True): (3, 4, 1, 32), + (256, 256, 1024, 16, 16, False, True, True): (3, 32, 5, 2), + (256, 256, 1024, 16, 16, True, False, True): (2, 32, 5, 2), + (256, 256, 1024, 32, 32, False, True, True): (1, 32, 4, 4), + (256, 256, 1024, 32, 32, True, False, True): (1, 32, 5, 4), + (256, 256, 1024, 64, 64, False, True, True): (4, 16, 3, 8), + (256, 256, 1024, 64, 64, True, False, True): (1, 16, 3, 8), + (256, 256, 1024, 128, 128, False, True, True): (1, 8, 1, 32), + (256, 256, 1024, 128, 128, True, False, True): (3, 8, 1, 32), + (256, 256, 2048, 16, 16, False, True, True): (3, 32, 3, 4), + (256, 256, 2048, 16, 16, True, False, True): (1, 64, 3, 2), + (256, 256, 2048, 32, 32, False, True, True): (1, 64, 3, 4), + (256, 256, 2048, 32, 32, True, False, True): (1, 64, 3, 4), + (256, 256, 2048, 64, 64, False, True, True): (2, 32, 1, 8), + (256, 256, 2048, 64, 64, True, False, True): (2, 32, 1, 8), + (256, 256, 2048, 128, 128, False, True, True): (4, 16, 1, 32), + (256, 256, 2048, 128, 128, True, False, True): (4, 16, 1, 32), + (256, 256, 4096, 16, 16, False, True, True): (1, 32, 2, 4), + (256, 256, 4096, 16, 16, True, False, True): (1, 32, 3, 4), + (256, 256, 4096, 32, 32, False, True, True): (1, 128, 2, 4), + (256, 256, 4096, 32, 32, True, False, True): (1, 128, 2, 4), + (256, 256, 4096, 64, 64, False, True, True): (2, 64, 4, 8), + (256, 256, 4096, 64, 64, True, False, True): (3, 64, 2, 8), + (256, 256, 4096, 128, 128, False, True, True): (3, 32, 1, 32), + (256, 256, 4096, 128, 128, True, False, True): (2, 32, 1, 32), + (256, 256, 8192, 16, 16, False, True, True): (1, 64, 3, 4), + (256, 256, 8192, 16, 16, True, False, True): (2, 128, 3, 2), + (256, 256, 8192, 32, 32, False, True, True): (3, 128, 3, 4), + (256, 256, 8192, 32, 32, True, False, True): (1, 128, 3, 4), + (256, 256, 8192, 64, 64, False, True, True): (3, 128, 1, 4), + (256, 256, 8192, 64, 64, True, False, True): (4, 128, 2, 8), + (256, 256, 8192, 128, 128, False, True, True): (6, 64, 1, 32), + (256, 256, 8192, 128, 128, True, False, True): (2, 64, 1, 32), + (256, 256, 16384, 16, 16, False, True, True): (4, 128, 3, 4), + (256, 256, 16384, 16, 16, True, False, True): (3, 128, 3, 4), + (256, 256, 16384, 32, 32, False, True, True): (4, 256, 3, 4), + (256, 256, 16384, 32, 32, True, False, True): (2, 256, 3, 4), + (256, 256, 16384, 64, 64, False, True, True): (3, 256, 1, 4), + (256, 256, 16384, 64, 64, True, False, True): (2, 256, 2, 4), + (256, 256, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (256, 256, 16384, 128, 128, True, False, True): (3, 128, 1, 32), + (256, 256, 32768, 16, 16, False, True, True): (1, 256, 3, 4), + (256, 256, 32768, 16, 16, True, False, True): (2, 128, 3, 4), + (256, 256, 32768, 32, 32, False, True, True): (2, 512, 3, 4), + (256, 256, 32768, 32, 32, True, False, True): (4, 512, 3, 4), + (256, 256, 32768, 64, 64, False, True, True): (1, 512, 1, 8), + (256, 256, 32768, 64, 64, True, False, True): (1, 512, 2, 4), + (256, 256, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (256, 256, 32768, 128, 128, True, False, True): (1, 256, 1, 32), + (256, 256, 65536, 16, 16, False, True, True): (2, 512, 3, 4), + (256, 256, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (256, 256, 65536, 32, 32, False, True, True): (1, 1024, 3, 4), + (256, 256, 65536, 32, 32, True, False, True): (2, 1024, 3, 4), + (256, 256, 65536, 64, 64, False, True, True): (1, 1024, 2, 4), + (256, 256, 65536, 64, 64, True, False, True): (1, 1024, 2, 4), + (256, 256, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (256, 256, 65536, 128, 128, True, False, True): (2, 512, 1, 32), + (256, 256, 131072, 16, 16, False, True, True): (1, 1024, 3, 4), + (256, 256, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (256, 256, 131072, 32, 32, False, True, True): (1, 2048, 3, 4), + (256, 256, 131072, 32, 32, True, False, True): (1, 2048, 3, 4), + (256, 256, 131072, 64, 64, False, True, True): (1, 2048, 1, 8), + (256, 256, 131072, 64, 64, True, False, True): (1, 2048, 2, 4), + (256, 256, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (256, 256, 131072, 128, 128, True, False, True): (4, 1024, 1, 32), + (384, 384, 256, 16, 16, False, True, True): (1, 8, 3, 4), + (384, 384, 256, 16, 16, True, False, True): (1, 8, 3, 4), + (384, 384, 256, 32, 32, False, True, True): (2, 8, 3, 8), + (384, 384, 256, 32, 32, True, False, True): (1, 8, 3, 4), + (384, 384, 256, 64, 64, False, True, True): (1, 4, 4, 8), + (384, 384, 256, 64, 64, True, False, True): (2, 4, 3, 8), + (384, 384, 512, 16, 16, False, True, True): (3, 16, 3, 2), + (384, 384, 512, 16, 16, True, False, True): (3, 16, 3, 2), + (384, 384, 512, 32, 32, False, True, True): (2, 8, 3, 4), + (384, 384, 512, 32, 32, True, False, True): (1, 8, 3, 4), + (384, 384, 512, 64, 64, False, True, True): (2, 8, 3, 8), + (384, 384, 512, 64, 64, True, False, True): (2, 8, 4, 8), + (384, 384, 1024, 16, 16, False, True, True): (3, 16, 3, 2), + (384, 384, 1024, 16, 16, True, False, True): (4, 32, 3, 2), + (384, 384, 1024, 32, 32, False, True, True): (1, 32, 3, 4), + (384, 384, 1024, 32, 32, True, False, True): (2, 16, 3, 4), + (384, 384, 1024, 64, 64, False, True, True): (2, 16, 3, 8), + (384, 384, 1024, 64, 64, True, False, True): (4, 16, 4, 8), + (384, 384, 2048, 16, 16, False, True, True): (3, 16, 3, 4), + (384, 384, 2048, 16, 16, True, False, True): (1, 32, 3, 4), + (384, 384, 2048, 32, 32, False, True, True): (3, 64, 2, 4), + (384, 384, 2048, 32, 32, True, False, True): (1, 64, 3, 4), + (384, 384, 2048, 64, 64, False, True, True): (4, 32, 4, 8), + (384, 384, 2048, 64, 64, True, False, True): (5, 32, 4, 8), + (384, 384, 4096, 16, 16, False, True, True): (1, 32, 3, 4), + (384, 384, 4096, 16, 16, True, False, True): (3, 32, 3, 4), + (384, 384, 4096, 32, 32, False, True, True): (2, 64, 3, 4), + (384, 384, 4096, 32, 32, True, False, True): (2, 64, 3, 4), + (384, 384, 4096, 64, 64, False, True, True): (2, 64, 3, 8), + (384, 384, 4096, 64, 64, True, False, True): (2, 64, 3, 8), + (384, 384, 8192, 16, 16, False, True, True): (1, 128, 3, 2), + (384, 384, 8192, 16, 16, True, False, True): (1, 128, 3, 2), + (384, 384, 8192, 32, 32, False, True, True): (1, 128, 3, 4), + (384, 384, 8192, 32, 32, True, False, True): (1, 128, 3, 4), + (384, 384, 8192, 64, 64, False, True, True): (3, 128, 3, 4), + (384, 384, 8192, 64, 64, True, False, True): (2, 128, 3, 4), + (384, 384, 16384, 16, 16, False, True, True): (1, 256, 3, 2), + (384, 384, 16384, 16, 16, True, False, True): (1, 64, 3, 4), + (384, 384, 16384, 32, 32, False, True, True): (2, 256, 3, 4), + (384, 384, 16384, 32, 32, True, False, True): (4, 256, 3, 4), + (384, 384, 16384, 64, 64, False, True, True): (2, 256, 3, 4), + (384, 384, 16384, 64, 64, True, False, True): (1, 256, 3, 4), + (384, 384, 32768, 16, 16, False, True, True): (1, 128, 3, 4), + (384, 384, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (384, 384, 32768, 32, 32, False, True, True): (1, 512, 3, 4), + (384, 384, 32768, 32, 32, True, False, True): (1, 512, 2, 4), + (384, 384, 32768, 64, 64, False, True, True): (1, 512, 3, 4), + (384, 384, 32768, 64, 64, True, False, True): (1, 512, 3, 4), + (384, 384, 65536, 16, 16, False, True, True): (1, 256, 3, 4), + (384, 384, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (384, 384, 65536, 32, 32, False, True, True): (1, 1024, 3, 4), + (384, 384, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (384, 384, 65536, 64, 64, False, True, True): (1, 1024, 3, 4), + (384, 384, 65536, 64, 64, True, False, True): (1, 1024, 3, 4), + (384, 384, 131072, 16, 16, False, True, True): (1, 512, 3, 4), + (384, 384, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (384, 384, 131072, 32, 32, False, True, True): (1, 1024, 3, 4), + (384, 384, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (384, 384, 131072, 64, 64, False, True, True): (1, 2048, 3, 4), + (384, 384, 131072, 64, 64, True, False, True): (1, 2048, 3, 4), + (512, 512, 256, 16, 16, False, True, True): (1, 8, 4, 4), + (512, 512, 256, 16, 16, True, False, True): (1, 8, 3, 2), + (512, 512, 256, 32, 32, False, True, True): (4, 8, 3, 4), + (512, 512, 256, 32, 32, True, False, True): (4, 8, 3, 4), + (512, 512, 256, 64, 64, False, True, True): (3, 4, 3, 8), + (512, 512, 256, 64, 64, True, False, True): (5, 4, 3, 8), + (512, 512, 256, 128, 128, False, True, True): (1, 2, 1, 32), + (512, 512, 256, 128, 128, True, False, True): (3, 2, 1, 32), + (512, 512, 512, 16, 16, False, True, True): (2, 16, 3, 2), + (512, 512, 512, 16, 16, True, False, True): (1, 8, 4, 4), + (512, 512, 512, 32, 32, False, True, True): (3, 16, 3, 4), + (512, 512, 512, 32, 32, True, False, True): (5, 16, 2, 4), + (512, 512, 512, 64, 64, False, True, True): (1, 8, 3, 8), + (512, 512, 512, 64, 64, True, False, True): (3, 8, 3, 8), + (512, 512, 512, 128, 128, False, True, True): (1, 4, 1, 32), + (512, 512, 512, 128, 128, True, False, True): (3, 4, 1, 16), + (512, 512, 1024, 16, 16, False, True, True): (1, 16, 3, 4), + (512, 512, 1024, 16, 16, True, False, True): (3, 16, 3, 4), + (512, 512, 1024, 32, 32, False, True, True): (3, 32, 3, 4), + (512, 512, 1024, 32, 32, True, False, True): (3, 32, 2, 4), + (512, 512, 1024, 64, 64, False, True, True): (1, 16, 3, 8), + (512, 512, 1024, 64, 64, True, False, True): (4, 16, 3, 8), + (512, 512, 1024, 128, 128, False, True, True): (4, 8, 1, 32), + (512, 512, 1024, 128, 128, True, False, True): (4, 8, 1, 32), + (512, 512, 2048, 16, 16, False, True, True): (5, 16, 3, 4), + (512, 512, 2048, 16, 16, True, False, True): (5, 16, 3, 4), + (512, 512, 2048, 32, 32, False, True, True): (1, 32, 3, 4), + (512, 512, 2048, 32, 32, True, False, True): (1, 32, 4, 4), + (512, 512, 2048, 64, 64, False, True, True): (4, 32, 3, 8), + (512, 512, 2048, 64, 64, True, False, True): (4, 32, 3, 8), + (512, 512, 2048, 128, 128, False, True, True): (3, 16, 1, 32), + (512, 512, 2048, 128, 128, True, False, True): (3, 16, 1, 32), + (512, 512, 4096, 16, 16, False, True, True): (4, 32, 3, 4), + (512, 512, 4096, 16, 16, True, False, True): (4, 64, 3, 2), + (512, 512, 4096, 32, 32, False, True, True): (3, 64, 3, 4), + (512, 512, 4096, 32, 32, True, False, True): (3, 64, 3, 4), + (512, 512, 4096, 64, 64, False, True, True): (4, 64, 2, 4), + (512, 512, 4096, 64, 64, True, False, True): (1, 64, 2, 4), + (512, 512, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (512, 512, 4096, 128, 128, True, False, True): (1, 32, 1, 32), + (512, 512, 8192, 16, 16, False, True, True): (1, 64, 3, 4), + (512, 512, 8192, 16, 16, True, False, True): (4, 64, 3, 4), + (512, 512, 8192, 32, 32, False, True, True): (2, 128, 3, 4), + (512, 512, 8192, 32, 32, True, False, True): (3, 128, 3, 4), + (512, 512, 8192, 64, 64, False, True, True): (1, 128, 2, 4), + (512, 512, 8192, 64, 64, True, False, True): (1, 128, 2, 4), + (512, 512, 8192, 128, 128, False, True, True): (6, 64, 1, 32), + (512, 512, 8192, 128, 128, True, False, True): (4, 64, 1, 32), + (512, 512, 16384, 16, 16, False, True, True): (1, 128, 3, 4), + (512, 512, 16384, 16, 16, True, False, True): (1, 64, 3, 4), + (512, 512, 16384, 32, 32, False, True, True): (1, 256, 3, 4), + (512, 512, 16384, 32, 32, True, False, True): (4, 256, 3, 4), + (512, 512, 16384, 64, 64, False, True, True): (1, 256, 2, 4), + (512, 512, 16384, 64, 64, True, False, True): (1, 256, 2, 4), + (512, 512, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (512, 512, 16384, 128, 128, True, False, True): (2, 128, 1, 32), + (512, 512, 32768, 16, 16, False, True, True): (1, 256, 3, 4), + (512, 512, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (512, 512, 32768, 32, 32, False, True, True): (1, 512, 3, 4), + (512, 512, 32768, 32, 32, True, False, True): (1, 512, 3, 4), + (512, 512, 32768, 64, 64, False, True, True): (1, 512, 2, 4), + (512, 512, 32768, 64, 64, True, False, True): (2, 512, 2, 4), + (512, 512, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (512, 512, 32768, 128, 128, True, False, True): (2, 256, 1, 32), + (512, 512, 65536, 16, 16, False, True, True): (1, 512, 3, 4), + (512, 512, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (512, 512, 65536, 32, 32, False, True, True): (1, 1024, 3, 4), + (512, 512, 65536, 32, 32, True, False, True): (1, 1024, 3, 4), + (512, 512, 65536, 64, 64, False, True, True): (1, 1024, 2, 4), + (512, 512, 65536, 64, 64, True, False, True): (1, 1024, 2, 4), + (512, 512, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (512, 512, 65536, 128, 128, True, False, True): (4, 512, 1, 32), + (512, 512, 131072, 16, 16, False, True, True): (1, 512, 3, 4), + (512, 512, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (512, 512, 131072, 32, 32, False, True, True): (1, 2048, 3, 4), + (512, 512, 131072, 32, 32, True, False, True): (1, 2048, 3, 4), + (512, 512, 131072, 64, 64, False, True, True): (1, 2048, 2, 4), + (512, 512, 131072, 64, 64, True, False, True): (1, 2048, 2, 4), + (512, 512, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (512, 512, 131072, 128, 128, True, False, True): (2, 1024, 1, 32), + (768, 768, 256, 16, 16, False, True, True): (1, 4, 5, 4), + (768, 768, 256, 16, 16, True, False, True): (3, 8, 3, 2), + (768, 768, 256, 32, 32, False, True, True): (2, 4, 3, 4), + (768, 768, 256, 32, 32, True, False, True): (3, 8, 4, 4), + (768, 768, 256, 64, 64, False, True, True): (1, 4, 4, 8), + (768, 768, 256, 64, 64, True, False, True): (3, 4, 3, 8), + (768, 768, 256, 128, 128, False, True, True): (3, 2, 1, 32), + (768, 768, 256, 128, 128, True, False, True): (2, 2, 2, 32), + (768, 768, 512, 16, 16, False, True, True): (2, 4, 5, 4), + (768, 768, 512, 16, 16, True, False, True): (2, 4, 4, 4), + (768, 768, 512, 32, 32, False, True, True): (1, 8, 3, 4), + (768, 768, 512, 32, 32, True, False, True): (3, 8, 4, 4), + (768, 768, 512, 64, 64, False, True, True): (2, 8, 3, 8), + (768, 768, 512, 64, 64, True, False, True): (5, 8, 3, 8), + (768, 768, 512, 128, 128, False, True, True): (2, 4, 1, 32), + (768, 768, 512, 128, 128, True, False, True): (2, 4, 2, 32), + (768, 768, 1024, 16, 16, False, True, True): (2, 16, 4, 2), + (768, 768, 1024, 16, 16, True, False, True): (4, 32, 3, 1), + (768, 768, 1024, 32, 32, False, True, True): (1, 32, 2, 4), + (768, 768, 1024, 32, 32, True, False, True): (1, 16, 5, 4), + (768, 768, 1024, 64, 64, False, True, True): (2, 16, 3, 8), + (768, 768, 1024, 64, 64, True, False, True): (2, 16, 3, 8), + (768, 768, 1024, 128, 128, False, True, True): (1, 8, 2, 32), + (768, 768, 1024, 128, 128, True, False, True): (1, 8, 1, 32), + (768, 768, 2048, 16, 16, False, True, True): (1, 16, 3, 4), + (768, 768, 2048, 16, 16, True, False, True): (1, 16, 3, 4), + (768, 768, 2048, 32, 32, False, True, True): (1, 32, 3, 4), + (768, 768, 2048, 32, 32, True, False, True): (5, 32, 3, 4), + (768, 768, 2048, 64, 64, False, True, True): (1, 32, 3, 8), + (768, 768, 2048, 64, 64, True, False, True): (1, 32, 3, 4), + (768, 768, 2048, 128, 128, False, True, True): (3, 16, 1, 32), + (768, 768, 2048, 128, 128, True, False, True): (4, 16, 1, 32), + (768, 768, 4096, 16, 16, False, True, True): (1, 64, 3, 2), + (768, 768, 4096, 16, 16, True, False, True): (3, 64, 3, 2), + (768, 768, 4096, 32, 32, False, True, True): (1, 64, 3, 4), + (768, 768, 4096, 32, 32, True, False, True): (1, 64, 3, 4), + (768, 768, 4096, 64, 64, False, True, True): (4, 64, 3, 4), + (768, 768, 4096, 64, 64, True, False, True): (4, 64, 3, 4), + (768, 768, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (768, 768, 4096, 128, 128, True, False, True): (1, 32, 2, 32), + (768, 768, 8192, 16, 16, False, True, True): (1, 128, 3, 2), + (768, 768, 8192, 16, 16, True, False, True): (2, 32, 3, 4), + (768, 768, 8192, 32, 32, False, True, True): (2, 128, 3, 4), + (768, 768, 8192, 32, 32, True, False, True): (1, 128, 2, 4), + (768, 768, 8192, 64, 64, False, True, True): (1, 128, 3, 4), + (768, 768, 8192, 64, 64, True, False, True): (2, 128, 3, 4), + (768, 768, 8192, 128, 128, False, True, True): (1, 64, 1, 32), + (768, 768, 8192, 128, 128, True, False, True): (2, 64, 1, 32), + (768, 768, 16384, 16, 16, False, True, True): (3, 64, 3, 4), + (768, 768, 16384, 16, 16, True, False, True): (1, 64, 3, 4), + (768, 768, 16384, 32, 32, False, True, True): (2, 256, 3, 4), + (768, 768, 16384, 32, 32, True, False, True): (4, 256, 2, 4), + (768, 768, 16384, 64, 64, False, True, True): (1, 256, 3, 4), + (768, 768, 16384, 64, 64, True, False, True): (1, 256, 3, 4), + (768, 768, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (768, 768, 16384, 128, 128, True, False, True): (2, 128, 1, 32), + (768, 768, 32768, 16, 16, False, True, True): (1, 128, 3, 4), + (768, 768, 32768, 16, 16, True, False, True): (2, 128, 3, 4), + (768, 768, 32768, 32, 32, False, True, True): (2, 256, 3, 4), + (768, 768, 32768, 32, 32, True, False, True): (1, 256, 3, 4), + (768, 768, 32768, 64, 64, False, True, True): (1, 512, 3, 4), + (768, 768, 32768, 64, 64, True, False, True): (1, 512, 3, 4), + (768, 768, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (768, 768, 32768, 128, 128, True, False, True): (1, 256, 1, 32), + (768, 768, 50432, 16, 16, False, True, True): (1, 197, 3, 4), + (768, 768, 50432, 32, 32, False, True, True): (1, 394, 3, 4), + (768, 768, 50432, 64, 64, False, True, True): (1, 788, 3, 4), + (768, 768, 50432, 128, 128, False, True, True): (3, 394, 1, 32), + (768, 768, 65536, 16, 16, False, True, True): (1, 256, 3, 4), + (768, 768, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (768, 768, 65536, 32, 32, False, True, True): (1, 512, 3, 4), + (768, 768, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (768, 768, 65536, 64, 64, False, True, True): (1, 1024, 3, 4), + (768, 768, 65536, 64, 64, True, False, True): (1, 1024, 3, 4), + (768, 768, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (768, 768, 65536, 128, 128, True, False, True): (1, 512, 1, 32), + (768, 768, 131072, 16, 16, False, True, True): (1, 512, 3, 4), + (768, 768, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (768, 768, 131072, 32, 32, False, True, True): (1, 1024, 3, 4), + (768, 768, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (768, 768, 131072, 64, 64, False, True, True): (1, 2048, 3, 4), + (768, 768, 131072, 64, 64, True, False, True): (1, 2048, 3, 4), + (768, 768, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (768, 768, 131072, 128, 128, True, False, True): (1, 1024, 1, 32), + (768, 3072, 256, 16, 16, False, True, True): (1, 2, 4, 4), + (768, 3072, 256, 16, 16, True, False, True): (1, 4, 3, 4), + (768, 3072, 256, 32, 32, False, True, True): (1, 4, 3, 4), + (768, 3072, 256, 32, 32, True, False, True): (3, 4, 3, 4), + (768, 3072, 256, 64, 64, False, True, True): (1, 4, 3, 8), + (768, 3072, 256, 64, 64, True, False, True): (1, 4, 3, 8), + (768, 3072, 256, 128, 128, False, True, True): (2, 2, 2, 32), + (768, 3072, 256, 128, 128, True, False, True): (2, 2, 1, 32), + (768, 3072, 512, 16, 16, False, True, True): (2, 4, 3, 4), + (768, 3072, 512, 16, 16, True, False, True): (1, 8, 3, 2), + (768, 3072, 512, 32, 32, False, True, True): (3, 8, 4, 4), + (768, 3072, 512, 32, 32, True, False, True): (3, 8, 3, 4), + (768, 3072, 512, 64, 64, False, True, True): (1, 8, 4, 8), + (768, 3072, 512, 64, 64, True, False, True): (1, 8, 3, 8), + (768, 3072, 512, 128, 128, False, True, True): (1, 4, 2, 32), + (768, 3072, 512, 128, 128, True, False, True): (1, 4, 1, 32), + (768, 3072, 1024, 16, 16, False, True, True): (4, 16, 3, 2), + (768, 3072, 1024, 16, 16, True, False, True): (4, 16, 3, 2), + (768, 3072, 1024, 32, 32, False, True, True): (4, 16, 5, 4), + (768, 3072, 1024, 32, 32, True, False, True): (4, 16, 5, 4), + (768, 3072, 1024, 64, 64, False, True, True): (2, 16, 3, 8), + (768, 3072, 1024, 64, 64, True, False, True): (2, 16, 3, 8), + (768, 3072, 1024, 128, 128, False, True, True): (1, 8, 1, 32), + (768, 3072, 1024, 128, 128, True, False, True): (1, 8, 1, 32), + (768, 3072, 2048, 16, 16, False, True, True): (2, 16, 3, 4), + (768, 3072, 2048, 16, 16, True, False, True): (2, 16, 3, 4), + (768, 3072, 2048, 32, 32, False, True, True): (4, 32, 5, 4), + (768, 3072, 2048, 32, 32, True, False, True): (2, 32, 3, 4), + (768, 3072, 2048, 64, 64, False, True, True): (2, 32, 3, 8), + (768, 3072, 2048, 64, 64, True, False, True): (2, 32, 3, 8), + (768, 3072, 2048, 128, 128, False, True, True): (1, 16, 1, 32), + (768, 3072, 2048, 128, 128, True, False, True): (2, 16, 1, 32), + (768, 3072, 4096, 16, 16, False, True, True): (1, 32, 5, 4), + (768, 3072, 4096, 16, 16, True, False, True): (3, 64, 3, 2), + (768, 3072, 4096, 32, 32, False, True, True): (5, 64, 3, 4), + (768, 3072, 4096, 32, 32, True, False, True): (5, 64, 3, 4), + (768, 3072, 4096, 64, 64, False, True, True): (1, 64, 3, 8), + (768, 3072, 4096, 64, 64, True, False, True): (5, 64, 3, 4), + (768, 3072, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (768, 3072, 4096, 128, 128, True, False, True): (1, 32, 1, 32), + (768, 3072, 8192, 16, 16, False, True, True): (1, 128, 3, 2), + (768, 3072, 8192, 16, 16, True, False, True): (1, 128, 3, 2), + (768, 3072, 8192, 32, 32, False, True, True): (1, 128, 3, 4), + (768, 3072, 8192, 32, 32, True, False, True): (1, 64, 3, 4), + (768, 3072, 8192, 64, 64, False, True, True): (3, 128, 3, 4), + (768, 3072, 8192, 64, 64, True, False, True): (3, 128, 3, 4), + (768, 3072, 8192, 128, 128, False, True, True): (4, 64, 2, 32), + (768, 3072, 8192, 128, 128, True, False, True): (2, 64, 1, 32), + (768, 3072, 16384, 16, 16, False, True, True): (1, 256, 2, 2), + (768, 3072, 16384, 16, 16, True, False, True): (1, 64, 3, 4), + (768, 3072, 16384, 32, 32, False, True, True): (8, 128, 3, 4), + (768, 3072, 16384, 32, 32, True, False, True): (1, 128, 3, 4), + (768, 3072, 16384, 64, 64, False, True, True): (1, 256, 3, 4), + (768, 3072, 16384, 64, 64, True, False, True): (3, 256, 3, 4), + (768, 3072, 16384, 128, 128, False, True, True): (3, 128, 1, 32), + (768, 3072, 16384, 128, 128, True, False, True): (2, 128, 2, 32), + (768, 3072, 32768, 16, 16, False, True, True): (1, 512, 3, 1), + (768, 3072, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (768, 3072, 32768, 32, 32, False, True, True): (1, 256, 3, 4), + (768, 3072, 32768, 32, 32, True, False, True): (1, 256, 3, 4), + (768, 3072, 32768, 64, 64, False, True, True): (2, 512, 3, 4), + (768, 3072, 32768, 64, 64, True, False, True): (1, 512, 3, 4), + (768, 3072, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (768, 3072, 32768, 128, 128, True, False, True): (2, 256, 2, 32), + (768, 3072, 50432, 16, 16, False, True, True): (1, 197, 3, 4), + (768, 3072, 50432, 16, 16, True, False, True): (1, 197, 3, 4), + (768, 3072, 50432, 32, 32, False, True, True): (1, 788, 2, 4), + (768, 3072, 50432, 32, 32, True, False, True): (1, 394, 3, 4), + (768, 3072, 50432, 64, 64, False, True, True): (1, 788, 3, 4), + (768, 3072, 50432, 64, 64, True, False, True): (2, 788, 3, 4), + (768, 3072, 50432, 128, 128, False, True, True): (1, 394, 1, 32), + (768, 3072, 50432, 128, 128, True, False, True): (2, 394, 2, 32), + (768, 3072, 65536, 16, 16, False, True, True): (1, 1024, 3, 1), + (768, 3072, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (768, 3072, 65536, 32, 32, False, True, True): (1, 512, 3, 4), + (768, 3072, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (768, 3072, 65536, 64, 64, False, True, True): (2, 1024, 3, 4), + (768, 3072, 65536, 64, 64, True, False, True): (5, 1024, 3, 4), + (768, 3072, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (768, 3072, 65536, 128, 128, True, False, True): (2, 512, 2, 32), + (768, 3072, 131072, 16, 16, False, True, True): (1, 2048, 3, 1), + (768, 3072, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (768, 3072, 131072, 32, 32, False, True, True): (1, 1024, 3, 4), + (768, 3072, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (768, 3072, 131072, 64, 64, False, True, True): (1, 2048, 3, 4), + (768, 3072, 131072, 64, 64, True, False, True): (2, 2048, 3, 4), + (768, 3072, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (768, 3072, 131072, 128, 128, True, False, True): (1, 1024, 2, 32), + (1024, 1024, 256, 16, 16, False, True, True): (4, 8, 3, 2), + (1024, 1024, 256, 16, 16, True, False, True): (2, 8, 3, 2), + (1024, 1024, 256, 32, 32, False, True, True): (1, 8, 3, 4), + (1024, 1024, 256, 32, 32, True, False, True): (1, 8, 3, 4), + (1024, 1024, 256, 64, 64, False, True, True): (1, 4, 3, 8), + (1024, 1024, 256, 64, 64, True, False, True): (2, 4, 3, 8), + (1024, 1024, 256, 128, 128, False, True, True): (3, 2, 1, 32), + (1024, 1024, 256, 128, 128, True, False, True): (5, 2, 1, 32), + (1024, 1024, 512, 16, 16, False, True, True): (3, 8, 3, 4), + (1024, 1024, 512, 16, 16, True, False, True): (3, 8, 3, 4), + (1024, 1024, 512, 32, 32, False, True, True): (1, 16, 3, 4), + (1024, 1024, 512, 32, 32, True, False, True): (3, 16, 3, 4), + (1024, 1024, 512, 64, 64, False, True, True): (6, 8, 3, 8), + (1024, 1024, 512, 64, 64, True, False, True): (8, 8, 3, 8), + (1024, 1024, 512, 128, 128, False, True, True): (1, 4, 1, 32), + (1024, 1024, 512, 128, 128, True, False, True): (1, 4, 1, 32), + (1024, 1024, 1024, 16, 16, False, True, True): (4, 8, 3, 4), + (1024, 1024, 1024, 16, 16, True, False, True): (1, 8, 3, 4), + (1024, 1024, 1024, 32, 32, False, True, True): (4, 16, 4, 4), + (1024, 1024, 1024, 32, 32, True, False, True): (5, 16, 3, 4), + (1024, 1024, 1024, 64, 64, False, True, True): (6, 16, 3, 8), + (1024, 1024, 1024, 64, 64, True, False, True): (3, 16, 2, 4), + (1024, 1024, 1024, 128, 128, False, True, True): (1, 8, 1, 32), + (1024, 1024, 1024, 128, 128, True, False, True): (2, 8, 1, 32), + (1024, 1024, 2048, 16, 16, False, True, True): (4, 16, 3, 4), + (1024, 1024, 2048, 16, 16, True, False, True): (1, 16, 3, 4), + (1024, 1024, 2048, 32, 32, False, True, True): (1, 32, 3, 4), + (1024, 1024, 2048, 32, 32, True, False, True): (2, 32, 3, 4), + (1024, 1024, 2048, 64, 64, False, True, True): (4, 32, 2, 4), + (1024, 1024, 2048, 64, 64, True, False, True): (8, 32, 2, 4), + (1024, 1024, 2048, 128, 128, False, True, True): (1, 16, 1, 32), + (1024, 1024, 2048, 128, 128, True, False, True): (1, 16, 1, 32), + (1024, 1024, 4096, 16, 16, False, True, True): (4, 32, 3, 4), + (1024, 1024, 4096, 16, 16, True, False, True): (1, 64, 3, 2), + (1024, 1024, 4096, 32, 32, False, True, True): (1, 64, 3, 4), + (1024, 1024, 4096, 32, 32, True, False, True): (1, 64, 3, 4), + (1024, 1024, 4096, 64, 64, False, True, True): (2, 64, 2, 4), + (1024, 1024, 4096, 64, 64, True, False, True): (2, 64, 2, 4), + (1024, 1024, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (1024, 1024, 4096, 128, 128, True, False, True): (4, 32, 1, 32), + (1024, 1024, 8192, 16, 16, False, True, True): (1, 128, 3, 1), + (1024, 1024, 8192, 16, 16, True, False, True): (1, 128, 3, 1), + (1024, 1024, 8192, 32, 32, False, True, True): (1, 128, 3, 4), + (1024, 1024, 8192, 32, 32, True, False, True): (1, 128, 3, 4), + (1024, 1024, 8192, 64, 64, False, True, True): (2, 128, 2, 4), + (1024, 1024, 8192, 64, 64, True, False, True): (2, 128, 2, 4), + (1024, 1024, 8192, 128, 128, False, True, True): (1, 64, 1, 32), + (1024, 1024, 8192, 128, 128, True, False, True): (4, 64, 1, 32), + (1024, 1024, 16384, 16, 16, False, True, True): (1, 128, 2, 4), + (1024, 1024, 16384, 16, 16, True, False, True): (4, 256, 3, 1), + (1024, 1024, 16384, 32, 32, False, True, True): (1, 256, 3, 4), + (1024, 1024, 16384, 32, 32, True, False, True): (1, 256, 3, 4), + (1024, 1024, 16384, 64, 64, False, True, True): (1, 256, 2, 4), + (1024, 1024, 16384, 64, 64, True, False, True): (1, 256, 2, 4), + (1024, 1024, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (1024, 1024, 16384, 128, 128, True, False, True): (4, 128, 1, 32), + (1024, 1024, 32768, 16, 16, False, True, True): (1, 256, 2, 4), + (1024, 1024, 32768, 16, 16, True, False, True): (4, 512, 3, 1), + (1024, 1024, 32768, 32, 32, False, True, True): (1, 512, 3, 4), + (1024, 1024, 32768, 32, 32, True, False, True): (1, 512, 3, 4), + (1024, 1024, 32768, 64, 64, False, True, True): (1, 512, 2, 4), + (1024, 1024, 32768, 64, 64, True, False, True): (1, 512, 2, 4), + (1024, 1024, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (1024, 1024, 32768, 128, 128, True, False, True): (1, 256, 1, 32), + (1024, 1024, 65536, 16, 16, False, True, True): (1, 512, 2, 4), + (1024, 1024, 65536, 16, 16, True, False, True): (1, 1024, 3, 1), + (1024, 1024, 65536, 32, 32, False, True, True): (1, 1024, 3, 4), + (1024, 1024, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (1024, 1024, 65536, 64, 64, False, True, True): (1, 1024, 2, 4), + (1024, 1024, 65536, 64, 64, True, False, True): (1, 1024, 2, 4), + (1024, 1024, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (1024, 1024, 65536, 128, 128, True, False, True): (1, 512, 1, 32), + (1024, 1024, 131072, 16, 16, False, True, True): (4, 2048, 3, 1), + (1024, 1024, 131072, 16, 16, True, False, True): (4, 2048, 3, 1), + (1024, 1024, 131072, 32, 32, False, True, True): (1, 2048, 3, 4), + (1024, 1024, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (1024, 1024, 131072, 64, 64, False, True, True): (1, 2048, 2, 4), + (1024, 1024, 131072, 64, 64, True, False, True): (1, 2048, 2, 4), + (1024, 1024, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (1024, 1024, 131072, 128, 128, True, False, True): (1, 1024, 1, 32), + (1280, 5120, 65792, 16, 16, False, True, True): (1, 1028, 3, 1), + (1280, 5120, 65792, 16, 16, True, False, True): (1, 257, 3, 4), + (1280, 5120, 65792, 32, 32, False, True, True): (1, 514, 3, 4), + (1280, 5120, 65792, 32, 32, True, False, True): (1, 514, 3, 4), + (1280, 5120, 65792, 64, 64, False, True, True): (2, 1028, 3, 4), + (1280, 5120, 65792, 64, 64, True, False, True): (1, 1028, 3, 4), + (1280, 5120, 65792, 128, 128, False, True, True): (2, 514, 2, 32), + (1280, 5120, 65792, 128, 128, True, False, True): (1, 514, 2, 32), + (1536, 1536, 256, 16, 16, False, True, True): (5, 4, 3, 2), + (1536, 1536, 256, 16, 16, True, False, True): (2, 2, 3, 4), + (1536, 1536, 256, 32, 32, False, True, True): (1, 8, 2, 4), + (1536, 1536, 256, 32, 32, True, False, True): (2, 4, 3, 4), + (1536, 1536, 256, 64, 64, False, True, True): (1, 4, 3, 8), + (1536, 1536, 256, 64, 64, True, False, True): (2, 4, 3, 8), + (1536, 1536, 256, 128, 128, False, True, True): (1, 2, 1, 32), + (1536, 1536, 256, 128, 128, True, False, True): (2, 2, 2, 32), + (1536, 1536, 512, 16, 16, False, True, True): (1, 8, 3, 2), + (1536, 1536, 512, 16, 16, True, False, True): (1, 8, 3, 2), + (1536, 1536, 512, 32, 32, False, True, True): (1, 16, 3, 4), + (1536, 1536, 512, 32, 32, True, False, True): (1, 16, 3, 4), + (1536, 1536, 512, 64, 64, False, True, True): (3, 8, 3, 8), + (1536, 1536, 512, 64, 64, True, False, True): (3, 8, 3, 8), + (1536, 1536, 512, 128, 128, False, True, True): (1, 4, 1, 32), + (1536, 1536, 512, 128, 128, True, False, True): (2, 4, 2, 32), + (1536, 1536, 1024, 16, 16, False, True, True): (2, 8, 3, 4), + (1536, 1536, 1024, 16, 16, True, False, True): (2, 8, 3, 4), + (1536, 1536, 1024, 32, 32, False, True, True): (1, 16, 3, 4), + (1536, 1536, 1024, 32, 32, True, False, True): (1, 16, 3, 4), + (1536, 1536, 1024, 64, 64, False, True, True): (2, 16, 3, 8), + (1536, 1536, 1024, 64, 64, True, False, True): (2, 16, 3, 8), + (1536, 1536, 1024, 128, 128, False, True, True): (3, 8, 1, 32), + (1536, 1536, 1024, 128, 128, True, False, True): (1, 8, 2, 32), + (1536, 1536, 2048, 16, 16, False, True, True): (1, 32, 3, 2), + (1536, 1536, 2048, 16, 16, True, False, True): (1, 32, 3, 2), + (1536, 1536, 2048, 32, 32, False, True, True): (3, 32, 2, 4), + (1536, 1536, 2048, 32, 32, True, False, True): (4, 32, 3, 4), + (1536, 1536, 2048, 64, 64, False, True, True): (1, 32, 3, 4), + (1536, 1536, 2048, 64, 64, True, False, True): (1, 32, 3, 4), + (1536, 1536, 2048, 128, 128, False, True, True): (1, 16, 1, 32), + (1536, 1536, 2048, 128, 128, True, False, True): (2, 16, 1, 32), + (1536, 1536, 4096, 16, 16, False, True, True): (1, 64, 3, 2), + (1536, 1536, 4096, 16, 16, True, False, True): (1, 16, 3, 4), + (1536, 1536, 4096, 32, 32, False, True, True): (1, 64, 2, 4), + (1536, 1536, 4096, 32, 32, True, False, True): (1, 64, 2, 4), + (1536, 1536, 4096, 64, 64, False, True, True): (1, 64, 3, 4), + (1536, 1536, 4096, 64, 64, True, False, True): (1, 64, 3, 4), + (1536, 1536, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (1536, 1536, 4096, 128, 128, True, False, True): (4, 32, 2, 32), + (1536, 1536, 8192, 16, 16, False, True, True): (1, 32, 3, 4), + (1536, 1536, 8192, 16, 16, True, False, True): (5, 32, 3, 4), + (1536, 1536, 8192, 32, 32, False, True, True): (1, 128, 2, 4), + (1536, 1536, 8192, 32, 32, True, False, True): (1, 128, 2, 4), + (1536, 1536, 8192, 64, 64, False, True, True): (1, 128, 3, 4), + (1536, 1536, 8192, 64, 64, True, False, True): (1, 128, 3, 4), + (1536, 1536, 8192, 128, 128, False, True, True): (1, 64, 1, 32), + (1536, 1536, 8192, 128, 128, True, False, True): (4, 64, 2, 32), + (1536, 1536, 16384, 16, 16, False, True, True): (1, 64, 3, 4), + (1536, 1536, 16384, 16, 16, True, False, True): (1, 64, 3, 4), + (1536, 1536, 16384, 32, 32, False, True, True): (1, 256, 2, 4), + (1536, 1536, 16384, 32, 32, True, False, True): (1, 128, 3, 4), + (1536, 1536, 16384, 64, 64, False, True, True): (1, 256, 3, 4), + (1536, 1536, 16384, 64, 64, True, False, True): (3, 256, 3, 4), + (1536, 1536, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (1536, 1536, 16384, 128, 128, True, False, True): (4, 128, 2, 32), + (1536, 1536, 32768, 16, 16, False, True, True): (1, 128, 3, 4), + (1536, 1536, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (1536, 1536, 32768, 32, 32, False, True, True): (1, 256, 3, 4), + (1536, 1536, 32768, 32, 32, True, False, True): (1, 256, 3, 4), + (1536, 1536, 32768, 64, 64, False, True, True): (1, 512, 3, 4), + (1536, 1536, 32768, 64, 64, True, False, True): (1, 512, 3, 4), + (1536, 1536, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (1536, 1536, 32768, 128, 128, True, False, True): (4, 256, 2, 32), + (1536, 1536, 65536, 16, 16, False, True, True): (5, 256, 3, 4), + (1536, 1536, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (1536, 1536, 65536, 32, 32, False, True, True): (1, 512, 3, 4), + (1536, 1536, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (1536, 1536, 65536, 64, 64, False, True, True): (1, 1024, 3, 4), + (1536, 1536, 65536, 64, 64, True, False, True): (1, 1024, 3, 4), + (1536, 1536, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (1536, 1536, 65536, 128, 128, True, False, True): (4, 512, 2, 32), + (1536, 1536, 131072, 16, 16, False, True, True): (3, 512, 3, 4), + (1536, 1536, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (1536, 1536, 131072, 32, 32, False, True, True): (1, 1024, 3, 4), + (1536, 1536, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (1536, 1536, 131072, 64, 64, False, True, True): (1, 2048, 3, 4), + (1536, 1536, 131072, 64, 64, True, False, True): (1, 2048, 3, 4), + (1536, 1536, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (1536, 1536, 131072, 128, 128, True, False, True): (4, 1024, 2, 32), + (2048, 2048, 256, 16, 16, False, True, True): (1, 4, 3, 4), + (2048, 2048, 256, 16, 16, True, False, True): (1, 4, 3, 4), + (2048, 2048, 256, 32, 32, False, True, True): (3, 8, 3, 4), + (2048, 2048, 256, 32, 32, True, False, True): (3, 8, 3, 4), + (2048, 2048, 256, 64, 64, False, True, True): (4, 4, 4, 8), + (2048, 2048, 256, 64, 64, True, False, True): (8, 4, 4, 8), + (2048, 2048, 256, 128, 128, False, True, True): (3, 2, 1, 32), + (2048, 2048, 256, 128, 128, True, False, True): (3, 2, 1, 32), + (2048, 2048, 512, 16, 16, False, True, True): (4, 8, 3, 2), + (2048, 2048, 512, 16, 16, True, False, True): (4, 8, 3, 2), + (2048, 2048, 512, 32, 32, False, True, True): (3, 8, 3, 4), + (2048, 2048, 512, 32, 32, True, False, True): (1, 16, 2, 4), + (2048, 2048, 512, 64, 64, False, True, True): (4, 8, 2, 4), + (2048, 2048, 512, 64, 64, True, False, True): (4, 8, 2, 4), + (2048, 2048, 512, 128, 128, False, True, True): (1, 4, 1, 32), + (2048, 2048, 512, 128, 128, True, False, True): (4, 4, 1, 32), + (2048, 2048, 1024, 16, 16, False, True, True): (4, 8, 3, 4), + (2048, 2048, 1024, 16, 16, True, False, True): (4, 8, 3, 4), + (2048, 2048, 1024, 32, 32, False, True, True): (4, 16, 3, 4), + (2048, 2048, 1024, 32, 32, True, False, True): (1, 16, 3, 4), + (2048, 2048, 1024, 64, 64, False, True, True): (2, 16, 2, 4), + (2048, 2048, 1024, 64, 64, True, False, True): (2, 16, 2, 4), + (2048, 2048, 1024, 128, 128, False, True, True): (8, 8, 1, 32), + (2048, 2048, 1024, 128, 128, True, False, True): (4, 8, 1, 32), + (2048, 2048, 2048, 16, 16, False, True, True): (4, 32, 3, 1), + (2048, 2048, 2048, 16, 16, True, False, True): (3, 32, 3, 2), + (2048, 2048, 2048, 32, 32, False, True, True): (1, 32, 3, 4), + (2048, 2048, 2048, 32, 32, True, False, True): (1, 32, 3, 4), + (2048, 2048, 2048, 64, 64, False, True, True): (2, 32, 2, 4), + (2048, 2048, 2048, 64, 64, True, False, True): (2, 32, 2, 4), + (2048, 2048, 2048, 128, 128, False, True, True): (6, 16, 1, 32), + (2048, 2048, 2048, 128, 128, True, False, True): (4, 16, 1, 32), + (2048, 2048, 4096, 16, 16, False, True, True): (4, 64, 3, 1), + (2048, 2048, 4096, 16, 16, True, False, True): (1, 64, 3, 1), + (2048, 2048, 4096, 32, 32, False, True, True): (1, 64, 3, 4), + (2048, 2048, 4096, 32, 32, True, False, True): (4, 64, 3, 4), + (2048, 2048, 4096, 64, 64, False, True, True): (2, 64, 2, 4), + (2048, 2048, 4096, 64, 64, True, False, True): (2, 64, 2, 4), + (2048, 2048, 4096, 128, 128, False, True, True): (4, 32, 1, 32), + (2048, 2048, 4096, 128, 128, True, False, True): (4, 32, 1, 32), + (2048, 2048, 8192, 16, 16, False, True, True): (4, 128, 3, 1), + (2048, 2048, 8192, 16, 16, True, False, True): (1, 128, 3, 1), + (2048, 2048, 8192, 32, 32, False, True, True): (4, 128, 3, 4), + (2048, 2048, 8192, 32, 32, True, False, True): (4, 64, 3, 4), + (2048, 2048, 8192, 64, 64, False, True, True): (1, 128, 2, 4), + (2048, 2048, 8192, 64, 64, True, False, True): (2, 128, 2, 4), + (2048, 2048, 8192, 128, 128, False, True, True): (1, 64, 1, 32), + (2048, 2048, 8192, 128, 128, True, False, True): (4, 64, 1, 32), + (2048, 2048, 16384, 16, 16, False, True, True): (4, 256, 3, 1), + (2048, 2048, 16384, 16, 16, True, False, True): (1, 256, 3, 1), + (2048, 2048, 16384, 32, 32, False, True, True): (1, 256, 3, 4), + (2048, 2048, 16384, 32, 32, True, False, True): (1, 128, 3, 4), + (2048, 2048, 16384, 64, 64, False, True, True): (1, 256, 2, 4), + (2048, 2048, 16384, 64, 64, True, False, True): (1, 256, 2, 4), + (2048, 2048, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (2048, 2048, 16384, 128, 128, True, False, True): (4, 128, 1, 32), + (2048, 2048, 32768, 16, 16, False, True, True): (8, 512, 3, 1), + (2048, 2048, 32768, 16, 16, True, False, True): (1, 512, 3, 1), + (2048, 2048, 32768, 32, 32, False, True, True): (1, 512, 3, 4), + (2048, 2048, 32768, 32, 32, True, False, True): (1, 256, 3, 4), + (2048, 2048, 32768, 64, 64, False, True, True): (1, 512, 2, 4), + (2048, 2048, 32768, 64, 64, True, False, True): (1, 512, 2, 4), + (2048, 2048, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (2048, 2048, 32768, 128, 128, True, False, True): (4, 256, 1, 32), + (2048, 2048, 65536, 16, 16, False, True, True): (4, 1024, 3, 1), + (2048, 2048, 65536, 16, 16, True, False, True): (1, 1024, 3, 1), + (2048, 2048, 65536, 32, 32, False, True, True): (1, 1024, 3, 4), + (2048, 2048, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (2048, 2048, 65536, 64, 64, False, True, True): (1, 1024, 2, 4), + (2048, 2048, 65536, 64, 64, True, False, True): (1, 1024, 2, 4), + (2048, 2048, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (2048, 2048, 65536, 128, 128, True, False, True): (4, 512, 1, 32), + (2048, 2048, 131072, 16, 16, False, True, True): (4, 2048, 3, 1), + (2048, 2048, 131072, 16, 16, True, False, True): (1, 2048, 3, 1), + (2048, 2048, 131072, 32, 32, False, True, True): (1, 2048, 3, 4), + (2048, 2048, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (2048, 2048, 131072, 64, 64, False, True, True): (1, 2048, 2, 4), + (2048, 2048, 131072, 64, 64, True, False, True): (1, 2048, 2, 4), + (2048, 2048, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (2048, 2048, 131072, 128, 128, True, False, True): (4, 1024, 1, 32), + (3072, 768, 256, 16, 16, False, True, True): (4, 4, 3, 2), + (3072, 768, 256, 16, 16, True, False, True): (1, 2, 6, 4), + (3072, 768, 256, 32, 32, False, True, True): (1, 4, 6, 4), + (3072, 768, 256, 32, 32, True, False, True): (5, 4, 3, 4), + (3072, 768, 256, 64, 64, False, True, True): (4, 4, 3, 8), + (3072, 768, 256, 64, 64, True, False, True): (4, 4, 3, 8), + (3072, 768, 256, 128, 128, False, True, True): (1, 2, 1, 32), + (3072, 768, 256, 128, 128, True, False, True): (5, 2, 1, 32), + (3072, 768, 512, 16, 16, False, True, True): (4, 4, 3, 4), + (3072, 768, 512, 16, 16, True, False, True): (1, 4, 3, 4), + (3072, 768, 512, 32, 32, False, True, True): (3, 8, 3, 4), + (3072, 768, 512, 32, 32, True, False, True): (3, 8, 3, 4), + (3072, 768, 512, 64, 64, False, True, True): (2, 8, 3, 8), + (3072, 768, 512, 64, 64, True, False, True): (2, 8, 3, 8), + (3072, 768, 512, 128, 128, False, True, True): (1, 4, 2, 32), + (3072, 768, 512, 128, 128, True, False, True): (1, 4, 1, 32), + (3072, 768, 1024, 16, 16, False, True, True): (1, 16, 3, 2), + (3072, 768, 1024, 16, 16, True, False, True): (3, 16, 3, 2), + (3072, 768, 1024, 32, 32, False, True, True): (1, 16, 3, 4), + (3072, 768, 1024, 32, 32, True, False, True): (3, 16, 3, 4), + (3072, 768, 1024, 64, 64, False, True, True): (4, 16, 3, 8), + (3072, 768, 1024, 64, 64, True, False, True): (4, 16, 3, 4), + (3072, 768, 1024, 128, 128, False, True, True): (5, 8, 1, 32), + (3072, 768, 1024, 128, 128, True, False, True): (5, 8, 1, 32), + (3072, 768, 2048, 16, 16, False, True, True): (4, 32, 3, 2), + (3072, 768, 2048, 16, 16, True, False, True): (1, 32, 3, 2), + (3072, 768, 2048, 32, 32, False, True, True): (1, 32, 3, 4), + (3072, 768, 2048, 32, 32, True, False, True): (1, 32, 2, 4), + (3072, 768, 2048, 64, 64, False, True, True): (2, 32, 3, 4), + (3072, 768, 2048, 64, 64, True, False, True): (4, 32, 3, 4), + (3072, 768, 2048, 128, 128, False, True, True): (1, 16, 1, 32), + (3072, 768, 2048, 128, 128, True, False, True): (1, 16, 1, 32), + (3072, 768, 4096, 16, 16, False, True, True): (3, 64, 3, 2), + (3072, 768, 4096, 16, 16, True, False, True): (1, 64, 3, 2), + (3072, 768, 4096, 32, 32, False, True, True): (1, 64, 3, 4), + (3072, 768, 4096, 32, 32, True, False, True): (1, 32, 3, 4), + (3072, 768, 4096, 64, 64, False, True, True): (2, 64, 3, 4), + (3072, 768, 4096, 64, 64, True, False, True): (2, 64, 3, 4), + (3072, 768, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (3072, 768, 4096, 128, 128, True, False, True): (1, 32, 1, 32), + (3072, 768, 8192, 16, 16, False, True, True): (4, 128, 3, 1), + (3072, 768, 8192, 16, 16, True, False, True): (1, 32, 3, 4), + (3072, 768, 8192, 32, 32, False, True, True): (1, 64, 3, 4), + (3072, 768, 8192, 32, 32, True, False, True): (1, 64, 3, 4), + (3072, 768, 8192, 64, 64, False, True, True): (2, 128, 3, 4), + (3072, 768, 8192, 64, 64, True, False, True): (2, 128, 3, 4), + (3072, 768, 8192, 128, 128, False, True, True): (1, 64, 1, 32), + (3072, 768, 8192, 128, 128, True, False, True): (1, 64, 1, 32), + (3072, 768, 16384, 16, 16, False, True, True): (4, 256, 3, 1), + (3072, 768, 16384, 16, 16, True, False, True): (1, 64, 3, 4), + (3072, 768, 16384, 32, 32, False, True, True): (1, 128, 3, 4), + (3072, 768, 16384, 32, 32, True, False, True): (1, 128, 3, 4), + (3072, 768, 16384, 64, 64, False, True, True): (2, 256, 3, 4), + (3072, 768, 16384, 64, 64, True, False, True): (2, 256, 3, 4), + (3072, 768, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (3072, 768, 16384, 128, 128, True, False, True): (1, 128, 1, 32), + (3072, 768, 32768, 16, 16, False, True, True): (4, 512, 3, 1), + (3072, 768, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (3072, 768, 32768, 32, 32, False, True, True): (1, 256, 3, 4), + (3072, 768, 32768, 32, 32, True, False, True): (1, 256, 3, 4), + (3072, 768, 32768, 64, 64, False, True, True): (2, 512, 3, 4), + (3072, 768, 32768, 64, 64, True, False, True): (2, 512, 3, 4), + (3072, 768, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (3072, 768, 32768, 128, 128, True, False, True): (1, 256, 1, 32), + (3072, 768, 50432, 16, 16, False, True, True): (4, 788, 3, 1), + (3072, 768, 50432, 16, 16, True, False, True): (1, 197, 3, 4), + (3072, 768, 50432, 32, 32, False, True, True): (1, 394, 3, 4), + (3072, 768, 50432, 32, 32, True, False, True): (1, 394, 3, 4), + (3072, 768, 50432, 64, 64, False, True, True): (1, 788, 3, 4), + (3072, 768, 50432, 64, 64, True, False, True): (2, 788, 3, 4), + (3072, 768, 50432, 128, 128, False, True, True): (1, 394, 1, 32), + (3072, 768, 50432, 128, 128, True, False, True): (1, 394, 1, 32), + (3072, 768, 65536, 16, 16, False, True, True): (4, 1024, 3, 1), + (3072, 768, 65536, 16, 16, True, False, True): (1, 256, 3, 4), + (3072, 768, 65536, 32, 32, False, True, True): (1, 512, 3, 4), + (3072, 768, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (3072, 768, 65536, 64, 64, False, True, True): (2, 1024, 3, 4), + (3072, 768, 65536, 64, 64, True, False, True): (2, 1024, 3, 4), + (3072, 768, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (3072, 768, 65536, 128, 128, True, False, True): (1, 512, 1, 32), + (3072, 768, 131072, 16, 16, False, True, True): (4, 2048, 3, 1), + (3072, 768, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (3072, 768, 131072, 32, 32, False, True, True): (1, 1024, 3, 4), + (3072, 768, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (3072, 768, 131072, 64, 64, False, True, True): (2, 2048, 3, 4), + (3072, 768, 131072, 64, 64, True, False, True): (2, 2048, 3, 4), + (3072, 768, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (3072, 768, 131072, 128, 128, True, False, True): (1, 1024, 1, 32), + (3072, 3072, 256, 16, 16, False, True, True): (1, 4, 5, 2), + (3072, 3072, 256, 16, 16, True, False, True): (1, 4, 3, 2), + (3072, 3072, 256, 32, 32, False, True, True): (1, 4, 4, 4), + (3072, 3072, 256, 32, 32, True, False, True): (1, 4, 3, 4), + (3072, 3072, 256, 64, 64, False, True, True): (2, 4, 3, 8), + (3072, 3072, 256, 64, 64, True, False, True): (2, 4, 3, 8), + (3072, 3072, 256, 128, 128, False, True, True): (6, 2, 1, 32), + (3072, 3072, 256, 128, 128, True, False, True): (8, 2, 2, 32), + (3072, 3072, 512, 16, 16, False, True, True): (2, 4, 3, 4), + (3072, 3072, 512, 16, 16, True, False, True): (2, 4, 3, 4), + (3072, 3072, 512, 32, 32, False, True, True): (2, 8, 3, 4), + (3072, 3072, 512, 32, 32, True, False, True): (2, 8, 3, 4), + (3072, 3072, 512, 64, 64, False, True, True): (2, 8, 3, 8), + (3072, 3072, 512, 64, 64, True, False, True): (2, 8, 3, 8), + (3072, 3072, 512, 128, 128, False, True, True): (5, 4, 1, 32), + (3072, 3072, 512, 128, 128, True, False, True): (5, 4, 2, 32), + (3072, 3072, 1024, 16, 16, False, True, True): (1, 16, 3, 2), + (3072, 3072, 1024, 16, 16, True, False, True): (1, 16, 3, 2), + (3072, 3072, 1024, 32, 32, False, True, True): (2, 16, 3, 4), + (3072, 3072, 1024, 32, 32, True, False, True): (1, 16, 3, 4), + (3072, 3072, 1024, 64, 64, False, True, True): (1, 16, 3, 4), + (3072, 3072, 1024, 64, 64, True, False, True): (1, 16, 3, 4), + (3072, 3072, 1024, 128, 128, False, True, True): (1, 8, 1, 32), + (3072, 3072, 1024, 128, 128, True, False, True): (3, 8, 2, 32), + (3072, 3072, 2048, 16, 16, False, True, True): (1, 32, 3, 2), + (3072, 3072, 2048, 16, 16, True, False, True): (1, 16, 2, 4), + (3072, 3072, 2048, 32, 32, False, True, True): (1, 32, 2, 4), + (3072, 3072, 2048, 32, 32, True, False, True): (1, 32, 3, 4), + (3072, 3072, 2048, 64, 64, False, True, True): (1, 32, 3, 4), + (3072, 3072, 2048, 64, 64, True, False, True): (1, 32, 3, 4), + (3072, 3072, 2048, 128, 128, False, True, True): (1, 16, 1, 32), + (3072, 3072, 2048, 128, 128, True, False, True): (4, 16, 2, 32), + (3072, 3072, 4096, 16, 16, False, True, True): (2, 16, 3, 4), + (3072, 3072, 4096, 16, 16, True, False, True): (2, 16, 3, 4), + (3072, 3072, 4096, 32, 32, False, True, True): (1, 64, 2, 4), + (3072, 3072, 4096, 32, 32, True, False, True): (1, 32, 3, 4), + (3072, 3072, 4096, 64, 64, False, True, True): (1, 64, 3, 4), + (3072, 3072, 4096, 64, 64, True, False, True): (1, 64, 3, 4), + (3072, 3072, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (3072, 3072, 4096, 128, 128, True, False, True): (2, 32, 2, 32), + (3072, 3072, 8192, 16, 16, False, True, True): (2, 32, 3, 4), + (3072, 3072, 8192, 16, 16, True, False, True): (2, 32, 3, 4), + (3072, 3072, 8192, 32, 32, False, True, True): (1, 64, 3, 4), + (3072, 3072, 8192, 32, 32, True, False, True): (1, 64, 3, 4), + (3072, 3072, 8192, 64, 64, False, True, True): (1, 128, 3, 4), + (3072, 3072, 8192, 64, 64, True, False, True): (1, 128, 3, 4), + (3072, 3072, 8192, 128, 128, False, True, True): (1, 64, 1, 32), + (3072, 3072, 8192, 128, 128, True, False, True): (4, 64, 2, 32), + (3072, 3072, 16384, 16, 16, False, True, True): (2, 64, 3, 4), + (3072, 3072, 16384, 16, 16, True, False, True): (1, 64, 3, 4), + (3072, 3072, 16384, 32, 32, False, True, True): (1, 128, 3, 4), + (3072, 3072, 16384, 32, 32, True, False, True): (1, 128, 3, 4), + (3072, 3072, 16384, 64, 64, False, True, True): (1, 256, 3, 4), + (3072, 3072, 16384, 64, 64, True, False, True): (1, 256, 3, 4), + (3072, 3072, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (3072, 3072, 16384, 128, 128, True, False, True): (4, 128, 2, 32), + (3072, 3072, 32768, 16, 16, False, True, True): (3, 128, 3, 4), + (3072, 3072, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (3072, 3072, 32768, 32, 32, False, True, True): (1, 256, 3, 4), + (3072, 3072, 32768, 32, 32, True, False, True): (1, 256, 3, 4), + (3072, 3072, 32768, 64, 64, False, True, True): (1, 512, 3, 4), + (3072, 3072, 32768, 64, 64, True, False, True): (1, 512, 3, 4), + (3072, 3072, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (3072, 3072, 32768, 128, 128, True, False, True): (4, 256, 2, 32), + (3072, 3072, 65536, 16, 16, False, True, True): (5, 256, 3, 4), + (3072, 3072, 65536, 16, 16, True, False, True): (2, 256, 3, 4), + (3072, 3072, 65536, 32, 32, False, True, True): (1, 512, 3, 4), + (3072, 3072, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (3072, 3072, 65536, 64, 64, False, True, True): (1, 1024, 3, 4), + (3072, 3072, 65536, 64, 64, True, False, True): (1, 1024, 3, 4), + (3072, 3072, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (3072, 3072, 65536, 128, 128, True, False, True): (4, 512, 2, 32), + (3072, 3072, 131072, 16, 16, False, True, True): (5, 512, 3, 4), + (3072, 3072, 131072, 16, 16, True, False, True): (1, 512, 3, 4), + (3072, 3072, 131072, 32, 32, False, True, True): (1, 1024, 3, 4), + (3072, 3072, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (3072, 3072, 131072, 64, 64, False, True, True): (1, 2048, 3, 4), + (3072, 3072, 131072, 64, 64, True, False, True): (1, 2048, 3, 4), + (3072, 3072, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (3072, 3072, 131072, 128, 128, True, False, True): (4, 1024, 2, 32), + (4096, 4096, 256, 16, 16, False, True, True): (1, 4, 3, 2), + (4096, 4096, 256, 16, 16, True, False, True): (1, 2, 3, 4), + (4096, 4096, 256, 32, 32, False, True, True): (4, 4, 4, 4), + (4096, 4096, 256, 32, 32, True, False, True): (4, 4, 4, 4), + (4096, 4096, 256, 64, 64, False, True, True): (1, 4, 3, 8), + (4096, 4096, 256, 64, 64, True, False, True): (4, 4, 2, 4), + (4096, 4096, 256, 128, 128, False, True, True): (1, 2, 1, 32), + (4096, 4096, 256, 128, 128, True, False, True): (3, 2, 1, 32), + (4096, 4096, 512, 16, 16, False, True, True): (1, 4, 3, 4), + (4096, 4096, 512, 16, 16, True, False, True): (5, 8, 3, 2), + (4096, 4096, 512, 32, 32, False, True, True): (4, 8, 3, 4), + (4096, 4096, 512, 32, 32, True, False, True): (4, 8, 3, 4), + (4096, 4096, 512, 64, 64, False, True, True): (1, 8, 2, 4), + (4096, 4096, 512, 64, 64, True, False, True): (1, 8, 2, 4), + (4096, 4096, 512, 128, 128, False, True, True): (4, 4, 1, 32), + (4096, 4096, 512, 128, 128, True, False, True): (4, 4, 1, 32), + (4096, 4096, 1024, 16, 16, False, True, True): (1, 8, 3, 4), + (4096, 4096, 1024, 16, 16, True, False, True): (1, 8, 3, 4), + (4096, 4096, 1024, 32, 32, False, True, True): (1, 16, 3, 4), + (4096, 4096, 1024, 32, 32, True, False, True): (1, 16, 3, 4), + (4096, 4096, 1024, 64, 64, False, True, True): (4, 16, 2, 4), + (4096, 4096, 1024, 64, 64, True, False, True): (4, 16, 2, 4), + (4096, 4096, 1024, 128, 128, False, True, True): (4, 8, 1, 32), + (4096, 4096, 1024, 128, 128, True, False, True): (4, 8, 1, 32), + (4096, 4096, 2048, 16, 16, False, True, True): (1, 32, 3, 1), + (4096, 4096, 2048, 16, 16, True, False, True): (6, 8, 3, 4), + (4096, 4096, 2048, 32, 32, False, True, True): (1, 32, 3, 4), + (4096, 4096, 2048, 32, 32, True, False, True): (1, 32, 3, 4), + (4096, 4096, 2048, 64, 64, False, True, True): (4, 32, 2, 4), + (4096, 4096, 2048, 64, 64, True, False, True): (4, 32, 2, 4), + (4096, 4096, 2048, 128, 128, False, True, True): (4, 16, 1, 32), + (4096, 4096, 2048, 128, 128, True, False, True): (4, 16, 1, 32), + (4096, 4096, 4096, 16, 16, False, True, True): (1, 16, 3, 4), + (4096, 4096, 4096, 16, 16, True, False, True): (1, 64, 3, 1), + (4096, 4096, 4096, 32, 32, False, True, True): (1, 64, 3, 4), + (4096, 4096, 4096, 32, 32, True, False, True): (1, 32, 3, 4), + (4096, 4096, 4096, 64, 64, False, True, True): (4, 64, 2, 4), + (4096, 4096, 4096, 64, 64, True, False, True): (4, 64, 2, 4), + (4096, 4096, 4096, 128, 128, False, True, True): (4, 32, 1, 32), + (4096, 4096, 4096, 128, 128, True, False, True): (4, 32, 1, 32), + (4096, 4096, 8192, 16, 16, False, True, True): (4, 128, 3, 1), + (4096, 4096, 8192, 16, 16, True, False, True): (1, 128, 3, 1), + (4096, 4096, 8192, 32, 32, False, True, True): (1, 128, 3, 4), + (4096, 4096, 8192, 32, 32, True, False, True): (1, 64, 3, 4), + (4096, 4096, 8192, 64, 64, False, True, True): (4, 128, 2, 4), + (4096, 4096, 8192, 64, 64, True, False, True): (4, 128, 2, 4), + (4096, 4096, 8192, 128, 128, False, True, True): (4, 64, 1, 32), + (4096, 4096, 8192, 128, 128, True, False, True): (4, 64, 1, 32), + (4096, 4096, 16384, 16, 16, False, True, True): (1, 64, 3, 4), + (4096, 4096, 16384, 16, 16, True, False, True): (1, 256, 3, 1), + (4096, 4096, 16384, 32, 32, False, True, True): (1, 256, 3, 4), + (4096, 4096, 16384, 32, 32, True, False, True): (1, 128, 3, 4), + (4096, 4096, 16384, 64, 64, False, True, True): (4, 256, 2, 4), + (4096, 4096, 16384, 64, 64, True, False, True): (4, 256, 2, 4), + (4096, 4096, 16384, 128, 128, False, True, True): (4, 128, 1, 32), + (4096, 4096, 16384, 128, 128, True, False, True): (4, 128, 1, 32), + (4096, 4096, 32768, 16, 16, False, True, True): (1, 128, 3, 4), + (4096, 4096, 32768, 16, 16, True, False, True): (1, 512, 3, 1), + (4096, 4096, 32768, 32, 32, False, True, True): (1, 512, 3, 4), + (4096, 4096, 32768, 32, 32, True, False, True): (1, 256, 3, 4), + (4096, 4096, 32768, 64, 64, False, True, True): (4, 512, 2, 4), + (4096, 4096, 32768, 64, 64, True, False, True): (4, 512, 2, 4), + (4096, 4096, 32768, 128, 128, False, True, True): (4, 256, 1, 32), + (4096, 4096, 32768, 128, 128, True, False, True): (4, 256, 1, 32), + (4096, 4096, 65536, 16, 16, False, True, True): (1, 256, 3, 4), + (4096, 4096, 65536, 16, 16, True, False, True): (1, 1024, 3, 1), + (4096, 4096, 65536, 32, 32, False, True, True): (1, 1024, 3, 4), + (4096, 4096, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (4096, 4096, 65536, 64, 64, False, True, True): (4, 1024, 2, 4), + (4096, 4096, 65536, 64, 64, True, False, True): (2, 1024, 2, 4), + (4096, 4096, 65536, 128, 128, False, True, True): (4, 512, 1, 32), + (4096, 4096, 65536, 128, 128, True, False, True): (4, 512, 1, 32), + (4096, 4096, 131072, 16, 16, False, True, True): (2, 2048, 3, 1), + (4096, 4096, 131072, 16, 16, True, False, True): (1, 2048, 3, 1), + (4096, 4096, 131072, 32, 32, False, True, True): (2, 2048, 3, 4), + (4096, 4096, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (4096, 4096, 131072, 64, 64, False, True, True): (2, 2048, 2, 4), + (4096, 4096, 131072, 64, 64, True, False, True): (2, 2048, 2, 4), + (4096, 4096, 131072, 128, 128, False, True, True): (4, 1024, 1, 32), + (4096, 4096, 131072, 128, 128, True, False, True): (4, 1024, 1, 32), + (5120, 1280, 65792, 16, 16, False, True, True): (2, 1028, 3, 1), + (5120, 1280, 65792, 16, 16, True, False, True): (1, 257, 3, 4), + (5120, 1280, 65792, 32, 32, False, True, True): (1, 514, 3, 4), + (5120, 1280, 65792, 32, 32, True, False, True): (1, 514, 3, 4), + (5120, 1280, 65792, 64, 64, False, True, True): (1, 1028, 3, 4), + (5120, 1280, 65792, 64, 64, True, False, True): (5, 1028, 3, 4), + (5120, 1280, 65792, 128, 128, False, True, True): (1, 514, 1, 32), + (5120, 1280, 65792, 128, 128, True, False, True): (4, 514, 2, 32), + (6144, 6144, 256, 16, 16, False, True, True): (2, 2, 3, 4), + (6144, 6144, 256, 16, 16, True, False, True): (2, 2, 3, 4), + (6144, 6144, 256, 32, 32, False, True, True): (2, 4, 3, 4), + (6144, 6144, 256, 32, 32, True, False, True): (2, 4, 3, 4), + (6144, 6144, 256, 64, 64, False, True, True): (1, 4, 3, 4), + (6144, 6144, 256, 64, 64, True, False, True): (1, 4, 3, 4), + (6144, 6144, 256, 128, 128, False, True, True): (1, 2, 1, 32), + (6144, 6144, 256, 128, 128, True, False, True): (5, 2, 2, 32), + (6144, 6144, 512, 16, 16, False, True, True): (4, 8, 3, 2), + (6144, 6144, 512, 16, 16, True, False, True): (4, 8, 3, 2), + (6144, 6144, 512, 32, 32, False, True, True): (2, 8, 3, 4), + (6144, 6144, 512, 32, 32, True, False, True): (2, 8, 3, 4), + (6144, 6144, 512, 64, 64, False, True, True): (1, 8, 3, 4), + (6144, 6144, 512, 64, 64, True, False, True): (1, 8, 3, 4), + (6144, 6144, 512, 128, 128, False, True, True): (1, 4, 1, 32), + (6144, 6144, 512, 128, 128, True, False, True): (4, 4, 2, 32), + (6144, 6144, 1024, 16, 16, False, True, True): (4, 16, 3, 2), + (6144, 6144, 1024, 16, 16, True, False, True): (4, 4, 3, 4), + (6144, 6144, 1024, 32, 32, False, True, True): (1, 16, 3, 4), + (6144, 6144, 1024, 32, 32, True, False, True): (1, 16, 3, 4), + (6144, 6144, 1024, 64, 64, False, True, True): (1, 16, 3, 4), + (6144, 6144, 1024, 64, 64, True, False, True): (1, 16, 3, 4), + (6144, 6144, 1024, 128, 128, False, True, True): (1, 8, 1, 32), + (6144, 6144, 1024, 128, 128, True, False, True): (4, 8, 2, 32), + (6144, 6144, 2048, 16, 16, False, True, True): (1, 8, 3, 4), + (6144, 6144, 2048, 16, 16, True, False, True): (4, 8, 3, 4), + (6144, 6144, 2048, 32, 32, False, True, True): (1, 16, 3, 4), + (6144, 6144, 2048, 32, 32, True, False, True): (1, 16, 3, 4), + (6144, 6144, 2048, 64, 64, False, True, True): (1, 32, 3, 4), + (6144, 6144, 2048, 64, 64, True, False, True): (3, 32, 3, 4), + (6144, 6144, 2048, 128, 128, False, True, True): (1, 16, 1, 32), + (6144, 6144, 2048, 128, 128, True, False, True): (1, 16, 2, 32), + (6144, 6144, 4096, 16, 16, False, True, True): (3, 16, 3, 4), + (6144, 6144, 4096, 16, 16, True, False, True): (4, 16, 3, 4), + (6144, 6144, 4096, 32, 32, False, True, True): (1, 32, 3, 4), + (6144, 6144, 4096, 32, 32, True, False, True): (1, 32, 3, 4), + (6144, 6144, 4096, 64, 64, False, True, True): (1, 64, 3, 4), + (6144, 6144, 4096, 64, 64, True, False, True): (1, 64, 3, 4), + (6144, 6144, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (6144, 6144, 4096, 128, 128, True, False, True): (4, 32, 2, 32), + (6144, 6144, 8192, 16, 16, False, True, True): (1, 32, 3, 4), + (6144, 6144, 8192, 16, 16, True, False, True): (4, 32, 3, 4), + (6144, 6144, 8192, 32, 32, False, True, True): (1, 64, 3, 4), + (6144, 6144, 8192, 32, 32, True, False, True): (1, 64, 3, 4), + (6144, 6144, 8192, 64, 64, False, True, True): (1, 128, 3, 4), + (6144, 6144, 8192, 64, 64, True, False, True): (1, 128, 3, 4), + (6144, 6144, 8192, 128, 128, False, True, True): (1, 64, 1, 32), + (6144, 6144, 8192, 128, 128, True, False, True): (4, 64, 2, 32), + (6144, 6144, 16384, 16, 16, False, True, True): (1, 64, 3, 4), + (6144, 6144, 16384, 16, 16, True, False, True): (4, 64, 3, 4), + (6144, 6144, 16384, 32, 32, False, True, True): (1, 128, 3, 4), + (6144, 6144, 16384, 32, 32, True, False, True): (1, 128, 3, 4), + (6144, 6144, 16384, 64, 64, False, True, True): (1, 256, 3, 4), + (6144, 6144, 16384, 64, 64, True, False, True): (1, 256, 3, 4), + (6144, 6144, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (6144, 6144, 16384, 128, 128, True, False, True): (4, 128, 2, 32), + (6144, 6144, 32768, 16, 16, False, True, True): (1, 128, 3, 4), + (6144, 6144, 32768, 16, 16, True, False, True): (4, 128, 3, 4), + (6144, 6144, 32768, 32, 32, False, True, True): (1, 256, 3, 4), + (6144, 6144, 32768, 32, 32, True, False, True): (1, 256, 3, 4), + (6144, 6144, 32768, 64, 64, False, True, True): (1, 512, 3, 4), + (6144, 6144, 32768, 64, 64, True, False, True): (1, 512, 3, 4), + (6144, 6144, 32768, 128, 128, False, True, True): (1, 256, 1, 32), + (6144, 6144, 32768, 128, 128, True, False, True): (4, 256, 2, 32), + (6144, 6144, 65536, 16, 16, False, True, True): (1, 256, 3, 4), + (6144, 6144, 65536, 16, 16, True, False, True): (2, 256, 3, 4), + (6144, 6144, 65536, 32, 32, False, True, True): (1, 512, 3, 4), + (6144, 6144, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (6144, 6144, 65536, 64, 64, False, True, True): (1, 1024, 3, 4), + (6144, 6144, 65536, 64, 64, True, False, True): (1, 1024, 3, 4), + (6144, 6144, 65536, 128, 128, False, True, True): (1, 512, 1, 32), + (6144, 6144, 65536, 128, 128, True, False, True): (4, 512, 2, 32), + (6144, 6144, 131072, 16, 16, False, True, True): (1, 512, 3, 4), + (6144, 6144, 131072, 16, 16, True, False, True): (2, 512, 3, 4), + (6144, 6144, 131072, 32, 32, False, True, True): (1, 1024, 3, 4), + (6144, 6144, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (6144, 6144, 131072, 64, 64, False, True, True): (1, 2048, 3, 4), + (6144, 6144, 131072, 64, 64, True, False, True): (1, 2048, 3, 4), + (6144, 6144, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (6144, 6144, 131072, 128, 128, True, False, True): (4, 1024, 2, 32), + (8192, 8192, 256, 16, 16, False, True, True): (2, 2, 4, 4), + (8192, 8192, 256, 16, 16, True, False, True): (1, 1, 3, 4), + (8192, 8192, 256, 32, 32, False, True, True): (2, 4, 3, 4), + (8192, 8192, 256, 32, 32, True, False, True): (2, 4, 3, 4), + (8192, 8192, 256, 64, 64, False, True, True): (4, 4, 2, 4), + (8192, 8192, 256, 64, 64, True, False, True): (4, 4, 2, 4), + (8192, 8192, 256, 128, 128, False, True, True): (1, 2, 1, 32), + (8192, 8192, 256, 128, 128, True, False, True): (4, 2, 1, 32), + (8192, 8192, 512, 16, 16, False, True, True): (1, 4, 3, 4), + (8192, 8192, 512, 16, 16, True, False, True): (3, 4, 3, 4), + (8192, 8192, 512, 32, 32, False, True, True): (1, 8, 3, 4), + (8192, 8192, 512, 32, 32, True, False, True): (6, 8, 3, 4), + (8192, 8192, 512, 64, 64, False, True, True): (4, 8, 2, 4), + (8192, 8192, 512, 64, 64, True, False, True): (4, 8, 2, 4), + (8192, 8192, 512, 128, 128, False, True, True): (4, 4, 1, 32), + (8192, 8192, 512, 128, 128, True, False, True): (4, 4, 1, 32), + (8192, 8192, 1024, 16, 16, False, True, True): (1, 4, 3, 4), + (8192, 8192, 1024, 16, 16, True, False, True): (1, 32, 3, 1), + (8192, 8192, 1024, 32, 32, False, True, True): (1, 16, 3, 4), + (8192, 8192, 1024, 32, 32, True, False, True): (1, 16, 3, 4), + (8192, 8192, 1024, 64, 64, False, True, True): (4, 16, 2, 4), + (8192, 8192, 1024, 64, 64, True, False, True): (4, 16, 2, 4), + (8192, 8192, 1024, 128, 128, False, True, True): (4, 8, 1, 32), + (8192, 8192, 1024, 128, 128, True, False, True): (4, 8, 1, 32), + (8192, 8192, 2048, 16, 16, False, True, True): (4, 8, 3, 4), + (8192, 8192, 2048, 16, 16, True, False, True): (1, 32, 3, 1), + (8192, 8192, 2048, 32, 32, False, True, True): (1, 32, 3, 4), + (8192, 8192, 2048, 32, 32, True, False, True): (1, 16, 4, 4), + (8192, 8192, 2048, 64, 64, False, True, True): (4, 32, 2, 4), + (8192, 8192, 2048, 64, 64, True, False, True): (4, 32, 2, 4), + (8192, 8192, 2048, 128, 128, False, True, True): (4, 16, 1, 32), + (8192, 8192, 2048, 128, 128, True, False, True): (4, 16, 1, 32), + (8192, 8192, 4096, 16, 16, False, True, True): (3, 16, 3, 4), + (8192, 8192, 4096, 16, 16, True, False, True): (2, 64, 3, 1), + (8192, 8192, 4096, 32, 32, False, True, True): (1, 64, 3, 4), + (8192, 8192, 4096, 32, 32, True, False, True): (1, 32, 3, 4), + (8192, 8192, 4096, 64, 64, False, True, True): (4, 64, 2, 4), + (8192, 8192, 4096, 64, 64, True, False, True): (2, 64, 2, 4), + (8192, 8192, 4096, 128, 128, False, True, True): (4, 32, 1, 32), + (8192, 8192, 4096, 128, 128, True, False, True): (4, 32, 1, 32), + (8192, 8192, 8192, 16, 16, False, True, True): (2, 128, 3, 1), + (8192, 8192, 8192, 16, 16, True, False, True): (2, 128, 3, 1), + (8192, 8192, 8192, 32, 32, False, True, True): (1, 128, 3, 4), + (8192, 8192, 8192, 32, 32, True, False, True): (1, 64, 3, 4), + (8192, 8192, 8192, 64, 64, False, True, True): (4, 128, 2, 4), + (8192, 8192, 8192, 64, 64, True, False, True): (2, 128, 2, 4), + (8192, 8192, 8192, 128, 128, False, True, True): (4, 64, 1, 32), + (8192, 8192, 8192, 128, 128, True, False, True): (4, 64, 1, 32), + (8192, 8192, 16384, 16, 16, False, True, True): (1, 64, 3, 4), + (8192, 8192, 16384, 16, 16, True, False, True): (1, 256, 3, 1), + (8192, 8192, 16384, 32, 32, False, True, True): (1, 256, 3, 4), + (8192, 8192, 16384, 32, 32, True, False, True): (1, 128, 3, 4), + (8192, 8192, 16384, 64, 64, False, True, True): (2, 256, 2, 4), + (8192, 8192, 16384, 64, 64, True, False, True): (2, 256, 2, 4), + (8192, 8192, 16384, 128, 128, False, True, True): (4, 128, 1, 32), + (8192, 8192, 16384, 128, 128, True, False, True): (4, 128, 1, 32), + (8192, 8192, 32768, 16, 16, False, True, True): (1, 512, 3, 1), + (8192, 8192, 32768, 16, 16, True, False, True): (1, 512, 3, 1), + (8192, 8192, 32768, 32, 32, False, True, True): (1, 512, 3, 4), + (8192, 8192, 32768, 32, 32, True, False, True): (1, 256, 3, 4), + (8192, 8192, 32768, 64, 64, False, True, True): (2, 512, 2, 4), + (8192, 8192, 32768, 64, 64, True, False, True): (2, 512, 2, 4), + (8192, 8192, 32768, 128, 128, False, True, True): (4, 256, 1, 32), + (8192, 8192, 32768, 128, 128, True, False, True): (4, 256, 1, 32), + (8192, 8192, 65536, 16, 16, False, True, True): (1, 256, 3, 4), + (8192, 8192, 65536, 16, 16, True, False, True): (1, 1024, 3, 1), + (8192, 8192, 65536, 32, 32, False, True, True): (1, 1024, 3, 4), + (8192, 8192, 65536, 32, 32, True, False, True): (1, 512, 3, 4), + (8192, 8192, 65536, 64, 64, False, True, True): (4, 1024, 2, 4), + (8192, 8192, 65536, 64, 64, True, False, True): (2, 1024, 2, 4), + (8192, 8192, 65536, 128, 128, False, True, True): (4, 512, 1, 32), + (8192, 8192, 65536, 128, 128, True, False, True): (4, 512, 1, 32), + (8192, 8192, 131072, 16, 16, False, True, True): (1, 2048, 3, 1), + (8192, 8192, 131072, 16, 16, True, False, True): (2, 2048, 3, 1), + (8192, 8192, 131072, 32, 32, False, True, True): (4, 2048, 3, 4), + (8192, 8192, 131072, 32, 32, True, False, True): (1, 1024, 3, 4), + (8192, 8192, 131072, 64, 64, False, True, True): (2, 2048, 2, 4), + (8192, 8192, 131072, 64, 64, True, False, True): (2, 2048, 2, 4), + (8192, 8192, 131072, 128, 128, False, True, True): (4, 1024, 1, 32), + (8192, 8192, 131072, 128, 128, True, False, True): (4, 1024, 1, 32), + (16384, 16384, 256, 16, 16, False, True, True): (1, 2, 3, 4), + (16384, 16384, 256, 16, 16, True, False, True): (1, 2, 3, 4), + (16384, 16384, 256, 32, 32, False, True, True): (1, 4, 3, 4), + (16384, 16384, 256, 32, 32, True, False, True): (1, 4, 3, 4), + (16384, 16384, 256, 64, 64, False, True, True): (2, 4, 2, 4), + (16384, 16384, 256, 64, 64, True, False, True): (2, 4, 2, 4), + (16384, 16384, 256, 128, 128, False, True, True): (2, 2, 1, 32), + (16384, 16384, 256, 128, 128, True, False, True): (2, 2, 1, 32), + (16384, 16384, 512, 16, 16, False, True, True): (1, 2, 3, 4), + (16384, 16384, 512, 16, 16, True, False, True): (5, 2, 3, 4), + (16384, 16384, 512, 32, 32, False, True, True): (1, 8, 3, 4), + (16384, 16384, 512, 32, 32, True, False, True): (1, 4, 3, 4), + (16384, 16384, 512, 64, 64, False, True, True): (4, 8, 2, 4), + (16384, 16384, 512, 64, 64, True, False, True): (4, 8, 2, 4), + (16384, 16384, 512, 128, 128, False, True, True): (4, 4, 1, 32), + (16384, 16384, 512, 128, 128, True, False, True): (4, 4, 1, 32), + (16384, 16384, 1024, 16, 16, False, True, True): (1, 4, 3, 4), + (16384, 16384, 1024, 16, 16, True, False, True): (2, 16, 3, 1), + (16384, 16384, 1024, 32, 32, False, True, True): (1, 16, 3, 4), + (16384, 16384, 1024, 32, 32, True, False, True): (1, 8, 3, 4), + (16384, 16384, 1024, 64, 64, False, True, True): (4, 16, 2, 4), + (16384, 16384, 1024, 64, 64, True, False, True): (4, 16, 2, 4), + (16384, 16384, 1024, 128, 128, False, True, True): (4, 8, 1, 32), + (16384, 16384, 1024, 128, 128, True, False, True): (4, 8, 1, 32), + (16384, 16384, 2048, 16, 16, False, True, True): (1, 8, 3, 4), + (16384, 16384, 2048, 16, 16, True, False, True): (2, 32, 3, 1), + (16384, 16384, 2048, 32, 32, False, True, True): (1, 32, 3, 4), + (16384, 16384, 2048, 32, 32, True, False, True): (1, 16, 3, 4), + (16384, 16384, 2048, 64, 64, False, True, True): (4, 32, 2, 4), + (16384, 16384, 2048, 64, 64, True, False, True): (2, 32, 2, 4), + (16384, 16384, 2048, 128, 128, False, True, True): (4, 16, 1, 32), + (16384, 16384, 2048, 128, 128, True, False, True): (4, 16, 1, 32), + (16384, 16384, 4096, 16, 16, False, True, True): (1, 16, 3, 4), + (16384, 16384, 4096, 16, 16, True, False, True): (2, 64, 3, 1), + (16384, 16384, 4096, 32, 32, False, True, True): (1, 64, 3, 4), + (16384, 16384, 4096, 32, 32, True, False, True): (1, 32, 3, 4), + (16384, 16384, 4096, 64, 64, False, True, True): (4, 64, 2, 4), + (16384, 16384, 4096, 64, 64, True, False, True): (2, 64, 2, 4), + (16384, 16384, 4096, 128, 128, False, True, True): (4, 32, 1, 32), + (16384, 16384, 4096, 128, 128, True, False, True): (4, 32, 1, 32), + (16384, 16384, 8192, 16, 16, False, True, True): (1, 128, 3, 1), + (16384, 16384, 8192, 16, 16, True, False, True): (2, 128, 3, 1), + (16384, 16384, 8192, 32, 32, False, True, True): (1, 128, 3, 4), + (16384, 16384, 8192, 32, 32, True, False, True): (1, 64, 3, 4), + (16384, 16384, 8192, 64, 64, False, True, True): (2, 128, 2, 4), + (16384, 16384, 8192, 64, 64, True, False, True): (2, 128, 2, 4), + (16384, 16384, 8192, 128, 128, False, True, True): (4, 64, 1, 32), + (16384, 16384, 8192, 128, 128, True, False, True): (4, 64, 1, 32), + (16384, 16384, 16384, 16, 16, False, True, True): (1, 64, 3, 4), + (16384, 16384, 16384, 16, 16, True, False, True): (2, 256, 3, 1), + (16384, 16384, 16384, 32, 32, False, True, True): (1, 256, 3, 4), + (16384, 16384, 16384, 32, 32, True, False, True): (1, 128, 3, 4), + (16384, 16384, 16384, 64, 64, False, True, True): (2, 256, 2, 4), + (16384, 16384, 16384, 64, 64, True, False, True): (2, 256, 2, 4), + (16384, 16384, 16384, 128, 128, False, True, True): (4, 128, 1, 32), + (16384, 16384, 16384, 128, 128, True, False, True): (4, 128, 1, 32), + (16384, 16384, 32768, 16, 16, False, True, True): (1, 512, 3, 1), + (16384, 16384, 32768, 16, 16, True, False, True): (1, 128, 3, 4), + (16384, 16384, 32768, 32, 32, False, True, True): (2, 512, 3, 4), + (16384, 16384, 32768, 32, 32, True, False, True): (1, 256, 4, 4), + (16384, 16384, 32768, 64, 64, False, True, True): (2, 512, 2, 4), + (16384, 16384, 32768, 64, 64, True, False, True): (2, 512, 2, 4), + (16384, 16384, 32768, 128, 128, False, True, True): (4, 256, 1, 32), + (16384, 16384, 32768, 128, 128, True, False, True): (4, 256, 1, 32), + (16384, 16384, 65536, 16, 16, False, True, True): (1, 256, 3, 4), + (16384, 16384, 65536, 16, 16, True, False, True): (1, 1024, 3, 1), + (16384, 16384, 65536, 32, 32, False, True, True): (1, 1024, 3, 4), + (16384, 16384, 65536, 32, 32, True, False, True): (1, 512, 4, 4), + (16384, 16384, 65536, 64, 64, False, True, True): (2, 1024, 2, 4), + (16384, 16384, 65536, 64, 64, True, False, True): (2, 1024, 2, 4), + (16384, 16384, 65536, 128, 128, False, True, True): (4, 512, 1, 32), + (16384, 16384, 65536, 128, 128, True, False, True): (4, 512, 1, 32), + (16384, 16384, 131072, 16, 16, False, True, True): (1, 1024, 4, 4), + (16384, 16384, 131072, 16, 16, True, False, True): (2, 2048, 3, 1), + (16384, 16384, 131072, 32, 32, False, True, True): (1, 1024, 2, 4), + (16384, 16384, 131072, 32, 32, True, False, True): (1, 1024, 2, 4), + (16384, 16384, 131072, 64, 64, False, True, True): (4, 2048, 2, 4), + (16384, 16384, 131072, 64, 64, True, False, True): (2, 2048, 2, 4), + (16384, 16384, 131072, 128, 128, False, True, True): (4, 1024, 1, 32), + (16384, 16384, 131072, 128, 128, True, False, True): (4, 1024, 1, 32), + }, + ("bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.float32, 0.56)): { + (192, 192, 256, 64, 64, False, True, True): (1, 4, 3, 8), + (192, 192, 256, 64, 64, True, False, True): (1, 4, 3, 8), + (192, 192, 512, 64, 64, False, True, True): (2, 8, 3, 8), + (192, 192, 512, 64, 64, True, False, True): (5, 8, 3, 8), + (192, 192, 1024, 64, 64, False, True, True): (2, 16, 4, 8), + (192, 192, 1024, 64, 64, True, False, True): (1, 16, 3, 8), + (192, 192, 2048, 64, 64, False, True, True): (3, 32, 3, 8), + (192, 192, 2048, 64, 64, True, False, True): (5, 32, 5, 8), + (192, 192, 4096, 64, 64, False, True, True): (3, 64, 2, 8), + (192, 192, 4096, 64, 64, True, False, True): (1, 64, 3, 8), + (192, 192, 8192, 64, 64, False, True, True): (3, 128, 3, 8), + (192, 192, 8192, 64, 64, True, False, True): (6, 128, 3, 4), + (192, 192, 16384, 64, 64, False, True, True): (1, 256, 1, 8), + (192, 192, 16384, 64, 64, True, False, True): (1, 256, 3, 4), + (192, 192, 32768, 64, 64, False, True, True): (1, 512, 1, 8), + (192, 192, 32768, 64, 64, True, False, True): (1, 512, 3, 4), + (192, 192, 65536, 64, 64, False, True, True): (1, 1024, 1, 8), + (192, 192, 65536, 64, 64, True, False, True): (1, 1024, 3, 4), + (192, 192, 131072, 64, 64, False, True, True): (1, 2048, 1, 8), + (192, 192, 131072, 64, 64, True, False, True): (3, 2048, 1, 4), + (384, 384, 256, 128, 128, False, True, True): (1, 2, 1, 32), + (384, 384, 256, 128, 128, True, False, True): (1, 2, 1, 32), + (384, 384, 512, 128, 128, False, True, True): (1, 4, 1, 32), + (384, 384, 512, 128, 128, True, False, True): (2, 4, 1, 32), + (384, 384, 1024, 128, 128, False, True, True): (1, 8, 1, 32), + (384, 384, 1024, 128, 128, True, False, True): (4, 8, 1, 32), + (384, 384, 2048, 128, 128, False, True, True): (1, 16, 1, 32), + (384, 384, 2048, 128, 128, True, False, True): (1, 16, 1, 32), + (384, 384, 4096, 128, 128, False, True, True): (1, 32, 1, 32), + (384, 384, 4096, 128, 128, True, False, True): (2, 32, 2, 32), + (384, 384, 8192, 128, 128, False, True, True): (1, 64, 1, 32), + (384, 384, 8192, 128, 128, True, False, True): (1, 64, 2, 32), + (384, 384, 16384, 128, 128, False, True, True): (1, 128, 1, 32), + (384, 384, 16384, 128, 128, True, False, True): (4, 128, 1, 32), + (384, 384, 32768, 128, 128, False, True, True): (3, 256, 1, 32), + (384, 384, 32768, 128, 128, True, False, True): (3, 256, 1, 32), + (384, 384, 65536, 128, 128, False, True, True): (3, 512, 1, 32), + (384, 384, 65536, 128, 128, True, False, True): (3, 512, 1, 32), + (384, 384, 131072, 128, 128, False, True, True): (1, 1024, 1, 32), + (384, 384, 131072, 128, 128, True, False, True): (3, 1024, 1, 32), + }, + ("bsr_dense_addmm", "NVIDIA A100-SXM4-80GB", (0, torch.int8, 0.5)): { + (1280, 5120, 65792, 32, 32, False, True, True): (1, 1028, 1, 8), + (1280, 5120, 65792, 32, 32, True, False, True): (1, 514, 3, 2), + (1280, 5120, 65792, 64, 64, False, True, True): (2, 514, 1, 4), + (1280, 5120, 65792, 64, 64, True, False, True): (1, 514, 3, 2), + (1280, 5120, 65792, 128, 128, False, True, True): (2, 514, 1, 8), + (1280, 5120, 65792, 128, 128, True, False, True): (1, 514, 2, 4), + (1280, 5120, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (1280, 5120, 65792, 256, 256, True, False, True): (1, 257, 1, 32), + (5120, 1280, 65792, 32, 32, False, True, True): (3, 1028, 1, 8), + (5120, 1280, 65792, 32, 32, True, False, True): (1, 514, 1, 2), + (5120, 1280, 65792, 64, 64, False, True, True): (1, 514, 1, 4), + (5120, 1280, 65792, 64, 64, True, False, True): (2, 514, 2, 2), + (5120, 1280, 65792, 128, 128, False, True, True): (2, 514, 1, 8), + (5120, 1280, 65792, 128, 128, True, False, True): (2, 514, 2, 4), + (5120, 1280, 65792, 256, 256, False, True, True): (1, 257, 1, 32), + (5120, 1280, 65792, 256, 256, True, False, True): (1, 257, 1, 32), + }, + ("scatter_mm", "NVIDIA A100-SXM4-80GB", (0, torch.bfloat16, 0.5)): { + (256, 256, 256, 16, 16): (1, 1, 16, 16, 1, 2), + (256, 256, 256, 32, 32): (1, 1, 16, 16, 1, 4), + (256, 256, 256, 64, 64): (1, 1, 16, 16, 1, 1), + (256, 256, 256, 128, 128): (2, 4, 16, 64, 1, 4), + (256, 256, 512, 16, 16): (1, 1, 16, 16, 1, 4), + (256, 256, 512, 32, 32): (1, 1, 16, 32, 1, 4), + (256, 256, 512, 64, 64): (1, 1, 16, 32, 1, 1), + (256, 256, 512, 128, 128): (1, 1, 32, 32, 1, 4), + (256, 256, 1024, 16, 16): (1, 1, 16, 16, 1, 4), + (256, 256, 1024, 32, 32): (1, 2, 16, 32, 1, 1), + (256, 256, 1024, 64, 64): (1, 1, 32, 32, 1, 2), + (256, 256, 1024, 128, 128): (1, 1, 32, 64, 1, 4), + (256, 256, 2048, 16, 16): (1, 1, 16, 64, 1, 8), + (256, 256, 2048, 32, 32): (2, 1, 32, 64, 1, 2), + (256, 256, 2048, 64, 64): (1, 1, 32, 32, 1, 1), + (256, 256, 2048, 128, 128): (1, 1, 64, 64, 1, 4), + (256, 256, 4096, 16, 16): (1, 1, 16, 64, 1, 1), + (256, 256, 4096, 32, 32): (2, 2, 32, 64, 1, 2), + (256, 256, 4096, 64, 64): (1, 1, 32, 128, 1, 4), + (256, 256, 4096, 128, 128): (1, 1, 64, 64, 1, 4), + (256, 256, 8192, 16, 16): (1, 2, 16, 64, 1, 2), + (256, 256, 8192, 32, 32): (1, 1, 32, 64, 1, 2), + (256, 256, 8192, 64, 64): (1, 1, 32, 64, 1, 2), + (256, 256, 8192, 128, 128): (1, 1, 64, 64, 1, 4), + (256, 256, 16384, 16, 16): (1, 1, 16, 64, 1, 2), + (256, 256, 16384, 32, 32): (1, 1, 32, 64, 1, 2), + (256, 256, 16384, 64, 64): (1, 1, 64, 64, 1, 2), + (256, 256, 16384, 128, 128): (2, 16, 64, 64, 1, 4), + (256, 256, 32768, 16, 16): (1, 1, 16, 128, 1, 2), + (256, 256, 32768, 32, 32): (1, 1, 32, 64, 1, 2), + (256, 256, 32768, 64, 64): (1, 1, 64, 64, 1, 2), + (256, 256, 32768, 128, 128): (2, 32, 64, 64, 1, 4), + (256, 256, 65536, 16, 16): (1, 1, 16, 64, 1, 1), + (256, 256, 65536, 32, 32): (1, 1, 32, 64, 1, 2), + (256, 256, 65536, 64, 64): (1, 1, 64, 32, 1, 1), + (256, 256, 65536, 128, 128): (2, 32, 64, 64, 1, 4), + (256, 256, 131072, 16, 16): (1, 1, 16, 64, 1, 1), + (256, 256, 131072, 32, 32): (1, 1, 32, 64, 1, 2), + (256, 256, 131072, 64, 64): (4, 1, 64, 32, 1, 1), + (256, 256, 131072, 128, 128): (2, 64, 64, 64, 1, 4), + (512, 512, 256, 16, 16): (1, 1, 16, 16, 1, 2), + (512, 512, 256, 32, 32): (1, 1, 16, 32, 1, 1), + (512, 512, 256, 64, 64): (1, 2, 16, 32, 1, 1), + (512, 512, 256, 128, 128): (2, 16, 64, 16, 2, 4), + (512, 512, 512, 16, 16): (1, 1, 16, 16, 1, 4), + (512, 512, 512, 32, 32): (1, 1, 16, 32, 1, 1), + (512, 512, 512, 64, 64): (1, 1, 32, 32, 1, 2), + (512, 512, 512, 128, 128): (2, 8, 32, 64, 1, 4), + (512, 512, 1024, 16, 16): (1, 1, 16, 64, 1, 8), + (512, 512, 1024, 32, 32): (1, 1, 32, 32, 3, 1), + (512, 512, 1024, 64, 64): (1, 4, 32, 64, 1, 2), + (512, 512, 1024, 128, 128): (1, 4, 64, 64, 1, 4), + (512, 512, 2048, 16, 16): (1, 1, 16, 64, 1, 2), + (512, 512, 2048, 32, 32): (1, 1, 32, 64, 1, 2), + (512, 512, 2048, 64, 64): (1, 1, 64, 64, 3, 4), + (512, 512, 2048, 128, 128): (1, 1, 64, 64, 1, 4), + (512, 512, 4096, 16, 16): (1, 1, 16, 64, 1, 2), + (512, 512, 4096, 32, 32): (2, 64, 32, 64, 1, 2), + (512, 512, 4096, 64, 64): (1, 1, 64, 64, 3, 4), + (512, 512, 4096, 128, 128): (1, 1, 64, 64, 1, 4), + (512, 512, 8192, 16, 16): (1, 2, 16, 128, 1, 2), + (512, 512, 8192, 32, 32): (1, 1, 32, 64, 1, 2), + (512, 512, 8192, 64, 64): (1, 1, 64, 64, 1, 2), + (512, 512, 8192, 128, 128): (1, 1, 64, 64, 1, 4), + (512, 512, 16384, 16, 16): (1, 2, 16, 128, 1, 2), + (512, 512, 16384, 32, 32): (1, 1, 32, 64, 1, 2), + (512, 512, 16384, 64, 64): (1, 1, 64, 64, 3, 2), + (512, 512, 16384, 128, 128): (2, 1, 64, 64, 1, 4), + (512, 512, 32768, 16, 16): (1, 2, 16, 128, 1, 2), + (512, 512, 32768, 32, 32): (1, 1, 32, 64, 1, 2), + (512, 512, 32768, 64, 64): (1, 1, 64, 64, 3, 4), + (512, 512, 32768, 128, 128): (2, 1, 64, 64, 1, 4), + (512, 512, 65536, 16, 16): (1, 2, 16, 128, 1, 2), + (512, 512, 65536, 32, 32): (1, 1, 32, 64, 1, 2), + (512, 512, 65536, 64, 64): (1, 1, 64, 64, 3, 4), + (512, 512, 65536, 128, 128): (2, 1, 64, 64, 1, 4), + (512, 512, 131072, 16, 16): (1, 1, 16, 64, 1, 1), + (512, 512, 131072, 32, 32): (1, 1, 32, 64, 1, 2), + (512, 512, 131072, 64, 64): (1, 1, 64, 64, 3, 4), + (512, 512, 131072, 128, 128): (2, 4, 64, 64, 1, 4), + (1024, 1024, 256, 16, 16): (1, 1, 16, 16, 1, 4), + (1024, 1024, 256, 32, 32): (2, 16, 32, 16, 3, 4), + (1024, 1024, 256, 64, 64): (1, 4, 32, 32, 1, 2), + (1024, 1024, 256, 128, 128): (1, 4, 128, 16, 3, 16), + (1024, 1024, 512, 16, 16): (1, 1, 16, 64, 1, 2), + (1024, 1024, 512, 32, 32): (2, 2, 32, 64, 1, 2), + (1024, 1024, 512, 64, 64): (2, 8, 64, 64, 3, 4), + (1024, 1024, 512, 128, 128): (1, 4, 64, 64, 1, 8), + (1024, 1024, 1024, 16, 16): (1, 1, 16, 64, 1, 2), + (1024, 1024, 1024, 32, 32): (1, 1, 32, 64, 1, 2), + (1024, 1024, 1024, 64, 64): (1, 8, 64, 64, 3, 4), + (1024, 1024, 1024, 128, 128): (1, 8, 64, 64, 1, 4), + (1024, 1024, 2048, 16, 16): (1, 2, 16, 64, 1, 2), + (1024, 1024, 2048, 32, 32): (1, 1, 32, 64, 1, 2), + (1024, 1024, 2048, 64, 64): (2, 16, 64, 64, 2, 2), + (1024, 1024, 2048, 128, 128): (2, 32, 64, 64, 1, 4), + (1024, 1024, 4096, 16, 16): (2, 16, 16, 128, 1, 2), + (1024, 1024, 4096, 32, 32): (1, 16, 32, 64, 3, 2), + (1024, 1024, 4096, 64, 64): (1, 1, 64, 64, 3, 4), + (1024, 1024, 4096, 128, 128): (2, 64, 128, 64, 1, 4), + (1024, 1024, 8192, 16, 16): (2, 16, 16, 128, 1, 2), + (1024, 1024, 8192, 32, 32): (1, 16, 32, 64, 3, 2), + (1024, 1024, 8192, 64, 64): (1, 1, 64, 64, 3, 4), + (1024, 1024, 8192, 128, 128): (2, 1, 64, 64, 1, 4), + (1024, 1024, 16384, 16, 16): (1, 2, 16, 128, 1, 2), + (1024, 1024, 16384, 32, 32): (1, 16, 32, 64, 3, 2), + (1024, 1024, 16384, 64, 64): (1, 1, 64, 64, 3, 4), + (1024, 1024, 16384, 128, 128): (2, 16, 128, 64, 1, 4), + (1024, 1024, 32768, 16, 16): (1, 1, 16, 128, 1, 2), + (1024, 1024, 32768, 32, 32): (1, 1, 32, 128, 1, 2), + (1024, 1024, 32768, 64, 64): (1, 32, 64, 32, 2, 1), + (1024, 1024, 32768, 128, 128): (2, 8, 128, 64, 1, 4), + (1024, 1024, 65536, 16, 16): (3, 2, 16, 128, 1, 2), + (1024, 1024, 65536, 32, 32): (1, 1, 32, 128, 1, 2), + (1024, 1024, 65536, 64, 64): (2, 4, 64, 32, 2, 1), + (1024, 1024, 65536, 128, 128): (2, 8, 128, 64, 1, 4), + (1024, 1024, 131072, 16, 16): (2, 1, 16, 128, 1, 2), + (1024, 1024, 131072, 32, 32): (1, 1, 32, 128, 1, 2), + (1024, 1024, 131072, 64, 64): (1, 4, 64, 32, 2, 1), + (1024, 1024, 131072, 128, 128): (4, 1, 128, 64, 1, 4), + (2048, 2048, 256, 16, 16): (1, 1, 16, 64, 1, 8), + (2048, 2048, 256, 32, 32): (1, 1, 32, 32, 3, 1), + (2048, 2048, 256, 64, 64): (1, 1, 32, 32, 2, 1), + (2048, 2048, 256, 128, 128): (1, 4, 64, 64, 1, 8), + (2048, 2048, 512, 16, 16): (1, 2, 16, 64, 1, 2), + (2048, 2048, 512, 32, 32): (1, 2, 32, 64, 1, 4), + (2048, 2048, 512, 64, 64): (1, 4, 64, 64, 1, 8), + (2048, 2048, 512, 128, 128): (1, 4, 64, 64, 1, 4), + (2048, 2048, 1024, 16, 16): (1, 2, 16, 128, 1, 2), + (2048, 2048, 1024, 32, 32): (1, 1, 32, 64, 1, 2), + (2048, 2048, 1024, 64, 64): (1, 8, 64, 64, 1, 4), + (2048, 2048, 1024, 128, 128): (1, 8, 128, 64, 1, 4), + (2048, 2048, 2048, 16, 16): (3, 4, 16, 128, 1, 2), + (2048, 2048, 2048, 32, 32): (1, 16, 32, 64, 5, 2), + (2048, 2048, 2048, 64, 64): (1, 1, 64, 64, 3, 4), + (2048, 2048, 2048, 128, 128): (1, 8, 128, 64, 1, 4), + (2048, 2048, 4096, 16, 16): (1, 2, 16, 128, 1, 2), + (2048, 2048, 4096, 32, 32): (1, 8, 32, 64, 3, 2), + (2048, 2048, 4096, 64, 64): (1, 1, 64, 64, 3, 4), + (2048, 2048, 4096, 128, 128): (1, 8, 128, 64, 1, 4), + (2048, 2048, 8192, 16, 16): (2, 4, 16, 128, 1, 2), + (2048, 2048, 8192, 32, 32): (1, 4, 32, 128, 3, 2), + (2048, 2048, 8192, 64, 64): (1, 8, 64, 64, 3, 2), + (2048, 2048, 8192, 128, 128): (1, 8, 128, 64, 1, 4), + (2048, 2048, 16384, 16, 16): (1, 2, 16, 128, 1, 2), + (2048, 2048, 16384, 32, 32): (1, 4, 32, 128, 3, 2), + (2048, 2048, 16384, 64, 64): (1, 8, 64, 64, 3, 2), + (2048, 2048, 16384, 128, 128): (1, 4, 128, 64, 1, 4), + (2048, 2048, 32768, 16, 16): (3, 2, 16, 128, 1, 2), + (2048, 2048, 32768, 32, 32): (1, 1, 32, 128, 3, 2), + (2048, 2048, 32768, 64, 64): (1, 1, 64, 64, 3, 2), + (2048, 2048, 32768, 128, 128): (1, 4, 128, 64, 1, 4), + (2048, 2048, 65536, 16, 16): (1, 2, 16, 128, 1, 2), + (2048, 2048, 65536, 32, 32): (1, 4, 32, 128, 1, 2), + (2048, 2048, 65536, 64, 64): (1, 1, 64, 64, 3, 2), + (2048, 2048, 65536, 128, 128): (1, 2, 128, 64, 1, 4), + (2048, 2048, 131072, 16, 16): (4, 2, 16, 128, 1, 2), + (2048, 2048, 131072, 32, 32): (1, 1, 32, 128, 3, 2), + (2048, 2048, 131072, 64, 64): (1, 1, 64, 64, 3, 2), + (2048, 2048, 131072, 128, 128): (1, 2, 128, 64, 1, 4), + (4096, 4096, 256, 16, 16): (1, 1, 16, 64, 1, 2), + (4096, 4096, 256, 32, 32): (1, 1, 32, 64, 3, 4), + (4096, 4096, 256, 64, 64): (1, 1, 64, 64, 3, 4), + (4096, 4096, 256, 128, 128): (3, 4, 128, 32, 1, 4), + (4096, 4096, 512, 16, 16): (1, 2, 16, 128, 1, 2), + (4096, 4096, 512, 32, 32): (1, 2, 32, 64, 3, 2), + (4096, 4096, 512, 64, 64): (1, 4, 64, 64, 1, 4), + (4096, 4096, 512, 128, 128): (1, 4, 128, 64, 1, 4), + (4096, 4096, 1024, 16, 16): (1, 2, 16, 128, 1, 2), + (4096, 4096, 1024, 32, 32): (1, 8, 32, 64, 3, 2), + (4096, 4096, 1024, 64, 64): (1, 4, 64, 64, 1, 4), + (4096, 4096, 1024, 128, 128): (2, 4, 128, 64, 1, 4), + (4096, 4096, 2048, 16, 16): (1, 1, 16, 128, 1, 2), + (4096, 4096, 2048, 32, 32): (1, 4, 32, 128, 1, 4), + (4096, 4096, 2048, 64, 64): (1, 1, 64, 64, 3, 4), + (4096, 4096, 2048, 128, 128): (1, 16, 128, 64, 1, 4), + (4096, 4096, 4096, 16, 16): (1, 1, 16, 64, 3, 1), + (4096, 4096, 4096, 32, 32): (1, 4, 32, 64, 3, 2), + (4096, 4096, 4096, 64, 64): (1, 1, 64, 64, 3, 4), + (4096, 4096, 4096, 128, 128): (5, 1, 128, 64, 1, 4), + (4096, 4096, 8192, 16, 16): (1, 1, 16, 128, 1, 2), + (4096, 4096, 8192, 32, 32): (1, 1, 32, 128, 3, 2), + (4096, 4096, 8192, 64, 64): (1, 1, 64, 64, 3, 4), + (4096, 4096, 8192, 128, 128): (2, 1, 128, 64, 1, 4), + (4096, 4096, 16384, 16, 16): (1, 1, 16, 128, 1, 2), + (4096, 4096, 16384, 32, 32): (1, 1, 32, 128, 3, 2), + (4096, 4096, 16384, 64, 64): (1, 1, 64, 64, 4, 4), + (4096, 4096, 16384, 128, 128): (2, 1, 128, 64, 1, 4), + (4096, 4096, 32768, 16, 16): (3, 1, 16, 128, 1, 2), + (4096, 4096, 32768, 32, 32): (1, 1, 32, 128, 3, 2), + (4096, 4096, 32768, 64, 64): (1, 1, 64, 64, 3, 4), + (4096, 4096, 32768, 128, 128): (2, 1, 128, 64, 1, 4), + (4096, 4096, 65536, 16, 16): (2, 2, 16, 128, 1, 2), + (4096, 4096, 65536, 32, 32): (1, 1, 32, 128, 4, 2), + (4096, 4096, 65536, 64, 64): (1, 1, 64, 64, 4, 4), + (4096, 4096, 65536, 128, 128): (2, 1, 128, 64, 1, 4), + (4096, 4096, 131072, 16, 16): (2, 1, 16, 128, 1, 2), + (4096, 4096, 131072, 32, 32): (1, 1, 32, 128, 3, 2), + (4096, 4096, 131072, 64, 64): (1, 1, 64, 64, 3, 4), + (4096, 4096, 131072, 128, 128): (2, 1, 128, 64, 1, 4), + (8192, 8192, 256, 16, 16): (1, 2, 16, 64, 1, 2), + (8192, 8192, 256, 32, 32): (1, 1, 32, 64, 1, 2), + (8192, 8192, 256, 64, 64): (1, 2, 64, 64, 1, 4), + (8192, 8192, 256, 128, 128): (3, 16, 128, 16, 1, 2), + (8192, 8192, 512, 16, 16): (1, 2, 16, 128, 1, 2), + (8192, 8192, 512, 32, 32): (1, 4, 32, 64, 3, 2), + (8192, 8192, 512, 64, 64): (2, 8, 64, 64, 4, 4), + (8192, 8192, 512, 128, 128): (1, 8, 128, 64, 1, 4), + (8192, 8192, 1024, 16, 16): (4, 2, 16, 128, 1, 2), + (8192, 8192, 1024, 32, 32): (1, 8, 32, 128, 1, 2), + (8192, 8192, 1024, 64, 64): (1, 16, 64, 64, 3, 2), + (8192, 8192, 1024, 128, 128): (2, 16, 128, 64, 2, 4), + (8192, 8192, 2048, 16, 16): (2, 1, 16, 64, 4, 1), + (8192, 8192, 2048, 32, 32): (1, 16, 32, 64, 5, 2), + (8192, 8192, 2048, 64, 64): (1, 16, 64, 64, 3, 2), + (8192, 8192, 2048, 128, 128): (2, 16, 128, 64, 2, 4), + (8192, 8192, 4096, 16, 16): (1, 1, 16, 64, 4, 1), + (8192, 8192, 4096, 32, 32): (1, 16, 32, 64, 5, 2), + (8192, 8192, 4096, 64, 64): (1, 16, 64, 64, 3, 2), + (8192, 8192, 4096, 128, 128): (2, 64, 128, 64, 2, 4), + (8192, 8192, 8192, 16, 16): (1, 1, 16, 64, 4, 1), + (8192, 8192, 8192, 32, 32): (1, 8, 32, 128, 5, 4), + (8192, 8192, 8192, 64, 64): (1, 8, 64, 64, 3, 2), + (8192, 8192, 8192, 128, 128): (2, 8, 128, 64, 1, 4), + (8192, 8192, 16384, 16, 16): (1, 1, 16, 64, 4, 1), + (8192, 8192, 16384, 32, 32): (1, 8, 32, 64, 5, 2), + (8192, 8192, 16384, 64, 64): (1, 8, 64, 64, 3, 2), + (8192, 8192, 16384, 128, 128): (1, 8, 128, 64, 1, 4), + (8192, 8192, 32768, 16, 16): (1, 1, 16, 64, 4, 1), + (8192, 8192, 32768, 32, 32): (1, 8, 32, 64, 5, 2), + (8192, 8192, 32768, 64, 64): (3, 8, 64, 64, 3, 2), + (8192, 8192, 32768, 128, 128): (2, 8, 128, 64, 1, 4), + (8192, 8192, 65536, 16, 16): (1, 1, 16, 64, 4, 1), + (8192, 8192, 65536, 32, 32): (5, 4, 32, 64, 3, 2), + (8192, 8192, 65536, 64, 64): (1, 8, 64, 64, 3, 2), + (8192, 8192, 65536, 128, 128): (2, 8, 128, 64, 1, 4), + (8192, 8192, 131072, 16, 16): (2, 1, 16, 64, 4, 1), + (8192, 8192, 131072, 32, 32): (1, 4, 32, 64, 5, 2), + (8192, 8192, 131072, 64, 64): (1, 4, 64, 128, 3, 4), + (8192, 8192, 131072, 128, 128): (2, 8, 128, 64, 1, 4), + (16384, 16384, 256, 16, 16): (1, 2, 16, 128, 1, 2), + (16384, 16384, 256, 32, 32): (1, 4, 32, 64, 3, 2), + (16384, 16384, 256, 64, 64): (2, 4, 64, 64, 4, 4), + (16384, 16384, 256, 128, 128): (1, 4, 128, 64, 1, 16), + (16384, 16384, 512, 16, 16): (1, 2, 16, 128, 3, 2), + (16384, 16384, 512, 32, 32): (1, 4, 32, 128, 5, 4), + (16384, 16384, 512, 64, 64): (1, 8, 64, 64, 3, 2), + (16384, 16384, 512, 128, 128): (2, 8, 128, 64, 1, 4), + (16384, 16384, 1024, 16, 16): (1, 2, 16, 128, 1, 2), + (16384, 16384, 1024, 32, 32): (1, 8, 32, 64, 5, 2), + (16384, 16384, 1024, 64, 64): (1, 16, 64, 64, 3, 2), + (16384, 16384, 1024, 128, 128): (5, 16, 128, 64, 2, 4), + (16384, 16384, 2048, 16, 16): (1, 2, 16, 128, 1, 2), + (16384, 16384, 2048, 32, 32): (1, 8, 32, 64, 5, 2), + (16384, 16384, 2048, 64, 64): (1, 16, 64, 64, 3, 2), + (16384, 16384, 2048, 128, 128): (4, 32, 128, 64, 2, 4), + (16384, 16384, 4096, 16, 16): (3, 2, 16, 128, 1, 2), + (16384, 16384, 4096, 32, 32): (1, 4, 32, 64, 5, 2), + (16384, 16384, 4096, 64, 64): (2, 16, 64, 64, 3, 2), + (16384, 16384, 4096, 128, 128): (3, 32, 128, 64, 2, 4), + (16384, 16384, 8192, 16, 16): (1, 2, 16, 128, 1, 2), + (16384, 16384, 8192, 32, 32): (1, 4, 32, 64, 5, 2), + (16384, 16384, 8192, 64, 64): (4, 8, 64, 64, 3, 2), + (16384, 16384, 8192, 128, 128): (5, 8, 128, 64, 1, 4), + (16384, 16384, 16384, 16, 16): (1, 2, 16, 128, 1, 2), + (16384, 16384, 16384, 32, 32): (1, 4, 32, 64, 5, 2), + (16384, 16384, 16384, 64, 64): (2, 4, 64, 128, 3, 4), + (16384, 16384, 16384, 128, 128): (4, 8, 128, 64, 1, 4), + (16384, 16384, 32768, 16, 16): (4, 2, 16, 128, 1, 2), + (16384, 16384, 32768, 32, 32): (1, 4, 32, 64, 5, 2), + (16384, 16384, 32768, 64, 64): (1, 8, 64, 64, 3, 2), + (16384, 16384, 32768, 128, 128): (2, 512, 128, 64, 2, 4), + (16384, 16384, 65536, 16, 16): (3, 2, 16, 128, 1, 2), + (16384, 16384, 65536, 32, 32): (1, 4, 32, 64, 5, 2), + (16384, 16384, 65536, 64, 64): (1, 4, 64, 128, 3, 4), + (16384, 16384, 65536, 128, 128): (2, 1024, 128, 64, 2, 4), + (16384, 16384, 131072, 16, 16): (1, 2, 16, 128, 1, 2), + (16384, 16384, 131072, 32, 32): (1, 4, 32, 64, 5, 2), + (16384, 16384, 131072, 64, 64): (3, 4, 64, 128, 3, 4), + (16384, 16384, 131072, 128, 128): (4, 2048, 128, 64, 2, 4), + }, + ("scatter_mm", "NVIDIA A100-SXM4-80GB", (0, torch.float16, 0.5)): { + (256, 256, 256, 16, 16): (5, 4, 16, 16, 1, 4), + (256, 256, 256, 32, 32): (5, 2, 32, 16, 1, 4), + (256, 256, 256, 64, 64): (4, 1, 32, 32, 1, 8), + (256, 256, 256, 128, 128): (2, 1, 32, 32, 1, 4), + (256, 256, 512, 16, 16): (2, 2, 16, 32, 1, 4), + (256, 256, 512, 32, 32): (4, 8, 32, 32, 1, 8), + (256, 256, 512, 64, 64): (4, 8, 32, 64, 1, 4), + (256, 256, 512, 128, 128): (4, 8, 32, 64, 1, 4), + (256, 256, 1024, 16, 16): (4, 2, 16, 64, 1, 2), + (256, 256, 1024, 32, 32): (4, 16, 32, 64, 1, 2), + (256, 256, 1024, 64, 64): (4, 16, 32, 64, 1, 4), + (256, 256, 1024, 128, 128): (4, 16, 64, 64, 1, 8), + (256, 256, 2048, 16, 16): (2, 16, 16, 64, 1, 8), + (256, 256, 2048, 32, 32): (4, 16, 32, 64, 1, 2), + (256, 256, 2048, 64, 64): (4, 16, 32, 64, 1, 4), + (256, 256, 2048, 128, 128): (4, 16, 64, 64, 1, 4), + (256, 256, 4096, 16, 16): (4, 32, 16, 64, 1, 1), + (256, 256, 4096, 32, 32): (2, 64, 32, 64, 1, 2), + (256, 256, 4096, 64, 64): (4, 64, 64, 64, 1, 4), + (256, 256, 4096, 128, 128): (4, 32, 64, 64, 1, 4), + (256, 256, 8192, 16, 16): (4, 64, 16, 64, 1, 1), + (256, 256, 8192, 32, 32): (4, 128, 32, 64, 1, 2), + (256, 256, 8192, 64, 64): (4, 64, 64, 64, 1, 4), + (256, 256, 8192, 128, 128): (4, 64, 64, 64, 1, 4), + (256, 256, 16384, 16, 16): (4, 128, 16, 64, 1, 1), + (256, 256, 16384, 32, 32): (2, 128, 32, 64, 1, 2), + (256, 256, 16384, 64, 64): (4, 32, 32, 128, 1, 4), + (256, 256, 16384, 128, 128): (4, 16, 64, 64, 1, 4), + (256, 256, 32768, 16, 16): (4, 64, 16, 64, 1, 1), + (256, 256, 32768, 32, 32): (2, 256, 32, 64, 1, 2), + (256, 256, 32768, 64, 64): (4, 32, 32, 128, 1, 4), + (256, 256, 32768, 128, 128): (4, 32, 64, 64, 1, 4), + (256, 256, 65536, 16, 16): (4, 128, 16, 64, 1, 1), + (256, 256, 65536, 32, 32): (4, 1, 32, 64, 1, 2), + (256, 256, 65536, 64, 64): (2, 1, 64, 64, 1, 2), + (256, 256, 65536, 128, 128): (4, 32, 64, 64, 1, 4), + (256, 256, 131072, 16, 16): (4, 64, 16, 64, 1, 1), + (256, 256, 131072, 32, 32): (2, 1, 32, 64, 1, 2), + (256, 256, 131072, 64, 64): (4, 32, 32, 128, 1, 4), + (256, 256, 131072, 128, 128): (4, 32, 64, 64, 1, 4), + (512, 512, 256, 16, 16): (4, 16, 16, 16, 1, 4), + (512, 512, 256, 32, 32): (2, 4, 32, 16, 1, 4), + (512, 512, 256, 64, 64): (2, 16, 64, 16, 3, 8), + (512, 512, 256, 128, 128): (4, 16, 64, 16, 1, 4), + (512, 512, 512, 16, 16): (1, 1, 16, 64, 1, 8), + (512, 512, 512, 32, 32): (2, 4, 16, 32, 1, 1), + (512, 512, 512, 64, 64): (2, 1, 32, 32, 1, 2), + (512, 512, 512, 128, 128): (4, 8, 32, 64, 1, 4), + (512, 512, 1024, 16, 16): (2, 8, 16, 64, 1, 8), + (512, 512, 1024, 32, 32): (4, 16, 32, 64, 1, 2), + (512, 512, 1024, 64, 64): (4, 16, 64, 64, 1, 4), + (512, 512, 1024, 128, 128): (2, 8, 64, 64, 1, 4), + (512, 512, 2048, 16, 16): (4, 16, 16, 64, 1, 4), + (512, 512, 2048, 32, 32): (4, 16, 32, 64, 1, 2), + (512, 512, 2048, 64, 64): (4, 16, 64, 64, 1, 8), + (512, 512, 2048, 128, 128): (4, 16, 64, 64, 1, 4), + (512, 512, 4096, 16, 16): (4, 32, 16, 128, 1, 2), + (512, 512, 4096, 32, 32): (4, 32, 32, 64, 1, 2), + (512, 512, 4096, 64, 64): (4, 32, 64, 64, 1, 4), + (512, 512, 4096, 128, 128): (4, 32, 64, 64, 1, 4), + (512, 512, 8192, 16, 16): (2, 32, 16, 128, 1, 2), + (512, 512, 8192, 32, 32): (4, 64, 32, 64, 1, 2), + (512, 512, 8192, 64, 64): (4, 128, 64, 64, 1, 2), + (512, 512, 8192, 128, 128): (4, 64, 64, 64, 1, 4), + (512, 512, 16384, 16, 16): (4, 32, 16, 64, 1, 1), + (512, 512, 16384, 32, 32): (4, 64, 32, 64, 1, 2), + (512, 512, 16384, 64, 64): (4, 16, 64, 64, 1, 4), + (512, 512, 16384, 128, 128): (4, 32, 64, 64, 1, 4), + (512, 512, 32768, 16, 16): (7, 16, 16, 128, 1, 2), + (512, 512, 32768, 32, 32): (4, 64, 32, 64, 1, 2), + (512, 512, 32768, 64, 64): (2, 32, 64, 64, 3, 2), + (512, 512, 32768, 128, 128): (2, 32, 64, 64, 1, 4), + (512, 512, 65536, 16, 16): (2, 32, 16, 64, 1, 1), + (512, 512, 65536, 32, 32): (4, 64, 32, 64, 1, 2), + (512, 512, 65536, 64, 64): (3, 32, 64, 64, 3, 2), + (512, 512, 65536, 128, 128): (4, 16, 64, 64, 1, 4), + (512, 512, 131072, 16, 16): (3, 32, 16, 128, 1, 2), + (512, 512, 131072, 32, 32): (4, 64, 32, 64, 1, 2), + (512, 512, 131072, 64, 64): (2, 32, 64, 64, 3, 2), + (512, 512, 131072, 128, 128): (3, 1, 64, 64, 1, 4), + (1024, 1024, 256, 16, 16): (4, 16, 16, 16, 1, 4), + (1024, 1024, 256, 32, 32): (4, 16, 32, 16, 1, 4), + (1024, 1024, 256, 64, 64): (4, 4, 64, 32, 1, 16), + (1024, 1024, 256, 128, 128): (4, 16, 64, 16, 1, 8), + (1024, 1024, 512, 16, 16): (2, 8, 16, 64, 1, 8), + (1024, 1024, 512, 32, 32): (3, 2, 32, 64, 1, 2), + (1024, 1024, 512, 64, 64): (4, 8, 32, 64, 1, 8), + (1024, 1024, 512, 128, 128): (4, 8, 64, 64, 1, 8), + (1024, 1024, 1024, 16, 16): (2, 2, 16, 64, 1, 2), + (1024, 1024, 1024, 32, 32): (2, 8, 32, 64, 1, 2), + (1024, 1024, 1024, 64, 64): (2, 8, 32, 128, 1, 4), + (1024, 1024, 1024, 128, 128): (2, 8, 64, 64, 1, 4), + (1024, 1024, 2048, 16, 16): (2, 16, 16, 128, 3, 2), + (1024, 1024, 2048, 32, 32): (4, 32, 32, 64, 1, 2), + (1024, 1024, 2048, 64, 64): (4, 16, 64, 64, 1, 4), + (1024, 1024, 2048, 128, 128): (4, 32, 64, 64, 1, 4), + (1024, 1024, 4096, 16, 16): (4, 16, 16, 128, 1, 2), + (1024, 1024, 4096, 32, 32): (3, 32, 32, 64, 1, 2), + (1024, 1024, 4096, 64, 64): (4, 32, 64, 64, 1, 4), + (1024, 1024, 4096, 128, 128): (4, 32, 64, 64, 1, 4), + (1024, 1024, 8192, 16, 16): (5, 16, 16, 128, 1, 2), + (1024, 1024, 8192, 32, 32): (2, 32, 32, 64, 3, 2), + (1024, 1024, 8192, 64, 64): (1, 16, 64, 64, 3, 2), + (1024, 1024, 8192, 128, 128): (4, 32, 64, 64, 1, 4), + (1024, 1024, 16384, 16, 16): (4, 16, 16, 128, 1, 2), + (1024, 1024, 16384, 32, 32): (1, 32, 32, 64, 3, 2), + (1024, 1024, 16384, 64, 64): (4, 16, 64, 64, 3, 2), + (1024, 1024, 16384, 128, 128): (4, 32, 128, 64, 1, 4), + (1024, 1024, 32768, 16, 16): (3, 16, 16, 128, 1, 2), + (1024, 1024, 32768, 32, 32): (1, 8, 32, 64, 3, 2), + (1024, 1024, 32768, 64, 64): (4, 16, 64, 64, 3, 2), + (1024, 1024, 32768, 128, 128): (4, 8, 128, 64, 2, 4), + (1024, 1024, 65536, 16, 16): (1, 2, 16, 128, 1, 2), + (1024, 1024, 65536, 32, 32): (2, 4, 32, 64, 3, 2), + (1024, 1024, 65536, 64, 64): (5, 16, 64, 64, 3, 2), + (1024, 1024, 65536, 128, 128): (5, 8, 128, 64, 2, 4), + (1024, 1024, 131072, 16, 16): (5, 2, 16, 128, 1, 2), + (1024, 1024, 131072, 32, 32): (1, 2, 32, 64, 3, 2), + (1024, 1024, 131072, 64, 64): (5, 16, 64, 64, 3, 2), + (1024, 1024, 131072, 128, 128): (2, 1, 128, 64, 2, 4), + (2048, 2048, 256, 16, 16): (4, 4, 16, 64, 1, 8), + (2048, 2048, 256, 32, 32): (4, 8, 32, 32, 1, 8), + (2048, 2048, 256, 64, 64): (4, 16, 64, 16, 1, 8), + (2048, 2048, 256, 128, 128): (4, 4, 128, 32, 3, 8), + (2048, 2048, 512, 16, 16): (2, 2, 16, 64, 1, 2), + (2048, 2048, 512, 32, 32): (2, 4, 32, 64, 3, 2), + (2048, 2048, 512, 64, 64): (4, 4, 64, 64, 1, 8), + (2048, 2048, 512, 128, 128): (4, 8, 64, 64, 1, 4), + (2048, 2048, 1024, 16, 16): (1, 8, 16, 64, 1, 2), + (2048, 2048, 1024, 32, 32): (2, 16, 32, 64, 3, 2), + (2048, 2048, 1024, 64, 64): (4, 8, 64, 64, 1, 4), + (2048, 2048, 1024, 128, 128): (4, 8, 128, 64, 1, 4), + (2048, 2048, 2048, 16, 16): (5, 4, 16, 128, 1, 2), + (2048, 2048, 2048, 32, 32): (1, 16, 32, 64, 3, 2), + (2048, 2048, 2048, 64, 64): (2, 8, 64, 64, 1, 4), + (2048, 2048, 2048, 128, 128): (2, 8, 128, 64, 1, 4), + (2048, 2048, 4096, 16, 16): (4, 2, 16, 128, 1, 2), + (2048, 2048, 4096, 32, 32): (2, 16, 32, 64, 3, 2), + (2048, 2048, 4096, 64, 64): (2, 8, 64, 64, 3, 2), + (2048, 2048, 4096, 128, 128): (4, 8, 128, 64, 1, 4), + (2048, 2048, 8192, 16, 16): (5, 4, 16, 128, 1, 2), + (2048, 2048, 8192, 32, 32): (2, 8, 32, 64, 3, 2), + (2048, 2048, 8192, 64, 64): (4, 8, 64, 64, 3, 2), + (2048, 2048, 8192, 128, 128): (4, 8, 128, 64, 1, 4), + (2048, 2048, 16384, 16, 16): (3, 2, 16, 128, 1, 2), + (2048, 2048, 16384, 32, 32): (2, 4, 32, 128, 3, 2), + (2048, 2048, 16384, 64, 64): (4, 8, 64, 64, 3, 2), + (2048, 2048, 16384, 128, 128): (4, 4, 128, 64, 1, 4), + (2048, 2048, 32768, 16, 16): (3, 2, 16, 128, 1, 2), + (2048, 2048, 32768, 32, 32): (3, 4, 32, 128, 3, 2), + (2048, 2048, 32768, 64, 64): (6, 4, 64, 64, 3, 2), + (2048, 2048, 32768, 128, 128): (3, 4, 128, 64, 1, 4), + (2048, 2048, 65536, 16, 16): (6, 2, 16, 128, 1, 2), + (2048, 2048, 65536, 32, 32): (1, 2, 32, 128, 1, 2), + (2048, 2048, 65536, 64, 64): (5, 4, 64, 64, 3, 2), + (2048, 2048, 65536, 128, 128): (5, 1, 128, 64, 2, 4), + (2048, 2048, 131072, 16, 16): (3, 2, 16, 128, 1, 2), + (2048, 2048, 131072, 32, 32): (2, 1, 32, 128, 3, 2), + (2048, 2048, 131072, 64, 64): (4, 1, 64, 64, 3, 2), + (2048, 2048, 131072, 128, 128): (3, 1, 128, 64, 2, 4), + (4096, 4096, 256, 16, 16): (5, 8, 16, 32, 1, 4), + (4096, 4096, 256, 32, 32): (4, 16, 32, 16, 2, 4), + (4096, 4096, 256, 64, 64): (2, 1, 64, 64, 3, 4), + (4096, 4096, 256, 128, 128): (4, 4, 128, 32, 1, 4), + (4096, 4096, 512, 16, 16): (4, 2, 16, 128, 1, 2), + (4096, 4096, 512, 32, 32): (4, 8, 32, 64, 1, 2), + (4096, 4096, 512, 64, 64): (4, 4, 64, 64, 1, 4), + (4096, 4096, 512, 128, 128): (4, 8, 128, 64, 2, 4), + (4096, 4096, 1024, 16, 16): (1, 2, 16, 128, 1, 2), + (4096, 4096, 1024, 32, 32): (6, 8, 32, 64, 3, 2), + (4096, 4096, 1024, 64, 64): (2, 16, 64, 64, 4, 4), + (4096, 4096, 1024, 128, 128): (2, 4, 128, 64, 2, 4), + (4096, 4096, 2048, 16, 16): (3, 1, 16, 128, 1, 2), + (4096, 4096, 2048, 32, 32): (1, 4, 32, 64, 5, 2), + (4096, 4096, 2048, 64, 64): (3, 16, 64, 64, 3, 2), + (4096, 4096, 2048, 128, 128): (4, 32, 128, 64, 2, 4), + (4096, 4096, 4096, 16, 16): (1, 2, 16, 128, 1, 2), + (4096, 4096, 4096, 32, 32): (1, 4, 32, 64, 3, 2), + (4096, 4096, 4096, 64, 64): (1, 1, 64, 64, 4, 4), + (4096, 4096, 4096, 128, 128): (2, 1, 128, 128, 1, 8), + (4096, 4096, 8192, 16, 16): (3, 1, 16, 128, 1, 2), + (4096, 4096, 8192, 32, 32): (2, 2, 32, 64, 5, 2), + (4096, 4096, 8192, 64, 64): (4, 16, 64, 64, 3, 2), + (4096, 4096, 8192, 128, 128): (4, 16, 128, 64, 2, 4), + (4096, 4096, 16384, 16, 16): (1, 2, 16, 128, 1, 2), + (4096, 4096, 16384, 32, 32): (4, 2, 32, 64, 5, 2), + (4096, 4096, 16384, 64, 64): (4, 16, 64, 64, 3, 2), + (4096, 4096, 16384, 128, 128): (4, 16, 128, 64, 2, 4), + (4096, 4096, 32768, 16, 16): (3, 1, 16, 128, 1, 2), + (4096, 4096, 32768, 32, 32): (3, 1, 32, 128, 1, 4), + (4096, 4096, 32768, 64, 64): (3, 1, 64, 64, 3, 4), + (4096, 4096, 32768, 128, 128): (5, 16, 128, 64, 2, 4), + (4096, 4096, 65536, 16, 16): (5, 1, 16, 128, 1, 2), + (4096, 4096, 65536, 32, 32): (5, 1, 32, 128, 1, 4), + (4096, 4096, 65536, 64, 64): (1, 1, 64, 64, 3, 4), + (4096, 4096, 65536, 128, 128): (3, 16, 128, 64, 2, 4), + (4096, 4096, 131072, 16, 16): (3, 1, 16, 128, 1, 2), + (4096, 4096, 131072, 32, 32): (3, 1, 32, 128, 3, 2), + (4096, 4096, 131072, 64, 64): (2, 1, 64, 64, 3, 4), + (4096, 4096, 131072, 128, 128): (1, 1, 128, 64, 1, 4), + (8192, 8192, 256, 16, 16): (4, 16, 16, 16, 1, 4), + (8192, 8192, 256, 32, 32): (1, 16, 32, 16, 4, 4), + (8192, 8192, 256, 64, 64): (4, 16, 64, 16, 3, 8), + (8192, 8192, 256, 128, 128): (4, 16, 128, 16, 1, 2), + (8192, 8192, 512, 16, 16): (2, 8, 16, 64, 1, 4), + (8192, 8192, 512, 32, 32): (4, 8, 32, 64, 3, 2), + (8192, 8192, 512, 64, 64): (2, 8, 64, 64, 4, 4), + (8192, 8192, 512, 128, 128): (4, 8, 128, 64, 2, 4), + (8192, 8192, 1024, 16, 16): (4, 16, 16, 64, 1, 8), + (8192, 8192, 1024, 32, 32): (2, 8, 32, 64, 5, 2), + (8192, 8192, 1024, 64, 64): (1, 16, 64, 64, 3, 2), + (8192, 8192, 1024, 128, 128): (5, 16, 128, 64, 2, 4), + (8192, 8192, 2048, 16, 16): (7, 2, 16, 128, 1, 2), + (8192, 8192, 2048, 32, 32): (1, 16, 32, 64, 5, 2), + (8192, 8192, 2048, 64, 64): (4, 16, 64, 64, 3, 2), + (8192, 8192, 2048, 128, 128): (6, 16, 128, 64, 2, 4), + (8192, 8192, 4096, 16, 16): (4, 2, 16, 128, 1, 2), + (8192, 8192, 4096, 32, 32): (2, 8, 32, 64, 5, 2), + (8192, 8192, 4096, 64, 64): (3, 16, 64, 64, 3, 2), + (8192, 8192, 4096, 128, 128): (3, 64, 128, 64, 2, 4), + (8192, 8192, 8192, 16, 16): (4, 2, 16, 128, 1, 2), + (8192, 8192, 8192, 32, 32): (1, 4, 32, 128, 5, 4), + (8192, 8192, 8192, 64, 64): (4, 4, 64, 64, 1, 4), + (8192, 8192, 8192, 128, 128): (2, 2, 128, 128, 3, 8), + (8192, 8192, 16384, 16, 16): (1, 2, 16, 128, 1, 2), + (8192, 8192, 16384, 32, 32): (4, 8, 32, 64, 5, 2), + (8192, 8192, 16384, 64, 64): (5, 8, 64, 64, 3, 2), + (8192, 8192, 16384, 128, 128): (3, 16, 128, 64, 2, 4), + (8192, 8192, 32768, 16, 16): (7, 2, 16, 128, 1, 2), + (8192, 8192, 32768, 32, 32): (3, 4, 32, 64, 3, 2), + (8192, 8192, 32768, 64, 64): (2, 8, 64, 64, 3, 2), + (8192, 8192, 32768, 128, 128): (6, 16, 128, 64, 2, 4), + (8192, 8192, 65536, 16, 16): (9, 2, 16, 128, 1, 2), + (8192, 8192, 65536, 32, 32): (7, 4, 32, 64, 5, 2), + (8192, 8192, 65536, 64, 64): (4, 8, 64, 64, 3, 2), + (8192, 8192, 65536, 128, 128): (3, 16, 128, 64, 2, 4), + (8192, 8192, 131072, 16, 16): (9, 2, 16, 128, 1, 2), + (8192, 8192, 131072, 32, 32): (1, 8, 32, 64, 5, 2), + (8192, 8192, 131072, 64, 64): (1, 8, 64, 64, 3, 2), + (8192, 8192, 131072, 128, 128): (4, 16, 128, 64, 2, 4), + (16384, 16384, 256, 16, 16): (5, 16, 16, 16, 1, 4), + (16384, 16384, 256, 32, 32): (4, 16, 32, 16, 4, 4), + (16384, 16384, 256, 64, 64): (4, 16, 64, 16, 3, 8), + (16384, 16384, 256, 128, 128): (4, 16, 128, 16, 1, 2), + (16384, 16384, 512, 16, 16): (2, 8, 16, 64, 1, 4), + (16384, 16384, 512, 32, 32): (1, 4, 32, 64, 5, 2), + (16384, 16384, 512, 64, 64): (4, 8, 64, 64, 1, 4), + (16384, 16384, 512, 128, 128): (3, 8, 128, 64, 2, 4), + (16384, 16384, 1024, 16, 16): (4, 2, 16, 128, 1, 2), + (16384, 16384, 1024, 32, 32): (4, 8, 32, 64, 5, 2), + (16384, 16384, 1024, 64, 64): (6, 16, 64, 64, 3, 2), + (16384, 16384, 1024, 128, 128): (3, 16, 128, 64, 2, 4), + (16384, 16384, 2048, 16, 16): (3, 2, 16, 128, 1, 2), + (16384, 16384, 2048, 32, 32): (1, 8, 32, 64, 5, 2), + (16384, 16384, 2048, 64, 64): (5, 16, 64, 64, 3, 2), + (16384, 16384, 2048, 128, 128): (2, 32, 128, 64, 2, 4), + (16384, 16384, 4096, 16, 16): (2, 2, 16, 128, 1, 2), + (16384, 16384, 4096, 32, 32): (1, 4, 32, 64, 3, 2), + (16384, 16384, 4096, 64, 64): (2, 8, 64, 64, 3, 2), + (16384, 16384, 4096, 128, 128): (3, 16, 128, 64, 2, 4), + (16384, 16384, 8192, 16, 16): (3, 2, 16, 128, 1, 2), + (16384, 16384, 8192, 32, 32): (2, 4, 32, 64, 5, 2), + (16384, 16384, 8192, 64, 64): (4, 8, 64, 64, 3, 2), + (16384, 16384, 8192, 128, 128): (8, 32, 128, 64, 2, 4), + (16384, 16384, 16384, 16, 16): (1, 2, 16, 256, 1, 4), + (16384, 16384, 16384, 32, 32): (1, 4, 32, 128, 3, 4), + (16384, 16384, 16384, 64, 64): (5, 4, 64, 64, 1, 4), + (16384, 16384, 16384, 128, 128): (4, 8, 128, 64, 2, 4), + (16384, 16384, 32768, 16, 16): (2, 2, 16, 128, 1, 2), + (16384, 16384, 32768, 32, 32): (1, 4, 32, 64, 3, 2), + (16384, 16384, 32768, 64, 64): (5, 4, 64, 64, 1, 4), + (16384, 16384, 32768, 128, 128): (5, 8, 128, 64, 2, 4), + (16384, 16384, 65536, 16, 16): (8, 2, 16, 128, 1, 2), + (16384, 16384, 65536, 32, 32): (6, 4, 32, 64, 5, 2), + (16384, 16384, 65536, 64, 64): (2, 4, 64, 64, 1, 4), + (16384, 16384, 65536, 128, 128): (4, 8, 128, 64, 2, 4), + (16384, 16384, 131072, 16, 16): (3, 1, 16, 128, 1, 2), + (16384, 16384, 131072, 32, 32): (1, 4, 32, 64, 3, 2), + (16384, 16384, 131072, 64, 64): (4, 4, 64, 64, 1, 4), + (16384, 16384, 131072, 128, 128): (1, 8, 128, 64, 2, 4), + (32768, 32768, 256, 16, 16): (4, 16, 16, 16, 1, 4), + (32768, 32768, 512, 16, 16): (4, 2, 16, 128, 1, 2), + (32768, 32768, 1024, 16, 16): (3, 2, 16, 128, 1, 2), + (32768, 32768, 2048, 16, 16): (4, 2, 16, 128, 1, 2), + (32768, 32768, 4096, 16, 16): (5, 4, 16, 64, 1, 1), + (32768, 32768, 8192, 16, 16): (4, 4, 16, 64, 1, 1), + (32768, 32768, 16384, 16, 16): (4, 4, 16, 64, 1, 1), + (32768, 32768, 32768, 16, 16): (5, 4, 16, 64, 1, 1), + }, + ("scatter_mm", "NVIDIA A100-SXM4-80GB", (0, torch.float32, 0.5)): { + (256, 256, 256, 16, 16): (1, 1, 16, 16, 1, 8), + (256, 256, 256, 32, 32): (1, 1, 16, 16, 1, 4), + (256, 256, 256, 64, 64): (1, 1, 16, 16, 1, 4), + (256, 256, 256, 128, 128): (1, 1, 16, 16, 1, 1), + (256, 256, 512, 16, 16): (1, 1, 16, 16, 1, 4), + (256, 256, 512, 32, 32): (1, 16, 16, 16, 1, 1), + (256, 256, 512, 64, 64): (1, 1, 16, 16, 1, 1), + (256, 256, 512, 128, 128): (1, 1, 32, 32, 1, 4), + (256, 256, 1024, 16, 16): (1, 1, 16, 32, 1, 2), + (256, 256, 1024, 32, 32): (1, 4, 16, 16, 1, 1), + (256, 256, 1024, 64, 64): (1, 1, 32, 32, 1, 4), + (256, 256, 1024, 128, 128): (1, 1, 32, 32, 1, 4), + (256, 256, 2048, 16, 16): (1, 2, 16, 32, 1, 2), + (256, 256, 2048, 32, 32): (1, 1, 16, 32, 1, 2), + (256, 256, 2048, 64, 64): (2, 1, 16, 32, 1, 2), + (256, 256, 2048, 128, 128): (1, 1, 16, 16, 1, 1), + (256, 256, 4096, 16, 16): (1, 1, 16, 32, 1, 2), + (256, 256, 4096, 32, 32): (1, 1, 16, 32, 1, 2), + (256, 256, 4096, 64, 64): (1, 1, 32, 32, 1, 4), + (256, 256, 4096, 128, 128): (3, 1, 32, 64, 1, 4), + (256, 256, 8192, 16, 16): (1, 32, 16, 64, 1, 2), + (256, 256, 8192, 32, 32): (1, 1, 32, 64, 1, 4), + (256, 256, 8192, 64, 64): (1, 1, 32, 64, 1, 4), + (256, 256, 8192, 128, 128): (2, 1, 64, 32, 1, 4), + (256, 256, 16384, 16, 16): (1, 1, 16, 64, 1, 2), + (256, 256, 16384, 32, 32): (1, 1, 32, 64, 1, 4), + (256, 256, 16384, 64, 64): (1, 128, 64, 64, 1, 4), + (256, 256, 16384, 128, 128): (2, 1, 64, 32, 1, 4), + (256, 256, 32768, 16, 16): (2, 128, 16, 64, 1, 1), + (256, 256, 32768, 32, 32): (1, 1, 32, 64, 1, 4), + (256, 256, 32768, 64, 64): (1, 128, 64, 64, 1, 4), + (256, 256, 32768, 128, 128): (2, 1, 64, 64, 1, 4), + (256, 256, 65536, 16, 16): (1, 1, 16, 64, 1, 2), + (256, 256, 65536, 32, 32): (1, 1, 32, 64, 1, 4), + (256, 256, 65536, 64, 64): (2, 1, 64, 64, 1, 4), + (256, 256, 65536, 128, 128): (1, 1, 128, 32, 1, 4), + (256, 256, 131072, 16, 16): (3, 128, 16, 64, 1, 1), + (256, 256, 131072, 32, 32): (1, 1, 32, 64, 1, 4), + (256, 256, 131072, 64, 64): (2, 1, 64, 64, 1, 4), + (256, 256, 131072, 128, 128): (1, 8192, 64, 16, 1, 4), + (512, 512, 256, 16, 16): (1, 2, 16, 16, 1, 1), + (512, 512, 256, 32, 32): (1, 4, 16, 16, 1, 1), + (512, 512, 256, 64, 64): (1, 16, 16, 16, 1, 1), + (512, 512, 256, 128, 128): (1, 1, 16, 32, 1, 4), + (512, 512, 512, 16, 16): (1, 8, 16, 32, 1, 2), + (512, 512, 512, 32, 32): (1, 8, 16, 32, 1, 2), + (512, 512, 512, 64, 64): (1, 2, 16, 32, 1, 2), + (512, 512, 512, 128, 128): (1, 1, 32, 32, 1, 4), + (512, 512, 1024, 16, 16): (1, 1, 16, 32, 1, 2), + (512, 512, 1024, 32, 32): (1, 1, 16, 32, 1, 2), + (512, 512, 1024, 64, 64): (1, 1, 16, 32, 1, 2), + (512, 512, 1024, 128, 128): (1, 1, 64, 32, 1, 4), + (512, 512, 2048, 16, 16): (1, 16, 16, 64, 1, 2), + (512, 512, 2048, 32, 32): (1, 1, 32, 32, 1, 4), + (512, 512, 2048, 64, 64): (1, 1, 32, 32, 1, 4), + (512, 512, 2048, 128, 128): (2, 1, 32, 32, 1, 4), + (512, 512, 4096, 16, 16): (2, 64, 16, 64, 1, 1), + (512, 512, 4096, 32, 32): (1, 64, 32, 64, 1, 4), + (512, 512, 4096, 64, 64): (1, 1, 32, 32, 1, 4), + (512, 512, 4096, 128, 128): (1, 1, 64, 32, 1, 4), + (512, 512, 8192, 16, 16): (2, 64, 16, 64, 1, 1), + (512, 512, 8192, 32, 32): (1, 256, 32, 32, 1, 1), + (512, 512, 8192, 64, 64): (1, 64, 64, 64, 1, 4), + (512, 512, 8192, 128, 128): (2, 1, 64, 32, 1, 8), + (512, 512, 16384, 16, 16): (2, 64, 16, 64, 1, 1), + (512, 512, 16384, 32, 32): (1, 128, 32, 32, 1, 1), + (512, 512, 16384, 64, 64): (1, 64, 64, 64, 1, 4), + (512, 512, 16384, 128, 128): (3, 1, 64, 32, 1, 8), + (512, 512, 32768, 16, 16): (2, 64, 16, 64, 1, 1), + (512, 512, 32768, 32, 32): (1, 128, 32, 32, 1, 1), + (512, 512, 32768, 64, 64): (1, 64, 64, 64, 1, 4), + (512, 512, 32768, 128, 128): (2, 1, 64, 32, 1, 8), + (512, 512, 65536, 16, 16): (2, 32, 16, 64, 1, 1), + (512, 512, 65536, 32, 32): (1, 128, 32, 32, 1, 1), + (512, 512, 65536, 64, 64): (1, 64, 64, 64, 1, 4), + (512, 512, 65536, 128, 128): (2, 1, 64, 32, 1, 8), + (512, 512, 131072, 16, 16): (2, 32, 16, 64, 1, 1), + (512, 512, 131072, 32, 32): (1, 128, 32, 32, 1, 1), + (512, 512, 131072, 64, 64): (3, 64, 64, 64, 1, 4), + (512, 512, 131072, 128, 128): (1, 8192, 64, 16, 1, 4), + (1024, 1024, 256, 16, 16): (1, 4, 16, 32, 1, 2), + (1024, 1024, 256, 32, 32): (1, 4, 16, 32, 1, 2), + (1024, 1024, 256, 64, 64): (1, 1, 16, 32, 1, 2), + (1024, 1024, 256, 128, 128): (1, 1, 16, 16, 1, 1), + (1024, 1024, 512, 16, 16): (1, 8, 16, 32, 1, 2), + (1024, 1024, 512, 32, 32): (1, 8, 16, 32, 1, 1), + (1024, 1024, 512, 64, 64): (1, 8, 32, 32, 1, 4), + (1024, 1024, 512, 128, 128): (2, 1, 32, 32, 1, 4), + (1024, 1024, 1024, 16, 16): (1, 16, 16, 32, 1, 2), + (1024, 1024, 1024, 32, 32): (1, 16, 32, 64, 1, 4), + (1024, 1024, 1024, 64, 64): (1, 16, 32, 64, 1, 4), + (1024, 1024, 1024, 128, 128): (1, 1, 32, 32, 1, 4), + (1024, 1024, 2048, 16, 16): (2, 32, 16, 64, 1, 1), + (1024, 1024, 2048, 32, 32): (1, 32, 32, 64, 1, 4), + (1024, 1024, 2048, 64, 64): (1, 32, 64, 64, 1, 4), + (1024, 1024, 2048, 128, 128): (1, 1, 32, 64, 1, 4), + (1024, 1024, 4096, 16, 16): (2, 16, 16, 64, 1, 1), + (1024, 1024, 4096, 32, 32): (1, 64, 32, 32, 1, 1), + (1024, 1024, 4096, 64, 64): (1, 64, 64, 64, 1, 4), + (1024, 1024, 4096, 128, 128): (2, 64, 64, 32, 1, 8), + (1024, 1024, 8192, 16, 16): (2, 16, 16, 64, 1, 1), + (1024, 1024, 8192, 32, 32): (1, 64, 32, 32, 1, 1), + (1024, 1024, 8192, 64, 64): (1, 64, 64, 64, 1, 4), + (1024, 1024, 8192, 128, 128): (4, 1, 32, 64, 1, 4), + (1024, 1024, 16384, 16, 16): (2, 16, 16, 64, 1, 1), + (1024, 1024, 16384, 32, 32): (1, 64, 32, 32, 1, 1), + (1024, 1024, 16384, 64, 64): (1, 32, 64, 64, 1, 4), + (1024, 1024, 16384, 128, 128): (2, 64, 64, 32, 1, 4), + (1024, 1024, 32768, 16, 16): (2, 16, 16, 64, 1, 1), + (1024, 1024, 32768, 32, 32): (1, 64, 32, 32, 1, 1), + (1024, 1024, 32768, 64, 64): (1, 32, 64, 64, 1, 4), + (1024, 1024, 32768, 128, 128): (4, 1, 32, 64, 1, 4), + (1024, 1024, 65536, 16, 16): (2, 16, 16, 64, 1, 1), + (1024, 1024, 65536, 32, 32): (1, 32, 32, 32, 1, 1), + (1024, 1024, 65536, 64, 64): (2, 32, 64, 64, 1, 4), + (1024, 1024, 65536, 128, 128): (4, 1, 64, 32, 1, 4), + (1024, 1024, 131072, 16, 16): (2, 16, 16, 64, 1, 1), + (1024, 1024, 131072, 32, 32): (1, 32, 32, 32, 1, 1), + (1024, 1024, 131072, 64, 64): (1, 16, 64, 64, 1, 4), + (1024, 1024, 131072, 128, 128): (1, 8192, 64, 16, 1, 4), + (2048, 2048, 256, 16, 16): (1, 4, 16, 32, 1, 2), + (2048, 2048, 256, 32, 32): (1, 8, 16, 32, 1, 1), + (2048, 2048, 256, 64, 64): (1, 8, 32, 32, 1, 4), + (2048, 2048, 256, 128, 128): (1, 4, 64, 64, 1, 8), + (2048, 2048, 512, 16, 16): (2, 8, 16, 32, 1, 2), + (2048, 2048, 512, 32, 32): (2, 8, 32, 64, 1, 4), + (2048, 2048, 512, 64, 64): (2, 4, 64, 64, 1, 4), + (2048, 2048, 512, 128, 128): (1, 8, 32, 64, 1, 4), + (2048, 2048, 1024, 16, 16): (2, 16, 16, 64, 3, 1), + (2048, 2048, 1024, 32, 32): (1, 32, 32, 32, 1, 1), + (2048, 2048, 1024, 64, 64): (1, 16, 64, 64, 1, 4), + (2048, 2048, 1024, 128, 128): (2, 4, 64, 64, 1, 8), + (2048, 2048, 2048, 16, 16): (2, 16, 16, 64, 1, 1), + (2048, 2048, 2048, 32, 32): (1, 32, 32, 32, 1, 1), + (2048, 2048, 2048, 64, 64): (1, 16, 64, 64, 1, 4), + (2048, 2048, 2048, 128, 128): (2, 32, 32, 64, 1, 4), + (2048, 2048, 4096, 16, 16): (3, 2, 16, 64, 1, 1), + (2048, 2048, 4096, 32, 32): (3, 4, 32, 32, 1, 1), + (2048, 2048, 4096, 64, 64): (1, 16, 64, 64, 1, 4), + (2048, 2048, 4096, 128, 128): (2, 32, 64, 32, 1, 4), + (2048, 2048, 8192, 16, 16): (3, 4, 16, 64, 1, 1), + (2048, 2048, 8192, 32, 32): (2, 4, 32, 32, 1, 1), + (2048, 2048, 8192, 64, 64): (2, 32, 64, 32, 1, 2), + (2048, 2048, 8192, 128, 128): (4, 1, 32, 64, 1, 4), + (2048, 2048, 16384, 16, 16): (3, 4, 16, 64, 1, 1), + (2048, 2048, 16384, 32, 32): (1, 4, 32, 32, 1, 1), + (2048, 2048, 16384, 64, 64): (2, 8, 64, 32, 1, 2), + (2048, 2048, 16384, 128, 128): (2, 8, 64, 32, 1, 4), + (2048, 2048, 32768, 16, 16): (2, 4, 16, 64, 1, 1), + (2048, 2048, 32768, 32, 32): (2, 8, 32, 32, 1, 1), + (2048, 2048, 32768, 64, 64): (1, 16, 64, 32, 1, 2), + (2048, 2048, 32768, 128, 128): (4, 1, 32, 64, 1, 4), + (2048, 2048, 65536, 16, 16): (3, 4, 16, 64, 1, 1), + (2048, 2048, 65536, 32, 32): (1, 8, 32, 32, 1, 1), + (2048, 2048, 65536, 64, 64): (1, 8, 64, 32, 1, 2), + (2048, 2048, 65536, 128, 128): (4, 1, 64, 32, 1, 4), + (2048, 2048, 131072, 16, 16): (2, 4, 16, 64, 1, 1), + (2048, 2048, 131072, 32, 32): (1, 8, 32, 32, 1, 1), + (2048, 2048, 131072, 64, 64): (3, 1, 64, 32, 1, 2), + (2048, 2048, 131072, 128, 128): (1, 8192, 128, 16, 1, 8), + (4096, 4096, 256, 16, 16): (2, 4, 16, 32, 1, 2), + (4096, 4096, 256, 32, 32): (1, 4, 32, 64, 1, 4), + (4096, 4096, 256, 64, 64): (1, 4, 64, 64, 1, 4), + (4096, 4096, 256, 128, 128): (1, 4, 32, 64, 1, 4), + (4096, 4096, 512, 16, 16): (2, 8, 16, 64, 3, 1), + (4096, 4096, 512, 32, 32): (2, 16, 32, 32, 1, 1), + (4096, 4096, 512, 64, 64): (1, 8, 64, 64, 1, 4), + (4096, 4096, 512, 128, 128): (1, 8, 32, 64, 1, 4), + (4096, 4096, 1024, 16, 16): (1, 8, 16, 64, 3, 1), + (4096, 4096, 1024, 32, 32): (1, 16, 32, 32, 1, 1), + (4096, 4096, 1024, 64, 64): (1, 16, 64, 32, 1, 2), + (4096, 4096, 1024, 128, 128): (1, 16, 32, 64, 1, 4), + (4096, 4096, 2048, 16, 16): (1, 16, 16, 64, 3, 1), + (4096, 4096, 2048, 32, 32): (1, 16, 32, 32, 1, 1), + (4096, 4096, 2048, 64, 64): (3, 16, 64, 32, 1, 2), + (4096, 4096, 2048, 128, 128): (4, 8, 32, 64, 1, 4), + (4096, 4096, 4096, 16, 16): (1, 8, 16, 64, 3, 1), + (4096, 4096, 4096, 32, 32): (1, 1, 32, 32, 1, 1), + (4096, 4096, 4096, 64, 64): (2, 16, 64, 32, 1, 2), + (4096, 4096, 4096, 128, 128): (4, 8, 32, 64, 1, 4), + (4096, 4096, 8192, 16, 16): (1, 8, 16, 64, 3, 1), + (4096, 4096, 8192, 32, 32): (2, 1, 32, 32, 1, 1), + (4096, 4096, 8192, 64, 64): (1, 16, 64, 32, 1, 2), + (4096, 4096, 8192, 128, 128): (2, 1, 32, 64, 1, 4), + (4096, 4096, 16384, 16, 16): (1, 8, 16, 64, 3, 1), + (4096, 4096, 16384, 32, 32): (1, 1, 32, 32, 1, 1), + (4096, 4096, 16384, 64, 64): (2, 8, 64, 32, 1, 2), + (4096, 4096, 16384, 128, 128): (2, 1, 32, 64, 1, 4), + (4096, 4096, 32768, 16, 16): (1, 8, 16, 64, 3, 1), + (4096, 4096, 32768, 32, 32): (1, 1, 32, 32, 1, 1), + (4096, 4096, 32768, 64, 64): (1, 8, 64, 32, 1, 2), + (4096, 4096, 32768, 128, 128): (2, 1, 32, 64, 1, 4), + (4096, 4096, 65536, 16, 16): (1, 8, 16, 64, 3, 1), + (4096, 4096, 65536, 32, 32): (3, 1, 32, 32, 1, 1), + (4096, 4096, 65536, 64, 64): (3, 4, 64, 32, 1, 2), + (4096, 4096, 65536, 128, 128): (2, 1, 32, 64, 1, 4), + (4096, 4096, 131072, 16, 16): (1, 8, 16, 64, 3, 1), + (4096, 4096, 131072, 32, 32): (1, 1, 32, 32, 1, 1), + (4096, 4096, 131072, 64, 64): (2, 8, 64, 32, 1, 2), + (4096, 4096, 131072, 128, 128): (1, 8192, 128, 16, 1, 8), + (8192, 8192, 256, 16, 16): (2, 4, 16, 64, 3, 1), + (8192, 8192, 256, 32, 32): (1, 8, 32, 32, 1, 1), + (8192, 8192, 256, 64, 64): (1, 4, 64, 64, 1, 4), + (8192, 8192, 256, 128, 128): (1, 4, 32, 64, 1, 4), + (8192, 8192, 512, 16, 16): (1, 4, 16, 64, 3, 1), + (8192, 8192, 512, 32, 32): (1, 16, 32, 32, 1, 1), + (8192, 8192, 512, 64, 64): (2, 4, 64, 64, 1, 4), + (8192, 8192, 512, 128, 128): (2, 1, 32, 64, 1, 4), + (8192, 8192, 1024, 16, 16): (3, 8, 16, 64, 3, 1), + (8192, 8192, 1024, 32, 32): (1, 16, 32, 32, 1, 1), + (8192, 8192, 1024, 64, 64): (1, 8, 64, 32, 1, 2), + (8192, 8192, 1024, 128, 128): (2, 4, 32, 64, 1, 4), + (8192, 8192, 2048, 16, 16): (1, 8, 16, 64, 3, 1), + (8192, 8192, 2048, 32, 32): (1, 16, 32, 32, 1, 1), + (8192, 8192, 2048, 64, 64): (2, 8, 64, 32, 1, 2), + (8192, 8192, 2048, 128, 128): (4, 1, 32, 64, 1, 4), + (8192, 8192, 4096, 16, 16): (1, 8, 16, 64, 3, 1), + (8192, 8192, 4096, 32, 32): (1, 16, 32, 32, 1, 1), + (8192, 8192, 4096, 64, 64): (1, 4, 64, 32, 1, 2), + (8192, 8192, 4096, 128, 128): (3, 1, 32, 64, 1, 4), + (8192, 8192, 8192, 16, 16): (1, 8, 16, 64, 3, 1), + (8192, 8192, 8192, 32, 32): (1, 8, 32, 32, 1, 1), + (8192, 8192, 8192, 64, 64): (1, 8, 64, 32, 1, 2), + (8192, 8192, 8192, 128, 128): (4, 1, 32, 64, 1, 4), + (8192, 8192, 16384, 16, 16): (3, 4, 16, 64, 3, 1), + (8192, 8192, 16384, 32, 32): (1, 8, 32, 32, 1, 1), + (8192, 8192, 16384, 64, 64): (2, 2, 64, 32, 1, 2), + (8192, 8192, 16384, 128, 128): (7, 1, 32, 64, 1, 4), + (8192, 8192, 32768, 16, 16): (1, 4, 16, 64, 3, 1), + (8192, 8192, 32768, 32, 32): (1, 8, 32, 32, 1, 1), + (8192, 8192, 32768, 64, 64): (3, 2, 64, 32, 1, 2), + (8192, 8192, 32768, 128, 128): (6, 1, 32, 64, 1, 4), + (8192, 8192, 65536, 16, 16): (1, 4, 16, 64, 3, 1), + (8192, 8192, 65536, 32, 32): (4, 8, 32, 32, 1, 1), + (8192, 8192, 65536, 64, 64): (1, 2, 64, 32, 1, 2), + (8192, 8192, 65536, 128, 128): (4, 1, 32, 64, 1, 4), + (8192, 8192, 131072, 16, 16): (1, 4, 16, 64, 3, 1), + (8192, 8192, 131072, 32, 32): (1, 8, 32, 32, 1, 1), + (8192, 8192, 131072, 64, 64): (5, 4, 64, 32, 1, 2), + (8192, 8192, 131072, 128, 128): (1, 4096, 128, 16, 1, 8), + (16384, 16384, 256, 16, 16): (1, 4, 16, 64, 3, 1), + (16384, 16384, 256, 32, 32): (1, 8, 32, 32, 1, 1), + (16384, 16384, 256, 64, 64): (1, 4, 64, 32, 1, 2), + (16384, 16384, 256, 128, 128): (1, 4, 32, 64, 1, 4), + (16384, 16384, 512, 16, 16): (1, 8, 16, 64, 3, 1), + (16384, 16384, 512, 32, 32): (1, 16, 32, 32, 1, 1), + (16384, 16384, 512, 64, 64): (1, 4, 64, 32, 1, 2), + (16384, 16384, 512, 128, 128): (3, 1, 32, 64, 1, 4), + (16384, 16384, 1024, 16, 16): (1, 8, 16, 64, 3, 1), + (16384, 16384, 1024, 32, 32): (1, 16, 32, 32, 1, 1), + (16384, 16384, 1024, 64, 64): (2, 4, 64, 32, 1, 2), + (16384, 16384, 1024, 128, 128): (1, 2, 32, 64, 1, 4), + (16384, 16384, 2048, 16, 16): (1, 4, 16, 64, 3, 1), + (16384, 16384, 2048, 32, 32): (1, 16, 32, 32, 1, 1), + (16384, 16384, 2048, 64, 64): (3, 4, 64, 32, 1, 2), + (16384, 16384, 2048, 128, 128): (2, 1, 32, 64, 1, 4), + (16384, 16384, 4096, 16, 16): (4, 8, 16, 64, 3, 1), + (16384, 16384, 4096, 32, 32): (5, 16, 32, 32, 1, 1), + (16384, 16384, 4096, 64, 64): (3, 2, 64, 32, 1, 2), + (16384, 16384, 4096, 128, 128): (2, 1, 32, 64, 1, 4), + (16384, 16384, 8192, 16, 16): (1, 4, 16, 64, 3, 1), + (16384, 16384, 8192, 32, 32): (1, 4, 32, 32, 1, 1), + (16384, 16384, 8192, 64, 64): (1, 2, 64, 32, 1, 2), + (16384, 16384, 8192, 128, 128): (2, 1, 32, 64, 1, 4), + (16384, 16384, 16384, 16, 16): (1, 8, 16, 64, 3, 1), + (16384, 16384, 16384, 32, 32): (1, 4, 32, 32, 1, 1), + (16384, 16384, 16384, 64, 64): (1, 2, 64, 32, 1, 2), + (16384, 16384, 16384, 128, 128): (3, 1, 32, 64, 1, 4), + (16384, 16384, 32768, 16, 16): (1, 4, 16, 64, 3, 1), + (16384, 16384, 32768, 32, 32): (1, 2, 32, 32, 1, 1), + (16384, 16384, 32768, 64, 64): (3, 2, 64, 32, 1, 2), + (16384, 16384, 32768, 128, 128): (3, 1, 32, 64, 1, 4), + (16384, 16384, 65536, 16, 16): (1, 8, 16, 64, 3, 1), + (16384, 16384, 65536, 32, 32): (1, 4, 32, 32, 1, 1), + (16384, 16384, 65536, 64, 64): (4, 4, 64, 32, 1, 2), + (16384, 16384, 65536, 128, 128): (5, 1, 32, 64, 1, 4), + (16384, 16384, 131072, 16, 16): (1, 2, 16, 64, 3, 1), + (16384, 16384, 131072, 32, 32): (1, 4, 32, 32, 1, 1), + (16384, 16384, 131072, 64, 64): (1, 2, 64, 32, 1, 2), + (16384, 16384, 131072, 128, 128): (1, 4096, 128, 16, 1, 8), + }, + # END GENERATED DATA +} + +if __name__ == "__main__": + for dtype in [torch.int8]: + for op in ["_int_bsr_dense_addmm"]: + main(op=op, force=False, dtype=dtype) + for dtype in [torch.float16, torch.bfloat16, torch.float32, torch.int8]: + for op in ["bsr_dense_addmm"]: + main(op=op, force=False, dtype=dtype) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/semi_structured.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/semi_structured.py new file mode 100644 index 0000000000000000000000000000000000000000..4d99dac9eadefcfb243fdc987023cccf6517edf8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/sparse/semi_structured.py @@ -0,0 +1,840 @@ +# mypy: allow-untyped-defs +import warnings +from collections import namedtuple +from collections.abc import Callable +from typing import Any + +import torch +from torch.sparse._semi_structured_conversions import ( + sparse_semi_structured_from_dense_cutlass, + sparse_semi_structured_to_dense_cutlass, +) +from torch.sparse._semi_structured_ops import ( + fallback_dispatcher, + semi_sparse_addmm, + semi_sparse_clone, + semi_sparse_detach, + semi_sparse_indices, + semi_sparse_linear, + semi_sparse_mm, + semi_sparse_scaled_mm, + semi_sparse_t, + semi_sparse_to, + semi_sparse_to_copy, + semi_sparse_values, + semi_sparse_view, +) + + +__all__ = [ + "SparseSemiStructuredTensor", + "SparseSemiStructuredTensorCUTLASS", + "SparseSemiStructuredTensorCUSPARSELT", + "to_sparse_semi_structured", +] + +_SEMI_STRUCTURED_SPARSE_CONFIG = namedtuple( + "_SEMI_STRUCTURED_SPARSE_CONFIG", + "sparse_min_rows sparse_min_cols dense_min_rows dense_min_cols", +) + + +class SparseSemiStructuredTensor(torch.Tensor): + """ + This class implements semi-structured sparsity as a Tensor subclass. + + Semi-structured sparsity describes a sparsity pattern where n in every 2n elements are sparse, + depending on the datatype. It is also referred to as 2:4 sparsity or fine-grained + structured sparsity. + + There are two backends available for semi_structred sparsity, either cuSPARSELt or CUTLASS. + This class is meant to serve as a base class for both implementations. SparseSemiStructuredCUTLASS + and SparseSemiStructuredCUSPARSELT both inherit from this class and define three backend-specific items. + Note that as such, this class cannot be instantiated directly. + + -`_DTYPE_SHAPE_CONSTRAINTS` - A dictionary holding backend specific dense/sparse min shape constraints + - `def from_dense()` - backend specific compression routines + - `def _mm()` - backend specific mm op (either torch._cslt_sparse_mm or torch._sparse_semi_structured_(mm|addmm)) + """ + + _DEFAULT_ALG_ID: int = 0 + _DTYPE_SHAPE_CONSTRAINTS: dict[torch.dtype, _SEMI_STRUCTURED_SPARSE_CONFIG] + _FORCE_CUTLASS: bool = False + _FUSE_TRANSPOSE: bool = False + _PROTOTYPE_WARNING_SHOWN: bool = False + + BACKEND: str + SPARSE_DISPATCH: dict[Callable, Callable] + + packed: torch.Tensor | None + meta: torch.Tensor | None + packed_t: torch.Tensor | None + meta_t: torch.Tensor | None + compressed_swizzled_bitmask: torch.Tensor | None + fuse_transpose_cusparselt: bool + alg_id_cusparselt: int + + __slots__ = ["packed", "meta", "packed_t", "meta_t", "compressed_swizzled_bitmask"] + + @staticmethod + def __new__( # noqa: PYI034 + cls, + shape: torch.Size, + packed: torch.Tensor | None, + meta: torch.Tensor | None, + packed_t: torch.Tensor | None, + meta_t: torch.Tensor | None, + compressed_swizzled_bitmask: torch.Tensor | None, + fuse_transpose_cusparselt: bool = False, + alg_id_cusparselt: int = 0, + requires_grad: bool = False, + ): + """ + Create a new instance of the tensor subclass from the compressed sparse representation. + + We have the option to create the subclass with the compressed representations of both X and X', for training. + For inference, we only need a single representation (either X or X'), while the corresponding other set will be None. + + Depending on the backend selected, certain fields will be set to None. (CUSPARSELT vs CUTLASS) + + Args: + shape: The shape of the original dense tensor + packed: The compressed representation of the original dense tensor + meta: The metadata of the original dense tensor, if it is stored separately + packed_t: The compressed representation of the transposed original dense tensor + meta_t: The metadata of the transposed original dense tensor, if it is stored separately + compressed_swizzled_bitmask: The masks used by the CUTLASS backend to determine which threads should + participate in the computation. Used for pointwise ops. + fuse_transpose_cusparselt: When running with cuSPARSELt, we have the option to fuse a transposition + with a matmul, which is useful in the case of 2:4 sparse training. + alg_id_cusparselt: The algorithm id to use when using cuSPARSELT, will have effect on performance + + Returns: + torch.Tensor: A torch.Tensor wrapper subclass. + + Raises: + ValueError: If all of the tensor arguments are None. + """ + if not cls._PROTOTYPE_WARNING_SHOWN: + warnings.warn( + ( + "The PyTorch API of SparseSemiStructuredTensor is in prototype stage " + "and will change in the near future. Please open a Github issue " + "for features requests and see our documentation on the torch.sparse " + "module for further information about the project." + ), + UserWarning, + stacklevel=2, + ) + cls._PROTOTYPE_WARNING_SHOWN = True + + # Because this only runs once, we also load the dispatch table here as well. + # We can't define the dispatch table explicitly because of torch.ops import errors, so we do this instead + # But this is useful since it allows users to overload the dispatch table for debugging / testing. + cls._load_dispatch_table() + + # we can also register the classes with dynamo when the warning is shown. + torch._dynamo.allow_in_graph(cls) + + if packed is not None: + previous_tensor = packed + elif packed_t is not None: + previous_tensor = packed_t + else: + raise ValueError("At least one of packed or packed_t must be provided") + + tensor = torch.Tensor._make_wrapper_subclass( + cls, + shape, + device=previous_tensor.device, + dtype=previous_tensor.dtype, + layout=previous_tensor.layout, + requires_grad=requires_grad, + ) + + tensor.packed = packed + tensor.meta = meta + tensor.packed_t = packed_t + tensor.meta_t = meta_t + tensor.compressed_swizzled_bitmask = compressed_swizzled_bitmask + tensor.fuse_transpose_cusparselt = fuse_transpose_cusparselt + tensor.alg_id_cusparselt = alg_id_cusparselt + return tensor + + def __repr__(self) -> str: # type: ignore[override] + if not hasattr(self, "shape"): + raise AssertionError("tensor has no shape attribute") + return f"{self.__class__.__name__}(shape={self.shape})" + + def __tensor_flatten__( + self, + ) -> tuple[list[str], tuple[torch.Size, bool, int, bool]]: + inner_tensors = list( + filter(lambda x: getattr(self, x) is not None, self.__slots__) + ) + tensor_meta = ( + self.shape, + self.fuse_transpose_cusparselt, + self.alg_id_cusparselt, + self.requires_grad, + ) + return inner_tensors, tensor_meta + + @classmethod + def __tensor_unflatten__( + cls, + inner_tensors, + tensor_meta: tuple[torch.Size, bool, int, bool], + outer_size, + outer_stride, + ) -> torch.Tensor: + shape, fuse_transpose_cusparselt, alg_id_cusparselt, requires_grad = tensor_meta + # pyrefly: ignore [no-matching-overload] + return cls( + shape=shape, + packed=inner_tensors.get("packed", None), + meta=inner_tensors.get("meta", None), + packed_t=inner_tensors.get("packed_t", None), + meta_t=inner_tensors.get("meta_t", None), + compressed_swizzled_bitmask=inner_tensors.get( + "compressed_swizzled_bitmask", None + ), + fuse_transpose_cusparselt=fuse_transpose_cusparselt, + alg_id_cusparselt=alg_id_cusparselt, + requires_grad=requires_grad, + ) + + __torch_function__ = torch._C._disabled_torch_function_impl # type: ignore[assignment] + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs) -> Any: # type: ignore[override] + if func._overloadpacket not in cls.SPARSE_DISPATCH: + raise NotImplementedError( + f"{cls.__name__} only supports a specific set of operations, " + f"can't perform requested op ({func.__name__})" + ) + return cls.SPARSE_DISPATCH[func._overloadpacket](func, types, args, kwargs) + + @classmethod + def _load_dispatch_table(cls, custom_dispatch_table=None) -> None: + """ + Loads the op overload sparse dispatch table for the current class. + """ + if getattr(cls, "SPARSE_DISPATCH", None) is None: + cls.SPARSE_DISPATCH = { + torch.ops.aten.values: semi_sparse_values, + torch.ops.aten.indices: semi_sparse_indices, + torch.ops.aten.is_same_size: fallback_dispatcher, + torch.ops.aten.detach_: fallback_dispatcher, + torch.ops.aten.detach: semi_sparse_detach, + torch.ops.aten.t: semi_sparse_t, + torch.ops.aten.view: semi_sparse_view, + torch.ops.aten.mm: semi_sparse_mm, + torch.ops.aten.matmul: semi_sparse_mm, + torch.ops.aten.addmm: semi_sparse_addmm, + torch.ops.aten.linear: semi_sparse_linear, + torch.ops.aten._to_copy: semi_sparse_to_copy, + torch.ops.aten._scaled_mm: semi_sparse_scaled_mm, + torch.ops.aten.clone: semi_sparse_clone, + torch.ops.aten.to: semi_sparse_to, + } + if custom_dispatch_table is not None: + cls.SPARSE_DISPATCH.update(custom_dispatch_table) + + @classmethod + def _validate_device_dim_dtype_shape(cls, original_tensor: torch.Tensor) -> None: + """ + Assert that the given tensor is valid for semi-structured sparse compression. + """ + # check device + if not original_tensor.is_cuda: + raise RuntimeError( + f"Error original_tensor.device= {original_tensor.device} is not supported! " + "Only CUDA tensors are currently supported." + ) + + # check dim + if original_tensor.dim() != 2: + raise RuntimeError( + f"Error original_tensor.dim = {original_tensor.dim()} is not supported! " + "Only 2d tensors are currently supported." + ) + + # check contiguous + if not original_tensor.is_contiguous(): + raise RuntimeError( + "Error original_tensor is not contiguous!" + "Only contiguous tensors are currently supported." + ) + + # check dtype + if original_tensor.dtype not in cls._DTYPE_SHAPE_CONSTRAINTS: + raise RuntimeError( + f"Error original_tensor.dtype {original_tensor.dtype} is not a supported dtype for {cls}!" + ) + + # check shape + m, n = original_tensor.shape + min_rows = cls._DTYPE_SHAPE_CONSTRAINTS[original_tensor.dtype].sparse_min_rows + min_cols = cls._DTYPE_SHAPE_CONSTRAINTS[original_tensor.dtype].sparse_min_cols + if m < min_rows or m % min_rows or n < min_cols or n % min_cols: + # TODO in the future we can add in padding to support sparse dimensions that aren't perfect multiples + raise RuntimeError( + f"Error original_tensor.shape {original_tensor.shape} is not supported! " + f"Both dimensions must be larger or equal than and a multiple of ({min_rows}, {min_cols})" + ) + + def to_dense(self): # type:ignore[override] + col = self.shape[-1] + return torch.mm(self, torch.eye(col, dtype=self.dtype, device=self.device)) + + @classmethod + def from_dense( + cls, + original_tensor: torch.Tensor, + alg_id: int = _DEFAULT_ALG_ID, + ) -> "SparseSemiStructuredTensor": + raise NotImplementedError + + def _mm( + self, + B: torch.Tensor, + *, + bias: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + raise NotImplementedError + + +def to_sparse_semi_structured( + original_tensor: torch.Tensor, + transposed: bool = False, + alg_id: int = SparseSemiStructuredTensor._DEFAULT_ALG_ID, +) -> SparseSemiStructuredTensor: + """ + This function converts a dense tensor into a sparse semi-structured tensor. + It will return a SparseSemiStructuredTensor, a subclass of torch.Tensor. + + This function will check to ensure the dense tensor has the right dtype, size, dims, and device. + We currently only support semi-structured sparse tensors for 2d CUDA tensors. + Additionally, your tensor must be a positive multiple of the minimum sparse block size, given in + `_DTYPE_TO_SHAPE_CONSTRAINTS` for each dtype (float32, float16, bfloat16, int8). + + Args: + original_tensor (Tensor): the dense tensor to convert + transposed (bool, optional): deprecated arg to be removed in another release. Do not use. + alg_id (int, optional): the algorithm id to use for cuSPARSELt matmul. Defaults to 0. + Can be obtained via ``torch._cslt_sparse_mm_search``. + Returns: + SparseSemiStructuredTensor: A sparse semi-structured tensor created from the given original_tensor + Raises: + None + Example: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> A = torch.Tensor([0, 0, 1, 1]).tile((128, 32)).half().cuda() + tensor([[0., 0., 1., ..., 0., 1., 1.], + [0., 0., 1., ..., 0., 1., 1.], + [0., 0., 1., ..., 0., 1., 1.], + ..., + [0., 0., 1., ..., 0., 1., 1.], + [0., 0., 1., ..., 0., 1., 1.], + [0., 0., 1., ..., 0., 1., 1.]], device='cuda:0', dtype=torch.float16) + >>> A_sparse = to_sparse_semi_structured(A) + SparseSemiStructuredTensor(shape=torch.Size([128, 128])) + >>> A_sparse.values() + tensor([[1., 1., 1., ..., 1., 1., 1.], + [1., 1., 1., ..., 1., 1., 1.], + [1., 1., 1., ..., 1., 1., 1.], + ..., + [1., 1., 1., ..., 1., 1., 1.], + [1., 1., 1., ..., 1., 1., 1.], + [1., 1., 1., ..., 1., 1., 1.]], device='cuda:0', dtype=torch.float16), + >>> A_sparse.indices() + tensor([[-4370, -4370, -4370, ..., -4370, -4370, -4370], + [-4370, -4370, -4370, ..., -4370, -4370, -4370], + [-4370, -4370, -4370, ..., -4370, -4370, -4370], + ..., + [-4370, -4370, -4370, ..., -4370, -4370, -4370], + [-4370, -4370, -4370, ..., -4370, -4370, -4370], + [-4370, -4370, -4370, ..., -4370, -4370, -4370]], device='cuda:0', dtype=torch.int16)) + """ + if transposed: + warnings.warn( + "Setting transpose from `to_sparse_semi_structured` is deprecated " + "and will be removed in a future release. " + "`SparseSemiStructuredTensor` only support contiguous input tensors.", + FutureWarning, + stacklevel=2, + ) + + # set from _FORCE_CUTLASS flag + SPARSE_SUBCLASS = ( + torch.sparse.SparseSemiStructuredTensorCUTLASS + if SparseSemiStructuredTensor._FORCE_CUTLASS + else torch.sparse.SparseSemiStructuredTensorCUSPARSELT + ) + + return SPARSE_SUBCLASS.from_dense(original_tensor, alg_id=alg_id) + + +class SparseSemiStructuredTensorCUTLASS(SparseSemiStructuredTensor): + """ + This class implements semi-structured sparsity for the CUTLASS backend. + + + In this implementation, the specified elements and metadata are stored separately, + in packed and meta respectively. + + When _FORCE_CUTLASS is set, or when cuSPARSELt is not available, this subclass calls into _sparse_semi_structured_(mm|addmm) and + sparse_semi_structured_from_dense for conversion to the compressed format. + """ + + BACKEND = "cutlass" + _DTYPE_SHAPE_CONSTRAINTS = { + torch.int8: _SEMI_STRUCTURED_SPARSE_CONFIG(16, 128, 16, 16), + torch.float16: _SEMI_STRUCTURED_SPARSE_CONFIG(32, 64, 8, 8), + torch.bfloat16: _SEMI_STRUCTURED_SPARSE_CONFIG(32, 64, 8, 8), + torch.float32: _SEMI_STRUCTURED_SPARSE_CONFIG(32, 32, 4, 4), + } + + @classmethod + def from_dense( + cls, + original_tensor: torch.Tensor, + alg_id: int = SparseSemiStructuredTensor._DEFAULT_ALG_ID, + ) -> "SparseSemiStructuredTensorCUTLASS": + cls._validate_device_dim_dtype_shape(original_tensor) + ( + sparse_tensor_cutlass, + meta_tensor_cutlass, + ) = sparse_semi_structured_from_dense_cutlass(original_tensor) + # pyrefly: ignore [no-matching-overload] + return cls( + original_tensor.shape, + packed=sparse_tensor_cutlass, + meta=meta_tensor_cutlass, + packed_t=None, + meta_t=None, + compressed_swizzled_bitmask=None, + requires_grad=original_tensor.requires_grad, + ) + + def to_dense(self): # type: ignore[override] + if self.meta is None or self.packed is None: + raise AssertionError("meta and packed must not be None") + return ( + sparse_semi_structured_to_dense_cutlass( + self.packed, + self.meta, + ) + if self.meta.ndim == 2 + else super().to_dense() + ) + + @classmethod + def prune_dense_static_sort( + cls, original_tensor: torch.Tensor, algorithm="" + ) -> "SparseSemiStructuredTensor": + """ + This function takes in a unpruned dense tensor and runs a (branchless) static sort across a 4x4 tile. + + It greedily picks the largest values in the tile, upholding the 2:4 sparsity constraint across both rows and columns. + The algorithm used to prune the matrix is implemented in `_sparse_semi_structured_tile`. + + Then it creates the packed and meta tensors for the compressed sparse representation of the pruned dense tensor. + It also calculates the packed_t and meta_t tensors for the compressed sparse representation of the transposed + pruned dense tensor. + Since we cannot transpose the compressed representations, we store both for the fw/bw pass respectively. + + Finally, this function also computes a compressed swizzled bitmask that encodes the sparsity pattern + This can be used in the backward pass to mask the gradients. + + [9 1 7 4] [9 0 7 0] + [1 2 3 0] [0 2 0 0] + [8 3 5 4] -> prune 4x4 tile -> [8 0 0 4] -> pack to CUTLASS semi-structured -> packed + [1 2 6 2] [0 0 6 2] -> metadata + + -> pack to transposed CUTLASS -> packed_t + semi-structured representation -> metadata_t + + -> compute swizzled bitmask -> compressed_swizzled_bitmask + + + The equivalent PyTorch code to create the same five outputs from the dense tensor can be found below: + ``` + from torch.sparse import SparseSemiStructuredTensorCUTLASS + from torch.sparse._semi_structured_conversions import ( + _sparse_semi_structured_tile, + _compute_compressed_swizzled_bitmask, + ) + + pruned = _sparse_semi_structured_tile(dense) + packed_cutlass, meta_cutlass = sparse_semi_structured_from_dense_cutlass(pruned) + packed_t_cutlass, meta_t_cutlass = sparse_semi_structured_from_dense_cutlass( + pruned.t().contiguous() + ) + bitmask = _compute_compressed_swizzled_bitmask(pruned) + + SparseSemiStructuredTensorCUTLASS( + dense.shape, + packed_cutlass, + meta_cutlass, + packed_t_cutlass, + meta_t_cutlass, + bitmask, + ) + ``` + """ + # We can either pack to the CUTLASS or cuSPARSELt representation, depending on the use_cutlass flag. + ( + packed, + meta, + packed_t, + meta_t, + compressed_swizzled_bitmask, + ) = torch._sparse_semi_structured_tile( + original_tensor, algorithm=algorithm, use_cutlass=True + ) + + # pyrefly: ignore [no-matching-overload] + return cls( + original_tensor.shape, + packed=packed, + meta=meta, + packed_t=packed_t, + meta_t=meta_t, + compressed_swizzled_bitmask=compressed_swizzled_bitmask, + requires_grad=False, + ) + + def _mm( + self, + B: torch.Tensor, + *, + bias: torch.Tensor | None = None, + should_transpose_dense: bool = False, + **kwargs, + ) -> torch.Tensor: + if isinstance(B, SparseSemiStructuredTensor): + raise ValueError( + "`SparseSemiStructuredTensor @ SparseSemiStructuredTensor` is not supported by the hardware" + ) + cls_name = self.__class__.__name__ + if self.ndim != 2 or B.ndim != 2: + raise NotImplementedError( + f"`{cls_name}` matmul: Broadcasting is not implemented" + ) + if self.packed is None or self.meta is None: + raise NotImplementedError( + f"`{cls_name}` matmul: operation is not supported" + ) + else: + _ensure_cutlass_mm_registered() + constraints = self._DTYPE_SHAPE_CONSTRAINTS[B.dtype] + return torch.ops.semi_structured.cutlass_mm( + B, + self.packed, + self.meta, + bias, + self.shape[0], + constraints.dense_min_rows, + constraints.dense_min_cols, + should_transpose_dense, + ) + + +class SparseSemiStructuredTensorCUSPARSELT(SparseSemiStructuredTensor): + """ + The cuSPARSELt backend expects the specified elements and the metadata to be stored in a single tensor: + packed = [ specified elements of original tensor | metadata ] + For an original tensor of size (m, k) we expect the first m * k // 2 elements to be the kept elements + The rest of the tensor is metadata. Since there is only one tensor, we only use the packed and packed_t + attributes respectively. + + cuSPARSELt also supports transposition fusion, which is necessary for performant 2:4 sparse training, as well + as specifying alg_id, a config that affects the performance of the matmul depending on matmul sizes. + """ + + BACKEND = "cusparselt" + _DTYPE_SHAPE_CONSTRAINTS = { + torch.float8_e4m3fn: _SEMI_STRUCTURED_SPARSE_CONFIG(32, 32, 16, 16), + torch.int8: _SEMI_STRUCTURED_SPARSE_CONFIG(32, 32, 16, 16), + torch.float16: _SEMI_STRUCTURED_SPARSE_CONFIG(16, 16, 8, 8), + torch.bfloat16: _SEMI_STRUCTURED_SPARSE_CONFIG(16, 16, 8, 8), + } + + @classmethod + def from_dense( + cls, + original_tensor: torch.Tensor, + alg_id: int = SparseSemiStructuredTensor._DEFAULT_ALG_ID, + ) -> "SparseSemiStructuredTensorCUSPARSELT": + cls._validate_device_dim_dtype_shape(original_tensor) + # pyrefly: ignore [no-matching-overload] + return cls( + shape=original_tensor.shape, + packed=torch._cslt_compress(original_tensor), + meta=None, + packed_t=None, + meta_t=None, + compressed_swizzled_bitmask=None, + fuse_transpose_cusparselt=SparseSemiStructuredTensor._FUSE_TRANSPOSE, + alg_id_cusparselt=alg_id, + requires_grad=original_tensor.requires_grad, + ) + + @classmethod + def prune_dense_static_sort( + cls, original_tensor: torch.Tensor, algorithm="" + ) -> "SparseSemiStructuredTensor": + """ + This function does the same thing as described in SparseSemiStructuredCUTLASS, but uses the cuSPARSELt metadata + layout and sparse matmul. + + The only functional difference is that cuSPARSELt stores `metadata` and `packed` together into a single tensor. + + [9 1 7 4] [9 0 7 0] + [1 2 3 0] [0 2 0 0] + [8 3 5 4] -> prune 4x4 tile -> [8 0 0 4] -> pack to cuSPARSELT semi-structured -> packed + [1 2 6 2] [0 0 6 2] + + -> pack to transposed cuSPARSELt -> packed_t + semi-structured representation + + -> compute swizzled bitmask -> compressed_swizzled_bitmask + + + The equivalent PyTorch code to create the same three outputs from the dense tensor can be found below: + ``` + from torch.sparse import SparseSemiStructuredTensorCUSPARSELT + from torch.sparse._semi_structured_conversions import ( + _sparse_semi_structured_tile, + _compute_compressed_swizzled_bitmask, + ) + + pruned = _sparse_semi_structured_tile(dense) + packed_cusparselt = torch._cslt_compress(pruned) + packed_t_cusparselt = torch._cslt_compress(pruned.t().contiguous()) + bitmask = _compute_compressed_swizzled_bitmask(pruned) + + SparseSemiStructuredTensorCUSPARSELT( + dense.shape, packed_cutlass, None, packed_t_cutlass, None, bitmask + ) + ``` + """ + ( + packed, + meta, + packed_t, + meta_t, + compressed_swizzled_bitmask, + ) = torch._sparse_semi_structured_tile( + original_tensor, algorithm=algorithm, use_cutlass=False + ) + + # Map this two 2-dim view of packed data. + # TODO: is this proper cuSPARSELt metadata? + packed = packed.view(original_tensor.shape[0], -1) + packed_t = packed_t.view(original_tensor.shape[1], -1) + + # pyrefly: ignore [no-matching-overload] + return cls( + original_tensor.shape, + packed=packed, + meta=meta, + packed_t=packed_t, + meta_t=meta_t, + compressed_swizzled_bitmask=compressed_swizzled_bitmask, + requires_grad=False, + ) + + def _mm( + self, + B: torch.Tensor, + *, + bias: torch.Tensor | None = None, + should_transpose_dense: bool = False, + **kwargs, + ) -> torch.Tensor: + if isinstance(B, SparseSemiStructuredTensor): + raise ValueError( + "`SparseSemiStructuredTensor @ SparseSemiStructuredTensor` is not supported by the hardware" + ) + if self.ndim != 2 or B.ndim != 2: + raise NotImplementedError( + f"`{self.__class__.__name__}` matmul: Broadcasting is not implemented" + ) + if B.dtype != self.dtype: + raise NotImplementedError( + f"`{self.__class__.__name__}` matmul: trying to do `A={tuple(self.shape)} @ B={tuple(B.shape)}`, " + f"with A.dtype={self.dtype} and B.dtype={B.dtype}. " + "This operation is only supported when A and B have the same data type." + ) + if bias is not None and bias.dtype != self.dtype: + raise NotImplementedError( + f"`{self.__class__.__name__}` matmul: trying to do `A={tuple(self.shape)} @ B={tuple(B.shape)} + C`, " + f"with A.dtype=B.dtype={self.dtype} and C.dtype={B.dtype}. " + "This operation is only supported when A, B and C have the same data type." + ) + # Force fp8 mm to error to be consistent with torch + if self.dtype == torch.float8_e4m3fn: + raise NotImplementedError( + f"`{self.__class__.__name__}` matmul: trying to do `A={tuple(self.shape)} @ B={tuple(B.shape)}`, " + f"with A.dtype=B.dtype={self.dtype}. " + "mm is not supported for float8_e4m3fn, please use `torch._scaled_mm` instead." + ) + if self.packed is None: + raise NotImplementedError( + f"`{self.__class__.__name__}` matmul: operation is not supported" + ) + else: + _ensure_cusparselt_mm_registered() + constraints = self._DTYPE_SHAPE_CONSTRAINTS[B.dtype] + return torch.ops.semi_structured.cusparselt_mm( + B, + self.packed, + bias, + self.shape[0], + constraints.dense_min_rows, + constraints.dense_min_cols, + self.fuse_transpose_cusparselt, + self.alg_id_cusparselt, + should_transpose_dense, + ) + + +_cutlass_mm_registered = False + + +def _ensure_cutlass_mm_registered(): + """Lazily register the cutlass_mm custom op. + + Registration is deferred to avoid importing torch.library at module load + time, since torch.sparse is imported early during ``import torch``. + """ + global _cutlass_mm_registered + if _cutlass_mm_registered: + return + _cutlass_mm_registered = True + + from torch.library import custom_op + + @custom_op("semi_structured::cutlass_mm", mutates_args=()) + def cutlass_mm( + dense: torch.Tensor, + packed: torch.Tensor, + meta: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + min_rows: int, + min_cols: int, + should_transpose_dense: bool, + ) -> torch.Tensor: + m, n = dense.shape + to_pad_m = (-m) % min_rows + to_pad_n = (-n) % min_cols + need_pad = to_pad_m != 0 or to_pad_n != 0 + dense_padded = dense + if need_pad: + dense_padded = torch.nn.functional.pad(dense, (0, to_pad_n, 0, to_pad_m)) + mm_input = dense_padded.t() if should_transpose_dense else dense_padded + if bias is None: + res = torch._sparse_semi_structured_mm(packed, meta, mm_input) + else: + res = torch._sparse_semi_structured_addmm(bias, packed, meta, mm_input) + if need_pad: + out_cols = m if should_transpose_dense else n + return ( + res[:out_features] + .narrow(1, 0, out_cols) + .clone(memory_format=torch.contiguous_format) + ) + return res.contiguous() + + @cutlass_mm.register_fake + def _cutlass_mm_fake( + dense: torch.Tensor, + packed: torch.Tensor, + meta: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + min_rows: int, + min_cols: int, + transpose_dense: bool, + ) -> torch.Tensor: + out_cols = dense.shape[0] if transpose_dense else dense.shape[1] + return torch.empty( + out_features, + out_cols, + dtype=dense.dtype, + device=dense.device, + ) + + +_cusparselt_mm_registered = False + + +def _ensure_cusparselt_mm_registered(): + """Lazily register the cusparselt_mm custom op.""" + global _cusparselt_mm_registered + if _cusparselt_mm_registered: + return + _cusparselt_mm_registered = True + + from torch.library import custom_op + + @custom_op("semi_structured::cusparselt_mm", mutates_args=()) + def cusparselt_mm( + dense: torch.Tensor, + packed: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + min_rows: int, + min_cols: int, + fuse_transpose: bool, + alg_id: int, + should_transpose_dense: bool = False, + ) -> torch.Tensor: + m, n = dense.shape + to_pad_m = (-m) % min_rows + to_pad_n = (-n) % min_cols + need_pad = to_pad_m != 0 or to_pad_n != 0 + dense_padded = dense + if need_pad: + dense_padded = torch.nn.functional.pad(dense, (0, to_pad_n, 0, to_pad_m)) + mm_input = dense_padded.t() if should_transpose_dense else dense_padded + res = torch._cslt_sparse_mm( + packed, + mm_input, + bias=bias, + transpose_result=fuse_transpose, + alg_id=alg_id, + ) + if fuse_transpose: + res = res.t() + if need_pad: + out_cols = m if should_transpose_dense else n + return res.narrow(1, 0, out_cols).clone( + memory_format=torch.contiguous_format + ) + return res.contiguous() + + @cusparselt_mm.register_fake + def _cusparselt_mm_fake( + dense: torch.Tensor, + packed: torch.Tensor, + bias: torch.Tensor | None, + out_features: int, + min_rows: int, + min_cols: int, + fuse_transpose: bool, + alg_id: int, + should_transpose_dense: bool, + ) -> torch.Tensor: + out_cols = dense.shape[0] if should_transpose_dense else dense.shape[1] + return torch.empty( + out_features, + out_cols, + dtype=dense.dtype, + device=dense.device, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/special/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/special/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dbc9314ad20875b6d37732e201a86f7bc66488dd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/special/__init__.py @@ -0,0 +1,1459 @@ +import torch +from torch._C import _add_docstr, _special # type: ignore[attr-defined] +from torch._torch_docs import common_args, multi_dim_common + + +__all__ = [ + "airy_ai", + "bessel_j0", + "bessel_j1", + "bessel_y0", + "bessel_y1", + "chebyshev_polynomial_t", + "chebyshev_polynomial_u", + "chebyshev_polynomial_v", + "chebyshev_polynomial_w", + "digamma", + "entr", + "erf", + "erfc", + "erfcx", + "erfinv", + "exp2", + "expit", + "expm1", + "gammainc", + "gammaincc", + "gammaln", + "hermite_polynomial_h", + "hermite_polynomial_he", + "i0", + "i0e", + "i1", + "i1e", + "laguerre_polynomial_l", + "legendre_polynomial_p", + "log1p", + "log_ndtr", + "log_softmax", + "logit", + "logsumexp", + "modified_bessel_i0", + "modified_bessel_i1", + "modified_bessel_k0", + "modified_bessel_k1", + "multigammaln", + "ndtr", + "ndtri", + "polygamma", + "psi", + "round", + "shifted_chebyshev_polynomial_t", + "shifted_chebyshev_polynomial_u", + "shifted_chebyshev_polynomial_v", + "shifted_chebyshev_polynomial_w", + "scaled_modified_bessel_k0", + "scaled_modified_bessel_k1", + "sinc", + "softmax", + "spherical_bessel_j0", + "xlog1py", + "xlogy", + "zeta", +] + +Tensor = torch.Tensor + +entr = _add_docstr( + _special.special_entr, + r""" +entr(input, *, out=None) -> Tensor +Computes the entropy on :attr:`input` (as defined below), elementwise. + +.. math:: + \begin{align} + \text{entr(x)} = \begin{cases} + -x * \ln(x) & x > 0 \\ + 0 & x = 0.0 \\ + -\infty & x < 0 + \end{cases} + \end{align} +""" + + """ + +Args: + input (Tensor): the input tensor. + +Keyword args: + out (Tensor, optional): the output tensor. + +Example:: + + >>> a = torch.arange(-0.5, 1, 0.5) + >>> a + tensor([-0.5000, 0.0000, 0.5000]) + >>> torch.special.entr(a) + tensor([ -inf, 0.0000, 0.3466]) +""", +) + +psi = _add_docstr( + _special.special_psi, + r""" +psi(input, *, out=None) -> Tensor + +Alias for :func:`torch.special.digamma`. +""", +) + +digamma = _add_docstr( + _special.special_digamma, + r""" +digamma(input, *, out=None) -> Tensor + +Computes the logarithmic derivative of the gamma function on `input`. + +.. math:: + \digamma(x) = \frac{d}{dx} \ln\left(\Gamma\left(x\right)\right) = \frac{\Gamma'(x)}{\Gamma(x)} +""" + + r""" +Args: + input (Tensor): the tensor to compute the digamma function on + +Keyword args: + {out} + +.. note:: This function is similar to SciPy's `scipy.special.digamma`. + +.. note:: From PyTorch 1.8 onwards, the digamma function returns `-Inf` for `0`. + Previously it returned `NaN` for `0`. + +Example:: + + >>> a = torch.tensor([1, 0.5]) + >>> torch.special.digamma(a) + tensor([-0.5772, -1.9635]) + +""".format(**common_args), +) + +gammaln = _add_docstr( + _special.special_gammaln, + r""" +gammaln(input, *, out=None) -> Tensor + +Computes the natural logarithm of the absolute value of the gamma function on :attr:`input`. + +.. math:: + \text{out}_{i} = \ln \Gamma(|\text{input}_{i}|) +""" + + """ +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.arange(0.5, 2, 0.5) + >>> torch.special.gammaln(a) + tensor([ 0.5724, 0.0000, -0.1208]) + +""".format(**common_args), +) + +polygamma = _add_docstr( + _special.special_polygamma, + r""" +polygamma(n, input, *, out=None) -> Tensor + +Computes the :math:`n^{th}` derivative of the digamma function on :attr:`input`. +:math:`n \geq 0` is called the order of the polygamma function. + +.. math:: + \psi^{(n)}(x) = \frac{d^{(n)}}{dx^{(n)}} \psi(x) + +.. note:: + This function is implemented only for nonnegative integers :math:`n \geq 0`. +""" + + """ +Args: + n (int): the order of the polygamma function + {input} + +Keyword args: + {out} + +Example:: + + >>> a = torch.tensor([1, 0.5]) + >>> torch.special.polygamma(1, a) + tensor([1.64493, 4.9348]) + >>> torch.special.polygamma(2, a) + tensor([ -2.4041, -16.8288]) + >>> torch.special.polygamma(3, a) + tensor([ 6.4939, 97.4091]) + >>> torch.special.polygamma(4, a) + tensor([ -24.8863, -771.4742]) +""".format(**common_args), +) + +erf = _add_docstr( + _special.special_erf, + r""" +erf(input, *, out=None) -> Tensor + +Computes the error function of :attr:`input`. The error function is defined as follows: + +.. math:: + \mathrm{erf}(x) = \frac{2}{\sqrt{\pi}} \int_{0}^{x} e^{-t^2} dt +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.erf(torch.tensor([0, -1., 10.])) + tensor([ 0.0000, -0.8427, 1.0000]) +""".format(**common_args), +) + +erfc = _add_docstr( + _special.special_erfc, + r""" +erfc(input, *, out=None) -> Tensor + +Computes the complementary error function of :attr:`input`. +The complementary error function is defined as follows: + +.. math:: + \mathrm{erfc}(x) = 1 - \frac{2}{\sqrt{\pi}} \int_{0}^{x} e^{-t^2} dt +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.erfc(torch.tensor([0, -1., 10.])) + tensor([ 1.0000, 1.8427, 0.0000]) +""".format(**common_args), +) + +erfcx = _add_docstr( + _special.special_erfcx, + r""" +erfcx(input, *, out=None) -> Tensor + +Computes the scaled complementary error function for each element of :attr:`input`. +The scaled complementary error function is defined as follows: + +.. math:: + \mathrm{erfcx}(x) = e^{x^2} \mathrm{erfc}(x) +""" + + r""" + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.erfcx(torch.tensor([0, -1., 10.])) + tensor([ 1.0000, 5.0090, 0.0561]) +""".format(**common_args), +) + +erfinv = _add_docstr( + _special.special_erfinv, + r""" +erfinv(input, *, out=None) -> Tensor + +Computes the inverse error function of :attr:`input`. +The inverse error function is defined in the range :math:`(-1, 1)` as: + +.. math:: + \mathrm{erfinv}(\mathrm{erf}(x)) = x +""" + + r""" + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.erfinv(torch.tensor([0, 0.5, -1.])) + tensor([ 0.0000, 0.4769, -inf]) +""".format(**common_args), +) + +logit = _add_docstr( + _special.special_logit, + r""" +logit(input, eps=None, *, out=None) -> Tensor + +Returns a new tensor with the logit of the elements of :attr:`input`. +:attr:`input` is clamped to [eps, 1 - eps] when eps is not None. +When eps is None and :attr:`input` < 0 or :attr:`input` > 1, the function will yields NaN. + +.. math:: + \begin{align} + y_{i} &= \ln(\frac{z_{i}}{1 - z_{i}}) \\ + z_{i} &= \begin{cases} + x_{i} & \text{if eps is None} \\ + \text{eps} & \text{if } x_{i} < \text{eps} \\ + x_{i} & \text{if } \text{eps} \leq x_{i} \leq 1 - \text{eps} \\ + 1 - \text{eps} & \text{if } x_{i} > 1 - \text{eps} + \end{cases} + \end{align} +""" + + r""" +Args: + {input} + eps (float, optional): the epsilon for input clamp bound. Default: ``None`` + +Keyword args: + {out} + +Example:: + + >>> a = torch.rand(5) + >>> a + tensor([0.2796, 0.9331, 0.6486, 0.1523, 0.6516]) + >>> torch.special.logit(a, eps=1e-6) + tensor([-0.9466, 2.6352, 0.6131, -1.7169, 0.6261]) +""".format(**common_args), +) + +logsumexp = _add_docstr( + _special.special_logsumexp, + r""" +logsumexp(input, dim, keepdim=False, *, out=None) + +Alias for :func:`torch.logsumexp`. +""".format(**multi_dim_common), +) + +expit = _add_docstr( + _special.special_expit, + r""" +expit(input, *, out=None) -> Tensor + +Computes the expit (also known as the logistic sigmoid function) of the elements of :attr:`input`. + +.. math:: + \text{out}_{i} = \frac{1}{1 + e^{-\text{input}_{i}}} +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> t = torch.randn(4) + >>> t + tensor([ 0.9213, 1.0887, -0.8858, -1.7683]) + >>> torch.special.expit(t) + tensor([ 0.7153, 0.7481, 0.2920, 0.1458]) +""".format(**common_args), +) + +exp2 = _add_docstr( + _special.special_exp2, + r""" +exp2(input, *, out=None) -> Tensor + +Computes the base two exponential function of :attr:`input`. + +.. math:: + y_{i} = 2^{x_{i}} + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.exp2(torch.tensor([0, math.log2(2.), 3, 4])) + tensor([ 1., 2., 8., 16.]) +""".format(**common_args), +) + +expm1 = _add_docstr( + _special.special_expm1, + r""" +expm1(input, *, out=None) -> Tensor + +Computes the exponential of the elements minus 1 +of :attr:`input`. + +.. math:: + y_{i} = e^{x_{i}} - 1 + +.. note:: This function provides greater precision than exp(x) - 1 for small values of x. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.expm1(torch.tensor([0, math.log(2.)])) + tensor([ 0., 1.]) +""".format(**common_args), +) + +xlog1py = _add_docstr( + _special.special_xlog1py, + r""" +xlog1py(input, other, *, out=None) -> Tensor + +Computes ``input * log1p(other)`` with the following cases. + +.. math:: + \text{out}_{i} = \begin{cases} + \text{NaN} & \text{if } \text{other}_{i} = \text{NaN} \\ + 0 & \text{if } \text{input}_{i} = 0.0 \text{ and } \text{other}_{i} != \text{NaN} \\ + \text{input}_{i} * \text{log1p}(\text{other}_{i})& \text{otherwise} + \end{cases} + +Similar to SciPy's `scipy.special.xlog1py`. + +""" + + r""" + +Args: + input (Number or Tensor) : Multiplier + other (Number or Tensor) : Argument + +.. note:: At least one of :attr:`input` or :attr:`other` must be a tensor. + +Keyword args: + {out} + +Example:: + + >>> x = torch.zeros(5,) + >>> y = torch.tensor([-1, 0, 1, float('inf'), float('nan')]) + >>> torch.special.xlog1py(x, y) + tensor([0., 0., 0., 0., nan]) + >>> x = torch.tensor([1, 2, 3]) + >>> y = torch.tensor([3, 2, 1]) + >>> torch.special.xlog1py(x, y) + tensor([1.3863, 2.1972, 2.0794]) + >>> torch.special.xlog1py(x, 4) + tensor([1.6094, 3.2189, 4.8283]) + >>> torch.special.xlog1py(2, y) + tensor([2.7726, 2.1972, 1.3863]) +""".format(**common_args), +) + +xlogy = _add_docstr( + _special.special_xlogy, + r""" +xlogy(input, other, *, out=None) -> Tensor + +Computes ``input * log(other)`` with the following cases. + +.. math:: + \text{out}_{i} = \begin{cases} + \text{NaN} & \text{if } \text{other}_{i} = \text{NaN} \\ + 0 & \text{if } \text{input}_{i} = 0.0 \\ + \text{input}_{i} * \log{(\text{other}_{i})} & \text{otherwise} + \end{cases} + +Similar to SciPy's `scipy.special.xlogy`. + +""" + + r""" + +Args: + input (Number or Tensor) : Multiplier + other (Number or Tensor) : Argument + +.. note:: At least one of :attr:`input` or :attr:`other` must be a tensor. + +Keyword args: + {out} + +Example:: + + >>> x = torch.zeros(5,) + >>> y = torch.tensor([-1, 0, 1, float('inf'), float('nan')]) + >>> torch.special.xlogy(x, y) + tensor([0., 0., 0., 0., nan]) + >>> x = torch.tensor([1, 2, 3]) + >>> y = torch.tensor([3, 2, 1]) + >>> torch.special.xlogy(x, y) + tensor([1.0986, 1.3863, 0.0000]) + >>> torch.special.xlogy(x, 4) + tensor([1.3863, 2.7726, 4.1589]) + >>> torch.special.xlogy(2, y) + tensor([2.1972, 1.3863, 0.0000]) +""".format(**common_args), +) + +i0 = _add_docstr( + _special.special_i0, + r""" +i0(input, *, out=None) -> Tensor + +Computes the zeroth order modified Bessel function of the first kind for each element of :attr:`input`. + +.. math:: + \text{out}_{i} = I_0(\text{input}_{i}) = \sum_{k=0}^{\infty} \frac{(\text{input}_{i}^2/4)^k}{(k!)^2} + +""" + + r""" +Args: + input (Tensor): the input tensor + +Keyword args: + {out} + +Example:: + + >>> torch.i0(torch.arange(5, dtype=torch.float32)) + tensor([ 1.0000, 1.2661, 2.2796, 4.8808, 11.3019]) + +""".format(**common_args), +) + +i0e = _add_docstr( + _special.special_i0e, + r""" +i0e(input, *, out=None) -> Tensor +Computes the exponentially scaled zeroth order modified Bessel function of the first kind (as defined below) +for each element of :attr:`input`. + +.. math:: + \text{out}_{i} = \exp(-|x|) * i0(x) = \exp(-|x|) * \sum_{k=0}^{\infty} \frac{(\text{input}_{i}^2/4)^k}{(k!)^2} + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.i0e(torch.arange(5, dtype=torch.float32)) + tensor([1.0000, 0.4658, 0.3085, 0.2430, 0.2070]) +""".format(**common_args), +) + +i1 = _add_docstr( + _special.special_i1, + r""" +i1(input, *, out=None) -> Tensor +Computes the first order modified Bessel function of the first kind (as defined below) +for each element of :attr:`input`. + +.. math:: + \text{out}_{i} = \frac{(\text{input}_{i})}{2} * \sum_{k=0}^{\infty} \frac{(\text{input}_{i}^2/4)^k}{(k!) * (k+1)!} + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.i1(torch.arange(5, dtype=torch.float32)) + tensor([0.0000, 0.5652, 1.5906, 3.9534, 9.7595]) +""".format(**common_args), +) + +i1e = _add_docstr( + _special.special_i1e, + r""" +i1e(input, *, out=None) -> Tensor +Computes the exponentially scaled first order modified Bessel function of the first kind (as defined below) +for each element of :attr:`input`. + +.. math:: + \text{out}_{i} = \exp(-|x|) * i1(x) = + \exp(-|x|) * \frac{(\text{input}_{i})}{2} * \sum_{k=0}^{\infty} \frac{(\text{input}_{i}^2/4)^k}{(k!) * (k+1)!} + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.i1e(torch.arange(5, dtype=torch.float32)) + tensor([0.0000, 0.2079, 0.2153, 0.1968, 0.1788]) +""".format(**common_args), +) + +ndtr = _add_docstr( + _special.special_ndtr, + r""" +ndtr(input, *, out=None) -> Tensor +Computes the area under the standard Gaussian probability density function, +integrated from minus infinity to :attr:`input`, elementwise. + +.. math:: + \text{ndtr}(x) = \frac{1}{\sqrt{2 \pi}}\int_{-\infty}^{x} e^{-\frac{1}{2}t^2} dt + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.ndtr(torch.tensor([-3., -2, -1, 0, 1, 2, 3])) + tensor([0.0013, 0.0228, 0.1587, 0.5000, 0.8413, 0.9772, 0.9987]) +""".format(**common_args), +) + +ndtri = _add_docstr( + _special.special_ndtri, + r""" +ndtri(input, *, out=None) -> Tensor +Computes the argument, x, for which the area under the Gaussian probability density function +(integrated from minus infinity to x) is equal to :attr:`input`, elementwise. + +.. math:: + \text{ndtri}(p) = \sqrt{2}\text{erf}^{-1}(2p - 1) + +.. note:: + Also known as quantile function for Normal Distribution. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.ndtri(torch.tensor([0, 0.25, 0.5, 0.75, 1])) + tensor([ -inf, -0.6745, 0.0000, 0.6745, inf]) +""".format(**common_args), +) + +log_ndtr = _add_docstr( + _special.special_log_ndtr, + r""" +log_ndtr(input, *, out=None) -> Tensor +Computes the log of the area under the standard Gaussian probability density function, +integrated from minus infinity to :attr:`input`, elementwise. + +.. math:: + \text{log\_ndtr}(x) = \log\left(\frac{1}{\sqrt{2 \pi}}\int_{-\infty}^{x} e^{-\frac{1}{2}t^2} dt \right) + +""" + + r""" +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> torch.special.log_ndtr(torch.tensor([-3., -2, -1, 0, 1, 2, 3])) + tensor([-6.6077 -3.7832 -1.841 -0.6931 -0.1728 -0.023 -0.0014]) +""".format(**common_args), +) + +log1p = _add_docstr( + _special.special_log1p, + r""" +log1p(input, *, out=None) -> Tensor + +Alias for :func:`torch.log1p`. +""", +) + +sinc = _add_docstr( + _special.special_sinc, + r""" +sinc(input, *, out=None) -> Tensor + +Computes the normalized sinc of :attr:`input.` + +.. math:: + \text{out}_{i} = + \begin{cases} + 1, & \text{if}\ \text{input}_{i}=0 \\ + \sin(\pi \text{input}_{i}) / (\pi \text{input}_{i}), & \text{otherwise} + \end{cases} +""" + + r""" + +Args: + {input} + +Keyword args: + {out} + +Example:: + + >>> t = torch.randn(4) + >>> t + tensor([ 0.2252, -0.2948, 1.0267, -1.1566]) + >>> torch.special.sinc(t) + tensor([ 0.9186, 0.8631, -0.0259, -0.1300]) +""".format(**common_args), +) + +round = _add_docstr( + _special.special_round, + r""" +round(input, *, out=None) -> Tensor + +Alias for :func:`torch.round`. +""", +) + +softmax = _add_docstr( + _special.special_softmax, + r""" +softmax(input, dim, *, dtype=None) -> Tensor + +Computes the softmax function. + +Softmax is defined as: + +:math:`\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}` + +It is applied to all slices along dim, and will re-scale them so that the elements +lie in the range `[0, 1]` and sum to 1. + +Args: + input (Tensor): input + dim (int): A dimension along which softmax will be computed. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is cast to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + +Examples:: + >>> t = torch.ones(2, 2) + >>> torch.special.softmax(t, 0) + tensor([[0.5000, 0.5000], + [0.5000, 0.5000]]) + +""", +) + +log_softmax = _add_docstr( + _special.special_log_softmax, + r""" +log_softmax(input, dim, *, dtype=None) -> Tensor + +Computes softmax followed by a logarithm. + +While mathematically equivalent to log(softmax(x)), doing these two +operations separately is slower and numerically unstable. This function +is computed as: + +.. math:: + \text{log\_softmax}(x_{i}) = \log\left(\frac{\exp(x_i) }{ \sum_j \exp(x_j)} \right) +""" + + r""" + +Args: + input (Tensor): input + dim (int): A dimension along which log_softmax will be computed. + dtype (:class:`torch.dtype`, optional): the desired data type of returned tensor. + If specified, the input tensor is cast to :attr:`dtype` before the operation + is performed. This is useful for preventing data type overflows. Default: None. + +Example:: + + >>> t = torch.ones(2, 2) + >>> torch.special.log_softmax(t, 0) + tensor([[-0.6931, -0.6931], + [-0.6931, -0.6931]]) +""", +) + +zeta = _add_docstr( + _special.special_zeta, + r""" +zeta(input, other, *, out=None) -> Tensor + +Computes the Hurwitz zeta function, elementwise. + +.. math:: + \zeta(x, q) = \sum_{k=0}^{\infty} \frac{1}{(k + q)^x} + +""" + + r""" +Args: + input (Tensor): the input tensor corresponding to `x`. + other (Tensor): the input tensor corresponding to `q`. + +.. note:: + The Riemann zeta function corresponds to the case when `q = 1` + +Keyword args: + {out} + +Example:: + + >>> x = torch.tensor([2., 4.]) + >>> torch.special.zeta(x, 1) + tensor([1.6449, 1.0823]) + >>> torch.special.zeta(x, torch.tensor([1., 2.])) + tensor([1.6449, 0.0823]) + >>> torch.special.zeta(2, torch.tensor([1., 2.])) + tensor([1.6449, 0.6449]) +""".format(**common_args), +) + +multigammaln = _add_docstr( + _special.special_multigammaln, + r""" +multigammaln(input, p, *, out=None) -> Tensor + +Computes the `multivariate log-gamma function +`_ with dimension +:math:`p` element-wise, given by + +.. math:: + \log(\Gamma_{p}(a)) = C + \displaystyle \sum_{i=1}^{p} \log\left(\Gamma\left(a - \frac{i - 1}{2}\right)\right) + +where :math:`C = \log(\pi) \cdot \frac{p (p - 1)}{4}` and :math:`\Gamma(-)` is the Gamma function. + +All elements must be greater than :math:`\frac{p - 1}{2}`, otherwise the behavior is undefined. +""" + + """ + +Args: + input (Tensor): the tensor to compute the multivariate log-gamma function + p (int): the number of dimensions + +Keyword args: + {out} + +Example:: + + >>> a = torch.empty(2, 3).uniform_(1, 2) + >>> a + tensor([[1.6835, 1.8474, 1.1929], + [1.0475, 1.7162, 1.4180]]) + >>> torch.special.multigammaln(a, 2) + tensor([[0.3928, 0.4007, 0.7586], + [1.0311, 0.3901, 0.5049]]) +""".format(**common_args), +) + +gammainc = _add_docstr( + _special.special_gammainc, + r""" +gammainc(input, other, *, out=None) -> Tensor + +Computes the regularized lower incomplete gamma function: + +.. math:: + \text{out}_{i} = \frac{1}{\Gamma(\text{input}_i)} \int_0^{\text{other}_i} t^{\text{input}_i-1} e^{-t} dt + +where both :math:`\text{input}_i` and :math:`\text{other}_i` are weakly positive +and at least one is strictly positive. +If both are zero or either is negative then :math:`\text{out}_i=\text{nan}`. +:math:`\Gamma(\cdot)` in the equation above is the gamma function, + +.. math:: + \Gamma(\text{input}_i) = \int_0^\infty t^{(\text{input}_i-1)} e^{-t} dt. + +See :func:`torch.special.gammaincc` and :func:`torch.special.gammaln` for related functions. + +Supports :ref:`broadcasting to a common shape ` +and float inputs. + +.. note:: + The backward pass with respect to :attr:`input` is not yet supported. + Please open an issue on PyTorch's Github to request it. + +""" + + r""" +Args: + input (Tensor): the first non-negative input tensor + other (Tensor): the second non-negative input tensor + +Keyword args: + {out} + +Example:: + + >>> a1 = torch.tensor([4.0]) + >>> a2 = torch.tensor([3.0, 4.0, 5.0]) + >>> a = torch.special.gammaincc(a1, a2) + tensor([0.3528, 0.5665, 0.7350]) + tensor([0.3528, 0.5665, 0.7350]) + >>> b = torch.special.gammainc(a1, a2) + torch.special.gammaincc(a1, a2) + tensor([1., 1., 1.]) + +""".format(**common_args), +) + +gammaincc = _add_docstr( + _special.special_gammaincc, + r""" +gammaincc(input, other, *, out=None) -> Tensor + +Computes the regularized upper incomplete gamma function: + +.. math:: + \text{out}_{i} = \frac{1}{\Gamma(\text{input}_i)} \int_{\text{other}_i}^{\infty} t^{\text{input}_i-1} e^{-t} dt + +where both :math:`\text{input}_i` and :math:`\text{other}_i` are weakly positive +and at least one is strictly positive. +If both are zero or either is negative then :math:`\text{out}_i=\text{nan}`. +:math:`\Gamma(\cdot)` in the equation above is the gamma function, + +.. math:: + \Gamma(\text{input}_i) = \int_0^\infty t^{(\text{input}_i-1)} e^{-t} dt. + +See :func:`torch.special.gammainc` and :func:`torch.special.gammaln` for related functions. + +Supports :ref:`broadcasting to a common shape ` +and float inputs. + +.. note:: + The backward pass with respect to :attr:`input` is not yet supported. + Please open an issue on PyTorch's Github to request it. + +""" + + r""" +Args: + input (Tensor): the first non-negative input tensor + other (Tensor): the second non-negative input tensor + +Keyword args: + {out} + +Example:: + + >>> a1 = torch.tensor([4.0]) + >>> a2 = torch.tensor([3.0, 4.0, 5.0]) + >>> a = torch.special.gammaincc(a1, a2) + tensor([0.6472, 0.4335, 0.2650]) + >>> b = torch.special.gammainc(a1, a2) + torch.special.gammaincc(a1, a2) + tensor([1., 1., 1.]) + +""".format(**common_args), +) + +airy_ai = _add_docstr( + _special.special_airy_ai, + r""" +airy_ai(input, *, out=None) -> Tensor + +Airy function :math:`\text{Ai}\left(\text{input}\right)`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +bessel_j0 = _add_docstr( + _special.special_bessel_j0, + r""" +bessel_j0(input, *, out=None) -> Tensor + +Bessel function of the first kind of order :math:`0`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +bessel_j1 = _add_docstr( + _special.special_bessel_j1, + r""" +bessel_j1(input, *, out=None) -> Tensor + +Bessel function of the first kind of order :math:`1`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +bessel_y0 = _add_docstr( + _special.special_bessel_y0, + r""" +bessel_y0(input, *, out=None) -> Tensor + +Bessel function of the second kind of order :math:`0`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +bessel_y1 = _add_docstr( + _special.special_bessel_y1, + r""" +bessel_y1(input, *, out=None) -> Tensor + +Bessel function of the second kind of order :math:`1`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +chebyshev_polynomial_t = _add_docstr( + _special.special_chebyshev_polynomial_t, + r""" +chebyshev_polynomial_t(input, n, *, out=None) -> Tensor + +Chebyshev polynomial of the first kind :math:`T_{n}(\text{input})`. + +If :math:`n = 0`, :math:`1` is returned. If :math:`n = 1`, :math:`\text{input}` +is returned. If :math:`n < 6` or :math:`|\text{input}| > 1` the recursion: + +.. math:: + T_{n + 1}(\text{input}) = 2 \times \text{input} \times T_{n}(\text{input}) - T_{n - 1}(\text{input}) + +is evaluated. Otherwise, the explicit trigonometric formula: + +.. math:: + T_{n}(\text{input}) = \text{cos}(n \times \text{arccos}(x)) + +is evaluated. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +chebyshev_polynomial_u = _add_docstr( + _special.special_chebyshev_polynomial_u, + r""" +chebyshev_polynomial_u(input, n, *, out=None) -> Tensor + +Chebyshev polynomial of the second kind :math:`U_{n}(\text{input})`. + +If :math:`n = 0`, :math:`1` is returned. If :math:`n = 1`, +:math:`2 \times \text{input}` is returned. If :math:`n < 6` or +:math:`|\text{input}| > 1`, the recursion: + +.. math:: + U_{n + 1}(\text{input}) = 2 \times \text{input} \times U_{n}(\text{input}) - U_{n - 1}(\text{input}) + +is evaluated. Otherwise, the explicit trigonometric formula: + +.. math:: + \frac{\text{sin}((n + 1) \times \text{arccos}(\text{input}))}{\text{sin}(\text{arccos}(\text{input}))} + +is evaluated. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +chebyshev_polynomial_v = _add_docstr( + _special.special_chebyshev_polynomial_v, + r""" +chebyshev_polynomial_v(input, n, *, out=None) -> Tensor + +Chebyshev polynomial of the third kind :math:`V_{n}^{\ast}(\text{input})`. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +chebyshev_polynomial_w = _add_docstr( + _special.special_chebyshev_polynomial_w, + r""" +chebyshev_polynomial_w(input, n, *, out=None) -> Tensor + +Chebyshev polynomial of the fourth kind :math:`W_{n}^{\ast}(\text{input})`. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +hermite_polynomial_h = _add_docstr( + _special.special_hermite_polynomial_h, + r""" +hermite_polynomial_h(input, n, *, out=None) -> Tensor + +Physicist's Hermite polynomial :math:`H_{n}(\text{input})`. + +If :math:`n = 0`, :math:`1` is returned. If :math:`n = 1`, :math:`\text{input}` +is returned. Otherwise, the recursion: + +.. math:: + H_{n + 1}(\text{input}) = 2 \times \text{input} \times H_{n}(\text{input}) - H_{n - 1}(\text{input}) + +is evaluated. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +hermite_polynomial_he = _add_docstr( + _special.special_hermite_polynomial_he, + r""" +hermite_polynomial_he(input, n, *, out=None) -> Tensor + +Probabilist's Hermite polynomial :math:`He_{n}(\text{input})`. + +If :math:`n = 0`, :math:`1` is returned. If :math:`n = 1`, :math:`\text{input}` +is returned. Otherwise, the recursion: + +.. math:: + He_{n + 1}(\text{input}) = 2 \times \text{input} \times He_{n}(\text{input}) - He_{n - 1}(\text{input}) + +is evaluated. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +laguerre_polynomial_l = _add_docstr( + _special.special_laguerre_polynomial_l, + r""" +laguerre_polynomial_l(input, n, *, out=None) -> Tensor + +Laguerre polynomial :math:`L_{n}(\text{input})`. + +If :math:`n = 0`, :math:`1` is returned. If :math:`n = 1`, :math:`\text{input}` +is returned. Otherwise, the recursion: + +.. math:: + L_{n + 1}(\text{input}) = 2 \times \text{input} \times L_{n}(\text{input}) - L_{n - 1}(\text{input}) + +is evaluated. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +legendre_polynomial_p = _add_docstr( + _special.special_legendre_polynomial_p, + r""" +legendre_polynomial_p(input, n, *, out=None) -> Tensor + +Legendre polynomial :math:`P_{n}(\text{input})`. + +If :math:`n = 0`, :math:`1` is returned. If :math:`n = 1`, :math:`\text{input}` +is returned. Otherwise, the recursion: + +.. math:: + P_{n + 1}(\text{input}) = 2 \times \text{input} \times P_{n}(\text{input}) - P_{n - 1}(\text{input}) + +is evaluated. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +modified_bessel_i0 = _add_docstr( + _special.special_modified_bessel_i0, + r""" +modified_bessel_i0(input, *, out=None) -> Tensor + +Modified Bessel function of the first kind of order :math:`0`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +modified_bessel_i1 = _add_docstr( + _special.special_modified_bessel_i1, + r""" +modified_bessel_i1(input, *, out=None) -> Tensor + +Modified Bessel function of the first kind of order :math:`1`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +modified_bessel_k0 = _add_docstr( + _special.special_modified_bessel_k0, + r""" +modified_bessel_k0(input, *, out=None) -> Tensor + +Modified Bessel function of the second kind of order :math:`0`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +modified_bessel_k1 = _add_docstr( + _special.special_modified_bessel_k1, + r""" +modified_bessel_k1(input, *, out=None) -> Tensor + +Modified Bessel function of the second kind of order :math:`1`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +scaled_modified_bessel_k0 = _add_docstr( + _special.special_scaled_modified_bessel_k0, + r""" +scaled_modified_bessel_k0(input, *, out=None) -> Tensor + +Scaled modified Bessel function of the second kind of order :math:`0`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +scaled_modified_bessel_k1 = _add_docstr( + _special.special_scaled_modified_bessel_k1, + r""" +scaled_modified_bessel_k1(input, *, out=None) -> Tensor + +Scaled modified Bessel function of the second kind of order :math:`1`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) + +shifted_chebyshev_polynomial_t = _add_docstr( + _special.special_shifted_chebyshev_polynomial_t, + r""" +shifted_chebyshev_polynomial_t(input, n, *, out=None) -> Tensor + +Chebyshev polynomial of the first kind :math:`T_{n}^{\ast}(\text{input})`. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +shifted_chebyshev_polynomial_u = _add_docstr( + _special.special_shifted_chebyshev_polynomial_u, + r""" +shifted_chebyshev_polynomial_u(input, n, *, out=None) -> Tensor + +Chebyshev polynomial of the second kind :math:`U_{n}^{\ast}(\text{input})`. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +shifted_chebyshev_polynomial_v = _add_docstr( + _special.special_shifted_chebyshev_polynomial_v, + r""" +shifted_chebyshev_polynomial_v(input, n, *, out=None) -> Tensor + +Chebyshev polynomial of the third kind :math:`V_{n}^{\ast}(\text{input})`. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +shifted_chebyshev_polynomial_w = _add_docstr( + _special.special_shifted_chebyshev_polynomial_w, + r""" +shifted_chebyshev_polynomial_w(input, n, *, out=None) -> Tensor + +Chebyshev polynomial of the fourth kind :math:`W_{n}^{\ast}(\text{input})`. + +""" + + r""" +Args: + {input} + n (Tensor): Degree of the polynomial. + +Keyword args: + {out} +""".format(**common_args), +) + +spherical_bessel_j0 = _add_docstr( + _special.special_spherical_bessel_j0, + r""" +spherical_bessel_j0(input, *, out=None) -> Tensor + +Spherical Bessel function of the first kind of order :math:`0`. + +""" + + r""" +Args: + {input} + +Keyword args: + {out} +""".format(**common_args), +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/storage.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/storage.py new file mode 100644 index 0000000000000000000000000000000000000000..9bb671c0bd5106d75a457a0cf62d08068f20e2fb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/storage.py @@ -0,0 +1,1558 @@ +# mypy: allow-untyped-defs + +from __future__ import annotations + +import collections +import copy +import functools +import io +import threading +import warnings +from typing import Any, cast, TYPE_CHECKING, TypeVar +from typing_extensions import Self + +import torch +from torch._utils import _to, _type +from torch.types import _bool, _int, Storage + + +if TYPE_CHECKING: + from torch._prims_common import DeviceLikeType + + +__all__ = ["TypedStorage", "UntypedStorage"] + + +try: + import numpy as np + + HAS_NUMPY = True +except ModuleNotFoundError: + HAS_NUMPY = False + np = None # type: ignore[assignment] + + +_share_memory_lock = threading.Lock() +_share_memory_map: dict[int, threading.RLock] = {} + +T = TypeVar("T", bound="_StorageBase | TypedStorage") + + +class _StorageBase: + _cdata: Any + is_sparse: _bool = False + is_sparse_csr: _bool = False + device: torch.device + # Used when + # (1) stashing FakeTensor device onto storage in torch.serialization.skip_data + # (2) stashing device onto storage to propagate to FakeTensor when torch.load under FakeTensorMode + _fake_device: torch.device | None = None + # Used when loading with FakeTensorMode to give information about offset of storage in torch.saved-file + _checkpoint_offset: int | None = None + + def __init__(self, *args, **kwargs): + pass + + def __len__(self) -> _int: + raise NotImplementedError + + def __getitem__(self, idx): + raise NotImplementedError + + def __setitem__(self, *args, **kwargs): + raise NotImplementedError + + def copy_(self, source: T, non_blocking: _bool | None = None) -> T: + raise NotImplementedError + + def new(self) -> _StorageBase | TypedStorage: + raise NotImplementedError + + def nbytes(self) -> _int: + raise NotImplementedError + + def size(self) -> _int: + return self.nbytes() + + def type( + self, dtype: str | None = None, non_blocking: _bool = False + ) -> _StorageBase | TypedStorage: + return _type(self, dtype, non_blocking) + + def cuda(self, device=None, non_blocking=False) -> _StorageBase | TypedStorage: + """Returns a copy of this object in CUDA memory. + + If this object is already in CUDA memory and on the correct device, then + no copy is performed and the original object is returned. + + Args: + device (int): The destination GPU id. Defaults to the current device. + non_blocking (bool): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. Otherwise, + the argument has no effect. + """ + device2 = torch.device("cuda", device) if device else torch.device("cuda") + return self.to(device=device2, non_blocking=non_blocking) + + def hpu(self, device=None, non_blocking=False) -> _StorageBase | TypedStorage: + """Returns a copy of this object in HPU memory. + + If this object is already in HPU memory and on the correct device, then + no copy is performed and the original object is returned. + + Args: + device (int): The destination HPU id. Defaults to the current device. + non_blocking (bool): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. Otherwise, + the argument has no effect. + """ + device2 = torch.device("hpu", device) if device else torch.device("hpu") + return self.to(device=device2, non_blocking=non_blocking) + + def element_size(self) -> _int: + raise NotImplementedError + + def get_device(self) -> _int: + return self.device.index + + def data_ptr(self) -> _int: + raise NotImplementedError + + def resizable(self) -> _bool: + raise NotImplementedError + + # Defined in torch/csrc/generic/StorageSharing.cpp + def _share_filename_cpu_(self, *args, **kwargs): + raise NotImplementedError + + def _share_fd_cpu_(self, *args, **kwargs): + raise NotImplementedError + + @classmethod + def _new_using_filename_cpu(cls, size: _int) -> Self: + raise NotImplementedError + + @classmethod + def _new_using_fd_cpu(cls, size: _int) -> Self: + raise NotImplementedError + + @classmethod + def from_buffer(cls, *args, **kwargs) -> Self: + raise NotImplementedError + + @classmethod + def _new_shared_filename_cpu( + cls, + manager, + obj, + size, + *, + device=None, + dtype=None, + ) -> Self: + raise NotImplementedError + + @classmethod + def _release_ipc_counter(cls, *args, device=None, **kwargs): + return cls._release_ipc_counter_cuda(*args, **kwargs) + + @classmethod + def _release_ipc_counter_cuda(cls, *args, **kwargs) -> Self: + raise NotImplementedError + + @classmethod + def _new_with_weak_ptr(cls, *args, **kwargs) -> Self: + raise NotImplementedError + + def _shared_decref(self) -> _StorageBase | TypedStorage: + raise NotImplementedError + + def _write_file(self, *args, **kwargs): + raise NotImplementedError + + def resize_(self, size: _int): + raise NotImplementedError + + def _weak_ref(self, *args, **kwargs) -> _StorageBase | TypedStorage: + raise NotImplementedError + + def _set_from_file(self, *args, **kwargs): + raise NotImplementedError + + def _set_cdata(self, *args, **kwargs): + raise NotImplementedError + + def _share_cuda_(self, *args, **kwargs): + raise NotImplementedError + + def is_shared(self) -> _bool: + raise NotImplementedError + + @classmethod + def _new_shared_cuda(cls, *args, **kwargs) -> Self: + raise NotImplementedError + + def _shared_incref(self, *args, **kwargs): + raise NotImplementedError + + @classmethod + def _free_weak_ref(cls, *args, **kwargs): + raise NotImplementedError + + @property + def is_cuda(self): + raise NotImplementedError + + @property + def is_hpu(self): + raise NotImplementedError + + @classmethod + def from_file(cls, filename, shared, nbytes) -> _StorageBase | TypedStorage: + raise NotImplementedError + + @classmethod + def _expired(cls, *args, **kwargs) -> _StorageBase | TypedStorage: + raise NotImplementedError + + def _byteswap(self, *args, **kwargs): + raise NotImplementedError + + def _get_filename(self, *args, **kwargs) -> str | None: + raise NotImplementedError + + def __repr__(self): + info_str = f"[{torch.typename(self)}(device={self.device}) of size {len(self)}]" + if self.device.type == "meta": + return "...\n" + info_str + data_str = " " + "\n ".join(str(self[i]) for i in range(self.size())) + return data_str + "\n" + info_str + + def __iter__(self): + return iter(self[i] for i in range(self.size())) + + def __copy__(self): + return self.clone() + + def __deepcopy__(self, memo): + memo = memo.setdefault("torch", {}) + if self._cdata in memo: + return memo[self._cdata] + new_storage = self.clone() + memo[self._cdata] = new_storage + return new_storage + + def __reduce__(self): + b = io.BytesIO() + torch.save(self, b, _use_new_zipfile_serialization=False) + return (_load_from_bytes, (b.getvalue(),)) + + def __sizeof__(self): + return super().__sizeof__() + self.size() + + def clone(self): + """Return a copy of this storage.""" + return type(self)(self.nbytes(), device=self.device).copy_(self) + + def tolist(self): + """Return a list containing the elements of this storage.""" + return list(self) + + def cpu(self): + """Return a CPU copy of this storage if it's not already on the CPU.""" + if self.device.type != "cpu": + return torch.UntypedStorage(self.size()).copy_(self, False) + return self + + def mps(self): + """Return a MPS copy of this storage if it's not already on the MPS.""" + if self.device.type != "mps": + return torch.UntypedStorage(self.size(), device="mps").copy_(self, False) + return self + + def _to(self, dtype): + if not isinstance(dtype, torch.dtype): + raise TypeError(f"Argument 'dtype' must be torch.dtype, not {type(dtype)}") + storage = ( + torch.tensor([], dtype=torch.uint8, device=self.device) + .set_(cast(Storage, self)) + .to(dtype) + ._typed_storage() + ) + if storage.data_ptr() == self.data_ptr(): + storage = storage.clone() + return storage + + def to(self, *, device: DeviceLikeType, non_blocking: _bool = False): + if not isinstance(device, torch.device): + device = torch.device(device) + return _to(self, device, non_blocking) + + def double(self): + """Casts this storage to double type.""" + return self._to(torch.double) + + def float(self): + """Casts this storage to float type.""" + return self._to(torch.float) + + def half(self): + """Casts this storage to half type.""" + return self._to(torch.half) + + def long(self): + """Casts this storage to long type.""" + return self._to(torch.long) + + def int(self): + """Casts this storage to int type.""" + return self._to(torch.int) + + def short(self): + """Casts this storage to short type.""" + return self._to(torch.short) + + def char(self): + """Casts this storage to char type.""" + return self._to(torch.int8) + + def byte(self): + """Casts this storage to byte type.""" + return self._to(torch.uint8) + + def bool(self): + """Casts this storage to bool type.""" + return self._to(torch.bool) + + def bfloat16(self): + """Casts this storage to bfloat16 type.""" + return self._to(torch.bfloat16) + + def complex_double(self): + """Casts this storage to complex double type.""" + return self._to(torch.cdouble) + + def complex_float(self): + """Casts this storage to complex float type.""" + return self._to(torch.cfloat) + + def float8_e5m2(self): + """Casts this storage to float8_e5m2 type""" + return self._to(torch.float8_e5m2) + + def float8_e4m3fn(self): + """Casts this storage to float8_e4m3fn type""" + return self._to(torch.float8_e4m3fn) + + def float8_e5m2fnuz(self): + """Casts this storage to float8_e5m2fnuz type""" + return self._to(torch.float8_e5m2fnuz) + + def float8_e4m3fnuz(self): + """Casts this storage to float8_e4m3fnuz type""" + return self._to(torch.float8_e4m3fnuz) + + def is_pinned(self, device: str | torch.device = "cuda"): + r"""Determine whether the CPU storage is already pinned on device. + + Args: + device (str or torch.device): The device to pin memory on (default: ``'cuda'``). + This argument is discouraged and subject to deprecated. + + Returns: + A boolean variable. + """ + return ( + torch.tensor([], dtype=torch.uint8, device=self.device) + .set_(cast(Storage, self)) + .is_pinned(device) + ) + + def pin_memory(self, device: str | torch.device = "cuda"): + r"""Copy the CPU storage to pinned memory, if it's not already pinned. + + Args: + device (str or torch.device): The device to pin memory on (default: ``'cuda'``). + This argument is discouraged and subject to deprecated. + + Returns: + A pinned CPU storage. + """ + if self.device.type != "cpu": + raise TypeError(f"cannot pin '{self.type()}' only CPU memory can be pinned") + + pinned_tensor = ( + torch.tensor([], dtype=torch.uint8, device=self.device) + .set_(cast(Storage, self)) + .pin_memory(device) + ) + return pinned_tensor.untyped_storage() + + def share_memory_(self): + """See :meth:`torch.UntypedStorage.share_memory_`""" + from torch.multiprocessing import get_sharing_strategy + + if self.device.type in ["cuda", torch._C._get_privateuse1_backend_name()]: + pass # CUDA or PrivateUse1 doesn't use POSIX shared memory + elif get_sharing_strategy() == "file_system": + self._share_filename_cpu_() + else: + self._share_fd_cpu_() + return self + + @classmethod + def _new_shared(cls, size, *, device="cpu"): + """Create a new storage in shared memory with the same data type.""" + from torch.multiprocessing import get_sharing_strategy + + device = torch.device(device) + if device.type in ["cuda", torch._C._get_privateuse1_backend_name(), "hpu"]: + return cls(size, device=device) + elif get_sharing_strategy() == "file_system": + return cls._new_using_filename_cpu(size) + else: + return cls._new_using_fd_cpu(size) + + def untyped(self): + return self + + def byteswap(self, dtype): + """Swap bytes in underlying data.""" + elem_size = torch._utils._element_size(dtype) + # for complex types, don't swap first and second numbers + if dtype.is_complex: + elem_size = max(int(elem_size / 2), 1) + self._byteswap(elem_size) + + +def _share_memory_lock_protected(fn): + @functools.wraps(fn) + def wrapper(self, *args, **kwargs): + to_free = None + to_wait = None + with _share_memory_lock: + key = self._cdata + if key in _share_memory_map: + to_wait = _share_memory_map[key] + else: + _share_memory_map[key] = threading.RLock() + _share_memory_map[key].acquire() + to_free = key + + # If we're already in the process of sharing the storage, wait + # for it to be done. + if to_wait is not None: + with to_wait: + pass + + try: + return fn(self, *args, **kwargs) + finally: + # If we acquired the storage lock here and we're done working on it + # we can now release it and free the entry. + if to_free is not None: + # Ensure that the cdata from the storage didn't change and only + # the data_ptr did. + if self._cdata != to_free: + raise AssertionError( + f"storage cdata changed unexpectedly: expected {to_free}, got {self._cdata}" + ) + with _share_memory_lock: + _share_memory_map[to_free].release() + del _share_memory_map[to_free] + + return wrapper + + +class UntypedStorage(torch._C.StorageBase, _StorageBase): + def __getitem__(self, *args, **kwargs): + if self.device.type == "meta": + raise NotImplementedError("Not available for 'meta' device type") + return super().__getitem__(*args, **kwargs) + + @property + def is_cuda(self): + return self.device.type == "cuda" + + @property + def is_hpu(self): + return self.device.type == "hpu" + + @property + def filename(self) -> str | None: + """Returns the file name associated with this storage. + + The file name will be a string if the storage is on CPU and was created via + :meth:`~torch.from_file()` with ``shared`` as ``True``. This attribute is ``None`` otherwise. + """ + return self._get_filename() + + @_share_memory_lock_protected + def share_memory_(self, *args, **kwargs): + """ + Moves the storage to shared memory. + + This is a no-op for storages already in shared memory and for CUDA + storages, which do not need to be moved for sharing across processes. + Storages in shared memory cannot be resized. + + Note that to mitigate issues like `this `_ + it is thread safe to call this function from multiple threads on the same object. + It is NOT thread safe though to call any other function on self without proper + synchronization. Please see :doc:`/notes/multiprocessing` for more details. + + .. note:: + When all references to a storage in shared memory are deleted, the associated shared memory + object will also be deleted. PyTorch has a special cleanup process to ensure that this happens + even if the current process exits unexpectedly. + + It is worth noting the difference between :meth:`share_memory_` and :meth:`from_file` with ``shared = True`` + + #. ``share_memory_`` uses `shm_open(3) `_ to create a + POSIX shared memory object while :meth:`from_file` uses + `open(2) `_ to open the filename passed by the user. + #. Both use an `mmap(2) call `_ with ``MAP_SHARED`` + to map the file/object into the current virtual address space + #. ``share_memory_`` will call ``shm_unlink(3)`` on the object after mapping it to make sure the shared memory + object is freed when no process has the object open. ``torch.from_file(shared=True)`` does not unlink the + file. This file is persistent and will remain until it is deleted by the user. + + Returns: + ``self`` + """ + return super().share_memory_(*args, **kwargs) + + @_share_memory_lock_protected + def _share_fd_cpu_(self, *args, **kwargs): + return super()._share_fd_cpu_(*args, **kwargs) + + @_share_memory_lock_protected + def _share_filename_cpu_(self, *args, **kwargs): + return super()._share_filename_cpu_(*args, **kwargs) + + +def _load_from_bytes(b): + return torch.load(io.BytesIO(b), weights_only=False) + + +@functools.cache +def _new_dtypes(): + # These are dtypes serialized as UntypedStorage unlike those in + # _dtype_to_storage_type_map + return { + torch.float8_e5m2, + torch.float8_e4m3fn, + torch.float8_e5m2fnuz, + torch.float8_e4m3fnuz, + torch.float8_e8m0fnu, + torch.float4_e2m1fn_x2, + torch.bits8, + torch.bits16, + torch.bits1x8, + torch.bits2x4, + torch.bits4x2, + torch.complex32, + torch.uint16, + torch.uint32, + torch.uint64, + } + + +@functools.cache +def _dtype_to_storage_type_map(): + # NOTE: We should no longer add dtypes to this map. This map + # is only used for BC/FC with older PyTorch versions. Going forward, + # new dtypes of TypedStorage should not translate to a legacy + # Storage class. Instead, new dtypes of TypedStorage should + # be serialized as an UntypedStorage paired with a torch.dtype + return { + torch.double: "DoubleStorage", + torch.float: "FloatStorage", + torch.half: "HalfStorage", + torch.long: "LongStorage", + torch.int: "IntStorage", + torch.int16: "ShortStorage", + torch.int8: "CharStorage", + torch.uint8: "ByteStorage", + torch.bool: "BoolStorage", + torch.bfloat16: "BFloat16Storage", + torch.cdouble: "ComplexDoubleStorage", + torch.cfloat: "ComplexFloatStorage", + torch.qint8: "QInt8Storage", + torch.qint32: "QInt32Storage", + torch.quint8: "QUInt8Storage", + torch.quint4x2: "QUInt4x2Storage", + torch.quint2x4: "QUInt2x4Storage", + } + + +@functools.cache +def _storage_type_to_dtype_map(): + dtype_map = {val: key for key, val in _dtype_to_storage_type_map().items()} + return dtype_map + + +def _get_storage_from_sequence(sequence, dtype, device): + if dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + interpret_dtypes = { + torch.quint8: torch.uint8, + torch.quint4x2: torch.uint8, + torch.quint2x4: torch.uint8, + torch.qint32: torch.int32, + torch.qint8: torch.int8, + } + tmp_tensor = torch.tensor( + sequence, dtype=interpret_dtypes[dtype], device=device + ) + + else: + tmp_tensor = torch.tensor(sequence, dtype=dtype, device=device) + + return tmp_tensor._typed_storage()._untyped_storage + + +def _isint(x): + if HAS_NUMPY: + return isinstance(x, (int, np.integer)) # pyrefly: ignore [missing-attribute] + else: + return isinstance(x, int) + + +_always_warn_typed_storage_removal = False + + +def _get_always_warn_typed_storage_removal(): + return _always_warn_typed_storage_removal + + +def _set_always_warn_typed_storage_removal(always_warn): + global _always_warn_typed_storage_removal + if not isinstance(always_warn, bool): + raise AssertionError( + f"always_warn must be bool, got {type(always_warn).__name__}" + ) + _always_warn_typed_storage_removal = always_warn + + +def _warn_typed_storage_removal(stacklevel=2): + global _always_warn_typed_storage_removal + + def is_first_time(): + if not hasattr(_warn_typed_storage_removal, "has_warned"): + return True + else: + return not _warn_typed_storage_removal.__dict__["has_warned"] + + if _get_always_warn_typed_storage_removal() or is_first_time(): + message = ( + "TypedStorage is deprecated. It will be removed in the future and " + "UntypedStorage will be the only storage class. This should only matter " + "to you if you are using storages directly. To access UntypedStorage " + "directly, use tensor.untyped_storage() instead of tensor.storage()" + ) + warnings.warn(message, UserWarning, stacklevel=stacklevel + 1) + _warn_typed_storage_removal.__dict__["has_warned"] = True + + +def _reset_warn_typed_storage_removal(): + _warn_typed_storage_removal.__dict__["has_warned"] = False + + +def _get_device_from_module(module: str): + last_part = module.rsplit(".", 1)[-1] + if last_part in ["cuda", torch._C._get_privateuse1_backend_name(), "hpu"]: + return last_part + else: + return "cpu" + + +class TypedStorage: + is_sparse: _bool = False + # Used when stashing FakeTensor device onto storage in torch.save(metadata_only=True) + _fake_device: torch.device | None = None + + dtype: torch.dtype + + @property + def _dtype(self): + return self.dtype + + @property + def filename(self) -> str | None: + """Returns the file name associated with this storage if the storage was memory mapped from a file. + or ``None`` if the storage was not created by memory mapping a file.""" + return self._untyped_storage.filename + + def fill_(self, value): + _warn_typed_storage_removal() + self._setitem(slice(0, self._size()), value) + return self + + def __new__( + cls, + *args, + wrap_storage=None, + dtype=None, + device=None, + _internal=False, + ): + if not _internal: + _warn_typed_storage_removal() + + if cls == torch.storage._LegacyStorage: + raise RuntimeError( + "Only child classes of _LegacyStorage can be instantiated" + ) + + if cls == TypedStorage: + return super().__new__(cls) + + else: + arg_error_msg = ( + f"{cls}.__new__ received an invalid combination " + f"of arguments. Expected one of:\n" + " * no arguments\n" + " * (int size)\n" + " * (Sequence data)\n" + " * (*, UntypedStorage wrap_storage)" + ) + + if device is not None: + raise RuntimeError( + arg_error_msg + "\nKeyword argument 'device' cannot be specified" + ) + + if dtype is not None: + raise RuntimeError( + arg_error_msg + "\nKeyword argument 'dtype' cannot be specified" + ) + + if wrap_storage is None: + if len(args) > 1: + raise RuntimeError( + arg_error_msg + "\nToo many positional arguments" + ) + + if ( + len(args) == 1 + and not _isint(args[0]) + and not isinstance(args[0], collections.abc.Sequence) + ): + raise TypeError( + arg_error_msg + + f"\nArgument type not recognized: {type(args[0])}" + ) + + return TypedStorage( + *args, + dtype=cls._dtype, + device=_get_device_from_module(cls.__module__), + _internal=True, + ) + + else: + if len(args) != 0: + raise RuntimeError( + arg_error_msg + + "\nNo positional arguments should be given when using " + "'wrap_storage'" + ) + + if not isinstance(wrap_storage, torch.UntypedStorage): + raise TypeError( + arg_error_msg + + f"\nArgument 'wrap_storage' must be UntypedStorage, but got {type(wrap_storage)}" + ) + + cls_device = _get_device_from_module(cls.__module__) + + if wrap_storage.device.type != cls_device: + raise RuntimeError( + arg_error_msg + + f"\nDevice of 'wrap_storage' must be {cls_device}" + f", but got {wrap_storage.device.type}" + ) + + return TypedStorage( + *args, + wrap_storage=wrap_storage, + dtype=cls.dtype, + _internal=True, + ) + + def __init__( + self, + *args, + device=None, + dtype=None, + wrap_storage=None, + _internal=False, + ): + if not _internal: + _warn_typed_storage_removal() + arg_error_msg = ( + "TypedStorage.__init__ received an invalid combination " + "of arguments. Expected one of:\n" + " * (*, torch.device device, torch.dtype dtype)\n" + " * (int size, *, torch.device device, torch.dtype dtype)\n" + " * (Sequence data, *, torch.device device, torch.dtype dtype)\n" + " * (*, UntypedStorage wrap_storage, torch.dtype dtype)" + ) + + if wrap_storage is not None: + if len(args) != 0: + raise RuntimeError( + arg_error_msg + + "\nNo positional arguments should be given when using " + "'wrap_storage'" + ) + + if dtype is None: + raise RuntimeError( + arg_error_msg + "\nArgument 'dtype' must be specified" + ) + + if not isinstance(dtype, torch.dtype): + raise TypeError( + arg_error_msg + + f"\nArgument 'dtype' must be torch.dtype, not {type(dtype)}" + ) + + if device is not None: + raise RuntimeError( + arg_error_msg + + "\nArgument 'device' should not be specified when 'wrap_storage' is given" + ) + + self.dtype = dtype + + if not isinstance(wrap_storage, torch.UntypedStorage): + raise TypeError( + arg_error_msg + + f"\nArgument 'wrap_storage' must be UntypedStorage, but got {type(wrap_storage)}" + ) + + self._untyped_storage = wrap_storage + + else: + self.dtype = torch.get_default_dtype() if dtype is None else dtype + device = torch.device("cpu" if device is None else device) + + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + if device.type == "cuda": + raise RuntimeError( + "Cannot create CUDA storage with quantized dtype" + ) + + if len(args) == 0: + self._untyped_storage = torch.UntypedStorage(device=device) + + elif len(args) == 1: + if _isint(args[0]): + self._untyped_storage = torch.UntypedStorage( + int(args[0]) * self._element_size(), device=device + ) + elif isinstance(args[0], collections.abc.Sequence): + self._untyped_storage = _get_storage_from_sequence( + args[0], self.dtype, device + ) + else: + raise TypeError( + arg_error_msg + + f"\nArgument type not recognized: {type(args[0])}" + ) + + else: + raise RuntimeError(arg_error_msg + "\nToo many positional arguments") + + @property + def is_cuda(self): + _warn_typed_storage_removal() + return self._untyped_storage.device.type == "cuda" + + @property + def is_hpu(self): + _warn_typed_storage_removal() + return self._untyped_storage.device.type == "hpu" + + def untyped(self): + """Return the internal :class:`torch.UntypedStorage`.""" + _warn_typed_storage_removal() + return self._untyped_storage + + def _new_wrapped_storage(self, untyped_storage) -> Self: + if type(untyped_storage) is not torch.UntypedStorage: + raise AssertionError( + f"expected UntypedStorage, got {type(untyped_storage).__name__}" + ) + + if type(self) is TypedStorage: + return cast( + Self, + TypedStorage( + wrap_storage=untyped_storage, dtype=self.dtype, _internal=True + ), + ) + else: + return type(self)(wrap_storage=untyped_storage) + + def __len__(self): + _warn_typed_storage_removal() + return self._size() + + def _maybe_wrap_index(self, idx, is_stop=False): + if idx is None: + if is_stop: + return self._size() + else: + return 0 + + else: + if type(idx) is not int: + raise TypeError(f"can't index a {type(self)} with {type(idx)}") + if is_stop: + if (idx > self._size()) or (idx < -self._size()): + raise IndexError( + f"index {idx} out of range for storage of size {self.size()}" + ) + if idx > 0: + return idx + else: + return idx % self._size() + else: + if (idx >= self._size()) or (idx < -self._size()): + raise IndexError( + f"index {idx} out of range for storage of size {self.size()}" + ) + return idx % self._size() + + def __setitem__(self, idx, value): + _warn_typed_storage_removal() + return self._setitem(idx, value) + + def _setitem(self, idx, value): + if not isinstance(idx, (int, slice)): + raise RuntimeError(f"can't index a {type(self)} with {type(idx)}") + if torch.is_storage(value): + raise RuntimeError(f"cannot set item with value type {type(value)}") + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + interpret_dtypes = { + torch.quint8: torch.uint8, + torch.quint4x2: torch.uint8, + torch.quint2x4: torch.uint8, + torch.qint32: torch.int32, + torch.qint8: torch.int8, + } + tmp_dtype = interpret_dtypes[self.dtype] + tmp_tensor = torch.tensor( + [], dtype=tmp_dtype, device=self._untyped_storage.device + ) + tmp_tensor.set_( + TypedStorage( + wrap_storage=self._untyped_storage, dtype=tmp_dtype, _internal=True + ) + ) + else: + tmp_tensor = torch.tensor( + [], dtype=self.dtype, device=self._untyped_storage.device + ).set_(self) + + tmp_tensor[idx] = value + + def __getitem__(self, idx): + _warn_typed_storage_removal() + return self._getitem(idx) + + def _getitem(self, idx): + if self._untyped_storage.device.type == "meta": + raise NotImplementedError("Not available for 'meta' device type") + + # NOTE: Before TypedStorage existed, indexing with a slice used to be + # possible for Storage objects. However, it would return + # a storage view, which would be a hassle to implement in TypedStorage, + # so it was disabled + if isinstance(idx, slice): + raise RuntimeError( + "slices are only supported in UntypedStorage.__getitem__" + ) + elif not isinstance(idx, int): + raise RuntimeError(f"can't index a {type(self)} with {type(idx)}") + + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + interpret_dtypes = { + torch.quint8: torch.uint8, + torch.quint4x2: torch.uint8, + torch.quint2x4: torch.uint8, + torch.qint32: torch.int32, + torch.qint8: torch.int8, + } + return TypedStorage( + wrap_storage=self._untyped_storage, + dtype=interpret_dtypes[self.dtype], + _internal=True, + )._getitem(idx) + + idx_wrapped = self._maybe_wrap_index(idx) + from torch._subclasses.fake_tensor import unset_fake_temporarily + + with unset_fake_temporarily(): + tmp_tensor = torch.tensor( + [], dtype=self.dtype, device=self._untyped_storage.device + ).set_(self) + return tmp_tensor[idx_wrapped].item() + + def copy_(self, source: T, non_blocking: bool | None = None): + _warn_typed_storage_removal() + if isinstance(source, TypedStorage): + self._untyped_storage.copy_(source._untyped_storage, non_blocking) + else: + self._untyped_storage.copy_(source, non_blocking) + return self + + def nbytes(self): + _warn_typed_storage_removal() + return self._nbytes() + + # For internal use only, to avoid deprecation warning + def _nbytes(self): + return self._untyped_storage.nbytes() + + def type( + self, + dtype: str | None = None, + non_blocking: bool = False, + ) -> _StorageBase | TypedStorage | str: + _warn_typed_storage_removal() + if dtype is None: + legacy_class = self._get_legacy_storage_class() + + if legacy_class is not None: + return legacy_class.__module__ + "." + legacy_class.__name__ + + return ".".join([self.__module__, type(self).__name__]) + + else: + return self._untyped_storage.type(dtype, non_blocking) + + def cuda(self, device=None, non_blocking=False) -> Self: + _warn_typed_storage_removal() + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + raise RuntimeError("Cannot create CUDA storage with quantized dtype") + cuda_storage = self._untyped_storage.cuda(device, non_blocking) + return self._new_wrapped_storage(cuda_storage) + + def hpu(self, device=None, non_blocking=False) -> Self: + _warn_typed_storage_removal() + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + raise RuntimeError("Cannot create HPU storage with quantized dtype") + hpu_storage = self._untyped_storage.hpu(device, non_blocking) + return self._new_wrapped_storage(hpu_storage) + + def to(self, *, device: DeviceLikeType, non_blocking: bool = False) -> Self: + _warn_typed_storage_removal() + if not isinstance(device, torch.device): + device = torch.device(device) + if self.dtype in [ + torch.quint8, + torch.quint4x2, + torch.quint2x4, + torch.qint32, + torch.qint8, + ]: + raise RuntimeError( + f"Cannot create {device.type.upper()} storage with quantized dtype" + ) + to_storage = self._untyped_storage.to(device=device, non_blocking=non_blocking) + return self._new_wrapped_storage(to_storage) + + def element_size(self): + _warn_typed_storage_removal() + return self._element_size() + + # For internal use only, to avoid deprecation warning + def _element_size(self): + return torch._utils._element_size(self.dtype) + + def get_device(self) -> _int: + _warn_typed_storage_removal() + return self._untyped_storage.get_device() + + def __str__(self): + _warn_typed_storage_removal() + info_str = ( + f"[{torch.typename(self)}(dtype={self.dtype}, " + f"device={self.device}) of size {len(self)}]" + ) + if self.device.type == "meta": + return "...\n" + info_str + else: + data_str = " " + "\n ".join(str(self[i]) for i in range(self.size())) + return data_str + "\n" + info_str + + def __repr__(self): + _warn_typed_storage_removal() + return str(self) + + def __iter__(self): + _warn_typed_storage_removal() + return iter(self[i] for i in range(self.size())) + + def __copy__(self): + _warn_typed_storage_removal() + return self._new_wrapped_storage(copy.copy(self._untyped_storage)) + + def __deepcopy__(self, memo): + _warn_typed_storage_removal() + return self._deepcopy(memo) + + # For internal use only, to avoid deprecation warning + def _deepcopy(self, memo): + return self._new_wrapped_storage(copy.deepcopy(self._untyped_storage, memo)) + + def __sizeof__(self): + _warn_typed_storage_removal() + return super().__sizeof__() + self.nbytes() + + def clone(self): + """Return a copy of this storage.""" + _warn_typed_storage_removal() + return self._new_wrapped_storage(self._untyped_storage.clone()) + + def tolist(self): + """Return a list containing the elements of this storage.""" + _warn_typed_storage_removal() + return list(self) + + def cpu(self): + """Return a CPU copy of this storage if it's not already on the CPU.""" + _warn_typed_storage_removal() + return self._new_wrapped_storage(self._untyped_storage.cpu()) + + def is_pinned(self, device: str | torch.device = "cuda"): + r"""Determine whether the CPU TypedStorage is already pinned on device. + + Args: + device (str or torch.device): The device to pin memory on (default: ``'cuda'``). + This argument is discouraged and subject to deprecated. + + Returns: + A boolean variable. + """ + _warn_typed_storage_removal() + return self._untyped_storage.is_pinned(device) + + def pin_memory(self, device: str | torch.device = "cuda"): + r"""Copy the CPU TypedStorage to pinned memory, if it's not already pinned. + + Args: + device (str or torch.device): The device to pin memory on (default: ``'cuda'``). + This argument is discouraged and subject to deprecated. + + Returns: + A pinned CPU storage. + """ + _warn_typed_storage_removal() + return self._new_wrapped_storage( + self._untyped_storage.pin_memory(device=device) + ) + + def share_memory_(self): + """See :meth:`torch.UntypedStorage.share_memory_`""" + _warn_typed_storage_removal() + return self._share_memory_() + + # For internal use only, to avoid deprecation warning + def _share_memory_(self): + self._untyped_storage.share_memory_() + return self + + def _new_shared(self, size, *, device=None): + """Create a new storage in shared memory with the same data type.""" + if device is None: + device = "cpu" + device = torch.device(device) + untyped_storage = torch.UntypedStorage._new_shared( + size * self._element_size(), device=device + ) + return TypedStorage( + wrap_storage=untyped_storage, dtype=self.dtype, _internal=True + ) + + @property + def _cdata(self): + return self._untyped_storage._cdata + + @property + def device(self): + _warn_typed_storage_removal() + return self._untyped_storage.device + + def size(self): + _warn_typed_storage_removal() + return self._size() + + # For internal use only, to avoid deprecation warning + def _size(self): + # NB: don't indirect through __len__, as that requires + # an int to be returned + return self._untyped_storage.nbytes() // self._element_size() + + def pickle_storage_type(self): + _warn_typed_storage_removal() + return self._pickle_storage_type() + + # For internal use only, to avoid deprecation warning + def _pickle_storage_type(self): + try: + return _dtype_to_storage_type_map()[self.dtype] + except KeyError as e: + raise KeyError(f"dtype {self.dtype} is not recognized") from e + + def __reduce__(self): + b = io.BytesIO() + torch.save(self, b, _use_new_zipfile_serialization=False) + return (_load_from_bytes, (b.getvalue(),)) + + def data_ptr(self): + _warn_typed_storage_removal() + return self._data_ptr() + + # For internal use only, to avoid deprecation warning + def _data_ptr(self): + return self._untyped_storage.data_ptr() + + def resizable(self): + _warn_typed_storage_removal() + return self._untyped_storage.resizable() + + def resize_(self, size): + _warn_typed_storage_removal() + self._resize_(size) + + # For internal use only, to avoid deprecation warning + def _resize_(self, size): + self._untyped_storage.resize_(size * self._element_size()) + + @classmethod + def _free_weak_ref(cls, *args, **kwargs): + return UntypedStorage._free_weak_ref(*args, **kwargs) + + def _weak_ref(self, *args, **kwargs): + return self._untyped_storage._weak_ref(*args, **kwargs) + + @classmethod + def from_buffer(cls, *args, **kwargs): + _warn_typed_storage_removal() + return cls._from_buffer(*args, **kwargs) + + @classmethod + def _from_buffer(cls, *args, dtype=None, device=None, **kwargs): + if cls == TypedStorage: + dtype = torch.get_default_dtype() if dtype is None else dtype + device = torch.device("cpu" if device is None else device) + if device.type != "cpu": + raise RuntimeError( + f"TypedStorage.from_buffer: Not available for device {device.type}" + ) + untyped_storage: torch.UntypedStorage = torch.UntypedStorage.from_buffer( + *args, dtype=dtype, **kwargs + ) + + else: + if dtype is not None or len(args) == 5: + raise RuntimeError( + "from_buffer: 'dtype' can only be specified in " + "UntypedStorage.from_buffer and TypedStorage.from_buffer" + ) + if device is not None: + raise RuntimeError( + "from_buffer: 'device' can only be specified in " + "UntypedStorage.from_buffer and TypedStorage.from_buffer" + ) + + dtype = cls._dtype + untyped_storage = torch.UntypedStorage.from_buffer( + *args, dtype=dtype, **kwargs + ) + + return TypedStorage(wrap_storage=untyped_storage, dtype=dtype, _internal=True) + + def _to(self, dtype): + if not isinstance(dtype, torch.dtype): + raise TypeError(f"Argument 'dtype' must be torch.dtype, not {type(dtype)}") + storage = ( + torch.tensor([], dtype=self.dtype, device=self.device) + .set_(self) + .to(dtype) + ._typed_storage() + ) + if storage.data_ptr() == self.data_ptr(): + storage = storage.clone() + return storage + + def double(self): + """Casts this storage to double type.""" + _warn_typed_storage_removal() + return self._to(torch.double) + + def float(self): + """Casts this storage to float type.""" + _warn_typed_storage_removal() + return self._to(torch.float) + + def half(self): + """Casts this storage to half type.""" + _warn_typed_storage_removal() + return self._to(torch.half) + + def long(self): + """Casts this storage to long type.""" + _warn_typed_storage_removal() + return self._to(torch.long) + + def int(self): + """Casts this storage to int type.""" + _warn_typed_storage_removal() + return self._to(torch.int) + + def short(self): + """Casts this storage to short type.""" + _warn_typed_storage_removal() + return self._to(torch.short) + + def char(self): + """Casts this storage to char type.""" + _warn_typed_storage_removal() + return self._to(torch.int8) + + def byte(self): + """Casts this storage to byte type.""" + _warn_typed_storage_removal() + return self._to(torch.uint8) + + def bool(self): + """Casts this storage to bool type.""" + _warn_typed_storage_removal() + return self._to(torch.bool) + + def bfloat16(self): + """Casts this storage to bfloat16 type.""" + _warn_typed_storage_removal() + return self._to(torch.bfloat16) + + def complex_double(self): + """Casts this storage to complex double type.""" + _warn_typed_storage_removal() + return self._to(torch.cdouble) + + def complex_float(self): + """Casts this storage to complex float type.""" + _warn_typed_storage_removal() + return self._to(torch.cfloat) + + def float8_e5m2(self): + """Casts this storage to float8_e5m2 type""" + _warn_typed_storage_removal() + return self._to(torch.float8_e5m2) + + def float8_e4m3fn(self): + """Casts this storage to float8_e4m3fn type""" + _warn_typed_storage_removal() + return self._to(torch.float8_e4m3fn) + + def float8_e5m2fnuz(self): + """Casts this storage to float8_e5m2fnuz type""" + _warn_typed_storage_removal() + return self._to(torch.float8_e5m2fnuz) + + def float8_e4m3fnuz(self): + """Casts this storage to float8_e4m3fnuz type""" + _warn_typed_storage_removal() + return self._to(torch.float8_e4m3fnuz) + + @classmethod + def from_file(cls, filename, shared, size): + """from_file(filename, shared=False, size=0) -> Storage + + Creates a CPU storage backed by a memory-mapped file. + + If ``shared`` is ``True``, then memory is shared between all processes. + All changes are written to the file. If ``shared`` is ``False``, then the changes on + the storage do not affect the file. + + ``size`` is the number of elements in the storage. If ``shared`` is ``False``, + then the file must contain at least ``size * sizeof(Type)`` bytes + (``Type`` is the type of storage). If ``shared`` is ``True`` the file will be created if needed. + + Args: + filename (str): file name to map + shared (bool): whether to share memory (whether ``MAP_SHARED`` or ``MAP_PRIVATE`` is passed to the + underlying `mmap(2) call `_) + size (int): number of elements in the storage + """ + _warn_typed_storage_removal() + if cls == TypedStorage: + raise RuntimeError("from_file can only be called on derived classes") + untyped_storage = UntypedStorage.from_file( + filename, shared, size * torch._utils._element_size(cls.dtype) + ) + storage = cls(wrap_storage=untyped_storage) + return storage + + @classmethod + def _expired(cls, *args, **kwargs): + return UntypedStorage._expired(*args, **kwargs) + + def _write_file(self, *args, **kwargs): + return self._untyped_storage._write_file(*args, **kwargs) + + def _set_from_file(self, *args, **kwargs): + return self._untyped_storage._set_from_file(*args, **kwargs) + + def _set_cdata(self, *args, **kwargs): + return self._untyped_storage._set_cdata(*args, **kwargs) + + def _share_cuda_(self, *args, **kwargs): + return self._untyped_storage._share_cuda_(*args, **kwargs) + + def is_shared(self): + _warn_typed_storage_removal() + return self._is_shared() + + # For internal use only, to avoid deprecation warning + def _is_shared(self): + return self._untyped_storage.is_shared() + + @classmethod + def _new_shared_cuda(cls, *args, **kwargs): + return torch.UntypedStorage._new_shared_cuda(*args, **kwargs) + + def _share_filename_cpu_(self, *args, **kwargs): + ( + manager_handle, + storage_handle, + size, + ) = self._untyped_storage._share_filename_cpu_(*args, **kwargs) + return manager_handle, storage_handle, size // self._element_size() + + def _shared_decref(self): + self._untyped_storage._shared_decref() + return self + + @classmethod + def _release_ipc_counter(cls, *args, device=None, **kwargs): + return torch.UntypedStorage._release_ipc_counter_cuda(*args, **kwargs) + + def _shared_incref(self, *args, **kwargs): + return self._untyped_storage._shared_incref(*args, **kwargs) + + def _share_fd_cpu_(self, *args, **kwargs): + fd, size = self._untyped_storage._share_fd_cpu_(*args, **kwargs) + return fd, size // self._element_size() + + def _get_legacy_storage_class(self): + if self.dtype not in _dtype_to_storage_type_map(): + return None + + storage_name = _dtype_to_storage_type_map()[self.dtype] + + if self.device.type not in [ + "cpu", + "cuda", + "hpu", + torch._C._get_privateuse1_backend_name(), + ]: + return None + + module = ( + torch if self.device.type == "cpu" else getattr(torch, self.device.type) + ) + + try: + return getattr(module, storage_name) + except AttributeError: + return None + + +TypedStorage.type.__doc__ = _type.__doc__ +TypedStorage.cuda.__doc__ = _StorageBase.cuda.__doc__ +TypedStorage.hpu.__doc__ = _StorageBase.hpu.__doc__ +TypedStorage.to.__doc__ = _to.__doc__ + + +class _LegacyStorageMeta(type): + dtype: torch.dtype + + def __instancecheck__(cls, instance): + if type(instance) is TypedStorage: + cls_device = _get_device_from_module(cls.__module__) + return (cls_device == instance.device.type) and ( + cls.dtype == instance.dtype + ) + return False + + +class _LegacyStorage(TypedStorage, metaclass=_LegacyStorageMeta): + @classmethod + def _new_shared(cls, size): # type: ignore[override] + """Create a new storage in shared memory with the same data type.""" + untyped_storage = torch.UntypedStorage._new_shared(size * cls()._element_size()) + return cls(wrap_storage=untyped_storage) + + @classmethod + def _release_ipc_counter(cls, *args, **kwargs): + return torch.UntypedStorage._release_ipc_counter_cuda(*args, **kwargs) + + @classmethod + def _new_shared_filename(cls, manager, obj, size): + bytes_size = size * torch._utils._element_size(cls.dtype) + return cls( + wrap_storage=torch.UntypedStorage._new_shared_filename_cpu( + manager, obj, bytes_size + ) + ) + + +def _get_dtype_from_pickle_storage_type(pickle_storage_type: str): + try: + return _storage_type_to_dtype_map()[pickle_storage_type] + except KeyError as e: + raise KeyError( + f'pickle storage type "{pickle_storage_type}" is not recognized' + ) from e diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6724bd3d523b09a582cef5ed256f4c14b281fa72 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/__init__.py @@ -0,0 +1,7 @@ +from torch._C import FileCheck as FileCheck + +from . import _utils + +# pyrefly: ignore [deprecated] +from ._comparison import assert_allclose, assert_close as assert_close +from ._creation import make_tensor as make_tensor diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_comparison.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_comparison.py new file mode 100644 index 0000000000000000000000000000000000000000..ef60dd19f7c3e4ecc34f5a9e8ab7d32d3b2c7022 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_comparison.py @@ -0,0 +1,1681 @@ +# mypy: allow-untyped-defs +import abc +import cmath +import collections.abc +import contextlib +from collections.abc import Callable, Collection, Sequence +from typing import Any, NoReturn +from typing_extensions import deprecated + +import torch + + +try: + import numpy as np + + HAS_NUMPY = True +except ModuleNotFoundError: + HAS_NUMPY = False + np = None # type: ignore[assignment] + +_HAS_DTENSOR = torch.distributed.is_available() + + +def _unwrap_dtensor_for_comparison(actual, expected): + """Handle DTensor inputs for assertEqual/assert_close.""" + if not _HAS_DTENSOR: + return actual, expected + from torch.distributed.tensor import DTensor + + actual_dt = isinstance(actual, DTensor) + expected_dt = isinstance(expected, DTensor) + if actual_dt and expected_dt: + if actual.placements != expected.placements: + raise AssertionError( + f"DTensor placements do not match: " + f"{actual.placements} != {expected.placements}" + ) + if actual.device_mesh != expected.device_mesh: + raise AssertionError( + f"DTensor device meshes do not match: " + f"{actual.device_mesh} != {expected.device_mesh}" + ) + return actual.to_local(), expected.to_local() + elif actual_dt != expected_dt: + raise TypeError( + "Comparing a DTensor to a non-DTensor is ambiguous. " + "Call .full_tensor() to compare the full logical tensor " + "or .to_local() to compare the local shard." + ) + return actual, expected + + +class ErrorMeta(Exception): + """Internal testing exception that makes that carries error metadata.""" + + def __init__( + self, type: type[Exception], msg: str, *, id: tuple[Any, ...] = () + ) -> None: + super().__init__( + "If you are a user and see this message during normal operation " + "please file an issue at https://github.com/pytorch/pytorch/issues. " + "If you are a developer and working on the comparison functions, please `raise ErrorMeta.to_error()` " + "for user facing errors." + ) + self.type = type + self.msg = msg + self.id = id + + def to_error(self, msg: str | Callable[[str], str] | None = None) -> Exception: + if not isinstance(msg, str): + generated_msg = self.msg + if self.id: + generated_msg += f"\n\nThe failure occurred for item {''.join(str([item]) for item in self.id)}" + + msg = msg(generated_msg) if callable(msg) else generated_msg + + return self.type(msg) + + +# Some analysis of tolerance by logging tests from test_torch.py can be found in +# https://github.com/pytorch/pytorch/pull/32538. +# {dtype: (rtol, atol)} +_DTYPE_PRECISIONS = { + torch.float16: (0.001, 1e-5), + torch.bfloat16: (0.016, 1e-5), + torch.float32: (1.3e-6, 1e-5), + torch.float64: (1e-7, 1e-7), + torch.complex32: (0.001, 1e-5), + torch.complex64: (1.3e-6, 1e-5), + torch.complex128: (1e-7, 1e-7), +} +# The default tolerances of torch.float32 are used for quantized dtypes, because quantized tensors are compared in +# their dequantized and floating point representation. For more details see `TensorLikePair._compare_quantized_values` +_DTYPE_PRECISIONS.update( + dict.fromkeys( + (torch.quint8, torch.quint2x4, torch.quint4x2, torch.qint8, torch.qint32), + _DTYPE_PRECISIONS[torch.float32], + ) +) + + +def default_tolerances( + *inputs: torch.Tensor | torch.dtype, + dtype_precisions: dict[torch.dtype, tuple[float, float]] | None = None, +) -> tuple[float, float]: + """Returns the default absolute and relative testing tolerances for a set of inputs based on the dtype. + + See :func:`assert_close` for a table of the default tolerance for each dtype. + + Returns: + (Tuple[float, float]): Loosest tolerances of all input dtypes. + """ + dtypes = [] + for input in inputs: + if isinstance(input, torch.Tensor): + dtypes.append(input.dtype) + elif isinstance(input, torch.dtype): + dtypes.append(input) + else: + raise TypeError( + f"Expected a torch.Tensor or a torch.dtype, but got {type(input)} instead." + ) + dtype_precisions = dtype_precisions or _DTYPE_PRECISIONS + rtols, atols = zip( + *[dtype_precisions.get(dtype, (0.0, 0.0)) for dtype in dtypes], strict=True + ) + return max(rtols), max(atols) + + +def get_tolerances( + *inputs: torch.Tensor | torch.dtype, + rtol: float | None, + atol: float | None, + id: tuple[Any, ...] = (), +) -> tuple[float, float]: + """Gets absolute and relative to be used for numeric comparisons. + + If both ``rtol`` and ``atol`` are specified, this is a no-op. If both are not specified, the return value of + :func:`default_tolerances` is used. + + Raises: + ErrorMeta: With :class:`ValueError`, if only ``rtol`` or ``atol`` is specified. + + Returns: + (Tuple[float, float]): Valid absolute and relative tolerances. + """ + if (rtol is None) ^ (atol is None): + # We require both tolerance to be omitted or specified, because specifying only one might lead to surprising + # results. Imagine setting atol=0.0 and the tensors still match because rtol>0.0. + raise ErrorMeta( + ValueError, + f"Both 'rtol' and 'atol' must be either specified or omitted, " + f"but got no {'rtol' if rtol is None else 'atol'}.", + id=id, + ) + elif rtol is not None and atol is not None: + return rtol, atol + else: + return default_tolerances(*inputs) + + +def _make_bitwise_mismatch_msg( + *, + default_identifier: str, + identifier: str | Callable[[str], str] | None = None, + extra: str | None = None, + first_mismatch_idx: tuple[int, ...] | None = None, +): + """Makes a mismatch error message for bitwise values. + + Args: + default_identifier (str): Default description of the compared values, e.g. "Tensor-likes". + identifier (Optional[Union[str, Callable[[str], str]]]): Optional identifier that overrides + ``default_identifier``. Can be passed as callable in which case it will be called with + ``default_identifier`` to create the description at runtime. + extra (Optional[str]): Extra information to be placed after the message header and the mismatch statistics. + first_mismatch_idx (Optional[tuple[int, ...]]): the index of the first mismatch, for each dimension. + """ + if identifier is None: + identifier = default_identifier + elif callable(identifier): + identifier = identifier(default_identifier) + + msg = f"{identifier} are not 'equal'!\n\n" + + if extra: + msg += f"{extra.strip()}\n" + if first_mismatch_idx is not None: + msg += f"The first mismatched element is at index {first_mismatch_idx}.\n" + return msg.strip() + + +def _make_mismatch_msg( + *, + default_identifier: str, + identifier: str | Callable[[str], str] | None = None, + extra: str | None = None, + abs_diff: float, + abs_diff_idx: int | tuple[int, ...] | None = None, + atol: float, + rel_diff: float, + rel_diff_idx: int | tuple[int, ...] | None = None, + rtol: float, +) -> str: + """Makes a mismatch error message for numeric values. + + Args: + default_identifier (str): Default description of the compared values, e.g. "Tensor-likes". + identifier (Optional[Union[str, Callable[[str], str]]]): Optional identifier that overrides + ``default_identifier``. Can be passed as callable in which case it will be called with + ``default_identifier`` to create the description at runtime. + extra (Optional[str]): Extra information to be placed after the message header and the mismatch statistics. + abs_diff (float): Absolute difference. + abs_diff_idx (Optional[Union[int, Tuple[int, ...]]]): Optional index of the absolute difference. + atol (float): Allowed absolute tolerance. Will only be added to mismatch statistics if it or ``rtol`` are + ``> 0``. + rel_diff (float): Relative difference. + rel_diff_idx (Optional[Union[int, Tuple[int, ...]]]): Optional index of the relative difference. + rtol (float): Allowed relative tolerance. Will only be added to mismatch statistics if it or ``atol`` are + ``> 0``. + """ + equality = rtol == 0 and atol == 0 + + def make_diff_msg( + *, + type: str, + diff: float, + idx: int | tuple[int, ...] | None, + tol: float, + ) -> str: + if idx is None: + msg = f"{type.title()} difference: {diff}" + else: + msg = f"Greatest {type} difference: {diff} at index {idx}" + if not equality: + msg += f" (up to {tol} allowed)" + return msg + "\n" + + if identifier is None: + identifier = default_identifier + elif callable(identifier): + identifier = identifier(default_identifier) + + msg = f"{identifier} are not {'equal' if equality else 'close'}!\n\n" + + if extra: + msg += f"{extra.strip()}\n" + + msg += make_diff_msg(type="absolute", diff=abs_diff, idx=abs_diff_idx, tol=atol) + msg += make_diff_msg(type="relative", diff=rel_diff, idx=rel_diff_idx, tol=rtol) + + return msg.strip() + + +def make_scalar_mismatch_msg( + actual: bool | int | float | complex, + expected: bool | int | float | complex, + *, + rtol: float, + atol: float, + identifier: str | Callable[[str], str] | None = None, +) -> str: + """Makes a mismatch error message for scalars. + + Args: + actual (Union[bool, int, float, complex]): Actual scalar. + expected (Union[bool, int, float, complex]): Expected scalar. + rtol (float): Relative tolerance. + atol (float): Absolute tolerance. + identifier (Optional[Union[str, Callable[[str], str]]]): Optional description for the scalars. Can be passed + as callable in which case it will be called by the default value to create the description at runtime. + Defaults to "Scalars". + """ + abs_diff = abs(actual - expected) + # pyrefly: ignore [bad-argument-type] + rel_diff = float("inf") if expected == 0 else abs_diff / abs(expected) + return _make_mismatch_msg( + default_identifier="Scalars", + identifier=identifier, + extra=f"Expected {expected} but got {actual}.", + abs_diff=abs_diff, + atol=atol, + rel_diff=rel_diff, + rtol=rtol, + ) + + +def make_tensor_mismatch_msg( + actual: torch.Tensor, + expected: torch.Tensor, + matches: torch.Tensor, + *, + rtol: float, + atol: float, + identifier: str | Callable[[str], str] | None = None, +): + """Makes a mismatch error message for tensors. + + Args: + actual (torch.Tensor): Actual tensor. + expected (torch.Tensor): Expected tensor. + matches (torch.Tensor): Boolean mask of the same shape as ``actual`` and ``expected`` that indicates the + location of matches. + rtol (float): Relative tolerance. + atol (float): Absolute tolerance. + identifier (Optional[Union[str, Callable[[str], str]]]): Optional description for the tensors. Can be passed + as callable in which case it will be called by the default value to create the description at runtime. + Defaults to "Tensor-likes". + """ + + def unravel_flat_index(flat_index: int) -> tuple[int, ...]: + if not matches.shape: + return () + + inverse_index = [] + for size in matches.shape[::-1]: + div, mod = divmod(flat_index, size) + flat_index = div + inverse_index.append(mod) + + return tuple(inverse_index[::-1]) + + number_of_elements = matches.numel() + total_mismatches = number_of_elements - int(torch.sum(matches)) + extra = ( + f"Mismatched elements: {total_mismatches} / {number_of_elements} " + f"({total_mismatches / number_of_elements:.1%})" + ) + if actual.dtype.is_floating_point and actual.dtype.itemsize == 1: + # skip checking for max_abs_diff and max_rel_diff for float8-like values + first_mismatch_idx = tuple(torch.nonzero(~matches, as_tuple=False)[0].tolist()) + return _make_bitwise_mismatch_msg( + default_identifier="Tensor-likes", + identifier=identifier, + extra=extra, + first_mismatch_idx=first_mismatch_idx, + ) + + actual_flat = actual.flatten() + expected_flat = expected.flatten() + matches_flat = matches.flatten() + + if not actual.dtype.is_floating_point and not actual.dtype.is_complex: + # TODO: Instead of always upcasting to int64, it would be sufficient to cast to the next higher dtype to avoid + # overflow + actual_flat = actual_flat.to(torch.int64) + expected_flat = expected_flat.to(torch.int64) + + abs_diff = torch.abs(actual_flat - expected_flat) + # Ensure that only mismatches are used for the max_abs_diff computation + abs_diff[matches_flat] = 0 + max_abs_diff, max_abs_diff_flat_idx = torch.max(abs_diff, 0) + + rel_diff = abs_diff / torch.abs(expected_flat) + # Ensure that only mismatches are used for the max_rel_diff computation + rel_diff[matches_flat] = 0 + max_rel_diff, max_rel_diff_flat_idx = torch.max(rel_diff, 0) + return _make_mismatch_msg( + default_identifier="Tensor-likes", + identifier=identifier, + extra=extra, + abs_diff=max_abs_diff.item(), + abs_diff_idx=unravel_flat_index(int(max_abs_diff_flat_idx)), + atol=atol, + rel_diff=max_rel_diff.item(), + rel_diff_idx=unravel_flat_index(int(max_rel_diff_flat_idx)), + rtol=rtol, + ) + + +class UnsupportedInputs(Exception): # noqa: B903 + """Exception to be raised during the construction of a :class:`Pair` in case it doesn't support the inputs.""" + + +class Pair(abc.ABC): + """ABC for all comparison pairs to be used in conjunction with :func:`assert_equal`. + + Each subclass needs to overwrite :meth:`Pair.compare` that performs the actual comparison. + + Each pair receives **all** options, so select the ones applicable for the subclass and forward the rest to the + super class. Raising an :class:`UnsupportedInputs` during constructions indicates that the pair is not able to + handle the inputs and the next pair type will be tried. + + All other errors should be raised as :class:`ErrorMeta`. After the instantiation, :meth:`Pair._make_error_meta` can + be used to automatically handle overwriting the message with a user supplied one and id handling. + """ + + def __init__( + self, + actual: Any, + expected: Any, + *, + id: tuple[Any, ...] = (), + **unknown_parameters: Any, + ) -> None: + self.actual = actual + self.expected = expected + self.id = id + self._unknown_parameters = unknown_parameters + + @staticmethod + def _inputs_not_supported() -> NoReturn: + raise UnsupportedInputs + + @staticmethod + def _check_inputs_isinstance(*inputs: Any, cls: type | tuple[type, ...]): + """Checks if all inputs are instances of a given class and raise :class:`UnsupportedInputs` otherwise.""" + if not all(isinstance(input, cls) for input in inputs): + Pair._inputs_not_supported() + + def _fail( + self, type: type[Exception], msg: str, *, id: tuple[Any, ...] = () + ) -> NoReturn: + """Raises an :class:`ErrorMeta` from a given exception type and message and the stored id. + + .. warning:: + + If you use this before the ``super().__init__(...)`` call in the constructor, you have to pass the ``id`` + explicitly. + """ + raise ErrorMeta(type, msg, id=self.id if not id and hasattr(self, "id") else id) + + @abc.abstractmethod + def compare(self) -> None: + """Compares the inputs and raises an :class`ErrorMeta` in case they mismatch.""" + + def extra_repr(self) -> Sequence[str | tuple[str, Any]]: + """Returns extra information that will be included in the representation. + + Should be overwritten by all subclasses that use additional options. The representation of the object will only + be surfaced in case we encounter an unexpected error and thus should help debug the issue. Can be a sequence of + key-value-pairs or attribute names. + """ + return [] + + def __repr__(self) -> str: + head = f"{type(self).__name__}(" + tail = ")" + body = [ + f" {name}={value!s}," + for name, value in [ + ("id", self.id), + ("actual", self.actual), + ("expected", self.expected), + *[ + (extra, getattr(self, extra)) if isinstance(extra, str) else extra + for extra in self.extra_repr() + ], + ] + ] + return "\n".join((head, *body, *tail)) + + +class ObjectPair(Pair): + """Pair for any type of inputs that will be compared with the `==` operator. + + .. note:: + + Since this will instantiate for any kind of inputs, it should only be used as fallback after all other pairs + couldn't handle the inputs. + + """ + + def compare(self) -> None: + try: + equal = self.actual == self.expected + except Exception as error: + # We are not using `self._raise_error_meta` here since we need the exception chaining + raise ErrorMeta( + ValueError, + f"{self.actual} == {self.expected} failed with:\n{error}.", + id=self.id, + ) from error + + if not equal: + self._fail(AssertionError, f"{self.actual} != {self.expected}") + + +class NonePair(Pair): + """Pair for ``None`` inputs.""" + + def __init__(self, actual: Any, expected: Any, **other_parameters: Any) -> None: + if not (actual is None or expected is None): + self._inputs_not_supported() + + super().__init__(actual, expected, **other_parameters) + + def compare(self) -> None: + if not (self.actual is None and self.expected is None): + self._fail( + AssertionError, f"None mismatch: {self.actual} is not {self.expected}" + ) + + +class BooleanPair(Pair): + """Pair for :class:`bool` inputs. + + .. note:: + + If ``numpy`` is available, also handles :class:`numpy.bool_` inputs. + + """ + + def __init__( + self, + actual: Any, + expected: Any, + *, + id: tuple[Any, ...], + **other_parameters: Any, + ) -> None: + actual, expected = self._process_inputs(actual, expected, id=id) + super().__init__(actual, expected, **other_parameters) + + @property + def _supported_types(self) -> tuple[type, ...]: + cls: list[type] = [bool] + if HAS_NUMPY: + # pyrefly: ignore [missing-attribute] + cls.append(np.bool_) + return tuple(cls) + + def _process_inputs( + self, actual: Any, expected: Any, *, id: tuple[Any, ...] + ) -> tuple[bool, bool]: + self._check_inputs_isinstance(actual, expected, cls=self._supported_types) + actual, expected = ( + self._to_bool(bool_like, id=id) for bool_like in (actual, expected) + ) + return actual, expected + + def _to_bool(self, bool_like: Any, *, id: tuple[Any, ...]) -> bool: + if isinstance(bool_like, bool): + return bool_like + # pyrefly: ignore [missing-attribute] + elif isinstance(bool_like, np.bool_): + return bool_like.item() + else: + raise ErrorMeta( + TypeError, f"Unknown boolean type {type(bool_like)}.", id=id + ) + + def compare(self) -> None: + if self.actual is not self.expected: + self._fail( + AssertionError, + f"Booleans mismatch: {self.actual} is not {self.expected}", + ) + + +class NumberPair(Pair): + """Pair for Python number (:class:`int`, :class:`float`, and :class:`complex`) inputs. + + .. note:: + + If ``numpy`` is available, also handles :class:`numpy.number` inputs. + + Kwargs: + rtol (Optional[float]): Relative tolerance. If specified ``atol`` must also be specified. If omitted, default + values based on the type are selected with the below table. + atol (Optional[float]): Absolute tolerance. If specified ``rtol`` must also be specified. If omitted, default + values based on the type are selected with the below table. + equal_nan (bool): If ``True``, two ``NaN`` values are considered equal. Defaults to ``False``. + check_dtype (bool): If ``True``, the type of the inputs will be checked for equality. Defaults to ``False``. + + The following table displays correspondence between Python number type and the ``torch.dtype``'s. See + :func:`assert_close` for the corresponding tolerances. + + +------------------+-------------------------------+ + | ``type`` | corresponding ``torch.dtype`` | + +==================+===============================+ + | :class:`int` | :attr:`~torch.int64` | + +------------------+-------------------------------+ + | :class:`float` | :attr:`~torch.float64` | + +------------------+-------------------------------+ + | :class:`complex` | :attr:`~torch.complex64` | + +------------------+-------------------------------+ + """ + + _TYPE_TO_DTYPE = { + int: torch.int64, + float: torch.float64, + complex: torch.complex128, + } + _NUMBER_TYPES = tuple(_TYPE_TO_DTYPE.keys()) + + def __init__( + self, + actual: Any, + expected: Any, + *, + id: tuple[Any, ...] = (), + rtol: float | None = None, + atol: float | None = None, + equal_nan: bool = False, + check_dtype: bool = False, + **other_parameters: Any, + ) -> None: + actual, expected = self._process_inputs(actual, expected, id=id) + super().__init__(actual, expected, id=id, **other_parameters) + + self.rtol, self.atol = get_tolerances( + *[self._TYPE_TO_DTYPE[type(input)] for input in (actual, expected)], + rtol=rtol, + atol=atol, + id=id, + ) + self.equal_nan = equal_nan + self.check_dtype = check_dtype + + @property + def _supported_types(self) -> tuple[type, ...]: + cls = list(self._NUMBER_TYPES) + if HAS_NUMPY: + # pyrefly: ignore [missing-attribute] + cls.append(np.number) + return tuple(cls) + + def _process_inputs( + self, actual: Any, expected: Any, *, id: tuple[Any, ...] + ) -> tuple[int | float | complex, int | float | complex]: + self._check_inputs_isinstance(actual, expected, cls=self._supported_types) + actual, expected = ( + self._to_number(number_like, id=id) for number_like in (actual, expected) + ) + return actual, expected + + def _to_number( + self, number_like: Any, *, id: tuple[Any, ...] + ) -> int | float | complex: + # pyrefly: ignore [missing-attribute] + if HAS_NUMPY and isinstance(number_like, np.number): + return number_like.item() + elif isinstance(number_like, self._NUMBER_TYPES): + return number_like # type: ignore[return-value] + else: + raise ErrorMeta( + TypeError, f"Unknown number type {type(number_like)}.", id=id + ) + + def compare(self) -> None: + if self.check_dtype and type(self.actual) is not type(self.expected): + self._fail( + AssertionError, + f"The (d)types do not match: {type(self.actual)} != {type(self.expected)}.", + ) + + if self.actual == self.expected: + return + + if self.equal_nan and cmath.isnan(self.actual) and cmath.isnan(self.expected): + return + + abs_diff = abs(self.actual - self.expected) + tolerance = self.atol + self.rtol * abs(self.expected) + + if cmath.isfinite(abs_diff) and abs_diff <= tolerance: + return + + self._fail( + AssertionError, + make_scalar_mismatch_msg( + self.actual, self.expected, rtol=self.rtol, atol=self.atol + ), + ) + + def extra_repr(self) -> Sequence[str]: + return ( + "rtol", + "atol", + "equal_nan", + "check_dtype", + ) + + +class TensorLikePair(Pair): + """Pair for :class:`torch.Tensor`-like inputs. + + Kwargs: + allow_subclasses (bool): + rtol (Optional[float]): Relative tolerance. If specified ``atol`` must also be specified. If omitted, default + values based on the type are selected. See :func:assert_close: for details. + atol (Optional[float]): Absolute tolerance. If specified ``rtol`` must also be specified. If omitted, default + values based on the type are selected. See :func:assert_close: for details. + equal_nan (bool): If ``True``, two ``NaN`` values are considered equal. Defaults to ``False``. + check_device (bool): If ``True`` (default), asserts that corresponding tensors are on the same + :attr:`~torch.Tensor.device`. If this check is disabled, tensors on different + :attr:`~torch.Tensor.device`'s are moved to the CPU before being compared. + check_dtype (bool): If ``True`` (default), asserts that corresponding tensors have the same ``dtype``. If this + check is disabled, tensors with different ``dtype``'s are promoted to a common ``dtype`` (according to + :func:`torch.promote_types`) before being compared. + check_layout (bool): If ``True`` (default), asserts that corresponding tensors have the same ``layout``. If this + check is disabled, tensors with different ``layout``'s are converted to strided tensors before being + compared. + check_stride (bool): If ``True`` and corresponding tensors are strided, asserts that they have the same stride. + """ + + def __init__( + self, + actual: Any, + expected: Any, + *, + id: tuple[Any, ...] = (), + allow_subclasses: bool = True, + rtol: float | None = None, + atol: float | None = None, + equal_nan: bool = False, + check_device: bool = True, + check_dtype: bool = True, + check_layout: bool = True, + check_stride: bool = False, + **other_parameters: Any, + ): + actual, expected = self._process_inputs( + actual, expected, id=id, allow_subclasses=allow_subclasses + ) + super().__init__(actual, expected, id=id, **other_parameters) + + self.rtol, self.atol = get_tolerances( + actual, expected, rtol=rtol, atol=atol, id=self.id + ) + self.equal_nan = equal_nan + self.check_device = check_device + self.check_dtype = check_dtype + self.check_layout = check_layout + self.check_stride = check_stride + + def _process_inputs( + self, actual: Any, expected: Any, *, id: tuple[Any, ...], allow_subclasses: bool + ) -> tuple[torch.Tensor, torch.Tensor]: + directly_related = isinstance(actual, type(expected)) or isinstance( + expected, type(actual) + ) + if not directly_related: + self._inputs_not_supported() + + if not allow_subclasses and type(actual) is not type(expected): + self._inputs_not_supported() + + actual, expected = (self._to_tensor(input) for input in (actual, expected)) + for tensor in (actual, expected): + self._check_supported(tensor, id=id) + return actual, expected + + def _to_tensor(self, tensor_like: Any) -> torch.Tensor: + if isinstance(tensor_like, torch.Tensor): + return tensor_like + + try: + return torch.as_tensor(tensor_like) + except Exception: + self._inputs_not_supported() + + def _check_supported(self, tensor: torch.Tensor, *, id: tuple[Any, ...]) -> None: + if tensor.layout not in { + torch.strided, + torch.jagged, + torch.sparse_coo, + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }: + raise ErrorMeta( + ValueError, f"Unsupported tensor layout {tensor.layout}", id=id + ) + + def compare(self) -> None: + actual, expected = self.actual, self.expected + + self._compare_attributes(actual, expected) + if any(input.device.type == "meta" for input in (actual, expected)): + return + + actual, expected = self._equalize_attributes(actual, expected) + self._compare_values(actual, expected) + + def _compare_attributes( + self, + actual: torch.Tensor, + expected: torch.Tensor, + ) -> None: + """Checks if the attributes of two tensors match. + + Always checks + + - the :attr:`~torch.Tensor.shape`, + - whether both inputs are quantized or not, + - and if they use the same quantization scheme. + + Checks for + + - :attr:`~torch.Tensor.layout`, + - :meth:`~torch.Tensor.stride`, + - :attr:`~torch.Tensor.device`, and + - :attr:`~torch.Tensor.dtype` + + are optional and can be disabled through the corresponding ``check_*`` flag during construction of the pair. + """ + + def raise_mismatch_error( + attribute_name: str, actual_value: Any, expected_value: Any + ) -> NoReturn: + self._fail( + AssertionError, + f"The values for attribute '{attribute_name}' do not match: {actual_value} != {expected_value}.", + ) + + if actual.shape != expected.shape: + raise_mismatch_error("shape", actual.shape, expected.shape) + + if actual.is_quantized != expected.is_quantized: + raise_mismatch_error( + "is_quantized", actual.is_quantized, expected.is_quantized + ) + elif actual.is_quantized and actual.qscheme() != expected.qscheme(): + raise_mismatch_error("qscheme()", actual.qscheme(), expected.qscheme()) + + if actual.layout != expected.layout: + if self.check_layout: + raise_mismatch_error("layout", actual.layout, expected.layout) + elif ( + actual.layout == torch.strided + and self.check_stride + and actual.stride() != expected.stride() + ): + raise_mismatch_error("stride()", actual.stride(), expected.stride()) + + if self.check_device and actual.device != expected.device: + raise_mismatch_error("device", actual.device, expected.device) + + if self.check_dtype and actual.dtype != expected.dtype: + raise_mismatch_error("dtype", actual.dtype, expected.dtype) + + def _equalize_attributes( + self, actual: torch.Tensor, expected: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + """Equalizes some attributes of two tensors for value comparison. + + If ``actual`` and ``expected`` are ... + + - ... not on the same :attr:`~torch.Tensor.device`, they are moved CPU memory. + - ... not of the same ``dtype``, they are promoted to a common ``dtype`` (according to + :func:`torch.promote_types`). + - ... not of the same ``layout``, they are converted to strided tensors. + + Args: + actual (Tensor): Actual tensor. + expected (Tensor): Expected tensor. + + Returns: + (Tuple[Tensor, Tensor]): Equalized tensors. + """ + # The comparison logic uses operators currently not supported by the MPS backends. + # See https://github.com/pytorch/pytorch/issues/77144 for details. + # TODO: Remove this conversion as soon as all operations are supported natively by the MPS backend + if actual.is_mps or expected.is_mps: # type: ignore[attr-defined] + actual = actual.cpu() + expected = expected.cpu() + + if actual.device != expected.device: + actual = actual.cpu() + expected = expected.cpu() + + if actual.dtype != expected.dtype: + actual_dtype = actual.dtype + expected_dtype = expected.dtype + # For uint64, this is not sound in general, which is why promote_types doesn't + # allow it, but for easy testing, we're unlikely to get confused + # by large uint64 overflowing into negative int64 + if actual_dtype in [torch.uint64, torch.uint32, torch.uint16]: + actual_dtype = torch.int64 + if expected_dtype in [torch.uint64, torch.uint32, torch.uint16]: + expected_dtype = torch.int64 + dtype = torch.promote_types(actual_dtype, expected_dtype) + actual = actual.to(dtype) + expected = expected.to(dtype) + + if actual.layout != expected.layout: + # These checks are needed, since Tensor.to_dense() fails on tensors that are already strided + actual = actual.to_dense() if actual.layout != torch.strided else actual + expected = ( + expected.to_dense() if expected.layout != torch.strided else expected + ) + + return actual, expected + + def _compare_values(self, actual: torch.Tensor, expected: torch.Tensor) -> None: + if actual.is_quantized: + compare_fn = self._compare_quantized_values + elif actual.is_sparse: + compare_fn = self._compare_sparse_coo_values + elif actual.layout in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }: + compare_fn = self._compare_sparse_compressed_values + elif actual.layout == torch.jagged: + actual, expected = actual.values(), expected.values() + compare_fn = self._compare_regular_values_close + elif actual.dtype.is_floating_point and actual.dtype.itemsize == 1: + + def bitwise_comp( + actual: torch.Tensor, + expected: torch.Tensor, + *, + rtol: float, + atol: float, + equal_nan: bool, + identifier: str | Callable[[str], str] | None = None, + ) -> None: + if rtol != 0.0 or atol != 0.0: + raise ErrorMeta( + AssertionError, + f"Rtol={rtol} and atol={atol} are not supported for bitwise comparison of low" + " dimensional floats. Please use rtol=0.0 and atol=0.0.", + ) + + return self._compare_regular_values_close( + actual, + expected, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + identifier=identifier, + ) + + compare_fn = bitwise_comp + else: + compare_fn = self._compare_regular_values_close + + compare_fn( + actual, expected, rtol=self.rtol, atol=self.atol, equal_nan=self.equal_nan + ) + + def _compare_quantized_values( + self, + actual: torch.Tensor, + expected: torch.Tensor, + *, + rtol: float, + atol: float, + equal_nan: bool, + ) -> None: + """Compares quantized tensors by comparing the :meth:`~torch.Tensor.dequantize`'d variants for closeness. + + .. note:: + + A detailed discussion about why only the dequantized variant is checked for closeness rather than checking + the individual quantization parameters for closeness and the integer representation for equality can be + found in https://github.com/pytorch/pytorch/issues/68548. + """ + return self._compare_regular_values_close( + actual.dequantize(), + expected.dequantize(), + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + identifier=lambda default_identifier: f"Quantized {default_identifier.lower()}", + ) + + def _compare_sparse_coo_values( + self, + actual: torch.Tensor, + expected: torch.Tensor, + *, + rtol: float, + atol: float, + equal_nan: bool, + ) -> None: + """Compares sparse COO tensors by comparing + + - the number of sparse dimensions, + - the number of non-zero elements (nnz) for equality, + - the indices for equality, and + - the values for closeness. + """ + if actual.sparse_dim() != expected.sparse_dim(): + self._fail( + AssertionError, + ( + f"The number of sparse dimensions in sparse COO tensors does not match: " + f"{actual.sparse_dim()} != {expected.sparse_dim()}" + ), + ) + + if actual._nnz() != expected._nnz(): + self._fail( + AssertionError, + ( + f"The number of specified values in sparse COO tensors does not match: " + f"{actual._nnz()} != {expected._nnz()}" + ), + ) + + self._compare_regular_values_equal( + actual._indices(), + expected._indices(), + identifier="Sparse COO indices", + ) + self._compare_regular_values_close( + actual._values(), + expected._values(), + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + identifier="Sparse COO values", + ) + + def _compare_sparse_compressed_values( + self, + actual: torch.Tensor, + expected: torch.Tensor, + *, + rtol: float, + atol: float, + equal_nan: bool, + ) -> None: + """Compares sparse compressed tensors by comparing + + - the number of non-zero elements (nnz) for equality, + - the plain indices for equality, + - the compressed indices for equality, and + - the values for closeness. + """ + format_name, compressed_indices_method, plain_indices_method = { + torch.sparse_csr: ( + "CSR", + torch.Tensor.crow_indices, + torch.Tensor.col_indices, + ), + torch.sparse_csc: ( + "CSC", + torch.Tensor.ccol_indices, + torch.Tensor.row_indices, + ), + torch.sparse_bsr: ( + "BSR", + torch.Tensor.crow_indices, + torch.Tensor.col_indices, + ), + torch.sparse_bsc: ( + "BSC", + torch.Tensor.ccol_indices, + torch.Tensor.row_indices, + ), + }[actual.layout] + + if actual._nnz() != expected._nnz(): + self._fail( + AssertionError, + ( + f"The number of specified values in sparse {format_name} tensors does not match: " + f"{actual._nnz()} != {expected._nnz()}" + ), + ) + + # Compressed and plain indices in the CSR / CSC / BSR / BSC sparse formats can be `torch.int32` _or_ + # `torch.int64`. While the same dtype is enforced for the compressed and plain indices of a single tensor, it + # can be different between two tensors. Thus, we need to convert them to the same dtype, or the comparison will + # fail. + actual_compressed_indices = compressed_indices_method(actual) + expected_compressed_indices = compressed_indices_method(expected) + indices_dtype = torch.promote_types( + actual_compressed_indices.dtype, expected_compressed_indices.dtype + ) + + self._compare_regular_values_equal( + actual_compressed_indices.to(indices_dtype), + expected_compressed_indices.to(indices_dtype), + identifier=f"Sparse {format_name} {compressed_indices_method.__name__}", + ) + self._compare_regular_values_equal( + plain_indices_method(actual).to(indices_dtype), + plain_indices_method(expected).to(indices_dtype), + identifier=f"Sparse {format_name} {plain_indices_method.__name__}", + ) + self._compare_regular_values_close( + actual.values(), + expected.values(), + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + identifier=f"Sparse {format_name} values", + ) + + def _compare_regular_values_equal( + self, + actual: torch.Tensor, + expected: torch.Tensor, + *, + equal_nan: bool = False, + identifier: str | Callable[[str], str] | None = None, + ) -> None: + """Checks if the values of two tensors are equal.""" + self._compare_regular_values_close( + actual, expected, rtol=0, atol=0, equal_nan=equal_nan, identifier=identifier + ) + + def _compare_regular_values_close( + self, + actual: torch.Tensor, + expected: torch.Tensor, + *, + rtol: float, + atol: float, + equal_nan: bool, + identifier: str | Callable[[str], str] | None = None, + ) -> None: + """Checks if the values of two tensors are close up to a desired tolerance.""" + matches = torch.isclose( + actual, expected, rtol=rtol, atol=atol, equal_nan=equal_nan + ) + if torch.all(matches): + return + + if actual.shape == torch.Size([]): + msg = make_scalar_mismatch_msg( + actual.item(), + expected.item(), + rtol=rtol, + atol=atol, + identifier=identifier, + ) + else: + msg = make_tensor_mismatch_msg( + actual, expected, matches, rtol=rtol, atol=atol, identifier=identifier + ) + self._fail(AssertionError, msg) + + def extra_repr(self) -> Sequence[str]: + return ( + "rtol", + "atol", + "equal_nan", + "check_device", + "check_dtype", + "check_layout", + "check_stride", + ) + + +def originate_pairs( + actual: Any, + expected: Any, + *, + pair_types: Sequence[type[Pair]], + sequence_types: tuple[type, ...] = (collections.abc.Sequence,), + mapping_types: tuple[type, ...] = (collections.abc.Mapping,), + id: tuple[Any, ...] = (), + **options: Any, + # pyrefly: ignore [bad-return] +) -> list[Pair]: + """Originates pairs from the individual inputs. + + ``actual`` and ``expected`` can be possibly nested :class:`~collections.abc.Sequence`'s or + :class:`~collections.abc.Mapping`'s. In this case the pairs are originated by recursing through them. + + Args: + actual (Any): Actual input. + expected (Any): Expected input. + pair_types (Sequence[Type[Pair]]): Sequence of pair types that will be tried to construct with the inputs. + First successful pair will be used. + sequence_types (Tuple[Type, ...]): Optional types treated as sequences that will be checked elementwise. + mapping_types (Tuple[Type, ...]): Optional types treated as mappings that will be checked elementwise. + id (Tuple[Any, ...]): Optional id of a pair that will be included in an error message. + **options (Any): Options passed to each pair during construction. + + Raises: + ErrorMeta: With :class`AssertionError`, if the inputs are :class:`~collections.abc.Sequence`'s, but their + length does not match. + ErrorMeta: With :class`AssertionError`, if the inputs are :class:`~collections.abc.Mapping`'s, but their set of + keys do not match. + ErrorMeta: With :class`TypeError`, if no pair is able to handle the inputs. + ErrorMeta: With any expected exception that happens during the construction of a pair. + + Returns: + (List[Pair]): Originated pairs. + """ + # We explicitly exclude str's here since they are self-referential and would cause an infinite recursion loop: + # "a" == "a"[0][0]... + if ( + isinstance(actual, sequence_types) + and not isinstance(actual, str) + and isinstance(expected, sequence_types) + and not isinstance(expected, str) + ): + actual_len = len(actual) # type: ignore[arg-type] + expected_len = len(expected) # type: ignore[arg-type] + if actual_len != expected_len: + raise ErrorMeta( + AssertionError, + f"The length of the sequences mismatch: {actual_len} != {expected_len}", + id=id, + ) + + pairs = [] + for idx in range(actual_len): + pairs.extend( + originate_pairs( + actual[idx], # type: ignore[index] + expected[idx], # type: ignore[index] + pair_types=pair_types, + sequence_types=sequence_types, + mapping_types=mapping_types, + id=(*id, idx), + **options, + ) + ) + return pairs + + elif isinstance(actual, mapping_types) and isinstance(expected, mapping_types): + actual_keys = set(actual.keys()) # type: ignore[attr-defined] + expected_keys = set(expected.keys()) # type: ignore[attr-defined] + if actual_keys != expected_keys: + missing_keys = expected_keys - actual_keys + additional_keys = actual_keys - expected_keys + raise ErrorMeta( + AssertionError, + ( + f"The keys of the mappings do not match:\n" + f"Missing keys in the actual mapping: {sorted(missing_keys)}\n" + f"Additional keys in the actual mapping: {sorted(additional_keys)}" + ), + id=id, + ) + + keys: Collection = actual_keys + # Since the origination aborts after the first failure, we try to be deterministic + with contextlib.suppress(Exception): + keys = sorted(keys) + + pairs = [] + for key in keys: + pairs.extend( + originate_pairs( + actual[key], # type: ignore[index] + expected[key], # type: ignore[index] + pair_types=pair_types, + sequence_types=sequence_types, + mapping_types=mapping_types, + id=(*id, key), + **options, + ) + ) + return pairs + + else: + for pair_type in pair_types: + try: + # pyrefly: ignore [bad-instantiation] + return [pair_type(actual, expected, id=id, **options)] + # Raising an `UnsupportedInputs` during origination indicates that the pair type is not able to handle the + # inputs. Thus, we try the next pair type. + except UnsupportedInputs: + continue + # Raising an `ErrorMeta` during origination is the orderly way to abort and so we simply re-raise it. This + # is only in a separate branch, because the one below would also except it. + except ErrorMeta: + raise + # Raising any other exception during origination is unexpected and will give some extra information about + # what happened. If applicable, the exception should be expected in the future. + except Exception as error: + raise RuntimeError( + f"Originating a {pair_type.__name__}() at item {''.join(str([item]) for item in id)} with\n\n" + f"{type(actual).__name__}(): {actual}\n\n" + f"and\n\n" + f"{type(expected).__name__}(): {expected}\n\n" + f"resulted in the unexpected exception above. " + f"If you are a user and see this message during normal operation " + "please file an issue at https://github.com/pytorch/pytorch/issues. " + "If you are a developer and working on the comparison functions, " + "please except the previous error and raise an expressive `ErrorMeta` instead." + ) from error + else: + raise ErrorMeta( + TypeError, + f"No comparison pair was able to handle inputs of type {type(actual)} and {type(expected)}.", + id=id, + ) + + +def not_close_error_metas( + actual: Any, + expected: Any, + *, + pair_types: Sequence[type[Pair]] = (ObjectPair,), + sequence_types: tuple[type, ...] = (collections.abc.Sequence,), + mapping_types: tuple[type, ...] = (collections.abc.Mapping,), + **options: Any, +) -> list[ErrorMeta]: + """Asserts that inputs are equal. + + ``actual`` and ``expected`` can be possibly nested :class:`~collections.abc.Sequence`'s or + :class:`~collections.abc.Mapping`'s. In this case the comparison happens elementwise by recursing through them. + + Args: + actual (Any): Actual input. + expected (Any): Expected input. + pair_types (Sequence[Type[Pair]]): Sequence of :class:`Pair` types that will be tried to construct with the + inputs. First successful pair will be used. Defaults to only using :class:`ObjectPair`. + sequence_types (Tuple[Type, ...]): Optional types treated as sequences that will be checked elementwise. + mapping_types (Tuple[Type, ...]): Optional types treated as mappings that will be checked elementwise. + **options (Any): Options passed to each pair during construction. + """ + # Hide this function from `pytest`'s traceback + __tracebackhide__ = True + + try: + pairs = originate_pairs( + actual, + expected, + pair_types=pair_types, + sequence_types=sequence_types, + mapping_types=mapping_types, + **options, + ) + except ErrorMeta as error_meta: + # Explicitly raising from None to hide the internal traceback + raise error_meta.to_error() from None # noqa: RSE102 + + error_metas: list[ErrorMeta] = [] + for pair in pairs: + try: + pair.compare() + except ErrorMeta as error_meta: + error_metas.append(error_meta) + # Raising any exception besides `ErrorMeta` while comparing is unexpected and will give some extra information + # about what happened. If applicable, the exception should be expected in the future. + except Exception as error: + raise RuntimeError( + f"Comparing\n\n" + f"{pair}\n\n" + f"resulted in the unexpected exception above. " + f"If you are a user and see this message during normal operation " + "please file an issue at https://github.com/pytorch/pytorch/issues. " + "If you are a developer and working on the comparison functions, " + "please except the previous error and raise an expressive `ErrorMeta` instead." + ) from error + + # [ErrorMeta Cycles] + # ErrorMeta objects in this list capture + # tracebacks that refer to the frame of this function. + # The local variable `error_metas` refers to the error meta + # objects, creating a reference cycle. Frames in the traceback + # would not get freed until cycle collection, leaking cuda memory in tests. + # We break the cycle by removing the reference to the error_meta objects + # from this frame as it returns. + # pyrefly: ignore [bad-assignment] + error_metas = [error_metas] + # pyrefly: ignore [bad-return] + return error_metas.pop() + + +def assert_close( + actual: Any, + expected: Any, + *, + allow_subclasses: bool = True, + rtol: float | None = None, + atol: float | None = None, + equal_nan: bool = False, + check_device: bool = True, + check_dtype: bool = True, + check_layout: bool = True, + check_stride: bool = False, + msg: str | Callable[[str], str] | None = None, +): + r"""Asserts that ``actual`` and ``expected`` are close. + + If ``actual`` and ``expected`` are strided, non-quantized, real-valued, and finite, they are considered close if + + .. math:: + + \lvert \text{actual} - \text{expected} \rvert \le \texttt{atol} + \texttt{rtol} \cdot \lvert \text{expected} \rvert + + Non-finite values (``-inf`` and ``inf``) are only considered close if and only if they are equal. ``NaN``'s are + only considered equal to each other if ``equal_nan`` is ``True``. + + In addition, they are only considered close if they have the same + + - :attr:`~torch.Tensor.device` (if ``check_device`` is ``True``), + - ``dtype`` (if ``check_dtype`` is ``True``), + - ``layout`` (if ``check_layout`` is ``True``), and + - stride (if ``check_stride`` is ``True``). + + If either ``actual`` or ``expected`` is a meta tensor, only the attribute checks will be performed. + + If ``actual`` and ``expected`` are sparse (either having COO, CSR, CSC, BSR, or BSC layout), their strided members are + checked individually. Indices, namely ``indices`` for COO, ``crow_indices`` and ``col_indices`` for CSR and BSR, + or ``ccol_indices`` and ``row_indices`` for CSC and BSC layouts, respectively, + are always checked for equality whereas the values are checked for closeness according to the definition above. + + If ``actual`` and ``expected`` are quantized, they are considered close if they have the same + :meth:`~torch.Tensor.qscheme` and the result of :meth:`~torch.Tensor.dequantize` is close according to the + definition above. + + ``actual`` and ``expected`` can be :class:`~torch.Tensor`'s or any tensor-or-scalar-likes from which + :class:`torch.Tensor`'s can be constructed with :func:`torch.as_tensor`. Except for Python scalars the input types + have to be directly related. In addition, ``actual`` and ``expected`` can be :class:`~collections.abc.Sequence`'s + or :class:`~collections.abc.Mapping`'s in which case they are considered close if their structure matches and all + their elements are considered close according to the above definition. + + .. note:: + + Python scalars are an exception to the type relation requirement, because their :func:`type`, i.e. + :class:`int`, :class:`float`, and :class:`complex`, is equivalent to the ``dtype`` of a tensor-like. Thus, + Python scalars of different types can be checked, but require ``check_dtype=False``. + + Args: + actual (Any): Actual input. + expected (Any): Expected input. + allow_subclasses (bool): If ``True`` (default) and except for Python scalars, inputs of directly related types + are allowed. Otherwise type equality is required. + rtol (Optional[float]): Relative tolerance. If specified ``atol`` must also be specified. If omitted, default + values based on the :attr:`~torch.Tensor.dtype` are selected with the below table. + atol (Optional[float]): Absolute tolerance. If specified ``rtol`` must also be specified. If omitted, default + values based on the :attr:`~torch.Tensor.dtype` are selected with the below table. + equal_nan (Union[bool, str]): If ``True``, two ``NaN`` values will be considered equal. + check_device (bool): If ``True`` (default), asserts that corresponding tensors are on the same + :attr:`~torch.Tensor.device`. If this check is disabled, tensors on different + :attr:`~torch.Tensor.device`'s are moved to the CPU before being compared. + check_dtype (bool): If ``True`` (default), asserts that corresponding tensors have the same ``dtype``. If this + check is disabled, tensors with different ``dtype``'s are promoted to a common ``dtype`` (according to + :func:`torch.promote_types`) before being compared. + check_layout (bool): If ``True`` (default), asserts that corresponding tensors have the same ``layout``. If this + check is disabled, tensors with different ``layout``'s are converted to strided tensors before being + compared. + check_stride (bool): If ``True`` and corresponding tensors are strided, asserts that they have the same stride. + msg (Optional[Union[str, Callable[[str], str]]]): Optional error message to use in case a failure occurs during + the comparison. Can also passed as callable in which case it will be called with the generated message and + should return the new message. + + Raises: + ValueError: If no :class:`torch.Tensor` can be constructed from an input. + ValueError: If only ``rtol`` or ``atol`` is specified. + AssertionError: If corresponding inputs are not Python scalars and are not directly related. + AssertionError: If ``allow_subclasses`` is ``False``, but corresponding inputs are not Python scalars and have + different types. + AssertionError: If the inputs are :class:`~collections.abc.Sequence`'s, but their length does not match. + AssertionError: If the inputs are :class:`~collections.abc.Mapping`'s, but their set of keys do not match. + AssertionError: If corresponding tensors do not have the same :attr:`~torch.Tensor.shape`. + AssertionError: If ``check_layout`` is ``True``, but corresponding tensors do not have the same + :attr:`~torch.Tensor.layout`. + AssertionError: If only one of corresponding tensors is quantized. + AssertionError: If corresponding tensors are quantized, but have different :meth:`~torch.Tensor.qscheme`'s. + AssertionError: If ``check_device`` is ``True``, but corresponding tensors are not on the same + :attr:`~torch.Tensor.device`. + AssertionError: If ``check_dtype`` is ``True``, but corresponding tensors do not have the same ``dtype``. + AssertionError: If ``check_stride`` is ``True``, but corresponding strided tensors do not have the same stride. + AssertionError: If the values of corresponding tensors are not close according to the definition above. + + The following table displays the default ``rtol`` and ``atol`` for different ``dtype``'s. In case of mismatching + ``dtype``'s, the maximum of both tolerances is used. + + +---------------------------+------------+----------+ + | ``dtype`` | ``rtol`` | ``atol`` | + +===========================+============+==========+ + | :attr:`~torch.float16` | ``1e-3`` | ``1e-5`` | + +---------------------------+------------+----------+ + | :attr:`~torch.bfloat16` | ``1.6e-2`` | ``1e-5`` | + +---------------------------+------------+----------+ + | :attr:`~torch.float32` | ``1.3e-6`` | ``1e-5`` | + +---------------------------+------------+----------+ + | :attr:`~torch.float64` | ``1e-7`` | ``1e-7`` | + +---------------------------+------------+----------+ + | :attr:`~torch.complex32` | ``1e-3`` | ``1e-5`` | + +---------------------------+------------+----------+ + | :attr:`~torch.complex64` | ``1.3e-6`` | ``1e-5`` | + +---------------------------+------------+----------+ + | :attr:`~torch.complex128` | ``1e-7`` | ``1e-7`` | + +---------------------------+------------+----------+ + | :attr:`~torch.quint8` | ``1.3e-6`` | ``1e-5`` | + +---------------------------+------------+----------+ + | :attr:`~torch.quint2x4` | ``1.3e-6`` | ``1e-5`` | + +---------------------------+------------+----------+ + | :attr:`~torch.quint4x2` | ``1.3e-6`` | ``1e-5`` | + +---------------------------+------------+----------+ + | :attr:`~torch.qint8` | ``1.3e-6`` | ``1e-5`` | + +---------------------------+------------+----------+ + | :attr:`~torch.qint32` | ``1.3e-6`` | ``1e-5`` | + +---------------------------+------------+----------+ + | other | ``0.0`` | ``0.0`` | + +---------------------------+------------+----------+ + + .. note:: + + :func:`~torch.testing.assert_close` is highly configurable with strict default settings. Users are encouraged + to :func:`~functools.partial` it to fit their use case. For example, if an equality check is needed, one might + define an ``assert_equal`` that uses zero tolerances for every ``dtype`` by default: + + >>> import functools + >>> assert_equal = functools.partial(torch.testing.assert_close, rtol=0, atol=0) + >>> assert_equal(1e-9, 1e-10) + Traceback (most recent call last): + ... + AssertionError: Scalars are not equal! + + Expected 1e-10 but got 1e-09. + Absolute difference: 9.000000000000001e-10 + Relative difference: 9.0 + + Examples: + >>> # tensor to tensor comparison + >>> expected = torch.tensor([1e0, 1e-1, 1e-2]) + >>> actual = torch.acos(torch.cos(expected)) + >>> torch.testing.assert_close(actual, expected) + + >>> # scalar to scalar comparison + >>> import math + >>> expected = math.sqrt(2.0) + >>> actual = 2.0 / math.sqrt(2.0) + >>> torch.testing.assert_close(actual, expected) + + >>> # numpy array to numpy array comparison + >>> import numpy as np + >>> expected = np.array([1e0, 1e-1, 1e-2]) + >>> actual = np.arccos(np.cos(expected)) + >>> torch.testing.assert_close(actual, expected) + + >>> # sequence to sequence comparison + >>> import numpy as np + >>> # The types of the sequences do not have to match. They only have to have the same + >>> # length and their elements have to match. + >>> expected = [torch.tensor([1.0]), 2.0, np.array(3.0)] + >>> actual = tuple(expected) + >>> torch.testing.assert_close(actual, expected) + + >>> # mapping to mapping comparison + >>> from collections import OrderedDict + >>> import numpy as np + >>> foo = torch.tensor(1.0) + >>> bar = 2.0 + >>> baz = np.array(3.0) + >>> # The types and a possible ordering of mappings do not have to match. They only + >>> # have to have the same set of keys and their elements have to match. + >>> expected = OrderedDict([("foo", foo), ("bar", bar), ("baz", baz)]) + >>> actual = {"baz": baz, "bar": bar, "foo": foo} + >>> torch.testing.assert_close(actual, expected) + + >>> expected = torch.tensor([1.0, 2.0, 3.0]) + >>> actual = expected.clone() + >>> # By default, directly related instances can be compared + >>> torch.testing.assert_close(torch.nn.Parameter(actual), expected) + >>> # This check can be made more strict with allow_subclasses=False + >>> torch.testing.assert_close( + ... torch.nn.Parameter(actual), expected, allow_subclasses=False + ... ) + Traceback (most recent call last): + ... + TypeError: No comparison pair was able to handle inputs of type + and . + >>> # If the inputs are not directly related, they are never considered close + >>> torch.testing.assert_close(actual.numpy(), expected) + Traceback (most recent call last): + ... + TypeError: No comparison pair was able to handle inputs of type + and . + >>> # Exceptions to these rules are Python scalars. They can be checked regardless of + >>> # their type if check_dtype=False. + >>> torch.testing.assert_close(1.0, 1, check_dtype=False) + + >>> # NaN != NaN by default. + >>> expected = torch.tensor(float("Nan")) + >>> actual = expected.clone() + >>> torch.testing.assert_close(actual, expected) + Traceback (most recent call last): + ... + AssertionError: Scalars are not close! + + Expected nan but got nan. + Absolute difference: nan (up to 1e-05 allowed) + Relative difference: nan (up to 1.3e-06 allowed) + >>> torch.testing.assert_close(actual, expected, equal_nan=True) + + >>> expected = torch.tensor([1.0, 2.0, 3.0]) + >>> actual = torch.tensor([1.0, 4.0, 5.0]) + >>> # The default error message can be overwritten. + >>> torch.testing.assert_close( + ... actual, expected, msg="Argh, the tensors are not close!" + ... ) + Traceback (most recent call last): + ... + AssertionError: Argh, the tensors are not close! + >>> # If msg is a callable, it can be used to augment the generated message with + >>> # extra information + >>> torch.testing.assert_close( + ... actual, expected, msg=lambda msg: f"Header\n\n{msg}\n\nFooter" + ... ) + Traceback (most recent call last): + ... + AssertionError: Header + + Tensor-likes are not close! + + Mismatched elements: 2 / 3 (66.7%) + Greatest absolute difference: 2.0 at index (1,) (up to 1e-05 allowed) + Greatest relative difference: 1.0 at index (1,) (up to 1.3e-06 allowed) + + Footer + """ + # Hide this function from `pytest`'s traceback + __tracebackhide__ = True + + actual, expected = _unwrap_dtensor_for_comparison(actual, expected) + + error_metas = not_close_error_metas( + actual, + expected, + pair_types=( + NonePair, + BooleanPair, + NumberPair, + TensorLikePair, + ), + allow_subclasses=allow_subclasses, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + check_device=check_device, + check_dtype=check_dtype, + check_layout=check_layout, + check_stride=check_stride, + msg=msg, + ) + + if error_metas: + # TODO: compose all metas into one AssertionError + raise error_metas[0].to_error(msg) + + +@deprecated( + "`torch.testing.assert_allclose()` is deprecated since 1.12 and will be removed in a future release. " + "Please use `torch.testing.assert_close()` instead. " + "You can find detailed upgrade instructions in https://github.com/pytorch/pytorch/issues/61844.", + category=FutureWarning, +) +def assert_allclose( + actual: Any, + expected: Any, + rtol: float | None = None, + atol: float | None = None, + equal_nan: bool = True, + msg: str = "", +) -> None: + """ + .. warning:: + + :func:`torch.testing.assert_allclose` is deprecated since ``1.12`` and will be removed in a future release. + Please use :func:`torch.testing.assert_close` instead. You can find detailed upgrade instructions + `here `_. + """ + if not isinstance(actual, torch.Tensor): + actual = torch.tensor(actual) + if not isinstance(expected, torch.Tensor): + expected = torch.tensor(expected, dtype=actual.dtype) + + if rtol is None and atol is None: + rtol, atol = default_tolerances( + actual, + expected, + dtype_precisions={ + torch.float16: (1e-3, 1e-3), + torch.float32: (1e-4, 1e-5), + torch.float64: (1e-5, 1e-8), + }, + ) + + torch.testing.assert_close( + actual, + expected, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + check_device=True, + check_dtype=False, + check_stride=False, + msg=msg or None, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_creation.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_creation.py new file mode 100644 index 0000000000000000000000000000000000000000..6b212d4a84b1d49811d7def6c1f0d0c1345dca55 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_creation.py @@ -0,0 +1,276 @@ +""" +This module contains tensor creation utilities. +""" + +import collections.abc +import functools +import math +import warnings +from typing import cast + +import torch + + +_INTEGRAL_TYPES = [ + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint16, + torch.uint32, + torch.uint64, +] +_FLOATING_TYPES = [torch.float16, torch.bfloat16, torch.float32, torch.float64] +_FLOATING_8BIT_TYPES = [ + torch.float8_e4m3fn, + torch.float8_e5m2, + torch.float8_e4m3fnuz, + torch.float8_e5m2fnuz, +] +_COMPLEX_TYPES = [torch.complex32, torch.complex64, torch.complex128] +_BOOLEAN_OR_INTEGRAL_TYPES = [torch.bool, *_INTEGRAL_TYPES] +_FLOATING_OR_COMPLEX_TYPES = [*_FLOATING_TYPES, *_COMPLEX_TYPES] + + +def _uniform_random_(t: torch.Tensor, low: float, high: float) -> torch.Tensor: + # uniform_ requires to-from <= std::numeric_limits::max() + # Work around this by scaling the range before and after the PRNG + if high - low >= torch.finfo(t.dtype).max: + return t.uniform_(low / 2, high / 2).mul_(2) + else: + return t.uniform_(low, high) + + +def make_tensor( + *shape: int | torch.Size | list[int] | tuple[int, ...], + dtype: torch.dtype, + device: str | torch.device, + low: float | None = None, + high: float | None = None, + requires_grad: bool = False, + noncontiguous: bool = False, + exclude_zero: bool = False, + memory_format: torch.memory_format | None = None, +) -> torch.Tensor: + r"""Creates a tensor with the given :attr:`shape`, :attr:`device`, and :attr:`dtype`, and filled with + values uniformly drawn from ``[low, high)``. + + If :attr:`low` or :attr:`high` are specified and are outside the range of the :attr:`dtype`'s representable + finite values then they are clamped to the lowest or highest representable finite value, respectively. + If ``None``, then the following table describes the default values for :attr:`low` and :attr:`high`, + which depend on :attr:`dtype`. + + +---------------------------+------------+----------+ + | ``dtype`` | ``low`` | ``high`` | + +===========================+============+==========+ + | boolean type | ``0`` | ``2`` | + +---------------------------+------------+----------+ + | unsigned integral type | ``0`` | ``10`` | + +---------------------------+------------+----------+ + | signed integral types | ``-9`` | ``10`` | + +---------------------------+------------+----------+ + | floating types | ``-9`` | ``9`` | + +---------------------------+------------+----------+ + | complex types | ``-9`` | ``9`` | + +---------------------------+------------+----------+ + + Args: + shape (Tuple[int, ...]): Single integer or a sequence of integers defining the shape of the output tensor. + dtype (:class:`torch.dtype`): The data type of the returned tensor. + device (Union[str, torch.device]): The device of the returned tensor. + low (Optional[Number]): Sets the lower limit (inclusive) of the given range. If a number is provided it is + clamped to the least representable finite value of the given dtype. When ``None`` (default), + this value is determined based on the :attr:`dtype` (see the table above). Default: ``None``. + high (Optional[Number]): Sets the upper limit (exclusive) of the given range. If a number is provided it is + clamped to the greatest representable finite value of the given dtype. When ``None`` (default) this value + is determined based on the :attr:`dtype` (see the table above). Default: ``None``. + + .. deprecated:: 2.1 + + Passing ``low==high`` to :func:`~torch.testing.make_tensor` for floating or complex types is deprecated + since 2.1 and will be removed in 2.3. Use :func:`torch.full` instead. + + requires_grad (Optional[bool]): If autograd should record operations on the returned tensor. Default: ``False``. + noncontiguous (Optional[bool]): If `True`, the returned tensor will be noncontiguous. This argument is + ignored if the constructed tensor has fewer than two elements. Mutually exclusive with ``memory_format``. + exclude_zero (Optional[bool]): If ``True`` then zeros are replaced with the dtype's small positive value + depending on the :attr:`dtype`. For bool and integer types zero is replaced with one. For floating + point types it is replaced with the dtype's smallest positive normal number (the "tiny" value of the + :attr:`dtype`'s :func:`~torch.finfo` object), and for complex types it is replaced with a complex number + whose real and imaginary parts are both the smallest positive normal number representable by the complex + type. Default ``False``. + memory_format (Optional[torch.memory_format]): The memory format of the returned tensor. Mutually exclusive + with ``noncontiguous``. + + Raises: + ValueError: If ``requires_grad=True`` is passed for integral `dtype` + ValueError: If ``low >= high``. + ValueError: If either :attr:`low` or :attr:`high` is ``nan``. + ValueError: If both :attr:`noncontiguous` and :attr:`memory_format` are passed. + TypeError: If :attr:`dtype` isn't supported by this function. + + Examples: + >>> # xdoctest: +SKIP + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA) + >>> from torch.testing import make_tensor + >>> # Creates a float tensor with values in [-1, 1) + >>> make_tensor((3,), device="cpu", dtype=torch.float32, low=-1, high=1) + >>> # xdoctest: +SKIP + tensor([ 0.1205, 0.2282, -0.6380]) + >>> # Creates a bool tensor on CUDA + >>> make_tensor((2, 2), device="cuda", dtype=torch.bool) + tensor([[False, False], + [False, True]], device='cuda:0') + """ + + def modify_low_high( + low: float | None, + high: float | None, + *, + lowest_inclusive: float, + highest_exclusive: float, + default_low: float, + default_high: float, + ) -> tuple[float, float]: + """ + Modifies (and raises ValueError when appropriate) low and high values given by the user (input_low, input_high) + if required. + """ + + def clamp(a: float, l: float, h: float) -> float: + return min(max(a, l), h) + + low = low if low is not None else default_low + high = high if high is not None else default_high + + if any(isinstance(value, float) and math.isnan(value) for value in [low, high]): + raise ValueError( + f"`low` and `high` cannot be NaN, but got {low=} and {high=}" + ) + elif low == high and dtype in _FLOATING_OR_COMPLEX_TYPES: + warnings.warn( + "Passing `low==high` to `torch.testing.make_tensor` for floating or complex types " + "is deprecated since 2.1 and will be removed in 2.3. " + "Use `torch.full(...)` instead.", + FutureWarning, + stacklevel=3, + ) + elif low >= high: + raise ValueError(f"`low` must be less than `high`, but got {low} >= {high}") + elif high < lowest_inclusive or low >= highest_exclusive: + raise ValueError( + f"The value interval specified by `low` and `high` is [{low}, {high}), " + f"but {dtype} only supports [{lowest_inclusive}, {highest_exclusive})" + ) + + low = clamp(low, lowest_inclusive, highest_exclusive) + high = clamp(high, lowest_inclusive, highest_exclusive) + + if dtype in _BOOLEAN_OR_INTEGRAL_TYPES: + # 1. `low` is ceiled to avoid creating values smaller than `low` and thus outside the specified interval + # 2. Following the same reasoning as for 1., `high` should be floored. However, the higher bound of + # `torch.randint` is exclusive, and thus we need to ceil here as well. + return math.ceil(low), math.ceil(high) + + return low, high + + if len(shape) == 1 and isinstance(shape[0], collections.abc.Sequence): + shape = shape[0] # type: ignore[assignment] + shape = cast(tuple[int, ...], tuple(shape)) + + if noncontiguous and memory_format is not None: + raise ValueError( + f"The parameters `noncontiguous` and `memory_format` are mutually exclusive, " + f"but got {noncontiguous=} and {memory_format=}" + ) + + if requires_grad and dtype in _BOOLEAN_OR_INTEGRAL_TYPES: + raise ValueError( + f"`requires_grad=True` is not supported for boolean and integral dtypes, but got {dtype=}" + ) + + noncontiguous = noncontiguous and functools.reduce(lambda x, y: x * y, shape, 1) > 1 + if noncontiguous: + # Double the size of the shape in the last dimension, so that we have + # non-identical values when we make the non-contiguous operation. + shape = cast(tuple[int, ...], (*shape[:-1], 2 * shape[-1])) + + if dtype is torch.bool: + low, high = cast( + tuple[int, int], + modify_low_high( + low, + high, + lowest_inclusive=0, + highest_exclusive=2, + default_low=0, + default_high=2, + ), + ) + result = torch.randint(low, high, shape, device=device, dtype=dtype) + elif dtype in _BOOLEAN_OR_INTEGRAL_TYPES: + low, high = cast( + tuple[int, int], + modify_low_high( + low, + high, + lowest_inclusive=torch.iinfo(dtype).min, + highest_exclusive=torch.iinfo(dtype).max + # In theory, `highest_exclusive` should always be the maximum value + 1. However, `torch.randint` + # internally converts the bounds to an int64 and would overflow. In other words: `torch.randint` cannot + # sample 2**63 - 1, i.e. the maximum value of `torch.int64` and we need to account for that here. + + (1 if dtype is not torch.int64 else 0), + # This is incorrect for `torch.uint8`, but since we clamp to `lowest`, i.e. 0 for `torch.uint8`, + # _after_ we use the default value, we don't need to special case it here + default_low=-9, + default_high=10, + ), + ) + result = torch.randint(low, high, shape, device=device, dtype=dtype) + elif dtype in _FLOATING_OR_COMPLEX_TYPES: + low, high = modify_low_high( + low, + high, + lowest_inclusive=torch.finfo(dtype).min, + highest_exclusive=torch.finfo(dtype).max, + default_low=-9, + default_high=9, + ) + result = torch.empty(shape, device=device, dtype=dtype) + _uniform_random_( + torch.view_as_real(result) if dtype in _COMPLEX_TYPES else result, low, high + ) + elif dtype in _FLOATING_8BIT_TYPES: + low, high = modify_low_high( + low, + high, + lowest_inclusive=torch.finfo(dtype).min, + highest_exclusive=torch.finfo(dtype).max, + default_low=-9, + default_high=9, + ) + result = torch.empty(shape, device=device, dtype=torch.float32) + _uniform_random_(result, low, high) + result = result.to(dtype) + else: + raise TypeError( + f"The requested dtype '{dtype}' is not supported by torch.testing.make_tensor()." + " To request support, file an issue at: https://github.com/pytorch/pytorch/issues" + ) + + if noncontiguous: + # Offset by 1 to also catch offsetting issues + result = result[..., 1::2] + elif memory_format is not None: + result = result.clone(memory_format=memory_format) + + if exclude_zero: + result[result == 0] = ( + 1 if dtype in _BOOLEAN_OR_INTEGRAL_TYPES else torch.finfo(dtype).tiny + ) + + if dtype in _FLOATING_OR_COMPLEX_TYPES: + result.requires_grad = requires_grad + + return result diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/autocast_test_lists.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/autocast_test_lists.py new file mode 100644 index 0000000000000000000000000000000000000000..2f0f02e507dd7047d4675fd6fb54f0ddd763ed23 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/autocast_test_lists.py @@ -0,0 +1,473 @@ +# mypy: ignore-errors + +import collections + +import torch +from torch.testing._internal.common_utils import TEST_WITH_ROCM +from torch.testing._internal.common_utils import TestCase + + +class AutocastTestLists: + def _rnn_cell_args(self, n, num_chunks, is_lstm, dev, dtype): + input = (torch.randn((n, n), device=dev, dtype=torch.float32),) + + hx = ((torch.randn((n, n), device=dev, dtype=torch.float32), + torch.randn((n, n), device=dev, dtype=torch.float32)) if is_lstm else + torch.randn((n, n), device=dev, dtype=torch.float32),) + + weights = (torch.randn((num_chunks * n, n), device=dev, dtype=torch.float32), # weight_ih + torch.randn((num_chunks * n, n), device=dev, dtype=torch.float32), # weight_hh + torch.randn((num_chunks * n), device=dev, dtype=torch.float32), # bias_ih + torch.randn((num_chunks * n), device=dev, dtype=torch.float32)) # bias_hh + + # returns args as a tuple + return input + hx + weights + + # Supplies ops and arguments for test_autocast_* in test/test_cuda.py + def __init__(self, dev): + super().__init__() + n = 8 + # Utility arguments, created as one-element tuples + pointwise0_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),) + pointwise1_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),) + pointwise2_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),) + mat0_fp16 = (torch.randn((n, n), dtype=torch.float16, device=dev),) + mat1_fp16 = (torch.randn((n, n), dtype=torch.float16, device=dev),) + mat2_fp16 = (torch.randn((n, n), dtype=torch.float16, device=dev),) + + dimsets = ((n, n, n), (n, n, n, n), (n, n, n, n, n)) + conv_args_fp32 = [(torch.randn(dimset, dtype=torch.float32, device=dev), + torch.randn(dimset, dtype=torch.float32, device=dev)) + for dimset in dimsets] + bias_fp32 = (torch.randn((n,), dtype=torch.float32, device=dev),) + element0_fp32 = (torch.randn(1, dtype=torch.float32, device=dev),) + pointwise0_fp32 = (torch.randn(n, dtype=torch.float32, device=dev),) + pointwise1_fp32 = (torch.randn(n, dtype=torch.float32, device=dev),) + mat0_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),) + mat1_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),) + mat2_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),) + mat3_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),) + + # The lists below organize ops that autocast needs to test. + # self.list_name corresponds to test_autocast_list_name in test/test_cuda.py. + # Each op is associated with a tuple of valid arguments. + # In addition, cudnn conv ops are not supported on ROCm and hence will + # be skipped by passing TEST_WITH_ROCM flag to those ops in self.torch_fp16 list. + + # Some ops implement built-in type promotion. These don't need autocasting, + # but autocasting relies on their promotion, so we include tests to double-check. + self.torch_expect_builtin_promote = [ + ("eq", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("ge", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("gt", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("le", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("lt", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("ne", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("add", pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("div", pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("mul", pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("cat", (pointwise0_fp16 + pointwise1_fp32,), torch.float32), + ("equal", pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("stack", (pointwise0_fp16 + pointwise1_fp32,), torch.float32), + ] + self.methods_expect_builtin_promote = [ + ("__eq__", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__ge__", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__gt__", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__le__", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__lt__", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__ne__", pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__add__", pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("__div__", pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("__mul__", pointwise0_fp32 + pointwise1_fp16, torch.float32), + ] + + # The remaining lists organize ops that autocast treats explicitly. + self.torch_fp16 = [ + # deprecated _convolution + ("_convolution", conv_args_fp32[1] + bias_fp32 + ((1, 1), (0, 0), (1, 1), False, + (0, 0), 1, False, True, True)), + # the current _convolution + ("_convolution", conv_args_fp32[1] + bias_fp32 + ((1, 1), (0, 0), (1, 1), False, + (0, 0), 1, False, True, True, True)), + ("conv1d", conv_args_fp32[0]), + ("conv2d", conv_args_fp32[1]), + ("conv3d", conv_args_fp32[2]), + ("conv_tbc", conv_args_fp32[0] + bias_fp32), + ("conv_transpose1d", conv_args_fp32[0]), + ("conv_transpose2d", conv_args_fp32[1]), + ("conv_transpose3d", conv_args_fp32[2]), + ("convolution", conv_args_fp32[1] + bias_fp32 + ((1, 1), (0, 0), (1, 1), False, (0, 0), 1)), + ("cudnn_convolution", conv_args_fp32[1] + ((0, 0), (1, 1), (1, 1), 1, False, True, True), TEST_WITH_ROCM), + ("cudnn_convolution_transpose", conv_args_fp32[1] + ((0, 0), (0, 0), (1, 1), + (1, 1), 1, False, True, True), TEST_WITH_ROCM), + ("prelu", pointwise0_fp32 + element0_fp32), + ("addmm", mat1_fp32 + mat2_fp32 + mat3_fp32), + ("addmv", pointwise0_fp32 + mat2_fp32 + pointwise1_fp32), + ("addr", mat0_fp32 + pointwise0_fp32 + pointwise1_fp32), + ("matmul", mat0_fp32 + mat1_fp32), + ("einsum", "bkhd,bqhd->bqkh", mat0_fp32 + mat1_fp32), + ("mm", mat0_fp32 + mat1_fp32), + ("mv", mat0_fp32 + pointwise0_fp32), + ("chain_matmul", mat0_fp32 + mat1_fp32 + mat2_fp32), + ("addbmm", mat0_fp32 + (torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32))), + ("baddbmm", (torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32))), + ("bmm", (torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32))), + # _thnn_fused_lstm_cell and _thnn_fused_gru_cell are not Python-exposed as far as I can tell. + # ("_thnn_fused_lstm_cell", mat0_fp32 + mat1_fp32 + mat2_fp32 + pointwise0_fp32 + pointwise1_fp32), + # ("_thnn_fused_gru_cell", mat0_fp32 + mat1_fp32 + mat2_fp32 + pointwise0_fp32 + pointwise1_fp32), + ("lstm_cell", self._rnn_cell_args(n, num_chunks=4, is_lstm=True, dev=dev, dtype=torch.float32)), + ("gru_cell", self._rnn_cell_args(n, num_chunks=3, is_lstm=False, dev=dev, dtype=torch.float32)), + ("rnn_tanh_cell", self._rnn_cell_args(n, num_chunks=1, is_lstm=False, dev=dev, dtype=torch.float32)), + ("rnn_relu_cell", self._rnn_cell_args(n, num_chunks=1, is_lstm=False, dev=dev, dtype=torch.float32)), + ] + self.torch_fp32 = [ + ("acos", (pointwise0_fp16[0].clamp(-.9, 0.9),)), + ("asin", (pointwise0_fp16[0].clamp(-.9, 0.9),)), + ("cosh", pointwise0_fp16), + ("erfinv", (pointwise0_fp16[0].clamp(-.9, .9),)), + ("exp", pointwise0_fp16), + ("expm1", pointwise0_fp16), + ("log", (pointwise0_fp16[0].clamp(0.1, 100.0),)), + ("log10", (pointwise0_fp16[0].clamp(0.1, 100.0),)), + ("log2", (pointwise0_fp16[0].clamp(0.1, 100.0),)), + ("log1p", (pointwise0_fp16[0].clamp(-0.9, 100.0),)), + ("reciprocal", pointwise0_fp16), + ("rsqrt", (pointwise0_fp16[0].clamp(0.0, 100.0),)), + ("sinh", pointwise0_fp16), + ("tan", (pointwise0_fp16[0].clamp(-3.1 / 2, 3.1 / 2),)), + ("pow", ((pointwise0_fp16[0] + 1.).clamp(0.0, 100.0),) + pointwise1_fp16), + ("pow", ((pointwise0_fp16[0] + 1.).clamp(0.0, 100.0),) + (1.7,)), + # ("pow", (1.7,) + pointwise0_fp16), # This variant has a backend, but is not documented in the API. + ("softmax", pointwise0_fp16 + (0,)), + ("log_softmax", pointwise0_fp16 + (0,)), + ("layer_norm", pointwise0_fp16 + ((pointwise0_fp16[0].numel(),),)), + ("rms_norm", pointwise0_fp16 + ((pointwise0_fp16[0].numel(),),)), + ("group_norm", mat0_fp16 + (1,)), + ("norm", pointwise0_fp16), + ("norm", pointwise0_fp16, {"dim": 0}), + # these need magma + # ("norm", mat0_fp16, {"p": "nuc"}), + # ("norm", mat0_fp16, {"p": "nuc", "dim": 0}), + ("norm", pointwise0_fp16, {"p": 1}), + ("norm", pointwise0_fp16, {"p": 1, "dim": 0}), + ("cosine_similarity", mat0_fp16 + mat1_fp16), + ("poisson_nll_loss", mat0_fp16 + mat1_fp16 + (True, False, 1.e-8, torch.nn._reduction.get_enum('mean'))), + ("cosine_embedding_loss", (torch.tensor([[1, 2, 3]], device=dev, dtype=torch.float16), + torch.tensor([[1, 3, 4]], device=dev, dtype=torch.float16), + torch.tensor([1], device=dev, dtype=torch.int))), + ("hinge_embedding_loss", mat0_fp16 + (torch.ones(n, device=dev, dtype=torch.int),)), + ("kl_div", mat0_fp16 + (torch.rand((n, n), device=dev, dtype=torch.float16),)), + ("margin_ranking_loss", mat0_fp16 + mat1_fp16 + (torch.ones((n,), device=dev, dtype=torch.float16),)), + ("triplet_margin_loss", mat0_fp16 + mat1_fp16 + mat2_fp16), + ("binary_cross_entropy_with_logits", mat0_fp16 + (torch.rand((n, n), device=dev, dtype=torch.float16),)), + ("cumprod", pointwise0_fp16 + (0,)), + ("cumsum", pointwise0_fp16 + (0,)), + ("dist", pointwise0_fp16 + pointwise1_fp16), + ("pdist", mat0_fp16), + ("cdist", mat0_fp16 + mat1_fp16), + ("prod", pointwise0_fp16), + ("prod", pointwise0_fp16 + (0,)), + ("renorm", mat0_fp16 + (2, 0, 1.0)), + ("sum", pointwise0_fp16), + ("sum", mat0_fp16 + (1,)), + ("logsumexp", mat0_fp16 + (1,)), + ] + self.torch_need_autocast_promote = [ + ("addcdiv", pointwise0_fp32 + pointwise1_fp16 + (pointwise2_fp16[0].clamp(0.1, 100),)), + ("addcmul", pointwise0_fp32 + pointwise1_fp16 + pointwise2_fp16), + ("atan2", pointwise0_fp32 + (pointwise1_fp16[0].clamp(0.1, 100),)), + ("bilinear", (torch.randn((1, 2), dtype=torch.float16, device=dev), + torch.randn((1, 2), dtype=torch.float32, device=dev), + torch.randn((1, 2, 2), dtype=torch.float16, device=dev), + torch.randn((1,), dtype=torch.float32, device=dev))), + ("cross", (torch.randn(3, dtype=torch.float32, device=dev), + torch.randn(3, dtype=torch.float16, device=dev))), + ("dot", pointwise0_fp16 + pointwise1_fp32), + ("vdot", pointwise0_fp16 + pointwise1_fp32), + ("grid_sampler", (torch.randn((2, 3, 33, 22), dtype=torch.float16, device=dev), + torch.randn((2, 22, 11, 2), dtype=torch.float32, device=dev), + 0, 0, False)), + ("index_put", pointwise0_fp32 + ((torch.tensor([1], device=dev, dtype=torch.long),), + torch.randn(1, device=dev, dtype=torch.float16))), + ("index_put", pointwise0_fp16 + ((torch.tensor([1], device=dev, dtype=torch.long),), + torch.randn(1, device=dev, dtype=torch.float32))), + ("tensordot", (torch.randn((2, 2, 2), dtype=torch.float32, device=dev), + torch.randn((2, 2, 2), dtype=torch.float16, device=dev))), + ("scatter_add", (torch.zeros(2, 2, 2, dtype=torch.float32, device=dev), + 0, + torch.randint(0, 2, (2, 2, 2), device=dev), + torch.randn((2, 2, 2), dtype=torch.float16, device=dev))), + ("scatter_add", (torch.zeros(2, 2, 2, dtype=torch.float16, device=dev), + 0, + torch.randint(0, 2, (2, 2, 2), device=dev), + torch.randn((2, 2, 2), dtype=torch.float32, device=dev))), + ] + self.nn_fp16 = [ + ("linear", mat0_fp32 + mat1_fp32 + mat2_fp32), + ] + self.nn_fp32 = [ + ("softplus", pointwise0_fp16), + ("nll_loss", (torch.rand((n, n), device=dev, dtype=torch.float), + torch.zeros((n,), device=dev, dtype=torch.long))), + ("nll_loss2d", (torch.rand((n, n, n, n), device=dev, dtype=torch.half), + torch.zeros((n, n, n), device=dev, dtype=torch.long))), + ("l1_loss", mat0_fp16 + mat1_fp16), + ("smooth_l1_loss", mat0_fp16 + mat1_fp16), + ("mse_loss", mat0_fp16 + mat1_fp16), + ("multilabel_margin_loss", mat0_fp16 + (torch.ones((n, n), device=dev, dtype=torch.long),)), + ("soft_margin_loss", mat0_fp16 + (torch.ones((n, n), device=dev, dtype=torch.long),)), + ("multi_margin_loss", mat0_fp16 + (torch.ones((n,), device=dev, dtype=torch.long),)), + ] + self.linalg_fp16 = [ + ("linalg_vecdot", mat0_fp32 + mat0_fp32), + ("linalg_multi_dot", (mat0_fp32 + mat1_fp32 + mat2_fp32,)), + ] + self.methods_fp16 = [ + ("__matmul__", mat0_fp32 + mat1_fp32) + ] + self.methods_fp32 = [ + ("__pow__", (torch.rand(n, device=dev, dtype=torch.float16), 1.5)), + ] + self.banned = [ + ("binary_cross_entropy", (torch.rand((n, n), device=dev, dtype=torch.float32), + torch.rand((n, n), device=dev, dtype=torch.float32)), torch._C._nn), + ] + + +class AutocastCPUTestLists: + # Supplies ops and arguments for test_autocast_* in test/test_cpu.py + def __init__(self, dev): + super().__init__() + n = 8 + # Utility arguments, created as one-element tuples + pointwise0_bf16 = (torch.randn(n, dtype=torch.bfloat16, device=dev),) + pointwise1_bf16 = (torch.randn(n, dtype=torch.bfloat16, device=dev),) + mat0_bf16 = (torch.randn((n, n), dtype=torch.bfloat16, device=dev),) + mat1_bf16 = (torch.randn((n, n), dtype=torch.bfloat16, device=dev),) + mat2_bf16 = (torch.randn((n, n), dtype=torch.bfloat16, device=dev),) + + pointwise0_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),) + pointwise1_fp16 = (torch.randn(n, dtype=torch.float16, device=dev),) + + dummy_dimsets = ((n,), (n, n), (n, n, n), (n, n, n, n), (n, n, n, n, n)) + + dummy_bf16 = [(torch.randn(dimset, dtype=torch.bfloat16, device=dev),) + for dimset in dummy_dimsets] + + dimsets = ((n, n, n), (n, n, n, n), (n, n, n, n, n)) + conv_args_fp32 = [(torch.randn(dimset, dtype=torch.float32, device=dev), + torch.randn(dimset, dtype=torch.float32, device=dev)) + for dimset in dimsets] + + element0_fp32 = (torch.randn(1, dtype=torch.float32, device=dev),) + pointwise0_fp32 = (torch.randn(n, dtype=torch.float32, device=dev),) + pointwise1_fp32 = (torch.randn(n, dtype=torch.float32, device=dev),) + mat0_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),) + mat1_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),) + mat2_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),) + mat3_fp32 = (torch.randn((n, n), dtype=torch.float32, device=dev),) + + dummy_fp32 = [ # noqa: F841 + (torch.randn(dimset, dtype=torch.float32, device=dev),) + for dimset in dummy_dimsets + ] + # The lists below organize ops that autocast needs to test. + # self.list_name corresponds to test_autocast_list_name in test/test_cpu.py. + # Each op is associated with a tuple of valid arguments. + + # Some ops implement built-in type promotion. These don't need autocasting, + # but autocasting relies on their promotion, so we include tests to double-check. + self.torch_expect_builtin_promote = [ + ("eq", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("ge", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("gt", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("le", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("lt", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("ne", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("add", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("div", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("mul", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32), + ] + + self.methods_expect_builtin_promote = [ + ("__eq__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__ge__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__gt__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__le__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__lt__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__ne__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.bool), + ("__add__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("__div__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32), + ("__mul__", pointwise0_fp32 + pointwise1_bf16, pointwise0_fp32 + pointwise1_fp16, torch.float32), + ] + # The remaining lists organize ops that autocast treats explicitly. + self.torch_16 = [ + ("conv1d", conv_args_fp32[0]), + ("conv2d", conv_args_fp32[1]), + ("conv3d", conv_args_fp32[2]), + ("bmm", (torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32))), + ("mm", mat0_fp32 + mat1_fp32), + ("matmul", mat0_fp32 + mat1_fp32), + ("baddbmm", (torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32))), + ("addmm", mat1_fp32 + mat2_fp32 + mat3_fp32), + ("_addmm_activation", mat1_fp32 + mat2_fp32 + mat3_fp32, {"beta": 1, "alpha": 1, "use_gelu": True}), + ("addbmm", mat0_fp32 + (torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32))), + ("conv_tbc", (torch.randn((10, 7, 3), device=dev, dtype=torch.float32), + torch.randn((5, 3, 5), device=dev, dtype=torch.float32), + torch.randn(5, device=dev, dtype=torch.float32), + 0)), + ("conv_transpose1d", conv_args_fp32[0]), + ("conv_transpose2d", conv_args_fp32[1]), + ("conv_transpose3d", conv_args_fp32[2]), + ("prelu", pointwise0_fp32 + element0_fp32), + ("_native_multi_head_attention", (torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32), + torch.randn((n, n, n), device=dev, dtype=torch.float32), + n, 4, torch.randn((3 * n, n), device=dev, dtype=torch.float32), + torch.randn((3 * n), device=dev, dtype=torch.float32), + torch.randn((n, n), device=dev, dtype=torch.float32), + torch.randn((n), device=dev, dtype=torch.float32))), + ] + self.torch_fp32 = [ + ("poisson_nll_loss", mat0_bf16 + mat1_bf16 + (True, False, 1.e-8, torch.nn._reduction.get_enum('mean'))), + ("cosine_embedding_loss", (torch.tensor([[1, 2, 3]], device=dev, dtype=torch.bfloat16), + torch.tensor([[1, 3, 4]], device=dev, dtype=torch.bfloat16), + torch.tensor([1], device=dev, dtype=torch.int))), + ("hinge_embedding_loss", mat0_bf16 + (torch.ones(n, device=dev, dtype=torch.int),)), + ("margin_ranking_loss", mat0_bf16 + mat1_bf16 + (torch.ones((n,), device=dev, dtype=torch.bfloat16),)), + ("triplet_margin_loss", mat0_bf16 + mat1_bf16 + mat2_bf16), + ("binary_cross_entropy_with_logits", mat0_bf16 + (torch.rand((n, n), device=dev, dtype=torch.bfloat16),)), + ] + self.nn_16 = [ + ("linear", mat0_fp32 + mat1_fp32, {}), + ] + self.nn_fp32 = [ + ("avg_pool3d", dummy_bf16[3], {"kernel_size": (3, 3, 3), "stride": (1, 1, 1)}), + ("binary_cross_entropy", (torch.rand((n, n), device=dev, dtype=torch.bfloat16),) + + (torch.rand((n, n), device=dev, dtype=torch.bfloat16),)), + ("reflection_pad1d", dummy_bf16[2], {"padding": (3, 3)}), + ("nll_loss", (torch.rand((n, n), device=dev, dtype=torch.bfloat16), + torch.zeros((n,), device=dev, dtype=torch.long))), + ("nll_loss2d", (torch.rand((n, n, n, n), device=dev, dtype=torch.bfloat16), + torch.zeros((n, n, n), device=dev, dtype=torch.long))), + ("l1_loss", mat0_bf16 + mat1_bf16), + ("smooth_l1_loss", mat0_bf16 + mat1_bf16), + ("mse_loss", mat0_bf16 + mat1_bf16), + ("multilabel_margin_loss", mat0_bf16 + (torch.ones((n, n), device=dev, dtype=torch.long),)), + ("soft_margin_loss", mat0_bf16 + (torch.ones((n, n), device=dev, dtype=torch.long),)), + ("multi_margin_loss", mat0_bf16 + (torch.ones((n,), device=dev, dtype=torch.long),)), + ("huber_loss", mat0_bf16 + mat1_bf16), + ] + self.torch_need_autocast_promote = [ + ("cat", (pointwise0_bf16 + pointwise1_fp32,), (pointwise0_fp16 + pointwise1_fp32,)), + ("stack", (pointwise0_bf16 + pointwise1_fp32,), (pointwise0_fp16 + pointwise1_fp32,)), + ] + + +class TestAutocast(TestCase): + def args_maybe_kwargs(self, op_with_args): + if len(op_with_args) == 2: + return op_with_args[0], op_with_args[1], {} + else: + return op_with_args[0], op_with_args[1], op_with_args[2] + + def _run_autocast_outofplace( + self, + op, + args, + run_as_type, + device, + out_type=None, + module=torch, + add_kwargs=None, + amp_dtype=torch.bfloat16, + ): + # helper to cast args + def cast(val, to_type): + if isinstance(val, torch.Tensor): + return val.to(to_type) if val.is_floating_point() else val + elif isinstance(val, collections.abc.Iterable): + return type(val)(cast(v, to_type) for v in val) + else: + return val + + if add_kwargs is None: + add_kwargs = {} + + self.assertFalse(torch.is_autocast_enabled(device_type=device)) + with torch.amp.autocast(device_type=device, dtype=amp_dtype): + self.assertTrue(torch.is_autocast_enabled(device_type=device)) + + out_type = out_type if out_type is not None else run_as_type + output = output_method = None + + # Try module.* variant, if requested: + if module is not None and hasattr(module, op): + output = getattr(module, op)(*args, **add_kwargs) + if isinstance(output, torch.Tensor): + self.assertTrue( + out_type == output.dtype, + f"autocast for torch.{op} produced {output.dtype}, should produce {out_type}", + ) + # Try Tensor.* variant: + if hasattr(torch.Tensor, op): + output_method = getattr(args[0], op)(*args[1:], **add_kwargs) + if isinstance(output_method, torch.Tensor): + self.assertTrue( + out_type == output_method.dtype, + f"autocast for torch.{op} produced {output_method.dtype}, should produce torch.{out_type}", + ) + + self.assertTrue( + (output is not None) or (output_method is not None), + f"{op} not found as an attribute on either Tensor or the requested module {module}", + ) + + # Accounts for ops that return Tensors, iterables, and other non-Tensors. + # For example, lstm_cell returns a tuple and equal returns bool. + def compare(first, second): + if isinstance(first, torch.Tensor): + return torch.equal(first, second) + elif isinstance(first, collections.abc.Iterable): + return all(compare(f, s) for f, s in zip(first, second, strict=False)) + else: + return first == second + + # If both torch.* and Tensor.* variants were found, check outputs are identical + if (output is not None) and (output_method is not None): + self.assertTrue(type(output) is type(output_method)) + comparison = compare(output, output_method) + self.assertTrue( + comparison, f"torch.{op} result did not match Tensor.{op} result" + ) + + # Compare numerics to Python-side "autocasting" that (we expect) does the same thing + # as the C++-side autocasting, and should be bitwise accurate. + output_to_compare = output if output is not None else output_method + with torch.amp.autocast(device_type=device, enabled=False): + self.assertFalse( + torch.is_autocast_enabled(device_type=device) + ) + + if module is not None and hasattr(module, op): + control = getattr(module, op)( + *cast(args, run_as_type), **add_kwargs + ) + else: + control = getattr(args[0].to(run_as_type), op)( + *cast(args[1:], run_as_type), **add_kwargs + ) + self.assertTrue(type(output_to_compare) is type(control)) + comparison = compare(output_to_compare, control) + self.assertTrue(comparison, f"torch.{op} result did not match control") + self.assertTrue(torch.is_autocast_enabled(device_type=device)) + self.assertFalse(torch.is_autocast_enabled(device_type=device)) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/autograd_function_db.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/autograd_function_db.py new file mode 100644 index 0000000000000000000000000000000000000000..3eba9f80603db8c9bfac1c91f0cd57be8909ca24 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/autograd_function_db.py @@ -0,0 +1,635 @@ +# mypy: ignore-errors + +import torch +from functools import partial +from torch.testing import make_tensor +from torch.testing._internal.opinfo.core import ( + OpInfo, + SampleInput, +) +from torch.testing._internal.common_dtype import all_types_and +import numpy as np + +# Note: [autograd.Function db] +# +# This is a collection of autograd.Function test cases written as OpInfos +# so they can easily be consumed by OpInfo-based tests to check if a subsystem +# supports autograd.Function. +# +# Axes: +# - saves {output, input, intermediate, non-tensor} +# - {inputs, output} x {single tensor, tensors, arbitrary objects} +# - Uses {mark_dirty, mark_non_differentiable, once_differentiable} + + +def to_numpy(tensor): + return tensor.cpu().numpy() + + +class NumpyCube(torch.autograd.Function): + @staticmethod + def forward(input): + input_np = to_numpy(input) + dinput = torch.tensor(3 * input_np ** 2, device=input.device) + return torch.tensor(input_np ** 3, device=input.device), dinput + + @staticmethod + def setup_context(ctx, inputs, output): + ctx.save_for_backward(inputs[0], output[1]) + ctx.save_for_forward(inputs[0], output[1]) + + @staticmethod + def backward(ctx, grad_output, grad_saved): + input, dinput = ctx.saved_tensors + return NumpyMul.apply(grad_output, dinput) + 6 * NumpyMul.apply(grad_saved, input) + + @staticmethod + def vmap(info, in_dims, input): + result = NumpyCube.apply(input) + return result, (in_dims[0], in_dims[0]) + + @staticmethod + def jvp(ctx, input_tangent): + input, dinput = ctx.saved_tensors + return NumpyMul.apply(input_tangent, dinput), 6 * NumpyMul.apply(input_tangent, input) + + +class CubeGenVmap(torch.autograd.Function): + generate_vmap_rule = True + + @staticmethod + def forward(x): + return x ** 3, 3 * x ** 2 + + @staticmethod + def setup_context(ctx, inputs, outputs): + ctx.save_for_backward(inputs[0], outputs[1]) + ctx.save_for_forward(inputs[0], outputs[1]) + + @staticmethod + def backward(ctx, grad_output, grad_saved): + _input, dinput = ctx.saved_tensors + result = grad_output * dinput + 6 * dinput + return result + + @staticmethod + def jvp(ctx, input_tangent): + input, dinput = ctx.saved_tensors + return MulGenVmap.apply(input_tangent, dinput), 6 * NumpyMul.apply(input_tangent, input) + + +def sample_inputs_numpy_cube(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg(1, low=0.8, high=2), args=()) + + +class NumpyCubeNotComposable(torch.autograd.Function): + @staticmethod + def forward(input): + input_np = to_numpy(input) + return torch.tensor(input_np ** 3, device=input.device), input_np + + @staticmethod + def setup_context(ctx, inputs, output): + _, input_np = output + ctx.input_np = input_np + ctx.device = inputs[0].device + + @staticmethod + @torch.autograd.function.once_differentiable + def backward(ctx, grad_output, grad_saved): + result_np = 3 * (ctx.input_np ** 2) + return torch.tensor(result_np, device=ctx.device) + + +class NumpyMul(torch.autograd.Function): + @staticmethod + def forward(x, y): + return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device) + + @staticmethod + def setup_context(ctx, inputs, output): + ctx.save_for_backward(*inputs) + ctx.save_for_forward(*inputs) + + @staticmethod + def backward(ctx, grad_output): + x, y = ctx.saved_tensors + gx = None + if ctx.needs_input_grad[0]: + gx = NumpyMul.apply(grad_output, y) + gy = None + if ctx.needs_input_grad[1]: + gy = NumpyMul.apply(grad_output, x) + return gx, gy + + @staticmethod + def vmap(info, in_dims, x, y): + x_bdim, y_bdim = in_dims + x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1) + y = y.movedim(y_bdim, -1) if y_bdim is not None else y.unsqueeze(-1) + result = NumpyMul.apply(x, y) + result = result.movedim(-1, 0) + return result, 0 + + @staticmethod + def jvp(ctx, x_tangent, y_tangent): + x, y = ctx.saved_tensors + return x_tangent * y + y_tangent * x + +def sample_inputs_numpy_mul(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + # Broadcasting + yield SampleInput(make_arg(4, low=0.9, high=2), args=(make_arg(3, 4, low=0.9, high=2),)) + +def sample_inputs_numpy_mul_scalar(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg(4, low=0.9, high=2), args=(), kwargs={"scalar": 3.14}) + +class MulGenVmap(torch.autograd.Function): + generate_vmap_rule = True + + @staticmethod + def forward(x, y): + return x * y + + @staticmethod + def setup_context(ctx, inputs, outputs): + ctx.save_for_backward(*inputs) + ctx.save_for_forward(*inputs) + + @staticmethod + def backward(ctx, grad_output): + x, y = ctx.saved_tensors + gx = None + if ctx.needs_input_grad[0]: + gx = MulGenVmap.apply(grad_output, y) + gy = None + if ctx.needs_input_grad[1]: + gy = MulGenVmap.apply(grad_output, x) + return gx, gy + + @staticmethod + def jvp(ctx, x_tangent, y_tangent): + x, y = ctx.saved_tensors + return x_tangent * y + y_tangent * x + + +class NumpyExp_(torch.autograd.Function): + @staticmethod + def forward(x): + x_np = to_numpy(x) + np.exp(x_np, x_np) + return x + + @staticmethod + def setup_context(ctx, inputs, output): + x, = inputs + ctx.mark_dirty(x) + ctx.save_for_backward(output) + ctx.save_for_forward(output) + + @staticmethod + def backward(ctx, grad_output): + output, = ctx.saved_tensors + return NumpyMul.apply(grad_output, output) + + @staticmethod + def vmap(info, in_dims, x): + NumpyExp_.apply(x) + return x, in_dims[0] + + @staticmethod + def jvp(ctx, x_tangent): + # Doesn't call numpy operations because I didn't want to write NumpyMul_ + output, = ctx.saved_tensors + x_tangent.mul_(output) + return x_tangent + +class NumpySort(torch.autograd.Function): + @staticmethod + def forward(x, dim): + device = x.device + x = to_numpy(x) + ind = np.argsort(x, axis=dim) + ind_inv = np.argsort(ind, axis=dim) + return ( + torch.tensor(x, device=device), + torch.tensor(ind, device=device), + torch.tensor(ind_inv, device=device), + ) + + @staticmethod + def setup_context(ctx, inputs, output): + _x, dim = inputs + _, ind, ind_inv = output + ctx.mark_non_differentiable(ind, ind_inv) + ctx.save_for_backward(ind, ind_inv) + ctx.save_for_forward(ind, ind_inv) + ctx.dim = dim + + @staticmethod + def backward(ctx, grad_output, _0, _1): + ind, ind_inv = ctx.saved_tensors + return NumpyTake.apply(grad_output, ind_inv, ind, ctx.dim), None + + @staticmethod + def vmap(info, in_dims, x, dim): + x_bdim, _ = in_dims + x = x.movedim(x_bdim, 0) + # wrap dim + dim = dim if dim >= 0 else dim + x.dim() - 1 + return NumpySort.apply(x, dim + 1), (0, 0, 0) + + @staticmethod + def jvp(ctx, x_tangent, _): + ind, ind_inv = ctx.saved_tensors + return NumpyTake.apply(x_tangent, ind, ind_inv, ctx.dim), None, None + +class SortGenVmap(torch.autograd.Function): + generate_vmap_rule = True + + @staticmethod + def forward(x, dim): + ind = torch.argsort(x, dim=dim) + ind_inv = torch.argsort(ind, axis=dim) + result = torch.take_along_dim(x, ind, dim=dim) + return result, ind, ind_inv + + @staticmethod + def setup_context(ctx, inputs, outputs): + x, dim = inputs + _, ind, ind_inv = outputs + ctx.mark_non_differentiable(ind, ind_inv) + ctx.save_for_backward(ind, ind_inv) + ctx.save_for_forward(ind, ind_inv) + ctx.dim = dim + + @staticmethod + def backward(ctx, grad_output, _0, _1): + ind, ind_inv = ctx.saved_tensors + return TakeGenVmap.apply(grad_output, ind_inv, ind, ctx.dim), None + + @staticmethod + def jvp(ctx, x_tangent, _): + ind, ind_inv = ctx.saved_tensors + return TakeGenVmap.apply(x_tangent, ind, ind_inv, ctx.dim), None, None + + +def sample_inputs_numpy_sort(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg(3, 5), args=(1,)) + + +def sample_inputs_numpy_take(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + tensor = make_arg(3, 5) + dim = 1 + _, ind, ind_inv = NumpySort.apply(tensor, 1) + yield SampleInput(tensor, args=(ind, ind_inv, dim)) + + +class NumpyTake(torch.autograd.Function): + @staticmethod + def forward(x, ind, ind_inv, dim): + device = x.device + x = to_numpy(x) + ind = to_numpy(ind) + return torch.tensor(np.take_along_axis(x, ind, dim), device=device) + + @staticmethod + def setup_context(ctx, inputs, output): + _x, ind, ind_inv, dim = inputs + ctx.save_for_backward(ind, ind_inv) + ctx.save_for_forward(ind, ind_inv) + ctx.dim = dim + + @staticmethod + def backward(ctx, grad_output): + ind, ind_inv = ctx.saved_tensors + result = NumpyTake.apply(grad_output, ind_inv, ind, ctx.dim) + return result, None, None, None + + @staticmethod + def vmap(info, in_dims, x, ind, ind_inv, dim): + x_bdim, ind_bdim, ind_inv_bdim, _ = in_dims + + # wrap dim + logical_dim = x.dim() if x_bdim is None else x_bdim - 1 + dim = dim if dim >= 0 else dim + logical_dim + + def expand_bdim(x, x_bdim): + if x_bdim is None: + return x.expand(info.batch_size, *x.shape) + return x.movedim(x_bdim, 0) + + x = expand_bdim(x, x_bdim) + ind = expand_bdim(ind, ind_bdim) + ind_inv = expand_bdim(ind_inv, ind_inv_bdim) + + return NumpyTake.apply(x, ind, ind_inv, dim + 1), 0 + + @staticmethod + def jvp(ctx, x_tangent, ind_tangent, ind_inv_tangent, _): + if ind_tangent is not None: + raise AssertionError("Expected ind_tangent to be None") + if ind_inv_tangent is not None: + raise AssertionError("Expected ind_inv_tangent to be None") + ind, ind_inv = ctx.saved_tensors + return NumpyTake.apply(x_tangent, ind, ind_inv, ctx.dim) + +class TakeGenVmap(torch.autograd.Function): + generate_vmap_rule = True + + @staticmethod + def forward(x, ind, ind_inv, dim): + return torch.take_along_dim(x, ind, dim) + + @staticmethod + def setup_context(ctx, inputs, outputs): + _x, ind, ind_inv, dim = inputs + ctx.save_for_backward(ind, ind_inv) + ctx.save_for_forward(ind, ind_inv) + ctx.dim = dim + + @staticmethod + def backward(ctx, grad_output): + ind, ind_inv = ctx.saved_tensors + result = TakeGenVmap.apply(grad_output, ind_inv, ind, ctx.dim) + return result, None, None, None + + @staticmethod + def jvp(ctx, x_tangent, ind_tangent, ind_inv_tangent, _): + ind, ind_inv = ctx.saved_tensors + return TakeGenVmap.apply(x_tangent, ind, ind_inv, ctx.dim) + +class Select(torch.autograd.Function): + @staticmethod + def forward(x, idx): + return x[idx] + + @staticmethod + def setup_context(ctx, inputs, output): + x, idx = inputs + ctx.x_shape = x.shape + ctx.idx = idx + + @staticmethod + def backward(ctx, grad_output): + result = grad_output.new_zeros(ctx.x_shape) + result[ctx.idx] = grad_output + return result, None + + @staticmethod + def vmap(info, in_dims, x, idx): + x_bdim, _ = in_dims + x = x.movedim(x_bdim, 1) + return Select.apply(x, idx), 0 + + @staticmethod + def jvp(ctx, x_tangent, _): + return Select.apply(x_tangent, ctx.idx) + +class SelectGenVmap(torch.autograd.Function): + generate_vmap_rule = True + + @staticmethod + def forward(x, idx): + return x[idx] + + @staticmethod + def setup_context(ctx, inputs, outputs): + x, idx = inputs + ctx.x_shape = x.shape + ctx.idx = idx + + @staticmethod + def backward(ctx, grad_output): + result = grad_output.new_zeros(ctx.x_shape) + result[ctx.idx] = grad_output + return result, None + + @staticmethod + def jvp(ctx, x_tangent, _): + return SelectGenVmap.apply(x_tangent, ctx.idx) + + +def sample_inputs_select(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg(3, 5), args=(2,)) + +class ScaleGradGenVmap(torch.autograd.Function): + generate_vmap_rule = True + scale = 3.14 + + @staticmethod + def forward(x): + return x.clone() + + @staticmethod + def setup_context(ctx, inputs, outputs): + pass + + @staticmethod + def backward(ctx, grad_output): + return grad_output * ScaleGradGenVmap.scale + + @staticmethod + def jvp(ctx, x_tangent): + return x_tangent * ScaleGradGenVmap.scale + +class ZeroGradientsGenVmap(torch.autograd.Function): + generate_vmap_rule = True + + @staticmethod + def forward(x, y): + return x.clone(), y.clone() + + @staticmethod + def setup_context(ctx, inputs, outputs): + pass + + @staticmethod + def backward(ctx, gx, gy): + # Intentionally returning torch.zeros instead of zeros_like or new_zeros. + # Also intentionally not None. + return ( + # Intentionally too-large gradient + torch.zeros(3, 4, *gx.shape, dtype=gx.dtype, device=gx.device), + torch.zeros(gy.shape, dtype=gy.dtype, device=gy.device), + ) + + @staticmethod + def jvp(ctx, gx, gy): + # Intentionally returning torch.zeros instead of zeros_like or new_zeros. + # Also intentionally not None. + return ( + torch.zeros(gx.shape, dtype=gx.dtype, device=gx.device), + torch.zeros(gy.shape, dtype=gy.dtype, device=gy.device), + ) + + +def sample_inputs_forward_default_args(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg(3, 5)) + + +class ForwardHasDefaultArgs(torch.autograd.Function): + @staticmethod + def forward(x, idx=(2,)): + return x[idx] + + @staticmethod + def setup_context(ctx, inputs, output): + x, idx = inputs + ctx.x_shape = x.shape + ctx.idx = idx + + @staticmethod + def backward(ctx, grad_output): + result = grad_output.new_zeros(ctx.x_shape) + result[ctx.idx] = grad_output + return result, None + + @staticmethod + def vmap(info, in_dims, x, idx): + x_bdim, _ = in_dims + x = x.movedim(x_bdim, 1) + return ForwardHasDefaultArgs.apply(x, idx), 0 + + @staticmethod + def jvp(ctx, x_tangent, _): + return ForwardHasDefaultArgs.apply(x_tangent, ctx.idx) + + +autograd_function_db = [ + OpInfo( + 'NumpyCubeAutogradFunction', + op=NumpyCube.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_numpy_cube, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'NumpyExpMarkDirtyAutogradFunction', + op=lambda x: NumpyExp_.apply(x.clone()), + inplace_variant=NumpyExp_.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_numpy_cube, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'NumpyMulAutogradFunction', + op=NumpyMul.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_numpy_mul, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'NumpyCubeNotComposableAutogradFunction', + op=lambda x: NumpyCubeNotComposable.apply(x)[0], + supports_forward_ad=False, + supports_fwgrad_bwgrad=False, + sample_inputs_func=sample_inputs_numpy_cube, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'NumpySortAutogradFunction', + op=NumpySort.apply, + supports_forward_ad=False, + supports_fwgrad_bwgrad=False, + sample_inputs_func=sample_inputs_numpy_sort, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + gradcheck_wrapper=lambda y, ind: y, + ), + OpInfo( + 'NumpyTakeAutogradFunction', + op=NumpyTake.apply, + supports_forward_ad=False, + supports_fwgrad_bwgrad=False, + sample_inputs_func=sample_inputs_numpy_take, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'SelectAutogradFunction', + op=Select.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_select, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'CubeGenVmapAutogradFunction', + op=CubeGenVmap.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_numpy_cube, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'MulGenVmapAutogradFunction', + op=MulGenVmap.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_numpy_mul, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'SortGenVmapAutogradFunction', + op=SortGenVmap.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_numpy_sort, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + gradcheck_wrapper=lambda y, ind: y, + ), + OpInfo( + 'SelectGenVmapAutogradFunction', + op=SelectGenVmap.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_select, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'ScaleGradGenVmapAutogradFunction', + op=ScaleGradGenVmap.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_numpy_cube, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'ZeroGradientsGenVmapAutogradFunction', + op=ZeroGradientsGenVmap.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_numpy_mul, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'ForwardHasDefaultArgsAutogradFunction', + op=ForwardHasDefaultArgs.apply, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_forward_default_args, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/check_kernel_launches.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/check_kernel_launches.py new file mode 100644 index 0000000000000000000000000000000000000000..d2219ef4ea56aa306dfdd3af18b7403af8384c78 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/check_kernel_launches.py @@ -0,0 +1,164 @@ +# mypy: ignore-errors + +import os +import re +import sys + +__all__ = [ + "check_code_for_cuda_kernel_launches", + "check_cuda_kernel_launches", +] + +# FILES TO EXCLUDE (match is done with suffix using `endswith`) +# You wouldn't drive without a seatbelt, though, so why would you +# launch a kernel without some safety? Use this as a quick workaround +# for a problem with the checker, fix the checker, then de-exclude +# the files in question. +exclude_files: list[str] = [] + +# Without using a C++ AST we can't 100% detect kernel launches, so we +# model them as having the pattern "<<>>(arguments);" +# We then require that `C10_CUDA_KERNEL_LAUNCH_CHECK` be +# the next statement. +# +# We model the next statement as ending at the next `}` or `;`. +# If we see `}` then a clause ended (bad) if we see a semi-colon then +# we expect the launch check just before it. +# +# Since the kernel launch can include lambda statements, it's important +# to find the correct end-paren of the kernel launch. Doing this with +# pure regex requires recursive regex, which aren't part of the Python +# standard library. To avoid an additional dependency, we build a prefix +# regex that finds the start of a kernel launch, use a paren-matching +# algorithm to find the end of the launch, and then another regex to +# determine if a launch check is present. + +# Finds potential starts of kernel launches +kernel_launch_start = re.compile( + r"^.*<<<[^>]+>>>\s*\(", flags=re.MULTILINE +) + +# This pattern should start at the character after the final paren of the +# kernel launch. It returns a match if the launch check is not the next statement +has_check = re.compile( + r"\s*;(?![^;}]*C10_CUDA_KERNEL_LAUNCH_CHECK\(\);)", flags=re.MULTILINE +) + +def find_matching_paren(s: str, startpos: int) -> int: + """Given a string "prefix (unknown number of characters) suffix" + and the position of the first `(` returns the index of the character + 1 past the `)`, accounting for paren nesting + """ + opening = 0 + for i, c in enumerate(s[startpos:]): + if c == '(': + opening += 1 + elif c == ')': + opening -= 1 + if opening == 0: + return startpos + i + 1 + + raise IndexError("Closing parens not found!") + + +def should_exclude_file(filename) -> bool: + for exclude_suffix in exclude_files: + if filename.endswith(exclude_suffix): + return True + return False + + +def check_code_for_cuda_kernel_launches(code, filename=None): + """Checks code for CUDA kernel launches without cuda error checks. + + Args: + filename - Filename of file containing the code. Used only for display + purposes, so you can put anything here. + code - The code to check + + Returns: + The number of unsafe kernel launches in the code + """ + if filename is None: + filename = "##Python Function Call##" + + # We break the code apart and put it back together to add + # helpful line numberings for identifying problem areas + code = enumerate(code.split("\n")) # Split by line breaks + code = [f"{lineno}: {linecode}" for lineno, linecode in code] # Number the lines + code = '\n'.join(code) # Put it back together + + num_launches_without_checks = 0 + for m in kernel_launch_start.finditer(code): + end_paren = find_matching_paren(code, m.end() - 1) + if has_check.match(code, end_paren): + num_launches_without_checks += 1 + context = code[m.start():end_paren + 1] + print(f"Missing C10_CUDA_KERNEL_LAUNCH_CHECK in '{filename}'. Context:\n{context}", file=sys.stderr) + + return num_launches_without_checks + + +def check_file(filename): + """Checks a file for CUDA kernel launches without cuda error checks + + Args: + filename - File to check + + Returns: + The number of unsafe kernel launches in the file + """ + if not (filename.endswith((".cu", ".cuh"))): + return 0 + if should_exclude_file(filename): + return 0 + with open(filename) as f: + contents = f.read() + unsafeCount = check_code_for_cuda_kernel_launches(contents, filename) + return unsafeCount + + +def check_cuda_kernel_launches(): + """Checks all pytorch code for CUDA kernel launches without cuda error checks + + Returns: + The number of unsafe kernel launches in the codebase + """ + torch_dir = os.path.dirname(os.path.realpath(__file__)) + torch_dir = os.path.dirname(torch_dir) # Go up to parent torch + torch_dir = os.path.dirname(torch_dir) # Go up to parent caffe2 + + kernels_without_checks = 0 + files_without_checks = [] + for root, dirnames, filenames in os.walk(torch_dir): + # `$BASE/build` and `$BASE/torch/include` are generated + # so we don't want to flag their contents + if root == os.path.join(torch_dir, "build") or root == os.path.join(torch_dir, "torch/include"): + # Curtail search by modifying dirnames and filenames in place + # Yes, this is the way to do this, see `help(os.walk)` + dirnames[:] = [] + continue + + for x in filenames: + filename = os.path.join(root, x) + file_result = check_file(filename) + if file_result > 0: + kernels_without_checks += file_result + files_without_checks.append(filename) + + if kernels_without_checks > 0: + count_str = f"Found {kernels_without_checks} instances in " \ + f"{len(files_without_checks)} files where kernel " \ + "launches didn't have checks." + print(count_str, file=sys.stderr) + print("Files without checks:", file=sys.stderr) + for x in files_without_checks: + print(f"\t{x}", file=sys.stderr) + print(count_str, file=sys.stderr) + + return kernels_without_checks + + +if __name__ == "__main__": + unsafe_launches = check_cuda_kernel_launches() + sys.exit(0 if unsafe_launches == 0 else 1) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/codegen/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/codegen/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1e3572cfc4c6a0ddc3d8fa2e1b056415204acdfa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/codegen/__init__.py @@ -0,0 +1 @@ +# mypy: ignore-errors diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_cuda.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_cuda.py new file mode 100644 index 0000000000000000000000000000000000000000..51c2f1fb7e21511863eeb37c4fe1be8ab485d19f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_cuda.py @@ -0,0 +1,497 @@ +# mypy: ignore-errors + +r"""This file is allowed to initialize CUDA context when imported.""" + +import functools +import torch +import torch.cuda +from torch.testing._internal.common_utils import LazyVal, TEST_NUMBA, TEST_WITH_ROCM, TEST_CUDA, IS_WINDOWS, IS_MACOS, TEST_XPU +import inspect +import contextlib +import os +import unittest + + +CUDA_ALREADY_INITIALIZED_ON_IMPORT = torch.cuda.is_initialized() + + +TEST_MULTIGPU = TEST_CUDA and torch.cuda.device_count() >= 2 +CUDA_DEVICE = torch.device("cuda:0") if TEST_CUDA else None +# note: if ROCm is targeted, TEST_CUDNN is code for TEST_MIOPEN +if TEST_WITH_ROCM: + TEST_CUDNN = LazyVal(lambda: TEST_CUDA) +else: + TEST_CUDNN = LazyVal(lambda: TEST_CUDA and torch.backends.cudnn.is_acceptable(torch.tensor(1., device=CUDA_DEVICE))) + +TEST_CUDNN_VERSION = LazyVal(lambda: torch.backends.cudnn.version() if TEST_CUDNN else 0) +ROCM_VERSION = LazyVal(lambda : tuple(int(v) for v in torch.version.hip.split('.')[:2]) if torch.version.hip else (0, 0)) + +SM53OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (5, 3)) +SM60OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (6, 0)) +SM70OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (7, 0)) +SM75OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (7, 5)) +SM80OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 0)) +SM89OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (8, 9)) +SM90OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (9, 0)) +SM100OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)) +SM120OrLater = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() >= (12, 0)) + +IS_THOR = LazyVal(lambda: torch.cuda.is_available() and torch.version.cuda is not None and + ((torch.cuda.get_device_capability() == (11, 0) and int(torch.version.cuda[:2]) >= 13) or + (torch.cuda.get_device_capability() == (10, 1) and int(torch.version.cuda[:2]) < 13))) +IS_JETSON = LazyVal(lambda: torch.cuda.is_available() and (torch.cuda.get_device_capability() in [(7, 2), (8, 7)] or IS_THOR)) +IS_SM89 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (8, 9)) +IS_SM90 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (9, 0)) +IS_SM100 = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0)) +IS_SM12X = LazyVal(lambda: torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 12) + +@contextlib.contextmanager +def blas_library_context(backend): + prev_backend = torch.backends.cuda.preferred_blas_library() + torch.backends.cuda.preferred_blas_library(backend) + try: + yield + finally: + torch.backends.cuda.preferred_blas_library(prev_backend) + +def evaluate_gfx_arch_within(arch_list): + if not torch.cuda.is_available(): + return False + gcn_arch_name = torch.cuda.get_device_properties('cuda').gcnArchName + effective_arch = os.environ.get('PYTORCH_DEBUG_FLASH_ATTENTION_GCN_ARCH_OVERRIDE', gcn_arch_name) + # gcnArchName can be complicated strings like gfx90a:sramecc+:xnack- + # Hence the matching should be done reversely + return any(arch in effective_arch for arch in arch_list) + +def CDNA3OrLater(): + return evaluate_gfx_arch_within(["gfx942", "gfx950"]) + +def CDNA2OrLater(): + return evaluate_gfx_arch_within(["gfx90a", "gfx942"]) + +def evaluate_platform_supports_flash_attention(): + if TEST_WITH_ROCM: + arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950"] + if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0": + arch_list += ["gfx1101", "gfx1102", "gfx1150", "gfx1151", "gfx1200"] + return evaluate_gfx_arch_within(arch_list) + if TEST_CUDA: + return not IS_WINDOWS and SM80OrLater + if TEST_XPU: + return True + return False + +def evaluate_platform_supports_ck_sdpa(): + if TEST_WITH_ROCM: + return torch.backends.cuda.is_ck_sdpa_available() + else: + return False + +def evaluate_platform_supports_efficient_attention(): + if TEST_WITH_ROCM: + arch_list = ["gfx90a", "gfx942", "gfx1100", "gfx1201", "gfx950"] + if os.environ.get("TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL", "0") != "0": + arch_list += ["gfx1101", "gfx1102", "gfx1150", "gfx1151", "gfx1200"] + return evaluate_gfx_arch_within(arch_list) + if TEST_CUDA: + return True + if TEST_XPU: + return True + return False + +def evaluate_platform_supports_cudnn_attention(): + return (not TEST_WITH_ROCM) and SM80OrLater and (TEST_CUDNN_VERSION >= 90000) + +def evaluate_platform_supports_green_context(): + if IS_WINDOWS: + return False + if not _get_torch_cuda_version() >= (12, 8): + return False + driver_version = torch.utils.collect_env.get_nvidia_driver_version(torch.utils.collect_env.run) + if driver_version is None: + return False + return int(driver_version.split('.')[0]) >= 570 + +PLATFORM_SUPPORTS_FLASH_ATTENTION: bool = LazyVal(lambda: evaluate_platform_supports_flash_attention()) +PLATFORM_SUPPORTS_MEM_EFF_ATTENTION: bool = LazyVal(lambda: evaluate_platform_supports_efficient_attention()) +PLATFORM_SUPPORTS_CUDNN_ATTENTION: bool = LazyVal(lambda: evaluate_platform_supports_cudnn_attention()) +# This condition always evaluates to PLATFORM_SUPPORTS_MEM_EFF_ATTENTION but for logical clarity we keep it separate +PLATFORM_SUPPORTS_FUSED_ATTENTION: bool = LazyVal(lambda: PLATFORM_SUPPORTS_FLASH_ATTENTION or + PLATFORM_SUPPORTS_CUDNN_ATTENTION or + PLATFORM_SUPPORTS_MEM_EFF_ATTENTION) + +PLATFORM_SUPPORTS_FUSED_SDPA: bool = TEST_CUDA and not TEST_WITH_ROCM + +PLATFORM_SUPPORTS_CK_SDPA: bool = LazyVal(lambda: evaluate_platform_supports_ck_sdpa()) + + +def evaluate_platform_supports_bf16(): + if torch.version.cuda: + return SM80OrLater + elif torch.version.hip: + return True + elif TEST_XPU: + return True + return False + + +def evaluate_platform_supports_bf16_atomics(): + if torch.version.cuda: + return SM80OrLater + elif torch.version.hip: + return ROCM_VERSION >= (8, 0) + return False + + +def evaluate_platform_supports_half_atomics(): + if torch.version.hip: + return ROCM_VERSION >= (8, 0) + return True + + +PLATFORM_SUPPORTS_BF16: bool = LazyVal(lambda: evaluate_platform_supports_bf16()) +PLATFORM_SUPPORTS_BF16_ATOMICS: bool = LazyVal(lambda: evaluate_platform_supports_bf16_atomics()) +PLATFORM_SUPPORTS_HALF_ATOMICS: bool = LazyVal(lambda: evaluate_platform_supports_half_atomics()) + +PLATFORM_SUPPORTS_GREEN_CONTEXT: bool = LazyVal(lambda: evaluate_platform_supports_green_context()) + +def evaluate_platform_supports_workqueue_config(): + if IS_WINDOWS: + return False + if not _get_torch_cuda_version() >= (13, 1): + return False + driver_version = torch.utils.collect_env.get_nvidia_driver_version(torch.utils.collect_env.run) + if driver_version is None: + return False + return int(driver_version.split('.')[0]) >= 590 + +PLATFORM_SUPPORTS_WORKQUEUE_CONFIG: bool = LazyVal(lambda: evaluate_platform_supports_workqueue_config()) + +def evaluate_platform_supports_fp8(): + if torch.cuda.is_available(): + if torch.version.hip: + archs = ['gfx94'] + if ROCM_VERSION >= (6, 3): + archs.extend(['gfx120']) + if ROCM_VERSION >= (6, 5): + archs.append('gfx95') + for arch in archs: + if arch in torch.cuda.get_device_properties(0).gcnArchName: + return True + return False + else: + return SM90OrLater or torch.cuda.get_device_capability() == (8, 9) + if torch.xpu.is_available(): + return True + # As CPU supports FP8 and is always available, return True. + return True + +def evaluate_platform_supports_fp8_grouped_gemm(): + if torch.cuda.is_available(): + if torch.version.hip: + if "USE_MSLK" not in torch.__config__.show(): + return False + archs = ['gfx942', 'gfx950'] + for arch in archs: + if arch in torch.cuda.get_device_properties(0).gcnArchName: + return True + else: + return SM90OrLater and not SM100OrLater + return False + +def evaluate_platform_supports_mx_gemm(): + if torch.cuda.is_available(): + if torch.version.hip: + if ROCM_VERSION >= (7, 0): + return 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName + else: + return SM100OrLater + return False + +def evaluate_platform_supports_mxfp8_grouped_gemm(): + if torch.cuda.is_available() and not torch.version.hip: + built_with_mslk = "USE_MSLK" in torch.__config__.show() + return built_with_mslk and IS_SM100 + return False + +def evaluate_platform_supports_fp8_sparse(): + if torch.cuda.is_available(): + if torch.version.hip: + return 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName + else: + return ( + (SM90OrLater or torch.cuda.get_device_capability() == (8, 9)) + and torch.backends.cusparselt.is_available() + and torch.backends.cusparselt.version() >= 602 + ) + return False + +PLATFORM_SUPPORTS_MX_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_mx_gemm()) +PLATFORM_SUPPORTS_FP8: bool = LazyVal(lambda: evaluate_platform_supports_fp8()) +PLATFORM_SUPPORTS_FP8_SPARSE: bool = LazyVal(lambda: evaluate_platform_supports_fp8_sparse()) +PLATFORM_SUPPORTS_FP8_GROUPED_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_fp8_grouped_gemm()) +PLATFORM_SUPPORTS_MXFP8_GROUPED_GEMM: bool = LazyVal(lambda: evaluate_platform_supports_mxfp8_grouped_gemm()) + +if TEST_NUMBA: + try: + import numba.cuda + TEST_NUMBA_CUDA = numba.cuda.is_available() + except (ImportError, RuntimeError, OSError): + TEST_NUMBA_CUDA = False + TEST_NUMBA = False +else: + TEST_NUMBA_CUDA = False + +# Used below in `initialize_cuda_context_rng` to ensure that CUDA context and +# RNG have been initialized. +__cuda_ctx_rng_initialized = False + + +# after this call, CUDA context and RNG must have been initialized on each GPU +def initialize_cuda_context_rng(): + global __cuda_ctx_rng_initialized + if not TEST_CUDA: + raise AssertionError('CUDA must be available when calling initialize_cuda_context_rng') + if not __cuda_ctx_rng_initialized: + # initialize cuda context and rng for memory tests + for i in range(torch.cuda.device_count()): + torch.randn(1, device=f"cuda:{i}") + __cuda_ctx_rng_initialized = True + + +@contextlib.contextmanager +def tf32_off(): + old_allow_tf32_matmul = torch.backends.cuda.matmul.allow_tf32 + try: + torch.backends.cuda.matmul.allow_tf32 = False + with torch.backends.cudnn.flags(enabled=None, benchmark=None, deterministic=None, allow_tf32=False): + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = old_allow_tf32_matmul + + +@contextlib.contextmanager +def tf32_on(self, tf32_precision=1e-5): + old_allow_tf32_matmul = torch.backends.cuda.matmul.allow_tf32 + old_precision = self.precision + try: + torch.backends.cuda.matmul.allow_tf32 = True + self.precision = tf32_precision + with torch.backends.cudnn.flags(enabled=None, benchmark=None, deterministic=None, allow_tf32=True): + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = old_allow_tf32_matmul + self.precision = old_precision + + +@contextlib.contextmanager +def tf32_enabled(): + """ + Context manager to temporarily enable TF32 for CUDA operations. + Restores the previous TF32 state after exiting the context. + """ + old_allow_tf32_matmul = torch.backends.cuda.matmul.allow_tf32 + try: + torch.backends.cuda.matmul.allow_tf32 = True + with torch.backends.cudnn.flags( + enabled=None, benchmark=None, deterministic=None, allow_tf32=True + ): + yield + finally: + torch.backends.cuda.matmul.allow_tf32 = old_allow_tf32_matmul + + +# This is a wrapper that wraps a test to run this test twice, one with +# allow_tf32=True, another with allow_tf32=False. When running with +# allow_tf32=True, it will use reduced precision as specified by the +# argument. For example: +# @dtypes(torch.float32, torch.float64, torch.complex64, torch.complex128) +# @tf32_on_and_off(0.005) +# def test_matmul(self, device, dtype): +# a = ...; b = ...; +# c = torch.matmul(a, b) +# self.assertEqual(c, expected) +# In the above example, when testing torch.float32 and torch.complex64 on CUDA +# on a CUDA >= 11 build on an >=Ampere architecture, the matmul will be running at +# TF32 mode and TF32 mode off, and on TF32 mode, the assertEqual will use reduced +# precision to check values. +# +# This decorator can be used for function with or without device/dtype, such as +# @tf32_on_and_off(0.005) +# def test_my_op(self) +# @tf32_on_and_off(0.005) +# def test_my_op(self, device) +# @tf32_on_and_off(0.005) +# def test_my_op(self, device, dtype) +# @tf32_on_and_off(0.005) +# def test_my_op(self, dtype) +# if neither device nor dtype is specified, it will check if the system has ampere device +# if device is specified, it will check if device is cuda +# if dtype is specified, it will check if dtype is float32 or complex64 +# tf32 and fp32 are different only when all the three checks pass +def tf32_on_and_off(tf32_precision=1e-5, *, only_if=True): + def with_tf32_disabled(self, function_call): + with tf32_off(): + function_call() + + def with_tf32_enabled(self, function_call): + with tf32_on(self, tf32_precision): + function_call() + + def wrapper(f): + params = inspect.signature(f).parameters + arg_names = tuple(params.keys()) + + @functools.wraps(f) + def wrapped(*args, **kwargs): + kwargs.update(zip(arg_names, args, strict=False)) + cond = torch.cuda.is_tf32_supported() and only_if + if 'device' in kwargs: + cond = cond and (torch.device(kwargs['device']).type == 'cuda') + if 'dtype' in kwargs: + cond = cond and (kwargs['dtype'] in {torch.float32, torch.complex64}) + if cond: + with_tf32_disabled(kwargs['self'], lambda: f(**kwargs)) + with_tf32_enabled(kwargs['self'], lambda: f(**kwargs)) + else: + f(**kwargs) + + return wrapped + return wrapper + +# This is a wrapper that wraps a test to run it with TF32 turned off. +# This wrapper is designed to be used when a test uses matmul or convolutions +# but the purpose of that test is not testing matmul or convolutions. +# Disabling TF32 will enforce torch.float tensors to be always computed +# at full precision. +def with_tf32_off(f): + @functools.wraps(f) + def wrapped(*args, **kwargs): + with tf32_off(): + return f(*args, **kwargs) + + return wrapped + +def _get_magma_version(): + if 'Magma' not in torch.__config__.show(): + return (0, 0) + position = torch.__config__.show().find('Magma ') + version_str = torch.__config__.show()[position + len('Magma '):].split('\n')[0] + return tuple(int(x) for x in version_str.split(".")) + +def _get_torch_cuda_version(): + if torch.version.cuda is None: + return (0, 0) + cuda_version = str(torch.version.cuda) + return tuple(int(x) for x in cuda_version.split(".")) + +def _get_torch_rocm_version(): + if not TEST_WITH_ROCM or torch.version.hip is None: + return (0, 0) + rocm_version = str(torch.version.hip) + rocm_version = rocm_version.split("-", maxsplit=1)[0] # ignore git sha + return tuple(int(x) for x in rocm_version.split(".")) + +def _get_torch_hipblaslt_version(): + if not TEST_WITH_ROCM: + return None + try: + # Access through direct C binding + # versionHipBLASLt returns: MAJOR * 10000 + MINOR * 100 + PATCH + version_int = torch._C._cuda_getHipblasltVersion() + if version_int is None or version_int == 0: + return None + major = version_int // 10000 + minor = (version_int % 10000) // 100 + patch = version_int % 100 + return (major, minor, patch) + except (AttributeError, RuntimeError): + return None + +def _check_cusparse_generic_available(): + return not TEST_WITH_ROCM + +def _check_hipsparse_generic_available(): + if not TEST_WITH_ROCM: + return False + if not torch.version.hip: + return False + + rocm_version = str(torch.version.hip) + rocm_version = rocm_version.split("-", maxsplit=1)[0] # ignore git sha + rocm_version_tuple = tuple(int(x) for x in rocm_version.split(".")) + return not (rocm_version_tuple is None or rocm_version_tuple < (5, 1)) + + +TEST_CUSPARSE_GENERIC = _check_cusparse_generic_available() +TEST_HIPSPARSE_GENERIC = _check_hipsparse_generic_available() + +# Shared by test_torch.py and test_multigpu.py +def _create_scaling_models_optimizers(device="cuda", optimizer_ctor=torch.optim.SGD, optimizer_kwargs=None): + # Create a module+optimizer that will use scaling, and a control module+optimizer + # that will not use scaling, against which the scaling-enabled module+optimizer can be compared. + mod_control = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)).to(device=device) + mod_scaling = torch.nn.Sequential(torch.nn.Linear(8, 8), torch.nn.Linear(8, 8)).to(device=device) + with torch.no_grad(): + for c, s in zip(mod_control.parameters(), mod_scaling.parameters(), strict=True): + s.copy_(c) + + kwargs = {"lr": 1.0} + if optimizer_kwargs is not None: + kwargs.update(optimizer_kwargs) + opt_control = optimizer_ctor(mod_control.parameters(), **kwargs) + opt_scaling = optimizer_ctor(mod_scaling.parameters(), **kwargs) + + return mod_control, mod_scaling, opt_control, opt_scaling + +# Shared by test_torch.py, test_cuda.py and test_multigpu.py +def _create_scaling_case(device="cuda", dtype=torch.float, optimizer_ctor=torch.optim.SGD, optimizer_kwargs=None): + data = [(torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)), + (torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)), + (torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device)), + (torch.randn((8, 8), dtype=dtype, device=device), torch.randn((8, 8), dtype=dtype, device=device))] + + loss_fn = torch.nn.MSELoss().to(device) + + skip_iter = 2 + + return _create_scaling_models_optimizers( + device=device, optimizer_ctor=optimizer_ctor, optimizer_kwargs=optimizer_kwargs, + ) + (data, loss_fn, skip_iter) + + +def xfailIfSM89(func): + return func if not IS_SM89 else unittest.expectedFailure(func) + +def xfailIfSM90(func): + return func if not IS_SM90 else unittest.expectedFailure(func) + +def xfailIfSM89PreCUDA13(func): + """xfail on SM89 only for CUDA < 13. On CUDA 13+, test should pass on all architectures.""" + if IS_SM89 and _get_torch_cuda_version() < (13, 0): + return unittest.expectedFailure(func) + return func + +def xfailIfSM100OrLater(func): + return func if not SM100OrLater else unittest.expectedFailure(func) + +def xfailIfSM120OrLater(func): + return func if not SM120OrLater else unittest.expectedFailure(func) + +def xfailIfSM12X(func): + return func if not IS_SM12X else unittest.expectedFailure(func) + +def xfailIfDistributedNotSupported(func): + return func if not (IS_MACOS or IS_JETSON) else unittest.expectedFailure(func) + +# When using nvcc from the CUDA toolkit its versuib must be at least the one from ptxas bundled with Triton +TRITON_PTXAS_VERSION = (12, 8) +requires_triton_ptxas_compat = unittest.skipIf(not torch.version.xpu + and torch.version.hip is None + and _get_torch_cuda_version() < TRITON_PTXAS_VERSION, + "Requires CUDA {}.{} to match Tritons ptxas version".format(*TRITON_PTXAS_VERSION)) + +# Importing this module should NOT eagerly initialize CUDA +if not CUDA_ALREADY_INITIALIZED_ON_IMPORT: + if torch.cuda.is_initialized(): + raise AssertionError("CUDA should not be initialized on import") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_device_type.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_device_type.py new file mode 100644 index 0000000000000000000000000000000000000000..8fc9f2cd9682fb6f0f267fdc0e28fe0f8b173130 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_device_type.py @@ -0,0 +1,2116 @@ +# mypy: ignore-errors + +import copy +import gc +import inspect +import os +import runpy +import sys +import threading +import unittest +from collections import namedtuple +from collections.abc import Callable, Iterable, Sequence +from enum import Enum +from functools import partial, wraps +from typing import Any, ClassVar, TypeVar +from typing_extensions import ParamSpec + +import torch +from torch._inductor.utils import GPU_TYPES +from torch._utils import _is_privateuse1_backend_available +from torch.testing._internal.common_cuda import ( + _get_torch_cuda_version, + _get_torch_hipblaslt_version, + _get_torch_rocm_version, + TEST_CUSPARSE_GENERIC, + TEST_HIPSPARSE_GENERIC, +) +from torch.testing._internal.common_dtype import get_all_dtypes +from torch.testing._internal.common_utils import ( + _TestParametrizer, + clear_tracked_input, + compose_parametrize_fns, + dtype_name, + get_tracked_input, + IS_FBCODE, + IS_MACOS, + IS_REMOTE_GPU, + IS_S390X, + IS_SANDCASTLE, + IS_WINDOWS, + NATIVE_DEVICES, + PRINT_REPRO_ON_FAILURE, + skipCUDANonDefaultStreamIf, + skipIfTorchDynamo, + TEST_HPU, + TEST_MKL, + TEST_MPS, + TEST_WITH_ASAN, + TEST_WITH_MIOPEN_SUGGEST_NHWC, + TEST_WITH_MTIA, + TEST_WITH_ROCM, + TEST_WITH_TORCHINDUCTOR, + TEST_WITH_TSAN, + TEST_WITH_UBSAN, + TEST_XPU, + TestCase, +) + + +_T = TypeVar("_T") +_P = ParamSpec("_P") + +try: + import psutil # type: ignore[import] + + HAS_PSUTIL = True +except ModuleNotFoundError: + HAS_PSUTIL = False + psutil = None + +# Note [Writing Test Templates] +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# This note was written shortly after the PyTorch 1.9 release. +# If you notice it's out-of-date or think it could be improved then please +# file an issue. +# +# PyTorch has its own framework for instantiating test templates. That is, for +# taking test classes that look similar to unittest or pytest +# compatible test classes and optionally doing the following: +# +# - instantiating a version of the test class for each available device type +# (often the CPU, CUDA, and META device types) +# - further instantiating a version of each test that's always specialized +# on the test class's device type, and optionally specialized further +# on datatypes or operators +# +# This functionality is similar to pytest's parametrize functionality +# (see https://docs.pytest.org/en/6.2.x/parametrize.html), but with considerable +# additional logic that specializes the instantiated test classes for their +# device types (see CPUTestBase and CUDATestBase below), supports a variety +# of composable decorators that allow for test filtering and setting +# tolerances, and allows tests parametrized by operators to instantiate +# only the subset of device type x dtype that operator supports. +# +# This framework was built to make it easier to write tests that run on +# multiple device types, multiple datatypes (dtypes), and for multiple +# operators. It's also useful for controlling which tests are run. For example, +# only tests that use a CUDA device can be run on platforms with CUDA. +# Let's dive in with an example to get an idea for how it works: +# +# -------------------------------------------------------- +# A template class (looks like a regular unittest TestCase) +# class TestClassFoo(TestCase): +# +# # A template test that can be specialized with a device +# # NOTE: this test case is not runnable by unittest or pytest because it +# # accepts an extra positional argument, "device", that they do not understand +# def test_bar(self, device): +# pass +# +# # Function that instantiates a template class and its tests +# instantiate_device_type_tests(TestCommon, globals()) +# -------------------------------------------------------- +# +# In the above code example we see a template class and a single test template +# that can be instantiated with a device. The function +# instantiate_device_type_tests(), called at file scope, instantiates +# new test classes, one per available device type, and new tests in those +# classes from these templates. It actually does this by removing +# the class TestClassFoo and replacing it with classes like TestClassFooCPU +# and TestClassFooCUDA, instantiated test classes that inherit from CPUTestBase +# and CUDATestBase respectively. Additional device types, like XLA, +# (see https://github.com/pytorch/xla) can further extend the set of +# instantiated test classes to create classes like TestClassFooXLA. +# +# The test template, test_bar(), is also instantiated. In this case the template +# is only specialized on a device, so (depending on the available device +# types) it might become test_bar_cpu() in TestClassFooCPU and test_bar_cuda() +# in TestClassFooCUDA. We can think of the instantiated test classes as +# looking like this: +# +# -------------------------------------------------------- +# # An instantiated test class for the CPU device type +# class TestClassFooCPU(CPUTestBase): +# +# # An instantiated test that calls the template with the string representation +# # of a device from the test class's device type +# def test_bar_cpu(self): +# test_bar(self, 'cpu') +# +# # An instantiated test class for the CUDA device type +# class TestClassFooCUDA(CUDATestBase): +# +# # An instantiated test that calls the template with the string representation +# # of a device from the test class's device type +# def test_bar_cuda(self): +# test_bar(self, 'cuda:0') +# -------------------------------------------------------- +# +# These instantiated test classes ARE discoverable and runnable by both +# unittest and pytest. One thing that may be confusing, however, is that +# attempting to run "test_bar" will not work, despite it appearing in the +# original template code. This is because "test_bar" is no longer discoverable +# after instantiate_device_type_tests() runs, as the above snippet shows. +# Instead "test_bar_cpu" and "test_bar_cuda" may be run directly, or both +# can be run with the option "-k test_bar". +# +# Removing the template class and adding the instantiated classes requires +# passing "globals()" to instantiate_device_type_tests(), because it +# edits the file's Python objects. +# +# As mentioned, tests can be additionally parametrized on dtypes or +# operators. Datatype parametrization uses the @dtypes decorator and +# require a test template like this: +# +# -------------------------------------------------------- +# # A template test that can be specialized with a device and a datatype (dtype) +# @dtypes(torch.float32, torch.int64) +# def test_car(self, device, dtype) +# pass +# -------------------------------------------------------- +# +# If the CPU and CUDA device types are available this test would be +# instantiated as 4 tests that cover the cross-product of the two dtypes +# and two device types: +# +# - test_car_cpu_float32 +# - test_car_cpu_int64 +# - test_car_cuda_float32 +# - test_car_cuda_int64 +# +# The dtype is passed as a torch.dtype object. +# +# Tests parametrized on operators (actually on OpInfos, more on that in a +# moment...) use the @ops decorator and require a test template like this: +# -------------------------------------------------------- +# # A template test that can be specialized with a device, dtype, and OpInfo +# @ops(op_db) +# def test_car(self, device, dtype, op) +# pass +# -------------------------------------------------------- +# +# See the documentation for the @ops decorator below for additional details +# on how to use it and see the note [OpInfos] in +# common_methods_invocations.py for more details on OpInfos. +# +# A test parametrized over the entire "op_db", which contains hundreds of +# OpInfos, will likely have hundreds or thousands of instantiations. The +# test will be instantiated on the cross-product of device types, operators, +# and the dtypes the operator supports on that device type. The instantiated +# tests will have names like: +# +# - test_car_add_cpu_float32 +# - test_car_sub_cuda_int64 +# +# The first instantiated test calls the original test_car() with the OpInfo +# for torch.add as its "op" argument, the string 'cpu' for its "device" argument, +# and the dtype torch.float32 for is "dtype" argument. The second instantiated +# test calls the test_car() with the OpInfo for torch.sub, a CUDA device string +# like 'cuda:0' or 'cuda:1' for its "device" argument, and the dtype +# torch.int64 for its "dtype argument." +# +# In addition to parametrizing over device, dtype, and ops via OpInfos, the +# @parametrize decorator is supported for arbitrary parametrizations: +# -------------------------------------------------------- +# # A template test that can be specialized with a device, dtype, and value for x +# @parametrize("x", range(5)) +# def test_car(self, device, dtype, x) +# pass +# -------------------------------------------------------- +# +# See the documentation for @parametrize in common_utils.py for additional details +# on this. Note that the instantiate_device_type_tests() function will handle +# such parametrizations; there is no need to additionally call +# instantiate_parametrized_tests(). +# +# Clever test filtering can be very useful when working with parametrized +# tests. "-k test_car" would run every instantiated variant of the test_car() +# test template, and "-k test_car_add" runs every variant instantiated with +# torch.add. +# +# It is important to use the passed device and dtype as appropriate. Use +# helper functions like make_tensor() that require explicitly specifying +# the device and dtype so they're not forgotten. +# +# Test templates can use a variety of composable decorators to specify +# additional options and requirements, some are listed here: +# +# - @deviceCountAtLeast() +# Passes a list of strings representing all available devices of +# the test class's device type as the test template's "device" argument. +# If there are fewer devices than the value passed to the decorator +# the test is skipped. +# - @dtypes() +# In addition to accepting multiple dtypes, the @dtypes decorator +# can accept a sequence of tuple pairs of dtypes. The test template +# will be called with each tuple for its "dtype" argument. +# - @onlyNativeDeviceTypes +# Skips the test if the device is not a native device type (currently CPU, CUDA, Meta) +# - @onlyCPU +# Skips the test if the device is not a CPU device +# - @onlyCUDA +# Skips the test if the device is not a CUDA device +# - @onlyMPS +# Skips the test if the device is not a MPS device +# - @skipCPUIfNoLapack +# Skips the test if the device is a CPU device and LAPACK is not installed +# - @skipCPUIfNoMkl +# Skips the test if the device is a CPU device and MKL is not installed +# - @skipCUDAIfNoMagma +# Skips the test if the device is a CUDA device and MAGMA is not installed +# - @skipCUDAIfRocm +# Skips the test if the device is a CUDA device and ROCm is being used + + +# Note [Adding a Device Type] +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# To add a device type: +# +# (1) Create a new "TestBase" extending DeviceTypeTestBase. +# See CPUTestBase and CUDATestBase below. +# (2) Define the "device_type" attribute of the base to be the +# appropriate string. +# (3) Add logic to this file that appends your base class to +# device_type_test_bases when your device type is available. +# (4) (Optional) Write setUpClass/tearDownClass class methods that +# instantiate dependencies (see MAGMA in CUDATestBase). +# (5) (Optional) Override the "instantiate_test" method for total +# control over how your class creates tests. +# +# setUpClass is called AFTER tests have been created and BEFORE and ONLY IF +# they are run. This makes it useful for initializing devices and dependencies. + + +def _dtype_test_suffix(dtypes): + """Returns the test suffix for a dtype, sequence of dtypes, or None.""" + if isinstance(dtypes, (list, tuple)): + if len(dtypes) == 0: + return "" + return "_" + "_".join(dtype_name(d) for d in dtypes) + elif dtypes: + return f"_{dtype_name(dtypes)}" + else: + return "" + + +def _update_param_kwargs(param_kwargs, name, value): + """Adds a kwarg with the specified name and value to the param_kwargs dict.""" + # Make name plural (e.g. devices / dtypes) if the value is composite. + plural_name = f"{name}s" + + # Clear out old entries of the arg if any. + if name in param_kwargs: + del param_kwargs[name] + if plural_name in param_kwargs: + del param_kwargs[plural_name] + + if isinstance(value, (list, tuple)): + param_kwargs[plural_name] = value + elif value is not None: + param_kwargs[name] = value + + # Leave param_kwargs as-is when value is None. + + +class DeviceTypeTestBase(TestCase): + device_type: str = "generic_device_type" + + # When True, @onlyOn-based decorators (@onlyCUDA, @onlyMPS, etc.) will not + # skip tests for this device type. This is a pragmatic short-term solution to + # allow PrivateUse1 backends to run tests that are currently gated behind + # device-specific decorators. It is intended to be used together with the + # skip mechanism (see https://github.com/pytorch/pytorch/issues/177253). + # In the longer term, we are incrementally migrating accelerator tests to be + # device-generic and removing @onlyCUDA on tests that should be device-generic. + bypass_device_restrictions: bool = False + + # Flag to disable test suite early due to unrecoverable error such as CUDA error. + _stop_test_suite = False + + # Precision is a thread-local setting since it may be overridden per test + _tls = threading.local() + _tls.precision = TestCase._precision + _tls.rel_tol = TestCase._rel_tol + + @property + def precision(self): + return self._tls.precision + + @precision.setter + def precision(self, prec): + self._tls.precision = prec + + @property + def rel_tol(self): + return self._tls.rel_tol + + @rel_tol.setter + def rel_tol(self, prec): + self._tls.rel_tol = prec + + # Returns a string representing the device that single device tests should use. + # Note: single device tests use this device exclusively. + @classmethod + def get_primary_device(cls): + return cls.device_type + + @classmethod + def _init_and_get_primary_device(cls): + try: + return cls.get_primary_device() + except Exception: + # For CUDATestBase, XPUTestBase, XLATestBase, and possibly others, the primary device won't be available + # until setUpClass() sets it. Call that manually here if needed. + if hasattr(cls, "setUpClass"): + cls.setUpClass() + return cls.get_primary_device() + + # Returns a list of strings representing all available devices of this + # device type. The primary device must be the first string in the list + # and the list must contain no duplicates. + # Note: UNSTABLE API. Will be replaced once PyTorch has a device generic + # mechanism of acquiring all available devices. + @classmethod + def get_all_devices(cls): + return [cls.get_primary_device()] + + # Returns the dtypes the test has requested. + # Prefers device-specific dtype specifications over generic ones. + @classmethod + def _get_dtypes(cls, test): + if not hasattr(test, "dtypes"): + return None + + default_dtypes = test.dtypes.get("all") + msg = f"@dtypes is mandatory when using @dtypesIf however '{test.__name__}' didn't specify it" + if default_dtypes is None: + raise AssertionError(msg) + + return test.dtypes.get(cls.device_type, default_dtypes) + + def _get_precision_override(self, test, dtype): + if not hasattr(test, "precision_overrides"): + return self.precision + return test.precision_overrides.get(dtype, self.precision) + + def _get_tolerance_override(self, test, dtype): + if not hasattr(test, "tolerance_overrides"): + return self.precision, self.rel_tol + return test.tolerance_overrides.get(dtype, tol(self.precision, self.rel_tol)) + + def _apply_precision_override_for_test(self, test, param_kwargs): + dtype = param_kwargs.get("dtype") + dtype = param_kwargs.get("dtypes", dtype) + if dtype: + self.precision = self._get_precision_override(test, dtype) + self.precision, self.rel_tol = self._get_tolerance_override(test, dtype) + + # Creates device-specific tests. + @classmethod + def instantiate_test(cls, name, test, *, generic_cls=None): + def instantiate_test_helper( + cls, name, *, test, param_kwargs=None, decorator_fn=lambda _: [] + ): + # Add the device param kwarg if the test needs device or devices. + param_kwargs = {} if param_kwargs is None else param_kwargs + test_sig_params = inspect.signature(test).parameters + if "device" in test_sig_params or "devices" in test_sig_params: + device_arg: str = cls._init_and_get_primary_device() + if hasattr(test, "num_required_devices"): + device_arg = cls.get_all_devices() + _update_param_kwargs(param_kwargs, "device", device_arg) + + # Apply decorators based on param kwargs. + for decorator in decorator_fn(param_kwargs): + test = decorator(test) + + # Constructs the test + @wraps(test) + def instantiated_test(self, param_kwargs=param_kwargs): + # Sets precision and runs test + # Note: precision is reset after the test is run + guard_precision = self.precision + guard_rel_tol = self.rel_tol + try: + self._apply_precision_override_for_test(test, param_kwargs) + result = test(self, **param_kwargs) + except RuntimeError as rte: + # check if rte should stop entire test suite. + self._stop_test_suite = self._should_stop_test_suite() + # Check if test has been decorated with `@expectedFailure` + # Using `__unittest_expecting_failure__` attribute, see + # https://github.com/python/cpython/blob/ffa505b580464/Lib/unittest/case.py#L164 + # In that case, make it fail with "unexpected success" by suppressing exception + if ( + getattr(test, "__unittest_expecting_failure__", False) + and self._stop_test_suite + ): + import sys + + print( + "Suppressing fatal exception to trigger unexpected success", + file=sys.stderr, + ) + return + # raise the runtime error as is for the test suite to record. + raise rte + finally: + self.precision = guard_precision + self.rel_tol = guard_rel_tol + + return result + + if hasattr(cls, name): + raise AssertionError(f"Redefinition of test {name}") + setattr(cls, name, instantiated_test) + + def default_parametrize_fn(test, generic_cls, device_cls): + # By default, no parametrization is needed. + yield (test, "", {}, lambda _: []) + + # Parametrization decorators set the parametrize_fn attribute on the test. + parametrize_fn = getattr(test, "parametrize_fn", default_parametrize_fn) + + # If one of the @dtypes* decorators is present, also parametrize over the dtypes set by it. + dtypes = cls._get_dtypes(test) + if dtypes is not None: + + def dtype_parametrize_fn(test, generic_cls, device_cls, dtypes=dtypes): + for dtype in dtypes: + param_kwargs: dict[str, Any] = {} + _update_param_kwargs(param_kwargs, "dtype", dtype) + + # Note that an empty test suffix is set here so that the dtype can be appended + # later after the device. + yield (test, "", param_kwargs, lambda _: []) + + parametrize_fn = compose_parametrize_fns( + dtype_parametrize_fn, parametrize_fn + ) + + # Instantiate the parametrized tests. + for ( + test, # noqa: B020 + test_suffix, + param_kwargs, + decorator_fn, + ) in parametrize_fn(test, generic_cls, cls): + test_suffix = "" if test_suffix == "" else "_" + test_suffix + cls_device_type = ( + cls.device_type + if cls.device_type != "privateuse1" + else torch._C._get_privateuse1_backend_name() + ) + device_suffix = "_" + cls_device_type + + # Note: device and dtype suffix placement + # Special handling here to place dtype(s) after device according to test name convention. + dtype_kwarg = None + if "dtype" in param_kwargs or "dtypes" in param_kwargs: + dtype_kwarg = ( + param_kwargs["dtypes"] + if "dtypes" in param_kwargs + else param_kwargs["dtype"] + ) + test_name = ( + f"{name}{test_suffix}{device_suffix}{_dtype_test_suffix(dtype_kwarg)}" + ) + + instantiate_test_helper( + cls=cls, + name=test_name, + test=test, + param_kwargs=param_kwargs, + decorator_fn=decorator_fn, + ) + + def run(self, result=None): + super().run(result=result) + # Early terminate test if _stop_test_suite is set. + if self._stop_test_suite: + result.stop() + + +class CPUTestBase(DeviceTypeTestBase): + device_type = "cpu" + + # No critical error should stop CPU test suite + def _should_stop_test_suite(self): + return False + + +class CUDATestBase(DeviceTypeTestBase): + device_type = "cuda" + _do_cuda_memory_leak_check = True + _do_cuda_non_default_stream = True + primary_device: ClassVar[str] + cudnn_version: ClassVar[Any] + no_magma: ClassVar[bool] + no_cudnn: ClassVar[bool] + + def has_cudnn(self): + return not self.no_cudnn + + @classmethod + def get_primary_device(cls): + return cls.primary_device + + @classmethod + def get_all_devices(cls): + primary_device_idx = int(cls.get_primary_device().split(":")[1]) + num_devices = torch.cuda.device_count() + + prim_device = cls.get_primary_device() + cuda_str = "cuda:{0}" + non_primary_devices = [ + cuda_str.format(idx) + for idx in range(num_devices) + if idx != primary_device_idx + ] + return [prim_device] + non_primary_devices + + @classmethod + def setUpClass(cls): + # has_magma shows up after cuda is initialized + t = torch.ones(1).cuda() + cls.no_magma = not torch.cuda.has_magma + + # Determines if cuDNN is available and its version + cls.no_cudnn = not torch.backends.cudnn.is_acceptable(t) + cls.cudnn_version = None if cls.no_cudnn else torch.backends.cudnn.version() + + # Acquires the current device as the primary (test) device + cls.primary_device = f"cuda:{torch.cuda.current_device()}" + + +# See Note [Lazy Tensor tests in device agnostic testing] +lazy_ts_backend_init = False + + +class LazyTestBase(DeviceTypeTestBase): + device_type = "lazy" + + def _should_stop_test_suite(self): + return False + + @classmethod + def setUpClass(cls): + import torch._lazy + import torch._lazy.metrics + import torch._lazy.ts_backend + + global lazy_ts_backend_init + if not lazy_ts_backend_init: + # Need to connect the TS backend to lazy key before running tests + torch._lazy.ts_backend.init() + lazy_ts_backend_init = True + + +class MPSTestBase(DeviceTypeTestBase): + device_type = "mps" + primary_device: ClassVar[str] + + @classmethod + def get_primary_device(cls): + return cls.primary_device + + @classmethod + def get_all_devices(cls): + # currently only one device is supported on MPS backend + prim_device = cls.get_primary_device() + return [prim_device] + + @classmethod + def setUpClass(cls): + cls.primary_device = "mps:0" + + def _should_stop_test_suite(self): + return False + + +class XPUTestBase(DeviceTypeTestBase): + device_type = "xpu" + primary_device: ClassVar[str] + + @classmethod + def get_primary_device(cls): + return cls.primary_device + + @classmethod + def get_all_devices(cls): + # currently only one device is supported on MPS backend + primary_device_idx = int(cls.get_primary_device().split(":")[1]) + num_devices = torch.xpu.device_count() + + prim_device = cls.get_primary_device() + xpu_str = "xpu:{0}" + non_primary_devices = [ + xpu_str.format(idx) + for idx in range(num_devices) + if idx != primary_device_idx + ] + return [prim_device] + non_primary_devices + + @classmethod + def setUpClass(cls): + cls.primary_device = f"xpu:{torch.xpu.current_device()}" + + def _should_stop_test_suite(self): + return False + + +class HPUTestBase(DeviceTypeTestBase): + device_type = "hpu" + primary_device: ClassVar[str] + + @classmethod + def get_primary_device(cls): + return cls.primary_device + + @classmethod + def setUpClass(cls): + cls.primary_device = "hpu:0" + + +class PrivateUse1TestBase(DeviceTypeTestBase): + primary_device: ClassVar[str] + device_mod = None + device_type = "privateuse1" + bypass_device_restrictions = False + + @classmethod + def get_primary_device(cls): + return cls.primary_device + + @classmethod + def get_all_devices(cls): + primary_device_idx = int(cls.get_primary_device().split(":")[1]) + num_devices = cls.device_mod.device_count() + prim_device = cls.get_primary_device() + device_str = f"{cls.device_type}:{{0}}" + non_primary_devices = [ + device_str.format(idx) + for idx in range(num_devices) + if idx != primary_device_idx + ] + return [prim_device] + non_primary_devices + + @classmethod + def setUpClass(cls): + cls.device_type = torch._C._get_privateuse1_backend_name() + cls.device_mod = getattr(torch, cls.device_type, None) + if cls.device_mod is None: + raise AssertionError( + f"torch has no module of `{cls.device_type}`, you should register " + "a module by `torch._register_device_module`." + ) + cls.primary_device = f"{cls.device_type}:{cls.device_mod.current_device()}" + + +# Adds available device-type-specific test base classes +def get_device_type_test_bases(): + # set type to List[Any] due to mypy list-of-union issue: + # https://github.com/python/mypy/issues/3351 + test_bases: list[Any] = [] + + if IS_SANDCASTLE or IS_FBCODE: + if IS_REMOTE_GPU: + # Skip if sanitizer is enabled or we're on MTIA machines + if ( + not TEST_WITH_ASAN + and not TEST_WITH_TSAN + and not TEST_WITH_UBSAN + and not TEST_WITH_MTIA + ): + test_bases.append(CUDATestBase) + else: + test_bases.append(CPUTestBase) + else: + test_bases.append(CPUTestBase) + if torch.cuda.is_available(): + test_bases.append(CUDATestBase) + + if _is_privateuse1_backend_available(): + test_bases.append(PrivateUse1TestBase) + # Disable MPS testing in generic device testing temporarily while we're + # ramping up support. + # elif torch.backends.mps.is_available(): + # test_bases.append(MPSTestBase) + + return test_bases + + +device_type_test_bases = get_device_type_test_bases() + + +def filter_desired_device_types(device_type_test_bases, except_for=None, only_for=None): + # device type cannot appear in both except_for and only_for + intersect = set(except_for if except_for else []) & set( + only_for if only_for else [] + ) + if intersect: + raise AssertionError( + f"device ({intersect}) appeared in both except_for and only_for" + ) + + # Replace your privateuse1 backend name with 'privateuse1' + # This handles the case where PrivateUse1TestBase.device_type has been + # changed from "privateuse1" to the actual backend name (e.g., "openreg") + # by setUpClass being called during previous instantiate_device_type_tests calls + if _is_privateuse1_backend_available(): + privateuse1_backend_name = torch._C._get_privateuse1_backend_name() + + def func_replace(x: str) -> str: + return x.replace(privateuse1_backend_name, "privateuse1") + + except_for = ( + ([func_replace(x) for x in except_for] if except_for is not None else None) + if not isinstance(except_for, str) + else func_replace(except_for) + ) + only_for = ( + ([func_replace(x) for x in only_for] if only_for is not None else None) + if not isinstance(only_for, str) + else func_replace(only_for) + ) + else: + + def func_replace(x: str) -> str: + return x + + if except_for: + device_type_test_bases = filter( + lambda x: func_replace(x.device_type) not in except_for, + device_type_test_bases, + ) + if only_for: + device_type_test_bases = filter( + lambda x: func_replace(x.device_type) in only_for, + device_type_test_bases, + ) + + return list(device_type_test_bases) + + +# Note [How to extend DeviceTypeTestBase to add new test device] +# The following logic optionally allows downstream projects like pytorch/xla to +# add more test devices. +# Instructions: +# - Add a python file (e.g. pytorch/xla/test/pytorch_test_base.py) in downstream project. +# - Inside the file, one should inherit from `DeviceTypeTestBase` class and define +# a new DeviceTypeTest class (e.g. `XLATestBase`) with proper implementation of +# `instantiate_test` method. +# - DO NOT import common_device_type inside the file. +# `runpy.run_path` with `globals()` already properly setup the context so that +# `DeviceTypeTestBase` is already available. +# - Set a top-level variable `TEST_CLASS` equal to your new class. +# E.g. TEST_CLASS = XLATensorBase +# - To run tests with new device type, set `TORCH_TEST_DEVICE` env variable to path +# to this file. Multiple paths can be separated by `:`. +# See pytorch/xla/test/pytorch_test_base.py for a more detailed example. +_TORCH_TEST_DEVICES = os.environ.get("TORCH_TEST_DEVICES", None) +if _TORCH_TEST_DEVICES: + for path in _TORCH_TEST_DEVICES.split(":"): + # runpy (a stdlib module) lacks annotations + mod = runpy.run_path(path, init_globals=globals()) # type: ignore[func-returns-value] + device_type_test_bases.append(mod["TEST_CLASS"]) + + +PYTORCH_CUDA_MEMCHECK = os.getenv("PYTORCH_CUDA_MEMCHECK", "0") == "1" + +PYTORCH_TESTING_DEVICE_ONLY_FOR_KEY = "PYTORCH_TESTING_DEVICE_ONLY_FOR" +PYTORCH_TESTING_DEVICE_EXCEPT_FOR_KEY = "PYTORCH_TESTING_DEVICE_EXCEPT_FOR" +PYTORCH_TESTING_DEVICE_FOR_CUSTOM_KEY = "PYTORCH_TESTING_DEVICE_FOR_CUSTOM" + + +def get_desired_device_type_test_bases( + except_for=None, only_for=None, include_lazy=False, allow_mps=False, allow_xpu=False +): + # allow callers to specifically opt tests into being tested on MPS, similar to `include_lazy` + test_bases = device_type_test_bases.copy() + if allow_mps and TEST_MPS and MPSTestBase not in test_bases: + test_bases.append(MPSTestBase) + if allow_xpu and TEST_XPU and XPUTestBase not in test_bases: + test_bases.append(XPUTestBase) + if TEST_HPU and HPUTestBase not in test_bases: + test_bases.append(HPUTestBase) + # Filter out the device types based on user inputs + desired_device_type_test_bases = filter_desired_device_types( + test_bases, except_for, only_for + ) + if include_lazy: + # Note [Lazy Tensor tests in device agnostic testing] + # Right now, test_view_ops.py runs with LazyTensor. + # We don't want to opt every device-agnostic test into using the lazy device, + # because many of them will fail. + # So instead, the only way to opt a specific device-agnostic test file into + # lazy tensor testing is with include_lazy=True + if IS_FBCODE: + print( + "TorchScript backend not yet supported in FBCODE/OVRSOURCE builds", + file=sys.stderr, + ) + else: + desired_device_type_test_bases.append(LazyTestBase) + + def split_if_not_empty(x: str): + return x.split(",") if x else [] + + # run some cuda testcases on other devices if available + # Usage: + # export PYTORCH_TESTING_DEVICE_FOR_CUSTOM=privateuse1 + env_custom_only_for = split_if_not_empty( + os.getenv(PYTORCH_TESTING_DEVICE_FOR_CUSTOM_KEY, "") + ) + if env_custom_only_for: + desired_device_type_test_bases += filter( + lambda x: x.device_type in env_custom_only_for, test_bases + ) + desired_device_type_test_bases = list(set(desired_device_type_test_bases)) + + # Filter out the device types based on environment variables if available + # Usage: + # export PYTORCH_TESTING_DEVICE_ONLY_FOR=cuda,cpu + # export PYTORCH_TESTING_DEVICE_EXCEPT_FOR=xla + env_only_for = split_if_not_empty( + os.getenv(PYTORCH_TESTING_DEVICE_ONLY_FOR_KEY, "") + ) + env_except_for = split_if_not_empty( + os.getenv(PYTORCH_TESTING_DEVICE_EXCEPT_FOR_KEY, "") + ) + + return filter_desired_device_types( + desired_device_type_test_bases, env_except_for, env_only_for + ) + + +# Adds 'instantiated' device-specific test cases to the given scope. +# The tests in these test cases are derived from the generic tests in +# generic_test_class. This function should be used instead of +# instantiate_parametrized_tests() if the test class contains +# device-specific tests (NB: this supports additional @parametrize usage). +# +# See note "Writing Test Templates" +# TODO: remove "allow_xpu" option after Interl GPU support all test case instantiate by this function. +def instantiate_device_type_tests( + generic_test_class, + scope, + except_for=None, + only_for=None, + include_lazy=False, + allow_mps=False, + allow_xpu=False, +): + # Removes the generic test class from its enclosing scope so its tests + # are not discoverable. + del scope[generic_test_class.__name__] + + generic_members = set(generic_test_class.__dict__.keys()) + generic_tests = [x for x in generic_members if x.startswith("test")] + + # Creates device-specific test cases + for base in get_desired_device_type_test_bases( + except_for, only_for, include_lazy, allow_mps, allow_xpu + ): + class_name = generic_test_class.__name__ + base.device_type.upper() + + # type set to Any and suppressed due to unsupported runtime class: + # https://github.com/python/mypy/wiki/Unsupported-Python-Features + device_type_test_class: Any = type(class_name, (base, generic_test_class), {}) + + # Arrange for setUpClass and tearDownClass methods defined both in the test template + # class and in the generic base to be called. This allows device-parameterized test + # classes to support setup and teardown. + # NB: This should be done before instantiate_test() is called as that invokes setup. + @classmethod + def _setUpClass(cls): + # This should always be called, whether or not the test class invokes + # super().setUpClass(), to set the primary device. + base.setUpClass() + # We want to call the @classmethod defined in the generic base, but pass + # it the device-specific class object (cls), hence the __func__ call. + generic_test_class.setUpClass.__func__(cls) + + @classmethod + def _tearDownClass(cls): + # We want to call the @classmethod defined in the generic base, but pass + # it the device-specific class object (cls), hence the __func__ call. + generic_test_class.tearDownClass.__func__(cls) + base.tearDownClass() + + device_type_test_class.setUpClass = _setUpClass + device_type_test_class.tearDownClass = _tearDownClass + + for name in generic_members: + if name in generic_tests: # Instantiates test member + test = getattr(generic_test_class, name) + # XLA-compat shim (XLA's instantiate_test takes doesn't take generic_cls) + sig = inspect.signature(device_type_test_class.instantiate_test) + if len(sig.parameters) == 3: + # Instantiates the device-specific tests + device_type_test_class.instantiate_test( + name, copy.deepcopy(test), generic_cls=generic_test_class + ) + else: + device_type_test_class.instantiate_test(name, copy.deepcopy(test)) + # Ports non-test member. Setup / teardown have already been handled above + elif name not in device_type_test_class.__dict__: + nontest = getattr(generic_test_class, name) + setattr(device_type_test_class, name, nontest) + + # Mimics defining the instantiated class in the caller's file + # by setting its module to the given class's and adding + # the module to the given scope. + # This lets the instantiated class be discovered by unittest. + device_type_test_class.__module__ = generic_test_class.__module__ + scope[class_name] = device_type_test_class + + # Delete the generic form of the test functions (e.g. TestFoo.test_bar()) so they're + # not discoverable. This mutates the original class (TestFoo), which was removed from + # scope above. At this point, device-specific tests (e.g. TestFooCUDA.test_bar_cuda) + # have already been created and the generic forms are no longer needed. + for name in generic_tests: + delattr(generic_test_class, name) + + +# Category of dtypes to run an OpInfo-based test for +# Example use: @ops(dtype=OpDTypes.supported) +# +# There are 7 categories: +# - supported: Every dtype supported by the operator. Use for exhaustive +# testing of all dtypes. +# - unsupported: Run tests on dtypes not supported by the operator. e.g. for +# testing the operator raises an error and doesn't crash. +# - supported_backward: Every dtype supported by the operator's backward pass. +# - unsupported_backward: Run tests on dtypes not supported by the operator's backward pass. +# - any_one: Runs a test for one dtype the operator supports. Prioritizes dtypes the +# operator supports in both forward and backward. +# - none: Useful for tests that are not dtype-specific. No dtype will be passed to the test +# when this is selected. +# - any_common_cpu_cuda_one: Pick a dtype that supports both CPU and CUDA. +class OpDTypes(Enum): + supported = 0 # Test all supported dtypes (default) + unsupported = 1 # Test only unsupported dtypes + supported_backward = 2 # Test all supported backward dtypes + unsupported_backward = 3 # Test only unsupported backward dtypes + any_one = 4 # Test precisely one supported dtype + none = 5 # Instantiate no dtype variants (no dtype kwarg needed) + any_common_cpu_cuda_one = ( + 6 # Test precisely one supported dtype that is common to both cuda and cpu + ) + + +# Arbitrary order +ANY_DTYPE_ORDER = ( + torch.float32, + torch.float64, + torch.complex64, + torch.complex128, + torch.float16, + torch.bfloat16, + torch.long, + torch.int32, + torch.int16, + torch.int8, + torch.uint8, + torch.bool, + torch.float8_e4m3fn, + torch.float8_e5m2, +) + + +def _serialize_sample(sample_input): + # NB: For OpInfos, SampleInput.summary() prints in a cleaner way. + if getattr(sample_input, "summary", None) is not None: + return sample_input.summary() + return str(sample_input) + + +# Decorator that defines the OpInfos a test template should be instantiated for. +# +# Example usage: +# +# @ops(unary_ufuncs) +# def test_numerics(self, device, dtype, op): +# +# +# This will instantiate variants of test_numerics for each given OpInfo, +# on each device the OpInfo's operator supports, and for every dtype supported by +# that operator. There are a few caveats to the dtype rule, explained below. +# +# The @ops decorator can accept two +# additional arguments, "dtypes" and "allowed_dtypes". If "dtypes" is specified +# then the test variants are instantiated for those dtypes, regardless of +# what the operator supports. If given "allowed_dtypes" then test variants +# are instantiated only for the intersection of allowed_dtypes and the dtypes +# they would otherwise be instantiated with. That is, allowed_dtypes composes +# with the options listed above and below. +# +# The "dtypes" argument can also accept additional values (see OpDTypes above): +# OpDTypes.supported - the test is instantiated for all dtypes the operator +# supports +# OpDTypes.unsupported - the test is instantiated for all dtypes the operator +# doesn't support +# OpDTypes.supported_backward - the test is instantiated for all dtypes the +# operator's gradient formula supports +# OpDTypes.unsupported_backward - the test is instantiated for all dtypes the +# operator's gradient formula doesn't support +# OpDTypes.any_one - the test is instantiated for one dtype the +# operator supports. The dtype supports forward and backward if possible. +# OpDTypes.none - the test is instantiated without any dtype. The test signature +# should not include a dtype kwarg in this case. +# OpDTypes.any_common_cpu_cuda_one - the test is instantiated for a dtype +# that supports both CPU and CUDA. +# +# These options allow tests to have considerable control over the dtypes +# they're instantiated for. + + +class ops(_TestParametrizer): + def __init__( + self, + op_list, + *, + dtypes: OpDTypes | Sequence[torch.dtype] = OpDTypes.supported, + allowed_dtypes: Sequence[torch.dtype] | None = None, + skip_if_dynamo=True, + ): + self.op_list = list(op_list) + self.opinfo_dtypes = dtypes + self.allowed_dtypes = ( + set(allowed_dtypes) if allowed_dtypes is not None else None + ) + self.skip_if_dynamo = skip_if_dynamo + + def _parametrize_test(self, test, generic_cls, device_cls): + """Parameterizes the given test function across each op and its associated dtypes.""" + if device_cls is None: + raise RuntimeError( + "The @ops decorator is only intended to be used in a device-specific " + "context; use it with instantiate_device_type_tests() instead of " + "instantiate_parametrized_tests()" + ) + + op = check_exhausted_iterator = object() + for op in self.op_list: + # Determine the set of dtypes to use. + dtypes: set[torch.dtype] | set[None] + if isinstance(self.opinfo_dtypes, Sequence): + dtypes = set(self.opinfo_dtypes) + elif self.opinfo_dtypes == OpDTypes.unsupported_backward: + dtypes = set(get_all_dtypes()).difference( + op.supported_backward_dtypes(device_cls.device_type) + ) + elif self.opinfo_dtypes == OpDTypes.supported_backward: + dtypes = op.supported_backward_dtypes(device_cls.device_type) + elif self.opinfo_dtypes == OpDTypes.unsupported: + dtypes = set(get_all_dtypes()).difference( + op.supported_dtypes(device_cls.device_type) + ) + elif self.opinfo_dtypes == OpDTypes.supported: + dtypes = set(op.supported_dtypes(device_cls.device_type)) + elif self.opinfo_dtypes == OpDTypes.any_one: + # Tries to pick a dtype that supports both forward or backward + supported = set(op.supported_dtypes(device_cls.device_type)) + supported_backward = op.supported_backward_dtypes( + device_cls.device_type + ) + supported_both = supported.intersection(supported_backward) + dtype_set = supported_both if len(supported_both) > 0 else supported + for dtype in ANY_DTYPE_ORDER: + if dtype in dtype_set: + dtypes = {dtype} + break + else: + dtypes = {} + elif self.opinfo_dtypes == OpDTypes.any_common_cpu_cuda_one: + # Tries to pick a dtype that supports both CPU and CUDA + supported = set(op.dtypes).intersection(op.dtypesIfCUDA) + if supported: + dtypes = { + next(dtype for dtype in ANY_DTYPE_ORDER if dtype in supported) + } + else: + dtypes = {} + + elif self.opinfo_dtypes == OpDTypes.none: + dtypes = {None} + else: + raise RuntimeError(f"Unknown OpDType: {self.opinfo_dtypes}") + + if self.allowed_dtypes is not None: + dtypes = dtypes.intersection(self.allowed_dtypes) + + # Construct the test name; device / dtype parts are handled outside. + # See [Note: device and dtype suffix placement] + test_name = op.formatted_name + + # Filter sample skips / xfails to only those that apply to the OpInfo. + # These are defined on the test function via decorators. + sample_skips_and_xfails = getattr(test, "sample_skips_and_xfails", None) + if sample_skips_and_xfails is not None: + sample_skips_and_xfails = [ + rule + for rule in sample_skips_and_xfails + if rule.op_match_fn(device_cls.device_type, op) + ] + + for dtype in dtypes: + # Construct parameter kwargs to pass to the test. + param_kwargs = {"op": op} + _update_param_kwargs(param_kwargs, "dtype", dtype) + + # NOTE: test_wrapper exists because we don't want to apply + # op-specific decorators to the original test. + # Test-specific decorators are applied to the original test, + # however. + try: + + @wraps(test) + def test_wrapper(*args, **kwargs): + try: + return test(*args, **kwargs) + except unittest.SkipTest as e: + raise e + except Exception as e: + tracked_input = get_tracked_input() + if PRINT_REPRO_ON_FAILURE and tracked_input is not None: + e_tracked = Exception( # noqa: TRY002 + f"{str(e)}\n\nCaused by {tracked_input.type_desc} " + f"at index {tracked_input.index}: " + f"{_serialize_sample(tracked_input.val)}" + ) + e_tracked._tracked_input = tracked_input # type: ignore[attr] + raise e_tracked from e + raise e + finally: + clear_tracked_input() + + if self.skip_if_dynamo and not TEST_WITH_TORCHINDUCTOR: + test_wrapper = skipIfTorchDynamo( + "Policy: we don't run OpInfo tests w/ Dynamo" + )(test_wrapper) + + # Initialize info for the last input seen. This is useful for tracking + # down which inputs caused a test failure. Note that TrackedInputIter is + # responsible for managing this. + test.tracked_input = None + + decorator_fn = partial( + op.get_decorators, + generic_cls.__name__, + test.__name__, + device_cls.device_type, + dtype, + ) + + if sample_skips_and_xfails is not None: + test_wrapper.sample_skips_and_xfails = sample_skips_and_xfails + + yield (test_wrapper, test_name, param_kwargs, decorator_fn) + except Exception as ex: + # Provides an error message for debugging before rethrowing the exception + print(f"Failed to instantiate {test_name} for op {op.name}!") + raise ex + if op is check_exhausted_iterator: + raise ValueError( + "An empty op_list was passed to @ops. " + "Note that this may result from reuse of a generator." + ) + + +# Decorator that skips a test if the given condition is true. +# Notes: +# (1) Skip conditions stack. +# (2) Skip conditions can be bools or strings. If a string the +# test base must have defined the corresponding attribute to be False +# for the test to run. If you want to use a string argument you should +# probably define a new decorator instead (see below). +# (3) Prefer the existing decorators to defining the 'device_type' kwarg. +class skipIf: + def __init__(self, dep, reason, device_type=None): + self.dep = dep + self.reason = reason + self.device_type = device_type + + def __call__(self, fn): + @wraps(fn) + def dep_fn(slf, *args, **kwargs): + if ( + self.device_type is None + or self.device_type == slf.device_type + or ( + isinstance(self.device_type, Iterable) + and slf.device_type in self.device_type + ) + ): + if (isinstance(self.dep, str) and getattr(slf, self.dep, True)) or ( + isinstance(self.dep, bool) and self.dep + ): + raise unittest.SkipTest(self.reason) + + return fn(slf, *args, **kwargs) + + return dep_fn + + +# Skips a test on CPU if the condition is true. +class skipCPUIf(skipIf): + def __init__(self, dep, reason): + super().__init__(dep, reason, device_type="cpu") + + +# Skips a test on CUDA if the condition is true. +class skipCUDAIf(skipIf): + def __init__(self, dep, reason): + super().__init__(dep, reason, device_type="cuda") + + +# Skips a test on XPU if the condition is true. +class skipXPUIf(skipIf): + def __init__(self, dep, reason): + super().__init__(dep, reason, device_type="xpu") + + +# Skips a test on XPU or CUDA if the condition is true. +class skipGPUIf(skipIf): + def __init__(self, dep, reason): + super().__init__(dep, reason, device_type=GPU_TYPES) + + +# Skips a test on Lazy if the condition is true. +class skipLazyIf(skipIf): + def __init__(self, dep, reason): + super().__init__(dep, reason, device_type="lazy") + + +# Skips a test on Meta if the condition is true. +class skipMetaIf(skipIf): + def __init__(self, dep, reason): + super().__init__(dep, reason, device_type="meta") + + +# Skips a test on MPS if the condition is true. +class skipMPSIf(skipIf): + def __init__(self, dep, reason): + super().__init__(dep, reason, device_type="mps") + + +class skipHPUIf(skipIf): + def __init__(self, dep, reason): + super().__init__(dep, reason, device_type="hpu") + + +# Skips a test on XLA if the condition is true. +class skipXLAIf(skipIf): + def __init__(self, dep, reason): + super().__init__(dep, reason, device_type="xla") + + +class skipPRIVATEUSE1If(skipIf): + def __init__(self, dep, reason): + device_type = torch._C._get_privateuse1_backend_name() + super().__init__(dep, reason, device_type=device_type) + + +def _has_sufficient_memory(device, size): + device_ = torch.device(device) + device_type = device_.type + if device_type in ["cuda", "xpu"]: + acc = torch.accelerator.current_accelerator() + # Case 1: no accelerator found + if not acc: + return False + # Case 2: accelerator found but not matching device type + if acc.type != device_type: + return True + # Case 3: accelerator found and matching device type but not available + if not torch.accelerator.is_available(): + return False + # Case 4: accelerator found and matching device type and available + gc.collect() + torch.accelerator.empty_cache() + + if device_.index is None: + device_ = torch.device(device_type, 0) + + if device_type == "cuda": + return ( + torch.cuda.memory.mem_get_info(device_)[0] + * torch.cuda.memory.get_per_process_memory_fraction(device_) + ) >= size + + if device_type == "xpu": + return torch.xpu.memory.mem_get_info(device_)[0] >= size + + if device_type == "xla": + raise unittest.SkipTest("TODO: Memory availability checks for XLA?") + + if device_type != "cpu": + raise unittest.SkipTest("Unknown device type") + + # CPU + if not HAS_PSUTIL: + raise unittest.SkipTest("Need psutil to determine if memory is sufficient") + + # The sanitizers have significant memory overheads + if TEST_WITH_ASAN or TEST_WITH_TSAN or TEST_WITH_UBSAN: + effective_size = size * 10 + else: + effective_size = size + + # don't try using all RAM on s390x, leave some for service processes + if IS_S390X: + effective_size = effective_size * 2 + + if psutil.virtual_memory().available < effective_size: + gc.collect() + return psutil.virtual_memory().available >= effective_size + + +def largeTensorTest(size, device=None, inductor=TEST_WITH_TORCHINDUCTOR): + """Skip test if the device has insufficient memory to run the test + + size may be a number of bytes, a string of the form "N GB", or a callable + + If the test is a device generic test, available memory on the primary device will be checked. + It can also be overridden by the optional `device=` argument. + In other tests, the `device=` argument needs to be specified. + """ + if isinstance(size, str): + if not size.endswith(("GB", "gb")): + raise AssertionError(f"only bytes or GB supported, got {size!r}") + size = 1024**3 * int(size[:-2]) + + def inner(fn): + @wraps(fn) + def dep_fn(self, *args, **kwargs): + size_bytes: int = size(self, *args, **kwargs) if callable(size) else size + _device = device + if _device is None: + if hasattr(self, "get_primary_device"): + _device = self.get_primary_device() + else: + _device = self.device + + # If this is running with GPU cpp_wrapper, the autotuning step will generate + # an additional array of the same size as the input. + if inductor and torch._inductor.config.cpp_wrapper and _device != "cpu": + size_bytes *= 2 + if not _has_sufficient_memory(_device, size_bytes): + raise unittest.SkipTest(f"Insufficient {_device} memory") + + return fn(self, *args, **kwargs) + + return dep_fn + + return inner + + +class expectedFailure: + def __init__(self, device_type, dtype=None): + self.device_type = device_type + self.dtype = dtype + + def __call__(self, fn): + @wraps(fn) + def efail_fn(slf, *args, **kwargs): + if ( + not hasattr(slf, "device_type") + and hasattr(slf, "device") + and isinstance(slf.device, str) + ): + target_device_type = slf.device + else: + target_device_type = slf.device_type + + target_dtype = kwargs.get("dtype", getattr(slf, "dtype", None)) + device_matches = ( + self.device_type is None or self.device_type == target_device_type + ) + dtype_matches = self.dtype is None or self.dtype == target_dtype + + if device_matches and dtype_matches: + try: + fn(slf, *args, **kwargs) + except Exception: + return + else: + slf.fail("expected test to fail, but it passed") + + return fn(slf, *args, **kwargs) + + return efail_fn + + +class onlyOn: + def __init__(self, device_type: str | list): + self.device_type = device_type + + def __call__(self, fn): + @wraps(fn) + def only_fn(slf, *args, **kwargs): + if slf.device_type not in self.device_type: + if getattr(slf, "bypass_device_restrictions", False): + return fn(slf, *args, **kwargs) + reason = f"Only runs on {self.device_type}" + if IS_SANDCASTLE or IS_FBCODE: + print( + f"Skipping {fn.__name__} on sandcastle for following reason: {reason}", + file=sys.stderr, + ) + return + raise unittest.SkipTest(reason) + + return fn(slf, *args, **kwargs) + + return only_fn + + +# Decorator that provides all available devices of the device type to the test +# as a list of strings instead of providing a single device string. +# Skips the test if the number of available devices of the variant's device +# type is less than the 'num_required_devices' arg. +class deviceCountAtLeast: + def __init__(self, num_required_devices): + self.num_required_devices = num_required_devices + + def __call__(self, fn): + if hasattr(fn, "num_required_devices"): + raise AssertionError(f"deviceCountAtLeast redefinition for {fn.__name__}") + fn.num_required_devices = self.num_required_devices + + @wraps(fn) + def multi_fn(slf, devices, *args, **kwargs): + if len(devices) < self.num_required_devices: + reason = f"fewer than {self.num_required_devices} devices detected" + if IS_SANDCASTLE or IS_FBCODE: + print( + f"Skipping {fn.__name__} on sandcastle for following reason: {reason}", + file=sys.stderr, + ) + return + raise unittest.SkipTest(reason) + + return fn(slf, devices, *args, **kwargs) + + return multi_fn + + +# Only runs the test on the native device type (currently CPU, CUDA, Meta and PRIVATEUSE1) +def onlyNativeDeviceTypes(fn: Callable[_P, _T]) -> Callable[_P, _T]: + @wraps(fn) + def only_fn(self, *args: _P.args, **kwargs: _P.kwargs) -> _T: + if self.device_type not in NATIVE_DEVICES: + reason = f"onlyNativeDeviceTypes: doesn't run on {self.device_type}" + raise unittest.SkipTest(reason) + + return fn(self, *args, **kwargs) + + return only_fn + + +# Only runs the test on the native device types and devices specified in the devices list +def onlyNativeDeviceTypesAnd(devices=None): + def decorator(fn): + @wraps(fn) + def only_fn(self, *args, **kwargs): + if ( + self.device_type not in NATIVE_DEVICES + and self.device_type not in devices + ): + reason = f"onlyNativeDeviceTypesAnd {devices} : doesn't run on {self.device_type}" + raise unittest.SkipTest(reason) + + return fn(self, *args, **kwargs) + + return only_fn + + return decorator + + +# Specifies per-dtype precision overrides. +# Ex. +# +# @precisionOverride({torch.half : 1e-2, torch.float : 1e-4}) +# @dtypes(torch.half, torch.float, torch.double) +# def test_X(self, device, dtype): +# ... +# +# When the test is instantiated its class's precision will be set to the +# corresponding override, if it exists. +# self.precision can be accessed directly, and it also controls the behavior of +# functions like self.assertEqual(). +# +# Note that self.precision is a scalar value, so if you require multiple +# precisions (or are working with multiple dtypes) they should be specified +# explicitly and computed using self.precision (e.g. +# self.precision *2, max(1, self.precision)). +class precisionOverride: + def __init__(self, d): + if not isinstance(d, dict): + raise AssertionError( + "precisionOverride not given a dtype : precision dict!" + ) + for dtype in d: + if not isinstance(dtype, torch.dtype): + raise AssertionError(f"precisionOverride given unknown dtype {dtype}") + + self.d = d + + def __call__(self, fn): + fn.precision_overrides = self.d + return fn + + +# Specifies per-dtype tolerance overrides tol(atol, rtol). It has priority over +# precisionOverride. +# Ex. +# +# @toleranceOverride({torch.float : tol(atol=1e-2, rtol=1e-3}, +# torch.double : tol{atol=1e-4, rtol = 0}) +# @dtypes(torch.half, torch.float, torch.double) +# def test_X(self, device, dtype): +# ... +# +# When the test is instantiated its class's tolerance will be set to the +# corresponding override, if it exists. +# self.rtol and self.precision can be accessed directly, and they also control +# the behavior of functions like self.assertEqual(). +# +# The above example sets atol = 1e-2 and rtol = 1e-3 for torch.float and +# atol = 1e-4 and rtol = 0 for torch.double. +tol = namedtuple("tol", ["atol", "rtol"]) + + +class toleranceOverride: + def __init__(self, d): + if not isinstance(d, dict): + raise AssertionError("toleranceOverride not given a dtype : tol dict!") + for dtype, prec in d.items(): + if not isinstance(dtype, torch.dtype): + raise AssertionError(f"toleranceOverride given unknown dtype {dtype}") + if not isinstance(prec, tol): + raise AssertionError("toleranceOverride not given a dtype : tol dict!") + + self.d = d + + def __call__(self, fn): + fn.tolerance_overrides = self.d + return fn + + +# Decorator that instantiates a variant of the test for each given dtype. +# Notes: +# (1) Tests that accept the dtype argument MUST use this decorator. +# (2) Can be overridden for CPU or CUDA, respectively, using dtypesIfCPU +# or dtypesIfCUDA. +# (3) Can accept an iterable of dtypes or an iterable of tuples +# of dtypes. +# Examples: +# @dtypes(torch.float32, torch.float64) +# @dtypes((torch.long, torch.float32), (torch.int, torch.float64)) +class dtypes: + def __init__(self, *args, device_type="all"): + if len(args) > 0 and isinstance(args[0], (list, tuple)): + for arg in args: + if not isinstance(arg, (list, tuple)): + raise AssertionError( + "When one dtype variant is a tuple or list, " + "all dtype variants must be. " + f"Received non-list non-tuple dtype {str(arg)}" + ) + if not all(isinstance(dtype, torch.dtype) for dtype in arg): + raise AssertionError(f"Unknown dtype in {str(arg)}") + else: + if not all(isinstance(arg, torch.dtype) for arg in args): + raise AssertionError(f"Unknown dtype in {str(args)}") + + self.args = args + self.device_type = device_type + + def __call__(self, fn): + d = getattr(fn, "dtypes", {}) + if self.device_type in d: + raise AssertionError(f"dtypes redefinition for {self.device_type}") + d[self.device_type] = self.args + fn.dtypes = d + return fn + + +# Overrides specified dtypes on the CPU. +class dtypesIfCPU(dtypes): + def __init__(self, *args): + super().__init__(*args, device_type="cpu") + + +# Overrides specified dtypes on CUDA. +class dtypesIfCUDA(dtypes): + def __init__(self, *args): + super().__init__(*args, device_type="cuda") + + +# Overrides specified dtypes on Intel GPU. +class dtypesIfXPU(dtypes): + def __init__(self, *args): + super().__init__(*args, device_type="xpu") + + +class dtypesIfMPS(dtypes): + def __init__(self, *args): + super().__init__(*args, device_type="mps") + + +class dtypesIfHPU(dtypes): + def __init__(self, *args): + super().__init__(*args, device_type="hpu") + + +class dtypesIfPRIVATEUSE1(dtypes): + def __init__(self, *args): + super().__init__(*args, device_type=torch._C._get_privateuse1_backend_name()) + + +def onlyCPU(fn): + return onlyOn("cpu")(fn) + + +def onlyCUDA(fn): + return onlyOn("cuda")(fn) + + +def onlyMPS(fn): + return onlyOn("mps")(fn) + + +def onlyXPU(fn): + return onlyOn("xpu")(fn) + + +def onlyHPU(fn): + return onlyOn("hpu")(fn) + + +def onlyPRIVATEUSE1(fn): + device_type = torch._C._get_privateuse1_backend_name() + device_mod = getattr(torch, device_type, None) + if device_mod is None: + reason = f"Skip as torch has no module of {device_type}" + return unittest.skip(reason)(fn) + return onlyOn(device_type)(fn) + + +def onlyCUDAAndPRIVATEUSE1(fn): + @wraps(fn) + def only_fn(self, *args, **kwargs): + if self.device_type not in ("cuda", torch._C._get_privateuse1_backend_name()): + reason = f"onlyCUDAAndPRIVATEUSE1: doesn't run on {self.device_type}" + raise unittest.SkipTest(reason) + + return fn(self, *args, **kwargs) + + return only_fn + + +def disablecuDNN(fn): + @wraps(fn) + def disable_cudnn(self, *args, **kwargs): + if self.device_type == "cuda" and self.has_cudnn(): + with torch.backends.cudnn.flags(enabled=False): + return fn(self, *args, **kwargs) + return fn(self, *args, **kwargs) + + return disable_cudnn + + +def disableMkldnn(fn): + @wraps(fn) + def disable_mkldnn(self, *args, **kwargs): + if torch.backends.mkldnn.is_available(): + with torch.backends.mkldnn.flags(enabled=False): + return fn(self, *args, **kwargs) + return fn(self, *args, **kwargs) + + return disable_mkldnn + + +def expectedFailureCPU(fn): + return expectedFailure("cpu")(fn) + + +def expectedFailureCUDA(fn): + return expectedFailure("cuda")(fn) + + +def expectedFailureXPU(fn): + return expectedFailure("xpu")(fn) + + +def expectedFailureMeta(fn): + return skipIfTorchDynamo()(expectedFailure("meta")(fn)) + + +def expectedFailureXLA(fn): + return expectedFailure("xla")(fn) + + +def expectedFailureHPU(fn): + return expectedFailure("hpu")(fn) + + +def expectedFailureMPS(fn): + return expectedFailure("mps")(fn) + + +def expectedFailureMPSComplex(fn): + return expectedFailure("mps", torch.complex64)(fn) + + +def expectedFailureMPSPre15(fn): + import platform + + version = float(".".join(platform.mac_ver()[0].split(".")[:2]) or -1) + if not version or version < 1.0: # cpu or other unsupported device + return fn + if version < 15.0: + return expectedFailure("mps")(fn) + return fn + + +def expectedFailureMPSPre14(fn): + import platform + + version = float(".".join(platform.mac_ver()[0].split(".")[:2]) or -1) + if not version or version < 1.0: # cpu or other unsupported device + return fn + if version < 14.0: + return expectedFailure("mps")(fn) + return fn + + +# Skips a test on CPU if LAPACK is not available. +def skipCPUIfNoLapack(fn): + return skipCPUIf(not torch._C.has_lapack, "PyTorch compiled without Lapack")(fn) + + +# Skips a test on CPU if FFT is not available. +def skipCPUIfNoFFT(fn): + return skipCPUIf(not torch._C.has_spectral, "PyTorch is built without FFT support")( + fn + ) + + +# Skips a test on CPU if MKL is not available. +def skipCPUIfNoMkl(fn): + return skipCPUIf(not TEST_MKL, "PyTorch is built without MKL support")(fn) + + +# Skips a test on CPU if MKL Sparse is not available (it's not linked on Windows). +def skipCPUIfNoMklSparse(fn): + return skipCPUIf( + IS_WINDOWS or not TEST_MKL, "PyTorch is built without MKL support" + )(fn) + + +# Skips a test on CPU if mkldnn is not available. +def skipCPUIfNoMkldnn(fn): + return skipCPUIf( + not torch.backends.mkldnn.is_available(), + "PyTorch is built without mkldnn support", + )(fn) + + +# Skips a test on CUDA if MAGMA is not available. +def skipCUDAIfNoMagma(fn): + return skipCUDAIf("no_magma", "no MAGMA library detected")( + skipCUDANonDefaultStreamIf(True)(fn) + ) + + +def has_cusolver(): + return not TEST_WITH_ROCM + + +def has_hipsolver(): + rocm_version = _get_torch_rocm_version() + # hipSOLVER is disabled on ROCM < 5.3 + return rocm_version >= (5, 3) + + +# Skips a test on CUDA if cuSOLVER is not available, +# and on ROCm if MAGMA is not available. +def skipCUDAIfNoCusolverROCMIfNoMagma(fn): + if TEST_WITH_ROCM: + return skipCUDAIfNoMagma(fn) + else: + return skipCUDAIfNoCusolver(fn) + + +# Skips a test on CUDA/ROCM if cuSOLVER/hipSOLVER is not available +def skipCUDAIfNoCusolver(fn): + return skipCUDAIf( + not has_cusolver() and not has_hipsolver(), "cuSOLVER not available" + )(fn) + + +# Skips a test if both cuSOLVER and MAGMA are not available +def skipCUDAIfNoMagmaAndNoCusolver(fn): + if has_cusolver(): + return fn + else: + # cuSolver is disabled on cuda < 10.1.243, tests depend on MAGMA + return skipCUDAIfNoMagma(fn) + + +# Skips a test if both cuSOLVER/hipSOLVER and MAGMA are not available +def skipCUDAIfNoMagmaAndNoLinalgsolver(fn): + if has_cusolver() or has_hipsolver(): + return fn + else: + # cuSolver is disabled on cuda < 10.1.243, tests depend on MAGMA + return skipCUDAIfNoMagma(fn) + + +# Skips a test on CUDA when using ROCm. +def skipCUDAIfRocm(func=None, *, msg="test doesn't currently work on the ROCm stack"): + def dec_fn(fn): + reason = f"skipCUDAIfRocm: {msg}" + return skipCUDAIf(TEST_WITH_ROCM, reason=reason)(fn) + + if func: + return dec_fn(func) + return dec_fn + + +# Skips a test on CUDA when not using ROCm. +def skipCUDAIfNotRocm(fn): + return skipCUDAIf( + not TEST_WITH_ROCM, "test doesn't currently work on the CUDA stack" + )(fn) + + +# Skips a test on CUDA if ROCm is unavailable or its version is lower than requested. +def skipCUDAIfRocmVersionLessThan(version=None): + def dec_fn(fn): + @wraps(fn) + def wrap_fn(self, *args, **kwargs): + if self.device_type == "cuda": + if not TEST_WITH_ROCM: + reason = "ROCm not available" + raise unittest.SkipTest(reason) + rocm_version_tuple = _get_torch_rocm_version() + if ( + rocm_version_tuple is None + or version is None + or rocm_version_tuple < tuple(version) + ): + reason = ( + f"ROCm {rocm_version_tuple} is available but {version} required" + ) + raise unittest.SkipTest(reason) + + return fn(self, *args, **kwargs) + + return wrap_fn + + return dec_fn + + +# Skips a test on CUDA if ROCm hipBLASLt is unavailable or its version is lower than requested. +def skipCUDAIfRocmHipBlasltVersionLessThan(version=None): + def dec_fn(fn): + @wraps(fn) + def wrap_fn(self, *args, **kwargs): + if self.device_type == "cuda": + if not TEST_WITH_ROCM: + reason = "ROCm not available" + raise unittest.SkipTest(reason) + hipblaslt_version_tuple = _get_torch_hipblaslt_version() + if ( + hipblaslt_version_tuple is None + or version is None + or hipblaslt_version_tuple < tuple(version) + ): + reason = f"hipBLASLt {hipblaslt_version_tuple} is available but {version} required" + raise unittest.SkipTest(reason) + + return fn(self, *args, **kwargs) + + return wrap_fn + + return dec_fn + + +# Skips a test on CUDA when using ROCm. +def skipCUDAIfNotMiopenSuggestNHWC(fn): + return skipCUDAIf( + not TEST_WITH_MIOPEN_SUGGEST_NHWC, + "test doesn't currently work without MIOpen NHWC activation", + )(fn) + + +# Skips a test for specified CUDA versions, given in the form of a list of [major, minor]s. +def skipCUDAVersionIn(versions: list[tuple[int, int]] | None = None): + def dec_fn(fn): + @wraps(fn) + def wrap_fn(self, *args, **kwargs): + version = _get_torch_cuda_version() + if version == (0, 0): # cpu or rocm + return fn(self, *args, **kwargs) + if version in (versions or []): + reason = f"test skipped for CUDA version {version}" + raise unittest.SkipTest(reason) + return fn(self, *args, **kwargs) + + return wrap_fn + + return dec_fn + + +# Skips a test for CUDA versions less than specified, given in the form of [major, minor]. +def skipCUDAIfVersionLessThan(versions: tuple[int, int] | None = None): + def dec_fn(fn): + @wraps(fn) + def wrap_fn(self, *args, **kwargs): + version = _get_torch_cuda_version() + if version == (0, 0): # cpu or rocm + return fn(self, *args, **kwargs) + if version < versions: + reason = f"test skipped for CUDA versions < {version}" + raise unittest.SkipTest(reason) + return fn(self, *args, **kwargs) + + return wrap_fn + + return dec_fn + + +# Skips a test on CUDA if cuDNN is unavailable or its version is lower than requested. +def skipCUDAIfCudnnVersionLessThan(version=0): + def dec_fn(fn): + @wraps(fn) + def wrap_fn(self, *args, **kwargs): + if self.device_type == "cuda": + if self.no_cudnn: + reason = "cuDNN not available" + raise unittest.SkipTest(reason) + if self.cudnn_version is None or self.cudnn_version < version: + reason = f"cuDNN version {self.cudnn_version} is available but {version} required" + raise unittest.SkipTest(reason) + + return fn(self, *args, **kwargs) + + return wrap_fn + + return dec_fn + + +# Skips a test on CUDA if cuSparse generic API is not available +def skipCUDAIfNoCusparseGeneric(fn): + return skipCUDAIf(not TEST_CUSPARSE_GENERIC, "cuSparse Generic API not available")( + fn + ) + + +def skipCUDAIfNoHipsparseGeneric(fn): + return skipCUDAIf( + not TEST_HIPSPARSE_GENERIC, "hipSparse Generic API not available" + )(fn) + + +def skipCUDAIfNoSparseGeneric(fn): + return skipCUDAIf( + not (TEST_CUSPARSE_GENERIC or TEST_HIPSPARSE_GENERIC), + "Sparse Generic API not available", + )(fn) + + +def skipCUDAIfNoCudnn(fn): + return skipCUDAIfCudnnVersionLessThan(0)(fn) + + +def skipCUDAIfMiopen(fn): + return skipCUDAIf(torch.version.hip is not None, "Marked as skipped for MIOpen")(fn) + + +def skipCUDAIfNoMiopen(fn): + return skipCUDAIf(torch.version.hip is None, "MIOpen is not available")( + skipCUDAIfNoCudnn(fn) + ) + + +def skipLazy(fn): + return skipLazyIf(True, "test doesn't work with lazy tensors")(fn) + + +def skipMeta(fn): + return skipMetaIf(True, "test doesn't work with meta tensors")(fn) + + +def skipXLA(fn): + return skipXLAIf(True, "Marked as skipped for XLA")(fn) + + +def skipMPS(fn): + return skipMPSIf(True, "test doesn't work on MPS backend")(fn) + + +def skipHPU(fn): + return skipHPUIf(True, "test doesn't work on HPU backend")(fn) + + +def skipXPU(fn): + return skipXPUIf(True, "test doesn't work on XPU backend")(fn) + + +def skipPRIVATEUSE1(fn): + return skipPRIVATEUSE1If(True, "test doesn't work on privateuse1 backend")(fn) + + +# TODO: the "all" in the name isn't true anymore for quite some time as we have also have for example XLA and MPS now. +# This should probably enumerate all available device type test base classes. +def get_all_device_types() -> list[str]: + return ["cpu"] if not torch.cuda.is_available() else ["cpu", "cuda"] + + +# skip since currently flex attention requires at least `avx2` support on CPU. +IS_FLEX_ATTENTION_CPU_PLATFORM_SUPPORTED = ( + not IS_MACOS + and torch.cpu._is_avx2_supported() + and os.getenv("ATEN_CPU_CAPABILITY") != "default" +) +IS_FLEX_ATTENTION_XPU_PLATFORM_SUPPORTED = ( + torch.xpu.is_available() and torch.utils._triton.has_triton() +) +IS_FLEX_ATTENTION_CUDA_PLATFORM_SUPPORTED = ( + torch.cuda.is_available() + and torch.utils._triton.has_triton() + and torch.cuda.get_device_capability() >= (8, 0) +) +flex_attention_supported_platform = unittest.skipUnless( + IS_FLEX_ATTENTION_XPU_PLATFORM_SUPPORTED + or ( + IS_FLEX_ATTENTION_CPU_PLATFORM_SUPPORTED + and not torch.xpu.is_available() + and not torch.cuda.is_available() + ) + or IS_FLEX_ATTENTION_CUDA_PLATFORM_SUPPORTED, + "Requires CUDA and Triton, Intel GPU and triton, or CPU with avx2 and later", +) +if ( + torch.version.hip + and torch.cuda.device_count() > 0 + and "gfx94" in torch.cuda.get_device_properties(0).gcnArchName +): + e4m3_type = torch.float8_e4m3fnuz + e5m2_type = torch.float8_e5m2fnuz + E4M3_MAX_POS = torch.finfo(torch.float8_e4m3fnuz).max + E5M2_MAX_POS = torch.finfo(torch.float8_e5m2fnuz).max +else: + e4m3_type = torch.float8_e4m3fn + e5m2_type = torch.float8_e5m2 + E4M3_MAX_POS = torch.finfo(torch.float8_e4m3fn).max + E5M2_MAX_POS = torch.finfo(torch.float8_e5m2).max diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_dist_composable.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_dist_composable.py new file mode 100644 index 0000000000000000000000000000000000000000..fd14b85a21915ddf8ab415f3bf5dc6e79db14dfc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_dist_composable.py @@ -0,0 +1,112 @@ +# mypy: ignore-errors + +# Owner(s): ["oncall: distributed"] + + +import torch +import torch.nn as nn + + +class UnitModule(nn.Module): + def __init__(self, device: torch.device): + super().__init__() + self.l1 = nn.Linear(100, 100, device=device) + self.seq = nn.Sequential( + nn.ReLU(), + nn.Linear(100, 100, device=device), + nn.ReLU(), + ) + self.l2 = nn.Linear(100, 100, device=device) + + def forward(self, x): + return self.l2(self.seq(self.l1(x))) + + +class CompositeModel(nn.Module): + def __init__(self, device: torch.device): + super().__init__() + self.l1 = nn.Linear(100, 100, device=device) + self.u1 = UnitModule(device) + self.u2 = UnitModule(device) + self.l2 = nn.Linear(100, 100, device=device) + + def forward(self, x): + return self.l2(self.u2(self.u1(self.l1(x)))) + + +class UnitParamModule(nn.Module): + def __init__(self, device: torch.device): + super().__init__() + self.l = nn.Linear(100, 100, device=device) + self.seq = nn.Sequential( + nn.ReLU(), + nn.Linear(100, 100, device=device), + nn.ReLU(), + ) + self.p = nn.Parameter(torch.randn((100, 100), device=device)) + + def forward(self, x): + return torch.mm(self.seq(self.l(x)), self.p) + + +class CompositeParamModel(nn.Module): + def __init__(self, device: torch.device): + super().__init__() + self.l = nn.Linear(100, 100, device=device) + self.u1 = UnitModule(device) + self.u2 = UnitModule(device) + self.p = nn.Parameter(torch.randn((100, 100), device=device)) + self.register_buffer( + "buffer", torch.randn((100, 100), device=device), persistent=True + ) + + def forward(self, x): + a = self.u2(self.u1(self.l(x))) + b = self.p + return torch.mm(a, b) + + +class FakeSequential(nn.Module): + # Define this class to achieve a desired nested wrapping using the module + # wrap policy with `nn.Sequential` + def __init__(self, *modules: tuple[nn.Module, ...]) -> None: + super().__init__() + self._module_sequence = list(modules) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for module in self._module_sequence: + x = module(x) + return x + + +class NestedSequentialModel(nn.Module): + def __init__(self, device: torch.device) -> None: + super().__init__() + # This nested structure exercises traversal order to catch differences + # between valid traversals (e.g. BFS and DFS variations). + self.seq1 = nn.Sequential( + nn.Linear(1, 1, device=device), + FakeSequential( + nn.Linear(1, 1, device=device), + nn.ReLU(), + FakeSequential( + nn.Linear(1, 1, device=device), + ), + nn.ReLU(), + ), + nn.Linear(1, 2, device=device), + ) + self.lin = nn.Linear(2, 2, device=device) + self.seq2 = nn.Sequential( + nn.ReLU(), + nn.Linear(2, 3, device=device), + FakeSequential( + nn.Linear(3, 2, bias=False, device=device), + nn.Linear(2, 4, bias=False, device=device), + ), + ) + + # FIXME(rec): forward() is not a method, it's a local function inside __init__ + # that is never used. It should probabkly be outdented by four spaces, or removed. + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.seq2(self.lin(self.seq1(x))) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_distributed.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..894acb6be19fe13734f112c5d59c3d0a4945a5d2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_distributed.py @@ -0,0 +1,2110 @@ +# mypy: ignore-errors + +import faulthandler +import functools +import inspect +import itertools +import logging +import multiprocessing +import operator +import os +import queue +import subprocess +import sys +import tempfile +import threading +import time +import traceback +import types +import unittest +from collections.abc import Callable +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import timedelta +from enum import Enum +from functools import partial, reduce, wraps +from io import StringIO +from typing import Any, NamedTuple +from unittest.mock import patch + +import torch +import torch._dynamo.test_case +import torch.cuda.nccl +import torch.distributed as c10d +import torch.nn as nn +from torch._C._autograd import DeviceType +from torch._C._distributed_c10d import _SymmetricMemory +from torch._logging._internal import trace_log +from torch.testing._internal import common_utils +from torch.testing._internal.common_utils import ( + FILE_SCHEMA, + find_free_port, + IS_SANDCASTLE, + LazyVal, + retry_on_connect_failures, + skip_but_pass_in_sandcastle, + skip_but_pass_in_sandcastle_if, + TEST_CUDA, + TEST_HPU, + TEST_WITH_ROCM, + TEST_WITH_TSAN, + TEST_XPU, + TestCase, +) +from torch.testing._internal.distributed.multi_threaded_pg import ( + _install_threaded_pg, + _uninstall_threaded_pg, + ProcessLocalGroup, +) + + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + +ACCELERATOR_DIST_BACKENDS = ["nccl", "xccl", "hccl"] +DDP_RANK_DEVICES = ["cuda", "xpu"] +HAS_ACCELERATOR = TEST_CUDA or TEST_HPU or TEST_XPU + + +class TestSkip(NamedTuple): + exit_code: int + message: str + + +TEST_SKIPS = { + "backend_unavailable": TestSkip( + 72, "Skipped because distributed backend is not available." + ), + "small_worldsize": TestSkip(73, "Skipped due to small world size."), + "odd_worldsize": TestSkip(87, "Skipped due to odd world size."), + "no_cuda": TestSkip(74, "CUDA is not available."), + "multi-gpu-1": TestSkip(75, "Need at least 1 CUDA device"), + "multi-gpu-2": TestSkip(77, "Need at least 2 CUDA devices"), + "multi-gpu-3": TestSkip(80, "Need at least 3 CUDA devices"), + "multi-gpu-4": TestSkip(81, "Need at least 4 CUDA devices"), + "multi-gpu-5": TestSkip(82, "Need at least 5 CUDA devices"), + "multi-gpu-6": TestSkip(83, "Need at least 6 CUDA devices"), + "multi-gpu-7": TestSkip(84, "Need at least 7 CUDA devices"), + "multi-gpu-8": TestSkip(85, "Need at least 8 CUDA devices"), + "nccl": TestSkip(76, "c10d not compiled with NCCL support"), + "skipIfRocm": TestSkip(78, "Test skipped for ROCm"), + "no_peer_access": TestSkip(79, "Test skipped because no GPU peer access"), + "generic": TestSkip( + 86, "Test skipped at subprocess level, look at subprocess log for skip reason" + ), + "importerror": TestSkip(88, "Test skipped due to missing import"), + "no_accelerator": TestSkip(89, "accelerator is not available."), +} + + +@dataclass +class DistTestCases: + # Backends that do not support a specific collective + skip_collective = {} + skip_collective["allgather_coalesced"] = {"nccl", "mpi", "ucc", "xccl"} + skip_collective["reduce"] = set() + skip_collective["sendrecv anysource"] = {"nccl", "ucc", "xccl"} + skip_collective["cpu barrier"] = {"nccl", "ucc", "xccl"} + + # Sets showing that something is implemented + backend_feature = {} + backend_feature["gpu"] = {"nccl", "gloo", "ucc"} + backend_feature["cuda"] = {"nccl", "gloo", "ucc"} + backend_feature["ddp"] = {"nccl", "gloo", "ucc"} + backend_feature["subgroup"] = {"nccl", "gloo", "ucc"} + backend_feature["plugin"] = set() + if TEST_HPU: + backend_feature["hpu"] = {"hccl"} + if TEST_XPU: + backend_feature["xpu"] = {"xccl"} + + +def requires_ddp_rank(device): + return device in DDP_RANK_DEVICES + + +def skip_if_no_gpu(func): + """Skips if the world size exceeds the number of GPUs, ensuring that if the + test is run, each rank has its own GPU via ``torch.cuda.device(rank)``.""" + + @wraps(func) + def wrapper(*args, **kwargs): + if not (TEST_CUDA or TEST_HPU or TEST_XPU): + sys.exit(TEST_SKIPS["no_cuda"].exit_code) + world_size = int(os.environ["WORLD_SIZE"]) + if TEST_CUDA and torch.cuda.device_count() < world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{world_size}"].exit_code) + if TEST_HPU and torch.hpu.device_count() < world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{world_size}"].exit_code) + if TEST_XPU and torch.xpu.device_count() < world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{world_size}"].exit_code) + + return func(*args, **kwargs) + + return wrapper + + +# TODO (kwen2501): what is the purpose of this decorator? Tests with this +# decorator were always skipped. So they may be outdated already. +# Oct 2024: bumping the small-world criteria to < 8, as we are increasing the +# number of GPUs in CI from 2 to 4, and we need to continue skipping those tests +# to keep CI green. But this is just a temporary solution. We should clean up +# those tests somehow. +def skip_if_small_worldsize(func): + @wraps(func) + def wrapper(*args, **kwargs): + if (os.environ["BACKEND"] != "mpi") and int(os.environ["WORLD_SIZE"]) < 8: + sys.exit(TEST_SKIPS["small_worldsize"].exit_code) + + return func(*args, **kwargs) + + return wrapper + + +def skip_if_odd_worldsize(func): + @wraps(func) + def wrapper(*args, **kwargs): + if (os.environ["BACKEND"] != "mpi") and int(os.environ["WORLD_SIZE"]) % 2 == 1: + sys.exit(TEST_SKIPS["odd_worldsize"].exit_code) + + return func(*args, **kwargs) + + return wrapper + + +def require_n_gpus_for_nccl_backend(n, backend): + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + if backend == "nccl" and torch.cuda.device_count() < n: + sys.exit(TEST_SKIPS[f"multi-gpu-{n}"].exit_code) + else: + return func(*args, **kwargs) + + return wrapper + + return decorator + + +def import_transformers_or_skip(): + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + try: + from transformers import AutoModelForMaskedLM, BertConfig # noqa: F401 + + return func(*args, **kwargs) + except ImportError: + sys.exit(TEST_SKIPS["importerror"].exit_code) + + return wrapper + + return decorator + + +def at_least_x_gpu(x): + if TEST_CUDA and torch.cuda.device_count() >= x: + return True + if TEST_HPU and torch.hpu.device_count() >= x: + return True + if TEST_XPU and torch.xpu.device_count() >= x: + return True + return False + + +def _maybe_handle_skip_if_lt_x_gpu(args, msg) -> bool: + _handle_test_skip = getattr(args[0], "_handle_test_skip", None) + if len(args) == 0 or _handle_test_skip is None: + return False + _handle_test_skip(msg) + return True + + +def skip_if_lt_x_gpu(x, *, allow_cpu=False): + """Skip if fewer than x accelerators available. + + Args: + x: Minimum number of accelerators required. + allow_cpu: If True, run the test on CPU-only machines (no accelerators). + """ + + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + if torch.cuda.is_available() and torch.cuda.device_count() >= x: + return func(*args, **kwargs) + if TEST_HPU and torch.hpu.device_count() >= x: + return func(*args, **kwargs) + if TEST_XPU and torch.xpu.device_count() >= x: + return func(*args, **kwargs) + if allow_cpu and not (torch.cuda.is_available() or TEST_HPU or TEST_XPU): + return func(*args, **kwargs) + test_skip = TEST_SKIPS[f"multi-gpu-{x}"] + if not _maybe_handle_skip_if_lt_x_gpu(args, test_skip.message): + sys.exit(test_skip.exit_code) + + return wrapper + + return decorator + + +def requires_world_size(n: int): + """ + Decorator to request a specific world size for a test. The test harness can + read this attribute to set the number of ranks to spawn. If there are fewer + than `n` CUDA devices available, the test should be skipped by the harness. + + Usage: + @require_world_size(3) + def test_something(self): + ... + """ + + def decorator(func): + func._required_world_size = n + available = torch.cuda.device_count() + return unittest.skipUnless( + available >= n, f"requires {n} GPUs, found {available}" + )(func) + + return decorator + + +def get_required_world_size(obj: Any, default: int) -> int: + """ + Returns the requested world size for the currently running unittest method on `obj` + if annotated via `@require_world_size(n)`, else returns `default`. + """ + try: + # Try MultiProcessTestCase helper first, then unittest fallback + test_name = ( + obj._current_test_name() # type: ignore[attr-defined] + if hasattr(obj, "_current_test_name") and callable(obj._current_test_name) + else obj._testMethodName + ) + fn = getattr(obj, test_name) + value = fn._required_world_size + return int(value) + except Exception: + return default + + +# This decorator helps avoiding initializing cuda while testing other backends +def nccl_skip_if_lt_x_gpu(backend, x): + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + if backend != "nccl": + return func(*args, **kwargs) + if torch.cuda.is_available() and torch.cuda.device_count() >= x: + return func(*args, **kwargs) + test_skip = TEST_SKIPS[f"multi-gpu-{x}"] + if not _maybe_handle_skip_if_lt_x_gpu(args, test_skip.message): + sys.exit(test_skip.exit_code) + + return wrapper + + return decorator + + +def verify_ddp_error_logged(model_DDP, err_substr): + # Verify error was logged in ddp_logging_data. + ddp_logging_data = model_DDP._get_ddp_logging_data() + if "iteration" not in ddp_logging_data: + raise AssertionError("Expected 'iteration' in ddp_logging_data") + if "has_error" not in ddp_logging_data: + raise AssertionError("Expected 'has_error' in ddp_logging_data") + if "error" not in ddp_logging_data: + raise AssertionError("Expected 'error' in ddp_logging_data") + logging_err = ddp_logging_data["error"] + # Remove C++ stacktrace if needed. + actual = ( + err_substr + if err_substr.find("\nException raised from ") == -1 + else err_substr.split("\nException raised from ")[0] + ) + if actual not in logging_err: + raise AssertionError( + f"Did not find expected {actual} in ddp logging data error: {logging_err}" + ) + + +def with_nccl_blocking_wait(func): + """ + Convenience decorator to set/unset TORCH_NCCL_BLOCKING_WAIT flag. Note that use of + this decorator will override the setting of TORCH_NCCL_ASYNC_ERROR_HANDLING for + the particular test. After the test, both TORCH_NCCL_BLOCKING_WAIT and + TORCH_NCCL_ASYNC_ERROR_HANDLING will be restored to their original values. + """ + + @wraps(func) + def wrapper(*args, **kwargs): + # Save and unset TORCH_NCCL_ASYNC_ERROR_HANDLING + try: + cached_nccl_async_error_handling: str | None = os.environ[ + "TORCH_NCCL_ASYNC_ERROR_HANDLING" + ] + del os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] + except KeyError: + # TORCH_NCCL_ASYNC_ERROR_HANDLING was unset + cached_nccl_async_error_handling = None + + # Save val of TORCH_NCCL_BLOCKING_WAIT and set it. + try: + cached_nccl_blocking_wait: str | None = os.environ[ + "TORCH_NCCL_BLOCKING_WAIT" + ] + except KeyError: + cached_nccl_blocking_wait = None + finally: + os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1" + + try: + ret = func(*args, **kwargs) + return ret + finally: + # restore old values. + if cached_nccl_async_error_handling is not None: + os.environ["TORCH_NCCL_ASYNC_ERROR_HANDLING"] = ( + cached_nccl_async_error_handling + ) + + if cached_nccl_blocking_wait is not None: + os.environ["TORCH_NCCL_BLOCKING_WAIT"] = cached_nccl_blocking_wait + + return wrapper + + +def with_dist_debug_levels(levels): + """ + Runs a test for each distributed debug level specified in levels. + """ + + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + old_level = os.environ.get("TORCH_DISTRIBUTED_DEBUG", None) + for level in levels: + os.environ["TORCH_DISTRIBUTED_DEBUG"] = level + c10d.set_debug_level_from_env() + ret = func(*args, **kwargs) + c10d.barrier() + if old_level is not None: + os.environ["TORCH_DISTRIBUTED_DEBUG"] = old_level + # Only returns test return for last test, but since these are + # unittests the return value is not really used and earlier tests + # would've raised had they failed. + return ret + + return wrapper + + return decorator + + +def requires_gloo(): + return skip_but_pass_in_sandcastle_if( + not c10d.is_gloo_available(), + "c10d was not compiled with the Gloo backend", + ) + + +def requires_nccl_version(version, msg): + if not TEST_CUDA: + return lambda f: f + if not c10d.is_nccl_available(): + return skip_but_pass_in_sandcastle( + "c10d was not compiled with the NCCL backend", + ) + else: + return skip_but_pass_in_sandcastle_if( + torch.cuda.nccl.version() < version, + f"Requires NCCL version greater than or equal to: {version}, found: {torch.cuda.nccl.version()}, reason: {msg}", + ) + + +def requires_nccl_shrink(): + """ + Require NCCL shrink support (NCCL available and version >= 2.27). + """ + return requires_nccl_version((2, 27), "Need NCCL 2.27+ for shrink_group") + + +def requires_nccl(): + return skip_but_pass_in_sandcastle_if( + not c10d.is_nccl_available(), + "c10d was not compiled with the NCCL backend", + ) + + +def requires_ucc(): + return skip_but_pass_in_sandcastle_if( + not c10d.is_ucc_available(), + "c10d was not compiled with the UCC backend", + ) + + +def requires_mpi(): + return skip_but_pass_in_sandcastle_if( + not c10d.is_mpi_available(), + "c10d was not compiled with the MPI backend", + ) + + +def requires_accelerator_dist_backend(backends=None): + """ + Decorator to skip tests if no accelerator communication backend (NCCL, XCCL, HCCL) is available. + + Args: + backends (Optional[List[str]]): Specific accelerator backends to check (e.g., ["nccl", "xccl", "hccl"]). + If None, checks all supported accelerator backends (NCCL, XCCL, HCCL). + + Returns: + callable: A decorator that skips the test if no specified accelerator backend is available. + """ + if backends is None: + backends = ACCELERATOR_DIST_BACKENDS + + backend_available = any( + { + "nccl": c10d.is_nccl_available, + "xccl": c10d.is_xccl_available, + "hccl": lambda: TEST_HPU, + }.get(backend, lambda: False)() + for backend in backends + ) + + return skip_but_pass_in_sandcastle_if( + not backend_available, + f"No accelerator communication backend available among {backends}", + ) + + +def requires_multicast_support(): + has_multicast_support = ( + torch.cuda.is_available() + and _SymmetricMemory.has_multicast_support(DeviceType.CUDA, 0) + ) + return skip_but_pass_in_sandcastle_if( + not has_multicast_support, + "multicast support is not available", + ) + + +def evaluate_platform_supports_symm_mem(): + if TEST_CUDA: + if TEST_WITH_ROCM: + arch_list = ["gfx942", "gfx950"] + for arch in arch_list: + if arch in torch.cuda.get_device_properties(0).gcnArchName: + return True + return False + else: + return True + else: + return False + + +PLATFORM_SUPPORTS_SYMM_MEM: bool = LazyVal( + lambda: evaluate_platform_supports_symm_mem() +) + + +def skip_if_rocm_multiprocess(func): + """Skips a test for ROCm multiprocess UTs""" + return unittest.skipIf(TEST_WITH_ROCM, TEST_SKIPS["skipIfRocm"].message)(func) + + +def skip_if_rocm_arch_multiprocess(arch: tuple[str, ...]): + """Skips a test for given ROCm archs - multiprocess UTs""" + + def decorator(func): + reason = None + if TEST_WITH_ROCM: + prop = torch.cuda.get_device_properties(0).gcnArchName.split(":")[0] + if prop in arch: + reason = f"skip_if_rocm_arch_multiprocess: test skipped on {arch}" + + return unittest.skipIf(reason is not None, reason)(func) + + return decorator + + +def skip_if_rocm_ver_lessthan_multiprocess(version=None): + """Skips a test for ROCm based on ROCm ver - multiprocess UTs""" + + def decorator(func): + reason = None + if TEST_WITH_ROCM: + rocm_version = str(torch.version.hip) + rocm_version = rocm_version.split("-", maxsplit=1)[0] # ignore git sha + rocm_version_tuple = tuple(int(x) for x in rocm_version.split(".")) + if ( + rocm_version_tuple is None + or version is None + or rocm_version_tuple < tuple(version) + ): + reason = f"skip_if_rocm_ver_lessthan_multiprocess: ROCm {rocm_version_tuple} is available but {version} required" + + return unittest.skipIf(reason is not None, reason)(func) + + return decorator + + +def skip_if_win32(): + return skip_but_pass_in_sandcastle_if( + sys.platform == "win32", + "This unit test case is not supported on Windows platform", + ) + + +def sm_is_or_higher_than(device: torch.device, major: int, minor: int) -> bool: + """ + Returns True if the device's compute capability is (major, minor) or higher. + Error out if the device is not a CUDA device. + Returns False if device is a RoCM device. + Returns True if device is a non-CUDA device. + """ + if device.type != "cuda": + return True + + if torch.version.hip is not None: + # ROCm devices may have different compute capability codes + return False + + return torch.cuda.get_device_capability(device) >= (major, minor) + + +@retry_on_connect_failures +def create_tcp_store( + addr="localhost", + world_size=1, + is_master=True, + timeout=timedelta(minutes=5), + wait_for_workers=True, + jit_class=False, + use_libuv=True, +): + """ + Creates a TCP store. Retries if the chosen port is already in use. + """ + port = find_free_port() + if jit_class: + timeout_millisecond = int(timeout / timedelta(milliseconds=1)) + return torch.classes.dist_c10d.TCPStore( + addr, port, world_size, is_master, timeout_millisecond + ) + else: + return c10d.TCPStore( + addr, + port, + world_size, + is_master, + wait_for_workers=wait_for_workers, + use_libuv=use_libuv, + ) + + +if TEST_WITH_TSAN: + # TSAN runs much slower. + TIMEOUT_DEFAULT = 500 +else: + TIMEOUT_DEFAULT = int(os.getenv("DISTRIBUTED_TESTS_DEFAULT_TIMEOUT", "300")) +TIMEOUT_OVERRIDE = { + "test_ddp_uneven_inputs": 400, + "test_DistributedDataParallel": 500, +} + + +# https://github.com/pytorch/pytorch/issues/75665 +if TEST_WITH_ROCM: + TIMEOUT_OVERRIDE["test_join_kwargs"] = 200 + + +def create_device(interface=None, lazy_init: bool = False): + if sys.platform == "win32" or interface is None: + return c10d.ProcessGroupGloo.create_device( + hostname="127.0.0.1", lazy_init=lazy_init + ) + else: + return c10d.ProcessGroupGloo.create_device( + interface=interface, lazy_init=lazy_init + ) + + +def get_timeout(test_id) -> int: + return TIMEOUT_OVERRIDE.get(test_id.split(".")[-1], TIMEOUT_DEFAULT) + + +@contextmanager +def captured_output(): + new_out, new_err = StringIO(), StringIO() + old_out, old_err = sys.stdout, sys.stderr + try: + sys.stdout, sys.stderr = new_out, new_err + yield sys.stdout, sys.stderr + finally: + sys.stdout, sys.stderr = old_out, old_err + + +def simple_sparse_reduce_tests(rank: int, world_size: int, num_inputs: int = 1): + """ + Generate a number of basic test cases for sparse reduction. + These cover tensors with a varying number of sparse dimensions and a varying + number of dense dimensions. The only reduction operation we support is sum. + """ + + def generate(rank: int, world_size: int, sparse_dims: int = 1, dense_dims: int = 0): + # First sparse dimension is [0..rank]. + # Subsequent dimensions are always 0, so we know there is + # a non-empty intersection between any two sparse tensors. + indices = torch.reshape(torch.arange(rank + 1), (1, rank + 1)) + shape = [world_size] + [2 for _ in range(dense_dims)] + for _ in range(sparse_dims - 1): + indices = torch.cat((indices, torch.zeros(1, rank + 1))) + shape.append(world_size) + values = torch.ones([rank + 1] + [2 for _ in range(dense_dims)]) + return torch.sparse_coo_tensor(indices, values, shape) + + def compute_sum(fn, world_size: int): + return reduce( + operator.add, [fn(rank, world_size) for rank in range(world_size)] + ) + + return [ + ( + [ + fn(num_inputs * rank + i, num_inputs * world_size) + for i in range(num_inputs) + ], + [compute_sum(fn, num_inputs * world_size) for i in range(num_inputs)], + ) + for fn in [ + partial(generate, sparse_dims=1), + partial(generate, sparse_dims=2), + partial(generate, sparse_dims=3), + partial(generate, dense_dims=1), + partial(generate, dense_dims=2), + partial(generate, dense_dims=3), + ] + ] + + +# HELPER FOR MULTIGPU TESTS +def init_multigpu_helper(world_size: int, backend: str): + """Multigpu tests are designed to simulate the multi nodes with multi + GPUs on each node. Nccl backend requires equal #GPUs in each process. + On a single node, all visible GPUs are evenly + divided to subsets, each process only uses a subset. + """ + nGPUs = torch.cuda.device_count() + if TEST_HPU: + nGPUs = torch.hpu.device_count() + if TEST_XPU: + nGPUs = torch.xpu.device_count() + visible_devices = range(nGPUs) + + # If rank is less than or equal to number of available GPU's + # then each rank can be mapped to corresponding GPU. + nGPUs_per_process = 1 + if world_size > nGPUs: + nGPUs_per_process = nGPUs // world_size + rank_to_GPU = { + i: list(visible_devices[i * nGPUs_per_process : (i + 1) * nGPUs_per_process]) + for i in range(world_size) + } + return rank_to_GPU + + +tmp_dir: tempfile.TemporaryDirectory | None = None + + +def initialize_temp_directories(init_method: str | None = None) -> None: + global tmp_dir + tmp_dir = tempfile.TemporaryDirectory() + os.environ["TEMP_DIR"] = tmp_dir.name + os.mkdir(os.path.join(tmp_dir.name, "barrier")) + os.mkdir(os.path.join(tmp_dir.name, "test_dir")) + init_dir_path = os.path.join(tmp_dir.name, "init_dir") + os.mkdir(init_dir_path) + # Set init method if specified. + if init_method is not None: + os.environ["INIT_METHOD"] = init_method + else: + os.environ["INIT_METHOD"] = FILE_SCHEMA + os.path.join( + init_dir_path, "shared_init_file" + ) + + +def cleanup_temp_dir() -> None: + if tmp_dir is not None: + tmp_dir.cleanup() + + +def retrieve_result_from_completion_queue( + process: torch.multiprocessing.Process, + completion_queue: torch.multiprocessing.Queue, + timeout: int | None = None, +) -> Any: + """Get result from the completion_queue associated with process. + + When the process finished without putting a result or the timeout expired an exception instance will be returned""" + queue_timeout = 120 if timeout is None else max(10, min(120, timeout // 4)) + start_time = time.time() + # Periodically check the process for liveness + while True: + try: + return completion_queue.get(timeout=queue_timeout) + except queue.Empty: + # If the process is no longer alive we cannot get a result from the queue unless it is there right now. + # This can happen if the timeout occurred just before the process put its result and terminated. + # So do a last check for emptiness before considering it as a failure. + if not process.is_alive() and completion_queue.empty(): + return RuntimeError(f"Exited with {process.exitcode}") + if timeout is not None: + elapsed = time.time() - start_time + if elapsed > timeout: + return RuntimeError(f"Process timed out out after {elapsed}s") + + +# Most tests operate with this worldsize +if TEST_WITH_ROCM: + DEFAULT_WORLD_SIZE = min(4, max(2, torch.cuda.device_count())) +else: + DEFAULT_WORLD_SIZE = 4 + +# [How does MultiProcessTestCase work?] +# Each MultiProcessTestCase instance uses 1 + `world_size()` processes, by +# default `world_size()` returns 4. Let's take `test_rpc_spawn.py` as an +# example which inherits from this class. Its `Setup()` methods calls into +# `MultiProcessTestCase._spawn_processes()` which spawns `world_size()` +# subprocesses. During the spawn, the main process passes the test name to +# subprocesses, and the name is acquired from self.id(). The subprocesses +# then use the provided test function name to retrieve the function attribute +# from the test instance and run it. The main process simply waits for all +# subprocesses to join. + + +class MultiProcessTestCase(TestCase): + MAIN_PROCESS_RANK = -1 + # This exit code is used to indicate that the test code had an error and + # exited abnormally. There are certain tests that might use sys.exit() to + # simulate failures and in those cases, we can't have an exit code of 0, + # but we still want to ensure we didn't run into any other errors. + TEST_ERROR_EXIT_CODE = 10 + + # do not early terminate for distributed tests. + def _should_stop_test_suite(self) -> bool: + return False + + # Many test cases init a process group but do not destroy it. This property + # determines whether this base test class should call + # `destroy_process_group` on behalf of the test. Its value is customizable + # by derived TestCase's but it is a pan-TestCase value (cannot be customized + # for each test). + @property + def destroy_pg_upon_exit(self) -> bool: + return True + + @property + def world_size(self) -> int: + return DEFAULT_WORLD_SIZE + + def join_or_run(self, fn): + @wraps(fn) + def wrapper(self): + if self.rank == self.MAIN_PROCESS_RANK: + self._join_processes(fn) + else: + fn() + + return types.MethodType(wrapper, self) + + # The main process spawns N subprocesses that run the test. + # Constructor patches current instance test method to + # assume the role of the main process and join its subprocesses, + # or run the underlying test function. + def __init__( + self, method_name: str = "runTest", methodName: str = "runTest" + ) -> None: + # methodName is the correct naming in unittest and testslide uses keyword arguments. + # So we need to use both to 1) not break BC and, 2) support testslide. + if methodName != "runTest": + method_name = methodName + super().__init__(method_name) + try: + fn = getattr(self, method_name) + setattr(self, method_name, self.join_or_run(fn)) + except AttributeError as e: + if methodName != "runTest": + # we allow instantiation with no explicit method name + # but not an *incorrect* or missing method name + raise ValueError( + f"no such test method in {self.__class__}: {methodName}" + ) from e + + def setUp(self) -> None: + super().setUp() + + # Used for tests that are expected to return a non-0 exit code, such as + # SIGABRT thrown by watchdog. + self.special_return_code_checks: dict = {} + + # Used for tests that may return any exit code, which makes it hard to + # check. This is rare, use with caution. + self.skip_return_code_checks: list = [] + + self.processes = [] # type: ignore[var-annotated] + self.rank = self.MAIN_PROCESS_RANK + with tempfile.NamedTemporaryFile(delete=False) as f: + self.file_name = f.name + # pid to pipe consisting of error message from process. + self.pid_to_pipe = {} # type: ignore[var-annotated] + + def tearDown(self) -> None: + super().tearDown() + for p in self.processes: + p.terminate() + # Each Process instance holds a few open file descriptors. The unittest + # runner creates a new TestCase instance for each test method and keeps + # it alive until the end of the entire suite. We must thus reset the + # processes to prevent an effective file descriptor leak. + self.processes = [] + + def _current_test_name(self) -> str: + # self.id() == e.g. '__main__.TestDistributed.TestAdditive.test_get_rank' + return self.id().split(".")[-1] + + def _start_processes(self, proc) -> None: + self.processes = [] + for rank in range(int(self.world_size)): + parent_conn, child_conn = torch.multiprocessing.Pipe() + process = proc( + target=self.__class__._run, + name="process " + str(rank), + args=( + rank, + self._current_test_name(), + self.file_name, + child_conn, + ), + kwargs={ + "fake_pg": getattr(self, "fake_pg", False), + }, + ) + process.start() + logger.info("Started process %s with pid %s", rank, process.pid) + self.pid_to_pipe[process.pid] = parent_conn + self.processes.append(process) + + def _spawn_processes(self) -> None: + try: + torch.multiprocessing.set_start_method("spawn") + except RuntimeError: + pass + + proc = torch.multiprocessing.get_context("spawn").Process + self._start_processes(proc) + + class Event(Enum): + GET_TRACEBACK = 1 + + @staticmethod + def _event_listener(parent_pipe, signal_pipe, rank: int): + logger.debug("Starting event listener thread for rank %s", rank) + while True: + ready_pipes = multiprocessing.connection.wait([parent_pipe, signal_pipe]) + + if parent_pipe in ready_pipes: + if parent_pipe.closed: + logger.debug( + "Pipe closed for process %s, stopping event listener thread", + rank, + ) + return + + event = parent_pipe.recv() + logger.info("Received event %s on process %s", event, rank) + + if event == MultiProcessTestCase.Event.GET_TRACEBACK: + # Return traceback to the parent process. + with tempfile.NamedTemporaryFile(mode="r+") as tmp_file: + faulthandler.dump_traceback(tmp_file) + # Flush buffers and seek to read from the beginning + tmp_file.flush() + tmp_file.seek(0) + parent_pipe.send(tmp_file.read()) + + logger.info("Process %s sent traceback", rank) + + if signal_pipe in ready_pipes: + return + + @classmethod + def _run( + cls, rank: int, test_name: str, file_name: str, parent_pipe, **kwargs + ) -> None: + self = cls(test_name) + self.rank = rank + self.file_name = file_name + self.run_test(test_name, parent_pipe) + + def run_test(self, test_name: str, parent_pipe) -> None: + # Start event listener thread. + signal_recv_pipe, signal_send_pipe = torch.multiprocessing.Pipe(duplex=False) + event_listener_thread = threading.Thread( + target=MultiProcessTestCase._event_listener, + args=(parent_pipe, signal_recv_pipe, self.rank), + daemon=True, + ) + event_listener_thread.start() + if sys.platform != "win32" and sys.platform != "darwin": + # Register signal handler to dump stack traces on FATALs. + # Windows and MacOS do not support the signal handlers. + torch._C._set_print_stack_traces_on_fatal_signal(True) + # Show full C++ stacktraces when a Python error originating from C++ is raised. + os.environ["TORCH_SHOW_CPP_STACKTRACES"] = "1" + common_utils.set_rng_seed() + + # self.id() == e.g. '__main__.TestDistributed.test_get_rank' + # We're retrieving a corresponding test and executing it. + try: + getattr(self, test_name)() + except unittest.SkipTest as se: + logger.info( # noqa: G200 + "Process %s skipping test %s for following reason: %s", + self.rank, + test_name, + se, + ) + sys.exit(TEST_SKIPS["generic"].exit_code) + except Exception: + logger.error( + "Caught exception: \n%s exiting process %s with exit code: %s", + traceback.format_exc(), + self.rank, + MultiProcessTestCase.TEST_ERROR_EXIT_CODE, + ) + # Send error to parent process. + parent_pipe.send(traceback.format_exc()) + sys.exit(MultiProcessTestCase.TEST_ERROR_EXIT_CODE) + finally: + if signal_send_pipe is not None: + signal_send_pipe.send(None) + + if event_listener_thread is None: + raise AssertionError("Expected event_listener_thread to not be None") + event_listener_thread.join() + # Close pipe after done with test. + parent_pipe.close() + + if self.destroy_pg_upon_exit: + try: + # Some tests do destroy the pgs, and destroy can't be called twice. + # This avoids spewing warnings about improperly shutting down. + c10d.destroy_process_group() + except (AssertionError, ValueError): + pass + + def _get_timedout_process_traceback(self) -> None: + pipes = [] + for i, process in enumerate(self.processes): + if process.exitcode is None: + pipe = self.pid_to_pipe[process.pid] + try: + pipe.send(MultiProcessTestCase.Event.GET_TRACEBACK) + pipes.append((i, pipe)) + except ConnectionError: + logger.exception( + "Encountered error while trying to get traceback for process %s", + i, + ) + + # Wait for results. + for rank, pipe in pipes: + try: + # Wait for traceback + if pipe.poll(5): + if pipe.closed: + logger.info( + "Pipe closed for process %s, cannot retrieve traceback", + rank, + ) + continue + + traceback = pipe.recv() + logger.error( + "Process %s timed out with traceback: \n\n%s", rank, traceback + ) + else: + logger.error( + "Could not retrieve traceback for timed out process: %s", rank + ) + except ConnectionError: + logger.exception( + "Encountered error while trying to get traceback for process %s", + rank, + ) + + def _join_processes(self, fn) -> None: + timeout = get_timeout(self.id()) + start_time = time.time() + subprocess_error = False + try: + while True: + # check to see if any subprocess exited with an error early. + for i, p in enumerate(self.processes): + # This is the exit code processes exit with if they + # encountered an exception. + if p.exitcode == MultiProcessTestCase.TEST_ERROR_EXIT_CODE: + print( + f"Process {i} terminated with exit code {p.exitcode}, terminating remaining processes." + ) + active_children = torch.multiprocessing.active_children() + for ac in active_children: + ac.terminate() + subprocess_error = True + break + if subprocess_error: + break + # All processes have joined cleanly if they all a valid exitcode + if all(p.exitcode is not None for p in self.processes): + break + # Check if we should time out the test. If so, we terminate each process. + elapsed = time.time() - start_time + if elapsed > timeout: + self._get_timedout_process_traceback() + print( + f"Timing out after {timeout} seconds and killing subprocesses." + ) + for p in self.processes: + p.terminate() + break + # Sleep to avoid excessive busy polling. + time.sleep(0.1) + + elapsed_time = time.time() - start_time + self._check_return_codes(fn, elapsed_time) + finally: + # Close all pipes + for pipe in self.pid_to_pipe.values(): + pipe.close() + + def _check_return_codes(self, fn, elapsed_time) -> None: + """ + Checks that the return codes of all spawned processes match, and skips + tests if they returned a return code indicating a skipping condition. + """ + # If no processes are spawned, there is nothing to check. + if not self.processes: + logger.warning( + "Note: no subprocesses were spawned, test was likely skipped." + ) + return + + first_process = self.processes[0] + # first, we check if there are errors in actual processes + # (via TEST_ERROR_EXIT CODE), and raise an exception for those. + # the reason we do this is to attempt to raise a more helpful error + # message than "Process x terminated/timed out" + # TODO: we should pipe the exception of the failed subprocess here. + # Currently, the actual exception is displayed as a logging output. + errored_processes = [ + (i, p) + for i, p in enumerate(self.processes) + if p.exitcode == MultiProcessTestCase.TEST_ERROR_EXIT_CODE + ] + if errored_processes: + error = "" + for i, process in errored_processes: + # Get error from pipe. + error_message = self.pid_to_pipe[process.pid].recv() + error += ( + f"Process {i} exited with error code {MultiProcessTestCase.TEST_ERROR_EXIT_CODE} " + f"and exception:\n{error_message}\n" + ) + + raise RuntimeError(error) + # If no process exited uncleanly, we check for timeouts, and then ensure + # each process exited cleanly. + for i, p in enumerate(self.processes): + if p.exitcode is None: + raise RuntimeError( + f"Process {i} terminated or timed out after {elapsed_time} seconds" + ) + + # Skip the test return code check + if fn in self.skip_return_code_checks: + return + + for skip in TEST_SKIPS.values(): + if first_process.exitcode == skip.exit_code: + if IS_SANDCASTLE: + # Don't use unittest.skip to skip the test on sandcastle + # since it creates tasks for skipped tests assuming there + # is some follow-up needed. Instead just "pass" the test + # with an appropriate message. + logger.info( + "Skipping %s on sandcastle for the following reason: %s", + self.id(), + skip.message, + ) + return + else: + raise unittest.SkipTest(skip.message) + + # In most cases, we expect test to return exit code 0, standing for success. + expected_return_code = 0 + # In some negative tests, we expect test to return non-zero exit code, + # such as watchdog throwing SIGABRT. + if fn in self.special_return_code_checks: + expected_return_code = self.special_return_code_checks[fn] + + self.assertEqual( + first_process.exitcode, + expected_return_code, + msg=f"Expected exit code {expected_return_code} but got {first_process.exitcode} for pid: {first_process.pid}", + ) + + @property + def is_master(self) -> bool: + return self.rank == 0 + + +# Utility base class for distributed Multi Process Test cases +# This abstracts the PG creation and deletion, the backends are selected based +# on device type. The tests functions can be instantiated per device type using +# common_device_type.instantiate_device_type_tests +# other backends can add entry in backend() function +class DistributedTestBase(MultiProcessTestCase): + def setUp(self): + super().setUp() + os.environ["WORLD_SIZE"] = str(self.world_size) + self._spawn_processes() + + def tearDown(self): + try: + torch.distributed.destroy_process_group() + except AssertionError: + pass + try: + os.remove(self.file_name) + except OSError: + pass + + def backend(self, device) -> str: + if "cuda" in device: + return "nccl" + elif "hpu" in device: # intel gaudi + return "hccl" + elif "xpu" in device: + return "xccl" + else: + return "gloo" + + def create_pg(self, device, world_size=None): + if world_size is None: + world_size = self.world_size + num_visible_devices = torch.get_device_module(device).device_count() + store = torch.distributed.FileStore(self.file_name, num_visible_devices) + torch.distributed.init_process_group( + backend=self.backend(device), + world_size=world_size, + rank=self.rank, + store=store, + ) + if "nccl" in self.backend(device) or "xccl" in self.backend(device): + torch.accelerator.set_device_index(self.rank) + return torch.distributed.distributed_c10d._get_default_group() + + def rank_to_device(self, device): + num_visible_devices = torch.get_device_module(device).device_count() + return {i: [i % num_visible_devices] for i in range(self.world_size)} + + +def run_subtests( + cls_inst, + subtest_config: dict[str, list[Any]], + test_fn: Callable, + *test_args, + **test_kwargs: Any, +): + """ + Runs a test function given by ``test_fn`` as a subtest according to the + configurations specified by ``subtest_config``. This amortizes the + costly setup overhead (including process spawn and initializing the + process group) over the subtests. + + Args: + subtest_config (Dict[str, List[Any]]): A mapping from subtest + keyword argument name to a list of its possible values. + test_fn (Callable): A callable that runs the actual test. + test_args: Positional arguments to pass to ``test_fn``. + test_kwargs: Keyword arguments to pass to ``test_fn``. + """ + # Convert the config mapping to a list to have a fixed order + subtest_config_items: list[tuple[str, list[Any]]] = list(subtest_config.items()) + subtest_config_keys: list[str] = [item[0] for item in subtest_config_items] + subtest_config_values: list[list[Any]] = [item[1] for item in subtest_config_items] + for values in itertools.product(*subtest_config_values): + # Map keyword to chosen value + subtest_kwargs = dict(zip(subtest_config_keys, values, strict=True)) + with cls_inst.subTest(**subtest_kwargs): + torch._dynamo.reset() + test_fn(*test_args, **test_kwargs, **subtest_kwargs) + torch._dynamo.reset() + c10d.barrier() + + +@functools.cache +def has_efa() -> bool: + """ + If shell command `fi_info -p efa -t FI_EP_RDM` returns exit code 0 then we assume that the machine has + Libfabric EFA interfaces and EFA software components installed, + see https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/efa-start.html. + """ + + try: + return ( + subprocess.run( + ["fi_info", "-p", "efa", "-t", "FI_EP_RDM"], check=False + ).returncode + == 0 + ) + except FileNotFoundError: + pass + return False + + +def tp_transports(): + """ + If the machine has Libfabric EFA interfaces and EFA software components installed it may cause + 'RuntimeError: In operator() at tensorpipe/common/ibv.h:172 "": Operation not supported' if tensorpipe + uses InfiniBand transport, so we exclude it from tensorpipe transports, + see https://github.com/pytorch/pytorch/issues/73885 and https://github.com/pytorch/pytorch/issues/65022 + """ + return ["shm", "uv"] if has_efa() else None + + +def spawn_threads_and_init_comms( + func=None, timeout=TIMEOUT_DEFAULT, world_size=DEFAULT_WORLD_SIZE +): + """ + Wrapper to use with a test method + """ + if func is None: + return partial( + spawn_threads_and_init_comms, timeout=timeout, world_size=world_size + ) + + def _run_test_method_with_multi_threads(world_size, callback): + world = _install_threaded_pg() + global_store = c10d.HashStore() + + def world_is_valid(): + return world == c10d.distributed_c10d._world + + def worker(rank, world_pg, store): + c10d.init_process_group( + backend="threaded", rank=rank, world_size=world_size, store=store + ) + try: + callback() + except BaseException as ex: # noqa: B036 + # Exceptions are handled in MultiThreadedTestCase + MultiThreadedTestCase.exception_queue.put((rank, sys.exc_info())) + ProcessLocalGroup.exception_handle( + ex + ) # trigger _terminate event and awaken worker threads + finally: + if world_is_valid(): + c10d.destroy_process_group() + + threads = [] + for rank in range(world_size): + t = threading.Thread(target=worker, args=(rank, world, global_store)) + t.start() + threads.append(t) + + return threads + + @wraps(func) + def wrapper(self, *args, **kwargs): + # TODO: get test name from kwargs + torch._C._distributed_c10d._set_thread_isolation_mode(True) + try: + threads = _run_test_method_with_multi_threads( + world_size, lambda: func(self, *args, **kwargs) + ) + # join and error handling + MultiThreadedTestCase._join_threads(threads, func) + finally: + torch._C._distributed_c10d._set_thread_isolation_mode(False) + + return wrapper + + +class MultiThreadedTestCase(TestCase): + """ + Test runner that runs all tests with the in-proc process group using + multiple threads with the threaded process group. + + Each test spawns world_size threads and run the test method in each thread. + + Difference from regular MultiProcess test runner: + Must explicitly defines SetUp and call self._spawn_threads() to run the tests. + Cannot use setUp / tearDown (must use perThreadSetup / perThreadShutdown) + to set up / tear down each thread when running each test. + No global state possible + How bad of a limitation is this? + """ + + exception_queue = queue.Queue() + + MAIN_THREAD_RANK = -1 + + def join_or_run(self, fn): + @wraps(fn) + def wrapper(self): + if self.rank == self.MAIN_THREAD_RANK: + self._join_threads(self.threads, fn) + else: + fn() + + return types.MethodType(wrapper, self) + + def __init__( + self, method_name: str = "runTest", methodName: str = "runTest" + ) -> None: + # methodName is the correct naming in unittest and testslide uses keyword arguments. + # So we need to use both to 1) not break BC and, 2) support testslide. + if methodName != "runTest": + method_name = methodName + super().__init__(method_name) + try: + fn = getattr(self, method_name) + setattr(self, method_name, self.join_or_run(fn)) + except AttributeError as e: + if methodName != "runTest": + # we allow instantiation with no explicit method name + # but not an *incorrect* or missing method name + raise ValueError( + f"no such test method in {self.__class__}: {methodName}" + ) from e + + def perThreadSetUp(self): + # super().setUp() # TestCase.setUp() calls torch.manual_seed() + pass + + def perThreadTearDown(self): + pass + + def setUp(self) -> None: + """ + setUp only set up things in the main thread, if you want to configure things + in the spawned threads, use perThreadSetUp + """ + super().setUp() + self.rank = self.MAIN_THREAD_RANK + self.threads = [] + # Show full C++ stacktraces when a Python error originating from C++ is raised. + os.environ["TORCH_SHOW_CPP_STACKTRACES"] = "1" + + def tearDown(self): + """ + tearDown only set up things in the main thread, if you want to configure things + in the spawned threads, use perThreadTearDown + """ + super().tearDown() + self.threads = [] + + def _spawn_threads(self): + """ + class method to spawn threads and run test, use this method in the SetUp of your TestCase + """ + torch._C._distributed_c10d._set_thread_isolation_mode(True) + test_name = self._current_test_name + # for each test case, we need to create thread local world, and a global store + world = _install_threaded_pg() + self.__class__.global_store = c10d.HashStore() + + def world_is_valid(): + return world == c10d.distributed_c10d._world + + if not world_is_valid(): + raise RuntimeError("Invalid world") + + for rank in range(self.world_size): + t = threading.Thread( + target=self.__class__._run, args=(test_name, rank, self.world_size) + ) + t.start() + self.threads.append(t) + + @classmethod + def _run(cls, test_name, rank, world_size, **kwargs): + self = cls(test_name) + self.rank = rank + + # precision/rel_tol is a thread-local setting since it may be overridden per test, need to make + # every thread have the same value. This would be relevant when we use op db tests, where it + # needs those states to be set i.e. using instantiate_device_type_tests() + # TODO: figure out a better way to do this + if hasattr(self, "_tls"): + self._tls = threading.local() + self._tls.precision = TestCase._precision + self._tls.rel_tol = TestCase._rel_tol + + self.run_test_with_threaded_pg(test_name, rank, world_size) + + def run_test_with_threaded_pg(self, test_name, rank, world_size): + """ + Run the current test associated with `test_name` using the threaded process group. + """ + c10d.init_process_group( + backend="threaded", + rank=rank, + world_size=world_size, + store=self.__class__.global_store, + ) + self.perThreadSetUp() + + try: + getattr(self, test_name)() + except BaseException as ex: # noqa: B036 + self.exception_queue.put((rank, sys.exc_info())) + ProcessLocalGroup.exception_handle( + ex + ) # trigger _terminate event and awaken worker threads + finally: + c10d.destroy_process_group() + self.perThreadTearDown() + + @classmethod + def _join_threads(cls, threads, fn): + timeout = TIMEOUT_DEFAULT + try: + for idx, thread in enumerate(threads): + thread.join(max(0, timeout)) + if thread.is_alive(): + MultiThreadedTestCase.exception_queue.put( + ( + idx, + ( + TimeoutError, + TimeoutError( + f"Rank failed to join in under {timeout} seconds" + ), + None, + ), + ) + ) + ProcessLocalGroup.reset() + failed_ranks = [] + while not cls.exception_queue.empty(): + failure = cls.exception_queue.get() + failed_ranks.append(failure) + finally: + _uninstall_threaded_pg() + torch._C._distributed_c10d._set_thread_isolation_mode(False) + + cls._check_return_codes(failed_ranks, timeout, fn) + + @classmethod + def _check_return_codes(cls, failed_ranks, timeout, fn): + # Print based on exceptions raised from threads + # SkipTest: print info for each thread + # TimeoutError: raise RuntimeError for any timed out thread + # Normal Exception: print error for each thread that raises exception + # and raise a RuntimeError + error_msg = "" + skip_code = -1 + for rank, exc_info in failed_ranks: + exc = exc_info[1] + if isinstance(exc, unittest.SkipTest): + logger.info( + "Thread %s skipping test %s for following reason: %s", + rank, + fn, + exc, + ) + if skip_code < 0: + skip_code = TEST_SKIPS["generic"].exit_code + elif isinstance(exc, TimeoutError): + msg = f"Thread {rank} terminated or timed out after {timeout} seconds\n" + logger.error(msg) + raise RuntimeError(msg) + elif isinstance(exc, Exception): + msg = "".join(traceback.format_exception(*exc_info)) + logger.error("Caught exception: \n%s exiting thread %s", msg, rank) + error_msg += f"Thread {rank} exited with exception:\n{msg}\n" + elif isinstance(exc, SystemExit): + if type(exc.code) is int and skip_code < 0: + skip_code = exc.code + + # check exceptions + if len(error_msg) > 0: + raise RuntimeError(error_msg) + # check skip + if skip_code > 0: + for skip in TEST_SKIPS.values(): + if skip_code == skip.exit_code: + if IS_SANDCASTLE: + # "pass" the test with an appropriate message. + logger.info( + "Skipping %s on sandcastle for the following reason: %s", + fn, + skip.message, + ) + return + else: + raise unittest.SkipTest(skip.message) + + @property + def world_size(self) -> int: + return DEFAULT_WORLD_SIZE + + @property + def _current_test_name(self) -> str: + # self.id() == e.g. '__main__.TestDistributed.TestAdditive.test_get_rank' + return self.id().split(".")[-1] + + def assertEqualOnRank(self, x, y, msg=None, *, rank=0): + """ + The reason why we have this util function instead of + self.assertEqual is all threads are sharing one CPU RNG + so the assertion result is only reliable on rank 0 + """ + if self.rank == rank: + self.assertEqual(x, y, msg) + + def assertNotEqualOnRank(self, x, y, msg=None, *, rank=0): + if self.rank == rank: + self.assertNotEqual(x, y) + + +class SaveForwardInputsModule(nn.Module): + def __init__( + self, + forward_inputs: dict[nn.Module, torch.Tensor], + cast_forward_inputs: bool, + ) -> None: + super().__init__() + self.l = nn.Linear(100, 100) + self.forward_inputs = forward_inputs + self.cast_forward_inputs = cast_forward_inputs + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.forward_inputs[self] = x + return self.l(x.to(self.l.weight.dtype) if self.cast_forward_inputs else x) + + +class SaveForwardInputsModel(nn.Module): + def __init__( + self, + forward_inputs: dict[nn.Module, torch.Tensor], + cast_forward_inputs: bool, + ) -> None: + super().__init__() + self.c1 = SaveForwardInputsModule(forward_inputs, cast_forward_inputs) + self.c2 = SaveForwardInputsModule(forward_inputs, cast_forward_inputs) + self.forward_inputs = forward_inputs + + def forward(self, x: torch.Tensor) -> torch.Tensor: + self.forward_inputs[self] = x + return self.c2(self.c1(x)) + + +@contextmanager +def _dynamo_dist_per_rank_init( + rank, world_size, backend=None, init_pg=True, fake_pg=False +): + # To avoid multiple inheritance from _dynamo.test_case.TestCase and MultiProcessTestCase, + # Just manually implement the most important part of the dynamo behavior to reset/clear. + if not fake_pg: + torch.accelerator.set_device_index(rank) + + device_type = ( + acc.type if (acc := torch.accelerator.current_accelerator()) else "cpu" + ) + if backend is None: + backend = c10d.get_default_backend_for_device(device_type) + + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = "6789" + if init_pg: + if fake_pg: + store = torch.testing._internal.distributed.fake_pg.FakeStore() + c10d.init_process_group( + backend="fake", + world_size=world_size, + rank=rank, + store=store, + ) + else: + c10d.init_process_group(backend=backend, rank=rank, world_size=world_size) + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + try: + yield + finally: + torch._dynamo.reset() + torch._dynamo.utils.counters.clear() + if init_pg: + c10d.destroy_process_group() + + +class DynamoDistributedSingleProcTestCase(torch._dynamo.test_case.TestCase): + """ + Test harness for single-process dynamo distributed tests, + initializes dist process group. + + Prefer this for simple tests, as it's easier to debug. + """ + + @classmethod + def setUpClass(cls): + super().setUpClass() + # _exit_stack is set up in TestCase + cls._exit_stack.enter_context( + patch.dict( + os.environ, + { + "MASTER_ADDR": "localhost", + "MASTER_PORT": "12355", + }, + ) + ) + cls.rank = 0 + device = torch.accelerator.current_accelerator().type + cls.device = f"{device}:{cls.rank}" + cls.device_ids = None if device in cls.device else [cls.rank] + c10d.init_process_group( + c10d.get_default_backend_for_device(device), rank=cls.rank, world_size=1 + ) + + @classmethod + def tearDownClass(cls): + c10d.destroy_process_group() + super().tearDownClass() + + +class DynamoDistributedMultiProcTestCase(DistributedTestBase): + """ + Use this for tests that actually run on multiple GPUs. + + Decorate tests with @skip_if_lt_x_gpu(ngpu) + + Note: MultiProcTestCase spawns processes per test and is slow. + Prefer MultiThreadedTestCase for most tests. Perhaps use this one + sparingly for integration tests. + """ + + @property + def world_size(self) -> int: + return torch.accelerator.device_count() + + @classmethod + def _run( + cls, rank: int, test_name: str, file_name: str, parent_pipe, **kwargs + ) -> None: + trace_log.addHandler(logging.NullHandler()) + + # The rest is copypasta from MultiProcessTestCase._run + self = cls(test_name) + self.rank = rank + self.file_name = file_name + self.run_test(test_name, parent_pipe) + + +class MultiProcContinuousTest(TestCase): + # Class variables: + MAIN_PROCESS_RANK = -1 + # number of test processes + world_size: int = -2 # unset state + # rank of the current process + rank: int = -2 # unset state + # Rendezvous file + rdvz_file: str | None = None + # timeout configured per class + timeout: timedelta = timedelta(seconds=120) + # Poison pill for rest of tests if one of them fails + poison_pill: bool = False + # Flag for lazy process spawning (to support instantiate_device_type_tests) + _processes_spawned: bool = False + + @classmethod + def backend_str(cls) -> str | None: + """ + ProcessGroup backend str. + To be customized by sub test classes, e.g. "nccl". + Otherwise we return None -- lazily decided by tensor. + """ + return None + + # Please override if you intend to test on specific device type + @classmethod + def device_type(cls) -> str: + curr_device = torch.accelerator.current_accelerator() + if curr_device is None: + return "cpu" + return curr_device.type + + @classmethod + def opts(cls, high_priority_stream=False): + """ + ProcessGroup init options. + To be customized by sub test classes, e.g. ProcessGroupNCCLOpTest + Here we return None. + """ + return None + + @classmethod + def _init_pg(cls, rank, world_size, rdvz_file): + if rdvz_file is None: + raise AssertionError("Expected rdvz_file to not be None") + # rank should be local_rank for tests running on <= 8 gpus which is how all these tests are designed + # and we expect LOCAL_RANK set by torchrun. Setting it lets init_device_mesh set the device without + # issuing a warning + os.environ["LOCAL_RANK"] = str(rank) + store = c10d.FileStore(rdvz_file, world_size) + # create nccl processgroup with opts + c10d.init_process_group( + backend=cls.backend_str(), + world_size=world_size, + rank=rank, + store=store, + pg_options=cls.opts(), + timeout=cls.timeout, + ) + cls.pg = c10d.distributed_c10d._get_default_group() + + @classmethod + def _run_test_given_id(cls, test_id: str, **kwargs) -> None: + # self.id() == e.g. '__main__.TestDistributed.TestAdditive.test_get_rank' + test_name = test_id.rsplit(".", maxsplit=1)[-1] + # Get the test function from the test class + self = cls(test_name) + self.rank = cls.rank + self.world_size = cls.world_size + test_fn = getattr(self, test_name) + + # Ensure all the ranks use the same seed. + common_utils.set_rng_seed() + + # Run the test function + test_fn(**kwargs) + + @classmethod + def _worker_loop(cls, rank, world_size, rdvz_file, task_queue, completion_queue): + raised_exception = False + # Sub tests are going to access these values, check first + if not (0 <= rank < world_size): + raise AssertionError( + f"Expected 0 <= rank < world_size, got rank={rank}, world_size={world_size}" + ) + # set class variables for the test class + cls.rank = rank + cls.world_size = world_size + + # Initialize the process group + init_skip_reason = None + try: + cls._init_pg(rank, world_size, rdvz_file) + except SystemExit as ex: + exit_code = getattr(ex, "code", None) + skip_entry = next( + (v for v in TEST_SKIPS.values() if v.exit_code == exit_code), + None, + ) + if skip_entry: + init_skip_reason = skip_entry.message + else: + raise + + # End of bootstrap + logger.debug("Setup complete") + + # Loop forever, waiting for a test name to run + while True: + test_id = task_queue.get() + logger.debug(f"Got test {test_id}") # noqa: G004 + # None means exit + if test_id is None: + break + + # If init failed with a skip, respond with SkipTest for all tests + if init_skip_reason is not None: + completion_queue.put(unittest.SkipTest(init_skip_reason)) + continue + + # Run the test + try: + cls._run_test_given_id(test_id) + completion_queue.put(test_id) + except BaseException as ex: # noqa: B036 + if isinstance(ex, SystemExit): + # Get exit code from the process + exit_code = getattr(ex, "code", None) + + # Look up exit code in TEST_SKIPS to see if it is a valid skip + skip_entry = next( + (v for v in TEST_SKIPS.values() if v.exit_code == exit_code), + None, + ) + + # If we found an entry, we want to skip the test and the object back to the main process + if skip_entry: + completion_queue.put(unittest.SkipTest(skip_entry.message)) + # Skip exception handling below, move to main thread for processing the skip + continue + + raised_exception = True + # Send the exception and stack trace back to the dispatcher + exc_info = sys.exc_info() + tb_str = "".join(traceback.format_exception(*exc_info)) + # Create a new exception with the original exception and traceback + enhanced_ex = RuntimeError(f"Exception in worker process:\n{tb_str}") + enhanced_ex.__cause__ = ex + completion_queue.put(enhanced_ex) + + # Termination + logger.debug("Terminating ...") + # Calling destroy_process_group when workers have exceptions + # while others are doing collectives will cause a deadlock since + # it waits for enqueued collectives to finish. + # Only call this on a clean exit path + if not raised_exception: + c10d.destroy_process_group() + + @classmethod + def _spawn_processes(cls, world_size) -> None: + cls.processes = [] + cls.task_queues = [] + cls.completion_queues = [] + # Need a rendezvous file for `init_process_group` purpose. + with tempfile.NamedTemporaryFile(delete=False) as f: + cls.rdvz_file = f.name + + # CUDA multiprocessing requires spawn instead of fork, to make sure + # child processes have their own memory space. + try: + torch.multiprocessing.set_start_method("spawn") + except RuntimeError: + # The start method has already been set + pass + + for rank in range(int(world_size)): + task_queue = torch.multiprocessing.Queue() + completion_queue = torch.multiprocessing.Queue() + process = torch.multiprocessing.Process( + target=cls._worker_loop, + name="process " + str(rank), + daemon=True, # so that child processes will exit if parent decides to terminate + args=(rank, world_size, cls.rdvz_file, task_queue, completion_queue), + ) + process.start() + cls.processes.append(process) + cls.task_queues.append(task_queue) + cls.completion_queues.append(completion_queue) + logger.debug("Started process %s with pid %s", rank, process.pid) # noqa: UP031 + + @classmethod + def _get_world_size(cls, device_type: str) -> int: + """ + Get world_size, handling both class variable and property definitions. + Properties are instance-level and need special handling in class methods. + """ + # Check if world_size is defined as a property (instance-level) + world_size_attr = inspect.getattr_static(cls, "world_size", None) + if isinstance(world_size_attr, property): + # Create a temporary instance to evaluate the property + # We use object.__new__ to avoid calling __init__ which may have side effects + temp_instance = object.__new__(cls) + world_size = world_size_attr.fget(temp_instance) + else: + world_size = cls.world_size + + # If world_size is not set (== -2), use device count + if world_size == -2: + world_size = torch.get_device_module(device_type).device_count() + if world_size == 0: + raise unittest.SkipTest(f"No {device_type} devices available") + + return world_size + + @classmethod + def setUpClass(cls): + """ + Class-scope test fixture. Run once for entire test class, before any test starts. + Note: Process spawning is deferred to setUp to support instantiate_device_type_tests, + which calls setUpClass during class creation before any tests run. + """ + super().setUpClass() + + @classmethod + def _ensure_processes_spawned(cls): + """ + Lazily spawn worker processes on first test run. + This supports instantiate_device_type_tests which calls setUpClass during + class creation (before any tests run), when spawning would be premature. + """ + if cls._processes_spawned: + return + + # Handle method, property, and string attribute for device_type + # (instantiate_device_type_tests sets device_type as a string attribute, + # making this compatible as a drop-in replacement for MultiProcessTestCase) + device_type_attr = cls.__dict__.get("device_type", cls.device_type) + if isinstance(device_type_attr, classmethod): + device_type = device_type_attr.__func__(cls) + elif isinstance(device_type_attr, property): + # Note: fget expects an instance but we pass cls since no instance + # exists yet. This works because DTensorTestMixin.device_type only + # accesses class-level attributes (world_size, module constants). + device_type = device_type_attr.fget(cls) + elif callable(device_type_attr): + device_type = device_type_attr() + else: + device_type = device_type_attr + + # Get world_size (handles both class variable and property) + cls.world_size = cls._get_world_size(device_type) + + # Check if the specified backend is available before spawning processes + backend = cls.backend_str() if callable(cls.backend_str) else cls.backend_str + if backend is not None: + backend_checks = { + "nccl": c10d.is_nccl_available, + "gloo": c10d.is_gloo_available, + "mpi": c10d.is_mpi_available, + "xccl": c10d.is_xccl_available, + } + check_fn = backend_checks.get(backend) + if check_fn is not None and not check_fn(): + raise unittest.SkipTest(f"Backend '{backend}' is not available") + + logger.info( + f"Testing class {cls.__name__} on {cls.world_size} {device_type}" # noqa: G004 + ) + + cls._spawn_processes(cls.world_size) + cls._processes_spawned = True + + @classmethod + def tearDownClass(cls): + """ + Class-scope test fixture. Run once for entire test class, after all tests finish. + Tear down the process group if spawned. + """ + # If processes were never spawned (e.g., all tests were skipped), nothing to tear down + if not cls._processes_spawned: + super().tearDownClass() + return + + logger.debug(f"Joining {cls.world_size} workers") # noqa: G004 + # Enqueue "None" to all workers to tell them to exit + for task_queue in cls.task_queues: + task_queue.put(None) + + # Wait for all workers to exit + for process in cls.processes: + process.join() + + # Clear up the rendezvous file + try: + os.remove(cls.rdvz_file) + except OSError: + pass + + logger.info(f"Class {cls.__name__} finished") # noqa: G004 + super().tearDownClass() + + def setUp(self) -> None: + """ + Test fixture. Run before each test. + """ + super().setUp() + + # Ensure processes are spawned (lazy initialization for instantiate_device_type_tests) + self.__class__._ensure_processes_spawned() + + # I am the dispatcher + self.rank = self.MAIN_PROCESS_RANK + + # If this test class hits an exception in one test, skip the rest of tests + if self.__class__.poison_pill: + raise unittest.SkipTest(f"Previous test failed, skipping {self.id()}") + + # Enqueue "current test" to all workers + for i, task_queue in enumerate(self.task_queues): + logger.debug(f"Sending Rank {i}: {self.id()}") # noqa: G004 + task_queue.put(self.id()) + + def _worker_run_main_wait(self, fn): + @wraps(fn) + def wrapper(self): + if self.rank == self.MAIN_PROCESS_RANK: + logger.debug(f"Waiting for workers to finish {self.id()}") # noqa: G004 + # Drain all completion queues before raising any exception, + # so stale results don't desync subsequent tests. + deferred_exception = None + for i, (p, completion_queue) in enumerate( + zip(self.processes, self.completion_queues) + ): + rv = retrieve_result_from_completion_queue( + p, completion_queue, timeout=get_timeout(self.id()) + ) + if deferred_exception is not None: + # Already captured an exception; just drain + continue + if isinstance(rv, unittest.SkipTest): + deferred_exception = rv + continue + if isinstance(rv, BaseException): + logger.warning( + f"Detected failure from Rank {i} in: {self.id()}, " # noqa: G004 + f"skipping rest of tests in Test class: {self.__class__.__name__}" # noqa: G004 + ) + self.__class__.poison_pill = True + deferred_exception = rv + continue + + # Success + if rv != self.id(): + raise AssertionError( + f"Expected rv == self.id(), got {rv} != {self.id()}" + ) + logger.debug( + f"Main proc detected rank {i} finished {self.id()}" # noqa: G004 + ) + + if deferred_exception is not None: + raise deferred_exception + else: + # Worker just runs the test + fn() + + return types.MethodType(wrapper, self) + + # The main process spawns N subprocesses that run the test. + # Constructor patches current instance test method to + # assume the role of the main process and join its subprocesses, + # or run the underlying test function. + def __init__( + self, method_name: str = "runTest", methodName: str = "runTest" + ) -> None: + # methodName is the correct naming in unittest and testslide uses keyword arguments. + # So we need to use both to 1) not break BC and, 2) support testslide. + if methodName != "runTest": + method_name = methodName + super().__init__(method_name) + try: + fn = getattr(self, method_name) + setattr(self, method_name, self._worker_run_main_wait(fn)) + except AttributeError as e: + if methodName != "runTest": + # we allow instantiation with no explicit method name + # but not an *incorrect* or missing method name + raise ValueError( + f"no such test method in {self.__class__}: {methodName}" + ) from e diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_dtype.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_dtype.py new file mode 100644 index 0000000000000000000000000000000000000000..6e62009b6b6013da46ba98a7ee1d121a113dd54a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_dtype.py @@ -0,0 +1,243 @@ +# mypy: ignore-errors + + +import torch + + +# Functions and classes for describing the dtypes a function supports +# NOTE: these helpers should correspond to PyTorch's C++ dispatch macros + + +# Verifies each given dtype is a torch.dtype +def _validate_dtypes(*dtypes): + for dtype in dtypes: + if not isinstance(dtype, torch.dtype): + raise AssertionError(f"Expected dtype to be torch.dtype, got {type(dtype)}") + return dtypes + + +# class for tuples corresponding to a PyTorch dispatch macro +class _dispatch_dtypes(tuple): + __slots__ = () + + def __add__(self, other): + if not isinstance(other, tuple): + raise AssertionError(f"Expected other to be a tuple, got {type(other)}") + return _dispatch_dtypes(tuple.__add__(self, other)) + + +_empty_types = _dispatch_dtypes(()) + + +def empty_types(): + return _empty_types + + +_floating_types = _dispatch_dtypes((torch.float32, torch.float64)) + + +def floating_types(): + return _floating_types + + +_floating_types_and_half = _floating_types + (torch.half,) + + +def floating_types_and_half(): + return _floating_types_and_half + + +def floating_types_and(*dtypes): + return _floating_types + _validate_dtypes(*dtypes) + + +_floating_and_complex_types = _floating_types + (torch.cfloat, torch.cdouble) + + +def floating_and_complex_types(): + return _floating_and_complex_types + + +def floating_and_complex_types_and(*dtypes): + return _floating_and_complex_types + _validate_dtypes(*dtypes) + + +_double_types = _dispatch_dtypes((torch.float64, torch.complex128)) + + +def double_types(): + return _double_types + + +# NB: Does not contain uint16/uint32/uint64 for BC reasons +_integral_types = _dispatch_dtypes( + (torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64) +) + + +def integral_types(): + return _integral_types + + +def integral_types_and(*dtypes): + return _integral_types + _validate_dtypes(*dtypes) + + +_all_types = _floating_types + _integral_types + + +def all_types(): + return _all_types + + +def all_types_and(*dtypes): + return _all_types + _validate_dtypes(*dtypes) + + +_complex_types = _dispatch_dtypes((torch.cfloat, torch.cdouble)) + + +def complex_types(): + return _complex_types + + +def complex_types_and(*dtypes): + return _complex_types + _validate_dtypes(*dtypes) + + +_all_types_and_complex = _all_types + _complex_types + + +def all_types_and_complex(): + return _all_types_and_complex + + +def all_types_and_complex_and(*dtypes): + return _all_types_and_complex + _validate_dtypes(*dtypes) + + +_all_types_and_half = _all_types + (torch.half,) + + +def all_types_and_half(): + return _all_types_and_half + + +_all_mps_types = ( + _dispatch_dtypes({torch.float, torch.half, torch.bfloat16}) + _integral_types +) + + +def all_mps_types(): + return _all_mps_types + + +def all_mps_types_and(*dtypes): + return _all_mps_types + _validate_dtypes(*dtypes) + + +_float8_types = _dispatch_dtypes( + ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + torch.float8_e5m2, + torch.float8_e5m2fnuz, + ) +) + + +def float8_types(): + return _float8_types + + +def float8_types_and(*dtypes): + return _float8_types + _validate_dtypes(*dtypes) + + +def all_types_complex_float8_and(*dtypes): + return _all_types + _complex_types + _float8_types + _validate_dtypes(*dtypes) + + +def custom_types(*dtypes): + """Create a list of arbitrary dtypes""" + return _empty_types + _validate_dtypes(*dtypes) + + +# The functions below are used for convenience in our test suite and thus have no corresponding C++ dispatch macro + + +# See AT_FORALL_SCALAR_TYPES_WITH_COMPLEX_AND_QINTS. +def get_all_dtypes( + include_half=True, + include_bfloat16=True, + include_bool=True, + include_complex=True, + include_complex32=False, + include_qint=False, +) -> list[torch.dtype]: + dtypes = get_all_int_dtypes() + get_all_fp_dtypes( + include_half=include_half, include_bfloat16=include_bfloat16 + ) + if include_bool: + dtypes.append(torch.bool) + if include_complex: + dtypes += get_all_complex_dtypes(include_complex32) + if include_qint: + dtypes += get_all_qint_dtypes() + return dtypes + + +def get_all_math_dtypes(device) -> list[torch.dtype]: + return ( + get_all_int_dtypes() + + get_all_fp_dtypes( + include_half=device.startswith("cuda"), include_bfloat16=False + ) + + get_all_complex_dtypes() + ) + + +def get_all_complex_dtypes(include_complex32=False) -> list[torch.dtype]: + return ( + [torch.complex32, torch.complex64, torch.complex128] + if include_complex32 + else [torch.complex64, torch.complex128] + ) + + +def get_all_int_dtypes() -> list[torch.dtype]: + return [torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64] + + +def get_all_fp_dtypes(include_half=True, include_bfloat16=True) -> list[torch.dtype]: + dtypes = [torch.float32, torch.float64] + if include_half: + dtypes.append(torch.float16) + if include_bfloat16: + dtypes.append(torch.bfloat16) + return dtypes + + +def get_all_qint_dtypes() -> list[torch.dtype]: + return [torch.qint8, torch.quint8, torch.qint32, torch.quint4x2, torch.quint2x4] + + +def highest_precision_float(device): + if torch.device(device).type == "mps": + return torch.float32 + else: + return torch.float64 + + +def highest_precision_complex(device): + if torch.device(device).type == "mps": + return torch.complex64 + else: + return torch.complex128 + + +float_to_corresponding_complex_type_map = { + torch.float16: torch.complex32, + torch.float32: torch.complex64, + torch.float64: torch.complex128, +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_fsdp.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_fsdp.py new file mode 100644 index 0000000000000000000000000000000000000000..cebeb1a4c01971bbcb3dffcd6499e3a3e57bfe86 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_fsdp.py @@ -0,0 +1,1751 @@ +# mypy: allow-untyped-defs +# Owner(s): ["oncall: distributed"] + +import contextlib +import os +import re +import sys +import time +import unittest +import warnings +from abc import ABC, abstractmethod +from collections.abc import Callable +from contextlib import nullcontext +from copy import deepcopy +from enum import auto, Enum +from functools import wraps +from typing import Any, cast, no_type_check +from unittest import mock + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed._composable import checkpoint +from torch.distributed.device_mesh import DeviceMesh +from torch.distributed.fsdp import ( + CPUOffload, + fully_shard, + FullyShardedDataParallel as FSDP, +) +from torch.distributed.fsdp._common_utils import TrainingState +from torch.distributed.fsdp._fully_shard._fsdp_param_group import ( + FSDPParamGroup, + RegisterPostBackwardFunction, +) +from torch.distributed.fsdp._init_utils import NO_RESHARD_AFTER_FORWARD_STRATEGIES +from torch.distributed.fsdp.fully_sharded_data_parallel import ( + BackwardPrefetch, + MixedPrecision, + ShardingStrategy, +) +from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler +from torch.distributed.fsdp.wrap import always_wrap_policy, ModuleWrapPolicy, wrap +from torch.distributed.tensor import distribute_tensor, DTensor, Shard +from torch.distributed.tensor.parallel import ( + ColwiseParallel, + parallelize_module, + RowwiseParallel, + SequenceParallel, +) +from torch.nn import TransformerDecoderLayer, TransformerEncoderLayer +from torch.nn.parallel.distributed import DistributedDataParallel as DDP +from torch.testing._internal.common_distributed import ( + MultiProcContinuousTest, + MultiProcessTestCase, + MultiThreadedTestCase, + run_subtests, + TEST_SKIPS, +) +from torch.testing._internal.common_utils import ( + FILE_SCHEMA, + get_cycles_per_ms, + set_rng_seed, + TEST_CUDA, + TEST_HPU, + TEST_WITH_ROCM, + TEST_XPU, +) +from torch.utils._triton import has_triton + + +if TEST_WITH_ROCM: + DEVICE_COUNT = min(4, max(2, torch.cuda.device_count())) +else: + DEVICE_COUNT = 4 + +if TEST_CUDA: + DEVICE_TYPE = "cuda" + DISTRIBUTED_BACKEND = "nccl" + DEVICE_COUNT = torch.cuda.device_count() +elif TEST_HPU: + DEVICE_TYPE = "hpu:0" + DISTRIBUTED_BACKEND = "hccl" +elif TEST_XPU: + DEVICE_TYPE = "xpu" + DISTRIBUTED_BACKEND = "xccl" + DEVICE_COUNT = torch.xpu.device_count() +else: + DEVICE_TYPE = "cpu" + DISTRIBUTED_BACKEND = "gloo" + DEVICE_COUNT = 1 + + +class FSDPInitMode(Enum): + # No FSDP wrapping + NO_FSDP = auto() + # FSDP recursive wrapping + RECURSIVE = auto() + # TODO: FSDP non-recursive wrapping + # NONRECURSIVE = auto() + + +class DEVICEInitMode(Enum): + # Move model to DEVICE before passing to the FSDP constructor + DEVICE_BEFORE = auto() + # Move model to DEVICE after passing to the FSDP constructor + DEVICE_AFTER = auto() + # Keep on CPU + DEVICE_NEVER = auto() + + +class FSDPTestModel(nn.Module, ABC): + """This defines the interface expected from all models used commonly for + FSDP unit tests.""" + + @abstractmethod + def get_input(self, device) -> tuple[torch.Tensor, ...]: + """Returns an input for the model as as tuple.""" + ... + + @abstractmethod + def get_loss(self, input, output) -> torch.Tensor: + """Returns the loss given the input and output.""" + ... + + @abstractmethod + def run_backward(self, loss) -> None: + """Runs the backward pass (e.g. including ``loss.backward()``).""" + ... + + @staticmethod + @abstractmethod + def init(*args: Any, **kwargs: Any) -> nn.Module: + """Initializes an instance of this model.""" + ... + + +def _assert_module_states( + model: nn.Module, + process_group: dist.ProcessGroup, + assert_fn: Callable, +): + """ + All-gathers module states across ranks and calls ``assert_fn`` on each pair + of corresponding states from rank 0 and a nonzero rank. For example, if + ``assert_fn`` is ``self.assertEqual()``, then this checks that all module + states are equal across ranks. + """ + # Include names for debugging convenience + named_module_states = [ + (param_name, param.detach().cpu()) + for param_name, param in model.named_parameters() + ] + named_module_states += [ + (buffer_name, buffer.detach().cpu()) + for buffer_name, buffer in model.named_buffers() + ] + world_size = dist.get_world_size(process_group) + olist = [None for _ in range(world_size)] + dist.all_gather_object(olist, named_module_states, group=process_group) + rank0_states = olist[0] + if rank0_states is None: + raise AssertionError("Expected rank0_states to not be None") # mypy + for state in olist[1:]: + if state is None: + raise AssertionError("Expected state to not be None") # mypy + for (_, p1), (_, p2) in zip(rank0_states, state, strict=True): + assert_fn(p1, p2) + + +def get_devtype(): + return torch.device(DEVICE_TYPE) + + +def _zero_model( + model: nn.Module, + zero_buffers: bool = False, + summon_full=True, +): + """Zeros the parameters and optionally buffers of ``model`` in place.""" + ctx = FSDP.summon_full_params(model) if summon_full else nullcontext() + with ctx: + for param in model.parameters(): + with torch.no_grad(): + param.zero_() + if zero_buffers: + for buffer in model.buffers(): + with torch.no_grad(): + buffer.zero_() + + +def _get_state_dict(model, cpu_offload=False, half=False): + if not cpu_offload: + model = model.to(DEVICE_TYPE) + if half: + model.half() + + return model.state_dict() + + +def subtest_name(test_name_mapping, *args): + return "_".join( + [test_name_mapping[str(s)] if s is not None else "none" for s in args] + ) + + +def _broadcast_state_dict(rank, state_dict): + # For non-FSDP roots, some parts of the model state on rank 0 may + # not be on CPU, so we move everything to CPU to avoid issues like: + # https://github.com/pytorch/pytorch/issues/77113. + for param_name, param in state_dict.items(): + if param.device != torch.device("cpu"): + state_dict[param_name] = param.cpu() + + olist = [state_dict if rank == 0 else None] + dist.broadcast_object_list(olist) + state_dict = cast(dict[str, torch.Tensor], olist[0]) + # Ensure that the state is on DEVICE + for param_name in state_dict: + state_dict[param_name] = state_dict[param_name].to(DEVICE_TYPE) + return state_dict + + +def get_full_params(model: nn.Module, recurse: bool = True): + """ + Returns the full unsharded parameters of ``model``. Any FSDP-managed + parameters offloaded to CPU are moved to GPU in the returned list. + + Args: + recurse (bool): If ``False``, only unshards the parameters immediate to + ``model``; if ``True``, recurses through the module hierarchy + rooted at ``model``. + """ + with FSDP.summon_full_params(model, recurse=recurse): + return deepcopy(list(model.parameters())) + + +def _move_to_device(model: nn.Module, move_to_device: bool): + return model.to(DEVICE_TYPE) if move_to_device else model + + +def _maybe_wrap_fsdp(model: nn.Module, wrap_fsdp: bool, *args, **kwargs): + return model if not wrap_fsdp else FSDP(model, *args, **kwargs) + + +class DummyProcessGroup: + def __init__(self, rank: int, size: int): + self._rank = rank + self._size = size + + def rank(self) -> int: + return self._rank + + def size(self) -> int: + return self._size + + def allreduce(self, *args, **kwargs): + dist_wait = mock.Mock() + + def get_future(): + future: torch.futures.Future = torch.futures.Future() + future.set_result(1) + return future + + dist_wait.get_future = get_future + return dist_wait + + +class TransformerWithSharedParams(FSDPTestModel): + def __init__( + self, + group: dist.ProcessGroup, + device_init_mode: DEVICEInitMode, + add_bn: bool, + deterministic: bool, + ): + super().__init__() + self.rank = group.rank() + self.world_size = group.size() + if deterministic: + torch.manual_seed(0) + d_vocab = 23 + d_model = 16 + + self.embed_tokens = nn.Embedding(d_vocab, d_model) + self.transformer = nn.Transformer( + d_model=d_model, + num_encoder_layers=2, + num_decoder_layers=2, + dim_feedforward=8, + dropout=0.1, + ) + self.output_proj = nn.Linear(d_model, d_vocab) + + # share the embedding and output projection weights + self.output_proj.weight = self.embed_tokens.weight + self.register_buffer( + "vocab_bias", self.embed_tokens.weight.new_ones((d_model,)) + ) + self.register_buffer( + "long_buffer", + torch.zeros_like(self.vocab_bias, dtype=torch.long), # type: ignore[arg-type] + ) # type: ignore[arg-type] + + self.bs = 2 + self.bn = torch.nn.BatchNorm1d(self.bs) if add_bn else torch.nn.Identity() + if device_init_mode == DEVICEInitMode.DEVICE_BEFORE: + self = self.to(DEVICE_TYPE) + if deterministic: + self.eval() + + def get_input(self, device): + torch.manual_seed(1 + self.rank) # keep everything deterministic + src = torch.arange(12, device=device).view(6, self.bs) # T x B + tgt = torch.arange(self.bs * 4, device=device).view(4, self.bs) # T x B + return (src, tgt) + + def forward(self, src_ids, tgt_ids): + src = self.embed_tokens(src_ids) + src = src + self.vocab_bias + self.long_buffer.type_as(src) # type: ignore[operator] + tgt = self.embed_tokens(tgt_ids) + tgt = self.bn(tgt) + x = self.transformer(src, tgt) + return self.output_proj(x) + + def get_loss(self, input, output): + _, tgt = input + return nn.functional.cross_entropy( + output.view(-1, output.size(-1)), tgt.view(-1), reduction="sum" + ) + + def run_backward(self, loss): + loss.backward() + + @staticmethod + def init( + group: dist.ProcessGroup, + fsdp_init_mode: FSDPInitMode, + device_init_mode: DEVICEInitMode, + fsdp_kwargs: dict[str, Any] | None = None, + deterministic: bool = False, + add_bn: bool = True, + ) -> nn.Module | FSDP: + """ + Initializes a :class:`TransformerWithSharedParams` instance. + + Args: + fsdp_init_mode (FSDPInitMode): If ``NO_FSDP``, then does not wrap + any modules with FSDP. If ``RECURSIVE``, then wraps with + top-level FSDP. By default, the top-level FSDP uses the + ``ModuleWrapPolicy`` for encoder and decoder layers, but a + different auto wrap policy may be specified via + ``fsdp_kwargs``. + device_init_mode (DEVICEInitMode): Determines model movement to DEVICE. + fsdp_kwargs (Optional[Dict[str, Any]]): Optional keyword arguments + forwarded to the FSDP constructor. + deterministic (bool): Whether to make the model deterministic + across constructions. + add_bn (bool): Whether to include batch norm in the model. + """ + + if fsdp_kwargs is None: + fsdp_kwargs = {} + if fsdp_init_mode == FSDPInitMode.NO_FSDP: + if isinstance(group, tuple): + pg = group[0] + else: + pg = group + return TransformerWithSharedParams( + pg, device_init_mode, add_bn, deterministic + ) + elif fsdp_init_mode == FSDPInitMode.RECURSIVE: + # Default to the `ModuleWrapPolicy` + if "auto_wrap_policy" not in fsdp_kwargs: + auto_wrap_policy = ModuleWrapPolicy( + { + TransformerEncoderLayer, + TransformerDecoderLayer, + } + ) + else: + auto_wrap_policy = fsdp_kwargs.pop("auto_wrap_policy") + + if ( + "sharding_strategy" in fsdp_kwargs + and fsdp_kwargs["sharding_strategy"] + in {ShardingStrategy.HYBRID_SHARD, ShardingStrategy._HYBRID_SHARD_ZERO2} + and not isinstance(group, tuple) + ): + fsdp_pg = None + else: + fsdp_pg = group + + if isinstance(group, tuple): + tformer_pg = group[0] + else: + tformer_pg = group + + m = TransformerWithSharedParams( + tformer_pg, device_init_mode, add_bn, deterministic + ) + fsdp_model = FSDP( + m, + fsdp_pg, + auto_wrap_policy=auto_wrap_policy, + **fsdp_kwargs, + ) + if device_init_mode == DEVICEInitMode.DEVICE_AFTER: + fsdp_model = fsdp_model.to(DEVICE_TYPE) + return fsdp_model + raise ValueError(f"Unsupported FSDP init mode: {fsdp_init_mode}") + + def get_ignored_modules(self): + return [self.transformer] + + +class NestedWrappedModule(FSDPTestModel): + def __init__( + self, + group: dist.ProcessGroup, + wrap_fsdp: bool, + device_init_mode: DEVICEInitMode, + deterministic: bool, + **fsdp_kwargs, + ): + super().__init__() + self.rank = group.rank() + self.world_size = group.size() + move_to_device = device_init_mode == DEVICEInitMode.DEVICE_BEFORE + + def _maybe_wrap(layer): + if wrap_fsdp: + return FSDP(layer, group, **fsdp_kwargs) + return layer + + if deterministic: + torch.manual_seed(0) + self.module = nn.Sequential( + _move_to_device(nn.Linear(8, 4), move_to_device), + _maybe_wrap( + nn.Sequential( + _maybe_wrap(_move_to_device(nn.Linear(4, 16), move_to_device)), + _move_to_device(nn.Linear(16, 16), move_to_device), + ), + ), + _maybe_wrap(_move_to_device(nn.Linear(16, 4), move_to_device)), + _move_to_device(nn.Linear(4, 8), move_to_device), + ) + + def get_input(self, device): + torch.manual_seed(1 + self.rank) # keep everything deterministic + return (torch.rand(4, 8, device=device),) + + def forward(self, x): + return self.module(x) + + def get_loss(self, input, output): + loss = output.sum() + return loss + + def run_backward(self, loss): + loss.backward() + + @staticmethod + def init( + group: dist.ProcessGroup, + fsdp_init_mode: FSDPInitMode, + device_init_mode: DEVICEInitMode, + fsdp_kwargs: dict[str, Any] | None = None, + deterministic: bool = False, + ) -> nn.Module: + """ + Initializes a :class:`NestedWrappedModule` instance. + + Args: + fsdp_init_mode (FSDPInitMode): If ``NO_FSDP``, then does not wrap + any modules with FSDP. If ``RECURSIVE``, then wraps some nested + modules with FSDP but not the top-level module. The model may + later be wrapped with a top-level FSDP external to this method + if desired. + device_init_mode (DEVICEInitMode): Determines model movement to DEVICE. + fsdp_kwargs (Optional[Dict[str, Any]]): Optional keyword arguments + forwarded to the FSDP constructor. + deterministic (bool): Whether to make the model deterministic + across constructions. + """ + if fsdp_kwargs is None: + fsdp_kwargs = {} + if fsdp_init_mode == FSDPInitMode.NO_FSDP: + return NestedWrappedModule( + group, + wrap_fsdp=False, + device_init_mode=device_init_mode, + deterministic=deterministic, + ) + elif fsdp_init_mode == FSDPInitMode.RECURSIVE: + # Does not wrap with top-level FSDP + fsdp_model = NestedWrappedModule( + group, + wrap_fsdp=True, + device_init_mode=device_init_mode, + deterministic=deterministic, + **fsdp_kwargs, + ) + if device_init_mode == DEVICEInitMode.DEVICE_AFTER: + fsdp_model = fsdp_model.to(DEVICE_TYPE) + return fsdp_model + raise ValueError(f"Unsupported FSDP init mode: {fsdp_init_mode}") + + +class AlwaysWrapNestedWrappedModule(NestedWrappedModule): + @staticmethod + def init( + group: dist.ProcessGroup, + fsdp_init_mode: FSDPInitMode, + device_init_mode: DEVICEInitMode, + fsdp_kwargs: dict[str, Any] | None = None, + deterministic: bool = False, + ): + """ + Initializes a :class:`NestedWrappedModule` instance, but unlike + :meth:`NestedWrappedModule.init`, for the ``RECURSIVE`` init mode, this + wraps with top-level FSDP and the ``always_wrap_policy()`` auto wrap + policy. + """ + model = super( + AlwaysWrapNestedWrappedModule, AlwaysWrapNestedWrappedModule + ).init( + group=group, + fsdp_init_mode=FSDPInitMode.NO_FSDP, + device_init_mode=device_init_mode, + fsdp_kwargs=fsdp_kwargs, + deterministic=deterministic, + ) + if fsdp_init_mode == FSDPInitMode.NO_FSDP: + return model + elif fsdp_init_mode == FSDPInitMode.RECURSIVE: + fsdp_kwargs = fsdp_kwargs or {} + fsdp_model = FSDP(model, auto_wrap_policy=always_wrap_policy, **fsdp_kwargs) + if device_init_mode == DEVICEInitMode.DEVICE_AFTER: + fsdp_model = fsdp_model.to(DEVICE_TYPE) + return fsdp_model + + +class NonUniformReqGradNWM(NestedWrappedModule): + def __init__( + self, + group: dist.ProcessGroup, + wrap_fsdp: bool, + device_init_mode: DEVICEInitMode, + deterministic: bool, + **fsdp_kwargs, + ): + super(NestedWrappedModule, self).__init__() + # This `__init__` only differs from `NestedWrappedModule.__init__` in that + # the last two `nn.Linear` layers are FSDP wrapped in a `nn.Sequential` + # container. This arrangement results in all elements of the last two parameters + # residing on a single rank. Freezing all parameters except those two allows us + # to verify that `ShardedGradScaler` accommodates situations where some ranks + # have no (non-zero sized) parameter shards. + self.rank = group.rank() + self.world_size = group.size() + move_to_device = device_init_mode == DEVICEInitMode.DEVICE_BEFORE + + def _maybe_wrap(layer): + if wrap_fsdp: + return FSDP(layer, group, **fsdp_kwargs) + return layer + + if deterministic: + torch.manual_seed(0) + self.module = nn.Sequential( + _move_to_device(nn.Linear(8, 4), move_to_device), + _maybe_wrap( + nn.Sequential( + _maybe_wrap(_move_to_device(nn.Linear(4, 16), move_to_device)), + _move_to_device(nn.Linear(16, 16), move_to_device), + ), + ), + _maybe_wrap( + nn.Sequential( + _move_to_device(nn.Linear(16, 4), move_to_device), + _move_to_device(nn.Linear(4, 8), move_to_device), + ), + ), + ) + + @staticmethod + def _set_nonuniform_req_grad(model, req_grad_mask) -> None: + for n, p in model.named_parameters(): + if not re.match(req_grad_mask, n): + p.requires_grad_(False) + + @staticmethod + def init( + group: dist.ProcessGroup, + fsdp_init_mode: FSDPInitMode, + device_init_mode: DEVICEInitMode, + fsdp_kwargs: dict[str, Any] | None = None, + deterministic: bool = False, + ): + """ + Initializes a :class:`NestedWrappedModule` instance, but unlike + :meth:`NestedWrappedModule.init`, it wraps a second :class:`torch.nn.Sequential` + container to enable the desired non-uniform ``requires_grad`` + ``use_orig_params=True`` tests. For both ``RECURSIVE`` and ``NO_FSDP`` + init modes, freezes all parameters except the last two to validate + ``ShardedGradScaler`` support for ranks with no (non-zero sized) local shards in + FSDP ``use_orig_params=True`` mode. + """ + # The parameters that should remain unfrozen are in `module.2.1`. The regex + # pattern below matches the relevant parameter names both with and without + # an interstitial FSDP module indicator (`_fsdp_wrapped_module`) present. + req_grad_pattern = re.compile(r"module\.2.*\.1.*") + if fsdp_init_mode == FSDPInitMode.NO_FSDP: + ddp_model = NonUniformReqGradNWM( + group, + wrap_fsdp=False, + device_init_mode=device_init_mode, + deterministic=deterministic, + ) + NonUniformReqGradNWM._set_nonuniform_req_grad(ddp_model, req_grad_pattern) + return ddp_model + elif fsdp_init_mode == FSDPInitMode.RECURSIVE: + if fsdp_kwargs is None: + fsdp_kwargs = {} + fsdp_model = NonUniformReqGradNWM( + group, + wrap_fsdp=True, + device_init_mode=device_init_mode, + deterministic=deterministic, + **fsdp_kwargs, + ) + if device_init_mode == DEVICEInitMode.DEVICE_AFTER: + fsdp_model = fsdp_model.to(DEVICE_TYPE) + NonUniformReqGradNWM._set_nonuniform_req_grad(fsdp_model, req_grad_pattern) + return fsdp_model + raise ValueError(f"Unsupported FSDP init mode: {fsdp_init_mode}") + + +class ModuleWithDelay(FSDPTestModel): + """This class wraps a :class:`FSDPTestModel` to optionally add a delay + after computing the loss and/or before the gradient reduction.""" + + def __init__( + self, + module: nn.Module, + delay_after_loss_ms: int, + delay_before_reduction_ms: int, + ): + super().__init__() + self.delay_after_loss_ms = delay_after_loss_ms + self.delay_before_reduction_ms = delay_before_reduction_ms + self.module = module + + def get_input(self, device): + return self.module.get_input(device) # type: ignore[operator] + + def forward(self, x): + return self.module(x) + + def get_loss(self, input, output): + loss = self.module.get_loss(input, output) # type: ignore[operator] + if self.delay_after_loss_ms > 0: + if TEST_HPU or TEST_XPU: + time.sleep(self.delay_after_loss_ms / 1000) + elif TEST_CUDA: + torch.cuda._sleep(int(self.delay_after_loss_ms * get_cycles_per_ms())) + + return loss + + def run_backward(self, loss): + orig_reduce_scatter = torch.distributed.reduce_scatter_tensor + + def _delayed_reduce_scatter(*args, **kwargs): + if self.delay_before_reduction_ms > 0: + if TEST_CUDA: + torch.cuda._sleep( + int(self.delay_before_reduction_ms * get_cycles_per_ms()) + ) + elif TEST_HPU or TEST_XPU: + time.sleep(self.delay_before_reduction_ms / 1000) + return orig_reduce_scatter(*args, **kwargs) + + with mock.patch( + "torch.distributed.reduce_scatter_tensor", _delayed_reduce_scatter + ): + self.module.run_backward(loss) # type: ignore[operator] + + @staticmethod + def init( + module_class: type[FSDPTestModel], + *model_args: Any, + delay_after_loss_ms: int, + delay_before_reduction_ms: int, + **model_kwargs: Any, + ): + """ + Args: + module_class (Type[FSDPTestModel]): Wrapped module class to which + to add delays. + model_args: Positional arguments forwarded to the ``module_class`` + ``init()``. + delay_after_loss_ms (int): Delay after computing the loss/before + the optimizer step (in ms). + delay_before_reduction_ms (int): Delay before reduce-scattering + gradients (in ms). + model_kwargs: Keyword arguments forwarded to the ``module_class`` + ``init()``. + """ + return ModuleWithDelay( + module_class.init(*model_args, **model_kwargs), + delay_after_loss_ms, + delay_before_reduction_ms, + ) + + +class NestedWrappedModuleWithDelay(ModuleWithDelay): + @staticmethod + def init( # type: ignore[override] + group: dist.ProcessGroup, + fsdp_init_mode: FSDPInitMode, + device_init_mode: DEVICEInitMode = DEVICEInitMode.DEVICE_AFTER, + fsdp_kwargs: dict[str, Any] | None = None, + deterministic: bool = False, + delay_after_loss_ms: int = 0, + delay_before_reduction_ms: int = 0, + ): + return ModuleWithDelay.init( + NestedWrappedModule, + group=group, + fsdp_init_mode=fsdp_init_mode, + device_init_mode=device_init_mode, + fsdp_kwargs=fsdp_kwargs, + deterministic=deterministic, + delay_after_loss_ms=delay_after_loss_ms, + delay_before_reduction_ms=delay_before_reduction_ms, + ) + + +class DummyDDP(nn.Module): + def __init__(self, module): + super().__init__() + self.module = module + + def forward(self, *args, **kwargs): + return self.module(*args, **kwargs) + + +class MixtureOfExperts(NestedWrappedModule): + def __init__( + self, + group: dist.ProcessGroup, + wrap_fsdp: bool, + device_init_mode: DEVICEInitMode, + delay_before_free_ms: int, + deterministic: bool, + **fsdp_kwargs, + ): + super().__init__( + group=group, + wrap_fsdp=wrap_fsdp, + device_init_mode=device_init_mode, + deterministic=deterministic, + ) + self.group = group + self.delay_before_free_ms = delay_before_free_ms + self.wrap_fsdp = wrap_fsdp + self.move_to_device = device_init_mode == DEVICEInitMode.DEVICE_BEFORE + if deterministic: + # Give each rank different expert parameters + torch.manual_seed(42 + self.rank) + d_expert = 23 + d_shared = 12 + d_input = 8 + expert = _move_to_device(nn.Linear(d_expert, d_shared), self.move_to_device) + + self.num_expert_params = sum(p.numel() for p in expert.parameters()) + for p in expert.parameters(): + p.expert = True # type: ignore[attr-defined] + + if deterministic: + # Keep all other parameters the same across ranks + torch.manual_seed(0) + + shared = _move_to_device(nn.Linear(d_shared, d_expert), self.move_to_device) + + if wrap_fsdp: + # we create a process group of size 1 for the expert params + expert_group = torch.distributed.new_group( + [group.rank()] + ) # world size 1 means no shard + expert = FSDP(expert, expert_group, **fsdp_kwargs) # type: ignore[assignment] + shared = FSDP(shared, group, **fsdp_kwargs) # type: ignore[assignment] + + self.module = nn.Sequential( + _move_to_device(nn.Linear(d_input, d_shared), self.move_to_device), + shared, + expert, + _move_to_device(nn.Linear(d_shared, d_input), self.move_to_device), + ) + + def forward(self, x): + if self.delay_before_free_ms > 0: + expert = self.module[2] + if isinstance(expert, FSDP): + orig_reshard = torch.distributed.fsdp._runtime_utils._reshard + + def _delayed_reshard(*args, **kwargs): + if TEST_CUDA: + torch.cuda._sleep( + int(self.delay_before_free_ms * get_cycles_per_ms()) + ) + elif TEST_HPU or TEST_XPU: + time.sleep(self.delay_before_free_ms / 1000) + + return orig_reshard(*args, **kwargs) + + # This patch covers any `import torch..._reshard` uses. + with mock.patch( + "torch.distributed.fsdp._runtime_utils._reshard", _delayed_reshard + ): + return self.module(x) + + return self.module(x) + + def run_backward(self, loss): + loss.backward() + # Manually reduce gradients if not wrapped in FullyShardedDataParallel + if not self.wrap_fsdp: + with torch.no_grad(): + for p in self.parameters(): + if hasattr(p, "expert"): + continue # these params don't need grad reduction + if p.grad is not None: + p.grad.div_(self.world_size) + torch.distributed.all_reduce(p.grad, group=self.group) + + @staticmethod + def init( + group: dist.ProcessGroup, + fsdp_init_mode: FSDPInitMode, + device_init_mode: DEVICEInitMode, + fsdp_kwargs: dict[str, Any] | None = None, + deterministic: bool = False, + delay_before_free_ms: int = 0, + ): + """ + Initializes a :class:`MixtureOfExperts` instance. + + Args: + fsdp_init_mode (FSDPInitMode): If ``NO_FSDP``, then does not wrap + any modules with FSDP. If ``RECURSIVE``, then wraps some nested + modules with FSDP, including the expert and shared layers, but + not the top-level module. The model may later be wrapped with a + top-level FSDP external to this method if desired. + device_init_mode (DEVICEInitMode): Determines model movement to DEVICE. + fsdp_kwargs (Optional[Dict[str, Any]]): Optional keyword arguments + forwarded to the FSDP constructor. + deterministic (bool): Whether to make the model deterministic + across constructions. + delay_before_free_ms (int): Delay before resharding expert + parameters in the forward pass (in ms). + """ + if fsdp_kwargs is None: + fsdp_kwargs = {} + if fsdp_init_mode == FSDPInitMode.NO_FSDP: + return MixtureOfExperts( + group, + wrap_fsdp=False, + device_init_mode=device_init_mode, + delay_before_free_ms=delay_before_free_ms, + deterministic=deterministic, + ) + elif fsdp_init_mode == FSDPInitMode.RECURSIVE: + # Does not wrap with top-level FSDP + fsdp_model = MixtureOfExperts( + group, + wrap_fsdp=True, + device_init_mode=device_init_mode, + delay_before_free_ms=delay_before_free_ms, + deterministic=deterministic, + **fsdp_kwargs, + ) + if device_init_mode == DEVICEInitMode.DEVICE_AFTER: + fsdp_model = fsdp_model.to(DEVICE_TYPE) + return fsdp_model + raise ValueError(f"Unsupported FSDP init mode: {fsdp_init_mode}") + + +class MLP(nn.Module): + def __init__( + self, + dim: int, + device: torch.device | None = None, + *, + bias: bool = True, + with_buffer: bool = False, + dim_multiplier: int = 4, + ): + super().__init__() + self.in_proj = nn.Linear(dim, dim_multiplier * dim, device=device, bias=bias) + self.out_proj = nn.Linear(dim_multiplier * dim, dim, device=device, bias=bias) + if with_buffer: + self.register_buffer("buffer", torch.randn((dim,), device=device)) + else: + self.buffer = None + + def forward(self, x: torch.Tensor) -> torch.Tensor: + z = self.in_proj(x) + z = F.relu(z) + z = self.out_proj(z) + z = F.relu(z) + if self.buffer is not None: + z = z + self.buffer + return z + + def reset_parameters(self): + if self.buffer is not None: + torch.nn.init.normal_(self.buffer) + + +class MLPStack(nn.Sequential): + def __init__(self, mlp_dim: int, *, with_seq_parallel: bool = False): + modules: list[nn.Module] = [ + # Use multiplier of 3 to exercise uneven case + MLP(mlp_dim, dim_multiplier=3), + MLP(mlp_dim), + MLP(mlp_dim, dim_multiplier=3), + ] + if with_seq_parallel: + modules.append(nn.LayerNorm(mlp_dim, bias=False)) + super().__init__(*modules) + self.with_seq_parallel = with_seq_parallel + + def parallelize( + self, + tp_mesh: DeviceMesh, + dp_mesh: DeviceMesh, + use_activation_checkpointing: bool, + **fsdp_kwargs, + ) -> "MLPStack": + parallelize_plan = { + # Pass `use_local_output=False` to keep as DTensor to preserve + # uneven activation dims + "0.in_proj": ColwiseParallel(use_local_output=False), + "0.out_proj": RowwiseParallel(use_local_output=False), + "1.in_proj": ColwiseParallel(use_local_output=False), + "1.out_proj": RowwiseParallel(use_local_output=False), + "2.in_proj": ColwiseParallel(use_local_output=False), + "2.out_proj": RowwiseParallel(output_layouts=Shard(1)) + if self.with_seq_parallel + else RowwiseParallel(), + } + if self.with_seq_parallel: + parallelize_plan["3"] = SequenceParallel(sequence_dim=1) + parallelize_module(self, device_mesh=tp_mesh, parallelize_plan=parallelize_plan) + for module in self: + if isinstance(module, nn.LayerNorm): + continue + if use_activation_checkpointing: + checkpoint(module) + fully_shard(module, mesh=dp_mesh, **fsdp_kwargs) + fully_shard(self, mesh=dp_mesh, **fsdp_kwargs) + return self + + +class DoubleLinear(nn.Module): + """ + This can be used for returning multiple outputs from a module + (``use_second_linear=True``) or for having an unused module (``False``). + """ + + def __init__(self, dim: int, use_second_linear: bool = True): + super().__init__() + self.lin1 = nn.Linear(dim, dim) + self.lin2 = nn.Linear(dim, dim) + self.relu = nn.ReLU() + self.use_second_linear = use_second_linear + + def forward( + self, x: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor] | torch.Tensor: + if self.use_second_linear: + return self.relu(self.lin1(x)), self.relu(self.lin2(x)) + return self.relu(self.lin1(x)) + + +# NOTE: For these patch methods, if we want safety under multi-threading (e.g. +# when using multi-threaded process group), then we want: +# (1) a barrier immediately after reading the original value to ensure that all +# threads see the same original value +# (2) a barrier immediately before restoring the original value to ensure that +# all threads use the patched value inside the context +@contextlib.contextmanager +def patch_all_gather(new_all_gather_into_tensor: Callable): + orig_all_gather = dist.all_gather_into_tensor + dist.barrier() + dist.all_gather_into_tensor = new_all_gather_into_tensor + try: + yield + finally: + dist.barrier() + dist.all_gather_into_tensor = orig_all_gather + + +@contextlib.contextmanager +def patch_foreach_all_gather(new_foreach_all_gather: Callable): + orig_foreach_all_gather = ( + torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_all_gather + ) + dist.barrier() + torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_all_gather = ( + new_foreach_all_gather + ) + try: + yield + finally: + dist.barrier() + torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_all_gather = ( + orig_foreach_all_gather + ) + + +@contextlib.contextmanager +def patch_foreach_reduce(new_foreach_reduce: Callable): + orig_foreach_foreach_reduce = ( + torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_reduce + ) + dist.barrier() + torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_reduce = ( + new_foreach_reduce + ) + try: + yield + finally: + dist.barrier() + torch.distributed.fsdp._fully_shard._fsdp_param_group.foreach_reduce = ( + orig_foreach_foreach_reduce + ) + + +@contextlib.contextmanager +def patch_reduce_scatter(new_reduce_scatter_tensor: Callable): + orig_reduce_scatter = dist.reduce_scatter_tensor + dist.barrier() + dist.reduce_scatter_tensor = new_reduce_scatter_tensor + try: + yield + finally: + dist.barrier() + dist.reduce_scatter_tensor = orig_reduce_scatter + + +@contextlib.contextmanager +def patch_all_reduce(new_all_reduce: Callable): + orig_all_reduce = dist.all_reduce + dist.barrier() + dist.all_reduce = new_all_reduce + try: + yield + finally: + dist.barrier() + dist.all_reduce = orig_all_reduce + + +@no_type_check +@contextlib.contextmanager +def patch_unshard(new_unshard: Callable): + orig_unshard = FSDPParamGroup.unshard + dist.barrier() + FSDPParamGroup.unshard = new_unshard + try: + yield + finally: + dist.barrier() + FSDPParamGroup.unshard = orig_unshard + + +@no_type_check +@contextlib.contextmanager +def patch_reshard(new_reshard: Callable): + orig_reshard = FSDPParamGroup.reshard + dist.barrier() + FSDPParamGroup.reshard = new_reshard + try: + yield + finally: + dist.barrier() + FSDPParamGroup.reshard = orig_reshard + + +@no_type_check +@contextlib.contextmanager +def patch_post_backward(new_post_backward: Callable): + orig_post_backward = FSDPParamGroup.post_backward + dist.barrier() + FSDPParamGroup.post_backward = new_post_backward + try: + yield + finally: + dist.barrier() + FSDPParamGroup.post_backward = orig_post_backward + + +@no_type_check +@contextlib.contextmanager +def patch_register_post_backward_hook_backward(new_backward: Callable): + orig_backward = RegisterPostBackwardFunction.backward + dist.barrier() + RegisterPostBackwardFunction.backward = new_backward + try: + yield + finally: + dist.barrier() + RegisterPostBackwardFunction.backward = orig_backward + + +def reduce_scatter_with_assert( + cls, + orig_reduce_scatter: Callable, + assert_fn: Callable, # `assert_fn(output: Tensor)` + *args: Any, + **kwargs: Any, +): + if len(args) > 0: + output = args[0] + elif "output" in kwargs: + output = kwargs["output"] + else: + raise AssertionError( + f"Cannot get reduce-scatter output from\nargs: {args}\nkwargs: {kwargs}" + ) + assert_fn(output) + return orig_reduce_scatter(*args, **kwargs) + + +def check_sharded_parity( + cls, # unit test class + replicated_module: nn.Module, + sharded_module: nn.Module, + prefixes_to_ignore: tuple[str, ...] = (), +): + for (replicated_name, replicated_param), (sharded_name, sharded_param) in zip( + replicated_module.named_parameters(), + sharded_module.named_parameters(), + strict=True, + ): + clean_sharded_name = sharded_name + for prefix in prefixes_to_ignore: + clean_sharded_name = clean_sharded_name.replace(prefix, "") + cls.assertEqual(replicated_name, clean_sharded_name) + cls.assertIsInstance(sharded_param, DTensor) + if not isinstance(sharded_param, DTensor): + raise AssertionError("Expected sharded_param to be a DTensor") # mypy + mesh, placements = sharded_param.device_mesh, sharded_param.placements + if tuple(placements) == (Shard(0), Shard(0)): + raise AssertionError( + "FSDP's (Shard(0), Shard(0)) layout differs from distribute_tensor(), " + "so we cannot check for equality using it" + ) + sharded_ref_param = distribute_tensor(replicated_param, mesh, placements) + cls.assertEqual(sharded_param.to_local(), sharded_ref_param.to_local()) + if replicated_param.grad is None: + cls.assertIsNone(sharded_param.grad) + continue + cls.assertIsNotNone(sharded_param.grad) + sharded_ref_grad = distribute_tensor(replicated_param.grad, mesh, placements) + cls.assertIsInstance(sharded_param.grad, DTensor) + if not isinstance(sharded_param.grad, DTensor): + raise AssertionError("Expected sharded_param.grad to be a DTensor") # mypy + cls.assertEqual(sharded_param.grad.to_local(), sharded_ref_grad.to_local()) + + +@unittest.skipIf(TEST_XPU, "not-support-multithread") +class FSDPTestMultiThread(MultiThreadedTestCase): + @property + def world_size(self): + return DEVICE_COUNT + + def setUp(self): + super().setUp() + self._spawn_threads() + + def run_subtests(self, *args, **kwargs): + return run_subtests(self, *args, **kwargs) + + def perThreadSetUp(self): + torch._dynamo.reset() + + def perThreadTearDown(self): + torch._dynamo.reset() + + +class FSDPTestMixin: + """ + Mixin class containing shared test utilities for FSDP tests. + Provides common helper methods for both FSDPTest and FSDPTestContinuous. + """ + + def _check_cpu_offload(self, fsdp_model, cpu_offload): + self.assertEqual(cpu_offload, fsdp_model.cpu_offload) + + def _check_backward_prefetch(self, fsdp_model, backward_prefetch): + self.assertEqual(backward_prefetch, fsdp_model.backward_prefetch) + + def _check_forward_prefetch(self, fsdp_model, forward_prefetch): + self.assertEqual(forward_prefetch, fsdp_model.forward_prefetch) + + def run_subtests(self, *args, **kwargs): + return run_subtests(self, *args, **kwargs) + + @classmethod + def _run(cls, rank, test_name, file_name, pipe, **kwargs): # type: ignore[override] + self = cls(test_name) + self.rank = rank + self.file_name = file_name + fake_pg = kwargs.get("fake_pg", False) + + print(f"dist init r={self.rank}, world={self.world_size}") + if DEVICE_TYPE != "cpu" and torch.accelerator.device_count() < self.world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code) + + # Specify gloo backend to make 'init_process_group()' succeed, + # Actual tests will be skipped if there is no enough GPUs. + try: + if fake_pg: + store = torch.testing._internal.distributed.fake_pg.FakeStore() + dist.init_process_group( + backend="fake", + world_size=self.world_size, + rank=rank, + store=store, + ) + else: + dist.init_process_group( + init_method=self.init_method, + backend=DISTRIBUTED_BACKEND, + world_size=int(self.world_size), + rank=self.rank, + ) + except RuntimeError as e: + if "recompile" in e.args[0]: + sys.exit(TEST_SKIPS["backend_unavailable"].exit_code) + + raise + + device_ids = None + device_id = self.rank % DEVICE_COUNT + if TEST_CUDA or TEST_XPU: + torch.accelerator.set_device_index(device_id) + device_ids = [device_id] + + # Execute barrier prior to running test to ensure that every process + # has finished initialization and that the following test + # immediately exiting due to a skip doesn't cause flakiness. + dist.barrier(device_ids=device_ids) + + torch._dynamo.reset() + set_rng_seed() + self.run_test(test_name, pipe) + torch._dynamo.reset() + + dist.barrier(device_ids=device_ids) + + dist.destroy_process_group() + + def _train_for_several_steps( + self, + model: nn.Module, + num_steps: int, + autocast: bool, + lr: float = 0.01, + fsdp_cpu_offload: CPUOffload | None = None, + save_model: bool = False, + mixed_precision: MixedPrecision | None = None, + enable_sharded_grad_scaler: bool = False, + use_pure_fp16: bool = False, + sharded_grad_scaler_kwargs: dict[str, Any] | None = None, + ): + cpu_offload_params = fsdp_cpu_offload and fsdp_cpu_offload.offload_params + + model_device = next(model.parameters()).device + if sharded_grad_scaler_kwargs is None: + sharded_grad_scaler_kwargs = {} + sharded_grad_scaler = ShardedGradScaler( + enabled=enable_sharded_grad_scaler, **sharded_grad_scaler_kwargs + ) + # use SGD with momentum instead of Adam, since Adam is scale invariant + # and this makes it bad for tests + optim = torch.optim.SGD(model.parameters(), lr=lr, momentum=0.9) + for _ in range(num_steps): + optim.zero_grad() + with torch.amp.autocast(DEVICE_TYPE, enabled=autocast): + # Inputs always cuda regardless of cpu offloading, or model.device + input = model.module.get_input(torch.device(DEVICE_TYPE)) # type: ignore[operator, union-attr] + if use_pure_fp16 or (mixed_precision and not isinstance(model, FSDP)): + if isinstance(input, torch.Tensor): + input = input.half() + else: + input = tuple(x.half() for x in input) + output = model(*input) + # Post-forward, if CPU offloading model param should be on CPU. + if ( + cpu_offload_params + and isinstance(model, FSDP) + # If not resharding after forward, the parameters are still + # exposed as unsharded views into the GPU flat parameter + and model.sharding_strategy + not in NO_RESHARD_AFTER_FORWARD_STRATEGIES + ): + for p in model.parameters(): + # Params should always be on CPU + self.assertEqual(p.device, torch.device("cpu")) + + loss = model.module.get_loss(input, output).to(model_device) # type: ignore[operator, union-attr] + loss = sharded_grad_scaler.scale(loss) + + if not mixed_precision and not use_pure_fp16: + if loss.dtype != torch.float32: + raise AssertionError( + "loss data type should be float32, as the original " + "parameter data type is float32." + ) + else: + if use_pure_fp16: + self.assertEqual(loss.dtype, torch.float16) + # FSDP loss is fp16, DDP AMP loss is fp32 + elif isinstance(model, FSDP): + if mixed_precision is None: + raise AssertionError( + "Expected mixed_precision to not be None" + ) # mypy + self.assertEqual(loss.dtype, mixed_precision.param_dtype) + else: + self.assertEqual(loss.dtype, torch.float32) + model.module.run_backward(loss) # type: ignore[operator, union-attr] + # Post-backward, if CPU offloading model params should be on CPU. + if cpu_offload_params and isinstance(model, FSDP): + for p in model.parameters(): + # Params should always be on CPU + self.assertEqual(p.device, torch.device("cpu")) + # Unscale the gradients and step + sharded_grad_scaler.step(optim) + # Update the scale factor + sharded_grad_scaler.update() + # if save_model, simulate save + load. + if save_model: + state_dict = {k: v.clone() for k, v in model.state_dict().items()} + # Zero params, if save/load state_dict did not work properly, this + # would break the parity test with DDP. + _zero_model(model) + model.load_state_dict(state_dict) + + if isinstance(model, FSDP): + model._assert_state(TrainingState.IDLE) + return loss.detach() # type: ignore[possibly-undefined] + + def _test_fsdp_parity( + self, + model_class: type[FSDPTestModel], + fsdp_init_mode: FSDPInitMode, + device_init_mode: DEVICEInitMode, + ref_init_fn: Callable | None = None, + num_iters: int = 2, + save_model: bool = True, + cpu_offload: CPUOffload = CPUOffload(), + backward_prefetch: BackwardPrefetch | None = None, + sharding_strategy: ShardingStrategy | None = None, + mixed_precision: MixedPrecision | None = None, + forward_prefetch: bool = False, + use_orig_params: bool = False, + enable_sharded_grad_scaler: bool = False, + use_pure_fp16: bool = False, + init_kwargs: dict[str, Any] | None = None, + sharded_grad_scaler_kwargs: dict[str, Any] | None = None, + **fsdp_kwargs, + ): + """ + Tests FSDP training against a reference, which defaults to DDP but + may be customized with ``ref_init_fn``. + + Args: + model_class (Type[FSDPTestModel]): A model class that inherits from + ``FSDPTestModel``, which defines the expected interface. + fsdp_init_mode (FSDPInitMode): The mode to initialize the + FSDP-wrapped model. This should not be ``NO_FSDP``. + ref_init_fn (Optional[Callable]): A callable to invoke that wraps a + non-wrapped model to construct the reference model, where this + wrapper should provide data parallel semantics. If ``None``, + then the callable defaults to the DDP constructor. + """ + if fsdp_init_mode == FSDPInitMode.NO_FSDP: + raise AssertionError("Expects an FSDP init mode that wraps with FSDP") + if init_kwargs is None: + init_kwargs = {} + lr = 1e-2 + rank = self.process_group.rank() + # Establish reference behavior with DDP + model = model_class.init( + self.process_group, + FSDPInitMode.NO_FSDP, + DEVICEInitMode.DEVICE_BEFORE, + deterministic=True, + **init_kwargs, + ) + if ref_init_fn is None: + if TEST_HPU: + ref_model = DDP( + model, device_ids=[DEVICE_TYPE], output_device=DEVICE_TYPE + ) + elif DEVICE_TYPE == "cpu": + ref_model = DDP(model) + else: + ref_model = DDP(model, device_ids=[rank], output_device=rank) + else: + ref_model = ref_init_fn(model) + if use_pure_fp16: + ref_model = ref_model.half() + ref_loss = self._train_for_several_steps( + ref_model, + num_iters, + autocast=mixed_precision is not None, + lr=lr, + fsdp_cpu_offload=cpu_offload, + mixed_precision=mixed_precision, + enable_sharded_grad_scaler=enable_sharded_grad_scaler, + use_pure_fp16=use_pure_fp16, + sharded_grad_scaler_kwargs=sharded_grad_scaler_kwargs, + ) + ddp_params = list(ref_model.parameters()) + # Check against FSDP behavior + fsdp_kwargs.update( + { + "cpu_offload": cpu_offload, + "backward_prefetch": backward_prefetch, + "sharding_strategy": sharding_strategy, + "mixed_precision": mixed_precision, + "forward_prefetch": forward_prefetch, + "use_orig_params": use_orig_params, + } + ) + try: + fsdp_model = model_class.init( + self.process_group, + fsdp_init_mode, + device_init_mode, + fsdp_kwargs, + deterministic=True, + **init_kwargs, + ) + except Exception as e: + raise ValueError(f"Initializing {model_class} raised error {str(e)}") from e + if not isinstance(fsdp_model, FSDP): + # Enforce that we wrap with top-level FSDP since we are comparing + # assuming a data parallel reference and some test models may not + # do so in their `init()` method + fsdp_model = FSDP(fsdp_model, self.process_group, **fsdp_kwargs) + if use_pure_fp16: + # Change the model parameter dtype after FSDP initialization + fsdp_model = fsdp_model.half() + if device_init_mode == DEVICEInitMode.DEVICE_AFTER: + fsdp_model = fsdp_model.to(DEVICE_TYPE) + offload_params = cpu_offload is not None and cpu_offload.offload_params + # Offloading parameters with `DEVICE_AFTER` should raise an error during + # lazy initialization due to the parameter devices not being CPU; + # otherwise, all parameter devices should be CPU + expects_device_error = ( + offload_params and device_init_mode == DEVICEInitMode.DEVICE_AFTER + ) + expects_cpu_device = ( + offload_params and device_init_mode != DEVICEInitMode.DEVICE_AFTER + ) + if expects_cpu_device: + cpu_device = torch.device("cpu") + for param in fsdp_model.parameters(): + self.assertEqual(param.device, cpu_device) + context = ( + self.assertRaisesRegex( + RuntimeError, + "An FSDP-managed module with parameter CPU offloading enabled " + f"has parameters on {DEVICE_TYPE}", + ) + if expects_device_error + else nullcontext() + ) + with context: + fsdp_loss = self._train_for_several_steps( + fsdp_model, + num_iters, + autocast=False, + lr=lr, + fsdp_cpu_offload=cpu_offload, + save_model=save_model, + mixed_precision=mixed_precision, + enable_sharded_grad_scaler=enable_sharded_grad_scaler, + use_pure_fp16=use_pure_fp16, + sharded_grad_scaler_kwargs=sharded_grad_scaler_kwargs, + ) + # No need to check for parameter and loss parity if expecting an error + if expects_device_error: + return + # Check parameter devices are CPU if offloading to CPU before calling + # `get_full_params()`, which will cast the parameters to FP32 + if offload_params: + cpu_device = torch.device("cpu") + for param in fsdp_model.parameters(): + self.assertEqual(param.device, cpu_device) + fsdp_loss = fsdp_loss.to(DEVICE_TYPE) + fsdp_unsharded_params = get_full_params(fsdp_model) + # Do not check dtype since the reference DDP loss may not be the same + # dtype as the FSDP loss in the case of mixed precision + torch.testing.assert_close(ref_loss, fsdp_loss, check_dtype=False) + # Do not check for parameter parity if using mixed precision since (1) + # the DDP parameters are in FP16 (from `half()`) while the FSDP + # parameters are in FP32 (from `summon_full_params()`) and (2) DDP runs + # the optimizer in FP16 while FSDP runs it in FP32 + # TODO: Disable checking the parameters for pure FP16 due to floating + # point inaccuracy. Note that this means that the backward pass is not + # checked: https://github.com/pytorch/pytorch/issues/90784 + if mixed_precision is None and not use_pure_fp16: + self.assertEqual( + ddp_params, + fsdp_unsharded_params, + exact_device=True, + msg="FSDP did not match DDP", + ) + + +class FSDPTest(FSDPTestMixin, MultiProcessTestCase): + def setUp(self): + super().setUp() + # Set TORCH_NCCL_DESYNC_DEBUG=0 to disable the NCCL `workCleanupLoop()`, + # which can cause unit test flakiness: + # https://github.com/pytorch/pytorch/issues/90848 + os.environ["TORCH_NCCL_DESYNC_DEBUG"] = "0" + self._spawn_processes() + + @property + def world_size(self): + return DEVICE_COUNT + + @property + def process_group(self): + return dist.distributed_c10d._get_default_group() + + @property + def destroy_pg_upon_exit(self) -> bool: + # Overriding base test class: do not auto destroy PG upon exit. + return False + + @property + def init_method(self): + return f"{FILE_SCHEMA}{self.file_name}" + + @classmethod + def _run(cls, rank, test_name, file_name, pipe, **kwargs): # type: ignore[override] + self = cls(test_name) + self.rank = rank + self.file_name = file_name + fake_pg = kwargs.get("fake_pg", False) + + print(f"dist init r={self.rank}, world={self.world_size}") + if torch.accelerator.device_count() < self.world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code) + + # Specify gloo backend to make 'init_process_group()' succeed, + # Actual tests will be skipped if there is no enough GPUs. + try: + if fake_pg: + store = torch.testing._internal.distributed.fake_pg.FakeStore() + dist.init_process_group( + backend="fake", + world_size=self.world_size, + rank=rank, + store=store, + ) + else: + dist.init_process_group( + init_method=self.init_method, + backend=DISTRIBUTED_BACKEND, + world_size=int(self.world_size), + rank=self.rank, + ) + except RuntimeError as e: + if "recompile" in e.args[0]: + sys.exit(TEST_SKIPS["backend_unavailable"].exit_code) + + raise + + device_ids = None + device_id = self.rank % DEVICE_COUNT + if TEST_CUDA or TEST_XPU: + torch.accelerator.set_device_index(device_id) + device_ids = [device_id] + + # Execute barrier prior to running test to ensure that every process + # has finished initialization and that the following test + # immediately exiting due to a skip doesn't cause flakiness. + dist.barrier(device_ids=device_ids) + + torch._dynamo.reset() + set_rng_seed() + self.run_test(test_name, pipe) + torch._dynamo.reset() + + dist.barrier(device_ids=device_ids) + + dist.destroy_process_group() + + +class FSDPTestContinuous(FSDPTestMixin, MultiProcContinuousTest): + """ + FSDP test base class using MultiProcContinuousTest for faster test execution. + This class reuses worker processes across tests, reducing process spawn overhead. + Use this for tests that don't require fresh process state between tests. + """ + + world_size: int = DEVICE_COUNT + + @classmethod + def backend_str(cls) -> str: + return DISTRIBUTED_BACKEND + + @classmethod + def device_type(cls) -> str: + return DEVICE_TYPE + + @classmethod + def _init_pg(cls, rank, world_size, rdvz_file): + # Set TORCH_NCCL_DESYNC_DEBUG=0 to disable the NCCL `workCleanupLoop()`, + # which can cause unit test flakiness: + # https://github.com/pytorch/pytorch/issues/90848 + os.environ["TORCH_NCCL_DESYNC_DEBUG"] = "0" + + if torch.accelerator.device_count() < world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{world_size}"].exit_code) + + device_id = rank % DEVICE_COUNT + if TEST_CUDA or TEST_XPU: + torch.accelerator.set_device_index(device_id) + + super()._init_pg(rank, world_size, rdvz_file) + + def setUp(self): + super().setUp() + # Barrier to synchronize workers before test, similar to FSDPTest._run(). + # This ensures all workers start the test together and prevents NCCL + # collective mismatches when the process group is reused across tests. + if self.rank != self.MAIN_PROCESS_RANK: + dist.barrier() + torch._dynamo.reset() + set_rng_seed() + + def tearDown(self): + # Barrier to synchronize workers after test, similar to FSDPTest._run(). + if self.rank != self.MAIN_PROCESS_RANK: + dist.barrier() + super().tearDown() + torch._dynamo.reset() + + @property + def process_group(self): + return self.__class__.pg + + +def compiled_fsdp_test(compile_compute_on_module: type | None = None): + def fully_shard_with_compiled_compute(*args, **kwargs): + torch.distributed.fsdp.fully_shard(*args, **kwargs) # type: ignore[operator] + if compile_compute_on_module is None or isinstance( + args[0], compile_compute_on_module + ): + args[0].compile() + + class FullyShardMode(Enum): + EAGER = auto() + COMPILED_COMPUTE = auto() + + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + original_fully_shard: Any = torch.distributed.fsdp.fully_shard + for mode in FullyShardMode: + if mode != FullyShardMode.EAGER and not has_triton(): + warnings.warn( + "Inductor on GPU needs Triton and recent GPU arch", stacklevel=2 + ) + continue + # barrier to ensure thread reading the same value + original_compile_threads = torch._inductor.config.compile_threads + torch.distributed.barrier() + + if mode == FullyShardMode.EAGER: + fully_shard_patch = original_fully_shard + elif mode == FullyShardMode.COMPILED_COMPUTE: + torch._inductor.config.compile_threads = 1 + fully_shard_patch = fully_shard_with_compiled_compute # type: ignore[assignment] + else: + raise NotImplementedError( + f"Need to implement FullyShardMode={mode}" + ) + + # fully_shard is imported as a global + # through `from ... import fully_shard` + func.__globals__[original_fully_shard.__name__] = fully_shard_patch + func(*args, **kwargs) + # other threads use patched func before this thread restores + torch.distributed.barrier() + func.__globals__[original_fully_shard.__name__] = original_fully_shard + torch._inductor.config.compile_threads = original_compile_threads + + return wrapper + + return decorator + + +class SkipModule(nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin = nn.Linear(10, 10, bias=False) + + def forward(self, x): + return self.lin(x) + + +class NestedLinear(nn.Module): + def __init__(self, fsdp_wrap): + super().__init__() + if fsdp_wrap: + self.nested_linear = wrap(nn.Linear(10, 10, bias=False).to(DEVICE_TYPE)) + else: + self.nested_linear = nn.Linear(10, 10, bias=False).to(DEVICE_TYPE) + + def forward(self, x): + return self.nested_linear(x) + + +class SkipModel(nn.Module): + def __init__(self, double_nest): + super().__init__() + self.linear = nn.Linear(10, 10, bias=False).to(DEVICE_TYPE) + self.linear_skip = SkipModule().to(DEVICE_TYPE) + self.nested_linear = wrap( + NestedLinear(fsdp_wrap=double_nest), device_id=DEVICE_TYPE + ) + + def forward(self, x): + x = self.linear(x) + x = self.linear_skip(x) + x = self.nested_linear(x) + return x diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_jit.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_jit.py new file mode 100644 index 0000000000000000000000000000000000000000..5381752d64142e9275bd01ccd9ac3f9230a499ea --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_jit.py @@ -0,0 +1,322 @@ +# mypy: ignore-errors + +# Torch +import torch +import torch.cuda +import torch.jit +import torch.jit._logging +import torch.jit.frontend +import torch.jit.quantized + +# Testing utils +from torch.testing._internal.common_dtype import floating_and_complex_types_and +from torch.testing._internal.common_utils import TestCase, \ + freeze_rng_state, TemporaryFileName, enable_profiling_mode_for_profiling_tests, is_iterable_of_tensors +from torch.testing._internal.common_utils import enable_profiling_mode # noqa: F401 + +# Standard library +from itertools import chain +from torch._C import TensorType + +import io + +def check_output_types(self, func, ref_outputs, args, kwargs): + graph = getattr(func, 'last_graph', None) + types = [o.type() for o in graph.outputs()] + self.assertTrue(len(types) == 1) + t = types[0] + torch._C._jit_assert_is_instance(ref_outputs, t) + +# Test names in this set are only checked for a single derivative +nn_functional_single_grad = frozenset('test_nn_' + name for name in [ + 'pdist', + 'multilabel_margin_loss', + 'max_unpool3d', + 'multi_margin_loss', + 'binary_cross_entropy', + 'binary_cross_entropy_size_average', + 'ctc_loss', + 'grid_sample', +]) + +def check_against_reference(self, func, reference_func, output_func, args, kwargs=None, + allow_unused=True, check_types=True, no_grad=False, no_gradgrad=False): + """Verifies a function performs identically to some reference implementation. + + Commonly, this is used to verify that a JIT implementation + (output_func) matches the behavior of the eager implementation + (reference_func). + """ + kwargs = kwargs if kwargs else {} + + def allSum(vs): + if isinstance(vs, torch.Tensor): + vs = (vs,) + return sum((i + 1) * v.sum().abs() if v.dtype.is_complex else (i + 1) * v.sum() + for i, v in enumerate(vs) + if v is not None and v.dtype in floating_and_complex_types_and(torch.half, torch.bfloat16)) + + def clone_tensor(t, preserve_requires_grad): + require_grad = preserve_requires_grad and t.requires_grad + return t.detach().clone().requires_grad_(require_grad) + + def clone_inputs(preserve_requires_grad: bool): + inputs: list[torch.Tensor | list[torch.Tensor]] = [] + + for arg in args: + if isinstance(arg, torch.Tensor): + inputs.append(clone_tensor(arg, preserve_requires_grad)) + elif is_iterable_of_tensors(arg): + inputs.append([clone_tensor(t, preserve_requires_grad) for t in arg]) + else: + inputs.append(arg) + + return inputs + + # Returns tensors in args that requires_grad, including tensors in TensorList args + def get_recording_tensors(args): + recording_tensors: list[torch.Tensor] = [] + + for arg in args: + if isinstance(arg, torch.Tensor) and arg.requires_grad: + recording_tensors.append(arg) + elif is_iterable_of_tensors(arg): + recording_tensors.extend(filter(lambda t: t.requires_grad, arg)) + + return recording_tensors + + # test no gradients case + nograd_inputs = clone_inputs(preserve_requires_grad=False) + outputs = self.runAndSaveRNG(reference_func, nograd_inputs, kwargs) + with enable_profiling_mode_for_profiling_tests(): + outputs_test = self.runAndSaveRNG(func, nograd_inputs, kwargs) + self.assertEqual(outputs, outputs_test) + + if check_types: + check_output_types(self, func, outputs_test, nograd_inputs, kwargs) + + if no_grad: + # skip grad tests + return + + with enable_profiling_mode_for_profiling_tests(): + # test single grad case + recording_inputs = clone_inputs(preserve_requires_grad=True) + recording_tensors = get_recording_tensors(recording_inputs) + outputs = output_func(self.runAndSaveRNG(reference_func, recording_inputs, kwargs)) + grads = torch.autograd.grad(allSum(outputs), recording_tensors, + allow_unused=allow_unused) + outputs_test = output_func(self.runAndSaveRNG(func, recording_inputs, kwargs)) + grads_test = torch.autograd.grad(allSum(outputs_test), recording_tensors, + allow_unused=allow_unused) + self.assertEqual(outputs, outputs_test) + self.assertEqual(grads, grads_test) + # test the grad grad case + if self._testMethodName in nn_functional_single_grad or no_gradgrad: + return + + outputs = output_func(self.runAndSaveRNG(reference_func, recording_inputs, kwargs)) + l1 = allSum(outputs) + grads = torch.autograd.grad(l1, recording_tensors, create_graph=True, + allow_unused=allow_unused) + + l2 = (allSum(grads) * l1) + grads2 = torch.autograd.grad(l2, recording_tensors, allow_unused=allow_unused) + recording_inputs = clone_inputs(preserve_requires_grad=True) + recording_tensors = get_recording_tensors(recording_inputs) + outputs_test = output_func(self.runAndSaveRNG(func, recording_inputs, kwargs)) + l1_test = allSum(outputs_test) + grads_test = torch.autograd.grad( + l1_test, recording_tensors, create_graph=True, allow_unused=allow_unused) + + l2_test = (allSum(grads_test) * l1_test) + grads2_test = torch.autograd.grad(l2_test, recording_tensors, allow_unused=allow_unused) + + self.assertEqual(outputs, outputs_test) + self.assertEqual(grads, grads_test) + for g2, g2_test in zip(grads2, grads2_test, strict=True): + if g2 is None and g2_test is None: + continue + self.assertEqual(g2, g2_test, atol=5e-4, rtol=1e-4) + +class JitCommonTestCase(TestCase): + def createFunctionFromGraph(self, trace): + graph = trace if isinstance(trace, torch._C.Graph) else trace.graph() + return torch._C._create_function_from_graph("forward", graph) + + def assertExportImport(self, trace, inputs): + m = self.createFunctionFromGraph(trace) + self.assertExportImportModule(m, inputs) + + def assertExportImportModule(self, m, inputs): + m_import = self.getExportImportCopy(m) + a = self.runAndSaveRNG(m, inputs) + b = self.runAndSaveRNG(m_import, inputs) + self.assertEqual(a, b, "Results of original model and " + "exported/imported version of model differed") + + def runAndSaveRNG(self, func, inputs, kwargs=None): + kwargs = kwargs if kwargs else {} + with freeze_rng_state(): + results = func(*inputs, **kwargs) + return results + + def getExportImportCopy(self, m, also_test_file=True, map_location=None): + buffer = io.BytesIO() + torch.jit.save(m, buffer) + buffer.seek(0) + imported = torch.jit.load(buffer, map_location=map_location) + + if not also_test_file: + return imported + + with TemporaryFileName() as fname: + torch.jit.save(imported, fname) + return torch.jit.load(fname, map_location=map_location) + + def autoDiffErrorMessage(self, should_autodiff_node, nodes_not_in_diff_graph, + fusion_nodes_not_found, non_fusible_nodes_being_fused, + fusion_nodes_found, nodes_in_diff_graph): + err_msg = "\nFailure in testing nodes' autodifferentiation. " + if should_autodiff_node: + err_msg += "One or more nodes were expected to be autodiffed, " \ + "but were not found in specified fusible/nonfusible " \ + "DifferentiableGraph groups. \nSpecifically:" + # The node is intended to appear in a differentiable graph but doesn't + diff_nodes_missing = [] + # The node is intended to appear in a differentiable graph + # outside of a fusion group but instead is in a fusion group + diff_nodes_in_fusion = [] + # The node is intended to appear in a fusion group but doesn't + fusion_nodes_missing = [] + # The node is intended to appear in a fusion group but instead + # is just in an outer differentiable graph + fusion_nodes_in_diff = [] + for node in nodes_not_in_diff_graph: + if node in non_fusible_nodes_being_fused: + diff_nodes_in_fusion.append(node) + else: + diff_nodes_missing.append(node) + for node in fusion_nodes_not_found: + if node in nodes_in_diff_graph: + fusion_nodes_in_diff.append(node) + else: + fusion_nodes_missing.append(node) + if len(diff_nodes_missing) > 0: + err_msg += f"\n {diff_nodes_missing} were not in one of the " \ + "DifferentiableGraphs when they were expected to be. " \ + "Did you intend for these nodes to be autodiffed? " \ + "If not, remove them from the list of nonfusible nodes." + if len(diff_nodes_in_fusion) > 0: + err_msg += f"\n {diff_nodes_in_fusion} were found in one of the FusionGroups " \ + "when they were expected to be just in a DifferentiableGraph. If it was " \ + "intended for these nodes to be in FusionGroups, reclassify these nodes as " \ + "fusible nodes. If these nodes were not intended to be fused, your " \ + "autodifferentiation logic might be wrong." + if len(fusion_nodes_missing) > 0: + err_msg += f"\n {fusion_nodes_missing} were not in one of the FusionGroups " \ + "of the DifferentiableGraphs when they were expected to be. " \ + "They were also not found in an outer DifferentiableGraph. Did you " \ + "intend for these nodes to be autodifferentiated? If not, you should " \ + "remove these nodes from the test's fusible nodes. Otherwise your " \ + "autodifferentiation logic might be wrong." + if len(fusion_nodes_in_diff) > 0: + err_msg += f"\n {fusion_nodes_in_diff} were not in one of the FusionGroups " \ + "of the DifferentiableGraphs when they were expected to be, " \ + "instead they were found just in an outer DifferentiableGraph. " \ + "Did you intend for these nodes to be fused? If not, you should " \ + "move these nodes into the test's nonfusible nodes. Otherwise your " \ + "autodifferentiation logic might be wrong." + else: + err_msg += "One or more nodes were not expected to be autodiffed " \ + "but were found in a DifferentiableGraph or in a FusionGroup " \ + "of a DifferentiableGraph. Did you intend for these nodes to be " \ + "autodiffed? If so, change this test to expect autodifferentiation. " \ + "\nSpecifically:" + if len(fusion_nodes_found) > 0: + err_msg += f"\n {fusion_nodes_found} were not expected to be in " \ + "one of the DifferentiableGraphs, but appeared in a FusionGroup " \ + "of a DifferentiableGraph. " + if len(nodes_in_diff_graph) > 0: + err_msg += f"\n {nodes_in_diff_graph} were not expected to " \ + "be in one of the DifferentiableGraphs but were." + return err_msg + + def assertAutodiffNode(self, graph, should_autodiff_node, nonfusible_nodes, fusible_nodes): + diff_nodes = graph.findAllNodes('prim::DifferentiableGraph') + diff_subgraphs = [node.g('Subgraph') for node in diff_nodes] + + # Note: currently no tests have fusible_nodes + fusion_nodes = list(chain.from_iterable([g.findAllNodes('prim::FusionGroup') for g in diff_subgraphs])) + fusion_subgraphs = [node.g('Subgraph') for node in fusion_nodes] + + # For any non-fusible node, it must show up in one of the DifferentiableGraphs. + nodes_in_diff_graph = [] + nodes_not_in_diff_graph = [] + non_fusible_nodes_being_fused = [] + for node in nonfusible_nodes: + if any(g.findNode(node) is not None for g in diff_subgraphs): + nodes_in_diff_graph.append(node) + else: + nodes_not_in_diff_graph.append(node) + if any(g.findNode(node) is not None for g in fusion_subgraphs): + non_fusible_nodes_being_fused.append(node) + found_all_nonfusible_nodes = len(nodes_in_diff_graph) == len(nonfusible_nodes) + + # For any fusible node, it must show up in one of the FusionGroups in one of the DifferentiableGraphs. + fusion_nodes_found = [] + fusion_nodes_not_found = [] + for node in fusible_nodes: + if any(g.findNode(node) is not None for g in fusion_subgraphs): + fusion_nodes_found.append(node) + else: + fusion_nodes_not_found.append(node) + found_all_fusible_nodes = len(fusion_nodes_found) == len(fusible_nodes) + + if should_autodiff_node is not None: + err_msg = self.autoDiffErrorMessage(should_autodiff_node, + nodes_not_in_diff_graph, + fusion_nodes_not_found, + non_fusible_nodes_being_fused, + fusion_nodes_found, + nodes_in_diff_graph) + self.assertEqual(should_autodiff_node, + found_all_nonfusible_nodes and found_all_fusible_nodes, err_msg) + + def checkShapeAnalysis(self, out_sizes: list[int] | list[list[int]], + traced_graph, assert_propagation, constant_prop=True): + # repropagte input shapes provided by tracing, + prev_symbolic_shapes_test_enabled = torch._C._jit_symbolic_shapes_test_mode_enabled() + for enable_test_mode in [True, False]: + # here we are testing allowing/disallowing substituting in complete shapes as constants, + # disallowing constants helps stress test partial eval and substitution pipeline + torch._C._jit_set_symbolic_shapes_test_mode(enable_test_mode) + torch._C._jit_erase_non_input_shape_information(traced_graph) + if constant_prop: + torch._C._jit_pass_constant_propagation(traced_graph) + torch._C._jit_pass_propagate_shapes_on_graph(traced_graph) + # Add sizes to default tensor type to avoid checking something out of scope + # and difficulties with tracer leaving in other parts of tensor type + output = next(traced_graph.outputs()).type() + + def test_type(type, actual_size): + sizes = type.symbolic_sizes() + out_type = TensorType.get().with_sizes(sizes) + actual_type = TensorType.get().with_sizes(actual_size) + + # always check actual shape is a subtype of the output + self.assertTrue(actual_type.isSubtypeOf(out_type)) + + # and then if assertion flag is provided, check shape analysis + # is successful + if assert_propagation: + self.assertEqual(out_type.sizes(), actual_size) + + if output.isSubtypeOf(torch._C.TensorType.get()): + test_type(output, out_sizes) + else: + tuple_elements = output.elements() + for i in range(len(tuple_elements)): + test_type(tuple_elements[i], out_sizes[i]) + + torch._C._jit_set_symbolic_shapes_test_mode(prev_symbolic_shapes_test_enabled) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_methods_invocations.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_methods_invocations.py new file mode 100644 index 0000000000000000000000000000000000000000..9b1665c67d0b3fb0d64a938906d2a861d2db2c38 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_methods_invocations.py @@ -0,0 +1,26902 @@ +# mypy: ignore-errors + +from functools import wraps, partial +from itertools import product, chain, islice +import itertools +import functools +import copy +import operator +import random +import unittest +import math +import enum + +import torch +import numpy as np +import numpy.typing as npt +from torch import inf, nan + +from typing import Any +from collections.abc import Sequence +from torch.testing import make_tensor +from torch.testing._internal.common_dtype import ( + _dispatch_dtypes, floating_types, floating_types_and, complex_types, floating_and_complex_types, + floating_and_complex_types_and, all_types_and_complex_and, all_types_and, all_types_and_complex, integral_types_and, + empty_types, complex_types_and, integral_types, custom_types, all_types_complex_float8_and, float8_types, + highest_precision_complex, + highest_precision_float, +) +from torch.testing._internal.common_device_type import ( + onlyCPU, onlyCUDA, onlyNativeDeviceTypes, disablecuDNN, skipCUDAIfNoMagma, skipCUDAIfNoMagmaAndNoCusolver, + skipCUDAIfNoCusolver, skipCPUIfNoLapack, skipCPUIfNoFFT, skipCUDAIf, precisionOverride, + skipCPUIfNoMklSparse, toleranceOverride, tol, skipXPU, e4m3_type, E4M3_MAX_POS, E5M2_MAX_POS, +) +from torch.testing._internal.common_cuda import ( + PLATFORM_SUPPORTS_FLASH_ATTENTION, PLATFORM_SUPPORTS_MEM_EFF_ATTENTION, + SM53OrLater, SM80OrLater, SM89OrLater, with_tf32_off, TEST_CUDNN, + _get_torch_cuda_version, +) +from torch.testing._internal.common_quantized import ( + _bfloat16_to_float4_e2m1fn_x2, +) +from torch.testing._internal.common_utils import ( + make_fullrank_matrices_with_distinct_singular_values, + IS_ARM64, + IS_CPU_EXT_SVE_SUPPORTED, + TEST_WITH_ROCM, IS_FBCODE, IS_WINDOWS, IS_MACOS, TEST_SCIPY, + torch_to_numpy_dtype_dict, numpy_to_torch_dtype, TEST_WITH_ASAN, + GRADCHECK_NONDET_TOL, slowTest, TEST_WITH_SLOW, + TEST_WITH_TORCHINDUCTOR, +) +from torch.testing._utils import wrapper_set_seed + +import torch._refs as refs # noqa: F401 +import torch._refs.nn.functional +import torch._refs.special +import torch._refs.linalg +import torch._prims as prims # noqa: F401 +from torch.utils import _pytree as pytree + + +from torch._vendor.packaging import version + +from torch.testing._internal.opinfo.core import ( # noqa: F401 + L, + M, + S, + XS, + _NOTHING, + _getattr_qual, + DecorateInfo, + SampleInput, + ErrorInput, + AliasInfo, + NumericsFilter, + OpInfo, + _generate_reduction_inputs, + _generate_reduction_kwargs, + sample_inputs_reduction, + ReductionOpInfo, + reference_inputs_elementwise_binary, + make_error_inputs_elementwise_binary, + generate_elementwise_binary_tensors, + generate_elementwise_binary_arbitrarily_strided_tensors, + generate_elementwise_binary_small_value_tensors, + generate_elementwise_binary_large_value_tensors, + generate_elementwise_binary_extremal_value_tensors, + generate_elementwise_binary_broadcasting_tensors, + generate_elementwise_binary_with_scalar_samples, + generate_elementwise_binary_with_scalar_and_type_promotion_samples, + generate_elementwise_binary_noncontiguous_tensors, + sample_inputs_elementwise_binary, + BinaryUfuncInfo, + sample_inputs_elementwise_unary, + generate_elementwise_unary_tensors, + generate_elementwise_unary_small_value_tensors, + generate_elementwise_unary_large_value_tensors, + generate_elementwise_unary_extremal_value_tensors, + reference_inputs_elementwise_unary, + UnaryUfuncInfo, + sample_inputs_spectral_ops, + SpectralFuncType, + SpectralFuncInfo, + ShapeFuncInfo, + sample_inputs_foreach, + ForeachFuncInfo, + gradcheck_wrapper_hermitian_input, + gradcheck_wrapper_ctc_loss, + gradcheck_wrapper_triangular_input, + gradcheck_wrapper_triangular_input_real_positive_diagonal, + gradcheck_wrapper_masked_operation, + gradcheck_wrapper_masked_pointwise_operation, + clone_sample, +) +from torch.testing._internal.opinfo.refs import ( # NOQA: F401 + _find_referenced_opinfo, + _inherit_constructor_args, + PythonRefInfo, + ReductionPythonRefInfo, + ElementwiseUnaryPythonRefInfo, + ElementwiseBinaryPythonRefInfo, +) +from torch.testing._internal.opinfo.utils import ( + np_unary_ufunc_integer_promotion_wrapper, + reference_reduction_numpy, + prod_numpy +) +from torch.testing._internal import opinfo +from torch.testing._internal.opinfo.definitions.linalg import ( + sample_inputs_linalg_cholesky, + sample_inputs_linalg_cholesky_inverse, + sample_inputs_cross, + sample_inputs_linalg_qr_geqrf, + sample_inputs_linalg_invertible, + sample_inputs_lu_solve, + sample_inputs_legacy_solve, + sample_inputs_svd, + sample_inputs_linalg_det_logdet_slogdet, + sample_inputs_linalg_lu, + sample_inputs_diagonal_diag_embed, + error_inputs_diagonal_diag_embed, +) +from torch.testing._internal.opinfo.definitions.special import ( + sample_inputs_i0_i1, + sample_inputs_polygamma, + reference_polygamma, +) +from torch.testing._internal.opinfo.definitions._masked import ( + sample_inputs_softmax_variant, +) +from torch.testing._internal.opinfo.definitions.sparse import ( + error_inputs_sparse_like_fns, + sample_inputs_sparse_like_fns, + error_inputs_sparse_mul, + sample_inputs_sparse_mul, + error_inputs_sparse_reduction_sum, + sample_inputs_sparse_reduction_sum +) + +if TEST_SCIPY: + from scipy import stats + import scipy.spatial + import scipy.special + + +def round_up(x: int, y: int) -> int: + return ((x + y - 1) // y) * y + + +# test if a tensor is close to an integer +def close_to_int(x, eps=0.1): + if x.is_complex(): + y = torch.abs(torch.view_as_complex(torch.frac(torch.view_as_real(x)))) + else: + y = torch.abs(torch.frac(x)) + return (y < eps) | (y > (1 - eps)) + + +def sample_inputs_slice(op_info, device, dtype, requires_grad, **kwargs): + + make_input = partial(make_tensor, device=device, dtype=dtype, + low=None, high=None, requires_grad=requires_grad) + + yield SampleInput(make_input(3), 0) + + yield SampleInput(make_input(20, 30, 40), dim=1, start=1, end=-2) + + yield SampleInput(make_input(20, 30, 40), dim=1, start=1, end=-2, step=3) + + yield SampleInput(make_input(20, 30, 40), dim=0, start=-10, end=-2, step=2) + + +def sample_inputs_tensor_split(op_info, device, dtype, requires_grad, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, + low=None, high=None, requires_grad=requires_grad) + + args_cases = ( + # Cases with tensor indices. + (torch.tensor([1, 2, 3]),), + (torch.tensor(1),), + (torch.tensor([1, 2, 3]), 1), + (torch.tensor([1, 4, 2, 5, 3, 6])[::2], 1), + # Cases with list of indices. + ((2, 4),), + ((2, 4), 1), + ((2, 4), -1), + # Cases with integer section. + (3,), + (3, 1), + (3, -1), + ) + + for args in args_cases: + yield SampleInput(make_input((S, S, S)), args=args) + + +def sample_inputs_hsplit(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, + low=None, high=None, requires_grad=requires_grad) + yield SampleInput(make_arg(6), 2) + yield SampleInput(make_arg(S, S, S), [1, 2, 3]) + +def sample_inputs_vsplit(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, + low=None, high=None, requires_grad=requires_grad) + yield SampleInput(make_arg(6, S), 2) + yield SampleInput(make_arg(S, S, S), [1, 2, 3]) + +def sample_inputs_dsplit(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, + low=None, high=None, requires_grad=requires_grad) + yield SampleInput(make_arg(S, S, S), [1, 2, 3]) + yield SampleInput(make_arg(S, S, 6), 2) + +def error_inputs_hsplit(op_info, device, **kwargs): + make_arg = partial(make_tensor, dtype=torch.float32, device=device) + err_msg1 = ("torch.hsplit requires a tensor with at least 1 dimension, " + "but got a tensor with 0 dimensions!") + yield ErrorInput(SampleInput(make_arg(()), 0), error_regex=err_msg1) + + err_msg2 = (f"torch.hsplit attempted to split along dimension 1, " + f"but the size of the dimension {S} " + f"is not divisible by the split_size 0!") + yield ErrorInput(SampleInput(make_arg((S, S, S)), 0), error_regex=err_msg2) + + # Incorrect type for indices_or_section argument + err_msg3 = ("received an invalid combination of arguments.") + yield ErrorInput( + SampleInput(make_arg((S, S, S)), "abc"), + error_type=TypeError, error_regex=err_msg3) + +def error_inputs_vsplit(op_info, device, **kwargs): + make_arg = partial(make_tensor, dtype=torch.float32, device=device) + err_msg1 = ("torch.vsplit requires a tensor with at least 2 dimension, " + "but got a tensor with 1 dimensions!") + yield ErrorInput(SampleInput(make_arg(S), 0), error_regex=err_msg1) + + err_msg2 = (f"torch.vsplit attempted to split along dimension 0, " + f"but the size of the dimension {S} " + f"is not divisible by the split_size 0!") + yield ErrorInput(SampleInput(make_arg(S, S, S), 0), + error_regex=err_msg2) + + # Incorrect type for indices_or_section argument + err_msg3 = ("received an invalid combination of arguments.") + yield ErrorInput(SampleInput(make_arg(S, S, S), "abc"), + error_type=TypeError, error_regex=err_msg3) + +def error_inputs_dsplit(op_info, device, **kwargs): + make_arg = partial(make_tensor, dtype=torch.float32, device=device) + err_msg1 = ("torch.dsplit requires a tensor with at least 3 dimension, " + "but got a tensor with 1 dimensions!") + yield ErrorInput(SampleInput(make_arg(S), 0), error_regex=err_msg1) + + err_msg2 = (f"torch.dsplit attempted to split along dimension 2, " + f"but the size of the dimension {S} " + f"is not divisible by the split_size 0!") + yield ErrorInput(SampleInput(make_arg(S, S, S), 0), error_regex=err_msg2) + + +def sample_inputs_as_strided(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # input shape, output shape, output stride, output storage offset + test_cases = ( + ((1,), (1,), (1,), 0), + ((3, 3), (2, 2), (1, 2), 0), + ((3, 3), (2, 2), (1, 2), 1), + ((16,), (2, 2, 2, 2), (1, 1, 1, 1), 0), + ((16,), (2, 1, 1, 2), (1, 7, 7, 1), 0), + ) + + for input_shape, output_shape, stride, storage_offset in test_cases: + input_t = make_arg(input_shape) + kwargs = dict(storage_offset=storage_offset) + yield SampleInput(input_t, args=(output_shape, stride), kwargs=kwargs) + +def sample_inputs_as_strided_partial_views(op_info, device, dtype, requires_grad, **kwargs): + def make_arg(): + base = make_tensor((20,), device=device, dtype=dtype) + return base[5:15].requires_grad_(requires_grad) + + # as_strided on offset, partial views + yield SampleInput(make_arg(), (2, 2), (1, 2)) + yield SampleInput(make_arg(), (2, 2), (1, 2), storage_offset=0) + yield SampleInput(make_arg(), (2, 2), (1, 2), storage_offset=10) + +def sample_inputs_as_strided_scatter(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # input shape, output shape, output stride, output storage offset + test_cases = [ + ((1,), (), (), 0), + ((1,), (1,), (1,), 0), + ((3, 3), (2, 2), (1, 2), 0), + ((3, 3), (2, 2), (1, 2), 1), + ((3, 3), (2, 2), (2, 1), 0), + # Scatter to larger dimensions + ((16,), (2, 2, 2, 2), (8, 4, 2, 1), 0), + # Scatter to larger dimensions with strides inverted + ((16,), (2, 1, 1, 2), (1, 2, 4, 8), 0), + ] + + for input_shape, output_shape, stride, storage_offset in test_cases: + input_t = make_arg(input_shape) + input_src = make_arg(output_shape) + yield SampleInput(input_t, input_src, output_shape, stride, storage_offset=storage_offset) + + +def error_inputs_as_strided_scatter(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32, requires_grad=False) + + # Create a small tensor and try to scatter it out of bounds + input_t = make_arg([4, 4]) + input_src = make_arg([2, 2]) + yield ErrorInput( + SampleInput(input_t, input_src, [2, 2], [200, 200], storage_offset=0), + error_regex="itemsize 4 requiring a storage size of 1604 are out of bounds for storage of size 64" + ) + + +def sample_inputs_combinations(op_info, device, dtype, requires_grad, **kwargs): + inputs = ( + (0,), + (0, 1), + (0, 1, 2, 3), + ) + + rvals = [1, 2, 4] + + products = product(inputs, rvals, [False, True]) + + for input_data, r, with_replacement in products: + input_t = torch.tensor(input_data, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(input_t, r=r, with_replacement=with_replacement) + +def sample_inputs_cartesian_prod(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(torch.tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # constructs 1-D tensors with varying number of elements + a = make_arg((0,)) + b = make_arg((0, 1)) + c = make_arg((0, 1, 2, 3)) + + # sample with only 1 tensor + yield SampleInput(a) + + # sample with 2 tensors + yield SampleInput(a, b) + + # sample with 3 tensors + yield SampleInput(a, b, c) + +def sample_inputs_cosine_similarity(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as input_shape, dict of dim and eps + cases: tuple[tuple, dict] = ( # type: ignore[assignment] + ((S, S), {'dim': 1}), + ((S, 2), {'dim': -1}), + ((S,), {'dim': 0, 'eps': 0.5}), + ((), {'dim': 0}), + ((S, S, M), {'dim': 2}), + ((S, S), {}) + ) + + for input_shape, kwargs in cases: + yield SampleInput(make_arg(input_shape), args=(make_arg(input_shape),), kwargs=kwargs) + # Test for Broadcasting + yield SampleInput(make_arg((1, 2, 3)), args=(make_arg((2, 1, 3)),), kwargs={'dim': -1}) + yield SampleInput(make_arg((1, 2, 3)), args=(make_arg((2, 1, 3)),), kwargs={'dim': -2}) + yield SampleInput(make_arg((2, 3)), args=(make_arg((2, 1, 3)),), kwargs={'dim': -1}) + + +def sample_inputs_item(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=False) + + cases = ( + (), + (()), + (1), + ((1,)), + ) + + for shape in cases: + yield SampleInput(make_arg(shape)) + +def error_inputs_item(op, device, **kwargs): + make_arg = partial(make_tensor, dtype=torch.float32, device=device, requires_grad=False) + + cases = ( + (M), + ((S,)), + (S, S), + (S, M, L), + ) + + for shape in cases: + yield ErrorInput( + SampleInput(make_arg(shape)), error_type=RuntimeError, + error_regex="elements cannot be converted to Scalar") + + +def sample_inputs_batch_norm(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_arg_without_requires_grad = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + # Ordered as: input shape, kwargs for training, momentum, eps + cases: tuple[tuple[int, ...], dict] = ( + ((S, S, S), {'training': True, 'momentum': 0.5, 'eps': 0.6}), + ((3, 2, 4), {'training': False, 'momentum': -1.2}), + ((3, 1), {'training': True, 'momentum': 0.0}), + ((0,), {'training': True}), + ((0,), {'training': False}), + ((3, 2, 3, 4), {'training': True, 'momentum': -1.0, 'eps': 0.5}), + ((3, 2, 3, 4), {'training': False, 'momentum': -1.0, 'eps': 0.5}), + ((2, 1), {}), + ) + + for input_shape, kwargs in cases: + # args: running mean, running var, weight and bias should necessarily be of shape: (channels,) + channels = input_shape[1] if len(input_shape) > 1 else 0 + weight = make_arg(channels) if channels > 0 else None + bias = make_arg(channels) if channels > 0 else None + running_mean = make_arg_without_requires_grad(channels, low=0) + running_var = make_arg_without_requires_grad(channels, low=0) + + yield SampleInput( + make_arg(input_shape), + args=( + running_mean, + running_var, + weight, + bias + ), + kwargs=kwargs + ) + + # Checking for permutations of weights and biases as `None` + is_training = [True, False, False] + + for training in is_training: + yield SampleInput( + make_arg(input_shape), + args=( + running_mean, + running_var, + make_arg(channels), + make_arg(channels) + ), + kwargs={'training': training} + ) + + # Test case for no optional kwargs + # running_mean and running_var are required in evaluation mode (training: False) but not in training mode + yield SampleInput(make_arg((1, 2, 3)), args=(None, None, None, None), kwargs={'training': True}) + +def sample_inputs_softmax_backward_data(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + cases = [ + ((S,), 0), + ((S, S), 0), + ((S, M, S), -1), + ] + input_dtypes = [dtype] + if dtype == torch.float and device == 'cuda': + input_dtypes += [torch.float16] + + for (shape, dim), input_dtype in product(cases, input_dtypes): + input = make_arg(shape) + output = torch.nn.functional.softmax(input, dim=dim, dtype=input_dtype) + yield SampleInput(make_arg(shape), output, dim, input_dtype) + +def sample_inputs_native_batch_norm(op_info, device, dtype, requires_grad, **kwargs): + samples = sample_inputs_batch_norm(op_info, device, dtype, requires_grad, **kwargs) + for sample in samples: + # torch.native_batch_norm does not support 0 numel tensors + # IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1) + if sample.input.numel() == 0: + continue + args = sample.args + training = sample.kwargs.get('training', True) + momentum = sample.kwargs.get('momentum', 0.5) + eps = sample.kwargs.get('eps', 1e-5) + yield SampleInput(sample.input, args=(args[2], args[3], args[0], args[1], training, momentum, eps)) + + +def sample_inputs__native_batch_norm_legit(op_info, device, dtype, requires_grad, **kwargs): + samples = sample_inputs_batch_norm(op_info, device, dtype, requires_grad, **kwargs) + for sample in samples: + # torch.native_batch_norm does not support 0 numel tensors + # IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1) + if sample.input.numel() == 0: + continue + args = sample.args + training = sample.kwargs.get('training', True) + momentum = sample.kwargs.get('momentum', 0.5) + eps = sample.kwargs.get('eps', 1e-5) + if args[0] is not None and args[1] is not None: + yield SampleInput(sample.input, args=(args[2], args[3], args[0], args[1], training, momentum, eps)) + else: + yield SampleInput(sample.input, args=(args[2], args[3], training, momentum, eps)) + +def sample_inputs__batch_norm_with_update(op_info, device, dtype, requires_grad, **kwargs): + samples = sample_inputs_batch_norm(op_info, device, dtype, requires_grad, **kwargs) + for sample in samples: + # torch.native_batch_norm does not support 0 numel tensors + # IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1) + if sample.input.numel() == 0: + continue + args = sample.args + momentum = sample.kwargs.get('momentum', 0.5) + eps = sample.kwargs.get('eps', 1e-5) + if any(args[i] is None for i in range(4)): + continue + yield SampleInput(sample.input, args=(args[2], args[3], args[0], args[1], momentum, eps)) + +def sample_inputs_nn_activation_relu(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases = ( + (()), + ((S, )), + ((S, S)), + ((S, M, S)) + ) + + for shape in cases: + yield SampleInput(make_arg(shape)) + +def sample_inputs_prelu(op_info, device, dtype, requires_grad, **kwargs): + op_kwargs = op_info.sample_kwargs(device, dtype, None)[0] + yield from sample_inputs_elementwise_unary(op_info, device, dtype, requires_grad, + op_kwargs=op_kwargs) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases = ( + (()), + ((S, )), + ((S, S)), + ((S, M, S)) + ) + + for shape in cases: + for weight in [-1., 0., 0.8, 1.]: + weight_tensor = torch.tensor(weight, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg(shape), args=(weight_tensor,)) + + channel_size = shape[1] if len(shape) >= 2 else 1 + yield SampleInput(make_arg(shape), args=(make_arg((channel_size,)),)) + + weight_tensor = torch.tensor(1., device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(make_arg((S, S)), kwargs=dict(weight=weight_tensor,)) + yield SampleInput(make_arg((S, S)), kwargs=dict(weight=make_arg((S,)),)) + +def reference_inputs_prelu(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_prelu(op, device, dtype, requires_grad, **kwargs) + yield from reference_inputs_elementwise_unary(op, device, dtype, requires_grad, **kwargs) + +def sample_kwargs_prelu_scalar_weight(device, dtype, input): + weight = torch.rand((), device=device, dtype=dtype) + # NumPy does not support bfloat16, so we default to float32 (only for NumPy) in that case + if dtype == torch.bfloat16: + weight_cpu = weight.to(dtype=torch.float32, device="cpu") + else: + weight_cpu = weight.cpu() + np_weight = weight_cpu.numpy() + return ({'weight': weight}, {'weight': np_weight}) + +def error_inputs_prelu(op, device): + # Weight has numel != 1, but self.ndim is zero-dim tensor + inp = make_tensor((), device=device, dtype=torch.float32) + weight = make_tensor((2,), device=device, dtype=torch.float32) + yield ErrorInput(SampleInput(inp, kwargs={'weight': weight}), + error_regex="Not allow zero-dim input tensor.") + + # Weight has numel != 1, but numel does not match channel size + inp = make_tensor((2, 8, 3), device=device, dtype=torch.float32) + weight = make_tensor((9,), device=device, dtype=torch.float32) + yield ErrorInput(SampleInput(inp, kwargs={'weight': weight}), + error_regex="Mismatch of parameter numbers and input channel size.") + + # Weight is neither a scalar nor 1-D tensor + inp = make_tensor((2, 8, 3), device=device, dtype=torch.float32) + weight = make_tensor((2, 4), device=device, dtype=torch.float32) + yield ErrorInput(SampleInput(inp, kwargs={'weight': weight}), + error_regex="prelu: Expected `weight` to be a scalar or 1D tensor, but got: ndim = 2") + + # src and index tensors must have the same # of dimensions +def sample_inputs_norm(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # ord = inf is tested in inputs_norm_inf as it fails on some tests + cases = [ + ((S, S), (2,), '2'), + ((S, S), (0,), '0'), + ((S, S), (0.5,), '0_5'), + ((S, S), (1,), '1'), + ((S, S), (3,), '3'), + ((S, S), (-1,), 'neg_1'), + ((S, S), (-2,), 'neg_2'), + ((S, S), (-0.5,), 'neg_0_5'), + ((S, S), (-1.5,), 'neg_1_5'), + ] + + cases_nonzero_input = ( + ((S, S, S), (1.5,), '1_5_default'), + ((S, S, S), (1.5, 1), '1_5_dim'), + ((S, S, S), (1.5, -1), '1_5_neg_dim'), + ((S, S, S), (1.5, 1, True), 'keepdim_1_5_dim'), + ((S, S, S), (1.5, -1, True), 'keepdim_1_5_neg_dim'), + ) + + cases_posdim = ( + ((S, S), (-2, 1,), 'neg_2_dim'), + ((S, S), (-1, 1,), 'neg_1_dim'), + ((S, S), (0, 1,), '0_dim'), + ((S, S), (1, 1,), '1_dim'), + ((S, S), (2, 1,), '2_dim'), + ((S, S), (3, 1,), '3_dim'), + ((S, S, S), (2, 1), '2_dim'), + ((S, S, S), (3, 1), '3_dim'), + ((S, S, S), (2, 1, True), 'keepdim_2_dim'), + ((S, S, S), (3, 1, True), 'keepdim_3_dim'), + ((), (2, 0), '2_dim_scalar'), + ((), (3, 0), '3_dim_scalar'), + ((), (2, 0, True), 'keepdim_2_dim_scalar'), + ((), (3, 0, True), 'keepdim_3_dim_scalar'), + ) + + cases_negdim = ((shape, args[:1] + (-args[1],) + args[2:], name.replace("_dim", "_neg_dim")) + for shape, args, name in cases_posdim) + + for shape, args, name in itertools.chain(cases, cases_posdim, cases_negdim): + yield SampleInput(make_arg(shape), args=args, name=name) + + for shape, args, name in cases_nonzero_input: + yield SampleInput(make_arg(shape, exclude_zero=True), args=args, name=name) + + +def sample_inputs_norm_fro(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases = ( + ((S, S), (), 'default'), + ((S, S), ('fro',), 'fro_default'), + ((S, S), ('fro', [0, 1],), 'fro'), + ) + + for shape, args, name in cases: + yield SampleInput(make_arg(shape), args=args, name=name) + + +def sample_inputs_norm_nuc(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases = ( + ((S, S), ('nuc',), 'nuc'), + ((S, S, S), ('nuc', [1, 2]), 'nuc_batched'), + ) + + for shape, args, name in cases: + yield SampleInput(make_arg(shape), args=args, name=name) + + +def sample_inputs_norm_inf(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases = ( + ((S, S), (-inf,), '-inf'), + ((S, S), (inf,), 'inf'), + ((S, S), (inf, 1,), 'inf_2_dim'), + ((S, S), (inf, -1,), 'inf_2_neg_dim'), + ) + + for shape, args, name in cases: + yield SampleInput(make_arg(shape), args=args, name=name) + + +def sample_inputs_equal(op, device, dtype, requires_grad, **kwargs): + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + shapes = ( + ((), ()), + ((S,), ()), + ((), (S,)), + ((S, 1), (S,)), + ((M, S), ()), + ((S, S), (S, S)) + ) + + for shape_lhs, shape_rhs in shapes: + lhs = make_arg(shape_lhs) + rhs = make_arg(shape_rhs) + broadcasts_input = shape_lhs != torch.broadcast_shapes(shape_lhs, shape_rhs) + + yield SampleInput(lhs, args=(rhs,), broadcasts_input=broadcasts_input) + if shape_lhs == shape_rhs: + yield SampleInput(lhs, args=(lhs.clone().detach_(),)) + + +def sample_inputs_jiterator(op, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + shapes = ( + ((), ()), + ((S,), ()), + ((S, 1), (S,)), + ((M, S), ()), + ((S, M, S), (M, S)), + ((S, M, S), (S, M, S)), + ((M, 1, S), (M, S)), + ((M, 1, S), (1, M, S)), + ((0, 1, 3), (0, 10, 3)) + ) + + num_inputs = kwargs.get('num_inputs') + sample_kwargs = kwargs.get('sample_kwargs', {}) + + for shape_lhs, shape_rhs in shapes: + lhs = make_arg(shape_lhs) + args = [make_arg(shape_rhs) for _ in range(num_inputs - 1)] + broadcasts_input = (shape_lhs != torch.broadcast_shapes(shape_lhs, shape_rhs)) + + yield SampleInput(lhs, args=tuple(args), kwargs=sample_kwargs, broadcasts_input=broadcasts_input) + +def sample_inputs_broadcast_shapes(op, device, dtype, requires_grad, **kwargs): + shapes = ( + ((), ()), + ((S,), ()), + ((S, 1), (S,)), + ((S, 1), S), + ((M, S), ()), + ((S, M, S), (M, S)), + ((S, M, S), (S, M, S)), + ((M, 1, S), (M, S)), + ((M, 1, S), (1, M, S)), + ((0, 1, 3), (0, 10, 3)) + ) + + for shape in shapes: + inp, *arg0 = shape + yield SampleInput(inp, args=tuple(arg0)) + +def sample_inputs_add_sub(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_elementwise_binary(op, device, dtype, requires_grad, **kwargs) + + # Adds alpha kwarg cases + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + lhs = make_arg((S, S), **op.lhs_make_tensor_kwargs) + rhs = make_arg((S, S), **op.rhs_make_tensor_kwargs) + if dtype is not torch.bool: + yield SampleInput(lhs, args=(rhs,), kwargs={'alpha': 2}) + else: + yield SampleInput(lhs, args=(rhs,), kwargs={'alpha': True}) + neg_alpha = -3.125 if (dtype.is_floating_point or dtype.is_complex) else -3 + lhs = make_arg((S, S), **op.lhs_make_tensor_kwargs) + rhs = make_arg((S, S), **op.rhs_make_tensor_kwargs) + if dtype is not torch.bool: + yield SampleInput(lhs, args=(rhs,), kwargs={'alpha': neg_alpha}) + else: + yield SampleInput(lhs, args=(rhs,), kwargs={'alpha': False}) + +def error_inputs_arange(op, device, **kwargs): + yield ErrorInput(SampleInput(0, args=(3, 0)), error_type=RuntimeError, error_regex='step must be nonzero') + yield ErrorInput(SampleInput(0, args=(-3, 2)), error_type=RuntimeError, + error_regex='upper bound and lower bound inconsistent with step sign') + yield ErrorInput(SampleInput(0, args=(3, -2)), error_type=RuntimeError, + error_regex='upper bound and lower bound inconsistent with step sign') + yield ErrorInput(SampleInput(1549556900, args=(1549556828, 1989724)), error_type=RuntimeError, + error_regex='upper bound and lower bound inconsistent with step sign') + yield ErrorInput(SampleInput(0, args=(float('inf'), 2)), error_type=RuntimeError, error_regex='unsupported range') + yield ErrorInput(SampleInput(float('-inf'), args=(1, 2)), error_type=RuntimeError, error_regex='unsupported range') + +def sample_inputs_arange(op, device, dtype, requires_grad, **kwargs): + int_samples = ( + # positive direction + (-1, 2, 2), + # negative direction + (2, -3, -1), + # start == end + (1, 1, 1), + (1, 1, -1), + # divides evenly + (0, -8, -4), + (1, 5, 2), + # bool + (False, True, True), + # default step + (0, 1, None), + # default start + (None, 3, None), + ) + + def to_float(start, end, step): + start = start + 0.1 if start is not None else None + end = end + 0.1 + step = float(step) if step is not None else None + return start, end, step + + float_samples = ( + # includes endpoint + (0., -8. - 1e-6, -4.), + (1., 5. + 1e-6, 2.), + (0., -8., -4.), + (1., 5., 2.), + *(to_float(start, end, step) for (start, end, step) in int_samples), + ) + + large_samples = ( + (0, 10000, None), + ) + + samples = int_samples + float_samples + if dtype not in (torch.int8, torch.uint8): + samples += large_samples + + for start, end, step in samples: + if start is None: + if step is not None: + raise AssertionError("Expected step to be None when start is None") + # Pass end as positional arg + yield SampleInput(end, kwargs={"dtype": dtype, "device": device}) + # (Similar to) calling torch.arange(end=3) + yield SampleInput(0, kwargs={"end": end, "dtype": dtype, "device": device}) + elif step is None: + yield SampleInput(start, args=(end,), kwargs={"dtype": dtype, "device": device}) + else: + yield SampleInput(start, args=(end, step), kwargs={"dtype": dtype, "device": device}) + + yield SampleInput(2) + yield SampleInput(1, args=(3, 1)) + +def sample_inputs_randn(op, device, dtype, requires_grad, **kwargs): + shapes = ( + (M,), + (S, S) + ) + + for shape in shapes: + yield SampleInput(input=shape, kwargs=dict(dtype=dtype, device=device, requires_grad=requires_grad)) + +def sample_inputs_normal(op, device, dtype, requires_grad, **kwargs): + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=False) + samples = ( + ((S, S), 0, 5), + ((S, S, S), -2, 0.5), + ) + for shape, mean, std in samples: + yield SampleInput(make_arg(shape), args=(mean, std)) + +def error_inputs_normal(op, device, **kwargs): + t = torch.zeros([10], device=device) + invalid_std = -1 + yield ErrorInput( + SampleInput(t, args=(0, invalid_std)), + error_type=RuntimeError, + error_regex=fr"normal expects std >= 0.0, but found std {invalid_std}", + ) + +def sample_inputs_cauchy(op, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=False) + samples = ( + ((M,), 0, 0.5), + ((S, S), 0, 1), + ((S, S, S), -2, 1), + ) + for shape, median, gamma in samples: + yield SampleInput(make_arg(shape), args=(median, gamma)) + + +def error_inputs_cauchy(op, device, **kwargs): + t = torch.zeros([10], device=device) + invalid_scale = 0 + yield ErrorInput( + SampleInput(t, args=(0, invalid_scale,)), + error_type=RuntimeError, + error_regex=fr"cauchy_ expects sigma > 0.0, but found sigma={invalid_scale}", + ) + + +def sample_inputs_exponential(op, device, dtype, requires_grad, **kwargs): + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=False) + samples = ( + ((M,), 0.5), + ((S, S), 1), + ((S, S, S), 1.5), + ) + for shape, rate in samples: + yield SampleInput(make_arg(shape), args=(rate,)) + + +def error_inputs_exponential(op, device, **kwargs): + t = torch.zeros([10], device=device) + invalid_rate = 0 + yield ErrorInput( + SampleInput(t, args=(invalid_rate,)), + error_type=RuntimeError, + error_regex=fr"exponential_ expects lambda > 0.0, but found lambda={invalid_rate}", + ) + + +def sample_inputs_geometric(op, device, dtype, requires_grad, **kwargs): + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=False) + samples = ( + ((M,), 0.2), + ((S, S), 0.5), + ((S, S, S), 0.8), + ) + for shape, rate in samples: + yield SampleInput(make_arg(shape), args=(rate,)) + + +def error_inputs_geometric(op, device, **kwargs): + t = torch.zeros([10], device=device) + neg_prob = -1 + yield ErrorInput( + SampleInput(t, args=(neg_prob,)), + error_type=RuntimeError, + error_regex=fr"geometric_ expects p to be in \(0, 1\), but got p={neg_prob}", + ) + + +def sample_inputs_log_normal(op, device, dtype, requires_grad, **kwargs): + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=False) + samples = ( + ((M,), 0, 0.25), + ((S, S), 0.5, 1), + ((S, S, S), 0, 0.5), + ) + for shape, mean, std in samples: + yield SampleInput(make_arg(shape), args=(mean, std)) + + +def error_inputs_log_normal(op, device, **kwargs): + t = torch.zeros([10], device=device) + invalid_std = 0 + yield ErrorInput( + SampleInput(t, args=(0, invalid_std)), + error_type=RuntimeError, + error_regex=fr"log_normal_ expects std > 0.0, but found std={invalid_std}", + ) + + +def sample_inputs_uniform(op, device, dtype, requires_grad, **kwargs): + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=False) + samples = ( + ((M,), -100, 100), + ((S, S), 0, 1), + ((S, S, S), 1, 2), + ) + for shape, hi, lo in samples: + yield SampleInput(make_arg(shape), args=(hi, lo)) + +def sample_inputs_ones_zeros(op, device, dtype, requires_grad, **kwargs): + # this is a bit messy, as we want the args to be tuples + # so if we pass size as a tuple, we have a tuple containing a tuple + sizes = ( + (M,), + (S, S), + ) + for size in sizes: + yield SampleInput(size, kwargs={'dtype': dtype, 'device': device}) + +def sample_inputs_full(op, device, dtype, requires_grad, **kwargs): + def get_val(dtype): + return make_tensor([], dtype=dtype, device="cpu").item() + + sizes = ( + (M,), + (S, S), + ) + fill_values = [get_val(dtype), get_val(torch.int)] + + for size, fill_value in product(sizes, fill_values): + yield SampleInput(size, fill_value, dtype=dtype, device=device) + + +def error_inputs_uniform(op, device, **kwargs): + t = torch.zeros([10], device=device) + yield ErrorInput( + SampleInput(t, args=(3, -1)), + error_type=RuntimeError, + error_regex=r"uniform_ expects to return a \[from, to\) range, but found from=3 > to=-1", + ) + + +def error_inputs_linspace(op, device, **kwargs): + yield ErrorInput(SampleInput(0, args=(3, -1)), error_type=RuntimeError, error_regex='number of steps must be non-negative') + yield ErrorInput( + SampleInput(0, args=(3, 1.)), + error_type=TypeError, + error_regex="received an invalid combination of arguments - got \\(int, int, float", + ) + yield ErrorInput( + SampleInput(torch.tensor([1, 1], device=device), args=(torch.tensor([3, 3], device=device), 1)), + error_type=RuntimeError, + error_regex="only supports 0-dimensional start and end tensors" + ) + + +def sample_inputs_linspace(op, device, dtype, requires_grad, **kwargs): + ends = (-3, 0, 1, 4, 50) + starts = (-2., 0, 4.3, 50) + nsteps = (0, 1, 50) + # Extra case to replicate off-by-one issue on CUDA + cases = list(product(starts, ends, nsteps)) + [(0, 7, 50)] + for start, end, nstep in cases: + if dtype == torch.uint8 and (end < 0 or start < 0): + continue + yield SampleInput(start, args=(end, nstep), kwargs={"dtype": dtype, "device": device}) + + yield SampleInput(1, args=(3, 1)) + + +def sample_inputs_linspace_tensor_overload(op, device, dtype, requires_grad, **kwargs): + ends = (-3, 0, 1, 4, 50) + starts = (-2., 0, 4.3, 50) + nsteps = (0, 1, 50) + is_start_end_tensors = ((True, True), (True, False), (False, True)) + make_arg = partial(torch.tensor, device=device, requires_grad=False) + + # Extra case to replicate off-by-one issue on CUDA + cases = list(product(starts, ends, nsteps, is_start_end_tensors)) + [(0, 7, 50, (True, True))] + for start, end, nstep, (is_start_tensor, is_end_tensor) in cases: + if dtype == torch.uint8 and (end < 0 or start < 0): + continue + + tensor_options = {"dtype": dtype, "device": device} + if is_start_tensor: + start = make_arg(start, dtype=torch.float32 if isinstance(start, float) else torch.int64) + if is_end_tensor: + end = make_arg(end, dtype=torch.float32 if isinstance(end, float) else torch.int64) + + yield SampleInput(start, args=(end, nstep), kwargs=tensor_options) + + yield SampleInput(1, args=(3, 1)) + + +def sample_inputs_logspace(op, device, dtype, requires_grad, **kwargs): + ends = (-3, 0, 1.2, 2, 4) + starts = (-2., 0, 1, 2, 4.3) + nsteps = (0, 1, 2, 4) + bases = (2., 1.1) if dtype in (torch.int8, torch.uint8) else (None, 2., 3., 1.1, 5.) + for start, end, nstep, base in product(starts, ends, nsteps, bases): + if dtype == torch.uint8 and end < 0 or start < 0: + continue + if nstep == 1 and isinstance(start, float) and not (dtype.is_complex or dtype.is_floating_point): + # https://github.com/pytorch/pytorch/issues/82242 + continue + if base is None: + yield SampleInput(start, args=(end, nstep), kwargs={"dtype": dtype, "device": device}) + else: + yield SampleInput(start, args=(end, nstep, base), kwargs={"dtype": dtype, "device": device}) + + yield SampleInput(1, args=(3, 1, 2.)) + + +def sample_inputs_logspace_tensor_overload(op, device, dtype, requires_grad, **kwargs): + ends = (-3, 0, 1.2, 2, 4) + starts = (-2., 0, 1, 2, 4.3) + nsteps = (0, 1, 2, 4) + bases = (2., 1.1) if dtype in (torch.int8, torch.uint8) else (None, 2., 3., 1.1, 5.) + is_start_end_tensors = ((True, True), (True, False), (False, True)) + make_arg = partial(torch.tensor, device=device) + for start, end, nstep, base, (is_start_tensor, is_end_tensor) in product(starts, ends, nsteps, bases, is_start_end_tensors): + if dtype == torch.uint8 and end < 0 or start < 0: + continue + if nstep == 1 and isinstance(start, float) and not (dtype.is_complex or dtype.is_floating_point): + # https://github.com/pytorch/pytorch/issues/82242 + continue + + tensor_options = {"dtype": dtype, "device": device} + + if (is_start_tensor): + start = make_arg(start, dtype=torch.float32 if isinstance(start, float) else torch.int64) + if (is_end_tensor): + end = make_arg(end, dtype=torch.float32 if isinstance(end, float) else torch.int64) + + if base is None: + yield SampleInput(start, args=(end, nstep), kwargs=tensor_options) + else: + yield SampleInput(start, args=(end, nstep, base), kwargs=tensor_options) + + yield SampleInput(1, args=(3, 1, 2.)) + + +def sample_inputs_isclose(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_elementwise_binary(op, device, dtype, requires_grad, **kwargs) + + # Creates additional inputs to test the rtol, atol, and equal_nan params + rtols = [0., 1e-7] + atols = [0., 1e-7] + equal_nans = [False, True] + + products = product(rtols, atols, equal_nans) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + for rtol, atol, equal_nan in products: + lhs = make_arg((S, S), **op.lhs_make_tensor_kwargs) + rhs = make_arg((S, S), **op.rhs_make_tensor_kwargs) + + yield SampleInput(lhs, args=(rhs,), + kwargs=dict(rtol=rtol, atol=atol, equal_nan=equal_nan)) + + +def error_inputs_isclose(op, device, **kwargs): + make_float_arg = partial(make_tensor, device=device, dtype=torch.float, requires_grad=False) + + yield ErrorInput(SampleInput(make_float_arg(()), args=(make_float_arg(()),), kwargs={'rtol': -0.4}), + error_type=RuntimeError, + error_regex='rtol must be greater than or equal to zero') + + yield ErrorInput(SampleInput(make_float_arg(()), args=(make_float_arg(()),), kwargs={'atol': -0.4}), + error_type=RuntimeError, + error_regex='atol must be greater than or equal to zero') + + +def sample_inputs_t(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg((1, 2))) + yield SampleInput(make_arg((2,))) + yield SampleInput(make_arg(())) + + +def sample_inputs_mm(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_arg_conj(size): + return make_arg(size).conj().requires_grad_(requires_grad) + + first_shape, second_shape = (S, M), (M, S) + + yield SampleInput(make_arg(first_shape), args=(make_arg(second_shape),)) + + if dtype.is_complex: + yield SampleInput(make_arg(first_shape), args=(make_arg_conj(second_shape),)) + + # Matmul of empty matrices + yield SampleInput(make_arg((0, S)), args=(make_arg(S, M),)) + yield SampleInput(make_arg((S, 0)), args=(make_arg(0, M),)) + + +def sample_inputs_addmm(op_info, device, dtype, requires_grad, **kwargs): + alpha_val = kwargs.get('alpha', 2 + 3j if dtype.is_complex else 0.6 if dtype.is_floating_point else 2) + beta_val = kwargs.get('beta', 1 + 2j if dtype.is_complex else 0.2 if dtype.is_floating_point else 3) + tests_list = [ + ((2, 3), (2, 2), (2, 3), False), + ((3, 3), (3, 3), (3, 3), False), + ] + tests_with_lhs_broadcasting = [ + ((1,), (2, 2), (2, 3), True), + ((), (2, 2), (2, 3), True), + ] + test_cases = tests_list + tests_with_lhs_broadcasting # type: ignore[operator] + + kwargs = dict(alpha=alpha_val, beta=beta_val) + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + for shape_a, shape_b, shape_c, broadcasts_input in test_cases: + yield SampleInput( + make_arg(shape_a), + make_arg(shape_b), + make_arg(shape_c), + **kwargs, + ).with_metadata(broadcasts_input=broadcasts_input) + + if dtype.is_complex: + shape = (3, 3) + yield SampleInput( + make_arg(shape), + make_arg(shape, requires_grad=False).mH.requires_grad_(requires_grad), + make_arg(shape), + **kwargs, + ) + yield SampleInput( + make_arg(shape), + make_arg(shape), + make_arg(shape, requires_grad=False).mH.requires_grad_(requires_grad), + **kwargs, + ) + # addmm of empty matrices + if dtype.is_floating_point: + yield SampleInput(make_arg(S, M), make_arg(S, 0), make_arg(0, M), **kwargs) + # empty matmul with broadcastable input + yield SampleInput(make_arg(M), make_arg(S, 0), make_arg(0, M), **kwargs).with_metadata(broadcasts_input=True) + +def sample_inputs_sparse_sampled_addmm(op_info, device, dtype, requires_grad, **kwargs): + alpha = 2 + 3j if dtype.is_complex else 0.6 + beta = 1 + 2j if dtype.is_complex else 0.2 + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # sparse.sampled_addmm performs: alpha * (A @ B) * sparse_ones_like(C) + beta * C + for m, n, k in itertools.product([0, 5], repeat=3): + yield SampleInput( + torch.eye(m, n, device=device, dtype=dtype) + .to_sparse_csr() + .requires_grad_(requires_grad), + make_arg((m, k)), + make_arg((k, n)), + alpha=alpha, + beta=beta, + ) + +def sample_inputs_sparse_mm_reduce(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + reductions = ["sum", "mean", "amax", "amin"] + for m, k, reduce in product([5, 7], [3, 11], reductions): + yield SampleInput( + torch.eye(m, m) + .to(device=device, dtype=dtype) + .to_sparse_csr() + .requires_grad_(requires_grad), + make_arg((m, k)), + reduce, + ) + + +def sample_inputs_mv(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, low=None, high=None, requires_grad=requires_grad) + yield SampleInput(make_arg(S, M), make_arg(M)) + +def sample_inputs_bmm(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, low=None, high=None, requires_grad=requires_grad) + yield SampleInput(make_arg(M, S, M), make_arg(M, M, S)) + +def sample_inputs_dot_vdot(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_arg_conj(size): + return make_arg(size).conj().requires_grad_(requires_grad) + + yield SampleInput(make_arg((S, )), make_arg((S, ))) + if dtype.is_complex: + # dot/vdot for (conj(input), conj(arg_tensor)) and (conj(input), arg_tensor) + # is tested in test_conj_view (which tests operations with only conjugated input tensor + # -- not conjugated arg tensors) + yield SampleInput(make_arg((S, )), make_arg_conj((S, ))) + + +def error_inputs_dot_vdot(op_info, device, is_ref=False, **kwargs): + make_input = partial(make_tensor, device=device, dtype=torch.float32) + + yield ErrorInput(SampleInput(make_input(1), args=(make_input(3, dtype=torch.float16),)), + error_regex='dot : expected both vectors to have same dtype') + yield ErrorInput(SampleInput(make_input(1, 1), args=(make_input(3),)), + error_regex='1D tensors expected') + yield ErrorInput(SampleInput(make_input(9), args=(make_input(3),)), + error_regex='inconsistent tensor size') + if device != "cpu" and not is_ref: + yield ErrorInput(SampleInput(make_input(3), args=(make_input(3, device="cpu"),)), + error_regex='Expected all tensors to be on the same device') + + +def sample_inputs_addmv(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + test_cases = (((S,), (S, M), (M,), 1, 1, False), + ((S,), (S, M), (M,), 0.2, 0.6, False), + ) + + test_cases_with_broadcast = (((1,), (S, M), (M,), 1, 1, True), + ((1,), (S, M), (M,), 0.2, 0.6, True), + ((), (S, M), (M,), 1, 1, True), + ((), (S, M), (M,), 0.2, 0.6, True), + ) + + cases = test_cases + test_cases_with_broadcast + + # addmv performs: beta * M + alpha * (mat @ vec) + for size, mat, vec, beta, alpha, broadcasts_input in cases: + yield SampleInput(make_arg(size), args=(make_arg(mat), make_arg(vec)), + kwargs=dict(beta=beta, alpha=alpha), broadcasts_input=broadcasts_input) + +def sample_inputs_addbmm(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # input_shape, batch1_shape, batch2_shape, beta_val, alpha_val, is_broadcasting + test_cases = [((S, M), (S, S, S), (S, S, M), 1, 1, False), + ((1,), (S, S, S), (S, S, M), 1, 1, True), + ((S, M), (S, S, S), (S, S, M), 0.6, 0.2, False), + ((1,), (S, S, S), (S, S, M), 0.6, 0.2, True), + ((), (S, S, S), (S, S, M), 1, 1, True), + ((), (S, S, S), (S, S, M), 0.6, 0.2, True), + ] + + for input_shape, batch1_shape, batch2_shape, beta, alpha, is_broadcasting in test_cases: + if dtype.is_complex: + beta_complex, alpha_complex = beta * (1 + 2j), alpha * (2 + 3j) + yield SampleInput(make_arg(input_shape), args=(make_arg(batch1_shape), make_arg(batch2_shape)), + kwargs=dict(beta=beta_complex, alpha=alpha_complex), broadcasts_input=is_broadcasting) + yield SampleInput(make_arg(input_shape), args=(make_arg(batch1_shape), make_arg(batch2_shape)), + kwargs=dict(beta=beta, alpha=alpha), broadcasts_input=is_broadcasting) + +def sample_inputs_addcmul_addcdiv(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + test_cases = [(((S, S), (S, S), (S, S)), False), + (((S, S), (S, 1), (1, S)), False), + (((1,), (S, S, 1), (1, S)), True), + (((), (), ()), False), + (((S, S), (), ()), True), + (((), (S, S, 1), (1, S)), True) + ] + + for input_args, broadcasts_input in test_cases: + # addcdiv should accept inputs with zero value + # Currently, it throws ZeroDivisionError when the denominator is zero + # TODO: exclude_zeros can be removed after https://github.com/pytorch/pytorch/issues/73638 is fixed + args = tuple(make_arg(arg, exclude_zero=True) if isinstance(arg, tuple) else arg + for arg in input_args) + yield SampleInput(*args).with_metadata(broadcasts_input=broadcasts_input) + + # addcdiv should accept inputs with zero value + # Currently, it throws ZeroDivisionError when the denominator is zero + # TODO: exclude_zeros can be removed after https://github.com/pytorch/pytorch/issues/73638 is fixed + args = tuple(make_arg(arg, exclude_zero=True) if isinstance(arg, tuple) else arg + for arg in input_args) + yield SampleInput( + *args, value=3.14 if dtype.is_floating_point or dtype.is_complex else 3 + ).with_metadata(broadcasts_input=broadcasts_input) + +def reference_inputs_addcmul_addcdiv(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_addcmul_addcdiv( + op_info, device, dtype, requires_grad, **kwargs) + + # type promotion cases + supported_dtypes = op_info.supported_dtypes(device) + make_arg = partial(make_tensor, device=device, requires_grad=requires_grad) + + types = ( + (torch.float64, torch.complex128), + (torch.bfloat16, torch.float32), + ) + + values = ( + None, + True, False, + 3.14, 3, + 1.0, 1, + 0.0, 0, + -3.14, -3, + 3.14 + 2.71j, + ) + + for (type2, type3), value in product(types, values): + if (type2 not in supported_dtypes or + type3 not in supported_dtypes): + continue + + # RuntimeError: value cannot be converted without overflow + if (type(value) is complex and + type2 is not torch.complex128): + continue + + arg1 = make_arg([5, 5], dtype=dtype) + arg2 = make_arg([5, 5], dtype=type2) + arg3 = make_arg([1, 5], dtype=type3) + + # TypeError: addcdiv(): argument 'value' must be Number, not NoneType + if value is not None: + yield SampleInput(arg1, args=(arg2, arg3), kwargs=dict(value=value)) + else: + yield SampleInput(arg1, args=(arg2, arg3)) + +def sample_inputs_baddbmm(op_info, device, dtype, requires_grad, **kwargs): + test_cases = [((S, S, M), (S, S, S), (S, S, M), 1, 1, False), + ((1,), (S, S, S), (S, S, M), 1, 1, True), + ((S, S, M), (S, S, S), (S, S, M), 0.6, 0.2, False), + ((1,), (S, S, S), (S, S, M), 0.6, 0.2, True), + ((), (S, S, S), (S, S, M), 1, 1, True), + ((), (S, S, S), (S, S, M), 0.6, 0.2, True), + ] + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, low=None, high=None) + for (input_shape, batch1_shape, batch2_shape, alpha, beta, broadcasts_input) in test_cases: + yield SampleInput( + make_arg(input_shape), + make_arg(batch1_shape), + make_arg(batch2_shape), + beta=beta, + alpha=alpha + ).with_metadata(broadcasts_input=broadcasts_input) + + if dtype.is_complex: + yield SampleInput( + make_arg(input_shape), + make_arg(batch1_shape), + make_arg(batch2_shape), + beta=beta * (1 + 2j), + alpha=alpha * (2 + 3j), + ).with_metadata(broadcasts_input=broadcasts_input) + + if dtype.is_complex: + shapes = [(S, S, S), (S, M, S), (S, S, M)] + args = tuple(make_arg(s) for s in shapes) + yield SampleInput( + args[0].transpose_(-1, 1), + args[1].transpose(-1, 1).conj().requires_grad_(requires_grad), + args[2].transpose(-1, 1).conj().requires_grad_(requires_grad), + beta=beta * (1 + 2j), + alpha=alpha * (2 + 3j), + ) + +# TODO: add reduction kwargs +def sample_inputs_multilabel_soft_margin_loss(op_info, device, dtype, requires_grad, **kwargs): + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + shapes = ( + (S,), + (S, S), + ) + + for shape in shapes: + # Produce one with weight and one without. + yield SampleInput(_make_tensor(shape), args=(_make_tensor(shape, requires_grad=False),), kwargs={}) + yield SampleInput(_make_tensor(shape), args=(_make_tensor(shape, requires_grad=False),), + kwargs={'weight': _make_tensor(shape, requires_grad=False)}) + +def sample_inputs_addr(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, low=None, high=None + ) + yield SampleInput(make_arg(S, M), make_arg(S), make_arg(M)) + + yield SampleInput(make_arg(), make_arg(S), make_arg(M)).with_metadata(broadcasts_input=True) + + if dtype.is_complex: + alpha, beta = 0.1 + 0.3j, 0.4 + 0.6j + elif dtype.is_floating_point: + alpha, beta = 0.2, 0.6 + else: + alpha, beta = 2, 3 + + yield SampleInput(make_arg(S, M), make_arg(S), make_arg(M), beta=beta, alpha=alpha) + + yield SampleInput( + make_arg(), + make_arg(S), + make_arg(M), + beta=beta, + alpha=alpha, + ).with_metadata(broadcasts_input=True) + + # These samples fail gradcheck + if dtype.is_floating_point and not requires_grad: + tensor_options = dict(device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput( + torch.tensor([[math.nan]], **tensor_options), + torch.tensor([0.0], **tensor_options), + torch.tensor([0.0], **tensor_options), + beta=0.0, + alpha=0.0, + ).with_metadata(broadcasts_input=True) + + yield SampleInput( + torch.tensor([[0.0]], **tensor_options), + torch.tensor([math.nan], **tensor_options), + torch.tensor([math.nan], **tensor_options), + beta=0.0, + alpha=0.0, + ).with_metadata(broadcasts_input=True) + +def sample_inputs_zero_(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases = ((), (S, S, S), (S,)) + + for shape in cases: + yield SampleInput(make_arg(shape)) + +def sample_inputs_multi_margin_loss(op_info, device, dtype, requires_grad, **kwargs): + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(_make_tensor, dtype=torch.long, requires_grad=False) + make_weight = partial(_make_tensor, requires_grad=False) + + inputs = ( + ((), make_target([], low=0, high=1), {}), + ((S,), make_target([], low=0, high=S), {"p": 1}), + ((S,), make_target([1], low=0, high=S), {"p": 2}), + ((S, M), make_target([S], low=0, high=M), {"margin": 1.0}), + ((S, M), make_target([S], low=0, high=M), {"margin": -3.14}), + ((M, S), make_target([M], low=0, high=S), {"weight": None}), + ((M, S), make_target([M], low=0, high=S), {"weight": make_weight([S], low=-10., high=10.)}), + ((M, S), make_target([M], low=0, high=S), {"reduction": "none"}), + ((M, S), make_target([M], low=0, high=S), {"reduction": "mean"}), + ((M, S), make_target([M], low=0, high=S), {"reduction": "sum"}), + ) + + for input_shape, target, kwargs in inputs: + yield SampleInput(_make_tensor(input_shape), args=(target,), kwargs=kwargs) + + +def reference_inputs_multi_margin_loss(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_multi_margin_loss(op_info, device, dtype, requires_grad, **kwargs) + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(_make_tensor, dtype=torch.long, requires_grad=False) + make_weight = partial(_make_tensor, requires_grad=False) + + inputs = ( + ((), make_target([], low=0, high=1)), + ((S,), make_target([], low=0, high=S)), + ((S,), make_target([1], low=0, high=S)), + ((M, S), make_target([M], low=0, high=S)), + ) + ps = (1, 2) + margins = (0, 7, -3.14) + weights = (False, True) + reductions = (None, "none", "mean", "sum") + + for (input_shape, target), p, margin, weight, reduction in product(inputs, ps, margins, weights, reductions): + input = _make_tensor(input_shape) + weight_shape = [input.size(-1)] if input.ndim > 0 else [1] + weight = make_weight(weight_shape, low=-10., high=10.) if weight else None + kwargs = {"p": p, "margin": margin, "weight": weight} + if reduction is not None: + kwargs["reduction"] = reduction + yield SampleInput(input, args=(target,), kwargs=kwargs) + + +def error_inputs_multi_margin_loss(op, device, **kwargs): + make_input = partial(make_tensor, device=device, dtype=torch.float32) + # invalid reduction + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5,),), kwargs={'reduction': 'abc'}), + error_type=ValueError, error_regex='abc is not a valid value for reduction') + # invalid input + yield ErrorInput(SampleInput(make_input(5, 0), args=(make_input(5,),), kwargs={}), + error_type=RuntimeError, + error_regex=r'Expected non-empty vector or matrix with optional 0-dim batch size, but got: \[5, 0\]') + yield ErrorInput(SampleInput(make_input(0,), args=(make_input(5,),), kwargs={}), + error_type=RuntimeError, + error_regex=r'Expected non-empty vector or matrix with optional 0-dim batch size, but got: \[0\]') + # invalid target + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5, 4),), kwargs={}), + error_type=RuntimeError, + error_regex=r'target tensor should be 1-D with size equal to.*Expected target size \[5\].*but got \[5, 4\]') + # invalid target dtype + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5,),), kwargs={}), + error_type=RuntimeError, error_regex='expected scalar type Long but found Float') + # invalid weight + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5,),), kwargs={'weight': make_input(())}), + error_type=ValueError, error_regex='weight must be one-dimensional') + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5,),), kwargs={'weight': make_input(5, 4)}), + error_type=ValueError, error_regex='weight must be one-dimensional') + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5,),), kwargs={'weight': make_input(5,)}), + error_type=RuntimeError, error_regex=r'inconsistent weight size, expected 4 but got \[5\]') + # invalid p + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5,),), kwargs={'p': 3}), + error_type=ValueError, error_regex='only p == 1 and p == 2 supported') + + +def sample_inputs_logsumexp(self, device, dtype, requires_grad, **kwargs): + inputs = ( + ((), (0,), True), + ((S, S), (1,), True), + ((S, S), (1,), False), + ((S, S), (-2,), False), + ((S, S), (0, 1), False), + ) + # Test large inputs to check numerical stability + lows = (None, 1e3, 1e6) if dtype in (torch.float32, torch.float64, torch.complex64, torch.complex128) else (None,) + for low in lows: + high = low * 2 if low is not None else None + for shape, dim, keepdim in inputs: + t = make_tensor(shape, dtype=dtype, device=device, + low=low, high=high, + requires_grad=requires_grad) + yield SampleInput(t, dim, keepdim) + +def reference_inputs_logsumexp(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_logsumexp(op, device, dtype, requires_grad, **kwargs) + + # https://github.com/pytorch/pytorch/issues/91843 + t = torch.tensor([20, 30, 100], dtype=dtype, device=device, requires_grad=requires_grad) + yield SampleInput(t, 0, False) + + t = torch.tensor((), dtype=dtype, device=device, requires_grad=requires_grad) + yield SampleInput(t, 0, False) + + # tests masking + # https://github.com/pytorch/pytorch/pull/91860#pullrequestreview-1241344073 + t = torch.tensor(float("inf")) + yield SampleInput(t, 0, True) + +def sample_inputs_like_fns(self, device, dtype, requires_grad, **kwargs): + inputs = [ + ((), {}), + ((S, S), {}), + ((0, S, 0), {}), + ((S,), {'dtype': dtype, 'device': device}), + # Hard-code some dtypes/devices. We want to test cases where the + # (dtype, device) is different from the input's (dtype, device) + ((S,), {'dtype': highest_precision_float(device)}), + ((S,), {'device': 'cpu'}), + ((S,), {'dtype': torch.double, 'device': 'cpu'}), + ] + if torch.cuda.is_available(): + inputs.append(((S,), {'device': 'cuda'})) + + for shape, kwargs in inputs: + t = make_tensor(shape, dtype=dtype, device=device, + low=None, high=None, + requires_grad=requires_grad) + yield SampleInput(t, **kwargs) + +def reference_inputs_like_fns(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_like_fns(op, device, dtype, requires_grad, **kwargs) + + # shape + cases = ( + (), (0,), (1, 0), (1, 1, 4, 5), (5, 3, 0, 1), (1, 4, 3, 1, 1) + ) + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + for shape in cases: + yield SampleInput(make_arg(shape)) + yield SampleInput(make_arg(shape).transpose(0, -1)) + yield SampleInput(make_arg(shape, noncontiguous=True)) + yield SampleInput(make_arg(shape, noncontiguous=True).transpose(0, -1)) + +def sample_inputs_multilabel_margin_loss(op_info, device, dtype, requires_grad, **kwargs): + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(_make_tensor, dtype=torch.long, requires_grad=False) + + inputs = ( + ([], make_target([], low=0, high=1), {}), + ([S], make_target([S], low=0, high=S), {}), + ([M, S], make_target([M, S], low=0, high=S), {}), + ([M, S], make_target([M, S], low=0, high=S), {"reduction": "none"}), + ([M, S], make_target([M, S], low=0, high=S), {"reduction": "mean"}), + ([M, S], make_target([M, S], low=0, high=S), {"reduction": "sum"}), + ) + + for shape, target, kwargs in inputs: + yield SampleInput(_make_tensor(shape), args=(target,), kwargs=kwargs) + + +def reference_inputs_multilabel_margin_loss(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_multilabel_margin_loss(op_info, device, dtype, requires_grad, **kwargs) + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(_make_tensor, dtype=torch.long, requires_grad=False) + make_target_tensor = partial(torch.tensor, device=device, dtype=torch.long, requires_grad=False) + + inputs = ( + # random tests including -1 target labels + ([], make_target([], low=-1, high=1)), + ([S], make_target([S], low=-1, high=S)), + ([M, S], make_target([M, S], low=-1, high=S)), + # repeated target labels and -1 (labels after the first -1 are ignored) + ([], make_target_tensor(-1)), + ([7], make_target_tensor([2, 0, 6, -1, 4, -1, 6])), + ([4, 5], make_target_tensor([[4, -1, 0, -1, 2], [0, 0, 4, 1, 4], [-1, 3, -1, 1, 0], [4, 3, 2, 1, 0]])), + ) + reductions = (None, "none", "mean", "sum") + + for (shape, target), reduction in product(inputs, reductions): + kwargs = {} + if reduction is not None: + kwargs["reduction"] = reduction + yield SampleInput(_make_tensor(shape), args=(target,), kwargs=kwargs) + + +def error_inputs_multilabel_margin_loss(op, device, **kwargs): + make_input = partial(make_tensor, device=device, dtype=torch.float32) + # invalid reduction + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5, 4),), kwargs={'reduction': 'abc'}), + error_type=ValueError, error_regex='abc is not a valid value for reduction') + # invalid input + yield ErrorInput(SampleInput(make_input(5, 0), args=(make_input(5, 4),), kwargs={}), + error_type=RuntimeError, + error_regex=r'Expected non-empty vector or matrix with optional 0-dim batch size, but got: \[5, 0\]') + yield ErrorInput(SampleInput(make_input(0,), args=(make_input(0,),), kwargs={}), + error_type=RuntimeError, + error_regex=r'Expected non-empty vector or matrix with optional 0-dim batch size, but got: \[0\]') + # invalid target + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(4,),), kwargs={}), + error_type=RuntimeError, + error_regex=r'inconsistent target size: \[4\] for input of size: \[5, 4\]') + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input((),),), kwargs={}), + error_type=RuntimeError, + error_regex=r'inconsistent target size: \[\] for input of size: \[5, 4\]') + + +def get_independent_tensor(tensor): + return tensor.clone().requires_grad_(tensor.requires_grad) + +def sample_inputs_randint(self, device, dtype, requires_grad, **kwargs): + low = 2 + high = 10 + + for sample in sample_inputs_like_fns(self, device, dtype, requires_grad, **kwargs): + sample.kwargs.setdefault('device', device) + # With high + yield SampleInput(high, sample.input.shape, *sample.args, **sample.kwargs) + # With low and high + yield SampleInput(low, high, sample.input.shape, *sample.args, **sample.kwargs) + +def sample_inputs_randint_like(self, device, dtype, requires_grad, **kwargs): + low = 2 + high = 10 + + for sample in sample_inputs_like_fns(self, device, dtype, requires_grad, **kwargs): + # With high + yield SampleInput( + sample.input, + high, + *sample.args, + **sample.kwargs) + # With low and high + yield SampleInput( + get_independent_tensor(sample.input), + low, + high, + *sample.args, + **sample.kwargs) + +def sample_inputs_margin_ranking_loss(op_info, device, dtype, requires_grad, **kwargs): + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + shapes = ( + (), + (S,), + (S, S), + (S, S, S), + ) + + margins = (0., 1.) + reductions = ('sum', 'mean', 'none') + + for shape in shapes: + for margin, reduction in product(margins, reductions): + kwargs = {'margin': margin, 'reduction': reduction} + yield SampleInput(_make_tensor(shape), + args=(_make_tensor(shape, requires_grad=False), + _make_tensor(shape, requires_grad=False)), + kwargs=kwargs) + +def reference_inputs_margin_ranking_loss(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_margin_ranking_loss(op, device, dtype, requires_grad, **kwargs) + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + for reduction in ('sum', 'mean', 'none'): + if dtype.is_floating_point: # only supports ints and floats + # NaN propagation + inp1 = make_input((10, )) + inp1[2] = float('nan') + inp2 = make_input((10, )) + inp2[4] = float('nan') + target = make_input((10, )) + inp2[9] = float('nan') + yield SampleInput(inp1, args=(inp2, target), kwargs={'reduction': reduction}) + + # Inf handling + inp1 = make_input((10, )) + inp2[1] = float('inf') + inp2 = make_input((10, )) + inp2[4] = float('inf') + target = make_input((10, )) + inp2[7] = float('inf') + yield SampleInput(inp1, args=(inp2, target), kwargs={'reduction': reduction}) + + # Broadcasting + inp1 = make_input((5, 2)) + inp2 = make_input((5, 1)) + target = make_input((1, 2)) + yield SampleInput(inp1, args=(inp2, target), kwargs={'reduction': reduction}) + +def error_inputs_margin_ranking_loss(op, device, **kwargs): + make_input = partial(make_tensor, device=device, dtype=torch.float32) + # invalid reduction value. + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5, 4), make_input(5, 4),), kwargs={'reduction': 'abc'}), + error_type=ValueError, error_regex='is not a valid value') + # invalid input shapes + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5, 4), make_input(5,),)), + error_regex='margin_ranking_loss : All input tensors should') + +def sample_inputs_new_fns(self, device, dtype, requires_grad, *, is_strided=False, **kwargs): + # input_shape, output_shape, strides, kwargs + # lengths of output_shape and strides must be equal + inputs = [ + ((), (), (), {}), + ((S, S), (2, 0), (3, 4), {}), + ((0, S, 0), (3, 2, 2), (1, 2, 3), {}), + ((S,), (2, 3), (7, 8), {'dtype': dtype, 'device': device}), + # Hard-code some dtypes/devices. We want to test cases where the + # (dtype, device) is different from the input's (dtype, device) + ((S,), (10,), (S,), {'dtype': highest_precision_float(device)}), + ((S,), (1, 1, 12), (S, L, M), {'device': 'cpu'}), + ((S,), (2, 2, 2), (L, M, S), {'dtype': torch.double, 'device': 'cpu'}), + ] + if torch.cuda.is_available(): + inputs.append(((S,), (7, 2), (3, 4), {'device': 'cuda'})) + + for input_shape, output_shape, strides, kwargs in inputs: + t = make_tensor(input_shape, dtype=dtype, device=device, + low=None, high=None, + requires_grad=requires_grad) + if is_strided: + yield SampleInput(t, output_shape, strides, **kwargs) + else: + yield SampleInput(t, output_shape, **kwargs) + +def sample_inputs_empty_strided(op, device, dtype, requires_grad=False, **kwargs): + + inputs = [ + ((), (), {'dtype': dtype, 'device': device}), + ((S,), (4,), {'dtype': dtype, 'device': device}), + ((S, S), (2, 1), {'dtype': dtype, 'device': device}), + ((S, S, S), (2, 0, 1), {'dtype': dtype, 'device': device}), + ] + + for shape, strides, kwargs in inputs: + yield SampleInput(shape, strides, requires_grad=requires_grad, **kwargs) + +def sample_inputs_empty(op, device, dtype, requires_grad, **kwargs): + # shape + cases = ( + (), (0,), (1,), (1, 3, 5), (5, 3, 1), (1, 0, 5, 1), + ) + + for case in cases: + yield SampleInput(case, device=device, dtype=dtype, requires_grad=requires_grad) + +def sample_inputs_empty_permuted(op, device, dtype, requires_grad, **kwargs): + # shape + cases = ( + (), (0,), (1,), (1, 3, 5), (5, 3, 1), (1, 0, 5, 1), + ) + + for case in cases: + for layout in itertools.permutations(range(len(case))): + yield SampleInput(case, layout, device=device, dtype=dtype, requires_grad=requires_grad) + +def error_inputs_empty_permuted(op_info, device, **kwargs): + yield ErrorInput( + SampleInput((2,), args=((0, 1),)), + error_type=RuntimeError, + error_regex="Number of dimensions in size does not match the length of the physical_layout" + ) + yield ErrorInput( + SampleInput((2,), args=((3,),)), + error_type=RuntimeError, + error_regex="Dimension out of range" + ) + yield ErrorInput( + SampleInput((2, 3), args=((0, 0),)), + error_type=RuntimeError, + error_regex="Duplicate dim not allowed" + ) + +def sample_inputs_scalar_tensor(op, device, dtype, requires_grad, **kwargs): + # Not including a scalar tensor in vals because meta tests start failing due to + # lack of meta support for _local_scalar_dense + # torch.tensor(2, device=device) + vals = (-5, 0, 1) + + for item in vals: + yield SampleInput(item, device=device, dtype=dtype, requires_grad=requires_grad) + +def sample_inputs_eye(op, device, dtype, requires_grad, **kwargs): + # only ints >= 0 are allowed for both arguments, unless m is omitted + sizes = (None, 0, 1, 2, 3, 4, 7, L, M, S) + + for n, m in product(sizes, sizes): + if n is None: + continue + + # TODO: no layout + _kwargs = {'device': device, 'dtype': dtype, 'requires_grad': requires_grad} + if m is None: + yield SampleInput(n, args=(), kwargs=_kwargs) + else: + yield SampleInput(n, args=(m,), kwargs=_kwargs) + +def error_inputs_eye(op_info, device, **kwargs): + # TODO: no layout + _kwargs = {'device': device, 'dtype': torch.float32} + + yield ErrorInput( + SampleInput(-1, args=(), kwargs=_kwargs), + error_regex="n must be greater or equal to 0, got -1" + ) + + yield ErrorInput( + SampleInput(-7, args=(42,), kwargs=_kwargs), + error_regex="n must be greater or equal to 0, got -7" + ) + + yield ErrorInput( + SampleInput(0, args=(-3,), kwargs=_kwargs), + error_regex="m must be greater or equal to 0, got -3" + ) + + +def sample_inputs_new_full(self, device, dtype, requires_grad, **kwargs): + def get_val(dtype): + return make_tensor([], dtype=dtype, device="cpu").item() + + for sample in sample_inputs_new_fns(self, device, dtype, requires_grad, **kwargs): + # The scalar we are passing to new_full must be the same dtype + # as the one of the resulting tensor + use_dtype = sample.kwargs.get('dtype', dtype) + yield SampleInput( + sample.input, *sample.args, get_val(use_dtype), **sample.kwargs) + +def sample_inputs_full_like(self, device, dtype, requires_grad, **kwargs): + def get_val(dtype): + return make_tensor([], dtype=dtype, device="cpu").item() + + double_dtype = highest_precision_float(device) + inputs = [ + ((), get_val(dtype), {}), + ((S, S), get_val(dtype), {}), + ((0, S, 0), get_val(dtype), {}), + ((S,), get_val(dtype), {'dtype': dtype, 'device': device}), + # Hard-code some dtypes/devices. We want to test cases where the + # (dtype, device) is different from the input's (dtype, device) + ((S,), get_val(double_dtype), {'dtype': double_dtype}), + ((S,), get_val(dtype), {'device': 'cpu'}), + ((S,), get_val(double_dtype), {'dtype': double_dtype, 'device': 'cpu'}), + ] + if torch.cuda.is_available(): + inputs.append(((S,), get_val(dtype), {'device': 'cuda'})) + + if torch.mps.is_available() and dtype not in [torch.float64, torch.complex128, torch.uint32, torch.uint16]: + inputs.append(((S,), get_val(dtype), {'device': 'mps'})) + + if not dtype.is_signed: + # For unsigned dtypes, negative values are converted. + inputs.append(((S,), -get_val(dtype), {})) + + for shape, fill_value, kwargs in inputs: + t = make_tensor(shape, dtype=dtype, device=device, + low=None, high=None, + requires_grad=requires_grad) + yield SampleInput(t, fill_value, **kwargs) + +def sample_inputs_multinomial(self, device, dtype, requires_grad, **kwargs): + cases = [ + ([3], 3, {}), + ([10], 3, {}), + ([3, 10], 3, {}), + ([3], 3, dict(replacement=False)), + ([3], 3, dict(replacement=True)), + ([3, 4], 4, dict(replacement=True)), + ([3, 4], 4, dict(replacement=False)), + ] + + for shape, num_samples, kwargs in cases: + t = make_tensor(shape, dtype=dtype, device=device, + low=0, high=None, + requires_grad=requires_grad) + yield SampleInput(t, num_samples, **kwargs) + +def sample_inputs_normal_common(self, device, dtype, requires_grad, cases, **kwargs): + def get_value_or_make_tensor(value_or_shape): + if isinstance(value_or_shape, list): + return make_tensor(value_or_shape, dtype=dtype, device=device, + low=0, high=None, + requires_grad=requires_grad) + return value_or_shape + + for value_or_mean_shape, value_or_std_shape, kwargs in cases: + mean = get_value_or_make_tensor(value_or_mean_shape) + std = get_value_or_make_tensor(value_or_std_shape) + yield SampleInput(mean, std, **kwargs) + +def sample_inputs_normal_tensor_first(self, device, dtype, requires_grad, **kwargs): + # value_or_size, value_or_size, kwargs + cases = [ + ([], [], {}), + ([3], [3], {}), + ([3, 4, 2], [3, 4, 2], {}), + ([2, 3], 1.1, {}), + ([1, 2, 3], [5, 2, 3], {}), # broadcasting + ] + + return sample_inputs_normal_common(self, device, dtype, requires_grad, cases, **kwargs) + +def sample_inputs_normal_tensor_second(self, device, dtype, requires_grad, **kwargs): + yield SampleInput(1.6, 0.3, [2, 3], dtype=dtype, device=device) + yield SampleInput(1.6, 0.3, [2, 2, 2], dtype=dtype, layout=torch.strided, device=device) + yield SampleInput(2.7, make_tensor([4, 3], dtype=dtype, device=device, low=0, high=None, requires_grad=requires_grad)) + +def sample_inputs_bernoulli(self, device, dtype, requires_grad, **kwargs): + shapes = [ + [3], + [], + [0, 3], + [2, 3, 4], + ] + + for shape in shapes: + t = make_tensor(shape, dtype=dtype, device=device, + low=0, high=1, + requires_grad=requires_grad) + yield SampleInput(t) + +def error_inputs_bernoulli(op_info, device, **kwargs): + # more than one element of the written-to tensor refers to a single memory location + x = torch.rand((1,), device=device).expand((6,)) + err_msg = 'unsupported operation' + yield ErrorInput(SampleInput(torch.rand_like(x), kwargs={'out': x}), + error_regex=err_msg) + +def sample_inputs_logcumsumexp(self, device, dtype, requires_grad, **kwargs): + inputs = ( + ((S, S, S), 0), + ((S, S, S), 1), + ((), 0), + ) + + for large_number in (True, False): + for shape, dim in inputs: + t = make_tensor(shape, dtype=dtype, device=device, + low=None, high=None, + requires_grad=requires_grad) + + if large_number and t.dim() > 0: + t[0] = 10000 + yield SampleInput(t, dim) + +def sample_inputs_trace(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, + requires_grad=requires_grad, low=None, high=None) + # Square, tall (rows > cols), wide (rows < cols), single row/col (#171704) + for shape in ((S, S), (S + 2, S), (S, S + 2), (1, S), (S, 1)): + yield SampleInput(make_arg(shape)) + + +def error_inputs_trace(op, device): + yield ErrorInput(SampleInput(make_tensor((3, 4, 5), dtype=torch.float32, device=device)), error_regex="expected a matrix") + + +def sample_inputs_renorm(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + cases = (((S, S, S), (2, 1, 0.5)), + ((S, S, S), (2, -1, 0.5)), + ((S, S, S), (1, 2, 3)), + ((S, S, S), (float('inf'), 2, 0.5)), + ) + + for shape, args in cases: + yield SampleInput(make_arg(shape), args=args) + + +def sample_inputs_transpose_swapdims(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + cases = (((1, 2, 3), (-1, -2)), + ((1, 2, 3), (-1, 2)), + ((1, 2, 3), (1, -2)), + ((1, 2, 3), (1, 2)), + ((), (0, 0)), + ((1, ), (0, 0)), + ((M, M), (0, 1)), + ((S, S, S), (2, 0)), ) + + for shape, args in cases: + yield SampleInput(make_arg(shape), args=args) + +def _numpy_ref_transpose(a, dim0, dim1): + if a.ndim <= 1: + return a + + return np.swapaxes(a, dim0, dim1) + +def sample_inputs_adjoint(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + shapes = ((1, 2, 3), (M, M), (S, S, S), (S, M, S), (M, S, M, S)) + return (SampleInput(make_arg(shape)) for shape in shapes) + +def sample_inputs_T(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + shapes = ((M, M), (M, L)) + return (SampleInput(make_arg(shape)) for shape in shapes) + +def error_inputs_T(self, device, has_ndims_error=False): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # Deprecated behavior in regular PyTorch, but throws an error in primTorch: + # https://github.com/pytorch/pytorch/issues/86968 + if has_ndims_error: + # ndims == 1 + yield ErrorInput(SampleInput(make_arg(M)), + error_regex=(r'The use of `x\.T` on tensors of dimension other than 0 or 2 ' + r'to reverse their shape is not supported\.')) + + # ndims > 2 + yield ErrorInput(SampleInput(make_arg(M, S, L)), + error_regex=(r'The use of `x\.T` on tensors of dimension other than 0 or 2 ' + r'to reverse their shape is not supported\.')) + + +def sample_inputs_singular_matrix_factors(op_info, device, dtype, requires_grad=False): + """ + This function produces two tensors of shape (*, m, k) and (*, n, k) with k <= min(m, n). + Their matrix product could be used to generate tensor of shape (*, m, n) of rank k. + """ + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + batches = [(), (2,)] + size = [3, 4] + for batch, m, n in product(batches, size, size): + k = 2 + a = make_arg((*batch, m, k)) + b = make_arg((*batch, n, k)) + yield a, b + + +def sample_inputs_svd_lowrank(op_info, device, dtype, requires_grad=False, **kwargs): + # Function that's well defined on the outputs for complex inputs + def fn(usv): + U, S, V = usv + return U @ V.mH, S + + for (a, b) in sample_inputs_singular_matrix_factors(op_info, device, dtype, requires_grad): + *batch, m, k = a.shape + n = b.shape[-2] + + # NOTE: since svd_lowrank relies on non rank-revealing SVD, + # it inherits the problem of unstable behavior with repeated + # singular values including zeros. + # Since we want to avoid (repeated) zeros as singular values, + # we can only use k for q. + # This issues could be resolved with using a rank-revealing SVD + # which does not include "zero" singular values. + yield SampleInput(a, b, q=k, M=None).with_metadata(output_process_fn_grad=fn) + + for (a, b) in sample_inputs_singular_matrix_factors(op_info, device, dtype, requires_grad): + *batch, m, k = a.shape + n = b.shape[-2] + M = make_tensor((*batch, m, n), dtype=dtype, device=device, requires_grad=requires_grad) + yield SampleInput(a, b, q=k, M=M).with_metadata(output_process_fn_grad=fn) + +def chunk_iter(iterable, size): + it = iter(iterable) + while True: + chunk = tuple(islice(it, size)) + if not chunk: + break + yield chunk + +def sample_inputs_pca_lowrank(op_info, device, dtype, requires_grad=False, **kwargs): + # we reuse samples from svd_lowrank which come in group of two with + # kwarg['M'] = None and with kwarg['M'] = + samples = sample_inputs_svd_lowrank(op_info, device, dtype, requires_grad, **kwargs) + for s1, s2 in chunk_iter(samples, 2): + del s1.kwargs['M'] + del s2.kwargs['M'] + s1.kwargs['center'] = False + s2.kwargs['center'] = True + yield s1 + yield s2 + +def np_sinc_with_fp16_as_fp32(x): + # Wraps numpy's sinc function so that fp16 values are promoted to fp32 + # before sinc is invoked. Context: numpy's sinc returns NaN when evaluated + # at 0 for fp16. + if x.dtype == np.float16: + return np.sinc(x.astype(np.float32)) + else: + return np.sinc(x) + +def sample_inputs_broadcast_to(op_info, device, dtype, requires_grad, **kwargs): + test_cases = ( + ((S, 1, 1), (S, S, S)), + ((S, 1, S), (S, S, S)), + ((S, 1), (S, S, S)), + ((1,), (S, S, S)), + ((1, S), (1, 1, S)), + ((), ()), + ((), (1, 3, 2)), + ) + + return ( + SampleInput( + make_tensor(size, dtype=dtype, device=device, low=None, high=None, requires_grad=requires_grad), + shape, + ) for size, shape in test_cases) + +def sample_inputs_broadcast_tensors(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + test_cases: tuple[tuple] = (((3,), (1, 2, 1), (1, 1), (5, 1, 1),),) + + for shape, *other_shapes in test_cases: + yield SampleInput(make_arg(shape), args=tuple(make_arg(s) for s in other_shapes)) + +def reference_inputs_broadcast_tensors(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_broadcast_tensors(op, device, dtype, requires_grad, **kwargs) + + m = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + n = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad, noncontiguous=True) + + cases = ( + ((), (1, 1), (1, 1, 7, 1), (3, 1, 1)), + ((3, 5, 6), (1, 3, 5, 6), (1, 1, 1, 1, 6), (8, 3, 5, 6)) + ) + + for a, b, c, d in cases: + yield SampleInput(m(a), args=(m(b), m(c), m(d))) + yield SampleInput(n(a), args=(n(b), n(c), n(d))) + +def sample_inputs_block_diag(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + test_cases: tuple[tuple] = ( + ((1, S), (2, S), (3, S),), + ((S, 1), (S, 2), (S, 3),), + ((1,), (2,), (3,),), + ((2, S), (S,)) + ) + + for shape, *other_shapes in test_cases: + yield SampleInput(make_arg(shape), args=tuple(make_arg(s) for s in other_shapes)) + # We also want to test mixed complex-non-complex inputs to block_diag + if dtype == torch.complex32 or dtype == torch.complex64: + non_complex_dtype = torch.float32 if dtype == torch.complex32 or device == 'mps:0' else torch.float64 + make_arg_non_complex = partial(make_tensor, dtype=non_complex_dtype, device=device, requires_grad=requires_grad) + yield SampleInput(make_arg_non_complex(shape), args=tuple(make_arg(s) for s in other_shapes)) + +def sample_inputs_cdist(op_info, device, dtype, requires_grad, **kwargs): + small_S = 2 + test_cases = ( + ((S, S, 2), (S, S + 1, 2)), + ((S, S), (S, S)), + ((S, S, S), (S, S, S)), + ((3, 5), (3, 5)), + ((2, 3, 5), (2, 3, 5)), + ((1, 2, 3), (1, 2, 3)), + ((1, 1), (S, 1)), + ((0, 5), (4, 5)), + ((4, 5), (0, 5)), + ((0, 4, 5), (3, 5)), + ((4, 5), (0, 3, 5)), + ((0, 4, 5), (1, 3, 5)), + ((1, 4, 5), (0, 3, 5)), + # Using S here would make this one test take 9s + ((small_S, small_S, small_S + 1, 2), (small_S, small_S, small_S + 2, 2)), + ((small_S, 1, 1, small_S), (1, small_S, small_S)), + ((1, 1, small_S), (small_S, 1, small_S, small_S)), + ) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + for cm in ['use_mm_for_euclid_dist', 'donot_use_mm_for_euclid_dist']: + # FIXME add an override for JIT and revert 0. back to 0 + # since it's accepted by eager + for p in [0., 1., 2., 3., 0.5, 1.5, 2.5, float("inf")]: + for t1_size, t2_size in test_cases: + # The args should never be non-contiguous as this is not supported in the backward + yield SampleInput(make_arg(t1_size), make_arg(t2_size), p, cm) + +def _fill_np(a, value): + a = a.copy() + a.fill(value) + return a + +def _fill_sample_kwargs(device, dtype, input): + if dtype is torch.bool: + value = True + else: + value = 3 + + return ({'value': value}, {'value': value}) + +def sample_inputs_comparison_ops(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_elementwise_binary(op, device, dtype, requires_grad, **kwargs) + + # Adds a sample input where both tensors have the same values + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + lhs = make_arg((S, S)) + yield SampleInput(lhs, args=(lhs.clone(),)) + +def sample_inputs_stack(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # shape x number of tensors + cases = ( + ((3, 4), 1), + ((1, 2, 1, 4), 3), + ((0, 1, 0), 2),) + + for shape, num_tensors in cases: + tensors = [make_arg(shape) for _ in range(num_tensors)] + for dim in range(-1, len(shape) - 1): + yield SampleInput(tensors, args=(dim,)) + + +def sample_inputs_chunk_cat(op_info, device, dtype, requires_grad, **kwargs): + # 1. If input tensors have different ndims, dim should be non-negative and be less than the ndims of every input tensors. + # If all input tensors have the same ndims, we support both negative and non-negative dim. + # 2. For wrapped_dim, all tensors should have the same size for 0,...,wrapped_dim-1 dimensions. + # No requirements for (wrapped_dim, ...)-th dimension. + # 3. Expect positive num_chunks + # 4. Expect non-empty input tensor list and each input tensor should have at least 1 element + # 5. Non-contiguous input tensors are allowed. + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + same_ndim_cases = ( + ( + [ + torch.Size([1, 2, 3]), + torch.Size([1, 2, 3]), + ], -1, 5 + ), + ( + [ + torch.Size([1, 2, 129]), + torch.Size([1, 2, 297]), + ], -1, 5 + ), + ( + [ + torch.Size([1, 2, 3]), + torch.Size([1, 2, 3]), + ], 1, 5 + ), + ( + [ + torch.Size([3, 3, 2, 1]), + torch.Size([1, 4, 2, 2]), + torch.Size([2, 1, 3, 3]), + ], 0, 2 + ), + ) + for sizes, dim, num_chunks in same_ndim_cases: + tensors = [make_arg(size) for size in sizes] + yield SampleInput(tensors, args=(dim, num_chunks)) + + different_ndim_case = [ + torch.Size([2, 3, 3]), + torch.Size([2, 3, 1, 2]), + torch.Size([2, 3]), + torch.Size([2, 3, 2]), + torch.Size([2, 3, 271]), + ] + max_dim, num_chunks = 2, 3 + for dim in range(max_dim): + tensors = [] + for size in different_ndim_case: + tensors.append(make_arg(size)) + yield SampleInput(tensors, args=(dim, num_chunks)) + + # non-contiguous + for dim in range(max_dim): + tensors = [] + for size in different_ndim_case: + # make the last 2 dims column-major (i.e. non-contiguous) + t = make_arg(size).transpose(-2, -1).contiguous().transpose(-2, -1) + tensors.append(t) + yield SampleInput(tensors, args=(dim, num_chunks)) + +def error_inputs_chunk_cat(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # input tensors have different ndims but dim is negative + sizes, dim, num_chunks = [torch.Size([2, 3]), torch.Size([4,])], -1, 3 + tensors = [make_arg(size) for size in sizes] + yield ErrorInput( + SampleInput(tensors, args=(dim, num_chunks)), + error_regex='_chunk_cat expects non-negative dim when input tensors have different ndims', + ) + + # input tensors have different ndims but dim >= ndim of some input tensors + sizes, dim, num_chunks = [torch.Size([2, 3]), torch.Size([4,])], 1, 3 + tensors = [make_arg(size) for size in sizes] + yield ErrorInput( + SampleInput(tensors, args=(dim, num_chunks)), + error_regex='_chunk_cat expects dim < ndim for all input tensors', + ) + + # some tensors have different sizes for 0, ..., dim-1 dimensions. + sizes, dim, num_chunks = [torch.Size([2, 3, 4]), torch.Size([4, 3])], 1, 3 + tensors = [make_arg(size) for size in sizes] + yield ErrorInput( + SampleInput(tensors, args=(dim, num_chunks)), + error_regex='_chunk_cat expects same sizes of 0,...,dim-1 dimensions for all tensors', + ) + + # negative num_chunks + sizes, dim, num_chunks = [torch.Size([2,]), torch.Size([3,])], 0, -1 + tensors = [make_arg(size) for size in sizes] + yield ErrorInput( + SampleInput(tensors, args=(dim, num_chunks)), + error_regex='_chunk_cat expects positive num_chunks', + ) + + # zero as num_chunks + sizes, dim, num_chunks = [torch.Size([2,]), torch.Size([3,])], 0, 0 + tensors = [make_arg(size) for size in sizes] + yield ErrorInput( + SampleInput(tensors, args=(dim, num_chunks)), + error_regex='_chunk_cat expects positive num_chunks', + ) + + # empty input tensor list + dim, num_chunks = 0, 1 + yield ErrorInput( + SampleInput([], args=(dim, num_chunks)), + error_regex='_chunk_cat expects a non-empty input tensor list', + ) + + # empty input tensor with 0 elements + sizes, dim, num_chunks = [torch.Size([0,]), torch.Size([3,])], 0, 1 + tensors = [make_arg(size) for size in sizes] + yield ErrorInput( + SampleInput(tensors, args=(dim, num_chunks)), + error_regex='_chunk_cat expects non-empty tensor', + ) + + +def sample_inputs_cat_concat(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases: tuple[tuple, tuple, dict] = ( # type: ignore[assignment] + ((S, S), (S, S), {'dim': -1}), + ((S, S), (S, S), {'dim': 1}), + ((M, S), (S, S), {'dim': 0}), # different shapes + ((1, 2, 3), (1, 2, 3), {'dim': -2}), + ((0,), (0,), {'dim': 0}), # empty tensor + ((0,), (S, S), {'dim': 1}), # empty tensor with unempty and dim=1 (special case for legacy_cat_wrap_dim) + ((0, S), (S, S), {'dim': 0}), + ((1,), (1,), {}) # dim not passed, fallback to default + ) + + for input_shape1, input_shape2, kwargs in cases: + yield SampleInput([make_arg(input_shape1), make_arg(input_shape2)], kwargs=kwargs) + + # from coat_lite_mini + yield SampleInput([make_arg((2, 2, 2, 2), memory_format=torch.channels_last)], args=(1,),) + +def error_inputs_cat(op_info, device, **kwargs): + + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # error inputs for more than one element of the written-to tensor refer to a single memory location + yield ErrorInput(SampleInput([make_arg((S, S)), make_arg((S, S))], + kwargs={'out': make_arg((1, S)).expand((2 * S, S))}), + error_regex='unsupported operation') + + # error inputs for empty tensors + yield ErrorInput(SampleInput([], kwargs={'dim': 1}), + error_regex='non-empty list of Tensors', error_type=ValueError) + + # error inputs for different sizes + yield ErrorInput(SampleInput([make_arg((S, S, L, L)), make_arg((S, 0, L - 1, L))], kwargs={'dim': 1}), + error_regex='Sizes of tensors must match except in dimension') + yield ErrorInput(SampleInput([make_arg((S, 0, L - 1, L)), make_arg((S, S, L, L))], kwargs={'dim': 1}), + error_regex='Sizes of tensors must match except in dimension') + + # error inputs for different dimensions + yield ErrorInput(SampleInput([make_arg((S - 1, 0)), make_arg((S, 0, L - 1, L))], kwargs={'dim': 1}), + error_regex='Tensors must have same number of dimensions') + yield ErrorInput(SampleInput([make_arg((S, 0, L - 1, L)), make_arg((S - 1, 0))], kwargs={'dim': 1}), + error_regex='Tensors must have same number of dimensions') + + # error inputs for same memory locations + x = torch.zeros((0), device=device) + y = torch.randn((4, 6), device=device) + + err_msg = "the written-to tensor refer to a single memory location" + + yield ErrorInput(SampleInput((x, y), kwargs={'dim': 0, 'out': x}), + error_regex=err_msg) + yield ErrorInput(SampleInput((x, y), kwargs={'dim': 0, 'out': y}), + error_regex=err_msg) + + z = torch.zeros((4, 6), device=device) + yield ErrorInput(SampleInput((y, z), kwargs={'out': z[:2, :]}), + error_regex=err_msg) + + # error inputs for different devices + if torch.device(device).type == 'cuda': + x_cuda = make_tensor((3, 3), device=device, dtype=torch.float32) + y_cpu = make_tensor((3, 3), device='cpu', dtype=torch.float32) + yield ErrorInput(SampleInput((x_cuda, y_cpu)), + error_regex='Expected all tensors to be on the same device') + + # error inputs for different input sizes for more than 2 tensors + yield ErrorInput(SampleInput([make_arg((L, 1)), make_arg((L, 1, 1)), make_arg((L, 1, 1))]), + error_regex='Tensors must have same number of dimensions') + + yield ErrorInput(SampleInput([make_arg((S, 1, M)), make_arg((S, 1, 1)), make_arg((S, M, 1))], + kwargs={'dim': 1}), + error_regex='Sizes of tensors must match') + + # error inputs for None input + yield ErrorInput(SampleInput((make_arg((S, 1, 1)), None)), error_type=TypeError, + error_regex='got None') + + # error inputs for zero-dimensional tensors + yield ErrorInput(SampleInput([make_arg(()), make_arg(())]), + error_regex='zero-dimensional.*cannot be concatenated') + + # error inputs for different dtype of out tensors + d = make_tensor((2, 3), device=device, dtype=torch.double if not device.startswith("mps") else torch.float16) + x = make_tensor((2, 3), device=device, dtype=torch.float32) + yield ErrorInput(SampleInput(x, kwargs={'out': d}), error_type=TypeError, + error_regex='invalid combination of arguments') + +def reference_inputs_cat(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_cat_concat(op, device, dtype, requires_grad, **kwargs) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Noncontiguous type promoting tensors + a = make_arg((3, 4, 2)) + b = make_arg((3, 2, 2), noncontiguous=True, dtype=highest_precision_float(device)) + c = make_arg((3, 3, 2), dtype=torch.float16).permute(1, 0, 2) + + yield SampleInput((a, b, c), kwargs={'dim': 1}) + + # Special 1D tensor with dim length of 0 case + a = make_arg((0,)) + b = make_arg((3, 2, 2)) + + yield SampleInput((a, b, a)) + yield SampleInput((a, a, a)) + +def _elementwise_type_promo_np(*args, type_promotion_kind): + def _maybe_torch(x): + if isinstance(x, np.ndarray): + return torch.from_numpy(x) + return x + + flattened = pytree.arg_tree_leaves(*args) + transformed = tuple(_maybe_torch(a) for a in flattened) + result_dtype, _ = prims.utils.elementwise_dtypes( + *transformed, + type_promotion_kind=type_promotion_kind) + return torch_to_numpy_dtype_dict[result_dtype] + +def _cat_np(input_seq, dim=0): + inputs = tuple(a for a in input_seq if not (a.ndim == 1 and a.size == 0)) + + if len(inputs) == 0: + np_dtype = _elementwise_type_promo_np( + input_seq, + type_promotion_kind=prims.utils.ELEMENTWISE_TYPE_PROMOTION_KIND.NO_OPMATH) + return np.empty(0, dtype=np_dtype) + + return np.concatenate(inputs, axis=dim) + +def _floor_divide_np(a, b): + dtype = _elementwise_type_promo_np( + a, + b, + type_promotion_kind=prims.utils.ELEMENTWISE_TYPE_PROMOTION_KIND.DEFAULT) + if isinstance(a, np.ndarray): + a = a.astype(dtype) + if isinstance(b, np.ndarray): + b = b.astype(dtype) + return np.floor_divide(a, b) + +def sample_inputs_hstack_dstack_vstack(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + tensor_shapes = ( + # First Tensor being 1-D is special + # case for hstack + ((S,), (S,), (S,)), + ((S, S), (S, S), (S, S)), + ) + for s1, s2, s3 in tensor_shapes: + tensors = (make_arg(s1,), make_arg(s2,), make_arg(s3)) + yield SampleInput(tensors) + +def error_inputs_hstack_dstack_vstack(op, device): + make_arg = partial(make_tensor, dtype=torch.int32, device=device, requires_grad=False) + tensor_shapes = ( + ((S,), (S, S, S, S), (S,)), + ) + for s1, s2, s3 in tensor_shapes: + tensors = (make_arg(s1,), make_arg(s2,), make_arg(s3)) + # Different dimension tensor + yield ErrorInput(SampleInput(tensors), error_regex="Tensors must have same number of dimensions") + + # empty tensor list + yield ErrorInput(SampleInput(()), error_regex="expects a non-empty TensorList") + +def sample_inputs_unbind(op_info, device, dtype, requires_grad, **kwargs): + # Note: we don't do any tests where we unbind along 0-length dims + # because in that case unbind returns and empty tuple, and that breaks + # some assumptions in some backward tests in test_ops.py + shape_dims = (((S,), 0), + ((S, S), 0), + ((S, S), 1), + ((S, S), -1), + ((S, 0, S), 0), + ((S, S, S), 1), + ) + for shape, dim in shape_dims: + yield SampleInput(make_tensor(shape, dtype=dtype, device=device, + requires_grad=requires_grad), + args=(dim,)) + +def error_inputs_unbind(op_info, device): + make_arg = partial(make_tensor, dtype=torch.int32, device=device, requires_grad=False) + yield ErrorInput(SampleInput(make_arg(()), args=(0,)), error_type=IndexError, + error_regex="Dimension specified as 0 but tensor has no dimensions") + yield ErrorInput(SampleInput(make_arg((2,)), args=(2,)), error_type=IndexError, + error_regex="Dimension out of range") + +def reference_unbind(t, dim): + """A numpy implementation of torch.unbind""" + return tuple(s.squeeze(dim) for s in np.split(t, t.shape[dim], dim)) + +def sample_inputs_gather(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, low=None, high=None) + yield SampleInput( + make_arg((M, S)), + 0, + gather_variable((S, S), 1, M, True, device=device)) + yield SampleInput( + make_arg((M, S)), + 0, + gather_variable((S, S), 1, M, True, device=device).to(torch.int32)) + yield SampleInput( + make_arg((M, S)), + 1, + gather_variable((M, S // 2), 0, S, True, device=device)) + # Empty index tensor case, see: https://github.com/pytorch/pytorch/pull/65006 + yield SampleInput( + make_arg((S,)), + 0, + torch.tensor([], dtype=torch.uint8, device=device)) + yield SampleInput( + make_arg((S,)), + 0, + torch.tensor([[], []], dtype=torch.uint8, device=device)) + # 0D tensor case + yield SampleInput( + make_arg(()), + 0, + torch.tensor([0], dtype=torch.int64, device=device)) + yield SampleInput( + make_arg(()), + 0, + torch.tensor(0, dtype=torch.int64, device=device)) + +def _fill_indices(idx, dim, dim_size, elems_per_row, m, n, o): + for i in range(1 if dim == 0 else m): + for j in range(1 if dim == 1 else n): + for k in range(1 if dim == 2 else o): + ii = [i, j, k] + ii[dim] = slice(0, idx.size(dim) + 1) + idx[tuple(ii)] = torch.randperm(dim_size)[0:elems_per_row] + +def error_inputs_gather(op_info, device, **kwargs): + # src is [1, 2] + # [3, 4] + src = torch.tensor(((1, 2), (3, 4)), device=device, dtype=torch.float32) + + # idx is [0, 0] + # [1, 0] + idx = torch.tensor(((0, 0), (1, 0)), device=device, dtype=torch.long) + + # Index should be smaller than self except on dimension 1 + bad_src = make_tensor((1, 1), device=device, dtype=torch.float32) + yield ErrorInput(SampleInput(bad_src, args=(1, idx,)), + error_regex="Size does not match at dimension 0") + + # TODO: FIXME + # out.dtype must match src.dtype + # Creates new src & idx since SampleInputs can't share tensors + src = torch.tensor(((1, 2), (3, 4)), device=device, dtype=torch.float32) + idx = torch.tensor(((0, 0), (1, 0)), device=device, dtype=torch.long) + out = torch.empty((2, 2), device=device, dtype=torch.float16 if torch.device(device).type == 'mps' else torch.float64) + yield ErrorInput(SampleInput(src, args=(1, idx), kwargs={'out': out}), + error_regex="Expected out tensor to have dtype") + + # src and index tensors must have the same # of dimensions + # idx too few dimensions + src = torch.tensor(((1, 2), (3, 4)), device=device, dtype=torch.float32) + idx = torch.tensor((0, 0), device=device, dtype=torch.long) + yield ErrorInput(SampleInput(src, args=(1, idx)), + error_regex="Index tensor must have the same number of dimensions") + + # src too few dimensions + src = torch.tensor((1, 2), device=device, dtype=torch.float32) + idx = torch.tensor(((0, 0), (1, 0)), device=device, dtype=torch.long) + yield ErrorInput(SampleInput(src, args=(0, idx)), + error_regex="Index tensor must have the same number of dimensions") + + # index out of bounds + # NOTE: this ErrorInput is guarded because bounds checking does not occur on CUDA devices + if torch.device(device).type == 'cpu': + src = torch.tensor(((1, 2), (3, 4)), device=device, dtype=torch.float32) + idx = torch.tensor(((0, 23), (1, 0)), device=device, dtype=torch.long) + yield ErrorInput(SampleInput(src, args=(1, idx,)), + error_regex="index 23 is out of bounds for dimension") + + x = torch.rand((1,), device=device).expand((3,)) + src = torch.rand((6,), device=device) + ind = torch.tensor([2, 1, 0], device=device, dtype=torch.int64) + + yield ErrorInput(SampleInput(src, args=(0, ind,), kwargs=dict(out=x)), + error_type=RuntimeError, + error_regex='unsupported operation') + + yield ErrorInput(SampleInput(src, args=(0, ind,), kwargs=dict(out=src)), + error_type=RuntimeError, + error_regex='unsupported operation') + + yield ErrorInput(SampleInput(ind.clone(), args=(0, ind[1:],), kwargs=dict(out=ind[:1])), + error_type=RuntimeError, + error_regex='unsupported operation') + +def error_inputs_take(op_info, device, **kwargs): + x = torch.rand((1,), device=device).expand((3,)) + src = torch.rand((6,), device=device) + ind = torch.tensor([2, 1, 0], device=device, dtype=torch.int64) + + yield ErrorInput(SampleInput(src, args=(ind,), kwargs=dict(out=x)), + error_type=RuntimeError, + error_regex='unsupported operation') + + yield ErrorInput(SampleInput(src, args=(ind,), kwargs=dict(out=src)), + error_type=RuntimeError, + error_regex='unsupported operation') + + yield ErrorInput(SampleInput(ind.clone(), args=(ind[1:],), kwargs=dict(out=ind[:-1])), + error_type=RuntimeError, + error_regex='unsupported operation') + +# Error inputs for scatter +def error_inputs_scatter_and_scatter_add(op_info, device, **kwargs): + # Error when self.dtype != src.dtype (and src is not a scalar) + src = make_tensor((2, 5), device=device, dtype=torch.float32) + idx = torch.tensor(((0, 1), (1, 2)), device=device, dtype=torch.long) + dst = torch.zeros((3, 5), device=device, dtype=torch.float16 if torch.device(device).type == 'mps' else torch.double) + yield ErrorInput(SampleInput(dst, args=(0, idx, src)), + error_regex="Expected self.dtype to be equal to src.dtype") + + # Index and destination must have the same number of dimensions + src = make_tensor((2, 5), device=device, dtype=torch.float32) + idx = torch.tensor(((0, 1), (1, 2)), device=device, dtype=torch.long) + dst = torch.zeros((3, 5, 3), device=device, dtype=torch.float32) + yield ErrorInput(SampleInput(dst, args=(0, idx, src)), + error_regex="Index tensor must have the same number of dimensions as self tensor") + + # Index and src must have the same number of dimensions when src is not a scalar + src = make_tensor((2, 5, 2), device=device, dtype=torch.float32) + idx = torch.tensor(((34, 1), (1, 2)), device=device, dtype=torch.long) + dst = torch.zeros((3, 5), device=device, dtype=torch.float32) + yield ErrorInput(SampleInput(dst, args=(0, idx, src)), + error_regex="Index tensor must have the same number of dimensions as src tensor") + + # Index out of bounds + # NOTE: this ErrorInput is guarded because bounds checking does not occur on CUDA devices + if torch.device(device).type == 'cpu': + src = make_tensor((2, 5), device=device, dtype=torch.float32) + idx = torch.tensor(((34, 1), (1, 2)), device=device, dtype=torch.long) + dst = torch.zeros((3, 5), device=device, dtype=torch.float32) + yield ErrorInput(SampleInput(dst, args=(0, idx, src)), + error_regex="index 34 is out of bounds for dimension 0 with size 3") + +def error_inputs_renorm(op_info, device, **kwargs): + zero_d = torch.randn((), device=device) + yield ErrorInput(SampleInput(zero_d, args=(0.5, 0, 1.0)), error_type=RuntimeError, + error_regex="needs at least 2 dimensions, got 0 dimensions") + + +def error_inputs_ormqr(op_info, device, **kwargs): + zero_d = torch.randn((), device=device) + yield ErrorInput(SampleInput(zero_d, args=(zero_d, zero_d)), error_type=RuntimeError, + error_regex="input must have at least 2 dimensions") + + # https://github.com/pytorch/pytorch/issues/85218 + tensor_0 = torch.full((5, 0,), 1, device=device) + tensor_1 = torch.full((5,), 1, device=device) + tensor_2 = torch.full((5, 5,), 1, device=device) + bool_3 = True + bool_4 = True + yield ErrorInput(SampleInput(tensor_0, args=(tensor_1, tensor_2, bool_3, bool_4)), error_type=RuntimeError, + error_regex=r"tau.shape\[-1\] must be equal to min\(other.shape\[-2\], input.shape\[-1\]\)") + + +def error_inputs_diag(op_info, device, **kwargs): + zero_d = torch.randn((), device=device) + yield ErrorInput(SampleInput(zero_d, args=(0,)), error_type=RuntimeError, + error_regex="1D or 2D") + zero_d = torch.randn(1, 1, 1, device=device) + yield ErrorInput(SampleInput(zero_d, args=(0,)), error_type=RuntimeError, + error_regex="1D or 2D") + +def error_inputs_embedding(op_info, device, **kwargs): + indices = torch.rand(2, 2, device=device).long() + weights = [ + torch.tensor(1.0, device=device), + torch.tensor(1.0, device=device).reshape(1, 1, 1), + ] + + for weight in weights: + yield ErrorInput(SampleInput(weight, args=(indices,)), error_type=RuntimeError, + error_regex="'weight' must be 2-D") + + +def error_inputs_t(op_info, device, **kwargs): + yield ErrorInput( + SampleInput(torch.randn(2, 3, 4, 5, device=device)), + error_regex="expects a tensor with <= 2", + ) + + +def error_inputs_multinomial(op_info, device, **kwargs): + dtype = highest_precision_float(device) + x = torch.empty(1, 2, 3, dtype=dtype, device=device) + yield ErrorInput(SampleInput(x, args=(2,)), + error_regex="prob_dist must be 1 or 2 dim") + + x = torch.empty(1, 2, dtype=torch.long, device=device) + yield ErrorInput(SampleInput(x, args=(2,)), + error_regex="multinomial only supports floating-point dtypes for input") + + x = torch.empty(1, 2, dtype=dtype, device=device) + y = torch.empty(1, 2, dtype=dtype, device=device) + yield ErrorInput(SampleInput(x, args=(2,), kwargs=dict(out=y)), + error_regex="multinomial expects Long tensor out") + + x = torch.empty(2, dtype=dtype, device=device) + yield ErrorInput(SampleInput(x, args=(0,)), + error_regex="cannot sample n_sample <= 0 samples") + + x = torch.empty(2, dtype=dtype, device=device) + yield ErrorInput(SampleInput(x, args=(-1,)), + error_regex="cannot sample n_sample <= 0 samples") + + x = torch.empty(2, dtype=dtype, device=device) + yield ErrorInput(SampleInput(x, args=(3, False,)), + error_regex="cannot sample n_sample > prob_dist") + + x = torch.empty(16777217, dtype=dtype, device=device) + yield ErrorInput(SampleInput(x, args=(3,)), + error_regex="number of categories cannot exceed") + + inputs = ((1., -1., 1.), (1., inf, 1.), (1., -inf, 1.), (1., 1., nan)) + + err_msg1 = "probability tensor contains either `inf`, `nan` or element < 0" + err_msg2 = "invalid multinomial distribution" + + rep_arg = (False, True) if torch.device(device).type == 'cpu' else (False,) + + if torch.device(device).type == 'cpu': + for rep in rep_arg: + kwargs = {'num_samples': 2, 'replacement': rep} + + for shape in inputs: + # error case when input tensor contains `inf`, `nan` or negative element + yield ErrorInput(SampleInput(torch.tensor(shape), kwargs=kwargs), + error_regex=err_msg1 if rep is False else err_msg2) + + # error case for the invalid multinomial distribution (sum of probabilities <= 0), 1-D input + x = torch.zeros(3, device=device) + yield ErrorInput(SampleInput(x, kwargs=kwargs), + error_regex=err_msg2) + + # error case for the invalid multinomial distribution (sum of probabilities <= 0), 2-D input + x = torch.zeros(3, 3, device=device) + yield ErrorInput(SampleInput(x, kwargs=kwargs), + error_regex=err_msg2) + + # error case for the invalid multinomial distribution + x[1, :] = 1 + yield ErrorInput(SampleInput(x, kwargs=kwargs), + error_regex=err_msg2) + +def error_inputs_gradient(op_info, device, **kwargs): + for dtype in [torch.long, torch.float32, torch.complex64]: + t = torch.tensor([[1, 2, 3], [4, 5, 6], [7, 8, 9]], device=device, dtype=dtype) + + dim = (1, 0) + spacing = [0.1] + yield ErrorInput(SampleInput(t, kwargs=dict(spacing=spacing, dim=dim, edge_order=1)), + error_type=RuntimeError, + error_regex='torch.gradient expected spacing to be unspecified, a scalar ') + + yield ErrorInput(SampleInput(t, kwargs=dict(edge_order=3)), + error_type=RuntimeError, + error_regex='torch.gradient only supports edge_order=1 and edge_order=2.') + + dim = (1, 1) + spacing = 0.1 + yield ErrorInput(SampleInput(t, kwargs=dict(spacing=spacing, dim=dim, edge_order=1)), + error_type=RuntimeError, + error_regex='dim 1 appears multiple times in the list of dims') + + dim = (0, 1) + coordinates = [torch.tensor([1, 2, 4], device='cpu'), torch.tensor([1, 2, 4], device='meta')] + yield ErrorInput(SampleInput(t, kwargs=dict(spacing=coordinates, dim=dim, edge_order=1)), + error_type=RuntimeError, + error_regex='torch.gradient expected each tensor to be on the same device,') + + yield ErrorInput(SampleInput(t, kwargs=dict(dim=3)), + error_type=IndexError, error_regex='') + + t = torch.tensor([[1], [2], [3]]) + yield ErrorInput(SampleInput(t, kwargs=dict(edge_order=1)), + error_type=RuntimeError, + error_regex='torch.gradient expected each dimension size to be at least') + + t = torch.tensor([[1, 2], [3, 4]]) + yield ErrorInput(SampleInput(t, kwargs=dict(edge_order=2)), + error_type=RuntimeError, + error_regex='torch.gradient expected each dimension size to be at least') + +def sample_inputs_rrelu(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_elementwise_unary( + op_info, device, dtype, requires_grad, op_kwargs=dict(lower=0., upper=1., training=True)) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg(S)) + yield SampleInput(make_arg(S), training=False) + +def error_inputs_rrelu(op_info, device, **kwargs): + input = make_tensor((S, S), device=device, dtype=torch.float32) + yield ErrorInput(SampleInput(input, kwargs={'lower': 0.3, 'upper': 0.1}), + error_regex='Lower bound should be less than or equal to the upper bound') + +def error_inputs_masked_select(op_info, device, **kwargs): + x = torch.rand((1,), device=device).expand((3,)) + y = torch.rand((6,), device=device) + mask = torch.tensor([True, False, True, True, False, False], device=device) + + yield ErrorInput(SampleInput(y, args=(mask,), kwargs=dict(out=x)), + error_type=RuntimeError, + error_regex='unsupported operation') + + yield ErrorInput(SampleInput(y, args=(mask,), kwargs=dict(out=y)), + error_type=RuntimeError, + error_regex='unsupported operation') + + yield ErrorInput(SampleInput(mask.clone(), args=(mask,), kwargs=dict(out=mask)), + error_type=RuntimeError, + error_regex='unsupported operation') + +def error_inputs_median(op_info, device, **kwargs): + x = torch.tensor([[[[[[[[[[[[[[[[[[[[[[[[[nan], + [nan]]]]]]]]]]]]]]]]]]]]]]]]], device=device) + if device == 'cuda': + yield ErrorInput(SampleInput(x, kwargs=dict(dim=(-1))), + error_type=RuntimeError, + error_regex='CUDA Tensors cannot have more than 25 dimensions') + else: + return + + +def error_inputs_index_select(op_info, device, **kwargs): + x = torch.rand((1, 6), device=device).expand((2, 6)) + y = torch.rand((3, 6), device=device) + ind = torch.tensor([0, 1], dtype=torch.int64, device=device) + + yield ErrorInput(SampleInput(y, args=(1, ind,), kwargs=dict(out=x)), + error_type=RuntimeError, + error_regex='unsupported operation') + +def error_inputs_index_add(op_info, device, **kwargs): + result = torch.tensor([[1., 2.], [4., 5.], [7., 8.]]) + source = torch.tensor([2., 4.]) + + yield ErrorInput(SampleInput(result, args=(0, torch.tensor([0, 2]), source)), + error_type=RuntimeError, + error_regex=r'source tensor shape must match self tensor shape, ' + r'excluding the specified dimension. Got self.shape = \[3, 2\] source.shape = \[2\]') + +def error_inputs_logcumsumexp(op_info, device, **kwargs): + dim = 3 + srcs = [torch.randn(5, 2, device=device), torch.randn(0, 2, device=device)] + for src in srcs: + yield ErrorInput(SampleInput(src, args=(dim,)), + error_type=IndexError, + error_regex='Dimension out of range') + +def sample_inputs_take_along_dim(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, low=None, high=None) + yield SampleInput( + make_arg((S, S)), gather_variable((S, S), 1, S, True, device=device), 0) + + # `indices` broadcast + yield SampleInput( + make_arg((S, S)), gather_variable((1, S // 2), 0, S, True, device=device), 1) + + # `self` broadcast + yield SampleInput( + make_arg((1, S)), gather_variable((S, S // 2), 0, S, True, device=device), 1) + + # without `dim` arg + yield SampleInput( + make_arg((S, S)), gather_variable((S, S // 2), 0, S, True, device=device)) + + # Negative indices sample — guarded against python_ref + if not kwargs.get('is_python_ref', False): + neg_idx = gather_variable((S, S), 1, S, True, device=device) - S + yield SampleInput( + make_arg((S, S)), + neg_idx, + 1) + + +def error_inputs_aminmax_amax_amin(op_info, device, is_ref=False, **kwargs): + + # Error Inputs for zero-dim tensors, when 'dim' arg is not provided. + shape = (S, 0, S) + err_msg_amax_amin = "reduction" + err_msg_aminmax = "cannot compute aminmax over an empty dimension as the operation has no identity" + if op_info.name in ['amax', 'amin', '_refs.amax', '_refs.amin']: + yield ErrorInput(SampleInput(torch.rand(shape, device=device)), error_regex=err_msg_amax_amin) + elif op_info.name == 'aminmax': + yield ErrorInput(SampleInput(torch.rand(shape, device=device)), error_regex=err_msg_aminmax) + + # Error Inputs for tensors with more than 64 dimension + sizes = [1] * 65 + err_msg1 = "only tensors with up to 64 dims are supported" + yield ErrorInput(SampleInput(torch.randn(sizes, device=device), kwargs={'dim': -1}), + error_regex=err_msg1) + yield ErrorInput(SampleInput(torch.randn(sizes, device=device), kwargs={'dim': 64}), + error_regex=err_msg1) + + # Error Inputs for repeated 'dim' + if op_info.name in ['amax', 'amin', '_refs.amax', '_refs.amin']: + dims = [(0, 0), (0, -4)] + err_msg2 = "in the list of dims" + x = torch.randn(S, S, S, S, device=device) + for dim in dims: + yield ErrorInput(SampleInput(x, kwargs={'dim': dim}), error_regex=err_msg2) + + # Error Input for illegal dtype + input5 = torch.randn(L, L, dtype=torch.float32, device=device) + max_values = torch.empty(L, dtype=torch.float32, device=device) + min_values = torch.empty(L, dtype=torch.double, device=device) + illegal_values = torch.empty(L, dtype=torch.int, device=device) + + # Unlike regular PyTorch, amax and amin refs don't require input and out + # dtypes to match exactly: + # https://github.com/pytorch/pytorch/pull/87765#pullrequestreview-1162023824 + if is_ref: + err_msg_amax_amin2 = ("Attempting to cast from torch.float32 to out tensor with dtype " + "torch.int32, but this can't be cast because it is not safe!") + else: + err_msg_amax_amin2 = ("Expected the dtype for input and out to match, but got Float " + "for input's dtype and Int for out's dtype.") + err_msg_aminmax2 = "Expected out tensor to have dtype float, but got double instead" + + if op_info.name in ['amax', 'amin', '_refs.amax', '_refs.amin']: + yield ErrorInput(SampleInput(input5, kwargs={'dim': 0, 'out': illegal_values}), + error_regex=err_msg_amax_amin2) + elif op_info.name == 'aminmax': + yield ErrorInput(SampleInput(input5, kwargs={'dim': 0, 'out': (max_values, min_values)}), + error_regex=err_msg_aminmax2) + + # Error Inputs for functions to raise an error on specified zero'd dimension as reduction dim + err_msg3 = "reduction" + # FIXME: eager and ref impl throw different types of errors + error_type = IndexError if 'refs' not in op_info.name else RuntimeError + yield ErrorInput(SampleInput(torch.rand(shape, device=device), kwargs={'dim': 1}), + error_type=error_type, error_regex=err_msg3) + +def sample_inputs_aminmax(op_info, device, dtype, requires_grad, **kwargs): + test_cases: tuple[tuple, dict] = ( # type: ignore[assignment] + ((S, S, S), {}), + ((S, S, S), {'dim': 1}), + ((S, S, S), {'dim': 1, 'keepdim': True}), + ((), {'dim': 0}), + ((), {}), + ((), {'dim': 0, 'keepdim': True}), + ((S, 0, S), {'dim': 0}), + ) + + for shape, kwargs in test_cases: + yield SampleInput( + make_tensor(shape, dtype=dtype, device=device, requires_grad=requires_grad), + **kwargs) + +def error_inputs_diff(op_info, device, **kwargs): + t = torch.rand((1, 3), device=device) + n = -1 + yield ErrorInput(SampleInput(t, args=(n, ), kwargs=kwargs), + error_type=RuntimeError, + error_regex=f'order must be non-negative but got {n}') + +def sample_inputs_diff(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + test_cases = ( + ((1,), 0, None, None), + ((S,), 0, None, None), + ((S, 1), 0, None, None), + ((S, 1), 1, None, None), + ((S, S), 0, None, None), + ((S, S), 1, None, None), + ((S, S), 0, (1, S), (2, S)), + ((S, S), 0, None, (2, S)), + ((XS, XS, XS), 1, None, None), + ((XS, XS, XS), 2, None, None), + ((XS, XS, XS), 1, (XS, 1, XS), (XS, 1, XS)), + ((XS, XS, XS), 2, (XS, XS, 1), (XS, XS, 1)), + ((XS, XS, XS), 2, (XS, XS, XS), (XS, XS, XS)),) + + for size, dim, size_prepend, size_append in test_cases: + prepend_size = 0 if (size_prepend is None) else size_prepend[dim] + append_size = 0 if (size_append is None) else size_append[dim] + dim_size = size[dim] + prepend_size + append_size + for n in range(dim_size): + input_tensor = make_arg(size) + prepend = make_arg(size_prepend) if size_prepend else None + append = make_arg(size_append) if size_append else None + yield SampleInput(input_tensor, n, dim, prepend, append) + + # add some samples with n > dim_size + yield SampleInput(make_arg((XS, XS, XS)), S + 1, 1) + yield SampleInput(make_arg((XS, XS, XS)), S * 3 + 2, 2, make_arg((XS, XS, XS)), make_arg((XS, XS, XS))) + +def sample_inputs_histogram(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + sizes = ((), (S,), (S, S), (S, S, S), (S, 1, S), (S, 0, S)) + + for size, bin_ct, weighted, density in product(sizes, range(1, 5), [False, True], [False, True]): + input_tensor = make_arg(size) + weight_tensor = make_arg(size) if weighted else None + + yield SampleInput(input_tensor, bin_ct, + weight=weight_tensor, density=density) + + bins_tensor = make_arg((bin_ct + 1,)) + sorted_bins, _bins_indices = torch.sort(bins_tensor) + yield SampleInput(input_tensor, sorted_bins, + weight=weight_tensor, density=density) + +def sample_inputs_histogramdd(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + sizes = ((S, S), (S, S, S), (S, 1, S), (S, 0, S)) + bin_ct_patterns = ((1, 1, 1, 1, 1), (2, 3, 2, 3, 2), (3, 2, 3, 2, 3)) + + for size, bin_ct_pattern, weighted, density in product(sizes, bin_ct_patterns, [False, True], [False, True]): + input_tensor = make_arg(size) + bin_ct = bin_ct_pattern[:size[-1]] + weight_tensor = make_arg(size[:-1]) if weighted else None + + yield SampleInput(input_tensor, bin_ct, + weight=weight_tensor, density=density) + + bins_tensor = [make_arg(ct + 1) for ct in bin_ct] + yield SampleInput(input_tensor, bins_tensor, + weight=weight_tensor, density=density) + +def error_inputs_histogramdd(opinfo, device, **kwargs): + invalid_bins = [1, 1, 1, 1, 1] + make_arg = partial(make_tensor, dtype=torch.float, device=device, requires_grad=False) + msg = "histogramdd: The size of bins must be equal to the innermost dimension of the input." + yield ErrorInput(SampleInput(make_arg(5, 6), invalid_bins), error_regex=msg) + +def sample_inputs_histc(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + sizes = ((), (S,), (S, S), (S, S, S), (S, 1, S), (S, 0, S)) + + for size, min, max in product(sizes, [0, -10], [0, 10]): + # construct sample input omitting bins arg + yield SampleInput(make_arg(size), min=min, max=max) + + # construct sample inputs with a few different bins values + for bins in [1, 3, 10]: + yield SampleInput(make_arg(size), bins=bins, min=min, max=max) + +def sample_inputs_bincount(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + for size, weighted in product((S, M), [False, True]): + input_tensor = torch.randint(0, size, (size,), dtype=dtype, device=device) + weight_tensor = make_arg((size,)) if weighted else None + + max_val = int(input_tensor.max().item()) + + for minlength in [0, max_val // 2, max_val, 2 * max_val]: + yield SampleInput( + input_tensor, weights=weight_tensor, minlength=minlength) + +def sample_inputs_bucketize(op_info, device, dtype, requires_grad, reference_inputs_mode=False, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + sizes = (((), S), ((S,), S), ((S, S), S), ((S, S, S), S), ((S, 1, S), S), ((S, 0, S), S)) + + if reference_inputs_mode: + sizes += (((256,), 128), ((128,), 256), ((32, 32), 11), ((32, 4, 32), 33)) + + for (input_shape, nb), out_int32, right in product(sizes, [False, True], [False, True]): + input_tensor = make_arg(input_shape) + boundaries = make_arg(nb).msort() + + yield SampleInput(input_tensor, boundaries, + out_int32=out_int32, right=right) + +reference_inputs_bucketize = partial(sample_inputs_bucketize, reference_inputs_mode=True) + +def error_inputs_bucketize(opinfo, device, **kwargs): + make_arg = partial(make_tensor, dtype=torch.float, device=device, requires_grad=False) + yield ErrorInput(SampleInput(make_arg((S, S, S)), make_arg((S, S))), + error_regex="boundaries tensor must be 1 dimension") + +def sample_inputs_searchsorted(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + # (unsorted tensor size, (input sizes,), is_scalar) + sizes = ( + ((0,), ((0,),), False), + ((M,), ((), (M,), (M, M)), False), + ((0, 0), ((0, 0),), False), + ((M, M), ((M, M),), False), + ((0, 0, 0), ((0, 0, 0),), False), + ((M, M, M), ((M, M, M),), False), + ((L,), ((),), True), + ) + + for (size, input_sizes, is_scalar), noncontiguous, out_int32, right in product( + sizes, [False, True], [False, True], [False, True] + ): + unsorted_tensor = make_arg(size, noncontiguous=noncontiguous) + for input_size in input_sizes: + input = make_arg(input_size, noncontiguous=noncontiguous) + if is_scalar: + input = input.item() + if np.prod(size) == 0: + boundary_tensor = unsorted_tensor + sorter = make_tensor(size, dtype=torch.int64, device=device, noncontiguous=noncontiguous) + else: + boundary_tensor, sorter = torch.sort(unsorted_tensor) + side = "right" if right else "left" + + yield SampleInput(boundary_tensor, input, out_int32=out_int32, right=right) + yield SampleInput(boundary_tensor, input, out_int32=out_int32, side=side) + + yield SampleInput(unsorted_tensor, input, out_int32=out_int32, right=right, sorter=sorter) + yield SampleInput(unsorted_tensor, input, out_int32=out_int32, side=side, sorter=sorter) + +def sample_inputs_gradient(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, low=None, high=None) + test_cases_float = ( + ((S,), None, None, 1), + ((S,), 2., None, 1), + ((S, S), None, None, 2), + ((S, S), [2.0, 2.1], None, 1), + ((S, S), [2.0, 2.1], (0, 1), 1), + ((4, 4, 4), [2., 1.], (0, 1), 2), + ) + for size, spacing, dim, edge_order in test_cases_float: + t = make_arg(size) + yield SampleInput(t, dim=dim, spacing=spacing, edge_order=edge_order) + + test_cases_tensor = ( + ((3, 3, 3), ((1.1, 2.0, 3.5), (4.0, 2, 6.0)), (0, -1), 1), + ((3, 3, 3), ((1.0, 3.0, 2.0), (8.0, 6.0, 1.0)), (0, 1), 2), + ) + for size, coordinates, dim, edge_order in test_cases_tensor: + t = make_arg(size) + coordinates_tensor_list = [] + for coords in coordinates: + a = torch.tensor(coords, device=device, dtype=dtype) + coordinates_tensor_list.append(a) + yield SampleInput(t, dim=dim, spacing=coordinates_tensor_list, edge_order=edge_order) + +def sample_inputs_getitem(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + test_args = [ + ([1, 2],), + (slice(0, 3),), + ((slice(0, 3), 1),), + (([0, 2, 3], [1, 3, 3], [0, 0, 2]),), + (([0, 0, 3], [1, 1, 3], [0, 0, 2]),), + ((slice(None), slice(None), [0, 3]),), + ((slice(None), [0, 3], slice(None)),), + (([0, 3], slice(None), slice(None)),), + (([0, 3], [1, 2], slice(None)),), + (([0, 3], ),), + (([0, 3], slice(None)),), + (([0, 3], Ellipsis),), + (([0, 2, 3], [1, 3, 3], torch.LongTensor([0, 0, 2])),), + (index_variable(2, S, device=device),), + (mask_not_all_zeros((S,)),), + ] + + for args in test_args: + yield SampleInput(make_arg((S, S, S)), args=args) + + yield SampleInput(make_arg((S, S, S, S)), args=((slice(None), [0, 1], slice(None), [0, 1]),)) + +def sample_inputs_index_put(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + for accumulate in [False, True]: + # Test with indices arg + yield SampleInput( + make_arg((S, S,)), + # As defined in the docs, if accumulate is false, duplicate indices are not supported + (index_variable(2 if accumulate else 1, S, device=device),), + make_arg((2 if accumulate else 1, S)), + accumulate=accumulate) + + # Test with mask arg + mask = torch.zeros(S, dtype=torch.bool) if accumulate else mask_not_all_zeros((S,)) + yield SampleInput( + make_arg((S, S)), (mask, ), make_arg((S,)), accumulate=accumulate) + +def sample_inputs_sort(op_info, device, dtype, requires_grad, **kwargs): + def small_3d_unique(): + res = torch.randperm(S * S * S, dtype=torch.int64, device=device).view(S, S, S) + res = res.to(dtype).requires_grad_(requires_grad) + return res + + def large_1d_unique(): + res = torch.randperm(L * L * L, dtype=torch.int64, device=device) + res = res.to(dtype).requires_grad_(requires_grad) + return res + + # Test case for large tensor. + yield SampleInput(large_1d_unique()) + + # Test cases for small 3d tensors. + # Imitates legacy tests from test/test_torch.py + dims = range(-3, 3) + flag = [True, False] + for dim, descending, stable in product(dims, flag, flag): + # default schema without stable sort + if not (dtype == torch.bool and torch.device(device).type == 'cuda'): + # bool and cuda requires stable sort for stable results, at least + # for the return index + yield SampleInput(small_3d_unique(), dim, descending) + # schema with stable sort, no CUDA support yet + if torch.device(device).type == 'cpu': + yield SampleInput( + small_3d_unique(), dim=dim, descending=descending, stable=stable) + + # Test cases for scalar tensor + tensor_opt = dict(dtype=dtype, device=device, requires_grad=requires_grad) + yield SampleInput(torch.tensor(1, **tensor_opt)) + yield SampleInput(torch.tensor(1, **tensor_opt), 0) + yield SampleInput(torch.tensor(1, **tensor_opt), 0, True) + + # Test cases for empty tensor + yield SampleInput(torch.tensor((), **tensor_opt)) + yield SampleInput(torch.tensor((), **tensor_opt), 0) + yield SampleInput(torch.tensor((), **tensor_opt), 0, True) + + # Test cases for stable sort + yield SampleInput(small_3d_unique(), stable=True) + yield SampleInput(small_3d_unique(), dim=0, stable=True) + yield SampleInput(small_3d_unique(), dim=0, descending=True, stable=True) + +def sample_inputs_threshold(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + sizes = ((), (S,), (S, S), (S, S, S)) + for x_size in sizes: + # threshold and values args must be numbers + yield SampleInput(make_arg(x_size), make_arg(()).item(), make_arg(()).item()) + +def sample_inputs_unique(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + sizes = ((), (S,), (S, S), (S, S, S), (S, 1, S), (S, 0, S)) + + for shape, sorted, return_inverse, return_counts, dim in \ + product(sizes, [False, True], [False, True], [False, True], [None, -2, -1, 0, 1, 2]): + # torch.unique cannot be called if the input tensor has a zero dimension which isn't the selected dim + if 0 in shape and shape.index(0) is not dim: + continue + + # skip invalid dim args + if dim is not None and (dim < -len(shape) or dim >= len(shape)): + continue + + kwargs = dict(sorted=sorted, return_inverse=return_inverse, return_counts=return_counts, dim=dim) + + # construct a test case with only one distinct value + input_t = torch.zeros(shape, dtype=dtype, device=device, requires_grad=requires_grad) + yield SampleInput(input_t, **kwargs) + + # construct a test case with mixed 0s and 1s + input_t = make_arg(shape, dtype=torch.bool, requires_grad=False)\ + .to(dtype).requires_grad_(requires_grad) + yield SampleInput(input_t, **kwargs) + + # construct a test case with many different values + yield SampleInput(make_arg(shape), **kwargs) + +def sample_inputs_unique_consecutive(*args, **kwargs): + for sample_input in sample_inputs_unique(*args, **kwargs): + if not sample_input.kwargs["sorted"]: + sample_input.kwargs.pop("sorted") + yield sample_input + +def sample_inputs_adaptive_avg_pool1d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as (input shape, output size) + cases = ( + ((0, 8, 8), (5,)), + ((3, 8, 8), 5), + ((3, 8, 8), 1) + ) + + for input_shape, output_size in cases: + # Batched + yield SampleInput(make_arg(input_shape), args=(output_size,)) + # Unbatched + yield SampleInput(make_arg(input_shape[1:]), args=(output_size,)) + + +def error_inputs_adaptive_avg_pool1d(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # error inputs for empty output + yield ErrorInput(SampleInput(make_arg((1, 2, 3)), output_size=()), + error_regex="'output_size' should contain one int") + + # error inputs for output_size lesser than 0 + yield ErrorInput(SampleInput(make_arg((1, 1, 1)), output_size=(-1,)), + error_regex="elements of output_size must be greater than or equal to 0") + + +def sample_inputs_adaptive_avg_pool2d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as (input shape, output size) + cases = ( + ((1, 8, 8, 8), (5, 7)), + ((2, 8, 8, 8), (None, 7)), + ((1, 8, 4, 3), (5, None)), + ((1, 8, 4, 3), (None, None)), + ((1, 8, 4, 3), (5)), + ) + + for input_shape, output_size in cases: + # Batched + yield SampleInput(make_arg(input_shape), args=(output_size,)) + # Unbatched + yield SampleInput(make_arg(input_shape[1:]), args=(output_size,)) + + +def error_inputs_adaptive_avg_pool2d(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # error inputs for incorrect input dimension + yield ErrorInput(SampleInput(make_arg((2, 2)), output_size=(2, 2)), + error_type=ValueError, error_regex="Input dimension should be at least 3") + + # error inputs for empty output + yield ErrorInput(SampleInput(make_arg((1, 2, 3, 4)), output_size=()), + error_regex="output_size must be 2") + + # error inputs for output_size lesser than 0 + yield ErrorInput(SampleInput(make_arg((1, 1, 1, 1)), output_size=(-1, 0)), + error_regex="elements of output_size must be greater than or equal to 0") + + +def sample_inputs_adaptive_avg_pool3d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as (input shape, output size) + cases = ( + ((0, 8, 8, 8, 8), (5, 7, 4)), + ((1, 8, 4, 3, 7), (None, None, None)), + ((1, 8, 4, 3, 7), (1, 1, 1)), + ((3, 3, 8, 8, 6), (5, 7, None)), + ((1, 3, 8, 8, 6), (5, None, 2)), + ((3, 3, 8, 8, 6), (None, 3, 2)), + ) + + for input_shape, output_size in cases: + # Batched + yield SampleInput(make_arg(input_shape), args=(output_size,)) + # Unbatched + yield SampleInput(make_arg(input_shape[1:]), args=(output_size,)) + + +def error_inputs_adaptive_avg_pool3d(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # error inputs for incorrect input dimension + yield ErrorInput(SampleInput(make_arg((2, 2, 2)), output_size=(2, 2, 2)), + error_type=ValueError, error_regex="Input dimension should be at least 4") + + # error inputs for empty output + yield ErrorInput(SampleInput(make_arg((1, 2, 3, 4)), output_size=()), + error_regex="output_size must be 3") + + # error inputs for output_size lesser than 0 + yield ErrorInput(SampleInput(make_arg((1, 1, 1, 1, 1)), output_size=(-1, 0, 2)), + error_regex="elements of output_size must be greater than or equal to 0") + + +def sample_inputs_adaptive_max_pool1d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as (input shape, output size) + cases = ( + # ((0, 8, 8), (5,)), + # 0 batch size doesn't work, cannot reshape tensor of 0 elements into shape [0, 8, -1] + ((3, 4, 4), 3), + ((3, 4, 4), 1) + ) + + for shapes, return_idx in product(cases, (True, False)): + # Batched + yield SampleInput(make_arg(shapes[0]), args=(shapes[1], return_idx)) + # Unbatched + yield SampleInput(make_arg(shapes[0][1:]), args=(shapes[1], return_idx)) + + +def error_inputs_adaptive_max_pool1d(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # error inputs for empty output + yield ErrorInput(SampleInput(make_arg((1, 2, 3)), output_size=()), + error_regex="'output_size' should contain one int") + + # error inputs for output_size lesser than 0 + yield ErrorInput(SampleInput(make_arg((1, 1, 1)), output_size=(-1,)), + error_regex="Trying to create tensor with negative dimension") + +def sample_inputs_adaptive_max_pool2d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as (input shape, output size) + cases = ( + # ((0, 8, 8, 8), (5, 7)), + # 0 batch size doesn't work, cannot reshape tensor of 0 elements into shape [0, 8, -1] + ((1, 4, 4, 4), (2, 3)), + ((2, 4, 4, 4), (None, 3)), + ((2, 4, 4, 4), (1, 1)), + ((1, 4, 4, 3), (3, None)), + ((1, 4, 4, 3), (None, None)), + ((1, 4, 4, 3), (3)), + ) + + for shapes, return_idx in product(cases, (True, False)): + # Batched + yield SampleInput(make_arg(shapes[0]), args=(shapes[1], return_idx)) + # Unbatched + yield SampleInput(make_arg(shapes[0][1:]), args=(shapes[1], return_idx)) + +def error_inputs_adaptive_max_pool2d(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # error inputs for incorrect input dimension + yield ErrorInput(SampleInput(make_arg((2, 2)), output_size=(2, 2)), + error_type=ValueError, error_regex="Input dimension should be at least 3") + + # error inputs for empty output + yield ErrorInput(SampleInput(make_arg((1, 2, 3, 4)), output_size=()), + error_regex="internal error") + + # error inputs for output_size lesser than 0 + yield ErrorInput(SampleInput(make_arg((1, 1, 1, 1)), output_size=(-1, 0)), + error_regex="Trying to create tensor with negative dimension") + + +def sample_inputs_adaptive_max_pool3d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as (input shape, output size) + cases = ( + # ((0, 8, 8, 8, 8), (5, 7, 4)), + # 0 batch size doesn't work, cannot reshape tensor of 0 elements into shape [0, 8, -1] + ((1, 4, 4, 3, 5), (None, None, None)), + ((1, 4, 4, 3, 5), (1, 1, 1)), + ((3, 3, 4, 4, 6), (2, 3, None)), + ((1, 3, 4, 4, 6), (3, None, 2)), + ((3, 3, 4, 4, 6), (None, 3, 2)), + ) + + for shapes, return_idx in product(cases, (True, False)): + # Batched + yield SampleInput(make_arg(shapes[0]), args=(shapes[1], return_idx)) + # Unbatched + yield SampleInput(make_arg(shapes[0][1:]), args=(shapes[1], return_idx)) + +def error_inputs_adaptive_max_pool3d(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # error inputs for incorrect input dimension + yield ErrorInput(SampleInput(make_arg((2, 2, 2)), output_size=(2, 2, 2)), + error_type=ValueError, error_regex="Input dimension should be at least 4") + + # error inputs for empty output + yield ErrorInput(SampleInput(make_arg((1, 2, 3, 4)), output_size=()), + error_regex="internal error") + + # error inputs for output_size lesser than 0 + yield ErrorInput(SampleInput(make_arg((1, 1, 1, 1, 1)), output_size=(-1, 0, 2)), + error_regex="Trying to create tensor with negative dimension") + + +class _TestParamsMaxPoolBase: + + def __init__(self) -> None: + self.kwargs = { + 'kernel_size': [3], + 'stride': [2, None], + 'ceil_mode': [True, False], + 'padding': [0, 1], + 'dilation': [1], + 'return_indices': [True, False] + } + + self.shapes = [ + [1, 2, None], # batch + [2], # channels + [3, 6] # signal + ] + + def _gen_shape(self): + for shape in product(*self.shapes): + # shape[0] is None indicates missing batch dimension + if shape[0] is None: + shape = shape[1:] + + yield shape, torch.contiguous_format + # only 2d (N, C, H, W) rank 4 tensors support channels_last memory format + if len(self.shapes) == 4 and len(shape) == 4: + yield shape, torch.channels_last + + def _gen_kwargs(self): + keys = self.kwargs.keys() + for values in product(*self.kwargs.values()): + yield dict(zip(keys, values, strict=True)) + + def gen_input_params(self): + yield from product(self._gen_shape(), self._gen_kwargs()) + +class _TestParamsMaxPool1d(_TestParamsMaxPoolBase): + + def __init__(self) -> None: + super().__init__() + self.kwargs['kernel_size'] += [(3,)] + self.kwargs['stride'] += [(2,)] + self.kwargs['padding'] += [(1,)] + self.kwargs['dilation'] += [(1,)] + +class _TestParamsMaxPool2d(_TestParamsMaxPoolBase): + + def __init__(self) -> None: + super().__init__() + self.kwargs['kernel_size'] += [(3, 2)] + self.kwargs['stride'] += [(2, 1)] + self.kwargs['padding'] += [(1, 1)] + self.kwargs['dilation'] += [(1, 2)] + + self.shapes.append([6]) + +class _TestParamsMaxPool3d(_TestParamsMaxPoolBase): + + def __init__(self) -> None: + super().__init__() + self.kwargs['kernel_size'] += [(3, 2, 3)] + self.kwargs['stride'] += [(2, 1, 2)] + self.kwargs['dilation'] += [(1, 2, 1)] + + self.shapes.append([6]) + self.shapes.append([5]) + +def sample_inputs_max_pool(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + params_generator_type_dict = { + 'nn.functional.max_pool1d': _TestParamsMaxPool1d, + 'nn.functional.max_pool2d': _TestParamsMaxPool2d, + 'nn.functional.max_pool3d': _TestParamsMaxPool3d, + 'max_pool2d_with_indices_backward': _TestParamsMaxPool2d, + } + + params_generator = params_generator_type_dict[op_info.name]() + for (shape, memory_format), kwargs in params_generator.gen_input_params(): + arg = make_arg(shape).to(memory_format=memory_format).requires_grad_(requires_grad) + yield SampleInput(arg, kwargs=kwargs) + +def max_pool2d_backward(*args, kernel_size=(), stride=(), padding=(0,), dilation=(1,), ceil_mode=False, **kwargs): + out, indices = torch.nn.functional.max_pool2d_with_indices( + *args, kernel_size=kernel_size, stride=stride, padding=padding, dilation=dilation, ceil_mode=ceil_mode, return_indices=True) + grad_out = torch.ones_like(out) + if stride is None: + stride = kernel_size + out_b = torch.ops.aten.max_pool2d_with_indices_backward.default( + grad_out, *args, kernel_size, stride, padding, dilation, ceil_mode, indices) + return out_b + +def error_inputs_max_pool1d(op_info, device, **kwargs): + # Toggle requires_grad because `max_pool1d` has different path + # based on whether `requires_grad` is set or not. + for requires_grad in (True, False): + make_arg = partial(make_tensor, device=device, dtype=torch.float, requires_grad=requires_grad) + # error inputs when pad is negative + x = make_arg((0, 1, 49)) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': -1, 'return_indices': True}), + error_regex='pad must be non-negative') + + # error inputs when pad > kernel_size / 2 + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': 4, 'return_indices': True}), + error_regex='pad should be at most half of effective kernel size') + + # error inputs when pad > ((kernel_size - 1) * dilation + 1) / 2, when dilation is not default + yield ErrorInput(SampleInput(x, + kwargs={'kernel_size': 3, 'dilation': 2, 'stride': 1, 'padding': 3, 'return_indices': True}), + error_regex='pad should be at most half of effective kernel size') + + # error inputs for input tensor + error_msg = r'Expected 2D or 3D \(batch mode\) tensor with optional 0 dim batch size for input' + yield ErrorInput(SampleInput(make_arg((), requires_grad=requires_grad), kwargs={'kernel_size': 1}), + error_regex=error_msg) + + # error inputs for empty input + yield ErrorInput(SampleInput(torch.tensor([], device=device, requires_grad=requires_grad), + kwargs={'kernel_size': 1}), + error_regex=error_msg) + + # error: unbatched input with 0 sized non-batch dims. + yield ErrorInput(SampleInput(make_arg((0, 10), requires_grad=requires_grad), + kwargs={'kernel_size': 1}), + error_regex=error_msg) + + # error: batched input with 0 sized non-batch dims. + yield ErrorInput(SampleInput(make_arg((1, 10, 0), requires_grad=requires_grad), + kwargs={'kernel_size': 1}), + error_regex=error_msg) + + # error inputs for empty input with stride=0 + error_msg = 'stride must be greater than zero, but got 0' + yield ErrorInput(SampleInput(make_arg((3, 3, 3)), kwargs={'kernel_size': 1, 'stride': 0}), + error_regex=error_msg) + + # error inputs for empty input with dilation=0 + error_msg = 'dilation must be greater than zero, but got 0' + yield ErrorInput(SampleInput(make_arg((3, 3, 3)), + kwargs={'kernel_size': 1, 'stride': 1, 'padding': 0, 'dilation': 0}), + error_regex=error_msg) + + # error inputs for invalid output size + error_msg = 'Invalid computed output size: -2' + yield ErrorInput(SampleInput(make_arg((2, 2, 2)), + kwargs={'kernel_size': 5, 'stride': 1, 'padding': 0, 'dilation': 1}), + error_regex=error_msg) + + # error inputs when kernel_size=0 + error_msg = 'kernel_size must be greater than zero' + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 0}), + error_regex=error_msg) + + # error inputs for strides > 0 + error_msg = 'stride must be greater than zero' + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 0}), + error_regex=error_msg) + + +def error_inputs_max_pool2d(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float, requires_grad=False) + # error inputs when pad is negative + x = make_arg((0, 1, 49)) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': -1, 'return_indices': True}), + error_regex='pad must be non-negative') + # 2-dimensional kernel + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (3, 2), 'stride': 50, 'padding': -1, 'return_indices': True}), + error_regex='pad must be non-negative') + + # error inputs when pad > kernel_size / 2 (kernel_size : int) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': 4, 'return_indices': True}), + error_regex='pad should be at most half of effective kernel size') + + # error inputs when pad > kernel_size / 2 (kernel_size : tuple) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (3, 2), 'stride': 50, 'padding': 4, 'return_indices': True}), + error_regex='pad should be at most half of effective kernel size') + + # error: unbatched input with 0 sized non-batch dims. + err_msg = r'Expected 3D or 4D \(batch mode\) tensor with optional 0 dim batch size for input' + yield ErrorInput(SampleInput(make_arg((1, 0, 10)), + kwargs={'kernel_size': 1}), + error_regex=err_msg) + + # error: batched input with 0 sized non-batch dims. + yield ErrorInput(SampleInput(make_arg((2, 1, 10, 0)), + kwargs={'kernel_size': 1}), + error_regex=err_msg) + + # error: inputs when kernel size too large for input + yield ErrorInput(SampleInput(make_arg((1, 1, 4)), + kwargs={'kernel_size': 2}), + error_regex='Output size is too small') + + +def error_inputs_max_pool3d(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float, requires_grad=False) + # error inputs when pad is negative + x = make_arg((0, 1, 49, 50)) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': -1, 'return_indices': True}), + error_regex='pad must be non-negative') + # 3-dimensional kernel + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (3, 2, 2), 'stride': 50, + 'padding': -1, 'return_indices': True}), + error_regex='pad must be non-negative') + + # error inputs when pad > kernel_size / 2 (kernel_size: int) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': 4, 'return_indices': True}), + error_regex='pad should be at most half of effective kernel size') + + # error inputs when pad > kernel_size / 2 (kernel_size: tuple) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (3, 2, 2), 'stride': 50, + 'padding': 4, 'return_indices': True}), + error_regex='pad should be at most half of effective kernel size') + + # error: unbatched input with 0 sized non-batch dims. + err_msg = r'Expected input\'s non-batch dimensions to have positive length' + yield ErrorInput(SampleInput(make_arg((0, 1, 2, 10)), + kwargs={'kernel_size': 1}), + error_regex=err_msg) + + # error: batched inputs with 0 sized non-batch dims. + yield ErrorInput(SampleInput(make_arg((2, 1, 0, 1, 2)), + kwargs={'kernel_size': 1}), + error_regex=err_msg) + + # error: inputs when kernel size too large for input + yield ErrorInput(SampleInput(make_arg((1, 1, 1, 4, 4)), + kwargs={'kernel_size': 2}), + error_regex='Output size is too small') + + + +def sample_inputs_normalize(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, low=-1, high=1, device=device, dtype=dtype, requires_grad=requires_grad) + + cases: tuple[tuple[int, ...], dict] = ( + ((2, 1, 4, 5), {'p': 1., 'dim': 2}), + ((2, 3, 4, 5), {'p': 2., 'dim': 1}), + ((1, 2, 4, 5), {'p': 0.5, 'dim': 0}), + ((1, 3, 4, 5), {'p': -1., 'dim': 1}), + ((1, 3, 4, 5), {'p': 0., 'dim': -1}), + ((), {'p': 1.2, 'dim': 0}), + ((2, 3, 4, 5), {}), + ((2, 3, 4, 5), {'eps': 1e-4})) + + for input_shape, kwargs in cases: + yield SampleInput(make_arg(input_shape), kwargs=kwargs) + + +def complex_conv(fn, input_size, weight, grad_output, stride, padding, dilation, groups): + # conv(W, x, b) = conv(Wr, xr, br) - conv(Wi, xi, 0) + i(conv(Wi, xr, bi) + conv(Wr, xi, 0)) + # a = conv(Wr, xr, br), + # b = conv(Wi, xi, 0), + # c = conv(Wr + Wi, xr + xi, br + bi) + # conv(W, x, b) = a - b + i(c - a - b) + + grad_output_ = torch.view_as_real(grad_output) + grad_output_r = grad_output_[..., 0] + grad_output_i = grad_output_[..., 1] + + weight_ = torch.view_as_real(weight) + weight_r = weight_[..., 0] + weight_i = weight_[..., 1] + + a = fn(input_size, weight_r, grad_output_r, stride, padding, dilation, groups) + b = fn(input_size, weight_i, grad_output_i, stride, padding, dilation, groups) + c = fn(input_size, weight_r + weight_i, grad_output_r + grad_output_i, stride, padding, dilation, groups) + + return (a - b) + 1j * (c - a - b) + + +def conv_transpose_ref(input, weight, bias, stride=1, padding=0, + output_padding=0, dilation=1, groups=1, + fn=None): + # Derivative of `conv` is `conv_transpose`. + # To verify the correctness of `conv_transpose`, + # we rely `torch.nn.grad` implementation (which is tested in test_nn.py) + # for floating dtypes. + + if fn is None: + raise AssertionError("Expected fn to not be None") + + grad_fn_map = {torch.nn.functional.conv_transpose1d: torch.nn.grad.conv1d_input, + torch.nn.functional.conv_transpose2d: torch.nn.grad.conv2d_input, + torch.nn.functional.conv_transpose3d: torch.nn.grad.conv3d_input} + batched_dim_map = {torch.nn.functional.conv_transpose1d: 3, + torch.nn.functional.conv_transpose2d: 4, + torch.nn.functional.conv_transpose3d: 5} + + # Input for `ref` is ndarray. + input, weight = torch.from_numpy(input), torch.from_numpy(weight) + + is_batched = len(input.shape) == batched_dim_map[fn] + if not is_batched: + input = input.unsqueeze(0) + + if bias is not None: + bias = torch.from_numpy(bias) + unsqueeze_dims = input.ndim - 2 + for _ in range(unsqueeze_dims): + bias = bias.unsqueeze(1) + + grad_output = input + # Get the input shape for grad_fn. + conv_transpose_output = fn(grad_output.to('meta'), weight.to('meta'), None, + stride=stride, padding=padding, output_padding=output_padding, + groups=groups, dilation=dilation) + input_size = conv_transpose_output.shape + + grad_fn = grad_fn_map[fn] + if weight.dtype.is_complex: + out = complex_conv(grad_fn, input_size, weight, grad_output, stride, padding, dilation, groups) + else: # Floating + out = grad_fn(input_size, weight, grad_output, stride, padding, dilation, groups) + + if bias is not None: + out = out + bias + + return out.squeeze(0) if not is_batched else out + + +def sample_inputs_conv_transpose1d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as shapes for input, weight, bias + # and a dict of values of (stride, padding, output_padding, groups, dilation) + cases: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], dict] = ( + ((1, 3, 4), (3, 3, 3), (3,), + {'stride': (2,), 'padding': 2, 'output_padding': (1,), 'groups': 1}), + ((2, 2, 4), (2, 2, 4), (4,), + {'stride': (3,), 'padding': (1,), 'output_padding': (2,), 'groups': 2, 'dilation': (4,)}), + ((1, 1, 4), (1, 1, 4), (1,), + {'stride': 2, 'padding': 1, 'output_padding': 1, 'groups': 1, 'dilation': (2,)}), + ((1, 1, 4), (1, 2, 3), None, + {'stride': 2, 'padding': 1, 'output_padding': 1, 'groups': 1}), + ((1, 4, 5), (4, 8, 3), None, + {}) + ) + + for input_shape, weight, bias, kwargs in cases: + # Batched + yield SampleInput(make_arg(input_shape), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + # Unbatched + yield SampleInput(make_arg(input_shape[1:]), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + + +def sample_inputs_conv_transpose2d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as shapes for input, weight, bias + # and a dict of values of (stride, padding, output_padding, groups, dilation) + cases: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], dict] = ( + ((1, 3, 4, 4), (3, 3, 3, 3), (3,), + {'stride': (2, 2), 'padding': 2, 'output_padding': (1, 1), 'groups': 1}), + ((2, 2, 4, 4), (2, 2, 4, 5), (4,), + {'stride': (3, 2), 'padding': (1, 2), 'output_padding': (2, 3), 'groups': 2, 'dilation': (4, 4)}), + ((1, 1, 4, 5), (1, 1, 4, 3), (1,), + {'stride': 2, 'padding': 1, 'output_padding': 1, 'groups': 1, 'dilation': (2, 3)}), + ((1, 1, 4, 3), (1, 2, 3, 4), None, + {'stride': 2, 'padding': 1, 'output_padding': 1, 'groups': 1}), + ((2, 4, 4, 4), (4, 1, 3, 3), None, {'groups': 4}), + ((1, 2, 5, 5), (2, 4, 3, 3), None, {}) + ) + + for input_shape, weight, bias, kwargs in cases: + # Batched + yield SampleInput(make_arg(input_shape), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + # Unbatched + yield SampleInput(make_arg(input_shape[1:]), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + +def sample_inputs_conv_transpose3d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as shapes for input, weight, bias + # and a dict of values of (stride, padding, output_padding, groups, dilation) + cases: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], dict] = ( + ((1, 3, 4, 4, 4), (3, 3, 3, 3, 3), (3,), + {'stride': (2, 2, 2), 'padding': 2, 'output_padding': (1, 1, 1), 'groups': 1}), + ((2, 2, 4, 4, 4), (2, 2, 4, 5, 6), (4,), + {'stride': (3, 2, 1), 'padding': (1, 2, 3), 'output_padding': (2, 3, 1), 'groups': 2, 'dilation': (4, 4, 4)}), + ((1, 1, 4, 5, 2), (1, 1, 4, 3, 1), (1,), + {'stride': 2, 'padding': 1, 'output_padding': 1, 'groups': 1, 'dilation': (2, 3, 2)}), + ((1, 1, 4, 3, 4), (1, 2, 3, 4, 5), None, + {'stride': 2, 'padding': 1, 'output_padding': 1, 'groups': 1}), + ((1, 4, 5, 5, 5), (4, 8, 3, 3, 3), None, + {}) + ) + + for input_shape, weight, bias, kwargs in cases: + # Batched + yield SampleInput(make_arg(input_shape), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + # Unbatched + yield SampleInput(make_arg(input_shape[1:]), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + + +def sample_inputs_conv1d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as shapes for input, weight, bias, + # and a dict of values of (stride, padding, dilation, groups) + cases: tuple = ( + ((1, 3, 4), (3, 3, 3), (3,), {'stride': (2,), 'padding': 2, 'groups': 1}), + ((2, 4, 8), (2, 2, 3), (2,), {'stride': 3, 'padding': 1, 'groups': 2, 'dilation': 2}), + ((1, 4, 5), (1, 4, 3), None, {'stride': (2,), 'padding': 'valid'}), + ((2, 2, 4), (2, 1, 4), (2,), {'stride': (1,), 'padding': 'same', 'groups': 2, 'dilation': (2,)}), + # With defaults + ((1, 4, 5), (3, 4, 3), None, {}), + ) + + for input_shape, weight, bias, kwargs in cases: + # Batched + yield SampleInput(make_arg(input_shape), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + # Unbatched + yield SampleInput(make_arg(input_shape[1:]), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + + +def error_inputs_conv1d(opinfo, device, **kwargs): + dtype = highest_precision_float(device) + cdtype = highest_precision_complex(device) + make_arg = partial(make_tensor, device=device, dtype=dtype) + make_int_arg = partial(make_tensor, device=device, dtype=torch.int64) + make_complex_arg = partial(make_tensor, device=device, dtype=cdtype) + + # error inputs for different dtypes of input tensor and bias + yield ErrorInput( + SampleInput(make_int_arg((1, 1, 4)), args=(make_int_arg((1, 1, 2)), make_arg((1,)))), + error_regex="should be the same") + + # error inputs for different dtypes of input tensor and bias + yield ErrorInput( + SampleInput(make_arg((1, 1, 4)), args=(make_arg((1, 1, 2)), make_complex_arg((1,)))), + error_regex="should be the same") + + # error inputs for negative strides + yield ErrorInput( + SampleInput(make_arg((1, 1, 4)), args=(make_arg((1, 2, 2)), make_arg((1,))), + kwargs={'stride': (-1,)}), error_regex="non-positive stride is not supported") + + # error inputs for negative padding + yield ErrorInput( + SampleInput(make_arg((1, 1, 4)), args=(make_arg((1, 2, 2)), make_arg((1,))), + kwargs={'padding': (-1,)}), error_regex="negative padding is not supported") + + # error inputs for negative dilation + yield ErrorInput( + SampleInput(make_arg((1, 1, 4)), args=(make_arg((1, 1, 2)), make_arg((1,))), + kwargs={'dilation': (-1,)}), error_regex="dilation should be greater than zero") + + # FIXME: https://github.com/pytorch/pytorch/issues/85656 + # error inputs for bias shape not equal to the output channels + # yield ErrorInput(SampleInput(make_arg((1, 1, 4)), args=(make_arg((1, 1, 3)), make_arg((2,)))), + # error_regex="expected bias to be 1-dimensional with 1 elements") + + # error inputs for input.ndim != weight.ndim + yield ErrorInput(SampleInput(make_arg((1, 1, 4)), args=(make_arg((1, 2)), make_arg((1,)))), + error_regex="weight should have at least three dimensions") + + # error inputs for the weight[0] are less than the number of groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 4)), args=(make_arg((2, 2, 2)), make_arg((2,))), + kwargs={'padding': 'same', 'groups': 3}), error_regex="expected weight to be at least 3 at dimension 0") + + # error inputs for the weight[0] are less than the number of groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 4)), args=(make_arg((2, 2, 2)), make_arg((2,))), + kwargs={'groups': 3}), error_regex="expected weight to be at least 3 at dimension 0") + + # error inputs for invalid groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 4)), args=(make_arg((2, 2, 2)), make_arg((2,))), + kwargs={'padding': 'same', 'groups': -1}), error_regex="non-positive groups is not supported") + + # error inputs for invalid groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 4)), args=(make_arg((2, 2, 2)), make_arg((2,))), + kwargs={'padding': 'same', 'groups': 0}), error_regex="non-positive groups is not supported") + + +def error_inputs_conv2d(opinfo, device, **kwargs): + dtype = highest_precision_float(device) + cdtype = highest_precision_complex(device) + make_arg = partial(make_tensor, device=device, dtype=dtype) + make_int_arg = partial(make_tensor, device=device, dtype=torch.int64) + make_complex_arg = partial(make_tensor, device=device, dtype=cdtype) + + # error inputs for different dtypes of input tensor and bias + yield ErrorInput( + SampleInput(make_int_arg((2, 4, 4)), args=(make_int_arg((3, 2, 3, 3)), make_arg((3,)))), + error_regex="should be the same") + + # error inputs for different dtypes of input tensor and bias + yield ErrorInput( + SampleInput(make_arg((2, 4, 4)), args=(make_arg((3, 2, 3, 3)), make_complex_arg((3,)))), + error_regex="should be the same") + + # error inputs for negative strides + yield ErrorInput( + SampleInput(make_arg((1, 1, 4, 4)), args=(make_arg((1, 2, 2, 3)), make_arg((1,))), + kwargs={'stride': (-1,)}), error_regex="non-positive stride is not supported") + + # error inputs for negative padding + yield ErrorInput( + SampleInput(make_arg((1, 1, 4, 3)), args=(make_arg((1, 2, 2, 4)), make_arg((1,))), + kwargs={'padding': (-1,)}), error_regex="negative padding is not supported") + + # error inputs for negative dilation + yield ErrorInput( + SampleInput(make_arg((1, 1, 4, 2)), args=(make_arg((1, 1, 2, 5)), make_arg((1,))), + kwargs={'dilation': (-1,)}), error_regex="dilation should be greater than zero") + + # FIXME: https://github.com/pytorch/pytorch/issues/85656 + # error inputs for bias shape not equal to the output channels + # yield ErrorInput(SampleInput(make_arg((1, 1, 4, 4)), args=(make_arg((1, 1, 3, 2)), make_arg((2,)))), + # error_regex="expected bias to be 1-dimensional with 1 elements") + + # error inputs for input.ndim != weight.ndim + yield ErrorInput( + SampleInput(make_arg((1, 1, 4, 3)), args=(make_arg((1, 2, 2)), make_arg((1,))), + kwargs={'padding': 'same'}), error_regex="Expected 3-dimensional input for 3-dimensional weight") + + # error inputs for the weight[0] are less than the number of groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 4, 3)), args=(make_arg((2, 2, 1, 3)), make_arg((2,))), + kwargs={'groups': 3}), error_regex="expected weight to be at least 3 at dimension 0") + + # error inputs for groups the weight[0] are less than the number of groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 4, 3)), args=(make_arg((2, 2, 1, 3)), make_arg((2,))), + kwargs={'padding': 'same', 'groups': 3}), error_regex="expected weight to be at least 3 at dimension 0") + + # error inputs for invalid groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 4, 5)), args=(make_arg((2, 2, 1, 4)), make_arg((2,))), + kwargs={'padding': 'same', 'groups': -1}), error_regex="non-positive groups is not supported") + + # error inputs for invalid groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 4, 3)), args=(make_arg((2, 2, 4, 3)), make_arg((2,))), + kwargs={'padding': 'same', 'groups': 0}), error_regex="non-positive groups is not supported") + + +def sample_inputs_conv2d(op_info, device, dtype, requires_grad, jit_fail_sample=False, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as shapes for input, weight, bias + # and a dict of values of (stride, padding, groups, dilation) + cases: tuple = ( + ((1, 3, 4, 4), (3, 3, 3, 3), (3,), + {'stride': (2, 2), 'padding': 2, 'groups': 1}), + ((2, 4, 8, 8), (2, 2, 3, 3), (2,), + {'stride': (3, 2), 'padding': (2, 1), 'groups': 2, 'dilation': (4, 4)}), + ((1, 4, 5, 5), (1, 4, 2, 3), (1,), + {'stride': 2, 'padding': 1, 'groups': 1, 'dilation': (2, 3)}), + ((1, 4, 5, 5), (1, 4, 2, 3), (1,), + {'stride': 2, 'padding': 1, 'groups': 1, 'dilation': (2, 3)}), + ((1, 2, 4, 3), (4, 2, 3, 4), None, + {'stride': 2, 'padding': 1, 'groups': 1}), + ((1, 4, 5, 5), (1, 4, 2, 3), (1,), + {'stride': 2, 'padding': "valid"}), + ((1, 4, 5, 5), (1, 4, 2, 3), (1,), + {'stride': 1, 'padding': "same", 'dilation': 3}), + # Below are the group related samples from common_nn.py + ((2, 4, 6, 6), (4, 1, 3, 3), (4,), {'groups': 4}), + ((2, 4, 6, 6), (8, 1, 3, 3), (8,), {'groups': 4}), + ((2, 4, 6, 6), (8, 1, 3, 3), None, {'groups': 4}), + ((2, 4, 6, 6), (4, 1, 3, 3), (4,), {'groups': 4, 'stride': (3, 2)}), + ((2, 4, 6, 6), (4, 1, 3, 3), (4,), {'groups': 4, 'padding': (1, 1)}), + ((2, 4, 5, 5), (4, 1, 2, 2), (4,), {'groups': 4, 'dilation': (2, 2)}), + ((2, 4, 6, 5), (6, 2, 3, 2), (6,), {'groups': 2}), + # With defaults + ((1, 4, 5, 5), (3, 4, 3, 3), None, {}), + ) + + for input_shape, weight, bias, kwargs in cases: + # Batched + yield SampleInput(make_arg(input_shape), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + # Unbatched + yield SampleInput(make_arg(input_shape[1:]), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + + +def sample_inputs_conv3d(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as shapes for input, weight, bias + # and dict of values of (stride, padding, dilation, groups) + cases: tuple = ( + ((1, 1, 4, 4, 4), (1, 1, 1, 1, 1), (1,), {'padding': 'same'}), + ((1, 1, 4, 4, 4), (1, 1, 4, 4, 4), (1,), {'stride': (2, 2, 2)}), + ((1, 1, 5, 5, 5), (1, 1, 3, 3, 3), (1,), {'dilation': 2}), + ((1, 1, 1, 1, 10), (1, 1, 1, 1, 4), None, {'padding': 'valid'}), + ((1, 1, 10, 11, 12), (1, 1, 1, 2, 5), None, {'padding': 'same'}), + ((1, 1, 10, 11, 12), (1, 1, 1, 2, 5), None, {'padding': 'same', 'dilation': 2}), + ((1, 1, 10, 11, 12), (1, 1, 4, 4, 4), None, {'padding': 'same', 'dilation': 3}), + ((1, 1, 1, 1, 10), (1, 1, 1, 1, 4), None, {'padding': 'valid'}), + ((3, 9, 3, 1, 9), (3, 3, 3, 1, 9), (3,), {'groups': 3}), + ((3, 9, 3, 1, 9), (3, 3, 3, 1, 9), (3,), {'stride': (2, 2, 2), 'dilation': 1, 'groups': 3}), + ) + + for input_shape, weight, bias, kwargs in cases: + # Batched + yield SampleInput(make_arg(input_shape), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + # Unbatched + yield SampleInput(make_arg(input_shape[1:]), args=( + make_arg(weight), + make_arg(bias) if bias is not None else bias + ), kwargs=kwargs) + + +def error_inputs_conv3d(opinfo, device, **kwargs): + dtype = highest_precision_float(device) + cdtype = highest_precision_complex(device) + make_arg = partial(make_tensor, device=device, dtype=dtype) + make_int_arg = partial(make_tensor, device=device, dtype=torch.int64) + make_complex_arg = partial(make_tensor, device=device, dtype=cdtype) + + # error inputs for different dtypes of input tensor and bias + yield ErrorInput( + SampleInput(make_int_arg((1, 1, 4, 4, 4)), args=(make_int_arg((1, 1, 2, 2, 2)), make_arg((1,)))), + error_regex="should be the same") + + # error inputs for different dtypes of input tensor and bias + yield ErrorInput( + SampleInput(make_arg((1, 1, 4, 4, 4)), args=(make_arg((1, 1, 2, 2, 2)), make_complex_arg((1,)))), + error_regex="should be the same") + + # error inputs for negative strides + yield ErrorInput( + SampleInput(make_arg((1, 1, 4, 4, 4)), args=(make_arg((1, 1, 2, 2, 2)), make_arg((1,))), + kwargs={'stride': (-1,)}), error_regex="non-positive stride is not supported") + + # error inputs for negative padding + yield ErrorInput( + SampleInput(make_arg((1, 1, 4, 4, 4)), args=(make_arg((1, 1, 2, 2, 2)), make_arg((1,))), + kwargs={'padding': (-1,)}), error_regex="negative padding is not supported") + + # error inputs for negative dilation + yield ErrorInput( + SampleInput(make_arg((1, 1, 4, 4, 4)), args=(make_arg((1, 1, 2, 2, 2)), make_arg((1,))), + kwargs={'dilation': (-1,)}), error_regex="dilation should be greater than zero") + + # FIXME: https://github.com/pytorch/pytorch/issues/85656 + # error inputs for bias shape not equal to the output channels + # yield ErrorInput(SampleInput(make_arg((1, 1, 4, 4, 4)), args=(make_arg((1, 1, 3, 3, 3)), make_arg((2,)))), + # error_regex="expected bias to be 1-dimensional with 1 elements") + + # error inputs for input.ndim != weight.ndim + yield ErrorInput( + SampleInput(make_arg((1, 1, 3, 4, 5)), args=(make_arg((1, 1, 4, 3)), make_arg((1,))), + kwargs={'padding': 'same'}), error_regex="Expected 4-dimensional input for 4-dimensional weight") + + # error inputs for the weight[0] are less than the number of groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 3, 4, 5)), args=(make_arg((2, 2, 4, 3, 3)), + make_arg((2,))), kwargs={'groups': 3}), + error_regex="expected weight to be at least 3 at dimension 0") + + # error inputs for the weight[0] are less than the number of groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 3, 4, 5)), args=(make_arg((2, 2, 4, 3, 3)), + make_arg((2,))), kwargs={'padding': 'same', 'groups': 3}), + error_regex="expected weight to be at least 3 at dimension 0") + + # error inputs for invalid groups + yield ErrorInput( + SampleInput(make_arg((2, 2, 3, 4, 5)), args=(make_arg((2, 2, 4, 3, 3)), + make_arg((2,))), kwargs={'padding': 'same', 'groups': 0}), + error_regex="non-positive groups is not supported") + + # error inputs for padding='same' not supported by strided convolutions + yield ErrorInput( + SampleInput(make_arg((18, 27, 9, 1, 9)), args=(make_arg((9, 9, 9, 1, 9)), + make_arg((9,))), kwargs={'stride': 2, 'padding': 'same', 'groups': 3}), + error_regex="padding='same' is not supported for strided convolutions") + + +def sample_inputs_group_norm(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as input shape, num groups, and kwargs for eps + cases: tuple[tuple[int, ...], int, float] = ( + ((1, 6, 3), 2, {'eps' : 0.5}), + ((2, 6, 3), 2, {'eps' : -0.5}), + ((1, 3), 1, {'eps' : 1e-5}), + ((0, 2), 1, {'eps' : 1e-5}), + ((S, S, S), 1, {'eps' : 0.5}), + ) + + # num_channels is inferred to be input.shape[1] dimension + for input_shape, num_groups, kwargs in cases: + # Shape of weight and bias should be the same as num_channels + channels = input_shape[1] if len(input_shape) > 1 else 0 + weight_tensor = make_arg(channels) + bias_tensor = make_arg(channels) + + # Checking for permutations of weights and biases as `None` + weights = [weight_tensor, None] + biases = [bias_tensor, None] + for weight, bias in itertools.product(weights, biases): + kwargs = { + 'weight': weight, + 'bias': bias, + **kwargs + } + yield SampleInput(make_arg(input_shape), num_groups, **kwargs) + + # Without any optional args + yield SampleInput(make_arg((1, 2)), args=(1,)) + +def reference_inputs_group_norm(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_group_norm( + op_info, device, dtype, requires_grad, **kwargs) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as input shape, num groups, and kwargs for eps + cases: tuple[tuple[int, ...], int, float] = ( + ((20, 6, 10, 10), 3, {'eps' : 1e-5}), + # equivalent with InstanceNorm + # GroupNorm(C, num_groups=C) == InstanceNorm(num_features=C) + ((20, 6, 10, 10), 6, {'eps' : 1e-5}), + # equivalent with LayerNorm + # GroupNorm(C, num_groups=1, affine=False) == LayerNorm(normalized_shape=[C, H, W], elementwise_affine=False) + ((20, 6, 10, 10), 1, {'eps' : 1e-5}), + ) + + # num_channels is inferred to be input.shape[1] dimension + for input_shape, num_groups, kwargs in cases: + # Shape of weight and bias should be the same as num_channels + channels = input_shape[1] if len(input_shape) > 1 else 0 + input_tensor = make_arg(input_shape) + weight_tensor = make_arg(channels) + bias_tensor = make_arg(channels) + + # Checking for permutations of weights and biases as `None` + weights = [weight_tensor, None] + biases = [bias_tensor, None] + for weight, bias in itertools.product(weights, biases): + kwargs = { + 'weight': weight, + 'bias': bias, + **kwargs + } + yield SampleInput(input_tensor, num_groups, **kwargs) + + +def sample_inputs_instance_norm(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_arg_without_requires_grad = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + # Ordered as: input shape, kwargs for momentum, eps + cases: tuple[tuple[int, ...], dict] = ( + ((S, S, S), {'momentum': 0.5, 'eps': 0.6}), + ((S, S, S), {'momentum': 0.5, 'eps': 0.6, 'use_input_stats': True}), + ((3, 2, 4), {'momentum': -1.2}), + ((3, 2, 4), {'momentum': 0.0}), + ((3, 2, 3, 4), {'momentum': -1.0, 'eps': 0.5}), + ((3, 2, 3, 4), {'momentum': -1.0, 'eps': 0.5}), + ) + + for input_shape, kwargs in cases: + # args: running mean, running var, weight and bias should necessarily be of shape: (channels,) + channels = input_shape[1] + weight = make_arg(channels) + bias = make_arg(channels) + running_mean = make_arg_without_requires_grad(channels, low=0) + running_var = make_arg_without_requires_grad(channels, low=0) + new_kwargs = { + 'running_mean': running_mean, + 'running_var': running_var, + 'weight': weight, + 'bias': bias, + **kwargs + } + + yield SampleInput( + make_arg(input_shape), + args=(), + kwargs=new_kwargs + ) + + # Checking for permutations of weights and biases as `None` + # instance_norm assumes that if there's a bias, there's a weight + weights = [channels, None] + biases = [None, None] + + for weight_channels, bias_channels in zip(weights, biases, strict=True): + running_mean = make_arg_without_requires_grad(channels, low=0) + running_var = make_arg_without_requires_grad(channels, low=0) + yield SampleInput( + make_arg(input_shape), + args=(), + kwargs={ + 'running_mean': running_mean, + 'running_var': running_var, + 'weight': make_arg(weight_channels) if weight_channels is not None else None, + 'bias': make_arg(bias_channels) if bias_channels is not None else None + } + ) + + # Test case for no optional kwargs + yield SampleInput(make_arg((1, 2, 3)), kwargs={}) + +def sample_inputs_safe_softmax(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + def make_bool_mask(*shape): + return torch.randint(0, 2, shape, device=device, dtype=torch.bool) + + def mask_two_rows(rows, cols): + mask_two_rows = torch.ones((rows, cols), dtype=torch.bool, device=device) + mask_two_rows[rows - 1] = False + mask_two_rows[rows - 3] = False + return mask_two_rows + + def convert_to_float_mask(mask: torch.Tensor) -> torch.Tensor: + return torch.where(~mask, float('-inf'), 0.0) + + def with_requires_grad(tensor): + return tensor.requires_grad_(requires_grad) + + def generate_input_from_mask(mask_shape, dim): + mask = make_bool_mask(*mask_shape) + input_tensor = make_arg(mask_shape) + masked_input = input_tensor + convert_to_float_mask(mask) + return SampleInput(with_requires_grad(masked_input), kwargs={'dim': dim}) + + samples = [ + # Basic 3D tensor with mask + generate_input_from_mask((2, 3, 4), dim=1), + # 2D tensor with mask, testing different dim + generate_input_from_mask((5, 5), dim=0), + # 4D tensor, testing with a different dim + generate_input_from_mask((2, 3, 4, 5), dim=2), + # Edge case: 1D tensor + generate_input_from_mask((10,), dim=0), + # Edge case: tensor with one dimension of size 1 + generate_input_from_mask((1, 5, 5), dim=1), + # Testing with all elements masked + SampleInput( + with_requires_grad( + make_arg((3, 3)) + + convert_to_float_mask( + torch.zeros((3, 3), dtype=torch.bool, device=device) + ) + ), + kwargs={"dim": 1}, + ), + # Testing with no elements masked + SampleInput( + with_requires_grad( + make_arg((3, 3)) + + convert_to_float_mask( + torch.ones((3, 3), dtype=torch.bool, device=device) + ) + ), + kwargs={"dim": 1}, + ), + # Testing with two rows masked + SampleInput( + with_requires_grad( + make_arg((6, 3)) + convert_to_float_mask(mask_two_rows(6, 3)) + ), + kwargs={"dim": 1}, + ), + ] + yield from samples + +def sample_inputs_layer_norm(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as input shape, normalized_shape and a kwarg dict for eps + cases: tuple[tuple[int, ...], tuple[int, ...], dict] = ( + ((1, 2, 3), (1, 2, 3), {'eps': 0.5}), + ((2, 2, 3), (2, 3), {'eps': -0.5}), + ((1,), (1,), {}), + ((1, 2), (2,), {}), + ((0, 1), (1,), {}), + ) + + for input_shape, normalized_shape, kwargs in cases: + # Shape of weight and bias should be the same as normalized_shape + weight = make_arg(normalized_shape) + bias = make_arg(normalized_shape) + yield SampleInput( + make_arg(input_shape), + args=(normalized_shape, weight, bias), + kwargs=kwargs + ) + # Without any optional args + yield SampleInput(make_arg((1, 2)), args=((2,),)) + + # TODO: @krshrimali, once to_numpy method in SampleInput class is modified to take None inputs, + # enable these inputs; see https://github.com/pytorch/pytorch/pull/63276#discussion_r691950400 + + # With weight and a `None` bias + # yield SampleInput(make_arg((1, 2)), args=((2,), make_arg((2,)), None)) + + # With `None` weight and bias (tests failing for this, see the link above) + # yield SampleInput(make_arg((1, 2)), args=((2,), None, make_arg((2,)))) + + +def sample_inputs_native_layer_norm(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as input shape, normalized_shape, eps + cases: tuple[tuple[int, ...], tuple[int, ...], float] = ( + ((1, 2, 3), (1, 2, 3), 0.5), + ((2, 2, 3), (2, 3), -0.5), + ((1,), (1,), 1e-5), + ((1, 2), (2,), 1e-5), + ((0, 1), (1,), 1e-5), + ) + + for input_shape, normalized_shape, eps in cases: + # Shape of weight and bias should be the same as normalized_shape + weight = make_arg(normalized_shape) + bias = make_arg(normalized_shape) + yield SampleInput( + make_arg(input_shape), + args=(normalized_shape, weight, bias, eps), + ) + yield SampleInput( + make_arg(input_shape), + args=(normalized_shape, None, bias, eps), + ) + yield SampleInput( + make_arg(input_shape), + args=(normalized_shape, weight, None, eps), + ) + yield SampleInput( + make_arg(input_shape), + args=(normalized_shape, None, None, eps), + ) + +def sample_inputs_rms_norm(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, high=1000) + + # Ordered as input shape, normalized_shape and a kwarg dict for eps + cases: tuple[tuple[int, ...], tuple[int, ...], dict] = ( + ((1, 2, 3), (1, 2, 3), {'eps': 0.5}), + ((2, 2, 3), (2, 3), {'eps': -0.5}), + ((1,), (1,), {}), + ((1, 2), (2,), {}), + ((0, 1), (1,), {}), + ) + + for input_shape, normalized_shape, kwargs in cases: + # Shape of weight and bias should be the same as normalized_shape + weight = make_arg(normalized_shape) + yield SampleInput( + make_arg(input_shape), + args=(normalized_shape, weight), + kwargs=kwargs + ) + # Without any optional args + yield SampleInput(make_arg((1, 2)), args=((2,),)) + +def error_inputs_group_norm(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32, requires_grad=False) + + # check that input has minimum number of dimensions + err_msg1 = "Expected at least 2 dimensions for input tensor but received" + s1 = SampleInput(make_arg(1), args=(1,)) + yield ErrorInput(s1, error_regex=err_msg1) + + # check that the channels dimension is compatible with number of groups + err_msg2 = "Expected number of channels in input to be divisible by num_groups, but got input of shape" + s2 = SampleInput(make_arg((2, 7, 4)), args=(2,)) + yield ErrorInput(s2, error_regex=err_msg2) + +def error_inputs_native_layer_norm(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32, requires_grad=False) + input_shape = (1, 2, 3) + + err_msg1 = "Expected normalized_shape to be at least 1-dimensional" + s1 = SampleInput( + make_arg(input_shape), args=((), None, None, 1e-5) + ) + yield ErrorInput(s1, error_regex=err_msg1) + + normalized_shape = (1, 2, 3) + weight = make_arg((1, 2)) + err_msg2 = "Expected weight to be of same shape as normalized_shape" + s2 = SampleInput( + make_arg(input_shape), args=(normalized_shape, weight, None, 1e-5) + ) + yield ErrorInput(s2, error_regex=err_msg2) + + bias = make_arg((1, 2)) + err_msg3 = "Expected bias to be of same shape as normalized_shape" + s3 = SampleInput( + make_arg(input_shape), args=(normalized_shape, None, bias, 1e-5) + ) + yield ErrorInput(s3, error_regex=err_msg3) + + err_msg4 = "Given normalized_shape=" + s4 = SampleInput( + make_arg((2, 2, 3)), args=((2, 2), None, None, 1e-5) + ) + yield ErrorInput(s4, error_regex=err_msg4) + +def error_inputs_rms_norm(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32, requires_grad=False) + input_shape = (1, 2, 3) + + err_msg1 = "Expected normalized_shape to be at least 1-dimensional" + s1 = SampleInput( + make_arg(input_shape), args=((), None, 1e-5) + ) + yield ErrorInput(s1, error_regex=err_msg1) + + normalized_shape = (1, 2, 3) + weight = make_arg((1, 2)) + err_msg2 = "Expected weight to be of same shape as normalized_shape" + s2 = SampleInput( + make_arg(input_shape), args=(normalized_shape, weight, 1e-5) + ) + yield ErrorInput(s2, error_regex=err_msg2) + + + err_msg4 = "Given normalized_shape=" + s4 = SampleInput( + make_arg((2, 2, 3)), args=((2, 2), None, 1e-5) + ) + yield ErrorInput(s4, error_regex=err_msg4) + + +def sample_inputs_local_response_norm(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Ordered as input shape, size and a kwarg dict for alpha, beta, and k + cases: tuple[tuple[int, ...], tuple[int, ...], dict] = ( + ((1, 6, 3), 2, {'alpha': 3e-05, 'beta': 0.5, 'k': 1.25}), + ((1, 6, 3), 2, {'beta': 0.5, 'k': 1.25}), + ((1, 6, 3), 2, {'alpha': 3e-05, 'k': 1.25}), + ((1, 6, 3), 2, {'alpha': 3e-05, 'beta': 0.5}), + ((1, 6, 3), 2, {'alpha': 3e-05}), + ((1, 6, 3), 2, {'beta': 0.5}), + ((1, 6, 3), 2, {'k': 1.25}), + ((1, 6, 3), 2, {}), + ((2, 6, 3), 2, {'alpha': 3e-05, 'beta': 0.5, 'k': 1.25}), + ((1, 1, 2), 1, {'alpha': 3e-05, 'beta': 0.5, 'k': 1.25}), + ((0, 1, 2), 1, {'alpha': 3e-05, 'beta': 0.5, 'k': 1.25}), + ) + + for input_shape, size, kwargs in cases: + yield SampleInput(make_arg(input_shape), args=(size,), kwargs=kwargs) + +def sample_inputs_hardswish(self, device, dtype, requires_grad, **kwargs): + N = 5 + # make sure we are testing -3 -> 3 range. default is -10 -> 10 so maybe unnecessary ? + make_arg = partial(make_tensor, device=device, dtype=dtype, + requires_grad=requires_grad, low=-5, high=5) + return (SampleInput(make_arg((N * 2, N * 2))) for _ in range(1, N)) + +def sample_inputs_linear(self, device, dtype, requires_grad, **kwargs): + features_options = [[3, 4], [8, 8]] + batch_options: list[list[int]] = [ + [], # no batch + [0], + [8], + [2, 3], + ] + create_tensor = partial(make_tensor, device=device, dtype=dtype, + requires_grad=requires_grad, low=-2, high=2) + + for has_bias, (in_feat, out_feat), batch_shape in \ + itertools.product([True, False], features_options, batch_options): + input_tensor = create_tensor(batch_shape + [in_feat]) + weight = create_tensor([out_feat, in_feat]) + if not has_bias: + yield SampleInput(input_tensor, weight) + continue + + bias = create_tensor([out_feat]) + yield SampleInput(input_tensor, weight, bias) + + # 5D tensor, used to crash on MPS, see https://github.com/pytorch/pytorch/issues/114942 + yield SampleInput(create_tensor(2, 1, 2, 1, 2), create_tensor(4, 2)) + yield SampleInput(create_tensor(2, 1, 2, 1, 2), create_tensor(4, 2), create_tensor(4)) + +def sample_inputs_bilinear(self, device, dtype, requires_grad, **kwargs): + features_options = [[3, 4, 5], [8, 8, 8]] + batch_options: list[list[int]] = [ + [], # no batch + [0], + [8], + [2, 3], + ] + create_tensor = partial(make_tensor, device=device, dtype=dtype, + requires_grad=requires_grad, low=-2, high=2) + + for has_bias, (in_feat1, in_feat2, out_feat), batch_shape in \ + itertools.product([True, False], features_options, batch_options): + input_tensor1 = create_tensor(batch_shape + [in_feat1]) + input_tensor2 = create_tensor(batch_shape + [in_feat2]) + weight = create_tensor([out_feat, in_feat1, in_feat2]) + if not has_bias: + yield SampleInput(input_tensor1, input_tensor2, weight) + continue + bias = create_tensor([out_feat]) + yield SampleInput(input_tensor1, input_tensor2, weight, bias) + +def sample_inputs_glu(self, device, dtype, requires_grad, **kwargs): + features_options = [[2], [2, 4], [8, 8], [3, 6, 8], [1, 4, 6, 7]] + batch_options: list[list[int]] = [ + [], # no batch + [0], + [8], + [2, 3], + ] + create_tensor = partial(make_tensor, device=device, dtype=dtype, + requires_grad=requires_grad, low=-2, high=2) + + for features, batch_shape in itertools.product(features_options, batch_options): + ndim = len(features) + len(batch_shape) + for dim in range(ndim): + input_tensor = create_tensor(batch_shape + features) + dim_size = input_tensor.size(dim) + if dim_size > 0 and dim_size % 2 == 0: + yield SampleInput(input_tensor, dim) + +def sample_inputs_interpolate(mode, self, device, dtype, requires_grad, **kwargs): + N, C = 2, 3 + D = 4 + S = 3 + L = 5 + + align_corners_options: tuple[Any, ...] = (None,) + if mode in ('linear', 'bilinear', 'bicubic', 'trilinear'): + align_corners_options = (True, False, None) + ranks_for_mode = { + 'nearest': [1, 2, 3], + 'nearest-exact': [1, 2, 3], + 'linear': [1], + 'bilinear': [2], + 'bicubic': [2], + 'trilinear': [3], + 'area': [1, 2, 3] + } + + def shape(size, rank, with_batch_channel=True): + if with_batch_channel: + return tuple([N, C] + ([size] * rank)) + return tuple([size] * rank) + + def uneven_shape(size, rank, with_batch_channel=True): + rc = list(shape(size, rank, with_batch_channel)) + rc[-1] += 1 + if rank > 2: + rc[-2] -= 1 + return tuple(rc) + + if mode in ('bilinear', 'bicubic') and dtype == torch.uint8: + make_arg = partial( + make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + # we pick more realistic upper bound 256 instead of default 10 for uint8 dtype + high=256 if dtype == torch.uint8 else None, + ) + # provide few samples for a more close to typical image processing usage + rank = 2 + for memory_format in [torch.contiguous_format, torch.channels_last]: + yield SampleInput( + make_arg(shape(270, rank), memory_format=memory_format), + shape(130, rank, False), + scale_factor=None, + mode=mode, + align_corners=False, + ) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + for align_corners in align_corners_options: + for rank in ranks_for_mode[mode]: + yield SampleInput( + make_arg(shape(D, rank)), + shape(S, rank, False), + scale_factor=None, + mode=mode, + align_corners=align_corners, + ) + yield SampleInput( + make_arg(shape(D, rank)), + shape(L, rank, False), + scale_factor=None, + mode=mode, + align_corners=align_corners, + ) + if rank > 1 and dtype.is_floating_point: + yield SampleInput( + make_arg(uneven_shape(D, rank)), + uneven_shape(S, rank, False), + scale_factor=None, + mode=mode, + align_corners=align_corners, + ) + yield SampleInput( + make_arg(uneven_shape(D, rank)), + uneven_shape(L, rank, False), + scale_factor=None, + mode=mode, + align_corners=align_corners, + ) + for recompute_scale_factor in [False, True]: + for scale_factor in [1.7, 0.6]: + yield SampleInput( + make_arg(shape(D, rank)), + size=None, + scale_factor=scale_factor, + mode=mode, + align_corners=align_corners, + recompute_scale_factor=recompute_scale_factor, + ) + +def reference_inputs_interpolate(mode, self, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_interpolate(mode, self, device, dtype, requires_grad, **kwargs) + + if mode in ('bilinear', 'bicubic'): + make_arg = partial( + make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + # we pick more realistic upper bound 256 instead of default 10 for uint8 dtype + high=256 if dtype == torch.uint8 else None, + ) + # provide few samples for more typical image processing usage + for memory_format in [torch.contiguous_format, torch.channels_last]: + for aa in [True, False]: + yield SampleInput( + make_arg((2, 3, 345, 456), memory_format=memory_format), + (270, 270), + scale_factor=None, + mode=mode, + align_corners=False, + antialias=aa, + ) + +def sample_inputs_upsample(mode, self, device, dtype, requires_grad, **kwargs): + N, C = 2, 3 + D = 4 + S = 3 + L = 5 + + ranks_for_mode = { + 'nearest': [1, 2, 3], + 'bilinear': [2], + } + + def shape(size, rank, with_batch_channel=True): + if with_batch_channel: + return torch.Size([N, C] + ([size] * rank)) + return torch.Size([size] * rank) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + for rank in ranks_for_mode[mode]: + yield SampleInput(make_arg(shape(D, rank)), size=shape(S, rank, False)) + yield SampleInput(make_arg(shape(D, rank)), size=shape(L, rank, False)) + yield SampleInput(make_arg(shape(D, rank)), scale_factor=1.7) + yield SampleInput(make_arg(shape(D, rank)), scale_factor=0.6) + +def reference_inputs_upsample(mode, self, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_upsample(mode, self, device, dtype, requires_grad, **kwargs) + + if mode == 'bilinear': + make_arg = partial( + make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + # we pick more realistic upper bound 256 instead of default 10 for uint8 dtype + high=256 if dtype == torch.uint8 else None, + ) + # provide a single sample for more typical image processing usage + for memory_format in [torch.contiguous_format, torch.channels_last]: + yield SampleInput( + make_arg((2, 3, 345, 456), memory_format=memory_format), + (270, 270), + ) + +def sample_inputs_upsample_aa(mode, self, device, dtype, requires_grad, **kwargs): + N = 6 + C = 3 + H = 10 + W = 20 + S = 3 + L = 5 + + input_tensor = make_tensor(torch.Size([N, C, H, W]), device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(input_tensor, output_size=torch.Size([S, S]), align_corners=False, scale_factors=None) + yield SampleInput(input_tensor, output_size=torch.Size([L, L]), align_corners=False, scale_factors=None) + yield SampleInput(input_tensor, output_size=None, align_corners=False, scale_factors=[1.7, 0.9]) + yield SampleInput(input_tensor, output_size=None, align_corners=True, scale_factors=[0.8, 1.0]) + + yield SampleInput(input_tensor, output_size=torch.Size([S, S]), align_corners=False, scales_h=None, scales_w=None) + yield SampleInput(input_tensor, output_size=torch.Size([S, S]), align_corners=False, scales_h=1.7, scales_w=0.9) + yield SampleInput(input_tensor, output_size=torch.Size([S, S]), align_corners=True, scales_h=1.7, scales_w=0.9) + +def sample_inputs_gelu(self, device, dtype, requires_grad, **kwargs): + N = 5 + for _ in range(1, N): + for approximate in ['none', 'tanh']: + yield SampleInput( + make_tensor((N * 2, N * 2), device=device, dtype=dtype, + requires_grad=requires_grad, low=-3, high=3), + approximate=approximate) + + +def error_inputs_gelu(op, device, **kwargs): + # Tests that gelu errors out when passed an approximation we don't know. + yield ErrorInput(SampleInput(make_tensor((), dtype=torch.float, device=device), kwargs={"approximate": "asdf"}), + error_regex="approximate argument must be either") + + +def sample_inputs_max_min_reduction_with_dim(op_info, device, dtype, requires_grad, **kwargs): + args_for_reduction_with_dim = ( + ((S, S, S), (1,),), + ((S, S, S), (1, True, ),), + ((), (0,),), + ((), (0, True,),), + ) + return ((SampleInput(make_tensor(input_tensor, dtype=dtype, device=device, + low=None, high=None, + requires_grad=requires_grad), + *args)) + for input_tensor, args in args_for_reduction_with_dim) + +def sample_inputs_max_min_reduction_no_dim(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, low=None, high=None) + yield SampleInput(make_arg((S, S, S))) + yield SampleInput(make_arg(())) + +def _generate_nan_reduction_inputs(device, dtype, requires_grad, **kwargs): + yield from _generate_reduction_inputs(device, dtype, requires_grad) + # NaN only exists for floating point numbers + if dtype.is_complex or dtype.is_floating_point: + yield torch.tensor([2, torch.nan, -1], device=device, dtype=dtype, requires_grad=requires_grad) + yield torch.tensor([[torch.nan, 2], [0, 1]], device=device, dtype=dtype, requires_grad=requires_grad) + +def sample_inputs_nan_reduction(supports_multiple_dims): + # Generates sample inputs for reduction ops that contain the input tensor + # and dim and keepdim kwargs. If a reduction op needs to test additional + # args/kwargs then create a separate sample_inputs function + def fn(op_info, device, dtype, requires_grad, **kwargs): + for t in _generate_nan_reduction_inputs(device, dtype, requires_grad): + # Add case without dim and keepdim kwargs + yield SampleInput(t.clone().requires_grad_(requires_grad)) + for kwargs in _generate_reduction_kwargs(t.ndim, supports_multiple_dims): + yield SampleInput(t.clone().requires_grad_(requires_grad), **kwargs) + + return fn + +def sample_inputs_reduction_quantile(op_info, device, dtype, requires_grad, **kwargs): + test_quantiles = (0.5, make_tensor((2,), dtype=dtype, device=device, low=0, high=1, requires_grad=requires_grad)) + test_interpolations = ['linear', 'midpoint'] + + for quantiles in test_quantiles: + for t in _generate_reduction_inputs(device, dtype, requires_grad): + # Add case without dim and keepdim kwargs + input = t.clone().requires_grad_(requires_grad) + yield SampleInput(input, quantiles) + for kwargs in _generate_reduction_kwargs(t.ndim, supports_multiple_dims=False): + # Interpolation kwarg for now is only supported when providing both dim and keepdim + kwargs.setdefault('dim', 0) + kwargs.setdefault('keepdim', False) + for interpolation in test_interpolations: + kwargs['interpolation'] = interpolation + input = t.clone().requires_grad_(requires_grad) + yield SampleInput(input, quantiles, **kwargs) + +def sample_inputs_reduction_count_nonzero(*args, **kwargs): + """Sample inputs for count_nonzero""" + # count_nonzero does not support keepdim yet + for sample in sample_inputs_reduction(*args, **kwargs): + sample.kwargs.pop('keepdim', None) + yield sample + +def sample_inputs_leaky_relu(op_info, device, dtype, requires_grad, **kwargs): + N = 10 + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + return (SampleInput(make_arg((N, N))) for _ in range(1, N)) + +def sample_inputs_fractional_max_pool2d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Order: input_shape, kernel_size + cases = (((1, 3, 9, 9), 3), + ((1, 3, 9, 9), (4, 4)), + ((1, 3, 9, 9), (6, 6)), + ((2, 3, 9, 9), (3, 3)), + ((1, 1, 4, 4), (2, 2)), + ((1, 2, 6, 6), (4, 4))) + + for input_shape, kernel_size in cases: + for return_indices in [False, True]: + # test case passing a single output size + yield SampleInput( + make_arg(input_shape), + kernel_size, + output_size=2, + return_indices=return_indices, + ) + + # test case passing a tuple output size + yield SampleInput( + make_arg(input_shape), + kernel_size, + output_size=(2, 3), + return_indices=return_indices, + ) + + # test case passing an output ratio + yield SampleInput( + make_arg(input_shape), + kernel_size, + output_ratio=(0.5, 0.5), + return_indices=return_indices, + ) + + yield SampleInput( + make_arg((1, 1, 16, 16)), + (1, 1), + output_ratio=(0.5, 0.5), + return_indices=True, + _random_samples=make_tensor((1, 1, 2), device=device, dtype=dtype, requires_grad=False), + ) + +def sample_inputs_fractional_max_pool3d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Order: input_shape, kernel_size + cases = (((2, 3, 5, 5, 5), (2, 2, 2)), + ((1, 2, 6, 5, 4), 2), + ((1, 2, 5, 6, 5), (2, 3, 2)), + ((1, 2, 6, 6, 6), (2, 3, 2)), + ((1, 1, 7, 6, 7), (2, 3, 4)), + ((1, 1, 4, 5, 4), (2, 2, 1)), + ((1, 1, 8, 7, 6), (4, 3, 2)), + ((0, 1, 4, 5, 4), (2, 2, 1))) + + for input_shape, kernel_size in cases: + for return_indices in [False, True]: + # test case passing a single output size + yield SampleInput( + make_arg(input_shape), + kernel_size, + output_size=2, + return_indices=return_indices, + ) + + # test case passing a tuple output size + yield SampleInput( + make_arg(input_shape), + kernel_size, + output_size=(2, 3, 2), + return_indices=return_indices, + ) + + # test case passing an output ratio + yield SampleInput( + make_arg(input_shape), + kernel_size, + output_ratio=(0.5, 0.5, 0.5), + return_indices=return_indices, + ) + +def sample_inputs_avgpool2d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Order: input_shape, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override + cases = (((1, 3, 9, 9), 3, 1, 1, True, False, 2), + ((1, 3, 9, 9), (4, 4), (2, 3), 1, True, False, 2), + ((1, 3, 9, 9), (6, 6), (3, 3), (2, 3), True, True, 2), + ((2, 3, 9, 9), (3, 3), (1, 1), (1, ), True, False, 2), + ((1, 1, 4, 4), (2, 2), (), (0, ), False, True, -2), + ((1, 2, 6, 6), (4, 4), (2, 2), (2, ), True, True, None)) + + for input_shape, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override in cases: + yield SampleInput(make_arg(input_shape), + args=(kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override)) + # Case with just input_shape and kernel_size + yield SampleInput(make_arg((1, 3, 9, 9)), args=((3, 3))) + +def sample_inputs_avgpool1d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Order: input_shape, kernel_size, kwargs + cases: list[tuple[tuple[int, ...], int | tuple[int, ...], dict]] = [ + ((2, 3, 9), (3,), {}), + ((1, 3, 9), 3, dict(stride=1, padding=1, ceil_mode=True, count_include_pad=False)), + ((1, 3, 9), (6,), dict(stride=(3,), padding=(2,), ceil_mode=True, count_include_pad=True)), + ((2, 3, 9), (3,), dict(stride=(1,), padding=(1,), ceil_mode=False, count_include_pad=True)), + ((0, 3, 9), (6,), dict(stride=(3,), padding=(2,), ceil_mode=False, count_include_pad=True)), + ((1, 2, 9), (7,), dict(stride=(3,), padding=(2,), ceil_mode=False)), + ((1, 2, 9), (7,), dict(stride=(3,), padding=(3,), ceil_mode=True)), + ((1, 2, 9), (7,), dict(stride=(3,), ceil_mode=False)), + ((1, 2, 9), (7,), dict(stride=(3,), ceil_mode=True)), + ] + + for input_shape, kernel_size, kwargs in cases: + yield SampleInput(make_arg(input_shape), args=(kernel_size,), kwargs=kwargs) + +def sample_inputs_avgpool3d(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Order: input_shape, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override + cases: list[tuple[tuple[int, ...], int | tuple[int, ...], dict]] = [ + ((2, 3, 3, 4, 4), (2, 2, 2), {}), + ((1, 2, 4, 4, 4), 2, dict(stride=1, padding=1, ceil_mode=True, + count_include_pad=False, divisor_override=2)), + ((1, 2, 5, 5, 5), (2, 3, 4), dict(stride=(1, 2, 2), padding=(0, 1, 2), ceil_mode=True, + count_include_pad=True, divisor_override=2)), + ((1, 2, 5, 5, 5), (2, 3, 4), dict(stride=(1, 2, 2), padding=(0, 1, 2), ceil_mode=False)), + ((1, 1, 7, 5, 7), (6, 3, 4), dict(stride=(2, 3, 2), padding=(3, 1, 0), ceil_mode=False, + count_include_pad=False, divisor_override=2)), + ((1, 1, 4, 5, 4), (2, 2, 3), dict(stride=(2, 2, 1), padding=0, ceil_mode=False, + count_include_pad=True, divisor_override=-2)), + ((1, 1, 6, 5, 6), (4, 5, 6), dict(stride=(2, 3, 2), padding=2, ceil_mode=True, + count_include_pad=True, divisor_override=None)), + ((0, 1, 4, 5, 4), (2, 3, 1), dict(stride=(2, 1, 2), padding=0, ceil_mode=False, + count_include_pad=True, divisor_override=None)), + ] + + for input_shape, kernel_size, kwargs in cases: + yield SampleInput(make_arg(input_shape), args=(kernel_size,), kwargs=kwargs) + +def error_inputs_avg_pool1d(op_info, device, **kwargs): + # error inputs when pad is negative + x = torch.rand([0, 1, 49], dtype=torch.float32) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': -1}), + error_regex='pad must be non-negative') + + # error inputs when pad > kernel_size / 2 + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': 4}), + error_regex='pad should be at most half of effective kernel size') + +def error_inputs_avg_pool2d(op_info, device, **kwargs): + # error inputs when pad is negative + x = torch.rand([0, 1, 49], dtype=torch.float32) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': -1}), + error_regex='pad must be non-negative') + # 2-dimensional kernel + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (3, 2), 'stride': 50, 'padding': -1}), + error_regex='pad must be non-negative') + + # error inputs when pad > kernel_size / 2 + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': 4}), + error_regex='pad should be at most half of effective kernel size') + # 2-dimensional kernel + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (3, 2), 'stride': 50, 'padding': 4}), + error_regex='pad should be at most half of effective kernel size') + + # error inputs for zero divisor + x = torch.zeros(3, 3, 3) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (2, 2), 'divisor_override': 0}), + error_regex='divisor must be not zero') + +def error_inputs_avg_pool3d(op_info, device, **kwargs): + # error inputs when pad is negative + x = torch.rand([0, 1, 49, 50], dtype=torch.float32) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': -1}), + error_regex='pad must be non-negative') + # 3-dimensional kernel + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (3, 2, 2), 'stride': 50, 'padding': -1}), + error_regex='pad must be non-negative') + + # error inputs when pad > kernel_size / 2 + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': 4}), + error_regex='pad should be at most half of effective kernel size') + # 3-dimensional kernel + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (3, 2, 2), 'stride': 50, 'padding': 4}), + error_regex='pad should be at most half of effective kernel size') + + # error inputs for zero divisor + x = torch.zeros(3, 3, 3, 3) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': (2, 2, 2), 'divisor_override': 0}), + error_regex='divisor must be not zero') + + # error inputs for invalid input dimension + x = torch.rand([0, 1, 49], dtype=torch.float32) + yield ErrorInput(SampleInput(x, kwargs={'kernel_size': 2, 'stride': 50, 'padding': 0}), + error_regex='non-empty 4D or 5D') + + +def sample_inputs_to(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + # test_multiple_devices_to_cuda would fail if we use a different device than given + devices = [device] + if torch.device(device).type == 'cpu': + devices = [torch.device('cpu'), torch.device('cuda:0')] if torch.cuda.is_available() else devices + memory_formats = [torch.preserve_format, torch.channels_last] + + # TODO: can't switch `to.device` overload to use positional arguments + # https://github.com/pytorch/pytorch/issues/84265 + # to.device overload + for device, nb, cp, mem_f in product(devices, [True, False], [True, False], memory_formats): + kwargs = { + "memory_format": mem_f, + } + yield SampleInput(make_arg((S, S, S, S)), args=(device, torch.float64, nb, cp), kwargs=kwargs) + + other_dtype = torch.bfloat16 if torch.device(device).type == 'mps' else torch.float64 + + # to.dtype overload + for nb, cp, mem_f in product([True, False], [True, False], memory_formats): + kwargs = { + "memory_format": mem_f, + } + yield SampleInput(make_arg((S, S, S, S)), args=(other_dtype, nb, cp), kwargs=kwargs) + + # to.other overload + for device, nb, cp, mem_f in product(devices, [True, False], [True, False], memory_formats): + kwargs = { + "memory_format": mem_f, + } + other = make_arg((S, S, S, S), dtype=other_dtype, device=device) + yield SampleInput(make_arg((S, S, S, S)), args=(other, nb, cp), kwargs=kwargs) + + +def sample_inputs_topk(op_info, device, dtype, requires_grad, **kwargs): + def get_tensor_input(size): + return make_tensor(size, dtype=dtype, device=device, requires_grad=requires_grad) + + yield SampleInput(get_tensor_input((S, M, S)), 3) + yield SampleInput(get_tensor_input((S, M, S)), 3, 1) + yield SampleInput(get_tensor_input((S, M, S)), 3, -2) + yield SampleInput(get_tensor_input((S, M, S)), 3, 1, True) + yield SampleInput(get_tensor_input((S, M, S)), 3, -2, True) + yield SampleInput(get_tensor_input((S, M, S)), 3, 1, True, True) + yield SampleInput(get_tensor_input((S, M, S)), 3, -2, True, True) + + yield SampleInput(get_tensor_input(()), 1) + yield SampleInput(get_tensor_input(()), 1, 0) + yield SampleInput(get_tensor_input(()), 1, -1) + yield SampleInput(get_tensor_input(()), 1, 0, True) + yield SampleInput(get_tensor_input(()), 1, -1, True) + yield SampleInput(get_tensor_input(()), 1, 0, True, True) + yield SampleInput(get_tensor_input(()), 1, -1, True, True) + +def sample_inputs_outer(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg(S), make_arg(M)) + +def sample_inputs_dist(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + sizes = ((S, S, S), (S,), (S, 1, S), (), (S, S)) + ps = (2, 4) + + for size_x, size_y, p in product(sizes, sizes, ps): + yield SampleInput(make_arg(size_x), args=(make_arg(size_y), p)) + +# Missing to test the nondeterminism of the operation +# https://github.com/pytorch/pytorch/issues/53352 +def sample_inputs_index(op_info, device, dtype, requires_grad, reference=False, **kwargs): + # target.index_add(dim, idx, source, *, alpha=1) + add = "index_add" in op_info.name + # target.index_copy(dim, idx, source) + copy = "index_copy" in op_info.name + # target.index_fill(dim, idx, value) + fill = "index_fill" in op_info.name + + # Extended reference inputs. We generate that exercise atomic adds / writing + # several times to one location + if reference: + make_arg = partial(torch.ones, device=device, dtype=dtype, requires_grad=requires_grad) + make_idx = partial(torch.zeros, device=device, dtype=torch.int64) + else: + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + # idx They need to be different for copy and add to be deterministic + if copy or add: + make_idx = partial(torch.randperm, device=device, dtype=torch.int64) + else: + def make_idx(n): + return make_tensor((n,), device=device, dtype=torch.int64, low=0, high=n) + + shapes = [(), (1,), (S, S)] + # extra parameter for add + if add: + if dtype == torch.bool: + alphas = (True, False) + else: + alphas = (-1, 0, 2) + else: + alphas = (None,) + + if fill: + # A weird number to catch errors. + # The former one tests `index_fill.int_Scalar`, and the latter one tests `index_fill.int_Tensor`. + values = (make_arg((1,)).item(), make_arg(())) + else: + values = (None,) + + for shape, alpha, value in product(shapes, alphas, values): + t = make_arg(shape) + args = [] + + # dim. We handle the scalar case + dim = -1 if t.ndim == 2 else 0 + args.append(dim) + + idx = make_idx(t.shape[dim] if t.ndim != 0 else 1) + args.append(idx) + + # source + if copy or add: + args.append(make_arg(shape)) + elif fill: + args.append(value) + + args = tuple(args) + kwargs = {} if alpha is None else {"alpha": alpha} + + yield SampleInput(t, args=args, kwargs=kwargs) + +def sample_inputs_index_reduce(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_idx(n, m): + return make_tensor((n,), device=device, dtype=torch.int64, low=0, high=m) + + shapes = [((), ()), ((1,), (1,)), ((S, S), (S, M)), ((S, S, S), (S, M, S))] + include_selfs = (True, False) + reduce = op_info.variant_test_name + if reduce not in ('prod', 'mean', 'amin', 'amax'): + raise AssertionError(f"Expected reduce to be one of 'prod', 'mean', 'amin', 'amax', got {reduce!r}") + + for shape, include_self in product(shapes, include_selfs): + self_shape, src_shape = shape + # dim. We handle the scalar case + dim = 1 if len(self_shape) >= 2 else 0 + idx = make_idx(src_shape[dim] if len(src_shape) != 0 else 1, + self_shape[dim] if len(self_shape) != 0 else 1) + args = (dim, idx, make_arg(src_shape), reduce) + yield SampleInput(make_arg(self_shape), + args=args, + kwargs={'include_self' : include_self}) + + # Sample inputs to test edge cases for backward + if requires_grad and reduce == 'prod': + # Check that gradients are propagated correctly for prod when zeros in self/src are reduced + # This sample tests gradients for the following cases + # (a) 1 zero reduced (from source (self[0, 1]), from self (self[0, 0])) + # (b) 2 zeros reduced (1 from src and 1 from self (self[1, 0], self[1, 1]) + # (c) no zeros reduced (self[2, 1], self[2, 2]) + # (d) 2 zeros reduced (both from src) is tested in test/test_autograd.py + # test_scatter_index_reduce_prod_gradgrad_error as this case is not supported for gradgrad + input = torch.tensor([[0, 13], [0, 0], [15, 19]], dtype=dtype, device=device, requires_grad=requires_grad) + src = torch.tensor([[2, 0], [0, 0], [2, 3], [2, 2]], dtype=dtype, device=device, requires_grad=requires_grad) + idx = torch.tensor([0, 1, 2, 0], dtype=torch.long, device=device) + + yield SampleInput(input, + args=(0, idx, src, reduce), + kwargs={'include_self': True}) + +def sample_inputs__unsafe_masked_index(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_idx(n, m, dim, d): + view_shape = [1] * dim + view_shape[d] = n + return make_tensor((n,), device=device, dtype=torch.int64, low=0, high=m).view(view_shape) + + cases = [ + ((S, S), S, M), + ((S, S), M, S), + ((S, S, S), S, M), + ] + + fill_value = make_tensor([], dtype=dtype, device="cpu").item() + + for c in cases: + self_shape, high, idx_size = c + dim = len(self_shape) + indices = [make_idx(idx_size, high, dim, d) for d in range(dim)] + masks = [torch.logical_and(idx >= 0, idx < self_shape[i]) for i, idx in enumerate(indices) if idx is not None] + mask = functools.reduce(torch.logical_and, masks) + yield SampleInput(make_arg(self_shape), mask, indices, fill_value) + + masks = [torch.logical_and(idx >= 1, idx < self_shape[i] - 1) for i, idx in enumerate(indices) if idx is not None] + mask = functools.reduce(torch.logical_and, masks) + yield SampleInput(make_arg(self_shape), mask, indices, fill_value) + +def sample_inputs__unsafe_masked_index_put_accumulate(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_idx(n, m, dim, d): + view_shape = [1] * dim + view_shape[d] = n + return make_tensor((n,), device=device, dtype=torch.int64, low=0, high=m).view(view_shape) + + cases = [ + ((S, S), S, (M, M)), + ((S, S), M, (S, S + 1)), + ((S, S, S), S, (M, M - 1, M + 1)), + ] + + for c in cases: + self_shape, high, idx_sizes = c + dim = len(self_shape) + indices = [make_idx(idx_sizes[d], high, dim, d) for d in range(dim)] + masks = [torch.logical_and(idx >= 0, idx < self_shape[i]) for i, idx in enumerate(indices) if idx is not None] + mask = functools.reduce(torch.logical_and, masks) + values = make_arg(idx_sizes) + if device == 'mps:0' and dtype in [torch.float16, torch.bfloat16]: + # TestConsistencyMPS.test_output_match compares CPU to MPS results + # Order of operations in GPU index_put_accumulate is not guaranteed, + # which can result in significant divergence between sequential and parallel execution + # Unless inputs are normalized + values = torch.nn.functional.normalize(values) + yield SampleInput(make_arg(self_shape), mask, indices, values) + + masks = [torch.logical_and(idx >= 1, idx < self_shape[i] - 1) for i, idx in enumerate(indices) if idx is not None] + mask = functools.reduce(torch.logical_and, masks) + yield SampleInput(make_arg(self_shape), mask, indices, values) + + +def sample_inputs_mode(op_info, device, dtype, requires_grad, **kwargs): + args = ( + ((S, S, S), (),), + ((S, S, S), (1, ),), + ((S, S, S), (1, True, ),), + ((), (),), + ((), (0,),), + ((), (0, True,),), + # Non-fused mode kernel on CUDA + ((3000,), ()), + ) + make_arg = partial(make_tensor, dtype=dtype, device=device, + requires_grad=requires_grad, low=None, high=None) + return (SampleInput(make_arg(input_tensor), *args) + for input_tensor, args in args) + +# Missing to test the nondeterminism of the operation +# https://github.com/pytorch/pytorch/issues/53352 +def sample_inputs_put(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + make_idx = partial(make_tensor, low=0, dtype=torch.int64, device=device, requires_grad=False) + + S = 3 + + # Generic inputs + idx = torch.randperm(S * S, device=device, dtype=torch.int64)[:S] + idx_list = [idx, -idx - 1] + for idx, acc in product(idx_list, (True, False)): + yield SampleInput(input=make_arg((S, S)), + args=(idx.clone(), + make_arg((S,)), + acc)) + + # Scalar cases + scalar_sizes = [(), (1,)] + tgt_gen = (make_arg(size) for size in scalar_sizes) + idx_gen = (make_idx(size, high=1) for size in scalar_sizes) + src_gen = (make_arg(size) for size in scalar_sizes) + for tgt, idx, src, acc in product(tgt_gen, idx_gen, src_gen, (True, False)): + yield SampleInput(input=tgt.clone().requires_grad_(requires_grad), + args=(idx.clone(), + src.clone().requires_grad_(requires_grad), + acc)) + + # Empty cases + tgt_sizes = [(0,), (), (1,), (3, 2)] + tgt_gen = (make_arg(size) for size in tgt_sizes) + idx = make_idx((0,), high=1) + src = make_arg((0,)) + for tgt, acc in product(tgt_gen, (True, False)): + yield SampleInput(input=tgt.clone().requires_grad_(requires_grad), + args=(idx.clone(), + src.clone().requires_grad_(requires_grad), + acc)) + +def sample_inputs_take(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + make_idx = partial(make_tensor, low=0, dtype=torch.int64, device=device, requires_grad=False) + + S = 3 + + # Generic inputs: take S elements out of S * S + index = make_idx((S,), high=(S * S)) + for idx in (index, -index - 1): + yield SampleInput(input=make_arg((S, S)), args=(idx,)) + + # Scalar cases + scalar_sizes = [(), (1,)] + src_gen = (make_arg(size) for size in scalar_sizes) + idx_gen = (make_idx(size, high=1) for size in scalar_sizes) + for src, idx in product(src_gen, idx_gen): + yield SampleInput(input=src.clone().requires_grad_(requires_grad), + args=(idx.clone(),)) + + # Empty cases + src_sizes = [(0,), (), (1,), (3, 2)] + src_gen = (make_arg(size) for size in src_sizes) + + idx = make_idx((0,), high=1) + for src in src_gen: + yield SampleInput(input=src.clone().requires_grad_(requires_grad), + args=(idx.clone(),)) + +def sample_movedim_moveaxis(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, low=None, high=None, requires_grad=requires_grad) + yield SampleInput(make_arg((4, 3, 2, 1)), [0, 1, 2, 3], [3, 2, 1, 0]) + yield SampleInput(make_arg((4, 3, 2, 1)), [0, -1, -2, -3], [-3, -2, -1, -0]) + +def reference_movedim_moveaxis(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_movedim_moveaxis(op_info, device, dtype, requires_grad, **kwargs) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # shape, source, destination + args = ( + # empty inputs + ((), (), ()), + # int inputs, negative + ((3, 5, 7, 2), -2, 1), + # swap bounds + ((3, 5, 7, 2), (-1, 0), (0, -1)), + # non-sequential, negative + ((2, 3, 4, 5, 6), (3, -3, 4), (1, 0, -1)), + # idempotence, negative + ((2, 3, 4, 5, 6), (-3, 4, 3, 1), (-3, 4, 3, 1)), + # reverse, sequential, positive + ((6, 2, 3, 5, 4), (4, 3, 2, 1, 0), (0, 1, 2, 3, 4)), + # reverse, non-sequential + ((6, 2, 3, 5, 4), (-3, -2, -4, -5, -1), (2, 1, 3, 4, 0)), + # reverse, sequential, negative + ((6, 2, 3, 5, 4), (4, -2, 2, -4, -5), (-5, 1, 2, -2, -1)), + ) + + for shape, source, destination in args: + yield SampleInput(make_arg(shape), args=(source, destination)) + +def error_movedim_moveaxis(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # source length < destination length + yield ErrorInput( + SampleInput(make_arg(2, 3, 4, 5, 6), args=((3, -3), (1, 0, -1))), + error_regex=(r"movedim: Invalid source or destination dims: source " + r"\(\[3, -3\] dims\) should contain the same number of " + r"dims as destination \(\[1, 0, -1\] dims\)"), + ) + + # source length > destination length + yield ErrorInput( + SampleInput(make_arg(2, 3, 4, 5, 6), args=((3, -3, 4), (1, 0))), + error_regex=(r"movedim: Invalid source or destination dims: source " + r"\(\[3, -3, 4\] dims\) should contain the same number of " + r"dims as destination \(\[1, 0\] dims\)"), + ) + + # repeated source dim, with negative indices + yield ErrorInput( + SampleInput(make_arg(2, 3, 4, 5, 6), args=((0, 4, -5), (1, 0, 2))), + error_regex=r"movedim: repeated dim in `source` \(\[0, 4, -5\]\)", + ) + + # repeated destination dim, with negative indices + yield ErrorInput( + SampleInput(make_arg(2, 3, 4, 5, 6), args=((1, 0, 2), (0, 4, -5))), + error_regex=r"movedim: repeated dim in `destination` \(\[0, 4, -5\]\)", + ) + + # repeated dim (both), with negative indices + yield ErrorInput( + SampleInput(make_arg(2, 3, 4, 5, 6), args=((1, 0, -4), (0, 4, -5))), + error_regex=r"movedim: repeated dim in `source` \(\[1, 0, -4\]\)", + ) + + # out of bounds source inputs, with negative indices + yield ErrorInput( + SampleInput(make_arg(2, 3, 4, 5, 6), args=((0, 1, -6), (1, 4, 2))), + error_regex=r"Dimension out of range \(expected to be in range of \[-5, 4\], but got -6\)", + error_type=IndexError, + ) + + # out of bounds destination inputs, with negative indices + yield ErrorInput( + SampleInput(make_arg(2, 3, 4, 5, 6), args=((1, 4, 2), (0, 1, -6))), + error_regex=r"Dimension out of range \(expected to be in range of \[-5, 4\], but got -6\)", + error_type=IndexError, + ) + + # out of bounds source input, int + yield ErrorInput( + SampleInput(make_arg(2, 3, 4, 5, 6), args=(-6, 1)), + error_regex=r"Dimension out of range \(expected to be in range of \[-5, 4\], but got -6\)", + error_type=IndexError, + ) + + # out of bounds destination input, int + yield ErrorInput( + SampleInput(make_arg(2, 3, 4, 5, 6), args=(3, -6)), + error_regex=r"Dimension out of range \(expected to be in range of \[-5, 4\], but got -6\)", + error_type=IndexError, + ) + +def sample_repeat_tile(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + rep_dims = ((), (0, ), (1, ), (0, 2), (1, 1), (2, 3), (2, 3, 2), (0, 2, 3), (2, 1, 1, 1),) + shapes = ((), (0,), (2,), (3, 0), (3, 2), (3, 0, 1)) + + if requires_grad: + # Tests for variant_consistency_jit, grad, gradgrad + # are slower. Use smaller bags of `rep_dims` and `shapes` + # in this case. + rep_dims = ((), (0, ), (0, 2), (1, 1), (2, 3), (1, 3, 2), (3, 1, 1)) # type: ignore[assignment] + shapes = ((), (0,), (2,), (3, 2)) # type: ignore[assignment] + + is_repeat_op = op_info.name in ['repeat', '_refs.repeat'] + for rep_dim, shape in product(rep_dims, shapes): + # `torch.repeat` errors for `len(rep_dims) < t.dim()`, + # so we filter such combinations. + if is_repeat_op and len(rep_dim) < len(shape): + continue + yield SampleInput(make_arg(shape), rep_dim) + + +def sample_inputs_narrow_narrow_copy(op_info, device, dtype, requires_grad, *, is_narrow, **kwargs): + shapes_and_args = ( + ((S, S, S), 1, 2, 2), + ((S, S, S), -1, 2, 2), + ((S, S, S), 1, 0, 0), + ((S, S, S), -1, 0, 0), + ((S, S, S), 2, 1, 2), + ) + + for shape, dim, start, length in shapes_and_args: + tensor = make_tensor(shape, dtype=dtype, device=device, low=None, high=None, + requires_grad=requires_grad) + yield SampleInput(tensor, dim, start, length) + # narrow also accepts the start argument being a Tensor + if is_narrow: + yield SampleInput(tensor, dim, torch.tensor(start), length) + +def reference_inputs_narrow_narrow_copy(op_info, device, dtype, requires_grad, *, is_narrow, **kwargs): + yield from sample_inputs_narrow_narrow_copy(op_info, device, dtype, requires_grad, is_narrow=is_narrow, **kwargs) + + shapes_and_args = ( + # 1-dim + ((M,), 0, 0, 0), # 0 elems from the left + ((M,), -1, -1, 0), # 0 elems from the right + ((M,), 0, 5, 3), # 3 elems from the left + ((M,), 0, -5, 2), # 2 elems from the right + ((M,), -1, 0, M), # M elems from the left + ((M,), 0, -M, M), # M elems from the right + + # 2-dim + ((M, S), 1, 0, 0), # dim 1, 0 elems from the left + ((S, M), -2, -1, 0), # dim 0, 0 elems from the right + ((L, S), 1, 2, 3), # dim 1, 3 elems from the left + ((L, S), -1, 3, 2), # dim 1, 2 elems from the left + ((M, L), 0, 0, M), # dim 0, M elems from the left + ((M, L), -1, -L, L), # dim 1, L elems from the right + + # 3-dim + ((L, M, S), 2, 0, 0), # dim 2, 0 elems from the left + ((M, S, L), -1, -1, 0), # dim 2, 0 elems from the right + ((S, L, M), 2, 0, M), # dim 2, M elems from the left + ((L, S, M), -1, -M, M), # dim 2, M elems from the right + ((S, L, M), 1, 0, 0), # dim 1, 0 elems from the left + ((S, L, M), 0, 2, 1), # dim 0, 1 elem from the left + ((M, S, M), -1, -5, 4), # dim 2, 4 elems from the right + ) + + for shape, dim, start, length in shapes_and_args: + tensor = make_tensor(shape, dtype=dtype, device=device, low=None, high=None, + requires_grad=requires_grad) + yield SampleInput(tensor, dim, start, length) + # narrow also accepts the start argument being a Tensor + if is_narrow: + yield SampleInput(tensor, dim, torch.tensor(start), length) + +def error_inputs_narrow_narrow_copy(op_info, device, *, is_narrow, is_ref): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # 0-dim + yield ErrorInput(SampleInput(make_arg(()), 0, 0, 1), + error_type=RuntimeError, + error_regex=r"narrow\(\) cannot be applied to a 0-dim tensor\.") + + # out of bounds dim + if not is_narrow and not is_ref and torch.device(device).type == 'cpu': + # narrow_copy_dense_cpu_out + yield ErrorInput(SampleInput(make_arg((M, S, L)), 3, 0, 0), + error_type=RuntimeError, + error_regex=r"Expected dim < static_cast\(self_sizes.size\(\)\) to be true, but got false\.") + else: + yield ErrorInput(SampleInput(make_arg((M, S, L)), 3, 0, 0), + error_type=IndexError, + error_regex=r"Dimension out of range \(expected to be in range of \[-3, 2\], but got 3\)") + # out of bounds dim (negative) + yield ErrorInput(SampleInput(make_arg((L, S, M)), -4, 0, 0), + error_type=IndexError, + error_regex=r"Dimension out of range \(expected to be in range of \[-3, 2\], but got -4\)") + + # out of bounds start + yield ErrorInput(SampleInput(make_arg((L, M, S)), 1, M + 1, 0), + error_type=IndexError, + error_regex=r"start out of range \(expected to be in range of \[-10, 10\], but got 11\)") + # out of bounds start (negative) + yield ErrorInput(SampleInput(make_arg((L, M, S)), 1, -M - 1, 0), + error_type=IndexError, + error_regex=r"start out of range \(expected to be in range of \[-10, 10\], but got -11\)") + + # out of bounds length + yield ErrorInput(SampleInput(make_arg((S, L, M)), 2, 0, M + 1), + error_type=RuntimeError, + error_regex=r"start \(0\) \+ length \(11\) exceeds dimension size \(10\)\.") + # out of bounds length (negative) + if not is_narrow and not is_ref and torch.device(device).type == 'cpu': + # narrow_copy_dense_cpu_out + yield ErrorInput(SampleInput(make_arg((M,)), 0, 0, -1), + error_type=RuntimeError, + error_regex=r"start \(0\) \+ length \(-1\) exceeds dimension size \(10\)\.") + else: + yield ErrorInput(SampleInput(make_arg((M,)), 0, 0, -1), + error_type=RuntimeError, + error_regex=r"narrow\(\): length must be non-negative\.") + + # Test Tensor overload that was added for XLA. Start must be an 0-dim + # integral Tensor. narrow_copy doesn't have this overload. + # https://github.com/pytorch/pytorch/issues/31558 + if is_narrow: + # *1-dim* integral Tensor + yield ErrorInput(SampleInput(make_arg((L, M, S)), 1, make_arg(S, dtype=torch.int), 2), + error_type=RuntimeError, + error_regex=r"start must be an 0-dim integral Tensor\.") + + # 0-dim *bool* Tensor (bools are not allowed) + yield ErrorInput(SampleInput(make_arg((L, M, S)), -3, make_arg((), dtype=torch.bool), 3), + error_type=RuntimeError, + error_regex=r"start must be an 0-dim integral Tensor\.") + + +def sample_trapezoid(op_info, device, dtype, requires_grad, **kwargs): + y_shape_x_shape_and_kwargs = [ + ((2, 3), (2, 3), {}), + ((2, 3), (2, 3), {'dim': 1}), + ((6,), (6,), {}), + ((6,), None, {}), + # When 'trapezoid' is called with an empty input, it does not produce an output with requires_grad + # See Issue #{61619} + # ((6,0), (6,0), {}), + ((2, 3), (1, 3), {}), + ((3, 3), (3, 3), {}), + ((3, 3), (3, 3), {'dim': -2}), + ((5,), None, {'dx': 2.0}), + ((2, 2), None, {'dx': 3.0}) + ] + make_arg = partial(make_tensor, dtype=dtype, device=device, low=None, high=None, + requires_grad=requires_grad) + for y_shape, x_shape, kwarg in y_shape_x_shape_and_kwargs: + y_tensor = make_arg(y_shape) + if x_shape is not None: + x_tensor = make_arg(x_shape) + yield SampleInput(y_tensor, x_tensor, **kwarg) + else: + yield SampleInput(y_tensor, **kwarg) + +def sample_cumulative_trapezoid(op_info, device, dtype, requires_grad, **kwargs): + + y_shape_x_shape_and_kwargs = [ + ((2, 3), (2, 3), {}), + ((2, 3), (2, 3), {'dim': 1}), + ((6,), (6,), {}), + ((6,), None, {}), + # When 'cumulative_trapezoid' is called with an empty input, it does not produce an output with requires_grad + # See Issue #{61619} + # ((6,0), (6,0), {}), + ((2, 3), (1, 3), {}), + ((3, 3), (3, 3), {}), + ((3, 3), (3, 3), {'dim': -2}), + ((5,), None, {'dx': 2.0}), + ((2, 2), None, {'dx': 3.0}) + ] + make_arg = partial(make_tensor, device=device, dtype=dtype, + requires_grad=requires_grad, low=None, high=None) + for y_shape, x_shape, kwarg in y_shape_x_shape_and_kwargs: + y_tensor = make_arg(y_shape) + if x_shape is not None: + x_tensor = make_arg(x_shape) + yield SampleInput(y_tensor, x_tensor, **kwarg) + else: + yield SampleInput(y_tensor, **kwarg) + +def sample_unsqueeze(op_info, device, dtype, requires_grad, **kwargs): + shapes_and_axes = [ + ((3, 4, 5), 0), + ((3, 4, 5), 1), + ((3, 4, 5), 3), + ((3, 4, 5), -1), + ((3, 4, 5), -3), + ((), 0), + ((), -1), + ((1,), 0), + ((1,), -1), + ] + + for shape, axis in shapes_and_axes: + tensor = make_tensor(shape, dtype=dtype, device=device, low=None, high=None, + requires_grad=requires_grad) + yield SampleInput(tensor, axis) + + +def sample_inputs_nn_unfold(op_info, device, dtype, requires_grad, **kwargs): + shapes = ((0, 1, 5, 5), (2, 3, 5, 5)) + kernel_sizes = (2, (2, 2), (2, 3)) + dilations = (1, 2, (1, 2)) + paddings = (0, 1, (1, 2)) + strides = (1, 2, (1, 2)) + + cases = product(shapes, kernel_sizes, dilations, paddings, strides) + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + for shape, kernel_size, dilation, padding, stride in cases: + tensor = make_arg(shape) + yield SampleInput(tensor, kernel_size, dilation, padding, stride) + + # With default args + yield SampleInput(make_arg((1, 1, 5, 5)), (3, 3)) + + +def sample_inputs_squeeze(op_info, device, dtype, requires_grad, **kwargs): + shapes_and_args = ( + ((S, 1, S, 1), ()), + ((1, 1, 1, 1), ()), + ((1, 1, 1, 1), (0,)), + ((S, 1, S, 1), (1,)), + ((S, 1, S, 1), (-1,)), + ((S, 1, S, 1), (2,)), + ((S, 1, S, 1), (-2,)), + ((), (0, )), + ) + + for shape, args in shapes_and_args: + tensor = make_tensor(shape, dtype=dtype, device=device, low=None, high=None, + requires_grad=requires_grad) + + yield SampleInput(tensor, args=args) + + +def sample_inputs_squeeze_multiple(op_info, device, dtype, requires_grad, **kwargs): + shapes_and_args = ( + ((1, 1, 1, 1), ()), + ((S, 1, S, 1), (1,)), + ((S, 1, S, 1), (-1,)), + ((S, 1, S, 1), (1, 3)), + ((S, 1, S, 1), (1, 2,)), + ((), (0,)), + ) + + for shape, dims in shapes_and_args: + tensor = make_tensor(shape, dtype=dtype, device=device, low=None, high=None, + requires_grad=requires_grad) + + yield SampleInput(tensor, dims) + + +def _squeeze_ref(x, axis=None): + # NumPy doesn't allow squeezing scalars + if x.ndim == 0: + return x + + if isinstance(axis, Sequence): + # Numpy doesn't allow specifying non-singular dimensions + axis = tuple(a for a in axis if x.shape[a] == 1) + + if isinstance(axis, int) and x.shape[axis] != 1: + return x + + return np.squeeze(x, axis) + +def sample_inputs_nn_pad(op_info, device, dtype, requires_grad, mode, **kwargs): + if mode not in ('constant', 'reflect', 'replicate', 'circular'): + raise AssertionError(f"Expected mode to be one of 'constant', 'reflect', 'replicate', 'circular', got {mode!r}") + if mode in ['reflect', 'replicate']: + cases: tuple = ( # ignore + ((1, 3), (1, 2)), + ((1, 3), (0, 1)), + ((0, 3, 3), (1, 2)), + ((0, 3, 3), (0, 1)), + ((1, 3, 3), (1, 2)), + ((1, 3, 3), (0, 1)), + ((1, 3, 3), (0, 2, 0, 1)), + ((0, 3, 3, 3), (0, 2, 0, 1)), + ((3, 3, 5, 5), (0, 2, 0, 1)), + ((3, 3, 5, 5), (1, 1, 1, 1, 1, 1)), + ((1, 3, 3, 3, 3), (1, 1, 1, 1, 1, 1)), + ((1, 3, 4, 4), (-1, 1, -2, 1)), + ) + elif mode == 'constant': + cases = ( + ((1, 3), (1, 2)), + ((1, 3), (0, 1)), + ((1, 3), (0, 2, 0, 1)), + ((5, 3), (-1, -2, 1, 1)), + ((0, 3, 3), (1, 2)), + ((0, 3, 3), (0, 1)), + ((0, 3, 3), (0, 2, 0, 1)), + ((0, 3, 3), (1, 1, 1, 1, 1, 1)), + ((1, 3, 3), (1, 2)), + ((1, 3, 3), (0, 1)), + ((1, 3, 3), (0, 2, 0, 1)), + ((1, 3, 3), (1, 1, 1, 1, 1, 1)), + ((0, 3, 3, 3), (1, 2)), + ((0, 3, 3, 3), (0, 1)), + ((0, 3, 3, 3), (0, 2, 0, 1)), + ((0, 3, 3, 3), (1, 1, 1, 1, 1, 1)), + ((3, 3, 5, 5), (1, 2)), + ((3, 3, 5, 5), (0, 1)), + ((3, 3, 5, 5), (0, 2, 0, 1)), + ((3, 3, 5, 5), (1, 1, 1, 1, 1, 1)), + ((1, 3, 3, 3, 3), (1, 2)), + ((1, 3, 3, 3, 3), (0, 1)), + ((1, 3, 3, 3, 3), (0, 2, 0, 1)), + ((1, 3, 3, 3, 3), (1, 1, 1, 1, 1, 1)), + ((1, 3, 4, 4), (-1, 1, -2, 1)), + ) + else: # mode == 'circular' + if dtype == torch.bool: + # test_dtypes fails on ASAN with for the case ab + # runtime error: load of value 190, which is not a valid value for type 'bool' + # Reference: https://github.com/pytorch/pytorch/pull/62814#issuecomment-894156562 + # Reference Issue: https://github.com/pytorch/pytorch/issues/63034 + cases = ( + ((2, 3, 3), (1, 2)), + ((1, 3, 3), (1, 2)), + ) + else: + cases = ( + ((0, 3, 3), (1, 2)), + ((0, 3, 3), (0, 1)), + ((1, 3, 3), (1, 2)), + ((1, 3, 3), (0, 1)), + ((0, 3, 3, 3), (0, 2, 0, 1)), + ((3, 3, 5, 5), (0, 2, 0, 1)), + ((1, 3, 3, 3, 3), (1, 1, 1, 1, 1, 1)), + ((1, 3, 4, 4), (-1, 1, -2, 1)), + ) + + make_inp = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + if mode == 'constant': + # Default args + yield SampleInput(make_inp((1, 3, 3)), args=((2, 2),)) + + if mode in ['reflect', 'replicate', 'circular']: + for shape, pad in cases: + yield SampleInput(make_inp(shape), args=(pad, mode)) + else: # mode == 'constant' + for pad_value in (1., 2.): + for shape, pad in cases: + yield SampleInput(make_inp(shape), args=(pad, mode, pad_value)) + +def sample_inputs_nn_pad_replicate_negative(op_info, device, dtype, requires_grad, **kwargs): + cases: tuple = ( + ((5, 3, 4, 4), (-4, 5, 0, 0)), + ((6, 2, 4, 4), (0, 0, 2, -4)), + ((5, 6, 4, 4), (5, -4, -4, 3)), + ((4, 2, 5, 5), (-2, -1, 4, 6)), + ((2, 6, 5, 5), (8, -1, -1, -3)), + ((8, 1, 5, 5), (-2, -1, -1, -3)), + ) + make_inp = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + for shape, pad in cases: + yield SampleInput(make_inp(shape), args=(pad, 'replicate')) + +def sample_inputs_constant_pad_nd(op_info, device, dtype, *args, **kwargs): + # Inherit sample inputs from nn.pad, but transform them to fit + # constant_pad_nd's interface + nn_samples = sample_inputs_nn_pad(op_info, device, dtype, *args, + mode='constant', **kwargs) + + # NOTE: primTorch is more strict about the type of the fill value argument + # So we must cast it to the correct dtype + from torch._prims_common import dtype_to_type + scalar_type = dtype_to_type(dtype) + + def drop_mode_argument(input, pad, mode=None, value=None): + if value is None: + return SampleInput(input, args=(pad,)) + else: + return SampleInput(input, args=(pad, scalar_type(value))) + + for sample in nn_samples: + yield drop_mode_argument(sample.input, *sample.args, **sample.kwargs) + +def sample_inputs_repeat_interleave(op_info, device, dtype, requires_grad, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(make_input(()), repeats=2) + yield SampleInput(make_input((2, 3, 4)), repeats=2) + yield SampleInput(make_input((2, 3, 4)), repeats=2, dim=1) + yield SampleInput(make_input((2, 3, 4)), repeats=torch.arange(3, device=device), dim=1) + yield SampleInput(make_input((4, 1)), repeats=torch.arange(4, device=device), dim=0, output_size=6) + + +def sample_inputs_stft(op_info, device, dtype, requires_grad, **kwargs): + def mt(shape, **kwargs): + return make_tensor(shape, device=device, dtype=dtype, + requires_grad=requires_grad, **kwargs) + + yield SampleInput(mt(100), n_fft=10, return_complex=True) + yield SampleInput(mt(100), n_fft=10, return_complex=False) + if dtype.is_complex: + yield SampleInput(mt(100), n_fft=10) + + for center in [False, True]: + yield SampleInput(mt(10), n_fft=7, center=center, return_complex=True) + yield SampleInput(mt((10, 100)), n_fft=16, hop_length=4, + center=center, return_complex=True) + + window = mt(16, low=.5, high=2.0) + yield SampleInput( + mt((2, 100)), kwargs=dict(n_fft=16, window=window, return_complex=True, center=center)) + yield SampleInput( + mt((3, 100)), kwargs=dict(n_fft=16, window=window, return_complex=True, center=center)) + if not dtype.is_complex: + yield SampleInput( + mt((10, 100)), n_fft=16, window=window, onesided=False, + return_complex=True) + + +def sample_inputs_istft(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def mt(shape, **kwargs): + real_shape = shape if dtype.is_complex else shape + (2,) + return make_arg(real_shape, **kwargs) + + yield SampleInput(mt((10, 2)), kwargs=dict(n_fft=10)) + yield SampleInput(mt((6, 3)), kwargs=dict(n_fft=6, onesided=False)) + yield SampleInput(mt((6, 4)), kwargs=dict(n_fft=10, onesided=True)) + + for center in [False, True]: + yield SampleInput(mt((10, 10, 6)), kwargs=dict(n_fft=10, center=center)) + yield SampleInput(mt((1, 9, 10)), kwargs=dict(n_fft=16, hop_length=4, center=center)) + + window = make_arg(10, low=.5, high=2.0) + yield SampleInput(mt((10, 10, 6)), kwargs=dict( + n_fft=10, window=window, center=center, return_complex=dtype.is_complex)) + yield SampleInput(mt((10, 10, 10)), kwargs=dict( + n_fft=10, window=window[:8], win_length=8, center=center, return_complex=True)) + + real_window = window if not dtype.is_complex else window.real + yield SampleInput(mt((10, 5, 6)), kwargs=dict(n_fft=8, window=real_window[:8], center=center)) + +def sample_inputs_ormqr(op_info, device, dtype, requires_grad, **kwargs): + # create a helper function wrapping `make_tensor` + make_input = partial(make_tensor, dtype=dtype, device=device, low=-1, high=1) + + batches = [(), (0, ), (2, ), (2, 1)] + ns = [5, 2, 0] + tf = [True, False] + for batch, (m, n), left, transpose in product(batches, product(ns, ns), tf, tf): + input = make_input((*batch, m, n)) + reflectors, tau = torch.geqrf(input) + reflectors.requires_grad_(requires_grad) + tau.requires_grad_(requires_grad) + other_matrix_shape = (m, n) if left else (n, m) + other = make_input((*batch, *other_matrix_shape), requires_grad=requires_grad) + yield SampleInput(reflectors, tau, other, left=left, transpose=transpose) + + +def sample_inputs_cholesky_solve(op_info, device, dtype, requires_grad=False, **kwargs): + cholesky_inverse_samples = sample_inputs_linalg_cholesky_inverse( + op_info, device, dtype, requires_grad=False + ) + + for sample in cholesky_inverse_samples: + psd_matrix = sample.input + sample.input = make_tensor(psd_matrix.shape, dtype=dtype, device=device, requires_grad=requires_grad, low=None, high=None) + sample.args = (psd_matrix.requires_grad_(requires_grad),) + yield sample + + +def sample_inputs_lu(op_info, device, dtype, requires_grad=False, **kwargs): + make_arg = partial(make_fullrank_matrices_with_distinct_singular_values, + dtype=dtype, device=device, requires_grad=requires_grad) + + # not needed once OpInfo tests support Iterables + batch_shapes = ((), (3,), (3, 3)) + for batch_shape, get_infos, size_delta in product(batch_shapes, (True, False), (-2, -1, 0, +1, +2)): + shape = batch_shape + (S + size_delta, S) + input = make_arg(*shape) + yield SampleInput(input, args=(True, get_infos)) + + +def sample_inputs_lu_unpack(op_info, device, dtype, requires_grad=False, **kwargs): + def out_fn(output): + return output[1], output[2] + + for lu_sample in sample_inputs_linalg_lu(op_info, device, dtype, requires_grad, **kwargs): + lu_data, pivots = torch.linalg.lu_factor(lu_sample.input) + lu_data.requires_grad_(requires_grad) + yield SampleInput(lu_data, pivots).with_metadata(output_process_fn_grad=out_fn) + + +def sample_inputs_roll(op_info, device, dtype, requires_grad=False, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + args = ((0, 0), (1, 2), (0, 2), (2, 0), (-1, 0), (10000, 1), (2,), ((1, 2, -1), (0, 1, 2))) + + for arg in args: + yield SampleInput(make_arg((0, 0, 0)), args=arg) + yield SampleInput(make_arg((S, S, S)), args=arg) + + # Scalar tensor + yield SampleInput(make_arg(()), args=(10, )) + +def error_inputs_roll(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + err_msg1 = "`shifts` required" + s1 = SampleInput(make_arg((S,)), ()) + yield ErrorInput(s1, error_regex=err_msg1) + + err_msg2 = ("shifts and dimensions must align") + s2 = SampleInput(make_arg((S, S)), (2, 1), 0) + yield ErrorInput(s2, error_regex=err_msg2) + + err_msg3 = ("out of range") + s3 = SampleInput(make_arg((S, )), 0, 2) + yield ErrorInput(s3, error_regex=err_msg3, error_type=IndexError) + + err_msg4 = ("Dimension specified as 0") + s4 = SampleInput(make_arg(()), 0, 0) + yield ErrorInput(s4, error_regex=err_msg4, error_type=IndexError) + +def sample_inputs_rot90(op_info, device, dtype, requires_grad=False, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + args = itertools.product(range(-5, 6), [(0, 1), (1, 2), (1, -1)]) + + yield SampleInput(make_arg((S, S, S))) + for arg in args: + yield SampleInput(make_arg((S, S, S)), args=arg) + + +def error_inputs_rot90(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + err_msg1 = "expected total rotation dims" + s1 = SampleInput(make_arg((S, S)), dims=(0,)) + yield ErrorInput(s1, error_regex=err_msg1) + + err_msg2 = "expected total dims >= 2" + s2 = SampleInput(make_arg((S,))) + yield ErrorInput(s2, error_regex=err_msg2) + + err_msg3 = "expected rotation dims to be different" + s3 = SampleInput(make_arg((S, S)), dims=(1, 1)) + yield ErrorInput(s3, error_regex=err_msg3) + + +def sample_inputs_std_var(op_info, device, dtype, requires_grad, **kwargs): + tensor_nd = partial(make_tensor, (S, S, S), device=device, dtype=dtype, + requires_grad=requires_grad) + tensor_1d = partial(make_tensor, (S,), device=device, dtype=dtype, + requires_grad=requires_grad) + + yield SampleInput(tensor_nd()) + yield SampleInput(tensor_nd(), dim=1) + yield SampleInput(tensor_nd(), dim=1, unbiased=True, keepdim=True) + yield SampleInput(tensor_1d(), dim=0, unbiased=True, keepdim=True) + yield SampleInput(tensor_1d(), dim=0, unbiased=False, keepdim=False) + + yield SampleInput(tensor_nd(), dim=(1,), correction=1.3) + yield SampleInput(tensor_nd(), dim=(1,), correction=S // 2) + yield SampleInput(tensor_nd(), dim=None, correction=0, keepdim=True) + yield SampleInput(tensor_nd(), dim=None, correction=None) + yield SampleInput(tensor_nd(), dim=None, correction=-1) + yield SampleInput(tensor_nd(), dim=None, correction=-5) + yield SampleInput(tensor_nd(), correction=0.5, keepdim=True) + yield SampleInput(tensor_nd(), correction=0, keepdim=True) + yield SampleInput(make_tensor(3, 4, 5, device=device, dtype=dtype, requires_grad=requires_grad), dim=-3) + + +def sample_inputs_std_var_unbiased(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, + requires_grad=requires_grad) + + # Test var_mean(Tensor self, bool unbiased=True) -> (Tensor, Tensor) + yield SampleInput(make_arg((S, S)), True) + yield SampleInput(make_arg((S,)), False) + + +def _generate_correlation_inputs(device, dtype, requires_grad, **kwargs): + shapes = [(2,), (1, 2), (3, 2), (2, 3)] + for shape in shapes: + yield make_tensor(shape, dtype=dtype, device=device, requires_grad=requires_grad) + + +def sample_inputs_corrcoef(op_info, device, dtype, requires_grad, **kwargs): + return (SampleInput(t) for t in _generate_correlation_inputs(device, dtype, requires_grad)) + +def sample_inputs_copysign(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_elementwise_binary(op_info, device, dtype, requires_grad, **kwargs) + if dtype.is_floating_point: + yield SampleInput(make_tensor(5, dtype=dtype, device=device, requires_grad=requires_grad), -3.14) + + +def sample_inputs_cov(op_info, device, dtype, requires_grad, **kwargs): + for t in _generate_correlation_inputs(device, dtype, requires_grad): + yield SampleInput(t) + num_observations = t.numel() if t.ndimension() < 2 else t.size(1) + fweights = make_tensor((num_observations,), dtype=torch.int, device=device, low=1, high=10) + aweights = make_tensor((num_observations,), dtype=torch.float, device=device, low=0, high=1, requires_grad=requires_grad) + for correction, fw, aw in product(range(num_observations), [None, fweights], [None, aweights]): + yield SampleInput(t.clone().requires_grad_(requires_grad), + correction=correction, fweights=fw, aweights=aw) + + +def error_inputs_cov(op_info, device, **kwargs): + a = torch.rand(S, device=device) + yield ErrorInput( + SampleInput(torch.rand(S, S, S, device=device)), + error_regex="expected input to have two or fewer dimensions") + yield ErrorInput( + SampleInput(a, fweights=torch.rand(S, S, device=device)), + error_regex="expected fweights to have one or fewer dimensions") + yield ErrorInput( + SampleInput(a, aweights=torch.rand(S, S, device=device)), + error_regex="expected aweights to have one or fewer dimensions") + yield ErrorInput( + SampleInput(a, fweights=torch.rand(S, device=device)), + error_regex="expected fweights to have integral dtype") + yield ErrorInput( + SampleInput(a, aweights=torch.tensor([1, 1], device=device)), + error_regex="expected aweights to have floating point dtype") + yield ErrorInput( + SampleInput(a, fweights=torch.tensor([1], device=device)), + error_regex="expected fweights to have the same numel") + yield ErrorInput( + SampleInput(a, aweights=torch.rand(1, device=device)), + error_regex="expected aweights to have the same numel") + yield ErrorInput( + SampleInput(a, fweights=torch.tensor([-1, -2, -3, -4 , -5], device=device)), + error_regex="fweights cannot be negative") + yield ErrorInput( + SampleInput(a, aweights=torch.tensor([-1., -2., -3., -4., -5.], device=device)), + error_regex="aweights cannot be negative") + + +def sample_inputs_permute(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases = [((1, 2, 3, 4), (0, 2, 3, 1)), + ((1, 2, 3, 4), (0, -2, -1, 1)), + ((), ()), + ((1, 2, 3, 4), (2, 1, 3, 0))] + + for shape, args in cases: + yield SampleInput(make_arg(shape), args=(args,)) + +def reference_inputs_permute(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_permute(op, device, dtype, requires_grad, **kwargs) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases = ( + ((), ()), + ((1,), (0,)), + ((2, 2), (1, 0)), + ((2, 2), (0, 1)), + ((2, 0, 1), (0, 2, 1)), + ((3, 4, 2), (2, 1, 0)), + ((3, 4, 2), (1, 0, 2)), + ((3, 4, 2), (0, 1, 2)), + ) + + # Adds tricky permutations and permutations with noncontiguity + for shape, permutation in cases: + for p in itertools.permutations(permutation): + a = make_arg(shape).permute(p) + yield SampleInput(a, args=(permutation,)) + + a = make_arg(shape, noncontiguous=True).permute(p) + yield SampleInput(a, args=(permutation,)) + +def error_inputs_softshrink(op, device, **kwargs): + yield ErrorInput(SampleInput(make_tensor((1,), dtype=torch.float, device=device), kwargs={"lambd": -0.5}), + error_regex=r"lambda must be in range \[0,.*input dtype.*found -0\.5") + +def sample_inputs_softshrink(op_info, device, dtype, requires_grad=False, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # The additional sample is to check additional values of lambd beyond the default + # value (what is already checked by sample_inputs_elementwise_unary) + for lbda in (0., 0.5): + yield SampleInput(make_arg(S, S), kwargs={"lambd": lbda}) + + yield from sample_inputs_elementwise_unary(op_info, device, dtype, requires_grad) + +def sample_inputs_hardshrink(op_info, device, dtype, requires_grad=False, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # The additional sample is to check additional values of lambd beyond the default + # value (what is already checked by sample_inputs_elementwise_unary) + # Note that unlike softshrink, lambd is allowed to be negative for hardshrink + for lbda in (-0.5, 0., 0.5): + yield SampleInput(make_arg(S, S), kwargs={"lambd": lbda}) + + yield from sample_inputs_elementwise_unary(op_info, device, dtype, requires_grad) + + +def sample_inputs_hardtanh(op_info, device, dtype, requires_grad=False, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # The additional sample is to check additional values of min_val and max_val beyond the default + # value (what is already checked by sample_inputs_elementwise_unary) + for max_val, min_val in ((0.5, -0.5), (0., 0.)): + yield SampleInput(make_arg(S, S), kwargs={"min_val": min_val, "max_val": max_val}) + + yield from sample_inputs_elementwise_unary(op_info, device, dtype, requires_grad) + +def error_inputs_hardtanh(op_info, device, **kwargs): + # Tests that hardtanh errors out when passed min_val > max_val. + yield ErrorInput(SampleInput(make_tensor((1,), dtype=torch.float, device=device), kwargs={"min_val": 0.5, "max_val": -0.5}), + error_type=ValueError, error_regex="min_val cannot be greater than max_val") + +def sample_inputs_einsum(op_info, device, dtype, requires_grad=False, **kwargs): + def c(t): + return t.clone().requires_grad_(requires_grad) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + x = make_arg((3,)) + y = make_arg((4,)) + A = make_arg((2, 3,)) + B = make_arg((1, 3,)) + C = make_arg((1, 2, 3,)) + D = make_arg((1, 3, 4,)) + E = make_arg((4, 4,)) + H = make_arg((3, 3,)) + I = make_arg((1, 3, 1,)) + + # Vector operations + yield SampleInput([c(x)], 'i->') # sum + yield SampleInput([c(x), c(y)], 'i,j->ij') # outer + + # Matrix operations + yield SampleInput([c(A)], "ij->i") # col sum + yield SampleInput([c(A), c(B)], "ij,kj->ik") # matmul + yield SampleInput([c(A), c(E)], "ij,Ab->ijAb") # matrix outer product + + # Tensor operations + yield SampleInput([c(C), c(D)], "aij,ajk->aik") # batch matmul + yield SampleInput([c(D), c(E)], "aij,jk->aik") # tensor matrix contraction + yield SampleInput([c(C), c(B)], "ijk,ik->j") # non contiguous + + # Test diagonals + yield SampleInput([c(I)], 'iji->j') # non-contiguous trace + + # Test ellipsis + yield SampleInput([c(H)], "i...->...") + yield SampleInput([c(C), c(x)], '...ik, ...j -> ij') + + +def sample_inputs_flip(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + sizes = ((S, M, S), (S, 0, M)) + all_dims = ((0, 1, 2), (0,), (0, 2), (-1,), ()) + + for size, dims in product(sizes, all_dims): + yield SampleInput(make_arg(size), kwargs={"dims": dims}) + +def sample_inputs_fliplr_flipud(op_info, device, dtype, requires_grad, **kwargs): + shapes = [ + (S, M, S), + (S, 0, M), + ] + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + return (SampleInput(make_arg(shape, low=None, high=None)) for shape in shapes) + +def error_inputs_fliplr(op, device, **kwargs): + yield ErrorInput(SampleInput(make_tensor((1,), dtype=torch.float, device=device)), + error_regex="Input must be >= 2-d.") + +def error_inputs_flipud(op, device, **kwargs): + yield ErrorInput(SampleInput(make_tensor((), dtype=torch.float, device=device)), + error_regex="Input must be >= 1-d.") + +def sample_inputs_clamp(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, low=None, high=None, requires_grad=requires_grad) + make_integral_arg = partial(make_tensor, dtype=torch.int32, device=device, low=None, high=None, requires_grad=False) + shape = (S, M, S) + + yield SampleInput(make_arg(shape), args=(make_arg(shape), make_arg(shape))) + yield SampleInput(make_arg(shape), args=(make_arg(shape[1:]), make_arg(shape[1:]))) + yield SampleInput(make_arg(shape), args=(make_arg((S, 1, S)),)) + yield SampleInput(make_arg(shape), args=(None, make_arg(shape))) + yield SampleInput(make_arg(shape), args=(make_arg(shape), None)) + # test type promotion + yield SampleInput(make_arg(shape), args=(make_integral_arg(shape), None)) + yield SampleInput(make_arg(shape), args=(make_arg(shape), make_integral_arg(shape))) + +def reference_inputs_elementwise_ternary(op, device, dtype, requires_grad, *, sample_inputs_func, supports_scalars=False, **kwargs): + yield from sample_inputs_func(op, device, dtype, requires_grad, **kwargs) + + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_scalar_tensor = partial(make_tensor, (), device='cpu', dtype=dtype, requires_grad=requires_grad) + supported_dtypes = op.supported_dtypes(device) + + # broadcasting and oncontiguous cases + cases = ( + ((4, 4), (4, 4), (4, 4)), + ((4, 4), (1, 4, 4), (4, 4)), + ((4, 4), (1, 4, 4), (4, 1, 4)), + ((4, 4, 1), (1, 4, 4), (4, 4)), + ((4, 1), (1, 4, 4), (1, 4)), + ((4, 4), (), (4, 4)), + ((4, 4), (), ()), + ((), (4, 4), (1, 4, 4)), + ) + + for a, b, c in cases: + yield SampleInput(make_arg(a), args=(make_arg(b), make_arg(c))) + yield SampleInput(make_arg(a, noncontiguous=True), + args=(make_arg(b).transpose(0, -1), make_arg(c, noncontiguous=True).transpose(0, -1))) + + # scalar cases + if supports_scalars: + cases = [ + ((), 1, 2,), + ((), 1., 2), + ((4, 4), 1., 2,), + ((3, 4), make_scalar_tensor(), make_scalar_tensor()), + ] + + if torch.complex64 in supported_dtypes: + cases.extend([ + ((3, 1, 4), complex(1, 2), 3.), + ]) + + for a, b, c in cases: + yield SampleInput(make_arg(a), args=(b, c)) + + # type promotion cases + # int x float + if torch.float in supported_dtypes and torch.long in supported_dtypes: + a = make_arg((), dtype=torch.long) + b = make_arg((1, 4), dtype=torch.float) + c = make_arg((3, 4)) + + cases = ( + (a, b, c), + (c, a, b), + ) + + for a, b, c in cases: + yield SampleInput(a, args=(b, c)) + + # NaN propagation + if dtype.is_floating_point or dtype.is_complex: + nan = float('nan') if dtype.is_floating_point else complex(float('nan'), float('nan')) + + a = make_arg((12,)) + a[4] = nan + a[7] = nan + b = make_arg((12,)) + b[1] = nan + b[7] = nan + c = make_arg((12,)) + c[9] = nan + + yield SampleInput(a, args=(b, c)) + + +def _clamp_min_numpy(a, min=None): + return np.maximum(a, min) + + +def _clamp_max_numpy(a, max=None): + return np.minimum(a, max) + + +def _clamp_numpy(a, min=None, max=None): + if min is None: + return np.minimum(a, max) + if max is None: + return np.maximum(a, min) + + return np.minimum(max, np.maximum(a, min)) + + +def sample_inputs_cumprod(op_info, device, dtype, requires_grad, **kwargs): + def make_arg(shape): + # shrink values to be in the interval [-1, +1] for better precision in gradgradcheck + return make_tensor(shape, dtype=dtype, device=device, low=-1, high=+1, requires_grad=requires_grad) + + def prod_zeros(dim_select): + if len(dim_select) != 2: + raise AssertionError(f"Expected len(dim_select) to be 2, got {len(dim_select)}") + result = make_arg(3 * (S,)) + result.narrow(dim_select[0], 0, 1).narrow(dim_select[1], 1, 1).zero_() + result.narrow(dim_select[0], 2, 1).narrow(dim_select[1], 3, 1).zero_() + result.narrow(dim_select[0], 4, 1).narrow(dim_select[1], 3, 1).zero_() + return result + + for dim in range(3): + yield SampleInput(make_arg((S, S, S)), args=(dim,)) + # Scalar tensors and empty tensor + for size in [(), (1,), (0,)]: + yield SampleInput(make_arg(size), args=(0,)) + + yield SampleInput(prod_zeros([0, 1]), args=(1,)) + yield SampleInput(prod_zeros([0, 2]), args=(1,)) + yield SampleInput(prod_zeros([1, 2]), args=(1,)) + + # test dtype kwarg + yield SampleInput(prod_zeros([1, 2]), args=(1,), kwargs={'dtype': dtype}) + +def sample_inputs_view_as_complex(op_info, device, dtype, requires_grad, **kwargs): + yield SampleInput(make_tensor((S, 2), dtype=dtype, device=device, requires_grad=requires_grad)) + +def sample_inputs_view_as_real(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + sizes = ((S, S), ()) + return (SampleInput(make_arg(size)) for size in sizes) + +def error_inputs_complex(op_info, device, is_ref=False, **kwargs): + make_arg = partial(make_tensor, dtype=torch.float32, device=device) + other_dtype = torch.float16 if device.startswith("mps") else torch.float64 + other_dtype_name = "Half" if device.startswith("mps") else "Double" + + if is_ref: + error_float = "Expected both inputs to be Half, Float or Double tensors but got torch.float32 and torch.int32" + error_dtype = "Expected object of scalar type torch.float32 but got scalar type torch.float64 for second argument" + error_out = "Expected out tensor to have dtype torch.complex128 but got torch.complex64 instead" + else: + error_float = "Expected both inputs to be Half, Float or Double tensors but got Float and Int" + error_dtype = f"Expected object of scalar type Float but got scalar type {other_dtype_name} for second argument" + error_out = f"Expected object of scalar type Complex{other_dtype_name} but got scalar type ComplexFloat for argument 'out'" + + yield ErrorInput(SampleInput(make_arg(M, S), make_arg(M, S, dtype=torch.int)), + error_type=RuntimeError, error_regex=error_float) + + yield ErrorInput(SampleInput(make_arg(M, S), make_arg(M, S, dtype=other_dtype)), + error_type=RuntimeError, error_regex=error_dtype) + + yield ErrorInput(SampleInput(make_arg(M, S, dtype=other_dtype), make_arg(M, S, dtype=other_dtype), + out=make_arg(M, S, dtype=torch.complex64)), + error_type=RuntimeError, error_regex=error_out) + +def sample_inputs_logaddexp(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + shape = (S, S) + yield SampleInput(make_arg(shape), make_arg(shape)) + +def sample_inputs_prod(op_info, device, dtype, requires_grad, **kwargs): + def make_arg(shape): + # shrink values to be in the interval [-1, +1] for better precision in gradgradcheck + return make_tensor(shape, dtype=dtype, device=device, low=-1, high=+1, requires_grad=requires_grad) + + def prod_single_zero(): + result = make_arg(2 * (S,)) + result[0, 1] = 0 + return result + + for sample in sample_inputs_cumprod(op_info, device, dtype, requires_grad): + # only Tensor, ignore other inputs + yield SampleInput(sample.input.clone().requires_grad_(requires_grad)) + yield sample + + # Generates samples with keepdim = True + for sample in sample_inputs_cumprod(op_info, device, dtype, requires_grad): + sample.kwargs['keepdim'] = True + yield sample + + yield SampleInput(prod_single_zero()) + yield SampleInput(make_arg((3, 3, 3)), args=(1,)) + yield SampleInput(make_arg((3, 3, 3)), args=(1,), kwargs={'keepdim': True}) + + yield SampleInput(make_arg((3, 0)), args=(1,)) + yield SampleInput(make_arg((3, 0)), args=(1,), kwargs={'keepdim': True}) + yield SampleInput(torch.tensor([2., 3, 0, 0], dtype=dtype, device=device, requires_grad=requires_grad)) + + # test zero scalar tensor + zero = make_arg(()) + zero.zero_() + yield SampleInput(zero.clone().requires_grad_(requires_grad)) + yield SampleInput(zero.clone().requires_grad_(requires_grad), args=(0,)) + yield SampleInput(zero.clone().requires_grad_(requires_grad), + args=(0,), + kwargs={'keepdim': True}) + +def error_inputs_neg(op_info, device, **kwargs): + si = SampleInput(torch.tensor((False, True), device=device)) + msg = ("Negation, the `\\-` operator, on a bool tensor is not supported." + " If you are trying to invert a mask, use the `\\~` or" + " `logical_not\\(\\)` operator instead.") + yield ErrorInput(si, error_regex=msg) + +def sample_inputs_diag(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, low=None, high=None) + yield SampleInput(make_arg(M)) + + tensors = ( + make_arg((M, M)), + make_arg((3, 5)), + make_arg((5, 3)), + ) + + args = ((), (2,), (-2,), (1,), (2,)) + + for tensor, arg in product(tensors, args): + yield SampleInput(tensor.clone().requires_grad_(requires_grad), *arg) + +def reference_inputs_diagonal_diag_embed(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_diagonal_diag_embed( + op_info, device, dtype, requires_grad, **kwargs) + + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + shapes1d = ((0,), (1,)) + shapes2d = ((L, M),) + shapes3d = ((L, M, S),) + + kwargs1d = {} + + kwargs2d = ( + # dim1 > dim2 is allowed + dict(dim1=1, dim2=0), + # negative dims are allowed + dict(dim1=-2, dim2=-1), + # one dim negative and the other nonnegative is allowed + dict(dim1=-1, dim2=0), + # out of bounds offset should return an empty tensor in diagonal and + # offset the diagonal in diag_embed + dict(offset=100), + ) + + kwargs3d = kwargs2d + ( + # make sure we can use non-sequential dims + dict(offset=-1, dim1=0, dim2=2), + ) + + samples1d = product(shapes1d, kwargs1d) + samples2d = product(shapes2d, kwargs2d) + samples3d = product(shapes3d, kwargs3d) + + for shape, kwargs in chain(samples1d, samples2d, samples3d): + if 'diagonal' in op_info.name: + # these are error inputs for diagonal + if shape in ((0,), (1,)): + continue + yield SampleInput(input=make_arg(shape), kwargs=kwargs) + + +def sample_inputs_diagonal_scatter(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + # Shapes for 2D Tensors + shapes_2d = ((M, M), (3, 5), (5, 3)) + + # Shapes for 3D Tensors + shapes_3d = ((M, M, M),) + + args_2d = ((), (2,), (-2,), (1,)) + args_3d = ((1, 1, 2), (2, 0, 1), (-2, 0, 1)) + + for input_shape, arg in chain(product(shapes_2d, args_2d), product(shapes_3d, args_3d)): + input_ = make_arg(input_shape) + # We can programmatically figure out the right shape for src: + # It should be the same size as input.diagonal(other_args...) + if not isinstance(arg, tuple): + arg_tuple = (arg,) + else: + arg_tuple = arg + src_shape = input_.diagonal(*arg_tuple).size() + src = make_arg(src_shape) + yield SampleInput(input_, args=(src, *arg_tuple)) + + +def sample_inputs_to_sparse(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(make_arg((S, S))).with_metadata(output_process_fn_grad=lambda x: x.to_dense()) + yield SampleInput(make_arg((S, S)), 1).with_metadata(output_process_fn_grad=lambda x: x.to_dense()) + +def sample_inputs_cross_entropy(op_info, device, dtype, requires_grad, **kwargs): + batch_size, num_classes = shape = (2, 3) + reductions = ("mean", "sum", "none") + + input_shape_and_kwargs: list[tuple[tuple[int, ...], dict[str, Any]]] = [ + (shape, {}), + ((*shape, 1), {}), + ((*shape, 1, 2), {}), + ((*shape, 1, 2, 3), {}), + *[(shape, dict(reduction=reduction)) for reduction in reductions], + *[ + ( + shape, + dict( + weight=make_tensor((num_classes,), device=device, dtype=dtype), + reduction=reduction, + ), + ) + for reduction in reductions + ], + (shape, dict(ignore_index=1)), + ] + + for (input_shape, kwargs), probabilities_target in itertools.product(input_shape_and_kwargs, (False, True)): + input = make_tensor(input_shape, device=device, dtype=dtype, requires_grad=requires_grad) + + if probabilities_target: + # ignore_index is not supported for probabilities target + if "ignore_index" in kwargs: + continue + + target = make_tensor( + input_shape, + low=0, + high=1, + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + else: + target = make_tensor( + (batch_size, *input_shape[2:]), + low=0, + high=num_classes, + device=device, + dtype=torch.long, + ) + + if "ignore_index" in kwargs and torch.all(target == kwargs["ignore_index"]): + # make sure at least one item in target is not ignored + target[0] = random.sample(sorted(set(range(num_classes)) - {kwargs["ignore_index"]}), 1)[0] + + yield SampleInput(input, target, **kwargs) + + +def sample_inputs_logit(op_info, device, dtype, requires_grad, **kwargs): + low, high = op_info.domain + + # Note: Operator is very sensitive at points near the + # start and end of domain and leads to NaN for float16 + # if domain_eps is 1e-5. + if dtype.is_floating_point or dtype.is_complex: + domain_eps = op_info._domain_eps if dtype != torch.float16 else 3e-2 + + low = low + domain_eps + high = high - domain_eps + + make_arg = partial(make_tensor, dtype=dtype, device=device, low=low, high=high, requires_grad=requires_grad) + + yield SampleInput(make_arg((S, S, S))) + yield SampleInput(make_arg((S, S, S)), 0.2) + yield SampleInput(make_arg(())) + yield SampleInput(make_arg(()), 0.2) + +def sample_inputs_isin(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + # isin has two paths based on the size of elements and test_elements. + # if elements.numel() < 10 * pow(test_elements.numel(), 0.145): + yield SampleInput(make_arg((L,)), args=(make_arg((S,)),)) + # else: + yield SampleInput(make_arg((S,)), args=(make_arg((L,)),)) + +def sample_inputs_masked_scatter(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(make_arg((S, S)), args=(torch.randn(S, S, device=device) > 0, make_arg((S, S)))) + yield SampleInput(make_arg((S, S)), args=(torch.randn((S,), device=device) > 0, make_arg((S, S)))) + yield SampleInput(make_arg((S, S)), args=(bernoulli_scalar().to(device), make_arg((S, S)))) + yield SampleInput(make_arg((S,)), + args=(torch.randn(S, S, device=device) > 0, make_arg((S, S))), + broadcasts_input=True) + +def error_inputs_masked_scatter(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float) + for mask_dtype in [torch.float, torch.uint8]: + yield ErrorInput(SampleInput(make_arg(1, 3), args=(torch.ones(1, 3, device=device, dtype=mask_dtype), + make_arg(3, 4))), + error_regex=r"masked_scatter_ only supports boolean masks") + +def sample_inputs_masked_fill(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(make_arg((S, S)), args=(torch.randn(S, S, device=device) > 0, 10)) + yield SampleInput(make_arg((S, S)), args=(torch.randn(S, S, device=device) > 0, make_arg(()))) + yield SampleInput(make_arg((S, S)), args=(torch.randn(S, device=device) > 0, 10)) + yield SampleInput(make_arg(()), args=(torch.randn((), device=device) > 0, 10)) + yield SampleInput(make_arg(()), args=(torch.randn((), device=device) > 0, make_arg(()))) + yield SampleInput(make_arg((S, S)), args=(torch.randn((), device=device) > 0, 10)) + + yield SampleInput(make_arg((S,)), + args=(torch.randn(S, S, device=device) > 0, make_arg(())), + broadcasts_input=True) + yield SampleInput(make_arg((S,)), + args=(torch.randn(S, S, device=device) > 0, 10), + broadcasts_input=True) + + if torch.device(device).type == 'cuda': + # `self` and `mask` on CUDA but `value` is a CPU scalar tensor. + yield SampleInput(make_arg((S, S)), + args=(torch.randn(S, S, device=device) > 0, + make_tensor((), device="cpu", dtype=dtype))) + +def error_inputs_masked_fill(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float, requires_grad=False) + # `value` is not a 0-D tensor. + yield ErrorInput(SampleInput(make_arg((2, 2)), args=(make_arg(()) > 0, make_arg((1,)))), + error_regex="only supports a 0-dimensional value tensor, but got tensor with 1 dimension") + # downcasting complex value (scalar overload) + yield ErrorInput(SampleInput(make_arg((2, 2)), args=(make_arg(()) > 0, 1j)), + error_regex=r"value cannot be converted to type .* without overflow") + # downcasting complex value (tensor overload) + yield ErrorInput(SampleInput(torch.ones(2, dtype=torch.long, device=device), + args=(make_arg(()) > 0, torch.tensor(1j, device=device))), + error_regex=r"value cannot be converted to type .* without overflow") + + if torch.device(device).type == 'cuda': + # `self` and `mask` on CPU but `value` is a CUDA scalar tensor. + yield ErrorInput(SampleInput(torch.randn((S, S), device='cpu'), + args=(torch.randn(S, S, device='cpu') > 0, + torch.randn((), device='cuda'))), + error_regex=r"to be on same device") + + +def sample_inputs_masked_select(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, low=None, high=None) + + yield SampleInput(make_arg((M, M)), torch.randn(M, M, device=device) > 0) + + yield SampleInput(make_arg((M, M)), torch.randn((M,), device=device) > 0) + yield SampleInput(make_arg((M,)), torch.randn((M, M), device=device) > 0) + + yield SampleInput(make_arg((M, 1, M)), torch.randn((M, M), device=device) > 0) + + yield SampleInput(make_arg(()), torch.tensor(1, device=device, dtype=torch.bool)) + + yield SampleInput(make_arg((M, M)), torch.tensor(1, device=device, dtype=torch.bool)) + + yield SampleInput(make_arg(()), torch.randn((M, M), device=device) > 0) + +def sample_inputs_matrix_exp(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(make_arg((S, S))) + yield SampleInput(make_arg((S, S, S))) + +def sample_inputs_matmul(op_info, device, dtype, requires_grad, is_rmatmul=False, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, low=None, + high=None, requires_grad=requires_grad) + test_cases = (((L,), (L,)), + ((S, M), (M,)), + ((M,), (M, S)), + ((S, M), (M, S)), + ((S, 0), (0, M)), + ((S, S, M), (M,)), + ((S, S, M), (M, S)), + ((S, S, 0), (0, S)), + ((M,), (S, M, S)), + ((S, M), (S, M, S)), + ((0, 0), (S, 0, 0)), + ((S, S, M, M), (S, S, M, S)), + ((S, S, M, M), (M,)), + ((M,), (S, S, M, S)), + ((S, S, S), (1, S, S)) + ) + for lhs_shape, rhs_shape in test_cases: + lhs = make_arg(lhs_shape) + rhs = make_arg(rhs_shape) + if not is_rmatmul: + yield SampleInput(lhs, rhs) + else: + yield SampleInput(rhs, lhs) + + +def sample_inputs_meshgrid(op_info: OpInfo, device: torch.device, dtype: torch.dtype, + requires_grad: bool, + *, variant: str, **kwargs) -> list[SampleInput]: + if variant == 'variadic': + def make_inputs( + tensors: list[torch.Tensor]) -> tuple[torch.Tensor | list[torch.Tensor], + tuple[torch.Tensor, ...]]: + return tensors + elif variant == 'list': + def make_inputs( + tensors: list[torch.Tensor]) -> tuple[torch.Tensor | list[torch.Tensor], + tuple[torch.Tensor, ...]]: + return [tensors] + else: + raise ValueError( + 'Unsupported variant, must be one of {"variadic", "list"}. ' + f'Got "{variant}".') + + SCALAR = torch.Size([]) + VECTOR = torch.Size([3]) + test_cases: list[list[torch.Size]] = [ + [SCALAR], + [VECTOR], + [VECTOR, SCALAR], + [VECTOR, SCALAR, VECTOR], + [VECTOR, SCALAR, VECTOR, SCALAR], + ] + + for shapes, indexing in itertools.product(test_cases, {'xy', 'ij'}): + args = make_inputs( + [make_tensor(shape, dtype=dtype, device=device, requires_grad=requires_grad) + for shape in shapes]) + yield SampleInput(*args, indexing=indexing) + + +def sample_inputs_mvlgamma(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + tensor_shapes = ((S, S), ()) + ns = (1, 2, 3, 4, 5) + + # Since the accepted lower bound for input + # to mvlgamma depends on `p` argument, + # the following function computes the lower bound + # which we pass to `make_tensor`. + def compute_min_val(p): + return (p - 1.) / 2 + + for shape, n in product(tensor_shapes, ns): + min_val = compute_min_val(n) + if not dtype.is_floating_point: + # Round-up minimum value for integral dtypes + min_val += 1 + else: + min_val += 2 * torch.finfo(dtype).eps + yield SampleInput(make_arg(shape, low=min_val), args=(n,)) + + +# Since `mvlgamma` has multiple entries, +# there are multiple common skips for the additional +# entries. Following function is a helper to that end. +def skips_mvlgamma(skip_redundant=False): + skips = ( + # outside domain values are hard error for mvlgamma op. + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_float_domains'), + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', + 'test_reference_numerics_extremal'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=(torch.float16, torch.int8)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_small', + dtypes=(torch.int8,)), + ) + if skip_redundant: + # Redundant tests + skips = skips + ( # type: ignore[assignment] + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestJit'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon'), + ) + return skips + + +# To test reference numerics against multiple values of argument `p`, +# we make multiple OpInfo entries with each entry corresponding to different value of p. +# We run the op tests from test_ops.py only for `p=1` to avoid redundancy in testing. +def make_mvlgamma_opinfo(variant_test_name, domain, skips, sample_kwargs): + return UnaryUfuncInfo('mvlgamma', + ref=reference_mvlgamma if TEST_SCIPY else None, + aliases=('special.multigammaln',), + variant_test_name=variant_test_name, + domain=domain, + decorators=(precisionOverride({torch.float16: 5e-2}),), + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_mvlgamma, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + skips=skips, + sample_kwargs=sample_kwargs) + + +def sample_inputs_cumulative_ops(op_info, device, dtype, requires_grad, supports_dtype_kwargs=True, **kwargs): + def _make_tensor_helper(shape, low=None, high=None): + return make_tensor(shape, dtype=dtype, device=device, low=low, high=high, requires_grad=requires_grad) + + yield SampleInput(_make_tensor_helper((S, S, S)), 0) + yield SampleInput(_make_tensor_helper((S, S, S)), 1) + yield SampleInput(_make_tensor_helper(()), 0) + + if supports_dtype_kwargs: + # NOTE: if `dtype` is not same as input, then inplace variants fail with + # `provided dtype must match the dtype of self tensor in cumsum` + yield SampleInput(_make_tensor_helper((S, S, S)), 1, dtype=dtype) + + +def sample_inputs_unfold(op_info, device, dtype, requires_grad, **kwargs): + test_cases = ( + ((), (0, 1, 1)), + ((S, S, S, S), (0, 3, 1)), + ((S, S, S, S), (1, 3, 1)), + ((S, S, S, S), (2, 3, 1)), + ((S, S, S, S), (3, 3, 1)), + ((S, S, S, S), (0, 3, 2)), + ((S, S, S, S), (1, 3, 2)), + ((S, S, S, S), (2, 3, 2)), + ((S, S, S, S), (3, 3, 2)), + ((S, S, S, S), (0, 4, 1)), + ((S, S, S, S), (1, 4, 1)), + ((S, S, S, S), (2, 4, 1)), + ((S, S, S, S), (3, 4, 1)), + ((M,), (0, 3, 1)), + ((M,), (0, 3, 2)), + ((M,), (0, 3, 3)), + ((1000,), (0, 3, 11)), + ((1000,), (0, 2, 27)), + ((10, 10), (0, 1, 2)), + ((10, 10), (1, 2, 3)), + ((10, 10), (1, 2, 2)), + ((S, S, S), (2, 3, 2)), + ) + + for shape, arguments in test_cases: + yield SampleInput(make_tensor(shape, dtype=dtype, device=device, + low=None, high=None, + requires_grad=requires_grad), + *arguments) + +def sample_inputs_split(op_info, device, dtype, requires_grad, *, list_args=False, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + if list_args: + cases = ( + ((S, S, S), (torch.Size([int(S / 3), S - int(S / 3) * 2, int(S / 3)]),)), + ((S, S, S), (torch.Size([int(S / 2), S - int(S / 2) * 2, int(S / 2)]), 2),), + ((S, S, S), (torch.Size([int(S / 2), S - int(S / 2) * 2, int(S / 2)]), -2),) + ) + else: + cases = ( # type: ignore[assignment] + ((S, S, S), (2,)), + ((S, S, S), (S, 1)), + ) + + for shape, args in cases: + yield SampleInput(make_arg(shape), args=args) + + +def sample_inputs_split_with_sizes(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases = (((S, S, S), (torch.Size([int(S / 3), S - int(S / 3) * 2, int(S / 3)]),)), + ((S, S, S), (torch.Size([int(S / 3), S - int(S / 3), 0]),)), + ((S, S, S), (torch.Size([int(S / 3), S - int(S / 3) * 2, int(S / 3)]), 2)), + ((S, S, S), (torch.Size([int(S / 3), S - int(S / 3) * 2, int(S / 3)]), -2)), + ) + + for shape, args in cases: + yield SampleInput(make_arg(shape), args=args) + + +def sample_inputs_msort(op_info, device, dtype, requires_grad, **kwargs): + def apply_grad(t): + if dtype in floating_types_and(torch.float16, torch.bfloat16): + t.requires_grad_(requires_grad) + + def large_1d_unique(dtype, device): + res = torch.randperm(L * L * L, dtype=torch.int64, device=device) + res = res.to(dtype) + apply_grad(res) + return res + + # Test case for large tensor. + yield SampleInput(large_1d_unique(dtype, device)) + + yield SampleInput(make_tensor((S, M, S), dtype=dtype, device=device, + low=None, high=None, + requires_grad=requires_grad)) + +def sample_inputs_lerp(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + # no broadcast + yield SampleInput(make_arg((S, S)), make_arg((S, S)), 0.4) + # broadcast rhs + yield SampleInput(make_arg((S, S)), make_arg((S,)), 0.4) + # scalar tensor + yield SampleInput(make_arg(()), make_arg(()), 0.4) + # broadcast rhs scalar-tensor + yield SampleInput(make_arg((S, S)), make_arg(()), 0.4) + # broadcast rhs with weight tensor + yield SampleInput(make_arg((S, S)), make_arg((S,)), make_arg((S, S))) + # broadcast rhs and weight tensor + yield SampleInput(make_arg((S, S)), make_arg((S, 1)), make_arg((S,))) + # broadcast lhs + yield SampleInput(make_arg((S,)), make_arg((S, S)), 0.4).with_metadata(broadcasts_input=True) + # scalar broadcast_lhs + yield SampleInput(make_arg(()), make_arg((S, S)), 0.4).with_metadata(broadcasts_input=True) + # broadcast all + yield SampleInput(make_arg((S, 1)), make_arg((S, S)), 0.4).with_metadata(broadcasts_input=True) + # tensor broadcast all + yield SampleInput(make_arg((S, 1)), make_arg((S, S)), make_arg((S, 1))).with_metadata( + broadcasts_input=True) + # no broadcast with weight tensor + yield SampleInput(make_arg((S, S)), make_arg((S, S)), make_arg((S, S))) + # broadcast lhs with weight tensor + yield SampleInput(make_arg((S,)), make_arg((S, S)), make_arg((S, S))).with_metadata( + broadcasts_input=True) + # broadcast lhs and weight tensor + yield SampleInput(make_arg((S,)), make_arg((S, S, S)), make_arg((S, S))).with_metadata( + broadcasts_input=True) + # broadcast lhs and weight tensor variant + yield SampleInput(make_arg((S, S)), make_arg((S, S, S)), make_arg((S,))).with_metadata( + broadcasts_input=True) + + if dtype.is_complex: + # no broadcast + yield SampleInput(make_arg((S, S)), make_arg((S, S)), 0.4j) + yield SampleInput(make_arg((S, S)), make_arg((S, S)), 1.2 + 0.1j) + # broadcast rhs + yield SampleInput(make_arg((S, S)), make_arg((S,)), 0.4j) + yield SampleInput(make_arg((S, S)), make_arg((S, S)), 5.4 + 9j) + # scalar tensor + yield SampleInput(make_arg(()), make_arg(()), 0.4j) + yield SampleInput(make_arg(()), make_arg(()), 6.1 + 0.004j) + # broadcast rhs scalar-tensor + yield SampleInput(make_arg((S, S)), make_arg(()), 0.4j) + yield SampleInput(make_arg((S, S)), make_arg(()), 1 + 2j) + +def sample_inputs_tensordot(self, device, dtype, requires_grad, **kwargs): + cases = ( + ((2, 2, 2), (2, 2, 2), (2)), + ((2, 2, 1), (2, 1, 2), ([0, 1], [2, 0])), + ((1, 1, 1), (2, 1, 2), ([0, 1], [2, 0])), + ) + for first_shape, second_shape, dims in cases: + yield SampleInput(make_tensor(first_shape, dtype=dtype, device=device, + requires_grad=requires_grad, low=-1, high=+2), + make_tensor(second_shape, dtype=dtype, device=device, + requires_grad=requires_grad, low=-1, high=+2), + dims=dims) + +def sample_inputs_kron(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad, low=None, high=None) + test_cases = ( + ((S, S), (M, L)), + ) + + for input_shape, other_shape in test_cases: + input = make_arg(input_shape) + other = make_arg(other_shape) + yield SampleInput(input, other) + +def sample_inputs_inner(self, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + yield SampleInput(make_arg(S), make_arg(S)) + yield SampleInput(make_arg(), make_arg(S, S)) + +def sample_inputs_scatter(op_info, device, dtype, requires_grad, **kwargs): + def _tensor(shape, dtype=dtype, low=None, high=None): + return make_tensor(shape, dtype=dtype, device=device, low=low, high=high, requires_grad=requires_grad) + + def _gather(shape, index_dim, max_indices): + return gather_variable(shape, index_dim, max_indices, device=device) + + zero = torch.tensor(0, dtype=torch.long, device=device) + test_cases = ( + (_tensor((M, S)), (0, _gather((S, S), 1, M), _tensor((S, S)))), + (_tensor((M, S)), (0, _gather((S, S), 1, M).to(torch.int32), _tensor((S, S)))), + (_tensor((M, S)), (1, _gather((S, S), 0, S), _tensor((S, S)))), + (_tensor((M, S)), (-1, _gather((S, S), 0, S), _tensor((S, S)))), + (_tensor((M, S)), (0, _gather((M, S // 2), 1, M), _tensor((M, S // 2)))), + (_tensor((M, S)), (1, _gather((M, S // 2), 0, S), _tensor((M, S // 2)))), + (_tensor((M, S)), (-1, _gather((M, S // 2), 0, S), _tensor((M, S // 2)))), + (_tensor(()), (0, zero.detach().clone(), _tensor(()))), + (_tensor(()), (0, zero.detach().clone(), 2.5)), + ) + + for tensor, args in test_cases: + yield SampleInput(tensor, *args) + + if not requires_grad: + yield SampleInput(tensor.detach().clone(), *args, reduce='add') + + if dtype.is_floating_point: + yield SampleInput(tensor.detach().clone(), *args, reduce='multiply') + +def sample_inputs_scatter_add(op_info, device, dtype, requires_grad, **kwargs): + def _tensor(shape, dtype=dtype, low=None, high=None): + return make_tensor(shape, dtype=dtype, device=device, low=low, high=high, requires_grad=requires_grad) + + def _gather(shape, index_dim, max_indices): + return gather_variable(shape, index_dim, max_indices, device=device) + + zero = torch.tensor(0, dtype=torch.long, device=device) + yield SampleInput(_tensor((M, S)), 0, _gather((S, S), 1, M), _tensor((S, S))) + yield SampleInput(_tensor((M, S)), 1, _gather((S, S), 0, S), _tensor((S, S))) + yield SampleInput(_tensor((M, S)), -1, _gather((S, S), 0, S), _tensor((S, S))) + yield SampleInput(_tensor((M, S)), 0, _gather((M, S // 2), 1, M), _tensor((M, S // 2))) + yield SampleInput(_tensor((M, S)), 1, _gather((M, S // 2), 0, S), _tensor((M, S // 2))) + yield SampleInput(_tensor((M, S)), -1, _gather((M, S // 2), 0, S), _tensor((M, S // 2))) + yield SampleInput(_tensor(()), 0, zero.detach().clone(), _tensor(())) + +def sample_inputs_scatter_reduce(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + gather = partial(gather_variable, device=device) + + zero = torch.tensor(0, dtype=torch.long, device=device) + test_cases = ( + ((M, S), 0, gather((S, S), 1, M), (S, S)), + ((M, S), 1, gather((S, S), 0, S), (S, S)), + ((M, S), -1, gather((S, S), 0, S), (S, S)), + ((M, S), 0, gather((M, S // 2), 1, M), (M, S // 2)), + ((M, S), 1, gather((M, S // 2), 0, S), (M, S // 2)), + ((M, S), -1, gather((M, S // 2), 0, S), (M, S // 2)), + ((), 0, zero.detach().clone(), ()), + ) + + reduce = op_info.variant_test_name + for (inp_shape, dim, index, src_shape), include_self in product(test_cases, [False, True, False]): + yield SampleInput(make_arg(inp_shape), + args=(dim, index, make_arg(src_shape), reduce), + kwargs={'include_self': include_self}) + + + # Sample inputs to test edge cases for backward + # Check that gradients are propagated correctly for prod when zeros in self/src are reduced + if requires_grad and reduce == 'prod': + # This sample tests gradients for the following cases + # (a) 1 zero reduced (from src (self[0, 1], self[1, 1]), from self (self[0, 0], self[2, 0])) + # (b) 2 zeros reduced (1 from src and 1 from self (self[1, 0]) + # (c) no zeros reduced (self([2, 1])) + # (d) 2 zeros reduced (both from src) is tested in test/test_autograd.py + # test_scatter_index_reduce_prod_gradgrad_error as this case is not supported for gradgrad + input = torch.tensor([[0, 13], [0, 17], [0, 19]], dtype=dtype, device=device, requires_grad=requires_grad) + src = torch.tensor([[0, 1, 2, 3], [0, 4, 0, 1], [2, 3, 5, 6]], dtype=dtype, device=device, requires_grad=requires_grad) + idx = torch.tensor([[1, 1, 0, 0], [0, 0, 1, 1], [0, 0, 0, 1]], dtype=torch.long, device=device) + + yield SampleInput(input, + args=(1, idx, src, reduce), + kwargs={'include_self': True}) + +def sample_inputs_segment_reduce(op_info, device, dtype, requires_grad, *, mode='lengths', **kwargs): + def _tensor(shape, dtype=dtype, low=None, high=None): + return make_tensor(shape, dtype=dtype, device=device, low=low, high=high, requires_grad=requires_grad) + + test_cases = ( + # inp_shape, dim, lengths, unsafe + ((S,), 0, [0, 1, 2, 2], False), + ((S,), 0, [0, 1, 2, 2], True), + ((S,), 0, [2, 0, 3, 0], False), + ((S, S), 0, [0, 1, 2, 2], False), + # test when lengths do not sum to dim size + ((M, S, S), 0, [1, 2, 0, 6, 0], True), + # test for higher dimensions + ((S, S), 1, [[0, 1, 2, 2] for _ in range(S)], False), + ((S, S), 1, [[2, 0, 3, 0], [0, 1, 2, 2], [3, 0, 2, 0], [1, 1, 1, 2], [0, 1, 2, 2]], False), + ((S, S, S), 1, [[0, 1, 2, 2] for _ in range(S)], False), + ((S, S, S), 1, [[2, 0, 3, 0], [0, 1, 2, 2], [3, 0, 2, 0], [1, 1, 1, 2], [0, 1, 2, 2]], False), + ) + + reductions = ["max", "mean", "min", "sum", "prod"] + for args, reduce, initial in product(test_cases, reductions, [1, 2]): + inp_shape, dim, lengths, unsafe = args + lengths_t = torch.tensor(lengths, dtype=torch.long, device=device) + sample_input_kwargs = {'axis': dim, 'unsafe': unsafe, 'initial': initial} + if mode == 'lengths': + sample_input_kwargs['lengths'] = lengths_t + elif mode == 'offsets': + zeros_shape = list(lengths_t.shape) + zeros_shape[dim] = 1 + offsets_t = torch.cat((lengths_t.new_zeros(zeros_shape), lengths_t), dim).cumsum_(dim) + sample_input_kwargs['offsets'] = offsets_t + else: + raise RuntimeError(f"mode most be one of 'offsets' or 'lengths' got '{mode}'.") + yield SampleInput(_tensor(inp_shape), + args=(reduce,), + kwargs=sample_input_kwargs) + + +def sample_inputs_ravel(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, + low=None, high=None, requires_grad=requires_grad) + yield SampleInput(make_arg((S, S, S))) + yield SampleInput(make_arg(())) + yield SampleInput(make_arg((S, S, S), noncontiguous=True)) + +def sample_inputs_unravel_index(op_info, device, dtype, requires_grad, **kwargs): + yield SampleInput( + torch.tensor( + [[3, 8, 13], [0, 5, 10]], + device=device, + dtype=dtype), + (4, 5)) + yield SampleInput( + torch.tensor([[3, 8, 13], [0, 5, 10]], device=device, dtype=dtype), + (4, 2**30)) + yield SampleInput( + torch.tensor([[3, 8, 13], [0, 5, 10]], device=device, dtype=dtype), + (2**30, 4)) + yield SampleInput( + torch.tensor(2, device=device, dtype=dtype), + (2, 2)) + max_val = 2**(8 * dtype.itemsize - (1 if dtype.is_signed else 0)) - 1 + yield SampleInput( + torch.tensor(max_val - 1, device=device, dtype=dtype), + (1, max_val)) + yield SampleInput( + torch.tensor([22, 41, 37], device=device, dtype=dtype), + (7, 6)) + yield SampleInput( + torch.tensor(min(1621, max_val), device=device, dtype=dtype), + (6, 7, 8, 9)) + yield SampleInput( + torch.tensor([], device=device, dtype=dtype), + (10, 3, 5)) + yield SampleInput( + torch.tensor( + [[1, 0, 1, 2, 3, 4], [1, 6, 1, 3, 2, 0]], + device=device, + dtype=dtype), + (5, 8)) + yield SampleInput( + torch.tensor( + [[1, 0, 1, 2, 3, 4], [1, 6, 1, 3, 2, 0], [1, 3, 1, 0, 9, 5]], + device=device, + dtype=dtype), + (5, 8, 10)) + yield SampleInput( + torch.tensor(0, device=device, dtype=dtype), + ()) + + a = np.array([[2, 4, 5, 6], [7, 8, 1, 15]]) + b = np.array([[3, 2, 7, 6], [10, 12, 8, 9]]) + _, i1, i2 = np.intersect1d(a, b, assume_unique=True, return_indices=True) + yield SampleInput(torch.tensor(i1, device=device, dtype=dtype), a.shape) + yield SampleInput(torch.tensor(i2, device=device, dtype=dtype), b.shape) + + a = np.array([[2, 4, 5, 6, 6], [4, 7, 8, 7, 2]]) + b = np.array([[3, 2, 7, 7], [10, 12, 8, 7]]) + _, i1, i2 = np.intersect1d(a, b, return_indices=True) + yield SampleInput(torch.tensor(i1, device=device, dtype=dtype), a.shape) + yield SampleInput(torch.tensor(i2, device=device, dtype=dtype), b.shape) + + +def sample_inputs_tril_triu(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + cases = (((M, M), ()), + ((M, M), (2,),), + ((M, S), ()), + ((M, S), (-1,)), + ((M, M), (2,),), + ((S, M, S), ()), + ((S, M, S), (2,)), + ((3, 3, S, S), ()),) + + for shape, args in cases: + yield SampleInput(make_arg(shape), args=args) + +def error_inputs_tril_triu(opinfo, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + # error inputs for input.ndim <= 2 + yield ErrorInput(SampleInput(make_arg((4,))), error_regex="input tensor must have at least 2 dimensions") + +def sample_inputs_trilu_indices(op_info, device, dtype, requires_grad, **kwargs): + # (row, col, offset) + args_list = ((0, 0), + (20, 0), + (0, 20), + (20, 21, 0), + (20, 21, 7), + (20, 21, -7), + # Large test cases below are deliberately commented out to speed up CI + # tests and to avoid OOM error. When modifying implementations of + # tril_indices and triu_indices, please enable these tests and make sure + # they pass. + # (2, 68435455, 3), + # (5000, 5000), + # (5000, 5000, 1234), + # (5000, 5000, -1233), + ) + for args in args_list: + yield SampleInput(args[0], args=args[1:], kwargs={"dtype": dtype, "device": device}) + +def sample_inputs_clone_contiguous(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + yield SampleInput(make_arg((S, M, S))) + yield SampleInput(make_arg(())) + +def reference_inputs_clone_contiguous(op, device, dtype, requires_grad, **kwargs): + # NOTE: the default memory format for clone is torch.preserve_format, for contiguous it's torch.contiguous_format + # This exploits that default to test torch.preserve_format for clone, without causing an error when testing contiguous + yield from sample_inputs_clone_contiguous(op, device, dtype, requires_grad, **kwargs) + + shapes = ( + (3, 5, 6), + (1, 1, 3, 5, 6), + (1, 1, 3, 5, 6, 1, 1), + (1, 0, 3, 5, 0, 2), + (1, 0, 3, 5, 0, 0, 1, 1, 2), + (), + ) + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + for shape in shapes: + yield SampleInput(make_arg(shape)) + yield SampleInput(make_arg(shape).transpose(0, -1)) + yield SampleInput(make_arg(shape, noncontiguous=True)) + yield SampleInput(make_arg(shape, noncontiguous=True).transpose(0, -1)) + + yield SampleInput(make_arg(shape), kwargs={'memory_format': torch.contiguous_format}) + yield SampleInput(make_arg(shape).transpose(0, -1), kwargs={'memory_format': torch.contiguous_format}) + yield SampleInput(make_arg(shape, noncontiguous=True), kwargs={'memory_format': torch.contiguous_format}) + yield SampleInput(make_arg(shape, noncontiguous=True).transpose(0, -1), kwargs={'memory_format': torch.contiguous_format}) + + # shape, strides, offset + strided_cases = ( + ((5, 6, 2), (1, 1, 7), 2), + ((5, 5, 4), (1, 1, 7), 2), + ((5, 5, 2), (4, 5, 7), 3), + ((5, 5, 2), (5, 5, 7), 3), + ((5, 5, 2), (5, 5, 5), 3), + ((9, 5, 2), (0, 1, 7), 3), + ) + + for shape, strides, offset in strided_cases: + yield SampleInput(make_arg(500,).as_strided(shape, strides, offset)) + yield SampleInput(make_arg(500,).as_strided(shape, strides, offset), kwargs={'memory_format': torch.contiguous_format}) + + # channels last 2D + yield SampleInput(make_arg((2, 2, 2, 2)), kwargs={'memory_format': torch.channels_last}) + a = make_arg((2, 2, 2, 2)).permute(0, 3, 1, 2) + yield SampleInput(a, kwargs={'memory_format': torch.channels_last}) + + # channels last 3D + yield SampleInput(make_arg((2, 2, 2, 2, 2)), kwargs={'memory_format': torch.channels_last_3d}) + a = make_arg((2, 2, 2, 2, 2)).permute(0, 4, 1, 2, 3) + yield SampleInput(a, kwargs={'memory_format': torch.channels_last_3d}) + + +def sample_inputs_sum_to_size(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + # list of tuples (shape, shape) defining the shapes of the input and output tensors + sample_shapes = [ + ((), ()), + ((S,), (1,)), + ((S, S), (1, 1)), + ((S, S), (1, S)), + ((S, S), (S, S)), + ((S, S, S), (S, 1, S)), + ] + + for input_shape, output_shape in sample_shapes: + yield SampleInput(make_arg(input_shape), args=(output_shape,)) + if output_shape == (): + continue + yield SampleInput(make_arg(input_shape), args=(list(output_shape),)) + yield SampleInput(make_arg(input_shape), args=(*output_shape,)) + + +def error_inputs_sum_to_size(op_info, device, **kwargs): + shape = (M, S, M) + err_msg = "is not expandable to size" + si = SampleInput(make_tensor(shape, device=device, dtype=torch.float32), args=(M, M)) + yield ErrorInput(si, error_regex=err_msg) + + shape = (M + 1, S, S, M) + err_msg = "is not expandable to size" + si = SampleInput(make_tensor(shape, device=device, dtype=torch.float32), args=(M + 1, 1)) + yield ErrorInput(si, error_regex=err_msg) + + +def sample_inputs_resize_ops(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device) + cases = (((S, S, S), (S * S, S)), + ((), ()), + ((), (1, 1, 1)), + ) + + for shape, args_or_shape in cases: + # Update `args` based on operator + if op_info.name == 'resize_': + # resize_ takes shape/tuple of ints, + args = (args_or_shape, ) + elif op_info.name == 'resize_as_': + # resize_as_ takes another tensor + args = (make_arg(shape, requires_grad=False), ) # type:ignore[assignment] + else: + raise ValueError("sample_inputs_resize_ops is being used with incorrect operator") + + yield SampleInput(make_arg(shape, requires_grad=requires_grad), args=args) + +def sample_inputs_view_reshape(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + cases = ( + # a, b, is_tensor_supported + ((S, S, S), (S * S, S), True), + ((S * S, S), (S, S, S), True), + ((S * S, S), (S, -1, S), False), # neg index + ((S * S * 2, S), (S, -1), False), # neg index + ((S,), (S,), True), + ((), (), False), # empty + ((), (1,), True), + ) + + for a, b, is_tensor_supported in cases: + # skip unsupported cases + if kwargs.get("tensor_arg") and not is_tensor_supported: + continue + + # convert to tensor + if kwargs.get("tensor_arg"): + b = make_arg(b, requires_grad=False) + + yield SampleInput(make_arg(a), args=(b,)) + +def reference_inputs_view_reshape(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_view_reshape(op, device, dtype, requires_grad, **kwargs) + + cases = ( + # a, b, is_tensor_supported + ((125,), (25, 5), True), + ((25, 25), (1, 5, 5, 1, 5, 1, 5, 1), True), + ((16, 32), (2, 4, 1, 4, 4, 1, 4), True), + ((16, 12), (12, 16), True), + ((1, 16, 12), (12, 16), True), + ((1, 5, 1, 5), (25, 1), True), + ((2, 4, 2), (4, 4), True), + ((1, 4), (1, 1, 2, 1, 2), True), + ((3, 5, 7), (7, 5, 3), True), + ((1,), (), False), # empty + ((5, 0, 2, 3), (5, 0, 2, 3), True), + ((2, 1, 0, 3, 1), (5, 0), True), + ((1,), (), False), # empty + ((4, 5, 6), (4, 5, 6, 1, 1, 1), True), + ((), (1, 1, 1, 1), False), # empty + ) + + irreversible_cases = ( + ((), (-1,), False), # neg index, empty + ((4, 7, 9, 1, 1), (1, 4, 3, -1, 1), False), # neg index + ) + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + for a, b, is_tensor_supported in cases: + # skip unsupported cases + if kwargs.get("tensor_arg") and not is_tensor_supported: + continue + + if kwargs.get("tensor_arg"): + # convert to tensor + yield SampleInput(make_arg(a), args=(make_arg(b, requires_grad=False),)) + yield SampleInput(make_arg(b), args=(make_arg(a, requires_grad=False),)) + else: + yield SampleInput(make_arg(a), args=(b,)) + yield SampleInput(make_arg(b), args=(a,)) + + for a, b, is_tensor_supported in irreversible_cases: + # skip unsupported cases + if kwargs.get("tensor_arg") and not is_tensor_supported: + continue + + # convert to tensor + if kwargs.get("tensor_arg"): + b = make_arg(b, requires_grad=False) + + yield SampleInput(make_arg(a), args=(b,)) + +def error_inputs_view_reshape(op, device, **kwargs): + + cases = ( + # a, b, is_tensor_supported + # Reshape to different numel + ((2,), (), False), # empty + ((1, 3, 0), (), False), # empty + ((4, 3), (4, 2), True), + ((1, 3, 5), (5, 2, 2), True), + # No valid inference + ((1, 3, 5), (5, -1, 2), False), # neg index + # Two inferred shapes + ((1, 3, 5), (5, -1, -1), False), # neg index + ((1), (0, -1), False), # neg index + ((0, 5), (0, -1), False), # neg index + ) + + make_arg = partial(make_tensor, dtype=torch.float32, device=device, requires_grad=False) + for a, b, is_tensor_supported in cases: + # skip unsupported cases + if kwargs.get("tensor_arg") and not is_tensor_supported: + continue + + if b == (5, -1, -1): + error_regex = "only one dimension can be inferred" + elif a == (0, 5): + error_regex = (r"cannot reshape tensor of 0 elements into shape " + r"\[0, -1\] because the unspecified dimension size " + r"-1 can be any value and is ambiguous") + else: + # to avoid having issues with a regex + shape = ', '.join(map(str, b)) + size = a if type(a) is int else functools.reduce(operator.mul, a, 1) + error_regex = rf"shape '\[{shape}\]' is invalid for input of size {size}" + + # convert to tensor + if kwargs.get("tensor_arg"): + b = make_arg(b, requires_grad=False) + + yield ErrorInput(SampleInput(make_arg(a), args=(b,)), error_type=Exception, + error_regex=error_regex) + + +def sample_inputs_atleast1d2d3d(op_info, device, dtype, requires_grad, **kwargs): + shapes = ((S, S, S, S), (S, S, S), (S, S), (S, ), (),) + make_tensor_partial = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + for shape in shapes: + yield SampleInput(make_tensor_partial(shape)) + yield SampleInput([make_tensor_partial(shape) for shape in shapes]) + +def sample_inputs_column_stack(op_info, device, dtype, requires_grad, **kwargs): + cases: tuple[tuple, tuple] = ( # type: ignore[assignment] + ((S, 2, 1), (S, 3, 1)), + ((S), (S, 5)), ((), (1, S)) + ) + make_tensor_partial = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + for shape1, shape2 in cases: + yield SampleInput([make_tensor_partial(shape1), make_tensor_partial(shape2)]) + +def sample_inputs_flatten(op_info, device, dtype, requires_grad, **kwargs): + shapes = ((S, S, S), (S, S), (S, ), (),) + make_tensor_partial = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + for shape in shapes: + yield SampleInput(make_tensor_partial(shape)) + if len(shape) > 1: + yield SampleInput(make_tensor_partial(shape), start_dim=1, end_dim=-1) + +def reference_inputs_flatten(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_flatten(op, device, dtype, requires_grad, **kwargs) + + # shape x start_dim x end_dim + cases = ( + ((5, 4, 0, 1, 3, 7), 1, 3), + ((5, 4, 0, 1, 3, 7), 4, 5), + ((5, 4, 1, 1, 3, 7), 2, 3), + ((), 0, -1), + ((1,), 0, -1), + ((3, 7, 5), 1, 2), + ((4, 5), 1, 1), + ((1, 5, 5, 1, 5, 1, 5, 1), 0, 2), + ((1, 5, 5, 1, 5, 1, 5, 1), 3, -1), + ((1, 5, 5, 1, 5, 7, 5, 1), -2, -1), + ((2, 4, 2), 0, 1), + ((4, 2, 2), 1, 2), + ((0, 3, 4, 5), 1, 3), + ) + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + for shape, start, end in cases: + yield SampleInput(make_arg(shape), args=(start, end,)) + yield SampleInput(make_arg(shape, noncontiguous=True).transpose(0, -1), args=(start, end,)) + yield SampleInput(make_arg(shape).transpose(0, -1), args=(start, end,)) + +def sample_inputs_unflatten(op_info, device, dtype, requires_grad, **kwargs): + # in_shape, dim, sizes + args = (((8,), 0, (8,)), + ((8,), 0, (4, 2)), + ((8,), -1, (2, 2, 2)), + ((8,), -1, (-1, 2)), + ((3, 6, 2), 1, (2, 3)), + ((3, 6, 2), -2, (2, 3)), + ((3, 6, 2), -2, (-1, 3)), + ((3, 2, 12), 2, (3, 2, 2)), + ((4, 0), 0, (2, 2)), + ((4, 0), 1, (2, 0, 0, 0)), + ) + make_tensor_partial = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + for in_shape, dim, sizes in args: + yield SampleInput(make_tensor_partial(in_shape), args=(dim, sizes)) + + +def sample_inputs_select(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + cases = (((S, S, S), (1, 2)), + ((S, S, S), (-1, 2)), + ((S, S, S), (-1, -1)), + ((S, S, S), (1, -1)), + ((S, S), (-1, 2)), + ((S,), (0, 2)) + ) + + for shape, args in cases: + yield SampleInput(make_arg(shape), args=args) + + +def sample_inputs_select_scatter(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + cases = (((S, S, S), (S, S), (1, 2)), + ((S, S, S), (S, S), (-1, 2)), + ((S, S, S), (S, S), (-1, -1)), + ((S, S, S), (S, S), (1, -1)), + ((S,), (), (0, 2)) + ) + + for input_shape, src_shape, args in cases: + input_ = make_arg(input_shape) + src = make_arg(src_shape) + yield SampleInput(input_, args=(src, *args)) + + +def sample_inputs_slice_scatter(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + cases = (((L, L, L), (L, L, L,), (0, 0, L, 1)), + ((L, L, L), (L // 2, L, L,), (0, L // 2, L, 1)), + ((L, L, L), (L // 4, L, L,), (0, L // 2, L, 2)), + ((L, L, L), (L, L, L,), (1, 0, L, 1)), + ((L, L, L), (L, L // 2, L,), (1, L // 2, L, 1)), + ((L, L, L), (L, L // 4, L,), (1, L // 2, L, 2)), + ((L, L, L), (L, L, L,), (2, 0, L, 1)), + ((L, L, L), (L, L, L // 2,), (2, L // 2, L, 1)), + ((L, L, L), (L, L, L // 4,), (2, L // 2, L, 2)), + ) + + for input_shape, src_shape, args in cases: + input_ = make_arg(input_shape) + src = make_arg(src_shape) + yield SampleInput(input_, args=(src, *args)) + +def sample_inputs_expand(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + cases = (((S, 1, 1), (S, S, S)), + ((S, 1, S), (S, S, S)), + ((S, 1, S), (-1, S, -1)), + ((S, 1, S), (-1, S, S)), + ((S, 1), (S, S, S)), + ((1,), (S, S, S)), + ((1, S), (1, 1, S)), + ((), ()), + ((), (1, 3, 2)), + ) + + for case in cases: + shape, args = case + yield SampleInput(make_arg(shape), args=(args,)) + +def sample_inputs_conversion(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + shapes = ((), + (2, 3)) + memory_format_options = [None, torch.contiguous_format] + + for shape, memory_format in itertools.product(shapes, memory_format_options): + yield SampleInput(make_arg(shape), + kwargs={'memory_format': memory_format} if memory_format else {}) + yield SampleInput(make_arg((2, 3, 2, 3)), kwargs={'memory_format': torch.channels_last}) + +def sample_inputs_byte(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, low=0, high=255, requires_grad=requires_grad) + + shapes = ((), + (2, 3)) + memory_format_options = [None, torch.contiguous_format] + + for shape, memory_format in itertools.product(shapes, memory_format_options): + yield SampleInput(make_arg(shape), + kwargs={'memory_format': memory_format} if memory_format else {}) + yield SampleInput(make_arg((2, 3, 2, 3)), kwargs={'memory_format': torch.channels_last}) + +def sample_inputs_expand_as(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device) + + cases = (((S, 1, 1), (S, S, S)), + ((), ()), + ((), (1, 1)), + ) + + for shape, shape_other in cases: + yield SampleInput(make_arg(shape, requires_grad=requires_grad), + args=(make_arg(shape_other, requires_grad=False),)) + + +def sample_inputs_where(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + def make_bool_mask(shape): + # Make sure at least one element is nonzero, + # except for empty tensor + mask_t = make_tensor(shape, dtype=torch.bool, device=device, requires_grad=False) + + if mask_t.numel() == 0: + return mask_t + elif mask_t.numel() == 1: + mask_t.fill_(True) + return mask_t + + if mask_t.sum() == 0: + def random_index(shape): + return tuple(random.randrange(0, max_idx) for max_idx in shape) + + mask_t[random_index(mask_t.shape)] = True + return mask_t + + return mask_t + + cases = (((M, M), (M, M), (M, M), False), + ((M, 1, M), (M, M), (M, M, 1), True), + ((), (), (), False), + ((M, 1, M), (), (M, M, 1), True), + ((), (M, M), (), True), + ((), (2), (1, 1), True), + ) + + for shape, mask_shape, other_shape, broadcasts_input in cases: + yield SampleInput(make_arg(shape), + args=(make_bool_mask(mask_shape), make_arg(other_shape)), + broadcasts_input=broadcasts_input) + +# TODO: add reference inputs for where(condition) signature +def reference_inputs_where(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_where(op, device, dtype, requires_grad, **kwargs) + + make_cond = partial(make_tensor, dtype=torch.bool, device=device, requires_grad=requires_grad) + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + # noncontiguous + c = make_cond((10, 3), noncontiguous=True) + a = make_arg((10, 1), noncontiguous=True) + b = make_arg((3, 10, 3)).transpose(0, -1) + + # NOTE that the OpInfo for where takes samples of the form a, cond, b + yield SampleInput(a, args=(c, b)) + + # MPS does not support float64, which causes issues in the following tests + if torch.device(device).type == "mps": + return + + # type promoting + # FIXME(rec): shouldn't other_dtype be used two lines below? + other_dtype = torch.double if dtype is not torch.double else torch.long # noqa: F841 + c = make_cond((10, 3), noncontiguous=True) + a = make_arg((10, 1), dtype=torch.long) + b = make_arg((10, 1)) + + yield SampleInput(a, args=(c, b)) + + # two python scalars + c = make_cond((10, 3), noncontiguous=True) + a = make_arg((1,)).item() + b = make_arg((1,)).item() + + yield SampleInput(a, args=(c, b)) + + # NaN propagation + if dtype.is_floating_point or dtype.is_complex: + if dtype.is_floating_point: + nan = float('nan') + else: + # dtype.is_complex + nan = complex(float('nan'), float('nan')) + c = make_cond((1, 10, 3)) + a = make_arg((10, 3), noncontiguous=True) + a[2, 1] = nan + b = make_arg((1, 3)) + b[0, 2] = nan + + yield SampleInput(a, args=(c, b)) + + # Python scalars type promotion + for scalar in (0, 0.0, 2j, False): + yield SampleInput(scalar, args=(c, b)) + yield SampleInput(a, args=(c, scalar)) + + +def error_inputs_where(op_info, device, **kwargs): + shape = (S,) + err_msg = "Expected all tensors to be on the same device" + for devices in product(('cpu', device), repeat=3): + if len(set(devices)) == 2: + si = SampleInput(make_tensor(shape, device=devices[0], dtype=torch.float32), + args=(make_tensor(shape, dtype=torch.bool, device=devices[1]), + make_tensor(shape, device=devices[2], dtype=torch.float32))) + yield ErrorInput(si, error_regex=err_msg) + +def sample_inputs_nonzero(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + sizes = ((), (S,), (S, S), (S, S, S), (S, 1, S), (S, 0, S)) + + inputs = [] + for shape in sizes: + # construct input without any non-zero elements + zeros = torch.zeros(shape, dtype=dtype, device=device, requires_grad=requires_grad) + inputs.append(zeros) + + # construct input with mixed zero and non-zero elements + mixed = make_arg(shape).requires_grad_(False) + mask_t = make_tensor(shape, dtype=torch.bool, device=device, requires_grad=False) + mixed[mask_t] = 0 + inputs.append(mixed) + + for input_t, as_tuple in product(inputs, [False, True]): + yield SampleInput(input_t.clone().requires_grad_(requires_grad), + kwargs=dict(as_tuple=as_tuple)) + +def sample_inputs_nonzero_static(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + sizes = ((), (S,), (S, S), (S, S, S), (S, 1, S), (S, 0, S)) + + inputs = [] + for shape in sizes: + # construct input without any non-zero elements + zeros = torch.zeros(shape, dtype=dtype, device=device, requires_grad=requires_grad) + inputs.append(zeros) + + # construct input with mixed zero and non-zero elements + mixed = make_arg(shape).requires_grad_(False) + mask_t = make_tensor(shape, dtype=torch.bool, device=device, requires_grad=False) + mixed[mask_t] = 0 + inputs.append(mixed) + + nonzero_sizes = [0, 1, XS, S, M] + + for input_t, nonzero_size in product(inputs, nonzero_sizes): + yield SampleInput(input_t.clone().requires_grad_(requires_grad), + kwargs=dict(size=nonzero_size)) + +def sample_inputs_chunk(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + cases = (((S, S, S), (2,)), + ((S, S, S), (S, 1)), + ((S, S, S), (S, -1))) + + for case in cases: + shape, args = case + yield SampleInput(make_arg(shape), args=args) + +def reference_inputs_chunk(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_chunk(op, device, dtype, requires_grad, **kwargs) + + make_arg = partial(make_tensor, dtype=dtype, device=device, requires_grad=requires_grad) + + # shape x chunks x dim + cases = ( + ((13, 9, 11), 17, -1), + ((13, 9, 11), 11, -1), + ((13,), 12, -1), + ((15,), 12, -1), + ((15,), 7, 0), + ((15,), 9, 0), + ((3, 7), 9, 1), + ((3, 7), 9, 0), + ((3, 7), 2, 0), + ((3, 7), 3, 0), + ((3, 7), 1, 0), + ((3, 7), 1, 1), + ((4, 4), 2, 0), + ) + + for shape, chunks, dim in cases: + yield SampleInput(make_arg(shape), args=(chunks, dim)) + +def sample_inputs_kthvalue(op_info, device, dtype, requires_grad, **kwargs): + def _tensor(shape, dtype=dtype, low=None, high=None): + return make_tensor(shape, dtype=dtype, device=device, low=low, high=high, requires_grad=requires_grad) + + test_cases = [ + ((S, S, S), (2,)), + ((S, S, S), (2, 1,)), + ((S, S, S), (2, -1,)), + ((S, S, S), (2, 1, True,)), + ((S, S, S), (2, -1, True,)), + ((S,), (2, 0,)), + ((S,), (2, 0, True,)), + ((), (1,)), + ((), (1, 0,)), + ((), (1, 0, True)), + ] + + yield from (SampleInput(_tensor(tensor), *args) for tensor, args in test_cases) + +def error_inputs_kthvalue(op_info, device, **kwargs): + # tests overlapping output fails + t = make_tensor(10, dtype=torch.float32, device=device) + indices = torch.empty((), device=device, dtype=torch.long) + yield ErrorInput(SampleInput(t, 5, out=(t, indices)), + error_regex="unsupported operation") + + k_out_of_range_err = "selected number k out of range for dimension" + yield ErrorInput(SampleInput(torch.randn(2, 2, device=device), 3, 0), + error_regex=k_out_of_range_err) + yield ErrorInput(SampleInput(torch.randn(2, 2, device=device), 3), + error_regex=k_out_of_range_err) + yield ErrorInput(SampleInput(torch.tensor(2, device=device), 3), + error_regex=k_out_of_range_err) + +def sample_inputs_dropout(op_info, device, dtype, requires_grad, *, + train=None, valid_input_dim=None, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + if valid_input_dim: + cases = ((S,) * i for i in valid_input_dim) + else: + cases = ((S, S), (S,), ()) + p_vals = [0.0, 0.5, 1.0] + # This is to handle special case for feature_alpha_dropout which has different + # supported dtypes depending on `train` parameter + training_vals = [train] if train is not None else [True, False] + + for case, p, training in product(cases, p_vals, training_vals): + yield SampleInput(make_arg(case), p=p, training=training) + yield SampleInput(make_arg(case)) + +def sample_inputs_dropout_backward(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_mask = partial(make_tensor, device=device, dtype=torch.bool, requires_grad=False) + + cases = ((S, S, S, S), (S,), ()) + scale_vals = [0.0, 1.0, 2.0] + + for case, scale in product(cases, scale_vals): + yield SampleInput(make_arg(case), make_mask(case), scale) + +def sample_inputs_embedding_bag(op_info, device, dtype, requires_grad, **kwargs): + def make_input(shape): + return make_tensor(shape, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_long_input(shape, *, low, high, noncontiguous=False): + return make_tensor(shape, device=device, dtype=torch.long, low=low, high=high, + noncontiguous=noncontiguous) + + def make_per_sample_weight(flag, idx): + # a tensor of float / double weights, or None + # to indicate all weights should be taken to be 1 + if flag: + return make_input(idx.shape) + return None + + offsets = torch.tensor([0, 3], device=device, dtype=torch.long) + for generate_per_sample_weight in (True, False): + for mode in ('sum', 'mean', 'max'): + # per_sample_weights is only supported for mode='sum' (got mode='****') + if generate_per_sample_weight and mode in ('mean', 'max'): + continue + + # 1-D index tensor + idx = make_long_input((S,), low=0, high=M) + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(make_input((M, S)), args=(idx,), + kwargs={'offsets': offsets, 'mode': mode, + 'per_sample_weights': per_sample_weights}) + + idx = make_long_input((S,), low=0, high=M, noncontiguous=True) + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(make_input((M, S)), args=(idx,), + kwargs={'offsets': offsets, 'mode': mode, + 'per_sample_weights': per_sample_weights}) + + # bag with zero length + idx = make_long_input((S,), low=0, high=M, noncontiguous=True) + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(make_input((M, S)), args=(idx,), + kwargs={'offsets': torch.tensor([0, 0, 3], device=device, dtype=torch.long), + 'mode': mode, + 'per_sample_weights': per_sample_weights}) + + # 2-D index tensor + idx = make_long_input((S, S), low=0, high=M) + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(make_input((M, S)), args=(idx,), + kwargs={'mode': mode, 'per_sample_weights': per_sample_weights}) + + idx = make_long_input((S, S), low=0, high=M, noncontiguous=True) + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(make_input((M, S)), args=(idx,), + kwargs={'mode': mode, 'per_sample_weights': per_sample_weights}) + + # The gradient vector at `padding_idx` is not updated. + # Negative padding_idx + idx = make_long_input((6,), low=0, high=S) + idx[0] = 4 + idx[4] = 4 + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(make_input((S, S)), args=(idx,), + kwargs={'padding_idx': -1, 'offsets': offsets, + 'mode': mode, 'per_sample_weights': per_sample_weights},) + + idx = make_long_input((3, 3), low=0, high=S) + # Positive padding_idx + idx[0, 0] = 2 + idx[1, 1] = 2 + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(make_input((S, S)), args=(idx,), + kwargs={'padding_idx': 2, 'mode': mode, + 'per_sample_weights': per_sample_weights},) + + idx = make_long_input((6, ), low=0, high=S) + weights = make_input((S, S)) + offsets_ = torch.tensor([0, 3, 6], device=device, dtype=torch.long) + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(weights, args=(idx,), + kwargs={'mode': mode, 'offsets': offsets_, 'include_last_offset': True},) + + if not requires_grad: + # Following inputs return different gradient from the numerical gradient. + # This is expected and relevant tests are present in `test_nn.py`. + + # Due to inplace renorming of weight, the numerical gradient doesn't match the + # analytical gradient. + idx = make_long_input((2, 2), low=0, high=S) + weights = make_input((S, S)) * 2 + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(weights, args=(idx,), + kwargs={'max_norm': 1., 'mode': mode, + 'per_sample_weights': per_sample_weights},) + + idx = make_long_input((6, ), low=0, high=S) + weights = make_input((S, S)) * 2 + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(weights, args=(idx,), + kwargs={'max_norm': 1., 'norm_type': 1.0, + 'mode': mode, 'offsets': offsets, + 'per_sample_weights': per_sample_weights},) + + if mode != 'max': + # Scale the gradient based on the inverse frequency of a particular index. + # Note : smax mode does not support sparse weights + idx = make_long_input((2, 2), low=0, high=S) + idx[0, 0] = 1 + idx[0, 1] = 1 + weights = make_input((S, S)) + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(weights, args=(idx,), + kwargs={'scale_grad_by_freq': True, 'mode': mode, + 'per_sample_weights': per_sample_weights},) + + # gradcheck not implemented for sparse tensors. + # Note : max mode does not support sparse weights + idx = make_long_input((6, ), low=0, high=S) + weights = make_input((S, S)) + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(weights, args=(idx,), + kwargs={'sparse': True, 'offsets': offsets, + 'mode': mode, 'per_sample_weights': per_sample_weights}) + + idx = make_long_input((6, ), low=0, high=S) + idx[0] = 1 # freq more than 1 + idx[1] = 1 # freq more than 1 + idx[3] = 0 # padding_idx + weights = make_input((S, S)) * 2 + per_sample_weights = make_per_sample_weight(generate_per_sample_weight, idx) + yield SampleInput(weights, args=(idx,), + kwargs={'sparse': True, 'scale_grad_by_freq': True, 'padding_idx': 0, + 'max_norm': 1., 'offsets': offsets, + 'mode': mode, 'per_sample_weights': per_sample_weights}) + + +def sample_inputs_embedding(op_info, device, dtype, requires_grad, **kwargs): + def make_input(shape): + return make_tensor(shape, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_long_input(shape, *, low, high): + return make_tensor(shape, device=device, dtype=torch.long, low=low, high=high) + + # 0-D index tensor + idx = make_long_input((), low=0, high=M) + yield SampleInput(make_input((M, S)), args=(idx,),) + + # 1-D index tensor + idx = make_long_input((S,), low=0, high=M) + yield SampleInput(make_input((M, S)), args=(idx,),) + + # 2-D index tensor + idx = make_long_input((S, S), low=0, high=M) + yield SampleInput(make_input((M, S)), args=(idx,),) + + if not requires_grad: + # Following inputs return different gradient from the numerical gradient. + # This is expected and relevant tests are present in `test_nn.py`. + + # The gradient vector at `padding_idx` is not updated. + idx = make_long_input((2, 2), low=0, high=S) + idx[0, 0] = 2 + idx[1, 1] = 2 + yield SampleInput(make_input((S, S)), args=(idx,), kwargs={'padding_idx': 2},) + + idx = make_long_input((2, 2), low=0, high=S) + idx[0, 0] = 4 + idx[1, 1] = 4 + yield SampleInput(make_input((S, S)), args=(idx,), kwargs={'padding_idx': -1},) + + # Due to inplace renorming of weight, the numerical gradient doesn't match the + # analytical gradient. + idx = make_long_input((2, 2), low=0, high=S) + weights = make_input((S, S)) * 2 + yield SampleInput(weights, args=(idx,), kwargs={'max_norm': 1.},) + + idx = make_long_input((2, 2), low=0, high=S) + weights = make_input((S, S)) * 2 + yield SampleInput(weights, args=(idx,), kwargs={'max_norm': 1., 'norm_type': 1.0},) + + # Scale the gradient based on the inverse frequency of a particular index. + idx = make_long_input((2, 2), low=0, high=S) + idx[0, 0] = 1 + idx[0, 1] = 1 + weights = make_input((S, S)) + yield SampleInput(weights, args=(idx,), kwargs={'scale_grad_by_freq': True},) + + # gradcheck not implemented for sparse tensors. + idx = make_long_input((2, 2), low=0, high=S) + weights = make_input((S, S)) + yield SampleInput(weights, args=(idx,), kwargs={'sparse': True}) + + idx = make_long_input((3, 3), low=0, high=S) + idx[0, 0] = 1 # freq more than 1 + idx[0, 1] = 1 # freq more than 1 + idx[1, 0] = 0 # padding_idx + weights = make_input((S, S)) * 2 + yield SampleInput(weights, args=(idx,), + kwargs={'sparse': True, 'scale_grad_by_freq': True, + 'padding_idx': 0, 'max_norm': 1.}) + + +def sample_inputs_one_hot(op_info, device, dtype, requires_grad, **kwargs): + def make_input(shape, *, low, high): + return make_tensor(shape, device=device, dtype=dtype, low=low, high=high, requires_grad=requires_grad) + + shapes = ((), (S,), (L, M, S)) + num_classess = (-1, 10) + + return ( + SampleInput( + make_input( + shape, + low=0, + high=10 if num_classes == -1 else num_classes // 2, + ), + kwargs=dict(num_classes=num_classes), + ) + for shape, num_classes in itertools.product(shapes, num_classess) + ) + + +def sample_inputs_loss(op_info, device, dtype, requires_grad, **kwargs): + rhs_requires_grad = kwargs.get('rhs_requires_grad', requires_grad) + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + # Although most losses also support the reduce and size_average combination instead of reduce, the former is + # deprecated since 0.4.1 and thus is not tested + shapes_and_kwargs = ( + ((), None), + ((S,), dict(reduction="mean")), + ((S,), dict(reduction="sum")), + ((S,), dict(reduction="none")), + ((S, S), None), + ((S, S, S), None), + ) + + for shape, kwargs in shapes_and_kwargs: + yield SampleInput(_make_tensor(shape), + args=(_make_tensor(shape, requires_grad=rhs_requires_grad),), + kwargs=kwargs) + +def sample_inputs_grid_sample(op_info, device, dtype, requires_grad, **kwargs): + # We get better tests if we change the range of the values to something like [-2,2] + # because for grid (second tensor argument) the "useful" range is [-1,1] and this way + # you get a better combination of out-of-range and in-range test cases + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, + low=-2, high=2) + + batch_size = 2 + num_channels = 3 + modes = ("bilinear", "nearest") + align_cornerss = (False, True) + padding_modes = ("zeros", "border", "reflection") + + for dim in (2, 3): + + modes_ = (*modes, "bicubic") if dim == 2 else modes + + for mode, padding_mode, align_corners in itertools.product(modes_, padding_modes, align_cornerss): + yield SampleInput( + _make_tensor((batch_size, num_channels, *[S] * dim)), + _make_tensor((batch_size, *[S] * dim, dim)), + mode=mode, + padding_mode=padding_mode, + align_corners=align_corners, + ) + +def reference_inputs_grid_sample(op_info, device, dtype, requires_grad, **kwargs): + + batch_size = 2 + num_channels = 3 + height = 345 + width = 456 + modes = ("bilinear", "nearest", "bicubic") + align_cornerss = (False, True) + padding_modes = ('zeros', 'border', 'reflection') + + # Create an affine transformation matrix + a = torch.deg2rad(torch.tensor(45.0)) + ca, sa = torch.cos(a), torch.sin(a) # rotation angles + s1, s2 = 1.23, 1.34 # scales + + theta = torch.tensor([[ + [ca / s1, sa, 0.0], + [-sa, ca / s2, 0.0], + ]], dtype=dtype, device=device) + theta = theta.expand(batch_size, 2, 3).contiguous() + + x = torch.arange(batch_size * num_channels * height * width, device=device) + x = x.reshape(batch_size, num_channels, height, width).to(torch.uint8) + x = x.to(dtype=dtype) + x.requires_grad_(requires_grad) + + for mode, padding_mode, align_corners in itertools.product(modes, padding_modes, align_cornerss): + grid = torch.nn.functional.affine_grid( + theta, size=(batch_size, num_channels, height, width), align_corners=align_corners + ) + yield SampleInput( + x, + grid, + mode, + padding_mode, + align_corners, + ) + +def sample_inputs_grid_sampler_2d(op_info, device, dtype, requires_grad, **kwargs): + # We get better tests if we change the range of the values to something like [-2,2] + # because for grid (second tensor argument) the "useful" range is [-1,1] and this way + # you get a better combination of out-of-range and in-range test cases + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, + low=-2, high=2) + + batch_size = 2 + num_channels = 3 + modes = (0, 1, 2) + align_cornerss = (False, True) + padding_modes = (0, 1, 2) + + for mode, padding_mode, align_corners in itertools.product(modes, padding_modes, align_cornerss): + yield SampleInput( + _make_tensor((batch_size, num_channels, S, L)), + _make_tensor((batch_size, M + 3, M, 2)), + mode, + padding_mode, + align_corners, + ) + +def sample_inputs_grid_sampler_3d(op_info, device, dtype, requires_grad, **kwargs): + _make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, + low=-1, high=1) + # Test both out-of-range and in-range grid values + _make_grid = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad, + low=-4, high=4) + + modes = (0,) + padding_modes = (0, 1, 2) + align_cornerss = (False, True) + shape_pairs = [ + # [input_shape, grid_shape] + [(1, 1, 2, 2, 2), (1, 1, 1, 1, 3)], + [(2, 3, S, L, L), (2, M + 2, M + 1, M, 3)], + [(L, L + 1, L + 2, L + 3, L + 4), (L, M + 2, M + 1, M, 3)], + [(M, M + 1, M + 2, M + 3, M + 4), (M, L + 3, L + 2, L + 1, 3)], + [(L, M + 1, M + 2, M + 3, M + 4), (L, L + 3, L + 2, L + 1, 3)], + ] + + params_prod = itertools.product(modes, padding_modes, align_cornerss, shape_pairs) + + for mode, padding_mode, align_corners, (input_shape, grid_shape) in params_prod: + yield SampleInput( + _make_input(input_shape), + _make_grid(grid_shape), + mode, + padding_mode, + align_corners, + ) + +def sample_inputs_cosine_embedding_loss(op_info, device, dtype, requires_grad, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_target(shape): + shape = () if len(shape) == 1 else (shape[0], ) + t = torch.randint(0, 2, shape, device=device, dtype=torch.long) + # Label with -1 or 1 + t = t * 2 - 1 + target = t.to(dtype=dtype).detach_().requires_grad_(requires_grad) + return target + + shapes = ((S, S), (S,)) + reductions = ('none', 'mean', 'sum') + for s, r in product(shapes, reductions): + yield SampleInput( + make_input(s), + args=(make_input(s), make_target(s)), + kwargs=dict(reduction=r, margin=random.uniform(-1, 1)) + ) + +def sample_inputs_ctc_loss(op_info, device, dtype, requires_grad, **kwargs): + input_length = 50 + batch = 16 + num_char = 20 + target_length = 30 + + def make_log_probs(s): + t = make_tensor(s, device=device, dtype=dtype) + log_probs = t.log_softmax(2).to(device=device, dtype=dtype).detach().requires_grad_(requires_grad=requires_grad) + return log_probs + + reductions = ('none', 'mean', 'sum') + zero_inf = (True, False) + lengths_type = (list, torch.Tensor) + for r, z, lt in product(reductions, zero_inf, lengths_type): + log_probs = make_log_probs((input_length, batch, num_char)) + targets = torch.randint(1, num_char, (batch, target_length), dtype=torch.long, device=device) + input_lengths = torch.full((batch, ), input_length, dtype=torch.long, device=device) + target_lengths = torch.randint(10, target_length, (batch, ), dtype=torch.long, device=device) + + # Dont generate int[] types if reduction = "Mean" since this results in non composite compliant calls + # to ctc_loss.IntList since a tensor needs to be created from the target lengths. + # Creating such a tensor requires the use of pointers to copy data from int[] -> torch.Tensor + # e.g. via std::copy. Similarly symbolic/real tracing with fx will also not work + if lt is list and r in ["none", "sum"]: + input_lengths = input_lengths.tolist() + target_lengths = target_lengths.tolist() + + yield SampleInput(log_probs, args=(targets, input_lengths, target_lengths,), + kwargs=dict(reduction=r, zero_infinity=z)) + + +def sample_inputs_nll_loss(op_info, device, dtype, requires_grad, **kwargs): + shape = (2, 3) + num_classes = shape[1] + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + # FIXME: Derivative wrt. weight not implemented + make_weight = partial(make_tensor, num_classes, device=device, dtype=dtype, requires_grad=False) + + def make_target(shape, zeros=False): + s = (shape[0], *shape[2:]) if len(shape) > 1 else () + if zeros: + return torch.zeros(s, device=device, dtype=torch.long) + else: + return make_tensor(s, + low=0, + high=shape[1] if len(shape) > 1 else shape[0], + device=device, + dtype=torch.long) + + + def gen_shape_kwargs(): + # Batched, non-batched and 2d + shapes = (shape, (num_classes,), shape + (2, 2)) + reductions = ('none', 'mean', 'sum') + for reduction, s in product(reductions, shapes): + yield make_input(s), make_target(s), dict(reduction=reduction) + yield make_input(s), make_target(s), dict(weight=make_weight(), reduction=reduction) + yield make_input(s), make_target(s), dict(weight=make_weight(low=0), reduction=reduction) + if dtype.is_floating_point or dtype.is_complex: + yield make_input(s), make_target(s), dict(weight=make_weight(high=0), reduction=reduction) + t = make_target(s) + ignore = num_classes // 2 + # If "mean", nll returns NaN, so it's not differentiable at those points + if t.eq(ignore).all() and reduction == "mean": + t.fill_(0) + yield make_input(s), t, dict(ignore_index=num_classes // 2, reduction=reduction) + yield make_input(s), t, dict(ignore_index=num_classes // 2, reduction=reduction, weight=make_weight()) + # Test ignoring all the targets + # If "mean", nll returns NaN, so it's not differentiable at those points + if reduction != "mean": + yield make_input(s), make_target(s, zeros=True), dict(ignore_index=0, reduction=reduction) + + for input, target, kwargs in gen_shape_kwargs(): + yield SampleInput(input, args=(target,), kwargs=kwargs) + + target = torch.tensor([-1, 2], device=device, dtype=torch.long) + yield SampleInput(make_input(shape), args=(target,), kwargs={'ignore_index': -1}) + + +def sample_inputs_binary_cross_entropy_with_logits( + op_info, device, dtype, requires_grad, **kwargs +): + make = partial(make_tensor, device=device, dtype=dtype) + make_prob = partial(make, low=0, high=1) + reductions = ("mean", "sum", "none") + + def make_weight_shape_kwargs(): + kwargs = [] + for shape in ((1,), (1, S), (S), (S, S)): + kwargs.extend([((S, S), dict(reduction=reduction, weight=make(shape))) for reduction in reductions]) + return kwargs + + shapes_and_kwargs = [ + *[(shape, None) for shape in ((), (1,), (S,), (S, S), (S, S, S))], + *[((S, S), dict(reduction=reduction)) for reduction in reductions], + *make_weight_shape_kwargs(), + *[((S, S), dict(reduction=reduction, pos_weight=make((S,), low=0))) for reduction in reductions], + *[((S, S), dict(reduction=reduction, weight=make((S, S)), pos_weight=make((S,), low=0))) for reduction in reductions], + ] + + for shape, kwargs in shapes_and_kwargs: + yield SampleInput( + make(shape, requires_grad=requires_grad), + args=(make_prob(shape, requires_grad=requires_grad),), + kwargs=kwargs, + ) + +def sample_inputs_argwhere(op_info, device, dtype, requires_grad, **kwargs): + yield SampleInput(torch.tensor([1, 0, 2, 0], dtype=dtype, device=device, requires_grad=requires_grad)) + mask = torch.tensor([[0, 1, 0, 1, 0], + [1, 1, 1, 1, 0], + [0, 0, 0, 1, 0], + [1, 0, 1, 1, 0], + [1, 0, 0, 1, 0]], dtype=torch.bool, device=device) + t = make_tensor((S, S), dtype=dtype, device=device, requires_grad=requires_grad) + t[mask] = 0 + yield SampleInput(t) + + t = make_tensor((S, S), dtype=dtype, device=device, requires_grad=requires_grad, noncontiguous=True) + t[mask] = 0 + yield SampleInput(t) + + t = make_tensor((S, 0), dtype=dtype, device=device, requires_grad=requires_grad) + yield SampleInput(t) + + yield SampleInput(torch.zeros((S,), dtype=dtype, device=device, requires_grad=requires_grad)) + yield SampleInput(make_tensor((), dtype=dtype, device=device, requires_grad=requires_grad)) + +def _generate_sample_shape_reduction(): + shapes = ((S,), (S, S), (S, S, S)) + reductions = ('none', 'mean', 'sum') + yield from product(shapes, reductions) + +def sample_inputs_gaussian_nll_loss(op_info, device, dtype, requires_grad, **kwargs): + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + # Set low slightly above 0 so gradcheck doesn't accidentally dip below 0 + make_var = partial(make_tensor, low=0.1, device=device, dtype=dtype, requires_grad=requires_grad) + + def gen_shape(shape): + yield shape + # Broadcast + yield (*shape[:-1], 1) + yield shape[:-1] + + def gen_shape_kwargs(): + for s, r in _generate_sample_shape_reduction(): + for t_s, v_s in product(gen_shape(s), gen_shape(s)): + yield _make_tensor(s), _make_tensor(t_s), make_var(v_s), dict(reduction=r) + yield ( + _make_tensor(s), _make_tensor(t_s), make_var(v_s), + dict(full=True, reduction=r) + ) + yield ( + _make_tensor(s), _make_tensor(t_s), make_var(v_s), + dict(eps=random.uniform(1e-6, 1e-3), reduction=r) + ) + yield ( + _make_tensor(s), _make_tensor(t_s), make_var(v_s), + dict(full=True, eps=random.uniform(1e-6, 1e-3), reduction=r) + ) + + for input, target, var, kwargs in gen_shape_kwargs(): + yield SampleInput(input, args=(target, var, ), kwargs=kwargs) + +def error_inputs_gaussian_nll_loss(op_info, device, **kwargs): + _make = partial(make_tensor, device=device, dtype=torch.float32) + + # invalid reduction value + yield ErrorInput(SampleInput(_make(10, 2, 3), _make(10, 2, 3), _make((10, 2, 3), low=0), reduction="abc"), + error_type=ValueError, error_regex="abc is not valid") + + # var is of incorrect shape + yield ErrorInput(SampleInput(_make(10, 2, 3), _make(10, 2, 3), _make((10, 2, 2), low=0)), + error_type=ValueError, error_regex="var is of incorrect size") + + # target is of incorrect shape + yield ErrorInput(SampleInput(_make(10, 2, 3), _make(10, 2, 2), _make((10, 2, 3), low=0)), + error_type=RuntimeError, + error_regex=(r"The size of tensor a \(3\) must match the size of tensor b \(2\) " + r"at non-singleton dimension 2")) + +def _generate_sample_inputs_nn_loss(op_info, device, dtype, requires_grad, **kwargs): + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + for s, r in _generate_sample_shape_reduction(): + yield _make_tensor(s), _make_tensor(s), dict(reduction=r) + +def sample_inputs_hinge_embedding_loss(op_info, device, dtype, requires_grad, **kwargs): + for input, target, d in _generate_sample_inputs_nn_loss(op_info, device, dtype, requires_grad, **kwargs): + # target should contain either 1 or -1 as per docs + mask = torch.rand_like(target) > 0.5 + target[mask] = 1 + target[~mask] = -1 + d['margin'] = random.uniform(-9, 9) + yield SampleInput(input, args=(target, ), kwargs=d) + + # scalar input and target. + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(_make_tensor(()), args=(_make_tensor(()), )) + +def error_inputs_hinge_embedding_loss(op, device, **kwargs): + make_input = partial(make_tensor, device=device, dtype=torch.float32) + # invalid reduction value + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5, 4),), kwargs={'reduction': 'abc'}), + error_type=ValueError, error_regex='is not a valid value') + +def reference_inputs_hinge_embedding_loss(op, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_hinge_embedding_loss(op, device, dtype, requires_grad, **kwargs) + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + for reduction in ('sum', 'mean', 'none'): + if dtype.is_floating_point: # only supports ints and floats + # NaN propagation + inp = make_input((10, )) + inp[2] = float('nan') + target = make_input((10, )) + # target should contain either 1 or -1 as per docs + mask = torch.rand_like(target) > 0.5 + target[mask] = -1 + target[~mask] = 1 + yield SampleInput(inp, args=(target,), kwargs={'reduction': reduction}) + + # Inf Handling + inp = make_input((10, )) + inp[4] = float('inf') + target = make_input((10, )) + mask = torch.rand_like(target) > 0.5 + target[mask] = -1 + target[~mask] = 1 + yield SampleInput(inp, args=(target,), kwargs={'reduction': reduction}) + + # Broadcasting + inp = make_input((5, 5)) + target = make_input((1, 5)) + mask = torch.rand_like(target) > 0.5 + target[mask] = -1 + target[~mask] = 1 + yield SampleInput(inp, args=(target,), kwargs={'reduction': reduction}) + +def sample_inputs_huber_loss(op_info, device, dtype, requires_grad, **kwargs): + for input, target, d in _generate_sample_inputs_nn_loss(op_info, device, dtype, requires_grad, **kwargs): + d['delta'] = random.uniform(1e-3, 9) + yield SampleInput(input, args=(target, ), kwargs=d) + +def error_inputs_huber_loss(op, device, **kwargs): + make_input = partial(make_tensor, device=device, dtype=torch.float32) + # invalid reduction value + err = 'is not a valid value for reduction' + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5, 4),), kwargs={'reduction': 'abc'}), + error_type=ValueError, error_regex=err) + # delta <= 0 + for delta in (0, -1): + err = 'huber_loss does not support non-positive values for delta.' + yield ErrorInput(SampleInput(make_input(5, 4), args=(make_input(5, 4),), kwargs={'delta': delta}), + error_type=RuntimeError, error_regex=err) + +def sample_inputs_poisson_nll_loss(op_info, device, dtype, requires_grad, **kwargs): + _make_tensor = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def gen_shape_kwargs(): + for s, r in _generate_sample_shape_reduction(): + for li in (True, False): + for f in (True, False): + i1 = _make_tensor(s) + i2 = _make_tensor(s) + # For Poisson NLL Loss, + # target is assumed to be from + # Poisson Distribution which + # always has positive samples + t1 = _make_tensor(s, low=0) + t2 = _make_tensor(s, low=0) + + if not li: + i1.abs_() + i2.abs_() + t1.abs_() + t2.abs_() + + yield ( + i1, t1, + dict(log_input=li, full=f, reduction=r) + ) + yield ( + i2, t2, + dict(log_input=li, full=f, + eps=random.uniform(1e-8, 1e-3), + reduction=r) + ) + + for input, target, kwargs in gen_shape_kwargs(): + yield SampleInput(input, args=(target, ), kwargs=kwargs) + + # test INT_TO_FLOAT promotion + if dtype.is_complex: + for d in (torch.bool, torch.int64): + yield SampleInput(_make_tensor(dtype=dtype), args=(_make_tensor(dtype=d),)) + yield SampleInput(_make_tensor(dtype=d), args=(_make_tensor(dtype=dtype),)) + +def error_inputs_poisson_nll_loss(op_info, device, **kwargs): + make = partial(make_tensor, device=device, dtype=torch.float32) + + # invalid reduction value + yield ErrorInput(SampleInput(make(5, 4), args=(make(5, 4),), + kwargs={'reduction': 'abc'}), + error_type=ValueError, + error_regex='abc is not a valid value for reduction') + # invalid input shapes + yield ErrorInput(SampleInput(make(5, 4), args=(make(5,),)), + error_regex=(r'(Attempting to broadcast a dimension of length|' + r'The size of tensor a \(5\) must match the ' + r'size of tensor b \(4\) at non-singleton ' + r'dimension 1)')) + +def error_inputs_soft_margin_loss(op_info, device, **kwargs): + make = partial(make_tensor, device=device, dtype=torch.float32) + + # invalid reduction value + yield ErrorInput(SampleInput(make(5, 4), args=(make(5, 4),), + kwargs={'reduction': 'abc'}), + error_type=ValueError, + error_regex='abc is not a valid value for reduction') + # invalid input shapes + yield ErrorInput(SampleInput(make(5, 4), args=(make(5,),)), + error_regex=(r'(Attempting to broadcast a dimension of length|' + r'The size of tensor a \(4\) must match the ' + r'size of tensor b \(5\) at non-singleton ' + r'dimension 1)')) + +def sample_inputs_triplet_margin_loss(op_info, device, dtype, requires_grad, with_distance=False, **kwargs): + make = partial(make_tensor, (S, M), device=device, dtype=dtype, requires_grad=requires_grad) + + kwargss = ( + *[dict(margin=margin) for margin in (1e-6, 1.0, 10.0)], + dict(swap=True), + *[dict(reduction=reduction) for reduction in ("mean", "sum", "none")], + ) + + for kwargs in kwargss: + input = make() + args = (make(), make()) + if with_distance: + kwargs["distance_function"] = torch.nn.PairwiseDistance() + yield SampleInput(input, args=args, kwargs=kwargs) + +def error_inputs_triplet_margin_loss(op_info, device, **kwargs): + make_input = partial(make_tensor, device=device, dtype=torch.float32) + + samples = ( + # input, args, kwargs, error_type, error_regex + # invalid reduction + (make_input(3, 4), (make_input(3, 4), make_input(3, 4)), + dict(reduction="abc"), + ValueError, "abc is not a valid value for reduction"), + + # invalid margin + (make_input(3, 4), (make_input(3, 4), make_input(3, 4)), + dict(margin=-1.0), + ValueError, "margin must be greater than 0, got -1.0"), + + # shape mismatch + (make_input(3, 5), (make_input(3, 4), make_input(3, 4)), + {}, + RuntimeError, + (r'(Attempting to broadcast a dimension of length|' + r"The size of tensor a \(5\) must match the size of tensor b \(4\) " + r"at non-singleton dimension 1)")), + (make_input(3, 4), (make_input(3, 5), make_input(3, 4)), + {}, + RuntimeError, + (r'(Attempting to broadcast a dimension of length|' + r"The size of tensor a \(4\) must match the size of tensor b \(5\) " + r"at non-singleton dimension 1)")), + (make_input(3, 4), (make_input(3, 4), make_input(3, 5)), + {}, + RuntimeError, + (r'(Attempting to broadcast a dimension of length|' + r"The size of tensor a \(4\) must match the size of tensor b \(5\) " + r"at non-singleton dimension 1)")), + + # different dimensions + (make_input(3,), (make_input(3, 4), make_input(3, 4)), + {}, + RuntimeError, + (r"The anchor, positive, and negative tensors are expected to have " + r"the same number of dimensions, but got: anchor 1D, positive 2D, " + r"and negative 2D inputs")), + (make_input(3, 4), (make_input(3,), make_input(3, 4)), + {}, + RuntimeError, + (r"The anchor, positive, and negative tensors are expected to have " + r"the same number of dimensions, but got: anchor 2D, positive 1D, " + r"and negative 2D inputs")), + (make_input(3, 4), (make_input(3, 4), make_input(3,)), + {}, + RuntimeError, + (r"The anchor, positive, and negative tensors are expected to have " + r"the same number of dimensions, but got: anchor 2D, positive 2D, " + r"and negative 1D inputs")), + ) + + for input, args, kwargs, error_type, error_regex in samples: + yield ErrorInput(SampleInput(input, args=args, kwargs=kwargs), + error_type=error_type, error_regex=error_regex) + + +def sample_inputs_scaled_mm(op_info, device, dtype, requires_grad, **kwargs): + def to_fp8_saturated(x: torch.Tensor, fp8_dtype: torch.dtype) -> torch.Tensor: + max_val = E4M3_MAX_POS if fp8_dtype == e4m3_type else E5M2_MAX_POS + x = x.clamp(min=-1 * max_val, max=max_val) + return x.to(fp8_dtype) + + def amax_to_scale(amax: torch.Tensor, float8_dtype: torch.dtype) -> torch.Tensor: + EPS = 1e-12 + max_pos = E4M3_MAX_POS if float8_dtype == e4m3_type else E5M2_MAX_POS + scale_val = max_pos / torch.clamp(amax, min=EPS) + return scale_val.to(dtype=torch.float32, device=device) + + def make_scale(x: float, float8_dtype: torch.dtype, dim=None) -> torch.Tensor: + if dim is None: + amax = torch.tensor(abs(x), dtype=torch.float32, device=device) + else: + amax = torch.max( + torch.abs(torch.tensor(x, device=device)), dim=dim, keepdim=True + ).values + return amax_to_scale(amax, float8_dtype) + + def make_mat(size: tuple[int], scale: float, fp8_dtype: torch.dtype) -> torch.Tensor: + mat = torch.randn(size, device=device, dtype=torch.float32) + return to_fp8_saturated(mat * scale, fp8_dtype) + + M, N, K = 15, 32, 16 + samples = [] + + # Case 1: Both matrices e4m3 + scale1 = random.random() + scale2 = random.random() + mat1 = make_mat((M, K), scale1, torch.float8_e4m3fn) + mat2 = make_mat((K, N), scale2, torch.float8_e4m3fn).t().contiguous().t() + scale_tensor1 = make_scale(scale1, torch.float8_e4m3fn) + scale_tensor2 = make_scale(scale2, torch.float8_e4m3fn) + samples.append(SampleInput(mat1, mat2, scale_tensor1, scale_tensor2)) + + # Case 2: mat1 e4m3, mat2 e5m2 + scale1 = random.random() + scale2 = random.random() + mat1 = make_mat((M, K), scale1, torch.float8_e4m3fn) + mat2 = make_mat((K, N), scale2, torch.float8_e5m2).t().contiguous().t() + scale_tensor1 = make_scale(scale1, torch.float8_e4m3fn) + scale_tensor2 = make_scale(scale2, torch.float8_e5m2) + samples.append(SampleInput(mat1, mat2, scale_tensor1, scale_tensor2)) + + # Case 3: mat1 e5m2, mat2 e4m3 + scale1 = random.random() + scale2 = random.random() + mat1 = make_mat((M, K), scale1, torch.float8_e5m2) + mat2 = make_mat((K, N), scale2, torch.float8_e4m3fn).t().contiguous().t() + scale_tensor1 = make_scale(scale1, torch.float8_e5m2) + scale_tensor2 = make_scale(scale2, torch.float8_e4m3fn) + samples.append(SampleInput(mat1, mat2, scale_tensor1, scale_tensor2)) + + # Case 4: MXFP4 (float4_e2m1fn_x2) with E8M0 blockwise scaling + # Regression test: E8M0 blockwise scale size validation must account for + # packed FP4 format where self.size(1) = K/2. + # Only supported on MI350 (gfx950). + if device == 'cuda' and torch.version.hip: + if 'gfx950' in torch.cuda.get_device_properties(0).gcnArchName: + mxfp4_M, mxfp4_K, mxfp4_N = 256, 256, 256 + block_size_k = 32 + block_size_mn = 128 + num_k_blocks = math.ceil(mxfp4_K / block_size_k) + padded_num_k_blocks = math.ceil(num_k_blocks / 4) * 4 + scale_a_size = block_size_mn * math.ceil(mxfp4_M / block_size_mn) * padded_num_k_blocks + scale_b_size = block_size_mn * math.ceil(mxfp4_N / block_size_mn) * padded_num_k_blocks + mat1 = _bfloat16_to_float4_e2m1fn_x2( + torch.randn(mxfp4_M, mxfp4_K, device=device, dtype=torch.bfloat16) + ) + mat2 = _bfloat16_to_float4_e2m1fn_x2( + torch.randn(mxfp4_N, mxfp4_K, device=device, dtype=torch.bfloat16) + ).t() + scale_tensor1 = torch.ones(scale_a_size, device=device, dtype=torch.float8_e8m0fnu) + scale_tensor2 = torch.ones(scale_b_size, device=device, dtype=torch.float8_e8m0fnu) + samples.append(SampleInput( + mat1, mat2, scale_tensor1, scale_tensor2, + out_dtype=torch.bfloat16, + )) + + yield from samples + +def sample_inputs_scaled_mm_v2(op_info, device, dtype, requires_grad, **kwargs): + from torch.nn.functional import ScalingType, SwizzleType + make_mat_e4m3 = partial(make_tensor, device=device, dtype=torch.float8_e4m3fn, requires_grad=requires_grad) + + make_scale = partial(make_tensor, device=device, dtype=torch.float, requires_grad=False) + + M, N, K = 15, 32, 16 + samples = [] + # two e4m3 tensorwise + mat1 = make_mat_e4m3((M, K)) + mat2 = make_mat_e4m3((K, N)).t().contiguous().t() + scale1 = make_scale((1,)) + scale2 = make_scale((1,)) + samples.append( + SampleInput( + mat1, + mat2, + [scale1, ], + [ScalingType.TensorWise, ], + [SwizzleType.NO_SWIZZLE, ], + [scale2, ], + [ScalingType.TensorWise, ], + [SwizzleType.NO_SWIZZLE, ], + None, # bias + torch.bfloat16, # out_dtype + ) + ) + # two e4m3 rowwise + mat1 = make_mat_e4m3((M, K)) + mat2 = make_mat_e4m3((K, N)).t().contiguous().t() + scale1 = make_scale((M, 1)) + scale2 = make_scale((1, N)) + samples.append( + SampleInput( + mat1, + mat2, + [scale1, ], + [ScalingType.RowWise, ], + [SwizzleType.NO_SWIZZLE, ], + [scale2, ], + [ScalingType.RowWise, ], + [SwizzleType.NO_SWIZZLE, ], + None, # bias + torch.bfloat16, # out_dtype + ) + ) + M, K, N = 256, 512, 768 + mat1 = make_mat_e4m3((M, K)) + mat2 = make_mat_e4m3((K, N)).t().contiguous().t() + + if torch.device(device).type == "cuda": + dmajor, dminor = torch.cuda.get_device_capability() + + # Blockwise scaling requires cublasLt >= 12.9 (CUDA 12.9+) + # See: https://github.com/pytorch/pytorch/issues/172227 + cuda_version = _get_torch_cuda_version() + if dmajor == 9 and not torch.version.hip and cuda_version >= (12, 9): + # 1x128 x 1x128 + scale1 = make_scale((K // 128, M)).t() + scale2 = make_scale((K // 128, N)).t() + samples.append( + SampleInput( + mat1, + mat2, + [scale1, ], + [ScalingType.BlockWise1x128, ], + [SwizzleType.NO_SWIZZLE, ], + [scale2, ], + [ScalingType.BlockWise1x128, ], + [SwizzleType.NO_SWIZZLE, ], + None, # bias + torch.bfloat16, # out_dtype + ) + ) + # 128x128 x 1x128 + L4 = round_up(K // 128, 4) + scale1 = make_scale((M // 128, L4)).t() + scale2 = make_scale((K // 128, N)).t() + samples.append( + SampleInput( + mat1, + mat2, + [scale1, ], + [ScalingType.BlockWise128x128, ], + [SwizzleType.NO_SWIZZLE, ], + [scale2, ], + [ScalingType.BlockWise1x128, ], + [SwizzleType.NO_SWIZZLE, ], + None, # bias + torch.bfloat16, # out_dtype + ) + ) + # 1x128 x 128x128 + L4 = round_up(K // 128, 4) + scale1 = make_scale((K // 128, M)).t() + scale2 = make_scale((N // 128, L4)).t() + samples.append( + SampleInput( + mat1, + mat2, + [scale1, ], + [ScalingType.BlockWise1x128, ], + [SwizzleType.NO_SWIZZLE, ], + [scale2, ], + [ScalingType.BlockWise128x128, ], + [SwizzleType.NO_SWIZZLE, ], + None, # bias + torch.bfloat16, # out_dtype + ) + ) + + if dmajor >= 10: + # MXFP8 + scale1 = make_scale((M, K // 32)).to(torch.float8_e8m0fnu) + scale2 = make_scale((K // 32, N)).to(torch.float8_e8m0fnu) + samples.append( + SampleInput( + mat1, + mat2, + [scale1, ], + [ScalingType.BlockWise1x32, ], + [SwizzleType.SWIZZLE_32_4_4, ], + [scale2, ], + [ScalingType.BlockWise1x32, ], + [SwizzleType.SWIZZLE_32_4_4, ], + None, # bias + torch.bfloat16, # out_dtype + ) + ) + # Single-level NVFP4 + # [M, K] -> [M, K // 2] + # [K, N] -> [K // 2, N] + mat1_fp4 = _bfloat16_to_float4_e2m1fn_x2(mat1.to(torch.bfloat16)) + mat2_fp4 = _bfloat16_to_float4_e2m1fn_x2(mat2.to(torch.bfloat16).t()).t() + scale1 = make_scale((M, K // 16)).to(torch.float8_e4m3fn) + scale2 = make_scale((K // 16, N)).to(torch.float8_e4m3fn) + samples.append( + SampleInput( + mat1_fp4, + mat2_fp4, + [scale1, ], + [ScalingType.BlockWise1x16, ], + [SwizzleType.SWIZZLE_32_4_4, ], + [scale2, ], + [ScalingType.BlockWise1x16, ], + [SwizzleType.SWIZZLE_32_4_4, ], + None, # bias + torch.bfloat16, # out_dtype + ) + ) + + # NVFP4 + # [M, K] -> [M, K // 2] + # [K, N] -> [K // 2, N] + mat1_fp4 = _bfloat16_to_float4_e2m1fn_x2(mat1.to(torch.bfloat16)) + mat2_fp4 = _bfloat16_to_float4_e2m1fn_x2(mat2.to(torch.bfloat16).t()).t() + scale1 = make_scale((M, K // 16)).to(torch.float8_e4m3fn) + global_scale1 = make_scale((1, )) + scale2 = make_scale((K // 16, N)).to(torch.float8_e4m3fn) + global_scale2 = make_scale((1, )) + samples.append( + SampleInput( + mat1_fp4, + mat2_fp4, + [scale1, global_scale1], + [ScalingType.BlockWise1x16, ScalingType.TensorWise], + [SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE], + [scale2, global_scale2], + [ScalingType.BlockWise1x16, ScalingType.TensorWise], + [SwizzleType.SWIZZLE_32_4_4, SwizzleType.NO_SWIZZLE], + None, # bias + torch.bfloat16, # out_dtype + ) + ) + + + yield from samples + +def sample_inputs_scaled_dot_product_attention(op_info, device, dtype, requires_grad, **kwargs): + make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + batch, seq_q, seq_kv, num_heads, head_dim = 4, 3, 6, 4, 8 + + dim_3_q_shape = (batch, seq_q, head_dim) + dim_3_kv_shape = (batch, seq_kv, head_dim) + dim_4_q_shape = (batch, num_heads, seq_q, head_dim) + dim_4_kv_shape = (batch, num_heads, seq_kv, head_dim) + + broadcast_tuple = ((num_heads, seq_q, head_dim), (batch, num_heads, seq_kv, head_dim)) + + qkv_shapes = [(dim_3_q_shape, dim_3_kv_shape), (dim_4_q_shape, dim_4_kv_shape), broadcast_tuple] + samples = [] + gqa_options = [True, False] + causal_options = [True, False] + for qkv_shape, is_causal, dropout_p, _enable_gqa in product( + qkv_shapes, causal_options, [0.0, 0.5], gqa_options): + shape_q, shape_kv = qkv_shape + samples.append(SampleInput( + make(shape_q), + make(shape_kv), + make(shape_kv), + is_causal=is_causal, + dropout_p=dropout_p + )) + + # Add non standard shapes + # FIXME(rec): should diff_v_head_dim be appended to samples? + diff_v_head_dim = SampleInput( # noqa: F841 + make((batch, num_heads, seq_q, head_dim)), + make((batch, num_heads, seq_kv, head_dim)), + make((batch, num_heads, seq_kv, head_dim + 8)), + is_causal=is_causal, + dropout_p=dropout_p + ) + + # Add an attn_mask + samples.append( + SampleInput( + make((batch, num_heads, seq_q, head_dim)), + make((batch, num_heads, seq_kv, head_dim)), + make((batch, num_heads, seq_kv, head_dim)), + attn_mask=make((seq_q, seq_kv)), + is_causal=False, + dropout_p=0.0) + ) + + yield from samples + + +def sample_inputs_efficient_attention_forward(op_info, device, dtype, requires_grad, **kwargs): + make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + batch, num_heads, head_dim = 4, 4, 8 + seq_q = 11 + seq_kv = 32 + + dim_4_q_shape = (batch, num_heads, seq_q, head_dim) + dim_4_kv_shape = (batch, num_heads, seq_kv, head_dim) + + qkv_shapes = [(dim_4_q_shape, dim_4_kv_shape)] + samples = [] + mask_types = [1, 2] # UpperLeft, LowerRight + scales = [None, 1.0] + + for qkv_shape, _is_causal, dropout_p, mask_type, scale in product( + qkv_shapes, [True, False], [0.0, 0.5], mask_types, scales): + shape_q, shape_kv = qkv_shape + samples.append(SampleInput( + make(shape_q).transpose(1, 2), + make(shape_kv).transpose(1, 2), + make(shape_kv).transpose(1, 2), + bias=None, + cu_seqlens_q=None, + cu_seqlens_k=None, + max_seqlen_q=None, + max_seqlen_k=None, + dropout_p=dropout_p, + custom_mask_type=mask_type, + compute_log_sumexp=requires_grad, + scale=scale, + seqlen_k=None + )) + + # Add non standard shapes + # FIXME(rec): should diff_v_head_dim be appended to samples? + diff_v_head_dim = SampleInput( # noqa: F841 + make((batch, seq_q, num_heads, head_dim)), + make((batch, seq_kv, num_heads, head_dim)), + make((batch, seq_kv, num_heads, head_dim + 8)), + bias=None, + cu_seqlens_q=None, + cu_seqlens_k=None, + max_seqlen_q=None, + max_seqlen_k=None, + dropout_p=dropout_p, + custom_mask_type=0, # No Mask + compute_log_sumexp=requires_grad, + scale=None, + seqlen_k=None + ) + + # Add an attn_mask + samples.append( + SampleInput( + make((batch, seq_q, num_heads, head_dim)), + make((batch, seq_kv, num_heads, head_dim)), + make((batch, seq_kv, num_heads, head_dim)), + bias=make(batch, num_heads, seq_q, seq_kv), + cu_seqlens_q=None, + cu_seqlens_k=None, + max_seqlen_q=None, + max_seqlen_k=None, + dropout_p=dropout_p, + custom_mask_type=0, # No Mask + compute_log_sumexp=requires_grad, + scale=None, + seqlen_k=None + ) + ) + + # jagged (with query/keys offsets) + cu_seqlens_k = torch.arange(-1, 32 * 2 + 1, 2, dtype=torch.int32, device=device) + cu_seqlens_k[-1] = 62 + cu_seqlens_k[0] = 0 + samples.append( + SampleInput( + make((32, 2, 64)).view(-1, 8, 8).unsqueeze(0), + make((64, 64)).view(-1, 8, 8).unsqueeze(0), + make((64, 64)).view(-1, 8, 8).unsqueeze(0), + bias=None, + cu_seqlens_q=torch.arange(0, 32 * 2 + 2, 2, dtype=torch.int32, device=device), + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=2, + max_seqlen_k=2, + dropout_p=0.0, + custom_mask_type=0, # No Mask + compute_log_sumexp=requires_grad, + scale=None, + seqlen_k=None, + ) + ) + + yield from samples + +def sample_inputs_flash_attention_forward(op_info, device, dtype, requires_grad, **kwargs): + make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + batch, num_heads, head_dim = 4, 4, 8 + seq_q = 11 + seq_kv = 32 + + dim_4_q_shape = (batch, num_heads, seq_q, head_dim) + dim_4_kv_shape = (batch, num_heads, seq_kv, head_dim) + + qkv_shapes = [(dim_4_q_shape, dim_4_kv_shape)] + samples = [] + scales = [None, 1.0] + + for qkv_shape, is_causal, dropout_p, scale in product( + qkv_shapes, [True, False], [0.0, 0.5], scales): + shape_q, shape_kv = qkv_shape + samples.append(SampleInput( + make(shape_q).transpose(1, 2), + make(shape_kv).transpose(1, 2), + make(shape_kv).transpose(1, 2), + cum_seq_q=None, + cum_seq_k=None, + max_q=seq_q, + max_k=seq_kv, + dropout_p=dropout_p, + is_causal=is_causal, + return_debug_mask=False, + scale=scale, + )) + + yield from samples + +def sample_inputs_pairwise_distance(op_info, device, dtype, requires_grad, **kwargs): + make = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + shape = (3,) + batched_shape = (2, *shape) + shapes_and_kwargs = [ + (shape, None), + (batched_shape, None), + (shape, dict(keepdim=True)), + (batched_shape, dict(keepdim=True)), + (shape, dict(p=5.0)), + (shape, dict(p=-1.0)), + (shape, dict(eps=1.0)), + ] + + return ( + SampleInput(make(shape), args=(make(shape),), kwargs=kwargs) for shape, kwargs in shapes_and_kwargs + ) + +def sample_inputs_pixel_shuffle(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield from ( + SampleInput(make_arg((1, 9, 2, 2)), upscale_factor=upscale_factor) + for upscale_factor in (1, 3) + ) + yield from ( + SampleInput(make_arg(shape), upscale_factor=1) + for shape in [ + (1, 0, 1, 1), + (1, 1, 0, 1), + (1, 1, 1, 0), + ] + ) + +def sample_inputs_pixel_unshuffle(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + yield from ( + SampleInput(make_arg((1, 1, 6, 6)), downscale_factor=downscale_factor) + for downscale_factor in (1, 3) + ) + yield from ( + SampleInput(make_arg(shape), downscale_factor=1) + for shape in [ + (1, 0, 1, 1), + (1, 1, 0, 1), + (1, 1, 1, 0), + ] + ) + +def sample_inputs_channel_shuffle(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + shapes_groups = [ + ((1, 4, 10, 10), 2), + ((2, 6, 8, 8), 3), + ((2, 8, 5, 5), 4), + ] + + yield from ( + SampleInput(make_arg(shape), args=(groups,)) + for shape, groups in shapes_groups + ) + +def sample_inputs_binary_cross_entropy(op_info, device, dtype, requires_grad, logits=False, **kwargs): + make = partial(make_tensor, device=device, dtype=dtype) + # Lower bounds must be greater than 'eps' defined in gradcheck.py::gradgradcheck() -> eps + # otherwise perturbation calculation causes Tensor value to become negative triggering + # a device-side hardware assertion + make_prob = partial(make, low=1e-6, high=1) + + reductions = ("mean", "sum", "none") + + shapes_and_kwargs = [ + *[(shape, None) for shape in ((), (1,), (S,), (S, S), (S, S, S))], + *[((S, S), dict(reduction=reduction)) for reduction in reductions], + *[((S, S), dict(reduction=reduction, weight=make((S, S)))) for reduction in reductions], + ] + + if logits: + shapes_and_kwargs.extend( + [((S, S), dict(reduction=reduction, pos_weight=make((S,), low=0))) for reduction in reductions] + ) + + for shape, kwargs in shapes_and_kwargs: + yield SampleInput( + (make if logits else make_prob)(shape, requires_grad=requires_grad), + args=(make_prob(shape, requires_grad=requires_grad),), + kwargs=kwargs, + ) + +def sample_inputs_allclose(op_info, device, dtype, requires_grad, **kwargs): + sample_shapes = [(), (S), (S, S, S)] + atols = [1e-2, 1e-16] + rtols = [1e-1, 0.5] + for s, rtol, atol in product(sample_shapes, rtols, atols): + # close sample + t = make_tensor(s, device=device, dtype=dtype, requires_grad=requires_grad) + close = (t + atol).detach().requires_grad_(requires_grad) + yield SampleInput(t, close, rtol=rtol, atol=atol) + + # random sample + a = make_tensor(s, device=device, dtype=dtype, requires_grad=requires_grad) + b = make_tensor(s, device=device, dtype=dtype, requires_grad=requires_grad) + yield SampleInput(a, b, rtol=rtol, atol=atol) + + +def sample_inputs_l1_loss(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_loss(op_info, device, dtype, requires_grad, **kwargs) + + # test COMPLEX_TO_FLOAT promotion + if dtype.is_complex: + make = partial(make_tensor, (), device=device, requires_grad=requires_grad) + other_dtype = highest_precision_float(device) + yield SampleInput(make(dtype=dtype), args=(make(dtype=other_dtype),)) + yield SampleInput(make(dtype=other_dtype), args=(make(dtype=dtype),)) + +def error_inputs_l1_loss(op_info, device, **kwargs): + make = partial(make_tensor, device=device, dtype=torch.float32) + + # invalid reduction value + yield ErrorInput(SampleInput(make(5, 4), args=(make(5, 4),), + kwargs={'reduction': 'abc'}), + error_type=ValueError, + error_regex='abc is not a valid value for reduction') + # invalid input shapes + yield ErrorInput(SampleInput(make(5, 4), args=(make(5,),)), + error_regex=(r'(Attempting to broadcast a dimension of length|' + r'The size of tensor a \(4\) must match the ' + r'size of tensor b \(5\) at non-singleton ' + r'dimension 1)') + ) + +def sample_inputs_smooth_l1_loss(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_loss(op_info, device, dtype, requires_grad, **kwargs) + + make = partial(make_tensor, (S, S), device=device, dtype=dtype, requires_grad=requires_grad) + + # This test case always triggers the smooth condition, since absolute difference of input and target + # is smaller than beta + yield SampleInput(make(low=0, high=2), args=(make(low=-2, high=0),), kwargs=dict(beta=5)) + yield SampleInput(make(), args=(make(),), kwargs=dict(beta=0)) + +def sample_inputs_kl_div(op_info, device, dtype, requires_grad, **kwargs): + # kl_div works with inputs in [0, 1] (aka the pdf of a probability measure) + # Then log [0, 1] = (-inf, 0], so this is the log space + make_arg = partial(make_tensor, low=0., device=device, dtype=dtype, requires_grad=requires_grad) + + def make_log(shape): + out = torch.nn.functional.log_softmax(make_arg(shape), -1) + out.requires_grad_(requires_grad) + return out + + def make_prob(shape): + out = torch.nn.functional.softmax(make_arg(shape), -1) + out.requires_grad_(requires_grad) + return out + + shapes = ((2,), (2, 3)) + reductions = ("none", "mean", "batchmean", "sum") + for shape, reduction, log_target in product(shapes, reductions, (True, False)): + input = make_log(shape) + target = make_log(shape) if log_target else make_prob(shape) + yield SampleInput(input, args=(target,), kwargs=dict(reduction=reduction, log_target=log_target)) + +def sample_inputs_pdist(op_info, device, dtype, requires_grad, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + yield from (SampleInput(make_input((n, m))) for n, m in itertools.product((1, S), repeat=2)) + yield from (SampleInput(make_input((S, S)), kwargs=dict(p=p)) for p in (0.0, 1.0, 2.0, 10.0, float("inf"))) + +def reference_pdist(input, p=2): + pdist = scipy.spatial.distance.pdist + if p == 0: + output = pdist(input, "hamming") * input.shape[1] + elif p == float("inf"): + output = pdist(input, lambda x, y: np.abs(x - y).max()) + else: + output = pdist(input, "minkowski", p=p) + return output.astype(input.dtype) + +def sample_inputs_diagflat(op_info, device, dtype, requires_grad, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(make_input(())) + yield SampleInput(make_input((2,))) + yield SampleInput(make_input((2, 2))) + yield SampleInput(make_input((2,)), offset=1) + yield SampleInput(make_input((2,)), offset=-1) + + +_UNPOOL_NAME_TO_DIM = { + 'nn.functional.max_unpool1d': 1, + 'nn.functional.max_unpool2d': 2, + 'nn.functional.max_unpool3d': 3 +} + + +def error_inputs_max_unpool(op_info, device, **kwargs): + """Error inputs for max_unpool: shape mismatch between input and indices.""" + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + pool_dim = _UNPOOL_NAME_TO_DIM[op_info.name] + + # Create mismatched shapes for input and indices + kwargs_dict = {'kernel_size': 3, 'stride': 2, 'padding': 0} + if pool_dim == 1: + input_shape = (8, 8) + indices_shape = (8, 7) + elif pool_dim == 2: + input_shape = (1, 1, 4, 4) + indices_shape = (1, 1, 4, 1) + else: # pool_dim == 3 + input_shape = (1, 1, 4, 4, 4) + indices_shape = (1, 1, 4, 4, 1) + + yield ErrorInput( + SampleInput( + make_arg(input_shape), + args=(torch.zeros(indices_shape, device=device, dtype=torch.long),), + kwargs=kwargs_dict + ), + error_type=RuntimeError, + error_regex='Expected shape of indices to be' + ) + + +def sample_inputs_max_unpool(op_info, device, dtype, requires_grad, **kwargs): + unpool_name_to_pool_method_dict = { + 'nn.functional.max_unpool1d': torch.nn.functional.max_pool1d, + 'nn.functional.max_unpool2d': torch.nn.functional.max_pool2d, + 'nn.functional.max_unpool3d': torch.nn.functional.max_pool3d + } + + unpool_to_pool_name_dict = {k: f'nn.functional.{v.__name__}' for k, v in unpool_name_to_pool_method_dict.items()} + + pool_dim = _UNPOOL_NAME_TO_DIM[op_info.name] + pool_method = unpool_name_to_pool_method_dict[op_info.name] + + pool_op_info = copy.copy(op_info) + pool_op_info.name = unpool_to_pool_name_dict[op_info.name] + + for sample in sample_inputs_max_pool(pool_op_info, device, dtype, requires_grad, **kwargs): + # shapes (C, ...) do not work as of now, + # see https://github.com/pytorch/pytorch/issues/68337 + # TODO: remove once the issue is resolved + if sample.input.dim() != pool_dim + 2: + continue + + # No dilation > 1 for max_unpool, + # see https://github.com/pytorch/pytorch/issues/68420 + if sample.kwargs['dilation'] != 1: + continue + + # Can't unpool without indices + if sample.kwargs['return_indices']: + pool, indices = pool_method(sample.input, **sample.kwargs) + # arg has to be a leaf + arg = pool.detach().requires_grad_(requires_grad) + sample_kwargs = { + 'kernel_size': sample.kwargs['kernel_size'], + 'stride': sample.kwargs['stride'], + 'padding': sample.kwargs['padding'], + # output_size could be None but we specify it explicitly + # to compensate for the information lose in pool due + # to the floor/ceil operation used to compute the shapes + 'output_size': sample.input.size() + } + + yield SampleInput(arg, args=(indices,), kwargs=sample_kwargs) + +def sample_inputs_max_unpool_grad(op_info, device, dtype, requires_grad, **kwargs): + for sample in sample_inputs_max_unpool(op_info, device, dtype, requires_grad, **kwargs): + indices = sample.args[0] + # The samples for max_unpool are generated with max_pool. + # It could be that a single element from the max_pool's + # input is mapped to several locations in its output. + # This situation leads to failed gradchecks because + # the finite difference algorithm perturbs the elements + # of the output one by one, and not in classes of + # equivalences determined by whether two elements + # in the output are coming from the same location in the + # input (simply put, they have the same corresponding index). + # So, there are two ways to resolve this issue: + # 1. Extract a perturbation for one element and apply it all + # the elements from the same equivalence class, or + # 2. Make sure that the equivalence classes are all singletons, + # i.e. the index tensor has to be comprised of only unique + # indices. + # Here we go with the solution 2, the easiest of all. + if indices.unique().numel() == indices.numel(): + yield sample + +def sample_inputs_multi_head_attention_forward(opinfo, device, dtype, requires_grad, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + if requires_grad: + # backward tests would take too long to complete, causing the job timeout. + bsz = 2 + is_batcheds = (True,) + use_separate_proj_weights = (False,) + emb_sizes = (2,) + src_lens = (XS,) + tgt_lens = (XS,) + heads = (2,) + dropouts = (0.5,) + mask_types = ("2d",) + else: + bsz = 2 + is_batcheds = (False, True) + use_separate_proj_weights = (False, True) + emb_sizes = (2, 4) + src_lens = (XS,) + tgt_lens = (XS, S) + heads = (1, 2) + dropouts = (0.0, 0.5) + mask_types = (None, "2d", "3d") + + for is_batched, use_separate_proj_weight, mask_type, emb_size, src_len, tgt_len, num_heads, dropout_p in itertools.product( + is_batcheds, use_separate_proj_weights, mask_types, emb_sizes, src_lens, tgt_lens, heads, dropouts + ): + attn_mask = None + if mask_type == "2d": + attn_mask = make_input(src_len, tgt_len) + elif mask_type == "3d": + attn_mask = make_input((bsz if is_batched else 1) * num_heads, src_len, tgt_len) + + if is_batched: + q = make_input(src_len, bsz, emb_size) + k = make_input(tgt_len, bsz, emb_size) + v = make_input(tgt_len, bsz, emb_size) + else: + q = make_input(src_len, emb_size) + k = make_input(tgt_len, emb_size) + v = make_input(tgt_len, emb_size) + if use_separate_proj_weight: + in_proj_weight = None + q_proj_weight = make_input(emb_size, emb_size) + k_proj_weight = make_input(emb_size, emb_size) + v_proj_weight = make_input(emb_size, emb_size) + else: + in_proj_weight = make_input(emb_size * 3, emb_size) + q_proj_weight = None + k_proj_weight = None + v_proj_weight = None + + bias_k = make_input(emb_size) + bias_v = make_input(emb_size) + in_proj_bias = make_input(emb_size * 3) + out_proj_weight = make_input(emb_size, emb_size) + out_proj_bias = make_input(emb_size) + sample_args = ( + k, v, emb_size, num_heads, in_proj_weight, + in_proj_bias, bias_k, bias_v, False, + dropout_p, out_proj_weight, out_proj_bias + ) + sample_kwargs = { + "q_proj_weight" : q_proj_weight, + "k_proj_weight" : k_proj_weight, + "v_proj_weight" : v_proj_weight, + "attn_mask" : attn_mask, + "training" : dropout_p > 0.0, + "use_separate_proj_weight" : use_separate_proj_weight + } + + yield SampleInput(q, args=sample_args, kwargs=sample_kwargs) + + +# Includes some values such that N * N won't be a multiple of 4, +# which should ensure we test the vectorized and non-vectorized +# kernel code paths. +NUM_SIZE0_TENSORS = 10000 +foreach_num_tensors = [20, 23] if not TEST_WITH_SLOW else [23, 30, 300] +_foreach_inputs_default_kwargs = {"noncontiguous": False, "same_size": False, "low": None, "high": None} + + +class ForeachRightmostArgType(enum.Enum): + TensorList = enum.auto() + ScalarList = enum.auto() + Scalar = enum.auto() + Tensor = enum.auto() + + +class ForeachSampleInput(SampleInput): + # For TensorList Scalar/Tensor, we compute the reference + # by converting it into TensorList ScalarList/TensorList and + # then converting into multiple Tensor Scalar/Tensor. + # ref_args contains the args converted to TensorList ScalarList/TensorList + ref_args: Any + disable_fastpath: bool + + def __init__(self, *args, disable_fastpath=False, ref_args=None, **kwargs): + super().__init__(*args, **kwargs) + self.ref_args = ref_args or self.args + self.disable_fastpath = disable_fastpath + + +class foreach_inputs_sample_func: + def __init__( + self, + arity: int, + rightmost_supports_scalar: bool, + rightmost_supports_scalarlist: bool, + rightmost_supports_tensor: bool = False, + ) -> None: + self.arity = arity + self._set_rightmost_arg_types( + rightmost_supports_scalar, rightmost_supports_scalarlist, rightmost_supports_tensor, + ) + self._intersperse_empty = (True, False) + + def _set_rightmost_arg_types( + self, + rightmost_supports_scalar: bool, + rightmost_supports_scalarlist: bool, + rightmost_supports_tensor: bool, + ) -> None: + self._rightmost_arg_types = [ForeachRightmostArgType.TensorList] + if self.arity > 1: + if rightmost_supports_scalar: + self._rightmost_arg_types.append(ForeachRightmostArgType.Scalar) + if rightmost_supports_scalarlist: + self._rightmost_arg_types.append(ForeachRightmostArgType.ScalarList) + if rightmost_supports_tensor: + self._rightmost_arg_types.append(ForeachRightmostArgType.Tensor) + + def _sample_rightmost_arg( + self, + opinfo, + rightmost_arg_type, + device, + dtype, + num_tensors, + allow_higher_dtype_scalars, + **_foreach_inputs_kwargs, + ): + if rightmost_arg_type == ForeachRightmostArgType.TensorList: + return [sample_inputs_foreach(None, device, dtype, num_tensors, **_foreach_inputs_kwargs)] + if rightmost_arg_type == ForeachRightmostArgType.Tensor: + return [make_tensor( + (), device=device, dtype=dtype, + noncontiguous=_foreach_inputs_kwargs["noncontiguous"], + requires_grad=_foreach_inputs_kwargs.get("requires_grad", False), + )] + should_use_simpler_scalars = opinfo.name == "_foreach_pow" and dtype in (torch.float16, torch.bfloat16) + + def sample_float(): + s = random.random() + if should_use_simpler_scalars: + return 1.0 if s > 0.5 else 2.0 + else: + return 1.0 - s + + high = 2 if should_use_simpler_scalars else 9 + if rightmost_arg_type == ForeachRightmostArgType.ScalarList: + scalarlist_list = [] + scalarlist_list.append([random.randint(0, high) + 1 for _ in range(num_tensors)]) + + if allow_higher_dtype_scalars or dtype.is_floating_point: + scalarlist_list.append([sample_float() for _ in range(num_tensors)]) + if allow_higher_dtype_scalars or dtype.is_complex: + scalarlist_list.append([complex(sample_float(), sample_float()) for _ in range(num_tensors)]) + scalarlist_list.append([1, 2.0, 3.0 + 4.5j] + [3.0 for _ in range(num_tensors - 3)]) + scalarlist_list.append([True, 1, 2.0, 3.0 + 4.5j] + [3.0 for _ in range(num_tensors - 4)]) + return scalarlist_list + if rightmost_arg_type == ForeachRightmostArgType.Scalar: + scalars = [] + scalars.append(random.randint(1, high + 1)) + if allow_higher_dtype_scalars or dtype.is_floating_point: + scalars.append(sample_float()) + if allow_higher_dtype_scalars or dtype.is_complex: + scalars.append(complex(sample_float(), sample_float())) + scalars.append(True) + return scalars + raise AssertionError(f"Invalid rightmost_arg_type of {rightmost_arg_type}") + + def _should_disable_fastpath(self, opinfo, rightmost_arg, rightmost_arg_type, dtype): + if self.arity == 1: + if "foreach_abs" in opinfo.name and dtype in complex_types(): + return True + if "foreach_clone" in opinfo.name: + return False + # unary + if opinfo.ref in (torch.abs, torch.neg): + return False + if opinfo.ref_inplace == torch.Tensor.zero_: + return False + return dtype in integral_types_and(torch.bool) + if self.arity < 2 or rightmost_arg_type == ForeachRightmostArgType.Tensor: + return None + if "foreach_pow" in opinfo.name and dtype in integral_types_and(torch.bool): + return True + if any( + foreach_name in opinfo.name + for foreach_name in ("foreach_clamp_max", "foreach_clamp_min", "foreach_maximum", "foreach_minimum") + ) and dtype in integral_types_and(torch.bool): + return True + if rightmost_arg_type == ForeachRightmostArgType.TensorList: + disable_fastpath = "foreach_div" in opinfo.name and dtype in integral_types_and(torch.bool) + if "foreach_add" in opinfo.name and dtype == torch.bool: + disable_fastpath = True + return disable_fastpath + elif rightmost_arg_type == ForeachRightmostArgType.Scalar: + disable_fastpath = "foreach_div" in opinfo.name and dtype in integral_types_and(torch.bool) + if isinstance(rightmost_arg, bool): + disable_fastpath |= dtype == torch.bool + if opinfo.ref in (torch.add, torch.mul): + disable_fastpath = False + elif isinstance(rightmost_arg, int): + disable_fastpath |= dtype == torch.bool + elif isinstance(rightmost_arg, float): + disable_fastpath |= dtype in integral_types_and(torch.bool) + elif isinstance(rightmost_arg, complex): + disable_fastpath |= dtype not in complex_types() + else: + raise AssertionError(f"Invalid scalar of type {rightmost_arg_type} - {rightmost_arg}") + return disable_fastpath + elif rightmost_arg_type == ForeachRightmostArgType.ScalarList: + disable_fastpath = opinfo.ref == torch.div and dtype in integral_types_and(torch.bool) + elmt_t = type(rightmost_arg[0]) + has_same_type = all(isinstance(v, elmt_t) for v in rightmost_arg) + if not has_same_type: + return dtype not in complex_types() + if isinstance(rightmost_arg[0], bool): + if ("foreach_add" in opinfo.name or "foreach_mul" in opinfo.name) and dtype == torch.bool: + disable_fastpath = False + elif isinstance(rightmost_arg[0], int): + disable_fastpath |= dtype == torch.bool + elif isinstance(rightmost_arg[0], float): + disable_fastpath |= dtype in integral_types_and(torch.bool) + elif isinstance(rightmost_arg[0], complex): + disable_fastpath |= dtype not in complex_types() + else: + raise AssertionError(f"Invalid scalarlist of {rightmost_arg}") + return disable_fastpath + else: + raise AssertionError(f"Invalid rightmost_arg_type of {rightmost_arg_type}") + + def _sample_kwargs(self, opinfo, rightmost_arg, rightmost_arg_type, dtype): + kwargs = {} + if rightmost_arg_type == ForeachRightmostArgType.TensorList and opinfo.supports_alpha_param: + if dtype in integral_types_and(torch.bool): + kwargs["alpha"] = 3 + elif dtype.is_complex: + kwargs["alpha"] = complex(3, 3) + else: + kwargs["alpha"] = 3.14 + if self.arity > 1: + kwargs["disable_fastpath"] = self._should_disable_fastpath(opinfo, rightmost_arg, rightmost_arg_type, dtype) + return kwargs + + def sample_zero_size_tensor_inputs(self, opinfo, device, dtype, requires_grad, **kwargs): + if "num_input_tensors" in kwargs: + raise AssertionError("num_input_tensors should not be in kwargs") + _foreach_inputs_kwargs = {k: kwargs.pop(k, v) for k, v in _foreach_inputs_default_kwargs.items()} + _foreach_inputs_kwargs["requires_grad"] = requires_grad + allow_higher_dtype_scalars = kwargs.pop("allow_higher_dtype_scalars", False) + for _rightmost_arg_type in self._rightmost_arg_types: + zero_size_foreach_inputs_kwargs = copy.deepcopy(_foreach_inputs_kwargs) + zero_size_foreach_inputs_kwargs["zero_size"] = True + input = sample_inputs_foreach(None, device, dtype, NUM_SIZE0_TENSORS, **zero_size_foreach_inputs_kwargs) + if self.arity > 1: + args = [ + sample_inputs_foreach(None, device, dtype, NUM_SIZE0_TENSORS, **zero_size_foreach_inputs_kwargs) + for _ in range(self.arity - 2) + ] + args.append( + self._sample_rightmost_arg( + opinfo, + ForeachRightmostArgType.TensorList, + device, + dtype, + NUM_SIZE0_TENSORS, + allow_higher_dtype_scalars=allow_higher_dtype_scalars, + **zero_size_foreach_inputs_kwargs, + )[0]) + kwargs = self._sample_kwargs( + opinfo, args[-1], ForeachRightmostArgType.TensorList, dtype) + else: + args = [] + kwargs = {} + if opinfo.ref in (torch.abs, torch.neg): + kwargs["disable_fastpath"] = False + else: + kwargs["disable_fastpath"] = dtype in integral_types_and(torch.bool) + yield ForeachSampleInput(input, *args, **kwargs) + + def __call__(self, opinfo, device, dtype, requires_grad, **kwargs): + num_input_tensors_specified = "num_input_tensors" in kwargs + num_input_tensors = kwargs.pop("num_input_tensors") if num_input_tensors_specified else foreach_num_tensors + if not isinstance(num_input_tensors, list): + raise AssertionError(f"Expected num_input_tensors to be a list, got {type(num_input_tensors)}") + _foreach_inputs_kwargs = {k: kwargs.pop(k, v) for k, v in _foreach_inputs_default_kwargs.items()} + _foreach_inputs_kwargs["requires_grad"] = requires_grad + _foreach_inputs_kwargs["zero_size"] = False + allow_higher_dtype_scalars = kwargs.pop("allow_higher_dtype_scalars", False) + + # add empty tensor interspersion to test fully fixing #100701 + for num_tensors, rightmost_arg_type, intersperse_empty_tensors in itertools.product( + num_input_tensors, self._rightmost_arg_types, self._intersperse_empty): + if intersperse_empty_tensors and (num_tensors != max(num_input_tensors) or str(device) == 'cpu'): + # generate interspersed empty tensors for only 1 N on non-cpu device to lessen redundancy + continue + _foreach_inputs_kwargs["intersperse_empty_tensors"] = intersperse_empty_tensors + input = sample_inputs_foreach( + None, device, dtype, num_tensors, **_foreach_inputs_kwargs) + args = [] + if self.arity > 1: + args = [ + sample_inputs_foreach( + None, device, dtype, num_tensors, **_foreach_inputs_kwargs) + for _ in range(self.arity - 2) + ] + rightmost_arg_list = self._sample_rightmost_arg( + opinfo, rightmost_arg_type, device, dtype, num_tensors, allow_higher_dtype_scalars, + **_foreach_inputs_kwargs) + for rightmost_arg in rightmost_arg_list: + args.append(rightmost_arg) + kwargs = self._sample_kwargs(opinfo, rightmost_arg, rightmost_arg_type, dtype) + ref_args = args + if rightmost_arg_type in (ForeachRightmostArgType.Scalar, ForeachRightmostArgType.Tensor): + ref_args = args[:-1] + [[args[-1] for _ in range(num_tensors)]] + sample = ForeachSampleInput(input, *args, ref_args=ref_args, **kwargs) + yield sample + args.pop() + else: + yield ForeachSampleInput( + input, + *args, + disable_fastpath=self._should_disable_fastpath(opinfo, None, None, dtype), + ) + + +class foreach_max_sample_func(foreach_inputs_sample_func): + def __init__( + self, + arity: int, + rightmost_supports_scalar: bool, + rightmost_supports_scalarlist: bool, + rightmost_supports_tensor: bool = False, + ) -> None: + super().__init__(arity, rightmost_supports_scalar, rightmost_supports_scalarlist, rightmost_supports_tensor) + self._intersperse_empty = (False,) + + def sample_zero_size_tensor_inputs(self, opinfo, device, dtype, requires_grad, **kwargs): + return [] + + def _should_disable_fastpath(self, opinfo, rightmost_arg, rightmost_arg_type, dtype): + return False + + +class foreach_norm_sample_func(foreach_inputs_sample_func): + def sample_zero_size_tensor_inputs(self, opinfo, device, dtype, requires_grad, **kwargs): + if "num_input_tensors" in kwargs: + raise AssertionError("num_input_tensors should not be in kwargs") + _foreach_inputs_kwargs = {k: kwargs.pop(k, v) for k, v in _foreach_inputs_default_kwargs.items()} + _foreach_inputs_kwargs["requires_grad"] = requires_grad + for ord in (0, 1, 2, -1, -2, float('inf'), float('-inf')): + input = sample_inputs_foreach(None, device, dtype, NUM_SIZE0_TENSORS, zero_size=True, **_foreach_inputs_kwargs) + disable_fastpath = True + if ord in (0, 1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16): + disable_fastpath = False + yield ForeachSampleInput(input, ord=ord, disable_fastpath=disable_fastpath) + + def __call__(self, opinfo, device, dtype, requires_grad, **kwargs): + num_input_tensors = kwargs.pop("num_input_tensors", foreach_num_tensors) + if not isinstance(num_input_tensors, list): + raise AssertionError(f"Expected num_input_tensors to be a list, got {type(num_input_tensors)}") + _foreach_inputs_kwargs = {k: kwargs.pop(k, v) for k, v in _foreach_inputs_default_kwargs.items()} + _foreach_inputs_kwargs["requires_grad"] = requires_grad + _allow_higher_dtype_scalars = kwargs.pop("allow_higher_dtype_scalars", False) + + for num_tensors, ord, out_dtype, intersperse_empty_tensors in product( + num_input_tensors, + (0, 1, 2, -1, -2, float('inf'), float('-inf')), + (None,) + (highest_precision_complex(device),) if dtype in complex_types() else (highest_precision_float(device),), + (True, False), + ): + # inf norm and negative norms on empty tensors is not supported by our reference func vector norm: + # linalg.vector_norm cannot compute the inf norm on an empty tensor because the operation does not have an identity + if (ord in [float('inf'), float('-inf')] or ord < 0) and intersperse_empty_tensors: + continue + + _foreach_inputs_kwargs["intersperse_empty_tensors"] = intersperse_empty_tensors + input = sample_inputs_foreach(None, device, dtype, num_tensors, zero_size=False, **_foreach_inputs_kwargs) + disable_fastpath = True + if ord in (0, 1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16): + disable_fastpath = False + yield ForeachSampleInput(input, ord=ord, disable_fastpath=disable_fastpath, dtype=out_dtype) + + # Also test nan propagation with a single tensor, but skip autograd testing + if not requires_grad: + nan_inputs = [ + [float('nan')], + [float('nan'), 1.0], + [1.0, float('nan')], + [1.0, 2.0, 3.0, float('nan'), float('nan'), 7.0, float('nan'), float('nan'), -1.5, 6.0], + [7.0, 3.0, float('nan'), float('nan'), -1.5, 6.0], + [3.0, float('nan'), float('nan'), -1.5, 6.0], + ] + for input in nan_inputs: + x = torch.tensor(input, device=device) + disable_fastpath = True + if ord in (0, 1, 2, float('inf')) and dtype in floating_types_and(torch.half, torch.bfloat16): + disable_fastpath = False + yield ForeachSampleInput([x], ord=ord, disable_fastpath=disable_fastpath) + + +class foreach_pointwise_sample_func(foreach_inputs_sample_func): + + def __init__( + self, + arity: int = 3, + rightmost_supports_scalar: bool = False, + rightmost_supports_scalarlist: bool = False, + ): + super().__init__(arity, rightmost_supports_scalar, rightmost_supports_scalarlist) + + def _should_disable_fastpath(self, opinfo, rightmost_arg, rightmost_arg_type, dtype): + return dtype in integral_types_and(torch.bool) and opinfo.ref == torch.addcmul + + def sample_zero_size_tensor_inputs(self, opinfo, device, dtype, requires_grad, **kwargs): + if "num_input_tensors" in kwargs: + raise AssertionError("num_input_tensors should not be in kwargs") + _foreach_inputs_kwargs = {k: kwargs.pop(k, v) for k, v in _foreach_inputs_default_kwargs.items()} + _foreach_inputs_kwargs["requires_grad"] = requires_grad + # zero_size tensor + input = sample_inputs_foreach(None, device, dtype, NUM_SIZE0_TENSORS, zero_size=True, **_foreach_inputs_kwargs) + args = [ + sample_inputs_foreach(None, device, dtype, NUM_SIZE0_TENSORS, zero_size=True, **_foreach_inputs_kwargs) + for _ in range(2) + ] + kwargs.pop("scalars", None) + kwargs.update(self._sample_kwargs(opinfo, args[-1], ForeachRightmostArgType.TensorList, dtype)) + yield ForeachSampleInput(input, *args, **kwargs) + + def __call__(self, opinfo, device, dtype, requires_grad, **kwargs): + num_input_tensors_specified = "num_input_tensors" in kwargs + num_input_tensors = kwargs.pop("num_input_tensors") if num_input_tensors_specified else foreach_num_tensors + if not isinstance(num_input_tensors, list): + raise AssertionError(f"Expected num_input_tensors to be a list, got {type(num_input_tensors)}") + _foreach_inputs_kwargs = {k: kwargs.pop(k, v) for k, v in _foreach_inputs_default_kwargs.items()} + _foreach_inputs_kwargs["requires_grad"] = requires_grad + allow_higher_dtype_scalars = kwargs.pop("allow_higher_dtype_scalars", False) + + for num_tensors, rightmost_arg_type, intersperse_empty_tensors in itertools.product( + num_input_tensors, self._rightmost_arg_types, (True, False)): + _foreach_inputs_kwargs["intersperse_empty_tensors"] = intersperse_empty_tensors + input = sample_inputs_foreach(None, device, dtype, num_tensors, zero_size=False, **_foreach_inputs_kwargs) + args = [ + sample_inputs_foreach(None, device, dtype, num_tensors, zero_size=False, **_foreach_inputs_kwargs) + for _ in range(2 - int(rightmost_arg_type == ForeachRightmostArgType.TensorList)) + ] + rightmost_arg_list = self._sample_rightmost_arg( + opinfo, + rightmost_arg_type, + device, + dtype, + num_tensors, + zero_size=False, + allow_higher_dtype_scalars=False if intersperse_empty_tensors else allow_higher_dtype_scalars, + **_foreach_inputs_kwargs, + ) + for rightmost_arg in rightmost_arg_list: + kwargs = {} + if rightmost_arg_type == ForeachRightmostArgType.TensorList: + args.append(rightmost_arg) + elif rightmost_arg_type in [ForeachRightmostArgType.Tensor, ForeachRightmostArgType.ScalarList]: + kwargs["scalars"] = rightmost_arg + else: + kwargs["value"] = rightmost_arg + kwargs.update(self._sample_kwargs(opinfo, rightmost_arg, rightmost_arg_type, dtype)) + if len(args) != 2: + raise AssertionError(f"Expected len(args) == 2, got {len(args)}") + sample = ForeachSampleInput(input, *args, **kwargs) + yield sample + if rightmost_arg_type == ForeachRightmostArgType.TensorList: + args.pop() + + +foreach_unary_op_db: list[OpInfo] = [ + ForeachFuncInfo( + 'exp', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32), + backward_requires_result=True, + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'acos', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'asin', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'atan', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'cos', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'cosh', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'log', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'log10', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'log2', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'tan', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + backward_requires_result=True, + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + # due to https://github.com/pytorch/pytorch/pull/102427 enabling jiterator for complex + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + toleranceOverride( + { + torch.complex64: tol(atol=3e-04, rtol=2e-05) + } + ), + 'TestForeach', + 'test_parity', + device_type='cuda' + ), + ), + ), + ForeachFuncInfo( + 'tanh', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + backward_requires_result=True, + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + toleranceOverride( + {torch.complex64: tol(atol=5e-03, rtol=1e-04)} + ), + 'TestForeach', + 'test_parity', + device_type='cuda' + ), + ), + ), + ForeachFuncInfo( + 'sin', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'sinh', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + ), + ), + ForeachFuncInfo( + 'neg', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_unary_op_tensors_on_different_devices", + device_type="cuda", + dtypes=(torch.bool,), + ), + ), + ), + ForeachFuncInfo( + 'sqrt', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + backward_requires_result=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + ), + ), + ForeachFuncInfo( + 'rsqrt', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + backward_requires_result=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + ), + ), + ForeachFuncInfo( + 'ceil', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + ), + ), + ForeachFuncInfo( + 'erf', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + ), + ), + ForeachFuncInfo( + 'erfc', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + ), + ), + ForeachFuncInfo( + 'expm1', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + backward_requires_result=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + ), + ), + ForeachFuncInfo( + 'floor', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + ), + ), + ForeachFuncInfo( + 'log1p', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + ), + ), + ForeachFuncInfo( + 'round', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + ), + ), + ForeachFuncInfo( + 'frac', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=integral_types_and(torch.bool) + complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + ), + ), + ForeachFuncInfo( + 'reciprocal', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + backward_requires_result=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + ), + ), + ForeachFuncInfo( + 'sigmoid', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + backward_requires_result=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + ), + ), + ForeachFuncInfo( + 'trunc', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=complex_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + ), + ), + ForeachFuncInfo( + 'abs', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + device_type="cpu", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + device_type="cpu", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + device_type="cpu", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + device_type="cpu", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + device_type="cpu", + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + device_type="cpu", + dtypes=(torch.bool,), + ), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", dtypes=complex_types()), + ), + ), + ForeachFuncInfo( + 'zero', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + supports_out=False, + ), + ForeachFuncInfo( + 'sign', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + ), + ), + ForeachFuncInfo( + 'lgamma', + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo(unittest.skip("In-place lgamma not supported for integral tensors"), "TestMeta", + "test_dispatch_symbolic_meta_inplace", dtypes=integral_types_and(torch.bool)), + # DecorateInfo(unittest.skip("In-place lgamma not supported for integral tensors"), "TestMeta", + # "test_dispatch_meta_inplace", dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.skip("In-place lgamma not supported for integral tensors"), "TestMeta", + "test_meta_inplace", dtypes=integral_types_and(torch.bool)), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=complex_types() + integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types() + integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=complex_types() + integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=complex_types(), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + ), + ), + ForeachFuncInfo( + "clone", + sample_inputs_func=foreach_inputs_sample_func(1, False, False), + supports_forward_ad=True, + supports_autograd=True, + supports_inplace_autograd=True, + ), +] + +foreach_binary_op_db: list[OpInfo] = [ + ForeachFuncInfo( + "add", + sample_inputs_func=foreach_inputs_sample_func(2, True, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16, torch.int32), + supports_alpha_param=True, + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + # These tests fail with aten._local_scalar_dense not being implemented. + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_outplace"), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16)), + # Samples have complex types and inplace only works if the dtype is complex. + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=integral_types() + complex_types_and(torch.bool, torch.bfloat16, torch.float16, torch.float64)), + ), + ), + ForeachFuncInfo( + "sub", + sample_inputs_func=foreach_inputs_sample_func(2, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_alpha_param=True, + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace"), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace"), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace"), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_outplace"), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace"), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_outplace"), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides"), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides"), + DecorateInfo(unittest.skip("consistently fails internally and causes other tests to appear flaky"), + "TestForeach", "test_parity", dtypes=(torch.complex128,), + active_if=lambda kwargs: IS_FBCODE and not kwargs["noncontiguous"]), + ), + ), + ForeachFuncInfo( + "mul", + sample_inputs_func=foreach_inputs_sample_func(2, True, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + # Samples have complex types and inplace only works if the dtype is complex. + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=(torch.bool,)), + DecorateInfo(unittest.skip("consistently fails internally and causes other tests to appear flaky"), + "TestForeach", "test_parity", dtypes=(torch.complex128,), + active_if=lambda kwargs: IS_FBCODE and not kwargs["noncontiguous"]), + ), + ), + ForeachFuncInfo( + "div", + sample_inputs_func=foreach_inputs_sample_func(2, True, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16, torch.int32, torch.int8), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + # Samples have complex types and inplace only works if the dtype is complex. + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", + dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=integral_types_and(torch.bool)), + ), + ), + ForeachFuncInfo( + "clamp_min", + sample_inputs_func=foreach_inputs_sample_func(2, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16, torch.int64, torch.int32, torch.int8, torch.bool), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides", + dtypes=complex_types_and(torch.bool)), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_binary_op_scalar_with_overlapping_tensors", + dtypes=complex_types(), + ), + ), + ), + ForeachFuncInfo( + "clamp_max", + sample_inputs_func=foreach_inputs_sample_func(2, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16, torch.int64, torch.int32, torch.int8, torch.bool), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides", + dtypes=complex_types_and(torch.bool)), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_binary_op_scalar_with_overlapping_tensors", + dtypes=complex_types(), + ), + ), + ), + # note(crcrpar): forward ad not implemented. + ForeachFuncInfo( + "minimum", + sample_inputs_func=foreach_inputs_sample_func(2, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_autograd=True, + supports_inplace_autograd=False, + supports_forward_ad=False, + decorators=( + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides", + dtypes=complex_types_and(torch.bool)), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_binary_op_scalar_with_overlapping_tensors", + dtypes=complex_types(), + ), + ), + ), + # note(crcrpar): forward ad not implemented. + ForeachFuncInfo( + "maximum", + sample_inputs_func=foreach_inputs_sample_func(2, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_autograd=True, + supports_forward_ad=False, + supports_inplace_autograd=False, + decorators=( + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_outplace", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=complex_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides", + dtypes=complex_types_and(torch.bool)), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + device_type="cuda", + dtypes=(torch.complex128,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_binary_op_scalar_with_overlapping_tensors", + dtypes=complex_types(), + ), + ), + ), + ForeachFuncInfo( + "pow", + supports_alpha_param=False, + supports_scalar_self_arg=True, + sample_inputs_func=foreach_inputs_sample_func(2, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16, torch.int32, torch.int8, torch.bool), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_outplace", dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_outplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_outplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace", + dtypes=(torch.bool,),), + DecorateInfo(unittest.skip("flaky"), "TestForeach", "test_parity", device_type="cpu", dtypes=(torch.complex64,)), + DecorateInfo( + unittest.skip("failed starting on ROCm 6.2"), + "TestForeach", + "test_parity", + device_type="cuda", + dtypes=(torch.complex64,), + active_if=TEST_WITH_ROCM), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_binary_op_with_scalar_self_support", + device_type="cuda", + dtypes=(torch.bool,), + active_if=lambda kwargs: kwargs["is_fastpath"], + ), + ), + backward_requires_result=True, + ), + ForeachFuncInfo( + "copy", + sample_inputs_func=foreach_inputs_sample_func(2, False, False), + supports_out=False, + supports_forward_ad=False, + supports_autograd=False, + supports_inplace_autograd=False, + ) +] + +foreach_pointwise_op_db: list[ForeachFuncInfo] = [ + ForeachFuncInfo( + "addcmul", + sample_inputs_func=foreach_pointwise_sample_func(4, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_outplace", dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_outplace", dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides", + dtypes=(torch.bool,)), + # # Samples have complex types and inplace only works if the dtype is complex. + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", dtypes=(torch.bool,)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=integral_types() + complex_types_and(torch.bool)), + ), + ), + ForeachFuncInfo( + "addcdiv", + sample_inputs_func=foreach_pointwise_sample_func(4, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + # Samples have complex types and inplace only works if the dtype is complex. + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace", + dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=integral_types() + complex_types_and(torch.bool)), + # fails with div_cpu is not implemented with ComplexHalf + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_outplace", + dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace", + dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_outplace", + dtypes=integral_types_and(torch.bool)), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides", + dtypes=integral_types() + complex_types_and(torch.bool)), + ), + ), +] + +foreach_reduce_op_db: list[ForeachFuncInfo] = [ + ForeachFuncInfo( + "max", + sample_inputs_func=foreach_max_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + # no complex support for ordering ops like max + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_autodiff", + dtypes=(torch.complex128, torch.complex64), + ), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_foreach_reduce_large_input", + dtypes=(torch.complex128, torch.complex64), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=(torch.complex128, torch.complex64), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=(torch.complex128, torch.complex64), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=(torch.complex128, torch.complex64), + ), + ), + ), + ForeachFuncInfo( + "norm", + sample_inputs_func=foreach_norm_sample_func(1, False, False), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_meta_inplace"), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace"), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_meta_inplace"), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_inplace_all_strides"), + DecorateInfo( + unittest.expectedFailure, + "TestForeach", + "test_foreach_reduce_large_input", + device_type="cuda", + dtypes=integral_types_and(torch.bool), + ), + # AssertionError: RuntimeError not raised : Expected RuntimeError when doing an unsafe + # cast from a result of dtype torch.float32 into an out= with dtype torch.long + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)), + ), + ), +] + +foreach_other_op_db: list[ForeachFuncInfo] = [ + ForeachFuncInfo( + "lerp", + sample_inputs_func=foreach_inputs_sample_func(3, True, True), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_autograd=True, + supports_inplace_autograd=True, + supports_forward_ad=True, + decorators=( + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_meta_outplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_inplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_meta_outplace", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_inplace_all_strides", + dtypes=integral_types_and(torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, + "TestMeta", + "test_dispatch_symbolic_meta_outplace_all_strides", + dtypes=integral_types_and(torch.bool), + ), + ), + ), +] + +def reference_sign(x): + if x.dtype == np.bool_: + # `np.sign` doesn't support `bool`. + # >>> np.sign(True) + # ufunc 'sign' did not contain a loop + # with signature matching types dtype('bool') -> dtype('bool') + return np.sign(x, dtype=np.uint8).astype(np.bool_) + return np.sign(x) + + +def reference_sgn(x): + # NumPy doesn't have an equivalent to `torch.sgn` when the dtype is complex. + # For complex inputs, `np.sign` returns sign(x.real) + 0j if x.real != 0 else sign(x.imag) + 0j. + # while `torch.sgn` returns, 0 if abs(input) == 0 else input/abs(input) + if x.dtype not in [np.complex64, np.complex128]: + return reference_sign(x) + + out = (x / np.abs(x)) + if out.ndim == 0: + # Handle x == 0 case + if (x == 0): + # Can't assign to np.complex object + # So make a new one. + return np.array(complex(0, 0), dtype=x.dtype) + return out + + # Handle x == 0 case + mask = (x == 0) + out[mask] = complex(0, 0) + return out + + +def reference_sigmoid(x): + # 'scipy.special.expit' not supported for the input types + if x.dtype in [np.complex64, np.complex128]: + return (1 / (1 + np.exp(-x))) + return scipy.special.expit(x) + + +def reference_logsigmoid(x): + return np.where( + x < 0, + x - np.log1p(np.exp(x)), + -np.log1p(np.exp(-x))) + + +def reference_hardsigmoid(x): + intermediate = x / 6 + 0.5 + y = np.clip(intermediate, 0, None) + return np.where(y > 1, 1, y).astype(x.dtype) + + +def reference_lgamma(x): + # scipy.special.gammaln returns `-inf` when input is `-inf`. + # While Pytorch, C and C++, all return `inf` when input is `-inf`. + # Reference: + # https://en.cppreference.com/w/cpp/numeric/math/lgamma + # https://en.cppreference.com/w/c/numeric/math/lgamma + + # To handle the above discrepancy, + # we replace -inf with inf so values + # that were originally -inf map to inf as expected + if x.dtype.kind == 'f': + x = np.where(x == float('-inf'), np.array(float('inf'), dtype=x.dtype), x) + + out = scipy.special.gammaln(x) + + if x.dtype == np.float16: + # `scipy.special.gammaln` returns output of float32 when input is float16, + # while `torch.lgamma` preserves `float16`. But due to smaller range of float16, + # Pytorch version outputs `inf` while SciPy returns finite values. + out = out.astype(np.float16) + + return out + + +def reference_mvlgamma(x, d): + if x.dtype == np.float16: + return scipy.special.multigammaln(x, d).astype(np.float16) + + return scipy.special.multigammaln(x, d) + +def reference_softplus(input, beta=1, threshold=20): + non_linear = input * beta <= threshold + output = input.copy() + output[non_linear] = np.log(1 + np.exp(beta * input[non_linear])) / beta + return output + +def reference_gelu(X, *, approximate='none'): + def _gelu_ref(X): + return X * stats.norm.cdf(X) + + def _tanh_gelu_ref(X): + M_SQRT_2_PI = math.sqrt(2 / math.pi) + Z = M_SQRT_2_PI * (X + 0.044715 * np.power(X, 3.0)) + return 0.5 * X * (1.0 + np.tanh(Z)) + + if approximate == 'tanh': + return _tanh_gelu_ref(X) + else: + return _gelu_ref(X) + + +def reference_one_hot(a: npt.NDArray, num_classes: int = -1) -> npt.NDArray: + if num_classes == -1: + num_classes = int(np.amax(a) + 1) + + idcs = a.reshape(-1) + np.arange(0, a.size, dtype=np.int64) * num_classes + one_hot = np.zeros((a.size, num_classes), dtype=a.dtype) + np.put(one_hot, idcs, 1) + return one_hot.reshape(*a.shape, -1) + + +def reference_mse_loss(input, target, reduction="mean"): + se = (input - target) ** 2 + if reduction == "mean": + return np.mean(se) + elif reduction == "sum": + return np.sum(se) + else: # reduction == "none" + return se + + +def reference_layer_norm(inp: npt.NDArray, normalized_shape: tuple[int, ...], weight=None, bias=None, eps=1e-5): + return reference_native_layer_norm(inp, normalized_shape, weight, bias, eps)[0] + + +def reference_native_layer_norm(inp: npt.NDArray, normalized_shape: tuple[int, ...], weight, bias, eps): + feature_size = np.prod(normalized_shape) + inp_view = inp.reshape(-1, feature_size) # type: ignore[call-overload] + mean = inp_view.mean(axis=-1, keepdims=True) + var = inp_view.var(axis=-1, ddof=0, keepdims=True) + Y = (inp_view - mean) / np.sqrt(var + eps) + if weight is None and bias is not None: + Y = Y + bias.reshape(-1) + elif weight is not None and bias is None: + Y = Y * weight.reshape(-1) + elif weight is not None and bias is not None: + Y = Y * weight.reshape(-1) + bias.reshape(-1) + axis = inp.ndim - len(normalized_shape) + stat_shape = inp.shape[:axis] + (1,) * len(normalized_shape) + return Y.reshape(*inp.shape), mean.reshape(stat_shape), (1.0 / np.sqrt(var + eps)).reshape(stat_shape) + + +def reference_rms_norm(inp: npt.NDArray, normalized_shape: tuple[int, ...], weight=None, eps=None): + if eps is None: + eps = torch.finfo(numpy_to_torch_dtype(inp.dtype)).eps + feature_size = np.prod(normalized_shape) + inp_view = inp.reshape(-1, feature_size) # type: ignore[call-overload] + rms = np.sqrt((inp_view**2).mean(axis=-1, keepdims=True) + eps) + Y = inp_view / rms + if weight is not None: + Y = Y * weight.reshape(-1) + return Y.reshape(*inp.shape) + + +def reference_group_norm(inp: npt.NDArray, num_groups: int, weight=None, bias=None, eps=1e-5): + inp_view = inp + if np.prod(inp.shape) != 0: + inp_view = inp.reshape((inp.shape[0], num_groups, -1)) + mean = inp_view.mean(axis=-1, keepdims=True) + var = inp_view.var(axis=-1, ddof=0, keepdims=True) + Y = (inp_view - mean) / np.sqrt(var + eps) + Y = Y.reshape(inp.shape) + if weight is not None: + # weight is a vector of length equal to the channel + if len(Y.shape) > 2: + weight = np.expand_dims(weight, [0] + [idx + 2 for idx in range(inp.ndim - 2)]) + Y = Y * weight + if bias is not None: + # bias is a vector of length equal to the channel + if len(Y.shape) > 2: + bias = np.expand_dims(bias, [0] + [idx + 2 for idx in range(inp.ndim - 2)]) + Y = Y + bias + return Y + + +# using a custom reference function since numpy only has a string side arg (instead of right and side) and doesn't +# have an out_int32 arg. Additionally, numpy doesn't support searchsorted with ND arrays, so this splits those into +# stacked 1D cases +def reference_searchsorted(sorted_sequence, boundary, out_int32=False, right=False, side='left', sorter=None): + side = 'right' if (right or side == 'right') else 'left' + if len(sorted_sequence.shape) == 1 : + ret = np.searchsorted(sorted_sequence, boundary, side=side, sorter=sorter) + return ret.astype(np.int32) if out_int32 else ret + elif sorted_sequence.shape[0] == 0: + if sorter is not None: + sorter = sorter.flatten() + ret = np.searchsorted(sorted_sequence.flatten(), boundary.flatten(), side=side, sorter=sorter) + ret = ret.astype(np.int32) if out_int32 else ret + return ret.reshape(boundary.shape) + else: + # numpy searchsorted only supports 1D inputs so we split up ND inputs + orig_shape = boundary.shape + num_splits = np.prod(sorted_sequence.shape[:-1]) + splits = range(num_splits) + sorted_sequence, boundary = sorted_sequence.reshape(num_splits, -1), boundary.reshape(num_splits, -1) + if sorter is not None: + sorter = sorter.reshape(num_splits, -1) + + split_sequence = [sorted_sequence[i] for i in splits] + split_boundary = [boundary[i] for i in splits] + split_sorter = [sorter[i] if (sorter is not None) else None for i in splits] + + split_ret = [np.searchsorted(s_seq, b, side=side, sorter=s_sort) + for (s_seq, b, s_sort) in zip(split_sequence, split_boundary, split_sorter, strict=True)] + split_ret = [i.astype(np.int32) for i in split_ret] if out_int32 else split_ret + return np.stack(split_ret).reshape(orig_shape) + +def reference_hash_tensor(tensor, dim=(), keepdim=False, mode=0): + if mode != 0: + raise AssertionError(f"Only mode=0 (xor_sum) is supported right now, got mode={mode}") + + dtype = tensor.dtype + if dtype.kind == 'f': + tensor = tensor.astype(np.float64).view(np.uint64) + else: + tensor = tensor.astype(np.uint64) + + + if dim == (): + result = np.bitwise_xor.reduce(tensor.flatten(), keepdims=keepdim) + else: + if isinstance(dim, list): + dim = tuple(dim) + result = np.bitwise_xor.reduce(tensor, axis=dim, keepdims=keepdim) + + return result + + +def loss_reference_reduction_wrapper(fn): + def wrapper(input, target, *, size_average=None, reduce=None, reduction="mean", **other_kwargs): + if size_average is not None or reduce is not None: + raise RuntimeError( + "The keyword arguments 'size_average' and 'reduce' are deprecated and not supported by this wrapper" + ) + output = fn(input, target, **other_kwargs) + if reduction == "mean": + return np.mean(output) + elif reduction == "sum": + return np.sum(output) + else: # reduction == "none" + return output + + return wrapper + +@loss_reference_reduction_wrapper +def reference_smooth_l1_loss(input, target, beta=1.0): + diff = input - target + abs_diff = np.abs(diff) + above_threshold = abs_diff >= beta + + loss = np.empty_like(input) + loss[above_threshold] = abs_diff[above_threshold] - 0.5 * beta + loss[~above_threshold] = diff[~above_threshold] ** 2 / (2 * beta) + + return loss + +def reference_std_var(f): + """Forwards unbiased/correction kwargs as NumPy's equivalent ddof""" + g = reference_reduction_numpy(f) + + @wraps(g) + def wrapper(x: npt.NDArray, *args, **kwargs): + if 'unbiased' in kwargs and 'correction' in kwargs: + raise AssertionError("Cannot specify both 'unbiased' and 'correction' in kwargs") + + if 'unbiased' in kwargs: + kwargs['ddof'] = int(kwargs.pop('unbiased')) + elif 'correction' in kwargs: + kwargs['ddof'] = kwargs.pop('correction') + + return g(x, *args, **kwargs) + + return wrapper + +def generate_std_var_kwargs(t: torch.Tensor, **kwargs): + """Generates unbiased/correction kwargs for std/var operators""" + yield ((), {'unbiased': True}) + yield ((), {'unbiased': False}) + + # Currently, calling std with correction is only enabled when + # both dim and keepdim are provided. + if 'dim' in kwargs and 'keepdim' in kwargs: + yield ((), {'correction': 0}) + yield ((), {'correction': 1}) + + numel = torch.tensor(t.shape)[kwargs.get('dim')].prod() + yield ((), {'correction': numel // 2}) + +def error_inputs_mean(op_info, device, is_ref=False, **kwargs): + if is_ref: + err_msg1 = (r"mean\(\): could not infer output dtype. " + r"Input dtype must be either a floating point or complex dtype. " + r"Got: torch.int64") + else: + err_msg1 = (r"mean\(\): could not infer output dtype. " + r"Input dtype must be either a floating point or complex dtype. " + r"Got: Long") + yield ErrorInput( + SampleInput(make_tensor((3, 4, 5), dtype=torch.int64, device=device), []), + error_regex=err_msg1, + ) + + if is_ref: + err_msg2 = (r"mean\(\): could not infer output dtype. " + r"Optional dtype must be either a floating point or complex dtype. " + r"Got: torch.int64") + else: + err_msg2 = (r"mean\(\): could not infer output dtype. " + r"Optional dtype must be either a floating point or complex dtype. " + r"Got: Long") + yield ErrorInput( + SampleInput( + make_tensor((3, 4, 5), dtype=torch.float32, device=device), + [], + dtype=torch.int64), + error_regex=err_msg2 + ) + +# numpy implementation of torch.flatten +# unfortunately there's no np.flatten. we figure out the desired shape and call np.reshape +def reference_flatten(input, start_dim=0, end_dim=-1): + in_shape = input.shape + in_rank = len(in_shape) + for d in start_dim, end_dim: + if not ((in_rank == 0 and d in (-1, 0)) or -in_rank <= d < in_rank): + raise IndexError(f"Dimension out of range (expected to be in range of [{-in_rank}, {in_rank - 1}], but got {d}") + end_dim = end_dim if end_dim >= 0 else in_rank + end_dim + start_dim = start_dim if start_dim >= 0 else in_rank + start_dim + if in_rank == 0: + end_dim = start_dim + if end_dim < start_dim: + raise RuntimeError("flatten() has invalid args: start_dim cannot come after end_dim") + flatten_bit_dim = functools.reduce(operator.mul, in_shape[start_dim:end_dim + 1], 1) + out_shape = in_shape[:start_dim] + (flatten_bit_dim,) + in_shape[end_dim + 1:] + return np.reshape(input, out_shape) + + +def sample_inputs_alias_copy(op_info, device, dtype, requires_grad, **kwargs): + yield SampleInput(make_tensor((S,), dtype=dtype, device=device, requires_grad=requires_grad)) + yield SampleInput(make_tensor((), dtype=dtype, device=device, requires_grad=requires_grad)) + +def sample_inputs_abs(op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs): + yield from sample_inputs_elementwise_unary(op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs) + if dtype == torch.cfloat: + yield SampleInput(torch.tensor( + [ + 1e-30 + 1e-30j, + 1e30 + 1e30j, + 1e-30 + 1e30j, + 1e30 + 1e-30j, + ], + device=device, + dtype=dtype, + requires_grad=requires_grad, + )) + +# Operator database (sorted alphabetically) +op_db: list[OpInfo] = [ + UnaryUfuncInfo('abs', + aliases=('absolute', ), + ref=np.abs, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16, torch.chalf), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=sample_inputs_abs, + skips=( + DecorateInfo(unittest.skip("In-place abs not supported for complex tensors"), 'TestBwdGradients', + 'test_inplace_grad', dtypes=(torch.cdouble,)), + DecorateInfo(unittest.skip("In-place abs not supported for complex tensors"), 'TestBwdGradients', + 'test_inplace_gradgrad', dtypes=(torch.cdouble,)), + DecorateInfo(unittest.skip("In-place abs not supported for complex tensors"), 'TestFwdGradients', + 'test_inplace_forward_mode_AD', dtypes=(torch.cdouble,)), + DecorateInfo(unittest.skip("In-place abs not supported for complex tensors"), "TestSparseUnaryUfuncs", + "test_inplace", dtypes=(torch.cdouble, torch.cfloat, torch.chalf)), + # Reference: https://github.com/pytorch/pytorch/issues/49224 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', + dtypes=[torch.int8], active_if=TEST_WITH_ASAN), + # TODO: Fix test_out_arg_all_dtypes as torch.empty_like(expected_output) where expected_output=op(input) + # We can break the logic of the loop over all possible types but it is OK. + # https://github.com/pytorch/pytorch/blob/master/test/test_unary_ufuncs.py#L440-L449 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_out_arg_all_dtypes', + dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_meta_inplace', + dtypes=(torch.cdouble, torch.cfloat, torch.chalf)), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_meta_inplace', + dtypes=(torch.cdouble, torch.cfloat, torch.chalf)), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_inplace', + dtypes=(torch.cdouble, torch.cfloat, torch.chalf)), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_inplace_all_strides', + dtypes=(torch.cdouble, torch.cfloat, torch.chalf)), + # ValueError: Expected 2D tensor but got tensor with dimension: 1. + DecorateInfo( + unittest.expectedFailure, 'TestSparseCSR', 'test_sparse_csr_unary_inplace', + device_type='cpu', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestSparseCSR', 'test_sparse_csr_unary_out', + device_type='cpu', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestSparseCSR', 'test_sparse_csr_unary_inplace', + device_type='cuda', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestSparseCSR', 'test_sparse_csr_unary_out', + device_type='cuda', dtypes=(torch.complex64,) + ), + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, 'TestOutputConsistencyFullGraph', 'test_complex_output_match_opinfo_', + device_type='cpu', dtypes=(torch.complex64,) + ), + ), + supports_fwgrad_bwgrad=True, + assert_autodiffed=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_forward_ad=True), + # NOTE: CPU complex acos produces incorrect outputs (https://github.com/pytorch/pytorch/issues/42952) + UnaryUfuncInfo('acos', + aliases=('arccos', ), + ref=np.arccos, + domain=(-1, 1), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + decorators=(precisionOverride({torch.float16: 1e-2, + torch.bfloat16: 1e-1, + torch.complex64: 1e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_normal', + device_type='cuda', dtypes=[torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], active_if=IS_WINDOWS), + # Failing with wrong imaginary sign on at least some Windows jobs + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + # Failing with wrong imaginary sign on at least some Windows jobs + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_grad', + dtypes=[torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_method_grad', + dtypes=[torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_inplace_grad', + dtypes=[torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD', + dtypes=[torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_inplace_forward_mode_AD', + dtypes=[torch.cdouble], active_if=IS_WINDOWS),)), + # NOTE: the derivative for inplace acosh is not implemented + UnaryUfuncInfo('acosh', + aliases=('arccosh', ), + ref=np.arccosh, + domain=(1, None), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + decorators=(precisionOverride({torch.bfloat16: 5e-2}),), + supports_inplace_autograd=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_normal', + device_type='cuda', dtypes=[torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + # Failing with wrong imaginary sign on at least some Windows jobs + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + ), + # acosh is not defined at x < 1 (real) + reference_numerics_filter=NumericsFilter( + condition=lambda x: (x < 1 if not x.is_complex() else torch.zeros_like(x, dtype=torch.bool)), + safe_val=2)), + BinaryUfuncInfo('add', + # NumPy has no builtin reference for the alpha kwarg, but it is easy enough to emulate + ref=lambda input, other, *, alpha=1: ( + np.add(input, other) + if alpha == 1 + else np.add(input, np.multiply(alpha, other)) + ), + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, + torch.float16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + assert_autodiffed=True, + sample_inputs_func=sample_inputs_add_sub, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + supports_two_python_scalars=True, + decorators=( + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=1e-2, rtol=0)}), + 'TestBinaryUfuncs', 'test_reference_numerics'), + ), + skips=( + # boolean alpha not handled properly + DecorateInfo(unittest.expectedFailure, + 'TestNNCOpInfo', + 'test_nnc_correctness', + dtypes=(torch.bool,)), + DecorateInfo(unittest.skip("Skipped!"), + 'TestCommon', + 'test_numpy_refs', + dtypes=(torch.complex128,)), + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics_extremal_values', + dtypes=(torch.complex64, torch.complex128)), + )), + OpInfo('item', + op=lambda inp, *args, **kwargs: wrapper_set_seed(torch.Tensor.item, inp, *args, **kwargs), + ref=np.ndarray.item, + method_variant=None, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.chalf, torch.bool), + dtypesIfHpu=custom_types(torch.float32), + supports_out=False, + supports_autograd=False, + error_inputs_func=error_inputs_item, + sample_inputs_func=sample_inputs_item, + skips=( + # Error testing item function variant + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.float32, torch.complex64)), + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # RuntimeError: Composite compliance check failed with the above error. + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_operator'), + # Booleans mismatch: AssertionError: False is not true + DecorateInfo(unittest.expectedFailure, 'TestFakeTensor', 'test_fake_autocast'), + # Booleans mismatch: AssertionError: False is not true + DecorateInfo(unittest.expectedFailure, 'TestFakeTensor', 'test_fake'), + )), + OpInfo('arange', + dtypes=all_types_and(torch.bfloat16, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + supports_out=True, + supports_autograd=False, + is_factory_function=True, + error_inputs_func=error_inputs_arange, + sample_inputs_func=sample_inputs_arange, + skips=( + # https://github.com/pytorch/pytorch/issues/81774 + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + + # Lazy tensor failures + DecorateInfo(unittest.expectedFailure, 'TestLazyOpInfo', 'test_dispatched_to_lazy'), + DecorateInfo(unittest.skip("Skipped!"), 'TestLazyOpInfo', 'test_correctness'), + DecorateInfo(unittest.skip("Skipped!"), 'TestLazyOpInfo', 'test_correctness_with_reusing_ir'), + + # Exception raised from analyzeImpl at ../torch/csrc/jit/ir/alias_analysis.cpp:608 + # We don't have an op for aten::arange but it isn't a special case. + # Argument types: bool, bool, bool, int, int, Device, boo + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness'), + + # Captured graph does not contain aten::arange (succeeds on complex!) + # g: graph(): + # %25 : Long(1, strides=[1], requires_grad=0, device=cpu) = prim::Constant[value={1}]() + # return (%25) + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + )), + OpInfo('cauchy', + op=lambda inp, *args, **kwargs: wrapper_set_seed(torch.Tensor.cauchy_, inp, *args, **kwargs), + inplace_variant=torch.Tensor.cauchy_, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_autograd=False, + allow_cow_input_materialize_forward=[0], + sample_inputs_func=sample_inputs_cauchy, + error_inputs_func=error_inputs_cauchy, + skips=( + # Tests that assume input tensor has a meaningful effect on output tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + + # vmap: calling random operator not supported + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + + DecorateInfo(unittest.skip("make_traced() doesn't set seed properly!"), 'TestCommon', 'test_python_ref_executor'), + + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_quick'), + )), + OpInfo('exponential', + op=lambda inp, *args, **kwargs: wrapper_set_seed(torch.Tensor.exponential_, inp, *args, **kwargs), + inplace_variant=torch.Tensor.exponential_, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_out=False, + supports_autograd=False, + allow_cow_input_materialize_forward=[0], + sample_inputs_func=sample_inputs_exponential, + error_inputs_func=error_inputs_exponential, + skips=( + # Tests that assume input tensor has a meaningful effect on output tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + + # vmap: calling random operator not supported + DecorateInfo(unittest.expectedFailure, "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + DecorateInfo(unittest.expectedFailure, "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_quick'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + OpInfo('geometric', + op=lambda inp, *args, **kwargs: wrapper_set_seed(torch.Tensor.geometric_, inp, *args, **kwargs), + inplace_variant=torch.Tensor.geometric_, + dtypes=floating_types_and(torch.float16, torch.bfloat16, torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_out=False, + supports_autograd=False, + allow_cow_input_materialize_forward=[0], + sample_inputs_func=sample_inputs_geometric, + error_inputs_func=error_inputs_geometric, + skips=( + # Tests that assume input tensor has a meaningful effect on output tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + + # vmap: calling random operator not supported + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_quick'), + )), + OpInfo('log_normal', + op=lambda inp, *args, **kwargs: wrapper_set_seed(torch.Tensor.log_normal_, inp, *args, **kwargs), + inplace_variant=torch.Tensor.log_normal_, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_out=False, + supports_autograd=False, + allow_cow_input_materialize_forward=[0], + sample_inputs_func=sample_inputs_log_normal, + error_inputs_func=error_inputs_log_normal, + skips=( + # Tests that assume input tensor has a meaningful effect on output tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + + # vmap: calling random operator not supported + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_quick'), + )), + OpInfo('normal', + variant_test_name='in_place', + op=lambda inp, *args, **kwargs: wrapper_set_seed(torch.Tensor.normal_, inp, *args, **kwargs), + inplace_variant=torch.Tensor.normal_, + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_out=False, + supports_autograd=False, + allow_cow_input_materialize_forward=[0], + sample_inputs_func=sample_inputs_normal, + error_inputs_func=error_inputs_normal, + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.skip("Test expects tensor input"), "TestCommon", "test_noncontiguous_samples"), + + # Tests that assume input tensor has a meaningful effect on output tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_quick'), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # vmap: calling random operator not supported + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + )), + OpInfo('uniform', + op=lambda inp, *args, **kwargs: wrapper_set_seed(torch.Tensor.uniform_, inp, *args, **kwargs), + method_variant=None, + inplace_variant=torch.Tensor.uniform_, + dtypes=floating_and_complex_types_and(torch.bfloat16, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_out=False, + supports_autograd=False, + is_factory_function=False, + allow_cow_input_materialize_forward=[0], + sample_inputs_func=sample_inputs_uniform, + error_inputs_func=error_inputs_uniform, + skips=( + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # Tests that assume input tensor has a meaningful effect on output tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # aten.uniform was not decomposed + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_quick'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + BinaryUfuncInfo('clamp_max', + ref=_clamp_max_numpy, + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + supports_forward_ad=True, + supports_rhs_python_scalar=False, + supports_fwgrad_bwgrad=True, + rhs_make_tensor_kwargs=dict(exclude_zero=False), + skips=( + # RuntimeError: "max_elementwise_cuda" not implemented for 'ComplexFloat' + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='xpu'), + # dispatch to lazy test failed + DecorateInfo(unittest.expectedFailure, 'TestLazyOpInfo', 'test_dispatched_to_lazy'), + # test error disabled since rhs non-tensor python scalar is supported + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors'), + )), + BinaryUfuncInfo('clamp_min', + ref=_clamp_min_numpy, + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + supports_forward_ad=True, + supports_rhs_python_scalar=False, + supports_fwgrad_bwgrad=True, + rhs_make_tensor_kwargs=dict(exclude_zero=False), + skips=( + # RuntimeError: "min_elementwise_cuda" not implemented for 'ComplexFloat' + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='xpu'), + # dispatch to lazy test failed + DecorateInfo(unittest.expectedFailure, 'TestLazyOpInfo', 'test_dispatched_to_lazy'), + # test error disabled since rhs non-tensor python scalar is supported + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors'), + )), + BinaryUfuncInfo('mul', + aliases=('multiply',), + dtypes=all_types_and_complex_and(torch.chalf, torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_two_python_scalars=True, + error_inputs_sparse_func=error_inputs_sparse_mul, + skips=( + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors_sparse', device_type='mps'), + ), + sample_inputs_sparse_coo_func=partial(sample_inputs_sparse_mul, layout=torch.sparse_coo), + sample_inputs_sparse_csr_func=partial(sample_inputs_sparse_mul, layout=torch.sparse_csr), + sample_inputs_sparse_csc_func=partial(sample_inputs_sparse_mul, layout=torch.sparse_csc), + sample_inputs_sparse_bsr_func=partial(sample_inputs_sparse_mul, layout=torch.sparse_bsr), + sample_inputs_sparse_bsc_func=partial(sample_inputs_sparse_mul, layout=torch.sparse_bsc)), + BinaryUfuncInfo('sub', + # NumPy has no builtin reference for the alpha kwarg, but it is easy enough to emulate + ref=lambda input, other, *, alpha=1: np.subtract(input, np.multiply(alpha, other)), + aliases=('subtract',), + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_add_sub, + supports_two_python_scalars=True, + decorators=( + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-2, rtol=0), + torch.bfloat16: tol(atol=1e-5, rtol=5e-3), + torch.complex32: tol(atol=1e-5, rtol=1e-3)}), + 'TestBinaryUfuncs', 'test_reference_numerics'), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=1e-2, rtol=0)}), + 'TestCommon', 'test_complex_half_reference_testing', device_type='cpu'), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=5e-3, rtol=0)}), + 'TestDecomp', 'test_comprehensive', device_type='cpu'), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=5e-3, rtol=0)}), + 'TestDecomp', 'test_quick', device_type='cpu'), + ), + skips=( + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics', + dtypes=(torch.uint8,)), + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.uint8,)), + )), + OpInfo('addmm', + # This addmm OpInfo is for when alpha and beta are not both equal to 1. + # alpha=beta=1 is tested in the following opinfo, because that special case will + # trigger addmm being decomposed by a jit pass. + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfROCM=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + sample_inputs_func=sample_inputs_addmm, + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=2e-3)}), + "TestConsistency", "test_output_grad_match", device_type="mps"), + # RuntimeError: value cannot be converted to type double without overflow + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + OpInfo('addmm', + # When alpha=beta=1 as compile-time constants, JIT will decompose addmm into mm and add. + variant_test_name='decomposed', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + autodiff_nonfusible_nodes=['aten::add', 'aten::mm'], + sample_inputs_func=partial(sample_inputs_addmm, alpha=1, beta=1), + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + # https://github.com/pytorch/pytorch/issues/71784 + DecorateInfo(unittest.skip('Skipped!'), 'TestNNCOpInfo', 'test_nnc_correctness', + device_type='cpu', dtypes=(torch.float16,)), + )), + OpInfo('addmv', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.complex64, torch.complex128, + torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[ + DecorateInfo( + toleranceOverride({torch.half: tol(atol=1e-5, rtol=3e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cpu'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-5, rtol=3e-6)}), + "TestConsistency", "test_output_match", device_type="mps"), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-5, rtol=3e-6)}), + "TestConsistency", "test_output_grad_match", device_type="mps"), + ], + sample_inputs_func=sample_inputs_addmv), + OpInfo('addbmm', + ref=lambda M, batch1, batch2, beta=1, alpha=1: np.add(np.multiply(np.asarray(beta, dtype=M.dtype), M), + np.multiply(np.asarray(alpha, dtype=batch1.dtype), + np.sum(np.matmul(batch1, batch2), axis=0))), + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, + *[torch.bfloat16] + if SM53OrLater or TEST_WITH_ROCM else []), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1.3e-05, rtol=1.3e-05), + torch.complex64: tol(atol=1e-05, rtol=1.2e-03)}), + 'TestCommon', 'test_numpy_refs'), + # MPS has slightly worse precision. Is this acceptable? + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1.3e-04, rtol=1.3e-04), + torch.complex64: tol(atol=1e-05, rtol=1.2e-03)}), + 'TestCommon', 'test_numpy_ref_mps'), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-5, rtol=1e-5), + torch.bfloat16: tol(atol=2e-1, rtol=6e-1)}), + 'TestConsistency', + 'test_output_match', + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1.5e-05, rtol=1e-05)}), + 'TestCommon', 'test_out'), + DecorateInfo( + toleranceOverride({torch.half: tol(atol=6e-3, rtol=1e-2)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cpu'), + ], + skips=( + # NVIDIA only assures that bfloat16 is supported by bmm if SM >= 5.3 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_dtypes', device_type='cuda', active_if=not SM53OrLater), + # addbmm does not correctly warn when resizing out= inputs + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # https://github.com/pytorch/pytorch/issues/55907 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + # AssertionError: RuntimeError not raised : Expected RuntimeError + # when doing an unsafe cast from a result of dtype torch.float32 + # into an out= with dtype torch.long + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', dtypes=(torch.float32,), device_type='mps'), + ), + sample_inputs_func=sample_inputs_addbmm), + OpInfo('baddbmm', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.complex64, torch.complex128, + torch.bfloat16), + backward_dtypesIfCUDA=floating_types_and(torch.float16, + *[torch.bfloat16] if SM53OrLater or TEST_WITH_ROCM else [], + torch.complex64, torch.complex128), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[ + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-05, rtol=1.2e-03)}), + 'TestCommon', 'test_variant_consistency_eager', device_type='cuda'), + # Higher differences starting with Zen3 or Alder Lake + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=4e-05, rtol=4e-06)}), + 'TestDecomp', 'test_quick', device_type='cpu'), + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-05, rtol=1.2e-03)}), + 'TestMathBits', 'test_conj_view', device_type='cuda'), + ], + sample_inputs_func=sample_inputs_baddbmm, + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + )), + OpInfo('dot', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfMPS=all_types_and_complex_and(torch.float16, torch.bfloat16, torch.bool), + assert_autodiffed=True, + sample_inputs_func=sample_inputs_dot_vdot, + error_inputs_func=error_inputs_dot_vdot, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + )), + OpInfo('vdot', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfMPS=all_types_and_complex_and(torch.float16, torch.bfloat16, torch.bool), + sample_inputs_func=sample_inputs_dot_vdot, + error_inputs_func=error_inputs_dot_vdot, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + )), + OpInfo('bmm', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, + *[torch.bfloat16] + if SM53OrLater or TEST_WITH_ROCM else []), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + assert_jit_shape_analysis=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # NVIDIA only assures that bfloat16 is supported by bmm if SM >= 5.3 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_dtypes', device_type='cuda', active_if=not SM53OrLater), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-5, rtol=1e-5)}), + "TestCommon", "test_out"), + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_out', + device_type='mps', dtypes=(torch.float32,) + ), + ), + sample_inputs_func=sample_inputs_bmm), + OpInfo('mv', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_mv), + OpInfo('addr', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + # Reference: https://github.com/pytorch/pytorch/issues/50747 + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/50747 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16)), + # FIXME: AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + # FIXME: AssertionError: RuntimeError not raised : Expected RuntimeError when doing an unsafe cast + # from a result of dtype torch.float32 into an out= with dtype torch.long + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + # FIXME: RuntimeError: value cannot be converted to type double without overflow + DecorateInfo( + unittest.expectedFailure, + 'TestCommon', + 'test_noncontiguous_samples', + device_type='mps', + dtypes=(torch.complex64,)), + # RuntimeError: MPS device does not support addr for non-float input + DecorateInfo( + unittest.expectedFailure, + 'TestCommon', + 'test_noncontiguous_samples', + device_type='mps', + dtypes=(torch.int64,)), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + ), + sample_inputs_func=sample_inputs_addr, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), + OpInfo('addcmul', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # TODO: update sample inputs with for_inplace_variant kwarg to support this test + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + # AssertionError: The supported dtypes for addcmul on device type mps are incorrect! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + ), + sample_inputs_func=sample_inputs_addcmul_addcdiv, + reference_inputs_func=partial( + reference_inputs_elementwise_ternary, sample_inputs_func=reference_inputs_addcmul_addcdiv)), + OpInfo('addcdiv', + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # TODO: update sample inputs with for_inplace_variant kwarg to support this test + DecorateInfo(unittest.expectedFailure, + 'TestCommon', + 'test_variant_consistency_eager'), + ), + sample_inputs_func=sample_inputs_addcmul_addcdiv, + reference_inputs_func=partial( + reference_inputs_elementwise_ternary, sample_inputs_func=reference_inputs_addcmul_addcdiv)), + UnaryUfuncInfo('asin', + aliases=('arcsin', ), + ref=np.arcsin, + domain=(-1, 1), + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-05, rtol=1e-03)}), + 'TestUnaryUfuncs', device_type='cuda' + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=8e-5, rtol=4e-5)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cuda' + ), + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=5e-05, rtol=2e-05)}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cpu' + ), + precisionOverride({torch.bfloat16: 1e-2}), + ], + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + )), + # NOTE: derivative for inplace asinh is not implemented + UnaryUfuncInfo('asinh', + aliases=('arcsinh', ), + ref=np.arcsinh, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + decorators=(precisionOverride({torch.bfloat16: 5e-2}),), + supports_inplace_autograd=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_normal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4), + torch.float16: tol(atol=1e-4, rtol=1e-4), + torch.bfloat16: tol(atol=1e-4, rtol=1e-4), + torch.complex64: tol(atol=1e-4, rtol=1e-4)}), + "TestConsistency", "test_output_match", device_type="mps"), + )), + UnaryUfuncInfo('atan', + aliases=('arctan', ), + ref=np.arctan, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + promotes_int_to_float=True, + decorators=(precisionOverride({torch.bfloat16: 1e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.cfloat, torch.cdouble], active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=[torch.cfloat, torch.cdouble], active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', + dtypes=[torch.cfloat, torch.cdouble], active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + )), + BinaryUfuncInfo('atan2', + aliases=('arctan2',), + dtypes=all_types_and(torch.bool, torch.bfloat16, torch.half), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + supports_rhs_python_scalar=False, + skips=( + # Incorrectly attempts to use a scalar for the second argument + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_jit_alias_remapping'), + )), + UnaryUfuncInfo('atanh', + aliases=('arctanh', ), + ref=np.arctanh, + domain=(-1, 1), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + decorators=[ + precisionOverride({torch.bfloat16: 1e-2}), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=9e-3, rtol=8e-5)}), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda" + ), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=2e-3)}), + "TestConsistency", "test_output_grad_match", device_type="mps"), + ], + supports_inplace_autograd=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cfloat], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + )), + OpInfo('allclose', + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + ref=np.allclose, + supports_autograd=False, + supports_forward_ad=False, + sample_inputs_func=sample_inputs_allclose, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCudaFuserOpInfo'), + ), + supports_out=False), + OpInfo('broadcast_to', + ref=np.broadcast_to, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_broadcast_to), + OpInfo('broadcast_shapes', + op=torch.broadcast_shapes, + ref=np.broadcast_shapes if np.lib.NumpyVersion(np.__version__) >= '1.20.0' else None, + dtypes=_dispatch_dtypes((torch.float32,)), + supports_out=False, + supports_gradgrad=False, + assert_autodiffed=False, + supports_autograd=False, + supports_scripting=False, + sample_inputs_func=sample_inputs_broadcast_shapes, + skips=( + # https://github.com/pytorch/pytorch/issues/64997 + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # skip dtype tests since broadcast_shape is not device dependent. + # having dtypes limited to torch.float32 would cause test_dtypes to report unexpected success + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_dtypes'), + # skip these tests since we have non tensor input + DecorateInfo(unittest.skip('Skipped!'), "TestCommon", "test_noncontiguous_samples"), + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip('Skipped!'), 'TestJit', 'test_variant_consistency_jit'), + )), + OpInfo('broadcast_tensors', + ref=np.broadcast_arrays, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=sample_inputs_broadcast_tensors, + reference_inputs_func=reference_inputs_broadcast_tensors, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + skips=( + # https://github.com/pytorch/pytorch/issues/64997 + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # JIT does not support variadic tensors. + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":252, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=[torch.float32]), + )), + OpInfo('block_diag', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # Default batching rule in core doesn't work for ops with TensorList args + check_batched_forward_grad=False, + skips=( + # https://github.com/pytorch/pytorch/issues/64997 + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # JIT does not support variadic tensors. + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":252, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=[torch.float32]), + ), + sample_inputs_func=sample_inputs_block_diag), + UnaryUfuncInfo('bitwise_not', + ref=np.bitwise_not, + dtypes=integral_types_and(torch.bool), + dtypesIfHpu=custom_types(torch.bool), + operator_variant=operator.invert, + supports_autograd=False), + BinaryUfuncInfo('bitwise_left_shift', + op=torch.bitwise_left_shift, + dtypes=integral_types(), + dtypesIfCUDA=integral_types(), + dtypesIfHpu=custom_types(torch.int32, torch.int8, torch.bool), + operator_variant=operator.lshift, + inplace_operator_variant=operator.ilshift, + supports_autograd=False, + supports_one_python_scalar=True, + rhs_make_tensor_kwargs=dict(low=0), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_type_promotion'), + # https://github.com/pytorch/pytorch/issues/70904 + DecorateInfo(unittest.skip("Some inputs produce undefined outputs"), 'TestCommon', 'test_compare_cpu'), + # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + BinaryUfuncInfo('bitwise_right_shift', + op=torch.bitwise_right_shift, + dtypes=integral_types(), + dtypesIfCUDA=integral_types(), + dtypesIfHpu=custom_types(torch.int32, torch.int8, torch.bool), + operator_variant=operator.rshift, + inplace_operator_variant=operator.irshift, + supports_autograd=False, + supports_one_python_scalar=True, + rhs_make_tensor_kwargs=dict(low=0), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_type_promotion'), + # https://github.com/pytorch/pytorch/issues/70904 + DecorateInfo(unittest.skip("Some inputs produce undefined outputs"), 'TestCommon', 'test_compare_cpu'), + # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + OpInfo('combinations', + op=torch.combinations, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + supports_out=False, + sample_inputs_func=sample_inputs_combinations), + OpInfo('cartesian_prod', + op=torch.cartesian_prod, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_cartesian_prod, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":270 + DecorateInfo(unittest.expectedFailure, + 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + )), + OpInfo('cdist', + dtypes=floating_types(), + dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_gradgrad=False, + assert_autodiffed=False, + sample_inputs_func=sample_inputs_cdist, + skips=( + # NotImplementedError: The operator 'aten::_cdist_backward' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + DecorateInfo(toleranceOverride( + { + torch.float16: tol(atol=1e-3, rtol=2e-2), + torch.float32: tol(atol=1.6e-5, rtol=1.6e-6), + }), + 'TestConsistency', device_type='mps'), + )), + UnaryUfuncInfo('ceil', + ref=np.ceil, + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + dtypesIfMPS=all_types_and(torch.half, torch.bfloat16, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo(unittest.expectedFailure, + 'TestNNCOpInfo', + 'test_nnc_correctness', + dtypes=tuple(t for t in integral_types() if t != torch.uint8)), + ), + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + assert_autodiffed=True), + OpInfo('cholesky', + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_cholesky, + gradcheck_wrapper=gradcheck_wrapper_hermitian_input, + skips=( + # linalg.solve.triangular(); Only float is supported! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error', + device_type='mps', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=(torch.complex64,) + ), + ), + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack],), + OpInfo('cholesky_inverse', + dtypes=floating_and_complex_types(), + backward_dtypes=floating_and_complex_types(), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + check_batched_gradgrad=True, + sample_inputs_func=sample_inputs_linalg_cholesky_inverse, + gradcheck_wrapper=gradcheck_wrapper_triangular_input_real_positive_diagonal, + decorators=[ + skipCUDAIfNoMagma, + skipCPUIfNoLapack, + DecorateInfo( + toleranceOverride({ + torch.float32: tol(atol=5e-03, rtol=1e-04) + }), + 'TestCommon', device_type='cpu', + ), + DecorateInfo( + toleranceOverride({ + torch.float32: tol(atol=5e-03, rtol=1e-04) + }), + 'TestEagerFusionOpInfo', device_type='cpu', + ), + ], + skips=( + # Strides are not the same! Original strides were ((4, 2, 1),) and strides are now ((4, 1, 2),) + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # torch._C._LinAlgError: linalg.cholesky: (Batch element 0): The factorization could not be completed + # because the input is not positive-definite (the leading minor of order 2 is not positive-definite). + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=(torch.complex64,) + ), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + # RuntimeError: cholesky_inverse: MPS only supports float type! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error', + device_type='mps', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.complex64,) + ),), + ), + OpInfo('cholesky_solve', + op=torch.cholesky_solve, + dtypes=floating_and_complex_types(), + dtypesIfMPS=custom_types(torch.float32), + backward_dtypesIfMPS=custom_types(torch.float32), + sample_inputs_func=sample_inputs_cholesky_solve, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_wrapper=lambda *args, **kwargs: gradcheck_wrapper_triangular_input(*args, idx=1, **kwargs), + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack]), + OpInfo('chunk', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + sample_inputs_func=sample_inputs_chunk, + reference_inputs_func=reference_inputs_chunk, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + OpInfo('unsafe_chunk', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + sample_inputs_func=sample_inputs_chunk, + check_batched_forward_grad=False, + reference_inputs_func=reference_inputs_chunk, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + OpInfo('clone', + ref=np.copy, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + sample_inputs_func=sample_inputs_clone_contiguous, + reference_inputs_func=reference_inputs_clone_contiguous, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + skips=( + # TypeError: _copy_dispatcher() got an unexpected keyword argument 'memory_format' + # (NumPy reference needs to be extended with memory_format) + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref_mps'), + ),), + OpInfo('contiguous', + op=lambda x, *args, **kwargs: x.contiguous(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + sample_inputs_func=sample_inputs_clone_contiguous, + reference_inputs_func=reference_inputs_clone_contiguous, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + autodiff_fusible_nodes=['aten::contiguous'], + assert_jit_shape_analysis=True, + supports_out=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + )), + OpInfo('sum_to_size', + op=lambda x, *args, **kwargs: x.sum_to_size(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_sum_to_size, + error_inputs_func=error_inputs_sum_to_size, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float,)), + )), + OpInfo('clamp', + aliases=('clip',), + ref=_clamp_numpy, + dtypes=all_types_and(torch.bfloat16, torch.half), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + sample_inputs_func=sample_inputs_clamp, + reference_inputs_func=partial(reference_inputs_elementwise_ternary, sample_inputs_func=sample_inputs_clamp), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # NNC appear to not handle boolean clamp + DecorateInfo(unittest.expectedFailure, + 'TestNNCOpInfo', + 'test_nnc_correctness', + dtypes=(torch.bool,)), + # MPS does not support float64, while numpy does internal computations in float64. + # See https://github.com/pytorch/pytorch/blob/3c1cf03fde145bdbe1f5ffb81765d076c10b4c04/test/test_ops.py#L260-L264 + DecorateInfo(unittest.expectedFailure, + 'TestCommon', + 'test_numpy_ref_mps'), + # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + UnaryUfuncInfo('positive', + ref=np.positive, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + ), + UnaryUfuncInfo('conj', + ref=np.conj, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, + torch.half, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.int32), + supports_sparse=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + supports_out=False), + UnaryUfuncInfo('conj_physical', + decomp_aten_name='_conj_physical', + ref=np.conj, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, + torch.half, torch.chalf), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + skips=( + # RuntimeError: inputSet && outputSet + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":118, + # please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32, )), + DecorateInfo(unittest.skip("Skipped! conj_physical_ not implemented for sparse"), + 'TestSparseUnaryUfuncs', 'test_inplace'), + # RuntimeError: false INTERNAL ASSERT FAILED at + # "/Users/kurtamohler/develop/pytorch-1/aten/src/ATen/native/DispatchStub.cpp":276 + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=(torch.complex64,) + ), + # RuntimeError: Expected self.is_complex() to be true, but got false. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + + )), + OpInfo('resolve_conj', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=sample_inputs_view_as_real, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + ), + OpInfo('resolve_neg', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=sample_inputs_view_as_real, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + ), + OpInfo('view_as_real', + dtypes=complex_types(), + supports_forward_ad=True, + supports_out=False, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_view_as_real, + test_conjugated_samples=False, + ), + OpInfo('view_as_complex', + dtypes=floating_types_and(torch.half), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + test_neg_view=False, + sample_inputs_func=sample_inputs_view_as_complex, + skips=( + # RuntimeError: Tensor must have a last dimension with stride 1 + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_noncontiguous_samples"), + # RuntimeError: "eq_cpu" not implemented for 'ComplexHalf' + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness', dtypes=(torch.half,)), + # RuntimeError: view size is not compatible with input tensor's size and stride + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides"), + )), + BinaryUfuncInfo('complex', + dtypes=floating_types_and(torch.half), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_rhs_python_scalar=False, + error_inputs_func=error_inputs_complex, + skips=( + # Tests don't account for complex's type promotion semantics + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_binary_ufuncs_mixed_dtype'),)), + BinaryUfuncInfo('copysign', + sample_inputs_func=sample_inputs_copysign, + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + promotes_int_to_float=True, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True), + OpInfo('corrcoef', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_corrcoef, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + ), + supports_out=False), + UnaryUfuncInfo('cos', + ref=np.cos, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + handles_large_floats=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + decorators=(precisionOverride({torch.bfloat16: 1e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', active_if=IS_WINDOWS), + # This fails on CUDA but passes on ROCm + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=(torch.cdouble,), device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_MACOS), + # AssertionError: Tensor-likes are not close! + # Greatest absolute difference: nan at index (700,) (up to 1e-05 allowed) + # Greatest relative difference: nan at index (700,) (up to 0.001 allowed) + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cuda', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + )), + UnaryUfuncInfo('cosh', + ref=np_unary_ufunc_integer_promotion_wrapper(np.cosh), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/48641 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.int8]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=[torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_MACOS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_MACOS), + # AssertionError: Tensor-likes are not close! + # Greatest absolute difference: nan at index (6000,) (up to 1e-05 allowed) + # Greatest relative difference: nan at index (6000,) (up to 0.001 allowed) + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cuda', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + )), + OpInfo('cov', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_cov, + error_inputs_func=error_inputs_cov, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + # Float did not match double + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_fn_grad'), + # Jacobian mismatch + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_fn_gradgrad'), + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_forward_mode_AD'), + DecorateInfo(unittest.skip("Barely fails"), 'TestFwdGradients', 'test_fn_fwgrad_bwgrad'), + # JIT test not working for tensor kwargs (https://github.com/pytorch/pytorch/issues/58507) + # RuntimeError: + # undefined value tensor: + # File "", line 3 + # def the_method(i0): + # return torch.cov(i0, correction=0, fweights=None, aweights=tensor([0.0518, 0.4681], dtype=torch.float32, requires_grad=True)) # noqa: B950 + # ~~~~~~ <--- HERE + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=8e-3, rtol=1.4e-3)}), + "TestInductorOpInfo", "test_comprehensive", device_type="cpu"), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=3e-4, rtol=1e-4)}), + "TestConsistency", "test_output_grad_match", device_type="mps"), + )), + OpInfo('cross', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + sample_inputs_func=sample_inputs_cross, + supports_fwgrad_bwgrad=True, + supports_out=True, + supports_forward_ad=True, + ), + OpInfo('cumsum', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # cumsum does not handle correctly out= dtypes + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bool} + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + ), + sample_inputs_func=sample_inputs_cumulative_ops), + OpInfo('cumprod', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # cumprod does not handle correctly out= dtypes + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bool} + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + ), + # gradgradcheck fails in fast_mode=True: #56275 + sample_inputs_func=sample_inputs_cumprod, + gradcheck_fast_mode=False), + OpInfo('cummax', + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_cumulative_ops, supports_dtype_kwargs=False), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), + OpInfo('cummin', + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_cumulative_ops, supports_dtype_kwargs=False), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), + UnaryUfuncInfo('deg2rad', + ref=np.radians, + decorators=(precisionOverride({torch.bfloat16: 7e-1, + torch.float16: 7e-1}),), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + promotes_int_to_float=True), + OpInfo('diff', + op=torch.diff, + # np.diff has np._NoValue as default values for prepend and append, compare_with_reference breaks if prepend/append + # are set as None when converting to numpy + ref=lambda input, n=1, dim=-1, prepend=np._NoValue, append=np._NoValue: ( + np.diff(input, n, dim, np._NoValue if prepend is None else prepend, np._NoValue if append is None else append) + ), + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_diff, + error_inputs_func=error_inputs_diff, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + skips=( + )), + BinaryUfuncInfo('div', + aliases=('divide',), + variant_test_name='no_rounding_mode', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + promotes_int_to_float=True, + supports_fwgrad_bwgrad=True, + supports_two_python_scalars=True, + assert_autodiffed=True, + rhs_make_tensor_kwargs=dict(exclude_zero=True),), + BinaryUfuncInfo('div', + aliases=('divide',), + variant_test_name='trunc_rounding', + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + sample_kwargs=lambda device, dtype, input: + ({"rounding_mode": "trunc"}, {"rounding_mode": "trunc"}), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_two_python_scalars=True, + assert_autodiffed=True, + rhs_make_tensor_kwargs=dict(exclude_zero=True), + decorators=( + # See https://github.com/pytorch/pytorch/issues/111126 + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + ), + skips=( + # RuntimeError: MALFORMED INPUT: Unhandled node kind (in computeValue): aten::div + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_working'), + # FIXME: + # torch.autograd.gradcheck.GradcheckError: Jacobian mismatch for + # output 0 with respect to input 1, + # numerical:tensor(-17746.9307, dtype=torch.float64) + # analytical:tensor(0., dtype=torch.float64) + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', + 'test_fn_grad', device_type='cpu', + dtypes=(torch.float64,)), + # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + BinaryUfuncInfo('div', + aliases=('divide',), + variant_test_name='floor_rounding', + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + sample_kwargs=lambda device, dtype, input: + ({"rounding_mode": "floor"}, {"rounding_mode": "floor"}), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_two_python_scalars=True, + assert_autodiffed=True, + rhs_make_tensor_kwargs=dict(exclude_zero=True), + decorators=( + # See https://github.com/pytorch/pytorch/issues/111126 + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + ), + skips=( + # RuntimeError: MALFORMED INPUT: Unhandled node kind (in computeValue): aten::div + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_working'), + # FIXME: + # torch.autograd.gradcheck.GradcheckError: Jacobian mismatch for + # output 0 with respect to input 1, + # numerical:tensor(-17746.9307, dtype=torch.float64) + # analytical:tensor(0., dtype=torch.float64) + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', + 'test_fn_grad', + dtypes=(torch.float64,), + device_type='cpu'), + # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + BinaryUfuncInfo('true_divide', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_forward_ad=True, + promotes_int_to_float=True, + supports_fwgrad_bwgrad=True, + supports_two_python_scalars=True, + rhs_make_tensor_kwargs=dict(exclude_zero=True)), + OpInfo('equal', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + ref=lambda input, other: (input == other).all(), + sample_inputs_func=sample_inputs_equal, + supports_autograd=False, + supports_tracing=False, + ), + UnaryUfuncInfo('exp', + ref=np_unary_ufunc_integer_promotion_wrapper(np.exp), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/48010 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), + ), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True), + OpInfo('expand', + op=lambda self, shape: self.expand(shape), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + sample_inputs_func=sample_inputs_expand, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + supports_out=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + )), + OpInfo('expand_as', + op=lambda self, other: self.expand_as(other), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_expand_as, + supports_out=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'),), + ), + OpInfo('expand_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_expand, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + supports_out=True, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + )), + OpInfo('diag', + ref=np.diag, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_diag, + error_inputs_func=error_inputs_diag), + OpInfo('diag_embed', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + supports_out=False, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_diagonal_diag_embed, + reference_inputs_func=reference_inputs_diagonal_diag_embed, + error_inputs_func=error_inputs_diagonal_diag_embed), + OpInfo('diagonal', + aten_backward_name='diagonal_backward', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_diagonal_diag_embed, + reference_inputs_func=reference_inputs_diagonal_diag_embed, + error_inputs_func=error_inputs_diagonal_diag_embed), + OpInfo('diagonal_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_diagonal_diag_embed, + reference_inputs_func=reference_inputs_diagonal_diag_embed, + error_inputs_func=error_inputs_diagonal_diag_embed), + OpInfo('diagonal_scatter', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_diagonal_scatter), + OpInfo('alias_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + sample_inputs_func=sample_inputs_alias_copy, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=True), + BinaryUfuncInfo('eq', + ref=np.equal, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + always_returns_bool=True, + supports_autograd=False, + sample_inputs_func=sample_inputs_comparison_ops, + ), + BinaryUfuncInfo('fmax', + op=torch.fmax, + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_rhs_python_scalar=False, + skips=( + # RuntimeError: "max_elementwise_cuda" not implemented for 'ComplexFloat' + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_type_promotion'), + )), + BinaryUfuncInfo('fmin', + op=torch.fmin, + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_rhs_python_scalar=False, + skips=( + # RuntimeError: "min_elementwise_cuda" not implemented for 'ComplexFloat' + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_type_promotion'), + )), + BinaryUfuncInfo('fmod', + ref=np.fmod, + dtypes=all_types_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + dtypesIfMPS=all_types_and(torch.float16, torch.bfloat16, torch.bool), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=None, + rhs_make_tensor_kwargs={'exclude_zero': True}, + decorators=( + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_contig_vs_every_other', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_non_contig', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.uint8,)), + # FIXME: + # torch.autograd.gradcheck.GradcheckError: Jacobian mismatch for + # output 0 with respect to input 1, + # numerical:tensor(101.6283, dtype=torch.float64) + # analytical:tensor(-18.3575, dtype=torch.float64) + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', + 'test_fn_grad', + dtypes=(torch.float64,), + device_type='cpu'), + )), + BinaryUfuncInfo('remainder', + ref=np.remainder, + dtypes=all_types_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + dtypesIfMPS=all_types_and(torch.float16, torch.bfloat16, torch.bool), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=None, + operator_variant=operator.mod, + inplace_operator_variant=operator.imod, + supports_one_python_scalar=True, + rhs_make_tensor_kwargs={'exclude_zero': True}, + decorators=( + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_contig_vs_every_other', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_non_contig', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.uint8,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', + 'test_nnc_correctness', + dtypes=(torch.bfloat16,)), + # Fails on XLA + # False is not true : Tensors failed to compare as equal! + # Attempted to compare equality of tensors with different dtypes + DecorateInfo(unittest.skip("Skipped!"), 'TestOpInfo', device_type='xla', dtypes=(torch.long,)), + # FIXME: + # torch.autograd.gradcheck.GradcheckError: Jacobian mismatch for + # output 0 with respect to input 1, + # numerical:tensor(102.4676, dtype=torch.float64) + # analytical:tensor(-17.5182, dtype=torch.float64) + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', + 'test_fn_grad', device_type='cpu', + dtypes=(torch.float64,)), + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=5e-4, rtol=3e-3), + }), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda" + ), + )), + UnaryUfuncInfo('frac', + ref=lambda x: np.modf(x)[0], + dtypes=floating_types_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=(torch.bfloat16, torch.float16, torch.float32, torch.float64)), + # 76047 + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness', + dtypes=(torch.bfloat16, torch.float32, torch.float64)), + )), + OpInfo('stft', + decorators=[ + skipCPUIfNoFFT, + DecorateInfo(unittest.skip("Skipped! stft does not match the native function"), + 'TestJit', 'test_variant_consistency_jit'), + ], + dtypes=floating_and_complex_types(), + dtypesIfMPS=floating_and_complex_types_and(torch.float16), + sample_inputs_func=sample_inputs_stft, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_out=False, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + ), + OpInfo('istft', + dtypes=complex_types(), + sample_inputs_func=sample_inputs_istft, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_out=False, + decorators=( + DecorateInfo(unittest.skip("Skipped! istft does not match the native function"), + 'TestJit', 'test_variant_consistency_jit'), + ), + skips=( + skipCPUIfNoFFT, + # gradcheck fails on ROCm (gh-68429) + # grad is computed improperly (probably for weights tensor) + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_fn_grad'), + # Pre-existing condition (calls .item); needs to be fixed + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_backward'), + )), + UnaryUfuncInfo('floor', + ref=np.floor, + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + dtypesIfMPS=all_types_and(torch.half, torch.bfloat16, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo(unittest.expectedFailure, + 'TestNNCOpInfo', + 'test_nnc_correctness', + dtypes=tuple(t for t in integral_types() if t != torch.uint8)), + ), + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + assert_autodiffed=True), + OpInfo('flip', + op=torch.flip, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + sample_inputs_func=sample_inputs_flip, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + OpInfo('fliplr', + op=torch.fliplr, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_fliplr_flipud, + error_inputs_func=error_inputs_fliplr, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + OpInfo('flipud', + op=torch.flipud, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_fliplr_flipud, + error_inputs_func=error_inputs_flipud, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + OpInfo('sparse.sampled_addmm', + dtypes=floating_and_complex_types(), + supports_autograd=True, + sample_inputs_func=sample_inputs_sparse_sampled_addmm, + decorators=[ + skipCPUIfNoMklSparse, + skipXPU], + skips=( + # NotImplementedError: Tensors of type SparseCsrTensorImpl do not have is_contiguous + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples'), + # RuntimeError: Sparse CSR tensors do not have strides. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out'), + DecorateInfo(unittest.skip("Skipped!"), 'TestTags', 'test_tags'), + # RuntimeError: sampled_addmm: Expected result to have sparse csr layout, but got Strided + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out_warning'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_operator'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_backward'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError: unsupported memory format option Preserve + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError: sparse_mask does not support automatic differentiation for outputs with complex dtype + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_fn_fwgrad_bwgrad'), + # ValueError: Sparse output is not supported at gradcheck yet. Please call to_dense(masked_grad=...) ... + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_grad'), + # RuntimeError: sparse_mask does not support automatic differentiation for outputs with complex dtype. + # RuntimeError: Sparse CSR tensors do not have is_contiguous + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_gradgrad'), + # ValueError: Sparse output is not supported at gradcheck yet. Please call to_dense(masked_grad=...) ... + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD'), + # NotImplementedError: Could not run 'aten::sparse_sampled_addmm' with arguments from the 'SparseCsrMeta' backend. + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_dispatch_meta_outplace'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_dispatch_symbolic_meta_outplace'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_meta_outplace'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_dispatch_symbolic_meta_outplace_all_strides'), + DecorateInfo(unittest.skip("Skipped!"), 'TestFakeTensor', 'test_fake_crossref_backward_no_amp'), + # NotImplementedError: The operator 'aten::_to_sparse_csr' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + )), + OpInfo('sparse.mm', + dtypes=floating_types_and(torch.bfloat16, torch.float16), + variant_test_name='reduce', + supports_autograd=True, + supports_out=False, + supports_gradgrad=False, + supports_forward_ad=False, + sample_inputs_func=sample_inputs_sparse_mm_reduce, + decorators=[onlyCPU], + skips=( + # NotImplementedError: Tensors of type SparseCsrTensorImpl do not have is_contiguous + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples'), + # RuntimeError: Sparse CSR tensors do not have strides. + DecorateInfo(unittest.skip("Skipped!"), 'TestTags', 'test_tags'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_operator'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_backward'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError: unsupported memory format option Preserve + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # ValueError: Sparse output is not supported at gradcheck yet. Please call to_dense(masked_grad=...) ... + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_fn_fwgrad_bwgrad'), + # RuntimeError: Sparse CSR tensors do not have is_contiguou + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_grad'), + # ValueError: Sparse output is not supported at gradcheck yet. Please call to_dense(masked_grad=...) ... + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_gradgrad'), + # RuntimeError: Sparse CSR tensors do not have strides + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD'), + # ValueError: Sparse output is not supported at gradcheck yet. Please call to_dense(masked_grad=...) ... + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_fail_gradgrad'), + # NotImplementedError: Could not run 'aten::_sparse_mm_reduce_impl' with arguments from the 'SparseCsrMeta' backend + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_dispatch_meta_outplace'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_dispatch_symbolic_meta_outplace'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_meta_outplace'), + )), + UnaryUfuncInfo('i0', + ref=np_unary_ufunc_integer_promotion_wrapper( + scipy.special.i0) if TEST_SCIPY else None, + aliases=('special.i0',), + decorators=(precisionOverride({torch.bfloat16: 3e-1, + torch.float16: 5e-1}),), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + sample_inputs_func=sample_inputs_i0_i1, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=(torch.int8,)), + )), + BinaryUfuncInfo('floor_divide', + ref=_floor_divide_np, + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + dtypesIfMPS=all_types_and(torch.half, torch.bfloat16, torch.bool), + + supports_autograd=False, + rhs_make_tensor_kwargs=dict(exclude_zero=True), + supports_two_python_scalars=True, + skips=( + # AssertionError: Results of original model and exported/imported version of model differed + DecorateInfo(unittest.skip('Skipped!'), 'TestJit', 'test_variant_consistency_jit'), + # bfloat16 floor_divide compared with a float32 reference works inconsistently + DecorateInfo(unittest.skip('Skipped!'), 'TestBinaryUfuncs', + dtypes=(torch.bfloat16,)), + # int8 floor divide has different results for -128 // -1 vs. NumPy + DecorateInfo(unittest.skip('Skipped!'), 'TestBinaryUfuncs', 'test_reference_numerics_small_values', + dtypes=(torch.int8,)), + # The following tests fails on some jobs + DecorateInfo(unittest.skip('Skipped!'), 'TestBinaryUfuncs', 'test_reference_numerics_extremal_values', + dtypes=(torch.float16,)), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=5e-3)}), + 'TestBinaryUfuncs', 'test_reference_numerics'), + )), + UnaryUfuncInfo('frexp', + op=torch.frexp, + ref=np.frexp, + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + # skip testing torch.frexp as it is not supported by ROCm platform yet + decorators=[], + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # skips below tests as torch.frexp returns tuple-like (mantissa, exponent) as outputs, + # while these tests currently requires output to a single tensor. + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_batch_vs_slicing'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_contig_vs_every_other'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_contig_vs_transposed'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_non_contig_expand'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_variant_consistency'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_out_arg_all_dtypes'), + + # skips test_reference_numerics due to error in Windows CI. + # The np.frexp returns exponent as np.intc dtype on Windows platform, + # and np.intc does not have the correspond torch dtype + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + active_if=IS_WINDOWS), + # Error: The operator 'aten::frexp.Tensor_out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + )), + UnaryUfuncInfo('log1p', + ref=np.log1p, + aliases=('special.log1p',), + domain=(-1, None), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + decorators=(precisionOverride({torch.bfloat16: 1e-1}),), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + assert_autodiffed=True, + promotes_int_to_float=True), + BinaryUfuncInfo('ge', + ref=np.greater_equal, + aliases=('greater_equal',), + dtypes=all_types_and(torch.bool, torch.bfloat16, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + always_returns_bool=True, + supports_autograd=False, + skips=( + # MPS: The following dtypes worked in forward but are not listed by the OpInfo: {torch.complex64}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + OpInfo('geqrf', + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_qr_geqrf, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + supports_autograd=False, + skips=( + # FIXME: geqrf can't forward with complex inputs that require grad + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes'), + # Strides are not the same! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # NotImplementedError: The operator 'aten::geqrf' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + )), + BinaryUfuncInfo('gt', + ref=np.greater, + aliases=('greater',), + dtypes=all_types_and(torch.bool, torch.bfloat16, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + always_returns_bool=True, + supports_autograd=False, + skips=( + # MPS: The following dtypes worked in forward but are not listed by the OpInfo: {torch.complex64}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + UnaryUfuncInfo('imag', + ref=np.imag, + dtypes=complex_types_and(torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/issues/66357 + # RuntimeError: view_as_real doesn't work on unresolved conjugated tensors. + check_batched_forward_grad=False, + skips=( + # Skip since real and imag don't have out variants. + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_out_arg_all_dtypes'), + )), + OpInfo('gradient', + dtypes=floating_and_complex_types_and(torch.int8, torch.int16, + torch.int32, torch.int64, + torch.bfloat16, torch.half), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # following tests give a runtime error with undefined value tensor + # see discussion : https://github.com/pytorch/pytorch/issues/56660 + # RuntimeError: + # Arguments for call are not valid. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32, torch.complex64)), # noqa: B950 + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCudaFuserOpInfo'), + ), + supports_inplace_autograd=False, + sample_inputs_func=sample_inputs_gradient, + error_inputs_func=error_inputs_gradient), + OpInfo('isin', + dtypes=all_types_and(torch.bfloat16, torch.half), + supports_autograd=False, + sample_inputs_func=sample_inputs_isin), + OpInfo('kthvalue', + dtypes=all_types_and(torch.bfloat16, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + dtypesIfMPS=all_types_and(torch.bfloat16, torch.float16, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_kthvalue, + skips=( + # Exception: "kthvalue_cpu" not implemented for 'Bool' + DecorateInfo(unittest.expectedFailure, 'TestConsistency', device_type='mps', dtypes=(torch.bool,)), + ), + error_inputs_func=error_inputs_kthvalue), + BinaryUfuncInfo('le', + ref=np.less_equal, + aliases=('less_equal',), + dtypes=all_types_and(torch.bool, torch.bfloat16, torch.float16), + always_returns_bool=True, + supports_autograd=False, + skips=( + # MPS: The following dtypes worked in forward but are not listed by the OpInfo: {torch.complex64}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + OpInfo('linspace', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + dtypesIfMPS=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.bool), + is_factory_function=True, + supports_out=True, + supports_autograd=False, + error_inputs_func=error_inputs_linspace, + sample_inputs_func=sample_inputs_linspace, + skips=( + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + + # Same failure as arange: cannot find linspace in captured graph + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # UserWarning: CUDA caching allocator reports a memory leak not verified by the driver API + # in __main__.TestJitCUDA.test_variant_consistency_jit_logspace_cuda_complex64! + # Caching allocator allocated memory was 0 and is now reported as 307200 on device 0. + # CUDA driver allocated memory was 1254555648 and is now 1242955776. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.cfloat,), device_type="cuda"), + # NotImplementedError: "linspace_cpu" not implemented for 'Bool' + DecorateInfo(unittest.expectedFailure, 'TestConsistency', device_type='mps', dtypes=(torch.bool,)), + )), + OpInfo('linspace', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + dtypesIfMPS=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.bool), + is_factory_function=True, + supports_out=True, + supports_autograd=False, + error_inputs_func=error_inputs_linspace, + sample_inputs_func=sample_inputs_linspace_tensor_overload, + variant_test_name="tensor_overload", + skips=( + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # TypeError: 'int' object is not subscriptable + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + + # Same failure as arange: cannot find linspace in captured graph + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # UserWarning: CUDA caching allocator reports a memory leak not verified by the driver API + # in __main__.TestJitCUDA.test_variant_consistency_jit_logspace_cuda_complex64! + # Caching allocator allocated memory was 0 and is now reported as 307200 on device 0. + # CUDA driver allocated memory was 1254555648 and is now 1242955776. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.cfloat,), device_type="cuda"), + # NotImplementedError: "linspace_cpu" not implemented for 'Bool' + DecorateInfo(unittest.expectedFailure, 'TestConsistency', device_type='mps', dtypes=(torch.bool,)), + )), + OpInfo('logspace', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + is_factory_function=True, + supports_out=True, + supports_autograd=False, + error_inputs_func=error_inputs_linspace, + sample_inputs_func=sample_inputs_logspace, + skips=( + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # Same failure as arange: cannot find linspace in captured graph + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + + # Off-by-one issue when casting floats to ints + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_quick', + dtypes=(torch.int16, torch.int32, torch.int64), device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_comprehensive', + dtypes=(torch.int16, torch.int32, torch.int64), device_type="cuda"), + # UserWarning: CUDA caching allocator reports a memory leak not verified by the driver API + # in __main__.TestJitCUDA.test_variant_consistency_jit_logspace_cuda_complex64! + # Caching allocator allocated memory was 0 and is now reported as 307200 on device 0. + # CUDA driver allocated memory was 1254555648 and is now 1242955776. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.cfloat,), device_type="cuda"), + # NotImplementedError: The operator 'aten::logspace.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + )), + OpInfo('logspace', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + is_factory_function=True, + supports_out=True, + supports_autograd=False, + error_inputs_func=error_inputs_linspace, + sample_inputs_func=sample_inputs_logspace_tensor_overload, + variant_test_name="tensor_overload", + skips=( + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # TypeError: 'int' object is not subscriptable + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + # Same failure as arange: cannot find linspace in captured graph + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + + # Off-by-one issue when casting floats to ints + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_quick', + dtypes=(torch.int16, torch.int32, torch.int64), device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_comprehensive', + dtypes=(torch.int16, torch.int32, torch.int64), device_type="cuda"), + # UserWarning: CUDA caching allocator reports a memory leak not verified by the driver API + # in __main__.TestJitCUDA.test_variant_consistency_jit_logspace_cuda_complex64! + # Caching allocator allocated memory was 0 and is now reported as 307200 on device 0. + # CUDA driver allocated memory was 1254555648 and is now 1242955776. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.cfloat,), device_type="cuda"), + # NotImplementedError: The operator 'aten::logspace.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + )), + UnaryUfuncInfo('log', + ref=np.log, + domain=(0, None), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + decorators=(precisionOverride({torch.bfloat16: 5e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + ), + # log(z)->-inf for |z|->0 + reference_numerics_filter=NumericsFilter(condition=lambda x: torch.abs(x) < 0.1, safe_val=1)), + UnaryUfuncInfo('log10', + ref=np.log10, + domain=(0, None), + decorators=(precisionOverride({torch.bfloat16: 5e-2}),), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + ), + # log10(z)->-inf for |z|->0 + reference_numerics_filter=NumericsFilter(condition=lambda x: torch.abs(x) < 0.1, safe_val=1)), + UnaryUfuncInfo('log2', + ref=np.log2, + domain=(0, None), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + decorators=(precisionOverride({torch.bfloat16: 1e-1}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.cfloat, torch.cdouble]), + ), + # log2(z)->-inf for |z|->0 + reference_numerics_filter=NumericsFilter(condition=lambda x: torch.abs(x) < 0.1, safe_val=1)), + BinaryUfuncInfo('ldexp', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_inplace_autograd=False, + promotes_int_to_float=True, + supports_out=True, + supports_rhs_python_scalar=False, + skips=( + # RuntimeError: mul(): functions with out=... arguments don't support + # automatic differentiation, but one of the arguments requires grad + # https://github.com/pytorch/pytorch/issues/68966 + # Eager tests pass but there is an error in the inductor tests. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + ), + decorators=[ + DecorateInfo( + toleranceOverride({ + torch.complex64: tol(atol=1e-05, rtol=1e-05) + }), + 'TestCommon', device_type='cpu', + ), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=4e-5, rtol=6e-6)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ], ), + BinaryUfuncInfo('logaddexp', + dtypes=floating_and_complex_types_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_and_complex_types_and(torch.bfloat16, torch.float16, torch.complex32), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfXPU=floating_types_and(torch.bfloat16, torch.float16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_rhs_python_scalar=False, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_type_promotion', + device_type='xpu'), + )), + OpInfo('logaddexp2', + dtypes=floating_types_and(torch.bfloat16, torch.half), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_logaddexp), + UnaryUfuncInfo('logical_not', + ref=np.logical_not, + decorators=(precisionOverride({torch.bfloat16: 7e-1, + torch.float16: 5e-1}),), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int8, torch.bool), + supports_autograd=False, + skips=( + # The function variant always returns BoolTensor + # while the inplace variant preserves the input dtype. + # >>> t = torch.randn(3) + # >>> torch.logical_not(t) + # tensor([False, False, False]) + # >>> torch.logical_not(t).dtype + # torch.bool + # >>> t.logical_not_().dtype + # torch.float32 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_variant_consistency', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16)), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16)), + # AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + # AssertionError: RuntimeError not raised : Expected RuntimeError when calling with + # input.device=mps:0 and out.device=cpu. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + )), + BinaryUfuncInfo('lt', + ref=np.less, + aliases=('less',), + dtypes=all_types_and(torch.bool, torch.bfloat16, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int8, torch.int32), + always_returns_bool=True, + supports_autograd=False, + skips=( + # MPS: The following dtypes worked in forward but are not listed by the OpInfo: {torch.complex64}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + OpInfo('lu_unpack', + op=torch.lu_unpack, + dtypes=floating_and_complex_types(), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + skipCPUIfNoLapack, + # RuntimeError: linalg.lu_factor(): MPS doesn't support complex types. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), + ), + sample_inputs_func=sample_inputs_lu_unpack), + OpInfo('lu', + op=torch.lu, + dtypes=floating_and_complex_types(), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_lu, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + # we skip jit tests because `lu` is a torch function + # RuntimeError: + # 'Tensor (inferred)' object has no attribute or method 'lu'.: + # File "", line 3 + # def the_method(i0): + # return i0.lu(True, True) + # ~~~~~ <--- HERE + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError not raised: Expected RuntimeError when calling with input.device=cpu and out.device=cuda + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # Exception: linalg.lu_factor(): MPS doesn't support complex types. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), + )), + OpInfo('lu_solve', + op=torch.lu_solve, + dtypes=floating_and_complex_types(), + supports_forward_ad=True, + # See https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_lu_solve, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out', + device_type='mps', dtypes=[torch.float32]), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=[torch.float32]), + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', + device_type='mps', dtypes=[torch.float32]), + # RuntimeError: The size of tensor a (5) must match the size of tensor b (4) at non-singleton dimension 1 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + # Exception: linalg.solve.triangular(); Only float is supported! + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), + DecorateInfo(unittest.skip("Tests different backward paths"), + "TestCommon", "test_floating_inputs_are_differentiable"),), + decorators=[skipCPUIfNoLapack, skipCUDAIfNoMagmaAndNoCusolver]), + OpInfo('masked_fill', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int8, torch.bool, torch.int32), + sample_inputs_func=sample_inputs_masked_fill, + error_inputs_func=error_inputs_masked_fill, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + supports_out=False), + OpInfo('masked_scatter', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int8, torch.bool, torch.int32), + sample_inputs_func=sample_inputs_masked_scatter, + error_inputs_func=error_inputs_masked_scatter, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + supports_out=False, + skips=( + # Compiler issue on ROCm. Regression started in ROCm 6.4. + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + # Exception: RuntimeError not raised : inplace variant either incorrectly allowed resizing + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + # Exception: "masked_scatter_ only supports boolean masks" does not match + # "masked_scatter: expected BoolTensor or ByteTensor for mask" + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='mps'), + )), + OpInfo('masked_select', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_masked_select, + error_inputs_func=error_inputs_masked_select, + skips=( + # Compiler issue on ROCm. Might need to skip until ROCm5.5 + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + )), + OpInfo('matrix_exp', + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + aliases=('linalg.matrix_exp',), + sample_inputs_func=sample_inputs_matrix_exp, + # Needs to construct a 2nx2n matrix by copy_ ing into it + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + skips=( + # mexp does not support bf16 and fp16 + DecorateInfo(unittest.skip('Skipped!'), 'TestInductorOpInfo', 'test_comprehensive', + dtypes=[torch.half], device_type="cpu"), + # The operator 'aten::linalg_matrix_exp' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + supports_out=False, + ), + OpInfo('matmul', + aliases=('linalg.matmul',), + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, + *[torch.bfloat16] + if SM53OrLater or TEST_WITH_ROCM else []), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + assert_autodiffed=True, + assert_jit_shape_analysis=True, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + sample_inputs_func=partial(sample_inputs_matmul, is_rmatmul=False), + decorators=[ + # NVIDIA only assures that bfloat16 is supported by bmm if SM >= 5.3 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_dtypes', device_type='cuda', active_if=not SM53OrLater), + # ROCm intermittently fails the test with standard atol/rtol + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=0)}), + 'TestCommon', 'test_noncontiguous_samples', device_type='cuda', + active_if=TEST_WITH_ROCM), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=0)}), + 'TestCommon', 'test_out', device_type='cuda', + active_if=TEST_WITH_ROCM), + # mv for the sample with shapes (S, S, M, M), (M,) has some variance in the + # backward on CPU + DecorateInfo(toleranceOverride({torch.float32: tol(atol=0, rtol=1e-5)}), + 'TestCommon', 'test_noncontiguous_samples', + device_type='cpu'), + DecorateInfo( + toleranceOverride({ + torch.float32: tol(atol=1e-5, rtol=1e-5), + torch.complex64: tol(atol=1e-5, rtol=1e-5), + }), + "TestDecomp", "test_comprehensive", device_type="cuda", + ), + ], + skips=( + # Strides are not the same! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # https://github.com/pytorch/pytorch/issues/67470 + DecorateInfo(unittest.skip("67470!"), + 'TestCommon', 'test_noncontiguous_samples', + device_type='cpu', dtypes=(torch.long,)), + # AssertionError: False is not true : Tensors failed to compare as equal! + DecorateInfo(unittest.skip("Skipped!"), 'TestOpInfo', + device_type='xla', dtypes=(torch.long,)), + # https://github.com/pytorch/pytorch/issues/71774 + DecorateInfo(unittest.skip('Skipped!'), 'TestNNCOpInfo', 'test_nnc_correctness', + device_type='cpu', dtypes=(torch.long,)), + )), + OpInfo('max', + variant_test_name='reduction_with_dim', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + sample_inputs_func=sample_inputs_max_min_reduction_with_dim, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True), + OpInfo('max', + variant_test_name='reduction_no_dim', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_max_min_reduction_no_dim, + skips=( + # NotImplementedError: The operator 'aten::max.unary_out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + )), + OpInfo('median', + dtypes=all_types_and(torch.bfloat16, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + dtypesIfMPS=all_types_and(torch.bfloat16, torch.float16, torch.bool), + # TODO: some signatures of median do support out + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + error_inputs_func=error_inputs_median, + sample_inputs_func=partial(sample_inputs_reduction, supports_multiple_dims=False)), + OpInfo('nanmedian', + dtypes=all_types_and(torch.bfloat16, torch.float16), + dtypesIfMPS=all_types_and(torch.bfloat16, torch.float16, torch.bool), + # TODO: some signatures of nanmedian do support out + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=partial(sample_inputs_reduction, supports_multiple_dims=False)), + OpInfo('var_mean', + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfMPS=floating_and_complex_types_and( + torch.half, torch.bfloat16, torch.int32, torch.uint8, torch.bool, torch.int8, torch.int16 + ), + sample_inputs_func=sample_inputs_std_var, + # TODO: some signatures of var_mean do support out + supports_out=False, + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo(toleranceOverride({torch.float64: tol(atol=2e-7, rtol=2e-7)}), + "TestDecomp", "test_comprehensive", device_type="cuda"), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=2e-3)}), + "TestInductorOpInfo", "test_comprehensive", device_type="cuda"), + )), + OpInfo('var_mean', + variant_test_name='unbiased', + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfMPS=floating_and_complex_types_and( + torch.half, torch.bfloat16, torch.int32, torch.uint8, torch.bool, torch.int8, torch.int16 + ), + sample_inputs_func=sample_inputs_std_var_unbiased, + # TODO: some signatures of var_mean do support out + supports_out=False, + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo(toleranceOverride({torch.float64: tol(atol=2e-7, rtol=2e-7)}), + "TestDecomp", "test_comprehensive", device_type="cuda"), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=2e-3)}), + "TestInductorOpInfo", "test_comprehensive", device_type="cuda"), + )), + OpInfo('std_mean', + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfMPS=floating_and_complex_types_and( + torch.half, torch.bfloat16, torch.uint8, torch.bool, torch.int8, torch.int16, torch.int32 + ), + sample_inputs_func=sample_inputs_std_var, + # TODO: some signatures of std_mean do support out + supports_out=False, + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo(toleranceOverride({torch.float64: tol(atol=2e-7, rtol=2e-7)}), + "TestDecomp", "test_comprehensive", device_type="cuda"), + )), + OpInfo('std_mean', + variant_test_name='unbiased', + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfMPS=floating_and_complex_types_and( + torch.half, torch.bfloat16, torch.uint8, torch.bool, torch.int8, torch.int16, torch.int32 + ), + sample_inputs_func=sample_inputs_std_var_unbiased, + # TODO: some signatures of var_mean do support out + supports_out=False, + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=4e-5, rtol=9e-3), + torch.float64: tol(atol=2e-7, rtol=2e-7), + }), + "TestDecomp", + "test_comprehensive", + device_type="cuda" + ), + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=4e-5, rtol=9e-3), + torch.float64: tol(atol=2e-7, rtol=2e-7), + }), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda" + ), + )), + OpInfo('meshgrid', + variant_test_name='variadic_tensors', + ref=np.meshgrid, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.bool, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_meshgrid, variant='variadic'), + skips=[ + # JIT does not support variadic tensors. + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":252, + # please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # meshgrid is defined in torch.functional to take a + # variadic list of tensors. Variadic parameters are not + # compatible with the normalize operator tests. + DecorateInfo(unittest.skip("Skipped!"), 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # Skip operator schema test because this is a functional and not an operator + DecorateInfo(unittest.skip("Skipped!"), 'TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), + ], + supports_out=False, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False,), + OpInfo('meshgrid', + variant_test_name='list_of_tensors', + # Unlike the variant above, we do not use np.meshgrid as a + # ref since it does not officially support list of numpy + # arrays. + dtypes=all_types_and_complex_and(torch.bfloat16, torch.bool, torch.float16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_meshgrid, variant='list'), + skips=[ + # meshgrid is defined in torch.functional to take a + # variadic list of tensors. Variadic parameters are not + # compatible with the normalize operator tests. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + ], + assert_autodiffed=True, + supports_out=False, + autodiff_nonfusible_nodes=[], + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False,), + OpInfo('min', + variant_test_name='reduction_with_dim', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + sample_inputs_func=sample_inputs_max_min_reduction_with_dim, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + ), + OpInfo('min', + variant_test_name='reduction_no_dim', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_max_min_reduction_no_dim, + skips=( + # NotImplementedError: The operator 'aten::min.unary_out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + )), + OpInfo('quantile', + dtypes=floating_types(), + sample_inputs_func=sample_inputs_reduction_quantile, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/issues/66357 + # Relies on copy_ to broadcast, but the forward AD path calls broadcast_to which + # does not have a batching rule in core + check_batched_forward_grad=False), + OpInfo('nanquantile', + dtypes=floating_types(), + sample_inputs_func=sample_inputs_reduction_quantile, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/issues/66357 + # Relies on copy_ to broadcast, but the forward AD path calls broadcast_to which + # does not have a batching rule in core + check_batched_forward_grad=False), + BinaryUfuncInfo( + 'max', + aliases=('maximum',), + variant_test_name='binary', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True, + ref=np.maximum, + supports_rhs_python_scalar=False, + skips=( + # Incorrectly attempts to use a scalar for the second argument + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_jit_alias_remapping'), + # TODO: FIXME: RuntimeError: "max_elementwise_cuda" not implemented for 'ComplexFloat' + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion', device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion', device_type='xpu'), + )), + BinaryUfuncInfo( + 'maximum', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ref=np.maximum, + supports_rhs_python_scalar=False, + skips=( + # TODO: FIXME: RuntimeError: "max_elementwise_cuda" not implemented for 'ComplexFloat' + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion', device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion', device_type='xpu'), + )), + BinaryUfuncInfo( + 'min', + aliases=('minimum',), + variant_test_name='binary', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True, + ref=np.minimum, + supports_rhs_python_scalar=False, + skips=( + # Incorrectly attempts to use a scalar for the second argument + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_jit_alias_remapping'), + # TODO: FIXME: RuntimeError: "min_elementwise_cuda" not implemented for 'ComplexFloat' + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='xpu'), + )), + BinaryUfuncInfo( + 'minimum', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ref=np.minimum, + supports_rhs_python_scalar=False, + skips=( + # TODO: FIXME: RuntimeError: "min_elementwise_cuda" not implemented for 'ComplexFloat' + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='xpu'), + ), + ), + BinaryUfuncInfo('logical_and', + ref=np.logical_and, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + supports_autograd=False, + always_returns_bool=True, + skips=( + # AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + # AssertionError: RuntimeError not raised : Expected RuntimeError when calling with + # input.device=mps:0 and out.device=cpu. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + ), + supports_rhs_python_scalar=False), + BinaryUfuncInfo('logical_or', + ref=np.logical_or, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int8, torch.bool), + supports_autograd=False, + always_returns_bool=True, + skips=( + # AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + # AssertionError: RuntimeError not raised : Expected RuntimeError when calling with + # input.device=mps:0 and out.device=cpu. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + ), + supports_rhs_python_scalar=False), + BinaryUfuncInfo('logical_xor', + ref=np.logical_xor, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int8, torch.bool), + supports_autograd=False, + always_returns_bool=True, + supports_rhs_python_scalar=False, + skips=( + # AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + # AssertionError: RuntimeError not raised : Expected RuntimeError when calling with + # input.device=mps:0 and out.device=cpu. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + )), + BinaryUfuncInfo('bitwise_and', + ref=np.bitwise_and, + dtypes=integral_types_and(torch.bool), + dtypesIfHpu=custom_types(torch.bool), + operator_variant=operator.and_, + inplace_operator_variant=operator.iand, + supports_autograd=False, + supports_one_python_scalar=True, + skips=( + # RuntimeError: "bitwise_and_cuda" not implemented for 'Half' + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', + 'test_type_promotion', device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', + 'test_type_promotion', device_type='xpu'), + )), + BinaryUfuncInfo('bitwise_or', + ref=np.bitwise_or, + dtypes=integral_types_and(torch.bool), + dtypesIfHpu=custom_types(torch.bool), + operator_variant=operator.or_, + inplace_operator_variant=operator.ior, + supports_autograd=False, + supports_one_python_scalar=True, + skips=( + # TODO: FIXME: RuntimeError: "bitwise_or_cuda" not implemented for 'Half' + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='xpu'), + )), + BinaryUfuncInfo('bitwise_xor', + ref=np.bitwise_xor, + dtypes=integral_types_and(torch.bool), + dtypesIfHpu=custom_types(torch.bool), + operator_variant=operator.xor, + inplace_operator_variant=operator.ixor, + supports_autograd=False, + supports_one_python_scalar=True, + skips=( + # TODO: FIXME: RuntimeError: "bitwise_xor_cuda" not implemented for 'Half' + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion', + device_type='xpu'), + )), + BinaryUfuncInfo('heaviside', + ref=lambda a, b: ( + # necessary because np.heaviside incorrectly returns float64 when passed args of dtype int64 + np.int64(np.heaviside(a, b)) if a.dtype == np.int64 and b.dtype == np.int64 else np.heaviside(a, b) + ), + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32), + supports_autograd=False, + supports_rhs_python_scalar=False, + skips=( + # RuntimeError: heaviside is not yet implemented for tensors with different dtypes. + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_binary_ufuncs_mixed_dtype'), + # PyTorch's heaviside does not appear to propagate NaNs + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics_extremal_values'), + # NotImplementedError: The operator 'aten::heaviside.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + BinaryUfuncInfo('lcm', + ref=np.lcm, + dtypes=integral_types_and(), + supports_autograd=False, + skips=( + # The operator 'aten::lcm.out' is not currently implemented for the MPS device. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ), + supports_rhs_python_scalar=False), + BinaryUfuncInfo('gcd', + ref=np.gcd, + dtypes=integral_types_and(), + supports_autograd=False, + supports_rhs_python_scalar=False, + skips=( + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.int8,)), + )), + BinaryUfuncInfo('isclose', + ref=np.isclose, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_isclose, + error_inputs_func=error_inputs_isclose, + supports_autograd=False, + supports_out=False, + supports_rhs_python_scalar=False, + skips=( + DecorateInfo(unittest.expectedFailure, + 'TestCommon', + 'test_numpy_refs', dtypes=(torch.complex128,)), + # RuntimeError: Short did not match Int + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_type_promotion'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_binary_ufuncs_mixed_dtype'), + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics_extremal_values'), + )), + # `softmax` supports different dtypes based on whether `dtype` argument, + # is passed or not. Hence two OpInfo entries, one with dtype and other without. + # https://github.com/pytorch/pytorch/issues/68752 + OpInfo('softmax', + aliases=('special.softmax', 'nn.functional.softmax',), + aten_name='softmax', + aten_backward_name='_softmax_backward_data', + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=sample_inputs_softmax_variant, + assert_jit_shape_analysis=True, + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=True), + OpInfo('softmax', + aliases=('special.softmax', 'nn.functional.softmax',), + variant_test_name="with_dtype", + aten_name='softmax', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_softmax_variant, with_dtype=True), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, 'TestConsistency', 'test_output_grad_match', + device_type='mps', dtypes=(torch.float32,) + ), + ), + supports_out=True), + OpInfo( + '_softmax_backward_data', + op=torch.ops.aten._softmax_backward_data, + aten_name='_softmax_backward_data', + dtypes=floating_types_and(torch.bfloat16, torch.float16), + sample_inputs_func=sample_inputs_softmax_backward_data, + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + ), + ), + # `softmin` supports different dtypes based on whether `dtype` argument, + # is passed or not. Hence two OpInfo entries, one with dtype and other without. + # https://github.com/pytorch/pytorch/issues/68752 + OpInfo('nn.functional.softmin', + aten_name='softmin', + dtypes=floating_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_softmax_variant, + assert_jit_shape_analysis=False, + assert_autodiffed=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + OpInfo('nn.functional.softmin', + variant_test_name="with_dtype", + aten_name='softmin', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_softmax_variant, with_dtype=True), + assert_autodiffed=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, 'TestConsistency', 'test_output_grad_match', + device_type='mps', dtypes=(torch.float32,) + ), + ), + supports_out=False), + OpInfo( + "nn.functional.cross_entropy", + dtypes=floating_types_and(torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_cross_entropy, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=3e-3, rtol=1e-3)}), + "TestJit", + "test_variant_consistency_jit", + device_type="cpu", + ), + ), + skips=( + # AssertionError: False is not true : Scalars failed to compare as equal! 0 != 1536 + # test_ops.TestJitCUDA.test_variant_consistency_jit_nn_functional_cross_entropy_cuda_float32 leaked + # 1536 bytes CUDA memory on device 0 + DecorateInfo( + unittest.expectedFailure, + "TestJit", + "test_variant_consistency_jit", + device_type="cuda", + ), + DecorateInfo(unittest.skip("FP16 corss_entropy cases have not been enabled on MPS yet"), + dtypes=(torch.half,), device_type="mps"), + + ) + ), + OpInfo('nn.functional.normalize', + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_normalize, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True), + OpInfo('aminmax', + ref=lambda x, dim=None, keepdim=False: (np.amin(x, axis=dim, keepdims=keepdim), np.amax(x, axis=dim, keepdims=keepdim)), + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8), + decorators=(onlyNativeDeviceTypes,), + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_aminmax, + skips=( + # Exception: MPS supports tensors with dimensions <= 16, but got 65. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='mps'), + ), + error_inputs_func=error_inputs_aminmax_amax_amin), + OpInfo('as_strided', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + sample_inputs_func=sample_inputs_as_strided, + skips=( + # Note: This xfail is fine -- it's inherent to how as_strided works + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples'), + # AssertionError: False is not true : Scalars failed to compare as equal! + DecorateInfo(unittest.skip("Errors when storage_offset is included"), + 'TestCommon', 'test_variant_consistency_eager'), + # Not close + DecorateInfo(unittest.skip("Errors when storage_offset is included"), + 'TestCommon', 'test_complex_half_reference_testing'), + # Not close + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Numerous errors"), 'TestFwdGradients'), + DecorateInfo(unittest.skip("Numerous errors"), 'TestBwdGradients'), + )), + OpInfo('as_strided', + variant_test_name='partial_views', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.int32, torch.int8, torch.bool), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + sample_inputs_func=sample_inputs_as_strided_partial_views, + skips=( + # Note: This xfail is fine -- it's inherent to how as_strided works + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples'), + # These fail because the test changes the input's in-memory layout + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_fn_fwgrad_bwgrad', + dtypes=(torch.complex64, torch.complex128)), + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_forward_mode_AD'), + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_inplace_forward_mode_AD'), + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_inplace_grad'), + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_inplace_gradgrad'), + DecorateInfo(unittest.expectedFailure, 'TestProxyTensorOpInfo', + 'test_make_fx_symbolic_exhaustive_inplace'), + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness'), + # Fail but are also flaky + DecorateInfo(unittest.skip("Test changes in memory layout"), 'TestMathBits'), + DecorateInfo(unittest.skip("Modifies input strides and storage_offset"), 'TestCommon', + 'test_non_standard_bool_values'), + # RuntimeError: setStorage: sizes [2, 2], strides [1, 2], storage offset 10, and itemsize 2 requiring a + # storage size of 28 are out of bounds for storage of size 20 + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_meta_inplace'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_meta_inplace'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_inplace'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_inplace_all_strides'), + )), + OpInfo('as_strided_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + sample_inputs_func=sample_inputs_as_strided, + skips=( + # Note: This xfail is fine -- it's inherent to how as_strided works + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples'), + # AssertionError: False is not true : Scalars failed to compare as equal! + DecorateInfo(unittest.skip("Errors when storage_offset is included"), + 'TestCommon', 'test_variant_consistency_eager'), + # Not close + DecorateInfo(unittest.skip("Errors when storage_offset is included"), + 'TestCommon', 'test_complex_half_reference_testing'), + # Not close + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Numerous errors"), 'TestFwdGradients'), + DecorateInfo(unittest.skip("Numerous errors"), 'TestBwdGradients'), + DecorateInfo(unittest.expectedFailure, 'TestDTensorOps', 'test_dtensor_op_db'), + )), + OpInfo('as_strided_scatter', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + sample_inputs_func=sample_inputs_as_strided_scatter, + error_inputs_func=error_inputs_as_strided_scatter, + skips=( + DecorateInfo(unittest.skip('Works for int64, fails for everything else'), 'TestCommon', 'test_noncontiguous_samples'), # noqa: B950 + DecorateInfo(unittest.skip('Fails in most cases, passes on LAZY for some reason'), 'TestCommon', 'test_variant_consistency_eager'), # noqa: B950 + DecorateInfo(unittest.skip('Fails on cuda'), 'TestCommon', 'test_complex_half_reference_testing', + active_if=not TEST_WITH_ROCM), + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_fn_grad'), + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_forward_mode_AD'), + DecorateInfo(unittest.skip('Passes on complex128 and float64 only'), 'TestFwdGradients', 'test_fn_fwgrad_bwgrad'), + # AssertionError: Tensor-likes are not close! (new_empty_strided.default) + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), 'TestDecomp', 'test_comprehensive'),)), + OpInfo('native_layer_norm', + aten_name='native_layer_norm', + ref=reference_native_layer_norm, + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_out=False, + assert_jit_shape_analysis=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_native_layer_norm, + error_inputs_func=error_inputs_native_layer_norm, + skips=( + # IndexError: tuple index out of range + DecorateInfo(unittest.skip('Skipped!'), 'TestFwdGradients', 'test_forward_mode_AD'), + # Tests fail when weight=None and bias is defined + # https://github.com/pytorch/pytorch/issues/79705 + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_fn_gradgrad'), + # JIT test also tries to compute double backward, which fails + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("Unsupported on MPS for now"), 'TestCommon', 'test_numpy_ref_mps'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-03, rtol=5e-03)}), + "TestDecomp", "test_comprehensive", device_type="cpu"), + # See https://github.com/pytorch/pytorch/issues/173525 + DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-4, rtol=2e-5)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + )), + OpInfo('native_batch_norm', + aten_name='native_batch_norm', + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.uint8, torch.bool, torch.int8, torch.int16, torch.int32 + ), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + allow_cow_input_materialize_forward=[3, 4], + allow_cow_input_materialize_backward=[3, 4], + sample_inputs_func=sample_inputs_native_batch_norm, + skips=( + # NotImplementedError: Could not run + # 'aten::native_batch_norm.out' with arguments from the 'CPU' backend. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type="cpu"), + # RuntimeError: out_invstd.dim() == 1 && out_invstd.is_contiguous() && out_invstd.sizes()[0] + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type="cuda"), + # Problem with _get_numerical_jacobian + # IndexError: tuple index out of range + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD'), + # RuntimeError: deepEquals(input.iValue, deepCopiedInput) INTERNAL ASSERT FAILED + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # https://github.com/pytorch/pytorch/issues/85960 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_compare_cpu'), + # AssertionError: Booleans mismatch: True is not False + DecorateInfo(unittest.skip("Skipped!"), 'TestFakeTensor', 'test_fake_autocast'), + DecorateInfo(unittest.skip("Skipped!"), 'TestFakeTensor', 'test_fake'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-5)}), + "TestCompositeCompliance", "test_forward_ad"), + # AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 5, 5]) != torch.Size([0]). + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)), + # AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 5, 5]) != torch.Size([5, 5, 6]). + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ) + ), + OpInfo('_native_batch_norm_legit', + aten_name='_native_batch_norm_legit', + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + allow_cow_input_materialize_forward=[3, 4], + allow_cow_input_materialize_backward=[3, 4], + sample_inputs_func=sample_inputs__native_batch_norm_legit, + skips=( + # NotImplementedError: Could not run + # 'aten::native_batch_norm.out' with arguments from the 'CPU' backend. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type="cpu"), + # RuntimeError: out_invstd.dim() == 1 && out_invstd.is_contiguous() && out_invstd.sizes()[0] + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type="cuda"), + # Problem with _get_numerical_jacobian + # IndexError: tuple index out of range + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD'), + # RuntimeError: deepEquals(input.iValue, deepCopiedInput) INTERNAL ASSERT FAILED + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # https://github.com/pytorch/pytorch/issues/85960 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_compare_cpu'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-5)}), + "TestCompositeCompliance", "test_forward_ad"), + # The following dtypes worked in forward but are not listed by + # the OpInfo: {torch.uint8, torch.bool, torch.int8, torch.int16, + # torch.int32}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + # FIXME: AssertionError: The values for attribute 'shape' do not match + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ) + ), + OpInfo('_batch_norm_with_update', + op=torch.ops.aten._batch_norm_with_update, + aten_name='_batch_norm_with_update', + dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + allow_cow_input_materialize_forward=[3, 4], + allow_cow_input_materialize_backward=[3, 4], + sample_inputs_func=sample_inputs__batch_norm_with_update, + skips=( + # NotImplementedError: Could not run + # 'aten::native_batch_norm.out' with arguments from the 'CPU' backend. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type="cpu"), + # RuntimeError: out_invstd.dim() == 1 && out_invstd.is_contiguous() && out_invstd.sizes()[0] + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type="cuda"), + # Problem with _get_numerical_jacobian + # IndexError: tuple index out of range + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD'), + # RuntimeError: deepEquals(input.iValue, deepCopiedInput) INTERNAL ASSERT FAILED + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-5)}), + "TestCompositeCompliance", "test_forward_ad"), + # _batch_norm_with_update expects contiguous inputs for cudnn and miopen + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type="cuda"), + DecorateInfo(unittest.expectedFailure, + 'TestMeta', 'test_dispatch_symbolic_meta_outplace_all_strides', device_type="cuda"), + # _batch_norm_with_update does not have python bindings + DecorateInfo(unittest.skip("Skipped!"), 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # aten out variants do not accept out= kwarg, only python out variants + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # The following dtypes worked in forward but are not listed by + # the OpInfo: {torch.uint8, torch.bool, torch.int8, torch.int16, + # torch.int32}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + ) + ), + OpInfo('nn.functional.cosine_similarity', + aten_name="cosine_similarity", + dtypes=floating_types_and(torch.half, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1.3e-5, rtol=2e-2)}), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda" + ), + ], + sample_inputs_func=sample_inputs_cosine_similarity), + OpInfo('nn.functional.adaptive_avg_pool1d', + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + error_inputs_func=error_inputs_adaptive_avg_pool1d, + skips=( + # RuntimeError: Adaptive pool MPS: input sizes must be divisible + # by output sizes. Non-divisible input sizes are not implemented + # on MPS device yet + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + ), + sample_inputs_func=sample_inputs_adaptive_avg_pool1d), + OpInfo('nn.functional.adaptive_avg_pool2d', + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16), + decorators=( + # RuntimeError: + # adaptive_avg_pool2d(Tensor input, int[2] output_size) -> (Tensor): + # Expected a value of type 'List[int]' for argument 'output_size' but + # instead found type 'Tuple[NoneType, int]'. : + # File "", line 3 + # def the_method(i0): + # return torch.nn.functional.adaptive_avg_pool2d(i0, (None, 7)) + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + skips=( + # RuntimeError: Adaptive pool MPS: input sizes must be divisible + # by output sizes. Non-divisible input sizes are not implemented + # on MPS device yet + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + ), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + error_inputs_func=error_inputs_adaptive_avg_pool2d, + sample_inputs_func=sample_inputs_adaptive_avg_pool2d), + OpInfo('nn.functional.adaptive_avg_pool3d', + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16), + decorators=( + # RuntimeError: + # adaptive_avg_pool3d(Tensor input, int[3] output_size) -> (Tensor): + # Expected a value of type 'List[int]' for argument 'output_size' but + # instead found type 'Tuple[NoneType, NoneType, NoneType]'. : + # File "", line 3 + # + # def the_method(i0): + # return torch.nn.functional.adaptive_avg_pool3d(i0, (None, None, None)) + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE + # + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + skips=( + # NotImplementedError: The operator 'aten::_adaptive_avg_pool3d' + # is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + ), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + error_inputs_func=error_inputs_adaptive_avg_pool3d, + sample_inputs_func=sample_inputs_adaptive_avg_pool3d), + OpInfo('nn.functional.adaptive_max_pool1d', + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.half, torch.bfloat16, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool, torch.int8 + ), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # got: Batching rule not implemented for aten::flatten.using_ints + check_batched_forward_grad=False, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + error_inputs_func=error_inputs_adaptive_max_pool1d, + sample_inputs_func=sample_inputs_adaptive_max_pool1d), + OpInfo('nn.functional.adaptive_max_pool2d', + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.half, torch.bfloat16, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool, torch.int8 + ), + decorators=( + # RuntimeError: + # adaptive_max_pool2d(Tensor input, int[2] output_size) -> (Tensor): + # Expected a value of type 'List[int]' for argument 'output_size' but + # instead found type 'Tuple[NoneType, int]'. : + # File "", line 3 + # def the_method(i0): + # return torch.nn.functional.adaptive_max_pool2d(i0, (None, 7)) + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # got: Batching rule not implemented for aten::flatten.using_ints + check_batched_forward_grad=False, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + error_inputs_func=error_inputs_adaptive_max_pool2d, + sample_inputs_func=sample_inputs_adaptive_max_pool2d), + OpInfo('nn.functional.adaptive_max_pool3d', + dtypes=floating_types_and(torch.bfloat16, torch.half), + decorators=( + # RuntimeError: + # adaptive_max_pool3d(Tensor input, int[3] output_size) -> (Tensor): + # Expected a value of type 'List[int]' for argument 'output_size' but + # instead found type 'Tuple[NoneType, NoneType, NoneType]'. : + # File "", line 3 + # + # def the_method(i0): + # return torch.nn.functional.adaptive_max_pool3d(i0, (None, None, None)) + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE + # + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + skips=( + # NotImplementedError: The operator + # 'aten::adaptive_max_pool3d.out' is not currently implemented + # for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + ), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # got: Batching rule not implemented for aten::flatten.using_ints + check_batched_forward_grad=False, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + error_inputs_func=error_inputs_adaptive_max_pool3d, + sample_inputs_func=sample_inputs_adaptive_max_pool3d), + OpInfo('nn.functional.avg_pool1d', + aten_name='avg_pool1d', + supports_autograd=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.int64, torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16), + dtypesIfMPS=floating_types_and( + torch.int64, torch.float16, torch.bfloat16, torch.int8, torch.int16, torch.int32, torch.uint8, torch.bool + ), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + error_inputs_func=error_inputs_avg_pool1d, + sample_inputs_func=sample_inputs_avgpool1d), + OpInfo('nn.functional.avg_pool3d', + aten_name='avg_pool3d', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.int64), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16), + dtypesIfMPS=floating_types_and( + torch.int64, torch.float16, torch.bfloat16, torch.int8, torch.int16, torch.int32, torch.uint8, torch.bool + ), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + error_inputs_func=error_inputs_avg_pool3d, + sample_inputs_func=sample_inputs_avgpool3d, + decorators=( + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-04), }), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ), + skips=( + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='cpu'), + # AssertionError: Scalars are not equal! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)), + )), + OpInfo( + "nn.functional.binary_cross_entropy_with_logits", + aten_name="binary_cross_entropy_with_logits", + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + sample_inputs_func=sample_inputs_binary_cross_entropy_with_logits, + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + 'TestJit', + 'test_variant_consistency_jit', + dtypes=(torch.float32,) + ), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-5, rtol=3e-6)}), + "TestConsistency", "test_output_match", device_type="mps"), + ), + ), + UnaryUfuncInfo( + 'nn.functional.relu', + aten_name="relu", + ref=lambda a: np.where(a <= 0, 0, a), + supports_autograd=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16, torch.float16), + dtypesIfMPS=all_types_and(torch.half, torch.bfloat16, torch.bool), + sample_inputs_func=sample_inputs_nn_activation_relu, + supports_out=False, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True), + OpInfo('nn.functional.conv_transpose1d', + # `ref` for this function is backward of + # corresponding `conv*d` + ref=partial(conv_transpose_ref, fn=torch.nn.functional.conv_transpose1d), + aten_name='conv_transpose1d', + aliases=('conv_transpose1d',), + dtypes=floating_and_complex_types_and(torch.int64, torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.chalf, + torch.bfloat16), + sample_inputs_func=sample_inputs_conv_transpose1d, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + decorators=( + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1.3e-06), }), + 'TestCommon', 'test_variant_consistency_eager', device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=5e-2, rtol=5e-2), }), + 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo( + toleranceOverride({torch.float: tol(atol=1.5e-5, rtol=1.5e-5), }), + 'TestCommon', 'test_numpy_ref_mps'), + DecorateInfo( + toleranceOverride({torch.half: tol(atol=1e-3, rtol=5e-3), }), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cpu'), + ), + skips=( + # Reason for Skip: https://github.com/pytorch/pytorch/pull/79694#issuecomment-1186949486 + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.complex64,)), + # RuntimeError: UNSUPPORTED DTYPE: complex + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness', + dtypes=(torch.complex64, torch.complex128)), + # RuntimeError: !lhs.isAliasOf(rhs)INTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":104, please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.float,)), + # RuntimeError: "slow_conv2d_cpu_grad_input" not implemented for 'Long' + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref', + dtypes=(torch.int64,)), + # RuntimeError: Convolution is supported only for Floating types + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.int64,) + ), + ), + supports_out=False,), + OpInfo('nn.functional.conv_transpose2d', + aten_name='conv_transpose2d', + aliases=('conv_transpose2d',), + # `ref` for this function is backward of + # corresponding `conv*d` + ref=partial(conv_transpose_ref, fn=torch.nn.functional.conv_transpose2d), + dtypes=floating_and_complex_types_and(torch.int64, torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.chalf, + torch.bfloat16), + sample_inputs_func=sample_inputs_conv_transpose2d, + # Runs very slowly on slow-gradcheck for complex. + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1.3e-06), }), + 'TestCommon', 'test_variant_consistency_eager', device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=2e-05, rtol=5e-05), }), + 'TestCommon', 'test_noncontiguous_samples', device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=8e-2, rtol=8e-2), }), + 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo( + toleranceOverride({torch.half: tol(atol=1e-3, rtol=4e-3), }), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cpu')], + skips=( + # RuntimeError: !lhs.isAliasOf(rhs)INTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":104, please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError: UNSUPPORTED DTYPE: complex + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness', + dtypes=(torch.complex64, torch.complex128)), + # RuntimeError: "slow_conv2d_cpu_grad_input" not implemented for 'Long' + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref', + dtypes=(torch.int64,)), + # Reference: https://github.com/pytorch/pytorch/issues/86356 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref', + dtypes=(torch.double, torch.cdouble)), + DecorateInfo(unittest.skip("Unsupported on MPS for now"), 'TestCommon', 'test_numpy_ref_mps'), + # AssertionError: None mismatch: torch.complex64 is not None + DecorateInfo(unittest.expectedFailure, 'TestDtypeCustomRules', 'test_custom_rules', + dtypes=(torch.complex64, torch.complex128)), + # RuntimeError: Convolution is supported only for Floating types + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.int64,) + ), + ), + supports_out=False,), + OpInfo('nn.functional.conv_transpose3d', + aten_name='conv_transpose3d', + aliases=('conv_transpose3d',), + # `ref` for this function is backward of + # corresponding `conv*d` + ref=partial(conv_transpose_ref, fn=torch.nn.functional.conv_transpose3d), + dtypes=floating_and_complex_types_and(torch.int64, torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and( + torch.float16, torch.chalf, torch.bfloat16), + sample_inputs_func=sample_inputs_conv_transpose3d, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + # Runs very slowly on slow-gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=5e-2, rtol=5e-2), }), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-04, rtol=1.3e-06), + torch.complex64: tol(atol=1.3e-04, rtol=1.3e-05)}), + 'TestCommon', 'test_variant_consistency_eager', device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=2e-04, rtol=2e-04), }), + 'TestCompositeCompliance', 'test_operator', device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1.3e-04, rtol=1.3e-06), + torch.complex64: tol(atol=1.3e-04, rtol=1.3e-05)}), + 'TestCommon', 'test_noncontiguous_samples', device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-04, rtol=2e-05), }), + 'TestCompositeCompliance', 'test_forward_ad', device_type='cuda', + active_if=TEST_CUDNN), + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-4, rtol=1e-4)}), + "TestMathBits", "test_conj_view", device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=9e-2, rtol=9e-2), }), + 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}), + 'TestOperators', 'test_vjpvmap', device_type='cuda' + ), + DecorateInfo( + toleranceOverride({torch.half: tol(atol=9e-3, rtol=2e-1), }), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cpu')], + skips=( + # RuntimeError: !lhs.isAliasOf(rhs)INTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":104, please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError: "slow_conv3d_cpu_grad_input" not implemented for 'Long' + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref', + dtypes=(torch.int64,)), + # Reference: https://github.com/pytorch/pytorch/issues/86356 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_numpy_ref', + dtypes=(torch.double, torch.cdouble)), + DecorateInfo(unittest.skip("Unsupported on MPS for now"), 'TestCommon', 'test_numpy_ref_mps'), + # RuntimeError: UNSUPPORTED DTYPE: complex + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness', + dtypes=(torch.complex64, torch.complex128)), + DecorateInfo(unittest.skip('Skipped for ROCm!'), 'TestCommon', 'test_complex_half_reference_testing', + dtypes=[torch.complex32], active_if=TEST_WITH_ROCM), + # RuntimeError: Convolution is supported only for Floating types + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.int64,) + ), + ), + supports_out=False,), + OpInfo('nn.functional.conv1d', + aliases=('conv1d',), + aten_name='conv1d', + dtypes=floating_and_complex_types_and(torch.int64, torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.chalf, + torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=sample_inputs_conv1d, + error_inputs_func=error_inputs_conv1d, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + decorators=( + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=1e-2, rtol=5e-2)}), + 'TestCommon', 'test_complex_half_reference_testing' + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=2e-3, rtol=1e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cuda', + ), + ), + skips=( + # RuntimeError: !lhs.isAliasOf(rhs)INTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":103, please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # Ref: https://github.com/pytorch/pytorch/issues/75309 + # AssertionError: None mismatch: torch.complex128 is not None + DecorateInfo(unittest.expectedFailure, 'TestDtypeCustomRules', + 'test_custom_rules', dtypes=(torch.complex64, torch.complex128)), + # Ref: https://github.com/pytorch/pytorch/issues/75309 + # RuntimeError: UNSUPPORTED DTYPE: complex + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', + 'test_nnc_correctness', dtypes=(torch.complex64, torch.complex128)), + # The following dtypes did not work in forward but are listed by the OpInfo: {torch.int64}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + # RuntimeError: Convolution is supported only for Floating types + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.int64,) + ), + ), + supports_expanded_weight=True, + supports_out=False,), + OpInfo('nn.functional.conv2d', + aliases=('conv2d',), + aten_name='conv2d', + dtypes=floating_and_complex_types_and(torch.int64, torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.chalf, + torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_conv2d), + error_inputs_func=error_inputs_conv2d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + decorators=( + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=6e-2, rtol=5e-2)}), + 'TestCommon', 'test_complex_half_reference_testing', + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=5e-3, rtol=1e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-5, rtol=1e-5)}), + 'TestCommon', 'test_noncontiguous_samples', device_type='mps', + ), + ), + skips=( + # RuntimeError: !lhs.isAliasOf(rhs)INTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":103, please report a bug to PyTorch. + DecorateInfo(unittest.skip("Works on some configs!"), 'TestJit', 'test_variant_consistency_jit'), + # Ref: https://github.com/pytorch/pytorch/issues/75309 + # AssertionError: None mismatch: torch.complex128 is not None + DecorateInfo(unittest.expectedFailure, 'TestDtypeCustomRules', + 'test_custom_rules', dtypes=(torch.complex64, torch.complex128)), + # RuntimeError: UNSUPPORTED DTYPE: complex + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', + 'test_nnc_correctness', dtypes=(torch.complex64, torch.complex128)), + # RuntimeError: Convolution is supported only for Floating types + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', + device_type='mps', dtypes=(torch.int64,) + ), + ), + supports_expanded_weight=True, + supports_out=False,), + OpInfo('nn.functional.conv3d', + aliases=('conv3d',), + aten_name='conv3d', + dtypes=floating_and_complex_types_and(torch.int64, torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.chalf, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=sample_inputs_conv3d, + error_inputs_func=error_inputs_conv3d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=6e-2, rtol=5e-2)}), + 'TestCommon', 'test_complex_half_reference_testing', + ), + # TF32 + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=5e-3, rtol=1e-3), + torch.complex64: tol(atol=5e-3, rtol=1e-3)}), + 'TestCommon', 'test_noncontiguous_samples', + ), + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=2e-5, rtol=3e-6)}), + 'TestCommon', 'test_variant_consistency_eager', + ), + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=5e-5, rtol=5e-6)}), + 'TestMathBits', 'test_conj_view', + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}), + 'TestOperators', 'test_vjpvmap', + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}), + 'TestOperators', 'test_jvpvjp', device_type="cuda" + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}), + 'TestOperators', 'test_vjp', device_type="cuda" + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=5e-5, rtol=5e-6)}), + 'TestCompositeCompliance', 'test_backward', device_type="cuda" + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=5e-3, rtol=1e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', + ), + ), + skips=( + # RuntimeError: !lhs.isAliasOf(rhs) INTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":103, please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError: UNSUPPORTED DTYPE: complex + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', + 'test_nnc_correctness', dtypes=(torch.complex64, torch.complex128)), + # AssertionError: Tensor-likes are not close! + # break slow tests + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_compare_cpu'), + # RuntimeError: Convolution is supported only for Floating types + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + # RuntimeError: Convolution is supported only for Floating types + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.int64,) + ), + ), + supports_expanded_weight=True, + supports_out=False,), + OpInfo('nn.functional.group_norm', + aten_name='group_norm', + aliases=('group_norm',), + ref=reference_group_norm, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.int32, torch.int16, torch.int8, torch.uint8), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + error_inputs_func=error_inputs_group_norm, + decorators=[ + # RuntimeError: Cannot insert a Tensor that requires grad as a constant. + # Consider making it a parameter or input, or detaching the gradient + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=5e-05, rtol=3e-03)}), + "TestDecomp", + "test_comprehensive", + device_type="cpu" + ), + # MPS supports int8/uint8 but CPU does not, so consistency test cannot run + DecorateInfo(unittest.expectedFailure, 'TestConsistency', 'test_output_match', + device_type='mps', dtypes=(torch.int8, torch.uint8)), + ], + sample_inputs_func=sample_inputs_group_norm, + reference_inputs_func=reference_inputs_group_norm, + supports_expanded_weight=True,), + OpInfo('nn.functional.instance_norm', + # no ref because instance_norm will often have numerical instability (large numbers or nan) + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + allow_cow_input_materialize_forward=['running_mean', 'running_var'], + decorators=[ + # RuntimeError: Cannot insert a Tensor that requires grad as a constant. + # Consider making it a parameter or input, or detaching the gradient + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + ], + sample_inputs_func=sample_inputs_instance_norm, + supports_expanded_weight=True,), + OpInfo('nn.functional.layer_norm', + aten_name='layer_norm', + aten_backward_name='layer_norm_backward', + aliases=('layer_norm',), + ref=reference_layer_norm, + dtypes=floating_types_and(torch.half, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-05, rtol=1e-03)}), + 'TestCommon', 'test_numpy_refs' + ), + DecorateInfo(unittest.skip("Bug in MPS backend!"), 'TestCommon', 'test_numpy_ref_mps'), + # See https://github.com/pytorch/pytorch/issues/173525 + DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-4, rtol=2e-5)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ], + sample_inputs_func=sample_inputs_layer_norm, + supports_expanded_weight=True,), + OpInfo('nn.functional.rms_norm', + aten_name='rms_norm', + aliases=('rms_norm',), + ref=reference_rms_norm, + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_rms_norm, + error_inputs_func=error_inputs_rms_norm,), + OpInfo('nn.functional.local_response_norm', + dtypes=floating_types_and(torch.int64, torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.int64, torch.float16, torch.bfloat16, torch.int8, torch.int16, torch.int32, torch.uint8, torch.bool + ), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[ + # RuntimeError: falseINTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + ], + sample_inputs_func=sample_inputs_local_response_norm,), + OpInfo('constant_pad_nd', + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half), + sample_inputs_func=sample_inputs_constant_pad_nd, + supports_out=False, + skips=( + # bool can't be passed to Scalar arguments in JIT tracer because + # BoolType is not a subtype of ScalarType. + DecorateInfo( + unittest.expectedFailure, 'TestNNCOpInfo', + 'test_nnc_correctness', dtypes=(torch.bool,)), + )), + OpInfo('nn.functional.pad', + variant_test_name='constant', + aten_name='constant_pad_nd', + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half), + sample_inputs_func=partial(sample_inputs_nn_pad, mode='constant'), + supports_out=False), + OpInfo('nn.functional.pad', + variant_test_name='reflect', + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half), + dtypesIfMPS=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), + sample_inputs_func=partial(sample_inputs_nn_pad, mode='reflect'), + skips=( + # Doesn't have a corresponding aten operator. + # RuntimeError: falseINTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + ), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + supports_out=False), + OpInfo('nn.functional.pad', + variant_test_name='replicate', + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfMPS=all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool), + sample_inputs_func=partial(sample_inputs_nn_pad, mode='replicate'), + skips=( + # Doesn't have a corresponding aten operator. + # RuntimeError: falseINTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + ), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + supports_out=False), + OpInfo('nn.functional.pad', + variant_test_name='replicate_negative', + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfMPS=all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool), + sample_inputs_func=sample_inputs_nn_pad_replicate_negative, + skips=( + # Doesn't have a corresponding aten operator. + # RuntimeError: falseINTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + # Some negative padding cases cause a segfault on MPS + DecorateInfo(unittest.skip("Not fully supported on MPS"), 'TestConsistency'), + # RuntimeError: start == 0 || start < input_size + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.float32, torch.complex64) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=(torch.float32, torch.complex64) + ), + ), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + supports_out=False), + OpInfo('nn.functional.pad', + variant_test_name='circular', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half), + sample_inputs_func=partial(sample_inputs_nn_pad, mode='circular'), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_grad=False, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + skips=( + # Doesn't have a corresponding aten operator. + # RuntimeError: falseINTERNAL ASSERT FAILED at + # "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + # Difference from is larger with decomposition new_empty_strided.default than original on output 0 + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), 'TestDecomp', 'test_comprehensive'), + ), + supports_out=False), + OpInfo('nn.functional.hardswish', + aten_name="hardswish", + aten_backward_name='hardswish_backward', + supports_autograd=True, + assert_autodiffed=True, + sample_inputs_func=sample_inputs_hardswish, + dtypes=floating_types_and(torch.bfloat16, torch.half), + supports_gradgrad=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + autodiff_nonfusible_nodes=["aten::hardswish"]), + OpInfo('nn.functional.unfold', + aten_name='im2col', + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16, torch.bool), + dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16, torch.bool), + sample_inputs_func=sample_inputs_nn_unfold, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + skips=( + # NOTE: this failure may not reproduce consistently on different systems + # false INTERNAL ASSERT FAILED at "...torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185 + DecorateInfo(unittest.skip("Internal assert failed!"), 'TestJit', 'test_variant_consistency_jit'), + # Compiler issue on ROCm. Regression started in ROCm 6.4. + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + # RuntimeError: Failed to create function state object for: col2im_kernel_float2 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=0.004, rtol=0.002)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + )), + OpInfo('nn.functional.interpolate', + aten_name="interpolate", + variant_test_name='nearest', + supports_autograd=True, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + dtypes=floating_types_and(torch.uint8, torch.half, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_interpolate, 'nearest'), + skips=( + # RuntimeError: false + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + supports_out=False), + OpInfo('nn.functional.interpolate', + aten_name="interpolate", + variant_test_name='nearest-exact', + supports_autograd=True, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + dtypes=floating_types_and(torch.half, torch.bfloat16, torch.uint8), + sample_inputs_func=partial(sample_inputs_interpolate, 'nearest-exact'), + skips=( + # RuntimeError: false + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError: aten::_upsample_nearest_exact*d hit the vmap fallback which is currently disabled + DecorateInfo(unittest.expectedFailure, 'TestOperators', 'test_vmapjvpall_has_batch_rule'), + DecorateInfo(unittest.expectedFailure, 'TestOperators', 'test_vmapvjp_has_batch_rule'), + DecorateInfo(unittest.expectedFailure, 'TestVmapOperatorsOpInfo', 'test_op_has_batch_rule'), + ), + supports_out=False), + OpInfo('nn.functional.interpolate', + aten_name="interpolate", + variant_test_name='linear', + supports_autograd=True, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + dtypes=floating_types_and(torch.half, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_interpolate, 'linear'), + skips=( + # RuntimeError: false + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + supports_out=False), + OpInfo('nn.functional.interpolate', + aten_name="interpolate", + variant_test_name='bilinear', + supports_fwgrad_bwgrad=True, + supports_autograd=True, + supports_forward_ad=True, + dtypes=floating_types_and(torch.uint8, torch.half, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + sample_inputs_func=partial(sample_inputs_interpolate, 'bilinear'), + reference_inputs_func=partial(reference_inputs_interpolate, 'bilinear'), + skips=( + # RuntimeError: false + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + supports_out=False), + OpInfo('nn.functional.interpolate', + aten_name="interpolate", + variant_test_name='bicubic', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.uint8, torch.half, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_interpolate, 'bicubic'), + reference_inputs_func=partial(reference_inputs_interpolate, 'bicubic'), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + # RuntimeError: false + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # torch.uint8 - Failed to create function state object for: upsample_bicubic2d_uchar + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + ), + supports_out=False), + OpInfo('nn.functional.interpolate', + aten_name="interpolate", + variant_test_name='trilinear', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.half, torch.bfloat16), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + sample_inputs_func=partial(sample_inputs_interpolate, 'trilinear'), + skips=( + # RuntimeError: false + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # torch.uint8 - Failed to create function state object for: upsample_bicubic2d_uchar + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + ), + supports_out=False), + OpInfo('nn.functional.interpolate', + aten_name="interpolate", + variant_test_name='area', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + sample_inputs_func=partial(sample_inputs_interpolate, 'area'), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + # RuntimeError: false + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError: Adaptive pool MPS: input sizes must be divisible by output sizes + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + ), + supports_out=False), + OpInfo('nn.functional.upsample_bilinear', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.uint8, torch.half, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + sample_inputs_func=partial(sample_inputs_upsample, 'bilinear'), + reference_inputs_func=partial(reference_inputs_upsample, 'bilinear'), + skips=( + # RuntimeError: false + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + supports_out=False), + OpInfo('_upsample_bilinear2d_aa', + op=torch.ops.aten._upsample_bilinear2d_aa, + aten_name='_upsample_bilinear2d_aa', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.uint8), + dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + sample_inputs_func=partial(sample_inputs_upsample_aa, 'bilinear'), + supports_out=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.expectedFailure, 'TestInductorOpInfo', 'test_comprehensive'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + # Error: The operator 'aten::_upsample_bilinear2d_aa_backward.grad_input' + # is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + )), + OpInfo( + "nn.functional.soft_margin_loss", + dtypes=floating_types_and(torch.half, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + # doesn't support grad on target + sample_inputs_func=partial(sample_inputs_loss, rhs_requires_grad=False), + error_inputs_func=error_inputs_soft_margin_loss, + ), + OpInfo('nn.functional.upsample_nearest', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.uint8, torch.half, torch.bfloat16), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + sample_inputs_func=partial(sample_inputs_upsample, 'nearest'), + skips=( + # RuntimeError: false + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":185, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + supports_out=False), + OpInfo( + "nn.functional.margin_ranking_loss", + dtypes=all_types_and(torch.half, torch.bfloat16), + supports_out=False, + sample_inputs_func=sample_inputs_margin_ranking_loss, + error_inputs_func=error_inputs_margin_ranking_loss, + reference_inputs_func=reference_inputs_margin_ranking_loss, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True), + OpInfo( + "nn.functional.multi_margin_loss", + dtypes=floating_types(), + dtypesIfCUDA=floating_types_and(torch.bfloat16, torch.float16), + supports_out=False, + supports_gradgrad=False, + sample_inputs_func=sample_inputs_multi_margin_loss, + reference_inputs_func=reference_inputs_multi_margin_loss, + error_inputs_func=error_inputs_multi_margin_loss, + skips=( + # NotImplementedError: The operator 'aten::multi_margin_loss' is not + # currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + decorators=( + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4)}), + "TestJit", + "test_variant_consistency_jit", + ), + ), + ), + OpInfo( + "nn.functional.multilabel_margin_loss", + dtypes=floating_types(), + dtypesIfCUDA=floating_types_and(torch.bfloat16, torch.float16), + supports_out=False, + supports_gradgrad=False, + sample_inputs_func=sample_inputs_multilabel_margin_loss, + reference_inputs_func=reference_inputs_multilabel_margin_loss, + error_inputs_func=error_inputs_multilabel_margin_loss, + skips=( + # NotImplementedError: The operator 'aten::multilabel_margin_loss' is not + # currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + ), + OpInfo('nn.functional.leaky_relu', + aliases=None, + aten_name="leaky_relu", + aten_backward_name='leaky_relu_backward', + sample_inputs_func=sample_inputs_leaky_relu, + dtypes=floating_types_and(torch.bfloat16, torch.float16), + inplace_variant=lambda x, negative_slope=0.01: + torch.nn.functional.leaky_relu(x, negative_slope, inplace=True), + supports_autograd=True, + assert_autodiffed=True, + supports_gradgrad=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + autodiff_nonfusible_nodes=["aten::leaky_relu"]), + OpInfo( + "nn.functional.multilabel_soft_margin_loss", + supports_out=False, + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfMPS=floating_types_and(torch.half, torch.bfloat16, torch.int8, torch.uint8, torch.int32, torch.int16), + sample_inputs_func=sample_inputs_multilabel_soft_margin_loss, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4)}), + "TestJit", + "test_variant_consistency_jit", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=4e-3, rtol=1.3e-3)}), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda" + ), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=4e-5, rtol=0.013)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ), + skips=( + # AssertionError: False is not true : Scalars failed to compare as equal! 0 != 4096 + # __main__.TestJitCUDA.test_variant_consistency_jit_nn_functional_multilabel_soft_margin_loss_cuda_float32 + # leaked 4096 bytes CUDA memory on device 0 + DecorateInfo( + # Skip instead of expectedFailure because this fails + # locally for me but passes in CI. + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="cuda", + ), + ), + ), + OpInfo('nn.functional.avg_pool2d', + aten_name='avg_pool2d', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.int64, torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.int64, torch.float16, torch.bfloat16, torch.int32, torch.uint8, torch.bool, torch.int8, torch.int16 + ), + error_inputs_func=error_inputs_avg_pool2d, + sample_inputs_func=sample_inputs_avgpool2d, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='cuda'), + # AssertionError: Scalars are not equal! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + )), + OpInfo('nn.functional.fractional_max_pool2d', + supports_autograd=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.fractional_max_pool2d, input, *args, **kwargs), + # vmap does not support random operations + check_batched_forward_grad=False, + dtypes=floating_types_and(torch.bfloat16, torch.float16), + test_neg_view=False, + sample_inputs_func=sample_inputs_fractional_max_pool2d, + decorators=( + # FIXME: AssertionError: False is not true : Tensors failed to compare as equal! + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":270 + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit')), + skips=( + # Exception: The operator 'aten::fractional_max_pool2d.output' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'),)), + OpInfo('nn.functional.fractional_max_pool3d', + supports_autograd=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.fractional_max_pool3d, input, *args, **kwargs), + # vmap does not support random operations + check_batched_forward_grad=False, + dtypes=floating_types_and(torch.bfloat16, torch.float16), + test_neg_view=False, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + sample_inputs_func=sample_inputs_fractional_max_pool3d, + decorators=( + # FIXME: both derivatives are implemented incorrectly + # https://github.com/pytorch/pytorch/issues/69322 + # FIXME: AssertionError: False is not true : Tensors failed to compare as equal! + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":270 + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit')), + skips=( + # Exception: The operator 'aten::fractional_max_pool3d.output' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'),)), + OpInfo('nn.functional.max_pool1d', + aten_name='max_pool1d', + supports_autograd=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # got: Batching rule not implemented for aten::flatten.using_ints + check_batched_forward_grad=False, + # TODO: add shape checks + assert_jit_shape_analysis=False, + dtypes=floating_types_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.bfloat16, torch.float16, torch.uint8, torch.bool, torch.int8, torch.int16, torch.int32, torch.int64 + ), + skips=( + # Pre-existing condition; Needs to be fixed + DecorateInfo(unittest.skip("Works on some configs"), 'TestNNCOpInfo', + 'test_nnc_correctness', dtypes=(torch.bfloat16,)), + # RuntimeError: The tensor has a non-zero number of elements, but its data is not allocated yet. + # Caffe2 uses a lazy allocation, so you will need to call mutable_data() or raw_mutable_data() + # to actually allocate memory + DecorateInfo(unittest.skip("Skipped!"), 'TestTags', 'test_tags'), + ), + error_inputs_func=error_inputs_max_pool1d, + sample_inputs_func=sample_inputs_max_pool), + OpInfo('nn.functional.max_pool2d', + aten_name='max_pool2d', + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + # Vmap is not happy with non-contiguous (channels_last) inputs + check_batched_gradgrad=False, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # got: Batching rule not implemented for aten::flatten.using_ints + check_batched_forward_grad=False, + assert_jit_shape_analysis=True, + dtypes=all_types_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=all_types_and(torch.float16, torch.bfloat16, torch.bool), + error_inputs_func=error_inputs_max_pool2d, + sample_inputs_func=sample_inputs_max_pool), + OpInfo('max_pool2d_with_indices_backward', + op=max_pool2d_backward, + # We've defined a custom op, so there's no corresponding aten op + aten_name=None, + method_variant=None, + inplace_variant=None, + operator_variant=None, + inplace_operator_variant=None, + check_batched_gradgrad=False, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + assert_jit_shape_analysis=False, + dtypes=floating_types_and(torch.bfloat16, torch.float16), + dtypesIfMPS=floating_types_and( + torch.bfloat16, torch.float16, torch.int8, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool + ), + sample_inputs_func=sample_inputs_max_pool, + skips=( + # We've defined a custom op here, and we don't handle the case where we receive an out kwarg + DecorateInfo(unittest.skip("Skipped!"), "TestCommon", "test_out"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # object has no attribute max_pool2d_with_indices_backward (It's not available on torch -- so expected) + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit') + )), + OpInfo('nn.functional.max_pool3d', + aten_name='max_pool3d', + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # got: Batching rule not implemented for aten::flatten.using_ints + check_batched_forward_grad=False, + # TODO: add shape checks + assert_jit_shape_analysis=False, + dtypes=all_types_and(torch.bfloat16, torch.float16), + dtypesIfMPS=all_types_and(torch.bfloat16, torch.float16, torch.bool), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + # TODO: investigate nondeterminism + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + error_inputs_func=error_inputs_max_pool3d, + sample_inputs_func=sample_inputs_max_pool), + OpInfo('nn.functional.max_unpool1d', + aten_name='max_unpool1d', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + assert_jit_shape_analysis=False, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool, torch.int8 + ), + sample_inputs_func=sample_inputs_max_unpool, + error_inputs_func=error_inputs_max_unpool, + skips=( + # Gradients are tested in `variant_test_name=grad` below. + # We skip tests here because there is non-determinism in backward + # with gather, when there are writes into the same memory location, + # and if there are several indices pointing to the same memory, + # gradcheck is oblivious about that and cannot perturb them all at once + # (see sample_inputs_max_unpool_grad to find out more). + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_grad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_gradgrad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD', + active_if=(not IS_MACOS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_forward_ad', + device_type='cpu'), + DecorateInfo(unittest.skip("Skipped!"), 'TestDecomp', 'test_quick_core_backward'), + )), + OpInfo('nn.functional.max_unpool1d', + variant_test_name='grad', + aten_name='max_unpool1d', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + assert_jit_shape_analysis=False, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool, torch.int8 + ), + sample_inputs_func=sample_inputs_max_unpool_grad), + OpInfo('nn.functional.max_unpool2d', + aten_name='max_unpool2d', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + assert_jit_shape_analysis=False, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool, torch.int8 + ), + sample_inputs_func=sample_inputs_max_unpool, + error_inputs_func=error_inputs_max_unpool, + skips=( + # Gradients are tested in `variant_test_name=grad` below. + # We skip tests here because there is non-determinism in backward + # with gather, when there are writes into the same memory location, + # and if there are several indices pointing to the same memory, + # gradcheck is oblivious about that and cannot perturb them all at once + # (see sample_inputs_max_unpool_grad to find out more). + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD', + active_if=(not IS_MACOS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_gradgrad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_grad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_forward_ad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestDecomp', 'test_quick_core_backward'), + )), + OpInfo('nn.functional.max_unpool2d', + variant_test_name='grad', + aten_name='max_unpool2d', + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # Vmap is not happy with non-contiguous (channels_last) inputs + check_batched_grad=False, + supports_out=False, + assert_jit_shape_analysis=False, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool, torch.int8 + ), + sample_inputs_func=sample_inputs_max_unpool_grad), + OpInfo('nn.functional.max_unpool3d', + aten_name='max_unpool3d', + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + assert_jit_shape_analysis=False, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool, torch.int8 + ), + sample_inputs_func=sample_inputs_max_unpool, + error_inputs_func=error_inputs_max_unpool, + skips=( + # Gradients are tested in `variant_test_name=grad` below. + # We skip tests here because there is non-determinism in backward + # with gather, when there are writes into the same memory location, + # and if there are several indices pointing to the same memory, + # gradcheck is oblivious about that and cannot perturb them all at once + # (see sample_inputs_max_unpool_grad to find out more). + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD', + active_if=(not IS_MACOS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_gradgrad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_grad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_forward_ad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestDecomp', 'test_quick_core_backward'), + )), + OpInfo('nn.functional.max_unpool3d', + variant_test_name='grad', + aten_name='max_unpool3d', + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + assert_jit_shape_analysis=False, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool, torch.int8 + ), + sample_inputs_func=sample_inputs_max_unpool_grad), + OpInfo('nn.functional.linear', + aten_name='linear', + supports_autograd=True, + supports_gradgrad=True, + sample_inputs_func=sample_inputs_linear, + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfROCM=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + # linear calls mm under the hood which is nondeterministic on CUDA + # https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html#torch.use_deterministic_algorithms + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + supports_expanded_weight=True, + skips=( + # RuntimeError: MPS device does not support linear for non-float inputs + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.int64,)), + ), + decorators=( + # Strides are not the same! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + )), + OpInfo('nn.functional.bilinear', + aten_name='bilinear', + supports_autograd=True, + sample_inputs_func=sample_inputs_bilinear, + dtypes=all_types_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.float16, + *[torch.bfloat16] if SM53OrLater or TEST_WITH_ROCM else []), + decorators=( + DecorateInfo(toleranceOverride({torch.float16: tol(atol=2e-03, rtol=1.3e-03)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cpu'), + ), + skips=( + # NVIDIA only assures that bfloat16 is supported by bmm if SM >= 5.3 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_dtypes', device_type='cuda', active_if=not SM53OrLater), + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness', dtypes=(torch.bfloat16,)), + ), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + OpInfo('nn.functional.glu', + aten_name='glu', + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + sample_inputs_func=sample_inputs_glu, + dtypes=floating_types_and(torch.bfloat16, torch.float16), + dtypesIfMPS=floating_types_and( + torch.bfloat16, torch.float16, torch.int32, torch.uint8, torch.bool, torch.int8, torch.int16 + ), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + UnaryUfuncInfo( + 'nn.functional.elu', + aten_backward_name='elu_backward', + ref=lambda x, alpha=1.0, inplace=False: + np.maximum(0., x) + np.minimum(0., alpha * (np.exp(x) - 1)), + dtypes=floating_types_and(torch.bfloat16, torch.float16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_autograd=True, + assert_autodiffed=False, + supports_gradgrad=True, + supports_out=False, + sample_kwargs=lambda device, dtype, input: + ({'alpha': 0.8}, {'alpha': 0.8}), + inplace_variant=lambda x, alpha=1.0: + torch.nn.functional.elu(x, alpha, inplace=True), + decorators=[ + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=1e-03, rtol=1.2e-03), + torch.bfloat16: tol(atol=1e-03, rtol=1.2e-03) + }), + 'TestUnaryUfuncs', device_type='cuda', + ), ], + ), + # Marked as a Unary function because it has some rather odd broadcasting semantics in its + # second argument + UnaryUfuncInfo( + 'nn.functional.prelu', + aten_backward_name='_prelu_kernel_backward', + ref=lambda x, weight: + np.maximum(0., x) + np.minimum(0., x) * + (weight if x.ndim == 1 else weight.reshape([weight.size if i == 1 else 1 for i in range(x.ndim)])), + dtypes=floating_types_and(torch.bfloat16, torch.float16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_autograd=True, + assert_autodiffed=False, + supports_gradgrad=True, + supports_out=False, + # test_reference_numerics only tests the case when the weight tensor is a scalar + sample_kwargs=sample_kwargs_prelu_scalar_weight, + error_inputs_func=error_inputs_prelu, + sample_inputs_func=sample_inputs_prelu, + reference_inputs_func=reference_inputs_prelu, + decorators=[ + # RuntimeError: Cannot insert a Tensor that requires grad as a constant. + # Consider making it a parameter or input, or detaching the gradient + # https://github.com/pytorch/pytorch/issues/68752 + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), ], + ), + UnaryUfuncInfo( + 'nn.functional.celu', + ref=lambda x, alpha=1.0, inplace=False: + np.maximum(0., x) + np.minimum(0., alpha * (np.exp(x / alpha) - 1)), + dtypes=floating_types_and(torch.bfloat16, torch.float16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_autograd=True, + assert_autodiffed=False, + supports_gradgrad=True, + supports_out=False, + sample_kwargs=lambda device, dtype, input: + ({'alpha': 0.8}, {'alpha': 0.8}), + inplace_variant=lambda x, alpha=1.0: + torch.nn.functional.celu(x, alpha, inplace=True), + decorators=[ + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=1e-03, rtol=1.2e-03), + torch.bfloat16: tol(atol=1e-03, rtol=1.2e-03) + }), + 'TestUnaryUfuncs', device_type='cuda', + ), ], + ), + UnaryUfuncInfo( + 'nn.functional.rrelu', + aten_backward_name='rrelu_with_noise_backward', + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.rrelu, input, *args, **kwargs), + inplace_variant=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.rrelu, input, *args, inplace=True, **kwargs), + dtypes=floating_types_and(torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + gradcheck_wrapper=wrapper_set_seed, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + sample_kwargs=lambda device, dtype, input: + (dict(lower=0., upper=1., training=True), dict(lower=0., upper=1., training=True)), + sample_inputs_func=sample_inputs_rrelu, + error_inputs_func=error_inputs_rrelu, + decorators=( + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=1e-03, rtol=1.2e-03), + torch.bfloat16: tol(atol=1e-03, rtol=1.2e-03) + }), + 'TestUnaryUfuncs', device_type='cuda', + ),), + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # In-place operations do not play well with forward AD + # https://github.com/pytorch/pytorch/issues/77447 + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', + 'test_inplace_forward_mode_AD'), + # NotImplementedError: The operator 'aten::rrelu_with_noise' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + # The noise vector that's generated in these tests is not the same elementwise + DecorateInfo(unittest.skip("Different noise"), 'TestUnaryUfuncs', 'test_batch_vs_slicing'), + DecorateInfo(unittest.skip("Different noise"), 'TestUnaryUfuncs', 'test_contig_vs_every_other'), + DecorateInfo(unittest.skip("Different noise"), 'TestUnaryUfuncs', 'test_non_contig_expand'), + DecorateInfo(unittest.skip("Different noise"), 'TestUnaryUfuncs', 'test_contig_vs_transposed'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu')), + skip_correctness_check_compile_vs_eager=True, + ), + UnaryUfuncInfo( + 'nn.functional.selu', + ref=lambda x, inplace=False: + 1.0507009873554804934193349852946 * ( + np.maximum(0., x) + np.minimum(0., 1.6732632423543772848170429916717 * (np.exp(x) - 1)) + ), + dtypes=floating_types_and(torch.bfloat16, torch.float16), + supports_forward_ad=True, # depends on 'elu' + supports_fwgrad_bwgrad=True, + supports_autograd=True, + assert_autodiffed=False, + supports_gradgrad=True, + supports_out=False, + inplace_variant=lambda x: torch.nn.functional.selu(x, inplace=True), + decorators=[ + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=1e-2, rtol=1.8e-2), + torch.bfloat16: tol(atol=1e-2, rtol=1.8e-2) + }), + 'TestUnaryUfuncs', device_type='cuda', + ), ], + ), + OpInfo( + 'torch._scaled_mm_v2', + sample_inputs_func=sample_inputs_scaled_mm_v2, + dtypes=float8_types(), + dtypesIfCUDA=empty_types() + (torch.float8_e4m3fn,), + supports_out=True, + supports_forward_ad=False, + supports_autograd=False, + decorators=[onlyCUDA, skipCUDAIf(not SM89OrLater or TEST_WITH_ROCM, 'Requires CUDA SM >= 8.9')], + skips=( + # Sample inputs isn't really parametrized on dtype + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_dtypes'), + # _scaled_mm_v2 is CUDA-only, no CPU implementation + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_compare_cpu'), + # "add_stub" not implemented for 'Float8_e4m3fn' + # "ufunc_add_CUDA" not implemented for 'Float8_e4m3fn' + # https://github.com/pytorch/pytorch/issues/107256 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out'), + # "mul_cuda" not implemented for float8_e4m3fn + # "mul_cpu_reduced_float" not implemented for 'Float8_e4m3fn' + # https://github.com/pytorch/pytorch/issues/107256 + DecorateInfo(unittest.skip("Skipped!"), 'TestSchemaCheckModeOpInfo', 'test_schema_correctness'), + # aten::_scaled_mm hit the vmap fallback which is currently disabled + DecorateInfo(unittest.skip("Skipped!"), "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + DecorateInfo(unittest.skip("Skipped!"), "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness', + dtypes=(torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz)), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_compare_cpu'), + ) + ), + OpInfo( + 'torch._scaled_mm', + sample_inputs_func=sample_inputs_scaled_mm, + dtypes=float8_types(), + dtypesIfCUDA=empty_types() + (torch.float8_e4m3fn,), + supports_out=True, + supports_forward_ad=False, + supports_autograd=False, + decorators=[skipXPU, skipCUDAIf(not SM89OrLater or TEST_WITH_ROCM, 'Requires CUDA SM >= 8.9')], + skips=( + # Sample inputs isn't really parametrized on dtype + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_dtypes'), + # "add_stub" not implemented for 'Float8_e4m3fn' + # "ufunc_add_CUDA" not implemented for 'Float8_e4m3fn' + # https://github.com/pytorch/pytorch/issues/107256 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out'), + # "mul_cuda" not implemented for float8_e4m3fn + # "mul_cpu_reduced_float" not implemented for 'Float8_e4m3fn' + # https://github.com/pytorch/pytorch/issues/107256 + DecorateInfo(unittest.skip("Skipped!"), 'TestSchemaCheckModeOpInfo', 'test_schema_correctness'), + # aten::_scaled_mm hit the vmap fallback which is currently disabled + DecorateInfo(unittest.skip("Skipped!"), "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + DecorateInfo(unittest.skip("Skipped!"), "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness', + dtypes=(torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz)), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_compare_cpu'), + # TypeError: Trying to convert Float8_e4m3fn to the MPS backend but it does not have support for that dtype. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ) + ), + OpInfo( + 'torch.ops.aten._safe_softmax.default', + dtypes=all_types_and(torch.half, torch.bfloat16, torch.bool), + sample_inputs_func=sample_inputs_safe_softmax, + assert_jit_shape_analysis=True, + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + supports_cow_input_no_materialize_backward=False, + decorators=[], + skips=( + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + ), + OpInfo( + 'nn.functional.scaled_dot_product_attention', + op=lambda *args, **kwargs: + wrapper_set_seed(torch.nn.functional.scaled_dot_product_attention, *args, **kwargs), + sample_inputs_func=sample_inputs_scaled_dot_product_attention, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=False, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + decorators=[DecorateInfo(toleranceOverride( + {torch.float32: tol(atol=5e-05, rtol=5e-6)}), 'TestCommon',), ], + skips=( + # When attn mask is a composite tensor this fails backward by returning a none + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_backward', device_type='cuda'), + # This is only failing on Linux Bionic 3.10 Cuda 11.6 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_dtypes', + device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples', + dtypes=(torch.float32,)), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # Forward works for dtype=float64 which is the math path + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD'), + # Not implemented for Forward AD + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_fn_fwgrad_bwgrad', + device_type='cpu'), + # Not implemented for backward derivative + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients', 'test_fn_gradgrad', + device_type='cpu'), + # CPU and CUDA have inconsistencies for intermediate outputs + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_dispatch_meta_outplace', + device_type='cpu'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_dispatch_symbolic_meta_outplace', + device_type='cpu'), + # When changing input from Tensor to CompositeCompliantTensor, input.requires_grad() changes from true to false + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_backward', + device_type='cpu'), + # OpInfo was implemented with a lambda + DecorateInfo(unittest.skip("Skipped!"), 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # TODO Need to understand what this is testing and why it doesn't work + DecorateInfo(unittest.skip("Skipped"), 'TestDecomp', 'test_comprehensive'), + DecorateInfo(unittest.skip('output is non-deterministic (when dropout_p > 0)'), 'TestCommon', 'test_compare_cpu'), + # TODO skip this for now since we can't skip on runtime arch support + DecorateInfo(unittest.skip('This is '), 'TestInductorOpInfo', 'test_comprehensive'), + # skip for sm < 80 + DecorateInfo(unittest.skip("Skipped!"), 'TestSchemaCheckModeOpInfo', 'test_schema_correctness', + device_type='cuda', dtypes=(torch.bfloat16,), active_if=not SM80OrLater), + # FIXME + DecorateInfo(unittest.skip('test_cow_input does not work with efficient attention on ROCM'), + 'TestCompositeCompliance', 'test_cow_input', + device_type='cuda', dtypes=(torch.bfloat16, torch.float16, torch.float32), + active_if=TEST_WITH_ROCM and PLATFORM_SUPPORTS_MEM_EFF_ATTENTION),), + ), + OpInfo( + 'torch.ops.aten._flash_attention_forward', + sample_inputs_func=sample_inputs_flash_attention_forward, + dtypes=empty_types(), + dtypesIfCUDA=custom_types(torch.float16) + if not SM80OrLater + else custom_types(torch.float16, torch.bfloat16), + supports_out=False, + supports_autograd=True, + supports_fwgrad_bwgrad=False, + supports_forward_ad=False, + check_batched_forward_grad=False, + decorators=[skipCUDAIf(not PLATFORM_SUPPORTS_FLASH_ATTENTION, "This platform doesn't support Flash Attention")], + skips=( + # Checking the scalar value of the philox seed and offset + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_operator', device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', device_type='cuda'), + # None Mismatch Tensor + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_backward', device_type='cuda'), + ) + ), + OpInfo( + 'torch.ops.aten._efficient_attention_forward', + sample_inputs_func=sample_inputs_efficient_attention_forward, + dtypes=empty_types(), + dtypesIfCUDA=custom_types(torch.float16, torch.float32) + if not SM80OrLater + else custom_types(torch.float16, torch.float32, torch.bfloat16), + supports_out=False, + supports_autograd=True, + supports_fwgrad_bwgrad=False, + supports_forward_ad=False, + check_batched_forward_grad=False, + # TODO: Skip because it produces a CUDA illegal memory access for some reason + skip_cow_input_backward=True, + # FIXME: mask_type == 2 (LowerRight) + decorators=[ + skipCUDAIf(not PLATFORM_SUPPORTS_MEM_EFF_ATTENTION, "This platform doesn't support efficient attention"), + skipCUDAIf(TEST_WITH_ROCM, "Efficient attention on ROCM doesn't support custom_mask_type==2"), + skipXPU], + skips=( + # Checking the scaler value of the philox seed and offset + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_operator', device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', device_type='cuda'), + # None Mismatch Tensor + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_backward', device_type='cuda'), + ) + ), + UnaryUfuncInfo( + 'nn.functional.silu', + aten_backward_name='silu_backward', + ref=lambda x, inplace=False: x / (1 + np.exp(-x)), + dtypes=floating_types_and(torch.bfloat16, torch.float16), + dtypesIfMPS=floating_types_and( + torch.bfloat16, torch.float16, torch.int16, torch.int32, torch.uint8, torch.bool, torch.int8 + ), + supports_forward_ad=True, + supports_autograd=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True, + supports_out=False, + inplace_variant=lambda x: torch.nn.functional.silu(x, inplace=True), + decorators=[ + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=1e-3, rtol=1e-3), + torch.bfloat16: tol(atol=1e-4, rtol=1e-4) + }), + 'TestUnaryUfuncs', device_type='cuda', + ), ], + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_normal', + dtypes=(torch.cfloat,), device_type='cpu'), + ), + autodiff_nonfusible_nodes=["aten::silu"], + ), + # TODO: combine this with the nn.functional.silu OpInfo when + # complex autodiff for silu is supported or when + # the forward bug is fixed + # Note: silu errors when given inputs that require grad + # but it doesn't support grad in their dtype + # This is why the dtypes list above passes test_dtypes, + # because it's getting lucky and failing in forward + # because test_dtypes sets requires_grad to True + # THIS IS A BUG + UnaryUfuncInfo( + 'nn.functional.silu', + variant_test_name='complex', + ref=lambda x, inplace=False: + x / (1 + np.exp(-x)), + dtypes=complex_types(), + dtypesIfCUDA=complex_types(), + supports_forward_ad=False, + supports_autograd=False, + assert_autodiffed=False, + supports_out=False, + inplace_variant=lambda x: torch.nn.functional.silu(x, inplace=True), + decorators=[ + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=1e-3, rtol=1e-3), + torch.bfloat16: tol(atol=1e-4, rtol=1e-4) + }), + 'TestUnaryUfuncs', device_type='cuda', + ), ], + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_normal', + dtypes=(torch.cfloat,)), + # FIXME: intentionally misreports dtypes + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes'), + # FIXME: numpy reference diverges: Comparing (nan+nanj) and (-0+0j) + DecorateInfo(unittest.skip("Skipped!"), + 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=(torch.complex64, torch.cdouble)), + DecorateInfo(unittest.skip("Skipped!"), + 'TestUnaryUfuncs', 'test_reference_numerics_small', + dtypes=(torch.complex64,)), + DecorateInfo(unittest.skip("Skipped!"), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=(torch.complex64,)))), + UnaryUfuncInfo( + 'nn.functional.hardsigmoid', + aten_backward_name='hardsigmoid_backward', + ref=reference_hardsigmoid, + dtypes=floating_types_and(torch.bfloat16, torch.float16), + supports_autograd=True, + assert_autodiffed=False, + supports_gradgrad=False, + supports_forward_ad=True, + supports_out=False, + inplace_variant=partial(torch.nn.functional.hardsigmoid, inplace=True), + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-04, rtol=0.001)}), 'TestUnaryUfuncs', device_type='cuda',), ], + skips=[ + # still want to test that first derivative works though second derivative isn't supported + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', "test_inplace_gradgrad")] + ), + UnaryUfuncInfo( + 'nn.functional.logsigmoid', + aten_name="log_sigmoid", + aten_backward_name='log_sigmoid_backward', + ref=reference_logsigmoid, + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.half, torch.bfloat16, torch.int16, torch.int32, torch.uint8, torch.bool, torch.int8 + ), + supports_autograd=True, + assert_autodiffed=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_gradgrad=True, + # autodiff_nonfusible_nodes=["aten::log_sigmoid"], + decorators=[ + DecorateInfo( + precisionOverride({torch.float16: 1e-2, torch.bfloat16: 5e-3}), + 'TestUnaryUfuncs', 'test_reference_numerics_small'), + DecorateInfo( + precisionOverride({torch.float16: 1e-2, torch.bfloat16: 5e-3}), + 'TestUnaryUfuncs', 'test_reference_numerics_large'), + DecorateInfo( + precisionOverride({torch.float16: 1e-2, torch.bfloat16: 5e-3}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal'), + ], + skips=( + # Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='cpu'), + # RuntimeError not raised : Expected RuntimeError when doing an + # unsafe cast from a result of dtype torch.float32 into an out= with + # dtype torch.long + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)), + # UserWarning not triggered : Resized a non-empty tensor but did not + # warn about it + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ), + ), + UnaryUfuncInfo( + 'nn.functional.mish', + aten_backward_name='mish_backward', + ref=lambda x: x * np.tanh(reference_softplus(x)), + dtypes=floating_types_and(torch.bfloat16, torch.float16), + dtypesIfMPS=floating_types_and( + torch.bfloat16, torch.float16, torch.int32, torch.uint8, torch.bool, torch.int8, torch.int16 + ), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_autograd=True, + assert_autodiffed=False, + supports_gradgrad=True, + supports_out=False, + inplace_variant=partial(torch.nn.functional.mish, inplace=True), + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-03)}), 'TestUnaryUfuncs',), ], + ), + UnaryUfuncInfo( + 'nn.functional.softsign', + ref=lambda x: x / (np.abs(x) + 1), + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.float16, torch.bfloat16, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_autograd=True, + assert_autodiffed=False, + supports_gradgrad=True, + supports_out=False, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-03, rtol=1.3e-04)}), 'TestUnaryUfuncs',), ], + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', + dtypes=(torch.int, torch.int8)),), + ), + UnaryUfuncInfo( + 'nn.functional.tanhshrink', + ref=lambda x: x - np.tanh(x), + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_autograd=True, + assert_autodiffed=False, + supports_gradgrad=True, + supports_out=False, + decorators=[ + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_normal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo( + toleranceOverride({torch.bfloat16: tol(atol=1e-02, rtol=1.6e-02)}), 'TestUnaryUfuncs',), + DecorateInfo(toleranceOverride({torch.complex64: tol(atol=6e-04, rtol=1e-05), + torch.bfloat16: tol(atol=1e-02, rtol=1.6e-02)}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cuda'), + ], + skips=( + # in each case, pytorch will produce a nan while numpy will not + DecorateInfo(unittest.skip("Fails on some jobs works on others!"), + 'TestUnaryUfuncs', "test_reference_numerics_large", + dtypes=(torch.complex64, torch.complex128), active_if=(IS_MACOS)), + DecorateInfo(unittest.skip("Fails on some jobs works on others!"), + 'TestUnaryUfuncs', "test_reference_numerics_extremal", + dtypes=(torch.complex64, torch.complex128), device_type='cpu', + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + ), + # tan(j * pi/2 * odd_number) is nan which also make tanhshrink nan. + reference_numerics_filter=NumericsFilter( + condition=lambda x: (close_to_int(x / (math.pi * 0.5j)) + if x.is_complex() else x.new_tensor(False, dtype=torch.bool)), + safe_val=0) + ), + UnaryUfuncInfo( + 'nn.functional.threshold', + ref=lambda x, threshold, value: np.where(x <= threshold, value, x).astype(x.dtype), + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfMPS=all_types_and(torch.half, torch.bfloat16, torch.bool), + inplace_variant=lambda x, threshold, value: + torch.nn.functional.threshold(x, threshold, value, inplace=True), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=False, + supports_gradgrad=True, + supports_out=False, + sample_kwargs=lambda device, dtype, input: ({'threshold': float.fromhex('0x1.3ap-3'), + 'value': -9}, + {'threshold': float.fromhex('0x1.3ap-3'), + 'value': -9}), + # TODO(whc) should not need sample_inputs_func, but without it + # kwargs aren't being hooked up properly + sample_inputs_func=sample_inputs_threshold, + ), + OpInfo( + "nn.functional.triplet_margin_loss", + sample_inputs_func=sample_inputs_triplet_margin_loss, + error_inputs_func=error_inputs_triplet_margin_loss, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ), + OpInfo( + "nn.functional.triplet_margin_with_distance_loss", + sample_inputs_func=partial(sample_inputs_triplet_margin_loss, with_distance=True), + error_inputs_func=error_inputs_triplet_margin_loss, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # This test cannot handle a callable passed to `distance_function`. If we would use + # `distance_function=None`, the test would pass fine. + DecorateInfo( + unittest.expectedFailure, + "TestJit", + "test_variant_consistency_jit", + ), + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + ), + ), + BinaryUfuncInfo('nextafter', + dtypes=floating_types_and(torch.bfloat16, torch.half), + supports_autograd=False, + supports_rhs_python_scalar=False), + OpInfo( + "to", + op=lambda x, *args, **kwargs: x.to(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + sample_inputs_func=sample_inputs_to, + skips=( + # RuntimeError: undefined value cpu + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="cpu", + ), + # NotImplementedError: Cannot copy out of meta tensor; no data! + DecorateInfo( + unittest.skip("Skipped!"), + "TestMeta", + "test_meta_outplace", + ), + # https://github.com/pytorch/pytorch/issues/84335 + DecorateInfo( + unittest.skip("Skipped!"), + "TestProxyTensorOpInfo", + "test_make_fx_symbolic_exhaustive", + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + ), + ), + OpInfo('topk', + dtypes=all_types_and(torch.bfloat16, torch.float16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + sample_inputs_func=sample_inputs_topk), + # Multiple variants for batch_norm to test with and without cuDNN disabled + # See https://github.com/pytorch/pytorch/pull/63218#discussion_r688549391 for more details + OpInfo('nn.functional.batch_norm', + aten_name='batch_norm', + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.uint8, torch.bool, torch.int8, torch.int16, torch.int32 + ), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + allow_cow_input_materialize_forward=[1, 2], + allow_cow_input_materialize_backward=[1, 2], + sample_inputs_func=sample_inputs_batch_norm, + skips=( + # see https://github.com/pytorch/pytorch/issues/71286 + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness'), + DecorateInfo(unittest.skip('Skipped!'), 'TestNNCOpInfo', 'test_nnc_correctness', + device_type='cpu', dtypes=(torch.bfloat16, torch.float16)), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=5e-05, rtol=1e-05)}), + 'TestCompositeCompliance', 'test_forward_ad', device_type="cpu"), + )), + # This variant tests batch_norm with cuDNN disabled only on CUDA devices + OpInfo('nn.functional.batch_norm', + variant_test_name='without_cudnn', + aten_name='batch_norm', + dtypes=empty_types(), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + allow_cow_input_materialize_forward=[1, 2], + allow_cow_input_materialize_backward=[1, 2], + decorators=[onlyCUDA, disablecuDNN], + skips=( + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-03, rtol=1e-04)}), + 'TestJit', 'test_variant_consistency_jit'), + ), + sample_inputs_func=sample_inputs_batch_norm), + OpInfo( + "nn.functional.binary_cross_entropy", + aten_backward_name='binary_cross_entropy_backward', + sample_inputs_func=sample_inputs_binary_cross_entropy, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + gradcheck_fast_mode=False, + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=( + # RuntimeError: expected int at position 0, but got: Tensor + DecorateInfo( + unittest.skip("Skipped!"), + "TestCudaFuserOpInfo", + ), + # RuntimeError: expected int at position 0, but got: Tensor + DecorateInfo( + unittest.skip("Skipped!"), + "TestNNCOpInfo", + "test_nnc_correctness", + ), + # Fails for unknown reason: https://github.com/pytorch/pytorch/issues/120783 + DecorateInfo( + unittest.skip("Skipped!"), + "TestCompositeCompliance", + "test_cow_input", + device_type='cuda', + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-3, rtol=1e-3)}), + "TestJit", + "test_variant_consistency_jit", + ), + # RuntimeError: output with shape [] doesn't match the broadcast shape [5, 5] + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_meta_outplace'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_outplace'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_outplace_all_strides'), + ), + skips=( + # RuntimeError: expected int at position 0, but got: Tensor + DecorateInfo( + unittest.expectedFailure, + "TestJit", + "test_variant_consistency_jit", + ), + ), + ), + # We have to add 2 OpInfo entry for `igamma` and `igammac`.First is the + # standard entry, second is to run gradcheck tests on the second argument. + BinaryUfuncInfo('igamma', + dtypes=floating_types_and(torch.bfloat16, torch.float16), + aliases=('torch.special.gammainc',), + dtypesIfCUDA=floating_types(), + # TODO: FIXME + supports_rhs_python_scalar=False, + supports_autograd=False, + skips=( + # FIXME: incorrectly tries to pass a rhs scalar + DecorateInfo(unittest.expectedFailure, 'TestJit', + 'test_jit_alias_remapping'), + )), + # TODO: FIXME, ideally by implemented grad for both inputs + # BinaryUfuncInfo('igamma', + # variant_test_name='grad_other', + # # Since autograd formula is implemented only for other and + # # gradcheck test verifies the formula for input in SampleInput, + # # we permute the arguments. + # op=lambda self, other, **kwargs: torch.igamma(other, self, **kwargs), + # inplace_variant=None, + # method_variant=None, + # supports_rhs_python_scalar=False, + # rhs_make_tensor_kwargs=dict(requires_grad=False), + # dtypes=floating_types_and(torch.bfloat16, torch.float16), + # backward_dtypesIfCPU=floating_types_and(torch.bfloat16), + # dtypesIfCUDA=floating_types(), + # backward_dtypesIfCUDA=floating_types(), + # supports_inplace_autograd=False, + # skips=( + # # Derivative wrt first tensor not implemented + # DecorateInfo(unittest.expectedFailure, "TestCommon", + # "test_floating_inputs_are_differentiable"),"), + # # test does not work with passing lambda for op + # # AssertionError: False is not true : Tensors failed to compare as equal! + # DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # # test fails are we permute the arguments function variant + # # but not for inplace or method. + # DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + # # TypeError: igamma(): argument 'input' (position 1) must be Tensor, not float + # DecorateInfo(unittest.skip('Skipped!'), 'TestBinaryUfuncs'), + # )), + BinaryUfuncInfo('igammac', + dtypes=floating_types_and(torch.bfloat16, torch.float16), + aliases=('torch.special.gammaincc',), + dtypesIfCUDA=floating_types(), + supports_autograd=False, + supports_rhs_python_scalar=False, + skips=( + # FIXME: incorrectly tries to pass a rhs scalar + DecorateInfo(unittest.expectedFailure, 'TestJit', + 'test_jit_alias_remapping'), + )), + # TODO: FIXME, ideally by implementing grad for both inputs + # BinaryUfuncInfo('igammac', + # variant_test_name='grad_other', + # # Since autograd formula is implemented only for other and + # # gradcheck test verifies the formula for input in SampleInput, + # # we permute the arguments + # op=lambda self, other, **kwargs: torch.igammac(other, self, **kwargs), + # inplace_variant=None, + # method_variant=None, + # supports_rhs_python_scalar=False, + # rhs_make_tensor_kwargs=dict(requires_grad=False), + # dtypes=floating_types_and(torch.bfloat16, torch.float16), + # backward_dtypesIfCPU=floating_types_and(torch.bfloat16), + # dtypesIfCUDA=floating_types(), + # backward_dtypesIfCUDA=floating_types(), + # supports_inplace_autograd=False, + # decorators=[ + # # Derivative wrt first tensor not implemented + # DecorateInfo(unittest.expectedFailure, "TestCommon", + # "test_floating_inputs_are_differentiable"), + # ], + # skips=( + # # test does not work with passing lambda for op + # # AssertionError: False is not true : Tensors failed to compare as equal! + # DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # # test fails are we permute the arguments function variant + # # but not for inplace or method. + # DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + # # TypeError: igammac(): argument 'input' (position 1) must be Tensor, not float + # DecorateInfo(unittest.skip('Skipped!'), 'TestBinaryUfuncs'), + # )), + UnaryUfuncInfo('nn.functional.softshrink', + aten_name="softshrink", + aten_backward_name='softshrink_backward', + dtypes=floating_types_and(torch.bfloat16, torch.float16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=False, + sample_inputs_func=sample_inputs_softshrink, + error_inputs_func=error_inputs_softshrink), + UnaryUfuncInfo('nn.functional.hardshrink', + aten_name="hardshrink", + aten_backward_name='hardshrink_backward', + dtypes=floating_types_and(torch.bfloat16, torch.float16), + assert_autodiffed=True, + sample_inputs_func=sample_inputs_hardshrink, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + autodiff_nonfusible_nodes=["aten::hardshrink"]), + UnaryUfuncInfo('nn.functional.hardtanh', + aten_name="hardtanh", + aten_backward_name='hardtanh_backward', + dtypes=floating_types_and(torch.int8, torch.int16, torch.int32, torch.int64, torch.half, torch.bfloat16), + backward_dtypes=all_types_and(torch.half, torch.bfloat16), + backward_dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + assert_autodiffed=True, + sample_inputs_func=sample_inputs_hardtanh, + error_inputs_func=error_inputs_hardtanh, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + autodiff_nonfusible_nodes=["aten::hardtanh"]), + OpInfo('nn.functional.gelu', + aten_name="gelu", + aten_backward_name='gelu_backward', + ref=reference_gelu if TEST_SCIPY else None, + error_inputs_func=error_inputs_gelu, + supports_autograd=True, + assert_autodiffed=True, + sample_inputs_func=sample_inputs_gelu, + dtypes=floating_types_and(torch.bfloat16, torch.half), + supports_gradgrad=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + autodiff_nonfusible_nodes=["aten::gelu"], + skips=( + # AssertionError: Tensor-likes are not close! + # May not replicate in CI + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out'), + DecorateInfo(unittest.skip("Unsupported on MPS for now"), 'TestCommon', 'test_numpy_ref_mps'), + )), + UnaryUfuncInfo('nn.functional.relu6', + aten_name="relu6", + dtypes=all_types_and(torch.half, torch.bfloat16), + backward_dtypes=floating_types_and(torch.half, torch.bfloat16), + assert_autodiffed=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + autodiff_nonfusible_nodes=["aten::relu6"]), + OpInfo('mm', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_mm, + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + )), + OpInfo('mode', + op=torch.mode, + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Resized a non-empty tensor but did not warn about it + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # FIXME: + # Expected 2114 but got 1123. + # Absolute difference: 991 (up to 0.001 allowed) + # Relative difference: 0.46877956480605487 (up to 0.001 allowed) + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_compare_cpu", + dtypes=(torch.float32,), + device_type="cuda", + ), + # The operator 'aten::mode' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + sample_inputs_func=sample_inputs_mode,), + make_mvlgamma_opinfo(variant_test_name='mvlgamma_p_1', + domain=(1, None), + skips=skips_mvlgamma(), + sample_kwargs=lambda device, dtype, input: ({'p': 1}, {'d': 1})), + make_mvlgamma_opinfo(variant_test_name='mvlgamma_p_3', + domain=(2, None), + skips=skips_mvlgamma(), + sample_kwargs=lambda device, dtype, input: ({'p': 3}, {'d': 3})), + make_mvlgamma_opinfo(variant_test_name='mvlgamma_p_5', + domain=(3, None), + skips=skips_mvlgamma(), + sample_kwargs=lambda device, dtype, input: ({'p': 5}, {'d': 5})), + BinaryUfuncInfo('ne', + ref=np.not_equal, + aliases=('not_equal',), + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + always_returns_bool=True, + supports_autograd=False, + ), + OpInfo('narrow', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=partial(sample_inputs_narrow_narrow_copy, is_narrow=True), + reference_inputs_func=partial(reference_inputs_narrow_narrow_copy, is_narrow=True), + error_inputs_func=partial(error_inputs_narrow_narrow_copy, is_narrow=True, is_ref=False), + skips=( + # Use of .item() + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_operator'), + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_backward'), + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_forward_ad'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + )), + OpInfo('narrow_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/86931 + sample_inputs_func=partial(sample_inputs_narrow_narrow_copy, is_narrow=False), + reference_inputs_func=partial(reference_inputs_narrow_narrow_copy, is_narrow=False), + error_inputs_func=partial(error_inputs_narrow_narrow_copy, is_narrow=False, is_ref=False), + skips=( + # https://github.com/pytorch/pytorch/issues/84577 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # Could not run 'aten::narrow_copy.out' with arguments from the 'CUDA' backend + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_meta_outplace', + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_meta_outplace', + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_outplace', + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_outplace_all_strides'), + )), + OpInfo('view_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + ref=lambda x, newshape: np.reshape(x, newshape).copy(), + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_autograd=True, + sample_inputs_func=sample_inputs_view_reshape, + error_inputs_func=error_inputs_view_reshape, + skips=( + # RuntimeError: view size is not compatible with input tensor's size and stride + # (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead. + DecorateInfo( + unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides" + ), + )), + UnaryUfuncInfo('neg', + aliases=('negative', ), + ref=np.negative, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16, torch.chalf), + error_inputs_func=error_inputs_neg, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + assert_autodiffed=True), + OpInfo('dist', + op=torch.dist, + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + # torch.autograd.gradcheck.GradcheckError: While computing batched gradients, got: + # Could not allocate memory to change Tensor SizesAndStrides! + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_dist), + OpInfo('outer', + op=torch.outer, + aliases=('ger', ), + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_outer,), + OpInfo('ormqr', + op=torch.ormqr, + dtypes=floating_and_complex_types(), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=False, + supports_fwgrad_bwgrad=False, + sample_inputs_func=sample_inputs_ormqr, + error_inputs_func=error_inputs_ormqr, + decorators=[skipCUDAIfNoCusolver, skipCPUIfNoLapack], + skips=( + # Strides are not the same! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # Error: The operator 'aten::geqrf' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + )), + OpInfo('permute', + ref=np.transpose, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + assert_autodiffed=True, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + assert_jit_shape_analysis=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_varargs=True, + sample_inputs_func=sample_inputs_permute, + reference_inputs_func=reference_inputs_permute), + OpInfo('permute_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=True, + assert_autodiffed=True, + assert_jit_shape_analysis=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_varargs=False, # torch.permute is also not varargs + sample_inputs_func=sample_inputs_permute, + reference_inputs_func=reference_inputs_permute, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=(torch.float32,)), + )), + BinaryUfuncInfo('pow', + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.half, torch.bfloat16, torch.chalf), + dtypesIfMPS=all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool), + ref=np.power, + # Due to AVX2 currently not being fully supported for Float16, log_vml_cpu can't be enabled + # for Float16, causing this test to fail. pow's autograd for Float16 is thus currently + # unsupported on CPU. + backward_dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.bfloat16, torch.half, torch.chalf), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_inplace_autograd=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True, + supports_one_python_scalar=True, + # Integer types do not support negative exponentes + rhs_make_tensor_kwargs=dict(low=0), + # Raising negative real numbers to fractional powers is not supported + lhs_make_tensor_kwargs=dict(low=0), + decorators=( + DecorateInfo(toleranceOverride({torch.complex64: tol(atol=1e-4, rtol=1.3e-05)}), + 'TestBinaryUfuncs', 'test_reference_numerics'), + DecorateInfo(toleranceOverride({torch.complex64: tol(atol=1e-4, rtol=1.3e-05), + torch.complex128: tol(atol=1e-4, rtol=1.3e-05)}), + 'TestBinaryUfuncs', 'test_scalar_support'), + ), + skips=( + # Skipping integers because they are being raised to negative powers causing an error + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_reference_numerics_small_values', + dtypes=[torch.int8, torch.int16, torch.int32, torch.int64]), + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_reference_numerics_large_values', + dtypes=[torch.int16, torch.int32, torch.int64]), + # FIXME Complex values error with: Greatest absolute difference: nan at index + # Ref: https://github.com/pytorch/pytorch/issues/76853 + # For `chalf`, reference computation in `numpy` is computed in `cfloat`. + # Output of `chalf` saturates to `inf` quicker than reference due to its small range + # which leads to failure of this test. + DecorateInfo(unittest.skip("Skipped!"), 'TestDecomp', 'test_quick', + dtypes=(torch.complex32,), active_if=TEST_WITH_ROCM), + # FIXME: + # Mismatched elements: 1 / 500 (0.2%) + # Greatest absolute difference: nan at index (7, 9, 0) (up to 1e-05 allowed) + # Greatest relative difference: nan at index (7, 9, 0) (up to 0.001 allowed) + DecorateInfo(unittest.skip("Skipped!"), 'TestDecomp', 'test_comprehensive', + dtypes=(torch.complex32,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_complex_half_reference_testing', + dtypes=(torch.complex32,), active_if=TEST_WITH_ROCM), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_batch_vs_slicing', + dtypes=(torch.complex32,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_non_contig', + dtypes=(torch.complex32,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_reference_numerics', + dtypes=(torch.complex32,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_reference_numerics_small_values', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_reference_numerics_large_values', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_reference_numerics_extremal_values', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + )), + BinaryUfuncInfo('float_power', + ref=np.float_power, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16, torch.bool), + promotes_int_to_float=True, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_one_python_scalar=True, + # Integer types do not support negative exponentes + rhs_make_tensor_kwargs=dict(low=0), + # Raising negative real numbers to fractional powers is not supported + lhs_make_tensor_kwargs=dict(low=0), + decorators=( + DecorateInfo(toleranceOverride({torch.complex64: tol(atol=1e-4, rtol=1.3e-05), + torch.complex128: tol(atol=1e-4, rtol=1.3e-05)}), + 'TestBinaryUfuncs', 'test_scalar_support'), + ), + skips=( + # FIXME + # AssertionError: Object comparison failed: torch.float64 != torch.float32 + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_type_promotion'), + # -3.43399e+38 is outside the range of representable values of type 'float' + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # Complex values error with: Greatest absolute difference: nan at index + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_reference_numerics_small_values', + dtypes=[torch.complex64, torch.complex128]), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_reference_numerics_large_values', + dtypes=[torch.complex64, torch.complex128]), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_reference_numerics_extremal_values', + dtypes=[torch.complex64, torch.complex128]), + # Inplace always promotes to double and thus other floating dtypes are not supported + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_meta_inplace', + dtypes=[torch.bfloat16, torch.float16, torch.float32]), + # TypeError: Cannot convert a MPS Tensor to float64 dtype + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_requires_grad_error', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_promotes_int_to_float', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + )), + OpInfo('qr', + op=torch.qr, + dtypes=floating_and_complex_types(), + dtypesIfMPS=floating_types(), + sample_inputs_func=sample_inputs_linalg_qr_geqrf, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # In-place ops + check_batched_gradgrad=False, + decorators=[skipCUDAIfNoCusolver, skipCPUIfNoLapack], + ), + UnaryUfuncInfo('rad2deg', + ref=np.degrees, + decorators=(precisionOverride({torch.bfloat16: 7e-1, + torch.float16: 7e-1}),), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + promotes_int_to_float=True), + UnaryUfuncInfo('real', + ref=np.real, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + skips=( + # Skip since real and imag don't have out variants. + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_out_arg_all_dtypes'), + )), + OpInfo( + "roll", + ref=np.roll, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf), + error_inputs_func=error_inputs_roll, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_roll, + decorators=(onlyNativeDeviceTypes,), + ), + OpInfo( + "rot90", + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half), + error_inputs_func=error_inputs_rot90, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_rot90, + ), + # To test reference numerics against multiple values of argument `decimals`, + # we make multiple OpInfo entries with each entry corresponding to different value of decimals. + UnaryUfuncInfo('round', + ref=np.round, + aliases=('special.round',), + dtypes=all_types_and(torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo(unittest.expectedFailure, + 'TestNNCOpInfo', + 'test_nnc_correctness', + dtypes=tuple(t for t in integral_types() if t != torch.uint8)), + DecorateInfo(unittest.skip("Skipped!"), + 'TestNNCOpInfo', + 'test_nnc_correctness', + dtypes=(torch.bfloat16,)), + ), + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + assert_autodiffed=True, + ), + UnaryUfuncInfo('round', + ref=np.round, + variant_test_name='decimals_0', + aliases=('special.round',), + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfMPS=all_types_and(torch.half, torch.bfloat16), + sample_kwargs=lambda device, dtype, input: ({'decimals': 0}, {'decimals': 0}), + sample_inputs_func=partial(sample_inputs_elementwise_unary, op_kwargs={'decimals': 0}), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=False, + supports_sparse_csr=False), + UnaryUfuncInfo('round', + ref=np.round, + variant_test_name='decimals_3', + aliases=('special.round',), + dtypes=floating_types_and(torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + sample_kwargs=lambda device, dtype, input: ({'decimals': 3}, {'decimals': 3}), + sample_inputs_func=partial(sample_inputs_elementwise_unary, op_kwargs={'decimals': 3}), + skips=( + # test_ops already tested for this overload with `decimals_0` opinfo entry + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon'), + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestJit'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits'), + DecorateInfo(toleranceOverride({torch.bfloat16: tol(atol=1e-3, rtol=0.016)}), + "TestUnaryUfuncs", "test_reference_numerics_extremal", + device_type="cuda"), + DecorateInfo(toleranceOverride({torch.bfloat16: tol(atol=1e-3, rtol=0.016)}), + "TestUnaryUfuncs", "test_reference_numerics_normal", + device_type="cuda"), + ), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=False, + supports_sparse_csr=False), + UnaryUfuncInfo('round', + ref=np.round, + variant_test_name='decimals_neg_3', + aliases=('special.round',), + dtypes=floating_types_and(torch.bfloat16), + dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + sample_kwargs=lambda device, dtype, input: ({'decimals': -3}, {'decimals': -3}), + sample_inputs_func=partial(sample_inputs_elementwise_unary, op_kwargs={'decimals': -3}), + skips=( + # test_ops already tested for this overload with `decimals_0` opinfo entry + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon'), + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestJit'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits'), + ), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=False, + supports_sparse_csr=False), + UnaryUfuncInfo('sin', + ref=np.sin, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + assert_autodiffed=True, + handles_large_floats=False, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + skips=( + # Fails on CUDA but passes on ROCm + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=(torch.cdouble,), device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=(torch.cfloat, torch.cdouble), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=2e-3)}), + "TestConsistency", "test_output_grad_match", device_type="mps"), + ), + decorators=(precisionOverride({torch.bfloat16: 1e-2}),)), + UnaryUfuncInfo('sinc', + ref=np_sinc_with_fp16_as_fp32, + aliases=('special.sinc',), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + handles_large_floats=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True), + UnaryUfuncInfo('sinh', + ref=np_unary_ufunc_integer_promotion_wrapper(np.sinh), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + promotes_int_to_float=True, + decorators=(precisionOverride({torch.float16: 1e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=(torch.cdouble,)), + # Reference: https://github.com/pytorch/pytorch/issues/48641 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.int8]), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + )), + UnaryUfuncInfo('sign', + ref=reference_sign, + dtypes=all_types_and(torch.bool, torch.bfloat16, torch.half), + dtypesIfCUDA=all_types_and(torch.bool, torch.bfloat16, torch.half), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/41245 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.bfloat16, torch.float16, torch.float32, torch.float64]), + )), + UnaryUfuncInfo('sgn', + ref=reference_sgn, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf), + dtypesIfMPS=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf, torch.complex32), + backward_dtypes=floating_and_complex_types_and(torch.bfloat16, torch.half), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.bfloat16, torch.half, torch.chalf), + backward_dtypesIfMPS=floating_and_complex_types_and(torch.bfloat16, torch.half, torch.complex32), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/41245 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.bfloat16, torch.float16, torch.float32, torch.float64]), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + )), + OpInfo('split', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool, torch.chalf), + sample_inputs_func=partial(sample_inputs_split, list_args=False), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + assert_autodiffed=True), + OpInfo('split', + # Cannot declare this aten_name because of + # test_variant_consistency_jit_split_list_args_cpu_float32 + decomp_aten_name='split_with_sizes', + variant_test_name='list_args', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), + sample_inputs_func=partial(sample_inputs_split, list_args=True), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + # `unsafe_split` supports only `int` for split_size argument + OpInfo('unsafe_split', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool, torch.chalf), + sample_inputs_func=partial(sample_inputs_split, list_args=False), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + assert_autodiffed=True, + check_batched_forward_grad=False), + OpInfo('split_with_sizes', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool, torch.chalf), + sample_inputs_func=sample_inputs_split_with_sizes, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True), + OpInfo('split_with_sizes_copy', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool, torch.chalf), + sample_inputs_func=sample_inputs_split_with_sizes, + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # No error raised + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out_requires_grad_error"), + )), + BinaryUfuncInfo('__radd__', + op=torch.Tensor.__radd__, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), + supports_out=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit',), + + ), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + autodiff_nonfusible_nodes=['aten::add'],), + BinaryUfuncInfo('__rdiv__', + op=torch.Tensor.__rdiv__, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), + promotes_int_to_float=True, + lhs_make_tensor_kwargs={'exclude_zero': True}, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + skips=( + # https://github.com/pytorch/pytorch/issues/76806 + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit',), + ), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True, + autodiff_nonfusible_nodes=['aten::mul', 'aten::reciprocal'],), + BinaryUfuncInfo('__rmul__', + op=torch.Tensor.__rmul__, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), + supports_out=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit',), + ), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + autodiff_nonfusible_nodes=['aten::mul'],), + BinaryUfuncInfo('__rand__', + op=torch.Tensor.__rand__, + dtypes=integral_types_and(torch.bool), + supports_out=False, + supports_autograd=False, + supports_forward_ad=True, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + )), + BinaryUfuncInfo('__ror__', + op=torch.Tensor.__ror__, + dtypes=integral_types_and(torch.bool), + supports_out=False, + supports_autograd=False, + supports_forward_ad=True, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + )), + BinaryUfuncInfo('__rxor__', + op=torch.Tensor.__rxor__, + dtypes=integral_types_and(torch.bool), + supports_out=False, + supports_autograd=False, + supports_forward_ad=True, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + )), + OpInfo('__rmatmul__', + op=torch.Tensor.__rmatmul__, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, + *[torch.bfloat16] + if SM53OrLater or TEST_WITH_ROCM else []), + assert_autodiffed=True, + sample_inputs_func=partial(sample_inputs_matmul, is_rmatmul=True), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + decorators=( + # NVIDIA only assures that bfloat16 is supported by bmm if SM >= 5.3 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_dtypes', device_type='cuda', active_if=not SM53OrLater), + DecorateInfo(toleranceOverride({torch.complex64: tol(atol=1e-05, rtol=1.2e-03)}), + 'TestMathBits', 'test_conj_view'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-05, rtol=1.2e-03)}), + 'TestCommon', 'test_noncontiguous_samples'), + DecorateInfo(toleranceOverride({torch.complex64: tol(atol=1e-05, rtol=1e-05)}), + "TestDecomp", "test_comprehensive", device_type="cuda", + active_if=TEST_WITH_ROCM), + ), + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit',), + # https://github.com/pytorch/pytorch/issues/67470 + DecorateInfo(unittest.skip("67470!"), + 'TestCommon', 'test_noncontiguous_samples', + device_type='cpu', dtypes=(torch.long,)), + # Fails on XLA. + # AssertionError: False is not true : Tensors failed to compare as equal + DecorateInfo(unittest.skip("Skipped!"), 'TestOpInfo', device_type='xla', dtypes=(torch.long,)), + # https://github.com/pytorch/pytorch/issues/71774 + DecorateInfo(unittest.skip('Skipped!'), 'TestNNCOpInfo', 'test_nnc_correctness', + device_type='cpu', dtypes=(torch.long,)), + )), + BinaryUfuncInfo('__rmod__', + op=torch.Tensor.__rmod__, + dtypes=floating_types_and(torch.bfloat16, torch.half,), + dtypesIfCUDA=all_types_and(torch.bfloat16, torch.half), + dtypesIfMPS=all_types_and(torch.bfloat16, torch.half, torch.bool), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_one_python_scalar=True, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit',), + # NotImplementedError: "remainder_cpu" not implemented for * + DecorateInfo( + unittest.expectedFailure, 'TestConsistency', device_type='mps', + dtypes=(torch.bool, torch.int16, torch.uint8, torch.int8, torch.int32, torch.int64) + ), + ), + # Support autograd after torch.remainder(Tensor, Tensor) supports + # autograd of the second argument. + # https://github.com/pytorch/pytorch/pull/58476/files#r637167630 + # supports_autograd=False, + assert_autodiffed=True, + autodiff_nonfusible_nodes=['aten::remainder'],), + BinaryUfuncInfo('__rpow__', + op=torch.Tensor.__rpow__, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half), + dtypesIfMPS=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), + # Reference: https://github.com/pytorch/pytorch/issues/54774 + # "log2" "_vml_cpu" not implemented for Half + backward_dtypes=all_types_and_complex_and(torch.bfloat16, torch.half), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_one_python_scalar=True, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit',), + # TODO: FIXME tolerance is too high + DecorateInfo(unittest.skip('Skipped!'), 'TestFwdGradients'), + DecorateInfo(unittest.skip('Skipped!'), 'TestBwdGradients'), + ), + assert_autodiffed=True, + autodiff_nonfusible_nodes=['aten::pow'],), + BinaryUfuncInfo('__rsub__', + op=torch.Tensor.__rsub__, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + supports_one_python_scalar=True, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit',), + ), + assert_autodiffed=True, + autodiff_nonfusible_nodes=['aten::rsub'],), + BinaryUfuncInfo('rsub', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + supports_inplace_autograd=False, + assert_autodiffed=None, + sample_inputs_func=sample_inputs_add_sub), + OpInfo('select', + aten_backward_name='select_backward', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool, torch.chalf), + sample_inputs_func=sample_inputs_select, + assert_jit_shape_analysis=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + OpInfo('select_scatter', + dtypes=all_types_and(torch.bfloat16, torch.half, torch.bool), + sample_inputs_func=sample_inputs_select_scatter, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False), + OpInfo('slice', + op=torch.ops.aten.slice.Tensor, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool, torch.chalf), + sample_inputs_func=sample_inputs_slice, + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_scripting=False, + supports_inplace_autograd=False, + supports_out=False), + OpInfo('slice_scatter', + dtypes=all_types_and(torch.bfloat16, torch.half, torch.bool), + sample_inputs_func=sample_inputs_slice_scatter, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=True), + UnaryUfuncInfo('signbit', + ref=np.signbit, + dtypes=all_types_and(torch.bool, torch.bfloat16, torch.half), + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_autograd=False,), + UnaryUfuncInfo('tan', + ref=np.tan, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + decorators=(DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-04, rtol=1e-05)}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cuda'),), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + # FIXME: + # Mismatched elements: 2 / 400 (0.5%) + # Greatest absolute difference: inf at index (7, 16) (up to 1e-05 allowed) + # Greatest relative difference: nan at index (7, 16) (up to 0.001 allowed) + DecorateInfo( + unittest.skip("Skipped!"), + "TestInductorOpInfo", + "test_comprehensive", + dtypes=(torch.float16,), + device_type="cuda", + ), + DecorateInfo(toleranceOverride({torch.complex64: tol(atol=3e-5, rtol=7e-6)}), + "TestConsistency", "test_output_match", device_type="mps"), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=2e-3)}), + "TestConsistency", "test_output_grad_match", device_type="mps"), + ), + # tan(pi/2 * odd_number) is nan + reference_numerics_filter=NumericsFilter( + condition=lambda x: close_to_int(x / (math.pi * 0.5)), safe_val=math.pi)), + UnaryUfuncInfo('tanh', + ref=np.tanh, + aten_backward_name='tanh_backward', + aliases=('nn.functional.tanh',), + decorators=(precisionOverride({torch.bfloat16: 1e-2}), + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-04, rtol=2e-05)}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cuda'),), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + assert_autodiffed=True, + assert_jit_shape_analysis=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='xpu', dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + DecorateInfo(toleranceOverride({torch.complex64: tol(atol=3e-5, rtol=7e-6)}), + "TestConsistency", "test_output_match", device_type="mps"), + ), + # tan(j * pi/2 * odd_number) is nan + reference_numerics_filter=NumericsFilter( + condition=lambda x: (close_to_int(x / (math.pi * 0.5j)) + if x.is_complex() else x.new_tensor(False, dtype=torch.bool)), + safe_val=0)), + OpInfo('tensor_split', + ref=np.array_split, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Pre-existing condition; Needs to be fixed + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_operator'), + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_backward'), + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_forward_ad'), + ), + sample_inputs_func=sample_inputs_tensor_split,), + OpInfo('hsplit', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.bfloat16, torch.float16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_hsplit, + error_inputs_func=error_inputs_hsplit,), + OpInfo('vsplit', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.bfloat16, torch.float16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_vsplit, + error_inputs_func=error_inputs_vsplit,), + OpInfo('dsplit', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.bfloat16, torch.float16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_dsplit, + error_inputs_func=error_inputs_dsplit,), + OpInfo('triangular_solve', + op=torch.triangular_solve, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_legacy_solve, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_wrapper=lambda *args, **kwargs: gradcheck_wrapper_triangular_input(*args, idx=1, **kwargs), + decorators=[ + skipCUDAIfNoMagma, + skipCPUIfNoLapack, + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=3e-5, rtol=3e-6)}), + 'TestConsistency', 'test_output_match', device_type='cpu', + ), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=6e-5, rtol=3e-6)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ], + skips=( + # AssertionError: Scalars are not equal! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # Gradcheck fails + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_fn_fwgrad_bwgrad', + dtypes=floating_and_complex_types()), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out', + device_type='mps', dtypes=[torch.float32]), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=[torch.float32]), + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', + device_type='mps', dtypes=[torch.float32]), + # RuntimeError: linalg.solve.triangular(); Only float is supported! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), + # see https://github.com/pytorch/pytorch/issues/177251 + DecorateInfo( + unittest.expectedFailure, + 'TestOperators', + 'test_jvp', + device_type='cpu', + dtypes=[torch.float32], + active_if=IS_ARM64 and IS_CPU_EXT_SVE_SUPPORTED, + ), + )), + UnaryUfuncInfo('trunc', + aliases=('fix', ), + ref=np.trunc, + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfMPS=all_types_and(torch.half, torch.bfloat16, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + skips=( + DecorateInfo(unittest.expectedFailure, + 'TestNNCOpInfo', + 'test_nnc_correctness', + dtypes=tuple(t for t in integral_types() if t != torch.uint8)), + ), + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + assert_autodiffed=True), + UnaryUfuncInfo('exp2', + aliases=('special.exp2', ), + ref=np_unary_ufunc_integer_promotion_wrapper(np.exp2), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=[torch.cdouble]), + # Reference: https://github.com/pytorch/pytorch/issues/48010 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + )), + UnaryUfuncInfo('expm1', + aliases=('special.expm1', ), + ref=np_unary_ufunc_integer_promotion_wrapper(np.expm1), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + promotes_int_to_float=True, + assert_autodiffed=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.complex128]), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + )), + UnaryUfuncInfo('nan_to_num', + ref=np.nan_to_num, + dtypes=all_types_and(torch.half, torch.bool, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.half, torch.bool, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + skips=( + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + # AssertionError: The values for attribute 'shape' do not match: torch.Size([20]) != torch.Size([0]). + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ), + # Passing numpy_kwargs via sample_kwargs, as numpy does comparison + # with BFloat16 in float, since it currently doesn't support BFloat16. + # Ref: https://github.com/pytorch/pytorch/issues/57982#issuecomment-839150556 + sample_kwargs=lambda device, dtype, input: ({}, + {'posinf': torch.finfo(torch.bfloat16).max, + 'neginf': torch.finfo(torch.bfloat16).min}) + if dtype is torch.bfloat16 else ({}, {})), + UnaryUfuncInfo('reciprocal', + ref=np_unary_ufunc_integer_promotion_wrapper(np.reciprocal), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/45690 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.cfloat, torch.cdouble]), + )), + UnaryUfuncInfo('rsqrt', + ref=lambda x: np.reciprocal(np.sqrt(x)), + domain=(0, None), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + decorators=(precisionOverride({torch.half: 5e-2}),), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=(torch.cfloat, torch.cdouble), device_type='cuda'), + # AssertionError: Tensor-likes are not close! + # Greatest absolute difference: nan at index (700,) (up to 0.01 allowed) + # Greatest relative difference: nan at index (700,) (up to 0.001 allowed) + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=(torch.chalf,), device_type='cuda'), + )), + UnaryUfuncInfo('sqrt', + ref=np.sqrt, + supports_sparse=True, + domain=(0, None), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + decorators=( + precisionOverride({torch.bfloat16: 7e-2}), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=1e-2, rtol=0)}), + 'TestUnaryUfuncs', 'test_reference_numerics_large'), + ), + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/47358 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + device_type='cpu', dtypes=(torch.cfloat, torch.cdouble), + active_if=IS_MACOS), + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + DecorateInfo(toleranceOverride({torch.complex64: tol(atol=2e-5, rtol=3e-6)}), + "TestConsistency", "test_output_match", device_type="mps"), + )), + UnaryUfuncInfo('square', + ref=np.square, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + decorators=(precisionOverride({torch.complex64: 3e-4, torch.bfloat16: 3e-1}),), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/52549 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=[torch.cfloat, torch.cdouble]), + # >>> t = torch.tensor(complex(-0.01, float("inf"))) + # >>> np.square(t.numpy()) + # (-inf-infj) + # >>> t.square() + # tensor(-inf-infj) + # >>> t.cuda().square() + # tensor(inf+nanj, device='cuda:0') + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_meta_inplace', + dtypes=[torch.bool]), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_meta_inplace', + dtypes=[torch.bool]), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_inplace', + dtypes=[torch.bool]), + ),), + OpInfo('lerp', + dtypes=floating_and_complex_types_and(torch.bfloat16, torch.half), + dtypesIfCUDA=floating_and_complex_types_and(torch.chalf, torch.half, torch.bfloat16), + dtypesIfMPS=all_types_and_complex_and(torch.bfloat16, torch.half, torch.bool), + sample_inputs_func=sample_inputs_lerp, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # NotImplementedError: "lerp_kernel_scalar" not implemented for 'Bool' + DecorateInfo( + unittest.expectedFailure, 'TestConsistency', device_type='mps', + dtypes=(torch.bool, torch.uint8, torch.int8, torch.int64, torch.int32, torch.int16) + ), + ), + assert_autodiffed=True), + UnaryUfuncInfo('angle', + ref=np.angle, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool), + decorators=(precisionOverride({torch.float16: 1e-2, + torch.bfloat16: 1e-2}),), + backward_dtypes=floating_and_complex_types_and(torch.bfloat16, torch.float16), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.chalf), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_complex_to_float=True, + skips=( + # Ref: https://github.com/pytorch/pytorch/issues/78413 + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', 'test_reference_numerics_small', + dtypes=(torch.bfloat16, torch.float16, torch.float32, torch.float64),), + )), + UnaryUfuncInfo('isfinite', + ref=np.isfinite, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + supports_out=False, + supports_autograd=False), + UnaryUfuncInfo('isinf', + ref=np.isinf, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + supports_out=False, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_autograd=False), + UnaryUfuncInfo('isposinf', + ref=np.isposinf, + dtypes=all_types_and(torch.bool, torch.bfloat16, torch.float16), + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_autograd=False), + UnaryUfuncInfo('isneginf', + ref=np.isneginf, + dtypes=all_types_and(torch.bool, torch.bfloat16, torch.float16), + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_autograd=False), + UnaryUfuncInfo('isreal', + ref=np.isreal, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + supports_out=False, + supports_autograd=False), + UnaryUfuncInfo('isnan', + ref=np.isnan, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + supports_out=False, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_autograd=False), + OpInfo('einsum', + # we need this lambda because SampleInput expects tensor input as the first argument + # TODO(@heitorschueroff) update SampleInput to handle such cases + op=lambda tensors, equation: torch.einsum(equation, tensors), + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + # See https://github.com/pytorch/pytorch/issues/66357 + sample_inputs_func=sample_inputs_einsum, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # test does not work with passing lambda for op + # there's a test `test_einsum` in `test_jit.py` to handle this case + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + )), + OpInfo('svd', + op=torch.svd, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_svd, + # Runs very slowly on slow-gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + # We're using at::allclose, which does not have a batching rule + check_batched_grad=False, + check_batched_gradgrad=False, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack, with_tf32_off], + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + )), + OpInfo('svd_lowrank', + op=lambda *args, **kwargs: wrapper_set_seed( + lambda a, b, **kwargs: torch.svd_lowrank(a @ b.mT, **kwargs), + *args, **kwargs + ), + dtypes=floating_and_complex_types(), + dtypesIfMPS=floating_types(), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + # Due to the use of randomness + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + sample_inputs_func=sample_inputs_svd_lowrank, + decorators=[skipCUDAIfNoCusolver, skipCPUIfNoLapack, with_tf32_off, + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-03, rtol=1e-03), + torch.complex64: tol(atol=1e-02, rtol=1e-02)}), + 'TestCommon', 'test_noncontiguous_samples'), + # FIXME This should be the following, but the toleranceOverride does not seem to do anything! + # DecorateInfo(toleranceOverride({torch.complex128: tol(atol=1e-04, rtol=1e-04)}), + # 'TestFwdGradients', 'test_fn_fwgrad_bwgrad'), + DecorateInfo(unittest.skip("See comment above"), + 'TestFwdGradients', + 'test_fn_fwgrad_bwgrad', + dtypes=[torch.complex128]), + ], + skips=( + # test does not work with passing lambda for op + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo(unittest.expectedFailure, 'TestSchemaCheckModeOpInfo', 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + DecorateInfo(slowTest, 'TestCompositeCompliance', 'test_forward_ad'), + )), + OpInfo('pca_lowrank', + op=lambda *args, **kwargs: wrapper_set_seed( + lambda a, b, **kwargs: torch.pca_lowrank(a @ b.mT, **kwargs), + *args, **kwargs + ), + dtypes=floating_and_complex_types(), + dtypesIfMPS=floating_types(), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + check_batched_forward_grad=False, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_pca_lowrank, + decorators=[skipCUDAIfNoCusolver, skipCPUIfNoLapack, with_tf32_off, + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-03, rtol=1e-03), + torch.complex64: tol(atol=4e-02, rtol=4e-02)}), + 'TestCommon', 'test_noncontiguous_samples'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-05, rtol=5e-05)}), + 'TestOperators', 'test_grad'), + # FIXME This should be the following, but the toleranceOverride does not seem to do anything! + # DecorateInfo(toleranceOverride({torch.complex128: tol(atol=1e-04, rtol=1e-04)}), + # 'TestFwdGradients', 'test_fn_fwgrad_bwgrad'), + DecorateInfo(unittest.skip("See comment above"), + 'TestFwdGradients', + 'test_fn_fwgrad_bwgrad', + dtypes=[torch.complex128]), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=3e-5, rtol=1e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=5e-5, rtol=3e-4)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ], + skips=( + # test does not work with passing lambda for op + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo(unittest.expectedFailure, 'TestSchemaCheckModeOpInfo', 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + BinaryUfuncInfo('polar', + dtypes=floating_types(), + dtypesIfMPS=floating_types_and(torch.float16), + # this function is undefined if 'abs' values are <0 + supports_forward_ad=True, + lhs_make_tensor_kwargs=dict(low=0), + supports_rhs_python_scalar=False, + skips=( + # RuntimeError: Expected object of scalar type Float but got scalar type Double for second argument + DecorateInfo(unittest.skip('Skipped!'), 'TestBinaryUfuncs', 'test_type_promotion'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_binary_ufuncs_mixed_dtype'), + # GradcheckError: Jacobian computed with forward mode mismatch for output 0 with respect to input 0 + # Numerical: + # tensor([[0.]], dtype=torch.float64) + # Analytical: + # tensor([[-0.0047]], dtype=torch.float64, grad_fn=) + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_fn_fwgrad_bwgrad'), + # Comparing fp32 CPU grads with fp16 MPS ones leads to high errors + DecorateInfo(toleranceOverride({torch.float16: tol(atol=2e-5, rtol=5e-3)}), + 'TestConsistency', 'test_output_match', device_type='mps'), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=5e-3, rtol=5e-2)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + )), + # TODO(@kshitij12345): Refactor similar to `mvlgamma` entries. + # To test reference numerics against multiple values of argument `n`, + # we make multiple OpInfo entries with each entry corresponding to different value of n (currently 0 to 4). + # We run the op tests from test_ops.py only for `n=0` to avoid redundancy in testing. + UnaryUfuncInfo('polygamma', + op=lambda x, n, **kwargs: torch.polygamma(n, x, **kwargs), + variant_test_name='polygamma_n_0', + ref=reference_polygamma if TEST_SCIPY else None, + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + sample_inputs_func=sample_inputs_polygamma, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + ), + sample_kwargs=lambda device, dtype, input: ({'n': 0}, {'n': 0}), + # polygamma functions have multiple singularities at x having non-positive integer value + reference_numerics_filter=NumericsFilter(condition=lambda x: (x < 0.1) & ((x - x.round()).abs() < 1e-4), + safe_val=1)), + *(UnaryUfuncInfo('polygamma', + op=lambda x, n, **kwargs: torch.polygamma(n, x, **kwargs), + variant_test_name=f'polygamma_n_{n_}', + ref=reference_polygamma if TEST_SCIPY else None, + dtypes=all_types_and(torch.bool, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + sample_inputs_func=sample_inputs_polygamma, + decorators=( + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-3)}), 'TestUnaryUfuncs'), + DecorateInfo(toleranceOverride({torch.bfloat16: tol(atol=1e1, rtol=1e-1), + torch.float32: tol(atol=1e-4, rtol=1e-2)}), + 'TestUnaryUfuncs', 'test_reference_numerics_normal', + active_if=IS_WINDOWS), + ), + skips=( + # Redundant tests + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestJit'), + DecorateInfo(unittest.skip("Skipped!"), 'TestNormalizeOperators'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon'), + # Mismatch: https://github.com/pytorch/pytorch/issues/55357 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large'), + ), + sample_kwargs=lambda device, dtype, input: ({'n': n_}, {'n': n_}), + # polygamma functions have multiple singularities at x having non-positive integer value + reference_numerics_filter=NumericsFilter(condition=lambda x: (x < 0.1) & ((x - x.round()).abs() < 1e-4), + safe_val=1)) + for n_ in (1, 2, 3, 4)), + OpInfo('ravel', + ref=np.ravel, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_ravel, + ), + OpInfo('unravel_index', + ref=np.unravel_index, + dtypes=integral_types_and(), + supports_out=False, + supports_autograd=False, + sample_inputs_func=sample_inputs_unravel_index, + ), + OpInfo('reshape', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + sample_inputs_func=sample_inputs_view_reshape, + reference_inputs_func=reference_inputs_view_reshape, + error_inputs_func=error_inputs_view_reshape, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ), + OpInfo('reshape_as', + op=lambda x, other: x.reshape_as(other), + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + sample_inputs_func=partial(sample_inputs_view_reshape, tensor_arg=True), + reference_inputs_func=partial(reference_inputs_view_reshape, tensor_arg=True), + error_inputs_func=partial(error_inputs_view_reshape, tensor_arg=True), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + )), + OpInfo('view', + op=lambda x, shape: x.view(shape), + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + sample_inputs_func=sample_inputs_view_reshape, + reference_inputs_func=reference_inputs_view_reshape, + error_inputs_func=error_inputs_view_reshape, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: view size is not compatible with input tensor's size and stride + # (at least one dimension spans across two contiguous subspaces). Use .reshape(...) instead. + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides"), + )), + OpInfo('view_as', + op=lambda x, other: x.view_as(other), + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=partial(sample_inputs_view_reshape, tensor_arg=True), + reference_inputs_func=partial(reference_inputs_view_reshape, tensor_arg=True), + error_inputs_func=partial(error_inputs_view_reshape, tensor_arg=True), + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: view size is not compatible with input tensor's size and stride + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides") + )), + OpInfo('atleast_1d', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_atleast1d2d3d, + skips=( + # JIT does not support variadic tensors. + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":252, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=[torch.float32]), + ), + ), + OpInfo('atleast_2d', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=[torch.float32]), + ), + sample_inputs_func=sample_inputs_atleast1d2d3d, + ), + OpInfo('atleast_3d', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', dtypes=[torch.float32]), + ), + sample_inputs_func=sample_inputs_atleast1d2d3d, + ), + OpInfo('flatten', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + ref=reference_flatten, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_flatten, + reference_inputs_func=reference_inputs_flatten, + ), + OpInfo('unflatten', + op=torch.unflatten, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_unflatten, + ), + OpInfo('column_stack', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_column_stack,), + OpInfo('pinverse', + op=torch.pinverse, + dtypes=floating_and_complex_types(), + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + supports_out=False, + sample_inputs_func=sample_inputs_linalg_invertible, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=[torch.float32]), + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', + device_type='mps', dtypes=[torch.float32]), + )), + OpInfo('gather', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_gather, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + error_inputs_func=error_inputs_gather, + ), + OpInfo('index_fill', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.complex32), + inplace_variant=torch.Tensor.index_fill_, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + skips=( + # RuntimeError: Mismatch on aten._unique.default: Shapes torch.Size([2]) and torch.Size([1]) are not equal! + DecorateInfo(unittest.expectedFailure, 'TestFakeTensor', 'test_fake_crossref_backward_no_amp'), + # RuntimeError: Mismatch on aten._unique.default: Shapes torch.Size([2]) and torch.Size([1]) are not equal! + DecorateInfo(unittest.expectedFailure, 'TestFakeTensor', 'test_fake_crossref_backward_amp'), + ), + sample_inputs_func=sample_inputs_index, + reference_inputs_func=partial(sample_inputs_index, reference=True)), + OpInfo('index_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.complex32), + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_index, + reference_inputs_func=partial(sample_inputs_index, reference=True), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), + OpInfo('index_select', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16, torch.chalf), + sample_inputs_func=sample_inputs_index, + reference_inputs_func=partial(sample_inputs_index, reference=True), + error_inputs_func=error_inputs_index_select, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + skips=( + # AssertionError: RuntimeError not raised : Expected RuntimeError when calling with + # input.device=mps:0 and out.device=cpu. + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_out', + device_type='mps', dtypes=(torch.float32,) + ), + ), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), + OpInfo('index_add', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + inplace_variant=torch.Tensor.index_add_, + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_index, + reference_inputs_func=partial(sample_inputs_index, reference=True), + error_inputs_func=error_inputs_index_add, + skips=( + # boolean alpha not handled properly + DecorateInfo(unittest.expectedFailure, + 'TestNNCOpInfo', + 'test_nnc_correctness', + dtypes=(torch.bool,)), + ), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL), + *(OpInfo('index_reduce', + variant_test_name=reduction_type, + dtypes=all_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=custom_types( + torch.float32, torch.bfloat16, torch.float16, torch.int32, + torch.int16, torch.int8, torch.uint8, + ), + skips=( + DecorateInfo(toleranceOverride({torch.float16: tol(atol=2e-3, rtol=3e-3)}), + 'TestInductorOpInfo', 'test_comprehensive'), + ), + supports_out=True, + sample_inputs_func=sample_inputs_index_reduce, + ) for reduction_type in ('mean', 'prod', 'amin', 'amax')), + OpInfo('_unsafe_masked_index', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16, torch.bool), + supports_out=False, + supports_inplace_autograd=False, + supports_scripting=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs__unsafe_masked_index, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + DecorateInfo(slowTest, 'TestDecomp', 'test_quick_core_backward', + dtypes=(torch.float64,), active_if=IS_WINDOWS), + ),), + OpInfo('_unsafe_masked_index_put_accumulate', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16, torch.bool), + supports_out=False, + supports_inplace_autograd=False, + supports_scripting=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=2e-3, rtol=3e-2)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cpu' + ), + ), + sample_inputs_func=sample_inputs__unsafe_masked_index_put_accumulate, + skips=( + DecorateInfo(slowTest, 'TestDecomp', 'test_quick_core_backward', + dtypes=(torch.float64,), active_if=IS_WINDOWS), + ),), + OpInfo('__getitem__', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_inplace_autograd=False, + supports_scripting=False, + op=torch.Tensor.__getitem__, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # AssertionError: False is not true : Scalars failed to compare as equal! 0 != 104448 + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit', device_type='cuda'),), + sample_inputs_func=sample_inputs_getitem), + OpInfo('index_put', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + supports_inplace_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + test_neg_view=False, + sample_inputs_func=sample_inputs_index_put, + skips=( + DecorateInfo(unittest.skip("Skipped"), 'TestBwdGradients', 'test_fn_grad', dtypes=[torch.float64], + device_type='cuda', active_if=(TEST_WITH_ROCM and TEST_WITH_TORCHINDUCTOR)), + )), + OpInfo('sort', + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_sort, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], device_type='cuda', active_if=not TEST_WITH_ROCM), + )), + OpInfo('unique', + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16, torch.uint16, torch.uint32, torch.uint64), + sample_inputs_func=sample_inputs_unique, + supports_out=False, + supports_autograd=False, + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Output order is undefined when sorted=False'), 'TestCommon', 'test_compare_cpu'), + )), + OpInfo('unique_consecutive', + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_unique_consecutive, + supports_out=False, + supports_autograd=False, + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + )), + OpInfo('put', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + check_batched_gradgrad=False, # vmap complains of the sizes + sample_inputs_func=sample_inputs_put, + skips=( + # NotImplementedError: The operator 'aten::put_' is not currently + # implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + )), + + OpInfo('take', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + check_batched_grad=False, # vmap complains of the sizes + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_take, + skips=( + # The operator 'aten::take' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + error_inputs_func=error_inputs_take), + OpInfo('scatter', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_scatter, + error_inputs_func=error_inputs_scatter_and_scatter_add, + skips=( + # Compiler issue on ROCm. Regression started in ROCm 6.4. + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + )), + UnaryUfuncInfo( + 'bfloat16', + op=lambda x, *args, **kwargs: x.bfloat16(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + skips=( + # autograd tests don't handle operators that change dtype + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients'), + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients'), + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + )), + UnaryUfuncInfo( + 'bool', + op=lambda x, *args, **kwargs: x.bool(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attributis not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + )), + UnaryUfuncInfo( + 'byte', + op=lambda x, *args, **kwargs: x.byte(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_out=False, + sample_inputs_func=sample_inputs_byte, + # The autograd test runner cannot handle functions that change dtype + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + )), + UnaryUfuncInfo( + 'char', + op=lambda x, *args, **kwargs: x.char(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + # The autograd test runner cannot handle functions that change dtype + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + )), + UnaryUfuncInfo( + 'double', + op=lambda x, *args, **kwargs: x.double(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Error: Cannot convert a MPS Tensor to float64 dtype + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + )), + UnaryUfuncInfo( + 'float', + op=lambda x, *args, **kwargs: x.float(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + skips=( + # autograd tests don't handle operators that change dtype + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients'), + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients'), + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + )), + UnaryUfuncInfo( + 'half', + op=lambda x, *args, **kwargs: x.half(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + supports_autograd=True, + skips=( + # autograd tests don't handle operators that change dtype + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients'), + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients'), + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + )), + UnaryUfuncInfo( + 'int', + op=lambda x, *args, **kwargs: x.int(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + )), + UnaryUfuncInfo( + 'long', + op=lambda x, *args, **kwargs: x.long(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + )), + UnaryUfuncInfo( + 'short', + op=lambda x, *args, **kwargs: x.short(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + )), + UnaryUfuncInfo( + 'cdouble', + op=torch.Tensor.cdouble, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + # Error: Undefined type ComplexDouble + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + )), + UnaryUfuncInfo( + 'cfloat', + op=torch.Tensor.cfloat, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + skips=( + # autograd tests don't handle operators that change dtype + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients'), + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients'), + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # RuntimeError: attribute lookup is not defined on builtin + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + )), + UnaryUfuncInfo( + 'chalf', + op=lambda x, *args, **kwargs: x.chalf(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_conversion, + skips=( + # autograd tests don't handle operators that change dtype + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients'), + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients'), + # use of lambda doesn't work with test_normalize_operator_exhaustive + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # RuntimeError: "sum_cpu" not implemented for 'ComplexHalf' + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', + device_type='cpu'), + # TypeError: 'int' object is not iterable + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # RuntimeError: "sum_cpu" not implemented for 'ComplexHalf' + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view', + device_type='cpu'), + # RuntimeError: "sum_cpu" not implemented for 'ComplexHalf' + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view', + device_type='cpu'), + # RuntimeError: "sum_cpu" not implemented for 'ComplexHalf' + # RuntimeError: "neg_conj_cuda" not implemented for 'ComplexHalf' + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + ) + ), + OpInfo('empty_like', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_like_fns, + reference_inputs_func=reference_inputs_like_fns, + supports_autograd=False, + skips=( + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), + "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCudaFuserOpInfo'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_complex_half_reference_testing'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_non_standard_bool_values'), + DecorateInfo(unittest.skip("Expected: empty_like is not comparable"), 'TestCompositeCompliance', + 'test_operator'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + OpInfo('zeros_like', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_like_fns, + supports_autograd=False, + error_inputs_sparse_func=error_inputs_sparse_like_fns, + sample_inputs_sparse_coo_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_coo), + sample_inputs_sparse_csr_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_csr), + sample_inputs_sparse_csc_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_csc), + sample_inputs_sparse_bsr_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_bsr), + sample_inputs_sparse_bsc_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_bsc), + skips=( + )), + OpInfo('ones_like', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_like_fns, + supports_autograd=False, + skips=( + )), + OpInfo('randn', + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16, torch.complex32), + op=lambda *args, **kwargs: wrapper_set_seed(torch.randn, *args, **kwargs), + supports_out=True, + sample_inputs_func=sample_inputs_randn, + supports_autograd=False, + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.skip("Test expects tensor input"), "TestCommon", "test_noncontiguous_samples"), + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + # CPU randn generates different values based on the strides of out tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='cpu'), + # randn fails to warn when resizing its out tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # Tests that assume input tensor has a meaningful effect on output tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.expectedFailure, 'TestDecomp', 'test_quick'), + )), + OpInfo('randn_like', + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16, torch.complex32), + op=lambda inp, *args, **kwargs: + wrapper_set_seed(torch.randn_like, inp, *args, **kwargs), + supports_out=False, + sample_inputs_func=sample_inputs_like_fns, + supports_autograd=False, + error_inputs_sparse_func=error_inputs_sparse_like_fns, + sample_inputs_sparse_coo_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_coo), + sample_inputs_sparse_csr_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_csr), + sample_inputs_sparse_csc_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_csc), + sample_inputs_sparse_bsr_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_bsr), + sample_inputs_sparse_bsc_func=partial(sample_inputs_sparse_like_fns, layout=torch.sparse_bsc), + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("Expected: randn_like is not comparable between dtypes"), + 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + OpInfo('rand_like', + dtypes=floating_types_and(torch.half, torch.bfloat16, torch.complex32, torch.complex64, torch.complex128), + op=lambda inp, *args, **kwargs: + wrapper_set_seed(torch.randn_like, inp, *args, **kwargs), + supports_out=False, + sample_inputs_func=sample_inputs_like_fns, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("Expected: randn_like is not comparable between dtypes"), + 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + OpInfo('randint', + dtypes=all_types_and(torch.half, torch.bfloat16), + dtypesIfMPS=all_types_and(torch.half, torch.bfloat16, torch.complex64, torch.bool), + op=lambda *args, **kwargs: + wrapper_set_seed(torch.randint, *args, **kwargs), + supports_out=False, + sample_inputs_func=sample_inputs_randint, + supports_autograd=False, + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.skip("Test expects tensor input"), "TestCommon", "test_noncontiguous_samples"), + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + DecorateInfo(unittest.skip("Test expects tensor input"), "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + # CPU randint generates different values based on the strides of out tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # randint fails to warn when resizing its out tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # Tests that assume input tensor has a meaningful effect on output tensor + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Might need to skip until ROCm5.5 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_multiple_devices', + dtypes=[torch.float32, torch.int64], active_if=TEST_WITH_ROCM), + )), + OpInfo('randint_like', + dtypes=all_types_and(torch.half, torch.bfloat16), + op=lambda inp, *args, **kwargs: + wrapper_set_seed(torch.randint_like, inp, *args, **kwargs), + supports_out=False, + sample_inputs_func=sample_inputs_randint_like, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + OpInfo('full_like', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, + torch.uint16, torch.uint32), + supports_out=False, + sample_inputs_func=sample_inputs_full_like, + supports_autograd=False, + ), + OpInfo('new_zeros', + op=lambda x, *args, **kwargs: x.new_zeros(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_new_fns, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + ), + supports_autograd=False), + OpInfo('new_ones', + op=lambda x, *args, **kwargs: x.new_ones(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_new_fns, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + ), + supports_autograd=False), + OpInfo('ones', + op=torch.ones, + supports_autograd=False, + supports_varargs=True, + is_factory_function=True, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=True, + sample_inputs_func=sample_inputs_ones_zeros, + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + + # Same failure as arange: cannot find linspace in captured graph + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + )), + OpInfo('zeros', + op=torch.zeros, + supports_autograd=False, + is_factory_function=True, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=True, + sample_inputs_func=sample_inputs_ones_zeros, + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + + # Same failure as arange: cannot find linspace in captured graph + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + )), + OpInfo('full', + op=torch.full, + supports_autograd=False, + is_factory_function=True, + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=True, + sample_inputs_func=sample_inputs_full, + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # Same failure as arange: cannot find linspace in captured graph + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # RuntimeError: UNSUPPORTED DTYPE: bool + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness', dtypes=(torch.bool,)), + )), + OpInfo('new_empty', + op=lambda x, *args, **kwargs: x.new_empty(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_new_fns, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCudaFuserOpInfo'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_non_standard_bool_values'), + DecorateInfo(unittest.skip("Expected: new_empty is not comparable"), 'TestCompositeCompliance', + 'test_operator'), + DecorateInfo(unittest.skip("Expected: new_empty is not comparable"), + 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + ), + supports_autograd=False), + OpInfo('new_empty_strided', + op=lambda x, *args, **kwargs: x.new_empty_strided(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=partial(sample_inputs_new_fns, is_strided=True), + supports_autograd=False, + skips=( + # FX failed to normalize op + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # Lazy tensor failures + DecorateInfo(unittest.skip("Skipped!"), 'TestLazyOpInfo', 'test_correctness'), + DecorateInfo(unittest.skip("Skipped!"), 'TestLazyOpInfo', 'test_correctness_with_reusing_ir'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestCommon', 'test_noncontiguous_samples'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestMathBits', 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestCommon', 'test_non_standard_bool_values'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestCompositeCompliance', 'test_operator'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestDecomp', 'test_comprehensive'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestDecomp', 'test_quick'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestProxyTensorOpInfo', 'test_make_fx_exhaustive'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestProxyTensorOpInfo', 'test_make_fx_fake_exhaustive'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestProxyTensorOpInfo', 'test_make_fx_symbolic_exhaustive'), + DecorateInfo(unittest.skip("Expected: new_empty_strided is not comparable"), + 'TestNNCOpInfo', 'test_nnc_correctness'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + OpInfo('empty_strided', + op=lambda inp, *args, **kwargs: wrapper_set_seed(torch.empty_strided, inp, *args, **kwargs), + dtypes=all_types_and_complex_and(torch.bfloat16, torch.bool, torch.half), + supports_out=False, + supports_autograd=False, + sample_inputs_func=sample_inputs_empty_strided, + skips=( + # FX failed to normalize op - add the op to the op_skip list. + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_non_standard_bool_values'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), 'TestCompositeCompliance', 'test_operator'), + # Lazy tensor failures + DecorateInfo(unittest.skip("Expected: empty is not comparable"), 'TestLazyOpInfo'), + # RuntimeError: unsupported operation: more than one element of the written-to tensor refers to a single + # memory location. Please clone() the tensor before performing the operation. + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_meta_outplace'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_outplace'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_dispatch_symbolic_meta_outplace_all_strides'), + )), + OpInfo('empty', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + sample_inputs_func=sample_inputs_empty, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCudaFuserOpInfo'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_non_standard_bool_values'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), 'TestCompositeCompliance', + 'test_operator'), + # requires_grad doesn't exist in the jit schema + DecorateInfo(unittest.expectedFailure, 'TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestLazyOpInfo'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + OpInfo('eye', + dtypes=all_types_complex_float8_and(torch.bool, torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_eye, + error_inputs_func=error_inputs_eye, + supports_out=True, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # TODO: same as this? + # https://github.com/pytorch/pytorch/issues/81774 + # also see: arange, new_full + # fails to match any schemas despite working in the interpreter + DecorateInfo(unittest.expectedFailure, 'TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), + # fails to match any schemas despite working in the interpreter + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # skip these tests since we have non tensor input + DecorateInfo(unittest.skip('Skipped!'), "TestCommon", "test_noncontiguous_samples"), + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # "mul_cpu_reduced_float" not implemented for 'Float8_e4m3fn' + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness', + dtypes=(torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz)), + )), + OpInfo('empty_permuted', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + sample_inputs_func=sample_inputs_empty_permuted, + error_inputs_func=error_inputs_empty_permuted, + supports_out=False, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestNNCOpInfo', 'test_nnc_correctness'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCudaFuserOpInfo'), + # Empty tensor data is garbage so it's hard to make comparisons with it. + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_non_standard_bool_values'), + DecorateInfo(unittest.skip("Expected: empty_permuted is not comparable"), 'TestCompositeCompliance', + 'test_operator'), + # requires_grad doesn't exist in the jit schema + DecorateInfo(unittest.expectedFailure, 'TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), + DecorateInfo(unittest.skip("Expected: empty_permuted is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: empty_permuted is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: empty_permuted is not comparable"), + 'TestLazyOpInfo'), + DecorateInfo(unittest.skip("Expected: empty_permuted is not comparable"), + 'TestCommon', 'test_complex_half_reference_testing'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + )), + OpInfo('scalar_tensor', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + sample_inputs_func=sample_inputs_scalar_tensor, + supports_autograd=False, + supports_out=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # fails to match any schemas despite working in the interpreter + DecorateInfo(unittest.expectedFailure, 'TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), + # fails to match any schemas despite working in the interpreter + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # skip these tests since we have non tensor input + DecorateInfo(unittest.skip('Skipped!'), "TestCommon", "test_noncontiguous_samples"), + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + )), + OpInfo('new_full', + op=lambda x, *args, **kwargs: x.new_full(*args, **kwargs), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_out=False, + sample_inputs_func=sample_inputs_new_full, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + ), + supports_autograd=False), + OpInfo('multinomial', + op=lambda inp, *args, **kwargs: + wrapper_set_seed(torch.multinomial, inp, *args, **kwargs), + method_variant=lambda inp, *args, **kwargs: + wrapper_set_seed(torch.Tensor.multinomial, inp, *args, **kwargs), + dtypes=floating_types_and(torch.bfloat16, torch.half), + supports_out=True, + sample_inputs_func=sample_inputs_multinomial, + error_inputs_func=error_inputs_multinomial, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # Strides are not the same! + # This may not be reproducible in CI + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out'), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu')), + supports_autograd=False), + OpInfo('normal', + op=lambda inp, *args, **kwargs: + wrapper_set_seed(torch.normal, inp, *args, **kwargs), + # The inplace variant (Tensor.normal_) is different from torch.normal + inplace_variant=None, + dtypes=floating_types_and(torch.bfloat16, torch.half), + dtypesIfCUDA=floating_types_and(torch.bfloat16, torch.half), + supports_out=True, + sample_inputs_func=sample_inputs_normal_tensor_first, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # Tensor-likes are not close! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # Computed gradient is incorrect -- would be an exfail but gradgrad somehow passes + DecorateInfo(unittest.skip("Gradients are incorrect!"), 'TestFwdGradients'), + DecorateInfo(unittest.skip("Gradients are incorrect!"), 'TestBwdGradients'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + # RuntimeError: Difference from {dtype} is larger with decomposition + DecorateInfo(unittest.skip("Skipped!"), 'TestDecomp', 'test_comprehensive'), + DecorateInfo(unittest.skip("Skipped!"), 'TestDecomp', 'test_quick'), + # RuntimeError: normal: mean and std must have same number of elements + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=(torch.float32,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.float32,) + ), + # The inplace variant (Tensor.normal_) is different from torch.normal + # inplace variant Tensor.normal_ is decomposed using randn_like() + DecorateInfo(unittest.skip("Skipped!"), 'TestMeta', 'test_dispatch_symbolic_meta_outplace_all_strides'))), + OpInfo('normal', + # This has its own variant b/c OpInfos assume the first arg is a Tensor but it is not here + variant_test_name='number_mean', + op=lambda std, mean, *args, **kwargs: + wrapper_set_seed(torch.normal, mean, std, *args, **kwargs), + # The inplace variant (Tensor.normal_) is different from torch.normal + inplace_variant=None, + dtypes=floating_types_and(torch.bfloat16, torch.half), + dtypesIfCUDA=floating_types_and(torch.bfloat16, torch.half), + supports_out=True, + sample_inputs_func=sample_inputs_normal_tensor_second, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out_warning'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_backward'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBwdGradients'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.skip("Skipped!"), 'TestEagerFusionOpInfo'), + DecorateInfo(unittest.skip("Skipped!"), 'TestOperators'), + # AssertionError + DecorateInfo(unittest.skip("Skipped!"), 'TestDecomp', 'test_comprehensive'), + # AssertionError + DecorateInfo(unittest.skip("Skipped!"), 'TestDecomp', 'test_quick'), + # AssertionError in CUDA variant + DecorateInfo(unittest.skip("Skipped!"), 'TestFakeTensor', device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestDeviceUtils', 'test_device_mode_ops'))), + OpInfo('bernoulli', + op=lambda inp, *args, **kwargs: + wrapper_set_seed(torch.bernoulli, inp, *args, **kwargs), + # The inplace variant (Tensor.bernoulli_) is different from torch.bernoulli + inplace_variant=None, + method_variant=lambda inp, *args, **kwargs: + wrapper_set_seed(torch.Tensor.bernoulli, inp, *args, **kwargs), + dtypes=floating_types_and(torch.bfloat16, torch.half), + dtypesIfMPS=all_types_and(torch.bfloat16, torch.half, torch.bool), + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_bernoulli, + error_inputs_func=error_inputs_bernoulli, + skips=( + # vmap: We do not yet support calling random operations inside of vmap + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_forward_mode_AD'), + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Expected RuntimeError when doing an unsafe cast from a result of + # dtype torch.float32 into an out= with dtype torch.lon + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # NotImplementedError: "bernoulli_tensor_cpu_p_" not implemented for * + DecorateInfo( + unittest.expectedFailure, 'TestConsistency', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64) + ), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'))), + OpInfo('scatter_add', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + inplace_variant=torch.Tensor.scatter_add_, + sample_inputs_func=sample_inputs_scatter_add, + error_inputs_func=error_inputs_scatter_and_scatter_add, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Compiler issue on ROCm. Regression started in ROCm 6.4. + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + )), + OpInfo('stack', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_stack, + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # https://github.com/pytorch/pytorch/issues/77046 + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + )), + OpInfo('_chunk_cat', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_chunk_cat, + error_inputs_func=error_inputs_chunk_cat, + supports_autograd=False, + supports_out=True, + ), + OpInfo('hstack', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_hstack_dstack_vstack, + error_inputs_func=error_inputs_hstack_dstack_vstack, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ), + BinaryUfuncInfo('hypot', + dtypes=floating_types_and(torch.bfloat16, torch.half), + dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_rhs_python_scalar=False), + OpInfo('histogram', + dtypes=floating_types(), + dtypesIfCUDA=_dispatch_dtypes(), # histogram is only implemented on CPU + sample_inputs_func=sample_inputs_histogram, + supports_autograd=False, + skips=( + # JIT tests don't work with Tensor keyword arguments + # https://github.com/pytorch/pytorch/issues/58507 + # RuntimeError: + # undefined value tensor: + # File "", line 3 + # def the_method(i0): + # return torch.histogram(i0, 1, weight=tensor(-0.5735, dtype=torch.float32), density=False) + # ~~~~~~ <--- HERE + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Not Implemented on XLA. + DecorateInfo(unittest.skip("Skipped!"), 'TestOpInfo', device_type='xla'), + )), + OpInfo('histogramdd', + dtypes=floating_types(), + dtypesIfCUDA=_dispatch_dtypes(), # histogramdd is only implemented on CPU + sample_inputs_func=sample_inputs_histogramdd, + error_inputs_func=error_inputs_histogramdd, + supports_autograd=False, + skips=( + # Not implemented on CUDA + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='cuda'), + # JIT tests don't work with Tensor keyword arguments + # https://github.com/pytorch/pytorch/issues/58507 + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + )), + OpInfo('histc', + dtypes=floating_types_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_types_and(torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64), + sample_inputs_func=sample_inputs_histc, + supports_out=True, + supports_autograd=False, + skips=( + # CUDA histc returns a float tensor but does not correctly warn when passed an integral out tensor + # "AssertionError: RuntimeError not raised : Expected RuntimeError when doing an unsafe cast + # from a result of dtype torch.float32 into an out= with dtype torch.long" + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='cuda'), + )), + OpInfo('bincount', + dtypes=integral_types_and(), + sample_inputs_func=sample_inputs_bincount, + supports_out=False, + supports_autograd=False, + skips=( + # JIT tests don't work with Tensor keyword arguments + # https://github.com/pytorch/pytorch/issues/58507 + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + )), + OpInfo('bucketize', + dtypes=all_types_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bfloat16, torch.float16), + sample_inputs_func=sample_inputs_bucketize, + reference_inputs_func=reference_inputs_bucketize, + error_inputs_func=error_inputs_bucketize, + supports_autograd=False, + skips=( + # JIT tests don't work with Tensor keyword arguments + DecorateInfo(unittest.skip("Expected failure!"), 'TestJit', 'test_variant_consistency_jit'), + # AssertionError: RuntimeError not raised : Expected RuntimeError when calling with + # input.device=mps:0 and out.device=cpu. + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_out', + device_type='mps', dtypes=(torch.float32,) + ), + )), + OpInfo('searchsorted', + dtypes=all_types_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=all_types_and(torch.bfloat16, torch.float16), + sample_inputs_func=sample_inputs_searchsorted, + supports_autograd=False, + ref=reference_searchsorted, + skips=( + # JIT tests don't work with Tensor keyword arguments + # https://github.com/pytorch/pytorch/issues/58507 + DecorateInfo(unittest.skip("Expected failure!"), 'TestJit', 'test_variant_consistency_jit'), + # AssertionError: RuntimeError not raised : Expected RuntimeError when calling with + # input.device=mps:0 and out.device=cpu. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + )), + OpInfo('cat', + ref=_cat_np, + aliases=('concat', 'concatenate'), + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.complex32), + sample_inputs_func=sample_inputs_cat_concat, + reference_inputs_func=reference_inputs_cat, + error_inputs_func=error_inputs_cat, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + assert_autodiffed=True, + skips=( + # RuntimeError: Arguments for call not valid. + # Expected a value of type 'List[Tensor]' for argument + # 'tensors' but instead found type 'Tensor (inferred)'. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_jit_alias_remapping'), + # see https://github.com/pytorch/pytorch/issues/71286 + DecorateInfo(unittest.expectedFailure, 'TestNNCOpInfo', 'test_nnc_correctness'), + # see https://github.com/pytorch/pytorch/issues/99806 + # RuntimeError: The size of tensor a (25) must match the size of tensor b (0) at non-singleton dimension 0. + DecorateInfo(unittest.expectedFailure, 'TestBwdGradients', 'test_fn_gradgrad'), + )), + OpInfo('unbind', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + ref=reference_unbind, + sample_inputs_func=sample_inputs_unbind, + error_inputs_func=error_inputs_unbind, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_gradgrad=True, + supports_out=False, + ), + OpInfo('unbind_copy', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + ref=reference_unbind, + sample_inputs_func=sample_inputs_unbind, + error_inputs_func=error_inputs_unbind, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_gradgrad=True, + supports_out=True, + check_batched_grad=False, + ), + OpInfo('vstack', + aliases=('row_stack',), + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_hstack_dstack_vstack, + error_inputs_func=error_inputs_hstack_dstack_vstack, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # RuntimeError: _fn() Expected a value of type + # 'Tensor (inferred)' for argument 't0' but instead found type 'tuple'. + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_jit_alias_remapping'),)), + OpInfo('dstack', + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_hstack_dstack_vstack, + error_inputs_func=error_inputs_hstack_dstack_vstack, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + ), + OpInfo('unfold', + op=lambda x, *args: x.unfold(*args), + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + backward_dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_gradgrad=False, + # See https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # Skip operator schema test because this is a functional and not an operator + DecorateInfo(unittest.expectedFailure, 'TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), + ), + sample_inputs_func=sample_inputs_unfold), + OpInfo('unfold_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + backward_dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_gradgrad=False, + # See https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_unfold), + OpInfo('msort', + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bool, torch.float16, torch.bfloat16), + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_msort), + OpInfo('movedim', + aliases=('moveaxis',), + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_movedim_moveaxis, + reference_inputs_func=reference_movedim_moveaxis, + error_inputs_func=error_movedim_moveaxis), + OpInfo('renorm', + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_renorm, + error_inputs_func=error_inputs_renorm, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # RuntimeError: Difference from float64 is larger with decomposition + # linalg_vector_norm.default than original on output 0. + # Original max diff: 2.560596747969157e-07, + # Decomp max diff: 1.8187482915266173e-06 + DecorateInfo(unittest.skip("Inconsistent accuracy"), 'TestDecomp', 'test_comprehensive', + device_type='cpu', dtypes=(torch.float16,)), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=3e-4, rtol=3e-6)}), + "TestConsistency", "test_output_match", device_type="mps"), + # RuntimeError: Failed to create function state object for: renorm_float2 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), + )), + ShapeFuncInfo('repeat', + op=lambda x, dims: x.repeat(dims), + ref=np.tile, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_repeat_tile, + skips=( + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + )), + OpInfo('squeeze', + ref=_squeeze_ref, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + assert_autodiffed=True, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + assert_jit_shape_analysis=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_squeeze), + OpInfo('squeeze', + ref=_squeeze_ref, + variant_test_name="multiple", + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + assert_autodiffed=True, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_squeeze_multiple), + OpInfo('squeeze_copy', + ref=_squeeze_ref, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=True, + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_squeeze, + skips=( + DecorateInfo( + unittest.expectedFailure, + 'TestJit', + 'test_variant_consistency_jit', + dtypes=(torch.float32,), + ), + )), + UnaryUfuncInfo( + 'fill', + ref=_fill_np, + method_variant=None, + sample_kwargs=_fill_sample_kwargs, + sample_inputs_func=partial(sample_inputs_elementwise_unary, op_kwargs={'value': True}), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + dtypes=all_types_and_complex_and(torch.complex32, torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + skips=( + # JIT has issue when op is passed as lambda + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip("No fill_ op"), 'TestCudaFuserOpInfo'), + DecorateInfo(unittest.skip("No fill_ op"), 'TestNNCOpInfo'), + )), + OpInfo('resize_', + op=lambda x, shape: x.clone().resize_(shape), + method_variant=None, + inplace_variant=torch.Tensor.resize_, + # the test fails because resize_ doesn't work with imag views as expected by the test + # https://github.com/pytorch/pytorch/issues/65945 + test_neg_view=False, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_autograd=False, + skips=( + # Cannot resize variables that require grad + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes'), + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.skip("Allowed exception"), 'TestCompositeCompliance', 'test_operator'), + ), + sample_inputs_func=sample_inputs_resize_ops), + OpInfo('resize_as_', + op=lambda x, other: torch.resize_as_(x.clone(), other), + method_variant=None, + inplace_variant=torch.Tensor.resize_as_, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_autograd=False, + skips=( + # Cannot resize variables that require grad + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes'), + DecorateInfo(unittest.skip('Allowed exemption'), 'TestCompositeCompliance', 'test_operator'), + ), + sample_inputs_func=sample_inputs_resize_ops), + OpInfo('take_along_dim', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_inplace_autograd=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_take_along_dim, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + decorators=( + # RuntimeError: view size is not compatible with input tensor's size and stride + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides"), + )), + ShapeFuncInfo('tile', + ref=np.tile, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_repeat_tile), + OpInfo('trapz', # TODO: in the future, 'trapz' should be made a proper alias of 'trapezoid' + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + decorators=[ + DecorateInfo( + toleranceOverride({torch.half: tol(atol=9e-4, rtol=4.3e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cuda' + ), + ], + sample_inputs_func=sample_trapezoid), + OpInfo('trapezoid', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + decorators=[ + DecorateInfo( + toleranceOverride({torch.half: tol(atol=9e-4, rtol=4.3e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cuda' + ), + ], + sample_inputs_func=sample_trapezoid), + OpInfo('cumulative_trapezoid', + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + supports_out=False, + decorators=( + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=4e-3, rtol=4e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', + ), + ), + sample_inputs_func=sample_cumulative_trapezoid,), + OpInfo('unsqueeze', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + assert_jit_shape_analysis=True, + assert_autodiffed=True, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + sample_inputs_func=sample_unsqueeze), + OpInfo('unsqueeze_copy', + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + assert_jit_shape_analysis=True, + assert_autodiffed=True, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + sample_inputs_func=sample_unsqueeze, + skips=( + DecorateInfo( + unittest.expectedFailure, + 'TestJit', + 'test_variant_consistency_jit', + dtypes=(torch.float32,), + ), + )), + BinaryUfuncInfo('xlogy', + aliases=('special.xlogy',), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + promotes_int_to_float=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_one_python_scalar=True, + # We don't test 0 as the gradient will be NaN and it'll break + rhs_make_tensor_kwargs=dict(low=0.01)), + OpInfo('zero_', + op=lambda x: torch.zero_(x.clone()), + method_variant=None, + inplace_variant=torch.Tensor.zero_, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_gradgrad=True, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + ), + sample_inputs_func=sample_inputs_zero_), + OpInfo('logsumexp', + aliases=('special.logsumexp',), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_fast_mode=False, + sample_inputs_func=sample_inputs_logsumexp, + reference_inputs_func=reference_inputs_logsumexp), + OpInfo('trace', + dtypes=all_types_and_complex(), + dtypesIfMPS=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.chalf, torch.bool, torch.half, torch.bfloat16), + error_inputs_func=error_inputs_trace, + supports_inplace_autograd=False, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_trace), + OpInfo('transpose', + ref=_numpy_ref_transpose, + aliases=('swapdims', 'swapaxes'), + assert_jit_shape_analysis=True, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + sample_inputs_func=sample_inputs_transpose_swapdims), + OpInfo('transpose_copy', + assert_jit_shape_analysis=True, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf), + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + sample_inputs_func=sample_inputs_transpose_swapdims, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestDTensorOps', 'test_dtensor_op_db'), + DecorateInfo( + unittest.expectedFailure, + 'TestJit', + 'test_variant_consistency_jit', + dtypes=(torch.float32,) + ), + )), + OpInfo('T', + op=lambda x: x.T, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, "TestJit", "test_variant_consistency_jit"),), + sample_inputs_func=sample_inputs_T, + error_inputs_func=error_inputs_T), + OpInfo('H', + op=lambda x: x.H, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, "TestJit", "test_variant_consistency_jit"),), + sample_inputs_func=sample_inputs_T), + OpInfo('mT', + op=lambda x: x.mT, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, "TestJit", "test_variant_consistency_jit"),), + sample_inputs_func=sample_inputs_adjoint), + OpInfo('mH', + op=lambda x: x.mH, + aliases=('adjoint',), + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half, torch.chalf), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, "TestJit", "test_variant_consistency_jit"),), + sample_inputs_func=sample_inputs_adjoint), + OpInfo('tril', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + error_inputs_func=error_inputs_tril_triu, + sample_inputs_func=sample_inputs_tril_triu, + skips=( + # Compiler issue on ROCm. Regression started in ROCm 6.4. + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + )), + OpInfo('triu', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + error_inputs_func=error_inputs_tril_triu, + sample_inputs_func=sample_inputs_tril_triu, + skips=( + # Compiler issue on ROCm. Regression started in ROCm 6.4. + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + )), + OpInfo('triu_indices', + dtypes=_dispatch_dtypes((torch.int32, torch.int64)), + sample_inputs_func=sample_inputs_trilu_indices, + ref=lambda h, w, ofs=0, dtype=torch.long, device='cpu' : np.array(np.triu_indices(h, ofs, w), dtype=dtype), + supports_out=False, + supports_autograd=False, + skips=( + # skip these tests since we have non tensor input + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_noncontiguous_samples'), + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip('Skipped!'), 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Skipped!'), 'TestMathBits', 'test_neg_view'), + )), + OpInfo('tril_indices', + dtypes=_dispatch_dtypes((torch.int32, torch.int64)), + sample_inputs_func=sample_inputs_trilu_indices, + ref=lambda h, w, ofs=0, dtype=torch.long, device='cpu' : np.array(np.tril_indices(h, ofs, w), dtype=dtype), + supports_out=False, + supports_autograd=False, + skips=( + # skip these tests since we have non tensor input + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_noncontiguous_samples'), + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip('Skipped!'), 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Skipped!'), 'TestMathBits', 'test_neg_view'), + )), + OpInfo('kron', + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_inplace_autograd=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_kron, + decorators=( + # RuntimeError: view size is not compatible with input tensor's size and stride + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides"), + )), + OpInfo('inner', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfROCM=floating_and_complex_types_and(torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_inner, + ), + OpInfo('tensordot', + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + dtypesIfROCM=floating_and_complex_types_and(torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_tensordot, + skips=( + # Skip operator schema test because this is a functional and not an operator. + # Reference: https://github.com/pytorch/pytorch/issues/54574 + DecorateInfo(unittest.skip("Skipped!"), 'TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), + ) + ), + OpInfo('to_sparse', + op=lambda x, *args: x.to_sparse(*args), + sample_inputs_func=sample_inputs_to_sparse, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + backward_dtypes=floating_types(), + backward_dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_sparse_csr=True, + supports_sparse_csc=True, + check_batched_grad=False, + check_batched_gradgrad=False, + skips=( + # NotImplementedError: Could not run 'aten::normal_' with arguments from the 'SparseCPU' backend + DecorateInfo(unittest.skip(""), 'TestCommon', 'test_noncontiguous_samples'), + # TODO: FIXME: complex inputs requiring grad error in forward + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_dtypes'), + # lambda impl + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # Allowed exception: sparse tensors don't have strides + DecorateInfo(unittest.skip("Allowed exception"), 'TestCompositeCompliance', 'test_operator'), + DecorateInfo(unittest.skip("Allowed exception"), 'TestCompositeCompliance', 'test_backward'), + DecorateInfo(unittest.skip("Allowed exception"), 'TestTags', 'test_tags'), + # TODO: implement csr.to_sparse(sample_dim) where sampled_dim is 1. + DecorateInfo(unittest.skip("csr.to_sparse(1) not implemented. Skipped!"), + 'TestSparseCSR', 'test_sparse_csr_consistency'), + # Compiler issue on ROCm. Might need to skip until ROCm5.5 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + ) + ), + OpInfo('logcumsumexp', + dtypes=floating_and_complex_types_and(torch.bfloat16, torch.half), + backward_dtypes=floating_and_complex_types_and(torch.bfloat16, torch.half), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='cuda'), + # RuntimeError: "max_values_cpu" not implemented for 'ComplexDouble' + # Falling back to non-numerically stabilized exp, causing nan in the results. + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_forward_mode_AD', dtypes=[torch.complex128]), + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_fn_fwgrad_bwgrad', dtypes=[torch.complex128]), + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=7e-5, rtol=6e-3), + }), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda" + ), + # AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ), + sample_inputs_func=sample_inputs_logcumsumexp, + error_inputs_func=error_inputs_logcumsumexp), + UnaryUfuncInfo('sigmoid', + aliases=('special.expit', 'nn.functional.sigmoid'), + aten_backward_name='sigmoid_backward', + ref=reference_sigmoid if TEST_SCIPY else None, + decorators=(precisionOverride({torch.float16: 1e-2, + torch.complex64: 1e-1, + torch.bfloat16: 1e-2}),), + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/56012 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.complex64, torch.cdouble], device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=[torch.chalf, torch.complex64, torch.cdouble], device_type='cuda')), + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.complex32, torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + assert_autodiffed=True, + # sigmoid(z) = 1 / (1 + exp(-z)), at z = j * pi * odd_number, the denominator is zero + reference_numerics_filter=NumericsFilter( + condition=lambda x: (close_to_int(x / (math.pi * 1j)) + if x.is_complex() else x.new_tensor(False, dtype=torch.bool)), + safe_val=0)), + UnaryUfuncInfo('digamma', + ref=scipy.special.digamma if TEST_SCIPY else None, + aliases=('special.psi', 'special.digamma',), + decorators=(precisionOverride({torch.float16: 5e-1}),), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # AssertionError: Tensor-likes are not close! + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-5, rtol=3.1e-6)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ), + promotes_int_to_float=True), + UnaryUfuncInfo('erf', + ref=scipy.special.erf if TEST_SCIPY else None, + aliases=('special.erf', ), + decorators=(precisionOverride({torch.float16: 1e-2, + torch.bfloat16: 1e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped! sparse backward not supported"), + 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + + ), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + assert_autodiffed=True, + assert_jit_shape_analysis=True, + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True), + UnaryUfuncInfo('erfc', + ref=scipy.special.erfc if TEST_SCIPY else None, + aliases=('special.erfc', ), + decorators=(precisionOverride({torch.float16: 1e-2, + torch.bfloat16: 1e-2}),), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + assert_autodiffed=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True), + UnaryUfuncInfo('erfinv', + ref=scipy.special.erfinv if TEST_SCIPY else None, + aliases=('special.erfinv', ), + decorators=(precisionOverride({torch.float16: 1e-2, + torch.bfloat16: 1e-2, + torch.float32: 1e-4}),), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_sparse=True, + supports_sparse_csr=True, + supports_sparse_csc=True, + supports_sparse_bsr=True, + supports_sparse_bsc=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + domain=(-1, 1), + skips=( + # Reference: https://github.com/pytorch/pytorch/pull/49155#issuecomment-742664611 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + active_if=TEST_SCIPY and version.parse(scipy.__version__) < version.parse("1.4.0")), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + active_if=TEST_SCIPY and version.parse(scipy.__version__) < version.parse("1.4.0")), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_small', + active_if=TEST_SCIPY and version.parse(scipy.__version__) < version.parse("1.4.0")), + DecorateInfo(unittest.expectedFailure, 'TestSparseUnaryUfuncs', 'test_sparse_fn_grad'), + )), + OpInfo("nn.functional.smooth_l1_loss", + ref=reference_smooth_l1_loss, + sample_inputs_func=sample_inputs_smooth_l1_loss, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + backward_dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # RuntimeError: input->type()->kind() == TypeKind::OptionalTypeINTERNAL ASSERT FAILED + # at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":270, please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, "TestJit", "test_variant_consistency_jit"),)), + OpInfo( + "nn.functional.l1_loss", + ref=loss_reference_reduction_wrapper(lambda input, target: np.abs(input - target)), + sample_inputs_func=sample_inputs_l1_loss, + error_inputs_func=error_inputs_l1_loss, + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # RuntimeError: input->type()->kind() == TypeKind::OptionalTypeINTERNAL ASSERT FAILED + # at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":270, please report a bug to PyTorch. + DecorateInfo( + unittest.expectedFailure, + "TestJit", + "test_variant_consistency_jit", + dtypes=(torch.float32,), + ), + ), + ), + UnaryUfuncInfo('lgamma', + ref=reference_lgamma if TEST_SCIPY else None, + aliases=('special.gammaln', ), + decorators=(precisionOverride({torch.float16: 7e-1}), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=2e-4, rtol=4e-6)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + skips=( + # Reference: https://github.com/pytorch/pytorch/pull/50140#issuecomment-756150214 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.float32, torch.float64], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=[torch.float32, torch.float64], active_if=IS_WINDOWS), + ), + # lgamma have multiple singularities at x <= 0 + reference_numerics_filter=NumericsFilter(condition=lambda x: x < 0.1, safe_val=1)), + OpInfo( + 'logdet', + dtypes=floating_and_complex_types(), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_linalg_det_logdet_slogdet, + skips=( + # Exception: linalg.lu_factor(): MPS doesn't support complex types. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), + ), + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack]), + # `log_softmax` supports different dtypes based on whether `dtype` argument, + # is passed or not. Hence two OpInfo entries, one with dtype and other without. + OpInfo( + 'log_softmax', + aliases=('special.log_softmax', 'nn.functional.log_softmax'), + supports_out=True, + aten_backward_name='_log_softmax_backward_data', + dtypes=floating_types_and(torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_softmax_variant, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # The following dtypes worked in forward but are not listed by the + # OpInfo: {torch.int16, torch.int8, torch.uint8, torch.int32}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + ), + assert_autodiffed=True), + OpInfo( + 'log_softmax', + variant_test_name='with_dtype', + aliases=('special.log_softmax', 'nn.functional.log_softmax'), + supports_out=True, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + sample_inputs_func=partial(sample_inputs_softmax_variant, with_dtype=True), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True), + UnaryUfuncInfo('logit', + aten_backward_name='logit_backward', + ref=scipy.special.logit if TEST_SCIPY else None, + domain=(0, 1), + aliases=('special.logit', ), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + decorators=(precisionOverride({torch.bfloat16: 5e-1, + torch.float16: 5e-1}),), + skips=( + # AssertionError: RuntimeError not raised : Expected RuntimeError when doing an unsafe cast from a result + # of dtype torch.float32 into an out= with dtype torch.long + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + # AssertionError: UserWarning not triggered : Resized a non-empty tensor but did not warn about it. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + # NotImplementedError: The operator 'aten::logit_' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + ), + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_logit), + OpInfo('where', + # Currently only the `input` is tested in gradcheck. + # If we pass `condition` first, none of the input which supports + # autograd will be tested. Hence the following lambda. + op=lambda self, condition, other, **kwargs: torch.where(condition, self, other, **kwargs), + ref=lambda self, condition, other: np.where(condition, self, other), + sample_inputs_func=sample_inputs_where, + reference_inputs_func=reference_inputs_where, + error_inputs_func=error_inputs_where, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=( + DecorateInfo(onlyCUDA, "TestCommon", 'test_errors'),), + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + ), + dtypes=all_types_and_complex_and(torch.bool, torch.half, torch.bfloat16, torch.chalf)), + OpInfo('nonzero', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + sample_inputs_func=sample_inputs_nonzero, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # nonzero(): argument 'out' must be Tensor, not tuple + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # https://github.com/pytorch/pytorch/issues/67458 + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # nonzero is not raising a warning when the out is resized + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + # Can't find schemas for this operator for some reason + DecorateInfo(unittest.expectedFailure, 'TestOperatorSignatures', 'test_get_torch_func_signature_exhaustive'), + # Compiler issue on ROCm. Might need to skip until ROCm5.5 + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + )), + OpInfo('nonzero_static', + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16, torch.chalf), + sample_inputs_func=sample_inputs_nonzero_static, + supports_out=False, + supports_autograd=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning'), + DecorateInfo(unittest.expectedFailure, 'TestDTensorOps', 'test_dtensor_op_db'), + DecorateInfo(unittest.expectedFailure, 'TestInductorOpInfo', 'test_comprehensive'), + DecorateInfo(unittest.expectedFailure, 'TestVmapOperatorsOpInfo', 'test_op_has_batch_rule'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + )), + # Following tests are for jiterator's python interface + # Jiterator can be used to author elementwise CUDA kernel + # jiterator._create_jit_fn returns a callable that behaves like a regular pytorch op + # See create_jit_fn in jiterator.py for more information + UnaryUfuncInfo( + 'jiterator_unary', + op=torch.cuda.jiterator._create_jit_fn("template T unary(T x) { return x * x + x; }"), + ref=lambda x: x * x + x, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.bool), + supports_out=False, + supports_autograd=False, # jiterator ops doesn't have backward defined + decorators=[ + onlyCUDA, + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-02)}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal'), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-02)}), + 'TestUnaryUfuncs', 'test_reference_numerics_hard'), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-02)}), + 'TestUnaryUfuncs', 'test_reference_numerics_normal'), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-02)}), + 'TestUnaryUfuncs', 'test_reference_numerics_small'), + ], + skips=( + # Jiterator ops doesn't support neg or conj view + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # Jiterator ops doesn't support CompositeCompliantTensor + # Following test should expectedFailure, but it's causing cascading failures in CUDA, thus skipped + DecorateInfo(unittest.skip("skip"), 'TestCompositeCompliance', 'test_operator'), + # Skip reference_numerics tests for bool type, as the defined function doesn't work for bool + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + dtypes=[torch.bool]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_hard', + dtypes=[torch.bool]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_normal', + dtypes=[torch.bool]), + # ROCm generates -inf+infj instead of nan+infj for complex64 for some of the results + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=[torch.complex64], active_if=TEST_WITH_ROCM), + # Newer numpy generates -inf+infj instead of nan+infj for complex64 for some of the results + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_large', + dtypes=[torch.complex64], device_type='cuda'), + # Expected failure: torch.jiterator_unary is not a valid op + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Skip Nvfuser + DecorateInfo(unittest.skip('Skipped!'), 'TestCudaFuserOpInfo'), + ) + ), + BinaryUfuncInfo( + 'jiterator_binary', + op=torch.cuda.jiterator._create_jit_fn( + "template T binary(T x, T y, T alpha) { return x + alpha * y; }", alpha=1), + ref=lambda input, other, *, alpha=1: ( + np.add(input, other) + if alpha == 1 + else np.add(input, np.multiply(alpha, other)) + ), + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.bool), + sample_inputs_func=partial(sample_inputs_jiterator, num_inputs=2, alpha=-3.14), + supports_out=False, + supports_autograd=False, # jiterator ops doesn't have backward defined + supports_rhs_python_scalar=False, + decorators=[onlyCUDA], + skips=( + # Jiterator ops doesn't support neg or conj view + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # Jiterator ops doesn't support CompositeCompliantTensor + # Following test should expectedFailure, but it's causing cascading failures in CUDA, thus skipped + DecorateInfo(unittest.skip("skip"), 'TestCompositeCompliance', 'test_operator'), + # Expected failure: torch.jiterator_binary is not a valid op + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Skip Nvfuser + DecorateInfo(unittest.skip('Skipped!'), 'TestCudaFuserOpInfo'), + ) + ), + OpInfo( + 'jiterator_4inputs_with_extra_args', + op=torch.cuda.jiterator._create_jit_fn( + "template T binary(T i0, T i1, T i2, T i3, T alpha, T beta) { return alpha * i0 + beta * i1 + i2 + i3; }", + alpha=1, beta=1), + ref=lambda i0, i1, i2, i3, *, alpha=1, beta=1: alpha * i0 + beta * i1 + i2 + i3, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.bool), + sample_inputs_func=partial(sample_inputs_jiterator, num_inputs=4, alpha=3.14, beta=-4.20), + supports_out=False, + supports_autograd=False, # jiterator ops doesn't have backward defined + decorators=[onlyCUDA], + skips=( + # Jiterator ops doesn't support neg or conj view + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # Jiterator ops doesn't support CompositeCompliantTensor + # Following test should expectedFailure, but it's causing cascading failures in CUDA, thus skipped + DecorateInfo(unittest.skip("skip"), 'TestCompositeCompliance', 'test_operator'), + # Expected failure: torch.jiterator_4inputs_with_extra_args is not a valid op + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Skip Nvfuser + DecorateInfo(unittest.skip('Skipped!'), 'TestCudaFuserOpInfo'), + ) + ), + BinaryUfuncInfo( + 'jiterator_binary_return_by_ref', + op=torch.cuda.jiterator._create_multi_output_jit_fn( + """ + template + void binary_return_by_ref(T i0, T i1, T& out0) { + out0 = i0 + i1; + } + """, + num_outputs=1), + ref=operator.add, + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.bool), + sample_inputs_func=partial(sample_inputs_jiterator, num_inputs=2, alpha=-0.42), + supports_out=False, + supports_autograd=False, # jiterator ops doesn't have backward defined + supports_rhs_python_scalar=False, + decorators=[onlyCUDA], + skips=( + # Jiterator ops doesn't support neg or conj view + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # Jiterator ops doesn't support CompositeCompliantTensor + # Following test should expectedFailure, but it's causing cascading failures in CUDA, thus skipped + DecorateInfo(unittest.skip("skip"), 'TestCompositeCompliance', 'test_operator'), + # Expected failure: torch.jiterator_4inputs_with_extra_args is not a valid op + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Skip Nvfuser + DecorateInfo(unittest.skip('Skipped!'), 'TestCudaFuserOpInfo'), + ) + ), + OpInfo( + 'jiterator_2inputs_2outputs', + op=torch.cuda.jiterator._create_multi_output_jit_fn( + """ + template + void binary_2outputs(T i0, T i1, T& out0, T& out1) { + out0 = i0 + i1; + out1 = i0 - i1; + } + """, + num_outputs=2), + ref=lambda i0, i1, *, alpha=1: (i0 + i1, i0 - i1), + dtypes=all_types_and_complex_and(torch.bfloat16, torch.float16, torch.bool), + sample_inputs_func=partial(sample_inputs_jiterator, num_inputs=2), + supports_out=False, + supports_autograd=False, # jiterator ops doesn't have backward defined + decorators=[onlyCUDA], + skips=( + # Jiterator ops doesn't support neg or conj view + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # Jiterator ops doesn't support CompositeCompliantTensor + # Following test should expectedFailure, but it's causing cascading failures in CUDA, thus skipped + DecorateInfo(unittest.skip("skip"), 'TestCompositeCompliance', 'test_operator'), + # Expected failure: torch.jiterator_4inputs_with_extra_args is not a valid op + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # Skip Nvfuser + DecorateInfo(unittest.skip('Skipped!'), 'TestCudaFuserOpInfo'), + ) + ), + # `torch.norm` has multiple code paths depending on the value of `p`. + # These paths have different dtype support. Also JIT supports, + # most variants but not all of them. So we split the OpInfo entries, + # for `norm` based on the code-paths and JIT support. + OpInfo( + "norm", + sample_inputs_func=sample_inputs_norm, + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16, torch.chalf), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + # TODO Benchmark again with the new implementation + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + check_batched_forward_grad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Dispatches in Python to vector_norm. Not sure how to make this test happy + # Happens to pass on complex64. Also a mystery + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.float32,)), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=0.2, rtol=0.002)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ) + ), + OpInfo('norm', + variant_test_name='nuc', + sample_inputs_func=sample_inputs_norm_nuc, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + check_batched_gradgrad=False, + # torch.autograd.gradcheck.GradcheckError: While computing batched gradients + # got: Could not allocate memory to change Tensor SizesAndStrides! + check_batched_forward_grad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_and_complex_types(), + dtypesIfCUDA=floating_and_complex_types(), + skips=( + # NotImplementedError: The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + # Dispatches in Python to matrix_norm. Not sure how to make this test happy + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.complex64, torch.float32,)),) + ), + OpInfo('norm', + variant_test_name='fro', + sample_inputs_func=sample_inputs_norm_fro, + dtypes=floating_and_complex_types_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + supports_forward_ad=True, + # torch.autograd.gradcheck.GradcheckError: While computing batched gradients + # got: Could not allocate memory to change Tensor SizesAndStrides! + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + skips=( + # MPS has some mild accuracy issues for float16. We divide the tolerances by 10 + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-4, rtol=0.01)}), + 'TestConsistency', + 'test_output_match', + + ), + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + 'TestSchemaCheckModeOpInfo', + 'test_schema_correctness', + dtypes=(torch.complex64, torch.complex128)), + # Dispatches in Python to vector_norm. Not sure how to make this test happy + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.complex64, torch.float32,)),) + ), + OpInfo( + "norm", + variant_test_name="inf", + sample_inputs_func=sample_inputs_norm_inf, + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16, torch.chalf), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16), + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + # fast gradcheck produces NaNs + gradcheck_fast_mode=False, + skips=( + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=2e-3, rtol=1e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type='cuda', + ), + # Dispatches in Python to vector_norm. Not sure how to make this test happy + # Happens to pass on complex64. Also a mystery + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit', + dtypes=(torch.float32,)) + ), + ), + OpInfo('t', + sample_inputs_func=sample_inputs_t, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + assert_autodiffed=True, + error_inputs_func=error_inputs_t), + OpInfo('t_copy', + sample_inputs_func=sample_inputs_t, + supports_out=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + # vmap does not support inplace views + check_inplace_batched_forward_grad=False, + autodiff_fusible_nodes=[], # aliases inputs, shouldn't be fused + autodiff_nonfusible_nodes=[], # aliases inputs, shouldn't be fused + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + assert_autodiffed=True, + error_inputs_func=error_inputs_t), + OpInfo( + "nn.functional.dropout", + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.dropout, input, *args, **kwargs), + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.complex64), + skips=( + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # Probably because we have used lambda for the op here + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # inplace variant dispatches to dropout kernel, while on CUDA + # the op dispatches to _fused_dropout (with a few more conditions) + # hence, different values and this skip here + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view', device_type='cuda'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu')), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + supports_out=False, + sample_inputs_func=sample_inputs_dropout, + inplace_variant=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.dropout, input, *args, **kwargs, inplace=True)), + OpInfo( + "native_dropout_backward", + op=torch.ops.aten.native_dropout_backward.default, + aten_name="native_dropout_backward", + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfCUDA=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=all_types_and(torch.float16, torch.bfloat16, torch.bool, torch.complex64), + supports_out=False, + sample_inputs_func=sample_inputs_dropout_backward, + skips=( + DecorateInfo(unittest.skip('Skipped!'), 'TestJit', 'test_variant_consistency_jit'), + # Lazy tensor failures + DecorateInfo(unittest.skip('Skipped!'), 'TestLazyOpInfo', 'test_dispatched_to_lazy'), + # These tests fail only when built with ASAN + DecorateInfo(unittest.skip("Fails with ASAN"), 'TestLazyOpInfo', 'test_correctness', active_if=TEST_WITH_ASAN), + DecorateInfo( + unittest.skip("Fails with ASAN"), + 'TestLazyOpInfo', + 'test_correctness_with_reusing_ir', + active_if=TEST_WITH_ASAN + ), + ), + ), + OpInfo( + "nn.functional.dropout2d", + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.dropout2d, input, *args, **kwargs), + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.complex64), + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu')), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + check_batched_forward_grad=False, + # As per the docs, valid input dims are (3, 4) + sample_inputs_func=partial(sample_inputs_dropout, valid_input_dim=(3, 4)), + inplace_variant=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.dropout2d, input, *args, **kwargs, inplace=True)), + OpInfo( + "nn.functional.dropout3d", + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.dropout3d, input, *args, **kwargs), + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.complex64), + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu')), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + check_batched_forward_grad=False, + # As per the docs, valid input dims are (4, 5) + sample_inputs_func=partial(sample_inputs_dropout, valid_input_dim=(4, 5)), + inplace_variant=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.dropout3d, input, *args, **kwargs, inplace=True)), + OpInfo( + "nn.functional.alpha_dropout", + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.alpha_dropout, input, *args, **kwargs), + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.complex64), + gradcheck_wrapper=wrapper_set_seed, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + sample_inputs_func=sample_inputs_dropout, + check_batched_forward_grad=False, + inplace_variant=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.alpha_dropout, input, *args, **kwargs, inplace=True), + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # AssertionError: Tensor-likes are not close! + # Fails in cuda11.7 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_compare_cpu', device_type='cuda'), + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_compare_cpu', device_type='xpu'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'),),), + # In training mode, feature_alpha_dropout currently doesn't support inputs of complex dtype + # unlike when `train=False`, it supports complex inputs, hence 2 OpInfos to cover all cases + OpInfo( + "nn.functional.feature_alpha_dropout", + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.feature_alpha_dropout, input, *args, **kwargs), + variant_test_name="with_train", + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and(torch.float16, torch.bfloat16, torch.complex64), + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + # torch.autograd.gradcheck.GradcheckError: While computing batched gradients, got: + # vmap: We do not yet support calling random operations inside of vmap. + # Please perform random operations outside of vmap as a workaround + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', "test_forward_mode_AD"), + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', "test_inplace_forward_mode_AD"), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu')), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + # As per the docs, valid input dims are (4, 5) + sample_inputs_func=partial(sample_inputs_dropout, train=True, valid_input_dim=(4, 5)), + inplace_variant=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.feature_alpha_dropout, input, *args, **kwargs, inplace=True)), + OpInfo( + "nn.functional.feature_alpha_dropout", + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.feature_alpha_dropout, input, *args, **kwargs), + variant_test_name="without_train", + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'),), + gradcheck_wrapper=wrapper_set_seed, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + sample_inputs_func=partial(sample_inputs_dropout, train=False), + inplace_variant=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.feature_alpha_dropout, input, *args, **kwargs, inplace=True)), + OpInfo( + "nn.functional.one_hot", + ref=reference_one_hot, + supports_out=False, + dtypes=_dispatch_dtypes((torch.int64,)), + sample_inputs_func=sample_inputs_one_hot, + ), + OpInfo( + "nn.functional.embedding", + aten_backward_name="embedding_dense_backward", + # We use lambda to reshuffle the positional arguments. + # This is because currently only the `input` field of SampleInput + # is tested in gradient tests. + op=lambda weight, idx, **kwargs: torch.nn.functional.embedding(idx, weight, **kwargs), + dtypes=floating_types_and(torch.bfloat16, torch.float16), + sample_inputs_func=sample_inputs_embedding, + allow_cow_input_materialize_forward=[0], + error_inputs_func=error_inputs_embedding, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # Fails on CI https://github.com/pytorch/pytorch/issues/85377 + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_compare_cpu'), + # Reference: https://github.com/pytorch/pytorch/issues/67084 + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view', device_type='cuda'), + # Not a problem: embedding does weird stuff to its input (it renormalizes) + DecorateInfo(unittest.skip('Allowed exemption'), 'TestCompositeCompliance', 'test_operator'), + # Fails due to non-determinism (see issue #74679) + # TODO: Investigate why more granular skips in the test don't work in CI + DecorateInfo(unittest.skip('Skipped!'), + 'TestExpandedWeightFunctional', + 'test_expanded_weight_forward'), + ), + supports_expanded_weight=True, + supports_out=False, + ), + OpInfo( + "nn.functional.embedding_bag", + # We use lambda to reshuffle the positional arguments. + # This is because currently only the `input` field of SampleInput + # is tested in gradient tests. + op=lambda weight, idx, **kwargs: torch.nn.functional.embedding_bag(idx, weight, **kwargs), + dtypes=floating_types_and(torch.bfloat16, torch.float16), + dtypesIfCUDA=floating_types_and(torch.bfloat16, torch.float16), + # backward is not supported for mode `max` and dtype `bfloat16` + backward_dtypesIfCUDA=floating_types_and(torch.float16), + sample_inputs_func=sample_inputs_embedding_bag, + skips=( + # lambda impl + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.expectedFailure, 'TestNormalizeOperators', 'test_normalize_operator_exhaustive'), + # Not a problem: embedding_bag does weird stuff to its input (it renormalizes) + DecorateInfo(unittest.skip('Allowed exemption'), 'TestCompositeCompliance', 'test_operator'), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=0.006, rtol=0.047)}), + 'TestConsistency', 'test_output_grad_match', device_type='mps'), + ), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + supports_out=False, + supports_gradgrad=False, + allow_cow_input_materialize_forward=[0], + ), + OpInfo( + "nn.functional.multi_head_attention_forward", + op=lambda input, *args, **kwargs: + wrapper_set_seed(torch.nn.functional.multi_head_attention_forward, input, *args, **kwargs), + dtypes=floating_types_and(torch.bfloat16, torch.float16), + sample_inputs_func=sample_inputs_multi_head_attention_forward, + skips=( + # Tensor-likes are not close + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_noncontiguous_samples', dtypes=(torch.float32,)), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=5e-3, rtol=0)}), 'TestDecomp', 'test_comprehensive'), + + # TODO skip this for now since we can't skip on runtime arch support (taken from scaled_dot_product_attention) + DecorateInfo(unittest.skip("Skipped!"), 'TestInductorOpInfo', 'test_comprehensive'), + # randomness + DecorateInfo(unittest.skip("Skipped!"), 'TestFwdGradients', 'test_forward_mode_AD'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + # lambda impl + # AssertionError: JIT Test does not execute any logic + DecorateInfo(unittest.expectedFailure, 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.expectedFailure, "TestNormalizeOperators", "test_normalize_operator_exhaustive"), + # tests running very slowly break slow tests, so we skip them instead of using `slowTest`. + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_forward_ad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCompositeCompliance', 'test_operator'), + DecorateInfo( + unittest.skip("Skipped - baddbmm decomp does not have enough precision for 16-bit float"), + 'TestDecomp', + 'test_comprehensive', + dtypes=(torch.bfloat16, torch.float16), + ), + DecorateInfo( + unittest.skip("Skipped - baddbmm decomp does not have enough precision for 16-bit float"), + 'TestDecomp', + 'test_quick', + dtypes=(torch.bfloat16, torch.float16))), + supports_out=False, + supports_gradgrad=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + ), + UnaryUfuncInfo( + "nn.functional.softplus", + aten_backward_name='softplus_backward', + ref=reference_softplus, + sample_kwargs=lambda device, dtype, input: ({'beta': 3, 'threshold': .2}, {'beta': 3, 'threshold': .2}), + sample_inputs_func=partial(sample_inputs_elementwise_unary, op_kwargs={'beta': 3, 'threshold': .2}), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.bfloat16, torch.float16), + dtypesIfMPS=floating_types_and( + torch.bfloat16, torch.float16, torch.int32, torch.uint8, torch.bool, torch.int8, torch.int16 + ), + decorators=( + DecorateInfo( + toleranceOverride + ({ + torch.half: tol(atol=1e-2, rtol=1e-2), + torch.bfloat16: tol(atol=1e-2, rtol=1e-2), + }), + 'TestUnaryUfuncs'), + ), + ), + OpInfo( + "nn.functional.mse_loss", + aten_backward_name='mse_loss_backward', + ref=loss_reference_reduction_wrapper(lambda input, target: (input - target) ** 2), + sample_inputs_func=sample_inputs_loss, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + skips=( + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":252, + # please report a bug to PyTorch. + DecorateInfo(unittest.expectedFailure, "TestJit", "test_variant_consistency_jit", dtypes=(torch.float32,),), + ), + ), + OpInfo( + "nn.functional.grid_sample", + dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + sample_inputs_func=sample_inputs_grid_sample, + reference_inputs_func=reference_inputs_grid_sample, + supports_gradgrad=True, + skips=( + # Exception: The operator 'aten::grid_sampler_2d_backward' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + ), + gradcheck_nondet_tol=1e-15), + # TODO: delete this OpInfo once we add meta support for grid_sampler_3d + OpInfo( + "grid_sampler_2d", + dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + sample_inputs_func=sample_inputs_grid_sampler_2d, + supports_gradgrad=True, + gradcheck_nondet_tol=1e-15, + skips=( + DecorateInfo(slowTest, 'TestDecomp', 'test_comprehensive', dtypes=(torch.float32, torch.float64), + active_if=IS_WINDOWS), + # Exception: The operator 'aten::grid_sampler_2d_backward' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + ),), + # TODO: Remove grid_sampler_3d tests once `nn.functional.grid_sample` has + # MPS support for all cases. + OpInfo( + "grid_sampler_3d", + dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + sample_inputs_func=sample_inputs_grid_sampler_3d, + supports_gradgrad=True, + gradcheck_nondet_tol=1e-15, + skips=( + # NOTE: Only run on MPS + DecorateInfo(unittest.skip('Skipped!'), device_type='cpu'), + DecorateInfo(unittest.skip('Skipped!'), device_type='cuda'), + DecorateInfo(unittest.skip('Skipped!'), device_type='xpu'), + DecorateInfo(unittest.skip('Skipped!'), device_type='meta'), + # Error: The operator 'aten::grid_sampler_3d_backward' is not currently implemented for the MPS device. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', device_type='mps'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4), + torch.float16: tol(atol=1e-4, rtol=1e-4), + torch.bfloat16: tol(atol=1e-4, rtol=1e-4)}), + "TestConsistency", "test_output_match", device_type="mps"), + ),), + OpInfo( + "argwhere", + ref=np.argwhere, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_autograd=False, + sample_inputs_func=sample_inputs_argwhere, + skips=( + # Compiler issue on ROCm. Might need to skip until ROCm5.5 + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + ), + ), + ReductionOpInfo( + 'all', + identity=True, + supports_autograd=False, + result_dtype=torch.bool, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + ref=reference_reduction_numpy(np.all), + skips=( + # FIXME: uint8 input returns uint8 instead of bool + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_result_dtype', dtypes=[torch.uint8]), + ), + ), + ReductionOpInfo( + 'any', + identity=False, + supports_autograd=False, + result_dtype=torch.bool, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + ref=reference_reduction_numpy(np.any), + skips=( + # FIXME: uint8 input returns uint8 instead of bool + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_result_dtype', dtypes=[torch.uint8]), + ), + ), + ReductionOpInfo( + 'amax', + nan_policy='propagate', + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + ref=reference_reduction_numpy(np.amax), + skips=( + # FIXME: reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.int64]), + # RuntimeError: MPS supports tensors with dimensions <= 16, but got 65. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='mps'), + ), + error_inputs_func=error_inputs_aminmax_amax_amin, + ), + ReductionOpInfo( + 'amin', + nan_policy='propagate', + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + ref=reference_reduction_numpy(np.amin), + skips=( + # FIXME: reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), + 'TestReductions', + 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.int64], + ), + # RuntimeError: MPS supports tensors with dimensions <= 16, but got 65. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_errors', device_type='mps'), + ), + error_inputs_func=error_inputs_aminmax_amax_amin, + ), + ReductionOpInfo( + 'argmax', + supports_multiple_dims=False, + supports_autograd=False, + assert_jit_shape_analysis=True, + result_dtype=torch.int64, + dtypes=all_types_and(torch.float16, torch.bfloat16), + ref=reference_reduction_numpy(np.argmax, supports_keepdims=False), + skips=( + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), + 'TestReductions', + 'test_ref_small_input', + device_type='xpu', + dtypes=floating_types_and( + torch.int64, torch.int8, torch.int16, torch.int32, torch.float16 + ), + ), + ), + ), + ReductionOpInfo( + 'argmin', + supports_multiple_dims=False, + supports_autograd=False, + result_dtype=torch.int64, + dtypes=all_types_and(torch.float16, torch.bfloat16), + ref=reference_reduction_numpy(np.argmin, supports_keepdims=False), + skips=( + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), + 'TestReductions', + 'test_ref_small_input', + device_type='xpu', + dtypes=floating_types_and( + torch.int64, torch.int8, torch.int16, torch.int32, torch.float16 + ), + ), + ), + ), + ReductionOpInfo( + 'count_nonzero', + identity=0, + supports_out=False, + supports_autograd=False, + result_dtype=torch.int64, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_reduction_count_nonzero, + ref=reference_reduction_numpy(np.count_nonzero), + skips=( + # FIXME: count_nonzero does not accept keepdim kwarg + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_default_keepdim'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_none_keepdim'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_single_keepdim'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_multi_keepdim'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_multi_unsorted_keepdim'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_offbounds_keepdim'), + # FIXME: dim=[] reduces all dimensions + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + ), + ), + ReductionOpInfo( + 'mean', + nan_policy='propagate', + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # FIXME: mean needs 'dim' parameter when using the 'out' overload. + # Adding it with 'generate_args_kwargs' does not work, since these also get passed + # onto the reference implementations. + supports_out=True, + assert_autodiffed=True, + assert_jit_shape_analysis=True, + promotes_int_to_float=True, + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + ref=reference_reduction_numpy(np.mean), + error_inputs_func=error_inputs_mean, + skips=( + # AssertionError: RuntimeError not raised : Expected RuntimeError when doing an unsafe cast from a result + # of dtype torch.float32 into an out= with dtype torch.long + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_out', device_type='cuda', dtypes=[torch.float32]), + # FIXME: mean does not support passing keepdim without passing dim + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_default_keepdim'), + # FIXME: mean reduces all dimensions when dim=[] + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: improve precision + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input', + dtypes=[torch.float16]), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_extremal_values', + device_type='cuda', dtypes=[torch.complex64]), + # Skipped on XPU because complex mean with extremal values (Inf/NaN) exhibits backend-dependent + # IEEE-754 behavior that differs from the CPU reference. + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_extremal_values', + device_type='xpu', dtypes=[torch.complex64]), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), + 'TestReductions', + 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.complex128], + ), + # AssertionError: RuntimeError not raised : Expected RuntimeError when doing an unsafe cast from a + # result of dtype torch.float32 into an out= with dtype torch.long + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + ), + ), + ReductionOpInfo( + 'nanmean', + nan_policy='omit', + assert_autodiffed=True, + promotes_int_to_float=True, + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16, torch.chalf), + sample_inputs_func=sample_inputs_nan_reduction(supports_multiple_dims=True), + ref=reference_reduction_numpy(np.nanmean), + skips=( + # AssertionError: False is not true : + # Failure in testing nodes' autodifferentiation. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # FIXME: prod reduces all dimensions when dim=[] + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: improve precision + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input', + dtypes=[torch.float16]), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_duplicate_values', + device_type='cuda', dtypes=[torch.float16]), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_extremal_values', + device_type='cuda', dtypes=[torch.complex64]), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=2e-5, rtol=4e-2)}), + "TestConsistency", "test_output_match", device_type="mps"), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.complex128]), + ), + ), + ReductionOpInfo( + 'std', + nan_policy='propagate', + supports_out=True, + complex_to_real=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True, + promotes_int_to_float=True, + check_batched_forward_grad=False, + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_std_var, + ref=reference_std_var(np.std), + generate_args_kwargs=generate_std_var_kwargs, + skips=( + # FIXME: cannot specify keepdim without dim + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_default_keepdim'), + # FIXME: dim=[] reduces all dimensions + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: improve precision + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input', + dtypes=(torch.float16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_duplicate_values', + dtypes=(torch.float16,)), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.float64]), + # MPS: std does not support automatic differentiation for outputs with complex dtype + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.complex64,) + ), + # The operator 'aten::std.correction_out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ), + ), + ReductionOpInfo( + 'std', + variant_test_name='unbiased', + nan_policy='propagate', + supports_out=False, + complex_to_real=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True, + promotes_int_to_float=True, + check_batched_forward_grad=False, + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_std_var_unbiased, + skips=( + # FIXME: dim=[] reduces all dimensions + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + # MPS: std does not support automatic differentiation for outputs with complex dtype + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.complex64,) + ), + ), + ), + ReductionOpInfo( + 'var', + nan_policy='propagate', + supports_out=True, + assert_autodiffed=True, + promotes_int_to_float=True, + complex_to_real=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_std_var, + ref=reference_std_var(np.var), + generate_args_kwargs=generate_std_var_kwargs, + skips=( + # FIXME: cannot specify keepdim without dim + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_default_keepdim'), + # FIXME: dim=[] reduces all dimensions + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: improve precision + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_duplicate_values'), + # NumPy is giving NaN for this + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_large_input'), + # RuntimeError: var does not support automatic differentiation for outputs with complex dtype. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_variant_consistency_eager', + device_type='mps', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_noncontiguous_samples', + device_type='mps', dtypes=(torch.complex64,) + ), + # NotImplementedError: The operator 'aten::var.correction_out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ), + ), + ReductionOpInfo( + 'var', + variant_test_name='unbiased', + nan_policy='propagate', + supports_out=False, + complex_to_real=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_autodiffed=True, + promotes_int_to_float=True, + check_batched_forward_grad=False, + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_std_var_unbiased, + skips=( + # FIXME: dim=[] reduces all dimensions + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + # RuntimeError: var does not support automatic differentiation for outputs with complex dtype. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.complex64,)), + ), + ), + ReductionOpInfo( + 'prod', + identity=1, + nan_policy='propagate', + supports_multiple_dims=False, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_int64=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + sample_inputs_func=sample_inputs_prod, + ref=prod_numpy, + skips=( + # FIXME: prod does not support passing keepdim without passing dim + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_default_keepdim'), + # FIXME: prod reduces all dimensions when dim=[] + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: prod does not support passing None to dim + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_none'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_none_keepdim'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input', + dtypes=[torch.float16, torch.complex64]), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_duplicate_values', + dtypes=[torch.uint8, torch.float16, torch.complex64]), + # FIXME: ValueError: The data in MaskedTensor a and Tensor b do not match + DecorateInfo(unittest.skip("Skipped!"), 'TestOperators', 'test_reduction_all', + dtypes=[torch.float16]), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.complex128, torch.int8, torch.int16, torch.int32, torch.int64]), + ), + ), + ReductionOpInfo( + 'sum', + identity=0, + nan_policy='propagate', + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_int64=True, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + ref=reference_reduction_numpy(np.sum), + error_inputs_sparse_func=error_inputs_sparse_reduction_sum, + sample_inputs_sparse_coo_func=partial(sample_inputs_sparse_reduction_sum, layout=torch.sparse_coo), + sample_inputs_sparse_csr_func=partial(sample_inputs_sparse_reduction_sum, layout=torch.sparse_csr), + sample_inputs_sparse_csc_func=partial(sample_inputs_sparse_reduction_sum, layout=torch.sparse_csc), + sample_inputs_sparse_bsr_func=partial(sample_inputs_sparse_reduction_sum, layout=torch.sparse_bsr), + sample_inputs_sparse_bsc_func=partial(sample_inputs_sparse_reduction_sum, layout=torch.sparse_bsc), + skips=( + # FIXME: sum does not support passing keepdim without passing dim + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_default_keepdim'), + # FIXME: sum reduces all dimensions when dim=[] + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: improve precision + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input', + dtypes=[torch.float16]), + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_duplicate_values', + dtypes=[torch.float16]), + DecorateInfo(unittest.skip("Skipped!"), 'TestOperators', 'test_reduction_all', + dtypes=[torch.float32]), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.complex128]), + ), + ), + ReductionOpInfo( + 'nansum', + identity=0, + nan_policy='omit', + supports_out=True, + promotes_int_to_int64=True, + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + sample_inputs_func=sample_inputs_nan_reduction(supports_multiple_dims=True), + ref=reference_reduction_numpy(np.nansum), + skips=( + # please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), 'TestJit', 'test_variant_consistency_jit'), + # FIXME: nansum reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: flaky test so skipped instead of xfailed + # possibly bad low precision reference in numpy + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input', + dtypes=[torch.float16]), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=3e-3, rtol=4e-2)}), + "TestConsistency", "test_output_match", device_type="mps"), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.complex128]), + # AssertionError: RuntimeError not raised : Expected RuntimeError when doing an unsafe cast from a result + # of dtype torch.float32 into an out= with dtype torch.long + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps', dtypes=(torch.float32,)), + ), + ), + ReductionOpInfo( + 'hash_tensor', + result_dtype=torch.uint64, + supports_autograd=False, + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=reference_hash_tensor, + skips=( + # hash_tensor reduces all dimensions when dim=[] (as do sum, prod etc.) + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo(unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # aten::hash_tensor hit the vmap fallback which is currently disabled + DecorateInfo(unittest.skip("Skipped!"), "TestVmapOperatorsOpInfo", "test_op_has_batch_rule"), + DecorateInfo(unittest.skip("Skipped!"), "TestVmapOperatorsOpInfo", "test_vmap_exhaustive"), + # NYI + DecorateInfo(unittest.expectedFailure, 'TestInductorOpInfo', 'test_comprehensive'), + # Sharding strategy NYI + DecorateInfo(unittest.expectedFailure, 'TestDTensorOps', 'test_dtensor_op_db'), + # Error: The operator 'aten::hash_tensor.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ) + ), + OpInfo( + "nn.functional.ctc_loss", + dtypes=floating_types(), + supports_out=False, + sample_inputs_func=sample_inputs_ctc_loss, + # gradcheck_wrapper, see https://github.com/pytorch/pytorch/issues/52241 + gradcheck_wrapper=gradcheck_wrapper_ctc_loss, + skips=( + # RuntimeError: derivative for aten::_ctc_loss_backward is not implemented + DecorateInfo( + unittest.expectedFailure, + "TestBwdGradients", + "test_fn_gradgrad", + dtypes=(torch.float64,), + ), + # RuntimeError: derivative for aten::_ctc_loss_backward is not implemented + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + dtypes=(torch.float32,), + ), + # Ref: https://github.com/pytorch/pytorch/issues/85231 + DecorateInfo(unittest.skip("Fails with ASAN"), + 'TestProxyTensorOpInfo', + 'test_make_fx_fake_exhaustive', active_if=TEST_WITH_ASAN), + # NotImplementedError: The operator 'aten::_ctc_loss' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + ), + OpInfo( + "nn.functional.cosine_embedding_loss", + dtypes=all_types_and(torch.half, torch.bfloat16, torch.bool), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-4, rtol=2e-3)}), + 'TestInductorOpInfo', 'test_comprehensive', device_type="cuda", + ), + ], + sample_inputs_func=sample_inputs_cosine_embedding_loss, + ), + OpInfo( + "nn.functional.nll_loss", + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.int16, torch.int32, torch.int64, torch.uint8, torch.bool, torch.int8 + ), + supports_out=False, + sample_inputs_func=sample_inputs_nll_loss, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + assert_jit_shape_analysis=True, + skips=( + # RuntimeError: + # undefined value tensor: + # File "", line 3 + # def the_method(i0, i1): + # return torch.nn.functional.nll_loss(i0, i1, weight=tensor([8.4784, 1.7658, 4.3228], dtype=torch.float32)) + # ~~~~~~ <--- HERE + DecorateInfo(unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit", dtypes=(torch.float32,),), + # Fails for unknown reason: https://github.com/pytorch/pytorch/issues/120782 + DecorateInfo( + unittest.skip("Skipped!"), + "TestCompositeCompliance", + "test_cow_input", + device_type='cuda', + ), + DecorateInfo(unittest.skip("FP16 nll_loss cases have not been enabled on MPS yet"), + dtypes=(torch.half,), device_type="mps"), + ), + ), + OpInfo( + "nn.functional.gaussian_nll_loss", + dtypes=floating_types_and(torch.half, torch.bfloat16), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_gaussian_nll_loss, + error_inputs_func=error_inputs_gaussian_nll_loss, + skips=( + # Pre-existing condition (calls .item); needs to be fixed + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_backward'), + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_forward_ad'), + # Pre-existing condition (calls .item); needs to be fixed + DecorateInfo(unittest.expectedFailure, 'TestCompositeCompliance', 'test_operator'), + # JIT does not support variadic tensors. + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":270, + # please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit", dtypes=(torch.float32,),), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=8e-3, rtol=2e-3)}), + "TestConsistency", "test_output_match", device_type="mps"), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=8e-3, rtol=2e-3)}), + "TestConsistency", "test_output_grad_match", device_type="mps"), + ), + ), + OpInfo( + "nn.functional.hinge_embedding_loss", + dtypes=floating_types_and(torch.half, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_hinge_embedding_loss, + error_inputs_func=error_inputs_hinge_embedding_loss, + reference_inputs_func=reference_inputs_hinge_embedding_loss, + ), + OpInfo( + "nn.functional.huber_loss", + aten_backward_name='huber_loss_backward', + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.float16, torch.bfloat16, torch.uint8, torch.bool, torch.int8, torch.int16, torch.int32 + ), + supports_out=False, + supports_forward_ad=True, + sample_inputs_func=sample_inputs_huber_loss, + error_inputs_func=error_inputs_huber_loss, + skips=( + # JIT does not support variadic tensors. + # RuntimeError: input->type()->kind() == TypeKind::OptionalType + # INTERNAL ASSERT FAILED at "../torch/csrc/jit/passes/utils/check_alias_annotation.cpp":270, + # please report a bug to PyTorch. + DecorateInfo(unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit", dtypes=(torch.float32,),), + ) + ), + OpInfo( + "nn.functional.pdist", + ref=reference_pdist, + sample_inputs_func=sample_inputs_pdist, + dtypes=floating_types(), + supports_out=False, + supports_gradgrad=False, + skips=( + DecorateInfo(unittest.skip("Unsupported on MPS for now"), 'TestCommon', 'test_numpy_ref_mps'), + # NotImplementedError: The operator 'aten::_pdist_forward' is not + # currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ) + ), + OpInfo( + "nn.functional.poisson_nll_loss", + dtypes=all_types_and(torch.half, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_poisson_nll_loss, + error_inputs_func=error_inputs_poisson_nll_loss, + ), + OpInfo( + "argsort", + dtypes=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.bool, torch.float16, torch.bfloat16), + sample_inputs_func=sample_inputs_sort, + supports_out=False, + supports_autograd=False, + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + dtypes=(torch.float32,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_non_standard_bool_values", + dtypes=[torch.bool], + device_type='cuda', + active_if=not TEST_WITH_ROCM + ), + ), + ), + OpInfo( + "repeat_interleave", + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16, torch.chalf), + backward_dtypesIfCUDA=floating_and_complex_types_and(torch.float16, torch.bfloat16, torch.chalf), + sample_inputs_func=sample_inputs_repeat_interleave, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + dtypes=(torch.float32, torch.complex64), + ), + ), + ), + OpInfo( + "nn.functional.pairwise_distance", + ref=lambda a, b, p=2.0, eps=1e-6, keepdim=False: ( + np.sum(np.abs(a - b + eps) ** p, axis=-1, keepdims=keepdim) ** (1 / p) + ), + sample_inputs_func=sample_inputs_pairwise_distance, + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + dtypes=(torch.float32, torch.complex64), + ), + ), + ), + OpInfo( + "nn.functional.pixel_shuffle", + sample_inputs_func=sample_inputs_pixel_shuffle, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + dtypes=(torch.float32, torch.complex64), + ), + ), + ), + OpInfo( + "nn.functional.pixel_unshuffle", + sample_inputs_func=sample_inputs_pixel_unshuffle, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + dtypes=(torch.float32, torch.complex64), + ), + ), + ), + OpInfo( + "nn.functional.channel_shuffle", + sample_inputs_func=sample_inputs_channel_shuffle, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + allow_cow_input_materialize_forward=[0], + allow_cow_input_materialize_backward=[0, 'output grad 0'], + skips=( + # Skip due to NotImplementedError for MPS device. + DecorateInfo(unittest.expectedFailure, 'TestConsistency'), + DecorateInfo(unittest.expectedFailure, "TestMeta", "test_dispatch_symbolic_meta_outplace_all_strides"), + # NotImplementedError: The operator 'aten::channel_shuffle' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + ), + OpInfo( + "nn.functional.kl_div", + sample_inputs_func=sample_inputs_kl_div, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ), + OpInfo( + "diagflat", + ref=lambda input, offset=0: np.diagflat(input, k=offset), + sample_inputs_func=sample_inputs_diagflat, + dtypes=all_types_and_complex_and(torch.bool, torch.bfloat16, torch.float16), + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + ), + OpInfo( + 'scatter_reduce', + variant_test_name='sum', + inplace_variant=torch.Tensor.scatter_reduce_, + # complex not added to dtypes as complex gradients are not properly handled + # and scatter_reduce hasn't been added to the whitelist in gen_variable_type yet + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_scatter_reduce, + skips=( + # Compiler issue on ROCm. Regression started in ROCm 6.4. + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_non_standard_bool_values', + dtypes=[torch.bool], active_if=TEST_WITH_ROCM), + ), + ), + OpInfo( + 'scatter_reduce', + variant_test_name='prod', + # complex not added to dtypes as complex gradients are not properly handled + # and scatter_reduce hasn't been added to the whitelist in gen_variable_type yet + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfCUDA=all_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + sample_inputs_func=sample_inputs_scatter_reduce, + skips=( + # Not implemented + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_forward_mode_AD'), + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_inplace_forward_mode_AD'), + DecorateInfo(unittest.expectedFailure, 'TestFwdGradients', 'test_fn_fwgrad_bwgrad'), + ), + ), + OpInfo( + 'scatter_reduce', + variant_test_name='mean', + # complex not added to dtypes as complex gradients are not properly handled + # and scatter_reduce hasn't been added to the whitelist in gen_variable_type yet + dtypes=all_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + dtypesIfCUDA=all_types_and(torch.float16, torch.bfloat16), + dtypesIfMPS=all_types_and(torch.float16, torch.bfloat16, torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_scatter_reduce, + ), + OpInfo( + 'scatter_reduce', + variant_test_name='amin', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfCUDA=all_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_scatter_reduce, + skips=( + # MPS: not supported for torch.int64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.int64,)), + ), + ), + OpInfo( + 'scatter_reduce', + variant_test_name='amax', + dtypes=all_types_and(torch.float16, torch.bfloat16, torch.bool), + dtypesIfCUDA=all_types_and(torch.float16, torch.bfloat16), + dtypesIfHpu=custom_types(torch.float32, torch.bfloat16), + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_scatter_reduce, + skips=( + # MPS: not supported for torch.int64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.int64,)), + ), + ), + OpInfo( + '_segment_reduce', + aten_name='segment_reduce', + variant_test_name='lengths', + dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + # RuntimeError: derivative for aten::_segment_reduce_backward is not implemented + supports_gradgrad=False, + sample_inputs_func=sample_inputs_segment_reduce, + skips=( + # FIXME: CUDA driver API confirmed a leak in + # __main__.TestJitCUDA.test_variant_consistency_jit_segment_reduce_cuda_float32 + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="cuda", + ), + # Error: The operator 'aten::segment_reduce' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type='mps'), + ), + ), + OpInfo( + '_segment_reduce', + aten_name='segment_reduce', + variant_test_name='offsets', + dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + # RuntimeError: derivative for aten::_segment_reduce_backward is not implemented + supports_gradgrad=False, + sample_inputs_func=partial(sample_inputs_segment_reduce, mode='offsets'), + skips=( + # FIXME: CUDA driver API confirmed a leak in + # __main__.TestJitCUDA.test_variant_consistency_jit_segment_reduce_cuda_float32 + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="cuda", + ), + # Error: The operator 'aten::segment_reduce' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type='mps'), + ), + ), +] +op_db += opinfo.definitions.op_db + + +# Separate registry for experimental Python Reference OpInfos. +python_ref_db = [ + # + # Elementwise Unary OpInfos + # + ElementwiseUnaryPythonRefInfo( + "_refs.abs", + torch_opinfo_name="abs", + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/49224 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_small', + dtypes=[torch.int8], active_if=TEST_WITH_ASAN), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.acos", + torch_opinfo_name="acos", + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_normal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + # Failing with wrong imaginary sign on at least some Windows jobs + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_small', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + # Failing with wrong imaginary sign on at least some Windows jobs + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs.acosh", + torch_opinfo_name="acosh", + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_normal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + # Failing with wrong imaginary sign on at least some Windows jobs + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_small', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.asin", + torch_opinfo_name="asin", + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-05, rtol=1e-03)}), + 'TestUnaryUfuncs', device_type='cuda'), + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=5e-05, rtol=2e-05)}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cpu' + ), + precisionOverride({torch.bfloat16: 1e-2}), + ], + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.asinh", + torch_opinfo_name="asinh", + decorators=(precisionOverride({torch.bfloat16: 5e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_small', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_normal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cdouble], + active_if=IS_WINDOWS), + ), + ), + PythonRefInfo( + "_refs.lerp", + torch_opinfo_name="lerp", + skips=( + # Exception: Dtypes torch.float32 and * are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.bool, torch.int16, torch.int32, torch.int64, torch.int8, torch.uint8) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.bool, torch.int16, torch.int32, torch.int64, torch.int8, torch.uint8) + ), + # RuntimeError: Failed to create function state object for: abs_dense_bool_bool + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', dtypes=(torch.bool,)), + ), + ), + PythonRefInfo( + "_refs.ones", + torch_opinfo_name="ones", + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + ), + ), + PythonRefInfo( + "_refs.zeros", + torch_opinfo_name="zeros", + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + ), + ), + PythonRefInfo( + "_refs.cauchy", + torch_opinfo_name="cauchy", + decorators=( + # TODO: RuntimeError: no _refs support for torch.rand_like + DecorateInfo(unittest.skip("TODO: RuntimeError: no _refs support for torch.rand_like"), + 'TestCommon', + 'test_python_ref'), + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: cauchy is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: cauchy is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor'), + DecorateInfo(unittest.skip("Expected: cauchy is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + ) + ), + PythonRefInfo( + "_refs.exponential", + torch_opinfo_name="exponential", + supports_out=True, + decorators=( + # dtypes that do not support check_uniform_bounds of rand_like + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', + dtypes=(torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64)), + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_dtypes'), + + # TODO: RuntimeError: no _refs support for torch.rand_like + DecorateInfo(unittest.skip("TODO: RuntimeError: no _refs support for torch.rand_like"), + 'TestCommon', + 'test_python_ref'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: exponential is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: exponential is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor'), + DecorateInfo(unittest.skip("Expected: exponential is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + ) + ), + PythonRefInfo( + "_refs.geometric", + torch_opinfo_name="geometric", + supports_out=True, + decorators=( + # dtypes that do not support check_uniform_bounds of rand_like + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_dtypes'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', + dtypes=(torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64)), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64)), + + # TODO: RuntimeError: no _refs support for torch.rand_like + DecorateInfo(unittest.skip("TODO: RuntimeError: no _refs support for torch.rand_like"), + 'TestCommon', + 'test_python_ref'), + DecorateInfo(unittest.skip("Expected: geometric is not comparable"), + 'TestCommon', + 'test_python_ref_executor', device_type='cuda'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: geometric is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: geometric is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: geometric is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + ) + ), + PythonRefInfo( + "_refs.log_normal", + torch_opinfo_name="log_normal", + supports_out=True, + decorators=( + # TODO: RuntimeError: no _refs support for torch.rand_like + DecorateInfo(unittest.skip("TODO: RuntimeError: no _refs support for torch.rand_like"), + 'TestCommon', + 'test_python_ref'), + DecorateInfo(unittest.skip("Expected: log_normal is not comparable"), + 'TestCommon', + 'test_python_ref_executor', device_type='cuda'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: log_normal is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: log_normal is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: log_normal is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + ) + ), + PythonRefInfo( + "_refs.normal", + torch_opinfo_name="normal", + supports_out=True, + decorators=( + # TODO: RuntimeError: no _refs support for torch.rand_like + DecorateInfo(unittest.skip("TODO: RuntimeError: no _refs support for torch.rand_like"), + 'TestCommon', + 'test_python_ref'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: normal is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: normal is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: normal is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: normal is not comparable"), 'TestDecomp', 'test_comprehensive'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.skip("make_traced() doesn't set seed properly!"), 'TestCommon', 'test_python_ref_executor'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + ) + ), + PythonRefInfo( + "_refs.normal", + torch_opinfo_name="normal", + torch_opinfo_variant_name="number_mean", + supports_out=True, + decorators=( + # TODO: RuntimeError: no _refs support for torch.rand_like + DecorateInfo(unittest.skip("TODO: RuntimeError: no _refs support for torch.rand_like"), + 'TestCommon', + 'test_python_ref'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: normal is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: normal is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: normal is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: normal is not comparable"), 'TestDecomp', 'test_comprehensive'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.skip("make_traced() doesn't set seed properly!"), 'TestCommon', 'test_python_ref_executor'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + ) + ), + PythonRefInfo( + "_refs.normal_", + op=torch.Tensor.normal_, + torch_opinfo_name="normal", + torch_opinfo_variant_name="in_place", + supports_out=False, + decorators=( + # TODO: RuntimeError: no _refs support for torch.rand_like + DecorateInfo(unittest.skip("TODO: RuntimeError: no _refs support for torch.rand_like"), + 'TestCommon', + 'test_python_ref'), + + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: normal is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: normal is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: normal is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: normal is not comparable"), 'TestDecomp', 'test_comprehensive'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + DecorateInfo(unittest.skip("make_traced() doesn't set seed properly!"), 'TestCommon', 'test_python_ref_executor'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + ) + ), + PythonRefInfo( + "_refs.arange", + torch_opinfo_name="arange", + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + # RuntimeError: value cannot be converted to type uint8_t without overflow + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.uint8,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.uint8,) + ), + ), + ), + PythonRefInfo( + "_refs.linspace", + torch_opinfo_name="linspace", + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + + # cpu implementation is wrong on some integral types + # https://github.com/pytorch/pytorch/issues/81996 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64), device_type="cpu"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64), device_type="cpu"), + + # cuda implementation is off-by-one on some inputs due to precision issues + # https://github.com/pytorch/pytorch/issues/82230 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64), + device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64), + device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64), + device_type="cuda"), + # AssertionError: Tensor-likes are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int64, torch.int32, torch.int16, torch.bool)), + # ValueError: value argument of type cannot be safely cast to type ! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int64, torch.int32, torch.int16, torch.bool) + ), + ), + ), + PythonRefInfo( + "_refs.linspace", + torch_opinfo_name="linspace", + torch_opinfo_variant_name="tensor_overload", + skips=( + # TypeError: 'int' object is not subscriptable + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + + # cpu implementation is wrong on some integral types + # https://github.com/pytorch/pytorch/issues/81996 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64), device_type="cpu"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.int8, torch.uint8, torch.int16, torch.int32, torch.int64), device_type="cpu"), + + # cuda implementation is off-by-one on some inputs due to precision issues + # https://github.com/pytorch/pytorch/issues/82230 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64), + device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64), + device_type="cuda"), + # TODO torch.ops.aten.copy is not in _refs + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.float32, torch.float64, torch.float16, torch.complex64, torch.complex128, torch.bfloat16), + device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.float32, torch.float64, torch.float16, torch.complex64, torch.complex128, torch.bfloat16), + device_type="cpu"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.int32, torch.int64), + device_type="cuda"), + # RuntimeError: no _refs support for aten.copy.default + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'), + # AssertionError: Tensor-likes are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.bool, torch.int16, torch.int32, torch.int64, torch.int8, torch.uint8) + ), + ), + ), + PythonRefInfo( + "_refs.logspace", + torch_opinfo_name="logspace", + skips=( + # Tests that assume input is a tensor or sequence of tensors + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_conj_view'), + + # Off-by-one issue when casting floats to ints + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.int16, torch.int32, torch.int64), + device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.int16, torch.int32, torch.int64), + device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor', + dtypes=(torch.int16, torch.int32, torch.int64), + device_type="cuda"), + # Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'), + ), + ), + PythonRefInfo( + "_refs.logspace", + torch_opinfo_name="logspace", + torch_opinfo_variant_name="tensor_overload", + skips=( + # TypeError: 'int' object is not subscriptable + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.expectedFailure, 'TestMathBits', 'test_conj_view'), + + # Off-by-one issue when casting floats to ints + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.int16, torch.int32, torch.int64), + device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.int16, torch.int32, torch.int64), + device_type="cuda"), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor', + dtypes=(torch.int16, torch.int32, torch.int64), + device_type="cuda"), + # TODO copy doesn't have prim refs + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=( + torch.float32, torch.float64, torch.float16, torch.complex64, + torch.complex128, torch.bfloat16, torch.int8, torch.uint8 + ), + device_type="cuda" + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=( + torch.float32, torch.float64, torch.float16, + torch.complex64, torch.complex128, torch.bfloat16, + torch.int16, torch.int32, torch.int64, torch.int8, torch.uint8 + ), + device_type="cpu"), + # Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'), + ), + ), + PythonRefInfo( + "_refs.meshgrid", + torch_opinfo_name="meshgrid", + torch_opinfo_variant_name="variadic_tensors", + ), + PythonRefInfo( + "_refs.take_along_dim", + torch_opinfo_name="take_along_dim", + skips=( + DecorateInfo(unittest.expectedFailure, + 'TestCommon', + 'test_python_ref'), + ), + ), + PythonRefInfo( + "_refs.to", + torch_opinfo_name="to", + skips=( + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'), + ), + ), + PythonRefInfo( + "_refs.triu", + torch_opinfo_name="triu", + ), + PythonRefInfo( + "_refs.tril", + torch_opinfo_name="tril", + ), + PythonRefInfo( + "_refs.triu_indices", + torch_opinfo_name="triu_indices", + # the implementation uses torch.stack that violates view consistency + validate_view_consistency=False, + skips=( + # skip these tests since we have non tensor input + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_noncontiguous_samples'), + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip('Skipped!'), 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Skipped!'), 'TestMathBits', 'test_neg_view'), + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.int64, torch.int32)), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + )), + PythonRefInfo( + "_refs.tril_indices", + torch_opinfo_name="tril_indices", + # the implementation uses torch.stack that violates view consistency + validate_view_consistency=False, + skips=( + # skip these tests since we have non tensor input + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_noncontiguous_samples'), + DecorateInfo(unittest.skip('Skipped!'), 'TestCommon', 'test_variant_consistency_eager'), + DecorateInfo(unittest.skip('Skipped!'), 'TestJit', 'test_variant_consistency_jit'), + DecorateInfo(unittest.skip('Skipped!'), 'TestMathBits', 'test_neg_view'), + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.int64, torch.int32)), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + )), + PythonRefInfo( + "_refs.meshgrid", + torch_opinfo_name="meshgrid", + torch_opinfo_variant_name="list_of_tensors", + ), + PythonRefInfo( + "_refs.movedim", + aliases=('moveaxis',), + torch_opinfo_name="movedim", + ), + PythonRefInfo( + "_refs.bucketize", + torch_opinfo_name="bucketize", + skips=( + # RuntimeError: It appears that you're trying to get value out of a tracing tensor with + # aten._local_scalar_dense.default - erroring out! [...] + # triggered by mid_val = boundaries[mid] + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_python_ref_executor"), + ) + ), + PythonRefInfo( + "_refs.equal", + torch_opinfo_name="equal", + skips=( + # RuntimeError: Cannot cast FakeTensor to number + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta',), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs.atan", + torch_opinfo_name="atan", + decorators=(precisionOverride({torch.bfloat16: 1e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_small', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.atanh", + torch_opinfo_name="atanh", + decorators=(precisionOverride({torch.bfloat16: 1e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_small', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cuda', dtypes=[torch.cfloat], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.bitwise_not", + torch_opinfo_name="bitwise_not", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.ceil", + torch_opinfo_name="ceil", + # Fails on int32 + # https://github.com/pytorch/pytorch/issues/85258 + ), + PythonRefInfo( + "_refs.item", + torch_opinfo_name="item", + skips=( + # RuntimeError: Cannot cast FakeTensor(FakeTensor(..., device='meta', size=()), cpu) to number + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta'), + # ValueError: Can't convert a tensor with 10 elements to a number! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'),), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.conj_physical", + torch_opinfo_name="conj_physical", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.cos", + torch_opinfo_name="cos", + decorators=(precisionOverride({torch.bfloat16: 1e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', + active_if=IS_WINDOWS), + # This fails on CUDA but passes on ROCm + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=(torch.cdouble,), device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_MACOS), + # AssertionError: Tensor-likes are not close! + # Greatest absolute difference: nan at index (700,) (up to 1e-05 allowed) + # Greatest relative difference: nan at index (700,) (up to 0.001 allowed) + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cuda', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.cosh", + torch_opinfo_name="cosh", + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/48641 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.int8]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=[torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_MACOS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', + dtypes=[torch.cfloat, torch.cdouble], active_if=IS_MACOS), + # AssertionError: Tensor-likes are not close! + # Greatest absolute difference: nan at index (6000,) (up to 1e-05 allowed) + # Greatest relative difference: nan at index (6000,) (up to 0.001 allowed) + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cuda', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.digamma", + torch_opinfo_name="digamma", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.erf", + torch_opinfo_name="erf", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.erfinv", + torch_opinfo_name="erfinv", + decorators=(precisionOverride({torch.float16: 1e-2, + torch.bfloat16: 1e-2, + torch.float32: 1e-4}),), + skips=( + # Reference: https://github.com/pytorch/pytorch/pull/49155#issuecomment-742664611 + DecorateInfo( + unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + active_if=TEST_SCIPY and version.parse(scipy.__version__) < version.parse("1.4.0")), + DecorateInfo( + unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + active_if=TEST_SCIPY and version.parse(scipy.__version__) < version.parse("1.4.0")), + DecorateInfo( + unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_small', + active_if=TEST_SCIPY and version.parse(scipy.__version__) < version.parse("1.4.0")), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.erfc", + torch_opinfo_name="erfc", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.exp", + torch_opinfo_name="exp", + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/48010 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', dtypes=[torch.chalf], + active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.expm1", + torch_opinfo_name="expm1", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.exp2", + torch_opinfo_name="exp2", + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=[torch.cdouble]), + # Reference: https://github.com/pytorch/pytorch/issues/48010 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.fill", + torch_opinfo_name="fill", + supports_out=True, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.floor", + torch_opinfo_name="floor", + # Fails on int32 + # https://github.com/pytorch/pytorch/issues/85258 + ), + ElementwiseUnaryPythonRefInfo( + "_refs.frexp", + torch_opinfo_name="frexp", + # Skipped due to numerical failures on Windows CI. + # This is also skipped in frexp earlier in the file. + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', 'test_reference_numerics_extremal', + active_if=IS_WINDOWS), + # The operator 'aten::frexp.Tensor_out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.frac", + torch_opinfo_name="frac", + skips=( + DecorateInfo( + unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=(torch.bfloat16, torch.float16, torch.float32, torch.float64)), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.imag", + torch_opinfo_name="imag", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.isfinite", + torch_opinfo_name="isfinite", + supports_out=True, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.isinf", + torch_opinfo_name="isinf", + supports_out=True, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.isposinf", + torch_opinfo_name="isposinf", + supports_out=True, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.isneginf", + torch_opinfo_name="isneginf", + supports_out=True, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.isnan", + torch_opinfo_name="isnan", + supports_out=True, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.isreal", + torch_opinfo_name="isreal", + supports_out=True, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.i0", + torch_opinfo_name="i0", + decorators=(precisionOverride({torch.bfloat16: 3e-1, + torch.float16: 5e-1}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), + 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=(torch.int8,)), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.lgamma", + torch_opinfo_name="lgamma", + decorators=(precisionOverride({torch.float16: 7e-1}),), + skips=( + # Reference: https://github.com/pytorch/pytorch/pull/50140#issuecomment-756150214 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=[torch.float32, torch.float64], active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=[torch.float32, torch.float64], active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.multigammaln", + torch_opinfo_name="mvlgamma", + torch_opinfo_variant_name="mvlgamma_p_1", + skips=skips_mvlgamma(), + decorators=( + DecorateInfo(torch.testing._internal.common_utils.markDynamoStrictTest, 'TestUnaryUfuncs', + 'test_reference_numerics_large'), + DecorateInfo(torch.testing._internal.common_utils.xfailIfTorchDynamo, 'TestUnaryUfuncs', + 'test_reference_numerics_large'), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.multigammaln", + torch_opinfo_name="mvlgamma", + torch_opinfo_variant_name="mvlgamma_p_3", + skips=skips_mvlgamma(), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.multigammaln", + torch_opinfo_name="mvlgamma", + torch_opinfo_variant_name="mvlgamma_p_5", + skips=skips_mvlgamma(), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.log", + torch_opinfo_name="log", + decorators=(precisionOverride({torch.bfloat16: 5e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.log1p", + torch_opinfo_name="log1p", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.log10", + torch_opinfo_name="log10", + decorators=(precisionOverride({torch.bfloat16: 5e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.log2", + torch_opinfo_name="log2", + decorators=(precisionOverride({torch.bfloat16: 1e-1}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=[torch.cfloat, torch.cdouble]), + ), + ), + PythonRefInfo( + "_refs.logsumexp", + torch_opinfo_name="logsumexp", + # When keepdim=False logsumexp function uses squeeze operation + # that is not yet exposed in nvFuser's Python API. + ), + PythonRefInfo( + "_refs.log_softmax", + torch_opinfo_name="log_softmax", + torch_opinfo_variant_name="with_dtype", + skips=( + # AssertionError: Tensor-likes are not close! + # RuntimeError: softmax only supported for floating types + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.float32, torch.complex64, torch.complex32) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.float32, torch.complex64, torch.complex32) + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nan_to_num", + torch_opinfo_name="nan_to_num", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.neg", + torch_opinfo_name="neg", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.positive", + torch_opinfo_name="positive", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.real", + torch_opinfo_name="real", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.reciprocal", + torch_opinfo_name="reciprocal", + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/45690 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=[torch.cfloat, torch.cdouble]), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.round", + torch_opinfo_name="round", + # Fails on int32 + # https://github.com/pytorch/pytorch/issues/85258 + skips=( + DecorateInfo(toleranceOverride({torch.bfloat16: tol(atol=1e-3, rtol=0.016)}), + "TestUnaryUfuncs", "test_reference_numerics_extremal", + device_type="cuda"), + DecorateInfo(toleranceOverride({torch.bfloat16: tol(atol=1e-3, rtol=0.016)}), + "TestUnaryUfuncs", "test_reference_numerics_normal", + device_type="cuda"), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.rsqrt", + torch_opinfo_name="rsqrt", + decorators=(precisionOverride({torch.half: 5e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=(torch.cfloat, torch.cdouble), device_type='cuda'), + # AssertionError: Tensor-likes are not close! + # Greatest absolute difference: nan at index (700,) (up to 0.01 allowed) + # Greatest relative difference: nan at index (700,) (up to 0.001 allowed) + DecorateInfo(unittest.expectedFailure, 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=(torch.chalf,), device_type='cuda'), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.sigmoid", + torch_opinfo_name="sigmoid", + aliases=('_refs.special.expit',), + # Reference: https://github.com/pytorch/pytorch/issues/56012 + handles_complex_extremal_values=False, + handles_large_floats=False, + decorators=(precisionOverride({torch.float16: 1e-2, + torch.complex64: 1e-1, + torch.bfloat16: 1e-2}),), + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/56012 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=[torch.complex64, torch.cdouble], device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=[torch.chalf, torch.complex64, torch.cdouble], device_type='cuda') + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.sign", + torch_opinfo_name="sign", + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/41245 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=[torch.bfloat16, torch.float16, torch.float32, + torch.float64]), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.sgn", + torch_opinfo_name="sgn", + # This is an issue with the vectorised abs on CPU + handles_complex_extremal_values=False, + handles_large_floats=False, + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/41245 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=[torch.bfloat16, torch.float16, torch.float32, + torch.float64]), + # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype. + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.complex32,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.complex32,) + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.signbit", + torch_opinfo_name="signbit", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.sin", + torch_opinfo_name="sin", + decorators=(precisionOverride({torch.bfloat16: 1e-2}),), + skips=( + # Fails on CUDA but passes on ROCm + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=(torch.cdouble,), device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=(torch.cfloat, torch.cdouble,), device_type='cpu', + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.sinc", + torch_opinfo_name="sinc", + decorators=(precisionOverride({torch.bfloat16: 1e-2, + torch.float16: 1e-2}),), + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/49133 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_small', + dtypes=[torch.cfloat]), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.sinh", + torch_opinfo_name="sinh", + decorators=(precisionOverride({torch.float16: 1e-2}),), + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=(torch.cdouble,)), + # Reference: https://github.com/pytorch/pytorch/issues/48641 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.int8]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + ), + ), + PythonRefInfo( + "_refs.softmax", + torch_opinfo_name="softmax", + torch_opinfo_variant_name="with_dtype", + skips=( + # AssertionError: Tensor-likes are not close! + # RuntimeError: softmax only supported for floating types + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.float32, torch.complex64), + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.float32, torch.complex64), + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.sqrt", + torch_opinfo_name="sqrt", + decorators=( + precisionOverride({torch.bfloat16: 7e-2}), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=1e-2, rtol=0)}), + 'TestUnaryUfuncs', 'test_reference_numerics_large'), + ), + skips=( + # Reference: https://github.com/pytorch/pytorch/issues/47358 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=(torch.cfloat, torch.cdouble), + active_if=IS_MACOS), + # Reference: https://github.com/pytorch/pytorch/pull/47293#issuecomment-721774436 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=(torch.bfloat16,)), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.square", + torch_opinfo_name="square", + decorators=(precisionOverride({torch.complex64: 3e-4, torch.bfloat16: 3e-1}),), + skips=( + # AssertionError: Reference result was farther (2.2417024338305655e-07) from the precise computation + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref_executor', dtypes=(torch.complex64,)), + # Reference: https://github.com/pytorch/pytorch/issues/52549 + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cuda', dtypes=[torch.cfloat, torch.cdouble]), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.tan", + torch_opinfo_name="tan", + decorators=[ + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-04, rtol=1e-05)}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cuda'), + ], + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs.tanh", + torch_opinfo_name="tanh", + decorators=[ + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-04, rtol=2e-05)}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cuda'), + ], + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_large', + device_type='xpu', + dtypes=(torch.chalf,), active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.trunc", + torch_opinfo_name="trunc", + # Fails on int32 + # https://github.com/pytorch/pytorch/issues/85258 + ), + PythonRefInfo( + "_refs.special.log_softmax", + torch_opinfo_name="log_softmax", # alias + torch_opinfo_variant_name="with_dtype", + supports_out=False, + skips=( + # AssertionError: Tensor-likes are not close! + # NotImplementedError: log_softmax for complex is not supported for MPS + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.float32, torch.complex64, torch.complex32) + ), + ), + ), + PythonRefInfo( + "_refs.special.softmax", + torch_opinfo_name="softmax", # alias + torch_opinfo_variant_name="with_dtype", + supports_out=False, + skips=( + # Exception: softmax only supported for floating types + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.float32, torch.complex64), + ), + ), + ), + # + # Elementwise Unary Special OpInfos + # + ElementwiseUnaryPythonRefInfo( + "_refs.special.logit", + torch_opinfo_name="logit", + ), + # + # Elementwise Unary nn.functional OpInfos + # + PythonRefInfo( + "_refs.nn.functional.alpha_dropout", + torch_opinfo_name="nn.functional.alpha_dropout", + decorators=( + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestCommon', + 'test_python_ref'), + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestCommon', + 'test_python_ref_executor', device_type='cuda'), + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestMathBits', + 'test_neg_view'), + # AssertionError: Tensor-likes are not close! + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestCommon', + 'test_compare_cpu'), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.celu", + torch_opinfo_name="nn.functional.celu", + supports_out=True, + ), + PythonRefInfo( + "_refs.nn.functional.channel_shuffle", + torch_opinfo_name="nn.functional.channel_shuffle", + supports_out=True, + skips=( + # NotImplementedError: The operator 'aten::channel_shuffle' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.threshold", + torch_opinfo_name="nn.functional.threshold", + supports_out=True, + skips=( + # RuntimeError: [srcBuf length] > 0 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'), + ), + ), + PythonRefInfo( + "_refs.nn.functional.dropout", + torch_opinfo_name="nn.functional.dropout", + decorators=( + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestCommon', + 'test_python_ref'), + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestMathBits', + 'test_conj_view'), + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestMathBits', + 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Expected: dropout is not comparable"), + 'TestMathBits', + 'test_neg_view'), + # dropout is not comparable + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.elu", + torch_opinfo_name="nn.functional.elu", + supports_out=True, + decorators=[ + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=1e-03, rtol=1.2e-03), + torch.bfloat16: tol(atol=1e-03, rtol=1.2e-03) + }), + 'TestUnaryUfuncs', device_type='cuda', + ), ], + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.hardtanh", + torch_opinfo_name="nn.functional.hardtanh", + supports_out=True, + ), + PythonRefInfo( # TODO: Port this to an UnaryOpInfo + "_refs.nn.functional.gelu", + torch_opinfo_name="nn.functional.gelu", + ), + PythonRefInfo( + "_refs.nn.functional.layer_norm", + torch_opinfo_name="nn.functional.layer_norm", + skips=( + # Reference result was farther (3.5762786809723224e-07) from the precise computation + # than the torch result was (2.5068410824946596e-07)! + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref', + dtypes=(torch.float32,), device_type='cpu'), + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.float32,) + ), + ), + ), + PythonRefInfo( + "_refs.nn.functional.glu", + torch_opinfo_name="nn.functional.glu", + supports_out=True, + skips=( + # Exception: Tensor-likes are not equal! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=( + torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool + )), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', dtypes=( + torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool + )), + ), + ), + PythonRefInfo( + "_refs.nn.functional.pairwise_distance", + torch_opinfo_name="nn.functional.pairwise_distance", + supports_out=True, + ), + PythonRefInfo( + "_refs.nn.functional.pdist", + torch_opinfo_name="nn.functional.pdist", + supports_out=True, + skips=( + # RunTimeError: no _refs support for torch.Tensor.index_select + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref'), + # Reference result was farther (1.946091651916504e-05) from the precise + # computation than the torch result was (1.1920928955078125e-06)! + DecorateInfo( + unittest.expectedFailure, + 'TestCommon', + 'test_python_ref_torch_fallback', + dtypes=(torch.float32,), + device_type='cpu', + ), + # NotImplementedError: The operator 'aten::_pdist_forward' is not + # currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'), + )), + PythonRefInfo( + "_refs.nn.functional.leaky_relu", + torch_opinfo_name="nn.functional.leaky_relu", + supports_out=True, + ), + PythonRefInfo( + "_refs.nn.functional.log_softmax", + torch_opinfo_name="log_softmax", # alias + torch_opinfo_variant_name="with_dtype", + supports_out=False, + skips=( + # NotImplementedError: log_softmax for complex is not supported for MPS + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.float32, torch.complex64, torch.complex32) + ), + ), + ), + PythonRefInfo( + "_refs.nn.functional.pixel_shuffle", + torch_opinfo_name="nn.functional.pixel_shuffle", + ), + PythonRefInfo( + "_refs.nn.functional.pixel_unshuffle", + torch_opinfo_name="nn.functional.pixel_unshuffle", + ), + PythonRefInfo( + "_refs.nn.functional.poisson_nll_loss", + torch_opinfo_name="nn.functional.poisson_nll_loss", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.prelu", + torch_opinfo_name="nn.functional.prelu", + skips=( + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.float16, torch.bfloat16, torch.float32) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.float16, torch.bfloat16, torch.float32) + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.relu", + torch_opinfo_name="nn.functional.relu", + supports_out=True, + skips=( + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.float16, torch.bfloat16, torch.float32) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.float16, torch.bfloat16, torch.float32) + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.relu6", + torch_opinfo_name="nn.functional.relu6", + supports_out=True, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.mish", + torch_opinfo_name="nn.functional.mish", + supports_out=True, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-03)}), + 'TestUnaryUfuncs',), ], + skips=( + # AssertionError: Tensor-likes are not equal! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=( + torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool + )), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', dtypes=( + torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool + )), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.selu", + torch_opinfo_name="nn.functional.selu", + supports_out=True, + decorators=[ + DecorateInfo( + toleranceOverride({ + torch.float16: tol(atol=1e-2, rtol=1.8e-2), + torch.bfloat16: tol(atol=1e-2, rtol=1.8e-2) + }), + 'TestUnaryUfuncs', device_type='cuda', + ), ], + ), + PythonRefInfo( + "_refs.nn.functional.softmax", + torch_opinfo_name="softmax", # alias + torch_opinfo_variant_name="with_dtype", + supports_out=False, + skips=( + # RuntimeError: softmax only supported for floating types + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.float32, torch.complex64), + ), + ), + ), + PythonRefInfo( + "_refs.nn.functional.softmin", + torch_opinfo_name="nn.functional.softmin", + torch_opinfo_variant_name="with_dtype", + supports_out=False, + skips=( + # RuntimeError: softmax only supported for floating types + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.complex64, torch.float32), + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.softplus", + torch_opinfo_name="nn.functional.softplus", + skips=( + # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bool}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + # AssertionError: Tensor-likes are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.bool, torch.int8,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', + device_type='mps', dtypes=(torch.bool,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.bool, torch.int8,) + ), + ), + ), + PythonRefInfo( + "_refs.nn.functional.l1_loss", + torch_opinfo_name="nn.functional.l1_loss", + ), + PythonRefInfo( + "_refs.nn.functional.margin_ranking_loss", + torch_opinfo_name="nn.functional.margin_ranking_loss", + ), + PythonRefInfo( + "_refs.nn.functional.mse_loss", + torch_opinfo_name="nn.functional.mse_loss", + ), + PythonRefInfo( + "_refs.nn.functional.smooth_l1_loss", + torch_opinfo_name="nn.functional.smooth_l1_loss", + ), + PythonRefInfo( + "_refs.nn.functional.hinge_embedding_loss", + torch_opinfo_name="nn.functional.hinge_embedding_loss", + ), + PythonRefInfo( + "_refs.nn.functional.nll_loss", + torch_opinfo_name="nn.functional.nll_loss", + # The corresponding PyTorch op doesn't support out. But the ref is + # registered as a decomp and ATen has an out variant. + supports_out=True, + # For simpler indexing, we flatten target indices, then reshape the result tensor. + # This creates inconsistent view state with reference impl. + validate_view_consistency=False, + skips=( + # RuntimeError: It appears that you're trying to get value out of a tracing tensor - erroring out! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor', device_type="cuda" + ), + # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bool}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + # AssertionError: Scalars are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.uint8, torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', + device_type='mps', dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.uint8, torch.bool,), + ), + ), + ), + PythonRefInfo( + "_refs.nn.functional.huber_loss", + torch_opinfo_name="nn.functional.huber_loss", + # The corresponding PyTorch op doesn't support out. But the ref is + # registered as a decomp and ATen has an out variant. + supports_out=True, + skips=( + # torch.bool - Subtraction, the `-` operator, with two bool tensors is not supported. + # Use the `^` or `logical_xor()` operator instead. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + # AssertionError: Tensor-likes are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool), + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', + dtypes=(torch.bool,), + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool), + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.tanhshrink", + torch_opinfo_name="nn.functional.tanhshrink", + decorators=[ + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_normal', + device_type='cpu', dtypes=[torch.cfloat, torch.cdouble]), + DecorateInfo( + toleranceOverride({torch.bfloat16: tol(atol=1e-02, rtol=1.6e-02), + torch.complex64: tol(atol=6e-04, rtol=1e-05)}), + 'TestUnaryUfuncs', 'test_reference_numerics_extremal', device_type='cuda'), + ], + skips=( + # in each case, pytorch will produce a nan while numpy will not + DecorateInfo(unittest.skip("Fails on some jobs works on others!"), + 'TestUnaryUfuncs', "test_reference_numerics_large", + dtypes=(torch.complex64, torch.complex128), + active_if=(IS_MACOS)), + DecorateInfo(unittest.skip("Fails on some jobs works on others!"), + 'TestUnaryUfuncs', "test_reference_numerics_extremal", + dtypes=(torch.complex64, torch.complex128), + device_type='cpu', + active_if=(IS_MACOS or IS_WINDOWS)), + DecorateInfo(unittest.skip("Skipped!"), 'TestUnaryUfuncs', + 'test_reference_numerics_extremal', + device_type='xpu', dtypes=[torch.cfloat, torch.cdouble], + active_if=IS_WINDOWS), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.hardshrink", + torch_opinfo_name="nn.functional.hardshrink", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.nn.functional.softshrink", + torch_opinfo_name="nn.functional.softshrink", + skips=( + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.float16, torch.bfloat16, torch.float32) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.float16, torch.bfloat16, torch.float32) + ), + ), + ), + # + # Elementwise Binary Reference OpInfos + # + ElementwiseBinaryPythonRefInfo( + "_refs.add", + torch_opinfo_name="add", + # https://github.com/pytorch/pytorch/issues/76944 + supports_two_python_scalars=True, + supports_one_python_scalar=True, + decorators=( + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=1e-2, rtol=0)}), + 'TestBinaryUfuncs', 'test_reference_numerics'), + ), + skips=( + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics_extremal_values', + dtypes=(torch.complex64, torch.complex128)), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.atan2", + torch_opinfo_name="atan2", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.bitwise_and", + torch_opinfo_name="bitwise_and", + skips=( + # AssertionError: Tensor-likes are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.bool) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int16) + ), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.bitwise_left_shift", + torch_opinfo_name="bitwise_left_shift", + skips=( + # https://github.com/pytorch/pytorch/issues/70904 + DecorateInfo(unittest.skip("Some inputs produce undefined outputs"), 'TestCommon', 'test_compare_cpu'), + # RuntimeError: value cannot be converted to type * without overflow + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.uint8, torch.int8) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', + dtypes=(torch.uint8, torch.int8) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.uint8, torch.int8) + ), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.bitwise_right_shift", + torch_opinfo_name="bitwise_right_shift", + skips=( + # # https://github.com/pytorch/pytorch/issues/70904 + DecorateInfo(unittest.skip("Skipped some inputs produce undefined outputs"), 'TestCommon', 'test_compare_cpu'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.bitwise_or", + torch_opinfo_name="bitwise_or", + skips=( + # AssertionError: Tensor-likes are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.bool) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int16) + ), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.bitwise_xor", + torch_opinfo_name="bitwise_xor", + skips=( + # AssertionError: Tensor-likes are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int16, torch.bool) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int16) + ), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.copysign", + torch_opinfo_name="copysign", + skips=( + # RuntimeError: Expected divisor (b) to be on the same device (cuda:0) as dividend (a), but it is found on cpu! + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', 'test_type_promotion'), + # FIXME output 0: meta disagrees with real impl + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_binary_ufuncs_mixed_dtype'), + ) + ), + ElementwiseBinaryPythonRefInfo( + "_refs.div", + torch_opinfo_name="div", + torch_opinfo_variant_name="no_rounding_mode", + # https://github.com/pytorch/pytorch/issues/76944 + supports_two_python_scalars=True, + supports_one_python_scalar=True, + skips=( + # NotImplementedError: argument of type: + DecorateInfo( + unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref_executor', + dtypes=(torch.complex32, torch.complex64, torch.complex128,) + ), + # Reference result was farther (0.7433461727239705) from the precise + # computation than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.complex32,), device_type="cuda" + ), + # Reference result was farther (0.7433461727239705) from the precise + # computation than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.complex32,), device_type="cuda" + ), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.div", + torch_opinfo_name="div", + torch_opinfo_variant_name="trunc_rounding", + # https://github.com/pytorch/pytorch/issues/76944 + supports_two_python_scalars=True, + supports_one_python_scalar=True, + decorators=( + # See https://github.com/pytorch/pytorch/issues/111126 + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.div", + torch_opinfo_name="div", + torch_opinfo_variant_name="floor_rounding", + # https://github.com/pytorch/pytorch/issues/76944 + supports_two_python_scalars=True, + supports_one_python_scalar=True, + decorators=( + # See https://github.com/pytorch/pytorch/issues/111126 + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.eq", + torch_opinfo_name="eq", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.float_power", + torch_opinfo_name="float_power", + skips=( + # Test doesn't account for float -> double type promotion + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + # Complex values error with: Greatest absolute difference: nan at index + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=[torch.complex64, torch.complex128]), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_large_values', + dtypes=[torch.complex64, torch.complex128]), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_extremal_values', + dtypes=[torch.complex64, torch.complex128]), + # TypeError: Cannot convert a MPS Tensor to float64 dtype + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out_warning', device_type='mps'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.logaddexp", + torch_opinfo_name="logaddexp", + skips=( + # failure due to mismatch in edge cases, which boils down to what torch.exp(inf + infj) should be + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + ), + ), + PythonRefInfo( + "_refs.logaddexp2", + torch_opinfo_name="logaddexp2", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.floor_divide", + torch_opinfo_name="floor_divide", + rhs_make_tensor_kwargs=dict(exclude_zero=True), + # https://github.com/pytorch/pytorch/issues/76944 + supports_two_python_scalars=True, + supports_one_python_scalar=True, + # bfloat16 floor_divide compared with a float32 reference works inconsistently + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.bfloat16,)), + # bfloat16 floor_divide compared with a float32 reference works inconsistently + DecorateInfo(unittest.skip('Skipped!'), 'TestBinaryUfuncs', + dtypes=(torch.bfloat16,)), + # int8 floor divide has different results for -128 // -1 vs. NumPy + DecorateInfo(unittest.skip('Skipped!'), 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.int8,)), + # The following tests fails on some jobs + DecorateInfo(unittest.skip('Skipped!'), 'TestBinaryUfuncs', + 'test_reference_numerics_extremal_values', + dtypes=(torch.float16,)), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-3, rtol=5e-3)}), + 'TestBinaryUfuncs', 'test_reference_numerics'), + # FIXME output 0: meta disagrees with real impl + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_binary_ufuncs_mixed_dtype'), + # The following dtypes did not work in forward but are listed by the OpInfo: {torch.bool}. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + # RuntimeError: torch.bool not supported for floor_divide + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps', dtypes=(torch.bool,)), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.fmax", + torch_opinfo_name="fmax", + supports_rhs_python_scalar=False, + ), + ElementwiseBinaryPythonRefInfo( + "_refs.fmin", + torch_opinfo_name="fmin", + supports_rhs_python_scalar=False, + ), + ElementwiseBinaryPythonRefInfo( + "_refs.fmod", + torch_opinfo_name="fmod", + rhs_make_tensor_kwargs={'exclude_zero': True}, + supports_rhs_python_scalar=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref', + dtypes=(torch.bfloat16,), device_type='cpu'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.bfloat16,), device_type='cpu'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_contig_vs_every_other', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_non_contig', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.uint8,)), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.gcd", + torch_opinfo_name="gcd", + skips=( + DecorateInfo(unittest.expectedFailure, + 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.int8,)), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.ge", + torch_opinfo_name="ge", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.gt", + torch_opinfo_name="gt", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.heaviside", + torch_opinfo_name="heaviside", + supports_rhs_python_scalar=False, + skips=( + # PyTorch's heaviside does not appear to propagate NaNs + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics_extremal_values'), + # NotImplementedError: The operator 'aten::heaviside.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.hypot", + torch_opinfo_name="hypot", + supports_rhs_python_scalar=False, + ), + ElementwiseBinaryPythonRefInfo( + "_refs.igamma", + torch_opinfo_name="igamma", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.igammac", + torch_opinfo_name="igammac", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.isclose", + torch_opinfo_name="isclose", + skips=( + # Intentional xfail -- isclose does not type promote + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_binary_ufuncs_mixed_dtype'), + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics_extremal_values'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.lcm", + torch_opinfo_name="lcm", + skips=( + # The operator 'aten::lcm.out' is not currently implemented for the MPS device. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.le", + torch_opinfo_name="le", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.logical_and", + torch_opinfo_name="logical_and", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.logical_not", + torch_opinfo_name="logical_not", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.logical_or", + torch_opinfo_name="logical_or", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.logical_xor", + torch_opinfo_name="logical_xor", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.lt", + torch_opinfo_name="lt", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.maximum", + torch_opinfo_name="maximum", + skips=( + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.minimum", + torch_opinfo_name="minimum", + skips=( + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.mul", + torch_opinfo_name="mul", + # https://github.com/pytorch/pytorch/issues/76944 + supports_two_python_scalars=True, + supports_one_python_scalar=True, + skips=( + # Reference result was farther (0.0) from the precise computation + # than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor', + dtypes=(torch.complex32,), + ), + # Reference result was farther (0.0) from the precise computation + # than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.complex32,), device_type='cuda' + ), + # Reference result was farther (0.0) from the precise computation + # than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.complex32,), device_type='cuda' + ), + # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.complex32,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.complex32,) + ), + ) + ), + ElementwiseBinaryPythonRefInfo( + "_refs.ne", + torch_opinfo_name="ne", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.nextafter", + torch_opinfo_name="nextafter", + ), + ElementwiseBinaryPythonRefInfo( + "_refs.pow", + torch_opinfo_name="pow", + decorators=( + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-4, rtol=1.3e-05)}), + 'TestBinaryUfuncs', 'test_reference_numerics'), + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-4, rtol=1.3e-05), + torch.complex128: tol(atol=1e-4, rtol=1.3e-05)}), + 'TestBinaryUfuncs', 'test_scalar_support'), + ), + skips=( + # Reference result was farther (inf) from the precise + # computation than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor', + dtypes=(torch.complex32,), + ), + # Reference result was farther (inf) from the precise + # computation than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.complex32,), device_type="cuda" + ), + # Reference result was farther (inf) from the precise + # computation than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.complex32,), device_type="cuda" + ), + # Skipping integers because they are being raised to negative powers causing an error + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=[torch.int8, torch.int16, torch.int32, torch.int64]), + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', + 'test_reference_numerics_large_values', + dtypes=[torch.int16, torch.int32, torch.int64]), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics', + dtypes=(torch.complex32,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_large_values', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_extremal_values', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + # torch._subclasses.fake_tensor.MetadataMismatchError: Dtypes torch.int64 and torch.bool are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.bool,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.bool,) + ), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.remainder", + torch_opinfo_name="remainder", + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref', + dtypes=(torch.bfloat16,), device_type='cpu'), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.bfloat16,), device_type='cpu'), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics', + dtypes=(torch.bfloat16,)), + DecorateInfo(unittest.skip("Skipped!"), 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.uint8,)), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.rsub", + torch_opinfo_name="rsub", + # https://github.com/pytorch/pytorch/issues/76944 + skips=( + # Reference result was farther (nan) from the precise computation than + # the torch result was (nan)! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.chalf,), device_type='cpu'), + # Reference result was farther (nan) from the precise computation than + # the torch result was (nan)! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.chalf,), device_type='cpu'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.sub", + torch_opinfo_name="sub", + # https://github.com/pytorch/pytorch/issues/76944 + supports_two_python_scalars=True, + supports_one_python_scalar=True, + decorators=( + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-2, rtol=0), + torch.bfloat16: tol(atol=1e-5, rtol=5e-3), + torch.complex32: tol(atol=1e-5, rtol=1e-3)}), + 'TestBinaryUfuncs', 'test_reference_numerics'), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=1e-2, rtol=0)}), + 'TestCommon', 'test_complex_half_reference_testing', device_type='cpu'), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=5e-3, rtol=0)}), + 'TestDecomp', 'test_comprehensive', device_type='cpu'), + DecorateInfo( + toleranceOverride({torch.chalf: tol(atol=5e-3, rtol=0)}), + 'TestDecomp', 'test_quick', device_type='cpu'), + ), + skips=( + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics', + dtypes=(torch.uint8,)), + DecorateInfo(unittest.skip("Skipped!"), + 'TestBinaryUfuncs', + 'test_reference_numerics_small_values', + dtypes=(torch.uint8,)), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.true_divide", + torch_opinfo_name="true_divide", + # https://github.com/pytorch/pytorch/issues/76944 + supports_two_python_scalars=True, + supports_one_python_scalar=True, + skips=( + # Reference result was farther (0.7433461727239705) from the precise + # computation than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_executor', + dtypes=(torch.complex32,), + ), + # Reference result was farther (0.7433461727239705) from the precise + # computation than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + dtypes=(torch.complex32,), device_type="cuda" + ), + # Reference result was farther (0.7433461727239705) from the precise + # computation than the torch result was (nan)! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.complex32,), device_type="cuda" + ), + ), + ), + # + # Elementwise Ternary Reference OpInfos + # + PythonRefInfo( + "_refs.addcdiv", + torch_opinfo_name="addcdiv", + ), + PythonRefInfo( + "_refs.addcmul", + torch_opinfo_name="addcmul", + skips=( + # Reference result was farther (1.3343989849090576e-05) + # from the precise computation than the torch result + # was (9.592622518539429e-06)! + # FIXME: enable dtype-based tolerances in test_ops.py:TestCommon._ref_test_helper + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref', + dtypes=(torch.float16,), device_type="cpu"), + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref_torch_fallback', + dtypes=(torch.float16,), device_type="cpu"), + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=( + torch.uint8, torch.int8, torch.int64, torch.int32, + torch.int16, + ) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', dtypes=( + torch.uint8, torch.int8, torch.int64, torch.int32, + torch.int16, + ) + ), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.clamp_min", + torch_opinfo_name="clamp_min", + skips=( + # test error disabled since rhs non-tensor python scalar is supported + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + ), + ), + ElementwiseBinaryPythonRefInfo( + "_refs.clamp_max", + torch_opinfo_name="clamp_max", + skips=( + # test error disabled since rhs non-tensor python scalar is supported + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + ), + ), + PythonRefInfo( + "_refs.clamp", + torch_opinfo_name="clamp", + ), + PythonRefInfo( + "_refs.nn.functional.triplet_margin_loss", + torch_opinfo_name="nn.functional.triplet_margin_loss", + supports_out=False, + # TODO: Uses minimum and clamp + skips=( + # AssertionError: Tensor-likes are not close! + # Greatest absolute difference: 6.103515625e-05 at index (4,) (up to 1e-05 allowed) + # Greatest relative difference: 8.519846983548175e-06 at index (4,) (up to 1.3e-06 allowed) + DecorateInfo(unittest.skip("Skipped!"), 'TestCommon', 'test_python_ref', + dtypes=(torch.uint8,), device_type="cpu"), + ) + ), + ElementwiseBinaryPythonRefInfo( + "_refs.xlogy", + torch_opinfo_name="xlogy", + supports_one_python_scalar=True, + ), + # + # Elementwise Binary Special OpInfos + # + ElementwiseBinaryPythonRefInfo( + "_refs.special.xlog1py", + torch_opinfo_name="special.xlog1py", + supports_one_python_scalar=True, + ), + # + # Data Conversion & Data Movement Opinfos + # + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.bfloat16", + torch_opinfo_name="bfloat16", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.bool", + torch_opinfo_name="bool", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.byte", + torch_opinfo_name="byte", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + skips=( + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.char", + torch_opinfo_name="char", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + skips=( + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + ) + ), + ElementwiseBinaryPythonRefInfo( + "_refs._conversions.complex", + torch_opinfo_name="complex", + error_inputs_func=partial(error_inputs_complex, is_ref=True), + skips=( + # Tests don't account for complex's type promotion semantics + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_binary_ufuncs_mixed_dtype'), + # RuntimeError: Expected object of scalar type torch.float32 but got scalar type torch.float16 for second argument + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors', device_type='mps', + ), + ) + ), + ElementwiseBinaryPythonRefInfo( + "_refs._conversions.polar", + torch_opinfo_name="polar", + skips=( + # Tests don't account for complex's type promotion semantics + DecorateInfo(unittest.expectedFailure, 'TestBinaryUfuncs', 'test_type_promotion'), + DecorateInfo(unittest.expectedFailure, 'TestMeta', 'test_binary_ufuncs_mixed_dtype'), + # Seems to fail on M2 but not M4: + # Exception: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo( + unittest.skip("Platform-specific error"), 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.float16,) + ), + DecorateInfo( + unittest.skip("Platform-specific error"), 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.float16,) + ), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.double", + torch_opinfo_name="double", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + skips=( + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.float", + torch_opinfo_name="float", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.half", + torch_opinfo_name="half", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.int", + torch_opinfo_name="int", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + skips=( + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.long", + torch_opinfo_name="long", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + skips=( + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.short", + torch_opinfo_name="short", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + skips=( + DecorateInfo(unittest.skip('Overflow when downcasting signed type is undefined'), 'TestCommon', 'test_compare_cpu'), + ) + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.chalf", + torch_opinfo_name="chalf", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.cfloat", + torch_opinfo_name="cfloat", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + ), + ElementwiseUnaryPythonRefInfo( + "_refs._conversions.cdouble", + torch_opinfo_name="cdouble", + # TODO: If self already has the correct dtype and device, then self is + # returned ignoring memory_format. + # https://github.com/pytorch/pytorch/issues/86558 + validate_view_consistency=False, + skips=( + # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype + DecorateInfo(unittest.expectedFailure, 'TestCommon', device_type='mps'), + ), + ), + PythonRefInfo( + "_refs.clone", + torch_opinfo_name="clone", + ), + # + # View & Shape OpInfos + # + PythonRefInfo( + "_refs.alias_copy", + torch_opinfo_name="alias_copy", + supports_out=True, + ), + PythonRefInfo( + "_refs.atleast_1d", + torch_opinfo_name="atleast_1d", + validate_view_consistency=False, + ), + PythonRefInfo( + "_refs.atleast_2d", + torch_opinfo_name="atleast_2d", + validate_view_consistency=False, + ), + PythonRefInfo( + "_refs.atleast_3d", + torch_opinfo_name="atleast_3d", + validate_view_consistency=False, + ), + PythonRefInfo( + "_refs.as_strided", + torch_opinfo_name="as_strided", + # FIXME: doesn't support chalf + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + skips=( + # cloned_mutable_input.is_same(returned_output) INTERNAL ASSERT FAILED + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_neg_conj_view'), + ), + ), + PythonRefInfo( + "_refs.as_strided_copy", + torch_opinfo_name="as_strided_copy", + supports_out=True, + # FIXME: doesn't support chalf + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + skips=( + # cloned_mutable_input.is_same(returned_output) INTERNAL ASSERT FAILED + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_neg_conj_view'), + # The view function this decompose into does not have a ref + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_python_ref"), + ), + ), + PythonRefInfo( + "_refs.as_strided", + torch_opinfo_name="as_strided", + torch_opinfo_variant_name="partial_views", + # FIXME: doesn't support chalf + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + skips=( + # cloned_mutable_input.is_same(returned_output) INTERNAL ASSERT FAILED + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Errors when storage_offset is included"), 'TestMathBits', 'test_neg_conj_view'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_compare_cpu'), + ), + ), + PythonRefInfo( + "_refs.as_strided_scatter", + torch_opinfo_name="as_strided_scatter", + # returns a view of an intermediate tensor (as_strided) + validate_view_consistency=False, + ), + PythonRefInfo( + "_refs.block_diag", + torch_opinfo_name="block_diag", + skips=( + # RuntimeError: Failed to create function state object for: cat_int32_t_float_float2 + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', + device_type='mps', dtypes=(torch.complex64, torch.complex32) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.complex64, torch.complex32) + ), + ), + ), + PythonRefInfo( + "_refs.broadcast_shapes", + torch_opinfo_name="broadcast_shapes", + ), + PythonRefInfo( + "_refs.broadcast_tensors", + torch_opinfo_name="broadcast_tensors", + ), + PythonRefInfo( + "_refs.broadcast_to", + torch_opinfo_name="broadcast_to", + ), + PythonRefInfo( + "_refs.cat", + torch_opinfo_name="cat", + skips=( + # FIXME: AssertionError: RuntimeError not raised + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + # RuntimeError: Failed to create function state object for: cat_int32_t_* + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.complex64, torch.complex32) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.complex64, torch.complex32) + ), + ), + ), + PythonRefInfo( + "_refs.chunk", + torch_opinfo_name="chunk", + ), + PythonRefInfo( + "_refs.column_stack", + torch_opinfo_name="column_stack", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.conj", + torch_opinfo_name="conj", + ), + PythonRefInfo( + "_refs.constant_pad_nd", + torch_opinfo_name="constant_pad_nd", + ), + PythonRefInfo( + "_refs.contiguous", + torch_opinfo_name="contiguous", + ), + ElementwiseUnaryPythonRefInfo( + "_refs.deg2rad", + torch_opinfo_name="deg2rad", + decorators=(precisionOverride({torch.bfloat16: 7e-1, + torch.float16: 7e-1}),), + ), + PythonRefInfo( + "_refs.dsplit", + torch_opinfo_name="dsplit", + ), + PythonRefInfo( + "_refs.diag", + torch_opinfo_name="diag", + ), + PythonRefInfo( + "_refs.diagonal", + torch_opinfo_name="diagonal", + ), + PythonRefInfo( + "_refs.diagonal_copy", + torch_opinfo_name="diagonal_copy", + supports_out=True, + ), + PythonRefInfo( + "_refs.diagonal_scatter", + torch_opinfo_name="diagonal_scatter", + supports_out=True, + # returns a view of an intermediate tensor (as_strided) + validate_view_consistency=False, + ), + PythonRefInfo( + "_refs.diag_embed", + torch_opinfo_name="diag_embed", + supports_out=True, + skips=( + # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype. + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.complex32,) + ), + ), + ), + PythonRefInfo( + "_refs.dstack", + torch_opinfo_name="dstack", + skips=( + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + ), + ), + PythonRefInfo( + "_refs.expand", + torch_opinfo_name="expand", + ), + PythonRefInfo( + "_refs.expand_as", + torch_opinfo_name="expand_as", + ), + PythonRefInfo( + "_refs.expand_copy", + torch_opinfo_name="expand_copy", + supports_out=True, + ), + PythonRefInfo( + "_refs.flatten", + torch_opinfo_name="flatten", + ), + PythonRefInfo( + "_refs.flip", + torch_opinfo_name="flip", + ), + PythonRefInfo( + "_refs.fliplr", + torch_opinfo_name="fliplr", + ), + PythonRefInfo( + "_refs.flipud", + torch_opinfo_name="flipud", + ), + PythonRefInfo( + "_refs.hstack", + torch_opinfo_name="hstack", + skips=( + # https://github.com/pytorch/pytorch/issues/78613 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + ), + ), + PythonRefInfo( + "_refs.narrow", + torch_opinfo_name="narrow", + error_inputs_func=partial(error_inputs_narrow_narrow_copy, is_narrow=True, is_ref=True), + ), + PythonRefInfo( + "_refs.narrow_copy", + torch_opinfo_name="narrow_copy", + supports_out=True, + error_inputs_func=partial(error_inputs_narrow_narrow_copy, is_narrow=False, is_ref=True), + skips=( + # The view function this decompose into does not have a ref + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_python_ref"), + ), + ), + PythonRefInfo( + "_refs.nn.functional.group_norm", + torch_opinfo_name="nn.functional.group_norm", + validate_view_consistency=False, + skips=( + # RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.int32, torch.int16, torch.int8, torch.uint8) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', + dtypes=(torch.int32, torch.int16, torch.int8, torch.uint8) + ), + ) + ), + PythonRefInfo( + "_refs.native_layer_norm", + torch_opinfo_name="native_layer_norm", + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCommon", "test_python_ref", + device_type="cpu", dtypes=(torch.float32,)), + DecorateInfo(unittest.skip("Skipped!"), "TestCommon", "test_python_ref_torch_fallback", + device_type="cpu", dtypes=(torch.float32,)), + # TypeError: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(torch.float32,)), + # Exception: Dtypes torch.float32 and torch.float16 are not equal! + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.float16, torch.bfloat16) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.float16, torch.bfloat16) + ), + ), + ), + PythonRefInfo( + "_refs.permute", + torch_opinfo_name="permute", + ), + PythonRefInfo( + "_refs.permute_copy", + torch_opinfo_name="permute_copy", + supports_out=True, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.rad2deg", + torch_opinfo_name="rad2deg", + decorators=(precisionOverride({torch.bfloat16: 7e-1, + torch.float16: 7e-1}),), + ), + PythonRefInfo( + "_refs.ravel", + torch_opinfo_name="ravel", + ), + PythonRefInfo( + "_refs.renorm", + torch_opinfo_name="renorm", + skips=( + # RuntimeError: Failed to create function state object for: renorm_float2 + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.complex64,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.complex64,) + ), + ), + ), + PythonRefInfo( + "_refs.repeat", + torch_opinfo_name="repeat", + validate_view_consistency=False, + ), + PythonRefInfo( + "_refs.reshape", + torch_opinfo_name="reshape", + ), + PythonRefInfo( + "_refs.reshape_as", + torch_opinfo_name="reshape_as", + ), + PythonRefInfo( + "_refs.roll", + torch_opinfo_name="roll", + validate_view_consistency=False, + ), + PythonRefInfo( + "_refs.rot90", + torch_opinfo_name="rot90", + validate_view_consistency=False, + ), + PythonRefInfo( + "_refs.select_scatter", + torch_opinfo_name="select_scatter", + ), + PythonRefInfo( + "_refs.stack", + torch_opinfo_name="stack", + validate_view_consistency=False, + ), + PythonRefInfo( + "_refs.squeeze", + torch_opinfo_name="squeeze", + ), + PythonRefInfo( + "_refs.squeeze_copy", + torch_opinfo_name="squeeze_copy", + supports_out=True, + ), + PythonRefInfo( + "_refs.squeeze", + torch_opinfo_name="squeeze", + torch_opinfo_variant_name="multiple", + ), + PythonRefInfo( + "_refs.tensor_split", + torch_opinfo_name="tensor_split", + skips=( + # RuntimeError: no _refs support for torch.Tensor.tolist + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref'), + ), + ), + PythonRefInfo( + "_refs.hsplit", + torch_opinfo_name="hsplit", + ), + PythonRefInfo( + "_refs.vsplit", + torch_opinfo_name="vsplit", + ), + PythonRefInfo( + "_refs.dot", + torch_opinfo_name="dot", + error_inputs_func=partial(error_inputs_dot_vdot, is_ref=True), + # .conj() does not set ._is_view() correctly in ATen + validate_view_consistency=False, + skips=( + # RuntimeError: no _refs support for torch.Tensor.is_conj + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', dtypes=[torch.complex64, torch.complex128]), + ), + ), + PythonRefInfo( + "_refs.vdot", + torch_opinfo_name="vdot", + error_inputs_func=partial(error_inputs_dot_vdot, is_ref=True), + # .conj() does not set ._is_view() correctly in ATen + validate_view_consistency=False, + skips=( + # RuntimeError: no _refs support for torch.Tensor.is_conj + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', dtypes=[torch.complex64, torch.complex128]), + # Exception: Conj mismatch! is_conj is set to True and False + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=[torch.complex64] + ), + ), + ), + PythonRefInfo( + "_refs.transpose", + torch_opinfo_name="transpose", + ), + PythonRefInfo( + "_refs.transpose_copy", + torch_opinfo_name="transpose_copy", + supports_out=True, + ), + PythonRefInfo( + "_refs.t", + torch_opinfo_name="t", + ), + PythonRefInfo( + "_refs.t_copy", + torch_opinfo_name="t_copy", + supports_out=True, + ), + PythonRefInfo( + "_refs.T", + torch_opinfo_name="T", + error_inputs_func=partial(error_inputs_T, has_ndims_error=True), + ), + PythonRefInfo( + "_refs.unbind_copy", + torch_opinfo_name="unbind_copy", + ), + PythonRefInfo( + "_refs.unfold", + torch_opinfo_name="unfold", + ), + PythonRefInfo( + "_refs.unfold_copy", + torch_opinfo_name="unfold_copy", + supports_out=True, + ), + PythonRefInfo( + "_refs.unsqueeze", + torch_opinfo_name="unsqueeze", + ), + PythonRefInfo( + "_refs.unsqueeze_copy", + torch_opinfo_name="unsqueeze_copy", + supports_out=True, + ), + PythonRefInfo( + "_refs.view", + torch_opinfo_name="view", + ), + PythonRefInfo( + "_refs.view_as", + torch_opinfo_name="view_as", + ), + PythonRefInfo( + "_refs.view_copy", + torch_opinfo_name="view_copy", + supports_out=True, + ), + PythonRefInfo( + "_refs.vstack", + torch_opinfo_name="vstack", + skips=( + # https://github.com/pytorch/pytorch/issues/78613 + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + ), + ), + PythonRefInfo( + "_refs.unflatten", + torch_opinfo_name="unflatten", + ), + PythonRefInfo( + "_refs.unbind", + torch_opinfo_name="unbind", + ), + # + # Reduction Reference OpInfos + # + ReductionPythonRefInfo( + "_refs.all", + torch_opinfo_name="all", + skips=( + # FIXME: uint8 input returns uint8 instead of bool + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_result_dtype', + dtypes=[torch.uint8]), + ), + ), + ReductionPythonRefInfo( + "_refs.amax", + torch_opinfo_name="amax", + error_inputs_func=partial(error_inputs_aminmax_amax_amin, is_ref=True), + skips=( + # FIXME: reduces all dimensions when dim=[] + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.int64]), + # RuntimeError: MPS supports tensors with dimensions <= 16, but got 65. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors', device_type='mps'), + ), + ), + ReductionPythonRefInfo( + "_refs.amin", + torch_opinfo_name="amin", + error_inputs_func=partial(error_inputs_aminmax_amax_amin, is_ref=True), + skips=( + # FIXME: reduces all dimensions when dim=[] + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.int64]), + # RuntimeError: MPS supports tensors with dimensions <= 16, but got 65. + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors', device_type='mps'), + ), + ), + ReductionPythonRefInfo( + "_refs.any", + torch_opinfo_name="any", + skips=( + # FIXME: uint8 input returns uint8 instead of bool + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_result_dtype', + dtypes=[torch.uint8]), + ), + ), + ReductionPythonRefInfo( + "_refs.count_nonzero", + torch_opinfo_name="count_nonzero", + skips=( + # FIXME: count_nonzero does not accept keepdim kwarg + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', + 'test_dim_default_keepdim'), + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_dim_none_keepdim'), + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_dim_single_keepdim'), + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_dim_multi_keepdim'), + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', + 'test_dim_multi_unsorted_keepdim'), + # FIXME: dim=[] reduces all dimensions + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + ), + ), + ReductionPythonRefInfo( + "_refs.mean", + torch_opinfo_name="mean", + supports_out=True, + error_inputs_func=partial(error_inputs_mean, is_ref=True), + skips=( + # FIXME: reduces all dimensions when dim=[] + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.complex128]), + ), + ), + ReductionPythonRefInfo( + "_refs.std", + torch_opinfo_name="std", + supports_out=True, + skips=( + # FIXME: reduces all dimensions when dim=[] + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: improve precision + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input', + dtypes=(torch.float16,)), + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', + 'test_ref_duplicate_values', + dtypes=(torch.float16,)), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.float64]), + # Exception: Dtypes torch.float32 and torch.complex64 are not equal! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(torch.complex64,)), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.complex64,) + ), + ), + ), + # std_mean and var_mean are not ReductionInfos + PythonRefInfo( + "_refs.std_mean", + torch_opinfo_name="std_mean", + skips=( + # Exception: Dtypes torch.float32 and torch.complex64 are not equal! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(torch.complex64,)), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.complex64,) + ), + # RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', + dtypes=(torch.uint8, torch.bool, torch.int8, torch.int16, torch.int32) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.uint8, torch.bool, torch.int8, torch.int16, torch.int32) + ), + ), + ), + ReductionPythonRefInfo( + "_refs.sum", + torch_opinfo_name="sum", + supports_out=True, + skips=( + # FIXME: doesn't test out behavior properly for this operator + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # FIXME: mean reduces all dimensions when dim=[] + DecorateInfo(unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty'), + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: improve precision + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input', + dtypes=[torch.float16]), + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', + 'test_ref_duplicate_values', + dtypes=[torch.float16]), + DecorateInfo( + unittest.skip("Skipped!"), 'TestOperators', 'test_reduction_all', + dtypes=[torch.float32]), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.complex128]), + + ), + ), + PythonRefInfo( + "_refs.cumsum", + torch_opinfo_name="cumsum", + supports_out=True, + skips=( + # doesn't test out behavior properly for this operator + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + ), + ), + PythonRefInfo( + "_refs.cumprod", + torch_opinfo_name="cumprod", + supports_out=True, + skips=( + # doesn't test out behavior properly for this operator + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + ), + ), + PythonRefInfo( + "_refs.sum_to_size", + torch_opinfo_name="sum_to_size", + validate_view_consistency=False, + ), + ReductionPythonRefInfo( + "_refs.prod", + torch_opinfo_name="prod", + supports_out=True, + supports_multiple_dims=True, + skips=( + # FIXME: doesn't test out behavior properly for this operator + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_out'), + # FIXME: reduces all dimensions when dim=[] + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: improve precision + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input', + dtypes=[torch.float16, torch.complex64]), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip('Skipped!'), 'TestReductions', 'test_ref_small_input', + device_type='xpu', + dtypes=[torch.int64, torch.int8, torch.int16, torch.int32, torch.complex128]), + ), + ), + ReductionPythonRefInfo( + "_refs.var", + torch_opinfo_name="var", + supports_out=True, + skips=( + # FIXME: reduces all dimensions when dim=[] + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty'), + DecorateInfo( + unittest.expectedFailure, 'TestReductions', 'test_dim_empty_keepdim'), + # FIXME: improve precision + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_ref_small_input'), + DecorateInfo( + unittest.skip("Skipped!"), 'TestReductions', 'test_ref_duplicate_values'), + # torch._subclasses.fake_tensor.MetadataMismatchError: Dtypes torch.float32 and torch.complex64 are not equal! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(torch.complex64,)), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.complex64,) + ), + ), + ), + PythonRefInfo( + "_refs.var_mean", + torch_opinfo_name="var_mean", + validate_view_consistency=False, + skips=( + # torch._subclasses.fake_tensor.MetadataMismatchError: Dtypes torch.float32 and torch.complex64 are not equal! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', dtypes=(torch.complex64,)), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.complex64,) + ), + # RuntimeError: mean(): could not infer output dtype. Input dtype must be either a floating point or complex dtype + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps', + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool) + ), + ), + ), + # + # Linear Algebra Operators + # + PythonRefInfo( + "_refs.addr", + torch_opinfo_name="addr", + decorators=( + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref',), + # RuntimeError: MPS device does not support addr for non-float input + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=( + torch.uint8, torch.int8, torch.int64, torch.int32, + torch.int16, torch.complex64, torch.bool, + ) + ), + ), + ), + PythonRefInfo( + "_refs.trace", + torch_opinfo_name="trace", + ), + PythonRefInfo( + "_refs.norm", + torch_opinfo_name="norm", + supports_out=True, + # Uses vector_norm inside and vector_norm is affected by + # https://github.com/pytorch/pytorch/issues/77216 + validate_view_consistency=False, + ), + # + # Tensor Creation Reference OpInfos + # + PythonRefInfo( + "_refs.empty", + torch_opinfo_name="empty", + skips=( + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_python_ref'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestMathBits', + 'test_conj_view'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestMathBits', + 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestMathBits', + 'test_neg_view'), + # FIXME: shouldn't check empty results + DecorateInfo(unittest.skip("Can't check result for empty"), 'TestCommon', 'test_python_ref_executor'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + ), + ), + PythonRefInfo( + "_refs.empty_like", + torch_opinfo_name="empty_like", + skips=( + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_python_ref'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestMathBits', + 'test_conj_view'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestMathBits', + 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestMathBits', + 'test_neg_view'), + # FIXME: should not compare results of empty_like + DecorateInfo(unittest.skip("Can't check result for empty_like"), 'TestCommon', 'test_python_ref_executor'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + ), + ), + PythonRefInfo( + "_refs.randn", + torch_opinfo_name="randn", + op=lambda *args, **kwargs: wrapper_set_seed(refs.randn, *args, **kwargs), + skips=( + # see https://github.com/pytorch/pytorch/issues/85121 + DecorateInfo(unittest.skip("make_traced() doesn't set seed properly!"), + 'TestCommon', + 'test_python_ref_executor'), + # These tests expect the input to be a tensor or a sequence of tensors + DecorateInfo(unittest.skip("Test expects tensor input"), "TestCommon", "test_noncontiguous_samples"), + DecorateInfo(unittest.skip("Test expects tensor input"), 'TestMathBits', 'test_neg_view'), + DecorateInfo(unittest.skip("Test expects tensor input"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Test expects tensor input"), 'TestMathBits', 'test_neg_conj_view'), + ), + ), + PythonRefInfo( + "_refs.eye", + torch_opinfo_name="eye", + skips=( + # skip these tests since we have non tensor input + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_conj_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Skipped!"), 'TestMathBits', 'test_neg_view'), + # TypeError: Trying to convert Float8_* to the MPS backend but it does not have support for that dtype. + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref', device_type='mps', + dtypes=(torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', device_type='mps', + dtypes=(torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', device_type='mps', + dtypes=(torch.float8_e4m3fn, torch.float8_e4m3fnuz, torch.float8_e5m2, torch.float8_e5m2fnuz,) + ), + ), + ), + PythonRefInfo( + "_refs.new_empty", + torch_opinfo_name="new_empty", + skips=( + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_python_ref'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_out'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestCommon', + 'test_out_warning'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestMathBits', + 'test_conj_view'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestMathBits', + 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Expected: empty is not comparable"), + 'TestMathBits', + 'test_neg_view'), + # FIXME: should not compare results of empty_like + DecorateInfo(unittest.skip("Can't check result for new_empty"), 'TestCommon', 'test_python_ref_executor'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + ), + ), + PythonRefInfo( + "_refs.new_empty_strided", + torch_opinfo_name="new_empty_strided", + skips=( + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestCommon', + 'test_python_ref'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestMathBits', + 'test_conj_view'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestMathBits', + 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestMathBits', + 'test_neg_view'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestCommon', + 'test_python_ref_executor'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + + ), + ), + PythonRefInfo( + "_refs.empty_strided", + torch_opinfo_name="empty_strided", + skips=( + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestCommon', + 'test_python_ref'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestCommon', + 'test_python_ref_torch_fallback'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestMathBits', + 'test_conj_view'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestMathBits', + 'test_neg_conj_view'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestMathBits', + 'test_neg_view'), + DecorateInfo(unittest.skip("Expected: empty_strided is not comparable"), + 'TestCommon', + 'test_python_ref_executor'), + DecorateInfo(unittest.skip('output is non-deterministic'), 'TestCommon', 'test_compare_cpu'), + ), + ), + PythonRefInfo( + "_refs.new_full", + torch_opinfo_name="new_full", + ), + PythonRefInfo( + "_refs.new_ones", + torch_opinfo_name="new_ones", + ), + PythonRefInfo( + "_refs.new_zeros", + torch_opinfo_name="new_zeros", + ), + # + # Conditional Reference OpInfos + # + PythonRefInfo( + "_refs.masked_fill", + torch_opinfo_name="masked_fill", + skips=( + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + # Seems to fail on M4 but not M2: + # TypeError: Trying to convert ComplexDouble to the MPS backend but it does not have support for that dtype. + DecorateInfo( + unittest.skip("Platform-specific error"), 'TestCommon', 'test_python_ref', + device_type='mps', dtypes=(torch.complex32,) + ), + ), + ), + PythonRefInfo( + "_refs.where", + torch_opinfo_name="where", + op=lambda self, condition, other: refs.where(condition, self, other), + supports_out=False, + skips=( + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors', device_type='cuda'), + # RuntimeError: Tensor on device mps:0 is not on the expected device cpu! + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors', device_type='mps'), + ), + ), + PythonRefInfo( + "_refs.index_select", + torch_opinfo_name="index_select", + # empty_strided + skips=( + # no _refs support for Tensor.__setitem__ + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref'), + # Sample out= with a stride of zero. This _out operation checks that the input has no + # inner overlap + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'),) + ), + PythonRefInfo( + "_refs.index_copy", + torch_opinfo_name="index_copy", + # empty_strided + skips=( + # no _refs support for Tensor.__setitem__ + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref'), + ), + ), + PythonRefInfo( + "_refs.index_add", + torch_opinfo_name="index_add", + # empty_strided + skips=( + # no _refs support for Tensor.__setitem__ + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref'), + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref_errors'), + ), + ), + PythonRefInfo( + "_refs.index_fill", + torch_opinfo_name="index_fill", + # empty_strided + skips=( + # no _refs support for Tensor.__setitem__ + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_python_ref'),) + ), + # + # Test-related functions + # + PythonRefInfo( + "_refs.allclose", + torch_opinfo_name="allclose", + ), + # + # Misc functions + # + PythonRefInfo( + "_refs.stft", + torch_opinfo_name="stft", + skips=[ + # RuntimeError: no _refs support for aten.pad + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref' + ), + # RuntimeError: Unsupported dtype Half + DecorateInfo(unittest.expectedFailure, 'TestCommon', 'test_dtypes', device_type='mps'), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_meta', + device_type='mps', dtypes=(torch.float16,) + ), + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref_torch_fallback', + device_type='mps', dtypes=(torch.float16,) + ), + ], + ), + PythonRefInfo( + "_refs.istft", + torch_opinfo_name="istft", + skips=[ + # RuntimeError: no _refs support for aten.unfold_backward + DecorateInfo( + unittest.expectedFailure, 'TestCommon', 'test_python_ref' + ), + DecorateInfo( + unittest.skip("Expected: unfold_backward() got an unexpected keyword argument 'input_sizes'"), + 'TestCommon', + 'test_python_ref_executor', + dtypes=(torch.complex64, torch.complex128), + ), + ], + ), + PythonRefInfo( + "_refs.view_as_complex", + torch_opinfo_name="view_as_complex", + ), + PythonRefInfo( + "_refs.split_with_sizes", + torch_opinfo_name="split_with_sizes", + ), +] +python_ref_db += opinfo.definitions.python_ref_db + +# Common operator groupings +ops_and_refs = op_db + python_ref_db +unary_ufuncs = [op for op in ops_and_refs if isinstance(op, UnaryUfuncInfo)] +binary_ufuncs = [op for op in ops_and_refs if isinstance(op, BinaryUfuncInfo)] +spectral_funcs = [op for op in ops_and_refs if isinstance(op, SpectralFuncInfo)] +sparse_unary_ufuncs = [op for op in op_db if isinstance(op, UnaryUfuncInfo) and op.supports_sparse] +sparse_csr_unary_ufuncs = [op for op in op_db if isinstance(op, UnaryUfuncInfo) and op.supports_sparse_csr] +sparse_reduction_ops = [op for op in op_db if isinstance(op, ReductionOpInfo) and op.supports_sparse] +shape_funcs = [op for op in ops_and_refs if isinstance(op, ShapeFuncInfo)] +reduction_ops = [op for op in ops_and_refs if isinstance(op, ReductionOpInfo)] +reference_filtered_ops = [op for op in reduction_ops if op.ref is not None] +reference_masked_ops = [op for op in reference_filtered_ops if op.name.startswith('masked.')] +sparse_masked_reduction_ops = [op for op in sparse_reduction_ops if op.name.startswith('masked.')] + +def index_variable(shape, max_indices, device=torch.device('cpu')): + if not isinstance(shape, tuple): + shape = (shape,) + return torch.testing.make_tensor(*shape, dtype=torch.long, device=device, low=0, high=max_indices) + +def gather_variable(shape, index_dim, max_indices, duplicate=False, device=torch.device('cpu')): + if len(shape) != 2: + raise AssertionError(f"Expected len(shape) == 2, got {len(shape)}") + if index_dim >= 2: + raise AssertionError(f"Expected index_dim < 2, got {index_dim}") + batch_dim = 1 - index_dim + index = torch.zeros(*shape, dtype=torch.long, device=device) + for i in range(shape[index_dim]): + index.select(index_dim, i).copy_( + torch.randperm(max_indices, device=device)[:shape[batch_dim]]) + if duplicate: + index.select(batch_dim, 0).copy_(index.select(batch_dim, 1)) + return index + +def bernoulli_scalar(): + return torch.tensor(0, dtype=torch.bool).bernoulli_() + +def mask_not_all_zeros(shape): + if len(shape) <= 0: + raise AssertionError(f"Expected len(shape) > 0, got {len(shape)}") + while True: + result = torch.randn(shape).gt(0) + if result.sum() > 0: + return result + +# Copied from functorch +def xfail(op_name, variant_name='', *, device_type=None, dtypes=None): + return (op_name, variant_name, device_type, dtypes, True) + + +def skip(op_name, variant_name='', *, device_type=None, dtypes=None): + return (op_name, variant_name, device_type, dtypes, False) + + +def skipOps(test_case_name, base_test_name, to_skip): + all_opinfos = op_db + for xfail in to_skip: + op_name, variant_name, device_type, dtypes, expected_failure = xfail + matching_opinfos = [o for o in all_opinfos + if o.name == op_name and o.variant_test_name == variant_name] + if len(matching_opinfos) < 1: + raise AssertionError(f"Couldn't find OpInfo for {xfail}") + for op in matching_opinfos: + decorators = list(op.decorators) + if expected_failure: + decorator = DecorateInfo(unittest.expectedFailure, + test_case_name, base_test_name, + device_type=device_type, dtypes=dtypes) + decorators.append(decorator) + else: + decorator = DecorateInfo(unittest.skip("Skipped!"), + test_case_name, base_test_name, + device_type=device_type, dtypes=dtypes) + decorators.append(decorator) + op.decorators = tuple(decorators) + + # This decorator doesn't modify fn in any way + def wrapped(fn): + return fn + return wrapped diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_mkldnn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_mkldnn.py new file mode 100644 index 0000000000000000000000000000000000000000..3e99f66820751fe754e0a4bb4f2ad0e8ab35e1b8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_mkldnn.py @@ -0,0 +1,113 @@ +# mypy: ignore-errors + +import contextlib +import functools +import inspect + +import torch + + +def bf32_is_not_fp32(): + if not torch.backends.mkldnn.is_available(): + return False + if not torch.ops.mkldnn._is_mkldnn_bf16_supported(): + return False + return True + + +def tf32_is_not_fp32(): + if not torch.backends.mkldnn.is_available(): + return False + if not torch.cpu._is_amx_fp16_supported(): + return False + return True + + +@contextlib.contextmanager +def reduced_f32_off(): + old_matmul_precision = torch.backends.mkldnn.matmul.fp32_precision + old_conv_precision = torch.backends.mkldnn.conv.fp32_precision + try: + torch.backends.mkldnn.matmul.fp32_precision = "ieee" + torch.backends.mkldnn.conv.fp32_precision = "ieee" + yield + finally: + torch.backends.mkldnn.matmul.fp32_precision = old_matmul_precision + torch.backends.mkldnn.conv.fp32_precision = old_conv_precision + + +@contextlib.contextmanager +def bf32_on(self, bf32_precision=1e-2): + old_matmul_precision = torch.backends.mkldnn.matmul.fp32_precision + old_conv_precision = torch.backends.mkldnn.conv.fp32_precision + old_precision = self.precision + try: + torch.backends.mkldnn.matmul.fp32_precision = "bf16" + torch.backends.mkldnn.conv.fp32_precision = "bf16" + self.precision = bf32_precision + yield + finally: + torch.backends.mkldnn.matmul.fp32_precision = old_matmul_precision + torch.backends.mkldnn.conv.fp32_precision = old_conv_precision + self.precision = old_precision + + +@contextlib.contextmanager +def tf32_on(self, tf32_precision=1e-5): + old_matmul_precision = torch.backends.mkldnn.matmul.fp32_precision + old_conv_precision = torch.backends.mkldnn.conv.fp32_precision + old_precision = self.precision + try: + torch.backends.mkldnn.matmul.fp32_precision = "tf32" + torch.backends.mkldnn.conv.fp32_precision = "tf32" + self.precision = tf32_precision + yield + finally: + torch.backends.mkldnn.matmul.fp32_precision = old_matmul_precision + torch.backends.mkldnn.conv.fp32_precision = old_conv_precision + self.precision = old_precision + + +# This is a wrapper that wraps a test to run this test three times, one with +# reduced_f32 OFF, the others with reduced_f32 ON (including bf32 ON and tf32 +# ON). When running with reduced_f32 ON, it will use reduced precision (bf16/ +# tf32) as specified by the argument. +def reduced_f32_on_and_off(bf32_precision=1e-2, tf32_precision=1e-5): + def with_reduced_f32_disabled(self, function_call): + with reduced_f32_off(): + function_call() + + def with_bf32_enabled(self, function_call): + with bf32_on(self, bf32_precision): + function_call() + + def with_tf32_enabled(self, function_call): + with tf32_on(self, tf32_precision): + function_call() + + def wrapper(f): + params = inspect.signature(f).parameters + arg_names = tuple(params.keys()) + + @functools.wraps(f) + def wrapped(*args, **kwargs): + kwargs.update(zip(arg_names, args, strict=False)) + cond = True + if "device" in kwargs: + cond = cond and (torch.device(kwargs["device"]).type == "cpu") + if "dtype" in kwargs: + cond = cond and (kwargs["dtype"] == torch.float) + bf32_cond = cond and bf32_is_not_fp32() + tf32_cond = cond and tf32_is_not_fp32() + if bf32_cond or tf32_cond: + with_reduced_f32_disabled(kwargs["self"], lambda: f(**kwargs)) + if bf32_cond: + with_bf32_enabled(kwargs["self"], lambda: f(**kwargs)) + if tf32_cond: + with_tf32_enabled(kwargs["self"], lambda: f(**kwargs)) + else: + f(**kwargs) + + return wrapped + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_modules.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_modules.py new file mode 100644 index 0000000000000000000000000000000000000000..3e5d4af0d2a5801f320426a2148cf13a6d154883 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_modules.py @@ -0,0 +1,4761 @@ +# mypy: ignore-errors + +import torch +import unittest +from copy import deepcopy +from enum import Enum +from functools import wraps, partial +from itertools import chain, product +import itertools +import math +import torch.nn.functional as F +from torch.nn.utils.rnn import pack_padded_sequence +from torch.testing import make_tensor +from torch.testing._internal.common_cuda import TEST_CUDNN +from torch.testing._internal.common_dtype import ( + floating_types, floating_and_complex_types_and, get_all_fp_dtypes) +from torch.testing._internal.common_device_type import ( + _TestParametrizer, _update_param_kwargs, expectedFailureMPS, toleranceOverride, tol, + precisionOverride, skipMeta, skipMPS) +from torch.testing._internal.common_methods_invocations import DecorateInfo +from torch.testing._internal.common_nn import ( + cosineembeddingloss_reference, cross_entropy_loss_reference, ctcloss_reference, + hingeembeddingloss_reference, huberloss_reference, kldivloss_reference, + marginrankingloss_reference, multimarginloss_reference, multilabelmarginloss_reference, + nllloss_reference, nlllossNd_reference, smoothl1loss_reference, softmarginloss_reference, get_reduction) +from torch.testing._internal.common_utils import ( + freeze_rng_state, skipIfMPS, GRADCHECK_NONDET_TOL, TEST_WITH_ROCM, IS_WINDOWS, + skipIfTorchDynamo, skipIfXpu) +from types import ModuleType +import operator + +# List of all namespaces containing modules to test. +MODULE_NAMESPACES: list[ModuleType] = [ + torch.nn.modules, + torch.ao.nn.qat.modules, + torch.ao.nn.quantizable.modules, + torch.ao.nn.quantized.modules, + torch.ao.nn.quantized.modules, +] + +# Modules that shouldn't be tested for one reason or another. +MODULES_TO_SKIP: set[type] = { + torch.nn.Module, # abstract base class + torch.nn.Container, # deprecated + torch.nn.NLLLoss2d, # deprecated + torch.ao.nn.quantized.MaxPool2d, # aliases to nn.MaxPool2d + torch.ao.nn.quantized.MaxPool2d, # aliases to nn.MaxPool2d +} + +# List of all module classes to test. +MODULE_CLASSES: list[type] = [*chain.from_iterable([ + [getattr(namespace, module_name) for module_name in namespace.__all__] # type: ignore[attr-defined] + for namespace in MODULE_NAMESPACES])] +MODULE_CLASSES = [cls for cls in MODULE_CLASSES if cls not in MODULES_TO_SKIP] + +# Dict of module class -> common name. Useful for making test names more intuitive. +# Example: torch.nn.modules.linear.Linear -> "nn.Linear" +MODULE_CLASS_NAMES: dict[type, str] = {} +for namespace in MODULE_NAMESPACES: + for module_name in namespace.__all__: # type: ignore[attr-defined] + module_cls = getattr(namespace, module_name) + namespace_name = namespace.__name__.replace('torch.', '').replace('.modules', '') + + # Deal with any aliases by preferring earlier names. + if module_cls not in MODULE_CLASS_NAMES: + MODULE_CLASS_NAMES[module_cls] = f'{namespace_name}.{module_name}' + + +# Specifies the modes (i.e. train, eval) to test over. +TrainEvalMode = Enum('TrainEvalMode', ('train_only', 'eval_only', 'train_and_eval')) + + +class modules(_TestParametrizer): + """ PROTOTYPE: Decorator for specifying a list of modules over which to run a test. """ + + def __init__(self, module_info_iterable, allowed_dtypes=None, + train_eval_mode=TrainEvalMode.train_and_eval, skip_if_dynamo=True): + self.module_info_list = list(module_info_iterable) + self.allowed_dtypes = set(allowed_dtypes) if allowed_dtypes is not None else None + self.train_eval_mode = train_eval_mode + self.skip_if_dynamo = skip_if_dynamo + + def _get_training_flags(self, module_info): + training_flags = [] + if (self.train_eval_mode == TrainEvalMode.train_only or + self.train_eval_mode == TrainEvalMode.train_and_eval): + training_flags.append(True) + + if (self.train_eval_mode == TrainEvalMode.eval_only or + self.train_eval_mode == TrainEvalMode.train_and_eval): + training_flags.append(False) + + # If train and eval modes don't differ for the module, don't bother using more than one. + if not module_info.train_and_eval_differ: + training_flags = training_flags[:1] + + return training_flags + + def _parametrize_test(self, test, generic_cls, device_cls): + if device_cls is None: + raise RuntimeError('The @modules decorator is only intended to be used in a device-specific ' + 'context; use it with instantiate_device_type_tests() instead of ' + 'instantiate_parametrized_tests()') + + for module_info in self.module_info_list: + dtypes = set(module_info.supported_dtypes(device_cls.device_type)) + if self.allowed_dtypes is not None: + dtypes = dtypes.intersection(self.allowed_dtypes) + + training_flags = self._get_training_flags(module_info) + for (training, dtype) in product(training_flags, dtypes): + # Construct the test name; device / dtype parts are handled outside. + # See [Note: device and dtype suffix placement] + test_name = module_info.formatted_name + if len(training_flags) > 1: + test_name += f"_{'train_mode' if training else 'eval_mode'}" + + # Construct parameter kwargs to pass to the test. + param_kwargs = {'module_info': module_info} + _update_param_kwargs(param_kwargs, 'dtype', dtype) + _update_param_kwargs(param_kwargs, 'training', training) + + try: + + @wraps(test) + def test_wrapper(*args, **kwargs): + return test(*args, **kwargs) + + if self.skip_if_dynamo and not torch.testing._internal.common_utils.TEST_WITH_TORCHINDUCTOR: + test_wrapper = skipIfTorchDynamo("Policy: we don't run ModuleInfo tests w/ Dynamo")(test_wrapper) + + decorator_fn = partial(module_info.get_decorators, generic_cls.__name__, + test.__name__, device_cls.device_type, dtype) + + yield (test_wrapper, test_name, param_kwargs, decorator_fn) + except Exception as ex: + # Provides an error message for debugging before rethrowing the exception + print(f"Failed to instantiate {test_name} for module {module_info.name}!") + raise ex + + +def get_module_common_name(module_cls): + if module_cls in MODULE_CLASS_NAMES: + # Example: "nn.Linear" + return MODULE_CLASS_NAMES[module_cls] + else: + return module_cls.__name__ + + +class FunctionInput: + """ Contains args and kwargs to pass as input to a function. """ + __slots__ = ['args', 'kwargs'] + + def __init__(self, *args, **kwargs): + self.args = args + self.kwargs = kwargs + + +class ModuleInput: + """ Contains args / kwargs for module instantiation + forward pass. """ + __slots__ = ['constructor_input', 'forward_input', 'desc', 'reference_fn'] + + def __init__(self, constructor_input, forward_input=None, desc='', reference_fn=None): + self.constructor_input = constructor_input # Inputs to pass during construction + self.forward_input = forward_input # Inputs to pass to forward() + self.desc = desc # Description for this set of inputs + self.reference_fn = reference_fn # Reference with signature: reference_fn(module, parameters, *args, **kwargs) + + if reference_fn is not None: + + @wraps(reference_fn) + def copy_reference_fn(m, *args, **kwargs): + # Copy inputs to avoid undesired side effects from calling the reference. + args, kwargs = deepcopy(args), deepcopy(kwargs) + + # Note that module parameters are passed in for convenience. + return reference_fn(m, list(m.parameters()), *args, **kwargs) + + self.reference_fn = copy_reference_fn + +class ModuleErrorEnum(Enum): + """ Enumerates when error is raised when testing modules. """ + CONSTRUCTION_ERROR = 0 + FORWARD_ERROR = 1 + +class ErrorModuleInput: + """ + A ModuleInput that will cause the operation to throw an error plus information + about the resulting error. + """ + + __slots__ = ["module_error_input", "error_on", "error_type", "error_regex"] + + def __init__(self, + module_error_input, + *, + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=RuntimeError, + error_regex): + self.module_error_input = module_error_input + self.error_on = error_on + self.error_type = error_type + self.error_regex = error_regex + + +class ModuleInfo: + """ Module information to be used in testing. """ + + def __init__(self, + module_cls, # Class object for the module under test + *, + module_inputs_func, # Function to generate module inputs + skips=(), # Indicates which tests to skip + decorators=None, # Additional decorators to apply to generated tests + dtypes=floating_types(), # dtypes this function is expected to work with + dtypesIfMPS=(torch.float16, torch.float32,), # dtypes this function is expected to work with on MPS + dtypesIfHpu=(torch.bfloat16, torch.float32,), + supports_gradgrad=True, # whether the op supports second order gradients + gradcheck_nondet_tol=0.0, # tolerance for nondeterminism while performing gradcheck + module_memformat_affects_out=False, # whether converting module to channels last will generate + # channels last output + train_and_eval_differ=False, # whether the module has differing behavior between train and eval + module_error_inputs_func=None, # Function to generate module inputs that error + gradcheck_fast_mode=None, # Whether to use the fast implementation for gradcheck/gradgradcheck. + # When set to None, defers to the default value provided by the wrapper + # function around gradcheck (testing._internal.common_utils.gradcheck) + ): + self.module_cls = module_cls + self.module_inputs_func = module_inputs_func + self.decorators = (*(decorators if decorators else []), *(skips if skips else [])) + self.dtypes = dtypes + self.dtypesIfMPS = dtypesIfMPS + self.dtypesIfHpu = dtypesIfHpu + self.supports_gradgrad = supports_gradgrad + self.gradcheck_nondet_tol = gradcheck_nondet_tol + self.module_memformat_affects_out = module_memformat_affects_out + self.train_and_eval_differ = train_and_eval_differ + self.module_error_inputs_func = module_error_inputs_func + self.gradcheck_fast_mode = gradcheck_fast_mode + self.is_lazy = issubclass(module_cls, torch.nn.modules.lazy.LazyModuleMixin) + + def get_decorators(self, test_class, test_name, device, dtype, param_kwargs): + result = [] + for decorator in self.decorators: + if isinstance(decorator, DecorateInfo): + if decorator.is_active(test_class, test_name, device, dtype, param_kwargs): + result.extend(decorator.decorators) + else: + result.append(decorator) + return result + + def supported_dtypes(self, device_type): + if device_type == 'mps': + return self.dtypesIfMPS + elif device_type == 'hpu': + return self.dtypesIfHpu + else: + return self.dtypes + + @property + def name(self): + return get_module_common_name(self.module_cls) + + @property + def formatted_name(self): + return self.name.replace('.', '_') + +# Start of module inputs functions. + +def module_inputs_torch_nn_Linear(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + module_inputs = [ + ModuleInput(constructor_input=FunctionInput(10, 8), + forward_input=FunctionInput(input=make_input((4, 10))), + reference_fn=lambda m, p, input: torch.mm(input, p[0].t()) + p[1].view(1, -1).expand(4, 8)), + ModuleInput(constructor_input=FunctionInput(10, 8, bias=False), + forward_input=FunctionInput(make_input((4, 10))), + desc='no_bias', + reference_fn=lambda m, p, i: torch.mm(i, p[0].t())), + ModuleInput(constructor_input=FunctionInput(3, 5), + forward_input=FunctionInput(make_input(3)), + desc='no_batch_dim', + reference_fn=lambda m, p, i: torch.mm(i.view(1, -1), p[0].t()).view(-1) + p[1]) + ] + + return module_inputs + + +def module_inputs_torch_nn_Bilinear(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def bilinear_reference_fn(m, p, x1, x2, bias=True): + result = torch.einsum('bn,anm,bm->ba', x1, p[0], x2) + if bias: + if x1.shape[0] == 1: + result = result.view(-1) + p[1] + else: + result = result + p[1].view(1, -1).expand(x1.shape[0], p[0].shape[0]) + return result + + module_inputs = [ + ModuleInput(constructor_input=FunctionInput(2, 3, 4), + forward_input=FunctionInput(make_input((8, 2)), make_input((8, 3))), + reference_fn=bilinear_reference_fn), + ModuleInput(constructor_input=FunctionInput(2, 3, 4, bias=False), + forward_input=FunctionInput(make_input((8, 2)), make_input((8, 3))), + desc='no_bias', + reference_fn=lambda m, p, x1, x2: bilinear_reference_fn(m, p, x1, x2, bias=False)), + ModuleInput(constructor_input=FunctionInput(2, 3, 4), + forward_input=FunctionInput(make_input(2), make_input(3)), + desc='no_batch_dim', + reference_fn=lambda m, p, x1, x2: bilinear_reference_fn(m, p, x1.view(1, -1), x2.view(1, -1))), + ] + + return module_inputs + + +def module_inputs_torch_nn_KLDivLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_batchmean', {'reduction': 'batchmean'}), + ('reduction_none', {'reduction': 'none'}), + ('log_target', {'log_target': True}) + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return kldivloss_reference(i, t, **constructor_kwargs) + + input = make_input((10, 10)).log() + target = make_input((10, 10)) if kwargs.get('log_target', False) else make_input((10, 10)).log() + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(input, target), + desc=desc, + reference_fn=reference_fn) + ) + + scalar_input = make_input(()).log() + # FIXME(rec): scalar_target is unused, perhaps should be argument to FunctionInput? + scalar_target = ( # noqa: F841 + make_input(()) if kwargs.get('log_target', False) else make_input(()).log() + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(scalar_input, scalar_input), + desc='scalar_' + desc, + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_NLLLoss(module_info, device, dtype, requires_grad, training, **kwargs): + def make_input(shape, device=device, dtype=dtype, requires_grad=requires_grad): + return make_tensor(shape, device=device, dtype=dtype, + requires_grad=False).log_softmax(dim=1).requires_grad_(requires_grad) + make_weight = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_none', {'reduction': 'none'}), + ('ignore_index', {'ignore_index': 2}), + ('weights', {'weight': make_weight(4).abs()}), + ('weights_ignore_index', {'weight': make_weight(4).abs(), 'ignore_index': 2}), + ('weights_ignore_index_neg', {'weight': make_weight(4).abs(), 'ignore_index': -1}) + ] + + # TODO: Uncomment when negative weights is supported. + # negative_weight = make_weight(10) + # negative_weight[0] = -1 + # cases.append(('weights_negative', {'weight': negative_weight})) + module_inputs = [] + for desc, constructor_kwargs in cases: + + def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return nllloss_reference(i, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((15, 4)), + torch.empty(15, device=device).uniform_().mul(4).floor().long()), + desc=desc, + reference_fn=reference_fn) + ) + + def nd_reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return nlllossNd_reference(i, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput( + make_input((2, 4, 5, 5)), + torch.empty(2, 5, 5, device=device).uniform_().mul(4).floor().long()), + desc=f"nd_{desc}", + reference_fn=nd_reference_fn) + ) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput( + make_input((2, 4, 5, 5, 2, 2)), + torch.empty(2, 5, 5, 2, 2, device=device).uniform_().mul(4).floor().long()), + desc=f"higher_dim_{desc}", + reference_fn=nd_reference_fn) + ) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput( + make_input((2, 4, 5)), + torch.empty(2, 5, device=device).uniform_().mul(4).floor().long()), + desc=f"3d_{desc}", + reference_fn=nd_reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_GaussianNLLLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('homoscedastic', {'homoscedastic': True}), + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + homoscedastic = constructor_kwargs.pop('homoscedastic', False) + var_input = make_input(1, 3).abs() if homoscedastic else make_input(4, 1).abs() + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input(4, 3), + make_target(4, 3), + var_input), + desc=desc, + reference_fn=no_batch_dim_reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_PoissonNLLLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('full', {'full': True}), + ('no_log_input', {'log_input': False}), + ('full_no_log_input', {'full': True, 'log_input': False}), + ] + + def poissonnllloss_reference_fn(i, t, log_input=True, full=False, reduction='mean', eps=1e-8): + if log_input: + result = i.exp() - t.mul(i) + else: + result = i - t.mul((i + eps).log()) + + if full: + result += (t.mul(t.log()) - t + 0.5 * (2. * math.pi * t).log()).masked_fill(t <= 1, 0) + + if reduction == 'none': + return result + elif reduction == 'mean': + return result.sum() / i.numel() + else: + return result.sum() + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return poissonnllloss_reference_fn(i, t, **constructor_kwargs) + + log_input = constructor_kwargs.get('log_input', True) + input = make_input((2, 3, 4, 5)) if log_input else make_input((2, 3, 4, 5)).abs().add(0.001) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(input, + make_target((2, 3, 4, 5)).floor_().abs_()), + desc=desc, + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_MSELoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ] + + def mse_loss_reference_fn(m, p, i, t, reduction='mean'): + if reduction == 'none': + return (i - t).pow(2) + elif reduction == 'mean': + return (i - t).pow(2).sum() / i.numel() + else: + return (i - t).pow(2).sum() + + module_inputs = [] + for desc, constructor_kwargs in cases: + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((2, 3, 4, 5)), + make_target((2, 3, 4, 5))), + desc=desc, + reference_fn=partial(mse_loss_reference_fn, **constructor_kwargs)) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input(()), + make_target(())), + desc=f'{desc}_scalar', + reference_fn=partial(mse_loss_reference_fn, **constructor_kwargs)) + ) + + return module_inputs + + +def no_batch_dim_reference_fn(m, p, *args, **kwargs): + """Reference function for modules supporting no batch dimensions. + + Unbatched inputs are unsqueezed to form a + single batch input before passing them to the module. + The output is squeezed to compare with the + output of unbatched input to the module. + + Currently it only supports modules which return a single Tensor as output. + You can bind the following kwargs. + Kwargs: + batch_first[bool] : If True, all the Tensors in `args` while be unsqueezed at dim `0` . + and output will be squeezed at dim `0` else dim `1` for both. + kwargs_to_batchify[dict] : Dictionary specifying the name of the argument and dimension to unsqueeze. + Useful if there are few arguments whose batch dimension are different + from the ones selected by `batch_first`. + is_criterion[bool] : Specify if the module is a criterion and handle the reduction for output accordingly. + """ + def get_and_pop(key, default): + v = kwargs.get(key, default) + if key in kwargs: + kwargs.pop(key) + return v + + batch_dim = 0 if get_and_pop('batch_first', True) else 1 + kwargs_to_batchify = get_and_pop('kwargs_to_batchify', None) + is_criterion = get_and_pop('is_criterion', False) + + if kwargs_to_batchify is not None: + if not isinstance(kwargs_to_batchify, dict): + raise AssertionError(f"Expected kwargs_to_batchify to be a dict, got {type(kwargs_to_batchify)}") + for k, v in kwargs.items(): + if k in kwargs_to_batchify and v is not None: + bdim = kwargs_to_batchify[k] + kwargs[k] = v.unsqueeze(bdim) + + single_batch_input_args = [input.unsqueeze(batch_dim) for input in args] + with freeze_rng_state(): + output = m(*single_batch_input_args, **kwargs).squeeze(batch_dim) + + if is_criterion: + reduction = get_reduction(m) + if reduction == 'none': + return output.squeeze(0) + return output + + +def no_batch_dim_reference_mha(m, p, *args, **kwargs): + """Reference function for MultiheadAttention supporting no batch dimensions. + + Unbatched inputs are unsqueezed to form a + single batch input before passing them to the module. + The output is squeezed to compare with the + output of unbatched input to the module. + """ + batch_dim = 0 if kwargs.get('batch_first', True) else 1 + if 'batch_first' in kwargs: + kwargs.pop('batch_first') + if 'key_padding_mask' in kwargs and kwargs['key_padding_mask'] is not None: + kwargs['key_padding_mask'] = kwargs['key_padding_mask'].unsqueeze(0) + single_batch_input_args = [input.unsqueeze(batch_dim) for input in args] + with freeze_rng_state(): + output = m(*single_batch_input_args, **kwargs) + return (output[0].squeeze(batch_dim), output[1].squeeze(0)) + + +def no_batch_dim_reference_rnn_gru(m, p, *args, **kwargs): + """Reference function for RNN and GRU supporting no batch dimensions. + + Unbatched inputs are unsqueezed to form a + single batch input before passing them to the module. + The output is squeezed to compare with the + output of unbatched input to the module. + """ + if len(args) == 1: + inp, = args + h = None + elif len(args) == 2: + inp, h = args + h = h.unsqueeze(1) + + batch_dim = 0 if kwargs['batch_first'] else 1 + kwargs.pop('batch_first') + inp = inp.unsqueeze(batch_dim) + single_batch_input_args = (inp, h) + with freeze_rng_state(): + output = m(*single_batch_input_args, **kwargs) + return (output[0].squeeze(batch_dim), output[1].squeeze(1)) + + +def no_batch_dim_reference_lstm(m, p, *args, **kwargs): + """Reference function for LSTM supporting no batch dimensions. + + Unbatched inputs are unsqueezed to form a + single batch input before passing them to the module. + The output is squeezed to compare with the + output of unbatched input to the module. + """ + if len(args) == 1: + inp, = args + h = None + elif len(args) == 2: + inp, h = args + h = (h[0].unsqueeze(1), h[1].unsqueeze(1)) + + batch_dim = 0 if kwargs['batch_first'] else 1 + kwargs.pop('batch_first') + inp = inp.unsqueeze(batch_dim) + single_batch_input_args = (inp, h) + with freeze_rng_state(): + output = m(*single_batch_input_args, **kwargs) + return (output[0].squeeze(batch_dim), (output[1][0].squeeze(1), output[1][1].squeeze(1))) + + +def no_batch_dim_reference_lstmcell(m, p, *args, **kwargs): + """Reference function for LSTMCell supporting no batch dimensions. + + The module is passed the input and target in batched form with a single item. + The output is squeezed to compare with the no-batch input. + """ + inp, (h, c) = args + single_batch_input_args = (inp.unsqueeze(0), (h.unsqueeze(0), c.unsqueeze(0))) + with freeze_rng_state(): + output = m(*single_batch_input_args, **kwargs) + return (output[0].squeeze(0), output[1].squeeze(0)) + + +def generate_regression_criterion_inputs(make_input): + return [ + ModuleInput( + constructor_input=FunctionInput(reduction=reduction), + forward_input=FunctionInput(make_input((4, )), make_input(4,)), + reference_fn=partial(no_batch_dim_reference_fn, is_criterion=True), + desc=f'no_batch_dim_{reduction}' + ) for reduction in ['none', 'mean', 'sum']] + + +def module_inputs_torch_nn_AvgPool1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(kernel_size=2), + forward_input=FunctionInput(make_input((3, 6))), + desc='no_batch_dim', + reference_fn=no_batch_dim_reference_fn), + ModuleInput(constructor_input=FunctionInput(2), + forward_input=FunctionInput(make_input((2, 3, 6)))), + ModuleInput(constructor_input=FunctionInput((2,), (2,)), + forward_input=FunctionInput(make_input((2, 3, 6))), + desc='stride'), + ModuleInput(constructor_input=FunctionInput(2, 2, 1), + forward_input=FunctionInput(make_input((2, 3, 6))), + desc='stride_pad')] + + +def module_inputs_torch_nn_AvgPool2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput((2, 2)), + forward_input=FunctionInput(make_input((3, 6, 6))), + desc='no_batch_dim', + reference_fn=no_batch_dim_reference_fn), + ModuleInput(constructor_input=FunctionInput((2, 2)), + forward_input=FunctionInput(make_input((2, 3, 6, 6)))), + ModuleInput(constructor_input=FunctionInput((2, 2), (2, 2)), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='stride'), + ModuleInput(constructor_input=FunctionInput((2, 2), (2, 2), (1, 1)), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='stride_pad'), + ModuleInput(constructor_input=FunctionInput((2, 2), divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='divisor'), + ModuleInput(constructor_input=FunctionInput((2, 2), (2, 2), divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='divisor_stride'), + ModuleInput(constructor_input=FunctionInput((2, 2), (2, 2), (1, 1), divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='divisor_stride_pad')] + + + +def module_inputs_torch_nn_AvgPool3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput((2, 2, 2)), + forward_input=FunctionInput(make_input((3, 4, 4, 4))), + desc='no_batch_dim', + reference_fn=no_batch_dim_reference_fn), + ModuleInput(constructor_input=FunctionInput((2, 2, 2)), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4)))), + ModuleInput(constructor_input=FunctionInput(2, (2, 2, 2)), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5))), + desc='stride'), + ModuleInput(constructor_input=FunctionInput(2, 2, (1, 1, 1)), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5))), + desc='stride_pad'), + ModuleInput(constructor_input=FunctionInput(4, 2, (1, 2, 1)), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5))), + desc='stride_pad_gpu_fixedkw_output'), + ModuleInput(constructor_input=FunctionInput((2, 4, 8), 1, (1, 1, 2)), + forward_input=FunctionInput(make_input((2, 3, 2, 4, 8))), + desc='stride_pad_gpu_general_output'), + ModuleInput(constructor_input=FunctionInput(3, 1, 0), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='stride1_pad0_gpu_input'), + ModuleInput(constructor_input=FunctionInput(2, 2, (1, 1, 1)), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='stride_pad_gpu_input_nooverlap'), + ModuleInput(constructor_input=FunctionInput((2, 2, 2), divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='divisor'), + ModuleInput(constructor_input=FunctionInput(2, (2, 2, 2), divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5))), + desc='divisor_stride'), + ModuleInput(constructor_input=FunctionInput(2, 2, (1, 1, 1), divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5))), + desc='divisor_stride_pad'), + ModuleInput(constructor_input=FunctionInput(4, 2, (1, 2, 1), divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5))), + desc='divisor_stride_pad_gpu_fixedkw_output'), + ModuleInput(constructor_input=FunctionInput((2, 4, 8), 1, (1, 1, 2), divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 2, 4, 8))), + desc='divisor_stride_pad_gpu_general_output'), + ModuleInput(constructor_input=FunctionInput(3, 1, 0, divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='divisor_stride1_pad0_gpu_input'), + ModuleInput(constructor_input=FunctionInput(2, 2, (1, 1, 1), divisor_override=1), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='divisor_stride_pad_gpu_input_nooverlap')] + + + +def module_inputs_torch_nn_AdaptiveAvgPool1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((1, 3, 5))), + desc='single'), + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((3, 5))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput(1,), + forward_input=FunctionInput(make_input((1, 3, 5))), + desc='one_output')] + + +def module_inputs_torch_nn_AdaptiveAvgPool2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((1, 3, 5, 6))), + desc='single'), + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((3, 5, 6))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput(1,), + forward_input=FunctionInput(make_input((1, 3, 5, 6))), + desc='single_1x1output'), + ModuleInput(constructor_input=FunctionInput((3, 4)), + forward_input=FunctionInput(make_input((1, 3, 5, 6))), + desc='tuple'), + ModuleInput(constructor_input=FunctionInput((3, None)), + forward_input=FunctionInput(make_input((1, 3, 5, 6))), + desc='tuple_none')] + +def module_inputs_torch_nn_AdaptiveAvgPool3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((2, 3, 5, 2, 7))), + desc='single'), + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((3, 5, 2, 7))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput((3, 4, 5)), + forward_input=FunctionInput(make_input((2, 3, 5, 3, 7))), + desc='tuple'), + ModuleInput(constructor_input=FunctionInput((None, 4, 5)), + forward_input=FunctionInput(make_input((2, 3, 5, 3, 7))), + desc='tuple_none'), + ModuleInput(constructor_input=FunctionInput((3, 2, 2)), + forward_input=FunctionInput(make_input((1, 1, 3, 2, 6))), + desc='last_dim')] + + +def module_inputs_torch_nn_AdaptiveMaxPool1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((1, 3, 5))), + desc='single'), + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((3, 5))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_AdaptiveMaxPool2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((1, 3, 5, 6))), + desc='single'), + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((3, 5, 6))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput((3, 4)), + forward_input=FunctionInput(make_input((1, 3, 5, 6))), + desc='tuple'), + ModuleInput(constructor_input=FunctionInput((3, None)), + forward_input=FunctionInput(make_input((1, 3, 5, 6))), + desc='tuple_none')] + + +def module_inputs_torch_nn_AdaptiveMaxPool3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((2, 3, 5, 6, 7))), + desc='single'), + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((3, 5, 6, 7))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput((3, 4, 5)), + forward_input=FunctionInput(make_input((2, 3, 5, 6, 7))), + desc='tuple'), + ModuleInput(constructor_input=FunctionInput((3, None, 5)), + forward_input=FunctionInput(make_input((2, 3, 5, 6, 7))), + desc='tuple_none'), + ModuleInput(constructor_input=FunctionInput(3), + forward_input=FunctionInput(make_input((2, 3, 12, 9, 3))), + desc='single_nonatomic'), + ModuleInput(constructor_input=FunctionInput((3, 4, 5)), + forward_input=FunctionInput(make_input((2, 3, 6, 4, 10))), + desc='tuple_nonatomic')] + + +def module_inputs_torch_nn_BatchNorm1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(10,), + forward_input=FunctionInput(make_input((4, 10))), + desc='affine'), + ModuleInput(constructor_input=FunctionInput(5,), + forward_input=FunctionInput(make_input((4, 5, 3))), + desc='3d_input'), + ModuleInput(constructor_input=FunctionInput(10, 1e-3, None), + forward_input=FunctionInput(make_input((4, 10))), + desc='affine_simple_average'), + ModuleInput(constructor_input=FunctionInput(10, 1e-3, 0.3, False), + forward_input=FunctionInput(make_input((4, 10))), + desc='not_affine'), + ModuleInput(constructor_input=FunctionInput(10, 1e-3, 0.3, True, False), + forward_input=FunctionInput(make_input((4, 10))), + desc='not_tracking_stats'), + ModuleInput(constructor_input=FunctionInput(5, 1e-3, 0.3, False), + forward_input=FunctionInput(make_input((4, 5, 3))), + desc='3d_input_not_affine'), + ModuleInput(constructor_input=FunctionInput(5, 1e-3, 0.3, False), + forward_input=FunctionInput(make_input((0, 5, 9))), + desc='zero_batch'), + ModuleInput(constructor_input=FunctionInput(10, bias=False), + forward_input=FunctionInput(make_input((4, 10))), + desc='affine_not_bias'),] + + +def module_inputs_torch_nn_BatchNorm2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((2, 3, 6, 6)))), + ModuleInput(constructor_input=FunctionInput(3, 1e-3, None), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='2d_simple_average'), + ModuleInput(constructor_input=FunctionInput(3, 1e-3, 0.8), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='momentum'), + ModuleInput(constructor_input=FunctionInput(3, 1e-3, 0.8, False), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='not_affine'), + ModuleInput(constructor_input=FunctionInput(3, 1e-3, 0.8, True, False), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='not_tracking_stats'), + ModuleInput(constructor_input=FunctionInput(5, 1e-3, 0.3, False), + forward_input=FunctionInput(make_input((0, 5, 2, 2))), + desc='zero_batch'), + ModuleInput(constructor_input=FunctionInput(3, bias=False), + forward_input=FunctionInput(make_input((2, 3, 6, 6))), + desc='affine_not_bias'),] + + +def module_inputs_torch_nn_BatchNorm3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4)))), + ModuleInput(constructor_input=FunctionInput(3, 1e-3, None), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='3d_simple_average'), + ModuleInput(constructor_input=FunctionInput(3, 1e-3, 0.7), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='momentum'), + ModuleInput(constructor_input=FunctionInput(3, 1e-3, 0.7, False), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='not_affine'), + ModuleInput(constructor_input=FunctionInput(3, 1e-3, 0.7, True, False), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='not_tracking_stats'), + ModuleInput(constructor_input=FunctionInput(5, 1e-3, 0.3, False), + forward_input=FunctionInput(make_input((0, 5, 2, 2, 2))), + desc='zero_batch'), + ModuleInput(constructor_input=FunctionInput(3, bias=False), + forward_input=FunctionInput(make_input((2, 3, 4, 4, 4))), + desc='affine_not_bias'),] + + +def module_error_inputs_torch_nn_BatchNorm1d_2d_3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + if module_info.module_cls == torch.nn.BatchNorm1d: + input_shape = (2, 10) + elif module_info.module_cls == torch.nn.BatchNorm2d: + input_shape = (2, 10, 5, 5) + else: + input_shape = (2, 10, 4, 4, 4) + + return [ + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, eps=-1.0), + forward_input=FunctionInput(make_input(input_shape)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=ValueError, + error_regex="eps must be positive" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, eps=0.0), + forward_input=FunctionInput(make_input(input_shape)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=ValueError, + error_regex="eps must be positive" + ), + ] + + +def module_inputs_torch_nn_ConvNd(module_info, device, dtype, requires_grad, training, **kwargs): + N = kwargs['N'] + lazy = kwargs.get('lazy', False) + transposed = kwargs.get('transposed', False) + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + conv_kwargs_list = [{}] if transposed else [{}, {'padding': 'same'}] + kernel_size, C_in, C_out = 3, 4, 5 + input_no_batch_shape = (C_in,) + tuple(i + 3 for i in range(N)) + input_batch_shape = (2,) + input_no_batch_shape + return [ + ModuleInput(constructor_input=(FunctionInput(C_out, kernel_size, **conv_kwargs) if lazy else + FunctionInput(C_in, C_out, kernel_size, **conv_kwargs)), + forward_input=FunctionInput(make_input( + input_batch_shape if with_batch else input_no_batch_shape)), + desc=('' if with_batch else 'no_batch_dim'), + reference_fn=(None if with_batch else no_batch_dim_reference_fn)) + for with_batch, conv_kwargs in itertools.product([True, False], conv_kwargs_list) + ] + + +def module_inputs_torch_nn_CosineEmbeddingLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('margin', {'margin': 0.7}) + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i1, i2, t, constructor_kwargs=constructor_kwargs): + return cosineembeddingloss_reference(i1, i2, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((15, 10)), make_input((15, 10)), + make_target((15,)).sign()), + desc=desc, + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_ELU(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(alpha=2.), + forward_input=FunctionInput(make_input((3, 2, 5))), + reference_fn=lambda m, p, i: torch.where(i >= 0, i, 2 * (i.exp() - 1))), + ModuleInput(constructor_input=FunctionInput(alpha=2.), + forward_input=FunctionInput(make_input(())), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((3,))), + desc='no_batch_dim', + reference_fn=no_batch_dim_reference_fn), + ModuleInput(constructor_input=FunctionInput(alpha=2.), + forward_input=FunctionInput(make_input((2, 3, 2, 5))), + desc='4d_input')] + + +def module_inputs_torch_nn_CELU(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(alpha=2.), + forward_input=FunctionInput(make_input((3, 2, 5))), + reference_fn=lambda m, p, i: torch.where(i >= 0, i, 2. * ((.5 * i).exp() - 1))), + ModuleInput(constructor_input=FunctionInput(alpha=2.), + forward_input=FunctionInput(make_input(())), + reference_fn=lambda m, p, i: torch.where(i >= 0, i, 2. * ((.5 * i).exp() - 1)), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(alpha=2.), + forward_input=FunctionInput(make_input((3,))), + desc='no_batch_dim', + reference_fn=no_batch_dim_reference_fn)] + + +def module_inputs_torch_nn_GLU(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((5, 6)))), + ModuleInput(constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((5, 6, 7))), + desc='dim'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((4,))), + desc='no_batch_dim', + reference_fn=no_batch_dim_reference_fn)] + + +def module_inputs_torch_nn_GELU(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput('none'), + forward_input=FunctionInput(make_input(())), + reference_fn=lambda m, p, x, *_: x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0))), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput('none'), + forward_input=FunctionInput(make_input((3, 2, 5))), + reference_fn=lambda m, p, x, *_: x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((3,))), + desc='no_batch_dim', + reference_fn=no_batch_dim_reference_fn)] + + +def module_inputs_torch_nn_ReLU(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4, 5))), + desc='channels_last_mem_format'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 3, 4, 5))), + desc='channels_last_3d_mem_format')] + + +def module_inputs_torch_nn_ReLU6(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4, 5))), + desc='channels_last_mem_format'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 3, 4, 5))), + desc='channels_last_3d_mem_format')] + + +def module_inputs_torch_nn_LeakyReLU(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((3, 2, 5)))), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput(0.5), + forward_input=FunctionInput(make_input((3, 2, 5))), + desc='with_negval'), + ModuleInput(constructor_input=FunctionInput(0.0), + forward_input=FunctionInput(make_input((10, 10))), + desc='with_zero_negval'), + ModuleInput(constructor_input=FunctionInput(0.5), + forward_input=FunctionInput(make_input(())), + desc='with_negval_scalar')] + + +def module_inputs_torch_nn_PReLU(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4))), + reference_fn=lambda m, p, i: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0], + desc='1d'), + ModuleInput(constructor_input=FunctionInput(3), + forward_input=FunctionInput(make_input((2, 3, 4))), + reference_fn=lambda m, p, i: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0], + desc='1d_multiparam'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4, 5))), + reference_fn=lambda m, p, i: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0], + desc='2d'), + ModuleInput(constructor_input=FunctionInput(3), + forward_input=FunctionInput(make_input((2, 3, 4, 5))), + reference_fn=lambda m, p, i: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0], + desc='2d_multiparam'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4, 5, 6))), + reference_fn=lambda m, p, i: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0], + desc='3d'), + ModuleInput(constructor_input=FunctionInput(3), + forward_input=FunctionInput(make_input((2, 3, 4, 5, 6))), + reference_fn=lambda m, p, i: torch.clamp(i, min=0) + torch.clamp(i, max=0) * p[0][0], + desc='3d_multiparam')] + + +def module_inputs_torch_nn_SELU(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((3, 2, 5)))), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + desc='scalar')] + + +def module_inputs_torch_nn_SiLU(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + reference_fn=lambda m, p, x, *_: x * torch.sigmoid(x), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((5, 6, 7))), + reference_fn=lambda m, p, x, *_: x * torch.sigmoid(x))] + + +def module_inputs_torch_nn_Softmax(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((10, 20))), + reference_fn=lambda m, p, i: torch.exp(i).div(torch.exp(i).sum(1, True).expand(10, 20))), + ModuleInput(constructor_input=FunctionInput(0), + forward_input=FunctionInput(make_input(())), + reference_fn=lambda m, p, i: torch.exp(i).div(torch.exp(i).sum(0, True)), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(-1), + forward_input=FunctionInput(make_input((4, 5))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_Softmax2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((1, 3, 10, 20))), + reference_fn=lambda m, p, i: torch.exp(i).div(torch.exp(i).sum(1, False))), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((3, 4, 5))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_LogSoftmax(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((10, 20))), + reference_fn=lambda m, p, i: torch.exp(i).div_(torch.exp(i).sum(1, True).expand(10, 20)).log_()), + ModuleInput(constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((1, 3, 10, 20))), + reference_fn=lambda m, p, i: torch.exp(i).div_(torch.exp(i).sum(1, False)).log_(), + desc='multiparam'), + ModuleInput(constructor_input=FunctionInput(0), + forward_input=FunctionInput(make_input(())), + reference_fn=lambda m, p, i: torch.exp(i).div_(torch.exp(i).sum(0, False)).log_(), + desc='multiparam_scalar'), + ModuleInput(constructor_input=FunctionInput(-1), + forward_input=FunctionInput(make_input((4, 5))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_Softmin(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((10, 20)))), + ModuleInput(constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((2, 3, 5, 10))), + desc='multidim'), + ModuleInput(constructor_input=FunctionInput(0), + forward_input=FunctionInput(make_input(())), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(-1), + forward_input=FunctionInput(make_input((3, 4, 10))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_Softplus(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((10, 20))), + reference_fn=lambda m, p, i: torch.log1p(torch.exp(i))), + ModuleInput(constructor_input=FunctionInput(2), + forward_input=FunctionInput(make_input((10, 20))), + reference_fn=lambda m, p, i: 1. / 2. * torch.log1p(torch.exp(2 * i)), + desc='beta'), + ModuleInput(constructor_input=FunctionInput(2, -100), + forward_input=FunctionInput(make_input((10, 20))), + reference_fn=( + lambda m, p, i: ((i * 2) > -100).type_as(i) * i + + ((i * 2) <= -100).type_as(i) * 1. / 2. * torch.log1p(torch.exp(2 * i))), + desc='beta_threshold'), + ModuleInput(constructor_input=FunctionInput(2, -100), + forward_input=FunctionInput(make_input(())), + reference_fn=( + lambda m, p, i: ((i * 2) > -100).type_as(i) * i + + ((i * 2) <= -100).type_as(i) * 1. / 2. * torch.log1p(torch.exp(2 * i))), + desc='beta_threshold_scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_Softshrink(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((3, 2, 5)))), + ModuleInput(constructor_input=FunctionInput(1,), + forward_input=FunctionInput(make_input((3, 2, 5))), + desc='lambda'), + ModuleInput(constructor_input=FunctionInput(1,), + forward_input=FunctionInput(make_input(())), + desc='lambda_scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_Softsign(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((3, 2, 5))), + reference_fn=lambda m, p, i: i.div(1 + torch.abs(i))), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + reference_fn=lambda m, p, i: i.div(1 + torch.abs(i)), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_Tanh(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4, 5)))), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + + +def module_inputs_torch_nn_Tanhshrink(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4, 5)))), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_Threshold(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(2., 1.), + forward_input=FunctionInput(make_input((2, 3, 4, 5))), + desc='threshold_value'), + ModuleInput(constructor_input=FunctionInput(2., 10.), + forward_input=FunctionInput(make_input((2, 3, 4, 5))), + desc='large_value'), + ModuleInput(constructor_input=FunctionInput(2., 1.), + forward_input=FunctionInput(make_input(())), + desc='threshold_value_scalar'), + ModuleInput(constructor_input=FunctionInput(2., 1.), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_Mish(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((5, 6, 7))), + reference_fn=lambda m, p, i: i * torch.tanh(F.softplus(i))), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + reference_fn=lambda m, p, i: i * torch.tanh(F.softplus(i)), + desc='scalar'), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim')] + + +def module_inputs_torch_nn_L1Loss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4)), + make_input((2, 3, 4))), + reference_fn=lambda m, p, i, t: 1. / i.numel() * sum((a - b).abs().sum() + for a, b in zip(i, t, strict=True))), + ModuleInput(constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(()), make_input(())), + reference_fn=lambda m, p, i, t: 1. / i.numel() * (i - t).abs().sum(), + desc='scalar')] + generate_regression_criterion_inputs(make_input) + + +def module_inputs_torch_nn_SmoothL1Loss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return smoothl1loss_reference(i, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((5, 10)), + make_input((5, 10))), + desc=desc, + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input(()), + make_input(())), + desc=f'scalar_{desc}', + reference_fn=reference_fn) + ) + + return module_inputs + + + +def module_inputs_torch_nn_BCELoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + make_weight = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('weights', {'weight': make_weight((10,))}), + ] + + def bce_loss_reference_fn(m, p, i, t, reduction='mean', weight=None): + result = -(t * i.log() + (1 - t) * (1 - i).log()) + + if weight is not None: + result = result * weight + + if reduction == 'none': + return result + elif reduction == 'mean': + return result.sum() / i.numel() + else: + return result.sum() + + module_inputs = [] + for desc, constructor_kwargs in cases: + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((15, 10), low=1e-2, high=1 - 1e-2), + make_target((15, 10)).gt(0).to(dtype)), + desc=desc, + reference_fn=partial(bce_loss_reference_fn, **constructor_kwargs)) + ) + + scalar_weight = make_weight(()) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(weight=scalar_weight), + forward_input=FunctionInput(make_input((), low=1e-2, high=1 - 1e-2), + make_target(()).gt(0).to(dtype)), + desc='scalar_weight', + reference_fn=partial(bce_loss_reference_fn, weight=scalar_weight)) + ) + + return module_inputs + + +def module_inputs_torch_nn_BCEWithLogitsLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + make_weight = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('weights', {'weight': make_weight((10,))}), + ('scalar_weights', {'weight': make_weight(())}) + ] + + def bce_withlogitsloss_reference_fn(m, p, i, t, reduction='mean', weight=None): + # TODO: add pos_weight to the definition here and corresponding SampleInputs + max_val = (-i).clamp(min=0) + result = (1 - t).mul_(i).add_(max_val).add_((-max_val).exp_().add_((-i - max_val).exp_()).log_()) + + if weight is not None: + result = result * weight + + if reduction == 'none': + return result + elif reduction == 'mean': + return result.sum() / i.numel() + else: + return result.sum() + + module_inputs = [] + for desc, constructor_kwargs in cases: + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((15, 10), low=1e-2, high=1 - 1e-2), + make_target((15, 10)).gt(0).to(dtype)), + desc=desc, + reference_fn=partial(bce_withlogitsloss_reference_fn, **constructor_kwargs)) + ) + + return module_inputs + + +def module_inputs_torch_nn_CrossEntropyLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=torch.long, requires_grad=False) + make_weight = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + reductions: list[str] = ['mean', 'sum', 'none'] + cases: list[tuple[str, dict]] = [ + ('', {}), + ('weights', {'weight': make_weight((3,))}), + ('ignore_index', {'ignore_index': 1}), + ('label_smoothing', {'label_smoothing': 0.15}), + ('ignore_index_label_smoothing', {'ignore_index': 1, 'label_smoothing': 0.15}) + ] + + module_inputs = [] + for reduction, (desc, constructor_kwargs) in product(reductions, cases): + def reference_fn(m, p, i, t, reduction=reduction, constructor_kwargs=constructor_kwargs): + return cross_entropy_loss_reference(i, t, reduction=reduction, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(reduction=reduction, **constructor_kwargs), + forward_input=FunctionInput(make_input((2, 3, 5, 5)), + make_target((2, 5, 5), low=0, high=3)), + desc=f"4d_{desc}_{reduction}", + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(reduction=reduction, **constructor_kwargs), + forward_input=FunctionInput(make_input((2, 3, 5)), + make_target((2, 5), low=0, high=3)), + desc=f"3d_{desc}_{reduction}", + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(reduction=reduction, **constructor_kwargs), + forward_input=FunctionInput(make_input((2, 3)), + make_target((2), low=0, high=3)), + desc=f"2d_{desc}_{reduction}", + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(reduction=reduction, **constructor_kwargs), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 2, 2)), + make_target((2, 5, 5, 2, 2), low=0, high=3)), + desc=f"higher_dim_{desc}_{reduction}", + reference_fn=reference_fn) + ) + + if constructor_kwargs.get('ignore_index', None) is None: + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(reduction=reduction, **constructor_kwargs), + forward_input=FunctionInput(make_input((5, 3, 4, 2)), + make_input((5, 3, 4, 2)).softmax(dim=1)), + desc=f"4d_prob_target_{desc}_{reduction}", + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(reduction=reduction, **constructor_kwargs), + forward_input=FunctionInput(make_input((5, 3, 4)), + make_input((5, 3, 4)).softmax(dim=1)), + desc=f"3d_prob_target_{desc}_{reduction}", + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(reduction=reduction, **constructor_kwargs), + forward_input=FunctionInput(make_input((5, 3)), + make_input((5, 3)).softmax(dim=1)), + desc=f"2d_prob_target_{desc}_{reduction}", + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(reduction=reduction, **constructor_kwargs), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 2, 2)), + make_input((2, 3, 5, 5, 2, 2)).softmax(dim=1)), + desc=f"higher_dim_prob_target_{desc}_{reduction}", + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(reduction=reduction, **constructor_kwargs), + forward_input=FunctionInput(make_input((3,)), + make_target((), low=0, high=3)), + desc=f"no_batch_dim_{desc}_{reduction}", + reference_fn=partial(no_batch_dim_reference_fn, is_criterion=True)) + ) + + return module_inputs + + +def module_error_inputs_torch_nn_CrossEntropyLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_weight = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + return [ + # Non-floating-point target with same shape as input (soft-target path) + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput( + make_input((2, 3)), + torch.zeros((2, 3), device=device, dtype=torch.long), + ), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="Expected floating point type for target with class probabilities", + ), + # Non-long target with different shape from input (class-index path) + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput( + make_input((2, 3)), + make_input((2,)), + ), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="expected target dtype to be Long or Byte", + ), + # Batch size mismatch + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput( + make_input((2, 3)), + torch.tensor([0, 1, 2], device=device, dtype=torch.long), + ), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=ValueError, + error_regex="Expected input batch_size .* to match target batch_size", + ), + # Weight wrong shape + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(weight=make_weight((2,))), + forward_input=FunctionInput( + make_input((2, 3)), + torch.tensor([0, 1], device=device, dtype=torch.long), + ), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="weight tensor should be defined either for all .* classes or no classes", + ), + # Prob target with ignore_index + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(ignore_index=1), + forward_input=FunctionInput( + make_input((2, 3)), + make_input((2, 3)).softmax(dim=1).detach(), + ), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="ignore_index is not supported for floating point target", + ), + # label_smoothing out of range + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(label_smoothing=1.5), + forward_input=FunctionInput( + make_input((2, 3)), + torch.tensor([0, 1], device=device, dtype=torch.long), + ), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="label_smoothing must be between 0.0 and 1.0", + ), + # Target spatial size mismatch (3D input) + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput( + make_input((2, 3, 5)), + torch.zeros((2, 4), device=device, dtype=torch.long), + ), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="Expected target size", + ), + ] + + +def module_inputs_torch_nn_CTCLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('blank', {'blank': 14}) + ] + target_dtypes = [torch.int, torch.long] + + module_inputs = [] + for target_dtype, (desc, constructor_kwargs) in product(target_dtypes, cases): + def reference_fn(m, p, i, t, il, tl, constructor_kwargs=constructor_kwargs): + return ctcloss_reference(i, t, il, tl, **constructor_kwargs) + + blank = constructor_kwargs.get('blank', 0) + low = 0 if blank == 14 else 1 + high = 14 if blank == 14 else 15 + + module_inputs.append( + ModuleInput( + constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((50, 3, 15)).log_softmax(2), + make_target((3, 30), dtype=target_dtype, low=low, high=high), + (50, 50, 50), (30, 25, 20)), + desc=f'{desc}_lengths_intlists', + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput( + constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((50, 3, 15)).log_softmax(2), + make_target((3, 30), dtype=target_dtype, low=low, high=high), + torch.tensor((50, 50, 50), device=device), + torch.tensor((30, 25, 20), device=device)), + desc=f'{desc}_lengths_tensors', + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput( + constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((50, 3, 15)).log_softmax(2), + make_target((30 + 25 + 20,), dtype=target_dtype, low=low, high=high), + (50, 50, 50), (30, 25, 20)), + desc=f'{desc}_1d_target_lengths_intlists', + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput( + constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((50, 3, 15)).log_softmax(2), + make_target((30 + 25 + 20,), dtype=target_dtype, low=low, high=high), + torch.tensor((50, 50, 50), device=device), + torch.tensor((30, 25, 20), device=device)), + desc=f'{desc}_1d_target_lengths_tensors', + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_GroupNorm(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(3, 6, 1e-3), + forward_input=FunctionInput(make_input((4, 6, 5))), + desc='1d_affine'), + ModuleInput( + constructor_input=FunctionInput(3, 6, 1e-3, bias=False), + forward_input=FunctionInput(make_input((4, 6, 5))), + desc='1d_affine_not_bias'), + ModuleInput( + constructor_input=FunctionInput(3, 12, 1e-3), + forward_input=FunctionInput(make_input((4, 12))), + desc='1d_affine_GN'), + ModuleInput( + constructor_input=FunctionInput(1, 6, 1e-3), + forward_input=FunctionInput(make_input((150, 6))), + desc='1d_affine_large_batch'), + ModuleInput( + constructor_input=FunctionInput(5, 5, 1e-3, False), + forward_input=FunctionInput(make_input((4, 5, 5))), + desc='1d_no_affine_IN'), + ModuleInput( + constructor_input=FunctionInput(1, 10, 1e-3, False), + forward_input=FunctionInput(make_input((4, 10))), + desc='1d_no_affine_LN'), + ModuleInput( + constructor_input=FunctionInput(3, 6, 1e-3), + forward_input=FunctionInput(make_input((4, 6, 2, 3))), + desc='2d_affine'), + ModuleInput( + constructor_input=FunctionInput(3, 9, 1e-3, bias=False), + forward_input=FunctionInput(make_input((4, 9, 2, 3))), + desc='2d_affine_not_bias'), + ModuleInput( + constructor_input=FunctionInput(3, 3, 1e-3, False), + forward_input=FunctionInput(make_input((4, 3, 2, 3))), + desc='2d_no_affine_IN'), + ModuleInput( + constructor_input=FunctionInput(1, 3, 1e-3, False), + forward_input=FunctionInput(make_input((4, 3, 2, 3))), + desc='2d_no_affine_LN'),] + + +def module_error_inputs_torch_nn_GroupNorm(module_info, device, dtype, requires_grad, training, **kwargs): + """ + Error inputs for GroupNorm that test error messages include actual values. + """ + return [ + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(3, 10), # num_groups=3, num_channels=10 + forward_input=FunctionInput(), # Not needed for construction error + ), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=ValueError, + error_regex=r"num_channels \(10\) must be divisible by num_groups \(3\)" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(5, 13), # num_groups=5, num_channels=13 + forward_input=FunctionInput(), # Not needed for construction error + ), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=ValueError, + error_regex=r"num_channels \(13\) must be divisible by num_groups \(5\)" + ), + ] + + +def module_inputs_torch_nn_Hardshrink(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(2.), + forward_input=FunctionInput(make_input((4, 3, 2, 4))), + ), + ModuleInput( + constructor_input=FunctionInput(2.), + forward_input=FunctionInput(make_input(())), + desc='scalar', + ), + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim', + ) + ] + + +def module_inputs_torch_nn_Hardswish(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim', + ), + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 2, 5))), + desc='4d_input') + ] + + +def module_inputs_torch_nn_Hardtanh(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((3, 2, 5))), + reference_fn=lambda m, p, i: i.clamp(-1, 1), + ), + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + reference_fn=lambda m, p, i: i.clamp(-1, 1), + desc='scalar', + ), + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim', + ) + ] + + +def module_inputs_torch_nn_HingeEmbeddingLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('margin', {'margin': 0.5}) + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return hingeembeddingloss_reference(i, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((10,)), + make_target((10,)).gt(0).to(dtype).mul_(2).sub_(1)), + desc=desc, + reference_fn=reference_fn) + ) + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input(()), + make_target(()).gt(0).to(dtype).mul_(2).sub_(1)), + desc=f'scalar_{desc}', + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_HuberLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return huberloss_reference(i, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((5, 10)), + make_input((5, 10))), + desc=desc, + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_InstanceNormNd(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + lazy = kwargs.get('lazy', False) + N = kwargs['N'] + num_features, eps, momentum, affine, track_running_stats = 3, 1e-3, 0.3, False, True + input_no_batch_shape_dict = {1: (3, 15), 2: (3, 6, 6), 3: (3, 4, 4, 4)} + input_no_batch_shape = input_no_batch_shape_dict[N] + input_batch_shape = (4,) + input_no_batch_shape + + return [ + ModuleInput( + constructor_input=( + FunctionInput(eps, momentum) if lazy else FunctionInput(num_features, eps, momentum) + ), + forward_input=FunctionInput(make_input(input_batch_shape))), + ModuleInput( + constructor_input=( + FunctionInput(eps, momentum, affine, track_running_stats) if lazy else + FunctionInput(num_features, eps, momentum, affine, track_running_stats) + ), + forward_input=FunctionInput(make_input(input_batch_shape)), + desc='tracking_stats'), + ModuleInput( + constructor_input=( + FunctionInput(eps, momentum) if lazy else FunctionInput(num_features, eps, momentum) + ), + forward_input=FunctionInput(make_input(input_no_batch_shape)), + reference_fn=no_batch_dim_reference_fn, + desc='tracking_stats_no_batch_dim'), + ModuleInput( + constructor_input=( + FunctionInput(eps, momentum, affine, track_running_stats) if lazy else + FunctionInput(num_features, eps, momentum, affine, track_running_stats) + ), + forward_input=FunctionInput(make_input(input_no_batch_shape)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput( + constructor_input=( + FunctionInput(eps, momentum, affine=True) if lazy else + FunctionInput(num_features, eps, momentum, affine=True) + ), + forward_input=FunctionInput(make_input(input_batch_shape)), + desc='affine'), + ModuleInput( + constructor_input=( + FunctionInput(eps, momentum, affine=True, bias=False) if lazy else + FunctionInput(num_features, eps, momentum, affine=True, bias=False) + ), + forward_input=FunctionInput(make_input(input_batch_shape)), + desc='affine_not_bias'),] + +def module_inputs_torch_nn_LayerNorm(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput([5], 1e-3), + forward_input=FunctionInput(make_input((4, 5, 5))), + desc='1d_elementwise_affine'), + ModuleInput( + constructor_input=FunctionInput([5], 1e-3), + forward_input=FunctionInput(make_input((128, 5, 5))), + desc='1d_elementwise_affine_large_batch'), + ModuleInput( + constructor_input=FunctionInput([5], 1e-3, False), + forward_input=FunctionInput(make_input((4, 5, 5))), + desc='1d_no_elementwise_affine'), + ModuleInput( + constructor_input=FunctionInput([2, 2, 5], 1e-3), + forward_input=FunctionInput(make_input((4, 2, 2, 5))), + desc='3d_elementwise_affine'), + ModuleInput( + constructor_input=FunctionInput([2, 2, 5], 1e-3, False), + forward_input=FunctionInput(make_input((4, 2, 2, 5))), + desc='3d_no_elementwise_affine'), + ModuleInput( + constructor_input=FunctionInput([5], 1e-3), + forward_input=FunctionInput(make_input((0, 5))), + desc='1d_empty_elementwise_affine'), + ModuleInput( + constructor_input=FunctionInput([2, 2, 5], 1e-3, elementwise_affine=True, bias=False), + forward_input=FunctionInput(make_input((4, 2, 2, 5))), + desc='3d_elementwise_affine_no_bias'), + ] + +def module_inputs_torch_nn_RMSNorm(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def rms_norm_reference_fn(m, p, i): + eps = m.eps + if eps is None: + eps = torch.finfo(i.dtype).eps + ndim = i.ndim + normalized_shape = m.normalized_shape + weight = m.weight + dims = [ndim - i - 1 for i in range(len(normalized_shape))] + upcasted_i = i.float() + result = upcasted_i * torch.rsqrt(upcasted_i.pow(2).mean(dim=dims, keepdim=True) + m.eps) + if weight is not None: + result *= weight + return result.type_as(i) + + return [ + ModuleInput( + constructor_input=FunctionInput([5], 1e-3), + forward_input=FunctionInput(make_input((4, 5, 5))), + desc='1d_elementwise_affine', + reference_fn=rms_norm_reference_fn), + ModuleInput( + constructor_input=FunctionInput([5], 1e-3), + forward_input=FunctionInput(make_input((128, 5, 5))), + desc='1d_elementwise_affine_large_batch', + reference_fn=rms_norm_reference_fn), + ModuleInput( + constructor_input=FunctionInput([5], 1e-3, False), + forward_input=FunctionInput(make_input((4, 5, 5))), + desc='1d_no_elementwise_affine', + reference_fn=rms_norm_reference_fn), + ModuleInput( + constructor_input=FunctionInput([2, 2, 5], 1e-3), + forward_input=FunctionInput(make_input((4, 2, 2, 5))), + desc='3d_elementwise_affine', + reference_fn=rms_norm_reference_fn), + ModuleInput( + constructor_input=FunctionInput([2, 2, 5], 1e-3, False), + forward_input=FunctionInput(make_input((4, 2, 2, 5))), + desc='3d_no_elementwise_affine', + reference_fn=rms_norm_reference_fn), + ModuleInput( + constructor_input=FunctionInput([5], 1e-3), + forward_input=FunctionInput(make_input((0, 5))), + desc='1d_empty_elementwise_affine', + reference_fn=rms_norm_reference_fn), + ] + + +def module_inputs_torch_nn_LocalResponseNorm(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(3,), + forward_input=FunctionInput(make_input((1, 5, 7))), + desc='1d'), + ModuleInput( + constructor_input=FunctionInput(2,), + forward_input=FunctionInput(make_input((1, 5, 7, 7))), + desc='2d_uneven_pad'), + ModuleInput( + constructor_input=FunctionInput(1, 1., 0.5, 2.), + forward_input=FunctionInput(make_input((1, 5, 7, 7, 7))), + desc='3d_custom_params'), + ] + + +def module_inputs_torch_nn_LPPool1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1.5, 2), + forward_input=FunctionInput(make_input((1, 3, 7))), + desc='norm'), + ModuleInput( + constructor_input=FunctionInput(2, 2, 3), + forward_input=FunctionInput(make_input((1, 3, 7)))), + ModuleInput( + constructor_input=FunctionInput(2, 2, 3), + forward_input=FunctionInput(make_input((3, 7))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ] + + + +def module_inputs_torch_nn_LPPool2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(2, 2, 2), + forward_input=FunctionInput(make_input((1, 3, 7, 7)))), + ModuleInput( + constructor_input=FunctionInput(2, 2, 2), + forward_input=FunctionInput(make_input((3, 7, 7))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput( + constructor_input=FunctionInput(1.5, 2), + forward_input=FunctionInput(make_input((1, 3, 7, 7))), + desc='norm'), + ] + + +def module_inputs_torch_nn_LPPool3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(2, 2, 2), + forward_input=FunctionInput(make_input((1, 3, 7, 7, 7)))), + ModuleInput( + constructor_input=FunctionInput(2, 2, 2), + forward_input=FunctionInput(make_input((3, 7, 7, 7))), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim'), + ModuleInput( + constructor_input=FunctionInput(1.5, 2), + forward_input=FunctionInput(make_input((1, 3, 7, 7, 7))), + desc='norm'), + ] + + +def module_inputs_torch_nn_MaxPool1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(4), + forward_input=FunctionInput(make_input((2, 10, 4))), + desc='3d_input'), + ModuleInput( + constructor_input=FunctionInput(4, 4), + forward_input=FunctionInput(make_input((2, 10, 4))), + desc='stride'), + ModuleInput( + constructor_input=FunctionInput(4, return_indices=True), + forward_input=FunctionInput(make_input((2, 10, 4))), + desc='return_indices'), + ] + + +def module_inputs_torch_nn_MaxPool2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput((3, 3), (2, 2), (1, 1)), + forward_input=FunctionInput(make_input((3, 7, 7))), + desc='3d_input'), + ModuleInput( + constructor_input=FunctionInput((3, 3), (2, 2), (1, 1)), + forward_input=FunctionInput(make_input((1, 3, 7, 7))), + desc='4d_input'), + ModuleInput( + constructor_input=FunctionInput((3, 3), (2, 2), (1, 1), return_indices=True), + forward_input=FunctionInput(make_input((1, 3, 7, 7))), + desc='return_indices'), + ] + + +def module_error_inputs_torch_nn_MaxPool2d(module_info, device, dtype, requires_grad, training, **kwargs): + """ + Error inputs for MaxPool2d that test error messages for invalid inputs. + """ + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + # Wrong input dimensions: 2D input instead of 3D/4D + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(2), + forward_input=FunctionInput(make_input((3, 4))), # 2D input + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"non-empty 3D or 4D \(batch mode\) tensor expected for input" + ), + # Wrong input dimensions: 5D input + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(2), + forward_input=FunctionInput(make_input((1, 2, 3, 4, 5))), # 5D input + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"non-empty 3D or 4D \(batch mode\) tensor expected for input" + ), + # Invalid padding: padding > kernel_size / 2 + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(3, padding=5), # kernel=3, pad=5 > 3/2 + forward_input=FunctionInput(make_input((1, 1, 10, 10))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"pad should be at most half of effective kernel size" + ), + ] + + +def module_inputs_torch_nn_MaxPool3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput((2, 2, 2)), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5)))), + ModuleInput( + constructor_input=FunctionInput(2, (2, 2, 2)), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5))), + desc='stride'), + ModuleInput( + constructor_input=FunctionInput(2, 2, (1, 1, 1)), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5))), + desc='stride_padding'), + ModuleInput( + constructor_input=FunctionInput(2, 2, (1, 1, 1), return_indices=True), + forward_input=FunctionInput(make_input((2, 3, 5, 5, 5))), + desc='return_indices'), + ] + + +def module_inputs_torch_nn_FractionalMaxPool2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_random_samples(): + return torch.empty((1, 3, 2), dtype=torch.double, device=device).uniform_() + + return [ + ModuleInput( + constructor_input=FunctionInput(2, output_ratio=0.5, _random_samples=make_random_samples()), + forward_input=FunctionInput(make_input((1, 3, 5, 7))), + desc='ratio'), + ModuleInput( + constructor_input=FunctionInput((2, 3), output_size=(4, 3), _random_samples=make_random_samples()), + forward_input=FunctionInput(make_input((1, 3, 7, 6))), + desc='size'), + ModuleInput( + constructor_input=FunctionInput( + 2, output_ratio=0.5, _random_samples=make_random_samples(), return_indices=True + ), + forward_input=FunctionInput(make_input((1, 3, 5, 7))), + desc='ratio_return_indices'), + ModuleInput( + constructor_input=FunctionInput(2, output_ratio=0.5, _random_samples=make_random_samples()), + forward_input=FunctionInput(make_input((3, 5, 7))), + reference_fn=no_batch_dim_reference_fn, + desc='ratio_no_batch_dim'), + ModuleInput( + constructor_input=FunctionInput((2, 3), output_size=(4, 3), _random_samples=make_random_samples()), + forward_input=FunctionInput(make_input((3, 7, 6))), + reference_fn=no_batch_dim_reference_fn, + desc='size_no_batch_dim'), + ] + + +def module_inputs_torch_nn_FractionalMaxPool3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def make_random_samples(): + return torch.empty((2, 4, 3), dtype=torch.double, device=device).uniform_() + + return [ + ModuleInput( + constructor_input=FunctionInput(2, output_ratio=0.5, _random_samples=make_random_samples()), + forward_input=FunctionInput(make_input((2, 4, 5, 5, 5))), + desc='ratio'), + ModuleInput( + constructor_input=FunctionInput((2, 2, 2), output_size=(4, 4, 4), _random_samples=make_random_samples()), + forward_input=FunctionInput(make_input((2, 4, 7, 7, 7))), + desc='size'), + ModuleInput( + constructor_input=FunctionInput((4, 2, 3), output_size=(10, 3, 2), _random_samples=make_random_samples()), + forward_input=FunctionInput(make_input((2, 4, 16, 7, 5))), + desc='asymsize'), + ModuleInput( + constructor_input=FunctionInput( + 2, output_ratio=0.5, _random_samples=make_random_samples(), return_indices=True + ), + forward_input=FunctionInput(make_input((2, 4, 5, 5, 5))), + desc='ratio_return_indices'), + ModuleInput( + constructor_input=FunctionInput(2, output_ratio=0.5, _random_samples=make_random_samples()), + forward_input=FunctionInput(make_input((4, 5, 5, 5))), + reference_fn=no_batch_dim_reference_fn, + desc='ratio_no_batch_dim'), + ModuleInput( + constructor_input=FunctionInput((2, 2, 2), output_size=(4, 4, 4), _random_samples=make_random_samples()), + forward_input=FunctionInput(make_input((4, 7, 7, 7))), + reference_fn=no_batch_dim_reference_fn, + desc='size_no_batch_dim'), + ] + + +def module_inputs_torch_nn_Sigmoid(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + desc='scalar' + ), + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim', + ), + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4, 5))), + desc='channels_last_mem_format' + ), + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 3, 4, 5))), + desc='channels_last_3d_mem_format' + ) + ] + + +def module_inputs_torch_nn_LogSigmoid(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(())), + reference_fn=lambda m, p, i: i.sigmoid().log(), + desc='scalar' + ), + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input((2, 3, 4))), + reference_fn=lambda m, p, i: i.sigmoid().log(), + ), + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput(make_input(4)), + reference_fn=no_batch_dim_reference_fn, + desc='no_batch_dim', + ), + ] + + +def module_inputs_torch_nn_MarginRankingLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=torch.long, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('margin', {'margin': 0.5}) + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i1, i2, t, constructor_kwargs=constructor_kwargs): + return marginrankingloss_reference(i1, i2, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((50,)), make_input((50,)), + make_target((50,)).sign()), + desc=desc, + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_MultiLabelMarginLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=torch.long, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return multilabelmarginloss_reference(i, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((10,)), + make_target((10), low=0, high=10)), + desc=f'1d_{desc}', + reference_fn=reference_fn) + ) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((5, 10)), + make_target((5, 10), low=0, high=10)), + desc=desc, + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_MultiMarginLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=torch.long, requires_grad=False) + make_weight = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('p', {'p': 2}), + ('margin', {'margin': 0.5}), + ('weights', {'weight': make_weight(10)}) + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return multimarginloss_reference(i, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((5, 10)), + make_target((5), low=0, high=10)), + desc=desc, + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_error_inputs_torch_nn_MultiMarginLoss(module_info, device, dtype, requires_grad, training, **kwargs): + """ + Error inputs for MultiMarginLoss that test the improved error message + for inconsistent target size. + Regression test for issue #106251. + """ + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + # Test: target size doesn't match batch size (5 vs 3) + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(), + forward_input=FunctionInput( + make_input((5, 10)), # input: batch_size=5, num_classes=10 + torch.tensor([0, 1, 2], device=device, dtype=torch.long), # target: wrong size (3) + ), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"target tensor should be 1-D with size equal to.*Expected target size \[5\].*but got \[3\]" + ), + ] + + +def module_inputs_torch_nn_MultiLabelSoftMarginLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=torch.long, requires_grad=False) + make_weight = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ('weight', {'weight': make_weight(10)}), + ] + + def multilabelsoftmargin_loss_reference_fn(m, p, i, t, reduction='mean', weight=None): + result = t * i.sigmoid().log() + (1 - t) * (-i).sigmoid().log() + if weight is not None: + result *= weight + result = (-result).sum(i.dim() - 1) / i.size(-1) + + if reduction == 'none': + return result + elif reduction == 'mean': + return result.mean() + else: + return result.sum() + + module_inputs = [] + for desc, constructor_kwargs in cases: + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((5, 10)), + make_target((5, 10), low=0, high=2)), + desc=desc, + reference_fn=partial(multilabelsoftmargin_loss_reference_fn, **constructor_kwargs)) + ) + + return module_inputs + + +def module_inputs_torch_nn_SoftMarginLoss(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + make_target = partial(make_tensor, device=device, dtype=dtype, requires_grad=False) + + cases: list[tuple[str, dict]] = [ + ('', {}), + ('reduction_sum', {'reduction': 'sum'}), + ('reduction_mean', {'reduction': 'mean'}), + ('reduction_none', {'reduction': 'none'}), + ] + + module_inputs = [] + for desc, constructor_kwargs in cases: + def reference_fn(m, p, i, t, constructor_kwargs=constructor_kwargs): + return softmarginloss_reference(i, t, **constructor_kwargs) + + module_inputs.append( + ModuleInput(constructor_input=FunctionInput(**constructor_kwargs), + forward_input=FunctionInput(make_input((5, 5)), + make_target((5, 5)).sign()), + desc=desc, + reference_fn=reference_fn) + ) + + return module_inputs + + +def module_inputs_torch_nn_TransformerEncoder(module_info, device, dtype, requires_grad, training, **kwargs): + # Reuse the TransformerEncoderLayer samples since the forward args are nearly the same. + samples = [] + for layer_module_input in module_inputs_torch_nn_TransformerEncoderLayer( + None, device, dtype, requires_grad, training): + # Construct a TransformerEncoderLayer object to pass to TransformerEncoder. + l_args, l_kwargs = (layer_module_input.constructor_input.args, + layer_module_input.constructor_input.kwargs) + l_kwargs['device'] = device + l_kwargs['dtype'] = dtype + encoder_layer = torch.nn.TransformerEncoderLayer(*l_args, **l_kwargs) + num_layers = 2 + # Note: TransformerEncoderLayer takes a "src_mask" while + # TransformerEncoder takes a "mask"; rename kwarg appropriately. + forward_input = layer_module_input.forward_input + if 'src_mask' in forward_input.kwargs: + forward_input.kwargs['mask'] = forward_input.kwargs['src_mask'] + del forward_input.kwargs['src_mask'] + samples.append(ModuleInput( + constructor_input=FunctionInput(encoder_layer, num_layers), + forward_input=forward_input, + desc=layer_module_input.desc + )) + return samples + +def module_inputs_torch_nn_TransformerEncoderLayer(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + samples = [ + ModuleInput( + constructor_input=FunctionInput(4, 2, 16, 0.0), + forward_input=FunctionInput( + make_input((2, 3, 4)) + ), + desc='relu_activation' + ), + ModuleInput( + constructor_input=FunctionInput(4, 2, 8, 0.0, F.gelu), + forward_input=FunctionInput( + make_input((2, 3, 4)) + ), + desc='gelu_activation' + ), + ModuleInput( + constructor_input=FunctionInput(4, 2, 8, 0.0, bias=False), + forward_input=FunctionInput( + make_input((2, 3, 4)) + ), + desc='no_bias' + ), ] + + # Samples below are for validating the no-batch-dim support. + key_padding_masks = (None, torch.tensor([False, False, True], device=device, dtype=torch.bool)) + attn_masks = (None, torch.tensor([False, False, True], device=device, dtype=torch.bool).expand((3, 3))) + for src_mask, src_key_padding_mask, norm_first, batch_first, bias in \ + itertools.product(attn_masks, key_padding_masks, (True, False), (True, False), (True, False)): + samples.append( + ModuleInput( + constructor_input=FunctionInput(d_model=4, nhead=2, dim_feedforward=8, + dropout=0.0, batch_first=batch_first, + norm_first=norm_first, bias=bias), + forward_input=FunctionInput( + make_input((3, 4)), src_mask=src_mask, src_key_padding_mask=src_key_padding_mask + ), + reference_fn=partial(no_batch_dim_reference_fn, + batch_first=batch_first, kwargs_to_batchify={'src_key_padding_mask': 0}), + desc=f'no_batch_dim_batch_first_{batch_first}' + )) + + # Samples below where we pass reference_fn are for validating the fast path, + # since the fast path requires no_grad mode, we run the fast path in .eval() + # and no_grad() in the reference_fn and verify that against the results in train mode. + def fast_path_reference_fn(module, parameters, *args, **kwargs): + if not module.training: + raise AssertionError("Expected module.training to be True") + module.train(False) + with torch.no_grad(): + output = module(*args, **kwargs) + module.train(True) + return output + + if training: + for norm_first, bias in itertools.product((True, False), (True, False)): + samples.append( + ModuleInput( + constructor_input=FunctionInput( + 4, 2, 8, dropout=0.0, batch_first=True, norm_first=norm_first, bias=bias + ), + forward_input=FunctionInput( + make_input((2, 3, 4)), + ), + # fastpath doesn't run when bias=False + reference_fn=fast_path_reference_fn if bias else None, + desc=f'fastpath_{bias}_norm_first_{norm_first}' + ) + ) + + return samples + + +def module_inputs_torch_nn_TransformerDecoderLayer(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + samples = [ + ModuleInput( + constructor_input=FunctionInput(4, 2, 16, 0.0), + forward_input=FunctionInput( + make_input((2, 3, 4)), make_input((2, 3, 4)) + ), + desc='relu_activation' + ), + ModuleInput( + constructor_input=FunctionInput(4, 2, 8, 0.0, F.gelu), + forward_input=FunctionInput( + make_input((2, 3, 4)), make_input((2, 3, 4)) + ), + desc='gelu_activation' + ), + ModuleInput( + constructor_input=FunctionInput(4, 2, 8, 0.0, bias=False), + forward_input=FunctionInput( + make_input((2, 3, 4)), make_input((2, 3, 4)) + ), + desc='no_bias' + ), ] + + key_padding_masks = (None, torch.tensor([False, False, True], device=device, dtype=torch.bool)) + attn_masks = (None, torch.tensor([False, False, True], device=device, dtype=torch.bool).expand((3, 3))) + for tgt_mask, tgt_key_padding_mask, norm_first, bias, batch_first in \ + itertools.product(attn_masks, key_padding_masks, (True, False), (True, False), (True, False)): + # Using same mask for tgt and memory + memory_mask = tgt_mask + memory_key_padding_mask = tgt_key_padding_mask + samples.append( + ModuleInput( + constructor_input=FunctionInput(d_model=4, nhead=2, dim_feedforward=8, + dropout=0.0, batch_first=batch_first, + norm_first=norm_first, bias=bias), + forward_input=FunctionInput( + make_input((3, 4)), make_input((3, 4)), tgt_mask=tgt_mask, memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask + ), + reference_fn=partial(no_batch_dim_reference_fn, + batch_first=batch_first, + kwargs_to_batchify={'tgt_key_padding_mask': 0, 'memory_key_padding_mask': 0}), + desc=f'no_batch_dim_batch_first_{batch_first}' + )) + src, tgt = make_input((2, 3, 4)), make_input((2, 3, 4)) + if not batch_first: + src, tgt = src.transpose(0, 1), tgt.transpose(0, 1) + if tgt_key_padding_mask is not None: + memory_key_padding_mask, tgt_key_padding_mask = (tgt_key_padding_mask.expand(2, 3),) * 2 + samples.append( + ModuleInput( + constructor_input=FunctionInput(d_model=4, nhead=2, dim_feedforward=8, + dropout=0.0, batch_first=batch_first, + norm_first=norm_first, bias=bias), + forward_input=FunctionInput( + src, tgt, tgt_mask=tgt_mask, memory_mask=memory_mask, + tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask + ), + desc=f'norm_first_{norm_first}_batch_first_{batch_first}_bias_{bias}' + )) + + return samples + + +def module_inputs_torch_nn_Transformer(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + samples = [] + # Samples below are for validating the no-batch-dim support. + key_padding_masks = (None, torch.tensor([False, False, True], device=device, dtype=torch.bool)) + attn_masks = (None, torch.tensor([False, False, True], device=device, dtype=torch.bool).expand((3, 3))) + for mask, key_padding_mask, norm_first, bias, batch_first in \ + itertools.product(attn_masks, key_padding_masks, (True, False), (True, False), (True, False)): + # Using same mask for tgt and memory + src_mask , tgt_mask = (mask,) * 2 + src_key_padding_mask, tgt_key_padding_mask = (key_padding_mask,) * 2 + samples.append( + ModuleInput( + constructor_input=FunctionInput(d_model=4, nhead=2, dim_feedforward=8, + num_encoder_layers=1, num_decoder_layers=1, + dropout=0.0, batch_first=batch_first, norm_first=norm_first, bias=bias), + forward_input=FunctionInput( + make_input((3, 4)), make_input((3, 4)), tgt_mask=tgt_mask, src_mask=src_mask, + tgt_key_padding_mask=tgt_key_padding_mask, src_key_padding_mask=src_key_padding_mask + ), + reference_fn=partial(no_batch_dim_reference_fn, + batch_first=batch_first, + kwargs_to_batchify={'tgt_key_padding_mask': 0, 'src_key_padding_mask': 0}), + desc=f'no_batch_dim_batch_first_{batch_first}' + )) + + src, tgt = make_input((2, 3, 4)), make_input((2, 3, 4)) + if not batch_first: + src = src.transpose(0, 1) + tgt = tgt.transpose(0, 1) + if key_padding_mask is not None: + src_key_padding_mask, tgt_key_padding_mask = (key_padding_mask.expand(2, 3),) * 2 + + samples.append( + ModuleInput( + constructor_input=FunctionInput(d_model=4, nhead=2, dim_feedforward=8, + num_encoder_layers=1, num_decoder_layers=1, + dropout=0.0, batch_first=batch_first, norm_first=norm_first, bias=bias), + forward_input=FunctionInput( + src, tgt, tgt_mask=tgt_mask, src_mask=src_mask, + tgt_key_padding_mask=tgt_key_padding_mask, src_key_padding_mask=src_key_padding_mask + ), + )) + return samples + + +def module_inputs_torch_nn_Embedding(module_info, device, dtype, requires_grad, training, **kwargs): + make_empty = partial(torch.empty, device=device, dtype=torch.long, requires_grad=False) + return [ + ModuleInput( + constructor_input=FunctionInput(num_embeddings=4, embedding_dim=3), + forward_input=FunctionInput(make_empty(2, 3).random_(4)) + ), + ModuleInput( + constructor_input=FunctionInput(num_embeddings=4, embedding_dim=3), + forward_input=FunctionInput(make_empty(1, 512).random_(4).expand(7, 512)), + desc='discontiguous' + ), + ] + + +def module_error_inputs_torch_nn_Embedding(module_info, device, dtype, requires_grad, training, **kwargs): + """ + Error inputs for Embedding that test error messages for invalid inputs. + """ + samples = [] + + # Out of range indices: index exceeds num_embeddings + # Only test on CPU - CUDA triggers kernel assertion instead of Python exception + if torch.device(device).type == 'cpu': + samples.append( + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(num_embeddings=10, embedding_dim=3), + forward_input=FunctionInput(torch.tensor([0, 5, 15], device=device, dtype=torch.long)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=IndexError, + error_regex=r"index out of range in self" + ) + ) + + # Float indices: wrong dtype for indices (works on all devices) + samples.append( + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(num_embeddings=10, embedding_dim=3), + forward_input=FunctionInput(torch.tensor([1.5, 2.5], device=device, dtype=torch.float32)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"Expected tensor for argument.*indices.*to have.*scalar type.*Long.*Int" + ) + ) + + # Negative num_embeddings (construction error, device-independent) + samples.append( + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(num_embeddings=-1, embedding_dim=3), + forward_input=FunctionInput(), + ), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=RuntimeError, + error_regex=r"Trying to create tensor with negative dimension" + ) + ) + + return samples + + + +def module_inputs_torch_nn_MultiheadAttention(module_info, device, dtype, requires_grad, training, **kwargs): + # Currently all samples below are for validating the no-batch-dim support. + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + samples = [] + bool_vals = (True, False) + key_padding_masks = (None, torch.tensor([False, False, True], device=device, dtype=torch.bool)) + attn_masks = (None, torch.tensor([False, False, True], device=device, dtype=torch.bool).expand((3, 3, 3))) + products = itertools.product(bool_vals, bool_vals, bool_vals, key_padding_masks, attn_masks) + for bias, add_bias_kv, add_zero_attn, key_padding_mask, attn_mask in products: + samples.append( + ModuleInput( + constructor_input=FunctionInput(embed_dim=3, num_heads=3, batch_first=True, + bias=bias, add_bias_kv=add_bias_kv, add_zero_attn=add_zero_attn), + forward_input=FunctionInput(make_input((3, 3)), make_input((3, 3)), make_input((3, 3)), + key_padding_mask=key_padding_mask, attn_mask=attn_mask), + reference_fn=no_batch_dim_reference_mha, + ) + ) + samples.append( + ModuleInput( + constructor_input=FunctionInput(embed_dim=3, num_heads=3, batch_first=False, + bias=bias, add_bias_kv=add_bias_kv, add_zero_attn=add_zero_attn), + forward_input=FunctionInput(make_input((3, 3)), make_input((3, 3)), make_input((3, 3)), + key_padding_mask=key_padding_mask, attn_mask=attn_mask), + reference_fn=partial(no_batch_dim_reference_mha, batch_first=False), + ) + ) + + return samples + + +def module_inputs_torch_nn_RNN_GRU_Cell(module_info, device, dtype, requires_grad, training, **kwargs): + # Currently all samples below are for validating the no-batch-dim support. + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + samples = [ + ModuleInput( + constructor_input=FunctionInput(5, 10), + forward_input=FunctionInput(make_input(5), make_input(10)), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput(5, 10, bias=True), + forward_input=FunctionInput(make_input(5), make_input(10)), + reference_fn=no_batch_dim_reference_fn, + ) + ] + + is_rnn = kwargs.get('is_rnn', False) + if is_rnn: + # RNN also supports `nonlinearity` argument. + # `tanh` is the default, so we check with `relu` + samples.append( + ModuleInput( + constructor_input=FunctionInput(5, 10, bias=True, nonlinearity='relu'), + forward_input=FunctionInput(make_input(5), make_input(10)), + reference_fn=no_batch_dim_reference_fn, + ) + ) + + return samples + + +def module_inputs_torch_nn_LSTMCell(module_info, device, dtype, requires_grad, training, **kwargs): + # Currently all samples below are for validating the no-batch-dim support. + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + samples = ( + ModuleInput( + constructor_input=FunctionInput(5, 10), + forward_input=FunctionInput(make_input(5), (make_input(10), make_input(10))), + reference_fn=no_batch_dim_reference_lstmcell, + ), + ModuleInput( + constructor_input=FunctionInput(5, 10, bias=True), + forward_input=FunctionInput(make_input(5), (make_input(10), make_input(10))), + reference_fn=no_batch_dim_reference_lstmcell, + ), + ) + + return samples + +def make_packed_sequence(inp, batch_sizes): + required_grad = inp.requires_grad + inp.requires_grad_(False) # user won't have access to inp so won't be able to get its grads + seq = pack_padded_sequence(inp, batch_sizes) + seq.data.requires_grad_(required_grad) + return seq + + +def module_inputs_torch_nn_RNN_GRU(module_info, device, dtype, requires_grad, training, with_packed_sequence=False, **kwargs): + # Currently all samples below are for validating the no-batch-dim support. + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + is_rnn = kwargs['is_rnn'] + nonlinearity = ('relu', 'tanh') + bias = (False, True) + batch_first = (False, True) + bidirectional = (False, True) + + samples = [] + if is_rnn: + prod_gen = product(nonlinearity, bias, batch_first, bidirectional) + else: + prod_gen = product(bias, batch_first, bidirectional) + + for args in prod_gen: + if is_rnn: + nl, b, b_f, bidir = args + else: + b, b_f, bidir = args + + cons_args = {'input_size': 2, 'hidden_size': 2, 'num_layers': 2, + 'batch_first': b_f, 'bias': b, 'bidirectional': bidir} + cons_args_hidden = {'input_size': 2, 'hidden_size': 3, 'num_layers': 2, + 'batch_first': b_f, 'bias': b, 'bidirectional': bidir} + + if is_rnn: + cons_args['nonlinearity'] = nl + cons_args_hidden['nonlinearity'] = nl + samples.append( + ModuleInput( + constructor_input=FunctionInput(**cons_args), + forward_input=FunctionInput(make_input((3, 2))), + reference_fn=partial(no_batch_dim_reference_rnn_gru, batch_first=b_f), + ) + ) + samples.append( + ModuleInput( + constructor_input=FunctionInput(**cons_args_hidden), + forward_input=FunctionInput(make_input((3, 2)), make_input((4 if bidir else 2, 3))), + reference_fn=partial(no_batch_dim_reference_rnn_gru, batch_first=b_f), + ) + ) + if with_packed_sequence: + samples.append( + ModuleInput( + constructor_input=FunctionInput(**cons_args), + forward_input=FunctionInput(make_packed_sequence(make_input((5, 2, 2)), torch.tensor([5, 3]))), + reference_fn=partial(no_batch_dim_reference_rnn_gru, batch_first=b_f), + ) + ) + samples.append( + ModuleInput( + constructor_input=FunctionInput(**cons_args), + forward_input=FunctionInput(make_packed_sequence(make_input((5, 5, 2)), torch.tensor([5, 3, 3, 2, 2]))), + reference_fn=partial(no_batch_dim_reference_rnn_gru, batch_first=b_f), + ) + ) + + return samples + + +def module_inputs_torch_nn_LSTM(module_info, device, dtype, requires_grad, training, **kwargs): + # Currently all samples below are for validating the no-batch-dim support. + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + bias = (False, True) + batch_first = (False, True) + bidirectional = (False, True) + proj_sizes = (0, 2) + + samples = [] + prod_gen = product(bias, batch_first, bidirectional, proj_sizes) + + for args in prod_gen: + b, b_f, bidir, proj_size = args + hidden_size = 3 + cons_args = {'input_size': 2, 'hidden_size': hidden_size, 'num_layers': 2, 'proj_size': proj_size, + 'batch_first': b_f, 'bias': b, 'bidirectional': bidir} + cons_args_hidden = {'input_size': 2, 'hidden_size': hidden_size, 'num_layers': 2, 'proj_size': proj_size, + 'batch_first': b_f, 'bias': b, 'bidirectional': bidir} + + samples.append( + ModuleInput( + constructor_input=FunctionInput(**cons_args), + forward_input=FunctionInput(make_input((2, 2))), + reference_fn=partial(no_batch_dim_reference_lstm, batch_first=b_f), + ) + ) + + h_out = proj_size if proj_size > 0 else hidden_size + hx = (make_input((4 if bidir else 2, h_out)), make_input((4 if bidir else 2, hidden_size))) + samples.append( + ModuleInput( + constructor_input=FunctionInput(**cons_args_hidden), + forward_input=FunctionInput(make_input((3, 2)), hx), + reference_fn=partial(no_batch_dim_reference_lstm, batch_first=b_f), + ) + ) + + + return samples + + + +def module_inputs_torch_nn_ReflectionPad1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((2, 3))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2)), + forward_input=FunctionInput(make_input((2, 3, 4))), + ), + ] + +def module_inputs_torch_nn_ReflectionPad2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((3, 4, 5))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 3, 4)), + forward_input=FunctionInput(make_input((3, 4, 5, 6))), + ), + ] + +def module_inputs_torch_nn_ReflectionPad3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((2, 3, 4, 5))), + reference_fn=no_batch_dim_reference_fn + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 1, 2, 1, 2)), + forward_input=FunctionInput(make_input((3, 3, 3, 3, 3))), + ), + ] + +def module_inputs_torch_nn_ReplicationPad1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((3, 4))), + reference_fn=no_batch_dim_reference_fn + ), + ModuleInput( + constructor_input=FunctionInput((1, 2)), + forward_input=FunctionInput(make_input((3, 4, 5))), + ), + ] + +def module_inputs_torch_nn_ReplicationPad2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((3, 4, 5))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 3, 4)), + forward_input=FunctionInput(make_input((3, 4, 5, 6))), + ), + ] + +def module_inputs_torch_nn_ReplicationPad3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((3, 4, 5, 6))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 3, 4, 5, 6)), + forward_input=FunctionInput(make_input((3, 4, 5, 6, 7))), + ), + ] + +def module_inputs_torch_nn_ZeroPad1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((3, 4))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2)), + forward_input=FunctionInput(make_input((3, 4, 5))), + ), + ] + +def module_inputs_torch_nn_ZeroPad2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((1, 2, 3))), + reference_fn=no_batch_dim_reference_fn + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 3, 4)), + forward_input=FunctionInput(make_input((1, 2, 3, 4))), + ), + ] + +def module_inputs_torch_nn_ZeroPad3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((3, 4, 5, 6))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 3, 4, 5, 6)), + forward_input=FunctionInput(make_input((1, 2, 3, 4, 5))), + ), + ] + +def module_inputs_torch_nn_ConstantPad1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1, 2), + forward_input=FunctionInput(make_input((3, 4))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2), 3), + forward_input=FunctionInput(make_input((3, 4, 5))), + ), + ] + +def module_inputs_torch_nn_ConstantPad2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1, 3), + forward_input=FunctionInput(make_input((3, 4, 5))), + reference_fn=no_batch_dim_reference_fn + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 3, 4), 5), + forward_input=FunctionInput(make_input((1, 2, 3, 4))), + ), + ] + +def module_inputs_torch_nn_ConstantPad3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + return [ + ModuleInput( + constructor_input=FunctionInput(1, 3), + forward_input=FunctionInput(make_input((3, 4, 5, 6))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 3, 4, 5, 6), 7), + forward_input=FunctionInput(make_input((1, 2, 1, 2, 1))), + ), + ] + +def module_inputs_torch_nn_CircularPad1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def padding1d_circular_ref(inp, pad): + r""" input: + [[[0., 1., 2.], + [3., 4., 5.]]] + pad: (1, 2) + output: + [[[2., 0., 1., 2., 0., 1.], + [5., 3., 4., 5., 3., 4.]]] + """ + return torch.cat([inp[:, :, -pad[0]:], inp, inp[:, :, :pad[1]]], dim=2) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((3, 4))), + reference_fn=no_batch_dim_reference_fn + ), + ModuleInput( + constructor_input=FunctionInput((1, 2)), + forward_input=FunctionInput(make_input((1, 2, 3))), + reference_fn=lambda m, p, i: padding1d_circular_ref(i, m.padding), + ), + ModuleInput( + constructor_input=FunctionInput((3, 1)), + forward_input=FunctionInput(make_input((1, 2, 3))), + reference_fn=lambda m, p, i: padding1d_circular_ref(i, m.padding), + ), + ModuleInput( + constructor_input=FunctionInput((3, 3)), + forward_input=FunctionInput(make_input((1, 2, 3))), + reference_fn=lambda m, p, i: padding1d_circular_ref(i, m.padding), + ), + ] + +def module_inputs_torch_nn_CircularPad2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + def padding2d_circular_ref(inp, pad): + r"""input: + [[[[0., 1., 2], + [3., 4., 5.]]]] + pad: (1, 2, 2, 1) + output: + [[[[2., 0., 1., 2., 0., 1.], + [5., 3., 4., 5., 3., 4.], + [2., 0., 1., 2., 0., 1.], + [5., 3., 4., 5., 3., 4.], + [2., 0., 1., 2., 0., 1.]]]] + """ + inp = torch.cat([inp[:, :, -pad[2]:], inp, inp[:, :, :pad[3]]], dim=2) + return torch.cat([inp[:, :, :, -pad[0]:], inp, inp[:, :, :, :pad[1]]], dim=3) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((3, 4, 5))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 2, 1)), + forward_input=FunctionInput(make_input((1, 1, 2, 3))), + reference_fn=lambda m, p, i: padding2d_circular_ref(i, m.padding), + ), + ModuleInput( + constructor_input=FunctionInput((2, 3, 2, 2)), + forward_input=FunctionInput(make_input((1, 1, 2, 3))), + reference_fn=lambda m, p, i: padding2d_circular_ref(i, m.padding), + ), + ModuleInput( + constructor_input=FunctionInput((3, 3, 3, 1)), + forward_input=FunctionInput(make_input((1, 1, 3, 3))), + reference_fn=lambda m, p, i: padding2d_circular_ref(i, m.padding), + ), + ] + +def module_inputs_torch_nn_CircularPad3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + + def padding3d_circular_ref(inp, pad): + r"""input: + [[[[[ 0., 1., 2.], + [ 3., 4., 5.]], + [[ 6., 7., 8.], + [ 9., 10., 11.]]]]] + pad: (1, 2, 2, 1, 1, 2) + output: [[[[[ 8., 6., 7., 8., 6., 7.], + [11., 9., 10., 11., 9., 10.], + [ 8., 6., 7., 8., 6., 7.], + [11., 9., 10., 11., 9., 10.], + [ 8., 6., 7., 8., 6., 7.]], + + [[ 2., 0., 1., 2., 0., 1.], + [ 5., 3., 4., 5., 3., 4.], + [ 2., 0., 1., 2., 0., 1.], + [ 5., 3., 4., 5., 3., 4.], + [ 2., 0., 1., 2., 0., 1.]], + + [[ 8., 6., 7., 8., 6., 7.], + [11., 9., 10., 11., 9., 10.], + [ 8., 6., 7., 8., 6., 7.], + [11., 9., 10., 11., 9., 10.], + [ 8., 6., 7., 8., 6., 7.]], + + [[ 2., 0., 1., 2., 0., 1.], + [ 5., 3., 4., 5., 3., 4.], + [ 2., 0., 1., 2., 0., 1.], + [ 5., 3., 4., 5., 3., 4.], + [ 2., 0., 1., 2., 0., 1.]], + + [[ 8., 6., 7., 8., 6., 7.], + [11., 9., 10., 11., 9., 10.], + [ 8., 6., 7., 8., 6., 7.], + [11., 9., 10., 11., 9., 10.], + [ 8., 6., 7., 8., 6., 7.]]]]] + """ + inp = torch.cat([inp[:, :, -pad[4]:], inp, inp[:, :, :pad[5]]], dim=2) + inp = torch.cat([inp[:, :, :, -pad[2]:], inp, inp[:, :, :, :pad[3]]], dim=3) + return torch.cat([inp[:, :, :, :, -pad[0]:], inp, inp[:, :, :, :, :pad[1]]], dim=4) + + return [ + ModuleInput( + constructor_input=FunctionInput(1), + forward_input=FunctionInput(make_input((3, 4, 5, 6))), + reference_fn=no_batch_dim_reference_fn, + ), + ModuleInput( + constructor_input=FunctionInput((1, 2, 1, 2, 1, 2)), + forward_input=FunctionInput(make_input((1, 1, 2, 2, 3))), + reference_fn=lambda m, p, i: padding3d_circular_ref(i, m.padding) + ), + ModuleInput( + constructor_input=FunctionInput((3, 2, 2, 1, 1, 2)), + forward_input=FunctionInput(make_input((1, 1, 2, 2, 3))), + reference_fn=lambda m, p, i: padding3d_circular_ref(i, m.padding) + ), + ModuleInput( + constructor_input=FunctionInput((3, 3, 2, 1, 2, 2)), + forward_input=FunctionInput(make_input((1, 1, 2, 2, 3))), + reference_fn=lambda m, p, i: padding3d_circular_ref(i, m.padding) + ), + ] + + +# All these operators share similar issues on cuDNN and MIOpen +rnn_gru_lstm_module_info_decorators = ( + # RuntimeError: Batching rule not implemented for aten::_cudnn_rnn_backward. + # We could not generate a fallback + DecorateInfo( + unittest.expectedFailure, "TestModule", "test_grad", + active_if=(TEST_CUDNN and not TEST_WITH_ROCM), device_type='cuda' + ), + # NotImplementedError: the derivative for '_cudnn_rnn_backward' is not implemented. + # Double backwards is not supported for CuDNN RNNs due to limitations in the CuDNN API + DecorateInfo( + unittest.expectedFailure, "TestModule", "test_gradgrad", + active_if=(TEST_CUDNN and not TEST_WITH_ROCM), device_type='cuda' + ), + # CUDNN GRU doesn't accept non-contiguous hx + DecorateInfo( + unittest.expectedFailure, "TestModule", "test_non_contiguous_tensors", + active_if=(TEST_CUDNN and not TEST_WITH_ROCM), device_type='cuda' + ), + # MIOPEN GRU doesn't accept non-contiguous hx (this is dispatched to miopen only for float). + DecorateInfo( + unittest.expectedFailure, "TestModule", "test_non_contiguous_tensors", + active_if=(TEST_CUDNN and TEST_WITH_ROCM), dtypes=(torch.float,), device_type='cuda' + ) +) + +# Start of module error inputs functions. + +def module_error_inputs_torch_nn_NLLLoss(module_info, device, dtype, requires_grad, training, **kwargs): + """ + Error inputs for NLLLoss that test weight dtype must match input dtype. + Regression test for device parity between CPU and CUDA with empty tensors. + """ + input_t = torch.tensor([], device=device, dtype=dtype).reshape((0, 0)) + weight_dtype = torch.float32 if dtype == torch.float16 else torch.float16 + weight_t = torch.tensor([], device=device, dtype=weight_dtype) + target_t = torch.tensor([], device=device, dtype=torch.long) + + return [ + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(weight=weight_t), + forward_input=FunctionInput(input_t, target_t), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex=r"expected scalar type \w+ but found \w+" + ), + ] + +def module_error_inputs_torch_nn_RNN_GRU_Cell(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + samples = [ + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20), + forward_input=FunctionInput(make_input(3, 11), make_input(3, 20)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="input has inconsistent input_size: got 11 expected 10" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20), + forward_input=FunctionInput(make_input(3, 10), make_input(3, 21)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="hidden0 has inconsistent hidden_size: got 21, expected 20" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20), + forward_input=FunctionInput(make_input(3, 10), make_input(5, 20)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="Input batch size 3 doesn't match hidden0 batch size 5" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20), + forward_input=FunctionInput(make_input(3, 10), make_input(3, 1, 1, 20)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=ValueError, + error_regex="Expected hidden to be 1D or 2D, got 4D instead" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20, 'relu'), + forward_input=FunctionInput(make_input(3, 10), make_input(3, 21)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="hidden0 has inconsistent hidden_size: got 21, expected 20" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20, 'tanh'), + forward_input=FunctionInput(make_input(3, 10), make_input(3, 21)), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="hidden0 has inconsistent hidden_size: got 21, expected 20" + ), + ] + return samples + +def module_error_inputs_torch_nn_LSTMCell(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + samples = [ + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20), + forward_input=FunctionInput(make_input(3, 11), (make_input(3, 20), make_input(3, 20))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="input has inconsistent input_size: got 11 expected 10" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20), + forward_input=FunctionInput(make_input(3, 10), (make_input(3, 21), make_input(3, 21))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="hidden0 has inconsistent hidden_size: got 21, expected 20" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20), + forward_input=FunctionInput(make_input(3, 10), (make_input(5, 20), make_input(5, 20))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=RuntimeError, + error_regex="Input batch size 3 doesn't match hidden0 batch size 5" + ), + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(10, 20), + forward_input=FunctionInput(make_input(3, 10), (make_input(3, 1, 1, 20), make_input(3, 1, 1, 20))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=ValueError, + error_regex="Expected hx\\[0\\] to be 1D or 2D, got 4D instead" + ), + ] + return samples + + +def module_error_inputs_torch_nn_RNN_GRU(module_info, device, dtype, requires_grad, training, **kwargs): + # use float64 for dtype mismatch test if current dtype is float32, otherwise use float32 + # MPS doesn't support float64, so use float16 instead + # Extract device type from device string (e.g., 'mps:0' -> 'mps') + device_type = device.split(':')[0] if isinstance(device, str) else device.type + if dtype == torch.float32: + mismatched_dtype = torch.float16 if device_type == 'mps' else torch.float64 + else: + mismatched_dtype = torch.float32 + make_input = partial(make_tensor, device=device, dtype=mismatched_dtype, requires_grad=requires_grad) + + samples = [ + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(10, 0, 1)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=ValueError, + error_regex="hidden_size must be greater than zero" + ), + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(10, 10, 0)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=ValueError, + error_regex="num_layers must be greater than zero" + ), + # Test dtype mismatch error message + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(3, 5, dtype=dtype, device=device), + forward_input=FunctionInput(make_input((2, 4, 3))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=ValueError, + error_regex=(r"RNN input dtype .* does not match weight dtype .* " + r"Convert input: input\.to\(.*\), or convert model: model\.to\(.*\)") + ), + # Test bias parameter type validation + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(3, 5, bias=0)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=TypeError, + error_regex="bias should be of type bool, got: int" + ), + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(3, 5, bias=1)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=TypeError, + error_regex="bias should be of type bool, got: int" + ), + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(3, 5, bias="True")), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=TypeError, + error_regex="bias should be of type bool, got: str" + ), + # Test batch_first parameter type validation + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(3, 5, batch_first=0)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=TypeError, + error_regex="batch_first should be of type bool, got: int" + ), + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(3, 5, batch_first=1)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=TypeError, + error_regex="batch_first should be of type bool, got: int" + ), + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(3, 5, batch_first="False")), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=TypeError, + error_regex="batch_first should be of type bool, got: str" + ), + # Test input_size parameter type validation + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(3.0, 5)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=TypeError, + error_regex="input_size should be of type int, got: float" + ), + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput("10", 5)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=TypeError, + error_regex="input_size should be of type int, got: str" + ), + # Test input_size parameter value validation + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(0, 5)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=ValueError, + error_regex="input_size must be greater than zero" + ), + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(-1, 5)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=ValueError, + error_regex="input_size must be greater than zero" + ), + # Test hidden_size parameter type validation + ErrorModuleInput( + ModuleInput(constructor_input=FunctionInput(3, 5.0)), + error_on=ModuleErrorEnum.CONSTRUCTION_ERROR, + error_type=TypeError, + error_regex="hidden_size should be of type int, got: float" + ), + ] + return samples + +def module_error_inputs_torch_nn_Pad1d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + is_constant = kwargs.get('is_constant', False) + + return [ + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(1, 3) if is_constant else FunctionInput(3), + forward_input=FunctionInput(make_input((2, 3, 4, 5))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=ValueError, + error_regex=r"expected 2D or 3D input \(got 4D input\)", + + ), + ] + +def module_error_inputs_torch_nn_Pad2d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + is_constant = kwargs.get('is_constant', False) + + return [ + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(1, 3) if is_constant else FunctionInput(3), + forward_input=FunctionInput(make_input((2, 3))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=ValueError, + error_regex=r"expected 3D or 4D input \(got 2D input\)", + + ), + ] + +def module_error_inputs_torch_nn_Pad3d(module_info, device, dtype, requires_grad, training, **kwargs): + make_input = partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + + is_constant = kwargs.get('is_constant', False) + + return [ + ErrorModuleInput( + ModuleInput( + constructor_input=FunctionInput(1, 3) if is_constant else FunctionInput(3), + forward_input=FunctionInput(make_input((2, 3))), + ), + error_on=ModuleErrorEnum.FORWARD_ERROR, + error_type=ValueError, + error_regex=r"expected 4D or 5D input \(got 2D input\)", + + ), + ] + + +_macos15_or_newer = torch.backends.mps.is_available() and torch.backends.mps.is_macos_or_newer(15, 0) + + +# Database of ModuleInfo entries in alphabetical order. +module_db: list[ModuleInfo] = [ + ModuleInfo(torch.nn.AdaptiveAvgPool1d, + module_inputs_func=module_inputs_torch_nn_AdaptiveAvgPool1d, + skips=( + # Fails on MPS backend if input/output sizes are not divisible + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.AdaptiveAvgPool2d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_inputs_func=module_inputs_torch_nn_AdaptiveAvgPool2d, + skips=( + # Fails on MPS backend if input/output sizes are not divisible + DecorateInfo(skipMPS), + # Fails on backward check if output size is 1x1 + DecorateInfo( + unittest.expectedFailure, + 'TestModule', + 'test_memory_format', + active_if=operator.itemgetter('training'), + ),) + ), + ModuleInfo(torch.nn.AdaptiveAvgPool3d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_inputs_func=module_inputs_torch_nn_AdaptiveAvgPool3d, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # not supported on MPS backend + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.AdaptiveMaxPool1d, + module_inputs_func=module_inputs_torch_nn_AdaptiveMaxPool1d, + ), + ModuleInfo(torch.nn.AdaptiveMaxPool2d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_inputs_func=module_inputs_torch_nn_AdaptiveMaxPool2d, + ), + ModuleInfo(torch.nn.AdaptiveMaxPool3d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_inputs_func=module_inputs_torch_nn_AdaptiveMaxPool3d, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # not supported on MPS backend + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.AvgPool1d, + module_inputs_func=module_inputs_torch_nn_AvgPool1d, + ), + ModuleInfo(torch.nn.AvgPool2d, + module_inputs_func=module_inputs_torch_nn_AvgPool2d, + skips=( + # The difference between channels last backward and + # channels first backward of AvgPool2d on CUDA is too large + # See https://github.com/pytorch/pytorch/issues/107201 + DecorateInfo( + unittest.expectedFailure, + 'TestModule', + 'test_memory_format', + active_if=operator.itemgetter('training'), + device_type='cuda',), + ),), + ModuleInfo(torch.nn.AvgPool3d, + module_inputs_func=module_inputs_torch_nn_AvgPool3d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + # No channels_last support for AvgPool1d as it does not take 4D inputs + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # backward not supported on MPS backend + DecorateInfo(skipMPS, 'TestModule', 'test_non_contiguous_tensors'),) + ), + ModuleInfo(torch.nn.BatchNorm1d, + train_and_eval_differ=True, + module_inputs_func=module_inputs_torch_nn_BatchNorm1d, + module_error_inputs_func=module_error_inputs_torch_nn_BatchNorm1d_2d_3d, + skips=( + # tracking here rather than in the list in test_aotdispatch.py as eval mode passes + # RuntimeError: tried to get Double out of SymInt + DecorateInfo( + unittest.expectedFailure, 'TestEagerFusionModuleInfo', + 'test_aot_autograd_symbolic_module_exhaustive', + active_if=operator.itemgetter('training') + ), + # torch._subclasses.fake_tensor.DataDependentOutputException: aten._local_scalar_dense.default + DecorateInfo( + unittest.expectedFailure, 'TestEagerFusionModuleInfo', + 'test_aot_autograd_module_exhaustive', + active_if=operator.itemgetter('training') + )) + ), + ModuleInfo(torch.nn.BatchNorm2d, + train_and_eval_differ=True, + module_inputs_func=module_inputs_torch_nn_BatchNorm2d, + module_error_inputs_func=module_error_inputs_torch_nn_BatchNorm1d_2d_3d, + skips=( + # See https://github.com/pytorch/pytorch/issues/134580 + DecorateInfo(expectedFailureMPS, 'TestModule', 'test_memory_format', active_if=operator.itemgetter('training')), + # tracking here rather than in the list in test_aotdispatch.py as eval mode passes + # RuntimeError: tried to get Double out of SymInt + DecorateInfo( + unittest.expectedFailure, 'TestEagerFusionModuleInfo', + 'test_aot_autograd_symbolic_module_exhaustive', + active_if=operator.itemgetter('training') + ), + # torch._subclasses.fake_tensor.DataDependentOutputException: aten._local_scalar_dense.default + DecorateInfo( + unittest.expectedFailure, 'TestEagerFusionModuleInfo', + 'test_aot_autograd_module_exhaustive', + active_if=operator.itemgetter('training') + ),) + ), + ModuleInfo(torch.nn.BatchNorm3d, + train_and_eval_differ=True, + module_inputs_func=module_inputs_torch_nn_BatchNorm3d, + module_error_inputs_func=module_error_inputs_torch_nn_BatchNorm1d_2d_3d, + skips=( + # not supported on MPS backend + DecorateInfo(skipMPS), + # tracking here rather than in the list in test_aotdispatch.py as eval mode passes + # RuntimeError: tried to get Double out of SymInt + DecorateInfo( + unittest.expectedFailure, 'TestEagerFusionModuleInfo', + 'test_aot_autograd_symbolic_module_exhaustive', + active_if=operator.itemgetter('training') + ), + # torch._subclasses.fake_tensor.DataDependentOutputException: aten._local_scalar_dense.default + DecorateInfo( + unittest.expectedFailure, 'TestEagerFusionModuleInfo', + 'test_aot_autograd_module_exhaustive', + active_if=operator.itemgetter('training') + ),) + ), + ModuleInfo(torch.nn.CELU, + module_inputs_func=module_inputs_torch_nn_CELU, + # not MPS specific, will be xfailed for all devices in next PR + skips=( + DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_check_inplace', + device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.Conv1d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=1, lazy=False), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.Conv2d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=2, lazy=False), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + skips=( + # This was wrongly being skipped before and needs investigation. + # See https://github.com/pytorch/pytorch/issues/80247 + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format", + device_type='cuda', dtypes=[torch.float64]), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.Conv3d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=3, lazy=False), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + skips=( + # Conv3d is not supported on MPS backend + DecorateInfo(skipMPS, device_type="mps"), + # This was wrongly being skipped before and needs investigation. + # See https://github.com/pytorch/pytorch/issues/80247 + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format"), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.ConvTranspose1d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=1, lazy=False, transposed=True), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + dtypes=floating_and_complex_types_and(torch.chalf), + skips=( + # Not implemented for chalf on CPU + DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_cpu_gpu_parity', + dtypes=(torch.chalf,), device_type='cuda'), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + DecorateInfo(precisionOverride({torch.chalf: 5e-03}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.ConvTranspose2d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=2, lazy=False, transposed=True), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + dtypes=floating_and_complex_types_and(torch.chalf), + skips=( + # Fails on backward check because ViewAsRealBackward apply contiguous for grad + DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_memory_format', + dtypes=(torch.complex32, torch.complex64, torch.complex128)), + # This was wrongly being skipped before and needs investigation. + # See https://github.com/pytorch/pytorch/issues/80247 + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format", device_type='cuda', + dtypes=[torch.float64, torch.complex128]), + # Not implemented for chalf on CPU + DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_cpu_gpu_parity', + dtypes=(torch.chalf,), device_type='cuda'), + DecorateInfo(skipIfXpu, 'TestModule', 'test_cpu_gpu_parity', + dtypes=(torch.chalf,), device_type='xpu'), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + DecorateInfo(precisionOverride({torch.chalf: 5e-03}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.ConvTranspose3d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=3, lazy=False, transposed=True), + dtypes=floating_and_complex_types_and(torch.chalf), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + skips=( + # ConvTranspose3d is not supported on MPS backend + DecorateInfo(skipMPS), + # This was wrongly being skipped before and needs investigation. + # See https://github.com/pytorch/pytorch/issues/80247 + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format"), + # Not implemented for chalf on CPU + DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_cpu_gpu_parity', + dtypes=(torch.chalf,), device_type='cuda'), + DecorateInfo(skipIfXpu, 'TestModule', 'test_cpu_gpu_parity', + dtypes=(torch.chalf,), device_type='xpu'), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + DecorateInfo(precisionOverride({torch.complex64: 1e-04}), 'TestModule', 'test_cpu_gpu_parity'), + DecorateInfo(precisionOverride({torch.chalf: 5e-03}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.CosineEmbeddingLoss, + module_inputs_func=module_inputs_torch_nn_CosineEmbeddingLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.ELU, + module_inputs_func=module_inputs_torch_nn_ELU, + # not MPS specific, will be xfailed for all devices in next PR + skips=( + DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_check_inplace', + device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.FractionalMaxPool2d, + module_inputs_func=module_inputs_torch_nn_FractionalMaxPool2d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + # not supported on MPS backend + DecorateInfo(skipMPS), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.FractionalMaxPool3d, + module_inputs_func=module_inputs_torch_nn_FractionalMaxPool3d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + # not supported on MPS backend + DecorateInfo(skipMPS), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.L1Loss, + module_inputs_func=module_inputs_torch_nn_L1Loss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.SmoothL1Loss, + module_inputs_func=module_inputs_torch_nn_SmoothL1Loss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # See #119108: input types 'tensor' and 'tensor<15x10xf16>' are not broadcast compatible + # NS: Still fails on MacOS15.1 + DecorateInfo(skipIfMPS, 'TestModule', 'test_non_contiguous_tensors', + dtypes=[torch.float16], device_type='mps'),), + ), + ModuleInfo(torch.nn.LazyConv1d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=1, lazy=True), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + skips=( + # Lazy modules don't currently play well with ModuleInfo tests on the meta device. + # See https://github.com/pytorch/pytorch/issues/70505 for more info. + DecorateInfo(skipMeta), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.LazyConv2d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=2, lazy=True), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + skips=( + # Lazy modules don't currently play well with ModuleInfo tests on the meta device. + # See https://github.com/pytorch/pytorch/issues/70505 for more info. + DecorateInfo(skipMeta), + # This was wrongly being skipped before and needs investigation. + # See https://github.com/pytorch/pytorch/issues/80247 + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format", + device_type='cuda', dtypes=[torch.float64]), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.LazyConv3d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=3, lazy=True), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + skips=( + # Lazy modules don't currently play well with ModuleInfo tests on the meta device. + # See https://github.com/pytorch/pytorch/issues/70505 for more info. + DecorateInfo(skipMeta), + # LazyConv3d is not supported on MPS backend + DecorateInfo(skipMPS), + # This was wrongly being skipped before and needs investigation. + # See https://github.com/pytorch/pytorch/issues/80247 + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format"), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.LazyConvTranspose1d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=1, lazy=True, transposed=True), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + skips=( + # Lazy modules don't currently play well with ModuleInfo tests on the meta device. + # See https://github.com/pytorch/pytorch/issues/70505 for more info. + DecorateInfo(skipMeta), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.LazyConvTranspose2d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=2, lazy=True, transposed=True), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + skips=( + # Lazy modules don't currently play well with ModuleInfo tests on the meta device. + # See https://github.com/pytorch/pytorch/issues/70505 for more info. + DecorateInfo(skipMeta), + # This was wrongly being skipped before and needs investigation. + # See https://github.com/pytorch/pytorch/issues/80247 + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format", device_type='cuda', + dtypes=[torch.float64]), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.LazyConvTranspose3d, + module_inputs_func=partial(module_inputs_torch_nn_ConvNd, N=3, lazy=True, transposed=True), + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + module_memformat_affects_out=True, + skips=( + # Lazy modules don't currently play well with ModuleInfo tests on the meta device. + # See https://github.com/pytorch/pytorch/issues/70505 for more info. + DecorateInfo(skipMeta), + # LazyConvTranspose3d is not supported on MPS backend + DecorateInfo(skipMPS), + # This was wrongly being skipped before and needs investigation. + # See https://github.com/pytorch/pytorch/issues/80247 + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format"), + ), + decorators=( + DecorateInfo(precisionOverride({torch.float32: 1e-04}), 'TestModule', 'test_memory_format'), + )), + ModuleInfo(torch.nn.Linear, + module_inputs_func=module_inputs_torch_nn_Linear, + skips=( + # No channels_last support for Linear currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.Bilinear, + module_inputs_func=module_inputs_torch_nn_Bilinear, + decorators=[ + DecorateInfo( + toleranceOverride({ + torch.float32: tol(atol=1e-4, rtol=1e-4), + torch.float64: tol(atol=1e-4, rtol=1e-4)}), + 'TestModule', 'test_forward', device_type='cpu'), + ], + skips=( + # No channels_last support for Bilinear currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # See #119108: tolerance issue + DecorateInfo(unittest.expectedFailure, "TestModule", "test_forward", + device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.LPPool1d, + module_inputs_func=module_inputs_torch_nn_LPPool1d, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_grad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad'),) + ), + ModuleInfo(torch.nn.LPPool2d, + module_inputs_func=module_inputs_torch_nn_LPPool2d, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_grad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad'), + # Fails on backward check on MPS + # See https://github.com/pytorch/pytorch/issues/107214 + DecorateInfo( + unittest.expectedFailure, + 'TestModule', + 'test_memory_format', + active_if=operator.itemgetter('training') and not _macos15_or_newer, + device_type='mps', + ),) + ), + ModuleInfo(torch.nn.LPPool3d, + module_inputs_func=module_inputs_torch_nn_LPPool3d, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_grad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + DecorateInfo(skipIfMPS, device_type='mps'),) + ), + ModuleInfo(torch.nn.MaxPool1d, + module_inputs_func=module_inputs_torch_nn_MaxPool1d, + ), + ModuleInfo(torch.nn.MaxPool2d, + module_inputs_func=module_inputs_torch_nn_MaxPool2d, + module_error_inputs_func=module_error_inputs_torch_nn_MaxPool2d, + ), + ModuleInfo(torch.nn.MaxPool3d, + module_inputs_func=module_inputs_torch_nn_MaxPool3d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + ), + ModuleInfo(torch.nn.KLDivLoss, + module_inputs_func=module_inputs_torch_nn_KLDivLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # https://github.com/pytorch/pytorch/issues/115588 + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_cpu_gpu_parity'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_grad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad'),) + ), + ModuleInfo(torch.nn.MSELoss, + module_inputs_func=module_inputs_torch_nn_MSELoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # See #119108: tolerance issue + DecorateInfo(unittest.expectedFailure, "TestModule", "test_forward", + device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.MarginRankingLoss, + module_inputs_func=module_inputs_torch_nn_MarginRankingLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.MultiLabelMarginLoss, + module_inputs_func=module_inputs_torch_nn_MultiLabelMarginLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # 'aten::multilabel_margin_loss_forward' is not currently implemented for the MPS device. + DecorateInfo(skipIfMPS, 'TestModule', device_type='mps'), + # derivative for aten::multilabel_margin_loss_backward is not implemented + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad'),) + ), + ModuleInfo(torch.nn.MultiMarginLoss, + module_inputs_func=module_inputs_torch_nn_MultiMarginLoss, + module_error_inputs_func=module_error_inputs_torch_nn_MultiMarginLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # 'aten::multi_margin_loss' is not currently implemented for the MPS device. + DecorateInfo(skipIfMPS, 'TestModule', device_type='mps'), + # RuntimeError: derivative for aten::multi_margin_loss_backward is not implemented + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad'),) + ), + ModuleInfo(torch.nn.SoftMarginLoss, + module_inputs_func=module_inputs_torch_nn_SoftMarginLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # See #119108: tolerance issue + DecorateInfo(unittest.expectedFailure, "TestModule", "test_forward", + device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.MultiLabelSoftMarginLoss, + module_inputs_func=module_inputs_torch_nn_MultiLabelSoftMarginLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.NLLLoss, + module_inputs_func=module_inputs_torch_nn_NLLLoss, + module_error_inputs_func=module_error_inputs_torch_nn_NLLLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # See #119108: tolerance issue + DecorateInfo(unittest.expectedFailure, "TestModule", "test_forward", + device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.GaussianNLLLoss, + module_inputs_func=module_inputs_torch_nn_GaussianNLLLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),)), + ModuleInfo(torch.nn.PoissonNLLLoss, + module_inputs_func=module_inputs_torch_nn_PoissonNLLLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),)), + ModuleInfo(torch.nn.HingeEmbeddingLoss, + module_inputs_func=module_inputs_torch_nn_HingeEmbeddingLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.HuberLoss, + module_inputs_func=module_inputs_torch_nn_HuberLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # See #119108: seemingly incorrect output dtype + DecorateInfo(unittest.expectedFailure, "TestModule", "test_forward", + device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.BCELoss, + module_inputs_func=module_inputs_torch_nn_BCELoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # error: input types 'tensor' and 'tensor<15x10xf16>' are not broadcast compatible + DecorateInfo(skipIfMPS, 'TestModule', dtypes=[torch.float16], device_type='mps'),) + ), + ModuleInfo(torch.nn.BCEWithLogitsLoss, + module_inputs_func=module_inputs_torch_nn_BCEWithLogitsLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # see #119108: tolerance issue + DecorateInfo(skipIfMPS, 'TestModule', dtypes=[torch.float16], device_type='mps'),) + ), + ModuleInfo(torch.nn.CrossEntropyLoss, + module_inputs_func=module_inputs_torch_nn_CrossEntropyLoss, + module_error_inputs_func=module_error_inputs_torch_nn_CrossEntropyLoss, + dtypes=get_all_fp_dtypes(include_half=True, include_bfloat16=False), + decorators=( + # No channels_last support for loss functions. + DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_memory_format'), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=3e-2, rtol=1e-3)}), "TestModule", + "test_forward", dtypes=[torch.float16], device_type='cpu'), + DecorateInfo(unittest.expectedFailure, "TestModule", "test_cpu_gpu_parity", dtypes=[torch.float16], + device_type='cuda'), + DecorateInfo(unittest.expectedFailure, "TestModule", "test_cpu_gpu_parity", dtypes=[torch.float16], + device_type='xpu'),), + ), + ModuleInfo(torch.nn.CTCLoss, + module_inputs_func=module_inputs_torch_nn_CTCLoss, + skips=( + # No channels_last support for loss functions. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # The operator aten::_ctc_loss is not currently implemented for the MPS device. + DecorateInfo(skipIfMPS, 'TestModule', device_type='mps',), + # derivative for aten::_ctc_loss_backward is not implemented + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_grad'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad'), + # https://github.com/pytorch/pytorch/issues/115585 + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_non_contiguous_tensors'),) + ), + ModuleInfo(torch.nn.GELU, + module_inputs_func=module_inputs_torch_nn_GELU, + skips=( + # See #119108: tolerance issue + DecorateInfo(unittest.expectedFailure, "TestModule", "test_forward", + device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.GLU, + module_inputs_func=module_inputs_torch_nn_GLU, + ), + ModuleInfo(torch.nn.GroupNorm, + module_inputs_func=module_inputs_torch_nn_GroupNorm, + module_error_inputs_func=module_error_inputs_torch_nn_GroupNorm, + dtypes=get_all_fp_dtypes(include_bfloat16=True, include_half=True), + skips=( + # Tracking at https://github.com/pytorch/pytorch/issues/98089 + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_cpu_gpu_parity'), + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4)}), + 'TestModule', 'test_memory_format', device_type='cpu'), + # No channels_last support for GroupNorm currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', device_type='mps'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', device_type='xpu'), + DecorateInfo(unittest.skip("Skipped!"), "TestModule", "test_grad", + active_if=TEST_WITH_ROCM, device_type='cuda'),) + ), + ModuleInfo(torch.nn.Hardshrink, + module_inputs_func=module_inputs_torch_nn_Hardshrink, + ), + ModuleInfo(torch.nn.Hardswish, + module_inputs_func=module_inputs_torch_nn_Hardswish, + supports_gradgrad=False), + ModuleInfo(torch.nn.Hardtanh, + module_inputs_func=module_inputs_torch_nn_Hardtanh, + ), + ModuleInfo(torch.nn.InstanceNorm1d, + module_inputs_func=partial(module_inputs_torch_nn_InstanceNormNd, N=1), + train_and_eval_differ=True, + skips=( + # No channels_last support for InstanceNorm1d currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.InstanceNorm2d, + module_inputs_func=partial(module_inputs_torch_nn_InstanceNormNd, N=2), + train_and_eval_differ=True, + skips=( + # No channels_last support for InstanceNorm2d currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.InstanceNorm3d, + module_inputs_func=partial(module_inputs_torch_nn_InstanceNormNd, N=3), + train_and_eval_differ=True, + skips=( + # not supported on MPS backend + DecorateInfo(expectedFailureMPS, 'TestModuleMPS', 'test_memory_format'), + DecorateInfo(expectedFailureMPS, 'TestModuleMPS', 'test_non_contiguous_tensors'), + DecorateInfo(expectedFailureMPS, 'TestModuleMPS', 'test_forward'), + DecorateInfo(expectedFailureMPS, 'TestModuleMPS', 'test_non_contiguous'), + DecorateInfo(expectedFailureMPS, 'TestModuleMPS', 'test_save_load'), + # No channels_last support for InstanceNorm3d currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.LocalResponseNorm, + module_inputs_func=module_inputs_torch_nn_LocalResponseNorm, + ), + ModuleInfo(torch.nn.LayerNorm, + module_inputs_func=module_inputs_torch_nn_LayerNorm, + skips=( + # No channels_last support for LayerNorm currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.RMSNorm, + module_inputs_func=module_inputs_torch_nn_RMSNorm, + ), + # TransformerEncoder takes the same inputs as TransformerEncoderLayer + ModuleInfo(torch.nn.TransformerEncoder, + train_and_eval_differ=True, + module_inputs_func=module_inputs_torch_nn_TransformerEncoder, + decorators=[ + # Not implemented for SDPA backward derivative + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad', + device_type='cpu'), + ], + skips=( + # No channels_last support for TransformerEncoderLayer currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # Doesn't support device / dtype kwargs directly because it is just a + # container of TransformerEncoderLayers. + DecorateInfo(unittest.expectedFailure, 'TestModule', 'test_factory_kwargs'),) + ), + ModuleInfo(torch.nn.TransformerEncoderLayer, + train_and_eval_differ=True, + module_inputs_func=module_inputs_torch_nn_TransformerEncoderLayer, + decorators=[ + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4)}), + 'TestModule', 'test_non_contiguous_tensors', + device_type='cpu', active_if=IS_WINDOWS), + DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-4, rtol=2e-3)}), + 'TestModule', 'test_forward', + device_type='mps'), + # Not implemented for SDPA backward derivative + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad', + device_type='cpu'), + ], + skips=( + # No channels_last support for TransformerEncoderLayer currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.TransformerDecoderLayer, + module_inputs_func=module_inputs_torch_nn_TransformerDecoderLayer, + decorators=[ + # Not implemented for SDPA backward derivative + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad', + device_type='cpu'), + ], + skips=( + # No channels_last support for TransformerDecoderLayer currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.Transformer, + module_inputs_func=module_inputs_torch_nn_Transformer, + # Inputs are too large to run with slow gradcheck + # https://github.com/pytorch/pytorch/issues/117140 + gradcheck_fast_mode=True, + decorators=[ + # Not implemented for SDPA backward derivative + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_gradgrad', + device_type='cpu'), + ], + skips=( + # No channels_last support for Transformer currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.MultiheadAttention, + train_and_eval_differ=True, + module_inputs_func=module_inputs_torch_nn_MultiheadAttention, + skips=( + # No channels_last support for MultiheadAttention currently. + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.Embedding, + module_inputs_func=module_inputs_torch_nn_Embedding, + module_error_inputs_func=module_error_inputs_torch_nn_Embedding, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + decorators=[ + DecorateInfo(toleranceOverride({torch.float32: tol(atol=1e-4, rtol=1e-4)}), + 'TestModule', 'test_non_contiguous_tensors', + device_type='mps')], + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.ReLU, + module_inputs_func=module_inputs_torch_nn_ReLU, + skips=None if _macos15_or_newer else ( + # Fails on backward check on MPS + # See https://github.com/pytorch/pytorch/issues/107214 + DecorateInfo( + unittest.expectedFailure, + 'TestModule', + 'test_memory_format', + active_if=operator.itemgetter('training'), + device_type='mps', + ),) + ), + ModuleInfo(torch.nn.LeakyReLU, + module_inputs_func=module_inputs_torch_nn_LeakyReLU, + ), + ModuleInfo(torch.nn.ReLU6, + module_inputs_func=module_inputs_torch_nn_ReLU6, + skips=( + # test fails on MPS backend and is being investigated. + # See https://github.com/pytorch/pytorch/issues/100914 + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.PReLU, + module_inputs_func=module_inputs_torch_nn_PReLU, + skips=( + # test fails on MPS backend and is being investigated. + # See https://github.com/pytorch/pytorch/issues/100914 + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.RNNCell, + module_inputs_func=partial(module_inputs_torch_nn_RNN_GRU_Cell, is_rnn=True), + module_error_inputs_func=module_error_inputs_torch_nn_RNN_GRU_Cell, + ), + ModuleInfo(torch.nn.GRUCell, + module_inputs_func=module_inputs_torch_nn_RNN_GRU_Cell, + module_error_inputs_func=module_error_inputs_torch_nn_RNN_GRU_Cell, + ), + ModuleInfo(torch.nn.LSTMCell, + module_inputs_func=module_inputs_torch_nn_LSTMCell, + module_error_inputs_func=module_error_inputs_torch_nn_LSTMCell, + ), + ModuleInfo(torch.nn.Sigmoid, + module_inputs_func=module_inputs_torch_nn_Sigmoid, + skips=None if _macos15_or_newer else ( + # Fails on backward check on MPS + # See https://github.com/pytorch/pytorch/issues/107214 + DecorateInfo( + unittest.expectedFailure, + 'TestModule', + 'test_memory_format', + active_if=operator.itemgetter('training'), + device_type='mps', + ),) + ), + ModuleInfo(torch.nn.LogSigmoid, + module_inputs_func=module_inputs_torch_nn_LogSigmoid, + skips=( + # See #119108: tolerance issue + DecorateInfo(unittest.expectedFailure, "TestModule", "test_forward", device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.SiLU, + module_inputs_func=module_inputs_torch_nn_SiLU, + ), + ModuleInfo(torch.nn.Softmax, + module_inputs_func=module_inputs_torch_nn_Softmax, + ), + ModuleInfo(torch.nn.Softmax2d, + module_inputs_func=module_inputs_torch_nn_Softmax2d, + skips=( + # no channels last support for Softmax2d currently + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # See #119108: tolerance issue + DecorateInfo(unittest.expectedFailure, "TestModule", "test_forward", device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.LogSoftmax, + module_inputs_func=module_inputs_torch_nn_LogSoftmax, + skips=( + # no channels last support for LogSoftmax currently + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'), + # See #119108: inf nan error + DecorateInfo(unittest.expectedFailure, "TestModule", "test_forward", device_type='mps', dtypes=[torch.float16]),) + ), + ModuleInfo(torch.nn.Softmin, + module_inputs_func=module_inputs_torch_nn_Softmin, + skips=( + # no channels last support for Softmin currently + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format'),) + ), + ModuleInfo(torch.nn.Softplus, + module_inputs_func=module_inputs_torch_nn_Softplus, + skips=( + # test fails on MPS backend and is being investigated. + # See https://github.com/pytorch/pytorch/issues/100914 + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.Softshrink, + module_inputs_func=module_inputs_torch_nn_Softshrink, + skips=( + # not supported on MPS backend + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.Softsign, + module_inputs_func=module_inputs_torch_nn_Softsign, + ), + ModuleInfo(torch.nn.Tanh, + module_inputs_func=module_inputs_torch_nn_Tanh, + skips=None if _macos15_or_newer else ( + # Fails on backward check on MPS + # See https://github.com/pytorch/pytorch/issues/107214 + DecorateInfo( + unittest.expectedFailure, + 'TestModule', + 'test_memory_format', + active_if=operator.itemgetter('training'), + device_type='mps', + ),) + ), + ModuleInfo(torch.nn.Tanhshrink, + module_inputs_func=module_inputs_torch_nn_Tanhshrink, + skips=None if _macos15_or_newer else ( + # Fails on backward check on MPS + # See https://github.com/pytorch/pytorch/issues/107214 + DecorateInfo( + unittest.expectedFailure, + 'TestModule', + 'test_memory_format', + active_if=operator.itemgetter('training'), + device_type='mps', + ),) + ), + ModuleInfo(torch.nn.Threshold, + module_inputs_func=module_inputs_torch_nn_Threshold, + skips=( + # test fails on MPS backend and is being investigated. + # See https://github.com/pytorch/pytorch/issues/100914 + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.Mish, + module_inputs_func=module_inputs_torch_nn_Mish, + skips=( + # not supported on MPS backend + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.RNN, + train_and_eval_differ=True, + module_inputs_func=partial(module_inputs_torch_nn_RNN_GRU, is_rnn=True), + module_error_inputs_func=module_error_inputs_torch_nn_RNN_GRU, + decorators=rnn_gru_lstm_module_info_decorators + ), + ModuleInfo(torch.nn.GRU, + train_and_eval_differ=True, + module_inputs_func=partial(module_inputs_torch_nn_RNN_GRU, is_rnn=False), + module_error_inputs_func=module_error_inputs_torch_nn_RNN_GRU, + decorators=rnn_gru_lstm_module_info_decorators), + ModuleInfo(torch.nn.LSTM, + train_and_eval_differ=True, + module_inputs_func=module_inputs_torch_nn_LSTM, + module_error_inputs_func=module_error_inputs_torch_nn_RNN_GRU, + skips=( + # LSTM with projections is not currently supported with MPS + DecorateInfo(skipMPS),), + decorators=rnn_gru_lstm_module_info_decorators), + ModuleInfo(torch.nn.ReflectionPad1d, + module_inputs_func=module_inputs_torch_nn_ReflectionPad1d, + ), + ModuleInfo(torch.nn.ReflectionPad2d, + module_inputs_func=module_inputs_torch_nn_ReflectionPad2d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='mps'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='xpu'),) + ), + ModuleInfo(torch.nn.ReflectionPad3d, + module_inputs_func=module_inputs_torch_nn_ReflectionPad3d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='mps'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='xpu'),) + ), + ModuleInfo(torch.nn.ReplicationPad1d, + module_inputs_func=module_inputs_torch_nn_ReplicationPad1d, + ), + ModuleInfo(torch.nn.ReplicationPad2d, + module_inputs_func=module_inputs_torch_nn_ReplicationPad2d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='mps'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='xpu'),) + ), + ModuleInfo(torch.nn.ReplicationPad3d, + module_inputs_func=module_inputs_torch_nn_ReplicationPad3d, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='cuda'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='mps'), + DecorateInfo(unittest.skip("Skipped!"), 'TestModule', 'test_memory_format', + device_type='xpu'),) + ), + ModuleInfo(torch.nn.SELU, + module_inputs_func=module_inputs_torch_nn_SELU, + skips=( + # test fails on MPS backend and is being investigated. + # See https://github.com/pytorch/pytorch/issues/100914 + DecorateInfo(skipMPS),) + ), + ModuleInfo(torch.nn.ZeroPad1d, + module_inputs_func=module_inputs_torch_nn_ZeroPad1d, + ), + ModuleInfo(torch.nn.ZeroPad2d, + module_inputs_func=module_inputs_torch_nn_ZeroPad2d, + skips=( + # Fails with channels last test on MPS backend + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format", device_type='mps'),) + ), + ModuleInfo(torch.nn.ZeroPad3d, + module_inputs_func=module_inputs_torch_nn_ZeroPad3d, + skips=( + # Fails with channels last test on MPS backend + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format", device_type='mps'),) + ), + ModuleInfo(torch.nn.CircularPad1d, + module_inputs_func=module_inputs_torch_nn_CircularPad1d, + module_error_inputs_func=module_error_inputs_torch_nn_Pad1d, + ), + ModuleInfo(torch.nn.CircularPad2d, + module_inputs_func=module_inputs_torch_nn_CircularPad2d, + module_error_inputs_func=module_error_inputs_torch_nn_Pad2d, + ), + ModuleInfo(torch.nn.CircularPad3d, + module_inputs_func=module_inputs_torch_nn_CircularPad3d, + module_error_inputs_func=module_error_inputs_torch_nn_Pad3d, + skips=( + # Fails with channels last test on MPS backend + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format"),) + ), + ModuleInfo(torch.nn.ConstantPad1d, + module_inputs_func=module_inputs_torch_nn_ConstantPad1d, + ), + ModuleInfo(torch.nn.ConstantPad2d, + module_inputs_func=module_inputs_torch_nn_ConstantPad2d, + skips=( + # Fails with channels last test on MPS backend + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format", device_type='mps'),) + ), + ModuleInfo(torch.nn.ConstantPad3d, + module_inputs_func=module_inputs_torch_nn_ConstantPad3d, + skips=( + # Fails with channels last test on MPS backend + DecorateInfo(unittest.expectedFailure, "TestModule", "test_memory_format", device_type='mps'),) + ) +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_mps.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_mps.py new file mode 100644 index 0000000000000000000000000000000000000000..f8fc0f4437d8b8a81438a82400ddd15359d94fed --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_mps.py @@ -0,0 +1,1058 @@ +import unittest +from collections.abc import Sequence + +import torch + +from .common_utils import MACOS_VERSION +from .opinfo.core import DecorateInfo, OpInfo + + +if torch.backends.mps.is_available(): + + def mps_ops_modifier( + ops: Sequence[OpInfo], + device_type: str = "mps", + xfail_exclusion: list[str] | None = None, + sparse: bool = False, + ) -> Sequence[OpInfo]: + if xfail_exclusion is None: + xfail_exclusion = [] + + # Supported complex OPS + SUPPORTED_COMPLEX_OPS = { + "__radd__", + "__rmul__", + "__rsub__", + "__getitem__", + "_unsafe_masked_index", + "_unsafe_masked_index_put_accumulate", + "abs", + "add", + "addbmm", + "alias_copy", + "argwhere", + "atleast_1d", + "atleast_2d", + "atleast_3d", + "as_strided", + "as_strided_copy", + "as_strided_scatter", + "asin", + "asinh", + "acos", + "atan", + "baddbmm", + "block_diag", + "broadcast_tensors", + "broadcast_to", + "chalf", + "cfloat", + "chunk", + "clone", + "conj", + "conj_physical", + "contiguous", + "cos", + "cosh", + "cross", + "cumsum", + "cumprod", + "cumulative_trapezoid", + "diag", + "diag_embed", + "diagflat", + "diagonal", + "diagonal_copy", + "diagonal_scatter", + "dist", + "divno_rounding_mode", + "dsplit", + "empty", + "empty_permuted", + "empty_strided", + "exp", + "expm1", + "exp2", + "expand", + "expand_as", + "expand_copy", + "gather", + "flatten", + "fill", + "full", + "full_like", + "H", + "hsplit", + "imag", + "index_add", + "index_copy", + "index_fill", + "index_select", + "index_put", + "isfinite", + "isinf", + "isreal", + "istft", + "item", + "kron", + "linalg.cross", + "linalg.diagonal", + "linalg.householder_product", + "linalg.svd", + "linalg.vander", + "linalg.vecdot", + "linalg.vector_norm", + "log10", + "log1p", + "log2", + "log", + "logaddexp", + "logaddexp2", + "logcumsumexp", + "mH", + "mT", + "masked_fill", + "masked_scatter", + "masked_select", + "meshgridlist_of_tensors", + "meshgridvariadic_tensors", + "movedim", + "mul", + "narrow", + "narrow_copy", + "neg", + "new_full", + "new_ones", + "new_zeros", + "nn.functional.conv1d", + "nn.functional.conv2d", + "nn.functional.conv_transpose1d", + "nn.functional.conv_transpose2d", + "nn.functional.conv_transpose3d", + "nn.functional.feature_alpha_dropoutwithout_train", + "nn.functional.l1_loss", + "nn.functional.linear", + "nn.functional.normalize", + "nn.functional.padcircular", + "nn.functional.pairwise_distance", + "nn.functional.softminwith_dtype", + "nn.functional.softsign", + "nn.functional.tanhshrink", + "nn.functional.triplet_margin_loss", + "nn.functional.triplet_margin_with_distance_loss", + "nn.functional.unfold", + "nonzero", + "nonzero_static", + "norm", + "normfro", + "norminf", + "ones", + "ones_like", + "outer", + "permute", + "permute_copy", + "positive", + "randn", + "ravel", + "real", + "repeat", + "repeat_interleave", + "reshape_as", + "reshape", + "resolve_conj", + "resolve_neg", + "rsqrt", + "rsub", + "scalar_tensor", + "scatter", + "scatter_add", + "select", + "sgn", + "sigmoid", + "sin", + "sinc", + "sinh", + "slice", + "softmaxwith_dtype", + "special.spherical_bessel_j0", + "special.entr", + "special.xlog1py", + "special.zeta", + "split", + "split_with_sizes", + "split_with_sizes_copy", + "splitlist_args", + "sqrt", + "squeeze", + "squeeze_copy", + "squeezemultiple", + "sub", + "svd", + "t", + "t_copy", + "take_along_dim", + "tanh", + "tan", + "tensor_split", + "tile", + "transpose", + "transpose_copy", + "tril", + "triu", + "true_divide", + "T", + "unbind", + "unbind_copy", + "unflatten", + "unfold", + "unfold_copy", + "unsafe_chunk", + "unsafe_split", + "unsqueeze", + "unsqueeze_copy", + "view_as", + "view_as_real", + "view", + "view_copy", + "vsplit", + "zero_", + "zeros", + "zeros_like", + "__rdiv__", + "__rmatmul__", + "_chunk_cat", + "acosh", + "all", + "allclose", + "angle", + "any", + "addcdiv", + "addcmul", + "addmmdecomposed", + "addmv", + "atanh", + "bfloat16", + "bmm", + "bool", + "cartesian_prod", + "cat", + "char", + "column_stack", + "combinations", + "corrcoef", + "constant_pad_nd", + "cov", + "count_nonzero", + "diff", + "div", + "dot", + "dstack", + "einsum", + "eq", + "equal", + "eye", + "fft.fft", + "fft.fft2", + "fft.fftn", + "fft.fftshift", + "fft.ifft", + "fft.ifft2", + "fft.ifftn", + "fft.ifftshift", + "fft.irfftn", + "fft.irfft2", + "fft.irfft", + "fft.hfftn", + "fft.hfft2", + "fft.hfft", + "flip", + "fliplr", + "flipud", + "float", + "gradient", + "half", + "hstack", + "inner", + "int", + "isclose", + "isnan", + "ldexp", + "lerp", + "linalg.multi_dot", + "linalg.pinv", + "linspace", + "linspacetensor_overload", + "logical_and", + "logical_not", + "logical_or", + "logical_xor", + "logsumexp", + "long", + "masked.cumsum", + "masked.cumprod", + "masked.mean", + "masked.normalize", + "masked.prod", + "masked.std", + "masked.sum", + "masked.var", + "masked.logsumexp", + "matmul", + "mean", + "mm", + "mv", + "ne", + "nn.functional.padconstant", + "nn.functional.padreflect", + "nn.functional.padreplicate", + "nn.functional.pixel_shuffle", + "nn.functional.pixel_unshuffle", + "nn.functional.rms_norm", + "pinverse", + "prod", + "reciprocal", + "roll", + "rot90", + "short", + "square", + "stack", + "stft", + "sum", + "sum_to_size", + "tensordot", + "trace", + "trapz", + "trapezoid", + "vdot", + "vstack", + "where", + "byte", + } + + MACOS_BEFORE_14_4_XFAILLIST = { + # These ops work fine in 14.4 but fail in 14.2 or 13.x + "fft.hfft2": [torch.complex64], + } + + # Those ops are not expected to work + UNIMPLEMENTED_XFAILLIST: dict[str, list | None] = { + # Failures due to lack of op implementation on MPS backend + "logspace": None, + "logspacetensor_overload": None, + "linalg.eig": None, + "linalg.eigvals": None, + "put": None, + "frexp": None, + "geqrf": None, + "nn.functional.grid_sample": None, # Unsupported Border padding mode + "hash_tensor": None, + "heaviside": None, + # "kthvalue": None, + "lcm": None, + "linalg.cond": None, + "linalg.eigh": None, + "linalg.eigvalsh": None, + "linalg.ldl_factor": None, + "linalg.ldl_factor_ex": None, + "linalg.ldl_solve": None, + "linalg.lstsq": None, + "linalg.lstsqgrad_oriented": None, + "linalg.matrix_norm": [torch.float32], + "linalg.norm": [torch.float32], + "linalg.normsubgradients_at_zero": [torch.float32], + "linalg.svdvals": None, + "masked.median": None, + "matrix_exp": None, + "max_pool2d_with_indices_backward": [ + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + ], + "median": [torch.bool], + "mode": None, + "nanmedian": [torch.bool], + "native_batch_norm": [ + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + torch.int32, + ], + "normnuc": None, + "nn.functional.avg_pool1d": [ + torch.int16, + torch.int32, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.avg_pool2d": [ + torch.int16, + torch.int32, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.avg_pool3d": [ + torch.int16, + torch.int32, + torch.uint8, + torch.int8, + ], + "nn.functional.batch_norm": [ + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + torch.int32, + ], + "nn.functional.fractional_max_pool2d": None, + "nn.functional.fractional_max_pool3d": None, + "nn.functional.group_norm": [torch.int16, torch.int32], + "nn.functional.glu": [ + torch.int32, + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + ], + "nn.functional.huber_loss": [ + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + torch.int32, + ], + "nn.functional.adaptive_avg_pool3d": None, + "nn.functional.adaptive_max_pool1d": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.adaptive_max_pool2d": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.adaptive_max_pool3d": None, + "nn.functional.interpolatearea": None, + "nn.functional.interpolatebicubic": [torch.uint8], + "nn.functional.ctc_loss": None, + "nn.functional.local_response_norm": [ + torch.int8, + torch.int16, + torch.int32, + torch.uint8, + torch.bool, + ], + "nn.functional.logsigmoid": [ + torch.int16, + torch.int32, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.max_pool1d": [ + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + torch.int32, + torch.int64, + ], + "nn.functional.max_pool2d": [torch.bool], + "nn.functional.max_pool3d": [torch.bool], + "nn.functional.max_unpool1d": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.max_unpool1dgrad": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.max_unpool2d": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.max_unpool2dgrad": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.max_unpool3d": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.max_unpool3dgrad": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.mish": [ + torch.int32, + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + ], + "nn.functional.multi_margin_loss": None, + "nn.functional.multilabel_margin_loss": [ + torch.int8, + torch.uint8, + torch.int32, + torch.int16, + torch.float32, + ], + "nn.functional.multilabel_soft_margin_loss": [ + torch.int8, + torch.uint8, + torch.int32, + torch.int16, + ], + "nn.functional.nll_loss": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.bool, + torch.int8, + ], + "nn.functional.padreplicate_negative": [torch.bool], + "nn.functional.pdist": None, + "nn.functional.relu": [torch.bool], + "nn.functional.rrelu": None, + "nn.functional.silu": [ + torch.int16, + torch.int32, + torch.uint8, + torch.int8, + ], + "nn.functional.softplus": [ + torch.int32, + torch.uint8, + torch.bool, + torch.int8, + torch.int16, + ], + "nn.functional.norm": None, + "ormqr": None, + "rounddecimals_0": [ + torch.uint8, + torch.int8, + torch.int64, + torch.int32, + torch.int16, + ], + "scatter_reduceamax": [torch.int32, torch.int64] + if MACOS_VERSION < 15.0 + else [torch.int64], + "scatter_reduceamin": [torch.int32, torch.int64] + if MACOS_VERSION < 15.0 + else [torch.int64], + "scatter_reducemean": [torch.bool], + "segment_reduce": None, + "_segment.reduce": None, + "segment.reduce": None, + "segment_reduce_offsets": None, + "_segment_reduce_offsets": None, + "_segment_reduce_lengths": None, + "_segment_reducelengths": None, + "_segment_reduceoffsets": None, + "sparse.mm": None, + "sparse.sampled_addmm": None, + "sparse.mmreduce": None, + "special.airy_ai": None, + "special.laguerre_polynomial_l": None, + "special.legendre_polynomial_p": None, + "special.log_ndtr": None, + "special.ndtri": None, + "stft": [torch.float16, torch.bfloat16], + "svd_lowrank": None, + "symeig": None, + "take": None, + "to": None, + "var_meanunbiased": [ + torch.uint8, + torch.int8, + torch.int32, + torch.int16, + torch.bool, + ], + "var_mean": [torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool], + "std_mean": [torch.uint8, torch.int8, torch.int32, torch.int16, torch.bool], + "std_meanunbiased": [ + torch.uint8, + torch.int8, + torch.int32, + torch.int16, + torch.bool, + ], + "segment_reduce_": None, + "_upsample_bilinear2d_aa": [torch.uint8], # uint8 is for CPU only + "_upsample_bicubic2d_aa": [torch.uint8], # uint8 is for CPU only + "cdouble": None, + "double": None, + "log_softmaxwith_dtype": [ + torch.uint8, + torch.int8, + torch.int32, + torch.int16, + torch.int64, + torch.float32, + ], + "float_power": None, + "linalg.matrix_rankhermitian": None, + "linalg.pinvhermitian": None, + # MPS: input sizes must be divisible by output sizes + "nn.functional.adaptive_avg_pool1d": None, + "nn.functional.adaptive_avg_pool2d": None, + # Convolution for integral types is not supported on MPS + "nn.functional.conv1d": [torch.int64], + "nn.functional.conv2d": [torch.int64], + "nn.functional.conv3d": [torch.int64], + "nn.functional.conv_transpose1d": [torch.int64], + "nn.functional.conv_transpose2d": [torch.int64, torch.bfloat16], + "nn.functional.conv_transpose3d": [ + torch.int64, + torch.bfloat16, + torch.float16, + ], + # Unsupported dtypes + # GEMM on MPS is not supported for integral types + "nn.functional.linear": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.int8, + ], + "mat": [torch.int16, torch.int32, torch.int64, torch.uint8, torch.int8], + # returned output on CPU is float64 + "bincount": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.int8, + ], + } + UNIMPLEMENTED_XFAILLIST_SPARSE: dict[str, list | None] = { + "logspace": None, + "logspacetensor_overload": None, + "linalg.eig": None, + "linalg.eigvals": None, + "put": None, + } + + if MACOS_VERSION < 15.0: + UNIMPLEMENTED_XFAILLIST.update( + { + "quantile": None, + "nanquantile": None, + } + ) + if sparse: + UNIMPLEMENTED_XFAILLIST.update(UNIMPLEMENTED_XFAILLIST_SPARSE) + + UNDEFINED_XFAILLIST: dict[str, list | None] = { + # Top 60 operators + # topk fails with duplicate indices + "topk": [ + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.int8, + ], + # Failures due to random output that they generate using + # Philox engine causing mismatch with CPU results + "multinomial": [ + torch.float16, + torch.float32, + torch.bfloat16, + ], # random results + "uniform": [torch.float16, torch.float32, torch.bfloat16], + "rand_like": [torch.float16, torch.float32, torch.bfloat16], + "randint": None, + "randint_like": None, + "randn": None, + "randn_like": None, + "bernoulli": [torch.float16, torch.float32, torch.bfloat16], + "exponential": [torch.float16, torch.float32, torch.bfloat16], + "log_normal": [torch.float16, torch.float32, torch.bfloat16], + "cauchy": [torch.float16, torch.float32, torch.bfloat16], + "geometric": [ + torch.float16, + torch.float32, + torch.bfloat16, + torch.int32, + torch.int16, + torch.int64, + torch.int8, + torch.uint8, + ], + "nn.functional.feature_alpha_dropoutwith_train": [ + torch.float16, + torch.float32, + torch.bfloat16, + ], + "normal": [torch.float16, torch.float32, torch.bfloat16], + "normalin_place": [torch.float16, torch.float32, torch.bfloat16], + "normalnumber_mean": [torch.float16, torch.float32, torch.bfloat16], + "nn.functional.alpha_dropout": [ + torch.float16, + torch.float32, + torch.bfloat16, + ], + "nn.functional.dropout": [ + torch.float16, + torch.float32, + torch.bfloat16, + torch.complex64, + ], + "nn.functional.dropout2d": [torch.float16, torch.float32, torch.bfloat16], + "nn.functional.dropout3d": [torch.float16, torch.float32, torch.bfloat16], + # See https://github.com/pytorch/pytorch/issues/111479 + "nn.functional.multi_head_attention_forward": [ + torch.float32, + torch.float16, + torch.bfloat16, + ], + # zero to negative integer powers are undefined + "__rpow__": [torch.int8, torch.int16, torch.int32, torch.int64], + "resize_": [torch.float16, torch.float32, torch.bfloat16], + "resize_as_": [torch.float16, torch.float32, torch.bfloat16], + # CPU Errors: + "addr": [ + torch.bool, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + torch.int8, + ], # "addmv_impl_cpu" not implemented for 'Half' + "as_stridedpartial_views": None, # cpu result off, showing random values + # random results + # mps vs cpu: + # Mismatched elements: 40 / 96 (41.7%) + # Greatest absolute difference: 17.892311096191406 at index (1, 0, 2) (up to 1e-05 allowed) + # Greatest relative difference: inf at index (1, 0, 0) (up to 1.3e-06 allowed) + # cuda(2.0.0.dev20230301+cu117) vs cpu: + # Mismatched elements: 56 / 96 (58.3%) + # Greatest absolute difference: 17.892311096191406 at index (1, 0, 2) (up to 1e-05 allowed) + # Greatest relative difference: inf at index (1, 0, 0) (up to 1.3e-06 allowed) + "nn.functional.scaled_dot_product_attention": [ + torch.float32, + torch.float16, + torch.bfloat16, + ], + } + + ON_MPS_XFAILLIST: dict[str, list | None] = { + # Failures due to lack of implementation of downstream functions on MPS backend + # TODO: remove these once downstream function 'aten::_linalg_svd.U' have been implemented + "linalg.matrix_rank": None, + # Exception: Caused by `torch.arange(-8.001, -4.0, dtype=torch.uint8, device="mps")` + "arange": [torch.uint8], + # Failure due to precision issue for fp16 + # on both cpu and mps there are test cases that might produce inf result + # 'nn.functional.pairwise_distance': [torch.float16], + # test blow pass on macOS 12 as it falls back to cpu + # Argsort case using duplicate indices (undefined behaviour): + # - CPU output: tensor([2546, 6917, 3181, ..., 7128, 5133, 30], device='cpu') + # - MPS output: tensor([2546, 6917, 3181, ..., 7128, 30, 5133], device='mps:0') + # Elements from index 30 and 5133 are both equal. + # Since CPU is not using argsort with stable=True, these cases result in undefined behaviour. + "argsort": [ + torch.float16, + torch.int8, + torch.uint8, + torch.bool, + torch.bfloat16, + ], + # Same issue as `argsort` with duplicate indices. This test checks both the sorted values and the indices. + # The values of the sorted tensor match the CPU, + # but in case of the returned indices this results in undefined behaviour. + "sort": [ + torch.int8, + torch.uint8, + torch.bool, + torch.float16, + torch.bfloat16, + ], + } + + EMPTY_OPS_SKIPLIST = { + # Fill tensors with uninitialized data, causing mismatch with CPU. + # They occasionally match, thus skipping them. + # See https://github.com/pytorch/pytorch/issues/100175 + "new_empty": None, + "new_empty_strided": None, + "empty_strided": None, + # CPU: empty is returning all 0's and there is a mismatch with MPS + # allocation (MacOS 13). According to + # https://pytorch.org/docs/2.0/generated/torch.empty.html + "empty": None, + "empty_like": None, + "empty_permuted": None, + } + + SKIPLIST = { + # Unsupported + # This doesn't work on M1, but is partially working on M2 with the exception of torch.float16 + "nn.functional.conv3d": None, + } + + def addDecorator(op: OpInfo, d: DecorateInfo) -> None: + if device_type is not None: + d.device_type = device_type + + op.decorators = op.decorators + (d,) + + for op in ops: + key = op.name + op.variant_test_name + addDecorator( + op, + DecorateInfo( + unittest.expectedFailure, + dtypes=[ + torch.double, + torch.cdouble, + ], + ), + ) + if sparse: + # Skipped due to test_sparse_zero_dims test in test_sparse.py which allocates empty tensor + # which leads to unexpected success with it + addDecorator( + op, + DecorateInfo( + unittest.skip( + "Skipped due to MPS not supporting complex128 tensors" + ), + dtypes=[ + torch.complex128, + ], + ), + ) + if key in EMPTY_OPS_SKIPLIST: + addDecorator( + op, + DecorateInfo( + unittest.skip("Skipping empty ops."), + dtypes=EMPTY_OPS_SKIPLIST[key], + ), + ) + if key in SKIPLIST: + addDecorator( + op, DecorateInfo(unittest.skip("Skipped!"), dtypes=SKIPLIST[key]) + ) + for xfaillist in [ + UNIMPLEMENTED_XFAILLIST, + UNDEFINED_XFAILLIST, + ON_MPS_XFAILLIST, + ]: + if key in xfaillist and key not in xfail_exclusion: + addDecorator( + op, + DecorateInfo(unittest.expectedFailure, dtypes=xfaillist[key]), + ) + + if ( + key in MACOS_BEFORE_14_4_XFAILLIST + and key not in xfail_exclusion + and (MACOS_VERSION < 14.4) + ): + addDecorator( + op, + DecorateInfo( + unittest.expectedFailure, + dtypes=MACOS_BEFORE_14_4_XFAILLIST[key], + ), + ) + + # If ops is not supported for complex types, expect it to fail + if key not in SUPPORTED_COMPLEX_OPS: + addDecorator( + op, + DecorateInfo( + unittest.expectedFailure, + dtypes=[torch.complex32, torch.complex64], + ), + ) + + return ops + + def mps_ops_grad_modifier(ops: Sequence[OpInfo]) -> Sequence[OpInfo]: + XFAILLIST_GRAD = { + # Unimplemented ops + "_segment_reduce": [torch.float16, torch.float32], + "_chunk_cat": [torch.float16, torch.float32], + "_upsample_bilinear2d_aa": None, # `_upsample_bilinear2d_aa_backward_out` not implemented for MPS + "_upsample_bicubic2d_aa": None, # `_upsample_bilinear2d_aa_backward_out` not implemented for MPS + "sparse.mmreduce": [torch.float32], # csr not supported + "linalg.householder_product": None, + "unique_consecutive": [torch.float16, torch.float32], + "scalar_tensor": [torch.float16, torch.float32], + "cdist": None, + "masked.scatter": [torch.float16, torch.float32], + "grid_sampler_2d": None, + "grid_sampler_3d": None, + "igamma": None, # currently not supported for any device + "igammac": None, # currently not supported for any device + "special.i1": [torch.float16], # "i1_backward" not implemented for 'Half' + "special.i1e": [torch.float16], # "i1e_backward" not implemented for 'Half' + # Correctness issues + # Same issue as `argsort` and `sort` with duplicate elements (undefined behaviour). + # Forward pass is passing since `msort` doesn't return the indices, just the values, which match the CPU. + # On the backward pass for `sort` both are used (values and indices), thus resulting in a issmatch between CPU and MPS. + # Running `msort` with stable `sort` passes. + "msort": [torch.float16], + # Random output + "exponential": [torch.float16, torch.float32], + "log_normal": [torch.float16, torch.float32], + "cauchy": [torch.float16, torch.float32], + "geometric": [torch.float16, torch.float32], + # CPU errors + # derivative for zeta is not implemented + "special.zeta": None, + # derivative for aten::nextafter is not implemented on CPU + "nextafter": None, + # derivative for aten::floor_divide is not implemented on CPU + "floor_divide": [torch.float16, torch.float32], + # derivative for aten::_histogramdd_from_bin_cts is not implemented on CPU + "histogramdd": [torch.float16, torch.float32], + # derivative for aten::histogram is not implemented + "histogram": [torch.float16, torch.float32], + # 'bool' object is not iterable + "allclose": [torch.float16, torch.float32], + "equal": [torch.float16, torch.float32], + # 'float' object is not iterable + "item": [torch.float16, torch.float32], + # cpu error: grad requires non-empty inputs + "randn": [torch.float16, torch.float32], + "signal.windows.bartlett": [torch.float32], + "signal.windows.blackman": [torch.float32], + "signal.windows.cosine": [torch.float32], + "signal.windows.exponential": [torch.float32], + "signal.windows.gaussian": [torch.float32], + "signal.windows.general_cosine": [torch.float32], + "signal.windows.general_hamming": [torch.float32], + "signal.windows.hamming": [torch.float32], + "signal.windows.hann": [torch.float32], + "signal.windows.kaiser": [torch.float32], + "signal.windows.nuttall": [torch.float32], + "eye": [torch.float16, torch.float32], + # Could not run 'aten::uniform_' with arguments from the 'SparseCPU' backend + "to_sparse": None, + # Exception: the derivative for '_unique2' is not implemented. + "unique": None, + } + + SKIPLIST_GRAD = { + "nn.functional.pairwise_distance": [torch.float16], + # failed assertion `destination datatype must be fp32' + "nn.functional.conv1d": [torch.float16], + "nn.functional.conv2d": [torch.float16], + "nn.functional.conv3d": [torch.float16], + "nn.functional.conv_transpose1d": [torch.float16], + "nn.functional.conv_transpose2d": [torch.float16], + "nn.functional.conv_transpose3d": [torch.float16], + } + + ON_MPS_XFAILLIST = { + # Failures due to lack of implementation of downstream functions on MPS backend + # TODO: remove these once downstream function 'aten::_linalg_svd.U' have been implemented + "linalg.matrix_rank": None, + # Exception: Caused by sample input at index 3 on MPS + "nn.functional.conv3d": [torch.float32], + } + + def addDecorator(op: OpInfo, d: DecorateInfo) -> None: + op.decorators = op.decorators + (d,) + + for op in ops: + key = op.name + op.variant_test_name + if key in XFAILLIST_GRAD: + addDecorator( + op, + DecorateInfo(unittest.expectedFailure, dtypes=XFAILLIST_GRAD[key]), + ) + + if key in SKIPLIST_GRAD: + addDecorator(op, DecorateInfo(unittest.skip, dtypes=SKIPLIST_GRAD[key])) + + if key in ON_MPS_XFAILLIST: + addDecorator( + op, + DecorateInfo( + unittest.expectedFailure, dtypes=ON_MPS_XFAILLIST[key] + ), + ) + + return ops + + def mps_ops_error_inputs_modifier(ops: Sequence[OpInfo]) -> Sequence[OpInfo]: + # Error input samples do not take a dtype argument. + XFAILLIST = { + # Exceptions are not raised + "__rmod__", + "__rsub__", + "__rpow__", + "clamp_max", + "clamp_min", + "masked_scatter", + # MPS does not support tensor dimensions > 16 + "amax", + "amin", + "aminmax", + } + + def addDecorator(op: OpInfo, d: DecorateInfo) -> None: + op.decorators = op.decorators + (d,) + + for op in ops: + key = op.name + op.variant_test_name + if key in XFAILLIST: + addDecorator(op, DecorateInfo(unittest.expectedFailure)) + + return ops +else: + + def mps_ops_modifier( + ops: Sequence[OpInfo], + device_type: str = "mps", + xfail_exclusion: list[str] | None = None, + sparse: bool = False, + ) -> Sequence[OpInfo]: + return ops diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_nn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_nn.py new file mode 100644 index 0000000000000000000000000000000000000000..7697a438f94b73ffe9d3e50696449f7b1c34628c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_nn.py @@ -0,0 +1,4013 @@ +# mypy: ignore-errors + +from abc import abstractmethod +import tempfile +import unittest + +from copy import deepcopy +from functools import reduce, partial +from itertools import product +from operator import mul + + +import torch +import torch.cuda +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import _reduction as _Reduction +from torch.testing._internal import common_utils +from torch.testing._internal.common_utils import TestCase, to_gpu, freeze_rng_state, is_iterable, \ + gradcheck, gradgradcheck, set_default_dtype, skipIfTorchDynamo, TEST_WITH_ROCM +from torch.testing._internal.common_cuda import TEST_CUDA, SM90OrLater +from torch.autograd.gradcheck import _get_numerical_jacobian, _iter_tensors +from torch.autograd import Variable +from torch.types import _TensorOrTensors +import torch.backends.cudnn + +from typing import Any +from collections.abc import Callable +from collections.abc import Sequence + +TemporaryFile = tempfile.TemporaryFile +PRECISION = 1e-5 + + +def get_reduction(m): + result = getattr(m, 'reduction', None) + if result is None: + result = _Reduction.legacy_get_string(getattr(m, 'sizeAverage', None), True, emit_warning=False) + if result is None: + raise AssertionError("Expected result to not be None") + return result + + +def get_weight(m): + result = getattr(m, 'weight', None) + if result is not None: + return result + return getattr(m, 'weights', None) + +# NOTE [How to check NN module / functional API parity between Python and C++ frontends] +# +# The way to check API parity is to add parity tests for the NN module / functional of interest. +# Here are the detailed steps: +# +# For NN module: +# 1. Make sure you already have a test dict with the module configuration you want to test. +# 2. Add `cpp_constructor_args` entry to the test dict, with its value exactly matching +# the Python module constructor arguments. For example, if in the test dict we pass +# `(10, 8)` to `torch.nn.Linear` constructor, then we should pass `torch::nn::LinearOptions(10, 8)` +# as the corresponding C++ constructor argument to `torch::nn::Linear`. +# 3. If in the process of performing the above step you referenced any variables +# in the `cpp_constructor_args` entry, you must add `cpp_var_map` entry +# to the test dict to make sure that those variables are populated with the right Python values. +# For example, if the Python constructor call is +# `torch.nn.FractionalMaxPool2d(2, output_ratio=0.5, _random_samples=random_samples)`, +# the corresponding C++ constructor argument is +# `torch::nn::FractionalMaxPool2dOptions(2).output_ratio(0.5)._random_samples(random_samples)`, +# and the `cpp_var_map` entry must be +# `{'random_samples': random_samples}` in order to populate the C++ variable `random_samples` +# used in the C++ constructor argument with the Python tensor value `random_samples`. +# +# For NN functional: +# 1. Make sure you already have a test dict with the functional configuration you want to test. +# 2. If the test dict's `constructor` entry looks like `wrap_functional(F.some_functional_name, ...)`, +# then you must add `cpp_options_args` entry to the test dict, with its value exactly matching the Python +# functional optional arguments. For example, if the test dict's `constructor` entry is +# `wrap_functional(F.interpolate, size=12, scale_factor=None, mode='nearest')`, +# then the `cpp_options_args` entry should be +# "F::InterpolateFuncOptions().size(std::vector({12})).scale_factor(std::nullopt).mode(torch::kNearest)". +# 3. Otherwise, if the test dict's `constructor` entry looks like +# `wrap_functional(lambda i: F.some_functional_name(...))`, +# then you must add `cpp_function_call` entry to the test dict, with its value exactly matching the Python +# functional function call. For example, if the test dict's `constructor` entry is +# `wrap_functional(lambda i: F.poisson_nll_loss(i, t.type_as(i), reduction='none'))`, +# then the `cpp_function_call` entry should be +# "F::poisson_nll_loss(i, t.to(i.options()), F::PoissonNLLLossFuncOptions().reduction(torch::kNone))". +# 4. If in the process of performing the above two steps you referenced any variables +# in the `cpp_options_args` or `cpp_function_call` entry, you must +# add `cpp_var_map` entry to the test dict to make sure that those variables +# are populated with the right Python values. For example, if the test dict's `constructor` entry is +# `wrap_functional(lambda i: F.poisson_nll_loss(i, t.type_as(i), reduction='none'))`, +# then the `cpp_function_call` entry should be +# "F::poisson_nll_loss(i, t.to(i.options()), F::PoissonNLLLossFuncOptions().reduction(torch::kNone))". +# Notice that there are two variables `i` and `t` that need to have their values provided, +# and the way to do so is to add a `cpp_var_map` entry: `cpp_var_map={'i': '_get_input()', 't': t}`. +# (Note that for `i`, since we want it to take the Python input value, we pass '_get_input()' string as value +# and the C++ parity test mechanism will populate `i` with the Python input value correctly.) +# +# There are also a few optional flags in the test dict to control the C++ parity test behavior: +# +# - `test_cpp_api_parity`: if `False`, skips the C++ parity test for this test dict. Default: True. +# - `has_parity`: if `False`, expects this test dict to fail the C++ parity test. Default: True. + + +module_tests = [ + dict( + module_name='Linear', + constructor_args=(10, 8), + cpp_constructor_args='torch::nn::LinearOptions(10, 8)', + input_size=(4, 10), + reference_fn=lambda i, p, _: torch.mm(i, p[0].t()) + p[1].view(1, -1).expand(4, 8), + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Linear', + constructor_args=(10, 8, False), + cpp_constructor_args='torch::nn::LinearOptions(10, 8).bias(false)', + input_size=(4, 10), + desc='no_bias', + reference_fn=lambda i, p, _: torch.mm(i, p[0].t()), + with_tf32=True, + tf32_precision=0.005, + # ROCM: skipping tf32 test on gfx94 archs due to tolerance issue. + test_cuda=not (TEST_WITH_ROCM and "gfx94" in torch.cuda.get_device_properties(0).gcnArchName), + default_dtype=torch.double, + ), + dict( + module_name='RReLU', + input_size=(1, 2, 2), + test_cuda=False, + default_dtype=torch.double, + ), + dict( + module_name='RReLU', + constructor_args=(0.1, 0.9), + cpp_constructor_args='torch::nn::RReLUOptions().lower(0.1).upper(0.9)', + input_size=(4, 4, 5), + desc='with_up_down', + test_cuda=False, + default_dtype=torch.double, + ), + dict( + module_name='Flatten', + input_size=(2, 3, 4, 5), + reference_fn=lambda i, *_: torch.flatten(i, 1), + default_dtype=torch.double, + ), + # TODO: reference function + dict( + module_name='CrossMapLRN2d', + constructor_args=(5, 5e-3, 1e-3, 2), + cpp_constructor_args='torch::nn::CrossMapLRN2dOptions(5).alpha(5e-3).beta(1e-3).k(2)', + input_size=(2, 3, 6, 6), + check_gradgrad=False, + # TODO(#50743): Figure out the error. "RuntimeError: Unrecognized tensor type ID: Batched" + check_batched_grad=False, + default_dtype=torch.double, + ), +] + + +# Generates rand tensor with non-equal values. This ensures that duplicate +# values won't be causing test failure for modules like MaxPooling. +# size should be small, otherwise randperm fails / long overflows. +def _rand_tensor_non_equal(*size): + total = reduce(mul, size, 1) + return torch.randperm(total).view(*size).double() + + +def wrap_functional(fn, **kwargs): + class FunctionalModule(nn.Module): + def forward(self, *args): + return fn(*args, **kwargs) + return FunctionalModule + + +def poissonnllloss_no_reduce_test(): + t = torch.randn(10, 10) + return dict( + fullname='PoissonNLLLoss_no_reduce', + constructor=wrap_functional( + lambda i: F.poisson_nll_loss(i, t.type_as(i), reduction='none')), + cpp_function_call='F::poisson_nll_loss(' + 'i, t.to(i.options()), F::PoissonNLLLossFuncOptions().reduction(torch::kNone))', + input_fn=lambda: torch.rand(10, 10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: i.exp() - t.mul(i), + pickle=False, + default_dtype=torch.double) + + +def bceloss_no_reduce_test(): + t = Variable(torch.randn(15, 10).gt(0).to(torch.double)) + return dict( + fullname='BCELoss_no_reduce', + constructor=wrap_functional( + lambda i: F.binary_cross_entropy(i, t.type_as(i), reduction='none')), + cpp_function_call='F::binary_cross_entropy(' + 'i, t.to(i.options()), F::BinaryCrossEntropyFuncOptions().reduction(torch::kNone))', + input_fn=lambda: torch.rand(15, 10).clamp_(2.8e-2, 1 - 2.8e-2), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: -(t * i.log() + (1 - t) * (1 - i).log()), + pickle=False, + precision=7e-4, + default_dtype=torch.double) + + +def bceloss_no_reduce_scalar_test(): + t = torch.randn(()).gt(0).to(torch.double) + return dict( + fullname='BCELoss_no_reduce_scalar', + constructor=wrap_functional( + lambda i: F.binary_cross_entropy(i, t.type_as(i), reduction='none')), + cpp_function_call='F::binary_cross_entropy(' + 'i, t.to(i.options()), F::BinaryCrossEntropyFuncOptions().reduction(torch::kNone))', + input_fn=lambda: torch.rand(()).clamp_(2.8e-2, 1 - 2.8e-2), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: -(t * i.log() + (1 - t) * (1 - i).log()), + pickle=False, + default_dtype=torch.double) + + +def bceloss_weights_no_reduce_test(): + t = Variable(torch.randn(15, 10, dtype=torch.double).gt(0).to(torch.double)) + weights = torch.rand(10, dtype=torch.double) + return dict( + fullname='BCELoss_weights_no_reduce', + constructor=wrap_functional( + lambda i: F.binary_cross_entropy(i, t.type_as(i), + weight=weights.type_as(i), reduction='none')), + cpp_function_call='F::binary_cross_entropy(' + 'i, t.to(i.options()), ' + 'F::BinaryCrossEntropyFuncOptions().weight(weights.to(i.options())).reduction(torch::kNone))', + input_fn=lambda: torch.rand(15, 10).clamp_(2.8e-2, 1 - 2.8e-2), + cpp_var_map={'i': '_get_input()', 't': t, 'weights': weights}, + reference_fn=lambda i, p, m: -(t * i.log() + (1 - t) * (1 - i).log()) * weights, + pickle=False, + precision=3e-4, + default_dtype=torch.double, + ) + + +def bceloss_weights_no_reduce_scalar_test(): + t = torch.randn(()).gt(0).to(torch.double) + weights = torch.rand((), dtype=torch.double) + return dict( + fullname='BCELoss_weights_no_reduce_scalar', + constructor=wrap_functional( + lambda i: F.binary_cross_entropy(i, t.type_as(i), + weight=weights.type_as(i), reduction='none')), + cpp_function_call='''F::binary_cross_entropy( + i, t.to(i.options()), + F::BinaryCrossEntropyFuncOptions().weight(weights.to(i.options())).reduction(torch::kNone))''', + cpp_var_map={'i': '_get_input()', 't': t, 'weights': weights}, + input_fn=lambda: torch.rand(()).clamp_(2.8e-2, 1 - 2.8e-2), + reference_fn=lambda i, *_: -(t * i.log() + (1 - t) * (1 - i).log()) * weights, + pickle=False, + default_dtype=torch.double, + ) + + +def bce_with_logistic_legacy_enum_test(): + t = Variable(torch.randn(15, 10).gt(0).to(torch.double)) + sigmoid = nn.Sigmoid() + return dict( + fullname='BCEWithLogitsLoss_legacy_enum', + constructor=wrap_functional( + lambda i: F.binary_cross_entropy_with_logits(i, t.type_as(i), reduce=False)), + cpp_function_call='''F::binary_cross_entropy_with_logits( + i, t.to(i.options()), F::BinaryCrossEntropyWithLogitsFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.rand(15, 10).clamp_(2.8e-2, 1 - 2.8e-2), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: -(t * sigmoid(i).log() + (1 - t) * (1 - sigmoid(i)).log()), + check_gradgrad=False, + pickle=False, + default_dtype=torch.double, + ) + + +def bce_with_logistic_no_reduce_test(): + t = Variable(torch.randn(15, 10).gt(0).to(torch.double)) + sigmoid = nn.Sigmoid() + return dict( + fullname='BCEWithLogitsLoss_no_reduce', + constructor=wrap_functional( + lambda i: F.binary_cross_entropy_with_logits(i, t.type_as(i), reduction='none')), + cpp_function_call='''F::binary_cross_entropy_with_logits( + i, t.to(i.options()), F::BinaryCrossEntropyWithLogitsFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.rand(15, 10).clamp_(2.8e-2, 1 - 2.8e-2), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: -(t * sigmoid(i).log() + (1 - t) * (1 - sigmoid(i)).log()), + check_gradgrad=False, + pickle=False, + default_dtype=torch.double, + ) + + +def bce_with_logistic_no_reduce_scalar_test(): + t = torch.randn(()).gt(0).to(torch.double) + sigmoid = nn.Sigmoid() + return dict( + fullname='BCEWithLogitsLoss_no_reduce_scalar', + constructor=wrap_functional( + lambda i: F.binary_cross_entropy_with_logits(i, t.type_as(i), reduction='none')), + cpp_function_call='''F::binary_cross_entropy_with_logits( + i, t.to(i.options()), F::BinaryCrossEntropyWithLogitsFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.rand(()).clamp_(2.8e-2, 1 - 2.8e-2), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: -(t * sigmoid(i).log() + (1 - t) * (1 - sigmoid(i)).log()), + check_gradgrad=False, + pickle=False, + default_dtype=torch.double, + ) + + +def kldivloss_with_target_no_reduce_test(): + t = torch.rand(10, 10, dtype=torch.double) + return dict( + fullname='KLDivLoss_with_target_no_reduce', + constructor=wrap_functional( + lambda i: F.kl_div(i, t.type_as(i), reduction='none')), + cpp_function_call='F::kl_div(i, t.to(i.options()), F::KLDivFuncOptions().reduction(torch::kNone))', + input_fn=lambda: torch.rand(10, 10).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['KLDivLoss'](i, t.type_as(i), reduction='none'), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def kldivloss_no_reduce_test(): + t = torch.rand(10, 10, dtype=torch.double) + return dict( + fullname='KLDivLoss_no_reduce', + constructor=wrap_functional( + lambda i: F.kl_div(i, t.type_as(i), reduction='none')), + cpp_function_call='F::kl_div(i, t.to(i.options()), F::KLDivFuncOptions().reduction(torch::kNone))', + input_fn=lambda: torch.rand(10, 10).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['KLDivLoss'](i, t.type_as(i), reduction='none'), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double, + ) + + +def kldivloss_no_reduce_scalar_test(): + t = torch.rand((), dtype=torch.double) + return dict( + fullname='KLDivLoss_no_reduce_scalar', + constructor=wrap_functional( + lambda i: F.kl_div(i, t.type_as(i), reduction='none')), + cpp_function_call='F::kl_div(i, t.to(i.options()), F::KLDivFuncOptions().reduction(torch::kNone))', + input_fn=lambda: torch.rand(()).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['KLDivLoss'](i, t.type_as(i), reduction='none'), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def kldivloss_with_log_target_no_reduce_test(): + t = torch.rand(10, 10, dtype=torch.double).log() + return dict( + fullname='KLDivLoss_with_log_target_no_reduce', + constructor=wrap_functional( + lambda i: F.kl_div(i, t.type_as(i), reduction='none', log_target=True)), + cpp_function_call='F::kl_div(i, t.to(i.options()), F::KLDivFuncOptions().reduction(torch::kNone).log_target(true))', + input_fn=lambda: torch.rand(10, 10).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['KLDivLoss_log_target'](i, t.type_as(i), reduction='none'), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def kldivloss_no_reduce_log_target_test(): + t = torch.rand(10, 10, dtype=torch.double).log() + return dict( + fullname='KLDivLoss_no_reduce_log_target', + constructor=wrap_functional( + lambda i: F.kl_div(i, t.type_as(i), reduction='none', log_target=True)), + cpp_function_call='F::kl_div(i, t.to(i.options()), F::KLDivFuncOptions().reduction(torch::kNone).log_target(true))', + input_fn=lambda: torch.rand(10, 10).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['KLDivLoss_log_target'](i, t.type_as(i), reduction='none'), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double, + ) + + +def kldivloss_no_reduce_scalar_log_target_test(): + t = torch.rand((), dtype=torch.double).log() + return dict( + fullname='KLDivLoss_no_reduce_scalar_log_target', + constructor=wrap_functional( + lambda i: F.kl_div(i, t.type_as(i), reduction='none', log_target=True)), + cpp_function_call='F::kl_div(i, t.to(i.options()), F::KLDivFuncOptions().reduction(torch::kNone).log_target(true))', + input_fn=lambda: torch.rand(()).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['KLDivLoss_log_target'](i, t.type_as(i), reduction='none'), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def l1loss_no_reduce_test(): + t = torch.randn(2, 3, 4, dtype=torch.double) + return dict( + fullname='L1Loss_no_reduce', + constructor=wrap_functional( + lambda i: F.l1_loss(i, t.type_as(i), reduction='none')), + cpp_function_call='F::l1_loss(i, t.to(i.options()), F::L1LossFuncOptions().reduction(torch::kNone))', + input_fn=lambda: torch.randn(2, 3, 4), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: (i - t.type_as(i)).abs(), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def l1loss_no_reduce_complex_test(): + t = torch.randn(2, 3, 4, dtype=torch.cdouble) + return dict( + fullname='L1Loss_no_reduce_complex', + constructor=wrap_functional( + lambda i: F.l1_loss(i, t.type_as(i), reduction='none')), + cpp_function_call='F::l1_loss(i, t.to(i.options()), F::L1LossFuncOptions().reduction(torch::kNone))', + input_fn=lambda: torch.randn(2, 3, 4, dtype=torch.cdouble), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: (i - t.type_as(i)).abs(), + supports_forward_ad=True, + pickle=False) + + +def l1loss_no_reduce_scalar_test(): + t = torch.randn((), dtype=torch.double) + return dict( + fullname='L1Loss_no_reduce_scalar', + constructor=wrap_functional( + lambda i: F.l1_loss(i, t.type_as(i), reduction='none')), + cpp_function_call='F::l1_loss(i, t.to(i.options()), F::L1LossFuncOptions().reduction(torch::kNone))', + input_fn=lambda: torch.randn(()), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: (i - t.type_as(i)).abs(), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def mseloss_no_reduce_test(): + input_size = (2, 3, 4, 5) + target = torch.randn(*input_size, dtype=torch.double) + return dict( + fullname='MSELoss_no_reduce', + constructor=wrap_functional( + lambda i: F.mse_loss(i, target.type_as(i), reduction='none')), + cpp_function_call='F::mse_loss(i, target.to(i.options()), F::MSELossFuncOptions().reduction(torch::kNone))', + input_size=input_size, + cpp_var_map={'i': '_get_input()', 'target': target}, + reference_fn=lambda i, *_: (i - target).pow(2), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def mseloss_no_reduce_scalar_test(): + input_size = () + target = torch.randn(input_size, dtype=torch.double) + return dict( + fullname='MSELoss_no_reduce_scalar', + constructor=wrap_functional( + lambda i: F.mse_loss(i, target.type_as(i), reduction='none')), + cpp_function_call='F::mse_loss(i, target.to(i.options()), F::MSELossFuncOptions().reduction(torch::kNone))', + input_size=input_size, + cpp_var_map={'i': '_get_input()', 'target': target}, + reference_fn=lambda i, *_: (i - target).pow(2), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def nllloss_no_reduce_test(): + t = Variable(torch.empty(15).uniform_().mul(10).floor().long()) + kwargs = {'reduction': 'none'} + return dict( + fullname='NLLLoss_no_reduce', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), reduction=kwargs['reduction'])), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), F::NLLLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.rand(15, 10).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLoss'](i, t.type_as(i).long(), **kwargs), + pickle=False, + default_dtype=torch.double) + + +def nllloss_no_reduce_ignore_index_test(): + t = Variable(torch.empty(15).uniform_().mul(10).floor().long()) + kwargs: dict[str, int | str] = {'ignore_index': 2, 'reduction': 'none'} + return dict( + fullname='NLLLoss_no_reduce_ignore_index', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), ignore_index=int(kwargs['ignore_index']), + reduction=str(kwargs['reduction']))), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), F::NLLLossFuncOptions().ignore_index(2).reduction(torch::kNone))''', + input_fn=lambda: torch.rand(15, 10).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLoss'](i, t.type_as(i).long(), **kwargs), + pickle=False, + default_dtype=torch.double) + + +def nllloss_no_reduce_weights_test(): + t = Variable(torch.empty(15).uniform_().mul(10).floor().long()) + weight = torch.rand(10) + + def kwargs(i): + return {'weight': weight.type_as(i), 'reduction': 'none'} + + return dict( + fullname='NLLLoss_no_reduce_weights', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), **kwargs(i))), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), + F::NLLLossFuncOptions().weight(weight.to(i.options())).reduction(torch::kNone))''', + input_fn=lambda: torch.rand(15, 10).add(1e-2).log(), + cpp_var_map={'i': '_get_input()', 't': t, 'weight': weight}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLoss'](i, t.type_as(i).long(), **kwargs(i)), + pickle=False, + default_dtype=torch.double) + + +def nllloss_no_reduce_weights_ignore_index_test(): + t = Variable(torch.empty(15).uniform_().mul(10).floor().long()) + weight = torch.rand(10) + + def kwargs(i): + return {'weight': weight.type_as(i), 'reduction': 'none', + 'ignore_index': 2} + + return dict( + fullname='NLLLoss_no_reduce_weights_ignore_index', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), **kwargs(i.data))), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), + F::NLLLossFuncOptions().weight(weight.to(i.options())).reduction(torch::kNone).ignore_index(2))''', + input_fn=lambda: torch.rand(15, 10).add(1e-2).log(), + cpp_var_map={'i': '_get_input()', 't': t, 'weight': weight}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLoss'](i, t.type_as(i).long(), **kwargs(i)), + pickle=False, + default_dtype=torch.double) + + +def nllloss_no_reduce_weights_ignore_index_neg_test(): + t = Variable(torch.empty(15).uniform_().mul(10).floor().long()) + weight = torch.rand(10) + + def kwargs(i): + return {'weight': weight.type_as(i), 'reduction': 'none', + 'ignore_index': -1} + + return dict( + fullname='NLLLoss_no_reduce_weights_ignore_index_neg', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), **kwargs(i))), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), + F::NLLLossFuncOptions().weight(weight.to(i.options())).reduction(torch::kNone).ignore_index(-1))''', + input=torch.rand(15, 10, dtype=torch.double).add(1e-2).log(), + cpp_var_map={'i': '_get_input()', 't': t, 'weight': weight}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLoss'](i, t.type_as(i).long(), **kwargs(i)), + pickle=False, + default_dtype=torch.double) + + +def nllloss2d_no_reduce_test(): + t = Variable(torch.rand(2, 5, 5).mul(3).floor().long()) + kwargs = {'reduction': 'none'} + return dict( + fullname='NLLLoss2d_no_reduce', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), reduction=kwargs['reduction'])), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), F::NLLLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.rand(2, 3, 5, 5).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLossNd'](i, t.type_as(i).long(), **kwargs), + pickle=False, + default_dtype=torch.double) + + +def nllloss2d_no_reduce_ignore_index_test(): + t = Variable(torch.rand(2, 5, 5).mul(3).floor().long()) + kwargs: dict[str, int | str] = {'ignore_index': 1, 'reduction': 'none'} + return dict( + fullname='NLLLoss2d_no_reduce_ignore_index', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), ignore_index=int(kwargs['ignore_index']), + reduction=str(kwargs['reduction']))), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), F::NLLLossFuncOptions().ignore_index(1).reduction(torch::kNone))''', + input_fn=lambda: torch.rand(2, 3, 5, 5).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLossNd'](i, t.type_as(i).long(), **kwargs), + pickle=False, + default_dtype=torch.double) + + +def nllloss2d_no_reduce_weights_test(): + t = Variable(torch.rand(2, 5, 5).mul(3).floor().long()) + weight = torch.rand(3) + + def kwargs(i): + return {'weight': weight.type_as(i), 'reduction': 'none'} + + return dict( + fullname='NLLLoss2d_no_reduce_weights', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), **kwargs(i))), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), + F::NLLLossFuncOptions().weight(weight.to(i.options())).reduction(torch::kNone))''', + input_fn=lambda: torch.rand(2, 3, 5, 5).log(), + cpp_var_map={'i': '_get_input()', 't': t, 'weight': weight}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLossNd'](i, t.type_as(i).long(), **kwargs(i)), + pickle=False, + default_dtype=torch.double) + + +def nlllossNd_no_reduce_test(): + t = Variable(torch.rand(2, 5, 5, 2, 2).mul(3).floor().long()) + kwargs = {'reduction': 'none'} + return dict( + fullname='NLLLossNd_no_reduce', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), reduction=kwargs['reduction'])), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), F::NLLLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.rand(2, 3, 5, 5, 2, 2).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLossNd'](i, t.type_as(i).long(), **kwargs), + pickle=False, + default_dtype=torch.double) + + +def nlllossNd_no_reduce_ignore_index_test(): + t = Variable(torch.rand(2, 5, 5, 2, 2).mul(3).floor().long()) + kwargs: dict[str, int | str] = {'ignore_index': 1, 'reduction': 'none'} + return dict( + fullname='NLLLossNd_no_reduce_ignore_index', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), ignore_index=int(kwargs['ignore_index']), + reduction=str(kwargs['reduction']))), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), F::NLLLossFuncOptions().ignore_index(1).reduction(torch::kNone))''', + input_fn=lambda: torch.rand(2, 3, 5, 5, 2, 2).log(), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLossNd'](i, t.type_as(i).long(), **kwargs), + pickle=False, + default_dtype=torch.double) + + +def nlllossNd_no_reduce_weights_test(): + t = Variable(torch.rand(2, 5, 5, 2, 2).mul(3).floor().long()) + weight = torch.rand(3) + + def kwargs(i): + return {'weight': weight.type_as(i), 'reduction': 'none'} + + return dict( + fullname='NLLLossNd_no_reduce_weights', + constructor=wrap_functional( + lambda i: F.nll_loss(i, t.type_as(i).long(), **kwargs(i))), + cpp_function_call='''F::nll_loss( + i, t.to(i.options()).to(torch::kLong), + F::NLLLossFuncOptions().weight(weight.to(i.options())).reduction(torch::kNone))''', + input_fn=lambda: torch.rand(2, 3, 5, 5, 2, 2).log(), + cpp_var_map={'i': '_get_input()', 't': t, 'weight': weight}, + reference_fn=lambda i, *_: + loss_reference_fns['NLLLossNd'](i, t.type_as(i).long(), **kwargs(i)), + pickle=False, + default_dtype=torch.double) + + +def smoothl1loss_no_reduce_test(): + t = torch.randn(2, 3, 4, dtype=torch.double) + return dict( + fullname='SmoothL1Loss_no_reduce', + constructor=wrap_functional( + lambda i: F.smooth_l1_loss(i, t.type_as(i), reduction='none')), + cpp_function_call='''F::smooth_l1_loss( + i, t.to(i.options()), F::SmoothL1LossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(2, 3, 4), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['SmoothL1Loss'](i, t.type_as(i), reduction='none'), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def smoothl1loss_no_reduce_scalar_test(): + t = torch.randn((), dtype=torch.double) + return dict( + fullname='SmoothL1Loss_no_reduce_scalar', + constructor=wrap_functional( + lambda i: F.smooth_l1_loss(i, t.type_as(i), reduction='none')), + cpp_function_call='''F::smooth_l1_loss( + i, t.to(i.options()), F::SmoothL1LossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(()), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['SmoothL1Loss'](i, t.type_as(i), reduction='none'), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def smoothl1loss_beta_test(): + t = torch.randn(2, 3, 4, dtype=torch.double) + return dict( + fullname='SmoothL1Loss_beta', + constructor=wrap_functional( + lambda i: F.smooth_l1_loss(i, t.type_as(i), reduction='none', beta=0.5)), + cpp_function_call='''F::smooth_l1_loss( + i, t.to(i.options()), F::SmoothL1LossFuncOptions().reduction(torch::kNone), 0.5)''', + input_fn=lambda: torch.randn(2, 3, 4), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['SmoothL1Loss'](i, t.type_as(i), reduction='none', beta=0.5), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def smoothl1loss_zero_beta_test(): + t = torch.randn(2, 3, 4, dtype=torch.double) + return dict( + fullname='SmoothL1Loss_zero_beta', + constructor=wrap_functional( + lambda i: F.smooth_l1_loss(i, t.type_as(i), reduction='none', beta=0)), + cpp_function_call='''F::smooth_l1_loss( + i, t.to(i.options()), F::SmoothL1LossFuncOptions().reduction(torch::kNone), 0)''', + input_fn=lambda: torch.randn(2, 3, 4), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['SmoothL1Loss'](i, t.type_as(i), reduction='none', beta=0), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def huberloss_delta_test(): + t = torch.randn(2, 3, 4) + return dict( + fullname='HuberLoss_delta', + constructor=wrap_functional( + lambda i: F.huber_loss(i, t.type_as(i), reduction='none', delta=0.5)), + cpp_function_call='''F::huber_loss( + i, t.to(i.options()), F::HuberLossFuncOptions().reduction(torch::kNone).delta(0.5))''', + input_fn=lambda: torch.randn(2, 3, 4), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['HuberLoss'](i, t.type_as(i), reduction='none', delta=0.5), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def multilabelmarginloss_0d_no_reduce_test(): + t = torch.zeros(()).long() + return dict( + fullname='MultiLabelMarginLoss_0d_no_reduce', + constructor=wrap_functional( + lambda i: F.multilabel_margin_loss(i, t.type_as(i).long(), reduction='none')), + cpp_function_call='''F::multilabel_margin_loss( + i, t.to(i.options()).to(torch::kLong), F::MultilabelMarginLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(()), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiLabelMarginLoss'](i, t.data.type_as(i).long(), reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False) + + +def multilabelmarginloss_1d_no_reduce_test(): + t = Variable(torch.rand(10).mul(10).floor().long()) + return dict( + fullname='MultiLabelMarginLoss_1d_no_reduce', + constructor=wrap_functional( + lambda i: F.multilabel_margin_loss(i, t.type_as(i).long(), reduction='none')), + cpp_function_call='''F::multilabel_margin_loss( + i, t.to(i.options()).to(torch::kLong), F::MultilabelMarginLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiLabelMarginLoss'](i, t.data.type_as(i).long(), reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def multilabelmarginloss_index_neg_test(): + t = Variable(torch.clamp(torch.rand(5, 10).add(-.5).mul(20).floor().long(), min=-1)) + return dict( + fullname='MultiLabelMarginLoss_index_neg', + constructor=wrap_functional( + lambda i: F.multilabel_margin_loss(i, t.type_as(i).long(), reduction='none')), + cpp_function_call='''F::multilabel_margin_loss( + i, t.to(i.options()).to(torch::kLong), F::MultilabelMarginLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(5, 10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiLabelMarginLoss'](i, t.data.type_as(i).long(), reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def multilabelmarginloss_no_reduce_test(): + t = Variable(torch.rand(5, 10).mul(10).floor().long()) + return dict( + fullname='MultiLabelMarginLoss_no_reduce', + constructor=wrap_functional( + lambda i: F.multilabel_margin_loss(i, t.type_as(i).long(), reduction='none')), + cpp_function_call='''F::multilabel_margin_loss( + i, t.to(i.options()).to(torch::kLong), F::MultilabelMarginLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(5, 10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiLabelMarginLoss'](i, t.data.type_as(i).long(), reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def hingeembeddingloss_no_reduce_test(): + t = Variable(torch.randn(10).gt(0).to(torch.double).mul_(2).sub(1)) + return dict( + fullname='HingeEmbeddingLoss_no_reduce', + constructor=wrap_functional( + lambda i: F.hinge_embedding_loss(i, t.type_as(i), reduction='none')), + cpp_function_call='''F::hinge_embedding_loss( + i, t.to(i.options()), F::HingeEmbeddingLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['HingeEmbeddingLoss'](i, t.type_as(i), reduction='none'), + check_sum_reduction=True, + pickle=False, + default_dtype=torch.double) + + +def hingeembeddingloss_margin_no_reduce_test(): + t = Variable(torch.randn(10).gt(0).to(torch.double).mul_(2).sub(1)) + return dict( + fullname='HingeEmbeddingLoss_margin_no_reduce', + constructor=wrap_functional( + lambda i: F.hinge_embedding_loss(i, t.type_as(i), margin=0.5, reduction='none')), + cpp_function_call='''F::hinge_embedding_loss( + i, t.to(i.options()), F::HingeEmbeddingLossFuncOptions().margin(0.5).reduction(torch::kNone))''', + input_fn=lambda: torch.randn(10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['HingeEmbeddingLoss'](i, t.type_as(i), margin=0.5, reduction='none'), + check_sum_reduction=True, + pickle=False, + default_dtype=torch.double) + + +def softmarginloss_no_reduce_test(): + t = torch.randn(5, 5, dtype=torch.double) + return dict( + fullname='SoftMarginLoss_no_reduce', + constructor=wrap_functional( + lambda i: F.soft_margin_loss(i, t.type_as(i), reduction='none')), + cpp_function_call='''F::soft_margin_loss( + i, t.to(i.options()), F::SoftMarginLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(5, 5), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['SoftMarginLoss'](i, t.type_as(i), reduction='none'), + supports_forward_ad=True, + pickle=False, + default_dtype=torch.double) + + +def multilabelsoftmarginloss_no_reduce_test(): + t = torch.rand(5, 10).mul(2).floor() + return dict( + fullname='MultiLabelSoftMarginLoss_no_reduce', + constructor=wrap_functional( + lambda i: F.multilabel_soft_margin_loss(i, t.type_as(i), reduction='none')), + cpp_function_call='''F::multilabel_soft_margin_loss( + i, t.to(i.options()), F::MultilabelSoftMarginLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(5, 10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + (-(t * i.sigmoid().log() + (1 - t) * (-i).sigmoid().log())).sum(dim=1) / i.size(1), + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def multilabelsoftmarginloss_weights_no_reduce_test(): + t = torch.rand(5, 10).mul(2).floor() + weights = torch.rand(10) + return dict( + fullname='MultiLabelSoftMarginLoss_weights_no_reduce', + constructor=wrap_functional( + lambda i: F.multilabel_soft_margin_loss(i, t.type_as(i), + weight=weights.type_as(i), reduction='none')), + cpp_function_call='''F::multilabel_soft_margin_loss( + i, t.to(i.options()), + F::MultilabelSoftMarginLossFuncOptions().weight(weights.to(i.options())).reduction(torch::kNone))''', + input_fn=lambda: torch.randn(5, 10), + cpp_var_map={'i': '_get_input()', 't': t, 'weights': weights}, + reference_fn=lambda i, *_: + (-(t * i.sigmoid().log() + (1 - t) * (-i).sigmoid().log()) * weights).sum(dim=1) / i.size(1), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def multimarginloss_no_reduce_test(): + t = torch.rand(5).mul(8).floor().long() + return dict( + fullname='MultiMarginLoss_no_reduce', + constructor=wrap_functional( + lambda i: F.multi_margin_loss(i, t.type_as(i).long(), reduction='none')), + cpp_function_call='''F::multi_margin_loss( + i, t.to(i.options()).to(torch::kLong), F::MultiMarginLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(5, 10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiMarginLoss'](i, t.data.type_as(i).long(), reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def multimarginloss_1d_no_reduce_test(): + t = torch.rand(1).mul(8).floor().long() + return dict( + fullname='MultiMarginLoss_1d_no_reduce', + constructor=wrap_functional( + lambda i: F.multi_margin_loss(i, t.type_as(i).long(), reduction='none')), + cpp_function_call='''F::multi_margin_loss( + i, t.to(i.options()).to(torch::kLong), F::MultiMarginLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiMarginLoss'](i, t.data.type_as(i).long(), reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def multimarginloss_1d_input_0d_target_no_reduce_test(): + t = torch.rand(()).mul(8).floor().long() + return dict( + fullname='multimarginloss_1d_input_0d_target_no_reduce', + constructor=wrap_functional( + lambda i: F.multi_margin_loss(i, t.type_as(i).long(), reduction='none')), + cpp_function_call='''F::multi_margin_loss( + i, t.to(i.options()).to(torch::kLong), F::MultiMarginLossFuncOptions().reduction(torch::kNone))''', + input_fn=lambda: torch.randn(10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiMarginLoss'](i, t.data.type_as(i).long(), reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def multimarginloss_p_no_reduce_test(): + t = torch.rand(5).mul(8).floor().long() + return dict( + fullname='MultiMarginLoss_p_no_reduce', + constructor=wrap_functional( + lambda i: F.multi_margin_loss(i, t.type_as(i).long(), p=2, reduction='none')), + cpp_function_call='''F::multi_margin_loss( + i, t.to(i.options()).to(torch::kLong), F::MultiMarginLossFuncOptions().p(2).reduction(torch::kNone))''', + input_fn=lambda: torch.randn(5, 10).clamp_(1e-2, 1 - 1e-2), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiMarginLoss'](i, t.data.type_as(i).long(), p=2, reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def multimarginloss_margin_no_reduce_test(): + t = torch.rand(5).mul(8).floor().long() + return dict( + fullname='MultiMarginLoss_margin_no_reduce', + constructor=wrap_functional( + lambda i: F.multi_margin_loss(i, t.type_as(i).long(), margin=0.5, reduction='none')), + cpp_function_call='''F::multi_margin_loss( + i, t.to(i.options()).to(torch::kLong), + F::MultiMarginLossFuncOptions().margin(0.5).reduction(torch::kNone))''', + input_fn=lambda: torch.randn(5, 10), + cpp_var_map={'i': '_get_input()', 't': t}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiMarginLoss'](i, t.data.type_as(i).long(), + margin=0.5, reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def multimarginloss_weights_no_reduce_test(): + t = torch.rand(5).mul(8).floor().long() + weights = torch.rand(10, dtype=torch.double) + return dict( + fullname='MultiMarginLoss_weights_no_reduce', + constructor=wrap_functional( + lambda i: F.multi_margin_loss(i, t.type_as(i).long(), weight=weights.type_as(i), + reduction='none')), + cpp_function_call='''F::multi_margin_loss( + i, t.to(i.options()).to(torch::kLong), + F::MultiMarginLossFuncOptions().weight(weights.to(i.options())).reduction(torch::kNone))''', + input_fn=lambda: torch.randn(5, 10), + cpp_var_map={'i': '_get_input()', 't': t, 'weights': weights}, + reference_fn=lambda i, *_: + loss_reference_fns['MultiMarginLoss'](i, t.data.type_as(i).long(), + weight=weights, reduction='none'), + check_sum_reduction=True, + check_gradgrad=False, + pickle=False, + default_dtype=torch.double) + + +def single_batch_reference_fn(input, parameters, module): + """Reference function for modules supporting no batch dimensions. + + The module is passed the input and target in batched form with a single item. + The output is squeezed to compare with the no-batch input. + """ + def unsqueeze_inp(inp): + if isinstance(inp, (list, tuple)): + return [t.unsqueeze(0) for t in inp] + return inp.unsqueeze(0) + + single_batch_input = unsqueeze_inp(input) + single_batch_input = [single_batch_input] if isinstance(single_batch_input, torch.Tensor) else single_batch_input + with freeze_rng_state(): + return module(*single_batch_input).squeeze(0) + + +def get_new_module_tests(): + common_utils.set_rng_seed() + new_module_tests = [ + poissonnllloss_no_reduce_test(), + bceloss_no_reduce_test(), + bceloss_weights_no_reduce_test(), + bce_with_logistic_legacy_enum_test(), + bce_with_logistic_no_reduce_test(), + bceloss_no_reduce_scalar_test(), + bceloss_weights_no_reduce_scalar_test(), + bce_with_logistic_no_reduce_scalar_test(), + kldivloss_with_target_no_reduce_test(), + kldivloss_no_reduce_test(), + kldivloss_no_reduce_scalar_test(), + kldivloss_with_log_target_no_reduce_test(), + kldivloss_no_reduce_log_target_test(), + kldivloss_no_reduce_scalar_log_target_test(), + l1loss_no_reduce_test(), + l1loss_no_reduce_complex_test(), + l1loss_no_reduce_scalar_test(), + mseloss_no_reduce_test(), + mseloss_no_reduce_scalar_test(), + nllloss_no_reduce_test(), + nllloss_no_reduce_ignore_index_test(), + nllloss_no_reduce_weights_test(), + nllloss_no_reduce_weights_ignore_index_test(), + nllloss_no_reduce_weights_ignore_index_neg_test(), + nllloss2d_no_reduce_test(), + nllloss2d_no_reduce_weights_test(), + nllloss2d_no_reduce_ignore_index_test(), + nlllossNd_no_reduce_test(), + nlllossNd_no_reduce_weights_test(), + nlllossNd_no_reduce_ignore_index_test(), + smoothl1loss_no_reduce_test(), + smoothl1loss_no_reduce_scalar_test(), + smoothl1loss_beta_test(), + smoothl1loss_zero_beta_test(), + huberloss_delta_test(), + multilabelmarginloss_0d_no_reduce_test(), + multilabelmarginloss_1d_no_reduce_test(), + multilabelmarginloss_index_neg_test(), + multilabelmarginloss_no_reduce_test(), + hingeembeddingloss_no_reduce_test(), + hingeembeddingloss_margin_no_reduce_test(), + softmarginloss_no_reduce_test(), + multilabelsoftmarginloss_no_reduce_test(), + multilabelsoftmarginloss_weights_no_reduce_test(), + multimarginloss_no_reduce_test(), + multimarginloss_1d_no_reduce_test(), + multimarginloss_1d_input_0d_target_no_reduce_test(), + multimarginloss_p_no_reduce_test(), + multimarginloss_margin_no_reduce_test(), + multimarginloss_weights_no_reduce_test(), + dict( + module_name='Conv1d', + constructor_args=(4, 5, 3), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 3)', + input_size=(2, 4, 10), + cudnn=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv1d', + constructor_args=(4, 5, 3, 2), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 3).stride(2)', + input_size=(2, 4, 10), + cudnn=True, + desc='stride', + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv1d', + constructor_args=(4, 5, 3, 1, 1), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 3).stride(1).padding(1)', + input_size=(2, 4, 10), + cudnn=True, + desc='pad1', + with_tf32=True, + tf32_precision=0.01, + default_dtype=torch.double, + ), + dict( + module_name='Conv1d', + constructor_args=(4, 5, 5, 1, 2), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 5).stride(1).padding(2)', + input_size=(2, 4, 10), + cudnn=True, + desc='pad2', + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv1d', + constructor_args=(4, 4, 3, 1, 1), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 4, 3).stride(1).padding(1)', + input_size=(1, 4, 1), + cudnn=True, + desc='pad1size1', + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv1d', + constructor_args=(4, 4, 5, 1, 2), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 4, 5).stride(1).padding(2)', + input_size=(1, 4, 1), + cudnn=True, + desc='pad2size1', + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv1d', + constructor_args=(4, 5, 3), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 3)', + input_size=(0, 4, 10), + cudnn=True, + desc='zero_batch', + with_tf32=True, + tf32_precision=0.005, + ), + dict( + fullname='Conv1d_dilated', + constructor=lambda: nn.Conv1d(4, 5, kernel_size=3, dilation=2), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 3).dilation(2)', + input_size=(2, 4, 10), + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv1d_groups', + constructor=lambda: nn.Conv1d(4, 6, kernel_size=3, groups=2), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 6, 3).groups(2)', + input_size=(2, 4, 6), + cudnn=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv1d_pad_valid', + constructor=lambda: nn.Conv1d(4, 5, 3, padding="valid"), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 3).padding(torch::kValid)', + input_size=(2, 4, 10), + cudnn=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv1d_pad_same', + constructor=lambda: nn.Conv1d(4, 5, 3, padding="same"), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 3).padding(torch::kSame)', + input_size=(2, 4, 10), + cudnn=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv1d_pad_same2', + constructor=lambda: nn.Conv1d(4, 5, 4, padding="same"), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 4).padding(torch::kSame)', + input_size=(2, 4, 10), + cudnn=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv1d_pad_same_dilated', + constructor=lambda: nn.Conv1d(4, 5, 4, padding="same", dilation=2), + cpp_constructor_args='torch::nn::Conv1dOptions(4, 5, 3).padding(torch::kSame).dilation(2)', + input_size=(2, 4, 10), + cudnn=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='ConvTranspose1d', + constructor=lambda: nn.ConvTranspose1d(3, 4, kernel_size=3, stride=(3,), padding=1, output_padding=(1,)), + cpp_constructor_args='torch::nn::ConvTranspose1dOptions(3, 4, 3).stride(3).padding(1).output_padding(1)', + cudnn=True, + input_size=(1, 3, 7), + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='ConvTranspose1d', + constructor_args=(3, 4, 3, 2, 1, 1, 1, False), + cpp_constructor_args='''torch::nn::ConvTranspose1dOptions(3, 4, 3) + .stride(2).padding(1).output_padding(1).groups(1).bias(false)''', + input_size=(1, 3, 6), + cudnn=True, + desc='no_bias', + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='ConvTranspose1d', + constructor_args=(3, 4, 3, 2, 1, 1, 1, True, 2), + cpp_constructor_args='''torch::nn::ConvTranspose1dOptions(3, 4, 3) + .stride(2).padding(1).output_padding(1).groups(1).bias(true).dilation(2)''', + input_size=(1, 3, 6), + cudnn=True, + desc='dilated', + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='ConvTranspose1d_groups', + constructor=lambda: nn.ConvTranspose1d(4, 6, 3, stride=(3,), padding=1, output_padding=(1,), groups=2), + cpp_constructor_args='''torch::nn::ConvTranspose1dOptions(4, 6, 3) + .stride(3).padding(1).output_padding(1).groups(2)''', + cudnn=True, + input_size=(2, 4, 7), + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv2d', + constructor_args=(3, 4, (3, 2)), + cpp_constructor_args='torch::nn::Conv2dOptions(3, 4, {3, 2})', + input_size=(2, 3, 7, 5), + cudnn=True, + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv2d', + constructor_args=(3, 4, (3, 3), (2, 2)), + cpp_constructor_args='torch::nn::Conv2dOptions(3, 4, {3, 3}).stride({2, 2})', + input_size=(2, 3, 6, 6), + cudnn=True, + desc='strided', + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv2d', + constructor_args=(3, 4, (3, 3), (2, 2), (1, 1)), + cpp_constructor_args='torch::nn::Conv2dOptions(3, 4, {3, 3}).stride({2, 2}).padding({1, 1})', + input_size=(2, 3, 6, 6), + cudnn=True, + desc='padding', + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv2d', + constructor_args=(3, 2, (3, 3), (2, 2), (1, 1), (2, 2)), + cpp_constructor_args='torch::nn::Conv2dOptions(3, 2, {3, 3}).stride({2, 2}).padding({1, 1}).dilation({2, 2})', + input_size=(2, 3, 8, 8), + cudnn=True, + desc='dilated', + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv2d', + constructor_args=(3, 4, (3, 2), 1, 0, 1, 1, False), + cpp_constructor_args='''torch::nn::Conv2dOptions(3, 4, {3, 2}) + .stride(1).padding(0).dilation(1).groups(1).bias(false)''', + input_size=(2, 3, 6, 5), + cudnn=True, + desc='no_bias', + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.015, + default_dtype=torch.double, + ), + dict( + module_name='Conv2d', + constructor_args=(3, 4, (3, 2)), + cpp_constructor_args='torch::nn::Conv2dOptions(3, 4, {3, 2})', + input_size=(0, 3, 7, 5), + cudnn=True, + desc='zero_batch', + check_with_long_tensor=True, + with_tf32=True, + ), + dict( + fullname='Conv2d_groups', + constructor=lambda: nn.Conv2d(4, 6, (3, 2), groups=2), + cpp_constructor_args='torch::nn::Conv2dOptions(4, 6, {3, 2}).groups(2)', + input_size=(2, 4, 6, 5), + cudnn=True, + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.015, + default_dtype=torch.double, + ), + dict( + fullname='Conv2d_groups_thnn', + constructor=lambda: nn.Conv2d(4, 6, (3, 2), groups=2), + cpp_constructor_args='torch::nn::Conv2dOptions(4, 6, {3, 2}).groups(2)', + input_size=(2, 4, 6, 5), + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.015, + default_dtype=torch.double, + ), + dict( + fullname='Conv2d_pad_valid', + constructor=lambda: nn.Conv2d(2, 4, (3, 4), padding="valid"), + cpp_constructor_args='torch::nn::Conv2dOptions(2, 4, {3, 4}).padding(torch::kValid)', + input_size=(2, 2, 6, 5), + cudnn=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv2d_pad_same', + constructor=lambda: nn.Conv2d(2, 4, (3, 4), padding="same"), + cpp_constructor_args='torch::nn::Conv2dOptions(2, 4, {3, 4}).padding(torch::kSame)', + input_size=(2, 2, 6, 5), + cudnn=True, + with_tf32=True, + tf32_precision=0.01, + default_dtype=torch.double, + ), + dict( + fullname='Conv2d_pad_same_dilated', + constructor=lambda: nn.Conv2d(2, 4, (3, 4), padding="same", dilation=2), + cpp_constructor_args='torch::nn::Conv2dOptions(2, 4, {3, 4}).padding(torch::kSame).dilation(2)', + input_size=(2, 2, 6, 5), + cudnn=True, + with_tf32=True, + tf32_precision=0.01, + default_dtype=torch.double, + ), + dict( + module_name='ConvTranspose2d', + constructor_args=(3, 4, 3, (3, 2), 1, (1, 1)), + cpp_constructor_args='''torch::nn::ConvTranspose2dOptions(3, 4, 3) + .stride({3, 2}).padding(1).output_padding({1, 1})''', + cudnn=True, + input_size=(1, 3, 7, 6), + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.01, + default_dtype=torch.double, + ), + dict( + module_name='ConvTranspose2d', + constructor_args=(3, 4, 3, (2, 3), 1, (1, 1), 1, False, (2, 2)), + cpp_constructor_args='''torch::nn::ConvTranspose2dOptions(3, 4, 3) + .stride({2, 3}) + .padding(1) + .output_padding({1, 1}) + .groups(1) + .bias(false) + .dilation({2, 2})''', + input_size=(1, 3, 6, 7), + cudnn=True, + desc='dilated', + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.01, + default_dtype=torch.double, + ), + dict( + module_name='ConvTranspose2d', + constructor_args=(3, 4, 3, (2, 3), 1, (1, 1), 1, False), + cpp_constructor_args='''torch::nn::ConvTranspose2dOptions(3, 4, 3) + .stride({2, 3}).padding(1).output_padding({1, 1}).groups(1).bias(false)''', + input_size=(1, 3, 6, 7), + cudnn=True, + desc='no_bias', + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.01, + default_dtype=torch.double, + ), + dict( + fullname='ConvTranspose2d_groups', + constructor=lambda: nn.ConvTranspose2d(2, 4, (2, 3), groups=2), + cpp_constructor_args='torch::nn::ConvTranspose2dOptions(2, 4, {2, 3}).groups(2)', + input_size=(1, 2, 4, 5), + cudnn=True, + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.01, + default_dtype=torch.double, + ), + dict( + fullname='Conv2d_depthwise', + constructor=lambda: nn.Conv2d(4, 4, (3, 3), groups=4), + cpp_constructor_args='torch::nn::Conv2dOptions(4, 4, {3, 3}).groups(4)', + input_size=(2, 4, 6, 6), + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv2d_depthwise_with_multiplier', + constructor=lambda: nn.Conv2d(4, 8, (3, 3), groups=4), + cpp_constructor_args='torch::nn::Conv2dOptions(4, 8, {3, 3}).groups(4)', + input_size=(2, 4, 6, 6), + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv2d_depthwise_strided', + constructor=lambda: nn.Conv2d(4, 4, (3, 3), stride=(2, 2), groups=4), + cpp_constructor_args='torch::nn::Conv2dOptions(4, 4, {3, 3}).stride({2, 2}).groups(4)', + input_size=(2, 4, 6, 6), + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv2d_depthwise_padded', + constructor=lambda: nn.Conv2d(4, 4, (3, 3), padding=(1, 1), groups=4), + cpp_constructor_args='torch::nn::Conv2dOptions(4, 4, {3, 3}).padding({1, 1}).groups(4)', + input_size=(2, 4, 6, 6), + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv2d_depthwise_dilated', + constructor=lambda: nn.Conv2d(4, 4, (2, 2), dilation=(2, 2), groups=4), + cpp_constructor_args='torch::nn::Conv2dOptions(4, 4, {2, 2}).dilation({2, 2}).groups(4)', + input_size=(2, 4, 5, 5), + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Conv3d', + constructor_args=(2, 3, (2, 3, 2)), + cpp_constructor_args='torch::nn::Conv3dOptions(2, 3, {2, 3, 2})', + input_size=(1, 2, 4, 5, 4), + cudnn=True, + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='Conv3d', + constructor_args=(2, 3, (2, 3, 4), 1, 0, 1, 1, False), + cpp_constructor_args='''torch::nn::Conv3dOptions(2, 3, {2, 3, 4}) + .stride(1).padding(0).dilation(1).groups(1).bias(false)''', + input_size=(1, 2, 3, 4, 5), + cudnn=True, + desc='no_bias', + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='Conv3d', + constructor_args=(2, 3, (1, 1, 1), 1, 0, 1, 1, False), + cpp_constructor_args='''torch::nn::Conv3dOptions(2, 3, {2, 3, 4}) + .stride(1).padding(0).dilation(1).groups(1).bias(false)''', + input_size=(1, 2, 3, 4, 5), + cudnn=True, + desc='1x1x1_no_bias', + check_with_long_tensor=False, + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='Conv3d', + constructor_args=(3, 4, 2, 2), + cpp_constructor_args='torch::nn::Conv3dOptions(3, 4, 2).stride(2)', + input_size=(2, 3, 5, 5, 5), + cudnn=True, + desc='stride', + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='Conv3d', + constructor_args=(3, 4, 2, 2, 1), + cpp_constructor_args='torch::nn::Conv3dOptions(3, 4, 2).stride(2).padding(1)', + input_size=(2, 3, 5, 5, 5), + cudnn=True, + desc='stride_padding', + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='Conv3d', + constructor_args=(3, 4, (2, 3, 4)), + cpp_constructor_args='torch::nn::Conv3dOptions(3, 4, {2, 3, 4})', + input_size=(0, 3, 3, 4, 5), + cudnn=True, + check_with_long_tensor=True, + desc='zero_batch', + with_tf32=True, + ), + dict( + fullname='Conv3d_groups', + constructor=lambda: nn.Conv3d(2, 4, kernel_size=3, groups=2), + cpp_constructor_args='torch::nn::Conv3dOptions(2, 4, 3).groups(2)', + input_size=(1, 2, 4, 5, 4), + cudnn=True, + check_with_long_tensor=True, + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + fullname='Conv3d_dilated', + constructor=lambda: nn.Conv3d(3, 4, kernel_size=2, dilation=2), + cpp_constructor_args='torch::nn::Conv3dOptions(3, 4, 2).dilation(2)', + input_size=(2, 3, 5, 5, 5), + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + fullname='Conv3d_dilated_strided', + constructor=lambda: nn.Conv3d(3, 4, kernel_size=2, dilation=2, stride=2), + cpp_constructor_args='torch::nn::Conv3dOptions(3, 4, 2).dilation(2).stride(2)', + input_size=(2, 3, 5, 5, 5), + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + fullname='Conv3d_pad_valid', + constructor=lambda: nn.Conv3d(3, 4, (2, 3, 4), padding="valid"), + cpp_constructor_args='torch::nn::Conv3dOptions(3, 4, {2, 3, 4}).padding(torch::kValid)', + input_size=(2, 3, 6, 5, 4), + cudnn=True, + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + fullname='Conv3d_pad_same', + constructor=lambda: nn.Conv3d(3, 4, (2, 3, 4), padding="same"), + cpp_constructor_args='torch::nn::Conv3dOptions(3, 4, {2, 3, 4}).padding(torch::kSame)', + input_size=(2, 3, 6, 5, 4), + cudnn=True, + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + fullname='Conv3d_pad_same_dilated', + constructor=lambda: nn.Conv3d(3, 4, (2, 3, 4), padding="same", dilation=2), + cpp_constructor_args='torch::nn::Conv3dOptions(3, 4, {2, 3, 4}).padding(torch::kSame).dilation(2)', + input_size=(2, 3, 6, 5, 4), + cudnn=True, + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='ConvTranspose3d', + constructor_args=(2, 3, (2, 3, 2)), + cpp_constructor_args='torch::nn::ConvTranspose3dOptions(2, 3, {2, 3, 2})', + cudnn=True, + input_size=(1, 2, 4, 5, 4), + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='ConvTranspose3d', + constructor_args=(2, 3, (2, 3, 2), 1, 0, 0, 1, True, (2, 2, 2)), + cpp_constructor_args='''torch::nn::ConvTranspose3dOptions(2, 3, {2, 3, 2}) + .stride(1).padding(0).output_padding(0).groups(1).bias(true).dilation({2, 2, 2})''', + cudnn=True, + input_size=(1, 2, 4, 5, 4), + desc='dilated', + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='ReplicationPad3d', + constructor_args=((1, 2, 3, 3, 2, 1),), + cpp_constructor_args='torch::nn::ReplicationPad3dOptions({1, 2, 3, 3, 2, 1})', + input_size=(2, 3, 2, 2, 2), + default_dtype=torch.double, + ), + dict( + module_name='ReplicationPad3d', + constructor_args=((1, 2, 3, 3, 2, 1),), + cpp_constructor_args='torch::nn::ReplicationPad3dOptions({1, 2, 3, 3, 2, 1})', + input_size=(3, 2, 2, 2), + reference_fn=single_batch_reference_fn, + desc='no_batch_dim', + default_dtype=torch.double, + ), + dict( + module_name='ReplicationPad3d', + constructor_args=((1, 2, 3, 3, 2, 1),), + cpp_constructor_args='torch::nn::ReplicationPad3dOptions({1, 2, 3, 3, 2, 1})', + input_fn=lambda: torch.rand(2, 3, 2, 2, 2, dtype=torch.complex128, requires_grad=True), + skip_half=True, + desc='complex' + ), + dict( + module_name='Embedding', + constructor_args=(4, 3), + cpp_constructor_args='torch::nn::EmbeddingOptions(4, 3)', + input_fn=lambda: torch.empty(2, 3, dtype=torch.long).random_(4), + check_gradgrad=False, + default_dtype=torch.double, + decorator=skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/117971") + ), + dict( + module_name='Embedding', + constructor_args=(4, 3), + cpp_constructor_args='torch::nn::EmbeddingOptions(4, 3)', + input_fn=lambda: torch.empty(1, 512, dtype=torch.long).random_(4).expand(7, 512), + check_gradgrad=False, + desc='discontiguous', + default_dtype=torch.double, + decorator=skipIfTorchDynamo("https://github.com/pytorch/pytorch/issues/117971") + ), + dict( + module_name='EmbeddingBag', + constructor_args=(4, 3), + cpp_constructor_args='torch::nn::EmbeddingBagOptions(4, 3)', + input_fn=lambda: torch.empty(2, 3, dtype=torch.long).random_(4), + check_gradgrad=False, + desc='mean', + default_dtype=torch.double, + ), + dict( + module_name='EmbeddingBag', + constructor_args=(4, 3), + cpp_constructor_args='torch::nn::EmbeddingBagOptions(4, 3)', + input_fn=lambda: torch.empty(1, 512, dtype=torch.long).random_(4).expand(7, 512), + check_gradgrad=False, + desc='discontiguous', + default_dtype=torch.double, + ), + dict( + module_name='EmbeddingBag', + constructor_args=(4, 3, None, 2., False, 'sum'), + cpp_constructor_args='''torch::nn::EmbeddingBagOptions(4, 3) + .max_norm(std::nullopt).norm_type(2.).scale_grad_by_freq(false).mode(torch::kSum)''', + input_fn=lambda: torch.empty(2, 3, dtype=torch.long).random_(4), + check_gradgrad=False, + desc='sum', + default_dtype=torch.double, + ), + dict( + module_name='EmbeddingBag', + constructor_args=(4, 3, None, 2., False, 'max'), + cpp_constructor_args='''torch::nn::EmbeddingBagOptions(4, 3) + .max_norm(std::nullopt).norm_type(2.).scale_grad_by_freq(false).mode(torch::kMax)''', + input_fn=lambda: torch.empty(2, 3, dtype=torch.long).random_(4), + check_gradgrad=False, + desc='max', + default_dtype=torch.double, + ), + dict( + fullname='EmbeddingBag_mean_padding_idx', + constructor=lambda: nn.EmbeddingBag(4, 3, padding_idx=1), + cpp_constructor_args='torch::nn::EmbeddingBagOptions(4, 3).padding_idx(1)', + input_fn=lambda: torch.stack([torch.randperm(3), torch.randperm(3)]), + check_gradgrad=False, + default_dtype=torch.double, + ), + dict( + fullname='EmbeddingBag_sum_padding_idx', + constructor=lambda: nn.EmbeddingBag(4, 3, None, 2., False, 'sum', padding_idx=1), + cpp_constructor_args='''torch::nn::EmbeddingBagOptions(4, 3) + .max_norm(std::nullopt).norm_type(2.).scale_grad_by_freq(false).mode(torch::kSum).padding_idx(1)''', + input_fn=lambda: torch.stack([torch.randperm(3), torch.randperm(3)]), + check_gradgrad=False, + default_dtype=torch.double, + ), + dict( + fullname='EmbeddingBag_max_padding_idx', + constructor=lambda: nn.EmbeddingBag(4, 3, None, 2., False, 'max', padding_idx=1), + cpp_constructor_args='''torch::nn::EmbeddingBagOptions(4, 3) + .max_norm(std::nullopt).norm_type(2.).scale_grad_by_freq(false).mode(torch::kMax).padding_idx(1)''', + input_fn=lambda: torch.stack([torch.randperm(3), torch.randperm(3)]), + check_gradgrad=False, + default_dtype=torch.double, + ), + dict( + fullname='EmbeddingBag_sparse', + constructor=lambda: nn.EmbeddingBag(4, 3, sparse=True, dtype=torch.double), + cpp_constructor_args='''torch::nn::EmbeddingBagOptions(4, 3) + .sparse(true)._weight(torch::rand({4, 3}).to(torch::kFloat64))''', + input_fn=lambda: torch.randperm(2).repeat(1, 2), + check_gradgrad=False, + has_sparse_gradients=True, + ), + dict( + constructor=lambda: nn.Embedding(4, 3, dtype=torch.double, sparse=True), + cpp_constructor_args='torch::nn::EmbeddingOptions(4, 3).sparse(true)._weight(torch::rand({4, 3}).to(torch::kFloat64))', + input_fn=lambda: torch.randperm(2).repeat(1, 2), + fullname='Embedding_sparse', + check_gradgrad=False, + has_sparse_gradients=True, + ), + dict( + module_name='PixelShuffle', + constructor_args=(3,), + cpp_constructor_args='torch::nn::PixelShuffleOptions(3)', + input_size=(1, 9, 4, 4), + default_dtype=torch.double, + ), + dict( + module_name='PixelUnshuffle', + constructor_args=(3,), + cpp_constructor_args='torch::nn::PixelUnshuffleOptions(3)', + input_size=(1, 1, 12, 12), + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12})).scale_factor(std::nullopt).mode(torch::kNearest)''', + input_size=(1, 2, 4), + fullname='interpolate_nearest_1d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12})).scale_factor(std::nullopt).mode(torch::kNearest)''', + input_size=(0, 2, 4), + fullname='interpolate_nearest_1d_zero_dim', + pickle=False, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(12, ), scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12})).scale_factor(std::nullopt).mode(torch::kNearest)''', + input_size=(1, 2, 3), + fullname='interpolate_nearest_tuple_1d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=4., mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt).scale_factor(std::vector({4.})).mode(torch::kNearest)''', + input_size=(1, 2, 4), + fullname='interpolate_nearest_scale_1d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='linear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12})) + .scale_factor(std::nullopt) + .mode(torch::kLinear) + .align_corners(false)''', + input_size=(1, 2, 4), + fullname='interpolate_linear_1d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(4, ), scale_factor=None, mode='linear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({4})) + .scale_factor(std::nullopt) + .mode(torch::kLinear) + .align_corners(false)''', + input_size=(1, 2, 3), + fullname='interpolate_linear_tuple_1d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=4., mode='linear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({4.})) + .mode(torch::kLinear) + .align_corners(false)''', + input_size=(1, 2, 4), + fullname='interpolate_linear_scale_1d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='linear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12})) + .scale_factor(std::nullopt) + .mode(torch::kLinear) + .align_corners(false)''', + input_size=(0, 2, 4), + fullname='interpolate_linear_1d_zero_dim', + pickle=False, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='linear', align_corners=True), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12})) + .scale_factor(std::nullopt) + .mode(torch::kLinear) + .align_corners(true)''', + input_size=(1, 2, 4), + fullname='interpolate_linear_1d_align_corners', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=4., mode='linear', align_corners=True), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({4.})) + .mode(torch::kLinear) + .align_corners(true)''', + input_size=(1, 2, 4), + fullname='interpolate_linear_scale_1d_align_corners', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=2, scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({2, 2})) + .scale_factor(std::nullopt) + .mode(torch::kNearest)''', + input_size=(1, 128, 1, 1), + fullname='interpolate_nearest_2d_launch_configs', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kNearest)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_nearest_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(12, 16), scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 16})) + .scale_factor(std::nullopt) + .mode(torch::kNearest)''', + input_size=(1, 2, 3, 4), + fullname='interpolate_nearest_tuple_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=4., mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({4., 4.})) + .mode(torch::kNearest)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_nearest_scale_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kNearest)''', + input_size=(0, 2, 4, 4), + fullname='interpolate_nearest_2d_zero_dim', + pickle=False, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='bilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kBilinear) + .align_corners(false)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bilinear_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='bilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kBilinear) + .align_corners(false)''', + input_size=(0, 2, 4, 4), + fullname='interpolate_bilinear_2d_zero_dim', + pickle=False, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(4, 6), scale_factor=None, + mode='bilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({4, 6})) + .scale_factor(std::nullopt) + .mode(torch::kBilinear) + .align_corners(false)''', + input_size=(1, 2, 2, 3), + fullname='interpolate_bilinear_tuple_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=4., + mode='bilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({4., 4.})) + .mode(torch::kBilinear) + .align_corners(false)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bilinear_scale_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=(2., 2.), + mode='bilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({2., 2.})) + .mode(torch::kBilinear) + .align_corners(false)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bilinear_scale_tuple_shared_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=(2., 1.), + mode='bilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({2., 1.})) + .mode(torch::kBilinear) + .align_corners(false)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bilinear_scale_tuple_skewed_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(4, 6), scale_factor=None, mode='bilinear', align_corners=True), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({4, 6})) + .scale_factor(std::nullopt) + .mode(torch::kBilinear) + .align_corners(true)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bilinear_tuple_2d_align_corners', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=(2., 1.), + mode='bilinear', align_corners=True), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({2., 1.})) + .mode(torch::kBilinear) + .align_corners(true)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bilinear_scale_tuple_skewed_2d_align_corners', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='bicubic', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kBicubic) + .align_corners(false)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bicubic_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='bicubic', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kBicubic) + .align_corners(false)''', + input_size=(0, 2, 4, 4), + fullname='interpolate_bicubic_2d_zero_dim', + pickle=False, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(4, 6), scale_factor=None, + mode='bicubic', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({4, 6})) + .scale_factor(std::nullopt) + .mode(torch::kBicubic) + .align_corners(false)''', + input_size=(1, 2, 2, 3), + fullname='interpolate_bicubic_tuple_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=4., mode='bicubic', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({4., 4.})) + .mode(torch::kBicubic) + .align_corners(false)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bicubic_scale_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=(2., 2.), + mode='bicubic', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({2., 2.})) + .mode(torch::kBicubic) + .align_corners(false)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bicubic_scale_tuple_shared_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=(2., 1.), + mode='bicubic', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({2., 1.})) + .mode(torch::kBicubic) + .align_corners(false)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bicubic_scale_tuple_skewed_2d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(4, 6), scale_factor=None, mode='bicubic', align_corners=True), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({4, 6})) + .scale_factor(std::nullopt) + .mode(torch::kBicubic) + .align_corners(true)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bicubic_tuple_2d_align_corners', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=(2., 1.), + mode='bicubic', align_corners=True), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({2., 1.})) + .mode(torch::kBicubic) + .align_corners(true)''', + input_size=(1, 2, 4, 4), + fullname='interpolate_bicubic_scale_tuple_skewed_2d_align_corners', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kNearest)''', + input_size=(1, 2, 4, 4, 4), + fullname='interpolate_nearest_3d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kNearest)''', + input_size=(0, 2, 4, 4, 4), + fullname='interpolate_nearest_3d_zero_dim', + pickle=False, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(12, 16, 16), scale_factor=None, mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 16, 16})) + .scale_factor(std::nullopt) + .mode(torch::kNearest)''', + input_size=(1, 2, 3, 4, 4), + fullname='interpolate_nearest_tuple_3d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=4., mode='nearest'), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({4., 4., 4.})) + .mode(torch::kNearest)''', + input_size=(1, 2, 4, 4, 4), + fullname='interpolate_nearest_scale_3d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='trilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kTrilinear) + .align_corners(false)''', + input_size=(1, 2, 4, 4, 4), + fullname='interpolate_trilinear_3d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=12, scale_factor=None, mode='trilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({12, 12, 12})) + .scale_factor(std::nullopt) + .mode(torch::kTrilinear) + .align_corners(false)''', + input_size=(0, 2, 4, 4, 4), + fullname='interpolate_trilinear_3d_zero_dim', + pickle=False, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(4, 6, 6), + scale_factor=None, mode='trilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({4, 6, 6})) + .scale_factor(std::nullopt) + .mode(torch::kTrilinear) + .align_corners(false)''', + input_size=(1, 2, 2, 3, 3), + fullname='interpolate_trilinear_tuple_3d', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=3., mode='trilinear', align_corners=False), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({3., 3., 3.})) + .mode(torch::kTrilinear) + .align_corners(false)''', + input_size=(1, 2, 3, 4, 5), + fullname='interpolate_trilinear_scale_3d', + # See https://github.com/pytorch/pytorch/issues/5006 + precision=3e-4, + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.interpolate, size=(4, 6, 6), scale_factor=None, + mode='trilinear', align_corners=True), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::vector({4, 6, 6})) + .scale_factor(std::nullopt) + .mode(torch::kTrilinear) + .align_corners(true)''', + input_size=(1, 2, 2, 3, 3), + fullname='interpolate_trilinear_tuple_3d_align_corners', + pickle=False, + default_dtype=torch.double + ), + dict( + constructor=wrap_functional(F.interpolate, size=None, scale_factor=3., mode='trilinear', align_corners=True), + cpp_options_args='''F::InterpolateFuncOptions() + .size(std::nullopt) + .scale_factor(std::vector({3., 3., 3.})) + .mode(torch::kTrilinear) + .align_corners(true)''', + input_size=(1, 2, 3, 4, 4), + fullname='interpolate_trilinear_scale_3d_align_corners', + # See https://github.com/pytorch/pytorch/issues/5006 + precision=3e-4, + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.softmax, dim=-1), + cpp_options_args='F::SoftmaxFuncOptions(-1)', + input_size=(2, 128), # trigger the last-dim algo in CUDA + fullname='softmax_lastdim', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.softmax, dim=1, dtype=torch.float64), + cpp_options_args='F::SoftmaxFuncOptions(1).dtype(torch::kFloat64)', + input_size=(2, 128), + fullname='softmax_lastdim_dtype', + pickle=False, + test_cuda=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.softmax, dim=1), + cpp_options_args='F::SoftmaxFuncOptions(1)', + input_size=(2, 128, 2, 2), # trigger special case of spatial CUDA algo + fullname='softmax_spatial_special', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.softmax, dim=1), + cpp_options_args='F::SoftmaxFuncOptions(1)', + input_size=(2, 2, 4, 4), # regular spatial algorithm + fullname='softmax_spatial', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.softmax, dim=1, dtype=torch.float64), + cpp_options_args='F::SoftmaxFuncOptions(1).dtype(torch::kFloat64)', + input_size=(2, 2, 4, 4), # regular spatial algorithm + fullname='softmax_spatial_dtype', + pickle=False, + test_cuda=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.softmax, dim=0), + cpp_options_args='F::SoftmaxFuncOptions(0)', + input_size=(2, 3, 4, 5), + fullname='softmax_functional_dim0', + test_cuda=False, + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.softmax, dim=3), + cpp_options_args='F::SoftmaxFuncOptions(3)', + input_size=(2, 3, 4, 5), + fullname='softmax_functional_dim3', + test_cuda=False, + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.softmax, dim=-1), + cpp_options_args='F::SoftmaxFuncOptions(-1)', + input_size=(), + fullname='softmax_functional_scalar', + test_cuda=False, + pickle=False, + ), + dict( + constructor=wrap_functional(F.log_softmax, dim=-1), + cpp_options_args='F::LogSoftmaxFuncOptions(-1)', + input_size=(2, 128), # trigger the last-dim algo in CUDA + fullname='log_softmax_lastdim', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.log_softmax, dim=1), + cpp_options_args='F::LogSoftmaxFuncOptions(1)', + input_size=(2, 128, 2, 2), # trigger special case of spatial CUDA algo + fullname='log_softmax_spatial_special', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.log_softmax, dim=1), + cpp_options_args='F::LogSoftmaxFuncOptions(1)', + input_size=(2, 2, 4, 4), # regular spatial algorithm + fullname='log_softmax_spatial', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.log_softmax, dim=0), + cpp_options_args='F::LogSoftmaxFuncOptions(0)', + input_size=(2, 3, 4, 5), + fullname='log_softmax_dim0', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.log_softmax, dim=3), + cpp_options_args='F::LogSoftmaxFuncOptions(3)', + input_size=(2, 3, 4, 5), + fullname='log_softmax_dim3', + pickle=False, + default_dtype=torch.double, + ), + dict( + constructor=wrap_functional(F.log_softmax, dim=0), + cpp_options_args='F::LogSoftmaxFuncOptions(0)', + input_size=(), + fullname='log_softmax_scalar', + pickle=False, + ), + dict( + fullname='Unfold', + constructor=lambda: nn.Unfold((2, 2), (1, 1), (0, 0), (1, 1)), + cpp_constructor_args='torch::nn::UnfoldOptions({2, 2}).dilation({1, 1}).padding({0, 0}).stride({1, 1})', + input_size=(2, 4, 3, 3), + check_gradgrad=False, + test_cuda=True, + default_dtype=torch.double, + ), + dict( + fullname='Fold', + constructor=lambda: nn.Fold((3, 3), (2, 2), (1, 1), (0, 0), (1, 1)), + cpp_constructor_args='torch::nn::FoldOptions({3, 3}, {2, 2}).dilation({1, 1}).padding({0, 0}).stride({1, 1})', + input_size=(2, 16, 4), + check_gradgrad=False, + test_cuda=True, + default_dtype=torch.double, + ), + dict( + fullname='Fold_no_batch_dim_input', + constructor=lambda: nn.Fold((3, 3), (2, 2), (1, 1), (0, 0), (1, 1)), + cpp_constructor_args='torch::nn::FoldOptions({3, 3}, {2, 2}).dilation({1, 1}).padding({0, 0}).stride({1, 1})', + input_size=(16, 4), + check_gradgrad=False, + ref=single_batch_reference_fn, + test_cuda=True, + default_dtype=torch.double, + ), + dict( + fullname='Unfold_int_input', + constructor=lambda: nn.Unfold(2, 1, 0, 1), + cpp_constructor_args='torch::nn::UnfoldOptions(2).dilation(1).padding(0).stride(1)', + input_size=(2, 4, 3, 3), + check_gradgrad=False, + test_cuda=True, + default_dtype=torch.double, + ), + dict( + fullname='Fold_int_input', + constructor=lambda: nn.Fold(3, 2, 1, 0, 1), + cpp_constructor_args='torch::nn::FoldOptions(3, 2).dilation(1).padding(0).stride(1)', + input_size=(2, 16, 4), + check_gradgrad=False, + test_cuda=True, + default_dtype=torch.double, + ), + dict( + fullname='Fold_no_batch_dim_int_input', + constructor=lambda: nn.Fold(3, 2, 1, 0, 1), + cpp_constructor_args='torch::nn::FoldOptions(3, 2).dilation(1).padding(0).stride(1)', + input_size=(16, 4), + ref=single_batch_reference_fn, + check_gradgrad=False, + test_cuda=True, + default_dtype=torch.double, + ), + dict( + module_name='RReLU', + constructor_args=(0.1, 0.9), + cpp_constructor_args='torch::nn::RReLUOptions().lower(0.1).upper(0.9)', + input_size=(), + desc='with_up_down_scalar', + test_cuda=False, + default_dtype=torch.double, + ), + dict( + module_name='PairwiseDistance', + input_fn=lambda: (torch.randn(10, 8), torch.randn(10, 8)), + default_dtype=torch.double, + ), + dict( + module_name='PairwiseDistance', + input_fn=lambda: (torch.randn(10, 1), torch.randn(10, 8)), + desc='broadcast_lhs', + default_dtype=torch.double, + ), + dict( + module_name='PairwiseDistance', + input_fn=lambda: (torch.randn(10, 8), torch.randn(1, 8)), + desc='broadcast_rhs', + default_dtype=torch.double, + ), + dict( + module_name='PairwiseDistance', + constructor_args=(1.5, 1e-05, True), + cpp_constructor_args='torch::nn::PairwiseDistanceOptions().p(1.5).eps(1e-05).keepdim(true)', + input_fn=lambda: (torch.randn(10, 8), torch.randn(10, 8)), + desc='with_non_default_args', + default_dtype=torch.double, + ), + dict( + module_name='PairwiseDistance', + input_fn=lambda: (torch.randn(8), torch.randn(8)), + reference_fn=single_batch_reference_fn, + desc='no_batch_dim', + default_dtype=torch.double, + ), + dict( + module_name='TransformerEncoderLayer', + constructor_args=(4, 2, 16, 0.0), + cpp_constructor_args='''torch::nn::TransformerEncoderLayerOptions(4, 2) + .dim_feedforward(16) + .dropout(0.0)''', + input_size=(2, 3, 4), + desc='relu_activation', + with_tf32=True, + tf32_precision=0.1, + # TODO(#50743): figure out the error + # RuntimeError: The size of tensor a (6) must match the size of tensor b (4) + # at non-singleton dimension 2 + check_batched_grad=False, + check_gradgrad=False, + default_dtype=torch.double, + ), + dict( + module_name='TransformerEncoderLayer', + constructor_args=(4, 2, 8, 0.0, F.gelu), + cpp_constructor_args='''torch::nn::TransformerEncoderLayerOptions(4, 2) + .dim_feedforward(8) + .dropout(0.0) + .activation(torch::kGELU)''', + input_size=(2, 3, 4), + check_gradgrad=False, + desc='gelu_activation', + with_tf32=True, + tf32_precision=0.08 if SM90OrLater else 0.05, + default_dtype=torch.double, + ), + dict( + module_name='TransformerDecoderLayer', + constructor_args=(4, 2, 8, 0.0), + cpp_constructor_args='''torch::nn::TransformerDecoderLayerOptions(4, 2) + .dim_feedforward(8) + .dropout(0.0)''', + input_fn=lambda: (torch.rand(3, 3, 4), torch.rand(2, 3, 4)), + check_gradgrad=False, + desc='relu_activation', + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='TransformerDecoderLayer', + constructor_args=(4, 2, 8, 0.0, F.gelu), + cpp_constructor_args='''torch::nn::TransformerDecoderLayerOptions(4, 2) + .dim_feedforward(8) + .dropout(0.0) + .activation(torch::kGELU)''', + input_fn=lambda: (torch.rand(3, 3, 4), torch.rand(2, 3, 4)), + check_gradgrad=False, + desc='gelu_activation', + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + dict( + module_name='Transformer', + constructor_args=(4, 2, 2, 2, 8, 0.0, F.relu), + cpp_constructor_args='''torch::nn::TransformerOptions() + .d_model(4) + .nhead(2) + .num_encoder_layers(2) + .num_decoder_layers(2) + .dim_feedforward(8) + .dropout(0.0) + .activation(torch::kReLU)''', + input_fn=lambda: (torch.rand(3, 3, 4), torch.rand(2, 3, 4), torch.rand(3, 3)), + check_gradgrad=False, + desc='multilayer_coder', + with_tf32=True, + tf32_precision=0.05 if SM90OrLater else 0.03, + default_dtype=torch.double, + ), + dict( + module_name='Linear', + constructor_args=(3, 5), + cpp_constructor_args='torch::nn::LinearOptions(3, 5)', + input_fn=lambda: torch.rand(3), + reference_fn=lambda i, p, _: torch.mm(i.view(1, -1), p[0].t()).view(-1) + p[1], + desc="no_batch_dim", + with_tf32=True, + tf32_precision=0.005, + default_dtype=torch.double, + ), + dict( + module_name='Flatten', + cpp_constructor_args='torch::nn::FlattenOptions().start_dim(-3).end_dim(-1)', + constructor_args=(-3, -1), + input_size=(3, 4, 5), + reference_fn=single_batch_reference_fn, + desc="no_batch_dim", + default_dtype=torch.double, + ), + dict( + module_name='Unflatten', + cpp_constructor_args='torch::nn::UnflattenOptions(-2, {2, 2})', + constructor_args=(-2, torch.Size([2, 2])), + input_size=(3, 4, 5), + reference_fn=single_batch_reference_fn, + desc="no_batch_dim", + default_dtype=torch.double, + ), + dict( + module_name='LayerNorm', + constructor_args=([56, 56, 56], 1e-5, False), + cpp_constructor_args='torch::nn::LayerNormOptions({56, 56, 56}).eps(1e-5).elementwise_affine(false)', + input_size=(4, 56, 56, 56), + cudnn=True, + check_eval=True, + gradcheck_fast_mode=True, + check_half=True, + desc='3d_no_affine_large_feature', + ), + ] + + # add conv padding mode tests: + for padding_mode, cpp_padding_mode in zip( + ['reflect', 'circular', 'replicate', 'zeros'], + ['torch::kReflect', 'torch::kCircular', 'torch::kReplicate', 'torch::kZeros'], strict=True): + # conv signature: + # in_channels, out_channels, kernel_size, stride=1, + # padding=0, dilation=1, groups=1, + # bias=True, padding_mode='zeros' + for d in (1, 2, 3): + if d == 3 and padding_mode == 'reflect': + # FIXME: remove after implementing reflection pad 3d + # https://github.com/pytorch/pytorch/issues/27655 + continue + padding = tuple(range(1, d + 1)) + cpp_padding = '{' + ', '.join(map(str, padding)) + '}' + input_size = (2, 2) + (4,) * d + output_size = (2, 3) + tuple(p + 1 for p in padding) # simplified from `(4 + 2 * p - 3) // 2 + 1` + new_module_tests.append( + dict( + module_name=f'Conv{d}d', + constructor_args=(2, 3, 3, 2, padding, 1, 1, True, padding_mode), + cpp_constructor_args=f'''torch::nn::Conv{d}dOptions(2, 3, 3) + .stride(2) + .padding({cpp_padding}) + .dilation(1) + .groups(1) + .bias(true) + .padding_mode({cpp_padding_mode})''', + input_size=input_size, + output_size=output_size, + cudnn=True, + desc=f'{padding_mode}_stride2_pad2', + with_tf32=True, + tf32_precision=0.05, + default_dtype=torch.double, + ), + ) + + # Check that non linear activations work with no batch dimensions + non_linear_activations_no_batch = [ + 'ELU', 'Hardshrink', 'Hardsigmoid', 'Hardtanh', 'Hardswish', 'LeakyReLU', + 'LogSigmoid', 'PReLU', 'ReLU', 'ReLU6', 'RReLU', 'SELU', 'CELU', 'GELU', 'GLU', + 'Sigmoid', 'SiLU', 'Mish', 'Softplus', 'Softshrink', 'Softsign', 'Tanh', + 'Tanhshrink', 'Threshold' + ] + non_linear_activations_extra_info: dict[str, dict] = { + 'CELU': {'constructor_args': (2.,), 'default_dtype': torch.double}, + 'Threshold': {'constructor_args': (2., 1.)}, + 'Hardsigmoid': {'check_gradgrad': False, 'check_jit': False, 'default_dtype': torch.double}, + 'Hardswish': {'check_gradgrad': False, 'check_jit': False, 'default_dtype': torch.double}, + # For RRelu, test that compare CPU and GPU results fail because RNG + # is different between CPU and GPU + 'RReLU': {'test_cuda': False, 'default_dtype': torch.double}, + 'ELU': {'default_dtype': torch.double}, + 'GELU': {'default_dtype': torch.double}, + 'GLU': {'default_dtype': torch.double}, + 'Hardshrink': {'default_dtype': torch.double}, + 'Hardtanh': {'default_dtype': torch.double}, + 'LeakyReLU': {'default_dtype': torch.double}, + 'LogSigmoid': {'default_dtype': torch.double}, + 'Mish': {'default_dtype': torch.double}, + 'PReLU': {'default_dtype': torch.double}, + 'ReLU6': {'default_dtype': torch.double}, + 'ReLU': {'default_dtype': torch.double}, + 'SELU': {'default_dtype': torch.double}, + 'SiLU': {'default_dtype': torch.double}, + 'Sigmoid': {'default_dtype': torch.double}, + 'Softplus': {'default_dtype': torch.double}, + 'Softshrink': {'default_dtype': torch.double}, + 'Softsign': {'default_dtype': torch.double}, + 'Tanh': {'default_dtype': torch.double}, + 'Tanhshrink': {'default_dtype': torch.double}, + } + for non_linear_activation in non_linear_activations_no_batch: + activation_test_info = dict( + module_name=non_linear_activation, + input_size=(4,), + reference_fn=single_batch_reference_fn, + desc='no_batch_dim', + test_cpp_api_parity=False, + ) + extra_info = non_linear_activations_extra_info.get(non_linear_activation, {}) + activation_test_info.update(extra_info) + new_module_tests.append(activation_test_info) + + + return new_module_tests + + +def kldivloss_reference(input, target, reduction='mean', log_target=False): + if log_target: + result = torch.exp(target) * (target - input) + else: + result = target * (target.log() - input) + if reduction == 'mean': + return result.mean() + elif reduction == 'sum': + return result.sum() + elif reduction == 'batchmean' and result.dim() != 0: + return result.sum() / result.size(0) + return result + + +def nlllossNd_reference(input, target, weight=None, ignore_index=-100, + reduction='mean'): + if input.dim() < 3: + raise AssertionError(f"Expected input.dim() >= 3, got {input.dim()}") + N = input.size(0) + C = input.size(1) + out_size = (N,) + input.size()[2:] + output = torch.zeros(out_size).type_as(input) + + if weight is None: + weight = torch.ones(C).type_as(input) + total_weight = 0 + for tup in product(*[range(size) for size in out_size]): + t_nx = target[tup] + norm = 0. if ignore_index == t_nx else weight[t_nx].item() + input_index = list(tup) + input_index.insert(1, t_nx) + output[tup] = -input[tuple(input_index)] * norm + total_weight += norm + + if reduction == 'mean': + return output.sum() / total_weight + elif reduction == 'sum': + return output.sum() + return output + + +def cross_entropy_loss_prob_target_reference(input, target, weight=None, reduction='mean', + label_smoothing=0.0): + if input.dim() < 2: + raise AssertionError(f"Expected input.dim() >= 2, got {input.dim()}") + + input = torch.log_softmax(input, 1) + C = input.size(1) + if weight is None: + weight = torch.ones(C).type_as(input) + weight = weight.view(1, C, *(1 for _ in input.shape[2:])) + + if label_smoothing > 0.0: + if label_smoothing > 1.0: + raise AssertionError(f"Expected label_smoothing <= 1.0, got {label_smoothing}") + target = (target * (1 - label_smoothing) + label_smoothing / C) + + output = -(input * target * weight).sum(dim=1) + if reduction == 'mean': + return output.mean() + elif reduction == 'sum': + return output.sum() + return output + + +def cross_entropy_loss_indices_target_reference(input, target, weight=None, ignore_index=-100, + reduction='mean', label_smoothing=0.0): + log_softmax_input = torch.log_softmax(input, 1) + nllloss = F.nll_loss( + log_softmax_input, + target, + weight, + ignore_index=ignore_index, + reduction=reduction) + + if label_smoothing == 0.0: + return nllloss + + if not (0.0 < label_smoothing <= 1.0): + raise AssertionError(f"Expected 0.0 < label_smoothing <= 1.0, got {label_smoothing}") + + input = torch.log_softmax(input, 1) + C = input.size(1) + if weight is not None: + input = input * weight.view(1, C, *(1 for _ in input.shape[2:])) + + smooth_loss = -torch.sum(input, 1) + + ignore_mask = target == ignore_index + smooth_loss.masked_fill_(ignore_mask, 0.0) + + if reduction == 'mean': + if weight is not None: + # TODO: This code can path can be removed if #61309 is resolved + # loss is normalized by the weights to be consistent with nll_loss_nd + ret = torch.sum(smooth_loss) / weight.gather(0, target.masked_select(ignore_mask.logical_not()).flatten()).sum() + else: + ret = torch.mean(smooth_loss.masked_select(ignore_mask.logical_not())) + elif reduction == 'sum': + ret = torch.sum(smooth_loss) + else: + ret = smooth_loss + + return (1 - label_smoothing) * nllloss + ret * (label_smoothing / C) + + +def cross_entropy_loss_reference(input, target, weight=None, ignore_index=-100, reduction='mean', + label_smoothing=0.0): + if input.shape == target.shape: + return cross_entropy_loss_prob_target_reference( + input, + target, + weight=weight, + reduction=reduction, + label_smoothing=label_smoothing) + else: + return cross_entropy_loss_indices_target_reference( + input, target, weight=weight, reduction=reduction, + ignore_index=ignore_index, label_smoothing=label_smoothing + ) + + +def nllloss_reference(input, target, weight=None, ignore_index=-100, + reduction='mean'): + + def nll_loss_helper(input, target, weight, ignore_index): + if target == ignore_index: + return (0, 0) + norm = 1 if weight is None else weight[target] + result = -input[target] * norm + return (result, norm) + + losses_and_weights = [nll_loss_helper(i, t, weight, ignore_index) + for i, t in zip(input, target, strict=True)] + losses, weights = zip(*losses_and_weights, strict=True) + losses_tensor = input.new_tensor(losses) + if reduction == 'mean': + return sum(losses_tensor) / sum(weights) + elif reduction == 'sum': + return sum(losses_tensor) + else: + return losses_tensor + + +def smoothl1loss_reference(input, target, reduction='mean', beta=1.0): + abs_diff = (input - target).abs() + ge_beta_mask = (abs_diff >= beta).type_as(abs_diff) + lt_beta_mask = (abs_diff < beta).type_as(abs_diff) + # when beta <= 0 we should just use l1_loss + if beta == 0: + output = abs_diff + else: + output = ge_beta_mask * (abs_diff - 0.5 * beta) + lt_beta_mask * 0.5 * (abs_diff ** 2) / beta + if reduction == 'mean': + return output.mean() + elif reduction == 'sum': + return output.sum() + return output + + +def huberloss_reference(input, target, reduction='mean', delta=1.0): + abs_diff = (input - target).abs() + ge_delta_mask = (abs_diff >= delta) + lt_delta_mask = (abs_diff < delta) + output = ge_delta_mask * delta * (abs_diff - 0.5 * delta) + lt_delta_mask * 0.5 * (abs_diff ** 2) + if reduction == 'mean': + return output.mean() + elif reduction == 'sum': + return output.sum() + return output + + +def _multilabelmarginloss_reference(input, target): + targets = [] + for target_index in target: + if target_index < 0: + break + targets.append(target_index) + + sum = 0 + for target_index in targets: + for i in range(len(input)): + if i not in targets: + sum += max(0, 1 - input[target_index] + input[i]) + + return sum + + +def multilabelmarginloss_reference(input, target, reduction='mean'): + # make everything 2-dimensional + input_dim = input.dim() + if input.dim() < 2: + if target.dim() >= 2: + raise AssertionError(f"Expected target.dim() < 2, got {target.dim()}") + input = input.unsqueeze(0) if input.dim() == 1 else input.unsqueeze(0).unsqueeze(0) + target = target.unsqueeze(0) if target.dim() == 1 else target.unsqueeze(0).unsqueeze(0) + + n = input.size(0) + dim = input.size(1) + output = input.new(n).zero_() + for i in range(n): + output[i] = _multilabelmarginloss_reference(input[i], target[i]) + + if reduction == 'mean': + return output.mean() / dim + elif reduction == 'sum': + return output.sum() / dim + elif input_dim < 2: + # we know we have (1, C) X (1, C) -> (1,), so squeeze will get us + # back to correct dimensionality + return output.squeeze() / dim + else: + return output / dim + + +def hingeembeddingloss_reference(input, target, margin=1.0, reduction='mean'): + margin_clamp = (margin - input).clamp(min=0).type_as(input) + output = torch.where(target == 1, input, margin_clamp) + + if reduction == 'mean': + return output.mean() + elif reduction == 'sum': + return output.sum() + return output + + +def softmarginloss_reference(input, target, reduction='mean'): + output = (1 + (-input * target).exp()).log() + + if reduction == 'mean': + return output.mean() + elif reduction == 'sum': + return output.sum() + return output + + +def _multimarginloss_reference(input, target_idx, p, margin, weight): + if weight is None: + weight = input.new(len(input)).fill_(1) + + output = 0 + for i in range(len(input)): + if i != target_idx: + output += weight[target_idx] * (max(0, (margin - input[target_idx] + input[i])) ** p) + return output + + +def multimarginloss_reference(input, target, p=1, margin=1, weight=None, reduction='mean'): + if input.dim() < 2: + input = input.unsqueeze(0) if input.dim() == 1 else input.unsqueeze(0).unsqueeze(0) + + target_dim = target.dim() + if target.dim() == 0: + target = target.unsqueeze(0) + + n = input.size(0) + dim = input.size(1) + output = input.new(n) + for x in range(n): + output[x] = _multimarginloss_reference(input[x], target[x], p, margin, weight) + + if reduction == 'mean': + return output.mean() / dim + elif reduction == 'sum': + return output.sum() / dim + elif target_dim == 0: + return output.squeeze(0) / dim + return output / dim + + +def cosineembeddingloss_reference(input1, input2, target, margin=0, reduction='mean'): + def _cos(a, b): + cos = a.new(a.size(0)) + for i in range(a.size(0)): + cos[i] = (a[i] * b[i]).sum() / ((((a[i] * a[i]).sum() + 1e-12) * ((b[i] * b[i]).sum() + 1e-12)) ** 0.5) + return cos + + output = torch.where(target == 1, 1 - _cos(input1, input2), (_cos(input1, input2) - margin).clamp(min=0)) + + if reduction == 'mean': + return output.mean() + elif reduction == 'sum': + return output.sum() + return output + + +def tripletmarginloss_reference(anchor, positive, negative, margin=1.0, p=2, eps=1e-6, swap=False, + reduction='mean'): + d_p = torch.pairwise_distance(anchor, positive, p, eps) + d_n = torch.pairwise_distance(anchor, negative, p, eps) + if swap: + d_s = torch.pairwise_distance(positive, negative, p, eps) + d_n = torch.min(d_n, d_s) + + output = torch.clamp(margin + d_p - d_n, min=0.0) + if reduction == 'mean': + return output.mean() + elif reduction == 'sum': + return output.sum() + return output + + +def marginrankingloss_reference(input1, input2, target, margin=0, reduction='mean'): + output = (-target * (input1 - input2) + margin).clamp(min=0) + if reduction == 'mean': + return output.mean() + elif reduction == 'sum': + return output.sum() + return output + + +# this directly follows Graves et al.'s paper, in contrast to the production implementation, it does not use log-space +def ctcloss_reference(log_probs, targets, input_lengths, target_lengths, blank=0, reduction='mean'): + input_lengths = torch.as_tensor(input_lengths, dtype=torch.long) + target_lengths = torch.as_tensor(target_lengths, dtype=torch.long) + dt = log_probs.dtype + log_probs = log_probs.double() # we need the accuracy as we are not in logspace + targets = targets.long() + cum_target_lengths = target_lengths.cumsum(0) + losses = [] + for i in range(log_probs.size(1)): + input_length = input_lengths[i].item() + target_length = target_lengths[i].item() + cum_target_length = cum_target_lengths[i].item() + targets_prime = targets.new_full((2 * target_length + 1,), blank) + if targets.dim() == 2: + targets_prime[1::2] = targets[i, :target_length] + else: + targets_prime[1::2] = targets[cum_target_length - target_length:cum_target_length] + probs = log_probs[:input_length, i].exp() + alpha = log_probs.new_zeros((target_length * 2 + 1,)) + alpha[0] = probs[0, blank] + alpha[1] = probs[0, targets_prime[1]] + mask_third = (targets_prime[:-2] != targets_prime[2:]) + for t in range(1, input_length): + alpha_next = alpha.clone() + alpha_next[1:] += alpha[:-1] + alpha_next[2:] += torch.where(mask_third, alpha[:-2], alpha.new_zeros(1)) + alpha = probs[t, targets_prime] * alpha_next + losses.append(-alpha[-2:].sum().log()[None]) + output = torch.cat(losses, 0) + if reduction == 'mean': + output = (output / target_lengths.to(dtype=output.dtype, device=output.device)).mean() + elif reduction == 'sum': + output = output.sum() + output = output.to(dt) + return output + + +loss_reference_fns: dict['str', Callable] = { + 'KLDivLoss': kldivloss_reference, + 'KLDivLoss_log_target': partial(kldivloss_reference, log_target=True), + 'NLLLoss': nllloss_reference, + 'NLLLossNd': nlllossNd_reference, + 'SmoothL1Loss': smoothl1loss_reference, + 'HuberLoss': huberloss_reference, + 'MultiLabelMarginLoss': multilabelmarginloss_reference, + 'HingeEmbeddingLoss': hingeembeddingloss_reference, + 'SoftMarginLoss': softmarginloss_reference, + 'MultiMarginLoss': multimarginloss_reference, + 'CosineEmbeddingLoss': cosineembeddingloss_reference, + 'TripletMarginLoss': tripletmarginloss_reference, + 'MarginRankingLoss': marginrankingloss_reference, + 'CTCLoss': ctcloss_reference, + 'CrossEntropyLoss': cross_entropy_loss_reference +} + + +criterion_tests = [] + + +def single_batch_reference_criterion_fn(*args): + """Reference function for criterion supporting no batch dimensions. + + The criterion is passed the input and target in batched form with a single item. + The output is squeezed to compare with the no-batch input. + """ + criterion = args[-1] + + def unsqueeze_inp(inp): + if isinstance(inp, (list, tuple)): + return [t.unsqueeze(0) for t in inp] + return inp.unsqueeze(0) + + def flatten(xs): + result = [] + if isinstance(xs, (list, tuple)): + for x in xs: + result.extend(flatten(x)) + else: + result.append(xs) + return result + + single_batch_input_args = flatten([unsqueeze_inp(input) for input in args[:-1]]) + + output = criterion(*single_batch_input_args) + reduction = get_reduction(criterion) + + if reduction == 'none': + return output.squeeze(0) + # reduction is 'sum' or 'mean' which results in a scalar + return output + + +# Check that regression criterion work with no batch dimensions +regression_criterion_no_batch = [ + 'L1Loss', 'MSELoss', 'PoissonNLLLoss', 'HuberLoss', 'SmoothL1Loss' +] +reductions = ['none', 'mean', 'sum'] +for name, reduction in product(regression_criterion_no_batch, reductions): + regression_test_info = dict( + fullname=f"{name}_no_batch_dim_{reduction}", + constructor=lambda *args, name=name: getattr(nn, name)(reduction=reduction), + input_size=(3, ), + target_size=(3, ), + reference_fn=single_batch_reference_criterion_fn, + test_cpp_api_parity=False, + default_dtype=torch.double, + ) + criterion_tests.append(regression_test_info) + + +for reduction in reductions: + regression_test_info = dict( + fullname=f"KLDivLoss_no_batch_dim_{reduction}", + constructor=lambda: nn.KLDivLoss(reduction=reduction), + input_fn=lambda: torch.rand((3,)).log(), + target_fn=lambda: torch.rand((3,)), + reference_fn=single_batch_reference_criterion_fn, + test_cpp_api_parity=False, + default_dtype=torch.double, + ) + criterion_tests.append(regression_test_info) + + +# Check that classification criterion work with no batch dimensions +# List of tuples of (name, input_fn, target_fn) +classification_criterion_no_batch = [ + ( + 'BCELoss', + lambda: torch.sigmoid(torch.randn(9, dtype=torch.double)), + lambda: torch.randn(9, dtype=torch.double).gt(0).to(torch.double) + ), + ('BCEWithLogitsLoss', lambda: torch.randn(9, dtype=torch.double), lambda: torch.randn(9, dtype=torch.double)), + ('HingeEmbeddingLoss', lambda: torch.randn(9, dtype=torch.double), lambda: torch.tensor([-1, 1, 1] * 3)), + ('MultiLabelMarginLoss', lambda: torch.randn(4, dtype=torch.double), lambda: torch.tensor([3, 0, -1, 1])), + ('SoftMarginLoss', lambda: torch.randn(9, dtype=torch.double), lambda: torch.tensor([-1, 1, 1] * 3)), + ('NLLLoss', lambda: F.log_softmax(torch.randn(3, dtype=torch.double), dim=0), lambda: torch.tensor(1)), + ( + 'CosineEmbeddingLoss', + lambda: (torch.randn(9, dtype=torch.double), torch.randn(9, dtype=torch.double)), + lambda: torch.tensor(1, dtype=torch.double) + ), + # For MarginRankingLoss, input_fn : (x1, x2) and target_fn : target + ('MarginRankingLoss', lambda: (torch.randn(()), torch.randn(())), lambda: torch.randn(()).sign()), + # For TripletMarginLoss, input_fn : (anchor, positive) and target_fn : negative + ( + 'TripletMarginLoss', + lambda: (torch.randn(9, dtype=torch.double), torch.randn(9, dtype=torch.double)), + lambda: torch.randn(9, dtype=torch.double) + ), + ('MultiLabelSoftMarginLoss', lambda: torch.randn(9, dtype=torch.double), lambda: torch.randn(9)), +] +classification_criterion_no_batch_extra_info: dict[str, dict] = { + 'MultiLabelMarginLoss': {'check_gradgrad': False}, +} +# TODO : Fix these discrepancies +classification_cpp_parity = { + 'BCELoss': False, + 'BCEWithLogitsLoss': False, + 'HingeEmbeddingLoss': False, + 'NLLLoss': False, + 'SoftMarginLoss': False, +} +reductions = ['none', 'mean', 'sum'] +for (name, input_fn, target_fn), reduction in product(classification_criterion_no_batch, + reductions): + classification_test_info = dict( + fullname=f"{name}_no_batch_dim_{reduction}", + constructor=lambda *args, name=name: getattr(nn, name)(reduction=reduction), + input_fn=lambda f=input_fn: f(), + target_fn=lambda f=target_fn: f(), + reference_fn=single_batch_reference_criterion_fn, + test_cpp_api_parity=True, + has_parity=classification_cpp_parity.get(name, True) + ) + extra_info = classification_criterion_no_batch_extra_info.get(name, {}) + classification_test_info.update(extra_info) + criterion_tests.append(classification_test_info) + + +class NNTestCase(TestCase): + + # _forward is defined in classes inheriting from NNTestCase + @abstractmethod + def _forward(self, *args, **kwargs): + raise NotImplementedError + + @abstractmethod + def _get_parameters(self, module: nn.Module) -> tuple[list[nn.Parameter], list[nn.Parameter]]: + raise NotImplementedError + + @abstractmethod + def _zero_grad_parameters(self, module: nn.Module) -> None: + raise NotImplementedError + + @abstractmethod + def _backward(self, module: nn.Module, + input: _TensorOrTensors, output: torch.Tensor, + grad_output: torch.Tensor | Sequence[torch.Tensor], + create_graph: bool = False): + raise NotImplementedError + + def _jacobian(self, input, num_out): + if isinstance(input, tuple): + return tuple(self._jacobian(elem, num_out) for elem in input) + elif isinstance(input, list): + return [self._jacobian(elem, num_out) for elem in input] + else: + return torch.zeros(input.nelement(), num_out) + + def _flatten_tensors(self, x): + if isinstance(x, torch.Tensor): + if x.is_sparse: + return x.to_dense().view(-1) + else: + return x.view(-1) + else: + return tuple(self._flatten_tensors(a) for a in x) + + def _zero_grad_input(self, input): + if isinstance(input, torch.Tensor): + if input.requires_grad and input.grad is not None: + input.grad.zero_() + input.grad.detach_() + else: + for i in input: + self._zero_grad_input(i) + + def _analytical_jacobian(self, module, input: _TensorOrTensors, jacobian_input=True, jacobian_parameters=True): + output = self._forward(module, input) + output_size = output.nelement() + + if jacobian_input: + jacobian_inp = self._jacobian(input, output_size) + flat_jacobian_input = list(_iter_tensors(jacobian_inp)) + + if jacobian_parameters: + num_param = sum(p.numel() for p in self._get_parameters(module)[0]) + jacobian_param = torch.zeros(num_param, output_size) + + for i in range(output_size): + param, d_param = self._get_parameters(module) + # make non grad zeros + d_param = [torch.zeros_like(p) if d is None else d for (p, d) in zip(param, d_param, strict=True)] + + d_out = torch.zeros_like(output) + flat_d_out = d_out.view(-1) + flat_d_out[i] = 1 + + if jacobian_parameters: + self._zero_grad_parameters(module) + # Tensors will accumulate gradient from multiple steps + if jacobian_input: + self._zero_grad_input(input) + d_input = self._backward(module, input, output, d_out) + + if jacobian_input: + for jacobian_x, d_x in zip(flat_jacobian_input, _iter_tensors(d_input), strict=True): + jacobian_x[:, i] = d_x.contiguous().view(-1) + if jacobian_parameters: + jacobian_param[:, i] = torch.cat(self._flatten_tensors(d_param), 0) + + res: tuple[torch.Tensor, ...] = () + if jacobian_input: + res += jacobian_inp, + if jacobian_parameters: + res += jacobian_param, + + return res + + def _numerical_jacobian(self, module, input: _TensorOrTensors, jacobian_input=True, jacobian_parameters=True): + def fw(*input): + return self._forward(module, input).detach() + + res: tuple[torch.Tensor, ...] = () + if jacobian_input: + res += _get_numerical_jacobian(fw, input, eps=1e-6), + if jacobian_parameters: + param, _ = self._get_parameters(module) + to_cat = [] + for p in param: + jacobian = _get_numerical_jacobian(fw, input, target=p, eps=1e-6) + # get_numerical_jacobian returns a list of tuples but we require a tensor + to_cat.append(jacobian[0][0]) + res += (torch.cat(to_cat, 0),) + return res + + def check_jacobian(self, module, input: _TensorOrTensors, jacobian_input=True): + jacobian_parameters = bool(self._get_parameters(module)[0]) + analytical = self._analytical_jacobian(module, input, jacobian_input, jacobian_parameters) + numerical = self._numerical_jacobian(module, input, jacobian_input, jacobian_parameters) + analytical_t = list(_iter_tensors(analytical)) + numerical_t = list(_iter_tensors(numerical)) + + differences = [] + for a, n in zip(analytical_t, numerical_t, strict=True): + if a.numel() != 0: + differences.append(a.add(n, alpha=-1).abs().max()) + # TODO: compare structure (ensure analytic jacobian has correct shape) + if len(differences) > 0: + self.assertLessEqual(max(differences), PRECISION) # type: ignore[type-var] + + +class TestBase: + + _required_arg_names = {'constructor_args', 'input', 'extra_args'} + + def __init__(self, constructor, desc='', reference_fn=None, fullname=None, **kwargs): + self.desc = desc + self.fullname = fullname + self.constructor = constructor + self.reference_fn = reference_fn + for name in self._required_arg_names: + if name not in kwargs and name + '_fn' not in kwargs and name + '_size' not in kwargs: + if name in {'constructor_args', 'extra_args'}: + kwargs[name] = () + else: + raise ValueError(f"{self.get_name()}: Specify {name} by a value, a function to generate it, or it's size!") + self._extra_kwargs = kwargs + self._arg_cache = {} + + def get_name(self): + if self.fullname is not None: + return 'test_' + self.fullname + + test_name = 'test_' + self.constructor.__name__ + if self.desc: + test_name += '_' + self.desc + return test_name + + def _unpack(self, value): + if isinstance(value, torch.Tensor): + return value + elif is_iterable(value): + return type(value)(self._unpack(v) for v in value) + else: + return value + + @property + def constructor_args(self): + return self._get_arg('constructor_args', True) + + @property + def extra_args(self): + return self._get_arg('extra_args', True) + + def _get_arg(self, name, unpack): + if name not in self._required_arg_names: + raise AssertionError(f"Expected name '{name}' to be in required arg names") + + if name not in self._arg_cache: + fn_name = name + '_fn' + size_name = name + '_size' + + if name in self._extra_kwargs: + self._arg_cache[name] = self._extra_kwargs[name] + elif fn_name in self._extra_kwargs: + self._arg_cache[name] = self._extra_kwargs[fn_name]() + else: + if size_name not in self._extra_kwargs: + raise AssertionError( + f"Missing `{name}`, `{size_name}` or `{fn_name}` for {self.get_name()}" + ) + + def map_tensor_sizes(sizes): + if isinstance(sizes, list): + return [map_tensor_sizes(s) for s in sizes] + elif isinstance(sizes, torch.Tensor): + return sizes.double() + else: + return torch.randn(sizes) + + self._arg_cache[name] = map_tensor_sizes(self._extra_kwargs[size_name]) + + return self._unpack(self._arg_cache[name]) if unpack else self._arg_cache[name] + + def _get_input(self, unpack=True): + return self._get_arg('input', unpack) + + def __call__(self, test_case): + raise NotImplementedError + + +class ModuleTest(TestBase): + + @abstractmethod + def _do_test(self, test_case: Any, module: nn.Module, input: Any) -> Any: + raise NotImplementedError + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.jacobian_input = kwargs.get('jacobian_input', True) + self.should_test_cuda = kwargs.get('test_cuda', True) + self.should_test_pickle = kwargs.get('pickle', True) + self.check_gradgrad = kwargs.get('check_gradgrad', True) + self.FIXME_no_cuda_gradgrad_comparison = \ + kwargs.get('FIXME_no_cuda_gradgrad_comparison', False) + self.precision = kwargs.get('precision', 2e-4) + self.check_forward_only = kwargs.get('check_forward_only', False) + self.default_dtype = kwargs.get('default_dtype') + if self.default_dtype is None: + self.default_dtype = torch.get_default_dtype() + + def __call__(self, test_case): + with set_default_dtype(self.default_dtype): + module = self.constructor(*self.constructor_args) + input = self._get_input() + + if self.reference_fn is not None: + out = test_case._forward(module, input) + ref_input = deepcopy(input) + ref_module = deepcopy(module) + expected_out = self.reference_fn(ref_input, test_case._get_parameters(module)[0], ref_module) + test_case.assertEqual(out, expected_out, exact_dtype=False) + if self.check_forward_only: + return + self.test_noncontig(test_case, module, input) + + if self.should_test_pickle: + # TODO: do this with in-memory files as soon as torch.save will support it + with tempfile.TemporaryFile() as f: + test_case._forward(module, input) + torch.save(module, f) + f.seek(0) + # weights_only=False as this is legacy code that saves the model + module_copy = torch.load(f, weights_only=False) + test_case.assertEqual(test_case._forward(module, input), test_case._forward(module_copy, input)) + + self._do_test(test_case, module, input) + + def noncontiguize(self, obj): + if isinstance(obj, list): + return [self.noncontiguize(o) for o in obj] + elif isinstance(obj, tuple): + return tuple(self.noncontiguize(o) for o in obj) + tensor = obj + ndim = tensor.dim() + # Always making only the last dimension noncontiguous is easy to hide + # bugs because .view(-1) will still work. So try to find a dim with size + # > 1 and make that non-contiguous, i.e., stack + select on the + # dimension directly after that. + dim = ndim + for d in range(ndim): + if tensor.size(d) > 1: + dim = d + 1 + break + noncontig = torch.stack([torch.empty_like(tensor), tensor], dim).select(dim, 1).detach() + if not (noncontig.numel() == 1 or noncontig.numel() == 0 or not noncontig.is_contiguous()): + raise AssertionError( + f"Expected noncontig to be non-contiguous or have numel <= 1, got numel={noncontig.numel()}" + ) + noncontig.requires_grad = tensor.requires_grad + return noncontig + + def test_noncontig(self, test_case, module, input): + # check no scalars, can't make non-contig + if isinstance(input, torch.Tensor) and input.dim() == 0: + return + if any(i.dim() == 0 for i in input if isinstance(i, torch.Tensor)): + return + + test_case._zero_grad_parameters(module) + test_case._zero_grad_input(input) + with freeze_rng_state(): + output = test_case._forward(module, input) + if getattr(module, "return_indices", False): + output = output[0] + grad_output = output.new(output.shape).normal_() + output = output.clone() + d_input = deepcopy(test_case._backward(module, input, output, grad_output)) + d_param = deepcopy(test_case._get_parameters(module)[1]) + + nc_input = self.noncontiguize(input) + nc_grad_output = self.noncontiguize(grad_output) + for contig_i, contig_g in product((True, False), repeat=2): + i = input if contig_i else nc_input + # Some ops, e.g., nn.Flatten, return gradient that shares + # storage with the grad_output. Hence we copy here. + go = deepcopy(grad_output if contig_g else nc_grad_output) + test_case._zero_grad_parameters(module) + test_case._zero_grad_input(i) + with freeze_rng_state(): + out = test_case._forward(module, i) + if getattr(module, "return_indices", False): + out = out[0] + grad = test_case._backward(module, i, out, go) + + test_case.assertEqual(out, output) + test_case.assertEqual(grad, d_input, atol=1e-4, rtol=0) + test_case.assertEqual(test_case._get_parameters(module)[1], d_param) + + def test_cuda(self, test_case): + if not TEST_CUDA or not self.should_test_cuda: + raise unittest.SkipTest('Excluded from CUDA tests') + + with set_default_dtype(self.default_dtype): + cpu_input = self._get_input() + + type_map = {torch.double: torch.float} + cpu_input_tuple = cpu_input if isinstance(cpu_input, tuple) else (cpu_input,) + + is_any_input_complex = any(isinstance(t, torch.Tensor) and t.dtype.is_complex for t in cpu_input_tuple) + + gpu_input_tuple = to_gpu(cpu_input_tuple, type_map=type_map) + + cpu_module = self.constructor(*self.constructor_args) + gpu_module = self.constructor(*self.constructor_args).float().cuda() + cpu_param = test_case._get_parameters(cpu_module) + gpu_param = test_case._get_parameters(gpu_module) + for cpu_p, gpu_p in zip(cpu_param[0], gpu_param[0], strict=True): + gpu_p.data.copy_(cpu_p) + + test_case._zero_grad_input(cpu_input_tuple) + test_case._zero_grad_input(gpu_input_tuple) + test_case._zero_grad_parameters(cpu_module) + test_case._zero_grad_parameters(gpu_module) + cpu_output = test_case._forward(cpu_module, cpu_input_tuple) + gpu_output = test_case._forward(gpu_module, gpu_input_tuple) + if getattr(cpu_module, "return_indices", False): + cpu_output = cpu_output[0] + gpu_output = gpu_output[0] + test_case.assertEqual(cpu_output, gpu_output, atol=self.precision, rtol=0, exact_dtype=False) + + # Run backwards on CPU and GPU and compare results + for _ in range(5): + cpu_gradOutput = cpu_output.clone().normal_() + gpu_gradOutput = cpu_gradOutput.type_as(gpu_output) + cpu_gradInput = test_case._backward(cpu_module, cpu_input_tuple, cpu_output, cpu_gradOutput) + gpu_gradInput = test_case._backward(gpu_module, gpu_input_tuple, gpu_output, gpu_gradOutput) + test_case.assertEqual(cpu_gradInput, gpu_gradInput, atol=self.precision, rtol=0, exact_dtype=False) + for cpu_d_p, gpu_d_p in zip(cpu_param[1], gpu_param[1], strict=True): + test_case.assertEqual(cpu_d_p, gpu_d_p, atol=self.precision, rtol=0) + + # Run double-backwards on CPU and GPU and compare results + if self.check_gradgrad and not self.FIXME_no_cuda_gradgrad_comparison: + cpu_output = cpu_module(*cpu_input_tuple) + gpu_output = gpu_module(*gpu_input_tuple) + if getattr(cpu_module, "return_indices", False): + cpu_output = cpu_output[0] + gpu_output = gpu_output[0] + + cpu_gradOutput = torch.randn_like(cpu_output, requires_grad=True) + gpu_gradOutput = cpu_gradOutput.type_as(gpu_output).detach() + gpu_gradOutput.requires_grad = True + + cpu_gradInputs = torch.autograd.grad( + cpu_output, + cpu_input_tuple + tuple(cpu_module.parameters()), + cpu_gradOutput, + create_graph=True) + gpu_gradInputs = torch.autograd.grad( + gpu_output, + gpu_input_tuple + tuple(gpu_module.parameters()), + gpu_gradOutput, + create_graph=True) + + for cpu_d_i, gpu_d_i in zip(cpu_gradInputs, gpu_gradInputs, strict=True): + test_case.assertEqual(cpu_d_i, gpu_d_i, atol=self.precision, rtol=0, exact_dtype=False) + + # We mix output into the second backwards computation so that + # torch.autograd.grad doesn't complain that some inputs + # are unreachable (which can happen if you differentiate + # only on the gradient. + if is_any_input_complex: + outputs_cpu = cpu_output.sum().abs() + sum(x.sum().abs() for x in cpu_gradInputs) + outputs_gpu = gpu_output.sum().abs() + sum(x.sum().abs() for x in gpu_gradInputs) + else: + outputs_cpu = cpu_output.sum() + sum(x.sum() for x in cpu_gradInputs) + outputs_gpu = gpu_output.sum() + sum(x.sum() for x in gpu_gradInputs) + + cpu_gg = torch.autograd.grad( + outputs_cpu, + cpu_input_tuple + (cpu_gradOutput,) + tuple(cpu_module.parameters()), + retain_graph=True) + gpu_gg = torch.autograd.grad( + outputs_gpu, + gpu_input_tuple + (gpu_gradOutput,) + tuple(gpu_module.parameters()), + retain_graph=True) + test_case.assertEqual(cpu_gradInput, gpu_gradInput, atol=self.precision, rtol=0, exact_dtype=False) + for cpu_d_p, gpu_d_p in zip(cpu_gg, gpu_gg, strict=True): + test_case.assertEqual(cpu_d_p, gpu_d_p, atol=self.precision, rtol=0, exact_dtype=False) + + self.test_noncontig(test_case, gpu_module, gpu_input_tuple) + + +class InputVariableMixin: + def _get_input(self): + input = TestBase._get_input(self, False) # type: ignore[arg-type] + + def map_variables(i): + if isinstance(i, torch.Tensor): + if i.is_floating_point() or i.is_complex(): + i.requires_grad = True + return i + else: + return type(i)(map_variables(elem) for elem in i) + + return map_variables(input) + + +class NewModuleTest(InputVariableMixin, ModuleTest): # type: ignore[misc] + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.cudnn = kwargs.get('cudnn', False) + self.check_inplace = kwargs.get('check_inplace', False) + self.check_gradgrad = kwargs.get('check_gradgrad', True) + self.skip_double = kwargs.get('skip_double', False) + self.skip_half = kwargs.get('skip_half', False) + self.with_tf32 = kwargs.get('with_tf32', False) + self.tf32_precision = kwargs.get('tf32_precision', 0.001) + self.test_cpu = kwargs.get('test_cpu', True) + self.has_sparse_gradients = kwargs.get('has_sparse_gradients', False) + self.check_batched_grad = kwargs.get('check_batched_grad', True) + self.gradcheck_fast_mode = kwargs.get('gradcheck_fast_mode') + self.supports_forward_ad = kwargs.get('supports_forward_ad', False) + self.supports_fwgrad_bwgrad = kwargs.get('supports_fwgrad_bwgrad', False) + + def _check_gradients(self, test_case, module, input_tuple): + params = tuple(x for x in module.parameters()) + num_inputs = len(input_tuple) + + def fn_to_gradcheck(*inputs_and_params, **kwargs): + if kwargs: + raise AssertionError(f"Expected no kwargs, got {kwargs}") + return test_case._forward(module, inputs_and_params[:num_inputs]) + + # gradcheck doesn't support operators that take in dense inputs but + # return sparse parameters. This only happens in the case of nn.Embedding + # and nn.EmbeddingBag. Instead, we call `self.check_jacobian`, which + # is a slightly different version of gradcheck that can handle this. + if self.has_sparse_gradients: + if num_inputs != 1: + raise AssertionError(f"Expected num_inputs == 1, got {num_inputs}") + test_input_jacobian = torch.is_floating_point(input_tuple[0]) + test_case.check_jacobian(module, input_tuple[0], test_input_jacobian) + else: + test_case.assertTrue(gradcheck(fn_to_gradcheck, input_tuple + params, + check_batched_grad=self.check_batched_grad, + fast_mode=self.gradcheck_fast_mode, + check_forward_ad=self.supports_forward_ad)) + + if self.check_gradgrad: + test_case.assertTrue(gradgradcheck(fn_to_gradcheck, input_tuple + params, + check_batched_grad=self.check_batched_grad, + fast_mode=self.gradcheck_fast_mode, + check_fwd_over_rev=self.supports_fwgrad_bwgrad)) + + def _do_test(self, test_case, module, input): + num_threads = torch.get_num_threads() + torch.set_num_threads(1) + input_tuple = input if isinstance(input, tuple) else (input,) + + self._check_gradients(test_case, module, input_tuple) + + # check if module can be printed + module.__repr__() + + if self.check_inplace: + # check if the inplace variant of the module gives the same result + # as the out-of-place + + # check_inplace doesn't support multiple input tensors, since we don't have any modules + # that modify the inputs in-place and that accept more than one input + if len(input_tuple) != 1: + raise AssertionError(f"Expected len(input_tuple) == 1, got {len(input_tuple)}") + input = input_tuple[0] + + module_ip = self.constructor(*self.constructor_args, inplace=True) + + input_version = input._version + with freeze_rng_state(): + output = module(input) + test_case.assertEqual(input._version, input_version) + + input_ip = deepcopy(input) + input_ip_clone = input_ip.clone() + with freeze_rng_state(): + output_ip = module_ip(input_ip_clone) + test_case.assertNotEqual(input_ip_clone._version, input_version) + test_case.assertEqual(output, output_ip) + grad = output.data.clone().normal_() + if input.grad is not None: + with torch.no_grad(): + input.grad.zero_() + if input_ip.grad is not None: + with torch.no_grad(): + input_ip.grad.zero_() + output.backward(grad) + output_ip.backward(grad) + test_case.assertEqual(input.grad, input_ip.grad) + + def assert_module_parameters_are(tensor_type, device_id=None): + for p in module.parameters(): + test_case.assertIsInstance(p, tensor_type) + if device_id is not None: + test_case.assertEqual(p.get_device(), device_id) + + if all(isinstance(t, torch.LongTensor) for t in input_tuple) and TEST_CUDA: + # check that cuda() moves module parameters to correct GPU device, + # and that float() casts parameters correctly + input_tuple = tuple(t.cuda() for t in input_tuple) + module.float().cuda() + module(*input_tuple) + assert_module_parameters_are(torch.cuda.FloatTensor, 0) # type: ignore[attr-defined] + + if torch.cuda.device_count() > 1: + input_tuple = tuple(t.cuda(1) for t in input_tuple) + module.cuda(1) + with torch.cuda.device(1): + module(*input_tuple) + assert_module_parameters_are(torch.cuda.FloatTensor, 1) # type: ignore[attr-defined] + else: + # check that float()/double() casters work correctly + def to_type(tensor, real, complex): + if tensor.is_complex(): + return tensor.to(complex) + elif tensor.is_floating_point(): + return tensor.to(real) + else: + return tensor + + def to_half(x): + # TODO: torch.complex32 when properly supported + return to_type(x, torch.float16, None) + + def to_single(x): + return to_type(x, torch.float32, torch.complex64) + + def to_double(x): + return to_type(x, torch.float64, torch.complex128) + + # to float + input_tuple = tuple(to_single(t) for t in input_tuple) + module.float() + module(*input_tuple) + assert_module_parameters_are(torch.FloatTensor) + + # and back to double + input_tuple = tuple(to_double(t) for t in input_tuple) + module.double() + module(*input_tuple) + assert_module_parameters_are(torch.DoubleTensor) + + if TEST_CUDA and self.should_test_cuda: + # check that cuda() moves module parameters to correct GPU device, + # and that float() casts parameters correctly + + # to GPU0 + input_tuple = tuple(to_single(t).cuda() for t in input_tuple) + module.float().cuda() + module(*input_tuple) + assert_module_parameters_are(torch.cuda.FloatTensor, 0) # type: ignore[attr-defined] + + # to CPU + input_tuple = tuple(t.cpu() for t in input_tuple) + module.cpu() + module(*input_tuple) + assert_module_parameters_are(torch.FloatTensor) + + # back to GPU0 + input_tuple = tuple(t.cuda() for t in input_tuple) + module.cuda() + module(*input_tuple) + assert_module_parameters_are(torch.cuda.FloatTensor, 0) # type: ignore[attr-defined] + + # test that forwards of module runs correctly without cuDNN + if self.cudnn: + with torch.backends.cudnn.flags(enabled=False): + module(*input_tuple) + assert_module_parameters_are(torch.cuda.FloatTensor, 0) # type: ignore[attr-defined] + + if torch.cuda.device_count() >= 2: + # test cross-GPU transfer works + # to GPU1 + input_tuple = tuple(t.cuda(1) for t in input_tuple) + module.cuda(1) + with torch.cuda.device(1): + module(*input_tuple) + assert_module_parameters_are(torch.cuda.FloatTensor, 1) # type: ignore[attr-defined] + + if not self.skip_double: + # test double() + input_tuple = tuple(to_double(t).cuda() for t in input_tuple) + module.double().cuda() + module(*input_tuple) + assert_module_parameters_are(torch.cuda.DoubleTensor, 0) # type: ignore[attr-defined] + + # test half() + if not self.skip_half: + input_tuple = tuple(to_half(t).cuda() for t in input_tuple) + module.half().cuda() + module(*input_tuple) + assert_module_parameters_are(torch.cuda.HalfTensor, 0) # type: ignore[attr-defined] + torch.set_num_threads(num_threads) + + def _get_target(self): + return self._get_arg('target', False) + + @property + def constructor_args(self): + return self._get_arg('constructor_args', False) + + +class CriterionTest(InputVariableMixin, TestBase): # type: ignore[misc] + # TODO: check that criterions don't ignore grad_output + + _required_arg_names = TestBase._required_arg_names.union({'target'}) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.should_test_cuda = kwargs.get('test_cuda', True) + self.check_forward_only = kwargs.get('check_forward_only', False) + self.check_gradgrad = kwargs.get('check_gradgrad', True) + self.check_half = kwargs.get('check_half', True) + self.check_bfloat16 = kwargs.get('check_bfloat16', False) + self.check_complex = kwargs.get('check_complex', False) + self.test_cpu = kwargs.get('test_cpu', True) + self.with_tf32 = kwargs.get('with_tf32', True) + self.tf32_precision = kwargs.get('tf32_precision', 0.001) + self.check_batched_grad = kwargs.get('check_batched_grad', True) + self.default_dtype = kwargs.get('default_dtype') + if self.default_dtype is None: + self.default_dtype = torch.get_default_dtype() + + def __call__(self, test_case): + with set_default_dtype(self.default_dtype): + module = self.constructor(*self.constructor_args) + input = self._get_input() + + # Check that these methods don't raise errors + module.__repr__() + str(module) + + target = self._get_target() + + if self.reference_fn is not None: + out = test_case._forward_criterion(module, input, target, extra_args=self.extra_args) + ref_args = (deepcopy(input), deepcopy(target)) + self.extra_args + (module,) + expected_out = self.reference_fn(*ref_args) + test_case.assertEqual(out, expected_out) + + if self.check_forward_only: + return + + params = tuple(x for x in module.parameters()) + if not isinstance(input, tuple): + inputs = (input,) + params + (target,) + + def apply_fn(input, target, *params): + return module(input, target) + else: + inputs = input + params + (target,) + + def apply_fn(input1, input2, target, *params): # type: ignore[misc] + return module(input1, input2, target) + + gradcheck(apply_fn, inputs, check_batched_grad=self.check_batched_grad) + + if self.check_gradgrad: + gradgradcheck(apply_fn, inputs, check_batched_grad=self.check_batched_grad) + + def test_cuda(self, test_case, dtype, extra_args=None): + def convert_dtype(obj, dtype, requires_grad=False): + if isinstance(obj, torch.Tensor): + return obj.detach().to(dtype=dtype).requires_grad_(requires_grad) + elif isinstance(obj, tuple): + return tuple(convert_dtype(o, dtype, requires_grad) for o in obj) + else: + return obj + + if not TEST_CUDA or not self.should_test_cuda: + raise unittest.SkipTest('Excluded from CUDA tests') + + with set_default_dtype(self.default_dtype): + cpu_input = self._get_input() + cpu_target = self._get_target() + cpu_module = self.constructor(*self.constructor_args) + gpu_module = self.constructor(*self.constructor_args) + + # Convert input, target and module parameters to dtype + cpu_input = convert_dtype(cpu_input, dtype, True) + if cpu_target.is_floating_point() or cpu_target.is_complex(): + cpu_target = convert_dtype(cpu_target, dtype) + cpu_module.type(dtype) + gpu_module.type(dtype) + + # GPU setup + gpu_input = to_gpu(cpu_input) + gpu_target = to_gpu(cpu_target) + gpu_module.cuda() + + # torch.HalfTensor doesn't support most operations, converting back to default + if dtype in {torch.half, torch.bfloat16}: + cpu_input = self._get_input() + cpu_target = self._get_target() + # Loss modules with weights require consistent input/module weight types + cpu_module = self.constructor(*self.constructor_args) + + cpu_output = test_case._forward_criterion(cpu_module, cpu_input, cpu_target, extra_args=extra_args) + gpu_output = test_case._forward_criterion(gpu_module, gpu_input, gpu_target, extra_args=extra_args) + # dtype used to be able to be None, so set precision in this way instead of a precision map + test_case.assertEqual(cpu_output, gpu_output, + atol=1e-1 if dtype in {torch.half, torch.bfloat16} else 4e-4, rtol=0, exact_dtype=False) + + cpu_gradInput = test_case._backward_criterion( + cpu_module, cpu_input, cpu_output, cpu_target, extra_args=extra_args) + gpu_gradInput = test_case._backward_criterion( + gpu_module, gpu_input, gpu_output, gpu_target, extra_args=extra_args) + # dtype used to be able to be None, so set precision in this way instead of a precision map + test_case.assertEqual(cpu_gradInput, gpu_gradInput, + atol=1e-1 if dtype in {torch.half, torch.bfloat16} else 4e-4, rtol=0, exact_dtype=False) + + def _get_target(self): + return self._get_arg('target', False) + + @property + def constructor_args(self): + return self._get_arg('constructor_args', False) + + @property + def extra_args(self): + return self._get_arg('extra_args', False) + + +def _test_bfloat16_ops(test_case, op, device, inp_dims=(), prec=1e-2, scale_factor=None): + # fp32 compute + input1 = torch.randn(inp_dims, dtype=torch.float32, device=device, requires_grad=True) + if scale_factor is not None: + input1 = (torch.rand(inp_dims, dtype=torch.bfloat16, device=device) * scale_factor).float().requires_grad_() + out1 = op(input1) + grad_input1 = torch.randn_like(out1, device=device) + out1.backward(grad_input1) + + # bfloat16 compute + op_bfp16 = op.bfloat16() + input2 = input1.detach().bfloat16().requires_grad_() + grad_input2 = grad_input1.bfloat16() + out2 = op_bfp16(input2) + out2.backward(grad_input2) + + test_case.assertEqual(out1, out2, atol=prec, rtol=prec, exact_dtype=False) + test_case.assertEqual(input1.grad.data, input2.grad.data, atol=prec, rtol=prec, exact_dtype=False) + +def _test_module_empty_input(test_case, module, inp, check_size=True, inference=False): + if not inference: + inp.requires_grad_(True) + out = module(inp) + if not inference: + gO = torch.rand_like(out) + out.backward(gO) + if check_size: + test_case.assertEqual(out.size(), inp.size()) + if not inference: + for p in module.parameters(): + if p.requires_grad: + test_case.assertEqual(p.grad, torch.zeros_like(p.grad)) + test_case.assertEqual(inp.grad, torch.zeros_like(inp)) + + +def _create_basic_net(): + class Layer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.layer_dummy_param = nn.Parameter(torch.empty(3, 5)) + self.layer_dummy_buf = nn.Buffer(torch.zeros(1, 3, 3, 7)) + + class Net(nn.Module): + def __init__(self) -> None: + super().__init__() + self.l1 = Layer() + self.dummy_param = nn.Parameter(torch.empty(3, 5)) + self.dummy_buf = nn.Buffer(torch.zeros(7, 3, 3, 1)) + + l = Layer() + n = Net() + s = nn.Sequential(n, n) + + return l, n, s diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_ops_unbacked.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_ops_unbacked.py new file mode 100644 index 0000000000000000000000000000000000000000..766126de0437556cfaa5477d507b4e33b0c44e98 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_ops_unbacked.py @@ -0,0 +1,258 @@ +# Owner(s): ["oncall: pt2"] + +""" +Shared xfail lists for unbacked symint tests. + +These lists are used by both test_ops_unbacked.py (base tensor tests) +and test_dtensor_ops.py (DTensor tests with unbacked dimensions). +""" + + +def xfail(op_name, variant_name="", *, device_type=None, dtypes=None): + return (op_name, variant_name, device_type, dtypes, True) + + +def skip(op_name, variant_name="", *, device_type=None, dtypes=None): + return (op_name, variant_name, device_type, dtypes, False) + + +# Ops that have data-dependent errors with unbacked dimensions. +# These fail at the base tensor level (not DTensor-specific). +ops_dde_xfail = { + xfail("_chunk_cat"), + xfail("_unsafe_masked_index_put_accumulate"), + xfail("_upsample_bilinear2d_aa"), + xfail("addmv"), + xfail("allclose"), + xfail("as_strided_scatter"), + xfail("baddbmm"), + xfail("bernoulli"), + xfail("cauchy"), + xfail("cdist"), + xfail("cholesky"), + xfail("chunk"), + xfail("combinations"), + xfail("corrcoef"), + xfail("cov"), + xfail("cross"), + xfail("cummax"), + xfail("cummin"), + xfail("cumulative_trapezoid"), + xfail("diagonal_scatter"), + xfail("diff"), + xfail("dsplit"), + xfail("equal"), + xfail("exponential"), + xfail("fft.fft"), + xfail("fft.fft2"), + xfail("fft.fftn"), + xfail("fft.fftshift"), + xfail("fft.hfft"), + xfail("fft.hfft2"), + xfail("fft.hfftn"), + xfail("fft.ifft"), + xfail("fft.ifft2"), + xfail("fft.ifftn"), + xfail("fft.ifftshift"), + xfail("fft.ihfft"), + xfail("fft.ihfft2"), + xfail("fft.ihfftn"), + xfail("fft.irfft"), + xfail("fft.irfft2"), + xfail("fft.irfftn"), + xfail("fft.rfft"), + xfail("fft.rfft2"), + xfail("fft.rfftn"), + xfail("float"), + xfail("geometric"), + xfail("geqrf"), + xfail("gradient"), + xfail("grid_sampler_2d"), + xfail("hash_tensor"), + xfail("histogram"), + xfail("histogramdd"), + xfail("hsplit"), + xfail("index_fill"), + xfail("inner"), + xfail("kron"), + xfail("linalg.cond"), + xfail("linalg.cross"), + xfail("linalg.householder_product"), + xfail("linalg.ldl_solve"), + xfail("linalg.lstsq"), + xfail("linalg.lstsq", "grad_oriented"), + xfail("linalg.lu_solve"), + xfail("linalg.matrix_norm"), + xfail("linalg.matrix_power"), + xfail("linalg.matrix_rank"), + xfail("linalg.matrix_rank", "hermitian"), + xfail("linalg.multi_dot"), + xfail("linalg.norm"), + xfail("linalg.norm", "subgradients_at_zero"), + xfail("linalg.pinv"), + xfail("linalg.pinv", "hermitian"), + xfail("linalg.pinv", "singular"), + xfail("linalg.qr"), + xfail("linalg.solve"), + xfail("linalg.solve_ex"), + xfail("linalg.solve_triangular"), + xfail("linalg.tensorinv"), + xfail("linalg.tensorsolve"), + xfail("linalg.vander"), + xfail("log_normal"), + xfail("logdet"), + xfail("logsumexp"), + xfail("lu_solve"), + xfail("lu_unpack"), + xfail("masked.amax"), + xfail("masked.amin"), + xfail("masked.argmax"), + xfail("masked.argmin"), + xfail("masked.cumprod"), + xfail("masked.cumsum"), + xfail("masked.log_softmax"), + xfail("masked.logaddexp"), + xfail("masked.logsumexp"), + xfail("masked.mean"), + xfail("masked.median"), + xfail("masked.norm"), + xfail("masked.prod"), + xfail("masked.softmax"), + xfail("masked.softmin"), + xfail("masked.std"), + xfail("masked.sum"), + xfail("masked.var"), + xfail("max_pool2d_with_indices_backward"), + xfail("multinomial"), + xfail("nn.functional.adaptive_avg_pool1d"), + xfail("nn.functional.adaptive_avg_pool2d"), + xfail("nn.functional.adaptive_avg_pool3d"), + xfail("nn.functional.adaptive_max_pool1d"), + xfail("nn.functional.adaptive_max_pool2d"), + xfail("nn.functional.adaptive_max_pool3d"), + xfail("nn.functional.alpha_dropout"), + xfail("nn.functional.avg_pool1d"), + xfail("nn.functional.avg_pool2d"), + xfail("nn.functional.avg_pool3d"), + xfail("nn.functional.batch_norm"), + xfail("nn.functional.bilinear"), + xfail("nn.functional.binary_cross_entropy"), + xfail("nn.functional.binary_cross_entropy_with_logits"), + xfail("nn.functional.channel_shuffle"), + xfail("nn.functional.cosine_similarity"), + xfail("nn.functional.cross_entropy"), + xfail("nn.functional.ctc_loss"), + xfail("nn.functional.dropout"), + xfail("nn.functional.dropout2d"), + xfail("nn.functional.dropout3d"), + xfail("nn.functional.embedding"), + xfail("nn.functional.embedding_bag"), + xfail("nn.functional.feature_alpha_dropout", "with_train"), + xfail("nn.functional.feature_alpha_dropout", "without_train"), + xfail("nn.functional.fractional_max_pool2d"), + xfail("nn.functional.fractional_max_pool3d"), + xfail("nn.functional.gaussian_nll_loss"), + xfail("nn.functional.grid_sample"), + xfail("nn.functional.group_norm"), + xfail("nn.functional.huber_loss"), + xfail("nn.functional.instance_norm"), + xfail("nn.functional.interpolate", "area"), + xfail("nn.functional.interpolate", "bicubic"), + xfail("nn.functional.interpolate", "bilinear"), + xfail("nn.functional.interpolate", "linear"), + xfail("nn.functional.interpolate", "trilinear"), + xfail("nn.functional.l1_loss"), + xfail("nn.functional.local_response_norm"), + xfail("nn.functional.max_pool1d"), + xfail("nn.functional.max_pool2d"), + xfail("nn.functional.max_pool3d"), + xfail("nn.functional.max_unpool1d"), + xfail("nn.functional.max_unpool1d", "grad"), + xfail("nn.functional.max_unpool2d"), + xfail("nn.functional.max_unpool2d", "grad"), + xfail("nn.functional.max_unpool3d"), + xfail("nn.functional.max_unpool3d", "grad"), + xfail("nn.functional.mse_loss"), + xfail("nn.functional.multi_head_attention_forward"), + xfail("nn.functional.multilabel_margin_loss"), + xfail("nn.functional.nll_loss"), + xfail("nn.functional.pad", "circular"), + xfail("nn.functional.pad", "reflect"), + xfail("nn.functional.pad", "replicate"), + xfail("nn.functional.pad", "replicate_negative"), + xfail("nn.functional.pdist"), + xfail("nn.functional.pixel_shuffle"), + xfail("nn.functional.prelu"), + xfail("nn.functional.rrelu"), + xfail("nn.functional.scaled_dot_product_attention"), + xfail("nn.functional.smooth_l1_loss"), + xfail("nn.functional.unfold"), + xfail("nn.functional.upsample_bilinear"), + xfail("normal"), + xfail("normal", "in_place"), + xfail("normal", "number_mean"), + xfail("ormqr"), + xfail("pca_lowrank"), + xfail("pinverse"), + xfail("qr"), + xfail("rand_like"), + xfail("randint_like"), + xfail("randn_like"), + xfail("repeat_interleave"), + xfail("resize_"), + xfail("resize_as_"), + xfail("roll"), + xfail("searchsorted"), + xfail("sparse.mm", "reduce"), + xfail("split"), + xfail("stft"), + xfail("svd_lowrank"), + xfail("sum_to_size"), + xfail("take"), + xfail("take_along_dim"), + xfail("tensordot"), + xfail("tensor_split"), + xfail("to_sparse"), + xfail("trapezoid"), + xfail("trapz"), + xfail("unbind"), + xfail("unbind_copy"), + xfail("uniform"), + xfail("unsafe_chunk"), + xfail("unsafe_split"), + xfail("vsplit"), +} + +# Ops that skip for unbacked tests (no valid samples with markable dims) +ops_unbacked_skip = { + skip("arange"), + skip("broadcast_shapes"), + skip("empty"), + skip("empty_permuted"), + skip("empty_strided"), + skip("eye"), + skip("full"), + skip("item"), + skip("linspace"), + skip("linspace", "tensor_overload"), + skip("logspace"), + skip("logspace", "tensor_overload"), + skip("ones"), + skip("randint"), + skip("randn"), + skip("scalar_tensor"), + skip("signal.windows.bartlett"), + skip("signal.windows.blackman"), + skip("signal.windows.cosine"), + skip("signal.windows.exponential"), + skip("signal.windows.gaussian"), + skip("signal.windows.general_cosine"), + skip("signal.windows.general_hamming"), + skip("signal.windows.hamming"), + skip("signal.windows.hann"), + skip("signal.windows.kaiser"), + skip("signal.windows.nuttall"), + skip("zeros"), + # Sparse ops that can't be deepcopied + skip("sparse.sampled_addmm"), +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_optimizers.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_optimizers.py new file mode 100644 index 0000000000000000000000000000000000000000..fb67c4d845021358a4a68f4da7f3d2dcd2611e75 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_optimizers.py @@ -0,0 +1,2316 @@ +# mypy: ignore-errors + +import functools +import itertools +import sys +import unittest +from copy import deepcopy +from enum import Enum +from typing import Any + +import torch +from torch import Tensor +from torch.nn import Parameter +from torch.optim import ( + Adadelta, + Adafactor, + Adagrad, + Adam, + Adamax, + AdamW, + ASGD, + LBFGS, + Muon, + NAdam, + Optimizer, + RAdam, + RMSprop, + Rprop, + SGD, + SparseAdam, +) +from torch.optim.lr_scheduler import ( + ConstantLR, + ExponentialLR, + LinearLR, + PolynomialLR, + ReduceLROnPlateau, + StepLR, +) +from torch.testing._internal.common_device_type import tol, toleranceOverride +from torch.testing._internal.common_methods_invocations import DecorateInfo +from torch.testing._internal.common_utils import ( + _TestParametrizer, + skipIfMPS, + skipIfTorchDynamo, + TEST_WITH_TORCHDYNAMO, +) +from torch.utils._foreach_utils import _get_foreach_kernels_supported_devices + + +CUDA_CONFIG_GPUS = ["cuda", "xpu"] + + +class OptimizerInput: + """Contains args / kwargs to be passed to an optimizer constructor.""" + + __slots__ = ["params", "kwargs", "desc"] + + def __init__( + self, + params: list[Parameter] | list[Tensor] | dict[Any, Any] | list[dict[str, Any]], + kwargs: dict[str, Any], + desc: str = "", + ): + # params can be a list of Tensors OR param_groups OR None + self.params = params + self.kwargs = kwargs + self.desc = desc + + def __repr__(self): + return f"params={self.params}, kwargs={self.kwargs}, desc={self.desc}" + + +class OptimizerErrorEnum(Enum): + """Enumerates when an error is raised when testing optimizers.""" + + CONSTRUCTION_ERROR = 0 + STEP_ERROR = 1 + + +class ErrorOptimizerInput: + """ + An OptimizerInput that will cause the optimizer to throw an error when constructed. + Includes the type and string of the resulting error. + """ + + __slots__ = ["optimizer_error_input", "error_on", "error_type", "error_regex"] + + def __init__( + self, + optimizer_error_input, + *, + error_on=OptimizerErrorEnum.CONSTRUCTION_ERROR, + error_type=RuntimeError, + error_regex="", + ): + self.optimizer_error_input = optimizer_error_input + self.error_on = error_on + self.error_type = error_type + self.error_regex = error_regex + + +class OptimizerInfo: + """Optimizer information to be used in testing.""" + + def __init__( + self, + optim_cls: Optimizer, # Class object for the Optimizer under test + *, + # Function to generate optimizer inputs EXCLUDING params. We delegate params responsibility + # to the test using the OptimizerInfo. OptimizerInput.params is likely None. + # Can optionally take in device to filter out certain unsupported configs + optim_inputs_func, + # Tuple of lambdas to generate LRScheduler instances to run with the optimizer for the + # LRScheduler tests like test_forloop_goes_right_direction with_lrsched. + # We DO NOT expect to thoroughly test LRSchedulers through the optimizers, so not every + # LRScheduler configuration will be included. See test_lrscheduler.py for that instead. + # A few optimizers like SGD and Adam will test more LRSchedulers. + scheduler_inputs=( + [ + lambda opt: StepLR(opt, gamma=0.9, step_size=10), + lambda opt: ReduceLROnPlateau(opt), + ], + ), + # A subset of the global-cliquey flags (fused, foreach, differentiable) the optimizer + # supports. See NOTE: [optimizer kwarg categories] for what global-cliquey means. + supported_impls: tuple[str, ...] = ("foreach", "differentiable"), + # A subset of all flags, signifying which ones were only supported after the + # original optimizer had already been released. aka impls where we need to check BC. + not_og_supported_flags: tuple[str, ...] = ( + "foreach", + "differentiable", + "maximize", + "capturable", + ), + # the optim supports passing in sparse gradients as well as dense grads + supports_sparse: bool = False, + # the optimizer constructor supports passing in capturable as a kwarg + has_capturable_arg: bool = False, + # the optim only supports one config: sparse grads w/ dense params, see SparseAdam + only_supports_sparse_grads: bool = False, + # Tuple of (optimizer kwargs, schedulers_constructors) specifically for sparse tests, + # with especially tuned hyperparameters. These only apply if the optimizer supports + # sparse parameters or grads. + metadata_for_sparse=({}, []), + # the optim supports complex parameters + supports_complex: bool = True, + # whether the optimizer.step() function requires a closure to be passed + step_requires_closure: bool = False, + # whether the optimizer supports per-param options with parameter groups + supports_param_groups: bool = True, + # whether the optimizer supports parameters on multiple devices + supports_multiple_devices: bool = True, + skips=(), # Indicates which tests to skip + decorators=None, # Additional decorators to apply to generated tests + optim_error_inputs_func=None, # Function to generate optim inputs that error + supports_fused_on: tuple[str, ...] = (), + ): + self.optim_cls = optim_cls + self.optim_inputs_func = optim_inputs_func + self.scheduler_inputs = scheduler_inputs + self.supported_impls = supported_impls + self.not_og_supported_flags = not_og_supported_flags + self.supports_sparse = supports_sparse + self.has_capturable_arg = has_capturable_arg + self.metadata_for_sparse = metadata_for_sparse + self.only_supports_sparse_grads = only_supports_sparse_grads + self.supports_complex = supports_complex + self.step_requires_closure = step_requires_closure + self.supports_param_groups = supports_param_groups + self.supports_multiple_devices = supports_multiple_devices + self.decorators = ( + *(decorators if decorators else []), + *(skips if skips else []), + ) + self.optim_error_inputs_func = optim_error_inputs_func + self.supports_fused_on = supports_fused_on + + def get_decorators(self, test_class, test_name, device, dtype, param_kwargs): + result = [] + for decorator in self.decorators: + if isinstance(decorator, DecorateInfo): + if decorator.is_active( + test_class, test_name, device, dtype, param_kwargs + ): + result.extend(decorator.decorators) + else: + result.append(decorator) + return result + + @property + def name(self): + return self.optim_cls.__name__ + + +class optims(_TestParametrizer): + """Decorator for specifying a list of optimizers over which to run a test.""" + + def __init__(self, optim_info_iterable, dtypes=None): + self.optim_info_list = list(optim_info_iterable) + + # optimizers aren't limited to be one dtype as parameters can have different dtypes + # We default to torch.float32, but dtypes should be specified through passed in + # parameters. + self.dtypes = dtypes if dtypes is not None else [torch.float32] + + def _parametrize_test(self, test, generic_cls, device_cls): + if device_cls is None: + raise RuntimeError( + "The @optims decorator is only intended to be used in a device-specific " + "context; use it with instantiate_device_type_tests() instead of " + "instantiate_parametrized_tests()" + ) + + for optim_info, dtype in itertools.product(self.optim_info_list, self.dtypes): + # Construct the test name; device / dtype parts are handled outside. + # See [Note: device and dtype suffix placement] + test_name = optim_info.name + + # Construct parameter kwargs to pass to the test. + param_kwargs = {"optim_info": optim_info, "dtype": dtype} + + try: + + @functools.wraps(test) + def test_wrapper(*args, **kwargs): + return test(*args, **kwargs) + + decorator_fn = functools.partial( + optim_info.get_decorators, + generic_cls.__name__, + test.__name__, + device_cls.device_type, + dtype, + ) + + yield (test_wrapper, test_name, param_kwargs, decorator_fn) + except Exception as ex: + # Provides an error message for debugging before rethrowing the exception + print( + f"Failed to instantiate {test_name} for module {optim_info.name}!" + ) + raise ex + + +# Helper function for generating error inputs for all optimizers, used below. +def get_error_inputs_for_all_optims(device, dtype): + if _get_device_type(device) == "cpu": + # Creating 2D parameters for compatibility with Muon. + sample_param = Parameter(torch.randn(1, 1, device=device, dtype=dtype)) + sample_param2 = Parameter(torch.randn(1, 1, device=device, dtype=dtype)) + return [ + ErrorOptimizerInput( + OptimizerInput( + params=sample_param, + kwargs={}, + desc="invalid param type", + ), + error_type=TypeError, + error_regex="params argument given to the optimizer should be an iterable of Tensors or dicts", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[sample_param, sample_param], + kwargs={}, + desc="a param group cannot have duplicate parameters", + ), + error_type=UserWarning, + error_regex=".*a parameter group with duplicate parameters.*", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[{"params": sample_param}, {"params": sample_param}], + kwargs={}, + desc="duplicate parameters should not occur across param groups either", + ), + error_type=ValueError, + error_regex="some parameters appear in more than one parameter group", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=torch.tensor([0.001, 0.001])), + desc="Tensor lr must be 1-element", + ), + error_type=ValueError, + error_regex="Tensor lr must be 1-element", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[("weight", sample_param), sample_param2], + kwargs={}, + desc="all optimizer params should be with/without names", + ), + error_type=ValueError, + error_regex="all optimizer params should be with/without names. Some param names are missing", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[ + {"params": [sample_param], "lr": 1e-2}, + {"params": [("weight", sample_param2)]}, + ], + kwargs={}, + desc="all optimizer param groups should be with/without names.", + ), + error_type=ValueError, + error_regex="all optimizer param groups should be with/without names. " + "cannot add param group with names to the optimizer", + ), + ] + else: + return [] + + +# ------------------------------------------------------------------------------------------ +# NOTE: [optimizer kwarg categories] +# We categorize optimizer kwargs as 3 types: +# 1. optimizer-specific flags are like amsgrad or rho or beta, flags that are specific to +# algorithms and thus only show up for certain optimizers. There are many of these, so I +# do not bother gathering them all and listing them here. The converse to these would be +# global flags that every optimizer ideally _should_ support. We break global flags into +# 2 further categories and list them all below. +# 2. global-friendly = ["lr", "weight_decay", "maximize", "capturable"] +# global-friendly flags are global flags who play nicely with all other global flags, +# i.e., are mutually exclusive in function. This means that any pair of the following +# flags can be toggled at once (e.g., maximize and weight_decay). Furthermore, any of the +# following flags theoretically can be enabled with ANY other global flag, including the +# cliquey ones (e.g, capturable and foreach). +# 3. global-cliquey = ["foreach", "fused", "differentiable"] +# global-cliquey flags are global flags that do NOT coexist with other cliquey flags, +# usually because they contradict each other in function. For example, one should not flip +# both foreach AND fused to True, because they are two differing performance optimizations +# in which you can only opt into one. +# +# The following optim_inputs_func_* sampling functions only return constructor combinations of +# optimizer-specific and global-friendly flags. This is because we are confident they would mesh +# well with additional kwargs. On the flip side of the same coin, we reserve setting the +# global-cliquey flags to individual tests and fully expect tests to edit OptimizerInput.kwargs. + + +def optim_inputs_func_adadelta(device, dtype=None): + cuda_supported_configs = [ + OptimizerInput(params=None, kwargs={"capturable": True}, desc="capturable"), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "capturable": True}, + desc="capturable with weight decay", + ), + OptimizerInput( + params=None, + kwargs={"lr": torch.tensor(0.001), "capturable": True}, + desc="Tensor lr with capturable", + ), + ] + + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 0.01}, desc="non-default lr"), + OptimizerInput( + params=None, kwargs={"weight_decay": 0.1}, desc="nonzero weight_decay" + ), + OptimizerInput(params=None, kwargs={"maximize": True}, desc="maximize"), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True}, + desc="maximize, weight_decay", + ), + OptimizerInput( + params=None, kwargs={"rho": 0.95, "weight_decay": 0.9}, desc="rho" + ), + ] + (cuda_supported_configs if _get_device_type(device) in CUDA_CONFIG_GPUS else []) + + +def optim_error_inputs_func_adadelta(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, rho=1.1), + desc="rho should be between 0 and 1", + ), + error_type=ValueError, + error_regex="Invalid rho value: 1.1", + ), + ] + return error_inputs + + +def optim_inputs_func_adafactor(device, dtype=None): + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "lr": 0.01}, + desc="nonzero weight_decay", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True}, + desc="maximize", + ), + OptimizerInput( + params=None, + kwargs={"beta2_decay": -1.0}, + desc="non-default beta2_decay", + ), + OptimizerInput( + params=None, + kwargs={"d": 1.5}, + desc="non-default clipping threshold d", + ), + ] + + +def optim_error_inputs_func_adafactor(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + complex_param = torch.rand(2, 3, device=device, dtype=torch.complex64) + complex_param.grad = torch.rand_like(complex_param) + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(eps=(-1e-30, 1e-3)), + desc="epsilon1 should be >= 0", + ), + error_type=ValueError, + error_regex="epsilon1 should be >= 0", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(d=0.0), + desc="invalid d", + ), + error_type=ValueError, + error_regex="Clipping threshold d should be >= 1", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(beta2_decay=0.8), + desc="invalid beta2_decay", + ), + error_type=ValueError, + error_regex="beta2_decay should be <= 0", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[complex_param], + kwargs=dict(), + desc="does not support complex parameters", + ), + error_type=RuntimeError, + error_regex="Adafactor does not support complex parameters", + error_on=OptimizerErrorEnum.STEP_ERROR, + ), + ] + return error_inputs + + +def optim_inputs_func_adagrad(device, dtype=None): + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput( + params=None, kwargs={"weight_decay": 0.1}, desc="nonzero weight_decay" + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True}, + desc="maximize", + ), + OptimizerInput(params=None, kwargs={"lr": 0.1}, desc="non-default lr"), + OptimizerInput( + params=None, + kwargs={"initial_accumulator_value": 0.1, "weight_decay": 0.1}, + desc="initial_accumulator_value", + ), + OptimizerInput( + params=None, + kwargs={"lr": 0.1, "lr_decay": 0.5, "weight_decay": 0.1}, + desc="lr_decay", + ), # TODO: Move out to testing in param_group? + OptimizerInput( + params=None, + kwargs={"lr": torch.tensor(0.001)}, + desc="Tensor lr", + ), + ] + + +def optim_error_inputs_func_adagrad(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, lr_decay=-0.5), + desc="lr_decay must be bigger than 0", + ), + error_type=ValueError, + error_regex="Invalid lr_decay value: -0.5", + ), + ] + if _get_device_type(device) == "cuda": + sample_tensor = torch.empty((), device=device, dtype=dtype) + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=[sample_tensor], + kwargs={"foreach": True, "fused": True}, + desc="`fused` and `foreach` cannot be `True` together", + ), + error_type=RuntimeError, + error_regex="`fused` and `foreach` cannot be `True` together", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[sample_tensor], + kwargs={"fused": True, "differentiable": True}, + desc="`fused` does not support `differentiable`", + ), + error_type=RuntimeError, + error_regex="`fused` does not support `differentiable`", + ), + ] + return error_inputs + + +# TODO: consider tensor LR! See multi_tensor_optimizer_configs in test_optim.py --> tensor LR should work +# with all implementation code paths... +def optim_inputs_func_adam(device, dtype=None): + cuda_supported_configs = [ + OptimizerInput(params=None, kwargs={"capturable": True}, desc="capturable"), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "amsgrad": True, "capturable": True}, + desc="capturable, amsgrad", + ), + OptimizerInput( + params=None, + kwargs={"lr": torch.tensor(0.001), "amsgrad": True, "capturable": True}, + desc="Tensor lr with capturable and amsgrad", + ), + OptimizerInput( + params=None, + kwargs={ + "lr": torch.tensor(0.001), + "betas": (torch.tensor([[[0.9]]]), torch.tensor([[0.99]])), + "amsgrad": True, + "capturable": True, + }, + desc="Tensor lr, Tensor betas, with capturable and amsgrad", + ), + OptimizerInput( + params=None, + kwargs={ + "lr": torch.tensor(0.001), + "betas": (torch.tensor(0.9), torch.tensor(0.99)), + "amsgrad": False, + "capturable": True, + }, + desc="Tensor lr, Tensor betas, with capturable", + ), + ] + mps_supported_configs = [ + OptimizerInput( + params=None, kwargs={"lr": torch.tensor(0.01)}, desc="Tensor lr" + ), + ] + + total = ( + [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 0.01}, desc="non-default lr"), + OptimizerInput( + params=None, kwargs={"weight_decay": 0.1}, desc="nonzero weight_decay" + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True}, + desc="maximize", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "amsgrad": True}, + desc="amsgrad", + ), + ] + + ( + cuda_supported_configs + if _get_device_type(device) in CUDA_CONFIG_GPUS + else [] + ) + + (mps_supported_configs if _get_device_type(device) == "mps" else []) + ) + if dtype == torch.float16: + for input in total: + """ + Too small eps will make denom to be zero for low precision dtype + denom = (exp_avg_sq.sqrt() / bias_correction2_sqrt).add_(eps) + For example, + >>> a + tensor([0.], dtype=torch.float16) + >>> a + 1e-8 + tensor([0.], dtype=torch.float16) + """ + input.kwargs["eps"] = 0.1 + return total + + +def optim_error_inputs_func_adam(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, betas=(1.0, 0.0)), + desc="beta1 should be between 0 and 1", + ), + error_type=ValueError, + error_regex="Invalid beta parameter at index 0: 1.0", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, weight_decay=-1), + desc="weight_decay should > 0", + ), + error_type=ValueError, + error_regex="Invalid weight_decay value: -1", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=torch.tensor(0.001), foreach=True), + desc="lr as Tensor doesn't work with foreach & not capturable", + ), + error_type=ValueError, + error_regex="lr as a Tensor is not supported for capturable=False and foreach=True", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, betas=(0.9, torch.tensor(0.99))), + desc="betas must be either both floats or both Tensors", + ), + error_type=ValueError, + error_regex="betas must be either both floats or both Tensors", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, betas=(torch.tensor(0.9), 0.99)), + desc="betas must be either both floats or both Tensors", + ), + error_type=ValueError, + error_regex="betas must be either both floats or both Tensors", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict( + lr=1e-2, + betas=(torch.tensor(0.9), torch.tensor(0.99)), + foreach=True, + ), + desc=r"betas\[0\] as a Tensor is not supported for capturable=False and foreach=True", + ), + error_type=ValueError, + error_regex=r"betas\[0\] as a Tensor is not supported for capturable=False and foreach=True", + ), + ] + if _get_device_type(device) in CUDA_CONFIG_GPUS: + sample_tensor = torch.empty((), device=device, dtype=dtype) + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=[sample_tensor], + kwargs={"foreach": True, "fused": True}, + desc="`fused` and `foreach` cannot be `True` together", + ), + error_type=RuntimeError, + error_regex="`fused` and `foreach` cannot be `True` together", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[sample_tensor], + kwargs={"fused": True, "differentiable": True}, + desc="`fused` does not support `differentiable`", + ), + error_type=RuntimeError, + error_regex="`fused` does not support `differentiable`", + ), + ] + return error_inputs + + +def optim_inputs_func_adamax(device, dtype=None): + cuda_supported_configs = [ + OptimizerInput(params=None, kwargs={"capturable": True}, desc="capturable"), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.9, "maximize": True, "capturable": True}, + desc="capturable, maximize, weight_decay", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0, "maximize": True, "capturable": True}, + desc="capturable, maximize", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.9, "maximize": False, "capturable": True}, + desc="capturable, weight_decay", + ), + OptimizerInput( + params=None, + kwargs={ + "lr": torch.tensor(0.001), + "weight_decay": 0.9, + "maximize": False, + "capturable": True, + }, + desc="capturable, weight_decay, tensor LR", + ), + ] + + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 0.1}, desc="non-default lr"), + OptimizerInput( + params=None, kwargs={"weight_decay": 0.1}, desc="nonzero weight_decay" + ), + OptimizerInput( + params=None, + kwargs={"maximize": True}, + desc="maximize", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True}, + desc="maximize, weight_decay", + ), + ] + (cuda_supported_configs if _get_device_type(device) in CUDA_CONFIG_GPUS else []) + + +def optim_error_inputs_func_adamax(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, betas=(0.0, 1.0)), + desc="beta2 should be between 0 and 1", + ), + error_type=ValueError, + error_regex="Invalid beta parameter at index 1: 1.0", + ), + ] + return error_inputs + + +def optim_inputs_func_adamw(device, dtype=None): + return optim_inputs_func_adam(device, dtype) + + +def optim_error_inputs_func_adamw(device, dtype): + return optim_error_inputs_func_adam(device, dtype) + + +def optim_inputs_func_asgd(device, dtype=None): + cuda_supported_configs = [ + OptimizerInput(params=None, kwargs={"capturable": True}, desc="capturable"), + OptimizerInput( + params=None, + kwargs={"maximize": True, "capturable": True}, + desc="maximize, capturable", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "capturable": True}, + desc="weight_decay, capturable", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True, "capturable": True}, + desc="maximize, weight_decay, capturable", + ), + OptimizerInput( + params=None, + kwargs={ + "lr": torch.tensor(0.001), + "weight_decay": 0.1, + "maximize": True, + "capturable": True, + }, + desc="maximize, weight_decay, capturable, tensor LR", + ), + ] + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lambd": 0.1}, desc="non-default lambd"), + OptimizerInput(params=None, kwargs={"lr": 0.02}, desc="non-default lr"), + OptimizerInput(params=None, kwargs={"t0": 100}, desc="t0"), + OptimizerInput(params=None, kwargs={"maximize": True}, desc="maximize"), + OptimizerInput( + params=None, kwargs={"weight_decay": 0.1}, desc="nonzero weight_decay" + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True}, + desc="maximize, nonzero weight_decay", + ), + ] + (cuda_supported_configs if _get_device_type(device) in CUDA_CONFIG_GPUS else []) + + +def optim_error_inputs_func_asgd(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, weight_decay=-0.5), + desc="weight_decay should > 0", + ), + error_type=ValueError, + error_regex="Invalid weight_decay value: -0.5", + ), + ] + return error_inputs + + +def optim_inputs_func_lbfgs(device, dtype=None): + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 0.01}, desc="non-default lr"), + OptimizerInput( + params=None, kwargs={"lr": torch.tensor(0.001)}, desc="Tensor lr" + ), + OptimizerInput( + params=None, kwargs={"tolerance_grad": 1e-6}, desc="tolerance_grad" + ), + OptimizerInput( + params=None, + kwargs={"line_search_fn": "strong_wolfe"}, + desc="strong_wolfe", + ), + ] + + +def optim_error_inputs_func_lbfgs(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + return error_inputs + + +def optim_inputs_func_muon(device, dtype=None): + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 0.01}, desc="non-default lr"), + OptimizerInput( + params=None, kwargs={"lr": torch.tensor(0.001)}, desc="Tensor lr" + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.2}, + desc="non-default weight_decay", + ), + OptimizerInput( + params=None, + kwargs={"momentum": 0.8}, + desc="non-default momentum", + ), + OptimizerInput( + params=None, + kwargs={"ns_steps": 6}, + desc="passing alternative ns_steps", + ), + OptimizerInput( + params=None, + kwargs={ + "ns_coefficients": (3.4, -4.7, 2.0), + }, + desc="passing alternative ns_coefficients", + ), + ] + + +def optim_error_inputs_func_muon(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + complex_param = torch.rand(2, 3, device=device, dtype=torch.complex64) + complex_param.grad = torch.rand_like(complex_param) + non_2d_param = torch.rand(2, 3, 4, device=device, dtype=dtype) + non_2d_param.grad = torch.rand_like(non_2d_param) + param = torch.rand(2, 3, device=device, dtype=dtype) + param.grad = torch.rand_like(param) + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=[non_2d_param], + kwargs=dict(), + desc="only support 2D parameters", + ), + error_type=ValueError, + error_regex="Muon only supports 2D parameters", + error_on=OptimizerErrorEnum.CONSTRUCTION_ERROR, + ), + ErrorOptimizerInput( + OptimizerInput( + params=[param], + kwargs={"adjust_lr_fn": "arbitrary"}, + desc="only support `original` and `match_rms_adamw`", + ), + error_type=ValueError, + error_regex="Adjust learning rate function arbitrary is not supported", + error_on=OptimizerErrorEnum.CONSTRUCTION_ERROR, + ), + ErrorOptimizerInput( + OptimizerInput( + params=[complex_param], + kwargs=dict(), + desc="does not support complex parameters", + ), + error_type=RuntimeError, + error_regex="Muon does not support complex parameters", + error_on=OptimizerErrorEnum.STEP_ERROR, + ), + ] + return error_inputs + + +def optim_inputs_func_nadam(device, dtype=None): + cuda_supported_configs = [ + OptimizerInput(params=None, kwargs={"capturable": True}, desc="capturable"), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.9, "momentum_decay": 6e-3, "capturable": True}, + desc="weight_decay, capturable", + ), + OptimizerInput( + params=None, + kwargs={ + "weight_decay": 0.9, + "momentum_decay": 6e-3, + "decoupled_weight_decay": True, + "capturable": True, + }, + desc="decoupled_weight_decay, capturable", + ), + OptimizerInput( + params=None, + kwargs={ + "lr": torch.tensor(0.001), + "weight_decay": 0.9, + "momentum_decay": 6e-3, + "decoupled_weight_decay": True, + "capturable": True, + }, + desc="decoupled_weight_decay, capturable", + ), + ] + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 1e-3}, desc="non-default lr"), + OptimizerInput( + params=None, + kwargs={"momentum_decay": 6e-3}, + desc="non-zero momentum_decay", + ), + OptimizerInput( + params=None, + kwargs={ + "weight_decay": 0.1, + }, + desc="weight_decay", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "momentum_decay": 6e-3}, + desc="weight_decay, momentum_decay", + ), + OptimizerInput( + params=None, + kwargs={ + "weight_decay": 0.1, + "momentum_decay": 6e-3, + "decoupled_weight_decay": True, + }, + desc="decoupled_weight_decay", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True}, + desc="maximize", + ), + ] + (cuda_supported_configs if _get_device_type(device) in CUDA_CONFIG_GPUS else []) + + +def optim_error_inputs_func_nadam(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, betas=(1.0, 0.0)), + desc="beta1 should be between 0 and 1", + ), + error_type=ValueError, + error_regex="Invalid beta parameter at index 0: 1.0", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, momentum_decay=-0.2), + desc="momentum_decay should > 0", + ), + error_type=ValueError, + error_regex="Invalid momentum_decay value: -0.2", + ), + ] + return error_inputs + + +# Weird story bro, NAdam and RAdam do not have maximize. +def optim_inputs_func_radam(device=None, dtype=None): + cuda_supported_configs = [ + OptimizerInput(params=None, kwargs={"capturable": True}, desc="capturable"), + OptimizerInput( + params=None, + kwargs={ + "capturable": True, + "weight_decay": 0.1, + }, + desc="capturable, weight_decay", + ), + OptimizerInput( + params=None, + kwargs={ + "capturable": True, + "weight_decay": 0.1, + "decoupled_weight_decay": True, + }, + desc="capturable, weight_decay, decoupled_weight_decay", + ), + OptimizerInput( + params=None, + kwargs={ + "lr": torch.tensor(0.001), + "capturable": True, + "weight_decay": 0.1, + "decoupled_weight_decay": True, + }, + desc="capturable, weight_decay, decoupled_weight_decay, tensor LR", + ), + ] + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 2e-3}, desc="non-default lr"), + OptimizerInput(params=None, kwargs={"eps": 1e-6}, desc="non-default eps"), + OptimizerInput( + params=None, kwargs={"weight_decay": 0.1}, desc="nonzero weight_decay" + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "decoupled_weight_decay": True}, + desc="decoupled_weight_decay", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True}, + desc="maximize", + ), + ] + (cuda_supported_configs if _get_device_type(device) in CUDA_CONFIG_GPUS else []) + + +def optim_error_inputs_func_radam(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, betas=(1.0, 0.0)), + desc="beta1 should be between 0 and 1", + ), + error_type=ValueError, + error_regex="Invalid beta parameter at index 0: 1.0", + ), + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, weight_decay=-1), + desc="weight_decay should > 0", + ), + error_type=ValueError, + error_regex="Invalid weight_decay value: -1", + ), + ] + return error_inputs + + +def optim_inputs_func_rmsprop(device, dtype=None): + cuda_supported_configs = [ + OptimizerInput(params=None, kwargs={"capturable": True}, desc="capturable"), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True, "capturable": True}, + desc="capturable, maximize", + ), + OptimizerInput( + params=None, + kwargs={"lr": torch.tensor(0.001), "capturable": True}, + desc="Tensor lr with capturable", + ), + ] + + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 1e-3}, desc="non-default lr"), + OptimizerInput( + params=None, kwargs={"weight_decay": 0.1}, desc="nonzero weight_decay" + ), + OptimizerInput( + params=None, + kwargs={ + "maximize": True, + }, + desc="maximize", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "centered": True}, + desc="centered", + ), + OptimizerInput( + params=None, + kwargs={ + "maximize": True, + "weight_decay": 0.1, + }, + desc="maximize, weight_decay", + ), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "centered": True, "momentum": 0.1}, + desc="momentum", + ), + OptimizerInput( + params=None, + kwargs={ + "weight_decay": 0.1, + "centered": True, + "momentum": 0.1, + "maximize": True, + }, + desc="maximize, centered, weight_decay, w/ momentum", + ), + ] + (cuda_supported_configs if _get_device_type(device) in CUDA_CONFIG_GPUS else []) + + +def optim_error_inputs_func_rmsprop(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, momentum=-1.0), + desc="momentum should be between 0 and 1", + ), + error_type=ValueError, + error_regex="Invalid momentum value: -1.0", + ), + ] + return error_inputs + + +def optim_inputs_func_rprop(device, dtype=None): + cuda_supported_configs = [ + OptimizerInput(params=None, kwargs={"capturable": True}, desc="capturable"), + OptimizerInput( + params=None, + kwargs={"lr": torch.tensor(0.001), "capturable": True}, + desc="Tensor lr with capturable", + ), + ] + + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 2e-4}, desc="non-default lr"), + OptimizerInput( + params=None, kwargs={"etas": (0.5, 1.5)}, desc="non-default etas" + ), + OptimizerInput( + params=None, + kwargs={"step_sizes": (2e-6, 100)}, + desc="non-default step_sizes", + ), + OptimizerInput(params=None, kwargs={"maximize": True}, desc="maximize"), + ] + (cuda_supported_configs if _get_device_type(device) in CUDA_CONFIG_GPUS else []) + + +def optim_error_inputs_func_rprop(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, etas=(1.0, 0.5)), + desc="0 < eta1 < 1 < eta2", + ), + error_type=ValueError, + error_regex="Invalid eta values: 1.0, 0.5", + ), + ] + return error_inputs + + +def optim_inputs_func_sgd(device, dtype=None): + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput(params=None, kwargs={"lr": 1e-2}, desc="non-default lr"), + OptimizerInput( + params=None, kwargs={"lr": torch.tensor(0.001)}, desc="tensor lr" + ), + OptimizerInput( + params=None, kwargs={"weight_decay": 0.5}, desc="non-zero weight_decay" + ), + OptimizerInput(params=None, kwargs={"momentum": 0.9}, desc="momentum"), + OptimizerInput( + params=None, + kwargs={"weight_decay": 0.1, "maximize": True}, + desc="maximize", + ), + OptimizerInput( + params=None, + kwargs={"momentum": 0.9, "dampening": 0.5}, + desc="dampening", + ), + OptimizerInput( + params=None, + kwargs={"momentum": 0.9, "weight_decay": 0.1}, + desc="weight_decay w/ momentum", + ), + OptimizerInput( + params=None, + kwargs={"momentum": 0.9, "nesterov": True, "weight_decay": 0.1}, + desc="nesterov", + ), + ] + + +def optim_error_inputs_func_sgd(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, momentum=-0.5), + desc="momentum should be between 0 and 1", + ), + error_type=ValueError, + error_regex="Invalid momentum value: -0.5", + ), + ] + return error_inputs + + +def optim_inputs_func_sparseadam(device, dtype=None): + return [ + OptimizerInput(params=None, kwargs={}, desc="default"), + OptimizerInput( + params=None, kwargs={"lr": 0.01}, desc="non-default lr" + ), # TODO: Move out to testing in param_group? + OptimizerInput( + params=None, kwargs={"lr": torch.tensor(0.001)}, desc="Tensor lr" + ), + OptimizerInput(params=None, kwargs={"maximize": True}, desc="maximize"), + ] + + +def optim_error_inputs_func_sparseadam(device, dtype): + error_inputs = get_error_inputs_for_all_optims(device, dtype) + + if _get_device_type(device) == "cpu": + error_inputs += [ + ErrorOptimizerInput( + OptimizerInput( + params=None, + kwargs=dict(lr=1e-2, betas=(1.0, 0.0)), + desc="beta1 should be between 0 and 1", + ), + error_type=ValueError, + error_regex="Invalid beta parameter at index 0: 1.0", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[ + torch.zeros( + 3, layout=torch.sparse_coo, device=device, dtype=dtype + ) + ], + kwargs={}, + desc="dense params required", + ), + error_type=ValueError, + error_regex="SparseAdam requires dense parameter tensors", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[ + { + "params": [ + torch.zeros( + 3, + layout=torch.sparse_coo, + device=device, + dtype=dtype, + ) + ] + } + ], + kwargs={}, + desc="dense params required in param_groups", + ), + error_type=ValueError, + error_regex="SparseAdam requires dense parameter tensors", + ), + ErrorOptimizerInput( + OptimizerInput( + params=[torch.rand(2, 3, device=device, dtype=torch.complex64)], + kwargs={}, + desc="complex not supported", + ), + error_type=ValueError, + error_regex="SparseAdam does not support complex parameters", + ), + ] + return error_inputs + + +def _get_device_type(device: str | torch.device) -> str: + # Returns the device type as a string, e.g., "cpu" or "cuda" + if isinstance(device, torch.device): + device = str(device.type) + if not isinstance(device, str): + raise AssertionError(f"Expected device to be a str, got {type(device)}") + return device.split(":")[0] + + +def _get_optim_inputs_including_global_cliquey_kwargs( + device, dtype, optim_info, skip=() +) -> list[OptimizerInput]: + """ + Return a list of all configs for a given optimizer as a list of OptimizerInputs, + including configs that have supported global cliquey kwargs (foreach, fused, + differentiable) based on optim_info.supported_impls. + + The configs (optim_inputs) returned by optim_info.optim_inputs_func(...) + intentionally do NOT include global cliquey kwargs to give flexibility to tests. + For example, testing correctness between toggling foreach on and off is now + trivial. That said, we sometimes want to test for all possible configs on an + optimizer including all supported flags, so this helper returns all optim inputs. + """ + if not all(x in ["foreach", "fused", "differentiable"] for x in skip): + raise AssertionError( + "skip must be a subset of ['foreach', 'fused', 'differentiable']" + ) + + optim_inputs = optim_info.optim_inputs_func(device) + + supported_impls = tuple( + x + for x in optim_info.supported_impls + if x not in skip + and (_get_device_type(device) in optim_info.supports_fused_on or x != "fused") + and ( + _get_device_type(device) in _get_foreach_kernels_supported_devices() + or x != "foreach" + ) + ) + + all_optim_inputs = [] + for optim_input in optim_inputs: + # Add the base config where all the flags are False + base_kwargs = deepcopy(optim_input.kwargs) + if len(supported_impls) != 0: + for flag in supported_impls: + base_kwargs[flag] = False + all_optim_inputs.append( + OptimizerInput(params=None, kwargs=base_kwargs, desc=optim_input.desc) + ) + else: + all_optim_inputs.append(optim_input) + # Add a config for when each of the global cliquey kwargs is True + # Note that in [optimizer kwarg categories], these kwargs are mutually + # exclusive, so we do not need to product them together. + for flag in supported_impls: + new_kwargs = deepcopy(base_kwargs) + new_kwargs[flag] = True + all_optim_inputs.append( + OptimizerInput( + params=None, kwargs=new_kwargs, desc=f"{optim_input.desc} & {flag}" + ) + ) + return all_optim_inputs + + +# Database of OptimizerInfo entries in alphabetical order. +optim_db: list[OptimizerInfo] = [ + OptimizerInfo( + Adadelta, + optim_inputs_func=optim_inputs_func_adadelta, + optim_error_inputs_func=optim_error_inputs_func_adadelta, + supported_impls=("foreach", "differentiable"), + has_capturable_arg=True, + skips=( + DecorateInfo( + skipIfTorchDynamo("See #116028"), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + # Note on tolerances: + # test_correctness_Adadelta_cuda_float32 + # Mismatched elements: 10 / 100 (10.0%) + # Greatest absolute difference: 4.838220775127411e-05 at index (7, 4) (up to 1e-05 allowed) + # Greatest relative difference: 0.007270356640219688 at index (7, 2) (up to 1e-05 allowed) + # This is due to floating point ordering error + usage of sqrt + DecorateInfo( + toleranceOverride( + { + torch.float32: tol( + rtol=5.5e-4, + atol=5e-5, + ) + } + ), + "CompiledOptimizerParityTests", + "test_correctness", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + Adafactor, + optim_inputs_func=optim_inputs_func_adafactor, + optim_error_inputs_func=optim_error_inputs_func_adafactor, + supported_impls=("foreach",), + not_og_supported_flags=("foreach",), + supports_complex=False, + skips=( + DecorateInfo( + unittest.skip("See #133268 regarding dtype being None"), + "CompiledOptimizerParityTests", + "test_correctness", + device_type="cuda", + active_if=lambda kwargs: kwargs.get("use_closure", False), + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_can_load_older_state_dict", + device_type="cuda", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_deepcopy_copies_all_public_attrs", + device_type="cuda", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_foreach_large_tensor", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_foreach_matches_forloop", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_load_nontensor_step", + device_type="cuda", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_mixed_device_dtype", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_param_groups_lr", + device_type="cuda", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_param_groups_weight_decay", + device_type="cuda", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_peak_memory_foreach", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_save_load_equality_with_weights_only", + device_type="cuda", + ), + DecorateInfo( + skipIfTorchDynamo("See #116028 regarding copy not supported"), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_state_dict_deterministic", + device_type="cuda", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_step_is_noop_for_zero_grads", + device_type="cuda", + ), + DecorateInfo( + unittest.skip("See #133268 regarding dtype being None"), + "CompiledOptimizerParityTests", + "test_correctness", + device_type="xpu", + active_if=lambda kwargs: kwargs.get("use_closure", False), + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_can_load_older_state_dict", + device_type="xpu", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_deepcopy_copies_all_public_attrs", + device_type="xpu", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_load_nontensor_step", + device_type="xpu", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_param_groups_lr", + device_type="xpu", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_param_groups_weight_decay", + device_type="xpu", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_save_load_equality_with_weights_only", + device_type="xpu", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_state_dict_deterministic", + device_type="xpu", + ), + DecorateInfo( + skipIfTorchDynamo("See #133268 regarding dtype being None"), + "TestOptimRenewed", + "test_step_is_noop_for_zero_grads", + device_type="xpu", + ), + ), + ), + OptimizerInfo( + Adagrad, + optim_inputs_func=optim_inputs_func_adagrad, + optim_error_inputs_func=optim_error_inputs_func_adagrad, + supported_impls=("foreach", "differentiable", "fused"), + not_og_supported_flags=( + "foreach", + "differentiable", + "fused", + "maximize", + "capturable", + ), + supports_fused_on=("cpu", "cuda"), + supports_sparse=True, + metadata_for_sparse=( + {"lr": 0.1, "weight_decay": 0, "lr_decay": 0}, + [ + lambda opt: StepLR(opt, gamma=1 - 1e-5, step_size=500), + lambda opt: ReduceLROnPlateau(opt, threshold=1e-4), + ], + ), + decorators=( + DecorateInfo( + # Note on tolerances: + # difference comes from the fact that the non fused kernel have + # more dtype cast operations. We have another test test_fused_cpu_matches_cuda + # to make sure there is no discrepancies between cuda fused kernel + # and cpu fused kernel + toleranceOverride( + { + torch.bfloat16: tol(atol=5e-3, rtol=5e-3), + torch.float16: tol(atol=5e-3, rtol=5e-3), + } + ), + "TestOptimRenewed", + "test_fused_matches_forloop", + ), + DecorateInfo( + toleranceOverride( + { # https://github.com/pytorch/pytorch/issues/116202 + torch.float32: tol(atol=5e-04, rtol=0.015), + } + ), + "TestOptimRenewed", + "test_mixed_device_dtype", + active_if=TEST_WITH_TORCHDYNAMO, + ), + ), + skips=( + DecorateInfo( + skipIfTorchDynamo("See #116028"), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + Adam, + optim_inputs_func=optim_inputs_func_adam, + scheduler_inputs=( + [lambda opt: ExponentialLR(opt, gamma=0.9)], + [lambda opt: LinearLR(opt, start_factor=0.4, total_iters=4)], + [ + lambda opt: ConstantLR(opt, factor=0.4, total_iters=4), + lambda opt: ExponentialLR(opt, gamma=0.9), + ], + [ + lambda opt: ExponentialLR(opt, gamma=0.9), + lambda opt: ReduceLROnPlateau(opt), + ], + [lambda opt: ConstantLR(opt, factor=0.4, total_iters=4)], + [lambda opt: PolynomialLR(opt, power=0.9, total_iters=4)], + [ + lambda opt: StepLR(opt, gamma=0.9, step_size=10), + lambda opt: ReduceLROnPlateau(opt), + ], + ), + optim_error_inputs_func=optim_error_inputs_func_adam, + supported_impls=("foreach", "differentiable", "fused"), + has_capturable_arg=True, + not_og_supported_flags=( + "foreach", + "differentiable", + "fused", + "maximize", + "capturable", + ), + supports_fused_on=("cpu", "cuda", "xpu", "mps"), + decorators=( + # Expected floating point error between fused and compiled forloop + DecorateInfo( + toleranceOverride({torch.float64: tol(atol=4.5e-7, rtol=2.2e-6)}), + "TestOptimRenewed", + "test_fused_matches_forloop", + active_if=lambda kwargs: TEST_WITH_TORCHDYNAMO + and kwargs["dtype"] == torch.float64, + ), + DecorateInfo( + # Note on tolerances: + # difference comes from the fact that the non fused kernel have + # more dtype cast operations. We have another test test_fused_cpu_matches_cuda + # to make sure there is no discrepancies between cuda fused kernel + # and cpu fused kernel + toleranceOverride( + { + torch.bfloat16: tol(atol=5e-3, rtol=5e-3), + torch.float16: tol(atol=5e-3, rtol=5e-3), + } + ), + "TestOptimRenewed", + "test_fused_matches_forloop", + ), + ), + skips=( + DecorateInfo( + skipIfTorchDynamo( + "Errors w/ Global state changed, see https://github.com/pytorch/pytorch/issues/116028" + ), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + Adamax, + optim_inputs_func=optim_inputs_func_adamax, + optim_error_inputs_func=optim_error_inputs_func_adamax, + supported_impls=("foreach", "differentiable"), + has_capturable_arg=True, + skips=( + DecorateInfo( + skipIfTorchDynamo("See #116028"), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + unittest.skip("Uses too much memory, even for H100, surprisingly."), + "TestOptimRenewed", + "test_foreach_large_tensor", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + AdamW, + optim_inputs_func=optim_inputs_func_adamw, + optim_error_inputs_func=optim_error_inputs_func_adamw, + supported_impls=("foreach", "differentiable", "fused"), + not_og_supported_flags=( + "foreach", + "differentiable", + "fused", + "maximize", + "capturable", + ), + supports_fused_on=("cpu", "cuda", "mps"), + has_capturable_arg=True, + decorators=( + # Expected error between compiled forloop and fused optimizers + DecorateInfo( + toleranceOverride({torch.float64: tol(atol=4.5e-7, rtol=2.2e-6)}), + "TestOptimRenewed", + "test_fused_matches_forloop", + active_if=lambda kwargs: TEST_WITH_TORCHDYNAMO + and kwargs["dtype"] == torch.float64, + ), + DecorateInfo( + toleranceOverride( + # Note on tolerances: + # difference comes from the fact that the non fused kernel have + # more dtype cast operations. We have another test test_fused_cpu_matches_cuda + # to make sure there is no discrepancies between cuda fused kernel + # and cpu fused kernel + { + torch.bfloat16: tol(atol=5e-3, rtol=5e-3), + torch.float16: tol(atol=5e-3, rtol=5e-3), + } + ), + "TestOptimRenewed", + "test_fused_matches_forloop", + ), + ), + skips=( + DecorateInfo( + skipIfTorchDynamo( + "Errors w/ Global state changed, see https://github.com/pytorch/pytorch/issues/116028" + ), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + ASGD, + optim_inputs_func=optim_inputs_func_asgd, + optim_error_inputs_func=optim_error_inputs_func_asgd, + supported_impls=("foreach", "differentiable"), + has_capturable_arg=True, + skips=( + DecorateInfo( + skipIfTorchDynamo( + "Errors w/ Global state changed, see https://github.com/pytorch/pytorch/issues/116028" + ), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + toleranceOverride( + { + torch.float32: tol(atol=1.5e-5, rtol=1e-5), + } + ), + "TestOptimRenewed", + "test_step_is_noop_for_zero_grads", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + DecorateInfo( + unittest.skip( + "ASGD internally changes the weights even with zero grad" + ), + "TestOptimRenewed", + "test_step_is_noop_for_zero_grads", + ), + ), + ), + OptimizerInfo( + LBFGS, + optim_inputs_func=optim_inputs_func_lbfgs, + optim_error_inputs_func=optim_error_inputs_func_lbfgs, + supported_impls=(), + step_requires_closure=True, + supports_param_groups=False, + supports_multiple_devices=False, + skips=( + # Fails on MacOS 13.2.1 in CI https://github.com/pytorch/pytorch/issues/117094 + DecorateInfo( + skipIfMPS, + "TestOptimRenewed", + "test_can_load_older_state_dict", + device_type="mps", + ), + DecorateInfo( + toleranceOverride( + { + torch.complex64: tol( + rtol=4.5e-5, + atol=5e-5, + ) + } + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + unittest.skip("Does not support param groups"), + "TestOptimRenewed", + "test_param_groups_lr", + ), + DecorateInfo( + unittest.skip("Does not support param groups"), + "TestOptimRenewed", + "test_param_groups_weight_decay", + ), + DecorateInfo( + unittest.skip("LBFGS doesn't support multidevice"), + "TestOptimRenewed", + "test_forloop_goes_right_direction_multigpu", + ), + DecorateInfo( + unittest.skip("Does not support param groups"), + "TestOptimRenewed", + "test_param_group_with_lrscheduler_goes_right_direction", + ), + # https://github.com/pytorch/pytorch/issues/131398 + DecorateInfo( + unittest.expectedFailure, + "CompiledOptimizerParityTests", + "test_correctness", + active_if=lambda kwargs: sys.platform == "darwin" + and kwargs["use_closure"], + ), + ), + ), + OptimizerInfo( + Muon, + optim_inputs_func=optim_inputs_func_muon, + optim_error_inputs_func=optim_error_inputs_func_muon, + supported_impls=(), + not_og_supported_flags=(), + supports_complex=False, + skips=( + # Note on numerical differences: `compile` applies different matmul tuning, + # which leads to deviations compared to eager mode. In the Newton-Schulz + # iteration for orthogonalization, computations are done in bfloat16, further + # amplifying these numerical differences. + DecorateInfo( + unittest.skip( + "Expect high difference between compiled and eager due to bfloat16 and iterative process." + ), + "CompiledOptimizerParityTests", + "test_correctness", + ), + ), + ), + OptimizerInfo( + NAdam, + optim_inputs_func=optim_inputs_func_nadam, + optim_error_inputs_func=optim_error_inputs_func_nadam, + supported_impls=("foreach", "differentiable"), + has_capturable_arg=True, + skips=( + DecorateInfo( + skipIfTorchDynamo( + "Errors w/ Global state changed, see https://github.com/pytorch/pytorch/issues/116028" + ), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + skipIfTorchDynamo( + "Errors, https://github.com/pytorch/pytorch/issues/117150" + ), + "TestOptimRenewed", + "test_load_nontensor_step", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + RAdam, + optim_inputs_func=optim_inputs_func_radam, + optim_error_inputs_func=optim_error_inputs_func_radam, + supported_impls=("foreach", "differentiable"), + has_capturable_arg=True, + skips=( + DecorateInfo( + skipIfTorchDynamo( + "Errors w/ Global state changed, see https://github.com/pytorch/pytorch/issues/116028" + ), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + toleranceOverride( + { + # previously atol=1e-7, rtol=1e-7 + torch.float64: tol(atol=1.5e-7, rtol=1.1e-7) + } + ), + "TestOptimRenewed", + "test_foreach_matches_forloop", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + RMSprop, + optim_inputs_func=optim_inputs_func_rmsprop, + optim_error_inputs_func=optim_error_inputs_func_rmsprop, + supported_impls=("foreach", "differentiable"), + has_capturable_arg=True, + skips=( + DecorateInfo( + skipIfTorchDynamo("See #116028"), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + toleranceOverride( + { # previously atol=5-05, rtol=0.001, https://github.com/pytorch/pytorch/issues/116202 + torch.float32: tol(atol=5e-04, rtol=0.01), + } + ), + "TestOptimRenewed", + "test_mixed_device_dtype", + active_if=TEST_WITH_TORCHDYNAMO, + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + Rprop, + optim_inputs_func=optim_inputs_func_rprop, + optim_error_inputs_func=optim_error_inputs_func_rprop, + supported_impls=("foreach", "differentiable"), + has_capturable_arg=True, + skips=( + DecorateInfo( + skipIfTorchDynamo("See #116028"), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + SGD, + optim_inputs_func=optim_inputs_func_sgd, + scheduler_inputs=( + [lambda opt: StepLR(opt, gamma=0.9, step_size=10)], + [ + lambda opt: LinearLR( + opt, start_factor=0.4, end_factor=0.8, total_iters=4 + ) + ], + [ + lambda opt: StepLR(opt, gamma=0.9, step_size=10), + lambda opt: LinearLR( + opt, start_factor=0.4, end_factor=0.6, total_iters=4 + ), + ], + [ + lambda opt: StepLR(opt, gamma=0.99, step_size=10), + lambda opt: ExponentialLR(opt, gamma=0.99), + lambda opt: ReduceLROnPlateau(opt), + ], + [lambda opt: ConstantLR(opt, factor=0.4, total_iters=4)], + [lambda opt: PolynomialLR(opt, power=0.9, total_iters=4)], + [ + lambda opt: StepLR(opt, gamma=0.9, step_size=10), + lambda opt: ReduceLROnPlateau(opt), + ], + ), + optim_error_inputs_func=optim_error_inputs_func_sgd, + supported_impls=("foreach", "differentiable", "fused"), + not_og_supported_flags=( + "foreach", + "differentiable", + "fused", + "maximize", + "capturable", + ), + supports_sparse=True, + metadata_for_sparse=( + { + "lr": 4.8e-3, + "maximize": False, + "momentum": 0, + "nesterov": False, + "weight_decay": 0, + }, + [lambda opt: StepLR(opt, gamma=0.99999, step_size=300)], + ), + supports_fused_on=( + "cpu", + "cuda", + "xpu", + "mps", + ), + skips=( + DecorateInfo( + skipIfTorchDynamo( + "Errors w/ Global state changed, see https://github.com/pytorch/pytorch/issues/116028" + ), + "TestOptimRenewed", + "test_set_default_dtype_works_with_foreach", + ), + DecorateInfo( + skipIfTorchDynamo( + "Accessing grad.real errors, see https://github.com/pytorch/pytorch/issues/117184" + ), + "TestOptimRenewed", + "test_complex_2d", + ), + DecorateInfo( + skipIfTorchDynamo( + "This test uses mocks, which dynamo does not support" + ), + "TestOptimRenewed", + "test_defaults_changed_to_foreach", + ), + ), + ), + OptimizerInfo( + SparseAdam, + optim_inputs_func=optim_inputs_func_sparseadam, + optim_error_inputs_func=optim_error_inputs_func_sparseadam, + supported_impls=(), + only_supports_sparse_grads=True, + metadata_for_sparse=({"lr": 4e-2}, []), + supports_complex=False, # Missing complex support, see #118153 + skips=( + DecorateInfo( + skipIfMPS, # SparseAdam does not support MPS + "TestOptimRenewed", + device_type="mps", + ), + DecorateInfo( + skipIfTorchDynamo("cannot call to_sparse on p.grad, see #117184"), + "TestOptimRenewed", + "test_param_groups_lr", + ), + DecorateInfo( + skipIfTorchDynamo("cannot call to_sparse on p.grad, see #117184"), + "TestOptimRenewed", + "test_tensor_lr", + ), + DecorateInfo( + unittest.skip( + "SparseAdam does not support dense gradients, see #116507" + ), + "TestOptimRenewed", + "test_can_load_older_state_dict", + ), + DecorateInfo( + skipIfTorchDynamo("cannot call to_sparse on p.grad, see #117184"), + "TestOptimRenewed", + "test_load_nontensor_step", + ), + DecorateInfo( + skipIfTorchDynamo("cannot call to_sparse on p.grad, see #117184"), + "TestOptimRenewed", + "test_forloop_goes_right_direction", + ), + DecorateInfo( + skipIfTorchDynamo("cannot call to_sparse on p.grad, see #117184"), + "TestOptimRenewed", + "test_forloop_goes_right_direction_multigpu", + ), + DecorateInfo( + skipIfTorchDynamo("cannot call to_sparse on p.grad, see #117184"), + "TestOptimRenewed", + "test_param_group_with_lrscheduler_goes_right_direction", + ), + DecorateInfo( + skipIfTorchDynamo("cannot call to_sparse on p.grad, see #117184"), + "TestOptimRenewed", + "test_state_dict_with_cuda_params", + ), + DecorateInfo( + skipIfTorchDynamo("cannot call to_sparse on p.grad, see #117184"), + "TestOptimRenewed", + "test_deepcopy_copies_all_public_attrs", + ), + ), + ), +] + + +class TensorTracker: + """ + A utility to track tensor clones in a list, with the expectation of popping them later (in + order) to make fair comparisons between two multi-step computation. The intended use case is + usually when comparing two supposed equal computations, such as an optimizer step that each + individually consists of multiple steps, where numerical deviation could multiply. + + The goal is to be able to compare and align numbers at every milestone so as to minimize + numerical discrepancies, and so when the test fails, it is likely a real problem. + """ + + def __init__(self, assert_eq_kwargs=None): + if assert_eq_kwargs is None: + assert_eq_kwargs = {} + self.assert_eq_kwargs = assert_eq_kwargs + self.tensors = [] + + def add(self, tensor): + """ + Add a detach().clone()'d version of the tensor + """ + self.tensors.append(tensor.detach().clone()) + + # pops from beginning, like a queue and not a stack! + def pop_check_set(self, tensor_to_set, testcase, override_assert_eq_kwargs=None): + """ + Pop the first element in the tensor tracker, assert equality between the popped tensor and + the input tensor, and then set the input tensor to have the same values as the popped tensor + (with copy_). + """ + testcase.assertGreater(len(self.tensors), 0, "no tensors to pop") + ref = self.tensors.pop(0) + + testcase.assertTrue(isinstance(ref, Tensor), f"{type(ref)=}") + assert_eq_kwargs = ( + override_assert_eq_kwargs + if override_assert_eq_kwargs is not None + else self.assert_eq_kwargs + ) + testcase.assertEqual(tensor_to_set, ref, **assert_eq_kwargs) + + with torch.no_grad(): + tensor_to_set.copy_(ref) + + def all_popped(self): + return len(self.tensors) == 0 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_pruning.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_pruning.py new file mode 100644 index 0000000000000000000000000000000000000000..13cd86e05bd6f7b4e9515cf102cc1e6d3b49781d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_pruning.py @@ -0,0 +1,385 @@ +# Owner(s): ["module: unknown"] + +from typing import Any +from torch.ao.pruning import BaseSparsifier +import torch +import torch.nn.functional as F +from torch import nn + +class ImplementedSparsifier(BaseSparsifier): + def __init__(self, **kwargs: dict[str, Any]) -> None: + super().__init__(defaults=kwargs) + + def update_mask(self, module: nn.Module, tensor_name: str, **kwargs: dict[str, Any]) -> None: + module.parametrizations.weight[0].mask[0] = 0 # type: ignore[index, union-attr] + linear_state = self.state['linear1.weight'] + linear_state['step_count'] = linear_state.get('step_count', 0) + 1 + + +class MockSparseLinear(nn.Linear): + """ + This class is a MockSparseLinear class to check convert functionality. + It is the same as a normal Linear layer, except with a different type, as + well as an additional from_dense method. + """ + @classmethod + def from_dense(cls, mod: nn.Linear) -> 'MockSparseLinear': + """ + """ + linear = cls(mod.in_features, + mod.out_features) + return linear + + +def rows_are_subset(subset_tensor: torch.Tensor, superset_tensor: torch.Tensor) -> bool: + """ + Checks to see if all rows in subset tensor are present in the superset tensor + """ + i = 0 + for row in subset_tensor: + while i < len(superset_tensor): + if not torch.equal(row, superset_tensor[i]): + i += 1 + else: + break + else: + return False + return True + + +class SimpleLinear(nn.Module): + r"""Model with only Linear layers without biases, some wrapped in a Sequential, + some following the Sequential. Used to test basic pruned Linear-Linear fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Linear(7, 5, bias=False), + nn.Linear(5, 6, bias=False), + nn.Linear(6, 4, bias=False), + ) + self.linear1 = nn.Linear(4, 4, bias=False) + self.linear2 = nn.Linear(4, 10, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.linear1(x) + x = self.linear2(x) + return x + + +class LinearBias(nn.Module): + r"""Model with only Linear layers, alternating layers with biases, + wrapped in a Sequential. Used to test pruned Linear-Bias-Linear fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Linear(7, 5, bias=True), + nn.Linear(5, 6, bias=False), + nn.Linear(6, 3, bias=True), + nn.Linear(3, 3, bias=True), + nn.Linear(3, 10, bias=False), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + return x + + +class LinearActivation(nn.Module): + r"""Model with only Linear layers, some with bias, some in a Sequential and some following. + Activation functions modules in between each Linear in the Sequential, and each outside layer. + Used to test pruned Linear(Bias)-Activation-Linear fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Linear(7, 5, bias=True), + nn.ReLU(), + nn.Linear(5, 6, bias=False), + nn.Tanh(), + nn.Linear(6, 4, bias=True), + ) + self.linear1 = nn.Linear(4, 3, bias=True) + self.act1 = nn.ReLU() + self.linear2 = nn.Linear(3, 10, bias=False) + self.act2 = nn.Tanh() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.linear1(x) + x = self.act1(x) + x = self.linear2(x) + x = self.act2(x) + return x + + +class LinearActivationFunctional(nn.Module): + r"""Model with only Linear layers, some with bias, some in a Sequential and some following. + Activation functions modules in between each Linear in the Sequential, and functional + activationals are called in between each outside layer. + Used to test pruned Linear(Bias)-Activation-Linear fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Linear(7, 5, bias=True), + nn.ReLU(), + nn.Linear(5, 6, bias=False), + nn.ReLU(), + nn.Linear(6, 4, bias=True), + ) + self.linear1 = nn.Linear(4, 3, bias=True) + self.linear2 = nn.Linear(3, 8, bias=False) + self.linear3 = nn.Linear(8, 10, bias=False) + self.act1 = nn.ReLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.linear1(x) + x = F.relu(x) + x = self.linear2(x) + x = F.relu(x) + x = self.linear3(x) + x = F.relu(x) + return x + + +class SimpleConv2d(nn.Module): + r"""Model with only Conv2d layers, all without bias, some in a Sequential and some following. + Used to test pruned Conv2d-Conv2d fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Conv2d(1, 32, 3, 1, bias=False), + nn.Conv2d(32, 64, 3, 1, bias=False), + ) + self.conv2d1 = nn.Conv2d(64, 48, 3, 1, bias=False) + self.conv2d2 = nn.Conv2d(48, 52, 3, 1, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.conv2d1(x) + x = self.conv2d2(x) + return x + + +class Conv2dBias(nn.Module): + r"""Model with only Conv2d layers, some with bias, some in a Sequential and some outside. + Used to test pruned Conv2d-Bias-Conv2d fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Conv2d(1, 32, 3, 1, bias=True), + nn.Conv2d(32, 32, 3, 1, bias=True), + nn.Conv2d(32, 64, 3, 1, bias=False), + ) + self.conv2d1 = nn.Conv2d(64, 48, 3, 1, bias=True) + self.conv2d2 = nn.Conv2d(48, 52, 3, 1, bias=False) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.conv2d1(x) + x = self.conv2d2(x) + return x + + +class Conv2dActivation(nn.Module): + r"""Model with only Conv2d layers, some with bias, some in a Sequential and some following. + Activation function modules in between each Sequential layer, functional activations called + in-between each outside layer. + Used to test pruned Conv2d-Bias-Activation-Conv2d fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Conv2d(1, 32, 3, 1, bias=True), + nn.ReLU(), + nn.Conv2d(32, 64, 3, 1, bias=True), + nn.Tanh(), + nn.Conv2d(64, 64, 3, 1, bias=False), + nn.ReLU(), + ) + self.conv2d1 = nn.Conv2d(64, 48, 3, 1, bias=False) + self.conv2d2 = nn.Conv2d(48, 52, 3, 1, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.conv2d1(x) + x = F.relu(x) + x = self.conv2d2(x) + x = F.hardtanh(x) + return x + + +class Conv2dPadBias(nn.Module): + r"""Model with only Conv2d layers, all with bias and some with padding > 0, + some in a Sequential and some following. Activation function modules in between each layer. + Used to test that bias is propagated correctly in the special case of + pruned Conv2d-Bias-(Activation)Conv2d fusion, when the second Conv2d layer has padding > 0.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Conv2d(1, 32, 3, 1, padding=1, bias=True), + nn.ReLU(), + nn.Conv2d(32, 32, 3, 1, bias=False), + nn.ReLU(), + nn.Conv2d(32, 32, 3, 1, padding=1, bias=True), + nn.ReLU(), + nn.Conv2d(32, 32, 3, 1, padding=1, bias=True), + nn.ReLU(), + nn.Conv2d(32, 64, 3, 1, bias=True), + nn.Tanh(), + ) + self.conv2d1 = nn.Conv2d(64, 48, 3, 1, padding=1, bias=True) + self.act1 = nn.ReLU() + self.conv2d2 = nn.Conv2d(48, 52, 3, 1, padding=1, bias=True) + self.act2 = nn.Tanh() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.conv2d1(x) + x = self.act1(x) + x = self.conv2d2(x) + x = self.act2(x) + return x + + +class Conv2dPool(nn.Module): + r"""Model with only Conv2d layers, all with bias, some in a Sequential and some following. + Activation function modules in between each layer, Pool2d modules in between each layer. + Used to test pruned Conv2d-Pool2d-Conv2d fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Conv2d(1, 32, kernel_size=3, padding=1, bias=True), + nn.MaxPool2d(kernel_size=2, stride=2, padding=1), + nn.ReLU(), + nn.Conv2d(32, 64, kernel_size=3, padding=1, bias=True), + nn.Tanh(), + nn.AvgPool2d(kernel_size=2, stride=2, padding=1), + ) + self.conv2d1 = nn.Conv2d(64, 48, kernel_size=3, padding=1, bias=True) + self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2, padding=1) + self.af1 = nn.ReLU() + self.conv2d2 = nn.Conv2d(48, 52, kernel_size=3, padding=1, bias=True) + self.conv2d3 = nn.Conv2d(52, 52, kernel_size=3, padding=1, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.conv2d1(x) + x = self.maxpool(x) + x = self.af1(x) + x = self.conv2d2(x) + x = F.avg_pool2d(x, kernel_size=2, stride=2, padding=1) + x = F.relu(x) + x = self.conv2d3(x) + return x + + +class Conv2dPoolFlattenFunctional(nn.Module): + r"""Model with Conv2d layers, all with bias, some in a Sequential and some following, and then a Pool2d + and a functional Flatten followed by a Linear layer. + Activation functions and Pool2ds in between each layer also. + Used to test pruned Conv2d-Pool2d-Flatten-Linear fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Conv2d(1, 3, kernel_size=3, padding=1, bias=True), + nn.MaxPool2d(kernel_size=2, stride=2, padding=1), + nn.ReLU(), + nn.Conv2d(3, 5, kernel_size=3, padding=1, bias=True), + nn.Tanh(), + nn.AvgPool2d(kernel_size=2, stride=2, padding=1), + ) + self.conv2d1 = nn.Conv2d(5, 7, kernel_size=3, padding=1, bias=True) + self.af1 = nn.ReLU() + self.conv2d2 = nn.Conv2d(7, 11, kernel_size=3, padding=1, bias=True) + self.avg_pool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = nn.Linear(11, 13, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.conv2d1(x) + x = F.max_pool2d(x, kernel_size=2, stride=2, padding=1) + x = self.af1(x) + x = self.conv2d2(x) + x = self.avg_pool(x) + x = torch.flatten(x, 1) # test functional flatten + x = self.fc(x) + return x + + +class Conv2dPoolFlatten(nn.Module): + r"""Model with Conv2d layers, all with bias, some in a Sequential and some following, and then a Pool2d + and a Flatten module followed by a Linear layer. + Activation functions and Pool2ds in between each layer also. + Used to test pruned Conv2d-Pool2d-Flatten-Linear fusion.""" + + def __init__(self) -> None: + super().__init__() + self.seq = nn.Sequential( + nn.Conv2d(1, 3, kernel_size=3, padding=1, bias=True), + nn.MaxPool2d(kernel_size=2, stride=2, padding=1), + nn.ReLU(), + nn.Conv2d(3, 5, kernel_size=3, padding=1, bias=True), + nn.Tanh(), + nn.AvgPool2d(kernel_size=2, stride=2, padding=1), + ) + self.conv2d1 = nn.Conv2d(5, 7, kernel_size=3, padding=1, bias=True) + self.af1 = nn.ReLU() + self.conv2d2 = nn.Conv2d(7, 11, kernel_size=3, padding=1, bias=True) + self.avg_pool = nn.AdaptiveAvgPool2d((2, 2)) + self.flatten = nn.Flatten() + self.fc = nn.Linear(44, 13, bias=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.seq(x) + x = self.conv2d1(x) + x = F.max_pool2d(x, kernel_size=2, stride=2, padding=1) + x = self.af1(x) + x = self.conv2d2(x) + x = self.avg_pool(x) + x = self.flatten(x) + x = self.fc(x) + return x + + +class LSTMLinearModel(nn.Module): + """Container module with an encoder, a recurrent module, and a linear.""" + + def __init__( + self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int + ) -> None: + super().__init__() + self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers) + self.linear = nn.Linear(hidden_dim, output_dim) + + def forward(self, input: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + output, _hidden = self.lstm(input) + decoded = self.linear(output) + return decoded, output + + +class LSTMLayerNormLinearModel(nn.Module): + """Container module with an LSTM, a LayerNorm, and a linear.""" + + def __init__( + self, input_dim: int, hidden_dim: int, output_dim: int, num_layers: int + ) -> None: + super().__init__() + self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers) + self.norm = nn.LayerNorm(hidden_dim) + self.linear = nn.Linear(hidden_dim, output_dim) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + x, state = self.lstm(x) + x = self.norm(x) + x = self.linear(x) + return x, state diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_quantization.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_quantization.py new file mode 100644 index 0000000000000000000000000000000000000000..abd69aa09e28c0f4b6b2770de06ce437ce190885 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_quantization.py @@ -0,0 +1,3323 @@ +# mypy: ignore-errors + +r"""Importing this file includes common utility methods and base classes for +checking quantization api and properties of resulting modules. +""" + +import torch +import torch.ao.nn.intrinsic.quantized.dynamic as nniqd +import torch.ao.nn.quantized as nnq +import torch.ao.nn.quantized.dynamic as nnqd +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from functorch.experimental import control_flow +from torch.ao.nn.intrinsic import _FusedModule +from torch.ao.quantization import ( + convert, + default_dynamic_qat_qconfig, + default_dynamic_qconfig, + default_dynamic_quant_observer, + default_embedding_qat_qconfig, + default_observer, + default_per_channel_qconfig, + default_qconfig, + default_symmetric_qnnpack_qat_qconfig, + default_weight_observer, + DeQuantStub, + float_qparams_weight_only_qconfig, + get_default_qat_qconfig, + get_default_qat_qconfig_mapping, + get_default_qconfig, + get_default_qconfig_mapping, + PerChannelMinMaxObserver, + propagate_qconfig_, + QConfig, + QConfigMapping, + quantize, + quantize_dynamic_jit, + quantize_jit, + QuantStub, + QuantType, + QuantWrapper, +) +from torch.ao.quantization.quantization_mappings import ( + get_default_dynamic_quant_module_mappings, + get_default_qat_module_mappings, + get_default_qconfig_propagation_list, +) + +from torch.jit.mobile import _load_for_lite_interpreter +from torch.testing._internal.common_quantized import override_quantized_engine +from torch.testing._internal.common_utils import TEST_WITH_ROCM, TestCase + +try: + from torch.ao.ns.fx.ns_types import NSSingleResultValuesType, NSSubgraph + + # graph mode quantization based on fx + from torch.ao.quantization.quantize_fx import ( + convert_fx, + convert_to_reference_fx, + prepare_fx, + prepare_qat_fx, + ) + from torch.fx import GraphModule + from torch.fx.graph import Node + + HAS_FX = True +except ImportError: + HAS_FX = False + +import copy +import functools +import io +import os + +import unittest +from typing import Any +from collections.abc import Callable + +import numpy as np +import torch._dynamo as torchdynamo +from torch.testing import FileCheck + + +class NodeSpec: + """Used for checking GraphModule Node""" + + def __init__(self, op, target): + """ + op: call_function | call_module + target: + for call_function, target would be a function + for call_module, target would be the type of PyTorch module + """ + self.op = op + self.target = target + + @classmethod + def call_function(cls, target): + return NodeSpec("call_function", target) + + @classmethod + def call_method(cls, target): + return NodeSpec("call_method", target) + + @classmethod + def call_module(cls, target): + return NodeSpec("call_module", target) + + def __hash__(self): + return hash((self.op, self.target)) + + def __eq__(self, other): + if not isinstance(other, NodeSpec): + return NotImplemented + + return self.op == other.op and self.target == other.target + + def __repr__(self): + return repr(self.op) + " " + repr(self.target) + + +def get_supported_device_types(): + return ( + ["cpu", "cuda"] if torch.cuda.is_available() and not TEST_WITH_ROCM else ["cpu"] + ) + + +def test_only_eval_fn(model, calib_data): + r""" + Default evaluation function takes a torch.utils.data.Dataset or a list of + input Tensors and run the model on the dataset + """ + for inp in calib_data: + model(*inp) + + +_default_loss_fn = torch.nn.CrossEntropyLoss() + + +def test_only_train_fn(model, train_data, loss_fn=_default_loss_fn): + r""" + Default train function takes a torch.utils.data.Dataset and train the model + on the dataset + """ + optimizer = torch.optim.Adam(model.parameters(), lr=0.001) + train_loss, correct, total = 0, 0, 0 + for _ in range(10): + model.train() + + for data, target in train_data: + optimizer.zero_grad() + output = model(data) + loss = loss_fn(output, target) + loss.backward() + optimizer.step() + train_loss += loss.item() + _, predicted = torch.max(output, 1) + total += target.size(0) + correct += (predicted == target).sum().item() + return train_loss, correct, total + + +class AverageMeter: + """Computes and stores the average and current value""" + + def __init__(self, name, fmt=":f"): + self.name = name + self.fmt = fmt + self.reset() + + def reset(self): + self.val = 0 + self.avg = 0 + self.sum = 0 + self.count = 0 + + def update(self, val, n=1): + self.val = val + self.sum += val * n + self.count += n + self.avg = self.sum / self.count + + def __str__(self): + fmtstr = "{name} {val" + self.fmt + "} ({avg" + self.fmt + "})" + return fmtstr.format(**self.__dict__) + + +def accuracy(output, target, topk=(1,)): + """Computes the accuracy over the k top predictions for the specified values of k""" + with torch.no_grad(): + maxk = max(topk) + batch_size = target.size(0) + + _, pred = output.topk(maxk, 1, True, True) + pred = pred.t() + correct = pred.eq(target.view(1, -1).expand_as(pred)) + + res = [] + for k in topk: + correct_k = correct[:k].view(-1).float().sum(0, keepdim=True) + res.append(correct_k.mul_(100.0 / batch_size)) + return res + + +def train_one_epoch(model, criterion, optimizer, data_loader, device, ntrain_batches): + model.train() + for cnt, (image, target) in enumerate(data_loader, start=1): + print(".", end="") + image, target = image.to(device), target.to(device) + output = model(image) + loss = criterion(output, target) + optimizer.zero_grad() + loss.backward() + optimizer.step() + accuracy(output, target, topk=(1, 5)) + if cnt >= ntrain_batches: + return + return + + +def ddp_setup(rank, world_size): + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = "12355" + + # initialize the process group + dist.init_process_group("gloo", rank=rank, world_size=world_size) + + +def ddp_cleanup(): + dist.destroy_process_group() + + +def run_ddp(rank, world_size, prepared): + ddp_setup(rank, world_size) + prepared.cuda() + prepared = torch.nn.parallel.DistributedDataParallel(prepared, device_ids=[rank]) + prepared.to(rank) + model_with_ddp = prepared + optimizer = torch.optim.SGD(model_with_ddp.parameters(), lr=0.0001) + train_one_epoch(model_with_ddp, criterion, optimizer, dataset, rank, 1) # noqa: F821 + ddp_cleanup() + + +def convert_dynamic(module): + convert(module, get_default_dynamic_quant_module_mappings(), inplace=True) + + +def prepare_dynamic(model, qconfig_dict=None): + propagate_qconfig_(model, qconfig_dict) + + +def _make_conv_test_input( + batch_size, + in_channels_per_group, + input_feature_map_size, + out_channels_per_group, + groups, + kernel_size, + X_scale, + X_zero_point, + W_scale, + W_zero_point, + use_bias, + use_channelwise, +): + in_channels = in_channels_per_group * groups + out_channels = out_channels_per_group * groups + + (X_value_min, X_value_max) = (0, 4) + X_init = torch.randint( + X_value_min, + X_value_max, + ( + batch_size, + in_channels, + ) + + input_feature_map_size, + ) + X = X_scale * (X_init - X_zero_point).float() + X_q = torch.quantize_per_tensor( + X, scale=X_scale, zero_point=X_zero_point, dtype=torch.quint8 + ) + + W_scale = W_scale * out_channels + W_zero_point = W_zero_point * out_channels + # Resize W_scale and W_zero_points arrays equal to out_channels + W_scale = W_scale[:out_channels] + W_zero_point = W_zero_point[:out_channels] + # For testing, we use small values for weights and for activations so that + # no overflow occurs in vpmaddubsw instruction. If the overflow occurs in + # qconv implementation and if there is no overflow. + # In reference we can't exactly match the results with reference. + # Please see the comment in qconv implementation file + # aten/src/ATen/native/quantized/cpu/qconv.cpp for more details. + (W_value_min, W_value_max) = (-5, 5) + # The operator expects them in the format + # (out_channels, in_channels/groups,) + kernel_size + W_init = torch.randint( + W_value_min, + W_value_max, + ( + out_channels, + in_channels_per_group, + ) + + kernel_size, + ) + b_init = torch.randint(0, 10, (out_channels,)) + + if use_channelwise: + W_shape = (-1, 1) + (1,) * len(kernel_size) + W_scales_tensor = torch.tensor(W_scale, dtype=torch.float) + W_zero_points_tensor = torch.tensor(W_zero_point, dtype=torch.float) + W = ( + W_scales_tensor.reshape(*W_shape) + * (W_init.float() - W_zero_points_tensor.reshape(*W_shape)).float() + ) + b = X_scale * W_scales_tensor * b_init.float() + W_q = torch.quantize_per_channel( + W, + W_scales_tensor.double(), + W_zero_points_tensor.long(), + 0, + dtype=torch.qint8, + ) + else: + W = W_scale[0] * (W_init - W_zero_point[0]).float() + b = X_scale * W_scale[0] * b_init.float() + W_q = torch.quantize_per_tensor( + W, scale=W_scale[0], zero_point=W_zero_point[0], dtype=torch.qint8 + ) + + return (X, X_q, W, W_q, b if use_bias else None) + + +def _make_conv_add_extra_input_tensor(scale, zero_point, sizes): + (X_value_min, X_value_max) = (0, 4) + X_init = torch.randint( + X_value_min, + X_value_max, + sizes, # Infer the size of tensor to do the add + ) + X = scale * (X_init - zero_point).float() + X_q = torch.quantize_per_tensor( + X, scale=scale, zero_point=zero_point, dtype=torch.quint8 + ) + return X, X_q + + +def skipIfNoFBGEMM(fn): + reason = "Quantized operations require FBGEMM. FBGEMM is only optimized for CPUs with instruction set support AVX2 or newer." + if isinstance(fn, type): + if "fbgemm" not in torch.backends.quantized.supported_engines: + fn.__unittest_skip__ = True + fn.__unittest_skip_why__ = reason + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if "fbgemm" not in torch.backends.quantized.supported_engines: + raise unittest.SkipTest(reason) + else: + fn(*args, **kwargs) + + return wrapper + + +def skipIfNoQNNPACK(fn): + reason = "Quantized operations require QNNPACK." + if isinstance(fn, type): + if "qnnpack" not in torch.backends.quantized.supported_engines: + fn.__unittest_skip__ = True + fn.__unittest_skip_why__ = reason + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if "qnnpack" not in torch.backends.quantized.supported_engines: + raise unittest.SkipTest(reason) + else: + fn(*args, **kwargs) + + return wrapper + + +def withQNNPACKBackend(fn): + # TODO(future PR): consider combining with skipIfNoQNNPACK, + # will require testing of existing callsites + reason = "Quantized operations require QNNPACK." + if isinstance(fn, type): + if "qnnpack" not in torch.backends.quantized.supported_engines: + fn.__unittest_skip__ = True + fn.__unittest_skip_why__ = reason + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if "qnnpack" not in torch.backends.quantized.supported_engines: + raise unittest.SkipTest(reason) + with override_quantized_engine("qnnpack"): + fn(*args, **kwargs) + + return wrapper + + +def skipIfNoONEDNN(fn): + reason = "Quantized operations require ONEDNN." + if isinstance(fn, type): + if "onednn" not in torch.backends.quantized.supported_engines: + fn.__unittest_skip__ = True + fn.__unittest_skip_why__ = reason + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if "onednn" not in torch.backends.quantized.supported_engines: + raise unittest.SkipTest(reason) + else: + fn(*args, **kwargs) + + return wrapper + + +def skipIfNoONEDNNBF16(fn): + reason = "Quantized operations require BF16 support." + if isinstance(fn, type): + if not torch.ops.mkldnn._is_mkldnn_bf16_supported(): + fn.__unittest_skip__ = True + fn.__unittest_skip_why__ = reason + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if not torch.ops.mkldnn._is_mkldnn_bf16_supported(): + raise unittest.SkipTest(reason) + else: + fn(*args, **kwargs) + + return wrapper + + +def skipIfNoX86(fn): + reason = "Quantized operations require X86." + if isinstance(fn, type): + if "x86" not in torch.backends.quantized.supported_engines: + fn.__unittest_skip__ = True + fn.__unittest_skip_why__ = reason + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if "x86" not in torch.backends.quantized.supported_engines: + raise unittest.SkipTest(reason) + else: + fn(*args, **kwargs) + + return wrapper + + +def skipIfNoDynamoSupport(fn): + reason = "dynamo doesn't support." + if isinstance(fn, type): + if not torchdynamo.is_dynamo_supported(): + fn.__unittest_skip__ = True + fn.__unittest_skip_why__ = reason + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if not torchdynamo.is_dynamo_supported(): + raise unittest.SkipTest(reason) + else: + fn(*args, **kwargs) + + return wrapper + + +def skipIfNoInductorSupport(fn): + reason = "inductor doesn't support." + if isinstance(fn, type): + if not torchdynamo.is_inductor_supported(): + fn.__unittest_skip__ = True + fn.__unittest_skip_why__ = reason + return fn + + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if not torchdynamo.is_inductor_supported(): + raise unittest.SkipTest(reason) + else: + fn(*args, **kwargs) + + return wrapper + + +try: + import torchvision # noqa: F401 + + HAS_TORCHVISION = True +except ImportError: + HAS_TORCHVISION = False +skip_if_no_torchvision = unittest.skipIf(not HAS_TORCHVISION, "no torchvision") + + +def get_script_module(model, tracing, data): + return torch.jit.trace(model, data) if tracing else torch.jit.script(model) + + +def lengths_to_offsets(t, offset_type=np.int64, use_begin_offset=True): + """ + Convert lengths to offsets for embedding_bag + """ + tt = np.zeros((t.shape[0] + 1,), dtype=offset_type) + tt[1:] = t + tt = torch.from_numpy(np.cumsum(tt, dtype=offset_type)) + if use_begin_offset: + return tt[:-1] + return tt[1:] + + +def _group_quantize_tensor(w, n_bit=4, q_group_size=16): + if w.dim() != 2: + raise AssertionError(f"expected w.dim() == 2, got {w.dim()}") + w = w.transpose(0, 1).contiguous() + if q_group_size <= 1: + raise AssertionError(f"expected q_group_size > 1, got {q_group_size}") + if w.shape[-1] % q_group_size != 0: + raise AssertionError( + f"expected w.shape[-1] % q_group_size == 0, got w.shape[-1]={w.shape[-1]}, q_group_size={q_group_size}" + ) + + to_quant = w.reshape(-1, q_group_size) + if torch.isnan(to_quant).sum() != 0: + raise AssertionError("to_quant contains NaN values") + + max_val = to_quant.amax(dim=1, keepdim=True) + min_val = to_quant.amin(dim=1, keepdim=True) + max_int = 2**n_bit - 1 + min_int = 0 + scales = (max_val - min_val).clamp(min=1e-6) / max_int + if torch.isnan(scales).sum() != 0: + raise AssertionError("scales contains NaN values") + + zeros = min_val + scales * (2 ** (n_bit - 1)) + if torch.isnan(zeros).sum() != 0: + raise AssertionError("zeros contains NaN values") + + out = to_quant.sub(min_val).div(scales).round().clamp_(min_int, max_int) + if torch.isnan(out).sum() != 0: + raise AssertionError("out contains NaN values") + + out = out.to(dtype=torch.int32).reshape(w.shape) + if out.device != torch.device("cpu"): + out = (out[::, ::2] << 4 | out[::, 1::2]).to(torch.uint8) + + # Scales and zeros for the same q-group should be contiguous, so we can + # load as a 32-bit word + scales = scales.view(w.shape[0], -1) + zeros = zeros.view(w.shape[0], -1) + scales_and_zeros = ( + torch.cat( + [ + scales.reshape(scales.size(0), scales.size(1), 1), + zeros.reshape(zeros.size(0), zeros.size(1), 1), + ], + 2, + ) + .transpose(0, 1) + .contiguous() + ) + + return out, scales_and_zeros + + +def _group_quantize_tensor_symmetric(w, n_bit=4, groupsize=32): + # W is of shape [K x N] + # We transpose W as Quantization is applied on [N x K] + w = w.transpose(0, 1).contiguous() + if w.dim() != 2: + raise AssertionError(f"Expected w.dim() == 2, got {w.dim()}") + if groupsize <= 1: + raise AssertionError(f"Expected groupsize > 1, got {groupsize}") + if w.shape[-1] % groupsize != 0: + raise AssertionError(f"Expected w.shape[-1] % groupsize == 0, got {w.shape[-1]} % {groupsize}") + # Calculate scale and zeros + to_quant = w.reshape(-1, groupsize) + max_val = to_quant.abs().amax(dim=1, keepdim=True) + eps = torch.finfo(max_val.dtype).eps + max_int = 2 ** (n_bit - 1) - 1 # For 4-bit, this is 7 + scales = max_val.clamp(min=eps) / max_int + zeros = torch.zeros_like(scales) + + # Quantize the weight + scales = scales.to(torch.float32).reshape(w.shape[0], -1) + zeros = zeros.to(torch.float32).reshape(w.shape[0], -1) + scales = scales.reshape(-1, 1) + zeros = zeros.reshape(-1, 1) + max_int = 2**n_bit - 1 + w_int8 = to_quant.div(scales).add(8.5).to(torch.int8).clamp(max=max_int) + # We pack 2 signed int4 values in unsigned uint8 container. + # This reduces the weight size by half and improves load perf + out_uint8 = (w_int8[::, 1::2] << 4 | w_int8[::, ::2]).to(torch.uint8) + + scales_and_zeros = scales.squeeze().contiguous() + + return out_uint8, scales_and_zeros + + +def _dynamically_quantize_per_channel(x, quant_min, quant_max, target_dtype): + # source: https://github.com/meta-pytorch/gpt-fast/blob/main/quantize.py + # default setup for affine quantization of activations + x_dtype = x.dtype + x = x.float() + eps = torch.finfo(torch.float32).eps + + # get min and max + min_val, max_val = torch.aminmax(x, dim=1) + + # calculate scales and zero_points based on min and max + # reference: https://fburl.com/code/srbiybme + min_val_neg = torch.min(min_val, torch.zeros_like(min_val)) + max_val_pos = torch.max(max_val, torch.zeros_like(max_val)) + device = min_val_neg.device + + # reference: https://fburl.com/code/4wll53rk + max_val_pos = torch.max(-min_val_neg, max_val_pos) + scales = max_val_pos / (float(quant_max - quant_min) / 2) + # ensure scales is the same dtype as the original tensor + scales = torch.clamp(scales, min=eps).to(x.dtype) + zero_points = torch.zeros(min_val_neg.size(), dtype=torch.int64, device=device) + + # quantize based on qmin/qmax/scales/zp + x_div = x / scales.unsqueeze(-1) + x_round = torch.round(x_div) + x_zp = x_round + zero_points.unsqueeze(-1) + quant = torch.clamp(x_zp, quant_min, quant_max).to(target_dtype) + + return quant, scales.to(x_dtype), zero_points + + +# QuantizationTestCase used as a base class for testing quantization on modules +class QuantizationTestCase(TestCase): + def setUp(self): + super().setUp() + self.calib_data = [[torch.rand(2, 5, dtype=torch.float)] for _ in range(2)] + self.train_data = [ + [ + torch.rand(2, 5, dtype=torch.float), + torch.randint(0, 1, (2,), dtype=torch.long), + ] + for _ in range(2) + ] + self.img_data_1d = [[torch.rand(2, 3, 10, dtype=torch.float)] for _ in range(2)] + self.img_data_2d = [ + [torch.rand(1, 3, 10, 10, dtype=torch.float)] for _ in range(2) + ] + self.img_data_3d = [ + [torch.rand(1, 3, 5, 5, 5, dtype=torch.float)] for _ in range(2) + ] + self.img_data_1d_train = [ + [ + torch.rand(2, 3, 10, dtype=torch.float), + torch.randint(0, 1, (1,), dtype=torch.long), + ] + for _ in range(2) + ] + self.img_data_2d_train = [ + [ + torch.rand(1, 3, 10, 10, dtype=torch.float), + torch.randint(0, 1, (1,), dtype=torch.long), + ] + for _ in range(2) + ] + self.img_data_3d_train = [ + [ + torch.rand(1, 3, 5, 5, 5, dtype=torch.float), + torch.randint(0, 1, (1,), dtype=torch.long), + ] + for _ in range(2) + ] + + self.img_data_dict = { + 1: self.img_data_1d, + 2: self.img_data_2d, + 3: self.img_data_3d, + } + + # Quant types that produce statically quantized ops + self.static_quant_types = [QuantType.STATIC, QuantType.QAT] + # All quant types for (fx based) graph mode quantization + self.all_quant_types = [QuantType.DYNAMIC, QuantType.STATIC, QuantType.QAT] + + def checkNoPrepModules(self, module): + r"""Checks the module does not contain child + modules for quantization preparation, e.g. + quant, dequant and observer + """ + self.assertFalse(hasattr(module, "quant")) + self.assertFalse(hasattr(module, "dequant")) + + def checkNoQconfig(self, module): + r"""Checks the module does not contain qconfig""" + self.assertFalse(hasattr(module, "qconfig")) + + for child in module.children(): + self.checkNoQconfig(child) + + def checkHasPrepModules(self, module): + r"""Checks the module contains child + modules for quantization preparation, e.g. + quant, dequant and observer + """ + self.assertTrue(hasattr(module, "module")) + self.assertTrue(hasattr(module, "quant")) + self.assertTrue(hasattr(module, "dequant")) + + def checkObservers( + self, module, propagate_qconfig_list=None, prepare_custom_config_dict=None + ): + r"""Checks the module or module's leaf descendants + have observers in preparation for quantization + """ + if propagate_qconfig_list is None: + propagate_qconfig_list = get_default_qconfig_propagation_list() + if prepare_custom_config_dict is None: + prepare_custom_config_dict = {} + float_to_observed_module_class_mapping = prepare_custom_config_dict.get( + "float_to_observed_custom_module_class", {} + ) + + # check if a module is a leaf module, ignoring activation_post_process attribute + def is_leaf_module(module): + submodule_name_count = 0 + for name, _ in module.named_children(): + if name != "activation_post_process": + submodule_name_count += 1 + return submodule_name_count == 0 + + if ( + hasattr(module, "qconfig") + and module.qconfig is not None + and ( + ( + is_leaf_module(module) + and not isinstance(module, torch.nn.Sequential) + and type(module) in propagate_qconfig_list + ) + or type(module) in float_to_observed_module_class_mapping + ) + and not isinstance(module, torch.ao.quantization.DeQuantStub) + ): + self.assertTrue( + hasattr(module, "activation_post_process"), + "module: " + str(type(module)) + " do not have observer", + ) + # we don't need to check observers for child modules of the + # qat modules + if ( + type(module) not in get_default_qat_module_mappings().values() + and type(module) not in float_to_observed_module_class_mapping.values() + and not isinstance(module, _FusedModule) + ): + for child in module.children(): + if type(child) is nn.Dropout: + continue + self.checkObservers( + child, propagate_qconfig_list, prepare_custom_config_dict + ) + + def checkQuantDequant(self, mod): + r"""Checks that mod has nn.Quantize and + nn.DeQuantize submodules inserted + """ + self.assertEqual(type(mod.quant), nnq.Quantize) + self.assertEqual(type(mod.dequant), nnq.DeQuantize) + + def checkWrappedQuantizedLinear(self, mod): + r"""Checks that mod has been swapped for an nnq.Linear + module, the bias is qint32, and that the module + has Quantize and DeQuantize submodules + """ + self.assertEqual(type(mod.module), nnq.Linear) + self.checkQuantDequant(mod) + + def checkQuantizedLinear(self, mod): + self.assertEqual(type(mod), nnq.Linear) + + def checkDynamicQuantizedLinear(self, mod, dtype): + r"""Checks that mod has been swapped for an nnqd.Linear + module, the bias is float. + """ + self.assertEqual(type(mod), nnqd.Linear) + self.assertEqual(mod._packed_params.dtype, dtype) + + def checkDynamicQuantizedLinearRelu(self, mod, dtype): + r"""Checks that mod has been swapped for an nnqd.Linear + module, the bias is float. + """ + self.assertEqual(type(mod), nniqd.LinearReLU) + self.assertEqual(mod._packed_params.dtype, dtype) + + def check_eager_serialization(self, ref_model, loaded_model, x): + # Check state dict serialization and torch.save APIs + model_dict = ref_model.state_dict() + b = io.BytesIO() + torch.save(model_dict, b) + b.seek(0) + # weights_only=False as we sometimes get a ScriptObject here (weird) + loaded_dict = torch.load(b, weights_only=False) + loaded_model.load_state_dict(loaded_dict) + ref_out = ref_model(*x) + load_out = loaded_model(*x) + + def check_outputs(ref_out, load_out): + self.assertEqual(ref_out[0], load_out[0]) + if isinstance(ref_out[1], tuple): + self.assertEqual(ref_out[1][0], load_out[1][0]) + self.assertEqual(ref_out[1][1], load_out[1][1]) + else: + self.assertEqual(ref_out[1], load_out[1]) + + check_outputs(ref_out, load_out) + b = io.BytesIO() + torch.save(ref_model, b) + b.seek(0) + # weights_only=False as this is legacy code that saves the model + loaded = torch.load(b, weights_only=False) + load_out = loaded(*x) + check_outputs(ref_out, load_out) + + def check_weight_bias_api(self, ref_model, weight_keys, bias_keys): + weight = ref_model.get_weight() + bias = ref_model.get_bias() + self.assertEqual(weight_keys ^ weight.keys(), set()) + self.assertEqual(bias_keys ^ bias.keys(), set()) + + def checkDynamicQuantizedLSTM(self, mod, reference_module_type, dtype): + r"""Checks that mod has been swapped for an nnqd.LSTM type + module, the bias is float. + """ + wt_dtype_map = { + torch.qint8: "quantized_dynamic", + torch.float16: "quantized_fp16", + } + self.assertEqual(type(mod), reference_module_type) + for packed_params in mod._all_weight_values: + self.assertEqual( + packed_params.param.__getstate__()[0][0], wt_dtype_map[dtype] + ) + + def checkLinear(self, mod): + self.assertEqual(type(mod), torch.nn.Linear) + + def checkDynamicQuantizedModule(self, mod, reference_module_type, dtype): + r"""Checks that mod has been swapped for an nnqd.Linear + module, the bias is float. + """ + wt_dtype_map = { + torch.qint8: "quantized_dynamic", + torch.float16: "quantized_fp16", + } + self.assertEqual(type(mod), reference_module_type) + if hasattr(mod, "_all_weight_values"): + for packed_params in mod._all_weight_values: + self.assertEqual( + packed_params.param.__getstate__()[0][0], wt_dtype_map[dtype] + ) + + def checkScriptable(self, orig_mod, calib_data, check_save_load=False): + scripted = torch.jit.script(orig_mod) + self._checkScriptable(orig_mod, scripted, calib_data, check_save_load) + + # Use first calib_data entry as trace input + traced = torch.jit.trace(orig_mod, calib_data[0]) + self._checkScriptable(orig_mod, traced, calib_data, check_save_load) + + # Call this twice: once for a scripted module and once for a traced module + def _checkScriptable(self, orig_mod, script_mod, calib_data, check_save_load): + self._checkModuleCorrectnessAgainstOrig(orig_mod, script_mod, calib_data) + + # Test save/load + buffer = io.BytesIO() + torch.jit.save(script_mod, buffer) + + buffer.seek(0) + loaded_mod = torch.jit.load(buffer) + # Pending __get_state_ and __set_state__ support + # See tracking task https://github.com/pytorch/pytorch/issues/23984 + if check_save_load: + self._checkModuleCorrectnessAgainstOrig(orig_mod, loaded_mod, calib_data) + + def _checkModuleCorrectnessAgainstOrig(self, orig_mod, test_mod, calib_data): + for inp in calib_data: + ref_output = orig_mod(*inp) + scripted_output = test_mod(*inp) + self.assertEqual(scripted_output, ref_output) + + def checkGraphModeOp( + self, + module, + inputs, + quantized_op, + tracing=False, + debug=False, + check=True, + eval_mode=True, + dynamic=False, + qconfig=None, + ): + if debug: + print("Testing:", str(module)) + qconfig_dict = {"": get_default_qconfig(torch.backends.quantized.engine)} + + if eval_mode: + module = module.eval() + if dynamic: + qconfig_dict = {"": default_dynamic_qconfig if qconfig is None else qconfig} + model = get_script_module(module, tracing, inputs[0]).eval() + if debug: + print("input graph:", model.graph) + models = {} + outputs = {} + for debug in [True, False]: + if dynamic: + models[debug] = quantize_dynamic_jit(model, qconfig_dict, debug=debug) + # make sure it runs + outputs[debug] = models[debug](inputs) + else: + # module under test can contain in-place ops, and we depend on + # input data staying constant for comparisons + inputs_copy = copy.deepcopy(inputs) + models[debug] = quantize_jit( + model, + qconfig_dict, + test_only_eval_fn, + [inputs_copy], + inplace=False, + debug=debug, + ) + # make sure it runs + outputs[debug] = models[debug](*inputs[0]) + + if debug: + print("debug graph:", models[True].graph) + print("non debug graph:", models[False].graph) + + if check: + # debug and non-debug option should have the same numerics + self.assertEqual(outputs[True], outputs[False]) + + # non debug graph should produce quantized op + FileCheck().check(quantized_op).run(models[False].graph) + + return models[False] + + def checkGraphModuleNodes( + self, + graph_module, + expected_node=None, + expected_node_occurrence=None, + expected_node_list=None, + ): + """Check if GraphModule contains the target node + Args: + graph_module: the GraphModule instance we want to check + expected_node, expected_node_occurrence, expected_node_list: + see docs for checkGraphModeFxOp + """ + nodes_in_graph = {} + node_list = [] + modules = dict(graph_module.named_modules(remove_duplicate=False)) + for node in graph_module.graph.nodes: + n = None + if node.op == "call_function" or node.op == "call_method": + n = NodeSpec(node.op, node.target) + elif node.op == "call_module": + n = NodeSpec(node.op, type(modules[node.target])) + + if n is not None: + node_list.append(n) + if n in nodes_in_graph: + nodes_in_graph[n] += 1 + else: + nodes_in_graph[n] = 1 + + if expected_node is not None: + self.assertTrue( + expected_node in nodes_in_graph, + "node:" + str(expected_node) + " not found in the graph module", + ) + + if expected_node_occurrence is not None: + for expected_node, occurrence in expected_node_occurrence.items(): + if occurrence != 0: + self.assertTrue( + expected_node in nodes_in_graph, + "Check failed for node:" + str(expected_node) + " not found", + ) + self.assertTrue( + nodes_in_graph[expected_node] == occurrence, + "Check failed for node:" + + str(expected_node) + + " Expected occurrence:" + + str(occurrence) + + " Found occurrence:" + + str(nodes_in_graph[expected_node]), + ) + else: + self.assertTrue( + expected_node not in nodes_in_graph, + "Check failed for node:" + + str(expected_node) + + " expected no occurrence but found", + ) + + if expected_node_list is not None: + cur_index = 0 + for n in node_list: + if cur_index == len(expected_node_list): + return + if n == expected_node_list[cur_index]: + cur_index += 1 + self.assertTrue( + cur_index == len(expected_node_list), + "Check failed for graph:" + + self.printGraphModule(graph_module, print_str=False) + + "Expected ordered list:" + + str(expected_node_list), + ) + + def printGraphModule(self, graph_module, print_str=True): + modules = dict(graph_module.named_modules(remove_duplicate=False)) + node_infos = [] + for n in graph_module.graph.nodes: + node_info = " ".join(map(repr, [n.op, n.name, n.target, n.args, n.kwargs])) + if n.op == "call_module": + node_info += " module type: " + repr(type(modules[n.target])) + node_infos.append(node_info) + str_to_print = "\n".join(node_infos) + if print_str: + print(str_to_print) + return str_to_print + + if HAS_FX: + + def assert_types_for_matched_subgraph_pairs( + self, + matched_subgraph_pairs: dict[str, tuple[NSSubgraph, NSSubgraph]], + expected_types: dict[ + str, tuple[tuple[Callable, Callable], tuple[Callable, Callable]] + ], + gm_a: GraphModule, + gm_b: GraphModule, + ) -> None: + """ + Verifies that the types specified in expected_types match + the underlying objects pointed to by the nodes in matched_subgraph_pairs. + + An example successful test case: + + matched_subgraph_pairs = {'x0': (graph_a_conv_0_node, graph_b_conv_0_node)} + expected_types = {'x0': (nn.Conv2d, nnq.Conv2d)} + + The function tests for key equivalence, and verifies types with + instance checks. + """ + + def _get_underlying_op_type( + node: Node, gm: GraphModule + ) -> Callable | str: + if node.op == "call_module": + mod = getattr(gm, node.target) + return type(mod) + else: + if node.op not in ("call_function", "call_method"): + raise AssertionError(f"Expected node.op in ('call_function', 'call_method'), got {node.op!r}") + return node.target + + self.assertTrue( + len(matched_subgraph_pairs) == len(expected_types), + f"Expected length of results to match, but got {len(matched_subgraph_pairs)} and {len(expected_types)}", + ) + for k, v in expected_types.items(): + expected_types_a, expected_types_b = v + exp_type_start_a, exp_type_end_a = expected_types_a + exp_type_start_b, exp_type_end_b = expected_types_b + subgraph_a, subgraph_b = matched_subgraph_pairs[k] + + act_type_start_a = _get_underlying_op_type(subgraph_a.start_node, gm_a) + act_type_start_b = _get_underlying_op_type(subgraph_b.start_node, gm_b) + act_type_end_a = _get_underlying_op_type(subgraph_a.end_node, gm_a) + act_type_end_b = _get_underlying_op_type(subgraph_b.end_node, gm_b) + types_match = ( + (exp_type_start_a is act_type_start_a) + and (exp_type_end_a is act_type_end_a) + and (exp_type_start_b is act_type_start_b) + and (exp_type_end_b is act_type_end_b) + ) + self.assertTrue( + types_match, + f"Type mismatch at {k}: expected {(exp_type_start_a, exp_type_end_a, exp_type_start_b, exp_type_end_b)}, " + f"got {(act_type_start_a, act_type_end_a, act_type_start_b, act_type_end_b)}", + ) + + def assert_ns_compare_dict_valid( + self, + act_compare_dict: dict[str, dict[str, dict[str, Any]]], + ) -> None: + """ + Verifies that the act_compare_dict (output of Numeric Suite APIs) is valid: + 1. for each layer, results are recorded for two models + 2. number of seen tensors match + 3. shapes of each pair of seen tensors match + """ + for layer_name, result_type_to_data in act_compare_dict.items(): + for result_type, layer_data in result_type_to_data.items(): + self.assertTrue( + len(layer_data) == 2, + f"Layer {layer_name} does not have exactly two model results.", + ) + model_name_0, model_name_1 = layer_data.keys() + for res_idx in range(len(layer_data[model_name_0])): + layer_data_0 = layer_data[model_name_0][res_idx] + layer_data_1 = layer_data[model_name_1][res_idx] + self.assertTrue( + layer_data_0["type"] == layer_data_0["type"], + f"Layer {layer_name}, {model_name_0} and {model_name_1} do not have the same type.", + ) + + self.assertTrue( + len(layer_data_0["values"]) == len(layer_data_1["values"]), + f"Layer {layer_name}, {model_name_0} and {model_name_1} do not have the same number of seen Tensors.", + ) + + # F.conv1d weight has rank 3, and toq.conv1d unpacked weight + # has rank 4. For now, skip the length check for conv1d only. + is_weight_functional_conv1d = ( + result_type == NSSingleResultValuesType.WEIGHT.value + and ( + "conv1d" in layer_data_0["prev_node_target_type"] + or "conv1d" in layer_data_1["prev_node_target_type"] + ) + ) + if not is_weight_functional_conv1d: + for idx in range(len(layer_data_0["values"])): + values_0 = layer_data_0["values"][idx] + values_1 = layer_data_1["values"][idx] + if isinstance(values_0, torch.Tensor): + self.assertTrue( + values_0.shape == values_1.shape, + f"Layer {layer_name}, {model_name_0} and {model_name_1} " + + f"have a shape mismatch at idx {idx}.", + ) + elif isinstance(values_0, list): + values_0 = values_0[0] + values_1 = values_1[0] + self.assertTrue( + values_0.shape == values_1.shape, + f"Layer {layer_name}, {model_name_0} and {model_name_1} " + + f"have a shape mismatch at idx {idx}.", + ) + else: + if not isinstance(values_0, tuple): + raise AssertionError(f"unhandled type {type(values_0)}") + if len(values_0) != 2: + raise AssertionError(f"Expected len(values_0) == 2, got {len(values_0)}") + if len(values_0[1]) != 2: + raise AssertionError(f"Expected len(values_0[1]) == 2, got {len(values_0[1])}") + if values_0[0].shape != values_1[0].shape: + raise AssertionError( + f"Expected values_0[0].shape == values_1[0].shape, " + f"got {values_0[0].shape} != {values_1[0].shape}" + ) + if values_0[1][0].shape != values_1[1][0].shape: + raise AssertionError( + f"Expected values_0[1][0].shape == values_1[1][0].shape, " + f"got {values_0[1][0].shape} != {values_1[1][0].shape}" + ) + if values_0[1][1].shape != values_1[1][1].shape: + raise AssertionError( + f"Expected values_0[1][1].shape == values_1[1][1].shape, " + f"got {values_0[1][1].shape} != {values_1[1][1].shape}" + ) + + # verify that ref_node_name is valid + ref_node_name_0 = layer_data_0["ref_node_name"] + ref_node_name_1 = layer_data_1["ref_node_name"] + prev_node_name_0 = layer_data_0["prev_node_name"] + prev_node_name_1 = layer_data_1["prev_node_name"] + if ( + layer_data_0["type"] + == NSSingleResultValuesType.NODE_OUTPUT.value + ): + self.assertTrue(ref_node_name_0 == prev_node_name_0) + self.assertTrue(ref_node_name_1 == prev_node_name_1) + elif ( + layer_data_0["type"] + == NSSingleResultValuesType.NODE_INPUT.value + ): + self.assertTrue(ref_node_name_0 != prev_node_name_0) + self.assertTrue(ref_node_name_1 != prev_node_name_1) + + def checkGraphModeFxOp( + self, + model, + inputs, + quant_type, + expected_node=None, + expected_node_occurrence=None, + expected_node_list=None, + is_reference=False, + print_debug_info=False, + custom_qconfig_dict=None, + prepare_expected_node=None, + prepare_expected_node_occurrence=None, + prepare_expected_node_list=None, + prepare_custom_config=None, + backend_config=None, + ): + """Quantizes model with graph mode quantization on fx and check if the + quantized model contains the quantized_node + + Args: + model: floating point torch.nn.Module + inputs: one positional sample input arguments for model + expected_node: NodeSpec + e.g. NodeSpec.call_function(torch.quantize_per_tensor) + expected_node_occurrence: a dict from NodeSpec to + expected number of occurrences (int) + e.g. {NodeSpec.call_function(torch.quantize_per_tensor) : 1, + NodeSpec.call_method('dequantize'): 1} + expected_node_list: a list of NodeSpec, used to check the order + of the occurrence of Node + e.g. [NodeSpec.call_function(torch.quantize_per_tensor), + NodeSpec.call_module(nnq.Conv2d), + NodeSpec.call_function(F.hardtanh_), + NodeSpec.call_method('dequantize')] + is_reference: if True, enables reference mode + print_debug_info: if True, prints debug info + custom_qconfig_dict: overrides default qconfig_dict + prepare_expected_node: same as expected_node, but for prepare + prepare_expected_node_occurrence: same as + expected_node_occurrence, but for prepare + prepare_expected_node_list: same as expected_node_list, but + for prepare + + Returns: + A dictionary with the following structure: + { + "prepared": ..., # the prepared model + "quantized": ..., # the quantized non-reference model + "quantized_reference": ..., # the quantized reference model + "result": ..., # the result for either quantized or + # quantized_reference model depending on the + # is_reference argument + } + """ + # TODO: make img_data a single example instead of a list + if type(inputs) is list: + inputs = inputs[0] + + if quant_type == QuantType.QAT: + qconfig_mapping = get_default_qat_qconfig_mapping( + torch.backends.quantized.engine + ) + model.train() + elif quant_type == QuantType.STATIC: + qconfig_mapping = get_default_qconfig_mapping( + torch.backends.quantized.engine + ) + model.eval() + else: + qconfig = default_dynamic_qconfig + qconfig_mapping = QConfigMapping().set_global(qconfig) + model.eval() + + if quant_type == QuantType.QAT: + prepare = prepare_qat_fx + else: + prepare = prepare_fx + + # overwrite qconfig_dict with custom_qconfig_dict + if custom_qconfig_dict is not None: + if type(custom_qconfig_dict) not in (QConfigMapping, dict): + raise AssertionError("custom_qconfig_dict should be a QConfigMapping or a dict") + if isinstance(custom_qconfig_dict, QConfigMapping): + qconfig_mapping = custom_qconfig_dict + else: + qconfig_mapping = QConfigMapping.from_dict(custom_qconfig_dict) + prepared = prepare( + model, + qconfig_mapping, + example_inputs=inputs, + prepare_custom_config=prepare_custom_config, + backend_config=backend_config, + ) + if quant_type != QuantType.DYNAMIC: + prepared(*inputs) + + if print_debug_info: + print() + print("quant type:\n", quant_type) + print("original model:\n", model) + print() + print("prepared model:\n", prepared) + + self.checkGraphModuleNodes( + prepared, + prepare_expected_node, + prepare_expected_node_occurrence, + prepare_expected_node_list, + ) + + prepared_copy = copy.deepcopy(prepared) + qgraph = convert_fx(copy.deepcopy(prepared)) + qgraph_reference = convert_to_reference_fx(copy.deepcopy(prepared)) + result = qgraph(*inputs) + result_reference = qgraph_reference(*inputs) + qgraph_copy = copy.deepcopy(qgraph) + qgraph_reference_copy = copy.deepcopy(qgraph_reference) + + qgraph_to_check = qgraph_reference if is_reference else qgraph + if print_debug_info: + print() + print("quantized model:\n", qgraph_to_check) + self.printGraphModule(qgraph_to_check) + print() + self.checkGraphModuleNodes( + qgraph_to_check, + expected_node, + expected_node_occurrence, + expected_node_list, + ) + return { + "prepared": prepared_copy, + "quantized": qgraph_copy, + "quantized_reference": qgraph_reference_copy, + "quantized_output": result, + "quantized_reference_output": result_reference, + } + + def checkEmbeddingSerialization( + self, + qemb, + num_embeddings, + embedding_dim, + indices, + offsets, + set_qconfig, + is_emb_bag, + dtype=torch.quint8, + ): + # Test serialization of dynamic EmbeddingBag module using state_dict + if is_emb_bag: + inputs = [indices, offsets] + else: + inputs = [indices] + emb_dict = qemb.state_dict() + b = io.BytesIO() + torch.save(emb_dict, b) + b.seek(0) + loaded_dict = torch.load(b) + embedding_unpack = torch.ops.quantized.embedding_bag_unpack + # Check unpacked weight values explicitly + for key in emb_dict: + if isinstance(emb_dict[key], torch._C.ScriptObject): + if not isinstance(loaded_dict[key], torch._C.ScriptObject): + raise AssertionError(f"Expected loaded_dict[{key!r}] to be ScriptObject") + emb_weight = embedding_unpack(emb_dict[key]) + loaded_weight = embedding_unpack(loaded_dict[key]) + self.assertEqual(emb_weight, loaded_weight) + + # Check state dict serialization and torch.save APIs + if is_emb_bag: + loaded_qemb = nnq.EmbeddingBag( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + include_last_offset=True, + mode="sum", + dtype=dtype, + ) + else: + loaded_qemb = nnq.Embedding( + num_embeddings=num_embeddings, embedding_dim=embedding_dim, dtype=dtype + ) + self.check_eager_serialization(qemb, loaded_qemb, inputs) + + loaded_qemb.load_state_dict(loaded_dict) + self.assertEqual( + embedding_unpack(qemb._packed_params._packed_weight), + embedding_unpack(loaded_qemb._packed_params._packed_weight), + ) + + # Test JIT serialization + self.checkScriptable(qemb, [inputs], check_save_load=True) + + # Test from_float call + if is_emb_bag: + float_embedding = torch.nn.EmbeddingBag( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + include_last_offset=True, + scale_grad_by_freq=False, + mode="sum", + ) + else: + float_embedding = torch.nn.Embedding( + num_embeddings=num_embeddings, embedding_dim=embedding_dim + ) + + if set_qconfig: + float_qparams_observer = PerChannelMinMaxObserver.with_args( + dtype=dtype, qscheme=torch.per_channel_affine_float_qparams, ch_axis=0 + ) + float_embedding.qconfig = QConfig( + activation=default_dynamic_quant_observer, weight=float_qparams_observer + ) + + prepare_dynamic(float_embedding) + + float_embedding(*inputs) + if is_emb_bag: + q_embeddingbag = nnq.EmbeddingBag.from_float(float_embedding) + expected_name = "QuantizedEmbeddingBag" + else: + q_embeddingbag = nnq.Embedding.from_float(float_embedding) + expected_name = "QuantizedEmbedding" + + q_embeddingbag(*inputs) + + self.assertTrue(expected_name in str(q_embeddingbag)) + + +class QuantizationLiteTestCase(QuantizationTestCase): + def _create_quantized_model(self, model_class: type[torch.nn.Module], **kwargs): + # Creates quantized model for testing mobile script modules + qengine = "qnnpack" + with override_quantized_engine(qengine): + # FIXME(rec): shouldn't qconfig be passed to quantize? + qconfig = torch.ao.quantization.get_default_qconfig(qengine) # noqa: F841 + model = model_class(**kwargs) + model = quantize(model, test_only_eval_fn, [self.calib_data]) + + return model + + def _compare_script_and_mobile(self, model: torch.nn.Module, input: torch.Tensor): + # Compares the numerical outputs for script and lite modules + qengine = "qnnpack" + with override_quantized_engine(qengine): + script_module = torch.jit.script(model) + script_module_result = script_module(input) + + max_retry = 5 + for retry in range(1, max_retry + 1): + # retries `max_retry` times; breaks iff succeeds else throws exception + try: + buffer = io.BytesIO( + script_module._save_to_buffer_for_lite_interpreter() + ) + buffer.seek(0) + mobile_module = _load_for_lite_interpreter(buffer) + + mobile_module_result = mobile_module(input) + + torch.testing.assert_close( + script_module_result, mobile_module_result + ) + mobile_module_forward_result = mobile_module.forward(input) + torch.testing.assert_close( + script_module_result, mobile_module_forward_result + ) + + mobile_module_run_method_result = mobile_module.run_method( + "forward", input + ) + torch.testing.assert_close( + script_module_result, mobile_module_run_method_result + ) + except AssertionError as e: + if retry == max_retry: + raise e + else: + continue + break + + +# Below are a series of toy models to use in testing quantization + + +class SingleLayerLinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Linear(5, 5).to(dtype=torch.float) + + def forward(self, x): + x = self.fc1(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class AnnotatedSingleLayerLinearModel(torch.nn.Module): + def __init__(self, qengine="fbgemm"): + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qconfig(qengine) + self.fc1 = QuantWrapper(torch.nn.Linear(5, 5).to(dtype=torch.float)) + + def forward(self, x): + x = self.fc1(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class SingleLayerLinearDynamicModel(torch.nn.Module): + def __init__(self, qengine="fbgemm"): + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qconfig(qengine) + self.fc1 = torch.nn.Linear(5, 5).to(dtype=torch.float) + + def forward(self, x): + x = self.fc1(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class LinearAddModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Linear(5, 8).to(dtype=torch.float) + self.fc2 = torch.nn.Linear(8, 5).to(dtype=torch.float) + + def forward(self, x): + x = self.fc1(x) + x = torch.add(x, 5) + x = self.fc2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class RNNDynamicModel(torch.nn.Module): + def __init__(self, mod_type): + super().__init__() + self.qconfig = default_dynamic_qconfig + if mod_type == "GRU": + self.mod = torch.nn.GRU(2, 2).to(dtype=torch.float) + if mod_type == "LSTM": + self.mod = torch.nn.LSTM(2, 2).to(dtype=torch.float) + + def forward(self, x): + x = self.mod(x) + return x + + +class RNNCellDynamicModel(torch.nn.Module): + def __init__(self, mod_type): + super().__init__() + self.qconfig = default_dynamic_qconfig + if mod_type == "GRUCell": + self.mod = torch.nn.GRUCell(2, 2).to(dtype=torch.float) + if mod_type == "LSTMCell": + self.mod = torch.nn.LSTMCell(2, 2).to(dtype=torch.float) + if mod_type == "RNNReLU": + self.mod = torch.nn.RNNCell(2, 2, nonlinearity="relu").to(dtype=torch.float) + if mod_type == "RNNTanh": + self.mod = torch.nn.RNNCell(2, 2, nonlinearity="tanh").to(dtype=torch.float) + + def forward(self, x): + x = self.mod(x) + return x + + +class LSTMwithHiddenDynamicModel(torch.nn.Module): + def __init__(self, qengine="fbgemm"): + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qconfig(qengine) + self.lstm = torch.nn.LSTM(2, 2).to(dtype=torch.float) + + def forward(self, x, hid): + x, hid = self.lstm(x, hid) + return x, hid + + +class ConvModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float) + + def forward(self, x): + x = self.conv(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class ConvTransposeModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.ConvTranspose2d(3, 5, 3, bias=False).to(dtype=torch.float) + + def forward(self, x): + x = self.conv(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class AnnotatedConvModel(torch.nn.Module): + def __init__(self, qengine): + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qconfig(qengine) + self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float) + self.quant = QuantStub() + self.dequant = DeQuantStub() + + def forward(self, x): + x = self.quant(x) + x = self.conv(x) + x = self.dequant(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class AnnotatedConvTransposeModel(torch.nn.Module): + def __init__(self, qengine): + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qconfig(qengine) + self.conv = torch.nn.ConvTranspose2d(3, 5, 3, bias=False).to(dtype=torch.float) + self.quant = QuantStub() + self.dequant = DeQuantStub() + + def forward(self, x): + x = self.quant(x) + x = self.conv(x) + x = self.dequant(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class ConvBnModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float) + self.bn = torch.nn.BatchNorm2d(5).to(dtype=torch.float) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class AnnotatedConvBnModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.qconfig = default_qconfig + self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float) + self.bn = torch.nn.BatchNorm2d(5).to(dtype=torch.float) + self.quant = QuantStub() + self.dequant = DeQuantStub() + + def forward(self, x): + x = self.quant(x) + x = self.conv(x) + x = self.bn(x) + x = self.dequant(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class ConvBnReLUModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float) + self.bn = torch.nn.BatchNorm2d(5).to(dtype=torch.float) + self.relu = nn.ReLU(inplace=True) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + x = self.relu(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class AnnotatedConvBnReLUModel(torch.nn.Module): + def __init__(self, qengine="fbgemm"): + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qconfig(qengine) + self.conv = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float) + self.bn = torch.nn.BatchNorm2d(5).to(dtype=torch.float) + self.relu = nn.ReLU(inplace=True) + self.quant = QuantStub() + self.dequant = DeQuantStub() + + def forward(self, x): + x = self.quant(x) + x = self.conv(x) + x = self.bn(x) + x = self.relu(x) + x = self.dequant(x) + return x + + def fuse_model(self): + # TODO: remove this check and define two fuse_modules function on this module + if self.training: + torch.ao.quantization.fuse_modules_qat( + self, [["conv", "bn", "relu"]], inplace=True + ) + else: + torch.ao.quantization.fuse_modules( + self, [["conv", "bn", "relu"]], inplace=True + ) + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class TwoLayerConvModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 5, 3, bias=False).to(dtype=torch.float) + self.conv2 = torch.nn.Conv2d(5, 5, 1, bias=False).to(dtype=torch.float) + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class TwoLayerLinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Linear(5, 8).to(dtype=torch.float) + self.fc2 = torch.nn.Linear(8, 5).to(dtype=torch.float) + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class LinearModelWithSubmodule(nn.Module): + def __init__(self) -> None: + super().__init__() + self.subm = TwoLayerLinearModel() + self.fc = nn.Linear(5, 5) + + def forward(self, x): + x = self.subm(x) + x = self.fc(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.subm.get_example_inputs() + + +class AnnotatedTwoLayerLinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Linear(5, 8).to(dtype=torch.float) + self.fc2 = QuantWrapper(torch.nn.Linear(8, 5).to(dtype=torch.float)) + self.fc2.qconfig = torch.ao.quantization.get_default_qconfig("fbgemm") + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class ActivationsTestModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qconfig("fbgemm") + self.quant = torch.ao.quantization.QuantStub() + self.hardswish = torch.nn.Hardswish().to(dtype=torch.float) + self.elu = torch.nn.ELU().to(dtype=torch.float) + self.dequant = torch.ao.quantization.DeQuantStub() + + def forward(self, x): + x = self.quant(x) + x = self.hardswish(x) + x = self.elu(x) + x = self.dequant(x) + return x + + +class LinearReluModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc = torch.nn.Linear(5, 5).to(dtype=torch.float) + self.relu = torch.nn.ReLU() + + def forward(self, x): + x = self.relu(self.fc(x)) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class LinearReluLinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Linear(5, 8).to(dtype=torch.float) + self.relu = torch.nn.ReLU() + self.fc2 = torch.nn.Linear(8, 5).to(dtype=torch.float) + + def forward(self, x): + x = self.fc1(x) + x = self.relu(x) + x = self.fc2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class LinearReluAddModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Linear(5, 5).to(dtype=torch.float) + self.relu = torch.nn.ReLU() + self.fc2 = torch.nn.Linear(5, 5).to(dtype=torch.float) + + def forward(self, x): + x = self.fc1(x) + x = self.relu(x) + x = torch.add(x, 5) + x = self.fc2(x) + self.relu = torch.nn.ReLU() + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class LinearBnLeakyReluModel(torch.nn.Module): + def __init__(self, with_bn=True): + super().__init__() + self.linear = nn.Linear(5, 5) + self.bn1d = nn.BatchNorm1d(5) + self.leaky_relu = nn.LeakyReLU(0.01) + self.with_bn = with_bn + + def forward(self, x): + x = self.linear(x) + if self.with_bn: + x = self.bn1d(x) + x = self.leaky_relu(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class LinearTanhModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(5, 5) + self.tanh = nn.Tanh() + + def forward(self, x): + x = self.linear(x) + x = self.tanh(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class ConvBnAddReluModel(torch.nn.Module): + def __init__( + self, + with_bn=True, + with_relu=True, + left_conv=True, + two_conv=True, + use_torch_add=True, + ): + super().__init__() + self.conv = nn.Conv2d(5, 5, (2, 2)) + self.conv2 = nn.Conv2d(5, 5, (2, 2)) + self.bn = nn.BatchNorm2d(5) + self.relu = nn.ReLU() + self.with_bn = with_bn + self.with_relu = with_relu + self.two_conv = two_conv + self.left_conv = left_conv + self.use_torch_add = use_torch_add + + def forward(self, x1, x2): + if self.two_conv: + if self.use_torch_add: + if self.with_bn: + x = torch.add(self.bn(self.conv(x1)), self.conv2(x1)) + else: + x = torch.add(self.conv(x1), self.conv2(x1)) + else: + if self.with_bn: + x = self.bn(self.conv(x1)) + self.conv2(x1) + else: + x = self.conv(x1) + self.conv2(x1) + else: + if self.use_torch_add: + if self.left_conv: + if self.with_bn: + x = torch.add(self.bn(self.conv(x1)), x2) + else: + x = torch.add(self.conv(x1), x2) + else: + if self.with_bn: + x = torch.add(x2, self.bn(self.conv(x1))) + else: + x = torch.add(x2, self.conv(x1)) + else: + if self.left_conv: + if self.with_bn: + x = self.bn(self.conv(x1)) + x2 + else: + x = self.conv(x1) + x2 + else: + if self.with_bn: + x = x2 + self.bn(self.conv(x1)) + else: + x = x2 + self.conv(x1) + if self.with_relu: + x = self.relu(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5, 3, 3), torch.rand(1, 5, 2, 2)) + + +# TODO: self.fc should be self.conv +class ConvReluModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc = torch.nn.Conv2d(3, 5, 3).to(dtype=torch.float) + self.relu = torch.nn.ReLU() + + def forward(self, x): + x = self.relu(self.fc(x)) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +# TODO: self.fc should be self.conv +class ConvReluConvModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Conv2d(3, 5, 3).to(dtype=torch.float) + self.relu = torch.nn.ReLU() + self.fc2 = torch.nn.Conv2d(5, 5, 1).to(dtype=torch.float) + + def forward(self, x): + x = self.fc1(x) + x = self.relu(x) + x = self.fc2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +# TODO: self.fc should be self.conv +class ConvReluAddModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Conv2d(3, 5, 3).to(dtype=torch.float) + self.relu = torch.nn.ReLU() + self.fc2 = torch.nn.Conv2d(5, 5, 1).to(dtype=torch.float) + + def forward(self, x): + x = self.fc1(x) + x = self.relu(x) + x = torch.add(x, 5) + x = self.fc2(x) + self.relu = torch.nn.ReLU() + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class NormalizationTestModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.quant = torch.ao.quantization.QuantStub() + self.fc1 = torch.nn.Linear(5, 8).to(dtype=torch.float) + self.layer_norm = torch.nn.LayerNorm(8) + self.group_norm = torch.nn.GroupNorm(2, 8) + self.instance_norm1d = torch.nn.InstanceNorm1d(8) + self.instance_norm2d = torch.nn.InstanceNorm2d(8) + self.instance_norm3d = torch.nn.InstanceNorm3d(8) + + def forward(self, x): + x = self.quant(x) + x = self.fc1(x) + x = self.layer_norm(x) + x = self.group_norm(x.unsqueeze(-1).repeat(1, 1, 3)) + x = self.instance_norm1d(x) + x = self.instance_norm2d(x.unsqueeze(-1)) + x = self.instance_norm3d(x.unsqueeze(-1)) + return x + + +class NestedModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.sub1 = LinearReluModel() + self.sub2 = TwoLayerLinearModel() + self.fc3 = torch.nn.Linear(5, 5).to(dtype=torch.float) + + def forward(self, x): + x = self.sub1(x) + x = self.sub2(x) + x = self.fc3(x) + return x + + +class AnnotatedNestedModel(torch.nn.Module): + def __init__(self, qengine): + super().__init__() + self.sub1 = LinearReluModel() + self.sub2 = TwoLayerLinearModel() + self.fc3 = QuantWrapper(torch.nn.Linear(5, 5).to(dtype=torch.float)) + self.fc3.qconfig = default_qconfig + self.sub2.fc1 = QuantWrapper(self.sub2.fc1) + if qengine == "fbgemm": + self.sub2.fc1.qconfig = default_per_channel_qconfig + else: + self.sub2.fc1.qconfig = default_qconfig + + def forward(self, x): + x = self.sub1(x) + x = self.sub2(x) + x = self.fc3(x) + return x + + +class AnnotatedSubNestedModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.sub1 = LinearReluModel() + self.sub2 = QuantWrapper(TwoLayerLinearModel()) + self.fc3 = QuantWrapper(torch.nn.Linear(5, 5).to(dtype=torch.float)) + self.fc3.qconfig = default_qconfig + self.sub2.qconfig = default_qconfig + + def forward(self, x): + x = self.sub1(x) + x = self.sub2(x) + x = self.fc3(x) + return x + + +class AnnotatedCustomConfigNestedModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.sub1 = LinearReluModel() + self.sub2 = TwoLayerLinearModel() + self.fc3 = QuantWrapper(torch.nn.Linear(5, 5).to(dtype=torch.float)) + self.fc3.qconfig = default_qconfig + self.sub2.qconfig = default_qconfig + + custom_options = {"dtype": torch.quint8, "qscheme": torch.per_tensor_affine} + custom_qconfig = QConfig( + activation=default_observer.with_args(**custom_options), + weight=default_weight_observer, + ) + self.sub2.fc1.qconfig = custom_qconfig + + self.sub2.fc1 = QuantWrapper(self.sub2.fc1) + self.sub2.fc2 = QuantWrapper(self.sub2.fc2) + + def forward(self, x): + x = self.sub1(x) + x = self.sub2(x) + x = self.fc3(x) + return x + + +class QuantSubModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.sub1 = LinearReluModel() + self.sub2 = QuantWrapper(TwoLayerLinearModel()) + self.sub2.qconfig = default_qconfig + self.fc3 = torch.nn.Linear(5, 5).to(dtype=torch.float) + self.fc3.qconfig = default_qconfig + + def forward(self, x): + x = self.sub1(x) + x = self.sub2(x) + x = self.fc3(x) + return x + + +class InnerModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = torch.nn.Linear(5, 8).to(dtype=torch.float) + self.relu1 = torch.nn.ReLU() + self.fc2 = torch.nn.Linear(8, 5).to(dtype=torch.float) + self.relu2 = torch.nn.ReLU() + + def forward(self, x): + return self.relu2(self.fc2(self.relu1(self.fc1(x)))) + + def fuse_modules(self): + fusable_layers = [] + named_children = list(self.named_children()) + for idx, (current_name, layer) in enumerate(named_children): + if isinstance(layer, torch.nn.Linear): + if idx >= len(named_children) - 1: + break + if isinstance(named_children[idx + 1][1], torch.nn.ReLU): + fusable_layers.append([current_name, named_children[idx + 1][0]]) + # TODO: remove this check and define two fuse_modules function on this module + if self.training: + torch.ao.quantization.fuse_modules_qat(self, fusable_layers, inplace=True) + else: + torch.ao.quantization.fuse_modules(self, fusable_layers, inplace=True) + + +class FunctionalLinear(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.rand((5, 5)) + self.bias = torch.zeros(5) + + def forward(self, x): + return F.linear(x, self.weight, self.bias) + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 5),) + + +class SingleLayerFunctionalLinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear1 = FunctionalLinear() + + def forward(self, x): + x = self.linear1(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.linear1.get_example_inputs() + + +class TwoLayerFunctionalLinearModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear1 = FunctionalLinear() + self.linear2 = FunctionalLinear() + + def forward(self, x): + x = self.linear1(x) + x = self.linear2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.linear1.get_example_inputs() + + +class FunctionalLinearAddModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear1 = FunctionalLinear() + self.linear2 = FunctionalLinear() + + def forward(self, x): + x = self.linear1(x) + x = torch.add(x, 5) + x = self.linear2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.linear1.get_example_inputs() + + +class FunctionalLinearReluModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear = FunctionalLinear() + + def forward(self, x): + x = self.linear(x) + x = F.relu(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.linear.get_example_inputs() + + +class FunctionalLinearReluLinearModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear1 = FunctionalLinear() + self.relu = nn.ReLU() + self.linear2 = FunctionalLinear() + + def forward(self, x): + x = self.linear1(x) + x = self.relu(x) + x = self.linear2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.linear1.get_example_inputs() + + +class FunctionalConv2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.weight = torch.rand(3, 3, 3, 3) + self.bias = torch.rand(3) + self.stride = (1, 1) + self.padding = (0, 0) + self.dilation = (1, 1) + self.groups = 1 + + def forward(self, x): + return F.conv2d( + x, + self.weight, + self.bias, + self.stride, + self.padding, + self.dilation, + self.groups, + ) + + def get_example_inputs(self) -> tuple[Any, ...]: + return (torch.rand(1, 3, 5, 5),) + + +class SingleLayerFunctionalConvModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = FunctionalConv2d() + + def forward(self, x): + x = self.conv1(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.conv1.get_example_inputs() + + +class TwoLayerFunctionalConvModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = FunctionalConv2d() + self.conv2 = FunctionalConv2d() + + def forward(self, x): + x = self.conv1(x) + x = self.conv2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.conv1.get_example_inputs() + + +class FunctionalConvReluModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = FunctionalConv2d() + + def forward(self, x): + x = self.conv(x) + x = F.relu(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.conv.get_example_inputs() + + +class FunctionalConvReluConvModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = FunctionalConv2d() + self.relu = nn.ReLU() + self.conv2 = FunctionalConv2d() + + def forward(self, x): + x = self.conv1(x) + x = self.relu(x) + x = self.conv2(x) + return x + + def get_example_inputs(self) -> tuple[Any, ...]: + return self.conv1.get_example_inputs() + + +class SkipQuantModel(torch.nn.Module): + r"""We can skip quantization by explicitly + setting qconfig of a submodule to None + """ + + def __init__(self) -> None: + super().__init__() + self.sub = InnerModule() + self.fc = torch.nn.Linear(5, 5).to(dtype=torch.float) + + def forward(self, x): + return self.fc(self.sub(x)) + + def fuse_modules(self): + self.sub.fuse_modules() + + +class AnnotatedSkipQuantModel(torch.nn.Module): + r"""We can skip quantization by explicitly + setting qconfig of a submodule to None + """ + + def __init__(self, qengine): + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qconfig(qengine) + self.sub = QuantWrapper(InnerModule()) + self.fc = torch.nn.Linear(5, 5).to(dtype=torch.float) + # don't quantize this fc + self.fc.qconfig = None + + def forward(self, x): + return self.fc(self.sub(x)) + + def fuse_modules(self): + self.sub.module.fuse_modules() + + +class QuantStubModel(torch.nn.Module): + r"""A Module with manually inserted `QuantStub` and `DeQuantStub`""" + + def __init__(self) -> None: + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qconfig("qnnpack") + self.quant = QuantStub() + self.dequant = DeQuantStub() + self.fc = torch.nn.Linear(5, 5).to(dtype=torch.float) + + def forward(self, x): + x = self.quant(x) + x = self.fc(x) + return self.dequant(x) + + +class ManualLinearQATModel(torch.nn.Module): + r"""A Module with manually inserted `QuantStub` and `DeQuantStub`""" + + def __init__(self, qengine): + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qat_qconfig(qengine) + self.quant = QuantStub() + self.dequant = DeQuantStub() + self.fc1 = torch.nn.Linear(5, 1).to(dtype=torch.float) + self.fc2 = torch.nn.Linear(1, 10).to(dtype=torch.float) + + def forward(self, x): + x = self.quant(x) + x = self.fc1(x) + x = self.fc2(x) + return self.dequant(x) + + +class ManualDropoutQATModel(torch.nn.Module): + r"""A Module with manually inserted `QuantStub` and `DeQuantStub`""" + + def __init__(self, qengine): + super().__init__() + self.qconfig = torch.ao.quantization.get_default_qat_qconfig(qengine) + self.quant = QuantStub() + self.dequant = DeQuantStub() + self.fc1 = torch.nn.Linear(5, 1).to(dtype=torch.float) + self.dropout = torch.nn.Dropout(0.5) + + def forward(self, x): + x = self.quant(x) + x = self.fc1(x) + x = self.dropout(x) + return self.dequant(x) + + +class ManualLinearDynamicQATModel(torch.nn.Module): + r"""A Module that uses a dynamic QAT by default.""" + + def __init__(self, qconfig=None): + super().__init__() + self.qconfig = qconfig or default_dynamic_qat_qconfig + self.fc1 = torch.nn.Linear(5, 1).to(dtype=torch.float) + self.fc2 = torch.nn.Linear(1, 10).to(dtype=torch.float) + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + return x + + +class ManualConvLinearQATModel(torch.nn.Module): + r"""A module with manually inserted `QuantStub` and `DeQuantStub` + and contains both linear and conv modules + """ + + def __init__(self, qconfig=None): + super().__init__() + self.qconfig = ( + qconfig + if qconfig + else torch.ao.quantization.get_default_qat_qconfig("qnnpack") + ) + self.quant = QuantStub() + self.dequant = DeQuantStub() + self.conv = torch.nn.Conv2d(3, 1, kernel_size=3).to(dtype=torch.float) + self.fc1 = torch.nn.Linear(64, 10).to(dtype=torch.float) + self.fc2 = torch.nn.Linear(10, 10).to(dtype=torch.float) + + def forward(self, x): + x = self.quant(x) + x = self.conv(x) + x = x.view(-1, 64).contiguous() + x = self.fc1(x) + x = self.fc2(x) + return self.dequant(x) + + +class ManualConvLinearSymmQATModel(ManualConvLinearQATModel): + r"""Same as ManualConvLinearQATModule but with Symmetric Quantization. + Supported only with qnnpack. + """ + + def __init__(self) -> None: + super().__init__(default_symmetric_qnnpack_qat_qconfig) + + +class ManualEmbeddingBagLinear(nn.Module): + def __init__(self) -> None: + super().__init__() + self.emb = nn.EmbeddingBag(num_embeddings=10, embedding_dim=12, mode="sum") + self.emb.qconfig = default_embedding_qat_qconfig + self.quant = QuantStub() + self.dequant = DeQuantStub() + self.linear = nn.Linear(12, 1).to(dtype=torch.float) + self.qconfig = get_default_qat_qconfig("qnnpack") + + def forward( + self, + input: torch.Tensor, + offsets: torch.Tensor | None = None, + per_sample_weights: torch.Tensor | None = None, + ): + x = self.emb(input, offsets, per_sample_weights) + x = self.quant(x) + x = self.linear(x) + return self.dequant(x) + + +class DeFusedEmbeddingBagLinear(nn.Module): + r"""A module to simulate QAT embedding bag with a linear layer, + this module uses a separate embedding and bagging op, similar + to that which is described in the EmbeddingBag documentation. + + https://pytorch.org/docs/stable/generated/torch.nn.EmbeddingBag.html + """ + + def __init__(self) -> None: + super().__init__() + self.emb = nn.Embedding(num_embeddings=10, embedding_dim=12) + self.emb.qconfig = default_embedding_qat_qconfig + self.bagging_op = torch.sum + self.quant = QuantStub() + self.dequant = DeQuantStub() + self.linear = nn.Linear(12, 1).to(dtype=torch.float) + self.qconfig = get_default_qat_qconfig("qnnpack") + + def forward(self, input: torch.Tensor) -> torch.Tensor: + x = self.bagging_op(self.emb(input), dim=1) + x = self.quant(x) + x = self.linear(x) + return self.dequant(x) + + +class SubModelForFusion(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv2d(2, 2, 1, bias=None).to(dtype=torch.float) + self.bn = nn.BatchNorm2d(2).to(dtype=torch.float) + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + return x + + +class SubModelWithoutFusion(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv2d(2, 2, 1, bias=None).to(dtype=torch.float) + self.relu = nn.ReLU(inplace=False).to(dtype=torch.float) + + def forward(self, x): + return self.relu(self.conv(x)) + + +class ModelForFusion(nn.Module): + def __init__(self, qconfig): + super().__init__() + self.conv1 = nn.Conv2d(3, 2, 1, bias=None).to(dtype=torch.float) + self.bn1 = nn.BatchNorm2d(2).to(dtype=torch.float) + self.relu1 = nn.ReLU(inplace=True).to(dtype=torch.float) + self.sub1 = SubModelForFusion() + self.sub2 = SubModelWithoutFusion() + self.fc = nn.Linear(36, 10).to(dtype=torch.float) + self.quant = QuantStub() + self.dequant = DeQuantStub() + self.qconfig = qconfig + self.conv2 = nn.Conv3d(3, 2, (1, 1, 1), bias=None).to(dtype=torch.float) + self.relu2 = nn.ReLU(inplace=False).to(dtype=torch.float) + self.bn2 = nn.BatchNorm3d(2).to(dtype=torch.float) + self.relu3 = nn.ReLU(inplace=True).to(dtype=torch.float) + self.conv3 = nn.Conv1d(3, 3, 2).to(dtype=torch.float) + self.bn3 = nn.BatchNorm1d(3).to(dtype=torch.float) + self.relu4 = nn.ReLU(inplace=True).to(dtype=torch.float) + # don't quantize sub2 + self.sub2.qconfig = None + self.fc.qconfig = None + + def forward(self, x): + x = x.squeeze(2) + x = self.quant(x) + x = self.conv3(x) + x = self.bn3(x) + x = self.relu4(x) + x = x.unsqueeze(2) + y = x.unsqueeze(2) + x = self.conv1(x) + x = self.bn1(x) + x = self.relu1(x) + x = self.sub1(x) + x = self.dequant(x) + x = self.sub2(x) + x = x.reshape(-1, 36).contiguous() + x = self.fc(x) + y = self.conv2(y) + y = self.relu2(y) + y = self.bn2(y) + y = self.relu3(y) + y = self.dequant(y) + return x + + +class ConvBNReLU(nn.Sequential): + def __init__(self) -> None: + super().__init__( + nn.Conv2d(3, 3, 1, 1, bias=False), nn.BatchNorm2d(3), nn.ReLU(inplace=False) + ) + + +class ModelWithSequentialFusion(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = nn.Conv2d(3, 3, 1) + self.relu1 = nn.ReLU(inplace=False) + layers = [ConvBNReLU() for _ in range(3)] + self.features = nn.Sequential(*layers) + head = [nn.Linear(300, 10), nn.ReLU(inplace=False)] + self.classifier = nn.Sequential(*head) + self.seq = nn.Sequential() + self.quant = QuantStub() + self.dequant = DeQuantStub() + + def forward(self, x): + x = self.quant(x) + x = self.conv1(x) + x = self.relu1(x) + x = self.features(x) + x = torch.reshape(x, (-1, 3 * 10 * 10)) + x = self.classifier(x) + x = self.seq(x) + x = self.dequant(x) + return x + + +class ModelForFusionWithBias(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = nn.Conv2d(3, 2, 5, bias=True).to(dtype=torch.float) + self.bn1 = nn.BatchNorm2d(2).to(dtype=torch.float) + self.relu1 = nn.ReLU(inplace=True).to(dtype=torch.float) + self.conv2 = nn.Conv2d(2, 2, 1, bias=True).to(dtype=torch.float) + self.bn2 = nn.BatchNorm2d(2).to(dtype=torch.float) + self.quant = QuantStub() + self.dequant = DeQuantStub() + + def forward(self, x): + x = self.quant(x) + x = self.conv1(x) + x = self.bn1(x) + x = self.relu1(x) + x = self.conv2(x) + x = self.bn2(x) + x = self.dequant(x) + return x + + +class ModelForLinearBNFusion(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc = nn.Linear(20, 10) + self.bn = nn.BatchNorm1d(10) + nn.init.uniform_(self.bn.weight) + nn.init.uniform_(self.bn.bias) + + def forward(self, x): + return self.bn(self.fc(x)) + + +class DummyObserver(torch.nn.Module): + def calculate_qparams(self): + return 1.0, 0 + + def forward(self, x): + return x + + +class ModelForConvTransposeBNFusion(nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = nn.ConvTranspose1d(3, 3, 1) + self.bn1 = nn.BatchNorm1d(3) + self.conv2 = nn.ConvTranspose2d(3, 3, 1) + self.bn2 = nn.BatchNorm2d(3) + self.conv3 = nn.ConvTranspose3d(3, 3, 1) + self.bn3 = nn.BatchNorm3d(3) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = x.unsqueeze(2) + x = self.conv2(x) + x = self.bn2(x) + x = x.unsqueeze(2) + x = self.conv3(x) + x = self.bn3(x) + return x + + +class ModelWithFunctionals(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.mycat = nnq.FloatFunctional() + self.myadd = nnq.FloatFunctional() + self.myadd_relu = nnq.FloatFunctional() + self.mymatmul = nnq.FloatFunctional() + # Tracing doesn't work yet for c10 ops with scalar inputs + # https://github.com/pytorch/pytorch/issues/27097 + # self.my_scalar_add = nnq.FloatFunctional() + # self.my_scalar_mul = nnq.FloatFunctional() + + def forward(self, x): + y = self.mycat.cat([x, x, x]) + z = self.myadd.add(y, y) + w = self.myadd_relu.add_relu(z, z) + u = self.mymatmul.matmul(w, w.T) + # Tracing doesn't work yet for c10 ops with scalar inputs + # https://github.com/pytorch/pytorch/issues/27097 + # w = self.my_scalar_add.add_scalar(w, -0.5) + # w = self.my_scalar_mul.mul_scalar(w, 0.5) + return u + + +class ResNetBase(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + norm_layer = nn.BatchNorm2d + inplanes = 3 + self.conv1 = nn.Conv2d(inplanes, inplanes, (1, 1), bias=False) + self.bn1 = norm_layer(inplanes) + self.relu1 = nn.ReLU() + self.relu2 = nn.ReLU() + self.downsample = torch.nn.Identity() + self.myop = nn.quantized.FloatFunctional() + self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) + self.fc = torch.nn.Linear(inplanes, 1) + + def forward(self, x): + out = self.conv1(x) + out = self.bn1(out) + out = self.relu1(out) + identity = self.downsample(x) + out = self.myop.add(out, identity) + out = self.relu2(out) + out = self.avgpool(out) + out = torch.flatten(out, 1) + out = self.fc(out) + return out + + def fuse_model(self): + # TODO: remove this check and define two fuse_model function on this module + if self.training: + torch.ao.quantization.fuse_modules_qat( + self, [["conv1", "bn1", "relu1"]], inplace=True + ) + else: + torch.ao.quantization.fuse_modules( + self, [["conv1", "bn1", "relu1"]], inplace=True + ) + + +class ModelMultipleOps(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + norm_layer = nn.BatchNorm2d + inplanes = 3 + self.conv1 = nn.Conv2d(inplanes, inplanes, (1, 1), bias=False) + self.conv2 = nn.Conv2d(inplanes, inplanes, (1, 1), bias=False) + self.bn1 = norm_layer(inplanes) + self.relu1 = nn.ReLU() + self.relu2 = nn.ReLU() + self.downsample = torch.nn.Identity() + self.skip_add = nn.quantized.FloatFunctional() + self.cat = nn.quantized.FloatFunctional() + self.avgpool = nn.AdaptiveAvgPool2d((4, 4)) + self.fc = nn.Linear(12, 6) + + def forward(self, x): + out = self.conv1(x) + out = self.bn1(out) + out = self.relu1(out) + identity = self.downsample(x) + out = self.skip_add.add(out, identity) + out = self.relu2(out) + out = self.avgpool(out) + out = self.conv2(out) + out = torch.nn.functional.max_pool2d(out, 2, 2) + out = self.cat.cat([out, out]) + out = out.reshape(-1, 3 * 2 * 2) + out = self.fc(out) + return out + + +# Model to ensure consistency of fake quant with true quant +# Average pooling and mean operations are not modelled +# accurately with fake-quant so this model does not +# contain those operations +class ModelMultipleOpsNoAvgPool(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + norm_layer = nn.BatchNorm2d + inplanes = 3 + self.conv1 = nn.Conv2d(inplanes, inplanes, (1, 1), bias=False) + self.conv2 = nn.Conv2d(inplanes, inplanes, (1, 1), bias=False) + self.bn1 = norm_layer(inplanes) + self.relu1 = nn.ReLU() + self.relu2 = nn.ReLU() + self.skip_add = nn.quantized.FloatFunctional() + self.cat = nn.quantized.FloatFunctional() + self.maxpool = nn.MaxPool2d((4, 4)) + self.fc = nn.Linear(12, 6) + + def forward(self, x): + out = self.conv1(x) + out = self.bn1(out) + out = self.relu1(out) + skip = self.conv2(x) + out = self.skip_add.add(out, skip) + out = self.relu2(out) + out = self.maxpool(out) + out = self.conv2(out) + out = torch.nn.functional.max_pool2d(out, 2, 2) + out = self.cat.cat([out, out]) + out = out.reshape(-1, 3 * 2 * 2) + out = self.fc(out) + return out + + +class EmbeddingBagModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.emb = torch.nn.EmbeddingBag( + num_embeddings=10, + embedding_dim=12, + include_last_offset=True, + scale_grad_by_freq=False, + mode="sum", + ) + + def forward(self, indices, offsets, per_sample_weights): + return self.emb(indices, offsets, per_sample_weights) + + +class EmbeddingModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.emb = torch.nn.Embedding(num_embeddings=10, embedding_dim=12) + + def forward(self, indices): + return self.emb(indices) + + +class EmbeddingWithStaticLinear(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.emb = torch.nn.EmbeddingBag(num_embeddings=10, embedding_dim=12) + self.fc = torch.nn.Linear(4, 2) + self.emb.qconfig = float_qparams_weight_only_qconfig + self.qconfig = default_qconfig + self.quant = QuantStub() + self.dequant = DeQuantStub() + + def forward(self, indices, offsets, linear_in): + emb = self.emb(indices, offsets) + q_x = self.quant(linear_in) + fc = self.fc(q_x) + fc = self.dequant(fc) + features = torch.cat([fc] + [emb], dim=1) + return features + + +class DenseTopMLP(nn.Module): + def __init__( + self, dense_dim, dense_out, embedding_dim, top_out_in, top_out_out + ) -> None: + super().__init__() + + self.dense_mlp = nn.Sequential( + nn.Linear(dense_dim, dense_out), + ) + self.top_mlp = nn.Sequential( + nn.Linear(dense_out + embedding_dim, top_out_in), + nn.Linear(top_out_in, top_out_out), + ) + + def forward( + self, + sparse_feature: torch.Tensor, + dense: torch.Tensor, + ) -> torch.Tensor: + dense_feature = self.dense_mlp(dense) + features = torch.cat([dense_feature] + [sparse_feature], dim=1) + + out = self.top_mlp(features) + return out + + +# thin wrapper around embedding bag, because tracing inside nn.Embedding +# bag is not supported at the moment and this is top level +class EmbBagWrapper(nn.Module): + def __init__(self, num_embeddings, embedding_dim): + super().__init__() + self.emb_bag = nn.EmbeddingBag(num_embeddings, embedding_dim, mode="sum") + + def forward(self, indices, offsets): + return self.emb_bag(indices, offsets) + + +class SparseNNModel(nn.Module): + _NUM_EMBEDDINGS = 10 + _EMBEDDING_DIM = 5 + _DENSE_DIM = 4 + _DENSE_OUTPUT = 2 + _TOP_OUT_IN = 2 + _TOP_OUT_OUT = 2 + _TOP_MLP_DIM = 1 + + def __init__(self) -> None: + super().__init__() + + self.model_sparse = EmbBagWrapper(self._NUM_EMBEDDINGS, self._EMBEDDING_DIM) + self.dense_top = DenseTopMLP( + self._DENSE_DIM, + self._DENSE_OUTPUT, + self._EMBEDDING_DIM, + self._TOP_OUT_IN, + self._TOP_OUT_OUT, + ) + + def forward( + self, + sparse_indices: torch.Tensor, + sparse_offsets: torch.Tensor, + dense: torch.Tensor, + ) -> torch.Tensor: + sparse_feature = self.model_sparse(sparse_indices, sparse_offsets) + out = self.dense_top(sparse_feature, dense) + + return out + + +class TestHelperModules: + class ControlFlow(torch.nn.Module): + def forward( + self, + xs: torch.Tensor, + pred1: torch.Tensor, + pred2: torch.Tensor, + y: torch.Tensor, + ) -> torch.Tensor: + def true_nested(y: torch.Tensor) -> torch.Tensor: + y = y + y + y = torch.mm(y, y) + return y + + def false_nested(y: torch.Tensor) -> torch.Tensor: + return torch.mm(y, y) + + def true_fn(x: torch.Tensor, pred2: torch.Tensor) -> torch.Tensor: + z = control_flow.cond(pred2, true_nested, false_nested, [x]) + return x + z + + def false_fn(x: torch.Tensor, _) -> torch.Tensor: + return x.cos() + + def map_fn( + x: torch.Tensor, + pred1: torch.Tensor, + pred2: torch.Tensor, + y: torch.Tensor, + ) -> torch.Tensor: + x = x.cos() + y = control_flow.cond(pred1, true_fn, false_fn, [y, pred2]) + x = x + y + return x.sin() + + y = torch.mm(y, y) + return control_flow.map(map_fn, xs, pred1, pred2, y) + + def example_inputs(self): + return ( + torch.ones(2, 2), + torch.tensor([False]), + torch.tensor([False]), + torch.ones(2, 2), + ) + + class Conv2dPropAnnotaton(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 3, 3) + self.linear = torch.nn.Linear(3, 3) + + def forward(self, x): + x = self.conv(x) + x = x.view(-1, 3) + x = torch.nn.functional.hardtanh(x, -0.5, 0.5) + x = self.linear(x) + return x + + class Conv2dWithObsSharingOps(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 3, 3) + self.hardtanh = torch.nn.Hardtanh() + self.adaptive_avg_pool2d = torch.nn.AdaptiveAvgPool2d((1, 1)) + + def forward(self, x): + x = self.conv(x) + x = self.adaptive_avg_pool2d(x) + x = self.hardtanh(x) + x = torch.mean(x) + return x + + class Conv2dWithTwoLinearPermute(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 16, 3) + self.linear1 = torch.nn.Linear(16, 8, bias=False) + self.linear2 = torch.nn.Linear(8, 8) + + def forward(self, x): + conv_out = self.conv(x) + permute_out = torch.permute(conv_out, (0, 2, 3, 1)) + return self.linear2(self.linear1(permute_out)) + + class Conv2dWithTwoLinear(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 16, 3) + self.linear1 = torch.nn.Linear(64, 8, bias=False) + self.linear2 = torch.nn.Linear(8, 8) + + def forward(self, x): + conv_out = self.conv(x) + reshape_out = torch.reshape(conv_out, (2, 64)) + return self.linear2(self.linear1(reshape_out)) + + class ConvLinearWPermute(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3) + self.linear1 = torch.nn.Linear(8, 8) + + def forward(self, x): + conv_out = self.conv(x) + permute_out = torch.permute(conv_out, (0, 2, 3, 1)) + return self.linear1(permute_out) + + class TwoLinearModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.linear1 = torch.nn.Linear(8, 16, bias=False) + self.linear2 = torch.nn.Linear(16, 8) + + def forward(self, x): + return self.linear2(self.linear1(x)) + + def example_inputs(self): + return (torch.randn(2, 8),) + + class ConvMaxPool2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(2, 2, 1) + self.pool = torch.nn.MaxPool2d(1, 1) + + def forward(self, x): + x = self.conv(x) + x = self.pool(x) + return x + + class ConvWithAdaptiveAvgPool2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 3, 3) + self.adaptive_avg_pool2d = torch.nn.AdaptiveAvgPool2d((1, 1)) + + def forward(self, x): + x = self.conv(x) + x = self.adaptive_avg_pool2d(x) + return x + + class ConvWithBNRelu(torch.nn.Module): + def __init__(self, relu, dim=2, bn=True, bias=True, padding=0): + super().__init__() + convs = {1: torch.nn.Conv1d, 2: torch.nn.Conv2d, 3: torch.nn.Conv3d} + bns = { + 1: torch.nn.BatchNorm1d, + 2: torch.nn.BatchNorm2d, + 3: torch.nn.BatchNorm3d, + } + self.conv = convs[dim](3, 3, 3, bias=bias, padding=padding) + + if bn: + self.bn = bns[dim](3) + else: + self.bn = torch.nn.Identity() + if relu: + self.relu = torch.nn.ReLU() + else: + self.relu = torch.nn.Identity() + + def forward(self, x): + x = self.conv(x) + x = self.bn(x) + return self.relu(x) + + class ConvTWithBNRelu(torch.nn.Module): + def __init__(self, relu, dim=2, bn=True, bias=True): + super().__init__() + convts = {1: torch.nn.ConvTranspose1d, 2: torch.nn.ConvTranspose2d} + bns = {1: torch.nn.BatchNorm1d, 2: torch.nn.BatchNorm2d} + self.convt = convts[dim](3, 3, 3, bias=bias) + + if bn: + self.bn = bns[dim](3) + else: + self.bn = torch.nn.Identity() + if relu: + self.relu = torch.nn.ReLU() + else: + self.relu = torch.nn.Identity() + + def forward(self, x): + x = self.convt(x) + x = self.bn(x) + return self.relu(x) + + class Conv2dThenConv1d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1d = torch.nn.Conv1d(3, 3, 3) + self.conv2d = torch.nn.Conv2d(3, 3, 3) + + def forward(self, x): + x = self.conv2d(x) + x = x.squeeze(0) + x = self.conv1d(x) + return x + + def example_inputs(self): + return (torch.randn(1, 3, 5, 5),) + + class Conv2dWithCat(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 3, 3) + self.conv2 = torch.nn.Conv2d(3, 3, 3) + + def forward(self, x, y): + x = self.conv1(x) + y = self.conv2(y) + z = torch.cat([x, y], dim=1) + return z + + class Conv2dWithTwoCat(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 3, 3) + self.conv2 = torch.nn.Conv2d(3, 3, 3) + + def forward(self, x1, x2, x3, x4): + x1 = self.conv1(x1) + x2 = self.conv2(x2) + y = torch.cat([x1, x2], dim=1) + z = x3 + x4 + w = torch.cat([z, y]) + return w + + class Conv2dWithSplit(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 3, 3) + self.conv2 = torch.nn.Conv2d(3, 3, 3) + + def forward(self, x): + x = self.conv1(x) + # use split so we get a list of Tensors + x1, x2 = torch.split(x, 2, dim=1) + y = torch.cat([x1, x2], dim=1) + return y + + def example_inputs(self): + return (torch.randn(1, 3, 16, 16),) + + class ThreeAdd(torch.nn.Module): + def forward(self, x1, x2, x3, x4): + y = x1 + x2 + z = x3 + x4 + w = y + z + return w + + class EmbeddingModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.emb = torch.nn.Embedding(num_embeddings=10, embedding_dim=12) + + def forward(self, indices): + return self.emb(indices) + + class EmbeddingConvLinearModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.emb = torch.nn.Embedding(num_embeddings=10, embedding_dim=8) + self.conv = torch.nn.Conv2d(8, 16, (1, 3)) + self.linear = torch.nn.Linear(16, 8) + + def forward(self, indices): + embeddings = self.emb(indices) + embeddings = torch.unsqueeze(embeddings, dim=0) + embeddings = torch.permute(embeddings, (0, 3, 1, 2)) + conv_out = self.conv(embeddings) + conv_out = torch.permute(conv_out, (0, 2, 3, 1)) + conv_out = torch.squeeze(conv_out, dim=0) + return self.linear(conv_out) + + class AddInplaceAdd(torch.nn.Module): + def forward(self, x, y): + x = x + y + x += y + return x + + class MulInplaceMul(torch.nn.Module): + def forward(self, x, y): + x = x * y + x *= y + return x + + class AddMulScalar(torch.nn.Module): + def forward(self, x): + x = x + 3 + x = x * 3 + x += 3 + x *= 3 + return x + + class ConvBnReLU2dAndLinearReLU(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv_bn_relu = TestHelperModules.ConvWithBNRelu(relu=True) + self.linear = torch.nn.Linear(3, 8, bias=False) + self.relu = torch.nn.ReLU() + + def forward(self, x): + x = self.conv_bn_relu(x) + permute_out = torch.permute(x, (0, 2, 3, 1)) + linear_out = self.linear(permute_out) + return linear_out + + class GroupwiseConv2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(4, 4, 3, groups=2) + + def forward(self, x): + return self.conv(x) + + def example_inputs(self): + return (torch.randn(2, 4, 10, 10),) + + class LinearReluModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc = torch.nn.Linear(5, 5).to(dtype=torch.float) + self.relu = torch.nn.ReLU() + + def forward(self, x): + x = self.relu(self.fc(x)) + return x + +def _static_reference_quantized_linear_module(N, K, bias, example_input): + """ + Generate a linear module with quantize-dequantize (reference quantized) + with static quantization parameters (no choose_qparams at runtime). + A simulation to PT2E quantization in Torchao. + It is used to test fusion and lowering passes in Inductor for X86 CPU. + Input quantization limit is 0-127 to avoid overflow on old platforms. + Params: + N: output feature dimension + K: input feature dimension + bias: boolean flag to indicate whether linear module has bias + example_input: example input tensor to get scale/zero point + Return: + An instance of the reference quantized linear module + """ + class Model(torch.nn.Module): + def __init__(self, N, K, bias, example_input): + super().__init__() + self.x_scale, self.x_zp = torch.ops.quantized_decomposed.choose_qparams.tensor( + example_input, + quant_min=0, + quant_max=127, + eps=torch.finfo(torch.float32).eps, + dtype=torch.uint8, + ) + self.x_scale, self.x_zp = self.x_scale.detach().item(), self.x_zp.detach().item() + self.linear = torch.nn.Linear(K, N, bias) + self.w_scales, self.w_zps = torch.ops.quantized_decomposed.choose_qparams_per_token( + self.linear.weight, dtype=torch.int8 + ) + self.w_scales = self.w_scales.detach().to(torch.float32).squeeze() + self.w_zps = self.w_zps.detach().to(torch.int64).squeeze() + self.qw = torch.ops.quantized_decomposed.quantize_per_channel.default( + self.linear.weight, + self.w_scales, + self.w_zps, + axis=0, + quant_min=-128, + quant_max=127, + dtype=torch.int8, + ) + + def forward(self, x): + dqw = torch.ops.quantized_decomposed.dequantize_per_channel.default( + self.qw, + self.w_scales, + self.w_zps, + axis=0, + quant_min=-128, + quant_max=127, + dtype=torch.int8, + ) + quantize_per_tensor_default = torch.ops.quantized_decomposed.quantize_per_tensor.default( + x, + self.x_scale, + self.x_zp, + quant_min=0, + quant_max=127, + dtype=torch.uint8, + ) + dequantize_per_tensor_default = torch.ops.quantized_decomposed.dequantize_per_tensor.default( + quantize_per_tensor_default, + self.x_scale, + self.x_zp, + quant_min=0, + quant_max=127, + dtype=torch.uint8, + ) + linear = torch.ops.aten.linear.default(dequantize_per_tensor_default, dqw, self.linear.bias) + return linear + + return Model(N, K, bias, example_input).eval() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_quantized.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_quantized.py new file mode 100644 index 0000000000000000000000000000000000000000..829485a1d8334b94e0f8ec4080a48e6b1f011b99 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_quantized.py @@ -0,0 +1,693 @@ +# mypy: ignore-errors + +r"""Importing this file includes common utility methods for checking quantized +tensors and modules. +""" +import numpy as np +import torch +from torch import Tensor +from contextlib import contextmanager +from torch.testing._internal.common_utils import TEST_WITH_TSAN, IS_PPC, IS_MACOS, IS_WINDOWS, IS_ARM64 + +supported_qengines = list(torch.backends.quantized.supported_engines) +# Note: We currently do not run QNNPACK tests on WINDOWS and MACOS as it is flaky. Issue #29326 +# QNNPACK is not supported on PPC +if 'qnnpack' in supported_qengines and any([IS_PPC, TEST_WITH_TSAN, IS_MACOS, IS_WINDOWS]): + supported_qengines.remove('qnnpack') +# FBGEMM and x86 engines require x86 architecture with AVX2/AVX512 support +# They are not supported on ARM64 architectures +if IS_ARM64: + supported_qengines = [qe for qe in supported_qengines if qe not in ('fbgemm', 'x86')] + +def _conv_output_shape(input_size, kernel_size, padding, stride, dilation, + output_padding=0): + """Computes the output shape given convolution parameters.""" + return np.floor((input_size + 2 * padding - kernel_size - (kernel_size - 1) + * (dilation - 1)) / stride) + 2 * output_padding + 1 + +# Quantization references +def _quantize(x, scale, zero_point, qmin=None, qmax=None, dtype=np.uint8): + """Quantizes a numpy array.""" + if qmin is None: + qmin = np.iinfo(dtype).min + if qmax is None: + qmax = np.iinfo(dtype).max + qx = np.round(x / scale + zero_point).astype(np.int64) + qx = np.clip(qx, qmin, qmax) + qx = qx.astype(dtype) + return qx + + +def _dequantize(qx, scale, zero_point): + """Dequantizes a numpy array.""" + x = (qx.astype(float) - zero_point) * scale + return x + + +def _requantize(x, multiplier, zero_point, qmin=0, qmax=255, qtype=np.uint8): + """Requantizes a numpy array, i.e., intermediate int32 or int16 values are + converted back to given type""" + qx = (x * multiplier).round() + zero_point + qx = np.clip(qx, qmin, qmax).astype(qtype) + return qx + +def _calculate_dynamic_qparams(X, dtype, reduce_range=False, qscheme=torch.per_tensor_affine): + """Calculate the dynamic quantization parameters (scale, zero_point) + according to the min and max element of the tensor""" + if qscheme not in (torch.per_tensor_affine, torch.per_tensor_symmetric): + raise AssertionError( + f"Expected qscheme to be per_tensor_affine or per_tensor_symmetric, got {qscheme}" + ) + if qscheme == torch.per_tensor_symmetric: + if dtype != torch.qint8: + raise AssertionError( + f"Expected dtype to be torch.qint8 for symmetric qscheme, got {dtype}" + ) + if isinstance(X, torch.Tensor): + X = X.numpy() + if dtype == torch.qint8: + if reduce_range: + qmin, qmax = -64, 63 + else: + qmin, qmax = -128, 127 + else: # dtype == torch.quint8 + if reduce_range: + qmin, qmax = 0, 127 + else: + qmin, qmax = 0, 255 + min_val = X.min() + max_val = X.max() + is_symmetric = (qscheme == torch.per_tensor_symmetric) + if min_val == max_val: + scale = 1.0 + zero_point = 0 + else: + if is_symmetric: + max_val = max(max_val, -min_val) + min_val = -max_val + scale = (max_val - min_val) / (qmax - qmin) + scale = max(scale, np.finfo(np.float32).eps) + zero_point = 0 + else: + max_val = max(max_val, 0.0) + min_val = min(min_val, 0.0) + scale = (max_val - min_val) / (qmax - qmin) + scale = max(scale, np.finfo(np.float32).eps) + zero_point = qmin - round(min_val / scale) + zero_point = max(qmin, zero_point) + zero_point = min(qmax, zero_point) + return [float(scale), int(zero_point)] + +def _calculate_dynamic_per_channel_qparams(X, dtype): + """Calculate the dynamic quantization parameters (scale, zero_point) + according to the min and max element of the tensor""" + if isinstance(X, torch.Tensor): + X = X.numpy() + qmin, qmax = torch.iinfo(dtype).min, torch.iinfo(dtype).max + n_levels = qmax - qmin + scale = np.zeros(X.shape[0], dtype=np.float64) + zero_point = np.zeros(X.shape[0], dtype=np.int64) + for i in range(zero_point.shape[0]): + min_val = X.min() + max_val = X.max() + if min_val == max_val: + scale[i] = 1.0 + zero_point[i] = 0 + else: + max_val = max(max_val, 0.0) + min_val = min(min_val, 0.0) + scale[i] = (max_val - min_val) / n_levels + scale[i] = max(scale[i], np.finfo(np.float32).eps) + zero_point[i] = qmin - round(min_val / scale[i]) + zero_point[i] = max(qmin, zero_point[i]) + zero_point[i] = min(qmax, zero_point[i]) + + return scale, zero_point + +def _snr(x, x_hat): + """Calculates the signal to noise ratio and returns the signal and noise + power, as well as the SNR in dB. + If the input is a list/tuple this function is called recursively on each + element. The result will have the same nested structure as the inputs. + + Args: + x, x_hat: Either a tensor or a nested list/tuple of tensors. + Returns: + signal, noise, SNR(in dB): Either floats or a nested list of floats + """ + if isinstance(x, (list, tuple)): + if len(x) != len(x_hat): + raise AssertionError(f"Expected len(x) == len(x_hat), got {len(x)} != {len(x_hat)}") + res = [_snr(x[idx], x_hat[idx]) for idx in range(len(x))] + return res + if x_hat.is_quantized: + x_hat = x_hat.dequantize() + if x.is_quantized: + x = x.dequantize() + noise = (x - x_hat).norm() + if noise == 0: + return 0.0, float('inf'), float('inf') + signal = x.norm() + snr = signal / noise + snr_db = 20 * snr.log10() + return signal, noise, snr_db + +@contextmanager +def override_quantized_engine(qengine): + previous = torch.backends.quantized.engine + torch.backends.quantized.engine = qengine + try: + yield + finally: + torch.backends.quantized.engine = previous + +@contextmanager +def override_cpu_allocator_for_qnnpack(qengine_is_qnnpack): + try: + if qengine_is_qnnpack: + torch._C._set_default_mobile_cpu_allocator() + yield + finally: + if qengine_is_qnnpack: + torch._C._unset_default_mobile_cpu_allocator() + +# TODO: Update all quantization tests to use this decorator. +# Currently for some of the tests it seems to have inconsistent params +# for fbgemm vs qnnpack. +def override_qengines(qfunction): + def test_fn(*args, **kwargs): + for qengine in supported_qengines: + with override_quantized_engine(qengine): + # qfunction should not return anything. + qfunction(*args, **kwargs) + return test_fn + +def qengine_is_fbgemm(): + return torch.backends.quantized.engine == 'fbgemm' +def qengine_is_qnnpack(): + return torch.backends.quantized.engine == 'qnnpack' +def qengine_is_onednn(): + return torch.backends.quantized.engine == 'onednn' +def qengine_is_x86(): + return torch.backends.quantized.engine == 'x86' + +# Helper function used to simulate per-channel fake-quant against any axis +def _permute_to_axis_zero(X, axis): + new_axis_list = list(range(X.dim())) + new_axis_list[axis] = 0 + new_axis_list[0] = axis + y = X.permute(tuple(new_axis_list)) + return y, new_axis_list + +# Reference method for fake quantize +# Note: because scale/zero_point are left as float in the actual kernel, this mimics how fake_quant works for float16/64 +def _fake_quantize_per_channel_affine_reference(X, per_channel_scale, per_channel_zero_point, axis, quant_min, quant_max): + dtype = X.dtype + X, permute_axis_list = _permute_to_axis_zero(X.to(torch.float32), axis) + res = torch.zeros_like(X) + + for i in range(X.size()[0]): + res[i] = (torch.clamp(torch.round(X[i] * (1.0 / per_channel_scale[i]) + + per_channel_zero_point[i]), quant_min, quant_max) - per_channel_zero_point[i]) * per_channel_scale[i] + + out = res.permute(tuple(permute_axis_list)) + return out.to(dtype) + +# Reference method for the gradient of the fake quantize operator +# Note: because scale/zero_point are left as float in the actual kernel, this mimics how fake_quant works for float16/64 +def _fake_quantize_per_channel_affine_grad_reference(dY, X, per_channel_scale, per_channel_zero_point, axis, quant_min, quant_max): + dtype = X.dtype + X, permute_axis_list = _permute_to_axis_zero(X.to(torch.float32), axis) + Xq = torch.zeros_like(X) + for i in range(X.size()[0]): + Xq[i] = torch.round(X[i] * (1.0 / per_channel_scale[i]) + per_channel_zero_point[i]) + Xq = Xq.permute(tuple(permute_axis_list)) + mask = (Xq >= quant_min) * (Xq <= quant_max) + res = torch.zeros_like(dY) + res[mask] = dY[mask] + return res.to(dtype) + +def to_tensor(X, device): + if not isinstance(X, torch.Tensor): + X = torch.tensor(X) + else: + X = X.detach().clone() + return X.to(device=torch.device(device), dtype=torch.float32) + +# copy-pasted from +# https://github.com/pytorch/ao/blob/bc4f51da86956275da7db0da6e420c506df97820/torchao/prototype/custom_fp_utils.py#L27C1-L142C29 +def _n_ones(n: int) -> int: + return (1 << n) - 1 + +EBITS_F32, MBITS_F32 = 8, 23 +F32_EXP_BIAS = _n_ones(EBITS_F32 - 1) + +# copy-pasted from +# https://github.com/pytorch/ao/blob/bc4f51da86956275da7db0da6e420c506df97820/torchao/prototype/custom_fp_utils.py#L27C1-L142C29 +def _f32_to_floatx_unpacked(x: Tensor, ebits: int, mbits: int) -> Tensor: + """Convert FP32 numbers to sub-byte floating point numbers with the given + number of exponent and mantissa bits. + + Input: torch.Tensor of dtype torch.float + Output: torch.Tensor of dtype torch.uint8, where the bit encoding is stored + in the least significant bits. e.g. + fp4: bits 0-3 empty and bits 4-7 in fp4_e2m1 encoding + fp6: bits 0-1 empty and bits 2-7 in fp6_e2m3 or fp6_e3m2 encoding + + Note: there are no special values (NaN, inf) support in this code. Values + outside the representable range of Floatx after rounding are clamped to the + maximum Floatx magnitude (sign is preserved). + + Code below is an adaptation of https://fburl.com/code/ciwofcg4 + + Background 1: last answer in https://stackoverflow.com/q/8981913 + Background 2: Computer Organization and Design, RISC-V edition, Chapter 3.5 + """ + if x.dtype != torch.float: + raise AssertionError(f"Expected x.dtype to be torch.float, got {x.dtype}") + if 1 + ebits + mbits > 8: + raise AssertionError(f"Expected 1 + ebits + mbits <= 8, got {1 + ebits + mbits}") + + # calculate constants + exp_bias = _n_ones(ebits - 1) + max_int = _n_ones(ebits + mbits) + sign_mask = 1 << (ebits + mbits) + + # TODO document this better + magic_adder = _n_ones(MBITS_F32 - mbits - 1) + + # all E bits and M bits are 1s + max_normal = 2 ** (_n_ones(ebits) - exp_bias) * (_n_ones(mbits + 1) / (2**mbits)) + + # E bits = 1, M bits = 0 + min_normal = 2 ** (1 - exp_bias) + + denorm_exp = ( + # exp bias conversion between formats + (F32_EXP_BIAS - exp_bias) + # mantissa length difference between formats + + (MBITS_F32 - mbits) + # add one to encoded exponent for denormalized numbers + + 1 + ) + denorm_mask_int = denorm_exp << MBITS_F32 + + # reinterpret int32 as float32 + denorm_mask_float = torch.tensor(denorm_mask_int, dtype=torch.int32).view( + torch.float32 + ) + + # save the sign + # Note that we have torch.uint32, but some ops like cpu bit shifts + # do not work on it. So, we stay in int32. + x = x.view(torch.int32) + sign = x & 0x80000000 + + # set everything to positive, will add sign back at the end + x = x ^ sign + + # TODO: can the branch floating point comparisons below be done without + # converting to float? probably but need to verify + x = x.view(torch.float) + + # rewrite saturate/denorm/norm branches without explicit data dependent + # control flow, to be more compiler friendly + saturate_mask = x >= max_normal + denormal_mask = torch.logical_and(torch.logical_not(saturate_mask), x < min_normal) + normal_mask = torch.logical_not(torch.logical_or(saturate_mask, denormal_mask)) + + # + # branch 1: saturate to max val - handled later in the code which combines + # the branches + # + + # + # branch 2: to conversion to denormal as well as rounding up to normal + # + denormal_x = x + denorm_mask_float + denormal_x = denormal_x.view(torch.int32) + denormal_x -= denorm_mask_int + denormal_x = denormal_x.to(torch.uint8) + + # + # branch 3: stay in normal range, adjust the exponent and round + # + normal_x = x.view(torch.int32) + # resulting mantissa is odd + mant_odd = (normal_x >> (MBITS_F32 - mbits)) & 1 + # update exponent, rounding bias part 1 + val_to_add = ((exp_bias - F32_EXP_BIAS) << MBITS_F32) + magic_adder + normal_x += val_to_add + # rounding bias part 2 + normal_x += mant_odd + # take the bits! + normal_x = normal_x >> (MBITS_F32 - mbits) + normal_x = normal_x.to(torch.uint8) + + # + # combine the branches + # + x = torch.full_like(x, max_int, dtype=torch.uint8) + x = torch.where(denormal_mask, denormal_x, x) + x = torch.where(normal_mask, normal_x, x) + + # add sign back + sign_lp = sign >> (MBITS_F32 + EBITS_F32 - mbits - ebits) + sign_lp = sign_lp.to(torch.uint8) + # Right shift of a negative signed integer can fill the least significant + # bits with either 1s or 0s, depending on the implementation. Since PyTorch + # doesn't have an uint32 dtype, we mask out these bits to get just the + # f4 sign bit + sign_lp = sign_lp & sign_mask + x = x | sign_lp + + return x.to(torch.uint8) + + +# copy-pasted from +# https://github.com/pytorch/ao/blob/29488018d99af7f7339f06353c6b5bbeae8a1493/torchao/prototype/custom_fp_utils.py#L147 +def _floatx_unpacked_to_f32(x: Tensor, ebits: int, mbits: int) -> Tensor: + """Convert sub-byte floating point numbers with the given number of exponent + and mantissa bits to FP32. + + Input: torch.Tensor of dtype uint8, where the bit encoding is stored + in the least significant bits. e.g. + fp4: bits 0-3 empty and bits 4-7 in fp4_e2m1 encoding + fp6: bits 0-1 empty and bits 2-7 in fp6_e2m3 or fp6_e3m2 encoding + Output: torch.Tensor of dtype fp32 with the dequantized value + """ + if x.dtype != torch.uint8: + raise AssertionError(f"Expected x.dtype to be torch.uint8, got {x.dtype}") + if 1 + ebits + mbits > 8: + raise AssertionError(f"Expected 1 + ebits + mbits <= 8, got {1 + ebits + mbits}") + + sign_mask = 1 << (ebits + mbits) + exp_bias = _n_ones(ebits - 1) + mantissa_mask = _n_ones(mbits) + + # save the sign + sign_lp = x & sign_mask + + # set everything to positive, will add sign back at the end + x_pos = x ^ sign_lp + + # + # 1. Calculate zero mask + # + zero_mask = x_pos == 0 + + # + # 2. Calculate the denormal path mask + # + denormal_mask = torch.logical_and((x_pos > 0), ((x_pos >> mbits) == 0)) + + # + # 3. Calculate the normal path + # + + # calculate the new exponent and shift it to bits 2:9 of the result + exp_biased_lp = x_pos >> mbits + exp_biased_f32 = exp_biased_lp - exp_bias + F32_EXP_BIAS + exp_biased_f32 = exp_biased_f32.to(torch.int32) << MBITS_F32 + + # shift the mantissa to bits 10:32 of the result + mantissa_lp_int32 = (x_pos & mantissa_mask).to(torch.int32) + mantissa_f32 = mantissa_lp_int32 << (MBITS_F32 - mbits) + result = exp_biased_f32 | mantissa_f32 + + # + # 4. Add the zero and denormal casts to the already casted normal path + # + result[zero_mask] = 0 + + denormal_exp_biased = 1 - exp_bias + F32_EXP_BIAS + + # fast path. + # without this, performance for FP4_E2M1 is slower by 2x + if mbits == 1: + result[denormal_mask] = (denormal_exp_biased - mbits) << MBITS_F32 + + else: + # iterate over all possible values of mantissa + # i=0, j=1 + # i=1, j=10,11 + # i=2, j=100,101,110,111 + # and so on + for i in range(mbits): + for mantissa_cmp in range(1 << i, 1 << (i + 1)): + # left shift mantissa until it overflows (create an implicit 1) + # subtract exponent by the same amount + left_shift = mbits - i + mantissa_f32 = (mantissa_cmp - (1 << i)) << ( + left_shift + MBITS_F32 - mbits + ) + exp_biased_f32 = (denormal_exp_biased - left_shift) << MBITS_F32 + + # we can update this in-place since the values won't overlap + # torch.compile() may complain unsupported operand type(s) for |: 'SymInt' and 'int' + # thus we use + instead of | here + mantissa_lp_int32[mantissa_lp_int32 == mantissa_cmp] = ( + exp_biased_f32 + mantissa_f32 + ) + + result = torch.where(denormal_mask, mantissa_lp_int32, result) + + # add sign back + sign_f32 = sign_lp.to(torch.int32) << (MBITS_F32 - mbits + EBITS_F32 - ebits) + result = result | sign_f32 + + return result.view(torch.float) + +# copied from https://github.com/drisspg/transformer_nuggets/blob/main/transformer_nuggets/mx/to_blocked.py +def ceil_div(a, b): + return (a + b - 1) // b + +# NVIDIA Blackwell HW requires scales for MX/NV blocked formats to be in a 128x4 tile layout, +# with a weird 32x4x4 internal layout of that tile. If we want to take swizzled scales and use them +# for non-gemm purposes (like testing), we need to de-swizzle them, then they can be applied much +# more naturally. +def from_blocked(input, input_scales, blocksize) -> torch.Tensor: + # Matrix is in a 128x4 pattern, internally blocked as 32x4x4 nonsense. + # Output should be [input.size(0, input.size(1) // blocksize] scales + output_scales = torch.zeros( + (input.size(0), input.size(1) // blocksize), + device=input.device, + dtype=input_scales.dtype, + ) + + # Swizzled scales are padded to tiles of 128x4, we need to replicate how that padding + # happened for offset purposes. + # There are K//blocksize scales, padded to groups of 4. + num_col_tiles = ceil_div(ceil_div(input.size(1), blocksize), 4) + + # (Very) slow reference implementation using horrifying loops. + for i in range(input.size(0)): + for j in range(input.size(1) // blocksize): + # which 128x4 tile of scaling factors am I in + scale_tile_h = i // 128 + scale_tile_w = j // 4 + + # There are (padded) input_scales.size(1) // 4 tiles along the w dim. + # So offset is 512 * (h_tile * tiles_per_row + tile_in_row) + tile_offset = 512 * (scale_tile_h * num_col_tiles + scale_tile_w) + + # indices within the tile - use nomenclature directly from cublas docs + outer = i % 128 # "outer" in cublas docs + inner = j % 4 # "inner" in cublas docs + + # Note: "offset" is given in terms of bytes, in cublas docs, but our scales are e8m0, + # anyway, and so 1B == 1 value => use offset directly. + # Formula directly from cublas docs in 3.1.4.3.2 + offset = tile_offset + (outer % 32) * 16 + (outer // 32) * 4 + inner + + output_scales[i, j] = input_scales[offset] + + return output_scales + +def from_blocked_format(x_mxfp8, scales_unswizzled, blocksize=32): + # expand scales + scales = torch.repeat_interleave(scales_unswizzled, blocksize, dim=1) + + # de-scale and convert + x_f32 = x_mxfp8.to(torch.float) * scales.to(torch.float) + return x_f32.to(torch.bfloat16) + +def to_blocked(input_matrix) -> torch.Tensor: + """ + Rearrange a large matrix by breaking it into blocks and applying the rearrangement pattern. + + See: + https://docs.nvidia.com/cuda/cublas/index.html#d-block-scaling-factors-layout + + Args: + input_matrix: Input tensor of shape (H, W) + + Returns: + Rearranged tensor of shape (32*ceil_div(H,128), 16*ceil_div(W,4)) + """ + rows, cols = input_matrix.shape + n_row_blocks = ceil_div(rows, 128) + n_col_blocks = ceil_div(cols, 4) + + # Calculate the padded shape + padded_rows = n_row_blocks * 128 + padded_cols = n_col_blocks * 4 + + padded = input_matrix + # Ideally we would use torch.nn.pad but it doesn't support float8_e8m0fnu for now + if (rows, cols) != (padded_rows, padded_cols): + padded = torch.zeros((padded_rows, padded_cols), device=input_matrix.device, dtype=input_matrix.dtype) + padded[:rows, :cols] = input_matrix + + # Rearrange the blocks + blocks = padded.view(n_row_blocks, 128, n_col_blocks, 4).permute(0, 2, 1, 3) + rearranged = blocks.reshape(-1, 4, 32, 4).transpose(1, 2).reshape(-1, 32, 16) + + return rearranged.flatten() + + +def down_size(size): + if size[-1] % 2 != 0: + raise AssertionError(f"{size} last dim not divisible by two") + return (*size[:-1], size[-1] // 2) + + +def pack_uint4(uint8_data) -> torch.Tensor: + # converting to uint8 for operations + shape = uint8_data.shape + if shape[-1] % 2 != 0: + raise AssertionError(f"Expected shape[-1] to be divisible by 2, got {shape[-1]}") + uint8_data = uint8_data.contiguous().view(-1) + return (uint8_data[1::2] << 4 | uint8_data[::2]).view(down_size(shape)) + + +# exponent and mantissa bits of `torch.float4_e2m1fn_x2` +FP4_EBITS, FP4_MBITS = 2, 1 + + +def _bfloat16_to_float4_e2m1fn_x2(x): + if x.dtype != torch.bfloat16: + raise AssertionError(f"Expected x.dtype to be torch.bfloat16, got {x.dtype}") + x = _f32_to_floatx_unpacked(x.float(), FP4_EBITS, FP4_MBITS) + x = pack_uint4(x) + x = x.view(torch.float4_e2m1fn_x2) + return x + + +# This function is extracted from https://github.com/pytorch/ao/blob/v0.12.0/torchao/prototype/mx_formats/mx_tensor.py#L142 +def to_mxfp( + data_hp: torch.Tensor, + block_size: int = 32, + format: str = "mxfp8", +): + if data_hp.dtype not in (torch.bfloat16, torch.float): + raise AssertionError(f"{data_hp.dtype} is not supported yet") + if data_hp.shape[-1] % block_size != 0: + raise AssertionError( + f"the last dimension of shape {data_hp.shape} must be divisible by block_size {block_size}" + ) + if not data_hp.is_contiguous(): + raise AssertionError("unsupported: data_hp must be contiguous") + + orig_shape = data_hp.shape + data_hp = data_hp.reshape( + *orig_shape[:-1], orig_shape[-1] // block_size, block_size + ) + + max_abs = torch.amax(torch.abs(data_hp), -1).unsqueeze(-1) + + data_hp = data_hp.to(torch.float32) + max_abs = max_abs.to(torch.float32) + + if format == "mxfp8": + F8E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max # 448.0 + max_pos = F8E4M3_MAX + elif format == "mxfp4": + F4E2M1_MAX = 6. + max_pos = F4E2M1_MAX + + # RCEIL + def _to_mx_rceil( + data_hp: torch.Tensor, + max_abs: torch.Tensor, + max_pos: float, + ) -> tuple[torch.Tensor, torch.Tensor]: + E8M0_EXPONENT_BIAS = 127 + descale = max_abs / max_pos + exponent = torch.where( + torch.isnan(descale), + 0xFF, # Handle biased exponent for nan + # NOTE: descale < (torch.finfo(torch.float32).smallest_normal / 2) is handled through clamping + ( + torch.clamp( + torch.ceil(torch.log2(descale)), + min=-E8M0_EXPONENT_BIAS, + max=E8M0_EXPONENT_BIAS, + ) + + E8M0_EXPONENT_BIAS + ).to(torch.uint8), + ) + + descale_fp = torch.where( + exponent == 0, + 1.0, + torch.exp2(E8M0_EXPONENT_BIAS - exponent.to(torch.float32)), + ) + + # scale and saturated cast the data elements to max of target dtype + data_lp = torch.clamp(data_hp * descale_fp, min=-1 * max_pos, max=max_pos) + return exponent, data_lp + + scale_e8m0_biased, data_lp = _to_mx_rceil(data_hp, max_abs, max_pos) + + # cast to target dtype + if format == "mxfp8": + data_lp = data_lp.to(torch.float8_e4m3fn) + # need to reshape at the end to help inductor fuse things + data_lp = data_lp.reshape(orig_shape) + elif format == "mxfp4": + data_lp = _bfloat16_to_float4_e2m1fn_x2(data_lp.to(torch.bfloat16)) + final_shape = list(orig_shape) + final_shape[-1] //= 2 + data_lp = data_lp.reshape(final_shape) + + scale_e8m0_biased = scale_e8m0_biased.view(torch.float8_e8m0fnu) + scale_e8m0_biased = scale_e8m0_biased.squeeze(-1) + return scale_e8m0_biased, data_lp + +# Source: https://github.com/pytorch/ao/blob/568c1932a16ae9f30d48da214a88dc0013e98ed8/torchao/prototype/moe_training/utils.py#L310 +def generate_jagged_offs(E, M, multiple_of=16, dtype=torch.int32, device="cuda"): + """ + Utility function for tests and benchmarks. + + Generates a tensor of length E, containing random values divisible by `multiple_of`, + from 0 to M, in sorted order, and where the final value in the tensor is always M. + Args: + E (int): The length of the tensor. + M (int): The maximum value in the tensor. + Returns: + torch.Tensor: A tensor of length E with the specified properties. + """ + import random + + # Ensure M is divisible by 16 + if M % multiple_of != 0: + raise ValueError(f"M must be divisible by {multiple_of}") + + # Generate a list of possible values + possible_values = list(range(multiple_of, M + 1, multiple_of)) + + # If E is larger than the number of possible values, raise an error + if E > len(possible_values): + raise ValueError("E cannot be larger than the number of possible values") + + # Randomly select E - 1 values from the possible values (excluding M) + selected_values = torch.tensor(random.sample(possible_values[:-1], E - 1)) + + # Append M to the selected values + selected_values = torch.cat((selected_values, torch.tensor([M]))) + + # Sort the selected values + selected_values, _ = torch.sort(selected_values) + + return selected_values.to(dtype).to(device) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_subclass.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_subclass.py new file mode 100644 index 0000000000000000000000000000000000000000..b596256569fb385e10c3428f7684f003206cdd12 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_subclass.py @@ -0,0 +1,345 @@ +# mypy: ignore-errors + +import torch +from copy import deepcopy +from torch.utils._pytree import tree_map +import torch.utils._pytree as pytree + + +# TODO: Move LoggingTensor here. +from torch.testing._internal.logging_tensor import LoggingTensor + + +# Base class for wrapper-style tensors. +class WrapperTensor(torch.Tensor): + @staticmethod + def __new__(cls, *args, **kwargs): + t, kwargs = cls.get_wrapper_properties(*args, **kwargs) + if "size" not in kwargs: + size = t.size() + else: + size = kwargs["size"] + del kwargs["size"] + if "dtype" not in kwargs: + kwargs["dtype"] = t.dtype + if "layout" not in kwargs: + kwargs["layout"] = t.layout + if "device" not in kwargs: + kwargs["device"] = t.device + if "requires_grad" not in kwargs: + kwargs["requires_grad"] = False + # Ignore memory_format and pin memory for now as I don't know how to + # safely access them on a Tensor (if possible??) + + wrapper = torch.Tensor._make_wrapper_subclass(cls, size, **kwargs) + wrapper._validate_methods() + return wrapper + + @classmethod + def get_wrapper_properties(cls, *args, **kwargs): + # Should return both an example Tensor and a dictionary of kwargs + # to override any of that example Tensor's properly. + # This is very similar to the `t.new_*(args)` API + raise NotImplementedError("You need to implement get_wrapper_properties") + + def _validate_methods(self): + # Skip this if not in debug mode? + # Changing these on the python side is wrong as it would not be properly reflected + # on the c++ side + # This doesn't catch attributes set in the __init__ + forbidden_overrides = ["size", "stride", "dtype", "layout", "device", "requires_grad"] + for el in forbidden_overrides: + if getattr(self.__class__, el) is not getattr(torch.Tensor, el): + raise RuntimeError(f"Subclass {self.__class__.__name__} is overwriting the " + f"property {el} but this is not allowed as such change would " + "not be reflected to c++ callers.") + + +class WrapperTensorWithCustomSizes(WrapperTensor): + @classmethod + def get_wrapper_properties(cls, t, requires_grad=False): + return t, {"requires_grad": requires_grad, "dispatch_sizes_strides_policy": "sizes"} + + def __init__(self, t, requires_grad=False): + self.t = t + + @classmethod + def __torch_dispatch__(cls, func, types, args=(), kwargs=None): + if not all(issubclass(cls, t) for t in types): + return NotImplemented + + if kwargs is None: + kwargs = {} + + def unwrap(e): + return e.t if isinstance(e, WrapperTensorWithCustomSizes) else e + + def wrap(e): + return WrapperTensorWithCustomSizes(e) if isinstance(e, torch.Tensor) else e + + rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs or {}))) + return rs + + def __repr__(self): + return super().__repr__(tensor_contents=f"t={self.t}") + + +class WrapperTensorWithCustomStrides(WrapperTensor): + @classmethod + def get_wrapper_properties(cls, t, requires_grad=False): + return t, {"requires_grad": requires_grad, "dispatch_sizes_strides_policy": "strides"} + + def __init__(self, t, requires_grad=False): + self.t = t + + @classmethod + def __torch_dispatch__(cls, func, types, args=(), kwargs=None): + if not all(issubclass(cls, t) for t in types): + return NotImplemented + + if kwargs is None: + kwargs = {} + + def unwrap(e): + return e.t if isinstance(e, WrapperTensorWithCustomStrides) else e + + def wrap(e): + return WrapperTensorWithCustomStrides(e) if isinstance(e, torch.Tensor) else e + + rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs or {}))) + return rs + + def __repr__(self): + return super().__repr__(tensor_contents=f"t={self.t}") + + +class DiagTensorBelow(WrapperTensor): + @classmethod + def get_wrapper_properties(cls, diag, requires_grad=False): + if diag.ndim != 1: + raise AssertionError(f"Expected diag.ndim == 1, got {diag.ndim}") + return diag, {"size": diag.size() + diag.size(), "requires_grad": requires_grad} + + def __init__(self, diag, requires_grad=False): + self.diag = diag + + handled_ops = {} + + @classmethod + def __torch_dispatch__(cls, func, types, args=(), kwargs=None): + if not all(issubclass(cls, t) for t in types): + return NotImplemented + + # For everything else, call the handler: + fn = cls.handled_ops.get(func.__name__, None) + if fn: + return fn(*args, **(kwargs or {})) + else: + # Note that here, because we don't need to provide the autograd formulas + # we can have a default "fallback" that creates a plain Tensor based + # on the diag elements and calls the func again. + + def unwrap(e): + return e.diag.diag() if isinstance(e, DiagTensorBelow) else e + + def wrap(e): + if isinstance(e, torch.Tensor) and e.ndim == 1: + return DiagTensorBelow(e) + if isinstance(e, torch.Tensor) and e.ndim == 2 and e.count_nonzero() == e.diag().count_nonzero(): + return DiagTensorBelow(e.diag()) + return e + + rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs or {}))) + return rs + + def __repr__(self): + return super().__repr__(tensor_contents=f"diag={self.diag}") + + +class SparseTensor(WrapperTensor): + @classmethod + def get_wrapper_properties(cls, size, values, indices, requires_grad=False): + if values.device != indices.device: + raise AssertionError(f"Expected values.device == indices.device, got {values.device} != {indices.device}") + return values, {"size": size, "requires_grad": requires_grad} + + def __init__(self, size, values, indices, requires_grad=False): + self.values = values + self.indices = indices + + def __repr__(self): + return super().__repr__(tensor_contents=f"values={self.values}, indices={self.indices}") + + def sparse_to_dense(self): + res = torch.zeros(self.size(), dtype=self.values.dtype) + res[self.indices.unbind(1)] = self.values + return res + + @staticmethod + def from_dense(t): + indices = t.nonzero() + values = t[indices.unbind(1)] + return SparseTensor(t.size(), values, indices) + + @classmethod + def __torch_dispatch__(cls, func, types, args=(), kwargs=None): + func_name = f"{func.__module__}.{func.__name__}" + + res = cls._try_call_special_impl(func_name, args, kwargs) + if res is not NotImplemented: + return res + + # Otherwise, use a default implementation that construct dense + # tensors and use that to compute values + def unwrap(e): + return e.sparse_to_dense() if isinstance(e, SparseTensor) else e + + # Wrap back all Tensors into our custom class + def wrap(e): + # Check for zeros and use that to get indices + return SparseTensor.from_dense(e) if isinstance(e, torch.Tensor) else e + + rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs or {}))) + return rs + + + _SPECIAL_IMPLS = {} + + @classmethod + def _try_call_special_impl(cls, func, args, kwargs): + if func not in cls._SPECIAL_IMPLS: + return NotImplemented + return cls._SPECIAL_IMPLS[func](args, kwargs) + + +# Example non-wrapper subclass that stores extra state. +class NonWrapperTensor(torch.Tensor): + def __new__(cls, data): + t = torch.Tensor._make_subclass(cls, data) + t.extra_state = { + 'last_func_called': None + } + return t + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + result = super().__torch_function__(func, types, args, kwargs) + + if isinstance(result, cls): + # Do something with the extra state. For the example here, just store the name of the + # last function called (skip for deepcopy so the copy has the same extra state). + if func is torch.Tensor.__deepcopy__: + result.extra_state = deepcopy(args[0].extra_state) + else: + result.extra_state = { + 'last_func_called': func.__name__, + } + + return result + + # new_empty() must be defined for deepcopy to work + def new_empty(self, shape): + return type(self)(torch.empty(shape)) + + +# Class used to store info about subclass tensors used in testing. +class SubclassInfo: + + __slots__ = ['name', 'create_fn', 'closed_under_ops'] + + def __init__(self, name, create_fn, closed_under_ops=True): + self.name = name + self.create_fn = create_fn # create_fn(shape) -> tensor instance + self.closed_under_ops = closed_under_ops + + +# Helper function to create a subclass of the given class and possibly cache sizes / strides. +def _create_and_access_shape(cls, shape): + sub = cls(torch.randn(shape)) + # NB: Wrapper subclasses with custom dispatched sizes / strides cache this info + # on the first call via non-serializable PyCapsules. We purposefully trigger cache + # population here for serialization / deepcopy tests to verify that the presence of this + # cache info doesn't cause problems. + sub.size() + sub.stride() + return sub + + +subclass_db = { + torch.Tensor: SubclassInfo( + 'base_tensor', create_fn=torch.randn + ), + NonWrapperTensor: SubclassInfo( + 'non_wrapper_tensor', + create_fn=lambda shape: NonWrapperTensor(torch.randn(shape)) + ), + LoggingTensor: SubclassInfo( + 'logging_tensor', + create_fn=lambda shape: LoggingTensor(torch.randn(shape)) + ), + SparseTensor: SubclassInfo( + 'sparse_tensor', + create_fn=lambda shape: SparseTensor.from_dense(torch.randn(shape).relu()) + ), + DiagTensorBelow: SubclassInfo( + 'diag_tensor_below', + create_fn=lambda shape: DiagTensorBelow(torch.randn(shape)), + closed_under_ops=False # sparse semantics + ), + WrapperTensorWithCustomSizes: SubclassInfo( + 'wrapper_with_custom_sizes', + create_fn=lambda shape: _create_and_access_shape(WrapperTensorWithCustomSizes, shape), + closed_under_ops=False, + ), + WrapperTensorWithCustomStrides: SubclassInfo( + 'wrapper_with_custom_strides', + create_fn=lambda shape: _create_and_access_shape(WrapperTensorWithCustomStrides, shape), + closed_under_ops=False, + ), +} + +class SubclassWithTensorFactory(torch.Tensor): + @staticmethod + def __new__(cls, src): + shape = src.shape + kwargs = {} + kwargs["strides"] = src.stride() + kwargs["storage_offset"] = src.storage_offset() + kwargs["device"] = src.device + kwargs["layout"] = src.layout + kwargs["requires_grad"] = src.requires_grad + kwargs["dtype"] = src.dtype + out = torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) + return out + + def __init__(self, src): + self.src = src + + def __repr__(self): + return f"{self.__class__.__name__}" + + def __tensor_flatten__(self): + return ["src"], None + + @classmethod + def __tensor_unflatten__(cls, inner_tensors, meta, outer_size, outer_stride): + src = inner_tensors["src"] + return cls(src) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + if kwargs is None: + kwargs = {} + + def _fn(x): + return x.src * torch.ones(x.src.shape) if x.src.dtype == torch.float32 else x.src + + _args = pytree.tree_map_only(cls, _fn, args) + _kwargs = pytree.tree_map_only(cls, _fn, kwargs) + + _out = func(*_args, **_kwargs) + + _out_flat, _out_spec = pytree.tree_flatten(_out) + + out_flat = [cls(o) if isinstance(o, torch.Tensor) else o for o in _out_flat] + return pytree.tree_unflatten(out_flat, _out_spec) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..193d625c5e9ce66acd59587b873b811061fc2779 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_utils.py @@ -0,0 +1,6142 @@ +# mypy: allow-untyped-defs + +r"""Importing this file must **not** initialize CUDA context. test_distributed +relies on this assumption to properly run. This means that when this is imported +no CUDA calls shall be made, including torch.cuda.device_count(), etc. + +torch.testing._internal.common_cuda.py can freely initialize CUDA context when imported. +""" + +import sysconfig +import argparse +import contextlib +import copy +import ctypes +import errno +import functools +import gc +import hashlib +import inspect +import io +import json +import logging +import math +import operator +import os +import pathlib +import platform +import random +import re +import shutil +import signal +import socket +import subprocess +import sys +import tempfile +import threading +import time +import types +import unittest +import warnings +from collections.abc import Mapping, Sequence +from contextlib import closing, contextmanager +from copy import deepcopy +from dataclasses import dataclass +from enum import Enum +from functools import partial, wraps +from itertools import product, chain +from pathlib import Path +from statistics import mean +from typing import ( + Any, + TypeVar, +) +from collections.abc import Callable +from collections.abc import Iterable, Iterator +from unittest.mock import MagicMock + +import expecttest +import numpy as np + +import __main__ # type: ignore[import] +import torch +import torch.backends.cudnn +import torch.backends.mkl +import torch.backends.mps +import torch.backends.xnnpack +import torch.cuda +from torch import Tensor +from torch._C import ScriptDict, ScriptList # type: ignore[attr-defined] +from torch._utils_internal import get_writable_path +from torch._logging.scribe import open_source_signpost +from torch.nn import ( + ModuleDict, + ModuleList, + ParameterDict, + ParameterList, + Sequential, +) +from torch.onnx import ( + register_custom_op_symbolic, + unregister_custom_op_symbolic, +) +from torch.testing import make_tensor +from torch.testing._comparison import ( + _unwrap_dtensor_for_comparison, + BooleanPair, + NonePair, + not_close_error_metas, + NumberPair, + Pair, + TensorLikePair, +) +from torch.testing._internal.common_dtype import get_all_dtypes +from torch.utils._import_utils import _check_module_exists +import torch.utils._pytree as pytree +from torch.utils import cpp_extension +from torch._utils import _is_privateuse1_backend_available +try: + import pytest # type: ignore[import-not-found] + has_pytest = True +except ImportError: + has_pytest = False + +SEED = 1234 +MI350_ARCH = ("gfx950",) +MI300_ARCH = ("gfx942",) +MI200_ARCH = ("gfx90a",) +NAVI_ARCH = ("gfx1030", "gfx1100", "gfx1101", "gfx1200", "gfx1201") +NAVI3_ARCH = ("gfx1100", "gfx1101") +NAVI4_ARCH = ("gfx1200", "gfx1201") + +class ProfilingMode(Enum): + LEGACY = 1 + SIMPLE = 2 + PROFILING = 3 + +# Set by parse_cmd_line_args() if called +DISABLED_TESTS_FILE = "" +GRAPH_EXECUTOR : ProfilingMode | None = None +LOG_SUFFIX = "" +PYTEST_SINGLE_TEST = "" +REPEAT_COUNT = 0 +RERUN_DISABLED_TESTS = False +RUN_PARALLEL = 0 +SHOWLOCALS = False +SLOW_TESTS_FILE = "" +TEST_BAILOUTS = False +TEST_DISCOVER = False +TEST_IN_SUBPROCESS = False +TEST_SAVE_XML = "" +UNITTEST_ARGS : list[str] = [] +USE_PYTEST = False + +def is_navi3_arch(): + if torch.cuda.is_available(): + prop = torch.cuda.get_device_properties(0) + gfx_arch = prop.gcnArchName.split(":")[0] + if gfx_arch in NAVI3_ARCH: + return True + return False + +def freeze_rng_state(*args, **kwargs): + return torch.testing._utils.freeze_rng_state(*args, **kwargs) + + +# Class to keep track of test flags configurable by environment variables. +# Flags set here are intended to be read-only and should not be modified after +# definition. +# TODO: Expand this class to handle arbitrary settings in addition to boolean flags? +class TestEnvironment: + # Set of env vars to set for the repro command that is output on test failure. + # Specifically, this includes env vars that are set to non-default values and + # are not implied. Maps from env var name -> value (int) + repro_env_vars: dict = {} + + # Defines a flag usable throughout the test suite, determining its value by querying + # the specified environment variable. + # + # Args: + # name (str): The name of the flag. A global variable with this name will be set + # for convenient access throughout the test suite. + # env_var (str): The name of the primary environment variable from which to + # determine the value of this flag. If this is None or the environment variable + # is unset, the default value will be used unless otherwise implied (see + # implied_by_fn). Default: None + # default (bool): The default value to use for the flag if unset by the environment + # variable and unimplied. Default: False + # include_in_repro (bool): Indicates whether this flag should be included in the + # repro command that is output on test failure (i.e. whether it is possibly + # relevant to reproducing the test failure). Default: True + # enabled_fn (Callable): Callable returning whether the flag should be enabled + # given the environment variable value and the default value. Default: Lambda + # requiring "0" to disable if on by default OR "1" to enable if off by default. + # implied_by_fn (Callable): Thunk returning a bool to imply this flag as enabled + # by something outside of its primary environment variable setting. For example, + # this can be useful if the value of another environment variable implies the flag + # as enabled. Default: Lambda returning False to indicate no implications. + @staticmethod + def def_flag( + name, + env_var=None, + default=False, + include_in_repro=True, + enabled_fn=lambda env_var_val, default: ( + (env_var_val != "0") if default else (env_var_val == "1")), + implied_by_fn=lambda: False, + ): + enabled = default + env_var_val = None + if env_var is not None: + env_var_val = os.getenv(env_var) + enabled = enabled_fn(env_var_val, default) + implied = implied_by_fn() + enabled = enabled or implied + if include_in_repro and (env_var is not None) and (enabled != default) and not implied: + TestEnvironment.repro_env_vars[env_var] = env_var_val + + # export flag globally for convenience + if name in globals(): + raise AssertionError(f"duplicate definition of flag '{name}'") + globals()[name] = enabled + return enabled + + # Defines a setting usable throughout the test suite, determining its value by querying + # the specified environment variable. This differs from a flag in that it's not restricted + # to a boolean value. + # + # Args: + # name (str): The name of the setting. A global variable with this name will be set + # for convenient access throughout the test suite. + # env_var (str): The name of the primary environment variable from which to + # determine the value of this setting. If this is None or the environment variable + # is unset, the default value will be used. Default: None + # default (Any): The default value to use for the setting if unset by the environment + # variable. Default: None + # include_in_repro (bool): Indicates whether this setting should be included in the + # repro command that is output on test failure (i.e. whether it is possibly + # relevant to reproducing the test failure). Default: True + # parse_fn (Callable): Callable parsing the env var string. Default value just uses + # the string itself. + @staticmethod + def def_setting( + name, + env_var=None, + default=None, + include_in_repro=True, + parse_fn=lambda maybe_val_str: maybe_val_str, + ): + value = default if env_var is None else os.getenv(env_var) + value = parse_fn(value) + if include_in_repro and (value != default): + TestEnvironment.repro_env_vars[env_var] = value + + # export setting globally for convenience + if name in globals(): + raise AssertionError(f"duplicate definition of setting '{name}'") + globals()[name] = value + return value + + # Returns a string prefix usable to set environment variables for any test + # settings that should be explicitly set to match this instantiation of the + # test suite. + # Example: "PYTORCH_TEST_WITH_ASAN=1 PYTORCH_TEST_WITH_ROCM=1" + @staticmethod + def repro_env_var_prefix() -> str: + return " ".join([f"{env_var}={value}" + for env_var, value in TestEnvironment.repro_env_vars.items()]) + + +log = logging.getLogger(__name__) +torch.backends.disable_global_flags() + +FILE_SCHEMA = "file://" +if sys.platform == 'win32': + FILE_SCHEMA = "file:///" + +# NB: This flag differs semantically from others in that setting the env var to any +# non-empty value will cause it to be true: +# CI=1, CI="true", CI=0, etc. all set the flag to be true. +# CI= and an unset CI set the flag to be false. +# GitHub sets the value to CI="true" to enable it. +IS_CI: bool = TestEnvironment.def_flag( + "IS_CI", + env_var="CI", + include_in_repro=False, + enabled_fn=lambda env_var_value, _: bool(env_var_value), +) +IS_SANDCASTLE: bool = TestEnvironment.def_flag( + "IS_SANDCASTLE", + env_var="SANDCASTLE", + implied_by_fn=lambda: os.getenv("TW_JOB_USER") == "sandcastle", + include_in_repro=False, +) +IN_RE_WORKER: bool = os.environ.get("INSIDE_RE_WORKER") is not None + +_is_fbcode_default = ( + hasattr(torch._utils_internal, "IS_FBSOURCE") and + torch._utils_internal.IS_FBSOURCE +) + +IS_FBCODE: bool = TestEnvironment.def_flag( + "IS_FBCODE", + env_var="PYTORCH_TEST_FBCODE", + default=_is_fbcode_default, + include_in_repro=False, +) +IS_REMOTE_GPU: bool = TestEnvironment.def_flag( + "IS_REMOTE_GPU", + env_var="PYTORCH_TEST_REMOTE_GPU", + include_in_repro=False, +) + +DISABLE_RUNNING_SCRIPT_CHK: bool = TestEnvironment.def_flag( + "DISABLE_RUNNING_SCRIPT_CHK", + env_var="PYTORCH_DISABLE_RUNNING_SCRIPT_CHK", + include_in_repro=False, +) +# NB: enabled by default unless in an fbcode context. +PRINT_REPRO_ON_FAILURE: bool = TestEnvironment.def_flag( + "PRINT_REPRO_ON_FAILURE", + env_var="PYTORCH_PRINT_REPRO_ON_FAILURE", + default=(not IS_FBCODE), + include_in_repro=False, +) + +# possibly restrict OpInfo tests to a single sample input +OPINFO_SAMPLE_INPUT_INDEX: int | None = TestEnvironment.def_setting( + "OPINFO_SAMPLE_INPUT_INDEX", + env_var="PYTORCH_OPINFO_SAMPLE_INPUT_INDEX", + default=None, + # Don't include the env var value in the repro command because the info will + # be queried from the tracked sample input instead + include_in_repro=False, + parse_fn=lambda val: None if val is None else int(val), +) + +DEFAULT_DISABLED_TESTS_FILE = '.pytorch-disabled-tests.json' +DEFAULT_SLOW_TESTS_FILE = 'slow_tests.json' + +disabled_tests_dict = {} +slow_tests_dict = {} + +def maybe_load_json(filename): + if os.path.isfile(filename): + with open(filename) as fp: + return json.load(fp) + log.warning("Attempted to load json file '%s' but it does not exist.", filename) + return {} + +# set them here in case the tests are running in a subprocess that doesn't call run_tests +if os.getenv("SLOW_TESTS_FILE", ""): + slow_tests_dict = maybe_load_json(os.getenv("SLOW_TESTS_FILE", "")) +if os.getenv("DISABLED_TESTS_FILE", ""): + disabled_tests_dict = maybe_load_json(os.getenv("DISABLED_TESTS_FILE", "")) + +NATIVE_DEVICES = ('cpu', 'cuda', 'xpu', 'meta', 'mps', 'mtia', torch._C._get_privateuse1_backend_name()) + +# used for managing devices testing for torch profiler UTs +# for now cpu, cuda and xpu are added for testing torch profiler UTs +DEVICE_LIST_SUPPORT_PROFILING_TEST = ('cpu', 'cuda', 'xpu') +ALLOW_XPU_PROFILING_TEST = True + +check_names = ['orin', 'concord', 'galen', 'xavier', 'nano', 'jetson', 'tegra', 'thor'] +IS_JETSON = any(name in platform.platform() for name in check_names) + +def gcIfJetson(fn): + # Irregular Jetson host/device memory setup requires cleanup to avoid tests being killed + @functools.wraps(fn) + def wrapper(*args, **kwargs): + if IS_JETSON: + gc.collect() + torch.cuda.empty_cache() + fn(*args, **kwargs) + return wrapper + +# Tries to extract the current test function by crawling the stack. +# If unsuccessful, return None. +def extract_test_fn() -> Callable | None: + try: + stack = inspect.stack() + for frame_info in stack: + frame = frame_info.frame + if "self" not in frame.f_locals: + continue + self_val = frame.f_locals["self"] + if isinstance(self_val, unittest.TestCase): + test_id = self_val.id() + *_, cls_name, test_name = test_id.rsplit('.', 2) + if cls_name == type(self_val).__name__ and test_name.startswith("test"): + test_fn = getattr(self_val, test_name).__func__ + return test_fn + except Exception: + pass + return None + +# Contains tracked input data useful for debugging purposes +@dataclass +class TrackedInput: + index: int + val: Any + type_desc: str + +# Attempt to pull out tracked input information from the test function. +# A TrackedInputIter is used to insert this information. +def get_tracked_input() -> TrackedInput | None: + test_fn = extract_test_fn() + if test_fn is None: + return None + return getattr(test_fn, "tracked_input", None) + +def clear_tracked_input() -> None: + test_fn = extract_test_fn() + if test_fn is None: + return + if not hasattr(test_fn, "tracked_input"): + return + test_fn.tracked_input = None # type: ignore[attr-defined] + +# Wraps an iterator and tracks the most recent value the iterator produces +# for debugging purposes. Tracked values are stored on the test function. +class TrackedInputIter: + def __init__( + self, + child_iter, + input_type_desc, + item_callback=None, + track_callback=None, + set_seed=True, + restrict_to_index=None + ): + self.child_iter = enumerate(child_iter) + # Input type describes the things we're tracking (e.g. "sample input", "error input"). + self.input_type_desc = input_type_desc + # NB: The two types of callbacks below exist because the thing we want to track isn't + # always the same as the thing we want returned from the iterator. An example of this + # is ErrorInput, which we want returned from the iterator, but which contains a + # SampleInput that we want to track. + # Item callback is run on each (iterated thing, index) to get the thing to return. + self.item_callback = item_callback + if self.item_callback is None: + self.item_callback = lambda x, i: x + # Track callback is run on each iterated thing to get the thing to track. + self.track_callback = track_callback + if self.track_callback is None: + self.track_callback = lambda x: x + self.test_fn = extract_test_fn() + # Indicates whether the random seed should be set before each call to the iterator + self.set_seed = set_seed + # Indicates that iteration should be restricted to only the provided index. + # If None, no restriction is done + self.restrict_to_index = restrict_to_index + + def __iter__(self): + return self + + def __next__(self): + while True: + if self.set_seed: + # use a test-name-specific hash for the seed if possible + seed = ( + int.from_bytes(hashlib.sha256( + self.test_fn.__qualname__.encode("utf-8")).digest()[:4], 'little') + if self.test_fn is not None else SEED + ) + set_rng_seed(seed) + + # allow StopIteration to bubble up + input_idx, input_val = next(self.child_iter) + if (self.restrict_to_index is None) or (input_idx == self.restrict_to_index): + break + + self._set_tracked_input( + TrackedInput( + index=input_idx, val=self.track_callback(input_val), type_desc=self.input_type_desc + ) + ) + return self.item_callback(input_val, input_idx) + + def _set_tracked_input(self, tracked_input: TrackedInput): + if self.test_fn is None: + return + if not hasattr(self.test_fn, "tracked_input"): + return + self.test_fn.tracked_input = tracked_input # type: ignore[attr-defined] + +class _TestParametrizer: + """ + Decorator class for parametrizing a test function, yielding a set of new tests spawned + from the original generic test, each specialized for a specific set of test inputs. For + example, parametrizing a test across the set of ops will result in a test function per op. + + The decision of how to parametrize / what to parametrize over is intended to be implemented + by each derived class. + + In the details, the decorator adds a 'parametrize_fn' property to the test function. This function + is intended to be called later by one of: + * Device-specific test instantiation via instantiate_device_type_tests(). Note that for this + case there is no need to explicitly parametrize over device type, as that is handled separately. + * Device-agnostic parametrized test instantiation via instantiate_parametrized_tests(). + + If the decorator is applied to a test function that already has a 'parametrize_fn' property, a new + composite 'parametrize_fn' will be created that generates tests with the product of the parameters + generated by the old and new parametrize_fns. This allows for convenient composability of decorators. + """ + def _parametrize_test(self, test, generic_cls, device_cls): + """ + Parametrizes the given test function across whatever dimension is specified by the derived class. + Tests can be parametrized over any arbitrary dimension or combination of dimensions, such as all + ops, all modules, or all ops + their associated dtypes. + + Args: + test (fn): Test function to parametrize over + generic_cls (class): Generic test class object containing tests (e.g. TestFoo) + device_cls (class): Device-specialized test class object (e.g. TestFooCPU); set to None + if the tests are not part of a device-specific set + + Returns: + Generator object returning 4-tuples of: + test (fn): Parametrized test function; must support a device arg and args for any params + test_name (str): Parametrized suffix for the test (e.g. opname_int64); will be appended to + the base name of the test + param_kwargs (dict): Param kwargs to pass to the test (e.g. {'op': 'add', 'dtype': torch.int64}) + decorator_fn (callable): Callable[[Dict], List] for list of decorators to apply given param_kwargs + """ + raise NotImplementedError + + def __call__(self, fn): + if hasattr(fn, 'parametrize_fn'): + # Do composition with the product of args. + old_parametrize_fn = fn.parametrize_fn + new_parametrize_fn = self._parametrize_test + fn.parametrize_fn = compose_parametrize_fns(old_parametrize_fn, new_parametrize_fn) + else: + fn.parametrize_fn = self._parametrize_test + return fn + + +def compose_parametrize_fns(old_parametrize_fn, new_parametrize_fn): + """ + Returns a parametrize_fn that parametrizes over the product of the parameters handled + by the given parametrize_fns. Each given parametrize_fn should each have the signature + f(test, generic_cls, device_cls). + + The test names will be a combination of the names produced by the parametrize_fns in + "_" order. This order is done to match intuition for constructed names + when composing multiple decorators; the names will be built in top to bottom order when stacking + parametrization decorators. + + Args: + old_parametrize_fn (callable) - First parametrize_fn to compose. + new_parametrize_fn (callable) - Second parametrize_fn to compose. + """ + + def composite_fn(test, generic_cls, device_cls, + old_parametrize_fn=old_parametrize_fn, + new_parametrize_fn=new_parametrize_fn): + old_tests = list(old_parametrize_fn(test, generic_cls, device_cls)) + for (old_test, old_test_name, old_param_kwargs, old_dec_fn) in old_tests: + for (new_test, new_test_name, new_param_kwargs, new_dec_fn) in \ + new_parametrize_fn(old_test, generic_cls, device_cls): + redundant_params = set(old_param_kwargs.keys()).intersection(new_param_kwargs.keys()) + if redundant_params: + raise RuntimeError('Parametrization over the same parameter by multiple parametrization ' + f'decorators is not supported. For test "{test.__name__}", the following parameters ' + f'are handled multiple times: {redundant_params}') + full_param_kwargs = {**old_param_kwargs, **new_param_kwargs} + merged_test_name = '{}{}{}'.format(new_test_name, + '_' if old_test_name != '' and new_test_name != '' else '', + old_test_name) + + def merged_decorator_fn(param_kwargs, old_dec_fn=old_dec_fn, new_dec_fn=new_dec_fn): + return list(old_dec_fn(param_kwargs)) + list(new_dec_fn(param_kwargs)) + + yield (new_test, merged_test_name, full_param_kwargs, merged_decorator_fn) + + return composite_fn + + +def instantiate_parametrized_tests(generic_cls): + """ + Instantiates tests that have been decorated with a parametrize_fn. This is generally performed by a + decorator subclass of _TestParametrizer. The generic test will be replaced on the test class by + parametrized tests with specialized names. This should be used instead of + instantiate_device_type_tests() if the test class contains device-agnostic tests. + + You can also use it as a class decorator. E.g. + + ``` + @instantiate_parametrized_tests + class TestFoo(TestCase): + ... + ``` + + Args: + generic_cls (class): Generic test class object containing tests (e.g. TestFoo) + """ + for attr_name in tuple(dir(generic_cls)): + class_attr = getattr(generic_cls, attr_name) + if not hasattr(class_attr, 'parametrize_fn'): + continue + + # Remove the generic test from the test class. + delattr(generic_cls, attr_name) + + # Add parametrized tests to the test class. + def instantiate_test_helper(cls, name, test, param_kwargs): + @wraps(test) + def instantiated_test(self, param_kwargs=param_kwargs): + test(self, **param_kwargs) + + if hasattr(generic_cls, name): + raise AssertionError(f"Redefinition of test {name}") + setattr(generic_cls, name, instantiated_test) + + for (test, test_suffix, param_kwargs, decorator_fn) in class_attr.parametrize_fn( + class_attr, generic_cls=generic_cls, device_cls=None): + full_name = f'{test.__name__}_{test_suffix}' + + # Apply decorators based on full param kwargs. + for decorator in decorator_fn(param_kwargs): + test = decorator(test) + + instantiate_test_helper(cls=generic_cls, name=full_name, test=test, param_kwargs=param_kwargs) + return generic_cls + + +class subtest: + """ + Explicit subtest case for use with test parametrization. + Allows for explicit naming of individual subtest cases as well as applying + decorators to the parametrized test. + + Args: + arg_values (iterable): Iterable of arg values (e.g. range(10)) or + tuples of arg values (e.g. [(1, 2), (3, 4)]). + name (str): Optional name to use for the test. + decorators (iterable): Iterable of decorators to apply to the generated test. + """ + __slots__ = ['arg_values', 'name', 'decorators'] + + def __init__(self, arg_values, name=None, decorators=None): + self.arg_values = arg_values + self.name = name + self.decorators = decorators if decorators else [] + + +class parametrize(_TestParametrizer): + """ + Decorator for applying generic test parametrizations. + + The interface for this decorator is modeled after `@pytest.mark.parametrize`. + Basic usage between this decorator and pytest's is identical. The first argument + should be a string containing comma-separated names of parameters for the test, and + the second argument should be an iterable returning values or tuples of values for + the case of multiple parameters. + + Beyond this basic usage, the decorator provides some additional functionality that + pytest does not. + + 1. Parametrized tests end up as generated test functions on unittest test classes. + Since this differs from how pytest works, this decorator takes on the additional + responsibility of naming these test functions. The default test names consists of + the test's base name followed by each parameter name + value (e.g. "test_bar_x_1_y_foo"), + but custom names can be defined using `name_fn` or the `subtest` structure (see below). + + 2. The decorator specially handles parameter values of type `subtest`, which allows for + more fine-grained control over both test naming and test execution. In particular, it can + be used to tag subtests with explicit test names or apply arbitrary decorators (see examples + below). + + Examples:: + + @parametrize("x", range(5)) + def test_foo(self, x): + ... + + @parametrize("x,y", [(1, 'foo'), (2, 'bar'), (3, 'baz')]) + def test_bar(self, x, y): + ... + + @parametrize("x,y", [(1, 'foo'), (2, 'bar'), (3, 'baz')], + name_fn=lambda x, y: '{}_{}'.format(x, y)) + def test_bar_custom_names(self, x, y): + ... + + @parametrize("x, y", [subtest((1, 2), name='double'), + subtest((1, 3), name='triple', decorators=[unittest.expectedFailure]), + subtest((1, 4), name='quadruple')]) + def test_baz(self, x, y): + ... + + To actually instantiate the parametrized tests, one of instantiate_parametrized_tests() or + instantiate_device_type_tests() should be called. The former is intended for test classes + that contain device-agnostic tests, while the latter should be used for test classes that + contain device-specific tests. Both support arbitrary parametrizations using the decorator. + + Args: + arg_str (str): String of arg names separate by commas (e.g. "x,y"). + arg_values (iterable): Iterable of arg values (e.g. range(10)) or + tuples of arg values (e.g. [(1, 2), (3, 4)]). + name_fn (Callable): Optional function that takes in parameters and returns subtest name. + """ + def __init__(self, arg_str, arg_values, name_fn=None): + self.arg_names: list[str] = [s.strip() for s in arg_str.split(',') if s != ''] + self.arg_values = arg_values + self.name_fn = name_fn + + def _formatted_str_repr(self, idx, name, value): + """ Returns a string representation for the given arg that is suitable for use in test function names. """ + if isinstance(value, torch.dtype): + return dtype_name(value) + elif isinstance(value, torch.device): + return str(value) + # Can't use isinstance as it would cause a circular import + elif type(value).__name__ in {'OpInfo', 'ModuleInfo'}: + return value.formatted_name + elif isinstance(value, (int, float, str)): + return f"{name}_{str(value).replace('.', '_')}" + else: + return f"{name}{idx}" + + def _default_subtest_name(self, idx, values): + return '_'.join([self._formatted_str_repr(idx, a, v) for a, v in zip(self.arg_names, values, strict=True)]) + + def _get_subtest_name(self, idx, values, explicit_name=None): + if explicit_name: + subtest_name = explicit_name + elif self.name_fn: + subtest_name = self.name_fn(*values) + else: + subtest_name = self._default_subtest_name(idx, values) + return subtest_name + + def _parametrize_test(self, test, generic_cls, device_cls): + if len(self.arg_names) == 0: + # No additional parameters needed for the test. + test_name = '' + yield (test, test_name, {}, lambda _: []) + else: + # Each "values" item is expected to be either: + # * A tuple of values with one for each arg. For a single arg, a single item is expected. + # * A subtest instance with arg_values matching the previous. + values = check_exhausted_iterator = object() + for idx, values in enumerate(self.arg_values): + maybe_name = None + + decorators: list[Any] = [] + if isinstance(values, subtest): + sub = values + values = sub.arg_values + maybe_name = sub.name + + @wraps(test) + def test_wrapper(*args, **kwargs): + return test(*args, **kwargs) + + decorators = sub.decorators + gen_test = test_wrapper + else: + gen_test = test + + values = list(values) if len(self.arg_names) > 1 else [values] # type: ignore[call-overload] + if len(values) != len(self.arg_names): + raise RuntimeError(f'Expected # values == # arg names, but got: {len(values)} ' + f'values and {len(self.arg_names)} names for test "{test.__name__}"') + + param_kwargs = dict(zip(self.arg_names, values, strict=True)) + + test_name = self._get_subtest_name(idx, values, explicit_name=maybe_name) + + def decorator_fn(_, decorators=decorators): + return decorators + + yield (gen_test, test_name, param_kwargs, decorator_fn) + + if values is check_exhausted_iterator: + raise ValueError(f'{test}: An empty arg_values was passed to @parametrize. ' + 'Note that this may result from reuse of a generator.') + + +class reparametrize(_TestParametrizer): + """ + Decorator for adjusting the way an existing parametrizer operates. This class runs + the given adapter_fn on each parametrization produced by the given parametrizer, + allowing for on-the-fly parametrization more flexible than the default, + product-based composition that occurs when stacking parametrization decorators. + + If the adapter_fn returns None for a given test parametrization, that parametrization + will be excluded. Otherwise, it's expected that the adapter_fn returns an iterable of + modified parametrizations, with tweaked test names and parameter kwargs. + + Examples:: + + def include_is_even_arg(test_name, param_kwargs): + x = param_kwargs["x"] + is_even = x % 2 == 0 + new_param_kwargs = dict(param_kwargs) + new_param_kwargs["is_even"] = is_even + is_even_suffix = "_even" if is_even else "_odd" + new_test_name = f"{test_name}{is_even_suffix}" + yield (new_test_name, new_param_kwargs) + + ... + + @reparametrize(parametrize("x", range(5)), include_is_even_arg) + def test_foo(self, x, is_even): + ... + + def exclude_odds(test_name, param_kwargs): + x = param_kwargs["x"] + is_even = x % 2 == 0 + yield None if not is_even else (test_name, param_kwargs) + + ... + + @reparametrize(parametrize("x", range(5)), exclude_odds) + def test_bar(self, x): + ... + + """ + def __init__(self, parametrizer, adapter_fn): + self.parametrizer = parametrizer + self.adapter_fn = adapter_fn + + def _parametrize_test(self, test, generic_cls, device_cls): + for (gen_test, test_name, param_kwargs, decorator_fn) in \ + self.parametrizer._parametrize_test(test, generic_cls, device_cls): + adapted = self.adapter_fn(test_name, param_kwargs) + if adapted is not None: + for adapted_item in adapted: + if adapted_item is not None: + new_test_name, new_param_kwargs = adapted_item + yield (gen_test, new_test_name, new_param_kwargs, decorator_fn) + + +class decorateIf(_TestParametrizer): + """ + Decorator for applying parameter-specific conditional decoration. + Composes with other test parametrizers (e.g. @modules, @ops, @parametrize, etc.). + + Examples:: + + @decorateIf(unittest.skip, lambda params: params["x"] == 2) + @parametrize("x", range(5)) + def test_foo(self, x): + ... + + @parametrize("x,y", [(1, 'foo'), (2, 'bar'), (3, 'baz')]) + @decorateIf( + unittest.expectedFailure, + lambda params: params["x"] == 3 and params["y"] == "baz" + ) + def test_bar(self, x, y): + ... + + @decorateIf( + unittest.expectedFailure, + lambda params: params["op"].name == "add" and params["dtype"] == torch.float16 + ) + @ops(op_db) + def test_op_foo(self, device, dtype, op): + ... + + @decorateIf( + unittest.skip, + lambda params: params["module_info"].module_cls is torch.nn.Linear and \ + params["device"] == "cpu" + ) + @modules(module_db) + def test_module_foo(self, device, dtype, module_info): + ... + + Args: + decorator: Test decorator to apply if the predicate is satisfied. + predicate_fn (Callable): Function taking in a dict of params and returning a boolean + indicating whether the decorator should be applied or not. + """ + def __init__(self, decorator, predicate_fn): + self.decorator = decorator + self.predicate_fn = predicate_fn + + def _parametrize_test(self, test, generic_cls, device_cls): + + # Leave test as-is and return the appropriate decorator_fn. + def decorator_fn(params, decorator=self.decorator, predicate_fn=self.predicate_fn): + if predicate_fn(params): + return [decorator] + else: + return [] + + @wraps(test) + def test_wrapper(*args, **kwargs): + return test(*args, **kwargs) + + test_name = '' + yield (test_wrapper, test_name, {}, decorator_fn) + + +def cppProfilingFlagsToProfilingMode(): + old_prof_exec_state = torch._C._jit_set_profiling_executor(True) + old_prof_mode_state = torch._C._get_graph_executor_optimize(True) + torch._C._jit_set_profiling_executor(old_prof_exec_state) + torch._C._get_graph_executor_optimize(old_prof_mode_state) + + if old_prof_exec_state: + if old_prof_mode_state: + return ProfilingMode.PROFILING + else: + return ProfilingMode.SIMPLE + else: + return ProfilingMode.LEGACY + +@contextmanager +def enable_profiling_mode_for_profiling_tests(): + old_prof_exec_state = False + old_prof_mode_state = False + if not GRAPH_EXECUTOR: + raise AssertionError("GRAPH_EXECUTOR must be set") + if GRAPH_EXECUTOR == ProfilingMode.PROFILING: + old_prof_exec_state = torch._C._jit_set_profiling_executor(True) + old_prof_mode_state = torch._C._get_graph_executor_optimize(True) + try: + yield + finally: + if GRAPH_EXECUTOR == ProfilingMode.PROFILING: + torch._C._jit_set_profiling_executor(old_prof_exec_state) + torch._C._get_graph_executor_optimize(old_prof_mode_state) + +@contextmanager +def enable_profiling_mode(): + old_prof_exec_state = torch._C._jit_set_profiling_executor(True) + old_prof_mode_state = torch._C._get_graph_executor_optimize(True) + try: + yield + finally: + torch._C._jit_set_profiling_executor(old_prof_exec_state) + torch._C._get_graph_executor_optimize(old_prof_mode_state) + +@contextmanager +def num_profiled_runs(num_runs): + old_num_runs = torch._C._jit_set_num_profiled_runs(num_runs) + try: + yield + finally: + torch._C._jit_set_num_profiled_runs(old_num_runs) + +func_call = torch._C.ScriptFunction.__call__ +meth_call = torch._C.ScriptMethod.__call__ + +def prof_callable(callable, *args, **kwargs): + if 'profile_and_replay' in kwargs: + del kwargs['profile_and_replay'] + if not GRAPH_EXECUTOR: + raise AssertionError("GRAPH_EXECUTOR must be set") + if GRAPH_EXECUTOR == ProfilingMode.PROFILING: + with enable_profiling_mode_for_profiling_tests(): + callable(*args, **kwargs) + return callable(*args, **kwargs) + + return callable(*args, **kwargs) + +def raise_on_run_directly(file_to_call): + raise RuntimeError("This test file is not meant to be run directly, " + f"use:\n\n\tpython {file_to_call} TESTNAME\n\n" + "instead.") + +def prof_func_call(*args, **kwargs): + return prof_callable(func_call, *args, **kwargs) + +def prof_meth_call(*args, **kwargs): + return prof_callable(meth_call, *args, **kwargs) + +torch._C.ScriptFunction.__call__ = prof_func_call # type: ignore[method-assign] +torch._C.ScriptMethod.__call__ = prof_meth_call # type: ignore[method-assign] + +def _get_test_report_path(): + # allow users to override the test file location. We need this + # because the distributed tests run the same test file multiple + # times with different configurations. + override = os.environ.get('TEST_REPORT_SOURCE_OVERRIDE') + test_source = override if override is not None else 'python-unittest' + return os.path.join('test-reports', test_source) + +def parse_cmd_line_args(): + global DISABLED_TESTS_FILE + global GRAPH_EXECUTOR + global LOG_SUFFIX + global PYTEST_SINGLE_TEST + global REPEAT_COUNT + global RERUN_DISABLED_TESTS + global RUN_PARALLEL + global SHOWLOCALS + global SLOW_TESTS_FILE + global TEST_BAILOUTS + global TEST_DISCOVER + global TEST_IN_SUBPROCESS + global TEST_SAVE_XML + global UNITTEST_ARGS + global USE_PYTEST + + is_running_via_run_test = "run_test.py" in getattr(__main__, "__file__", "") + parser = argparse.ArgumentParser(add_help=not is_running_via_run_test, allow_abbrev=False) + parser.add_argument('--subprocess', action='store_true', + help='whether to run each test in a subprocess') + parser.add_argument('--accept', action='store_true') + parser.add_argument('--jit-executor', '--jit_executor', type=str) + parser.add_argument('--repeat', type=int, default=1) + parser.add_argument('--test-bailouts', '--test_bailouts', action='store_true') + parser.add_argument('--use-pytest', action='store_true') + parser.add_argument('--save-xml', nargs='?', type=str, + const=_get_test_report_path(), + default=_get_test_report_path() if IS_CI else None) + parser.add_argument('--discover-tests', action='store_true') + parser.add_argument('--log-suffix', type=str, default="") + parser.add_argument('--run-parallel', type=int, default=1) + parser.add_argument('--import-slow-tests', type=str, nargs='?', const=DEFAULT_SLOW_TESTS_FILE) + parser.add_argument('--import-disabled-tests', type=str, nargs='?', const=DEFAULT_DISABLED_TESTS_FILE) + parser.add_argument('--rerun-disabled-tests', action='store_true') + parser.add_argument('--pytest-single-test', type=str, nargs=1) + parser.add_argument('--showlocals', action=argparse.BooleanOptionalAction, default=False) + +# Only run when -h or --help flag is active to display both unittest and parser help messages. + def run_unittest_help(argv): + unittest.main(argv=argv) + + if '-h' in sys.argv or '--help' in sys.argv: + help_thread = threading.Thread(target=run_unittest_help, args=(sys.argv,)) + help_thread.start() + help_thread.join() + + args, remaining = parser.parse_known_args() + if args.jit_executor == 'legacy': + GRAPH_EXECUTOR = ProfilingMode.LEGACY + elif args.jit_executor == 'profiling': + GRAPH_EXECUTOR = ProfilingMode.PROFILING + elif args.jit_executor == 'simple': + GRAPH_EXECUTOR = ProfilingMode.SIMPLE + else: + # infer flags based on the default settings + GRAPH_EXECUTOR = cppProfilingFlagsToProfilingMode() + + RERUN_DISABLED_TESTS = args.rerun_disabled_tests + + SLOW_TESTS_FILE = args.import_slow_tests + DISABLED_TESTS_FILE = args.import_disabled_tests + LOG_SUFFIX = args.log_suffix + RUN_PARALLEL = args.run_parallel + TEST_BAILOUTS = args.test_bailouts + USE_PYTEST = args.use_pytest + PYTEST_SINGLE_TEST = args.pytest_single_test + TEST_DISCOVER = args.discover_tests + TEST_IN_SUBPROCESS = args.subprocess + TEST_SAVE_XML = args.save_xml + REPEAT_COUNT = args.repeat + SHOWLOCALS = args.showlocals + if not getattr(expecttest, "ACCEPT", False): + expecttest.ACCEPT = args.accept + UNITTEST_ARGS = [sys.argv[0]] + remaining + + set_rng_seed() + + +def wait_for_process(p, timeout=None): + try: + return p.wait(timeout=timeout) + except KeyboardInterrupt: + # Give `p` a chance to handle KeyboardInterrupt. Without this, + # `pytest` can't print errors it collected so far upon KeyboardInterrupt. + exit_status = p.wait(timeout=5) + if exit_status is not None: + return exit_status + else: + p.kill() + raise + except subprocess.TimeoutExpired: + # send SIGINT to give pytest a chance to make xml + p.send_signal(signal.SIGINT) + exit_status = None + try: + exit_status = p.wait(timeout=5) + # try to handle the case where p.wait(timeout=5) times out as well as + # otherwise the wait() call in the finally block can potentially hang + except subprocess.TimeoutExpired: + pass + if exit_status is not None: + return exit_status + else: + p.kill() + raise + except: # noqa: B001,E722, copied from python core library + p.kill() + raise + finally: + # Always call p.wait() to ensure exit + p.wait() + +def shell(command, cwd=None, env=None, stdout=None, stderr=None, timeout=None): + sys.stdout.flush() + sys.stderr.flush() + # The following cool snippet is copied from Py3 core library subprocess.call + # only the with + # 1. `except KeyboardInterrupt` block added for SIGINT handling. + # 2. In Py2, subprocess.Popen doesn't return a context manager, so we do + # `p.wait()` in a `final` block for the code to be portable. + # + # https://github.com/python/cpython/blob/71b6c1af727fbe13525fb734568057d78cea33f3/Lib/subprocess.py#L309-L323 + if isinstance(command, str): + raise AssertionError("Command to shell should be a list or tuple of tokens") + p = subprocess.Popen(command, universal_newlines=True, cwd=cwd, env=env, stdout=stdout, stderr=stderr) + return wait_for_process(p, timeout=timeout) + + +def retry_shell( + command, + cwd=None, + env=None, + stdout=None, + stderr=None, + timeout=None, + retries=1, + was_rerun=False, +) -> tuple[int, bool]: + # Returns exicode + whether it was rerun + if not (retries >= 0): + raise AssertionError( + f"Expecting non negative number for number of retries, got {retries}" + ) + try: + exit_code = shell( + command, cwd=cwd, env=env, stdout=stdout, stderr=stderr, timeout=timeout + ) + if exit_code == 0 or retries == 0: + return exit_code, was_rerun + print( + f"Got exit code {exit_code}, retrying (retries left={retries})", + file=stdout, + flush=True, + ) + except subprocess.TimeoutExpired: + if retries == 0: + print( + f"Command took >{timeout // 60}min, returning 124", + file=stdout, + flush=True, + ) + return 124, was_rerun + print( + f"Command took >{timeout // 60}min, retrying (retries left={retries})", + file=stdout, + flush=True, + ) + return retry_shell( + command, + cwd=cwd, + env=env, + stdout=stdout, + stderr=stderr, + timeout=timeout, + retries=retries - 1, + was_rerun=True, + ) + + +def discover_test_cases_recursively(suite_or_case): + if isinstance(suite_or_case, unittest.TestCase): + return [suite_or_case] + rc = [] + for element in suite_or_case: + print(element) + rc.extend(discover_test_cases_recursively(element)) + return rc + +def get_test_names(test_cases): + return ['.'.join(case.id().split('.')[-2:]) for case in test_cases] + +def _print_test_names(): + suite = unittest.TestLoader().loadTestsFromModule(__main__) + test_cases = discover_test_cases_recursively(suite) + for name in get_test_names(test_cases): + print(name) + +def chunk_list(lst, nchunks): + return [lst[i::nchunks] for i in range(nchunks)] + +# sanitize filename e.g., distributed/pipeline/sync/skip/test_api.py -> distributed.pipeline.sync.skip.test_api +def sanitize_test_filename(filename): + strip_py = re.sub(r'.py$', '', filename) + return re.sub('/', r'.', strip_py) + +def lint_test_case_extension(suite): + succeed = True + for test_case_or_suite in suite: + test_case = test_case_or_suite + if isinstance(test_case_or_suite, unittest.TestSuite): + first_test = test_case_or_suite._tests[0] if len(test_case_or_suite._tests) > 0 else None + if first_test is not None and isinstance(first_test, unittest.TestSuite): + return succeed and lint_test_case_extension(test_case_or_suite) + test_case = first_test + + if test_case is not None: + if not isinstance(test_case, TestCase): + test_class = test_case.id().split('.', 1)[1].split('.')[0] + err = "This test class should extend from torch.testing._internal.common_utils.TestCase but it doesn't." + print(f"{test_class} - failed. {err}") + succeed = False + return succeed + + +def get_report_path(argv=None, pytest=False): + if argv is None: + argv = UNITTEST_ARGS + test_filename = sanitize_test_filename(argv[0]) + test_report_path = TEST_SAVE_XML + LOG_SUFFIX + test_report_path = os.path.join(test_report_path, test_filename) + if pytest: + test_report_path = test_report_path.replace('python-unittest', 'python-pytest') + os.makedirs(test_report_path, exist_ok=True) + test_report_path = os.path.join(test_report_path, f"{test_filename}-{os.urandom(8).hex()}.xml") + return test_report_path + os.makedirs(test_report_path, exist_ok=True) + return test_report_path + + +def sanitize_pytest_xml(xml_file: str): + # pytext xml is different from unittext xml, this function makes pytest xml more similar to unittest xml + # consider somehow modifying the XML logger in conftest to do this instead + import xml.etree.ElementTree as ET + tree = ET.parse(xml_file) + for testcase in tree.iter('testcase'): + full_classname = testcase.attrib.get("classname") + if full_classname is None: + continue + # The test prefix is optional + regex_result = re.search(r"^(test\.)?(?P.*)\.(?P[^\.]*)$", full_classname) + if regex_result is None: + continue + classname = regex_result.group("classname") + file = regex_result.group("file").replace(".", "/") + testcase.set("classname", classname) + testcase.set("file", f"{file}.py") + tree.write(xml_file) + + +def get_pytest_test_cases(argv: list[str]) -> list[str]: + class TestCollectorPlugin: + def __init__(self) -> None: + self.tests: list[Any] = [] + + def pytest_collection_finish(self, session): + for item in session.items: + self.tests.append(session.config.cwd_relative_nodeid(item.nodeid)) + + test_collector_plugin = TestCollectorPlugin() + import pytest + pytest.main( + [arg for arg in argv if arg != '-vv'] + ['--collect-only', '-qq', '--use-main-module'], + plugins=[test_collector_plugin] + ) + return test_collector_plugin.tests + + +def run_tests(argv=None): + parse_cmd_line_args() + if argv is None: + argv = UNITTEST_ARGS + + # import test files. + if SLOW_TESTS_FILE: + if os.path.exists(SLOW_TESTS_FILE): + with open(SLOW_TESTS_FILE) as fp: + global slow_tests_dict + slow_tests_dict = json.load(fp) + # use env vars so pytest-xdist subprocesses can still access them + os.environ['SLOW_TESTS_FILE'] = SLOW_TESTS_FILE + else: + warnings.warn(f'slow test file provided but not found: {SLOW_TESTS_FILE}', stacklevel=2) + if DISABLED_TESTS_FILE: + if os.path.exists(DISABLED_TESTS_FILE): + with open(DISABLED_TESTS_FILE) as fp: + global disabled_tests_dict + disabled_tests_dict = json.load(fp) + os.environ['DISABLED_TESTS_FILE'] = DISABLED_TESTS_FILE + else: + warnings.warn(f'disabled test file provided but not found: {DISABLED_TESTS_FILE}', stacklevel=2) + # Determine the test launch mechanism + if TEST_DISCOVER: + _print_test_names() + return + + # Before running the tests, lint to check that every test class extends from TestCase + suite = unittest.TestLoader().loadTestsFromModule(__main__) + if not lint_test_case_extension(suite): + sys.exit(1) + + if SHOWLOCALS: + argv = [ + argv[0], + *(["--showlocals", "--tb=long", "--color=yes"] if USE_PYTEST else ["--locals"]), + *argv[1:], + ] + + if TEST_IN_SUBPROCESS: + other_args = [] + if DISABLED_TESTS_FILE: + other_args.append("--import-disabled-tests") + if SLOW_TESTS_FILE: + other_args.append("--import-slow-tests") + if USE_PYTEST: + other_args.append("--use-pytest") + if RERUN_DISABLED_TESTS: + other_args.append("--rerun-disabled-tests") + if TEST_SAVE_XML: + other_args += ['--save-xml', TEST_SAVE_XML] + + test_cases = ( + get_pytest_test_cases(argv) if USE_PYTEST else + [case.id().split('.', 1)[1] for case in discover_test_cases_recursively(suite)] + ) + + failed_tests = [] + + for test_case_full_name in test_cases: + + cmd = ( + [sys.executable] + [argv[0]] + other_args + argv[1:] + + (["--pytest-single-test"] if USE_PYTEST else []) + + [test_case_full_name] + ) + string_cmd = " ".join(cmd) + + timeout = None if RERUN_DISABLED_TESTS else 15 * 60 + + exitcode, _ = retry_shell(cmd, timeout=timeout, retries=0 if RERUN_DISABLED_TESTS else 1) + + if exitcode != 0: + # This is sort of hacky, but add on relevant env variables for distributed tests. + if 'TestDistBackendWithSpawn' in test_case_full_name: + backend = os.environ.get("BACKEND", "") + world_size = os.environ.get("WORLD_SIZE", "") + env_prefix = f"BACKEND={backend} WORLD_SIZE={world_size}" + string_cmd = env_prefix + " " + string_cmd + # Log the command to reproduce the failure. + print(f"Test exited with non-zero exitcode {exitcode}. Command to reproduce: {string_cmd}") + failed_tests.append(test_case_full_name) + + if len(failed_tests) != 0: + raise AssertionError( + "{} unit test(s) failed:\n\t{}".format( + len(failed_tests), '\n\t'.join(failed_tests) + ) + ) + + elif RUN_PARALLEL > 1: + test_cases = discover_test_cases_recursively(suite) + test_batches = chunk_list(get_test_names(test_cases), RUN_PARALLEL) + processes = [] + for i in range(RUN_PARALLEL): + command = [sys.executable] + argv + [f'--log-suffix=-shard-{i + 1}'] + test_batches[i] + processes.append(subprocess.Popen(command, universal_newlines=True)) + failed = False + for p in processes: + failed |= wait_for_process(p) != 0 + if failed: + raise AssertionError("Some test shards have failed") + elif USE_PYTEST: + pytest_args = argv + ["--use-main-module"] + test_report_path = "" + if TEST_SAVE_XML: + test_report_path = get_report_path(pytest=True) + print(f'Test results will be stored in {test_report_path}') + pytest_args.append(f'--junit-xml-reruns={test_report_path}') + if PYTEST_SINGLE_TEST: + pytest_args = PYTEST_SINGLE_TEST + pytest_args[1:] + + import pytest + os.environ["NO_COLOR"] = "1" + exit_code = pytest.main(args=pytest_args) + if TEST_SAVE_XML: + sanitize_pytest_xml(test_report_path) + + # exitcode of 5 means no tests were found, which happens since some test configs don't + # run tests from certain files + sys.exit(0 if exit_code == 5 else exit_code) + elif TEST_SAVE_XML: + # import here so that non-CI doesn't need xmlrunner installed + import xmlrunner # type: ignore[import] + from xmlrunner.result import _XMLTestResult # type: ignore[import] + + class XMLTestResultVerbose(_XMLTestResult): + """ + Adding verbosity to test outputs: + by default test summary prints 'skip', + but we want to also print the skip reason. + GH issue: https://github.com/pytorch/pytorch/issues/69014 + + This works with unittest_xml_reporting<=3.2.0,>=2.0.0 + (3.2.0 is latest at the moment) + """ + + def addSkip(self, test, reason): + super().addSkip(test, reason) + for c in self.callback.__closure__: + if isinstance(c.cell_contents, str) and c.cell_contents == 'skip': + # this message is printed in test summary; + # it stands for `verbose_str` captured in the closure + c.cell_contents = f"skip: {reason}" + + def printErrors(self) -> None: + super().printErrors() + self.printErrorList("XPASS", self.unexpectedSuccesses) + test_report_path = get_report_path() + verbose = '--verbose' in argv or '-v' in argv + if verbose: + print(f'Test results will be stored in {test_report_path}') + unittest.main(argv=argv, testRunner=xmlrunner.XMLTestRunner( + output=test_report_path, + verbosity=2 if verbose else 1, + resultclass=XMLTestResultVerbose)) + elif REPEAT_COUNT > 1: + for _ in range(REPEAT_COUNT): + if not unittest.main(exit=False, argv=argv).result.wasSuccessful(): + sys.exit(-1) + else: + unittest.main(argv=argv) + +IS_LINUX = sys.platform == "linux" +IS_WINDOWS = sys.platform == "win32" +IS_MACOS = sys.platform == "darwin" +IS_PPC = platform.machine() == "ppc64le" +IS_X86 = platform.machine() in ('x86_64', 'i386') +IS_ARM64 = platform.machine() in ('arm64', 'aarch64') +IS_S390X = platform.machine() == "s390x" +IS_AVX512_VNNI_SUPPORTED = torch.cpu.get_capabilities().get("avx512_vnni", False) +IS_CPU_EXT_SVE_SUPPORTED = torch.cpu.get_capabilities().get("sve", False) +IS_CPU_CAPABILITY_SVE256 = torch._C._get_cpu_capability() == "SVE256" + +if IS_WINDOWS: + @contextmanager + def TemporaryFileName(*args, **kwargs): + # Ideally we would like to not have to manually delete the file, but NamedTemporaryFile + # opens the file, and it cannot be opened multiple times in Windows. To support Windows, + # close the file after creation and try to remove it manually + if 'delete' in kwargs: + if kwargs['delete'] is not False: + raise UserWarning("only TemporaryFileName with delete=False is supported on Windows.") + else: + kwargs['delete'] = False + f = tempfile.NamedTemporaryFile(*args, **kwargs) # noqa:SIM115 + try: + f.close() + yield f.name + finally: + os.unlink(f.name) +else: + @contextmanager # noqa: T484 + def TemporaryFileName(*args, **kwargs): + with tempfile.NamedTemporaryFile(*args, **kwargs) as f: + yield f.name + +if IS_WINDOWS: + @contextmanager + def TemporaryDirectoryName(suffix=None): + # On Windows the directory created by TemporaryDirectory is likely to be removed prematurely, + # so we first create the directory using mkdtemp and then remove it manually + try: + dir_name = tempfile.mkdtemp(suffix=suffix) + yield dir_name + finally: + shutil.rmtree(dir_name) +else: + @contextmanager # noqa: T484 + def TemporaryDirectoryName(suffix=None): + with tempfile.TemporaryDirectory(suffix=suffix) as d: + yield d + + +def make_lazy_class(cls): + + def lazy_init(self, cb): + self._cb = cb + self._value = None + + cls.__init__ = lazy_init + + for basename in [ + "add", "sub", "mul", "truediv", "floordiv", "mod", "divmod", "pow", + "lshift", "rshift", "and", "or", "xor", "neg", "pos", "abs", "invert", + "eq", "ne", "lt", "le", "gt", "ge", "bool", "int", "index", + ]: + name = f"__{basename}__" + + def inner_wrapper(name): + use_operator = basename not in ("bool", "int") + + def wrapped(self, *args, **kwargs): + if self._cb is not None: + self._value = self._cb() + self._cb = None + if not use_operator: + return getattr(self._value, name)(*args, **kwargs) + else: + return getattr(operator, name)(self._value, *args, **kwargs) + return wrapped + + setattr(cls, name, inner_wrapper(name)) + + return cls + + +@make_lazy_class +class LazyVal: + pass + + +IS_FILESYSTEM_UTF8_ENCODING = sys.getfilesystemencoding() == 'utf-8' + +TEST_NUMPY = _check_module_exists('numpy') +TEST_FAIRSEQ = _check_module_exists('fairseq') +TEST_SCIPY = _check_module_exists('scipy') +TEST_MKL = torch.backends.mkl.is_available() +TEST_ACL = torch.backends.mkldnn.is_available() and torch.ops.mkldnn._is_mkldnn_acl_supported() +TEST_MPS = torch.backends.mps.is_available() +MACOS_VERSION = float('.'.join(platform.mac_ver()[0].split('.')[:2]) or -1) +TEST_XPU = torch.xpu.is_available() +TEST_HPU = bool(hasattr(torch, "hpu") and torch.hpu.is_available()) +TEST_CUDA = torch.cuda.is_available() +TEST_ACCELERATOR = LazyVal(lambda: torch.accelerator.is_available()) # type: ignore[call-arg] +TEST_MULTIACCELERATOR = LazyVal(lambda: torch.accelerator.device_count() > 1) # type: ignore[call-arg] +custom_device_mod = getattr(torch, torch._C._get_privateuse1_backend_name(), None) +TEST_PRIVATEUSE1 = _is_privateuse1_backend_available() +TEST_PRIVATEUSE1_DEVICE_TYPE = torch._C._get_privateuse1_backend_name() +TEST_NUMBA = _check_module_exists('numba') +TEST_TRANSFORMERS = _check_module_exists('transformers') +TEST_DILL = _check_module_exists('dill') + +TEST_LIBROSA = _check_module_exists('librosa') and not IS_ARM64 + +TEST_OPT_EINSUM = _check_module_exists('opt_einsum') + +TEST_Z3 = _check_module_exists('z3') + +# DSL availability (lazy evaluation to avoid import overhead) +class LazyDSLCheck: + """Lazy DSL availability checker to avoid import-time overhead""" + def __init__(self): + self._registry = None + self._import_attempted = False + + def _get_registry(self): + if not self._import_attempted: + self._import_attempted = True + try: + from torch._native.dsl_registry import dsl_registry + self._registry = dsl_registry + except ImportError: + self._registry = None + return self._registry + + def is_available(self, dsl_name: str) -> bool: + """Check if specific DSL is available""" + registry = self._get_registry() + return registry.is_dsl_available(dsl_name) if registry is not None else False + + def list_available(self) -> list[str]: + """Get list of available DSLs""" + registry = self._get_registry() + return list(registry.list_available_dsls()) if registry is not None else [] + + def list_all(self) -> list[str]: + """Get list of all registered DSLs""" + registry = self._get_registry() + return list(registry.list_all_dsls()) if registry is not None else [] + +_dsl_checker = LazyDSLCheck() + +# Lazy constants to avoid import-time overhead +TEST_TRITON_DSL = LazyVal(lambda: _dsl_checker.is_available('triton')) +TEST_CUTEDSL = LazyVal(lambda: _dsl_checker.is_available('cutedsl')) + +def split_if_not_empty(x: str): + return x.split(",") if len(x) != 0 else [] + +NOTEST_CPU = "cpu" in split_if_not_empty(os.getenv('PYTORCH_TESTING_DEVICE_EXCEPT_FOR', '')) + +skipIfNoDill = unittest.skipIf(not TEST_DILL, "no dill") + +# DSL skip decorators (following existing pattern) +skipIfNoTritonDSL = unittest.skipIf(not TEST_TRITON_DSL, "Triton DSL not available") +skipIfNoCuteDSL = unittest.skipIf(not TEST_CUTEDSL, "CuTeDSL not available") + +def skipIfDSLUnavailable(dsl_name: str, reason: str | None = None): + """Skip test if specific DSL is not available""" + available = _dsl_checker.is_available(dsl_name) + msg = reason or f"{dsl_name} DSL not available" + return unittest.skipIf(not available, msg) + +def skipUnlessDSLAvailable(dsl_name: str, reason: str | None = None): + """Skip test unless specific DSL is available""" + available = _dsl_checker.is_available(dsl_name) + msg = reason or f"{dsl_name} DSL required" + return unittest.skipUnless(available, msg) + +def get_available_dsls() -> list[str]: + """Get list of available DSL names for test parameterization""" + return _dsl_checker.list_available() + +def is_dsl_available(dsl_name: str) -> bool: + """Check if specific DSL is available for conditional testing""" + return _dsl_checker.is_available(dsl_name) + +def get_all_dsls() -> list[str]: + """Get all registered DSL names (available or not) for comprehensive testing""" + return _dsl_checker.list_all() + + +NO_MULTIPROCESSING_SPAWN: bool = False +TEST_WITH_ASAN: bool = TestEnvironment.def_flag( + "TEST_WITH_ASAN", + env_var="PYTORCH_TEST_WITH_ASAN", +) +TEST_WITH_DEV_DBG_ASAN: bool = TestEnvironment.def_flag( + "TEST_WITH_DEV_DBG_ASAN", + env_var="PYTORCH_TEST_WITH_DEV_DBG_ASAN", +) +TEST_WITH_TSAN: bool = TestEnvironment.def_flag( + "TEST_WITH_TSAN", + env_var="PYTORCH_TEST_WITH_TSAN", +) +TEST_WITH_UBSAN: bool = TestEnvironment.def_flag( + "TEST_WITH_UBSAN", + env_var="PYTORCH_TEST_WITH_UBSAN", +) +TEST_WITH_ROCM: bool = TestEnvironment.def_flag( + "TEST_WITH_ROCM", + env_var="PYTORCH_TEST_WITH_ROCM", + implied_by_fn=lambda: torch.version.hip is not None, +) +TEST_WITH_MTIA: bool = TestEnvironment.def_flag( + "TEST_WITH_MTIA", + env_var="PYTORCH_TEST_WITH_MTIA", +) + +# TODO: Remove PYTORCH_MIOPEN_SUGGEST_NHWC once ROCm officially supports NHWC in MIOpen +# See #64427 +TEST_WITH_MIOPEN_SUGGEST_NHWC = os.getenv('PYTORCH_MIOPEN_SUGGEST_NHWC', '0') == '1' +# Enables tests that are slow to run (disabled by default) +TEST_WITH_SLOW: bool = TestEnvironment.def_flag( + "TEST_WITH_SLOW", + env_var="PYTORCH_TEST_WITH_SLOW", +) + +# Disables non-slow tests (these tests enabled by default) +# This is usually used in conjunction with TEST_WITH_SLOW to +# run *only* slow tests. (I could have done an enum, but +# it felt a little awkward. +TEST_SKIP_FAST: bool = TestEnvironment.def_flag( + "TEST_SKIP_FAST", + env_var="PYTORCH_TEST_SKIP_FAST", +) + +# Enables crossref tests, in addition to standard tests which +# are being run. crossref tests work by installing a torch +# function mode that runs extra compute alongside the regular +# computation that happens with the test. After both computations +# are done, we cross-reference them (thus the name) to check for +# correction, before throwing out the extra compute and proceeding +# as we had before. By default, we don't run these tests. +TEST_WITH_CROSSREF: bool = TestEnvironment.def_flag( + "TEST_WITH_CROSSREF", + env_var="PYTORCH_TEST_WITH_CROSSREF", +) + +TEST_SKIP_CUDAGRAPH: bool = TestEnvironment.def_flag( + "TEST_SKIP_CUDAGRAPH", + env_var="PYTORCH_TEST_SKIP_CUDAGRAPH", +) +TEST_CUDA_GRAPH = TEST_CUDA and (not TEST_SKIP_CUDAGRAPH) and ( + torch.version.cuda or + (torch.version.hip and float(".".join(torch.version.hip.split(".")[0:2])) >= 5.3) +) + +TEST_CUDA_CUDSS = TEST_CUDA and (torch.version.cuda and int(torch.version.cuda.split(".")[0]) >= 12) +TEST_CUDA_GRAPH_CONDITIONAL_NODES = TEST_CUDA_GRAPH and ( + torch.version.cuda and ( + (int(torch.version.cuda.split(".")[0]) >= 12 and int(torch.version.cuda.split(".")[1]) >= 4) or + (int(torch.version.cuda.split(".")[0]) >= 13) + ) +) + +TEST_CUDA_PYTHON_BINDINGS = _check_module_exists("cuda.bindings") and ( + torch.version.cuda and int(torch.version.cuda.split(".")[0]) >= 12 +) + +if TEST_CUDA_PYTHON_BINDINGS: + def cuda_python_error_check(function_call_output): + """Makes calls to cuda-python's cuda runtime functions more + pythonic by throwing an exception if they return a status + which is not cudaSuccess + """ + import cuda.bindings # type: ignore[import] + + error, *others = function_call_output + if error != cuda.bindings.runtime.cudaError_t.cudaSuccess: + raise ValueError(f"CUDA failure! {error}") + else: + return tuple(others) +else: + cuda_python_error_check = None # type: ignore[assignment] + +def allocator_option_enabled_fn(allocator_config, _, option): + if allocator_config is None: + return False + allocator_config = allocator_config.split(',') if ',' in allocator_config else [allocator_config] + mapping = dict([var.split(':') for var in allocator_config]) + + if option in mapping and mapping[option] == 'True': + return True + else: + return False + +EXPANDABLE_SEGMENTS: bool = TestEnvironment.def_flag( + "EXPANDABLE_SEGMENTS", + env_var="PYTORCH_CUDA_ALLOC_CONF", + enabled_fn=functools.partial(allocator_option_enabled_fn, option='expandable_segments'), +) + +if TEST_CUDA and 'NUM_PARALLEL_PROCS' in os.environ: + num_procs = int(os.getenv("NUM_PARALLEL_PROCS", "2")) + gb_available = torch.cuda.mem_get_info()[1] / 2 ** 30 + # other libraries take up about a little under 1 GB of space per process + torch.cuda.set_per_process_memory_fraction(round((gb_available - num_procs * .85) / gb_available / num_procs, 2)) + +requires_cuda = unittest.skipUnless(torch.cuda.is_available(), "Requires CUDA") + +def skipIfCrossRef(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if TEST_WITH_CROSSREF: + raise unittest.SkipTest("test doesn't currently with crossref") + else: + fn(*args, **kwargs) + return wrapper + +class CrossRefMode(torch.overrides.TorchFunctionMode): + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + r = func(*args, **kwargs) + return r + +# Run PyTorch tests with TorchDynamo +TEST_WITH_TORCHINDUCTOR: bool = TestEnvironment.def_flag( + "TEST_WITH_TORCHINDUCTOR", + env_var="PYTORCH_TEST_WITH_INDUCTOR", +) +# AOT_EAGER not tested in ci, useful for debugging +TEST_WITH_AOT_EAGER: bool = TestEnvironment.def_flag( + "TEST_WITH_AOT_EAGER", + env_var="PYTORCH_TEST_WITH_AOT_EAGER", +) +TEST_WITH_TORCHDYNAMO: bool = TestEnvironment.def_flag( + "TEST_WITH_TORCHDYNAMO", + env_var="PYTORCH_TEST_WITH_DYNAMO", + implied_by_fn=lambda: TEST_WITH_TORCHINDUCTOR or TEST_WITH_AOT_EAGER, +) +TEST_WITHOUT_COMPILED_AUTOGRAD: bool = TestEnvironment.def_flag( + "TEST_WITHOUT_COMPILED_AUTOGRAD", + env_var="PYTORCH_TEST_WITHOUT_COMPILED_AUTOGRAD", +) + +if TEST_WITH_TORCHDYNAMO: + import torch._dynamo + # Do not spend time on helper functions that are called with different inputs + torch._dynamo.config.accumulated_recompile_limit = 64 + # Do not log compilation metrics from unit tests + torch._dynamo.config.log_compilation_metrics = False + # Silence 3.13.0 guard performance warnings + torch._dynamo.config.issue_3_13_0_warning = False + if TEST_WITH_TORCHINDUCTOR: + import torch._inductor.config + torch._inductor.config.fallback_random = True + else: + # only dynamo for now + torch._dynamo.config.compiled_autograd = not TEST_WITHOUT_COMPILED_AUTOGRAD + + +# seems like this is only used in test/torch_np +def xpassIfTorchDynamo_np(func): + # numpy 2.0+ is causing issues + if TEST_WITH_TORCHDYNAMO and np.__version__[0] == '2': + return unittest.skip("skipping numpy 2.0+ dynamo-wrapped test")(func) + return func if TEST_WITH_TORCHDYNAMO else unittest.expectedFailure(func) + + +def xfailIfACL(func): + return unittest.expectedFailure(func) if TEST_ACL else func + + +def xfailIfTorchDynamo(func): + return unittest.expectedFailure(func) if TEST_WITH_TORCHDYNAMO else func + + +def xfailIfPy312Plus(func): + return unittest.expectedFailure(func) if sys.version_info >= (3, 12) else func + + +def xfailIfLinux(func): + return unittest.expectedFailure(func) if IS_LINUX and not TEST_WITH_ROCM and not IS_FBCODE else func + + +def xfailIfWindows(func): + return unittest.expectedFailure(func) if IS_WINDOWS else func + + +def xfailIfROCm(func): + return unittest.expectedFailure(func) if torch.version.hip is not None else func + + +def skipIfFreeThreaded(msg="Test doesn't work with free-threaded python"): + if not isinstance(msg, str): + raise AssertionError("Are you using skipIfFreeThreaded correctly?") + return unittest.skipIf(sysconfig.get_config_var("Py_GIL_DISABLED") == 1, msg) + + +def skipIfTorchDynamo(msg="test doesn't currently work with dynamo"): + """ + Usage: + @skipIfTorchDynamo(msg) + def test_blah(self): + ... + """ + if not isinstance(msg, str): + raise AssertionError("Are you using skipIfTorchDynamo correctly?") + + def decorator(fn): + if not isinstance(fn, type): + @wraps(fn) + def wrapper(*args, **kwargs): + if TEST_WITH_TORCHDYNAMO: + raise unittest.SkipTest(msg) + else: + fn(*args, **kwargs) + return wrapper + + if not isinstance(fn, type): + raise AssertionError(f"expected fn to be a type, got {type(fn)}") + if TEST_WITH_TORCHDYNAMO: + fn.__unittest_skip__ = True # type: ignore[attr-defined] + fn.__unittest_skip_why__ = msg # type: ignore[attr-defined] + + return fn + + return decorator + +def skipIfTorchInductor(msg="test doesn't currently work with torchinductor", + condition=TEST_WITH_TORCHINDUCTOR): + def decorator(fn): + if not isinstance(fn, type): + @wraps(fn) + def wrapper(*args, **kwargs): + if condition: + raise unittest.SkipTest(msg) + else: + fn(*args, **kwargs) + return wrapper + + if not isinstance(fn, type): + raise AssertionError(f"expected fn to be a type, got {type(fn)}") + if condition: + fn.__unittest_skip__ = True # type: ignore[attr-defined] + fn.__unittest_skip_why__ = msg # type: ignore[attr-defined] + + return fn + + return decorator + +def runWithoutCompiledAutograd(msg="test doesn't currently work with compiled autograd"): + """ + Usage: + @runWithoutCompiledAutograd(msg) + def test_blah(self): + ... + """ + if not isinstance(msg, str): + raise AssertionError(f"expected msg to be str, got {type(msg)}") + + def decorator(func): + @wraps(func) + def wrapper(*args, **kwargs): + with torch._dynamo.compiled_autograd._disable(): + func(*args, **kwargs) + return wrapper + + return decorator + +def serialTest(condition=True): + """ + Decorator for running tests serially. Requires pytest + """ + # If one apply decorator directly condition will be callable + # And test will essentially be essentially skipped, which is undesirable + if type(condition) is not bool: + raise AssertionError(f"expected condition to be bool, got {type(condition)}") + + def decorator(fn): + if has_pytest and condition: + return pytest.mark.serial(fn) + return fn + return decorator + +def unMarkDynamoStrictTest(cls=None): + def decorator(cls): + cls.dynamo_strict = False + return cls + + if cls is None: + return decorator + else: + return decorator(cls) + + +def markDynamoStrictTest(cls_or_func=None, nopython=False): + """ + Marks the test as 'strict'. In strict mode, we reset before and after the + test, and run without suppress errors. + + Args: + - nopython: if we should run torch._dynamo.optimize with nopython={True/False}. + """ + def decorator(cls_or_func): + if inspect.isclass(cls_or_func): + cls_or_func.dynamo_strict = True + cls_or_func.dynamo_strict_nopython = nopython + return cls_or_func + + fn = cls_or_func + + @wraps(fn) + def wrapper(*args, **kwargs): + torch._dynamo.reset() + with unittest.mock.patch("torch._dynamo.config.suppress_errors", False): + fn(*args, **kwargs) + torch._dynamo.reset() + return wrapper + + if cls_or_func is None: + return decorator + else: + return decorator(cls_or_func) + + +def skipRocmIfTorchInductor(msg="test doesn't currently work with torchinductor on the ROCm stack"): + return skipIfTorchInductor(msg=msg, condition=TEST_WITH_ROCM and TEST_WITH_TORCHINDUCTOR) + +def skipIfLegacyJitExecutor(msg="test doesn't currently work with legacy JIT executor"): + def decorator(fn): + if not isinstance(fn, type): + @wraps(fn) + def wrapper(*args, **kwargs): + if not GRAPH_EXECUTOR: + raise AssertionError("GRAPH_EXECUTOR must be set") + if GRAPH_EXECUTOR == ProfilingMode.LEGACY: + raise unittest.SkipTest(msg) + else: + fn(*args, **kwargs) + return wrapper + + if not isinstance(fn, type): + raise AssertionError(f"expected fn to be a type, got {type(fn)}") + if GRAPH_EXECUTOR == ProfilingMode.LEGACY: + fn.__unittest_skip__ = True # type: ignore[attr-defined] + fn.__unittest_skip_why__ = msg # type: ignore[attr-defined] + + return fn + + + return decorator + + +def make_dynamo_test( + fn: Callable[..., Any] | None = None +) -> Callable[..., Any]: + """ + Decorator function to create a dynamo test case. A function annotate with + this decorator takes as input a unittest object. + """ + from torch._dynamo.testing import CompileCounter, reset, optimize_assert + if fn is None: + return lambda fn: make_dynamo_test(fn) + + def standard_test( + self: Any, + fn: Callable[..., Any], + kwargs, + ) -> None: + def dummy() -> None: + fn(self, **kwargs) + + actual = CompileCounter() + + dummy() + reset() + opt_fn = optimize_assert(actual)(dummy) + opt_fn() + reset() + + @functools.wraps(fn) + def test_fn(self: Any, **kwargs) -> None: + return standard_test( + self, + fn=fn, + kwargs=kwargs, + ) + + return test_fn + + +# Run PyTorch tests with translation validation on. +TEST_WITH_TV = os.getenv('PYTORCH_TEST_WITH_TV') == '1' + +if TEST_WITH_TV: + torch.fx.experimental._config.translation_validation = True + +# Determine whether to enable cuda memory leak check. +# CUDA mem leak check is expensive and thus we don't want to execute it on every +# test case / configuration. +# If this is True then CUDA memory leak checks are skipped. If this is false +# then CUDA memory leak checks are performed. +# See: https://github.com/pytorch/pytorch/pull/59402#issuecomment-858811135 +TEST_CUDA_MEM_LEAK_CHECK: bool = TestEnvironment.def_flag( + "TEST_CUDA_MEM_LEAK_CHECK", + env_var="PYTORCH_TEST_CUDA_MEM_LEAK_CHECK", +) + + +# Dict of NumPy dtype -> torch dtype (when the correspondence exists) +numpy_to_torch_dtype_dict = { + np.bool_ : torch.bool, + np.uint8 : torch.uint8, + np.uint16 : torch.uint16, + np.uint32 : torch.uint32, + np.uint64 : torch.uint64, + np.int8 : torch.int8, + np.int16 : torch.int16, + np.int32 : torch.int32, + np.int64 : torch.int64, + np.float16 : torch.float16, + np.float32 : torch.float32, + np.float64 : torch.float64, + np.complex64 : torch.complex64, + np.complex128 : torch.complex128 +} + + +# numpy dtypes like np.float64 are not instances, but rather classes. This leads to rather absurd cases like +# np.float64 != np.dtype("float64") but np.float64 == np.dtype("float64").type. +# Especially when checking against a reference we can't be sure which variant we get, so we simply try both. +def numpy_to_torch_dtype(np_dtype): + try: + return numpy_to_torch_dtype_dict[np_dtype] + except KeyError: + return numpy_to_torch_dtype_dict[np_dtype.type] + + +def has_corresponding_torch_dtype(np_dtype): + try: + numpy_to_torch_dtype(np_dtype) + return True + except KeyError: + return False + + +if IS_WINDOWS: + # Size of `np.intc` is platform defined. + # It is returned by functions like `bitwise_not`. + # On Windows `int` is 32-bit + # https://docs.microsoft.com/en-us/cpp/cpp/data-type-ranges?view=msvc-160 + numpy_to_torch_dtype_dict[np.intc] = torch.int + +# Dict of torch dtype -> NumPy dtype +torch_to_numpy_dtype_dict = {value : key for (key, value) in numpy_to_torch_dtype_dict.items()} +torch_to_numpy_dtype_dict.update({ + torch.bfloat16: np.float32, + torch.complex32: np.complex64 +}) + +def skipIfNNModuleInlined( + msg="test doesn't currently work with nn module inlining", + condition=True, +): + def decorator(fn): + if not isinstance(fn, type): + + @wraps(fn) + def wrapper(*args, **kwargs): + if condition: + raise unittest.SkipTest(msg) + else: + fn(*args, **kwargs) + + return wrapper + + if not isinstance(fn, type): + raise AssertionError(f"expected fn to be a type, got {type(fn)}") + if condition: + fn.__unittest_skip__ = True # type: ignore[attr-defined] + fn.__unittest_skip_why__ = msg # type: ignore[attr-defined] + + return fn + + return decorator + +def skipIfRocm(func=None, *, msg="test doesn't currently work on the ROCm stack"): + def dec_fn(fn): + reason = f"skipIfRocm: {msg}" + + @wraps(fn) + def wrapper(*args, **kwargs): + if TEST_WITH_ROCM: + raise unittest.SkipTest(reason) + else: + return fn(*args, **kwargs) + return wrapper + if func: + return dec_fn(func) + return dec_fn + +def getRocmArchName(device_index: int = 0): + return torch.cuda.get_device_properties(device_index).gcnArchName + +def isRocmArchAnyOf(arch: tuple[str, ...]): + if not torch.version.hip: + return False + rocmArch = getRocmArchName() + return any(x in rocmArch for x in arch) + +def skipIfRocmArch(arch: tuple[str, ...]): + def dec_fn(fn): + @wraps(fn) + def wrap_fn(self, *args, **kwargs): + if TEST_WITH_ROCM and isRocmArchAnyOf(arch): + reason = f"skipIfRocm: test skipped on {arch}" + raise unittest.SkipTest(reason) + return fn(self, *args, **kwargs) + return wrap_fn + return dec_fn + +def runOnRocm(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if TEST_WITH_ROCM: + fn(*args, **kwargs) + else: + raise unittest.SkipTest("test currently only works on the ROCm stack") + return wrapper + +def runOnRocmArch(arch: tuple[str, ...]): + def dec_fn(fn): + @wraps(fn) + def wrap_fn(self, *args, **kwargs): + if TEST_WITH_ROCM and not isRocmArchAnyOf(arch): + reason = f"skipIfRocm: test only runs on {arch}" + raise unittest.SkipTest(reason) + return fn(self, *args, **kwargs) + return wrap_fn + return dec_fn + +def xfailIfS390X(func): + return unittest.expectedFailure(func) if IS_S390X else func + +def xfailIf(condition): + def wrapper(func): + if condition: + return unittest.expectedFailure(func) + else: + return func + return wrapper + +def skipIfXpu(func=None, *, msg="test doesn't currently work on the XPU stack"): + def dec_fn(fn): + reason = f"skipIfXpu: {msg}" + + @wraps(fn) + def wrapper(*args, **kwargs): + if TEST_XPU: + raise unittest.SkipTest(reason) + else: + return fn(*args, **kwargs) + return wrapper + if func: + return dec_fn(func) + return dec_fn + +def skipIfMPS(fn): + sig = inspect.signature(fn) + has_device_arg = "device" in sig.parameters + + if not has_device_arg: + warnings.warn( + f"skipIfMPS applied to {fn.__qualname__} which has no 'device' parameter. " + "Consider using device-generic tests with instantiate_device_type_tests instead.", + stacklevel=2, + ) + + @wraps(fn) + def wrapper(*args, **kwargs): + if has_device_arg: + # For device-generic tests, only skip when actually running on MPS + slf = args[0] if args else None + if slf is not None: + device_type = getattr(slf, "device_type", None) or getattr( + slf, "device", None + ) + if isinstance(device_type, str) and device_type == "mps": + raise unittest.SkipTest("test doesn't currently work with MPS") + elif TEST_MPS: + raise unittest.SkipTest("test doesn't currently work with MPS") + return fn(*args, **kwargs) + + return wrapper + + +def skipIfHpu(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if TEST_HPU: + raise unittest.SkipTest("test doesn't currently work with HPU") + else: + fn(*args, **kwargs) + return wrapper + +def getRocmVersion() -> tuple[int, int]: + from torch.testing._internal.common_cuda import _get_torch_rocm_version + rocm_version = _get_torch_rocm_version() + return (rocm_version[0], rocm_version[1]) + +# Skips a test on CUDA if ROCm is available and its version is lower than requested. +def skipIfRocmVersionLessThan(version=None): + def dec_fn(fn): + @wraps(fn) + def wrap_fn(self, *args, **kwargs): + if TEST_WITH_ROCM: + rocm_version_tuple = getRocmVersion() + if rocm_version_tuple is None or version is None or rocm_version_tuple < tuple(version): + reason = f"ROCm {rocm_version_tuple} is available but {version} required" + raise unittest.SkipTest(reason) + return fn(self, *args, **kwargs) + return wrap_fn + return dec_fn + +def skipIfNotMiopenSuggestNHWC(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if not TEST_WITH_MIOPEN_SUGGEST_NHWC: + raise unittest.SkipTest("test doesn't currently work without MIOpen NHWC activation") + else: + fn(*args, **kwargs) + return wrapper + +def skipIfWindows(func=None, *, msg="test doesn't currently work on the Windows stack"): + def dec_fn(fn): + reason = f"skipIfWindows: {msg}" + + @wraps(fn) + def wrapper(*args, **kwargs): + if IS_WINDOWS: # noqa: F821 + raise unittest.SkipTest(reason) + else: + return fn(*args, **kwargs) + return wrapper + if func: + return dec_fn(func) + return dec_fn + +def skipIfWindowsXPU(func=None, *, msg="test doesn't currently work on the Windows stack"): + def dec_fn(fn): + reason = f"skipIfWindowsXPU: {msg}" + + @wraps(fn) + def wrapper(*args, **kwargs): + if IS_WINDOWS and torch.xpu.is_available(): # noqa: F821 + raise unittest.SkipTest(reason) + else: + return fn(*args, **kwargs) + return wrapper + if func: + return dec_fn(func) + return dec_fn + +def requires_cuda_p2p_access(): + cuda_p2p_access_available = ( + torch.cuda.is_available() + and torch.cuda.get_device_capability() >= (8, 0) + and torch.cuda.device_count() >= 2 + ) + num_devices = torch.cuda.device_count() + for i in range(num_devices - 1): + for j in range(i + 1, num_devices): + if not torch.cuda.can_device_access_peer(i, j): + cuda_p2p_access_available = False + break + if not cuda_p2p_access_available: + break + + return skip_but_pass_in_sandcastle_if( + not cuda_p2p_access_available, + "cuda p2p access is not available", + ) + +# Reverts the linalg backend back to default to make sure potential failures in one +# test do not affect other tests +def setLinalgBackendsToDefaultFinally(fn): + @wraps(fn) + def _fn(*args, **kwargs): + _preferred_backend = torch.backends.cuda.preferred_linalg_library() + try: + fn(*args, **kwargs) + finally: + torch.backends.cuda.preferred_linalg_library(_preferred_backend) + return _fn + + +# Reverts the blas backend back to default to make sure potential failures in one +# test do not affect other tests +def setBlasBackendsToDefaultFinally(fn): + @wraps(fn) + def _fn(*args, **kwargs): + _preferred_backend = torch.backends.cuda.preferred_blas_library() + try: + fn(*args, **kwargs) + finally: + torch.backends.cuda.preferred_blas_library(_preferred_backend) + return _fn + + +# Context manager for setting deterministic flag and automatically +# resetting it to its original value +class DeterministicGuard: + def __init__(self, deterministic, *, warn_only=False, fill_uninitialized_memory=True): + self.deterministic = deterministic + self.warn_only = warn_only + self.fill_uninitialized_memory = fill_uninitialized_memory + + @classmethod + def _current_state(cls): + return cls( + torch.are_deterministic_algorithms_enabled(), + warn_only=torch.is_deterministic_algorithms_warn_only_enabled(), + fill_uninitialized_memory=torch.utils.deterministic.fill_uninitialized_memory, # type: ignore[attr-defined] + ) + + def _update(self): + torch.use_deterministic_algorithms(self.deterministic, warn_only=self.warn_only) + torch.utils.deterministic.fill_uninitialized_memory = self.fill_uninitialized_memory # type: ignore[attr-defined] + + def __enter__(self): + self._restore = self._current_state() + self._update() + + def __exit__(self, exception_type, exception_value, traceback): + self._restore._update() + +class AlwaysWarnTypedStorageRemoval: + def __init__(self, always_warn): + if not isinstance(always_warn, bool): + raise AssertionError(f"expected always_warn to be bool, got {type(always_warn)}") + self.always_warn = always_warn + + def __enter__(self): + self.always_warn_restore = torch.storage._get_always_warn_typed_storage_removal() + torch.storage._set_always_warn_typed_storage_removal(self.always_warn) + + def __exit__(self, exception_type, exception_value, traceback): + torch.storage._set_always_warn_typed_storage_removal(self.always_warn_restore) + +# Context manager for setting cuda sync debug mode and reset it +# to original value +# we are not exposing it to the core because sync debug mode is +# global and thus not thread safe +class CudaSyncGuard: + def __init__(self, sync_debug_mode): + self.mode = sync_debug_mode + + def __enter__(self): + self.debug_mode_restore = torch.cuda.get_sync_debug_mode() + torch.cuda.set_sync_debug_mode(self.mode) + + def __exit__(self, exception_type, exception_value, traceback): + torch.cuda.set_sync_debug_mode(self.debug_mode_restore) + +# Context manager for setting torch.__future__.set_swap_module_params_on_conversion +# and automatically resetting it to its original value +class SwapTensorsGuard: + def __init__(self, use_swap_tensors): + self.use_swap_tensors = use_swap_tensors + + def __enter__(self): + self.swap_tensors_restore = torch.__future__.get_swap_module_params_on_conversion() + if self.use_swap_tensors is not None: + torch.__future__.set_swap_module_params_on_conversion(self.use_swap_tensors) + + def __exit__(self, exception_type, exception_value, traceback): + torch.__future__.set_swap_module_params_on_conversion(self.swap_tensors_restore) + +# This decorator can be used for API tests that call +# torch.use_deterministic_algorithms(). When the test is finished, it will +# restore the previous deterministic flag setting. +# +# If CUDA >= 10.2, this will set the environment variable +# CUBLAS_WORKSPACE_CONFIG=:4096:8 so that the error associated with that +# setting is not thrown during the test unless the test changes that variable +# on purpose. The previous CUBLAS_WORKSPACE_CONFIG setting will also be +# restored once the test is finished. +# +# Note that if a test requires CUDA to actually register the changed +# CUBLAS_WORKSPACE_CONFIG variable, a new subprocess must be created, because +# CUDA only checks the variable when the runtime initializes. Tests can be +# run inside a subprocess like so: +# +# import subprocess, sys, os +# script = ''' +# # Test code should go here +# ''' +# try: +# subprocess.check_output( +# [sys.executable, '-c', script], +# stderr=subprocess.STDOUT, +# cwd=os.path.dirname(os.path.realpath(__file__)), +# env=os.environ.copy()) +# except subprocess.CalledProcessError as e: +# error_message = e.output.decode('utf-8') +# # Handle exceptions raised by the subprocess here +# +def wrapDeterministicFlagAPITest(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + with DeterministicGuard( + torch.are_deterministic_algorithms_enabled(), + warn_only=torch.is_deterministic_algorithms_warn_only_enabled()): + class CuBLASConfigGuard: + cublas_var_name = 'CUBLAS_WORKSPACE_CONFIG' + + def __enter__(self): + self.cublas_config_restore = os.environ.get(self.cublas_var_name) + os.environ[self.cublas_var_name] = ':4096:8' + + def __exit__(self, exception_type, exception_value, traceback): + cur_cublas_config = os.environ.get(self.cublas_var_name) + if self.cublas_config_restore is None: + if cur_cublas_config is not None: + del os.environ[self.cublas_var_name] + else: + os.environ[self.cublas_var_name] = self.cublas_config_restore + with CuBLASConfigGuard(): + fn(*args, **kwargs) + return wrapper + +# This decorator can be used for API tests that want to safely call +# torch.__future__.set_swap_module_params_on_conversion. `swap` can be set to +# True, False or None where None indicates that the context manager does not +# set the flag. When the test is finished, it will restore the previous swap +# flag setting. +def wrapSwapTensorsTest(swap=None): + def dec_fn(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + with SwapTensorsGuard(swap): + fn(*args, **kwargs) + return wrapper + return dec_fn + +# test parametrizer for swapping +class swap(_TestParametrizer): + def __init__(self, swap_values): + super().__init__() + self.swap_values = swap_values + + def _parametrize_test(self, test, generic_cls, device_cls): + for swap in self.swap_values: + yield wrapSwapTensorsTest(swap)(test), f'swap_{swap}', {}, lambda _: [] + +def skipIfCompiledWithoutNumpy(fn): + # Even if the numpy module is present, if `USE_NUMPY=0` is used during the + # build, numpy tests will fail + numpy_support = TEST_NUMPY + if numpy_support: + try: + # The numpy module is present, verify that PyTorch is compiled with + # numpy support + torch.from_numpy(np.array([2, 2])) + except RuntimeError: + numpy_support = False + + @wraps(fn) + def wrapper(*args, **kwargs): + if not numpy_support: + raise unittest.SkipTest("PyTorch was compiled without numpy support") + else: + fn(*args, **kwargs) + return wrapper + +def _test_function(fn, device): + def run_test_function(self): + return fn(self, device) + return run_test_function + +def skipIfNoXNNPACK(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if not torch.backends.xnnpack.enabled: # type: ignore[attr-defined] + raise unittest.SkipTest('XNNPACK must be enabled for these tests. Please build with USE_XNNPACK=1.') + else: + fn(*args, **kwargs) + return wrapper + +def skipIfNoLapack(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if not torch._C.has_lapack: + raise unittest.SkipTest('PyTorch compiled without Lapack') + else: + fn(*args, **kwargs) + return wrapper + +def skipIfNotRegistered(op_name, message): + """Wraps the decorator to hide the import of the `core`. + + Args: + op_name: Check if this op is registered in `core._REGISTERED_OPERATORS`. + message: message to fail with. + + Usage: + @skipIfNotRegistered('MyOp', 'MyOp is not linked!') + This will check if 'MyOp' is in the caffe2.python.core + """ + return unittest.skip("Pytorch is compiled without Caffe2") + +def skipIfNoSciPy(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if not TEST_SCIPY: + raise unittest.SkipTest("test require SciPy, but SciPy not found") + else: + fn(*args, **kwargs) + return wrapper + +def skip_if_pytest(fn): + @wraps(fn) + def wrapped(*args, **kwargs): + if "PYTEST_CURRENT_TEST" in os.environ: + raise unittest.SkipTest("does not work under pytest") + return fn(*args, **kwargs) + + return wrapped + +def skipIfNoXPU(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if not TEST_XPU: + raise unittest.SkipTest("test required PyTorched compiled with XPU") + else: + fn(*args, **kwargs) + return wrapper + +def slowTest(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + if not TEST_WITH_SLOW: + raise unittest.SkipTest("test is slow; run with PYTORCH_TEST_WITH_SLOW to enable test") + else: + fn(*args, **kwargs) + wrapper.__dict__['slow_test'] = True + return wrapper + + +def slowTestIf(condition): + return slowTest if condition else lambda fn: fn + + +def skipCUDAMemoryLeakCheckIf(condition): + def dec(fn): + if getattr(fn, '_do_cuda_memory_leak_check', True): # if current True + fn._do_cuda_memory_leak_check = not condition + return fn + return dec + +def skipCUDANonDefaultStreamIf(condition): + def dec(fn): + if getattr(fn, '_do_cuda_non_default_stream', True): # if current True + fn._do_cuda_non_default_stream = not condition + return fn + return dec + +def suppress_warnings(fn): + @wraps(fn) + def wrapper(*args, **kwargs): + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + fn(*args, **kwargs) + return wrapper + + +def to_gpu(obj, type_map=None): + if type_map is None: + type_map = {} + if isinstance(obj, torch.Tensor): + if not obj.is_leaf: + raise AssertionError("expected obj to be a leaf tensor") + t = type_map.get(obj.dtype, obj.dtype) + with torch.no_grad(): + res = obj.to(dtype=t, device="cuda", copy=True) + res.requires_grad = obj.requires_grad + return res + elif torch.is_storage(obj): + return obj.new().resize_(obj.size()).copy_(obj) # type: ignore[attr-defined, union-attr] + elif isinstance(obj, list): + return [to_gpu(o, type_map) for o in obj] + elif isinstance(obj, tuple): + return tuple(to_gpu(o, type_map) for o in obj) + else: + return deepcopy(obj) + + +def get_function_arglist(func): + return inspect.getfullargspec(func).args + + +def set_rng_seed(seed=None): + if seed is None: + seed = SEED + torch.manual_seed(seed) + random.seed(seed) + if TEST_NUMPY: + np.random.seed(seed) + + +@contextlib.contextmanager +def set_default_dtype(dtype): + saved_dtype = torch.get_default_dtype() + torch.set_default_dtype(dtype) + try: + yield + finally: + torch.set_default_dtype(saved_dtype) + +@contextlib.contextmanager +def set_default_tensor_type(tensor_type): + saved_tensor_type = torch.tensor([]).type() + torch.set_default_tensor_type(tensor_type) + try: + yield + finally: + torch.set_default_tensor_type(saved_tensor_type) + +def iter_indices(tensor): + if tensor.dim() == 0: + return range(0) + if tensor.dim() == 1: + return range(tensor.size(0)) + return product(*(range(s) for s in tensor.size())) + + +def is_iterable(obj): + try: + iter(obj) + return True + except TypeError: + return False + + +def is_iterable_of_tensors(iterable, include_empty=False): + """ Returns True if iterable is an iterable of tensors and False o.w. + + If the iterable is empty, the return value is :attr:`include_empty` + """ + # Tensor itself is iterable so we check this first + if isinstance(iterable, torch.Tensor): + return False + + try: + if len(iterable) == 0: + return include_empty + + for t in iter(iterable): + if not isinstance(t, torch.Tensor): + return False + + except TypeError: + return False + + return True + + +class CudaNonDefaultStream: + def __enter__(self): + # Before starting CUDA test save currently active streams on all + # CUDA devices and set new non default streams to all CUDA devices + # to ensure CUDA tests do not use default stream by mistake. + beforeDevice = torch.cuda.current_device() + self.beforeStreams = [] + for d in range(torch.cuda.device_count()): + self.beforeStreams.append(torch.cuda.current_stream(d)) + deviceStream = torch.cuda.Stream(device=d) + self.beforeStreams[-1].synchronize() + torch._C._cuda_setStream(stream_id=deviceStream.stream_id, + device_index=deviceStream.device_index, + device_type=deviceStream.device_type) + torch._C._cuda_setDevice(beforeDevice) + + def __exit__(self, exc_type, exc_value, traceback): + # After completing CUDA test load previously active streams on all + # CUDA devices. + beforeDevice = torch.cuda.current_device() + for d in range(torch.cuda.device_count()): + torch._C._cuda_setStream(stream_id=self.beforeStreams[d].stream_id, + device_index=self.beforeStreams[d].device_index, + device_type=self.beforeStreams[d].device_type) + torch._C._cuda_setDevice(beforeDevice) + +class CudaMemoryLeakCheck: + def __init__(self, testcase, name=None): + self.name = testcase.id() if name is None else name + self.testcase = testcase + + # initialize context & RNG to prevent false positive detections + # when the test is the first to initialize those + from torch.testing._internal.common_cuda import initialize_cuda_context_rng + initialize_cuda_context_rng() + + # Stores CUDA memory data provided by PyTorch's caching allocator and + # the CUDA driver. + # + # NOTE: The undocumented torch.cuda.mem_get_info() returns + # (#free bytes, #total bytes available) on the GPU + def __enter__(self): + self.caching_allocator_befores = [] + self.driver_befores = [] + + # Performs a gc if required (required if any CUDA memory is held) + num_devices = torch.cuda.device_count() + for i in range(num_devices): + caching_allocator_mem_allocated = torch.cuda.memory_allocated(i) + # NOTE: gc is based exclusively on caching allocator memory + # because the driver will always have some bytes in use (context size?) + if caching_allocator_mem_allocated > 0: + gc.collect() + torch._C._cuda_clearCublasWorkspaces() + torch.cuda.empty_cache() + break + + # Acquires caching allocator and driver statistics before the test is run + for i in range(num_devices): + self.caching_allocator_befores.append(torch.cuda.memory_allocated(i)) + bytes_free, bytes_total = torch.cuda.mem_get_info(i) + driver_mem_allocated = bytes_total - bytes_free + self.driver_befores.append(driver_mem_allocated) + + def __exit__(self, exc_type, exc_value, traceback): + # Don't check for leaks if an exception was thrown + if exc_type is not None: + return + + # Compares caching allocator before/after statistics + # An increase in allocated memory is a discrepancy indicating a possible + # memory leak + discrepancy_detected = False + num_devices = torch.cuda.device_count() + for i in range(num_devices): + # avoid counting cublasWorkspace allocations + torch._C._cuda_clearCublasWorkspaces() + caching_allocator_mem_allocated = torch.cuda.memory_allocated(i) + + if caching_allocator_mem_allocated > self.caching_allocator_befores[i]: + discrepancy_detected = True + break + + # Short-circuits if no discrepancy detected + if not discrepancy_detected: + return + + # Validates the discrepancy persists after garbage collection and + # is confirmed by the driver API + + # NOTE: driver API iscrepancies alone are ignored because with the jiterator + # some tests may permanently increase the CUDA context size and + # that will appear as a driver memory leak but is the expected behavior. + + # GCs and clears the cache + gc.collect() + torch.cuda.empty_cache() + + for i in range(num_devices): + + discrepancy_detected = True + + # Query memory multiple items to ensure leak was not transient + for _ in range(3): + caching_allocator_mem_allocated = torch.cuda.memory_allocated(i) + bytes_free, bytes_total = torch.cuda.mem_get_info(i) + driver_mem_allocated = bytes_total - bytes_free + + caching_allocator_discrepancy = False + driver_discrepancy = False + + if caching_allocator_mem_allocated > self.caching_allocator_befores[i]: + caching_allocator_discrepancy = True + + if driver_mem_allocated > self.driver_befores[i]: + driver_discrepancy = True + + if not (caching_allocator_discrepancy or driver_discrepancy): + # Leak was false positive, exit loop + discrepancy_detected = False + break + + if not discrepancy_detected: + continue + + if caching_allocator_discrepancy and not driver_discrepancy: # type: ignore[possibly-undefined] + # Just raises a warning if the leak is not validated by the + # driver API + # NOTE: this may be a problem with how the caching allocator collects its + # statistics or a leak too small to trigger the allocation of an + # additional block of memory by the CUDA driver + msg = ("CUDA caching allocator reports a memory leak not " # type: ignore[possibly-undefined] + f"verified by the driver API in {self.name}! " + f"Caching allocator allocated memory was {self.caching_allocator_befores[i]} " + f"and is now reported as {caching_allocator_mem_allocated} " # type: ignore[possibly-undefined] + f"on device {i}. " + f"CUDA driver allocated memory was {self.driver_befores[i]} and is now {driver_mem_allocated}.") # type: ignore[possibly-undefined] + warnings.warn(msg, stacklevel=2) + elif caching_allocator_discrepancy and driver_discrepancy: # type: ignore[possibly-undefined] + # A caching allocator discrepancy validated by the driver API is a + # failure (except on ROCm, see below) + msg = (f"CUDA driver API confirmed a leak in {self.name}! " # type: ignore[possibly-undefined] + f"Caching allocator allocated memory was {self.caching_allocator_befores[i]} " + f"and is now reported as {caching_allocator_mem_allocated} " # type: ignore[possibly-undefined] + f"on device {i}. " + f"CUDA driver allocated memory was {self.driver_befores[i]} and is now {driver_mem_allocated}.") # type: ignore[possibly-undefined] + + raise RuntimeError(msg) + +@contextmanager +def skip_exception_type(exc_type): + try: + yield + except exc_type as e: + raise unittest.SkipTest(f"not implemented: {e}") from e + +@contextmanager +def print_repro_on_failure(repro_parts): + try: + yield + except unittest.SkipTest: + raise + except Exception as e: + # Get the index of the sample input that failed the test if possible. + sample_isolation_prefix = "" + tracked_input = getattr(e, "_tracked_input", None) + if tracked_input is not None: + sample_isolation_prefix = f"PYTORCH_OPINFO_SAMPLE_INPUT_INDEX={tracked_input.index}" + + repro_str = " ".join(filter(None, (sample_isolation_prefix, *repro_parts))) + + open_source_signpost( + subsystem="test_repros", + name="test_failure", + parameters=json.dumps( + { + "repro": " ".join(filter(None, (sample_isolation_prefix, *repro_parts))), + } + ), + ) + + repro_msg = f""" +To execute this test, run the following from the base repo dir: + {repro_str} + +This message can be suppressed by setting PYTORCH_PRINT_REPRO_ON_FAILURE=0""" + + # NB: Hacking the exception args is the cleanest way I've found to append + # failure reproduction info without poisoning the stack trace. + if len(e.args) >= 1: + e.args = (f"{e.args[0]}\n{repro_msg}", *e.args[1:]) + raise + +# "min_satisfying_examples" setting has been deprecated in hypothesis +# 3.56.0 and removed in hypothesis 4.x +try: + import hypothesis + + def settings(*args, **kwargs): + if 'min_satisfying_examples' in kwargs and hypothesis.version.__version_info__ >= (3, 56, 0): + kwargs.pop('min_satisfying_examples') + return hypothesis.settings(*args, **kwargs) + + + hypothesis.settings.register_profile( + "pytorch_ci", + settings( + derandomize=True, + suppress_health_check=[hypothesis.HealthCheck.too_slow], + database=None, + max_examples=50, + verbosity=hypothesis.Verbosity.normal)) + hypothesis.settings.register_profile( + "dev", + settings( + suppress_health_check=[hypothesis.HealthCheck.too_slow], + database=None, + max_examples=10, + verbosity=hypothesis.Verbosity.normal)) + hypothesis.settings.register_profile( + "debug", + settings( + suppress_health_check=[hypothesis.HealthCheck.too_slow], + database=None, + max_examples=1000, + verbosity=hypothesis.Verbosity.verbose)) + + hypothesis.settings.load_profile( + "pytorch_ci" if IS_CI else os.getenv('PYTORCH_HYPOTHESIS_PROFILE', 'dev') + ) +except ImportError: + warnings.warn('Fail to import hypothesis in common_utils, tests are not derandomized', ImportWarning, stacklevel=2) + +# Used in check_if_enable to see if a test method should be disabled by an issue, +# sanitizes a test method name from appended suffixes by @dtypes parametrization. +# e.g., an issue with title "DISABLED test_bitwise_ops (__main__.TestBinaryUfuncs)" should +# disabled ALL parametrized test_bitwise_ops tests, such test_bitwise_ops_cuda_int32 +def remove_device_and_dtype_suffixes(test_name: str) -> str: + # import statement is localized to avoid circular dependency issues with common_device_type.py + from torch.testing._internal.common_device_type import get_device_type_test_bases + device_suffixes = [x.device_type for x in get_device_type_test_bases()] + dtype_suffixes = [str(dt)[len("torch."):] for dt in get_all_dtypes()] + + test_name_chunks = test_name.split("_") + if len(test_name_chunks) > 0 and test_name_chunks[-1] in dtype_suffixes: + if len(test_name_chunks) > 1 and test_name_chunks[-2] in device_suffixes: + return "_".join(test_name_chunks[0:-2]) + return "_".join(test_name_chunks[0:-1]) + return test_name + + +def check_if_enable(test: unittest.TestCase): + classname = str(test.__class__).split("'")[1].split(".")[-1] + sanitized_testname = remove_device_and_dtype_suffixes(test._testMethodName) + + def matches_test(target: str): + target_test_parts = target.split() + if len(target_test_parts) < 2: + # poorly formed target test name + return False + target_testname = target_test_parts[0] + target_classname = target_test_parts[1][1:-1].split(".")[-1] + # if test method name or its sanitized version exactly matches the disabled + # test method name AND allow non-parametrized suite names to disable + # parametrized ones (TestSuite disables TestSuiteCPU) + return classname.startswith(target_classname) and (target_testname in (test._testMethodName, sanitized_testname)) + + if any(matches_test(x) for x in slow_tests_dict): + getattr(test, test._testMethodName).__dict__['slow_test'] = True + if not TEST_WITH_SLOW: + raise unittest.SkipTest("test is slow; run with PYTORCH_TEST_WITH_SLOW to enable test") + + if not IS_SANDCASTLE: + should_skip = False + skip_msg = "" + + for disabled_test, (issue_url, platforms) in disabled_tests_dict.items(): + if matches_test(disabled_test): + platform_to_conditional: dict = { + "mac": IS_MACOS, + "macos": IS_MACOS, + "win": IS_WINDOWS, + "windows": IS_WINDOWS, + "linux": IS_LINUX, + "rocm": TEST_WITH_ROCM, + "xpu": TEST_XPU, + "asan": TEST_WITH_ASAN, + "dynamo": TEST_WITH_TORCHDYNAMO, + "dynamo_wrapped": TEST_WITH_TORCHDYNAMO, + "inductor": TEST_WITH_TORCHINDUCTOR, + "slow": TEST_WITH_SLOW, + } + + invalid_platforms = list(filter(lambda p: p not in platform_to_conditional, platforms)) + if len(invalid_platforms) > 0: + invalid_plats_str = ", ".join(invalid_platforms) + valid_plats = ", ".join(platform_to_conditional.keys()) + + print(f"Test {disabled_test} is disabled for some unrecognized ", + f"platforms: [{invalid_plats_str}]. Please edit issue {issue_url} to fix the platforms ", + 'assigned to this flaky test, changing "Platforms: ..." to a comma separated ', + f"subset of the following (or leave it blank to match all platforms): {valid_plats}") + + # Sanitize the platforms list so that we continue to disable the test for any valid platforms given + platforms = list(filter(lambda p: p in platform_to_conditional, platforms)) + + if platforms == [] or any(platform_to_conditional[platform] for platform in platforms): + should_skip = True + skip_msg = f"Test is disabled because an issue exists disabling it: {issue_url}" \ + f" for {'all' if platforms == [] else ''}platform(s) {', '.join(platforms)}. " \ + "If you're seeing this on your local machine and would like to enable this test, " \ + "please make sure CI is not set and you are not using the flag --import-disabled-tests." + break + + if should_skip and not RERUN_DISABLED_TESTS: + # Skip the disabled test when not running under --rerun-disabled-tests verification mode + raise unittest.SkipTest(skip_msg) + + if not should_skip and RERUN_DISABLED_TESTS: + # Probably test has disable issue but not for this platform + skip_msg = "Test is enabled but --rerun-disabled-tests verification mode is set, so only" \ + " disabled tests are run" + raise unittest.SkipTest(skip_msg) + + if TEST_SKIP_FAST: + if hasattr(test, test._testMethodName) and not getattr(test, test._testMethodName).__dict__.get('slow_test', False): + raise unittest.SkipTest("test is fast; we disabled it with PYTORCH_TEST_SKIP_FAST") + + +# `TestCase.assertEqual` is very permissive and coerced the inputs into a format that could be compared. This is very +# convenient when writing tests, but not so much while reviewing them. By default, the comparison `Pair` framework of +# `torch.testing._comparison.are_equal`, used for example by the public testing function +# `torch.testing.assert_close`, is more strict. In order to use the same framework and thus reduce the divergence +# between internal and external comparison logic as much as possible, we define some "relaxed" pairs here. They only +# change the supported inputs, but the comparison logic is the same. +# TODO: Revisit the relaxed pairs and check how much work it is to fix the tests that would fail without the relaxation. + +class RelaxedBooleanPair(BooleanPair): + """Pair for boolean-like inputs. + + In contrast to the builtin :class:`BooleanPair`, this class also supports one input being a number or a single + element tensor-like. + """ + _supported_number_types = NumberPair(0, 0)._supported_types + + def _process_inputs(self, actual, expected, *, id): + # We require only one of the inputs of the inputs to be a boolean and the other can also be a boolean, a + # number, or a single element tensor or array, whereas in default BooleanPair both inputs have to be booleans. + tensor_or_array_types: tuple[type, ...] = (torch.Tensor, np.ndarray) + other_supported_types = (*self._supported_types, *self._supported_number_types, *tensor_or_array_types) + if not ( + (isinstance(actual, self._supported_types) and isinstance(expected, other_supported_types)) + or (isinstance(expected, self._supported_types) and isinstance(actual, other_supported_types)) + ): + self._inputs_not_supported() + + return [self._to_bool(input, id=id) for input in (actual, expected)] + + def _to_bool(self, bool_like, *, id): + if isinstance(bool_like, np.number): + return bool(bool_like.item()) + elif type(bool_like) in self._supported_number_types: + return bool(bool_like) + elif isinstance(bool_like, (torch.Tensor, np.ndarray)): + numel = bool_like.numel() if isinstance(bool_like, torch.Tensor) else bool_like.size + if numel > 1: + self._fail( + ValueError, + f"Only single element tensor-likes can be compared against a boolean. " + f"Got {numel} elements instead.", + id=id + ) + + return bool(bool_like.item()) + else: + return super()._to_bool(bool_like, id=id) + + +class RelaxedNumberPair(NumberPair): + """Pair for number-like inputs. + + In contrast to the builtin :class:`NumberPair`, this class also supports one input being a single element + tensor-like or a :class:`enum.Enum`. (D)Type checks are disabled, meaning comparing 1 to 1.0 succeeds even when + ``check_dtype=True`` is passed. + + In addition, this class uses looser default tolerances for :class:`float` and :class:`complex` inputs. Also + supports overriding the absolute and relative tolerance through the ``@precisionOverride`` and + ``@toleranceOverride`` decorators. + """ + _TYPE_TO_DTYPE = { + int: torch.int64, + float: torch.float32, + complex: torch.complex64, + } + + def __init__( + self, actual, expected, *, rtol_override=0.0, atol_override=0.0, check_dtype=None, **other_parameters + ) -> None: + super().__init__(actual, expected, check_dtype=False, **other_parameters) + self.rtol = max(self.rtol, rtol_override) + self.atol = max(self.atol, atol_override) + + def _process_inputs(self, actual, expected, *, id): + # We require only one of the inputs of the inputs to be a number and the other can also be a number or a single + # element tensor or array, whereas in default NumberPair both inputs have to be numbers. + tensor_or_array_types: tuple[type, ...] = (torch.Tensor, np.ndarray) + other_supported_types = (*self._supported_types, *tensor_or_array_types) + if not ( + (isinstance(actual, self._supported_types) and isinstance(expected, other_supported_types)) + or (isinstance(expected, self._supported_types) and isinstance(actual, other_supported_types)) + ): + self._inputs_not_supported() + + return [self._to_number(input, id=id) for input in (actual, expected)] + + def _to_number(self, number_like, *, id): + if isinstance(number_like, (torch.Tensor, np.ndarray)): + numel = number_like.numel() if isinstance(number_like, torch.Tensor) else number_like.size + if numel > 1: + self._fail( + ValueError, + f"Only single element tensor-likes can be compared against a number. " + f"Got {numel} elements instead.", + id=id + ) + number = number_like.item() + if isinstance(number, bool): + number = int(number) + + return number + elif isinstance(number_like, Enum): + return int(number_like) # type: ignore[call-overload] + else: + number = super()._to_number(number_like, id=id) + if type(number) not in self._TYPE_TO_DTYPE: + self._inputs_not_supported() + return number + + +class TensorOrArrayPair(TensorLikePair): + """Pair for tensor-like inputs. + + On the one hand this class is stricter than the builtin :class:`TensorLikePair` since it only allows instances of + :class:`torch.Tensor` and :class:`numpy.ndarray` rather than allowing any tensor-like than can be converted into a + tensor. On the other hand this class is looser since it converts all inputs into tensors with no regard of their + relationship, e.g. comparing a :class:`torch.Tensor` to :class:`numpy.ndarray` is fine. + + In addition, this class supports overriding the absolute and relative tolerance through the ``@precisionOverride`` + and ``@toleranceOverride`` decorators. + """ + def __init__(self, actual, expected, *, rtol_override=0.0, atol_override=0.0, **other_parameters): + super().__init__(actual, expected, **other_parameters) + self.rtol = max(self.rtol, rtol_override) + self.atol = max(self.atol, atol_override) + + def _process_inputs(self, actual, expected, *, id, allow_subclasses): + self._check_inputs_isinstance(actual, expected, cls=(torch.Tensor, np.ndarray)) + + actual, expected = (self._to_tensor(input) for input in (actual, expected)) + for tensor in (actual, expected): + self._check_supported(tensor, id=id) + return actual, expected + + +class TypedStoragePair(TensorLikePair): + """Pair for :class:`torch.storage.TypedStorage` inputs.""" + def __init__(self, actual, expected, *, rtol_override=0.0, atol_override=0.0, **other_parameters): + self._check_inputs_isinstance(actual, expected, cls=torch.storage.TypedStorage) + super().__init__(actual, expected, **other_parameters) + self.rtol = max(self.rtol, rtol_override) + self.atol = max(self.atol, atol_override) + + def _to_tensor(self, typed_storage): + return torch.tensor( + typed_storage._untyped_storage, + dtype={ + torch.quint8: torch.uint8, + torch.quint4x2: torch.uint8, + torch.quint2x4: torch.uint8, + torch.qint32: torch.int32, + torch.qint8: torch.int8 + }.get(typed_storage.dtype, typed_storage.dtype), + device=typed_storage.device, + ) + + +class UnittestPair(Pair): + """Fallback ABC pair that handles non-numeric inputs. + + To avoid recreating the mismatch messages of :meth:`unittest.TestCase.assertEqual`, this pair simply wraps it in + order to use it with the :class:`Pair` "framework" from :func:`are_equal`. + + Define the :attr:`UnittestPair.CLS` in a subclass to indicate which class(es) of the inputs the pair should support. + """ + CLS: type | tuple[type, ...] + TYPE_NAME: str | None = None + + def __init__(self, actual, expected, **other_parameters): + self._check_inputs_isinstance(actual, expected, cls=self.CLS) + super().__init__(actual, expected, **other_parameters) + + def compare(self): + test_case = unittest.TestCase() + + try: + return test_case.assertEqual(self.actual, self.expected) + except test_case.failureException as error: + msg = str(error) + + type_name = self.TYPE_NAME or (self.CLS if isinstance(self.CLS, type) else self.CLS[0]).__name__ + self._fail(AssertionError, f"{type_name.title()} comparison failed: {msg}") + + +class StringPair(UnittestPair): + CLS = (str, bytes) + TYPE_NAME = "string" + + +class SetPair(UnittestPair): + CLS = set + + +class TypePair(UnittestPair): + CLS = type + + +class ObjectPair(UnittestPair): + CLS = object + + +# This implements a variant of assertRaises/assertRaisesRegex where we first test +# if the exception is NotImplementedError, and if so just skip the test instead +# of failing it. +# +# This is implemented by inheriting from the (private) implementation of +# assertRaises from unittest.case, and slightly tweaking it for this new +# behavior. The year is 2021: this private class hierarchy hasn't changed since +# 2010, seems low risk to inherit from. +class AssertRaisesContextIgnoreNotImplementedError(unittest.case._AssertRaisesContext): + def __exit__(self, exc_type, exc_value, tb): + if exc_type is not None and issubclass(exc_type, NotImplementedError): + self.test_case.skipTest(f"not_implemented: {exc_value}") # type: ignore[attr-defined] + return super().__exit__(exc_type, exc_value, tb) + + +@contextmanager +def set_warn_always_context(new_val: bool): + old_val = torch.is_warn_always_enabled() + torch.set_warn_always(new_val) + try: + yield + finally: + torch.set_warn_always(old_val) + + +class NoTest: + # causes pytest to not recognize this class as a test + __test__ = False + + +class TestCase(expecttest.TestCase): + # NOTE: "precision" lets classes and generated tests set minimum + # atol values when comparing tensors. Used by @precisionOverride and @toleranceOverride, for + # example. + # NOTE: "rel_tol" lets classes and generated tests set minimum + # rtol values when comparing tensors. Used by @toleranceOverride, for example. + _precision: float = 0 + _rel_tol: float = 0 + + # Toggles whether to assert that `torch.get_default_dtype()` returns + # `torch.float` when `setUp` and `tearDown` are called. + _default_dtype_check_enabled: bool = False + + # Always use difflib to print diffs on multi line equality. + # Undocumented feature in unittest + _diffThreshold = sys.maxsize + maxDiff = None + + # checker to early terminate test suite if unrecoverable failure occurs. + def _should_stop_test_suite(self): + if torch.cuda.is_initialized(): + # CUDA device side error will cause subsequence test cases to fail. + # stop entire test suite if catches RuntimeError during torch.cuda.synchronize(). + try: + torch.cuda.synchronize() + except RuntimeError as rte: + print("TEST SUITE EARLY TERMINATION due to torch.cuda.synchronize() failure", file=sys.stderr) + print(str(rte), file=sys.stderr) + return True + return False + else: + return False + + @property + def precision(self) -> float: + return self._precision + + @precision.setter + def precision(self, prec: float) -> None: + self._precision = prec + + @property + def rel_tol(self) -> float: + return self._rel_tol + + @rel_tol.setter + def rel_tol(self, prec: float) -> None: + self._rel_tol = prec + + _do_cuda_memory_leak_check = False + _do_cuda_non_default_stream = False + + # When True, if a test case raises a NotImplementedError, instead of failing + # the test, skip it instead. + _ignore_not_implemented_error = False + + def __init__(self, method_name='runTest', methodName='runTest'): + # methodName is the correct naming in unittest and testslide uses keyword arguments. + # So we need to use both to 1) not break BC and, 2) support testslide. + if methodName != "runTest": + method_name = methodName + super().__init__(method_name) + + test_method = getattr(self, method_name, None) + if test_method is not None: + # Wraps the tested method if we should do CUDA memory check. + if TEST_CUDA_MEM_LEAK_CHECK: + self._do_cuda_memory_leak_check &= getattr(test_method, '_do_cuda_memory_leak_check', True) + # FIXME: figure out the flaky -1024 anti-leaks on windows. See #8044 + if self._do_cuda_memory_leak_check and not IS_WINDOWS: + self.wrap_with_cuda_policy(method_name, self.assertLeaksNoCudaTensors) + + # Wraps the tested method if we should enforce non default CUDA stream. + self._do_cuda_non_default_stream &= getattr(test_method, '_do_cuda_non_default_stream', True) + if self._do_cuda_non_default_stream and not IS_WINDOWS: + self.wrap_with_cuda_policy(method_name, self.enforceNonDefaultStream) + + if self._ignore_not_implemented_error: + self.wrap_with_policy(method_name, lambda: skip_exception_type(NotImplementedError)) + + if PRINT_REPRO_ON_FAILURE: + try: + def _get_rel_test_path(abs_test_path): + # Attempt to get relative path based on the "test" dir. + # In CI, the working dir is not guaranteed to be the base repo dir so + # we can't just compute relative path from that. + parts = Path(abs_test_path).parts + for i, part in enumerate(parts): + if part == "test": + base_dir = os.path.join(*parts[:i]) if i > 0 else '' + return os.path.relpath(abs_test_path, start=base_dir) + + # Can't determine containing dir; just return the test filename. + # The path isn't strictly correct but it's arguably better than nothing. + return os.path.split(abs_test_path)[1] + + abs_test_path = inspect.getfile(type(self)) + test_filename = _get_rel_test_path(abs_test_path) + class_name = type(self).__name__ + test_run_cmd = f"python {test_filename} {class_name}.{method_name}" + env_var_prefix = TestEnvironment.repro_env_var_prefix() + repro_parts = [env_var_prefix, test_run_cmd] + self.wrap_with_policy( + method_name, + lambda repro_parts=repro_parts: print_repro_on_failure(repro_parts)) + except Exception as e: + # Don't fail entirely if we can't get the test filename + log.info("could not print repro string", extra=str(e)) # type: ignore[arg-type] + + def assertLeaksNoCudaTensors(self, name=None): + name = self.id() if name is None else name + return CudaMemoryLeakCheck(self, name) + + def enforceNonDefaultStream(self): + return CudaNonDefaultStream() + + def _remove_ansi_escape(self, input): + # 7-bit C1 ANSI sequences + ansi_escape = re.compile(r''' + \x1B # ESC + (?: # 7-bit C1 Fe (except CSI) + [@-Z\\-_] + | # or [ for CSI, followed by a control sequence + \[ + [0-?]* # Parameter bytes + [ -/]* # Intermediate bytes + [@-~] # Final byte + ) + ''', re.VERBOSE) + return ansi_escape.sub('', input) + + def remove_comment_lines(self, input_string): + lines = input_string.split('\n') + filtered_lines = [line for line in lines if not line.strip().startswith('#')] + return '\n'.join(filtered_lines) + + def remove_empty_lines(self, input_string): + lines = input_string.split('\n') + filtered_lines = [line for line in lines if line.strip() != ''] + return '\n'.join(filtered_lines) + + # ignore comments will ignore lines that starts with # after being stripped + def assertExpectedInline(self, actual, expect, skip=0, ignore_comments=False, ignore_empty_lines=False): + actual = actual if isinstance(actual, str) else str(actual) + actual = self._remove_ansi_escape(actual) + expect = self._remove_ansi_escape(expect) + if ignore_comments: + actual = self.remove_comment_lines(actual) + expect = self.remove_comment_lines(expect) + + if ignore_empty_lines: + actual = self.remove_empty_lines(actual) + expect = self.remove_empty_lines(expect) + + return super().assertExpectedInline(actual if isinstance(actual, str) else str(actual), expect, skip + 1) + + # Munges exceptions that internally contain stack traces, using munge_exc + def assertExpectedInlineMunged( + self, exc_type, callable, expect, *, skip=0, suppress_suffix=True, post_munge=None, + ): + try: + callable() + except exc_type as e: + munged = munge_exc(e, suppress_suffix=suppress_suffix, skip=skip + 1) + if post_munge: + munged = post_munge(munged) + self.assertExpectedInline( + munged, expect, skip=skip + 1 + ) + return + self.fail(msg="Did not raise when expected to") + + def assertLogs(self, logger=None, level=None): + if logger is None: + logger = logging.getLogger("torch") + return super().assertLogs(logger, level) + + def assertNoLogs(self, logger=None, level=None): + if logger is None: + logger = logging.getLogger("torch") + return super().assertNoLogs(logger, level) + + def wrap_with_cuda_policy(self, method_name, policy): + test_method = getattr(self, method_name) + # the import below may initialize CUDA context, so we do it only if + # self._do_cuda_memory_leak_check or self._do_cuda_non_default_stream + # is True. + # TODO: sure looks like we unconditionally initialize the context here + # -- ezyang + from torch.testing._internal.common_cuda import TEST_CUDA + fullname = self.id().lower() # class_name.method_name + if TEST_CUDA and ('gpu' in fullname or 'cuda' in fullname): + setattr(self, method_name, self.wrap_method_with_policy(test_method, policy)) + + def wrap_with_policy(self, method_name, policy): + test_method = getattr(self, method_name) + setattr(self, method_name, self.wrap_method_with_policy(test_method, policy)) + + # A policy is a zero-argument function that returns a context manager. + # We don't take the context manager directly as it may be necessary to + # construct it once per test method + def wrap_method_with_policy(self, method, policy): + # Assumes that `method` is the tested function in `self`. + # NOTE: Python Exceptions (e.g., unittest.Skip) keeps objects in scope + # alive, so this cannot be done in setUp and tearDown because + # tearDown is run unconditionally no matter whether the test + # passes or not. For the same reason, we can't wrap the `method` + # call in try-finally and always do the check. + @wraps(method) + def wrapper(self, *args, **kwargs): + with policy(): + method(*args, **kwargs) + return types.MethodType(wrapper, self) + + def wrap_with_cuda_memory_check(self, method): + return self.wrap_method_with_policy(method, self.assertLeaksNoCudaTensors) + + def _dynamo_test_key(self): + return f"{self.__class__.__name__}.{self._testMethodName}" + + def compile_fn(self, fn, backend, nopython): + # Allows subclasses to control compilation + return torch._dynamo.optimize(backend, nopython=nopython)(fn) + + def _run_custom(self, result=None): + using_unittest = isinstance(result, unittest.TestResult) + + super_run = super().run + test_cls = super_run.__self__ # type: ignore[attr-defined] + + # Are we compiling? + compiled = TEST_WITH_TORCHDYNAMO or TEST_WITH_AOT_EAGER or TEST_WITH_TORCHINDUCTOR + # Is the class strict and compiling? + strict_default = False + should_reset_dynamo = False + + # We disable size_asserts for test_ops since some tests fail + # due to mismatch of strides returned from eager v.s. meta kernels + # Only some of the ops has this problem, but since tests in + # test_op.py are parametrized, it's hard to do this specifically + # for the affected ops. + # It's not a big deal since these problems are captured by + # test_torchinductor_opinfo.py as well. + should_disable_size_asserts = False + if compiled: + try: + path = inspect.getfile(type(test_cls)) + full_path = os.path.abspath(path) + match = re.match(r".*/test/(.*).py", full_path) + if match is not None: + filename = match.group(1) + if TEST_WITH_TORCHINDUCTOR: + from .dynamo_test_failures import FIXME_inductor_non_strict + strict_default = filename not in FIXME_inductor_non_strict + should_reset_dynamo = True + + if filename == "test_ops": + should_disable_size_asserts = True + else: + strict_default = True + # inspect.getfile can fail with these + except (OSError, TypeError): + pass + if "STRICT_DEFAULT" in os.environ: + if os.environ["STRICT_DEFAULT"] == "1": + strict_default = True + + strict_mode = False + if compiled: + test_method = getattr(self, self._testMethodName) + if hasattr(test_method, "dynamo_strict"): + strict_mode = test_method.dynamo_strict + elif hasattr(test_cls, "dynamo_strict"): + strict_mode = test_cls.dynamo_strict + else: + strict_mode = strict_default + nopython = getattr(test_cls, "dynamo_strict_nopython", False) and compiled + + if strict_mode or should_reset_dynamo: + torch._dynamo.reset() + + torch.compiler.set_stance("default") + + # TODO: Remove this; this is grandfathered in because we suppressed errors + # on test suite previously + # When strict mode is False, suppress_errors is True + if compiled: + suppress_errors = not strict_mode + else: + suppress_errors = torch._dynamo.config.suppress_errors + + maybe_disable_size_asserts = ( + torch._inductor.config.patch(size_asserts=False) + if should_disable_size_asserts + else contextlib.nullcontext() + ) + + with unittest.mock.patch("torch._dynamo.config.suppress_errors", suppress_errors), maybe_disable_size_asserts: + if TEST_WITH_AOT_EAGER: + super_run = self.compile_fn(super_run, "aot_eager_decomp_partition", nopython) + elif TEST_WITH_TORCHDYNAMO or TEST_WITH_TORCHINDUCTOR: + if TEST_WITH_TORCHINDUCTOR: + super_run = self.compile_fn(super_run, "inductor", nopython) + else: + # Assume eager-generated GraphModules will not error out. + # If we do, this is probably a Dynamo bug! + super_run = self.compile_fn(super_run, "eager_noexcept", nopython) + + key = self._dynamo_test_key() + + def expect_failure(f, file_name): + @wraps(f) + def wrapper(*args, **kwargs): + try: + f(*args, **kwargs) + except BaseException as e: # noqa: B036 + self.skipTest(e) + raise RuntimeError(f"Unexpected success, please remove `{file_name}`") + return wrapper + + if TEST_WITH_TORCHINDUCTOR: + subdir = "test/inductor_expected_failures" + from .dynamo_test_failures import inductor_expected_failures as expected_failures + else: + subdir = "test/dynamo_expected_failures" + from .dynamo_test_failures import dynamo_expected_failures as expected_failures + + if key in expected_failures: + method = getattr(self, self._testMethodName) + file_name = os.path.join(subdir, key) + setattr(self, self._testMethodName, expect_failure(method, file_name)) + + def ignore_failure(f, file_name): + @wraps(f) + def wrapper(*args, **kwargs): + try: + f(*args, **kwargs) + except BaseException as e: # noqa: B036 + self.skipTest(e) + method = getattr(self, self._testMethodName) + if getattr(method, "__unittest_expecting_failure__", False): + self.skipTest("unexpected success") + else: + self.skipTest(f"This test passed, maybe we can remove `{file_name}`") + return wrapper + + if TEST_WITH_TORCHINDUCTOR: + subdir = "test/inductor_skips" + from .dynamo_test_failures import inductor_skips as skips + else: + subdir = "test/dynamo_skips" + from .dynamo_test_failures import dynamo_skips as skips + + if key in skips: + method = getattr(self, self._testMethodName) + file_name = os.path.join(subdir, key) + setattr(self, self._testMethodName, ignore_failure(method, file_name)) + + from .dynamo_test_failures import compiled_autograd_skips + if torch._dynamo.config.compiled_autograd and key in compiled_autograd_skips: + # Still run the test, but with compiled autograd disabled + super_run = runWithoutCompiledAutograd()(super_run) + + super_run(result=result) + + if strict_mode or should_reset_dynamo: + torch._dynamo.reset() + elif torch._dynamo.config.compiled_autograd: + torch._dynamo.compiled_autograd.reset() + + # Early terminate test if necessary. If using pytest, use the -x flag instead + if using_unittest and self._should_stop_test_suite(): + if result.wasSuccessful(): + case = TestCase() + if TEST_SAVE_XML is not None: + # This is a big hacky, XMLRunner modifies expected type from TestCase to TestInfo + # Create dummy TestInfo to record results correctly + from xmlrunner.result import _TestInfo # type: ignore[import] + case = _TestInfo(result, case) + case.output = _TestInfo.ERROR # type: ignore[attr-defined] + case.elapsed_time = 0.0 # type: ignore[attr-defined] + case.test_description = "TestSuiteEarlyFailure" # type: ignore[attr-defined] + # This shouldn't really happen, but if does add fake failure + # For more details see https://github.com/pytorch/pytorch/issues/71973 + result.failures.append((case, "TestSuite execution was aborted early")) + if result.wasSuccessful() is not False: + raise AssertionError("expected result.wasSuccessful() to be False after adding failure") + result.stop() + + + def run(self, result=None): + with contextlib.ExitStack() as stack: + if TEST_WITH_CROSSREF: + stack.enter_context(CrossRefMode()) + self._run_custom( + result=result, + ) + + def setUp(self): + check_if_enable(self) + set_rng_seed() + + # Save global check sparse tensor invariants state that can be + # restored from tearDown: + self._check_invariants = torch.sparse.check_sparse_tensor_invariants.is_enabled() + + # Enable invariant checks for all sparse tensors constructions + # including the unsafe ones. If this is not desired for some + # test case, use check_invariants=False optional argument to + # sparse tensor constructors or + # @torch.sparse.check_sparse_tensor_invariants(False) + # decorator to disable the invariant checks. + torch.sparse.check_sparse_tensor_invariants.enable() + + if self._default_dtype_check_enabled: + if torch.get_default_dtype() != torch.float: + raise AssertionError( + f"expected default dtype to be torch.float, got {torch.get_default_dtype()}" + ) + + # attempt to reset some global state at the end of the test + self._prev_grad_state = torch.is_grad_enabled() + + def tearDown(self): + # There exists test cases that override TestCase.setUp + # definition, so we cannot assume that _check_invariants + # attribute is defined in general. + if hasattr(self, '_check_invariants'): + # Restore the global check sparse tensor invariants state + if self._check_invariants: + torch.sparse.check_sparse_tensor_invariants.enable() + else: + torch.sparse.check_sparse_tensor_invariants.disable() + + if self._default_dtype_check_enabled: + if torch.get_default_dtype() != torch.float: + raise AssertionError( + f"expected default dtype to be torch.float, got {torch.get_default_dtype()}" + ) + + # attribute may not be defined, per above + if hasattr(self, '_prev_grad_state'): + torch.set_grad_enabled(self._prev_grad_state) + + @staticmethod + def _make_crow_indices(n_rows, n_cols, nnz, + *, device, dtype, random=True): + """Return crow_indices of a CSR tensor with size (n_rows, n_cols) and + the number of specified elements nnz. + + If random is True, the column counts of rows are in random + order. Otherwise, the column counts of rows are defined by the + used sampling method. + + Sampling method + --------------- + + The used sampling method was introduced in + https://pearu.github.io/csr_sampling.html, and here we give + only an overall description of the method. + + Notice that crow_indices can be defined as cumsum(counts) + where counts is a sequence of non-negative integers satisfying + the following conditions: + + len(counts) == n_rows + 1 + counts.max() <= n_cols + + while counts[i + 1] is interpreted as the number of specified + elements in the i-th row. + + The used sampling method aims at increasing the diversity of + CSR samples, that is, a CSR sample should contain (i) rows + that are all filled, (ii) rows with no elements at all, and + (iii) rows that are partially filled. At the same time and for + the given total number of specified elements (nnz), there + should be minimal preference to rows with a given number of + elements. To achieve this, the sampling method is built-up on + using a sawteeth model for counts. In the simplest case, we + would have + + counts = arange(n_rows + 1) % (n_cols + 1) + + that has equal number of all possible column counts per row. + This formula can be used only for specific input values of + n_rows, n_cols, and nnz. To generalize this model to any + combinations of inputs, the counts model above is extended + with an incomplete sawtooth, and the right and lower + rectangular parts that will guarantee that + + counts.sum() == nnz + + for any combination of n_rows, n_cols, and nnz. Basically, + we'll find a maximal window in (n_rows + 1, n_cols + 1)-grid + that is able to hold a sequence of sawteeth and so-called + final correction, while the external part of the window is + filled with counts to meet the nnz constraint exactly. + """ + if not (0 <= nnz <= n_rows * n_cols): + raise AssertionError( + f"nnz out of bounds: expected 0 <= nnz <= n_rows * n_cols, got nnz={nnz}, n_rows={n_rows}, n_cols={n_cols}" + ) + + def sawteeth(n, m): + # return the total number of counts in the sequence of + # sawteeth where n and m define a window in (n_rows+1, + # n_cols+1) rectangle where the sequence of sawteeth + # perfectly fit. + M = (n_cols - m) * (n_cols - m + 1) // 2 + K = (n_rows - n) % (n_cols - m + 1) + return M * ((n_rows - n) // (n_cols - m + 1)) + K * (K - 1) // 2 + + # Different from the original method description, here counts + # has leading 0 required by crow_indices: + counts = torch.zeros(n_rows + 1, dtype=dtype, device=torch.device('cpu')) + + n = m = 0 + N = sawteeth(n, m) + if N and nnz >= max(N, n_cols): + # determine the width of the sawteeth window. We use bisection to solve + # N(n, 0) == 0 or nnz - n * n_cols < max(N(n, 0), n_cols) + # for n + n_left = n + n_right = n_rows - 1 + N_right = sawteeth(n_right, m) + while n_right - n_left > 1: + n_middle = (n_left + n_right) // 2 + N_middle = sawteeth(n_middle, m) + if N_middle == 0 or nnz - n_middle * n_cols < max(N_middle, n_cols): + n_right, N_right = n_middle, N_middle + else: + n_left = n_middle + n, N = n_right, N_right + # fill the right rectangle with counts: + if not n: + raise AssertionError("n must be non-zero") + counts[-n:].fill_(n_cols) + + if N and nnz - n * n_cols >= max(N, n_rows - n): + # determine the height of the sawteeth window. We use bisection to solve + # N(n, m) == 0 or nnz - n * n_cols - m * (n_rows - n) < max(N(n, m), n_rows - n) + # for m. + m_left = m + m_right = n_cols - 1 + N_right = sawteeth(n, m_right) + while m_right - m_left > 1: + m_middle = (m_left + m_right) // 2 + N_middle = sawteeth(n, m_middle) + if N_middle == 0 or nnz - n * n_cols - m_middle * (n_rows - n) < max(N_middle, n_rows - n): + m_right, N_right = m_middle, N_middle + else: + m_left = m_middle + m, N = m_right, N_right + # fill the bottom rectangle with counts: + if not m: + raise AssertionError("m must be non-zero") + counts[1:n_rows - n + 1].fill_(m) + + if N: + # fill the sawteeth window with counts + q, r = divmod(nnz - n * n_cols - m * (n_rows - n), + (n_cols - m) * (n_cols - m + 1) // 2) + p = 1 + q * (n_cols - m + 1) + k = math.isqrt(2 * r) + if k * (k + 1) > 2 * r: + k -= 1 + corr = r - k * (k + 1) // 2 + if (p > 1) and (m > 0): + raise AssertionError( + f"full sawteeth are never on top of a bottom rectangle: p={p}, m={m}" + ) + # sequence of full sawteeth: + counts[1:p] = torch.arange(p - 1, dtype=dtype, device=counts.device) % (n_cols - m + 1) + # incomplete sawtooth: + counts[p:p + k + 1] += torch.arange(k + 1, dtype=dtype, device=counts.device) + else: + # given input does not support sawteeth + p = 1 + corr = nnz - n * n_cols - m * (n_rows - n) + + # correction that will guarantee counts.sum() == nnz: + counts[p] += corr + + if random: + # randomize crow_indices by shuffling the sawteeth + # sequence: + perm = torch.randperm(n_rows, device=counts.device) + counts[1:] = counts[1:][perm] + + # compute crow_indices: + crow_indices = counts + crow_indices.cumsum_(dim=0) + return crow_indices.to(device=device) + + def genSparseCompressedTensor(self, size, nnz, *, layout, device, dtype, index_dtype, blocksize=(), dense_dims=0): + from operator import mul + from functools import reduce + sparse_dim = 2 + if not (all(size[d] > 0 for d in range(len(size))) or nnz == 0): + raise AssertionError(f"invalid arguments: size={size}, nnz={nnz}") + if len(size) < sparse_dim: + raise AssertionError(f"expected len(size) >= {sparse_dim}, got {len(size)}") + if blocksize: + if len(blocksize) != 2: + raise AssertionError(f"expected len(blocksize) == 2, got size={size}, blocksize={blocksize}") + if size[-2 - dense_dims] % blocksize[0] != 0: + raise AssertionError( + f"size[-2 - dense_dims] must be divisible by blocksize[0]: size={size}, blocksize={blocksize}" + ) + if size[-1 - dense_dims] % blocksize[1] != 0: + raise AssertionError( + f"size[-1 - dense_dims] must be divisible by blocksize[1]: size={size}, blocksize={blocksize}" + ) + blocksize0, blocksize1 = blocksize + else: + blocksize0 = blocksize1 = 1 + + size = tuple(size) + dense_size = size[(len(size) - dense_dims):] + + def random_sparse_compressed(n_compressed_dims, n_plain_dims, nnz): + compressed_indices = self._make_crow_indices(n_compressed_dims, n_plain_dims, nnz, device=device, dtype=index_dtype) + plain_indices = torch.zeros(nnz, dtype=index_dtype, device=device) + for i in range(n_compressed_dims): + count = compressed_indices[i + 1] - compressed_indices[i] + plain_indices[compressed_indices[i]:compressed_indices[i + 1]], _ = torch.sort( + torch.randperm(n_plain_dims, dtype=index_dtype, device=device)[:count]) + low = -1 if dtype != torch.uint8 else 0 + high = 1 if dtype != torch.uint8 else 2 + values = make_tensor((nnz,) + blocksize + dense_size, device=device, dtype=dtype, low=low, high=high) + return values, compressed_indices, plain_indices + + batch_shape = size[:-2 - dense_dims] + n_batch = reduce(mul, batch_shape, 1) + + if layout in {torch.sparse_csr, torch.sparse_bsr}: + n_compressed_dims, n_plain_dims = size[-2 - dense_dims] // blocksize0, size[-1 - dense_dims] // blocksize1 + else: + n_compressed_dims, n_plain_dims = size[-1 - dense_dims] // blocksize1, size[-2 - dense_dims] // blocksize0 + blocknnz = nnz // (blocksize0 * blocksize1) + sparse_tensors = [random_sparse_compressed(n_compressed_dims, n_plain_dims, blocknnz) for _ in range(n_batch)] + sparse_tensors_it = map(list, zip(*sparse_tensors, strict=True)) + + values = torch.stack(next(sparse_tensors_it)).reshape(*batch_shape, blocknnz, *blocksize, *dense_size) + compressed_indices = torch.stack(next(sparse_tensors_it)).reshape(*batch_shape, -1) + plain_indices = torch.stack(next(sparse_tensors_it)).reshape(*batch_shape, -1) + return torch.sparse_compressed_tensor(compressed_indices, plain_indices, + values, size=size, dtype=dtype, layout=layout, device=device) + + def genSparseCSRTensor(self, size, nnz, *, device, dtype, index_dtype, dense_dims=0): + return self.genSparseCompressedTensor(size, nnz, layout=torch.sparse_csr, device=device, + dtype=dtype, index_dtype=index_dtype, blocksize=(), dense_dims=dense_dims) + + def genSparseCSCTensor(self, size, nnz, *, device, dtype, index_dtype, dense_dims=0): + return self.genSparseCompressedTensor(size, nnz, layout=torch.sparse_csc, device=device, + dtype=dtype, index_dtype=index_dtype, blocksize=(), dense_dims=0) + + def genSparseBSRTensor(self, size, blocksize, nnz, *, device, dtype, index_dtype, dense_dims=0): + if len(blocksize) != 2: + raise AssertionError(f"expected len(blocksize) == 2, got {len(blocksize)}") + return self.genSparseCompressedTensor(size, nnz, layout=torch.sparse_bsr, device=device, + dtype=dtype, index_dtype=index_dtype, blocksize=blocksize, dense_dims=dense_dims) + + def genSparseBSCTensor(self, size, blocksize, nnz, *, device, dtype, index_dtype, dense_dims=0): + if len(blocksize) != 2: + raise AssertionError(f"expected len(blocksize) == 2, got {len(blocksize)}") + return self.genSparseCompressedTensor(size, nnz, layout=torch.sparse_bsc, device=device, + dtype=dtype, index_dtype=index_dtype, blocksize=blocksize, dense_dims=dense_dims) + + def genSparseTensor(self, size, sparse_dim, nnz, is_uncoalesced, device, dtype): + # Assert not given impossible combination, where the sparse dims have + # empty numel, but nnz > 0 makes the indices containing values. + if not (all(size[d] > 0 for d in range(sparse_dim)) or nnz == 0): + raise AssertionError(f"invalid arguments: size={size}, sparse_dim={sparse_dim}, nnz={nnz}") + + v_size = [nnz] + list(size[sparse_dim:]) + v = make_tensor(v_size, device=device, dtype=dtype, low=-1, high=1) + i = torch.rand(sparse_dim, nnz, device=device) + i.mul_(torch.tensor(size[:sparse_dim]).unsqueeze(1).to(i)) + i = i.to(torch.long) + if is_uncoalesced: + i1 = i[:, :(nnz // 2), ...] + i2 = i[:, :((nnz + 1) // 2), ...] + i = torch.cat([i1, i2], 1) + x = torch.sparse_coo_tensor(i, v, torch.Size(size), dtype=dtype, device=device) + + if not is_uncoalesced: + x = x.coalesce() + else: + # FIXME: `x` is a sparse view of `v`. Currently rebase_history for + # sparse views is not implemented, so this workaround is + # needed for inplace operations done on `x`, e.g., copy_(). + # Remove after implementing something equivalent to CopySlice + # for sparse views. + # NOTE: We do clone() after detach() here because we need to be able to change size/storage of x afterwards + x = x.detach().clone()._coalesced_(False) + return x, x._indices().clone(), x._values().clone() + + def generate_simple_inputs(self, layout, + device=None, + dtype=None, + index_dtype=None, + pin_memory=None, + members_pin_memory=None, + enable_batch=True, + enable_hybrid=True, + enable_zero_sized=True, + enable_non_contiguous_indices=True, + enable_non_contiguous_values=True, + enable_batch_variable_nse=False, + output_tensor=True, + patterns=None): + """Generator of simple inputs for tensor constructors of the given layout. + + The generated tensor inputs have the following properties: + + - tensor shapes are minimal but not trivial + - tensor values are sorted sequences for COO and CSR formats, e.g. [1, 2, 3, 4] + - the generated tensors represent the same mathematical tensor for all layouts + - the generated tensors include regular, zero-sized, and optionally, batched or/and hybrid tensors. + - the generated tensors include contiguous or non-contiguous tensors both in indices and values + + If output_tensor is True, yield tensors with the given + layout. Otherwise, yield inputs to the corresponding tensor + constructors: + + - sparse compressed input is defined as + (compressed_indices, plain_indices, values), dict(size=expected_size_from_shape_inference, device=device, dtype=dtype, + pin_memory=pin_memory) + + - sparse COO input is defined as + (indices, values), dict(size=expected_size_from_shape_inference, device=device, dtype=dtype, pin_memory=pin_memory) + + - strided input is defined as + (values,), dict(device=device, dtype=dtype) + """ + if index_dtype is None: + index_dtype = torch.int64 + + is_compressed_sparse_layout = layout in {torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc} + + if output_tensor: + for args, kwargs in self.generate_simple_inputs(layout, device=device, dtype=dtype, index_dtype=index_dtype, + pin_memory=pin_memory, + enable_batch=enable_batch, enable_hybrid=enable_hybrid, + enable_zero_sized=enable_zero_sized, + enable_non_contiguous_indices=enable_non_contiguous_indices, + enable_non_contiguous_values=enable_non_contiguous_values, + enable_batch_variable_nse=enable_batch_variable_nse, + output_tensor=False): + if members_pin_memory: + args = tuple(a.pin_memory() for a in args) + if layout is torch.strided: + if len(args) != 1: + raise AssertionError(f"expected len(args) == 1 for strided layout, got {len(args)}") + size = kwargs.pop('size', None) # to ensure that a zero-sized tensor has the desired shape + if size is None: + raise AssertionError("size must not be None for strided layout") + if pin_memory: + yield args[0].reshape(size).pin_memory() + else: + yield args[0].reshape(size) + elif layout is torch.sparse_coo: + yield torch.sparse_coo_tensor(*args, **kwargs) + elif is_compressed_sparse_layout: + kwargs.update(layout=layout) + yield torch.sparse_compressed_tensor(*args, **kwargs) + else: + raise AssertionError(f"unreachable: unexpected layout {layout}") + return + + def get_blockpattern(pattern, blocksize): + basesize = pattern.shape + if basesize[0] % blocksize[0] != 0: + raise AssertionError( + f"basesize[0] must be divisible by blocksize[0]: basesize={basesize}, blocksize={blocksize}" + ) + if basesize[1] % blocksize[1] != 0: + raise AssertionError( + f"basesize[1] must be divisible by blocksize[1]: basesize={basesize}, blocksize={blocksize}" + ) + blockpattern = pattern.reshape(-1, + blocksize[0], + basesize[1] // blocksize[1], + blocksize[1]).transpose(-3, -2).any(-1).any(-1) + block_ids = torch.arange(1, blockpattern.numel() + 1).reshape(blockpattern.shape) + return (blockpattern != 0) * block_ids + + def get_sparse_data(pattern): + basesize = pattern.shape + if len(basesize) != 2: + raise AssertionError(f"pattern is expected to be a matrix, got shape {basesize}") + + # We cannot use `torch.sparse_xyz_tensor(pattern)` to + # compute the sparse layout indices and values because + # generate_simple_inputs is used to generate the inputs to + # test `torch.sparse_xyz_tensor` factory functions, so + # we'll compute the indices and values independently of + # the factory functions. + + indices = torch.where(pattern != 0) + coo_indices = torch.stack(indices) + crow_indices = torch.zeros(basesize[0] + 1, dtype=torch.int64) + crow_indices[1:] = torch.cumsum(coo_indices[0].bincount(minlength=basesize[0]), 0) + col_indices = coo_indices[1] + strided_values = torch.zeros(basesize, dtype=torch.int64) + + # the property of `values == range(1, 1+nnz)` is used in + # get_sparse_data_with_block to relate BSR and BSC values, + # so, don't change the following line: + values = torch.arange(1, 1 + len(indices[0]), dtype=torch.int64) + strided_values[indices] = values + + indices_T = torch.where(pattern.transpose(0, 1) != 0) + coo_indices_T = torch.stack(indices_T) + ccol_indices = torch.zeros(basesize[1] + 1, dtype=torch.int64) + ccol_indices[1:] = torch.cumsum(coo_indices_T[0].bincount(minlength=basesize[1]), 0) + row_indices = coo_indices_T[1] + csc_values = strided_values.transpose(0, 1)[indices_T] + + return {torch.sparse_coo: (coo_indices, values), + torch.sparse_csr: (crow_indices, col_indices, values), + torch.sparse_csc: (ccol_indices, row_indices, csc_values), + torch.strided: (strided_values,)} + + def get_sparse_data_with_block(pattern, blocksize): + nonblock_data = get_sparse_data(pattern) + blockpattern = get_blockpattern(pattern, blocksize) + block_data = get_sparse_data(blockpattern) + + strided_values = nonblock_data[torch.strided][0] + block_indices = block_data[torch.sparse_coo][0] + bsr_values = torch.stack([strided_values[bi * blocksize[0]:(bi + 1) * blocksize[0], + bj * blocksize[1]:(bj + 1) * blocksize[1]] + for bi, bj in block_indices.transpose(0, 1)]) + + # here we use the property `values == range(1, 1+nnz)` and + # `values` relation to `csc_values` (see get_sparse_data) + # to get BSC blocks via reordering the BSR blocks: + bsc_values = bsr_values[block_data[torch.sparse_csc][2] - 1] + + return {torch.sparse_bsr: (*block_data[torch.sparse_csr][:2], bsr_values), + torch.sparse_bsc: (*block_data[torch.sparse_csc][:2], bsc_values), + **nonblock_data} + + def get_batch_sparse_data(pattern, blocksize): + size = pattern.shape + if len(size) <= 2: # non-batch + return get_sparse_data_with_block(pattern, blocksize) + + # batch data is created recursively: + batch_data = {} # type: ignore[var-annotated] + for i, item in enumerate(pattern): + for layout, d in get_batch_sparse_data(item, blocksize).items(): + target = batch_data.get(layout) + if layout is torch.sparse_coo: + # a "batch COO" means a COO with the leading + # sparse dimensions interpreted as batch + # dimensions + ext_coo_indices1 = torch.cat((torch.full((1, len(d[1])), i, dtype=torch.int64), d[0])) + if target is None: + target = batch_data[layout] = (ext_coo_indices1, d[1]) + else: + target[0].set_(torch.cat((target[0], ext_coo_indices1), 1)) # type: ignore[call-overload] + target[1].set_(torch.cat((target[1], d[1]))) + else: + if target is None: + target = batch_data[layout] = tuple(d[j].unsqueeze(0) for j in range(len(d))) + else: + for j in range(len(d)): + target[j].set_(torch.cat((target[j], d[j].unsqueeze(0)))) # type: ignore[call-overload] + return batch_data + + def generate_values(base, densesize): + """Generates a tensor of shape densesize with values equal to + + base + i_1 * 10^0 + ... + i_d * 10^{d - 1} + + at indices i_1, ..., i_d (with 0 <= i_j < densesize[j] for any 1 <= j <= + len(densesize)) + + This mapping produces unique values as long as + densesize[i] < 10 for all i in range(len(densesize)). + """ + + if not densesize: + return base + if not isinstance(base, int) and base.ndim > 0: + return torch.stack([generate_values(b, densesize) for b in base]) + if base == 0: + return torch.zeros(densesize, dtype=torch.int64) + r = torch.arange(densesize[0], dtype=torch.int64) + for i, d in enumerate(densesize[1:]): + y = torch.arange(d, dtype=torch.int64) * (10 ** (i + 1)) + r = r[..., None] + y[None, ...] + r.add_(base) + return r + + if patterns is None: + # A pattern is a 3-tuple with the following items: + # + # - a list of integers with the depth of two or more. The + # integers define the sparsity patterns of the generated + # inputs: zero values correspond to unspecified + # elements/blocks, and non-zero values to the specified + # elements. + # + # For debugging convenience, the elements with the same + # value typically belong to the same block. However, it + # is not a hard requirement: as long as the shape of a + # pattern divides with block sizes, the pattern will be + # a valid one. + # + # If the depth of the list is larger than two, inputs + # with batch dimensions will be generated. + # + # - a list of 2-tuples of block sizes, used to generate + # BSR/BSC tensors with various block size parameters + # + # - a list of tuples of dense dimensions, used to generate + # hybrid tensors with various dense dimensions + # + patterns = [ + # a simple 3 x 2 tensor: non-hybrid, hybrid with 1 and 2 dense dimensions + ([[1, 2, 0], + [1, 0, 3]], [(2, 1), (1, 3)], [(), (2,), (4, 5)]), + # 2 x 3 batch of 3 x 2 tensors: non-hybrid and hybrid with 2 dense dimensions + ([[[[1, 2, 0], + [1, 0, 3]], + [[1, 2, 3], + [1, 0, 0]], + [[1, 0, 0], + [1, 2, 3]]], + [[[0, 2, 0], + [1, 2, 3]], + [[1, 0, 3], + [1, 2, 0]], + [[1, 2, 3], + [0, 2, 0]]]], [(2, 1), (2, 3)], [(), (2,)]), + # tensor with non-trivial blocksize + ([[0, 1, 0, 2, 0, 2], + [0, 1, 0, 0, 2, 0], + [3, 3, 3, 0, 0, 0], + [0, 0, 0, 0, 0, 0], + [0, 5, 0, 6, 6, 6], + [5, 0, 5, 6, 6, 6], + [0, 0, 0, 0, 8, 8], + [7, 7, 7, 0, 8, 8]], [(2, 3)], [(), (4, 5)]), + # batch tensor with variable NSE + # Requires https://github.com/pytorch/pytorch/pull/84843 or similar. + ([[[1, 2], + [3, 4]], + [[1, 0], + [0, 0]]], [(1, 1)], ([()] if enable_batch_variable_nse else []))] + + def non_contiguous_copy(t, dim=-1, offset=0): + # return a copy of t that is non-contiguous along the + # given dimension and with the given storage offset + self.assertTrue(t.is_contiguous()) + if dim < 0: + dim = dim + t.ndim + if not (dim >= 0 and dim < t.ndim): + raise AssertionError(f"dim out of range: expected 0 <= dim < {t.ndim}, got dim={dim}") + step = max(2, offset + 1) + tmp = torch.zeros((*t.shape[:dim], t.shape[dim] * step, *t.shape[dim + 1:]), dtype=t.dtype, device=t.device) + dim_slices = (*((slice(None),) * dim), slice(offset, None, step)) + r = tmp[dim_slices].copy_(t) + self.assertFalse(r.is_contiguous()) + self.assertEqual(t, r) + return r + + # the main loop of the method: + for pattern, blocksizes, densesizes in patterns: + if not enable_hybrid: + densesizes = [s for s in densesizes if not s] + if not (densesizes and blocksizes): + continue + pattern = torch.tensor(pattern, dtype=torch.int64) + if not enable_batch and pattern.ndim > 2: + continue + for blocksize in blocksizes: + data = get_batch_sparse_data(pattern, blocksize)[layout] + for densesize in densesizes: + indices = [a.to(device=device, dtype=index_dtype) for a in data[:-1]] + values = generate_values(data[-1], densesize).to(device=device, dtype=dtype) + kwargs = dict(device=device, dtype=dtype, size=pattern.shape + densesize) + if pin_memory is not None: + kwargs.update(pin_memory=pin_memory) + + yield (*indices, values), kwargs.copy() + if enable_non_contiguous_indices and pattern.ndim > 2: + # sparse compressed indices can be sliced only along batch dimensions + for (dim, offset) in {(0, 1), (-2, 0)}: + indices_copy = [non_contiguous_copy(a, dim=dim, offset=offset) for a in indices] + yield (*indices_copy, values), kwargs.copy() + + if enable_non_contiguous_values: + values_copy = non_contiguous_copy(values, dim=-1, offset=1) + yield (*indices_copy, values_copy), kwargs.copy() + + if enable_non_contiguous_values: + values_copy = non_contiguous_copy(values, dim=-1, offset=1) + yield (*indices, values_copy), kwargs.copy() + + # zero-sized tensor inputs, non-batch, non-hybrid/hybrid + if enable_zero_sized: + for basesize, blocksizes, densesizes in [ + ((2, 0), [(1, 2)], [(), (2,), (2, 3)] if enable_hybrid else [()]), + ((0, 2), [(1, 2), (2, 1), (3, 2)], [()]), + ((0, 0), [(1, 2)], [()]), + ]: + for blocksize in blocksizes: + for densesize in densesizes: # type: ignore[attr-defined] + if layout == torch.strided: + indices = () # type: ignore[assignment] + values = torch.empty((basesize + densesize), device=device, dtype=dtype) + elif layout == torch.sparse_coo: + indices = (torch.empty(len(basesize), 0, device=device, dtype=index_dtype),) # type: ignore[assignment] + values = torch.empty((0, *densesize), device=device, dtype=dtype) + elif layout == torch.sparse_csr: + crow_indices = torch.tensor([0] * (basesize[0] + 1), device=device, dtype=index_dtype) + col_indices = torch.empty(0, device=device, dtype=index_dtype) + indices = (crow_indices, col_indices) # type: ignore[assignment] + values = torch.empty((0, *densesize), device=device, dtype=dtype) + elif layout == torch.sparse_csc: + ccol_indices = torch.tensor([0] * (basesize[1] + 1), device=device, dtype=index_dtype) + row_indices = torch.empty(0, device=device, dtype=index_dtype) + indices = (ccol_indices, row_indices) # type: ignore[assignment] + values = torch.empty((0, *densesize), device=device, dtype=dtype) + elif layout == torch.sparse_bsr: + crow_indices = torch.tensor([0] * (basesize[0] // blocksize[0] + 1), device=device, dtype=index_dtype) + col_indices = torch.empty(0, device=device, dtype=index_dtype) + indices = (crow_indices, col_indices) # type: ignore[assignment] + values = torch.empty((0, *blocksize, *densesize), device=device, dtype=dtype) + elif layout == torch.sparse_bsc: + ccol_indices = torch.tensor([0] * (basesize[1] // blocksize[1] + 1), device=device, dtype=index_dtype) + row_indices = torch.empty(0, device=device, dtype=index_dtype) + indices = (ccol_indices, row_indices) # type: ignore[assignment] + values = torch.empty((0, *blocksize, *densesize), device=device, dtype=dtype) + else: + raise AssertionError(f"unreachable: unexpected layout {layout}") + kwargs = dict(device=device, dtype=dtype, size=basesize + densesize) + if pin_memory is not None: + kwargs.update(pin_memory=pin_memory) + yield (*indices, values), kwargs + + def safeToDense(self, t): + # coalesce is only implemented for COO + if t.layout == torch.sparse_coo: + t = t.coalesce() + return t.to_dense() + + # Compares a torch function with a reference function for a given sample input (object of SampleInput) + # Note: only values are compared, type comparison is not done here + def compare_with_reference(self, torch_fn, ref_fn, sample_input, **kwargs): + numpy_sample = sample_input.numpy() + n_inp, n_args, n_kwargs = numpy_sample.input, numpy_sample.args, numpy_sample.kwargs + t_inp, t_args, t_kwargs = sample_input.input, sample_input.args, sample_input.kwargs + + actual = torch_fn(t_inp, *t_args, **t_kwargs) + expected = ref_fn(n_inp, *n_args, **n_kwargs) + + self.assertEqual(actual, expected, exact_device=False, **kwargs) + + # Compares the given Torch and NumPy functions on the given tensor-like object. + # NOTE: both torch_fn and np_fn should be functions that take a single + # tensor (array). If the torch and/or NumPy function require additional + # arguments then wrap the function in a lambda or pass a partial function. + # TODO: add args/kwargs for passing to assertEqual (e.g. rtol, atol) + def compare_with_numpy(self, torch_fn, np_fn, tensor_like, + device=None, dtype=None, **kwargs): + if not TEST_NUMPY: + raise AssertionError("TEST_NUMPY must be True to use compare_with_numpy") + + if isinstance(tensor_like, torch.Tensor): + if device is not None: + raise AssertionError("device must be None when tensor_like is a Tensor") + if dtype is not None: + raise AssertionError("dtype must be None when tensor_like is a Tensor") + t_cpu = tensor_like.detach().cpu() + if t_cpu.dtype is torch.bfloat16: + t_cpu = t_cpu.float() + a = t_cpu.numpy() + t = tensor_like + else: + d = copy.copy(torch_to_numpy_dtype_dict) + d[torch.bfloat16] = np.float32 + a = np.array(tensor_like, dtype=d[dtype]) + t = torch.tensor(tensor_like, device=device, dtype=dtype) + + np_result = np_fn(a) + torch_result = torch_fn(t).cpu() + + # Converts arrays to tensors + if isinstance(np_result, np.ndarray): + try: + np_result = torch.from_numpy(np_result) + except Exception: + # NOTE: copying an array before conversion is necessary when, + # for example, the array has negative strides. + np_result = torch.from_numpy(np_result.copy()) + if t.dtype is torch.bfloat16 and torch_result.dtype is torch.bfloat16 and np_result.dtype is torch.float: + torch_result = torch_result.to(torch.float) + + self.assertEqual(np_result, torch_result, **kwargs) + + def assertEqualIgnoreType(self, *args, **kwargs) -> None: + # If you are seeing this function used, that means test is written wrongly + # and deserves detailed investigation + return self.assertEqual(*args, exact_dtype=False, **kwargs) + + def assertEqualBroadcasting(self, x, y, *args, **kwargs) -> None: + r"""Tests if tensor x equals to y, if y to be broadcast to x.shape. + """ + if not isinstance(y, Iterable): + # int, float, etc. or different shape tensors + y = torch.ones_like(x) * y + if not isinstance(y, torch.Tensor): + # iterable, but not a tensor + y = torch.ones_like(x) * torch.tensor(y) + return self.assertEqual(x, y, *args, **kwargs) + + def assertEqual( + self, + x, + y, + msg: str | Callable[[str], str] | None = None, + *, + atol: float | None = None, + rtol: float | None = None, + equal_nan=True, + exact_dtype=True, + # TODO: default this to True + exact_device=False, + exact_layout=False, + exact_stride=False, + exact_is_coalesced=False + ): + # Hide this function from `pytest`'s traceback + __tracebackhide__ = True + + # numpy's dtypes are a superset of what PyTorch supports. In case we encounter an unsupported dtype, we fall + # back to an elementwise comparison. Note that this has to happen here and not for example in + # `TensorOrArrayPair`, since at that stage we can no longer split the array into its elements and perform + # multiple comparisons. + if any( + isinstance(input, np.ndarray) and not has_corresponding_torch_dtype(input.dtype) for input in (x, y) + ): + def to_list(input): + return input.tolist() if isinstance(input, (torch.Tensor, np.ndarray)) else list(input) + + x = to_list(x) + y = to_list(y) + # When comparing a sequence of numbers to a tensor, we need to convert the sequence to a tensor here. + # Otherwise, the pair origination of `are_equal` will fail, because the sequence is recognized as container + # that should be checked elementwise while the tensor is not. + elif isinstance(x, torch.Tensor) and isinstance(y, Sequence): + y = torch.as_tensor(y, dtype=x.dtype, device=x.device) + elif isinstance(x, Sequence) and isinstance(y, torch.Tensor): + x = torch.as_tensor(x, dtype=y.dtype, device=y.device) + + # unbind NSTs to compare them; don't do this for NJTs + if isinstance(x, torch.Tensor) and x.is_nested and x.layout == torch.strided: + x = x.unbind() + if isinstance(y, torch.Tensor) and y.is_nested and y.layout == torch.strided: + y = y.unbind() + + x, y = _unwrap_dtensor_for_comparison(x, y) + + error_metas = not_close_error_metas( + x, + y, + pair_types=( + NonePair, + RelaxedBooleanPair, + RelaxedNumberPair, + TensorOrArrayPair, + TypedStoragePair, + StringPair, + SetPair, + TypePair, + ObjectPair, + ), + sequence_types=( + Sequence, + Sequential, + ModuleList, + ParameterList, + ScriptList, + torch.utils.data.dataset.Subset, + ), + mapping_types=(Mapping, ModuleDict, ParameterDict, ScriptDict), + rtol=rtol, + rtol_override=self.rel_tol, + atol=atol, + atol_override=self.precision, + equal_nan=equal_nan, + check_device=exact_device, + check_dtype=exact_dtype, + check_layout=exact_layout, + check_stride=exact_stride, + check_is_coalesced=exact_is_coalesced, + ) + + if error_metas: + # See [ErrorMeta Cycles] + error_metas = [error_metas] # type: ignore[list-item] + # TODO: compose all metas into one AssertionError + raise error_metas.pop()[0].to_error( # type: ignore[index] + # This emulates unittest.TestCase's behavior if a custom message passed and + # TestCase.longMessage (https://docs.python.org/3/library/unittest.html#unittest.TestCase.longMessage) + # is True (default) + (lambda generated_msg: f"{generated_msg}\n{msg}") if isinstance(msg, str) and self.longMessage else msg + ) + + def assertNotEqual(self, x, y, msg: str | None = None, *, # type: ignore[override] + atol: float | None = None, rtol: float | None = None, **kwargs) -> None: + with self.assertRaises(AssertionError, msg=msg): + self.assertEqual(x, y, msg, atol=atol, rtol=rtol, **kwargs) + + def assertEqualTypeString(self, x, y) -> None: + # This API is used simulate deprecated x.type() is y.type() + self.assertEqual(x.device, y.device) + self.assertEqual(x.dtype, y.dtype) + self.assertEqual(x.is_sparse, y.is_sparse) + + def assertObjectIn(self, obj: Any, iterable: Iterable[Any]) -> None: + for elem in iterable: + if id(obj) == id(elem): + return + raise AssertionError("object not found in iterable") + + # Reimplemented to provide special behavior when + # _ignore_not_implemented_error is True + def assertRaises(self, expected_exception, *args, **kwargs): + if self._ignore_not_implemented_error: + context: AssertRaisesContextIgnoreNotImplementedError | None = \ + AssertRaisesContextIgnoreNotImplementedError(expected_exception, self) # type: ignore[call-arg] + try: + return context.handle('assertRaises', args, kwargs) # type: ignore[union-attr, arg-type] + finally: + # see https://bugs.python.org/issue23890 + context = None + else: + return super().assertRaises(expected_exception, *args, **kwargs) + + # Reimplemented to provide special behavior when + # _ignore_not_implemented_error is True + def assertRaisesRegex(self, expected_exception, expected_regex, *args, **kwargs): + # Verifies that an exception with the type expected_exception and message + # matching the regular expression defined by expected_regex is thrown. + # If the test is instantiated for a non-native device type (like XLA) + # then the message is not validated. + + # Checks whether the test is instantiated for a device type by testing + # if the test class has defined the device_type attribute and, + # if so, tests whether the instantiated device type is native or not + if hasattr(self, 'device_type') and self.device_type not in NATIVE_DEVICES and self.device_type != "mps": # type: ignore[attr-defined] + # empty string matches any string + expected_regex = '' + + if self._ignore_not_implemented_error: + context = AssertRaisesContextIgnoreNotImplementedError( # type: ignore[call-arg] + expected_exception, self, expected_regex) + return context.handle('assertRaisesRegex', args, kwargs) # type: ignore[attr-defined, arg-type] + else: + return super().assertRaisesRegex(expected_exception, expected_regex, *args, **kwargs) + + # Verifies that no unraisable exceptions are raised by callable. Unlike regular + # exceptions, these do not actually propagate to the caller and are + # suppressed. We must test for them specially. + def assertNoUnraisable(self, callable, *args, **kwargs): + raised = None + + def record_unraisable(unraisable): + nonlocal raised + raised = unraisable + + # Disable GC when running the callable to prevent spurious flakiness + # from unlucky GCs inside the callable + prev = gc.isenabled() + gc.disable() + try: + with unittest.mock.patch("sys.unraisablehook", record_unraisable): + callable(*args, **kwargs) + finally: + if prev: + gc.enable() + + self.assertIsNone(raised) + + # TODO: Support context manager interface + # NB: The kwargs forwarding to callable robs the 'subname' parameter. + # If you need it, manually apply your callable in a lambda instead. + def assertExpectedRaises(self, exc_type, callable, *args, **kwargs): + subname = None + if 'subname' in kwargs: + subname = kwargs['subname'] + del kwargs['subname'] + try: + callable(*args, **kwargs) + except exc_type as e: + self.assertExpected(str(e), subname) + return + # Don't put this in the try block; the AssertionError will catch it + self.fail(msg="Did not raise when expected to") + + def assertNotWarn(self, callable, msg=''): + r""" + Test if :attr:`callable` does not raise a warning. + """ + with warnings.catch_warnings(record=True) as ws: + warnings.simplefilter("always") # allow any warning to be raised + with set_warn_always_context(True): + callable() + self.assertTrue(len(ws) == 0, msg) + + @contextmanager + def assertWarnsOnceRegex(self, category, regex=''): + """Context manager for code that *must always* warn + + This filters expected warnings from the test and fails if + the expected warning is not caught. It uses set_warn_always() to force + TORCH_WARN_ONCE to behave like TORCH_WARN + """ + pattern = re.compile(regex) + with warnings.catch_warnings(record=True) as ws: + warnings.simplefilter("always") # allow any warning to be raised + with set_warn_always_context(True): + yield + if len(ws) == 0: + self.fail('no warning caught') + self.assertTrue(any(type(w.message) is category for w in ws)) + self.assertTrue( + any(re.match(pattern, str(w.message)) for w in ws), + f'{pattern}, {[w.message for w in ws if type(w.message) is category]}') + + def assertExpected(self, s, subname=None): + r""" + Test that a string matches the recorded contents of a file + derived from the name of this test and subname. This file + is placed in the 'expect' directory in the same directory + as the test script. You can automatically update the recorded test + output using --accept. + + If you call this multiple times in a single function, you must + give a unique subname each time. + """ + if not isinstance(s, str): + raise TypeError("assertExpected is strings only") + + def remove_prefix(text, prefix): + if text.startswith(prefix): + return text[len(prefix):] + return text + # NB: we take __file__ from the module that defined the test + # class, so we place the expect directory where the test script + # lives, NOT where test/common_utils.py lives. This doesn't matter in + # PyTorch where all test scripts are in the same directory as + # test/common_utils.py, but it matters in onnx-pytorch + module_id = self.__class__.__module__ + munged_id = remove_prefix(self.id(), module_id + ".") + test_file = os.path.realpath(sys.modules[module_id].__file__) # type: ignore[type-var] + expected_file = os.path.join(os.path.dirname(test_file), # type: ignore[type-var, arg-type] + "expect", + munged_id) + + subname_output = "" + if subname: + expected_file += "-" + subname + subname_output = f" ({subname})" + expected_file += ".expect" + expected = None + + def accept_output(update_type): + print(f"Accepting {update_type} for {munged_id}{subname_output}:\n\n{s}") + with open(expected_file, 'w') as f: + # Adjust for producer_version, leave s unmodified + s_tag = re.sub(r'(producer_version): "[0-9.]*"', + r'\1: "CURRENT_VERSION"', s) + f.write(s_tag) + + try: + with open(expected_file) as f: + expected = f.read() + except OSError as e: + if e.errno != errno.ENOENT: + raise + elif expecttest.ACCEPT: + return accept_output("output") + else: + raise RuntimeError( + f"I got this output for {munged_id}{subname_output}:\n\n{s}\n\n" + "No expect file exists; to accept the current output, run:\n" + f"python {__main__.__file__} {munged_id} --accept") from None + + # a hack for JIT tests + if IS_WINDOWS: + expected = re.sub(r'CppOp\[(.+?)\]', 'CppOp[]', expected) + s = re.sub(r'CppOp\[(.+?)\]', 'CppOp[]', s) + + # Adjust for producer_version + expected = expected.replace( + 'producer_version: "CURRENT_VERSION"', + f'producer_version: "{torch.onnx.producer_version}"' + ) + if expecttest.ACCEPT: + if expected != s: + return accept_output("updated output") + else: + if hasattr(self, "assertMultiLineEqual"): + # Python 2.7 only + # NB: Python considers lhs "old" and rhs "new". + self.assertMultiLineEqual(expected, s) + else: + self.assertEqual(s, expected) + + def assertExpectedStripMangled(self, s, subname=None): + s = re.sub(r'__torch__[^ ]+', '', s) + self.assertExpected(s, subname) + + def assertGreaterAlmostEqual(self, first, second, places=None, msg=None, delta=None): + """Assert that ``first`` is greater than or almost equal to ``second``. + + The equality of ``first`` and ``second`` is determined in a similar way to + the ``assertAlmostEqual`` function of the standard library. + """ + if delta is not None and places is not None: + raise TypeError("specify delta or places not both") + + if first >= second: + return + + diff = second - first + if delta is not None: + if diff <= delta: + return + + standardMsg = f"{first} not greater than or equal to {second} within {delta} delta" + else: + if places is None: + places = 7 + + if round(diff, places) == 0: + return + + standardMsg = f"{first} not greater than or equal to {second} within {places} places" + + msg = self._formatMessage(msg, standardMsg) + raise self.failureException(msg) + + def assertAtenOp(self, onnx_model, operator, overload_name=""): + all_aten_nodes = [p for p in onnx_model.graph.node + if p.op_type == "ATen" and p.domain == "org.pytorch.aten"] + self.assertTrue(all_aten_nodes) + + for op in all_aten_nodes: + attrs = {attr.name: attr.s.decode() for attr in op.attribute} + if attrs.get("operator") == operator: + break + + self.assertEqual(attrs["operator"], operator) # type: ignore[possibly-undefined] + self.assertEqual(attrs.get("overload_name", ""), overload_name) + + def check_nondeterministic_alert(self, fn, caller_name, should_alert=True): + '''Checks that an operation produces a nondeterministic alert when + expected while `torch.use_deterministic_algorithms(True)` is set. + + Args: + fn (callable): Function to check for a nondeterministic alert + + caller_name (str): Name of the operation that produces the + nondeterministic alert. This name is expected to appear at the + beginning of the error/warning message. + + should_alert (bool, optional): If True, then the check will only pass + if calling `fn` produces a nondeterministic error/warning with the + expected message. If False, then the check will only pass if + calling `fn` does not produce an error. Default: `True`. + ''' + + alert_message = '^' + caller_name + ' does not have a deterministic implementation, but you set' + + # Check that errors are thrown correctly + with DeterministicGuard(True): + if should_alert: + with self.assertRaisesRegex( + RuntimeError, + alert_message, + msg='expected a non-deterministic error, but it was not raised'): + fn() + + else: + # If a nondeterministic error is not expected, make sure + # that it is not raised + try: + fn() + except RuntimeError as e: + if 'does not have a deterministic implementation' in str(e): + self.fail( + 'did not expect non-deterministic error message, ' + + 'but got one anyway: "' + str(e) + '"') + # Reraise exceptions unrelated to nondeterminism + raise + + # Check that warnings are thrown correctly + with DeterministicGuard(True, warn_only=True): + if should_alert: + with self.assertWarnsRegex( + UserWarning, + alert_message): + fn() + else: + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + fn() + for warning in w: + if isinstance(warning, UserWarning): + self.assertTrue(re.search(alert_message, str(warning)) is None) + + # run code in subprocess and capture exceptions. + @staticmethod + def run_process_no_exception(code, env=None): + import subprocess + + with subprocess.Popen( + [sys.executable, "-c", code], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + env=env, + ) as p: + (stdout, stderr) = p.communicate() + return (stdout, stderr) + + # returns captured stderr + @staticmethod + def runWithPytorchAPIUsageStderr(code): + env = os.environ.copy() + env["PYTORCH_API_USAGE_STDERR"] = "1" + # remove CI flag since this is a wrapped test process. + # CI flag should be set in the parent process only. + env.pop("CI", None) + env.pop("TEST_SHOWLOCALS", None) + _stdout, stderr = TestCase.run_process_no_exception(code, env=env) + return stderr.decode('ascii') + + def _attempt_load_from_subprocess( + self, + file: pathlib.Path, + import_string: str, + expected_failure_message: str | None = None + ) -> None: + """ + Attempts weights_only `torch.load` in a subprocess. This is used to test that + weights_only `torch.load` works as expected without global imports. + + Args: + file (pathlib.Path): The path to the checkpoint to load. + import_string (str): import string to add to the script + exected_failure_message (str, optional): The expected failure message if the + checkpoint fails to load. If None, the test will pass + """ + script = f"import torch;{import_string}torch.load(r'{file}', weights_only=True)" + cm = ( + self.assertRaisesRegex(RuntimeError, re.escape(expected_failure_message)) + if expected_failure_message else contextlib.nullcontext() + ) + with cm: + try: + subprocess.check_output( + [sys.executable, "-c", script], + # On Windows, opening the subprocess with the default CWD makes `import torch` + # fail, so just set CWD to this script's directory + cwd=os.path.dirname(os.path.realpath(__file__)), + stderr=subprocess.STDOUT, + ) + except subprocess.CalledProcessError as e: + raise RuntimeError(e.output.decode("utf-8")) from None + + +class TestCaseBase(TestCase): + # Calls to super() in dynamically created classes are a bit odd. + # See https://github.com/pytorch/pytorch/pull/118586 for more info + # Subclassing this class and then calling super(TestCaseBase) will run + # TestCase's setUp, tearDown etc functions + pass + + +def download_file(url, binary=True): + from urllib.parse import urlsplit + from urllib import request, error + + filename = os.path.basename(urlsplit(url)[2]) + data_dir = get_writable_path(os.path.join(os.path.dirname(__file__), 'data')) + path = os.path.join(data_dir, filename) + + if os.path.exists(path): + return path + try: + with request.urlopen(url, timeout=15) as f1, open(path, 'wb' if binary else 'w') as f2: + data = f1.read() + f2.write(data) + return path + except error.URLError as e: + msg = f"could not download test file '{url}'" + warnings.warn(msg, RuntimeWarning, stacklevel=2) + raise unittest.SkipTest(msg) from e + +def find_free_port(): + """ + Finds an available port and returns that port number. + + NOTE: If this function is being used to allocate a port to Store (or + indirectly via init_process_group or init_rpc), it should be used + in conjunction with the `retry_on_connect_failures` decorator as there is a potential + race condition where the allocated port may become unavailable before it can be used + """ + with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock: + sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + sock.bind(('localhost', 0)) + _, port = sock.getsockname() + return port + +# Errors that we can get in c10d initialization for which we should retry tests for. +ADDRESS_IN_USE = "Address already in use" +CONNECT_TIMEOUT = "connect() timed out." + +def retry_on_connect_failures(func=None, connect_errors=(ADDRESS_IN_USE)): + """Reruns a test if the test returns a RuntimeError and the exception + contains one of the strings in connect_errors.""" + # This if block is executed when using this function as a decorator with arguments. + if func is None: + return partial(retry_on_connect_failures, connect_errors=connect_errors) + + @wraps(func) + def wrapper(*args, **kwargs): + n_retries = 10 + tries_remaining = n_retries + while True: + try: + return func(*args, **kwargs) + except RuntimeError as error: + if any(connect_error in str(error) for connect_error in connect_errors): + tries_remaining -= 1 + if tries_remaining == 0: + raise RuntimeError(f"Failing after {n_retries} retries with error: {str(error)}") from error + time.sleep(random.random()) + continue + raise + return wrapper + + +# Decorator to retry upon certain Exceptions. +def retry(ExceptionToCheck, tries=3, delay=3, skip_after_retries=False): + def deco_retry(f): + @wraps(f) + def f_retry(*args, **kwargs): + mtries, mdelay = tries, delay + while mtries > 1: + try: + return f(*args, **kwargs) + except ExceptionToCheck as e: + msg = f"{e}, Retrying in {mdelay:d} seconds..." + print(msg) + time.sleep(mdelay) + mtries -= 1 + try: + return f(*args, **kwargs) + except ExceptionToCheck as e: + raise unittest.SkipTest(f"Skipping after {tries} consecutive {str(e)}") from e if skip_after_retries else e + return f_retry # true decorator + return deco_retry + + +# FIXME: modernize these to be consistent with make_tensor +# and review including them in torch.testing +# Methods for matrix generation + +def random_square_matrix_of_rank(l, rank, dtype=torch.double, device='cpu'): + if rank > l: + raise AssertionError(f"rank must be <= l, got rank={rank}, l={l}") + A = torch.randn(l, l, dtype=dtype, device=device) + u, s, vh = torch.linalg.svd(A, full_matrices=False) + for i in range(l): + if i >= rank: + s[i] = 0 + elif s[i] == 0: + s[i] = 1 + return (u * s.to(dtype).unsqueeze(-2)) @ vh + +def random_well_conditioned_matrix(*shape, dtype, device, mean=1.0, sigma=0.001): + """ + Returns a random rectangular matrix (batch of matrices) + with singular values sampled from a Gaussian with + mean `mean` and standard deviation `sigma`. + The smaller the `sigma`, the better conditioned + the output matrix is. + """ + primitive_dtype = { + torch.float: torch.float, + torch.double: torch.double, + torch.cfloat: torch.float, + torch.cdouble: torch.double + } + x = torch.rand(shape, dtype=dtype, device=device) + m = x.size(-2) + n = x.size(-1) + u, _, vh = torch.linalg.svd(x, full_matrices=False) + s = (torch.randn(*(shape[:-2] + (min(m, n),)), dtype=primitive_dtype[dtype], device=device) * sigma + mean) \ + .sort(-1, descending=True).values.to(dtype) + return (u * s.unsqueeze(-2)) @ vh + +# Returns a noncontiguous (tensor with the same shape and values as t +# The noncontiguous tensor is constructed such that elements in the innermost +# dimension are separated by zeros or (whenever possible) nans +# TODO: consider more complicated noncontiguity schemes +def noncontiguous_like(t): + # Short-circuits if t is already noncontiguous + if not t.is_contiguous(): + return t + + # Choose a "weird" value that won't be accessed + if t.dtype.is_floating_point or t.dtype.is_complex: + value = math.nan + elif t.dtype == torch.bool: + value = True + else: + value = 12 + + result = t.new_empty(t.shape + (2,)) + result[..., 0] = value + result[..., 1] = t.detach() + result = result[..., 1] + result.requires_grad_(t.requires_grad) + return result + +# TODO: remove this (prefer make_symmetric_matrices below) +def random_symmetric_matrix(l, *batches, **kwargs): + dtype = kwargs.get('dtype', torch.double) + device = kwargs.get('device', 'cpu') + A = torch.randn(*(batches + (l, l)), dtype=dtype, device=device) + A = (A + A.mT).div_(2) + return A + +# Creates a symmetric matrix or batch of symmetric matrices +# Shape must be a square matrix or batch of square matrices +def make_symmetric_matrices(*shape, device, dtype): + if shape[-1] != shape[-2]: + raise AssertionError(f"expected square matrix, got shape[-1]={shape[-1]} != shape[-2]={shape[-2]}") + t = make_tensor(shape, device=device, dtype=dtype) + t = (t + t.mT).div_(2) + return t + +def random_hermitian_matrix(l, *batches, **kwargs): + dtype = kwargs.get('dtype', torch.double) + device = kwargs.get('device', 'cpu') + A = torch.randn(*(batches + (l, l)), dtype=dtype, device=device) + A = (A + A.mH).div_(2) + return A + + +def random_symmetric_psd_matrix(l, *batches, **kwargs): + """ + Returns a batch of random symmetric positive-semi-definite matrices. + The shape of the result is batch_dims + (matrix_size, matrix_size) + The following example creates a tensor of size 2 x 4 x 3 x 3 + >>> # xdoctest: +SKIP("undefined variables") + >>> matrices = random_symmetric_psd_matrix(3, 2, 4, dtype=dtype, device=device) + """ + dtype = kwargs.get('dtype', torch.double) + device = kwargs.get('device', 'cpu') + A = torch.randn(*(batches + (l, l)), dtype=dtype, device=device) + return A @ A.mT + + +def random_hermitian_psd_matrix(matrix_size, *batch_dims, dtype=torch.double, device='cpu'): + """ + Returns a batch of random Hermitian positive-semi-definite matrices. + The shape of the result is batch_dims + (matrix_size, matrix_size) + The following example creates a tensor of size 2 x 4 x 3 x 3 + >>> # xdoctest: +SKIP("undefined variables") + >>> matrices = random_hermitian_psd_matrix(3, 2, 4, dtype=dtype, device=device) + """ + A = torch.randn(*(batch_dims + (matrix_size, matrix_size)), dtype=dtype, device=device) + return A @ A.mH + + +# TODO: remove this (prefer make_symmetric_pd_matrices below) +def random_symmetric_pd_matrix(matrix_size, *batch_dims, **kwargs): + dtype = kwargs.get('dtype', torch.double) + device = kwargs.get('device', 'cpu') + A = torch.randn(*(batch_dims + (matrix_size, matrix_size)), + dtype=dtype, device=device) + return torch.matmul(A, A.mT) \ + + torch.eye(matrix_size, dtype=dtype, device=device) * 1e-5 + + +# Creates a symmetric positive-definite matrix or batch of +# such matrices +def make_symmetric_pd_matrices(*shape, device, dtype): + if shape[-1] != shape[-2]: + raise AssertionError(f"expected square matrix, got shape[-1]={shape[-1]} != shape[-2]={shape[-2]}") + t = make_tensor(shape, device=device, dtype=dtype) + i = torch.eye(shape[-1], device=device, dtype=dtype) * 1e-5 + return t @ t.mT + i + +def random_hermitian_pd_matrix(matrix_size, *batch_dims, dtype, device): + """ + Returns a batch of random Hermitian positive-definite matrices. + The shape of the result is batch_dims + (matrix_size, matrix_size) + The following example creates a tensor of size 2 x 4 x 3 x 3 + >>> # xdoctest: +SKIP("undefined variables") + >>> matrices = random_hermitian_pd_matrix(3, 2, 4, dtype=dtype, device=device) + """ + A = torch.randn(*(batch_dims + (matrix_size, matrix_size)), + dtype=dtype, device=device) + return A @ A.mH + torch.eye(matrix_size, dtype=dtype, device=device) + +# Creates a full rank matrix with distinct singular values or +# a batch of such matrices +def make_fullrank_matrices_with_distinct_singular_values(*shape, device, dtype, requires_grad=False): + with torch.no_grad(): + t = make_tensor(shape, device=device, dtype=dtype) + u, _, vh = torch.linalg.svd(t, full_matrices=False) + real_dtype = t.real.dtype if t.dtype.is_complex else t.dtype + k = min(shape[-1], shape[-2]) + # We choose the singular values to be "around one" + # This is to make the matrix well conditioned + # s = [2, 3, ..., k+1] + s = torch.arange(2, k + 2, dtype=real_dtype, device=device) + # s = [2, -3, 4, ..., (-1)^k k+1] + s[1::2] *= -1. + # 1 + 1/s so that the singular values are in the range [2/3, 3/2] + # This gives a condition number of 9/4, which should be good enough + s.reciprocal_().add_(1.) + # Note that the singular values need not be ordered in an SVD so + # we don't need need to sort S + x = (u * s.to(u.dtype)) @ vh + x.requires_grad_(requires_grad) + return x + +def random_matrix(rows, columns, *batch_dims, **kwargs): + """Return rectangular matrix or batches of rectangular matrices. + + Parameters: + dtype - the data type + device - the device kind + singular - when True, the output will be singular + """ + dtype = kwargs.get('dtype', torch.double) + device = kwargs.get('device', 'cpu') + silent = kwargs.get("silent", False) + singular = kwargs.get("singular", False) + if silent and not torch._C.has_lapack: + return torch.ones(rows, columns, dtype=dtype, device=device) + + A = torch.randn(batch_dims + (rows, columns), dtype=dtype, device=device) + if A.numel() == 0: + return A + u, _, vh = torch.linalg.svd(A, full_matrices=False) + k = min(rows, columns) + s = torch.linspace(1 / (k + 1), 1, k, dtype=dtype, device=device) + if singular: + # make matrix singular + s[k - 1] = 0 + if k > 2: + # increase the order of singularity so that the pivoting + # in LU factorization will be non-trivial + s[0] = 0 + return (u * s.unsqueeze(-2)) @ vh + + +def random_matrix_with_scaled_reduction_dim(rows, columns, *batch_dims, **kwargs): + """Return rectangular matrix or batches of rectangular matrices + with entries being iid and sampled from N(0, sigma^2) such that + the variance of (A @ A.T)[..., i, j] is 1 if reduction_dim=-1, or + the variance of (A.T @ A)[..., i, j] is 1 if reduction_dim=-2. + + Parameters: + dtype - the data type + device - the device kind + requires_grad - whether output requires grad + reduction_dim - the row/column dimension to re-scale. + Expected to be either -1 (columns) or -2 (rows). + """ + dtype = kwargs.get('dtype', torch.double) + device = kwargs.get('device', 'cpu') + requires_grad = kwargs.get('requires_grad', False) + reduction_dim = kwargs.get('reduction_dim', -1) + + shape = (*batch_dims, rows, columns) + red_scale = math.sqrt(shape[reduction_dim]) + res = torch.randn(*shape, dtype=dtype, device=device) / red_scale + res.requires_grad_(requires_grad) + return res + + +def random_lowrank_matrix(rank, rows, columns, *batch_dims, **kwargs): + """Return rectangular matrix or batches of rectangular matrices with + given rank. + """ + B = random_matrix(rows, rank, *batch_dims, **kwargs) + C = random_matrix(rank, columns, *batch_dims, **kwargs) + return B.matmul(C) + + +def _generate_indices_prefer_all_rows(rows: int, cols: int, num_indices: int) -> torch.Tensor: + """Generate indices for a row x cols matrix, preferring at least one index per row if possible.""" + indices = [] # type: ignore[var-annotated] + n_per_row = math.ceil(num_indices / rows) + col_indices = list(range(cols)) + + for r in range(rows): + # Note that this can yield overlapping indices + indices.extend((r, c) for c in random.choices(col_indices, k=n_per_row)) + + return torch.tensor(indices[:num_indices]) + + +def random_sparse_matrix(rows, columns, density=0.01, **kwargs): + """Return rectangular random sparse matrix within given density. + + The density of the result approaches to given density as the size + of the matrix is increased and a relatively small value of density + is specified but higher than min(rows, columns)/(rows * columns) + for non-singular matrices. + """ + dtype = kwargs.get('dtype', torch.double) + device = kwargs.get('device', 'cpu') + + nonzero_elements = max(min(rows, columns), int(rows * columns * density)) + indices = _generate_indices_prefer_all_rows(rows, columns, nonzero_elements) + values = torch.randn(nonzero_elements, dtype=dtype, device=device) + + # ensure that the diagonal dominates + values *= torch.tensor([-float(i - j)**2 for i, j in indices], dtype=dtype, device=device).exp() + A = torch.sparse_coo_tensor(indices.t(), values, (rows, columns), device=device) + return A.coalesce() + + +def random_sparse_pd_matrix(matrix_size, density=0.01, **kwargs): + """Return random sparse positive-definite matrix with given density. + + The eigenvalues of the matrix are defined as:: + arange(1, matrix_size+1)/matrix_size + + Algorithm: + A = diag(arange(1, matrix_size+1)/matrix_size) + while : + + R = + A = R^T A R + """ + import math + torch = kwargs.get('torch', globals()['torch']) + dtype = kwargs.get('dtype', torch.double) + device = kwargs.get('device', 'cpu') + data = {(i, i): float(i + 1) / matrix_size + for i in range(matrix_size)} + + + def multiply(data, N, i, j, cs, sn, left=True): + for k in range(N): + if left: + ik, jk = (k, i), (k, j) + else: + ik, jk = (i, k), (j, k) + aik, ajk = data.get(ik, 0), data.get(jk, 0) + aik, ajk = cs * aik + sn * ajk, -sn * aik + cs * ajk + if aik: + data[ik] = aik + else: + data.pop(ik, None) + if ajk: + data[jk] = ajk + else: + data.pop(jk, None) + + target_nnz = density * matrix_size * matrix_size + while len(data) < target_nnz: + i = random.randint(0, matrix_size - 1) + j = random.randint(0, matrix_size - 1) + if i != j: + theta = random.uniform(0, 2 * math.pi) + cs = math.cos(theta) + sn = math.sin(theta) + multiply(data, matrix_size, i, j, cs, sn, left=True) + multiply(data, matrix_size, i, j, cs, sn, left=False) + icoords, jcoords, values = [], [], [] + for (i, j), v in sorted(data.items()): + icoords.append(i) + jcoords.append(j) + values.append(v) + indices_tensor = torch.tensor([icoords, jcoords]) + return torch.sparse_coo_tensor(indices_tensor, values, (matrix_size, matrix_size), dtype=dtype, device=device) + +# FIXME: remove this by updating test suites using it +def do_test_dtypes(self, dtypes, layout, device): + for dtype in dtypes: + if dtype != torch.float16: + out = torch.zeros((2, 3), dtype=dtype, layout=layout, device=device) + self.assertIs(dtype, out.dtype) + self.assertIs(layout, out.layout) + self.assertEqual(device, out.device) + +# FIXME: remove this by updating test suites using it +def do_test_empty_full(self, dtypes, layout, device): + shape = torch.Size([2, 3]) + + def check_value(tensor, dtype, layout, device, value, requires_grad): + self.assertEqual(shape, tensor.shape) + self.assertIs(dtype, tensor.dtype) + self.assertIs(layout, tensor.layout) + self.assertEqual(tensor.requires_grad, requires_grad) + if tensor.is_cuda and device is not None: + self.assertEqual(device, tensor.device) + if value is not None: + fill = tensor.new(shape).fill_(value) + self.assertEqual(tensor, fill) + + def get_int64_dtype(dtype): + module = '.'.join(str(dtype).split('.')[1:-1]) + if not module: + return torch.int64 + return operator.attrgetter(module)(torch).int64 + + default_dtype = torch.get_default_dtype() + check_value(torch.empty(shape), default_dtype, torch.strided, -1, None, False) + check_value(torch.full(shape, -5.), default_dtype, torch.strided, -1, None, False) + for dtype in dtypes: + for rg in {dtype.is_floating_point, False}: + int64_dtype = get_int64_dtype(dtype) + v = torch.empty(shape, dtype=dtype, device=device, layout=layout, requires_grad=rg) + check_value(v, dtype, layout, device, None, rg) + out = v.new() + check_value(torch.empty(shape, out=out, device=device, layout=layout, requires_grad=rg), + dtype, layout, device, None, rg) + check_value(v.new_empty(shape), dtype, layout, device, None, False) + check_value(v.new_empty(shape, dtype=int64_dtype, device=device, requires_grad=False), + int64_dtype, layout, device, None, False) + check_value(torch.empty_like(v), dtype, layout, device, None, False) + check_value(torch.empty_like(v, dtype=int64_dtype, layout=layout, device=device, requires_grad=False), + int64_dtype, layout, device, None, False) + + if dtype is not torch.float16 and layout != torch.sparse_coo: + fv = 3 + v = torch.full(shape, fv, dtype=dtype, layout=layout, device=device, requires_grad=rg) + check_value(v, dtype, layout, device, fv, rg) + check_value(v.new_full(shape, fv + 1), dtype, layout, device, fv + 1, False) + out = v.new() + check_value(torch.full(shape, fv + 2, out=out, device=device, layout=layout, requires_grad=rg), + dtype, layout, device, fv + 2, rg) + check_value(v.new_full(shape, fv + 3, dtype=int64_dtype, device=device, requires_grad=False), + int64_dtype, layout, device, fv + 3, False) + check_value(torch.full_like(v, fv + 4), dtype, layout, device, fv + 4, False) + check_value(torch.full_like(v, fv + 5, + dtype=int64_dtype, layout=layout, device=device, requires_grad=False), + int64_dtype, layout, device, fv + 5, False) + +# FIXME: improve load_tests() documentation here +running_script_path = None # type: ignore[var-annotated] +def set_running_script_path(): + global running_script_path + try: + running_file = os.path.abspath(os.path.realpath(sys.argv[0])) + if running_file.endswith('.py'): # skip if the running file is not a script + running_script_path = running_file + except Exception: + pass + +def check_test_defined_in_running_script(test_case): + if running_script_path is None: + return + test_case_class_file = os.path.abspath(os.path.realpath(inspect.getfile(test_case.__class__))) + if test_case_class_file != running_script_path: + raise AssertionError( + f'Class of loaded TestCase "{test_case.id()}" ' + f'is not defined in the running script "{running_script_path}", but in "{test_case_class_file}". Did you ' + "accidentally import a unittest.TestCase from another file?" + ) + +def load_tests(loader, tests, pattern): + set_running_script_path() + test_suite = unittest.TestSuite() + for test_group in tests: + if not DISABLE_RUNNING_SCRIPT_CHK: + for test in test_group: + check_test_defined_in_running_script(test) + if test_group._tests: + test_suite.addTest(test_group) + return test_suite + +# FIXME: document this and move it to test_serialization +class BytesIOContext(io.BytesIO): + def __enter__(self): + return self + + def __exit__(self, *args): + pass + +# Tentative value for nondet_tol for gradcheck when backward implementation +# relies on nondeterministic operations, i.e., those listed here: +# https://pytorch.org/docs/stable/generated/torch.use_deterministic_algorithms.html +# +# For more information see https://github.com/pytorch/pytorch/issues/56202 +GRADCHECK_NONDET_TOL = 1e-12 + +TEST_WITH_SLOW_GRADCHECK: bool = TestEnvironment.def_flag( + "TEST_WITH_SLOW_GRADCHECK", + env_var="PYTORCH_TEST_WITH_SLOW_GRADCHECK", +) + +skipIfSlowGradcheckEnv = unittest.skipIf( + TEST_WITH_SLOW_GRADCHECK, + "Tests that don't use gradcheck don't need to run on slow_gradcheck CI", +) + + +def gradcheck(fn, inputs, **kwargs): + # Wrapper around gradcheck that enables certain keys by default. + # Use this testing-internal gradcheck instead of autograd.gradcheck so that new features like vmap and + # forward-mode AD are tested by default. We create this wrapper because we'd like to keep new checks + # to be disabled to default for the public-facing api to avoid breaking user code. + # + # All PyTorch devs doing testing should use this wrapper instead of autograd.gradcheck. + default_values = { + "check_batched_grad": True, + "fast_mode": True, + } + + if TEST_WITH_SLOW_GRADCHECK: + default_values["fast_mode"] = False + + for key, value in default_values.items(): + # default value override values explicitly set to None + k = kwargs.get(key) + kwargs[key] = k if k is not None else value + + return torch.autograd.gradcheck(fn, inputs, **kwargs) + +def gradgradcheck(fn, inputs, grad_outputs=None, **kwargs): + # Wrapper around gradgradcheck that enables certain keys by default + # See gradcheck above for an explanation of why we need something like this. + # + # All PyTorch devs doing testing should use this wrapper instead of autograd.gradgradcheck + default_values = { + "check_batched_grad": True, + "fast_mode": True, + } + + if TEST_WITH_SLOW_GRADCHECK: + default_values["fast_mode"] = False + + for key, value in default_values.items(): + # default value override values explicitly set to None + k = kwargs.get(key) + kwargs[key] = k if k is not None else value + + return torch.autograd.gradgradcheck(fn, inputs, grad_outputs, **kwargs) + + +def _assertGradAndGradgradChecks(test_case, apply_fn, inputs, **kwargs): + # call assert function rather than returning a bool since it's nicer + # if we get whether this failed on the gradcheck or the gradgradcheck. + test_case.assertTrue(gradcheck(apply_fn, inputs, **kwargs)) + test_case.assertTrue(gradgradcheck(apply_fn, inputs, **kwargs)) + + +@contextmanager +def set_cwd(path: str) -> Iterator[None]: + old_cwd = os.getcwd() + try: + os.chdir(path) + yield + finally: + os.chdir(old_cwd) + + +# FIXME: delete this +# Using @toleranceOverride specific to your test is the recommended way +# of doing this. These are just some values that worked for test_nn. +dtype2prec_DONTUSE = {torch.float: 1e-5, + torch.double: 1e-5, + torch.half: 1e-2, + torch.bfloat16: 1e-1} + +# FIXME: move to test_sparse or sparse utils +# This is a wrapper that wraps a test to run this test twice, one with +# coalesced=True, another with coalesced=False for coalesced/uncoalesced sparse tensors. +def coalescedonoff(f): + @wraps(f) + def wrapped(self, *args, **kwargs): + f(self, *args, **kwargs, coalesced=True) + f(self, *args, **kwargs, coalesced=False) + return wrapped + + +def is_coalesced_indices(s): + indices = s._indices() + hash_coeffs = (1,) + s.shape[s.sparse_dim() - 1:0:-1] + hash_indices = torch.tensor(hash_coeffs, device=s.device).cumprod(-1).flip(-1) + if s.sparse_dim() > 1: + hash_indices.unsqueeze_(-1) + hash_indices = (indices * hash_indices).sum(0) + else: + hash_indices = indices * hash_indices + + # check if indices are sorted + res = torch.allclose(hash_indices, hash_indices.sort()[0]) + + # check if there are no repeated indices + res = res and torch.allclose(hash_indices, hash_indices.unique()) + + return res + + +@contextlib.contextmanager +def disable_gc(): + if gc.isenabled(): + try: + gc.disable() + yield + finally: + gc.enable() + else: + yield + + +def find_library_location(lib_name: str) -> Path: + # return the shared library file in the installed folder if exist, + # else the file in the build folder + torch_root = Path(torch.__file__).resolve().parent + path = torch_root / 'lib' / lib_name + if os.path.exists(path): + return path + torch_root = Path(__file__).resolve().parents[2] + return torch_root / 'build' / 'lib' / lib_name + +def skip_but_pass_in_sandcastle(reason): + """ + Similar to unittest.skip, however in the sandcastle environment it just + "passes" the test instead to avoid creating tasks complaining about tests + skipping continuously. + """ + def decorator(func): + if not IS_SANDCASTLE: + func.__unittest_skip__ = True + func.__unittest_skip_why__ = reason + return func + + @wraps(func) + def wrapper(*args, **kwargs): + print(f'Skipping {func.__name__} on sandcastle for following reason: {reason}', file=sys.stderr) + return + return wrapper + + return decorator + +def mock_wrapper(method): + """ + Returns a function that calls the real implementation of a method + in addition to passing args to a mock object. + """ + mock = MagicMock() + + @wraps(method) + def wrapper(self, *args, **kwargs): + mock(*args, **kwargs) + return method(self, *args, **kwargs) + wrapper.mock = mock # type: ignore[attr-defined] + return wrapper + +def get_tensors_from(args, kwargs): + """ Returns a set of all Tensor objects in the given args and kwargs. """ + return set([arg for arg in args if isinstance(arg, Tensor)] + + [v for v in kwargs.values() if isinstance(v, Tensor)]) + + +# Returns scalar tensor representation of a list of integer byte values +def bytes_to_scalar(byte_list: list[int], dtype: torch.dtype, device: torch.device): + dtype_to_ctype: dict[torch.dtype, Any] = { + torch.int8: ctypes.c_int8, + torch.uint8: ctypes.c_uint8, + torch.uint16: ctypes.c_uint16, + torch.uint32: ctypes.c_uint32, + torch.uint64: ctypes.c_uint64, + torch.int16: ctypes.c_int16, + torch.int32: ctypes.c_int32, + torch.int64: ctypes.c_int64, + torch.bool: ctypes.c_bool, + torch.float32: ctypes.c_float, + torch.complex64: ctypes.c_float, + torch.float64: ctypes.c_double, + torch.complex128: ctypes.c_double, + } + ctype = dtype_to_ctype[dtype] + num_bytes = ctypes.sizeof(ctype) + + def check_bytes(byte_list): + for byte in byte_list: + if not (0 <= byte <= 255): + raise AssertionError(f"byte value out of range: expected 0 <= byte <= 255, got {byte}") + + if dtype.is_complex: + if len(byte_list) != (num_bytes * 2): + raise AssertionError( + f"expected len(byte_list) == {num_bytes * 2} for complex dtype, got {len(byte_list)}" + ) + check_bytes(byte_list) + real = ctype.from_buffer((ctypes.c_byte * num_bytes)( + *byte_list[:num_bytes])).value + imag = ctype.from_buffer((ctypes.c_byte * num_bytes)( + *byte_list[num_bytes:])).value + res = real + 1j * imag + else: + if len(byte_list) != num_bytes: + raise AssertionError( + f"expected len(byte_list) == {num_bytes}, got {len(byte_list)}" + ) + check_bytes(byte_list) + res = ctype.from_buffer((ctypes.c_byte * num_bytes)( + *byte_list)).value + + return torch.tensor(res, device=device, dtype=dtype) + + +def copy_func(f): + """Based on http://stackoverflow.com/a/6528148/190597 (Glenn Maynard)""" + g = types.FunctionType(f.__code__, f.__globals__, name=f.__name__, + argdefs=f.__defaults__, + closure=f.__closure__) + g = functools.update_wrapper(g, f) + g.__kwdefaults__ = f.__kwdefaults__ # type: ignore[attr-defined] + return g + + +def xfail_inherited_tests(tests): + """ + Given a list of test names which are defined by a superclass of the + class this decorates, mark them as expected failure. This is useful + if you are doing poor man's parameterized tests by subclassing a generic + test class. + """ + def deco(cls): + for t in tests: + # NB: expectedFailure operates by mutating the method in question, + # which is why you have to copy the function first + setattr(cls, t, unittest.expectedFailure(copy_func(getattr(cls, t)))) + return cls + return deco + + +def skip_but_pass_in_sandcastle_if(condition, reason): + """ + Similar to unittest.skipIf, however in the sandcastle environment it just + "passes" the test instead to avoid creating tasks complaining about tests + skipping continuously. + """ + def decorator(func): + if condition: + if IS_SANDCASTLE: + @wraps(func) + def wrapper(*args, **kwargs): + print(f'Skipping {func.__name__} on sandcastle for following reason: {reason}', file=sys.stderr) + return wrapper + else: + func.__unittest_skip__ = True + func.__unittest_skip_why__ = reason + + return func + + return decorator + +def dtype_name(dtype): + """ Returns the pretty name of the dtype (e.g. torch.int64 -> int64). """ + return str(dtype).split('.')[1] + + +def _cpu_sleep(cycles: int) -> None: + """Spin-wait for approximately the given number of cycles.""" + for _ in range(cycles): + pass + + +def device_sleep(device: str, cycles: int) -> None: + """Sleep for the given number of cycles on the specified device. + + For CPU, temporarily patches torch.cpu._sleep if needed. + For CUDA/other devices, uses torch.get_device_module(device)._sleep. + """ + if device == "cpu": + orig = getattr(torch.cpu, "_sleep", None) + torch.cpu._sleep = _cpu_sleep + try: + torch.cpu._sleep(cycles) + finally: + if orig is None: + delattr(torch.cpu, "_sleep") + else: + torch.cpu._sleep = orig + else: + torch.get_device_module(device)._sleep(cycles) + + +@functools.lru_cache +def get_cycles_per_ms(device: str = "cuda") -> float: + """Measure and return approximate number of cycles per millisecond for device _sleep. + + Args: + device: Device type to measure cycles for ("cuda" or "cpu"). + + Works for both CUDA (torch.cuda._sleep) and CPU (torch.cpu._sleep). + """ + test_cycles = 1000000 + + if device == "cpu": + import time + + def measure() -> float: + start = time.perf_counter() + _cpu_sleep(test_cycles) + end = time.perf_counter() + elapsed_ms = (end - start) * 1000 + cycles_per_ms = test_cycles / elapsed_ms if elapsed_ms > 0 else 1000000 + return cycles_per_ms + else: + def measure() -> float: + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + torch.cuda._sleep(test_cycles) + end.record() + end.synchronize() + cycles_per_ms = test_cycles / start.elapsed_time(end) + return cycles_per_ms + + # Get 10 values and remove the 2 max and 2 min and return the avg. + # This is to avoid system disturbance that skew the results, e.g. + # the very first cuda call likely does a bunch of init, which takes + # much longer than subsequent calls. + # + # Tested on both Tesla V100, Quadro GP100, Titan RTX, RTX 3090 GPUs + # and seems to return stable values. Therefore, we enable caching + # using lru_cache decorator above. + num = 10 + vals = [measure() for _ in range(num)] + vals = sorted(vals) + return mean(vals[2 : num - 2]) + + +# OpInfo utils + +T = TypeVar('T') +def first_sample(self: unittest.TestCase, samples: Iterable[T]) -> T: + """ + Returns the first sample from an iterable of samples, like those returned by OpInfo. + The test will be skipped if no samples are available. + """ + try: + return next(iter(samples)) + except StopIteration as e: + raise unittest.SkipTest('Skipped! Need at least 1 sample input') from e + +# this helper method is to recursively +# clone the tensor-type input of operators tested by OpInfo +def clone_input_helper(input): + if isinstance(input, torch.Tensor): + return torch.clone(input) + + if isinstance(input, Sequence): + return tuple(map(clone_input_helper, input)) + + return input + +@contextmanager +def custom_op(opname, symbolic_fn, opset_version): + """Context manager/decorator to test ONNX export with custom operator""" + try: + register_custom_op_symbolic(opname, symbolic_fn, opset_version) + yield + finally: + unregister_custom_op_symbolic(opname, opset_version) + + +def outs_and_grads(fn, graph_inps, inps): + outs = fn(*graph_inps) + for out in pytree.tree_leaves(outs): + if isinstance(out, torch.Tensor) and out.requires_grad: + out.sum().backward(retain_graph=True) + grads = [inp.grad for inp in pytree.tree_leaves(inps) if isinstance(inp, torch.Tensor)] + for inp in pytree.tree_leaves(inps): + if isinstance(inp, torch.Tensor): + inp.grad = None + return outs, grads + +def compare_equal_outs_and_grads(test, m1, m2, inps): + r1, g1 = outs_and_grads(m1, inps, inps) + r2, g2 = outs_and_grads(m2, inps, inps) + test.assertEqual(r1, r2) + test.assertEqual(g1, g2) + +class TestGradients(TestCase): + exact_dtype = True + + # Copies inputs to inplace operations to avoid inplace modifications + # to leaves requiring gradient + def _get_safe_inplace(self, inplace_variant): + @wraps(inplace_variant) + def _fn(t, *args, **kwargs): + return inplace_variant(t.clone(), *args, **kwargs) + + return _fn + + def _check_helper(self, device, dtype, op, variant, check, *, check_forward_ad=False, check_backward_ad=True, + check_batched_grad=None, check_batched_forward_grad=False): + if check not in ('gradcheck', 'bwgrad_bwgrad', 'fwgrad_bwgrad'): + raise AssertionError( + f"check must be one of ('gradcheck', 'bwgrad_bwgrad', 'fwgrad_bwgrad'), got {check!r}" + ) + # NB: check_backward_ad does not affect gradgradcheck (always True) + if variant is None: + self.skipTest("Skipped! Variant not implemented.") + if not op.supports_dtype(dtype, torch.device(device).type): + self.skipTest(f"Skipped! {op.name} does not support dtype {str(dtype)}") + + def is_inplace(variant): + if hasattr(variant, "__wrapped__"): + return variant.__wrapped__ is op.get_inplace() + return variant is op.get_inplace() + + include_conjugated_inputs = op.test_conjugated_samples and dtype.is_complex + + samples = op.sample_inputs(device, dtype, requires_grad=True, include_conjugated_inputs=include_conjugated_inputs, + small_inputs_only=TEST_WITH_SLOW_GRADCHECK) + + for sample in samples: + if sample.broadcasts_input and is_inplace(variant): + continue + + # Gradcheck expects tensors as its input, but autograd actually supports tensorlists + # and tensors passed as kwargs. The following creates a function that accepts just + # the tensors that require grad as varargs, and then recomposes them back into the + # original input. + + # Creates gradcheck inputs by identifying tensors requiring grad + all_args = None + if is_iterable_of_tensors(sample.input): + all_args = chain(sample.input, sample.args, sample.kwargs.values()) + else: + all_args = tuple(chain((sample.input,), sample.args, sample.kwargs.values())) # type: ignore[assignment] + gradcheck_args = tuple(x for x in all_args if (isinstance(x, torch.Tensor) and x.requires_grad)) # type: ignore[union-attr] + + # Verifies sample input tensors should have no grad + # This may happen if the same tensor is used in two different SampleInputs + for t in gradcheck_args: + self.assertIsNone(t.grad, + "A sampled input has a gradient before running autograd. " + "This usually means that (at least) one input tensor is reused " + "across different SampleInputs. " + "Please create a new tensor for each SampleInput.") + + def _input_recomposition_helper(inputs, inp, input_idx): + if is_iterable_of_tensors(inp): + tensor_list = [] + for x in inp: + if isinstance(x, torch.Tensor) and x.requires_grad: + tensor_list.append(inputs[input_idx]) + input_idx = input_idx + 1 + else: + tensor_list.append(x) + return tensor_list, input_idx + elif isinstance(inp, torch.Tensor) and inp.requires_grad: + return inputs[input_idx], input_idx + 1 + else: + return inp, input_idx + + def fn(*inputs): + # Puts inputs back into sample properly + positional_args = [] + input_idx = 0 + inp, input_idx = _input_recomposition_helper(inputs, sample.input, input_idx) + positional_args.append(inp) + + for x in sample.args: + inp, input_idx = _input_recomposition_helper(inputs, x, input_idx) + positional_args.append(inp) + + # Recreates kwargs + kwargs = {} + for k, v in sample.kwargs.items(): + inp, input_idx = _input_recomposition_helper(inputs, v, input_idx) + kwargs[k] = inp + + output = op.gradcheck_wrapper(variant, *positional_args, **kwargs) + if sample.output_process_fn_grad is not None: + return sample.output_process_fn_grad(output) + return output + + if check == 'gradcheck': + if check_batched_grad is None: + check_batched_grad = op.check_batched_grad + self.assertTrue(gradcheck(fn, gradcheck_args, + check_batched_grad=check_batched_grad, + check_grad_dtypes=True, + nondet_tol=op.gradcheck_nondet_tol, + fast_mode=op.gradcheck_fast_mode, + check_forward_ad=check_forward_ad, + check_backward_ad=check_backward_ad, + check_undefined_grad=True, + check_batched_forward_grad=check_batched_forward_grad)) + elif check in ('bwgrad_bwgrad', 'fwgrad_bwgrad'): # gradgrad check + self.assertFalse(check_forward_ad, msg="Cannot run forward AD check for gradgradcheck") + for gen_non_contig_grad_outputs in (False, True): + kwargs = { + "gen_non_contig_grad_outputs": gen_non_contig_grad_outputs, + "check_batched_grad": op.check_batched_gradgrad, + "check_grad_dtypes": True, + "nondet_tol": op.gradcheck_nondet_tol, + "fast_mode": op.gradcheck_fast_mode + } + if check == "fwgrad_bwgrad": + kwargs["check_fwd_over_rev"] = True + kwargs["check_rev_over_rev"] = False + kwargs["check_batched_grad"] = False + kwargs["check_undefined_grad"] = False + + self.assertTrue(gradgradcheck(fn, gradcheck_args, **kwargs)) + else: + self.assertTrue(False, msg="Unknown check requested!") + + def _grad_test_helper(self, device, dtype, op, variant, *, check_forward_ad=False, check_backward_ad=True, + check_batched_grad=None, check_batched_forward_grad=False): + return self._check_helper(device, dtype, op, variant, 'gradcheck', check_forward_ad=check_forward_ad, + check_backward_ad=check_backward_ad, check_batched_grad=check_batched_grad, + check_batched_forward_grad=check_batched_forward_grad) + + def _skip_helper(self, op, device, dtype): + if dtype not in op.supported_backward_dtypes(torch.device(device).type): + self.skipTest("Skipped! Op doesn't support autograd for this dtype.") + if not op.supports_autograd and not op.supports_forward_ad: + self.skipTest("Skipped! autograd not supported.") + + + + +# Base TestCase for NT tests; used to define common helpers, etc. +class NestedTensorTestCase(TestCase): + def assertEqualIgnoringNestedInts(self, a, b): + # unbinding NJTs allows us to compare them as essentially equal without + # caring about exact nested int comparison + def _unbind_njts(x): + if isinstance(x, torch.Tensor) and x.is_nested and x.layout == torch.jagged: + return x.unbind() + else: + return x + + self.assertEqual(pytree.tree_map(_unbind_njts, a), pytree.tree_map(_unbind_njts, b)) + + def assertEqualNoncontigAware(self, a, b): + # assertEqual() doesn't take into account lengths, so hack around this + # by comparing unbound components and shapes + self.assertEqualIgnoringNestedInts(a, b) + + def _get_njt_shapes(x): + return ( + x.shape + if isinstance(x, torch.Tensor) and x.is_nested + else None + ) + + a_shapes = pytree.tree_map(_get_njt_shapes, a) + b_shapes = pytree.tree_map(_get_njt_shapes, b) + self.assertEqual(a_shapes, b_shapes) + + @contextlib.contextmanager + def branch_nested_state(self): + """Context manager to branch and restore the nested tensor state.""" + nested_tensor_module = torch.nested._internal.nested_tensor + original_tensor_symint_registry = nested_tensor_module._tensor_symint_registry.copy() + original_tensor_id_counter = nested_tensor_module._tensor_id_counter + try: + yield + finally: + nested_tensor_module._tensor_id_counter = original_tensor_id_counter + nested_tensor_module._tensor_symint_registry = original_tensor_symint_registry + + +def munge_exc(e, *, suppress_suffix=True, suppress_prefix=True, file=None, skip=0): + from torch._dynamo.trace_rules import _as_posix_path + + if file is None: + file = inspect.stack()[1 + skip].filename # skip one frame + + file = _as_posix_path(file) + s = _as_posix_path(str(e)) + + # Remove everything that looks like stack frames in NOT this file + def repl_frame(m): + if m.group(2) != file: + return "" + # Don't accept top-level, even for this script, these will wobble + # depending on how the testing script was invoked + if m.group(3) == "": + return "" + + return m.group(0) + + s = re.sub( + r'( *)File "([^"]+)", line \d+, in (.+)\n(\1 .+\n( +[~^]+ *\n)?)+', + repl_frame, + s, + ) + s = re.sub(r"line \d+", "line N", s) + s = re.sub(r".py:\d+", ".py:N", s) + s = re.sub(r'https:/([a-zA-Z0-9_.-]+)', r'https://\1', s) + s = re.sub(file, _as_posix_path(os.path.basename(file)), s) + s = re.sub(_as_posix_path(os.path.join(os.path.dirname(torch.__file__), "")), "", s) + # 3.10 CALL_FUNCTION bytecode compatibility for dynamo graph break messages + s = re.sub( + r"attempting to trace CALL_FUNCTION:.*$", + "attempting to trace CALL: a function call, e.g. f(x, y):", + s, + flags=re.MULTILINE, + ) + if suppress_suffix: + s = re.sub(r"\n*Set TORCH_LOGS.+", "", s, flags=re.DOTALL) + s = re.sub(r"\n*You can suppress this exception.+", "", s, flags=re.DOTALL) + s = re.sub(r"\n*Set TORCHDYNAMO_VERBOSE=1.+", "", s, flags=re.DOTALL) + if suppress_prefix: + s = re.sub(r"Cannot export model.+\n\n", "", s) + s = re.sub(r" +$", "", s, flags=re.MULTILINE) + return s + + +@contextmanager +def check_leaked_tensors(limit=1, matched_type=torch.Tensor): + """Wrap around operations you want to ensure are not leaking tensor memory. + + This code intentionally ignores other reference cycles, which can be benign and which we have plenty + of in pytorch code. It focuses on any reference cycles that directly or indirectly result holding a Tensor alive, + since this is likely a more serious leak than typical python refcycles. + + limit specifies how many tensors to dump debug graphs for (default=1) + """ + def match_obj(obj): + return isinstance(obj, matched_type) + + try: + gc.collect() + gc.set_debug(gc.DEBUG_SAVEALL) + garbage_objs = [] # type: ignore[var-annotated] + + # run the user code, after cleaning any existing refcycles, and then check for new ones + # also allow usercode to check the garbage objs (e.g. for assertion) after exiting ctxmgr + yield garbage_objs + + gc.collect() + garbage_objs.extend(filter(match_obj, gc.garbage)) + num_garbage_objs = len(garbage_objs) + if num_garbage_objs > 0: + warnings.warn( + f"{num_garbage_objs} tensors were found in the garbage. Did you introduce a reference cycle?", stacklevel=2 + ) + try: + import objgraph # type: ignore[import-not-found,import-untyped] + warnings.warn( + f"Dumping first {limit} objgraphs of leaked {matched_type}s rendered to png", stacklevel=2 + ) + for g in garbage_objs[:limit]: + objgraph.show_backrefs([g], max_depth=10) + except ImportError: + warnings.warn("`pip install objgraph` to enable memory leak debugging", stacklevel=2) + + finally: + gc.set_debug(0) + + +def remove_cpp_extensions_build_root(): + """ + Removes the default root folder under which extensions are built. + """ + default_build_root = cpp_extension.get_default_build_root() + if os.path.exists(default_build_root): + if IS_WINDOWS: + # rmtree returns permission error: [WinError 5] Access is denied + # on Windows, this is a workaround + subprocess.run(["rm", "-rf", default_build_root], stdout=subprocess.PIPE) + else: + shutil.rmtree(default_build_root, ignore_errors=True) + + +def install_cpp_extension(extension_root): + # Wipe the build / install dirs if they exist + build_dir = os.path.join(extension_root, "build") + install_dir = os.path.join(extension_root, "install") + for d in (build_dir, install_dir): + if os.path.exists(d): + shutil.rmtree(d) + + # Build the extension + cmd = [sys.executable, "-m", "pip", "install", extension_root, "-v", "--no-build-isolation", "--root", install_dir] + return_code = shell(cmd, cwd=extension_root, env=os.environ) + if return_code != 0: + raise RuntimeError(f"build failed for cpp extension at {extension_root}") + + mod_install_dir = None + # install directory is the one that is named site-packages + for root, directories, _ in os.walk(install_dir): + for directory in directories: + if "-packages" in directory: + mod_install_dir = os.path.join(root, directory) + + if mod_install_dir is None: + raise RuntimeError(f"installation failed for cpp extension at {extension_root}") + + if mod_install_dir not in sys.path: + sys.path.insert(0, mod_install_dir) + + +# Decorator to provide a helper to load inline extensions to a temp directory +def scoped_load_inline(func): + + @wraps(func) + def wrapper(*args, **kwargs): + def load_inline(*args, **kwargs): + if IS_WINDOWS: + # TODO(xmfan): even using TemporaryDirectoryName will result in permission error + return cpp_extension.load_inline(*args, **kwargs) + + if "build_directory" in kwargs: + raise AssertionError("build_directory should not be specified when using scoped_load_inline") + with TemporaryDirectoryName() as temp_dir_name: + if kwargs.get("verbose", False): + print(f'Using temporary extension directory {temp_dir_name}...', file=sys.stderr) + kwargs["build_directory"] = temp_dir_name + return cpp_extension.load_inline(*args, **kwargs) + + return func(*args, load_inline=load_inline, **kwargs) + return wrapper + +def recover_orig_fp32_precision(fn): + @contextlib.contextmanager + def recover(): + old_mkldnn_conv_p = torch.backends.mkldnn.conv.fp32_precision # type: ignore[attr-defined] + old_mkldnn_rnn_p = torch.backends.mkldnn.rnn.fp32_precision # type: ignore[attr-defined] + old_mkldnn_matmul_p = torch.backends.mkldnn.matmul.fp32_precision # type: ignore[attr-defined] + old_cudnn_conv_p = torch.backends.cudnn.conv.fp32_precision # type: ignore[attr-defined] + old_cudnn_rnn_p = torch.backends.cudnn.rnn.fp32_precision # type: ignore[attr-defined] + old_cuda_matmul_p = torch.backends.cuda.matmul.fp32_precision + try: + yield + finally: + torch.backends.mkldnn.conv.fp32_precision = old_mkldnn_conv_p # type: ignore[attr-defined] + torch.backends.mkldnn.rnn.fp32_precision = old_mkldnn_rnn_p # type: ignore[attr-defined] + torch.backends.mkldnn.matmul.fp32_precision = old_mkldnn_matmul_p # type: ignore[attr-defined] + torch.backends.cudnn.conv.fp32_precision = old_cudnn_conv_p # type: ignore[attr-defined] + torch.backends.cudnn.rnn.fp32_precision = old_cudnn_rnn_p # type: ignore[attr-defined] + torch.backends.cuda.matmul.fp32_precision = old_cuda_matmul_p + + return recover()(fn) + +def skipIfPythonVersionMismatch(predicate): + vi = sys.version_info + + def dec_fn(fn): + @wraps(fn) + def wrap_fn(self, *args, **kwargs): + if predicate(vi.major, vi.minor, vi.micro): + return fn(self, *args, **kwargs) + else: + raise unittest.SkipTest("Python version mismatch") + return wrap_fn + return dec_fn + +# Decorator to patch multiple test class members for the duration of the subtest +def patch_test_members(updates: dict[str, Any]): + def decorator(test_func): + @wraps(test_func) + def wrapper(self, *args, **kwargs): + # Store the original values of the specified members + original_values = {member: getattr(self, member) for member in updates} + + # Update the members before running the subtest + for member, value in updates.items(): + setattr(self, member, value) + + # Run the test function, allowing subtests to run + try: + return test_func(self, *args, **kwargs) + finally: + # Restore the original values of the specified members after the subtest finishes + for member, original_value in original_values.items(): + setattr(self, member, original_value) + + return wrapper + return decorator + +def get_gcc_major_version(): + """ + Return GCC major version as int, or None if GCC is not available. + """ + try: + out = subprocess.check_output( + ["gcc", "-dumpfullversion", "-dumpversion"], + stderr=subprocess.STDOUT, + text=True, + ).strip() + return int(out.split(".")[0]) + except Exception: + return None diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_xpu.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_xpu.py new file mode 100644 index 0000000000000000000000000000000000000000..ee4f8c921eac5b84ce9e01c5e9dc64578b6abe3a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/common_xpu.py @@ -0,0 +1,88 @@ +import enum +import functools + +import torch +import torch.xpu +from torch.testing._internal.common_utils import IS_WINDOWS, LazyVal, TEST_XPU + + +XPU_ALREADY_INITIALIZED_ON_IMPORT = torch.xpu.is_initialized() + + +class XPUCodename(enum.Enum): + PVC = "PVC" # Intel® Data Center GPU Max Series + BMG = "BMG" # Intel® Arc™ Pro Battlemage Graphics + + +class XPUArch(enum.IntEnum): + Unknown = 0 + Xe = 1 # Xe HPC + Xe2 = 2 + + +# device_id -> GPU codename +# From https://github.com/intel/intel-graphics-compiler/blob/master/inc/common/igfxfmid.h +_DEVICE_ID_TO_CODENAME = { + 0x0BD0: XPUCodename.PVC, + 0x0BD4: XPUCodename.PVC, + 0x0BD5: XPUCodename.PVC, + 0x0BD6: XPUCodename.PVC, + 0x0BD7: XPUCodename.PVC, + 0x0BD8: XPUCodename.PVC, + 0x0BD9: XPUCodename.PVC, + 0x0BDA: XPUCodename.PVC, + 0x0BDB: XPUCodename.PVC, + 0x0B69: XPUCodename.PVC, + 0x0B6E: XPUCodename.PVC, + 0xE202: XPUCodename.BMG, + 0xE20B: XPUCodename.BMG, + 0xE20C: XPUCodename.BMG, + 0xE20D: XPUCodename.BMG, + 0xE210: XPUCodename.BMG, + 0xE212: XPUCodename.BMG, + 0xE215: XPUCodename.BMG, + 0xE216: XPUCodename.BMG, + 0xE220: XPUCodename.BMG, + 0xE221: XPUCodename.BMG, + 0xE222: XPUCodename.BMG, + 0xE223: XPUCodename.BMG, +} + +# GPU codename -> architecture +_CODENAME_TO_ARCH = { + XPUCodename.PVC: XPUArch.Xe, + XPUCodename.BMG: XPUArch.Xe2, +} + + +@functools.lru_cache(1) +def get_xpu_codename() -> XPUCodename | None: + device_id = torch.xpu.get_device_capability()["device_id"] + return _DEVICE_ID_TO_CODENAME.get(device_id) + + +@functools.lru_cache(1) +def get_xpu_arch() -> XPUArch | None: + codename = get_xpu_codename() + return _CODENAME_TO_ARCH.get(codename, XPUArch.Unknown) + + +Xe2_Or_Later = LazyVal( + lambda: torch.xpu.is_available() and get_xpu_arch() >= XPUArch.Xe2 +) + + +def evaluate_platform_supports_flash_attention(): + if TEST_XPU: + return not IS_WINDOWS and Xe2_Or_Later + return False + + +PLATFORM_SUPPORTS_FLASH_ATTENTION_XPU: bool = LazyVal( + lambda: evaluate_platform_supports_flash_attention() +) + +# Importing this module should NOT eagerly initialize XPU +if not XPU_ALREADY_INITIALIZED_ON_IMPORT: + if torch.xpu.is_initialized(): + raise AssertionError("XPU should not be initialized on import") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/composite_compliance.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/composite_compliance.py new file mode 100644 index 0000000000000000000000000000000000000000..3a0d33887bb4107465e2fe6b643fa355de05e098 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/composite_compliance.py @@ -0,0 +1,616 @@ +# mypy: ignore-errors + +import torch +from torch import Tensor +import itertools + +from torch.utils._python_dispatch import TorchDispatchMode +from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten +from torch.utils import _pytree as pytree +from functools import partial +from torch.utils._mode_utils import no_dispatch, all_same_mode +import torch.autograd.forward_ad as fwAD +from collections.abc import Callable +import re + + +def check_attr_consistency(wrapper_tensor, metadata_name, metadata_accessor): + elem = wrapper_tensor.elem + metadata_wrapper_tensor = metadata_accessor(wrapper_tensor) + metadata_elem = metadata_accessor(elem) + if metadata_wrapper_tensor == metadata_elem: + return + raise RuntimeError( + f"This operator is not Composite Compliant: the " + f"{metadata_name} of the tensor was modified directly without " + f"going through the PyTorch dispatcher.") + +def check_metadata_consistency(wrapper_tensor, CCT): + # CCT: CompositeCompliantTensor class which is generated using generate_cct + if not isinstance(wrapper_tensor, CCT): + return + things_to_check = { + 'shape': Tensor.size, + 'dtype': lambda x: x.dtype, + 'device': lambda x: x.device, + 'numel': Tensor.numel, + 'stride': Tensor.stride, + 'storage_offset': Tensor.storage_offset, + } + for metadata_name, metadata_accessor in things_to_check.items(): + check_attr_consistency(wrapper_tensor, metadata_name, metadata_accessor) + +def is_view_fn(func): + return func.overloadpacket.__name__ in { + 'as_strided', + 'detach', + 'diagonal', + 'expand', + 'expand_as', + 'movedim', + 'narrow', + 'permute', + 'select', + 'squeeze', + 'transpose', + 't', + 'real', + 'imag', + 'view_as_real', + 'view_as_complex', + 'unflatten', + 'unfold', + 'unsqueeze', + 'view', + 'view_as', + 'unbind', + 'split', + 'split_with_sizes', + 'vsplit', + 'hsplit', + 'tensor_split', + 'chunk', + 'swapaxes', + 'slice', + '_reshape_alias', + '_unsafe_view', + '_conj', + 'alias', + } + +# manually populated from native_functions that have inplace_view: True. +# In the future we will probably be able to grab that list directly +def is_inplace_view_fn(func): + return func.overloadpacket.__name__ in { + 'as_strided_', + 'detach_', + 'squeeze_', + 'swapaxes_', + 'swapdims_', + 't_', + 'transpose_', + 'unsqueeze_', + } + + +# Introspection please save us +def is_inplace(func): + name = func.overloadpacket.__name__ + if re.match('__i.+__', name): + return True + if re.match('__.+__', name): + return False + return name[-1] == '_' + + +def generate_cct_and_mode(autograd_view_consistency=True): + # This function returns a new class CompositeCompliantTensor + # The two arguments control the behaviour described below. + + # autograd_view_consistency: + # If True, alias result using `set_` if func returns a view + # (See Note [Alias Result]). + # Since Forward AD doesn't work with `set_` + # we disable it by setting alias to False. + + class CompositeCompliantTensor(torch.Tensor): + elem: torch.Tensor + + __slots__ = ['elem'] + + @staticmethod + def __new__(cls, elem, mode, *args, **kwargs): + if type(elem) is cls: + raise AssertionError( + "Wrapping a CompositeCompliantTensor in a CompositeCompliantTensor is not supported" + ) + + # The storage of CompositeCompliantTensor should never be used directly + # by a Composite operation; if the Composite + # operator attempts to read from the storage without dispatching then it'll + # raise a RuntimeError due to it being a meta storage. + r = torch.Tensor._make_wrapper_subclass( + cls, elem.size(), + dtype=elem.dtype, layout=elem.layout, + device=elem.device, requires_grad=elem.requires_grad, + strides=elem.stride(), storage_offset=elem.storage_offset()) + + if elem.requires_grad: + # CompositeCompliantTensor steals the "requires_grad"-ness. + # Why a new copy of `elem`? Because sometimes OpInfo shares inputs between tests... + tmp = torch.empty( + (), + dtype=elem.dtype, + device=elem.device, + layout=elem.layout, + requires_grad=False, + ) + # Use set_ rather than empty_strided() + copy_ so that we can preserve + # things like storage_offset. + tmp.set_( + source=elem.untyped_storage().clone(), + storage_offset=elem.storage_offset(), + size=elem.size(), + stride=elem.stride(), + ) + r.elem = tmp + else: + r.elem = elem + + if r.stride() != r.elem.stride(): + raise AssertionError(f"Expected r.stride() == r.elem.stride(), got {r.stride()} != {r.elem.stride()}") + + # Propagate conjugate bits to the wrapper tensor + # Ref: https://github.com/albanD/subclass_zoo/issues/24 + # Ref: https://github.com/albanD/subclass_zoo/issues/21 + torch._C._set_conj(r, r.elem.is_conj()) + torch._C._set_neg(r, r.elem.is_neg()) + + r.mode = mode + return r + + def __repr__(self): + return f"CompositeCompliantTensor({self.elem})" + + @classmethod + def __torch_dispatch__(cls, func, types, args=(), kwargs=None): + all_args = pytree.arg_tree_leaves(*args, **(kwargs or {})) + modes = tuple(e.mode for e in all_args if isinstance(e, CompositeCompliantTensor)) + if not all_same_mode(modes): + raise RuntimeError("Multiple CompositeCompliantTensorModes NYI") + with modes[0]: + return func(*args, **kwargs) + + class CompositeCompliantTensorMode(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + def unwrap(e): + return e.elem if isinstance(e, CompositeCompliantTensor) else e + + def wrap(e): + return CompositeCompliantTensor(e, self) if isinstance(e, torch.Tensor) else e + + if func is torch.ops.aten._local_scalar_dense.default: + raise RuntimeError( + ".item() is not allowed to be called inside of composite " + "functions in the PyTorch library because not all backends " + "and/or Tensor subclasses (e.g. vmap, ProxyTensor) support them.") + + if func.overloadpacket.__name__ in ('set_', 'resize_'): + raise RuntimeError( + f"{func.__name__} is not allowed to be called inside of " + f"Composite operators.") + + if is_inplace(func): + # NB: We are making an assumption that if the function is in-place, + # then the first argument is being written to. Introspection please save us! + mutated_argument = args[0] + if not isinstance(mutated_argument, CompositeCompliantTensor) and \ + any(isinstance(a, CompositeCompliantTensor) for a in args[1:]): + raise RuntimeError( + 'Not composite compliant: performing in-place operation ' + f'{func.__name__} where the Tensor being written to is ' + 'regular Tensor but the other tensors are Tensor Subclasses. ' + 'Please try to avoid this in-place operation.') + + unwrapped_args = tree_map(unwrap, args) + unwrapped_kwargs = tree_map(unwrap, kwargs) + unwrapped_rs = func(*unwrapped_args, **unwrapped_kwargs) + rs = tree_map(wrap, unwrapped_rs) + + if is_view_fn(func) and autograd_view_consistency: + # Note [Alias Result] + # Autograd asserts that for B = A.view_fn(...), B and A's storages + # are the same. Here we try to make B alias A to avoid those asserts. + # See https://github.com/pytorch/pytorch/issues/65339 for more information + # about the issue. + with no_dispatch(): + # Idea: this is a weird way of getting a storage that aliases the input. + # This is a workaround for #65339. + # 1. under no_dispatch, all of the wrapper tensors look like regular + # tensors with special storage (the storage is nullptr and + # advertises CPU/CUDA device. + # 2. we run func, which ends up running the view operation + # 3. All view operations reuse the input's storage and return + # result Tensor(s) with new sizes/strides/offset that alias + # the input. + # 4. we set the storage (and sizes/strides/offset) of the wrapper + # tensor results to be that of the tensors that alias the input + result = func(*args, **kwargs) + if isinstance(result, (tuple, list)): + for a, b in zip(rs, result, strict=True): + a.set_(b) + else: + rs.set_(result) + + # Some operations are allowed to in-place modify the metadata of the + # inputs. The only ones are the "inplace view functions"; when we + # run into these, we manually modify the metadata of the input. + with no_dispatch(): + if is_inplace_view_fn(func): + func(*args, **kwargs) + + # For each CompositeCompliantTensor t, we check that t and t.elem + # have consistent metadata. If they don't have consistent metadata, + # that means the operator did something fishy. + check = partial(check_metadata_consistency, CCT=CompositeCompliantTensor) + pytree.tree_map_(check, args) + pytree.tree_map_(check, kwargs) + pytree.tree_map_(check, rs) + return rs + + return CompositeCompliantTensor, CompositeCompliantTensorMode() + +def is_tensorlist(lst): + if not isinstance(lst, list) and not isinstance(lst, tuple): + return False + if len(lst) == 0: + return False + all_tensors = all(isinstance(elt, torch.Tensor) for elt in lst) + if all_tensors: + return True + exists_one_tensor = all(isinstance(elt, torch.Tensor) for elt in lst) + if exists_one_tensor: + raise RuntimeError('This test assumes that PyTorch APIs cannot take ' + 'mixed lists of Tensor and other things') + return False + + +def maybe_map(fn, should_map, arg): + return fn(arg) if should_map else arg + + +def wrap(arg, CCT, cct_mode): + # CCT: CompositeCompliantTensor class which is generated using generate_cct_and_mode + if isinstance(arg, torch.Tensor): + return CCT(arg, cct_mode) + if is_tensorlist(arg): + return [CCT(a, cct_mode) for a in arg] + raise RuntimeError("wrap assumes that the input can be wrapped") + + +# Given a list of flat arguments, some of which may be Tensors, return all +# possible ways some of the arguments could be CompositeCompliantTensors (CCT). +# For example, given Tensors A, B, C and flat_args = [A, 1, B], +# We would return the following 4 options: +# [CCT(A), 1, CCT(B)] +# [CCT(A), 1, B] +# [A, 1, CCT(B)] +# [A, 1, B] +# NB: Yes, this is exponential. No, we don't care too much because PyTorch ops +# don't accept that many input Tensors. +def generate_subclass_choices(flat_args, CCT, cct_mode): + # CCT: CompositeCompliantTensor class which is generated using generate_cct_and_mode + is_tensor_likes = [isinstance(arg, torch.Tensor) or is_tensorlist(arg) for arg in flat_args] + subclass_options = [[False, True] if is_tensor_like else [False] for is_tensor_like in is_tensor_likes] + + for which_args_are_wrapped in itertools.product(*subclass_options): + + result = [maybe_map(partial(wrap, CCT=CCT, cct_mode=cct_mode), should_wrap_arg, arg) + for should_wrap_arg, arg in zip(which_args_are_wrapped, flat_args, strict=True)] + yield result, which_args_are_wrapped + + +# For an operation f(*args, **kwargs), each Tensor argument may either be +# a regular Tensor or a Tensor Subclass. This iterator iterates through +# all of those options. +def generate_subclass_choices_args_kwargs(args, kwargs, CCT, cct_mode): + # CCT: CompositeCompliantTensor class which is generated using generate_cct_and_mode + flat_kwargs, spec = tree_flatten(kwargs) + flat_args_kwargs = list(args) + list(flat_kwargs) + for choice, debug_metadata in generate_subclass_choices(flat_args_kwargs, CCT, cct_mode): + new_args = choice[:len(args)] + new_kwargs = tree_unflatten(choice[len(args):], spec) + which_args_are_wrapped = debug_metadata[:len(args)] + which_kwargs_are_wrapped = tree_unflatten(debug_metadata[len(args):], spec) + yield new_args, new_kwargs, which_args_are_wrapped, which_kwargs_are_wrapped + + +def raise_composite_compliance_error(err, additional_info=''): + raise RuntimeError( + "Composite compliance check failed with " + "the above error.\n" + f"{additional_info}" + "If you are adding an OpInfo of an " + "existing operator, please feel free to skip this test " + "because the problem was pre-existing and file an issue. " + "Otherwise, if you added a new operator, please read " + "through the Composite Compliance section in " + "aten/src/ATen/native/README.md for how to resolve this. " + ) from err + + +# This test checks ALL possible permutations of calling `op` with arguments +# that are individually either a regular Tensor or a Tensor subclass. +# +# The general strategy is to wrap some Tensor args and kwargs in +# CompositeCompliantTensor wrappers and call the operation. + +# If some composite operation does any non-compliant behavior, +# CompositeCompliantTensor will raise an error. +def check_all_permutations(op, args, kwargs, assert_equal_fn): + CCT, cct_mode = generate_cct_and_mode() + expected = op(*args, **kwargs) + for choice in generate_subclass_choices_args_kwargs(args, kwargs, CCT, cct_mode): + new_args, new_kwargs, which_args_are_wrapped, which_kwargs_are_wrapped = choice + + try: + actual = op(*new_args, **new_kwargs) + # NOTE: [What errors are Composite Compliance trying to catch?] + # + # There's two things we want to catch: + # - errors that would raise within the torch_dispatch impl + # - data_ptr accesses + # The first is easy to filter for (we could make the error a different + # error class), the second is always going to be a RuntimeError due to + # how it is implemented (if you try to access the data_ptr of the + # wrapper Tensor, it raises you some internal RuntimeError). + # + # So the most general thing to catch here was RuntimeError. If you + # are here and debugging why your test failed, it's plausible that + # the operator itself is broken and that there are other tests failing. + except RuntimeError as err: + raise_composite_compliance_error( + err, + f"- wrapped_args: {which_args_are_wrapped}\n" + f"- wrapped_kwargs: {which_kwargs_are_wrapped}\n" + ) + + def unwrap(e): + return e.elem if isinstance(e, CCT) else e + + assert_equal_fn(tree_map(unwrap, actual), expected) + +# Checks via the usage of torch dispatch mode certain anti-patterns that +# are not composite compliant. +# +# In particular, the anti-pattern we are trying to prevent is a user +# creating an empty tensor and then resize_-ing it. Torch Dispatch Mode helps +# here because all factory functions will create tensors that are +# CompositeCompliantTensor. +# +# The general strategy is to wrap all Tensor args and kwargs in +# CompositeCompliantTensor wrappers. If an operator that is +# Composite does any non-compliant behavior, +# CompositeCompliantTensor will raise an error. +def check_with_mode(op, args, kwargs, assert_equal_fn): + CCT, cct_mode = generate_cct_and_mode() + + def wrap(e): + return CCT(e, cct_mode) if isinstance(e, torch.Tensor) else e + + expected = op(*args, **kwargs) + + args = tree_map(wrap, args) + kwargs = tree_map(wrap, kwargs) + try: + with cct_mode: + actual = op(*args, **kwargs) + # see NOTE: [What errors are Composite Compliance trying to catch?] + except RuntimeError as err: + raise_composite_compliance_error(err) + + def unwrap(e): + return e.elem if isinstance(e, CCT) else e + + assert_equal_fn(tree_map(unwrap, actual), expected) + +def gather_leaf_tensors(args, kwargs): + leaf_tensors = [] + args, _args_spec = tree_flatten(args) + kwargs, _kwargs_spec = tree_flatten(kwargs) + args = args + kwargs + for arg in args: + if not isinstance(arg, torch.Tensor): + continue + if arg.requires_grad: + leaf_tensors.append(arg) + return leaf_tensors + + +def compute_expected_grads(op, args, kwargs, output_process_fn_grad=None, gradcheck_wrapper=None): + if gradcheck_wrapper is None: + results = op(*args, **kwargs) + else: + results = gradcheck_wrapper(op, *args, **kwargs) + + if output_process_fn_grad is not None: + results = output_process_fn_grad(results) + + flat_results = pytree.tree_leaves(results) + flat_results = [r for r in flat_results if isinstance(r, torch.Tensor)] + flat_diff_results = [r for r in flat_results if r.requires_grad] + if len(flat_diff_results) <= 0: + raise AssertionError("Expected len(flat_diff_results) > 0") + + grads = [torch.ones(r.shape, device=r.device, dtype=r.dtype) for r in flat_diff_results] + leaf_tensors = gather_leaf_tensors(args, kwargs) + if len(leaf_tensors) <= 0: + raise AssertionError("Expected len(leaf_tensors) > 0") + return torch.autograd.grad(flat_diff_results, leaf_tensors, + grads, allow_unused=True, retain_graph=True) + + +# Checks if the backward formula is composite compliant by testing +# all possible permutations of {inputs, grad_outputs} being +# CompositeCompliantTensor or regular Tensors. +# +# NB: it is important that op is accepted as a Callable and not an OpInfo, +# this means we can apply check_backward_formula to things that aren't OpInfos +# while debugging. +def check_backward_formula(op: Callable, args, kwargs, + output_process_fn_grad=None, + gradcheck_wrapper=None, assert_equal_fn=None): + CCT, cct_mode = generate_cct_and_mode() + + expected = compute_expected_grads(op, args, kwargs, output_process_fn_grad, gradcheck_wrapper) + + for choice in generate_subclass_choices_args_kwargs(args, kwargs, CCT, cct_mode): + new_args, new_kwargs, which_args_are_wrapped, which_kwargs_are_wrapped = choice + leaf_tensors = gather_leaf_tensors(new_args, new_kwargs) + if len(leaf_tensors) <= 0: + raise AssertionError("Expected len(leaf_tensors) > 0") + + try: + if gradcheck_wrapper is None: + results = op(*new_args, **new_kwargs) + else: + results = gradcheck_wrapper(op, *new_args, **new_kwargs) + if output_process_fn_grad is not None: + results = output_process_fn_grad(results) + # see NOTE: [What errors are Composite Compliance trying to catch?] + except RuntimeError as err: + raise_composite_compliance_error( + err, + f"- wrapped_args: {which_args_are_wrapped}\n" + f"- wrapped_kwargs: {which_kwargs_are_wrapped}\n" + ) + + flat_results = pytree.tree_leaves(results) + flat_results = [r for r in flat_results if isinstance(r, torch.Tensor)] + flat_diff_results = [r for r in flat_results if r.requires_grad] + if len(flat_diff_results) <= 0: + raise AssertionError("Expected len(flat_diff_results) > 0") + + # NB: ones, not ones_like, so we get a regular Tensor here + grads = [torch.ones(r.shape, device=r.device, dtype=r.dtype) + for r in flat_diff_results] + for flat_new_grads, which_grad_is_batched in generate_subclass_choices(grads, CCT, cct_mode): + try: + actual = torch.autograd.grad(flat_diff_results, leaf_tensors, flat_new_grads, + allow_unused=True, retain_graph=True) + # see NOTE: [What errors are Composite Compliance trying to catch?] + except RuntimeError as err: + raise_composite_compliance_error( + err, + f"- wrapped_args: {which_args_are_wrapped}\n" + f"- wrapped_kwargs: {which_kwargs_are_wrapped}\n" + f"- wrapped_grads: {which_grad_is_batched}\n" + ) + + def unwrap(e): + return e.elem if isinstance(e, CCT) else e + + assert_equal_fn(tuple(map(unwrap, actual)), expected, equal_nan=True) + +# Checks if the forward AD formula is composite compliant by testing +# all possible permutations of {primals, tangents} being +# CompositeCompliantTensor or regular Tensors. +# +# NB: it is important that op is accepted as a Callable and not an OpInfo, +# this means we can apply check_forward_ad_formula to things that aren't OpInfos +# while debugging. +def check_forward_ad_formula(op: Callable, args, kwargs, gradcheck_wrapper=None, assert_equal_fn=None): + CCT, cct_mode = generate_cct_and_mode(autograd_view_consistency=False) + + def maybe_tangent(t): + if type(t) is CCT: + raise AssertionError("Expected type(t) is not CCT") + # Generate `tangent` tensor + # if given object is a Tensor and requires grad is set. + if isinstance(t, torch.Tensor) and t.requires_grad: + return torch.randn_like(t) + elif is_tensorlist(t): + return [torch.randn_like(e) if e.requires_grad else None for e in t] + return None + + tangent_args = tuple(maybe_tangent(arg) for arg in args) + flat_kwargs, spec = tree_flatten(kwargs) + flat_tangent_kwargs = tuple(maybe_tangent(arg) for arg in flat_kwargs) + tangent_kwargs = tree_unflatten(flat_tangent_kwargs, spec) + + with fwAD.dual_level(): + def maybe_make_dual(dual): + # Returns dual tensor if primal is a tensor/tensor subclass + # with requires_grad set. + primal, tangent = dual + if isinstance(primal, torch.Tensor) and primal.requires_grad: + return fwAD.make_dual(primal.detach(), tangent) + elif is_tensorlist(primal): + return tuple(fwAD.make_dual(pri.detach(), tang) if tang is not None else pri + for pri, tang in zip(primal, tangent, strict=True)) + return primal + + def compute_expected_grad(args, tangent_args, kwargs, tangent_kwargs): + op_args = tuple(map(maybe_make_dual, zip(args, tangent_args, strict=True))) + op_kwargs = {k: maybe_make_dual((v, tangent_kwargs[k])) for k, v in kwargs.items()} + + if gradcheck_wrapper is None: + return op(*op_args, **op_kwargs) + return gradcheck_wrapper(op, *op_args, **op_kwargs) + + expected = compute_expected_grad(args, tangent_args, kwargs, tangent_kwargs) + expected = tree_map(fwAD.unpack_dual, expected) + expected_primals = tree_map( + lambda x: x.primal, + expected, + is_leaf=lambda x: type(x) is fwAD.UnpackedDualTensor, + ) + expected_tangents = tree_map( + lambda x: x.tangent, + expected, + is_leaf=lambda x: type(x) is fwAD.UnpackedDualTensor, + ) + + # Permutations of arg and kwargs in CCT. + for choice in generate_subclass_choices_args_kwargs(args, kwargs, CCT, cct_mode): + new_args, new_kwargs, which_args_are_wrapped, which_kwargs_are_wrapped = choice + + # Permutations tangent arg and tangent kwargs in CCT. + for tang_choice in generate_subclass_choices_args_kwargs(tangent_args, tangent_kwargs, CCT, cct_mode): + new_tang_args, new_tang_kwargs, \ + which_tang_args_are_wrapped, which_tang_kwargs_are_wrapped = tang_choice + + op_args = tuple(map(maybe_make_dual, zip(new_args, new_tang_args, strict=True))) + op_kwargs = {k: maybe_make_dual((v, new_tang_kwargs[k])) for k, v in new_kwargs.items()} + + try: + if gradcheck_wrapper is None: + actual = op(*op_args, **op_kwargs) + else: + actual = gradcheck_wrapper(op, *op_args, **op_kwargs) + # see NOTE: [What errors are Composite Compliance trying to catch?] + except RuntimeError as err: + raise_composite_compliance_error( + err, + f"- wrapped_args: {which_args_are_wrapped}\n" + f"- wrapped_kwargs: {which_kwargs_are_wrapped}\n" + f"- wrapped_tangent_args: {which_tang_args_are_wrapped}\n" + f"- wrapped_tangent_kwargs: {which_tang_kwargs_are_wrapped}\n" + ) + + def unwrap(e): + return e.elem if isinstance(e, CCT) else e + + actual = tree_map(fwAD.unpack_dual, actual) + actual_primals = tree_map( + lambda x: unwrap(x.primal), + actual, + is_leaf=lambda x: type(x) is fwAD.UnpackedDualTensor, + ) + actual_tangents = tree_map( + lambda x: unwrap(x.tangent), + actual, + is_leaf=lambda x: type(x) is fwAD.UnpackedDualTensor, + ) + assert_equal_fn(actual_primals, expected_primals, equal_nan=True) + assert_equal_fn(actual_tangents, expected_tangents, equal_nan=True) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/custom_op_db.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/custom_op_db.py new file mode 100644 index 0000000000000000000000000000000000000000..d2b9931a84c4412941070369c43aa2fe4b5de6c3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/custom_op_db.py @@ -0,0 +1,603 @@ +# mypy: allow-untyped-defs +import torch +import functools +from torch.testing import make_tensor +from torch.testing._internal.opinfo.core import ( + OpInfo, + SampleInput, +) +from torch.testing._internal.common_dtype import all_types_and +import numpy as np +from torch.testing._internal.autograd_function_db import ( + sample_inputs_numpy_cube, + sample_inputs_numpy_mul, + sample_inputs_numpy_mul_scalar, + sample_inputs_numpy_sort, + sample_inputs_numpy_take, +) +from torch import Tensor +from torch.types import Number +from typing import * # noqa: F403 + +# Note: [custom op db] +# +# This is a collection of custom operator test cases written as OpInfos +# so they can easily be consumed by OpInfo-based tests to check if subsystems +# support them correctly. + +def to_numpy(tensor): + return tensor.cpu().numpy() + +@torch.library.custom_op("_torch_testing::numpy_cube", mutates_args=()) +def numpy_cube(x: Tensor) -> tuple[Tensor, Tensor]: + x_np = to_numpy(x) + dx = torch.tensor(3 * x_np ** 2, device=x.device) + return torch.tensor(x_np ** 3, device=x.device), dx + +@numpy_cube.register_fake +def _(x): + return x.clone(), x.clone() + +def numpy_cube_setup_context(ctx, inputs, output): + x, = inputs + _cube, dx = output + ctx.save_for_backward(x, dx) + +def numpy_cube_backward(ctx, grad_out, grad_dx): + x, dx = ctx.saved_tensors + grad_x = numpy_mul(grad_out, dx) + 6 * numpy_mul(grad_dx, x) + return grad_x + +numpy_cube.register_autograd(numpy_cube_backward, setup_context=numpy_cube_setup_context) + +def numpy_cube_vmap(info, in_dims, x): + result = numpy_cube(x) + return result, (in_dims[0], in_dims[0]) + +numpy_cube.register_vmap(numpy_cube_vmap) + +@torch.library.custom_op("_torch_testing::numpy_mul", mutates_args=()) +def numpy_mul(x: Tensor, y: Tensor) -> Tensor: + return torch.tensor(to_numpy(x) * to_numpy(y), device=x.device) + +@numpy_mul.register_fake +def _(x, y): + if x.device != y.device: + raise AssertionError(f"x.device={x.device} != y.device={y.device}") + return (x * y).contiguous() + +def numpy_mul_setup_context(ctx, inputs, output): + ctx.save_for_backward(*inputs) + +def numpy_mul_backward(ctx, grad_out): + x, y = ctx.saved_tensors + grad_x = grad_out * y if ctx.needs_input_grad[0] else None + grad_y = grad_out * x if ctx.needs_input_grad[1] else None + return grad_x, grad_y + +numpy_mul.register_autograd(numpy_mul_backward, setup_context=numpy_mul_setup_context) + +def numpy_mul_vmap(info, in_dims, x, y): + x_bdim, y_bdim = in_dims + x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1) + y = y.movedim(y_bdim, -1) if y_bdim is not None else y.unsqueeze(-1) + result = x * y + result = result.movedim(-1, 0) + return result, 0 + +numpy_mul.register_vmap(numpy_mul_vmap) + +@torch.library.custom_op("_torch_testing::numpy_mul_scalar", mutates_args=()) +def numpy_mul_scalar(x: Tensor, *, scalar: float) -> Tensor: + return torch.tensor(to_numpy(x) * scalar, device=x.device) + +@numpy_mul_scalar.register_fake +def _(x, *, scalar): + return (x * scalar).contiguous() + +def numpy_mul_scalar_setup_context(ctx, inputs, keyword_only_inputs, output): + ctx.scalar = keyword_only_inputs["scalar"] + +def numpy_mul_scalar_backward(ctx, grad_out): + grad_x = grad_out * ctx.scalar + return grad_x + +numpy_mul_scalar.register_autograd(numpy_mul_scalar_backward, setup_context=numpy_mul_scalar_setup_context) + +def numpy_mul_scalar_vmap(info, in_dims, x, *, scalar): + x_bdim, = in_dims + x = x.movedim(x_bdim, -1) if x_bdim is not None else x.unsqueeze(-1) + result = x * scalar + result = result.movedim(-1, 0) + return result, 0 + +numpy_mul_scalar.register_vmap(numpy_mul_scalar_vmap) + +@torch.library.custom_op("_torch_testing::numpy_sort", mutates_args=()) +def numpy_sort(x: Tensor, dim: int) -> tuple[Tensor, Tensor, Tensor]: + device = x.device + x = to_numpy(x) + ind = np.argsort(x, axis=dim) + ind_inv = np.argsort(ind, axis=dim) + result = np.take_along_axis(x, ind, axis=dim) + return ( + torch.tensor(result, device=device), + torch.tensor(ind, device=device), + torch.tensor(ind_inv, device=device), + ) + +@numpy_sort.register_fake +def _(x, dim): + return torch.empty_like(x), torch.empty_like(x, dtype=torch.long), torch.empty_like(x, dtype=torch.long) + +def numpy_sort_setup_context(ctx, inputs, output): + _out, ind, ind_inv = output + ctx.dim = inputs[1] + ctx.save_for_backward(ind, ind_inv) + ctx.mark_non_differentiable(ind, ind_inv) + +def numpy_sort_backward(ctx, grad_out, grad_ind, grad_ind_inv): + ind, ind_inv = ctx.saved_tensors + return numpy_take(grad_out, ind_inv, ind, ctx.dim), None + +numpy_sort.register_autograd(numpy_sort_backward, setup_context=numpy_sort_setup_context) + +def numpy_sort_vmap(info, in_dims, x, dim): + x_bdim, _ = in_dims + x = x.movedim(x_bdim, 0) + dim = dim if dim >= 0 else dim + x.dim() - 1 + result = numpy_sort(x, dim + 1) + return result, (0, 0, 0) + +numpy_sort.register_vmap(numpy_sort_vmap) + +@torch.library.custom_op("_torch_testing::numpy_take", mutates_args=()) +def numpy_take(x: Tensor, ind: Tensor, ind_inv: Tensor, dim: int) -> Tensor: + device = x.device + x = to_numpy(x) + ind = to_numpy(ind) + return torch.tensor(np.take_along_axis(x, ind, dim), device=device) + +@numpy_take.register_fake +def _(x, ind, ind_inv, dim): + if x.device != ind.device: + raise AssertionError(f"x.device={x.device} != ind.device={ind.device}") + if x.device != ind_inv.device: + raise AssertionError(f"x.device={x.device} != ind_inv.device={ind_inv.device}") + if ind.dtype != torch.long: + raise AssertionError(f"ind.dtype must be torch.long, got {ind.dtype}") + if ind_inv.dtype != torch.long: + raise AssertionError(f"ind_inv.dtype must be torch.long, got {ind_inv.dtype}") + return torch.empty_like(x) + +def numpy_take_setup_context(ctx, inputs, output): + _x, ind, ind_inv, dim = inputs + ctx.dim = dim + ctx.save_for_backward(ind, ind_inv) + +def numpy_take_backward(ctx, grad_out): + ind, ind_inv = ctx.saved_tensors + grad_x = numpy_take(grad_out, ind_inv, ind, ctx.dim) + return grad_x, None, None, None + +numpy_take.register_autograd(numpy_take_backward, setup_context=numpy_take_setup_context) + +def numpy_take_vmap(info, in_dims, x, ind, ind_inv, dim): + x_bdim, ind_bdim, ind_inv_bdim, _ = in_dims + + # wrap dim + logical_dim = x.dim() if x_bdim is None else x_bdim - 1 + dim = dim if dim >= 0 else dim + logical_dim + + def expand_bdim(x, x_bdim): + if x_bdim is None: + return x.expand(info.batch_size, *x.shape) + return x.movedim(x_bdim, 0) + + x = expand_bdim(x, x_bdim) + ind = expand_bdim(ind, ind_bdim) + ind_inv = expand_bdim(ind_inv, ind_inv_bdim) + + return numpy_take(x, ind, ind_inv, dim + 1), 0 + +numpy_take.register_vmap(numpy_take_vmap) + +@torch.library.custom_op("_torch_testing::numpy_nonzero", mutates_args=()) +def numpy_nonzero(x: Tensor) -> Tensor: + x_np = to_numpy(x) + res = np.stack(np.nonzero(x_np), axis=1) + if res.shape[0] <= 1: + raise RuntimeError("not supported") + return torch.tensor(res, device=x.device) + +@numpy_nonzero.register_fake +def _(x): + ctx = torch._custom_op.impl.get_ctx() + i0 = ctx.create_unbacked_symint() + shape = [i0, x.dim()] + result = x.new_empty(shape, dtype=torch.long) + return result + +def sample_inputs_numpy_nonzero(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + shape = 10 + result = make_arg(shape, low=0.9, high=2) + mask = make_tensor(shape, low=0, high=2, device=device, dtype=torch.long) + with torch.no_grad(): + result *= mask + + yield SampleInput(result, args=()) + +def numpy_nonzero_vmap(info, in_dims, x): + raise NotImplementedError("Operator is data-dependent and cannot be vmapped.") + +numpy_nonzero.register_vmap(numpy_nonzero_vmap) + +@torch.library.custom_op("_torch_testing::numpy_view_copy", mutates_args=()) +def numpy_view_copy(x: Tensor, shape: Sequence[int]) -> Tensor: + return torch.tensor(np.copy(to_numpy(x).reshape(shape)), device=x.device) + +@numpy_view_copy.register_fake +def _(x, shape) -> Tensor: + return x.clone().view(shape).clone() + +def numpy_view_copy_setup_context(ctx, inputs, output) -> None: + ctx.x_shape = inputs[0].shape + +def numpy_view_copy_backward(ctx, grad_out): + return torch.ops._torch_testing.numpy_view_copy(grad_out, ctx.x_shape), None + +numpy_view_copy.register_autograd(numpy_view_copy_backward, setup_context=numpy_view_copy_setup_context) + +def numpy_view_copy_vmap(info, in_dims, x, shape): + x_bdim, _ = in_dims + x = x.movedim(x_bdim, 0) + x_shape = x.shape[0] + batch_shape = (x_shape, *shape) + result = numpy_view_copy(x, batch_shape) + return result, 0 + +numpy_view_copy.register_vmap(numpy_view_copy_vmap) + +def sample_inputs_numpy_view_copy(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + result = make_arg(2, 3, 4, low=0.9, high=2) + yield SampleInput(result, args=([2, 12],)) + +@torch.library.custom_op('_torch_testing::numpy_cat', mutates_args=()) +def numpy_cat(xs: Sequence[Tensor], dim: int) -> Tensor: + if len(xs) == 0: + raise AssertionError("xs must not be empty") + if not all(x.device == xs[0].device for x in xs): + raise AssertionError("All tensors must be on the same device") + if not all(x.dtype == xs[0].dtype for x in xs): + raise AssertionError("All tensors must have the same dtype") + np_xs = [to_numpy(x) for x in xs] + np_out = np.concatenate(np_xs, axis=dim) + return torch.tensor(np_out, device=xs[0].device) + +@numpy_cat.register_fake +def _(xs, dim): + if len(xs) == 0: + raise AssertionError("xs must not be empty") + if not all(x.device == xs[0].device for x in xs): + raise AssertionError("All tensors must be on the same device") + if not all(x.dtype == xs[0].dtype for x in xs): + raise AssertionError("All tensors must have the same dtype") + return torch.cat(xs, dim=dim) + +def numpy_cat_setup_context(ctx, inputs, output): + xs, dim = inputs + ctx.dim_sizes = [x.shape[dim] for x in xs] + ctx.dim = dim + +def numpy_cat_backward(ctx, grad_out): + dim_sizes = ctx.dim_sizes + dim = ctx.dim + + splits = list(np.cumsum(dim_sizes)[:-1]) + grad_xs = torch.ops._torch_testing.numpy_split_copy(grad_out, splits, dim) + return grad_xs, None + +numpy_cat.register_autograd(numpy_cat_backward, setup_context=numpy_cat_setup_context) + +def numpy_cat_vmap(info, in_dims, x, dim): + x_bdim, = in_dims + result = numpy_cat(x, dim) + return result, x_bdim + +numpy_cat.register_vmap(numpy_cat_vmap) + +def sample_inputs_numpy_cat(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + r0 = make_arg(2, 3, 4, low=0.9, high=2) + r1 = make_arg(4, 3, 4, low=0.9, high=2) + r2 = make_arg(5, 3, 4, low=0.9, high=2) + yield SampleInput([r0, r1, r2], args=(0,)) + +@torch.library.custom_op('_torch_testing::numpy_split_copy', mutates_args=()) +def numpy_split_copy(x: Tensor, splits: Sequence[int], dim: int) -> List[Tensor]: + x_np = to_numpy(x) + arrs = np.split(x_np, splits, axis=dim) + return [torch.tensor(arr, device=x.device, dtype=x.dtype) for arr in arrs] + +@numpy_split_copy.register_fake +def _(x, splits, dim): + return [xi.clone() for xi in torch.tensor_split(x, splits, dim)] + +def numpy_split_copy_setup_context(ctx, inputs, output): + _, _, dim = inputs + ctx.dim = dim + +def numpy_split_copy_backward(ctx, grad_out): + result = torch.ops._torch_testing.numpy_cat(grad_out, dim=ctx.dim) + return result, None, None + +numpy_split_copy.register_autograd(numpy_split_copy_backward, setup_context=numpy_split_copy_setup_context) + +def numpy_split_copy_vmap(info, in_dims, x, splits, dim): + x_bdim, _ , _ = in_dims + x = x.movedim(x_bdim, 0) + result = numpy_split_copy(x, splits, dim + 1) + return result, 0 + +numpy_split_copy.register_vmap(numpy_split_copy_vmap) + +def sample_inputs_numpy_split_copy(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial(make_tensor, device=device, dtype=dtype, requires_grad=requires_grad) + x = make_arg(2, 9, low=0.9, high=2) + yield SampleInput(x, args=([1, 3, 6], 1)) + +@torch.library.custom_op('_torch_testing::numpy_split_copy_with_int', mutates_args=()) +def numpy_split_copy_with_int(x: Tensor, splits: Sequence[int], dim: int) -> tuple[List[Tensor], int]: + x_np = to_numpy(x) + arrs = np.split(x_np, splits, axis=dim) + return [torch.tensor(arr, device=x.device, dtype=x.dtype) for arr in arrs], len(splits) + +@numpy_split_copy_with_int.register_fake +def _(x, splits, dim): + return [xi.clone() for xi in torch.tensor_split(x, splits, dim)], len(splits) + +def numpy_split_copy_with_int_setup_context(ctx, inputs, output): + _, _, dim = inputs + ctx.dim = dim + +def numpy_split_copy_with_int_backward(ctx, grad_out, _): + return torch.ops._torch_testing.numpy_cat(grad_out, dim=ctx.dim), None, None + +numpy_split_copy_with_int.register_autograd( + numpy_split_copy_with_int_backward, + setup_context=numpy_split_copy_with_int_setup_context) + +def numpy_split_copy_with_int_vmap(info, in_dims, x, splits, dim): + x_bdim, _ , _ = in_dims + x = x.movedim(x_bdim, 0) + result, len_split = numpy_split_copy_with_int(x, splits, dim + 1) + return (result, len_split), ([0 for _ in range(len(result))], None) + +numpy_split_copy_with_int.register_vmap(numpy_split_copy_with_int_vmap) + +@torch.library.custom_op("_torch_testing::numpy_nms", mutates_args=()) +def numpy_nms(boxes: Tensor, scores: Tensor, iou_threshold: Number) -> Tensor: + # Adapted from Ross Girshick's fast-rcnn implementation at + # https://github.com/rbgirshick/fast-rcnn/blob/master/lib/utils/nms.py + if boxes.device != scores.device: + raise AssertionError(f"boxes.device={boxes.device} != scores.device={scores.device}") + device = boxes.device + + boxes = to_numpy(boxes) + scores = to_numpy(scores) + + N = boxes.shape[0] + if boxes.shape != (N, 4): + raise AssertionError(f"boxes.shape must be (N, 4), got {boxes.shape}") + if scores.shape != (N,): + raise AssertionError(f"scores.shape must be (N,), got {scores.shape}") + + x1 = boxes[:, 0] + y1 = boxes[:, 1] + x2 = boxes[:, 2] + y2 = boxes[:, 3] + + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + order = scores.argsort()[::-1] + + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + xx1 = np.maximum(x1[i], x1[order[1:]]) + yy1 = np.maximum(y1[i], y1[order[1:]]) + xx2 = np.minimum(x2[i], x2[order[1:]]) + yy2 = np.minimum(y2[i], y2[order[1:]]) + + w = np.maximum(0.0, xx2 - xx1 + 1) + h = np.maximum(0.0, yy2 - yy1 + 1) + inter = w * h + ovr = inter / (areas[i] + areas[order[1:]] - inter) + + inds = np.where(ovr <= iou_threshold)[0] + order = order[inds + 1] + + result = torch.tensor(np.stack(keep), device=device) + # Needed for data-dependent condition :( + if result.size(0) < 2: + raise AssertionError(f"result.size(0) must be >= 2, got {result.size(0)}") + return result + +@numpy_nms.register_fake +def _(boxes, scores, iou_threshold): + if boxes.device != scores.device: + raise AssertionError(f"boxes.device={boxes.device} != scores.device={scores.device}") + N = boxes.shape[0] + if boxes.shape != (N, 4): + raise AssertionError(f"boxes.shape must be (N, 4), got {boxes.shape}") + if scores.shape != (N,): + raise AssertionError(f"scores.shape must be (N,), got {scores.shape}") + + ctx = torch._custom_op.impl.get_ctx() + i0 = ctx.create_unbacked_symint() + result = boxes.new_empty([i0], dtype=torch.int64) + return result + +def numpy_nms_vmap(info, in_dims, boxes, scores, iou_threshold): + raise NotImplementedError("Operator is data-dependent and cannot be vmapped.") + +numpy_nms.register_vmap(numpy_nms_vmap) + +def sample_inputs_numpy_nms(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial(make_tensor, device=device, dtype=dtype) + N = 64 + xs = make_arg([N], low=0, high=28) + dx = make_arg([N], low=0, high=4) + ys = make_arg([N], low=0, high=28) + dy = make_arg([N], low=0, high=4) + boxes = torch.stack([xs, ys, xs + dx, ys + dy], dim=1).requires_grad_(requires_grad) + scores = make_arg([N], low=0, high=1, requires_grad=requires_grad) + iou_threshold = make_arg([], low=0, high=1).item() + + yield SampleInput(boxes, args=(scores, iou_threshold)) + +custom_op_db = [ + OpInfo( + 'NumpyCubeCustomOp', + op=numpy_cube._opoverload, + sample_inputs_func=sample_inputs_numpy_cube, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'NumpyMulCustomOp', + op=numpy_mul._opoverload, + sample_inputs_func=sample_inputs_numpy_mul, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'NumpyMulScalarCustomOp', + op=numpy_mul_scalar._opoverload, + sample_inputs_func=sample_inputs_numpy_mul_scalar, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'NumpySortCustomOp', + op=numpy_sort._opoverload, + sample_inputs_func=sample_inputs_numpy_sort, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'NumpyTakeCustomOp', + op=numpy_take._opoverload, + sample_inputs_func=sample_inputs_numpy_take, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + ), + OpInfo( + 'NumpyNonzeroCustomOp', + op=numpy_nonzero._opoverload, + sample_inputs_func=sample_inputs_numpy_nonzero, + dtypes=all_types_and(torch.bool, torch.half), + supports_autograd=False, + supports_out=False, + ), + OpInfo( + 'NumpyNMSCustomOp', + op=torch.ops._torch_testing.numpy_nms, + sample_inputs_func=sample_inputs_numpy_nms, + dtypes=all_types_and(torch.bool, torch.half), + supports_autograd=False, + supports_out=False, + ), + OpInfo( + 'NumpyViewCopyCustomOp', + op=torch.ops._torch_testing.numpy_view_copy, + sample_inputs_func=sample_inputs_numpy_view_copy, + dtypes=all_types_and(torch.bool, torch.half), + supports_autograd=True, + supports_out=False, + ), + OpInfo( + 'NumpyCatCustomOp', + op=torch.ops._torch_testing.numpy_cat, + sample_inputs_func=sample_inputs_numpy_cat, + dtypes=all_types_and(torch.bool, torch.half), + supports_autograd=True, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_out=False, + ), + OpInfo( + 'NumpySplitCopyCustomOp', + op=torch.ops._torch_testing.numpy_split_copy, + sample_inputs_func=sample_inputs_numpy_split_copy, + dtypes=all_types_and(torch.bool, torch.half), + supports_autograd=True, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_out=False, + ), + OpInfo( + 'NumpySplitCopyWithIntCustomOp', + op=torch.ops._torch_testing.numpy_split_copy_with_int, + sample_inputs_func=sample_inputs_numpy_split_copy, + dtypes=all_types_and(torch.bool, torch.half), + gradcheck_wrapper=lambda op, *args, **kwargs: op(*args, **kwargs)[0], + supports_autograd=True, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_out=False, + ), +] + + +# ============================================================== +# some mechanical test cases +# ============================================================== + +lib = torch.library.Library("_torch_testing", "FRAGMENT") # noqa: TOR901 + +lib.define("source0(Tensor x) -> Tensor") + +@torch.library.register_fake("_torch_testing::source0", lib=lib) +def _(x): + return x.clone() + +lib.define("source1(Tensor x) -> Tensor") + +def source1_fake(x): + return x.clone() + +torch.library.register_fake("_torch_testing::source1", source1_fake, lib=lib) + +lib.define("source2(Tensor x) -> Tensor") + +@torch.library.register_fake("_torch_testing::source2", lib=lib) +def _(x): + return x.clone() + +lib.define("source3(Tensor x) -> Tensor") + +def source3_fake(x): + return x.clone() + +torch.library.register_fake("_torch_testing::source3", source3_fake, lib=lib) + + +@torch.library.custom_op("_torch_testing::source4", mutates_args=()) +def source4(x: Tensor) -> Tensor: + return x.clone() + +@source4.register_fake +def _(x): + return x.clone() + +@torch.library.custom_op("_torch_testing::source5", mutates_args=()) +def source5(x: Tensor) -> Tensor: + return x.clone() + +def source5_fake(x): + return x.clone() + +source5.register_fake(source5_fake) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/custom_tensor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/custom_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..8da079d34eed925c9fcad73e5e815bd3dfb1beed --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/custom_tensor.py @@ -0,0 +1,161 @@ +# mypy: ignore-errors + + +from collections import namedtuple + +import torch +import torch.utils._pytree as pytree +from torch.utils._python_dispatch import return_and_correct_aliasing + + +FancyNamedTuple = namedtuple("FancyNamedTuple", ["foo", "bar"]) + + +# A simple tensor subclass that holds a tensor with custom metadata and custom method +class ConstantExtraMetadataTensor(torch.Tensor): + @staticmethod + def __new__(cls, elem): + shape = elem.shape + kwargs = {} + kwargs["strides"] = elem.stride() + kwargs["storage_offset"] = elem.storage_offset() + kwargs["device"] = elem.device + kwargs["layout"] = elem.layout + kwargs["requires_grad"] = elem.requires_grad + kwargs["dtype"] = elem.dtype + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) + + def __init__(self, elem): + self.elem = elem + self.constant_attribute = 4 + + def __repr__(self): + inner_repr = repr(self.elem) + return f"CustomTensor({inner_repr})" + + def get_complicated_metadata(self): + return FancyNamedTuple(self.constant_attribute, self.constant_attribute) + + def __tensor_flatten__(self): + return ["elem"], self.constant_attribute + + def add_constant(self, a): + self.constant_attribute += a + + @staticmethod + def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride): + if meta is None: + raise AssertionError("Expected meta to not be None") + elem = inner_tensors["elem"] + out = ConstantExtraMetadataTensor(elem) + out.constant_attribute = meta + return out + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + if kwargs is None: + kwargs = {} + args_inner = pytree.tree_map_only( + ConstantExtraMetadataTensor, lambda x: x.elem, args + ) + + kwargs_inner = pytree.tree_map_only( + ConstantExtraMetadataTensor, lambda x: x.elem, kwargs + ) + + out_inner = func(*args_inner, **kwargs_inner) + out_inner_flat, spec = pytree.tree_flatten(out_inner) + # for aten ops that return non-tensors, just assume that + # our cust inner tensors return the same value + out_flat = [ + ConstantExtraMetadataTensor(o_inner) + if isinstance(o_inner, torch.Tensor) + else o_inner + for o_inner in out_inner_flat + ] + out = pytree.tree_unflatten(out_flat, spec) + return return_and_correct_aliasing(func, args, kwargs, out) + + +# A simple tensor subclass that always returns plain tensor during __torch_dispatch__ +# It is similar to TwoTensor and is used to simulate torchao quantized tensors +class CustomTensorPlainOut(torch.Tensor): + @staticmethod + def __new__(cls, elem1, elem2): + shape = elem1.shape + kwargs = {} + kwargs["strides"] = elem1.stride() + kwargs["storage_offset"] = elem1.storage_offset() + kwargs["device"] = elem1.device + kwargs["layout"] = elem1.layout + kwargs["requires_grad"] = elem1.requires_grad + kwargs["dtype"] = elem1.dtype + return torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) + + def __init__(self, elem1, elem2): + self.elem1 = elem1 + self.elem2 = elem2 + + def get_elem(self): + return self.elem1 + + def __repr__(self): + inner_repr_1 = repr(self.elem1) + inner_repr_2 = repr(self.elem2) + return f"CustomTensorPlainOut({inner_repr_1}, {inner_repr_2})" + + def __tensor_flatten__(self): + return ["elem1", "elem2"], None + + @staticmethod + def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride): + elem1 = inner_tensors["elem1"] + elem2 = inner_tensors["elem2"] + out = CustomTensorPlainOut(elem1, elem2) + return out + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + # Don't use this tensor with view ops + if kwargs is None: + kwargs = {} + args_inner_1 = pytree.tree_map_only( + CustomTensorPlainOut, lambda x: x.elem1, args + ) + + kwargs_inner_1 = pytree.tree_map_only( + CustomTensorPlainOut, lambda x: x.elem1, kwargs + ) + + args_inner_2 = pytree.tree_map_only( + CustomTensorPlainOut, lambda x: x.elem2, args + ) + + kwargs_inner_2 = pytree.tree_map_only( + CustomTensorPlainOut, lambda x: x.elem2, kwargs + ) + + out_inner_1 = func(*args_inner_1, **kwargs_inner_1) + out_inner_2 = func(*args_inner_2, **kwargs_inner_2) + + out_inner_flat_1, spec = pytree.tree_flatten(out_inner_1) + out_inner_flat_2, spec = pytree.tree_flatten(out_inner_2) + + if func.is_view: + new_out = pytree.tree_unflatten( + ( + CustomTensorPlainOut(tensor1, tensor2) + for tensor1, tensor2 in zip( + out_inner_flat_1, out_inner_flat_2, strict=True + ) + ), + spec, + ) + return return_and_correct_aliasing(func, args, kwargs, new_out) + + out_new = ( + out_inner_flat_1[ix] + out_inner_flat_2[ix] + for ix in range(len(out_inner_flat_1)) + ) + + return pytree.tree_unflatten(out_new, spec) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/data/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1e3572cfc4c6a0ddc3d8fa2e1b056415204acdfa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/data/__init__.py @@ -0,0 +1 @@ +# mypy: ignore-errors diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/data/network1.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/data/network1.py new file mode 100644 index 0000000000000000000000000000000000000000..8755643a78cca80668988df9e9db3de75778b5db --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/data/network1.py @@ -0,0 +1,10 @@ +# mypy: ignore-errors + +import torch.nn as nn + + +class Net(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(10, 20) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/data/network2.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/data/network2.py new file mode 100644 index 0000000000000000000000000000000000000000..19b0b8ee53d3b530aa33978c7a13da4e5fee4ebd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/data/network2.py @@ -0,0 +1,11 @@ +# mypy: ignore-errors + +import torch.nn as nn + + +class Net(nn.Module): + + def __init__(self) -> None: + super().__init__() + self.linear = nn.Linear(10, 20) + self.relu = nn.ReLU() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/dist_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/dist_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..45af2552cf25cef03a517f5b136c1a2e61c3a61d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/dist_utils.py @@ -0,0 +1,199 @@ +# mypy: ignore-errors + +import re +import sys +import time +from functools import partial, wraps + +import torch.distributed as dist +import torch.distributed.rpc as rpc +from torch.distributed.rpc import _rref_context_get_debug_info +from torch.testing._internal.common_utils import FILE_SCHEMA, TEST_WITH_TSAN + + +if not dist.is_available(): + print("c10d not available, skipping tests", file=sys.stderr) + sys.exit(0) + + +INIT_METHOD_TEMPLATE = FILE_SCHEMA + "{file_name}" + +def dist_init( + old_test_method=None, + setup_rpc: bool = True, + clean_shutdown: bool = True, + faulty_messages=None, + messages_to_delay=None, +): + """ + We use this decorator for setting up and tearing down state since + MultiProcessTestCase runs each `test*` method in a separate process and + each process just runs the `test*` method without actually calling + 'setUp' and 'tearDown' methods of unittest. + + Note: pass the string representation of MessageTypes that should be used + with the faulty agent's send function. By default, all retriable messages + ("RREF_FORK_REQUEST", "RREF_CHILD_ACCEPT", "RREF_USER_DELETE", + "CLEANUP_AUTOGRAD_CONTEXT_REQ") will use the faulty send (this default is + set from faulty_rpc_agent_test_fixture.py). + """ + # If we use dist_init without arguments (ex: @dist_init), old_test_method is + # appropriately set and we return the wrapper appropriately. On the other + # hand if dist_init has arguments (ex: @dist_init(clean_shutdown=False)), + # old_test_method is None and we return a functools.partial which is the real + # decorator that is used and as a result we recursively call dist_init with + # old_test_method and the rest of the arguments appropriately set. + if old_test_method is None: + return partial( + dist_init, + setup_rpc=setup_rpc, + clean_shutdown=clean_shutdown, + faulty_messages=faulty_messages, + messages_to_delay=messages_to_delay, + ) + + @wraps(old_test_method) + def new_test_method(self, *arg, **kwargs): + # Setting _ignore_rref_leak to make sure OwnerRRefs are properly deleted + # in tests. + import torch.distributed.rpc.api as api + + api._ignore_rref_leak = False + self.worker_id = self.rank + self.setup_fault_injection(faulty_messages, messages_to_delay) + + rpc_backend_options = self.rpc_backend_options + if setup_rpc: + if TEST_WITH_TSAN: + # TSAN runs much slower. + rpc_backend_options.rpc_timeout = rpc.constants.DEFAULT_RPC_TIMEOUT_SEC * 5 + rpc.constants.DEFAULT_SHUTDOWN_TIMEOUT = 60 + + rpc.init_rpc( + name=f"worker{self.rank:d}", + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + + return_value = old_test_method(self, *arg, **kwargs) + + if setup_rpc: + rpc.shutdown(graceful=clean_shutdown) + + return return_value + + return new_test_method + + +def noop() -> None: + pass + + +def wait_until_node_failure(rank: int, expected_error_regex: str = ".*") -> str: + """ + Loops until an RPC to the given rank fails. This is used to + indicate that the node has failed in unit tests. + Args: + rank (int): Rank of the node expected to fail + expected_error_regex (optional, str): Regex of exception message expected. Useful to ensure a specific failure + occurs, not just any. + """ + while True: + try: + rpc.rpc_sync(f"worker{rank}", noop, args=()) + time.sleep(0.1) + except Exception as e: + if re.search(pattern=expected_error_regex, string=str(e)): + return str(e) + + +def wait_until_pending_futures_and_users_flushed(timeout: int = 20) -> None: + """ + The RRef protocol holds forkIds of rrefs in a map until those forks are + confirmed by the owner. The message confirming the fork may arrive after + our tests check whether this map is empty, which leads to failures and + flaky tests. to_here also does not guarantee that we have finished + processind the owner's confirmation message for the RRef. This function + loops until the map is empty, which means the messages have been received + as processed. Call this function before asserting the map returned by + _get_debug_info is empty. + """ + start = time.time() + while True: + debug_info = _rref_context_get_debug_info() + num_pending_futures = int(debug_info["num_pending_futures"]) + num_pending_users = int(debug_info["num_pending_users"]) + if num_pending_futures == 0 and num_pending_users == 0: + break + time.sleep(0.1) + if time.time() - start > timeout: + raise ValueError( + f"Timed out waiting to flush pending futures and users, " + f"had {num_pending_futures} pending futures and {num_pending_users} pending users" + ) + + +def get_num_owners_and_forks() -> tuple[str, str]: + """ + Retrieves number of OwnerRRefs and forks on this node from + _rref_context_get_debug_info. + """ + rref_dbg_info = _rref_context_get_debug_info() + num_owners = rref_dbg_info["num_owner_rrefs"] + num_forks = rref_dbg_info["num_forks"] + return num_owners, num_forks + + +def wait_until_owners_and_forks_on_rank( + num_owners: int, num_forks: int, rank: int, timeout: int = 20 +) -> None: + """ + Waits until timeout for num_forks and num_owners to exist on the rank. Used + to ensure proper deletion of RRefs in tests. + """ + start = time.time() + while True: + num_owners_on_rank, num_forks_on_rank = rpc.rpc_sync( + worker_name(rank), get_num_owners_and_forks, args=(), timeout=5 + ) + num_owners_on_rank = int(num_owners_on_rank) + num_forks_on_rank = int(num_forks_on_rank) + if num_owners_on_rank == num_owners and num_forks_on_rank == num_forks: + return + time.sleep(1) + if time.time() - start > timeout: + raise ValueError( + f"Timed out waiting {timeout} sec for {num_owners} owners and {num_forks} forks on rank," + f" had {num_owners_on_rank} owners and {num_forks_on_rank} forks" + ) + + +def initialize_pg(init_method, rank: int, world_size: int) -> None: + # This is for tests using `dist.barrier`. + if not dist.is_initialized(): + dist.init_process_group( + backend="gloo", + init_method=init_method, + rank=rank, + world_size=world_size, + ) + + +def worker_name(rank: int) -> str: + return f"worker{rank}" + + +def get_function_event(function_events, partial_event_name): + """ + Returns the first event that matches partial_event_name in the provided + function_events. These function_events should be the output of + torch.autograd.profiler.function_events(). + + Args: + function_events: function_events returned by the profiler. + event_name (str): partial key that the event was profiled with. + """ + event = [event for event in function_events if partial_event_name in event.name][0] # noqa: RUF015 + return event diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..acc7005c6b9e3d64d1ca50714839b0732d41b5a5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/__init__.py @@ -0,0 +1 @@ +# mypy: allow-untyped-defs diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..60c744ac1a84cfb9220221a583a4849b6039c353 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/__init__.py @@ -0,0 +1,103 @@ +# mypy: allow-untyped-defs + +import sys +from functools import partial, wraps + +import torch +import torch.distributed as dist +from torch.distributed import rpc +from torch.testing._internal.common_distributed import ( + MultiProcessTestCase, + TEST_SKIPS, + tp_transports, +) + + +TEST_GPU_NUM = 4 + + +class ShardedTensorTestBase(MultiProcessTestCase): + @property + def world_size(self): + return TEST_GPU_NUM + + def init_pg(self, backend="nccl"): + if backend not in ["nccl", "gloo", "mpi", "hccl"]: + raise RuntimeError(f"Backend {backend} not supported!") + + dist.init_process_group( + backend=backend, + world_size=self.world_size, + rank=self.rank, + init_method=f"file://{self.file_name}", + ) + + # set device for nccl pg for collectives + if backend == "nccl": + torch.cuda.set_device(self.rank) + + def init_rpc(self): + rpc_backend_options = rpc.TensorPipeRpcBackendOptions( + _transports=tp_transports() + ) + rpc_backend_options.init_method = f"file://{self.file_name}" + for rank in range(self.world_size): + rpc_backend_options.set_device_map( + f"worker{rank}", {rank: self.rank, self.rank: rank} + ) + + rpc.init_rpc( + name=f"worker{self.rank:d}", + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + + def init_comms(self, init_rpc=True, backend="nccl"): + if init_rpc: + self.init_rpc() + self.init_pg(backend=backend) + + def destroy_comms(self, destroy_rpc=True): + # Wait for all ranks to reach here before starting shutdown. + dist.barrier() + + if destroy_rpc: + rpc.shutdown() + dist.destroy_process_group() + + def setUp(self) -> None: + super().setUp() + self._spawn_processes() + + def assert_sharded_tensor_equal(self, st1, st2): + st1_local_shards = st1.local_shards() + st2_local_shards = st2.local_shards() + self.assertEqual(len(st1_local_shards), len(st2_local_shards)) + for i, st1_local_shard in enumerate(st1_local_shards): + self.assertEqual(st1_local_shard.tensor, st2_local_shards[i].tensor) + self.assertEqual(st1_local_shard.metadata, st2_local_shards[i].metadata) + + self.assertEqual(st1.metadata(), st2.metadata()) + self.assertEqual(st1.sharding_spec(), st2.sharding_spec()) + self.assertEqual(len(st1.remote_shards()), len(st2.remote_shards())) + + +# wrapper to initialize comms (processgroup + rpc) +def with_comms(func=None, init_rpc=True, backend="nccl"): + if func is None: + return partial( + with_comms, + init_rpc=init_rpc, + backend=backend, + ) + + @wraps(func) + def wrapper(self, *args, **kwargs): + if backend == "nccl" and torch.cuda.device_count() < self.world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code) + self.init_comms(init_rpc=init_rpc, backend=backend) + func(self, *args, **kwargs) + self.destroy_comms(destroy_rpc=init_rpc) + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_ops_common.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_ops_common.py new file mode 100644 index 0000000000000000000000000000000000000000..e83bc3a35102a051d42587352c2dcb7967510903 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_ops_common.py @@ -0,0 +1,137 @@ +# mypy: allow-untyped-defs + +import builtins + +import torch +from torch.distributed._shard.sharding_spec import ( + ChunkShardingSpec, + EnumerableShardingSpec, + ShardMetadata, +) +from torch.distributed._shard.sharding_spec._internals import ( + get_chunked_dim_size, + get_split_size, +) + + +def generate_chunk_sharding_specs_for_test(sharding_dim): + return [ + ChunkShardingSpec( + dim=sharding_dim, + placements=[ + "rank:0/cuda:0", + "rank:1/cuda:1", + "rank:2/cuda:2", + "rank:3/cuda:3", + ], + ), + # Test different ordering. (Case 1) + ChunkShardingSpec( + dim=sharding_dim, + placements=[ + "rank:2/cuda:2", + "rank:3/cuda:3", + "rank:0/cuda:0", + "rank:1/cuda:1", + ], + ), + # Test different ordering. (Case 2) + ChunkShardingSpec( + dim=sharding_dim, + placements=[ + "rank:3/cuda:3", + "rank:0/cuda:0", + "rank:1/cuda:1", + "rank:2/cuda:2", + ], + ), + ] + + +def generate_enumerable_sharding_specs_for_test(): + return [ + EnumerableShardingSpec( + [ + ShardMetadata( + shard_offsets=[0, 0], + shard_sizes=[5, 5], + placement="rank:0/cuda:0", + ), + ShardMetadata( + shard_offsets=[5, 0], + shard_sizes=[5, 5], + placement="rank:1/cuda:1", + ), + ShardMetadata( + shard_offsets=[0, 5], + shard_sizes=[5, 5], + placement="rank:2/cuda:2", + ), + ShardMetadata( + shard_offsets=[5, 5], + shard_sizes=[5, 5], + placement="rank:3/cuda:3", + ), + ] + ) + ] + + +def generate_local_weight_sharding_params_for_test( + local_weight, sharded_dim, gpu_num, spec, rank +): + """ + Shard the local weight based the given spec, so we can compare against + the one from sharded tensor. + + Args: + local_weight: weight matrix to be sharded. + sharded_dim: The dimension which we shard on. + gpu_num: number of ranks. + spec: sharding spec. + rank: # of cuda process. + + Returns: + start_pos: start position of sharded weight on the given rank. + chunk_size: chunk size of sharded weight on the given rank. + """ + sharding_dim_size = local_weight.size(sharded_dim) + split_size = get_split_size(sharding_dim_size, gpu_num) + current_offsets = 0 + start_pos = current_offsets + for idx, placement in enumerate(spec.placements): + chunk_size = get_chunked_dim_size(sharding_dim_size, split_size, idx) + if rank == placement.rank(): + start_pos = current_offsets + break + current_offsets += chunk_size + return start_pos, chunk_size + + +def clone_module_parameter(module, param_name): + """ + Clone a parameter from a given existing module. + + Args: + module (:class:`torch.nn.Module`): Module whose parameter needs to be cloned. + param_name (str): Name of the parameter of ``module`` that needs to be cloned. + + Returns: cloned tensor as :class:`torch.nn.Parameter`. + """ + tensor = getattr(module, param_name) + return torch.nn.Parameter(tensor.detach().clone()) + + +def gen_binary_op_func(python_op, inplace=False): + src_lines = ["def f(lhs, rhs):"] + if "torch" in python_op: + src_lines.append(f" return {python_op}(lhs, rhs)\n") + elif inplace: + src_lines.append(f" lhs {python_op}= rhs\n return lhs\n") + else: + src_lines.append(f" return lhs {python_op} rhs\n") + + code_str = "\n".join(src_lines) + g = {"torch": torch} + builtins.exec(code_str, g) + return g["f"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_st_common.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_st_common.py new file mode 100644 index 0000000000000000000000000000000000000000..1fe82a8dc43f8f876cb4c8d0c000cda9a32d46fb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/sharded_tensor/_test_st_common.py @@ -0,0 +1,56 @@ +# mypy: allow-untyped-defs + +import copy +import random + +import torch +from torch.distributed._shard import sharded_tensor +from torch.distributed._shard.sharding_spec import ChunkShardingSpec + + +PLACEMENTS = [ + "rank:0/cuda:0", + "rank:1/cuda:1", + "rank:2/cuda:2", + "rank:3/cuda:3", +] + +DEFAULT_GPU_NUM = 4 + + +def _chunk_sharding_specs_list_for_test(sharding_dims, seed=0): + spec_list = [] + for i in range(len(sharding_dims)): + random.Random(seed + i).shuffle(PLACEMENTS) + spec_list.append( + ChunkShardingSpec( + dim=sharding_dims[i], + placements=copy.deepcopy(PLACEMENTS), + ) + ) + return spec_list + + +class MyShardedModel2(torch.nn.Module): + def __init__(self, spec=None, group=None, init_rrefs=True) -> None: + super().__init__() + if spec is not None: + self.sharded_tensor2 = sharded_tensor.rand( + spec, 10, 20, process_group=group, init_rrefs=init_rrefs + ) + else: + self.sharded_tensor2 = None + self.random_tensor2 = torch.nn.Parameter(torch.rand(2, 2)) + + +class MyShardedModel1(torch.nn.Module): + def __init__(self, spec=None, group=None, init_rrefs=True) -> None: + super().__init__() + if spec is not None: + self.sharded_tensor1 = sharded_tensor.rand( + spec, 10, 20, process_group=group, init_rrefs=init_rrefs + ) + else: + self.sharded_tensor1 = None + self.random_tensor1 = torch.nn.Parameter(torch.rand(2, 2)) + self.submodule = MyShardedModel2(spec, group, init_rrefs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/test_common.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/test_common.py new file mode 100644 index 0000000000000000000000000000000000000000..e9390da489851872ec1d0715a0b3e46275e5752b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_shard/test_common.py @@ -0,0 +1,41 @@ +# mypy: allow-untyped-defs + +import torch +import torch.nn as nn +from torch.distributed._shard.sharded_tensor import ShardedTensor + + +class SimpleMegatronLM(nn.Module): + def __init__(self, linear_size, rank=None, dtype=torch.float32): + super().__init__() + self.fc1 = nn.Linear(*linear_size[0], dtype=dtype) + self.gelu = nn.GELU() + self.fc2 = nn.Linear(*linear_size[1], dtype=dtype) + if rank is not None: + self.fc1.cuda(rank) + self.fc2.cuda(rank) + + def forward(self, inp): + return self.fc2(self.gelu(self.fc1(inp))) + + def get_weights(self): + if isinstance(self.fc1.weight, ShardedTensor): + weight1 = self.fc1.weight.local_tensor() + else: + weight1 = self.fc1.weight + + if isinstance(self.fc2.weight, ShardedTensor): + weight2 = self.fc2.weight.local_tensor() + else: + weight2 = self.fc2.weight + + return (weight1, weight2) + + def get_biases(self): + return (self.fc1.bias, self.fc2.bias) + + def get_weight_grads(self): + return (self.fc1.weight.grad, self.fc2.weight.grad) + + def get_bias_grads(self): + return (self.fc1.bias.grad, self.fc2.bias.grad) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_tensor/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_tensor/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_tensor/common_dtensor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_tensor/common_dtensor.py new file mode 100644 index 0000000000000000000000000000000000000000..68c6baa863e98826e6bc74db40af1e09c77ab296 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/_tensor/common_dtensor.py @@ -0,0 +1,1509 @@ +# mypy: allow-untyped-defs + +# Copyright (c) Meta Platforms, Inc. and affiliates + +import contextlib +import copy +import functools +import itertools +import sys +import threading +import types +import unittest +from collections.abc import Callable, Iterator, Sequence +from dataclasses import dataclass +from functools import partial, wraps +from typing import Any, cast, TypeVar + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from torch.distributed._functional_collectives import ( + all_gather_tensor_autograd, + reduce_scatter_tensor_autograd, +) +from torch.distributed._local_tensor import ( + LocalIntNode, + LocalTensor, + LocalTensorMode, + maybe_disable_local_tensor_mode, + maybe_run_for_local_tensor, +) +from torch.distributed.tensor import ( + DeviceMesh, + distribute_module, + distribute_tensor, + DTensor, + init_device_mesh, + Partial, + Placement, + Replicate, + Shard, +) +from torch.distributed.tensor._dtensor_spec import ShardOrderEntry +from torch.distributed.tensor._redistribute import redistribute_local_tensor +from torch.distributed.tensor.parallel import ( + ColwiseParallel, + parallelize_module, + ParallelStyle, + PrepareModuleInput, + RowwiseParallel, + SequenceParallel, +) +from torch.testing._internal.common_distributed import ( + ACCELERATOR_DIST_BACKENDS, + MultiProcContinuousTest, + MultiProcessTestCase, + MultiThreadedTestCase, + run_subtests, + skip_if_lt_x_gpu, + TEST_SKIPS, +) +from torch.testing._internal.common_utils import ( + TEST_CUDA, + TEST_HPU, + TEST_PRIVATEUSE1, + TEST_WITH_ROCM, + TEST_XPU, +) +from torch.testing._internal.distributed.fake_pg import FakeStore +from torch.utils._pytree import tree_flatten, tree_unflatten, TreeSpec + + +DEVICE_COUNT: int + +if TEST_CUDA or TEST_XPU or TEST_HPU or TEST_PRIVATEUSE1: + DEVICE_TYPE = torch.accelerator.current_accelerator().type + DEVICE_COUNT = torch.accelerator.device_count() + PG_BACKEND = dist.Backend.default_device_backend_map[DEVICE_TYPE] +else: + DEVICE_TYPE = "cpu" + PG_BACKEND = "gloo" + +if TEST_WITH_ROCM: + NUM_DEVICES = min(4, max(2, torch.cuda.device_count())) +else: + NUM_DEVICES = 4 + +# We use this as a proxy for "multiple GPUs exist" +if (TEST_CUDA or TEST_XPU or TEST_HPU or TEST_PRIVATEUSE1) and DEVICE_COUNT > 1: + # when we actually have multiple GPUs, relax the requirement to smaller counts. + NUM_DEVICES = min(NUM_DEVICES, DEVICE_COUNT) + +T = TypeVar("T") + + +# simple RMSNorm layer for testing +class RMSNormPython(torch.nn.Module): + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.eps = eps + self.weight = torch.nn.Parameter(torch.ones(dim)) + + def _norm(self, x): + return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) + + def forward(self, x): + output = self._norm(x) + return output * self.weight + + +class MLPModule(nn.Module): + def __init__(self, device, bias: bool = True): + super().__init__() + torch.manual_seed(5) + self.net1 = nn.Linear(10, 16, bias=bias, device=device) + self.relu = nn.ReLU() + self.net2 = nn.Linear(16, 10, bias=bias, device=device) + + def forward(self, x): + return self.net2(self.relu(self.net1(x))) + + def reset_parameters(self): + self.net1.reset_parameters() + self.net2.reset_parameters() + + +class MLPStacked(nn.Module): + def __init__(self, device, n_layers: int = 2): + super().__init__() + self.layers = nn.ModuleList([MLPModule(device) for i in range(n_layers)]) + + def forward(self, x): + for layer in self.layers: + x = layer(x) + return x + + +@dataclass +class ModelArgs: + n_layers: int = 2 + vocab_size: int = 8 + max_seq_len: int = 16 + dim: int = 16 + n_heads: int = 4 + dropout_p: float = 0.1 + use_attn_mask: bool = True + weight_tying: bool = True + checkpoint_activations: bool = False + num_experts: int = 0 + + +class Attention(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + if args.dim % args.n_heads != 0: + raise AssertionError( + f"Expected args.dim % args.n_heads == 0, got {args.dim} % {args.n_heads}" + ) + self.head_dim = args.dim // args.n_heads + self.n_heads = args.n_heads + self.dropout_p = args.dropout_p + self.resid_dropout = nn.Dropout(args.dropout_p) + self.use_attn_mask = args.use_attn_mask + + self.wq = nn.Linear(args.dim, args.dim, bias=False) + self.wk = nn.Linear(args.dim, args.dim, bias=False) + self.wv = nn.Linear(args.dim, args.dim, bias=False) + self.wo = nn.Linear(args.dim, args.dim, bias=False) + + def forward(self, x): + bsz, seq_len, _ = x.size() + queries, keys, values = self.wq(x), self.wk(x), self.wv(x) + queries = queries.view(bsz, seq_len, self.n_heads, self.head_dim) + keys = keys.view(bsz, seq_len, self.n_heads, self.head_dim) + values = values.view(bsz, seq_len, self.n_heads, self.head_dim) + + queries = queries.transpose(1, 2) # (bsz, n_heads, seq_len, head_dim) + keys = keys.transpose(1, 2) # (bsz, n_heads, seq_len, head_dim) + values = values.transpose(1, 2) # (bsz, n_heads, seq_len, head_dim) + + output = F.scaled_dot_product_attention( + queries, + keys, + values, + None, + self.dropout_p if self.training else 0, + self.use_attn_mask, + ) + output = output.transpose(1, 2).contiguous().view(bsz, seq_len, -1) + return self.resid_dropout(self.wo(output)) + + +class FeedForward(nn.Module): + def __init__(self, dim, hidden_dim, dropout_p): + super().__init__() + self.w1 = nn.Linear(dim, hidden_dim) + self.gelu = nn.GELU() + self.w2 = nn.Linear(hidden_dim, dim) + self.resid_dropout = nn.Dropout(dropout_p) + + def forward(self, x): + return self.resid_dropout(self.w2(self.gelu(self.w1(x)))) + + +class Experts(nn.Module): + def __init__(self, dim: int, hidden_dim: int, num_experts: int): + super().__init__() + self.num_experts = num_experts + self.w1 = nn.Parameter(torch.empty(num_experts, hidden_dim, dim)) + self.w2 = nn.Parameter(torch.empty(num_experts, dim, hidden_dim)) + self.reset_parameters() + self.gelu = nn.GELU() + + def reset_parameters(self): + nn.init.normal_(self.w1, std=0.02) + nn.init.normal_(self.w2, std=0.02) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Weights are DTensors (sharded by EP/TP) but x is a plain tensor + # (dispatched by EP hooks), so extract local shards for bmm. + if isinstance(self.w1, DTensor): + w1, w2 = self.w1.to_local(), self.w2.to_local() + else: + w1, w2 = self.w1, self.w2 + E = w1.shape[0] + x_exp = x.unsqueeze(0).expand(E, -1, -1) + h = self.gelu(torch.bmm(x_exp, w1.transpose(-2, -1))) + out = torch.bmm(h, w2.transpose(-2, -1)) + return out.sum(dim=0) + + +class ExpertLayer(nn.Module): + def __init__(self, num_experts: int, dim: int, hidden_dim: int): + super().__init__() + self.num_experts = num_experts + self.experts = Experts(dim=dim, hidden_dim=hidden_dim, num_experts=num_experts) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + bs, slen, dim = x.shape + x_flat = x.view(-1, dim) + expert_out = self.experts(x_flat) + expert_out = expert_out / self.num_experts + return expert_out.view(bs, slen, dim) + + +class TensorParallelForExpert(ParallelStyle): + """TP for Experts: shard w1 colwise (Shard(1)), w2 rowwise (Shard(2)). + + For seq parallel, set input_layouts=Shard(0) and output_layouts=Shard(0) + to all-gather tokens before computation and reduce-scatter after. + """ + + def __init__(self, *, input_layouts=None, output_layouts=None): + super().__init__() + self.input_layouts = input_layouts + self.output_layouts = output_layouts or Replicate() + + def _partition_fn(self, name, mod, device_mesh): + for pn, p in mod.named_parameters(recurse=False): + if pn == "w1": + placement = [Shard(1)] + elif pn == "w2": + placement = [Shard(2)] + else: + continue + mod.register_parameter( + pn, nn.Parameter(distribute_tensor(p, device_mesh, placement)) + ) + + def _input_fn(self, mod, inputs, device_mesh): + if self.input_layouts is not None: + x = inputs[0] + x = ( + DTensor.from_local(x, device_mesh, [self.input_layouts]) + .redistribute(device_mesh, [Replicate()]) + .to_local() + ) + return (x,) + return inputs + + def _output_fn(self, mod, outputs, device_mesh): + return ( + DTensor.from_local(outputs, device_mesh, [Partial()]) + .redistribute(device_mesh, [self.output_layouts]) + .to_local() + ) + + def _apply(self, module, device_mesh): + return distribute_module( + module, + device_mesh, + partition_fn=self._partition_fn, + input_fn=self._input_fn, + output_fn=self._output_fn, + ) + + +class ExpertParallel(ParallelStyle): + """Distributes experts across ranks with Shard(0) on the expert dimension. + + Dispatch: all-gather tokens so every rank sees all tokens. + Combine: reduce-scatter (sum) expert outputs back to token owners. + """ + + def _partition_fn(self, name: str, mod: nn.Module, device_mesh: DeviceMesh) -> None: + for param_name, param in mod.named_parameters(recurse=False): + dist_param = nn.Parameter(distribute_tensor(param, device_mesh, [Shard(0)])) + mod.register_parameter(param_name, dist_param) + + def _token_dispatch( + self, mod: nn.Module, inputs: tuple, device_mesh: DeviceMesh + ) -> tuple[torch.Tensor]: + (x,) = inputs + x_gathered = all_gather_tensor_autograd( + x, + gather_dim=0, + group=device_mesh.get_group(), + ) + x_gathered = torch.ops._c10d_functional.wait_tensor(x_gathered) + return (x_gathered,) + + def _token_combine( + self, mod: nn.Module, output: torch.Tensor, device_mesh: DeviceMesh + ) -> torch.Tensor: + result = reduce_scatter_tensor_autograd( + output, + "sum", + scatter_dim=0, + group=device_mesh.get_group(), + ) + result = torch.ops._c10d_functional.wait_tensor(result) + return result + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + return distribute_module( + module, + device_mesh, + partition_fn=self._partition_fn, + input_fn=self._token_dispatch, + output_fn=self._token_combine, + ) + + +class ExpertParallelWithTP(ParallelStyle): + """Combined EP + TP for experts. + + Applied to ExpertLayer. Distributes expert params on a 2D (ep, tp) mesh + with [Shard(0), Shard(1/2)]. Token dispatch/combine hooks are registered + on ExpertLayer (outer), TP reduction hook on Experts (inner), so forward + execution is: EP dispatch -> TP input -> forward -> TP reduce -> EP combine. + """ + + def __init__( + self, + ep_mesh: DeviceMesh, + tp_mesh: DeviceMesh, + ): + super().__init__() + self.ep_mesh = ep_mesh + self.tp_mesh = tp_mesh + self.ep_tp_mesh = DeviceMesh._concatenate([ep_mesh, tp_mesh]) + + def _apply(self, module: nn.Module, device_mesh: DeviceMesh) -> nn.Module: + experts = module.experts # type: ignore[attr-defined] + + # Partition expert weights on 2D (ep, tp) mesh + for pn, p in experts.named_parameters(recurse=False): + if pn == "w1": + placements = [Shard(0), Shard(1)] + elif pn == "w2": + placements = [Shard(0), Shard(2)] + else: + continue + experts.register_parameter( + pn, nn.Parameter(distribute_tensor(p, self.ep_tp_mesh, placements)) + ) + + # EP dispatch/combine hooks on ExpertLayer (outer module) + ep_mesh = self.ep_mesh + + def ep_dispatch(mod, inputs): + (x,) = inputs + x = all_gather_tensor_autograd(x, gather_dim=0, group=ep_mesh.get_group()) + return (torch.ops._c10d_functional.wait_tensor(x),) + + def ep_combine(mod, inputs, output): + out = reduce_scatter_tensor_autograd( + output, "sum", scatter_dim=0, group=ep_mesh.get_group() + ) + return torch.ops._c10d_functional.wait_tensor(out) + + module.register_forward_pre_hook(ep_dispatch) + module.register_forward_hook(ep_combine) + + # TP reduction hook on Experts (inner module) + tp_mesh = self.tp_mesh + + def tp_allreduce_input_grad(mod, inputs): + (x,) = inputs + return ( + DTensor.from_local(x, tp_mesh, [Replicate()]).to_local( + grad_placements=[Partial()] + ), + ) + + experts.register_forward_pre_hook(tp_allreduce_input_grad) + + def tp_reduce(mod, inputs, output): + return ( + DTensor.from_local(output, tp_mesh, [Partial()]) + .redistribute(tp_mesh, [Replicate()]) + .to_local() + ) + + experts.register_forward_hook(tp_reduce) + return module + + +class TransformerBlock(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + self.attention_norm = nn.LayerNorm(args.dim) + self.attention = Attention(args) + self.ffn_norm = nn.LayerNorm(args.dim) + + self.has_experts = args.num_experts > 0 + if self.has_experts: + self.expert_layer = ExpertLayer( + args.num_experts, + dim=args.dim, + hidden_dim=4 * args.dim, + ) + else: + self.feed_forward = FeedForward( + args.dim, + hidden_dim=4 * args.dim, + dropout_p=args.dropout_p, + ) + + def forward(self, x): + h = x + self.attention(self.attention_norm(x)) + if self.has_experts: + out = h + self.expert_layer(self.ffn_norm(h)) + else: + out = h + self.feed_forward(self.ffn_norm(h)) + return out + + +# A toy transformer model, partly inspired by the nanoGPT model: +# https://github.com/karpathy/nanoGPT. +class Transformer(nn.Module): + def __init__(self, args: ModelArgs): + super().__init__() + if args.vocab_size is None: + raise AssertionError("Expected args.vocab_size to not be None") + if args.max_seq_len is None: + raise AssertionError("Expected args.max_seq_len to not be None") + self.model_args = args + self.max_seq_len = args.max_seq_len + self.tok_embeddings = nn.Embedding(args.vocab_size, args.dim) + self.pos_embeddings = nn.Embedding(args.max_seq_len, args.dim) + self.dropout = nn.Dropout(args.dropout_p) + self.layers = nn.ModuleList() + for _ in range(args.n_layers): + self.layers.append(TransformerBlock(args)) + self.norm = nn.LayerNorm(args.dim) + self.output = nn.Linear(args.dim, args.vocab_size, bias=False) + if args.weight_tying: + self.output.weight = self.tok_embeddings.weight + self.checkpoint_activations = args.checkpoint_activations + + def forward(self, tokens): + _bsz, seq_len = tokens.size() + if seq_len > self.max_seq_len: + raise AssertionError( + f"Expected seq_len <= max_seq_len, got {seq_len} > {self.max_seq_len}" + ) + h = self.tok_embeddings(tokens) + pos = torch.arange(0, seq_len, device=tokens.device) + p = self.pos_embeddings(pos) # positional embeddings of shape (seq_len, dim) + h = h + p + h = self.dropout(h) + for layer in self.layers: + if self.checkpoint_activations: + h = torch.utils.checkpoint.checkpoint(layer, h, use_reentrant=False) + else: + h = layer(h) + h = self.norm(h) + output = self.output(h).float() + return output + + @staticmethod + def parallelize( + module: "Transformer", + tp_mesh: DeviceMesh | None, + use_seq_parallel: bool, + local_output_for_attn: bool = False, + ep_mesh: DeviceMesh | None = None, + ) -> nn.Module: + if not isinstance(module, Transformer): + raise AssertionError(f"Requires Transformer but got {module}") + if tp_mesh is None and ep_mesh is None: + raise ValueError("At least one of tp_mesh or ep_mesh must be provided") + + # Parallelize the root submodules with TP. + if tp_mesh is not None: + if use_seq_parallel: + root_plan = { + "tok_embeddings": RowwiseParallel( + input_layouts=Replicate(), output_layouts=Shard(1) + ), + "pos_embeddings": RowwiseParallel( + input_layouts=Replicate(), output_layouts=Shard(0) + ), + "norm": SequenceParallel(), + } + else: + root_plan = { + "tok_embeddings": RowwiseParallel( + input_layouts=Replicate(), output_layouts=Replicate() + ), + "pos_embeddings": RowwiseParallel( + input_layouts=Replicate(), output_layouts=Replicate() + ), + } + parallelize_module(module, tp_mesh, root_plan) + + # Parallelize the attention and feed forward submodules. + for layer in module.layers: + if tp_mesh is not None: + layer_parallelize_plan = {} + if use_seq_parallel: + layer_parallelize_plan["attention"] = PrepareModuleInput( + input_layouts=Shard(1), + desired_input_layouts=Replicate(), + ) + # shard the RMSNorms + layer_parallelize_plan["attention_norm"] = SequenceParallel() + layer_parallelize_plan["ffn_norm"] = SequenceParallel() + layer_parallelize_plan["attention.wq"] = ColwiseParallel( + use_local_output=local_output_for_attn + ) + layer_parallelize_plan["attention.wk"] = ColwiseParallel( + use_local_output=local_output_for_attn + ) + layer_parallelize_plan["attention.wv"] = ColwiseParallel( + use_local_output=local_output_for_attn + ) + layer_parallelize_plan["attention.wo"] = ( + RowwiseParallel(output_layouts=Shard(1)) + if use_seq_parallel + else RowwiseParallel() + ) + + if not layer.has_experts: + layer_parallelize_plan["feed_forward.w1"] = ( + ColwiseParallel(input_layouts=Shard(1)) + if use_seq_parallel + else ColwiseParallel() + ) + layer_parallelize_plan["feed_forward.w2"] = ( + RowwiseParallel(output_layouts=Shard(1)) + if use_seq_parallel + else RowwiseParallel() + ) + elif ep_mesh is None: + # No EP mesh provided, use TP for experts + layer_parallelize_plan["expert_layer.experts"] = ( + TensorParallelForExpert( + input_layouts=Shard(0), + output_layouts=Shard(0), + ) + if use_seq_parallel + else TensorParallelForExpert() + ) + + parallelize_module(layer, tp_mesh, layer_parallelize_plan) + + # EP (+ optional TP) for experts + if ep_mesh is not None and layer.has_experts: + if tp_mesh is not None: + parallelize_module( + layer.expert_layer, + ep_mesh, + ExpertParallelWithTP(ep_mesh, tp_mesh), + ) + else: + parallelize_module( + layer.expert_layer.experts, ep_mesh, ExpertParallel() + ) + + if tp_mesh is not None: + # Parallelize the output submodule. If weight tying is enabled, + # we need to make sure output.weight is sharded consistently as + # tok_embeddings.weight, at the cost of the all_reduce operation + # using RowwiseParallel. + output_parallelize_plan = ( + ColwiseParallel( + input_layouts=Shard(1), + output_layouts=Replicate(), + ) + if use_seq_parallel + else ColwiseParallel(output_layouts=Replicate()) + ) + parallelize_module(module.output, tp_mesh, output_parallelize_plan) + + if local_output_for_attn: + for layer in module.layers: + layer.attention.n_heads = ( + module.model_args.n_heads // tp_mesh.size() + ) + + # Manually set output.weight so that parameters and gradients + # are shared. + if module.model_args.weight_tying: + module.output.weight = module.tok_embeddings.weight + + return module + + +def skip_unless_torch_gpu(method: T) -> T: + """ + Test decorator which skips the test unless there's a GPU available to torch. + + >>> # xdoctest: +SKIP + >>> @skip_unless_torch_gpu + >>> def test_some_method(self) -> None: + >>> ... + """ + # The builtin @skip_if_no_gpu relies on os.environ['WORLD_SIZE'] being set. + return cast(T, skip_if_lt_x_gpu(NUM_DEVICES)(method)) + + +class DTensorTestMixin: + """Shared test helpers for DTensorTestBase and DTensorContinuousTestBase.""" + + @property + def is_local_tensor_enabled(self) -> bool: + return False + + @property + def device_type(self) -> str: + if ( + not (TEST_CUDA or TEST_XPU or TEST_HPU or TEST_PRIVATEUSE1) + or DEVICE_COUNT < self.world_size + ): + return "cpu" + else: + return DEVICE_TYPE + + def build_device_mesh(self) -> DeviceMesh: + return init_device_mesh(self.device_type, (self.world_size,)) + + def init_manual_seed_for_rank(self) -> None: + torch.manual_seed(self.rank) + + def _test_op_on_dtensor(self, op_call, *args, **kwargs) -> None: + """ + Checks ``op_call(dtensor).full_tensor() == op_call(dtensor.full_tensor())``. + Unlike _test_op where the DTensor sharding is generated by DTensorConverter, + this function takes in DTensor object directly as argument and test the equality + of calling op on full_tensor() and DTensor. + """ + args_flattened, args_spec = tree_flatten(args) + full_tensor_args_flattened = tuple( + arg.full_tensor().detach().clone() if isinstance(arg, DTensor) else arg + for arg in args_flattened + ) + full_tensor_args = tree_unflatten(full_tensor_args_flattened, args_spec) + full_tensor_kwargs = { + k: v.full_tensor() if isinstance(v, DTensor) else v + for k, v in kwargs.items() + } + + out_flattened, _ = tree_flatten( + op_call(*full_tensor_args, **full_tensor_kwargs) + ) + d_out_flattened, _ = tree_flatten(op_call(*args, **kwargs)) + d_out_full_tensor_flattened = [dt.full_tensor() for dt in d_out_flattened] + self.assertEqual(out_flattened, d_out_full_tensor_flattened) + + # pyre-ignore[2]: + def _test_op(self, mesh: DeviceMesh, op_call, *args, **kwargs) -> None: + out = op_call(*args, **kwargs) + dtc = DTensorConverter(mesh, args, kwargs) + for d_args, d_kwargs in dtc: + # pyre can't find assertTrue anymore? + self.assertEqual(dtc.successful(), True) + d_out = op_call(*d_args, **d_kwargs) + self.assertEqual(d_out.full_tensor(), out) + + def run_subtests(self, *args, **kwargs): + return run_subtests(self, *args, **kwargs) + + +class DTensorContinuousTestBase(DTensorTestMixin, MultiProcContinuousTest): + @classmethod + def backend_str(cls) -> str: + backend = dist.get_default_backend_for_device(DEVICE_TYPE) + return backend + + @classmethod + def _init_pg(cls, rank, world_size, rdvz_file): + # Set device before initializing process group to ensure + # each rank is bound to the correct GPU. However, if world_size > device_count, + # we skip the test. + if torch.accelerator.is_available(): + if world_size > torch.accelerator.device_count(): + sys.exit(TEST_SKIPS[f"multi-gpu-{world_size}"].exit_code) + else: + torch.accelerator.set_device_index(rank) + + # Call parent's _init_pg to do the actual process group initialization + super()._init_pg(rank, world_size, rdvz_file) + + +class LocalDTensorContinuousTestBase(DTensorContinuousTestBase): + @property + def is_local_tensor_enabled(self) -> bool: + return True + + def _handle_test_skip(self, msg: str) -> None: + self.skipTest(msg) + + def _get_local_tensor_mode(self): + return LocalTensorMode(frozenset(range(self.world_size))) + + @classmethod + def _ensure_processes_spawned(cls): + if cls._processes_spawned: + return + if cls.world_size == -2: + cls.world_size = NUM_DEVICES + store = FakeStore() + dist.init_process_group( + backend="fake", + world_size=cls.world_size, + rank=0, + store=store, + ) + cls.processes = [] + cls.task_queues = [] + cls.completion_queues = [] + cls._processes_spawned = True + + @classmethod + def tearDownClass(cls): + if cls._processes_spawned: + dist.destroy_process_group() + cls._processes_spawned = False + unittest.TestCase.tearDownClass() + + def setUp(self): + unittest.TestCase.setUp(self) + self.__class__._ensure_processes_spawned() + torch.autograd._enable_record_function(False) + + def tearDown(self): + from torch.distributed.tensor import _random as random + + random._rng_tracker = None + unittest.TestCase.tearDown(self) + torch.autograd._enable_record_function(True) + + def __init__(self, method_name="runTest", methodName="runTest"): + if methodName != "runTest": + method_name = methodName + unittest.TestCase.__init__(self, method_name) + + @property + def rank(self): + return torch.SymInt(LocalIntNode({r: r for r in range(self.world_size)})) + + @rank.setter + def rank(self, rank): + pass + + def build_device_mesh(self) -> DeviceMesh: + with maybe_disable_local_tensor_mode(): + return super().build_device_mesh() + + def init_manual_seed_for_rank(self) -> None: + torch.manual_seed(0) + + +class DTensorTestBase(DTensorTestMixin, MultiProcessTestCase): + @property + def world_size(self) -> int: + return NUM_DEVICES + + @property + def backend(self) -> str: + backend = dist.get_default_backend_for_device(self.device_type) + return backend + + def init_pg(self, eager_init, backend: str | None = None) -> None: + if backend is None: + backend = self.backend + + requires_gpu = any( + gpu_backend in backend for gpu_backend in ACCELERATOR_DIST_BACKENDS + ) + if requires_gpu and torch.accelerator.device_count() < self.world_size: + sys.exit(TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code) + + curr_backend = dist.get_default_backend_for_device(self.device_type) + + if backend not in [ + "nccl", + "gloo", + "mpi", + f"cpu:gloo,{self.device_type}:{curr_backend}", + "cpu:gloo,cuda:ncclx", + "cuda:ncclx", + "hccl", + "xccl", + "fake", + "cpu:gloo,xpu:xccl", + ]: + raise RuntimeError(f"Backend {backend} not supported!") + + device_id = None + if "nccl" in backend or "xccl" in backend: + # set device for nccl pg for collectives + # TODO: if users want to enable testing across hosts, we may need + # to change this part. + torch.accelerator.set_device_index(self.rank) + # we only need to set device_id for nccl backend with eager init + device_id = ( + torch.device(f"{self.device_type}:{self.rank}") if eager_init else None + ) + + # For nccl backend, bind the device to the process if device_id is not None + # so the nccl communicator is immediately formed and we can use `ncclCommSplit` + # for form subgroup to avoid unnecessary overhead. + dist.init_process_group( + backend=backend, + world_size=self.world_size, + rank=self.rank, # pyre-ignore[16] + init_method=f"file://{self.file_name}", # pyre-ignore[16] + device_id=device_id, + ) + + def destroy_pg(self, device_id: int | None = None) -> None: + # Wait for all ranks to reach here before starting shutdown. + # FIXME dist.barrier deadlocks with multiple threads and NCCL: https://github.com/pytorch/pytorch/issues/95895 + # dist.all_reduce(torch.zeros((1,), device="cuda" if TEST_CUDA else "cpu")) + # FIXME can't use the above all_reduce as it causes hangs on bionic and focal. It hangs: + # test_dtensor.py -- DTensorMeshTest.test_dtensor_device_mesh_device_conversion + if device_id is None: + device_id = ( + torch.cuda.current_device() if self.device_type == "cuda" else self.rank + ) + + if self.device_type == "cpu": + # NOTE: when `device_id` is not None, barrier() will choose the accelerator + # of the most pripority, which means if the test specifies to use CPU for + # testing while CUDA is available on the host, the barrier() will use CUDA. + # To avoid this and better respect `self.device_type`, we add this branch to + # enforce barrier() to use CPU when `self.device_type` is CPU and other + # accelerator is also available. + dist.barrier() + else: + dist.barrier(device_ids=[device_id]) + + dist.destroy_process_group() + + def setUp(self) -> None: + super().setUp() + self._spawn_processes() + + +TestFunc = Callable[[...], object] + + +# wrapper to initialize comms (processgroup) +def with_comms( + eager_init: TestFunc | bool = False, + backend: str | None = None, +) -> TestFunc: + def decorator(func, eager_init: bool = False, backend: str | None = None): + @wraps(func) # pyre-ignore[6] + def wrapper( + self, + *args: tuple[object], + **kwargs: dict[str, Any], # type: ignore[misc] + ) -> None: + # just passthrough if harness doesn't + # support init_pg e.g., DTensorOpTestBase + if not hasattr(self, "init_pg"): + func(self, *args, **kwargs) + return + + self.init_pg(eager_init, backend) + + try: + func(self, *args, **kwargs) # type: ignore[misc] + except Exception as e: + dist.destroy_process_group() + raise e + + self.destroy_pg() + + return wrapper + + return ( + decorator(func=eager_init) + if callable(eager_init) + else partial(decorator, eager_init=eager_init, backend=backend) + ) + + +class DTensorOpTestBase(MultiThreadedTestCase): + @property + def world_size(self) -> int: + return NUM_DEVICES + + @property + def device_type(self) -> str: + return DEVICE_TYPE + + def build_device_mesh(self): + return init_device_mesh(self.device_type, (self.world_size,)) + + def setUp(self) -> None: + super().setUp() + # Enable thread-safe lock for ShardingPropagator since we run + # multi-threaded tests. + from torch.distributed.tensor._sharding_prop import ShardingPropagator + + self._orig_fake_mode_lock = ShardingPropagator._fake_mode_lock + ShardingPropagator._fake_mode_lock = threading.Lock() + self._spawn_threads() + + def tearDown(self) -> None: + # Restore the original (no-op) lock + from torch.distributed.tensor._sharding_prop import ShardingPropagator + + ShardingPropagator._fake_mode_lock = self._orig_fake_mode_lock + super().tearDown() + + +# This is a class for converting args/kwargs of an op into distributed args/kwargs +class DTensorConverter: + def __init__( + self, + mesh: DeviceMesh, + args: tuple[object, ...], + kwargs: dict[str, object], + replicate_only: bool = False, + ) -> None: + self.hit = 0 + self.miss = 0 + self.mesh = mesh + self.args = args + self.kwargs = kwargs + self.replicate_only = replicate_only + flatten_args, flatten_args_spec = tree_flatten(args) + flatten_kwargs, flatten_kwargs_spec = tree_flatten(kwargs) + + self.flatten_args: list[object] = flatten_args + self.flatten_args_spec: TreeSpec = flatten_args_spec + self.flatten_kwargs: list[object] = flatten_kwargs + self.flatten_kwargs_spec: TreeSpec = flatten_kwargs_spec + + choices_for_args = [ + self.gen_sharding_choices_for_arg(arg) + for arg in self.flatten_args + if isinstance(arg, torch.Tensor) + ] + + choices_for_args.extend( + self.gen_sharding_choices_for_arg(arg) + for arg in self.flatten_kwargs + if isinstance(arg, torch.Tensor) + ) + + self.sharding_combs: Iterator[Sequence[Placement]] = iter( + itertools.product(*choices_for_args) + ) + + def successful(self) -> bool: + return self.hit > 0 and self.miss == 0 + + def is_supported_tensor(self, t: torch.Tensor) -> bool: + # TODO: dist tensor need to support quantized and sparse + # tensors, quantized tensor might be relatively easy, but + # sparse tensor have special layouts that we need to possibly + # deal with, until we are clear about them, we don't officially + # support them. + return not any( + [ + t.is_sparse_csr, + t.is_sparse, + t.is_mkldnn, + t.is_quantized, + t.is_nested, + torch._is_functional_tensor(t), + t.is_neg(), + t.is_conj(), + t.device.type in ("lazy", "meta"), + # We need a way to test if a tensor is batched but there + # is no official APi to do it + # torch._C._is_batched(t), + ] + ) + + def gen_sharding_choices_for_arg(self, arg: torch.Tensor) -> Sequence[Placement]: + # If replicate_only is set, only use Replicate placement + if self.replicate_only: + return [Replicate()] + + mesh_size = self.mesh.size() + sharding_choices: list[Placement] = [Replicate()] + # c10d collective does not support bool tensor + # for bool tensor we treat it as replicated + if arg.dtype != torch.bool: + # only generating choices with: replicate, or sharding + # evenly on a dimension that could be sharded + sharding_choices = sharding_choices + [ + Shard(i) + for i, s in enumerate(arg.shape) + if s > 1 and s % mesh_size == 0 + ] + # TODO: add multi mesh choices + # all_choices = itertools.product( + # *(self.mesh.ndim * [sharding_choices]) + # ) + return sharding_choices + + def __iter__(self) -> "DTensorConverter": + return self + + def __next__(self) -> tuple[tuple[object, ...], dict[str, object]]: + try: + next_sharding_choices = next(self.sharding_combs) + idx = 0 + + new_args: list[object] = [] + for arg in self.flatten_args: + if isinstance(arg, torch.Tensor): + new_args.append( + self.to_dist_tensor( + arg, self.mesh, [next_sharding_choices[idx]] + ) + ) + idx += 1 + else: + new_args.append(arg) + + new_kwargs: list[object] = [] + for arg in self.flatten_kwargs: + if isinstance(arg, torch.Tensor): + new_kwargs.append( + self.to_dist_tensor( + arg, self.mesh, [next_sharding_choices[idx]] + ) + ) + idx += 1 + else: + new_kwargs.append(arg) + + return ( + tree_unflatten(new_args, self.flatten_args_spec), + tree_unflatten(new_kwargs, self.flatten_kwargs_spec), + ) + except StopIteration as e: + raise StopIteration from e + + def to_dist_tensor( + self, t: torch.Tensor, mesh: DeviceMesh, placements: list[Placement] + ) -> torch.Tensor: + if type(t) is torch.Tensor or type(t) is nn.Parameter or type(t) is LocalTensor: + if self.is_supported_tensor(t): + self.hit += 1 + if t.ndim == 0: + # scalar tensor by default will be replicated + r = distribute_tensor(t, mesh, [Replicate()] * mesh.ndim) + else: + # distribute non-scalar tensors + r = distribute_tensor(t, mesh, placements) + if isinstance(t, nn.Parameter): + r = nn.Parameter( # type: ignore[assignment] + r, requires_grad=r.requires_grad + ) + return r + else: + self.miss += 1 + return t + elif torch.overrides.is_tensor_like(t): + # Blindly converting tensor subclasses to dist tensor can cause + # unpredictable problems, we explicitly disable this conversion + # for now (i.e. we don't support DTensor holding tensor subclass + # until there's a strong reason later). + self.miss += 1 + return t + else: + raise RuntimeError(f"Trying to convert to DTensor, but got {type(t)}") + + +class LocalDTensorOpTestBase(DTensorOpTestBase): + @property + def is_local_tensor_enabled(self) -> bool: + return True + + def _handle_test_skip(self, msg: str) -> None: + self.skipTest(msg) + + def _get_local_tensor_mode(self): + return LocalTensorMode(frozenset(range(self.world_size))) + + def setUp(self) -> None: + super().setUp() + torch.autograd._enable_record_function(False) + + def tearDown(self) -> None: + from torch.distributed.tensor import _random as random + + random._rng_tracker = None + super().tearDown() + torch.autograd._enable_record_function(True) + + @property + def rank(self): + return torch.SymInt(LocalIntNode({r: r for r in range(self.world_size)})) + + @rank.setter + def rank(self, rank): + pass + + def join_or_run(self, fn): + @wraps(fn) + def wrapper(self): + fn() + + return types.MethodType(wrapper, self) + + def build_device_mesh(self) -> DeviceMesh: + with maybe_disable_local_tensor_mode(): + return super().build_device_mesh() + + def init_pg(self, eager_init, backend: str | None = None) -> None: + dist.init_process_group("fake", rank=0, world_size=self.world_size) + self._pg = dist.distributed_c10d._get_default_group() + + def destroy_pg(self, device_id: int | None = None) -> None: + dist.destroy_process_group(self._pg) + self._pg = None + + def _spawn_processes(self) -> None: + pass + + def _spawn_threads(self) -> None: + pass + + def run_test(self, test_name: str, parent_pipe) -> None: + getattr(self, test_name)() + + def init_manual_seed_for_rank(self) -> None: + torch.manual_seed(0) + + +class LocalDTensorTestBase(DTensorTestBase): + @property + def is_local_tensor_enabled(self) -> bool: + return True + + def _handle_test_skip(self, msg: str) -> None: + self.skipTest(msg) + + def _get_local_tensor_mode(self): + return LocalTensorMode(frozenset(range(self.world_size))) + + def setUp(self) -> None: + super().setUp() + torch.autograd._enable_record_function(False) + + def tearDown(self) -> None: + from torch.distributed.tensor import _random as random + + random._rng_tracker = None + super().tearDown() + torch.autograd._enable_record_function(True) + + @property + def rank(self): + return torch.SymInt(LocalIntNode({r: r for r in range(self.world_size)})) + + @rank.setter + def rank(self, rank): + pass + + def join_or_run(self, fn): + @wraps(fn) + def wrapper(self): + fn() + + return types.MethodType(wrapper, self) + + def build_device_mesh(self) -> DeviceMesh: + with maybe_disable_local_tensor_mode(): + return super().build_device_mesh() + + def init_pg(self, eager_init, backend: str | None = None) -> None: + dist.init_process_group("fake", rank=0, world_size=self.world_size) + self._pg = dist.distributed_c10d._get_default_group() + + def destroy_pg(self, device_id: int | None = None) -> None: + dist.destroy_process_group(self._pg) + self._pg = None + + def _spawn_processes(self) -> None: + pass + + def run_test(self, test_name: str, parent_pipe) -> None: + getattr(self, test_name)() + + def init_manual_seed_for_rank(self) -> None: + torch.manual_seed(0) + + +def make_wrapped(fn, ctxs): + @functools.wraps(fn) + def wrapped(self): + torch._dynamo.reset() + stack = contextlib.ExitStack() + for ctx in ctxs: + if callable(ctx): + stack.enter_context(ctx(self)) + else: + stack.enter_context(ctx) + try: + out = fn(self) + finally: + stack.close() + return out + + return wrapped + + +def create_local_tensor_test_class( + orig_cls, skipped_tests=None, base_class=LocalDTensorTestBase +): + if skipped_tests is None: + skipped_tests = [] + + dct = orig_cls.__dict__.copy() + for name in list(dct.keys()): + fn = dct[name] + if not callable(fn): + continue + elif name in skipped_tests: + dct[name] = lambda self: self.skipTest("Skipped test") + elif name.startswith("test_"): + ctxs = [ + lambda test: test._get_local_tensor_mode(), + ] + dct[name] = make_wrapped(fn, ctxs) + + cls = type( + orig_cls.__name__ + "WithLocalTensor", + (base_class,) + orig_cls.__bases__, + dct, + ) + cls.__file__ = __file__ + return cls + + +@maybe_run_for_local_tensor +def map_local_tensor_for_rank(tensor, rank, func): + return func(tensor, rank) + + +@maybe_run_for_local_tensor +def map_local_for_rank(rank, func): + return func(rank) + + +def reduce_local_int(val, func): + return func(val.node._local_ints) + + +def _convert_shard_order_dict_to_ShardOrder(shard_order): + """Convert shard_order dict to ShardOrder""" + return tuple( + ShardOrderEntry(tensor_dim=tensor_dim, mesh_dims=tuple(mesh_dims)) + for tensor_dim, mesh_dims in shard_order.items() + ) + + +# TODO(zpcore): remove once the native redistribute supports shard_order arg +def redistribute( + dtensor_input, + device_mesh, + placements, + shard_order, + use_graph_based_transform=True, +): + """ + wrapper function to support shard_order for redistribution + This is a simpler version of Redistribute, only considers the forward. + """ + if placements is None: + placements = shard_order_to_placement(shard_order, device_mesh) + placements = tuple(placements) + old_spec = dtensor_input._spec + new_spec = copy.deepcopy(old_spec) + new_spec.placements = placements + if shard_order is not None: + new_spec.shard_order = shard_order + else: + new_spec.shard_order = () + if old_spec == new_spec: + return dtensor_input + dtensor_input = DTensor.from_local( + redistribute_local_tensor( + dtensor_input.to_local(), + old_spec, + new_spec, + use_graph_based_transform=use_graph_based_transform, + ), + device_mesh, + ) + dtensor_input._spec = copy.deepcopy(new_spec) + return dtensor_input # returns DTensor + + +# TODO(zpcore): remove once the native distribute_tensor supports +# shard_order arg +def patched_distribute_tensor( + input_tensor, + device_mesh, + placements, + shard_order, + use_graph_based_transform=True, + src_data_rank: int | None = 0, +): + """wrapper function to support shard_order for tensor distribution""" + if placements is None: + placements = shard_order_to_placement(shard_order, device_mesh) + placements = tuple(placements) + tensor_dt = distribute_tensor( + input_tensor, device_mesh, placements, src_data_rank=src_data_rank + ) + # Do not consider _StridedShard to express shard order + tensor_dt._spec.use_strided_shard_as_shard_order = False + tensor_dt._spec.__post_init__() + # fix the shard order + return redistribute( + tensor_dt, device_mesh, placements, shard_order, use_graph_based_transform + ) + + +# TODO(zpcore): remove once the native redistribute supports shard_order arg +def make_full_tensor(dtensor_input): + """wrapper function to support DTensor.full_tensor""" + return redistribute( + dtensor_input, dtensor_input.device_mesh, placements=None, shard_order=() + ).to_local() + + +def shard_order_to_placement(shard_order, mesh): + """convert shard_order to placement with only Replicate() and Shard()""" + placements: list[Any] = [Replicate() for _ in range(mesh.ndim)] + if shard_order is not None: + for entry in shard_order: + tensor_dim = entry.tensor_dim + mesh_dims = entry.mesh_dims + for mesh_dim in mesh_dims: + placements[mesh_dim] = Shard(tensor_dim) + return tuple(placements) + + +def generate_shard_orders(mesh, tensor_rank): + # Generate all possible sharding placement of tensor with rank + # `tensor_rank` over mesh. + def _split_list(lst: list, N: int): + def compositions(n: int, k: int): + # yields lists of length k, positive ints summing to n + for cuts in itertools.combinations(range(1, n), k - 1): + # add 0 and n as sentinels, then take consecutive differences + yield [b - a for a, b in itertools.pairwise((0, *cuts, n))] + + length = len(lst) + for comp in compositions(length, N): + result = [] + start = 0 + for size in comp: + result.append(lst[start : start + size]) + start += size + yield result + + all_mesh = list(range(mesh.ndim)) + all_device_order = list(itertools.permutations(all_mesh)) + for device_order in all_device_order: + # split on device orders, and assign each device order segment to a tensor dim + for num_split in range(1, mesh.ndim + 1): + for splitted_list in _split_list(list(range(mesh.ndim)), num_split): + for tensor_dims in itertools.combinations( + range(tensor_rank), len(splitted_list) + ): + shard_order = {} + if len(tensor_dims) != len(splitted_list): + raise AssertionError( + f"Expected len(tensor_dims) == len(splitted_list), " + f"got {len(tensor_dims)} != {len(splitted_list)}" + ) + for tensor_dim, mesh_dims in zip(tensor_dims, splitted_list): + shard_order[tensor_dim] = device_order[ + mesh_dims[0] : mesh_dims[-1] + 1 + ] + yield _convert_shard_order_dict_to_ShardOrder(shard_order) + + +def validate_sharding_rule_sample( + op, full_args, full_kwargs, input_placements, output_placements, device_mesh +): + from torch.utils import _pytree as pytree + + # Extract tensors from args in order, pair with placements + full_tensors = [ + a for a in pytree.tree_leaves(full_args) if isinstance(a, torch.Tensor) + ] + full_tensors += [ + a for a in pytree.tree_leaves(full_kwargs) if isinstance(a, torch.Tensor) + ] + + dtensors = [ + distribute_tensor(t, device_mesh, (p,)) + for t, p in zip(full_tensors, input_placements) + ] + + # Build sharded args by replacing tensors with their sharded local versions + dtensor_idx = 0 + + def _to_local_shard(a): + nonlocal dtensor_idx + if isinstance(a, torch.Tensor): + local = dtensors[dtensor_idx].to_local() + dtensor_idx += 1 + return local + return a + + local_args, local_kwargs = pytree.tree_map( + _to_local_shard, (full_args, full_kwargs) + ) + + # run and compare + ref_output = op(*full_args, **full_kwargs) + local_output = op(*local_args, **local_kwargs) + + ref_tensors = [ + t for t in pytree.tree_leaves(ref_output) if isinstance(t, torch.Tensor) + ] + local_tensors = [ + t for t in pytree.tree_leaves(local_output) if isinstance(t, torch.Tensor) + ] + + for ref, local, plc in zip(ref_tensors, local_tensors, output_placements): + dt = DTensor.from_local(local, device_mesh, (plc,)) + full = dt.redistribute(device_mesh, (Replicate(),)).to_local() + if ref.shape != full.shape or not torch.allclose( + ref, full, atol=1e-5, rtol=1e-5, equal_nan=True + ): + return False + return True + + +@contextlib.contextmanager +def op_strategy_context(op_overload, strategy_func, schema_info=None): + """ + Context manager for setting and clearing op strategies. + Args: + op_overload: The operator overload to set or clear the strategy for. + strategy_func: The strategy function to set for the operator overload. + schema_info: Optional schema information for the operator overload. + Yields: + None + """ + from torch.distributed.tensor._ops.utils import register_op_strategy + from torch.distributed.tensor.debug import _clear_sharding_prop_cache + + propagator = DTensor._op_dispatcher.sharding_propagator + _origin_op_strategy_funcs = None + _origin_op_strategy_schema = None + try: + # register the op strategy + if op_overload in propagator.op_strategy_funcs: + _origin_op_strategy_funcs = propagator.op_strategy_funcs[op_overload] + del propagator.op_strategy_funcs[op_overload] + if op_overload in propagator.op_to_schema_info: + _origin_op_strategy_schema = propagator.op_to_schema_info[op_overload] + del propagator.op_to_schema_info[op_overload] + register_op_strategy(op_overload, schema_info=schema_info)(strategy_func) + yield + finally: + # clear this op strategy cache + if _origin_op_strategy_funcs is None: + if op_overload in propagator.op_strategy_funcs: + del propagator.op_strategy_funcs[op_overload] + else: + propagator.op_strategy_funcs[op_overload] = _origin_op_strategy_funcs + if _origin_op_strategy_schema is None: + if op_overload in propagator.op_to_schema_info: + del propagator.op_to_schema_info[op_overload] + else: + propagator.op_to_schema_info[op_overload] = _origin_op_strategy_schema + _clear_sharding_prop_cache() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/checkpoint_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/checkpoint_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..91e79a5dbf1c2d0d78fa197958d414494dd11305 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/checkpoint_utils.py @@ -0,0 +1,196 @@ +# mypy: allow-untyped-defs + +# Copyright (c) Meta Platforms, Inc. and affiliates + +import io +import logging +import os +import shutil +import tempfile +from collections.abc import Callable +from functools import wraps +from typing import Any, cast, IO + +# introduced as collections.abc.Buffer in Python 3.12 +from typing_extensions import Buffer + +import torch.distributed as dist +from torch.distributed.checkpoint._extension import ( + ExtensionRegistry, + StreamTransformExtension, +) + + +class Rot13Example(StreamTransformExtension): + """ + This is an example stream transform extension which just does rot13 on each + alphanumeric character of the stream. It is mainly intended as a demonstration + and for testing; there isn't a production use case for this. + """ + + def __init__(self, chunk_size: int = io.DEFAULT_BUFFER_SIZE) -> None: + super().__init__() + self._chunk_size = chunk_size + + @staticmethod + def from_descriptor(version: str) -> "Rot13Example": + if version.partition(".")[0] != "1": + raise ValueError(f"Unknown extension {version=}") + return Rot13Example() + + @staticmethod + def registry_name() -> str: + return "stream.rot13" + + def get_descriptor(self) -> str: + return f"{self.registry_name()}/1" + + @staticmethod + def _rot13bytes(b: Buffer, count: int) -> None: + b = memoryview(b) + for i in range(count): + ch = b[i] + if ch >= ord("A") and ch <= ord("Z"): + ch += ord("a") - ord("A") + elif ch >= ord("a") and ch <= ord("z"): + ch += ord("A") - ord("a") + b[i] = ch + + def transform_to(self, output: IO[bytes]) -> IO[bytes]: + class Writer(io.RawIOBase): + def __init__(self, output: IO[bytes]) -> None: + self.output = output + + def writeable(self) -> bool: + return True + + def write(self, b: Buffer) -> int | None: + # Don't mutate the input + chunk = bytearray(b) + Rot13Example._rot13bytes(chunk, len(chunk)) + return self.output.write(chunk) + + def flush(self) -> None: + self.output.flush() + + return cast(IO[bytes], Writer(output)) + + def transform_from(self, input: IO[bytes]) -> IO[bytes]: + class Reader(io.RawIOBase): + def __init__(self, input: IO[bytes]) -> None: + self.input = input + + def readable(self) -> bool: + return True + + def readinto(self, b: Buffer) -> int | None: + if hasattr(self.input, "readinto"): + count = self.input.readinto(b) + else: + # It's possible self.input is an IO[bytes] with no readinto method. + # In that case, we emulate with a read and copy. In practice, + # all of the current concrete extensions have readinto. + view = memoryview(b) + r = self.input.read(len(view)) + if r is None: + count = None + else: + count = len(r) + view[:count] = r + if count == 0 or count is None: + return count + + Rot13Example._rot13bytes(b, count) + return count + + def seekable(self) -> bool: + return self.input.seekable() + + def seek(self, offset: int, whence: int = os.SEEK_SET) -> int: + return self.input.seek(offset, whence) + + def tell(self) -> int: + return self.input.tell() + + return cast(IO[bytes], Reader(input)) + + +def get_test_extension_registry() -> ExtensionRegistry: + registry = ExtensionRegistry() + registry.register(Rot13Example) + return registry + + +def with_temp_dir( + func: Callable | None = None, +) -> Callable | None: + """ + Wrapper to initialize temp directory for distributed checkpoint. + """ + if func is None: + raise AssertionError("Expected func to not be None") + + @wraps(func) + def wrapper(self, *args: tuple[object], **kwargs: dict[str, Any]) -> None: + if dist.is_initialized(): + # Only create temp_dir when rank is 0 + if dist.get_rank() == 0: + temp_dir = tempfile.mkdtemp() + print(f"Using temp directory: {temp_dir}") + else: + temp_dir = "" + object_list = [temp_dir] + + # Broadcast temp_dir to all the other ranks + os.sync() + dist.broadcast_object_list(object_list) + self.temp_dir = object_list[0] + os.sync() + else: + temp_dir = tempfile.mkdtemp() + print(f"No process group initialized, using temp directory: {temp_dir}") + self.temp_dir = temp_dir + + try: + func(self, *args, **kwargs) + finally: + if dist.is_initialized() and dist.get_rank() == 0: + shutil.rmtree(self.temp_dir, ignore_errors=True) + else: + shutil.rmtree(self.temp_dir, ignore_errors=True) + + return wrapper + + +def with_checkpoint_logging( + func: Callable | None = None, + logger_name: str = "torch.distributed.checkpoint", + level: int = logging.INFO, +) -> Callable | None: + """ + Wrapper to configure checkpoint logging for distributed tests. + + Args: + func: The test function to wrap + logger_name: Name of the logger to configure (default: 'torch.distributed.checkpoint') + level: Logging level to set (default: logging.INFO) + """ + if func is None: + raise AssertionError("Expected func to not be None") + + @wraps(func) + def wrapper(self, *args: tuple[object], **kwargs: dict[str, Any]) -> None: + # Get the logger and store original level + target_logger = logging.getLogger(logger_name) + original_level = target_logger.level + + # Set the desired logging level + target_logger.setLevel(level) + + try: + func(self, *args, **kwargs) + finally: + # Restore original logging level + target_logger.setLevel(original_level) + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/common_state_dict.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/common_state_dict.py new file mode 100644 index 0000000000000000000000000000000000000000..a78e312306ba2500afa3722d6271c645d25f97cf --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/common_state_dict.py @@ -0,0 +1,170 @@ +# mypy: allow-untyped-defs + +# Owner(s): ["oncall: distributed"] + +import copy +from itertools import chain +from typing import Any + +import torch +import torch.nn as nn +from torch.distributed._sharded_tensor import ShardedTensor +from torch.distributed._state_dict_utils import _gather_state_dict +from torch.distributed.checkpoint.state_dict import ( + _PG, + _STATE, + set_state_dict, + StateDictOptions, +) +from torch.distributed.tensor import DTensor + + +class VerifyStateDictMixin: + def _compare_tensor(self, orig_tensor, dist_tensor, offload_to_cpu=False): + if isinstance(dist_tensor, (DTensor, ShardedTensor)): + dist_tensor = _gather_state_dict({"mykey": dist_tensor}).pop("mykey") + + if offload_to_cpu: + orig_tensor = orig_tensor.cpu() + dist_tensor = dist_tensor.cpu() + self.assertTrue(isinstance(dist_tensor, torch.Tensor)) + self.assertTrue(torch.allclose(orig_tensor, dist_tensor)) + + def _verify_msd( + self, + msd: dict[str, Any], + dist_msd: dict[str, Any], + options: StateDictOptions = StateDictOptions(), + offload_to_cpu=False, + ) -> None: + if not options.ignore_frozen_params: + self.assertEqual(len(msd), len(dist_msd)) + for fqn, param in msd.items(): + dist_param = dist_msd.get(fqn) + if not options.ignore_frozen_params: + self.assertIsNotNone(dist_param, f"{fqn=}") + try: + self._compare_tensor(param, dist_param, offload_to_cpu) + except AssertionError as e: + raise AssertionError( + f"{fqn} has mismatched value {param} {dist_param}" + ) from e + elif dist_param is None: + self.assertFalse(param.requires_grad, f"{fqn=}") + + def _verify_osd( + self, + model: nn.Module, + optim: torch.optim.Optimizer, + osd: dict[str, Any], + dist_osd: dict[str, Any], + ) -> None: + params = list(chain.from_iterable(g["params"] for g in optim.param_groups)) + param_pid_mapping = dict(zip(params, range(len(params)), strict=True)) + fqn_pid_mapping = {} + for fqn, param in model.named_parameters(): + pid = param_pid_mapping[param] + fqn_pid_mapping[fqn] = pid + fqn_pid_mapping[pid] = fqn + # Check optimizer_state_dict state + + self.assertEqual(len(osd[_STATE]), len(dist_osd[_STATE])) + for pid, states in osd[_STATE].items(): + fqn = fqn_pid_mapping[pid] + dist_states = dist_osd[_STATE].get(fqn, None) + self.assertIsNotNone(dist_states, fqn) + self.assertEqual(len(states), len(dist_states)) + for key, state in states.items(): + dist_state = states.get(key, None) + self.assertIsNotNone(dist_state) + self._compare_tensor(state, dist_state) + + # Check optimizer_state_dict param_group + old_dist_osd_pg = dist_osd[_PG] + if len(osd[_PG]) != len(dist_osd[_PG]): + self.assertTrue(len(dist_osd[_PG]) > len(osd[_PG])) + new_pg = copy.deepcopy(dist_osd[_PG][0]) + new_pg["params"] = [] + for dist_group in dist_osd[_PG]: + new_pg["params"].extend(dist_group["params"]) + dist_osd[_PG] = [new_pg] + + self.assertEqual(len(osd[_PG]), len(dist_osd[_PG])) + for group, dist_group in zip(osd[_PG], dist_osd[_PG], strict=True): + self.assertEqual(len(group), len(dist_group)) + for key, value in group.items(): + # Below doesn't work because param_groups can have None + # values. + # dist_value = dist_group.get(key, None) + # self.assertIsNotNone(dist_value, (dist_group, group)) + dist_value = dist_group[key] + if key == "params": + fqns = [fqn_pid_mapping[pid] for pid in value] + self.assertEqual(sorted(fqns), sorted(dist_value)) + else: + self.assertEqual(value, dist_value) + dist_osd[_PG] = old_dist_osd_pg + + def _verify_osd_by_load( + self, + model: nn.Module, + optim: torch.optim.Optimizer, + new_optim: torch.optim.Optimizer, + dist_osd: dict[str, Any], + ) -> None: + new_dist_osd = _gather_state_dict(dist_osd) + set_state_dict( + model, + optimizers=new_optim, + model_state_dict={}, + optim_state_dict=new_dist_osd, + ) + self.assertEqual(optim.state_dict(), new_optim.state_dict()) + + +class FusionEmbedding(nn.Module): + def __init__(self, vocab_size: int, fusion_vocab_size: int, embed_dim: int) -> None: + super().__init__() + self.embedding = nn.Embedding(vocab_size, embed_dim) + self.fusion_embedding = nn.Embedding(fusion_vocab_size, embed_dim) + + +class FusionEmbeddingWithHook(nn.Module): + def __init__(self, vocab_size: int, fusion_vocab_size: int, embed_dim: int) -> None: + super().__init__() + self.embedding = nn.Embedding(vocab_size, embed_dim) + self.fusion_embedding = nn.Embedding(fusion_vocab_size, embed_dim) + self._register_state_dict_hook(FusionEmbeddingWithHook._state_dict_hook) + self._register_load_state_dict_pre_hook( + FusionEmbeddingWithHook._load_state_dict_hook, with_module=True + ) + + def _state_dict_hook(self, destination, prefix, keep_vars): + """Remove "embedding" from the original embedding in the state_dict + name. This keeps the original state dict name for the embedding + from before fusing with the FusionEmbedding. + """ + key = prefix + "embedding.weight" + new_key = prefix + "weight" + destination[new_key] = destination[key] + del destination[key] + + def _load_state_dict_hook(self, state_dict, prefix, *args, **kwargs): + """Apply extra "embedding" prefix to the state_dict key to + account for the FusionEmbedding wrapping. + """ + if state_dict: + key = prefix + "weight" + new_key = prefix + "embedding.weight" + state_dict[new_key] = state_dict[key] + del state_dict[key] + + +class FusionEmbeddingWithModifier(FusionEmbeddingWithHook): + # _fqn_modifiers is a private function as a contract between DSD. When users change the state_dict + # keys, they need to provide a mapping from the new key to the original key. This is used to ensure + # consistency between the state_dict keys and fqn. + def _fqn_modifiers(self) -> dict[str, str]: + return { + "weight": "embedding", + } diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/ddp_under_dist_autograd_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/ddp_under_dist_autograd_test.py new file mode 100644 index 0000000000000000000000000000000000000000..853d394d18449f16877630143f16827697d033bf --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/ddp_under_dist_autograd_test.py @@ -0,0 +1,755 @@ +# mypy: allow-untyped-defs + +import contextlib +import enum +import logging +import os +import threading +from typing import NamedTuple + +import torch +import torch.distributed as dist +import torch.distributed.autograd as dist_autograd +import torch.nn as nn +from torch.distributed import rpc +from torch.distributed.nn import RemoteModule +from torch.nn.parallel import DistributedDataParallel +from torch.testing._internal.common_distributed import ( + requires_gloo, + requires_nccl, + skip_if_lt_x_gpu, + skip_if_rocm_multiprocess, +) +from torch.testing._internal.dist_utils import dist_init, INIT_METHOD_TEMPLATE +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +NUM_EM_ROW = 2 +D_SPARSE = 3 +D_DENSE = 2 +D_HID = 3 +D_OUT = 1 +NUM_TRAINERS = 4 +# Trainers + the master + the remote worker +WORLD_SIZE = NUM_TRAINERS + 2 +TRAINER_RANKS = list(range(NUM_TRAINERS)) +REMOTE_WORKER_RANK = TRAINER_RANKS[-1] + 1 +MASTER_RANK = REMOTE_WORKER_RANK + 1 + + +class DdpMode(enum.Enum): + # Don't apply DDP + NONE = enum.auto() + # Apply DDP to the top level nn.Module + OUTSIDE = enum.auto() + # Embed DDP inside the top level nn.Module + INSIDE = enum.auto() + + +def init_logger(): + logger = logging.getLogger(__name__) + level = logging.DEBUG if "debug" in os.environ else logging.INFO + logger.setLevel(level) + console = logging.StreamHandler() + formatter = logging.Formatter( + "%(asctime)s %(filename)s:%(lineno)s %(levelname)s p:%(processName)s t:%(threadName)s: %(message)s" + ) + console.setFormatter(formatter) + console.setLevel(level) + # add the handlers to the logger + logger.addHandler(console) + logger.propagate = False + return logger + + +gLogger = init_logger() + + +class FeatureSet(NamedTuple): + """A feature set has 2 types of features""" + + dense_features: torch.Tensor + sparse_features: torch.LongTensor + values: torch.Tensor + + +def _call_method(method, rref, *args, **kwargs): + return method(rref.local_value(), *args, **kwargs) + + +def _remote_method(method, rref, *args, **kwargs): + args_tup = tuple([method, rref] + list(args)) + return rpc.rpc_sync(rref.owner(), _call_method, args=args_tup, kwargs=kwargs) + + +def _remote_method_async(method, rref, *args, **kwargs): + args_tup = tuple([method, rref] + list(args)) + return rpc.rpc_async(rref.owner(), _call_method, args=args_tup, kwargs=kwargs) + + +class RemoteEM(nn.Module): + def __init__(self, num_embeddings: int, embedding_dim: int): + gLogger.info("Initing RemoteEM with %s %s", num_embeddings, embedding_dim) + super().__init__() + init_em = [0.5] * embedding_dim + self.em = nn.EmbeddingBag( + num_embeddings, + embedding_dim, + _weight=torch.tensor([init_em] * num_embeddings), + ) + + def forward(self, input: torch.Tensor): + gLogger.debug("Running RemoteEM.forward() on: %s", input) + return self.em(input, offsets=torch.LongTensor(range(input.shape[0]))) + + +# Return a linear module with predefined parameters. +def getLinear(d_in, d_out): + l = nn.Linear(d_in, d_out, bias=False) + w = torch.ones((d_out, d_in)) + w[0][0] = -1 + w.requires_grad_() + l.weight.data = w + return l + + +class RemoteNet(nn.Module): + def __init__(self, d_in: int, d_out: int): + gLogger.info("Initing RemoteNet with %s %s", d_in, d_out) + super().__init__() + self.fc = getLinear(d_in, d_out) + self.relu = nn.ReLU() + + def forward(self, input: torch.Tensor): + gLogger.debug("Running RemoteNet.forward() on: %s", input) + return self.relu(self.fc(input)) + + +class HybridModel(nn.Module): + def __init__( + self, + remote_em_rref: rpc.RRef, + remote_net_rref: rpc.RRef, + process_group_for_ddp: dist.ProcessGroup = None, + ): + super().__init__() + self.remote_em_rref = remote_em_rref + self.remote_net_rref = remote_net_rref + self.fc1 = getLinear(D_DENSE, D_DENSE) + self.fc2 = getLinear(D_HID, D_OUT) + + self.non_ddp_params = tuple(self.fc1.parameters()) + tuple( + self.fc2.parameters() + ) + self.ddp_params = () + + if process_group_for_ddp is not None: + self.non_ddp_params, self.ddp_params = ( + tuple(self.fc1.parameters()), + tuple(self.fc2.parameters()), + ) + gLogger.info("Use DDP for the second local net.") + self.fc2 = DistributedDataParallel( + self.fc2, check_reduction=True, process_group=process_group_for_ddp + ) + + gLogger.info( + "HybridModel has %s groups of parameters.", len(list(self.parameters())) + ) + + def forward(self, input: FeatureSet): + gLogger.debug("Running HybridModel.forward on %s", input) + sparse = _remote_method( + RemoteEM.forward, self.remote_em_rref, input.sparse_features + ) + # The same size of mini batch. + if sparse.shape[0] != input.dense_features.shape[0]: + raise AssertionError( + f"Expected sparse.shape[0] == input.dense_features.shape[0], " + f"got {sparse.shape[0]} != {input.dense_features.shape[0]}" + ) + dense = self.fc1(input.dense_features) + x = torch.cat((dense, sparse), 1) + gLogger.debug("Concatenated feature: %s", x) + x = _remote_method(RemoteNet.forward, self.remote_net_rref, x) + return self.fc2(x) + + +class Trainer: + def __init__( + self, + remote_em_rref: rpc.RRef, + remote_net_rref: rpc.RRef, + ddp_mode: DdpMode, + rank: int, + ): + self.rank = rank + self.trainer_group = ( + dist.new_group(TRAINER_RANKS) + if ddp_mode in (DdpMode.INSIDE, DdpMode.OUTSIDE) + else None + ) + self.remote_em_rref = remote_em_rref + self.remote_net_rref = remote_net_rref + self.hybrid_module = HybridModel( + self.remote_em_rref, + self.remote_net_rref, + self.trainer_group if ddp_mode == DdpMode.INSIDE else None, + ) + self.ddp_params, self.non_ddp_params = ( + self.hybrid_module.ddp_params, + self.hybrid_module.non_ddp_params, + ) + if ddp_mode == DdpMode.OUTSIDE: + gLogger.info("Wrapping the whole hybrid module into DDP.") + self.ddp_params += self.non_ddp_params + self.non_ddp_params = () + self.hybrid_module = DistributedDataParallel( + self.hybrid_module, + check_reduction=True, + process_group=self.trainer_group, + ) + gLogger.info( + "Succeeded in creating a HybridModel instance with " + "%s ddp params and %s other local params.", + len(self.ddp_params), + len(self.non_ddp_params), + ) + + def destroy_pg(self): + if self.trainer_group: + dist.destroy_process_group(self.trainer_group) + + def train_batch( + self, + mini_batch: FeatureSet, + trainer_has_less_inputs: bool, + simulate_uneven_inputs: bool, + ): + grads_dict = None + + if not simulate_uneven_inputs: + input_batches = [mini_batch] + else: + # Split into microbatches, and trim to simulate uneven inputs. + dense_features = mini_batch.dense_features + sparse_features = mini_batch.sparse_features + values = mini_batch.values + + dense_microbatch = torch.split(dense_features, 2) + sparse_microbatch = torch.split(sparse_features, 2) + values_microbatch = torch.split(values, 2) + batches = [] + for d, s, v in zip( + dense_microbatch, sparse_microbatch, values_microbatch, strict=True + ): + feature_set = FeatureSet(dense_features=d, sparse_features=s, values=v) + batches.append(feature_set) + + if trainer_has_less_inputs: + input_batches = batches[: len(batches) // 2] + gLogger.info( + "Trainer reduced input patches from %s " + "to %s to simulate uneven inputs.", + len(batches), + len(input_batches), + ) + else: + input_batches = batches + + with ( + self.hybrid_module.join() + if simulate_uneven_inputs + else contextlib.nullcontext() + ): + for b in input_batches: + with dist_autograd.context() as context_id: + output = self.hybrid_module.forward(b) + loss = (output * mini_batch.values).sum() + dist_autograd.backward(context_id, [loss]) + grads_dict = dist_autograd.get_gradients(context_id) + gLogger.info( + "Loss is %s for mini batch: %s. Grads dict has %s entries: %s", + loss, + mini_batch, + len(grads_dict), + grads_dict, + ) + return ( + tuple(grads_dict[param] for param in self.ddp_params), + tuple(grads_dict[param] for param in self.non_ddp_params), + ) + + +def get_training_examples(): + n = 16 + training_examples = FeatureSet( + dense_features=torch.zeros((n, D_DENSE)), + sparse_features=torch.zeros(n, dtype=torch.long), + values=torch.zeros(n), + ) + idx = 0 + # Every example has another one that has exactly the same features but an + # opposite value. Therefore, their grads cancel each other in all-reduce. + for value in (-1, 1): + for x in (-1.0 * value, 1.0 * value): + for y in (1.0 * value, -1.0 * value): + for z in (0, 1): + training_examples.dense_features[idx, :] = torch.tensor((x, y)) + training_examples.sparse_features[idx] = z + training_examples.values[idx] = value + idx += 1 + + # Split the examples among NUM_TRAINERS trainers + if n % NUM_TRAINERS != 0: + raise AssertionError( + f"Expected n % NUM_TRAINERS == 0, got {n} % {NUM_TRAINERS} = {n % NUM_TRAINERS}" + ) + examples_per_trainer = int(n / NUM_TRAINERS) + return [ + FeatureSet( + dense_features=training_examples.dense_features[ + start : start + examples_per_trainer, : + ], + sparse_features=training_examples.sparse_features[ + start : start + examples_per_trainer + ], + values=training_examples.values[start : start + examples_per_trainer], + ) + for start in range(0, n, examples_per_trainer) + ] + + +shutdown_signal = threading.Condition() + + +def set_shutdown_signal(): + global shutdown_signal + with shutdown_signal: + shutdown_signal.notify() + + +class DdpUnderDistAutogradTest(RpcAgentTestFixture): + @property + def world_size(self) -> int: + return WORLD_SIZE + + def remote_worker_name(self) -> str: + # The name has to be consistent with that in 'dist_init' decorator. + return f"worker{REMOTE_WORKER_RANK}" + + def trainer_name(self, rank): + # The name has to be consistent with that in 'dist_init' decorator. + return f"worker{rank}" + + def _remote_worker_process(self, ddp_mode): + gLogger.info("The remote worker is running.") + dist.init_process_group( + backend="gloo", + init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name), + world_size=self.world_size, + rank=self.rank, + ) + + if ddp_mode in (DdpMode.INSIDE, DdpMode.OUTSIDE): + # new_group needs to be called on ranks. + dist.new_group(TRAINER_RANKS) + + global shutdown_signal + with shutdown_signal: + shutdown_signal.wait() + gLogger.info("Exiting remote worker.") + dist.destroy_process_group() + + def _trainer_process(self, rank: int): + gLogger.info("Running the trainer #%s...", rank) + gLogger.info( + "Initing trainer process group by trainer #%s with ranks %s", + rank, + TRAINER_RANKS, + ) + dist.init_process_group( + backend="gloo", + init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name), + world_size=self.world_size, + rank=self.rank, + ) + + gLogger.info("Waiting for shutdown signal on trainer #%s...", rank) + + global shutdown_signal + with shutdown_signal: + shutdown_signal.wait() + gLogger.info("Exiting the trainer #%s...", rank) + dist.destroy_process_group() + + def _master_process(self, ddp_mode: DdpMode, simulate_uneven_inputs: bool): + gLogger.info("Running the master process...") + dist.init_process_group( + backend="gloo", + init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name), + world_size=self.world_size, + rank=self.rank, + ) + + remote_em_rref = rpc.remote( + self.remote_worker_name(), RemoteEM, args=(NUM_EM_ROW, D_SPARSE) + ) + remote_net_rref = rpc.remote( + self.remote_worker_name(), RemoteNet, args=(D_DENSE + D_SPARSE, D_HID) + ) + gLogger.info("Created remote rrefs on master") + self.do_test_on_master( + ddp_mode, simulate_uneven_inputs, remote_em_rref, remote_net_rref + ) + + def do_test_on_master( + self, + ddp_mode: DdpMode, + simulate_uneven_inputs: bool, + remote_em_rref: rpc.RRef, + remote_net_rref: rpc.RRef, + ): + if simulate_uneven_inputs: + gLogger.info( + "Running DDP + RPC test with simulating uneven inputs across trainers." + ) + + trainer_rrefs = [] + for rank in TRAINER_RANKS: + trainer = self.trainer_name(rank) + trainer_rrefs.append( + rpc.remote( + trainer, + Trainer, + args=(remote_em_rref, remote_net_rref, ddp_mode, rank), + ) + ) + + if ddp_mode in (DdpMode.INSIDE, DdpMode.OUTSIDE): + # new_group needs to be called on ranks. + dist.new_group(TRAINER_RANKS) + + training_examples = get_training_examples() + for _ in range(3): + futures = [] + num_trainers = len(trainer_rrefs) + for idx, trainer_rref in enumerate(trainer_rrefs): + # Half the trainers will deplete inputs earlier than the rest. + trainer_has_less_inputs = ( + simulate_uneven_inputs and idx < num_trainers // 2 + ) + futures.append( + _remote_method_async( + Trainer.train_batch, + trainer_rref, + training_examples[idx], + trainer_has_less_inputs, + simulate_uneven_inputs, + ) + ) + + for future in futures: + ddp_grads, non_ddp_grads = future.wait() + # When there are uneven inputs, it is not necessary that grads + # cancel each other out, since some trainers contribute 0 grad. + if not simulate_uneven_inputs: + for grad in ddp_grads: + self.assertEqual( + grad, + torch.zeros_like(grad), + msg=f"The grad for any ddp parameter should be zeros, because " + "the training examples' grads cancel each other. Received " + f"gradient {grad}", + ) + for grad in non_ddp_grads: + self.assertNotEqual( + grad, + torch.zeros_like(grad), + msg="The grad for any non-ddp parameter shouldn't be zeros", + ) + + # Destroy process groups + for trainer_rref in trainer_rrefs: + _remote_method_async(Trainer.destroy_pg, trainer_rref).wait() + + # Send shutdown signals. + for rank in TRAINER_RANKS: + trainer = self.trainer_name(rank) + rpc.rpc_sync(trainer, set_shutdown_signal, args=()) + + rpc.rpc_sync(self.remote_worker_name(), set_shutdown_signal, args=()) + + def _do_test(self, ddp_mode, simulate_uneven_inputs=False): + if self.rank == MASTER_RANK: + self._master_process(ddp_mode, simulate_uneven_inputs) + elif self.rank == REMOTE_WORKER_RANK: + self._remote_worker_process(ddp_mode) + elif self.rank in TRAINER_RANKS: + self._trainer_process(self.rank) + else: + raise RuntimeError(f"Unknown process rank: {self.rank}") + + @requires_gloo() + @dist_init + def test_backward_no_ddp(self): + self._do_test(DdpMode.NONE) + + @requires_gloo() + @dist_init + def test_backward_ddp_outside(self): + self._do_test(DdpMode.OUTSIDE) + + @requires_gloo() + @dist_init + def test_backward_ddp_outside_uneven_inputs(self): + self._do_test(DdpMode.OUTSIDE, simulate_uneven_inputs=True) + + @requires_gloo() + @dist_init + def test_backward_ddp_inside(self): + self._do_test(DdpMode.INSIDE) + + +# Common utils for both CPU and CUDA test suites +class CommonDdpComparisonTest(RpcAgentTestFixture): + @property + def world_size(self) -> int: + return NUM_TRAINERS + + def trainer_name(self, rank): + # The name has to be consistent with that in 'dist_init' decorator. + return f"worker{rank}" + + @staticmethod + def get_remote_grads(rref, context_id): + return dist_autograd.get_gradients(context_id)[rref.local_value().weight] + + +class DdpComparisonTest(CommonDdpComparisonTest): + def _run_test_ddp_comparision(self, simulate_uneven_inputs=False): + gLogger.info("Running trainer rank: %s", self.rank) + # Each trainer uses a different random seed. Otherwise, they are going + # to have exactly the same initial model parameters, input, and + # therefore grads. That means the grads will be the same before and + # after DDP's all-reduce. + torch.manual_seed(self.rank) + dist.init_process_group( + backend="gloo", + # Postfix file_name with "pg" since file_name is also used by RPC agent + init_method=INIT_METHOD_TEMPLATE.format(file_name=f"{self.file_name}_pg"), + world_size=self.world_size, + rank=self.rank, + ) + net = nn.Linear(2, 3) + ddp_net = DistributedDataParallel(net) + + # Odd ranks join early if simulate_uneven_inputs. + num_inputs = 1 + if simulate_uneven_inputs: + if self.rank % 2 == 0: + num_inputs += 2 + inputs_list = [torch.rand((3, 2)) for _ in range(num_inputs)] + + if simulate_uneven_inputs: + gLogger.info( + "Rank %s training with %s inputs.", self.rank, len(inputs_list) + ) + + # Use distributed autograd. The gradients will be in RPC context map. + grads_dict = {} + with ddp_net.join(simulate_uneven_inputs): + for i, inputs in enumerate(inputs_list): + with dist_autograd.context() as context_id: + loss = ddp_net(inputs).norm() + dist_autograd.backward(context_id, [loss]) + grads_dict = dist_autograd.get_gradients(context_id) + gLogger.info("Trainer #%s got grad dict: %s", self.rank, grads_dict) + + # Use local autograd. The gradients will be in each variable's '.grad'. + ddp_net.zero_grad() + loss = ddp_net(inputs).norm() + loss.backward() + + # The gradients should be the same + for param in net.parameters(): + self.assertTrue( + param in grads_dict, + msg=f"Param {param} is not in dist_auto grad dict {grads_dict} for iteration {i}", + ) + self.assertEqual( + grads_dict[param], + param.grad, + msg=f"The grads for param {param} are different under local " + f"and dist autograd: {param.grad} \n---\n {grads_dict[param]} for iteration {i}", + ) + dist.destroy_process_group() + + @requires_gloo() + @dist_init + def test_ddp_comparison(self): + self._run_test_ddp_comparision() + + @requires_gloo() + @dist_init + def test_ddp_comparison_uneven_inputs(self): + # test with simulating uneven inputs in DDP + self._run_test_ddp_comparision(simulate_uneven_inputs=True) + + @requires_gloo() + @dist_init + def test_ddp_dist_autograd_sparse_grads(self): + # Each trainer uses a different random seed. Otherwise, they are going + # to have exactly the same initial model parameters, input, and + # therefore grads. That means the grads will be the same before and + # after DDP's all-reduce. + torch.manual_seed(self.rank) + dist.init_process_group( + backend="gloo", + init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name), + world_size=self.world_size, + rank=self.rank, + ) + + model = nn.EmbeddingBag(10, 3, sparse=True) + ddp_model = DistributedDataParallel(model) + + # Different inputs for each + input = torch.LongTensor(10).random_(0, 10) + offsets = torch.LongTensor([0, 4]) + + # Run local. + loss = ddp_model(input, offsets).sum() + loss.backward() + + with dist_autograd.context() as context_id: + loss = ddp_model(input, offsets).sum() + dist_autograd.backward(context_id, [loss]) + grads_dict = dist_autograd.get_gradients(context_id) + self.assertEqual(1, len(grads_dict)) + self.assertEqual(model.weight.grad, grads_dict[model.weight]) + + @requires_gloo() + @dist_init + def test_ddp_dist_autograd_local_vs_remote(self): + # Each trainer uses a different random seed. Otherwise, they are going + # to have exactly the same initial model parameters, input, and + # therefore grads. That means the grads will be the same before and + # after DDP's all-reduce. + torch.manual_seed(self.rank) + dist.init_process_group( + backend="gloo", + init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name), + world_size=self.world_size, + rank=self.rank, + ) + + # Use two different remote device input string, w/ and w/o the default + # device string "cpu", respectively. + for remote_device in ["worker0/cpu", "worker0"]: + remote_layer1 = RemoteModule( + remote_device=remote_device, module_cls=nn.Linear, args=(10, 5, False) + ) + layer1 = nn.Linear(10, 5, False) + # Start with the same parameters for remote and local + layer1.weight = remote_layer1.module_rref.to_here().weight + + # Run local case. + layer2 = nn.Linear(5, 1) + inputs = torch.rand((10, 10)) + ddp_model = DistributedDataParallel(layer2) + loss = ddp_model(layer1(inputs)).sum() + loss.backward() + + # Run remote case. + with dist_autograd.context() as context_id: + loss = ddp_model(remote_layer1(inputs)).sum() + dist_autograd.backward(context_id, [loss]) + grads_dict = dist_autograd.get_gradients(context_id) + dist.barrier() + self.assertEqual(layer2.weight.grad, grads_dict[layer2.weight]) + self.assertEqual( + layer1.weight.grad, + rpc.rpc_sync( + "worker0", + CommonDdpComparisonTest.get_remote_grads, + args=(remote_layer1.module_rref, context_id), + ), + ) + + +class CudaDdpComparisonTest(CommonDdpComparisonTest): + @skip_if_lt_x_gpu(NUM_TRAINERS) + @requires_nccl() + @dist_init + @skip_if_rocm_multiprocess + def test_ddp_dist_autograd_local_vs_remote_gpu(self): + # Each trainer uses a different random seed. Otherwise, they are going + # to have exactly the same initial model parameters, input, and + # therefore grads. That means the grads will be the same before and + # after DDP's all-reduce. + torch.manual_seed(self.rank) + dist.init_process_group( + backend="gloo", + init_method=INIT_METHOD_TEMPLATE.format(file_name=self.file_name), + world_size=self.world_size, + rank=self.rank, + ) + + remote_layer1 = RemoteModule( + remote_device="worker0/cpu", module_cls=nn.Linear, args=(10, 7, False) + ) + layer1 = nn.Linear(10, 7, False) + # Start with the same parameters for remote and local + layer1.weight = remote_layer1.module_rref.to_here().weight + + layer2 = nn.Linear(7, 5).cuda(self.rank) + ddp_layer2 = DistributedDataParallel(layer2, device_ids=[self.rank]) + + remote_layer3 = RemoteModule( + remote_device="worker0/cpu", module_cls=nn.Linear, args=(5, 3, False) + ) + layer3 = nn.Linear(5, 3, False) + # Start with the same parameters for remote and local + layer3.weight = remote_layer3.module_rref.to_here().weight + + layer4 = nn.Linear(3, 1).cuda(self.rank) + ddp_layer4 = DistributedDataParallel(layer4, device_ids=[self.rank]) + + # Run local case. + inputs = torch.rand((10, 10)) + loss = ddp_layer4( + layer3(ddp_layer2(layer1(inputs).cuda(self.rank)).cpu()).cuda(self.rank) + ).sum() + loss.backward() + + # Run remote case. + with dist_autograd.context() as context_id: + loss = ddp_layer4( + remote_layer3( + ddp_layer2(remote_layer1(inputs).cuda(self.rank)).cpu() + ).cuda(self.rank) + ).sum() + dist_autograd.backward(context_id, [loss]) + grads_dict = dist_autograd.get_gradients(context_id) + dist.barrier() + self.assertEqual( + layer1.weight.grad, + rpc.rpc_sync( + "worker0", + CommonDdpComparisonTest.get_remote_grads, + args=(remote_layer1.module_rref, context_id), + ), + ) + self.assertEqual(layer2.weight.grad, grads_dict[layer2.weight]) + self.assertEqual( + layer3.weight.grad, + rpc.rpc_sync( + "worker0", + CommonDdpComparisonTest.get_remote_grads, + args=(remote_layer3.module_rref, context_id), + ), + ) + self.assertEqual(layer4.weight.grad, grads_dict[layer4.weight]) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/distributed_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/distributed_test.py new file mode 100644 index 0000000000000000000000000000000000000000..3f9695617097687fec843c0a759ca5fe5f0acf80 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/distributed_test.py @@ -0,0 +1,10836 @@ +# mypy: allow-untyped-defs + +import copy +import itertools +import json +import math +import operator +import os +import random +import re +import sys +import tempfile +import time +import unittest +from collections import defaultdict, namedtuple, OrderedDict +from collections.abc import Callable +from contextlib import contextmanager, nullcontext +from dataclasses import dataclass +from datetime import timedelta +from functools import reduce +from typing import Any, NamedTuple + +import numpy as np + +import torch +import torch.cuda +import torch.distributed as dist +import torch.distributed.algorithms.model_averaging.averagers as averagers +import torch.distributed.algorithms.model_averaging.hierarchical_model_averager as hierarchicalSGD +import torch.distributed.algorithms.model_averaging.utils as model_averaging_utils +import torch.distributed.optim.post_localSGD_optimizer as post_localSGD_optimizer +import torch.nn as nn +import torch.nn.functional as F +from torch._utils_internal import ( + TEST_MASTER_ADDR as MASTER_ADDR, + TEST_MASTER_PORT as MASTER_PORT, +) +from torch.autograd import DeviceType +from torch.cuda.amp import autocast, GradScaler +from torch.distributed.algorithms.ddp_comm_hooks import ( + default_hooks as default, + post_localSGD_hook as post_localSGD, + powerSGD_hook as powerSGD, + quantization as quantization_hooks, +) +from torch.distributed.distributed_c10d import ( + _get_default_group, + _get_pg_config, + get_world_size, +) +from torch.distributed.optim import _apply_optimizer_in_backward +from torch.distributed.utils import ( + _sync_module_states, + _verify_param_shape_across_processes, +) +from torch.nn.parallel import DistributedDataParallel +from torch.nn.parallel.distributed import ( + _BucketCapacityConfig, + _DEFAULT_BUCKET_CAP_MB, + _dump_DDP_relevant_env_vars, + _MB_TO_BYTES, + _MixedPrecision, +) +from torch.profiler import ExecutionTraceObserver, ProfilerActivity +from torch.testing._internal.common_distributed import ( + captured_output, + cleanup_temp_dir, + DistTestCases, + init_multigpu_helper, + initialize_temp_directories, + MultiProcessTestCase, + nccl_skip_if_lt_x_gpu, + require_n_gpus_for_nccl_backend, + requires_nccl_version, + simple_sparse_reduce_tests, + skip_if_lt_x_gpu, + skip_if_no_gpu, + skip_if_odd_worldsize, + skip_if_rocm_multiprocess, + skip_if_small_worldsize, + TEST_SKIPS, + verify_ddp_error_logged, + with_dist_debug_levels, + with_nccl_blocking_wait, +) +from torch.testing._internal.common_utils import ( + FILE_SCHEMA, + instantiate_parametrized_tests, + IS_FBCODE, + IS_MACOS, + IS_SANDCASTLE, + IS_WINDOWS, + skip_but_pass_in_sandcastle, + skip_but_pass_in_sandcastle_if, + TemporaryFileName, +) +from torch.utils._python_dispatch import TorchDispatchMode +from torch.utils.data.distributed import DistributedSampler + + +try: + import torchvision + + HAS_TORCHVISION = True +except Exception: # Covering both ImportError and RuntimeError + HAS_TORCHVISION = False + +if sys.platform == "win32": + import msvcrt +else: + import fcntl + + +class NetWithBuffers(nn.Module): + def __init__(self) -> None: + super().__init__() + self.a = nn.Linear(10, 10, bias=False) + self.b = nn.Linear(10, 1, bias=False) + self.register_buffer("buffer", torch.randn(1, 2)) + + def forward(self, x): + self.buffer.add_(1) + return self.b(self.a(x)) + + +class Foo: + def __init__(self, x): + # Can be tensor or int + self.x = x + + def __eq__(self, other): + def eq(value, other): + if isinstance(value, torch.Tensor): + return torch.equal(value, other) + return value == other + + for attr, value in self.__dict__.items(): + other_value = other.__dict__[attr] + if not eq(value, other_value): + return False + return True + + +f = Foo(10) +f.bar = 1 + + +# Defer instantiation until the seed is set so that randn() returns the same +# values in all processes. +def create_collectives_object_test_list(): + return [ + {"key1": 3, "key2": 4, "key3": {"nested": True}}, + f, + Foo(torch.randn(3, 3)), + "foo", + [1, 2, True, "string", [4, 5, "nested"]], + ] + + +# Allowlist of distributed backends where profiling collectives is supported. +PROFILING_SUPPORTED_BACKENDS = [ + dist.Backend.NCCL, + dist.Backend.GLOO, + dist.Backend.MPI, + dist.Backend.UCC, +] + +# Allowlist of distributed backends where profiling is supported with use_cuda=True +CUDA_PROFILING_SUPPORTED_BACKENDS = [ + dist.Backend.GLOO, + dist.Backend.MPI, + dist.Backend.NCCL, + dist.Backend.UCC, +] + +# Allowlist of distributed backends where profiling is supported for p2p ops +SEND_RECV_PROFILING_SUPPORTED_BACKENDS = [ + dist.Backend.MPI, + dist.Backend.GLOO, + dist.Backend.NCCL, + dist.Backend.UCC, +] + +# Dummy NamedTuple data structures to test DDP support for NamedTuple types. +EXPECTED_FIELDS = ("a", "b") +TestNamedTupleInput_0 = namedtuple("NamedTuple", EXPECTED_FIELDS) + + +class TestNamedTupleInput_1(NamedTuple): + a: torch.tensor + b: torch.tensor + + +skipIfNoTorchVision = skip_but_pass_in_sandcastle_if( + not HAS_TORCHVISION, "no torchvision" +) + +BACKEND = os.environ["BACKEND"] +INIT_METHOD = os.getenv("INIT_METHOD", "env://") + + +def get_profiling_event(event_name, profiler, dedup_gpu_user_annotation=False): + event_list = ( + profiler.events() + if isinstance(profiler, torch.profiler.profile) + else profiler.function_events + ) + return [ + event + for event in event_list + if ( + (event.name.endswith(event_name) or event.name.startswith(event_name)) + and ( + not dedup_gpu_user_annotation + or event.device_type not in [DeviceType.CUDA, DeviceType.XPU] + ) + ) + ] + + +def get_profiler_nccl_meta(prof): + """Torch profiler includes nccl metadata in an inserted operator called "record_param_comms" + We will need to test metadata obtained from profiler here""" + with TemporaryFileName(mode="w+t", suffix=".json") as trace_file: + prof.export_chrome_trace(trace_file) + with open(trace_file) as f: + events = json.load(f)["traceEvents"] + print(f"Trace saved to {trace_file}") + + return [e for e in events if e.get("name") == "record_param_comms"] + + +# Base error message substring on unfinished reductions. +ddp_prev_reduction_unfinished_str = ( + "Expected to have finished reduction in the prior iteration" +) +# Error message substring when find_unused_parameters=True has not been passed +ddp_recommend_find_unused_params_str = ( + "passing the keyword argument `find_unused_parameters=True`" +) +# Error message substring when find_unused_parameters=True is enabled +ddp_find_unused_params_enabled_str = "Since `find_unused_parameters=True` is enabled" +# Error message substring for possibility of not all model outputs being used +# in loss computation +ddp_outputs_not_used_in_loss_str = ( + "`forward` function outputs participate in calculating loss" +) +# Error message substring suggesting to use TORCH_DISTRIBUTED_DEBUG +ddp_suggest_debug_mode_str = ( + "set the environment variable TORCH_DISTRIBUTED_DEBUG to either INFO or DETAIL" +) + + +class DDPUnevenTestInput(NamedTuple): + name: str + model: nn.Module + inp: torch.Tensor | tuple + sync_interval: int + throw_on_early_termination: bool = False + hook: Callable = None + state: Any = None + + +class _FC2(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc = nn.Linear(10, 50, bias=True) + self.fc.bias.requires_grad = False + + def forward(self, x): + x = self.fc(x) + return x + + +class Net(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(2, 10, bias=False) + self.fc2 = _FC2() + self.fc3 = nn.Linear(50, 4, bias=False) + self.relu = nn.ReLU() + self.no_grad_param = nn.Parameter( + torch.tensor([2, 2]).long(), requires_grad=False + ) + + def forward(self, x): + x = self.relu(self.fc1(x)) + x = self.relu(self.fc2(x)) + x = self.fc3(x) + return F.softmax(x, dim=1) + + +class LargeNet(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(1000, 2000, bias=False) + self.fc2 = nn.Linear(2000, 500, bias=False) + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + return x + + +class Task(nn.Module): + def __init__(self) -> None: + super().__init__() + self.p = nn.Parameter(torch.ones(2, 2)) + + def forward(self, x): + return self.p + x + + +class BatchNormNet(nn.Module): + def __init__(self, affine=True): + super().__init__() + self.fc1 = nn.Linear(2, 40, bias=False) + self.bn = nn.BatchNorm1d(4, affine=affine) + self.fc2 = nn.Linear(40, 4, bias=False) + + def forward(self, x): + x = torch.reshape(self.fc1(x), (-1, 4, 10)) + x = self.bn(x) + x = torch.reshape(x, (-1, 40)) + x = self.fc2(x) + return F.softmax(x, dim=1) + + +class UnusedParamTwoLinLayerNet(nn.Module): + def __init__(self) -> None: + super().__init__() + self.a = nn.Linear(10, 10, bias=False) + self.b = nn.Linear(10, 10, bias=False) + self.c = nn.Linear(5, 5, bias=False) + + def forward(self, x): + a = self.a(x) + b = self.b(x) + return (a, b) + + +class DictOutputModule(nn.Module): + def __init__(self) -> None: + super().__init__() + self.module = UnusedParamTwoLinLayerNet() + + def forward(self, x): + predictions = self.module(x) + loss = (predictions[0] + predictions[1]).sum() + return { + "predictions": predictions, + "loss": loss, + } + + +class TwoLinLayerNet(nn.Module): + def __init__(self) -> None: + super().__init__() + self.a = nn.Linear(10, 10, bias=False) + self.b = nn.Linear(10, 1, bias=False) + + def forward(self, x): + a = self.a(x) + b = self.b(x) + return (a, b) + + +class EmbeddingNetDifferentParams(nn.Module): + """ + A module containing an embedding with different dimension or different # of + parameters depending on the rank. + """ + + def __init__(self, rank, diff_num_params=False): + super().__init__() + embedding_dim = 500 if diff_num_params or rank == 0 else 50 + self.embedding = nn.Embedding(num_embeddings=10, embedding_dim=embedding_dim) + self.lin = nn.Linear(embedding_dim, 1) + if diff_num_params: + self.lin2 = nn.Linear(1, 1, bias=False) + + def forward(self, x): + x = self.embedding(x) + return self.lin(x) + + +class ControlFlowToyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin1 = nn.Linear(10, 10, bias=False) + self.lin2 = nn.Linear(10, 10, bias=False) + + def forward(self, x): + # Second layer is used dependent on input x. + use_second_layer = torch.equal(x, torch.ones(20, 10, device=x.device)) + if use_second_layer: + return self.lin2(F.relu(self.lin1(x))) + else: + return F.relu(self.lin1(x)) + + +default_pg_timeout = 60 + +CUSTOM_PG_TIMEOUT = { + # This test runs slowly and needs additional time to complete, otherwise can + # be taken down by TORCH_NCCL_ASYNC_ERROR_HANDLING + "test_ddp_uneven_inputs": 300, + # This test has a short timeout since it tests being taken down by + # TORCH_NCCL_ASYNC_ERROR_HANDLING which we want to happen quickly. + "test_ddp_model_diff_across_ranks": 5, + # This test has a short timeout since it tests being taken down by + # TORCH_NCCL_ASYNC_ERROR_HANDLING which we want to happen quickly. + "test_ddp_has_finalized": 5, +} + + +def require_backend_is_available(backends): + def check(backend): + if backend == dist.Backend.GLOO: + return dist.is_gloo_available() + if backend == dist.Backend.NCCL: + return dist.is_nccl_available() + if backend == dist.Backend.MPI: + return dist.is_mpi_available() + if backend == dist.Backend.UCC: + return dist.is_ucc_available() + if backend in DistTestCases.backend_feature["plugin"]: + return True + return False + + if BACKEND not in backends: + return skip_but_pass_in_sandcastle( + f"Test requires backend {BACKEND} to be one of {backends}" + ) + + if not check(dist.Backend(BACKEND)): + return skip_but_pass_in_sandcastle( + f"Test requires backend {BACKEND} to be available" + ) + return lambda func: func + + +def require_world_size(world_size): + if int(os.environ["WORLD_SIZE"]) < world_size: + return skip_but_pass_in_sandcastle( + f"Test requires world size of {world_size:d}" + ) + return lambda func: func + + +def require_exact_world_size(world_size): + if int(os.environ["WORLD_SIZE"]) != world_size: + return skip_but_pass_in_sandcastle( + f"Test requires an exact world size of {world_size:d}" + ) + return lambda func: func + + +@contextmanager +def _lock(): + TEMP_DIR = os.environ["TEMP_DIR"] + lockfile = os.path.join(TEMP_DIR, "lockfile") + with open(lockfile, "w") as lf: + try: + if sys.platform == "win32": + msvcrt.locking(lf.fileno(), msvcrt.LK_RLCK, 1) + yield + else: + fcntl.flock(lf.fileno(), fcntl.LOCK_EX) + yield + finally: + if sys.platform == "win32": + msvcrt.locking(lf.fileno(), msvcrt.LK_UNLCK, 1) + else: + fcntl.flock(lf.fileno(), fcntl.LOCK_UN) + lf.close() + + +@contextmanager +def _rank_temp_file(): + if dist.get_rank() == 0: + fd, name = tempfile.mkstemp() + os.close(fd) + else: + name = None + object_list = [name] + dist.broadcast_object_list(object_list) + name = object_list[0] + try: + yield name + finally: + if dist.get_rank() == 0: + os.remove(name) + + +def _build_tensor(size, value=None, dtype=torch.float, device_id=None): + if value is None: + value = size + if device_id is None: + return torch.empty(size, size, size, dtype=dtype).fill_(value) + else: + return torch.empty(size, size, size, dtype=dtype).fill_(value).cuda(device_id) + + +def _build_multidim_tensor(dim, dim_size, value=None, dtype=torch.float): + if value is None: + value = dim + return torch.empty(size=[dim_size for _ in range(dim)], dtype=dtype).fill_(value) + + +def _create_autograd_profiler(): + return torch.autograd.profiler.profile(record_shapes=True) + + +def _create_torch_profiler(): + return torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + ], + record_shapes=True, + ) + + +class Barrier: + barrier_id = 0 + + @classmethod + def init(cls): + cls.barrier_id = 0 + barrier_dir = os.path.join(os.environ["TEMP_DIR"], "barrier") + for f_name in os.listdir(barrier_dir): + os.unlink(os.path.join(barrier_dir, f_name)) + + @classmethod + def sync(cls, wait_for=None, timeout=10): + if wait_for is None: + wait_for = dist.get_world_size() + cls.barrier_id += 1 + barrier_dir = os.path.join(os.environ["TEMP_DIR"], "barrier") + pid = str(os.getpid()) + barrier_file = os.path.join(barrier_dir, pid) + with _lock(): + with open(barrier_file, "w") as f: + f.write(str(cls.barrier_id)) + + start_time = time.time() + while True: + arrived = 0 + with _lock(): + for f_name in os.listdir(barrier_dir): + with open(os.path.join(barrier_dir, f_name)) as f: + data = f.read() + if int(data) >= cls.barrier_id: + arrived += 1 + if arrived == wait_for: + break + + if time.time() - start_time > timeout: + raise RuntimeError("barrier timeout") + time.sleep(0.1) + + +class TestDistBackend(MultiProcessTestCase): + @classmethod + def setUpClass(cls): + os.environ["MASTER_ADDR"] = str(MASTER_ADDR) + # Not setting MASTER_PORT and get a random free port + super().setUpClass() + + def setUp(self): + super().setUp() + # initialize temp directories + initialize_temp_directories() + # initialize Barrier + Barrier.init() + # Skip return code checking for following tests as they are expected to + # crash a process due to TORCH_NCCL_ASYNC_ERROR_HANDLING. + self.skip_return_code_checks = [self.test_ddp_has_finalized.__wrapped__] + + def tearDown(self): + cleanup_temp_dir() + super().tearDown() + + @property + def init_method(self): + return f"{FILE_SCHEMA}{self.file_name}" + + @property + def destroy_pg_upon_exit(self) -> bool: + # Overriding base test class: do not auto destroy PG upon exit. + return False + + @classmethod + def _run(cls, rank, test_name, file_name, pipe, **kwargs): + if BACKEND == "nccl" and not torch.cuda.is_available(): + sys.exit(TEST_SKIPS["no_cuda"].exit_code) + self = cls(test_name) + self.rank = rank + self.file_name = file_name + + if torch.cuda.is_available() and torch.cuda.device_count() < int( + self.world_size + ): + sys.exit(TEST_SKIPS[f"multi-gpu-{self.world_size}"].exit_code) + try: + pg_timeout_seconds = CUSTOM_PG_TIMEOUT.get(test_name, default_pg_timeout) + timeout = timedelta(seconds=pg_timeout_seconds) + dist.init_process_group( + init_method=self.init_method, + backend=BACKEND, + world_size=int(self.world_size), + rank=self.rank, + timeout=timeout, + ) + except RuntimeError as e: + if "recompile" in e.args[0]: + sys.exit(TEST_SKIPS["backend_unavailable"].exit_code) + + raise + + # Execute barrier prior to running test to ensure that every process + # has finished initialization and that the following test + # immediately exiting due to a skip doesn't cause flakiness. + self._barrier() + + self.run_test(test_name, pipe) + self._barrier() + dist.destroy_process_group() + sys.exit(0) + + # Needed since MultiProcessTestCase assumes a world_size of 4, but we + # run these tests under other various world_sizes. + @property + def world_size(self): + return os.environ["WORLD_SIZE"] + + +class DistributedTest: + class _DistTestBase: + def _barrier(self, *args, **kwargs): + Barrier.sync(*args, **kwargs) + + def _init_group_test(self, **kwargs): + group = [1, 2] + group_id = dist.new_group(group, **kwargs) + rank = dist.get_rank() + if rank not in group: + return ([], None, rank) + + return (group, group_id, rank) + + def _init_full_group_test(self, **kwargs): + group = list(range(dist.get_world_size())) + group_id = dist.new_group(**kwargs) + rank = dist.get_rank() + return (group, group_id, rank) + + def _init_global_test(self): + group = list(range(dist.get_world_size())) + group_id = dist.group.WORLD + rank = dist.get_rank() + return (group, group_id, rank) + + def _verify_buffers_equal(self, m1, m2): + # verify buffers across models + m1_buf_dict = dict(m1.module.named_buffers()) + for name, buf in m2.module.named_buffers(): + self.assertEqual(buf, m1_buf_dict[name]) + + # Verify buffers across ranks. + m1_buffers = list(m1.buffers()) + m2_buffers = list(m2.buffers()) + for buf1, buf2 in zip(m1_buffers, m2_buffers, strict=True): + gathered_bufs = [ + torch.empty_like(buf1) for _ in range(dist.get_world_size()) + ] + dist.all_gather(gathered_bufs, buf1) + gathered_bufs_m2 = [ + torch.empty_like(buf2) for _ in range(dist.get_world_size()) + ] + for b in gathered_bufs: + self.assertEqual(b, buf1) + dist.all_gather(gathered_bufs_m2, buf2) + for b in gathered_bufs_m2: + self.assertEqual(b, buf2) + + def _sanity_check_profiler_nccl_meta(self, nccl_meta_events): + """Torch profiler includes nccl metadata in an inserted operator called "record_param_comms" + We test for basic fields in this profiler event that correspond to the nccl communication + collectives""" + per_coll_meta = defaultdict(list) + for e in nccl_meta_events: + args = e.get("args", {}) + collname = args.get("Collective name", "") + self.assertNotEqual(collname, "") + self.assertNotEqual(args.get("dtype", ""), "") + + per_coll_meta[collname].append(args) + if collname == "wait": + continue + + self.assertEqual(args["Process Group Description"], "default_pg") + self.assertNotEqual(args["Process Group Ranks"], "") + + self.assertGreaterEqual(args.get("In msg nelems", -1), 0) + self.assertGreaterEqual(args.get("Out msg nelems", -1), 0) + self.assertGreaterEqual(args.get("Group size", -1), 0) + self.assertGreaterEqual(args.get("Global rank start", -1), 0) + self.assertGreaterEqual(args.get("Global rank stride", -1), 0) + + # print(per_coll_meta) + return per_coll_meta + + def test_dump_DDP_relevant_env_vars(self): + with captured_output() as (out, _): + _dump_DDP_relevant_env_vars() + lines = out.getvalue().splitlines() + + def format_line(var): + return f"env:{var}={os.environ.get(var, 'N/A')}" + + # Check relevant env vars + vars = [ + "MASTER_ADDR", + "MASTER_PORT", + "WORLD_SIZE", + "NCCL_TOPO_DUMP_FILE", # N/A + "TORCH_NCCL_ASYNC_ERROR_HANDLING", + ] + for var in vars: + line = format_line(var) + self.assertIn(line, lines) + # Check irrelevant env vars + vars = [ + "xxx", + "yyy", + "zzz", + ] + for var in vars: + line = format_line(var) + self.assertNotIn(line, lines) + + # GET RANK + def test_get_rank(self): + test_dir = os.path.join(os.environ["TEMP_DIR"], "test_dir") + pid = str(os.getpid()) + num_processes = dist.get_world_size() + with open(os.path.join(test_dir, pid), "w") as f: + f.write(str(dist.get_rank())) + + self._barrier() + + all_ranks = set() + for f_name in os.listdir(test_dir): + with open(os.path.join(test_dir, f_name)) as f: + all_ranks.add(int(f.read())) + self.assertEqual(len(all_ranks), num_processes) + + self._barrier() + + if dist.get_rank() == 0: + for f_name in os.listdir(test_dir): + os.unlink(os.path.join(test_dir, f_name)) + + self._barrier() + + def test_get_backend(self): + if dist.get_world_size() > 2: + group = [1, 2] + else: + group = [0, 1] + group_id = dist.new_group(group) + backend_str = BACKEND.lower() + self.assertEqual(dist.get_backend(), backend_str) + if dist.get_rank() in group: + self.assertEqual(dist.get_backend(group_id), backend_str) + else: + with self.assertRaisesRegex( + ValueError, "Invalid process group specified" + ): + dist.get_backend(group_id) + + def test_Backend_enum_class(self): + # test parsing + backend = BACKEND.lower() + self.assertEqual(dist.Backend(BACKEND.upper()), backend) + self.assertEqual(dist.Backend(BACKEND), backend) + with self.assertRaises(ValueError): + dist.Backend(None) + with self.assertRaises(ValueError): + dist.Backend(3) + with self.assertRaises(ValueError): + dist.Backend(["gloo"]) + + # Test destroy + def test_destroy_group(self): + if dist.get_world_size() > 2: + group = [1, 2] + else: + group = [0, 1] + group_id = dist.new_group(group) + self._barrier() + dist.destroy_process_group(group_id) + + # Test get rank and size of group + def test_get_rank_size_group(self): + if dist.get_world_size() > 2: + group = [1, 2] + else: + group = [0, 1] + group_id = dist.new_group(group) + if dist.get_rank() in group: + self.assertEqual(dist.get_world_size(group_id), 2) + self.assertTrue(dist.get_rank(group_id) in list(range(2))) + else: + self.assertEqual(dist.get_world_size(group_id), -1) + self.assertEqual(dist.get_rank(group_id), -1) + + # Test destroy full groups + def test_destroy_full_group(self): + _, group_id, _ = self._init_full_group_test() + self._barrier() + dist.destroy_process_group(group_id) + + # Test get rank and size of full group + def test_get_rank_size_full_group(self): + _, group_id, _ = self._init_full_group_test() + self.assertEqual(dist.get_world_size(group_id), dist.get_world_size()) + self.assertEqual(dist.get_rank(group_id), dist.get_rank()) + + def _test_barrier_timeout(self, group_id, timeout): + local_rank = dist.get_rank(group_id) + + # Only execute barrier on rank == 0, causing it to timeout + if local_rank == 0: + expected_time = time.time() + timeout.total_seconds() + # In debug mode, we execute a monitored_barrier before the + # collective, so assert on that. + if dist.get_debug_level() == dist.DebugLevel.DETAIL: + exception_ctx = self.assertRaisesRegex( + Exception, "failed to pass monitoredBarrier" + ) + else: + exception_ctx = self.assertRaisesRegex( + Exception, " (Timed out|closed|timeout) " + ) + with exception_ctx: + dist.barrier(group_id) + self.assertGreaterAlmostEqual(time.time(), expected_time, delta=0.1) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo", "Only gloo backend supports timeouts" + ) + @skip_but_pass_in_sandcastle_if( + not INIT_METHOD.startswith("file://"), + "Requires file:// initialization method. " + + "Both tcp:// and env:// rely on the TCP store for which " + "reinitialization has proven racy.", + ) + def test_barrier_timeout_global(self): + dist.destroy_process_group() + + # Explicitly pass world size to the barrier because we've + # just destroyed any state in torch.distributed. + self._barrier(wait_for=int(os.environ["WORLD_SIZE"])) + + # Reinitialize global process group + timeout = timedelta(seconds=1) + dist.init_process_group( + init_method=INIT_METHOD, + backend=BACKEND, + world_size=int(os.environ["WORLD_SIZE"]), + rank=self.rank, + timeout=timeout, + ) + self._test_barrier_timeout(dist.group.WORLD, timeout) + + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo", "Only gloo backend supports timeouts" + ) + def test_barrier_timeout_group(self): + timeout = timedelta(seconds=5) + _, group_id, _ = self._init_group_test(timeout=timeout) + if group_id is not None: + self._test_barrier_timeout(group_id, timeout) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo", "Only gloo backend supports timeouts" + ) + def test_barrier_timeout_full_group(self): + timeout = timedelta(seconds=1) + _, group_id, _ = self._init_full_group_test(timeout=timeout) + if group_id is not None: + self._test_barrier_timeout(group_id, timeout) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo", "Only gloo backend supports timeouts" + ) + def test_barrier_timeout_arg(self): + """Test that the timeout argument to barrier() overrides PG default. + + Create a PG with a large default timeout, then have only rank 0 + call barrier with a tiny timeout. The barrier should time out using + the per-call timeout (1ms) rather than the PG default (300s). + """ + pg = dist.new_group(timeout=timedelta(seconds=300)) + + if dist.get_rank() == 0: + with self.assertRaisesRegex(RuntimeError, "Timed out waiting 1ms"): + dist.barrier(group=pg, timeout=timedelta(seconds=0.001)) + + dist.destroy_process_group(pg) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @require_world_size(4) + @skip_if_lt_x_gpu(2) + def test_new_subgroups(self): + subgroup_size = 2 + cur_subgroup, subgroups = dist.new_subgroups(subgroup_size) + + world_size = dist.get_world_size() + self.assertEqual(cur_subgroup.size(), subgroup_size) + self.assertEqual(len(subgroups), world_size / subgroup_size) + self.assertFalse(dist._rank_not_in_group(cur_subgroup)) + + for subgroup in subgroups: + dist.destroy_process_group(subgroup) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @require_exact_world_size(4) + def test_new_subgroups_with_group_param(self): + # Initialize global test environment + self._init_global_test() + # Set up GPU devices for each rank + init_multigpu_helper(dist.get_world_size(), BACKEND) + # Create two subgroups: one with ranks [0,2] and another with ranks [1,3] + cur_subgroup, subgroups = dist.new_subgroups_by_enumeration( + ranks_per_subgroup_list=[[0, 2], [1, 3]] + ) + + # Further divide the current subgroup into sub-subgroups of size 1 + cur_sub_subgroup, sub_subgroups = dist.new_subgroups( + group_size=1, group=cur_subgroup + ) + # Verify we have 2 sub-subgroups (one for each rank in the original subgroup) + self.assertEqual(len(sub_subgroups), 2) + # Verify the current process's sub-subgroup has size 1 + self.assertEqual(cur_sub_subgroup.size(), 1) + # Verify the current process is in its assigned sub-subgroup + self.assertFalse(dist._rank_not_in_group(group=cur_sub_subgroup)) + + # Clean up by destroying all created process groups + for sub_subgroup in sub_subgroups: + dist.destroy_process_group(sub_subgroup) + + for subgroup in subgroups: + dist.destroy_process_group(subgroup) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @skip_if_no_gpu + def test_new_subgroups_group_size_exceeds_world_size(self): + with self.assertRaisesRegex(ValueError, "must not exceed"): + dist.new_subgroups(100) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @require_world_size(4) + @skip_if_lt_x_gpu(4) + def test_new_subgroups_world_size_not_divisible_by_group_size(self): + expected_msg = f"The world size ({dist.get_world_size()}) must be divisible by 'group_size=3'" + with self.assertRaisesRegex( + ValueError, + re.escape(expected_msg), + ): + dist.new_subgroups(3) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @require_world_size(4) + @skip_if_lt_x_gpu(4) + def test_new_subgroups_by_enumeration(self): + _group, _group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + cur_subgroup, subgroups = dist.new_subgroups_by_enumeration( + ranks_per_subgroup_list=[[0, 2], [1, 3]] + ) + if device_id >= 4: + self.assertIsNone(cur_subgroup) + else: + self.assertEqual(cur_subgroup.size(), 2) + self.assertEqual(len(subgroups), 2) + if device_id == 0 or device_id == 2: + self.assertEqual(cur_subgroup, subgroups[0]) + else: + self.assertEqual(cur_subgroup, subgroups[1]) + + for subgroup in subgroups: + dist.destroy_process_group(subgroup) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @require_world_size(4) + @skip_if_lt_x_gpu(4) + def test_new_subgroups_by_enumeration_input_rank_exceeds_world_size(self): + _group, group_id, _rank = self._init_global_test() + init_multigpu_helper(dist.get_world_size(), BACKEND) + world_size = get_world_size(group_id) + + with self.assertRaisesRegex( + ValueError, + f"Rank {world_size} is out of range", + ): + dist.new_subgroups_by_enumeration( + ranks_per_subgroup_list=[[0, 1], [world_size, 2]] + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @skip_if_no_gpu + def test_new_subgroups_by_enumeration_negative_input_rank(self): + self._init_global_test() + + with self.assertRaisesRegex( + ValueError, + r"Rank -\d+ is out of range", + ): + dist.new_subgroups_by_enumeration( + ranks_per_subgroup_list=[[-1, -2], [-3, -4]] + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @require_world_size(4) + @skip_if_lt_x_gpu(4) + def test_new_subgroups_overlap_not_allowed(self): + with self.assertRaisesRegex( + ValueError, "Rank 1 has appeared in both subgroup" + ): + dist.new_subgroups_by_enumeration( + ranks_per_subgroup_list=[[0], [1, 2], [1, 3]] + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @skip_if_lt_x_gpu(2) + def test_average_parameters(self): + rank = dist.get_rank() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + + model = nn.Sequential( + nn.Conv2d(3, 3, kernel_size=3, padding=1), + nn.ReLU(), + nn.Linear(1, 5, bias=False), + ).cuda(device_id) + # Test global model averaging + for p in model.parameters(): + p.data = torch.ones_like(p.data) + model_averaging_utils.average_parameters( + params=model.parameters(), process_group=None + ) + # Every element will be the same as the input. + for p in model.parameters(): + self.assertEqual(p.data, torch.ones_like(p.data)) + + # Test partial model averaging + for p in model.parameters(): + p.data = torch.ones_like(p.data) * rank + group_nccl = dist.new_group(ranks=[0, 1], backend="nccl") + model_averaging_utils.average_parameters( + params=model.parameters(), process_group=group_nccl + ) + if not dist._rank_not_in_group(group_nccl): + # Every element on device 0 or 1 should be the average of 0 and 1, i.e., 0.5. + for p in model.parameters(): + self.assertEqual(p.data, torch.ones_like(p.data) * 0.5) + else: + # Every element on device not in the subgroup should remain the same. + for p in model.parameters(): + self.assertEqual(p.data, torch.ones_like(p.data) * rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @skip_if_lt_x_gpu(2) + def test_periodic_model_averager(self): + rank = dist.get_rank() + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + device_id = rank_to_GPU[rank][0] + + model = nn.Linear(1, 5, bias=False).cuda(device_id) + param = next(model.parameters()) + tensor = torch.ones_like(param.data) * rank + expected_avg_tensor = ( + torch.ones_like(param.data) * sum(range(world_size)) / world_size + ) + period = 4 + for warmup_steps in [12, 13, 14, 15]: + averager = averagers.PeriodicModelAverager( + period=period, warmup_steps=warmup_steps + ) + for step in range(20): + # Reset the parameters at every step. + param.data = copy.deepcopy(tensor) + for params in model.parameters(): + # mock grad + params.grad = torch.ones_like(param.data) + averager.average_parameters(model.parameters()) + if step >= warmup_steps and (step - warmup_steps) % period == 0: + self.assertEqual(param.data, expected_avg_tensor) + else: + # No model averaging, so the parameters are not updated. + self.assertEqual(param.data, tensor) + + @skip_if_lt_x_gpu(2) + def test_periodic_model_averager_param_group(self): + rank = dist.get_rank() + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + device_id = rank_to_GPU[rank][0] + + model = nn.Linear(1, 5, bias=False).cuda(device_id) + param = next(model.parameters()) + opt = torch.optim.SGD(model.parameters(), lr=0.1) + + period = 4 + for warmup_steps in [12, 13, 14, 15]: + averager = averagers.PeriodicModelAverager( + period=period, warmup_steps=warmup_steps + ) + for step in range(20): + # Reset the parameters at every step. + for param_group in opt.param_groups: + for params in param_group["params"]: + # mock grad + params.grad = torch.ones_like(param.data) * rank + params.data = torch.ones_like(param.data) * rank + averager.average_parameters(opt.param_groups) + if step >= warmup_steps and (step - warmup_steps) % period == 0: + for param_group in opt.param_groups: + for params in param_group["params"]: + if params.grad is None: + continue + self.assertEqual( + param.data, + torch.ones_like(param.data) + * sum(range(world_size)) + / world_size, + ) + else: + # No model averaging, so the parameters are not updated. + for param_group in opt.param_groups: + for params in param_group["params"]: + if params.grad is None: + continue + self.assertEqual( + param.data, torch.ones_like(param.data) * rank + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @skip_if_lt_x_gpu(2) + def test_1_level_hierarchical_model_averager_equivalent_to_periodic_model_averager( + self, + ): + rank = dist.get_rank() + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + device_id = rank_to_GPU[rank][0] + + model = nn.Linear(1, 5, bias=False).cuda(device_id) + param = next(model.parameters()) + tensor = torch.ones_like(param.data) * rank + expected_avg_tensor = ( + torch.ones_like(param.data) * sum(range(world_size)) / world_size + ) + period = 4 + for warmup_steps in [12, 13, 14, 15]: + averager = hierarchicalSGD.HierarchicalModelAverager( + # Run the global averaging at a period of 4, + # which is equivalent to the above periodic model averaging test case. + period_group_size_dict=OrderedDict([(period, world_size)]), + warmup_steps=warmup_steps, + ) + + averager = averagers.PeriodicModelAverager( + period=period, warmup_steps=warmup_steps + ) + for step in range(20): + # Reset the parameters at every step. + param.data = copy.deepcopy(tensor) + for params in model.parameters(): + # mock grad + params.grad = torch.ones_like(param.data) + averager.average_parameters(model.parameters()) + if step >= warmup_steps and (step - warmup_steps) % period == 0: + self.assertEqual(param.data, expected_avg_tensor) + else: + # No model averaging, so the parameters are not updated. + self.assertEqual(param.data, tensor) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["subgroup"], + f"The {BACKEND} backend does not support creating subgroups on CUDA devices", + ) + @require_exact_world_size(4) + @skip_if_lt_x_gpu(4) + def test_3_level_hierarchical_model_averager(self): + rank = dist.get_rank() + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + device_id = rank_to_GPU[rank][0] + + model = nn.Linear(1, 5, bias=False).cuda(device_id) + param = next(model.parameters()) + tensor = torch.ones_like(param.data) * rank + # Set up such a hierarchical model averaging as follows: + # after the first 10 warmup steps, + # run model averaging every 2 steps within each subgroup of size 2, + # run model averaging every 4 steps within each subgroup of size 3, + # and run the global model averaging every 8 steps. + # If there is a conflict in model averaging at a step, only run the highest-level model averaging. + warmup_steps = 10 + subgroup_size1 = 2 + subgroup_avg_period1 = 2 + subgroup_size2 = 4 + subgroup_avg_period2 = 4 + global_avg_period = 8 + period_group_size_dict = OrderedDict( + [ + (subgroup_avg_period1, subgroup_size1), + (subgroup_avg_period2, subgroup_size2), + (global_avg_period, world_size), + ] + ) + averager = hierarchicalSGD.HierarchicalModelAverager( + period_group_size_dict=period_group_size_dict, warmup_steps=warmup_steps + ) + self.assertEqual(dist.get_pg_count(), len(period_group_size_dict)) + + subgroup1 = averager.period_process_group_dict[subgroup_avg_period1] + subgroup2 = averager.period_process_group_dict[subgroup_avg_period2] + real_group_ranks_res1 = _get_pg_config(subgroup1)["ranks"] + real_group_ranks_res2 = _get_pg_config(subgroup2)["ranks"] + + expect_group_ranks_res1 = ( + rank // subgroup_size1 * subgroup_size1 + + np.array(list(range(subgroup_size1))) + ).tolist() + expect_group_ranks_res2 = ( + rank // subgroup_size2 * subgroup_size2 + + np.array(list(range(subgroup_size2))) + ).tolist() + self.assertEqual(real_group_ranks_res1, expect_group_ranks_res1) + self.assertEqual(real_group_ranks_res2, expect_group_ranks_res2) + + expected_avg_tensor_within_subgroup1 = ( + torch.ones_like(param.data) + * sum(real_group_ranks_res1) + / subgroup_size1 + ) + expected_avg_tensor_within_subgroup2 = ( + torch.ones_like(param.data) + * sum(real_group_ranks_res2) + / subgroup_size2 + ) + expected_global_avg_tensor = ( + torch.ones_like(param.data) * sum(range(world_size)) / world_size + ) + for step in range(25): + # Reset the parameters at every step. + param.data = copy.deepcopy(tensor) + for params in model.parameters(): + # mock grad + params.grad = torch.ones_like(param.data) + averager.average_parameters(model.parameters()) + if step == 16 or step == 24: + # Run global model averaging when `step` can be divided by 8. + self.assertEqual(param.data, expected_global_avg_tensor) + elif step == 12 or step == 20: + # Run model averaging within subgroup when `step` can be divided by 4 but not by 8. + self.assertEqual(param.data, expected_avg_tensor_within_subgroup2) + elif step == 10 or step == 14 or step == 18 or step == 22: + # Run model averaging within subgroup when `step` can be divided by 2 but not by 4 or 8. + self.assertEqual(param.data, expected_avg_tensor_within_subgroup1) + else: + # No model averaging, so the parameters are not updated. + self.assertEqual(param.data, tensor) + + # Coalescing manager (sync mode) + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl" or IS_FBCODE or IS_SANDCASTLE, + "Coalescing manager currently tests with NCCL only; internal test flaky", + ) + def test_coalescing_manager(self): + self._barrier() + rank = dist.get_rank() + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + num_colls = 2 + size_per_coll = 8 + small_tensors = [ + torch.ones(size_per_coll, device=device_id) for _ in range(num_colls) + ] + + with dist._coalescing_manager(): + for i in range(num_colls): + dist.all_reduce(small_tensors[i]) + + big_tensor = torch.ones(num_colls * size_per_coll, device=device_id) + dist.all_reduce(big_tensor) + + for i in range(num_colls): + self.assertEqual( + small_tensors[i], + big_tensor[i * size_per_coll : (i + 1) * size_per_coll], + ) + + self._barrier() + + # Coalescing manager (async mode) + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl" or IS_FBCODE or IS_SANDCASTLE, + "Coalescing manager currently tests with NCCL only; internal test flaky", + ) + def test_coalescing_manager_async(self): + self._barrier() + rank = dist.get_rank() + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + num_colls = 2 + size_per_coll = 8 + small_tensors = [ + torch.ones(size_per_coll, device=device_id) for _ in range(num_colls) + ] + + with dist._coalescing_manager(async_ops=True) as cm: + for i in range(num_colls): + dist.all_reduce(small_tensors[i]) + cm.wait() + + big_tensor = torch.ones(num_colls * size_per_coll, device=device_id) + dist.all_reduce(big_tensor) + + for i in range(num_colls): + self.assertEqual( + small_tensors[i], + big_tensor[i * size_per_coll : (i + 1) * size_per_coll], + ) + + self._barrier() + + # NCCL Batch SEND RECV + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Batch Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def test_batch_isend_irecv_nccl(self): + self._barrier() + rank = dist.get_rank() + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + p2p_op_list = [] + recv_tensors = [None for _ in range(world_size)] + expected_tensors = [None for _ in range(world_size)] + + for val in ["1", "0"]: + os.environ["TORCH_NCCL_BLOCKING_WAIT"] = val + for src in range(world_size): + send_tensor = _build_tensor(rank + 1, device_id=device_id).fill_( + src + ) + recv_tensors[src] = _build_tensor( + src + 1, value=-1, device_id=device_id + ).fill_(-1) + expected_tensors[src] = _build_tensor( + src + 1, value=-1, device_id=device_id + ).fill_(rank) + recv_op = dist.P2POp(dist.irecv, recv_tensors[src], src) + p2p_op_list.append(recv_op) + send_op = dist.P2POp(dist.isend, send_tensor, src) + p2p_op_list.append(send_op) + + reqs = dist.batch_isend_irecv(p2p_op_list) + for req in reqs: + req.wait() + + for src in range(world_size): + self.assertEqual(recv_tensors[src], expected_tensors[src]) + + self._barrier() + + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Batch Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def test_batch_isend_irecv_ring_exchange_nccl(self): + self._barrier() + rank = dist.get_rank() + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + + send_tensor = _build_tensor(world_size, device_id=device_id) + recv_tensor = _build_tensor(world_size, value=-1, device_id=device_id) + send_op = dist.P2POp(dist.isend, send_tensor, (rank + 1) % world_size) + recv_op = dist.P2POp( + dist.irecv, recv_tensor, (rank - 1 + world_size) % world_size + ) + reqs = dist.batch_isend_irecv([send_op, recv_op]) + for req in reqs: + req.wait() + + self._barrier() + + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Batch Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def test_batch_isend_irecv_self_nccl(self): + self._barrier() + # Ensure the process group has been fully initialized (needed by + # the first sub-group batch_isend_irecv call) + dist.barrier() + rank = dist.get_rank() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + p2p_op_list = [] + + if rank == 0: + send_tensor = _build_tensor(rank + 1, device_id=device_id) + recv_tensor = _build_tensor(rank + 1, value=-1, device_id=device_id) + recv_op = dist.P2POp(dist.irecv, recv_tensor, 0) + p2p_op_list.append(recv_op) + send_op = dist.P2POp(dist.isend, send_tensor, 0) + p2p_op_list.append(send_op) + + reqs = dist.batch_isend_irecv(p2p_op_list) + for req in reqs: + req.wait() + + self._barrier() + + @skip_if_no_gpu + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Batch Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def test_batch_isend_irecv_no_rank_zero_nccl(self): + self._barrier() + # Ensure the process group has been fully initialized (needed by + # the first sub-group batch_isend_irecv call) + dist.barrier() + rank = dist.get_rank() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + p2p_op_list = [] + + if rank == 1: + peer = 2 + elif rank == 2: + peer = 1 + + if rank in [1, 2]: + send_tensor = _build_tensor(rank + 1, device_id=device_id) + recv_tensor = _build_tensor(peer + 1, value=-1, device_id=device_id) + recv_op = dist.P2POp(dist.irecv, recv_tensor, peer) + p2p_op_list.append(recv_op) + send_op = dist.P2POp(dist.isend, send_tensor, peer) + p2p_op_list.append(send_op) + + reqs = dist.batch_isend_irecv(p2p_op_list) + for req in reqs: + req.wait() + + self._barrier() + + # GLOO Batch SEND RECV CPU + @skip_but_pass_in_sandcastle_if(BACKEND != "gloo", "GLOO Batch Send Recv CPU") + def test_batch_isend_irecv_gloo(self): + self._barrier() + rank = dist.get_rank() + p2p_op_list = [] + + for src in range(dist.get_world_size()): + if src == rank: + continue + send_tensor = _build_tensor(rank + 1) + recv_tensor = _build_tensor(src + 1, value=-1) + recv_op = dist.P2POp(dist.irecv, recv_tensor, src) + p2p_op_list.append(recv_op) + send_op = dist.P2POp(dist.isend, send_tensor, src) + p2p_op_list.append(send_op) + + reqs = dist.batch_isend_irecv(p2p_op_list) + for req in reqs: + req.wait() + + self._barrier() + + # GLOO Batch SEND RECV CPU with provided tags + @skip_but_pass_in_sandcastle_if(BACKEND != "gloo", "GLOO Batch Send Recv CPU") + def test_batch_isend_irecv_gloo_tags(self): + self._barrier() + rank = dist.get_rank() + p2p_op_list = [] + + for src in range(dist.get_world_size()): + if src == rank: + continue + send_tensor = _build_tensor(rank + 1) + recv_tensor = _build_tensor(src + 1, value=-1) + recv_op = dist.P2POp(dist.irecv, recv_tensor, src, tag=src) + p2p_op_list.append(recv_op) + send_op = dist.P2POp(dist.isend, send_tensor, src, tag=rank) + p2p_op_list.append(send_op) + + reqs = dist.batch_isend_irecv(p2p_op_list) + for req in reqs: + req.wait() + + self._barrier() + + # NCCL Batch SEND RECV Op Error + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Batch Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def test_batch_isend_irecv_op_err(self): + self._barrier() + rank = dist.get_rank() + if rank == 0: + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + with self.assertRaisesRegex(ValueError, "^Invalid ``op``"): + send_tensor = _build_tensor(rank + 1, device_id=device_id) + send_op = dist.P2POp(dist.broadcast, send_tensor, 1) + dist.batch_isend_irecv([send_op]) + + # NCCL Batch SEND RECV p2p_op_list Error + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Batch Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def test_batch_isend_irecv_op_list_err(self): + self._barrier() + rank = dist.get_rank() + if rank == 0: + with self.assertRaisesRegex(ValueError, "^Invalid ``p2p_op_list``"): + dist.batch_isend_irecv([1, 2]) + + # NCCL Batch SEND RECV Mixed Backend Error + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Batch Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def test_batch_isend_irecv_mixed_backend_err(self): + self._barrier() + rank = dist.get_rank() + init_multigpu_helper(dist.get_world_size(), BACKEND) + group_gloo = dist.new_group(ranks=[0, 1], backend="gloo") + group_nccl = dist.new_group(ranks=[0, 1], backend="nccl") + if rank == 0: + with self.assertRaisesRegex( + ValueError, "All ops need to use the same group" + ): + send_tensor = _build_tensor(rank + 1) + send_op_gloo = dist.P2POp(dist.isend, send_tensor, 1, group_gloo) + send_op_nccl = dist.P2POp(dist.isend, send_tensor, 1, group_nccl) + dist.batch_isend_irecv([send_op_gloo, send_op_nccl]) + + # NCCL SEND RECV + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def _test_send_recv_nccl(self, profiler_ctx=None): + # TODO: now that nccl send/recv is supported, there does not seem to + # be a need to have nccl send/recv be tested separately. + rank = dist.get_rank() + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + + tensor = _build_tensor(rank + 1, device_id=device_id) + profiler_cls = profiler_ctx if profiler_ctx is not None else nullcontext() + with profiler_cls as prof: + for src in range(world_size): + if src == rank: + # Send mode + for dst in range(world_size): + if dst == rank: + continue + dist.send(tensor, dst) + else: + # Recv mode + expected_tensor = _build_tensor(src + 1) + output_tensor = _build_tensor( + src + 1, value=-1, device_id=device_id + ) + dist.recv(output_tensor, src) + self.assertEqual(output_tensor, expected_tensor) + + self._barrier() + + if profiler_ctx is not None: + backend = dist.get_backend() + if backend in SEND_RECV_PROFILING_SUPPORTED_BACKENDS: + for event_name in [f"{backend}:send", f"{backend}:recv"]: + events = get_profiling_event( + event_name, prof, dedup_gpu_user_annotation=True + ) + self.assertTrue(events) + # Event order is not deterministic, so simply assert their shape + # is found in the following list. + expected_shapes = [ + [[rank + 1] * 3] for rank in range(dist.get_world_size()) + ] + for event in events: + self.assertTrue(event.input_shapes in expected_shapes) + + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def test_send_recv_nccl(self): + self._test_send_recv_nccl() + + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + def test_send_recv_nccl_autograd_profiler(self): + profiler_ctx = torch.autograd.profiler.profile(record_shapes=True) + self._test_send_recv_nccl(profiler_ctx) + + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if(BACKEND != "nccl", "NCCL Send Recv Only") + @requires_nccl_version((2, 7, 0), "Need NCCL 2.7+ for send/recv") + @skip_but_pass_in_sandcastle_if(IS_FBCODE, "Kineto in fbcode causes hang") + @skip_but_pass_in_sandcastle_if( + IS_MACOS or IS_WINDOWS, + "torch.profiler not enabled for mac/windows: https://github.com/pytorch/pytorch/pull/56124", + ) + def test_send_recv_nccl_torch_profiler(self): + profiler_ctx = torch.profiler.profile( + activities=[ + torch.profiler.ProfilerActivity.CPU, + torch.profiler.ProfilerActivity.CUDA, + ], + record_shapes=True, + ) + self._test_send_recv_nccl(profiler_ctx) + + # SEND RECV + def _test_send_recv(self, profiler_ctx): + rank = dist.get_rank() + send_size = rank + 1 + tensor = _build_tensor(send_size) + ctx = profiler_ctx if profiler_ctx is not None else nullcontext() + with ctx as prof: + for src in range(dist.get_world_size()): + if src == rank: + # Send mode + for dst in range(dist.get_world_size()): + if dst == rank: + continue + dist.send(tensor, dst) + else: + # Recv mode + recv_size = src + 1 + expected_tensor = _build_tensor(recv_size) + output_tensor = _build_tensor(recv_size, value=-1) + dist.recv(output_tensor, src) + self.assertEqual(output_tensor, expected_tensor) + + if profiler_ctx is not None: + backend = dist.get_backend() + if backend in SEND_RECV_PROFILING_SUPPORTED_BACKENDS: + for event_name in [f"{backend}:send", f"{backend}:recv"]: + events = get_profiling_event(event_name, prof) + # Each rank sends/recvs from all other ranks. + event_count = sum(e.count for e in events) + expected_event_count = dist.get_world_size() - 1 + self.assertEqual(event_count, expected_event_count) + # Event order is not deterministic, so simply assert their shape + # is found in the following list. + expected_shapes = [ + [[rank + 1] * 3] for rank in range(dist.get_world_size()) + ] + for event in events: + self.assertTrue(event.is_async) + self.assertTrue(event.input_shapes in expected_shapes) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl send/recv tested by test_send_recv_nccl" + ) + def test_send_recv(self): + self._test_send_recv(profiler_ctx=None) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "NCCL send/recv tested by test_send_recv_nccl" + ) + def test_send_recv_autograd_profiler(self): + autograd_profiler_ctx = _create_autograd_profiler() + self._test_send_recv(profiler_ctx=autograd_profiler_ctx) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "NCCL send/recv tested by test_send_recv_nccl" + ) + @skip_but_pass_in_sandcastle_if(IS_FBCODE, "Kineto in fbcode causes hang") + @skip_but_pass_in_sandcastle_if( + IS_MACOS or IS_WINDOWS, + "torch.profiler not enabled for mac/windows: https://github.com/pytorch/pytorch/pull/56124", + ) + def test_send_recv_torch_profiler(self): + torch_profiler_ctx = _create_torch_profiler() + return self._test_send_recv(profiler_ctx=torch_profiler_ctx) + + # SEND RECV ANY SOURCE + def _test_send_recv_any_source(self, profiler_ctx): + rank = dist.get_rank() + send_recv_size = 10 + tensor = _build_tensor(send_recv_size, value=rank) + recv_ranks = [] + irecv_ranks = [] + + ctx = profiler_ctx if profiler_ctx is not None else nullcontext() + with ctx as prof: + for dst in range(dist.get_world_size()): + if dst == rank: + # Recv mode + for dst in range(dist.get_world_size()): + if dst == rank: + continue + + for recv in ["recv", "irecv"]: + output_tensor = _build_tensor(send_recv_size, value=-1) + + if recv == "recv": + sender = dist.recv(output_tensor) + recv_ranks.append(sender) + elif recv == "irecv": + work = dist.irecv(output_tensor) + work.wait() + sender = work._source_rank() + irecv_ranks.append(sender) + + # Assert the scalar value "sender" that should be + # equal to the rank of the sender is equal to all + # values in the received tensor. + self.assertTrue(output_tensor.eq(sender).all()) + else: + # Send mode + dist.send(tensor, dst) # recv + dist.send(tensor, dst) # irecv + + if profiler_ctx is not None: + backend = dist.get_backend() + if backend in SEND_RECV_PROFILING_SUPPORTED_BACKENDS: + for event_name in [f"{backend}:send", f"{backend}:recvAnySource"]: + events = get_profiling_event(event_name, prof) + # Each rank sends/recvs from other rank twice. + self.assertEqual( + sum(event.count for event in events), + 2 * (dist.get_world_size() - 1), + ) + for event in events: + self.assertTrue(event.is_async) + self.assertEqual(event.input_shapes, [[send_recv_size] * 3]) + + # Each rank would have 2 * (world_size - 1) sends, verify that + # globally we receive the same amount on the other end. + recv_ranks_tensor = torch.cat( + (torch.tensor(recv_ranks), torch.tensor(irecv_ranks)), 0 + ) + global_recv_ranks = [ + torch.empty_like(recv_ranks_tensor) + for _ in range(dist.get_world_size()) + ] + dist.all_gather(global_recv_ranks, recv_ranks_tensor) + global_recv_ranks_list = [] + for tensor in global_recv_ranks: + global_recv_ranks_list += tensor.tolist() + + from itertools import groupby + + global_recv_ranks_list.sort() + frequency = [ + len(list(group)) for key, group in groupby(global_recv_ranks_list) + ] + self.assertEqual(dist.get_world_size(), len(frequency)) + self.assertEqual( + [2 * (dist.get_world_size() - 1)] * dist.get_world_size(), frequency + ) + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["sendrecv anysource"], + f"{BACKEND} does not support send/recv from any source", + ) + def test_send_recv_any_source(self): + self._test_send_recv_any_source(profiler_ctx=None) + + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["sendrecv anysource"], + f"{BACKEND} does not support send/recv from any source", + ) + def test_send_recv_any_source_autograd_profiler(self): + autograd_profiler_ctx = _create_autograd_profiler() + self._test_send_recv_any_source(profiler_ctx=autograd_profiler_ctx) + + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["sendrecv anysource"], + f"{BACKEND} does not support send/recv from any source", + ) + @skip_but_pass_in_sandcastle_if(IS_FBCODE, "Kineto in fbcode code causes hang") + @skip_but_pass_in_sandcastle_if( + IS_MACOS or IS_WINDOWS, + "torch.profiler not enabled for mac/windows: https://github.com/pytorch/pytorch/pull/56124", + ) + def test_send_recv_any_source_torch_profiler(self): + torch_profiler_ctx = _create_torch_profiler() + return self._test_send_recv_any_source(profiler_ctx=torch_profiler_ctx) + + # SEND RECV WITH TAG + def _test_send_recv_with_tag(self, profiler_ctx): + rank = dist.get_rank() + world_size = dist.get_world_size() + send_recv_size = 10 + tensor = _build_tensor(send_recv_size, value=rank) + ctx = profiler_ctx if profiler_ctx is not None else nullcontext() + with ctx as prof: + for dst in range(world_size): + if dst == rank: + # Recv mode + for src in range(world_size): + if src == rank: + continue + output_tensor = _build_tensor(send_recv_size, value=-1) + dist.recv(output_tensor, src, tag=src) + self.assertTrue(output_tensor.eq(src).all()) + else: + # Send mode + dist.send(tensor, dst, tag=rank) + + if profiler_ctx is not None: + backend = dist.get_backend() + if backend in SEND_RECV_PROFILING_SUPPORTED_BACKENDS: + for event_name in [f"{backend}:send", f"{backend}:recv"]: + events = get_profiling_event(event_name, prof) + # Each rank sends/recvs from all other ranks + event_count = sum(e.count for e in events) + expected_event_count = dist.get_world_size() - 1 + self.assertEqual(event_count, expected_event_count) + for event in events: + self.assertTrue(event.is_async) + self.assertEqual(event.name, event_name) + self.assertEqual(event.input_shapes, [[send_recv_size] * 3]) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "NCCL send/recv tested by test_send_recv_nccl" + ) + def test_send_recv_with_tag(self): + self._test_send_recv_with_tag(profiler_ctx=None) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "NCCL send/recv tested by test_send_recv_nccl" + ) + def test_send_recv_with_tag_autograd_profiler(self): + autograd_profiler_ctx = _create_autograd_profiler() + return self._test_send_recv_with_tag(profiler_ctx=autograd_profiler_ctx) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "NCCL send/recv tested by test_send_recv_nccl" + ) + @skip_but_pass_in_sandcastle_if(IS_FBCODE, "Kineto in fbcode code causes hang") + @skip_but_pass_in_sandcastle_if( + IS_MACOS or IS_WINDOWS, + "torch.profiler not enabled for mac/windows: https://github.com/pytorch/pytorch/pull/56124", + ) + def test_send_recv_with_tag_torch_profiler(self): + torch_profiler_ctx = _create_torch_profiler() + return self._test_send_recv_with_tag(profiler_ctx=torch_profiler_ctx) + + # ISEND + def _test_isend(self, profiler_ctx): + rank = dist.get_rank() + world_size = dist.get_world_size() + ctx = profiler_ctx if profiler_ctx is not None else nullcontext() + with ctx as prof: + if rank == 0: + requests = [ + dist.isend(_build_tensor(dest, 10), dest) + for dest in range(1, world_size) + ] + for request in requests: + request.wait() + self.assertTrue(request.is_completed()) + else: + tensor = _build_tensor(rank, -1) + dist.recv(tensor, 0) + self.assertEqual(tensor, _build_tensor(rank, 10)) + + self._barrier() + + if profiler_ctx is not None: + backend = dist.get_backend() + if backend in SEND_RECV_PROFILING_SUPPORTED_BACKENDS: + expected_event_name = ( + f"{backend}:send" if rank == 0 else f"{backend}:recv" + ) + events = get_profiling_event(expected_event_name, prof) + event_count = sum(e.count for e in events) + expected_count = dist.get_world_size() - 1 if rank == 0 else 1 + self.assertEqual(expected_count, event_count) + # Event ordering is not guaranteed, so simply ensure the shapes are + # found in the following map. + expected_shapes = { + r: [[r] * 3] for r in range(1, dist.get_world_size()) + } + for event in events: + self.assertTrue(event.is_async) + self.assertEqual(event.name, expected_event_name) + if rank == 0: + self.assertTrue( + event.input_shapes in expected_shapes.values() + ) + else: + self.assertEqual(event.input_shapes, expected_shapes[rank]) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support isend" + ) + def test_isend(self): + self._test_isend(profiler_ctx=None) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support isend" + ) + def test_isend_autograd_profiler(self): + autograd_profiler_ctx = _create_autograd_profiler() + self._test_isend(profiler_ctx=autograd_profiler_ctx) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support isend" + ) + @skip_but_pass_in_sandcastle_if(IS_FBCODE, "Kineto in fbcode code causes hang") + @skip_but_pass_in_sandcastle_if( + IS_MACOS or IS_WINDOWS, + "torch.profiler not enabled for mac/windows: https://github.com/pytorch/pytorch/pull/56124", + ) + def test_isend_torch_profiler(self): + torch_profiler_ctx = _create_torch_profiler() + self._test_isend(profiler_ctx=torch_profiler_ctx) + + # IRECV + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support irecv" + ) + def test_irecv(self): + rank = dist.get_rank() + world_size = dist.get_world_size() + + if rank == 0: + expected_tensors = [ + _build_tensor(src, -1) for src in range(1, world_size) + ] + requests = [ + dist.irecv(expected_tensors[src - 1], src) + for src in range(1, world_size) + ] + + for src in range(1, world_size): + requests[src - 1].wait() + self.assertTrue(requests[src - 1].is_completed()) + self.assertEqual(expected_tensors[src - 1], _build_tensor(src, 10)) + else: + tensor = _build_tensor(rank, 10) + dist.send(tensor, 0) + + self._barrier() + + # BROADCAST + def _test_broadcast_helper( + self, + group, + group_id, + rank, + cuda=False, + rank_to_GPU=None, + with_options=False, + ): + for dtype, value, requires_cuda in [ + (torch.float, -1e-10, False), + (torch.double, -1e-100, False), + (torch.half, -0.1, True), + (torch.int8, -2, False), + (torch.uint8, 129, False), + (torch.int, -1e5, False), + (torch.long, -1e15, False), + ]: + if requires_cuda and not cuda: + continue + for src in group: + expected_tensor = _build_tensor(src + 1, value, dtype) + if cuda: + expected_tensor = expected_tensor.cuda(rank_to_GPU[rank][0]) + if rank == src: + if with_options: + opts = dist.BroadcastOptions() + opts.rootTensor = 0 + opts.rootRank = src + self.call_dist_op( + ":broadcast", + True, + group_id.broadcast, + [expected_tensor], + opts, + ) + else: + self.call_dist_op( + ":broadcast", + False, + dist.broadcast, + expected_tensor, + src, + group_id, + ) + else: + tensor = _build_tensor(src + 1, -1, dtype) + if cuda: + tensor = tensor.cuda(rank_to_GPU[rank][0]) + if with_options: + opts = dist.BroadcastOptions() + opts.rootTensor = 0 + opts.rootRank = src + self.call_dist_op( + ":broadcast", True, group_id.broadcast, [tensor], opts + ) + else: + self.call_dist_op( + ":broadcast", + False, + dist.broadcast, + tensor, + src, + group_id, + ) + self.assertEqual(tensor.size(), expected_tensor.size()) + self.assertEqual( + tensor.ne(expected_tensor).max(), torch.tensor(False) + ) + + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_broadcast(self): + group, group_id, rank = self._init_global_test() + self._test_broadcast_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo" and BACKEND != "nccl", + "Only Gloo and Nccl backend supports CUDA allReduce", + ) + @skip_if_no_gpu + def test_broadcast_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + self._test_broadcast_helper(group, group_id, rank, True, rank_to_GPU) + + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_broadcast_group(self): + group, group_id, rank = self._init_group_test() + self._test_broadcast_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_broadcast_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_broadcast_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", + "Only NCCL backend supports high priority stream", + ) + @skip_if_no_gpu + def test_nccl_high_priority_stream(self): + group, _, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + + new_port = str(MASTER_PORT + 1) + os.environ["MASTER_PORT"] = new_port + gen_iterator = dist.rendezvous("env://", rank, dist.get_world_size()) + store, rank, size = next(gen_iterator) + store = dist.PrefixStore(new_port, store) + + opts = dist.ProcessGroupNCCL.Options() + opts.is_high_priority_stream = False + group_id = dist.ProcessGroupNCCL(store, rank, size, opts) + + self._test_broadcast_helper(group, group_id, rank, True, rank_to_GPU, True) + + # REDUCE + def _test_reduce_helper( + self, + group, + group_id, + rank, + op, + master_value, + worker_value, + expected_value, + cuda=False, + rank_to_GPU=None, + ): + for src in group: + tensor = _build_tensor(src + 1).fill_( + master_value if rank == src else worker_value + ) + if cuda: + tensor = tensor.cuda(rank_to_GPU[rank][0]) + self.call_dist_op( + ":reduce", + False, + dist.reduce, + tensor, + src, + op, + group_id, + tensor_shapes=[tensor.shape], + ) + if rank == src: + self.assertEqual(tensor, _build_tensor(src + 1, expected_value)) + + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + def test_reduce_sum(self): + group, group_id, rank = self._init_global_test() + self._test_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA reduce" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + @skip_if_no_gpu + def test_reduce_sum_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + self._test_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + 10 * (len(group) - 1), + True, + rank_to_GPU, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + def test_reduce_product(self): + group, group_id, rank = self._init_global_test() + self._test_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.PRODUCT, + 2, + 10, + reduce(operator.mul, [10] * (len(group) - 1), 2), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + def test_reduce_min(self): + group, group_id, rank = self._init_global_test() + self._test_reduce_helper( + group, group_id, rank, dist.ReduceOp.MIN, 1010, 1, 1 + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + def test_reduce_max(self): + group, group_id, rank = self._init_global_test() + self._test_reduce_helper( + group, group_id, rank, dist.ReduceOp.MAX, -1, 10, 10 + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + @skip_if_small_worldsize + def test_reduce_group_sum(self): + group, group_id, rank = self._init_group_test() + self._test_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + @skip_if_small_worldsize + def test_reduce_group_product(self): + group, group_id, rank = self._init_group_test() + self._test_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.PRODUCT, + 2, + 10, + reduce(operator.mul, [10] * (len(group) - 1), 2), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + @skip_if_small_worldsize + def test_reduce_group_min(self): + group, group_id, rank = self._init_group_test() + self._test_reduce_helper( + group, group_id, rank, dist.ReduceOp.MIN, 1010, 1, 1 + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + @skip_if_small_worldsize + def test_reduce_group_max(self): + group, group_id, rank = self._init_group_test() + self._test_reduce_helper( + group, group_id, rank, dist.ReduceOp.MAX, -1, 10, 10 + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + def test_reduce_full_group_sum(self): + group, group_id, rank = self._init_full_group_test() + self._test_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + def test_reduce_full_group_product(self): + group, group_id, rank = self._init_full_group_test() + self._test_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.PRODUCT, + 2, + 10, + reduce(operator.mul, [10] * (len(group) - 1), 2), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + def test_reduce_full_group_min(self): + group, group_id, rank = self._init_full_group_test() + self._test_reduce_helper( + group, group_id, rank, dist.ReduceOp.MIN, 1010, 1, 1 + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + def test_reduce_full_group_max(self): + group, group_id, rank = self._init_full_group_test() + self._test_reduce_helper( + group, group_id, rank, dist.ReduceOp.MAX, -1, 10, 10 + ) + + # REDUCE TWICE + def _test_reduce_twice_helper( + self, + group, + group_id, + rank, + op, + master_value, + worker_value, + expected_value, + cuda=False, + rank_to_GPU=None, + ): + for src in group: + tensors = [ + _build_tensor(src + 1).fill_( + master_value if rank == src else worker_value + ) + for i in range(2) + ] + if cuda: + for i in range(2): + tensors[i] = tensors[i].cuda(rank_to_GPU[rank][0]) + self.call_dist_op( + ":reduce", + False, + dist.reduce, + tensors[0], + src, + op, + group_id, + secondary_op_call=lambda: dist.reduce( + tensors[1], src, op, group_id + ), + tensor_shapes=[tensors[0].shape], + ) + if rank == src: + for tensor in tensors: + self.assertEqual(tensor, _build_tensor(src + 1, expected_value)) + + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + def test_reduce_sum_twice(self): + group, group_id, rank = self._init_global_test() + self._test_reduce_twice_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA reduce" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + @skip_if_no_gpu + def test_reduce_sum_cuda_twice(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + torch.cuda.set_device(device_id) + self._test_reduce_twice_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + 10 * (len(group) - 1), + True, + rank_to_GPU, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports reduce_scatter_v" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["reduce"], + f"{BACKEND} does not support reduce", + ) + @skip_if_no_gpu + def test_reduce_scatter_v_cuda(self): + self._barrier() + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + + input_split_sizes = [src + 1 for src in group] + start_len = sum(input_split_sizes[:rank]) + end_len = start_len + input_split_sizes[rank] + sum_len = sum(input_split_sizes) + master_value = 2 + worker_value = 10 + + for async_val in [True, False]: + tensor = _build_tensor(sum_len, worker_value, device_id=device_id) + tensor[start_len:end_len].fill_(master_value) + out_tensor = ( + torch.empty( + input_split_sizes[rank], sum_len, sum_len, dtype=torch.float + ) + .fill_(-1) + .cuda(device_id) + ) + + req = dist.reduce_scatter( + out_tensor, + list(torch.split(tensor, input_split_sizes)), + dist.ReduceOp.SUM, + group_id, + async_val, + ) + if async_val: + req.wait() + + expected_value = 2 + (10 * (len(group) - 1)) + expected_tensor = torch.empty( + input_split_sizes[rank], sum_len, sum_len, dtype=torch.float + ) + expected_tensor = expected_tensor.fill_(expected_value).cuda(device_id) + + self.assertEqual(out_tensor, expected_tensor) + self._barrier() + + # Test reduce_scatter_tensor accepting single tensor as input + def _reduce_scatter_tensor_helper( + self, tensor_out, tensor_in, group_id, rank, cuda=True, rank_to_GPU=None + ): + if cuda: + tensor_in = tensor_in.cuda(rank_to_GPU[rank][0]) + tensor_out = tensor_out.cuda(rank_to_GPU[rank][0]) + tensor_shapes = [tensor_out.shape] + self.call_dist_op( + ":reduce_scatter_tensor", + False, + dist.reduce_scatter_tensor, + tensor_out, + tensor_in, + dist.ReduceOp.SUM, + group_id, + False, + expect_event=False, + tensor_shapes=tensor_shapes, + ) + return tensor_out + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA reduce_scatter_tensor" + ) + @skip_if_no_gpu + def test_reduce_scatter_tensor_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + size = 2 + tensor_out = torch.zeros(size, dtype=torch.int64) + + # Concatenated input + tensor_in = torch.arange(len(group) * size) + tensor_out = self._reduce_scatter_tensor_helper( + tensor_out, tensor_in, group_id, rank, True, rank_to_GPU + ) + # Check result + expected_tensor = torch.arange(rank * size, (rank + 1) * size) * len(group) + self.assertEqual(tensor_out, expected_tensor) + self._barrier() + + # Stacked input + tensor_in = torch.reshape(tensor_in, (len(group), size)) + tensor_out = self._reduce_scatter_tensor_helper( + tensor_out, tensor_in, group_id, rank, True, rank_to_GPU + ) + # Check result + # Should be the same as the result in concatenated case + self.assertEqual(tensor_out, expected_tensor) + self._barrier() + + def call_dist_op( + self, + profiling_title_postfix, + is_async, + op, + *args, + expect_event=True, + secondary_op_call=None, + profile_cuda=False, + tensor_shapes=None, + **kwargs, + ): + op_calls = [lambda: op(*args, **kwargs)] + if secondary_op_call is not None: + op_calls.append(secondary_op_call) + + autograd_profiler_ctx = torch.autograd.profiler.profile( + use_cuda=profile_cuda, record_shapes=True + ) + + # TODO: move this test to use torch.profiler once kineto issues are + # fixed internally. + with autograd_profiler_ctx: + works = [op_call() for op_call in op_calls] + if is_async: + for work in works: + work.wait() + + if expect_event and dist.get_backend() in PROFILING_SUPPORTED_BACKENDS: + # We are only interested in the backend's implementation not the dispatcher wrapper. + events = get_profiling_event( + dist.get_backend() + profiling_title_postfix, autograd_profiler_ctx + ) + # DETAIL debug mode can use a pg wrapper that issues more collectives + # under the hood + if dist.get_debug_level() != dist.DebugLevel.DETAIL: + self.assertEqual(len(events), len(op_calls)) + for e in events: + self.assertTrue(e.is_async) + self.assertEqual(e.count, 1) + self.assertGreaterEqual(e.cpu_time, 0) + # Verify tensor shapes if given + # DETAIL debug mode can use a pg wrapper that issues more collectives + # under the hood + if ( + tensor_shapes is not None + and dist.get_debug_level() != dist.DebugLevel.DETAIL + ): + self.assertEqual( + e.input_shapes, + tensor_shapes, + f"event shape: {e.input_shapes} vs tensor {tensor_shapes}", + ) + + # ALL REDUCE + def _test_all_reduce_helper( + self, + group, + group_id, + rank, + op, + master_value, + worker_value, + expected_value, + cuda=False, + rank_to_GPU=None, + dtype=torch.float, + async_op=False, + ): + for src in group: + curr_value = master_value if rank == src else worker_value + + tensor = _build_tensor(src + 1, dtype=dtype).fill_(curr_value) + if cuda: + tensor = tensor.cuda(rank_to_GPU[rank][0]) + if tensor.dtype == torch.complex64: + tensor_shapes = [torch.view_as_real(tensor).shape] + else: + tensor_shapes = [tensor.shape] + self.call_dist_op( + ":all_reduce", + async_op, + dist.all_reduce, + tensor, + op, + group_id, + async_op=async_op, + tensor_shapes=tensor_shapes, + ) + # Currently, only Gloo backend has profiling tested with CUDA enabled. + # Only run cuda profiling test for one rank to speed up since + # running with different src_rank does not affect the correctness. + if ( + src == 0 + and cuda + and dist.get_backend() in CUDA_PROFILING_SUPPORTED_BACKENDS + ): + self.call_dist_op( + ":all_reduce", + async_op, + dist.all_reduce, + tensor, + op, + group_id, + async_op=async_op, + profile_cuda=True, + tensor_shapes=tensor_shapes, + ) + + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_sum(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_sum_async(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + async_op=True, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo" and BACKEND != "nccl", + "Only Gloo and NCCL backends will have CUDA allReduce tested", + ) + @skip_if_no_gpu + def test_all_reduce_sum_cuda(self): + torch.cuda.set_device(self.rank) + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + True, + rank_to_GPU, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo" and BACKEND != "nccl", + "Only Gloo and NCCL backends will have CUDA allReduce tested", + ) + @skip_if_no_gpu + def test_all_reduce_sum_cuda_async(self): + torch.cuda.set_device(self.rank) + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + True, + rank_to_GPU, + async_op=True, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_sum_complex(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + complex(2, 3), + complex(10, 11), + complex(2, 3) + (complex(10, 11) * (len(group) - 1)), + dtype=torch.cfloat, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_complex_unsupported_ops(self): + unsupported_ops = [ + dist.ReduceOp.MAX, + dist.ReduceOp.MIN, + dist.ReduceOp.PRODUCT, + dist.ReduceOp.BAND, + dist.ReduceOp.BOR, + dist.ReduceOp.BXOR, + ] + _group, group_id, _rank = self._init_global_test() + for unsupported_op in unsupported_ops: + with self.assertRaisesRegex(ValueError, "all_reduce does not support"): + dist.all_reduce( + _build_tensor(1, dtype=torch.cfloat), unsupported_op, group_id + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo" and BACKEND != "nccl", + "Only Gloo and NCCL backends will have CUDA allReduce tested", + ) + @skip_if_no_gpu + def test_all_reduce_sum_cuda_complex(self): + torch.cuda.set_device(self.rank) + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + complex(2, 3), + complex(10, 11), + complex(2, 3) + (complex(10, 11) * (len(group) - 1)), + True, + rank_to_GPU, + dtype=torch.cfloat, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_product(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.PRODUCT, + 2, + 10, + reduce(operator.mul, [10] * (len(group) - 1), 2), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_min(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_helper( + group, group_id, rank, dist.ReduceOp.MIN, 1010, 1, 1 + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_max(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_helper( + group, group_id, rank, dist.ReduceOp.MAX, -1, 10, 10 + ) + + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_group_sum(self): + group, group_id, rank = self._init_group_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + ) + + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_group_product(self): + group, group_id, rank = self._init_group_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.PRODUCT, + 2, + 10, + reduce(operator.mul, [10] * (len(group) - 1), 2), + ) + + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_group_min(self): + group, group_id, rank = self._init_group_test() + self._test_all_reduce_helper( + group, group_id, rank, dist.ReduceOp.MIN, 1010, 1, 1 + ) + + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_group_max(self): + group, group_id, rank = self._init_group_test() + self._test_all_reduce_helper( + group, group_id, rank, dist.ReduceOp.MAX, -1, 10, 10 + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_full_group_sum(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + 2, + 10, + 2 + (10 * (len(group) - 1)), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_full_group_product(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_helper( + group, + group_id, + rank, + dist.ReduceOp.PRODUCT, + 2, + 10, + reduce(operator.mul, [10] * (len(group) - 1), 2), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_full_group_min(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_helper( + group, group_id, rank, dist.ReduceOp.MIN, 1010, 1, 1 + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_full_group_max(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_helper( + group, group_id, rank, dist.ReduceOp.MAX, -1, 10, 10 + ) + + # SPARSE ALL REDUCE + def _test_sparse_all_reduce_sum(self, fn): + _group, group_id, rank = self._init_global_test() + + tests = simple_sparse_reduce_tests( + rank, dist.get_world_size(), num_inputs=1 + ) + for inputs, outputs in tests: + tensors = [fn(input) for input in inputs] + dist.all_reduce(tensors[0], dist.ReduceOp.SUM, group_id) + self.assertEqual(tensors[0], outputs[0]) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo", "Only Gloo backend support sparse all reduce" + ) + def test_sparse_all_reduce_sum(self): + self._test_sparse_all_reduce_sum(lambda t: t) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "gloo", "Only Gloo backend support sparse all reduce" + ) + @skip_if_no_gpu + def test_sparse_all_reduce_sum_cuda(self): + self._test_sparse_all_reduce_sum(lambda t: t.clone().cuda()) + + # ALL REDUCE - COALESCED + @staticmethod + def _all_reduce_coalesced_sum_test_cases(group_size): + return ( + [2, 3, complex(2, 3)], + [10, 11, complex(10, 11)], + [ + 2 + 10 * (group_size - 1), + 3 + 11 * (group_size - 1), + complex(2, 3) + complex(10, 11) * (group_size - 1), + ], + [torch.float, torch.float, torch.cfloat], + ) + + @staticmethod + def _all_reduce_coalesced_product_test_cases(group_size): + return ( + [1, 2], + [3, 4], + [1 * 3 ** (group_size - 1), 2 * 4 ** (group_size - 1)], + [torch.float, torch.float], + ) + + @staticmethod + def _all_reduce_coalesced_min_test_cases(group_size): + return ( + [1, 4], + [2, 3], + [1, 3], + [torch.float, torch.float], + ) + + @staticmethod + def _all_reduce_coalesced_max_test_cases(group_size): + return ( + [1, 4], + [2, 3], + [2, 4], + [torch.float, torch.float], + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_reduce_coalesced_max_complex_unsupported(self): + _group, group_id, _rank = self._init_global_test() + with self.assertRaisesRegex(ValueError, "all_reduce does not support"): + dist.all_reduce_coalesced( + [_build_tensor(1, dtype=torch.cfloat)], dist.ReduceOp.MAX, group_id + ) + + def _test_all_reduce_coalesced_helper( + self, + group, + group_id, + rank, + op, + cuda=False, + rank_to_GPU=None, + ): + test_case_func = { + dist.ReduceOp.SUM: self._all_reduce_coalesced_sum_test_cases, + dist.ReduceOp.PRODUCT: self._all_reduce_coalesced_product_test_cases, + dist.ReduceOp.MIN: self._all_reduce_coalesced_min_test_cases, + dist.ReduceOp.MAX: self._all_reduce_coalesced_max_test_cases, + }[op] + + master_values, worker_values, expected_values, dtypes = test_case_func( + len(group) + ) + + for src in group: + curr_values = master_values if rank == src else worker_values + tensors = [ + _build_tensor(src + 1, val, dtype=dtype) + for dtype, val in zip(dtypes, curr_values, strict=True) + ] + if cuda: + tensors = [t.cuda(rank_to_GPU[rank][0]) for t in tensors] + tensor_shapes = [] + for tensor in tensors: + if tensor.dtype == torch.complex64: + tensor_shapes.append(torch.view_as_real(tensor).shape) + else: + tensor_shapes.append(tensor.shape) + self.call_dist_op( + ":all_reduce", + False, + dist.all_reduce_coalesced, + tensors, + op, + group_id, + tensor_shapes=tensor_shapes, + ) + expected_tensors = [ + _build_tensor(src + 1, expected_value, dtype=dtype) + for dtype, expected_value in zip( + dtypes, expected_values, strict=True + ) + ] + self.assertEqual(tensors, expected_tensors) + + self._barrier() + + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_sum(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_coalesced_helper( + group, + group_id, + rank, + dist.ReduceOp.SUM, + cuda=False, + rank_to_GPU=None, + ) + + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_product(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_coalesced_helper( + group, + group_id, + rank, + dist.ReduceOp.PRODUCT, + cuda=False, + rank_to_GPU=None, + ) + + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_min(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_coalesced_helper( + group, + group_id, + rank, + dist.ReduceOp.MIN, + cuda=False, + rank_to_GPU=None, + ) + + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_max(self): + group, group_id, rank = self._init_global_test() + self._test_all_reduce_coalesced_helper( + group, group_id, rank, dist.ReduceOp.MAX, cuda=False, rank_to_GPU=None + ) + + @skip_if_small_worldsize + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_group_sum(self): + group, group_id, rank = self._init_group_test() + self._test_all_reduce_coalesced_helper( + group, group_id, rank, dist.ReduceOp.SUM, cuda=False, rank_to_GPU=None + ) + + @skip_if_small_worldsize + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_group_product(self): + group, group_id, rank = self._init_group_test() + self._test_all_reduce_coalesced_helper( + group, + group_id, + rank, + dist.ReduceOp.PRODUCT, + cuda=False, + rank_to_GPU=None, + ) + + @skip_if_small_worldsize + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_group_min(self): + group, group_id, rank = self._init_group_test() + self._test_all_reduce_coalesced_helper( + group, group_id, rank, dist.ReduceOp.MIN, cuda=False, rank_to_GPU=None + ) + + @skip_if_small_worldsize + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_group_max(self): + group, group_id, rank = self._init_group_test() + self._test_all_reduce_coalesced_helper( + group, group_id, rank, dist.ReduceOp.MAX, cuda=False, rank_to_GPU=None + ) + + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_full_group_sum(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_coalesced_helper( + group, group_id, rank, dist.ReduceOp.SUM, cuda=False, rank_to_GPU=None + ) + + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_full_group_product(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_coalesced_helper( + group, + group_id, + rank, + dist.ReduceOp.PRODUCT, + cuda=False, + rank_to_GPU=None, + ) + + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_full_group_min(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_coalesced_helper( + group, + group_id, + rank, + dist.ReduceOp.MIN, + cuda=False, + rank_to_GPU=None, + ) + + @require_backend_is_available({"gloo"}) + def test_all_reduce_coalesced_full_group_max(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_reduce_coalesced_helper( + group, group_id, rank, dist.ReduceOp.MAX, cuda=False, rank_to_GPU=None + ) + + # SCATTER + def _test_scatter_helper( + self, group, group_id, rank, cuda=False, rank_to_GPU=None, dtype=torch.float + ): + for dest in group: + tensor = _build_tensor(dest + 1, -1, dtype=dtype) + expected_tensor = _build_tensor(dest + 1, rank, dtype=dtype) + tensors = ( + [_build_tensor(dest + 1, i, dtype=dtype) for i in group] + if rank == dest + else [] + ) + if cuda: + tensor = tensor.cuda(rank_to_GPU[rank][0]) + tensors = [t.cuda(rank_to_GPU[rank][0]) for t in tensors] + if dtype == torch.complex64: + tensor_shapes = [torch.view_as_real(t).shape for t in tensors] + else: + tensor_shapes = [t.shape for t in tensors] + self.call_dist_op( + ":scatter", + False, + dist.scatter, + tensor, + src=dest, + scatter_list=tensors, + group=group_id, + expect_event=False, + tensor_shapes=tensor_shapes, + ) + self.assertEqual(tensor, expected_tensor) + + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + def test_scatter_checks(self): + group, _group_id, rank = self._init_global_test() + one = torch.ones([1]) + + # Specify scatter_list argument only on source rank. + output = one.clone() * -1 + if rank == 0: + scatter_list = [one.clone() * i for i in group] + dist.scatter(output, src=0, scatter_list=scatter_list) + else: + dist.scatter(output, src=0) + self.assertEqual(output, one * rank) + + # Don't specify src argument. + output = one.clone() * -1 + if rank == 0: + scatter_list = [one.clone() * i for i in group] + dist.scatter(output, scatter_list=scatter_list) + else: + dist.scatter(output) + self.assertEqual(output, one * rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + def test_scatter(self): + group, group_id, rank = self._init_global_test() + self._test_scatter_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA gather" + ) + @skip_if_no_gpu + def test_scatter_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_scatter_helper(group, group_id, rank, True, rank_to_GPU) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + def test_scatter_complex(self): + group, group_id, rank = self._init_global_test() + self._test_scatter_helper(group, group_id, rank, dtype=torch.cfloat) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA gather" + ) + @skip_if_no_gpu + def test_scatter_cuda_complex(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_scatter_helper( + group, group_id, rank, True, rank_to_GPU, dtype=torch.cfloat + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + @skip_if_small_worldsize + def test_scatter_group(self): + group, group_id, rank = self._init_group_test() + self._test_scatter_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + def test_scatter_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_scatter_helper(group, group_id, rank) + + # GATHER + def _test_gather_helper( + self, group, group_id, rank, cuda=False, rank_to_GPU=None + ): + for dest in group: + tensor = _build_tensor(dest + 1, rank) + tensors = ( + [_build_tensor(dest + 1, -1) for i in group] if rank == dest else [] + ) + if cuda: + tensor = tensor.cuda(rank_to_GPU[rank][0]) + tensors = [t.cuda(rank_to_GPU[rank][0]) for t in tensors] + self.call_dist_op( + ":gather", + False, + dist.gather, + tensor, + dst=dest, + gather_list=tensors, + group=group_id, + expect_event=False, + tensor_shapes=[tensors[0].shape] if len(tensors) > 0 else None, + ) + if rank == dest: + expected_tensors = [_build_tensor(dest + 1, i) for i in group] + for t1, t2 in zip(tensors, expected_tensors, strict=True): + self.assertEqual(t1, t2) + + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + def test_gather_checks(self): + group, _group_id, rank = self._init_global_test() + one = torch.ones([1]) + + # Specify gather_list argument only on destination rank. + if rank == 0: + gather_list = [one.clone() for _ in group] + dist.gather(one * rank, dst=0, gather_list=gather_list) + for i in group: + self.assertEqual(gather_list[i], one * i) + else: + dist.gather(one * rank, dst=0) + + # Don't specify dst argument. + if rank == 0: + gather_list = [one.clone() for _ in group] + dist.gather(one * rank, gather_list=gather_list) + for i in group: + self.assertEqual(gather_list[i], one * i) + else: + dist.gather(one * rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + def test_gather(self): + group, group_id, rank = self._init_global_test() + self._test_gather_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA gather" + ) + @skip_if_no_gpu + def test_gather_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_gather_helper(group, group_id, rank, True, rank_to_GPU) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + @skip_if_small_worldsize + def test_gather_group(self): + group, group_id, rank = self._init_group_test() + self._test_gather_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + def test_gather_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_gather_helper(group, group_id, rank) + + # ALL GATHER + def _test_all_gather_helper( + self, group, group_id, rank, cuda=False, rank_to_GPU=None, dtype=torch.float + ): + for dest in group: + tensor = _build_tensor(dest + 1, rank, dtype=dtype) + tensors = [_build_tensor(dest + 1, -1, dtype=dtype) for i in group] + allgather = dist.all_gather + if cuda: + tensor = tensor.cuda(rank_to_GPU[rank][0]) + tensors = [t.cuda(rank_to_GPU[rank][0]) for t in tensors] + if tensors[0].dtype == torch.complex64: + tensor_shapes = [torch.view_as_real(tensors[0]).shape] + else: + tensor_shapes = [tensors[0].shape] + self.call_dist_op( + ":all_gather", + False, + allgather, + tensors, + tensor, + group_id, + False, + tensor_shapes=tensor_shapes, + ) + + expected_tensors = [ + _build_tensor(dest + 1, i, dtype=dtype) for i in group + ] + for t1, t2 in zip(tensors, expected_tensors, strict=True): + self.assertEqual(t1, t2) + + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_gather(self): + group, group_id, rank = self._init_global_test() + self._test_all_gather_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all gather" + ) + @skip_if_no_gpu + def test_all_gather_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_gather_helper(group, group_id, rank, True, rank_to_GPU) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_gather_complex(self): + group, group_id, rank = self._init_global_test() + self._test_all_gather_helper(group, group_id, rank, dtype=torch.cfloat) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all gather" + ) + @skip_if_no_gpu + def test_all_gather_cuda_complex(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_gather_helper( + group, group_id, rank, True, rank_to_GPU, dtype=torch.cfloat + ) + + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_gather_group(self): + group, group_id, rank = self._init_group_test() + self._test_all_gather_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "Nccl does not support CPU tensors" + ) + def test_all_gather_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_gather_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports all_gather_v" + ) + @skip_if_no_gpu + def test_all_gather_v_cuda(self): + self._barrier() + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + + output_split_sizes = [dst + 1 for dst in group] + sum_len = sum(output_split_sizes) + value = 2 + + for async_val in [True, False]: + tensor = ( + torch.empty( + output_split_sizes[rank], sum_len, sum_len, dtype=torch.float + ) + .fill_(value) + .cuda(device_id) + ) + out_tensor = _build_tensor(sum_len, -1, device_id=device_id) + + req = dist.all_gather( + list(torch.split(out_tensor, output_split_sizes)), + tensor, + group_id, + async_val, + ) + if async_val: + req.wait() + + expected_value = value + expected_tensor = _build_tensor( + sum_len, expected_value, device_id=device_id + ) + + self.assertEqual(out_tensor, expected_tensor) + self._barrier() + + # Test all_gather accepting single tensor as output + def _all_gather_into_tensor_helper( + self, tensor_out, tensor_in, group_id, rank, cuda=True, rank_to_GPU=None + ): + if cuda: + tensor_in = tensor_in.cuda(rank_to_GPU[rank][0]) + tensor_out = tensor_out.cuda(rank_to_GPU[rank][0]) + if tensor_out.dtype == torch.complex64: + tensor_shapes = [torch.view_as_real(tensor_in).shape] + else: + tensor_shapes = [tensor_in.shape] + self.call_dist_op( + ":all_gather_into_tensor", + False, + dist.all_gather_into_tensor, + tensor_out, + tensor_in, + group_id, + False, + expect_event=False, + tensor_shapes=tensor_shapes, + ) + return tensor_out + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_gather_into_tensor" + ) + @skip_if_no_gpu + def test_all_gather_into_cat_tensor_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + size = 2 + tensor_in = torch.ones([size, size]) * rank + # Concatenated output + tensor_out = torch.ones([len(group) * size, size]) * (-1) + tensor_out = self._all_gather_into_tensor_helper( + tensor_out, tensor_in, group_id, rank, True, rank_to_GPU + ) + + # Check result + # Concatenate all blocks into a bigger tensor + expected_tensor = torch.cat([torch.ones([size, size]) * i for i in group]) + self.assertEqual(tensor_out, expected_tensor) + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_gather_into_tensor" + ) + @skip_if_no_gpu + def test_all_gather_into_stack_tensor_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + size = 2 + tensor_in = torch.ones([size, size]) * rank + # Stacked output + tensor_out = torch.ones([len(group), size, size]) * (-1) + tensor_out = self._all_gather_into_tensor_helper( + tensor_out, tensor_in, group_id, rank, True, rank_to_GPU + ) + + # Check result + # Stack all blocks into a bigger tensor + expected_tensor = torch.stack([torch.ones([size, size]) * i for i in group]) + self.assertEqual(tensor_out, expected_tensor) + self._barrier() + + def _run_all_gather_coalesced_and_verify( + self, output_tensor_lists, input_tensors, expected_tensors, group_id + ): + """ + Helper that runs all_gather_coalesced and returns true if output + matches expectations. + """ + tensor_shapes = [] + for input_tensor in input_tensors: + if input_tensor.dtype == torch.complex64: + tensor_shapes.append(torch.view_as_real(input_tensor).shape) + else: + tensor_shapes.append(input_tensor.shape) + self.call_dist_op( + ":all_gather", + False, + dist.all_gather_coalesced, + output_tensor_lists, + input_tensors, + group_id, + tensor_shapes=tensor_shapes, + ) + + for l1, l2 in zip(output_tensor_lists, expected_tensors, strict=True): + for t1, t2 in zip(l1, l2, strict=True): + if not torch.equal(t1, t2): + return False + return True + + def _test_all_gather_coalesced_helper( + self, group, group_id, rank, dtype=torch.float + ): + # TODO: Instead we should probably go through _rank_not_in_group + # mechanism to disable sending tensors + if group_id is not None: + for test_case_id in range(2, 5): + # Make sure we create tensors of incompatible sizes, e.g. + # [1], [2x2], [3x3x3] ... to be sent in one batch + input_tensors = [ + _build_multidim_tensor( + tensor_id, tensor_id, rank + tensor_id, dtype=dtype + ) + for tensor_id in range(1, test_case_id) + ] + output_tensor_lists = [ + [ + _build_multidim_tensor( + tensor_id, tensor_id, -1, dtype=dtype + ) + for tensor_id in range(1, test_case_id) + ] + for _ in group + ] + expected_tensors = [ + [ + _build_multidim_tensor( + tensor_id, tensor_id, rank_iter + tensor_id, dtype=dtype + ) + for tensor_id in range(1, test_case_id) + ] + for rank_iter in group + ] + if not self._run_all_gather_coalesced_and_verify( + output_tensor_lists, input_tensors, expected_tensors, group_id + ): + raise AssertionError( + "output tensors do not match expected outputs" + ) + + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["allgather_coalesced"], + f"{BACKEND} does not support all_gather_coalesced", + ) + def test_all_gather_coalesced_simple(self): + group, group_id, rank = self._init_global_test() + self._test_all_gather_coalesced_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["allgather_coalesced"], + f"{BACKEND} does not support all_gather_coalesced", + ) + def test_all_gather_coalesced_complex(self): + group, group_id, rank = self._init_global_test() + self._test_all_gather_coalesced_helper( + group, group_id, rank, dtype=torch.cfloat + ) + + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["allgather_coalesced"], + f"{BACKEND} does not support all_gather_coalesced", + ) + def test_all_gather_coalesced_group(self): + group, group_id, rank = self._init_group_test() + self._test_all_gather_coalesced_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["allgather_coalesced"], + f"{BACKEND} does not support all_gather_coalesced", + ) + def test_all_gather_coalesced_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_gather_coalesced_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["allgather_coalesced"], + f"{BACKEND} does not support all_gather_coalesced", + ) + def test_all_gather_coalesced_with_empty(self): + group, group_id, rank = self._init_global_test() + input_tensors = [ + rank * torch.ones([2, 2]), + torch.ones([0]), + (rank + 1) * torch.ones([3, 3]), + torch.ones([0]), + torch.ones([0]), + ] + output_tensors_lists = [ + [ + -1 * torch.ones([2, 2]), + -1 * torch.ones([0]), + -1 * torch.ones([3, 3]), + -1 * torch.ones([0]), + -1 * torch.ones([0]), + ] + for _ in group + ] + expected_tensors = [ + [ + r * torch.ones([2, 2]), + torch.ones([0]), + (r + 1) * torch.ones([3, 3]), + torch.ones([0]), + torch.ones([0]), + ] + for r in group + ] + if not self._run_all_gather_coalesced_and_verify( + output_tensors_lists, input_tensors, expected_tensors, group_id + ): + raise AssertionError("output tensors do not match expected outputs") + self._barrier() + + # AllToAll + def _test_all_to_all_single_equal_split_helper( + self, group, group_id, rank, cuda=False, rank_to_GPU=None, dtype=torch.float + ): + if group_id is not None: + size = len(group) + in_tensor = torch.ones([size, size], dtype=dtype) * rank + expected_tensor = torch.cat( + [torch.ones([1, size], dtype=dtype) * i for i in group] + ) + out_tensor = torch.ones([size, size], dtype=dtype) * -1 + if cuda: + in_tensor = in_tensor.cuda(rank_to_GPU[rank][0]) + expected_tensor = expected_tensor.cuda(rank_to_GPU[rank][0]) + out_tensor = out_tensor.cuda(rank_to_GPU[rank][0]) + if dtype == torch.complex64: + tensor_shapes = [torch.view_as_real(in_tensor).shape] + else: + tensor_shapes = [in_tensor.shape] + self.call_dist_op( + ":all_to_all", + False, + dist.all_to_all_single, + out_tensor, + in_tensor, + group=group_id, + tensor_shapes=tensor_shapes, + ) + self.assertEqual(out_tensor, expected_tensor) + self._barrier() + + def _test_all_to_all_single_unequal_split_helper( + self, group, group_id, rank, cuda=False, rank_to_GPU=None, dtype=torch.float + ): + if group_id is not None: + size = len(group) + in_splits = [i + 1 for i in group] + out_splits = [rank + 1 for _ in group] + in_tensor = torch.ones([sum(in_splits), size], dtype=dtype) * rank + out_tensor = torch.ones([(rank + 1) * size, size], dtype=dtype) + expected_tensor = torch.cat( + [torch.ones([rank + 1, size], dtype=dtype) * i for i in group] + ) + if cuda: + in_tensor = in_tensor.cuda(rank_to_GPU[rank][0]) + expected_tensor = expected_tensor.cuda(rank_to_GPU[rank][0]) + out_tensor = out_tensor.cuda(rank_to_GPU[rank][0]) + dist.all_to_all_single( + out_tensor, in_tensor, out_splits, in_splits, group=group_id + ) + self.assertEqual(out_tensor, expected_tensor) + self._barrier() + + def _test_all_to_all_helper( + self, + group, + group_id, + rank, + cuda=False, + rank_to_GPU=None, + dtype=torch.float, + ): + if group_id is not None: + size = len(group) + in_splits = [i + 1 for i in group] + in_tensors = [ + torch.ones([in_splits[i], size], dtype=dtype) * rank + for i, _ in enumerate(group) + ] + out_tensors = [ + torch.ones([(rank + 1), size], dtype=dtype) for _ in group + ] + expected_tensors = [ + torch.ones([rank + 1, size], dtype=dtype) * i for i in group + ] + if cuda: + in_tensors = [t.cuda(rank_to_GPU[rank][0]) for t in in_tensors] + expected_tensors = [ + t.cuda(rank_to_GPU[rank][0]) for t in expected_tensors + ] + out_tensors = [t.cuda(rank_to_GPU[rank][0]) for t in out_tensors] + dist.all_to_all(out_tensors, in_tensors, group=group_id) + for t1, t2 in zip(out_tensors, expected_tensors, strict=True): + self.assertEqual(t1, t2) + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports CPU all_to_all_single" + ) + def test_all_to_all_single_equal_split(self): + group, group_id, rank = self._init_global_test() + self._test_all_to_all_single_equal_split_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_to_all_single" + ) + @skip_if_no_gpu + def test_all_to_all_single_equal_split_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_single_equal_split_helper( + group, + group_id, + rank, + True, + rank_to_GPU, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports CPU all_to_all_single" + ) + def test_all_to_all_single_equal_split_complex(self): + group, group_id, rank = self._init_global_test() + self._test_all_to_all_single_equal_split_helper( + group, group_id, rank, dtype=torch.cfloat + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_to_all_single" + ) + @skip_if_no_gpu + def test_all_to_all_single_equal_split_cuda_complex(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_single_equal_split_helper( + group, group_id, rank, True, rank_to_GPU, dtype=torch.cfloat + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports CPU all_to_all_single" + ) + def test_all_to_all_single_unequal_split(self): + group, group_id, rank = self._init_global_test() + self._test_all_to_all_single_unequal_split_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_to_all_single" + ) + @skip_if_no_gpu + def test_all_to_all_single_unequal_split_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_single_unequal_split_helper( + group, + group_id, + rank, + True, + rank_to_GPU, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports CPU all_to_all_single" + ) + def test_all_to_all_single_unequal_split_complex(self): + group, group_id, rank = self._init_global_test() + self._test_all_to_all_single_unequal_split_helper( + group, group_id, rank, dtype=torch.cfloat + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_to_all_single" + ) + @skip_if_no_gpu + def test_all_to_all_single_unequal_split_cuda_complex(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_single_unequal_split_helper( + group, + group_id, + rank, + True, + rank_to_GPU, + dtype=torch.cfloat, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports all_to_all" + ) + def test_all_to_all(self): + group, group_id, rank = self._init_global_test() + self._test_all_to_all_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only NCCL supports CUDA all_to_all" + ) + @skip_if_rocm_multiprocess + def test_all_to_all_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_helper(group, group_id, rank, True, rank_to_GPU) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports all_to_all" + ) + def test_all_to_all_complex(self): + group, group_id, rank = self._init_global_test() + self._test_all_to_all_helper(group, group_id, rank, dtype=torch.cfloat) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only NCCL supports CUDA all_to_all" + ) + @skip_if_rocm_multiprocess + def test_all_to_all_cuda_complex(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_helper( + group, group_id, rank, True, rank_to_GPU, dtype=torch.cfloat + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports CPU all_to_all_single" + ) + @skip_if_small_worldsize + def test_all_to_all_single_equal_split_group(self): + group, group_id, rank = self._init_group_test() + self._test_all_to_all_single_equal_split_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_to_all_single" + ) + @skip_if_no_gpu + @skip_if_small_worldsize + def test_all_to_all_single_equal_split_group_cuda(self): + group, group_id, rank = self._init_group_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_single_equal_split_helper( + group, + group_id, + rank, + True, + rank_to_GPU, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports CPU all_to_all_single" + ) + @skip_if_small_worldsize + def test_all_to_all_single_unequal_split_group(self): + group, group_id, rank = self._init_group_test() + self._test_all_to_all_single_unequal_split_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_to_all_single" + ) + @skip_if_no_gpu + @skip_if_small_worldsize + def test_all_to_all_single_unequal_split_group_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_single_unequal_split_helper( + group, + group_id, + rank, + True, + rank_to_GPU, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports all_to_all" + ) + @skip_if_small_worldsize + def test_all_to_all_group(self): + group, group_id, rank = self._init_group_test() + self._test_all_to_all_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_to_all_single" + ) + @skip_if_small_worldsize + @skip_if_rocm_multiprocess + def test_all_to_all_group_cuda(self): + group, group_id, rank = self._init_group_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_helper(group, group_id, rank, True, rank_to_GPU) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports CPU all_to_all_single" + ) + def test_all_to_all_single_equal_split_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_to_all_single_equal_split_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_to_all_single" + ) + @skip_if_no_gpu + def test_all_to_all_single_equal_split_full_group_cuda(self): + group, group_id, rank = self._init_full_group_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_single_equal_split_helper( + group, + group_id, + rank, + True, + rank_to_GPU, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports CPU all_to_all_single" + ) + def test_all_to_all_single_unequal_split_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_to_all_single_unequal_split_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only Nccl supports CUDA all_to_all_single" + ) + @skip_if_no_gpu + def test_all_to_all_single_unequal_split_full_group_cuda(self): + group, group_id, rank = self._init_full_group_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_single_unequal_split_helper( + group, + group_id, + rank, + True, + rank_to_GPU, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi", "Only MPI supports all_to_all" + ) + def test_all_to_all_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_all_to_all_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl", "Only NCCL supports CUDA all_to_all" + ) + @skip_if_rocm_multiprocess + def test_all_to_all_full_group_cuda(self): + group, group_id, rank = self._init_full_group_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_all_to_all_helper(group, group_id, rank, True, rank_to_GPU) + + # BARRIER + def _test_barrier_helper( + self, group, group_id, rank, cuda=False, rank_to_GPU=None + ): + WAIT_TIME = 0.3 # seconds + + for dest in group: + expected_time = torch.DoubleTensor(1).fill_(0.0) + if cuda: + expected_time = expected_time.cuda(rank_to_GPU[rank][0]) + if dest == rank: + expected_time.fill_(time.time() + WAIT_TIME) + dist.broadcast(expected_time, dest, group_id) + time.sleep(WAIT_TIME + 0.1) # sleep a little bit longer + dist.barrier(group_id) + else: + dist.broadcast(expected_time, dest, group_id) + dist.barrier(group_id) + self.assertGreaterAlmostEqual( + float(time.time()), + float(expected_time[0]), + msg=f"destination rank: {dest:d}, my rank: {rank:d}" + + " (if you see this failure, please report in #14554)", + ) + + # Use higher timeout for the instance where the test runs + # against a subgroup and uses a CUDA tensor for expected time. + # The CUDA initialization for the participating processes can + # take long enough for the barrier timeout to trigger on the + # process that doesn't participate in the group. + self._barrier(timeout=20) + + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if( + BACKEND == "mpi", "MPI doesn't supports GPU barrier" + ) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc" and IS_SANDCASTLE, "Skipped internally" + ) + def test_barrier_cuda(self): + group, group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_barrier_helper(group, group_id, rank, True, rank_to_GPU) + + @skip_if_small_worldsize + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if( + BACKEND == "mpi", "MPI doesn't supports GPU barrier" + ) + def test_barrier_group_cuda(self): + group, group_id, rank = self._init_group_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_barrier_helper(group, group_id, rank, True, rank_to_GPU) + + @skip_if_small_worldsize + @skip_if_no_gpu + @skip_but_pass_in_sandcastle_if( + BACKEND == "mpi", "MPI doesn't supports GPU barrier" + ) + def test_barrier_full_group_cuda(self): + group, group_id, rank = self._init_full_group_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + self._test_barrier_helper(group, group_id, rank, True, rank_to_GPU) + + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["cpu barrier"], + f"{BACKEND} does not support CPU barrier", + ) + def test_barrier(self): + group, group_id, rank = self._init_global_test() + self._test_barrier_helper(group, group_id, rank) + + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["cpu barrier"], + f"{BACKEND} does not support CPU barrier", + ) + def test_barrier_group(self): + group, group_id, rank = self._init_group_test() + self._test_barrier_helper(group, group_id, rank) + + @skip_but_pass_in_sandcastle_if( + BACKEND in DistTestCases.skip_collective["cpu barrier"], + f"{BACKEND} does not support CPU barrier", + ) + def test_barrier_full_group(self): + group, group_id, rank = self._init_full_group_test() + self._test_barrier_helper(group, group_id, rank) + + def _model_step(self, model): + for param in model.parameters(): + if param.grad is not None: + with torch.no_grad(): + param += param.grad + param.grad = None + + def _model_step_with_zero_grad(self, model): + for param in model.parameters(): + if param.grad is not None: + with torch.no_grad(): + param += param.grad + param.grad.requires_grad_(False) + param.grad.zero_() + + def _prepare_dummy_data(self, local_bs): + # global_bs for DDP should be divisible by WORLD_SIZE + world_size = int(os.environ["WORLD_SIZE"]) + global_bs = world_size * local_bs + input_cpu = torch.randn(global_bs, 2) + target = torch.randn(global_bs, 4) + loss = nn.MSELoss() + return global_bs, input_cpu, target, loss + + # END TO END TEST FOR DISTRIBUTEDDATAPARALLEL + def _test_DDP_helper( + self, model, input_var, target, loss, scale_factor=1.0, memory_format=None + ): + model.train() + output = model(input_var) + l = loss(output, target) * scale_factor + l.backward() + if memory_format is not None: + self.assertTrue(output.is_contiguous(memory_format=memory_format)) + + def _assert_equal_param(self, param_gpu, param_DDP): + self.assertEqual(len(param_gpu), len(param_DDP)) + for p_gpu, p_DDP in zip(param_gpu, param_DDP, strict=True): + self.assertEqual(p_gpu, p_DDP) + + def _test_DDP_niter( + self, + model_base, + model_DDP, + input, + target, + loss, + local_bs, + rank, + batch_size, + test_save, + offset=None, + world_size=0, + zero_grad=False, + memory_format=None, + n_iter=5, + ): + for idx in range(n_iter): + # single cpu/gpu training + self._test_DDP_helper( + model_base, input, target, loss, memory_format=memory_format + ) + + if offset is None: + offset = rank * local_bs + + # DDP training, DDP scatters subsets of input_cpu to nodes/GPUs + self._test_DDP_helper( + model_DDP, + input[offset : offset + local_bs], + target[offset : offset + local_bs], + loss, + world_size * local_bs / batch_size if world_size != 0 else 1, + memory_format=memory_format, + ) + + # Update weights and run a second iteration to shake out errors + if zero_grad: + self._model_step_with_zero_grad(model_base) + self._model_step_with_zero_grad(model_DDP) + else: + self._model_step(model_base) + self._model_step(model_DDP) + self._assert_equal_param( + list(model_base.parameters()), list(model_DDP.module.parameters()) + ) + + # Shuffle the input so that DDP input is different + input = input[torch.randperm(batch_size)] + + # save the model in the middle and reload + if test_save and idx == 2 and INIT_METHOD.startswith("file://"): + with tempfile.NamedTemporaryFile() as tmp: + if sys.platform == "win32": + torch.save(model_DDP, tmp) + tmp.seek(0) + # weights_only=False as this is legacy code that saves the model + model_DDP = torch.load(tmp, weights_only=False) + else: + torch.save(model_DDP, tmp.name) + # weights_only=False as this is legacy code that saves the model + model_DDP = torch.load(tmp.name, weights_only=False) + + with tempfile.TemporaryFile() as tmp_file: + torch.save(model_DDP, tmp_file) + tmp_file.seek(0) + # weights_only=False as this is legacy code that saves the model + saved_model = torch.load(tmp_file, weights_only=False) + for k in model_DDP.state_dict(): + self.assertEqual(model_DDP.state_dict()[k], saved_model.state_dict()[k]) + + def _test_DistributedDataParallel( + self, + gpu_subset, + rank, + output_device=None, + gradient_as_bucket_view=False, + static_graph=False, + set_static_graph_twice=False, + ): + # Run a simple end to end DDP model, use result of single node model + # as baseline + + # cpu training setup + model = Net() + + # single gpu training setup + model_gpu = copy.deepcopy(model) + model_gpu.cuda(gpu_subset[0]) + + # DDP training setup + model_DDP = copy.deepcopy(model) + model_DDP.cuda(gpu_subset[0]) + model_DDP = nn.parallel.DistributedDataParallel( + model_DDP, + device_ids=gpu_subset, + gradient_as_bucket_view=gradient_as_bucket_view, + static_graph=static_graph, + ) + + if set_static_graph_twice: + model_DDP._set_static_graph() + + # test serializable/unserializable + with tempfile.NamedTemporaryFile() as tmp: + if sys.platform == "win32": + torch.save(model_DDP, tmp) + tmp.seek(0) + # weights_only=False as this is legacy code that saves the model + model_DDP = torch.load(tmp, weights_only=False) + else: + torch.save(model_DDP, tmp.name) + # weights_only=False as this is legacy code that saves the model + model_DDP = torch.load(tmp.name, weights_only=False) + + # dummy data initialization + local_bs = len(gpu_subset) + global_bs, input_cpu, target, loss = self._prepare_dummy_data(local_bs) + + # check two model parameters over 5 iterations + self._test_DDP_niter( + model_gpu, + model_DDP, + input_cpu.cuda(gpu_subset[0]), + target.cuda(gpu_subset[0]), + loss, + local_bs, + rank, + global_bs, + True, + ) + self._barrier() + + def _test_DistributedDataParallelCPU(self, gradient_as_bucket_view=False): + # Run a simple end to end DDP-CPU model, use result of single node + # model as baseline + _group, _group_id, rank = self._init_global_test() + + # cpu training setup + model_base = Net() + + # DDP-CPU training setup + model_DDP = copy.deepcopy(model_base) + model_DDP = nn.parallel.DistributedDataParallel( + model_DDP, gradient_as_bucket_view=gradient_as_bucket_view + ) + + # dummy data initialization + local_bs = 2 + global_bs, input_cpu, target, loss = self._prepare_dummy_data(local_bs) + + # check two model parameters over 5 iterations + self._test_DDP_niter( + model_base, + model_DDP, + input_cpu, + target, + loss, + local_bs, + rank, + global_bs, + False, + zero_grad=True, + ) + self._barrier() + + return model_DDP + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "nccl does not support DDP on CPU models" + ) + def test_DistributedDataParallelCPU(self): + self._test_DistributedDataParallelCPU() + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "nccl does not support DDP on CPU models" + ) + def test_DistributedDataParallelCPU_grad_is_view(self): + self._test_DistributedDataParallelCPU(gradient_as_bucket_view=True) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_DistributedDataParallel_requires_grad(self): + # a module without gradients shouldn't be accepted + self.assertRaises( + RuntimeError, lambda: nn.parallel.DistributedDataParallel(nn.Module()) + ) + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_ddp_zero_output_features(self): + class ToyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.net1 = nn.Linear(10, 10) + self.relu = nn.ReLU() + self.net2 = nn.Linear(10, 0) + + model = ToyModel().to(self.rank) + nn.parallel.DistributedDataParallel(model, device_ids=[self.rank]) + + @skip_but_pass_in_sandcastle_if(BACKEND == "nccl", "Gloo-only test") + def test_ddp_create_graph(self): + class Model(nn.Module): + def __init__(self) -> None: + super().__init__() + self.p = nn.Parameter(torch.tensor(1.0)) + + def forward(self): + return self.p.pow(2) + + model = Model() + ddp_model = torch.nn.parallel.DistributedDataParallel(model) + for _ in range(6): + # Verify DDP doesn't throw when ran with create_graph=True. + # Although we do warn about potential issues, please see + # https://github.com/pytorch/pytorch/issues/63929 for details. + ddp_model().backward(create_graph=True) + # grad tensors should require grad. + self.assertTrue( + all(param.requires_grad for param in ddp_model.parameters()) + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_DistributedDataParallel_non_default_stream(self): + stream = torch.cuda.Stream(self.rank) + rank = self.rank + with torch.cuda.stream(stream): + net = torch.nn.parallel.DistributedDataParallel( + torch.nn.Linear(1, 1, bias=False).cuda(rank), device_ids=[rank] + ) + for i in range(1000): + # Clear gradients manually + grad = net.module.weight.grad + if grad is not None: + grad.requires_grad_(False) + grad.zero_() + # Forward + BW + batch = torch.tensor([rank]).float().cuda(rank) + loss = net(batch).sum() + loss.backward() + # For each worker, the gradient on the weight should be worker_rank. + grad = net.module.weight.grad + avg = grad.clone() + # All-reducing the gradient averages should give us the gradient + # average. If not, then one of the workers has not correctly + # written back the averaged gradient before this all-reduce call. + dist.all_reduce(avg) + world_size = int(os.environ["WORLD_SIZE"]) + avg.div_(world_size) + expected_grad = sum(i for i in range(world_size)) / world_size + self.assertEqual( + avg[0, 0], + expected_grad, + msg=f"Expected gradient of {expected_grad} but got {avg} on rank {self.rank}", + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["cuda"], + f"The {BACKEND} backend does not support DDP communication hook on CUDA devices", + ) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_ddp_comm_hook_logging(self): + hooks = [ + default.allreduce_hook, + default.fp16_compress_hook, + powerSGD.powerSGD_hook, + powerSGD.batched_powerSGD_hook, + quantization_hooks.quantization_pertensor_hook, + quantization_hooks.quantization_perchannel_hook, + ] + + cpp_builtin_hooks = [ + dist.BuiltinCommHookType.ALLREDUCE, + dist.BuiltinCommHookType.FP16_COMPRESS, + ] + + for hook in hooks: + ddp_model = torch.nn.parallel.DistributedDataParallel( + torch.nn.Linear(1, 1, bias=False).cuda(self.rank), + device_ids=[self.rank], + ) + ddp_logging_data = ddp_model._get_ddp_logging_data() + # Hook not registered yet, so should be empty + self.assertEqual(ddp_logging_data.get("comm_hook"), None) + ddp_model.register_comm_hook(None, hook) + ddp_logging_data = ddp_model._get_ddp_logging_data() + self.assertEqual(ddp_logging_data.get("comm_hook"), hook.__qualname__) + + for hook in cpp_builtin_hooks: + ddp_model = torch.nn.parallel.DistributedDataParallel( + torch.nn.Linear(1, 1, bias=False).cuda(self.rank), + device_ids=[self.rank], + ) + ddp_logging_data = ddp_model._get_ddp_logging_data() + # Hook not registered yet, so should be empty + self.assertEqual(ddp_logging_data.get("comm_hook"), None) + ddp_model._register_builtin_comm_hook(hook) + ddp_logging_data = ddp_model._get_ddp_logging_data() + self.assertEqual(ddp_logging_data.get("comm_hook"), str(hook)) + + # No hook registered + ddp_model = torch.nn.parallel.DistributedDataParallel( + torch.nn.Linear(1, 1, bias=False).cuda(self.rank), + device_ids=[self.rank], + ) + ddp_logging_data = ddp_model._get_ddp_logging_data() + # Hook not registered yet, so should be empty + self.assertEqual(ddp_logging_data.get("comm_hook"), None) + # After second forward pass, hook should still be empty string + for _ in range(2): + inp = torch.ones(1, 1, device=self.rank) + loss = ddp_model(inp).sum() + loss.backward() + + ddp_logging_data = ddp_model._get_ddp_logging_data() + # Note: DETAIL debug mode logs DDP logging data to stdout and + # thus accesses std::map, which fills in a default value for the + # type if it didn't exist. + self.assertEqual(ddp_logging_data.get("comm_hook", ""), "") + + def _test_ddp_hook_with_optimizer_parity( + self, + grad_as_bucket_view, + static_graph, + optim_cls, + optimize_subset, + *functional_optim_args, + **functional_optim_kwargs, + ): + rank = self.rank + torch.cuda.set_device(rank) + torch.manual_seed(rank) + torch.cuda.manual_seed(rank) + models_to_test = [ + (LargeNet(), torch.randn(1, 1000).cuda()), + ] + if HAS_TORCHVISION: + models_to_test.append( + (torchvision.models.resnet50(), torch.randn(1, 3, 3, 1000).cuda()) + ) + for model, inp in models_to_test: + # Enable determinism in cudnn operators + with torch.backends.cudnn.flags( + enabled=True, deterministic=True, benchmark=False + ): + # Create DDP model that runs optimizer in fused fashion. + ddp_model_with_optimizer_hook = ( + torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(model).cuda(), + device_ids=[self.rank], + gradient_as_bucket_view=grad_as_bucket_view, + static_graph=static_graph, + ) + ) + + # Create DDP model with no hook that does optimizer after + # backward. + ddp_model_with_no_hook = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(model).cuda(), + device_ids=[self.rank], + gradient_as_bucket_view=grad_as_bucket_view, + static_graph=static_graph, + ) + hook_params = ddp_model_with_optimizer_hook.parameters() + no_hook_params = ddp_model_with_no_hook.parameters() + if optimize_subset: + hook_params = list(hook_params) + no_hook_params = list(no_hook_params) + self.assertGreater(len(hook_params), 0) + hook_params = [hook_params[0]] + no_hook_params = [no_hook_params[0]] + + # Register a fused optimizer that will run optimizer in step + # with allreduce. + + if optimize_subset: + # API where optim_params is specified. + ddp_model_with_optimizer_hook._register_fused_optim( + optim_cls, + *functional_optim_args, + optim_params=hook_params, + **functional_optim_kwargs, + ) + else: + # API where optim_params is omitted + ddp_model_with_optimizer_hook._register_fused_optim( + optim_cls, + *functional_optim_args, + **functional_optim_kwargs, + ) + + optimizer_no_hook = optim_cls( + no_hook_params, + *functional_optim_args, + **functional_optim_kwargs, + ) + + # Verify parameters are equal initially. + for hook_param, allreduce_param in zip( + ddp_model_with_optimizer_hook.parameters(), + ddp_model_with_no_hook.parameters(), + strict=True, + ): + self.assertEqual(hook_param, allreduce_param) + + # Save old parameters to later verify optimizer modified them. + opt_hook_init_params = copy.deepcopy( + list(ddp_model_with_optimizer_hook.parameters()) + ) + + # Run optimizer with hook model. + for _ in range(6): + ddp_model_with_optimizer_hook.zero_grad() + out = ddp_model_with_optimizer_hook(inp) + loss = out.sum() + loss.backward() + + dist.barrier() + + # Run regular model. + for _ in range(6): + ddp_model_with_no_hook.zero_grad() + out = ddp_model_with_no_hook(inp) + loss = out.sum() + loss.backward() + optimizer_no_hook.step() + + dist.barrier() + + # Now verify parameters are equal. + for hook_param, allreduce_param in zip( + ddp_model_with_optimizer_hook.parameters(), + ddp_model_with_no_hook.parameters(), + strict=True, + ): + self.assertEqual(hook_param, allreduce_param) + + # Verify optimizer modified appropriate parameter set, + # otherwise they'd be trivially equal above. + if optimize_subset: + self.assertNotEqual( + opt_hook_init_params[0], + next(iter(ddp_model_with_optimizer_hook.parameters())), + ) + # Untouched params should be equal + self.assertEqual( + opt_hook_init_params[1:], + list(ddp_model_with_optimizer_hook.parameters())[1:], + ) + else: + self.assertNotEqual( + opt_hook_init_params, + list(ddp_model_with_optimizer_hook.parameters()), + ) + dist.barrier() + + """ + # Commenting out the following 3 tests as they cause Sandcastle jobs to fail + # Failure signature: + # AttributeError: type object 'TestDistBackendWithSpawn' has no attribute 'test_ddp_hook_with_optimizer_parity_adamw + + from torch.testing._internal.common_utils import parametrize + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl" or BACKEND == "ucc", + "Issues with async error handling, see https://github.com/pytorch/pytorch/issues/73259", + ) + @skip_if_lt_x_gpu(2) + @parametrize("grad_as_bucket_view", [True, False]) + @parametrize("static_graph", [True, False]) + @parametrize("optimize_subset", [True, False]) + def test_ddp_hook_with_optimizer_parity_adamw( + self, + grad_as_bucket_view, + static_graph, + optimize_subset, + ): + adamw_lr = 1e-2 + adamw_betas = (0.9, 0.99) + adamw_eps = 1e-6 + self._test_ddp_hook_with_optimizer_parity( + grad_as_bucket_view, + static_graph, + torch.optim.AdamW, + optimize_subset, + adamw_lr, + betas=adamw_betas, + eps=adamw_eps, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl" or BACKEND == "ucc", + "Issues with async error handling, see https://github.com/pytorch/pytorch/issues/73259", + ) + @skip_if_lt_x_gpu(2) + @parametrize("optimize_subset", [True, False]) + def test_ddp_hook_with_optimizer_parity_adam(self, optimize_subset): + adam_lr = 1e-2 + adam_betas = (0.9, 0.99) + adam_eps = 1e-6 + self._test_ddp_hook_with_optimizer_parity( + True, # grad as bucket view + False, # static graph + torch.optim.Adam, + optimize_subset, + adam_lr, + betas=adam_betas, + eps=adam_eps, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl" or BACKEND == "ucc", + "Issues with async error handling, see https://github.com/pytorch/pytorch/issues/73259", + ) + @skip_if_lt_x_gpu(2) + @parametrize("optimize_subset", [True, False]) + def test_ddp_hook_with_optimizer_parity_sgd(self, optimize_subset): + sgd_lr = 1e-2 + sgd_momentum = 0.9 + sgd_weight_decay = 0.01 + # Not testing grad_as_bucket_view and static_graph as they are + # tested in AdamW test above. + self._test_ddp_hook_with_optimizer_parity( + True, # grad as bucket view + False, # static_graph + torch.optim.SGD, + optimize_subset, + sgd_lr, + momentum=sgd_momentum, + weight_decay=sgd_weight_decay, + ) + """ + + @skip_if_lt_x_gpu(2) + def test_get_data_parallel_params(self): + torch.cuda.set_device(self.rank) + model = TwoLinLayerNet().cuda() + # Parameters to ignore are in the format {module_name}.{param_name} + params_to_ignore = ["a.weight"] + torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model( + model, params_to_ignore + ) + torch.nn.parallel.DistributedDataParallel(model, device_ids=[self.rank]) + dp_params = ( + torch.nn.parallel.DistributedDataParallel._get_data_parallel_params( + model, named_params=True + ) + ) + for name, _ in dp_params: + self.assertNotEqual(f"module.{params_to_ignore[0]}", name) + + # test named_params=False, just check if returns the expected + # no of parameters. + num_ddp_params = len(list(model.parameters())) - 1 + count = 0 + dp_params = ( + torch.nn.parallel.DistributedDataParallel._get_data_parallel_params( + model, named_params=False + ) + ) + for _ in dp_params: + count += 1 + self.assertEqual(count, num_ddp_params) + + def _test_ddp_apply_optim_in_backward( + self, + optim_cls, + optim_kwargs, + init_before, + gradient_as_bucket_view=True, + ): + # Need to seed to ensure inputs are unique across rank. Otherwise, + # allreduce won't have any effect. + torch.manual_seed(self.rank) + torch.cuda.manual_seed(self.rank) + torch.cuda.set_device(self.rank) + + # Test a simple linear as well as a ResNet model. + models_to_test = [ + nn.Sequential(nn.Linear(3, 3), nn.Linear(3, 3), nn.Linear(3, 3)).cuda(), + # run model of at least 1M parameters to hit potential race conditions in + # stream semantics + nn.Sequential( + nn.Linear(3, 1024), nn.Linear(1024, 1024), nn.Linear(1024, 3) + ).cuda(), + ] + if HAS_TORCHVISION: + models_to_test.append(torchvision.models.resnet50().cuda()) + + for j, model in enumerate(models_to_test): + model_optim_in_bwd = copy.deepcopy(model) + model = nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + gradient_as_bucket_view=gradient_as_bucket_view, + ) + optim = optim_cls(model.parameters(), **optim_kwargs) + if init_before: + _apply_optimizer_in_backward( + optimizer_class=optim_cls, + params=model_optim_in_bwd.parameters(), + optimizer_kwargs=optim_kwargs, + ) + model_optim_in_bwd = nn.parallel.DistributedDataParallel( + model_optim_in_bwd, + device_ids=[self.rank], + gradient_as_bucket_view=gradient_as_bucket_view, + ) + if not init_before: + _apply_optimizer_in_backward( + optimizer_class=optim_cls, + params=model_optim_in_bwd.parameters(), + optimizer_kwargs=optim_kwargs, + ) + + for p1, p2 in zip( + model.parameters(), model_optim_in_bwd.parameters(), strict=True + ): + self.assertEqual(p1, p2, "Parameters not initially equal!") + # Enable determinism in cudnn operators + with torch.backends.cudnn.flags( + enabled=True, deterministic=True, benchmark=False + ): + for i in range(8): + inp = ( + torch.randn(1, 3, 1000, 1000, device="cuda") + if j == 2 + else torch.randn(10, 3, device="cuda") + ) + model(inp).sum().backward() + optim.step() + model_optim_in_bwd( + inp + ).sum().backward() # runs optimizer as well + for p1, p2 in zip( + model.parameters(), + model_optim_in_bwd.parameters(), + strict=True, + ): + self.assertEqual( + p1, p2, f"Params not equal at iteration {i}" + ) + self.assertTrue( + p2.grad is None, + f"Optim in backward grad is not None at {i}", + ) + + # set_to_none for regular optimizer to match in backward + # case. + optim.zero_grad(set_to_none=True) + + @skip_if_lt_x_gpu(2) + def test_ddp_apply_optim_in_backward(self): + for optim_cls, init_before in itertools.product( + [torch.optim.SGD, torch.optim.Adam], [True, False] + ): + with self.subTest(optim_cls=optim_cls): + self._test_ddp_apply_optim_in_backward( + optim_cls=optim_cls, + optim_kwargs={"lr": 0.03}, + init_before=init_before, + ) + + @skip_if_lt_x_gpu(2) + def test_ddp_apply_optim_in_backward_grad_as_bucket_view_false(self): + for init_before in [True, False]: + self._test_ddp_apply_optim_in_backward( + optim_cls=torch.optim.SGD, + optim_kwargs={"lr": 0.03}, + init_before=init_before, + gradient_as_bucket_view=False, + ) + + @skip_if_lt_x_gpu(2) + def test_ddp_apply_optim_in_backward_ignored_params(self): + torch.cuda.set_device(self.rank) + for init_before in [True, False]: + with self.subTest(init_before=init_before): + torch.manual_seed(self.rank) + torch.cuda.manual_seed(self.rank) + model = TwoLinLayerNet() + # Parameters to ignore are in the format {module_name}.{param_name} + params_to_ignore = ["a.weight"] + torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model( + model, params_to_ignore + ) + if init_before: + _apply_optimizer_in_backward( + optimizer_class=torch.optim.SGD, + params=model.parameters(), + optimizer_kwargs={"lr": 0.03}, + ) + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(self.rank), + device_ids=[self.rank], + ) + if not init_before: + _apply_optimizer_in_backward( + optimizer_class=torch.optim.SGD, + params=model.parameters(), + optimizer_kwargs={"lr": 0.03}, + ) + inp = torch.randn(1, 10) + a, b = net(inp) + (a.transpose(0, 1) @ b).sum().backward() + # a.weight did not go through allreduce, so optimizer acted on local + # gradient, which should be different across ranks. Remaining params + # should be equal. + models = [None for _ in range(dist.get_world_size())] + dist.all_gather_object(models, model) + rank0_model, remainder = models[0], models[1:] + for m in remainder: + self.assertNotEqual(rank0_model.a.weight, m.a.weight) + self.assertEqual( + list(rank0_model.b.parameters()), list(m.b.parameters()) + ) + self.assertEqual(rank0_model.a.bias, m.a.bias) + + def _get_fp16_config(self) -> _MixedPrecision: + return _MixedPrecision( + param_dtype=torch.float16, + reduce_dtype=torch.float16, + buffer_dtype=torch.float16, + ) + + @skip_if_lt_x_gpu(2) + def test_ddp_native_mixed_precision_ignored_params(self): + rank = self.rank + torch.manual_seed(rank) + torch.cuda.manual_seed(rank) + torch.cuda.set_device(rank) + model = TwoLinLayerNet() + model.register_buffer("buffer", torch.ones(5)) + # Parameters to ignore are in the format {module_name}.{param_name} + to_ignore = ["a.weight", "buffer"] + torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model( + model, + to_ignore, + ) + mp_config = self._get_fp16_config() + net = torch.nn.parallel.DistributedDataParallel( + model.to(rank), + device_ids=[rank], + mixed_precision=mp_config, + gradient_as_bucket_view=True, + ) + to_ignore = [f"module.{name}" for name in to_ignore] + expected_ignored = len(to_ignore) + n_ignored = 0 + # ignored params should not have _mp_param or _fp_param fields. + for n, p in itertools.chain(net.named_parameters(), net.named_buffers()): + if n in to_ignore: + n_ignored += 1 + self.assertFalse(hasattr(p, "_mp_param")) + self.assertFalse(hasattr(p, "_fp_param")) + else: + self.assertEqual(mp_config.param_dtype, p._mp_param.dtype) + self.assertEqual(torch.float32, p._fp_param.dtype) + + self.assertEqual(expected_ignored, n_ignored) + + def _test_ddp_native_mixed_precision( + self, gradient_as_bucket_view, set_grad_to_none + ): + rank = self.rank + torch.manual_seed(rank) + torch.cuda.manual_seed(rank) + torch.cuda.set_device(rank) + inp = torch.randn(10, 1) + mp_config = self._get_fp16_config() + + class MyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.m = torch.nn.Linear(1, 5) + self.register_buffer("buffer", torch.randn(1, 2)) + self.p = torch.nn.Parameter(torch.randn(10, 5), requires_grad=False) + + def forward(self_, x): # noqa: B902 + params = self_.m.parameters() + for p in params: + self.assertEqual(mp_config.param_dtype, p.dtype) + + self.assertEqual(self_.buffer.dtype, mp_config.buffer_dtype) + + self.assertEqual(mp_config.param_dtype, x.dtype) + return self_.m(x) + self_.p + + m = MyModel() + + net = torch.nn.parallel.DistributedDataParallel( + m.to(rank), + device_ids=[rank], + mixed_precision=mp_config, + gradient_as_bucket_view=gradient_as_bucket_view, + ) + # Buffers are casted in constructor. + self.assertEqual(net.module.buffer.dtype, mp_config.buffer_dtype) + # Each param should have an mp_param in the lower precision, and + # an fp_param in the higher precision. + for p in net.parameters(): + self.assertEqual(mp_config.param_dtype, p._mp_param.dtype) + self.assertEqual(torch.float32, p._fp_param.dtype) + + for _ in range(6): + loss = net(inp).sum() + loss.backward() + # Verify gradient synchronization and params and grads are fp32. + for n, param in net.named_parameters(): + self.assertEqual(param.dtype, torch.float32) + if param.grad is None: + if n != "module.p": # Only param that doesn't require grad + raise AssertionError(f"Expected n == 'module.p', got {n!r}") + else: + self.assertEqual(param.grad.dtype, torch.float32) + tensor_list = [ + torch.zeros_like(param.grad) + for _ in range(dist.get_world_size(net.process_group)) + ] + dist.all_gather(tensor_list, param.grad) + g, rest = tensor_list[0], tensor_list[1:] + self.assertEqual(g.dtype, torch.float32) + for g_ in rest: + self.assertEqual(g_.dtype, torch.float32) + self.assertEqual(g, g_) + net.zero_grad(set_to_none=set_grad_to_none) + + @skip_if_lt_x_gpu(2) + def test_ddp_native_mixed_precision_no_grad_as_bucket_view_no_set_grad_none( + self, + ): + self._test_ddp_native_mixed_precision( + gradient_as_bucket_view=False, + set_grad_to_none=False, + ) + + @skip_if_lt_x_gpu(2) + def test_ddp_native_mixed_precision_grad_as_bucket_view_no_set_grad_none(self): + self._test_ddp_native_mixed_precision( + gradient_as_bucket_view=True, + set_grad_to_none=False, + ) + + @skip_if_lt_x_gpu(2) + def test_ddp_native_mixed_precision_grad_as_bucket_view_set_grad_to_none(self): + self._test_ddp_native_mixed_precision( + gradient_as_bucket_view=True, set_grad_to_none=True + ) + + @skip_if_lt_x_gpu(2) + def test_ddp_native_mixed_precision_no_grad_as_bucket_view_set_grad_to_none( + self, + ): + self._test_ddp_native_mixed_precision( + gradient_as_bucket_view=True, set_grad_to_none=True + ) + + def _test_ddp_hook_parity(self, state, hook, num_validated_iters=100): + rank = self.rank + m = torch.nn.Linear(1, 5) + try: + process_group = state.process_group + except AttributeError: + process_group = state + + net_with_hook = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(m).to(rank), + device_ids=[rank], + process_group=process_group, + ) + net_with_hook.register_comm_hook(state=state, hook=hook) + net_without_hook = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(m).to(rank), + device_ids=[rank], + process_group=process_group, + ) + for i in range(100): + # Clear gradients manually. + for g in [ + net_without_hook.module.weight.grad, + net_with_hook.module.weight.grad, + ]: + if g is not None: + g.requires_grad_(False) + g.zero_() + # Forward + BW + batch = torch.tensor([rank]).float().cuda(rank) + loss = net_without_hook(batch).sum() + loss.backward() + # For each worker, the gradient on the weight should be worker_rank. + grad = net_without_hook.module.weight.grad + avg = grad.clone() + expected_grad = ( + sum(i for i in range(dist.get_world_size())) / dist.get_world_size() + ) + loss_hook = net_with_hook(batch).sum() + loss_hook.backward() + grad_hook = net_with_hook.module.weight.grad + avg_hook = grad_hook.clone() + + if i < num_validated_iters: + # Verify hook grad with expected. + self.assertEqual( + avg_hook[0, 0].item(), + expected_grad, + msg=f"Expected hook grad of {expected_grad} but got {avg_hook[0, 0]}", + ) + # Verify hook grad with vanilla allreduce + self.assertEqual( + avg_hook[0, 0], + avg[0, 0], + msg=f"Expected hook grad to be close to allreduce {avg[0, 0]}, but got {avg_hook[0, 0]}", + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["cuda"], + f"The {BACKEND} backend does not support DDP communication hook on CUDA devices", + ) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_ddp_hook_parity_allreduce(self): + self._test_ddp_hook_parity(state=None, hook=default.allreduce_hook) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["cuda"], + f"The {BACKEND} backend does not support DDP communication hook on CUDA devices", + ) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_ddp_hook_parity_allreduce_process_group(self): + # process_group is passed in to both DDP and comm. hook + world_size = dist.get_world_size() + rank_to_GPU = init_multigpu_helper(world_size, BACKEND) + gpus = [rank_to_GPU[int(r)][0] for r in range(world_size)] + process_group = torch.distributed.new_group(gpus) + self._test_ddp_hook_parity(state=process_group, hook=default.allreduce_hook) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["cuda"], + f"The {BACKEND} backend does not support DDP communication hook on CUDA devices", + ) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_ddp_hook_parity_powerSGD(self): + for warm_start in [True, False]: + powersgd_state = powerSGD.PowerSGDState( + process_group=None, + matrix_approximation_rank=1, + start_powerSGD_iter=2, + warm_start=warm_start, + ) + self._test_ddp_hook_parity( + state=powersgd_state, hook=powerSGD.powerSGD_hook + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["cuda"], + f"The {BACKEND} backend does not support DDP communication hook on CUDA devices", + ) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_ddp_hook_parity_post_localSGD(self): + # Although we start run local SGD at iteration 10, since we still use the global process group to run it, + # the post-LocalSGD actually still allreduces gradients globally for the remaining iterations. + state = post_localSGD.PostLocalSGDState( + process_group=None, subgroup=dist.group.WORLD, start_localSGD_iter=10 + ) + self._test_ddp_hook_parity( + state=state, hook=post_localSGD.post_localSGD_hook + ) + # Only validate the warmup iterations before local SGD is applied, + # because when `post_local_gradient_allreduce` is disabled, the gradients will not be synchronized at all. + # Note that in practice a model averager has to be applied to run model averaging, + # so local gradient averaging is not necessary. + start_localSGD_iter = 10 + state = post_localSGD.PostLocalSGDState( + process_group=None, + subgroup=dist.group.WORLD, + start_localSGD_iter=start_localSGD_iter, + post_local_gradient_allreduce=False, + ) + self._test_ddp_hook_parity( + state=state, + hook=post_localSGD.post_localSGD_hook, + num_validated_iters=start_localSGD_iter, + ) + + # When `subgroup` is None, it is equivalent to the subgroup on the each node. + # For this single-node test environment, the intra-node process group is equivalent to + # the global process group. + if self.world_size == dist.get_world_size(): + state = post_localSGD.PostLocalSGDState( + process_group=None, subgroup=None, start_localSGD_iter=10 + ) + self._test_ddp_hook_parity( + state=state, hook=post_localSGD.post_localSGD_hook + ) + + # Since we start local SGD later than the total number of 100 iterations, + # no local SGD actually is executed, and we don't even need to provide a subgroup for this case. + state = post_localSGD.PostLocalSGDState( + process_group=None, subgroup=None, start_localSGD_iter=1000 + ) + self._test_ddp_hook_parity( + state=state, hook=post_localSGD.post_localSGD_hook + ) + + def _prepare_single_device_module( + self, + rank, + process_group, + devices, + device_ids, + global_batch_size, + gradient_as_bucket_view=False, + ): + model = Net() + device = devices[0] if devices else torch.device(f"cuda:{rank:d}") + ddp_model = DistributedDataParallel( + copy.deepcopy(model).to(device), + device_ids=device_ids, + process_group=process_group, + bucket_cap_mb=0.001, + gradient_as_bucket_view=gradient_as_bucket_view, + ) + + model.to(device) + + input = torch.randn(global_batch_size, 2).to(device) + target = torch.randn(global_batch_size, 4).to(device) + + return model, ddp_model, input, target + + def _prepare_cpu_module( + self, + process_group, + global_batch_size, + gradient_as_bucket_view=False, + ): + model = Net() + ddp_model = DistributedDataParallel( + copy.deepcopy(model), + process_group=process_group, + bucket_cap_mb=0.001, + gradient_as_bucket_view=gradient_as_bucket_view, + ) + input = torch.randn(global_batch_size, 2) + target = torch.randn(global_batch_size, 4) + return model, ddp_model, input, target + + def _test_accumulate_gradients_no_sync( + self, num_iters=2, ddp_comm_hook=None, gradient_as_bucket_view=False + ): + """ + This is the recommended way to implement accumulate grads. + If ``ddp_comm_hook`` input was specified, it will also register that hook + to the ``ddp_model``. The hook fed into this function should not change + the resulting gradients. + """ + _group, group_id, rank = self._init_global_test() + world_size = get_world_size() + + # FIXME: Add testing for gloo/CUDA + if BACKEND == "mpi" or BACKEND == "gloo": + global_batch_size = world_size + local_batch_size = 1 + model, ddp_model, input, target = self._prepare_cpu_module( + group_id, global_batch_size, gradient_as_bucket_view + ) + + if BACKEND == "nccl": + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + int_devices = rank_to_GPU[rank][:1] + devices = [torch.device("cuda:" + str(i)) for i in int_devices] + global_batch_size = world_size + local_batch_size = len(devices) + model, ddp_model, input, target = self._prepare_single_device_module( + rank, + group_id, + devices, + devices, + global_batch_size, + gradient_as_bucket_view, + ) + + if ddp_comm_hook is not None: + ddp_model.register_comm_hook(group_id, ddp_comm_hook) + + def step_model(model, input, target): + model.train() + output = model(input) + loss = F.mse_loss(output, target.to(output.device)) + loss.backward() + + # ensure accumulate grads works with no_grad => no grads are accumulated. + with torch.no_grad(): + with ddp_model.no_sync(): + ddp_model.train() + ddp_model(input) + + # check two model parameters over num_iters iterations + for iteration in range(num_iters): + step_model(model, input, target) + + ddp_input = input[ + rank * local_batch_size : (rank + 1) * local_batch_size + ] + ddp_target = target[ + rank * local_batch_size : (rank + 1) * local_batch_size + ] + + if iteration % 2 == 0: + # accumulate grads locally + with ddp_model.no_sync(): + step_model(ddp_model, ddp_input, ddp_target) + else: + # sync grads + step_model(ddp_model, ddp_input, ddp_target) + + for i, j in zip( + model.parameters(), ddp_model.parameters(), strict=True + ): + if not i.requires_grad: + continue + if iteration % 2 == 0: + self.assertNotEqual(i.grad, j.grad) + else: + self.assertEqual(i.grad, j.grad) + + # Shuffle the input so that DDP input is different + torch.manual_seed(1337 + iteration) + input = input[torch.randperm(global_batch_size)] + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi" and BACKEND != "nccl" and BACKEND != "gloo", + "get_future is only supported on mpi, nccl and gloo", + ) + @nccl_skip_if_lt_x_gpu(BACKEND, 2) + def test_accumulate_gradients_no_sync(self): + """ + Runs _test_accumulate_gradients_no_sync using default inputs + """ + self._test_accumulate_gradients_no_sync() + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi" and BACKEND != "nccl" and BACKEND != "gloo", + "get_future is only supported on mpi, nccl and gloo", + ) + @nccl_skip_if_lt_x_gpu(BACKEND, 2) + def test_accumulate_gradients_no_sync_grad_is_view(self): + """ + Runs _test_accumulate_gradients_no_sync using default inputs + """ + self._test_accumulate_gradients_no_sync(gradient_as_bucket_view=True) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi" and BACKEND != "nccl" and BACKEND != "gloo", + "get_future is only supported on mpi, nccl and gloo", + ) + @nccl_skip_if_lt_x_gpu(BACKEND, 2) + def test_accumulate_gradients_no_sync_allreduce_hook(self): + """ + Runs multiple iterations on _test_accumulate_gradients_no_sync + using allreduce hook and validates whether future result was properly + passed as gradients in reducer. + """ + + world_size = get_world_size() + + def allreduce_hook( + group_id: object, bucket: dist.GradBucket + ) -> torch.futures.Future[torch.Tensor]: + tensors = [bucket.buffer() / world_size] + return ( + group_id.allreduce(tensors) + .get_future() + .then(lambda fut: fut.value()[0]) + ) + + self._test_accumulate_gradients_no_sync( + num_iters=4, ddp_comm_hook=allreduce_hook + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi" and BACKEND != "nccl" and BACKEND != "gloo", + "get_future is only supported on mpi, nccl and gloo", + ) + @nccl_skip_if_lt_x_gpu(BACKEND, 2) + def test_accumulate_gradients_no_sync_allreduce_with_then_hook(self): + """ + Runs multiple iterations on _test_accumulate_gradients_no_sync using allreduce + hook that also uses then callbacks. In first then callback result is multiplied + by 2, and the second callback divides the result by 2 * world_size. It validates + whether final result was properly passed as gradients in reducer. + """ + + world_size = get_world_size() + + def allreduce_with_then_hook( + group_id: object, bucket: dist.GradBucket + ) -> torch.futures.Future[torch.Tensor]: + fut = group_id.allreduce([bucket.buffer()]).get_future() + + def mult(fut): + # Multiply the result by 2. + return 2 * fut.wait()[0] + + def div(fut): + # Divide the result by 2 * world_size. + return fut.wait() / (2 * world_size) + + return fut.then(mult).then(div) + + self._test_accumulate_gradients_no_sync( + num_iters=4, ddp_comm_hook=allreduce_with_then_hook + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND != "mpi" and BACKEND != "nccl" and BACKEND != "gloo", + "get_future is only supported on mpi, nccl and gloo", + ) + @nccl_skip_if_lt_x_gpu(BACKEND, 2) + def test_get_future(self): + def mult(fut): + return [t * 3 for t in fut.wait()] + + def add(fut): + return [t + 1 for t in fut.wait()] + + group, group_id, rank = self._init_global_test() + input = _build_tensor(3, 2) + if BACKEND == "nccl": + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + device_id = rank_to_GPU[rank][0] + input = input.to(device_id) + fut = group_id.allreduce([input]).get_future() + res = fut.then(mult).then(add).wait() + expected = _build_tensor(3, 2 * len(group) * 3 + 1) + + self.assertEqual(res[0], expected) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_DistributedDataParallel(self): + _group, _group_id, rank = self._init_global_test() + rank_to_GPU = init_multigpu_helper(dist.get_world_size(), BACKEND) + gpus = list(rank_to_GPU[rank]) + + for use_bucket_view, static_graph in itertools.product( + (False, True), (False, True) + ): + self._test_DistributedDataParallel( + gpu_subset=gpus, + rank=rank, + gradient_as_bucket_view=use_bucket_view, + static_graph=static_graph, + ) + + # test set static graph twice + self._test_DistributedDataParallel( + gpu_subset=gpus, + rank=rank, + gradient_as_bucket_view=use_bucket_view, + static_graph=static_graph, + set_static_graph_twice=True, + ) + + # test output_device + self._test_DistributedDataParallel( + gpu_subset=gpus, + rank=rank, + output_device=torch.device("cuda"), + gradient_as_bucket_view=use_bucket_view, + static_graph=static_graph, + ) + + # test device_ids + gpus_list = [torch.device("cuda:" + str(i)) for i in gpus] + self._test_DistributedDataParallel( + gpu_subset=gpus_list, + rank=rank, + output_device=torch.device("cuda"), + gradient_as_bucket_view=use_bucket_view, + static_graph=static_graph, + ) + + def _test_DistributedDataParallel_with_amp(self, grad_is_view=False): + torch.manual_seed(31415) + # Creates model and optimizer in default precision + model = Net().cuda() + optimizer = torch.optim.SGD(model.parameters(), lr=0.03) + + # Creates a GradScaler once at the beginning of training. + scaler = GradScaler() + + ddp_model = nn.parallel.DistributedDataParallel( + model, device_ids=[self.rank], gradient_as_bucket_view=grad_is_view + ) + + input = torch.randn(dist.get_world_size() * 2, 2).cuda() + target = torch.randn(dist.get_world_size() * 2, 4).cuda() + loss_fn = nn.MSELoss() + + # verify grads are none before training + for p in ddp_model.parameters(): + self.assertTrue(p is not None) + self.assertTrue(p.grad is None) + + for idx in range(20): + optimizer.zero_grad() + # Runs the forward pass with autocasting. + with autocast(): + output = ddp_model(input) + loss = loss_fn(output, target) + + # Scales loss. Calls backward() on scaled loss to create scaled gradients. + # Backward passes under autocast are not recommended. + # Backward ops run in the same dtype autocast chose for corresponding forward ops. + scaler.scale(loss).backward() + + # verify grads are not none and are valid during training + for p in ddp_model.parameters(): + if p.requires_grad: + self.assertTrue(p.grad is not None) + self.assertFalse(p.grad.isnan().any()) + self.assertFalse(p.grad.isinf().any()) + + # scaler.step() first unscales the gradients of the optimizer's assigned params. + # If these gradients do not contain infs or NaNs, optimizer.step() is then called, + # otherwise, optimizer.step() is skipped. + scaler.step(optimizer) + + # Updates the scale for next iteration. + scaler.update() + + # Shuffle the input so that DDP input is different + torch.manual_seed(1337 + idx) + input = input[torch.randperm(dist.get_world_size() * 2)] + + return ddp_model + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_DistributedDataParallel_with_amp_and_grad_is_view(self): + torch.cuda.set_device(self.rank) + ddp_model_grad_not_view = self._test_DistributedDataParallel_with_amp( + grad_is_view=False + ) + ddp_model_grad_is_view = self._test_DistributedDataParallel_with_amp( + grad_is_view=True + ) + for i, j in zip( + ddp_model_grad_not_view.parameters(), + ddp_model_grad_is_view.parameters(), + strict=True, + ): + self.assertEqual(i, j) + + def _test_DistributedDataParallel_SyncBatchNorm( + self, + gpu_subset, + rank, + local_bs, + global_bs, + offset, + output_device=None, + affine=True, + ): + # Run a simple end to end DDP model, use result of single node model + # as baseline + + # cpu training setup + model = BatchNormNet() if affine else BatchNormNet(affine=False) + + # single gpu training setup + model_gpu = copy.deepcopy(model) + model_gpu.cuda(gpu_subset[0]) + + # DDP training setup + model_DDP = nn.SyncBatchNorm.convert_sync_batchnorm(copy.deepcopy(model)) + model_DDP.cuda(gpu_subset[0]) + model_DDP = nn.parallel.DistributedDataParallel( + model_DDP, device_ids=gpu_subset + ) + + # test serializable/unserializable + with tempfile.NamedTemporaryFile() as tmp: + if sys.platform == "win32": + torch.save(model_DDP, tmp) + tmp.seek(0) + # weights_only=False as this is legacy code that saves the model + model_DDP = torch.load(tmp, weights_only=False) + else: + torch.save(model_DDP, tmp.name) + # weights_only=False as this is legacy code that saves the model + model_DDP = torch.load(tmp.name, weights_only=False) + + # data initialization + input_cpu = torch.randn(global_bs, 2) + target = torch.randn(global_bs, 4) + loss = nn.MSELoss() + + # check two model parameters over 5 iterations + self._test_DDP_niter( + model_gpu, + model_DDP, + input_cpu.cuda(gpu_subset[0]), + target.cuda(gpu_subset[0]), + loss, + local_bs, + rank, + global_bs, + True, + offset, + dist.get_world_size(), + 5 if affine else 2, + ) + self._barrier() + + def _test_post_localSGD_optimizer_parity(self, create_averager, grad_is_view): + learning_rate = 0.03 + + DDP_NET = Net() + net = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(DDP_NET).cuda(), + device_ids=[self.rank], + gradient_as_bucket_view=grad_is_view, + ) + averager = create_averager() + opt = torch.optim.SGD(net.parameters(), lr=learning_rate) + + net_using_post_localSGD_opt = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(DDP_NET).cuda(), + device_ids=[self.rank], + gradient_as_bucket_view=grad_is_view, + ) + # Process group cannot be pickled in some environments, + # so cannot deep copy an averager. See: + # https://github.com/pytorch/pytorch/pull/74737#pullrequestreview-922487496 + averager2 = create_averager() + post_localSGD_opt = self._create_post_localSGD_optimizer( + net_using_post_localSGD_opt, learning_rate, averager2 + ) + + input = torch.randn(dist.get_world_size() * 2, 2).cuda() + target = torch.randn(dist.get_world_size() * 2, 4).cuda() + loss_fn = nn.MSELoss() + + for _ in range(20): + self._perform_a_train_step(opt, net, loss_fn, input, target) + averager.average_parameters(net.parameters()) + + self._perform_a_train_step( + post_localSGD_opt, + net_using_post_localSGD_opt, + loss_fn, + input, + target, + ) + for p1, p2 in zip( + net.parameters(), + net_using_post_localSGD_opt.parameters(), + strict=True, + ): + self.assertEqual(p1.data, p2.data) + + # Also check if the built-in step counters are the same to prevent a bug like #74737. + self.assertEqual(averager.step, averager2.step) + + def _create_periodic_model_averager(self): + return averagers.PeriodicModelAverager(period=4, warmup_steps=10) + + def _create_post_localSGD_optimizer(self, net, learning_rate, averager): + return post_localSGD_optimizer.PostLocalSGDOptimizer( + optim=torch.optim.SGD(net.parameters(), lr=learning_rate), + averager=averager, + ) + + def _perform_a_train_step(self, optimizer, net, loss_fn, input, target): + optimizer.zero_grad() + output = net(input) + loss = loss_fn(output, target) + loss.backward() + optimizer.step() + + def _test_post_localSGD_optimizer_step_reload( + self, create_averager, chkpt_file + ): + learning_rate = 0.03 + + net_using_post_localSGD_opt = torch.nn.parallel.DistributedDataParallel( + Net().cuda(), device_ids=[self.rank] + ) + + averager = create_averager() + post_localSGD_opt = self._create_post_localSGD_optimizer( + net_using_post_localSGD_opt, learning_rate, averager + ) + + averager2 = create_averager() + dummy_post_localSGD_opt = self._create_post_localSGD_optimizer( + net_using_post_localSGD_opt, learning_rate, averager2 + ) + + input = torch.randn(dist.get_world_size() * 2, 2).cuda() + target = torch.randn(dist.get_world_size() * 2, 4).cuda() + loss_fn = nn.MSELoss() + + for _ in range(20): + self._perform_a_train_step( + post_localSGD_opt, + net_using_post_localSGD_opt, + loss_fn, + input, + target, + ) + + if self.rank == 0: + torch.save( + {"optimizer_state_dict": post_localSGD_opt.state_dict()}, chkpt_file + ) + + dist.barrier() + map_location = {"cuda:0": f"cuda:{self.rank:d}"} + checkpoint = torch.load(chkpt_file, map_location=map_location) + dummy_post_localSGD_opt.load_state_dict(checkpoint["optimizer_state_dict"]) + + # Check that we didn't hit the trivial case + self.assertNotEqual(averager2.step, 0) + # Check if dummy averager was initialized to a correct value + self.assertEqual(averager.step, averager2.step) + + # Remove 'step' entry from a checkpoint. + # And make sure it is not in the state dictionary + del checkpoint["optimizer_state_dict"]["step"] + self.assertNotIn("step", checkpoint["optimizer_state_dict"]) + + # Check if checkpoint without a 'step' entry invokes a warning + with self.assertWarnsRegex( + expected_warning=UserWarning, + expected_regex="Loaded state dict does not contain a step counter for an averager. " + "Setting step counter to 0.", + ): + dummy_post_localSGD_opt.load_state_dict( + checkpoint["optimizer_state_dict"] + ) + + self.assertEqual(averager2.step, 0) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_post_localSGD_optimizer_parity(self): + torch.cuda.set_device(self.rank) + self._test_post_localSGD_optimizer_parity( + self._create_periodic_model_averager, + grad_is_view=False, + ) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_post_localSGD_optimizer_parity_grad_is_view(self): + torch.cuda.set_device(self.rank) + self._test_post_localSGD_optimizer_parity( + self._create_periodic_model_averager, + grad_is_view=True, + ) + + def _create_hierarchical_model_averager(self): + period_group_size_dict = OrderedDict([(2, 2), (4, dist.get_world_size())]) + return hierarchicalSGD.HierarchicalModelAverager( + period_group_size_dict=period_group_size_dict, warmup_steps=4 + ) + + @skip_if_lt_x_gpu(4) + @skip_if_odd_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_post_localSGD_optimizer_parity_with_hierarchical_sgd(self): + torch.cuda.set_device(self.rank) + self._test_post_localSGD_optimizer_parity( + self._create_hierarchical_model_averager, + grad_is_view=False, + ) + + @skip_if_lt_x_gpu(4) + @skip_if_odd_worldsize + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_post_localSGD_optimizer_parity_with_hierarchical_sgd_grad_is_view( + self, + ): + torch.cuda.set_device(self.rank) + self._test_post_localSGD_optimizer_parity( + self._create_hierarchical_model_averager, + grad_is_view=True, + ) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_post_localSGD_optimizer_step_reload(self): + torch.cuda.set_device(self.rank) + with _rank_temp_file() as tmp_file: + self._test_post_localSGD_optimizer_step_reload( + self._create_periodic_model_averager, tmp_file + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_DistributedDataParallel_SyncBatchNorm_Channels_Last(self): + self._test_DistributedDataParallel_SyncBatchNorm_with_memory_format( + torch.channels_last + ) + self._test_DistributedDataParallel_SyncBatchNorm_with_memory_format( + torch.channels_last_3d + ) + + def _test_DistributedDataParallel_SyncBatchNorm_with_memory_format( + self, memory_format + ): + _group, _group_id, rank = self._init_global_test() + num_processes = dist.get_world_size() + local_bs = 2 + bs_offset = int(rank * 2) + global_bs = int(num_processes * 2) + + model = nn.SyncBatchNorm(2, momentum=0.99) + model_gpu = copy.deepcopy(model).cuda(rank) + model_DDP = nn.parallel.DistributedDataParallel( + model_gpu, device_ids=[rank] + ) + + shapes = [global_bs, 2, 4, 4] + ( + [] if memory_format is torch.channels_last else [4] + ) + + input_gpu = ( + torch.randn(*shapes, dtype=torch.float) + .cuda(rank) + .to(memory_format=memory_format) + ) + target_gpu = ( + torch.randn(*shapes, dtype=torch.float) + .cuda(rank) + .to(memory_format=memory_format) + ) + loss = nn.MSELoss() + + # check two model parameters over 5 iterations + self._test_DDP_niter( + model_gpu, + model_DDP, + input_gpu, + target_gpu, + loss, + local_bs, + rank, + global_bs, + True, + bs_offset, + dist.get_world_size(), + memory_format=memory_format, + ) + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_DistributedDataParallel_SyncBatchNorm(self): + _group, _group_id, rank = self._init_global_test() + world_size = dist.get_world_size() + # DDP does not support replicating BN layers within a process, hence + # testing with one module replica per process + gpus = [rank] + + local_bs = 2 + bs_offset = int(rank * 2) + global_bs = int(world_size * 2) + + self._test_DistributedDataParallel_SyncBatchNorm( + gpu_subset=gpus, + rank=rank, + local_bs=local_bs, + global_bs=global_bs, + offset=bs_offset, + ) + + # test output_device + self._test_DistributedDataParallel_SyncBatchNorm( + gpu_subset=gpus, + rank=rank, + local_bs=local_bs, + global_bs=global_bs, + offset=bs_offset, + output_device=torch.device("cuda"), + ) + + # test device_ids + gpus = [torch.device("cuda:" + str(i)) for i in gpus] + self._test_DistributedDataParallel_SyncBatchNorm( + gpu_subset=gpus, + rank=rank, + local_bs=local_bs, + global_bs=global_bs, + offset=bs_offset, + output_device=torch.device("cuda"), + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_DistributedDataParallel_SyncBatchNorm_No_Affine(self): + _group, _group_id, rank = self._init_global_test() + world_size = dist.get_world_size() + # DDP does not support replicating BN layers within a process, hence + # testing with one module replica per process + gpus = [rank] + + local_bs = 2 + bs_offset = int(rank * 2) + global_bs = int(world_size * 2) + + self._test_DistributedDataParallel_SyncBatchNorm( + gpu_subset=gpus, + rank=rank, + local_bs=local_bs, + global_bs=global_bs, + offset=bs_offset, + affine=False, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_DistributedDataParallel_SyncBatchNorm_2D_Input(self): + _group, _group_id, rank = self._init_global_test() + # DDP does not support replicating BN layers within a process, hence + # testing with one module replica per process + gpus = [rank] + + model = nn.BatchNorm1d(2) + + # single gpu training setup + model_gpu = copy.deepcopy(model) + model_gpu.cuda(gpus[0]) + + # DDP training setup + model_DDP = nn.SyncBatchNorm.convert_sync_batchnorm(copy.deepcopy(model)) + model_DDP.cuda(gpus[0]) + model_DDP = nn.parallel.DistributedDataParallel(model_DDP, device_ids=gpus) + + local_bs = len(gpus) * 2 + global_bs = dist.get_world_size() * local_bs + input_cpu = torch.randn(global_bs, 2) + target = torch.randn(global_bs, 2) + loss = nn.MSELoss() + + # disabling cudnn. + # SyncBatchNorm goes through native_batch_norm kernel, this avoids the + # numerical issue created by the divergent code path. + with torch.backends.cudnn.flags(False): + # check two model parameters over 5 iterations + self._test_DDP_niter( + model_gpu, + model_DDP, + input_cpu.cuda(gpus[0]), + target.cuda(gpus[0]), + loss, + local_bs, + rank, + global_bs, + True, + ) + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + @require_world_size(2) + def test_DistributedDataParallel_SyncBatchNorm_Single_Input_Per_Process(self): + _group, _group_id, rank = self._init_global_test() + # DDP does not support replicating BN layers within a process, hence + # testing with one module replica per process + gpus = [rank] + + model = nn.BatchNorm1d(2) + + # single gpu training setup + model_gpu = copy.deepcopy(model) + model_gpu.cuda(gpus[0]) + + # DDP training setup + model_DDP = nn.SyncBatchNorm.convert_sync_batchnorm(copy.deepcopy(model)) + model_DDP.cuda(gpus[0]) + model_DDP = nn.parallel.DistributedDataParallel(model_DDP, device_ids=gpus) + + local_bs = 1 + global_bs = dist.get_world_size() + input_cpu = torch.randn(global_bs, 2) + target = torch.randn(global_bs, 2) + loss = nn.MSELoss() + + # disabling cudnn. + # SyncBatchNorm goes through native_batch_norm kernel, this avoids the + # numerical issue created by the divergent code path. + with torch.backends.cudnn.flags(False): + # check two model parameters over 5 iterations + self._test_DDP_niter( + model_gpu, + model_DDP, + input_cpu.cuda(gpus[0]), + target.cuda(gpus[0]), + loss, + local_bs, + rank, + global_bs, + True, + ) + self._barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_DistributedDataParallel_SyncBatchNorm_Diff_Input_Sizes_Running_Value( + self, + ): + ONLY_SBN_NET = nn.SyncBatchNorm(2, momentum=0.99) + _group, _group_id, rank = self._init_global_test() + model = nn.parallel.DistributedDataParallel( + ONLY_SBN_NET.cuda(rank), device_ids=[rank] + ) + + input_var = [] + for i in range(dist.get_world_size()): + input_var_rank = torch.cat( + [ + torch.ones(2, 1, 10 ** (i + 1)) * (0.1 ** (i - 1)), + torch.ones(2, 1, 10 ** (i + 1)) * (0.3 ** (i - 1)), + ], + dim=1, + ) + input_var.append(input_var_rank) + + all_input_var = torch.cat( + [ + x.permute(1, 0, 2).contiguous().view(ONLY_SBN_NET.num_features, -1) + for x in input_var + ], + dim=1, + ).cuda(rank) + + for _ in range(100): + y = model(input_var[rank].cuda(rank)) + y.mean().backward() + + running_mean, running_var = ( + model.module.running_mean, + model.module.running_var, + ) + torch.testing.assert_close(running_mean, all_input_var.mean(1)) + torch.testing.assert_close(running_var, all_input_var.var(1)) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_DistributedDataParallel_SyncBatchNorm_Diff_Input_Sizes_gradient(self): + _group, _group_id, rank = self._init_global_test() + # only do single GPU per process + gpus = [rank] + + # cpu training setup + num_processes = dist.get_world_size() + local_bs = rank + 2 + bs_offset = int((rank + 3) * rank / 2) + global_bs = int((num_processes + 3) * num_processes / 2) + + self._test_DistributedDataParallel_SyncBatchNorm( + gpu_subset=gpus, + rank=rank, + local_bs=local_bs, + global_bs=global_bs, + offset=bs_offset, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_DistributedDataParallel_SyncBatchNorm_half(self): + _group, _group_id, rank = self._init_global_test() + + model = BatchNormNet() + model = model.half() + model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) + model = nn.parallel.DistributedDataParallel( + model.cuda(rank), device_ids=[rank] + ) + inp = torch.randn(2, 2, dtype=torch.float16, device=torch.device(rank)) + # Check that forward/backward do not error with dtype mismatch + out = model(inp) + self.assertEqual(out.dtype, torch.float16) + out.sum().backward() + for param in model.parameters(): + self.assertEqual(param.grad.dtype, torch.float16) + + def _test_ddp_logging_data(self, is_gpu): + rank = dist.get_rank() + model_DDP = Net() + if is_gpu: + model_DDP = nn.parallel.DistributedDataParallel( + model_DDP.cuda(rank), device_ids=[rank] + ) + else: + model_DDP = nn.parallel.DistributedDataParallel(model_DDP) + + # dummy data initialization + local_bs = 2 + batch_size, input, target, loss = self._prepare_dummy_data(local_bs) + if is_gpu: + input = input.cuda(rank) + target = target.cuda(rank) + + model_DDP._set_ddp_runtime_logging_sample_rate(2) + + for idx in range(20): + offset = rank * local_bs + + # DDP training, DDP scatters subsets of input to nodes/GPUs + self._test_DDP_helper( + model_DDP, + input[offset : offset + local_bs], + target[offset : offset + local_bs], + loss, + 1, + ) + + self._model_step_with_zero_grad(model_DDP) + + # Verify DDP logging data is sampled as expected + # If it has ran more than 10 iterations and this is + # the sampled iteration for measuring run time stats, + # the run time stats for this idx-th iteration will not + # be zeros. + ddp_logging_data = model_DDP._get_ddp_logging_data() + if idx > 0 and (idx < 10 or idx % 2 == 0): + self.assertGreaterEqual( + ddp_logging_data.get("forward_compute_time"), 1 + ) + self.assertGreaterEqual( + ddp_logging_data.get("backward_compute_time"), 1 + ) + self.assertGreaterEqual( + ddp_logging_data.get("backward_comm_time"), 1 + ) + self.assertGreaterEqual( + ddp_logging_data.get("backward_compute_time"), + ddp_logging_data.get("backward_compute_comm_overlap_time"), + ) + self.assertGreaterEqual( + ddp_logging_data.get("backward_comm_time"), + ddp_logging_data.get("backward_compute_comm_overlap_time"), + ) + self.assertEqual(ddp_logging_data.get("iteration"), idx) + elif idx > 0: + # if the idx-th iteration is not sampled to set runtime stats, + # ddp_logging_data.iteration will not be updated to current + # iteration. + self.assertNotEqual(ddp_logging_data.get("iteration"), idx) + + # Shuffle the input so that DDP input is different + input = input[torch.randperm(batch_size)] + + return model_DDP + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "nccl does not support DDP on CPU models" + ) + def test_ddp_logging_data_cpu(self): + def parse_env(var): + return os.environ.get(var, "N/A") + + dist.set_debug_level(dist.DebugLevel.INFO) + _, group_id, _ = self._init_global_test() + model_DDP = self._test_ddp_logging_data(is_gpu=False) + + ddp_logging_data = model_DDP._get_ddp_logging_data() + self.assertEqual(ddp_logging_data.get("world_size"), dist.get_world_size()) + self.assertEqual(ddp_logging_data.get("rank"), dist.get_rank()) + self.assertEqual(ddp_logging_data.get("module_name"), "Net") + self.assertEqual(ddp_logging_data.get("device_ids"), "") + # output_device is -1 in default if it is not set, e.g. + # output_device of CPU training is -1. + self.assertEqual(ddp_logging_data.get("output_device"), -1) + self.assertEqual(ddp_logging_data.get("broadcast_buffers"), 1) + self.assertEqual(ddp_logging_data.get("bucket_cap_bytes"), 25 * 1024 * 1024) + self.assertEqual(ddp_logging_data.get("find_unused_parameters"), 0) + self.assertEqual(ddp_logging_data.get("gradient_as_bucket_view"), 0) + self.assertEqual( + ddp_logging_data.get("backend_name"), dist.get_backend(group_id) + ) + self.assertEqual(ddp_logging_data.get("iteration"), 18) + params = list(model_DDP.parameters()) + num_params = 0 + param_size = 0 + params = list(filter(lambda parameter: parameter.requires_grad, params)) + for p in params: + num_params += 1 + param_size += p.numel() * p.element_size() + self.assertEqual(ddp_logging_data.get("dtypes"), "float") + self.assertEqual( + ddp_logging_data.get("total_parameter_size_bytes"), param_size + ) + self.assertEqual(ddp_logging_data.get("num_parameter_tensors"), num_params) + self.assertEqual(ddp_logging_data.get("bucket_sizes"), str(param_size)) + self.assertEqual( + ddp_logging_data.get("master_port"), parse_env("MASTER_PORT") + ) + self.assertEqual( + ddp_logging_data.get("master_addr"), parse_env("MASTER_ADDR") + ) + self.assertEqual( + ddp_logging_data.get("torch_distributed_debug"), + parse_env("TORCH_DISTRIBUTED_DEBUG"), + ) + self.assertEqual( + ddp_logging_data.get("cuda_visible_devices"), + parse_env("CUDA_VISIBLE_DEVICES"), + ) + if ddp_logging_data.get("backend_name") == "gloo": + self.assertEqual( + ddp_logging_data.get("gloo_socket_ifname"), + parse_env("GLOO_SOCKET_IFNAME"), + ) + self.assertEqual( + ddp_logging_data.get("gloo_device_transport"), + parse_env("GLOO_DEVICE_TRANSPORT"), + ) + default_gloo_threads = 2 + self.assertEqual( + ddp_logging_data.get("gloo_num_threads"), + default_gloo_threads, + ) + + self.assertEqual(ddp_logging_data.get("nccl_socket_ifname"), None) + self.assertEqual(ddp_logging_data.get("nccl_blocking_wait"), None) + self.assertEqual(ddp_logging_data.get("nccl_async_error_handling"), None) + self.assertEqual(ddp_logging_data.get("nccl_debug"), None) + self.assertEqual(ddp_logging_data.get("nccl_nthreads"), None) + self.assertEqual(ddp_logging_data.get("nccl_ib_timeout"), None) + # test runtime logging fields + # Note: DETAIL debug mode logs DDP logging data to stdout and + # thus accesses std::map, which fills in a default value for the + # type if it didn't exist. + self.assertEqual(ddp_logging_data.get("unused_parameter_size", 0), 0) + self.assertEqual(ddp_logging_data.get("has_rebuilt_buckets"), 1) + self.assertEqual( + ddp_logging_data.get("rebuilt_bucket_sizes"), str(param_size) + ) + grad_ready_order = ddp_logging_data.get( + "prev_iteration_grad_ready_order_indices" + ) + expected_order = list(reversed([str(x) for x in range(3)])) + self.assertEqual(grad_ready_order, ", ".join(expected_order)) + bucket_indices = ddp_logging_data.get("rebuilt_per_bucket_param_indices") + self.assertEqual(bucket_indices, " ".join(expected_order)) + # It is hard to test accurate latency, but it can test whether the latency is + # a valid value and in the expected range. + self.assertGreaterEqual(ddp_logging_data.get("avg_forward_compute_time"), 1) + self.assertGreaterEqual( + ddp_logging_data.get("avg_backward_compute_time"), 1 + ) + self.assertGreaterEqual(ddp_logging_data.get("avg_backward_comm_time"), 1) + self.assertGreaterEqual( + ddp_logging_data.get("avg_backward_compute_time"), + ddp_logging_data.get("avg_backward_compute_comm_overlap_time"), + ) + self.assertGreaterEqual( + ddp_logging_data.get("avg_backward_comm_time"), + ddp_logging_data.get("avg_backward_compute_comm_overlap_time"), + ) + # Test host-side times are roughly in the order that we expect + fwd_host_side_time = ddp_logging_data.get("forward_compute_time_start") + bwd_comp_start_host_side_time = ddp_logging_data.get( + "backward_compute_time_start" + ) + bwd_comp_end_host_side_time = ddp_logging_data.get( + "backward_compute_time_end" + ) + bwd_comm_start_host_side_time = ddp_logging_data.get( + "backward_comm_time_start" + ) + bwd_comm_end_host_side_time = ddp_logging_data.get("backward_comm_time_end") + self.assertGreaterEqual( + bwd_comm_end_host_side_time, bwd_comm_start_host_side_time + ) + self.assertGreaterEqual( + bwd_comm_start_host_side_time, bwd_comp_start_host_side_time + ) + self.assertGreaterEqual( + bwd_comp_end_host_side_time, bwd_comp_start_host_side_time + ) + self.assertGreaterEqual(bwd_comp_start_host_side_time, fwd_host_side_time) + + # test larger net with mixed data types, verify multiple bucket sizes + model = LargeNet() + model.float() + model.fc1.double() + model_DDP = nn.parallel.DistributedDataParallel(model, bucket_cap_mb=1.5) + ddp_logging_data = model_DDP._get_ddp_logging_data() + params = list(model_DDP.parameters()) + self.assertEqual( + ddp_logging_data.get("bucket_cap_bytes"), int(1.5 * 1024 * 1024) + ) + bucket_sizes = [ + params[1].numel() * params[1].element_size(), + params[0].numel() * params[0].element_size(), + ] + self.assertEqual( + ddp_logging_data.get("bucket_sizes"), + ", ".join(str(x) for x in bucket_sizes), + ) + self.assertEqual(ddp_logging_data.get("dtypes"), "double, float") + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_no_gpu + def test_ddp_logging_data_gpu(self): + _group, _group_id, rank = self._init_global_test() + model_DDP = self._test_ddp_logging_data(is_gpu=True) + ddp_logging_data = model_DDP._get_ddp_logging_data() + self.assertEqual(ddp_logging_data.get("device_ids"), str(rank)) + self.assertEqual(ddp_logging_data.get("output_device"), rank) + grad_ready_order = ddp_logging_data.get( + "prev_iteration_grad_ready_order_indices" + ) + expected_order = list(reversed([str(x) for x in range(3)])) + self.assertEqual(grad_ready_order, ", ".join(expected_order)) + bucket_indices = ddp_logging_data.get("rebuilt_per_bucket_param_indices") + self.assertEqual(bucket_indices, " ".join(expected_order)) + # test runtime logging fields + # It is hard to test accurate latency, but it can test whether the latency is + # a valid value and in the expected range. + self.assertGreaterEqual(ddp_logging_data.get("avg_forward_compute_time"), 1) + self.assertGreaterEqual( + ddp_logging_data.get("avg_backward_compute_comm_overlap_time"), 1 + ) + self.assertGreaterEqual( + ddp_logging_data.get("avg_backward_compute_time"), + ddp_logging_data.get("avg_backward_compute_comm_overlap_time"), + ) + self.assertGreaterEqual( + ddp_logging_data.get("avg_backward_comm_time"), + ddp_logging_data.get("avg_backward_compute_comm_overlap_time"), + ) + # Test host-side times are roughly in the order that we expect + fwd_host_side_time = ddp_logging_data.get("forward_compute_time_start") + bwd_comp_start_host_side_time = ddp_logging_data.get( + "backward_compute_time_start" + ) + bwd_comp_end_host_side_time = ddp_logging_data.get( + "backward_compute_time_end" + ) + bwd_comm_start_host_side_time = ddp_logging_data.get( + "backward_comm_time_start" + ) + bwd_comm_end_host_side_time = ddp_logging_data.get("backward_comm_time_end") + self.assertGreaterEqual( + bwd_comm_end_host_side_time, bwd_comm_start_host_side_time + ) + self.assertGreaterEqual( + bwd_comm_start_host_side_time, bwd_comp_start_host_side_time + ) + self.assertGreaterEqual( + bwd_comp_end_host_side_time, bwd_comp_start_host_side_time + ) + self.assertGreaterEqual(bwd_comp_start_host_side_time, fwd_host_side_time) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_bucket_cap_mb_list(self): + """ + Test that bucket_cap_mb_list parameter is properly used for bucket rebuilding. + When bucket_cap_mb_list is provided, it should be used to rebuild buckets. + When not provided, original logic should be used. + """ + torch.cuda.set_device(self.rank) + + # Test 1: With bucket_cap_mb_list provided + model_with_list = LargeNet().cuda(self.rank) + bucket_cap_mb_list = [1, 2] # Two buckets with different sizes + ddp_with_list = nn.parallel.DistributedDataParallel( + model_with_list, + device_ids=[self.rank], + bucket_cap_mb_list=bucket_cap_mb_list, + ) + + # Run a forward/backward pass to trigger bucket rebuilding + input_tensor = torch.randn(10, 1000).cuda(self.rank) + loss = ddp_with_list(input_tensor).sum() + loss.backward() + + # Run another iteration to trigger rebuild + input_tensor = torch.randn(10, 1000).cuda(self.rank) + loss = ddp_with_list(input_tensor).sum() + loss.backward() + + # When bucket_cap_mb_list is provided, buckets are created with the correct + # structure from the start. No "rebuild" is needed since initial bucket + # creation uses the same limits as rebuild would. Hence has_rebuilt_buckets=0 + # is the expected behavior. + # Verify bucket_bytes_cap_list was properly set instead. + expected_bucket_bytes_cap_list = [ + int(cap * 1024 * 1024) for cap in bucket_cap_mb_list + ] + self.assertEqual( + ddp_with_list.bucket_bytes_cap_list, expected_bucket_bytes_cap_list + ) + + # Test 2: Without bucket_cap_mb_list (backward compatibility) + model_without_list = LargeNet().cuda(self.rank) + ddp_without_list = nn.parallel.DistributedDataParallel( + model_without_list, + device_ids=[self.rank], + bucket_cap_mb=25, # Standard bucket size + ) + + # Run a forward/backward pass + input_tensor = torch.randn(10, 1000).cuda(self.rank) + loss = ddp_without_list(input_tensor).sum() + loss.backward() + + # Run another iteration to trigger rebuild + input_tensor = torch.randn(10, 1000).cuda(self.rank) + loss = ddp_without_list(input_tensor).sum() + loss.backward() + + # Verify that DDP without bucket_cap_mb_list still works correctly + # and bucket_bytes_cap_list should be empty (using original logic) + self.assertEqual(ddp_without_list.bucket_bytes_cap_list, []) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "nccl", "nccl does not support DDP on CPU models" + ) + def test_static_graph_api_cpu(self): + model_DDP = nn.parallel.DistributedDataParallel(Net()) + expected_err = "should be called before training loop starts" + with self.assertRaisesRegex(RuntimeError, expected_err): + local_bs = 2 + _batch_size, input, target, loss = self._prepare_dummy_data(local_bs) + offset = dist.get_rank() * local_bs + + # DDP training, DDP scatters subsets of input to nodes/GPUs + self._test_DDP_helper( + model_DDP, + input[offset : offset + local_bs], + target[offset : offset + local_bs], + loss, + 1, + ) + model_DDP._set_static_graph() + + # Verify error was logged in ddp_logging_data. + verify_ddp_error_logged(model_DDP, expected_err) + + @skipIfNoTorchVision + def test_SyncBatchNorm_process_group(self): + # When adopting `convert_sync_batchnorm` to convert a `nn.modules`, + # it need to recursively pass the `process_group` in the module when the `SyncBatchNorm` + # is nested in a sub-module or sub-sub-module (e.g. resnet50 in torchvision.models). + + process_ids = 0 + process_group = torch.distributed.new_group([process_ids]) + res50_model = torchvision.models.resnet50() + res50_model_sync = nn.SyncBatchNorm.convert_sync_batchnorm( + copy.deepcopy(res50_model), process_group + ) + process_group_sync = res50_model_sync.layer1[0].bn1.process_group + self.assertEqual(process_group_sync, process_group) + + def _run_reduction_test( + self, tensor, expected_tensor, op, reduction_fn=dist.all_reduce, dst=None + ): + if reduction_fn is not dist.all_reduce and dst is None: + raise ValueError(f"Reduction fn {reduction_fn} must specify dst!") + if dst is not None: + reduction_fn(tensor, dst, op) + # Only destination rank tensor is expected to have final result. + if dist.get_rank() == dst: + self.assertEqual(tensor, expected_tensor) + else: + reduction_fn(tensor, op) + self.assertEqual(tensor, expected_tensor) + + @require_backend_is_available({"nccl"}) + @skip_if_lt_x_gpu(2) + def test_nccl_backend_bool_allreduce(self): + torch.cuda.set_device(self.rank) + # Run all_reduce with PRODUCT + element = self.rank % 2 == 0 + for op in [dist.ReduceOp.PRODUCT, dist.ReduceOp.MIN]: + input_tensor = torch.tensor([element, element]).to(self.rank) + self._run_reduction_test( + input_tensor, torch.tensor([False, False]).to(self.rank), op + ) + # Ensure that all ranks contributing True (cast to 1) results in the + # correct reduction. + input_tensor = torch.tensor([True, True]).to(self.rank) + expected_tensor = input_tensor.clone() + self._run_reduction_test(input_tensor, expected_tensor, op) + + # Run all_reduce with SUM + for op in [dist.ReduceOp.SUM, dist.ReduceOp.MAX]: + input_tensor = torch.tensor([element, element]).to(self.rank) + self._run_reduction_test( + input_tensor, torch.tensor([True, True]).to(self.rank), op + ) + # TODO: NCCL backend does not work correctly for bitwise reduction ops + # (see https://github.com/pytorch/pytorch/issues/41362). Add tests for + # these once it is supported. + + @require_backend_is_available({"nccl"}) + @skip_if_lt_x_gpu(2) + def test_nccl_backend_bool_allgather(self): + torch.cuda.set_device(self.rank) + inp = {0: [True, True], 1: [False, True]} + input_tensor = torch.tensor(inp[self.rank % 2]).to(self.rank) + # Preserve a copy of the tensor to compare against after allgather. + input_tensor_copy = input_tensor.clone() + tensor_list = [ + torch.tensor([False, False]).to(self.rank) + for _ in range(dist.get_world_size()) + ] + dist.all_gather(tensor_list, input_tensor) + + self.assertEqual(len(tensor_list), dist.get_world_size()) + for i, t in enumerate(tensor_list): + expected = torch.tensor(inp[i % 2]).to(self.rank) + self.assertEqual(t, expected) + # Ensure that the input tensor is not modified, since this collective + # does not modify its input. + self.assertEqual(input_tensor_copy, input_tensor) + + @require_backend_is_available({"nccl"}) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_nccl_backend_bool_reduce(self): + torch.cuda.set_device(self.rank) + inp = {0: [True, True], 1: [False, False]} + # Run reduce() with product op + for op in [dist.ReduceOp.PRODUCT, dist.ReduceOp.MIN]: + # make sure rank 0 gets False if WORLD_SIZE=1 to match expected tensor + input_tensor = torch.tensor(inp[(self.rank + 1) % 2]).to(self.rank) + expected = torch.tensor([False, False]).to(self.rank) + self._run_reduction_test(input_tensor, expected, op, dist.reduce, dst=0) + # Ensure that all ranks contributing True (cast to 1) results in the + # correct reduction. + input_tensor = torch.tensor([True, True]).to(self.rank) + expected_tensor = input_tensor.clone() + self._run_reduction_test( + input_tensor, expected_tensor, op, dist.reduce, dst=0 + ) + + for op in [dist.ReduceOp.SUM, dist.ReduceOp.MAX]: + input_tensor = torch.tensor(inp[self.rank % 2]).to(self.rank) + expected = ( + torch.tensor([True, True]).to(self.rank) + if self.rank == 0 + else input_tensor.clone() + ) + self._run_reduction_test(input_tensor, expected, op, dist.reduce, dst=0) + + @require_backend_is_available({"nccl"}) + @skip_if_lt_x_gpu(2) + def test_nccl_backend_bool_broadcast(self): + tensor_size = 10 + bcast_tensor = torch.tensor( + [ + (random.random() < 0.5 if self.rank == 0 else False) + for _ in range(tensor_size) + ] + ).to(self.rank) + dist.broadcast(bcast_tensor, src=0) + # Now allgather and ensure the tensors are equal. + tensor_list = [ + torch.tensor([False for _ in range(tensor_size)]).to(self.rank) + for _ in range(dist.get_world_size()) + ] + dist.all_gather(tensor_list, bcast_tensor) + expected = tensor_list[0] + for tensor in tensor_list[1:]: + self.assertEqual(tensor, expected) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_DistributedSampler_padding(self): + # Tests padding of distributed sampler. + world_size = dist.get_world_size() + + # Simulates the 'casual' dataset size + dataset_size = 100 + world_size + 1 + dataset = [torch.ones(1).to(self.rank) * i for i in range(dataset_size)] + + # Simulates the 'tiny' dataset size + dataset_tiny_size = max(world_size // 2 - 1, 1) + dataset_tiny = [ + torch.ones(1).to(self.rank) * i for i in range(dataset_tiny_size) + ] + + # Specifying drop_last=True will cause the tail of the data to be dropped. + dist_sampler = DistributedSampler(dataset=dataset, drop_last=True) + local_num_samples, local_dataset_size = ( + dist_sampler.num_samples, + dist_sampler.total_size, + ) + # The effective dataset size should be the greatest integer that is <= + # dataset_size that is divisible by the world_size. This is to ensure each + # rank processes the same number of samples. + effective_dataset_size = ( + math.ceil((dataset_size - world_size) / world_size) + if dataset_size % world_size != 0 + else dataset_size / world_size + ) + self.assertEqual(local_num_samples, effective_dataset_size) + self.assertEqual(local_dataset_size, local_num_samples * world_size) + indices_list = list(iter(dist_sampler)) + self.assertEqual(len(indices_list), local_num_samples) + + def validate_global_samples(local_num_samples): + # Ensure that each rank processes the same number of samples. + world_samples = [ + torch.LongTensor([0]).to(self.rank) for _ in range(world_size) + ] + dist.all_gather( + world_samples, torch.tensor([local_num_samples]).to(self.rank) + ) + world_samples = [sample.item() for sample in world_samples] + self.assertEqual(len(set(world_samples)), 1) + + validate_global_samples(local_num_samples) + + # drop_last=False is the default and will add additional indices to be sampled, + # increasing the effective dataset size. + dist_sampler_added_samples = DistributedSampler(dataset=dataset) + local_num_samples, local_dataset_size = ( + dist_sampler_added_samples.num_samples, + dist_sampler_added_samples.total_size, + ) + # The effective dataset size is the smallest integer that is >= dataset_size + # and divisible by the world size. + self.assertEqual(local_num_samples, math.ceil(dataset_size / world_size)) + self.assertEqual(local_dataset_size, local_num_samples * world_size) + indices_list = list(iter(dist_sampler_added_samples)) + self.assertEqual(len(indices_list), local_num_samples) + + # Ensure that each rank processes the same number of samples. + validate_global_samples(local_num_samples) + + # Ensure additional samples are padded even when + # the extremely small dataset is given. + dist_sampler_added_samples_tiny = DistributedSampler(dataset=dataset_tiny) + local_num_samples, local_dataset_size = ( + dist_sampler_added_samples_tiny.num_samples, + dist_sampler_added_samples_tiny.total_size, + ) + self.assertEqual( + local_num_samples, math.ceil(dataset_tiny_size / world_size) + ) + self.assertEqual(local_dataset_size, local_num_samples * world_size) + indices_list = list(iter(dist_sampler_added_samples_tiny)) + self.assertEqual(len(indices_list), local_num_samples) + validate_global_samples(local_num_samples) + + def _test_allgather_object(self, subgroup=None): + # Only set device for NCCL backend since it must use GPUs. + + gather_objects = create_collectives_object_test_list() + + backend = os.environ["BACKEND"] + if backend == "nccl": + # Case where rank != GPU device. + next_rank = (self.rank + 1) % int(self.world_size) + torch.cuda.set_device(next_rank) + + # If GPU test, add object with GPU tensor + if backend == "nccl": + gather_objects.append(Foo(torch.randn(3, 3, device=0))) + + output_gathered = [None for _ in range(dist.get_world_size())] + dist.all_gather_object( + output_gathered, + gather_objects[self.rank % len(gather_objects)], + group=subgroup, + ) + + for i, val in enumerate(output_gathered): + expected = gather_objects[i % len(gather_objects)] + self.assertEqual(val, expected) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @require_n_gpus_for_nccl_backend( + int(os.environ["WORLD_SIZE"]), os.environ["BACKEND"] + ) + @with_dist_debug_levels(levels=["OFF", "INFO", "DETAIL"]) + def test_all_gather_object_default_pg(self): + return self._test_allgather_object() + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @require_n_gpus_for_nccl_backend( + int(os.environ["WORLD_SIZE"]), os.environ["BACKEND"] + ) + @with_dist_debug_levels(levels=["DETAIL", "OFF", "INFO"]) + def test_all_gather_object_subgroup(self): + default = _get_default_group() + backend = dist.get_backend(default) + subgroup = dist.new_group(backend=backend) + return self._test_allgather_object(subgroup=subgroup) + + def _test_gather_object(self, pg=None): + # Ensure stateful objects can be gathered + gather_objects = create_collectives_object_test_list() + my_rank = dist.get_rank(pg) + + backend = os.environ["BACKEND"] + if backend == "nccl": + # Case where rank != GPU device. + next_rank = (self.rank + 1) % int(self.world_size) + torch.cuda.set_device(next_rank) + + # If GPU test, add object with GPU tensor + if backend == "nccl": + gather_objects.append(Foo(torch.randn(3, 3, device=my_rank))) + + output_gathered = [None for _ in range(dist.get_world_size(pg))] + gather_on_rank = 0 + dist.gather_object( + gather_objects[self.rank % len(gather_objects)], + object_gather_list=output_gathered + if my_rank == gather_on_rank + else None, + dst=gather_on_rank, + group=pg, + ) + if my_rank != gather_on_rank: + self.assertEqual( + output_gathered, [None for _ in range(dist.get_world_size())] + ) + else: + for i, val in enumerate(output_gathered): + expected = gather_objects[i % len(gather_objects)] + self.assertEqual(val, expected) + + # Validate errors when objects can't be pickled. + class Bar: + pass + + b = Bar() + gather_objects = [b for _ in range(dist.get_world_size())] + with self.assertRaises(AttributeError): + dist.all_gather_object( + [None for _ in range(dist.get_world_size())], + gather_objects[self.rank], + group=pg, + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @with_dist_debug_levels(levels=["DETAIL", "OFF", "INFO"]) + @require_exact_world_size(4) + def test_gather_object(self): + return self._test_gather_object() + + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc", "CPU tensor ops not supported by UCP TL" + ) + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @with_dist_debug_levels(levels=["DETAIL", "OFF", "INFO"]) + @require_exact_world_size(4) + def test_gather_object_subgroup(self): + default = _get_default_group() + backend = dist.get_backend(default) + subgroup = dist.new_group(backend=backend) + return self._test_gather_object(subgroup) + + def validate_net_equivalence(self, net): + # Helper to validate synchronization of nets across ranks. + net_module_states = list(net.module.state_dict().values()) + # Check that all tensors in module's state_dict() are equal. + for t in net_module_states: + tensor_list = [ + torch.zeros_like(t) for _ in range(dist.get_world_size()) + ] + dist.all_gather(tensor_list, t) + for tensor in tensor_list: + self.assertEqual(tensor, t) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_sync_module_states(self): + # Test that after calling _sync_module_states, models across ranks + # are the same and are equal to the model on the input rank. + dim = 2 + rank = self.rank + rank_to_broadcast = 1 + # Seed to ensure that ranks are initialized with different initial models. + torch.manual_seed(rank) + model = nn.Linear(dim, dim, bias=False) + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(rank), device_ids=[self.rank], bucket_cap_mb=1 + ) + new_model = nn.Linear(dim, dim, bias=False).cuda(rank) + net.module = copy.deepcopy(new_model) + # Assert params are different + net_module_states = list(net.module.state_dict().values()) + for t in net_module_states: + tensor_list = [ + torch.zeros_like(t) for _ in range(dist.get_world_size()) + ] + dist.all_gather(tensor_list, t) + for i, tensor in enumerate(tensor_list): + if i == rank: + self.assertEqual(t, tensor) + else: + # tensor from another rank should be different. + self.assertNotEqual(t, tensor) + + _sync_module_states( + module=net.module, + process_group=net.process_group, + broadcast_bucket_size=net.broadcast_bucket_size, + src=rank_to_broadcast, + params_and_buffers_to_ignore=net.parameters_to_ignore, + ) + # Now all model params should be the same. + self.validate_net_equivalence(net) + # Since the network params were broadcast from rank_to_broadcast, validate that + # they are the same as new_model on rank_to_broadcast. + if rank == rank_to_broadcast: + expected_states = new_model.state_dict().values() + for t, expected in zip(net_module_states, expected_states, strict=True): + self.assertEqual(t, expected) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_grad_div_uneven_inputs(self): + # Test gradient division during training with join() API. If + # divide_by_initial_world_size=False, we scale by the effective world + # size when allreducing grads. + dim = 5 + batch = 1 + grad_scale = 50 + rank = self.rank + model = nn.Linear(dim, dim, bias=False) + inp = torch.ones(batch, dim, device=self.rank) * grad_scale + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(rank), device_ids=[self.rank], bucket_cap_mb=1 + ) + n_iters = 3 + if self.rank > 0: + n_iters += 2 + + with net.join(divide_by_initial_world_size=False): + for _ in range(n_iters): + loss = net(inp).sum() + loss.backward() + # The grad is always expected_grad, since we divide by the number + # of currently active processes and inactive processes contribute + # zero gradient. If we kept dividing by static initial world + # size as processes leave, the grad would be smaller. + expected_grad = torch.ones(dim, dim, device=self.rank) * grad_scale + param = next(iter(net.parameters())) + self.assertEqual(expected_grad, param.grad) + # Avoid accumulating grads so that it's the same every iteration + net.zero_grad() + torch.cuda.synchronize(device=self.rank) + + # If divide_by_initial_world_size=True (default), we always scale grads + # by the initial world_size. + with net.join(divide_by_initial_world_size=True): + for i in range(n_iters): + loss = net(inp).sum() + loss.backward() + effective_ws = dist.get_world_size() + if i >= 3: + effective_ws -= 1 + expected_grad = ( + torch.ones(dim, dim, device=self.rank) + * grad_scale + * effective_ws + ) / dist.get_world_size() + param = next(iter(net.parameters())) + self.assertEqual(expected_grad, param.grad) + # Avoid accumulating grad so that it's the same every iteration. + net.zero_grad() + torch.cuda.synchronize(device=self.rank) + + def _test_ddp_profiling(self, profiler_ctx, profiler_ctx2=None): + """Runs DDP based model training and captures profiles. + This test will do two profiler runs. + 1. An initial basic run to check if profiler events are correctly captured. + 2. A second profiling pass after running some iterations of DDP, to check robustness of thread local state. + + args + profiler_ctx : Profiler context manager for pass 1 + profiler_ctx2 : Profiler context manager for pass 2. + This can be left out as None, in which case a deepcopy + of profiler_ctx is used. + Returns: + prof: Instantiated profiler object that can be used for post analysis. + """ + batch = 3 + dim = 10 + num_iters = 6 + torch.cuda.set_device(self.rank) + model = nn.Linear(dim, dim, bias=False) + inp = torch.rand(batch, dim, device=self.rank) + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(self.rank), + device_ids=[self.rank], + ) + if profiler_ctx2 is None: + profiler_ctx2 = copy.deepcopy(profiler_ctx) + + with profiler_ctx as prof: + for _ in range(num_iters): + loss = net(inp).sum() + loss.backward() + + all_reduce_event_name = f"{dist.get_backend()}:all_reduce" + events = get_profiling_event( + all_reduce_event_name, prof, dedup_gpu_user_annotation=True + ) + event_count = sum(e.count for e in events) + self.assertEqual(event_count, num_iters) + for event in events: + self.assertTrue(event.is_async) + self.assertEqual(event.name, all_reduce_event_name) + + broadcast_event_name = f"{dist.get_backend()}:broadcast" + broadcast_events = get_profiling_event( + broadcast_event_name, prof, dedup_gpu_user_annotation=True + ) + event_count = sum(e.count for e in broadcast_events) + # Broadcast is called during rebuild_buckets + self.assertGreaterEqual(event_count, 1) + for event in broadcast_events: + self.assertEqual(event.name, broadcast_event_name) + + # Run DDP with profiling for a few iterations, then enable profiling + # for a single pass, and ensure it is recorded. This tests that the + # thread local state is correctly updated. + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(self.rank), + device_ids=[self.rank], + find_unused_parameters=True, + ) + for _ in range(3): + loss = net(inp).sum() + loss.backward() + # Now enable the profiler. + with profiler_ctx2 as prof: + loss = net(inp).sum() + loss.backward() + + events = get_profiling_event( + all_reduce_event_name, prof, dedup_gpu_user_annotation=True + ) + self.assertGreaterEqual(len(events), 1) + self.assertGreaterEqual(events[0].count, 1) + self.assertEqual(events[0].name, all_reduce_event_name) + for event in events: + self.assertTrue(event.is_async) + # Ensure searching unused parameters was profiled + events = get_profiling_event("search_unused_parameters", prof) + self.assertEqual(len(events), 1) + + return prof + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle("Currently failing in NVIDIA internal CI") + def test_ddp_profiling_autograd_profiler(self): + autograd_profiler_ctx = torch.autograd.profiler.profile() + return self._test_ddp_profiling(profiler_ctx=autograd_profiler_ctx) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if(IS_FBCODE, "Kineto in fbcode code causes hang") + @skip_but_pass_in_sandcastle_if( + IS_MACOS or IS_WINDOWS, + "torch.profiler not enabled for mac/windows: https://github.com/pytorch/pytorch/pull/56124", + ) + def test_ddp_profiling_torch_profiler(self): + cpu_act = torch.profiler.ProfilerActivity.CPU + cuda_act = torch.profiler.ProfilerActivity.CUDA + torch_profiler_ctx = torch.profiler.profile(activities=[cpu_act, cuda_act]) + prof = self._test_ddp_profiling(profiler_ctx=torch_profiler_ctx) + + if dist.get_backend() != "nccl": + return + + # Note comment out the "os.remove(trace_file)" in `get_profiler_nccl_meta()` + # to debug any mismatches. + nccl_meta_events = get_profiler_nccl_meta(prof) + self.assertGreater(len(nccl_meta_events), 0) + + nccl_meta = self._sanity_check_profiler_nccl_meta(nccl_meta_events) + + # additionally check the specific collectives in this test case + self.assertEqual(len(nccl_meta["allreduce"]), 2) + self.assertEqual(len(nccl_meta["wait"]), 1) + + # check allreduce message sizes + a0 = nccl_meta["allreduce"][0] + self.assertEqual(a0["Out msg nelems"], 100, msg=f"{a0}") + self.assertEqual(a0["dtype"], "Float", msg=f"{a0}") + a1 = nccl_meta["allreduce"][1] + self.assertEqual(a1["Out msg nelems"], 1, msg=f"{a1}") + self.assertEqual(a1["dtype"], "Int", msg=f"{a1}") + + def _validate_execution_trace_nccl(self, et_file: str) -> None: + """Torch profiler includes nccl metadata in an inserted operator called "record_param_comms" + We test for basic fields in these nodes in the Execution Trace. + """ + with open(et_file) as f: + et = json.load(f) + pg_cfg_node = [ + n for n in et["nodes"] if n["name"] == "## process_group:init ##" + ] + self.assertGreaterEqual(len(pg_cfg_node), 1) + nccl_meta_nodes = [ + n for n in et["nodes"] if n["name"] == "record_param_comms" + ] + self.assertEqual(len(nccl_meta_nodes), 3) + per_coll_meta = defaultdict(list) + + # Sanity check NCCL metadata nodes + for n in nccl_meta_nodes: + attrs_list = n.get("attrs", []) + self.assertGreater(len(attrs_list), 0) + attrs = {a["name"]: a["value"] for a in attrs_list} + + collname = attrs.get("collective_name", "") + self.assertNotEqual(collname, "") + self.assertNotEqual(attrs.get("dtype", ""), "") + + per_coll_meta[collname].append(attrs) + if collname == "wait": + continue + + self.assertEqual(attrs["pg_name"], "0") # yes this is a string + self.assertEqual(attrs["pg_desc"], "default_pg") + self.assertEqual(attrs["pg_size"], 2) + + self.assertGreaterEqual(attrs.get("in_msg_nelems", -1), 0) + self.assertGreaterEqual(attrs.get("out_msg_nelems", -1), 0) + self.assertTrue("in_split_size" in attrs) + self.assertTrue("out_split_size" in attrs) + self.assertEqual(attrs.get("global_rank_start", -1), 0) + self.assertEqual(attrs.get("global_rank_stride", -1), 1) + + # print(per_coll_meta) + self.assertEqual(len(per_coll_meta["allreduce"]), 2) + self.assertEqual(len(per_coll_meta["wait"]), 1) + + # check allreduce message sizes + a0 = per_coll_meta["allreduce"][0] + self.assertEqual(a0["out_msg_nelems"], 100, msg=f"{a0}") + self.assertEqual(a0["dtype"], "Float", msg=f"{a0}") + a1 = per_coll_meta["allreduce"][1] + self.assertEqual(a1["out_msg_nelems"], 1, msg=f"{a1}") + self.assertEqual(a1["dtype"], "Int", msg=f"{a1}") + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if(IS_FBCODE, "Kineto in fbcode code causes hang") + @skip_but_pass_in_sandcastle_if( + IS_MACOS or IS_WINDOWS, + "torch.profiler not enabled for mac/windows: https://github.com/pytorch/pytorch/pull/56124", + ) + @unittest.skipIf(BACKEND != "nccl", "Tests nccl metadata primarily.") + def test_ddp_profiling_execution_trace(self): + self.assertEqual(dist.get_backend(), "nccl") + # Create a temp file to save execution trace data + with TemporaryFileName("w+t", suffix=".et.json") as et_file: + et = ExecutionTraceObserver().register_callback(et_file) + + # first profiler context need not have ET + torch_profiler_ctx1 = torch.profiler.profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + ) + # collect ET in second profiler pass + torch_profiler_ctx2 = torch.profiler.profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + execution_trace_observer=et, + ) + self._test_ddp_profiling( + profiler_ctx=torch_profiler_ctx1, + profiler_ctx2=torch_profiler_ctx2, + ) + + print(f"Execution trace saved at {et_file}") + self._validate_execution_trace_nccl(et_file) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_join_model_equivalence(self): + # Verifies equivalence with model training locally and with DDP under + # the join context manager. + batch = 3 + dim = 10 + learning_rate = 0.03 + model = nn.Linear(dim, dim, bias=False) + inp = torch.rand(batch, dim, device=self.rank) + local_model = copy.deepcopy(model) + local_model = local_model.cuda(self.rank) + rank_to_iter_mapping = { + rank: 2 * (rank + 1) for rank in range(dist.get_world_size()) + } + # run local model + local_iters = sum(rank_to_iter_mapping.values()) + local_optim = torch.optim.SGD(local_model.parameters(), lr=learning_rate) + for _ in range(local_iters): + local_optim.zero_grad() + out = local_model(inp) + loss = out.sum() + loss.backward() + local_optim.step() + + # run DDP model with join API + num_iters = rank_to_iter_mapping[self.rank] + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(self.rank), device_ids=[self.rank] + ) + ddp_optim = torch.optim.SGD( + model.parameters(), lr=learning_rate * dist.get_world_size() + ) + with net.join(): + for _ in range(num_iters): + ddp_optim.zero_grad() + out = net(inp) + loss = out.sum() + loss.backward() + torch.cuda.synchronize(device=self.rank) + ddp_optim.step() + + # Validate model state dicts are equal + for (_, local_tensor), (_, dist_tensor) in zip( + local_model.state_dict().items(), + net.module.state_dict().items(), + strict=True, + ): + self.assertEqual(local_tensor, dist_tensor) + + def _run_uneven_inputs_test( + self, + test_case, + iteration_mapping, + find_unused_params, + ): + model = test_case.model + inp = test_case.inp + rank = self.rank + sync_interval = test_case.sync_interval + torch.cuda.set_device(rank) + # Ensure all outstanding GPU work is completed so this test runs independently. + dist.barrier() + # Bucket_cap_mb is intentionally low to test allreduce scheduling when + # there are many buckets. + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(rank), + device_ids=[rank], + bucket_cap_mb=1, + find_unused_parameters=find_unused_params, + ) + # Register hook if specified + if test_case.hook is not None: + net.register_comm_hook(test_case.state, test_case.hook) + print(f"registered hook {test_case.hook}") + + # Determine num iters for this rank via the passed in mapping. + num_iters = iteration_mapping[rank] + # If we throw when earliest rank terminates, we should ensure + # that we iterate for that minimum number of times. + num_iters_tensor = torch.tensor( + [num_iters], device=torch.cuda.current_device() + ) + dist.all_reduce(num_iters_tensor, op=dist.ReduceOp.MIN) + min_num_iters = num_iters_tensor.item() + total_iters = 0 + if test_case.throw_on_early_termination: + if min_num_iters == num_iters: + # Early termination rank(s) + exception_ctx = self.assertRaisesRegex( + RuntimeError, f"Rank {self.rank} exhausted all inputs" + ) + else: + # Non early termination rank + exception_ctx = self.assertRaisesRegex( + RuntimeError, + "Detected at least one rank that exhausted inputs.", + ) + else: + exception_ctx = nullcontext() + with exception_ctx: + with net.join( + throw_on_early_termination=test_case.throw_on_early_termination + ): + for i in range(num_iters): + # Use model.no_sync() to disable grad synchronization every + # sync_interval. + if i % sync_interval != 0: + context = net.no_sync() + else: + context = nullcontext() + with context: + if isinstance(inp, tuple): + loss = net(*inp).sum() + else: + loss = net(inp).sum() + loss.backward() + self._model_step(net) + # Ensure completion of GPU kernels (including allreduce). If the + # join API is not properly implemented, then this should hang + # since the allreduce will hang. + torch.cuda.synchronize(device=rank) + total_iters += 1 + if test_case.throw_on_early_termination: + # Ensure we iterated min_num_iters times. + self.assertEqual(total_iters, min_num_iters) + else: + # Ensure we iterated at least min_num_iters times. + self.assertGreaterEqual(total_iters, min_num_iters) + + # Ensure completion of all GPU kernels. + torch.cuda.synchronize(device=rank) + # When throwing on early rank termination, we do not + # broadcast model state from an authoritative rank. All models + # should already be in sync. + if not test_case.throw_on_early_termination: + self.assertTrue(net._authoritative_rank) + # All ranks should have agreed on the same authoritative_rank! + final_rank_tensor = torch.tensor( + [net._authoritative_rank], device=self.rank + ) + tensor_list = [ + torch.zeros_like(final_rank_tensor) + for _ in range(dist.get_world_size()) + ] + dist.all_gather(tensor_list, final_rank_tensor) + max_rank = dist.get_world_size() - 1 + self.assertSetEqual( + {max_rank}, {tensor.item() for tensor in tensor_list} + ) + # Ensure that all models are the same across ranks after all have joined. + self.validate_net_equivalence(net) + # Ensure that running with DDP uneven inputs was logged. + ddp_logging_data = net._get_ddp_logging_data() + self.assertTrue(ddp_logging_data.get("join_uneven_inputs")) + dist.barrier() + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_uneven_inputs_stop_iteration_sync_bn(self): + # Tests that uneven inputs join handler correctly throws StopIteration + # for models with SyncBN or general collective comm when + # throw_on_early_termination=True. + class ModelWithComm(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin = nn.Linear(2, 40, bias=False) + + def forward(self, x): + x = self.lin(x) + dist.all_reduce(x) + return x + + torch.cuda.set_device(self.rank) + model_bn = BatchNormNet() + model_bn = nn.SyncBatchNorm.convert_sync_batchnorm( + copy.deepcopy(model_bn) + ).cuda(self.rank) + comm_model = ModelWithComm().cuda(self.rank) + model_input = torch.randn(10, 2).cuda(torch.cuda.current_device()) + + for model in [model_bn, comm_model]: + model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + ) + min_num_iters = 5 + if self.rank != 0: + # Early termination rank(s) + num_iters = min_num_iters + exception_ctx = self.assertRaisesRegex( + RuntimeError, f"Rank {self.rank} exhausted all inputs" + ) + else: + # Non early termination rank + num_iters = min_num_iters * 2 + exception_ctx = self.assertRaisesRegex( + RuntimeError, + "Detected at least one rank that exhausted inputs.", + ) + n = 0 + with exception_ctx: + with model.join(throw_on_early_termination=True): + for _ in range(num_iters): + loss = model(model_input).sum() + loss.backward() + self._model_step(model) + n += 1 + + self.assertEqual(n, min_num_iters) + # Verify model equivalence + self.validate_net_equivalence(model) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_uneven_inputs(self): + dim = 1000 + batch = 1 + # Create a variety of models to run uneven input tests on. + large_model = nn.Sequential( + nn.Conv2d(1, 20, 5), + nn.ReLU(), + nn.Conv2d(20, 32, 5), + nn.ReLU(), + nn.Conv2d(32, 256, 5), + nn.ReLU(), + ) + small_model = nn.Linear(dim, dim, bias=False) + bn_net = BatchNormNet() + + class UnusedParamModule(nn.Module): + def __init__(self, unused_params_rank): + super().__init__() + self.t0 = Task() + self.t1 = Task() + self.unused_params_rank = unused_params_rank + + def task_parameters(self): + return (self.t0.p, self.t1.p) + + def forward(self, x, rank): + return ( + self.t1(self.t0(x)) + if rank != self.unused_params_rank + else self.t1(x) + ) + + unjoined_rank_with_unused_params_model = UnusedParamModule(1) + joined_rank_with_unused_params_model = UnusedParamModule(0) + + rank = self.rank + models_to_test = [ + # Network with batchnorm + DDPUnevenTestInput( + name="batch_norm_net", + model=bn_net, + inp=torch.ones(batch, 2, device=rank), + sync_interval=1, + ), + DDPUnevenTestInput( + name="large_conv_model", + model=large_model, + inp=torch.ones(batch, batch, dim, dim, device=rank), + sync_interval=1, + ), + DDPUnevenTestInput( + name="small_model", + model=small_model, + inp=torch.ones(batch, dim, device=rank), + sync_interval=1, + ), + # Unused parameter test where rank that does not join early has unused params + DDPUnevenTestInput( + name="unjoined_rank_with_unused_params_model", + model=unjoined_rank_with_unused_params_model, + inp=(torch.ones(batch, 2, device=rank), rank), + sync_interval=1, + ), + # Unused parameter test where rank that does join early has unused params + DDPUnevenTestInput( + name="joined_rank_with_unused_params_model", + model=joined_rank_with_unused_params_model, + inp=(torch.ones(batch, 2, device=rank), rank), + sync_interval=1, + ), + ] + + # Test models that have hook installed. + models_with_hook = [ + DDPUnevenTestInput( + name="small_model_allreduce_hook", + model=small_model, + hook=default.allreduce_hook, + state=None, + inp=torch.ones(batch, dim, device=rank), + sync_interval=1, + ), + DDPUnevenTestInput( + name="small_model_power_sgd_hook", + model=small_model, + hook=powerSGD.powerSGD_hook, + state=powerSGD.PowerSGDState( + process_group=None, + matrix_approximation_rank=1, + # Config so that powerSGD runs immediately instead of + # allreduce. + start_powerSGD_iter=1, + warm_start=False, + use_error_feedback=False, + ), + inp=torch.ones(batch, dim, device=rank), + sync_interval=1, + ), + ] + models_to_test.extend(models_with_hook) + + # Add resnet model if we have torchvision installed. + if HAS_TORCHVISION: + resnet_model = torchvision.models.resnet50() + models_to_test.append( + DDPUnevenTestInput( + name="resnet_model", + model=resnet_model, + inp=torch.ones(1, 3, 1000, 1000), + sync_interval=1, + ) + ) + + # Test with no_sync every 2, 3, 4, ... iterations. + models_with_sync = [] + for i, test_input in enumerate(models_to_test): + models_with_sync.append( + DDPUnevenTestInput( + name=test_input.name, + model=test_input.model, + inp=test_input.inp, + sync_interval=i + 2, + ) + ) + + throw_on_early_term_tests = [] + for test_input in models_to_test: + throw_on_early_term_tests.append( + DDPUnevenTestInput( + name=test_input.name, + model=test_input.model, + inp=test_input.inp, + sync_interval=test_input.sync_interval, + throw_on_early_termination=True, + ) + ) + + models_to_test.extend(models_with_sync) + models_to_test.extend(throw_on_early_term_tests) + + # 0 iteration tests for when one process does not train model at all, so + # we must shadow the broadcast calls made when rebuilding buckets. + baseline_num_iters = [0, 5] + iteration_offsets = [2, 3, 10] + num_uneven_ranks = [1] + if dist.get_world_size() > 2: + num_uneven_ranks.append(2) + iteration_mappings = [] + # Generate rank : num_iters mappings for various uneven input scenarios. + # This includes cases where rank 0 joins early and all other ranks join + # later, and scenarios where multiple ranks join early, but at different + # iterations, and later ranks join later. + for num_early_join_ranks in num_uneven_ranks: + for baseline_iter in baseline_num_iters: + for offset in iteration_offsets: + mapping = dict.fromkeys( + range(num_early_join_ranks), baseline_iter + ) + # if num_early_join_ranks > 1, ranks > 0 that will join early + # iterate offset//2 more times than rank 0, to test nodes + # depleting inputs at different times. + if num_early_join_ranks > 1: + for rank in mapping: + if rank > 0: + mapping[rank] += offset // 2 + mapping.update( + dict.fromkeys( + range(num_early_join_ranks, dist.get_world_size()), + baseline_iter + offset, + ) + ) + iteration_mappings.append(mapping) + + for test_case, iteration_mapping in itertools.product( + models_to_test, iteration_mappings + ): + if self.rank == 0: + print( + f"""Running test: {test_case.name} sync interval + {test_case.sync_interval} with iteration mapping + {iteration_mapping}""" + ) + self._run_uneven_inputs_test( + test_case, + iteration_mapping, + find_unused_params=("unused_params_model" in test_case.name), + ) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_uneven_input_join_disable(self): + # tests that if net.join() with enable=False is specified, DDP works as + # expected with even inputs. + torch.manual_seed(self.rank) + net = torch.nn.parallel.DistributedDataParallel( + torch.nn.Linear(1, 1).cuda(self.rank), device_ids=[self.rank] + ) + inp = torch.ones(1) * self.rank + n_iters = 5 + world_size = dist.get_world_size() + with net.join(enable=False): + for _ in range(n_iters): + # Clear grads + grad = net.module.weight.grad + if grad is not None: + grad.requires_grad_(False) + grad.zero_() + out = net(inp) + loss = out.sum() + loss.backward() + # Validate gradients to ensure that we divide by the correct + # world_size when join mode is disabled. + expected_grad = sum(i for i in range(world_size)) / world_size + self.assertEqual(net.module.weight.grad.item(), expected_grad) + + join_config = net._join_config + self.assertFalse(join_config.enable) + self.validate_net_equivalence(net) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_uneven_input_exception(self): + # Tests that exceptions during training are correctly propagated by the + # context manager. + error_str = "Intentional error" + + class ExceptionModule(nn.Module): + def __init__(self) -> None: + super().__init__() + self.param = nn.Parameter(torch.ones(1, requires_grad=True)) + + def forward(self, _): + raise ValueError(error_str) + + exception_module = ExceptionModule() + net = torch.nn.parallel.DistributedDataParallel( + exception_module.cuda(self.rank), device_ids=[self.rank] + ) + inp = torch.ones(1) + with self.assertRaisesRegex(ValueError, error_str): + with net.join(): + out = net(inp) + loss = out.sum() + loss.backward() + + def _test_broadcast_object_list(self, group=None): + gather_objects = create_collectives_object_test_list() + + # Only set device for NCCL backend since it must use GPUs. + # Case where rank != GPU device. + next_rank = (self.rank + 1) % int(self.world_size) + backend = os.environ["BACKEND"] + if backend == "nccl": + torch.cuda.set_device(next_rank) + + src_rank = 0 + # If GPU test, add object with GPU tensor + if backend == "nccl": + gather_objects.append(Foo(torch.randn(3, 3, device=0))) + + if IS_FBCODE: + # Create Tensor with > 2^31 Bytes storage requirements + # Only on FBCODE as testing OOMs in OSS + gather_objects.append(Foo(torch.randn(3, 178956971))) + objects = ( + gather_objects + if self.rank == src_rank + else [None for _ in gather_objects] + ) + + # Single object test with device specified. Backend="gloo", device=cpu + if backend != "nccl": + single_obj_list = [objects[0]] + if self.rank != src_rank: + self.assertNotEqual(single_obj_list[0], gather_objects[0]) + dist.broadcast_object_list( + single_obj_list, src=0, group=group, device=torch.device("cpu") + ) + self.assertEqual(single_obj_list[0], gather_objects[0]) + + # Single object test with device specified. Backend="gloo", device=current_device+1 + # The test is gated by the fact GPU count is the same as world size to avoid the case + # when backend is gloo but there is no multiple GPU devices. + if backend != "nccl" and torch.cuda.device_count() == int(self.world_size): + single_obj_list = [objects[0]] + if self.rank != src_rank: + self.assertNotEqual(single_obj_list[0], gather_objects[0]) + dist.broadcast_object_list( + single_obj_list, src=0, group=group, device=torch.device(next_rank) + ) + self.assertEqual(single_obj_list[0], gather_objects[0]) + + # Single object test with device specified. Backend="nccl", device=current_device+1 + if backend == "nccl" and torch.cuda.device_count() == int(self.world_size): + single_obj_list = [objects[0]] + if self.rank != src_rank: + self.assertNotEqual(single_obj_list[0], gather_objects[0]) + dist.broadcast_object_list( + single_obj_list, src=0, group=group, device=torch.device(next_rank) + ) + self.assertEqual(single_obj_list[0], gather_objects[0]) + + # Single object test: backward compatibility with device unspecified + single_obj_list = [objects[0]] + if self.rank != src_rank: + self.assertNotEqual(single_obj_list[0], gather_objects[0]) + dist.broadcast_object_list(single_obj_list, src=0, group=group) + self.assertEqual(single_obj_list[0], gather_objects[0]) + + # Multiple input objects test + if self.rank != src_rank: + self.assertNotEqual(objects, gather_objects) + dist.broadcast_object_list(objects, src=0, group=group) + self.assertEqual(objects, gather_objects) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @require_n_gpus_for_nccl_backend( + int(os.environ["WORLD_SIZE"]), os.environ["BACKEND"] + ) + @with_dist_debug_levels(levels=["DETAIL"]) + @unittest.skip( + "Test is failing, see https://github.com/pytorch/pytorch/pull/113620" + ) + def test_broadcast_object_list(self): + return self._test_broadcast_object_list() + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @require_n_gpus_for_nccl_backend( + int(os.environ["WORLD_SIZE"]), os.environ["BACKEND"] + ) + @with_dist_debug_levels(levels=["DETAIL"]) + def _test_broadcast_object_list_subgroup(self): + default = _get_default_group() + backend = dist.get_backend(default) + subgroup = dist.new_group(backend=backend) + return self._test_broadcast_object_list(subgroup) + + def _test_ddp_ignore_params_arg(self, static_graph=False): + class TestModel(nn.Module): + def __init__(self, rank): + self.rank = rank + super().__init__() + self.fc1 = nn.Linear(1, 1, bias=False) + # Proxy that will be materialized to another architecture later. + # (after wrapping model with DDP) + if self.rank == 0: + self.fc2 = nn.Linear(1, 10, bias=False) + else: + self.fc2 = nn.Linear(10, 10, bias=False) + + def forward(self, x): + x = self.fc1(x) + x = self.fc2(x) + return x + + device_id = self.rank + # Ensure the test works for both find_unused_parameter and broadcast_buffer settings. + for find_unused, broadcast_buffers in itertools.product( + [False, True], [False, True] + ): + model = TestModel(self.rank).float().to(device_id) + # Note that the model can have different shape buffers if we pass + # them in to be ignored as well. + model.fc2.register_buffer( + "ignore_buffer", torch.zeros(5 + self.rank, device=self.rank) + ) + proxy_params = list(model.fc2.parameters()) + model_fc2_name = next( + module_name + for module_name, module in model.named_modules() + if module is model.fc2 + ) + proxy_param_names = [ + f"{model_fc2_name}.{param_name}" + for param_name, _ in model.fc2.named_parameters() + ] + proxy_buffer_names = [ + f"{model_fc2_name}.{buf_name}" + for buf_name, _ in model.fc2.named_buffers() + ] + # Specify that we should ignore proxy_params since it will be + # materialized later. + torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model( + model, proxy_param_names + proxy_buffer_names + ) + ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[device_id], + find_unused_parameters=find_unused, + broadcast_buffers=broadcast_buffers, + static_graph=static_graph, + ) + # Materialize new params. These are not registered in DDP and thus + # don't have autograd hooks installed on them. + ddp.module.fc2 = nn.Linear(1, 1, bias=False).to(device_id) + + # local model with the new materialized parameters. + local_model = copy.deepcopy(ddp.module).cuda(self.rank) + + inp = torch.ones(1, dtype=torch.float).to(device_id) * (self.rank + 1) + for _ in range(6): + ddp(inp).sum().backward() + + local_model(inp).sum().backward() + # materialized param grad is not touched by DDP, so its grad should + # be the same as if running locally. + for materialized_param, local_param in zip( + ddp.module.fc2.parameters(), + local_model.fc2.parameters(), + strict=True, + ): + self.assertEqual(materialized_param.grad, local_param.grad) + + # fc1 parameter grad should still be different, due to allreduce. + for synced_param, local_param in zip( + ddp.module.fc1.parameters(), + local_model.fc1.parameters(), + strict=True, + ): + self.assertFalse(synced_param.grad == local_param.grad) + + # Proxy module grad should not be touched + for proxy_param in proxy_params: + self.assertTrue(proxy_param.grad is None) + + # Synchronize since we run multiple iterations of this test, to + # isolate failure hangs. + torch.cuda.synchronize(device=self.rank) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_ignore_params_arg(self): + self._test_ddp_ignore_params_arg(static_graph=False) + self._test_ddp_ignore_params_arg(static_graph=True) + + @with_dist_debug_levels(levels=["OFF", "INFO", "DETAIL"]) + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_unused_params_rebuild_buckets_exception(self): + class ToyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.net1 = nn.Linear(10, 10, bias=False) + self.net2 = nn.Linear(10, 10, bias=False) + + def forward(self, x): + return self.net1(x) + + ddp = torch.nn.parallel.DistributedDataParallel( + ToyModel().cuda(self.rank), device_ids=[self.rank] + ) + for i in range(2): + inp = torch.rand(1, 10) + if i > 0: + # On 2nd iteration, this will fail during rebuild_buckets, + # but we should report an error regarding unused parameters + # since that is the underlying root cause. + try: + ddp(inp).sum().backward() + except RuntimeError as e: + msg = str(e) + verify_ddp_error_logged(ddp, msg) + expected_strs = [ + ddp_prev_reduction_unfinished_str, + ddp_recommend_find_unused_params_str, + ddp_outputs_not_used_in_loss_str, + ] + # In debug mode, should show parameters that weren't reduced. + # Without debug mode, should show suggestion to use debug mode. + if dist.get_debug_level() == dist.DebugLevel.OFF: + expected_strs.append(ddp_suggest_debug_mode_str) + else: + unreduced_params = ", ".join(["net2.weight"]) + expected_strs.append( + f"did not receive grad for rank {self.rank}: {unreduced_params}" + ) + for s in expected_strs: + self.assertTrue(s in msg, f"Expected {s} to be in {msg}") + self.assertFalse(ddp_find_unused_params_enabled_str in msg) + else: + self.assertFalse( + True, "DDP unused parameters error not raised." + ) + else: + ddp(inp).sum().backward() + + dist.barrier() + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_shared_grad_acc_unused_params(self): + # When find_unused_parameters=True, ensure we mark unused parameters + # even if they share gradient accumulators. + class ToyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + # net1, bias, and net1.bias are all unused params. + self.net1 = nn.Linear(10, 5, bias=False) + self.bias = nn.Parameter(torch.zeros(5)) + # net1.bias and self.bias are names for the same underlying + # parameter, so they share the same grad acc. This caused + # the bug reported in https://github.com/pytorch/pytorch/issues/41324. + self.net1.bias = self.bias + self.net2 = nn.Linear(10, 5) + + def forward(self, x): + return self.net2(x).sum() + + torch.cuda.set_device(self.rank) + model = ToyModel().to(torch.cuda.current_device()) + for static in [True, False]: + ddp_model = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(model), + device_ids=[self.rank], + find_unused_parameters=True, + static_graph=static, + ) + inp = torch.randn(20, 10, device=self.rank) + for _ in range(6): + loss = ddp_model(inp) + # To test https://github.com/pytorch/pytorch/issues/61982 + loss /= 10 + loss.backward() + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_device(self): + expected_len = 2 + + class TensorWrapper: + __slots__ = ["t", "moved_to_gpu"] + + def __init__(self, t): + self.t = t + self.moved_to_gpu = False + + # Handlers for specific types of validation we want to do based on + # the input type. + + def tuple_and_list_validator(x): + self.assertTrue(len(x), expected_len) + self.assertEqual(1, len({t.device for t in x})) + self.assertEqual(x[0].device.index, self.rank) + return x[0] + x[1] + + def namedtuple_validator(x): + self.assertEqual(x._fields, EXPECTED_FIELDS) + self.assertEqual(x.a.device.index, x.b.device.index) + self.assertEqual(x.a.device.index, self.rank) + return x.a + x.b + + def custom_type_validator(x): + self.assertTrue(x.moved_to_gpu or (str(x.t.device) == "cpu")) + x.t = x.t.to(self.rank) + x.moved_to_gpu = True + return x.t + + def dict_validator(x): + self.assertTrue(EXPECTED_FIELDS[0] in x) + self.assertTrue(EXPECTED_FIELDS[1] in x) + self.assertEqual(1, len({t.device for t in x.values()})) + self.assertEqual(x[EXPECTED_FIELDS[0]].device.index, self.rank) + return x[EXPECTED_FIELDS[0]] + x[EXPECTED_FIELDS[1]] + + validators = { + TensorWrapper: custom_type_validator, + tuple: tuple_and_list_validator, + list: tuple_and_list_validator, + TestNamedTupleInput_0: namedtuple_validator, + TestNamedTupleInput_1: namedtuple_validator, + dict: dict_validator, + } + + class ToyModel(torch.nn.Module): + def __init__(self_): # noqa: B902 + super().__init__() + self_.lin = nn.Linear(10, 10, bias=False) + + def forward(self_, x, expected_type): # noqa: B902 + # Similar to scatter, the recursive to in the single-device + # case does not move tensors if they are in a custom type. + self.assertTrue(isinstance(x, expected_type)) + fwd_tensor = validators[expected_type](x) + return self_.lin(fwd_tensor) + + model = torch.nn.parallel.DistributedDataParallel( + ToyModel().to(self.rank), device_ids=[self.rank] + ) + + def train_iter(inp, input_type): + for _ in range(4): + out = model(inp, input_type) + out.sum().backward() + + # CPU tuple input, should be moved to the proper device before call + # to forward. + inp = tuple(torch.randn(10, 10) for _ in range(expected_len)) + train_iter(inp, tuple) + + # List CPU input, should be moved to proper device before call to + # forward. + inp = [torch.randn(10, 10) for _ in range(expected_len)] + train_iter(inp, list) + # Custom type containing tensor. The type is maintained, but the + # device is not propagated (which is what happens with scatter too) + inp = TensorWrapper(torch.randn(10, 10)) + train_iter(inp, TensorWrapper) + # NamedTuple input. The type should be maintained and tensor inputs + # should be moved to the correct device as in scatter. + batch = 5 + dim = 10 + a = torch.rand(batch, dim) + b = torch.rand(batch, dim) + + inp = TestNamedTupleInput_0(a, b) + train_iter(inp, type(inp)) + + inp = TestNamedTupleInput_1(a, b) + train_iter(inp, type(inp)) + + # dictionary input. + inp = { + EXPECTED_FIELDS[0]: a, + EXPECTED_FIELDS[1]: b, + } + train_iter(inp, type(inp)) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_namedtuple(self): + batch = 5 + dim = 10 + + a = torch.rand(batch, dim, device=self.rank) + b = torch.rand(batch, dim, device=self.rank) + + class NamedTupleModule(torch.nn.Module): + def __init__(self_): # noqa: B902 + super().__init__() + self_.lin = nn.Linear(10, 1) + + def forward(self_, input, expected_type): # noqa: B902 + # Without NamedTuple support, this would be of type tuple. + self.assertTrue( + isinstance(input, expected_type), + f"Expected type {expected_type} but got {type(input)}", + ) + self.assertEqual(input._fields, EXPECTED_FIELDS) + self.assertEqual(a, input.a) + self.assertEqual(b, input.b) + return self_.lin(torch.mul(input.a, input.b)) + + model = torch.nn.parallel.DistributedDataParallel( + NamedTupleModule().cuda(self.rank), device_ids=[self.rank] + ) + inp = TestNamedTupleInput_0(a, b) + # The following would fail if DDP does not propagate NamedTuples correctly. + model(inp, type(inp)) + + inp = TestNamedTupleInput_1(a, b) + model(inp, type(inp)) + + @require_backend_is_available({"gloo"}) + def test_grads_same_across_ranks_with_no_sync(self): + _group, _group_id, rank = self._init_global_test() + world_size = dist.get_world_size() + if world_size < 2: + self.skipTest("This test requires at least two ranks.") + + class SimpleConditionalModel(nn.Module): + # if rank is 0, uses nn1 on the first pass and nn2 on the second pass. + # else, uses nn3 on the first pass and nn4 on the second pass. + + def __init__(self, rank): + super().__init__() + + self.rank = rank + self.nn1 = nn.Linear(1, 1) + self.nn2 = nn.Linear(1, 1) + self.nn3 = nn.Linear(1, 1) + self.nn4 = nn.Linear(1, 1) + self.state = 0 + + def forward(self, input): + if self.state == 0: + self.state = 1 + if self.rank == 0: + return self.nn1(input) + else: + return self.nn3(input) + else: + self.state = 0 + if self.rank == 0: + return self.nn2(input) + else: + return self.nn4(input) + + model = torch.nn.parallel.DistributedDataParallel( + SimpleConditionalModel(rank), find_unused_parameters=True + ) + mse_loss = nn.MSELoss() + grad_accumulation = 2 + + for microbatch_idx in range(grad_accumulation): + if microbatch_idx < grad_accumulation - 1: + context = model.no_sync + else: + context = nullcontext + + with context(): + input = torch.rand((1,)) + output = model.forward(input) + target = torch.rand((1,)) + + loss = mse_loss(output, target) + loss.backward() + + self.assertTrue( + not any(p.grad is None for p in model.parameters()), + "Gradients can't be None for any model parameter.", + ) + grads = torch.cat([p.grad.view(-1) for p in model.parameters()]) + + # Gather all gradients to rank 0. + if rank == 0: + gathered_grads = [torch.zeros_like(grads) for _ in range(world_size)] + else: + gathered_grads = [] + + dist.gather(grads, gather_list=gathered_grads, dst=0) + if rank == 0: + for g in gathered_grads[1:]: + self.assertTrue( + torch.allclose(gathered_grads[0], g), + "Gradients are not the same for all ranks.", + ) + + @with_dist_debug_levels(levels=["OFF", "INFO", "DETAIL"]) + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_control_flow_same_across_ranks(self): + # Control flow that is the same across ranks. + batch = 20 + dim = 10 + + world_size = dist.get_world_size() + torch.cuda.set_device(self.rank) + model = torch.nn.parallel.DistributedDataParallel( + ControlFlowToyModel().cuda(self.rank), + device_ids=[self.rank], + find_unused_parameters=True, + ) + random_input = torch.randn(batch, dim, device=self.rank) + ones_input = torch.ones(batch, dim, device=self.rank) + for i in range(6): + if i % 2 == 0: + out = model(random_input) + else: + out = model(ones_input) + loss = out.sum() + loss.backward() + # On even iterations, 2nd param goes unused, on odd iterations, + # it is used. + local_used_map = model.reducer._get_local_used_map() + if i % 2 == 0: + expected = torch.tensor( + [world_size, 0], device=self.rank, dtype=torch.int32 + ) + else: + expected = torch.tensor( + [world_size, world_size], device=self.rank, dtype=torch.int32 + ) + + # Validate parameter usage. + variable_usage_tensor = local_used_map + self.assertEqual(variable_usage_tensor, expected) + + # Validate appropriate error message when DDP is used with + # find_unused_parameters=False. + model = torch.nn.parallel.DistributedDataParallel( + ControlFlowToyModel().cuda(self.rank), + device_ids=[self.rank], + find_unused_parameters=False, + ) + for i in range(2): + if i == 0: + loss = model(random_input).sum() + loss.backward() + else: + try: + loss = model(random_input).sum() + loss.backward() + except RuntimeError as e: + msg = str(e) + verify_ddp_error_logged(model, msg) + # 2nd linear layer is unused + unused_param_index = 1 + expected_strs = [ + ddp_prev_reduction_unfinished_str, + ddp_recommend_find_unused_params_str, + ddp_outputs_not_used_in_loss_str, + f"Parameter indices which did not receive grad for rank {self.rank}: {unused_param_index}", + ] + # In debug mode, should show parameters that weren't reduced. + # Without debug mode, should show suggestion to use debug mode. + if dist.get_debug_level() == dist.DebugLevel.OFF: + expected_strs.append(ddp_suggest_debug_mode_str) + else: + unreduced_params = ", ".join(["lin2.weight"]) + expected_strs.append( + f"did not receive grad for rank {self.rank}: {unreduced_params}" + ) + for s in expected_strs: + self.assertTrue(s in msg, f"Expected {s} to be in {msg}") + self.assertFalse(ddp_find_unused_params_enabled_str in msg) + else: + self.assertFalse(True, "DDP error not raised") + + dist.barrier() + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_invalid_static_graph(self): + torch.cuda.set_device(self.rank) + model = torch.nn.parallel.DistributedDataParallel( + ControlFlowToyModel().cuda(self.rank), + device_ids=[self.rank], + static_graph=True, + ) + random_input = torch.randn(20, 10, device=self.rank) + ones_input = torch.ones(20, 10, device=self.rank) + # unused parameter in the first iteration got used + # in second iteration. + expected_err = "Your training graph has changed in this iteration" + with self.assertRaisesRegex(RuntimeError, expected_err): + for i in range(2): + if i % 2 == 0: + out = model(random_input) + else: + out = model(ones_input) + loss = out.sum() + loss.backward() + + verify_ddp_error_logged(model, expected_err) + + # used parameter in the first iteration got unused + # in second iteration. + with self.assertRaisesRegex( + RuntimeError, + "Expected to have finished reduction in the prior iteration " + "before starting a new one. This error indicates that your " + "training graph has changed in this iteration, " + "e.g., one parameter is used in first iteration, " + "but then got unused in the second iteration. " + "this is not compatible with static_graph set to True.\n" + "Parameter indices which did not receive grad for", + ): + for i in range(2): + if i % 2 != 0: + out = model(random_input) + else: + out = model(ones_input) + loss = out.sum() + loss.backward() + + verify_ddp_error_logged(model, "Expected to have finished reduction") + + @with_dist_debug_levels(levels=["OFF", "INFO", "DETAIL"]) + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_control_flow_different_across_ranks(self): + # Control flow that is different across ranks. + batch = 20 + dim = 10 + + class ToyModel(nn.Module): + def __init__(self, rank): + super().__init__() + self.lin1 = nn.Linear(10, 10, bias=False) + self.lin2 = nn.Linear(10, 10, bias=False) + self.rank = rank + + def forward(self, x): + # Control-flow that is rank and input dependent for the + # model. + use_second_layer = ( + torch.equal(x, torch.ones(batch, dim, device=x.device)) + and self.rank == 1 + ) + + if use_second_layer: + return self.lin2(F.relu(self.lin1(x))) + else: + return F.relu(self.lin1(x)) + + world_size = dist.get_world_size() + torch.cuda.set_device(self.rank) + model = torch.nn.parallel.DistributedDataParallel( + ToyModel(self.rank).cuda(self.rank), + device_ids=[self.rank], + find_unused_parameters=True, + ) + random_input = torch.randn(batch, dim, device=self.rank) + ones_input = torch.ones(batch, dim, device=self.rank) + for i in range(6): + if i % 2 == 0: + out = model(random_input) + else: + out = model(ones_input) + loss = out.sum() + loss.backward() + # On even iterations, 2nd param goes unused, on odd iterations, + # it is used only on rank 1. + local_used_map = model.reducer._get_local_used_map() + + if i % 2 == 0: + expected = torch.tensor( + [world_size, 0], device=self.rank, dtype=torch.int32 + ) + else: + expected = torch.tensor( + [world_size, 1], device=self.rank, dtype=torch.int32 + ) + + variable_usage_tensor = local_used_map + # Validate parameter usage. On odd iterations, 2nd param is only + # used on rank 1. + self.assertEqual(variable_usage_tensor, expected) + + # Validate appropriate error message when DDP is used with + # find_unused_parameters=False. + model = torch.nn.parallel.DistributedDataParallel( + ToyModel(self.rank).cuda(self.rank), + device_ids=[self.rank], + find_unused_parameters=False, + ) + for i in range(2): + if i == 0: + loss = model(random_input).sum() + loss.backward() + else: + try: + loss = model(random_input).sum() + loss.backward() + except RuntimeError as e: + msg = str(e) + verify_ddp_error_logged(model, msg) + unused_param_index = 1 + expected_strs = [ + ddp_prev_reduction_unfinished_str, + ddp_recommend_find_unused_params_str, + ddp_outputs_not_used_in_loss_str, + f"Parameter indices which did not receive grad for rank {self.rank}: {unused_param_index}", + ] + # In debug mode, should show parameters that weren't reduced. + # Without debug mode, should show suggestion to use debug mode. + if dist.get_debug_level() == dist.DebugLevel.OFF: + expected_strs.append(ddp_suggest_debug_mode_str) + else: + unreduced_params = ", ".join(["lin2.weight"]) + expected_strs.append( + f"did not receive grad for rank {self.rank}: {unreduced_params}" + ) + for s in expected_strs: + self.assertTrue(s in msg, f"Expected {s} to be in {msg}") + self.assertFalse(ddp_find_unused_params_enabled_str in msg) + else: + self.assertFalse(True, "DDP error not raised") + + dist.barrier() + + @require_backend_is_available({"gloo"}) + def test_scatter_object_list(self): + src_rank = 0 + collectives_object_test_list = create_collectives_object_test_list() + scatter_list = ( + collectives_object_test_list + if self.rank == src_rank + else [None for _ in collectives_object_test_list] + ) + world_size = dist.get_world_size() + scatter_list = scatter_list[:world_size] + i = 0 + while len(scatter_list) < world_size: + scatter_list.append(scatter_list[i]) + i += 1 + + output_obj_list = [None] + dist.scatter_object_list(output_obj_list, scatter_list, src=src_rank) + self.assertEqual( + output_obj_list[0], + collectives_object_test_list[ + self.rank % len(collectives_object_test_list) + ], + ) + # Ensure errors are raised upon incorrect arguments. + with self.assertRaisesRegex( + ValueError, + "Expected argument scatter_object_output_list to be a list of size at least 1.", + ): + dist.scatter_object_list([], scatter_list, src=src_rank) + + def _generate_sparse_tensors_for_bucket_assignment_test(self): + tensors = [ + torch.empty([50], dtype=torch.float), + torch.empty([25], dtype=torch.double), + torch.empty([50], dtype=torch.float), + torch.empty([25], dtype=torch.double), + torch.empty([50], dtype=torch.float), + torch.empty([25], dtype=torch.double), + ] + + tensors_sparse = [t.to_sparse() for t in tensors] + return tensors_sparse + + def _test_compute_bucket_assignment_by_size(self, use_logger): + group_gloo = dist.new_group( + timeout=timedelta(seconds=60), backend=dist.Backend.GLOO + ) + # Set TORCH_NCCL_BLOCKING_WAIT and use a new NCCL group to improve test + # determinism. + os.environ["TORCH_NCCL_BLOCKING_WAIT"] = "1" + group_to_use = dist.new_group( + backend=dist.get_backend(), timeout=timedelta(seconds=5) + ) + torch.cuda.set_device(self.rank) + + # Create a valid model. The constructor initializes the logger that we use later. + # We never actually use the rest of the model - we only need its logger. + net = EmbeddingNetDifferentParams(0) + net = torch.nn.parallel.DistributedDataParallel( + net.to(self.rank), + device_ids=[self.rank], + process_group=group_to_use, + ) + + # if we don't pass a logger then we can only check that an exception was thrown. + expected_err = "No support for sparse tensors." + with self.assertRaisesRegex(RuntimeError, expected_err): + tensors_sparse = ( + self._generate_sparse_tensors_for_bucket_assignment_test() + ) + if use_logger: + dist._compute_bucket_assignment_by_size( + tensors_sparse, [400], logger=net.logger + ) + else: + dist._compute_bucket_assignment_by_size(tensors_sparse, [400]) + if use_logger: + verify_ddp_error_logged(net, expected_err) + + # Perform gloo-based barrier to ensure one rank doesn't exit test + # early which causes failure with Barrier.sync. + dist.barrier(group_gloo) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_compute_bucket_assignment_by_size_sparse_error_without_logger(self): + self._test_compute_bucket_assignment_by_size(use_logger=False) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_compute_bucket_assignment_by_size_sparse_error_with_logger(self): + self._test_compute_bucket_assignment_by_size(use_logger=True) + + def _test_verify_model_across_rank(self, use_logger): + group_gloo = dist.new_group( + timeout=timedelta(seconds=60), backend=dist.Backend.GLOO + ) + group_to_use = dist.new_group( + backend=dist.get_backend(), timeout=timedelta(seconds=5) + ) + torch.cuda.set_device(self.rank) + + # Create a valid model. The constructor initializes the logger that we use later. + net = EmbeddingNetDifferentParams(0) + net = torch.nn.parallel.DistributedDataParallel( + net.to(self.rank), + device_ids=[self.rank], + process_group=group_to_use, + ) + + # Modify the model so that the number of parameters are different for each rank. + # This will cause a RuntimeError to be thrown below in _verify_param_shape_across_processes, + # so we can check if the correct error is thrown and is logged. + # We can't do this in the constructor above otherwise the logger will + # not be properly initialized. + net.module.lin = nn.Linear(100 if self.rank == 0 else 10, 1) + + # if we pass a logger we can verify that it was logged + caught = 0 + try: + if use_logger: + _verify_param_shape_across_processes( + net.process_group, list(net.parameters()), net.logger + ) + else: + _verify_param_shape_across_processes( + net.process_group, list(net.parameters()) + ) + except Exception: + caught = 1 + + # As long as there is one rank catching the exception + t = torch.Tensor([caught]) + dist.all_reduce(t, group=group_gloo) + self.assertGreater(t, 0) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc" and IS_SANDCASTLE, "Skipped internally" + ) + @skip_if_lt_x_gpu(2) + def test_verify_model_across_rank_with_logger(self): + self._test_verify_model_across_rank(use_logger=True) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc" and IS_SANDCASTLE, "Skipped internally" + ) + @skip_if_lt_x_gpu(2) + def test_verify_model_across_rank_without_logger(self): + self._test_verify_model_across_rank(use_logger=False) + + def _run_test_ddp_model_with_diff_params(self, net, ddp_group, group_gloo): + caught = 0 + try: + net = torch.nn.parallel.DistributedDataParallel( + net.to(self.rank), device_ids=[self.rank], process_group=ddp_group + ) + except Exception: + caught = 1 + + # As long as there is one rank catching the exception + t = torch.Tensor([caught]) + dist.all_reduce(t, group=group_gloo) + self.assertGreater(t, 0) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc" and IS_SANDCASTLE, "Skipped internally" + ) + @skip_if_lt_x_gpu(2) + def test_ddp_model_diff_shape_across_ranks(self): + group_gloo = dist.new_group( + timeout=timedelta(seconds=60), backend=dist.Backend.GLOO + ) + group_to_use = dist.new_group( + backend=dist.get_backend(), timeout=timedelta(seconds=10) + ) + torch.cuda.set_device(self.rank) + # Creates network with different sized embedding table on different + # ranks. This should throw an error during DDP init. + net = EmbeddingNetDifferentParams(self.rank) + self._run_test_ddp_model_with_diff_params(net, group_to_use, group_gloo) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_but_pass_in_sandcastle_if( + BACKEND == "ucc" and IS_SANDCASTLE, "Skipped internally" + ) + @skip_if_lt_x_gpu(2) + def test_ddp_model_diff_num_params_across_ranks(self): + group_gloo = dist.new_group( + timeout=timedelta(seconds=60), backend=dist.Backend.GLOO + ) + group_to_use = dist.new_group( + backend=dist.get_backend(), timeout=timedelta(seconds=10) + ) + torch.cuda.set_device(self.rank) + + # Creates network with diff # of param across ranks, reducer should + # recognize this and throw appropriate error. + net = EmbeddingNetDifferentParams( + self.rank, diff_num_params=(self.rank == 1) + ) + + self._run_test_ddp_model_with_diff_params( + net, + group_to_use, + group_gloo, + ) + + def _test_output_unused_in_loss(self, module_cls, gradient_as_bucket_view): + model = module_cls() + local_net = copy.deepcopy(model) + net = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(model).cuda(self.rank), + device_ids=[self.rank], + find_unused_parameters=True, + ) + + # Tests that certain parameters not getting gradient since the + # output is unused in loss computation is supported. Specifically, + # checks that the grads remain unchanged and are the same as local + # training. + inp = torch.randn(10, 10) + + # Ensure that if a param is not used in loss computation, its + # gradient is untouched, i.e. if it is None before it is None after, + # not zero. + if module_cls == DictOutputModule: + a, b = local_net(inp)["predictions"] + a_dist, b_dist = net(inp)["predictions"] + else: + a, b = local_net(inp) + a_dist, b_dist = net(inp) + + loss_dist = b_dist.sum() + loss_dist.backward() + + # Ensure that gradient corresponding to parameter "a" was not + # touched, i.e. it is None and matches the local grad. + if module_cls == DictOutputModule: + self.assertTrue(net.module.module.a.weight.grad is None) + self.assertEqual( + net.module.module.a.weight.grad, local_net.module.a.weight.grad + ) + else: + self.assertTrue(net.module.a.weight.grad is None) + self.assertEqual(net.module.a.weight.grad, local_net.a.weight.grad) + + saved_a_local_grad = None + saved_a_dist_grad = None + net.zero_grad() + local_net.zero_grad() + for i in range(6): + if module_cls == DictOutputModule: + a, b = local_net(inp)["predictions"] + a_dist, b_dist = net(inp)["predictions"] + else: + a, b = local_net(inp) + a_dist, b_dist = net(inp) + if i < 2: + # Use both params in loss computation. Later, "a" will go + # unused and we check to ensure DDP supports this and + # gradients remain the same as local training. + t = a @ b + t_dist = a_dist @ b_dist + loss = t.sum() + loss_dist = t_dist.sum() + else: + # Model output "a" unused in loss. + loss = b.sum() + loss_dist = b_dist.sum() + loss.backward() + loss_dist.backward() + if i == 1: + # Save grads to compare with them in next iterations. + if module_cls == DictOutputModule: + saved_a_local_grad = local_net.module.a.weight.grad + saved_a_dist_grad = net.module.module.a.weight.grad + else: + saved_a_local_grad = local_net.a.weight.grad + saved_a_dist_grad = net.module.a.weight.grad + self.assertEqual(saved_a_local_grad, saved_a_dist_grad) + elif i >= 2: + # parameter "a" of both models should be the same and not change + if module_cls == DictOutputModule: + self.assertEqual( + net.module.module.a.weight.grad, saved_a_dist_grad + ) + self.assertEqual( + local_net.module.a.weight.grad, saved_a_local_grad + ) + else: + self.assertEqual(net.module.a.weight.grad, saved_a_dist_grad) + self.assertEqual(local_net.a.weight.grad, saved_a_local_grad) + + # Verify grads are the same + for local_param, dist_param in zip( + local_net.parameters(), net.parameters(), strict=True + ): + local_grad = local_param.grad + dist_grad = dist_param.grad + self.assertEqual(local_grad, dist_grad) + + dist.barrier() + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_lt_x_gpu(2) + def test_output_unused_in_loss_tuple_module(self): + module_cls = UnusedParamTwoLinLayerNet + for grad_as_bucket_view in [True, False]: + self._test_output_unused_in_loss(module_cls, grad_as_bucket_view) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_lt_x_gpu(2) + def test_output_unused_in_loss_dict_module(self): + module_cls = DictOutputModule + for grad_as_bucket_view in [True, False]: + self._test_output_unused_in_loss(module_cls, grad_as_bucket_view) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_lt_x_gpu(2) + def test_undefined_grad_parity_unused_parameters(self): + # TODO: enable this for general training use cases: + # https://github.com/pytorch/pytorch/issues/58511. + x = torch.ones(1, 2).to(self.rank) + net = Net().to(self.rank) + local_net = copy.deepcopy(net) + net = torch.nn.parallel.DistributedDataParallel( + net, + device_ids=[self.rank], + find_unused_parameters=True, + ) + out = net(x).sum() + local_out = local_net(x).sum() + # Simulates undefined gradients. + torch._C._functions.UndefinedGrad()(out).backward() + torch._C._functions.UndefinedGrad()(local_out).backward() + for (dist_param_name, dist_param), (local_param_name, local_param) in zip( + net.named_parameters(), local_net.named_parameters(), strict=True + ): + dist_grad = dist_param.grad + local_grad = local_param.grad + self.assertEqual( + dist_grad, + local_grad, + f"""DDP param {dist_param_name} with grad {dist_grad} + does not match local param {local_param_name} with grad + {local_grad}""", + ) + + def _test_different_graph_across_ranks( + self, find_unused_parameters=False, static_graph=False + ): + class ToyModel(nn.Module): + def __init__(self, rank): + super().__init__() + self.lin1 = nn.Linear(10, 10, bias=False) + self.lin2 = nn.Linear(10, 10, bias=False) + self.rank = rank + + def forward(self, x): + if self.rank == 0: + return self.lin2(F.relu(self.lin1(x))) + else: + return F.relu(self.lin1(x)) + + torch.manual_seed(31415) + torch.cuda.set_device(self.rank) + model = ToyModel(self.rank).cuda(self.rank) + ddp_model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + find_unused_parameters=find_unused_parameters, + gradient_as_bucket_view=True, + static_graph=static_graph, + ) + random_input = torch.randn(20, 10, device=self.rank) + for _ in range(10): + out = ddp_model(random_input) + loss = out.sum() + loss.backward() + return ddp_model + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_different_graph_across_ranks(self): + base_model = self._test_different_graph_across_ranks( + find_unused_parameters=True + ) + self.assertFalse( + base_model._get_ddp_logging_data().get("has_rebuilt_buckets", 0) + ) + static_model = self._test_different_graph_across_ranks(static_graph=True) + self.assertTrue( + static_model._get_ddp_logging_data().get("has_rebuilt_buckets", 0) + ) + for i, j in zip( + base_model.parameters(), static_model.parameters(), strict=True + ): + self.assertEqual(i, j) + + @require_backend_is_available({"gloo"}) + @skip_but_pass_in_sandcastle_if( + IS_MACOS or IS_WINDOWS, + "MacOS uses uv transport which does not have as robust error handling as tcp transport", + ) + def test_monitored_barrier_gloo(self): + tensors = [torch.ones(10) * self.rank] + # Kick off some allreduce work on all ranks + for _ in range(10): + dist.all_reduce(torch.cat(tensors)) + # Run monitored barrier and ensure it passes + timeout = timedelta(seconds=2) + dist.monitored_barrier(timeout=timeout) + # Check monitored_barrier success with wait_all_ranks=True + for _ in range(10): + dist.all_reduce(torch.cat(tensors)) + dist.monitored_barrier(timeout=timeout, wait_all_ranks=True) + # All ranks besides 1 call into barrier, rank 0 should report failure + # while others report gloo error. + failed_rank = 1 + src_rank = 0 + if self.rank == src_rank: + with self.assertRaisesRegex( + RuntimeError, f"Rank {failed_rank} failed to pass monitoredBarrier" + ): + dist.monitored_barrier(timeout=timeout) + elif self.rank != failed_rank: + # Other ranks should not pass barrier since rank 0 failed. + err_regex = ( + f"Rank {self.rank} successfully reached monitoredBarrier," + f" but received errors while waiting for send/recv from rank" + f" {src_rank}" + ) + with self.assertRaisesRegex(RuntimeError, err_regex): + dist.monitored_barrier(timeout=timeout) + + # We need a barrier since otherwise failed_rank exits too early + # and cause a timeout. + self._barrier(timeout=30) + + @require_backend_is_available({"gloo"}) + def test_monitored_barrier_gloo_subgroup(self): + # Tests that monitored_barrier works as expected on non-default + # process groups. + failed_rank = 1 + timeout = 0.1 + subgroup = dist.new_group(ranks=[0, 1]) + + if self.rank == failed_rank: + return + + if self.rank == 0: + with self.assertRaisesRegex( + RuntimeError, f"Rank {failed_rank} failed to pass monitoredBarrier" + ): + dist.monitored_barrier(subgroup, timeout) + else: + # Other ranks call into monitored_barrier, but this should be a + # noop because they are not part of the subgroup. Verify that + # there are no errors here. + dist.monitored_barrier(subgroup, timeout) + + def _test_monitored_barrier_allreduce_hang(self, wait_all_ranks): + # tests expected behavior when nonzero rank hangs. + nccl_pg = dist.new_group( + ranks=list(range(int(self.world_size))), + # provide sufficient timeout so communicators + # can be initialized in ctor. + timeout=timedelta(seconds=15), + backend=dist.Backend.NCCL, + ) + gloo_pg = dist.new_group( + ranks=list(range(int(self.world_size))), + backend=dist.Backend.GLOO, + ) + tensors = [torch.ones(10, device=self.rank) * self.rank] + # Let all ranks call allreduce first to set up communicators etc. + # Directly simulating error here will run into store issue described + # in https://github.com/pytorch/pytorch/issues/54524. + nccl_pg.allreduce(tensors).wait(timedelta(seconds=5)) + # All ranks besides 0 call into allreduce. This is to simulate a + # desync across the world, where some ranks call into + # monitored_barrier() and others are stuck in collective comm. In + # practice, we don't need TORCH_NCCL_BLOCKING_WAIT, but we use it in this + # test to ensure it exits cleanly. + if self.rank != 0: + # Can get different errors here depending on whether gloo-based + # wrapper PG is enabled or not, since with wrapper pg, it will + # fail in a collective synchronization check and not actually + # call into the nccl pg. + if dist.get_debug_level() == dist.DebugLevel.DETAIL: + err_regex = "Timed out waiting" + else: + err_regex = "caught collective operation timeout" + with self.assertRaisesRegex(RuntimeError, err_regex): + nccl_pg.allreduce(tensors).wait(timedelta(seconds=0.1)) + else: + # Rank 0 should report first (in order) timed out rank or all ranks + # depending on wait_all_ranks flag passed into monitored_barrier. + if wait_all_ranks: + rank_str = ", ".join( + [str(i) for i in range(1, int(self.world_size))] + ) + err_regex = f"Ranks {rank_str} failed to pass monitoredBarrier" + else: + expected_first_fail_rank = 1 + err_regex = f"Rank {expected_first_fail_rank} failed to pass monitoredBarrier" + monitored_barrier_timeout_seconds = timedelta(seconds=0.1) + with self.assertRaisesRegex(RuntimeError, err_regex): + gloo_pg.monitored_barrier( + monitored_barrier_timeout_seconds, wait_all_ranks=wait_all_ranks + ) + + self._barrier(timeout=30) + + @with_nccl_blocking_wait + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_monitored_barrier_allreduce_hang(self): + # tests expected behavior when nonzero rank hangs and we want to + # report first timed out rank. + self._test_monitored_barrier_allreduce_hang(wait_all_ranks=False) + + @with_nccl_blocking_wait + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + def test_monitored_barrier_allreduce_hang_wait_all_ranks(self): + # Need to disable TORCH_NCCL_DUMP_ON_TIMEOUT otherwise this test times out + os.environ["TORCH_NCCL_DUMP_ON_TIMEOUT"] = "0" + # tests expected behavior when nonzero rank hangs and we want to + # report all timed out ranks. + self._test_monitored_barrier_allreduce_hang(wait_all_ranks=True) + + @require_backend_is_available({"gloo"}) + def test_monitored_barrier_gloo_rank_0_timeout(self): + # tests error when rank 0 exhausts its given timeout. + process_group = dist.new_group(ranks=list(range(int(self.world_size)))) + timeout = timedelta(seconds=0) + if self.rank == 0: + with self.assertRaisesRegex( + RuntimeError, f"Rank {self.rank} timed out in monitoredBarrier" + ): + process_group.monitored_barrier(timeout) + + @require_backend_is_available({"gloo"}) + @skip_if_small_worldsize + @skip_but_pass_in_sandcastle_if( + IS_MACOS or IS_WINDOWS, + "MacOS uses uv transport which does not have as robust error handling as tcp transport", + ) + def test_monitored_barrier_failure_order(self): + # Ensure that the first (in sorted order) rank is reported when + # multiple ranks fail to pass the monitored_barrier. + # TODO(#54879): Provide ability to wait and report all failed ranks + expected_first_failed_rank = 2 + timeout = timedelta(seconds=2) + src_rank = 0 + if self.rank == src_rank: + with self.assertRaisesRegex( + RuntimeError, f"Rank {expected_first_failed_rank}" + ): + dist.monitored_barrier(timeout=timeout) + elif self.rank == 1: + err_regex = ( + f"Rank {self.rank} successfully reached monitoredBarrier," + f" but received errors while waiting for send/recv from rank" + f" {src_rank}" + ) + with self.assertRaisesRegex(RuntimeError, err_regex): + dist.monitored_barrier(timeout=timeout) + + @require_backend_is_available({"gloo"}) + @skip_if_small_worldsize + def test_monitored_barrier_wait_all_ranks(self): + # Tests simple case where > 1 rank does not call into monitored + # barrier and verifies all ranks are reported by rank 0. + if self.rank == 0: + timeout = timedelta(seconds=0.1) + rank_str = ", ".join([str(i) for i in range(1, int(self.world_size))]) + err_regex = f"Ranks {rank_str} failed to pass monitoredBarrier" + with self.assertRaisesRegex(RuntimeError, err_regex): + dist.monitored_barrier(timeout=timeout, wait_all_ranks=True) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @with_dist_debug_levels(levels=["INFO"]) + @skip_if_lt_x_gpu(2) + def test_ddp_build_debug_param_to_name_mapping(self): + model = TwoLinLayerNet() + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(self.rank), + device_ids=[self.rank], + ) + expected_mapping = {0: "a.weight", 1: "b.weight"} + net_params, _ = net._build_params_for_reducer() + param_to_name_mapping = net._build_debug_param_to_name_mapping(net_params) + self.assertDictEqual(expected_mapping, param_to_name_mapping) + + # Test when DDP is used with ignored parameters. + model = TwoLinLayerNet() + # Parameters to ignore are in the format {module_name}.{param_name} + params_to_ignore = ["a.weight"] + torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model( + model, params_to_ignore + ) + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(self.rank), + device_ids=[self.rank], + ) + expected_mapping = {0: "b.weight"} + net_params, _ = net._build_params_for_reducer() + param_to_name_mapping = net._build_debug_param_to_name_mapping(net_params) + self.assertDictEqual(expected_mapping, param_to_name_mapping) + + # Test errors are raised when DDP and module parameters mismatch. + # This generally indicates a bug with DDP and is not expected to + # happen in user applications. + model = TwoLinLayerNet() + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(self.rank), + device_ids=[self.rank], + ) + net_params, _ = net._build_params_for_reducer() + if self.rank == 0: + print(type(net_params[0])) + + net_params.extend( + [ + torch.nn.Parameter(torch.ones(1)), + torch.nn.Parameter(torch.ones(1)), + ] + ) + + with self.assertRaisesRegex(ValueError, "Expected param to name mapping"): + net._build_debug_param_to_name_mapping(net_params) + + net_params = net_params[:-3] + with self.assertRaisesRegex(ValueError, "Param with name"): + net._build_debug_param_to_name_mapping(net_params) + + net_params.extend( + [ + torch.nn.Parameter(torch.ones(1)), + torch.nn.Parameter(torch.ones(1)), + ] + ) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @with_dist_debug_levels(levels=["INFO"]) + @skip_if_lt_x_gpu(2) + def test_ddp_build_debug_param_to_name_mapping_requires_grad(self): + class Net(nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin = nn.Linear(10, 10) + # Is not tracked by DDP and should not show up in param to + # name mapping. + self.lin.bias.requires_grad_(False) + + def forward(self, x): + return self.lin(x) + + model = Net() + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(self.rank), device_ids=[self.rank] + ) + expected_mapping = { + 0: "lin.weight", + } + net_params, _ = net._build_params_for_reducer() + param_to_name_mapping = net._build_debug_param_to_name_mapping(net_params) + self.assertEqual(param_to_name_mapping, expected_mapping) + + def _test_ddp_multiple_nested_unused_params_error(self, ignore_sparse): + debug_mode_off = dist.get_debug_level() == dist.DebugLevel.OFF + + class SubModule(nn.Module): + def __init__(self) -> None: + super().__init__() + self.embedding_net = EmbeddingNetDifferentParams(0) + self.lin = TwoLinLayerNet() + self.bn = BatchNormNet() + self.lin_layer = nn.Linear(4, 10, bias=False) + + def forward(self, x): + x = self.bn(x) + x = self.lin_layer(x) + x = self.lin.a(x) # self.lin.b param unused + # EmbeddingNetDifferentParams entirely unused: self.embedding_net.embedding and + # self.embedding_net.lin unused. + return x + + class MyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.sub_module = SubModule() + + def forward(self, x): + return self.sub_module(x) + + model = MyModel() + sparse_embedding_fqns = [] + if ignore_sparse: + for module_name, module in model.named_modules(): + if module == model.sub_module.embedding_net.embedding: + for parameter_name, _param in module.named_parameters( + recurse=False + ): + fqn = f"{module_name}.{parameter_name}" + sparse_embedding_fqns.append(fqn) + + torch.nn.parallel.DistributedDataParallel._set_params_and_buffers_to_ignore_for_model( + model, sparse_embedding_fqns + ) + unused_modules = [ + model.sub_module.embedding_net.lin, + model.sub_module.lin.b, + ] + else: + unused_modules = list(model.sub_module.embedding_net.modules()) + [ + model.sub_module.lin.b, + ] + + expected_unused_param_fqns = [] + used_param_fqns = [] # Validate that these don't mistakenly show up. + fqn_to_param_index = {} + index = 0 + for module_name, module in model.named_modules(): + for parameter_name, _param in module.named_parameters(recurse=False): + fqn = f"{module_name}.{parameter_name}" + fqn_to_param_index[fqn] = index + if fqn not in sparse_embedding_fqns: + index += 1 + if module in unused_modules: + expected_unused_param_fqns.append(fqn) + else: + if ( + not ignore_sparse + or module != model.sub_module.embedding_net.embedding + ): + used_param_fqns.append(fqn) + + net = torch.nn.parallel.DistributedDataParallel( + model.cuda(self.rank), + device_ids=[self.rank], + ) + batch, dim = 10, 2 + inp = torch.ones(batch, dim) + for i in range(2): + if i == 0: + out = net(inp) + loss = out.sum() + loss.backward() + else: + try: + out = net(inp) + loss = out.sum() + loss.backward() + except RuntimeError as e: + e = str(e) + + unused_param_substr = e[e.find("did not receive grad") :] + # Validate that each unused param fully qualified name + # shows up in error logs. We do this instead of + # constructing a joined string since order of parameters + # can be different in Reducer. In addition, validate + # param indices show up as well. + for unused_param_fqn in expected_unused_param_fqns: + self.assertTrue( + unused_param_fqn in unused_param_substr + or debug_mode_off + ) + self.assertTrue( + str(fqn_to_param_index[unused_param_fqn]) + in unused_param_substr, + f"Did not find index {fqn_to_param_index[unused_param_fqn]} for {unused_param_fqn}", + ) + + # Validate that used param fqns don't show up in error + # logs. + for used_param_fqn in used_param_fqns: + self.assertFalse(used_param_fqn in unused_param_substr) + # Validate that ignored param fqns don't show up as unused + # (since DDP does not track them) + for sparse_param_fqn in sparse_embedding_fqns: + self.assertFalse(sparse_param_fqn in unused_param_substr) + else: + self.assertTrue(False, "Expected error was not raised!") + + @with_dist_debug_levels(levels=["OFF", "INFO", "DETAIL"]) + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_multiple_nested_unused_params_error(self): + self._test_ddp_multiple_nested_unused_params_error(ignore_sparse=False) + + @with_dist_debug_levels(levels=["OFF", "INFO", "DETAIL"]) + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_multiple_nested_unused_params_err_ignore_params(self): + # Tests unused parameter reporting when DDP is configured to ignore + # certain parameters. + self._test_ddp_multiple_nested_unused_params_error(ignore_sparse=True) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_lt_x_gpu(2) + def test_ddp_inference(self): + # tests that DDP module can be run on a single node with no_grad + # or eval setting and there is no hang. + rank = self.rank + torch.cuda.set_device(rank) + model = Net().cuda() + local_model = copy.deepcopy(model) + model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[rank], + ) + syncbn_model = nn.SyncBatchNorm( + 2, momentum=0.99, track_running_stats=False + ).cuda() + local_syncbn_model = copy.deepcopy(syncbn_model) + syncbn_model = torch.nn.parallel.DistributedDataParallel( + syncbn_model, device_ids=[rank] + ) + inp = torch.randn(10, 2, device=rank) + inp_syncbn = torch.randn(10, 2, 4, 4, device=rank) + tests = [ + (model, local_model, inp), + (syncbn_model, local_syncbn_model, inp_syncbn), + ] + for test in tests: + test_model, test_local_model, test_inp = test + if self.rank == 0: + test_model.eval() + test_local_model.eval() + for _ in range(6): + self.assertEqual( + test_model(test_inp), test_local_model(test_inp) + ) + + # Barrier since only rank 0 runs inference. Test should be + # much faster than 30s, but this is to avoid flakiness. + self._barrier(timeout=30) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @skip_if_lt_x_gpu(2) + @unittest.skip( + "Test is failing, see https://github.com/pytorch/pytorch/pull/113620" + ) + def test_ddp_sync_bn_training_vs_eval(self): + rank = self.rank + torch.cuda.set_device(rank) + # Need to set track_running_stats=False, when track_running_stats=True, + # bn_training is False and sync could not occur in eval model. + model = nn.SyncBatchNorm(2, momentum=0.99, track_running_stats=False).cuda( + rank + ) + model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) + # Test sync occurs in training mode. + with torch.autograd.profiler.profile() as prof: + for _ in range(6): + inp = torch.randn(10, 2, 4, 4).cuda(rank) + out = model(inp) + loss = out.sum() + loss.backward() + + # SyncBN allgathers stats across all ranks, so verify call to + # all_gather in profiler. + if BACKEND == "nccl": + all_gather_calls = get_profiling_event("_all_gather_base", prof) + else: + all_gather_calls = get_profiling_event("all_gather", prof) + self.assertNotEqual([], all_gather_calls) + + # Only do inference on one rank. If SyncBN did collective stats sync, + # this would hang/error. + model_inference = model.module + if self.rank == 0: + model_inference.eval() + with torch.autograd.profiler.profile() as prof: + for _ in range(6): + inp = torch.randn(10, 2, 4, 4).cuda(rank) + out = model_inference(inp) + loss = out.sum() + loss.backward() + + # Ensure sync does not occur in eval() mode. + if BACKEND == "nccl": + all_gather_calls = get_profiling_event("_all_gather_base", prof) + else: + all_gather_calls = get_profiling_event("all_gather", prof) + self.assertEqual([], all_gather_calls) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_python_error_logged(self): + # Most python exceptions in DDP are raised during init before + # reducer is constructed, so we don't have a logger in those cases. + # However, the below is one example where a python error is thrown + # after reducer is constructed. + model = TwoLinLayerNet().cuda(self.rank) + model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + ) + expected_err = "must be callable" + with self.assertRaisesRegex(TypeError, expected_err): + model.register_comm_hook({}, {}) + + verify_ddp_error_logged(model, expected_err) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_static_graph_nested_types(self): + # Tests for static graph training when outputs are not just tensors + # but can be (nested) tuple, list, dict, etc. + rank = self.rank + torch.cuda.set_device(rank) + + class NestedOutputModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin = nn.Linear(100, 1, bias=False) + + def forward(self, inp, output_type): + if output_type == "tuple": + return ( + self.lin(inp), + ( + self.lin(inp), + self.lin(inp), + ), + ) + elif output_type == "list": + return [ + self.lin(inp), + [ + self.lin(inp), + self.lin(inp), + ], + ] + elif output_type == "dict": + return { + "a": self.lin(inp), + "b": { + "c": self.lin(inp), + }, + } + + def get_loss(model_output): + loss = 0.0 + if isinstance(model_output, torch.Tensor): + return model_output.sum() + elif isinstance(model_output, dict): + for value in model_output.values(): + loss += get_loss(value) + elif isinstance(model_output, (tuple, list)): + for x in model_output: + loss += get_loss(x) + else: + raise ValueError(f"Unknown model output type {type(model_output)}") + return loss + + model = NestedOutputModule().cuda(rank) + model_static_graph = copy.deepcopy(model) + model = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[rank], + ) + model_static_graph = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[rank], + static_graph=True, + ) + inp = torch.randn(10, 100) + type_mapping = { + "list": list, + "tuple": tuple, + "dict": dict, + } + for output_type in type_mapping: + for _ in range(6): + out = model(inp, output_type=output_type) + loss = get_loss(out) + loss.backward() + self._model_step(model) + out_static = model_static_graph(inp, output_type=output_type) + self.assertTrue(isinstance(out_static, type_mapping[output_type])) + loss_static = get_loss(out_static) + loss_static.backward() + self._model_step(model_static_graph) + for p, p_static in zip( + model.parameters(), model_static_graph.parameters(), strict=True + ): + self.assertEqual(p, p_static) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_returns_tensor_with_no_grad(self): + # Tests case where module returns tensor that does not require grad. + torch.cuda.set_device(self.rank) + + class MyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(10, 10, bias=False) + self.fc2 = nn.Linear(10, 10, bias=False) + + def forward(self, x): + x = self.fc2(F.relu(self.fc1(x))) + y = x.clone() + x = x.detach() + if x.requires_grad: + raise AssertionError("Expected x.requires_grad to be False") + return (x, y) + + model = MyModel().to(self.rank) + inp = torch.randn(1, 10, device=self.rank) + for find_unused, static_graph in itertools.product( + [True, False], [True, False] + ): + ddp = DistributedDataParallel( + model, + device_ids=[self.rank], + output_device=self.rank, + find_unused_parameters=find_unused, + static_graph=static_graph, + ) + for _ in range(6): + out = ddp(inp) + self.assertFalse(out[0].requires_grad) + o = (out[0] + out[1]).sum() + o.backward() + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_detect_ddp_is_actually_static(self): + class ToyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.net1 = nn.Linear(10, 10, bias=False) + self.net2 = nn.Linear(10, 10) + + def forward(self, x, find_unused, dynamic): + if find_unused: + if dynamic: + return self.net2(self.net1(x)) + else: + return self.net2(x) + else: + return self.net2(self.net1(x)) + + # Set of unused parameters don't change across iterations + torch.cuda.set_device(self.rank) + model = ToyModel().cuda() + for find_unused in [True, False]: + ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + find_unused_parameters=find_unused, + ) + inp = torch.randn(1, 10, device="cuda") + for _ in range(6): + out = ddp(inp, find_unused=find_unused, dynamic=False) + loss = out.sum() + loss.backward() + self.assertTrue(ddp.reducer._ddp_graph_static()) + + # Set of unused parameters dynamically change + ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + find_unused_parameters=True, + ) + inp = torch.randn(1, 10, device="cuda") + for i in range(6): + out = ddp(inp, find_unused=True, dynamic=i % 2 == 0) + loss = out.sum() + loss.backward() + self.assertFalse(ddp.reducer._ddp_graph_static()) + + def _test_ddp_new_tensor_in_fwd(self, static_graph): + # Test from https://github.com/pytorch/pytorch/issues/60733 + class MyModel(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(10, 10, bias=False) + self.fc2 = nn.Linear(10, 10, bias=False) + self.device = self.fc1.weight.device + + def __init_opt(self): + opt = torch.randn(1, 10, device=self.device) + return opt + + def forward(self, x, opt_1, opt_2, opt_nested): + x = F.relu(self.fc1(x)) + x = self.fc2(x) + if opt_1 is None: + opt_1 = self.__init_opt() + if opt_2 is None: + opt_2 = self.__init_opt() + if opt_nested is None or not torch.is_tensor(opt_nested): + opt_nested = self.__init_opt() + # Test multiple tensors as well as newly created tensors + # within a struct. + return x, opt_1, opt_2, {"tensor": opt_nested} + + model = MyModel().to(self.rank) + for find_unused in [True, False]: + ddp = DistributedDataParallel( + model, + device_ids=[self.rank], + output_device=self.rank, + broadcast_buffers=False, + find_unused_parameters=find_unused, + static_graph=static_graph, + ) + + opt = [None for _ in range(3)] + for i in range(2): + ddp.zero_grad() + x = torch.randn(1, 10, device=self.rank) + out, opt[0], opt[1], opt[2] = ddp( + x, opt_1=opt[0], opt_2=opt[1], opt_nested=opt[2] + ) + for i in range(len(opt)): + if torch.is_tensor(opt[i]): + self.assertEqual(opt[i].grad_fn, None) + else: + self.assertEqual(opt[i]["tensor"].grad_fn, None) + out.mean().backward() + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_new_tensor_in_fwd(self): + return self._test_ddp_new_tensor_in_fwd(static_graph=False) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_new_tensor_in_fwd_static_graph(self): + return self._test_ddp_new_tensor_in_fwd(static_graph=True) + + def _test_ddp_buffer_hook_allreduce(self, return_futures): + rank = self.rank + torch.cuda.set_device(rank) + torch.manual_seed(rank) + torch.cuda.manual_seed(rank) + + def buffer_comm_hook(ddp, named_buffers): + buffers = [buffer for (_, buffer) in named_buffers.items()] + futs = [ + dist.all_reduce( + buffer, group=ddp.process_group, async_op=True + ).get_future() + for buffer in buffers + ] + if return_futures: + return futs + else: + torch.futures.collect_all(futs).wait() + + hook_pre_fwd = ( + torch.nn.parallel.distributed._BufferCommHookLocation.PRE_FORWARD + ) + hook_post_fwd = ( + torch.nn.parallel.distributed._BufferCommHookLocation.POST_FORWARD + ) + for hook_run_location in [ + hook_pre_fwd, + hook_post_fwd, + ]: + model = NetWithBuffers().cuda(rank) + model_ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + ) + model_ddp._register_buffer_comm_hook( + model_ddp, buffer_comm_hook, hook_run_location + ) + model_ddp_no_hook = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(model), + device_ids=[self.rank], + broadcast_buffers=False, + ) + inp = torch.randn(2, 10, device=rank) + for _ in range(2): + loss_hook = model_ddp(inp).sum() + # Since buffer reduction is done pre-forward, simulate it for + # no hook case here. + # Simulate allreduce appropriately depending on hook location. + if hook_run_location == hook_pre_fwd: + model_no_hook_buffers = list(model_ddp_no_hook.module.buffers()) + for tensor in model_no_hook_buffers: + dist.all_reduce(tensor) + + loss_no_hook = model_ddp_no_hook(inp).sum() + if hook_run_location == hook_post_fwd: + model_no_hook_buffers = list(model_ddp_no_hook.module.buffers()) + for tensor in model_no_hook_buffers: + dist.all_reduce(tensor) + torch.cuda.synchronize() + + # if return_futures, they are only awaited on by DDP + # at the end of the backwards pass for maximum overlap. + if not return_futures: + self._verify_buffers_equal(model_ddp, model_ddp_no_hook) + loss_hook.backward() + loss_no_hook.backward() + # Note that when custom hooks return futures, this + # comparison is not expected to work when hook run location + # is pre-forward pass. This is because the hook does async + # communication and forward pass modifies the buffer without + # appropriate synchronization. Therefore, if returning + # futures from custom buffer hooks, it is advised to set + # hook run location to post forward. + if return_futures and hook_run_location == hook_post_fwd: + self._verify_buffers_equal(model_ddp, model_ddp_no_hook) + dist.barrier() + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_buffer_hook_allreduce_return_future(self): + self._test_ddp_buffer_hook_allreduce(return_futures=True) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_buffer_hook_allreduce(self): + self._test_ddp_buffer_hook_allreduce(return_futures=False) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_broadcast_buffer_via_hook(self): + # test that _distributed_broadcast_coalesced via registered hook is + # equivalent to DDP's default broadcast coalesced. + rank = self.rank + torch.cuda.set_device(rank) + torch.manual_seed(rank) + torch.cuda.manual_seed(rank) + + def buffer_comm_hook(ddp, named_buffers): + # named_buffers is a Dict[str, Tensor] representing a mapping + # from buffer name to buffer. + buffers = [buffer for (_, buffer) in named_buffers.items()] + ddp._default_broadcast_coalesced(buffers) + + model = NetWithBuffers().cuda(rank) + model_ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + ) + model_ddp._register_buffer_comm_hook(model_ddp, buffer_comm_hook) + model_ddp_no_hook = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(model), + device_ids=[self.rank], + ) + inp = torch.randn(2, 10, device=rank) + for _ in range(2): + loss_hook = model_ddp(inp).sum() + loss_no_hook = model_ddp_no_hook(inp).sum() + self._verify_buffers_equal(model_ddp, model_ddp_no_hook) + loss_hook.backward() + loss_no_hook.backward() + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_remove_autograd_hooks(self): + class SimulateError(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + def backward(ctx, grad_output): + raise RuntimeError + + class MyModel(nn.Module): + def __init__(self, device): + super().__init__() + self.error = True + self.fc1 = nn.Linear(10, 10).cuda(device) + + def forward(self, inp): + if self.error: + return self.fc1(SimulateError.apply(inp)) + else: + return self.fc1(inp) + + # Run with error to trigger backward pass that marks fc1 as being marked + # ready. If we don't remove autograd hooks before running below it would + # fail on the old autograd hook. + model = MyModel(self.rank) + input = torch.rand(10, 10, requires_grad=True).cuda(self.rank) + model_ddp1 = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + ) + + with self.assertRaises(RuntimeError): + model_ddp1(input).sum().backward() + + # Remove autograd hooks on old instance. + model_ddp1._remove_autograd_hooks() + + # Try another DDP instance without error now. + model.error = False + model_ddp2 = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + ) + model_ddp2(input).sum().backward() + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + @unittest.skip( + "Test is failing, tracking issue at https://github.com/pytorch/pytorch/issues/102751" + ) + def test_ddp_has_finalized(self): + @dataclass + class MyClass: + obj: torch.Tensor + + class MyModel(nn.Module): + def __init__(self, rank): + super().__init__() + self.rank = rank + self.fc1 = nn.Linear(1024, 1024).cuda(rank) + self.fc2 = nn.Linear(1024, 2 * 1024).cuda(rank) + + def forward(self, inp): + if self.rank == 0: + return self.fc1(inp), MyClass(self.fc2(inp)) + else: + return self.fc1(inp), self.fc2(inp) + + model = MyModel(self.rank) + input = torch.rand(10, 1024, requires_grad=True).cuda(self.rank) + ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + find_unused_parameters=True, + bucket_cap_mb=(1024 * 4 / 1024 / 1024), # One bucket per parameter. + ) + + if self.rank == 0: + out1, _ = ddp(input) + out1.sum().backward() + else: + out1, out2 = ddp(input) + (out1.sum() + out2.sum()).backward() + + if self.rank == 0: + with self.assertRaisesRegex( + RuntimeError, + "Expected to have finished reduction in the prior iteration", + ): + ddp._check_reducer_finalized() + + with self.assertRaisesRegex( + RuntimeError, + "Expected to have finished reduction in the prior iteration", + ): + ddp(input) + else: + ddp._check_reducer_finalized() + ddp(input) + + """ + # The set of "test_ddp_update_process_group..." below failed after + # upgrading CI from 2 GPUs to 4 GPUs. + # Commented out for now. + # Test purpose needs better documentation. + + def _run_ddp_update_process_group(self, new_pg): + def get_num_torch_recompiles(): + guard_failures = torch._dynamo.utils.guard_failures + num_recompiles = [len(guard_failures[code]) for code in guard_failures] + return 0 if len(num_recompiles) == 0 else max(num_recompiles) + + class SimulateError(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + def backward(ctx, grad_output): + raise RuntimeError + + class MyModel(torch.nn.Module): + def __init__(self, device): + super().__init__() + # 4MB for multiple buckets. + self.fc1 = torch.nn.Linear(1024, 1024).cuda(device) + self.fc2 = torch.nn.Linear(1024, 1024).cuda(device) + self.fc3 = torch.nn.Linear(1024, 1024).cuda(device) + + def forward(self, inp, error): + if error: + return self.fc3(self.fc2(self.fc1(SimulateError.apply(inp)))) + else: + return self.fc3(self.fc2(self.fc1(inp))) + + + input = torch.rand(10, 1024, requires_grad=True).cuda(self.rank) + ddp = torch.nn.parallel.DistributedDataParallel( + MyModel(self.rank), + device_ids=[self.rank], + find_unused_parameters=True, + bucket_cap_mb=1, + ) + model = torch.compile(ddp) + + def run_iteration(): + # Run regular iteration. + out = model(input, error=False) + out.sum().backward() + torch.cuda.synchronize() + + # Run with error. + with self.assertRaises(RuntimeError): + out = model(input, error=True) + out.sum().backward() + torch.cuda.synchronize() + + run_iteration() + if get_num_torch_recompiles() != 0: + raise AssertionError(f"Expected 0 torch recompiles, got {get_num_torch_recompiles()}") + + if new_pg: + # Now reduce world_size and run iteration. + group_size_2 = dist.new_group(ranks=[0, 1]) + ddp._update_process_group(group_size_2) + if self.rank in [0, 1]: + run_iteration() + + # Increase the world size and run iteration. + group_size_3 = dist.new_group(ranks=[1, 2, 3]) + ddp._update_process_group(group_size_3) + if self.rank in [1, 2, 3]: + run_iteration() + + # Back to default size. + ddp._update_process_group(_get_default_group()) + run_iteration() + else: + # Create default pg of smaller size. + dist.destroy_process_group() + + if self.rank in [1, 2, 3]: + dist.init_process_group( + init_method=self.init_method, + backend=BACKEND, + world_size=3, + rank=self.rank - 1, + timeout=timedelta(seconds=default_pg_timeout), + ) + ddp._update_process_group(_get_default_group()) + run_iteration() + dist.destroy_process_group() + + # Need a barrier here to ensure ranks 1, 2 and 3 are done. + self._barrier(wait_for=4) + + # Need to init pg again for "_barrier" to succeed. + dist.init_process_group( + init_method=self.init_method, + backend=BACKEND, + world_size=4, + rank=self.rank, + timeout=timedelta(seconds=default_pg_timeout), + ) + + # Validate no more recompiles. + if get_num_torch_recompiles() != 0: + raise AssertionError(f"Expected 0 torch recompiles, got {get_num_torch_recompiles()}") + + @skip_if_lt_x_gpu(4) + @require_world_size(4) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_update_process_group_new_group(self): + self._run_ddp_update_process_group(new_pg=True) + + @skip_if_lt_x_gpu(4) + @require_world_size(4) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_update_process_group_default_group(self): + self._run_ddp_update_process_group(new_pg=False) + + @skip_if_lt_x_gpu(4) + @require_world_size(4) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_update_process_group_grad_undefined(self): + class SimulateError(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + def backward(ctx, grad_output): + raise RuntimeError + + class MyModel(torch.nn.Module): + def __init__(self, device): + super().__init__() + self.fc1 = torch.nn.Linear(10, 10).cuda(device) + self.fc2 = torch.nn.Linear(10, 10).cuda(device) + self.fc3 = torch.nn.Linear(10, 10).cuda(device) + + def forward(self, inp, error): + if error: + return self.fc3(self.fc2(self.fc1(SimulateError.apply(inp)))) + else: + return self.fc2(self.fc1(inp)) + + + input = torch.rand(10, 10, requires_grad=True).cuda(self.rank) + ddp = torch.nn.parallel.DistributedDataParallel( + MyModel(self.rank), + device_ids=[self.rank], + find_unused_parameters=True, + bucket_cap_mb=1, + ) + + try: + ddp(input, True).sum().backward() + except RuntimeError: + ddp._update_process_group(_get_default_group()) + + # Reset grads. + for param in ddp.parameters(): + param.grad = None + + # Run ddp again. + ddp(input, False).sum().backward() + + @skip_if_lt_x_gpu(4) + @require_world_size(4) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_update_process_group_no_find_unused(self): + ddp = torch.nn.parallel.DistributedDataParallel( + torch.nn.Linear(10, 10).cuda(self.rank), + device_ids=[self.rank], + find_unused_parameters=False, + ) + ddp._update_process_group(_get_default_group()) + """ + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_broadcast_buffer(self): + rank = self.rank + torch.cuda.set_device(rank) + torch.manual_seed(rank) + torch.cuda.manual_seed(rank) + + class NetWithBuffers(nn.Module): + def __init__(self) -> None: + super().__init__() + self.a = nn.Linear(10, 10, bias=False) + self.b = nn.Linear(10, 1, bias=False) + self.register_buffer("buffer", torch.randn(1, 2)) + + def forward(self, x): + return self.b(self.a(x)) + + model = NetWithBuffers().cuda(rank) + model_ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + ) + inp = torch.randn(2, 10, device=rank) + for _ in range(2): + if rank == 0: + model_ddp.module.buffer = model_ddp.module.buffer + 1 + loss = model_ddp(inp).sum() + loss.backward() + # Ensure all buffers are synchronized. + bufs = [ + torch.empty_like(model_ddp.module.buffer) + for _ in range(dist.get_world_size()) + ] + dist.all_gather(bufs, model_ddp.module.buffer) + rank_0_buf = bufs[0] + for buf in bufs[1:]: + self.assertEqual(rank_0_buf, buf) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl" and BACKEND != "gloo", + "Only Nccl & Gloo backend support DistributedDataParallel", + ) + def test_static_graph_multi_forward(self): + class Net(nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin = nn.Linear(10, 10) + self.relu = nn.ReLU() + + def forward(self, x): + return self.relu(self.lin(x)) + + torch.cuda.set_device(self.rank) + torch.manual_seed(42 << 1337 % (self.rank + 1)) + model = Net().cuda(self.rank) + local_model = copy.deepcopy(model) + model = torch.nn.parallel.DistributedDataParallel( + model, device_ids=[self.rank], static_graph=True + ) + inp = torch.ones(2, 10, device="cuda") + for _ in range(3): + model.zero_grad() + local_model.zero_grad() + a = model(inp) + b = model(inp) + loss = a.sum() + b.sum() + loss.backward() + # Grads should be equal to a local model that ran through inp + # `world_size` times and averaged grads + if self.rank == 0: + inp_clone = inp.clone() + iters = dist.get_world_size() + for _ in range(iters): + a = local_model(inp_clone) + b = local_model(inp_clone) + loss = a.sum() + b.sum() + loss.backward() + + for p in local_model.parameters(): + p.grad.data = p.grad / iters + + for p_ddp, p_local in zip( + model.parameters(), local_model.parameters(), strict=True + ): + self.assertTrue( + torch.allclose(p_ddp.grad, p_local.grad), + f"{p_ddp.grad} vs {p_local.grad}", + ) + + dist.barrier() + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND != "nccl" and BACKEND != "gloo", + "Only Nccl & Gloo backend support DistributedDataParallel", + ) + def test_sync_bn_logged(self): + model = BatchNormNet() + rank = self.rank + # single gpu training setup + model_gpu = model.cuda(rank) + no_sync_bn = torch.nn.parallel.DistributedDataParallel( + copy.deepcopy(model_gpu), + device_ids=[self.rank], + ) + ddp_logging_data = no_sync_bn._get_ddp_logging_data() + sync_bn_logged = ddp_logging_data.get("has_sync_bn", True) + self.assertFalse(sync_bn_logged) + model_DDP = nn.SyncBatchNorm.convert_sync_batchnorm(model_gpu) + model_DDP = torch.nn.parallel.DistributedDataParallel( + model_DDP, + device_ids=[self.rank], + ) + ddp_logging_data = model_DDP._get_ddp_logging_data() + sync_bn_logged = ddp_logging_data.get("has_sync_bn", False) + self.assertTrue(sync_bn_logged) + + @skip_if_lt_x_gpu(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_stateless_api_with_ddp(self): + class MockModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.l1 = torch.nn.Linear(1, 1) + buffer = torch.ones(1) + self.register_buffer("buffer", buffer) + + def forward(self, x): + return self.l1(x) + self.buffer + + device = self.rank + module = MockModule().to(device) + module = torch.nn.parallel.DistributedDataParallel( + module, device_ids=[device] + ) + x = torch.rand((1, 1)).to(device) + weight = torch.tensor([[1.0]], device=device, requires_grad=True) + bias = torch.tensor([0.0], device=device, requires_grad=True) + buffer = torch.tensor([0.0], device=device) + parameters = { + "module.l1.weight": weight, + "module.l1.bias": bias, + "module.buffer": buffer, + } + prev_weight = module.module.l1.weight.clone() + prev_buffer = module.module.buffer.clone() + + res = torch.func.functional_call(module, parameters, x) + self.assertEqual(x, res) + # check that the weight remain unmodified + cur_weight = module.module.l1.weight + cur_buffer = module.module.buffer + self.assertEqual(cur_weight, prev_weight) + self.assertEqual(cur_buffer, prev_buffer) + # run a backward pass and check the gradients + res.backward() + self.assertIsNotNone(weight.grad) + self.assertIsNotNone(bias.grad) + # Gradient was not calculated for the module stated and buffers + self.assertIsNone(buffer.grad) + self.assertIsNone(module.module.l1.weight.grad) + self.assertIsNone(module.module.l1.bias.grad) + self.assertIsNone(module.module.buffer.grad) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_forward_backward_hook(self): + class DummyTestModel(nn.Module): + def __init__(self) -> None: + super().__init__() + torch.manual_seed(0) + self.fc = nn.Linear(2, 2) + + def forward(self, x): + return self.fc(x) + + def relu_hook(module, input): + return nn.functional.relu(input[0]) + + def gelu_hook(module, _input, output): + return nn.functional.gelu(output) + + def celu_hook(module, _input, output): + return (nn.functional.celu(output[0]),) + + local_model = DummyTestModel() + ddp_model = DummyTestModel() + local_model.fc.register_forward_pre_hook(relu_hook) + local_model.fc.register_forward_hook(gelu_hook) + ddp_model.fc.register_forward_pre_hook(relu_hook) + ddp_model.fc.register_forward_hook(gelu_hook) + local_model.fc.register_backward_hook(celu_hook) + ddp_model.fc.register_backward_hook(celu_hook) + ddp_model = DistributedDataParallel( + ddp_model.to(self.rank), device_ids=[self.rank] + ) + input_data = torch.rand(5, 2) + output_local = local_model(input_data) + output_ddp = ddp_model(input_data.to(self.rank)) + self.assertEqual(output_local, output_ddp) + output_local.sum().backward() + output_ddp.sum().backward() + ddp_grads = [p.grad for p in ddp_model.parameters()] + self.assertEqual(ddp_grads[0], local_model.fc.weight.grad) + self.assertEqual(ddp_grads[1], local_model.fc.bias.grad) + + def _test_hook_pickling(self, hook, hook_state): + torch.manual_seed(0) + learning_rate = 0.01 + chkpt_file = tempfile.gettempdir() + "/checkpoint.pt" + rank = self.rank + + input = torch.randn(7, 1, device=rank) + target = torch.randn(7, 5, device=rank) + net = torch.nn.Linear(1, 5).to(rank) + ddp_model = DistributedDataParallel(copy.deepcopy(net), device_ids=[rank]) + dummy_ddp_model = DistributedDataParallel( + copy.deepcopy(net), device_ids=[rank] + ) + optimizer = torch.optim.SGD(ddp_model.parameters(), lr=learning_rate) + ddp_model.register_comm_hook(hook_state, hook) + ddp_model.train() + + for _ in range(10): + optimizer.zero_grad() + out = ddp_model(input) + loss = F.mse_loss(out, target) + loss.backward() + optimizer.step() + + state = { + "state_dict": ddp_model.state_dict(), + "comm_hook": hook, + "comm_hook_state": hook_state, + } + + if rank == 0: + with self.assertLogs("torch.distributed") as captured: + torch.save(state, chkpt_file) + + # Check that the logger has only one entry + self.assertEqual(len(captured.records), 1) + # Check that the logger has an expected entry + self.assertEqual( + captured.records[0].getMessage(), + "NOTE: Process group is not serializable and excluded from a saved state.", + ) + + dist.barrier() + map_location = {"cuda:0": f"cuda:{rank:d}"} + with self.assertLogs("torch.distributed") as captured: + checkpoint = torch.load(chkpt_file, map_location=map_location) + + # Check that the logger has only one entry + self.assertEqual(len(captured.records), 1) + # Check that the logger has an expected entry + self.assertEqual( + captured.records[0].getMessage(), + "NOTE: Process group will be set to a default group (i.e. the world size).\ + If a different group is desired, please set `self.process_group` after PowerSGD state is loaded.", + ) + + dummy_ddp_model.load_state_dict(checkpoint["state_dict"]) + dummy_hook = checkpoint["comm_hook"] + dummy_hook_state = checkpoint["comm_hook_state"] + dummy_optimizer = torch.optim.SGD( + dummy_ddp_model.parameters(), lr=learning_rate + ) + + # Check that loaded function is correct + self.assertEqual(dummy_hook.__qualname__, hook.__qualname__) + + # Check that all slots' keys were restored correctly + self.assertEqual(hook_state.__slots__, dummy_hook_state.__slots__) + + # Check that all slots' attributes are restored correctly + # Excluding ``process_group`` and ``rng``. + for entry in dummy_hook_state.__slots__: + if entry != "process_group" and entry != "rng": + self.assertEqual( + getattr(dummy_hook_state, entry), getattr(hook_state, entry) + ) + + # Check that ``process_group`` was set to default + self.assertEqual(dummy_hook_state.process_group, _get_default_group()) + + # Check that a random state was restored properly: + # ``np.random.RandomState.get_state`` returns a tuple with entries: + # ``bit_generator`` - str, + # ``state.key`` - ndarray dtype[uint32], + # ``state.pos`` - int, + # ``has_gauss`` - int, + # ``gauss`` - float + # (refer to https://github.com/numpy/numpy/blob/266aad7478bc7fbcc55eea7f942a0d373b838396/numpy/random/mtrand.pyi) + # To make sure random state was restored properly, all entries should equal the original + for entry1, entry2 in zip( + hook_state.rng.get_state(), + dummy_hook_state.rng.get_state(), + strict=True, + ): + np.testing.assert_array_equal(entry1, entry2) + + dummy_ddp_model.register_comm_hook(dummy_hook_state, dummy_hook) + dummy_ddp_model.train() + + for _ in range(10): + optimizer.zero_grad() + dummy_optimizer.zero_grad() + out_origin = ddp_model(input) + out_dummy = dummy_ddp_model(input) + loss_origin = F.mse_loss(out_origin, target) + loss_dummy = F.mse_loss(out_dummy, target) + loss_origin.backward() + loss_dummy.backward() + optimizer.step() + dummy_optimizer.step() + + # Check that gradients after 10 epochs are the same + for orig_param, dummy_param in zip( + ddp_model.parameters(), dummy_ddp_model.parameters(), strict=True + ): + self.assertEqual(orig_param.grad, dummy_param.grad) + + dist.barrier() + if rank == 0: + os.remove(chkpt_file) + + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["cuda"], + f"The {BACKEND} backend does not support DDP communication hook on CUDA devices", + ) + @skip_if_lt_x_gpu(int(os.environ["WORLD_SIZE"])) + @skip_but_pass_in_sandcastle_if(True, "Skipped due to flakiness") + def test_ddp_hook_pickling_powerSGD(self): + hook = powerSGD.powerSGD_hook + powersgd_state = powerSGD.PowerSGDState( + process_group=None, + matrix_approximation_rank=1, + start_powerSGD_iter=4, + ) + self._test_hook_pickling(hook, powersgd_state) + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_ddp_device_mesh_initialization(self): + """ + Test DDP with device_mesh initialization. + """ + world_size = int(os.environ["WORLD_SIZE"]) + + from torch.distributed.device_mesh import init_device_mesh + + device_mesh = init_device_mesh("cuda", (world_size,)) + + pg = _get_default_group() + + torch.cuda.set_device(self.rank) + model = TwoLinLayerNet().cuda() + ddp_model = torch.nn.parallel.DistributedDataParallel( + model, device_mesh=device_mesh + ) + self.assertEqual(ddp_model.device_mesh, device_mesh) + + with self.assertRaisesRegex( + RuntimeError, + "Cannot specify both process_group and device_mesh arguments.", + ): + ddp_model = torch.nn.parallel.DistributedDataParallel( + model, process_group=pg, device_mesh=device_mesh + ) + + with self.assertRaisesRegex( + RuntimeError, "Only 1D device mesh is supported," + ): + device_mesh = init_device_mesh("cuda", (2, world_size // 2)) + ddp_model = torch.nn.parallel.DistributedDataParallel( + model, device_mesh=device_mesh + ) + + @skip_if_lt_x_gpu(2) + @require_world_size(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_compile_static_graph(self): + "Tests that DDP works with torch compile when static_graph=True" + model = torch.nn.Linear(10, 10).cuda(self.rank) + model_clone = copy.deepcopy(model) + ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + ) + ddp_static = torch.nn.parallel.DistributedDataParallel( + model_clone, device_ids=[self.rank], static_graph=True + ) + ddp = torch.compile(ddp) + ddp_static = torch.compile(ddp_static) + input = torch.rand(10, 10).cuda(self.rank) + # verify output and gradient parity + for _ in range(6): + out_ddp = ddp(input).sum() + out_ddp_static = ddp_static(input).sum() + self.assertEqual(out_ddp, out_ddp_static) + out_ddp.backward() + out_ddp_static.backward() + for p1, p2 in zip( + ddp.parameters(), ddp_static.parameters(), strict=True + ): + self.assertEqual(p1.grad, p2.grad) + + @skip_if_lt_x_gpu(2) + @require_world_size(2) + @skip_but_pass_in_sandcastle_if( + BACKEND not in DistTestCases.backend_feature["ddp"], + f"The {BACKEND} backend does not support DistributedDataParallel", + ) + def test_ddp_sink_noclone(self): + "Tests that we can configure DDP to avoid clone" + + class OpPatcher(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + func_packet = func._overloadpacket + if func_packet == torch.ops.aten.clone: + raise RuntimeError("clone encountered!") + kwargs = kwargs if kwargs else {} + return func(*args, **kwargs) + + class MyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc = torch.nn.Linear(10, 10) + + def forward(self, input): + return self.fc(input) + + model = MyModel().cuda(self.rank) + ddp = torch.nn.parallel.DistributedDataParallel( + model, + device_ids=[self.rank], + find_unused_parameters=True, + ) + ddp._set_ddp_sink_clone(False) + input = torch.rand(10, 10).cuda(self.rank) + + with OpPatcher(): + ddp(input).sum().backward() + + def _test_skip_all_reduce_unused_parameters( + self, + find_unused_parameters=False, + static_graph=False, + skip_all_reduce_unused_params=False, + ): + class LargeNet(nn.Module): + def __init__(self) -> None: + super().__init__() + self.fc1 = nn.Linear(100, 5000, bias=False) + # fc2 is unused + self.fc2 = nn.Linear(100, 100, bias=False) + + def forward(self, x): + y = self.fc1(x) + return y + + torch.manual_seed(31415) + torch.cuda.set_device(self.rank) + model = LargeNet().cuda(self.rank) + ddp_model = torch.nn.parallel.DistributedDataParallel( + model, + find_unused_parameters=find_unused_parameters, + static_graph=static_graph, + bucket_cap_mb=1.5, + skip_all_reduce_unused_params=skip_all_reduce_unused_params, + ) + random_input = torch.randn(20, 100, device=self.rank) + for _ in range(10): + out = ddp_model(random_input) + loss = out.sum() + loss.backward() + return ddp_model + + @require_backend_is_available(DistTestCases.backend_feature["gpu"]) + @skip_if_lt_x_gpu(2) + def test_skip_all_reduce_unused_parameters(self): + base_model = self._test_skip_all_reduce_unused_parameters( + find_unused_parameters=True, static_graph=False + ) + test_model_1 = self._test_skip_all_reduce_unused_parameters( + find_unused_parameters=True, + static_graph=False, + skip_all_reduce_unused_params=True, + ) + + self.assertEqual( + base_model._get_ddp_logging_data().get("num_buckets_reduced"), 2 + ) + self.assertEqual( + test_model_1._get_ddp_logging_data().get("num_buckets_reduced"), 1 + ) + + for i, j in zip( + base_model.parameters(), test_model_1.parameters(), strict=True + ): + self.assertEqual(i, j) + + +# ============================================================================= +# Unit tests for _BucketCapacityConfig dataclass +# These tests verify pure Python logic without requiring multi-process setup. +# ============================================================================= +class TestBucketCapacityConfig(unittest.TestCase): + """Unit tests for _BucketCapacityConfig dataclass.""" + + def test_create_with_default_bucket_cap(self): + """Test _BucketCapacityConfig.create() with default bucket_cap_mb (None).""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + expected_bytes = _DEFAULT_BUCKET_CAP_MB * _MB_TO_BYTES + self.assertEqual(config.bucket_bytes_cap, expected_bytes) + self.assertEqual(config.per_bucket_bytes_caps, ()) + self.assertEqual( + config.first_bucket_bytes_cap, dist._DEFAULT_FIRST_BUCKET_BYTES + ) + self.assertEqual(config.bucket_cap_mb, _DEFAULT_BUCKET_CAP_MB) + + def test_create_with_custom_bucket_cap_mb(self): + """Test _BucketCapacityConfig.create() with custom bucket_cap_mb.""" + custom_cap_mb = 100 + config = _BucketCapacityConfig.create( + bucket_cap_mb=custom_cap_mb, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + expected_bytes = custom_cap_mb * _MB_TO_BYTES + self.assertEqual(config.bucket_bytes_cap, expected_bytes) + self.assertEqual(config.per_bucket_bytes_caps, ()) + self.assertEqual(config.first_bucket_bytes_cap, expected_bytes) + self.assertEqual(config.bucket_cap_mb, custom_cap_mb) + + def test_create_with_bucket_cap_mb_list(self): + """Test _BucketCapacityConfig.create() with bucket_cap_mb_list.""" + cap_list = [10, 25, 50] + config = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=cap_list, + use_python_reducer=False, + ) + + expected_list = tuple(cap * _MB_TO_BYTES for cap in cap_list) + self.assertEqual(config.per_bucket_bytes_caps, expected_list) + expected_max_bytes = max(cap_list) * _MB_TO_BYTES + self.assertEqual(config.bucket_bytes_cap, expected_max_bytes) + self.assertEqual(config.first_bucket_bytes_cap, expected_max_bytes) + self.assertTrue(config.has_custom_per_bucket_caps) + + def test_create_with_bucket_cap_mb_list_overrides_bucket_cap_mb(self): + """Test that bucket_cap_mb_list takes precedence over bucket_cap_mb.""" + cap_list = [10, 25, 50] + config = _BucketCapacityConfig.create( + bucket_cap_mb=100, + bucket_cap_mb_list=cap_list, + use_python_reducer=False, + ) + + expected_max_bytes = max(cap_list) * _MB_TO_BYTES + self.assertEqual(config.bucket_bytes_cap, expected_max_bytes) + + def test_create_raises_for_python_reducer_with_list(self): + """Test that AssertionError is raised when using bucket_cap_mb_list with Python reducer.""" + with self.assertRaises(AssertionError) as context: + _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=[10, 25, 50], + use_python_reducer=True, + ) + + self.assertIn("python reducer is not supported", str(context.exception)) + + def test_has_custom_per_bucket_caps_property(self): + """Test has_custom_per_bucket_caps property.""" + config_no_list = _BucketCapacityConfig.create( + bucket_cap_mb=25, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + self.assertFalse(config_no_list.has_custom_per_bucket_caps) + + config_with_list = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=[10, 25], + use_python_reducer=False, + ) + self.assertTrue(config_with_list.has_custom_per_bucket_caps) + + def test_bucket_cap_mb_property(self): + """Test bucket_cap_mb property correctly derives MiB from bytes.""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=100, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + self.assertEqual(config.bucket_cap_mb, 100) + + config_default = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + self.assertEqual(config_default.bucket_cap_mb, _DEFAULT_BUCKET_CAP_MB) + + def test_immutability(self): + """Test that _BucketCapacityConfig is immutable (frozen dataclass).""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=25, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + with self.assertRaises(AttributeError): + config.bucket_bytes_cap = 100 + + with self.assertRaises(AttributeError): + config.per_bucket_bytes_caps = (100,) + + +class TestBucketCapacityConfigComputeLimits(unittest.TestCase): + """Unit tests for _BucketCapacityConfig.compute_bucket_size_limits().""" + + def test_compute_limits_with_custom_per_bucket_caps(self): + """Test compute_bucket_size_limits returns custom caps when provided.""" + cap_list = [10, 25, 50] + config = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=cap_list, + use_python_reducer=False, + ) + + limits, rebuild_limits = config.compute_bucket_size_limits( + static_graph=False, + find_unused_parameters=True, + ) + + expected_limits = [cap * _MB_TO_BYTES for cap in cap_list] + self.assertEqual(limits, expected_limits) + self.assertEqual(rebuild_limits, expected_limits) + + def test_compute_limits_static_graph_uses_maxsize(self): + """Test that static_graph=True uses sys.maxsize for initial bucketing.""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=25, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + limits, rebuild_limits = config.compute_bucket_size_limits( + static_graph=True, + find_unused_parameters=False, + ) + + self.assertEqual(limits, [sys.maxsize]) + self.assertEqual(rebuild_limits, []) + + def test_compute_limits_find_unused_false_uses_maxsize(self): + """Test that find_unused_parameters=False uses sys.maxsize.""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=25, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + limits, rebuild_limits = config.compute_bucket_size_limits( + static_graph=False, + find_unused_parameters=False, + ) + + self.assertEqual(limits, [sys.maxsize]) + self.assertEqual(rebuild_limits, []) + + def test_compute_limits_default_with_find_unused_true(self): + """Test default config with find_unused_parameters=True uses smaller first bucket.""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + limits, rebuild_limits = config.compute_bucket_size_limits( + static_graph=False, + find_unused_parameters=True, + ) + + expected_first = dist._DEFAULT_FIRST_BUCKET_BYTES + expected_main = _DEFAULT_BUCKET_CAP_MB * _MB_TO_BYTES + self.assertEqual(limits, [expected_first, expected_main]) + self.assertEqual(rebuild_limits, []) + + def test_compute_limits_custom_cap_with_find_unused_true(self): + """Test custom bucket_cap_mb with find_unused_parameters=True uses uniform bucket.""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=100, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + limits, rebuild_limits = config.compute_bucket_size_limits( + static_graph=False, + find_unused_parameters=True, + ) + + expected_bytes = 100 * _MB_TO_BYTES + self.assertEqual(limits, [expected_bytes]) + self.assertEqual(rebuild_limits, []) + + def test_compute_limits_empty_rebuild_for_non_list_config(self): + """Test that rebuild_limits is empty when not using bucket_cap_mb_list.""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=50, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + _, rebuild_limits = config.compute_bucket_size_limits( + static_graph=False, + find_unused_parameters=True, + ) + + self.assertEqual(rebuild_limits, []) + + +class TestBucketCapacityConfigEdgeCases(unittest.TestCase): + """Edge case tests for _BucketCapacityConfig.""" + + def test_single_element_bucket_cap_mb_list(self): + """Test bucket_cap_mb_list with a single element.""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=[50], + use_python_reducer=False, + ) + + self.assertEqual(config.per_bucket_bytes_caps, (50 * _MB_TO_BYTES,)) + self.assertEqual(config.bucket_bytes_cap, 50 * _MB_TO_BYTES) + self.assertTrue(config.has_custom_per_bucket_caps) + + def test_large_bucket_cap_mb(self): + """Test with large bucket_cap_mb value.""" + large_cap = 1000 + config = _BucketCapacityConfig.create( + bucket_cap_mb=large_cap, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + expected_bytes = large_cap * _MB_TO_BYTES + self.assertEqual(config.bucket_bytes_cap, expected_bytes) + self.assertEqual(config.first_bucket_bytes_cap, expected_bytes) + + def test_small_bucket_cap_mb(self): + """Test with small bucket_cap_mb value.""" + small_cap = 1 + config = _BucketCapacityConfig.create( + bucket_cap_mb=small_cap, + bucket_cap_mb_list=None, + use_python_reducer=False, + ) + + expected_bytes = small_cap * _MB_TO_BYTES + self.assertEqual(config.bucket_bytes_cap, expected_bytes) + + def test_empty_bucket_cap_mb_list(self): + """Test with empty bucket_cap_mb_list (should use default).""" + config = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=[], + use_python_reducer=False, + ) + + expected_bytes = _DEFAULT_BUCKET_CAP_MB * _MB_TO_BYTES + self.assertEqual(config.bucket_bytes_cap, expected_bytes) + self.assertEqual(config.per_bucket_bytes_caps, ()) + self.assertFalse(config.has_custom_per_bucket_caps) + + def test_bucket_cap_mb_list_with_duplicates(self): + """Test bucket_cap_mb_list with duplicate values.""" + cap_list = [25, 25, 25] + config = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=cap_list, + use_python_reducer=False, + ) + + expected_list = tuple(25 * _MB_TO_BYTES for _ in range(3)) + self.assertEqual(config.per_bucket_bytes_caps, expected_list) + self.assertEqual(config.bucket_bytes_cap, 25 * _MB_TO_BYTES) + + def test_bucket_cap_mb_list_ascending_order(self): + """Test bucket_cap_mb_list in ascending order.""" + cap_list = [5, 10, 15, 20] + config = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=cap_list, + use_python_reducer=False, + ) + + self.assertEqual(config.bucket_bytes_cap, 20 * _MB_TO_BYTES) + + def test_bucket_cap_mb_list_descending_order(self): + """Test bucket_cap_mb_list in descending order.""" + cap_list = [100, 50, 25, 10] + config = _BucketCapacityConfig.create( + bucket_cap_mb=None, + bucket_cap_mb_list=cap_list, + use_python_reducer=False, + ) + + self.assertEqual(config.bucket_bytes_cap, 100 * _MB_TO_BYTES) + + +class TestBucketCapacityConfigConstants(unittest.TestCase): + """Tests for module-level constants.""" + + def test_default_bucket_cap_mb_value(self): + """Test that _DEFAULT_BUCKET_CAP_MB is 25.""" + self.assertEqual(_DEFAULT_BUCKET_CAP_MB, 25) + + def test_mb_to_bytes_value(self): + """Test that _MB_TO_BYTES is 1024 * 1024.""" + self.assertEqual(_MB_TO_BYTES, 1024 * 1024) + + +instantiate_parametrized_tests(DistributedTest._DistTestBase) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/distributed_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/distributed_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..55a1e8f81a908a3d30aaf97b1eb32556907f0bc2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/distributed_utils.py @@ -0,0 +1,68 @@ +# mypy: allow-untyped-defs + +from contextlib import contextmanager +from datetime import timedelta +from functools import partial, wraps + +import torch.distributed as dist +import torch.distributed.distributed_c10d as c10d + + +class MockProcessGroup(dist.ProcessGroup): + def getBackendName(self): + return "mock_process_group" + + +def create_mock_pg(prefix_store, rank, world_size, timeout): + return MockProcessGroup(rank, world_size) + + +dist.Backend.register_backend("mock_process_group", create_mock_pg) + + +def mock_init_dist(rank, world_size): + # !!! WARNING !!! + # Kids don't try this at home, this is a cute pile of hacks that + # depends on a small mountain of c10d internals + if dist.is_initialized(): + raise AssertionError("Expected dist to not be initialized") + store = dist.HashStore() + # Trick _store_based_barrier into believing everyone else already checked-in + # Zero is the group index + store.add(f"{c10d.STORE_BASED_BARRIER_PREFIX}:0", world_size - 1) + dist.init_process_group( + backend="mock_process_group", + rank=rank, + world_size=world_size, + store=store, + group_name="fake", + timeout=timedelta(seconds=1), + ) + + +@contextmanager +def with_dist(rank=0, world_size=2): + """ + Context manager that initializer c10d with a fake process group. + """ + mock_init_dist(rank=rank, world_size=world_size) + try: + yield + finally: + dist.destroy_process_group() + + +def with_fake_comms(func=None, rank=0, world_size=2): + """ + Function wrapper that inits a fake process group designed for testing. + Right now only querying for world size is available + """ + if func is None: + return partial(with_fake_comms, rank=rank, world_size=world_size) + + @wraps(func) + def wrapper(self, *args, **kwargs): + with with_dist(rank, world_size): + func(self, *args, **kwargs) + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/fake_pg.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/fake_pg.py new file mode 100644 index 0000000000000000000000000000000000000000..af1917eeb0fd05c92be424104fd57f35a201694c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/fake_pg.py @@ -0,0 +1,35 @@ +# mypy: allow-untyped-defs + +import torch.distributed as dist +from torch._C._distributed_c10d import FakeProcessGroup + + +class FakeStore(dist.Store): + """ + A fake store is a fake Key-Value store simply for initialization usage + the of fake process group, one can either use FakeStore or HashStore. + """ + + +def _create_fake_pg(common_opts, backend_opts): + """ + A fake process group (not related to FakeTensor) is a process group which + doesn't actually do any communication, it just hallucinates some + communication. You can run a single rank with a fake process group + without needing multiple processes (simulates per-rank behavior) + + NOTE: This is not a real process group, and it would produce wrong results + for every collective. It should be used as a convenient tool when playing + with distributed but don't care about the actual data. + """ + return FakeProcessGroup._create_internal( + common_opts.group_rank, common_opts.group_size, backend_opts + ) + + +dist.Backend.register_backend( + dist.Backend.FAKE, + _create_fake_pg, + extended_api=True, + devices=["cpu", "cuda", "hpu", "xpu"], +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/multi_threaded_pg.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/multi_threaded_pg.py new file mode 100644 index 0000000000000000000000000000000000000000..a60a1479a91a6f788650ffa48e2429438e2cc1f3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/multi_threaded_pg.py @@ -0,0 +1,644 @@ +# mypy: allow-untyped-defs + +import sys +import threading +import weakref +from dataclasses import dataclass +from functools import partial, reduce + +import torch +import torch.distributed as dist +from torch._C._distributed_c10d import ( + _create_work_from_future, + AllgatherOptions, + AllreduceOptions, + AllToAllOptions, + BarrierOptions, + BroadcastOptions, + ReduceOp, + ReduceScatterOptions, + ScatterOptions, + Store, +) +from torch.distributed.distributed_c10d import _CollOp, _store_based_barrier, P2POp +from torch.futures import Future +from torch.utils import _pytree as pytree + + +""" +TODO: +Lots of missing collectives. +Collectives validation. +Make timeout robust by making collectives respect the test deadline. +Make tests robust by making collectives interruptible. +We need some synchronization around cleanup to ensure that timedout ranks don't cause spurious failures. + +""" + + +def flatten_list(lst): + return pytree.tree_leaves(lst) + + +def ret_work(ret): + fut = Future() + fut.set_result(ret) + return _create_work_from_future(fut) + + +def binop_reduce(tensors, op): + res = op(torch.stack(tensors), dim=0) + if isinstance(res, torch.Tensor): + return res + # min/max return a namedtuple + return res.values + + +def bitwise_reduce(tensors, op): + return reduce(op, tensors) + + +_reduce_ops = { + ReduceOp.SUM: partial(binop_reduce, op=torch.sum), + ReduceOp.AVG: partial(binop_reduce, op=torch.mean), + ReduceOp.PRODUCT: partial(binop_reduce, op=torch.prod), + ReduceOp.MIN: partial(binop_reduce, op=torch.min), + ReduceOp.MAX: partial(binop_reduce, op=torch.max), + ReduceOp.BAND: partial(bitwise_reduce, op=torch.bitwise_and), + ReduceOp.BOR: partial(bitwise_reduce, op=torch.bitwise_or), + ReduceOp.BXOR: partial(bitwise_reduce, op=torch.bitwise_xor), +} + + +# Note [Hide collectives mutation from autograd] +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# Threaded PG is intended to closely simulate the behavior of regular process +# groups. However, our regular PG implementations perform a dispatch through +# c10d, whereas Threaded PG does not for some reason (some superficial +# but not very convincing reasons include that Threaded PG is implemented +# in Python but you can't override Backend in Python, you can only override +# ProcessGroup in Python), thereby bypassing the dispatch step. Now we have +# a problem: c10d's signatures are LIES, they mutate their (output) tensor +# arguments but their annotations don't have mutations on them so we don't +# actually update any view metadata if you do differentiation. This +# ordinarily "doesn't matter" because distributed collectives aren't +# differentiable anyway, but it's possible to tickle this in testing if +# someone tries to touch the grad_fn of a Tensor. There a few ways to +# fix this, but the easiest way was to use the .detach() trick to hide +# the mutations from autograd. + + +class AllToAll: + @torch.no_grad() + def work(self, data): + world_size = len(data) + for dest_rank in range(world_size): + output_tensor_list, _ = data[dest_rank] + for src_rank in range(world_size): + _, input_tensor_list = data[src_rank] + # See Note [Hide collectives mutation from autograd] + output_tensor_list[src_rank].detach().copy_( + input_tensor_list[dest_rank] + ) + + +class AllToAllBase: + @torch.no_grad() + def work(self, data): + world_size = len(data) + for dest_rank in range(world_size): + output_buffer, _, output_split_sizes, _ = data[dest_rank] + + output_indexes = self._size_cumsum( + output_buffer.size(0), output_split_sizes, world_size + ) + + for src_rank in range(world_size): + _, input_buffer, _, input_split_sizes = data[src_rank] + input_indexes = self._size_cumsum( + input_buffer.size(0), input_split_sizes, world_size + ) + + # See Note [Hide collectives mutation from autograd] + output_buffer[ + output_indexes[src_rank] : output_indexes[src_rank + 1] + ].detach().copy_( + input_buffer[ + input_indexes[dest_rank] : input_indexes[dest_rank + 1] + ] + ) + + def _size_cumsum( + self, + buf_size: int, + sizes: torch.Tensor | list[int] | None, + world_size: int, + ) -> torch.Tensor: + if sizes is None or len(sizes) == 0: + sizes = torch.full((world_size,), buf_size // world_size, dtype=torch.int64) + if not isinstance(sizes, torch.Tensor): + sizes = torch.tensor(sizes, dtype=torch.int64) + if sizes.dtype != torch.int64: + raise AssertionError( + f"Expected sizes.dtype == torch.int64, got {sizes.dtype}" + ) + sizes = torch.cumsum( + torch.cat( + (torch.tensor([0], dtype=torch.int64, device=sizes.device), sizes), + dim=0, + ), + dim=0, + ) + return sizes + + +class AllReduce: + def __init__(self, op): + if op.op not in _reduce_ops: + raise NotImplementedError( + f"AllReduce op {op.op} not supported on multithreaded pg for now." + ) + self.op = op.op + + @torch.no_grad() + def work(self, data): + for i in range(len(data[0])): + # use rank0 as the device for sum + rank_0_device = data[0][i].device + # collect all data to the list and make them + # all on rank 0 device + tensors = [ + data[src_rank][i].to(rank_0_device) for src_rank in range(len(data)) + ] + + # now mimic reduce across all ranks + res = _reduce_ops[self.op](tensors) + + # copy all the reduced value to each rank + for src_rank in range(len(data)): + # See Note [Hide collectives mutation from autograd] + data[src_rank][i].detach().copy_(res.to(data[src_rank][i].device)) + + +class AllGather: + @torch.no_grad() + def work(self, data): + for src_rank in range(len(data)): + in_tensor_list = data[src_rank][1] + # Can't handle all_gather with multiple tensors + if len(in_tensor_list) != 1: + raise AssertionError( + f"Can't handle all_gather with multiple tensors, got {len(in_tensor_list)}" + ) + src_tensor = in_tensor_list[0] + + for dest in data: + dest_tensor = dest[0][0][src_rank] + # See Note [Hide collectives mutation from autograd] + dest_tensor.detach().copy_(src_tensor) + + +class Scatter: + def __init__(self, src): + self.src = src + + @torch.no_grad() + def work(self, data): + src_in_tensor_list = data[self.src][1] + # Can't handle scatter with multiple input tensor list + if len(src_in_tensor_list) != 1: + raise AssertionError( + f"Can't handle scatter with multiple input tensor list, got {len(src_in_tensor_list)}" + ) + src_in_tensors = src_in_tensor_list[0] + + for rank, each_rank_data in enumerate(data): + out_tensor_list = each_rank_data[0] + # Can't handle scatter with multiple output tensor + if len(out_tensor_list) != 1: + raise AssertionError( + f"Can't handle scatter with multiple output tensor, got {len(out_tensor_list)}" + ) + dest_tensor = out_tensor_list[0] + # See Note [Hide collectives mutation from autograd] + dest_tensor.detach().copy_(src_in_tensors[rank]) + + +class Gather: + def __init__(self, dst): + self.dst = dst + + @torch.no_grad() + def work(self, data): + # Can't handle gather with multiple tensor lists + if len(data[self.dst][0]) != 1: + raise AssertionError( + f"Can't handle gather with multiple tensor lists, got {len(data[self.dst][0])}" + ) + out_tensor_list = data[self.dst][0][0] + for rank, each_rank_data in enumerate(data): + src_in_tensor_list = each_rank_data[1] + # Can't handle gather with multiple tensor lists + if len(src_in_tensor_list) != 1: + raise AssertionError( + f"Can't handle gather with multiple tensor lists, got {len(src_in_tensor_list)}" + ) + dest_tensor = out_tensor_list[rank] + # See Note [Hide collectives mutation from autograd] + dest_tensor.detach().copy_(src_in_tensor_list[0]) + + +class ReduceScatter: + def __init__(self, op): + if op != dist.ReduceOp.SUM and op != dist.ReduceOp.AVG: + raise NotImplementedError(f"ReduceScatter does not support {op}") + self.op = op + + @torch.no_grad() + def work(self, data): + start_reduction = [False for _ in range(len(data))] + for each_rank_data in data: + # Can't handle reduce_scatter with multiple scatter list + if len(each_rank_data[1]) != 1: + raise AssertionError( + f"Can't handle reduce_scatter with multiple scatter list, got {len(each_rank_data[1])}" + ) + to_scatter = each_rank_data[1][0] + for i in range(len(to_scatter)): + dest_tensor_on_rank_i = data[i][0] + # Can't handle reduce_scatter with multiple output tensor + if len(dest_tensor_on_rank_i) != 1: + raise AssertionError( + f"Can't handle reduce_scatter with multiple output tensor, got {len(dest_tensor_on_rank_i)}" + ) + dst_tensor_device = dest_tensor_on_rank_i[0].device + if not start_reduction[i]: + # See Note [Hide collectives mutation from autograd] + dest_tensor_on_rank_i[0].detach().copy_( + to_scatter[i].to(dst_tensor_device) + ) + start_reduction[i] = True + else: + # See Note [Hide collectives mutation from autograd] + dest_tensor_on_rank_i[0].detach().add_( + to_scatter[i].to(dst_tensor_device) + ) + if self.op == dist.ReduceOp.AVG: + num_ranks = len(data) + for each_rank_data in data: + # See Note [Hide collectives mutation from autograd] + each_rank_data[0][0].detach().div_(num_ranks) + + +class Broadcast: + def __init__(self, src): + self.src = src + + @torch.no_grad() + def work(self, data): + in_tensor_list = flatten_list(data[self.src]) + for i in range(len(data)): + if i == self.src: + continue + out_tensor_list = flatten_list(data[i]) + for j in range(len(in_tensor_list)): + # See Note [Hide collectives mutation from autograd] + out_tensor_list[j].detach().copy_(in_tensor_list[j]) + + +class Collective: + def __init__(self, world_size, collective, pg): + self._world_size = world_size + self._collective = collective + + self._start_cond = threading.Condition() + self._done_cond = threading.Condition() + + self._data = [None] * world_size + self._count = 0 + self._done = False + + self._pg = pg + + def join(self, rank, data): + with self._start_cond: + self._data[rank] = data + self._count += 1 + + # notify rank 0 + if self._count == self._world_size: + if rank > 0: + self._start_cond.notify() + + if rank == 0: + self._start_cond.wait_for( + lambda: self._count == self._world_size + or self._pg._terminate.is_set() + ) + # SystemExit is not a subclass of Exception but BaseException + # and can be distinguished from normal exception raised from program errors + # so that we can hide it from the exception queue + if self._pg._terminate.is_set(): + sys.exit("Test termination event occurs.") + + with self._done_cond: + # wait for rank 0 to finish + if rank > 0: + self._done_cond.wait_for( + lambda: self._done or self._pg._terminate.is_set() + ) + if self._pg._terminate.is_set(): + sys.exit("Test termination event occurs.") + else: + # copy data around + self._collective.work(self._data) + self._done = True + self._done_cond.notify_all() + return ret_work(data) + + +class ProcessLocalGroup(dist.ProcessGroup): + _coll_lock = threading.Lock() + _cur_coll_on_pgs = {} + + _terminate = threading.Event() + + @classmethod + def _start_coll(cls, collective, pg): + with cls._coll_lock: + # pg_name is unique, we use that to record the mapping between pg and collective + if pg.pg_name not in cls._cur_coll_on_pgs: + cls._cur_coll_on_pgs[pg.pg_name] = Collective( + pg.size(), collective, cls + ) + return cls._cur_coll_on_pgs[pg.pg_name] + + @classmethod + def _end_coll(cls, collective, pg): + # This is racily called by all ranks, so only one will work + with cls._coll_lock: + if ( + pg.pg_name in cls._cur_coll_on_pgs + and cls._cur_coll_on_pgs[pg.pg_name] == collective + ): + cls._cur_coll_on_pgs.pop(pg.pg_name) + + @classmethod + def exception_handle(cls, exc): + cls._terminate.set() + for coll in cls._cur_coll_on_pgs.values(): + with coll._start_cond: + coll._start_cond.notify() + with coll._done_cond: + coll._done_cond.notify_all() + + @classmethod + def reset(cls): + with cls._coll_lock: + cls._cur_coll_on_pgs = {} + cls._terminate.clear() + + def alltoall_base( + self, + output_buffer: torch.Tensor, + input_buffer: torch.Tensor, + output_split_sizes: list[int] | None, + input_split_sizes: list[int] | None, + opts=AllToAllOptions(), + ) -> torch.Tensor: + coll = ProcessLocalGroup._start_coll(AllToAllBase(), self) + res = coll.join( + self._rank, + (output_buffer, input_buffer, output_split_sizes, input_split_sizes), + ) + ProcessLocalGroup._end_coll(coll, self) + return res + + def alltoall(self, output_tensor_list, input_tensor_list, opts=AllToAllOptions()): + coll = ProcessLocalGroup._start_coll(AllToAll(), self) + res = coll.join(self._rank, (output_tensor_list, input_tensor_list)) + ProcessLocalGroup._end_coll(coll, self) + return res + + def allreduce(self, tensor_list, opts=AllreduceOptions()): + coll = ProcessLocalGroup._start_coll(AllReduce(opts.reduceOp), self) + res = coll.join(self._rank, tensor_list) + ProcessLocalGroup._end_coll(coll, self) + return res + + def allreduce_coalesced(self, tensor_list, opts=AllreduceOptions()): + coll = ProcessLocalGroup._start_coll(AllReduce(opts.reduceOp), self) + res = coll.join(self._rank, tensor_list) + ProcessLocalGroup._end_coll(coll, self) + return res + + def barrier(self, opts=BarrierOptions()): + return self.allreduce(tensor_list=[torch.ones(1)]) + + def allgather(self, output_tensors, input_tensor, opts=AllgatherOptions()): + coll = ProcessLocalGroup._start_coll(AllGather(), self) + res = coll.join(self._rank, (output_tensors, input_tensor)) + ProcessLocalGroup._end_coll(coll, self) + return res + + def _allgather_base(self, output_tensor, input_tensor, opts=AllgatherOptions()): + tensor_list = list(torch.chunk(output_tensor, self._world_size)) + return self.allgather([tensor_list], [input_tensor], opts) + + def broadcast(self, tensor_list, opts=BroadcastOptions()): + coll = ProcessLocalGroup._start_coll(Broadcast(opts.rootRank), self) + res = coll.join(self._rank, tensor_list) + ProcessLocalGroup._end_coll(coll, self) + return res + + def scatter(self, output_tensors, input_tensors, opts=ScatterOptions()): + coll = ProcessLocalGroup._start_coll(Scatter(opts.rootRank), self) + res = coll.join(self._rank, (output_tensors, input_tensors)) + ProcessLocalGroup._end_coll(coll, self) + return res + + def gather(self, output_tensors, input_tensors, opts=ScatterOptions()): + coll = ProcessLocalGroup._start_coll(Gather(opts.rootRank), self) + res = coll.join(self._rank, (output_tensors, input_tensors)) + ProcessLocalGroup._end_coll(coll, self) + return res + + def reduce_scatter(self, output_tensor, scatter_list, opts=ReduceScatterOptions()): + coll = ProcessLocalGroup._start_coll(ReduceScatter(opts.reduceOp), self) + res = coll.join(self._rank, (output_tensor, scatter_list)) + ProcessLocalGroup._end_coll(coll, self) + return res + + def _reduce_scatter_base( + self, output_tensor, input_tensor, opts=ReduceScatterOptions() + ): + tensor_list = list(torch.chunk(input_tensor, self._world_size)) + return self.reduce_scatter([output_tensor], [tensor_list], opts) + + def reduce_scatter_tensor_coalesced( + self, output_tensors, input_tensors, opts=ReduceScatterOptions() + ): + works = [ + self._reduce_scatter_base(output_tensor, input_tensor, opts) + for output_tensor, input_tensor in zip( + output_tensors, input_tensors, strict=True + ) + ] + for work in works[:-1]: + work.wait() + return works[-1] + + def allgather_into_tensor_coalesced( + self, output_tensor_list, input_tensor_list, opts=AllgatherOptions() + ): + res = None + for o_t, i_t in zip(output_tensor_list, input_tensor_list, strict=True): + res = self._allgather_base(o_t, i_t) + return res + + def __init__(self, rank, world_size): + super().__init__(rank, world_size) + self._rank = rank + self._world_size = world_size + world = dist.distributed_c10d._world + if isinstance(world, ThreadLocalWorld): + world = world._get_world() + self._world = weakref.ref(world) + self._ctx = torch.autograd.set_multithreading_enabled(False) + + def size(self): + return self._world_size + + @property + def pg_name(self): + """ + return the global registered name of the current pg in the world + """ + return self._world().pg_names[self] + + @property + def group_name(self): + return self.pg_name + + def getBackendName(self): + return "threaded" + + def __repr__(self): + return f"ThreadedPG world_size:{self._world_size} rank:{self._rank}" + + +def _create_threaded_pg(prefix_store, rank, world_size, timeout): + pg = ProcessLocalGroup(rank, world_size) + # https://github.com/pytorch/pytorch/pull/103033 changed store based barrier to optional + # When device mesh involves sub groups while store based barrier is not enabled in c10d, + # even though threaded pg actual collectives are assumed to be single threaded, + # different threads may be initializing different groups, + # leading to race conditions. + # For example, if we have a mesh of [[0, 1], [2, 3]], the sub groups + # (dim 0 and 1) would be initialized in different threads independently. + # In this case we can no longer rely on class or global variables + # but have to rely on store based barrier to make sure each group + # is ready separately before we can invoke collectives in any of the groups. + + # the prefix store is already per group so we pass an empty name here + _store_based_barrier(rank, prefix_store, "", world_size, timeout) + return pg + + +dist.Backend.register_backend("threaded", _create_threaded_pg, devices=["cpu", "cuda"]) + + +@dataclass +class WorldData: + default_pg: dist.ProcessGroup + pg_map: dict[dist.ProcessGroup, tuple[str, Store | None]] + pg_names: dict[dist.ProcessGroup, str] + pg_group_ranks: dict[dist.ProcessGroup, dict[int, int]] + pg_backend_config: dict[dist.ProcessGroup, str] + group_count: int + tags_to_pg: dict[str, list[dist.ProcessGroup]] + pg_to_tag: dict[dist.ProcessGroup, str] + pg_coalesce_state: dict[dist.ProcessGroup, list[_CollOp | P2POp]] + comms: list + + +class ThreadLocalWorld: + _world = threading.local() + + def _get_world(self) -> WorldData: + if not hasattr(ThreadLocalWorld._world, "world"): + ThreadLocalWorld._world.world = WorldData( + None, {}, {}, {}, {}, 0, {}, {}, {}, [] + ) + return ThreadLocalWorld._world.world + + @property + def default_pg(self): + return self._get_world().default_pg + + @default_pg.setter + def default_pg(self, value): + self._get_world().default_pg = value + + @property + def pg_map(self): + return self._get_world().pg_map + + @property + def pg_names(self): + return self._get_world().pg_names + + @property + def pg_group_ranks(self): + return self._get_world().pg_group_ranks + + @property + def pg_backend_config(self): + return self._get_world().pg_backend_config + + @property + def group_count(self) -> int: + return self._get_world().group_count + + @group_count.setter + def group_count(self, value): + self._get_world().group_count = value + + @property + def tags_to_pg(self): + return self._get_world().tags_to_pg + + @property + def pg_to_tag(self): + return self._get_world().pg_to_tag + + @property + def pg_coalesce_state(self) -> dict[dist.ProcessGroup, list[_CollOp | P2POp]]: + return self._get_world().pg_coalesce_state + + @property + def comms(self): + return self._get_world().comms + + +_old_pg_world = None +_ctx_manager = None + + +def _install_threaded_pg(): + global _old_pg_world + global _ctx_manager + _old_pg_world = dist.distributed_c10d._world + dist.distributed_c10d._world = ThreadLocalWorld() + _ctx_manager = torch.autograd.set_multithreading_enabled(False) + + return dist.distributed_c10d._world + + +def _uninstall_threaded_pg(): + global _ctx_manager + dist.distributed_c10d._world = _old_pg_world + # Restore autograd multithreading state that was disabled in _install_threaded_pg + if _ctx_manager is not None: + _ctx_manager.__exit__(None, None, None) + _ctx_manager = None diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/nn/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/nn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/nn/api/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/nn/api/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/nn/api/remote_module_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/nn/api/remote_module_test.py new file mode 100644 index 0000000000000000000000000000000000000000..af136fb8722d17d70767718a0cd327f71d730fda --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/nn/api/remote_module_test.py @@ -0,0 +1,754 @@ +# mypy: allow-untyped-defs + +import enum + +import torch +import torch.distributed.rpc as rpc +import torch.testing._internal.dist_utils as dist_utils +from torch import nn, Tensor +from torch._jit_internal import Future +from torch.distributed.nn import RemoteModule +from torch.distributed.nn.api.remote_module import ( + _REMOTE_MODULE_PICKLED_ATTRIBUTES, + _RemoteModule, +) +from torch.testing._internal.common_distributed import skip_if_lt_x_gpu +from torch.testing._internal.common_utils import TemporaryFileName, TEST_WITH_ROCM +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +_PARAM_VAL = torch.nn.Parameter(torch.ones(1)) + + +# RPC handler for querying the device on the destination worker. +def remote_device(module_rref): + for param in module_rref.local_value().parameters(): + return param.device + + +# RPC handler for querying __dict__ on the destination worker. +def remote_module_attributes(remote_module): + return remote_module.__dict__ + + +# RPC handler for running forward on the destination worker. +def remote_forward(remote_module, args): + return remote_module.forward(*args) + + +# RPC handler for running forward_async on the destination worker. +def remote_forward_async(remote_module, args): + # Since future cannot be pickled and sent over the RPC layer, + # have to wait and behave just like ``forward_sync``. + return remote_module.forward_async(*args).wait() + + +# RPC handler for getting training mode on the destination worker. +def get_remote_training_arg(module_rref): + return module_rref.local_value().training + + +class ModuleCreationMode(enum.Enum): + MODULE_CTOR_WITH_INTERFACE = "module_ctor_with_interface" + MODULE_CTOR = "module_ctor" + + +@torch.jit.interface +class MyModuleInterface: + def forward( + self, tensor: Tensor, number: int, word: str = "default" + ) -> tuple[str, int, Tensor]: + # pyre-ignore[7]: Pyre and torch.jit.interface don't mix well + pass + + +@torch.jit.interface +class RemoteMyModuleInterface: + def forward( + self, tensor: Tensor, number: int, word: str = "default" + ) -> tuple[str, int, Tensor]: + # pyre-ignore[7]: Pyre and torch.jit.interface don't mix well + pass + + def forward_async( + self, tensor: Tensor, number: int, word: str = "default" + ) -> Future[tuple[str, int, Tensor]]: + pass + + +class MyModule(nn.Module): + def __init__(self, first_arg, first_kwarg=-1): + super().__init__() + self.param1 = _PARAM_VAL + + def forward( + self, tensor: Tensor, number: int, word: str = "default" + ) -> tuple[str, int, Tensor]: + return word, number, tensor + + +class BadModule: + def __init__(self, first_arg, first_kwarg=-1): + pass + + +def create_scripted_module(first_arg, first_kwarg=-1): + module = MyModule(first_arg, first_kwarg=first_kwarg) + scripted_module = torch.jit.script(module) + return scripted_module + + +# Common utils for both CPU and CUDA test suites +class CommonRemoteModuleTest(RpcAgentTestFixture): + @property + def world_size(self): # Override setting in RpcAgentTestFixture + return 2 + + @staticmethod + def _create_remote_module_iter(remote_device, modes=None): + if modes is None: + modes = ModuleCreationMode.__members__.values() + + args = (1,) + kwargs = dict(first_kwarg=2) + + if ModuleCreationMode.MODULE_CTOR in modes: + remote_module = RemoteModule(remote_device, MyModule, args, kwargs) + yield remote_module + + if ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE in modes: + remote_module = _RemoteModule( + remote_device, + create_scripted_module, + args, + kwargs, + _module_interface_cls=MyModuleInterface, + ) + scripted_remote_module = torch.jit.script(remote_module) + yield scripted_remote_module + + +class RemoteModuleTest(CommonRemoteModuleTest): + @dist_utils.dist_init + def test_bad_module(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + remote_device = f"{dst_worker_name}/cpu" + args = (1,) + kwargs = dict(first_kwarg=2) + + with self.assertRaisesRegex( + ValueError, + r"Expect `module_cls\(\*args, \*\*kwargs\)` returns an instance of ,", + ): + RemoteModule(remote_device, BadModule, args, kwargs).forward() + + with self.assertRaisesRegex( + ValueError, + r"Expect `module_cls\(\*args, \*\*kwargs\)` returns an instance of ,", + ): + RemoteModule(remote_device, BadModule, args, kwargs).forward() + + @dist_utils.dist_init + def test_forward_async(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + args = (torch.ones(1), 2, "3") + for remote_module in self._create_remote_module_iter(dst_worker_name): + ret_fut = remote_module.forward_async(*args) + ret = ret_fut.wait() + self.assertEqual(ret, tuple(reversed(args))) + + @dist_utils.dist_init + def test_forward_async_script(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + scripted_remote_module = next( + self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE] + ) + ) + + @torch.jit.script + def run_forward_async(scripted_remote_module: RemoteMyModuleInterface): + ret_fut = scripted_remote_module.forward_async(torch.ones(1), 2, "3") + ret = ret_fut.wait() + return ret + + ret = run_forward_async(scripted_remote_module) + + self.assertEqual(ret, ("3", 2, torch.ones(1))) + + @dist_utils.dist_init + def test_forward_sync(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + args = (torch.ones(1), 2, "3") + for remote_module in self._create_remote_module_iter(dst_worker_name): + ret = remote_module.forward(*args) + self.assertEqual(ret, tuple(reversed(args))) + + @dist_utils.dist_init + def test_forward_sync_script(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + scripted_remote_module = next( + self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE] + ) + ) + + @torch.jit.script + def run_forward(scripted_remote_module: MyModuleInterface): + ret = scripted_remote_module.forward(torch.ones(1), 2, "3") + return ret + + ret = run_forward(scripted_remote_module) + + self.assertEqual(ret, ("3", 2, torch.ones(1))) + + @dist_utils.dist_init + def test_forward_with_kwargs(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + args = (torch.ones(1), 2) + kwargs = dict(word="3") + # Only test Python nn.Module, because script module methods don't support taking kwargs. + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + ret_fut = remote_module.forward_async(*args, **kwargs) + ret = ret_fut.wait() + self.assertEqual(ret, tuple(reversed(args + ("3",)))) + + ret = remote_module.forward(*args, **kwargs) + self.assertEqual(ret, tuple(reversed(args + ("3",)))) + + @dist_utils.dist_init + def test_remote_parameters(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + # Only test Python nn.Module, because script module methods don't support ``remote_parameters``. + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + param_rrefs = remote_module.remote_parameters() + self.assertEqual(len(param_rrefs), 1) + self.assertTrue(torch.equal(param_rrefs[0].to_here(), _PARAM_VAL)) + + @dist_utils.dist_init + def test_get_module_rref(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + # Only test Python nn.Module, because script module methods don't support ``get_module_rref``. + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + rref = remote_module.get_module_rref() + self.assertEqual(rref, remote_module.module_rref) + for param in rref.to_here().parameters(): + self.assertTrue(torch.equal(param, _PARAM_VAL)) + + @dist_utils.dist_init + def test_train_eval(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + remote_module.train() + ret1 = rpc.rpc_sync( + dst_worker_name, + get_remote_training_arg, + args=(remote_module.get_module_rref(),), + ) + self.assertEqual(ret1, True) + + remote_module.eval() + ret2 = rpc.rpc_sync( + dst_worker_name, + get_remote_training_arg, + args=(remote_module.get_module_rref(),), + ) + self.assertEqual(ret2, False) + + @dist_utils.dist_init + def test_unsupported_methods(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + with self.assertRaisesRegex( + ValueError, r"Method ``register_buffer`` not supported for RemoteModule" + ): + remote_module.register_buffer("buffer", torch.ones(5)) + with self.assertRaisesRegex( + ValueError, + r"Method ``register_parameter`` not supported for RemoteModule", + ): + remote_module.register_parameter( + "param", torch.nn.Parameter(torch.ones(1)) + ) + with self.assertRaisesRegex( + ValueError, r"Method ``add_module`` not supported for RemoteModule" + ): + remote_module.add_module("empty", None) + + with self.assertRaisesRegex( + ValueError, r"Method ``apply`` not supported for RemoteModule" + ): + fn = torch.rand((3, 3), requires_grad=False) + remote_module.apply(fn) + + with self.assertRaisesRegex( + ValueError, r"Method ``cuda`` not supported for RemoteModule" + ): + remote_module.cuda() + with self.assertRaisesRegex( + ValueError, r"Method ``cpu`` not supported for RemoteModule" + ): + remote_module.cpu() + with self.assertRaisesRegex( + ValueError, r"Method ``type`` not supported for RemoteModule" + ): + remote_module.type(torch.FloatTensor) + with self.assertRaisesRegex( + ValueError, r"Method ``float`` not supported for RemoteModule" + ): + remote_module.float() + with self.assertRaisesRegex( + ValueError, r"Method ``double`` not supported for RemoteModule" + ): + remote_module.double() + with self.assertRaisesRegex( + ValueError, r"Method ``bfloat16`` not supported for RemoteModule" + ): + remote_module.bfloat16() + with self.assertRaisesRegex( + ValueError, r"Method ``to`` not supported for RemoteModule" + ): + remote_module.to("cpu", dtype=torch.int32) + + def hook(module, grad_input, grad_output): + pass + + with self.assertRaisesRegex( + ValueError, + r"Method ``register_backward_hook`` not supported for RemoteModule", + ): + remote_module.register_backward_hook(hook) + with self.assertRaisesRegex( + ValueError, + r"Method ``register_forward_pre_hook`` not supported for RemoteModule", + ): + remote_module.register_forward_pre_hook(hook) + with self.assertRaisesRegex( + ValueError, + r"Method ``register_forward_hook`` not supported for RemoteModule", + ): + remote_module.register_forward_hook(hook) + + with self.assertRaisesRegex( + ValueError, r"Method ``state_dict`` not supported for RemoteModule" + ): + remote_module.state_dict() + with self.assertRaisesRegex( + ValueError, r"Method ``load_state_dict`` not supported for RemoteModule" + ): + remote_module.load_state_dict({}) + + with self.assertRaisesRegex( + ValueError, + r"Method ``parameters`` not supported for RemoteModule. Please use ``remote_parameters`` instead.", + ): + remote_module.parameters() + with self.assertRaisesRegex( + ValueError, + r"Method ``named_parameters`` not supported for RemoteModule", + ): + remote_module.named_parameters() + with self.assertRaisesRegex( + ValueError, r"Method ``buffers`` not supported for RemoteModule" + ): + remote_module.buffers() + with self.assertRaisesRegex( + ValueError, r"Method ``named_buffers`` not supported for RemoteModule" + ): + remote_module.named_buffers() + with self.assertRaisesRegex( + ValueError, r"Method ``children`` not supported for RemoteModule" + ): + remote_module.children() + with self.assertRaisesRegex( + ValueError, r"Method ``named_children`` not supported for RemoteModule" + ): + remote_module.named_children() + with self.assertRaisesRegex( + ValueError, r"Method ``modules`` not supported for RemoteModule" + ): + remote_module.modules() + with self.assertRaisesRegex( + ValueError, r"Method ``named_modules`` not supported for RemoteModule" + ): + remote_module.named_modules() + + with self.assertRaisesRegex( + ValueError, r"Method ``requires_grad_`` not supported for RemoteModule" + ): + remote_module.requires_grad_() + with self.assertRaisesRegex( + ValueError, r"Method ``zero_grad`` not supported for RemoteModule" + ): + remote_module.zero_grad() + with self.assertRaisesRegex( + ValueError, r"Method ``share_memory`` not supported for RemoteModule" + ): + remote_module.share_memory() + with self.assertRaisesRegex( + ValueError, r"Method ``extra_repr`` not supported for RemoteModule" + ): + remote_module.extra_repr() + + @dist_utils.dist_init + def test_send_remote_module_with_a_new_attribute_not_pickled_over_the_wire(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + # If a new attribute is added to this RemoteModule after the initialization, + # and it will be sent over the wire by RPC, + # this new field will not be pickled, because it's not specified in _REMOTE_MODULE_PICKLED_ATTRIBUTES. + # Note that adding a new attribute out of constructor should rarely happen. + # If a new attribute is added to RemoteModule constructor, + # there is a sanity check to enforce developers to add this attribute to either + # _REMOTE_MODULE_PICKLED_ATTRIBUTES or _REMOTE_MODULE_ATTRIBUTES_IGNORE_FOR_PICKLING. + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + new_attr_name = "new_attr" + setattr(remote_module, new_attr_name, 1) + + attrs = rpc.rpc_sync( + dst_worker_name, remote_module_attributes, (remote_module,) + ) + self.assertNotIn(new_attr_name, attrs) + + @dist_utils.dist_init + def test_remote_module_py_pickle_not_supported(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + with TemporaryFileName() as fname: + with self.assertRaisesRegex( + RuntimeError, + "Cannot pickle RemoteModule in python pickler. RemoteModule can only be pickled when using RPC", + ): + torch.save(remote_module, fname) + + @dist_utils.dist_init + def test_remote_module_py_pickle_not_supported_script(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + for remote_module in self._create_remote_module_iter( + dst_worker_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE] + ): + with ( + TemporaryFileName() as fname, + self.assertRaisesRegex( + torch.jit.Error, "can only be pickled when using RPC" + ), + ): + torch.save(remote_module, fname) + + +class ThreeWorkersRemoteModuleTest(CommonRemoteModuleTest): + @property + def world_size(self): # Override setting in CommonRemoteModuleTest + return 3 + + @dist_utils.dist_init + def test_send_remote_module_over_the_wire(self): + if self.rank != 0: + return + dst_worker1_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + dst_worker2_name = dist_utils.worker_name((self.rank + 2) % self.world_size) + + # Unpickled attributes include both the inherent attributes of RemoteModule + # (not inherited from the superclass) and two installed methods. + expected_unpickled_attrs = list(_REMOTE_MODULE_PICKLED_ATTRIBUTES) + expected_unpickled_attrs.append("forward_async") + expected_unpickled_attrs.append("forward") + + # Create a remote module on worker1 and then pass it to worker2 over the RPC layer. + for remote_module in self._create_remote_module_iter( + dst_worker1_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + # Test querying some simple attributes from worker2. + attrs = rpc.rpc_sync( + dst_worker2_name, remote_module_attributes, (remote_module,) + ) + self.assertListEqual(list(attrs.keys()), expected_unpickled_attrs) + self.assertEqual(attrs["on"], "worker1") + self.assertEqual(attrs["device"], "cpu") + self.assertFalse(attrs["is_device_map_set"]) + self.assertFalse(attrs["is_scriptable"]) + + # Test the installed methods on worker1's can be initiated by worker2 over RPC layer. + # NOTE: In practice a remote module should be directly stored on the worker that runs ``forward``` or ``forward_async``, + # not have another worker to initiate forward over the RPC layer. + args = (torch.ones(1), 2, "3") + ret1 = rpc.rpc_sync(dst_worker2_name, remote_forward, (remote_module, args)) + self.assertEqual(ret1, tuple(reversed(args))) + ret2 = rpc.rpc_sync( + dst_worker2_name, remote_forward_async, (remote_module, args) + ) + self.assertEqual(ret2, tuple(reversed(args))) + + @dist_utils.dist_init + def test_send_remote_module_over_the_wire_script_not_supported(self): + if self.rank != 0: + return + dst_worker1_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + dst_worker2_name = dist_utils.worker_name((self.rank + 2) % self.world_size) + + # Unpickled attributes include both the inherent attributes of RemoteModule + # (not inherited from the superclass) and two installed methods. + expected_unpickled_attrs = list(_REMOTE_MODULE_PICKLED_ATTRIBUTES) + expected_unpickled_attrs.append("forward_async") + expected_unpickled_attrs.append("forward") + + with self.assertRaisesRegex( + RuntimeError, "Passing a script RemoteModule over RPC is not supported." + ): + # Create a remote module on worker1 and then pass it to worker2 over the RPC layer. + for remote_module in self._create_remote_module_iter( + dst_worker1_name, modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE] + ): + # Test querying some simple attributes from worker2. + rpc.rpc_sync( + dst_worker2_name, remote_module_attributes, (remote_module,) + ) + + @dist_utils.dist_init + def test_create_remote_module_from_module_rref(self): + if self.rank != 0: + return + dst_worker1_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + dst_worker2_name = dist_utils.worker_name((self.rank + 2) % self.world_size) + + # Create a remote module on worker1 and then pass its `module_rref` to worker2 over the RPC layer. + for remote_module in self._create_remote_module_iter( + dst_worker1_name, modes=[ModuleCreationMode.MODULE_CTOR] + ): + remote_module2 = rpc.rpc_sync( + dst_worker2_name, + RemoteModule.init_from_module_rref, + (dst_worker2_name, remote_module.get_module_rref()), + ) + + args = (torch.ones(1), 2, "3") + ret1 = rpc.rpc_sync(dst_worker1_name, remote_forward, (remote_module, args)) + ret2 = rpc.rpc_sync( + dst_worker2_name, remote_forward, (remote_module2, args) + ) + self.assertEqual(ret1, ret2) + + +class CudaRemoteModuleTest(CommonRemoteModuleTest): + @skip_if_lt_x_gpu(1) + @dist_utils.dist_init + def test_valid_device(self): + if self.rank != 0: + return + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = dist_utils.worker_name(dst_rank) + + for remote_module in self._create_remote_module_iter( + f"{dst_worker_name}/cuda:0", modes=[ModuleCreationMode.MODULE_CTOR] + ): + device = rpc.rpc_sync( + dst_worker_name, remote_device, (remote_module.module_rref,) + ) + self.assertEqual(device.type, "cuda") + self.assertEqual(device.index, 0) + + # Test rank works as well. + for remote_module in self._create_remote_module_iter( + f"rank:{dst_rank}/cuda:0", modes=[ModuleCreationMode.MODULE_CTOR] + ): + device = rpc.rpc_sync( + dst_worker_name, remote_device, (remote_module.module_rref,) + ) + self.assertEqual(device.type, "cuda") + self.assertEqual(device.index, 0) + + @skip_if_lt_x_gpu(1) + @dist_utils.dist_init + def test_invalid_devices(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + with self.assertRaisesRegex( + RuntimeError, + r"Expected one of .+ device type at start of device string", + ): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/foo", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + if TEST_WITH_ROCM: + errorString = ( + r"HIP error: invalid device ordinal\n" + r"HIP kernel errors might be asynchronously reported at some other API call, " + r"so the stacktrace below might be incorrect.\n" + r"For debugging consider passing AMD_SERIALIZE_KERNEL=3" + ) + else: + errorString = r"CUDA error: invalid device ordinal" + with self.assertRaisesRegex(RuntimeError, errorString): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/cuda:100", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex(RuntimeError, r"Invalid device string: 'cpu2'"): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/cpu2", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex(RuntimeError, r"Device string must not be empty"): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex( + ValueError, + r"Could not parse remote_device: worker1/cuda:0/cuda:1. The valid format is '/'", + ): + [ + m.forward() + for m in self._create_remote_module_iter( + f"{dst_worker_name}/cuda:0/cuda:1", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex( + ValueError, + r"Could not parse remote_device: /. The valid format is '/'", + ): + [ + m.forward() + for m in self._create_remote_module_iter( + "/", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + with self.assertRaisesRegex( + ValueError, + r"Could not parse remote_device: /cuda:0. The valid format is '/'", + ): + [ + m.forward() + for m in self._create_remote_module_iter( + "/cuda:0", + modes=[ModuleCreationMode.MODULE_CTOR], + ) + ] + + @skip_if_lt_x_gpu(1) + @dist_utils.dist_init + def test_input_moved_to_cuda_device(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + # These two CPU tensors (in args and kwargs) should be implicitly moved to an appropriate cuda device. + t1 = torch.ones(1) + args = (t1, 2) + t2 = t1 * 2 + kwargs = dict(word=t2) + + # Only test Python nn.Module, because script module methods don't support taking kwargs. + for remote_module in self._create_remote_module_iter( + f"{dst_worker_name}/cuda:0", modes=[ModuleCreationMode.MODULE_CTOR] + ): + ret_fut = remote_module.forward_async(*args, **kwargs) + ret = ret_fut.wait() + self.assertEqual(ret, tuple(reversed(args + (t2,)))) + # TODO: Once the RPC backend can support directly sending GPU tensors, the expected device type should be "cuda:0". + self.assertEqual(ret[0].device.type, "cpu") + self.assertEqual(ret[2].device.type, "cpu") + + ret = remote_module.forward(*args, **kwargs) + self.assertEqual(ret, tuple(reversed(args + (t2,)))) + # TODO: Once the RPC backend can support directly sending GPU tensors, the expected device type should be "cuda:0". + self.assertEqual(ret[0].device.type, "cpu") + self.assertEqual(ret[2].device.type, "cpu") + + @skip_if_lt_x_gpu(1) + @dist_utils.dist_init + def test_input_moved_to_cuda_device_script(self): + if self.rank != 0: + return + dst_worker_name = dist_utils.worker_name((self.rank + 1) % self.world_size) + + scripted_remote_module = next( + self._create_remote_module_iter( + f"{dst_worker_name}/cuda:0", + modes=[ModuleCreationMode.MODULE_CTOR_WITH_INTERFACE], + ) + ) + + @torch.jit.script + def run_forward(scripted_remote_module: MyModuleInterface): + ret = scripted_remote_module.forward(torch.ones(1), 2, "3") + return ret + + ret = run_forward(scripted_remote_module) + + self.assertEqual(ret, ("3", 2, torch.ones(1))) + # TODO: Once the RPC backend can support directly sending GPU tensors, the expected device type should be "cuda:0". + self.assertEqual(ret[2].device.type, "cpu") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/dist_autograd_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/dist_autograd_test.py new file mode 100644 index 0000000000000000000000000000000000000000..1c5392e2b6f7d01c0e358a1faaac17409371c968 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/dist_autograd_test.py @@ -0,0 +1,2767 @@ +# mypy: allow-untyped-defs + +import random +import sys +import threading +import time +from datetime import timedelta +from enum import Enum + +import torch +import torch.distributed as dist +import torch.distributed.autograd as dist_autograd +import torch.distributed.rpc as rpc +import torch.nn as nn +import torch.testing._internal.dist_utils +from torch.autograd import Function +from torch.autograd.function import once_differentiable +from torch.distributed.rpc import RRef +from torch.testing._internal.common_distributed import skip_if_lt_x_gpu +from torch.testing._internal.common_utils import ( + IS_MACOS, + skip_but_pass_in_sandcastle_if, +) +from torch.testing._internal.dist_utils import ( + dist_init, + initialize_pg, + wait_until_node_failure, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +# Right now we test up to 3-layer nested rpc calls. +# rpc_done[1] and ctx_ids[1] represent rpc is done in prev rank, and context id +# sent from prev rank respectively. +# rpc_done[2] and ctx_ids[2] represents for prev of prev rank. +# rpc_done[3] and ctx_ids[3] represents for prev of prev of prev rank. +# rpc_done[0] and ctx_ids[0] represents for current rank, but mostly not used. +rpc_done = [False, False, False, False] +ctx_ids = [-1, -1, -1, -1] + +known_context_ids = set() + +requires_grad_tensor = torch.ones(3, 3, requires_grad=True) + + +# Send rpc done info and context_id to +# dst_rank = (self.rank + rank_distance) % self.world_size +# we don't need a lock here since the GIL is held while executing remote +# python UDFs, so access is serialized across several workers. +def _set_rpc_done(ctx_id, rank_distance): + global rpc_done + global ctx_ids + global known_context_ids + rpc_done[rank_distance] = True + ctx_ids[rank_distance] = ctx_id + known_context_ids.add(ctx_id) + + +def _check_rpc_done(rank_distance): + while not rpc_done[rank_distance]: + time.sleep(0.1) + + +def _torch_ones(sizes, requires_grad=False): + return torch.ones(sizes, requires_grad=requires_grad) + + +# This method must be called on the rref owner, and verifies that the grad of +# rref tensor equals to the given grad. +def _compare_owner_value(context_id, rref, grad): + grads = dist_autograd.get_gradients(context_id) + x = grads[rref.local_value()] + if x.is_sparse: + if not grad.is_sparse: + raise AssertionError("Expected grad to be sparse") + x = x.to_dense() + grad = grad.to_dense() + else: + if grad.is_sparse: + raise AssertionError("Expected grad to not be sparse") + return torch.equal(x, grad) + + +def create_tensor(): + return torch.ones((3, 3), requires_grad=True) + + +def build_sparse_tensor(coalesce=False, requires_grad=True, dtype=torch.float32): + i = [[0, 1, 1], [2, 0, 2]] + v = [3.2, 4.1, 5.3] + tensor = torch.sparse_coo_tensor( + i, v, (3, 3), requires_grad=requires_grad, dtype=dtype + ) + if coalesce: + tensor = tensor.coalesce() + return tensor + + +@torch.jit.script +def create_torchscript_tensor() -> torch.Tensor: + return torch.ones((3, 3)).requires_grad_() + + +def my_py_add(t1, t2): + return torch.add(t1, t2) + + +def my_scalar_add(a, b): + return a + b + + +def my_rref_add(rref_t1, t2): + ret = torch.add(rref_t1.local_value(), t2) + return ret + + +@torch.jit.script +def my_script_add(t1, t2): + return torch.add(t1, t2) + + +@torch.jit.script +def my_script_ref_add(ref_t1: RRef[torch.Tensor], t2: torch.Tensor) -> torch.Tensor: + t1 = ref_t1.to_here() + return torch.add(t1, t2) + + +def my_nested_rref_add(dst, rref_t1, t2): + return rpc.rpc_sync(dst, my_rref_add, args=(rref_t1, t2)) + + +def ret_requires_grad(): + return requires_grad_tensor + + +def my_py_nested_call(t1, t2, dst, world_size, hops): + next_dst = (dst + 1) % world_size + if hops > 0: + return rpc.rpc_sync( + worker_name(next_dst), + my_py_nested_call, + args=(t1, t2, next_dst, world_size, hops - 1), + ) + else: + return rpc.rpc_sync(worker_name(next_dst), my_py_add, args=(t1, t2)) + + +# after dist autograd context is cleaned up, it should be cleaned up on other +# nodes. This helper allows timeout_seconds for those RPCs to be completed, and +# ensures that all the contexts have been cleaned up in that timeframe.any +def _all_contexts_cleaned_up(timeout_seconds=10): + global known_context_ids + start = time.time() + context_id_to_raised = set() + while ( + time.time() - start < timeout_seconds + and context_id_to_raised != known_context_ids + ): + for context_id in known_context_ids: + try: + dist_autograd._retrieve_context(context_id) + except RuntimeError: + context_id_to_raised.add(context_id) + # all contexts have been cleaned up if trying to retrieve any context resulted in a RuntimeError. + success = context_id_to_raised == known_context_ids + return success + + +# This function creates a dis autograd context, run rpc_sync on the given ps, +# and then blocks until the ps has verified the grads are correctly accumulated. +def _run_trainer(rref_t1, t2, ps, rank_diff, sparse): + with dist_autograd.context() as context_id: + ret = rpc.rpc_sync(ps, my_rref_add, args=(rref_t1, t2)) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + dist_autograd.backward(context_id, [loss]) + # prevent deleting dist autograd context + rpc.rpc_sync(ps, _set_rpc_done, args=(context_id, rank_diff)) + rpc.rpc_sync(ps, _check_rpc_done, args=(0,)) + + +# This function is the same as _run_trainer, except rpc calls torchscript +# function "my_script_ref_add" instead of python function "my_rref_add" +def _run_trainer_torchscript(rref_t1, t2, ps, rank_diff, sparse): + with dist_autograd.context() as context_id: + ret = rpc.rpc_sync(ps, my_script_ref_add, args=(rref_t1, t2)) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + dist_autograd.backward(context_id, [loss]) + # prevent deleting dist autograd context + rpc.rpc_sync(ps, _set_rpc_done, args=(context_id, rank_diff)) + rpc.rpc_sync(ps, _check_rpc_done, args=(0,)) + + +class SimulateBackwardError(Function): + _simulate_error = True + + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + @once_differentiable + def backward(ctx, input): + if SimulateBackwardError._simulate_error: + raise Exception("Simulate error on backward pass") # noqa: TRY002 + else: + return input + + +class ExecMode(Enum): + LOCAL = 1 # Run the operation locally. + RPC_SYNC = 2 # Run the operation using rpc_sync + REMOTE = 3 # Run the operation using remote. + RPC_ASYNC = 4 # Run the operation using rpc_async + + +# Common utils for both CPU and CUDA test suites +class CommonDistAutogradTest(RpcAgentTestFixture): + def _exec_func_with_dst(self, dst, exec_mode, method, *args): + if ExecMode.LOCAL == exec_mode: + if len(args) == 1 and isinstance(args[0], list): + return method(*args[0]) + return method(*args) + elif ExecMode.RPC_SYNC == exec_mode: + return rpc.rpc_sync(worker_name(dst), method, args=(args)) + elif ExecMode.REMOTE == exec_mode: + return rpc.remote(worker_name(dst), method, args=(args)).to_here() + elif ExecMode.RPC_ASYNC == exec_mode: + fut = rpc.rpc_async(worker_name(dst), method, args=(args)) + return fut.wait() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + def _exec_func(self, exec_mode, method, *args): + return self._exec_func_with_dst(self._next_rank(), exec_mode, method, *args) + + def _next_rank(self): + if hasattr(self, "dst_rank"): + self.dst_rank = (self.dst_rank + 1) % self.world_size + if self.dst_rank == self.rank: + return self._next_rank() + else: + self.dst_rank = (self.rank + 1) % self.world_size + return self.dst_rank + + def _check_rpc_done(self, rank_distance): + _check_rpc_done(rank_distance) + + def _verify_backwards(self, exec_mode, tensors, context_id, local_grads, *args): + if exec_mode == ExecMode.LOCAL: + torch.autograd.backward(tensors) + return [arg.grad for arg in args] + else: + self._verify_backwards_remote(tensors, context_id, local_grads, *args) + + def _verify_backwards_remote(self, tensors, context_id, local_grads, *args): + dist_autograd.backward(context_id, tensors) + + # Verify grads were accumulated appropriately. + grads = dist_autograd.get_gradients(context_id) + nargs = len(args) + ngrads = 0 + for i in range(nargs): + if local_grads[i] is not None: + self.assertIn(args[i], grads) + self.assertEqual(local_grads[i], grads[args[i]]) + ngrads += 1 + else: + self.assertNotIn(args[i], grads) + + self.assertEqual(ngrads, len(grads)) + + def _test_graph(self, fn, exec_mode, sparse): + dst_rank = (self.rank + 1) % self.world_size + + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + if sparse: + t1 = build_sparse_tensor() + t2 = build_sparse_tensor() + else: + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync(worker_name(dst_rank), fn, args=(t1, t2)) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote(worker_name(dst_rank), fn, args=(t1, t2)).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + rpc.rpc_sync(worker_name(dst_rank), _set_rpc_done, args=(context_id, 1)) + + # Verify graph for current context id. + ctx = dist_autograd._current_context() + self.assertEqual(context_id, ctx._context_id()) + send_functions = ctx._send_functions() + self.assertEqual(1, len(send_functions)) + recv_functions = ctx._recv_functions() + self.assertEqual(1, len(recv_functions)) + self._verify_graph_for_first_rpc_call( + next(iter(send_functions.values())), + next(iter(recv_functions.values())), + t1, + t2, + ret, + ) + + # Wait for the prev rank to be done with rpc. + self._check_rpc_done(1) + # Verify graph for previous context id. + ctx = dist_autograd._retrieve_context(ctx_ids[1]) + send_functions = ctx._send_functions() + self.assertEqual(1, len(send_functions)) + self._verify_graph_for_rpc_call_exec(next(iter(send_functions.values()))) + # this barrier is needed so one worker does not clean up their + # autograd context before another worker tries to access it. + dist.barrier() + + # autograd context should be cleaned up by now. + with self.assertRaises(RuntimeError): + ctx = dist_autograd._retrieve_context(context_id) + + # No autograd context available. + with self.assertRaises(RuntimeError): + ctx = dist_autograd._current_context() + + # 3-layer nested calls + def _test_graph_for_py_nested_call(self, exec_mode, sparse): + dst_rank = (self.rank + 1) % self.world_size + + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + if sparse: + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + else: + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + worker_name(dst_rank), + my_py_nested_call, + args=(t1, t2, dst_rank, self.world_size, 1), + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + worker_name(dst_rank), + my_py_nested_call, + args=(t1, t2, dst_rank, self.world_size, 1), + ).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + # Barrier to ensure all RPCs are done. + dist.barrier() + + for rd in [1, 2, 3]: + rpc.rpc_sync( + worker_name((self.rank + rd) % self.world_size), + _set_rpc_done, + args=(context_id, rd), + ) + + # Barrier to ensure all set_rpc_done have completed. + dist.barrier() + + # For self.rank, it has 4 graphs to verify + # One is for current context id when this rank send first rpc call. + # Second one is for prev context id when this rank make 1st nested + # call. + # Third one is for prev prev context id when this rank make + # 2nd nested call. + # Last one is for prev prev prev context id when this rank + # execute the torch.add() operator. + + # Verify first graph for current context id. + ctx = dist_autograd._current_context() + self.assertEqual(context_id, ctx._context_id()) + send_functions = ctx._send_functions() + self.assertEqual(1, len(send_functions)) + recv_functions = ctx._recv_functions() + self.assertEqual(1, len(recv_functions)) + self._verify_graph_for_first_rpc_call( + next(iter(send_functions.values())), + next(iter(recv_functions.values())), + t1, + t2, + ret, + ) + + # Verify second graph for 1st nested call. + ctx = dist_autograd._retrieve_context(ctx_ids[1]) + self._verify_graph_for_nested_rpc_call(ctx) + + # Verify third graph for 2nd nested call. + ctx = dist_autograd._retrieve_context(ctx_ids[2]) + self._verify_graph_for_nested_rpc_call(ctx) + + # verify last graph for rpc call execution. + ctx = dist_autograd._retrieve_context(ctx_ids[3]) + send_functions = ctx._send_functions() + self.assertEqual(1, len(send_functions)) + self._verify_graph_for_rpc_call_exec(next(iter(send_functions.values()))) + # this barrier is needed so one worker does not clean up their + # autograd context before another worker tries to access it. + dist.barrier() + + # Rank0->Rank1->Rank0 + def _test_graph_for_py_nested_call_itself(self, exec_mode, sparse): + dst_rank = (self.rank + 1) % self.world_size + + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + if sparse: + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + else: + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync( + worker_name(dst_rank), + my_py_nested_call, + args=( + t1, + t2, + (self.rank - 1 + self.world_size) % self.world_size, + self.world_size, + 0, + ), + ) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + worker_name(dst_rank), + my_py_nested_call, + args=( + t1, + t2, + (self.rank - 1 + self.world_size) % self.world_size, + self.world_size, + 0, + ), + ).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + rpc.rpc_sync( + worker_name((self.rank + 1) % self.world_size), + _set_rpc_done, + args=(context_id, 1), + ) + + # For self.rank, it has 2 graphs to verify. + # One is for current context id when this rank send first rpc + # call and execute the torch.add() operator. + # Another one is for prev context id when this rank make + # nested call. + ctx = dist_autograd._current_context() + self.assertEqual(context_id, ctx._context_id()) + send_functions = ctx._send_functions() + self.assertEqual(2, len(send_functions)) + recv_functions = ctx._recv_functions() + self.assertEqual(2, len(recv_functions)) + self._verify_graph_for_first_rpc_call( + next(iter(send_functions.values())), + list(recv_functions.values())[1], + t1, + t2, + ret, + ) + self._verify_graph_for_rpc_call_exec(list(send_functions.values())[1]) + + # Verify two pairs of send and recv functions for nested + # call + self._check_rpc_done(1) + ctx = dist_autograd._retrieve_context(ctx_ids[1]) + self._verify_graph_for_nested_rpc_call(ctx) + # this barrier is needed so one worker does not clean up their + # autograd context before another worker tries to access it. + dist.barrier() + + def _test_no_graph_with_tensors_not_require_grad(self, exec_mode, sparse): + initialize_pg(self.file_init_method, self.rank, self.world_size) + dst_rank = (self.rank + 1) % self.world_size + with dist_autograd.context() as context_id: + if sparse: + t1 = build_sparse_tensor(requires_grad=False) + t2 = build_sparse_tensor(requires_grad=False) + else: + t1 = torch.ones(3, 3, requires_grad=False) + t2 = torch.zeros(3, 3, requires_grad=False) + if ExecMode.RPC_SYNC == exec_mode: + rpc.rpc_sync(worker_name(dst_rank), torch.add, args=(t1, t2)) + elif ExecMode.REMOTE == exec_mode: + rpc.remote(worker_name(dst_rank), torch.add, args=(t1, t2)).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + rpc.rpc_sync(worker_name(dst_rank), _set_rpc_done, args=(context_id, 1)) + + ctx = dist_autograd._current_context() + send_functions = ctx._send_functions() + self.assertEqual(len(send_functions), 0) + recv_functions = ctx._recv_functions() + self.assertEqual(len(recv_functions), 0) + + # Wait for the prev rank to be done with rpc. + self._check_rpc_done(1) + # NB: RRef.to_here() always passes the autograd context to the + # the callee, as the caller does not know whether the return + # value would contain a requires_grad tensor or not. + # + # rpc/remote with udf (_set_rpc_done here) also always passes the + # autograd context to the callee due to the same reason. + self.assertNotEqual(-1, dist_autograd._retrieve_context(ctx_ids[1])) + dist.barrier() + + def _test_rpc_complex_args(self, exec_mode, sparse): + with dist_autograd.context(): + num_tensors = 10 + tensors = [] + for i in range(num_tensors): + if sparse: + tensor = build_sparse_tensor(requires_grad=(i % 2 == 0)) + else: + tensor = torch.ones(3, 3, requires_grad=(i % 2 == 0)) + tensors.append(tensor) + dst_rank = self._next_rank() + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync(worker_name(dst_rank), torch.stack, args=(tensors,)) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote( + worker_name(dst_rank), torch.stack, args=(tensors,) + ).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + self.assertEqual(torch.stack(tensors), ret) + + # Verify appropriate tensors have been attached the autograd graph. + next_funcs = next( + iter(dist_autograd._current_context()._send_functions().values()) + ).next_functions + for i in range(len(next_funcs)): + self.assertEqual( + "torch::autograd::AccumulateGrad", next_funcs[i][0].name() + ) + self.assertEqual(tensors[i], next_funcs[i][0].variable) + + # Verify that the worker id has been recorded in the context + ctx = dist_autograd._current_context() + worker_ids = ctx._known_worker_ids() + self.assertEqual(len(worker_ids), 1) + self.assertEqual(worker_ids, {dst_rank}) + + def context_cleanup_test_helper(self, rpc_args, func, nested=False): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + # test that in dist autograd, in the case that tensors communicated over RPC do + # NOT require grad, we still cleanup the dist autograd contexts created + # on other nodes. This is because the autograd context is still + # communicated over RPC even if tensor arguments do not require grad, as + # it is possible that the response could. + if nested: + dst_rank = (self.rank + 1) % self.world_size + nested_dst_rank = (dst_rank + 1) % self.world_size + dst_ranks = {dst_rank} + else: + dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} + + with dist_autograd.context() as context_id: + for dst_rank in dst_ranks: + rpc.rpc_sync(worker_name(dst_rank), func, args=rpc_args) + rpc.rpc_sync(worker_name(dst_rank), _set_rpc_done, args=(context_id, 1)) + if nested: + rpc.rpc_sync( + worker_name(nested_dst_rank), + _set_rpc_done, + args=(context_id, 2), + ) + # the thread's context id should be cleaned up + with self.assertRaises(RuntimeError): + dist_autograd._retrieve_context(context_id) + # Ensure all peers have finished mutating the + # `known_context_ids` set. + dist.barrier() + # check that all contexts have been cleaned up. + success = _all_contexts_cleaned_up() + self.assertTrue(success) + + def _backward_no_grad_on_tensor(self, t1, t2, sparse): + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + worker_name(self._next_rank()), torch.add, args=(t1, t2) + ) + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + dist_autograd.backward(context_id, [loss], retain_graph=True) + self.assertIsNone(t1.grad) + self.assertIsNone(t2.grad) + + # Now populate .grad with local autograd engine and + # verify dist autograd doesn't mess with it. + loss_local = torch.add(t1, t2) + if sparse: + loss_local = torch.sparse.sum(loss_local) + else: + loss_local = loss_local.sum() + loss_local.backward() + self.assertIsNotNone(t1.grad) + self.assertIsNotNone(t2.grad) + + t1_grad_before = t1.grad + t2_grad_before = t2.grad + dist_autograd.backward(context_id, [loss]) + self.assertEqual(t1_grad_before, t1.grad) + self.assertEqual(t2_grad_before, t2.grad) + + # The current rank first creates a tensor on the rref_owner, and then passes + # the rref with another tensor to the callee to run either my_rref_add or + # my_nested_rref_add, depending on whether the callee is the rref owner. + # The grad of tensor lives on the current rank, and the grad of the rref + # tensor lives on the rref owner. + def _backward_rref(self, callee, rref_owner, t1, t2, local_grads, sparse): + local_ret = torch.add(t1, t2) + if sparse: + local_ret = torch.sparse.sum(local_ret) + else: + local_ret = local_ret.sum() + local_ret.backward() + with dist_autograd.context() as context_id: + if sparse: + rref_t1 = rpc.remote( + rref_owner, + build_sparse_tensor, + args=( + False, + True, + ), + ) + else: + rref_t1 = rpc.remote( + rref_owner, + _torch_ones, + args=((3, 3),), + kwargs={"requires_grad": True}, + ) + if callee == rref_owner: + rref = rpc.remote(callee, my_rref_add, args=(rref_t1, t2)) + else: + rref = rpc.remote( + callee, my_nested_rref_add, args=(rref_owner, rref_t1, t2) + ) + ret = rref.to_here() + if sparse: + ret = torch.sparse.sum(ret) + else: + ret = ret.sum() + dist_autograd.backward(context_id, [ret]) + + # verify grads on caller + grads = dist_autograd.get_gradients(context_id) + self.assertIn(t2, grads) + self.assertEqual(grads[t2], t2.grad) + + # verify grads on rref owner + self.assertTrue( + rpc.rpc_sync( + rref_owner, + _compare_owner_value, + args=(context_id, rref_t1, t1.grad), + ) + ) + + # In this test, every rank will serve as a parameter server (ps) and a + # driver, and then kicks off trainers on the other three ranks. So, we have: + # ps = rank0 with trainers = rank1/2/3 + # ps = rank2 with trainers = rank2/3/0 + # ps = rank3 with trainers = rank3/0/1 + # ps = rank4 with trainers = rank0/1/2 + # + # These four test ps-trainer groups run on completely separate autograd + # graphs, but they share the same set of underlying RpcAgents. + def _test_trainer_ps(self, create_ref_fn, trainer_fn, sparse): + if sparse: + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + else: + t1 = torch.ones((3, 3), requires_grad=True) + t2 = torch.zeros((3, 3), requires_grad=True) + + local_ret = torch.add(t1, t2) + if sparse: + torch.sparse.sum(local_ret).backward() + else: + local_ret.sum().backward() + + # create rref on self + rref_t1 = rpc.remote(worker_name(self.rank), create_ref_fn, args=()) + + # kick off forward and backward pass on three other workers (trainers) + rank_diffs = [1, 2, 3] + futures = [ + rpc.rpc_async( + worker_name((self.rank + rank_diff) % self.world_size), + trainer_fn, + args=(rref_t1, t2, worker_name(self.rank), rank_diff, sparse), + ) + for rank_diff in rank_diffs + ] + + # check if the trainers have done with their backward pass + for rank_diff in rank_diffs: + self._check_rpc_done(rank_diff) + + # trainers are done and holding the context for verification + for rank_diff in rank_diffs: + # make sure grads are accumulated for the same tensors and values + # are all correct + ctx_id = ctx_ids[rank_diff] + grads = dist_autograd.get_gradients(ctx_id) + local_t1 = rref_t1.to_here() + self.assertIn(local_t1, grads) + self.assertEqual(grads[local_t1], t1.grad) + + # unblock trainers + _set_rpc_done(None, 0) + + # wait until all trainers are done + torch.futures.wait_all(futures) + + def _backward_multiple_round_trips(self, t1, t2, t3, t4, t5, local_grads, sparse): + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + # Multiple RPCs between different nodes. + val = self._exec_func(exec_mode, torch.add, t1, t2) + val = self._exec_func(exec_mode, torch.mul, t3, val) + s1 = self._exec_func(exec_mode, torch.stack, (t4, val)) + s2 = self._exec_func(exec_mode, torch.stack, (t5, val)) + if sparse: + val = self._exec_func(exec_mode, torch.mul, s1, s2) + val = self._exec_func(exec_mode, torch.mul, val, val) + loss = torch.sparse.sum(val) + else: + val = self._exec_func(exec_mode, torch.bmm, s1, s2) + val = self._exec_func(exec_mode, torch.matmul, val, val) + loss = val.sum() + + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2, t3, t4, t5 + ) + local_grads = ret if ret else local_grads + + def _backward_different_dtypes(self, t1, t2, sparse): + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + loss = self._exec_func(exec_mode, torch.add, t1, t2) + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + local_grads = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + + # Run the same code locally and with dist autograd and verify gradients + # are same. + def _backward_simple_python_udf(self, t1, t2, sparse): + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func(exec_mode, my_py_add, t1, t2) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + local_grads = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + + # Run the same code locally and with dist autograd and verify gradients + # are same. + def _backward_simple_script_call(self, t1, t2, sparse): + local_grads = None + for exec_mode in [ + ExecMode.LOCAL, + ExecMode.RPC_SYNC, + ExecMode.RPC_ASYNC, + ExecMode.REMOTE, + ]: + with dist_autograd.context() as context_id: + forward_ret = self._exec_func(exec_mode, my_script_add, t1, t2) + if sparse: + loss = torch.sparse.sum(forward_ret) + else: + loss = forward_ret.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + def _nested_backward_accumulate_grads(self, t1, t2, sparse): + with dist_autograd.context() as context_id: + ret = rpc.rpc_sync( + worker_name(self._next_rank()), + DistAutogradTest._test_nested_backward_accumulate_grads, + args=(t1, t2, self._next_rank()), + ) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + # Run backward twice. + dist_autograd.backward(context_id, [loss], retain_graph=True) + dist_autograd.backward(context_id, [loss]) + + def _backwards_nested_python_udf(self, t1, t2, sparse): + t3 = t1 * t2 + t4 = t1 + t2 + res = t3 + t4 + loss = t1 * t2 * t3 * t4 * res + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + torch.autograd.backward([loss]) + + # Now run distributed autograd. + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + worker_name(self._next_rank()), + DistAutogradTest._nested_python_udf, + args=(t1, t2, self._next_rank()), + ) + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + dist_autograd.backward(context_id, [loss]) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(t1.grad, grads[t1]) + self.assertEqual(t2.grad, grads[t2]) + + def _mixed_requires_grad(self, t1, t2, sparse): + for exec_mode in [ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func( + exec_mode, DistAutogradTest._mixed_requires_grad_operaton, t1, t2 + ) + self.assertEqual(t1 * t2, ret) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + dist_autograd.backward(context_id, [loss]) + self.assertTrue(t1.requires_grad) + self.assertFalse(t2.requires_grad) + grads = dist_autograd.get_gradients(context_id) + self.assertIn(t1, grads) + self.assertNotIn(t2, grads) + self.assertEqual(t2, grads[t1]) + + def _multiple_backward(self, t1, t2, sparse): + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + worker_name(self._next_rank()), torch.add, args=(t1, t2) + ) + if sparse: + loss = torch.sparse.sum(loss) + else: + loss = loss.sum() + # Run backward in a loop multiple times. + for _ in range(1000): + dist_autograd.backward(context_id, [loss], retain_graph=True) + + # For current context, this rank sends t1 and t2 tensors to dst_rank, + # then get t3 = torch.add(t1, t2) result tensor. + # For the current context in this rank, it expects graph like this: + # send function: + # rpcSendBackward + # / \ + # t1.AccumulateGrad t2.AccumulateGrad + # + # recv function: + # + # | + # t3.rpcRecvBackward + # + def _verify_graph_for_first_rpc_call( + self, send_function, recv_function, t1, t2, ret + ): + # Retrieve the next functions in the graph. + next_funcs = send_function.next_functions + self.assertEqual(2, len(next_funcs)) + + # We should now hit t1 and t2 in the autograd graph. + self.assertEqual("torch::autograd::AccumulateGrad", next_funcs[0][0].name()) + self.assertEqual(t1, next_funcs[0][0].variable) + self.assertEqual(0, next_funcs[0][1]) + self.assertEqual("torch::autograd::AccumulateGrad", next_funcs[1][0].name()) + self.assertEqual(t2, next_funcs[1][0].variable) + self.assertEqual(0, next_funcs[1][1]) + + # Test recv functions. + self.assertEqual(ret.grad_fn, recv_function) + + # Run the same code locally and with dist autograd and verify gradients + # are same. + def _backward_simple(self, dst, t1, t2, local_grads, sparse): + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func_with_dst(dst, exec_mode, torch.add, t1, t2) + if sparse: + loss = torch.sparse.sum(ret) + else: + loss = ret.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + # For a context passed from previous nested chain calls, this rank + # receives two tensors t1 and t2, executes torch.add(t1, t2) and sends + # result tensor t3 back. + # For this context in this rank, it expects graph like this: + # send and recv functions: + # rpcSendBackward + # | + # t3.AddBackward0 + # / \ + # t1.recvRpcBackward t2.recvRpcBackward + def _verify_graph_for_rpc_call_exec(self, send_function): + # Verify next function is AddBackward0 + next_funcs = send_function.next_functions + self.assertEqual(1, len(next_funcs)) + add_backward_fn = next_funcs[0][0] + self.assertEqual("AddBackward0", add_backward_fn.name()) + + # Verify the next two functions are the same recv backward function. + next_funcs = add_backward_fn.next_functions + self.assertEqual(2, len(next_funcs)) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[0][0].name() + ) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[1][0].name() + ) + self.assertEqual(next_funcs[0][0], next_funcs[1][0]) + + # For a context passed from previous nested chain calls, this rank + # receives two tensors t1 and t2, forwards t1 and t2 tensors using + # nested rpc call to next dst. In return route, receive result tensor t3 + # from next dst and forwarding t3 back to previous calls. + # For this context in this rank, it expects graph like this: + # send and recv functions for receiving and forwarding t1 and t2: + # rpcSendBackward + # / \ + # t1.recvRpcBackward t2.recvRpcBackward + # send and recv functions for receiving and forwarding t3: + # rpcSendBackward + # | + # t3.recvRpcBackward + def _verify_graph_for_nested_rpc_call(self, ctx): + send_functions = ctx._send_functions() + self.assertEqual(2, len(send_functions)) + + # For send function when making nest rpc call, + # next functions of the send function are two recv functions + # for received two tensors from previous call + next_funcs = next(iter(send_functions.values())).next_functions + self.assertEqual(2, len(next_funcs)) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[0][0].name() + ) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[1][0].name() + ) + self.assertEqual(next_funcs[0][0], next_funcs[1][0]) + + # For send function when returning response to previous call + # next function of the send function is the recv function + # for received tensor result returned from nested call + next_funcs = list(send_functions.values())[1].next_functions + self.assertEqual(1, len(next_funcs)) + self.assertEqual( + "torch::distributed::autograd::RecvRpcBackward", next_funcs[0][0].name() + ) + + +class TensorPipeAgentDistAutogradTest(CommonDistAutogradTest): + # Sparse tests only work with TensorPipeAgent. + @dist_init + def test_graph_for_builtin_call_sparse(self): + self._test_graph(torch.add, ExecMode.RPC_SYNC, True) + + @dist_init + def test_graph_for_python_call_sparse(self): + self._test_graph(my_py_add, ExecMode.RPC_SYNC, True) + + @dist_init + def test_graph_for_builtin_remote_call_sparse(self): + self._test_graph(torch.add, ExecMode.REMOTE, True) + + @dist_init + def test_graph_for_python_remote_call_sparse(self): + self._test_graph(my_py_add, ExecMode.REMOTE, True) + + @dist_init + def test_graph_for_py_nested_call_sparse(self): + self._test_graph_for_py_nested_call(ExecMode.RPC_SYNC, True) + + @dist_init + def test_graph_for_py_nested_remote_call_sparse(self): + self._test_graph_for_py_nested_call(ExecMode.REMOTE, True) + + @dist_init + def test_graph_for_py_nested_call_itself_sparse(self): + self._test_graph_for_py_nested_call_itself(ExecMode.RPC_SYNC, True) + + @dist_init + def test_graph_for_py_nested_remote_call_itself_sparse(self): + self._test_graph_for_py_nested_call_itself(ExecMode.REMOTE, True) + + @dist_init + def test_no_graph_with_tensors_not_require_grad_sparse(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.RPC_SYNC, True) + + @dist_init + def test_no_graph_with_tensors_not_require_grad_remote_sparse(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.REMOTE, True) + + @dist_init + def test_rpc_complex_args_sparse(self): + self._test_rpc_complex_args(ExecMode.RPC_SYNC, True) + + @dist_init + def test_remote_complex_args_sparse(self): + self._test_rpc_complex_args(ExecMode.REMOTE, True) + + @dist_init + def test_context_cleanup_tensor_with_grad_sparse(self): + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + self.context_cleanup_test_helper(rpc_args=(t1, t2), func=torch.add) + + @dist_init + def test_context_cleanup_tensor_no_grad_sparse(self): + t1 = build_sparse_tensor(requires_grad=False) + self.context_cleanup_test_helper(rpc_args=(t1, t1), func=torch.add) + + @dist_init + def test_context_cleanup_nested_rpc_sparse(self): + t1 = build_sparse_tensor(requires_grad=True) + t2 = build_sparse_tensor(requires_grad=True) + dst_rank = (self.rank + 1) % self.world_size + args = (t1, t2, dst_rank, self.world_size, 0) + self.context_cleanup_test_helper( + rpc_args=args, func=my_py_nested_call, nested=True + ) + + @dist_init + def test_backward_no_grad_on_tensor_sparse(self): + self._backward_no_grad_on_tensor( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True, + ) + + @dist_init + def test_backward_simple_sparse(self): + self._backward_simple( + self._next_rank(), + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True, + ) + + @dist_init + def test_backward_simple_self_sparse(self): + self._backward_simple( + self.rank, + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True, + ) + + @dist_init + def test_backward_rref_multi_sparse(self): + if self.rank > 0: + callee = "worker0" + rref_owner = callee + self._backward_rref( + callee, + rref_owner, + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True, + ) + + @dist_init + def test_backward_rref_sparse(self): + callee = worker_name(self._next_rank()) + rref_owner = callee + self._backward_rref( + callee, + rref_owner, + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True, + ) + + @dist_init + def test_backward_rref_nested_sparse(self): + callee = worker_name((self.rank + 1) % self.world_size) + rref_owner = worker_name((self.rank + 2) % self.world_size) + self._backward_rref( + callee, + rref_owner, + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + None, + True, + ) + + @dist_init + def test_trainer_ps_sparse(self): + self._test_trainer_ps(build_sparse_tensor, _run_trainer, True) + + @dist_init + def test_backward_multiple_round_trips_sparse(self): + self._backward_multiple_round_trips( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=False), + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=False), + build_sparse_tensor(requires_grad=True), + None, + True, + ) + + @dist_init + def test_backward_different_dtypes_sparse(self): + self._backward_different_dtypes( + build_sparse_tensor(requires_grad=True, dtype=torch.float32), + build_sparse_tensor(requires_grad=True, dtype=torch.float64), + True, + ) + + @dist_init + def test_backward_simple_python_udf_sparse(self): + self._backward_simple_python_udf( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True, + ) + + @dist_init + def test_backward_simple_script_call_sparse(self): + self._backward_simple_script_call( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True, + ) + + @dist_init + def test_nested_backward_accumulate_grads_sparse(self): + self._nested_backward_accumulate_grads( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True, + ) + + @dist_init + def test_backwards_nested_python_udf_sparse(self): + # Run equivalent of _nested_python_udf locally. + self._backwards_nested_python_udf( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True, + ) + + @dist_init + def test_mixed_requires_grad_sparse(self): + self._mixed_requires_grad( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=False), + True, + ) + + @dist_init + def test_multiple_backward_sparse(self): + self._multiple_backward( + build_sparse_tensor(requires_grad=True), + build_sparse_tensor(requires_grad=True), + True, + ) + + @dist_init + def test_embedding_bag_with_no_grad_tensors(self): + dst = self._next_rank() + remote_embedding = rpc.remote( + worker_name(dst), + torch.nn.EmbeddingBag, + args=(16, 16), + kwargs={"mode": "sum", "sparse": True}, + ) + local_embedding = torch.nn.EmbeddingBag(16, 16, mode="sum", sparse=True) + + input = torch.LongTensor([1, 2, 4, 5, 4, 3, 2, 9]) + # requires_grad = True to record send/recv functions + per_sample_weights = torch.rand((8), requires_grad=True) + offsets = torch.LongTensor([0, 4]) + + local_res = local_embedding(input, offsets, per_sample_weights) + + # Run backward twice. + torch.autograd.backward([local_res.sum()], retain_graph=True) + torch.autograd.backward([local_res.sum()]) + local_grad = local_embedding.weight.grad + + with dist_autograd.context() as context_id: + res = rpc.rpc_sync( + worker_name(dst), + DistAutogradTest._call_remote_embedding, + args=(remote_embedding, input, offsets, per_sample_weights), + ) + + # Run backward twice to test accumulation of sparse gradients. + dist_autograd.backward(context_id, [res.sum()], retain_graph=True) + dist_autograd.backward(context_id, [res.sum()]) + + remote_grad = rpc.rpc_sync( + worker_name(dst), + DistAutogradTest._get_grad, + args=(remote_embedding, context_id), + ) + + self.assertEqual(local_grad, remote_grad) + + +class DistAutogradTest(CommonDistAutogradTest): + @dist_init + def test_autograd_context(self): + # Verify max possible id. + max_auto_increment = 281474976710655 + self.assertEqual( + max_auto_increment + (self.worker_id << 48), dist_autograd._get_max_id() + ) + + context_ids = [] + for _ in range(200): + with dist_autograd.context() as context_id: + self.assertEqual( + context_id, + dist_autograd._retrieve_context(context_id)._context_id(), + ) + # First 16 bits should be worker_id. + self.assertEqual(self.worker_id, context_id >> 48) + context_ids.append(context_id) + + for context_id in context_ids: + with self.assertRaisesRegex( + RuntimeError, + f"Could not find autograd context with id: {context_id}", + ): + dist_autograd._retrieve_context(context_id) + + @dist_init + def test_nested_context(self): + with ( + dist_autograd.context(), + self.assertRaisesRegex( + RuntimeError, "Already have an autograd context id for this thread" + ), + dist_autograd.context(), + ): + pass + + @dist_init + def test_graph_for_builtin_call(self): + self._test_graph(torch.add, ExecMode.RPC_SYNC, False) + + @dist_init + def test_graph_for_python_call(self): + self._test_graph(my_py_add, ExecMode.RPC_SYNC, False) + + @dist_init + def test_graph_for_builtin_remote_call(self): + self._test_graph(torch.add, ExecMode.REMOTE, False) + + @dist_init + def test_graph_for_python_remote_call(self): + self._test_graph(my_py_add, ExecMode.REMOTE, False) + + @dist_init + def test_graph_for_py_nested_call(self): + self._test_graph_for_py_nested_call(ExecMode.RPC_SYNC, False) + + @dist_init + def test_graph_for_py_nested_remote_call(self): + self._test_graph_for_py_nested_call(ExecMode.REMOTE, False) + + @dist_init + def test_graph_for_py_nested_call_itself(self): + self._test_graph_for_py_nested_call_itself(ExecMode.RPC_SYNC, False) + + @dist_init + def test_graph_for_py_nested_remote_call_itself(self): + self._test_graph_for_py_nested_call_itself(ExecMode.REMOTE, False) + + @dist_init + def test_no_graph_with_tensors_not_require_grad(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.RPC_SYNC, False) + + @dist_init + def test_no_graph_with_tensors_not_require_grad_remote(self): + self._test_no_graph_with_tensors_not_require_grad(ExecMode.REMOTE, False) + + def _test_grad_only_on_return_value(self, exec_mode): + initialize_pg(self.file_init_method, self.rank, self.world_size) + dst_rank = (self.rank + 1) % self.world_size + with dist_autograd.context() as context_id: + if ExecMode.RPC_SYNC == exec_mode: + ret = rpc.rpc_sync(worker_name(dst_rank), ret_requires_grad) + elif ExecMode.REMOTE == exec_mode: + ret = rpc.remote(worker_name(dst_rank), ret_requires_grad).to_here() + else: + raise ValueError(f"Unrecognized ExecMode {exec_mode}") + + dist_autograd.backward(context_id, [ret.sum()]) + + rpc.rpc_sync(worker_name(dst_rank), _set_rpc_done, args=(context_id, 1)) + + # Wait for the prev rank to be done with rpc. + self._check_rpc_done(1) + grads = dist_autograd.get_gradients(ctx_ids[1]) + self.assertEqual(1, len(grads)) + self.assertIn(requires_grad_tensor, grads) + self.assertEqual(torch.ones_like(ret), grads[requires_grad_tensor]) + # due to the above get_gradients call, ensure that dist autograd + # contexts aren't cleaned up until all workers exit context managers + dist.barrier() + + @dist_init + def test_grad_only_on_return_value(self): + self._test_grad_only_on_return_value(ExecMode.RPC_SYNC) + + @dist_init + def test_grad_only_on_return_value_remote(self): + self._test_grad_only_on_return_value(ExecMode.REMOTE) + + @dist_init + def test_rpc_complex_args(self): + self._test_rpc_complex_args(ExecMode.RPC_SYNC, False) + + @dist_init + def test_remote_complex_args(self): + self._test_rpc_complex_args(ExecMode.REMOTE, False) + + @dist_init + def test_context_cleanup_tensor_with_grad(self): + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + self.context_cleanup_test_helper(rpc_args=(t1, t2), func=torch.add) + + @dist_init + def test_context_cleanup_tensor_no_grad(self): + t1 = torch.ones(3, 3, requires_grad=False) + self.context_cleanup_test_helper(rpc_args=(t1, t1), func=torch.add) + + @dist_init + def test_context_cleanup_no_tensors(self): + self.context_cleanup_test_helper(rpc_args=(1, 1), func=my_scalar_add) + + @dist_init + def test_context_cleanup_nested_rpc(self): + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + dst_rank = (self.rank + 1) % self.world_size + args = (t1, t2, dst_rank, self.world_size, 0) + self.context_cleanup_test_helper( + rpc_args=args, func=my_py_nested_call, nested=True + ) + + @dist_init + def test_worker_ids_recorded(self): + dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} + with dist_autograd.context() as context_id: + # if no tensors require grad, we should still record worker_ids, as + # the autograd context ID is still passed to other workers. + t1 = torch.ones(3, 3, requires_grad=False) + t2 = torch.zeros(3, 3, requires_grad=False) + for dst_rank in dst_ranks: + rpc.rpc_sync(worker_name(dst_rank), torch.add, args=(t1, t2)) + rpc.rpc_sync(worker_name(dst_rank), _set_rpc_done, args=(context_id, 1)) + # all worker_ids in dst_ranks should be recorded. + ctx = dist_autograd._current_context() + worker_ids = ctx._known_worker_ids() + self.assertEqual(worker_ids, dst_ranks) + + # worker_ids should be recorded when tensors do require grad + t1.requires_grad = True + t2.requires_grad = True + for dst_rank in dst_ranks: + rpc.rpc_sync(worker_name(dst_rank), torch.add, args=(t1, t2)) + rpc.rpc_sync(worker_name(dst_rank), _set_rpc_done, args=(context_id, 1)) + # all worker_ids in dst_ranks should be recorded. + worker_ids = ctx._known_worker_ids() + self.assertEqual(worker_ids, dst_ranks) + + @dist_init + def test_dist_autograd_profiling(self): + with dist_autograd.context() as context_id: + t1 = torch.rand(3, 3, requires_grad=True) + t2 = torch.rand(3, 3, requires_grad=True) + loss = rpc.rpc_sync( + worker_name(self._next_rank()), torch.add, args=(t1, t2) + ).sum() + with torch.autograd.profiler.profile() as p: + dist_autograd.backward(context_id, [loss]) + + function_events = p.function_events + + def get_event(partial_key): + return next(event for event in function_events if partial_key in event.name) + + send_event = get_event("SendRpcBackward") + recv_event = get_event("RecvRpcBackward") + backward_event = get_event("torch::distributed::autograd::backward") + # There should be at least 1 send and recv_events each, corresponding to send/recv functions executed. + self.assertEqual(send_event.count, 1) + self.assertEqual(recv_event.count, 1) + # The CPU total for backward event should be great than send and recv, since + # applying those functions in the backwards pass is a subset of the entire backward pass. + self.assertGreater(backward_event.cpu_time_total, send_event.cpu_time_total) + self.assertGreater(backward_event.cpu_time_total, recv_event.cpu_time_total) + + @dist_init + def test_error_in_context(self): + with dist_autograd.context(): + t1 = torch.rand(3, 3, requires_grad=True) + t2 = torch.rand(6, 6, requires_grad=True) + + with self.assertRaises(RuntimeError): + # This should throw an error since matrix sizes don't match. + rpc.rpc_sync( + worker_name(self._next_rank()), torch.matmul, args=(t1, t2) + ) + + @dist_init + def test_backward_no_grad_on_tensor(self): + self._backward_no_grad_on_tensor( + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + False, + ) + + @dist_init + def test_backward_simple(self): + self._backward_simple( + self._next_rank(), + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False, + ) + + @dist_init + def test_backward_simple_self(self): + self._backward_simple( + self.rank, + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False, + ) + + @dist_init + def test_backward_rref(self): + callee = worker_name(self._next_rank()) + rref_owner = callee + self._backward_rref( + callee, + rref_owner, + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False, + ) + + @dist_init + def test_backward_rref_multi(self): + if self.rank > 0: + callee = "worker0" + rref_owner = callee + self._backward_rref( + callee, + rref_owner, + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False, + ) + + @dist_init + def test_backward_rref_nested(self): + callee = worker_name((self.rank + 1) % self.world_size) + rref_owner = worker_name((self.rank + 2) % self.world_size) + self._backward_rref( + callee, + rref_owner, + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3), requires_grad=True), + None, + False, + ) + + @dist_init + def test_trainer_ps(self): + self._test_trainer_ps(create_tensor, _run_trainer, False) + + @dist_init + def test_trainer_ps_torchscript_functions(self): + # TODO, need more investigation + # there is rref leak when shutting down, suspect it is because + # ref as arg is passed to pybind boundary, and the ref is not garbage + # collected by python when calling shutdown() + import torch.distributed.rpc.api as api + + api._ignore_rref_leak = True + + self._test_trainer_ps( + create_torchscript_tensor, _run_trainer_torchscript, False + ) + + @dist_init + def test_backward_multiple_round_trips(self): + self._backward_multiple_round_trips( + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3)), + torch.rand((3, 3), requires_grad=True), + torch.rand((3, 3)), + torch.rand((3, 3), requires_grad=True), + None, + False, + ) + + @dist_init + def test_backward_different_tensor_dims(self): + local_grads = None + t1 = torch.rand((4, 6), requires_grad=True) + t2 = torch.rand((6, 5)) + t3 = torch.rand((5, 7), requires_grad=True) + t4 = torch.rand((7, 9)) + + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + val = self._exec_func(exec_mode, torch.matmul, t1, t2) + val = self._exec_func(exec_mode, torch.linalg.multi_dot, (val, t3, t4)) + loss = val.sum() + + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2, t2, t3, t4 + ) + local_grads = ret if ret else local_grads + + @dist_init + def test_backward_unused_tensors(self): + local_grads = None + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + t3 = torch.rand((3, 3), requires_grad=True) + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + s = self._exec_func(exec_mode, torch.stack, (t1, t2, t3)) + val = self._exec_func( + exec_mode, + torch.matmul, + torch.narrow(s, 0, 0, 1), + torch.narrow(s, 0, 2, 1), + ) + + loss = val.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2, t3 + ) + local_grads = ret if ret else local_grads + + @dist_init + def test_backward_multiple_output_tensors(self): + local_grads = None + t = torch.rand((10, 2), requires_grad=True) + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + tensor_list = self._exec_func(exec_mode, torch.split, t, 2) + t1 = tensor_list[0] + t2 = tensor_list[2] + t3 = tensor_list[4] + + val = self._exec_func(exec_mode, torch.linalg.multi_dot, (t1, t2, t3)) + + loss = val.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t + ) + local_grads = ret if ret else local_grads + + def _run_test_backward_unused_send_function_in_thread(self): + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + + # We don't use the result of an RPC function, as a result the + # backward pass would hang in the "FAST" mode. + rpc.rpc_sync(worker_name(self._next_rank()), torch.add, args=(t1, t2)) + + val = torch.mul(t1, t2) + + # Run backward, this would hang forever. + dist_autograd.backward(context_id, [val.sum()]) + + @dist_init + def test_backward_unused_send_function(self): + # Run the test in a thread which would never finish. + t = threading.Thread( + target=self._run_test_backward_unused_send_function_in_thread + ) + t.daemon = True + t.start() + t.join(10) # Wait for 10s. + + # Verify thread is still alive (indicating backward hasn't completed yet). + self.assertTrue(t.is_alive()) + + @dist_init + def test_backward_autograd_engine_error(self): + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + # Perform some ops before error simulation. + tmp = (t1 + t2) * (t1 + t2) + t3 = SimulateBackwardError.apply(tmp) + + # Run multiple round trips across different nodes and verify the + # original node receives an error thrown on a node deep in the chain. + val = rpc.rpc_sync(worker_name(self._next_rank()), torch.add, args=(t2, t3)) + val = rpc.rpc_sync( + worker_name(self._next_rank()), torch.mul, args=(val, t2) + ) + val = rpc.rpc_sync( + worker_name(self._next_rank()), torch.matmul, args=(val, t2) + ) + val = rpc.rpc_sync( + worker_name(self._next_rank()), torch.div, args=(val, t2) + ) + + with self.assertRaisesRegex( + RuntimeError, "Error on Node [0-9]+: Simulate error on backward pass" + ): + # Run backwards, and validate we receive an error. + dist_autograd.backward(context_id, [val.sum()]) + + @dist_init(clean_shutdown=False) + @skip_but_pass_in_sandcastle_if( + IS_MACOS, + "Test is flaky on MacOS since libuv error handling is not as robust as TCP", + ) + def test_backward_node_failure(self): + rpc._set_rpc_timeout(5) # 5 seconds + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + res = rpc.rpc_sync(worker_name(self._next_rank()), torch.add, args=(t1, t2)) + + # Wait for all RPCs to be done. + dist.barrier() + + # Kill all odd rank nodes. + if self.rank % 2 == 0: + shutdown_error_regex = self.get_shutdown_error_regex() + # Wait for all other nodes to die. + for rank in range(self.world_size): + if rank % 2 != 0: + wait_until_node_failure(rank, shutdown_error_regex) + + # Shutdown sequence is not very well defined and as a result + # we might see any error given by get_shutdown_error_regex() + with self.assertRaisesRegex(RuntimeError, shutdown_error_regex): + # Run backwards, and validate we receive an error since all + # other nodes are dead. + dist_autograd.backward(context_id, [res.sum()]) + else: + # Exit all other nodes. + pass + + @dist_init + def test_backward_without_context(self): + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + + context_id = 100 # dummy context_id + with self.assertRaisesRegex( + RuntimeError, + f"Could not find autograd context with id: {context_id}", + ): + res = rpc.rpc_sync(worker_name(self._next_rank()), torch.add, args=(t1, t2)) + dist_autograd.backward(context_id, [res.sum()]) + + @dist_init + def test_backward_without_rpc(self): + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + t3 = torch.add(t1, t2) + + dist_autograd.backward(context_id, [t3.sum()]) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(2, len(grads)) + self.assertIn(t1, grads) + self.assertIn(t2, grads) + self.assertEqual(torch.ones(3, 3), grads[t1]) + self.assertEqual(torch.ones(3, 3), grads[t2]) + + @dist_init + def test_backward_invalid_args(self): + with dist_autograd.context() as context_id: + with self.assertRaisesRegex(TypeError, "incompatible function arguments"): + dist_autograd.backward(context_id, None) + + with self.assertRaisesRegex(TypeError, "incompatible function arguments"): + dist_autograd.backward(None, None) + + with self.assertRaisesRegex( + RuntimeError, "No tensors provided for gradient computation" + ): + dist_autograd.backward(context_id, []) + + with self.assertRaisesRegex(RuntimeError, "requires_grad not set on"): + t = torch.rand(3, 3) + dist_autograd.backward(context_id, [t]) + + with self.assertRaisesRegex( + RuntimeError, "is not a scalar, all roots need to be scalar" + ): + t = torch.rand(3, 3, requires_grad=True) + dist_autograd.backward(context_id, [t]) + + with self.assertRaisesRegex( + RuntimeError, "does not have a valid gradient function" + ): + t = torch.rand(1, requires_grad=True) + dist_autograd.backward(context_id, [t]) + + @dist_init + def test_backward_multiple_roots(self): + local_grads = None + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC]: + with dist_autograd.context() as context_id: + r1 = self._exec_func(exec_mode, torch.add, t1, t2).sum() + r2 = self._exec_func(exec_mode, torch.mul, t1, t2).sum() + r3 = self._exec_func(exec_mode, torch.cos, t1).sum() + r4 = self._exec_func(exec_mode, torch.div, t1, t2).sum() + + local_grads = self._verify_backwards( + exec_mode, [r1, r2, r3, r4], context_id, local_grads, t1, t2 + ) + + @dist_init + def test_backward_different_dtypes(self): + self._backward_different_dtypes( + torch.rand((3, 3), requires_grad=True, dtype=torch.float32), + torch.rand((3, 3), requires_grad=True, dtype=torch.float64), + False, + ) + + @dist_init + def test_backward_simple_python_udf(self): + self._backward_simple_python_udf( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False, + ) + + @dist_init + def test_backward_simple_script_call(self): + self._backward_simple_script_call( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False, + ) + + @staticmethod + def _complex_python_udf(t1, t2): + t3 = torch.nn.functional.linear(t1, t2) + t4 = torch.nn.functional.linear(t2, t3) + t5 = torch.nn.functional.linear(t3, t4) + return torch.linalg.multi_dot([t1, t2, t3, t4, t5]) + + @dist_init + def test_backward_complex_python_udf(self): + # Run the same code locally and with dist autograd and verify gradients + # are same. + local_grads = None + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + for exec_mode in [ExecMode.LOCAL, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func( + exec_mode, DistAutogradTest._complex_python_udf, t1, t2 + ) + loss = ret.sum() + local_grads = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + + @staticmethod + def _python_udf_with_backward_error(t1, t2): + t3 = t1 + t2 + t4 = SimulateBackwardError.apply(t3) + return torch.linalg.multi_dot([t1, t2, t3, t4]) + + @staticmethod + def _nested_rpc_call_backward_error(t1, t2, dst): + t1 = t1 * t2 + t2 = t1 + t2 + res = rpc.rpc_sync( + worker_name(dst), + DistAutogradTest._python_udf_with_backward_error, + args=(t1, t2), + ) + return torch.linalg.multi_dot([t1, t2, res]) + + @dist_init + def test_backward_python_udf_error(self): + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + worker_name(self._next_rank()), + DistAutogradTest._nested_rpc_call_backward_error, + args=(t1, t2, self._next_rank()), + ) + with self.assertRaisesRegex( + RuntimeError, "Simulate error on backward pass" + ): + dist_autograd.backward(context_id, [loss.sum()]) + + _backward_done = False + + @dist_init(clean_shutdown=False) + @skip_but_pass_in_sandcastle_if( + IS_MACOS, + "Test is flaky on MacOS since libuv error handling is not as robust as TCP", + ) + def test_backward_node_failure_python_udf(self): + # Set a short timeout to quickly time out failed RPCs. + rpc._set_rpc_timeout(5) # 5 seconds + initialize_pg(self.file_init_method, self.rank, self.world_size) + + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + + dst = self._next_rank() + res = rpc.rpc_sync( + worker_name(dst), + my_py_nested_call, + args=(t1, t2, dst, self.world_size, 1), + ) + + dist.barrier() + + # Kill rank 2 (last hop of nested rpc) and verify rank 0 receives an error. + if self.rank == 2: + return + + store = dist.distributed_c10d._get_default_store() + if self.rank == 0: + # Wait for rank 2 to die. + shutdown_error_regex = self.get_shutdown_error_regex() + wait_until_node_failure(2, shutdown_error_regex) + # Shutdown sequence is not very well defined and as a result + # we might see any error given by get_shutdown_error_regex(). + with self.assertRaisesRegex(RuntimeError, shutdown_error_regex): + # Run backwards, and validate we receive an error since rank 2 is dead. + dist_autograd.backward(context_id, [res.sum()]) + + # Mark rank 0 is done in the store, since the RPC framework on + # some nodes might be broken at this point. + store.set("test_backward_node_failure_python_udf_rank0_done", "True") + else: + # Wait for backward to finish on rank 0. + store.wait( + ["test_backward_node_failure_python_udf_rank0_done"], + timedelta(seconds=10), + ) + + @staticmethod + def _nested_python_udf(t1, t2, dst): + t3 = t1 * t2 + t4 = t1 + t2 + res = rpc.rpc_sync(worker_name(dst), my_py_add, args=(t3, t4)) + return t1 * t2 * t3 * t4 * res + + @dist_init + def test_backwards_nested_python_udf(self): + # Run equivalent of _nested_python_udf locally. + self._backwards_nested_python_udf( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False, + ) + + _test_clean_context_backward_context_id = None + + class MyBackwardFunc(Function): + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + @once_differentiable + def backward(ctx, input): + if DistAutogradTest._test_clean_context_backward_context_id is None: + raise AssertionError( + "Expected _test_clean_context_backward_context_id to not be None" + ) + + # Release the context to simulate error (use barrier before releasing + # context to ensure all nodes execute the backward function). + dist.barrier() + dist_autograd._release_context( + DistAutogradTest._test_clean_context_backward_context_id + ) + + # Verify all contexts are cleaned up. + if not _all_contexts_cleaned_up(): + raise AssertionError("Expected all contexts to be cleaned up") + + return input + + @dist_init + def test_clean_context_during_backward(self): + """ + This test simulates the situation where the 'backward' call might throw + an exception locally which would lead to the autograd context being + cleaned up if we're using the context manager. As a result, the autograd + context might be cleaned up while some threads are still using the + autograd context. + + It is fine for the 'backward' call to throw an exception in this test, + but the process should not crash. + """ + initialize_pg(self.file_init_method, self.rank, self.world_size) + + context = dist_autograd._new_context() + context_id = context._context_id() + DistAutogradTest._test_clean_context_backward_context_id = context_id + + # Send the context id to all nodes. + for i in range(self.world_size): + if i != self.rank: + rank_distance = (i - self.rank + self.world_size) % self.world_size + rpc.rpc_sync( + worker_name(i), + _set_rpc_done, + args=(context_id, rank_distance), + ) + + dist.barrier() + + # Verify all context ids have been received. + self.assertEqual(self.world_size - 1, len(known_context_ids)) + + t1 = torch.rand((3, 3), requires_grad=True) + for _ in range(100): + dst = self._next_rank() + t1 = rpc.rpc_sync(worker_name(dst), torch.add, args=(t1, t1)) + + # Call MyBackwardFunc as the first op of the backward pass to + # ensure we release the context early in the backward pass. + t1 = DistAutogradTest.MyBackwardFunc.apply(t1) + self.assertEqual(100, len(context._send_functions())) + + context_id = 100 # dummy context_id + with self.assertRaisesRegex( + RuntimeError, + f"Could not find autograd context with id: {context_id}", + ): + dist_autograd.backward(context_id, [t1.sum()]) + + # HACK: Killing workers since otherwise the autograd engine gets stuck on + # other nodes. The proper fix would be addressing: + # https://github.com/pytorch/pytorch/issues/27643, which would inform + # other nodes about the failure. + # The autograd engine gets stuck on other nodes since they're waiting to + # receive gradients from the node that received an error (and as a + # result it didn't execute the rest of the graph). + dist.barrier() + rpc.shutdown(graceful=False) + sys.exit(0) + + @classmethod + def _call_remote_embedding(cls, embedding_rref, input, offsets, per_sample_weights): + embedding = embedding_rref.local_value() + return embedding(input, offsets, per_sample_weights) + + @classmethod + def _get_grad(cls, embedding_rref, context_id): + embedding = embedding_rref.local_value() + grad_map = dist_autograd.get_gradients(context_id) + return grad_map[embedding.weight] + + @classmethod + def _mixed_requires_grad_operaton(cls, t1, t2): + if t2.requires_grad: + return t1 - t2 + else: + return t1 * t2 + + @dist_init + def test_mixed_requires_grad(self): + self._mixed_requires_grad( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=False), + False, + ) + + class TestDebugInfoFunc(Function): + @staticmethod + def forward(ctx, input): + return input + + @staticmethod + @once_differentiable + def backward(ctx, input): + debug_info = dist_autograd._get_debug_info() + if debug_info is None: + raise AssertionError("Expected debug_info to not be None") + backward_passes = int(debug_info["num_current_backward_passes"]) + + # Hard to validate exact numbers because of the distributed nature. + # We can't use a barrier() here since that would block the single + # CPU thread available for autograd and can cause deadlocks. + if not (1 <= backward_passes <= 4): + raise AssertionError( + f"Expected 1 <= backward_passes <= 4, got {backward_passes}" + ) + return input + + @dist_init + def test_debug_info(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + i = 0 + res = {} + res[i] = t1 + for rank in range(self.world_size): + if rank != self.rank: + res[i + 1] = rpc.rpc_sync( + worker_name(rank), torch.add, args=(res[i], t2) + ) + i += 1 + + # Call custom function in middle of backward pass to ensure all + # nodes are still waiting on a backward(). + res[i + 1] = DistAutogradTest.TestDebugInfoFunc.apply(res[i]) + i += 1 + + for rank in range(self.world_size): + if rank != self.rank: + res[i + 1] = rpc.rpc_sync( + worker_name(rank), torch.add, args=(res[i], t2) + ) + i += 1 + + dist_autograd.backward(context_id, [res[i].sum()]) + + debug_info = dist_autograd._get_debug_info() + num_autograd_context = int(debug_info["num_autograd_contexts"]) + # Need at least one context and not more than 4. + self.assertTrue(num_autograd_context >= 1 and num_autograd_context <= 4) + + for rd in range(self.world_size - 1): + rpc.rpc_sync( + worker_name((self.rank + rd + 1) % self.world_size), + _set_rpc_done, + args=(context_id, rd + 1), + ) + + dist.barrier() + + # Validate information + debug_info = dist_autograd._get_debug_info() + if debug_info is None: + raise AssertionError("Expected debug_info to not be None") + self.assertEqual(0, int(debug_info["num_current_backward_passes"])) + # only have `num_current_backward_passes` and `num_autograd contexts` + self.assertTrue(len(debug_info) == 2) + + self.assertTrue(_all_contexts_cleaned_up()) + + # All contexts should be cleaned up. + debug_info = dist_autograd._get_debug_info() + self.assertEqual(0, int(debug_info["num_autograd_contexts"])) + + @staticmethod + def _workload_thread(): + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + t3 = rpc.rpc_sync("worker0", torch.add, args=(t1, t2)) + t4 = rpc.rpc_sync("worker0", torch.mul, args=(t2, t3)) + t5 = rpc.rpc_sync("worker0", torch.matmul, args=(t3, t4)) + t6 = rpc.rpc_sync("worker0", torch.add, args=(t4, t5)) + + dist_autograd.backward(context_id, [t6.sum()]) + + @dist_init + def test_async_dist_autograd(self): + """ + This test ensures async processing for distributed autograd works + appropriately. This is achieved by spawning multiple threads and + hammering a single node with a lot of backward() calls. + """ + + initialize_pg(self.file_init_method, self.rank, self.world_size) + if self.rank != 0: + # All other ranks schedule work on rank 0. + threads = [] + for _ in range(20): + t = threading.Thread(target=DistAutogradTest._workload_thread) + t.start() + threads.append(t) + + for thread in threads: + thread.join() + + dist.barrier() + + @dist_init + def test_backward_accumulate_grads(self): + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + t3 = torch.matmul(t1, t2) + # Run backward twice. + torch.autograd.backward([t3.sum()], retain_graph=True) + torch.autograd.backward([t3.sum()]) + + t3 = rpc.rpc_sync( + worker_name(self._next_rank()), torch.matmul, args=(t1, t2) + ) + # Run backward twice. + dist_autograd.backward(context_id, [t3.sum()], retain_graph=True) + dist_autograd.backward(context_id, [t3.sum()]) + + # Verify the gradients are same for local and remote execution. + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(2, len(grads)) + self.assertIn(t1, grads) + self.assertIn(t2, grads) + self.assertEqual(t1.grad, grads[t1]) + self.assertEqual(t2.grad, grads[t2]) + + @staticmethod + def _test_nested_backward_accumulate_grads(t1, t2, dst_rank): + return rpc.rpc_sync(worker_name(dst_rank), torch.add, args=(t1, t2)) + + @dist_init + def test_nested_backward_accumulate_grads(self): + self._nested_backward_accumulate_grads( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False, + ) + + @dist_init + def test_multiple_backward(self): + self._multiple_backward( + torch.rand(3, 3, requires_grad=True), + torch.rand(3, 3, requires_grad=True), + False, + ) + + @dist_init(clean_shutdown=False) + def test_multiple_backward_with_errors(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + with dist_autograd.context() as context_id: + loss = rpc.rpc_sync( + f"worker{self._next_rank()}", + DistAutogradTest._python_udf_with_backward_error, + args=(t1, t2), + ).sum() + + try: + # Run backward in a loop multiple times. + for i in range(100): + if i < 50: + with self.assertRaisesRegex( + RuntimeError, "Simulate error on backward pass" + ): + dist_autograd.backward( + context_id, [loss], retain_graph=True + ) + elif i > 50: + # Recovered from error. + dist_autograd.backward(context_id, [loss], retain_graph=True) + else: + dist.barrier() + SimulateBackwardError._simulate_error = False + dist.barrier() + finally: + # Sync before resetting flag. + dist.barrier() + + # Reset the flag. + SimulateBackwardError._simulate_error = True + + @dist_init + def test_backward_verify_hooks(self): + t1 = torch.ones((3, 3), requires_grad=True) + # Double the gradient. + t1.register_hook(lambda grad: grad * 2) + t2 = torch.ones((3, 3), requires_grad=True) + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC, ExecMode.REMOTE]: + with dist_autograd.context() as context_id: + ret = self._exec_func(exec_mode, torch.matmul, t1, t2) + loss = ret.sum() + ret = self._verify_backwards( + exec_mode, [loss], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + @dist_init + def test_no_grad_copy(self): + """ + Similar to test in test_autograd.py. + """ + + # create autograd function that saves grad pointer as class static + class MyFunc(Function): + static_grad_ptr = None + + @staticmethod + def forward(ctx, inp1, inp2): + return inp1 + inp2 + + @staticmethod + def backward(ctx, grad): + MyFunc.static_grad_ptr = grad.data_ptr() + return grad, grad + + class MyFuncSingleGrad(Function): + static_grad_ptr = None + + @staticmethod + def forward(ctx, inp): + return inp + + @staticmethod + def backward(ctx, grad): + MyFuncSingleGrad.static_grad_ptr = grad.data_ptr() + return grad + + class NonContGradFunc(Function): + @staticmethod + def forward(ctx, inp1): + ctx.size = inp1.size() + return torch.tensor([1.0]) + + @staticmethod + def backward(ctx, grad): + return torch.ones(1).expand(ctx.size) + + a = torch.randn(5, 6, requires_grad=True) + b = torch.randn(5, 6, requires_grad=True) + # non-contiguous grad should be copied + with dist_autograd.context() as context_id: + dist_autograd.backward( + context_id, [NonContGradFunc.apply(MyFunc.apply(a, b))] + ) + grads = dist_autograd.get_gradients(context_id) + self.assertFalse(grads[a].data_ptr() == MyFunc.static_grad_ptr) + self.assertFalse(grads[b].data_ptr() == MyFunc.static_grad_ptr) + + # test case that should trigger no copy for a + with dist_autograd.context() as context_id: + dist_autograd.backward(context_id, [MyFuncSingleGrad.apply(a)[1][0]]) + grads = dist_autograd.get_gradients(context_id) + p_g = MyFuncSingleGrad.static_grad_ptr + p_a = grads[a].data_ptr() + # Verify there was no clone. + self.assertTrue(p_a == p_g) + + # Test case that should trigger copy for both of a,b. This is + # different in the distributed autograd case since we hold + # a reference to all grads in a vector until all accumulation is done. + with dist_autograd.context() as context_id: + dist_autograd.backward(context_id, [MyFunc.apply(a, b)[1][0]]) + grads = dist_autograd.get_gradients(context_id) + p_g = MyFunc.static_grad_ptr + p_a = grads[a].data_ptr() + p_b = grads[b].data_ptr() + # check a,b uses different grad buffer + self.assertFalse(p_a == p_b) + # both should be copied. + self.assertFalse(grads[a].data_ptr() == MyFunc.static_grad_ptr) + self.assertFalse(grads[b].data_ptr() == MyFunc.static_grad_ptr) + + @dist_init + def test_no_grad_copy_sparse(self): + # create autograd function that saves grad pointer as class static + class MyFunc(Function): + static_grad_ptr = None + + @staticmethod + def forward(ctx, inp): + return inp + + @staticmethod + def backward(ctx, grad): + MyFunc.static_grad_ptr = grad._values().data_ptr() + return grad + + class NonContGradFunc(Function): + static_grad_ptr = None + + @staticmethod + def forward(ctx, inp1, inp2): + return inp1 + inp2 + + @staticmethod + def backward(ctx, grad): + # Create a sparse tensor with non-contiguous indices and values + # and return as grad. + v = torch.rand(1, 3) + i = torch.ones(1, 1, dtype=torch.long) + nv = v.expand(8, 3) + ni = i.expand(1, 8) + ngrad = torch.sparse_coo_tensor(ni, nv, (10, 3), dtype=torch.float32) + NonContGradFunc.static_grad_ptr = ngrad._values().data_ptr() + return ngrad, ngrad + + a = torch.randn(10, 3, requires_grad=True) + b = torch.randn(10, 3, requires_grad=True) + input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9]) + offsets = torch.tensor([0, 4]) + import torch.nn.functional as F + + # test case that should trigger no copy for a. + with dist_autograd.context() as context_id: + emb_matrix = MyFunc.apply(a) + loss = F.embedding_bag(emb_matrix, input, offsets, sparse=True).sum() + dist_autograd.backward(context_id, [loss], retain_graph=True) + grads = dist_autograd.get_gradients(context_id) + p_g = MyFunc.static_grad_ptr + p_a = grads[a]._values().data_ptr() + # check a uses the same buffer + self.assertTrue(p_a == p_g) + + # Run backwards multiple times. + for _ in range(10): + dist_autograd.backward(context_id, [loss], retain_graph=True) + + # non-contiguous indices and value, we should trigger a copy. + with dist_autograd.context() as context_id: + emb_matrix = NonContGradFunc.apply(a, b) + loss = F.embedding_bag(emb_matrix, input, offsets, sparse=True).sum() + dist_autograd.backward(context_id, [loss], retain_graph=True) + grads = dist_autograd.get_gradients(context_id) + p_g = NonContGradFunc.static_grad_ptr + p_a = grads[a]._values().data_ptr() + p_b = grads[b]._values().data_ptr() + # check a,b uses different grad buffer + self.assertFalse(p_a == p_b) + # Verify we cloned both grads. + self.assertFalse(p_a == p_g) + self.assertFalse(p_b == p_g) + + # Run backwards multiple times to verify accumulation. + for _ in range(10): + dist_autograd.backward(context_id, [loss], retain_graph=True) + + @dist_init + def test_grad_copy_sparse_indices_extra_ref(self): + # create autograd function that saves grad pointer as class static + class MyFunc(Function): + static_grad_ptr = None + static_grad_indices_ref = None + static_grad_values_ref = None + + @staticmethod + def forward(ctx, inp): + return inp + + @staticmethod + def backward(ctx, grad): + MyFunc.static_grad_ptr = grad._values().data_ptr() + # indices() and values() return views, so holding onto + # references of them would not increment refcount of indices + # and values inside the sparse tensor. + MyFunc.static_grad_indices_ref = grad._indices() + MyFunc.static_grad_values_ref = grad._values() + return grad + + a = torch.randn(10, 3, requires_grad=True) + input = torch.tensor([1, 2, 4, 5, 4, 3, 2, 9]) + offsets = torch.tensor([0, 4]) + import torch.nn.functional as F + + with dist_autograd.context() as context_id: + emb_matrix = MyFunc.apply(a) + loss = F.embedding_bag(emb_matrix, input, offsets, sparse=True).sum() + dist_autograd.backward(context_id, [loss], retain_graph=True) + grads = dist_autograd.get_gradients(context_id) + p_g = MyFunc.static_grad_ptr + p_a = grads[a]._values().data_ptr() + self.assertIsNotNone(MyFunc.static_grad_indices_ref) + self.assertIsNotNone(MyFunc.static_grad_values_ref) + # grad would be stolen, since static_grad_indices_ref and + # static_grad_values_ref are holding onto views and don't bump the + # refcount. + self.assertTrue(p_g == p_a) + + @dist_init + def test_post_hooks(self): + self.hook_called_times = 0 + + def post_hook_add_one(output_grads, input_grads): + self.hook_called_times += 1 + return output_grads + + def post_hook_add_two(output_grads, input_grads): + self.hook_called_times += 2 + return output_grads + + t = torch.rand(10, 10, requires_grad=True) + a = t + t + + # Register post hooks + accumulate_grad_0 = a.grad_fn.next_functions[0][0] + accumulate_grad_0.register_hook(post_hook_add_one) + accumulate_grad_0.register_hook(post_hook_add_two) + + accumulate_grad_1 = a.grad_fn.next_functions[1][0] + accumulate_grad_1.register_hook(post_hook_add_two) + + with dist_autograd.context() as context_id: + loss = a.sum() + dist_autograd.backward(context_id, [loss]) + self.assertEqual(5, self.hook_called_times) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(1, len(grads)) + self.assertTrue(t in grads) + + @staticmethod + def _slow_add(t1, t2): + time.sleep(1) + t3 = t1 + t2 + t3.requires_grad = True + return t3 + + @dist_init + def test_thread_local_context_id(self): + t1 = torch.rand((3, 3)) + t2 = torch.rand((3, 3)) + + t3 = t1 + t2 + t3.requires_grad = True + t3.sum().backward() + + dst = worker_name((self.rank + 1) % self.world_size) + rref = rpc.remote(dst, DistAutogradTest._slow_add, args=(t1, t2)) + + with dist_autograd.context() as context_id: + loss = rref.to_here().sum() + # due to slow add, the continuation of this backward pass will be + # invoked by the previous rpc.remote thread which does not have a + # valid context_id. So, this can test whether we propagate + # thread_local states properly when jumping across threads on the + # server side. + dist_autograd.backward(context_id, [loss]) + self.assertTrue( + rpc.rpc_sync( + dst, _compare_owner_value, args=(context_id, rref, t3.grad) + ) + ) + + +class CudaDistAutogradTest(CommonDistAutogradTest): + @skip_if_lt_x_gpu(1) + @dist_init + def test_gpu_simple(self): + t1 = torch.rand(3, 3, requires_grad=True, device="cuda:0") + t2 = torch.rand(3, 3, requires_grad=True, device="cuda:0") + (t1 + t2).sum().backward() + with dist_autograd.context() as context_id: + t3 = t1 + t2 + dist_autograd.backward(context_id, [t3.sum()]) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(2, len(grads)) + self.assertEqual(t1.grad, grads[t1]) + self.assertEqual(t2.grad, grads[t2]) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_gpu_to_cpu_continuation(self): + t1 = torch.rand(3, 3, requires_grad=True, device="cuda:0") + t2 = torch.rand(3, 3, requires_grad=True) + # Run a few iterations. + for _ in range(3): + t1.grad = None + t2.grad = None + # Root is CPU + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC]: + with dist_autograd.context() as context_id: + t3 = self._exec_func(exec_mode, torch.add, t2, t2) + t4 = t3.cuda(0) + t1 + t5 = self._exec_func(exec_mode, torch.add, t4.cpu(), t2) + t6 = t5.cuda(0) + t4 + t7 = self._exec_func(exec_mode, torch.add, t6.cpu(), t5) + # Autograd graph consists of CPU -> GPU -> CPU execution. + ret = self._verify_backwards( + exec_mode, [t7.sum()], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + @skip_if_lt_x_gpu(1) + @dist_init + def test_gpu_to_cpu_continuation_gpu_root(self): + t1 = torch.rand(3, 3, requires_grad=True, device="cuda:0") + t2 = torch.rand(3, 3, requires_grad=True) + # Run a few iterations. + for _ in range(3): + t1.grad = None + t2.grad = None + # Root is CPU + local_grads = None + for exec_mode in [ExecMode.LOCAL, ExecMode.RPC_SYNC]: + with dist_autograd.context() as context_id: + t3 = self._exec_func(exec_mode, torch.add, t2, t2) + t4 = t3.cuda(0) + t1 + t5 = self._exec_func(exec_mode, torch.add, t4.cpu(), t2) + t6 = t5.cuda(0) + t4 + # Autograd graph consists of CPU -> GPU -> CPU execution. + ret = self._verify_backwards( + exec_mode, [t6.sum()], context_id, local_grads, t1, t2 + ) + local_grads = ret if ret else local_grads + + +class FaultyAgentDistAutogradTest(RpcAgentTestFixture): + # Reusing a simplified helper function from DistAutogradTest to ensure + # autograd context is successfully cleaned up even when RPCs are failing. + def context_cleanup_test_helper(self, rpc_args, func): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + # test that in dist autograd, in the case that tensors communicated over RPC do + # NOT require grad, we still cleanup the dist autograd contexts created + # on other nodes. This is because the autograd context is still + # communicated over RPC even if tensor arguments do not require grad, as + # it is possible that the response could. + dst_ranks = {rank for rank in range(self.world_size) if rank != self.rank} + + with dist_autograd.context() as context_id: + for dst_rank in dst_ranks: + rpc.rpc_sync(worker_name(dst_rank), func, args=rpc_args) + rpc.rpc_sync(worker_name(dst_rank), _set_rpc_done, args=(context_id, 1)) + # the thread's context id should be cleaned up + with self.assertRaises(RuntimeError): + dist_autograd._retrieve_context(context_id) + # Ensure all peers have finished mutating the + # `known_context_ids` set. + dist.barrier() + # check that all contexts have been cleaned up. + success = _all_contexts_cleaned_up() + self.assertTrue(success) + + # no faulty_messages defined so this fails all retryable messages - see + # faulty_rpc_agent_test_fixture.py for the list of retryable messages. + @dist_init + def test_context_cleanup_tensor_with_grad(self): + t1 = torch.ones(3, 3, requires_grad=True) + t2 = torch.zeros(3, 3, requires_grad=True) + self.context_cleanup_test_helper(rpc_args=(t1, t2), func=torch.add) + + @dist_init + def test_verify_backend_options(self): + self.assertEqual( + self.rpc_backend, rpc.backend_registry.BackendType.FAULTY_TENSORPIPE + ) + self.assertEqual(self.rpc_backend_options.num_worker_threads, 8) + self.assertEqual(self.rpc_backend_options.num_fail_sends, 3) + self.assertEqual(len(self.rpc_backend_options.messages_to_fail), 4) + + +class WrapperModule(nn.Module): + def __init__(self, model, device): + super().__init__() + self.model = model.to(device) + + def forward(self, *args): + return self.model(*args) + + def gradients(self, ctx_id): + grads = dist_autograd.get_gradients(ctx_id) + return [grads[p] for p in self.model.parameters()] + + +class TensorPipeCudaDistAutogradTest(RpcAgentTestFixture): + @skip_if_lt_x_gpu(4) + def test_device_maps_backward_pass(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + + # The reverse of this device mapping should be used for the backward pass. + options.set_device_map(dst, {self.rank: (self.rank + 1) % self.world_size}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + t1 = torch.rand(10, device=self.rank, requires_grad=True) + t2 = torch.rand(10, device=self.rank, requires_grad=True) + with dist_autograd.context() as context_id: + res = rpc.rpc_sync(dst, torch.add, args=(t1, t2)) + dist_autograd.backward(context_id, [res.sum()]) + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(torch.ones(10), grads[t1]) + self.assertEqual(torch.ones(10), grads[t2]) + self.assertEqual(t1.device, grads[t1].device) + self.assertEqual(t2.device, grads[t2].device) + + rpc.shutdown() + + class MyRemoteCompute(torch.nn.Module): + def forward(self, input): + input = input * 2.0 + return input + + class MyLocalCompute(torch.nn.Module): + def __init__(self, next_stage): + super().__init__() + self.next_stage = next_stage + + def forward(self, input): + return self.next_stage.rpc_sync().forward(input) + + @skip_if_lt_x_gpu(4) + def test_dist_autograd_sync_streams(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + + # The reverse of this device mapping should be used for the backward pass. + options.set_device_map(dst, {self.rank: (self.rank + 1) % self.world_size}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + remote_compute = rpc.remote(dst, TensorPipeCudaDistAutogradTest.MyRemoteCompute) + local_compute = TensorPipeCudaDistAutogradTest.MyLocalCompute(remote_compute) + for _ in range(10): + input = torch.rand([1000, 10000], device=self.rank, requires_grad=True) + # Run local autograd + result = input * 2.0 + r = random.random() + loss = result.sum() * r + loss.backward() + + # Run distributed autograd + with dist_autograd.context() as context_id: + result = local_compute(input) + loss = result.sum() * r + dist_autograd.backward(context_id, [loss]) + + # Compare grads. + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(input.grad, grads[input]) + + rpc.shutdown() + + @skip_if_lt_x_gpu(4) + def test_gradients_synchronizations(self): + options = self.rpc_backend_options + for peer_rank in range(self.world_size): + options.set_device_map(worker_name(peer_rank), {self.rank: peer_rank}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + if self.rank == 0: + # this is master + layers = [nn.Linear(2000, 2000) for _ in range(self.world_size - 1)] + local_layers = [l.to(0) for l in layers] + remote_layers = [ + rpc.remote( + worker_name(rank), WrapperModule, args=(layers[rank - 1], rank) + ) + for rank in range(1, self.world_size) + ] + + x = torch.randn(5000, 2000).to(0) + # local iteration + local_model = nn.Sequential(*local_layers) + local_model(x).sum().backward() + + # remote iteration + with dist_autograd.context() as context_id: + for remote_layer in remote_layers: + x = remote_layer.rpc_sync().forward(x) + + dist_autograd.backward(context_id, [x.sum()]) + + futs = [] + for remote_layer in remote_layers: + futs.append(remote_layer.rpc_async().gradients(context_id)) + + for i in range(len(futs)): + local_gradients = [p.grad for p in local_layers[i].parameters()] + for g1, g2 in zip(futs[i].wait(), local_gradients, strict=True): + self.assertEqual(g1, g2) + + rpc.shutdown() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/dist_optimizer_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/dist_optimizer_test.py new file mode 100644 index 0000000000000000000000000000000000000000..2d335325f8364241dd14517da5c67c2a6e6a032b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/dist_optimizer_test.py @@ -0,0 +1,281 @@ +# mypy: allow-untyped-defs + + +import threading + +import torch +import torch.distributed.autograd as dist_autograd +import torch.distributed.rpc as rpc +from torch import optim +from torch.distributed.optim import DistributedOptimizer +from torch.testing._internal.dist_utils import dist_init +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +class MyModule: + lock = threading.Lock() + + def __init__(self, requires_grad=True): + # cannot directly use torch.manual_seed(0) as all threads share the same + # default generator. The race from multiple RPC threads could mess up + # the draw order from the default RNG instance, leading to + # non-deterministic behavior. Hence, create a dedicated RNG here. + g_cpu = torch.Generator() + g_cpu.manual_seed(0) + self.w = torch.rand((3, 3), requires_grad=requires_grad, generator=g_cpu) + + def forward(self, t1): + return torch.mm(self.w, t1) + + def get_w(self): + return self.w + + +class FailingOptimizer(optim.Optimizer): + def __init__(self, params): + super().__init__(params, {}) + + def step(self, closure=None): + raise ValueError("Error running optimizer.") + + +class OptimizerFailingOnConstructor(optim.Optimizer): + def __init__(self, params): + super().__init__(params, {}) + raise ValueError("Error creating optimizer.") + + def step(self, closure=None): + raise NotImplementedError + + +def _call_method(method, obj_rref, *args, **kwargs): + return method(obj_rref.local_value(), *args, **kwargs) + + +def remote_method(method, obj_rref, *args, **kwargs): + """ + Call rpc.remote on a method in a remote object. + + Args: + method: the method (for example, Class.method) + obj_rref (RRef): remote reference to the object + args: positional arguments to pass to the method + kwargs: keyword arguments to pass to the method + + Returns a RRef to the remote method call result. + """ + return rpc.remote( + obj_rref.owner(), + _call_method, + args=[method, obj_rref] + list(args), + kwargs=kwargs, + ) + + +def rpc_async_method(method, obj_rref, *args, **kwargs): + """ + Call rpc.rpc_async on a method in a remote object. + + Args: + method: the method (for example, Class.method) + obj_rref (RRef): remote reference to the object + args: positional arguments to pass to the method + kwargs: keyword arguments to pass to the method + + Returns a Future to the method call result. + """ + return rpc.rpc_async( + obj_rref.owner(), + _call_method, + args=[method, obj_rref] + list(args), + kwargs=kwargs, + ) + + +class DistOptimizerTest(RpcAgentTestFixture): + @dist_init() + def test_dist_optim_exception(self): + # distributed version + owner1 = f"worker{(self.rank + 1) % self.world_size:d}" + owner2 = f"worker{(self.rank + 2) % self.world_size:d}" + + remote_module1 = rpc.remote(owner1, MyModule) + remote_module2 = rpc.remote(owner2, MyModule) + remote_param1 = remote_method(MyModule.get_w, remote_module1) + remote_param2 = remote_method(MyModule.get_w, remote_module2) + + dist_optim = DistributedOptimizer( + FailingOptimizer, [remote_param1, remote_param2] + ) + + with dist_autograd.context() as context_id: + g_cpu = torch.Generator() + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = rpc_async_method(MyModule.forward, remote_module1, t2) + output2 = rpc_async_method(MyModule.forward, remote_module2, output1.wait()) + loss = torch.add(output2.wait(), t1).sum() + + dist_autograd.backward(context_id, [loss]) + with self.assertRaisesRegex(Exception, "Error running optimizer"): + dist_optim.step(context_id) + + @dist_init() + def test_dist_optim_exception_on_constructor(self): + # distributed version + owner1 = f"worker{(self.rank + 1) % self.world_size:d}" + owner2 = f"worker{(self.rank + 2) % self.world_size:d}" + + remote_module1 = rpc.remote(owner1, MyModule) + remote_module2 = rpc.remote(owner2, MyModule) + remote_param1 = remote_method(MyModule.get_w, remote_module1) + remote_param2 = remote_method(MyModule.get_w, remote_module2) + + with self.assertRaisesRegex(Exception, "Error creating optimizer."): + DistributedOptimizer( + OptimizerFailingOnConstructor, [remote_param1, remote_param2] + ) + + def _test_dist_optim_base(self, optim_cls, *args, **kwargs): + # local version + module1 = MyModule() + module2 = MyModule() + params = [module1.get_w(), module2.get_w()] + local_optim = optim_cls(params, *args, **kwargs) + + old_w1 = module1.w.detach().clone() + old_w2 = module2.w.detach().clone() + + g_cpu = torch.Generator() + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = module1.forward(t2) + output2 = module2.forward(output1) + loss = torch.add(output2, t1).sum() + + loss.backward() + local_optim.step() + + # distributed version + owner1 = f"worker{(self.rank + 1) % self.world_size:d}" + owner2 = f"worker{(self.rank + 2) % self.world_size:d}" + + remote_module1 = rpc.remote(owner1, MyModule) + remote_module2 = rpc.remote(owner2, MyModule) + remote_param1 = remote_method(MyModule.get_w, remote_module1) + remote_param2 = remote_method(MyModule.get_w, remote_module2) + + # sanity check: local and remote initial weights should match + self.assertEqual(old_w1, remote_param1.to_here()) + self.assertEqual(old_w2, remote_param2.to_here()) + + dist_optim = DistributedOptimizer( + optim_cls, [remote_param1, remote_param2], *args, **kwargs + ) + + with dist_autograd.context() as context_id: + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = rpc_async_method(MyModule.forward, remote_module1, t2) + output2 = rpc_async_method(MyModule.forward, remote_module2, output1.wait()) + loss = torch.add(output2.wait(), t1) + + dist_autograd.backward(context_id, [loss.sum()]) + dist_optim.step(context_id) + + new_w1 = rpc_async_method(MyModule.get_w, remote_module1).wait() + new_w2 = rpc_async_method(MyModule.get_w, remote_module2).wait() + + # ensure optimizer changed weights + self.assertNotEqual(old_w1, new_w1) + self.assertNotEqual(old_w2, new_w2) + # ensure local equals remote + self.assertEqual(new_w1, module1.get_w()) + self.assertEqual(new_w2, module2.get_w()) + + @dist_init() + def test_dist_optim(self): + self._test_dist_optim_base(optim.Adagrad, lr=0.05) + self._test_dist_optim_base(optim.Adam, lr=1e-2, amsgrad=True) + self._test_dist_optim_base(optim.AdamW, lr=0.05, amsgrad=True) + self._test_dist_optim_base(optim.SGD, lr=0.05) + self._test_dist_optim_base( + optim.SGD, lr=1e-3, momentum=1, weight_decay=1, nesterov=True + ) + self._test_dist_optim_base(optim.Adadelta, rho=0.95) + self._test_dist_optim_base(optim.RMSprop, lr=0.05) + self._test_dist_optim_base(optim.Adamax, lr=0.05) + self._test_dist_optim_base(optim.Rprop, lr=0.05) + + def _test_dist_optim_none_grads(self, optim_cls, *args, **kwargs): + # local version + module1 = MyModule() + module2 = MyModule(requires_grad=False) + params = [module1.get_w(), module2.get_w()] + local_optim = optim_cls(params, *args, **kwargs) + + old_w1 = module1.w.detach().clone() + old_w2 = module2.w.detach().clone() + + g_cpu = torch.Generator() + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = module1.forward(t2) + output2 = module2.forward(output1) + loss = torch.add(output2, t1).sum() + + loss.backward() + local_optim.step() + + # distributed version + owner1 = f"worker{(self.rank + 1) % self.world_size:d}" + owner2 = f"worker{(self.rank + 2) % self.world_size:d}" + + remote_module1 = rpc.remote(owner1, MyModule) + remote_module2 = rpc.remote(owner2, MyModule, args=(False,)) + remote_param1 = remote_module1.remote().get_w() + remote_param2 = remote_module2.remote().get_w() + + # sanity check: local and remote initial weights should match + self.assertEqual(old_w1, remote_param1.to_here()) + self.assertEqual(old_w2, remote_param2.to_here()) + + dist_optim = DistributedOptimizer( + optim_cls, [remote_param1, remote_param2], *args, **kwargs + ) + + with dist_autograd.context() as context_id: + g_cpu.manual_seed(0) + t1 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + t2 = torch.rand((3, 3), requires_grad=True, generator=g_cpu) + output1 = remote_module1.rpc_async().forward(t2) + output2 = remote_module2.rpc_async().forward(output1.wait()) + loss = torch.add(output2.wait(), t1) + + dist_autograd.backward(context_id, [loss.sum()]) + dist_optim.step(context_id) + + new_w1 = remote_module1.rpc_async().get_w().wait() + new_w2 = remote_module2.rpc_async().get_w().wait() + + # ensure optimizer changed weights for w1 + self.assertNotEqual(old_w1, new_w1) + + # ensure optimizer not changed weights for w2 + self.assertEqual(old_w2, new_w2) + # ensure local equals remote + self.assertEqual(new_w1, module1.get_w()) + self.assertEqual(new_w2, module2.get_w()) + + @dist_init() + def test_dist_optim_none_grads(self): + self._test_dist_optim_none_grads(optim.SGD, lr=0.05) + self._test_dist_optim_none_grads(optim.RMSprop, lr=0.05) + self._test_dist_optim_none_grads(optim.Rprop, lr=0.05) + self._test_dist_optim_none_grads(optim.Adadelta, rho=0.95) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/examples/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/examples/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/examples/parameter_server_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/examples/parameter_server_test.py new file mode 100644 index 0000000000000000000000000000000000000000..ad0b7fbe2207f8533da1eba8c23cda513f2bcf25 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/examples/parameter_server_test.py @@ -0,0 +1,140 @@ +# mypy: allow-untyped-defs + +# If you need to modify this file to make this test pass, please also apply same edits accordingly to +# https://github.com/pytorch/examples/blob/master/distributed/rpc/batch/parameter_server.py +# and https://pytorch.org/tutorials/intermediate/rpc_async_execution.html#batch-updating-parameter-server + +import threading +from datetime import datetime +from time import perf_counter + +import torch +import torch.distributed.rpc as rpc +import torch.nn as nn +from torch import optim +from torch.testing._internal.dist_utils import dist_init, worker_name +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +batch_size = 20 +in_features = 100 +out_features = 30 +num_batches = 4 + + +def timed_log(text): + print(f"{datetime.now().strftime('%H:%M:%S')} {text}") + + +class BatchUpdateParameterServer: + def __init__(self, batch_update_size): + self.model = nn.Linear(in_features, out_features) + self.lock = threading.Lock() + self.future_model = torch.futures.Future() + self.batch_update_size = batch_update_size + self.curr_update_size = 0 + self.optimizer = optim.SGD(self.model.parameters(), lr=0.001, momentum=0.9) + for p in self.model.parameters(): + p.grad = torch.zeros_like(p) + + def get_model(self): + return self.model + + @staticmethod + @rpc.functions.async_execution + def update_and_fetch_model(ps_rref, grads): + self = ps_rref.local_value() + for p, g in zip(self.model.parameters(), grads, strict=True): + if p.grad is None: + p.grad = g + else: + p.grad += g + with self.lock: + timed_log( + f"PS got {self.curr_update_size}/{self.batch_update_size} updates" + ) + self.curr_update_size += 1 + fut = self.future_model + + if self.curr_update_size >= self.batch_update_size: + for p in self.model.parameters(): + p.grad /= self.batch_update_size + self.curr_update_size = 0 + self.optimizer.step() + self.optimizer.zero_grad() + fut.set_result(self.model) + timed_log("PS updated model") + self.future_model = torch.futures.Future() + + return fut + + +class Trainer: + def __init__(self, ps_rref): + self.ps_rref = ps_rref + self.loss_fn = nn.L1Loss() + + def get_next_batch(self): + for _ in range(num_batches): + inputs = torch.randn(batch_size, in_features) + labels = torch.zeros(batch_size, out_features) + yield inputs, labels + + def train(self): + name = rpc.get_worker_info().name + m = self.ps_rref.rpc_sync().get_model() + for inputs, labels in self.get_next_batch(): + timed_log(f"{name} processing one batch") + self.loss_fn(m(inputs), labels).backward() + timed_log(f"{name} reporting grads") + m = rpc.rpc_sync( + self.ps_rref.owner(), + BatchUpdateParameterServer.update_and_fetch_model, + args=(self.ps_rref, [p.grad for p in m.cpu().parameters()]), + ) + timed_log(f"{name} got updated model") + + +def run_trainer(ps_rref): + trainer = Trainer(ps_rref) + trainer.train() + + +def run_ps(trainers): + timed_log("Start training") + start = perf_counter() + ps_rref = rpc.RRef(BatchUpdateParameterServer(len(trainers))) + futs = [ + rpc.rpc_async(trainer, run_trainer, args=(ps_rref,)) for trainer in trainers + ] + + torch.futures.wait_all(futs) + stop = perf_counter() + timed_log("Finish training") + timed_log(f"Time spent training: {stop - start}s") + + +class ParameterServerTest(RpcAgentTestFixture): + @dist_init(setup_rpc=False) + def test_batch_updating_parameter_server(self): + if self.rank != 0: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + else: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + run_ps([f"{worker_name(r)}" for r in range(1, self.world_size)]) + + rpc.shutdown() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/examples/reinforcement_learning_rpc_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/examples/reinforcement_learning_rpc_test.py new file mode 100644 index 0000000000000000000000000000000000000000..57008aed17dba34aacbc3b8a7a5b62c6dcbb5526 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/examples/reinforcement_learning_rpc_test.py @@ -0,0 +1,265 @@ +# mypy: allow-untyped-defs + +# If you need to modify this file to make this test pass, please also apply same edits accordingly to +# https://github.com/pytorch/examples/blob/master/distributed/rpc/rl/main.py +# and https://pytorch.org/tutorials/intermediate/rpc_tutorial.html + +import numpy as np + +import torch +import torch.distributed.rpc as rpc +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.distributed.rpc import remote, rpc_async, rpc_sync, RRef +from torch.distributions import Categorical +from torch.testing._internal.dist_utils import dist_init, worker_name +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +TOTAL_EPISODE_STEP = 5000 +GAMMA = 0.1 +SEED = 543 + + +def _call_method(method, rref, *args, **kwargs): + r""" + a helper function to call a method on the given RRef + """ + return method(rref.local_value(), *args, **kwargs) + + +def _remote_method(method, rref, *args, **kwargs): + r""" + a helper function to run method on the owner of rref and fetch back the + result using RPC + """ + args = [method, rref] + list(args) + return rpc_sync(rref.owner(), _call_method, args=args, kwargs=kwargs) + + +class Policy(nn.Module): + r""" + Borrowing the ``Policy`` class from the Reinforcement Learning example. + Copying the code to make these two examples independent. + See https://github.com/pytorch/examples/tree/master/reinforcement_learning + """ + + def __init__(self) -> None: + super().__init__() + self.affine1 = nn.Linear(4, 128) + self.dropout = nn.Dropout(p=0.6) + self.affine2 = nn.Linear(128, 2) + + self.saved_log_probs = [] + self.rewards = [] + + def forward(self, x): + x = self.affine1(x) + x = self.dropout(x) + x = F.relu(x) + action_scores = self.affine2(x) + return F.softmax(action_scores, dim=1) + + +class DummyEnv: + r""" + A dummy environment that implements the required subset of the OpenAI gym + interface. It exists only to avoid a dependency on gym for running the + tests in this file. It is designed to run for a set max number of iterations, + returning random states and rewards at each step. + """ + + def __init__(self, state_dim=4, num_iters=10, reward_threshold=475.0): + self.state_dim = state_dim + self.num_iters = num_iters + self.iter = 0 + self.reward_threshold = reward_threshold + + def seed(self, manual_seed): + torch.manual_seed(manual_seed) + + def reset(self): + self.iter = 0 + return torch.randn(self.state_dim) + + def step(self, action): + self.iter += 1 + state = torch.randn(self.state_dim) + reward = torch.rand(1).item() * self.reward_threshold + done = self.iter >= self.num_iters + info = {} + return state, reward, done, info + + +class Observer: + r""" + An observer has exclusive access to its own environment. Each observer + captures the state from its environment, and send the state to the agent to + select an action. Then, the observer applies the action to its environment + and reports the reward to the agent. + """ + + def __init__(self) -> None: + self.id = rpc.get_worker_info().id + self.env = DummyEnv() + self.env.seed(SEED) + + def run_episode(self, agent_rref, n_steps): + r""" + Run one episode of n_steps. + Arguments: + agent_rref (RRef): an RRef referencing the agent object. + n_steps (int): number of steps in this episode + """ + state, _ep_reward = self.env.reset(), 0 + for _ in range(n_steps): + # send the state to the agent to get an action + action = _remote_method(Agent.select_action, agent_rref, self.id, state) + + # apply the action to the environment, and get the reward + state, reward, done, _ = self.env.step(action) + + # report the reward to the agent for training purpose + _remote_method(Agent.report_reward, agent_rref, self.id, reward) + + if done: + break + + +class Agent: + def __init__(self, world_size): + self.ob_rrefs = [] + self.agent_rref = RRef(self) + self.rewards = {} + self.saved_log_probs = {} + self.policy = Policy() + self.optimizer = optim.Adam(self.policy.parameters(), lr=1e-2) + self.eps = np.finfo(np.float32).eps.item() + self.running_reward = 0 + self.reward_threshold = DummyEnv().reward_threshold + for ob_rank in range(1, world_size): + ob_info = rpc.get_worker_info(worker_name(ob_rank)) + self.ob_rrefs.append(remote(ob_info, Observer)) + self.rewards[ob_info.id] = [] + self.saved_log_probs[ob_info.id] = [] + + def select_action(self, ob_id, state): + r""" + This function is mostly borrowed from the Reinforcement Learning example. + See https://github.com/pytorch/examples/tree/master/reinforcement_learning + The main difference is that instead of keeping all probs in one list, + the agent keeps probs in a dictionary, one key per observer. + + NB: no need to enforce thread-safety here as GIL will serialize + executions. + """ + probs = self.policy(state.unsqueeze(0)) + m = Categorical(probs) + action = m.sample() + self.saved_log_probs[ob_id].append(m.log_prob(action)) + return action.item() + + def report_reward(self, ob_id, reward): + r""" + Observers call this function to report rewards. + """ + self.rewards[ob_id].append(reward) + + def run_episode(self, n_steps=0): + r""" + Run one episode. The agent will tell each observer to run n_steps. + """ + # make async RPC to kick off an episode on all observers + futs = [ + rpc_async( + ob_rref.owner(), + _call_method, + args=(Observer.run_episode, ob_rref, self.agent_rref, n_steps), + ) + for ob_rref in self.ob_rrefs + ] + + # wait until all observers have finished this episode + for fut in futs: + fut.wait() + + def finish_episode(self): + r""" + This function is mostly borrowed from the Reinforcement Learning example. + See https://github.com/pytorch/examples/tree/master/reinforcement_learning + The main difference is that it joins all probs and rewards from + different observers into one list, and uses the minimum observer rewards + as the reward of the current episode. + """ + + # joins probs and rewards from different observers into lists + R, probs, rewards = 0, [], [] + for ob_id in self.rewards: + probs.extend(self.saved_log_probs[ob_id]) + rewards.extend(self.rewards[ob_id]) + + # use the minimum observer reward to calculate the running reward + min_reward = min(sum(self.rewards[ob_id]) for ob_id in self.rewards) + self.running_reward = 0.05 * min_reward + (1 - 0.05) * self.running_reward + + # clear saved probs and rewards + for ob_id in self.rewards: + self.rewards[ob_id] = [] + self.saved_log_probs[ob_id] = [] + + policy_loss, returns = [], [] + for r in rewards[::-1]: + R = r + GAMMA * R + returns.insert(0, R) + returns = torch.tensor(returns) + returns = (returns - returns.mean()) / (returns.std() + self.eps) + for log_prob, R in zip(probs, returns, strict=True): + policy_loss.append(-log_prob * R) + self.optimizer.zero_grad() + policy_loss = torch.cat(policy_loss).sum() + policy_loss.backward() + self.optimizer.step() + return min_reward + + +def run_agent(agent, n_steps): + while True: + agent.run_episode(n_steps=n_steps) + agent.finish_episode() + + if agent.running_reward > agent.reward_threshold: + print(f"Solved! Running reward is now {agent.running_reward}!") + break + + +class ReinforcementLearningRpcTest(RpcAgentTestFixture): + @dist_init(setup_rpc=False) + def test_rl_rpc(self): + if self.rank == 0: + # Rank 0 is the agent. + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + agent = Agent(self.world_size) + run_agent(agent, n_steps=int(TOTAL_EPISODE_STEP / (self.world_size - 1))) + + # Ensure training was run. We don't really care about whether the task was learned, + # since the purpose of the test is to check the API calls. + self.assertGreater(agent.running_reward, 0.0) + else: + # Other ranks are observers that passively wait for instructions from the agent. + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + rpc.shutdown() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/faulty_agent_rpc_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/faulty_agent_rpc_test.py new file mode 100644 index 0000000000000000000000000000000000000000..747155e3e1cbce8f8e8c14756fe3f98bf22a8987 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/faulty_agent_rpc_test.py @@ -0,0 +1,337 @@ +# mypy: allow-untyped-defs + +import time + +import torch +import torch.distributed.rpc as rpc +from torch.distributed.rpc.api import _delete_all_user_and_unforked_owner_rrefs +from torch.testing._internal.dist_utils import ( + dist_init, + wait_until_owners_and_forks_on_rank, + wait_until_pending_futures_and_users_flushed, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +def my_sleep_func(seconds=1): + time.sleep(seconds) + return torch.mul(torch.tensor(1), torch.tensor(1)) + + +@torch.jit.script +def my_script_func(tensor): + return torch.add(tensor, tensor) + + +def add_rref_to_value(rref, value): + return rref.to_here() + value + + +class FaultyAgentRpcTest(RpcAgentTestFixture): + # no faulty_messages defined so this fails all retryable messages - see + # faulty_rpc_agent_test_fixture.py for the list of retryable messages. + @dist_init(messages_to_delay={}) + def test_check_failed_messages(self): + if self.rank == 0: + dst_worker_b = worker_name((self.rank + 1) % self.world_size) + dst_worker_c = worker_name((self.rank + 2) % self.world_size) + + # Worker0 sends RPC to Worker1 and creates an RRef there + rref = rpc.remote( + dst_worker_b, torch.add, args=(torch.ones(2, 2), torch.ones(2, 2)) + ) + # Worker0 sends an RPC to Worker2 with the RRef as an arg + rpc.remote(dst_worker_c, add_rref_to_value, args=(rref, torch.ones(2, 2))) + # check if the output is as expected + self.assertEqual( + rref.to_here(), torch.add(torch.ones(2, 2), torch.ones(2, 2)) + ) + # explicitly delete all User RRefs + _delete_all_user_and_unforked_owner_rrefs() + + @dist_init + def test_verify_backend_options(self): + self.assertEqual( + self.rpc_backend, rpc.backend_registry.BackendType.FAULTY_TENSORPIPE + ) + self.assertEqual(self.rpc_backend_options.num_worker_threads, 8) + self.assertEqual(self.rpc_backend_options.num_fail_sends, 3) + self.assertEqual(len(self.rpc_backend_options.messages_to_fail), 4) + self.assertEqual(len(self.rpc_backend_options.messages_to_delay), 2) + self.assertEqual( + self.rpc_backend_options.rpc_timeout, rpc.constants.DEFAULT_RPC_TIMEOUT_SEC + ) + + @dist_init(faulty_messages=["RREF_FORK_REQUEST", "RREF_CHILD_ACCEPT"]) + def test_custom_faulty_messages(self): + self.assertEqual( + {"RREF_FORK_REQUEST", "RREF_CHILD_ACCEPT"}, + set(self.rpc_backend_options.messages_to_fail), + ) + + @dist_init(faulty_messages=[]) + def test_no_faulty_messages(self): + self.assertEqual(len(self.rpc_backend_options.messages_to_fail), 0) + + @dist_init(messages_to_delay={"SCRIPT_CALL": 1.5}) + def test_custom_messages_to_delay(self): + self.assertEqual( + self.rpc_backend_options.messages_to_delay, {"SCRIPT_CALL": 1.5} + ) + + def _test_remote_message_dropped_pickle(self, dst=None): + if self.rank != 0: + return + dst_rank = dst if dst is not None else (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + # Since we fail python_remote_call messages synchronously, the future + # corresponding to this remote call will be marked with an error when + # this function returns. + rref = rpc.remote(dst_worker, my_sleep_func, args=(1,)) + # Call to ensure pending callbacks are run. + wait_until_pending_futures_and_users_flushed() + # Attempt to fork the RRef should raise an error indicating the rpc.remote timeout. + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref._serialize() + # Test that using RRef as arg over RPC (which forks) results in the same + # error + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rpc.rpc_async(dst_worker, add_rref_to_value, args=(rref, 1)) + + @dist_init(faulty_messages=["PYTHON_REMOTE_CALL"]) + def test_remote_message_dropped_pickle(self): + self._test_remote_message_dropped_pickle() + + @dist_init(faulty_messages=["PYTHON_REMOTE_CALL"]) + def test_remote_message_dropped_pickle_to_self(self): + self._test_remote_message_dropped_pickle(self.rank) + + def _test_remote_message_dropped_timeout(self, func, args, dst=None): + if self.rank != 0: + return + + # test the case where rpc.remote() message creation is completely dropped. + dst_rank = dst if dst is not None else (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + # Since we fail python_remote_call messages synchronously, the future + # corresponding to this remote call will be marked with an error when + # this function returns. + rref = rpc.remote(dst_worker, func, args=args) + # Call to ensure pending callbacks are run. + wait_until_pending_futures_and_users_flushed() + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref.to_here() + # Note: during shutdown, logs will indicate "Could not find OwnerRRef..." + # on the owning nodes, this is expected because the OwnerRRef was never + # successfully created. Therefore, delAllUsers will work as expected. + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_builtin_remote_message_dropped_timeout(self): + func = torch.add + args = (torch.tensor(1), torch.tensor(1)) + self._test_remote_message_dropped_timeout(func, args) + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_builtin_remote_message_dropped_timeout_to_self(self): + func = torch.add + args = (torch.tensor(1), torch.tensor(1)) + self._test_remote_message_dropped_timeout(func, args, dst=0) + + @dist_init(faulty_messages=["PYTHON_REMOTE_CALL"]) + def test_udf_remote_message_dropped_timeout(self): + func = my_sleep_func + args = (2,) + self._test_remote_message_dropped_timeout(func, args) + + @dist_init(faulty_messages=["PYTHON_REMOTE_CALL"]) + def test_udf_remote_message_dropped_timeout_to_self(self): + func = my_sleep_func + args = (2,) + self._test_remote_message_dropped_timeout(func, args, dst=0) + + def _test_remote_message_delay_timeout(self, func, args, dst=None): + if self.rank != 0: + return + # Test the case where remote message is eventually processed on the owner, + # but the future on the creator times out before the response comes back. + dst_rank = dst if dst is not None else (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + # 10 ms timeout + rref = rpc.remote(dst_worker, func, args=args, timeout=0.001) + # Future corresponding to the remote creation should time out. + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rref._get_future().wait() + + # Call to ensure pending callbacks are run. + wait_until_pending_futures_and_users_flushed() + # to_here() should now pick up that rpc.remote() creation has failed. + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref.to_here() + + # Test the case where rpc.remote() times out, but to_here() has already + # started blocking before. + # NOTE: we only test this when not sending to self, as to_here() calls + # calls localValue(), which does not send an RPC and thus does not have + # a timeout. This can be supported by allowing future.wait() to + # take in an optional timeout (https://github.com/pytorch/pytorch/issues/39280) + if dst_rank != self.rank: + slow_rref = rpc.remote(dst_worker, func, args=args, timeout=2) + + with self.assertRaisesRegex(RuntimeError, expected_error): + # to_here() should raise timeout error, since it does not know about the + # status of rpc.remote(). + slow_rref.to_here(0.001) + # Note: If we proceed with shutdown, UserRRef will send out a RRefUserDelete + # but this can be a noop since it may not exist on the owner yet. Later, + # the owner can process the RRef creation and wait for the delete message, + # thus leading to a timeout. + # Therefore, we wait until we get notification that pending owners have + # been confirmed before sending out RRefUserDeletes. + if dst_rank != self.rank: + wait_until_owners_and_forks_on_rank(2, 2, rank=dst_rank) + + @dist_init(faulty_messages=[], messages_to_delay={"PYTHON_REMOTE_CALL": 2}) + def test_udf_remote_message_delay_timeout(self): + func = my_sleep_func + args = (2,) + self._test_remote_message_delay_timeout(func, args) + + @dist_init(faulty_messages=[], messages_to_delay={"PYTHON_REMOTE_CALL": 2}) + def test_udf_remote_message_delay_timeout_to_self(self): + func = my_sleep_func + args = (1,) + self._test_remote_message_delay_timeout(func, args, dst=0) + + @dist_init( + faulty_messages=[], + messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1}, + ) + def test_remote_message_builtin_delay_timeout(self): + func = torch.add + args = (torch.tensor(1), torch.tensor(1)) + self._test_remote_message_delay_timeout(func, args) + + @dist_init( + faulty_messages=[], + messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1}, + ) + def test_remote_message_builtin_delay_timeout_to_self(self): + func = torch.add + args = (torch.tensor(1), torch.tensor(1)) + self._test_remote_message_delay_timeout(func, args, dst=0) + + @dist_init( + faulty_messages=[], + messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1}, + ) + def test_remote_message_script_delay_timeout(self): + func = my_script_func + args = (torch.tensor(1),) + self._test_remote_message_delay_timeout(func, args) + + @dist_init( + faulty_messages=[], + messages_to_delay={"SCRIPT_REMOTE_CALL": 2, "SCRIPT_RREF_FETCH_CALL": 1}, + ) + def test_remote_message_script_delay_timeout_to_self(self): + func = my_script_func + args = (torch.tensor(1),) + self._test_remote_message_delay_timeout(func, args, dst=0) + + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_RREF_FETCH_CALL": 1}) + def test_rref_to_here_timeout(self): + if self.rank != 0: + return + + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rref.to_here(0.01) + + rref.to_here() + + @dist_init(faulty_messages=[]) + def test_rpc_builtin_timeout(self): + next_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(next_rank) + expected_error = self.get_timeout_error_regex() + # PYTHON_CALL message types which correspond to Python UDF over RPC + # by default get a delay (see faulty_rpc_agent_test_fixture) + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.rpc_sync( + dst_worker, + torch.add, + args=(torch.tensor(1), torch.tensor(1)), + timeout=1, + ) + + fut = rpc.rpc_async( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)), timeout=1 + ) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure that the currently set default timeout is large enough such + # that RPCs with delays still complete. + fut = rpc.rpc_async( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + fut.wait() + + # Ensure timeout if we set a new default and don't override + rpc._set_rpc_timeout(0.001) + fut = rpc.rpc_async( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure run to completion if we specify timeout of 0 + fut = rpc.rpc_async( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)), timeout=0 + ) + fut.wait() + # Reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_CALL": 1.5}) + def test_rpc_script_timeout(self): + next_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(next_rank) + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.rpc_sync(dst_worker, my_script_func, args=(torch.tensor(1),), timeout=1) + + fut = rpc.rpc_async( + dst_worker, my_script_func, args=(torch.tensor(1),), timeout=1 + ) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure that the currently set default timeout is large enough such + # that RPCs with delays still complete. + fut = rpc.rpc_async(dst_worker, my_script_func, args=(torch.tensor(1),)) + fut.wait() + + # Ensure timeout if we set a new default and don't override + rpc._set_rpc_timeout(0.001) + fut = rpc.rpc_async(dst_worker, my_script_func, args=(torch.tensor(1),)) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure run to completion if we specify timeout of 0 + rpc._set_rpc_timeout(0.001) + fut = rpc.rpc_async( + dst_worker, my_script_func, args=(torch.tensor(1),), timeout=0 + ) + fut.wait() + # Reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/faulty_rpc_agent_test_fixture.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/faulty_rpc_agent_test_fixture.py new file mode 100644 index 0000000000000000000000000000000000000000..aff7d556d10621e7290c07ecb433b865d7133bb2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/faulty_rpc_agent_test_fixture.py @@ -0,0 +1,64 @@ +# mypy: allow-untyped-defs + +import torch.distributed.rpc as rpc +import torch.distributed.rpc._testing # noqa: F401 +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +# The following message types are currently retried in the RREF protocol and +# distributed autograd. Thus only these messages should be tested with the +# Faulty RPC Agent. +retryable_message_types = [ + "RREF_FORK_REQUEST", + "RREF_CHILD_ACCEPT", + "RREF_USER_DELETE", + "CLEANUP_AUTOGRAD_CONTEXT_REQ", +] + +# The following messages incur the corresponding delay in seconds while being +# processed in FaultyTensorPipeAgent's enqueueSend() function. +default_messages_to_delay = { + "PYTHON_CALL": 1.5, # Python UDF + "SCRIPT_CALL": 1.5, # Script/Builtin +} + + +class FaultyRpcAgentTestFixture(RpcAgentTestFixture): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.messages_to_fail = retryable_message_types + self.messages_to_delay = default_messages_to_delay + + @property + def rpc_backend(self): + return rpc.backend_registry.BackendType["FAULTY_TENSORPIPE"] + + @property + def rpc_backend_options(self): + return rpc.backend_registry.construct_rpc_backend_options( + self.rpc_backend, + init_method=self.init_method, + num_worker_threads=8, + num_fail_sends=3, + messages_to_fail=self.messages_to_fail, + messages_to_delay=self.messages_to_delay, + ) + + def setup_fault_injection(self, faulty_messages, messages_to_delay): + if faulty_messages is not None: + self.messages_to_fail = faulty_messages + if messages_to_delay is not None: + self.messages_to_delay = messages_to_delay + + def get_shutdown_error_regex(self): + error_regexes = [ + "Exception in thread pool task", + "Connection reset by peer", + "Connection closed by peer", + ] + return "|".join([f"({error_str})" for error_str in error_regexes]) + + def get_timeout_error_regex(self): + return "RPC ran for more than" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/dist_autograd_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/dist_autograd_test.py new file mode 100644 index 0000000000000000000000000000000000000000..fde1fe2355c2968e1b351b288d20c674835b0ca2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/dist_autograd_test.py @@ -0,0 +1,113 @@ +# mypy: allow-untyped-defs + + +import torch +import torch.distributed.autograd as dist_autograd +import torch.distributed.rpc as rpc +from torch import Tensor +from torch.distributed.rpc import rpc_async +from torch.testing import FileCheck +from torch.testing._internal.dist_utils import dist_init, worker_name +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +@torch.jit.script +def local_add(t1, t2): + return torch.add(t1, t2) + + +@torch.jit.script +def remote_add(t1, t2, dst: str): # noqa: E999 + return rpc_async(dst, local_add, (t1, t2)).wait() + + +@torch.jit.script +def fork_add(t1, t2, dst: str): + fut = torch.jit._fork(remote_add, t1, t2, dst) + return torch.jit._wait(fut) + + +class JitDistAutogradTest(RpcAgentTestFixture): + @dist_init + def test_get_gradients(self): + @torch.jit.script + def dist_get_gradients(context_id: int) -> dict[Tensor, Tensor]: + return dist_autograd.get_gradients(context_id) + + FileCheck().check("get_gradients").run(str(dist_get_gradients.graph)) + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + t3 = torch.add(t1, t2) + + dist_autograd.backward(context_id, [t3.sum()]) + grads = dist_get_gradients(context_id) + + self.assertEqual(2, len(grads)) + self.assertIn(t1, grads) + self.assertIn(t2, grads) + self.assertEqual(torch.ones(3, 3), grads[t1]) + self.assertEqual(torch.ones(3, 3), grads[t2]) + + @dist_init + def test_dist_backward(self): + if self.rank != 0: + return + + @torch.jit.script + def dist_backward_script(context_id: int, loss: torch.Tensor): + dist_autograd.backward(context_id, [loss]) + + FileCheck().check("dist_backward").run(str(dist_backward_script.graph)) + with dist_autograd.context() as context_id: + t1 = torch.rand(3, 3, requires_grad=True) + t2 = torch.rand(3, 3, requires_grad=True) + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + loss = rpc.rpc_sync(dst_worker_name, torch.add, args=(t1, t2)).sum() + dist_backward_script(context_id, loss) + + @dist_init + def test_jit_fork_within_context(self): + with dist_autograd.context() as context_id: + t1 = torch.rand((3, 3), requires_grad=True) + t2 = torch.rand((3, 3), requires_grad=True) + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + res = fork_add(t1, t2, dst_worker_name) + loss = res.sum() + dist_autograd.backward(context_id, [loss]) + + grads = dist_autograd.get_gradients(context_id) + self.assertEqual(2, len(grads)) + self.assertIn(t1, grads) + self.assertIn(t2, grads) + + @dist_init + def test_restore_context_after_swtich_to_jit_thread(self): + if self.rank != 0: + return + + @torch.jit.script + def forward_script( + context_id: int, dst_worker_name: str, t1: Tensor, t2: Tensor + ) -> tuple[Tensor, Tensor]: + res1_fut = rpc.rpc_async(dst_worker_name, local_add, (t1, t1)) + res1 = res1_fut.wait() # After this, the script runs in a new JIT thread. + loss1 = res1.sum() + + # SendRpcBackward is not attached, since DistAutogradContext is lost here. + res2_fut = rpc.rpc_async(dst_worker_name, local_add, (t2, t2)) + res2 = res2_fut.wait() + loss2 = res2.sum() + + return loss1, loss2 + + with dist_autograd.context() as context_id: + t1 = torch.ones((2, 3), requires_grad=True) + t2 = torch.ones((2, 3), requires_grad=True) + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + loss0, loss1 = forward_script(context_id, dst_worker_name, t1, t2) + dist_autograd.backward(context_id, [loss0, loss1]) + grad0, grad1 = dist_autograd.get_gradients(context_id) + self.assertEqual(grad0, grad1) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test.py new file mode 100644 index 0000000000000000000000000000000000000000..82a5d66e87f38672fe7076075b764a094bb81b4c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test.py @@ -0,0 +1,1384 @@ +# mypy: allow-untyped-defs + +import io +import time +from typing import Any + +import torch +import torch.distributed as dist +import torch.distributed.rpc as rpc +from torch import Tensor +from torch.autograd.profiler import record_function +from torch.autograd.profiler_legacy import profile as _profile +from torch.distributed.rpc import RRef +from torch.distributed.rpc.internal import _build_rpc_profiling_key, RPCExecMode +from torch.futures import Future +from torch.testing._internal.common_utils import TemporaryFileName +from torch.testing._internal.dist_utils import ( + dist_init, + get_function_event, + initialize_pg, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +def rref_isinstance(rref, cls_to_check): + return isinstance(rref.local_value(), cls_to_check) + + +def sleep(t): + time.sleep(t) + + +def rpc_return_rref(dst): + return rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 1)) + + +@torch.jit.script +def rref_local_value(rref: RRef[Tensor]) -> Tensor: + return rref.local_value() + + +@torch.jit.script +def list_create() -> list[int]: + global_list = [1, 2, 3] + return global_list + + +@torch.jit.script +def rref_list_mutate(rref: RRef[list[int]]) -> None: + rref.local_value().append(4) + rref.to_here().append(5) + rref.to_here(5.0).append(6) + + +def return_value(value: int) -> int: + return value + + +class RRefAPITest: + @dist_init + def test_rref_is_owner(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + rref_var = rpc_return_rref(dst_worker_name) + + @torch.jit.script + def rref_tensor_is_owner(rref_var: RRef[Tensor]) -> bool: + return rref_var.is_owner() + + res = rref_tensor_is_owner(rref_var) + self.assertEqual(res, False) + + @dist_init + def test_rref_local_value(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + rref = rpc_return_rref(dst_worker_name) + + with self.assertRaisesRegex( + RuntimeError, r"Can't call RRef.local_value\(\) on a non-owner RRef" + ): + rref_local_value(rref) + + ret = rpc.rpc_sync(dst_worker_name, rref_local_value, (rref,)) + self.assertEqual(ret, torch.add(torch.ones(2, 2), 1)) + + @dist_init + def test_local_rref_local_value(self): + if self.rank != 0: + return + + dst_worker_name = worker_name(self.rank) + rref = rpc.remote(dst_worker_name, return_value, (5,), {}) + + ret = rref_local_value(rref) + self.assertEqual(ret, 5) + + def _create_rref(self): + owner_rank = (self.rank + 2) % self.world_size + return rpc.remote( + worker_name(owner_rank), torch.add, args=(torch.zeros(2, 2), 1) + ) + + @dist_init + def test_user_rrefs_confirmed(self): + dst_rank = (self.rank + 1) % self.world_size + rref = self._create_rref() + ret = rpc.rpc_sync( + worker_name(dst_rank), script_check_rref_confirmed, args=(rref,) + ) + self.assertEqual(ret, True) + + @dist_init + def test_user_rrefs_confirmed_remote(self): + dst_rank = (self.rank + 1) % self.world_size + rref = self._create_rref() + ret_rref = rpc.remote( + worker_name(dst_rank), script_check_rref_confirmed, args=(rref,) + ) + self.assertEqual(ret_rref.to_here(), True) + + @dist_init + def test_rref_list_mutate(self): + dst = worker_name((self.rank + 1) % self.world_size) + list_rref = rpc.remote(dst, list_create) + + rpc.rpc_sync(dst, rref_list_mutate, args=(list_rref,)) + self.assertEqual(list_rref.to_here(), [1, 2, 3, 4, 5, 6]) + + +@torch.jit.script +def no_arg(): + return 0 + + +@torch.jit.script +def one_arg(value): + return value + 1 + + +@torch.jit.script +def script_add_ones(x): + return torch.add(x, torch.ones(1)) + + +@torch.jit.script +def script_add_ones_with_record_function(x, block: str): + with record_function(block): + return torch.add(x, torch.ones(1)) + + +@torch.jit.script +def record_function_on_caller_rpc_async(dst_worker_name: str, block: str) -> Tensor: + t: Tensor = torch.ones(1) + with record_function(block): + fut1 = rpc.rpc_async(dst_worker_name, script_add_ones, (t,)) + # Extra operator call to avoid de-duplication of the next async call + # see https://github.com/pytorch/pytorch/pull/62710#discussion_r694680279 + zero = torch.zeros_like(t) + fut2 = rpc.rpc_async(dst_worker_name, script_add_ones, (t,)) + res = fut1.wait() + fut2.wait() + zero + return res + + +@torch.jit.script +def script_fork_wait_udf(tensor): + fut = torch.jit._fork(script_add_ones, tensor) + x = torch.jit._wait(fut) + return x + + +@torch.jit.script +def rref_to_here(rref_var: RRef[Tensor]) -> Tensor: + return rref_var.to_here() + + +@torch.jit.script +def return_rref(rref_var: RRef[Tensor]) -> RRef[Tensor]: + return rref_var + + +@torch.jit.script +def script_raise_func(value): + if value.numel() == 2: + raise ValueError("Expected error") + return value + 1 + + +@torch.jit.script +def script_fork_wait_throw(invalue): + fut = torch.jit._fork(script_raise_func, invalue) + value = torch.jit._wait(fut) + return value + + +@torch.jit.script +def call_rpc_with_profiling( + record: torch.classes.profiler._RecordFunction, dst_worker_name: str +) -> Tensor: + # Call rpc_async from within ScriptFunction and ensure that we can attach + # profiling callbacks. Note that handle here is a Tensor representation of + # RecordFunction. + fut = rpc.rpc_async(dst_worker_name, one_arg, (torch.tensor(1),)) + torch.ops.profiler._call_end_callbacks_on_jit_fut(record, fut) + ret = fut.wait() + return ret + + +@torch.jit.script +def call_rpc_torchscript_with_record_function( + dst_worker_name: str, block: str +) -> Tensor: + fut = rpc.rpc_async( + dst_worker_name, script_add_ones_with_record_function, (torch.tensor(1), block) + ) + return fut.wait() + + +@torch.jit.script +def call_fork_with_profiling(record: torch.classes.profiler._RecordFunction) -> Tensor: + # Call fork from within ScriptFunction and ensure that we can attach profiling + # callbacks to the resulting future. Note that handle here is a Tensor + # representation of RecordFunction. + fut = torch.jit._fork(one_arg, torch.tensor(1)) + torch.ops.profiler._call_end_callbacks_on_jit_fut(record, fut) + ret = fut.wait() + return ret + + +class MyScriptModuleWithRRefs(torch.jit.ScriptModule): + def __init__(self, dst_worker): + super().__init__() + self.rrefs = [] + for _ in range(4): + self.rrefs.append(rpc_return_rref(dst_worker)) + + @torch.jit.script_method + def forward(self) -> Tensor: + res_tensor = torch.ones(2, 2) + for rref in self.rrefs: + res_tensor += rref.to_here() + + return res_tensor + + +@torch.jit.ignore +def rref_python_annotation(rref_var: RRef[Tensor]) -> RRef[Tensor]: + return rref_var + + +@torch.jit.script +def rref_script_annotation(rref_var: RRef[Tensor]) -> Tensor: + return rref_python_annotation(rref_var).to_here() + + +class RRefTypingTest: + @dist_init + def test_rref_as_arg_and_return(self): + n = self.rank + 1 + dst_rank = n % self.world_size + local_ret = one_arg(torch.ones(2, 2)) + + # create rref on current rank + rref = rpc.remote(worker_name(self.rank), one_arg, args=(torch.ones(2, 2),)) + + # pass rref to another user in rpc call + ret = rpc.rpc_sync(worker_name(dst_rank), rref_to_here, args=(rref,)) + self.assertEqual(ret, local_ret) + + # return rref in rpc call + rref1 = rpc.rpc_sync(worker_name(dst_rank), return_rref, args=(rref,)) + self.assertEqual(rref1.to_here(), local_ret) + + # pass rref to another user in remote call + rref2 = rpc.remote(worker_name(dst_rank), rref_to_here, args=(rref,)) + self.assertEqual(rref2.to_here(), local_ret) + + # return rref in remote call + rref3 = rpc.remote(worker_name(dst_rank), return_rref, args=(rref,)) + self.assertEqual(rref3.to_here().to_here(), local_ret) + + @dist_init + def test_my_script_module_with_rrefs(self): + n = self.rank + 1 + dst_rank = n % self.world_size + + module_with_rrefs = MyScriptModuleWithRRefs(worker_name(dst_rank)) + res = module_with_rrefs() + self.assertEqual(res, torch.ones(2, 2) * 9) + + @dist_init + def test_rref_python_annotation(self): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_var = rpc_return_rref(worker_name(dst_rank)) + + res = rref_script_annotation(rref_var) + self.assertEqual(res, torch.ones(2, 2) + 1) + + +class FutureTypingTest: + @dist_init + def test_future_passed_between_python_and_jit(self): + dst_rank = (self.rank + 1) % self.world_size + inputs = (torch.tensor([1, 1]), torch.tensor([2, 2])) + ret_fut = rpc.rpc_async(worker_name(dst_rank), two_args_two_kwargs, args=inputs) + expected_res = torch.tensor([10, 10]) + + @torch.jit.script + def future_wait_in_script(fut: Future[Tensor]) -> Tensor: + return fut.wait() + + self.assertEqual(future_wait_in_script(ret_fut), expected_res) + + @torch.jit.script + def future_return_to_python( + dst_rank: int, inputs: tuple[Tensor, Tensor] + ) -> Future[Tensor]: + return rpc.rpc_async(f"worker{dst_rank}", two_args_two_kwargs, inputs) + + fut_res = future_return_to_python(dst_rank, inputs) + self.assertEqual(fut_res.wait(), expected_res) + + @dist_init + def test_future_python_annotation(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + input_0 = torch.ones(2, 2) + input_1 = 1 + expected_res = torch.add(input_0, input_1) + + @torch.jit.ignore + def python_return_future() -> Future[Tensor]: + fut = rpc.rpc_async(dst_worker_name, torch.add, (input_0, input_1), {}) + return fut + + @torch.jit.script + def script_use_future() -> Tensor: + fut = python_return_future() + return fut.wait() + + res = script_use_future() + self.assertEqual(res, expected_res) + + +@torch.jit.script +class MyScriptClass: + def __init__(self, a: int): + self.a = a + + def get_value(self) -> int: + return self.a + + +@torch.jit.interface +class MyModuleInterface(torch.nn.Module): + def forward(self) -> Tensor: + # pyre-ignore[7]: Pyre and torch.jit.interface don't mix well + pass + + +class MyScriptModule(torch.jit.ScriptModule): + def __init__(self, rank): + super().__init__() + self.a = torch.ones(rank) + + @torch.jit.script_method + def forward(self) -> Tensor: + return self.a + + @torch.jit.script_method + def custom_func(self) -> Tensor: + return self.a + + +def owner_create_rref_my_script_class(a): + return rpc.RRef(MyScriptClass(a)) + + +def owner_create_rref_my_script_module(a): + return rpc.RRef(MyScriptModule(a), type_hint=MyModuleInterface) + + +@torch.jit.script +def script_rref_get_value_my_script_class(rref: RRef[MyScriptClass]) -> int: + return rref.to_here().get_value() + + +@torch.jit.script +def script_rref_run_forward_my_script_module(rref: RRef[MyModuleInterface]) -> Tensor: + return rref.to_here().forward() + + +class LocalRRefTest: + @dist_init + def test_create_local_script_class_rref_in_py(self): + if self.rank != 0: + return + + # Create a local RRef. + rref_script_class = rpc.RRef(MyScriptClass(self.rank)) + ret = rref_script_class.to_here().get_value() + self.assertEqual(ret, self.rank) + + @dist_init + def test_create_local_script_module_rref_in_py(self): + if self.rank != 0: + return + + # Create a local RRef. + rref_script_module = rpc.RRef(MyScriptModule(self.rank), MyModuleInterface) + ret = rref_script_module.to_here().forward() + self.assertEqual(ret, torch.ones(self.rank)) + + # Create a local RRef without type hint. + with self.assertRaisesRegex( + RuntimeError, + ( + "The RRef being created contains a ScriptModule, " + "must provide its ModuleInterface type hint." + ), + ): + rref_script_module = rpc.RRef(MyScriptModule(self.rank)) + + @dist_init + def test_return_local_script_class_rref_in_py_and_use_in_script(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + # Create a local RRef remotely in Python. + rref = rpc.rpc_sync( + dst_worker_name, owner_create_rref_my_script_class, args=(self.rank,) + ) + + def use_rref_on_owner(rref: RRef[MyScriptClass]) -> int: + args = (rref,) + kwargs: dict[str, Any] = {} + fut = rpc.rpc_async( + rref.owner(), script_rref_get_value_my_script_class, args, kwargs + ) + ret = fut.wait() + return ret + + # Use RRef in local Python RPC and remote Script run. + ret = use_rref_on_owner(rref) + self.assertEqual(ret, self.rank) + + # Use RRef in local Script RPC and remote Script run. + use_rref_on_owner_script = torch.jit.script(use_rref_on_owner) + ret = use_rref_on_owner_script(rref) + self.assertEqual(ret, self.rank) + + @dist_init + def test_return_local_script_module_rref_in_py_and_use_in_script(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + # Create a local RRef remotely in Python. + rref = rpc.rpc_sync( + dst_worker_name, owner_create_rref_my_script_module, args=(self.rank,) + ) + + def use_rref_on_owner(rref: RRef[MyModuleInterface]) -> Tensor: + args = (rref,) + kwargs: dict[str, Any] = {} + fut = rpc.rpc_async( + rref.owner_name(), + script_rref_run_forward_my_script_module, + args, + kwargs, + ) + ret = fut.wait() + return ret + + # Use RRef in local Python RPC and remote Script run. + ret = use_rref_on_owner(rref) + self.assertEqual(ret, torch.ones(self.rank)) + + # Use RRef in local Script RPC and remote Script run. + use_rref_on_owner_script = torch.jit.script(use_rref_on_owner) + ret = use_rref_on_owner_script(rref) + self.assertEqual(ret, torch.ones(self.rank)) + + +def python_function(): + return 0 + + +@torch.jit.script +def two_args_two_kwargs( + first_arg, + second_arg, + first_kwarg=torch.tensor([3, 3]), + second_kwarg=torch.tensor([4, 4]), +): + return first_arg + second_arg + first_kwarg + second_kwarg + + +@torch.jit.script +def assorted_types_args_kwargs( + tensor_arg: Tensor, # noqa: E999 + str_arg: str, + int_arg: int, + tensor_kwarg: Tensor = torch.tensor([2, 2]), + str_kwarg: str = "str_kwarg", + int_kwarg: int = 2, +): + return tensor_arg + tensor_kwarg, str_arg + str_kwarg, int_arg + int_kwarg + + +@torch.jit.script +def raise_script(): + raise RuntimeError("Expected error") + + +@torch.jit.script +def script_rpc_async_call( + dst_worker_name: str, args: tuple[Tensor, Tensor], kwargs: dict[str, Tensor] +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + + +@torch.jit.script +def script_rpc_sync_call( + dst_worker_name: str, args: tuple[Tensor, Tensor], kwargs: dict[str, Tensor] +): + res = rpc.rpc_sync(dst_worker_name, two_args_two_kwargs, args, kwargs) + return res + + +@torch.jit.script +def script_rpc_remote_call( + dst_worker_name: str, args: tuple[Tensor, Tensor], kwargs: dict[str, Tensor] +): + rref_res = rpc.remote(dst_worker_name, two_args_two_kwargs, args, kwargs) + return rref_res.to_here() + + +class JitRpcOpTest: + # Call functions remotely from Script. + @dist_init + def test_all_kwargs_are_populated_by_defaults(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = {} + + for script_op in [ + script_rpc_async_call, + script_rpc_sync_call, + script_rpc_remote_call, + ]: + ret = script_op(dst_worker_name, args, kwargs) + self.assertEqual(ret, torch.tensor([10, 10])) + + @dist_init + def test_some_kwargs_are_populated_by_defaults(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = {"first_kwarg": torch.tensor([2, 2])} + + for script_op in [ + script_rpc_async_call, + script_rpc_sync_call, + script_rpc_remote_call, + ]: + ret = script_op(dst_worker_name, args, kwargs) + self.assertEqual(ret, torch.tensor([9, 9])) + + @dist_init + def test_no_kwargs_are_populated_by_defaults(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = { + "first_kwarg": torch.tensor([2, 2]), + "second_kwarg": torch.tensor([3, 3]), + } + for script_op in [ + script_rpc_async_call, + script_rpc_sync_call, + script_rpc_remote_call, + ]: + ret = script_op(dst_worker_name, args, kwargs) + self.assertEqual(ret, torch.tensor([8, 8])) + + @dist_init + def test_args_and_kwargs_contain_different_types(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def script_rpc_async_call_with_assorted_types( + dst_worker_name: str, + ): + args = (torch.tensor([1, 1]), "str_arg", 1) + # Must annotate the value type as `Any`, because JIT type inference + # does not support multiple types when defining a Dict. + # The error JIT gives is, + # "Dict values must contain only a single type, " + # "expected: Tensor but found str instead." + kwargs: dict[str, Any] = { + "tensor_kwarg": torch.tensor([3, 3]), + "str_kwarg": "_str_kwarg", + "int_kwarg": 3, + } + fut = rpc.rpc_async( + dst_worker_name, assorted_types_args_kwargs, args, kwargs + ) + ret = fut.wait() + return ret + + ret = script_rpc_async_call_with_assorted_types(dst_worker_name) + self.assertEqual(ret, (torch.tensor([4, 4]), "str_arg_str_kwarg", 4)) + + @dist_init + def test_kwargs_not_passed(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def script_rpc_async_call_without_kwargs_passed( + dst_worker_name: str, + ): + args = () + fut = rpc.rpc_async(dst_worker_name, no_arg, args) + ret = fut.wait() + return ret + + ret = script_rpc_async_call_without_kwargs_passed(dst_worker_name) + self.assertEqual(ret, 0) + + @dist_init + def test_args_kwargs_are_neither_passed(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def script_rpc_async_call_without_args_kwargs_passed( + dst_worker_name: str, + ): + fut = rpc.rpc_async(dst_worker_name, no_arg) + ret = fut.wait() + return ret + + ret = script_rpc_async_call_without_args_kwargs_passed(dst_worker_name) + self.assertEqual(ret, 0) + + @dist_init + def test_less_than_needed_args_are_specified(self): + if self.rank != 0: + return + + # Notice, args matching happens during scripting. + with self.assertRaisesRegex(RuntimeError, "Argument second_arg not provided"): + + @torch.jit.script + def script_rpc_async_call_with_less_args( + dst_worker_name: str, # noqa: E999 + ): + args = (torch.tensor([1, 1]),) + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + + @dist_init + def test_more_than_needed_args_are_specified(self): + if self.rank != 0: + return + + # Notice, args matching happens during scripting. + with self.assertRaisesRegex( + RuntimeError, + "Expected at most 4 arguments but found 5 positional arguments", + ): + + @torch.jit.script + def script_rpc_async_call_with_more_args( + dst_worker_name: str, + ): + args = ( + torch.tensor([1, 1]), + torch.tensor([2, 2]), + torch.tensor([3, 3]), + torch.tensor([4, 4]), + torch.tensor([5, 5]), + ) + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + + @dist_init + def test_unexepected_kwarg_is_specified(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + # Notice, kwargs matching happens during execution. + @torch.jit.script + def script_rpc_async_call_with_unexpected_kwarg( + dst_worker_name: str, # noqa: E999 + ): + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = {"third_kwarg": torch.tensor([1, 1])} + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + + with self.assertRaisesRegex( + RuntimeError, "Unknown keyword argument 'third_kwarg'" + ): + ret = script_rpc_async_call_with_unexpected_kwarg(dst_worker_name) + self.assertEqual(ret, 0) + + @dist_init + def test_call_python_function_remotely_from_script_not_supported(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def rpc_async_call_remote_py_function_in_torchscript(dst_worker_name: str): + args = () + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, python_function, args, kwargs) + ret = fut.wait() + return ret + + with self.assertRaisesRegex( + RuntimeError, "attempted to get undefined function" + ): + ret = rpc_async_call_remote_py_function_in_torchscript(dst_worker_name) + self.assertEqual(ret, 0) + + @dist_init + def test_call_script_function_that_raises_remotely_from_script(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + # Notice, TorchScript always translates(emits) Python `raise` statement, + # as the exception message string, "Exception", + # no matter what exception type and exception message are in the statement, + @torch.jit.script + def rpc_async_call_remote_raising_torchscript_in_torchscript( + dst_worker_name: str, + ): + args = () + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, raise_script, args, kwargs) + ret = fut.wait() + return ret + + with self.assertRaisesRegex(RuntimeError, "Expected error"): + ret = rpc_async_call_remote_raising_torchscript_in_torchscript( + dst_worker_name + ) + self.assertEqual(ret, 0) + + @dist_init + def test_call_script_function_that_not_exists_remotely_from_script(self): + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + @torch.jit.script + def nonexisting_script(): + return 0 + + @torch.jit.script + def rpc_async_call_remote_nonexisting_torchscript_in_torchscript( + dst_worker_name: str, + ): + args = () + kwargs = {} + fut = rpc.rpc_async(dst_worker_name, nonexisting_script, args, kwargs) + ret = fut.wait() + return ret + + with self.assertRaisesRegex( + RuntimeError, "attempted to get undefined function nonexisting_script" + ): + ret = rpc_async_call_remote_nonexisting_torchscript_in_torchscript( + dst_worker_name + ) + self.assertEqual(ret, 0) + + +@torch.jit.ignore +def my_script_module_init(rank: int) -> MyModuleInterface: + return MyScriptModule(rank) + + +@torch.jit.script +def construct_my_script_module(rank: int) -> MyModuleInterface: + return my_script_module_init(rank) + + +@torch.jit.script +def run_ref_script_module( + ref_script_module: RRef[MyModuleInterface], t: Tensor +) -> Tensor: + module = ref_script_module.to_here() + return module.forward() + t + + +@torch.jit.script +def script_check_rref_confirmed(rref: RRef[Tensor]) -> bool: + return rref.confirmed_by_owner() + + +@torch.jit.script +def save_rref(rref_var: RRef[Tensor], fname: str) -> None: + torch.save(rref_var, fname) + + +@torch.jit.script +def script_add(x: Tensor, y: Tensor) -> Tensor: + return x + y + + +@rpc.functions.async_execution +@torch.jit.script +def async_add(to: str, x: Tensor, y: Tensor) -> Future[Tensor]: + return rpc.rpc_async(to, script_add, (x, y)) + + +@rpc.functions.async_execution +@torch.jit.script +def async_wrong_type() -> Tensor: + return torch.zeros(2) + + +def load_script_module_with_pickled_rref(pickled_script_module): + f = io.BytesIO(pickled_script_module) + m = torch.jit.load(f) + return m() + + +class JitRpcTest( + RRefAPITest, + RRefTypingTest, + LocalRRefTest, + JitRpcOpTest, + FutureTypingTest, + RpcAgentTestFixture, +): + @dist_init + def test_torchscript_function(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + local_ret = one_arg(torch.ones(2, 2)) + ret = rpc.rpc_sync(dst_worker_name, one_arg, args=(torch.ones(2, 2),)) + self.assertEqual(ret, local_ret) + rref = rpc.remote(dst_worker_name, one_arg, args=(torch.ones(2, 2),)) + self.assertEqual(rref.to_here(), local_ret) + # create rref to itself + local_rref = rpc.remote( + worker_name(self.rank), one_arg, args=(torch.ones(2, 2),) + ) + self.assertEqual(local_rref.to_here(), local_ret) + + @dist_init + def test_torchscript_function_exception(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + with self.assertRaisesRegex(RuntimeError, r"one_arg\(\) expected at most"): + rpc.rpc_sync(dst_worker_name, one_arg, args=(10, 20)) + + with self.assertRaisesRegex(RuntimeError, r"one_arg\(\) expected at most"): + rpc.remote(dst_worker_name, one_arg, args=(10, 20)) + + @dist_init + def test_torchscript_functions_not_supported(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + my_local_script_module = MyScriptModule(self.rank) + + # It is not thread safe to instantiate MyScriptModule in multiple threads, + # wait for local MyScriptModule instantiation to finish, + # otherwise it could instantiate MyScriptModule in parallel with + # server thread in the below + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + + # rpc_sync still accepts script class and run it in + # the same code path as python call. + rpc.rpc_sync(dst_worker_name, MyScriptClass, args=(self.rank,)) + + # rpc_sync does not accept script module method. + # Python 3.5 and Python 3.6 throw different error message, the only + # common word can be greped is "pickle". + with self.assertRaisesRegex(TypeError, "pickle"): + rpc.rpc_async(dst_worker_name, my_local_script_module.forward, args=()) + + @dist_init + def test_remote_script_module(self): + # TODO, need more investigation + # there is rref leak when shutting down, suspect it is because + # ref as arg is passed to pybind boundary, and the ref is not garbage + # collected by python when calling shutdown() + import torch.distributed.rpc.api as api + + api._ignore_rref_leak = True + + local_ret = torch.ones(self.rank) + torch.ones(self.rank) + + n = self.rank + 1 + dst_rank = n % self.world_size + remote_ref = rpc.remote( + worker_name(dst_rank), construct_my_script_module, args=(self.rank,) + ) + + # pass rref arg to owner + ret = rpc.rpc_sync( + worker_name(dst_rank), + run_ref_script_module, + args=(remote_ref, torch.ones(self.rank)), + ) + self.assertEqual(ret, local_ret) + + # pass rref arg to self/user + with self.assertRaisesRegex( + RuntimeError, + "is an RRef to a ScriptModule. It can't be sent through RPC from owner,", + ): + ret = rpc.rpc_sync( + worker_name(self.rank), + run_ref_script_module, + args=(remote_ref, torch.ones(self.rank)), + ) + + @dist_init + def test_create_script_module_on_remote(self): + dst_name = worker_name((self.rank + 1) % self.world_size) + # Construct on remote end with rpc_sync + created_script_module = rpc.rpc_sync( + dst_name, MyScriptModule, args=(self.rank,) + ) + # Forward should output a ones tensor of self.rank. + self.assertTrue(isinstance(created_script_module, torch.jit.ScriptModule)) + rank_ones_tensor = created_script_module() + self.assertEqual(torch.ones(self.rank), rank_ones_tensor) + + # Construct ScriptModule with rpc.remote. + remote_script_module = rpc.remote(dst_name, MyScriptModule, args=(self.rank,)) + # Verify it is an instance of ScriptModule on remote end. + remote_end_is_script = rpc.rpc_sync( + remote_script_module.owner(), + rref_isinstance, + args=(remote_script_module, torch.jit.ScriptModule), + ) + self.assertTrue(remote_end_is_script) + # Run forward pass remotely. + remote_forward_output = remote_script_module.rpc_sync().forward() + self.assertEqual(remote_forward_output, torch.ones(self.rank)) + # Run function defined on ScriptModule remotely. + remote_func_output = remote_script_module.rpc_sync().custom_func() + self.assertEqual(remote_func_output, torch.ones(self.rank)) + # Ensure we can transfer ScriptModule RRef to this rank and run + # forward pass. + local_script_module = remote_script_module.to_here() + self.assertTrue(isinstance(local_script_module, torch.jit.ScriptModule)) + rank_ones_tensor = local_script_module() + self.assertEqual(rank_ones_tensor, torch.ones(self.rank)) + local_script_func_output = local_script_module.custom_func() + self.assertEqual(local_script_func_output, torch.ones(self.rank)) + + @dist_init + def test_load_script_module_with_pickled_rref(self): + dst_name = worker_name((self.rank + 1) % self.world_size) + m1 = MyScriptModuleWithRRefs(dst_name) + m2 = MyScriptModuleWithRRefs(dst_name) + + f = io.BytesIO() + + rpc._enable_jit_rref_pickle() + torch.jit.save(m1, f) + rpc._disable_jit_rref_pickle() + + out1 = rpc.rpc_sync( + dst_name, load_script_module_with_pickled_rref, args=(f.getvalue(),) + ) + out2 = m2() + self.assertEqual(out1, out2) + + @dist_init + def test_rref_jit_pickle_not_supported(self): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_var = rpc_return_rref(worker_name(dst_rank)) + with ( + TemporaryFileName() as fname, + self.assertRaisesRegex( + RuntimeError, "RRef jit pickling is only allowed inside RPC calls" + ), + ): + save_rref(rref_var, fname) + + @dist_init + def test_remote_script_throw(self): + rref = rpc.remote( + worker_name((self.rank + 1) % self.world_size), + script_raise_func, + args=(torch.ones(2),), + ) + with self.assertRaisesRegex(Exception, ".*Expected error.*"): + rref.to_here() + + @dist_init + def test_remote_script_udf(self): + rref = rpc.remote( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_udf, + args=(torch.ones(2),), + ) + self.assertEqual(rref.to_here(), torch.ones(2) * 2) + + @dist_init + def test_async_script_udf(self): + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_udf, + args=(torch.ones(2),), + ) + self.assertEqual(future.wait(), torch.ones(2) * 2) + + @dist_init + def test_callback_simple(self): + def callback(fut): + return fut.wait() + 1 + + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_udf, + args=(torch.ones(2),), + ).then(callback) + self.assertEqual(future.wait(), torch.ones(2) * 2 + 1) + + @dist_init + def test_callback_chain(self): + n = self.rank + 1 + + def callback(fut): + return fut.wait() + 1 + + fut = rpc.rpc_async( + worker_name(n % self.world_size), one_arg, args=(torch.ones(n, n),) + ) + + num_cbs = 20 + for _ in range(num_cbs): + fut = fut.then(callback) + + self.assertEqual(fut.wait(), torch.ones(n, n) + 1 + num_cbs) + + @dist_init + def test_add_done_callback(self): + callback_called = None + + def callback(fut): + nonlocal callback_called + callback_called = fut.wait() * 2 + + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_udf, + args=(torch.ones(2),), + ) + + future.add_done_callback(callback) + future_then = future.then(lambda _: True) + + self.assertEqual(future.wait(), torch.ones(2) * 2) + + # We have no guarantee that the add_done_callback fn will execute before the test finishes. + # Adding a 'then' callback that runs afterwards to guarantee we wait for the first callback + future_then.wait() + self.assertEqual(callback_called, torch.ones(2) * 4) + + @dist_init + def test_async_script_throw(self): + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_throw, + args=(torch.ones(2),), + ) + with self.assertRaisesRegex(Exception, ".*Expected error.*"): + future.wait() + + @dist_init + def test_callback_with_exception(self): + def callback(fut): + with self.assertRaisesRegex(Exception, ".*Expected error.*"): + fut.wait() + raise RuntimeError("Another expected error") + + future = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + script_fork_wait_throw, + args=(torch.ones(2),), + ).then(callback) + + with self.assertRaisesRegex(RuntimeError, "Another expected error"): + future.wait() + + @dist_init + def test_call_rpc_with_profiling(self): + # Ensures that we can call torch.ops.profiler._call_end_callbacks_on_jit_fut on a jit + # future from within a script function that calls rpc_async + if self.rank == 0: + with _profile() as prof: + prof_key = _build_rpc_profiling_key( + RPCExecMode.ASYNC, + torch._jit_internal._qualified_name(one_arg), + "worker0", + "worker1", + ) + with torch.autograd.profiler.record_function(prof_key) as rf: + call_rpc_with_profiling(rf.record, "worker1") + # TODO: Can't get a reliable time for this profiling event since + # it's hard to estimate the execution time on the remote end for non-UDFs. + # This can be resolved by https://github.com/pytorch/pytorch/issues/36272. + # After that, this test should be modified to validate the function time. + events = prof.function_events + function_event = get_function_event(events, prof_key) + self.assertTrue( + torch._jit_internal._qualified_name(one_arg) in function_event.name + ) + + @dist_init + def test_rpc_async_jit_profiled(self): + # Tests that rpc_async calls made from within a TorchScript function are + # profiled. + if self.rank == 0: + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = worker_name(dst_rank) + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = {} + with _profile() as prof: + script_rpc_async_call(dst_worker_name, args, kwargs) + + # Ensure rpc_async call is profiled + function_events = prof.function_events + qual_name = torch._jit_internal._qualified_name(two_args_two_kwargs) + rpc_async_jit_event = [ + event + for event in function_events + if qual_name in event.name and event.node_id == self.rank + ] + self.assertEqual(len(rpc_async_jit_event), 1) + rpc_async_jit_event = rpc_async_jit_event[0] + profiled_name = _build_rpc_profiling_key( + RPCExecMode.ASYNC_JIT, + qual_name, + worker_name(self.rank), + dst_worker_name, + ) + self.assertEqual(profiled_name, rpc_async_jit_event.name) + remote_events = [event for event in function_events if event.is_remote] + # All remote events should have taken place on dst_rank + remote_event_node_ids = { + remote_event.node_id for remote_event in remote_events + } + self.assertEqual(remote_event_node_ids, {dst_rank}) + # script_rpc_async_call invokes add operator + # so we should see this as a remote event. + remote_add = next( + remote_event + for remote_event in remote_events + if "aten::add" in remote_event.name + ) + remote_add_profiled_name = f"{profiled_name}#remote_op: aten::add" + self.assertEqual(remote_add.name, remote_add_profiled_name) + + @dist_init + def test_record_function_on_caller_rpc_async(self): + if self.rank == 0: + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = worker_name(dst_rank) + block_scope = "foo" + with _profile() as prof: + # Runs 2 rpc_async calls within JIT under record_function. + record_function_on_caller_rpc_async(dst_worker_name, block_scope) + + # Ensure record_function event is profiled. + function_events = prof.function_events + record_function_scope_event = [ + event for event in function_events if event.name == block_scope + ] + self.assertEqual(1, len(record_function_scope_event)) + record_function_scope_event = record_function_scope_event[0] + # Ensure RPC future is profiled. + expected_key = _build_rpc_profiling_key( + RPCExecMode.ASYNC_JIT, + torch._jit_internal._qualified_name(script_add_ones), + worker_name(self.rank), + dst_worker_name, + ) + jit_rpc_events = [ + event for event in function_events if event.name == expected_key + ] + self.assertEqual(2, len(jit_rpc_events)) + # Validate that the record_function scope time is greater than both + # of the individual RPC async call times. The reason it is not necessarily + # greater than the sum is because the two can execute in parallel. + for jit_rpc_event in jit_rpc_events: + self.assertTrue( + record_function_scope_event.cpu_time_total + > jit_rpc_event.cpu_time_total + ) + + @dist_init + def test_rpc_torchscript_record_function(self): + # tests that torchscript functions can be profiled using with + # record_function(...) over RPC. + REMOTE_OP_STR = "#remote_op: " + if self.rank == 0: + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = worker_name(dst_rank) + block_scope = "foo" + with _profile() as prof: + call_rpc_torchscript_with_record_function(dst_worker_name, block_scope) + + # Need to call below to populate CPU children. + prof.key_averages() + function_events = prof.function_events + expected_key = ( + _build_rpc_profiling_key( + RPCExecMode.ASYNC_JIT, + torch._jit_internal._qualified_name( + script_add_ones_with_record_function + ), + worker_name(self.rank), + dst_worker_name, + ) + + REMOTE_OP_STR + + block_scope + ) + remote_record_function_event = next( + evt for evt in function_events if evt.name == expected_key + ) + self.assertTrue(block_scope in remote_record_function_event.name) + remote_children = remote_record_function_event.cpu_children + self.assertTrue("aten::add" in child.name for child in remote_children) + + def test_record_function_jit_end_callbacks_with_fork(self): + # Ensures that we can call rf._call_end_callbacks_on_future on a jit + # future in python eager mode with torch.jit.fork + sleep_interval = 1 + with _profile() as prof: + with torch.autograd.profiler.record_function("foo") as rf: + fut = torch.jit._fork(sleep, sleep_interval) + rf._call_end_callbacks_on_future(fut) + fut.wait() + + function_events = prof.function_events + sleep_event = get_function_event(function_events, "foo") + self.assertEqual(sleep_event.name, "foo") + # Validate that callbacks were fired at the right time by checking the + # profiling event cpu time + self.assertGreaterAlmostEqual(sleep_event.cpu_time * 1e-6, sleep_interval) + + def test_call_fork_in_jit_with_profiling(self): + # Ensures that we can call torch.ops.profiler._call_end_callbacks_on_jit_fut on a jit + # future from within a script function with torch.jit.fork + with _profile() as prof, torch.autograd.profiler.record_function("foo") as rf: + call_fork_with_profiling(rf.record) + + events = prof.function_events + function_event = get_function_event(events, "foo") + self.assertEqual(function_event.name, "foo") + + @dist_init + def test_async_function_simple(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + ret = rpc.rpc_sync( + dst1, async_add, args=(dst2, torch.ones(2, 2), torch.ones(2, 2)) + ) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + @dist_init + def test_async_function_wrong_return_type(self): + with self.assertRaisesRegex( + RuntimeError, + "Async functions must return an IValue of Future type, but got Tensor", + ): + rpc.rpc_sync( + worker_name((self.rank + 1) % self.world_size), async_wrong_type + ) + + @dist_init + def test_async_function_wrong_decorator_order(self): + # @torch.jit.script complains about undefined value rpc. Error is shown + # below. The reason for not checking error string is to avoid making + # JIT error handling code depend on RPC tests, as we don't have any + # restrictions on the error message here. + # + # RuntimeError: + # undefined value rpc: + # def async_wrong_decorator_order(to, x, y): + # # type: (str, Tensor, Tensor) -> Future[Tensor] + # return rpc.rpc_async(to, script_add, (x, y)) + # ~~~ <--- HERE + with self.assertRaises(RuntimeError): + + @torch.jit.script + @rpc.functions.async_execution + def async_wrong_decorator_order( + to: str, x: Tensor, y: Tensor + ) -> Future[Tensor]: + return rpc.rpc_async(to, script_add, (x, y)) + + @dist_init + def test_async_function_remote(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + rref = rpc.remote( + dst1, async_add, args=(dst2, torch.ones(2, 2), torch.ones(2, 2)) + ) + self.assertEqual(rref.to_here(), torch.ones(2, 2) + 1) + + @dist_init + def test_async_function_remote_multi(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + num = 20 + rrefs = [ + rpc.remote( + dst1, async_add, args=(dst2, torch.ones(2, 2), torch.ones(2, 2) * i) + ) + for i in range(num) + ] + + for i in range(num): + self.assertEqual(rrefs[i].to_here(), torch.ones(2, 2) + i) + + @dist_init + def test_async_function_wrong_return_type_remote(self): + rref = rpc.remote( + worker_name((self.rank + 1) % self.world_size), async_wrong_type + ) + + with self.assertRaisesRegex( + RuntimeError, + "Async functions must return an IValue of Future type, but got Tensor", + ): + rref.to_here() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test_faulty.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test_faulty.py new file mode 100644 index 0000000000000000000000000000000000000000..9bedaad32d0e904a9a7523f31eced9cef96e832d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/jit/rpc_test_faulty.py @@ -0,0 +1,219 @@ +# mypy: allow-untyped-defs + + +import torch +import torch.distributed.rpc as rpc +from torch import Tensor +from torch.distributed.rpc import RRef +from torch.testing._internal.dist_utils import ( + dist_init, + wait_until_pending_futures_and_users_flushed, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +@torch.jit.script +def two_args_two_kwargs( + first_arg, + second_arg, + first_kwarg=torch.tensor([3, 3]), + second_kwarg=torch.tensor([4, 4]), +): + return first_arg + second_arg + first_kwarg + second_kwarg + + +@torch.jit.script +def script_rpc_async_call( + dst_worker_name: str, args: tuple[Tensor, Tensor], kwargs: dict[str, Tensor] +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + ret = fut.wait() + return ret + + +@torch.jit.script +def rpc_async_call_with_timeout( + dst_worker_name: str, + args: tuple[Tensor, Tensor], + kwargs: dict[str, Tensor], + timeout: float, +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs, timeout) + ret = fut.wait() + return ret + + +@torch.jit.script +def rpc_async_call_with_timeout_future_ret( + dst_worker_name: str, + args: tuple[Tensor, Tensor], + kwargs: dict[str, Tensor], + timeout: float, +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs, timeout) + return fut + + +@torch.jit.script +def rpc_async_call_future_ret( + dst_worker_name: str, args: tuple[Tensor, Tensor], kwargs: dict[str, Tensor] +): + fut = rpc.rpc_async(dst_worker_name, two_args_two_kwargs, args, kwargs) + return fut + + +@torch.jit.script +def rref_to_here(rref_var: RRef[Tensor]) -> Tensor: + return rref_var.to_here() + + +@torch.jit.script +def rref_to_here_with_timeout(rref_var: RRef[Tensor], timeout: float) -> Tensor: + return rref_var.to_here(timeout) + + +@torch.jit.script +def rpc_async_with_rref_arg(dst_worker_name: str, args: tuple[RRef[Tensor]]) -> Tensor: + fut = rpc.rpc_async(dst_worker_name, rref_to_here, args) + ret = fut.wait() + return ret + + +class JitFaultyAgentRpcTest(RpcAgentTestFixture): + """ + Run tests for rpc_async in JIT under the faulty agent test fixture to test + arbitrary timeouts. + """ + + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_CALL": 1.5}) + def test_timeout_in_torchscript_function(self): + # Call rpc_async + fut.wait() in torchscript function and ensure that + # timeout is raised. + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = { + "first_kwarg": torch.tensor([2, 2]), + "second_kwarg": torch.tensor([3, 3]), + } + expected_error = self.get_timeout_error_regex() + # Ensure that we get a timeout if we override the default timeout and + # the RPC takes longer to execute. + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc_async_call_with_timeout(dst_worker_name, args, kwargs, 0.5) + + # Ensure that we timeout if we don't specify a timeout but the default + # is less than the RPC takes to execute. + rpc._set_rpc_timeout(0.001) + with self.assertRaisesRegex(RuntimeError, expected_error): + script_rpc_async_call(dst_worker_name, args, kwargs) + + # Ensure that we run to completion if zero timeout is specified. + ret = rpc_async_call_with_timeout(dst_worker_name, args, kwargs, 0) + self.assertEqual(ret, torch.tensor([8, 8])) + # reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_CALL": 1.5}) + def test_timeout_in_python(self): + # Ensures timeouts are raised if we call rpc_async from within a + # torchscript function, but wait on the future in python. + if self.rank != 0: + return + + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + args = (torch.tensor([1, 1]), torch.tensor([2, 2])) + kwargs = { + "first_kwarg": torch.tensor([2, 2]), + "second_kwarg": torch.tensor([3, 3]), + } + expected_error = self.get_timeout_error_regex() + + fut = rpc_async_call_with_timeout_future_ret(dst_worker_name, args, kwargs, 0.5) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure timeout if we don't specify but the default is less than the + # RPC takes to execute. + rpc._set_rpc_timeout(0.001) + fut = rpc_async_call_future_ret(dst_worker_name, args, kwargs) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure run to completion if zero timeout is specified + fut = rpc_async_call_with_timeout_future_ret(dst_worker_name, args, kwargs, 0) + result = fut.wait() + self.assertEqual(result, torch.tensor([8, 8])) + # reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_remote_timeout_to_here_in_jit(self): + # Test that calling to_here() in JIT will raise timeout error if + # rpc.remote failed. + if self.rank != 0: + return + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + # Will ensure error handling callbacks are run. + wait_until_pending_futures_and_users_flushed() + # Call to_here() within a ScriptFunction and ensure it raises + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref_to_here(rref) + + @dist_init(faulty_messages=[], messages_to_delay={"SCRIPT_RREF_FETCH_CALL": 1}) + def test_rref_to_here_timeout_in_jit(self): + if self.rank != 0: + return + + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rref_to_here_with_timeout(rref, 0.01) + + rref_to_here_with_timeout(rref, 100) + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_rref_timeout_pickle_in_jit(self): + if self.rank != 0: + return + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + # Will ensure error handling callbacks are run. + wait_until_pending_futures_and_users_flushed() + # Call RPC with RRef arg in JIT, which will go through JIT pickling and + # ensure error is raised. + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rpc_async_with_rref_arg(dst_worker, (rref,)) + + @dist_init(faulty_messages=["SCRIPT_REMOTE_CALL"]) + def test_rref_timeout_pickle_script_func(self): + # Similar to above test, but calls python rpc with script function. + if self.rank != 0: + return + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + rref = rpc.remote( + dst_worker, torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + # Will ensure error handling callbacks are run. + wait_until_pending_futures_and_users_flushed() + # Call RPC with script function that takes RRef, ensure timeout during pickling + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rpc.rpc_sync(dst_worker, rref_to_here, args=(rref,)) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/rpc_agent_test_fixture.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/rpc_agent_test_fixture.py new file mode 100644 index 0000000000000000000000000000000000000000..3a684b73d2f315a00465371fad3050a795251ddb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/rpc_agent_test_fixture.py @@ -0,0 +1,63 @@ +# mypy: allow-untyped-defs + +import os +from abc import ABC, abstractmethod + +import torch.testing._internal.dist_utils + + +class RpcAgentTestFixture(ABC): + @property + def world_size(self) -> int: + return 4 + + @property + def init_method(self): + use_tcp_init = os.environ.get("RPC_INIT_WITH_TCP", None) + if use_tcp_init == "1": + master_addr = os.environ["MASTER_ADDR"] + master_port = os.environ["MASTER_PORT"] + return f"tcp://{master_addr}:{master_port}" + else: + return self.file_init_method + + @property + def file_init_method(self): + return torch.testing._internal.dist_utils.INIT_METHOD_TEMPLATE.format( + file_name=self.file_name + ) + + @property + @abstractmethod + def rpc_backend(self): + pass + + @property + @abstractmethod + def rpc_backend_options(self): + pass + + def setup_fault_injection(self, faulty_messages, messages_to_delay): # noqa: B027 + """Method used by dist_init to prepare the faulty agent. + + Does nothing for other agents. + """ + + # Shutdown sequence is not well defined, so we may see any of the following + # errors when running tests that simulate errors via a shutdown on the + # remote end. + @abstractmethod + def get_shutdown_error_regex(self): + """ + Return various error message we may see from RPC agents while running + tests that check for failures. This function is used to match against + possible errors to ensure failures were raised properly. + """ + + @abstractmethod + def get_timeout_error_regex(self): + """ + Returns a partial string indicating the error we should receive when an + RPC has timed out. Useful for use with assertRaisesRegex() to ensure we + have the right errors during timeout. + """ diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/rpc_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/rpc_test.py new file mode 100644 index 0000000000000000000000000000000000000000..e3b2452d5eea69401ca2893e4ab87b4830f0c9f6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/rpc_test.py @@ -0,0 +1,6307 @@ +# mypy: allow-untyped-defs + +import concurrent.futures +import contextlib +import json +import operator +import os +import sys +import threading +import time +from collections import namedtuple +from functools import partial +from threading import Event, Lock +from unittest import mock + +import torch +import torch.distributed as dist +import torch.distributed.autograd as dist_autograd +import torch.distributed.rpc as rpc +import torch.nn as nn +from torch.autograd.profiler_legacy import profile as _profile +from torch.distributed.rpc import ( + _get_debug_info, + _rref_context_get_debug_info, + RRef, + WorkerInfo, +) +from torch.distributed.rpc.api import _thread_local_var, _use_rpc_pickler, _wait_all +from torch.distributed.rpc.internal import ( + _build_rpc_profiling_key, + _internal_rpc_pickler, + PythonUDF, + RPCExecMode, +) +from torch.futures import Future +from torch.testing._internal.common_distributed import ( + captured_output, + skip_if_lt_x_gpu, + tp_transports, +) +from torch.testing._internal.common_utils import ( + get_cycles_per_ms, + IS_MACOS, + load_tests, + skip_but_pass_in_sandcastle_if, + TemporaryFileName, +) +from torch.testing._internal.dist_utils import ( + dist_init, + get_function_event, + initialize_pg, + wait_until_node_failure, + wait_until_owners_and_forks_on_rank, + wait_until_pending_futures_and_users_flushed, + worker_name, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +def foo_add(): + return torch.add(torch.ones(1), torch.ones(1)) + + +def udf_with_torch_ops(device=-1, use_record_function=False): + device_ctx = contextlib.nullcontext() if device == -1 else torch.cuda.device(device) + record_function_ctx = ( + torch.autograd.profiler.record_function("##forward##") + if use_record_function + else contextlib.nullcontext() + ) + with device_ctx, record_function_ctx: + t1, t2 = torch.ones(1), torch.ones(1) + t = torch.add(t1, t2) + t = torch.mul(t, t) + t = t.relu() + t = t.sigmoid() + + +# Events (operator invocations) that are expected to be ran as part of the above +# function. +EXPECTED_REMOTE_EVENTS = [ + "aten::ones", + "aten::ones", + "aten::add", + "aten::mul", + "aten::relu", + "aten::clamp_min", + "aten::sigmoid", +] + +# Remote operations are prefixed with the following string for RPC profiling. +REMOTE_OP_STR = "#remote_op: " + + +VALUE_FUTURE = concurrent.futures.Future() +DONE_FUTURE = concurrent.futures.Future() + +FIFTY_MIL_CYCLES = 50000000 + +_rpc_barrier_count = 0 + + +def _increment_count(): + global _rpc_barrier_count + _rpc_barrier_count += 1 + + +def _reset_count(): + global _rpc_barrier_count + _rpc_barrier_count = 0 + + +class StubRpcAgent: + def __init__(self, world_size): + self.world_size = world_size + + def get_worker_infos(self): + return { + WorkerInfo(name=worker_name(rank), id=rank) + for rank in range(self.world_size) + } + + +def _stub_construct_rpc_backend_options_handler(**kwargs): + return mock.Mock() # RpcBackendOptions. + + +def _stub_init_rpc_backend_handler(store, name, rank, world_size, rpc_backend_options): + return StubRpcAgent(world_size=world_size) + + +def set_value(value): + VALUE_FUTURE.set_result(value) + + +def wait_for_value_future(): + return VALUE_FUTURE.result() + + +def set_and_check_done(value): + VALUE_FUTURE.set_result(value) + return DONE_FUTURE.result() + + +# it is used to test python user defined function over rpc +# classes and functions are used to test python user defined class and +# methods over rpc +TensorClass = namedtuple("TensorClass", ["tensors"]) + + +class MyPickleClass: + def __init__(self) -> None: + self.t = None + + def __getstate__(self): + (pickled_python_udf, tensors) = _internal_rpc_pickler.serialize( + PythonUDF(my_tensor_function, (torch.ones(2, 2), torch.ones(2, 2)), None) + ) + return (pickled_python_udf, tensors) + + def __setstate__(self, obj): + python_udf = _internal_rpc_pickler.deserialize(obj[0], obj[1]) + result = python_udf.func(python_udf.args[0], python_udf.args[1]) + self.t = result + + def set(self, val): + self.t = val + + +class SlowPickleClass: + def __init__(self, t): + self.t = t + + def __getstate__(self): + time.sleep(self.t) + return (self.t,) + + def __setstate__(self, obj): + self.t = obj[0] + time.sleep(self.t) + + +class MyClass: + def __init__(self, a, delay=False): + self.a = a + # delay initialization to simulate errors if specified + if delay: + time.sleep(2) + + def my_instance_method(self, b): + return self.a + b + + @classmethod + def my_class_method(cls, d, e): + return d + e + + @staticmethod + def my_static_method(f): + return f > 10 + + def increment_value(self, increment): + self.a += increment + + def get_value(self): + return self.a + + def my_slow_method(self, my_tensor_arg): + time.sleep(5) + return torch.add(self.a, my_tensor_arg) + + +def _call_method_on_rref(method, rref, *args, **kwargs): + return method(rref.local_value(), *args, **kwargs) + + +def get_rref_list(values): + return [RRef(MyClass(a)) for a in values] + + +def add_rref_to_value(rref, value): + return rref.to_here() + value + + +def run_nested_pickle(pickle_cls_instance, tensor): + return pickle_cls_instance.t + tensor + + +def build_sparse_tensor(coalesce=False): + i = [[0, 1, 1], [2, 0, 2]] + v = [3, 4, 5] + tensor = torch.sparse_coo_tensor(i, v, (2, 3)) + if coalesce: + tensor = tensor.coalesce() + return tensor + + +def build_complex_tensors(): + a = torch.ones(3, 3) + b = [a, a] + c = [b, b] + d = [a, b] + e = {a: d} + return [a, b, c, d, e] + + +def non_cont_test(t_view, t_cont): + if t_view.is_contiguous(): + raise Exception("t_view is contiguous!") # noqa: TRY002 + if not t_cont.is_contiguous(): + raise Exception("t_cont is not contiguous!") # noqa: TRY002 + if not torch.equal(t_view, t_cont): + raise Exception("t_view is not equal to t_cont!") # noqa: TRY002 + return t_view + + +def my_function(a, b, c): + return a + b + c + + +def my_tensor_function(a, b): + return a + b + + +def my_container_sum(a): + result = a[0] + for tensor in a[1:]: + result += tensor + return result + + +def my_sleep_func(seconds=1): + time.sleep(seconds) + return torch.mul(torch.tensor(1), torch.tensor(1)) + + +def my_complex_tensor_function(list_input, tensor_class_input, dict_input): + res = list_input[0] + for t in list_input: + res += t + for v in dict_input.values(): + res += v + complex_tensors = tensor_class_input.tensors + return (res, complex_tensors[0], complex_tensors[1], complex_tensors[2]) + + +def my_rref_function(rref_a, rref_b): + return rref_a.to_here() + rref_b.to_here() + + +def delayed_add(a, b, seconds=0.05): + time.sleep(seconds) + return a + b + + +def identity(a): + return a + + +def no_result(): + print("do nothing") + + +def raise_or_inc(value): + if value.numel() == 2: + raise ValueError("Expected error") + return value + 1 + + +def nested_rpc(dst): + return rpc.rpc_sync(dst, torch.add, args=(torch.ones(2, 2), 1)) + + +def nested_rpc_sparse(dst): + return rpc.rpc_sync( + dst, torch.add, args=(build_sparse_tensor(), build_sparse_tensor()) + ) + + +def multi_layer_nested_async_rpc(dst, world_size, ttl): + # this method returns immediately without blocking the callee, but will + # generate additional requests. + if ttl > 0: + current_dst = worker_name(dst) + next_dst = (dst + 1) % world_size + rpc.rpc_async( + current_dst, + multi_layer_nested_async_rpc, + args=(next_dst, world_size, ttl - 1), + ) + return 0 + + +def nested_rref(dst): + return ( + rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 1)), + rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 2)), + ) + + +def nested_rref_sparse(dst): + return ( + rpc.remote(dst, torch.add, args=(build_sparse_tensor(), build_sparse_tensor())), + rpc.remote(dst, torch.add, args=(build_sparse_tensor(), build_sparse_tensor())), + ) + + +def nested_remote(dst): + rref = rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 3)) + return rref.to_here() + + +def nested_remote_sparse(dst): + rref = rpc.remote( + dst, torch.add, args=(build_sparse_tensor(), build_sparse_tensor()) + ) + return rref.to_here() + + +def rref_forward_chain(dst, world_size, rref, ttl): + if ttl > 0: + current_dst = worker_name(dst) + next_dst = (dst + 1) % world_size + ret_rref = rpc.remote( + current_dst, rref_forward_chain, args=(next_dst, world_size, rref, ttl - 1) + ) + return [ret_rref] + else: + return rref.to_here() + + +def rpc_return_rref(dst): + return rpc.remote(dst, torch.add, args=(torch.ones(2, 2), 1)) + + +def light_rpc(): + return 0 + + +def heavy_rpc(tensor): + for i in range(1, 100): + tensor *= i + tensor /= i + 1 + return 0 + + +def heavy_rpc_sparse(tensor): + for i in range(1, 100): + tensor *= i + tensor = tensor / (i + 1) + return 0 + + +@torch.jit.script +def heavy_rpc_torchscript(tensor): + for i in range(1, 100): + tensor *= i + tensor /= i + 1 + return 0 + + +@torch.jit.script +def my_script_func(tensor): + return torch.add(tensor, tensor) + + +expected_err = "Expected error" + + +# Note that it needs to inherit from Exception, not BaseException. See comment +# in rpc/internal.py +class CustomException(Exception): + def __init__(self, bool, msg): + self.bool = bool + super().__init__(msg) + + +def raise_func(): + raise ValueError(expected_err) + + +def custom_raise_func(): + raise CustomException(True, "foo") + + +@torch.jit.script +def raise_func_script(expected_err: str) -> torch.Tensor: + raise ValueError(expected_err) + + +expected_err_escape = ( + "\nFirst line of error \n next line of error \n last line of error" +) + + +def raise_func_escape(): + raise ValueError(expected_err_escape) + + +global_rref = None + + +def set_global_rref(rref): + global global_rref + global_rref = rref + + +def clear_global_rref(): + global global_rref + global_rref = None + + +def check_rref_confirmed(rref): + return rref.confirmed_by_owner() + + +def get_rref_debug_info(): + return _rref_context_get_debug_info() + + +def add_use_future_cb(to, x, y, z): + out = concurrent.futures.Future() + + def callback(fut): + out.set_result(fut.wait() + z) + + fut = rpc.rpc_async(to, torch.add, args=(x, y)) + fut.then(callback) + return out.result() + + +def get_events_from_profile(profile_rref): + return profile_rref.local_value().process_global_function_events + + +def add_use_future_set_result(to, x, y, z): + out = torch.futures.Future() + fut = rpc.rpc_async(to, torch.add, args=(x, y)) + fut.then(lambda fut: out.set_result(fut.wait() + z)) + return out.wait() + + +def add_use_future_nested_cb(to, x, y, z): + out = torch.futures.Future() + + def callback(fut1): + fut2 = rpc.rpc_async(to, torch.add, args=(fut1.wait(), z)) + fut2.then(lambda fut2: out.set_result(fut2.wait())) + + fut1 = rpc.rpc_async(to, torch.add, args=(x, y)) + fut1.then(callback) + return out.wait() + + +def fail_on_fut(fut): + pass + + +@rpc.functions.async_execution +def async_raise_func(): + raise RuntimeError("Expected error") + + +@rpc.functions.async_execution +def async_wrong_type(): + return torch.zeros(2, 2) + + +@rpc.functions.async_execution +def async_add(to, x, y): + return rpc.rpc_async(to, torch.add, args=(x, y)) + + +def slow_add(x, y, device="cpu"): + time.sleep(1) + x = x.to(device) + y = y.to(device) + return torch.add(x, y).cpu() + + +@rpc.functions.async_execution +def slow_async_add(to, x, y, device="cpu"): + return rpc.rpc_async(to, slow_add, args=(x, y, device)) + + +@rpc.functions.async_execution +def async_add_with_future_ctor(to, x, y, z): + fut = torch.futures.Future() + rpc.rpc_async(to, torch.add, args=(x, y)).then( + lambda fut1: fut.set_result(fut1.wait() + z) + ) + return fut + + +@rpc.functions.async_execution +def async_add_chained(to, x, y, z): + return rpc.rpc_async(to, torch.add, args=(x, y)).then(lambda fut: fut.wait() + z) + + +@rpc.functions.async_execution +def async_add_chained_multi(to, x, num, step): + fut = rpc.rpc_async(to, torch.add, args=(x, 0)) + for _ in range(num): + fut = fut.then(lambda fut: fut.wait() + step) + return fut + + +@rpc.functions.async_execution +def async_add_nested(to, x, y, z): + return rpc.rpc_async(to, async_add, args=(to, x, y)).then( + lambda fut: fut.wait() + z + ) + + +@rpc.functions.async_execution +def async_add_multi_fanout(to, x, num, step): + futs = [] + for i in range(num): + if i == 0: + futs.append(rpc.rpc_async(to, torch.add, args=(x, step))) + else: + futs.append(rpc.rpc_async(to, torch.add, args=(0, step))) + + # TODO: use torch.futures.collect_all + lock = Lock() + state = {"cnt": 0, "ret": torch.zeros_like(x)} + ret_future = torch.futures.Future() + + def inc_and_set(fut): + with lock: + state["cnt"] += 1 + state["ret"] += fut.wait() + if state["cnt"] >= len(futs): + ret_future.set_result(state["ret"]) + + for fut in futs: + fut.then(inc_and_set) + + return ret_future + + +@rpc.functions.async_execution +def async_cuda_sleep_and_set_to_one(t): + device = t.device + original_stream = torch.cuda.current_stream(device) + new_stream = torch.cuda.Stream(device) + new_stream.wait_stream(original_stream) + with torch.cuda.stream(new_stream): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + t.fill_(1) + fut = Future(devices=[device]) + fut.set_result(t) + return fut + + +@rpc.functions.async_execution +def async_cuda_nested_add(to, x, y, z): + def cb(fut): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + return fut.value() + z + + return rpc.rpc_async(to, torch.add, args=(x, y)).then(cb) + + +# A custom Python class that contains a tensor, needed to see if we correctly +# use the Python pickler to extract tensors from non-IValue-convertible types. +class TensorWrapper: + __slots__ = ("tensor", "lock", "event", "thread") + + def __init__(self, t): + self.tensor = t + # Add one non-picklable field, to ensure it's ignored/skipped. + self.lock = Lock() + self.event = torch.cuda.Event(enable_timing=True) + self.thread = threading.Thread() + self.thread.start() + + def increase(self, v): + with self.lock: + self.tensor += v + + def sum(self): + with self.lock: + self.event.record() + return self.tensor.sum() + + +class AsyncExecutionClass: + @staticmethod + @rpc.functions.async_execution + def static_async_add(to, x, y, z): + return rpc.rpc_async(to, torch.add, args=(x, y)).then( + lambda fut: fut.wait() + z + ) + + @classmethod + @rpc.functions.async_execution + def class_async_add(cls, to, x, y, z): + ret_fut = torch.futures.Future() + rpc.rpc_async(to, torch.add, args=(x, y)).then( + lambda fut: ret_fut.set_result(fut.wait() + z) + ) + return ret_fut + + @rpc.functions.async_execution + def bound_async_add(self, to, x, y, z): + return rpc.rpc_async(to, torch.add, args=(x, y)).then( + lambda fut: fut.wait() + z + ) + + +def return_future(): + return torch.futures.Future() + + +class FooBackendOptions(rpc.RpcBackendOptions): + def __init__(self, init_method): + # Must call the __init__ of the superclass (and do so directly, + # without using super()) because... pybind. + rpc.RpcBackendOptions.__init__(self) + self.init_method = init_method + + +# load_tests from common_utils is used to automatically filter tests for +# sharding on sandcastle. This line silences flake warnings +load_tests = load_tests # noqa: PLW0127 + + +class MyEmbeddingBagModel(torch.nn.Module): + def __init__(self, sparse): + super().__init__() + self.eb = torch.nn.EmbeddingBag(10, 10, sparse=sparse) + + def forward(self, x): + return self.eb(x) + + +class MyParameterServer: + def __init__(self, trainers): + self.lock = Lock() + self.trainers = trainers + self.iteration = 0 + self.updates = 0 + self.futures = [] + self.total = None + self.gradient = None + + @staticmethod + def get_gradient(rref): + return rref.local_value().gradient + + @staticmethod + @rpc.functions.async_execution + def average(rref, riteration, tensor): + self = rref.local_value() + fut = torch.futures.Future() + with self.lock: + if riteration > self.iteration: + self.iteration = riteration + self.updates = 0 + self.futures.clear() + self.futures.append(fut) + if self.total is None: + self.total = tensor + else: + self.total += tensor + self.updates += 1 + if self.trainers == self.updates: + self.gradient = self.total / float(self.trainers) + for fut in self.futures: + result = self.total / float(self.trainers) + fut.set_result(result) + return fut + + +class MyConvNetForMNIST(nn.Module): + def __init__(self, device): + super().__init__() + self.net = nn.Sequential( + nn.Conv2d(1, 16, 3, 1), + nn.ReLU(), + nn.Conv2d(16, 32, 3, 1), + nn.ReLU(), + nn.MaxPool2d(2), + nn.Flatten(1), + nn.Linear(4608, 128), + nn.ReLU(), + nn.Linear(128, 10), + ).to(device) + self.device = device + + def forward(self, x, is_rref=False): + x = x.to_here() if is_rref else x + with torch.cuda.stream(torch.cuda.current_stream(self.device)): + # intentionally adding delay to current CUDA stream + torch.cuda._sleep(10 * FIFTY_MIL_CYCLES) + return self.net(x) + + def __getstate__(self): + # return an empty dict to avoid inspecting the model contents on the + # owner + return {} + + +class RpcTestCommon: + def _run_func_in_mode(self, to, fn, mode, args=None, kwargs=None): + if mode == RPCExecMode.SYNC: + return rpc.rpc_sync(to, fn, args=args, kwargs=kwargs) + elif mode == RPCExecMode.ASYNC: + return rpc.rpc_async(to, fn, args=args, kwargs=kwargs).wait() + elif mode == RPCExecMode.REMOTE: + return rpc.remote(to, fn, args=args, kwargs=kwargs).to_here() + + def _self_py_udf_remote(self, worker_info, x, y, z): + rref = rpc.remote(worker_info, my_function, args=(x, y, z)) + self.assertEqual(rref.to_here(), x + y + z) + + def _self_remote_rref_as_rpc_arg(self, dst, x, y, z): + self_worker_info = rpc.get_worker_info() + rref = rpc.remote(self_worker_info, my_function, args=(x, y, z)) + fut = rpc.rpc_async(dst, add_rref_to_value, args=(rref, x)) + ret = rpc.rpc_sync(dst, add_rref_to_value, args=(rref, x + y)) + self.assertEqual(ret, x + y + z + x + y) + self.assertEqual(fut.wait(), x + y + z + x) + + def _self_remote_rref_as_remote_arg(self, dst, x, y, z): + self_worker_info = rpc.get_worker_info() + rref = rpc.remote(self_worker_info, my_function, args=(x, y, z)) + ret_rref = rpc.remote(dst, add_rref_to_value, args=(rref, x)) + self.assertEqual(ret_rref.to_here(), x + y + z + x) + + def _world_size_one(self, a, b): + if self.rank == 0: + rpc.init_rpc( + name="me", + backend=self.rpc_backend, + rank=0, + world_size=1, + rpc_backend_options=self.rpc_backend_options, + ) + + def _rpc_sync(x, y): + expect = x * 2 + result = rpc.rpc_sync("me", my_tensor_function, args=(x, y)) + self.assertEqual(expect, result) + + def _rpc_async(x, y): + expect = x * 2 + result = rpc.rpc_async("me", my_tensor_function, args=(x, y)).wait() + self.assertEqual(expect, result) + + def _remote(x, y): + expect = x * 2 + result = rpc.remote("me", my_tensor_function, args=(x, y)).to_here() + self.assertEqual(expect, result) + + _rpc_sync(a, b) + _rpc_async(a, b) + _remote(a, b) + + rpc.shutdown() + + def _multi_rpc(self, sparse): + dst_rank = (self.rank + 1) % self.world_size + for i in range(20): + n = i + self.rank + 1 + if sparse: + x = build_sparse_tensor() * n + y = build_sparse_tensor() * n + else: + x = torch.ones(2, 2) + y = torch.ones(2, 2) + ret = rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(x, y), + ) + self.assertEqual(ret, x * 2) + + def _run_uneven_workload(self, f, x, num_repeat=30): + # worker0 drives and waits for worker1 and worker2 + # throughout the test. + if self.rank == 0: + self.assertTrue(self.world_size >= 3) + + # Phase 1: Only worker1 has workload. + dst = "worker1" + futs = [] + for _ in range(num_repeat): + fut = rpc.rpc_async(dst, f, args=(x,)) + futs.append(fut) + + for fut in torch.futures.collect_all(futs).wait(): + self.assertEqual(fut.wait(), 0) + + # Phase 2: Only worker2 has workload. + # If join is not correctly implemented, + # worker2 should be closed by now. + dst = "worker2" + futs = [] + for _ in range(num_repeat): + fut = rpc.rpc_async(dst, f, args=(x,)) + futs.append(fut) + + for val in torch.futures.wait_all(futs): + self.assertEqual(val, 0) + + def _wait_all_workers(self, f, x): + initialize_pg(self.file_init_method, self.rank, self.world_size) + rpc.init_rpc( + name=f"worker{self.rank:d}", + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + self._run_uneven_workload(f, x) + + # worker0 calls this at the end after waiting for RPC responses. + # worker1/2 calls this immediately and has some works after it. + # worker3 calls this immediately and has no more work. + rpc.api._wait_all_workers() + + # Wait before proceeding to shutdown to ensure worker0 RPCs make + # it through to other workers. + dist.barrier() + rpc.shutdown(graceful=False) + + def _wait_all_workers_twice(self, f, x): + initialize_pg(self.file_init_method, self.rank, self.world_size) + rpc.init_rpc( + name=f"worker{self.rank:d}", + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + self._run_uneven_workload(f, x) + + # worker0 calls this at the end after waiting for RPC responses. + # worker1/2 calls this immediately and has some works after it. + # worker3 calls this immediately and has no more work. + rpc.api._wait_all_workers() + rpc.api._wait_all_workers() + + # Wait before proceeding to shutdown to ensure worker0 RPCs make + # it through to other workers. + dist.barrier() + rpc.shutdown(graceful=False) + + def _nested_rpc(self, f, expected): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + f, + args=(worker_name(self.rank),), + ) + self.assertEqual(ret, expected) + + def _stress_test_rpc(self, f, repeat=1000, args=()): + n = self.rank + 1 + dst_rank = n % self.world_size + futs = [] + tik = time.time() + for _ in range(repeat): + fut = rpc.rpc_async(worker_name(dst_rank), f, args=args) + futs.append(fut) + + for val in torch.futures.wait_all(futs): + self.assertEqual(val, 0) + tok = time.time() + print( + f"Rank {self.rank} finished testing {repeat} times in {tok - tik} seconds." + ) + + def _builtin_remote_ret(self, x, y, expected): + n = self.rank + 1 + dst_rank = n % self.world_size + rref = rpc.remote( + worker_name(dst_rank), + torch.add, + args=(x, y), + ) + self.assertEqual(rref.to_here(), expected) + + def _builtin_remote_self(self, x, y, expected): + rref = rpc.remote( + worker_name(self.rank), + torch.add, + args=(x, y), + ) + self.assertEqual(rref.local_value(), expected) + + def _test_multi_remote_call( + self, fn, sparse, args_fn=lambda x, y: (), kwargs_fn=lambda x, y: {} + ): + m = 10 + n = self.rank + 1 + dst_rank = n % self.world_size + rrefs = [] + expected = [] + for i in range(m): + n = n + i + rrefs.append( + rpc.remote( + worker_name(dst_rank), + fn, + args=args_fn(n, sparse), + kwargs=kwargs_fn(n, sparse), + ) + ) + expected.append(fn(*args_fn(n, sparse), **kwargs_fn(n, sparse))) + + for i in range(m): + self.assertEqual(rrefs[i].to_here(), expected[i]) + + def _py_rref_args(self, a, b, x, y, expected): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_a = rpc.remote(worker_name(dst_rank), torch.add, args=(a, b)) + rref_b = rpc.remote(worker_name(dst_rank), torch.add, args=(x, y)) + rref_c = rpc.remote( + worker_name(dst_rank), my_rref_function, args=(rref_a, rref_b) + ) + self.assertEqual(rref_c.to_here(), expected) + + def _py_rref_args_user_share(self, a, b, c, x, y, z, expected): + n = self.rank + 1 + owner_rank = n % self.world_size + user_rank = (n + 1) % self.world_size + rref_a = rpc.remote(worker_name(owner_rank), my_function, args=(a, b, c)) + rref_b = rpc.remote(worker_name(owner_rank), my_function, args=(x, y, z)) + rref_c = rpc.remote( + worker_name(user_rank), my_rref_function, args=(rref_a, rref_b) + ) + self.assertEqual(rref_c.to_here(), expected) + + def _py_rpc_rref_args(self, a, b, c, x, y, z, expected): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_a = rpc.remote(worker_name(dst_rank), my_function, args=(a, b, c)) + rref_b = rpc.remote(worker_name(dst_rank), my_function, args=(x, y, z)) + + c = rpc.rpc_sync(worker_name(dst_rank), my_rref_function, args=(rref_a, rref_b)) + self.assertEqual(c, expected) + + def _nested_remote(self, f, expected): + n = self.rank + 1 + dst_rank1 = n % self.world_size + dst_rank2 = (n + 1) % self.world_size + + rref = rpc.remote( + worker_name(dst_rank1), + f, + args=(worker_name(dst_rank2),), + ) + self.assertEqual(rref.to_here(), expected) + + def _nested_rref(self, f, expected1, expected2): + n = self.rank + 1 + dst_rank1 = n % self.world_size + dst_rank2 = (n + 1) % self.world_size + rref_of_rrefs = rpc.remote( + worker_name(dst_rank1), + f, + args=(worker_name(dst_rank2),), + ) + + # Say C has 2 OwnerRRefs. + # B has 2 UserRRefs to those 2 OwnerRRefs, respectively. + # This call is effectively A asking B to share its 2 UserRRefs. + rrefs = rref_of_rrefs.to_here() + + self.assertEqual(len(rrefs), 2) + self.assertEqual(rrefs[0].to_here(), expected1) + self.assertEqual(rrefs[1].to_here(), expected2) + + def _nested_rref_stress(self, f, expected1, expected2): + n = self.rank + 1 + dst_rank1 = n % self.world_size + dst_rank2 = (n + 1) % self.world_size + all_rrefs = [ + rpc.remote( + worker_name(dst_rank1), + f, + args=(worker_name(dst_rank2),), + ) + for _ in range(20) + ] + + for i in range(20): + rref_of_rrefs = all_rrefs[i] + rrefs = rref_of_rrefs.to_here() + self.assertEqual(len(rrefs), 2) + self.assertEqual(rrefs[0].to_here(), expected1) + self.assertEqual(rrefs[1].to_here(), expected2) + + def _trainer_func(self, rref, sparse): + m = MyEmbeddingBagModel(sparse=sparse) + loss_fn = nn.MSELoss() + for i in range(10): + outputs = m(torch.rand(10, 10).long()) + loss_fn(outputs, torch.rand(10, 10)).backward() + gradient = next(iter(m.parameters())).grad + fut = rref.rpc_async().average(rref, i, gradient) + gradient = fut.wait() + if gradient.is_sparse: + gradient = gradient.to_dense().double() + ps_gradient = rref.rpc_sync().get_gradient(rref) + if ps_gradient.is_sparse: + ps_gradient = ps_gradient.to_dense().double() + self.assertTrue(torch.equal(gradient, ps_gradient)) + + def _my_parameter_server(self, sparse): + ps_rref = RRef(MyParameterServer(self.world_size - 1)) + futures = [ + rpc.rpc_async( + worker_name((self.rank + index) % self.world_size), + self._trainer_func, + args=(ps_rref, sparse), + ) + for index in range(1, self.world_size) + ] + torch.futures.wait_all(futures) + + def _test_cuda_future_extraction(self, wrapper, unwrapper, sparse_tensor): + # We check proper CUDA stream synchronization by adding to the tensor + # in one stream to get the expected value, and reading it from another stream. + future = Future(devices=["cuda:0"]) + with torch.cuda.device("cuda:0"): + stream = torch.cuda.Stream() + another_stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + if sparse_tensor: + tensor = build_sparse_tensor().to("cuda:0") + add_tensor = build_sparse_tensor().to("cuda:0") + expected_tensor = (tensor + add_tensor).coalesce() + else: + tensor = torch.zeros((100,), device="cuda:0") + add_tensor = torch.ones((100,), device="cuda:0") + expected_tensor = tensor + add_tensor + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + tensor += add_tensor + if sparse_tensor: + tensor = tensor.coalesce() + future.set_result(wrapper(tensor)) + with torch.cuda.stream(another_stream): + tensor = unwrapper(future.wait()) + if sparse_tensor: + self.assertTrue( + torch.eq(tensor.indices(), expected_tensor.indices()) + .all() + .item() + ) + self.assertTrue( + torch.eq(tensor.values(), expected_tensor.values()).all().item() + ) + self.assertEqual(tensor.size(), expected_tensor.size()) + else: + self.assertTrue(torch.eq(tensor, expected_tensor).all().item()) + + +class RpcTest(RpcAgentTestFixture, RpcTestCommon): + @dist_init + def test_worker_id(self): + n = self.rank + 1 + peer_rank = n % self.world_size + self_worker_info = rpc.get_worker_info() + peer_worker_info = rpc.get_worker_info(worker_name(peer_rank)) + + self.assertEqual(self_worker_info.name, worker_name(self.rank)) + self.assertEqual(peer_worker_info.name, worker_name(peer_rank)) + + with self.assertRaisesRegex(RuntimeError, "could not find destination"): + rpc.get_worker_info("WorkerUnknown") + + @dist_init + def test_get_worker_infos(self): + worker_infos = rpc.api._get_current_rpc_agent().get_worker_infos() + + worker_names = {worker_info.name for worker_info in worker_infos} + expected_worker_names = {worker_name(rank) for rank in range(self.world_size)} + self.assertEqual(worker_names, expected_worker_names) + + worker_ids = {worker_info.id for worker_info in worker_infos} + expected_worker_ids = set(range(self.world_size)) + self.assertEqual(worker_ids, expected_worker_ids) + + @dist_init + def test_self_add(self): + self_worker_info = rpc.get_worker_info() + fut = rpc.rpc_async(self_worker_info, torch.add, args=(torch.ones(2, 2), 1)) + ret = rpc.rpc_sync(self_worker_info, torch.add, args=(torch.ones(2, 2), 1)) + self.assertEqual(fut.wait(), torch.ones(2, 2) + 1) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + @dist_init + def test_send_to_rank(self): + dst_rank = (self.rank + 1) % self.world_size + + # Test dense tensor + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + ret = self._run_func_in_mode( + dst_rank, torch.add, exec_mode, args=(torch.ones(2, 2), 1) + ) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + # Test invalid ranks + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + with self.assertRaises(RuntimeError): + self._run_func_in_mode( + self.world_size + 1, + torch.add, + exec_mode, + args=(torch.ones(2, 2), 1), + ) + + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + with self.assertRaises(RuntimeError): + self._run_func_in_mode( + -1, torch.add, exec_mode, args=(torch.ones(2, 2), 1) + ) + + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + with self.assertRaises(ValueError): + self._run_func_in_mode( + dst_rank + 0.5, torch.add, exec_mode, args=(torch.ones(2, 2), 1) + ) + + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + with self.assertRaises(ValueError): + self._run_func_in_mode( + dst_rank - 0.5, torch.add, exec_mode, args=(torch.ones(2, 2), 1) + ) + + @dist_init + def test_self_py_udf_remote(self): + self._self_py_udf_remote(rpc.get_worker_info(), torch.ones(2, 2), 1, 3) + + @dist_init + def test_self_remote_rref_as_rpc_arg(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._self_remote_rref_as_rpc_arg(dst, torch.ones(2, 2), 1, 3) + + @dist_init + def test_self_remote_rref_as_self_rpc_arg(self): + self._self_remote_rref_as_rpc_arg(rpc.get_worker_info(), torch.ones(2, 2), 1, 3) + + @dist_init + def test_self_remote_rref_as_remote_arg(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._self_remote_rref_as_remote_arg(dst, torch.ones(2, 2), 1, 3) + + @dist_init + def test_self_remote_rref_as_self_remote_arg(self): + self._self_remote_rref_as_remote_arg( + rpc.get_worker_info(), torch.ones(2, 2), 1, 3 + ) + + @dist_init + def test_rref_proxy_non_exist(self): + dst = worker_name((self.rank + 1) % self.world_size) + rref = rpc.remote(dst, my_function, args=(torch.ones(2, 2), 1, 3)) + msg = "has no attribute 'non_exist'" + with self.assertRaisesRegex(AttributeError, msg): + rref.rpc_sync().non_exist() + + with self.assertRaisesRegex(AttributeError, msg): + rref.rpc_async().non_exist().wait() + + with self.assertRaisesRegex(AttributeError, msg): + rref.remote().non_exist() + + def _test_rref_proxy_tensor(self, dst): + rref = rpc.remote(dst, my_function, args=(torch.ones(2, 2), 1, 3)) + + expected = torch.ones(2, 2) + 1 + 3 + self.assertEqual(expected.size(), rref.rpc_sync().size()) + self.assertEqual(expected + 1, rref.rpc_async().add(1).wait()) + self.assertEqual(expected.view(1, 4), rref.remote().view(1, 4).to_here()) + + @dist_init + def test_rref_proxy_tensor(self): + self._test_rref_proxy_tensor(worker_name((self.rank + 1) % self.world_size)) + + @dist_init + def test_rref_proxy_tensor_self(self): + self._test_rref_proxy_tensor(rpc.get_worker_info()) + + @dist_init + def test_rref_proxy_reuse(self): + rref = rpc.remote( + worker_name((self.rank + 1) % self.world_size), + my_function, + args=(torch.ones(2, 2), 1, 3), + ) + expected = torch.ones(2, 2) + 1 + 3 + + proxy_rpc_sync = rref.rpc_sync() + proxy_rpc_async = rref.rpc_async() + proxy_remote = rref.remote() + + self.assertEqual(expected.size(), proxy_rpc_sync.size()) + self.assertEqual(expected + 1, proxy_rpc_sync.add(1)) + self.assertEqual(expected.view(1, 4), proxy_rpc_sync.view(1, 4)) + + self.assertEqual(expected.size(), proxy_rpc_async.size().wait()) + self.assertEqual(expected + 3, proxy_rpc_async.add(3).wait()) + self.assertEqual(expected.view(4, 1), proxy_rpc_async.view(4, 1).wait()) + + self.assertEqual(expected.size(), proxy_remote.size().to_here()) + self.assertEqual(expected + 5, proxy_remote.add(5).to_here()) + self.assertEqual(expected.view(-1), proxy_remote.view(-1).to_here()) + + def _test_rref_proxy_class(self, dst): + rref = rpc.remote(dst, MyClass, args=(7,)) + expected = MyClass(7) + self.assertEqual(expected.get_value(), rref.rpc_sync().get_value()) + self.assertEqual(expected.get_value(), rref.rpc_async().get_value().wait()) + self.assertEqual(expected.get_value(), rref.remote().get_value().to_here()) + + expected.increment_value(3) + self.assertEqual(None, rref.rpc_sync().increment_value(1)) + self.assertEqual(None, rref.rpc_async().increment_value(1).wait()) + self.assertEqual(None, rref.remote().increment_value(1).to_here()) + + self.assertEqual(expected.get_value(), rref.rpc_sync().get_value()) + self.assertEqual(expected.get_value(), rref.rpc_async().get_value().wait()) + self.assertEqual(expected.get_value(), rref.remote().get_value().to_here()) + + self.assertEqual( + expected.my_instance_method(2), rref.rpc_sync().my_instance_method(2) + ) + self.assertEqual( + expected.my_instance_method(3), + rref.rpc_async().my_instance_method(3).wait(), + ) + self.assertEqual( + expected.my_instance_method(4), + rref.remote().my_instance_method(4).to_here(), + ) + + self.assertEqual( + expected.my_static_method(9), rref.rpc_sync().my_static_method(9) + ) + self.assertEqual( + expected.my_static_method(10), rref.rpc_async().my_static_method(10).wait() + ) + self.assertEqual( + expected.my_static_method(11), rref.remote().my_static_method(11).to_here() + ) + + self.assertEqual( + expected.my_class_method(2, torch.zeros(2, 2)), + rref.rpc_sync().my_class_method(2, torch.zeros(2, 2)), + ) + self.assertEqual( + expected.my_class_method(2, torch.ones(3, 3)), + rref.rpc_async().my_class_method(2, torch.ones(3, 3)).wait(), + ) + self.assertEqual( + expected.my_class_method(2, torch.ones(4, 4)), + rref.remote().my_class_method(2, torch.ones(4, 4)).to_here(), + ) + + @dist_init + def test_rref_proxy_class(self): + self._test_rref_proxy_class(worker_name((self.rank + 1) % self.world_size)) + + @dist_init + def test_rref_proxy_class_self(self): + self._test_rref_proxy_class(rpc.get_worker_info()) + + @mock.patch.object(torch.distributed.autograd, "_init") + @mock.patch.object(torch.distributed.rpc.api, "_set_and_start_rpc_agent") + @dist_init(setup_rpc=False) + def test_register_rpc_backend_and_set_and_start_rpc_backend( + self, mock_rpc_agent, mock_dist_autograd_init + ): + backend_name = "stub_backend" + + backend = rpc.backend_registry.register_backend( + backend_name, + _stub_construct_rpc_backend_options_handler, + _stub_init_rpc_backend_handler, + ) + + with self.assertRaisesRegex( + RuntimeError, "^RPC backend .+: already registered$" + ): + backend = rpc.backend_registry.register_backend( + backend_name, + _stub_construct_rpc_backend_options_handler, + _stub_init_rpc_backend_handler, + ) + + rpc.init_rpc( + name="worker1", + backend=backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + @dist_init(setup_rpc=False) + def test_duplicate_name(self): + with self.assertRaisesRegex(RuntimeError, "is not unique"): + store, _, _ = next( + torch.distributed.rendezvous( + self.init_method, rank=self.rank, world_size=self.world_size + ) + ) + rpc._init_rpc_backend( + backend=self.rpc_backend, + store=store, + name="duplicate_name", + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + @dist_init(setup_rpc=False) + def test_duplicate_name_2(self): + with self.assertRaisesRegex(RuntimeError, "is not unique"): + rpc.init_rpc( + name=worker_name(self.rank % (self.world_size - 1)), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + @dist_init(setup_rpc=False) + def test_reinit(self): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + initialize_pg(self.file_init_method, self.rank, self.world_size) + # Wait for all init to complete. + dist.barrier() + + # TODO: with TCP init, rank 0 raises Address already in use because + # rank 0 is the start daemon and the store is created before checking if + # RPC is already initialized in init_rpc. + if os.environ.get("RPC_INIT_WITH_TCP", None) == "1" and self.rank == 0: + expected_reinit_err = "Address already in use" + else: + expected_reinit_err = "is already initialized" + + with self.assertRaisesRegex(RuntimeError, expected_reinit_err): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + rpc.shutdown() + + @dist_init(setup_rpc=False) + def test_pg_init_no_rpc_init(self): + dist.init_process_group( + backend="gloo", + init_method=self.file_init_method, + rank=self.rank, + world_size=self.world_size, + ) + + class MyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.lin = torch.nn.Linear(3, 4) + + def forward(self, x): + return self.lin(x) + + model = MyModel() + model.train() + model = torch.nn.parallel.DistributedDataParallel(model) + + with self.assertRaisesRegex( + RuntimeError, + "Current RPC agent is not set! Did you initialize the RPC framework", + ): + [RRef(param) for param in model.parameters()] + + def test_world_size_one(self): + self._world_size_one(torch.ones(2, 2), torch.ones(2, 2)) + + @dist_init(setup_rpc=False) + def test_invalid_names(self): + worker_id = 0 + with self.assertRaisesRegex(RuntimeError, "Worker name must match"): + WorkerInfo("abc*", worker_id) + + with self.assertRaisesRegex(RuntimeError, "Worker name must match"): + WorkerInfo(" ", worker_id) + + with self.assertRaisesRegex(RuntimeError, "must be non-empty"): + WorkerInfo("", worker_id) + + # If the number in the message does not match, it is likely that the + # value of MAX_NAME_LEN in RPC WorkerInfo has changed. + with self.assertRaisesRegex(RuntimeError, "shorter than 128"): + WorkerInfo("".join(["a" for i in range(500)]), worker_id) + + # Test that WorkerInfo can be pickled and sent in RPC call + @dist_init + def test_worker_info_pickle(self): + dst_rank = (self.rank + 1) % self.world_size + worker_info = rpc.api.get_worker_info() + ret = rpc.rpc_sync(worker_name(dst_rank), identity, args=(worker_info,)) + self.assertEqual(ret, worker_info) + + @dist_init + def test_add(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + self.assertEqual(ret, torch.ones(n, n) * 2) + + @staticmethod + def return_callee_id(): + return rpc.get_worker_info().id + + @dist_init + def test_int_callee(self): + dst_rank = (self.rank + 1) % self.world_size + ret = rpc.rpc_sync(dst_rank, RpcTest.return_callee_id) + self.assertEqual(ret, dst_rank) + + @dist_init + def test_add_with_id(self): + n = self.rank + 1 + dst_rank = n % self.world_size + workder_info = rpc.get_worker_info(worker_name(dst_rank)) + + ret = rpc.rpc_sync( + workder_info, torch.add, args=(torch.ones(n, n), torch.ones(n, n)) + ) + self.assertEqual(ret, torch.ones(n, n) * 2) + + @dist_init + def test_scalar_add(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync(worker_name(dst_rank), torch.add, args=(torch.ones(n, n), n)) + self.assertEqual(ret, (torch.ones(n, n) + n)) + + @dist_init + def test_async_add(self): + n = self.rank + 1 + dst_rank = n % self.world_size + fut = rpc.rpc_async( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + @dist_init + def test_nonzero(self): + n = self.rank + 1 + dst_rank = n % self.world_size + x = torch.ones(self.world_size, self.world_size) + x[self.rank][self.rank] = 0 + ret = rpc.rpc_sync(worker_name(dst_rank), torch.nonzero, args=(x,)) + self.assertEqual(ret, x.nonzero()) + + @dist_init + def test_multi_rpc(self): + self._multi_rpc(False) + + @dist_init + def test_future_wait_twice(self): + dst = worker_name((self.rank + 1) % self.world_size) + futs = [rpc.rpc_async(dst, raise_func) for _ in range(20)] + + with self.assertRaisesRegex(ValueError, "Expected error"): + torch.futures.wait_all(futs) + + for fut in futs: + with self.assertRaisesRegex(ValueError, "Expected error"): + fut.wait() + + @dist_init(setup_rpc=False) + def test_wait_all_workers_timeout(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + og_func = rpc.api._wait_all_workers + + def wait_all_workers_sleep(timeout): + rpc.api._all_gather(SlowPickleClass(0.5), timeout=timeout) + + rpc.api._wait_all_workers = wait_all_workers_sleep + + try: + with self.assertRaisesRegex(RuntimeError, ""): + rpc.shutdown(graceful=True, timeout=0.01) + finally: + rpc.api._wait_all_workers = og_func + dist.barrier() + + def test_wait_all_workers_dense(self): + self._wait_all_workers(heavy_rpc, torch.ones(100, 100)) + + def test_wait_all_workers_twice_dense(self): + self._wait_all_workers_twice(heavy_rpc, torch.ones(100, 100)) + + @dist_init + def test_all_gather(self): + info = rpc.get_worker_info() + results = rpc.api._all_gather(info.id) + expected = {} + for info in rpc._get_current_rpc_agent().get_worker_infos(): + expected[info.name] = info.id + + self.assertEqual(expected, results) + + @dist_init + def test_all_gather_timeout(self): + rpc._set_rpc_timeout(0.1) + + if self.rank == 0: + with self.assertRaisesRegex( + RuntimeError, "timed out in _all_gather after 0\\.10 seconds" + ): + rpc.api._all_gather(SlowPickleClass(0.5)) + else: + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.api._all_gather(SlowPickleClass(0.5)) + + def _test_barrier_helper(self, info, names, multi_threaded=False): + names = sorted(names) + leader = names[0] + rpc.rpc_sync(leader, _reset_count) + if not multi_threaded and info.name == leader: + self.assertEqual(_rpc_barrier_count, 0) + rpc.api._barrier(names) + rpc.rpc_sync(leader, _increment_count) + rpc.api._barrier(names) + if not multi_threaded and info.name == leader: + self.assertEqual(_rpc_barrier_count, len(names)) + + @dist_init + def test_rpc_barrier_all(self): + # Test rpc barrier when called with full list of workers + info = rpc.get_worker_info() + all_worker_info = rpc._get_current_rpc_agent().get_worker_infos() + names = [worker.name for worker in all_worker_info] + self._test_barrier_helper(info, names) + + @dist_init + def test_rpc_barrier_subset(self): + # Test rpc barrier when processes are called with different subsets of the full list + info = rpc.get_worker_info() + all_worker_info = rpc._get_current_rpc_agent().get_worker_infos() + if info.id % 2: + names = [worker.name for worker in all_worker_info if worker.id % 2] + else: + names = [worker.name for worker in all_worker_info if not worker.id % 2] + self._test_barrier_helper(info, names) + + @dist_init + def test_rpc_barrier_partial_subset(self): + # Test rpc barrier when some processes are not involved in the barrier + info = rpc.get_worker_info() + all_worker_info = rpc._get_current_rpc_agent().get_worker_infos() + if info.id % 2: + names = [worker.name for worker in all_worker_info if worker.id % 2] + else: + names = [f"worker{info.id}"] + self._test_barrier_helper(info, names) + + @dist_init + def test_rpc_barrier_multithreaded(self): + # This tests validates the implementation of barrier when multiple threads call into it + # We only need to check that it does not hang in this case + info = rpc.get_worker_info() + all_worker_info = rpc._get_current_rpc_agent().get_worker_infos() + names = [worker.name for worker in all_worker_info] + threads = [] + for _ in range(3): + th = threading.Thread( + target=self._test_barrier_helper, args=(info, names, True) + ) + threads.append(th) + th.start() + for th in threads: + th.join() + + @dist_init + def test_graceful_shutdown_with_uneven_workload(self): + """Test graceful termination.""" + self._run_uneven_workload(heavy_rpc, torch.ones(100, 100)) + + @dist_init(setup_rpc=False) + def test_shutdown_followed_by_rpc(self): + # Initialize RPC. + rpc.init_rpc( + name=f"worker{self.rank:d}", + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + self.assertEqual(ret, torch.ones(n, n) * 2) + rpc.shutdown() + + with self.assertRaisesRegex(RuntimeError, "^RPC has not been initialized"): + rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + + @dist_init + def test_expected_src(self): + dst_rank = (self.rank + 1) % self.world_size + expected_src_rank = (self.rank - 1) % self.world_size + rpc.rpc_sync(worker_name(dst_rank), set_value, args=(self.rank,)) + value = VALUE_FUTURE.result() + self.assertEqual(value, expected_src_rank) + + @dist_init + def test_py_built_in(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync(worker_name(dst_rank), min, args=(n, n + 1, n + 2)) + self.assertEqual(ret, min(n, n + 1, n + 2)) + + @dist_init + def test_py_user_defined(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + my_function, + kwargs={"a": n, "b": n + 1, "c": n + 2}, + ) + self.assertEqual(ret, my_function(n, n + 1, n + 2)) + + def test_build_rpc_profiling_key(self): + # Tests that the name that shows up as an Event in profiling RPCs has all + # the necessary information. + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + rpc_profiling_key = _build_rpc_profiling_key( + exec_mode, "foo", "worker0", "worker1" + ) + self.assertIn(exec_mode.value, rpc_profiling_key) + self.assertIn("foo", rpc_profiling_key) + self.assertIn("worker0", rpc_profiling_key) + self.assertIn("worker1", rpc_profiling_key) + + def check_profiling_info( + self, self_worker_name, dst_worker_name, func, rpc_event, rpc_exec_mode + ): + self.assertTrue(self_worker_name in rpc_event.name) + self.assertTrue(dst_worker_name in rpc_event.name) + if isinstance(func, torch.jit.ScriptFunction): + self.assertTrue(torch._jit_internal._qualified_name(func) in rpc_event.name) + else: + self.assertTrue(func.__name__ in rpc_event.name) + self.assertTrue(rpc_exec_mode.value in rpc_event.name) + self.assertEqual(rpc_event.count, 1) + + @dist_init + def test_profiler_rpc_record_shapes(self): + if self.rank != 1: + return + dst = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst) + t1, t2 = torch.ones(100), torch.ones(100) + with _profile(record_shapes=True) as prof: + rpc.rpc_sync(dst_worker, torch.add, args=(t1, t2)) + + function_events = prof.function_events + remote_events = [event for event in function_events if event.is_remote] + remote_add_event = next( + event for event in remote_events if "aten::add" in event.name + ) + remote_add_input_shapes = remote_add_event.input_shapes + # Run profiler on equivalent local op and validate shapes are the same. + with _profile(record_shapes=True) as prof: + torch.add(t1, t2) + + local_function_events = prof.function_events + local_add_event = next( + event for event in local_function_events if "aten::add" in event.name + ) + local_add_input_shapes = local_add_event.input_shapes + self.assertEqual(remote_add_input_shapes, local_add_input_shapes) + + @dist_init + def test_profiler_rpc_memory(self): + if self.rank != 1: + return + dst = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst) + with _profile(profile_memory=True) as p: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + function_events = p.function_events + event_cpu_mem_usages = {event.cpu_memory_usage for event in function_events} + # if cpu_memory_usage was not propagated over the wire, this set would + # only contain 0 (indicates no memory being profiled) + self.assertNotEqual({0}, event_cpu_mem_usages) + # No memory profiled if profile_memory=False + with _profile(profile_memory=False) as p: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + function_events = p.function_events + event_cpu_mem_usages = {event.cpu_memory_usage for event in function_events} + self.assertEqual({0}, event_cpu_mem_usages) + + @dist_init + def test_profiler_export_trace(self): + if self.rank != 1: + return + dst = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst) + with _profile() as p: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + with TemporaryFileName() as fname: + path = fname + p.export_chrome_trace(path) + with open(path) as f: + trace = json.load(f) + event_names = [event["name"] for event in trace] + for expected_event_name in EXPECTED_REMOTE_EVENTS + [ + RPCExecMode.ASYNC.value + ]: + event_exists = any( + expected_event_name in event_name for event_name in event_names + ) + self.assertTrue(event_exists) + + @dist_init + def test_profiler_rpc_key_names(self): + # tests that remote events are properly prefixed with the RPC profiling key. + if self.rank != 1: + return + + # Spawn multiple threads that send RPCs to ensure keys are correctly + # prefixed when there are multiple RPCs being created/in flight at the + # same time. + dst_ranks = [rank for rank in range(self.world_size) if rank != self.rank] + + def rpc_with_profiling(dst_worker): + with _profile() as prof: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + events = prof.function_events + remote_event_names = { + event.name: event for event in events if event.is_remote + } + rpc_profiling_key = _build_rpc_profiling_key( + RPCExecMode.ASYNC, + udf_with_torch_ops.__qualname__, + worker_name(self.rank), + dst_worker, + ) + + remote_event_name_set = set(EXPECTED_REMOTE_EVENTS) + for name, event in remote_event_names.items(): + # Ensure that we have the expected key as part of the remote + # event. + self.assertTrue(name.startswith(rpc_profiling_key)) + self.assertTrue(event.is_remote) + self.assertTrue(event.node_id == rpc.get_worker_info(dst_worker).id) + # Ensure that the remote event name also contains the operator. + operator_name_substr = name[len(rpc_profiling_key) :] + # Note: we don't assert that every remote event needs to be + # in the above set, the set is just a representative set of + # what we expect to see. The profiler can change and add more + # events, but we should always expect to see this representative + # set. + matching_event = { + remote_event_name + for remote_event_name in remote_event_name_set + if remote_event_name in operator_name_substr + } + remote_event_name_set -= matching_event + + # The set should be empty, otherwise its contained elements did + # not show up in the remote profiler output. + self.assertTrue( + remote_event_name_set == set(), + f"Expected {remote_event_name_set} to be included in remote profiler output.", + ) + + for dst in dst_ranks: + dst_worker = worker_name(dst) + num_parallel_rpcs = 2 + with concurrent.futures.ThreadPoolExecutor( + max_workers=num_parallel_rpcs + ) as executor: + futs = [ + executor.submit(rpc_with_profiling, dst_worker) + for _ in range(num_parallel_rpcs) + ] + # Wait for workers to finish test + for fut in futs: + fut.result() + + def _run_test_profiler_remote_events_profiled(self): + # Tests that we can successfully invoke the profiler on a remote node, + # and collect the remote events back in the local profiler. + if self.rank != 1: + return + + dst_ranks = [rank for rank in range(self.world_size) if rank != self.rank] + for dst in dst_ranks: + dst_worker = worker_name(dst) + with _profile() as prof: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=()) + fut.wait() + + events = prof.function_events + + rpc_event = get_function_event(events, RPCExecMode.ASYNC.value) + self.check_profiling_info( + worker_name(self.rank), + dst_worker, + udf_with_torch_ops, + rpc_event, + RPCExecMode.ASYNC, + ) + + remote_events = {event.name: event for event in events if event.is_remote} + rpc_profiling_key = _build_rpc_profiling_key( + RPCExecMode.ASYNC, + udf_with_torch_ops.__qualname__, + worker_name(self.rank), + worker_name(dst), + ) + + for expected_remote_event_name in EXPECTED_REMOTE_EVENTS: + expected_key = ( + rpc_profiling_key + REMOTE_OP_STR + expected_remote_event_name + ) + self.assertTrue(expected_key in remote_events) + remote_event = remote_events[expected_key] + # Remote event should have a node ID corresponding to the worker + # it ran on. + self.assertEqual(remote_event.node_id, dst) + + # Validate order remote events show up in profiling output. + def convert_remote_to_local(event_name): + remote_op_key = rpc_profiling_key + REMOTE_OP_STR + return event_name[event_name.find(remote_op_key) + len(remote_op_key) :] + + remote_events_list = [ + convert_remote_to_local(event.name) + for event in events + if convert_remote_to_local(event.name) in EXPECTED_REMOTE_EVENTS + ] + self.assertEqual( + set(remote_events_list), + set(EXPECTED_REMOTE_EVENTS), + f"Mismatch between profiled events: {set(remote_events_list)} and expected events: {set(EXPECTED_REMOTE_EVENTS)}", + ) + + @dist_init + def test_profiler_remote_events_profiled(self): + self._run_test_profiler_remote_events_profiled() + + @dist_init + def test_profiler_remote_events_profiled_single_threaded(self): + self._run_test_profiler_remote_events_profiled() + + def run_profiling_workload(self, dst): + fut = rpc.rpc_async( + worker_name(dst), + torch.mul, + args=( + torch.tensor(1.0, requires_grad=True), + torch.tensor(1.0, requires_grad=True), + ), + ) + fut.wait() + + def _run_rpc_profiling_async_function(self, device="cpu"): + if self.rank != 1: + return + + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + x = torch.ones(2) + y = torch.ones(2) + with _profile() as prof: + ret = rpc.rpc_async( + dst1, slow_async_add, args=(dst2, x, y, device), timeout=20 + ) + ret.wait() + + function_events = prof.function_events + # slow_async_add resulted in an RPC from dst1 -> dst2, so this should be + # recorded. + key_prefix = _build_rpc_profiling_key( + RPCExecMode.ASYNC, slow_async_add.__qualname__, worker_name(self.rank), dst1 + ) + + nested_rpc_key_prefix = _build_rpc_profiling_key( + RPCExecMode.ASYNC, slow_add.__qualname__, dst1, dst2 + ) + expected_key = key_prefix + REMOTE_OP_STR + nested_rpc_key_prefix + remote_events = [event for event in function_events if event.is_remote] + rpc_remote_event = [ + event for event in remote_events if event.name == expected_key + ] + self.assertEqual(1, len(rpc_remote_event)) + rpc_remote_event = rpc_remote_event[0] + self.assertEqual(rpc_remote_event.node_id, (self.rank + 1) % self.world_size) + # slow_async_add's RPC does an add on dst2, which should be reflected as well. + remote_add_key = ( + expected_key + REMOTE_OP_STR + torch.jit._builtins._find_builtin(torch.add) + ) + remote_add_event = [ + event for event in remote_events if event.name == remote_add_key + ] + self.assertEqual(1, len(remote_add_event)) + remote_add_event = remote_add_event[0] + # Validate that node_id is dst2. + self.assertEqual(remote_add_event.node_id, (self.rank + 2) % self.world_size) + + @dist_init + def test_rpc_profiling_async_function(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + self._run_rpc_profiling_async_function() + if torch.cuda.is_available(): + dist.barrier() + self._run_rpc_profiling_async_function(device="cuda:0") + + @dist_init + def test_rpc_profiling_async_function_single_threaded(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + self._run_rpc_profiling_async_function() + if torch.cuda.is_available(): + dist.barrier() + self._run_rpc_profiling_async_function(device="cuda:0") + + @dist_init + def test_rpc_profiling_remote_record_function(self): + # test that functions run over RPC with record_function show the expected + # profiled block. + if self.rank != 1: + return + dst_ranks = [i for i in range(self.world_size) if i != self.rank] + for dst_rank in dst_ranks: + dst_worker = worker_name(dst_rank) + with _profile() as prof: + fut = rpc.rpc_async(dst_worker, udf_with_torch_ops, args=(-1, True)) + fut.wait() + + function_events = prof.function_events + record_function_remote_event = [ + evt for evt in function_events if "##forward##" in evt.name + ] + self.assertEqual(1, len(record_function_remote_event)) + record_function_remote_event = record_function_remote_event[0] + self.assertEqual(record_function_remote_event.node_id, dst_rank) + # cpu_children only returns direct children, so here we get all + # children recursively. + + def get_cpu_children(event): + if not event.cpu_children: + return [] + cpu_children = event.cpu_children + for e in event.cpu_children: + cpu_children.extend(get_cpu_children(e)) + return cpu_children + + remote_children = get_cpu_children(record_function_remote_event) + # Get local children and verify parity. + with _profile() as prof: + udf_with_torch_ops(-1, True) + + local_function_events = prof.function_events + local_record_function_event = next( + evt for evt in local_function_events if "##forward##" in evt.name + ) + local_children = get_cpu_children(local_record_function_event) + local_children_names = [evt.name for evt in local_children] + + REMOTE_OP_STR = "#remote_op: " + + def convert_remote_to_local(event_name): + remote_op_key = REMOTE_OP_STR + return event_name[event_name.find(remote_op_key) + len(remote_op_key) :] + + for evt in remote_children: + local_name = convert_remote_to_local(evt.name) + self.assertTrue(local_name in local_children_names) + + def validate_profiling_workload(self, dst, prof): + def convert_remote_to_local(event_name): + return event_name[event_name.find(REMOTE_OP_STR) + len(REMOTE_OP_STR) :] + + events = prof.function_events + remote_events = { + convert_remote_to_local(event.name): event + for event in events + if event.is_remote + } + self.assertTrue("aten::mul" in remote_events) + remote_mul_event = remote_events["aten::mul"] + self.assertEqual(remote_mul_event.node_id, dst) + self.check_profiling_info( + worker_name(self.rank), + worker_name(dst), + torch.mul, + remote_mul_event, + RPCExecMode.ASYNC, + ) + + def _run_test_profiler_with_autograd_context(self): + dst = (self.rank + 1) % self.world_size + if self.rank == 1: + # Cases where we can double wrap messages with profiling information and autograd info. + with dist_autograd.context(), _profile() as prof: + self.run_profiling_workload(dst) + + self.validate_profiling_workload(dst, prof) + + # Ensure that flipped order of ctx managers results in events being + # recorded as expected. + with _profile() as prof, dist_autograd.context(): + self.run_profiling_workload(dst) + + self.validate_profiling_workload(dst, prof) + + @dist_init + def test_profiler_with_autograd_context_single_threaded(self): + self._run_test_profiler_with_autograd_context() + + @dist_init + def test_profiler_with_autograd_context(self): + self._run_test_profiler_with_autograd_context() + + def _profiler_test_with_rpc( + self, + rpc_exec_mode, + func, + args, + use_record_function=False, + dst=None, + kineto_profile=False, + ): + dst = dst if dst is not None else (self.rank + 1) % self.world_size + + # only run profiler on rank 1. + p = _profile if not kineto_profile else torch.profiler.profile # kineto + if self.rank == 1: + with p() as prof: + record_function_ctx_mgr = ( + contextlib.nullcontext() + if not use_record_function + else torch.autograd.profiler.record_function("foo") + ) + with record_function_ctx_mgr: + if rpc_exec_mode == RPCExecMode.SYNC: + rpc.rpc_sync(worker_name(dst), func, args=args) + elif rpc_exec_mode == RPCExecMode.ASYNC: + fut = rpc.rpc_async(worker_name(dst), func, args=args) + if kineto_profile: + # Ensure multiple async RPCs don't cause issues. + # Would have raised + # "RuntimeError: Cannot call + # RemoteProfilerManager::setCurrentKey when current + # key is already set." error if RPC profiling was + # not disabled properly for kineto. + fut2 = rpc.rpc_async(worker_name(dst), func, args=args) + fut2.wait() + fut.wait() + else: + self.assertTrue(rpc_exec_mode == RPCExecMode.REMOTE) + rref = rpc.remote(worker_name(dst), func, args=args) + rref.to_here() + # To avoid flakiness, wait for the RRef to be profiled. This + # means that we received the acknowledgement of successful + # creation on the owner and ran the callbacks responsible + # for recording the profiling event. + rref._get_profiling_future().wait() + + events = prof.function_events if not kineto_profile else prof.events() + if kineto_profile: + # RPC profiling is disabled so there should be no rpc related + # events. + with self.assertRaises(IndexError): + get_function_event(events, rpc_exec_mode.value) + + return + + rpc_event = get_function_event(events, rpc_exec_mode.value) + # verify Node ID for this rpc event. + self.assertEqual(rpc_event.node_id, self.rank) + # Ensure recording of remote events. + remote_events = {event for event in events if event.node_id == dst} - { + rpc_event + } + self.assertGreaterEqual(len(remote_events), 1) + for remote_event in remote_events: + self.assertEqual(remote_event.node_id, dst) + + if use_record_function: + scope_event = get_function_event(events, "foo") + # Since RPC call is within the scope, its CPU interval should be + # contained within foo's interval. + self.assertLessEqual( + scope_event.time_range.start, rpc_event.time_range.start + ) + self.assertGreaterEqual( + scope_event.time_range.end, rpc_event.time_range.end + ) + # the sender, dest worker, function run, and type of RPC should all + # be recorded. + self_worker_name = worker_name(self.rank) + dst_worker_name = worker_name(dst) + self.check_profiling_info( + self_worker_name, dst_worker_name, func, rpc_event, rpc_exec_mode + ) + if use_record_function: + # verify order by ensuring that the outer context comes + # before the rpc event. + foo_event_ix = next( + i for i, event in enumerate(events) if "foo" in event.name + ) + rpc_event_idx = next( + i + for i, event in enumerate(events) + if rpc_exec_mode.value in event.name + ) + self.assertLess(foo_event_ix, rpc_event_idx) + + def _run_test_profiler_with_sync_rpc_udf(self): + self._profiler_test_with_rpc(RPCExecMode.SYNC, my_sleep_func, args=(1,)) + self._profiler_test_with_rpc( + RPCExecMode.SYNC, my_sleep_func, args=(1,), use_record_function=True + ) + + @dist_init + def test_profiler_with_sync_rpc_udf(self): + self._run_test_profiler_with_sync_rpc_udf() + + @dist_init + def test_profiler_with_sync_rpc_udf_single_threaded(self): + self._run_test_profiler_with_sync_rpc_udf() + + def _run_test_profiler_with_sync_rpc_builtin(self): + self._profiler_test_with_rpc( + RPCExecMode.SYNC, torch.mul, args=(torch.ones(1), torch.ones(1)) + ) + self._profiler_test_with_rpc( + RPCExecMode.SYNC, + torch.mul, + args=(torch.ones(1), torch.ones(1)), + use_record_function=True, + ) + + @dist_init + def test_profiler_with_sync_rpc_builtin(self): + self._run_test_profiler_with_sync_rpc_builtin() + + @dist_init + def test_profiler_with_sync_rpc_builtin_single_threaded(self): + self._run_test_profiler_with_sync_rpc_builtin() + + def _run_test_profiler_with_async_rpc_udf(self): + self._profiler_test_with_rpc(RPCExecMode.ASYNC, my_sleep_func, args=(1,)) + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, my_sleep_func, args=(1,), use_record_function=True + ) + # Test to ensure that kineto profiler enabled in RPC does not enable + # RPC profiling (it is unsupported) and does not result in issues. + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, my_sleep_func, args=(1,), kineto_profile=True + ) + + @dist_init + def test_profiler_with_async_rpc_udf(self): + self._run_test_profiler_with_async_rpc_udf() + + @dist_init + def test_profiler_with_async_rpc_udf_single_threaded(self): + self._run_test_profiler_with_async_rpc_udf() + + def _run_test_profiler_with_async_rpc_builtin(self): + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, torch.mul, args=(torch.ones(1), torch.ones(1)) + ) + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, + torch.mul, + args=(torch.ones(1), torch.ones(1)), + use_record_function=True, + ) + + @dist_init + def test_profiler_with_async_rpc_builtin(self): + self._run_test_profiler_with_async_rpc_builtin() + + @dist_init + def test_profiler_with_async_rpc_builtin_single_threaded(self): + self._run_test_profiler_with_async_rpc_builtin() + + def _run_test_profiler_with_remote_udf(self): + self._profiler_test_with_rpc(RPCExecMode.REMOTE, my_sleep_func, args=(1,)) + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, my_sleep_func, args=(1,), use_record_function=True + ) + # test remote to self + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, my_sleep_func, args=(1,), dst=self.rank + ) + + @dist_init + def test_profiler_with_remote_udf(self): + self._run_test_profiler_with_remote_udf() + + @dist_init + def test_profiler_with_remote_udf_single_threaded(self): + self._run_test_profiler_with_remote_udf() + + def _run_test_profiler_with_remote_builtin(self): + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, torch.mul, args=(torch.ones(1), torch.ones(1)) + ) + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, + torch.mul, + args=(torch.ones(1), torch.ones(1)), + use_record_function=True, + ) + # test remote to self + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, + torch.mul, + args=(torch.ones(1), torch.ones(1)), + dst=self.rank, + ) + + @dist_init + def test_profiler_with_remote_builtin(self): + self._run_test_profiler_with_remote_builtin() + + @dist_init + def test_profiler_with_remote_builtin_single_threaded(self): + self._run_test_profiler_with_remote_builtin() + + def _run_test_profiler_with_script_async_rpc(self): + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, my_script_func, args=(torch.tensor(1),) + ) + self._profiler_test_with_rpc( + RPCExecMode.ASYNC, + my_script_func, + args=(torch.tensor(1),), + use_record_function=True, + ) + + @dist_init + def test_profiler_with_script_async_rpc(self): + self._run_test_profiler_with_script_async_rpc() + + @dist_init + def test_profiler_with_script_async_rpc_single_threaded(self): + self._run_test_profiler_with_script_async_rpc() + + def _run_test_profiler_with_script_sync_rpc(self): + self._profiler_test_with_rpc( + RPCExecMode.SYNC, my_script_func, args=(torch.tensor(1),) + ) + self._profiler_test_with_rpc( + RPCExecMode.SYNC, + my_script_func, + args=(torch.tensor(1),), + use_record_function=True, + ) + + @dist_init + def test_profiler_with_script_sync_rpc(self): + self._run_test_profiler_with_script_sync_rpc() + + @dist_init + def test_profiler_with_script_sync_rpc_single_threaded(self): + self._run_test_profiler_with_script_sync_rpc() + + def _run_test_profiler_with_script_remote_rpc(self): + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, my_script_func, args=(torch.tensor(1),) + ) + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, + my_script_func, + args=(torch.tensor(1),), + use_record_function=True, + ) + # test remote to self + self._profiler_test_with_rpc( + RPCExecMode.REMOTE, my_script_func, args=(torch.tensor(1),), dst=self.rank + ) + + @dist_init + def test_profiler_with_script_remote_rpc(self): + self._run_test_profiler_with_script_remote_rpc() + + @dist_init + def test_profiler_with_script_remote_rpc_single_threaded(self): + self._run_test_profiler_with_script_remote_rpc() + + def _assert_top_level_events( + self, process_global_events, expected_top_level_event_names + ): + top_level_event_names = [] + for thread_local_events in process_global_events: + # Get top-level events from all events happened on a thread. + last_end_time = 0 + for event in thread_local_events: + event_name = event.name + time_range = event.time_range + if time_range.start > last_end_time: + top_level_event_names.append(event_name) + last_end_time = time_range.end + top_level_event_names = sorted(top_level_event_names) + expected_top_level_event_names = sorted(expected_top_level_event_names) + self.assertEqual( + top_level_event_names, + expected_top_level_event_names, + f"Expected events {expected_top_level_event_names}, but got {top_level_event_names}", + ) + + @dist_init + def test_server_process_global_profiler(self): + if self.rank != 0: + return + + dst_rank = (self.rank + 1) % self.world_size + dst_worker_name = worker_name(dst_rank) + + x = torch.tensor(1) + y = torch.tensor(2) + + outer_profile_rref = rpc.remote( + dst_worker_name, rpc._server_process_global_profile + ) + outer_profile_rref.rpc_sync().__enter__() + rpc.rpc_sync(dst_worker_name, torch.add, (x, y)) + inner_profile_rref = rpc.remote( + dst_worker_name, rpc._server_process_global_profile + ) + inner_profile_rref.rpc_sync().__enter__() + rpc.rpc_sync(dst_worker_name, torch.sub, (x, y)) + inner_profile_rref.rpc_sync().__exit__(None, None, None) + outer_profile_rref.rpc_sync().__exit__(None, None, None) + + inner_events = rpc.rpc_sync( + dst_worker_name, get_events_from_profile, (inner_profile_rref,) + ) + expected_inner_events = ["aten::sub"] + expected_outer_events = expected_inner_events + ["aten::add"] + + self._assert_top_level_events(inner_events, expected_inner_events) + outer_events = rpc.rpc_sync( + dst_worker_name, get_events_from_profile, (outer_profile_rref,) + ) + self._assert_top_level_events(outer_events, expected_outer_events) + + inner_profile_rref.rpc_sync().key_averages() + outer_profile_rref.rpc_sync().key_averages() + + @dist_init + def test_async_record_function_double_end_callbacks(self): + num_sleep_seconds = 1 + if self.rank == 1: + # Validate that calling the function twice results in an error. + with _profile(): + with torch.autograd.profiler.record_function("foo") as rf: + fut = rpc.rpc_async( + worker_name(0), my_sleep_func, args=(num_sleep_seconds,) + ) + rf._call_end_callbacks_on_future(fut) + with self.assertRaisesRegex( + RuntimeError, "can only be called once." + ): + rf._call_end_callbacks_on_future(fut) + fut.wait() + + @dist_init + def test_async_record_function_legacy(self): + # Test the legacy _record_function ops work + # Note: These exist for backward compatibility with TorchScript + num_sleep_seconds = 1 + if self.rank == 1: + with _profile(): + try: + handle = torch.ops.profiler._record_function_enter("foo", None) + fut = rpc.rpc_async( + worker_name(0), my_sleep_func, args=(num_sleep_seconds,) + ) + torch.ops.profiler._call_end_callbacks_on_jit_fut(handle, fut) + finally: + torch.ops.profiler._record_function_exit(handle) + + fut.wait() + + @dist_init + def test_async_record_function_cbs_jit_call(self): + if self.rank == 1: + with _profile() as pf: + key = _build_rpc_profiling_key( + RPCExecMode.ASYNC, + torch._jit_internal._qualified_name(my_script_func), + "worker1", + "worker0", + ) + with torch.autograd.profiler.record_function(key) as rf: + fut = rpc.rpc_async( + worker_name(0), my_script_func, args=(torch.tensor(1),) + ) + # Intentionally calling record_function internals + fut = torch.ops.profiler._call_end_callbacks_on_jit_fut( + rf.record, fut + ) + result = fut.wait() + # Validate that the profiling future returns the same value as the RPC + # future. + expected = torch.add(torch.tensor(1), torch.tensor(1)) + self.assertEqual(result, expected) + events = pf.function_events + rpc_event = get_function_event( + events, torch._jit_internal._qualified_name(my_script_func) + ) + self.assertTrue( + torch._jit_internal._qualified_name(my_script_func) in rpc_event.name + ) + + @dist_init + def test_py_class_constructor(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync(worker_name(dst_rank), MyClass, args=(n,)) + self.assertEqual(ret.a, n) + + @dist_init + def test_py_class_instance_method(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), MyClass(2).my_instance_method, args=(n,) + ) + self.assertEqual(ret, MyClass(2).my_instance_method(n)) + + @dist_init + def test_py_class_method(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), MyClass.my_class_method, args=(n, n + 1) + ) + self.assertEqual(ret, MyClass.my_class_method(n, n + 1)) + + @dist_init + def test_py_class_static_method(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), MyClass.my_static_method, args=(n + 10,) + ) + self.assertEqual(ret, MyClass.my_static_method(n + 10)) + + @dist_init + def test_py_multi_async_call(self): + n = self.rank + 1 + dst_rank = n % self.world_size + dst_worker_info = rpc.get_worker_info(worker_name(dst_rank)) + fut1 = rpc.rpc_async(dst_worker_info, MyClass.my_static_method, args=(n + 10,)) + fut2 = rpc.rpc_async(dst_worker_info, min, args=(n, n + 1, n + 2)) + self.assertEqual(fut1.wait(), MyClass.my_static_method(n + 10)) + self.assertEqual(fut2.wait(), min(n, n + 1, n + 2)) + + @dist_init + def test_py_no_return_result(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync(worker_name(dst_rank), no_result) + self.assertEqual(ret, no_result()) + + @dist_init + def test_py_tensors(self): + n = self.rank + 1 + dst_rank = n % self.world_size + ret = rpc.rpc_sync( + worker_name(dst_rank), + my_tensor_function, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + self.assertEqual(ret, my_tensor_function(torch.ones(n, n), torch.ones(n, n))) + + @dist_init + def test_py_tensors_multi_async_call(self): + futs = [] + n = self.rank + 1 + dst_rank = n % self.world_size + for i in range(100): + fut = rpc.rpc_async( + worker_name(dst_rank), + my_tensor_function, + args=(torch.ones(i, i), torch.ones(i, i)), + ) + futs.append(fut) + + for j, val in enumerate(torch.futures.wait_all(futs)): + self.assertEqual( + val, my_tensor_function(torch.ones(j, j), torch.ones(j, j)) + ) + + @dist_init + def test_py_tensors_in_container(self): + n = self.rank + 1 + dst_rank = n % self.world_size + a = [torch.ones(n, n), torch.ones(n, n)] + b = TensorClass(build_complex_tensors()) + c = {"foo": torch.ones(n, n), "bar": torch.ones(n, n)} + ret = rpc.rpc_sync( + worker_name(dst_rank), my_complex_tensor_function, args=(a, b, c) + ) + self.assertEqual(ret, my_complex_tensor_function(a, b, c)) + + @dist_init + def test_py_nested_pickle(self): + n = self.rank + 1 + dst_rank = n % self.world_size + + ret = rpc.rpc_sync( + worker_name(dst_rank), + run_nested_pickle, + args=(MyPickleClass(), torch.ones(2, 2)), + ) + + m = MyPickleClass() + m.set(my_tensor_function(torch.ones(2, 2), torch.ones(2, 2))) + self.assertEqual(ret, run_nested_pickle(m, torch.ones(2, 2))) + + @dist_init + def test_py_function_exception(self): + n = self.rank + 1 + dst_rank = n % self.world_size + with self.assertRaises(TypeError): + rpc.rpc_sync(worker_name(dst_rank), no_result, args=(10,)) + + @dist_init + def test_py_raise_in_user_func(self): + with captured_output() as (_, err): + # This barrier prevents a race condition where the main thread has + # not entered the context manager when the remote function runs. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + n = self.rank + 1 + dst_rank = n % self.world_size + fut = rpc.rpc_async(worker_name(dst_rank), raise_func) + with self.assertRaisesRegex(ValueError, expected_err): + fut.wait() + # This barrier prevents a race condition where the main thread exits + # context manager before the remote function has ran. + dist.barrier() + + # Validate that trainers log errors when running functions. + stderr_lines = err.getvalue() + self.assertTrue(expected_err in stderr_lines) + + @dist_init + def test_py_raise_in_user_func_escaped_str(self): + n = self.rank + 1 + dst_rank = n % self.world_size + fut = rpc.rpc_async(worker_name(dst_rank), raise_func_escape) + try: + fut.wait() + except ValueError as e: + msg = str(e) + # Ensure newlines are unescaped to provide a better repr of error. + self.assertEqual(msg, msg.encode("utf-8").decode("unicode_escape")) + else: + self.assertTrue(False, "expected raise_func_escape to raise ValueError.") + + @dist_init + def test_nested_rpc(self): + self._nested_rpc(nested_rpc, torch.ones(2, 2) + 1) + + @dist_init + def test_stress_light_rpc(self): + self._stress_test_rpc(light_rpc) + + @dist_init + def test_stress_heavy_rpc(self): + self._stress_test_rpc(heavy_rpc, repeat=20, args=(torch.ones(100, 100),)) + + @dist_init + def test_stress_heavy_rpc_torchscript(self): + self._stress_test_rpc( + heavy_rpc_torchscript, repeat=20, args=(torch.ones(100, 100),) + ) + + @dist_init + def test_builtin_remote_ret(self): + self._builtin_remote_ret( + torch.ones(2, 2), torch.ones(2, 2), torch.ones(2, 2) * 2 + ) + + @dist_init + def test_builtin_remote_self(self): + self._builtin_remote_self( + torch.ones(2, 2), torch.ones(2, 2), torch.ones(2, 2) * 2 + ) + + @staticmethod + def _multi_args_fn(n, sparse=False): + if sparse: + return (build_sparse_tensor(), build_sparse_tensor()) + else: + return (torch.ones(n, n), torch.ones(n, n)) + + @dist_init + def test_multi_builtin_remote_ret(self): + self._test_multi_remote_call(torch.add, False, args_fn=RpcTest._multi_args_fn) + + @dist_init + def test_py_udf_remote(self): + n = self.rank + 1 + dst_rank = n % self.world_size + rref = rpc.remote( + worker_name(dst_rank), + my_function, + kwargs={"a": n, "b": n + 1, "c": n + 2}, + ) + self.assertEqual(rref.to_here(), my_function(n, n + 1, n + 2)) + + @staticmethod + def _multi_kwargs_fn(n, sparse=False): + if sparse: + return { + "a": build_sparse_tensor(), + "b": build_sparse_tensor(), + "c": build_sparse_tensor(), + } + else: + return {"a": torch.ones(n, n), "b": torch.ones(n, n), "c": torch.ones(n, n)} + + @dist_init + def test_multi_py_udf_remote(self): + self._test_multi_remote_call( + my_function, False, kwargs_fn=RpcTest._multi_kwargs_fn + ) + + @dist_init + def test_py_rref_args(self): + self._py_rref_args( + torch.ones(2, 2), 1, torch.ones(2, 2), 2, torch.ones(2, 2) * 2 + 3 + ) + + @dist_init + def test_py_rref_args_user_share(self): + self._py_rref_args_user_share( + torch.ones(2, 2), 1, 2, torch.ones(2, 2), 3, 4, torch.ones(2, 2) * 2 + 10 + ) + + @dist_init + def test_py_rpc_rref_args(self): + self._py_rpc_rref_args( + torch.ones(2, 2), 1, 2, torch.ones(2, 2), 3, 4, torch.ones(2, 2) * 2 + 10 + ) + + @dist_init + def test_nested_remote(self): + self._nested_remote(nested_remote, torch.ones(2, 2) + 3) + + @dist_init + def test_nested_rref(self): + self._nested_rref(nested_rref, torch.ones(2, 2) + 1, torch.ones(2, 2) + 2) + + @dist_init + def test_nested_rref_stress(self): + self._nested_rref_stress( + nested_rref, torch.ones(2, 2) + 1, torch.ones(2, 2) + 2 + ) + + @dist_init + def test_multi_layer_nested_async_rpc(self): + # This test will exit right away, but there will be a chain of async + # RPCs. The termination algorithm should detect those messages properly. + # Otherwise, some peer could exit early, leaving others to timeout + # errors or connection closed errors. + ttl = 20 + n = self.rank + 1 + dst_rank = n % self.world_size + + multi_layer_nested_async_rpc(dst_rank, self.world_size, ttl) + + @dist_init + def test_remote_with_exception(self): + n = self.rank + 1 + dst_rank = n % self.world_size + # check ref to other workers + rref = rpc.remote(worker_name(dst_rank), raise_func) + with self.assertRaises(ValueError): + rref.to_here() + # check ref to itself + rref = rpc.remote(worker_name(self.rank), no_result, args=(10,)) + with self.assertRaises(TypeError): + rref.to_here() + + @dist_init + def test_rpc_return_rref(self): + n = self.rank + 1 + dst_rank1 = n % self.world_size + dst_rank2 = (n + 1) % self.world_size + rref = rpc.rpc_sync( + worker_name(dst_rank1), + rpc_return_rref, + args=(worker_name(dst_rank2),), + ) + self.assertEqual(rref.to_here(), torch.ones(2, 2) + 1) + + @dist_init + def test_rref_forward_chain(self): + ttl = 8 + n = self.rank + 1 + dst_rank = n % self.world_size + + rref = rpc.remote(worker_name(dst_rank), torch.add, args=(torch.ones(n, n), 1)) + + ret_rref = rref_forward_chain(dst_rank, self.world_size, rref, ttl) + + for _ in range(ttl): + self.assertEqual(len(ret_rref), 1) + ret_rref = ret_rref[0].to_here() + + ret = ret_rref + self.assertEqual(ret, torch.add(torch.ones(n, n), 1)) + + @dist_init + def test_local_rref_no_fork(self): + local_rref = RRef(35) + self.assertEqual(local_rref.local_value(), 35) + + @dist_init + def test_local_value_not_on_owner(self): + # ensure that an error message is thrown if a user tries to call + # local_value() on a non-owning node. + next_rank = (self.rank + 1) % self.world_size + rref = rpc.remote( + worker_name(next_rank), torch.add, args=(torch.ones(1), torch.ones(1)) + ) + with self.assertRaisesRegex( + RuntimeError, + ( + rf"For UserRRef\(rref_id=GloballyUniqueId\(created_on={self.rank}, local_id=0\), " + rf"fork_id=GloballyUniqueId\(created_on={self.rank}, local_id=1\)\), " + r"can't call localValue\(\) on user " + rf"WorkerInfo\(id={self.rank}, name={worker_name(self.rank)}\). " + rf"Call it on owner WorkerInfo\(id={next_rank}, name={worker_name(next_rank)}\)" + ), + ): + rref.local_value() + + @dist_init + def test_return_local_rrefs(self): + n = self.rank + 1 + dst_rank = n % self.world_size + + rref_list = rpc.rpc_sync( + worker_name(dst_rank), get_rref_list, args=([1, 2, 3],) + ) + + for rref in rref_list: + rpc.rpc_sync( + rref.owner(), + _call_method_on_rref, + args=(MyClass.increment_value, rref, 10), + ) + + rets = [ + rpc.rpc_sync( + rref.owner(), _call_method_on_rref, args=(MyClass.get_value, rref) + ) + for rref in rref_list + ] + + self.assertEqual(rets, [11, 12, 13]) + + @dist_init + def _test_rref_type(self, blocking): + def launched_rpc(events): + expected_name = f"rpc_{RPCExecMode.ASYNC.value}#_rref_typeof_on_owner" + return any(e.name.startswith(expected_name) for e in events) + + dst = worker_name((self.rank + 1) % self.world_size) + rref = rpc.remote(dst, torch.add, args=(torch.ones(2), 1)) + + with _profile() as p: + t = rref._get_type(blocking=blocking) + if not blocking: + t = t.wait() + + self.assertTrue(launched_rpc(p.function_events)) + expected_type = type(torch.ones(2)) + self.assertEqual(t, expected_type) + + futs = [] + + def verify(fut): + self.assertEqual(fut.value(), expected_type) + + with _profile() as p: + for _ in range(10): + t = rref._get_type(blocking=blocking) + if not blocking: + futs.append(t) + t.add_done_callback(verify) + t = t.wait() + self.assertEqual(t, expected_type) + + if not blocking: + # Note that cached calls with blocking=False all return the same + # cached original future. + first_fut = futs[0] + for f in futs[1:]: + self.assertTrue(f is first_fut) + # Ensure we never launch another RPC, other than for the very + # first call. + self.assertFalse(launched_rpc(p.function_events)) + self.assertEqual(t, type(torch.ones(2))) + + rref = rpc.remote(dst, MyClass, args=(0,)) + rref_type = rref._get_type(blocking=blocking) + if not blocking: + rref_type = rref_type.wait() + self.assertEqual(rref_type, MyClass) + + def test_rref_type_blocking(self): + self._test_rref_type(blocking=True) + + def test_rref_type_non_blocking(self): + self._test_rref_type(blocking=False) + + @dist_init + def _test_rref_type_with_error(self, blocking): + dst = worker_name((self.rank + 1) % self.world_size) + # 10 ms timeout + rref = rpc.remote(dst, raise_func) + # Blocking: error raised inline + if blocking: + with self.assertRaisesRegex(ValueError, "Expected error"): + rref._get_type(blocking=blocking) + else: + # Non-blocking: Immediately return future, block on wait + fut = rref._get_type(blocking=blocking) + with self.assertRaisesRegex(ValueError, "Expected error"): + fut.wait() + + def test_rref_type_with_error_blocking(self): + self._test_rref_type_with_error(blocking=True) + + def test_rref_type_with_error_non_blocking(self): + self._test_rref_type_with_error(blocking=False) + + @dist_init + def _test_rref_type_owner(self, blocking): + rref = RRef(torch.ones(2) + 1) + rref_type = rref._get_type(blocking=blocking) + if not blocking: + rref_type = rref_type.wait() + self.assertEqual(rref_type, type(torch.ones(2))) + + rref = RRef(MyClass(0)) + rref_type = rref._get_type(blocking=blocking) + if not blocking: + rref_type = rref_type.wait() + self.assertEqual(rref_type, MyClass) + + def test_rref_type_owner_blocking(self): + self._test_rref_type_owner(blocking=True) + + def test_rref_type_owner_non_blocking(self): + self._test_rref_type_owner(blocking=False) + + @staticmethod + def _slow_add(x, y): + time.sleep(1) + return x + y + + @dist_init + def test_rref_type_slow_init(self): + dst = worker_name((self.rank + 1) % self.world_size) + rref = rpc.remote(dst, RpcTest._slow_add, args=(torch.ones(2), 1)) + self.assertEqual(rref._get_type(), type(torch.ones(2))) + + @dist_init + def test_owner_equality(self): + a = RRef(40) + b = RRef(50) + + other_rank = (self.rank + 1) % self.world_size + other_a = rpc.remote( + worker_name(other_rank), torch.add, args=(torch.ones(1), 1) + ) + other_b = rpc.remote( + worker_name(other_rank), torch.add, args=(torch.ones(1), 1) + ) + other_a.to_here() # to ensure clean termination + other_b.to_here() + + self.assertNotEqual(a.owner(), 23) + self.assertEqual(other_a.owner(), other_b.owner()) + self.assertNotEqual(a.owner(), other_a.owner()) + self.assertEqual(other_a.owner(), other_a.owner()) + self.assertEqual(other_a.owner(), other_b.owner()) + self.assertEqual(a.owner(), a.owner()) + self.assertEqual(a.owner(), b.owner()) + self.assertEqual(a.owner(), rpc.get_worker_info()) + x = {} + x[a.owner()] = a + x[other_a.owner()] = other_a + self.assertEqual(x[a.owner()], a) + self.assertEqual(x[b.owner()], a) + self.assertEqual(x[other_a.owner()], other_a) + self.assertEqual(x[other_b.owner()], other_a) + self.assertEqual(len(x), 2) + + @dist_init + def test_pass_local_rrefs(self): + n = self.rank + 1 + dst_rank = n % self.world_size + dst_worker = worker_name(dst_rank) + + rref = RRef(40) + self.assertEqual( + rpc.rpc_sync(dst_worker, add_rref_to_value, args=(rref, 50)), 90 + ) + self.assertEqual( + rpc.rpc_async(dst_worker, add_rref_to_value, args=(rref, 50)).wait(), 90 + ) + self.assertEqual( + rpc.remote(dst_worker, add_rref_to_value, args=(rref, 50)).to_here(), 90 + ) + + @dist_init + def test_remote_same_worker(self): + n = self.rank + 1 + dst_rank = n % self.world_size + rref_a = rpc.remote( + worker_name(dst_rank), torch.add, args=(torch.ones(n, n), 2) + ) + rref_b = rpc.remote( + worker_name(dst_rank), torch.add, args=(torch.ones(n, n), 1) + ) + rref_c = rpc.remote( + worker_name(dst_rank), my_rref_function, args=(rref_a, rref_b) + ) + self.assertEqual(rref_c.to_here(), torch.ones(n, n) + 4) + + @dist_init(setup_rpc=True) + def test_call_method_on_rref(self): + """ + Tests that it is possible to call an instance method on a remote object + by using rref.owner() as destination of the call. + """ + vals = [10, 2, 5, 7] + dst_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst_rank) + + # creates a remote object + rref = rpc.remote(dst_worker, MyClass, args=(vals[0],)) + + # modifies state of the remote object + rpc.rpc_sync( + rref.owner(), + _call_method_on_rref, + args=(MyClass.increment_value, rref, vals[1]), + ) + rpc.rpc_async( + rref.owner(), + _call_method_on_rref, + args=(MyClass.increment_value, rref, vals[2]), + ).wait() + rpc.remote( + rref.owner(), + _call_method_on_rref, + args=(MyClass.increment_value, rref, vals[3]), + ).to_here() + + # queries state of the remote object + result = rpc.rpc_sync( + dst_worker, _call_method_on_rref, args=(MyClass.get_value, rref) + ) + + self.assertEqual(result, sum(vals)) + + # Notice `rpc.api.shutdown()` accesses + # `_delete_all_user_and_unforked_owner_rrefs` through + # `torch.distributed.rpc.api`, so patching + # `torch.distributed.rpc._delete_all_user_and_unforked_owner_rrefs` will + # not help. + @mock.patch.object( + torch.distributed.rpc.api, "_delete_all_user_and_unforked_owner_rrefs" + ) + def _test_rref_leak( + self, _mock_delete_all_user_and_unforked_owner_rrefs, ignore_leak + ): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + initialize_pg(self.file_init_method, self.rank, self.world_size) + # Wait for all init to complete. + dist.barrier() + + rref = rpc.remote( # noqa: F841 + worker_name((self.rank + 1) % self.world_size), + torch.add, + args=(torch.ones(2, 2), 1), + ) + + import torch.distributed.rpc.api as api + + if ignore_leak: + api._ignore_rref_leak = True + rpc.shutdown(graceful=True) + else: + api._ignore_rref_leak = False + with self.assertRaisesRegex(RuntimeError, "Leaking RRef"): + rpc.shutdown(graceful=True) + + @dist_init(setup_rpc=False) + def test_rref_leak(self): + self._test_rref_leak(ignore_leak=False) + + @dist_init(setup_rpc=False) + def test_ignore_rref_leak(self): + self._test_rref_leak(ignore_leak=True) + + @dist_init + def test_rref_str(self): + rref1 = RRef(self.rank) + id_class = "GloballyUniqueId" + self.assertEqual( + f"OwnerRRef({id_class}(created_on={self.rank}, local_id=0))", + rref1.__str__(), + ) + + dst_rank = (self.rank + 1) % self.world_size + rref2 = rpc.remote(worker_name(dst_rank), torch.add, args=(torch.ones(2, 2), 1)) + self.assertEqual( + rref2.__str__(), + f"UserRRef(RRefId = {id_class}(created_on={self.rank}, local_id=1), " + f"ForkId = {id_class}(created_on={self.rank}, local_id=2))", + ) + + @dist_init + def test_rref_get_future(self): + # Tests that we can obtain the future corresponding to the creation of + # the RRef on remote end + if self.rank == 0: + # Builtin + rref = rpc.remote(worker_name(1), torch.add, args=(1, 1)) + rref.to_here() + fut = rref._get_future() + self.assertIsInstance(fut, torch._C.Future) + + # UDF + rref = rpc.remote(worker_name(1), foo_add, args=()) + rref.to_here() + fut = rref._get_future() + self.assertIsInstance(fut, torch._C.Future) + + # Script + rref = rpc.remote(worker_name(1), my_script_func, args=(torch.tensor(1),)) + rref.to_here() + fut = rref._get_future() + self.assertIsInstance(fut, torch._C.Future) + + @dist_init + def test_rref_context_debug_info(self): + # This test checks local states that are modified by remote workers. + # This means that we would need barrier before and after every check. + # The barrier before the check makes sure that all previous states are + # cleared globally, the barrier after ensures that no following states + # change gets into the current check. + initialize_pg(self.file_init_method, self.rank, self.world_size) + + # Check 1: local RRef does not update owners_ map or add a pending user. + ################################################# + + rref1 = RRef(self.rank) + + # don't need a barrier here as local RRef is handled by this thread + info = _rref_context_get_debug_info() + self.assertIn("num_owner_rrefs", info) + self.assertIn("num_pending_users", info) + # RRef on local value is not added to context until shared across RPC + self.assertEqual(0, int(info["num_owner_rrefs"])) + self.assertEqual(0, int(info["num_pending_users"])) + # barrier after the check 1 + dist.barrier() + + # Check 2: Sharing RRef as an arg should update owners_ map + ########################################################### + + dst_rank = (self.rank + 1) % self.world_size + rpc.rpc_sync(worker_name(dst_rank), set_global_rref, args=(rref1,)) + + # barrier before check 2 + wait_until_pending_futures_and_users_flushed() + dist.barrier() + + info = _rref_context_get_debug_info() + self.assertIn("num_owner_rrefs", info) + self.assertEqual(1, int(info["num_owner_rrefs"])) + # no pending users since the fork is finished + self.assertEqual(0, int(info["num_pending_users"])) + # barrier after check 2 + dist.barrier() + + # clear states for check 2 + rpc.rpc_sync(worker_name(dst_rank), clear_global_rref) + + # Wait for owner rref to be cleared. + while int(info["num_owner_rrefs"]) != 0: + info = _rref_context_get_debug_info() + time.sleep(0.1) + dist.barrier() + + # Check 3: rpc.remote call should update owners_ map + #################################################### + rref2 = rpc.remote(worker_name(dst_rank), torch.add, args=(torch.ones(2, 2), 1)) + rref3 = rpc.remote(worker_name(dst_rank), torch.add, args=(torch.ones(2, 2), 1)) + rref2.to_here() + rref3.to_here() + + # barrier before check 3 + wait_until_pending_futures_and_users_flushed() + dist.barrier() + + info = _rref_context_get_debug_info() + self.assertIn("num_owner_rrefs", info) + self.assertEqual(2, int(info["num_owner_rrefs"])) + # no pending users since the fork is finished + self.assertEqual(0, int(info["num_pending_users"])) + + # barrier after check 3 + dist.barrier() + + @dist_init + def test_disable_gil_profiling(self): + # test that rpc.enable_gil_profiling(false) will result in + # GIL wait time not being recorded. + + # GIL profiling should be disabled by default. + dst_rank = (self.rank + 1) % self.world_size + rpc.rpc_sync( + worker_name(dst_rank), torch.add, args=(torch.ones(1), torch.ones(1)) + ) + info = rpc.api._get_current_rpc_agent().get_debug_info() + self.assertRaises(KeyError, lambda: info["agent.gil_average_wait_time_us"]) + rpc.enable_gil_profiling(True) + rpc.rpc_sync( + worker_name(dst_rank), torch.add, args=(torch.ones(1), torch.ones(1)) + ) + info = rpc.api._get_current_rpc_agent().get_debug_info() + self.assertIn("agent.gil_average_wait_time_us", info) + + @dist_init(setup_rpc=False) + def test_local_shutdown(self): + # test that we can start RPC and then immediately locally shutdown + # without sending any messages. + rpc.init_rpc( + name=f"worker{self.rank:d}", + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + # pass in graceful=False to ensure that we don't wait for other workers. + rpc.shutdown(graceful=False) + + @dist_init + def test_debug_info(self): + # only test keys in this test case. Values should be covered by + # individual module debug info tests + import torch.distributed.autograd as dist_autograd + + info = _get_debug_info() + rref_info = _rref_context_get_debug_info() + agent_info = rpc.api._get_current_rpc_agent().get_debug_info() + autograd_info = dist_autograd._get_debug_info() + common_keys = rref_info.keys() & agent_info.keys() & autograd_info.keys() + self.assertEqual(0, len(common_keys)) + expected = {} + expected.update(rref_info) + expected.update(agent_info) + expected.update(autograd_info) + self.assertEqual(info.keys(), expected.keys()) + + @dist_init(setup_rpc=False) + @skip_but_pass_in_sandcastle_if( + IS_MACOS, + "Test is flaky on MacOS since libuv error handling is not as robust as TCP", + ) + def test_handle_send_exceptions(self): + # test that if a callee node has gone down, we raise an appropriate + # exception instead of just crashing. + rpc.init_rpc( + name=f"worker{self.rank:d}", + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + rpc._set_rpc_timeout(10) + # This barrier is needed to ensure that some workers do not exit before + # others have been brought up. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + if self.rank == 1: + dst_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst_rank) + # allow destination worker to exit without joining + error_str = self.get_shutdown_error_regex() + wait_until_node_failure(dst_rank, error_str) + fut = rpc.rpc_async(dst_worker, torch.add, args=(torch.ones(1), 3)) + # Shutdown sequence is not very well defined and as a result + # we can see any of the error messages defined in get_shutdown_error_regex. + with self.assertRaisesRegex(RuntimeError, error_str): + fut.wait() + # exit all workers non-gracefully. + rpc.shutdown(graceful=False) + + @dist_init + def test_deadlock(self): + # this test is copied from https://github.com/pytorch/pytorch/issues/45089 + if self.rank == 1: + dst1 = worker_name((self.rank + 1) % self.world_size) + x = torch.ones(2) + y = torch.ones(2) + rpc.rpc_async(dst1, RpcTest._slow_add, args=(x, y), timeout=15).wait() + + dist_initialized = dist.is_initialized() + if not dist_initialized: + dist.init_process_group( + backend="gloo", + init_method=self.file_init_method, + rank=self.rank, + world_size=self.world_size, + ) + + @dist_init(setup_rpc=False) + def test_local_shutdown_with_rpc(self): + # test that we can start RPC, send RPCs, and then run local shutdown. + rpc.init_rpc( + name=f"worker{self.rank:d}", + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + n = self.rank + 1 + dst_rank = n % self.world_size + rpc.rpc_sync( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + # A barrier is needed to ensure that all RPCs are processed. + # Otherwise, some RPCs can timeout since the receiving end + # has terminated. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + # pass in graceful=False to ensure that we don't wait for other workers. + rpc.shutdown(graceful=False) + + @dist_init(setup_rpc=False) + def test_set_and_get_default_rpc_timeout(self): + timeout = 0.5 + + # A new `RpcBackendOptions` is constructed + # when accessing `self.rpc_backend_options`. + rpc_backend_options = self.rpc_backend_options + rpc_backend_options.rpc_timeout = timeout + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + set_timeout = rpc.get_rpc_timeout() + self.assertEqual(timeout, set_timeout) + rpc.shutdown() + + @dist_init + def test_default_timeout_used(self): + """ + Tests that if no timeout is passed into rpc_async and rpc_sync, then the + default timeout is used. + """ + dst_rank = (self.rank + 1) % self.world_size + rpc._set_rpc_timeout(0.001) # 1 ms + # futures should time out and be marked with an exception indicating it as such. + futs = [ + rpc.rpc_async(worker_name(dst_rank), my_sleep_func, args=()) + for _ in range(10) + ] + expected_error = self.get_timeout_error_regex() + for fut in futs: + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # ensure that if a new timeout is set old futures don't time out but new ones do. + rpc._set_rpc_timeout(200) # 200 seconds + # create a longstanding RPC. + fut1 = rpc.rpc_async(worker_name(dst_rank), my_sleep_func, args=(1,)) + # now, set a short timeout. + rpc._set_rpc_timeout(0.001) + # fut2 should time out, fut1 should not. + fut2 = rpc.rpc_async(worker_name(dst_rank), my_sleep_func, args=(1,)) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut2.wait() + fut1.wait() + + # Zero timeout means infinity, so future should run to completion. + rpc._set_rpc_timeout(0) + rpc.rpc_async(worker_name(dst_rank), my_sleep_func, args=()).wait() + + # reset to default timeout so shutdown messages can process cleanly. + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + @dist_init + def test_rpc_timeouts(self): + # TODO: enable timeouts for rpc.remote/RRef (https://github.com/pytorch/pytorch/issues/33803) + dst_rank = (self.rank + 1) % self.world_size + dst_worker = worker_name(dst_rank) + timeout = 0.1 # 100 ms + expected_error = self.get_timeout_error_regex() + # Test async UDF + fut = rpc.rpc_async(dst_worker, my_sleep_func, args=(1,), timeout=timeout) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + + # Ensure run to completion if there is no timeout and we use the default + # RPC timeout. + rpc.rpc_async(dst_worker, my_sleep_func, args=(1,)).wait() + + # Test sync UDF + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,), timeout=timeout) + + # Ensure run to completion if there is no timeout and we use the default + # RPC timeout. + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,)) + + # If we set a default timeout for RPCs, it should be respected, though + # still overridden if we pass in a different timeout to the APIs. + rpc._set_rpc_timeout(0.001) + fut = rpc.rpc_async(dst_worker, my_sleep_func, args=(1,)) + with self.assertRaisesRegex(RuntimeError, expected_error): + fut.wait() + with self.assertRaisesRegex(RuntimeError, expected_error): + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,)) + + # The RPCs should run to completion since we override the timeout. + rpc.rpc_async(dst_worker, my_sleep_func, args=(1,), timeout=5).wait() + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,), timeout=5) + # Passing in a zero timeout should ensure that the RPC won't time out. + rpc.rpc_async(dst_worker, my_sleep_func, args=(1,), timeout=0).wait() + rpc.rpc_sync(dst_worker, my_sleep_func, args=(1,), timeout=0) + # Reset for clean shutdown + rpc._set_rpc_timeout(rpc.constants.DEFAULT_RPC_TIMEOUT_SEC) + + def test_dist_init_decorator(self): + @dist_init(setup_rpc=False) + def test_func(self): + return "expected result" + + self.assertEqual(test_func(self), "expected result") + + @dist_init + def test_func(self): + return "expected result" + + self.assertEqual(test_func(self), "expected result") + + def test_use_rpc_pickler(self): + class TestPickler: + pass + + test_pickler = TestPickler() + with _use_rpc_pickler(test_pickler): + self.assertTrue(torch.distributed.rpc.api._default_pickler is test_pickler) + self.assertTrue( + torch.distributed.rpc.api._default_pickler is _internal_rpc_pickler + ) + + @dist_init + def test_wait_all(self): + with _wait_all(): + self.assertTrue(_thread_local_var.future_list == []) + dst = worker_name((self.rank + 1) % self.world_size) + fut = rpc.rpc_async(dst, torch.add, (torch.ones(2, 2), 1)) + self.assertTrue(len(_thread_local_var.future_list) == 1) + self.assertTrue( + isinstance(_thread_local_var.future_list[0], torch._C.Future) + ) + self.assertTrue(fut.done()) + self.assertEqual(fut.wait(), torch.ones(2, 2) + 1) + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_wait_all_multiple_call(self): + with _wait_all(): + self.assertTrue(_thread_local_var.future_list == []) + dst = worker_name((self.rank + 1) % self.world_size) + for i in range(20): + fut = rpc.rpc_async(dst, torch.add, (torch.ones(i, i), 1)) + res = rpc.rpc_sync(dst, torch.add, (torch.ones(i, i), 1)) + self.assertEqual(res, torch.ones(i, i) + 1) + self.assertEqual(fut.wait(), torch.ones(i, i) + 1) + self.assertTrue(len(_thread_local_var.future_list) == 20) + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_wait_all_timeout(self): + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error), _wait_all(): + self.assertTrue(_thread_local_var.future_list == []) + dst = worker_name((self.rank + 1) % self.world_size) + timeout = 0.1 # 100 ms + rpc.rpc_async(dst, my_sleep_func, args=(1,), timeout=timeout) + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_wait_all_raise_in_user_func(self): + with self.assertRaises(ValueError), _wait_all(): + self.assertTrue(_thread_local_var.future_list == []) + dst = worker_name((self.rank + 1) % self.world_size) + rpc.rpc_async(dst, raise_func) + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_wait_all_raise_in_body(self): + with self.assertRaises(ValueError), _wait_all(): + raise_func() + self.assertFalse(hasattr(_thread_local_var, "future_list")) + + @dist_init + def test_custom_exception_throw_during_reconstruction(self): + """ + Test that we still throw info about the remote side exception even when + we cannot recreate it on client side. + """ + initialize_pg(self.file_init_method, self.rank, self.world_size) + if self.rank != 0: + exc_caught = False + dst = worker_name(0) + try: + rpc.rpc_sync(dst, custom_raise_func, args=()) + except RuntimeError as e: + exc_caught = True + msg = str(e) + print(f"Got msg {msg}") + self.assertTrue("Original exception on remote side was" in msg) + self.assertTrue("CustomException" in msg) + except BaseException as e: # noqa: B036 + raise RuntimeError(f"Failure - expected RuntimeError, got {e}") from e + finally: + self.assertTrue(exc_caught) + + dist.barrier() + + timed_out_rpc_event = None + + @staticmethod + def timed_out_rpc(): + RpcTest.timed_out_rpc_event.wait() + + @dist_init + def test_wait_all_exit_early_python(self): + # Initialize the event in the subprocess. + RpcTest.timed_out_rpc_event = Event() + + # Wait for all processes to initialize event. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + + dst = worker_name((self.rank + 1) % self.world_size) + fut1 = rpc.rpc_async(dst, RpcTest.timed_out_rpc) + fut2 = rpc.rpc_async(dst, raise_func) + fut3 = rpc.rpc_async(dst, raise_func) + + # We should receive the error from fut2 + with self.assertRaisesRegex(ValueError, expected_err): + torch.futures.wait_all([fut1, fut2, fut3]) + + # Unblock RPC thread for fut1 + RpcTest.timed_out_rpc_event.set() + + @dist_init + def test_wait_all_exit_early_builtin(self): + # Initialize the event in the subprocess. + RpcTest.timed_out_rpc_event = Event() + + # Wait for all processes to initialize event. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + + dst = worker_name((self.rank + 1) % self.world_size) + fut1 = rpc.rpc_async(dst, RpcTest.timed_out_rpc) + fut2 = rpc.rpc_async(dst, torch.add, args=(torch.rand(10), torch.rand(5))) + fut3 = rpc.rpc_async(dst, torch.add, args=(torch.rand(10), torch.rand(5))) + + # We should receive the error from fut2 + with self.assertRaisesRegex(RuntimeError, "size of tensor"): + torch.futures.wait_all([fut1, fut2, fut3]) + + # Unblock RPC thread for fut1 + RpcTest.timed_out_rpc_event.set() + + @dist_init + def test_wait_all_exit_early_script_function(self): + # Initialize the event in the subprocess. + RpcTest.timed_out_rpc_event = Event() + + # Wait for all processes to initialize event. + initialize_pg(self.file_init_method, self.rank, self.world_size) + dist.barrier() + + dst = worker_name((self.rank + 1) % self.world_size) + fut1 = rpc.rpc_async(dst, RpcTest.timed_out_rpc) + fut2 = rpc.rpc_async(dst, raise_func_script, args=(expected_err,)) + fut3 = rpc.rpc_async(dst, raise_func_script, args=(expected_err,)) + + # We should receive the error from fut2 + with self.assertRaisesRegex(RuntimeError, expected_err): + torch.futures.wait_all([fut1, fut2, fut3]) + + # Unblock RPC thread for fut1 + RpcTest.timed_out_rpc_event.set() + + @dist_init + def test_function_not_on_callee(self): + # test that if a function does not exist on a callee, we don't crash, + # instead we get an AttributeError indicating that the func does not exist. + this_module = sys.modules[__name__] + caller_worker = "worker0" + callee_worker = "worker1" + + if self.rank == 1: + # Use delattr to remove the binding of a func on this nodes + delattr(this_module, "foo_add") + # notify remote end that we have removed it. + rpc.rpc_sync(caller_worker, set_value, args=(self.rank,)) + + if self.rank == 0: + # func exists on caller, but not callee. + # wait for remote end to remove the binding of foo_add func. + wait_for_value_future() + # Ensure that we have the attribute on this module. Otherwise, the test could fail due to a caller-side pickling error. + self.assertTrue(hasattr(this_module, "foo_add")) + with self.assertRaisesRegex(RuntimeError, "RPC pickler does not serialize"): + rpc.rpc_sync(callee_worker, foo_add, args=()) + + @dist_init + def test_non_garbage_collected_user_rref_due_to_local_circular_dependency(self): + dst_worker_name = worker_name((self.rank + 1) % self.world_size) + + a = MyClass(1) + b = MyClass(2) + + # This is to make Python not garbage collect a and b. + a.other = b + b.other = a + + n = self.rank + a.rref = rpc.remote(dst_worker_name, torch.add, args=(torch.ones(n, n), 2)) + + @dist_init(setup_rpc=False) + def test_use_rref_after_shutdown(self): + rpc.init_rpc( + name=f"worker{self.rank:d}", + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + n = self.rank + 1 + dst_rank = n % self.world_size + rref = rpc.remote( + worker_name(dst_rank), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + # pass in graceful=True to ensure that local UserRRefs are deleted. + rpc.shutdown(graceful=True) + + with self.assertRaisesRegex( + RuntimeError, "Cannot call to_here\\(\\) on it after deletion." + ): + rref.to_here() + + with self.assertRaisesRegex( + RuntimeError, "Cannot call fork an UserRRef after deletion." + ): + import torch.distributed.rpc.internal as internal + + internal.serialize(rref) + + @staticmethod + def _return_gpu_tensor(): + return torch.rand(3, 3).cuda(0) + + @staticmethod + def _return_gpu_tensor_list(): + return [torch.rand(3, 3).cuda(0), torch.rand(3, 3).cuda(1)] + + @staticmethod + def _gpu_tensor_list_arg(tensor_list): + return torch.rand(3, 3) + + def _create_rref(self): + owner_rank = (self.rank + 2) % self.world_size + return rpc.remote( + worker_name(owner_rank), torch.add, args=(torch.zeros(2, 2), 1) + ) + + @dist_init + def test_user_rrefs_confirmed(self): + dst_rank = (self.rank + 1) % self.world_size + rref = self._create_rref() + ret = rpc.rpc_sync(worker_name(dst_rank), check_rref_confirmed, args=(rref,)) + self.assertEqual(ret, True) + + @dist_init + def test_user_rrefs_confirmed_remote(self): + dst_rank = (self.rank + 1) % self.world_size + rref = self._create_rref() + ret_rref = rpc.remote(worker_name(dst_rank), check_rref_confirmed, args=(rref,)) + self.assertEqual(ret_rref.to_here(), True) + + @dist_init + def test_rref_py_pickle_not_supported(self): + local_rref = RRef(35) + with ( + TemporaryFileName() as fname, + self.assertRaisesRegex( + RuntimeError, "Can not pickle rref in python pickler" + ), + ): + torch.save(local_rref, fname) + + @dist_init + def test_remote_throw(self): + rref = rpc.remote( + worker_name((self.rank + 1) % self.world_size), + raise_or_inc, + args=(torch.ones(2),), + ) + with self.assertRaisesRegex(Exception, ".*Expected error.*"): + rref.to_here() + + @dist_init + def test_non_cont_tensors(self): + if self.rank == 0: + # Create a non-contiguous tensor. + t = torch.rand(5, 5) + t_view = t.narrow(1, 2, 2) + self.assertFalse(t_view.is_contiguous()) + t_cont = t_view.contiguous() + self.assertTrue(t_cont.is_contiguous()) + self.assertEqual(t_view, t_cont) + + # Send non-cont tensor over RPC. + next_rank = (self.rank + 1) % self.world_size + t_ret = rpc.rpc_sync( + worker_name(next_rank), non_cont_test, args=(t_view, t_cont) + ) + + # Verify the returned tensor. + self.assertEqual(t_view, t_ret) + self.assertFalse(t_ret.is_contiguous()) + + @dist_init + def test_callback_simple(self): + set_by_cb = concurrent.futures.Future() + n = self.rank + 1 + + def callback(fut): + ret = fut.wait() + self.assertEqual(ret, torch.ones(n, n) * 2) + set_by_cb.set_result(ret.clone() + 1) + + fut = rpc.rpc_async( + worker_name(n % self.world_size), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + + fut.then(callback) + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + self.assertEqual(set_by_cb.result(), torch.ones(n, n) * 2 + 1) + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + @dist_init + def test_callback_wrong_arg_num(self): + n = self.rank + 1 + + fut = rpc.rpc_async( + worker_name(n % self.world_size), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + + cb_fut = fut.then(my_function) + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + with self.assertRaisesRegex( + RuntimeError, "my\\_function\\(\\) missing 2 required positional arguments" + ): + cb_fut.wait() + + @dist_init + def test_callback_wrong_arg_type(self): + dst = worker_name((self.rank + 1) % self.world_size) + + fut0 = rpc.rpc_async(dst, torch.add, args=(torch.ones(2, 2), 1)) + fut1 = fut0.then(lambda x: x + 1) + + with self.assertRaisesRegex( + RuntimeError, "unsupported operand type\\(s\\) for \\+" + ): + fut1.wait() + + @dist_init + def test_callback_multi(self): + num_cbs = 10 + n = self.rank + 1 + + def callback(idx, fut): + ret = fut.wait() + self.assertEqual(ret, torch.ones(n, n) * 2) + return ret + idx + + fut = rpc.rpc_async( + worker_name(n % self.world_size), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + + cb_futs = [fut.then(partial(callback, idx)) for idx in range(num_cbs)] + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + for idx in range(num_cbs): + self.assertEqual(cb_futs[idx].wait(), torch.ones(n, n) * 2 + idx) + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + @dist_init + def test_callback_chain(self): + n = self.rank + 1 + + def callback(fut): + return fut.wait() + 1 + + fut = rpc.rpc_async( + worker_name(n % self.world_size), torch.add, args=(torch.ones(n, n), 1) + ) + + num_cbs = 20 + for _ in range(num_cbs): + fut = fut.then(callback) + + self.assertEqual(fut.wait(), torch.ones(n, n) + 1 + num_cbs) + + @dist_init + def test_callback_in_rpc(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + ret = rpc.rpc_sync(dst1, add_use_future_cb, args=(dst2, torch.ones(2, 2), 1, 2)) + self.assertEqual(ret, torch.ones(2, 2) + 1 + 2) + + @dist_init + def test_callback_with_ret(self): + dst = worker_name((self.rank + 1) % self.world_size) + + def callback(fut0): + fut2 = rpc.rpc_async(dst, torch.add, args=(fut0.wait(), 1)).then( + lambda fut1: fut1.wait() + 1 + ) + + return fut2.wait() + + fut3 = rpc.rpc_async(dst, torch.add, args=(torch.ones(2, 2), 1)).then(callback) + + self.assertEqual(fut3.wait(), torch.ones(2, 2) + 3) + + @dist_init + def test_callback_with_error(self): + dst = worker_name((self.rank + 1) % self.world_size) + + def callback(fut0): + with self.assertRaisesRegex(ValueError, "Expected error"): + fut0.wait() + raise RuntimeError("Another expected error") + + fut1 = rpc.rpc_async(dst, raise_func).then(callback) + with self.assertRaisesRegex(RuntimeError, "Another expected error"): + fut1.wait() + + @dist_init + def test_callback_none(self): + dst = worker_name((self.rank + 1) % self.world_size) + with self.assertRaisesRegex(TypeError, "incompatible function arguments."): + rpc.rpc_async(dst, raise_func).then(None) + + @dist_init + def test_add_done_callback(self): + set_by_cb = False + n = self.rank + 1 + + def callback(fut): + nonlocal set_by_cb + fut.wait() + set_by_cb = True + + fut = rpc.rpc_async( + worker_name(n % self.world_size), + torch.add, + args=(torch.ones(n, n), torch.ones(n, n)), + ) + + fut.add_done_callback(callback) + fut_then = fut.then(lambda _: True) + + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + # We have no guarantee that the add_done_callback fn will execute before the test finishes. + # Adding a 'then' callback that runs afterwards to guarantee we wait for the first callback + fut_then.wait() + self.assertTrue(set_by_cb) + self.assertEqual(fut.wait(), torch.ones(n, n) * 2) + + @dist_init + def test_mark_future_twice(self): + fut = rpc.rpc_async( + worker_name((self.rank + 1) % self.world_size), + torch.add, + args=(torch.zeros(2, 2), 1), + ) + self.assertEqual(fut.wait(), torch.zeros(2, 2) + 1) + with self.assertRaisesRegex( + RuntimeError, "Future can only be marked completed once" + ): + fut.set_result(1) + + @dist_init + def test_pickle_future(self): + fut = torch.futures.Future() + errMsg = "Can not pickle torch.futures.Future" + + dst = worker_name((self.rank + 1) % self.world_size) + with TemporaryFileName(), self.assertRaisesRegex(RuntimeError, errMsg): + rpc.rpc_sync(dst, fail_on_fut, args=(fut,)) + + with TemporaryFileName(), self.assertRaisesRegex(RuntimeError, errMsg): + rpc.rpc_async(dst, fail_on_fut, args=(fut,)) + + with TemporaryFileName(), self.assertRaisesRegex(RuntimeError, errMsg): + rpc.remote(dst, fail_on_fut, args=(fut,)) + + @dist_init + def test_future_done(self): + dst = worker_name((self.rank + 1) % self.world_size) + fut = rpc.rpc_async(dst, torch.add, args=(torch.zeros(2), 1)) + fut.wait() + self.assertTrue(fut.done()) + + @dist_init + def test_future_done_exception(self): + dst = worker_name((self.rank + 1) % self.world_size) + fut = rpc.rpc_async(dst, raise_func) + with self.assertRaisesRegex(ValueError, "Expected error"): + fut.wait() + self.assertTrue(fut.done()) + + def _test_future_cb(self, func): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + ret = rpc.rpc_sync(dst1, func, args=(dst2, torch.ones(2, 2), 1, 2)) + self.assertEqual(ret, torch.ones(2, 2) + 1 + 2) + + @dist_init + def test_future_in_rpc(self): + self._test_future_cb(add_use_future_set_result) + + @dist_init + def test_future_nested_callback(self): + self._test_future_cb(add_use_future_nested_cb) + + def _test_async_function_raise(self, mode): + with self.assertRaisesRegex(RuntimeError, "Expected error"): + self._run_func_in_mode( + worker_name((self.rank + 1) % self.world_size), async_raise_func, mode + ) + + @dist_init + def test_async_function_raise(self): + self._test_async_function_raise(RPCExecMode.SYNC) + + @dist_init + def test_async_function_raise_async(self): + self._test_async_function_raise(RPCExecMode.ASYNC) + + @dist_init + def test_async_function_raise_remote(self): + self._test_async_function_raise(RPCExecMode.REMOTE) + + def _test_async_function_wrong_return_type(self, mode): + errMsg = ( + "Functions decorated with @rpc\\.async_function must return a " + "torch\\.futures\\.Future object," + ) + with self.assertRaisesRegex(RuntimeError, errMsg): + self._run_func_in_mode( + worker_name((self.rank + 1) % self.world_size), async_wrong_type, mode + ) + + @dist_init + def test_async_function_wrong_return_type(self): + self._test_async_function_wrong_return_type(RPCExecMode.SYNC) + + @dist_init + def test_async_function_wrong_return_type_async(self): + self._test_async_function_wrong_return_type(RPCExecMode.ASYNC) + + @dist_init + def test_async_function_wrong_return_type_remote(self): + self._test_async_function_wrong_return_type(RPCExecMode.REMOTE) + + @dist_init + def test_async_function_simple(self): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + ret = rpc.rpc_sync(dst1, async_add, args=(dst2, torch.ones(2, 2), 1)) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + def _test_async_function(self, fn, mode=RPCExecMode.SYNC): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + args = (dst2, torch.ones(2, 2), 1, 2) + ret = self._run_func_in_mode(dst1, fn, mode, args=args) + self.assertEqual(ret, torch.ones(2, 2) + 3) + + @dist_init + def test_async_function_with_future_ctor(self): + self._test_async_function(async_add_with_future_ctor) + + @dist_init + def test_async_function_with_future_ctor_remote(self): + self._test_async_function(async_add_with_future_ctor, RPCExecMode.REMOTE) + + @dist_init + def test_async_function_chained(self): + self._test_async_function(async_add_chained) + + @dist_init + def test_async_function_chained_remote(self): + self._test_async_function(async_add_chained, RPCExecMode.REMOTE) + + @dist_init + def test_async_function_nested(self): + self._test_async_function(async_add_nested) + + @dist_init + def test_async_function_nested_remote(self): + self._test_async_function(async_add_nested, RPCExecMode.REMOTE) + + @dist_init + def test_async_static_method(self): + self._test_async_function(AsyncExecutionClass.static_async_add) + + @dist_init + def test_async_static_method_remote(self): + self._test_async_function( + AsyncExecutionClass.static_async_add, RPCExecMode.REMOTE + ) + + @dist_init + def test_async_class_method(self): + self._test_async_function(AsyncExecutionClass.class_async_add) + + @dist_init + def test_async_class_method_remote(self): + self._test_async_function( + AsyncExecutionClass.class_async_add, RPCExecMode.REMOTE + ) + + def _test_test_async_class_rref_proxy(self, mode=RPCExecMode.SYNC): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + rref = rpc.remote(dst1, AsyncExecutionClass) + + x = torch.ones(2, 2) + y = torch.ones(2, 2) + 1 + if mode == RPCExecMode.SYNC: + ret = rref.rpc_sync().static_async_add(dst2, x, x, y) + ret += rref.rpc_sync().class_async_add(dst2, x, x, y) + ret += rref.rpc_sync().bound_async_add(dst2, x, x, y) + elif mode == RPCExecMode.ASYNC: + ret = rref.rpc_async().static_async_add(dst2, x, x, y).wait() + ret += rref.rpc_async().class_async_add(dst2, x, x, y).wait() + ret += rref.rpc_async().bound_async_add(dst2, x, x, y).wait() + elif mode == RPCExecMode.REMOTE: + ret = rref.remote().static_async_add(dst2, x, x, y).to_here() + ret += rref.remote().class_async_add(dst2, x, x, y).to_here() + ret += rref.remote().bound_async_add(dst2, x, x, y).to_here() + + self.assertEqual(ret, 3 * 4 * x) + + @dist_init + def test_async_class_rref_proxy(self): + self._test_test_async_class_rref_proxy() + + @dist_init + def test_async_class_rref_proxy_async(self): + self._test_test_async_class_rref_proxy(mode=RPCExecMode.ASYNC) + + @dist_init + def test_async_class_rref_proxy_remote(self): + self._test_test_async_class_rref_proxy(mode=RPCExecMode.REMOTE) + + def _test_async_function_multi(self, fn, mode=RPCExecMode.SYNC): + dst1 = worker_name((self.rank + 1) % self.world_size) + dst2 = worker_name((self.rank + 2) % self.world_size) + + num = 20 + step = 3 + args = (dst2, torch.ones(2, 2), num, step) + ret = self._run_func_in_mode(dst1, fn, mode, args=args) + self.assertEqual(ret, torch.ones(2, 2) + num * step) + + @dist_init + def test_async_function_multi_chained(self): + self._test_async_function_multi(async_add_chained_multi) + + @dist_init + def test_async_function_multi_chained_async(self): + self._test_async_function_multi(async_add_chained_multi, RPCExecMode.ASYNC) + + @dist_init + def test_async_function_multi_chained_remote(self): + self._test_async_function_multi(async_add_chained_multi, RPCExecMode.REMOTE) + + @dist_init + def test_async_function_multi_fanout(self): + self._test_async_function_multi(async_add_multi_fanout) + + @dist_init + def test_async_function_multi_fanout_async(self): + self._test_async_function_multi(async_add_multi_fanout, RPCExecMode.ASYNC) + + @dist_init + def test_async_function_multi_fanout_remote(self): + self._test_async_function_multi(async_add_multi_fanout, RPCExecMode.REMOTE) + + def _test_return_future(self, mode): + with self.assertRaisesRegex( + RuntimeError, "Can not pickle torch.futures.Future" + ): + self._run_func_in_mode( + worker_name((self.rank + 1) % self.world_size), return_future, mode + ) + + @dist_init + def test_return_future(self): + self._test_return_future(RPCExecMode.SYNC) + + @dist_init + def test_return_future_async(self): + self._test_return_future(RPCExecMode.ASYNC) + + @dist_init + def test_return_future_remote(self): + self._test_return_future(RPCExecMode.REMOTE) + + @dist_init + def test_rref_timeout(self): + # This test is similar to ones in FaultyProcessGroupTest, but is meant to be + # run with other backends besides ProcessGroup. + if self.rank != 0: + return + + dst_rank = (self.rank + 1) % self.world_size + dst_worker = f"worker{dst_rank}" + # 10 ms timeout + rref = rpc.remote(dst_worker, my_sleep_func, args=(2,), timeout=0.01) + # Future corresponding to the remote creation should time out. + expected_error = self.get_timeout_error_regex() + with self.assertRaisesRegex(RuntimeError, expected_error): + rref._get_future().wait() + # Call to ensure pending callbacks are run. + wait_until_pending_futures_and_users_flushed() + with self.assertRaisesRegex(RuntimeError, "RRef creation"): + rref.to_here() + + wait_until_owners_and_forks_on_rank(1, 1, rank=1) + + @dist_init(setup_rpc=False) + @skip_but_pass_in_sandcastle_if( + os.environ.get("RPC_INIT_WITH_TCP", None) == "1", + "init_pg_then_rpc does not work with TCP init, see https://github.com/pytorch/pytorch/issues/41614.", + ) + def test_init_pg_then_rpc(self): + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + # Test RPC. + next_rank = (self.rank + 1) % self.world_size + ret = rpc.rpc_sync( + worker_name(next_rank), torch.add, args=(torch.ones(2, 2), 1) + ) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + # Test PG + dist.barrier() + + rpc.shutdown() + + @dist_init(setup_rpc=False) + @skip_but_pass_in_sandcastle_if( + os.environ.get("RPC_INIT_WITH_TCP", None) == "1", + "init_rpc_then_pg does not work with TCP init, see https://github.com/pytorch/pytorch/issues/41614.", + ) + def test_init_rpc_then_pg(self): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + dist.init_process_group( + backend="gloo", + init_method=self.init_method, + rank=self.rank, + world_size=self.world_size, + ) + + # Test RPC. + next_rank = (self.rank + 1) % self.world_size + ret = rpc.rpc_sync( + worker_name(next_rank), torch.add, args=(torch.ones(2, 2), 1) + ) + self.assertEqual(ret, torch.ones(2, 2) + 1) + + # Test PG + dist.barrier() + + rpc.shutdown() + + @dist_init + def test_wait_all_with_exception(self): + dst = worker_name((self.rank + 1) % self.world_size) + futs = [rpc.rpc_async(dst, raise_func) for _ in range(10)] + + with self.assertRaisesRegex(ValueError, "Expected error"): + torch.futures.wait_all(futs) + + @dist_init + def test_wait_all_with_partial_exception(self): + dst = worker_name((self.rank + 1) % self.world_size) + futs = [ + rpc.rpc_async(dst, torch.add, args=(torch.ones(2), 1)) for _ in range(10) + ] + + futs.append(rpc.rpc_async(dst, raise_func)) + + with self.assertRaisesRegex(ValueError, "Expected error"): + torch.futures.wait_all(futs) + + @dist_init(setup_rpc=False) + @skip_but_pass_in_sandcastle_if( + os.environ.get("RPC_INIT_WITH_TCP", None) == "1", + "Test does not work with TCP init, see https://github.com/pytorch/pytorch/issues/46491", + ) + def test_init_rpc_twice(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + rpc.shutdown() + + # Wait for all init to complete. + dist.barrier() + + # Use a different file name for the next initialization + new_backend_options = self.rpc_backend_options + new_backend_options.init_method += "init_2" + + # Ensure rpc initialization works again. + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=new_backend_options, + ) + + # Verify RPCs work after re-init. + dst = worker_name((self.rank + 1) % self.world_size) + rpc.rpc_sync(dst, torch.add, args=(torch.ones(2, 2), 1)) + rpc.rpc_sync(dst, foo_add, args=()) + + rpc.shutdown() + + def test_wrong_types(self): + with self.assertRaisesRegex( + TypeError, + "Argument backend must be a member of BackendType", + ): + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + backend="TENSORPIPE", + ) + + with self.assertRaisesRegex( + TypeError, + "Argument rpc_backend_options must be an instance of RpcBackendOptions", + ): + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + backend=self.rpc_backend, + rpc_backend_options={"init_method": self.init_method}, + ) + + def test_cannot_infer_backend_from_options(self): + # An exception should be raised if the backend isn't specified but + # options are given which are not an instance of any of the known + # agents' option classes. + rpc_backend_options = FooBackendOptions(self.init_method) + + with self.assertRaisesRegex(TypeError, "Could not infer backend for options"): + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + # Do _not_ pass backend. + rpc_backend_options=rpc_backend_options, + ) + + @dist_init + def test_owner_rref_backward(self): + dst = worker_name((self.rank + 1) % self.world_size) + t1 = torch.rand(10, 10, requires_grad=True) + rref = rpc.RRef(t1.sum() + t1.sum()) + rref.backward() + expected_grad = torch.ones_like(t1) * 2 + self.assertEqual(expected_grad, t1.grad) + + with dist_autograd.context() as context_id: + t2 = rpc.rpc_sync(dst, torch.add, args=(t1, t1)) + rref = rpc.RRef(t2.sum()) + rref.backward(context_id) + self.assertEqual(expected_grad, dist_autograd.get_gradients(context_id)[t1]) + + # Double backward. + with dist_autograd.context() as context_id: + t2 = rpc.rpc_sync(dst, torch.add, args=(t1, t1)) + rref = rpc.RRef(t2.sum()) + rref.backward(context_id, retain_graph=True) + rref.backward(context_id) + self.assertEqual( + expected_grad * 2, dist_autograd.get_gradients(context_id)[t1] + ) + + # Test errors. + with self.assertRaisesRegex( + RuntimeError, "tensors does not require grad and does not have a grad_fn" + ): + rpc.RRef(torch.rand(10)).backward() + + with self.assertRaisesRegex( + RuntimeError, "grad can be implicitly created only for scalar outputs" + ): + rpc.RRef(torch.rand(10, requires_grad=True)).backward() + + with self.assertRaisesRegex( + RuntimeError, "Could not find autograd context with id: 100" + ): + rpc.RRef(torch.rand(10, requires_grad=True).sum()).backward(100) + + with self.assertRaisesRegex( + RuntimeError, "RRef should contain a tensor for .backward()" + ): + rpc.RRef("foo").backward() + + @staticmethod + def _sum(x): + return x.sum() + + @staticmethod + def _identity(x): + return x + + @dist_init + def test_user_rref_backward(self): + dst = worker_name((self.rank + 1) % self.world_size) + t = torch.rand(10, requires_grad=True) + with dist_autograd.context() as context_id: + rref = rpc.remote(dst, RpcTest._sum, args=(t,)) + rref.backward(context_id, retain_graph=True) + rref.backward(context_id) + self.assertEqual( + torch.ones_like(t) * 2, dist_autograd.get_gradients(context_id)[t] + ) + + with dist_autograd.context() as context_id: + rref = rpc.remote(dst, RpcTest._identity, args=("foo",)) + with self.assertRaisesRegex( + RuntimeError, "RRef should contain a tensor for .backward()" + ): + rref.backward(context_id) + + with self.assertRaisesRegex( + RuntimeError, + "User RRefs require 'dist_autograd_ctx_id' to be specified", + ): + rref.backward() + + @dist_init(setup_rpc=False) + def test_shutdown_errors(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + if self.rank != 0: + og_func = rpc.api._broadcast_to_followers + og_rref_func = rpc.api._delete_all_user_and_unforked_owner_rrefs + + # Monkey-patch _broadcast_to_followers to fail, which would ensure + # _all_gather on leader raises an exception. + def raise_error(sequence_id, objects_map): + og_func(sequence_id, objects_map) + raise RuntimeError("simulation") + + # Monkey-patch _delete_all_user_and_unforked_owner_rrefs to fail, + # which would ensure barrier is not called on followers. + def rref_error(): + raise RuntimeError("simulation rref") + + try: + rpc.api._broadcast_to_followers = raise_error + rpc.api._delete_all_user_and_unforked_owner_rrefs = rref_error + with self.assertRaisesRegex(RuntimeError, "simulation rref"): + rpc.shutdown() + finally: + rpc.api._broadcast_to_followers = og_func + rpc.api._delete_all_user_and_unforked_owner_rrefs = og_rref_func + else: + with self.assertRaisesRegex(RuntimeError, "timed out in _all_gather"): + rpc.shutdown() + + dist.barrier() + + @dist_init + def test_my_parameter_server(self): + self._my_parameter_server(False) + + +class CudaRpcTest(RpcAgentTestFixture): + @skip_if_lt_x_gpu(2) + @dist_init + def test_profiler_remote_cuda(self): + if self.rank != 1: + return + + dst_cuda_0 = (self.rank + 1) % self.world_size + dst_cuda_1 = (self.rank + 2) % self.world_size + dst_worker_cuda_0 = worker_name(dst_cuda_0) + dst_worker_cuda_1 = worker_name(dst_cuda_1) + + with _profile(use_cuda=True) as p: + fut1 = rpc.rpc_async(dst_worker_cuda_0, udf_with_torch_ops, args=(0,)) + fut2 = rpc.rpc_async(dst_worker_cuda_1, udf_with_torch_ops, args=(1,)) + fut1.wait() + fut2.wait() + + def get_name(event): + return event.name[event.name.find(REMOTE_OP_STR) + len(REMOTE_OP_STR) :] + + function_events = p.function_events + for event in function_events: + if event.is_async: + self.assertEqual(0, event.device_time_total) + self.assertEqual([], event.kernels) + self.assertEqual(0, event.device_time) + else: + if event.node_id == 1: + continue + self.assertTrue(event.node_id in [dst_cuda_0, dst_cuda_1]) + if get_name(event) in EXPECTED_REMOTE_EVENTS: + self.assertGreater(event.device_time_total, 0) + self.assertEqual(1, len(event.kernels)) + kernel = event.kernels[0] + if event.node_id == dst_cuda_0: + self.assertEqual(kernel.device, 0) + if event.node_id == dst_cuda_1: + self.assertEqual(kernel.device, 1) + self.assertGreater(event.device_time, 0) + + # Validate that EXPECTED_REMOTE_EVENTS is a subset of remotely profiled + # events. + remote_events = [event for event in function_events if event.is_remote] + remote_event_names = [ + get_name(event) + for event in remote_events + if get_name(event) in EXPECTED_REMOTE_EVENTS + ] + self.assertEqual(set(remote_event_names), set(EXPECTED_REMOTE_EVENTS)) + + +class TensorPipeAgentRpcTest(RpcAgentTestFixture, RpcTestCommon): + def test_mismatched_type_for_options(self): + # An exception should be raised if the options are not an instance of + # TensorPipeRpcBackendOptions. + rpc_backend_options = FooBackendOptions(self.init_method) + + with self.assertRaisesRegex( + TypeError, "`rpc_backend_options` must be a `TensorPipeRpcBackendOptions`" + ): + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + backend=rpc.BackendType.TENSORPIPE, + rpc_backend_options=rpc_backend_options, + ) + + def test_infer_backend_from_options(self): + rpc_backend_options = rpc.TensorPipeRpcBackendOptions( + init_method=self.init_method, _transports=tp_transports() + ) + + rpc.init_rpc( + name=worker_name(self.rank), + rank=self.rank, + world_size=self.world_size, + # Do _not_ pass backend. + rpc_backend_options=rpc_backend_options, + ) + + self.assertIsInstance(rpc.api._get_current_rpc_agent(), rpc.TensorPipeAgent) + + # FIXME Merge this test with the corresponding one in RpcTest. + @dist_init(setup_rpc=False) + def test_set_and_get_num_worker_threads(self): + NUM_THREADS = 27 + rpc_backend_options = rpc.TensorPipeRpcBackendOptions( + init_method=self.rpc_backend_options.init_method, + num_worker_threads=NUM_THREADS, + _transports=tp_transports(), + ) + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + + info = rpc.api._get_current_rpc_agent().get_debug_info() + self.assertEqual(int(info["agent.thread_pool_size"]), NUM_THREADS) + rpc.shutdown() + + # FIXME Merge this test with the corresponding one in RpcTest. + @dist_init(setup_rpc=False) + def test_tensorpipe_set_default_timeout(self): + # Set a high timeout since it doesn't affect test runtime and ensures + # the test doesn't erroneously timeout due to slow machines. + timeout = 100 + rpc_backend_options = rpc.TensorPipeRpcBackendOptions( + init_method=self.rpc_backend_options.init_method, + num_worker_threads=self.rpc_backend_options.num_worker_threads, + rpc_timeout=timeout, + _transports=tp_transports(), + ) + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc_backend_options, + ) + + default_timeout = rpc.get_rpc_timeout() + self.assertEqual(default_timeout, timeout) + rpc.shutdown() + + # FIXME Merge this test with the corresponding one in RpcTest. + @dist_init(setup_rpc=False) + def test_tensorpipe_options_throw_on_timedelta_timeout(self): + from datetime import timedelta + + timeout = timedelta() + # Ensure that constructing TensorPipeRpcBackendOptions with timedelta fails + with self.assertRaisesRegex(TypeError, "incompatible constructor arguments"): + rpc.TensorPipeRpcBackendOptions( + init_method=self.rpc_backend_options.init_method, + num_worker_threads=self.rpc_backend_options.num_worker_threads, + rpc_timeout=timeout, + ) + + @dist_init + def _test_rref_get_type_timeout(self, blocking): + # Test where we try to get the type of a RRef from an owner, but RRef + # creation is slower than timeout passed into _get_type. + dst_rank = (self.rank + 1) % self.world_size + dst = worker_name(dst_rank) + slow_rref = rpc.remote(dst, MyClass, args=(torch.ones(2, 2), True)) + timeout = 0.5 + expected_err = self.get_timeout_error_regex() + # Blocking: blocks on inline call + if blocking: + with self.assertRaisesRegex(RuntimeError, expected_err): + slow_rref._get_type(timeout=timeout, blocking=blocking) + # Non-blocking: blocks on wait + else: + fut = slow_rref._get_type(timeout=timeout, blocking=blocking) + with self.assertRaisesRegex(RuntimeError, expected_err): + fut.wait() + + # FIXME We wait until the remote completed creating the OwnerRRef + # because there's currently a race if we shut down RPC before that. + slow_rref.to_here() + + def test_rref_get_type_timeout_blocking(self): + self._test_rref_get_type_timeout(blocking=True) + + def test_rref_get_type_timeout_non_blocking(self): + self._test_rref_get_type_timeout(blocking=False) + + @dist_init + def test_op_with_invalid_args(self): + dst = worker_name((self.rank + 1) % self.world_size) + with self.assertRaisesRegex( + RuntimeError, + "Overloaded torch operator invoked from Python failed to match any schema", + ): + rpc.rpc_sync(dst, torch.add, args=()) + + def _test_rref_proxy_timeout(self, rref_proxy_api): + dst_rank = (self.rank + 1) % self.world_size + dst = worker_name(dst_rank) + rref = rpc.remote(dst, MyClass, args=(torch.ones(2, 2),)) + # Ensure RRef is created on remote node. + rref.to_here() + rref_api = getattr(rref, rref_proxy_api) + self.assertTrue( + rref_api is not None, f"Failed to get RRef proxy api: {rref_proxy_api}" + ) + expected_error = self.get_timeout_error_regex() + timeout = 2 + with self.assertRaisesRegex(RuntimeError, expected_error): + result = rref_api(timeout=timeout).my_slow_method(torch.ones(2, 2)) + if rref_api == rref.rpc_async: + result.wait() + elif rref_api == rref.remote: + result._get_future().wait() + + # Case where rpc.remote() is stuck and exceeds timeout + slow_rref = rpc.remote(dst, MyClass, args=(torch.ones(2, 2), True)) + timeout = 0.01 + rref_api = getattr(slow_rref, rref_proxy_api) + # Note that even when we call rref.rpc_async() in this case, we + # time out in future creation, not waiting for future. This is because + # rref proxy function calls rref._get_type before returning future, + # which blocks on the RRef being created on owner node, until the + # specified timeout. + with self.assertRaisesRegex(RuntimeError, expected_error): + result = rref_api(timeout=timeout).my_instance_method(torch.ones(2, 2)) + # rpc_async returns immediately and surface a timeout through wait() + if rref_api == slow_rref.rpc_async: + result.wait() + + # FIXME We wait until the remote completed creating the OwnerRRef + # because there's currently a race if we shut down RPC before that. + slow_rref.to_here() + + @dist_init + def test_rref_proxy_timeout(self): + for rpc_api in ["rpc_sync", "rpc_async", "remote"]: + self._test_rref_proxy_timeout(rpc_api) + + @dist_init + def test_send_to_rank_sparse(self): + dst_rank = (self.rank + 1) % self.world_size + + # Test sparse tensor + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + x = build_sparse_tensor() + y = build_sparse_tensor() + expected_tensor = x + y + ret = self._run_func_in_mode(dst_rank, torch.add, exec_mode, args=(x, y)) + self.assertEqual(expected_tensor, ret) + + for exec_mode in [RPCExecMode.SYNC, RPCExecMode.ASYNC, RPCExecMode.REMOTE]: + x = build_sparse_tensor(coalesce=True) + y = build_sparse_tensor(coalesce=True) + expected_tensor = x + y + ret = self._run_func_in_mode(dst_rank, torch.add, exec_mode, args=(x, y)) + self.assertEqual(expected_tensor, ret) + + @dist_init + def test_self_py_udf_remote_sparse(self): + self._self_py_udf_remote( + rpc.get_worker_info(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + ) + + @dist_init + def test_self_remote_rref_as_rpc_arg_sparse(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._self_remote_rref_as_rpc_arg( + dst, build_sparse_tensor(), build_sparse_tensor(), build_sparse_tensor() + ) + + @dist_init + def test_self_remote_rref_as_self_rpc_arg_sparse(self): + self._self_remote_rref_as_rpc_arg( + rpc.get_worker_info(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + ) + + @dist_init + def test_self_remote_rref_as_remote_arg_sparse(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._self_remote_rref_as_remote_arg( + dst, build_sparse_tensor(), build_sparse_tensor(), build_sparse_tensor() + ) + + @dist_init + def test_self_remote_rref_as_self_remote_arg_sparse(self): + self._self_remote_rref_as_remote_arg( + rpc.get_worker_info(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + ) + + def test_world_size_one_sparse(self): + self._world_size_one(build_sparse_tensor(), build_sparse_tensor()) + + @dist_init + def test_multi_rpc_sparse(self): + self._multi_rpc(True) + + def test_wait_all_workers_sparse(self): + self._wait_all_workers(heavy_rpc_sparse, build_sparse_tensor()) + + def test_wait_all_workers_twice_sparse(self): + self._wait_all_workers_twice(heavy_rpc_sparse, build_sparse_tensor()) + + @dist_init + def test_py_sparse_tensors_in_container(self): + n = self.rank + 1 + dst_rank = n % self.world_size + a = [build_sparse_tensor(), build_sparse_tensor()] + ret = rpc.rpc_sync(worker_name(dst_rank), my_container_sum, args=(a,)) + self.assertEqual(ret, my_container_sum(a)) + + @dist_init + def test_nested_rpc_sparse(self): + self._nested_rpc(nested_rpc_sparse, build_sparse_tensor() * 2) + + @dist_init + def test_stress_heavy_rpc_sparse(self): + self._stress_test_rpc( + heavy_rpc_sparse, repeat=20, args=(build_sparse_tensor(),) + ) + + @dist_init + def test_builtin_remote_ret_sparse(self): + self._builtin_remote_ret( + build_sparse_tensor(), build_sparse_tensor(), build_sparse_tensor() * 2 + ) + + @dist_init + def test_builtin_remote_self_sparse(self): + self._builtin_remote_self( + build_sparse_tensor(), build_sparse_tensor(), build_sparse_tensor() * 2 + ) + + @dist_init + def test_multi_builtin_remote_ret_sparse(self): + self._test_multi_remote_call(torch.add, True, args_fn=RpcTest._multi_args_fn) + + @dist_init + def test_multi_py_udf_remote_sparse(self): + self._test_multi_remote_call( + my_function, True, kwargs_fn=RpcTest._multi_kwargs_fn + ) + + @dist_init + def test_py_rref_args_sparse(self): + self._py_rref_args( + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() * 4, + ) + + @dist_init + def test_py_rref_args_user_share_sparse(self): + self._py_rref_args_user_share( + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() * 6, + ) + + @dist_init + def test_py_rpc_rref_args_sparse(self): + self._py_rpc_rref_args( + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor(), + build_sparse_tensor() * 6, + ) + + @dist_init + def test_nested_remote_sparse(self): + self._nested_remote( + nested_remote_sparse, build_sparse_tensor() + build_sparse_tensor() + ) + + @dist_init + def test_nested_rref_sparse(self): + self._nested_rref( + nested_rref_sparse, build_sparse_tensor() * 2, build_sparse_tensor() * 2 + ) + + @dist_init + def test_nested_rref_stress_sparse(self): + self._nested_rref_stress( + nested_rref_sparse, build_sparse_tensor() * 2, build_sparse_tensor() * 2 + ) + + @dist_init + def test_my_parameter_server_sparse(self): + self._my_parameter_server(True) + + # Test init_rpc without world_size argument + @dist_init(setup_rpc=False) + def test_dynamic_rpc_init_rpc(self): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + rpc.shutdown() + + # Dynamic RPC new ranks communicate with existing ranks + @dist_init(setup_rpc=False) + def test_dynamic_rpc_new_rank_can_communicated_with_existing_rank(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + if self.rank == 0: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + + # Rank 0 will be initialized with RPC after this barrier + dist.barrier() + + if self.rank != 0: + # Newly joined ranks will be able to communicate with rank 0, since that was created first + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + result = rpc.rpc_sync( + worker_name(0), torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + self.assertEqual(torch.add(torch.tensor(1), torch.tensor(1)), result) + + # Barrier to ensure that all rpc_sync calls are finished + dist.barrier() + rpc.shutdown() + + # Dynamic RPC existing ranks can communicate with new ranks + @dist_init(setup_rpc=False) + def test_dynamic_rpc_existing_rank_can_communicate_with_new_rank(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + if self.rank == 0: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + + # Rank 0 will be initialized with RPC after this barrier + dist.barrier() + + # Rest of ranks join after barrier + if self.rank != 0: + # Newly joined ranks will be able to communicate with rank 0, since that was created first + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + + dist.barrier() + if self.rank == 0: + for i in range(1, self.world_size): + result = rpc.rpc_sync( + worker_name(i), torch.add, args=(torch.tensor(1), torch.tensor(1)) + ) + self.assertEqual(torch.add(torch.tensor(1), torch.tensor(1)), result) + + # Barrier to ensure that all rpc_sync calls are finished + dist.barrier() + rpc.shutdown() + + # Dynamic RPC existing ranks can communicate with new ranks using CUDA rpc + @skip_if_lt_x_gpu(2) + @dist_init(setup_rpc=False) + def test_dynamic_rpc_existing_rank_can_communicate_with_new_rank_cuda(self): + initialize_pg(self.file_init_method, self.rank, self.world_size) + + if self.rank == 0: + options = self.rpc_backend_options + for i in range(1, self.world_size): + dst = worker_name(i) + options.set_device_map(dst, {1: 0}) + options.set_device_map(dst, {0: 1}) + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=options, + ) + + # Rank 0 will be initialized with RPC after this barrier + dist.barrier() + + # Rest of ranks join after barrier + if self.rank != 0: + # Newly joined ranks will be able to communicate with rank 0, since that was created first + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + + # TODO: Cuda RPC is failing due to: + # terminate called after throwing an instance of 'c10::Error' + # what(): 0 <= device && static_cast(device) < device_allocator.size() + # INTERNAL ASSERT FAILED at "../c10/cuda/CUDACachingAllocator.cpp":1937, + # please report a bug to PyTorch. Allocator not initialized for device 1: did you call init? + # dist.barrier() + # if self.rank == 0: + # for i in range(1, self.world_size): + # x = torch.ones(2) + # result_on_device_0 = rpc.rpc_sync(worker_name(i), torch.add, args=(x.to(0), 1)) + # result_on_device_1 = rpc.rpc_sync(worker_name(i), torch.add, args=(x.to(1), 1)) + # self.assertEqual(torch.add(torch.ones(2), 1), result_on_device_0) + # self.assertEqual(torch.device('cuda:0'), result_on_device_0.device) + # self.assertEqual(torch.add(torch.ones(2), 1), result_on_device_1) + # self.assertEqual(torch.device('cuda:1'), result_on_device_1.device) + + # Barrier to ensure that all rpc_sync calls are finished + dist.barrier() + rpc.shutdown() + + @dist_init(setup_rpc=False) + def test_dynamic_rpc_init_rpc_without_rank(self): + # default initialization uses file init + with self.assertRaisesRegex(ValueError, "rank parameter missing"): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rpc_backend_options=self.rpc_backend_options, + ) + + # env init + with self.assertRaisesRegex(ValueError, "environment variable RANK expected"): + rpc_backend_options = rpc.TensorPipeRpcBackendOptions(init_method="env://") + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rpc_backend_options=rpc_backend_options, + ) + + # tcp init + with self.assertRaisesRegex(ValueError, "rank parameter missing"): + rpc_backend_options = rpc.TensorPipeRpcBackendOptions( + init_method="tcp://127.0.0.1:23456" + ) + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rpc_backend_options=rpc_backend_options, + ) + + @dist_init(setup_rpc=False) + def test_dynamic_and_static_init_rpc_together(self): + # Initialize a static rpc group with size = self.world_size - 1 + dist.init_process_group( + backend="gloo", + init_method=self.file_init_method, + rank=self.rank, + world_size=self.world_size, + ) + + world_size_minus_one = self.world_size - 1 + if self.rank < world_size_minus_one: + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=world_size_minus_one, + rpc_backend_options=self.rpc_backend_options, + ) + + dist.barrier() + + # Attempt to add an additional dynamic group member + if self.rank == world_size_minus_one: + # Expect error message to be thrown + with self.assertRaisesRegex( + RuntimeError, + "RPC group mixes statically and dynamically\ + initialized members which is not supported.", + ): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + rpc_backend_options=self.rpc_backend_options, + ) + + +class TensorPipeAgentCudaRpcTest(RpcAgentTestFixture, RpcTestCommon): + def _test_device_maps(self, options, errMsg): + with self.assertRaisesRegex(ValueError, errMsg): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + self.assertFalse(rpc.api._is_current_rpc_agent_set()) + + @skip_if_lt_x_gpu(2) + def test_device_maps_wrong_worker_name(self): + options = self.rpc_backend_options + options.set_device_map("none_exist", {0: 1}) + + self._test_device_maps( + options, + errMsg="Node worker0 has invalid target node names in its device maps", + ) + + @skip_if_lt_x_gpu(1) + def test_device_maps_invalid_max_local_device(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {torch.cuda.device_count(): 0}) + + self._test_device_maps( + options, + errMsg="Node worker0 has source devices with invalid indices in its device map for worker1", + ) + + @skip_if_lt_x_gpu(1) + def test_device_maps_invalid_max_remote_device(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {0: torch.cuda.device_count()}) + + self._test_device_maps( + options, + errMsg="Node worker0 has target devices with invalid indices in its device map for worker1", + ) + + @skip_if_lt_x_gpu(2) + def test_device_maps_many_to_one(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {1: 0}) + options.set_device_map(dst, {0: 0}) + + self._test_device_maps( + options, + errMsg="Node worker0 has duplicated target devices in its device map for worker1", + ) + + @skip_if_lt_x_gpu(2) + def test_device_maps_one_to_many(self): + if self.rank == 0: + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {0: 1}) + with self.assertRaisesRegex( + ValueError, "`set_device_map` only supports 1-to-1 mapping" + ): + options.set_device_map(dst, {0: 0}) + + @skip_if_lt_x_gpu(1) + def test_device_maps_invalid_min_device(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + with self.assertRaisesRegex(RuntimeError, "Device index must not be negative"): + options.set_device_map(dst, {-1: 0}) + + with self.assertRaisesRegex(RuntimeError, "Device index must not be negative"): + options.set_device_map(dst, {0: -1}) + + @staticmethod + def _gpu_add(x, y): + if all([x.is_cuda, x.device.index == 1, y.is_cuda, y.device.index == 1]): + return (x + y).to(0) + else: + raise ValueError("Wrong device affinity") + + @skip_if_lt_x_gpu(2) + def test_device_maps_gpu(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {0: 1, 1: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + ret = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._gpu_add, + args=(torch.zeros(2).to(0), torch.ones(2).to(0)), + ) + self.assertEqual(ret.device, torch.device(1)) + self.assertEqual(ret, (torch.zeros(2) + torch.ones(2)).to(1)) + rpc.shutdown() + + @staticmethod + def _gpu_add_given_devices(x, y, x_to, y_to, z_to): + x_device = "cpu" if x.device.type == "cpu" else x.device.index + y_device = "cpu" if y.device.type == "cpu" else y.device.index + if x_device == x_to and y_device == y_to: + return x.to(z_to) + y.to(z_to) + else: + raise ValueError("Wrong device affinity") + + def _test_device_maps_gpu( + self, x_from, y_from, z_to, device_map, dst=None, fn=None + ): + fn = TensorPipeAgentCudaRpcTest._gpu_add_given_devices if fn is None else fn + x_to = device_map[x_from] + y_to = device_map[y_from] + + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) if dst is None else dst + options.set_device_map(dst, device_map) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + x = torch.zeros(2).to(x_from) + y = torch.ones(2).to(y_from) + + ret = rpc.rpc_sync(dst, fn, args=(x, y, x_to, y_to, z_to)) + + reverse_device_map = {device_map[k]: k for k in device_map} + z_from = reverse_device_map[z_to] + + ret_device = "cpu" if ret.device.type == "cpu" else ret.device.index + self.assertEqual(ret_device, z_from) + self.assertEqual(ret, torch.ones(2).to(z_from)) + + rpc.shutdown() + + def test_device_map_cpu(self): + self._test_device_maps_gpu( + x_from="cpu", + y_from="cpu", + z_to="cpu", + device_map={"cpu": "cpu"}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(1) + def test_device_map_cpu_to_gpu_default(self): + self._test_device_maps_gpu( + x_from="cpu", + y_from="cpu", + z_to=0, + device_map={"cpu": 0}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_cpu_to_gpu_non_default(self): + self._test_device_maps_gpu( + x_from="cpu", + y_from="cpu", + z_to=1, + device_map={"cpu": 1}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(1) + def test_device_map_gpu_to_cpu_default(self): + self._test_device_maps_gpu( + x_from=0, + y_from=0, + z_to="cpu", + device_map={0: "cpu"}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_to_cpu_non_default(self): + self._test_device_maps_gpu( + x_from=1, + y_from=1, + z_to="cpu", + device_map={1: "cpu"}, + fn=TensorPipeAgentCudaRpcTest._gpu_add_given_devices, + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_default(self): + self._test_device_maps_gpu(x_from=0, y_from=0, z_to=0, device_map={0: 0}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_non_default(self): + self._test_device_maps_gpu(x_from=1, y_from=1, z_to=1, device_map={1: 1}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_default_to_non_default(self): + self._test_device_maps_gpu(x_from=0, y_from=0, z_to=1, device_map={0: 1}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_non_default_to_default(self): + self._test_device_maps_gpu(x_from=1, y_from=1, z_to=0, device_map={1: 0}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_1(self): + self._test_device_maps_gpu(x_from=0, y_from=1, z_to=0, device_map={0: 0, 1: 1}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_2(self): + self._test_device_maps_gpu(x_from=0, y_from=1, z_to=1, device_map={0: 0, 1: 1}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_3(self): + self._test_device_maps_gpu(x_from=1, y_from=0, z_to=0, device_map={0: 0, 1: 1}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_4(self): + self._test_device_maps_gpu(x_from=1, y_from=0, z_to=1, device_map={0: 0, 1: 1}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_5(self): + self._test_device_maps_gpu(x_from=0, y_from=1, z_to=0, device_map={0: 1, 1: 0}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_6(self): + self._test_device_maps_gpu(x_from=0, y_from=1, z_to=1, device_map={0: 1, 1: 0}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_7(self): + self._test_device_maps_gpu(x_from=1, y_from=0, z_to=0, device_map={0: 1, 1: 0}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_8(self): + self._test_device_maps_gpu(x_from=1, y_from=0, z_to=1, device_map={0: 1, 1: 0}) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_1(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=0, + device_map={0: 0, 1: 1}, + dst=worker_name(self.rank), + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_2(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=1, + device_map={0: 0, 1: 1}, + dst=worker_name(self.rank), + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_3(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=0, + device_map={0: 0, 1: 1}, + dst=worker_name(self.rank), + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_4(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=1, + device_map={0: 0, 1: 1}, + dst=worker_name(self.rank), + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_5(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=0, + device_map={0: 1, 1: 0}, + dst=worker_name(self.rank), + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_6(self): + self._test_device_maps_gpu( + x_from=0, + y_from=1, + z_to=1, + device_map={0: 1, 1: 0}, + dst=worker_name(self.rank), + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_7(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=0, + device_map={0: 1, 1: 0}, + dst=worker_name(self.rank), + ) + + @skip_if_lt_x_gpu(2) + def test_device_map_gpu_mixed_self_8(self): + self._test_device_maps_gpu( + x_from=1, + y_from=0, + z_to=1, + device_map={0: 1, 1: 0}, + dst=worker_name(self.rank), + ) + + @staticmethod + def _gpu_add_multi_gpu(x, y): + if all([x.is_cuda, x.device.index == 1, y.is_cuda, y.device.index == 0]): + return x.to(0) + y, x - y.to(1) + else: + raise ValueError("Wrong device affinity") + + def _test_device_maps_multi_gpu(self, dst): + options = self.rpc_backend_options + options.set_device_map(dst, {0: 1}) + options.set_device_map(dst, {1: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + x = torch.zeros(2).to(0) + y = torch.ones(2).to(1) + rets = rpc.rpc_sync( + dst, TensorPipeAgentCudaRpcTest._gpu_add_multi_gpu, args=(x, y) + ) + + self.assertEqual(rets[0].device, torch.device(1)) + self.assertEqual(rets[1].device, torch.device(0)) + self.assertEqual(rets[0], (torch.zeros(2) + torch.ones(2)).to(1)) + self.assertEqual(rets[1], (torch.zeros(2) - torch.ones(2)).to(0)) + rpc.shutdown() + + @skip_if_lt_x_gpu(2) + def test_device_maps_multi_gpu(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._test_device_maps_multi_gpu(dst) + + @skip_if_lt_x_gpu(2) + def test_device_maps_multi_gpu_self(self): + dst = worker_name(self.rank) + self._test_device_maps_multi_gpu(dst) + + @staticmethod + def _gpu_add_return_to_gpu(x, y): + if x.device.type == "cpu" and y.device.type == "cpu": + return (x + y).to(0), (x - y).to(1), (x * y).to(2), (x / y).to(3) + else: + raise ValueError("Wrong device affinity") + + @skip_if_lt_x_gpu(2) + def test_device_maps_in_options(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=rpc.TensorPipeRpcBackendOptions( + init_method=options.init_method, + num_worker_threads=options.num_worker_threads, + device_maps={dst: {0: 1, 1: 0}}, + _transports=tp_transports(), + ), + ) + + rets = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._gpu_add_multi_gpu, + args=(torch.zeros(2).to(0), torch.ones(2).to(1)), + ) + self.assertEqual(rets[0].device, torch.device(1)) + self.assertEqual(rets[1].device, torch.device(0)) + self.assertEqual(rets[0], (torch.zeros(2) + torch.ones(2)).to(1)) + self.assertEqual(rets[1], (torch.zeros(2) - torch.ones(2)).to(0)) + rpc.shutdown() + + def _test_device_maps_return_to_gpu(self, dst): + options = self.rpc_backend_options + + options.set_device_map(dst, {0: 1}) + options.set_device_map(dst, {1: 2}) + options.set_device_map(dst, {2: 3}) + options.set_device_map(dst, {3: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rets = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._gpu_add_return_to_gpu, + args=(torch.zeros(2), torch.ones(2)), + ) + for i in range(len(rets)): + self.assertEqual(rets[i].device, torch.device((3 + i) % 4)) + self.assertEqual(rets[0], (torch.zeros(2) + torch.ones(2)).to(3)) + self.assertEqual(rets[1], (torch.zeros(2) - torch.ones(2)).to(0)) + self.assertEqual(rets[2], (torch.zeros(2) * torch.ones(2)).to(1)) + self.assertEqual(rets[3], (torch.zeros(2) / torch.ones(2)).to(2)) + rpc.shutdown() + + @skip_if_lt_x_gpu(4) + def test_device_maps_return_to_gpu(self): + dst = worker_name((self.rank + 1) % self.world_size) + self._test_device_maps_return_to_gpu(dst) + + @skip_if_lt_x_gpu(4) + def test_device_maps_return_to_gpu_self(self): + dst = worker_name(self.rank) + self._test_device_maps_return_to_gpu(dst) + + @staticmethod + def _add_to_gpu(x, y): + return (x + y).to(0) + + def _test_device_maps_missing_config(self, mode): + dst = worker_name((self.rank + 1) % self.world_size) + errMsg = ( + "TensorPipe RPC backend only supports CPU tensors by default.*" + "`set_device_map` on `TensorPipeRpcBackendOptions`" + ) + + with self.assertRaisesRegex(RuntimeError, errMsg): + if mode == RPCExecMode.SYNC: + rpc.rpc_sync(dst, torch.add, args=(torch.zeros(2).to(0), 1)) + elif mode == RPCExecMode.REMOTE: + rpc.remote(dst, torch.add, args=(torch.zeros(2).to(0), 1)).to_here() + else: + raise ValueError(f"unexpected mode {mode}") + + # make sure RPC is still functioning + ret = rpc.rpc_sync(dst, torch.add, args=(torch.ones(2), 1)) + self.assertEqual(ret, torch.ones(2) + 1) + + def _test_device_maps_missing_config_response(self, mode): + dst = worker_name((self.rank + 1) % self.world_size) + errMsg = "Response device mapping is not available" + + with self.assertRaisesRegex(RuntimeError, errMsg): + if mode == RPCExecMode.SYNC: + rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._add_to_gpu, + args=(torch.zeros(2), 1), + ) + elif mode == RPCExecMode.REMOTE: + rpc.remote( + dst, + TensorPipeAgentCudaRpcTest._add_to_gpu, + args=(torch.zeros(2), 1), + ).to_here() + else: + raise ValueError(f"unexpected mode {mode}") + + # make sure RPC is still functioning + ret = rpc.rpc_sync(dst, torch.add, args=(torch.ones(2), 1)) + self.assertEqual(ret, torch.ones(2) + 1) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config(self): + self._test_device_maps_missing_config(RPCExecMode.SYNC) + + @skip_if_lt_x_gpu(1) + def test_device_maps_missing_config_not_timeout(self): + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=self.rpc_backend_options, + ) + + timeout = rpc.get_rpc_timeout() + + tik = time.time() + self._test_device_maps_missing_config(RPCExecMode.SYNC) + rpc.shutdown() + tok = time.time() + + self.assertTrue(tok - tik < timeout) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_loop(self): + for _ in range(self.rpc_backend_options.num_worker_threads + 5): + self._test_device_maps_missing_config(RPCExecMode.SYNC) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_response(self): + self._test_device_maps_missing_config_response(RPCExecMode.SYNC) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_response_loop(self): + for _ in range(self.rpc_backend_options.num_worker_threads + 5): + self._test_device_maps_missing_config_response(RPCExecMode.SYNC) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_remote(self): + self._test_device_maps_missing_config(RPCExecMode.REMOTE) + + @skip_if_lt_x_gpu(1) + @dist_init + def test_device_maps_missing_config_remote_response(self): + self._test_device_maps_missing_config_response(RPCExecMode.REMOTE) + + @skip_if_lt_x_gpu(2) + def test_device_maps_remote(self): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, {1: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rref = rpc.remote( + dst, TensorPipeAgentCudaRpcTest._add_to_gpu, args=(torch.zeros(2), 1) + ) + + self.assertEqual(rref.to_here().device.index, 1) + self.assertEqual(rref.to_here(), torch.ones(2).to(1)) + + rpc.shutdown() + + @staticmethod + def _slow_add_on_user_stream(x, y): + s0 = torch.cuda.current_stream(x.device) + s1 = torch.cuda.Stream(device=x.device) + s1.wait_stream(s0) + x.record_stream(s1) + y.record_stream(s1) + with torch.cuda.stream(s1): + torch.cuda._sleep(10 * FIFTY_MIL_CYCLES) + z = x + y + s0.wait_stream(s1) + z.record_stream(s0) + return z + + def _test_custom_stream(self, fn, device_map): + options = self.rpc_backend_options + dst = worker_name((self.rank + 1) % self.world_size) + options.set_device_map(dst, device_map) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + fn(dst) + + rpc.shutdown() + + def _test_stream_sync(self, dst): + x = torch.ones(2, 2).to(0) + ret = rpc.rpc_sync( + dst, TensorPipeAgentCudaRpcTest._slow_add_on_user_stream, args=(x, x) + ) + self.assertEqual(ret, 2 * x) + + @skip_if_lt_x_gpu(2) + def test_custom_stream(self): + self._test_custom_stream(self._test_stream_sync, {"cuda:0": "cuda:1"}) + + def _test_stream_multi_async(self, dst): + futs = [] + for i in range(20): + x = torch.ones(2, 2).to(0) * i + futs.append( + rpc.rpc_async( + dst, + TensorPipeAgentCudaRpcTest._slow_add_on_user_stream, + args=(x, x), + ) + ) + + for i in range(20): + self.assertEqual(futs[i].wait(), 2 * torch.ones(2, 2).to(0) * i) + + @skip_if_lt_x_gpu(2) + def test_custom_stream_multi(self): + self._test_custom_stream(self._test_stream_multi_async, {"cuda:0": "cuda:1"}) + + @staticmethod + def _nested_slow_add_on_user_stream(dst, x, y, z): + ret = rpc.rpc_sync( + dst, TensorPipeAgentCudaRpcTest._slow_add_on_user_stream, args=(x, y) + ) + + return TensorPipeAgentCudaRpcTest._slow_add_on_user_stream(ret, z) + + def _test_stream_nested_sync(self, dst): + x = torch.ones(2, 2).to(0) + y = torch.ones(2, 2).to(0) * 2 + z = torch.ones(2, 2).to(0) * 3 + nested_dst = worker_name((self.rank + 2) % self.world_size) + ret = rpc.rpc_sync( + dst, + TensorPipeAgentCudaRpcTest._nested_slow_add_on_user_stream, + args=(nested_dst, x, y, z), + ) + self.assertEqual(ret, 6 * x) + + @skip_if_lt_x_gpu(2) + def test_custom_stream_nested(self): + self._test_custom_stream( + self._test_stream_nested_sync, {"cuda:0": "cuda:1", "cuda:1": "cuda:0"} + ) + + def _test_stream_nested_multi_async(self, dst): + if self.rank == 0: + futs = [] + n = 5 + xs, ys, zs = [], [], [] + for i in range(n): + x = torch.ones(2, 2).to(0) * (i - 1) + y = torch.ones(2, 2).to(0) * i + z = torch.ones(2, 2).to(0) * (i + 1) + xs.append(x) + ys.append(y) + zs.append(z) + nested_dst = worker_name((self.rank + 2) % self.world_size) + futs.append( + rpc.rpc_async( + dst, + TensorPipeAgentCudaRpcTest._nested_slow_add_on_user_stream, + args=(nested_dst, x, y, z), + ) + ) + + for i in range(n): + self.assertEqual(futs[i].wait(), xs[i] + ys[i] + zs[i]) + + @skip_if_lt_x_gpu(2) + def test_custom_stream_nested_multi(self): + self._test_custom_stream( + self._test_stream_nested_multi_async, + {"cuda:0": "cuda:1", "cuda:1": "cuda:0"}, + ) + + @staticmethod + def _gpu_add_wrong_gpus(x, y): + if x.is_cuda and y.is_cuda: + return x.cpu() + y.cuda() + else: + raise ValueError("Wrong device affinity") + + @skip_if_lt_x_gpu(1) + def test_device_mismatch(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {0: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + x = torch.zeros(2).to(0) + y = torch.ones(2).to(0) + + with self.assertRaisesRegex( + RuntimeError, + "Expected all tensors to be on the same device, but found at least two devices", + ): + rpc.rpc_sync( + dst, TensorPipeAgentCudaRpcTest._gpu_add_wrong_gpus, args=(x, y) + ) + + rpc.shutdown() + + def _test_rref_synchronization(self, local_device, remote_device): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {local_device: remote_device}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + if self.rank == 1: + # This test compares rref.rpc_sync().forward(x) vs rref.remote().forward(x).to_here() + # If to_here() is properly synchronized with forward(x) the results must be identical + # This test needs multiple iterations and significant batch size to simulate real + # training of a CNN of MNIST-like data. + # see https://github.com/pytorch/pytorch/issues/54771 + rref = rpc.remote(dst, MyConvNetForMNIST, args=(remote_device,)) + for _ in range(10): + x = torch.randn(200, 1, 28, 28).to(local_device) + actual = rref.remote().forward(x).to_here() + expected = rref.rpc_sync().forward(x) + self.assertEqual(actual, expected) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_rref_to_here_synchronization1(self): + self._test_rref_synchronization("cuda:0", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_to_here_synchronization2(self): + self._test_rref_synchronization("cuda:1", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_to_here_synchronization3(self): + self._test_rref_synchronization("cuda:1", "cuda:1") + + @skip_if_lt_x_gpu(2) + def test_rref_to_here_synchronization4(self): + self._test_rref_synchronization("cuda:0", "cuda:1") + + def _test_rref_as_arg_synchronization( + self, local_device, remote_device, devicesOptions=None + ): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {local_device: remote_device}) + + input_src = worker_name((self.rank - 1 + self.world_size) % self.world_size) + options.set_device_map(input_src, {remote_device: local_device}) + + if devicesOptions is not None: + options.set_devices(devicesOptions[self.rank]) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + if self.rank == 1: + # This test compares rref.rpc_sync().forward(x) vs rref.remote().forward(x).to_here() + # If to_here() is properly synchronized with forward(x) the results must be identical + # This test needs multiple iterations and significant batch size to simulate real + # training of a CNN of MNIST-like data. + # see https://github.com/pytorch/pytorch/issues/54771 + rref = rpc.remote(dst, MyConvNetForMNIST, args=(remote_device,)) + for _ in range(10): + rref_x = RRef(torch.randn(200, 1, 28, 28).to(local_device)) + actual = rref.remote().forward(rref_x, True).to_here() + expected = rref.rpc_sync().forward(rref_x, True) + self.assertEqual(actual, expected) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_rref_as_arg_synchronization1(self): + self._test_rref_as_arg_synchronization("cuda:0", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_as_arg_synchronization2(self): + self._test_rref_as_arg_synchronization("cuda:1", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_as_arg_synchronization3(self): + self._test_rref_as_arg_synchronization("cuda:1", "cuda:1") + + @skip_if_lt_x_gpu(2) + def test_rref_as_arg_synchronization4(self): + self._test_rref_as_arg_synchronization("cuda:0", "cuda:1") + + @skip_if_lt_x_gpu(1) + def test_rref_as_arg_synchronization5(self): + self._test_rref_as_arg_synchronization( + "cuda:0", + "cuda:0", + [["cuda:0"] for _ in range(4)], # devicesOptions + ) + + @staticmethod + def _rref_relay(rref): + return rref.to_here() + + def _test_rref_forward_synchronization(self, local_device, remote_device): + options = self.rpc_backend_options + + input_src = worker_name(0) + model_dst = worker_name(1) + out_relay = worker_name(2) + + if self.rank == 0: + # for 1) model construction 2) forward execution + options.set_device_map(model_dst, {local_device: remote_device}) + + # Forward output will be first copied to the relay node before + # returning to the worker. This is intentional, to test RRef + # forward CUDA stream synchronizations. + options.set_device_map(out_relay, {local_device: local_device}) + elif self.rank == 1: + # worker1 hosts the model and runs forward. The forward functions + # calls RRef.to_here(), hence needs to configure the device map + options.set_device_map(input_src, {remote_device: local_device}) + elif self.rank == 2: + # worker2 will get the out RRef and call to_here() and hence, needs + # to configure device map. + options.set_device_map(model_dst, {local_device: remote_device}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + if self.rank == 0: + # This test compares rref.rpc_sync().forward(x) vs rref.remote().forward(x).to_here() + # If to_here() is properly synchronized with forward(x) the results must be identical + # This test needs multiple iterations and significant batch size to simulate real + # training of a CNN of MNIST-like data. + # see https://github.com/pytorch/pytorch/issues/54771 + rref = rpc.remote(model_dst, MyConvNetForMNIST, args=(remote_device,)) + for _ in range(10): + rref_input = RRef(torch.randn(200, 1, 28, 28).to(local_device)) + rref_out = rref.remote().forward(rref_input, True) + out = rpc.remote( + out_relay, TensorPipeAgentCudaRpcTest._rref_relay, args=(rref_out,) + ).to_here() + expected = rref.rpc_sync().forward(rref_input, True) + self.assertEqual(out, expected) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_rref_forward_synchronization1(self): + self._test_rref_forward_synchronization("cuda:0", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_forward_synchronization2(self): + self._test_rref_forward_synchronization("cuda:0", "cuda:1") + + @skip_if_lt_x_gpu(2) + def test_rref_forward_synchronization3(self): + self._test_rref_forward_synchronization("cuda:1", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_rref_forward_synchronization4(self): + self._test_rref_forward_synchronization("cuda:1", "cuda:1") + + def _test_owner_rref_forward_synchronization(self, local_device, remote_device): + if self.rank == 0: + options = self.rpc_backend_options + options.set_device_map("w0", {local_device: remote_device}) + rpc.init_rpc("w0", rank=0, world_size=1, rpc_backend_options=options) + + model = ( + rpc.remote("w0", torch.nn.Linear, (2048, 20000)) + .remote() + .to(remote_device) + ) + for _ in range(30): + data = torch.rand(2048, 2048).to(local_device) + output = model.rpc_sync().forward(data) + # to_here() internally calls localValue as the caller is + # the owner of the RRef. + v0 = rpc.RRef(output).remote().sum().to_here().item() + v1 = output.sum().item() + self.assertEqual(v0, v1) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_owner_rref_forward_synchronization1(self): + self._test_owner_rref_forward_synchronization("cuda:0", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_owner_rref_forward_synchronization2(self): + self._test_owner_rref_forward_synchronization("cuda:0", "cuda:1") + + @skip_if_lt_x_gpu(2) + def test_owner_rref_forward_synchronization3(self): + self._test_owner_rref_forward_synchronization("cuda:1", "cuda:0") + + @skip_if_lt_x_gpu(2) + def test_owner_rref_forward_synchronization4(self): + self._test_owner_rref_forward_synchronization("cuda:1", "cuda:1") + + @staticmethod + def _return_tensor_view(i): + with torch.cuda.stream(torch.cuda.current_stream(0)): + x = torch.ones(1000, 200).cuda(0) * i + torch.cuda._sleep(10 * FIFTY_MIL_CYCLES) + # serialization of the return value will create a new tensor from the + # view, which is done outside of the user function. + return x.split(100)[0] + + @skip_if_lt_x_gpu(1) + def test_tensor_view_as_return_value(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {0: 0}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + futs = [ + rpc.rpc_async( + dst, TensorPipeAgentCudaRpcTest._return_tensor_view, args=(i,) + ) + for i in range(5) + ] + + for i in range(5): + self.assertEqual(torch.ones(100, 200) * i, futs[i].wait()) + + rpc.shutdown() + + @skip_if_lt_x_gpu(2) + def test_devices_option_mismatch(self): + with self.assertRaisesRegex( + ValueError, + "Node worker0 has unexpected source devices in its device map for worker1", + ): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {0: 0}) + options.set_devices([1]) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rpc.shutdown() + + @skip_if_lt_x_gpu(2) + def test_devices_option_mismatch_reverse(self): + with self.assertRaisesRegex( + ValueError, + "Node worker0 has unexpected target devices in its device map for worker1", + ): + dst = worker_name((self.rank + 1) % self.world_size) + + options = rpc.TensorPipeRpcBackendOptions( + init_method=self.rpc_backend_options.init_method, + num_worker_threads=self.rpc_backend_options.num_worker_threads, + device_maps={dst: {0: 1}}, + devices=[0], + ) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_cuda_future_device_as_int(self): + Future(devices=[0]) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_device_as_str(self): + Future(devices=["cuda:0"]) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_device_as_device(self): + Future(devices=[torch.device("cuda", 0)]) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_device_not_cuda(self): + with self.assertRaisesRegex( + ValueError, "Expected devices to have indices, got cpu" + ): + Future(devices=["cpu"]) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_cuda_tensor(self): + self._test_cuda_future_extraction( + wrapper=lambda t: t, unwrapper=lambda v: v, sparse_tensor=False + ) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_list_with_cuda_tensor(self): + self._test_cuda_future_extraction( + wrapper=lambda t: [t], unwrapper=operator.itemgetter(0), sparse_tensor=False + ) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_custom_class_with_cuda_tensor(self): + self._test_cuda_future_extraction( + wrapper=TensorWrapper, unwrapper=lambda v: v.tensor, sparse_tensor=False + ) + + @skip_if_lt_x_gpu(2) + def test_cuda_future_callback_changes_devices(self): + # We check proper CUDA stream synchronization by filling the tensor with + # the expected value in one stream, and reading it from another stream. + tensor0 = torch.zeros((100,), device="cuda:0") + tensor1 = torch.zeros((100,), device="cuda:1") + parent_future = Future(devices=["cuda:0", "cuda:1"]) + + def cb(fut): + t0 = fut.value() + tensor1.copy_(t0, non_blocking=True) + return tensor1 + + child_future = parent_future.then(cb) + with torch.cuda.device("cuda:0"): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + tensor0.fill_(1) + parent_future.set_result(tensor0) + with torch.cuda.device("cuda:1"): + another_stream = torch.cuda.Stream() + with torch.cuda.stream(another_stream): + self.assertTrue(torch.eq(child_future.wait(), 1).all().item()) + + @skip_if_lt_x_gpu(2) + def test_cuda_future_value_on_bad_device(self): + tensor0 = torch.zeros((100,), device="cuda:0") + tensor1 = torch.zeros((100,), device="cuda:1") + parent_future = Future(devices=["cuda:1"]) + + # As a plus, we test that futures still invoke callbacks even in case of + # error, and that the child futures are successful if those callbacks + # don't access the parent future. + def cb(fut): + with torch.cuda.device("cuda:1"): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + tensor1.fill_(1) + return tensor1 + + child_future = parent_future.then(cb) + with torch.cuda.device("cuda:0"): + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + torch.cuda._sleep(int(1000 * get_cycles_per_ms())) + tensor0.fill_(1) + parent_future.set_result(tensor0) + with self.assertRaisesRegex( + ValueError, + r"The result contained tensors residing on device\(s\) cuda:0 " + r"which are not among the expected device\(s\) cuda:1", + ): + parent_future.wait() + with torch.cuda.device("cuda:1"): + another_stream = torch.cuda.Stream() + with torch.cuda.stream(another_stream): + self.assertTrue(torch.eq(child_future.wait(), 1).all().item()) + + @skip_if_lt_x_gpu(1) + def test_async_execution_with_cuda_future(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {"cuda:0": "cuda:0"}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + t = torch.zeros((100,), device="cuda:0") + fut = rpc.rpc_async(dst, async_cuda_sleep_and_set_to_one, args=(t,)) + another_stream = torch.cuda.Stream("cuda:0") + with torch.cuda.stream(another_stream): + self.assertTrue(torch.eq(fut.wait(), 1).all().item()) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_async_execution_nested_with_cuda_future(self): + dst = worker_name((self.rank + 1) % self.world_size) + nested_dst = worker_name((self.rank + 2) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {"cuda:0": "cuda:0"}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + a = torch.ones((100,), device="cuda:0") + b = torch.ones((100,), device="cuda:0") + c = torch.ones((100,), device="cuda:0") + fut = rpc.rpc_async(dst, async_cuda_nested_add, args=(nested_dst, a, b, c)) + another_stream = torch.cuda.Stream("cuda:0") + with torch.cuda.stream(another_stream): + self.assertTrue(torch.eq(fut.wait(), 3).all().item()) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_cuda_future_modify_tensor_inplace(self): + tensor = torch.zeros((100,), device="cuda:0") + future = Future(devices=["cuda:0"]) + future.set_result(tensor) + # It's weird to modify the value of a future once it's complete, but + # technically possible. Currently this is considered undefined behavior + # (in practice the future will ignore the modification and still + # synchronize with the original value). We could one day add logic to + # detect and warn or throw in such cases, but for now we just check that + # this doesn't crash. + tensor.fill_(1) + future.wait() + + @skip_if_lt_x_gpu(1) + def test_cuda_future_replace_tensor(self): + tensor_list = [torch.zeros((100,), device="cuda:0")] + future = Future(devices=["cuda:0"]) + future.set_result(tensor_list) + # It's weird to modify the value of a future once it's complete, but + # technically possible. Currently this is considered undefined behavior + # (in practice the future will ignore the modification and still + # synchronize with the original value). We could one day add logic to + # detect and warn or throw in such cases, but for now we just check that + # this doesn't crash. + # We set things up so that the original tensor contained in the list + # gets deleted once we replace it with the other one. This will + # invalidate any cached information held by the future. + tensor_list[0] = torch.ones((100,), device="cuda:0") + future.wait() + + @skip_if_lt_x_gpu(1) + def test_rref_with_unpickleable_attributes(self): + dst = worker_name((self.rank + 1) % self.world_size) + options = self.rpc_backend_options + options.set_device_map(dst, {"cuda:0": "cuda:0"}) + + rpc.init_rpc( + name=worker_name(self.rank), + backend=self.rpc_backend, + rank=self.rank, + world_size=self.world_size, + rpc_backend_options=options, + ) + + rref = rpc.remote(dst, TensorWrapper, args=(torch.zeros(42, device="cuda:0"),)) + rref.rpc_sync().increase(1) + ret = rref.rpc_sync().sum() + self.assertEqual(ret, 42) + + rpc.shutdown() + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_cuda_sparse_tensor(self): + self._test_cuda_future_extraction( + wrapper=lambda t: t, unwrapper=lambda v: v, sparse_tensor=True + ) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_list_with_cuda_sparse_tensor(self): + self._test_cuda_future_extraction( + wrapper=lambda t: [t], unwrapper=operator.itemgetter(0), sparse_tensor=True + ) + + @skip_if_lt_x_gpu(1) + def test_cuda_future_can_extract_custom_class_with_cuda_sparse_tensor(self): + self._test_cuda_future_extraction( + wrapper=TensorWrapper, unwrapper=lambda v: v.tensor, sparse_tensor=True + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/tensorpipe_rpc_agent_test_fixture.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/tensorpipe_rpc_agent_test_fixture.py new file mode 100644 index 0000000000000000000000000000000000000000..021ae60468009d2fd4fa947c90455d99c1c6d54e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc/tensorpipe_rpc_agent_test_fixture.py @@ -0,0 +1,28 @@ +# mypy: allow-untyped-defs + +import torch.distributed.rpc as rpc +from torch.testing._internal.common_distributed import tp_transports +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) + + +class TensorPipeRpcAgentTestFixture(RpcAgentTestFixture): + @property + def rpc_backend(self): + return rpc.backend_registry.BackendType["TENSORPIPE"] + + @property + def rpc_backend_options(self): + return rpc.backend_registry.construct_rpc_backend_options( + self.rpc_backend, init_method=self.init_method, _transports=tp_transports() + ) + + def get_shutdown_error_regex(self): + # FIXME Once we consolidate the error messages returned by the + # TensorPipe agent put some more specific regex here. + error_regexes = [".*"] + return "|".join([f"({error_str})" for error_str in error_regexes]) + + def get_timeout_error_regex(self): + return "RPC ran for more than" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a24e4f97f05df22396dc08e3e6bc381085477882 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/distributed/rpc_utils.py @@ -0,0 +1,188 @@ +# mypy: allow-untyped-defs + +import os +import sys +import unittest + +from torch.testing._internal.common_distributed import MultiProcessTestCase +from torch.testing._internal.common_utils import ( + find_free_port, + IS_SANDCASTLE, + TEST_WITH_DEV_DBG_ASAN, +) +from torch.testing._internal.distributed.ddp_under_dist_autograd_test import ( + CudaDdpComparisonTest, + DdpComparisonTest, + DdpUnderDistAutogradTest, +) +from torch.testing._internal.distributed.nn.api.remote_module_test import ( + CudaRemoteModuleTest, + RemoteModuleTest, + ThreeWorkersRemoteModuleTest, +) +from torch.testing._internal.distributed.rpc.dist_autograd_test import ( + CudaDistAutogradTest, + DistAutogradTest, + FaultyAgentDistAutogradTest, + TensorPipeAgentDistAutogradTest, + TensorPipeCudaDistAutogradTest, +) +from torch.testing._internal.distributed.rpc.dist_optimizer_test import ( + DistOptimizerTest, +) +from torch.testing._internal.distributed.rpc.examples.parameter_server_test import ( + ParameterServerTest, +) +from torch.testing._internal.distributed.rpc.examples.reinforcement_learning_rpc_test import ( + ReinforcementLearningRpcTest, +) +from torch.testing._internal.distributed.rpc.faulty_agent_rpc_test import ( + FaultyAgentRpcTest, +) +from torch.testing._internal.distributed.rpc.jit.dist_autograd_test import ( + JitDistAutogradTest, +) +from torch.testing._internal.distributed.rpc.jit.rpc_test import JitRpcTest +from torch.testing._internal.distributed.rpc.jit.rpc_test_faulty import ( + JitFaultyAgentRpcTest, +) +from torch.testing._internal.distributed.rpc.rpc_agent_test_fixture import ( + RpcAgentTestFixture, +) +from torch.testing._internal.distributed.rpc.rpc_test import ( + CudaRpcTest, + RpcTest, + TensorPipeAgentCudaRpcTest, + TensorPipeAgentRpcTest, +) + + +def _check_and_set_tcp_init(): + # if we are running with TCP init, set main address and port + # before spawning subprocesses, since different processes could find + # different ports. + use_tcp_init = os.environ.get("RPC_INIT_WITH_TCP", None) + if use_tcp_init == "1": + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(find_free_port()) + + +def _check_and_unset_tcp_init(): + use_tcp_init = os.environ.get("RPC_INIT_WITH_TCP", None) + if use_tcp_init == "1": + del os.environ["MASTER_ADDR"] + del os.environ["MASTER_PORT"] + + +# The tests for the RPC module need to cover multiple possible combinations: +# - different aspects of the API, each one having its own suite of tests; +# - different agents (ProcessGroup, TensorPipe, ...); +# To avoid a combinatorial explosion in code size, and to prevent forgetting to +# add a combination, these are generated automatically by the code in this file. +# Here, we collect all the test suites that we need to cover. +# We then have one separate file for each agent, from which +# we call the generate_tests function of this file, passing to it a fixture for +# the agent, which then gets mixed-in with each test suite. + + +@unittest.skipIf( + TEST_WITH_DEV_DBG_ASAN, + "Skip ASAN as torch + multiprocessing spawn have known issues", +) +class SpawnHelper(MultiProcessTestCase): + def setUp(self): + super().setUp() + _check_and_set_tcp_init() + self._spawn_processes() + + def tearDown(self): + _check_and_unset_tcp_init() + super().tearDown() + + +# This list contains test suites that are agent-agnostic and that only verify +# compliance with the generic RPC interface specification. These tests should +# *not* make use of implementation details of a specific agent (options, +# attributes, ...). These test suites will be instantiated multiple times, once +# for each agent (except the faulty agent, which is special). +GENERIC_TESTS = [ + RpcTest, + ParameterServerTest, + DistAutogradTest, + DistOptimizerTest, + JitRpcTest, + JitDistAutogradTest, + RemoteModuleTest, + ThreeWorkersRemoteModuleTest, + DdpUnderDistAutogradTest, + DdpComparisonTest, + ReinforcementLearningRpcTest, +] +GENERIC_CUDA_TESTS = [ + CudaRpcTest, + CudaDistAutogradTest, + CudaRemoteModuleTest, + CudaDdpComparisonTest, +] + + +# This list contains test suites that will only be run on the TensorPipeAgent. +# These suites should be standalone, and separate from the ones in the generic +# list (not subclasses of those!). +TENSORPIPE_TESTS = [ + TensorPipeAgentRpcTest, + TensorPipeAgentDistAutogradTest, +] +TENSORPIPE_CUDA_TESTS = [ + TensorPipeAgentCudaRpcTest, + TensorPipeCudaDistAutogradTest, +] + + +# This list contains test suites that will only be run on the faulty RPC agent. +# That agent is special as it's only used to perform fault injection in order to +# verify the error handling behavior. Thus the faulty agent will only run the +# suites in this list, which were designed to test such behaviors, and not the +# ones in the generic list. +FAULTY_AGENT_TESTS = [ + FaultyAgentRpcTest, + FaultyAgentDistAutogradTest, + JitFaultyAgentRpcTest, +] + + +def generate_tests( + prefix: str, + mixin: type[RpcAgentTestFixture], + tests: list[type[RpcAgentTestFixture]], + module_name: str, +) -> dict[str, type[RpcAgentTestFixture]]: + """Mix in the classes needed to autogenerate the tests based on the params. + + Takes a series of test suites, each written against a "generic" agent (i.e., + derived from the abstract RpcAgentTestFixture class), as the `tests` args. + Takes a concrete subclass of RpcAgentTestFixture, which specializes it for a + certain agent, as the `mixin` arg. Produces all combinations of them. + Returns a dictionary of class names to class type + objects which can be inserted into the global namespace of the calling + module. The name of each test will be a concatenation of the `prefix` arg + and the original name of the test suite. + The `module_name` should be the name of the calling module so + that the classes can be fixed to make it look like they belong to it, which + is necessary for pickling to work on them. + """ + ret: dict[str, type[RpcAgentTestFixture]] = {} + for test_class in tests: + if IS_SANDCASTLE and TEST_WITH_DEV_DBG_ASAN: + print( + f"Skipping test {test_class} on sandcastle for the following reason: " + "Skip dev-asan as torch + multiprocessing spawn have known issues", + file=sys.stderr, + ) + continue + + name = f"{prefix}{test_class.__name__}" + class_ = type(name, (test_class, mixin, SpawnHelper), {}) + class_.__module__ = module_name + ret[name] = class_ + return ret diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/dynamo_pytree_test_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/dynamo_pytree_test_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..737b7d27a1561477c8a3781926453f90cf622c8c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/dynamo_pytree_test_utils.py @@ -0,0 +1,28 @@ +import torch +import torch._dynamo.test_case +import torch.utils._pytree as pytree + + +class PytreeRegisteringTestCase(torch._dynamo.test_case.TestCase): + """TestCase that prunes all temporary pytree registrations and resets Dynamo.""" + + def setUp(self) -> None: + super().setUp() + self._registered_pytree_nodes: list[type] = [] + self._registered_constant_nodes: list[type] = [] + + def tearDown(self) -> None: + for cls in reversed(self._registered_pytree_nodes): + pytree._deregister_pytree_node(cls) + for cls in reversed(self._registered_constant_nodes): + pytree._deregister_pytree_node(cls) + torch._dynamo.reset() + super().tearDown() + + def register_pytree_node(self, cls, *args, **kwargs) -> None: # type: ignore[no-untyped-def] + pytree.register_pytree_node(cls, *args, **kwargs) + self._registered_pytree_nodes.append(cls) + + def register_constant(self, cls: type) -> None: + pytree.register_constant(cls) + self._registered_constant_nodes.append(cls) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/dynamo_test_failures.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/dynamo_test_failures.py new file mode 100644 index 0000000000000000000000000000000000000000..f56a7b07f63a337cf5cb4b64b7dee5815c5bebde --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/dynamo_test_failures.py @@ -0,0 +1,144 @@ +""" +This file contains the list of tests that are known to fail under Dynamo + +We generate xFailIfTorchDynamo* for all tests in `dynamo_expected_failures` +We generate skipIfTorchDynamo* for all tests in `dynamo_skips` +We generate runWithoutCompiledAutograd for all tests in `compiled_autograd_skips` + +For an easier-than-manual way of generating and updating these lists, +see scripts/compile_tests/update_failures.py + +If you're adding a new test, and it's failing PYTORCH_TEST_WITH_DYNAMO=1, +either add the appropriate decorators to your test or add skips for them +via test/dynamo_skips and test/dynamo_expected_failures. + +*These are not exactly unittest.expectedFailure and unittest.skip. We'll +always execute the test and then suppress the signal, if necessary. +If your tests crashes, or is slow, please use @skipIfTorchDynamo instead. + +The expected failure and skip files are located in test/dynamo_skips and +test/dynamo_expected_failures. They're individual files rather than a list so +git will merge changes easier. +""" + +import logging +import os +import sys + + +def find_test_dir() -> str | None: + # Find the path to the dynamo expected failure and skip files. + from os.path import abspath, basename, dirname, exists, join, normpath + + if sys.platform == "win32": + return None + + # Check relative to this file (local build): + test_dir = normpath(join(dirname(abspath(__file__)), "../../../test")) + if exists(join(test_dir, "dynamo_expected_failures")): + return test_dir + + # Check relative to __main__ (installed builds relative to test file): + main = sys.modules["__main__"] + file = getattr(main, "__file__", None) + if file is None: + # Generated files do not have a module.__file__ + return None + test_dir = dirname(abspath(file)) + while dirname(test_dir) != test_dir: + if basename(test_dir) == "test" and exists( + join(test_dir, "dynamo_expected_failures") + ): + return test_dir + test_dir = dirname(test_dir) + + # Not found + return None + + +test_dir = find_test_dir() +if not test_dir: + logger = logging.getLogger(__name__) + logger.warning( + "test/dynamo_expected_failures directory not found - known dynamo errors won't be skipped." + ) + +# Tests that run without strict mode in PYTORCH_TEST_WITH_INDUCTOR=1. +# Please don't add anything to this list. +FIXME_inductor_non_strict = { + "test_modules", + "test_ops", + "test_ops_gradients", + "test_torch", +} + +# We generate unittest.expectedFailure for all of the following tests +# when run under PYTORCH_TEST_WITH_DYNAMO=1. +# see NOTE [dynamo_test_failures.py] for more details +# +# This lists exists so we can more easily add large numbers of failing tests, +if test_dir is None: + dynamo_expected_failures = set() + dynamo_skips = set() + + inductor_expected_failures = set() + inductor_skips = set() + + compiled_autograd_skips = set() +else: + dynamo_failures_directory = os.path.join(test_dir, "dynamo_expected_failures") + dynamo_skips_directory = os.path.join(test_dir, "dynamo_skips") + + dynamo_expected_failures = set(os.listdir(dynamo_failures_directory)) + dynamo_skips = set(os.listdir(dynamo_skips_directory)) + + inductor_failures_directory = os.path.join(test_dir, "inductor_expected_failures") + inductor_skips_directory = os.path.join(test_dir, "inductor_skips") + + inductor_expected_failures = set(os.listdir(inductor_failures_directory)) + inductor_skips = set(os.listdir(inductor_skips_directory)) + + compiled_autograd_skips_directory = os.path.join( + test_dir, "compiled_autograd_skips" + ) + compiled_autograd_skips = set(os.listdir(compiled_autograd_skips_directory)) + +# TODO: due to case sensitivity problems, for now list these files by hand +extra_dynamo_skips = { + "TestProxyTensorOpInfoCPU.test_make_fx_exhaustive_T_cpu_float32", + "TestProxyTensorOpInfoCPU.test_make_fx_exhaustive_t_cpu_float32", + "TestProxyTensorOpInfoCPU.test_make_fx_fake_exhaustive_T_cpu_float32", + "TestProxyTensorOpInfoCPU.test_make_fx_fake_exhaustive_t_cpu_float32", + "TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_T_cpu_float32", + "TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_t_cpu_float32", + "TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_inplace_T_cpu_float32", + "TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_inplace_t_cpu_float32", + "TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_out_T_cpu_float32", + "TestProxyTensorOpInfoCPU.test_make_fx_symbolic_exhaustive_out_t_cpu_float32", +} +dynamo_skips = dynamo_skips.union(extra_dynamo_skips) + + +# verify some invariants +for test in ( + dynamo_expected_failures + | dynamo_skips + | inductor_expected_failures + | inductor_skips +): + if len(test.split(".")) != 2: + raise AssertionError(f'Invalid test name: "{test}"') + +dynamo_intersection = dynamo_expected_failures.intersection(dynamo_skips) +if len(dynamo_intersection) > 0: + raise AssertionError( + "there should be no overlap between dynamo_expected_failures " + "and dynamo_skips, got " + str(dynamo_intersection) + ) + +inductor_intersection = inductor_expected_failures.intersection(inductor_skips) +if len(inductor_intersection) > 0: + raise AssertionError( + "there should be no overlap between inductor_expected_failures " + "and inductor_skips, got " + str(inductor_intersection) + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/fake_config_module.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/fake_config_module.py new file mode 100644 index 0000000000000000000000000000000000000000..0dc2a5d7df4064882028c11270c215c1e9934b57 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/fake_config_module.py @@ -0,0 +1,52 @@ +import sys + +from torch.utils._config_module import Config, install_config_module + + +e_bool = True +e_int = 1 +e_float = 1.0 +e_string = "string" +e_list = [1] +e_set = {1} +e_tuple = (1,) +e_dict = {1: 2} +e_none: bool | None = None +e_optional: bool | None = True +e_ignored = True +_e_ignored = True +magic_cache_config_ignored = True +# [@compile_ignored: debug] +e_compile_ignored = True +e_config: bool = Config(default=True) +e_jk: bool = Config(justknob="does_not_exist", default=True) +e_jk_false: bool = Config(justknob="does_not_exist", default=False) +e_env_default: bool = Config(env_name_default="ENV_TRUE", default=False) +e_env_default_FALSE: bool = Config(env_name_default="ENV_FALSE", default=True) +e_env_default_str: bool = Config(env_name_default="ENV_STR", default="default") +e_env_default_str_empty: bool = Config( + env_name_default="ENV_STR_EMPTY", default="default" +) +e_env_force: bool = Config(env_name_force="ENV_TRUE", default=False) +e_aliased_bool: bool = Config( + alias="torch.testing._internal.fake_config_module2.e_aliasing_bool" +) +e_deprecated: bool = Config( + default=True, deprecated=True, deprecation_message="is no longer needed" +) +e_not_deprecated: bool = Config(default=False) +e_deprecated_alias: bool = Config( + alias="torch.testing._internal.fake_config_module.e_not_deprecated", + deprecated=True, + deprecation_message="use something else instead", +) + + +class nested: + e_bool = True + + +_cache_config_ignore_prefix = ["magic_cache_config"] +_save_config_ignore = ["e_ignored"] + +install_config_module(sys.modules[__name__]) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/fake_config_module2.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/fake_config_module2.py new file mode 100644 index 0000000000000000000000000000000000000000..77c2e2baa4ddca7685adf734809488979c21ab63 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/fake_config_module2.py @@ -0,0 +1,13 @@ +import sys + +from torch.utils._config_module import Config, install_config_module + + +e_aliasing_bool = False + +e_env_default_multi: bool = Config( + env_name_default=["ENV_TRUE", "ENV_FALSE"], default=False +) +e_env_force_multi: bool = Config(env_name_force=["ENV_FAKE", "ENV_TRUE"], default=False) + +install_config_module(sys.modules[__name__]) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/fake_config_module3.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/fake_config_module3.py new file mode 100644 index 0000000000000000000000000000000000000000..41e2c3fddba2d8fec96656c950f1fcbd7225e9ea --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/fake_config_module3.py @@ -0,0 +1,11 @@ +import sys +from typing import Callable # noqa: UP035 + +from torch.utils._config_module import install_config_module + + +e_list = [1] +e_set = {1} +e_func: Callable | None = None + +install_config_module(sys.modules[__name__]) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/generated/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/generated/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/generated/annotated_fn_args.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/generated/annotated_fn_args.py new file mode 100644 index 0000000000000000000000000000000000000000..e3934bd525ea17229fc199134ec27927947ce3a6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/generated/annotated_fn_args.py @@ -0,0 +1,2923 @@ +""" +This file is needed for generating procedural tests required for +testing __torch_function__. See tests/test_overrides.py. +""" + +# flake8: noqa +import torch + +annotated_args = { + torch._C._VariableFunctions._cast_Byte: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._cast_Char: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._cast_Double: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._cast_Float: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._cast_Int: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._cast_Long: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._cast_Short: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._cast_Half: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._make_dual: [{'is_kwarg_only': 'False', 'name': 'primal', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tangent', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'level', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._unpack_dual: [{'is_kwarg_only': 'False', 'name': 'dual', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'level', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.align_tensors: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._assert_async: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._assert_async: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'assert_msg', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions._assert_scalar: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'assert_msg', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions._functional_assert_scalar: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'assert_msg', 'simple_type': 'c10::string_view'}, {'is_kwarg_only': 'False', 'name': 'dep_token', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._functional_assert_async: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'assert_msg', 'simple_type': 'c10::string_view'}, {'is_kwarg_only': 'False', 'name': 'dep_token', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._assert_tensor_metadata: [{'is_kwarg_only': 'False', 'name': 'a', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._print: [{'is_kwarg_only': 'False', 'name': 's', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.sym_constrain_range: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.sym_constrain_range_for_size: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._functional_sym_constrain_range: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'int64_t?'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'int64_t?'}, {'is_kwarg_only': 'False', 'name': 'dep_token', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._functional_sym_constrain_range_for_size: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'int64_t?'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'int64_t?'}, {'is_kwarg_only': 'False', 'name': 'dep_token', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._make_dep_token: [], + torch._C._VariableFunctions._use_cudnn_ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'blank', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._use_cudnn_ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'blank', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._cudnn_ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'blank', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'zero_infinity', 'simple_type': 'bool'}], + torch._C._VariableFunctions._cudnn_ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'blank', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'zero_infinity', 'simple_type': 'bool'}], + torch._C._VariableFunctions._use_cudnn_rnn_flatten_weight: [], + torch._C._VariableFunctions._cudnn_rnn_flatten_weight: [{'is_kwarg_only': 'False', 'name': 'weight_arr', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'weight_stride0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'input_size', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'hidden_size', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'proj_size', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}], + torch._C._VariableFunctions._cudnn_rnn: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'weight_stride0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'weight_buf', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'cx', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'hidden_size', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'proj_size', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'batch_sizes', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dropout_state', 'simple_type': 'Tensor?'}], + torch._C._VariableFunctions._cudnn_init_dropout_state: [{'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'dropout_seed', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._debug_has_internal_overlap: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._fused_dropout: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}], + torch._C._VariableFunctions._masked_scale: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'double'}], + torch._C._VariableFunctions.native_dropout: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool?'}], + torch._C._VariableFunctions._sobol_engine_draw: [{'is_kwarg_only': 'False', 'name': 'quasi', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'sobolstate', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dimension', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'num_generated', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType?'}], + torch._C._VariableFunctions._sobol_engine_ff_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'sobolstate', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dimension', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'num_generated', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._sobol_engine_scramble_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'ltm', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dimension', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._sobol_engine_initialize_state_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dimension', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._reshape_from_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'shape', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._shape_as_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.dropout: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}], + torch._C._VariableFunctions.dropout_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}], + torch._C._VariableFunctions.feature_dropout: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}], + torch._C._VariableFunctions.feature_dropout_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}], + torch._C._VariableFunctions.alpha_dropout: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}], + torch._C._VariableFunctions.alpha_dropout_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}], + torch._C._VariableFunctions.feature_alpha_dropout: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}], + torch._C._VariableFunctions.feature_alpha_dropout_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}], + torch._C._VariableFunctions.abs: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.abs: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.abs_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.absolute: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.absolute: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.angle: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.angle: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.view_as_real: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.view_as_complex: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sgn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sgn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.real: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.imag: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._conj: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conj: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._conj_physical: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conj_physical: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conj_physical: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conj_physical_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.resolve_conj: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.resolve_neg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._neg_view: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.acos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.acos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.acos_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arccos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arccos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arccos_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.avg_pool1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.adaptive_avg_pool1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.adaptive_max_pool1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._add_relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._add_relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._add_relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._add_relu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._add_relu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.addmv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addmv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addmv_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.affine_grid_generator: [{'is_kwarg_only': 'False', 'name': 'theta', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._VariableFunctions._is_all_true: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._is_any_true: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._test_check_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._test_functorch_fallback: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.allclose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arange: [{'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.arange: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.arange: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.arange: [{'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.arange: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._dim_arange: [{'is_kwarg_only': 'False', 'name': 'like', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.argmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.argmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.argmin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.argmin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.acosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.acosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.acosh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arccosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arccosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arccosh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.asinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.asinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.asinh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arcsinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arcsinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arcsinh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atanh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arctanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arctanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arctanh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.as_strided: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.as_strided_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.asin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.asin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.asin_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arcsin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arcsin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arcsin_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atan_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arctan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arctan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arctan_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atleast_1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atleast_1d: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.atleast_2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atleast_2d: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.atleast_3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atleast_3d: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.baddbmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.baddbmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.baddbmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.baddbmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.bartlett_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.bartlett_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'periodic', 'simple_type': 'bool'}], + torch._C._VariableFunctions.batch_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'cudnn_enabled', 'simple_type': 'bool'}], + torch._C._VariableFunctions.quantized_batch_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'var', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'output_scale', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'output_zero_point', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._batch_norm_impl_index: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'cudnn_enabled', 'simple_type': 'bool'}], + torch._C._VariableFunctions.bernoulli: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bernoulli: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bernoulli: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}], + torch._C._VariableFunctions.bilinear: [{'is_kwarg_only': 'False', 'name': 'input1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.binary_cross_entropy_with_logits: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bincount: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_not: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_not: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.copysign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.copysign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.copysign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.copysign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._lazy_clone: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logical_not: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logical_not: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logical_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logical_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logical_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logical_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logical_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logical_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.blackman_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.blackman_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'periodic', 'simple_type': 'bool'}], + torch._C._VariableFunctions.bmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.bmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.broadcast_tensors: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.broadcast_to: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions._sparse_broadcast_to: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.cat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.cat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.cat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.cat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.concat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.concat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.concat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.concat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.concatenate: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.concatenate: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.concatenate: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.concatenate: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.block_diag: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.ceil: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ceil: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ceil_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.chain_matmul: [{'is_kwarg_only': 'False', 'name': 'matrices', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.chain_matmul: [{'is_kwarg_only': 'False', 'name': 'matrices', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.unsafe_chunk: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'chunks', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.chunk: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'chunks', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.tensor_split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sections', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.tensor_split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.tensor_split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor_indices_or_sections', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.clamp_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.clamp_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp_max_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.clamp_max_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp_min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.clamp_min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp_min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.clamp_min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clamp_min_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.clamp_min_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clip: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clip: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clip: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clip: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clip_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.clip_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cudnn_is_acceptable: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.complex: [{'is_kwarg_only': 'False', 'name': 'real', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'imag', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.complex: [{'is_kwarg_only': 'False', 'name': 'real', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'imag', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.polar: [{'is_kwarg_only': 'False', 'name': 'abs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'angle', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.polar: [{'is_kwarg_only': 'False', 'name': 'abs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'angle', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.constant_pad_nd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'pad', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.convolution: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'transposed', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'output_padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions._convolution: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'transposed', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'output_padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'benchmark', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'cudnn_enabled', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'allow_tf32', 'simple_type': 'bool'}], + torch._C._VariableFunctions._convolution: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'transposed', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'output_padding', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'benchmark', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'cudnn_enabled', 'simple_type': 'bool'}], + torch._C._VariableFunctions._convolution_mode: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'c10::string_view'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.conv1d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conv1d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conv2d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conv2d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conv3d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conv3d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conv_tbc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conv_transpose1d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conv_transpose2d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.conv_transpose3d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._copy_from: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dst', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._copy_from_and_resize: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dst', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cos_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cosh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cosine_embedding_loss: [{'is_kwarg_only': 'False', 'name': 'input1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.count_nonzero: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.count_nonzero: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cov: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.corrcoef: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cudnn_affine_grid_generator: [{'is_kwarg_only': 'False', 'name': 'theta', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'N', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'C', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'H', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'W', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cudnn_batch_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'exponential_average_factor', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'epsilon', 'simple_type': 'double'}], + torch._C._VariableFunctions.cudnn_batch_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'exponential_average_factor', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'epsilon', 'simple_type': 'double'}], + torch._C._VariableFunctions.cudnn_convolution: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'benchmark', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'allow_tf32', 'simple_type': 'bool'}], + torch._C._VariableFunctions.cudnn_convolution: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'benchmark', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'allow_tf32', 'simple_type': 'bool'}], + torch._C._VariableFunctions.cudnn_convolution_transpose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'output_padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'benchmark', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'allow_tf32', 'simple_type': 'bool'}], + torch._C._VariableFunctions._mps_convolution_transpose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'output_padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.cudnn_convolution_relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.cudnn_convolution_add_relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'z', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'alpha', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.cudnn_grid_sampler: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'grid', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cummax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cummax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cummax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.cummax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions._cummax_helper: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cummin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cummin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cummin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.cummin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions._cummin_helper: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cumprod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cumprod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cumprod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.cumprod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.cumsum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cumsum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cumsum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.cumsum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.cumulative_trapezoid: [{'is_kwarg_only': 'False', 'name': 'y', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cumulative_trapezoid: [{'is_kwarg_only': 'False', 'name': 'y', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.diag_embed: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.diagflat: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.diagonal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.diagonal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.diff: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.diff: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gradient: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gradient: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'spacing', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.gradient: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.gradient: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'spacing', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions.gradient: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'spacing', 'simple_type': 'ScalarList'}, {'is_kwarg_only': 'True', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.gradient: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'spacing', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.gradient: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'spacing', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch._C._VariableFunctions.div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch._C._VariableFunctions.div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch._C._VariableFunctions.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch._C._VariableFunctions.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch._C._VariableFunctions.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch._C._VariableFunctions.true_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.true_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.true_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.dot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.dot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.vdot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.vdot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.einsum: [{'is_kwarg_only': 'False', 'name': 'equation', 'simple_type': 'c10::string_view'}, {'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.embedding: [{'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.embedding_renorm_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max_norm', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'norm_type', 'simple_type': 'double'}], + torch._C._VariableFunctions._embedding_bag_forward_only: [{'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._rowwise_prune: [{'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'compressed_indices_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.row_stack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.row_stack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.embedding_bag: [{'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.embedding_bag: [{'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_grad_by_freq', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'sparse', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'per_sample_weights', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'include_last_offset', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'padding_idx', 'simple_type': 'int64_t?'}], + torch._C._VariableFunctions._embedding_bag: [{'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.empty: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch._C._VariableFunctions.empty: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.empty: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.empty_permuted: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'physical_layout', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._empty_affine_quantized: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions._empty_per_channel_affine_quantized: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'scales', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'zero_points', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'axis', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._resize_output_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'device', 'simple_type': 'Device'}], + torch._C._VariableFunctions.empty_quantized: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'qtensor', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.empty_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.empty_strided: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.erf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.erf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.erf_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.erfc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.erfc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.erfc_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.exp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.exp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.exp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.exp2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.exp2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.exp2_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.expm1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.expm1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.expm1_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.eye: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.eye: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'm', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.eye: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.eye: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'm', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.flatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.flatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'start_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'end_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'out_dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.flatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'start_dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'end_dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'out_dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.flatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'DimnameList'}, {'is_kwarg_only': 'False', 'name': 'out_dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.unflatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'sizes', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.unflatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'sizes', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'names', 'simple_type': 'DimnameList'}], + torch._C._VariableFunctions.fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.floor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.floor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.floor_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.floor_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.floor_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.floor_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.frac: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.frac: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.frac_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.full: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'fill_value', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch._C._VariableFunctions.full: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'fill_value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.full: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'fill_value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.full_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'fill_value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.from_file: [{'is_kwarg_only': 'False', 'name': 'filename', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.gcd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gcd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gcd_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lcm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lcm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lcm_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.grid_sampler: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'grid', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'interpolation_mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'padding_mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._VariableFunctions.grid_sampler_2d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'grid', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'interpolation_mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'padding_mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._VariableFunctions._grid_sampler_2d_cpu_fallback: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'grid', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'interpolation_mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'padding_mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._VariableFunctions.grid_sampler_3d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'grid', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'interpolation_mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'padding_mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._VariableFunctions.hann_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.hann_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'periodic', 'simple_type': 'bool'}], + torch._C._VariableFunctions.hamming_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.hamming_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'periodic', 'simple_type': 'bool'}], + torch._C._VariableFunctions.hamming_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'periodic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'alpha', 'simple_type': 'double'}], + torch._C._VariableFunctions.hamming_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'periodic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'alpha', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'beta', 'simple_type': 'double'}], + torch._C._VariableFunctions.kaiser_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.kaiser_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'periodic', 'simple_type': 'bool'}], + torch._C._VariableFunctions.kaiser_window: [{'is_kwarg_only': 'False', 'name': 'window_length', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'periodic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'beta', 'simple_type': 'double'}], + torch._C._VariableFunctions.hinge_embedding_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.group_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'num_groups', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.native_group_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'N', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'C', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'HxW', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'group', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions._fft_r2c: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'normalization', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'onesided', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fft_r2c: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'normalization', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'onesided', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fft_c2r: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'normalization', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'last_dim_size', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions._fft_c2r: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'normalization', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'last_dim_size', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions._fft_c2c: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'normalization', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'forward', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fft_c2c: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'normalization', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'forward', 'simple_type': 'bool'}], + torch._C._VariableFunctions._validate_compressed_sparse_indices: [{'is_kwarg_only': 'False', 'name': 'is_crow', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'compressed_idx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'plain_idx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'cdim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'nnz', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._cufft_get_plan_cache_size: [{'is_kwarg_only': 'False', 'name': 'device_index', 'simple_type': 'DeviceIndex'}], + torch._C._VariableFunctions._cufft_get_plan_cache_max_size: [{'is_kwarg_only': 'False', 'name': 'device_index', 'simple_type': 'DeviceIndex'}], + torch._C._VariableFunctions._cufft_set_plan_cache_max_size: [{'is_kwarg_only': 'False', 'name': 'device_index', 'simple_type': 'DeviceIndex'}, {'is_kwarg_only': 'False', 'name': 'max_size', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._cufft_clear_plan_cache: [{'is_kwarg_only': 'False', 'name': 'device_index', 'simple_type': 'DeviceIndex'}], + torch._C._VariableFunctions._unsafe_index: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'c10::List<::std::optional>'}], + torch._C._VariableFunctions._unsafe_masked_index: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'c10::List<::std::optional>'}, {'is_kwarg_only': 'False', 'name': 'fill', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._unsafe_masked_index_put_accumulate: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'c10::List<::std::optional>'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_put_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'c10::List<::std::optional>'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_put: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'c10::List<::std::optional>'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._unsafe_index_put: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'c10::List<::std::optional>'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._index_put_impl_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'c10::List<::std::optional>'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.instance_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'use_input_stats', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'cudnn_enabled', 'simple_type': 'bool'}], + torch._C._VariableFunctions.isclose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isin: [{'is_kwarg_only': 'False', 'name': 'elements', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'test_elements', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isin: [{'is_kwarg_only': 'False', 'name': 'elements', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'test_elements', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isin: [{'is_kwarg_only': 'False', 'name': 'elements', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'test_element', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.isin: [{'is_kwarg_only': 'False', 'name': 'elements', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'test_element', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.isin: [{'is_kwarg_only': 'False', 'name': 'element', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'test_elements', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isin: [{'is_kwarg_only': 'False', 'name': 'element', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'test_elements', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isnan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.is_distributed: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.is_floating_point: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.is_complex: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.is_conj: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._is_zerotensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.is_neg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isreal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.is_nonzero: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.is_same_size: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.is_signed: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.is_inference: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.kl_div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.kron: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.kron: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.kthvalue: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.kthvalue: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.kthvalue: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.kthvalue: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.layer_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'normalized_shape', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.native_layer_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'normalized_shape', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions.rms_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'normalized_shape', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions._fused_rms_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'normalized_shape', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double?'}], + torch._C._VariableFunctions.nan_to_num: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nan_to_num: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nan_to_num_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mkldnn_linear_backward_weights: [{'is_kwarg_only': 'False', 'name': 'grad_output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias_defined', 'simple_type': 'bool'}], + torch._C._VariableFunctions._cslt_compress: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._cslt_sparse_mm: [{'is_kwarg_only': 'False', 'name': 'compressed_A', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dense_B', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._cslt_sparse_mm_search: [{'is_kwarg_only': 'False', 'name': 'compressed_A', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dense_B', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sparse_semi_structured_tile: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sparse_semi_structured_apply: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'thread_masks', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sparse_semi_structured_apply_dense: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'thread_masks', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sparse_semi_structured_linear: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'meta', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sparse_semi_structured_mm: [{'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1_meta', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sparse_semi_structured_addmm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1_meta', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._mixed_dtypes_linear: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fbgemm_linear_int8_weight_fp32_activation: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight_scale', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'weight_zero_point', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fbgemm_linear_int8_weight: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight_scale', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'weight_zero_point', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fbgemm_linear_quantize_weight: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fbgemm_pack_gemm_matrix_fp16: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._wrapped_linear_prepack: [{'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight_scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight_zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._wrapped_quantized_linear_prepacked: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'out_channel', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.fbgemm_linear_fp16_weight_fp32_activation: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}], + torch._C._VariableFunctions.fbgemm_linear_fp16_weight_fp32_activation: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fbgemm_linear_fp16_weight: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fbgemm_linear_fp16_weight: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fbgemm_pack_quantized_matrix: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fbgemm_pack_quantized_matrix: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'K', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'N', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.ldexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ldexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ldexp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.linspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.linspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.linspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.linspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.linspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.linspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.linspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.linspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.log: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log10: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log10: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log10_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log1p: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log1p: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log1p_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.log2_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logaddexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logaddexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logaddexp2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logaddexp2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.xlogy_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.xlogy_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.logspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logspace: [{'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'steps', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.log_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.log_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.log_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions._log_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'half_to_float', 'simple_type': 'bool'}], + torch._C._VariableFunctions._log_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'half_to_float', 'simple_type': 'bool'}], + torch._C._VariableFunctions._log_softmax_backward_data: [{'is_kwarg_only': 'False', 'name': 'grad_output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'input_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions._log_softmax_backward_data: [{'is_kwarg_only': 'False', 'name': 'grad_output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'input_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions._logcumsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._logcumsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logcumsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logcumsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.logcumsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.logcumsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.logsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.logsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.logsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.logsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.margin_ranking_loss: [{'is_kwarg_only': 'False', 'name': 'input1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.matmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.matmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.matrix_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.matrix_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.matrix_exp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._aminmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._aminmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.aminmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.aminmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._compute_linear_combination: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'coefficients', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._compute_linear_combination: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'coefficients', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.amax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.amax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.max_pool1d_with_indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.max_pool1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.max_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._VariableFunctions.mkldnn_max_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._VariableFunctions.mkldnn_max_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._VariableFunctions.quantized_max_pool1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.quantized_max_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._VariableFunctions.quantized_max_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._VariableFunctions.max_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._VariableFunctions.mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.nanmean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nanmean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.median: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.median: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.median: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.median: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.median: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.nanmedian: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nanmedian: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.nanmedian: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.nanmedian: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.nanmedian: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.amin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.amin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._mps_convolution: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.mkldnn_convolution: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.mkldnn_rnn_layer: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight0', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight3', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx_', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'cx_', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reverse', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'batch_sizes', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'hidden_size', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}], + torch._C._VariableFunctions.miopen_batch_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'exponential_average_factor', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'epsilon', 'simple_type': 'double'}], + torch._C._VariableFunctions.miopen_convolution: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'benchmark', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}], + torch._C._VariableFunctions.miopen_convolution_transpose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'output_padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'benchmark', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}], + torch._C._VariableFunctions.miopen_depthwise_convolution: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'benchmark', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}], + torch._C._VariableFunctions.miopen_convolution_relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.miopen_convolution_add_relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'z', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'alpha', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.miopen_rnn: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'weight_stride0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'cx', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'mode', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'hidden_size', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'batch_sizes', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dropout_state', 'simple_type': 'Tensor?'}], + torch._C._VariableFunctions._use_miopen_ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'blank', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._use_miopen_ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'blank', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.miopen_ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'blank', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'zero_infinity', 'simple_type': 'bool'}], + torch._C._VariableFunctions.miopen_ctc_loss: [{'is_kwarg_only': 'False', 'name': 'log_probs', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'targets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target_lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'blank', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'deterministic', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'zero_infinity', 'simple_type': 'bool'}], + torch._C._VariableFunctions.mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions._int_mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._int_mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._convert_weight_to_int4pack: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'innerKTiles', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._weight_int4pack_mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'qGroupSize', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'qScaleAndZeros', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._weight_int4pack_mm_with_scales_and_zeros: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'qGroupSize', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'qScale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'qZeros', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._convert_weight_to_int4pack_for_cpu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'innerKTiles', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._weight_int4pack_mm_for_cpu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'qGroupSize', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'qScaleAndZeros', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._dyn_quant_pack_4bit_weight: [{'is_kwarg_only': 'False', 'name': 'weights', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scales_zeros', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'block_size', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'in_features', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'out_features', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._dyn_quant_matmul_4bit: [{'is_kwarg_only': 'False', 'name': 'inp', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_weights', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'block_size', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'in_features', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'out_features', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._weight_int8pack_mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scales', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sparse_sparse_matmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mode: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mode: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mode: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.mode: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.mul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.multiply: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.multiply: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.multiply: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.mv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mvlgamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.mvlgamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.narrow_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'length', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.narrow_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'length', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.narrow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'length', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.narrow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'length', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.native_batch_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions.native_batch_norm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions._native_batch_norm_legit: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions._native_batch_norm_legit: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions._native_batch_norm_legit: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions._native_batch_norm_legit: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'training', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions._native_batch_norm_legit_no_training: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions.batch_norm_stats: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions.batch_norm_elemt: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'invstd', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions.batch_norm_elemt: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'invstd', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}], + torch._C._VariableFunctions.batch_norm_gather_stats: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'invstd', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'count', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.batch_norm_gather_stats_with_counts: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'invstd', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'counts', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.batch_norm_backward_reduce: [{'is_kwarg_only': 'False', 'name': 'grad_out', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'invstd', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'input_g', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'weight_g', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bias_g', 'simple_type': 'bool'}], + torch._C._VariableFunctions.batch_norm_backward_elemt: [{'is_kwarg_only': 'False', 'name': 'grad_out', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'invstd', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'sum_dy', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sum_dy_xmu', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'count', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.batch_norm_update_stats: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_mean', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'running_var', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'momentum', 'simple_type': 'double'}], + torch._C._VariableFunctions.is_vulkan_available: [], + torch._C._VariableFunctions._nnpack_available: [], + torch._C._VariableFunctions._nnpack_spatial_convolution: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._VariableFunctions.ones: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch._C._VariableFunctions.ones: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.ones: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.ones_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.pairwise_distance: [{'is_kwarg_only': 'False', 'name': 'x1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'x2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cdist: [{'is_kwarg_only': 'False', 'name': 'x1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'x2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._euclidean_dist: [{'is_kwarg_only': 'False', 'name': 'x1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'x2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.pdist: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cosine_similarity: [{'is_kwarg_only': 'False', 'name': 'x1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'x2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.permute: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.movedim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'destination', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.movedim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'destination', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.moveaxis: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'destination', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.moveaxis: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'destination', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.adjoint: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.pixel_shuffle: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'upscale_factor', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.pixel_unshuffle: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'downscale_factor', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.channel_shuffle: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.native_channel_shuffle: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'groups', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions._pin_memory: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.pinverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.poisson_nll_loss: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'log_input', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'full', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'reduction', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.rad2deg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rad2deg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rad2deg_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.deg2rad: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.deg2rad: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.deg2rad_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.scalar_tensor: [{'is_kwarg_only': 'False', 'name': 's', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.rand: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch._C._VariableFunctions.rand: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}, {'is_kwarg_only': 'True', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch._C._VariableFunctions.rand: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.rand: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.rand: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.rand: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.rand_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rand_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randint: [{'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.randint: [{'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randint: [{'is_kwarg_only': 'False', 'name': 'low', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.randint: [{'is_kwarg_only': 'False', 'name': 'low', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randint: [{'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.randint: [{'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randint: [{'is_kwarg_only': 'False', 'name': 'low', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.randint: [{'is_kwarg_only': 'False', 'name': 'low', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randint_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.randint_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randint_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.randint_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randint_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'low', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.randint_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'low', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'high', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randn: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.randn: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randn: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch._C._VariableFunctions.randn: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}, {'is_kwarg_only': 'True', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch._C._VariableFunctions.randn: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.randn: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randn_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.randn_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randperm: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.randperm: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.randperm: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.randperm: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'True', 'name': 'generator', 'simple_type': 'Generator?'}], + torch._C._VariableFunctions.ravel: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.reciprocal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.reciprocal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.reciprocal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.neg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.neg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.neg_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.negative: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.negative: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.negative_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.repeat_interleave: [{'is_kwarg_only': 'False', 'name': 'repeats', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.repeat_interleave: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'repeats', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.repeat_interleave: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'repeats', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.reshape: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'shape', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions._mkldnn_reshape: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'shape', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.round: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.round: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.round: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.round: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.round_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.round_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rrelu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rrelu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.relu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.prelu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._prelu_kernel: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.hardshrink: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.hardshrink: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rsqrt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rsqrt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rsqrt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.selu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.selu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.celu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.celu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sigmoid: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sigmoid: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sigmoid_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logit_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sin_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sinc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sinc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sinc_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sinh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.detach: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.detach_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.slice_inverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.slice_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.slice_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.select_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.diagonal_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.as_strided_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.smm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions._softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'half_to_float', 'simple_type': 'bool'}], + torch._C._VariableFunctions._softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'half_to_float', 'simple_type': 'bool'}], + torch._C._VariableFunctions._softmax_backward_data: [{'is_kwarg_only': 'False', 'name': 'grad_output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'input_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions._softmax_backward_data: [{'is_kwarg_only': 'False', 'name': 'grad_output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'input_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.unsafe_split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_size', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_size', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.unsafe_split_with_sizes: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_sizes', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.split_with_sizes: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_sizes', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.hsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sections', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.hsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.vsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sections', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.vsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.dsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sections', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.dsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.squeeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.squeeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.squeeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.squeeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.sspaddmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sspaddmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._chunk_cat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'num_chunks', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._chunk_cat: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'num_chunks', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.stack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.stack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._stack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._stack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.hstack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.hstack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.vstack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.vstack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.dstack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.dstack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.stft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n_fft', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.stft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n_fft', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.istft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n_fft', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.nansum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nansum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.hash_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.hash_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sqrt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sqrt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sqrt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.square: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.square: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.square_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.std_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.std_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.std_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.std_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.std_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.prod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.prod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.prod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.prod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.prod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.t: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tan_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tanh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tensordot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims_self', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dims_other', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.tensordot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims_self', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'dims_other', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.threshold: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'threshold', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.threshold: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'threshold', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.threshold_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'threshold', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.tile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.transpose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.transpose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions._mkldnn_transpose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._mkldnn_transpose_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.flip: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.fliplr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.flipud: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.roll: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'shifts', 'simple_type': 'SymIntArrayRef', 'size': 1}], + torch._C._VariableFunctions.rot90: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.trapezoid: [{'is_kwarg_only': 'False', 'name': 'y', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.trapezoid: [{'is_kwarg_only': 'False', 'name': 'y', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.trapz: [{'is_kwarg_only': 'False', 'name': 'y', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.trapz: [{'is_kwarg_only': 'False', 'name': 'y', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._transform_bias_rescale_qkv: [{'is_kwarg_only': 'False', 'name': 'qkv', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'qkv_bias', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'num_heads', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._nested_tensor_from_mask: [{'is_kwarg_only': 'False', 'name': 't', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_tensor_from_mask_left_aligned: [{'is_kwarg_only': 'False', 'name': 't', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_from_padded: [{'is_kwarg_only': 'False', 'name': 'padded', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'cpu_nested_shape_example', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_from_padded_and_nested_example: [{'is_kwarg_only': 'False', 'name': 'padded', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'nt_example', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_view_from_buffer: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'nested_size', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'nested_strides', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_view_from_buffer_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'nested_size', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'nested_strides', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_view_from_buffer_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'nested_size', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'nested_strides', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_view_from_jagged: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dummy', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_view_from_jagged_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dummy', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_view_from_jagged_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dummy', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_get_values: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_get_values_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_get_values_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_get_offsets: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_get_lengths: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_get_ragged_idx: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_get_min_seqlen: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_get_max_seqlen: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_get_jagged_dummy: [{'is_kwarg_only': 'False', 'name': 'any', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_compute_contiguous_strides_offsets: [{'is_kwarg_only': 'False', 'name': 'nested_size', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._trilinear: [{'is_kwarg_only': 'False', 'name': 'i1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'i2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'i3', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'expand1', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'expand2', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'expand3', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'sumdim', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.triplet_margin_loss: [{'is_kwarg_only': 'False', 'name': 'anchor', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'positive', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'negative', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.trunc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.trunc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.trunc_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fix: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fix: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fix_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._has_compatible_shallow_copy_type: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'from', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._unique: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.unique_dim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.unique_consecutive: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._unique2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.unsqueeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.vander: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.var_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.var_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch._C._VariableFunctions.var_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.var_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.var_mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.where: [{'is_kwarg_only': 'False', 'name': 'condition', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.where: [{'is_kwarg_only': 'False', 'name': 'condition', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.where: [{'is_kwarg_only': 'False', 'name': 'condition', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.where: [{'is_kwarg_only': 'False', 'name': 'condition', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.where: [{'is_kwarg_only': 'False', 'name': 'condition', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.where: [{'is_kwarg_only': 'False', 'name': 'condition', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.norm_except_dim: [{'is_kwarg_only': 'False', 'name': 'v', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._weight_norm: [{'is_kwarg_only': 'False', 'name': 'v', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'g', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._weight_norm_interface: [{'is_kwarg_only': 'False', 'name': 'v', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'g', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.zeros: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'True', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch._C._VariableFunctions.zeros: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.zeros: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions._efficientzerotensor: [{'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.zeros_like: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._standard_gamma_grad: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._standard_gamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._philox_key_split: [{'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'num_splits', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._philox_key_fold_in: [{'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'data', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._philox_normal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._philox_uniform_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._dirichlet_grad: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'alpha', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'total', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sample_dirichlet: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.poisson: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.binomial: [{'is_kwarg_only': 'False', 'name': 'count', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'prob', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.native_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.native_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}, {'is_kwarg_only': 'False', 'name': 'keepdim', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType?'}], + torch._C._VariableFunctions._sparse_sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sparse_sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions._sparse_sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions._sparse_sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions._sparse_csr_sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions._sparse_csr_prod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions._sparse_softmax_backward_data: [{'is_kwarg_only': 'False', 'name': 'grad_output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._sparse_log_softmax_backward_data: [{'is_kwarg_only': 'False', 'name': 'grad_output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}, {'is_kwarg_only': 'False', 'name': 'keepdim', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}, {'is_kwarg_only': 'False', 'name': 'keepdim', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}, {'is_kwarg_only': 'False', 'name': 'keepdim', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}, {'is_kwarg_only': 'False', 'name': 'keepdim', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch._C._VariableFunctions.frexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.frexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.frobenius_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.frobenius_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._VariableFunctions.nuclear_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nuclear_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nuclear_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._VariableFunctions.nuclear_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._VariableFunctions.clone: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.positive: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.resize_as_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'the_template', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.resize_as_sparse_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'the_template', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.zero_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sub: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sub: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.subtract: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.subtract: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.subtract: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.rsub: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rsub: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.heaviside: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.heaviside: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.addmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions._addmm_activation: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._addmm_activation: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._scaled_mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_a', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_b', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._scaled_mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_a', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_b', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._scaled_mm_v2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_a', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'recipe_a', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'swizzle_a', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'scale_b', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'recipe_b', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'swizzle_b', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType?'}], + torch._C._VariableFunctions._scaled_mm_v2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_a', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'recipe_a', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'swizzle_a', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'scale_b', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'recipe_b', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'swizzle_b', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'out_dtype', 'simple_type': 'ScalarType?'}], + torch._C._VariableFunctions._scaled_grouped_mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_a', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_b', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._scaled_grouped_mm_v2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_a', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'recipe_a', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'swizzle_a', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'scale_b', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'recipe_b', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'swizzle_b', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._grouped_mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._validate_sparse_coo_tensor_args: [{'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._validate_sparse_compressed_tensor_args: [{'is_kwarg_only': 'False', 'name': 'compressed_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'plain_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'layout', 'simple_type': 'Layout'}], + torch._C._VariableFunctions._validate_sparse_csr_tensor_args: [{'is_kwarg_only': 'False', 'name': 'crow_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._validate_sparse_csc_tensor_args: [{'is_kwarg_only': 'False', 'name': 'ccol_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'row_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._validate_sparse_bsr_tensor_args: [{'is_kwarg_only': 'False', 'name': 'crow_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._validate_sparse_bsc_tensor_args: [{'is_kwarg_only': 'False', 'name': 'ccol_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'row_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._to_cpu: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._coalesce: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.hspmm: [{'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.hspmm: [{'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.unbind: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.unbind: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions._to_sparse_semi_structured: [{'is_kwarg_only': 'False', 'name': 'dense', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.quantize_per_tensor_dynamic: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType'}, {'is_kwarg_only': 'False', 'name': 'reduce_range', 'simple_type': 'bool'}], + torch._C._VariableFunctions.quantize_per_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.quantize_per_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.quantize_per_tensor: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scales', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_points', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.quantize_per_channel: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scales', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_points', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'axis', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.dequantize: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.dequantize: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.q_scale: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.q_zero_point: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.q_per_channel_scales: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.q_per_channel_zero_points: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.q_per_channel_axis: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.int_repr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._make_per_tensor_quantized_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._make_per_channel_quantized_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'axis', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.fake_quantize_per_tensor_affine: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_min', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_max', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.fake_quantize_per_tensor_affine: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'quant_min', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_max', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._fake_quantize_per_tensor_affine_cachemask_tensor_qparams: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'fake_quant_enabled', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'quant_min', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_max', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._fake_quantize_learnable_per_tensor_affine: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'quant_min', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_max', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.fake_quantize_per_channel_affine: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'axis', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_min', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_max', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._fake_quantize_learnable_per_channel_affine: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'axis', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_min', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_max', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.fused_moving_avg_obs_fake_quant: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'observer_on', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'fake_quant_on', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_min', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_max', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'averaging_const', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'quant_min', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_max', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'ch_axis', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._fused_moving_avg_obs_fq_helper: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'observer_on', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'fake_quant_on', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_min', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'running_max', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'zero_point', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'averaging_const', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'quant_min', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'quant_max', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'ch_axis', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._choose_qparams_per_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._saturate_weight_to_fp16: [{'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.choose_qparams_optimized: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'numel', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'n_bins', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'ratio', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'bit_width', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.meshgrid: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.meshgrid: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'indexing', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.cartesian_prod: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.combinations: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.result_type: [{'is_kwarg_only': 'False', 'name': 'tensor', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.result_type: [{'is_kwarg_only': 'False', 'name': 'tensor', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.result_type: [{'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'tensor', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.result_type: [{'is_kwarg_only': 'False', 'name': 'scalar1', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'scalar2', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.can_cast: [{'is_kwarg_only': 'False', 'name': 'from_', 'simple_type': 'ScalarType'}, {'is_kwarg_only': 'False', 'name': 'to', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.promote_types: [{'is_kwarg_only': 'False', 'name': 'type1', 'simple_type': 'ScalarType'}, {'is_kwarg_only': 'False', 'name': 'type2', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions._lstm_mps: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'params', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}], + torch._C._VariableFunctions.lstm: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'params', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}], + torch._C._VariableFunctions.lstm: [{'is_kwarg_only': 'False', 'name': 'data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch_sizes', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'params', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}], + torch._C._VariableFunctions.gru: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'params', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}], + torch._C._VariableFunctions.gru: [{'is_kwarg_only': 'False', 'name': 'data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch_sizes', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'params', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}], + torch._C._VariableFunctions.rnn_tanh: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'params', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}], + torch._C._VariableFunctions.rnn_tanh: [{'is_kwarg_only': 'False', 'name': 'data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch_sizes', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'params', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}], + torch._C._VariableFunctions.rnn_relu: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'params', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}], + torch._C._VariableFunctions.rnn_relu: [{'is_kwarg_only': 'False', 'name': 'data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch_sizes', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'params', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'has_biases', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'num_layers', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dropout', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'train', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'bidirectional', 'simple_type': 'bool'}], + torch._C._VariableFunctions.lstm_cell: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'w_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_hh', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gru_cell: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_hh', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rnn_tanh_cell: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_hh', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.rnn_relu_cell: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_hh', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.quantized_lstm_cell: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'w_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_ih', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'scale_hh', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'zero_point_ih', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'zero_point_hh', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.quantized_gru_cell: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_ih', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'scale_hh', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'zero_point_ih', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'zero_point_hh', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.quantized_rnn_relu_cell: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_ih', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'scale_hh', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'zero_point_ih', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'zero_point_hh', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.quantized_rnn_tanh_cell: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'hx', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'w_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'packed_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets_ih', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_offsets_hh', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_ih', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'scale_hh', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'zero_point_ih', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'zero_point_hh', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._pack_padded_sequence: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'lengths', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}], + torch._C._VariableFunctions._pad_packed_sequence: [{'is_kwarg_only': 'False', 'name': 'data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch_sizes', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch_first', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'padding_value', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'total_length', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.masked_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.masked_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.masked_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._masked_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.put: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_reduce: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.index_reduce: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.index_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.index_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.index_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.scatter_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.scatter_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.scatter_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.scatter_reduce: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.scatter_reduce: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions.bitwise_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.__and__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.__and__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.__or__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.__or__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.__xor__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.__xor__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.__lshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.__lshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_left_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_left_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_left_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_left_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_left_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.__rshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.__rshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_right_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_right_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bitwise_right_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_right_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.bitwise_right_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addbmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addbmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.diag: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.diag: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cross: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cross: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.triu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.triu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tril: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tril: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.tril_indices: [{'is_kwarg_only': 'False', 'name': 'row', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'col', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.triu_indices: [{'is_kwarg_only': 'False', 'name': 'row', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'col', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.trace: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ne: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.ne: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.ne: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ne: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.not_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.not_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.not_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.not_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.eq: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.eq: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.eq: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.eq: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ge: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.ge: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.ge: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ge: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.greater_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.greater_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.greater_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.greater_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.le: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.le: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.le: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.le: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.less_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.less_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.less_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.less_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.gt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.gt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.greater: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.greater: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.greater: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.greater: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.lt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.lt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.less: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.less: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.less: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.less: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.take: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.take: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.take_along_dim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.take_along_dim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.index_select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.masked_select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.masked_select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nonzero_static: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'size', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.nonzero_static: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'size', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.argwhere: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gather: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gather: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gather: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.gather: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addcmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addcmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addcdiv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.addcdiv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.triangular_solve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.triangular_solve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_check_errors: [{'is_kwarg_only': 'False', 'name': 'info', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'api_name', 'simple_type': 'c10::string_view'}, {'is_kwarg_only': 'True', 'name': 'is_matrix', 'simple_type': 'bool'}], + torch._C._VariableFunctions.svd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.svd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.swapaxes: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'axis0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'axis1', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.swapdims: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.cholesky: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cholesky: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cholesky_solve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cholesky_solve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cholesky_inverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.cholesky_inverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.qr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.qr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.geqrf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.geqrf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.orgqr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.orgqr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ormqr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input3', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ormqr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input3', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._lu_with_info: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lu_solve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'LU_data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'LU_pivots', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lu_solve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'LU_data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'LU_pivots', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lu_unpack: [{'is_kwarg_only': 'False', 'name': 'LU_data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'LU_pivots', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lu_unpack: [{'is_kwarg_only': 'False', 'name': 'LU_data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'LU_pivots', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.multinomial: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'num_samples', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.multinomial: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'num_samples', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.lgamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lgamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.digamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.digamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.polygamma: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.polygamma: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.erfinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.erfinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.i0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.i0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.i0_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.signbit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.signbit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.dist: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atan2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.atan2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arctan2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.arctan2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lerp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.lerp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.lerp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.lerp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.histc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.histc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.histogram: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bins', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.histogram: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bins', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.histogram: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.histogram: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._histogramdd_bin_edges: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bins', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._histogramdd_from_bin_cts: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bins', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._histogramdd_from_bin_tensors: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bins', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.histogramdd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bins', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.histogramdd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bins', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.histogramdd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bins', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.fmod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.fmod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.fmod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fmod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.hypot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.hypot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.igamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.igamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.igammac: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.igammac: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nextafter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nextafter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.remainder: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.remainder: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.remainder: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.remainder: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.remainder: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fmin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fmin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.fmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.maximum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.maximum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.minimum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.minimum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.quantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.quantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.quantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'double'}], + torch._C._VariableFunctions.quantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'double'}], + torch._C._VariableFunctions.nanquantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nanquantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.nanquantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'double'}], + torch._C._VariableFunctions.nanquantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'double'}], + torch._C._VariableFunctions.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'stable', 'simple_type': 'bool?'}], + torch._C._VariableFunctions.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'stable', 'simple_type': 'bool?'}], + torch._C._VariableFunctions.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'stable', 'simple_type': 'bool?'}, {'is_kwarg_only': 'True', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'stable', 'simple_type': 'bool?'}, {'is_kwarg_only': 'True', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.msort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.msort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.argsort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.argsort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'stable', 'simple_type': 'bool'}], + torch._C._VariableFunctions.argsort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'stable', 'simple_type': 'bool'}], + torch._C._VariableFunctions.argsort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch._C._VariableFunctions.topk: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.topk: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.renorm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'maxnorm', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.renorm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'maxnorm', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.float_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.float_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.float_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.float_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.float_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.float_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.normal: [{'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.normal: [{'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.normal: [{'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'std', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.normal: [{'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'std', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.normal: [{'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'std', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.normal: [{'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'std', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.normal: [{'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'std', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.normal: [{'is_kwarg_only': 'False', 'name': 'mean', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'std', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions._amp_foreach_non_finite_check_and_unscale_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'found_inf', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'inv_scale', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._amp_update_scale_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'growth_tracker', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'found_inf', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'scale_growth_factor', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'scale_backoff_factor', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'growth_interval', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._foreach_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_add_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_add_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_add_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_add_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_sub: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_sub: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sub: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_sub_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_sub_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sub_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_mul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_mul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_mul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_mul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_mul_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_mul_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_mul_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_mul_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_div_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_div_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_div_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_div_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_clamp_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_clamp_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_clamp_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_clamp_max_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_clamp_max_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_clamp_max_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_clamp_min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_clamp_min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_clamp_min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_clamp_min_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_clamp_min_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_clamp_min_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_maximum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_maximum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_maximum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_maximum_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_maximum_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_maximum_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_minimum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_minimum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_minimum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_minimum_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalar', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_minimum_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_minimum_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_addcdiv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_addcdiv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_addcdiv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_addcdiv_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_addcdiv_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_addcdiv_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_addcmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_addcmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_addcmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_addcmul_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_addcmul_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_addcmul_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'scalars', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foreach_abs: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_abs_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_acos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_acos_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_asin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_asin_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_atan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_atan_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_ceil: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_ceil_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_cos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_cos_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_cosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_cosh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_erf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_erf_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_erfc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_erfc_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_exp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_exp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_expm1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_expm1_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_floor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_floor_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_frac: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_frac_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_lerp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensors1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'weights', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_lerp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensors1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_lerp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensors1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_lerp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensors1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'weights', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_lerp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensors1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_lerp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'tensors1', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_lgamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_lgamma_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_log: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_log_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_log10: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_log10_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_log1p: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_log1p_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_log2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_log2_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_neg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_neg_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_powsum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_pow_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_pow_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._foreach_pow_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'ScalarList'}], + torch._C._VariableFunctions._foreach_reciprocal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_reciprocal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_round: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_round_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_rsqrt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_rsqrt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sigmoid: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sigmoid_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sign_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sin_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sinh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sqrt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_sqrt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_tan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_tan_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_tanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_tanh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_trunc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_trunc_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_zero_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_clone: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._foreach_copy_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.bucketize: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'boundaries', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bucketize: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'boundaries', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.bucketize: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'boundaries', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.searchsorted: [{'is_kwarg_only': 'False', 'name': 'sorted_sequence', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.searchsorted: [{'is_kwarg_only': 'False', 'name': 'sorted_sequence', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.searchsorted: [{'is_kwarg_only': 'False', 'name': 'sorted_sequence', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions.searchsorted: [{'is_kwarg_only': 'False', 'name': 'sorted_sequence', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}], + torch._C._VariableFunctions._convert_indices_from_coo_to_csr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._convert_indices_from_coo_to_csr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._convert_indices_from_csr_to_coo: [{'is_kwarg_only': 'False', 'name': 'crow_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_indices', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._convert_indices_from_csr_to_coo: [{'is_kwarg_only': 'False', 'name': 'crow_indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'col_indices', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.mkldnn_adaptive_avg_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._VariableFunctions.mkldnn_adaptive_avg_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._VariableFunctions._adaptive_avg_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._VariableFunctions._adaptive_avg_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._VariableFunctions.column_stack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.column_stack: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions.isfinite: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isinf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isposinf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isposinf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isneginf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.isneginf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._add_batch_dim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'level', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._remove_batch_dim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'level', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'batch_size', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'out_dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._linalg_det: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_det: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.det: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_slogdet: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_slogdet: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.slogdet: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.slogdet: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.logdet: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_eigh: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_eigh: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.inverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.inverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.inner: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.inner: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.outer: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.outer: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ger: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ger: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_svd: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_svd: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_solve_ex: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._linalg_solve_ex: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._test_serialization_subcmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._test_parallel_materialize: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'num_parallel', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._test_autograd_multiple_dispatch: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._test_autograd_multiple_dispatch: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b', 'simple_type': 'bool'}], + torch._C._VariableFunctions._test_autograd_multiple_dispatch_view: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._test_autograd_multiple_dispatch_view_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._test_autograd_multiple_dispatch_view_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.segment_reduce: [{'is_kwarg_only': 'False', 'name': 'data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch._C._VariableFunctions._nested_tensor_from_tensor_list: [{'is_kwarg_only': 'False', 'name': 'list', 'simple_type': 'TensorList'}], + torch._C._VariableFunctions._fw_primal_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'level', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._fw_primal_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'level', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._make_dual_copy: [{'is_kwarg_only': 'False', 'name': 'primal', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tangent', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'level', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._make_dual_copy: [{'is_kwarg_only': 'False', 'name': 'primal', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tangent', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'level', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.view_as_real_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.view_as_real_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.view_as_complex_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.view_as_complex_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._conj_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._conj_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._neg_view_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._neg_view_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.as_strided_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.as_strided_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions._sparse_broadcast_to_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._sparse_broadcast_to_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.diagonal_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.diagonal_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.expand_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.expand_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.permute_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.permute_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions._reshape_alias_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions._reshape_alias_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.select_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.select_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.detach_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.detach_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.slice_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.slice_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.split_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_size', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.split_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_size', 'simple_type': 'SymInt'}], + torch._C._VariableFunctions.split_with_sizes_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_sizes', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.split_with_sizes_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_sizes', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.squeeze_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.squeeze_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.squeeze_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.squeeze_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.squeeze_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.squeeze_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch._C._VariableFunctions.t_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.t_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.transpose_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.transpose_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.unsqueeze_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.unsqueeze_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._values_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._values_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.values_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.values_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.crow_indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.crow_indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.col_indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.col_indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ccol_indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.ccol_indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.row_indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.row_indices_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.unbind_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.unbind_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.view_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.view_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.view_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch._C._VariableFunctions.view_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch._C._VariableFunctions.unfold_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dimension', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'step', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.unfold_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dimension', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'step', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions.alias_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions.alias_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_from_padded_tensor: [{'is_kwarg_only': 'False', 'name': 'padded', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'offsets', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dummy', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._nested_tensor_softmax_with_shape: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._safe_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._transformer_encoder_layer_fwd: [{'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'embed_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'num_heads', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'qkv_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'qkv_bias', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'proj_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'proj_bias', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'use_gelu', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'norm_first', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'norm_weight_1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'norm_bias_1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'norm_weight_2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'norm_bias_2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'ffn_weight_1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'ffn_bias_1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'ffn_weight_2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'ffn_bias_2', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._native_multi_head_attention: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'embed_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'num_head', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'qkv_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'qkv_bias', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'proj_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'proj_bias', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._fused_sdp_choice: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._scaled_dot_product_attention_math: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._scaled_dot_product_attention_math_for_mps: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._scaled_dot_product_flash_attention: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._scaled_dot_product_flash_attention: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q_descale', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'k_descale', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'v_descale', 'simple_type': 'Tensor?'}], + torch._C._VariableFunctions._scaled_dot_product_flash_attention_for_cpu: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._scaled_dot_product_efficient_attention: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'attn_bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'compute_log_sumexp', 'simple_type': 'bool'}], + torch._C._VariableFunctions._scaled_dot_product_cudnn_attention: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'attn_bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'compute_log_sumexp', 'simple_type': 'bool'}], + torch._C._VariableFunctions._flash_attention_forward_no_dropout_inplace: [{'is_kwarg_only': 'False', 'name': 'out', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'cum_seq_q', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'cum_seq_k', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'max_q', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'max_k', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'dropout_p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'is_causal', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'return_debug_mask', 'simple_type': 'bool'}], + torch._C._VariableFunctions._triton_scaled_dot_attention: [{'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'v', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._fill_mem_eff_dropout_mask_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dropout_p', 'simple_type': 'double'}, {'is_kwarg_only': 'False', 'name': 'seed', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'offset', 'simple_type': 'int64_t'}], + torch._C._VariableFunctions._triton_multi_head_attention: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'embed_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'num_head', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'qkv_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'qkv_bias', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'proj_weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'proj_bias', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._foobar: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._VariableFunctions._fused_adam_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'grads', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exp_avgs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exp_avg_sqs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'max_exp_avg_sqs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'state_steps', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'lr', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'beta1', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'beta2', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'weight_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'amsgrad', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'maximize', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fused_adam_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'grads', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exp_avgs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exp_avg_sqs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'max_exp_avg_sqs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'state_steps', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'lr', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'beta1', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'beta2', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'weight_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'amsgrad', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'maximize', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fused_adamw_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'grads', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exp_avgs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exp_avg_sqs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'max_exp_avg_sqs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'state_steps', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'lr', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'beta1', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'beta2', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'weight_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'amsgrad', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'maximize', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fused_adamw_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'grads', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exp_avgs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'exp_avg_sqs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'max_exp_avg_sqs', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'state_steps', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'lr', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'beta1', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'beta2', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'weight_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'amsgrad', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'maximize', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fused_sgd_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'grads', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'momentum_buffer_list', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'weight_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'lr', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'dampening', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'nesterov', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'maximize', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'is_first_step', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fused_sgd_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'grads', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'momentum_buffer_list', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'weight_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'momentum', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'lr', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'dampening', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'nesterov', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'maximize', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'is_first_step', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fused_adagrad_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'grads', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'state_sums', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'state_steps', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'lr', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'lr_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'weight_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'maximize', 'simple_type': 'bool'}], + torch._C._VariableFunctions._fused_adagrad_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'grads', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'state_sums', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'False', 'name': 'state_steps', 'simple_type': 'TensorList'}, {'is_kwarg_only': 'True', 'name': 'lr', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'lr_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'weight_decay', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'eps', 'simple_type': 'double'}, {'is_kwarg_only': 'True', 'name': 'maximize', 'simple_type': 'bool'}], + torch._C._VariableFunctions._propagate_xla_data: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output', 'simple_type': 'Tensor'}], + torch._C._nn.binary_cross_entropy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.binary_cross_entropy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.linear: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._nn.linear: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._nn.mkldnn_linear: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch._C._nn.relu6: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.relu6_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.gelu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.gelu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.gelu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.silu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.silu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.silu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.mish: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.mish: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.mish_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.one_hot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.mkldnn_reorder_conv2d_weight: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.mkldnn_reorder_conv3d_weight: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.cross_entropy_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.mse_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.mse_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.l1_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.multi_margin_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.multi_margin_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.multilabel_margin_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.multilabel_margin_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.nll_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.nll_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.nll_loss_nd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.nll_loss2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.nll_loss2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.smooth_l1_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.smooth_l1_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.huber_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.huber_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.soft_margin_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.soft_margin_loss: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'target', 'simple_type': 'Tensor'}], + torch._C._nn.elu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.elu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.elu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.glu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.glu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.hardsigmoid: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.hardsigmoid: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.hardsigmoid_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.hardtanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.hardtanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.hardtanh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.hardswish: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.hardswish: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.hardswish_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.leaky_relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.leaky_relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.leaky_relu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.log_sigmoid: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.log_sigmoid: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.rrelu_with_noise: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'noise', 'simple_type': 'Tensor'}], + torch._C._nn.rrelu_with_noise: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'noise', 'simple_type': 'Tensor'}], + torch._C._nn.rrelu_with_noise_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'noise', 'simple_type': 'Tensor'}], + torch._C._nn.softplus: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.softplus: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.softshrink: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.softshrink: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.adaptive_avg_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.adaptive_avg_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.adaptive_avg_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.adaptive_avg_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.adaptive_max_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn.adaptive_max_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn.adaptive_max_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._nn.adaptive_max_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._nn.avg_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn.avg_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn.avg_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._nn.avg_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._nn.fractional_max_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'random_samples', 'simple_type': 'Tensor'}], + torch._C._nn.fractional_max_pool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'random_samples', 'simple_type': 'Tensor'}], + torch._C._nn.fractional_max_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'random_samples', 'simple_type': 'Tensor'}], + torch._C._nn.fractional_max_pool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'IntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'random_samples', 'simple_type': 'Tensor'}], + torch._C._nn.max_pool2d_with_indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn.max_pool2d_with_indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn.max_pool3d_with_indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._nn.max_pool3d_with_indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._nn.max_unpool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.max_unpool2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.max_unpool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'IntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._nn.max_unpool3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'IntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'IntArrayRef', 'size': 3}], + torch._C._nn.reflection_pad1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.reflection_pad1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.reflection_pad2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 4}], + torch._C._nn.reflection_pad2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 4}], + torch._C._nn.reflection_pad3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 6}], + torch._C._nn.reflection_pad3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 6}], + torch._C._nn.replication_pad1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.replication_pad1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.replication_pad2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 4}], + torch._C._nn.replication_pad2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 4}], + torch._C._nn.replication_pad3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 6}], + torch._C._nn.replication_pad3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 6}], + torch._C._nn._pad_circular: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'pad', 'simple_type': 'SymIntArrayRef'}], + torch._C._nn._pad_enum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'pad', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'mode', 'simple_type': 'int64_t'}], + torch._C._nn.pad: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'pad', 'simple_type': 'SymIntArrayRef'}], + torch._C._nn.upsample_linear1d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn.upsample_linear1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 1}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn.upsample_linear1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 1}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn.upsample_bilinear2d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn.upsample_bilinear2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn.upsample_bilinear2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn._upsample_bilinear2d_aa: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn._upsample_bilinear2d_aa: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn._upsample_bilinear2d_aa: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn.upsample_trilinear3d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn.upsample_trilinear3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn.upsample_trilinear3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn.upsample_bicubic2d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn.upsample_bicubic2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn.upsample_bicubic2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn._upsample_bicubic2d_aa: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn._upsample_bicubic2d_aa: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn._upsample_bicubic2d_aa: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn._upsample_lanczos2d_aa: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn._upsample_lanczos2d_aa: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn._upsample_lanczos2d_aa: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'align_corners', 'simple_type': 'bool'}], + torch._C._nn.upsample_nearest1d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn.upsample_nearest1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 1}], + torch._C._nn.upsample_nearest1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 1}], + torch._C._nn._upsample_nearest_exact1d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn._upsample_nearest_exact1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 1}], + torch._C._nn._upsample_nearest_exact1d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 1}], + torch._C._nn.upsample_nearest2d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn.upsample_nearest2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.upsample_nearest2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn._upsample_nearest_exact2d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn._upsample_nearest_exact2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn._upsample_nearest_exact2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.upsample_nearest3d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn.upsample_nearest3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.upsample_nearest3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn._upsample_nearest_exact3d: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef?'}, {'is_kwarg_only': 'False', 'name': 'scale_factors', 'simple_type': 'ArrayRef?'}], + torch._C._nn._upsample_nearest_exact3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn._upsample_nearest_exact3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.slow_conv_transpose2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.slow_conv_transpose2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.slow_conv_transpose3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.slow_conv_transpose3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.thnn_conv2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.thnn_conv2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn._conv_depthwise2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn._conv_depthwise2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.conv_depthwise3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'bias', 'simple_type': 'Tensor?'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'SymIntArrayRef', 'size': 3}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.slow_conv3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.slow_conv3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.slow_conv_dilated2d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 2}], + torch._C._nn.slow_conv_dilated3d: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'SymIntArrayRef', 'size': 3}], + torch._C._nn.col2im: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn.col2im: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'output_size', 'simple_type': 'SymIntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn.im2col: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn.im2col: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'kernel_size', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'dilation', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'IntArrayRef', 'size': 2}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'IntArrayRef', 'size': 2}], + torch._C._nn._test_optional_intlist: [{'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'addends', 'simple_type': 'IntArrayRef?'}], + torch._C._nn._test_optional_filled_intlist: [{'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'addends', 'simple_type': 'IntArrayRef?', 'size': 2}], + torch._C._nn._test_optional_floatlist: [{'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'addends', 'simple_type': 'ArrayRef?'}], + torch._C._nn._test_string_default: [{'is_kwarg_only': 'False', 'name': 'dummy', 'simple_type': 'Tensor'}], + torch._C._nn._test_ambiguous_defaults: [{'is_kwarg_only': 'False', 'name': 'dummy', 'simple_type': 'Tensor'}], + torch._C._nn._test_ambiguous_defaults: [{'is_kwarg_only': 'False', 'name': 'dummy', 'simple_type': 'Tensor'}], + torch._C._nn._test_warn_in_autograd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._nn.pad_sequence: [{'is_kwarg_only': 'False', 'name': 'sequences', 'simple_type': 'TensorList'}], + torch._C._nn.flatten_dense_tensors: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._nn.unflatten_dense_tensors: [{'is_kwarg_only': 'False', 'name': 'flat', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._nn.scaled_dot_product_attention: [{'is_kwarg_only': 'False', 'name': 'query', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_diagonal: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_solve_triangular: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'upper', 'simple_type': 'bool'}], + torch._C._linalg.linalg_solve_triangular: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'upper', 'simple_type': 'bool'}], + torch._C._linalg.linalg_vander: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_cholesky_ex: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_cholesky_ex: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_cholesky: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_cholesky: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_cross: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_cross: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lu_factor: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lu_factor: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lu_factor_ex: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lu_factor_ex: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lu: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lu: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lu_solve: [{'is_kwarg_only': 'False', 'name': 'LU', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'pivots', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lu_solve: [{'is_kwarg_only': 'False', 'name': 'LU', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'pivots', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_det: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_det: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_ldl_factor_ex: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_ldl_factor_ex: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_ldl_factor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_ldl_factor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_ldl_solve: [{'is_kwarg_only': 'False', 'name': 'LD', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'pivots', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_ldl_solve: [{'is_kwarg_only': 'False', 'name': 'LD', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'pivots', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lstsq: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_lstsq: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'b', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_vecdot: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'y', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_vecdot: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'y', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matrix_exp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_slogdet: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_slogdet: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_eig: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_eig: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg._linalg_eigvals: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_eigvals: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_eigvals: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_eigh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_eigh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_eigvalsh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_eigvalsh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_householder_product: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tau', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_householder_product: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tau', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_inv_ex: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_inv_ex: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_inv: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_inv: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'ord', 'simple_type': 'c10::string_view'}], + torch._C._linalg.linalg_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'ord', 'simple_type': 'c10::string_view'}], + torch._C._linalg.linalg_vector_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_vector_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg__powsum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matrix_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'ord', 'simple_type': 'Scalar'}], + torch._C._linalg.linalg_matrix_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'ord', 'simple_type': 'Scalar'}], + torch._C._linalg.linalg_matrix_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matrix_norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_svd: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_svd: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_svdvals: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_svdvals: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_cond: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_cond: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_cond: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'c10::string_view'}], + torch._C._linalg.linalg_cond: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'c10::string_view'}], + torch._C._linalg.linalg_pinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_pinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_pinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_pinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_pinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'rcond', 'simple_type': 'double'}], + torch._C._linalg.linalg_pinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'rcond', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_pinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'rcond', 'simple_type': 'double'}], + torch._C._linalg.linalg_pinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'rcond', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_solve_ex: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_solve_ex: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_solve: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_solve: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'B', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_tensorinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_tensorinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_tensorsolve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_tensorsolve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_qr: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_qr: [{'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matrix_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch._C._linalg.linalg_matrix_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch._C._linalg.linalg_matrix_rank: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matrix_rank: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matrix_rank: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matrix_rank: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matrix_rank: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tol', 'simple_type': 'double'}], + torch._C._linalg.linalg_matrix_rank: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tol', 'simple_type': 'double'}], + torch._C._linalg.linalg_matrix_rank: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tol', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_matrix_rank: [{'is_kwarg_only': 'False', 'name': 'input', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tol', 'simple_type': 'Tensor'}], + torch._C._linalg.linalg_multi_dot: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._linalg.linalg_multi_dot: [{'is_kwarg_only': 'False', 'name': 'tensors', 'simple_type': 'TensorList'}], + torch._C._special.special_entr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_entr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_ndtri: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_ndtri: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_log_ndtr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_log_ndtr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_expm1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_expm1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_exp2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_exp2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_psi: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_psi: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_digamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_digamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_gammaln: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_gammaln: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_erf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_erf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_erfc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_erfc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_erfcx: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_erfcx: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_erfinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_erfinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_ndtr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_ndtr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_xlog1py: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_xlog1py: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_xlog1py: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._special.special_xlog1py: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_xlog1py: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_xlog1py: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._special.special_xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._special.special_xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._special.special_zeta: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_zeta: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_zeta: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._special.special_zeta: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_zeta: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_zeta: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch._C._special.special_i0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_i0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_i0e: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_i0e: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_i1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_i1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_i1e: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_i1e: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_logit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_logit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_polygamma: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_polygamma: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_logsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._special.special_logsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch._C._special.special_expit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_expit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_sinc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_sinc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_round: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_round: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_log1p: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_log1p: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_log_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._special.special_gammainc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_gammainc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_gammaincc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_gammaincc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch._C._special.special_multigammaln: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'int64_t'}], + torch._C._special.special_multigammaln: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'int64_t'}], + torch._C._special.special_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch._C._special.special_airy_ai: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._special.special_airy_ai: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._special.special_bessel_j0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_bessel_j0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_bessel_j1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_bessel_j1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_bessel_y0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_bessel_y0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_bessel_y1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_bessel_y1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_hermite_polynomial_h: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_hermite_polynomial_h: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_hermite_polynomial_h: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_hermite_polynomial_h: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_hermite_polynomial_h: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_hermite_polynomial_h: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_hermite_polynomial_he: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_hermite_polynomial_he: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_hermite_polynomial_he: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_hermite_polynomial_he: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_hermite_polynomial_he: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_hermite_polynomial_he: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_laguerre_polynomial_l: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_laguerre_polynomial_l: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_laguerre_polynomial_l: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_laguerre_polynomial_l: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_laguerre_polynomial_l: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_laguerre_polynomial_l: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_legendre_polynomial_p: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_legendre_polynomial_p: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_legendre_polynomial_p: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_legendre_polynomial_p: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_legendre_polynomial_p: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_legendre_polynomial_p: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_modified_bessel_i0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_modified_bessel_i0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_modified_bessel_i1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_modified_bessel_i1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_modified_bessel_k0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_modified_bessel_k0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_modified_bessel_k1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_modified_bessel_k1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._special.special_scaled_modified_bessel_k0: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._special.special_scaled_modified_bessel_k0: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._special.special_scaled_modified_bessel_k1: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._special.special_scaled_modified_bessel_k1: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_shifted_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_t: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_shifted_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_shifted_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_u: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_shifted_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_shifted_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_v: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_shifted_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_shifted_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Tensor'}], + torch._C._special.special_shifted_chebyshev_polynomial_w: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'Scalar'}], + torch._C._special.special_spherical_bessel_j0: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._special.special_spherical_bessel_j0: [{'is_kwarg_only': 'False', 'name': 'x', 'simple_type': 'Tensor'}], + torch._C._fft.fft_fft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_fft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ifft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ifft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_rfft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_rfft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_irfft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_irfft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_hfft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_hfft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ihfft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ihfft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_fft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_fft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ifft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ifft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_rfft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_rfft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_irfft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_irfft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_hfft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_hfft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ihfft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ihfft2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_fftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_fftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ifftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ifftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_rfftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_rfftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_irfftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_irfftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_hfftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_hfftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ihfftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ihfftn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_fftfreq: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch._C._fft.fft_fftfreq: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch._C._fft.fft_rfftfreq: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch._C._fft.fft_rfftfreq: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch._C._fft.fft_fftshift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch._C._fft.fft_ifftshift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.retain_grad: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.rename_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch.Tensor.rename: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'names', 'simple_type': 'DimnameList?'}], + torch.Tensor.align_to: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'names', 'simple_type': 'DimnameList'}], + torch.Tensor.align_to: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'order', 'simple_type': 'DimnameList'}, {'is_kwarg_only': 'False', 'name': 'ellipsis_idx', 'simple_type': 'int64_t'}], + torch.Tensor.align_as: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.refine_names: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'names', 'simple_type': 'DimnameList'}], + torch.Tensor.abs: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.abs_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.absolute: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.absolute_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.angle: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sgn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sgn_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.chalf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._conj: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.conj: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._conj_physical: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.conj_physical: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.conj_physical_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.resolve_conj: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.resolve_neg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._neg_view: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.acos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.acos_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arccos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arccos_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.add_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.addmv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec', 'simple_type': 'Tensor'}], + torch.Tensor.addmv_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec', 'simple_type': 'Tensor'}], + torch.Tensor.addr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch.Tensor.addr_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch.Tensor._is_all_true: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._is_any_true: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.all: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.allclose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.any: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.argmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.argmin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.acosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.acosh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arccosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arccosh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.asinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.asinh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arcsinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arcsinh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.atanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.atanh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arctanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arctanh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.as_strided: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.as_strided_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.asin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.asin_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arcsin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arcsin_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.atan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.atan_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arctan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.arctan_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.baddbmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}], + torch.Tensor.baddbmm_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}], + torch.Tensor.bernoulli: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.bernoulli: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}], + torch.Tensor.bernoulli_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Tensor'}], + torch.Tensor.bernoulli_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.bincount: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_not: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_not_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.copysign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.copysign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.copysign_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.copysign_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor._lazy_clone: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.logical_not: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.logical_not_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.logical_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.logical_xor_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.logical_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.logical_and_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.logical_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.logical_or_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch.Tensor.broadcast_to: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.ceil: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.ceil_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.unsafe_chunk: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'chunks', 'simple_type': 'int64_t'}], + torch.Tensor.chunk: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'chunks', 'simple_type': 'int64_t'}], + torch.Tensor.tensor_split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sections', 'simple_type': 'SymInt'}], + torch.Tensor.tensor_split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.tensor_split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor_indices_or_sections', 'simple_type': 'Tensor'}], + torch.Tensor.clamp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.clamp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.clamp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.clamp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.clamp_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Scalar'}], + torch.Tensor.clamp_max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Tensor'}], + torch.Tensor.clamp_max_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Scalar'}], + torch.Tensor.clamp_max_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'max', 'simple_type': 'Tensor'}], + torch.Tensor.clamp_min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Scalar'}], + torch.Tensor.clamp_min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Tensor'}], + torch.Tensor.clamp_min_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Scalar'}], + torch.Tensor.clamp_min_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'min', 'simple_type': 'Tensor'}], + torch.Tensor.clip: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.clip: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.clip_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.clip_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.cos: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.cos_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.cosh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.cosh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.count_nonzero: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch.Tensor.count_nonzero: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.cov: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.corrcoef: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.cummax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.cummax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.cummin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.cummin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.cumprod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.cumprod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.cumprod_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.cumprod_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.cumsum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.cumsum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.cumsum_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.cumsum_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.diag_embed: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.diagflat: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.diagonal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.diagonal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.fill_diagonal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'fill_value', 'simple_type': 'Scalar'}], + torch.Tensor.diff: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch.Tensor.div: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch.Tensor.div_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.div_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch.Tensor.div_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch.Tensor.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch.Tensor.divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch.Tensor.divide_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.divide_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.divide_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch.Tensor.divide_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'rounding_mode', 'simple_type': 'c10::string_view?'}], + torch.Tensor.true_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.true_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.true_divide_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.true_divide_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.dot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor', 'simple_type': 'Tensor'}], + torch.Tensor.vdot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.new_empty: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.new_empty_strided: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.new_full: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'fill_value', 'simple_type': 'Scalar'}], + torch.Tensor.new_zeros: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.new_ones: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.resize_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.erf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.erf_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.erfc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.erfc_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.exp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.exp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.exp2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.exp2_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.expm1: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.expm1_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.expand: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.expand_as: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.flatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.flatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'start_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'end_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'out_dim', 'simple_type': 'Dimname'}], + torch.Tensor.flatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'start_dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'end_dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'out_dim', 'simple_type': 'Dimname'}], + torch.Tensor.flatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'DimnameList'}, {'is_kwarg_only': 'False', 'name': 'out_dim', 'simple_type': 'Dimname'}], + torch.Tensor.unflatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'sizes', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.unflatten: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'sizes', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'names', 'simple_type': 'DimnameList'}], + torch.Tensor.fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch.Tensor.floor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.floor_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.floor_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.floor_divide: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.floor_divide_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.floor_divide_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.frac: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.frac_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.gcd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.gcd_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.lcm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.lcm_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.index_copy_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.index_copy_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.index_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.index_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.index_put_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'c10::List<::std::optional>'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch.Tensor.index_put: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'c10::List<::std::optional>'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch.Tensor.isclose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.isnan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_distributed: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_floating_point: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_complex: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_conj: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._is_zerotensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_neg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.isreal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_nonzero: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_same_size: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.is_signed: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_inference: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.kron: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.kthvalue: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'SymInt'}], + torch.Tensor.kthvalue: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.nan_to_num: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.nan_to_num_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.ldexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.ldexp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.log: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.log_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.log10: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.log10_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.log1p: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.log1p_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.log2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.log2_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.logaddexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.logaddexp2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.xlogy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.xlogy_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.xlogy_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.log_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.log_softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.logcumsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.logcumsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.logsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch.Tensor.logsumexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch.Tensor.matmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.matrix_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch.Tensor.matrix_exp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.aminmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.max: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.amax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch.Tensor.mean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch.Tensor.nanmean: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.median: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.median: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.median: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.nanmedian: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.nanmedian: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.nanmedian: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.min: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.amin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.mm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch.Tensor.mode: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.mode: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.mul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.mul_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.multiply: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.multiply: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.multiply_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.multiply_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.mv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec', 'simple_type': 'Tensor'}], + torch.Tensor.mvlgamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'int64_t'}], + torch.Tensor.mvlgamma_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'int64_t'}], + torch.Tensor.narrow_copy: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'length', 'simple_type': 'SymInt'}], + torch.Tensor.narrow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'SymInt'}, {'is_kwarg_only': 'False', 'name': 'length', 'simple_type': 'SymInt'}], + torch.Tensor.narrow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'start', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'length', 'simple_type': 'SymInt'}], + torch.Tensor.permute: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'IntArrayRef'}], + torch.Tensor.movedim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'destination', 'simple_type': 'IntArrayRef'}], + torch.Tensor.movedim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'destination', 'simple_type': 'int64_t'}], + torch.Tensor.moveaxis: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'destination', 'simple_type': 'IntArrayRef'}], + torch.Tensor.moveaxis: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'destination', 'simple_type': 'int64_t'}], + torch.Tensor.adjoint: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_pinned: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.pin_memory: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.pinverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.rad2deg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.rad2deg_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.deg2rad: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.deg2rad_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.ravel: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.reciprocal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.reciprocal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.neg: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.neg_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.negative: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.negative_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.repeat: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'repeats', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.repeat_interleave: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'repeats', 'simple_type': 'Tensor'}], + torch.Tensor.repeat_interleave: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'repeats', 'simple_type': 'SymInt'}], + torch.Tensor.reshape: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'shape', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.reshape_as: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.round: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.round: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.round_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.round_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.relu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.relu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.prelu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch.Tensor.hardshrink: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.rsqrt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.rsqrt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'int64_t'}], + torch.Tensor.select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'SymInt'}], + torch.Tensor.sigmoid: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sigmoid_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.logit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.logit_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sin_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sinc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sinc_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sinh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sinh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.detach: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.detach_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.slice_inverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.slice_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.select_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'SymInt'}], + torch.Tensor.diagonal_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.as_strided_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'stride', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.smm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch.Tensor.softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.softmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.unsafe_split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_size', 'simple_type': 'SymInt'}], + torch.Tensor.split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_size', 'simple_type': 'SymInt'}], + torch.Tensor.split: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_size', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.unsafe_split_with_sizes: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_sizes', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.split_with_sizes: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'split_sizes', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.hsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sections', 'simple_type': 'int64_t'}], + torch.Tensor.hsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'IntArrayRef'}], + torch.Tensor.vsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sections', 'simple_type': 'int64_t'}], + torch.Tensor.vsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'IntArrayRef'}], + torch.Tensor.dsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sections', 'simple_type': 'int64_t'}], + torch.Tensor.dsplit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'IntArrayRef'}], + torch.Tensor.squeeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.squeeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.squeeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.squeeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch.Tensor.squeeze_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.squeeze_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.squeeze_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef'}], + torch.Tensor.squeeze_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.sspaddmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch.Tensor.stft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n_fft', 'simple_type': 'int64_t'}], + torch.Tensor.stft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n_fft', 'simple_type': 'int64_t'}], + torch.Tensor.istft: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n_fft', 'simple_type': 'int64_t'}], + torch.Tensor.sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch.Tensor.sum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch.Tensor.nansum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.hash_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sum_to_size: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.sqrt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sqrt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.square: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.square_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch.Tensor.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch.Tensor.std: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch.Tensor.prod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.prod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.prod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.t: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.t_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.tan: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.tan_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.tanh: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.tanh_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.tile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.transpose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch.Tensor.transpose: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'Dimname'}], + torch.Tensor.transpose_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch.Tensor.flip: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dims', 'simple_type': 'IntArrayRef'}], + torch.Tensor.fliplr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.flipud: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.roll: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'shifts', 'simple_type': 'SymIntArrayRef', 'size': 1}], + torch.Tensor.rot90: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._nested_tensor_size: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._nested_tensor_strides: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._nested_tensor_storage_offsets: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.trunc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.trunc_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.fix: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.fix_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.type_as: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.unsqueeze: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.unsqueeze_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}], + torch.Tensor.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef?', 'size': 1}], + torch.Tensor.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch.Tensor.var: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch.Tensor.view_as: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.where: [{'is_kwarg_only': 'False', 'name': 'condition', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.where: [{'is_kwarg_only': 'False', 'name': 'condition', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor._philox_normal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}], + torch.Tensor._philox_uniform_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'key', 'simple_type': 'Tensor'}], + torch.Tensor.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch.Tensor.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}, {'is_kwarg_only': 'False', 'name': 'keepdim', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch.Tensor.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'IntArrayRef', 'size': 1}], + torch.Tensor.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}, {'is_kwarg_only': 'False', 'name': 'keepdim', 'simple_type': 'bool'}, {'is_kwarg_only': 'True', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch.Tensor.norm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar?'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'DimnameList', 'size': 1}], + torch.Tensor.frexp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.clone: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.positive: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.resize_as_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'the_template', 'simple_type': 'Tensor'}], + torch.Tensor.resize_as_sparse_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'the_template', 'simple_type': 'Tensor'}], + torch.Tensor.zero_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sub: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.sub_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.subtract: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.subtract: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.subtract_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.subtract_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.heaviside: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch.Tensor.heaviside_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'values', 'simple_type': 'Tensor'}], + torch.Tensor.addmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch.Tensor.addmm_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch.Tensor._addmm_activation: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mat2', 'simple_type': 'Tensor'}], + torch.Tensor.sparse_resize_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'sparse_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dense_dim', 'simple_type': 'int64_t'}], + torch.Tensor.sparse_resize_and_clear_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'IntArrayRef'}, {'is_kwarg_only': 'False', 'name': 'sparse_dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dense_dim', 'simple_type': 'int64_t'}], + torch.Tensor.sparse_mask: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}], + torch.Tensor._sparse_mask_projection: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}], + torch.Tensor.to_dense: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._to_dense: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sparse_dim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._dimI: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.dense_dim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._dimV: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._nnz: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.coalesce: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.is_coalesced: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._values: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._coalesced_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'coalesced', 'simple_type': 'bool'}], + torch.Tensor.indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.values: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.crow_indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.col_indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.ccol_indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.row_indices: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.unbind: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.unbind: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.to_sparse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sparse_dim', 'simple_type': 'int64_t'}], + torch.Tensor.to_sparse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._to_sparse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'sparse_dim', 'simple_type': 'int64_t'}], + torch.Tensor._to_sparse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.to_sparse_csr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._to_sparse_csr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.to_sparse_csc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._to_sparse_csc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.to_sparse_bsr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'blocksize', 'simple_type': 'IntArrayRef', 'size': 2}], + torch.Tensor._to_sparse_bsr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'blocksize', 'simple_type': 'IntArrayRef', 'size': 2}], + torch.Tensor.to_sparse_bsc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'blocksize', 'simple_type': 'IntArrayRef', 'size': 2}], + torch.Tensor._to_sparse_bsc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'blocksize', 'simple_type': 'IntArrayRef', 'size': 2}], + torch.Tensor.to_mkldnn: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.dequantize: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.q_scale: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.q_zero_point: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.q_per_channel_scales: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.q_per_channel_zero_points: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.q_per_channel_axis: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.int_repr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.qscheme: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor._autocast_to_reduced_precision: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'cuda_enabled', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'cpu_enabled', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'cuda_dtype', 'simple_type': 'ScalarType'}, {'is_kwarg_only': 'False', 'name': 'cpu_dtype', 'simple_type': 'ScalarType'}], + torch.Tensor._autocast_to_full_precision: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'cuda_enabled', 'simple_type': 'bool'}, {'is_kwarg_only': 'False', 'name': 'cpu_enabled', 'simple_type': 'bool'}], + torch.Tensor.is_set_to: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor', 'simple_type': 'Tensor'}], + torch.Tensor.masked_fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.masked_fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch.Tensor.masked_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.masked_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch.Tensor.masked_scatter_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.masked_scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.view: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'SymIntArrayRef'}], + torch.Tensor.view: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dtype', 'simple_type': 'ScalarType'}], + torch.Tensor.put_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.put: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.index_add_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.index_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.index_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}], + torch.Tensor.index_reduce_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch.Tensor.index_reduce: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'source', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch.Tensor.index_fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.index_fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch.Tensor.index_fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.index_fill_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch.Tensor.index_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.index_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch.Tensor.index_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.index_fill: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Tensor'}], + torch.Tensor.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch.Tensor.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'True', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch.Tensor.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.scatter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.scatter_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.scatter_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.scatter_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.scatter_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'value', 'simple_type': 'Scalar'}], + torch.Tensor.scatter_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.scatter_add: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.scatter_add_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}], + torch.Tensor.scatter_reduce: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch.Tensor.scatter_reduce_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'src', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'reduce', 'simple_type': 'c10::string_view'}], + torch.Tensor.eq_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.eq_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.bitwise_and: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_and_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.bitwise_and_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.__and__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__and__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.__iand__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__iand__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.bitwise_or: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_or_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.bitwise_or_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.__or__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__or__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.__ior__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__ior__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.bitwise_xor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_xor_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.bitwise_xor_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.__xor__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__xor__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.__ixor__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__ixor__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.__lshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__lshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.__ilshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__ilshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_left_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_left_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.bitwise_left_shift_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_left_shift_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__rshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__rshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.__irshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.__irshift__: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_right_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_right_shift: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.bitwise_right_shift_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.bitwise_right_shift_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.tril_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.triu_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.digamma_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.lerp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Scalar'}], + torch.Tensor.lerp_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch.Tensor.addbmm_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}], + torch.Tensor.addbmm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'batch2', 'simple_type': 'Tensor'}], + torch.Tensor.random_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'from', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'to', 'simple_type': 'int64_t?'}], + torch.Tensor.random_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'to', 'simple_type': 'int64_t'}], + torch.Tensor.random_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.uniform_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.cauchy_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.log_normal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.exponential_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.geometric_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'double'}], + torch.Tensor.diag: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.cross: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.triu: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.tril: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.trace: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.ne: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.ne: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.ne_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.ne_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.not_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.not_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.not_equal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.not_equal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.eq: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.eq: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.ge: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.ge: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.ge_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.ge_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.greater_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.greater_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.greater_equal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.greater_equal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.le: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.le: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.le_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.le_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.less_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.less_equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.less_equal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.less_equal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.gt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.gt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.gt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.gt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.greater: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.greater: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.greater_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.greater_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.lt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.lt: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.lt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.lt_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.less: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.less: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.less_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.less_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.take: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch.Tensor.take_along_dim: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'indices', 'simple_type': 'Tensor'}], + torch.Tensor.index_select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch.Tensor.index_select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch.Tensor.masked_select: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'mask', 'simple_type': 'Tensor'}], + torch.Tensor.nonzero_static: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'size', 'simple_type': 'SymInt'}], + torch.Tensor.argwhere: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.gather: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch.Tensor.gather: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}, {'is_kwarg_only': 'False', 'name': 'index', 'simple_type': 'Tensor'}], + torch.Tensor.addcmul: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'Tensor'}], + torch.Tensor.addcmul_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'Tensor'}], + torch.Tensor.addcdiv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'Tensor'}], + torch.Tensor.addcdiv_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor1', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'tensor2', 'simple_type': 'Tensor'}], + torch.Tensor.triangular_solve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'A', 'simple_type': 'Tensor'}], + torch.Tensor.svd: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.swapaxes: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'axis0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'axis1', 'simple_type': 'int64_t'}], + torch.Tensor.swapaxes_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'axis0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'axis1', 'simple_type': 'int64_t'}], + torch.Tensor.swapdims: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch.Tensor.swapdims_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim0', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'dim1', 'simple_type': 'int64_t'}], + torch.Tensor.cholesky: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.cholesky_solve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}], + torch.Tensor.cholesky_inverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.qr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.geqrf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.orgqr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}], + torch.Tensor.ormqr: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input2', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'input3', 'simple_type': 'Tensor'}], + torch.Tensor.lu_solve: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'LU_data', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'LU_pivots', 'simple_type': 'Tensor'}], + torch.Tensor.multinomial: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'num_samples', 'simple_type': 'SymInt'}], + torch.Tensor.lgamma_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.lgamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.digamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.polygamma: [{'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.polygamma_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'n', 'simple_type': 'int64_t'}], + torch.Tensor.erfinv: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.erfinv_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.i0: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.i0_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sign: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sign_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.signbit: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.dist: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.atan2_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.atan2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.arctan2: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.arctan2_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.lerp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Scalar'}], + torch.Tensor.lerp: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'end', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'weight', 'simple_type': 'Tensor'}], + torch.Tensor.histc: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.histogram: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'bins', 'simple_type': 'Tensor'}], + torch.Tensor.histogram: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.fmod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.fmod: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.fmod_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.fmod_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.hypot: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.hypot_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.igamma: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.igamma_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.igammac: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.igammac_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.nextafter: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.nextafter_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.remainder: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.remainder: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.remainder_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Scalar'}], + torch.Tensor.remainder_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.fmin: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.fmax: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.maximum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.minimum: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.quantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'Tensor'}], + torch.Tensor.quantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'double'}], + torch.Tensor.nanquantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'Tensor'}], + torch.Tensor.nanquantile: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'q', 'simple_type': 'double'}], + torch.Tensor.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'stable', 'simple_type': 'bool?'}], + torch.Tensor.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.sort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'stable', 'simple_type': 'bool?'}, {'is_kwarg_only': 'True', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.msort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.argsort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.argsort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'True', 'name': 'stable', 'simple_type': 'bool'}], + torch.Tensor.argsort: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'Dimname'}], + torch.Tensor.topk: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'k', 'simple_type': 'SymInt'}], + torch.Tensor.renorm: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'maxnorm', 'simple_type': 'Scalar'}], + torch.Tensor.renorm_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'p', 'simple_type': 'Scalar'}, {'is_kwarg_only': 'False', 'name': 'dim', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'maxnorm', 'simple_type': 'Scalar'}], + torch.Tensor.unfold: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'dimension', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'size', 'simple_type': 'int64_t'}, {'is_kwarg_only': 'False', 'name': 'step', 'simple_type': 'int64_t'}], + torch.Tensor.equal: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch.Tensor.pow: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch.Tensor.pow_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch.Tensor.pow_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch.Tensor.float_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch.Tensor.float_power: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch.Tensor.float_power_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Scalar'}], + torch.Tensor.float_power_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'exponent', 'simple_type': 'Tensor'}], + torch.Tensor.normal_: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.isfinite: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.isinf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.record_stream: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 's', 'simple_type': 'Stream'}], + torch.Tensor.isposinf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.isneginf: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.det: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.slogdet: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.logdet: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.inverse: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}], + torch.Tensor.inner: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'other', 'simple_type': 'Tensor'}], + torch.Tensor.outer: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch.Tensor.ger: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'vec2', 'simple_type': 'Tensor'}], + torch.Tensor.to_padded_tensor: [{'is_kwarg_only': 'False', 'name': 'self', 'simple_type': 'Tensor'}, {'is_kwarg_only': 'False', 'name': 'padding', 'simple_type': 'double'}], +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/hop_db.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/hop_db.py new file mode 100644 index 0000000000000000000000000000000000000000..80a249b8d909f5f0b76754dd5e880f62fb0e2f93 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/hop_db.py @@ -0,0 +1,684 @@ +# mypy: ignore-errors + +import functools +import unittest + +import torch +from functorch.experimental.control_flow import map +from torch._higher_order_ops.flex_attention import ( + flex_attention as flex_attention_hop, +) +from torch.nn.attention.flex_attention import ( + _create_empty_block_mask, + create_block_mask, + flex_attention, +) +from torch.testing import make_tensor +from torch._higher_order_ops.inline_asm_elementwise import inline_asm_elementwise +from torch.testing._internal.common_device_type import onlyCUDA +from torch.testing._internal.common_dtype import all_types_and, custom_types +from torch.testing._internal.opinfo.core import DecorateInfo, OpInfo, SampleInput +from torch._higher_order_ops.invoke_subgraph import mark_compile_region +from torch._higher_order_ops import InvokeQuant, invoke_quant_packed + + +def sample_inputs_map(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + yield SampleInput( + [make_arg(2, 2, 2, low=0.1, high=2), make_arg(2, 2, 2, low=0.1, high=2)], + args=(make_arg(1, low=0.1, high=2), make_arg(1, low=0.1, high=2)), + ) + + +def inner_f(x, y0, y1): + return [x[0].cos().add_(1.0) * y0, (x[1] + y1.sin()).cos_().view(x[1].size())] + + +def simple_map(xs, y0, y1): + def f(x, y0, y1): + return inner_f(x, y0, y1) + + return map(f, xs, y0, y1) + + +def nested_map(xs, y0, y1): + def f1(xx, y0, y1): + def f2(x, y0, y1): + return inner_f(x, y0, y1) + + return map(f2, xx, y0, y1) + + return map(f1, xs, y0, y1) + + +def triple_nested_map(xs, y0, y1): + def f0(xs, y0, y1): + def f1(xx, y0, y1): + def f2(x, y0, y1): + return inner_f(x, y0, y1) + + return map(f2, xx, y0, y1) + + return map(f1, xs, y0, y1) + + return map(f0, xs, y0, y1) + + +# PLEASE DON'T ADD ANYTHING NEW TO THIS LIST, +# and do add an OpInfo for your HOP. +# The OpInfo lets us do automated testing for the HOP to check that +# your HOP will work correctly with PyTorch! +# +# Your new HOP may fail some automated testing. That's OK. If you don't +# care about certain features (like torch.export), it's fine to xfail those +# failing tests. It is less fine to xfail a more critical check (like checking +# if torch.compile works with your HOP, or if your HOP has a docstring). +# If you don't know if a test is fine to xfail, please ask. +# +# There are legitimate reasons why something cannot be added to this list +# (e.g. it uses executorch which is not in PyTorch). If that's the case then +# please leave a comment. +FIXME_hop_that_doesnt_have_opinfo_test_allowlist = [ + "custom_function_call", + "autograd_function_apply", + "run_and_save_rng_state", + "run_with_rng_state", + "run_dtensor_rng_op", + "graphsafe_run_with_rng_state", + "out_dtype", + "trace_wrapped", + 'tag_activation_checkpoint', + 'executorch_call_delegate', + 'wrap', + 'wrap_with_set_grad_enabled', + 'auto_functionalized_v2', + 'associative_scan', + 'flat_apply', # is WIP, doesn't pass any of the tests yet + 'wrap_with_autocast', + 'wrap_activation_checkpoint', + 'run_const_graph', + 'auto_functionalized', + "map", # T183144629 + "map_impl", + "with_effects", + "strict_mode", + "_export_tracepoint", + "call_torchbind", + "triton_kernel_wrapper_mutation", + "triton_kernel_wrapper_functional", + "hints_wrapper", + "dynamo_bypassing_wrapper", # TODO(soulitzer) + "foreach_map", + "aoti_call_delegate", + "print", + "inductor_compiled_code", # Tested separately in test_inductor_wrap_inductor_compile_regions + "invoke_leaf_function", # Needs torch.compile, tested separately in test_leaf_function* +] + +torch.library.define( + "testlib::mutating_custom_op", + "(Tensor(a!) x, Tensor(b!) z) -> (Tensor, Tensor, Tensor)", + tags=torch.Tag.pt2_compliant_tag, +) + + +@torch.library.impl("testlib::mutating_custom_op", "cpu") +def foo_impl_cpu(x, z): + x.add_(5) + z.add_(5) + return x.clone(), z.clone(), x + z + + +@torch.library.impl("testlib::mutating_custom_op", "cuda") +def foo_impl_cuda(x, z): + x.add_(5) + z.add_(5) + return x.clone(), z.clone(), x + z + + +@torch.library.impl("testlib::mutating_custom_op", "xpu") +def foo_impl_xpu(x, z): + x.add_(5) + z.add_(5) + return x.clone(), z.clone(), x + z + + +@torch.library.register_fake("testlib::mutating_custom_op") +def foo_impl_abstract(x, z): + return x.clone(), z.clone(), x + z + + +def sample_inputs_cond(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2)) + + +def simple_cond(x): + return torch.cond(x.sum() > 2, lambda x: (x.cos(),), lambda x: (x.sin(),), [x]) + + +def sample_inputs_invoke_subgraph(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2)) + + +@mark_compile_region +def fn_for_invoke_subgraph(x): + return torch.sin(x) + + +def simple_invoke_subgraph(x): + return fn_for_invoke_subgraph(x) + + +def sample_inputs_auto_functionalize(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=False + ) + yield SampleInput( + make_arg(2, 2, 2, low=0.1, high=2), make_arg(2, 2, 2, low=0.1, high=2) + ) + + +def simple_auto_functionalize(x, z): + return torch.ops.testlib.mutating_custom_op(x, z) + + +def sample_inputs_flex_attention(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + + def score_mod(score, b, h, m, n): + return score + h + + q, k, v = (make_arg(2, 2, 128, 8, low=0.1, high=2) for _ in range(3)) + block_mask = _create_empty_block_mask(q, k) + yield SampleInput(q, k, v, score_mod, block_mask) + + +def sample_inputs_flex_attention_backward( + opinfo, device, dtype, requires_grad, **kwargs +): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=False + ) + + def score_mod(score, b, h, m, n): + return score + + def mask_mod(b, h, m, n): + return m >= n + + q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3)) + block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device) + scale = 1.0 / q.size(-1) ** 0.5 + out, logsumexp, _ = flex_attention_hop( + q, k, v, score_mod, block_mask.as_tuple(), scale, {}, + ) + yield SampleInput( + q, + args=( + k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None, + score_mod, None, block_mask.as_tuple(), + scale, {}, (), (), + ), + ) + + +def sample_inputs_flex_attention_backward_explicit_buffers( + opinfo, device, dtype, requires_grad, **kwargs +): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=False + ) + mask_offset = torch.full((), 128, device=device, dtype=torch.int32) + + def score_mod(score, b, h, m, n): + return score + + def mask_mod(b, h, m, n): + return m + mask_offset >= n + + q, k, v = (make_arg(2, 2, 128, 16, low=0.1, high=2) for _ in range(3)) + block_mask = create_block_mask(mask_mod, B=2, H=2, Q_LEN=128, KV_LEN=128, device=device) + scale = 1.0 / q.size(-1) ** 0.5 + out, logsumexp, _ = flex_attention_hop( + q, k, v, score_mod, block_mask.as_tuple(), scale, {}, + ) + yield SampleInput( + q, + args=( + k, v, out.detach(), logsumexp.detach(), torch.rand_like(out), None, + score_mod, None, block_mask.as_tuple(), + scale, {}, (), (), + ), + ) + + +def simple_flex_attention_backward( + query, + key, + value, + out, + logsumexp, + grad_out, + grad_logsumexp, + fw_graph, + joint_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, +): + return torch.ops.higher_order.flex_attention_backward( + query, + key, + value, + out, + logsumexp, + grad_out, + grad_logsumexp, + fw_graph, + joint_graph, + block_mask, + scale, + kernel_options, + score_mod_other_buffers, + mask_mod_other_buffers, + ) + + +def sample_inputs_while_loop(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=False + ) + yield SampleInput( + torch.tensor(3), + make_arg(2, 3, 4, low=0.1, high=2), + ) + + +def simple_while_loop(iter_t, x): + def cond_fn(iter_t, x): + return iter_t > 0 + + def body_fn(iter_t, x): + return iter_t - 1, x.cos() + + return torch._higher_order_ops.while_loop(cond_fn, body_fn, (iter_t, x)) + + +def simple_while_loop_stack_output(iter_t, x): + def cond_fn(iter_t, x): + return iter_t > 0 + + def body_fn(iter_t, x): + return iter_t - 1, x.cos() + + return torch._higher_order_ops.while_loop_stack_output( + cond_fn, body_fn, (iter_t, x), tuple() + ) + + +def sample_inputs_local_map_hop(opinfo, device, dtype, requires_grad, **kwargs): + # TODO: once HOPs support DTensor inputs, we should also test DTensors + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=False + ) + yield SampleInput( + make_arg(2, 3, 4, low=0.1, high=2), + make_arg(2, 3, 4, low=0.1, high=2), + ) + + +def simple_local_map_hop(inp1, inp2): + def body_gm(inp1, inp2): + return inp1.cos() + inp2.sin() + + gm = torch.fx.symbolic_trace(body_gm) + + if not torch.distributed.is_available(): + raise AssertionError("Expected torch.distributed to be available") + from torch.distributed.tensor.placement_types import Replicate + + gm.meta["local_map_kwargs"] = { + "in_placements": (Replicate(), Replicate(), Replicate()), + "out_placements": ((Replicate(), Replicate(), Replicate()),), + } + + # TODO: Dynamo would rewrite this op differently + return torch._higher_order_ops.local_map_hop(gm, inp1, inp2) + + +def sample_inputs_scan(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + yield SampleInput( + make_arg(2, 2, low=0.1, high=2), + make_arg(2, 2, 2, low=0.1, high=2), + ) + + +def simple_scan(init, xs): + def combine_fn(carry, x): + result = carry @ x + x + return result, carry.clone() + + return torch._higher_order_ops.scan(combine_fn, init, xs) + + +quant_tracer = InvokeQuant() + + +def simple_invoke_quant(x): + def fn(x, y): + return (torch.sin(x) * y,) + + return quant_tracer(fn, x, x)[0] * 2.0 + + +def simple_invoke_quant_packed(x): + def fn(x): + return (torch.sin(x),) + + return invoke_quant_packed(fn, x)[0] * 2.0 + + +def sample_inputs_inline_asm(opinfo, device, dtype, requires_grad, **kwargs): + make_arg = functools.partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + yield SampleInput(make_arg(2, 2, 2, low=0.1, high=2)) + + +def simple_inline_asm(x): + if torch.version.hip: + return inline_asm_elementwise( + x, + asm_str="v_mov_b32_e32 $0, $1", + constraints="=v, v", + dtype=torch.float32, + ) + + return inline_asm_elementwise( + x, asm_str="mov.f32 $0, $1;", constraints="=f,f", dtype=torch.float32 + ) + + +hop_db = [ + OpInfo( + name="scan", + variant_test_name="simple", + op=simple_scan, + sample_inputs_func=sample_inputs_scan, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=False, + # "torch.compile with aot_autograd does not currently support double backward." + supports_gradgrad=False, + ), + OpInfo( + name="invoke_subgraph", + variant_test_name="simple", + op=simple_invoke_subgraph, + sample_inputs_func=sample_inputs_invoke_subgraph, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=True, + # "torch.compile with aot_autograd does not currently support double backward." + supports_gradgrad=False, + ), + OpInfo( + name="map", + variant_test_name="simple", + op=simple_map, + sample_inputs_func=sample_inputs_map, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + ), + OpInfo( + name="map", + variant_test_name="nested", + op=nested_map, + sample_inputs_func=sample_inputs_map, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + ), + OpInfo( + name="map", + variant_test_name="triple_nested", + op=triple_nested_map, + sample_inputs_func=sample_inputs_map, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + ), + OpInfo( + name="cond", + variant_test_name="simple", + op=simple_cond, + sample_inputs_func=sample_inputs_cond, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=True, + # "torch.compile with aot_autograd does not currently support double backward." + supports_gradgrad=False, + ), + OpInfo( + name="invoke_quant", + variant_test_name="simple", + op=simple_invoke_quant, + sample_inputs_func=sample_inputs_invoke_subgraph, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=True, + # "torch.compile with aot_autograd does not currently support double backward." + skips=( + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"), + DecorateInfo( + unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export" + ), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"), + ), + # "torch.compile with aot_autograd does not currently support double backward." + supports_gradgrad=False, + ), + OpInfo( + name="invoke_quant_packed", + variant_test_name="simple", + op=simple_invoke_quant_packed, + sample_inputs_func=sample_inputs_invoke_subgraph, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=True, + # "torch.compile with aot_autograd does not currently support double backward." + supports_gradgrad=False, + ), + OpInfo( + name="while_loop", + variant_test_name="simple", + op=simple_while_loop, + sample_inputs_func=sample_inputs_while_loop, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=False, + ), + OpInfo( + name="while_loop_stack_output", + variant_test_name="simple", + op=simple_while_loop_stack_output, + sample_inputs_func=sample_inputs_while_loop, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=False, + ), + OpInfo( + name="auto_functionalize", + variant_test_name="simple", + op=simple_auto_functionalize, + sample_inputs_func=sample_inputs_auto_functionalize, + dtypes=all_types_and(torch.bool, torch.half), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=False, + ), + OpInfo( + name="flex_attention", + variant_test_name="simple", + op=flex_attention, + sample_inputs_func=sample_inputs_flex_attention, + dtypes=custom_types(torch.float16, torch.float32), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"), + DecorateInfo( + unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export" + ), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"), + ), + decorators=[onlyCUDA], + ), + OpInfo( + name="flex_attention_backward", + variant_test_name="simple", + op=simple_flex_attention_backward, + sample_inputs_func=sample_inputs_flex_attention_backward, + dtypes=custom_types(torch.float16, torch.float32), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=False, + supports_gradgrad=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"), + DecorateInfo( + unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export" + ), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"), + ), + decorators=[onlyCUDA], + ), + OpInfo( + name="flex_attention_backward", + variant_test_name="explicit_buffers", + op=simple_flex_attention_backward, + sample_inputs_func=sample_inputs_flex_attention_backward_explicit_buffers, + dtypes=custom_types(torch.float16, torch.float32), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=False, + supports_gradgrad=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"), + DecorateInfo( + unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export" + ), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"), + ), + decorators=[onlyCUDA], + ), + OpInfo( + name="local_map_hop", + variant_test_name="simple", + op=simple_local_map_hop, + sample_inputs_func=sample_inputs_local_map_hop, + dtypes=custom_types(torch.float16, torch.float32), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + skips=( + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_aot_export"), + DecorateInfo( + unittest.expectedFailure, "TestHOP", "test_pre_dispatch_export" + ), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_serialize_export"), + DecorateInfo(unittest.expectedFailure, "TestHOP", "test_retrace_export"), + ), + decorators=[ + onlyCUDA, + unittest.skipIf( + not torch.distributed.is_available(), "requires distributed build" + ), + ], + ), + OpInfo( + name="inline_asm_elementwise", + variant_test_name="simple", + op=simple_inline_asm, + sample_inputs_func=sample_inputs_inline_asm, + dtypes=custom_types(torch.float32), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + check_inplace_batched_forward_grad=False, + supports_autograd=False, + decorators=[onlyCUDA], + ), +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/hypothesis_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/hypothesis_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..33ba93d0d2debb26d972d5f954c526a6292cf33e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/hypothesis_utils.py @@ -0,0 +1,383 @@ +# mypy: ignore-errors + +from collections import defaultdict +from collections.abc import Iterable +import numpy as np +import torch + +import hypothesis +from functools import reduce +from importlib.metadata import version +from hypothesis import assume +from hypothesis import settings +from hypothesis import strategies as st +from hypothesis.extra import numpy as stnp +from hypothesis.strategies import SearchStrategy + +from torch.testing._internal.common_quantized import _calculate_dynamic_qparams, _calculate_dynamic_per_channel_qparams + +# Setup for the hypothesis tests. +# The tuples are (torch_quantized_dtype, zero_point_enforce), where the last +# element is enforced zero_point. If None, any zero_point point within the +# range of the data type is OK. + +# Tuple with all quantized data types. +_ALL_QINT_TYPES = ( + torch.quint8, + torch.qint8, + torch.qint32, +) + +# Enforced zero point for every quantized data type. +# If None, any zero_point point within the range of the data type is OK. +_ENFORCED_ZERO_POINT = defaultdict(lambda: None, { + torch.quint8: None, + torch.qint8: None, + torch.qint32: 0 +}) + +def _get_valid_min_max(qparams): + scale, zero_point, _quantized_type = qparams + adjustment = 1 + torch.finfo(torch.float).eps + _long_type_info = torch.iinfo(torch.long) + long_min, long_max = _long_type_info.min / adjustment, _long_type_info.max / adjustment + # make sure intermediate results are within the range of long + min_value = max((long_min - zero_point) * scale, (long_min / scale + zero_point)) + max_value = min((long_max - zero_point) * scale, (long_max / scale + zero_point)) + return np.float32(min_value), np.float32(max_value) + +# This wrapper wraps around `st.floats` and checks the version of `hypothesis`, if +# it is too old, removes the `width` parameter (which was introduced) +# in 3.67.0 +def _floats_wrapper(*args, **kwargs): + if 'width' in kwargs and hypothesis.version.__version_info__ < (3, 67, 0): + # As long as nan, inf, min, max are not specified, reimplement the width + # parameter for older versions of hypothesis. + no_nan_and_inf = ( + (('allow_nan' in kwargs and not kwargs['allow_nan']) or + 'allow_nan' not in kwargs) and + (('allow_infinity' in kwargs and not kwargs['allow_infinity']) or + 'allow_infinity' not in kwargs)) + min_and_max_not_specified = ( + len(args) == 0 and + 'min_value' not in kwargs and + 'max_value' not in kwargs + ) + if no_nan_and_inf and min_and_max_not_specified: + if kwargs['width'] == 16: + kwargs['min_value'] = torch.finfo(torch.float16).min + kwargs['max_value'] = torch.finfo(torch.float16).max + elif kwargs['width'] == 32: + kwargs['min_value'] = torch.finfo(torch.float32).min + kwargs['max_value'] = torch.finfo(torch.float32).max + elif kwargs['width'] == 64: + kwargs['min_value'] = torch.finfo(torch.float64).min + kwargs['max_value'] = torch.finfo(torch.float64).max + kwargs.pop('width') + return st.floats(*args, **kwargs) + +def floats(*args, **kwargs): + if 'width' not in kwargs: + kwargs['width'] = 32 + return _floats_wrapper(*args, **kwargs) + +"""Hypothesis filter to avoid overflows with quantized tensors. + +Args: + tensor: Tensor of floats to filter + qparams: Quantization parameters as returned by the `qparams`. + +Returns: + True + +Raises: + hypothesis.UnsatisfiedAssumption + +Note: This filter is slow. Use it only when filtering of the test cases is + absolutely necessary! +""" +def assume_not_overflowing(tensor, qparams): + min_value, max_value = _get_valid_min_max(qparams) + assume(tensor.min() >= min_value) + assume(tensor.max() <= max_value) + return True + +"""Strategy for generating the quantization parameters. + +Args: + dtypes: quantized data types to sample from. + scale_min / scale_max: Min and max scales. If None, set to 1e-3 / 1e3. + zero_point_min / zero_point_max: Min and max for the zero point. If None, + set to the minimum and maximum of the quantized data type. + Note: The min and max are only valid if the zero_point is not enforced + by the data type itself. + +Generates: + scale: Sampled scale. + zero_point: Sampled zero point. + quantized_type: Sampled quantized type. +""" +@st.composite +def qparams(draw, dtypes=None, scale_min=None, scale_max=None, + zero_point_min=None, zero_point_max=None): + if dtypes is None: + dtypes = _ALL_QINT_TYPES + if not isinstance(dtypes, (list, tuple)): + dtypes = (dtypes,) + quantized_type = draw(st.sampled_from(dtypes)) + + _type_info = torch.iinfo(quantized_type) + qmin, qmax = _type_info.min, _type_info.max + + # TODO: Maybe embed the enforced zero_point in the `torch.iinfo`. + _zp_enforced = _ENFORCED_ZERO_POINT[quantized_type] + if _zp_enforced is not None: + zero_point = _zp_enforced + else: + _zp_min = qmin if zero_point_min is None else zero_point_min + _zp_max = qmax if zero_point_max is None else zero_point_max + zero_point = draw(st.integers(min_value=_zp_min, max_value=_zp_max)) + + if scale_min is None: + scale_min = torch.finfo(torch.float).eps + if scale_max is None: + scale_max = torch.finfo(torch.float).max + scale = draw(floats(min_value=scale_min, max_value=scale_max, width=32)) + + return scale, zero_point, quantized_type + +"""Strategy to create different shapes. +Args: + min_dims / max_dims: minimum and maximum rank. + min_side / max_side: minimum and maximum dimensions per rank. + +Generates: + Possible shapes for a tensor, constrained to the rank and dimensionality. + +Example: + # Generates 3D and 4D tensors. + @given(Q = qtensor(shapes=array_shapes(min_dims=3, max_dims=4)) + some_test(self, Q):... +""" +@st.composite +def array_shapes(draw, min_dims=1, max_dims=None, min_side=1, max_side=None, max_numel=None): + """Return a strategy for array shapes (tuples of int >= 1).""" + if min_dims >= 32: + raise AssertionError(f"Expected min_dims < 32, got {min_dims}") + if max_dims is None: + max_dims = min(min_dims + 2, 32) + if max_dims >= 32: + raise AssertionError(f"Expected max_dims < 32, got {max_dims}") + if max_side is None: + max_side = min_side + 5 + candidate = st.lists(st.integers(min_side, max_side), min_size=min_dims, max_size=max_dims) + if max_numel is not None: + candidate = candidate.filter(lambda x: reduce(int.__mul__, x, 1) <= max_numel) + return draw(candidate.map(tuple)) + + +"""Strategy for generating test cases for tensors. +The resulting tensor is in float32 format. + +Args: + shapes: Shapes under test for the tensor. Could be either a hypothesis + strategy, or an iterable of different shapes to sample from. + elements: Elements to generate from for the returned data type. + If None, the strategy resolves to float within range [-1e6, 1e6]. + qparams: Instance of the qparams strategy. This is used to filter the tensor + such that the overflow would not happen. + +Generates: + X: Tensor of type float32. Note that NaN and +/-inf is not included. + qparams: (If `qparams` arg is set) Quantization parameters for X. + The returned parameters are `(scale, zero_point, quantization_type)`. + (If `qparams` arg is None), returns None. +""" +@st.composite +def tensor(draw, shapes=None, elements=None, qparams=None, dtype=np.float32): + if isinstance(shapes, SearchStrategy): + _shape = draw(shapes) + else: + _shape = draw(st.sampled_from(shapes)) + if qparams is None: + if elements is None: + elements = floats(-1e6, 1e6, allow_nan=False, width=32) + X = draw(stnp.arrays(dtype=dtype, elements=elements, shape=_shape)) + assume(not (np.isnan(X).any() or np.isinf(X).any())) + return X, None + qparams = draw(qparams) + if elements is None: + min_value, max_value = _get_valid_min_max(qparams) + elements = floats(min_value, max_value, allow_infinity=False, + allow_nan=False, width=32) + X = draw(stnp.arrays(dtype=dtype, elements=elements, shape=_shape)) + # Recompute the scale and zero_points according to the X statistics. + scale, zp = _calculate_dynamic_qparams(X, qparams[2]) + enforced_zp = _ENFORCED_ZERO_POINT.get(qparams[2], None) + if enforced_zp is not None: + zp = enforced_zp + return X, (scale, zp, qparams[2]) + +@st.composite +def per_channel_tensor(draw, shapes=None, elements=None, qparams=None): + if isinstance(shapes, SearchStrategy): + _shape = draw(shapes) + else: + _shape = draw(st.sampled_from(shapes)) + if qparams is None: + if elements is None: + elements = floats(-1e6, 1e6, allow_nan=False, width=32) + X = draw(stnp.arrays(dtype=np.float32, elements=elements, shape=_shape)) + assume(not (np.isnan(X).any() or np.isinf(X).any())) + return X, None + qparams = draw(qparams) + if elements is None: + min_value, max_value = _get_valid_min_max(qparams) + elements = floats(min_value, max_value, allow_infinity=False, + allow_nan=False, width=32) + X = draw(stnp.arrays(dtype=np.float32, elements=elements, shape=_shape)) + # Recompute the scale and zero_points according to the X statistics. + scale, zp = _calculate_dynamic_per_channel_qparams(X, qparams[2]) + enforced_zp = _ENFORCED_ZERO_POINT.get(qparams[2], None) + if enforced_zp is not None: + zp = enforced_zp + # Permute to model quantization along an axis + axis = int(np.random.randint(0, X.ndim, 1)) + permute_axes = np.arange(X.ndim) + permute_axes[0] = axis + permute_axes[axis] = 0 + X = np.transpose(X, permute_axes) + + return X, (scale, zp, axis, qparams[2]) + +"""Strategy for generating test cases for tensors used in Conv. +The resulting tensors is in float32 format. + +Args: + spatial_dim: Spatial Dim for feature maps. If given as an iterable, randomly + picks one from the pool to make it the spatial dimension + batch_size_range: Range to generate `batch_size`. + Must be tuple of `(min, max)`. + input_channels_per_group_range: + Range to generate `input_channels_per_group`. + Must be tuple of `(min, max)`. + output_channels_per_group_range: + Range to generate `output_channels_per_group`. + Must be tuple of `(min, max)`. + feature_map_range: Range to generate feature map size for each spatial_dim. + Must be tuple of `(min, max)`. + kernel_range: Range to generate kernel size for each spatial_dim. Must be + tuple of `(min, max)`. + max_groups: Maximum number of groups to generate. + elements: Elements to generate from for the returned data type. + If None, the strategy resolves to float within range [-1e6, 1e6]. + qparams: Strategy for quantization parameters. for X, w, and b. + Could be either a single strategy (used for all) or a list of + three strategies for X, w, b. +Generates: + (X, W, b, g): Tensors of type `float32` of the following drawen shapes: + X: (`batch_size, input_channels, H, W`) + W: (`output_channels, input_channels_per_group) + kernel_shape + b: `(output_channels,)` + groups: Number of groups the input is divided into +Note: X, W, b are tuples of (Tensor, qparams), where qparams could be either + None or (scale, zero_point, quantized_type) + + +Example: + @given(tensor_conv( + spatial_dim=2, + batch_size_range=(1, 3), + input_channels_per_group_range=(1, 7), + output_channels_per_group_range=(1, 7), + feature_map_range=(6, 12), + kernel_range=(3, 5), + max_groups=4, + elements=st.floats(-1.0, 1.0), + qparams=qparams() + )) +""" +@st.composite +def tensor_conv( + draw, spatial_dim=2, batch_size_range=(1, 4), + input_channels_per_group_range=(3, 7), + output_channels_per_group_range=(3, 7), feature_map_range=(6, 12), + kernel_range=(3, 7), max_groups=1, can_be_transposed=False, + elements=None, qparams=None +): + + # Resolve the minibatch, in_channels, out_channels, iH/iW, iK/iW + batch_size = draw(st.integers(*batch_size_range)) + input_channels_per_group = draw( + st.integers(*input_channels_per_group_range)) + output_channels_per_group = draw( + st.integers(*output_channels_per_group_range)) + groups = draw(st.integers(1, max_groups)) + input_channels = input_channels_per_group * groups + output_channels = output_channels_per_group * groups + + if isinstance(spatial_dim, Iterable): + spatial_dim = draw(st.sampled_from(spatial_dim)) + + feature_map_shape = [draw(st.integers(*feature_map_range)) for _ in range(spatial_dim)] + + kernels = [draw(st.integers(*kernel_range)) for _ in range(spatial_dim)] + + tr = False + weight_shape = (output_channels, input_channels_per_group) + tuple(kernels) + bias_shape = output_channels + if can_be_transposed: + tr = draw(st.booleans()) + if tr: + weight_shape = (input_channels, output_channels_per_group) + tuple(kernels) + bias_shape = output_channels + + # Resolve the tensors + if qparams is not None: + if isinstance(qparams, (list, tuple)): + if len(qparams) != 3: + raise AssertionError("Need 3 qparams for X, w, b") + else: + qparams = [qparams] * 3 + + X = draw(tensor(shapes=( + (batch_size, input_channels) + tuple(feature_map_shape),), + elements=elements, qparams=qparams[0])) + W = draw(tensor(shapes=(weight_shape,), elements=elements, + qparams=qparams[1])) + b = draw(tensor(shapes=(bias_shape,), elements=elements, + qparams=qparams[2])) + + return X, W, b, groups, tr + + +# We set the deadline in the currently loaded profile. +# Creating (and loading) a separate profile overrides any settings the user +# already specified. +hypothesis_version = tuple(map(int, version("hypothesis").split(".")[:3])) + +if (3, 16, 0) <= hypothesis_version < (3, 27, 0): + # Hypothesis 3.16 → 3.26: use `timeout` instead of `deadline` + settings.register_profile("no_deadline", timeout=hypothesis.unlimited) +else: + # Hypothesis >=3.27: use `deadline=None` + settings.register_profile("no_deadline", deadline=None) + +# Activate the profile +settings.load_profile("no_deadline") + + +def assert_deadline_disabled(): + """Check that deadlines are effectively disabled across Hypothesis versions.""" + if hypothesis_version < (3, 27, 0): + import warnings + + warning_message = ( + "Your version of hypothesis is outdated. " + "To avoid `DeadlineExceeded` errors, please update. " + f"Current hypothesis version: {hypothesis.__version__}" + ) + warnings.warn(warning_message, stacklevel=2) + else: + if settings().deadline is not None: + raise AssertionError("Expected settings().deadline to be None") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/inductor_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/inductor_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..dae4c952c6a81e31a6ce37b0933cb1b0fcbd2320 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/inductor_utils.py @@ -0,0 +1,470 @@ +# mypy: ignore-errors + +import contextlib +import functools +import logging +import os +import re +import sys +import unittest +from subprocess import CalledProcessError + +import torch +import torch._inductor.async_compile # noqa: F401 required to warm up AsyncCompile pools +import torch._inductor.config as config +from torch._inductor.codecache import CppCodeCache +from torch._inductor.codegen.common import ( + get_custom_backend_config_for_device, + get_custom_backend_pass_for_device, + get_scheduling_for_device, + get_wrapper_codegen_for_device, + init_backend_registration, + register_backend_for_device, +) +from torch._inductor.codegen.wrapper import PythonWrapperCodegen +from torch._inductor.compile_fx import shape_env_from_inputs +from torch._inductor.custom_graph_pass import CustomGraphModulePass, CustomGraphPass +from torch._inductor.graph import GraphLowering +from torch._inductor.utils import ( + get_gpu_shared_memory, + get_gpu_type, + GPU_TYPES, + is_big_gpu, + is_gpu, + OrderedSet, +) +from torch.fx.experimental.proxy_tensor import make_fx +from torch.utils._helion import has_helion +from torch.utils._pallas import has_pallas_package, has_tpu_pallas +from torch.utils._triton import has_triton +from torch.utils._config_module import ConfigModule +from torch.testing._internal.common_device_type import ( + get_desired_device_type_test_bases, +) +from torch.testing._internal.common_utils import ( + IS_CI, + IS_WINDOWS, + LazyVal, + TestCase, +) + +from collections.abc import Callable + +log: logging.Logger = logging.getLogger(__name__) + + +def test_cpu(): + try: + CppCodeCache.load("") + return True + except ( + CalledProcessError, + OSError, + torch._inductor.exc.InvalidCxxCompiler, + torch._inductor.exc.CppCompileError, + ): + return False + + +HAS_CPU = LazyVal(test_cpu) + +HAS_TRITON = has_triton() + +HAS_PALLAS = has_pallas_package() + +HAS_HELION = has_helion() + +if HAS_TRITON: + import triton + + TRITON_HAS_CPU = "cpu" in triton.backends.backends +else: + TRITON_HAS_CPU = False + + +HAS_CUDA_AND_TRITON = torch.cuda.is_available() and HAS_TRITON + +HAS_XPU_AND_TRITON = torch.xpu.is_available() and HAS_TRITON + +HAS_MPS = torch.mps.is_available() + +HAS_GPU = HAS_CUDA_AND_TRITON or HAS_XPU_AND_TRITON +HAS_GPU_AND_TRITON = HAS_GPU + +GPU_TYPE = get_gpu_type() + +HAS_MULTIGPU = any( + getattr(torch, gpu).is_available() and getattr(torch, gpu).device_count() >= 2 + for gpu in GPU_TYPES +) + +_desired_test_bases = get_desired_device_type_test_bases(allow_xpu=True) +RUN_GPU = HAS_GPU and any( + is_gpu(getattr(x, "device_type", "")) for x in _desired_test_bases +) + +RUN_CPU = HAS_CPU and any( + getattr(x, "device_type", "") == "cpu" for x in _desired_test_bases +) + +HAS_TPU = has_tpu_pallas() +# TPU is a privateuse1 backend that isn't in _desired_test_bases (it requires +# runtime initialization before the test base is registered). Check the env var +# directly, matching the same semantics as RUN_CPU/RUN_GPU: when the env var is +# unset, run if the hardware is available; when set, only run if "tpu" is listed. +_only_for = os.environ.get("PYTORCH_TESTING_DEVICE_ONLY_FOR", "") +RUN_TPU = HAS_TPU and ("tpu" in _only_for.split(",") if _only_for else True) + + +def _check_has_dynamic_shape( + self: TestCase, + code, +): + for_loop_found = False + has_dynamic = False + lines = code.split("\n") + for line in lines: + if "for(" in line: + for_loop_found = True + if re.search(r";.*ks.*;", line) is not None: + has_dynamic = True + break + self.assertTrue( + has_dynamic, msg=f"Failed to find dynamic for loop variable\n{code}" + ) + self.assertTrue(for_loop_found, f"Failed to find for loop\n{code}") + + +def skipDeviceIf(cond, msg, *, device): + if cond: + + def decorate_fn(fn): + @functools.wraps(fn) + def inner(self, *args, **kwargs): + if not hasattr(self, "device"): + warn_msg = ( + "Expect the test class to have attribute device but not found. " + ) + if hasattr(self, "device_type"): + warn_msg += "Consider using the skip device decorators in common_device_type.py" + log.warning(warn_msg) + if self.device == device: + raise unittest.SkipTest(msg) + return fn(self, *args, **kwargs) + + return inner + + else: + + def decorate_fn(fn): + return fn + + return decorate_fn + + +def skip_windows_ci(name: str, file: str) -> None: + if IS_WINDOWS and IS_CI: + module = os.path.basename(file).strip(".py") + sys.stderr.write( + f"Windows CI does not have necessary dependencies for {module} tests yet\n" + ) + if name == "__main__": + sys.exit(0) + raise unittest.SkipTest("requires sympy/functorch/filelock") + + +# TODO: Remove HAS_MPS condition when `HAS_GPU` includes HAS_MPS +requires_gpu = functools.partial( + unittest.skipIf, not (HAS_GPU or HAS_MPS), "requires gpu" +) +requires_triton = functools.partial(unittest.skipIf, not HAS_TRITON, "requires triton") +requires_helion = functools.partial(unittest.skipIf, not HAS_HELION, "requires helion") + + +def requires_gpu_with_enough_memory(min_mem_required): + def inner(fn): + total_memory = sys.maxsize + if torch.xpu.is_available(): + total_memory = torch.xpu.get_device_properties().total_memory + elif torch.cuda.is_available(): + total_memory = torch.cuda.get_device_properties().total_memory + if ( + not (torch.cuda.is_available() or torch.xpu.is_available()) + or total_memory < min_mem_required + ): + return unittest.skip( + f"Only if the GPU device has at least {min_mem_required / 1e9:.3f}GB memory to be safe" + )(fn) + else: + return fn + + return inner + + +skipCUDAIf = functools.partial(skipDeviceIf, device="cuda") +skipXPUIf = functools.partial(skipDeviceIf, device="xpu") +skipCPUIf = functools.partial(skipDeviceIf, device="cpu") + +IS_A100 = LazyVal(lambda: HAS_CUDA_AND_TRITON and get_gpu_shared_memory() == 166912) + +IS_H100 = LazyVal(lambda: HAS_CUDA_AND_TRITON and get_gpu_shared_memory() == 232448) + +IS_BIG_GPU = LazyVal(lambda: HAS_GPU_AND_TRITON and is_big_gpu()) + + +def dummy_graph() -> GraphLowering: + """ + Create a graph. This is useful for unit testing code which accesses + V.graph.sizevars. + """ + example_inputs = [torch.randn(10) for _ in range(2)] + gm = make_fx(torch.add, tracing_mode="fake")(*example_inputs) + shape_env = shape_env_from_inputs(example_inputs) + graph = GraphLowering( + gm, + shape_env=shape_env, + ) + + return graph + + +def maybe_skip_size_asserts(op): + """ + For certain ops, there meta and eager implementation returns different + strides. This cause size/strides assert fail. Skip adding those + asserts for now. + """ + if ( + op.aten_name + in ( + "fft_hfftn", + "fft_hfft", + "fft_hfft2", + "fft_ihfftn", + "fft_fft", + "fft_fft2", + "fft_fftn", + "fft_ifft", + "fft_ifft2", + "fft_ifftn", + "fft_irfft", + "fft_irfft2", + "fft_irfftn", + "fft_ihfft", + "fft_ihfft2", + "fft_rfft", + "fft_rfft2", + "fft_rfftn", + "linalg_eig", + "linalg_eigvals", + ) + and "TORCHINDUCTOR_SIZE_ASSERTS" not in os.environ + ): + return torch._inductor.config.patch(size_asserts=False) + else: + return contextlib.nullcontext() + + +def get_func_call() -> str: + return ( + "void inductor_entry_impl(" + if torch._inductor.config.cpp_wrapper + else "def call(" + ) + + +def get_kernel_launch() -> str: + return "call_triton_" if torch._inductor.config.cpp_wrapper else ".run(" + + +def clone_preserve_strides_offset(x, device=None): + if not isinstance(x, torch.Tensor): + return x + buffer = torch.as_strided( + x, (x.untyped_storage().size() // x.element_size(),), (1,), 0 + ) + if not device: + buffer = buffer.clone() + else: + buffer = buffer.to(device, copy=True) + out = torch.as_strided(buffer, x.size(), x.stride(), x.storage_offset()) + return out + + +# define the e4m3/e5m2 constants +E4M3_MAX_POS = torch.finfo(torch.float8_e4m3fn).max +E5M2_MAX_POS = torch.finfo(torch.float8_e5m2).max +E4M3FNUZ_MAX_POS = torch.finfo(torch.float8_e4m3fnuz).max +E5M2FNUZ_MAX_POS = torch.finfo(torch.float8_e5m2fnuz).max + +FP16_MAX_POS: float = torch.finfo(torch.float16).max +EPS: float = 1e-12 + +Tensor = torch.Tensor + + +def _to_fp8_saturated(x: Tensor, float8_dtype: torch.dtype) -> Tensor: + # The default behavior in PyTorch for casting to `float8_e4m3fn` + # and `e5m2` is to not saturate. In this context, we should saturate. + # A common case where we want to saturate is when the history of a + # tensor has a maximum value of `amax1`, and the current amax value + # is `amax2`, where `amax1 < amax2`. This is common when using delayed + # scaling. + if float8_dtype == torch.float8_e4m3fn: + x = x.clamp(min=-1 * E4M3_MAX_POS, max=E4M3_MAX_POS) + elif float8_dtype == torch.float8_e5m2: + x = x.clamp(min=-1 * E5M2_MAX_POS, max=E5M2_MAX_POS) + elif float8_dtype == torch.float8_e4m3fnuz: + x = x.clamp(min=-1 * E4M3FNUZ_MAX_POS, max=E4M3FNUZ_MAX_POS) + elif float8_dtype == torch.float8_e5m2fnuz: + x = x.clamp(min=-1 * E5M2FNUZ_MAX_POS, max=E5M2FNUZ_MAX_POS) + else: + raise TypeError(f"Unsupported float8_dtype: {float8_dtype}") + return x.to(float8_dtype) + + +@torch.no_grad() +def _amax_to_scale( + amax: torch.Tensor, float8_dtype: torch.dtype, orig_dtype: torch.dtype +) -> torch.Tensor: + # To make scale dtype to be fp32 for accuracy + amax = amax.float() + if float8_dtype == torch.float8_e4m3fn: + res = E4M3_MAX_POS / torch.clamp(amax, min=EPS) + else: # e5m2 + res = E5M2_MAX_POS / torch.clamp(amax, min=EPS) + + # Ensure that the scale is representable in float16, + # this helps when amax is small. We are assuming that we don't need + # to care about this for float32/bfloat16. + if orig_dtype is torch.float16: + res = torch.clamp(res, max=FP16_MAX_POS) + return res + + +def _quantize_tensorwise(x: Tensor, float8_dtype: torch.dtype): + amax = torch.max(torch.abs(x)) + scale = _amax_to_scale(amax, float8_dtype, x.dtype) + x_fp8 = _to_fp8_saturated(x * scale, float8_dtype) + inverse_scale = scale.reciprocal() + return x_fp8, inverse_scale + + +def _quantize_rowwise(x: Tensor, float8_dtype: torch.dtype): + amax = torch.max(torch.abs(x), dim=1, keepdim=True).values + scale = _amax_to_scale(amax, float8_dtype, x.dtype) + x_fp8 = _to_fp8_saturated(x * scale, float8_dtype) + inverse_scale = scale.reciprocal() + return x_fp8, inverse_scale + + +def _quantize_blockwise( + x: Tensor, float8_dtype: torch.dtype, block_outer: int, block_inner: int +): + min_outer = min(block_outer, x.shape[0]) + min_inner = min(block_inner, x.shape[1]) + x = x.unflatten(1, (-1, min_inner)).unflatten(0, (-1, min_outer)) + amax = x.abs().amax(dim=[1, 3], keepdim=True).float() + scale = _amax_to_scale(amax, float8_dtype, x.dtype) + x = x.flatten(2, 3).flatten(0, 1) + scale = scale.flatten(2, 3).flatten(0, 1) + scale_expanded = scale.repeat_interleave(min_outer, dim=0).repeat_interleave( + min_inner, dim=1 + ) + x_fp8 = _to_fp8_saturated( + x / scale_expanded, # Ensures that scaling doesn't cause inf/nan values + float8_dtype, + ) + inverse_scale = scale.reciprocal() + return x_fp8, inverse_scale + + +class MockGraphHandler(GraphLowering): + """Minimal mock graph handler for testing virtualized context.""" + + def __init__(self, name_to_buffer=None): + import torch._inductor.sizevars + + self.sizevars = torch._inductor.sizevars.SizeVarAllocator() + self.name_to_buffer = name_to_buffer or {} + self.graph_inputs = {} + self.mutated_buffers = OrderedSet() + self.removed_buffers = OrderedSet() + self.constants = {} + self.scheduler = None + + def get_dtype(self, buffer_name: str) -> torch.dtype: # noqa: ARG002 + """Return default dtype for any buffer (for testing).""" + return torch.float32 + + +@contextlib.contextmanager +def patch_inductor_backend( + device: str, + python_wrapper_codegen: PythonWrapperCodegen = None, + custom_pass: CustomGraphModulePass = None, + custom_backend_config: ConfigModule = None, +): + """ + Patch the inductor backend for a specific device. + """ + # Make sure the backend is already registered + init_backend_registration() + + # Get the original registration parameters + original_scheduling = get_scheduling_for_device(device) + original_python_wrapper = get_wrapper_codegen_for_device(device, False) + original_cpp_wrapper = get_wrapper_codegen_for_device(device, True) + original_fx_wrapper = get_wrapper_codegen_for_device(device, fx_wrapper=True) + original_custom_pass = get_custom_backend_pass_for_device(device) + original_custom_backend_config = get_custom_backend_config_for_device(device) + + try: + # Register modified backend for the device + register_backend_for_device( + device, + original_scheduling, + ( + python_wrapper_codegen + if python_wrapper_codegen is not None + else original_python_wrapper + ), + original_cpp_wrapper, + original_fx_wrapper, + custom_pass if custom_pass is not None else original_custom_pass, + ( + custom_backend_config + if custom_backend_config is not None + else original_custom_backend_config + ), + ) + yield + finally: + # Restore the original backend + register_backend_for_device( + device, + original_scheduling, + original_python_wrapper, + original_cpp_wrapper, + original_fx_wrapper, + original_custom_pass, + original_custom_backend_config, + ) + +def patch_custom_fallback_pass(predicate: Callable[[torch.fx.Node], bool]) -> contextlib.ContextDecorator: + """ + Create a custom pass which falls back based on the provided predicate. For example, + we could provide a predicate which returns True for all aten.add.default nodes. + Returns a context activating the pass. + """ + class Pass(CustomGraphPass): + def __call__(self, graph: torch.fx.Graph): + for node in graph.nodes: + if predicate(node): + node.meta["should_fallback"] = True + + def uuid(self): + return None + + + return config.patch(post_grad_custom_pre_pass=Pass()) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/jit_metaprogramming_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/jit_metaprogramming_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..43894c733829cc4e342c8e1533b0bdf8c3aa1590 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/jit_metaprogramming_utils.py @@ -0,0 +1,726 @@ +# mypy: ignore-errors + +# Torch +from torch.jit.annotations import BroadcastingList2, BroadcastingList3 # noqa: F401 +import torch.nn.functional as F +import torch +import torch.cuda +import torch.jit +import torch.jit._logging +import torch.jit.frontend +from torch.testing._internal.common_nn import module_tests, get_new_module_tests +from torch.testing._internal.common_utils import is_iterable_of_tensors, noncontiguous_like + +import collections +from copy import deepcopy +from typing import Any +import math # noqa: F401 + +# Testing utils +from torch import inf + +if torch.get_default_dtype() != torch.float32: + raise AssertionError(f"Expected torch.get_default_dtype() == torch.float32, got {torch.get_default_dtype()}") + +L = 20 +M = 10 +S = 5 + + +def unpack_variables(args): + if isinstance(args, tuple): + return tuple(unpack_variables(elem) for elem in args) + else: + return args + +class dont_convert(tuple): + __slots__ = () + +non_differentiable = collections.namedtuple('non_differentiable', ['tensor']) + +def create_input(call_args, requires_grad=True, non_contiguous=False, call_kwargs=None, dtype=torch.float, device=None): + if not isinstance(call_args, tuple): + call_args = (call_args,) + + def map_arg(arg): + def maybe_non_contig(tensor): + if not non_contiguous or tensor.numel() < 2: + return tensor.clone() + + return noncontiguous_like(tensor) + + def conjugate(tensor): + return tensor.conj() + + if isinstance(arg, (torch.Size, dont_convert)): + return arg + elif isinstance(arg, tuple) and len(arg) == 0: + var = conjugate(torch.randn((), dtype=dtype, device=device)) + var.requires_grad = requires_grad + return var + elif isinstance(arg, tuple) and not isinstance(arg[0], torch.Tensor): + return conjugate(maybe_non_contig(torch.randn(*arg, dtype=dtype, device=device))).requires_grad_(requires_grad) + # double check casting + elif isinstance(arg, non_differentiable): + if isinstance(arg.tensor, torch.Tensor): + return conjugate(maybe_non_contig(arg.tensor.to(device=device))) + return conjugate(maybe_non_contig(arg.tensor.to(device=device))) + elif isinstance(arg, torch.Tensor): + if arg.is_complex() != dtype.is_complex: + raise RuntimeError("User provided tensor is real for a test that runs with complex dtype, ", + "which is not supported for now") + # NOTE: We do clone() after detach() here because we need to be able to change size/storage of v afterwards + v = conjugate(maybe_non_contig(arg)).detach().to(device=device).clone() + v.requires_grad = requires_grad and (v.is_floating_point() or v.is_complex()) + return v + elif callable(arg): + return map_arg(arg(dtype=dtype, device=device)) + else: + return arg + args_out = tuple(map_arg(arg) for arg in call_args) + kwargs_out = {k: map_arg(v) for k, v in call_kwargs.items()} if call_kwargs else {} + return args_out, kwargs_out + +# NB: JIT script tests for all nn functional interfaces, script mode does +# not support in_place operations yet, so no inplace operation tests added. +# removed all the deprecated functions +# +# ( +# method name, +# input size/constructing fn, +# args (tuple represents shape of a tensor arg), +# test variant name(will be used at test name suffix, +# 'inplace' skips grad tests), // optional +# (True, nonfusible_nodes, fusible_nodes) for autodiff // optional +# fn to determine if test should be skipped, // optional +# fn mapping output to part that should be gradcheck'ed, // optional +# kwargs for function, // optional +# ) +def get_nn_functional_tests(): + nn_functional_tests = [ + ('conv1d', (S, S, S), ((S, S, S),)), + ('conv2d', (S, S, S, S), ((S, S, S, S),)), + ('conv3d', (S, S, S, S, S), ((S, S, S, S, S),)), + ('conv_transpose1d', (S, S, S), ((S, S, S),)), + ('conv_transpose2d', (S, S, S, S), ((S, S, S, S),)), + ('conv_transpose3d', (S, S, S, S, S), ((S, S, S, S, S),)), + ('conv_tbc', (S, S, S), ((S, S, S), (S,), 2)), + ('avg_pool1d', (S, S, S), (3,)), + ('avg_pool2d', (S, S, S, S), (3,), '', (True,)), + ('avg_pool3d', (S, S, S, S, S), (3,)), + ('fractional_max_pool2d', (S, S, S, S), (3, [2, 3],)), + ('max_pool1d', (S, S, S), (2, 1)), + ('max_pool1d', (S, S, S), (2, 1, 1, 1, False, True), 'with_indices'), + ('max_pool2d', (S, S, S, S), (2, 1), '', (True, 'aten::max_pool2d_with_indices')), + ('max_pool2d', (S, S, S, S), (2, 1, 1, 1, False, True), 'with_indices', (True, 'aten::max_pool2d_with_indices')), + ('max_pool3d', (S, S, S, S, S), (2, 1)), + ('max_unpool1d', torch.tensor([[[2., 4]]]), (torch.tensor([[[1, 3]]]), 2, 2, 0)), + ('max_unpool2d', torch.tensor([[[[2., 4]]]]), (torch.tensor([[[[1, 3]]]]), 2, 2, 0)), + ('max_unpool3d', torch.tensor([[[[[2., 4]]]]]), (torch.tensor([[[[[1, 3]]]]]), 2, 2, 0)), + ('lp_pool1d', (S, S, S), (2., 3, 2,)), + ('lp_pool2d', (S, S, S, S), (2., 3, 2,)), + ('lp_pool3d', (S, S, S, S, S), (2., 3, 2,)), + ('adaptive_max_pool1d', (S, S, S), (5,)), + ('adaptive_max_pool2d', (S, S, S, S), ([5, 7],)), + ('adaptive_max_pool3d', (S, S, S, S, S), ([3, 2, 2],)), + ('adaptive_avg_pool1d', (S, S, S), (5,), '', (True,)), + ('adaptive_avg_pool2d', (S, S, S, S), ([5, 7],), '', (True,)), + ('adaptive_avg_pool3d', (S, S, S, S, S), ([3, 2, 2],), '', (True,)), + ('dropout', (S, S, S), (0.5,), '', (True, 'aten::native_dropout')), + ('alpha_dropout', (S, S, S), (0.5,)), + ('dropout2d', (S, S, S), (0.5,)), + ('dropout2d', (S, S, S, S), (0.5,), 'batched'), + ('dropout3d', (S, S, S, S), (0.5,)), + ('dropout3d', (S, S, S, S, S), (0.5,), 'batched'), + ('feature_alpha_dropout', (S, S, S), (0.5,)), + ('threshold', (S, S, S), (0.1, 2.), '', (True,)), + ('threshold', (S, S, S), (0.1, 2., True), 'inplace'), + ('relu', (S, S, S), (), '', (True,)), + ('relu', (S, S, S), (), 'inplace'), + ('glu', (S - 1, S - 1, S - 1), (),), + ('hardtanh', (S, S, S), (-0.5, 0.5), '', (True,)), + ('hardtanh', (S, S, S), (-0.5, 0.5, True), 'inplace'), + ('relu6', (S, S, S), (), '', (True,)), + ('relu6', (S, S, S), (True), 'inplace'), + ('elu', (S, S, S), (0.9,),), + ('elu', (S, S, S), (0.9, True), 'inplace'), + ('selu', (S, S, S), (),), + ('selu', (S, S, S), (True), 'inplace'), + ('celu', (S, S, S), (0.9,),), + ('celu', (S, S, S), (0.9, True), 'inplace'), + ('leaky_relu', (S, S, S), (0.02,), '', (True,)), + ('leaky_relu', (S, S, S), (0.02,), 'inplace'), + ('rrelu', (S, S), (0.1, 0.3, False),), + ('rrelu', (S, S), (0.1, 0.3, False, True), 'inplace'), + ('hardshrink', (S, S, S), (0.4,), '', (True,)), + ('tanhshrink', (S, S, S), (),), + ('softsign', (S, S, S), (),), + ('softplus', (S, S, S), (), '', (True,)), + ('softmin', (S, S, S), (0,),), + ('softmax', (S, S, S), (0,), '', (True,)), + ('softmax', (S, S, S), (0, 3, torch.double), 'with_all_args', (True,)), + ('tanh', (S, S, S), (), '', (True,)), + ('sigmoid', (S, S, S), (), '', (True,)), + ('silu', (S, S, S), (), '', (True,)), + ('log_softmax', (S, S, S), (0,), '', (True,)), + ('linear', (S, S), ((M, S),), '', (True, ['aten::linear'])), + ('linear', (S, S), ((M, S), (M,)), 'addmm', (True, ['aten::linear'])), + ('bilinear', (S, S, S), ((S, S, M), torch.zeros(M, S, M),),), + ('embedding', torch.tensor([[1, 2, 4, 5], [4, 3, 2, 5]]), (torch.rand(6, 3), ), '', (True,)), + ('embedding_bag', torch.tensor([1, 2, 4, 2]), (torch.rand(5, 3), torch.tensor([0, 4]),),), + ('batch_norm', (S, S), + (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), None, None, True, ), + 'training', (True, 'aten::_batch_norm_impl_index')), + ('batch_norm', (0, S, S, S), + (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), + non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), True, ), + 'size_zero', (True, 'aten::_batch_norm_impl_index')), + ('batch_norm', (0, S, S, S), + (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), + non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), True, ), + 'size_zero_inference', (True, 'aten::_batch_norm_impl_index')), + ('batch_norm', (S, S), + (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), + non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), True, ), + 'with_weight_and_bias_training', (True, 'aten::_batch_norm_impl_index')), + ('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), + None, non_differentiable(torch.ones(S)), True, ), + 'with_only_bias_training', (True, 'aten::_batch_norm_impl_index')), + ('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), + non_differentiable(torch.randn(S)), None, True, ), + 'with_only_weight_training', (True, 'aten::_batch_norm_impl_index')), + ('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), + None, None, False, ), + 'inference', (True, 'aten::_batch_norm_impl_index')), + ('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), + non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), False, ), + 'with_weight_and_bias_inference', (True, 'aten::_batch_norm_impl_index')), + ('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), + None, non_differentiable(torch.ones(S)), False, ), + 'with_only_bias_inference', (True, 'aten::_batch_norm_impl_index')), + ('batch_norm', (S, S), (non_differentiable(torch.randn(S)), non_differentiable(torch.ones(S)), + non_differentiable(torch.randn(S)), None, False, ), + 'with_only_weight_inference', (True, 'aten::_batch_norm_impl_index')), + ('instance_norm', (S, S, S), (non_differentiable(torch.zeros(S)), non_differentiable(torch.ones(S))),), + ('layer_norm', (S, S, S, S), ([5],), '', + (False, ['aten::contiguous', 'aten::_batch_norm_impl_index'])), + ('layer_norm', (S, S, S, S), ([5], non_differentiable(torch.rand(S)),), 'with_only_weight', + (False, ['aten::contiguous', 'aten::_batch_norm_impl_index'])), + ('layer_norm', (S, S, S, S), ([5], None, non_differentiable(torch.rand(S)),), 'with_only_bias', + (False, ['aten::contiguous', 'aten::_batch_norm_impl_index'])), + ('layer_norm', (S, S, S, S), ([5], non_differentiable(torch.rand(S)), + non_differentiable(torch.rand(S))), 'with_weight_and_bias', + (False, ['aten::contiguous', 'aten::_batch_norm_impl_index', 'aten::addcmul'])), + ('group_norm', (S, S, S), (1, torch.rand(5),),), + ('local_response_norm', (S, S, S), (2, ),), + ('nll_loss', F.log_softmax(torch.randn(3, 5), dim=0), (torch.tensor([1, 0, 4]),), '',), + ('poisson_nll_loss', torch.rand(S, 2), (torch.rand(S, 2),),), + ('poisson_nll_loss', torch.rand(S, 2), (torch.rand(S, 2), True, True), 'full'), + ('kl_div', F.log_softmax(torch.randn(S, 10), 1), (F.softmax(torch.randn(S, 10), 1),),), + ('cross_entropy', (3, S), (torch.randint(S, (3,), dtype=torch.int64),),), + ('binary_cross_entropy_with_logits', (3,), (torch.empty(3).random_(2), ),), + ('smooth_l1_loss', (3, S), (non_differentiable(torch.rand(3, S)),),), + ('huber_loss', (3, S), (non_differentiable(torch.rand(3, S)),),), + ('l1_loss', (3, S), (non_differentiable(torch.rand(3, S)),),), + ('mse_loss', (3, S), (non_differentiable(torch.rand(3, S)),),), + ('smooth_l1_loss', (3, S), ((torch.rand(3, S)),), 'with_grad'), + ('huber_loss', (3, S), ((torch.rand(3, S)),), 'with_grad'), + ('l1_loss', (3, S), ((torch.rand(3, S)),), 'with_grad'), + ('mse_loss', (3, S), ((torch.rand(3, S)),), 'with_grad'), + ('margin_ranking_loss', (S,), ((S,), (S,)),), + ('hinge_embedding_loss', (3, S), (non_differentiable(torch.rand(3, S)),),), + ('soft_margin_loss', (3, S), (non_differentiable(torch.rand(3, S)),),), + ('multilabel_soft_margin_loss', (3, S), (non_differentiable(torch.rand(3, S)),),), + ('cosine_embedding_loss', (S, S), ((S, S), non_differentiable(torch.rand(S,))),), + ('pixel_shuffle', (1, 9, 4, 4), (3,),), + ('pixel_unshuffle', (1, 1, 12, 12), (3,),), + ('affine_grid', (S, 2, 3), (torch.Size([S, 1, 7, 7]),),), + ('pad', (3, 3, 4, 2), ([1, 1],),), + ('pairwise_distance', (S, S), ((S, S),),), + ('pdist', (S, S), (),), + ('cosine_similarity', (S, S), ((S, S),),), + ('triplet_margin_loss', (S, S), ((S, S), (S, S)),), + ('normalize', (S, S, S), (),), + ('unfold', (S, S, S, S), ([2, 3]),), + ('fold', (1, 3 * 2 * 2, 12), ([4, 5], [2, 2]),), + ('grid_sample', (S, S, S, S), (non_differentiable(torch.rand(S, S, S, 2)),),), + ('gumbel_softmax', (S, S), (2.,), '', (True, ['aten::softmax', 'aten::add', 'aten::div'], ['aten::neg'])), + ('gumbel_softmax', (S, S), (2., True,), 'hard', (True, ['aten::softmax', 'aten::add', 'aten::div'], ['aten::neg'])), + ('multilabel_margin_loss', torch.tensor([[0.2, -0.2, 0.07]]), (torch.tensor([[0, 0, 1]]),),), + ('multi_margin_loss', (S, S), (non_differentiable(torch.randint(S, (S, ), dtype=torch.int64)), + 1, 1., non_differentiable(torch.randn(S))),), + ('binary_cross_entropy', torch.randn(3, 2).sigmoid(), (non_differentiable(torch.rand(3, 2)), + non_differentiable(torch.randn(3, 2))),), + ('binary_cross_entropy', torch.randn(3, 2).sigmoid(), + (non_differentiable(torch.rand(3, 2)), + non_differentiable(torch.randn(3, 2)), None, None, 'mean'), 'size_average'), + ('ctc_loss', torch.rand(S, S, S).log_softmax(2).detach().requires_grad_(), + (torch.randint(1, S, (S, S), dtype=torch.long), torch.full((S,), S, dtype=torch.long), + torch.randint(1, S, (S,), dtype=torch.long))), + ('upsample', torch.randn(S, S, M, M), (None, 2.), 'with_scale'), + ('upsample', torch.randn(S, S, M, M), (4,), 'with_size'), + ('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2,), 'nearest_4d'), + ('interpolate', torch.randn(S, S, M, M), (None, 2.), 'nearest_4d_with_scale'), + ('interpolate', torch.randn(S, S, M, M), (4,), 'nearest_4d_with_size'), + ('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2,), 'area_4d'), + ('interpolate', torch.randn(S, S, M, M), (None, 2.), 'area_4d_with_scale'), + ('interpolate', torch.randn(S, S, M, M), (4,), 'area_4d_with_size'), + ('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2,), 'bilinear_4d'), + ('interpolate', torch.randn(S, S, M, M), (None, 2.), 'bilinear_4d_with_scale'), + ('interpolate', torch.randn(S, S, M, M), (4,), 'bilinear_4d_with_size'), + ('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2,), 'bicubic_4d'), + ('interpolate', torch.randn(S, S, M, M), (None, 2.), 'bicubic_4d_with_scale'), + ('interpolate', torch.randn(S, S, M, M), (4,), 'bicubic_4d_with_size'), + ('interpolate', torch.zeros(3, 3).view(1, 3, 3), (2,), 'nearest_3d'), + ('interpolate', torch.randn(S, M, M), (None, 2.), 'nearest_3d_with_scale'), + ('interpolate', torch.randn(S, M, M), (4,), 'nearest_3d_with_size'), + ('interpolate', torch.zeros(3, 3).view(1, 3, 3), (2,), 'area_3d'), + ('interpolate', torch.randn(S, M, M), (None, 2.), 'area_3d_with_scale'), + ('interpolate', torch.randn(S, M, M), (4,), 'area_3d_with_size'), + ('interpolate', torch.zeros(3, 3).view(1, 3, 3), (2,), 'linear_3d'), + ('interpolate', torch.randn(S, M, M), (None, 2.), 'linear_3d_with_scale'), + ('interpolate', torch.randn(S, M, M), (4,), 'linear_3d_with_size'), + ('interpolate', torch.randn(S, M, M, M, M), (None, 2.), 'nearest_5d_with_scale'), + ('interpolate', torch.randn(S, M, M, M, M), (4,), 'nearest_5d_with_size'), + ('interpolate', torch.zeros(3, 3, 3).view(1, 1, 3, 3, 3), (2,), 'area_5d'), + ('interpolate', torch.randn(S, M, M, M, M), (None, 2.), 'area_5d_with_scale'), + ('interpolate', torch.randn(S, M, M, M, M), (4,), 'area_5d_with_size'), + ('interpolate', torch.zeros(3, 3, 3).view(1, 1, 3, 3, 3), (2,), 'trilinear_5d'), + ('interpolate', torch.randn(S, M, M, M, M), (None, 2.), 'trilinear_5d_with_scale'), + ('interpolate', torch.randn(S, M, M, M, M), (4,), 'trilinear_5d_with_size'), + ('interpolate', torch.zeros(3, 3).view(1, 1, 3, 3), (2, None, 'nearest', None, False), + 'nearest_4d_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, S, M, M), (4, None, 'nearest', None, False), + 'nearest_4d_with_size_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, S, M, M), (None, 2., 'bilinear', None, False), + 'bilinear_4d_with_scale_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, S, M, M), (4, None, 'bilinear', None, False), + 'bilinear_4d_with_size_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, S, M, M), (None, 2., 'bicubic', None, False), + 'bicubic_4d_with_scale_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, S, M, M), (4, None, 'bicubic', None, False), + 'bicubic_4d_with_size_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, M, M), (None, 2., 'nearest', None, False), + 'nearest_3d_with_scale_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, M, M), (4, None, 'nearest', None, False), + 'nearest_3d_with_size_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, M, M), (None, 2., 'linear', None, False), + 'linear_3d_with_scale_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, M, M), (4, None, 'linear', None, False), + 'linear_3d_with_size_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, M, M, M, M), (None, 2., 'nearest', None, False), + 'nearest_5d_with_scale_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, M, M, M, M), (4, None, 'nearest', None, False), + 'nearest_5d_with_size_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, M, M, M, M), (None, 2., 'trilinear', None, False), + 'trilinear_5d_with_scale_not_recompute_scale_factor'), + ('interpolate', torch.randn(S, M, M, M, M), (4, None, 'trilinear', None, False), + 'trilinear_5d_with_size_not_recompute_scale_factor'), + ] + return nn_functional_tests + +script_template = ''' +def the_method({}): + return {} +''' + +def value_to_literal(value): + if isinstance(value, str): + # Quotes string and escapes special characters + return ascii(value) + if isinstance(value, torch.Tensor): + return 'torch.' + str(value) + else: + return str(value) + +def get_call(method_name, func_type, args, kwargs): + kwargs_str = ', '.join([k + '=' + value_to_literal(v) for k, v in kwargs.items()]) + self_arg = args[0] + if func_type == 'method': + args = args[1:] + + argument_str = ', '.join(args) + argument_str += ', ' if len(args) and len(kwargs) else '' + argument_str += kwargs_str + + if func_type == 'functional' or func_type == 'function': + call = f'torch.{method_name}({argument_str})' + elif func_type == 'method': + call = f'{self_arg}.{method_name}({argument_str})' + elif func_type == 'nn_functional': + call = f'torch.nn.functional.{method_name}({argument_str})' + else: + raise TypeError('Unsupported function type') + + return call + +def get_constant(x): + if x == inf: + return 'math.inf' + if x == -inf: + return '-math.inf' + return x + +def get_script_args(args): + formals: list[str] = [] + tensors: list[torch.Tensor | list[torch.Tensor]] = [] + actuals: list[str] = [] + for arg in args: + if isinstance(arg, torch.Tensor): + name = f'i{len(formals)}' + formals.append(name) + actuals.append(name) + tensors.append(arg) + elif is_iterable_of_tensors(arg): + name = f'i{len(formals)}' + formals.append(name + ': List[torch.Tensor]') + actuals.append(name) + tensors.append(list(arg)) + elif isinstance(arg, str): + actuals.append(f"'{arg}'") + else: + actuals.append(str(get_constant(arg))) + return (formals, tensors, actuals) + +# create a script function from (name, func_type, output_process_fn), +# and returns the compiled function and example inputs +def gen_script_fn_and_args(method_name, func_type, *args, **kwargs): + formals, tensors, actuals = get_script_args(args) + call = get_call(method_name, func_type, actuals, kwargs) + script = script_template.format(', '.join(formals), call) + CU = torch.jit.CompilationUnit(script) + return CU.the_method, tensors + +# create a script function from (name, func_type), +# returns a function takes in (args, kwargs) and runs the compiled function +def create_script_fn(self, method_name, func_type): + # function returns tuple containing original output and + # filtered output to be used in checking gradients + def script_fn(*args, **kwargs): + fn, tensors = gen_script_fn_and_args(method_name, func_type, *args, **kwargs) + self.assertExportImport(fn.graph, tensors) + output = fn(*tensors) + # skip type annotate function attributes for now, see: https://github.com/python/mypy/issues/2087 + script_fn.last_graph = fn.graph_for(*tensors) # type: ignore[attr-defined] + return output + return script_fn + +class SplitInputs: + all_tensors: list[Any] + tensor_args: list[Any] + nontensor_args: list[Any] + arg_types: list[str] + tensor_kwargs: dict[str, Any] + kwarg_order: list[str] + nontensor_kwargs: dict[str, Any] + kwarg_types: dict[str, Any] + + @staticmethod + def _is_tensor_input(arg): + return isinstance(arg, torch.Tensor) or is_iterable_of_tensors(arg) + + def __init__(self, args, kwargs): + self.arg_types = ['t' if self._is_tensor_input(arg) else 's' for arg in args] + self.kwarg_types = {k: 't' if self._is_tensor_input(v) else 's' for k, v in kwargs.items()} + self.tensor_args = [arg for arg in args if self._is_tensor_input(arg)] + self.nontensor_args = [arg for arg in args if not self._is_tensor_input(arg)] + self.tensor_kwargs = {k: v for k, v in kwargs.items() if self._is_tensor_input(v)} + self.nontensor_kwargs = {k: v for k, v in kwargs.items() if not self._is_tensor_input(v)} + self.all_tensors = [*self.tensor_args, *[v for k, v in self.tensor_kwargs.items()]] + self.kwarg_order = [k for k, v in kwargs.items()] + + def nontensors_match(self, other: 'SplitInputs'): + if self.arg_types != other.arg_types: + return False + if self.kwarg_types != other.kwarg_types: + return False + if self.kwarg_order != other.kwarg_order: + return False + if self.nontensor_args != other.nontensor_args: + return False + if self.nontensor_kwargs != other.nontensor_kwargs: + return False + return True + +# make a new function where all non-tensor arguments in 'args' have been partially +# applied, and all tensor arguments remain. +# used to trace functions when some arguments are not tensors +def partial_apply_nontensors(fn, args, kwargs): + inputs = SplitInputs(args, kwargs) + + def new_fn(*tensors_): + tensors = iter(tensors_) + full_args = [args[i] if s == 's' else next(tensors) for i, s in enumerate(inputs.arg_types)] + full_kwargs = {k: kwargs[k] if s == 's' else next(tensors) for k, s in inputs.kwarg_types.items()} + return fn(*full_args, **full_kwargs) + + return new_fn, inputs + +# create a trace function from input fn +def create_traced_fn(self, fn, cache_traced_fn=False): + def traced_fn(*inputs, **kwargs): + # `check_trace` is set to False because check_trace is run with @no_grad + # Also, `check_against_reference` already does all the checks + # against python function + fn_tensors, split_inputs = partial_apply_nontensors(fn, inputs, kwargs) + if not cache_traced_fn or not hasattr(traced_fn, 'traced'): + traced = torch.jit.trace(fn_tensors, split_inputs.all_tensors, check_trace=False) + self.assertExportImport(traced.graph, split_inputs.all_tensors) + output = traced(*split_inputs.all_tensors) + if cache_traced_fn: + traced_fn.traced = traced + traced_fn.split_inputs = split_inputs + else: + # Guard to check that nontensor inputs are the same as during tracing + self.assertTrue(traced_fn.split_inputs.nontensors_match(split_inputs)) + output = traced_fn.traced(*split_inputs.all_tensors) + traced = traced_fn.traced + # skip type annotate function attributes for now, see: https://github.com/python/mypy/issues/2087 + traced_fn.last_graph = traced.graph_for(*split_inputs.all_tensors) # type: ignore[attr-defined] + traced_fn.graph = traced.graph # type: ignore[attr-defined] + return output + return traced_fn + +# known to be failing in script +EXCLUDE_SCRIPT = { + 'test_norm_fro_default', + 'test_norm_fro_cpu', + 'test_norm_nuc', + 'test_norm_fro', + 'test_norm_nuc_batched', + + # aten op has additional cudnn argument + 'test_nn_unfold', + + # flaky test - TODO fix + 'test_nn_ctc_loss', + + # unknown builtin op + 'test_nn_fold', + + # jit doesn't support sparse tensors. + 'test_to_sparse', + 'test_to_sparse_dim', +} + +# generates a script function and set of example inputs +# from a specified test in the format of nn_functional_tests +def get_nn_functional_compiled_fn_and_inputs(name, self_size, args, variant_name='', *extra_args): + test_name = 'test_nn_' + name + + if variant_name != '': + test_name = test_name + '_' + variant_name + + self_variable = create_input((self_size,))[0][0] + + # need to record this because methods can change the size (e.g. unsqueeze) + args_variable, _kwargs_variable = create_input(args) + + self_tensor = deepcopy(self_variable.data) + args_tensor = deepcopy(unpack_variables(args_variable)) + + f_args_variable = (self_variable,) + args_variable + f_args_tensor = (self_tensor,) + args_tensor # noqa: F841 + with torch._jit_internal._disable_emit_hooks(): + script_fn, inputs = gen_script_fn_and_args(name, "nn_functional", *f_args_variable) + return script_fn, inputs + + + +EXCLUDE_SCRIPT_MODULES = { + 'test_nn_AdaptiveAvgPool2d_tuple_none', + 'test_nn_AdaptiveAvgPool3d_tuple_none', + 'test_nn_AdaptiveMaxPool2d_tuple_none', + 'test_nn_AdaptiveMaxPool3d_tuple_none', + + # Doesn't use future division, so this is not supported + 'test_nn_CrossMapLRN2d', + # Derivative for aten::_scaled_dot_product_flash_attention_backward is not implemented + 'test_nn_TransformerDecoderLayer_gelu_activation', + 'test_nn_TransformerDecoderLayer_relu_activation', + 'test_nn_TransformerEncoderLayer_gelu_activation', + 'test_nn_TransformerEncoderLayer_relu_activation', + 'test_nn_Transformer_multilayer_coder', +} + +script_method_template = ''' +def forward({}): + return {} +''' + +def create_script_module(self, nn_module, constructor_args, *args, **kwargs): + def script_module(*args, **kwargs): + _formals, tensors, actuals = get_script_args(args) + + method_args = ', '.join(['self'] + actuals) + call_args_str = ', '.join(actuals) + call = f"self.submodule({call_args_str})" + script = script_method_template.format(method_args, call) + + submodule_constants = [] + if kwargs.get('is_constant'): + submodule_constants = ['submodule'] + + # Create module to use the script method + class TheModule(torch.jit.ScriptModule): + __constants__ = submodule_constants + + def __init__(self) -> None: + super().__init__() + self.submodule = nn_module(*constructor_args) + + def make_module(script): + module = TheModule() + # check __repr__ + str(module) + module.define(script) + return module + + module = make_module(script) + if self: + self.assertExportImportModule(module, tensors) + module(*args) + # skip type annotate function attributes for now, see: https://github.com/python/mypy/issues/2087 + create_script_module.last_graph = module.graph # type: ignore[attr-defined] + return module + return script_module + +def check_alias_annotation(method_name, args, kwargs, *, aten_name, func_type='method'): + formals, tensors, actuals = get_script_args(args) + call = get_call(method_name, func_type, actuals, kwargs) + script = script_template.format(', '.join(formals), call) + CU = torch.jit.CompilationUnit(script) + # to clean up IR + torch._C._jit_pass_inline(CU.the_method.graph) + torch._C._jit_pass_constant_propagation(CU.the_method.graph) + torch._C._jit_check_alias_annotation(CU.the_method.graph, tuple(tensors), aten_name) + +def get_nn_module_name_from_kwargs(**kwargs): + if 'module_name' in kwargs: + return kwargs['module_name'] + elif 'fullname' in kwargs: + return kwargs['fullname'] + elif 'constructor' in kwargs: + return kwargs['constructor'].__name__ + +def get_nn_mod_test_name(**kwargs): + if 'fullname' in kwargs: + test_name = kwargs['fullname'] + else: + test_name = get_nn_module_name_from_kwargs(**kwargs) + if 'desc' in kwargs: + test_name = f"{test_name}_{kwargs['desc']}" + return f'test_nn_{test_name}' + +def get_nn_module_class_from_kwargs(**kwargs): + name = get_nn_module_name_from_kwargs(**kwargs) + index = name.find("_") + if index == -1: + return name + else: + return name[0:name.find("_")] + +def try_get_nn_module_compiled_mod_and_inputs(*args, **kwargs): + name = get_nn_module_name_from_kwargs(**kwargs) + + if 'desc' in kwargs and 'eval' in kwargs['desc']: + # eval() is not supported, so skip these tests + return + + test_name = name + if 'desc' in kwargs: + test_name = f"{test_name}_{kwargs['desc']}" + test_name = get_nn_mod_test_name(**kwargs) + + if test_name in EXCLUDE_SCRIPT_MODULES: + return + if 'constructor' in kwargs: + nn_module = kwargs['constructor'] + else: + nn_module = getattr(torch.nn, name) + + if "FunctionalModule" in str(nn_module): + return + + if 'constructor_args_fn' in kwargs: + constructor_args = kwargs['constructor_args_fn']() + else: + constructor_args = kwargs.get('constructor_args', ()) + + # Set up inputs from tuple of sizes or constructor fn + input_dtype = torch.double + if 'input_fn' in kwargs: + input = kwargs['input_fn']() + if isinstance(input, torch.Tensor): + input = (input,) + + if all(tensor.is_complex() for tensor in input): + input_dtype = torch.cdouble + else: + input = (kwargs['input_size'],) + + # Extra parameters to forward() + if 'extra_args' in kwargs: + input = input + kwargs['extra_args'] + + if 'target_size' in kwargs: + input = input + (kwargs['target_size'],) + elif 'target_fn' in kwargs: + if torch.is_tensor(input): + input = (input,) + input = input + (kwargs['target_fn'](),) + + args_variable, _kwargs_variable = create_input(input, dtype=input_dtype) + f_args_variable = deepcopy(unpack_variables(args_variable)) + out_var = deepcopy(f_args_variable) + + + _args, mod = f_args_variable, create_script_module( + None, nn_module, constructor_args, *f_args_variable + )(*f_args_variable) + + return mod, out_var + + +def get_all_nn_module_tests(): + # additional modules test + # TODO: delete this list once we make all nn_tests work + additional_module_tests = [ + { + 'module_name': 'Bilinear', + 'constructor_args': (S, S, M), + 'input_size': (S, S), + 'extra_args': ((S, S),) + }, + { + 'module_name': 'RNNCell', + 'constructor_args': (S, S), + 'input_size': (S, S), + }, + { + 'module_name': 'LSTMCell', + 'constructor_args': (S, S), + 'input_size': (S, S), + }, + { + 'module_name': 'GRUCell', + 'constructor_args': (S, S), + 'input_size': (S, S), + }, + { + 'module_name': 'MultiheadAttention', + 'constructor_args': (128, 8), + 'input_size': (10, 8, 128), + 'extra_args': (torch.randn(10, 8, 128), torch.randn(10, 8, 128)), + 'slowTest': True + }, + { + 'module_name': 'Transformer', + 'constructor_args': (1, 1, 1, 1, 2), + 'input_size': (3, 1, 1), + 'extra_args': (torch.randn(1, 1, 1),), + 'slowTest': True + } + ] + + return module_tests + get_new_module_tests() + additional_module_tests diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/jit_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/jit_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..7381e1583dcf2175bd27c4e44c792e6190cf14ab --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/jit_utils.py @@ -0,0 +1,897 @@ +# mypy: ignore-errors + +# Torch +from torch.autograd import Variable +from torch.autograd.function import _nested_map +from torch.jit.annotations import BroadcastingList2, BroadcastingList3 # noqa: F401 + +from torch.onnx import OperatorExportTypes +import torch +import torch.cuda +import torch.jit +import torch.jit._logging +import torch.jit.frontend +import torch.jit.quantized +import zipfile +import functools + +# Testing utils +from torch.testing import FileCheck +from torch.testing._internal.common_utils import IS_WINDOWS, \ + freeze_rng_state, enable_profiling_mode_for_profiling_tests, ProfilingMode, TEST_BAILOUTS, \ + is_iterable_of_tensors +from torch.testing._internal.common_jit import JitCommonTestCase +from torch.testing._internal.common_utils import enable_profiling_mode # noqa: F401 + +# Standard library +from contextlib import contextmanager +from functools import reduce +from io import StringIO +from collections import defaultdict + +import importlib.util +import inspect +import io +import math +import os +import pickle +import sys +import tempfile +import textwrap +from importlib.abc import Loader +from typing import Any + +RUN_CUDA = torch.cuda.is_available() +RUN_CUDA_MULTI_GPU = RUN_CUDA and torch.cuda.device_count() > 1 +RUN_CUDA_HALF = RUN_CUDA +# HIP supports half, no version check necessary +if torch.cuda.is_available() and not torch.version.hip: + CUDA_VERSION = torch._C._cuda_getCompiledVersion() + for d in range(torch.cuda.device_count()): + major = torch.cuda.get_device_capability(d)[0] + if (major < 6): + RUN_CUDA_HALF = False + +def execWrapper(code, glob, loc): + exec(code, glob, loc) + +def do_input_map(fn, input): + return _nested_map(lambda t: isinstance(t, torch.Tensor), fn)(input) + +def clear_class_registry(): + torch._C._jit_clear_class_registry() + torch.jit._recursive.concrete_type_store = torch.jit._recursive.ConcreteTypeStore() + torch.jit._state._clear_class_state() + +def get_execution_plan(graph_executor_state): + execution_plans = list(graph_executor_state.execution_plans.values()) + num_plans = len(execution_plans) + if num_plans != 1: + raise RuntimeError('This test assumes this GraphExecutor should ' + f'only have one execution plan, got: {num_plans}') + return execution_plans[0] + +class _AssertRaisesRegexWithHighlightContext: + """ + A context manager that is useful for checking that error messages highlight + the correct part of the source code. + """ + + def __init__(self, test_case, exception, regex, highlight): + self.test_case = test_case + self.exception_type = exception + self.regex = regex + self.highlight = highlight + + def __enter__(self): + return self + + def __exit__(self, type, value, traceback): + with self.test_case.assertRaisesRegex(self.exception_type, self.regex): + if type: + raise value + + if self.highlight: + FileCheck().check_source_highlighted(self.highlight).run(str(value)) + + return True + +FUSION_GROUP = "prim::TensorExprGroup" + +class JitTestCase(JitCommonTestCase): + _do_cuda_memory_leak_check = True + _restored_warnings = False + + class capture_stdout(list): + """ + Replace sys.stdout with a temporary StringIO + """ + def __enter__(self): + self.sys_stdout = sys.stdout + self.stringio = StringIO() + sys.stdout = self.stringio + return self + + def __exit__(self, *args): + self.append(str(self.stringio.getvalue())) + del self.stringio + sys.stdout = self.sys_stdout + + class capture_stderr(list): + """ + Replace sys.stderr with a temporary StringIO + """ + def __enter__(self): + self.sys_stderr = sys.stderr + self.stringio = StringIO() + sys.stderr = self.stringio + return self + + def __exit__(self, *args): + self.append(str(self.stringio.getvalue())) + del self.stringio + sys.stderr = self.sys_stderr + + def setHooks(self): + torch._C._jit_set_emit_hooks(self.emitModuleHook, self.emitFunctionHook) + + def clearHooks(self): + torch._C._jit_set_emit_hooks(None, None) + + def setUp(self): + super().setUp() + # unittest overrides all warning filters and forces all of them to show up + # after we install our own to silence those coming from inside PyTorch. + # This will ensure that our filter still takes precedence. + if not JitTestCase._restored_warnings: + torch.jit.TracerWarning.ignore_lib_warnings() + JitTestCase._restored_warnings = True + self.setHooks() + + def tearDown(self): + super().tearDown() + # needs to be cleared because python might be unloaded before + # the callback gets destructed + self.clearHooks() + clear_class_registry() + + def assertAllFused(self, graph, except_for=()): + + # note this helper collects nodes on 'fast path' only + # i.e. the true blocks of specialized checks + def get_nodes_and_parents_recursively(block, kind, acc): + for node in block.nodes(): + if node.kind() == kind: + acc[block].append(node) + elif node.kind() == 'prim::DifferentiableGraph': + get_nodes_and_parents_recursively(node.g('Subgraph'), kind, acc) + elif node.kind() == 'prim::If' and (node.inputs().__next__().node().kind() == 'aten::all' or + node.inputs().__next__().node().kind() == 'prim::TypeCheck' or + node.inputs().__next__().node().kind() == 'prim::RequiresGradCheck'): + get_nodes_and_parents_recursively(node.blocks().__next__(), kind, acc) + else: + for inner_block in node.blocks(): + get_nodes_and_parents_recursively(inner_block, kind, acc) + + allowed_nodes = {'prim::Constant', FUSION_GROUP, 'prim::BailoutTemplate', + 'prim::TupleConstruct', 'prim::If', 'prim::TypeCheck', 'prim::RequiresGradCheck'} | set(except_for) + + fusion_groups : dict[torch._C.Block, list[torch._C.Node]] = defaultdict(list) + get_nodes_and_parents_recursively(graph, FUSION_GROUP, fusion_groups) + self.assertTrue(len(fusion_groups) == 1, f'got {graph}') + (graph, fusion_nodes) = next(iter(fusion_groups.items())) + # the block contains one FUSION_GROUP and the rest of nodes are `allowed_nodes` + self.assertTrue(len(fusion_nodes) == 1, f'got {graph}') + self.assertTrue(all(node.kind() in allowed_nodes for node in graph.nodes()), + f'got {graph}') + + def _isHookExceptionOk(self, e): + se = str(e) + allowed = ("Could not export Python function", + "closures are not exportable") + for a in allowed: + if a in se: + return True + return False + + def _compared_saved_loaded(self, m): + def extract_files(buffer): + # crack open the zip format to get at the main module code + with zipfile.ZipFile(buffer) as archive: + # check that we have no duplicate names + self.assertEqual(len(set(archive.namelist())), len(archive.namelist())) + files = list(filter(lambda x: x.startswith('archive/code/'), archive.namelist())) + # unwrap all the code files into strings + code_files_str = filter(lambda x: x.endswith('.py'), files) + code_files = [] + for f in code_files_str: + with archive.open(f) as stream: + code_files.append("".join([line.decode() for line in stream])) + + # unpickled all the debug files + debug_files_str = filter(lambda f: f.endswith('.debug_pkl'), files) + debug_files = [] + for f in debug_files_str: + with archive.open(f) as stream: + debug_files.append(pickle.load(stream)) + return code_files, debug_files + + # disable the hook while we parse code, otherwise we will re-enter the hook + with torch._jit_internal._disable_emit_hooks(): + try: + # short-circuit if this is an empty function or module + if len(m.code) == 0: + return + if isinstance(m, torch._C.ScriptModule): + if len(m._method_names()) == 0: + return + + # save the module to a buffer + buffer = io.BytesIO() + torch.jit.save(m, buffer) + # copy the data in the buffer so we can restore it later. This + # is because py2 and py3 have different semantics with zipfile + # and it's easier to just work with a fresh copy each time. + buffer_copy = buffer.getvalue() + + code_files, _debug_files = extract_files(buffer) + + except RuntimeError as e: + if not self._isHookExceptionOk(e): + raise + else: + return + + # import the model again (from a the copy we made of the original) + buffer2 = io.BytesIO(buffer_copy) + imported = torch.jit.load(buffer2) + + # save it again + saved_module_buffer_2 = io.BytesIO() + torch.jit.save(imported, saved_module_buffer_2) + + saved_module_buffer_2.seek(0) + code_files_2, _debug_files_2 = extract_files(saved_module_buffer_2) + + for a, b in zip(code_files, code_files_2, strict=True): + self.assertMultiLineEqual(a, b) + + if isinstance(m, torch._C.ScriptModule): + self.assertTrue(torch._C._ivalue_tags_match(m, imported._c)) + + + def emitFunctionHook(self, func): + # func has invalid names for export, skip the jitter check + if func.name == "" or "aten::" in func.name: + return + self._compared_saved_loaded(func) + + def emitModuleHook(self, module): + self._compared_saved_loaded(module) + + + def getExportImportCopyWithPacking(self, m, also_test_file=True, map_location=None): + buffer = io.BytesIO() + m.apply(lambda s: s._pack() if s._c._has_method('_pack') else None) + torch.jit.save(m, buffer) + m.apply(lambda s: s._unpack() if s._c._has_method('_unpack') else None) + buffer.seek(0) + imported = torch.jit.load(buffer, map_location=map_location) + imported.apply(lambda s: s._unpack() if s._c._has_method('_unpack') else None) + + if not also_test_file: + return imported + + # Ideally we would like to not have to manually delete the file, but NamedTemporaryFile + # opens the file, and it cannot be opened multiple times in Windows. To support Windows, + # close the file after creation and try to remove it manually + with tempfile.NamedTemporaryFile(delete=False) as f: + try: + f.close() + imported.save(f.name) + result = torch.jit.load(f.name, map_location=map_location) + finally: + os.unlink(f.name) + + result.apply(lambda s: s._unpack() if s._c._has_method('_unpack') else None) + return result + + def assertGraphContains(self, graph, kind, consider_subgraphs=False): + + if consider_subgraphs: + strgraph = str(graph) + count = strgraph.count(kind) - strgraph.count(f'with {kind}') + self.assertTrue(count > 0) + return + + def nodes(block): + out = [] + for node in block.nodes(): + if node.kind() == kind: + out.append(node) + for block in node.blocks(): + out += nodes(block) + return out + + out_nodes = nodes(graph) + self.assertTrue(len(out_nodes) > 0) + + def assertGraphContainsExactly(self, graph, kind, num_kind_nodes, consider_subgraphs=False): + def perform_assert(graph, kind, actual, expected, consider_subgraphs): + if actual == expected: + return + subgraph = 'including' if consider_subgraphs else 'excluding' + raise AssertionError( + f'{graph}\nError: graph contains {actual} {kind} nodes ({subgraph} subgraphs) but expected {expected}') + + if consider_subgraphs: + strgraph = str(graph) + count = strgraph.count(kind) - strgraph.count(f'with {kind}') + perform_assert(graph, kind, count, num_kind_nodes, + consider_subgraphs) + return + + def nodes(block): + out = [] + for node in block.nodes(): + if node.kind() == kind: + out.append(node) + for block in node.blocks(): + out += nodes(block) + return out + + out_nodes = nodes(graph) + perform_assert(graph, kind, len(out_nodes), num_kind_nodes, + consider_subgraphs) + + def assertExpectedONNXGraph(self, g, *args, **kwargs): + g = torch.onnx._optimize_trace(g, operator_export_type=OperatorExportTypes.ONNX) + self.assertExpectedGraph(g, *args, **kwargs) + + def assertExpectedGraph(self, trace, *args, **kwargs): + if isinstance(trace, torch._C.Graph): + graph = trace + else: + graph = trace.graph() + + torch._C._jit_pass_lint(graph) + torch._C._jit_pass_dce(graph) + torch._C._jit_pass_lint(graph) + graph = torch._C._jit_pass_canonicalize(graph) + torch._C._jit_pass_lint(graph) + self.assertExpected(str(graph), *args, **kwargs) + + def run_pass(self, name, trace): + if isinstance(trace, torch._C.Graph): + graph = trace + set_graph = False + else: + set_graph = True + graph = trace.graph() + + torch._C._jit_pass_lint(graph) + result = getattr(torch._C, '_jit_pass_' + name)(graph) + if result is not None and not isinstance(result, bool): + graph = result + torch._C._jit_pass_lint(graph) + + if set_graph: + trace.set_graph(graph) + return graph + + def get_frame_vars(self, frames_up): + frame = inspect.currentframe() + if not frame: + raise RuntimeError("failed to inspect frame") + i = 0 + while i < frames_up + 1: + frame = frame.f_back + if not frame: + raise RuntimeError("failed to get frame") + i += 1 + defined_vars: dict[str, Any] = {} + defined_vars.update(frame.f_locals) + defined_vars.update(frame.f_globals) + return defined_vars + + def assertRaisesRegexWithHighlight(self, exception, regex, highlight): + return _AssertRaisesRegexWithHighlightContext(self, exception, regex, highlight) + + def checkScriptRaisesRegex(self, script, inputs, exception, regex, + name=None, outputs=None, capture_output=False, + frames_up=1, profiling=ProfilingMode.PROFILING): + """ + Checks that a given function will throw the correct exception, + when executed with normal python, the string frontend, and the + AST frontend. Logic taken from `checkScript` (see comments there + for details) + """ + with enable_profiling_mode_for_profiling_tests(): + # Normal Python + with self.assertRaisesRegex(exception, regex): + if isinstance(script, str): + frame = self.get_frame_vars(frames_up) + the_locals: dict[str, Any] = {} + execWrapper(script, glob=frame, loc=the_locals) + frame.update(the_locals) + + python_fn = frame[name] + else: + python_fn = script + + python_fn(*inputs) + + # String frontend + with self.assertRaisesRegex(exception, regex): + if isinstance(script, str): + cu = torch.jit.CompilationUnit(script, _frames_up=frames_up) + string_frontend = getattr(cu, name) + else: + source = textwrap.dedent(inspect.getsource(script)) + cu = torch.jit.CompilationUnit(source, _frames_up=frames_up) + string_frontend = getattr(cu, script.__name__) + + string_frontend(*inputs) + + # Python AST frontend + if not isinstance(script, str): + with self.assertRaisesRegex(exception, regex): + ge = torch.jit.script(python_fn) + ge(*inputs) + + def checkBailouts(self, model, inputs, expected): + state = model.get_debug_state() + plan = get_execution_plan(state) + num_bailouts = plan.code.num_bailouts() + for i in range(num_bailouts): + plan.code.request_bailout(i) + bailout_outputs = model(*inputs) + self.assertEqual(bailout_outputs, expected) + + def checkScript(self, + script, + inputs, + name='func', + optimize=True, + inputs_requires_grad=False, + capture_output=False, + frames_up=1, + profiling=ProfilingMode.PROFILING, + atol=None, + rtol=None): + """ + Checks that a given script generates the same output as the Python + version using the given inputs. + """ + with torch.jit.optimized_execution(optimize), enable_profiling_mode_for_profiling_tests(): + extra_profile_runs = any(isinstance(x, torch.Tensor) and x.requires_grad for x in inputs) + if isinstance(script, str): + # Compile the string to a Script function + # with enable_profiling_mode(): + cu = torch.jit.CompilationUnit(script, _frames_up=frames_up) + + # Execute the Python function so we can run it later and get its + # outputs + + frame = self.get_frame_vars(frames_up) + the_locals: dict[str, Any] = {} + execWrapper(script, glob=frame, loc=the_locals) + frame.update(the_locals) + + python_fn = frame[name] + scripted_fn = getattr(cu, name) + else: + + # Check the string frontend first + source = textwrap.dedent(inspect.getsource(script)) + self.checkScript( + source, + inputs, + script.__name__, + optimize=optimize, + inputs_requires_grad=inputs_requires_grad, + capture_output=capture_output, + profiling=profiling, + frames_up=2) + + # Continue checking the Python frontend + scripted_fn = torch.jit.script(script, _frames_up=1) + python_fn = script + + if inputs_requires_grad: + recording_inputs = do_input_map(lambda t: t.detach().requires_grad_(), inputs) + else: + recording_inputs = inputs + + if capture_output: + with self.capture_stdout() as script_stdout: + script_outputs = scripted_fn(*recording_inputs) + with self.capture_stdout(): + opt_script_outputs = scripted_fn(*recording_inputs) + with self.capture_stdout(): + python_outputs = python_fn(*inputs) + if not IS_WINDOWS: + self.assertExpected(script_stdout[0], subname='stdout') + self.assertEqual(python_outputs, opt_script_outputs, atol=atol, rtol=rtol) + else: + # profiling run + script_outputs = scripted_fn(*recording_inputs) + if inputs_requires_grad or extra_profile_runs: + opt_script_outputs = scripted_fn(*recording_inputs) + # optimized run + opt_script_outputs = scripted_fn(*recording_inputs) + if TEST_BAILOUTS: + self.checkBailouts(scripted_fn, inputs, opt_script_outputs) + python_outputs = python_fn(*inputs) + self.assertEqual(python_outputs, script_outputs, atol=atol, rtol=rtol) + self.assertEqual(script_outputs, opt_script_outputs, atol=atol, rtol=rtol) + return scripted_fn + + def checkTrace(self, func, reference_tensors, input_tensors=None, + drop=None, allow_unused=False, verbose=False, + inputs_require_grads=True, check_tolerance=1e-5, export_import=True, + _force_outplace=False, grad_atol=None, grad_rtol=None): + + # TODO: check gradients for parameters, not just inputs + def allSum(vs): + # drop allows us to remove some values from ever being used + # to test unused outputs + if drop is not None: + vs = vs[:-drop] + # we don't want all the grad for all the outputs to be the same + # so we multiply each by a constant + return sum(math.log(i + 2) * v.sum() for i, v in enumerate(vs) if v is not None) + if input_tensors is None: + input_tensors = reference_tensors + + def flatten_inputs(inputs): + def input_reduce(input, fn, acc): + if isinstance(input, torch.Tensor): + fn(input, acc) + elif isinstance(input, dict): + reduce(lambda acc, key: input_reduce(input[key], fn, acc), input, acc) + else: + reduce(lambda acc, val: input_reduce(val, fn, acc), input, acc) + return acc + return tuple(input_reduce(recording_inputs, lambda t, acc: acc.append(t), [])) + + nograd_inputs = reference_tensors + if inputs_require_grads: + recording_inputs = do_input_map(lambda t: t.clone().requires_grad_(), reference_tensors) + flattened_recording_inputs = flatten_inputs(recording_inputs) + else: + recording_inputs = reference_tensors + + # `check_trace` is set to False because check_trace is run with @no_grad + # Also, `checkTrace` already does all the checks + # against python function + ge = torch.jit.trace(func, input_tensors, check_tolerance=check_tolerance, + _force_outplace=_force_outplace, check_trace=False) + + if export_import: + ge = self.getExportImportCopy(ge) + + if verbose: + print(ge.graph) + + # test no gradients case + outputs = func(*nograd_inputs) + outputs_ge = ge(*nograd_inputs) + self.assertEqual(outputs, outputs_ge) + + # test gradients case + outputs = func(*recording_inputs) + if inputs_require_grads: + grads = torch.autograd.grad(allSum(outputs), flattened_recording_inputs, + allow_unused=allow_unused) + + outputs_ge = ge(*recording_inputs) + if inputs_require_grads: + grads_ge = torch.autograd.grad(allSum(outputs_ge), flattened_recording_inputs, + allow_unused=allow_unused) + self.assertEqual(outputs, outputs_ge) + if inputs_require_grads: + self.assertEqual(grads, grads_ge, atol=grad_atol, rtol=grad_rtol) + + # test the grad grad case + outputs = func(*recording_inputs) + l1 = allSum(outputs) + if inputs_require_grads: + grads = torch.autograd.grad(l1, flattened_recording_inputs, create_graph=True, + allow_unused=allow_unused) + if inputs_require_grads: + l2 = (allSum(grads) * l1) + grads2 = torch.autograd.grad(l2, flattened_recording_inputs, allow_unused=allow_unused) + + if inputs_require_grads: + recording_inputs = do_input_map(lambda t: Variable(t, requires_grad=True), reference_tensors) + flattened_recording_inputs = flatten_inputs(recording_inputs) + + outputs_ge = ge(*recording_inputs) + l1_ge = allSum(outputs_ge) + if inputs_require_grads: + grads_ge = torch.autograd.grad( + l1_ge, flattened_recording_inputs, create_graph=True, allow_unused=allow_unused) + + if inputs_require_grads: + l2_ge = (allSum(grads_ge) * l1_ge) + grads2_ge = torch.autograd.grad(l2_ge, flattened_recording_inputs, allow_unused=allow_unused) + + self.assertEqual(outputs, outputs_ge) + if inputs_require_grads: + self.assertEqual(grads, grads_ge, atol=grad_atol, rtol=grad_rtol) + for g2, g2_ge in zip(grads2, grads2_ge, strict=True): + if g2 is None and g2_ge is None: + continue + self.assertEqual(g2, g2_ge, atol=8e-4, rtol=8e-4) + + return ge + + def checkModule(self, nn_module, args): + """ + Check that a nn.Module's results in Script mode match eager and that it + can be exported + """ + sm = torch.jit.script(nn_module) + + with freeze_rng_state(): + eager_out = nn_module(*args) + + with freeze_rng_state(): + script_out = sm(*args) + + self.assertEqual(eager_out, script_out) + self.assertExportImportModule(sm, args) + + return sm + +class NoTracerWarnContextManager: + def __enter__(self): + self.prev = torch._C._jit_get_tracer_state_warn() + torch._C._jit_set_tracer_state_warn(False) + + def __exit__(self, *args): + torch._C._jit_set_tracer_state_warn(self.prev) + +@contextmanager +def inline_everything_mode(should_inline): + old = torch._C._jit_get_inline_everything_mode() + torch._C._jit_set_inline_everything_mode(should_inline) + try: + yield + finally: + torch._C._jit_set_inline_everything_mode(old) + +@contextmanager +def set_fusion_group_inlining(inlining): + old = torch._C._debug_get_fusion_group_inlining() + torch._C._debug_set_fusion_group_inlining(inlining) + try: + yield + finally: + torch._C._debug_set_fusion_group_inlining(old) + +# note: not re-entrant, use unnested only +@contextmanager +def disable_autodiff_subgraph_inlining(enabled=True): + torch._C._debug_set_autodiff_subgraph_inlining(not enabled) + try: + yield + finally: + torch._C._debug_set_autodiff_subgraph_inlining(True) + +def _inline_everything(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with inline_everything_mode(True): + fn(*args, **kwargs) + return wrapper + +# this exists for forward compatibility reasons temporarily. +# TODO(suo) remove +def _tmp_donotuse_dont_inline_everything(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + with inline_everything_mode(False): + fn(*args, **kwargs) + return wrapper + +# make it easy to quickly define/trace a function for these tests +def _trace(*args, **kwargs): + def wrapper(func): + return torch.jit.trace(func, args, **kwargs) + return wrapper + + +def enable_cpu_fuser(fn): + def wrapper(*args, **kwargs): + torch._C._jit_override_can_fuse_on_cpu_legacy(True) + torch._C._jit_override_can_fuse_on_cpu(True) + torch._C._jit_set_te_must_use_llvm_cpu(False) + try: + fn(*args, **kwargs) + finally: + torch._C._jit_override_can_fuse_on_cpu_legacy(False) + torch._C._jit_override_can_fuse_on_cpu(False) + torch._C._jit_set_te_must_use_llvm_cpu(True) + return wrapper + + +def enable_cpu_fuser_if(cond): + if cond: + return enable_cpu_fuser + else: + def noop_fuser(fn): + def wrapper(*args, **kwargs): + return fn(*args, **kwargs) + return wrapper + return noop_fuser + +def get_forward(c): + return c._get_method('forward') + +def get_forward_graph(c): + return c._get_method('forward').graph + +def get_module_method(m, module, method): + return m._c.getattr(module)._get_method(method) + +def attrs_with_prefix(module, prefix): + return [x for x, _ in module._modules._c.items() + if x.startswith(prefix)] + +def warmup_backward(f, *args): + profiling_count = 3 + results = [] + for _ in range(profiling_count): + if len(args) > 0: + r = torch.autograd.grad(f, *args) + results.append(r) + else: + f.backward(retain_graph=True) + + return results + +# TODO: Remove me once https://bugs.python.org/issue42666 is resolved +def make_global(*args): + for arg in args: + setattr(sys.modules[arg.__module__], arg.__name__, arg) + +# Helper function to eval Python3 code without causing a syntax error for +# this file under py2 +def _get_py3_code(code, fn_name): + with tempfile.TemporaryDirectory() as tmp_dir: + script_path = os.path.join(tmp_dir, 'script.py') + with open(script_path, 'w') as f: + f.write(code) + spec = importlib.util.spec_from_file_location(fn_name, script_path) + module = importlib.util.module_from_spec(spec) + loader = spec.loader + if not isinstance(loader, Loader): # Assert type to meet MyPy requirement + raise AssertionError(f"Expected loader to be Loader, got {type(loader)}") + loader.exec_module(module) + fn = getattr(module, fn_name) + return fn + +class TensorExprTestOptions: + def __init__(self) -> None: + self.old_profiling_executor = torch._C._jit_set_profiling_executor(True) + self.old_profiling_mode = torch._C._get_graph_executor_optimize(True) + + self.old_cpu_fuser_state = torch._C._jit_can_fuse_on_cpu() + self.old_gpu_fuser_state = torch._C._jit_can_fuse_on_gpu() + torch._C._jit_override_can_fuse_on_cpu(True) + torch._C._jit_override_can_fuse_on_gpu(True) + self.texpr_fuser_state = torch._C._jit_texpr_fuser_enabled() + torch._C._jit_set_texpr_fuser_enabled(True) + self.old_fusion_inlining = torch._C._debug_get_fusion_group_inlining() + torch._C._debug_set_fusion_group_inlining(False) + self.old_te_must_use_llvm_cpu = torch._C._jit_get_te_must_use_llvm_cpu() + torch._C._jit_set_te_must_use_llvm_cpu(False) + + def restore(self): + torch._C._jit_set_profiling_executor(self.old_profiling_executor) + torch._C._get_graph_executor_optimize(self.old_profiling_mode) + + torch._C._jit_set_texpr_fuser_enabled(self.texpr_fuser_state) + torch._C._jit_override_can_fuse_on_gpu(self.old_gpu_fuser_state) + torch._C._jit_override_can_fuse_on_cpu(self.old_cpu_fuser_state) + torch._C._debug_set_fusion_group_inlining(self.old_fusion_inlining) + torch._C._jit_set_te_must_use_llvm_cpu(self.old_te_must_use_llvm_cpu) + +def clone_inputs(args): + inputs: list[torch.Tensor | list[torch.Tensor]] = [] + + for arg in args: + if isinstance(arg, torch.Tensor): + inputs.append(arg.detach().clone()) + elif is_iterable_of_tensors(arg): + inputs.append([t.detach().clone() for t in arg]) + else: + inputs.append(arg) + + return inputs + +def get_traced_sample_variant_pairs(device, dtype, op): + # tuples of (variant, sample) + outputs: list[tuple[Any, Any]] = [] + + samples = op.sample_inputs(device, dtype) + + # Acquires variants to test + func = op.get_op() + method = op.get_method() + variants = { + # TODO: inplace tests currently fail, fix and add inplace variant + 'function': func, 'method': method, + } + + # TODO: find better way to standardize on op registration itself.. + has_fake_function = op.name in ["resize_", 'resize_as_'] + + if has_fake_function: + variants = {'method': getattr(torch.Tensor, op.name)} + + # In eager mode, these ops can take (Tensor, bool) args; but in + # JIT they can only take (Tensor, Scalar), and bool is not a + # scalar in the JIT type system. So to test these in JIT, the bool + # is converted to an int for the test. + ops_with_unsupported_bool_args = [ + { + "name": "div_floor_rounding", + "arg_idx": [0], + }, + { + "name": "div_no_rounding_mode", + "arg_idx": [0], + }, + { + "name": "div_trunc_rounding", + "arg_idx": [0], + }, + { + "name": "index_fill", + "arg_idx": [2], + }, + { + "name": "full_like", + "arg_idx": [0], + }, + { + "name": "mul", + "arg_idx": [0], + }, + { + "name": "new_full", + "arg_idx": [1], + }, + ] + + # doesn't support tracing + if has_fake_function: + return outputs + + for sample in samples: + for variant in variants.values(): + if variant is None: + continue + + if is_lambda(variant): + continue + + matching_ops = filter(lambda x: op.formatted_name == x["name"], ops_with_unsupported_bool_args) + for op_data in matching_ops: + for idx in op_data["arg_idx"]: + args = list(sample.args) + if len(sample.args) > idx and isinstance(sample.args[idx], bool): + args[idx] = int(args[idx]) + sample.args = tuple(args) + + outputs.append((variant, sample)) + + return outputs + +# types.LambdaType gave false positives +def is_lambda(lamb): + LAMBDA = lambda: 0 # noqa: E731 + return isinstance(lamb, type(LAMBDA)) and lamb.__name__ == LAMBDA.__name__ diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/logging_tensor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/logging_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..0edbe8ced6dd4be78a1080ea2713f18a284e28bc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/logging_tensor.py @@ -0,0 +1,167 @@ +# mypy: ignore-errors + +import torch +from torch.utils._pytree import tree_map +from collections.abc import Iterator +import logging +import contextlib +import itertools +from torch.utils._dtype_abbrs import dtype_abbrs as _dtype_abbrs +from torch.utils._python_dispatch import TorchDispatchMode +from torch.utils.weak import WeakTensorKeyDictionary +import functools +from torch._C._profiler import gather_traceback, symbolize_tracebacks + +logger = logging.getLogger("LoggingTensor") + +# How the chain of calls works for LoggingTensor: +# 1. Call torch.sin +# 2. Attempt __torch_function__. In LoggingTensor torch function is disabled so we bypass it entirely +# 3. Enter dispatcher, wind your way through Autograd +# 4. Hit Python dispatch key, call __torch_dispatch__ + +# This Tensor can work with autograd in two ways: +# - The wrapped Tensor does not require gradients. In that case, the LoggingTensor +# can require gradients if the user asks for it as a constructor kwarg. +# - The wrapped Tensor can require gradients. In that case autograd will be tracked +# for the wrapped Tensor and the LoggingTensor itself cannot require gradients. +# WARNING: We allow these two possibilities for testing purposes. You should NEVER use both in a single +# test or you might get surprising behavior. + +# TODO: TensorBase should work +class LoggingTensor(torch.Tensor): + elem: torch.Tensor + + __slots__ = ['elem'] + + context = contextlib.nullcontext + + @staticmethod + def __new__(cls, elem, *args, **kwargs): + # The wrapping tensor (LoggingTensor) shouldn't hold any + # memory for the class in question, but it should still + # advertise the same device as before + r = torch.Tensor._make_wrapper_subclass( + cls, elem.size(), + strides=elem.stride(), storage_offset=elem.storage_offset(), + # TODO: clone storage aliasing + dtype=elem.dtype, layout=elem.layout, + device=elem.device, requires_grad=kwargs.get("requires_grad", False) + ) + # ...the real tensor is held as an element on the tensor. + r.elem = elem.detach() if r.requires_grad else elem + return r + + def __repr__(self): + return super().__repr__(tensor_contents=f"{self.elem}") + + @classmethod + def __torch_dispatch__(cls, func, types, args=(), kwargs=None): + def unwrap(e): + return e.elem if isinstance(e, cls) else e + + def wrap(e): + return cls(e) if isinstance(e, torch.Tensor) else e + + with cls.context(): + rs = tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs))) + logging.getLogger("LoggingTensor").info(f"{func.__module__}.{func.__name__}", args, kwargs, rs) # noqa: G004 + return rs + +class LoggingTensorMode(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + rs = func(*args, **kwargs) + logging.getLogger("LoggingTensor").info(f"{func.__module__}.{func.__name__}", args, kwargs, rs) # noqa: G004 + return rs + +class LoggingTensorReentrant(LoggingTensor): + context = torch.overrides.enable_reentrant_dispatch + +# https://stackoverflow.com/questions/36408496/python-logging-handler-to-append-to-list +class LoggingTensorHandler(logging.Handler): + def __init__( + self, log_list: list[str], use_shortid_for_all_tensors: bool, + with_type: bool, tracebacks_list: list | None) -> None: + logging.Handler.__init__(self) + self.log_list = log_list + self.use_shortid_for_all_tensors = use_shortid_for_all_tensors + self.tracebacks_list = tracebacks_list + self.memo = WeakTensorKeyDictionary() + self.next_id = 0 + self.with_type = with_type + + def _shortid(self, t: torch.Tensor) -> int: + if t not in self.memo: + self.memo[t] = self.next_id + self.next_id += 1 + return self.memo[t] + + def _fmt(self, a: object, with_type: bool = False) -> str: + cond_cls = torch.Tensor if self.use_shortid_for_all_tensors else LoggingTensor + if isinstance(a, cond_cls): + maybe_type = "" + if with_type and self.with_type: + maybe_type = f": {_dtype_abbrs[a.dtype]}[{', '.join(map(str, a.shape))}]" + x = f"${self._shortid(a)}{maybe_type}" + return x + else: + return repr(a) + + def emit(self, record): + fmt_args = ", ".join( + itertools.chain( + (str(tree_map(self._fmt, a)) for a in record.args[0]), + (f"{k}={str(tree_map(self._fmt, v))}" for k, v in record.args[1].items()), + ) + ) + fmt_rets = tree_map(functools.partial(self._fmt, with_type=True), record.args[2]) + self.log_list.append(f'{fmt_rets} = {record.msg}({fmt_args})') + if self.tracebacks_list is not None: + self.tracebacks_list.append(record.traceback) + +def log_input(name: str, var: object) -> None: + logger.info("input", (name,), {}, var) # noqa: PLE1205 + +class GatherTraceback(logging.Filter): + def __init__(self, python=True, script=True, cpp=False): + self.python = python + self.script = script + self.cpp = cpp + + def filter(self, record): + record.traceback = gather_traceback(python=self.python, script=self.script, cpp=self.cpp) + return True + +@contextlib.contextmanager +def capture_logs(is_mode=False, python_tb=False, script_tb=False, cpp_tb=False) -> Iterator[list[str]]: + collect_traceback = python_tb or script_tb or cpp_tb + log_list: list[str] = [] + tracebacks_list: list[str] = [] + handler = LoggingTensorHandler( + log_list, + with_type=True, + use_shortid_for_all_tensors=is_mode, + tracebacks_list=tracebacks_list if collect_traceback else None + ) + logger.addHandler(handler) + logger.setLevel(logging.INFO) + logger.propagate = False + if collect_traceback: + logger.addFilter(GatherTraceback(python=python_tb, script=script_tb, cpp=cpp_tb)) + try: + if collect_traceback: + yield log_list, tracebacks_list + else: + yield log_list + finally: + symbolized_tracebacks = symbolize_tracebacks(tracebacks_list) + tracebacks_list.clear() + tracebacks_list.extend(symbolized_tracebacks) + logger.removeHandler(handler) + +@contextlib.contextmanager +def capture_logs_with_logging_tensor_mode(python_tb=False, script_tb=False, cpp_tb=False): + with LoggingTensorMode(), capture_logs(True, python_tb, script_tb, cpp_tb) as logs: + yield logs diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/logging_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/logging_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1e1ecf8f4f707c9b3712a6fb738fc9ce1467b835 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/logging_utils.py @@ -0,0 +1,243 @@ +# mypy: ignore-errors + +import torch._dynamo.test_case +import unittest.mock +import os +import contextlib +import torch._logging +import torch._logging._internal +from contextlib import AbstractContextManager +from collections.abc import Callable +from torch._dynamo.utils import LazyString +from torch._inductor import config as inductor_config +import logging +import io + +@contextlib.contextmanager +def preserve_log_state(): + prev_state = torch._logging._internal._get_log_state() + torch._logging._internal._set_log_state(torch._logging._internal.LogState()) + try: + yield + finally: + torch._logging._internal._set_log_state(prev_state) + torch._logging._internal._init_logs() + +def log_settings(settings): + exit_stack = contextlib.ExitStack() + settings_patch = unittest.mock.patch.dict(os.environ, {"TORCH_LOGS": settings}) + exit_stack.enter_context(preserve_log_state()) + exit_stack.enter_context(settings_patch) + torch._logging._internal._init_logs() + return exit_stack + +def log_api(**kwargs): + exit_stack = contextlib.ExitStack() + exit_stack.enter_context(preserve_log_state()) + torch._logging.set_logs(**kwargs) + return exit_stack + + +def kwargs_to_settings(**kwargs): + INT_TO_VERBOSITY = {10: "+", 20: "", 40: "-"} + + settings = [] + + def append_setting(name, level): + if isinstance(name, str) and isinstance(level, int) and level in INT_TO_VERBOSITY: + settings.append(INT_TO_VERBOSITY[level] + name) + return + else: + raise ValueError("Invalid value for setting") + + for name, val in kwargs.items(): + if isinstance(val, bool): + settings.append(name) + elif isinstance(val, int): + append_setting(name, val) + elif isinstance(val, dict) and name == "modules": + for module_qname, level in val.items(): + append_setting(module_qname, level) + else: + raise ValueError("Invalid value for setting") + + return ",".join(settings) + + +# Note on testing strategy: +# This class does two things: +# 1. Runs two versions of a test: +# 1a. patches the env var log settings to some specific value +# 1b. calls torch._logging.set_logs(..) +# 2. patches the emit method of each setup handler to gather records +# that are emitted to each console stream +# 3. passes a ref to the gathered records to each test case for checking +# +# The goal of this testing in general is to ensure that given some settings env var +# that the logs are setup correctly and capturing the correct records. +def make_logging_test(**kwargs): + def wrapper(fn): + @inductor_config.patch({"fx_graph_cache": False}) + def test_fn(self): + + torch._dynamo.reset() + records = [] + # run with env var + if len(kwargs) == 0: + with self._handler_watcher(records): + fn(self, records) + else: + with log_settings(kwargs_to_settings(**kwargs)), self._handler_watcher(records): + fn(self, records) + + # run with API + torch._dynamo.reset() + records.clear() + with log_api(**kwargs), self._handler_watcher(records): + fn(self, records) + + + return test_fn + + return wrapper + +def make_settings_test(settings): + def wrapper(fn): + def test_fn(self): + torch._dynamo.reset() + records = [] + # run with env var + with log_settings(settings), self._handler_watcher(records): + fn(self, records) + + return test_fn + + return wrapper + +class LoggingTestCase(torch._dynamo.test_case.TestCase): + @classmethod + def setUpClass(cls): + super().setUpClass() + cls._exit_stack.enter_context( + unittest.mock.patch.dict(os.environ, {"___LOG_TESTING": ""}) + ) + cls._exit_stack.enter_context( + unittest.mock.patch("torch._dynamo.config.suppress_errors", True) + ) + cls._exit_stack.enter_context( + unittest.mock.patch("torch._dynamo.config.verbose", False) + ) + + @classmethod + def tearDownClass(cls): + cls._exit_stack.close() + torch._logging._internal.log_state.clear() + torch._logging._init_logs() + + def hasRecord(self, records, m): + return any(m in r.getMessage() for r in records) + + def getRecord(self, records, m): + record = None + for r in records: + # NB: not r.msg because it looks like 3.11 changed how they + # structure log records + if m in r.getMessage(): + self.assertIsNone( + record, + msg=LazyString( + lambda: f"multiple matching records: {record} and {r} among {records}" + ), + ) + record = r + if record is None: + self.fail(f"did not find record with {m} among {records}") + return record + + # This patches the emit method of each handler to gather records + # as they are emitted + def _handler_watcher(self, record_list): + exit_stack = contextlib.ExitStack() + + def emit_post_hook(record): + nonlocal record_list + record_list.append(record) + + # registered logs are the only ones with handlers, so patch those + for log_qname in torch._logging._internal.log_registry.get_log_qnames(): + logger = logging.getLogger(log_qname) + num_handlers = len(logger.handlers) + self.assertLessEqual( + num_handlers, + 2, + "All pt2 loggers should only have at most two handlers (debug artifacts and messages above debug level).", + ) + + self.assertGreater(num_handlers, 0, "All pt2 loggers should have more than zero handlers") + + for handler in logger.handlers: + old_emit = handler.emit + + def new_emit(record): + old_emit(record) + emit_post_hook(record) + + exit_stack.enter_context( + unittest.mock.patch.object(handler, "emit", new_emit) + ) + + return exit_stack + + +def logs_to_string(module, log_option): + """Example: + logs_to_string("torch._inductor.compile_fx", "post_grad_graphs") + returns the output of TORCH_LOGS="post_grad_graphs" from the + torch._inductor.compile_fx module. + """ + log_stream = io.StringIO() + handler = logging.StreamHandler(stream=log_stream) + + @contextlib.contextmanager + def tmp_redirect_logs(): + try: + logger = torch._logging.getArtifactLogger(module, log_option) + logger.addHandler(handler) + yield + finally: + logger.removeHandler(handler) + + def ctx_manager(): + exit_stack = log_settings(log_option) + exit_stack.enter_context(tmp_redirect_logs()) + return exit_stack + + return log_stream, ctx_manager + + +def multiple_logs_to_string(module: str, *log_options: str) -> tuple[list[io.StringIO], Callable[[], AbstractContextManager[None]]]: + """Example: + multiple_logs_to_string("torch._inductor.compile_fx", "pre_grad_graphs", "post_grad_graphs") + returns the output of TORCH_LOGS="pre_graph_graphs, post_grad_graphs" from the + torch._inductor.compile_fx module. + """ + log_streams = [io.StringIO() for _ in range(len(log_options))] + handlers = [logging.StreamHandler(stream=log_stream) for log_stream in log_streams] + + @contextlib.contextmanager + def tmp_redirect_logs(): + loggers = [torch._logging.getArtifactLogger(module, option) for option in log_options] + try: + for logger, handler in zip(loggers, handlers, strict=True): + logger.addHandler(handler) + yield + finally: + for logger, handler in zip(loggers, handlers, strict=True): + logger.removeHandler(handler) + + def ctx_manager() -> AbstractContextManager[None]: + exit_stack = log_settings(", ".join(log_options)) + exit_stack.enter_context(tmp_redirect_logs()) + return exit_stack # type: ignore[return-value] + + return log_streams, ctx_manager diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..97c38f3560625213fbd59d09a9cfd22bad26ba04 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/__init__.py @@ -0,0 +1,4 @@ +# mypy: ignore-errors + +import torch.testing._internal.opinfo.core +import torch.testing._internal.opinfo.definitions diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/core.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/core.py new file mode 100644 index 0000000000000000000000000000000000000000..acd9cad91644136e7fb706b4b957d202b27e25d7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/core.py @@ -0,0 +1,3293 @@ +# mypy: ignore-errors + +import collections +import collections.abc +import contextlib +import logging +import math +import operator +import unittest +from abc import ABC, abstractmethod +from collections.abc import Callable, Iterable +from dataclasses import asdict, dataclass, field +from enum import Enum +from functools import partial +from itertools import product +from typing import Any, TypeVar + +import torch +from torch.testing import make_tensor +from torch.testing._internal.common_device_type import ( + skipCPUIfNoFFT, + tol, + toleranceOverride, +) +from torch.testing._internal.common_dtype import ( + _dispatch_dtypes, + floating_and_complex_types, + floating_and_complex_types_and, + floating_types, + get_all_dtypes, +) +from torch.testing._internal.common_utils import ( + extract_test_fn, + IS_FBCODE, + is_iterable_of_tensors, + noncontiguous_like, + OPINFO_SAMPLE_INPUT_INDEX, + TEST_WITH_ROCM, + torch_to_numpy_dtype_dict, + TrackedInputIter, + USE_PYTEST, +) +from torch.testing._internal.opinfo import utils +from torchgen.utils import dataclass_repr + + +# setup logging +log = logging.getLogger(__name__) + +# Reasonable testing sizes for dimensions +L = 20 +M = 10 +S = 5 +XS = 3 + +# Unique value to distinguish default from anything else +_NOTHING = object() + + +# Extension of getattr to support qualified names +# e.g. _getattr_qual(torch, 'linalg.norm') -> torch.linalg.norm +def _getattr_qual(obj, name, default=_NOTHING): + try: + for path in name.split("."): + obj = getattr(obj, path) + return obj + except AttributeError: + if default is not _NOTHING: + return default + else: + raise + + +class DecorateInfo: + """Describes which test, or type of tests, should be wrapped in the given + decorators when testing an operator. Any test that matches all provided + arguments will be decorated. The decorators will only be applied if the + active_if argument is True.""" + + __slots__ = [ + "decorators", + "cls_name", + "test_name", + "device_type", + "dtypes", + "active_if", + ] + + def __init__( + self, + decorators, + cls_name=None, + test_name=None, + *, + device_type=None, + dtypes=None, + active_if=True, + ): + self.decorators = ( + list(decorators) + if isinstance(decorators, collections.abc.Sequence) + else [decorators] + ) + self.cls_name = cls_name + self.test_name = test_name + self.device_type = device_type + self.dtypes = dtypes + self.active_if = active_if + + # Validate dtypes + if self.dtypes is not None: + for dtype in self.dtypes: + if not isinstance(dtype, torch.dtype): + raise AssertionError(f"Expected torch.dtype, got {type(dtype)}") + + def is_active(self, cls_name, test_name, device_type, dtype, param_kwargs): + return ( + self.active_if + and (self.cls_name is None or self.cls_name == cls_name) + and (self.test_name is None or self.test_name == test_name) + and (self.device_type is None or self.device_type == device_type) + and (self.dtypes is None or dtype in self.dtypes) + # Support callables over kwargs to determine if the decorator is active. + and ( + self.active_if(param_kwargs) + if isinstance(self.active_if, Callable) + else self.active_if + ) + ) + + +# FIXME +# Note: historically the 'input' kwarg had to be a Tensor or TensorList, but we are trying +# to support scalar inputs, too. Some tests still depend on 'input' being a Tensor +# or TensorList, however. +class SampleInput: + """Represents sample inputs to a function.""" + + __slots__ = [ + "input", + "args", + "kwargs", + "output_process_fn_grad", + "broadcasts_input", + "name", + ] + + def __init__( + self, + input, + *var_args, + args=None, + kwargs=None, + output_process_fn_grad=None, + broadcasts_input=None, + name=None, + **var_kwargs, + ): + # input is the first input to the op and is typically either a Tensor or TensorList (Sequence[Tensor]). + # This follows the typical pattern where for Tensor inputs op(t, ...) = t.op(...). + self.input = input + + # Allow calling either as SampleInput(input, args=args, kwargs=kwargs), or as + # SampleInput(input, *args, **kwargs) but not to mix the two forms + if args is not None or kwargs is not None: + if var_args or var_kwargs: + raise AssertionError( + "A SampleInput can be constructed 'naturally' with *args and **kwargs or by " + "explicitly setting the 'args' and 'kwargs' parameters, but the two " + "methods of construction cannot be mixed!" + ) + elif var_args or var_kwargs: + if not ( + output_process_fn_grad is None + and broadcasts_input is None + and name is None + ): + raise AssertionError( + "A SampleInput constructed 'naturally' with *args and **kwargs " + "cannot specify additional metadata in keyword arguments" + ) + + self.args = args if args is not None else var_args + if not isinstance(self.args, tuple): + raise AssertionError(f"Expected args to be tuple, got {type(self.args)}") + self.kwargs = kwargs if kwargs is not None else var_kwargs + if not isinstance(self.kwargs, dict): + raise AssertionError(f"Expected kwargs to be dict, got {type(self.kwargs)}") + + self.output_process_fn_grad = ( + output_process_fn_grad + if output_process_fn_grad is not None + else lambda x: x + ) + self.name = name if name is not None else "" + + # Specifies if `self.input` is broadcasted or not, + # given that the operator supports broadcasting. + # This field is used to verify the behavior for inplace variant. + # + # If a SampleInput is marked with `broadcasts_input=True`, + # it is verified that we get a `RuntimeError` with this sample, + # and inplace variant. Also inplace grad{grad} tests are skipped, + # for such inputs (as they will error out otherwise). + self.broadcasts_input = ( + broadcasts_input if broadcasts_input is not None else False + ) + + def with_metadata( + self, *, output_process_fn_grad=None, broadcasts_input=None, name=None + ): + if output_process_fn_grad is not None: + self.output_process_fn_grad = output_process_fn_grad + if broadcasts_input is not None: + self.broadcasts_input = broadcasts_input + if name is not None: + self.name = name + return self + + def _repr_helper(self, formatter): + # Helper function to return the details of the SampleInput as `str` + # It consolidates all the fields of SampleInput and allows, + # formatting the fields like `input`, `args`, etc with `formatter` + # callable to customize the representation. + # Look at `summary` method for example. + arguments = [ + f"input={formatter(self.input)}", + f"args={formatter(self.args)}", + f"kwargs={formatter(self.kwargs)}", + f"broadcasts_input={self.broadcasts_input}", + f"name={repr(self.name)}", + ] + + return f"SampleInput({', '.join(a for a in arguments if a is not None)})" + + def __repr__(self): + return self._repr_helper(lambda x: x) + + def summary(self): + # Returns the SampleInput details in a more + # friendly format. + # It formats `Tensor` and `TensorList` + # in a more condensed representation. + def formatter(arg): + # Format any instance of `Tensor` (standalone, in list, or in dict) + # by Tensor[TensorShape] + # Eg. Tensor with shape (3, 4) is formatted as Tensor[3, 4] + if isinstance(arg, torch.Tensor): + shape = str(tuple(arg.shape)) + dtype = str(arg.dtype) + device = str(arg.device) + contiguity_suffix = "" + # NB: sparse CSR tensors annoyingly return is_sparse=False + is_sparse = arg.is_sparse or arg.layout == torch.sparse_csr + if not is_sparse and not arg.is_contiguous(): + contiguity_suffix = ", contiguous=False" + return f'Tensor[size={shape}, device="{device}", dtype={dtype}{contiguity_suffix}]' + elif isinstance(arg, dict): + return {k: formatter(v) for k, v in arg.items()} + elif is_iterable_of_tensors(arg): + return "TensorList[" + ", ".join(map(formatter, arg)) + "]" + elif isinstance(arg, (list, tuple)): # Handle list, tuple + return "(" + ",".join(map(formatter, arg)) + ")" + + return repr(arg) + + return self._repr_helper(formatter) + + # Applies the transform f(t) -> t to each tensor and dtype in the SampleInput + def transform(self, f): + def tt(t): + def _tt(t): + with torch.no_grad(): + return f(t) + + if isinstance(t, torch.Tensor): + return _tt(t) + elif isinstance(t, torch.dtype): + return _tt(t) + elif isinstance(t, list): + return list(map(tt, t)) + elif isinstance(t, tuple): + return tuple(map(tt, t)) + elif isinstance(t, dict): + return {k: tt(v) for k, v in t.items()} + else: + return t + + sample_tt_input, tt_args, tt_kwargs = ( + tt(self.input), + tt(self.args), + tt(self.kwargs), + ) + + # Note the transformed SampleInput assumes metadata like output_process_fn_grad is still valid! + return SampleInput( + sample_tt_input, + args=tt_args, + kwargs=tt_kwargs, + output_process_fn_grad=self.output_process_fn_grad, + broadcasts_input=self.broadcasts_input, + name=self.name + "_transformed", + ) + + # Returns the NumPy version of the sample input object in the form of a tuple: (input, args, kwargs) + # Converts tensors to ndarrays by calling .detach().cpu().numpy() on them + # Converts dtypes by remapping them using torch_to_numpy_dtype_dict + def numpy(self): + def to_numpy(t): + if isinstance(t, torch.Tensor): + if t.dtype is torch.bfloat16: + return t.detach().cpu().to(torch.float32).numpy() + if t.dtype is torch.chalf: + return t.detach().cpu().to(torch.cfloat).numpy() + return t.detach().cpu().numpy() + elif isinstance(t, torch.dtype): + return torch_to_numpy_dtype_dict[t] + + return t + + return self.transform(to_numpy) + + def noncontiguous(self): + def to_noncontiguous(t): + if isinstance(t, torch.Tensor): + return noncontiguous_like(t) + elif isinstance(t, torch.dtype): + return t + + return t + + return self.transform(to_noncontiguous) + + +NumericsFilter = collections.namedtuple("NumericsFilter", ["condition", "safe_val"]) + + +class ErrorInput: + """ + A SampleInput that will cause the operation to throw an error plus information + about the resulting error. + """ + + __slots__ = ["sample_input", "error_type", "error_regex"] + + def __init__(self, sample_input, *, error_type=RuntimeError, error_regex): + self.sample_input = sample_input + self.error_type = error_type + self.error_regex = error_regex + + +class AliasInfo: + """Class holds alias information. For example, torch.abs -> + torch.absolute, torch.Tensor.absolute, torch.Tensor.absolute_ + """ + + def __init__(self, alias_name): + self.name = alias_name + self.op = _getattr_qual(torch, alias_name) + self.method_variant = getattr(torch.Tensor, alias_name, None) + self.inplace_variant = getattr(torch.Tensor, alias_name + "_", None) + + def __call__(self, *args, **kwargs): + return self.op(*args, **kwargs) + + +# Note [OpInfos] +# ~~~~~~~~~~~~~~ +# +# The majority of this note was written shortly after the PyTorch 1.9 release. +# If you notice it's out-of-date or think it could be improved then please +# file an issue. +# +# See also: the OpInfo tracker (https://github.com/pytorch/pytorch/issues/54261) +# See also: "Writing Test Templates" in common_device_type.py to learn how to +# parametrize a test template using OpInfos. +# See also: PyTorch's GitHub wiki on running and writing tests +# https://github.com/pytorch/pytorch/wiki/Running-and-writing-tests +# See also: ModuleInfos, OpInfo's sister class, defined in common_modules.py +# +# An OpInfo is a collection of metadata related to a PyTorch operator. This +# metadata is used to generate tests that validate properties of the operator, +# like if it implements the correct gradient formula. +# +# WHY OPINFOS? +# ~~~~~~~~~~~~ +# +# OpInfos are principally intended to do three things: +# +# 1) to allow systematic testing over all PyTorch's operators +# 2) to simplify operating testing by autogenerating many tests +# 3) to allow systems (like autograd, torchscript, fx, nnc...) to test +# against every PyTorch operator +# +# All these goals are still a work in progress. Not every operator has an +# OpInfo, and some operator tests that could be automatically generated +# still have to be written manually. +# +# It's helpful to understand that OpInfos are both about test simplification and +# modularity. PyTorch is a complicated framework with many interrelated systems, +# too many for any one person to keep track of. An OpInfo can be thought of as the +# interface between an operator implementer and those other systems. Instead of +# requiring the implementer of torch.foo understand how to test its forward +# mode AD or NNC support that's typically handled automatically just by +# defining an OpInfo. +# +# It's often surprising to OpInfo writers that just implementing an OpInfo +# typically can't verify an operator is actually implemented correctly: +# +# "If an OpInfo doesn't validate my op works as expected, what's the point +# of it?" +# +# But the point of is the above. OpInfos are intended to let you focus on testing +# the operator logic you're familiar with instead of having to write tests for +# how the operator interacts with each of PyTorch's many systems. +# +# And, OK, it turns out that SOMETIMES just writing an OpInfo DOES +# validate your op works as expected, but that's only in special +# cases. See below for details. +# +# WHAT'S AN OPINFO? +# ~~~~~~~~~~~~~~~~~ +# +# So what is an OpInfo? It's a Python class that describes an operator's properties, +# like which dtypes it supports on the CPU and whether it has any aliases. +# These properties can be divided into three categories: +# +# 1) Metadata describing the operator, like the operator's name and if it +# "supports" the out kwarg. +# 2) Test directives, like "skips" that tell the test suite to skip some +# tests. +# 3) A "sample inputs" function that generates valid inputs for the operator. +# +# OpInfo attributes are described in more detail below. +# +# THE SAMPLE INPUTS FUNCTION +# ~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# The "sample inputs" function merits special elaboration. This function is +# crucial to testing with OpInfos. A typical OpInfo test has to treat the operator +# as a black box. There's no structure for the test to understand or exploit. +# Without "sample inputs" it wouldn't even know how to call the OpInfo's +# operator. The sample input function saves the day by providing different +# "SampleInputs" that can be used to call the operator. A sample input +# function should have the following signature: +# +# def sample_inputs_foo(op_info, device, dtype, requires_grad, **kwargs): +# +# And should return an iterable of SampleInputs (see the class description +# above). Each SampleInput defines an "input", "args", "kwargs", an +# "output_process_fn_grad" function, the "broadcasts_input" bool and a +# "name". +# +# All the "sample_inputs" functions are invoked within a `torch.no_grad()` +# environment for efficiency and correctness. As such remember to set the +# "requires_grad" flag on the inputs **after** performing any transformations +# on them. +# +# The "input" is the first argument to the operator, or the tensor that +# the method or inplace variants of the operator should be called on, and +# should be on the requested device, of the requested dtype, and its +# requires_grad attribute should be set to the requires_grad argument. +# +# "args" should contain positional arguments, and "kwargs" keyword arguments. +# +# "output_process_fn_grad" has an interesting name. It's a function that maps +# the operator's output (when given the input, args, and kwargs) to the +# portion of the output to gradcheck. For example, consider an operator +# like torch.linalg.slogdet +# (https://pytorch.org/docs/main/generated/torch.linalg.slogdet.html). +# This operator returns a tuple of two tensors, but the first tensor +# cannot be backwarded through. Its "output_process_fn_grad" filters +# this output tuple to just the second argument, which we can call backward +# on. Functions that produce a single tensor can ignore this argument. +# +# "broadcasts_input" is a bool indicated if the SampleInput causes the operator +# to broadcast the "input" argument. This is important for tests to understand +# because inplace variants of operations throw a runtime error if they +# would broadcast their input arguments, so tests that work with inplace +# variants filter SampleInputs that broadcast their input. +# +# "name" is a string that's just used for debugging. It appears when printing +# the SampleInput. +# +# Sample inputs are designed to be used with many tests, some +# that are very time consuming, so they should be a small +# set with small tensors. An elaborated set of sample inputs +# can be specified using the "reference_inputs_func" attribute. +# The "reference inputs" for an operation are an extended +# set of sample inputs that can more exhaustively test an +# operator. They are used by only a few tests that are careful +# not to take too long to run. Adding reference inputs +# is highly encouraged! +# +# THE (OPTIONAL) ERROR INPUTS FUNCTION +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# OpInfos may optionally specify "error inputs" through an error function. If +# specified test_errors in test_ops.py will call the op with these inputs +# and validate that the desired error is thrown. +# +# Error inputs automate a common testing pattern where multiple inputs are +# passed to an operation and the errors they thrown are reviewed. Tests +# written in this style should be ported to the new OpInfo pattern. +# +# Error inputs are specified using the ErrorInputs class, which contains +# a SampleInput (see above) and data about the expected error. +# +# OPINFO FILE ORGANIZATION +# ~~~~~~~~~~~~~~~~~~~~~~~~ +# +# All OpInfos are currently defined in this file. Most OpInfo tests are defined +# in test_ops.py, but some system-specific tests are defined in those +# systems' test files, and subclass-specific tests are defined in the test +# file that corresponds to that subclass (see the below). +# Expect a reorganization in the future. +# +# WHAT'S TESTED? +# ~~~~~~~~~~~~~~ +# +# Every OpInfo in the op_db sequence has the following properties validated in +# test_ops.py: +# +# - that its supported dtypes are specified correctly +# - that the operation produces the same results when called with noncontiguous inputs +# - that it supports the out= argument properly (if it allows out=), +# see https://github.com/pytorch/pytorch/wiki/Developer-FAQ#how-does-out-work-in-pytorch +# - that it works with the conjugate view bit properly +# - that its function, method, and inplace variants perform the same operation +# (that is, that torch.add, torch.Tensor.add, and torch.Tensor.add_ all +# do the same thing). +# - that its inplace variant preserves the input's storage +# - that its gradient formula is implemented correctly, and that it supports +# gradgrad and complex grad and gradgrad and forward mode AD properly for +# the op's function and inplace variants (method variants are skipped +# to reduce test time). +# - that the operation performs the same operation when traced or scripted +# using the jit +# - that the operation is autodifferentiated by the jit as expected +# - that the operator's aliases, if any, perform the same operation and that +# the jit understands the alias +# - that the operator throws the correct errors (if error_inputs is defined) +# - that the operator produces the same results as a NumPy reference (if ref is defined) +# - that the operator produces the same results as a NumPy reference on an extended +# set of "reference inputs" (if both ref and reference_inputs_func are defined) +# (NOTE: elementwise unary and elementwise binary OpInfos do this even if only +# ref is defined, because they effectively autogenerate reference inputs) +# - that the operator works on different CUDA devices +# +# Additional OpInfo tests are in test_jit_fuser_te.py, test_fx_experimental.py, +# and test_fx.py. These tests validate that operators work with NNC and FX +# as expected. +# +# For performance, some of the above tests may only run on the first +# SampleInput returned by an OpInfo's sample input function. +# +# In addition to these tests, some subclasses (discussed in the next section) +# define additional tests. +# +# Critically, as mentioned above, what's not necessarily tested is that the operator +# works as expected. When implementing an OpInfo an engineer must still +# typically write one or more tests validating the operator's behavior. +# The exception to this is if reference testing is sufficient, or if +# the operation belongs to an OpInfo subclass that has more exhaustive +# operator testing. Elementwise unary and elementwise binary operators, +# in particular, usually don't require additional testing beyond +# writing an Opinfo. +# +# +# OPINFO (SUB)CLASSES +# ~~~~~~~~~~~~~~~~~~~ +# +# In addition to the OpInfo base class there are several specialized OpInfo +# subclasses. For example, the UnaryUfuncInfo subclass is used for +# unary elementwise operations. These operations have a common structure +# that test_unary_ufuncs.py exploits with additional automated testing. +# The automated testing in test_unary_ufuncs.py is so thorough, comparing +# the operator to a NumPy reference function on a plethora of values, that +# just implementing an OpInfo for a unary elementwise operation is often +# sufficient testing. +# +# The ForeachFuncInfo is another OpInfo subclass that is hyper-specialized to a +# very unique class of operations. These OpInfos aren't included in the +# op_db sequence and have their own tests. +# +# Other OpInfo subclasses, like SpectralFuncInfo, are just for convenience +# when writing OpInfos. +# +# TESTING A NEW OPERATOR +# ~~~~~~~~~~~~~~~~~~~~~~ +# +# If you're adding a new operator to any of the following namespaces: +# - torch +# - torch.fft +# - torch.linalg, +# - torch.special +# - torch.nn.functional +# then you should typically add an OpInfo for it. +# +# As mentioned a couple times above, implementing an OpInfo is not +# usually sufficient testing (unless the operator is a unary or binary elementwise +# operator). The OpInfo will only test the properties described in the +# "WHAT'S TESTED" section. It DOES NOT necessarily verify that the operator is +# implemented correctly. +# +# TIPS FOR WRITING AN OPINFO AND OPINFO TESTS +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# Writing an OpInfo can be a little daunting. Since the point of an OpInfo is to +# be consumed by a variety of systems it can be hard to understand how to +# deal with test failures or how to set the OpInfo metadata properly. +# +# Before adding an OpInfo it helps to look at other OpInfos. A sample inputs +# function must be defined, and the operator's dtypes must be specified. +# Once that's done you should run the operator's tests in test_ops.py +# (these can be filtered using the "-k" argument in pytest). Tests that +# fail should provide an error message that describes what to change about +# your OpInfo. You don't need to worry about changing an OpInfo's default +# values unless a test yells at you. +# +# Similarly, if you're writing a test that consumes OpInfos then it's critical +# your test provides a clear error message describing what to do when it +# fails. You should not assume the OpInfo implementer is familiar with your +# system. +# +# If you see a confusing error message while developing an OpInfo then please +# file an issue describing what happened. +# +# This trial-and-error approach to writing an OpInfo can be frustrating, +# but it's probably necessary as long as OpInfos don't require +# learning about all the systems that consume them. One thing that can help +# is the get_supported_dtypes() function defined in utils.py. This +# function can be used to programmatically specify the dtypes an operator +# supports, and is especially useful if writing an OpInfo on a machine +# without a CUDA device. See its documentation for more details. +# +# THE FUTURE OF OPINFOS AND OPINFO TESTING +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# In the future we expect OpInfo coverage to improve and cover +# the great majority of PyTorch's (public) operators. +# + + +# Classes and methods for the operator database +@dataclass +class OpInfo: + """Operator information and helper functions for acquiring it.""" + + # the string name of the function + name: str + + # An optional reference function that accepts ndarrays (AKA "NumPy arrays"). + # If given, the op will be compared with its reference on each of its sample inputs. + ref: Callable | None = None + + # the following metadata describes the operator, its variants, and its aliases, if any + + # iterable of aliases, e.g. ("absolute",) for torch.abs + aliases: Iterable = None + + # additional string to include in the test name + # this is useful when an op needs multiple OpInfos, + # like divide does, often because it's really several + # different ops behind the scenes + variant_test_name: str = "" + + # the function variant of the operation, populated as torch. if None + op: Callable = None + + # allows the method variant of this operation to be specified as follows: + # - if _NOTHING (default), then the OpInfo attempts to discover the variant using its name + # - if None, then the OpInfo explicitly specifies is has no associated method + # - if a Callable, then that callable should be the method associated with this operation + method_variant: Callable = _NOTHING + + # allows the inplace variant of this operation to be specified as follows: + # - if _NOTHING (default), then the OpInfo attempts to discover the variant using its name + # - if None, then the OpInfo explicitly specifies is has no associated inplace variant + # - if a Callable, then that callable should be the inplace variant associated with this operation + inplace_variant: Callable = _NOTHING + + # allows the operator variant of this operation to be specified as follows: + # - if _NOTHING (default), then the OpInfo attempts to discover the variant using its name + # - if None, then the OpInfo explicitly specifies is has no associated operator + # - if a Callable, then that callable should be the operator associated with this operation + operator_variant: Callable = _NOTHING + + # allows the inplace operator variant of this operation to be specified as follows: + # - if _NOTHING (default), then the OpInfo attempts to discover the variant using its name + # - if None, then the OpInfo explicitly specifies is has no associated inplace operator + # - if a Callable, then that callable should be the inplace operator associated with this operation + inplace_operator_variant: Callable = _NOTHING + + # the following metadata are test directives for skipping or modifying tests + + # information about which tests to skip + skips: tuple = () + + # decorators to apply to generated tests + decorators: tuple = () + + # the following are pointers to functions to generate certain classes of inputs + + # function to generate sample inputs with strided layouts + sample_inputs_func: Callable = None + + # function to generate a more thorough set of samples inputs with strided layouts + reference_inputs_func: Callable = None + + # function to generate inputs that will throw errors + error_inputs_func: Callable = None + + # function to generate sparse (coo, csr, csc, bsr, bsc) inputs that will throw errors + error_inputs_sparse_func: Callable = None + + # function to generate sample inputs with sparse coo layouts + sample_inputs_sparse_coo_func: Callable = None + + # function to generate sample inputs with sparse csr layouts + sample_inputs_sparse_csr_func: Callable = None + + # function to generate sample inputs with sparse csc layouts + sample_inputs_sparse_csc_func: Callable = None + + # function to generate sample inputs with sparse bsr layouts + sample_inputs_sparse_bsr_func: Callable = None + + # function to generate sample inputs with sparse bsc layouts + sample_inputs_sparse_bsc_func: Callable = None + + # the following metadata relates to dtype support and is tested for correctness in test_ops.py + + # dtypes this function works with on the CPU, + # inherited by other device types that don't specify their own dtypes + dtypes: _dispatch_dtypes = None + + # the following dtypesIf... options override the dtypes value on their respective device types + # I.e. instead of writing multiple `dtypesIfCUDA`, `dtypesIfROCM`, etc one can simply define a dict + # dtypesIf = { 'cuda': (torch.float, torch.double), 'rocm': (torch.half, torch.bfloat16) } + dtypesIf: dict[str, _dispatch_dtypes] = field(default_factory=dict) + + def __getattribute__(self, name: str) -> Any: + if name.startswith("dtypesIf") and name != "dtypesIf": + # TODO: Warn if used + dev_name = name.removeprefix("dtypesIf").lower() + return self.dtypesIf.get(dev_name) + return super().__getattribute__(name) + + def __setattr__(self, name: str, value: Any) -> None: + # TODO: After migration, start adding warnings here + if name.startswith("dtypesIf") and name != "dtypesIf": + if not isinstance(value, (_dispatch_dtypes, type(None))): + raise AssertionError( + f"Expected _dispatch_dtypes or None, got {type(value)}" + ) + dev_name = name.removeprefix("dtypesIf").lower() + self.dtypesIf[dev_name] = value + return + super().__setattr__(name, value) + + # dtypes this function is expected to work with on CUDA + dtypesIfCUDA: _dispatch_dtypes = None + + # dtypes this function is expected to work with on ROCM + dtypesIfROCM: _dispatch_dtypes = None + + # dtypes this function is expected to work with on MPS + dtypesIfMPS: _dispatch_dtypes = None + + dtypesIfHpu: _dispatch_dtypes = None + + # dtypes this function is expected to work with on XPU + dtypesIfXPU: _dispatch_dtypes = None + + # backward dtypes this function is expected to work with + backward_dtypes: _dispatch_dtypes = None + + # backward dtypes this function is expected to work with on CUDA + backward_dtypesIfCUDA: _dispatch_dtypes = None + + # backward dtypes this function is expected to work with on ROCM + backward_dtypesIfROCM: _dispatch_dtypes = None + + # backward dtypes this function is expected to work with on MPS + backward_dtypesIfMPS: _dispatch_dtypes = None + + backward_dtypesIfHpu: _dispatch_dtypes = None + + # the following metadata describes the operators out= support + + # whether the op supports the out kwarg + # defaults to True, if the op does not allow the out kwarg or + # supports it incorrectly then test_out in test_ops.py should fail + supports_out: bool = True + + # the following metadata relates to autograd support + # whether the operation supports backward mode AD + # if true, gradient correctness is tested in test_ops.py + # using the op's sample inputs + supports_autograd: bool = True + + # whether the op supports second order gradients + # if true, gradgrad correctness is tested in test_ops.py + # defaults to support_autograd's value + # TODO: rename this to supports_bwgrad_bwgrad to be consistent with below + supports_gradgrad: bool = None + + # whether the ops supports second order gradients via + # forward-over-reverse. If True, forward-over-reverse gradgrad correctness + # is tested. If False, test that forward grad is not implemented. + # Defaults to False. + supports_fwgrad_bwgrad: bool = False + + # whether the operation supports inplace autograd + # if true, tested in test_ops.py + # defaults to supports_autograd's value + supports_inplace_autograd: bool = None + + # Whether the operation support forward mode AD + # If the value is True, we check that the gradients are correct + # If the value is False, we test that forward grad is not implemented + supports_forward_ad: bool = False + + # Whether the operation has a varargs variant + # (e.g. functions like ones, zeros, methods like view, permute) + supports_varargs: bool = False + + # Whether the forward operation avoids materializing COW tensor inputs + supports_cow_input_no_materialize_forward: bool = True + + # Whether the backward operation avoids materializing COW tensor inputs + supports_cow_input_no_materialize_backward: bool = True + + # Whether to skip the backward part of the COW tensor input test + skip_cow_input_backward: bool = False + + # If `supports_cow_input_no_materialize_forward == True`, this list contains + # the arg indices or kwarg names of inputs that are expected to materialize + allow_cow_input_materialize_forward: list[int | str] = None + + # If `supports_cow_input_no_materialize_backward == True`, this list contains + # the arg indices or kwarg names of inputs that are expected to materialize + allow_cow_input_materialize_backward: list[int | str] = None + + # wrapper function for gradcheck + gradcheck_wrapper: Callable = lambda op, *args, **kwargs: op(*args, **kwargs) + + # whether to check batched grad when doing gradcheck + # defaults to support_autograd's value + check_batched_grad: bool = None + + # whether to check batched grad grad when doing gradgradcheck + # default's to support_gradgrad's value + check_batched_gradgrad: bool = None + + # whether to check batched forward grad when doing gradcheck + # defaults to the value of `supports_forward_ad` + check_batched_forward_grad: bool = None + + # whether to check batched forward grad when doing gradcheck + # defaults to the value of `check_batched_forward_grad` + check_inplace_batched_forward_grad: bool = None + + # tolerance for nondeterminism while performing gradcheck + gradcheck_nondet_tol: float = 0.0 + + # Whether to use the fast implementation for gradcheck/gradgradcheck. + # When set to None, defers to the default value provided by the wrapper + # function around gradcheck (testing._internal.common_utils.gradcheck) + gradcheck_fast_mode: bool = None + + # the following metadata relates to JIT support and is tested for correctness in test_ops.py + + # name of the corresponding aten:: operator + aten_name: str = None + + # if this is a composite implicit autograd op, the decomposed op + decomp_aten_name: str | None = None + + # name of the corresponding aten:: operator for backwards + aten_backward_name: str | None = None + + # if a op's aten::node is expected to be symbolically autodiffed + assert_autodiffed: bool = False + + # a list of strings with node names that are expected to be in a + # DifferentiableGraph when autodiffed. Ex: ['aten::add', 'aten::mm'], + # default is populated to be ['aten::(name of Python operator)'] + autodiff_nonfusible_nodes: list[str] = None + + # a list of strings with node names that are expected to be in FusionGroups + # inside of DifferentiableGraphs when this operation is autodiffed. + # Ex: ['aten::add', 'aten::mm'], defaults to an empty list + # Note: currently no ops use fusible nodes + autodiff_fusible_nodes: list[str] = None + + # the following metadata relates to sparse support and is used in test_sparse.py + + # whether the op supports sparse coo inputs, defaults to False + # TODO: rename supports_sparse to supports_sparse_coo + supports_sparse: bool = None + + # only run tracing tests + supports_scripting: bool = True + + # if the operator can be traced + supports_tracing: bool = True + + # the following metadata relates to sparse compressed support and + # is used in test_sparse_csr.py and test_sparse.py + + # whether the op supports sparse csr inputs, defaults to False + supports_sparse_csr: bool = None + # whether the op supports sparse csc inputs, defaults to False + supports_sparse_csc: bool = None + # whether the op supports sparse bsr inputs, defaults to False + supports_sparse_bsr: bool = None + # whether the op supports sparse bsc inputs, defaults to False + supports_sparse_bsc: bool = None + # whether the op supports nested jagged inputs, defaults to False + supports_njt: bool = None + + # whether the op promotes integer inputs to float + promotes_int_to_float: bool = False + + # the following metadata relates to complex support and is checked in test_ops.py + + test_conjugated_samples: bool = True + + test_neg_view: bool = True + + # assert that jit shape analysis fully propagates shape + assert_jit_shape_analysis: bool = False + + # the following metadata relates to ExpandedWeights support and is checked in test_expanded_weights.py + + supports_expanded_weight: bool = False + + is_factory_function: bool = False + + skip_correctness_check_compile_vs_eager: bool = False + + def __post_init__(self): + self._original_opinfo_args = asdict(self).copy() + + if self.dtypes is None: + raise AssertionError(f"OpInfo for {self.name} has no dtypes!") + + # Validates the dtypes are generated from the dispatch-related functions + for name, val in self.dtypesIf.items(): + if val is not None: + if not isinstance(val, _dispatch_dtypes): + raise AssertionError(f"Expected _dispatch_dtypes, got {type(val)}") + self.dtypesIf[name] = set(val) + + if self.aten_name is None: + self.aten_name = self.name + + # Attribute to verify dynamic_dtypes are used. + self.dynamic_dtypes = any( + isinstance(dtypes, utils._dynamic_dispatch_dtypes) + for dtypes in self.dtypesIf.values() + ) + + if self.dynamic_dtypes: + # Make sure `dtyesIfCUDA` is dynamic, if dynamic dispatch is used for CPU + # This is because, below we set dtypesIfCUDA to dtypes if they are None. + if not isinstance(self.dtypesIfCUDA, utils._dynamic_dispatch_dtypes): + raise AssertionError( + f"To use dynamic dtypes for operator {self.name}, " + "acquire the dtypes dynamically for argument `dtypesIfCUDA`. " + "This is to ensure that CUDA dtypes are acquired correctly as they " + "differ from CPU dtypes occasionally" + ) + if not isinstance(self.dtypesIfMPS, utils._dynamic_dispatch_dtypes): + raise AssertionError( + f"To use dynamic dtypes for operator {self.name}, " + "acquire the dtypes dynamically for argument `dtypesIfMPS`. " + "This is to ensure that MPS dtypes are acquired correctly as they " + "differ from CPU dtypes occasionally" + ) + + self.dtypes = set(self.dtypes) + + # NOTE: backward dtypes must be acquired before forward dtypes + # since they fallback to explicit (not implicit!) specifications of + # forward dtypes + self.backward_dtypesIfROCM = ( + set(self.backward_dtypesIfROCM) + if self.backward_dtypesIfROCM is not None + else ( + self.backward_dtypesIfCUDA + if self.backward_dtypesIfCUDA is not None + else self.backward_dtypes + if self.backward_dtypes is not None + else self.dtypesIfROCM + if self.dtypesIfROCM is not None + else self.dtypesIfCUDA + if self.dtypesIfCUDA is not None + else self.dtypes + ) + ) + self.backward_dtypesIfCUDA = ( + set(self.backward_dtypesIfCUDA) + if self.backward_dtypesIfCUDA is not None + else ( + self.backward_dtypes + if self.backward_dtypes is not None + else self.dtypesIfCUDA + if self.dtypesIfCUDA is not None + else self.dtypes + ) + ) + self.backward_dtypesIfMPS = ( + set(self.backward_dtypesIfMPS) - {torch.float64, torch.cdouble} + if self.backward_dtypesIfMPS is not None + else ( + set(self.backward_dtypes) - {torch.float64, torch.cdouble} + if self.backward_dtypes is not None + else set(self.dtypesIfMPS) - {torch.float64, torch.cdouble} + if self.dtypesIfMPS is not None + else set(self.dtypes) - {torch.float64, torch.cdouble} + ) + ) + self.backward_dtypesIfHpu = ( + set(self.backward_dtypesIfHpu) + if self.backward_dtypesIfHpu is not None + else ( + self.backward_dtypes + if self.backward_dtypes is not None + else self.dtypes + ) + ) + + self.backward_dtypes = ( + set(self.backward_dtypes) + if self.backward_dtypes is not None + else self.dtypes + ) + + # Inherit from cpu + for dev_type in ["cuda", "hpu"]: + if self.dtypesIf.get(dev_type) is None: + self.dtypesIf[dev_type] = self.dtypes + + # Inherit from CUDA + for dev_type in ["rocm", "xpu"]: + if self.dtypesIf.get(dev_type) is None: + self.dtypesIf[dev_type] = self.dtypesIf["cuda"] + + # Inherit from cpu + for dev_type in ["mps"]: + if self.dtypesIf.get(dev_type) is None: + # Double floats are not supported on MPS + self.dtypesIf[dev_type] = self.dtypes - {torch.float64, torch.cdouble} + else: + self.dtypesIf[dev_type] = self.dtypesIf[dev_type] - { + torch.float64, + torch.cdouble, + } + + # NOTE: if the op is unspecified it is assumed to be under the torch namespace + if not self.op: + self.op = _getattr_qual(torch, self.name) + + if self.method_variant is _NOTHING: + self.method_variant = getattr(torch.Tensor, self.name, None) + + # attributes like real, imag are not callable + if not callable(self.method_variant): + self.method_variant = None + + if self.inplace_variant is _NOTHING: + inplace_name = self.name + "_" + self.inplace_variant = getattr(torch.Tensor, inplace_name, None) + + if self.operator_variant is _NOTHING: + self.operator_variant = getattr(operator, self.name, None) + + if self.inplace_operator_variant is _NOTHING: + # Note: operator.i will use operator. and assign the result to the lhs when no + # __i__ method is found. This results in the appearance of an inplace operator variant which + # does not have the correct inplace behavior. To avoid this, we guard automatic detection of the inplace + # operator with a check that an inplace variant exists. + if self.inplace_variant is not None: + inplace_operator_name = "i" + self.name + self.inplace_operator_variant = getattr( + operator, inplace_operator_name, None + ) + else: + self.inplace_operator_variant = None + + self.decorators = (*self.decorators, *self.skips) + + # Specifying sample inputs function without specifying the + # corresponding layout support implies the layout support: + if self.supports_sparse is None: + self.supports_sparse = self.sample_inputs_sparse_coo_func is not None + if self.sample_inputs_sparse_coo_func is None: + self.sample_inputs_sparse_coo_func = self._sample_inputs_unspecified + + if self.supports_sparse_csr is None: + self.supports_sparse_csr = self.sample_inputs_sparse_csr_func is not None + if self.sample_inputs_sparse_csr_func is None: + self.sample_inputs_sparse_csr_func = self._sample_inputs_unspecified + + if self.supports_sparse_csc is None: + self.supports_sparse_csc = self.sample_inputs_sparse_csc_func is not None + if self.sample_inputs_sparse_csc_func is None: + self.sample_inputs_sparse_csc_func = self._sample_inputs_unspecified + + if self.supports_sparse_bsr is None: + self.supports_sparse_bsr = self.sample_inputs_sparse_bsr_func is not None + if self.sample_inputs_sparse_bsr_func is None: + self.sample_inputs_sparse_bsr_func = self._sample_inputs_unspecified + + if self.supports_sparse_bsc is None: + self.supports_sparse_bsc = self.sample_inputs_sparse_bsc_func is not None + if self.sample_inputs_sparse_bsc_func is None: + self.sample_inputs_sparse_bsc_func = self._sample_inputs_unspecified + + if self.supports_njt is None: + self.supports_njt = False + + # We run the sampling functions without tracking the gradiends of the creation of inputs + self.sample_inputs_func = torch.no_grad()(self.sample_inputs_func) + self.sample_inputs_sparse_coo_func = torch.no_grad()( + self.sample_inputs_sparse_coo_func + ) + self.sample_inputs_sparse_csr_func = torch.no_grad()( + self.sample_inputs_sparse_csr_func + ) + self.sample_inputs_sparse_csc_func = torch.no_grad()( + self.sample_inputs_sparse_csc_func + ) + self.sample_inputs_sparse_bsr_func = torch.no_grad()( + self.sample_inputs_sparse_bsr_func + ) + self.sample_inputs_sparse_bsc_func = torch.no_grad()( + self.sample_inputs_sparse_bsc_func + ) + if self.reference_inputs_func is not None: + self.reference_inputs_func = torch.no_grad()(self.reference_inputs_func) + + if not self.autodiff_fusible_nodes: + self.autodiff_fusible_nodes = [] + + if self.autodiff_nonfusible_nodes is None: + self.autodiff_nonfusible_nodes = ["aten::" + self.name] + + # Autograd support + + # Autograd flags that depend on backward AD only + # - If setting has been explicitly set, raise error if inconsistent + if self.supports_gradgrad is None: + self.supports_gradgrad = self.supports_autograd + else: + if self.supports_gradgrad and not self.supports_autograd: + raise AssertionError( + "supports_gradgrad refines the part of autograd is supported, so it should " + "not be set if supports_autograd is False" + ) + if self.check_batched_grad is None: + self.check_batched_grad = self.supports_autograd or self.supports_forward_ad + else: + if self.check_batched_grad and not ( + self.supports_autograd or self.supports_forward_ad + ): + raise AssertionError( + "check_batched_grad refines the part of autograd that will be checked (by gradcheck), so " + "it should not be set if supports_autograd is False" + ) + if self.check_batched_gradgrad is None: + self.check_batched_gradgrad = self.supports_gradgrad + else: + if self.check_batched_gradgrad and not self.supports_gradgrad: + raise AssertionError( + "check_batched_gradgrad refines the part of autograd that will be checked (by " + "gradgradcheck), so it should not be set if either supports_gradgrad or supports_autograd " + "is False." + ) + if self.check_batched_forward_grad is None: + self.check_batched_forward_grad = self.supports_forward_ad + else: + if self.check_batched_forward_grad and not self.supports_forward_ad: + raise AssertionError( + "check_batched_forward_grad should only be used when supports_forward_ad " + "is True. It is used to disable the test in the specific cases " + "where the op supports forward ad but fails to compute " + "batched forward grad." + ) + + if self.check_inplace_batched_forward_grad is None: + self.check_inplace_batched_forward_grad = self.check_batched_forward_grad + else: + if ( + self.check_inplace_batched_forward_grad + and not self.check_batched_forward_grad + ): + raise AssertionError( + "check_batched_forward_grad should only be used when check_batched_forward_grad " + "is True. It is used to disable the test in the specific cases " + "where the op supports batched forward grad but fails to compute batched forward " + "grad for the inplace variant of the op." + ) + + if self.supports_fwgrad_bwgrad and not self.supports_autograd: + raise AssertionError( + "supports_fwgrad_bwgrad enables forward-over-backward gradgrad checks and should only be " + f"True if backward ad is also checked, i.e., supports_forward_ad should be True. ({self.name})" + ) + + # Autograd flags that depend on both forward AD and backward AD + if self.supports_inplace_autograd is None: + self.supports_inplace_autograd = ( + self.supports_autograd or self.supports_forward_ad + ) + else: + if ( + self.supports_inplace_autograd + and not self.supports_autograd + and not self.supports_forward_ad + ): + raise AssertionError( + "supports_inplace_autograd refines the part of autograd that is supported, so " + "it should not be set if both supports_autograd and supports_forward_ad are False" + ) + + if self.aliases is not None: + self.aliases = tuple(AliasInfo(a) for a in self.aliases) # type: ignore[assignment] + else: + self.aliases = () + + def __call__(self, *args, **kwargs): + """Calls the function variant of the operator.""" + return self.op(*args, **kwargs) + + def __str__(self): + return dataclass_repr(self) + + def get_op(self): + """Returns the function variant of the operator, torch..""" + return self.op + + def get_method(self): + """Returns the method variant of the operator, torch.Tensor.. + Returns None if the operator has no method variant. + """ + return self.method_variant + + def get_inplace(self): + """Returns the inplace variant of the operator, torch.Tensor._. + Returns None if the operator has no inplace variant. + """ + return self.inplace_variant + + def get_operator(self): + """Returns operator variant of the operator, e.g. operator.neg + Returns None if the operator has no operator variant. + """ + return self.operator_variant + + def get_inplace_operator(self): + """Returns the inplace operator variant of the operator, e.g operator.iadd + Returns None if the operator has no inplace operator variant""" + return self.inplace_operator_variant + + # Returns a tuple of callables: + # (TestCase -> subtest context, TestCase -> skip / xfail context) + # I'd love to combine these into one but I haven't figured out how to do it + # in a way that works like it should, and I tried a LOT of things. + def _maybe_skip_or_xfail(self, rules, device, sample, idx): + def _subtest_fn(test_case, sample=sample.name, idx=idx): + return test_case.subTest(sample=sample, idx=idx) + + if rules is None or len(rules) == 0: + return (_subtest_fn, lambda _: contextlib.nullcontext()) + + # NB: match first rule only (order matters!) + for rule in rules: + if rule.sample_match_fn(device, sample): + log.debug( + "matched %s rule '%s': %s %s %s", + rule.type, + rule.name, + self.full_name, + device, + sample, + ) + + # Provide a context for the test case to run the sample input + # through as a subtest AND handle skip / xfail for it as needed. + return ( + _subtest_fn, + lambda test_case, rule=rule: rule.get_context(test_case), + ) + + log.debug("matched no rules: %s %s %s", self.full_name, device, sample) + return (_subtest_fn, lambda _: contextlib.nullcontext()) + + def _sample_callback_fn(self, use_subtests, device): + # Get sample-specific skips / xfails. + sample_skips_and_xfails = getattr( + extract_test_fn(), "sample_skips_and_xfails", None + ) + + if sample_skips_and_xfails is not None and not use_subtests: + raise RuntimeError( + """Sample-specific skips / xfails require use_subtests=True. +Please pass this to the sample generation function and run the test logic within the +returned contexts (NB: order matters!). For example: + +def test_foo(self, device, dtype, op): + for sample, subtest_ctx, skip_xfail_ctx in op.sample_inputs(..., use_subtests=True): + # these contexts handle running within subtests and skips / xfails + with subtest_ctx(self), skip_xfail_ctx(self): + # test logic here + ...""" + ) + + if not use_subtests: + # use the default callback that returns the sample without a subtest context + return None + + if USE_PYTEST: + try: + import pytest_subtests # noqa: F401 + except ModuleNotFoundError: + raise RuntimeError( + "Encountered an OpInfo test with use_subtests=True and pytest-subtests is " + "not installed. The feature will not work correctly within pytest without " + "this package; please install it." + ) from None + + def _f( + sample, + idx, + self=self, + device=device, + sample_skips_and_xfails=sample_skips_and_xfails, + use_subtests=use_subtests, + ): + # When subtests are enabled, also return a subtest context. This is required + # for xfails / skips to work properly. + return ( + sample, + *self._maybe_skip_or_xfail( + sample_skips_and_xfails, device, sample, idx + ), + ) + + return _f + + def conjugate_sample_inputs(self, device, dtype, requires_grad=False, **kwargs): + """Returns an iterable of SampleInputs but with the tensor input or first + tensor in a sequence input conjugated. + """ + + set_seed = kwargs.pop("set_seed", True) + use_subtests = kwargs.pop("use_subtests", False) + samples = self.sample_inputs_func(self, device, dtype, requires_grad, **kwargs) + conj_samples = list(samples) + + def conjugate(tensor): + _requires_grad = tensor.requires_grad + tensor = tensor.conj() + return tensor.requires_grad_(_requires_grad) + + for i, sample in enumerate(samples): + sample = conj_samples[i] + # Note: it is assumed that the input here is either a tensor or tensorlist + if isinstance(sample.input, torch.Tensor): + sample.input = conjugate(sample.input) + else: + sample.input[0] = conjugate(sample.input[0]) + + return TrackedInputIter( + iter(conj_samples), + "conjugate sample input", + item_callback=self._sample_callback_fn(use_subtests, device), + set_seed=set_seed, + restrict_to_index=OPINFO_SAMPLE_INPUT_INDEX, + ) + + def sample_inputs(self, device, dtype, requires_grad=False, **kwargs): + """ + Returns an iterable of SampleInputs. + + These samples should be sufficient to test the function works correctly + with autograd, TorchScript, etc. + """ + set_seed = kwargs.pop("set_seed", True) + use_subtests = kwargs.pop("use_subtests", False) + samples = self.sample_inputs_func(self, device, dtype, requires_grad, **kwargs) + + if kwargs.get("include_conjugated_inputs", False): + conj_samples = self.conjugate_sample_inputs( + device, dtype, requires_grad, **kwargs + ) + samples_list = list(samples) + samples_list.extend(conj_samples) + samples = tuple(samples_list) + + return TrackedInputIter( + iter(samples), + "sample input", + item_callback=self._sample_callback_fn(use_subtests, device), + set_seed=set_seed, + restrict_to_index=OPINFO_SAMPLE_INPUT_INDEX, + ) + + def reference_inputs(self, device, dtype, requires_grad=False, **kwargs): + """ + Returns an iterable of SampleInputs. + + Distinct from sample_inputs() above because this returns an expanded set + of inputs when reference_inputs_func is defined. If undefined this returns + the sample inputs. + """ + set_seed = kwargs.pop("set_seed", True) + use_subtests = kwargs.pop("use_subtests", False) + if self.reference_inputs_func is None: + samples = self.sample_inputs_func( + self, device, dtype, requires_grad, **kwargs + ) + return TrackedInputIter( + iter(samples), + "reference input", + item_callback=self._sample_callback_fn(use_subtests, device), + set_seed=set_seed, + restrict_to_index=OPINFO_SAMPLE_INPUT_INDEX, + ) + + if kwargs.get("include_conjugated_inputs", False): + raise NotImplementedError + + references = self.reference_inputs_func( + self, device, dtype, requires_grad, **kwargs + ) + return TrackedInputIter( + iter(references), + "reference input", + item_callback=self._sample_callback_fn(use_subtests, device), + set_seed=set_seed, + restrict_to_index=OPINFO_SAMPLE_INPUT_INDEX, + ) + + def error_inputs(self, device, **kwargs): + """ + Returns an iterable of ErrorInputs. + """ + set_seed = kwargs.pop("set_seed", True) + use_subtests = kwargs.pop("use_subtests", False) + errs = self.error_inputs_func(self, device, **kwargs) + + def _error_item_callback(e, i, use_subtests=use_subtests, device=device): + cb = self._sample_callback_fn(use_subtests, device) + # no rules to apply; just return the sample + if cb is None: + return e + + # adapt the callback call since ErrorInputs contain SampleInputs + _, subtest_ctx = cb(e.sample_input, i) + return (e, subtest_ctx) + + return TrackedInputIter( + iter(errs), + "error input", + track_callback=lambda e: e.sample_input, + item_callback=_error_item_callback, + set_seed=set_seed, + restrict_to_index=OPINFO_SAMPLE_INPUT_INDEX, + ) + + def error_inputs_sparse(self, device, layout, **kwargs): + """ + Returns an iterable of ErrorInputs that contain sparse sample + inputs with a specified layout. + """ + if not self.supports_sparse_layout(layout): + raise unittest.SkipTest("unsupported sparse layout") + return self.error_inputs_sparse_func(self, device, layout, **kwargs) + + def supports_sparse_layout(self, layout): + """Return True if OpInfo supports the specified sparse layout.""" + layout_name = str(layout).split(".")[-1] + # map torch.sparse_coo to OpInfo.supports_sparse: + layout_name = layout_name.replace("_coo", "") + return getattr(self, f"supports_{layout_name}") + + def sample_inputs_sparse( + self, layout, device, dtype, requires_grad=False, **kwargs + ): + """Returns an iterable of SampleInputs that contain inputs with a + specified sparse layout. + """ + layout_name = str(layout).split(".")[-1] + sample_inputs_mth = getattr(self, "sample_inputs_" + layout_name) + + def non_empty_sampler(op, generator): + found_sample = False + for sample in generator: + found_sample = True + yield sample + if not found_sample: + raise unittest.SkipTest("NO SAMPLES!") + + return non_empty_sampler( + self, + sample_inputs_mth(device, dtype, requires_grad=requires_grad, **kwargs), + ) + + def _sample_inputs_unspecified(self, *args, **kwargs): + """Raises an NotImplemented exception in a OpInfo instance creation + that specifies supports_sparse(|_csr|_csc|_bsr|_bsc)=True + without specifying the corresponding sample function as + sample_inputs_sparse_(coo|csr|csc|bsr|bsc)_func. + + To avoid this, either define the corresponding sample function, + or re-map unsupported samples to error inputs in an appropriate + + opinfo/definitions/sparse.py:_validate_sample_input_sparse_ + + function. + """ + raise NotImplementedError("no sample function specified") + + def sample_inputs_sparse_coo(self, device, dtype, requires_grad=False, **kwargs): + """Returns an iterable of SampleInputs that contain inputs with sparse + coo layout. + """ + return self.sample_inputs_sparse_coo_func( + self, device, dtype, requires_grad, **kwargs + ) + + def sample_inputs_sparse_csr(self, device, dtype, requires_grad=False, **kwargs): + """Returns an iterable of SampleInputs that contain inputs with sparse + csr layout. + """ + return self.sample_inputs_sparse_csr_func( + self, device, dtype, requires_grad, **kwargs + ) + + def sample_inputs_sparse_csc(self, device, dtype, requires_grad=False, **kwargs): + """Returns an iterable of SampleInputs that contain inputs with sparse + csc layout. + """ + return self.sample_inputs_sparse_csc_func( + self, device, dtype, requires_grad, **kwargs + ) + + def sample_inputs_sparse_bsr(self, device, dtype, requires_grad=False, **kwargs): + """Returns an iterable of SampleInputs that contain inputs with sparse + bsr layout. + """ + return self.sample_inputs_sparse_bsr_func( + self, device, dtype, requires_grad, **kwargs + ) + + def sample_inputs_sparse_bsc(self, device, dtype, requires_grad=False, **kwargs): + """Returns an iterable of SampleInputs that contain inputs with sparse + bsc layout. + """ + return self.sample_inputs_sparse_bsc_func( + self, device, dtype, requires_grad, **kwargs + ) + + def get_decorators(self, test_class, test_name, device, dtype, param_kwargs): + """Returns the decorators targeting the given test.""" + result = [] + for decorator in self.decorators: + if isinstance(decorator, DecorateInfo): + if decorator.is_active( + test_class, test_name, device, dtype, param_kwargs + ): + result.extend(decorator.decorators) + else: + result.append(decorator) + return result + + def supported_dtypes(self, device_type): + if device_type == "privateuse1": + device_type = torch._C._get_privateuse1_backend_name() + device_type = torch.device(device_type).type + if device_type == "cuda" and TEST_WITH_ROCM: + device_type = "rocm" + result = self.dtypesIf.get(device_type, self.dtypes) + return result + + def supported_backward_dtypes(self, device_type): + if not self.supports_autograd: + return set() + + if device_type == "privateuse1": + device_type = torch._C._get_privateuse1_backend_name() + device_type = torch.device(device_type).type + backward_dtypes = None + if device_type == "cuda": + backward_dtypes = ( + self.backward_dtypesIfROCM + if TEST_WITH_ROCM + else self.backward_dtypesIfCUDA + ) + elif device_type == "hpu": + backward_dtypes = self.backward_dtypesIfHpu + elif device_type == "mps": + backward_dtypes = self.backward_dtypesIfMPS + else: + backward_dtypes = self.backward_dtypes + + allowed_backward_dtypes = floating_and_complex_types_and( + torch.bfloat16, torch.float16, torch.complex32 + ) + return set(allowed_backward_dtypes).intersection(backward_dtypes) + + def supports_dtype(self, dtype, device_type) -> bool: + return dtype in self.supported_dtypes(device_type) + + @property + def full_name(self): + """Returns a full name that helps to uniquely identify this OpInfo.""" + variant = "." + self.variant_test_name if self.variant_test_name else "" + # example: "normal.in_place" where "normal" is the name and "in_place" is the variant + return f"{self.name}{variant}" + + @property + def formatted_name(self): + """Returns a formatted full name for this OpInfo that can be used in test names.""" + return self.full_name.replace(".", "_") + + +# Represents a skip / xfail rule matching a particular set of tests. It allows granularity +# at the device, dtype, op, and individual sample levels. This flexibility allows entire +# bugs to be represented by a single rule, even if this corresponds with multiple conceptual +# test cases across multiple ops. +@dataclass +class SampleRule(ABC): + # function to indicate whether the rule applies to this op; return True if so + # NB: str arg of callable is device_type + op_match_fn: Callable[[str, OpInfo], bool] = None + # function to indicate whether the rule applies to this sample; return True if so + sample_match_fn: Callable[[torch.device, SampleInput], bool] = None + # optional name for identifying the rule + name: str = "" + + def __post_init__(self): + if self.op_match_fn is None: + raise ValueError("must have op_match_fn set to be useful") + if self.sample_match_fn is None: + # by default, match for all samples + self.sample_match_fn = lambda device, sample: True + + # returns a string identifier of the rule type + @abstractmethod + def type(self) -> str: ... + + # returns an appropriate context that handles the xfail, skips, etc. + @abstractmethod + def get_context(self, test_case): ... + + +# useful for specifying xfails +@dataclass +class XFailRule(SampleRule): + # expected error type + error_type: TypeVar = Exception + # expected error message + error_msg: str = ".*" + + @property + def type(self) -> str: + return "xfail" + + def get_context(self, test_case): + return test_case.assertRaisesRegex( + # failing within torch.compile wraps within a BackendCompilerFailed + (self.error_type, torch._dynamo.exc.BackendCompilerFailed), + self.error_msg, + ) + + +# useful for specifying skips +@dataclass +class SkipRule(SampleRule): + @property + def type(self): + return "skip" + + def get_context(self, test_case): + @contextlib.contextmanager + def skipcontext(test_case=test_case): + test_case.skipTest("Skipped!") + yield + + return skipcontext() + + +# Decorator that defines skip / xfail rules for a given test function. If these are +# present, the @ops decorator will apply these for each op and place them onto the +# parametrized test functions for use by e.g. OpInfo.sample_inputs(). +class sample_skips_and_xfails: + def __init__(self, rules): + self.rules = rules + + def __call__(self, fn): + rules = getattr(fn, "sample_skips_and_xfails", None) + if rules is not None: + raise RuntimeError("Multiple sets of sample_skips_and_xfails defined") + + fn.sample_skips_and_xfails = self.rules + return fn + + +def _generate_reduction_inputs(device, dtype, requires_grad, **kwargs): + """Generates input tensors for testing reduction operators""" + yield make_tensor([], dtype=dtype, device=device, requires_grad=requires_grad) + yield make_tensor([2], dtype=dtype, device=device, requires_grad=requires_grad) + yield make_tensor([3, 5], dtype=dtype, device=device, requires_grad=requires_grad) + yield make_tensor( + [3, 2, 1, 2], dtype=dtype, device=device, requires_grad=requires_grad + ) + + +def _generate_reduction_kwargs(ndim, supports_multiple_dims=True): + """Generates a subset of all valid dim and keepdim kwargs given ndim that + is appropriate for testing reduction operators. + """ + + # Test default dim and keepdim + yield {} + + # Test reducing inner and outer most dimensions + yield {"dim": 0, "keepdim": True} + yield {"dim": -1, "keepdim": False} + + # Test reducing middle dimension + if ndim > 2: + yield {"dim": ndim // 2, "keepdim": True} + + if supports_multiple_dims: + # Test reducing all dimensions + yield {"dim": tuple(range(ndim)), "keepdim": False} + + # Test reducing both first and last dimensions + if ndim > 1: + yield {"dim": (0, -1), "keepdim": True} + + # Test reducing every other dimension starting with the second + if ndim > 3: + yield {"dim": tuple(range(1, ndim, 2)), "keepdim": False} + + +def sample_inputs_reduction(op_info, device, dtype, requires_grad, **kwargs): + """Sample inputs for reduction operators.""" + + # TODO(@heitorschueroff) Once all reduction operators are using + # ReductionOpInfo use op_info.supports_multiple_dims directly. + supports_multiple_dims: bool = kwargs.get("supports_multiple_dims", True) + + # TODO(@heitorschueroff) Once all reduction operators are using ReductionOpInfo + # use op_info.generate_args_kwargs directly. + generate_args_kwargs = kwargs.get( + "generate_args_kwargs", lambda *args, **kwargs: (yield (), {}) + ) + + for t in _generate_reduction_inputs(device, dtype, requires_grad): + for reduction_kwargs in _generate_reduction_kwargs( + t.ndim, supports_multiple_dims + ): + for args, kwargs in generate_args_kwargs(t, **reduction_kwargs): + kwargs.update(reduction_kwargs) + yield SampleInput( + t.detach().requires_grad_(requires_grad), args=args, kwargs=kwargs + ) + + +# NOTE [Reductions]: +# +# For testing purposes, we relax the definition of a reduction operator +# as defined in the docstring below. We do this to capture operators with +# a similar API so they can be tested automatically. However... +# +# Strictly speaking a reduction operator is an operator that can reduce an +# array to a single scalar value and that can be computed from the partial +# result of reducing subarrays. This usually means that the reduction operation +# should be commutative and associative. This definition is important when it +# comes to implementation as it determines how a reduction can be parallelized. +# +# For example, many summary statistics such as median, mode and quantile cannot +# be computed from partial results because these are sorting and counting based +# algorithms that need information that would be lost in the reduced value. +class ReductionOpInfo(OpInfo): + """Reduction operator information. + + An operator is a reduction operator if it reduces one or more dimensions of + the input tensor to a single value. Reduction operators must implement the + following signature: + + - `op(input, *args, *, dim=None, keepdim=False, **kwargs) -> Tensor` + + ReductionOpInfo tests that reduction operators implement a consistent API. + Optional features such as reducing over multiple dimensions are captured in + the optional keyword parameters of the ReductionOpInfo constructor. + + If a reduction operator does not yet implement the full required API of + reduction operators, this should be documented by xfailing the failing + tests rather than adding optional parameters to ReductionOpInfo. + + NOTE + The API for reduction operators has not yet been finalized and some + requirements may change. + + See tests in test/test_reductions.py + """ + + def __init__( + self, + name, + *, + # The identity value for the operator if it has one. + identity: Any | None = None, + # The nan policy for the operator if it implements one. + # - propagate: NaN values are propagated to the output + # - omit: NaN values are discarded during the reduction + nan_policy: str | None = None, + # Whether the operator supports reducing multiple dimensions. + supports_multiple_dims: bool = True, + # Whether the operator promotes integral to floating point dtypes. + promotes_int_to_float: bool = False, + # Whether the operator promotes all integral dtypes to int64. + promotes_int_to_int64: bool = False, + # If a specific dtype is given, then the operator always returns that + # dtype irrespective of the input dtype. If None, the operator returns + # the dtype according to the type promotion rules above. + result_dtype: torch.dtype | None = None, + # Casts complex results to real (e.g. linalg.norm or torch.var) + complex_to_real: bool = False, + # ReductionOpInfo tests generate their own input, dim and keepdim + # arguments and call this function to generate tuples of extra args and + # kwargs to use when calling the op. This is required for operators that + # have other required parameters besides the input tensor. + generate_args_kwargs: Callable = lambda t, dim=None, keepdim=False: ( + yield ( + (), + {}, + ) + ), + # Options from the OpInfo base class + **kwargs, + ): + self._original_reduction_args = locals().copy() + if nan_policy not in (None, "propagate", "omit"): + raise AssertionError( + f"nan_policy must be None, 'propagate', or 'omit', got {nan_policy}" + ) + + # These are mutually exclusive options + if result_dtype and promotes_int_to_float: + raise AssertionError( + "result_dtype and promotes_int_to_float are mutually exclusive" + ) + if result_dtype and promotes_int_to_int64: + raise AssertionError( + "result_dtype and promotes_int_to_int64 are mutually exclusive" + ) + if result_dtype and complex_to_real: + raise AssertionError( + "result_dtype and complex_to_real are mutually exclusive" + ) + if promotes_int_to_float and promotes_int_to_int64: + raise AssertionError( + "promotes_int_to_float and promotes_int_to_int64 are mutually exclusive" + ) + + # Default sample_inputs_func for ReductionOpInfo which augments sample + # inputs from sample_inputs_reduction with the args and kwargs from + # generate_args_kwargs. This is only used if sample_inputs_func is None. + def sample_inputs_func(*args, **kwargs): + kwargs["supports_multiple_dims"] = supports_multiple_dims + kwargs["generate_args_kwargs"] = generate_args_kwargs + yield from sample_inputs_reduction(*args, **kwargs) + + # Override OpInfo defaults and call base class __init__ + kwargs.setdefault("inplace_variant", None) + kwargs.setdefault("sample_inputs_func", sample_inputs_func) + super().__init__(name, promotes_int_to_float=promotes_int_to_float, **kwargs) + + self.identity = identity + self.nan_policy = nan_policy + self.supports_multiple_dims = supports_multiple_dims + self.promotes_int_to_int64 = promotes_int_to_int64 + self.complex_to_real = complex_to_real + self.result_dtype = result_dtype + self.generate_args_kwargs = generate_args_kwargs + + +# The base reference input generation for elementwise binary operations +def _reference_inputs_elementwise_binary( + op, device, dtype, requires_grad, exclude_zero, **kwargs +): + yield from op.sample_inputs_func(op, device, dtype, requires_grad, **kwargs) + yield from generate_elementwise_binary_tensors( + op, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + if dtype is not torch.bool: + yield from generate_elementwise_binary_small_value_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad + ) + if dtype not in (torch.bool, torch.uint8, torch.int8): + yield from generate_elementwise_binary_large_value_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad + ) + yield from generate_elementwise_binary_broadcasting_tensors( + op, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + yield from generate_elementwise_binary_with_scalar_samples( + op, device=device, dtype=dtype, requires_grad=requires_grad + ) + + yield from generate_elementwise_binary_with_scalar_and_type_promotion_samples( + op, device=device, dtype=dtype, requires_grad=requires_grad + ) + + if dtype.is_floating_point or dtype.is_complex: + yield from generate_elementwise_binary_extremal_value_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad + ) + + +# Note that these references inputs use scalars for the SampleInput.input value, +# and many tests require SampleInput.input be a tensor or a list of tensors +def reference_inputs_elementwise_binary(op, device, dtype, requires_grad, **kwargs): + if hasattr(op, "rhs_make_tensor_kwargs"): + exclude_zero = op.rhs_make_tensor_kwargs.get("exclude_zero", False) + + gen = partial( + _reference_inputs_elementwise_binary, + op, + device, + dtype, + requires_grad, + exclude_zero, + **kwargs, + ) + + # yields "normal" samples + yield from gen() + + # yields noncontiguous samples + for sample in gen(): + yield sample.noncontiguous() + + yield from generate_elementwise_binary_noncontiguous_tensors( + op, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + + yield from generate_elementwise_binary_arbitrarily_strided_tensors( + op, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + + +# A functional that extends an elementwise binary operator's bespoke error inputs +# with generic error inputs for the class of elementwise binary operations +def make_error_inputs_elementwise_binary(error_inputs_func): + def error_inputs_func_wrapper(op, device, **kwargs): + if error_inputs_func is not None: + yield from error_inputs_func(op, device, **kwargs) + + if not op.supports_rhs_python_scalar: + si = SampleInput(torch.tensor((1, 2, 3), device=device), args=(2,)) + yield ErrorInput(si, error_type=Exception, error_regex="") + + if not op.supports_one_python_scalar: + si = SampleInput(2, args=(torch.tensor((1, 2, 3), device=device),)) + yield ErrorInput(si, error_type=Exception, error_regex="") + + if ( + not kwargs.get("skip_two_python_scalars", False) + and not op.supports_two_python_scalars + ): + si = SampleInput(2, args=(3,)) + yield ErrorInput(si, error_type=Exception, error_regex="") + + return error_inputs_func_wrapper + + +# The following functions and classes are for testing elementwise binary operators. + + +# Returns a generator of pairs of contiguous tensors on the requested device +# and with the requested dtype. +# +# This function is intended to test the non-vectorized and vectorized code +# paths of elementwise binary functions, as well as their handling of odd tensor +# sizes (like zero-dim tensors and tensors with zero elements). +# +# Each iterable will include an a tensor with no elements, +# zero dim (scalar) tensors, small 1D tensors, a medium 1D tensor, and +# a large 2D tensor. +def generate_elementwise_binary_tensors( + op, *, device, dtype, requires_grad=False, exclude_zero=False +): + shapes = ( + # tensors with no elements + (0,), + (1, 0, 3), + # zero dim (scalar) tensor + (), + # small 1D tensor + (20,), + # medium 1D tensor + (812,), + # large 2D tensor + (1029, 917), + ) + + make_arg = partial( + make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + for shape in shapes: + lhs = make_arg(shape, **op.lhs_make_tensor_kwargs) + rhs = make_arg(shape, **op.rhs_make_tensor_kwargs) + yield SampleInput( + lhs, args=(rhs,), kwargs=op.sample_kwargs(device, dtype, lhs)[0] + ) + + +def generate_elementwise_binary_arbitrarily_strided_tensors( + op, *, device, dtype, requires_grad=False, exclude_zero=False +): + # shape, strides, offset + strided_cases = ( + ((5, 6, 2), (1, 1, 7), 2), + ((5, 5, 4), (1, 1, 7), 2), + ((5, 5, 2), (4, 5, 7), 3), + ((5, 5, 2), (5, 5, 7), 3), + ((5, 5, 2), (5, 5, 5), 3), + ((9, 5, 2), (0, 1, 7), 3), + ) + + make_arg = partial( + make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + for shape, strides, offset in strided_cases: + a = make_arg( + 500, + ).as_strided(shape, strides, offset) + b = make_arg(shape) + yield SampleInput(a, args=(b,), kwargs=op.sample_kwargs(device, dtype, a)[0]) + + +# Returns a generator of pairs of contiguous tensors on the requested device and with +# the requested dtype. +# +# Unlike the previous function, the values in these tensors are specified manually. +def generate_elementwise_binary_small_value_tensors( + op, *, device, dtype, requires_grad=False, exclude_zero=None +): + if exclude_zero is None: + if hasattr(op, "rhs_make_tensor_kwargs"): + exclude_zero = op.rhs_make_tensor_kwargs.get("exclude_zero", False) + + # defines interesting values + _unsigned_int_vals = (0, 1, 55, 127, 128, 190, 210, 220, 254) + _int_vals = (0, -1, 1, -55, 55, -127, 127, -128) + _float_vals = ( + 0.0, + -0.0, + -0.001, + 0.001, + -0.25, + 0.25, + -1.0, + 1.0, + -math.pi / 2, + math.pi / 2, + -math.pi + 0.00001, + math.pi - 0.00001, + -math.pi, + math.pi, + -math.pi - 0.00001, + math.pi + 0.00001, + ) + + l_vals = [] + r_vals = [] + + if dtype.is_floating_point: + prod = product(_float_vals, _float_vals) + elif dtype.is_complex: + complex_vals = product(_float_vals, _float_vals) + # Note the use of list is required here or the map generator will be + # emptied by the following product and it won't produce the desired cross-product + complex_vals = [complex(*x) for x in complex_vals] + prod = product(complex_vals, complex_vals) + elif dtype in (torch.int8, torch.int16, torch.int32, torch.int64): + prod = product(_int_vals, _int_vals) + elif dtype in (torch.uint8, torch.uint16, torch.uint32, torch.uint64): + prod = product(_unsigned_int_vals, _unsigned_int_vals) + else: + raise ValueError("Unsupported dtype!") + + for l, r in prod: + l_vals.append(l) + if r == 0 and exclude_zero: + r_vals.append(1) + else: + r_vals.append(r) + + lhs = torch.tensor(l_vals, device=device, dtype=dtype, requires_grad=requires_grad) + rhs = torch.tensor(r_vals, device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(lhs, args=(rhs,), kwargs=op.sample_kwargs(device, dtype, lhs)[0]) + + +def generate_elementwise_binary_large_value_tensors( + op, *, device, dtype, requires_grad=False +): + _large_int_vals = (-1113, 1113, -10701, 10701) + _large_float16_vals = (-501, 501, -1001.2, 1001.2, -13437.7, 13437.7) + _large_float_vals = _large_float16_vals + (-4988429.2, 4988429.2, -1e20, 1e20) + _large_uint_vals = (1113, 10701, 60000) + + l_vals = [] + r_vals = [] + + if dtype == torch.float16: + prod = product(_large_float16_vals, _large_float16_vals) + elif dtype.is_floating_point: + prod = product(_large_float_vals, _large_float_vals) + elif dtype.is_complex: + complex_vals = product(_large_float_vals, _large_float_vals) + # Note the use of list is required here or the map generator will be + # emptied by the following product and it won't produce the desired cross-product + complex_vals = [complex(*x) for x in complex_vals] + prod = product(complex_vals, complex_vals) + elif dtype in (torch.int16, torch.int32, torch.int64): + prod = product(_large_int_vals, _large_int_vals) + elif dtype in (torch.uint16, torch.uint32, torch.uint64): + prod = product(_large_uint_vals, _large_uint_vals) + else: + raise ValueError("Unsupported dtype!") + + for l, r in prod: + l_vals.append(l) + r_vals.append(r) + + lhs = torch.tensor(l_vals, device=device, dtype=dtype, requires_grad=requires_grad) + rhs = torch.tensor(r_vals, device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(lhs, args=(rhs,), kwargs=op.sample_kwargs(device, dtype, lhs)[0]) + + +def generate_elementwise_binary_extremal_value_tensors( + op, *, device, dtype, requires_grad=False +): + _float_extremals = (float("inf"), float("-inf"), float("nan")) + + l_vals = [] + r_vals = [] + + if dtype.is_floating_point: + prod = product(_float_extremals, _float_extremals) + elif dtype.is_complex: + complex_vals = product(_float_extremals, _float_extremals) + # Note the use of list is required here or the map generator will be + # emptied by the following product and it won't produce the desired cross-product + complex_vals = [complex(*x) for x in complex_vals] + prod = product(complex_vals, complex_vals) + else: + raise ValueError("Unsupported dtype!") + + for l, r in prod: + l_vals.append(l) + r_vals.append(r) + + lhs = torch.tensor(l_vals, device=device, dtype=dtype, requires_grad=requires_grad) + rhs = torch.tensor(r_vals, device=device, dtype=dtype, requires_grad=requires_grad) + + yield SampleInput(lhs, args=(rhs,), kwargs=op.sample_kwargs(device, dtype, lhs)[0]) + + # Test case for NaN propagation + nan = ( + float("nan") if dtype.is_floating_point else complex(float("nan"), float("nan")) + ) + lhs = make_tensor( + (128, 128), device=device, dtype=dtype, requires_grad=requires_grad + ) + lhs.view(-1)[::3] = nan + rhs = make_tensor( + (128, 128), device=device, dtype=dtype, requires_grad=requires_grad + ) + rhs.view(-1)[::3] = nan + + yield SampleInput(lhs, args=(rhs,), kwargs=op.sample_kwargs(device, dtype, lhs)[0]) + + +# Returns a generator of pairs of contiguous and noncontiguous tensors that +# require broadcasting +def generate_elementwise_binary_broadcasting_tensors( + op, *, device, dtype, requires_grad=False, exclude_zero=False +): + shapes = ( + ((1,), ()), + ((2,), ()), + ((1,), (2,)), + ((2, 1), (2,)), + ((1, 2), (2,)), + ((3, 2), (2,)), + ((1, 3, 2), (2,)), + ((1, 3, 2), (3, 2)), + ((3, 1, 2), (3, 2)), + ((2, 3, 2), ()), + ((3, 1, 2), (1, 3, 2)), + ) + + make_arg = partial( + make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + for shape, noncontiguous in product(shapes, [True, False]): + shape_lhs, shape_rhs = shape + lhs = make_arg( + shape_lhs, noncontiguous=noncontiguous, **op.lhs_make_tensor_kwargs + ) + rhs = make_arg( + shape_rhs, noncontiguous=noncontiguous, **op.rhs_make_tensor_kwargs + ) + + yield SampleInput( + lhs, + args=(rhs,), + broadcasts_input=True, + kwargs=op.sample_kwargs(device, dtype, lhs)[0], + ) + + +# Returns a generator of pairs of contiguous tensors and scalars +def generate_elementwise_binary_with_scalar_samples( + op, *, device, dtype, requires_grad=False +): + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + + shapes = ((), (3,), (5, 3), (0, 1, 3), (1, 5)) + if op.supports_rhs_python_scalar: + for shape in shapes: + lhs = make_arg(shape, **op.lhs_make_tensor_kwargs) + rhs = make_arg(shape, **op.rhs_make_tensor_kwargs) + lhs_scalar = make_arg((), **op.lhs_make_tensor_kwargs).item() + rhs_scalar = make_arg((), **op.rhs_make_tensor_kwargs).item() + + yield SampleInput( + lhs, args=(rhs_scalar,), kwargs=op.sample_kwargs(device, dtype, lhs)[0] + ) + + # Extends with scalar lhs + if op.supports_one_python_scalar: + yield SampleInput( + lhs_scalar, + args=(rhs,), + kwargs=op.sample_kwargs(device, dtype, lhs_scalar)[0], + ) + + if op.supports_two_python_scalars: + lhs_scalar = make_arg((), **op.lhs_make_tensor_kwargs).item() + rhs_scalar = make_arg((), **op.rhs_make_tensor_kwargs).item() + + yield SampleInput( + lhs_scalar, + args=(rhs_scalar,), + kwargs=op.sample_kwargs(device, dtype, lhs_scalar)[0], + ) + + +# Returns a generator of pairs of contiguous tensors and 0d tensors and scalars and type promotion +def generate_elementwise_binary_with_scalar_and_type_promotion_samples( + op, *, device, dtype, requires_grad=False +): + # add these samples only for logical and comparison ops, arithmetic ops are not happy about extremal scalars + if op.name in ( + "eq", + "ne", + "gt", + "ge", + "lt", + "le", + "logical_and", + "logical_or", + "logical_xor", + ): + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + shape = ( + 23, + ) # this shape is big enough to trigger vectorization, and has non-vectorized tail + values = (float("nan"), float("inf"), -float("inf")) + scalar_tensors = tuple(torch.tensor(val) for val in values) + if op.supports_rhs_python_scalar: + lhs = make_arg(shape, **op.lhs_make_tensor_kwargs) + rhs = make_arg(shape, **op.rhs_make_tensor_kwargs) + for scalar in values + scalar_tensors: + yield SampleInput( + lhs, args=(scalar,), kwargs=op.sample_kwargs(device, dtype, lhs)[0] + ) + # Extends with scalar lhs + if op.supports_one_python_scalar: + yield SampleInput( + scalar, + args=(rhs,), + kwargs=op.sample_kwargs(device, dtype, scalar)[0], + ) + + +# Returns a generator of pairs of noncontiguous tensors +def generate_elementwise_binary_noncontiguous_tensors( + op, *, device, dtype, requires_grad=False, exclude_zero=False +): + make_arg = partial( + make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + + # Generic noncontiguity + lhs = make_arg((1026,), noncontiguous=True, **op.lhs_make_tensor_kwargs) + rhs = make_arg((1026,), noncontiguous=True, **op.rhs_make_tensor_kwargs) + + yield SampleInput( + lhs.clone(), args=(rhs.clone(),), kwargs=op.sample_kwargs(device, dtype, lhs)[0] + ) + yield SampleInput( + lhs.contiguous(), args=(rhs,), kwargs=op.sample_kwargs(device, dtype, lhs)[0] + ) + + # Transposed + lhs = make_arg((789, 357), **op.lhs_make_tensor_kwargs) + rhs = make_arg((789, 357), **op.rhs_make_tensor_kwargs) + + yield SampleInput( + lhs.T, args=(rhs.T,), kwargs=op.sample_kwargs(device, dtype, lhs)[0] + ) + + # More noncontiguity + shapes = ((5, 7), (1024,)) + + for shape in shapes: + lhs = make_arg(shape, **op.lhs_make_tensor_kwargs) + rhs = make_arg(shape, **op.rhs_make_tensor_kwargs) + + lhs_non_contig = torch.empty(shape + (2,), device=device, dtype=dtype)[..., 0] + lhs_non_contig.copy_(lhs) + + rhs_non_contig = torch.empty(shape + (2,), device=device, dtype=dtype)[..., 0] + rhs_non_contig.copy_(rhs) + + yield SampleInput( + lhs_non_contig.clone(), + args=(rhs_non_contig.clone(),), + kwargs=op.sample_kwargs(device, dtype, lhs)[0], + ) + yield SampleInput( + lhs_non_contig.contiguous(), + args=(rhs_non_contig,), + kwargs=op.sample_kwargs(device, dtype, lhs)[0], + ) + + # Noncontiguous indices + shape = (2, 2, 1, 2) + lhs = make_arg(shape, **op.lhs_make_tensor_kwargs) + rhs = make_arg(shape, **op.rhs_make_tensor_kwargs) + + lhs_non_contig = lhs[:, 1, ...] + rhs_non_contig = rhs[:, 1, ...] + + yield SampleInput( + lhs_non_contig.clone(), + args=(rhs_non_contig.clone(),), + kwargs=op.sample_kwargs(device, dtype, lhs)[0], + ) + yield SampleInput( + lhs_non_contig.contiguous(), + args=(rhs_non_contig,), + kwargs=op.sample_kwargs(device, dtype, lhs)[0], + ) + + # Expanded tensors + shapes = ((1, 3), (1, 7), (5, 7)) + + for shape in shapes: + lhs = make_arg(shape, **op.lhs_make_tensor_kwargs) + rhs = make_arg(shape, **op.rhs_make_tensor_kwargs) + + lhs_non_contig = lhs.expand(3, -1, -1) + rhs_non_contig = rhs.expand(3, -1, -1) + + yield SampleInput( + lhs_non_contig, + args=(rhs_non_contig,), + kwargs=op.sample_kwargs(device, dtype, lhs)[0], + ) + + +# Sample inputs for elementwise binary operators, like add +def sample_inputs_elementwise_binary(op, device, dtype, requires_grad, **kwargs): + _M = S if kwargs.get("small_inputs_only", False) else M + _S = XS if kwargs.get("small_inputs_only", False) else S + + if hasattr(op, "rhs_make_tensor_kwargs"): + exclude_zero = op.rhs_make_tensor_kwargs.get("exclude_zero", False) + + make_arg = partial( + make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + + shapes = ( + ((), ()), + ((_S,), ()), + ((_S, 1), (_S,)), + ((_M, _S), ()), + ((_S, _M, _S), (_M, _S)), + ((_S, _M, _S), (_S, _M, _S)), + ((_M, 1, _S), (_M, _S)), + ((_M, 1, _S), (1, _M, _S)), + ((0, 1, XS), (0, _M, XS)), + ) + + for shape_lhs, shape_rhs in shapes: + lhs = make_arg(shape_lhs, **op.lhs_make_tensor_kwargs) + rhs = make_arg(shape_rhs, **op.rhs_make_tensor_kwargs) + broadcasts_input = shape_lhs != torch.broadcast_shapes(shape_lhs, shape_rhs) + + yield SampleInput( + lhs, + args=(rhs,), + kwargs=op.sample_kwargs(device, dtype, lhs)[0], + broadcasts_input=broadcasts_input, + ) + + +# Metadata class for binary "universal functions (ufuncs)" that accept two +# tensor and have common properties +class BinaryUfuncInfo(OpInfo): + """Operator information for 'universal binary functions (binary ufuncs).' + These are functions of two tensors with common properties like: + - they are elementwise functions + - the output shape is determined by the input shape + - they typically have method and inplace variants + - they typically support the out kwarg + - they typically have NumPy or SciPy references + See NumPy's universal function documentation + (https://numpy.org/doc/stable/reference/ufuncs.html) for more details + about the concept of ufuncs. + """ + + def __init__( + self, + name, + *, + sample_inputs_func=sample_inputs_elementwise_binary, + reference_inputs_func=reference_inputs_elementwise_binary, + sample_kwargs=lambda device, dtype, input: ({}, {}), + error_inputs_func=None, + lhs_make_tensor_kwargs=None, + rhs_make_tensor_kwargs=None, + always_returns_bool=False, # Set to true if the op always returns bool tensors + supports_rhs_python_scalar=True, # Whether the operator allows Tensor x scalar inputs + supports_one_python_scalar=False, # Whether the operator allows scalar x tensor and tensor x scalar inputs + supports_two_python_scalars=False, # Whether the operator allows scalar x scalar inputs + **kwargs, + ): + self._original_binary_ufunc_args = locals().copy() + + # Elementwise binary operations perform the equivalent of test_numpy_refs + # in test_binary_ufuncs, but with additional test granularity. So the + # generic test_ops.py test is skipped because it's redundant. + common_skips = ( + DecorateInfo( + unittest.skip("Skipping redundant test."), + "TestCommon", + "test_numpy_refs", + ), + ) + kwargs["skips"] = kwargs.get("skips", ()) + common_skips + super().__init__( + name, + sample_inputs_func=sample_inputs_func, + reference_inputs_func=reference_inputs_func, + error_inputs_func=make_error_inputs_elementwise_binary(error_inputs_func), + **kwargs, + ) + + self.sample_kwargs = sample_kwargs + + # [lr]hs_make_tensor_kwargs are part of the OpInfo to be able to dynamically generate valid samples later on. + if lhs_make_tensor_kwargs is None: + lhs_make_tensor_kwargs = {} + self.lhs_make_tensor_kwargs = lhs_make_tensor_kwargs + + if rhs_make_tensor_kwargs is None: + rhs_make_tensor_kwargs = {} + self.rhs_make_tensor_kwargs = rhs_make_tensor_kwargs + + self.always_returns_bool = always_returns_bool + self.supports_rhs_python_scalar = supports_rhs_python_scalar + self.supports_one_python_scalar = supports_one_python_scalar + self.supports_two_python_scalars = supports_two_python_scalars + + if self.supports_two_python_scalars: + self.supports_one_python_scalar = True + + if self.supports_one_python_scalar: + if not supports_rhs_python_scalar: + raise AssertionError( + "Can't support lhs and rhs Python scalars but not rhs scalars!" + ) + + +# The following functions and classes are for testing elementwise unary operators. +def sample_inputs_elementwise_unary( + op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs +): + if not op_kwargs: + op_kwargs = {} + + _L = S if kwargs.get("small_inputs_only", False) else L + + low, high = op_info.domain + is_floating = dtype.is_floating_point or dtype.is_complex + low = low if low is None or not is_floating else low + op_info._domain_eps + high = high if high is None or not is_floating else high - op_info._domain_eps + if ( + op_info.supports_sparse_csr + or op_info.supports_sparse_csc + or op_info.supports_sparse_bsr + or op_info.supports_sparse_bsc + ): + # Tensors with dim=2 for sparse compressed testing + yield SampleInput( + make_tensor( + (_L, _L), + device=device, + dtype=dtype, + low=low, + high=high, + requires_grad=requires_grad, + ), + kwargs=op_kwargs, + ) + else: + # Creates a 1D, empty, and scalar tensor + for shape in ((_L,), (1, 0, 3), ()): + yield SampleInput( + make_tensor( + shape, + device=device, + dtype=dtype, + low=low, + high=high, + requires_grad=requires_grad, + ), + kwargs=op_kwargs, + ) + + +# Replace values satisfying condition with a safe value. This is used to block +# out values the could cause singularity like tan(pi/2) +def _replace_values_in_tensor(tensor, condition, safe_value): + mask = condition(tensor) + tensor.masked_fill_(mask, safe_value) + + +# Helper to create a unary elementwise tensor with valid inputs +def _make_unary_elementwise_tensor(shape, *, op, dtype, **kwargs): + low, high = op.domain + is_floating = dtype.is_floating_point or dtype.is_complex + low = low if low is None or not is_floating else low + op._domain_eps + high = high if high is None or not is_floating else high - op._domain_eps + + a = make_tensor(shape, low=low, high=high, dtype=dtype, **kwargs) + + if op.reference_numerics_filter is not None and dtype is not torch.bool: + condition, safe_value = op.reference_numerics_filter + _replace_values_in_tensor(a, condition, safe_value) + + return a + + +# Restricts the values in the tensor to the domain of the +# given elementwise unary operator +def _filter_unary_elementwise_tensor(a, *, op): + # short-circuits for boolean tensors + if a.dtype is torch.bool: + return a + + low, high = op.domain + is_floating = a.dtype.is_floating_point or a.dtype.is_complex + low = low if low is None or not is_floating else low + op._domain_eps + high = high if high is None or not is_floating else high - op._domain_eps + + if a.dtype is torch.uint8 and low is not None: + low = max(low, 0) + + if not a.dtype.is_floating_point and not a.dtype.is_complex: + low = math.ceil(low) if low is not None else None + high = math.floor(high) if high is not None else None + + if op.reference_numerics_filter is not None: + condition, safe_value = op.reference_numerics_filter + _replace_values_in_tensor(a, condition, safe_value) + + if low is not None or high is not None: + if a.dtype.is_complex: + a.real.clamp_(low, high) + a.imag.clamp_(low, high) + else: + a.clamp_(min=low, max=high) + + return a + + +def generate_elementwise_unary_tensors(op, *, device, dtype, requires_grad, **kwargs): + # Special-cases bool + if dtype is torch.bool: + tensors = ( + torch.empty(0, device=device, dtype=torch.bool), + torch.tensor(True, device=device), + torch.tensor(False, device=device), + torch.tensor((True, False), device=device), + make_tensor((812,), device=device, dtype=dtype), + make_tensor((1029, 917), device=device, dtype=dtype), + ) + for a in tensors: + yield SampleInput(a, kwargs=op.sample_kwargs(device, dtype, a)[0]) + + shapes = ( + (1029, 917), + (812,), + # Empty sizes + (0,), + (0, 3, 3), + (1, 0, 5), + (6, 0, 0, 0), + (3, 0, 1, 0), + ) + + make_arg = partial( + _make_unary_elementwise_tensor, + op=op, + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + for shape in shapes: + a = make_arg(shape) + yield SampleInput(a, kwargs=op.sample_kwargs(device, dtype, a)[0]) + + +def generate_elementwise_unary_small_value_tensors( + op, *, device, dtype, requires_grad=False +): + for sample in generate_elementwise_binary_small_value_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad + ): + a = _filter_unary_elementwise_tensor(sample.input, op=op) + yield SampleInput(a, kwargs=op.sample_kwargs(device, dtype, a)[0]) + + +def generate_elementwise_unary_large_value_tensors( + op, *, device, dtype, requires_grad=False +): + for sample in generate_elementwise_binary_large_value_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad + ): + a = _filter_unary_elementwise_tensor(sample.input, op=op) + yield SampleInput(sample.input, kwargs=op.sample_kwargs(device, dtype, a)[0]) + + +def generate_elementwise_unary_extremal_value_tensors( + op, *, device, dtype, requires_grad=False +): + for sample in generate_elementwise_binary_extremal_value_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad + ): + yield SampleInput( + sample.input, kwargs=op.sample_kwargs(device, dtype, sample.input)[0] + ) + + +def generate_elementwise_unary_noncontiguous_tensors( + op, *, device, dtype, requires_grad=False +): + make_arg = partial( + _make_unary_elementwise_tensor, + op=op, + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + + # Generic noncontiguity + t = make_arg((1026,), noncontiguous=True) + yield SampleInput(t, kwargs=op.sample_kwargs(device, dtype, t)[0]) + + # Transposed + t = make_arg((1024, 1024)).T + yield SampleInput(t, kwargs=op.sample_kwargs(device, dtype, t)[0]) + + # Expanded tensors + shapes = ((1, 3), (1, 7), (5, 7)) + + for shape in shapes: + t = make_arg(shape) + t_non_contig = t.expand(3, -1, -1) + yield SampleInput( + t_non_contig, kwargs=op.sample_kwargs(device, dtype, t_non_contig)[0] + ) + + +def generate_elementwise_unary_arbitrarily_strided_tensors( + op, *, device, dtype, requires_grad=False +): + # shape, strides, offset + strided_cases = ( + ((5, 6, 2), (1, 1, 7), 2), + ((5, 5, 4), (1, 1, 7), 2), + ((5, 5, 2), (4, 5, 7), 3), + ((5, 5, 2), (5, 5, 7), 3), + ((5, 5, 2), (5, 5, 5), 3), + ((9, 5, 2), (0, 1, 7), 3), + ) + + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + for shape, strides, offset in strided_cases: + a = make_arg( + 500, + ).as_strided(shape, strides, offset) + yield SampleInput(a, kwargs=op.sample_kwargs(device, dtype, a)[0]) + + +# Reuses the elementwise binary generators for consistency +# TODO: in the future generalize the reference generators to handle n-ary elementwise operations +def _reference_inputs_elementwise_unary(op, device, dtype, requires_grad, **kwargs): + yield from op.sample_inputs_func(op, device, dtype, requires_grad, **kwargs) + + yield from generate_elementwise_unary_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad, **kwargs + ) + + if dtype is not torch.bool: + yield from generate_elementwise_unary_small_value_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad, **kwargs + ) + if dtype not in (torch.bool, torch.uint8, torch.int8) and ( + op.handles_large_floats + or (not dtype.is_floating_point and not dtype.is_complex) + ): + yield from generate_elementwise_unary_large_value_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad, **kwargs + ) + + if dtype.is_floating_point or ( + op.handles_complex_extremal_values and dtype.is_complex + ): + yield from generate_elementwise_unary_extremal_value_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad, **kwargs + ) + + +def reference_inputs_elementwise_unary(op, device, dtype, requires_grad, **kwargs): + gen = partial( + _reference_inputs_elementwise_unary, op, device, dtype, requires_grad, **kwargs + ) + + # yields "normal" samples + yield from gen() + + # yields noncontiguous samples + for sample in gen(): + yield sample.noncontiguous() + + yield from generate_elementwise_unary_noncontiguous_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad, **kwargs + ) + + yield from generate_elementwise_unary_arbitrarily_strided_tensors( + op, device=device, dtype=dtype, requires_grad=requires_grad, **kwargs + ) + + +# Metadata class for unary "universal functions (ufuncs)" that accept a single +# tensor and have common properties like: +class UnaryUfuncInfo(OpInfo): + """Operator information for 'universal unary functions (unary ufuncs).' + These are functions of a single tensor with common properties like: + - they are elementwise functions + - the input shape is the output shape + - they typically have method and inplace variants + - they typically support the out kwarg + - they typically have NumPy or SciPy references + See NumPy's universal function documentation + (https://numpy.org/doc/1.18/reference/ufuncs.html) for more details + about the concept of ufuncs. + """ + + def __init__( + self, + name, # the string name of the function + *, + dtypes=floating_types(), + domain=(None, None), # the [low, high) domain of the function + handles_complex_extremal_values=True, # whether the op correctly handles extremal values (like nan/inf) + handles_large_floats=True, # whether the op correctly handles large float values (like 1e20) + supports_complex_to_float=False, # op supports casting from complex input to real output safely eg. angle + sample_inputs_func=sample_inputs_elementwise_unary, + reference_inputs_func=reference_inputs_elementwise_unary, + sample_kwargs=lambda device, dtype, input: ({}, {}), + reference_numerics_filter=None, # Filters values in the range of the domain specified above but that should not be tested + **kwargs, + ): + self._original_unary_ufunc_args = locals().copy() + + super().__init__( + name, + dtypes=dtypes, + sample_inputs_func=sample_inputs_func, + reference_inputs_func=reference_inputs_func, + **kwargs, + ) + self.domain = domain + self.handles_complex_extremal_values = handles_complex_extremal_values + self.handles_large_floats = handles_large_floats + self.supports_complex_to_float = supports_complex_to_float + self.reference_numerics_filter = reference_numerics_filter + + # test_unary_ufuncs.py generates its own inputs to test the consistency + # of the operator on sliced tensors, non-contig tensors, etc. + # `sample_kwargs` is a utility function to provide kwargs + # along with those inputs if required (eg. clamp). + # It should return two dictionaries, first holding kwarg for + # torch operator and second one for reference NumPy operator. + self.sample_kwargs = sample_kwargs + + # Epsilon to ensure grad and gradgrad checks don't test values + # outside a function's domain. + self._domain_eps = 1e-5 + + +def sample_inputs_spectral_ops(self, device, dtype, requires_grad=False, **kwargs): + is_fp16_or_chalf = dtype == torch.complex32 or dtype == torch.half + if not is_fp16_or_chalf: + nd_tensor = partial( + make_tensor, + (S, S + 1, S + 2), + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + oned_tensor = partial( + make_tensor, (31,), device=device, dtype=dtype, requires_grad=requires_grad + ) + else: + # cuFFT supports powers of 2 for half and complex half precision + # NOTE: For hfft, hfft2, hfftn, irfft, irfft2, irfftn with default args + # where output_size n=2*(input_size - 1), we make sure that logical fft size is a power of two + low = None + high = None + if self.name in ["fft.hfft", "fft.irfft", "_refs.fft.hfft", "_refs.fft.irfft"]: + shapes = ((2, 9, 9), (33,)) + elif self.name in [ + "fft.hfft2", + "fft.irfft2", + "_refs.fft.hfft2", + "_refs.fft.irfft2", + ]: + shapes = ((2, 8, 9), (33,)) + elif self.name in [ + "fft.hfftn", + "fft.irfftn", + "_refs.fft.hfftn", + "_refs.fft.irfftn", + ]: + shapes = ((2, 2, 33), (33,)) + # Adjusting the limits because the test would be flaky due to over-saturation of float16 + # See: https://github.com/pytorch/pytorch/pull/81416 + low = -1.0 + high = 1.0 + else: + shapes = ((2, 8, 16), (32,)) + nd_tensor = partial( + make_tensor, + shapes[0], + device=device, + low=low, + high=high, + dtype=dtype, + requires_grad=requires_grad, + ) + oned_tensor = partial( + make_tensor, + shapes[1], + device=device, + low=low, + high=high, + dtype=dtype, + requires_grad=requires_grad, + ) + + if self.ndimensional == SpectralFuncType.ND: + yield SampleInput( + nd_tensor(), + s=(3, 10) if not is_fp16_or_chalf else (4, 8), + dim=(1, 2), + norm="ortho", + ) + yield SampleInput(nd_tensor(), norm="ortho") + yield SampleInput(nd_tensor(), s=(8,)) + yield SampleInput(oned_tensor()) + yield from (SampleInput(nd_tensor(), dim=dim) for dim in [-1, -2, -3, (0, -1)]) + elif self.ndimensional == SpectralFuncType.TwoD: + yield SampleInput( + nd_tensor(), + s=(3, 10) if not is_fp16_or_chalf else (4, 8), + dim=(1, 2), + norm="ortho", + ) + yield SampleInput(nd_tensor(), norm="ortho") + yield SampleInput(nd_tensor(), s=(6, 8) if not is_fp16_or_chalf else (4, 8)) + yield SampleInput(nd_tensor(), dim=0) + yield SampleInput(nd_tensor(), dim=(0, -1)) + yield SampleInput(nd_tensor(), dim=(-3, -2, -1)) + else: + yield SampleInput( + nd_tensor(), + n=10 if not is_fp16_or_chalf else 8, + dim=1, + norm="ortho", + ) + yield SampleInput(nd_tensor(), norm="ortho") + yield SampleInput(nd_tensor(), n=7 if not is_fp16_or_chalf else 8) + yield SampleInput(oned_tensor()) + yield from (SampleInput(nd_tensor(), dim=dim) for dim in [-1, -2, -3]) + + +SpectralFuncType = Enum("SpectralFuncType", ("OneD", "TwoD", "ND")) + + +# Metadata class for Fast Fourier Transforms in torch.fft. +class SpectralFuncInfo(OpInfo): + """Operator information for torch.fft transforms.""" + + def __init__( + self, + name, # the string name of the function + *, + ref=None, # Reference implementation (probably in np.fft namespace) + dtypes=floating_and_complex_types(), + ndimensional: SpectralFuncType, + sample_inputs_func=sample_inputs_spectral_ops, + decorators=None, + **kwargs, + ): + self._original_spectral_func_args = dict(locals()).copy() + self._original_spectral_func_args.update(kwargs) + + decorators = list(decorators) if decorators is not None else [] + decorators += [ + skipCPUIfNoFFT, + DecorateInfo( + toleranceOverride({torch.chalf: tol(4e-2, 4e-2)}), + "TestCommon", + "test_complex_half_reference_testing", + ), + ] + + super().__init__( + name=name, + dtypes=dtypes, + decorators=decorators, + sample_inputs_func=sample_inputs_func, + **kwargs, + ) + self.ref = ref + self.ndimensional = ndimensional + + +class ShapeFuncInfo(OpInfo): + """Early version of a specialized OpInfo for Shape manipulating operations like tile and roll""" + + def __init__( + self, + name, # the string name of the function + *, + ref, # a reference function + dtypes=floating_types(), + dtypesIfCUDA=None, + dtypesIfMPS=None, + dtypesIfROCM=None, + dtypesIfXPU=None, + sample_inputs_func=None, + **kwargs, + ): + super().__init__( + name, + dtypes=dtypes, + dtypesIfCUDA=dtypesIfCUDA, + dtypesIfMPS=dtypesIfMPS, + dtypesIfROCM=dtypesIfROCM, + dtypesIfXPU=dtypesIfXPU, + sample_inputs_func=sample_inputs_func, + **kwargs, + ) + self.ref = ref + + +def sample_inputs_foreach( + self, + device, + dtype, + N, + *, + noncontiguous=False, + same_size=False, + low=None, + high=None, + # zero_size means EVERY input is empty + zero_size: bool, + requires_grad: bool, + # mutually exclusive from same_size and zero_size, which are all or nothing + intersperse_empty_tensors: bool = False, +): + if zero_size: + return [torch.empty(0, dtype=dtype, device=device) for _ in range(N)] + if same_size: + return [ + make_tensor( + (N, N), + dtype=dtype, + device=device, + noncontiguous=noncontiguous, + low=low, + high=high, + requires_grad=requires_grad, + ) + for _ in range(N) + ] + else: + # interweave some empty tensors + have the last 2 tensors be empty (see #100701) + return [ + torch.empty(0, dtype=dtype, device=device, requires_grad=requires_grad) + if (i % 3 == 0 or i >= N - 2) and intersperse_empty_tensors + else make_tensor( + (N - i, N - i), + dtype=dtype, + device=device, + noncontiguous=noncontiguous, + low=low, + high=high, + requires_grad=requires_grad, + ) + for i in range(N) + ] + + +def get_foreach_method_names(name): + # get torch inplace reference function + op_name = "_foreach_" + name + inplace_op_name = op_name + "_" + + op = getattr(torch, op_name, None) + inplace_op = getattr(torch, inplace_op_name, None) + + ref = getattr(torch, name, None) + ref_inplace = getattr(torch.Tensor, name + "_", None) + return op, inplace_op, ref, ref_inplace + + +@dataclass +class ForeachFuncInfo(OpInfo): + """Early version of a specialized OpInfo for foreach functions + + The main differences from the parent class are (a) `dtypes`, `dtypesIfCUDA`, and `dtypesIfROCM` + are set to `get_all_dtypes(include_qint=False)`, and (b) the following arguments. + + ``supports_alpha_param=True`` means that the function supports a python scalar (``numbers.Number``) + as the last keyword argument such as `_foreach_add`. + ``supports_scalar_self_arg=True`` means that the function can take a python scalar as its first argument. + Currently only `_foreach_pow` supports this. + ``backward_requires_result=True``, which could sound self-explanatory, means that the function uses + the forward result for its backward computation. + """ + + supports_alpha_param: bool = False + supports_scalar_self_arg: bool = False + backward_requires_result: bool = False + + def __post_init__(self): + ( + foreach_method, + foreach_method_inplace, + torch_ref_method, + torch_ref_inplace, + ) = get_foreach_method_names(self.name) + if not self.supports_out: + # note(crcrpar): `foreach_method` for `"zero"` is `None` but `None` would call + # `_getattr_qual` in `OpInfo.__post_init__` which should fail since `_foreach_zero` + # is not defined at the moment. Thus to skip the qualification, set a similar torch + # function. + if foreach_method is not None: + raise AssertionError("foreach_method must be None") + if torch_ref_method is not None: + raise AssertionError("torch_ref_method must be None") + foreach_method = foreach_method_inplace + torch_ref_method = torch_ref_inplace + + # We disable all complex128 tests internally for foreach due to reported flakiness + # tracked in #139648 + supported_dtypes = get_all_dtypes(include_qint=False) + if IS_FBCODE: + supported_dtypes = [ + x for x in supported_dtypes if x is not torch.complex128 + ] + self.dtypes = _dispatch_dtypes(supported_dtypes) + + self.op = foreach_method + self.method_variant = foreach_method + self.ref = torch_ref_method + self.inplace_variant = foreach_method_inplace + self.ref_inplace = torch_ref_inplace + self.has_no_in_place = self.inplace_variant is None + + name = self.name + self.name = f"_foreach_{name}" + if name == "norm": + self.ref = torch.linalg.vector_norm + elif name == "minimum": + # because minimum ref does not support inplace or scalar + self.ref = torch.clamp_max + self.ref_inplace = torch.Tensor.clamp_max_ + elif name == "maximum": + # because maximum ref does not support inplace or scalar + self.ref = torch.clamp_min + self.ref_inplace = torch.Tensor.clamp_min_ + + # The following sets `dtypesIfCUDA` and `dtypesIfROCM` accordingly. + super().__post_init__() + + def sample_zero_size_inputs(self, device, dtype, requires_grad=False, **kwargs): + if not hasattr(self.sample_inputs_func, "sample_zero_size_tensor_inputs"): + return [] + return self.sample_inputs_func.sample_zero_size_tensor_inputs( + self, device, dtype, requires_grad, **kwargs + ) + + +def gradcheck_wrapper_hermitian_input(op, input, *args, **kwargs): + """Gradcheck wrapper for functions that take Hermitian matrices as input. + + They require a modified function because the finite-difference algorithm + for calculating derivatives does not preserve the Hermitian property of the input. + """ + return op(input + input.mH, *args, **kwargs) + + +def gradcheck_wrapper_ctc_loss(op, input, *args, **kwargs): + """Gradcheck wrapper for ctc loss to project onto log-simplex space.""" + # See https://github.com/pytorch/pytorch/issues/52241 + return op(input.log_softmax(dim=2), *args, **kwargs) + + +def gradcheck_wrapper_triangular_input(op, *args, upper=False, idx=0, **kwargs): + """Gradcheck wrapper for functions that take lower or upper triangular matrices as input. + + They require a modified function because the finite-difference algorithm + for calculating derivatives does not preserve the triangular property of the input. + `idx` is used to specific which `args[idx]` is to be triangularized. + """ + triangular_arg = args[idx].triu() if upper else args[idx].tril() + return op(*args[:idx], triangular_arg, *args[idx + 1 :], upper, **kwargs) + + +def gradcheck_wrapper_triangular_input_real_positive_diagonal( + op, *args, upper=False, idx=0, **kwargs +): + """Gradcheck wrapper for functions that take lower/upper triangular matrices + with real and positive diagonals, for example, cholesky-like operations. + """ + arg = args[idx] + arg_diag = arg.diagonal(0, -2, -1) + arg_diag_embed = torch.diag_embed(arg_diag) + id_diag_tensor = torch.ones_like(arg_diag) + id_tensor = torch.diag_embed(id_diag_tensor) + # new_arg = arg - diag(arg) + I + new_arg = arg - arg_diag_embed + id_tensor + return gradcheck_wrapper_triangular_input( + op, *args[:idx], new_arg, *args[idx + 1 :], upper=upper, idx=idx, **kwargs + ) + + +def gradcheck_wrapper_masked_operation(op, input, *args, **kwargs): + """Gradcheck wrapper for masked operations. + + When mask is specified, replaces masked-out elements with zeros. + + Use for operations that produce non-finite masked-out elements, + for instance, for minimum and maximum reductions. + """ + output = op(input, *args, **kwargs) + mask = kwargs.get("mask") + if mask is not None: + output_mask = torch.masked._output_mask(op, input, *args, **kwargs) + output = torch.where(output_mask, output, output.new_zeros([])) + return output + + +def gradcheck_wrapper_masked_pointwise_operation(op, input, *args, **kwargs): + """Gradcheck wrapper for masked pointwise operations. Assumes that the result + will be masked iff both tensors are masked at a specific index + + When mask is specified, replaces masked-out elements with zeros. + + Use for operations that produce non-finite masked-out elements, + for instance, for minimum and maximum reductions. + """ + output = op(input, *args, **kwargs) + input_mask = kwargs.get("input_mask") + other_mask = kwargs.get("other_mask") + if input_mask is not None and other_mask is not None: + combined_mask = torch.logical_and(input_mask, other_mask) + new_kwargs = dict(mask=combined_mask, **kwargs) + output_mask = torch.masked._input_mask(input, *args, **new_kwargs) + output = torch.where(output_mask, output, output.new_zeros([])) + return output + + +def clone_sample(sample, **kwargs): + """ + Given a SampleInput, this function analyzes its input, args and kwargs, + and produces a copy with each non-Tensor entry being copied by reference, + and with each Tensor entry cloned with `t.clone().requires_grad_(t.requires_grad)` + """ + + def clone_tensor(t): + if isinstance(t, torch.Tensor): + return t.detach().clone().requires_grad_(t.requires_grad) + else: + return t + + sample_kwargs = kwargs if kwargs else sample.kwargs + + return SampleInput( + clone_tensor(sample.input), + args=tuple(map(clone_tensor, sample.args)), + kwargs={k: clone_tensor(v) for k, v in sample_kwargs.items()}, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f26d3f402e741a54f21a5fca48beded5b0a58aec --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/__init__.py @@ -0,0 +1,26 @@ +# mypy: ignore-errors + +from torch.testing._internal.opinfo.core import OpInfo +from torch.testing._internal.opinfo.definitions import ( + _masked, + fft, + linalg, + signal, + special, +) + + +# Operator database +op_db: list[OpInfo] = [ + *fft.op_db, + *linalg.op_db, + *signal.op_db, + *special.op_db, + *_masked.op_db, +] + +python_ref_db: list[OpInfo] = [ + *fft.python_ref_db, + *linalg.python_ref_db, + *special.python_ref_db, +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/_masked.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/_masked.py new file mode 100644 index 0000000000000000000000000000000000000000..670b34c64ecdf17f04adff63f9006fba15691688 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/_masked.py @@ -0,0 +1,1386 @@ +# mypy: ignore-errors + +import unittest +from collections.abc import Sequence +from functools import partial + +import numpy as np + +import torch +from torch.testing import make_tensor +from torch.testing._internal.common_device_type import tol, toleranceOverride +from torch.testing._internal.common_dtype import ( + all_types_and, + all_types_and_complex_and, + complex_types, + floating_and_complex_types_and, + floating_types_and, + integral_types, +) +from torch.testing._internal.opinfo.core import ( + DecorateInfo, + gradcheck_wrapper_masked_operation, + gradcheck_wrapper_masked_pointwise_operation, + M, + OpInfo, + ReductionOpInfo, + S, + sample_inputs_reduction, + SampleInput, +) +from torch.testing._internal.opinfo.utils import prod_numpy, reference_reduction_numpy + + +# Used for log_softmax, softmax, softmin +def sample_inputs_softmax_variant( + op_info, + device, + dtype, + requires_grad, + with_dtype=False, + use_zero_dimensions=True, + **kwargs, +): + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + cases = [ + ((S,), (0,)), + ((S, S), (0,)), + ((S, S), (1,)), + ((S, S), (-1,)), + ((S, M, S), (2,)), + *([((S, 0, 0), (-1,))] if use_zero_dimensions else []), + ] + kwargs = ( + dict( + dtype=( + torch.bfloat16 if torch.device(device).type == "mps" else torch.float64 + ) + ) + if with_dtype + else None + ) + + # PyTorch on XLA throws an error when passed with dim argument for 0d tensor. + # See https://github.com/pytorch/xla/issues/3061 for more details. + if torch.device(device).type != "xla": + cases.append(((), (0,))) + + return ( + SampleInput(make_arg(shape), args=dim, kwargs=kwargs) for shape, dim in cases + ) + + +def _generate_masked_op_mask(input_shape, device, **kwargs): + make_arg = partial( + make_tensor, dtype=torch.bool, device=device, requires_grad=False + ) + yield None + yield make_arg(input_shape) + if len(input_shape) > 2: + # broadcast last mask dimension: + yield make_arg(input_shape[:-1] + (1,)) + # broadcast middle mask dimension: + yield make_arg(input_shape[:1] + (1,) + input_shape[2:]) + # broadcast first mask dimension: + yield make_arg((1,) + input_shape[1:]) + # mask.ndim < input.ndim + yield make_arg(input_shape[1:]) + # mask.ndim == 1 + yield make_arg(input_shape[-1:]) + # masks that require broadcasting of inputs (mask.ndim > + # input.ndim) will not be supported, however, we may + # reconsider this if there will be demand on this kind of + # degenerate cases. + + +def sample_inputs_masked_reduction(op_info, device, dtype, requires_grad, **kwargs): + """Sample inputs for masked reduction operators. + + Masked reduction operator is a reduction operator with trailing + mask optional argument. A mask is a bool tensor with the same + shape as input or a shape that is broadcastable to input shape. + """ + kwargs["supports_multiple_dims"] = op_info.supports_multiple_dims + + for sample_input in sample_inputs_reduction( + op_info, device, dtype, requires_grad, **kwargs + ): + for mask in _generate_masked_op_mask( + sample_input.input.shape, device, **kwargs + ): + sample_input_args, sample_input_kwargs = ( + sample_input.args, + dict(mask=mask, **sample_input.kwargs), + ) + yield SampleInput( + sample_input.input.detach().requires_grad_(requires_grad), + args=sample_input_args, + kwargs=sample_input_kwargs, + ) + if ( + not requires_grad + and dtype.is_floating_point + and sample_input.input.ndim == 2 + and mask is not None + and mask.shape == sample_input.input.shape + ): + for v in [torch.inf, -torch.inf, torch.nan]: + t = sample_input.input.detach() + t.diagonal(0, -2, -1).fill_(v) + yield SampleInput( + t.requires_grad_(requires_grad), + args=sample_input_args, + kwargs=sample_input_kwargs, + ) + + +def sample_inputs_sparse_coo_masked_reduction( + op_info, device, dtype, requires_grad, **kwargs +): + """Sample inputs for masked reduction operators that support inputs + with sparse coo layouts. + """ + if op_info.supports_sparse: + op_name = op_info.name.replace("masked.", "") + for sample_input in sample_inputs_masked_reduction( + op_info, device, dtype, requires_grad, **kwargs + ): + mask = sample_input.kwargs.get("mask") + if mask is not None: + sample_input_kwargs = sample_input.kwargs.copy() + sample_input_kwargs.update(mask=mask.to_sparse()) + yield SampleInput( + sample_input.input.to_sparse(), + args=sample_input.args, + kwargs=sample_input_kwargs, + ) + else: + if op_name in {"prod", "amax", "amin"}: + # FIXME: for now reductions with non-zero reduction identity and + # unspecified mask are not supported for sparse COO + # tensors, see torch.masked.prod implementation + # for details. + continue + yield SampleInput( + sample_input.input.to_sparse(), + args=sample_input.args, + kwargs=sample_input.kwargs, + ) + + +def sample_inputs_sparse_csr_masked_reduction( + op_info, device, dtype, requires_grad, **kwargs +): + """Sample inputs for masked reduction operators that support inputs + with sparse csr layouts. + """ + if op_info.supports_sparse_csr: + op_name = op_info.name.replace("masked.", "") + for sample_input in sample_inputs_masked_reduction( + op_info, device, dtype, requires_grad, **kwargs + ): + if not ( + sample_input.input.ndim == 2 and sample_input.kwargs.get("keepdim") + ): + # - sparse CSR tensors are always 2-D tensors + # - masked reduction on CSR tensors are defined only if keepdim is True. + continue + mask = sample_input.kwargs.get("mask") + if mask is not None: + sample_input_kwargs = sample_input.kwargs.copy() + sample_input_kwargs.update(mask=mask.to_sparse_csr()) + new_sample = SampleInput( + sample_input.input.to_sparse_csr(), + args=sample_input.args, + kwargs=sample_input_kwargs, + ) + else: + if op_name in ["prod", "amax", "amin", "mean"]: + # reductions with non-zero reduction identity and + # unspecified mask is not supported for sparse CSR + # tensors, see torch.masked.prod implementation + # for details. + continue + new_sample = SampleInput( + sample_input.input.to_sparse_csr(), + args=sample_input.args, + kwargs=sample_input.kwargs, + ) + yield new_sample + if sample_input.kwargs["dim"] == 0: + # Reductions of CSR tensors use different implementations for + # inner and/or outer dimensions. So, as a minimum of testing CSR + # implementations the following kwargs must be generated: + # dict(dim=0, keepdim=True) + # dict(dim=1, keepdim=True) + # dict(dim=(0, 1), keepdim=True) + # Here we generate the dim=1 case from the dim=0 case. + sample_input_kwargs = new_sample.kwargs.copy() + sample_input_kwargs.update(dim=1) + yield SampleInput( + new_sample.input.clone(), + args=sample_input.args, + kwargs=sample_input_kwargs, + ) + + +def sample_inputs_masked_norm(op_info, device, dtype, requires_grad, **kwargs): + """Sample inputs for masked norm.""" + for ord in [2.0, 1, float("inf"), float("-inf"), 0]: + for sample_input in sample_inputs_masked_reduction( + op_info, device, dtype, requires_grad, **kwargs + ): + sample_input_args, sample_input_kwargs = ( + (ord,) + sample_input.args, + sample_input.kwargs.copy(), + ) + yield SampleInput( + sample_input.input.clone().requires_grad_(requires_grad), + args=sample_input_args, + kwargs=sample_input_kwargs, + ) + + +def reference_masked_std_var( + numpy_fn, +): + ref = reference_reduction_numpy(numpy_fn) + + # Translate unbiased or correction arguments into ddof + def func( + input, + dim=None, + unbiased=None, + *, + correction=None, + **kwargs, + ): + ddof = 1 + if unbiased is not None: + ddof = 1 if unbiased else 0 + if correction is not None: + ddof = correction + + if isinstance(dim, Sequence): + dim = tuple(dim) + + return ref(input, dim, ddof=ddof, **kwargs) + + return func + + +def sample_inputs_masked_std_var(op_info, device, dtype, requires_grad, **kwargs): + """Sample inputs for masked std/var.""" + kwargs["supports_multiple_dims"] = op_info.supports_multiple_dims + from torch.testing._internal.common_methods_invocations import sample_inputs_std_var + + def masked_samples(): + for sample_input in sample_inputs_std_var( + op_info, device, dtype, requires_grad, **kwargs + ): + if len(sample_input.args) and isinstance(sample_input.args[0], bool): + continue # masked.{std, var} doesn't support `.var(unbiased)` + + for mask in _generate_masked_op_mask( + sample_input.input.shape, device, **kwargs + ): + sample_input_args, sample_input_kwargs = ( + sample_input.args, + dict(mask=mask, **sample_input.kwargs), + ) + yield SampleInput( + sample_input.input.detach().requires_grad_(requires_grad), + args=sample_input_args, + kwargs=sample_input_kwargs, + ) + if ( + not requires_grad + and dtype.is_floating_point + and sample_input.input.ndim == 2 + and mask is not None + and mask.shape == sample_input.input.shape + ): + for v in [torch.inf, -torch.inf, torch.nan]: + t = sample_input.input.detach() + t.diagonal(0, -2, -1).fill_(v) + yield SampleInput( + t.requires_grad_(requires_grad), + args=sample_input_args, + kwargs=sample_input_kwargs, + ) + + for sample_input in masked_samples(): + correction = sample_input.kwargs.get("correction") + if correction is None: + correction = int(sample_input.kwargs.get("unbiased", True)) + + dim = sample_input.kwargs.get("dim", None) + + if sample_input.kwargs.get("mask") is None: + orig_count = torch.masked.sum( + torch.ones(sample_input.input.shape, dtype=torch.int64), + dim, + keepdim=True, + ) + else: + inmask = torch.masked._input_mask( + sample_input.input, *sample_input.args, **sample_input.kwargs + ) + orig_count = torch.masked.sum( + inmask.new_ones(sample_input.input.shape, dtype=torch.int64), + dim, + keepdim=True, + mask=inmask, + ) + if orig_count.min() <= correction + 1: + # Skip samples that lead to nans in var computation + continue + + yield sample_input + + +def sample_inputs_masked_softmax( + op_info, device, dtype, requires_grad, with_dtype=False, **kwargs +): + """Sample inputs for masked softmax, log_softmax, and softmin. + + Masked normalization operator is a reduction operator with + trailing mask optional argument. A mask is a bool tensor with the + same shape as input or a shape that is broadcastable to input + shape. + """ + for sample_input in sample_inputs_softmax_variant( + op_info, device, dtype, requires_grad, with_dtype=with_dtype, **kwargs + ): + for mask in _generate_masked_op_mask( + sample_input.input.shape, device, **kwargs + ): + yield SampleInput( + sample_input.input.clone().requires_grad_(requires_grad), + *sample_input.args, + mask=mask, + **sample_input.kwargs, + ) + + +def sample_inputs_masked_cumops(op_info, device, dtype, requires_grad, **kwargs): + """Sample inputs for masked cumsum and cumprod.""" + for sample_input in sample_inputs_softmax_variant( + op_info, device, dtype, requires_grad, **kwargs + ): + for mask in _generate_masked_op_mask( + sample_input.input.shape, device, **kwargs + ): + if type(mask) is not torch.Tensor: + continue + sample_input_args, sample_input_kwargs = ( + sample_input.args, + dict(mask=mask, **sample_input.kwargs), + ) + if "keepdim" in sample_input_kwargs: + sample_input_kwargs.pop("keepdim") + # dimension is required + if sample_input_args: + dim = sample_input.args[0] + else: + if "dim" not in sample_input_kwargs: + continue + dim = sample_input_kwargs.pop("dim") + sample_input_args = (dim,) + yield SampleInput( + sample_input.input.clone().requires_grad_(requires_grad), + *sample_input_args, + **sample_input_kwargs, + ) + + +def sample_inputs_masked_logaddexp(op_info, device, dtype, requires_grad, **kwargs): + """Sample inputs for masked logaddexp.""" + shapes = [(S,), (S, S), (S, M, S)] + input_mask_lists = [ + list(_generate_masked_op_mask(shape, device, **kwargs)) for shape in shapes + ] + other_mask_lists = [ + list(_generate_masked_op_mask(shape, device, **kwargs)) for shape in shapes + ] + + make_arg = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad + ) + for shape, input_masks, other_masks in zip( + shapes, input_mask_lists, other_mask_lists, strict=True + ): + for input_mask, other_mask in zip(input_masks, other_masks, strict=True): + yield SampleInput( + make_arg(shape), + make_arg(shape), + input_mask=input_mask, + other_mask=other_mask, + ) + + +def sample_inputs_masked_normalize(op_info, device, dtype, requires_grad, **kwargs): + """Sample inputs for masked normalize.""" + for ord in [2.0, 1, float("inf"), float("-inf"), 0]: + for sample_input in sample_inputs_softmax_variant( + op_info, device, dtype, requires_grad, use_zero_dimensions=False, **kwargs + ): + yield SampleInput( + sample_input.input.clone().requires_grad_(requires_grad), + ord, + *sample_input.args, + **sample_input.kwargs, + ) + + +op_db: list[OpInfo] = [ + ReductionOpInfo( + "masked.sum", + ref=reference_reduction_numpy(np.sum), + method_variant=None, + identity=0, + nan_policy="propagate", + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + promotes_int_to_int64=True, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + skips=( + DecorateInfo( + unittest.skip("Failing on some jobs"), + "TestReductions", + "test_reference_masked", + dtypes=(torch.bool, torch.int8, torch.int16, torch.int32), + ), + # FIXME: improve precision + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_reference_masked", + dtypes=(torch.float16,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # FIXME: sum reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, "TestReductions", "test_dim_empty"), + DecorateInfo( + unittest.expectedFailure, "TestReductions", "test_dim_empty_keepdim" + ), + # RuntimeError: undefined value tensor + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_ref_small_input", + device_type="xpu", + dtypes=[torch.complex128], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_reference_masked", + device_type="xpu", + dtypes=[torch.complex128], + ), + ), + decorators=[ + DecorateInfo( + toleranceOverride( + { + torch.bfloat16: tol(atol=1e-03, rtol=5e-2), + torch.float16: tol(atol=1e-03, rtol=5e-3), + } + ), + "TestReductions", + "test_reference_masked", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-03)}), + "TestReductions", + "test_ref_small_input", + ), + DecorateInfo( + toleranceOverride( + { + torch.bfloat16: tol(atol=0.1, rtol=0.1), + torch.float16: tol(atol=5e-3, rtol=5e-3), + } + ), + "TestMasked", + "test_mask_layout", + ), + ], + sample_inputs_func=sample_inputs_masked_reduction, + sample_inputs_sparse_coo_func=sample_inputs_sparse_coo_masked_reduction, + sample_inputs_sparse_csr_func=sample_inputs_sparse_csr_masked_reduction, + ), + ReductionOpInfo( + "masked.prod", + ref=prod_numpy, + method_variant=None, + identity=1, + nan_policy="propagate", + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse=True, + supports_sparse_csr=True, + promotes_int_to_int64=True, + dtypes=all_types_and_complex_and(torch.bool, torch.float16, torch.bfloat16), + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + DecorateInfo( + unittest.skip("Failing on some jobs"), + "TestReductions", + "test_reference_masked", + dtypes=(torch.bool, torch.int8, torch.int16, torch.int32), + ), + DecorateInfo( + "TestReductions", + "test_ref_small_input", + dtypes=(torch.int8, torch.int16, torch.int32), + ), + # FIXME: "cuda_scatter_gather_base_kernel_func" not implemented for ... (used for sparse_coo inputs) + DecorateInfo( + unittest.skip("Skipped!"), + "TestMasked", + "test_mask_layout", + device_type="cuda", + dtypes=(torch.bool, *integral_types(), *complex_types()), + ), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_reference_masked", + device_type="xpu", + dtypes=(torch.complex128, torch.int64, torch.uint8), + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_ref_small_input", + device_type="xpu", + dtypes=( + torch.int64, + torch.int32, + torch.int16, + torch.int8, + torch.complex128, + ), + ), + ), + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-03, rtol=1e-02)}), + "TestReductions", + "test_reference_masked", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-03, rtol=1e-03)}), + "TestReductions", + "test_ref_duplicate_values", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-03, rtol=1e-03)}), + "TestReductions", + "test_ref_small_input", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1.5e-03)}), + "TestMasked", + "test_mask_layout", + device_type="cpu", + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-05, rtol=1e-05)}), + "TestOperators", + "test_jvp", + device_type="cuda", + ), + ], + sample_inputs_func=sample_inputs_masked_reduction, + sample_inputs_sparse_coo_func=sample_inputs_sparse_coo_masked_reduction, + sample_inputs_sparse_csr_func=sample_inputs_sparse_csr_masked_reduction, + ), + OpInfo( + "masked.cumsum", + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + method_variant=None, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # NotSupportedError: Compiled functions can't ... use keyword-only arguments with defaults + DecorateInfo( + unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit" + ), + # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + ), + # Can reuse the same inputs; dim is required in both + sample_inputs_func=sample_inputs_masked_cumops, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), + OpInfo( + "masked.cumprod", + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + method_variant=None, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # NotSupportedError: Compiled functions can't ... use keyword-only arguments with defaults + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # NotSupportedError: Compiled functions can't ... use keyword-only arguments with defaults + DecorateInfo( + unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit" + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-5, rtol=1e-5)}), + "TestCompositeCompliance", + "test_backward", + device_type="cuda", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-2, rtol=2.6e-3)}), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda", + ), + # The following dtypes worked in forward but are not listed by the OpInfo: {torch.bool}. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + ), + # Can reuse the same inputs; dim is required in both + sample_inputs_func=sample_inputs_masked_cumops, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), + ReductionOpInfo( + "masked.amax", + nan_policy="propagate", + supports_out=False, + dtypes=all_types_and(torch.float16, torch.bfloat16), + supports_sparse=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_sparse_csr=True, + ref=reference_reduction_numpy(np.amax), + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # FIXME: amax reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, "TestReductions", "test_dim_empty"), + DecorateInfo( + unittest.expectedFailure, "TestReductions", "test_dim_empty_keepdim" + ), + # RuntimeError: Unknown builtin op: aten::iinfo + DecorateInfo( + unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit" + ), + # FIXME: "cuda_scatter_gather_base_kernel_func" not implemented for ... (used for sparse_coo inputs) + # FIXME: "_segment_reduce_lengths_cpu/cuda" not implemented for ... (used for sparse_csr inputs) + DecorateInfo( + unittest.skip("Skipped!"), + "TestMasked", + "test_mask_layout", + dtypes=(torch.bool, *integral_types(), *complex_types()), + ), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_reference_masked", + device_type="xpu", + dtypes=[torch.int64], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_ref_small_input", + device_type="xpu", + dtypes=[torch.int64], + ), + ), + sample_inputs_func=sample_inputs_masked_reduction, + sample_inputs_sparse_coo_func=sample_inputs_sparse_coo_masked_reduction, + sample_inputs_sparse_csr_func=sample_inputs_sparse_csr_masked_reduction, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), + ReductionOpInfo( + "masked.amin", + nan_policy="propagate", + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + dtypes=all_types_and(torch.float16, torch.bfloat16), + supports_sparse=True, + supports_sparse_csr=True, + ref=reference_reduction_numpy(np.amin), + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # FIXME: amax reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, "TestReductions", "test_dim_empty"), + DecorateInfo( + unittest.expectedFailure, "TestReductions", "test_dim_empty_keepdim" + ), + # RuntimeError: Unknown builtin op: aten::iinfo + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # FIXME: "cuda_scatter_gather_base_kernel_func" not implemented for ... (used for sparse_coo inputs) + # FIXME: "_segment_reduce_lengths_cpu/cuda" not implemented for ... (used for sparse_csr inputs) + DecorateInfo( + unittest.skip("Skipped!"), + "TestMasked", + "test_mask_layout", + dtypes=(torch.bool, *integral_types(), *complex_types()), + ), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_reference_masked", + device_type="xpu", + dtypes=[torch.int64], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_ref_small_input", + device_type="xpu", + dtypes=[torch.int64], + ), + ), + sample_inputs_func=sample_inputs_masked_reduction, + sample_inputs_sparse_coo_func=sample_inputs_sparse_coo_masked_reduction, + sample_inputs_sparse_csr_func=sample_inputs_sparse_csr_masked_reduction, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), + ReductionOpInfo( + "masked.argmax", + supports_out=False, + supports_multiple_dims=False, + supports_autograd=False, + dtypes=all_types_and(torch.float16, torch.bfloat16), + ref=reference_reduction_numpy(np.argmax, supports_keepdims=False), + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # initial is not a keyword for argmax + DecorateInfo( + unittest.expectedFailure, "TestReductions", "test_reference_masked" + ), + # NotSupportedError: Compiled functions can't ... use keyword-only arguments with defaults + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_ref_small_input", + device_type="xpu", + dtypes=floating_types_and( + torch.int8, torch.int16, torch.int32, torch.int64, torch.float16 + ), + ), + ), + sample_inputs_func=sample_inputs_masked_reduction, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), + ReductionOpInfo( + "masked.argmin", + supports_out=False, + supports_multiple_dims=False, + supports_autograd=False, + dtypes=all_types_and(torch.float16, torch.bfloat16), + ref=reference_reduction_numpy(np.argmin, supports_keepdims=False), + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # initial is not a keyword for argmin + DecorateInfo( + unittest.expectedFailure, "TestReductions", "test_reference_masked" + ), + # NotSupportedError: Compiled functions can't ... use keyword-only arguments with defaults + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_ref_small_input", + device_type="xpu", + dtypes=floating_types_and( + torch.int8, torch.int16, torch.int32, torch.int64, torch.float16 + ), + ), + ), + sample_inputs_func=sample_inputs_masked_reduction, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), + ReductionOpInfo( + "masked.mean", + ref=reference_reduction_numpy(np.mean) + if np.lib.NumpyVersion(np.__version__) >= "1.20.2" + else None, + method_variant=None, + nan_policy="propagate", + supports_out=False, + supports_sparse_csr=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + promotes_int_to_float=True, + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # FIXME: sum reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, "TestReductions", "test_dim_empty"), + DecorateInfo( + unittest.expectedFailure, "TestReductions", "test_dim_empty_keepdim" + ), + # RuntimeError: undefined value tensor + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # FIXME: "_segment_reduce_lengths_cpu/cuda" not implemented for ... (used for sparse_csr inputs) + DecorateInfo( + unittest.skip("Skipped!"), + "TestMasked", + "test_mask_layout", + dtypes=(torch.bool, *integral_types(), *complex_types()), + ), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_ref_small_input", + device_type="xpu", + dtypes=[torch.complex128], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_reference_masked", + device_type="xpu", + dtypes=[torch.complex128], + ), + ), + decorators=[ + DecorateInfo( + toleranceOverride( + { + torch.bfloat16: tol(atol=1e-03, rtol=0.05), + torch.float16: tol(atol=1e-03, rtol=1e-03), + } + ), + "TestReductions", + "test_reference_masked", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-03, rtol=1e-03)}), + "TestReductions", + "test_ref_small_input", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-03, rtol=2e-03)}), + "TestSparseCompressed", + "test_consistency", + device_type="cuda", + ), + ], + sample_inputs_func=sample_inputs_masked_reduction, + sample_inputs_sparse_csr_func=sample_inputs_sparse_csr_masked_reduction, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), + OpInfo( + "masked.median", + dtypes=floating_types_and(torch.bfloat16, torch.float16), + method_variant=None, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # NotSupportedError: Compiled functions can't ... use keyword-only arguments with defaults + DecorateInfo( + unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit" + ), + ), + sample_inputs_func=partial( + sample_inputs_masked_softmax, use_zero_dimensions=False + ), + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), + ReductionOpInfo( + "masked.norm", + identity=0, + method_variant=None, + nan_policy="propagate", + supports_out=False, + promotes_int_to_float=True, + dtypes=floating_types_and(torch.float16, torch.bfloat16), + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # FIXME: sum reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, "TestReductions", "test_dim_empty"), + DecorateInfo( + unittest.expectedFailure, "TestReductions", "test_dim_empty_keepdim" + ), + # torch.jit.frontend.NotSupportedError: Compiled functions + # can't take variable number of arguments or use + # keyword-only arguments with defaults + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + ), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_masked_norm, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), + ReductionOpInfo( + "masked.var", + ref=reference_masked_std_var(np.var) + if np.lib.NumpyVersion(np.__version__) >= "1.20.2" + else None, + method_variant=None, + nan_policy="propagate", + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + promotes_int_to_float=True, + dtypes=all_types_and_complex_and(torch.float16, torch.bfloat16), + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + "TestSchemaCheckModeOpInfo", + "test_schema_correctness", + dtypes=(torch.complex64, torch.complex128), + ), + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # FIXME: sum reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, "TestReductions", "test_dim_empty"), + DecorateInfo( + unittest.expectedFailure, "TestReductions", "test_dim_empty_keepdim" + ), + # RuntimeError: undefined value tensor + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_ref_small_input", + device_type="xpu", + dtypes=[torch.complex128], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_reference_masked", + device_type="xpu", + dtypes=[torch.complex128], + ), + ), + decorators=[ + DecorateInfo( + toleranceOverride( + { + torch.float16: tol(atol=1e-02, rtol=1e-02), + torch.bfloat16: tol(atol=1e-03, rtol=1e-03), + } + ), + "TestReductions", + "test_reference_masked", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-02)}), + "TestReductions", + "test_ref_small_input", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-02)}), + "TestMasked", + "test_reference_masked", + ), + DecorateInfo( + toleranceOverride( + { + torch.float16: tol(atol=1e-02, rtol=1e-02), + torch.bfloat16: tol(atol=1e-03, rtol=1e-03), + } + ), + "TestMasked", + "test_reference_masked", + ), + DecorateInfo( + toleranceOverride( + { + torch.float16: tol(atol=4e-5, rtol=2e-2), + } + ), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda", + ), + ], + sample_inputs_func=sample_inputs_masked_std_var, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + check_batched_grad=True, + ), + ReductionOpInfo( + "masked.std", + ref=reference_masked_std_var(np.std) + if np.lib.NumpyVersion(np.__version__) >= "1.20.2" + else None, + method_variant=None, + nan_policy="propagate", + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + promotes_int_to_float=True, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + "TestSchemaCheckModeOpInfo", + "test_schema_correctness", + dtypes=(torch.complex64, torch.complex128), + ), + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # FIXME: sum reduces all dimensions when dim=[] + DecorateInfo(unittest.expectedFailure, "TestReductions", "test_dim_empty"), + DecorateInfo( + unittest.expectedFailure, "TestReductions", "test_dim_empty_keepdim" + ), + # RuntimeError: undefined value tensor + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # Driver issue of XPU, see https://github.com/intel/torch-xpu-ops/issues/2295 + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_ref_small_input", + device_type="xpu", + dtypes=[torch.complex128], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_reference_masked", + device_type="xpu", + dtypes=[torch.complex128], + ), + ), + decorators=[ + DecorateInfo( + toleranceOverride( + { + torch.bfloat16: tol(atol=1e-02, rtol=1e-02), + torch.float16: tol(atol=1e-02, rtol=1e-02), + } + ), + "TestReductions", + "test_reference_masked", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-02)}), + "TestReductions", + "test_ref_small_input", + ), + DecorateInfo( + toleranceOverride( + { + torch.float16: tol(atol=1e-02, rtol=1e-02), + torch.bfloat16: tol(atol=5e-03, rtol=5e-04), + } + ), + "TestMasked", + "test_reference_masked", + ), + ], + sample_inputs_func=sample_inputs_masked_std_var, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + check_batched_grad=True, + ), + OpInfo( + "masked.softmax", + method_variant=None, + dtypes=floating_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_masked_softmax, + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + ), + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + ), + OpInfo( + "masked.log_softmax", + method_variant=None, + dtypes=floating_types_and(torch.half, torch.bfloat16), + dtypesIfMPS=floating_types_and( + torch.half, + torch.bfloat16, + torch.uint8, + torch.int32, + torch.int16, + torch.int8, + ), + sample_inputs_func=sample_inputs_masked_softmax, + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # NotImplementedError: "log_softmax_lastdim_kernel_impl" not implemented for * + DecorateInfo( + unittest.expectedFailure, + "TestConsistency", + device_type="mps", + dtypes=( + torch.uint8, + torch.int32, + torch.int16, + torch.int8, + ), + ), + ), + decorators=[ + DecorateInfo( + toleranceOverride({torch.bfloat16: tol(atol=1e-02, rtol=1e-02)}), + "TestMasked", + "test_reference_masked", + ), + ], + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + ), + OpInfo( + "masked.softmin", + method_variant=None, + dtypes=floating_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_masked_softmax, + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # FIXME: + # Mismatched elements: 2 / 2 (100.0%) + # Greatest absolute difference: nan at index (0,) (up to 0.0001 allowed) + # Greatest relative difference: nan at index (0,) (up to 0.0001 allowed + DecorateInfo( + unittest.skip("Skipped!"), + "TestOperators", + "test_vmapvjpvjp", + device_type="cpu", + ), + ), + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + ), + OpInfo( + "masked.normalize", + method_variant=None, + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_masked_normalize, + decorators=[ + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=2e-5, rtol=6e-3)}), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda", + ), + ], + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + ), + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + ), + OpInfo( + "masked.logaddexp", + dtypes=floating_types_and(torch.float16, torch.bfloat16), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # NotSupportedError: Compiled functions can't ... use keyword-only arguments with defaults + DecorateInfo( + unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit" + ), + DecorateInfo( + unittest.skip("Skipped!"), "TestFwdGradients", "test_fn_gradgrad" + ), + DecorateInfo( + unittest.skip("Skipped!"), "TestBwdGradients", "test_fn_gradgrad" + ), + ), + sample_inputs_func=sample_inputs_masked_logaddexp, + gradcheck_wrapper=gradcheck_wrapper_masked_pointwise_operation, + ), + ReductionOpInfo( + "masked.logsumexp", + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + method_variant=None, + nan_policy="propagate", + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + # FIXME: reduces all dimensions when dim=[] + DecorateInfo(unittest.skip("Skipped!"), "TestReductions", "test_dim_empty"), + DecorateInfo( + unittest.skip("Skipped!"), "TestReductions", "test_dim_empty_keepdim" + ), + # Identity can't be -torch.inf without overflow + DecorateInfo( + unittest.skip("Skipped!"), + "TestReductions", + "test_empty_tensor_empty_slice", + ), + # NotSupportedError: Compiled functions can't ... use keyword-only arguments with defaults + DecorateInfo( + unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit" + ), + # all the values are the same except for -inf vs nan + DecorateInfo(unittest.skip("Skipped!"), "TestDecomp", "test_comprehensive"), + # FIXME: + # Mismatched elements: 2 / 12 (16.7%) + # Greatest absolute difference: 9223372034707292160 at index (0, 0, 0, 0) + # Greatest relative difference: 0.0 at index (0, 0, 0, 1) + DecorateInfo( + unittest.skip("Skipped!"), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cpu", + ), + ), + sample_inputs_func=sample_inputs_masked_reduction, + gradcheck_wrapper=gradcheck_wrapper_masked_operation, + ), +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/fft.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/fft.py new file mode 100644 index 0000000000000000000000000000000000000000..cb7d8a6507db9cc373aceff7cdc5cf00b992e5cc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/fft.py @@ -0,0 +1,939 @@ +# mypy: ignore-errors + +import unittest +from functools import partial + +import numpy as np + +import torch +from torch.testing import make_tensor +from torch.testing._internal.common_device_type import precisionOverride +from torch.testing._internal.common_dtype import ( + all_types_and, + all_types_and_complex_and, +) +from torch.testing._internal.common_utils import TEST_SCIPY, TEST_WITH_ROCM +from torch.testing._internal.opinfo.core import ( + DecorateInfo, + ErrorInput, + OpInfo, + sample_inputs_spectral_ops, + SampleInput, + SpectralFuncInfo, + SpectralFuncType, +) +from torch.testing._internal.opinfo.refs import ( + _find_referenced_opinfo, + _inherit_constructor_args, + PythonRefInfo, +) + + +has_scipy_fft = False +if TEST_SCIPY: + try: + import scipy.fft + + has_scipy_fft = True + except ModuleNotFoundError: + pass + + +class SpectralFuncPythonRefInfo(SpectralFuncInfo): + """ + An OpInfo for a Python reference of an elementwise unary operation. + """ + + def __init__( + self, + name, # the stringname of the callable Python reference + *, + op=None, # the function variant of the operation, populated as torch. if None + torch_opinfo_name, # the string name of the corresponding torch opinfo + torch_opinfo_variant="", + **kwargs, + ): # additional kwargs override kwargs inherited from the torch opinfo + self.torch_opinfo_name = torch_opinfo_name + self.torch_opinfo = _find_referenced_opinfo( + torch_opinfo_name, torch_opinfo_variant, op_db=op_db + ) + if not isinstance(self.torch_opinfo, SpectralFuncInfo): + raise AssertionError( + f"Expected torch_opinfo to be SpectralFuncInfo, got {type(self.torch_opinfo)}" + ) + + inherited = self.torch_opinfo._original_spectral_func_args + ukwargs = _inherit_constructor_args(name, op, inherited, kwargs) + + super().__init__(**ukwargs) + + +def error_inputs_fft(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + # Zero-dimensional tensor has no dimension to take FFT of + yield ErrorInput( + SampleInput(make_arg()), + error_type=IndexError, + error_regex="Dimension specified as -1 but tensor has no dimensions", + ) + + +def error_inputs_fftn(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + # Specifying a dimension on a zero-dimensional tensor + yield ErrorInput( + SampleInput(make_arg(), dim=(0,)), + error_type=IndexError, + error_regex="Dimension specified as 0 but tensor has no dimensions", + ) + + +def sample_inputs_fft_with_min( + op_info, device, dtype, requires_grad=False, *, min_size, **kwargs +): + yield from sample_inputs_spectral_ops( + op_info, device, dtype, requires_grad, **kwargs + ) + if TEST_WITH_ROCM: + # FIXME: Causes floating point exception on ROCm + return + + # Check the "Invalid number of data points" error isn't too strict + # https://github.com/pytorch/pytorch/pull/109083 + a = make_tensor(min_size, dtype=dtype, device=device, requires_grad=requires_grad) + yield SampleInput(a) + + +def sample_inputs_fftshift(op_info, device, dtype, requires_grad, **kwargs): + def mt(shape, **kwargs): + return make_tensor( + shape, device=device, dtype=dtype, requires_grad=requires_grad, **kwargs + ) + + yield SampleInput(mt((9, 10))) + yield SampleInput(mt((50,)), kwargs=dict(dim=0)) + yield SampleInput(mt((5, 11)), kwargs=dict(dim=(1,))) + yield SampleInput(mt((5, 6)), kwargs=dict(dim=(0, 1))) + yield SampleInput(mt((5, 6, 2)), kwargs=dict(dim=(0, 2))) + + +# Operator database +op_db: list[OpInfo] = [ + SpectralFuncInfo( + "fft.fft", + aten_name="fft_fft", + decomp_aten_name="_fft_c2c", + ref=np.fft.fft, + ndimensional=SpectralFuncType.OneD, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32), + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=1), + error_inputs_func=error_inputs_fft, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + ), + SpectralFuncInfo( + "fft.fft2", + aten_name="fft_fft2", + ref=np.fft.fft2, + decomp_aten_name="_fft_c2c", + ndimensional=SpectralFuncType.TwoD, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32), + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + decorators=[precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4})], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_complex_half_reference_testing", + device_type="cuda", + dtypes=[torch.complex32], + active_if=TEST_WITH_ROCM, + ), + # RuntimeError: [srcBuf length] > 0 INTERNAL ASSERT FAILED + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out", + device_type="mps", + ), + # AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 3, 10]) != torch.Size([5, 3, 11]). + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + ), + ), + SpectralFuncInfo( + "fft.fftn", + aten_name="fft_fftn", + decomp_aten_name="_fft_c2c", + ref=np.fft.fftn, + ndimensional=SpectralFuncType.ND, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and( + torch.bool, + torch.half, + torch.complex32, + ), + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + decorators=[precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4})], + skips=( + # RuntimeError: [srcBuf length] > 0 INTERNAL ASSERT FAILED + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out", + device_type="mps", + ), + # AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 3, 10]) != torch.Size([5, 3, 11]). + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + ), + ), + SpectralFuncInfo( + "fft.hfft", + aten_name="fft_hfft", + decomp_aten_name="_fft_c2r", + ref=np.fft.hfft, + ndimensional=SpectralFuncType.OneD, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and( + torch.bool, + torch.half, + torch.complex32, + ), + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=2), + error_inputs_func=error_inputs_fft, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + check_batched_gradgrad=False, + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + "TestSchemaCheckModeOpInfo", + "test_schema_correctness", + dtypes=(torch.complex64, torch.complex128), + ), + ), + ), + SpectralFuncInfo( + "fft.hfft2", + aten_name="fft_hfft2", + decomp_aten_name="_fft_c2r", + ref=scipy.fft.hfft2 if has_scipy_fft else None, + ndimensional=SpectralFuncType.TwoD, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and( + torch.bool, + torch.half, + torch.complex32, + ), + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(2, 2)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_gradgrad=False, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 2e-4, torch.cfloat: 2e-4}), + "TestFFT", + "test_reference_nd", + ), + ], + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + "TestSchemaCheckModeOpInfo", + "test_schema_correctness", + ), + # FIXME: errors are too large; needs investigation + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_complex_half_reference_testing", + device_type="cuda", + ), + ), + ), + SpectralFuncInfo( + "fft.hfftn", + aten_name="fft_hfftn", + decomp_aten_name="_fft_c2r", + ref=scipy.fft.hfftn if has_scipy_fft else None, + ndimensional=SpectralFuncType.ND, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32), + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(2, 2)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_gradgrad=False, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 2e-4, torch.cfloat: 2e-4}), + "TestFFT", + "test_reference_nd", + ), + ], + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + "TestSchemaCheckModeOpInfo", + "test_schema_correctness", + ), + ), + ), + SpectralFuncInfo( + "fft.rfft", + aten_name="fft_rfft", + decomp_aten_name="_fft_r2c", + ref=np.fft.rfft, + ndimensional=SpectralFuncType.OneD, + dtypes=all_types_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and(torch.bool, torch.half), + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=1), + error_inputs_func=error_inputs_fft, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_grad=False, + check_batched_gradgrad=False, + ), + SpectralFuncInfo( + "fft.rfft2", + aten_name="fft_rfft2", + decomp_aten_name="_fft_r2c", + ref=np.fft.rfft2, + ndimensional=SpectralFuncType.TwoD, + dtypes=all_types_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and(torch.bool, torch.half), + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_grad=False, + check_batched_gradgrad=False, + decorators=[ + precisionOverride({torch.float: 1e-4}), + ], + ), + SpectralFuncInfo( + "fft.rfftn", + aten_name="fft_rfftn", + decomp_aten_name="_fft_r2c", + ref=np.fft.rfftn, + ndimensional=SpectralFuncType.ND, + dtypes=all_types_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and(torch.bool, torch.half), + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_grad=False, + check_batched_gradgrad=False, + decorators=[ + precisionOverride({torch.float: 1e-4}), + ], + ), + SpectralFuncInfo( + "fft.ifft", + aten_name="fft_ifft", + decomp_aten_name="_fft_c2c", + ref=np.fft.ifft, + ndimensional=SpectralFuncType.OneD, + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=1), + error_inputs_func=error_inputs_fft, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32), + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out", + device_type="mps", + ), + ), + ), + SpectralFuncInfo( + "fft.ifft2", + aten_name="fft_ifft2", + decomp_aten_name="_fft_c2c", + ref=np.fft.ifft2, + ndimensional=SpectralFuncType.TwoD, + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32), + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + skips=( + # RuntimeError: [srcBuf length] > 0 INTERNAL ASSERT FAILED + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out", + device_type="mps", + ), + # AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 3, 10]) != torch.Size([5, 3, 11]). + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + ), + ), + SpectralFuncInfo( + "fft.ifftn", + aten_name="fft_ifftn", + decomp_aten_name="_fft_c2c", + ref=np.fft.ifftn, + ndimensional=SpectralFuncType.ND, + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and( + torch.bool, + torch.half, + torch.complex32, + ), + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + skips=( + # RuntimeError: [srcBuf length] > 0 INTERNAL ASSERT FAILED + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + # AssertionError: The values for attribute 'shape' do not match: torch.Size([5, 3, 10]) != torch.Size([5, 3, 11]). + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + ), + ), + SpectralFuncInfo( + "fft.ihfft", + aten_name="fft_ihfft", + decomp_aten_name="_fft_r2c", + ref=np.fft.ihfft, + ndimensional=SpectralFuncType.OneD, + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)), + error_inputs_func=error_inputs_fft, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + dtypes=all_types_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and(torch.bool, torch.half), + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + check_batched_grad=False, + ), + SpectralFuncInfo( + "fft.ihfft2", + aten_name="fft_ihfft2", + decomp_aten_name="_fft_r2c", + ref=scipy.fft.ihfftn if has_scipy_fft else None, + ndimensional=SpectralFuncType.TwoD, + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + dtypes=all_types_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and(torch.bool, torch.half), + check_batched_grad=False, + check_batched_gradgrad=False, + decorators=( + # The values for attribute 'shape' do not match: torch.Size([5, 6, 5]) != torch.Size([5, 6, 6]). + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out_warning"), + DecorateInfo( + precisionOverride({torch.float: 2e-4}), "TestFFT", "test_reference_nd" + ), + # Mismatched elements! + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out"), + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out_warnings"), + ), + ), + SpectralFuncInfo( + "fft.ihfftn", + aten_name="fft_ihfftn", + decomp_aten_name="_fft_r2c", + ref=scipy.fft.ihfftn if has_scipy_fft else None, + ndimensional=SpectralFuncType.ND, + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 1)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + dtypes=all_types_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archss + dtypesIfCUDA=all_types_and(torch.bool, torch.half), + check_batched_grad=False, + check_batched_gradgrad=False, + decorators=[ + # The values for attribute 'shape' do not match: torch.Size([5, 6, 5]) != torch.Size([5, 6, 6]). + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out_warning"), + # Mismatched elements! + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_out"), + DecorateInfo( + precisionOverride({torch.float: 2e-4}), "TestFFT", "test_reference_nd" + ), + ], + ), + SpectralFuncInfo( + "fft.irfft", + aten_name="fft_irfft", + decomp_aten_name="_fft_c2r", + ref=np.fft.irfft, + ndimensional=SpectralFuncType.OneD, + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 2)), + error_inputs_func=error_inputs_fft, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32), + check_batched_gradgrad=False, + ), + SpectralFuncInfo( + "fft.irfft2", + aten_name="fft_irfft2", + decomp_aten_name="_fft_c2r", + ref=np.fft.irfft2, + ndimensional=SpectralFuncType.TwoD, + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 2)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32), + check_batched_gradgrad=False, + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + ), + SpectralFuncInfo( + "fft.irfftn", + aten_name="fft_irfftn", + decomp_aten_name="_fft_c2r", + ref=np.fft.irfftn, + ndimensional=SpectralFuncType.ND, + sample_inputs_func=partial(sample_inputs_fft_with_min, min_size=(1, 2)), + error_inputs_func=error_inputs_fftn, + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + dtypes=all_types_and_complex_and(torch.bool), + # CUDA supports Half/ComplexHalf Precision FFT only on SM53 or later archs + dtypesIfCUDA=all_types_and_complex_and(torch.bool, torch.half, torch.complex32), + check_batched_gradgrad=False, + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + ), + OpInfo( + "fft.fftshift", + dtypes=all_types_and_complex_and( + torch.bool, torch.bfloat16, torch.half, torch.chalf + ), + sample_inputs_func=sample_inputs_fftshift, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ), + OpInfo( + "fft.ifftshift", + dtypes=all_types_and_complex_and( + torch.bool, torch.bfloat16, torch.half, torch.chalf + ), + sample_inputs_func=sample_inputs_fftshift, + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ), +] + +python_ref_db: list[OpInfo] = [ + SpectralFuncPythonRefInfo( + "_refs.fft.fft", + torch_opinfo_name="fft.fft", + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.ifft", + torch_opinfo_name="fft.ifft", + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.rfft", + torch_opinfo_name="fft.rfft", + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.irfft", + torch_opinfo_name="fft.irfft", + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.hfft", + torch_opinfo_name="fft.hfft", + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.ihfft", + torch_opinfo_name="fft.ihfft", + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.fftn", + torch_opinfo_name="fft.fftn", + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + ), + SpectralFuncPythonRefInfo( + "_refs.fft.ifftn", + torch_opinfo_name="fft.ifftn", + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + ), + SpectralFuncPythonRefInfo( + "_refs.fft.rfftn", + torch_opinfo_name="fft.rfftn", + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out", + device_type="mps", + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.irfftn", + torch_opinfo_name="fft.irfftn", + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.hfftn", + torch_opinfo_name="fft.hfftn", + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 2e-4, torch.cfloat: 2e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.ihfftn", + torch_opinfo_name="fft.ihfftn", + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 2e-4}), + "TestFFT", + "test_reference_nd", + ), + # AssertionError: Reference result was farther (0.09746177145360499) from the precise + # computation than the torch result was (0.09111555632069855) + # See https://github.com/pytorch/pytorch/pull/170856 for more details. + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_python_ref_torch_fallback", + dtypes=(torch.float16,), + device_type="cuda", + ), + # AssertionError: Reference result was farther (0.0953431016138116) from the precise + # computation than the torch result was (0.09305490684430734) + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_python_ref_executor", + dtypes=(torch.float16,), + device_type="cuda", + ), + ], + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.fft2", + torch_opinfo_name="fft.fft2", + ), + SpectralFuncPythonRefInfo( + "_refs.fft.ifft2", + torch_opinfo_name="fft.ifft2", + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + ), + SpectralFuncPythonRefInfo( + "_refs.fft.rfft2", + torch_opinfo_name="fft.rfft2", + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.irfft2", + torch_opinfo_name="fft.irfft2", + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 1e-4, torch.cfloat: 1e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.hfft2", + torch_opinfo_name="fft.hfft2", + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 2e-4, torch.cfloat: 2e-4}), + "TestFFT", + "test_reference_nd", + ) + ], + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + SpectralFuncPythonRefInfo( + "_refs.fft.ihfft2", + torch_opinfo_name="fft.ihfft2", + decorators=[ + DecorateInfo( + precisionOverride({torch.float: 2e-4}), + "TestFFT", + "test_reference_nd", + ), + # FIXME: + # Reference result was farther (0.0953431016138116) from the precise computation + # than the torch result was (0.09305490684430734)! + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_python_ref_executor", + device_type="cuda", + ), + ], + skips=( + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + ), + ), + PythonRefInfo( + "_refs.fft.fftshift", + op_db=op_db, + torch_opinfo_name="fft.fftshift", + ), + PythonRefInfo( + "_refs.fft.ifftshift", + op_db=op_db, + torch_opinfo_name="fft.ifftshift", + ), +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/linalg.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/linalg.py new file mode 100644 index 0000000000000000000000000000000000000000..b96a66af8db95321cbfc0cebf235c548621f950b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/linalg.py @@ -0,0 +1,2853 @@ +# mypy: ignore-errors + +import itertools +import random +import unittest +from collections.abc import Iterable +from functools import partial +from itertools import chain, product + +import numpy as np +from numpy import inf + +import torch +from torch.testing import make_tensor +from torch.testing._internal.common_cuda import _get_magma_version, with_tf32_off +from torch.testing._internal.common_device_type import ( + has_cusolver, + skipCPUIfNoLapack, + skipCUDAIfNoCusolver, + skipCUDAIfNoMagma, + skipCUDAIfNoMagmaAndNoCusolver, + skipCUDAIfNoMagmaAndNoLinalgsolver, + skipCUDAIfRocm, + tol, + toleranceOverride, +) +from torch.testing._internal.common_dtype import ( + all_types_and_complex, + all_types_and_complex_and, + floating_and_complex_types, + floating_and_complex_types_and, + floating_types, +) +from torch.testing._internal.common_utils import ( + GRADCHECK_NONDET_TOL, + IS_ARM64, + IS_CPU_CAPABILITY_SVE256, + IS_LINUX, + make_fullrank_matrices_with_distinct_singular_values, + skipIfSlowGradcheckEnv, + slowTest, + TEST_WITH_ROCM, + TEST_XPU, +) +from torch.testing._internal.opinfo.core import ( + clone_sample, + DecorateInfo, + ErrorInput, + gradcheck_wrapper_hermitian_input, + L, + M, + OpInfo, + ReductionOpInfo, + S, + SampleInput, +) +from torch.testing._internal.opinfo.refs import PythonRefInfo, ReductionPythonRefInfo + + +def sample_kwargs_vector_norm(t, **kwargs): + # orders with / without identity + def ords(): + has_id = (6, 4, 2, 1, 0, 0.9) + no_id = (inf, -2.1, -inf) + if t.numel() == 0: + dim = kwargs.get("dim") + if dim is None: + return has_id + if not isinstance(dim, Iterable): + dim = (dim,) + for d in dim: + if t.size(d) == 0: + return has_id + return has_id + no_id + + return (((), dict(ord=o)) for o in ords()) + + +def sample_inputs_svd(op_info, device, dtype, requires_grad=False, **kwargs): + make_fullrank = make_fullrank_matrices_with_distinct_singular_values + make_arg = partial( + make_fullrank, dtype=dtype, device=device, requires_grad=requires_grad + ) + + is_linalg_svd = "linalg.svd" in op_info.name + batches = [(), (0,), (3,)] + ns = [0, 3, 5] + + def uniformize(usv): + S = usv[1] + k = S.shape[-1] + U = usv[0][..., :k] + Vh = usv[2] if is_linalg_svd else usv[2].mH + Vh = Vh[..., :k, :] + return U, S, Vh + + def fn_U(usv): + U, _, _ = uniformize(usv) + return U.abs() + + def fn_S(usv): + return uniformize(usv)[1] + + def fn_Vh(usv): + # We also return S to test + _, S, Vh = uniformize(usv) + return S, Vh.abs() + + def fn_UVh(usv): + U, S, Vh = uniformize(usv) + return U @ Vh, S + + fns = (fn_U, fn_S, fn_Vh, fn_UVh) + + fullmat = "full_matrices" if is_linalg_svd else "some" + + for batch, n, k, fullmat_val, fn in product(batches, ns, ns, (True, False), fns): + shape = batch + (n, k) + yield SampleInput( + make_arg(*shape), kwargs={fullmat: fullmat_val}, output_process_fn_grad=fn + ) + + +def sample_inputs_cross(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad + ) + yield SampleInput(make_arg((S, 3)), args=(make_arg((S, 3)),)) + yield SampleInput( + make_arg((S, 3, S)), args=(make_arg((S, 3, S)),), kwargs=dict(dim=1) + ) + yield SampleInput(make_arg((1, 3)), args=(make_arg((S, 3)),), kwargs=dict(dim=-1)) + + +def error_inputs_cross(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + sample = SampleInput(input=make_arg((S, 3)), args=(make_arg((S, 1)),)) + err = "inputs dimension -1 must have length 3" + yield ErrorInput(sample, error_regex=err, error_type=RuntimeError) + + sample = SampleInput(input=make_arg((5, S, 3)), args=(make_arg((S, 3)),)) + err = "inputs must have the same number of dimensions" + yield ErrorInput(sample, error_regex=err, error_type=RuntimeError) + + sample = SampleInput(input=make_arg((S, 2)), args=(make_arg((S, 2)),)) + err = "must have length 3" + yield ErrorInput(sample, error_regex=err, error_type=RuntimeError) + + sample = SampleInput( + input=make_arg((S, 2)), args=(make_arg((S, 2)),), kwargs=dict(dim=2) + ) + err = "Dimension out of range" + yield ErrorInput(sample, error_regex=err, error_type=IndexError) + + +def sample_inputs_householder_product(op_info, device, dtype, requires_grad, **kwargs): + """ + This function generates input for torch.linalg.householder_product (torch.orgqr). + The first argument should be a square matrix or batch of square matrices, the second argument is a vector or batch of vectors. + Empty, square, rectangular, batched square and batched rectangular input is generated. + """ + make_arg = partial( + make_tensor, + device=device, + dtype=dtype, + requires_grad=requires_grad, + low=-2, + high=2, + ) + # Each column of the matrix is getting multiplied many times leading to very large values for + # the Jacobian matrix entries and making the finite-difference result of grad check less accurate. + # That's why gradcheck with the default range [-9, 9] fails and [-2, 2] is used here. + yield SampleInput(make_arg((S, S)), make_arg((S,))) + yield SampleInput(make_arg((S + 1, S)), make_arg((S,))) + yield SampleInput(make_arg((2, 1, S, S)), make_arg((2, 1, S))) + yield SampleInput(make_arg((2, 1, S + 1, S)), make_arg((2, 1, S))) + yield SampleInput( + make_arg((0, 0), low=None, high=None), + make_arg((0,), low=None, high=None), + ) + yield SampleInput(make_arg((S, S)), make_arg((0,), low=None, high=None)) + # m = n = S, k = S - 2 + yield SampleInput(make_arg((S, S)), make_arg((S - 2,), low=None, high=None)) + # m = S, n = S -1, k = S - 2 + yield SampleInput(make_arg((S, S - 1)), make_arg((S - 2,), low=None, high=None)) + + +def sample_inputs_linalg_matrix_power(op_info, device, dtype, requires_grad, **kwargs): + make_fullrank = make_fullrank_matrices_with_distinct_singular_values + make_arg = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad + ) + make_arg_fullrank = partial( + make_fullrank, dtype=dtype, device=device, requires_grad=requires_grad + ) + # (, ()) + test_sizes = [ + (1, ()), + (2, (0,)), + (2, (2,)), + ] + + for matrix_size, batch_sizes in test_sizes: + size = batch_sizes + (matrix_size, matrix_size) + for n in (0, 3, 5): + yield SampleInput(make_arg(size), args=(n,)) + for n in [-4, -2, -1]: + yield SampleInput(make_arg_fullrank(*size), args=(n,)) + + +def sample_inputs_linalg_det_logdet_slogdet( + op_info, device, dtype, requires_grad, **kwargs +): + make_fullrank = make_fullrank_matrices_with_distinct_singular_values + make_arg = partial( + make_fullrank, dtype=dtype, device=device, requires_grad=requires_grad + ) + batches = [(), (0,), (3,)] + ns = [0, 1, 5] + + is_logdet = op_info.name == "logdet" + + for ( + batch, + n, + ) in product(batches, ns): + shape = batch + (n, n) + A = make_arg(*shape) + # Need to make the matrices in A have positive determinant for autograd + # To do so, we multiply A by its determinant to flip the sign of its determinant + if is_logdet and not A.is_complex() and A.numel() > 0: + s = torch.linalg.slogdet(A).sign + A = A * s.unsqueeze(-1).unsqueeze(-1) + A.requires_grad_(requires_grad) + yield SampleInput(A) + + +def sample_inputs_lu_solve(op_info, device, dtype, requires_grad=False, **kwargs): + """Samples the inputs for both linalg.lu_solve and lu_solve""" + make_fn = make_fullrank_matrices_with_distinct_singular_values + make_a = partial(make_fn, dtype=dtype, device=device) + make_b = partial(make_tensor, dtype=dtype, device=device) + + def clone(X, requires_grad): + Y = X.clone() + Y.requires_grad_(requires_grad) + return Y + + is_linalg_lu_solve = op_info.name == "linalg.lu_solve" + + batches = ((), (0,), (2,)) + ns = (3, 1, 0) + nrhs = (4, 1, 0) + + for n, batch, rhs in product(ns, batches, nrhs): + A = make_a(*(batch + (n, n))) + if torch.device(device).type == "mps": + # TODO: Fix lu_factor for MPS, because it does not work for all of + # these cases. So we resort to the CPU impl here and move the + # outputs back to MPS. + LU, pivots = (x.to(device) for x in torch.linalg.lu_factor(A.cpu())) + else: + LU, pivots = torch.linalg.lu_factor(A) + + B = make_b(batch + (n, rhs)) + + grads = (False,) if not requires_grad else (True, False) + # we try all possible combinations of requires_grad for each input + for LU_grad, B_grad in product(grads, grads): + # when requires_grad == True, at least one input has to have requires_grad enabled + if requires_grad and not LU_grad and not B_grad: + continue + + if is_linalg_lu_solve: + for adjoint, left in product((True, False), repeat=2): + yield SampleInput( + clone(LU, LU_grad), + args=(pivots, clone(B if left else B.mT, B_grad)), + kwargs=dict(adjoint=adjoint, left=left), + ) + else: + yield SampleInput(clone(B, B_grad), args=(clone(LU, LU_grad), pivots)) + + +def sample_inputs_linalg_multi_dot(op_info, device, dtype, requires_grad, **kwargs): + # Each test case consists of the sizes in the chain of multiplications + # e.g. [2, 3, 4, 5] generates matrices (2, 3) @ (3, 4) @ (4, 5) + test_cases = [ + [1, 2, 1], + [2, 0, 2], + [0, 2, 2], + [2, 2, 2, 2], + [2, 3, 4, 5], + [5, 4, 0, 2], + [2, 4, 3, 5, 3, 2], + ] + + for sizes in test_cases: + tensors = [] + for size in itertools.pairwise(sizes): + t = make_tensor( + size, dtype=dtype, device=device, requires_grad=requires_grad + ) + tensors.append(t) + yield SampleInput(tensors) + + +def sample_inputs_linalg_matrix_norm(op_info, device, dtype, requires_grad, **kwargs): + low_precision_dtypes = (torch.float16, torch.bfloat16, torch.complex32) + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + + sizes = ((2, 2), (2, 3, 2)) + if dtype in low_precision_dtypes: + # svdvals not supported for low precision dtypes + ords = ("fro", inf, -inf, 1, -1) + else: + ords = ("fro", "nuc", inf, -inf, 1, -1, 2, -2) + dims = ((-2, -1), (-1, 0)) + + for size, ord, dim, keepdim in product(sizes, ords, dims, [True, False]): + yield SampleInput(make_arg(size), args=(ord, dim, keepdim)) + + +def sample_inputs_linalg_norm( + op_info, device, dtype, requires_grad, *, variant=None, **kwargs +): + if variant is not None and variant != "subgradient_at_zero": + raise ValueError( + f"Unsupported variant, expected variant to be 'subgradient_at_zero' but got: {variant}" + ) + + test_sizes = [ + (S,), + (0,), + (S, S), + (0, 0), + (S, 0), + (0, S), + (S, S, S), + (0, S, S), + (S, 0, S), + (0, 0, 0), + ] + + vector_ords = (None, 0, 0.5, 1, 2, 3.5, inf, -0.5, -1, -2, -3.5, -inf) + if dtype in {torch.float16, torch.bfloat16, torch.complex32}: + # svdvals not supported for low precision dtypes + matrix_ords = ("fro", inf, -inf, 1, -1) + else: + matrix_ords = (None, "fro", "nuc", inf, -inf, 1, -1, 2, -2) + + make_arg = partial( + make_tensor, + dtype=dtype, + device=device, + requires_grad=requires_grad, + low=None, + high=None, + ) + + for test_size in test_sizes: + is_vector_norm = len(test_size) == 1 + is_matrix_norm = len(test_size) == 2 + + # IndexError: amax(): Expected reduction dim 0 to have non-zero size. + is_valid_for_p2 = is_vector_norm or (test_size[-1] != 0 and test_size[-2] != 0) + + for keepdim in [False, True]: + if variant != "subgradient_at_zero" and is_valid_for_p2: + yield SampleInput(make_arg(test_size), keepdim=keepdim) + + if not (is_vector_norm or is_matrix_norm): + continue + + ords = vector_ords if is_vector_norm else matrix_ords + + for ord in ords: + if is_vector_norm and test_size[-1] == 0: + if ord == np.inf or (ord is not None and ord < 0): + # RuntimeError: linalg.vector_norm cannot compute the + # {ord} norm on an empty tensor because the operation + # does not have an identity + continue + elif is_matrix_norm: + dims_to_check = { + None: (0,), + -1: (1,), + -2: (0, 1), + -np.inf: (0,), + }.get(ord, ()) + + if any(test_size[d] == 0 for d in dims_to_check): + # IndexError: amax(): Expected reduction dim {dim} to + # have non-zero size. + continue + + no_grad_dims_to_check = { + np.inf: (0,), + 2: (0, 1), + 1: (1,), + }.get(ord, ()) + + if ( + any(test_size[d] == 0 for d in no_grad_dims_to_check) + and requires_grad + ): + continue + + if variant == "subgradient_at_zero": + yield SampleInput( + torch.zeros( + test_size, + dtype=dtype, + device=device, + requires_grad=requires_grad, + ), + ord, + keepdim=keepdim, + ) + else: + yield SampleInput(make_arg(test_size), ord, keepdim=keepdim) + + if ord in ["nuc", "fro"]: + yield SampleInput( + make_arg(test_size), ord=ord, keepdim=keepdim, dim=(0, 1) + ) + + +def sample_inputs_linalg_vecdot(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + batches = ((), (0,), (1,), (5,)) + ns = (0, 1, 3, 5) + for b, n in product(batches, ns): + shape = b + (n,) + yield SampleInput(make_arg(shape), args=(make_arg(shape),)) + for i in range(len(shape)): + yield SampleInput( + make_arg(shape), args=(make_arg(shape),), kwargs=dict(dim=i) + ) + + +def sample_inputs_linalg_invertible( + op_info, device, dtype, requires_grad=False, **kwargs +): + """ + This function generates invertible inputs for linear algebra ops + The input is generated as the itertools.product of 'batches' and 'ns'. + In total this function generates 8 SampleInputs + 'batches' cases include: + () - single input, + (0,) - zero batched dimension, + (2,) - batch of two matrices, + (1, 1) - 1x1 batch of matrices + 'ns' gives 0x0 and 5x5 matrices. + Zeros in dimensions are edge cases in the implementation and important to test for in order to avoid unexpected crashes. + """ + make_fn = make_fullrank_matrices_with_distinct_singular_values + make_arg = partial(make_fn, dtype=dtype, device=device, requires_grad=requires_grad) + + batches = [(), (0,), (2,), (1, 1)] + ns = [5, 0] + + for batch, n in product(batches, ns): + yield SampleInput(make_arg(*batch, n, n)) + + +def sample_inputs_matrix_rank(op_info, device, dtype, requires_grad=False, **kwargs): + """ + This function produces inputs for matrix rank that test + all possible combinations for atol and rtol + """ + + def make_tol_arg(kwarg_type, inp): + if kwarg_type == "none": + return None + if kwarg_type == "float": + return 1.0 + if kwarg_type != "tensor": + raise AssertionError(f"Expected kwarg_type == 'tensor', got {kwarg_type!r}") + return torch.ones(inp.shape[:-2], device=device) + + for tol_type in ["float", "tensor"]: + for atol_type, rtol_type in product(["none", tol_type], repeat=2): + if ( + not atol_type and not rtol_type + ): # default behavior, so skipped here so it's not tested 2 extra times + continue + for sample in sample_inputs_linalg_invertible( + op_info, device, dtype, requires_grad + ): + if sample.kwargs != {}: + raise AssertionError( + f"Expected sample.kwargs == {{}}, got {sample.kwargs}" + ) + sample.kwargs = { + "atol": make_tol_arg(atol_type, sample.input), + "rtol": make_tol_arg(rtol_type, sample.input), + } + yield sample + + # default kwargs + yield from sample_inputs_linalg_invertible(op_info, device, dtype, requires_grad) + + +def sample_inputs_linalg_pinv_singular( + op_info, device, dtype, requires_grad=False, **kwargs +): + """ + This function produces factors `a` and `b` to generate inputs of the form `a @ b.t()` to + test the backward method of `linalg_pinv`. That way we always preserve the rank of the + input no matter the perturbations applied to it by the gradcheck. + Note that `pinv` is Frechet-differentiable in a rank-preserving neighborhood. + """ + batches = [(), (0,), (2,), (1, 1)] + # the size of at least 30 is required to cause failures for the previous implicit implementation + # of the pinv's backward method, albeit it is slow. + size = [0, 3, 50] + + for batch, m, n in product(batches, size, size): + for k in range(min(3, m, n)): + # Note that by making the columns of `a` and `b` orthonormal we make sure that + # the product matrix `a @ b.t()` has condition number 1 when restricted to its image + a = ( + torch.rand(*batch, m, k, device=device, dtype=dtype) + .qr() + .Q.requires_grad_(requires_grad) + ) + b = ( + torch.rand(*batch, n, k, device=device, dtype=dtype) + .qr() + .Q.requires_grad_(requires_grad) + ) + yield SampleInput(a, args=(b,)) + + +def sample_inputs_linalg_cond(op_info, device, dtype, requires_grad=False, **kwargs): + make_arg = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad + ) + + # autograd is not supported for inputs with zero number of elements + shapes = ( + (S, S), + (2, S, S), + (2, 1, S, S), + ) + + for shape in shapes: + yield SampleInput(make_arg(shape)) + + +def sample_inputs_linalg_vander(op_info, device, dtype, requires_grad=False, **kwargs): + make_arg = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad + ) + + shapes = ( + (), + (1,), + (S,), + (2, S), + ) + + for shape in shapes: + if len(shape) > 0 and shape[-1] > 1: + yield SampleInput(make_arg(shape)) + n = shape[-1] if len(shape) > 0 else 1 + for i in range(3): + # n-1, n, n+1 + N = n + i - 1 + if N < 2: + continue + yield SampleInput(make_arg(shape), kwargs=dict(N=N)) + + +def np_vander_batched(x, N=None): + # Wrapper around np.vander that supports batches of 1 dimension (enough for the tests) + if x.ndim == 0: + x = x[np.newaxis] + if x.ndim == 1: + y = np.vander(x, N=N, increasing=True) + return y + else: + if N is None: + N = x.shape[-1] + y = np.vander(x.ravel(), N=N, increasing=True).reshape((*x.shape, N)) + return y + + +def sample_inputs_linalg_cholesky_inverse( + op_info, device, dtype, requires_grad=False, **kwargs +): + from torch.testing._internal.common_utils import random_well_conditioned_matrix + + # Cholesky factorization is for positive-definite matrices + single_well_conditioned_matrix = random_well_conditioned_matrix( + S, S, dtype=dtype, device=device + ) + batch_well_conditioned_matrices = random_well_conditioned_matrix( + 2, S, S, dtype=dtype, device=device + ) + single_pd = single_well_conditioned_matrix @ single_well_conditioned_matrix.mH + batch_pd = batch_well_conditioned_matrices @ batch_well_conditioned_matrices.mH + + inputs = ( + torch.zeros(0, 0, dtype=dtype, device=device), # 0x0 matrix + torch.zeros(0, 2, 2, dtype=dtype, device=device), # zero batch of matrices + single_pd, + batch_pd, + ) + test_cases = (torch.linalg.cholesky(a, upper=False) for a in inputs) + for l in test_cases: + # generated lower-triangular samples + l.requires_grad = requires_grad + yield SampleInput(l) # upper=False by default + yield SampleInput( + l.detach().clone().requires_grad_(requires_grad), kwargs=dict(upper=False) + ) + + # generate upper-triangular inputs + u = l.detach().clone().mT.contiguous().requires_grad_(requires_grad) + yield SampleInput(u, kwargs=dict(upper=True)) + + +def sample_inputs_linalg_ldl_factor( + op_info, device, dtype, requires_grad=False, **kwargs +): + from torch.testing._internal.common_utils import ( + random_hermitian_pd_matrix, + random_symmetric_pd_matrix, + ) + + device = torch.device(device) + + # Symmetric inputs + yield SampleInput( + random_symmetric_pd_matrix(S, dtype=dtype, device=device), + kwargs=dict(hermitian=False), + ) # single matrix + yield SampleInput( + random_symmetric_pd_matrix(S, 2, dtype=dtype, device=device), + kwargs=dict(hermitian=False), + ) # batch of matrices + yield SampleInput( + torch.zeros(0, 0, dtype=dtype, device=device), kwargs=dict(hermitian=False) + ) # 0x0 matrix + yield SampleInput( + torch.zeros(0, 2, 2, dtype=dtype, device=device), kwargs=dict(hermitian=False) + ) # zero batch of matrices + + # Hermitian inputs + # hermitian=True for complex inputs on CUDA is supported only with MAGMA 2.5.4+ + magma_254_available = device.type == "cuda" and _get_magma_version() >= (2, 5, 4) + if dtype.is_complex and (device.type == "cpu" or magma_254_available): + yield SampleInput( + random_hermitian_pd_matrix(S, dtype=dtype, device=device), + kwargs=dict(hermitian=True), + ) # single matrix + yield SampleInput( + random_hermitian_pd_matrix(S, 2, dtype=dtype, device=device), + kwargs=dict(hermitian=True), + ) # batch of matrices + + +def sample_inputs_linalg_ldl_solve( + op_info, device, dtype, requires_grad=False, **kwargs +): + # Generate LDL factors of symmetric (and Hermitian on CPU) matrices + from torch.testing._internal.common_utils import ( + random_hermitian_pd_matrix, + random_symmetric_pd_matrix, + ) + + device = torch.device(device) + symmetric_inputs = ( + random_symmetric_pd_matrix(S, dtype=dtype, device=device), # single matrix + random_symmetric_pd_matrix( + S, 2, dtype=dtype, device=device + ), # batch of matrices + torch.zeros(0, 0, dtype=dtype, device=device), # 0x0 matrix + torch.zeros(0, 2, 2, dtype=dtype, device=device), # zero batch of matrices + ) + hermitian_inputs = ( + ( + random_hermitian_pd_matrix(S, dtype=dtype, device=device), + random_hermitian_pd_matrix(S, 2, dtype=dtype, device=device), + ) + if device.type == "cpu" and dtype.is_complex + else () + ) + test_cases1 = ( + torch.linalg.ldl_factor_ex(a, hermitian=False) for a in symmetric_inputs + ) + test_cases2 = ( + torch.linalg.ldl_factor_ex(a, hermitian=True) for a in hermitian_inputs + ) + + # Symmetric case + make_arg = partial( + make_tensor, device=device, dtype=dtype, requires_grad=requires_grad + ) + for test_case in test_cases1: + factors, pivots, _ = test_case + factors.requires_grad = requires_grad + for B_batch_shape in ((), factors.shape[:-2]): + B = make_arg((*B_batch_shape, factors.shape[-1], S)) + yield SampleInput(factors, args=(pivots, B), kwargs=dict(hermitian=False)) + clone_factors = factors.detach().clone().requires_grad_(requires_grad) + yield SampleInput( + clone_factors, args=(pivots, B), kwargs=dict(hermitian=False) + ) + + # Hermitian case + for test_case in test_cases2: + factors, pivots, _ = test_case + factors.requires_grad = requires_grad + for B_batch_shape in ((), factors.shape[:-2]): + B = make_arg((*B_batch_shape, factors.shape[-1], S)) + yield SampleInput(factors, args=(pivots, B), kwargs=dict(hermitian=True)) + clone_factors = factors.detach().clone().requires_grad_(requires_grad) + yield SampleInput( + clone_factors, args=(pivots, B), kwargs=dict(hermitian=True) + ) + + +def sample_inputs_linalg_lstsq(op_info, device, dtype, requires_grad=False, **kwargs): + from torch.testing._internal.common_utils import random_well_conditioned_matrix + + device = torch.device(device) + + drivers: tuple[str, ...] + if device.type == "cuda": + drivers = ("gels",) + else: + drivers = ("gels", "gelsy", "gelss", "gelsd") + + # we generate matrices of shape (..., n + delta, n) + deltas: tuple[int, ...] + if device.type == "cpu" or has_cusolver(): + deltas = (-1, 0, +1) + # only square systems if Cusolver is not available + # because we solve a lstsq problem with a transposed matrix in the backward + else: + deltas = (0,) + + for batch, driver, delta in product(((), (3,), (3, 3)), drivers, deltas): + shape = batch + (3 + delta, 3) + a = random_well_conditioned_matrix(*shape, dtype=dtype, device=device) + a.requires_grad_(requires_grad) + b = make_tensor( + shape, + dtype=dtype, + device=device, + low=None, + high=None, + requires_grad=requires_grad, + ) + yield SampleInput(a, b, driver=driver) + + +def error_inputs_lstsq(op_info, device, **kwargs): + zero_d = torch.randn((), device=device) + yield ErrorInput( + SampleInput(zero_d, args=(zero_d,)), + error_type=RuntimeError, + error_regex="at least 2 dimensions", + ) + + +def error_inputs_lstsq_grad_oriented(op_info, device, **kwargs): + zero_d = torch.randn((), device=device) + yield ErrorInput( + SampleInput(zero_d, args=(zero_d, None)), + error_type=RuntimeError, + error_regex="at least 2 dimensions", + ) + + +def sample_inputs_diagonal_diag_embed(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad + ) + + # Shapes for 2D Tensors + shapes_2d = ((S, S), (3, 5), (5, 3)) + + # Shapes for 3D Tensors + shapes_3d = ((S, S, S),) + + kwargs_2d = ({}, dict(offset=2), dict(offset=2), dict(offset=1)) + kwargs_3d = ( + dict(offset=1, dim1=1, dim2=2), + dict(offset=2, dim1=0, dim2=1), + dict(offset=-2, dim1=0, dim2=1), + ) + + for shape, kwarg in chain( + product(shapes_2d, kwargs_2d), product(shapes_3d, kwargs_3d) + ): + yield SampleInput(make_arg(shape), kwargs=kwarg) + + +def error_inputs_diagonal_diag_embed(op_info, device, **kwargs): + make_arg = partial(make_tensor, device=device, dtype=torch.float32) + + shapes1d = (0, 1, (0,), (1,)) + shapes2d = ((M, L),) + shapes3d = ((M, S, L),) + + kwargs1d = {} + + kwargs2d = ( + # dim1 == dim2 is not allowed + dict(dim1=1, dim2=1), + # out of bounds dims are not allowed + dict(dim1=10000), + dict(dim2=10000), + ) + + kwargs3d = kwargs2d + + samples1d = product(shapes1d, kwargs1d) + samples2d = product(shapes2d, kwargs2d) + samples3d = product(shapes3d, kwargs3d) + + for shape, kwargs in chain(samples1d, samples2d, samples3d): + arg = make_arg(shape) + sample = SampleInput(input=arg, kwargs=kwargs) + + dim1 = kwargs.get("dim1") + dim2 = kwargs.get("dim2") + + if "diagonal" in op_info.name: + num_dim = arg.dim() + elif op_info.name in ("diag_embed", "_refs.diag_embed"): + # these are valid inputs for diag_embed + if shape in ((0,), (1,)): + continue + num_dim = arg.dim() + 1 + else: + raise RuntimeError("should be unreachable") + + bound1 = -num_dim + bound2 = num_dim - 1 + dim_range = range(bound1, bound2 + 1) + dim1_cond = dim1 and dim1 not in dim_range + dim2_cond = dim2 and dim2 not in dim_range + + if dim1 == dim2: + err = f"diagonal dimensions cannot be identical {dim1}, {dim2}" + yield ErrorInput(sample, error_regex=err, error_type=RuntimeError) + elif dim1_cond or dim2_cond: + err_dim = dim1 if dim1_cond else dim2 + err = ( + r"Dimension out of range \(expected to be in range of " + rf"\[{bound1}, {bound2}\], but got {err_dim}\)" + ) + yield ErrorInput(sample, error_regex=err, error_type=IndexError) + else: + raise RuntimeError("should be unreachable") + + +def sample_inputs_linalg_cholesky( + op_info, device, dtype, requires_grad=False, **kwargs +): + """ + This function generates always positive-definite input for torch.linalg.cholesky using + random_hermitian_pd_matrix. + The input is generated as the itertools.product of 'batches' and 'ns'. + In total this function generates 8 SampleInputs + 'batches' cases include: + () - single input, + (0,) - zero batched dimension, + (2,) - batch of two matrices, + (1, 1) - 1x1 batch of matrices + 'ns' gives 0x0 and 5x5 matrices. + Zeros in dimensions are edge cases in the implementation and important to test for in order to avoid unexpected crashes. + """ + from torch.testing._internal.common_utils import random_hermitian_pd_matrix + + batches = [(), (0,), (2,), (1, 1)] + ns = [5, 0] + for batch, n, upper in product(batches, ns, [True, False]): + a = random_hermitian_pd_matrix(n, *batch, dtype=dtype, device=device) + a.requires_grad = requires_grad + yield SampleInput(a, upper=upper) + + +def sample_inputs_linalg_eig(op_info, device, dtype, requires_grad=False, **kwargs): + """ + This function generates input for torch.linalg.eig + """ + + def out_fn(output): + return output[0], abs(output[1]) + + samples = sample_inputs_linalg_invertible(op_info, device, dtype, requires_grad) + for sample in samples: + sample.output_process_fn_grad = out_fn + yield sample + + +def sample_inputs_linalg_eigh(op_info, device, dtype, requires_grad=False, **kwargs): + """ + This function generates input for torch.linalg.eigh/eigvalsh with UPLO="U" or "L" keyword argument. + """ + + def out_fn(output): + if isinstance(output, tuple): + # eigh function + return output[0], abs(output[1]) + else: + # eigvalsh function + return output + + # Samples do not need to be Hermitian, as we're using gradcheck_wrapper_hermitian_input + samples = sample_inputs_linalg_invertible(op_info, device, dtype, requires_grad) + for sample in samples: + # Note: we cannot use np.random.choice here as TorchDynamo + # does not support tensors of strings. + sample.kwargs = {"UPLO": random.choice(["L", "U"])} + sample.output_process_fn_grad = out_fn + yield sample + + +def sample_inputs_linalg_pinv(op_info, device, dtype, requires_grad=False, **kwargs): + """ + This function generates input for torch.linalg.pinv with hermitian=False keyword argument. + """ + for o in sample_inputs_linalg_invertible( + op_info, device, dtype, requires_grad, **kwargs + ): + real_dtype = o.input.real.dtype if dtype.is_complex else dtype + # requires_grad path for rtol tensor is not implemented + for rtol in (None, 1.0, torch.tensor(1.0, dtype=real_dtype, device=device)): + o = clone_sample(o) + o.kwargs = {"rtol": rtol} + yield o + + +def sample_inputs_linalg_pinv_hermitian( + op_info, device, dtype, requires_grad=False, **kwargs +): + """ + This function generates input for torch.linalg.pinv with hermitian=True keyword argument. + """ + for o in sample_inputs_linalg_invertible( + op_info, device, dtype, requires_grad, **kwargs + ): + o.kwargs = {"hermitian": True} + yield o + + +def sample_inputs_linalg_solve( + op_info, device, dtype, requires_grad=False, vector_rhs_allowed=True, **kwargs +): + """ + This function generates always solvable input for torch.linalg.solve + We sample a fullrank square matrix (i.e. invertible) A + The first input to torch.linalg.solve is generated as the itertools.product of 'batches' and 'ns'. + The second input is generated as the product of 'batches', 'ns' and 'nrhs'. + In total this function generates 18 SampleInputs + 'batches' cases include: + () - single input, + (0,) - zero batched dimension, + (2,) - batch of two matrices. + 'ns' gives 0x0 and 5x5 matrices. + and 'nrhs' controls the number of vectors to solve for: + () - using 1 as the number of vectors implicitly + (1,) - same as () but explicit + (3,) - solve for 3 vectors. + Zeros in dimensions are edge cases in the implementation and important to test for in order to avoid unexpected crashes. + 'vector_rhs_allowed' controls whether to include nrhs = () to the list of SampleInputs. + torch.solve / triangular_solve / cholesky_solve (opposed to torch.linalg.solve) do not allow + 1D tensors (vectors) as the right-hand-side. + Once torch.solve / triangular_solve / cholesky_solve and its testing are removed, + 'vector_rhs_allowed' may be removed here as well. + """ + make_fullrank = make_fullrank_matrices_with_distinct_singular_values + make_a = partial( + make_fullrank, dtype=dtype, device=device, requires_grad=requires_grad + ) + make_b = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad + ) + + batches = [(), (0,), (2,), (2, 2)] + ns = [5, 0] + if vector_rhs_allowed: + nrhs = [(), (1,), (3,)] + else: + nrhs = [(1,), (3,)] + + for n, batch, rhs in product(ns, batches, nrhs): + yield SampleInput(make_a(*batch, n, n), args=(make_b(batch + (n,) + rhs),)) + + +def sample_inputs_linalg_solve_triangular( + op_info, device, dtype, requires_grad=False, **kwargs +): + make_arg = partial(make_tensor, dtype=dtype, device=device) + bs = (1, 2, 0) + ns = (3, 0) + ks = (1, 3, 0) + + for b, n, k, (left, upper, uni) in product( + bs, ns, ks, product((True, False), repeat=3) + ): + if b == 1: + A = make_arg((n, n)) if left else make_arg((k, k)) + B = make_arg((n, k)) + else: + A = make_arg((b, n, n)) if left else make_arg((b, k, k)) + B = make_arg((b, n, k)) + if uni: + # Not really necessary, but writing it for consistency + A.diagonal(0, -2, -1).fill_(1.0) + else: + d = A.diagonal(0, -2, -1) + d[d.abs() < 1e-6] = 1.0 + if upper: + A.triu_() + else: + A.tril_() + kwargs = {"upper": upper, "left": left, "unitriangular": uni} + if requires_grad: + for grad_A, grad_B in product((True, False), repeat=2): + # Either A or B needs to have a gradient + if not grad_A and not grad_B: + continue + yield SampleInput( + A.clone().requires_grad_(grad_A), + args=(B.clone().requires_grad_(grad_B),), + kwargs=kwargs, + ) + else: + yield SampleInput(A, args=(B,), kwargs=kwargs) + + +def sample_inputs_legacy_solve(op_info, device, dtype, requires_grad=False, **kwargs): + """ + This function generates always solvable input for legacy solve functions + (the ones that are not in torch.linalg module). + The difference from sample_inputs_linalg_solve is that here the right-hand-side of A x = b equation + should have b.ndim >= 2, vectors are not allowed. + Also the arguments order is swapped. + """ + out = sample_inputs_linalg_solve( + op_info, device, dtype, requires_grad=requires_grad, vector_rhs_allowed=False + ) + + def out_fn(output): + return output[0] + + # Reverses tensor order + for sample in out: + sample.input, sample.args = sample.args[0], (sample.input,) + if op_info.name == "solve": + sample.output_process_fn_grad = out_fn + yield sample + + +def sample_inputs_linalg_lu(op_info, device, dtype, requires_grad=False, **kwargs): + full_rank = op_info.name == "linalg.lu_factor" + make_fn = ( + make_tensor + if not full_rank + else make_fullrank_matrices_with_distinct_singular_values + ) + make_arg = partial(make_fn, dtype=dtype, device=device, requires_grad=requires_grad) + + def out_fn(output): + if op_info.name == "linalg.lu": + return output[1], output[2] + else: + return output + + batch_shapes = ((), (3,), (3, 3), (0,)) + # pivot=False only supported in CUDA + pivots = (True, False) if torch.device(device).type == "cuda" else (True,) + deltas = (-2, -1, 0, +1, +2) + for batch_shape, pivot, delta in product(batch_shapes, pivots, deltas): + shape = batch_shape + (S + delta, S) + # Insanely annoying that make_fullrank_blablabla accepts a *shape and not a tuple! + A = make_arg(shape) if not full_rank else make_arg(*shape) + yield SampleInput(A, kwargs={"pivot": pivot}, output_process_fn_grad=out_fn) + + +def sample_inputs_linalg_svdvals(op_info, device, dtype, requires_grad=False, **kwargs): + make_arg = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad + ) + + batches = [(), (0,), (2,), (1, 1)] + ns = [5, 2, 0] + + for batch, m, n in product(batches, ns, ns): + yield SampleInput(make_arg(batch + (m, n))) + + +def sample_inputs_linalg_qr_geqrf( + op_info, device, dtype, requires_grad=False, **kwargs +): + # QR is just well defined when the matrix is full rank + make_fullrank = make_fullrank_matrices_with_distinct_singular_values + make_arg = partial( + make_fullrank, dtype=dtype, device=device, requires_grad=requires_grad + ) + + batches = [(), (0,), (2,), (1, 1)] + ns = [5, 2, 0] + + for batch, (m, n) in product(batches, product(ns, ns)): + shape = batch + (m, n) + yield SampleInput(make_arg(*shape)) + + +def sample_inputs_tensorsolve(op_info, device, dtype, requires_grad, **kwargs): + a_shapes = [(2, 3, 6), (3, 4, 4, 3)] + # Zero-dim tensors are not supported in NumPy, so we skip them for now. + # NumPy is used in reference check tests. + # See https://github.com/numpy/numpy/pull/20482 for tracking NumPy bugfix. + # a_shapes += [(0, 0, 1, 2, 3, 0)] + dimss = [None, (0, 2)] + + make_arg = partial( + make_tensor, dtype=dtype, device=device, requires_grad=requires_grad + ) + for a_shape, dims in itertools.product(a_shapes, dimss): + a = make_arg(a_shape) + b = make_arg(a_shape[:2]) + yield SampleInput(a, b, dims=dims) + + +def sample_inputs_tensorinv(op_info, device, dtype, requires_grad, **kwargs): + make_arg = make_fullrank_matrices_with_distinct_singular_values + + def make_input(): + return make_arg(12, 12, device=device, dtype=dtype, requires_grad=requires_grad) + + # lhs / rhs shape can have any number of dimensions as long as their product equals 12 + shapes = [ + ((2, 2, 3), (12, 1)), + ((4, 3), (6, 1, 2)), + ] + + for shape_lhs, shape_rhs in shapes: + inp = make_input().reshape(*shape_lhs, *shape_rhs).detach() + inp.requires_grad_(requires_grad) + yield SampleInput(inp, ind=len(shape_lhs)) + + +op_db: list[OpInfo] = [ + OpInfo( + "linalg.cross", + ref=lambda x, y, dim=-1: np.cross(x, y, axis=dim), + op=torch.linalg.cross, + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + aten_name="linalg_cross", + sample_inputs_func=sample_inputs_cross, + error_inputs_func=error_inputs_cross, + supports_out=True, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + skips=( + DecorateInfo( + unittest.skip("Unsupported on MPS for now"), + "TestCommon", + "test_numpy_ref_mps", + ), + ), + ), + OpInfo( + "linalg.det", + aten_name="linalg_det", + op=torch.linalg.det, + aliases=("det",), + dtypes=floating_and_complex_types(), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_linalg_det_logdet_slogdet, + decorators=[skipCPUIfNoLapack, skipCUDAIfNoMagmaAndNoCusolver], + check_batched_gradgrad=False, + skips=( + # Exception: linalg.lu_factor(): MPS doesn't support complex types. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.diagonal", + aten_name="linalg_diagonal", + aten_backward_name="diagonal_backward", + dtypes=all_types_and_complex_and( + torch.bool, torch.bfloat16, torch.float16, torch.chalf + ), + supports_out=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_diagonal_diag_embed, + error_inputs_func=error_inputs_diagonal_diag_embed, + ), + OpInfo( + "linalg.cholesky", + aten_name="linalg_cholesky", + dtypes=floating_and_complex_types(), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_linalg_cholesky, + gradcheck_wrapper=gradcheck_wrapper_hermitian_input, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + # linalg.solve.triangular(); Only float is supported! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + # Exception: linalg.cholesky: The factorization could not be completed because the input is not positive-definite + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.cholesky_ex", + aten_name="linalg_cholesky_ex", + dtypes=floating_and_complex_types(), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_linalg_cholesky, + gradcheck_wrapper=gradcheck_wrapper_hermitian_input, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + # The following dtypes worked in forward but are not listed by the + # OpInfo: {torch.bfloat16, torch.float16}. The following dtypes did + # not work in backward but are listed by the OpInfo: + # {torch.complex64}. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + # RuntimeError: linalg.solve.triangular(); Only float is supported! + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.vecdot", + aten_name="linalg_vecdot", + ref=lambda x, y, *, dim=-1: (x.conj() * y).sum(dim), + dtypes=floating_and_complex_types_and(torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_linalg_vecdot, + check_batched_forward_grad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Issue with conj and torch dispatch, see https://github.com/pytorch/pytorch/issues/82479 + DecorateInfo( + unittest.skip("Skipped!"), + "TestSchemaCheckModeOpInfo", + "test_schema_correctness", + dtypes=(torch.complex64, torch.complex128), + ), + DecorateInfo( + unittest.skip("Unsupported on MPS for now"), + "TestCommon", + "test_numpy_ref_mps", + ), + DecorateInfo( + toleranceOverride({torch.half: tol(atol=1.2e-2, rtol=1.7e-2)}), + "TestInductorOpInfo", + "test_comprehensive", + device_type="cuda", + ), + ), + ), + OpInfo( + "linalg.cond", + aten_name="linalg_cond", + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_cond, + check_batched_gradgrad=False, + check_batched_forward_grad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack, with_tf32_off], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_no_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + # The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + OpInfo( + "linalg.eig", + aten_name="linalg_eig", + op=torch.linalg.eig, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_eig, + check_batched_forward_grad=False, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # AssertionError: Scalars are not equal! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="cpu" + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + # Exception: The operator 'aten::linalg_eig' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack, with_tf32_off], + ), + OpInfo( + "linalg.eigvals", + aten_name="linalg_eigvals", + op=torch.linalg.eigvals, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_invertible, + check_batched_forward_grad=False, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + # Exception: The operator 'aten::linalg_eig' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + OpInfo( + "linalg.eigh", + aten_name="linalg_eigh", + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_eigh, + gradcheck_wrapper=gradcheck_wrapper_hermitian_input, + check_batched_forward_grad=False, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack, with_tf32_off], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + # Exception: The operator 'aten::_linalg_eigh.eigenvalues' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + OpInfo( + "linalg.eigvalsh", + aten_name="linalg_eigvalsh", + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_eigh, + gradcheck_wrapper=gradcheck_wrapper_hermitian_input, + check_batched_forward_grad=False, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack], + skips=( + # Pre-existing condition; Needs to be fixed + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + # Exception: The operator 'aten::_linalg_eigh.eigenvalues' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + OpInfo( + "linalg.householder_product", + aten_name="linalg_householder_product", + op=torch.linalg.householder_product, + aliases=("orgqr",), + dtypes=floating_and_complex_types(), + dtypesIfMPS=floating_and_complex_types_and(torch.bfloat16, torch.float16), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + # TODO: backward uses in-place operations that vmap doesn't like + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_householder_product, + decorators=[ + skipCUDAIfNoCusolver, + skipCPUIfNoLapack, + DecorateInfo( + toleranceOverride({torch.complex64: tol(atol=1e-3, rtol=1e-3)}) + ), + DecorateInfo( + unittest.skip("Skipped! Flaky"), + "TestFwdGradients", + "test_fn_fwgrad_bwgrad", + device_type="cpu", + dtypes=(torch.complex128,), + ), + # Exception: "orgqr_cpu" not implemented for 'Half' + DecorateInfo( + unittest.expectedFailure, + "TestConsistency", + device_type="mps", + dtypes=(torch.bfloat16, torch.float16), + ), + skipCUDAIfRocm, # regression in ROCm 6.4 + ], + ), + OpInfo( + "linalg.ldl_factor", + aten_name="linalg_ldl_factor", + dtypes=floating_and_complex_types(), + supports_autograd=False, + sample_inputs_func=sample_inputs_linalg_ldl_factor, + decorators=[skipCUDAIfNoMagmaAndNoLinalgsolver, skipCPUIfNoLapack], + skips=( + # NotImplementedError: The operator 'aten::linalg_ldl_factor_ex.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + OpInfo( + "linalg.ldl_factor_ex", + aten_name="linalg_ldl_factor_ex", + dtypes=floating_and_complex_types(), + supports_autograd=False, + sample_inputs_func=sample_inputs_linalg_ldl_factor, + decorators=[skipCUDAIfNoMagmaAndNoLinalgsolver, skipCPUIfNoLapack], + skips=( + # NotImplementedError: The operator 'aten::linalg_ldl_factor_ex.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + OpInfo( + "linalg.ldl_solve", + aten_name="linalg_ldl_solve", + dtypes=floating_and_complex_types(), + supports_autograd=False, + sample_inputs_func=sample_inputs_linalg_ldl_solve, + decorators=[ + skipCUDAIfNoCusolver, + skipCUDAIfRocm, + skipCPUIfNoLapack, + ], + skips=( + # NotImplementedError: The operator 'aten::linalg_ldl_factor_ex.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + OpInfo( + "linalg.lstsq", + aten_name="linalg_lstsq", + dtypes=floating_and_complex_types(), + supports_out=True, + sample_inputs_func=sample_inputs_linalg_lstsq, + error_inputs_func=error_inputs_lstsq, + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack], + skips=( + # we skip gradient checks for this suite as they are tested in + # variant_test_name='grad_oriented' + DecorateInfo(unittest.skip("Skipped!"), "TestFwdGradients"), + DecorateInfo(unittest.skip("Skipped!"), "TestBwdGradients"), + # The values for attribute 'shape' do not match + DecorateInfo(unittest.skip("Skipped!"), "TestCommon", "test_out"), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + # Exception: The operator 'aten::linalg_lstsq.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + # see https://github.com/pytorch/pytorch/issues/177249 + DecorateInfo( + unittest.expectedFailure, + "TestJit", + "test_variant_consistency_jit", + device_type="cpu", + dtypes=[torch.complex64], + active_if=IS_LINUX and IS_ARM64 and not IS_CPU_CAPABILITY_SVE256, + ), + ), + ), + OpInfo( + "linalg.lstsq", + aten_name="linalg_lstsq", + variant_test_name="grad_oriented", + # gradchecks for forward AD fails with full output tuple + # works when taking [:2], which is (solution, residuals) + op=lambda a, b, driver: torch.linalg.lstsq(a, b, driver=driver)[:2], + supports_out=False, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_lstsq, + error_inputs_func=error_inputs_lstsq_grad_oriented, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_autograd=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack], + skips=( + # tests do not work with passing lambda for op + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + DecorateInfo( + unittest.expectedFailure, + "TestOperatorSignatures", + "test_get_torch_func_signature_exhaustive", + ), + # Exception: The operator 'aten::linalg_lstsq.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + OpInfo( + "linalg.matrix_power", + aliases=("matrix_power",), + aten_name="linalg_matrix_power", + dtypes=floating_and_complex_types(), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_inplace_autograd=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + check_batched_grad=False, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack, with_tf32_off], + skips=( + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=8e-5, rtol=2e-6)}), + "TestConsistency", + "test_output_grad_match", + device_type="mps", + ), + # RuntimeError: linalg_inv: not supported for complex types yet! + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + dtypes=(torch.complex64,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=(torch.complex64,), + ), + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out", + device_type="mps", + dtypes=(torch.float32,), + ), + ), + sample_inputs_func=sample_inputs_linalg_matrix_power, + ), + OpInfo( + "linalg.multi_dot", + # Need this lambda because gradcheck does not work with TensorList inputs + aten_name="linalg_multi_dot", + dtypes=all_types_and_complex_and(torch.half, torch.bfloat16), + dtypesIfCUDA=floating_and_complex_types_and(torch.half, torch.bfloat16), + supports_inplace_autograd=False, + # Batched grad checks fail for empty input tensors (see https://github.com/pytorch/pytorch/issues/53407) + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # https://github.com/pytorch/pytorch/issues/66357 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_linalg_multi_dot, + gradcheck_nondet_tol=GRADCHECK_NONDET_TOL, + skips=( + # https://github.com/pytorch/pytorch/issues/67470 + DecorateInfo( + unittest.skip("67470!"), "TestCommon", "test_noncontiguous_samples" + ), + # Fails on XLA. + # AssertionError: False is not true : Tensors failed to compare as equal! + DecorateInfo( + unittest.skip("Skipped!"), + "TestOpInfo", + device_type="xla", + dtypes=(torch.long,), + ), + # https://github.com/pytorch/pytorch/issues/71774 + DecorateInfo( + unittest.skip("Skipped!"), + "TestNNCOpInfo", + "test_nnc_correctness", + device_type="cpu", + dtypes=(torch.long,), + ), + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out", + device_type="mps", + dtypes=(torch.float32,), + ), + ), + ), + # NB: linalg.norm has two variants so that different skips can be used for different sample inputs + OpInfo( + "linalg.norm", + aten_name="linalg_norm", + op=torch.linalg.norm, + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack, with_tf32_off], + sample_inputs_func=sample_inputs_linalg_norm, + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + skips=( + DecorateInfo( + unittest.expectedFailure, "TestBwdGradients", "test_fn_gradgrad" + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_no_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + # NotImplementedError: The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref_torch_fallback", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + ), + ), + ), + OpInfo( + "linalg.norm", + op=torch.linalg.norm, + variant_test_name="subgradients_at_zero", + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack, with_tf32_off], + sample_inputs_func=partial( + sample_inputs_linalg_norm, variant="subgradient_at_zero" + ), + aten_name="linalg_norm", + supports_forward_ad=True, + # torch.autograd.gradcheck.GradcheckError: While computing batched gradients, got: + # Could not allocate memory to change Tensor SizesAndStrides! + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + skips=( + # [NEW] Skips specifically for sample inputs at zero + # norm's vjp/jvp are not well-conditioned near zero + DecorateInfo( + unittest.expectedFailure, "TestBwdGradients", "test_fn_gradgrad" + ), + DecorateInfo( + unittest.expectedFailure, "TestFwdGradients", "test_fn_fwgrad_bwgrad" + ), + DecorateInfo( + unittest.expectedFailure, "TestFwdGradients", "test_forward_mode_AD" + ), + DecorateInfo(unittest.expectedFailure, "TestBwdGradients", "test_fn_grad"), + # NotImplementedError: The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref_torch_fallback", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + ), + ), + ), + OpInfo( + "linalg.matrix_norm", + aten_name="linalg_matrix_norm", + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + supports_forward_ad=True, + check_batched_forward_grad=False, + check_batched_gradgrad=False, + supports_fwgrad_bwgrad=True, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack, with_tf32_off], + sample_inputs_func=sample_inputs_linalg_matrix_norm, + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_no_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + # NotImplementedError: The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + ), + ), + ), + OpInfo( + "linalg.qr", + aten_name="linalg_qr", + op=torch.linalg.qr, + dtypes=floating_and_complex_types(), + dtypesIfMPS=floating_types(), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # In-place ops + check_batched_gradgrad=False, + sample_inputs_func=sample_inputs_linalg_qr_geqrf, + decorators=[skipCUDAIfNoCusolver, skipCPUIfNoLapack], + ), + OpInfo( + "linalg.slogdet", + aten_name="linalg_slogdet", + op=torch.linalg.slogdet, + dtypes=floating_and_complex_types(), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_linalg_det_logdet_slogdet, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + # Exception: linalg.lu_factor(): MPS doesn't support complex types. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.vander", + aten_name="linalg_vander", + ref=np_vander_batched, + op=torch.linalg.vander, + dtypes=all_types_and_complex(), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_out=False, + sample_inputs_func=sample_inputs_linalg_vander, + skips=( + DecorateInfo( + unittest.skip("Unsupported on MPS for now"), + "TestCommon", + "test_numpy_ref_mps", + ), + ), + ), + ReductionOpInfo( + "linalg.vector_norm", + op=torch.linalg.vector_norm, + identity=0, + nan_policy="propagate", + supports_multiple_dims=True, + complex_to_real=True, + supports_forward_ad=True, + # torch.autograd.gradcheck.GradcheckError: While computing batched gradients + # got: Could not allocate memory to change Tensor SizesAndStrides! + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + dtypes=floating_and_complex_types_and(torch.float16, torch.bfloat16), + generate_args_kwargs=sample_kwargs_vector_norm, + aten_name="linalg_vector_norm", + ), + OpInfo( + "linalg.lu_factor", + aten_name="linalg_lu_factor", + op=torch.linalg.lu_factor, + dtypes=floating_and_complex_types(), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_linalg_lu, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + # linalg.lu_factor: LU without pivoting is not implemented on the CPU + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_compare_cpu", + active_if=(not TEST_XPU), + ), + # Exception: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + # RuntimeError: linalg.lu_factor(): MPS doesn't support complex types. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.lu_factor_ex", + aten_name="linalg_lu_factor_ex", + op=torch.linalg.lu_factor_ex, + dtypes=floating_and_complex_types(), + # https://github.com/pytorch/pytorch/issues/80411 + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_linalg_lu, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + # linalg.lu_factor: LU without pivoting is not implemented on the CPU + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_compare_cpu", + active_if=(not TEST_XPU), + ), + # Exception: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + # RuntimeError: linalg.lu_factor(): MPS doesn't support complex types. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.lu", + aten_name="linalg_lu", + op=torch.linalg.lu, + dtypes=floating_and_complex_types(), + # https://github.com/pytorch/pytorch/issues/80411 + # Runs very slowly on slow-gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_linalg_lu, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + # linalg.lu_factor: LU without pivoting is not implemented on the CPU + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_compare_cpu", + active_if=(not TEST_XPU), + ), + # AssertionError: Resizing an out= argument with no elements threw a resize warning! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + # Exception: linalg.lu_factor(): MPS doesn't support complex types. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.lu_solve", + op=torch.linalg.lu_solve, + aten_name="linalg_lu_solve", + dtypes=floating_and_complex_types(), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_lu_solve, + skips=( + DecorateInfo( + unittest.skip("Tests different backward paths"), + "TestCommon", + "test_floating_inputs_are_differentiable", + ), + # RuntimeError: The size of tensor a (5) must match the size of tensor b (4) at non-singleton dimension 1 + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + # AssertionError: The values for attribute 'shape' do not match: torch.Size([3, 4]) != torch.Size([0]). + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + # RuntimeError: linalg.solve.triangular(); Only float is supported! + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + decorators=[skipCPUIfNoLapack, skipCUDAIfNoMagmaAndNoCusolver], + ), + OpInfo( + "linalg.inv", + aten_name="linalg_inv", + op=torch.linalg.inv, + aliases=("inverse",), + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_invertible, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + # Exception: linalg_inv: not supported for complex types yet! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.inv_ex", + aten_name="linalg_inv_ex", + op=torch.linalg.inv_ex, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_invertible, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + # Exception: linalg_inv: not supported for complex types yet! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.solve", + aten_name="linalg_solve", + op=torch.linalg.solve, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_solve, + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[ + skipCUDAIfNoMagmaAndNoCusolver, + skipCPUIfNoLapack, + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1.3e-05, rtol=6e-04)}), + "TestCommon", + "test_noncontiguous_samples", + device_type="cpu", + ), + # Exception: linalg.lu_factor(): MPS only supports floats. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + ), + ), + OpInfo( + "linalg.solve_ex", + aten_name="linalg_solve_ex", + op=torch.linalg.solve_ex, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_solve, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[ + skipCUDAIfNoMagmaAndNoCusolver, + skipCPUIfNoLapack, + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1.3e-05, rtol=6e-04)}), + "TestCommon", + "test_noncontiguous_samples", + device_type="cpu", + ), + ], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + # Exception: linalg.lu_factor(): MPS only supports floats. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.solve_triangular", + aten_name="linalg_solve_triangular", + op=torch.linalg.solve_triangular, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_linalg_solve_triangular, + supports_fwgrad_bwgrad=True, + skips=( + skipCPUIfNoLapack, + # AssertionError: Tensor-likes are not close! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + # Exception: linalg.solve.triangular(); Only float is supported! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + # linalg.solve_triangular cannot be batched over because of a call to out.copy_(result); + supports_forward_ad=True, + ), + OpInfo( + "linalg.matrix_rank", + aten_name="linalg_matrix_rank", + dtypes=floating_and_complex_types(), + supports_autograd=False, + sample_inputs_func=sample_inputs_matrix_rank, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + # jit doesn't accept tensor inputs for matrix rank + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + dtypes=[torch.complex64, torch.float32], + ), + # NotImplementedError: The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + ), + ), + ), + OpInfo( + "linalg.matrix_rank", + aten_name="linalg_matrix_rank", + variant_test_name="hermitian", + dtypes=floating_and_complex_types(), + supports_autograd=False, + sample_inputs_func=sample_inputs_linalg_pinv_hermitian, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + # NotImplementedError: The operator 'aten::_linalg_eigh.eigenvalues' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + ), + ), + ), + OpInfo( + "linalg.pinv", + aten_name="linalg_pinv", + op=torch.linalg.pinv, + dtypes=floating_and_complex_types(), + # Runs very slowly on slow gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_linalg_pinv, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack], + skips=( + # errors with "leaked XXXX bytes CUDA memory on device 0" + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="cuda", + ), + ), + ), + OpInfo( + "linalg.pinv", + aten_name="linalg_pinv", + variant_test_name="singular", + # pinv is Frechet-differentiable in a rank-preserving neighborhood, + # so we feed inputs that are the products of two full-rank factors, + # to avoid any rank changes caused by the perturbations in the gradcheck + op=lambda a, b: torch.linalg.pinv(a @ b.mT), + dtypes=floating_and_complex_types(), + dtypesIfMPS=floating_types(), + supports_out=False, + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_linalg_pinv_singular, + # Only large tensors show issues with implicit backward used prior to + # explicit backward implementation. + decorators=[slowTest, skipCUDAIfNoCusolver, skipCPUIfNoLapack], + skips=( + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # CUDA runs out of memory + DecorateInfo( + unittest.skip("Skipped!"), + "TestFwdGradients", + "test_fn_fwgrad_bwgrad", + device_type="cuda", + dtypes=[torch.cdouble], + ), + # This test takes almost 2 hours to run! + DecorateInfo( + unittest.skip("Skipped!"), + "TestBwdGradients", + "test_fn_gradgrad", + device_type="cuda", + dtypes=[torch.cdouble], + ), + ), + ), + OpInfo( + "linalg.pinv", + aten_name="linalg_pinv", + variant_test_name="hermitian", + dtypes=floating_and_complex_types(), + check_batched_grad=False, + check_batched_gradgrad=False, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + sample_inputs_func=sample_inputs_linalg_pinv_hermitian, + gradcheck_wrapper=gradcheck_wrapper_hermitian_input, + decorators=[skipCUDAIfNoMagma, skipCPUIfNoLapack], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_out", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestJit", + "test_variant_consistency_jit", + device_type="mps", + dtypes=[torch.float32], + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-5, rtol=1e-5)}), + "TestCommon", + "test_noncontiguous_samples", + device_type="cuda", + ), + # This test is flaky under slow gradcheck, likely due to rounding issues + DecorateInfo( + skipIfSlowGradcheckEnv, + "TestFwdGradients", + "test_fn_fwgrad_bwgrad", + device_type="cuda", + ), + # NotImplementedError: The operator 'aten::_linalg_eigh.eigenvalues' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_requires_grad_error", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + ), + # see https://github.com/pytorch/pytorch/issues/177264 + DecorateInfo( + unittest.expectedFailure, + "TestEagerFusionOpInfo", + "test_aot_autograd_symbolic_exhaustive", + device_type="cpu", + dtypes=[torch.float32], + active_if=IS_ARM64 and IS_LINUX, + ), + ), + ), + OpInfo( + "linalg.svd", + op=torch.linalg.svd, + aten_name="linalg_svd", + decomp_aten_name="_linalg_svd", + dtypes=floating_and_complex_types(), + # Runs very slowly on slow-gradcheck - alternatively reduce input sizes + gradcheck_fast_mode=True, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + check_batched_forward_grad=False, + # We're using at::allclose, which does not have a batching rule + check_batched_grad=False, + check_batched_gradgrad=False, + sample_inputs_func=sample_inputs_svd, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack, with_tf32_off], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_no_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + ), + ), + OpInfo( + "linalg.svdvals", + op=torch.linalg.svdvals, + aten_name="linalg_svdvals", + decomp_aten_name="_linalg_svd", + dtypes=floating_and_complex_types(), + check_batched_forward_grad=False, + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + # We're using at::allclose, which does not have a batching rule + check_batched_gradgrad=False, + sample_inputs_func=sample_inputs_linalg_svdvals, + decorators=[skipCUDAIfNoMagmaAndNoCusolver, skipCPUIfNoLapack, with_tf32_off], + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestFakeTensor", + "test_fake_crossref_backward_no_amp", + device_type="cuda", + dtypes=[torch.float32], + active_if=TEST_WITH_ROCM, + ), + # The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + OpInfo( + "linalg.tensorinv", + ref=np.linalg.tensorinv, + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_tensorinv, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + # See https://github.com/pytorch/pytorch/pull/78358 + check_batched_forward_grad=False, + decorators=[skipCPUIfNoLapack, skipCUDAIfNoMagmaAndNoCusolver], + skips=( + DecorateInfo( + unittest.skip("Unsupported on MPS for now"), + "TestCommon", + "test_numpy_ref_mps", + ), + # Exception: linalg_inv: not supported for complex types yet! + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), + OpInfo( + "linalg.tensorsolve", + ref=lambda a, b, dims=None: np.linalg.tensorsolve(a, b, axes=dims), + dtypes=floating_and_complex_types(), + sample_inputs_func=sample_inputs_tensorsolve, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=[ + skipCUDAIfNoMagmaAndNoCusolver, + skipCPUIfNoLapack, + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=1e-03, rtol=1e-03)}), + "TestCommon", + "test_noncontiguous_samples", + device_type="cuda", + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=8e-04, rtol=7e-06)}), + "TestCommon", + "test_noncontiguous_samples", + device_type="cpu", + ), + DecorateInfo( + toleranceOverride({torch.float32: tol(atol=2e-04, rtol=3e-06)}), + "TestConsistency", + "test_output_match", + device_type="mps", + ), + ], + skips=( + # Exception: linalg.lu_factor(): MPS only supports floats. + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + device_type="mps", + dtypes=(torch.complex64,), + ), + ), + ), +] + +python_ref_db: list[OpInfo] = [ + # + # torch.linalg + # + PythonRefInfo( + "_refs.linalg.cross", + torch_opinfo_name="linalg.cross", + supports_out=True, + op_db=op_db, + skips=( + # TODO: is this really needed? + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_python_ref_errors" + ), + ), + ), + PythonRefInfo( + "_refs.linalg.diagonal", + torch_opinfo_name="linalg.diagonal", + supports_out=False, + op_db=op_db, + ), + PythonRefInfo( + "_refs.linalg.vecdot", + torch_opinfo_name="linalg.vecdot", + op_db=op_db, + ), + ReductionPythonRefInfo( + "_refs.linalg.vector_norm", + torch_opinfo_name="linalg.vector_norm", + supports_out=True, + op_db=op_db, + ), + PythonRefInfo( + "_refs.linalg.matrix_norm", + torch_opinfo_name="linalg.matrix_norm", + supports_out=True, + # Uses vector_norm inside and vector_norm is affected by + # https://github.com/pytorch/pytorch/issues/77216 + validate_view_consistency=False, + op_db=op_db, + skips=( + # NotImplementedError: The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref_torch_fallback", + device_type="mps", + dtypes=(torch.float32, torch.complex64), + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref", + device_type="mps", + dtypes=(torch.float32, torch.complex64), + ), + ), + ), + PythonRefInfo( + "_refs.linalg.norm", + torch_opinfo_name="linalg.norm", + supports_out=True, + # Uses vector_norm inside and vector_norm is affected by + # https://github.com/pytorch/pytorch/issues/77216 + validate_view_consistency=False, + op_db=op_db, + skips=( + # NotImplementedError: The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref", + device_type="mps", + dtypes=(torch.float32, torch.complex64), + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref_torch_fallback", + device_type="mps", + dtypes=(torch.float32, torch.complex64), + ), + ), + ), + PythonRefInfo( + "_refs.linalg.svd", + torch_opinfo_name="linalg.svd", + supports_out=True, + op_db=op_db, + ), + PythonRefInfo( + "_refs.linalg.svdvals", + torch_opinfo_name="linalg.svdvals", + supports_out=True, + op_db=op_db, + skips=( + # The operator 'aten::_linalg_svd.U' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref_torch_fallback", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref", + device_type="mps", + ), + ), + ), +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/nested.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/nested.py new file mode 100644 index 0000000000000000000000000000000000000000..bdb4d942aba065e0d71233e6688ad1382bfc2219 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/nested.py @@ -0,0 +1,1613 @@ +# mypy: ignore-errors + +import math +from copy import copy +from dataclasses import dataclass +from functools import partial + +import torch +from torch.fx.experimental.symbolic_shapes import is_nested_int +from torch.testing._internal.common_methods_invocations import op_db +from torch.testing._internal.opinfo.core import ( + BinaryUfuncInfo, + ReductionOpInfo, + SampleInput, + UnaryUfuncInfo, +) +from torch.utils._pytree import tree_flatten, tree_map + + +@dataclass +class ExtraOpData: + """ + Contains info on top of the typical OpInfo data that is useful for NJT test generation. + + The process that converts the standard op_db -> an NJT-compatible op_db will attach this + data onto each associated OpInfo entry. + """ + + # Indicates whether the associated op is a view op + is_view: bool = False + + # Specifies the names of any dim-related args that the op takes in. This is useful + # for NJT tests because there is often asymmetry across the supported set of dims for + # an op; it may make sense to operate over the batch dim but not the ragged dim, for + # example. The length of this list should match the number of relevant overloads. + # Each list item of the outer list should specify dim argnames. Ellipses should be used + # to indicate multi-dim support for a given overload. + # + # For example, squeeze() has both a dim and multi-dim overload, where the argname for + # each is simply "dim". Its entry should be: [["dim"], ["dim..."]]. + # + # If no overload of the op accepts dim-related args, this should be None. + dim_args: list[list[str]] = None + + # Helper function to extract names of dim-related args. + # Returns: tuple of (single dim argname if available, dim list argname if available) + # If the op doesn't support dim-related args at all OR this op only has overloads + # with multiple dim args (e.g. transpose()), then this returns (None, None). + def get_dim_argnames(self) -> tuple[str | None, str | None]: + if self.dim_args is None: + return (None, None) + + # name for the dim arg that supports a single dim + single_dim_argname = None + # name for the dim arg that supports a list of dims + dimlist_argname = None + for overload in self.dim_args: + # only consider overloads with a single dim-related arg + if len(overload) != 1: + continue + if overload[0].endswith("..."): + dimlist_argname = overload[0].replace("...", "") + if single_dim_argname is None: + single_dim_argname = dimlist_argname + else: + single_dim_argname = overload[0] + return (single_dim_argname, dimlist_argname) + + +# Mapping of OpInfo full names -> extra data to tack onto the OpInfo entry for use +# in test generation. +extra_op_data = { + "_segment_reduce.lengths": ExtraOpData(dim_args=[["axis0"]]), + "_segment_reduce.offsets": ExtraOpData(dim_args=[["axis0"]]), + "all": ExtraOpData(dim_args=[["dim"], ["dim..."]]), + "argmax": ExtraOpData(dim_args=[["dim"]]), + "argmin": ExtraOpData(dim_args=[["dim"]]), + "amax": ExtraOpData(dim_args=[["dim..."]]), + "amin": ExtraOpData(dim_args=[["dim..."]]), + "any": ExtraOpData(dim_args=[["dim"], ["dim..."]]), + "argsort": ExtraOpData(dim_args=[["dim"]]), + "broadcast_to": ExtraOpData(is_view=True), + "cat": ExtraOpData(dim_args=[["dim"]]), + "chunk": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "conj": ExtraOpData(is_view=True), + "contiguous": ExtraOpData(is_view=True), + "count_nonzero": ExtraOpData(dim_args=[["dim"], ["dim..."]]), + "cummax": ExtraOpData(dim_args=[["dim"]]), + "cummin": ExtraOpData(dim_args=[["dim"]]), + "cumprod": ExtraOpData(dim_args=[["dim"]]), + "cumsum": ExtraOpData(dim_args=[["dim"]]), + "cumulative_trapezoid": ExtraOpData(dim_args=[["dim"]]), + "diag_embed": ExtraOpData(dim_args=[["dim1", "dim2"]]), + "diagonal": ExtraOpData(is_view=True, dim_args=[["dim1", "dim2"]]), + "diagonal_copy": ExtraOpData(dim_args=[["dim1", "dim2"]]), + "diagonal_scatter": ExtraOpData(dim_args=[["dim1", "dim2"]]), + "diff": ExtraOpData(dim_args=[["dim"]]), + "expand": ExtraOpData(is_view=True), + "expand_as": ExtraOpData(is_view=True), + "fft.fft": ExtraOpData(dim_args=[["dim"]]), + "fft.hfft": ExtraOpData(dim_args=[["dim"]]), + "fft.ifft": ExtraOpData(dim_args=[["dim"]]), + "fft.ihfft": ExtraOpData(dim_args=[["dim"]]), + "fft.irfft": ExtraOpData(dim_args=[["dim"]]), + "fft.rfft": ExtraOpData(dim_args=[["dim"]]), + "flatten": ExtraOpData(is_view=True, dim_args=[["start_dim", "end_dim"]]), + "flip": ExtraOpData(dim_args=[["dims..."]]), + "gather": ExtraOpData(dim_args=[["dim"]]), + "hash_tensor": ExtraOpData(dim_args=[["dim..."]]), + "imag": ExtraOpData(is_view=True), + "index_add": ExtraOpData(dim_args=[["dim"]]), + "index_copy": ExtraOpData(dim_args=[["dim"]]), + "index_fill": ExtraOpData(dim_args=[["dim"]]), + "index_reduce.amax": ExtraOpData(dim_args=[["dim"]]), + "index_reduce.amin": ExtraOpData(dim_args=[["dim"]]), + "index_reduce.mean": ExtraOpData(dim_args=[["dim"]]), + "index_reduce.prod": ExtraOpData(dim_args=[["dim"]]), + "index_select": ExtraOpData(dim_args=[["dim"]]), + "kthvalue": ExtraOpData(dim_args=[["dim"]]), + "linalg.cross": ExtraOpData(dim_args=[["dim"]]), + "linalg.diagonal": ExtraOpData(is_view=True, dim_args=[["dim1", "dim2"]]), + "linalg.tensorsolve": ExtraOpData(dim_args=[["dims..."]]), + "linalg.vecdot": ExtraOpData(dim_args=[["dim"]]), + "linalg.vector_norm": ExtraOpData(dim_args=[["dim..."]]), + "log_softmax": ExtraOpData(dim_args=[["dim"]]), + "logcumsumexp": ExtraOpData(dim_args=[["dim"]]), + "masked.amax": ExtraOpData(dim_args=[["dim"]]), + "masked.amin": ExtraOpData(dim_args=[["dim"]]), + "masked.argmax": ExtraOpData(dim_args=[["dim"]]), + "masked.argmin": ExtraOpData(dim_args=[["dim"]]), + "masked.logsumexp": ExtraOpData(dim_args=[["dim"]]), + "masked.mean": ExtraOpData(dim_args=[["dim"]]), + "masked.norm": ExtraOpData(dim_args=[["dim"]]), + "masked.prod": ExtraOpData(dim_args=[["dim"]]), + "masked.std": ExtraOpData(dim_args=[["dim"]]), + "masked.sum": ExtraOpData(dim_args=[["dim"]]), + "masked.var": ExtraOpData(dim_args=[["dim"]]), + "max.reduction_with_dim": ExtraOpData(dim_args=[["dim"]]), + "median": ExtraOpData(dim_args=[["dim"]]), + "mean": ExtraOpData(dim_args=[["dim..."]]), + "min.reduction_with_dim": ExtraOpData(dim_args=[["dim"]]), + "mode": ExtraOpData(dim_args=[["dim"]]), + "movedim": ExtraOpData( + dim_args=[["source", "destination"], ["source...", "destination..."]] + ), + "nanmean": ExtraOpData(dim_args=[["dim..."]]), + "nanmedian": ExtraOpData(dim_args=[["dim"]]), + "nansum": ExtraOpData(dim_args=[["dim..."]]), + "narrow": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "narrow_copy": ExtraOpData(dim_args=[["dim"]]), + "nn.functional.cosine_similarity": ExtraOpData(dim_args=[["dim"]]), + "nn.functional.glu": ExtraOpData(dim_args=[["dim"]]), + "permute": ExtraOpData(is_view=True, dim_args=[["dims..."]]), + "positive": ExtraOpData(is_view=True), + "prod": ExtraOpData(dim_args=[["dim"]]), + "ravel": ExtraOpData(is_view=True), + "real": ExtraOpData(is_view=True), + "renorm": ExtraOpData(dim_args=[["dim"]]), + "reshape": ExtraOpData(is_view=True), + "reshape_as": ExtraOpData(is_view=True), + "roll": ExtraOpData(dim_args=[["dims..."]]), + "rot90": ExtraOpData(dim_args=[["dims..."]]), + "scatter": ExtraOpData(dim_args=[["dim"]]), + "scatter_add": ExtraOpData(dim_args=[["dim"]]), + "scatter_reduce.amax": ExtraOpData(dim_args=[["dim"]]), + "scatter_reduce.amin": ExtraOpData(dim_args=[["dim"]]), + "scatter_reduce.mean": ExtraOpData(dim_args=[["dim"]]), + "scatter_reduce.prod": ExtraOpData(dim_args=[["dim"]]), + "scatter_reduce.sum": ExtraOpData(dim_args=[["dim"]]), + "select": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "select_scatter": ExtraOpData(dim_args=[["dim"]]), + "slice": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "slice_scatter": ExtraOpData(dim_args=[["dim"]]), + "softmax": ExtraOpData(dim_args=[["dim"]]), + "sort": ExtraOpData(dim_args=[["dim"]]), + "split": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "split_with_sizes": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "split_with_sizes_copy": ExtraOpData(dim_args=[["dim"]]), + "squeeze": ExtraOpData(is_view=True, dim_args=[["dim"], ["dim..."]]), + "squeeze_copy": ExtraOpData(dim_args=[["dim"], ["dim..."]]), + "stack": ExtraOpData(dim_args=[["dim"]]), + "std": ExtraOpData(dim_args=[["dim..."]]), + "std.unbiased": ExtraOpData(dim_args=[["dim..."]]), + "sum": ExtraOpData(dim_args=[["dim..."]]), + "t": ExtraOpData(is_view=True), + "tensor_split": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "tensordot": ExtraOpData(dim_args=[["dims..."]]), + "tile": ExtraOpData(dim_args=[["dims..."]]), + "topk": ExtraOpData(dim_args=[["dim"]]), + "transpose": ExtraOpData(is_view=True, dim_args=[["dim0", "dim1"]]), + "transpose_copy": ExtraOpData(dim_args=[["dim0", "dim1"]]), + "trapezoid": ExtraOpData(dim_args=[["dim"]]), + "trapz": ExtraOpData(dim_args=[["dim"]]), + "unbind": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "unflatten": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "unfold": ExtraOpData(is_view=True, dim_args=[["dimension"]]), + "unfold_copy": ExtraOpData(dim_args=[["dimension"]]), + "unsafe_chunk": ExtraOpData(dim_args=[["dim"]]), + "unsafe_split": ExtraOpData(dim_args=[["dim"]]), + "unsqueeze": ExtraOpData(is_view=True, dim_args=[["dim"]]), + "unsqueeze_copy": ExtraOpData(dim_args=[["dim"]]), + "var": ExtraOpData(dim_args=[["dim..."]]), + "var.unbiased": ExtraOpData(dim_args=[["dim..."]]), + "view": ExtraOpData(is_view=True), + "view_as": ExtraOpData(is_view=True), + "view_as_complex": ExtraOpData(is_view=True), + "view_as_real": ExtraOpData(is_view=True), +} + + +# random integer used for sizes +def _rnd(): + return torch.randint(3, 8, ()).item() + + +def _raggedness_matches(nt1, nt2): + return ( + nt1.is_nested + and nt2.is_nested + and nt1._ragged_idx == nt2._ragged_idx + and nt1.shape[nt1._ragged_idx] == nt2.shape[nt2._ragged_idx] + ) + + +# Helper function to avoid reusing the exact same tensor / NJT across SampleInputs, +# as this causes autograd problems. +def _clone(t): + requires_grad = t.requires_grad + return t.detach().clone().requires_grad_(requires_grad) + + +# Helper function to update a sample with new kwargs / name +def _update_sample(sample, new_kwargs): + all_kwargs = dict(sample.kwargs) + all_kwargs.update(new_kwargs) + full_name = ", ".join([sample.name, *(f"{k}={v}" for (k, v) in new_kwargs.items())]) + return SampleInput( + _clone(sample.input), + args=sample.args, + kwargs=all_kwargs, + name=full_name, + ) + + +# Generates a random NT. +# dims should be something like [5, None, 10], with None indicating that a +# random ragged structure should be used +def random_nt_from_dims( + dims, device=None, dtype=None, layout=torch.strided, requires_grad=False +): + sizes = [[d if d is not None else _rnd() for d in dims[1:]] for d in range(dims[0])] + return torch.nested.nested_tensor( + [torch.randn(*size) for size in sizes], + device=device, + dtype=dtype, + layout=layout, + requires_grad=requires_grad, + ) + + +# Helper function to get a reasonable string representation of an NJT for use in +# SampleInput names. +def _describe_njt(njt) -> str: + contig_type = "_contig" if njt.is_contiguous() else "_noncontig" + if njt._lengths is not None and njt._offsets is not None: + contig_type += "_holes" + elif njt._ragged_idx != 1: + contig_type += "_transposed" + + cached_data = "_without_seqlen_cache" + if njt._max_seqlen_tensor is not None: + cached_data = "_with_seqlen_cache" + + return f"{njt.dim()}D{contig_type}{cached_data}" + + +# Helper function to get a reasonable string representation of a given dim wrt an NJT. +def _describe_dim(njt, dim): + if dim == 0: + return "batch_dim" + elif dim == njt._ragged_idx: + return "ragged_dim" + return "normal_dim" + + +# Helper function for generating a comprehensive set of NJT sample inputs. +def _sample_njts(device, dtype, requires_grad=False, dims=None): + if dims is None: + dims = [2, 3, 4] + if not isinstance(dims, (list, tuple)): + dims = [dims] + + # contiguous NJTs + for dim in dims: + # with min / max seqlen cached + shape = (_rnd(), None, *[_rnd() for _ in range(dim - 2)]) + nt = random_nt_from_dims( + shape, + device=device, + dtype=dtype, + requires_grad=requires_grad, + layout=torch.jagged, + ) + yield nt + + # without min / max seqlen cached + values = _clone(nt.values()) + offsets = _clone(nt.offsets()) + yield torch.nested.nested_tensor_from_jagged(values, offsets).requires_grad_( + requires_grad + ) + + # non-contiguous transposed NJT (not possible for 2D) + if dim > 2: + yield nt.transpose(-1, nt._ragged_idx) + + # non-contiguous with holes NJT + values = _clone(nt.values()) + offsets = _clone(nt.offsets()) + # subtract 1 to cause holes + lengths = _clone(offsets.diff() - 1) + yield torch.nested.nested_tensor_from_jagged( + values=values, + offsets=offsets, + lengths=lengths, + ).requires_grad_(requires_grad) + + +# Computes an unbind-based reference for a given OpInfo on a given SampleInput. +# This reference unbinds the input NJT and invokes the op on each of the components, +# optionally wrapping the result in an NJT. +def unbind_reference(op, sample, wrap_output_as_njt=True): + # first NJT in the arglist determines expected ragged structure + nt_inp = ( + sample.input + if sample.input.is_nested + # TODO: look in kwargs too? + else next(a for a in sample.args if a.is_nested) + ) + + out_ref_components = [] + for i in range(nt_inp.shape[0]): + + def _slice_input(t, i=i, inp=nt_inp): + # any NJT with the same ragged structure as the input should + # be sliced to pass to the reference + if isinstance(t, torch.Tensor) and _raggedness_matches(t, inp): + return t[i] + # allow the SampleInput to tell us how to slice it for ref calculation + elif isinstance(t, torch.Tensor) and hasattr(t, "_batch_dim"): + bdim = t._batch_dim # type: ignore[attr] + if t.shape[bdim] == 1: + return t[0] + else: + return t.select(bdim, i) + else: + return t + + inp = _slice_input(sample.input) + args = tree_map(_slice_input, sample.args) + kwargs = tree_map(_slice_input, sample.kwargs) + + # Handle indices in index_put + if "index_put" in op.full_name and "indices" in kwargs: + if len(kwargs["indices"]) > 1: + # If after unrolling we still have indices left, use them + kwargs["indices"] = [t[i] for t in kwargs["indices"][1:]] + else: + # If no indices are left, create them so they match the NJT implementation + sequence_put = kwargs["indices"][0].tolist() + if i in sequence_put: + kwargs["indices"] = [ + torch.tensor( + list(range(inp.shape[0])), + dtype=torch.int32, + device=kwargs["indices"][0].device, + ) + ] + else: + kwargs["indices"] = [ + torch.tensor( + [], dtype=torch.int32, device=kwargs["indices"][0].device + ) + ] + + from torch.nested._internal.ops import _outer_to_inner_dim + + # Need to adjust dims to apply on NJT component + if op._extra_op_data.dim_args is not None: + # get all possible dim-related argnames that could be encountered for this op + argnames = tree_map( + lambda a: a.replace("...", ""), + tree_flatten(op._extra_op_data.dim_args)[0], + ) + # for all dim-related args present, convert from outer -> inner dim space + for argname in {a for a in argnames if a in kwargs}: + # allow the SampleInput to tell us how to canonicalize the dim kwargs + ndim = nt_inp._ndim if hasattr(nt_inp, "_ndim") else nt_inp.dim() + kwargs[argname] = _outer_to_inner_dim( + ndim, kwargs[argname], nt_inp._ragged_idx, canonicalize=True + ) + + out_ref_component = op.op(inp, *args, **kwargs) + out_ref_components.append(out_ref_component) + + if wrap_output_as_njt: + # handle list / tuple of outputs + if len(out_ref_components) > 0 and isinstance( + out_ref_components[0], (list, tuple) + ): + num_returns = len(out_ref_components[0]) + # ensure we get the same number of returns for each invocation + if not all(len(o) == num_returns for o in out_ref_components): + raise AssertionError( + f"Expected all outputs to have {num_returns} returns" + ) + # construct NJTs from same index returns from each invocation + njt_returns = [ + torch.nested.as_nested_tensor( + [o[r] for o in out_ref_components], layout=torch.jagged + ) + for r in range(num_returns) + ] + return type(out_ref_components[0])(njt_returns) + return torch.nested.as_nested_tensor(out_ref_components, layout=torch.jagged) + + return out_ref_components + + +# Computes the reference value for a non-reduction unary op with dim-wise application. +def unary_dimwise_reference(op, sample, batchwise_reference=None): + # extract info about the dim args this op supports + if op._extra_op_data.dim_args is None: + raise AssertionError("Expected op._extra_op_data.dim_args to not be None") + single_dim_argname, dimlist_argname = op._extra_op_data.get_dim_argnames() + # only support a single non-list dim arg for now + if dimlist_argname is not None: + raise AssertionError("Expected dimlist_argname to be None") + if single_dim_argname is None: + raise AssertionError("Expected single_dim_argname to not be None") + if sample.kwargs[single_dim_argname] == 0: + # unbind reference won't work for batch-wise operation; handle this case here + if batchwise_reference is None: + raise AssertionError("Expected batchwise_reference to not be None") + return batchwise_reference(op, sample) + return unbind_reference(op, sample) + + +# Computes the reference value for a reduction op. +def reduction_reference(op, sample): + if not sample.input.is_nested: + raise AssertionError("Expected sample.input.is_nested to be True") + + # extract info about the dim args this op supports + if op._extra_op_data.dim_args is None: + raise AssertionError("Expected op._extra_op_data.dim_args to not be None") + single_dim_argname, dimlist_argname = op._extra_op_data.get_dim_argnames() + if single_dim_argname is None: + raise AssertionError("Expected single_dim_argname to not be None") + + dim = sample.kwargs.get( + dimlist_argname, sample.kwargs.get(single_dim_argname, None) + ) + keepdim = sample.kwargs.get("keepdim", False) + if dim == 0: + raise AssertionError("reductions over just the batch dim are not supported") + if isinstance(dim, (tuple, list)): + reduce_on_ragged = sample.input._ragged_idx in dim + reduce_on_batch = 0 in dim + else: + reduce_on_ragged = sample.input._ragged_idx == dim + reduce_on_batch = dim == 0 + + if dim is None: + # calculate reference value by running reduction on values buffer + return op.op(sample.input.values(), *sample.args, **sample.kwargs) + + if reduce_on_ragged and reduce_on_batch: + # run reference directly on buffer with dims converted to inner space + from torch.nested._internal.ops import _outer_to_inner_dim + + ref_kwargs = dict(sample.kwargs) + if dimlist_argname is None: + raise AssertionError("Expected dimlist_argname to not be None") + ref_kwargs[dimlist_argname] = _outer_to_inner_dim( + sample.input.dim(), dim, sample.input._ragged_idx, canonicalize=True + ) + out = op.op(sample.input.values(), *sample.args, **ref_kwargs) + if keepdim: + if isinstance(out, (tuple, list)): + # some ops return multiple things; unsqueeze all of them + out = type(out)(o.unsqueeze(0) for o in out) + else: + out = out.unsqueeze(0) + return out + + if reduce_on_ragged and not reduce_on_batch: + # calculate reference value by running an unbind reference and stacking + out_ref_components = unbind_reference(op, sample, wrap_output_as_njt=False) + if len(out_ref_components) > 0 and isinstance( + out_ref_components[0], (tuple, list) + ): + # some ops return multiple things; stack all of them + num_returns = len(out_ref_components[0]) + # ensure we get the same number of returns for each invocation + if not all(len(o) == num_returns for o in out_ref_components): + raise AssertionError( + f"Expected all outputs to have {num_returns} returns" + ) + # stack same index returns from each invocation + stacked_returns = [ + torch.stack([o[r] for o in out_ref_components], dim=0) + for r in range(num_returns) + ] + return type(out_ref_components[0])(stacked_returns) + return torch.stack(out_ref_components, dim=0) + + # unbind reference works for other reductions + return unbind_reference(op, sample) + + +def sample_inputs_elementwise_njt_unary( + op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs +): + if not op_kwargs: + op_kwargs = {} + + for njt in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[2, 3, 4] + ): + yield SampleInput(njt, kwargs=dict(op_kwargs), name=_describe_njt(njt)) + + +def sample_inputs_elementwise_njt_binary( + op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs +): + if not op_kwargs: + op_kwargs = {} + + for njt1 in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[2, 3, 4] + ): + njt_desc = _describe_njt(njt1) + njt2 = torch.randn_like(njt1) + yield SampleInput( + _clone(njt1), + args=(njt2,), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (NT, NT)", + ) + + # broadcasting case: (B, j0, ...) with (B, 1, ...) + dense_shape = list(njt1.shape) + dense_shape[njt1._ragged_idx] = 1 + t = torch.randn( + dense_shape, + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + t2 = _clone(t) + # used for slicing in unbind_reference() + t._batch_dim = 0 + t2._batch_dim = 0 + # (NT, T) + yield SampleInput( + _clone(njt1), + args=(t,), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (NT, T) broadcasting 1 over ragged", + ) + # (T, NT) + yield SampleInput( + t2, + args=(_clone(njt1),), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (T, NT) broadcasting 1 over ragged", + ) + + # broadcasting case: (B, j0, ...) with (1, 1...) + t = torch.randn( + [1 for _ in range(njt1.dim())], + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + t2 = _clone(t) + # used for slicing in unbind_reference() + t._batch_dim = 0 + t2._batch_dim = 0 + # (NT, T) + yield SampleInput( + _clone(njt1), + args=(t,), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (NT, T) broadcasting all 1s", + ) + # (T, NT) + yield SampleInput( + t2, + args=(_clone(njt1),), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (T, NT) broadcasting all 1s", + ) + + # broadcasting case: (B, j0, ...) with (...) + if njt1.dim() > njt1._ragged_idx + 1: + t = torch.randn( + njt1.shape[njt1._ragged_idx + 1 :], + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + # (NT, T) + yield SampleInput( + _clone(njt1), + args=(_clone(t),), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (NT, T) broadcasting normal dims", + ) + # (T, NT) + yield SampleInput( + _clone(t), + args=(_clone(njt1),), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (T, NT) broadcasting normal dims", + ) + + # broadcasting case: (B, j0, ...) with scalar + t = torch.randn((), device=device, dtype=dtype, requires_grad=requires_grad) + # (NT, T) + yield SampleInput( + _clone(njt1), + args=(_clone(t),), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (NT, T) broadcasting with scalar", + ) + # (T, NT) + yield SampleInput( + _clone(t), + args=(_clone(njt1),), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (T, NT) broadcasting with scalar", + ) + + # mixed broadcasting case: (B, j0, 1) with (B, 1, D) + B = 4 + D = 16 + njt = random_nt_from_dims( + (B, None, 1), + device=device, + dtype=dtype, + requires_grad=requires_grad, + layout=torch.jagged, + ) + njt_desc = _describe_njt(njt) + t = torch.randn(B, 1, D, device=device, dtype=dtype, requires_grad=requires_grad) + t2 = _clone(t) + # used for slicing in unbind_reference() + t._batch_dim = 0 + t2._batch_dim = 0 + + # (NT, T) + yield SampleInput( + _clone(njt), + args=(t,), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (NT, T) mixed broadcasting", + ) + # (T, NT) + yield SampleInput( + t2, + args=(_clone(njt),), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: (T, NT) mixed broadcasting", + ) + + +def sample_inputs_njt_reduction( + op_info, + device, + dtype, + requires_grad, + supports_keepdim=True, + op_kwargs=None, + **kwargs, +): + if not op_kwargs: + op_kwargs = {} + + # extract info about the dim args this op supports + if op_info._extra_op_data.dim_args is None: + raise AssertionError("Expected op_info._extra_op_data.dim_args to not be None") + ( + single_dim_argname, + dimlist_argname, + ) = op_info._extra_op_data.get_dim_argnames() + if single_dim_argname is None: + raise AssertionError("Expected single_dim_argname to not be None") + supports_dimlist = dimlist_argname is not None + + for njt in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[2, 3, 4] + ): + njt_desc = _describe_njt(njt) + keepdim_values = [False, True] if supports_keepdim else [None] + for keepdim in keepdim_values: + keepdim_suffix = f" with keepdim={keepdim}" if supports_keepdim else "" + # single dim-wise reduction; includes reduction over the ragged dim + # NB: reduction over the batch dim is not supported! + # TODO: Cover this in the set of error inputs + for dim in range(1, njt.dim()): + dim_desc = "normal" if dim != njt._ragged_idx else "ragged" + yield SampleInput( + _clone(njt), + kwargs={ + **op_kwargs, + single_dim_argname: dim, + **({"keepdim": keepdim} if supports_keepdim else {}), + }, + name=f"{njt_desc}: {dim_desc} dim reduction{keepdim_suffix}", + ) + + if supports_dimlist: + # reduce on both batch and ragged dims + yield SampleInput( + _clone(njt), + kwargs={ + **op_kwargs, + dimlist_argname: [0, njt._ragged_idx], + **({"keepdim": keepdim} if supports_keepdim else {}), + }, + name=f"{njt_desc}: batch+ragged reduction{keepdim_suffix}", + ) + + # reduce on batch, ragged, and other dims + for other_dim in range(njt._ragged_idx + 1, njt.dim()): + yield SampleInput( + _clone(njt), + kwargs={ + **op_kwargs, + dimlist_argname: [0, njt._ragged_idx, other_dim], + **({"keepdim": keepdim} if supports_keepdim else {}), + }, + name=( + f"{njt_desc}: batch+ragged+dim={other_dim} " + f"reduction{keepdim_suffix}" + ), + ) + + # reduce on two non-ragged, non-batch dims + if njt.dim() > 3 and njt._ragged_idx == 1: + yield SampleInput( + _clone(njt), + kwargs={ + **op_kwargs, + dimlist_argname: [njt.dim() - 2, njt.dim() - 1], + **({"keepdim": keepdim} if supports_keepdim else {}), + }, + name=f"{njt_desc}: two normal dim reduction{keepdim_suffix}", + ) + + # full reduction by specifying all dims + yield SampleInput( + _clone(njt), + kwargs={ + **op_kwargs, + dimlist_argname: list(range(njt.dim())), + **({"keepdim": keepdim} if supports_keepdim else {}), + }, + name=f"{njt_desc}: all dim reduction{keepdim_suffix}", + ) + + # TODO: Reducing on ragged dim and non-batch dim is not supported; + # cover this in the set of error inputs. + + # full reduction + yield SampleInput( + _clone(njt), + kwargs=dict(op_kwargs), + name=f"{njt_desc}: full reduction with keepdim={keepdim}", + ) + + +def unsupported_sample_inputs_func(op_name): + def _f(op_info, device, dtype, requires_grad, op_name=op_name, **kwargs): + raise RuntimeError( + f"OpInfo for {op_name} does not support NJT. Support can be added by modifying " + "torch/testing/_internal/opinfo/definitions/nested.py." + ) + + return _f + + +def unsupported_reference(op_name): + def _f(op, sample): + raise RuntimeError( + f"OpInfo for {op_name} does not define a ref() function. Support can be added by " + "modifying torch/testing/_internal/opinfo/definitions/nested.py." + ) + + return _f + + +# === BEGIN OP-SPECIFIC SAMPLE INPUTS FUNCS / REFERENCES === +def sample_inputs_unary_dimwise( + op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs +): + if op_kwargs is None: + op_kwargs = {} + + # only support a single non-list dim arg for now + if op_info._extra_op_data is None: + raise AssertionError("Expected op_info._extra_op_data to not be None") + single_dim_argname, dimlist_argname = op_info._extra_op_data.get_dim_argnames() + if single_dim_argname is None: + raise AssertionError("Expected single_dim_argname to not be None") + if dimlist_argname is not None: + raise AssertionError("Expected dimlist_argname to be None") + + for njt in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[2, 3, 4] + ): + for dim in range(njt.dim()): + kwargs = {single_dim_argname: dim} + kwargs.update(op_kwargs) + yield SampleInput( + _clone(njt), + kwargs=kwargs, + name=f"{_describe_njt(njt)}: {_describe_dim(njt, dim)}", + ) + + +def batchwise_reference_chunk(op, sample): + # reference for chunk() over dim=0 + B = sample.input.size(0) + num_chunks = sample.kwargs["chunks"] + chunk_size = math.ceil(B / num_chunks) + num_full_chunks = B // chunk_size + chunk_sizes = [chunk_size for _ in range(num_full_chunks)] + if B % chunk_size != 0: + # final chunk contains the leftovers + chunk_sizes.append(B % chunk_size) + + # split unbound components into chunks according to calculated sizes + components = list(sample.input.unbind()) + start = 0 + chunks = [] + for chunk_size in chunk_sizes: + chunks.append(components[start : start + chunk_size]) + start += chunk_size + + # rejoin into NJT outputs + return [torch.nested.as_nested_tensor(lst, layout=torch.jagged) for lst in chunks] + + +def batchwise_reference_narrow(op, sample): + # TODO: write this! + raise NotImplementedError + + +def batchwise_reference_select(op, sample): + # reference for select() over dim=0 + return sample.input.unbind()[sample.kwargs["index"]] + + +def batchwise_reference_split(op, sample): + # TODO: write this! + raise NotImplementedError + + +def batchwise_reference_split_with_sizes(op, sample): + # TODO: write this! + raise NotImplementedError + + +def batchwise_reference_unflatten(op, sample): + # TODO: write this! + raise NotImplementedError + + +def batchwise_reference_unsqueeze(op, sample): + raise ValueError("unsqueeze() is not intended to operate on the batch dim") + + +def sample_inputs_clone(op_info, device, dtype, requires_grad, **kwargs): + # non-contiguous NJTs + for njt in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[2, 3, 4] + ): + yield SampleInput(njt, name=_describe_njt(njt)) + + for memory_format in (torch.contiguous_format, torch.preserve_format): + # construct a "non-contiguous with holes" NJT + values = torch.randn( + 10, 5, device=device, dtype=dtype, requires_grad=requires_grad + ) + offsets = torch.tensor([0, 2, 4, 10], device=device, dtype=torch.int64) + lengths = torch.tensor([2, 1, 3], device=device, dtype=torch.int64) + njt = torch.nested.nested_tensor_from_jagged( + values, offsets=offsets, lengths=lengths + ) + + njt_desc = _describe_njt(njt) + yield SampleInput( + njt, + kwargs={"memory_format": memory_format}, + name=f"{njt_desc}: {memory_format})", + ) + + +def sample_inputs_fill(op_info, device, dtype, requires_grad, **kwargs): + # scalar case + unary_func = partial(sample_inputs_elementwise_njt_unary, op_kwargs={"value": 42.0}) + yield from unary_func(op_info, device, dtype, requires_grad) + + # TODO: add Tensor case + + +def sample_inputs_mvl_gamma(p): + return partial(sample_inputs_elementwise_njt_unary, op_kwargs={"p": p}) + + +def sample_inputs_polygamma_n(n): + return partial(sample_inputs_elementwise_njt_unary, op_kwargs={"n": n}) + + +def sample_inputs_special_polygamma_n(n): + return partial(sample_inputs_elementwise_njt_unary, op_kwargs={"n": n}) + + +def sample_inputs_to(op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs): + for njt in _sample_njts( + device=device, + dtype=dtype, + requires_grad=requires_grad, + dims=[2, 3, 4], + ): + other_dtypes = ( + d for d in (torch.float32, torch.half, torch.double) if d is not dtype + ) + for other_dtype in other_dtypes: + sample_name = f"{njt.dim()}D: {dtype} -> {other_dtype}" + yield SampleInput(_clone(njt), kwargs={"dtype": dtype}, name=sample_name) + + # only include device transfer for CUDA inputs + if "cuda" in device: + other_device = "cpu" + sample_name = f"{_describe_njt(njt)}: {device} -> {other_device}" + yield SampleInput( + _clone(njt), kwargs={"device": other_device}, name=sample_name + ) + + +def sample_inputs_bmm(op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs): + for njt_3d in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[3] + ): + # (B, j1, D) x (B, D, E) => (B, j1, E) + if njt_3d._ragged_idx == 1: + B, D = njt_3d.shape[0], njt_3d.shape[-1] + E = D + 2 + other = torch.randn(B, D, E, device=device, dtype=dtype) + # used for slicing in unbind_reference() + other._batch_dim = 0 + njt_desc = _describe_njt(njt_3d) + yield SampleInput( + _clone(njt_3d), + kwargs={"mat2": other}, + name=f"{njt_desc}: (B, j, D) x (B, D, E)", + ) + + # TODO (need factory functions): + # (B, D, j1) x (B, j1, E) => (B, D, E) + + +def reference_bmm(op, sample): + # unbind reduces a dim and bmm requires 3D, so use matmul as the reference + matmul_op = copy(op) + matmul_op.op = torch.matmul + # change arg name from mat2 -> other + modified_sample = copy(sample) + other = modified_sample.kwargs["mat2"] + del modified_sample.kwargs["mat2"] + modified_sample.kwargs["other"] = other + return unbind_reference(matmul_op, modified_sample) + + +def sample_inputs_chunk(op_info, device, dtype, requires_grad, **kwargs): + for sample_input in sample_inputs_unary_dimwise( + op_info, device, dtype, requires_grad, **kwargs + ): + # ragged dim chunking: test a single chunks value + if sample_input.kwargs["dim"] == sample_input.input._ragged_idx: + yield _update_sample(sample_input, {"chunks": 3}) + # other dim chunking: test different chunks values + else: + D = sample_input.input.size(sample_input.kwargs["dim"]) + for chunks in [1, D // 2, D - 1, D]: + yield _update_sample(sample_input, {"chunks": chunks}) + + +def sample_inputs_matmul( + op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs +): + # also run bmm samples through + for sample_input in sample_inputs_bmm(op_info, device, dtype, requires_grad): + # change arg name from mat2 -> other + other = sample_input.kwargs["mat2"] + del sample_input.kwargs["mat2"] + sample_input.kwargs["other"] = other + yield sample_input + + # 3D cases not covered by bmm + for njt_3d in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[3] + ): + # (B, j1, D) x (D, E) => (B, j1, E) + if njt_3d._ragged_idx == 1: + D = njt_3d.shape[-1] + E = D + 2 + njt_desc = _describe_njt(njt_3d) + yield SampleInput( + _clone(njt_3d), + kwargs={"other": torch.randn(D, E, device=device, dtype=dtype)}, + name=f"{njt_desc}: (B, j, D) x (D, E)", + ) + + # 4D cases + for njt_4d in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[4] + ): + # (B, j1, D, E) x (E, F) => (B, j1, D, F) + if njt_4d._ragged_idx == 1: + E = njt_4d.shape[-1] + F = E + 2 + njt_desc = _describe_njt(njt_4d) + yield SampleInput( + _clone(njt_4d), + kwargs={"other": torch.randn(E, F, device=device, dtype=dtype)}, + name=f"{njt_desc}: (B, j, D, E) x (E, F)", + ) + + # Dense x NJT cases + for njt_3d in _sample_njts( + device=device, + dtype=dtype, + requires_grad=requires_grad, + dims=[3], + ): + # (B, F, E) x (B, E, j1) => (B, F, j1) + if njt_3d._ragged_idx == 2: + B = njt_3d.shape[0] + E = njt_3d.shape[1] + F = E + 2 + njt_desc = _describe_njt(njt_3d) + dense_t = torch.randn( + B, F, E, device=device, dtype=dtype, requires_grad=requires_grad + ) + dense_t._batch_dim = 0 # for unbind_reference() + yield SampleInput( + dense_t, + args=(_clone(njt_3d),), + name=f"{njt_desc}: (B, F, E) x (B, E, j1)", + ) + + # NJT x NJT => Dense case + for njt_3d in _sample_njts( + device=device, + dtype=dtype, + requires_grad=requires_grad, + dims=[3], + ): + # (B, E, j1) x (B, j1, F) => (B, E, F) + if njt_3d._ragged_idx == 2 and njt_3d.is_contiguous(): + B, E, _ = njt_3d.shape + sum_j1 = len(njt_3d.values()) + other_cont = torch.randn( + sum_j1, E + 2, device=device, dtype=dtype, requires_grad=requires_grad + ) + other_njt = torch.nested.nested_tensor_from_jagged( + other_cont, njt_3d.offsets(), lengths=njt_3d._lengths + ) + njt_desc = _describe_njt(njt_3d) + yield SampleInput( + _clone(njt_3d), + kwargs={"other": _clone(other_njt)}, + name=f"{njt_desc}: (B, E, j1) x (B, j1, F)", + ) + + # TODO (need factory functions): + # (B, j1, D, E) x (B, j1, E, F) => (B, j1, D, F) + + +def sample_inputs_masked_select( + op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs +): + for njt in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[2] + ): + yield SampleInput( + njt, + kwargs={"mask": (torch.randn_like(njt, requires_grad=False) < 0.0)}, + name=_describe_njt(njt), + ) + + +def sample_inputs_narrow(op_info, device, dtype, requires_grad, **kwargs): + for sample_input in sample_inputs_unary_dimwise( + op_info, device, dtype, requires_grad, **kwargs + ): + # ragged dim narrowing: test a single start, length value + if sample_input.kwargs["dim"] == sample_input.input._ragged_idx: + yield _update_sample(sample_input, {"start": 1, "length": 2}) + # other dim narrowing: test different start, length values + else: + D = sample_input.input.size(sample_input.kwargs["dim"]) + for start, length in [(0, D), (0, D - 1), (1, D - 1), (D - 1, 1)]: + yield _update_sample(sample_input, {"start": start, "length": length}) + + +def sample_inputs_nn_functional_embedding( + op_info, device, dtype, requires_grad, **kwargs +): + indices = torch.nested.nested_tensor( + [ + torch.tensor([0, 2, 1, 3]), + torch.tensor([4, 2, 1]), + torch.tensor([6, 7, 5, 2, 4]), + ], + layout=torch.jagged, + dtype=torch.int64, + device=device, + ) + + NUM_EMBEDDINGS = 20 + EMBEDDING_DIM = 32 + weight = torch.randn(NUM_EMBEDDINGS, EMBEDDING_DIM, device=device, dtype=dtype) + + # NB: the OpInfo entry for embedding_bag expects weight first so the gradients + # can be checked + yield SampleInput( + _clone(weight).requires_grad_(), + args=(indices,), + ) + + yield SampleInput( + _clone(weight).requires_grad_(), + args=(indices,), + kwargs={"padding_idx": 1}, + ) + + +def sample_inputs_index_put( + op_info, device, dtype, requires_grad, op_kwargs=None, **kwargs +): + for njt in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[2, 3, 4] + ): + for dim in range(njt.dim()): + indices = [ + torch.tensor(list(range(njt.size(0))), device=njt.device), + *[ + torch.tensor([0] * njt.size(0), device=njt.device) + for _ in range(dim - 1) + ], + ] + njt_desc = _describe_njt(njt) + yield SampleInput( + _clone(njt), + kwargs={ + "indices": indices, + "values": torch.tensor(1.0, device=njt.device), + }, + name=f"{njt_desc}: up to dim {dim - 1}", + ) + + # Non-cont NJT for completeness + offsets = torch.tensor([0, 2, 5, 7], device=device) + lengths = torch.tensor([2, 2, 2], device=device) + indices = [ + torch.tensor([0, 1, 2], device=device), + torch.tensor([0, 1, 1], device=device), + torch.tensor([0, 0, 0], device=device), + ] + a = torch.nested.nested_tensor_from_jagged( + torch.zeros(7, 3, device=device), offsets, lengths + ).requires_grad_(requires_grad) + + njt_desc = _describe_njt(a) + yield SampleInput( + _clone(a), + kwargs={"indices": indices, "values": torch.tensor(1.0, device=a.device)}, + name=f"{njt_desc}: all dims", + ) + + +def sample_inputs_nn_functional_embedding_bag( + op_info, device, dtype, requires_grad, **kwargs +): + for generate_per_sample_weight in (True, False): + for mode in ("sum", "mean", "max"): + # per_sample_weights is only supported for mode='sum' + if mode != "sum" and generate_per_sample_weight: + continue + + NUM_EMBEDDINGS = 10 + EMBEDDING_DIM = 32 + weight = torch.randn( + NUM_EMBEDDINGS, EMBEDDING_DIM, dtype=dtype, device=device + ) + + njt = torch.nested.nested_tensor( + [ + torch.randint(0, NUM_EMBEDDINGS, size=(2,)), + torch.randint(0, NUM_EMBEDDINGS, size=(3,)), + torch.randint(0, NUM_EMBEDDINGS, size=(4,)), + ], + layout=torch.jagged, + dtype=torch.int64, + device=device, + ) + + per_sample_weights = None + if generate_per_sample_weight: + per_sample_weights = torch.randn_like(njt, dtype=dtype) + + # NB: the OpInfo entry for embedding_bag expects weight first so the gradients + # can be checked + yield SampleInput( + weight, + args=(njt,), + kwargs={ + "mode": mode, + "per_sample_weights": per_sample_weights, + }, + ) + + +def reference_nn_functional_embedding_bag(op, sample): + # run reference on a single bag at a time + new_kwargs = dict(sample.kwargs) + new_kwargs.update( + {"offsets": torch.tensor([0], dtype=torch.int64, device=sample.input.device)} + ) + # flip input / weight back to what unbind_reference() expects + sample = SampleInput(sample.args[0], args=(sample.input,), kwargs=new_kwargs) + old_op = op.op + op.op = torch.nn.functional.embedding_bag + output = unbind_reference(op, sample, wrap_output_as_njt=False) + op.op = old_op + # concat bag outputs to get final output + return torch.cat(output, dim=0) + + +def sample_inputs_nn_functional_linear(op_info, device, dtype, requires_grad, **kwargs): + for njt in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[3, 4, 5] + ): + # projection over a ragged dim is not currently supported + if is_nested_int(njt.size(-1)): + continue + + # with bias + NUM_OUTPUT = 10 + weight = torch.randn( + NUM_OUTPUT, + njt.size(-1), + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + bias = torch.randn( + NUM_OUTPUT, device=device, dtype=dtype, requires_grad=requires_grad + ) + yield SampleInput( + _clone(njt), + kwargs={ + "weight": _clone(weight), + "bias": _clone(bias), + }, + name=f"{_describe_njt(njt)}: with bias", + ) + + # without bias + yield SampleInput( + _clone(njt), + kwargs={ + "weight": _clone(weight), + }, + name=f"{_describe_njt(njt)}: without bias", + ) + + +def sample_inputs_nn_functional_prelu(op_info, device, dtype, requires_grad, **kwargs): + for njt in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[3, 4] + ): + # Second dim is interpreted as number of channels; this should be non-ragged for now + num_channels = njt.size(1) + if is_nested_int(num_channels): + continue + + # 1D weight + weight = torch.randn( + num_channels, + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + + yield SampleInput( + _clone(njt), + kwargs={ + "weight": _clone(weight), + }, + name=f"{_describe_njt(njt)}: 1D weight", + ) + + # scalar tensor weight + yield SampleInput( + _clone(njt), + kwargs={ + "weight": torch.tensor(4.2, device=device, dtype=dtype), + }, + name=f"{_describe_njt(njt)}: scalar tensor weight", + ) + + +def sample_inputs_nn_functional_rms_norm( + op_info, device, dtype, requires_grad, **kwargs +): + for njt in _sample_njts( + device=device, dtype=dtype, requires_grad=requires_grad, dims=[3, 4] + ): + # normalize over non-ragged dims + for start_dim in range(njt.dim()): + if start_dim <= njt._ragged_idx: + continue + + normalized_shape = njt.shape[start_dim:] + weight = torch.randn( + normalized_shape, + device=device, + dtype=dtype, + requires_grad=requires_grad, + ) + + yield SampleInput( + _clone(njt), + kwargs={ + "normalized_shape": normalized_shape, + "weight": weight, + }, + name=f"{_describe_njt(njt)}", + ) + + +sample_inputs_nn_functional_threshold = partial( + sample_inputs_elementwise_njt_unary, + op_kwargs={"threshold": float.fromhex("0x1.3ap-3"), "value": -9}, +) + + +def sample_inputs_select(op_info, device, dtype, requires_grad, **kwargs): + for sample_input in sample_inputs_unary_dimwise( + op_info, device, dtype, requires_grad, **kwargs + ): + # ragged dim chunking: test a single index + if sample_input.kwargs["dim"] == sample_input.input._ragged_idx: + yield _update_sample(sample_input, {"index": 0}) + # other dim chunking: test different indices + else: + D = sample_input.input.size(sample_input.kwargs["dim"]) + for index in [0, D // 2, D - 1]: + yield _update_sample(sample_input, {"index": index}) + + +def sample_inputs_split(op_info, device, dtype, requires_grad, **kwargs): + for sample_input in sample_inputs_unary_dimwise( + op_info, device, dtype, requires_grad, **kwargs + ): + # ragged dim chunking: test a single split size + if sample_input.kwargs["dim"] == sample_input.input._ragged_idx: + yield _update_sample(sample_input, {"split_size_or_sections": 3}) + # other dim chunking: test different split sizes + else: + D = sample_input.input.size(sample_input.kwargs["dim"]) + for split_size in [1, D // 2, D - 1, D]: + yield _update_sample( + sample_input, {"split_size_or_sections": split_size} + ) + + +def sample_inputs_split_with_sizes(op_info, device, dtype, requires_grad, **kwargs): + for sample_input in sample_inputs_unary_dimwise( + op_info, device, dtype, requires_grad, **kwargs + ): + # It will never make sense to operate on the ragged dim. + # TODO: Handle this with error_inputs + if sample_input.kwargs["dim"] == sample_input.input._ragged_idx: + continue + + D = sample_input.input.size(sample_input.kwargs["dim"]) + # splits should add up to D + split1 = torch.randint(0, D - 1, size=()).item() + split2 = D - split1 + yield _update_sample(sample_input, {"split_sizes": [split1, split2]}) + + +def sample_inputs_squeeze(op_info, device, dtype, requires_grad, **kwargs): + # squeeze-specific NJT generator (need to ensure there are some 1s in the shape) + def _get_njts(): + njt = random_nt_from_dims( + (4, None, 1, 3, 1), + device=device, + dtype=dtype, + requires_grad=requires_grad, + layout=torch.jagged, + ) + yield njt + # without min / max seqlen cached + values = njt.values().detach().clone() + offsets = njt.offsets().detach().clone() + yield torch.nested.nested_tensor_from_jagged(values, offsets) + # non-contiguous transposed + yield njt.transpose(1, 3) + # non-contiguous with holes + values = njt.values().detach().clone() + offsets = njt.offsets().detach().clone() + # subtract 1 to cause holes + lengths = (offsets.diff() - 1).detach().clone() + yield torch.nested.nested_tensor_from_jagged( + values=values, + offsets=offsets, + lengths=lengths, + ) + + for njt in _get_njts(): + # single dim operation + for dim in range(njt.dim()): + # Operation on batch / ragged dim is never expected to work. + # TODO: Handle these via error_inputs. + if dim == 0 or dim == njt._ragged_idx: + continue + + yield SampleInput( + _clone(njt), + kwargs={"dim": dim}, + name=f"{_describe_njt(njt)}: {_describe_dim(njt, dim)}", + ) + + # multiple dim operation (pass no args) + yield SampleInput( + _clone(njt), + kwargs={"dim": dim}, + name=f"{_describe_njt(njt)}: multiple dims", + ) + + +def sample_inputs_unflatten(op_info, device, dtype, requires_grad, **kwargs): + for sample_input in sample_inputs_unary_dimwise( + op_info, device, dtype, requires_grad, **kwargs + ): + # It will never make sense to operate on the ragged dim. + # TODO: Handle this with error_inputs + if sample_input.kwargs["dim"] == sample_input.input._ragged_idx: + continue + + D = sample_input.input.size(sample_input.kwargs["dim"]) + # sizes should multiply to be D + yield _update_sample(sample_input, {"sizes": [D, 1]}) + yield _update_sample(sample_input, {"sizes": [1, D]}) + if D % 2 == 0: + yield _update_sample(sample_input, {"sizes": [D // 2, 2]}) + yield _update_sample(sample_input, {"sizes": [2, D // 2]}) + + +def sample_inputs_unsqueeze(op_info, device, dtype, requires_grad, **kwargs): + for sample_input in sample_inputs_unary_dimwise( + op_info, device, dtype, requires_grad, **kwargs + ): + yield sample_input + + last_dim_sample = _update_sample(sample_input, {"dim": -1}) + last_dim_sample.name = ( + f"{_describe_njt(last_dim_sample.input)}: add dim to the end" + ) + # Tell the unbind reference how to canonicalize the dim kwargs + # This is necessary because unsqueeze() allows for a dim after + # the last dim to indicate an unsqueeze at the end. + last_dim_sample.input._ndim = last_dim_sample.input.dim() + 1 + yield last_dim_sample + + +def sample_inputs_where(op_info, device, dtype, requires_grad, **kwargs): + for sample in sample_inputs_elementwise_njt_binary( + op_info, device, dtype, requires_grad, **kwargs + ): + other = sample.args[0] + sample.args = () + sample.kwargs["other"] = other + sample.kwargs["condition"] = sample.input > 0.0 + sample.name = sample.name.replace("(", "(NT, ") + yield sample + + +# === END OP-SPECIFIC SAMPLE INPUTS FUNCS / REFERENCES === + + +# Mapping of OpInfo full names -> sample_inputs_funcs, which define the set of sample inputs +# (involving NJTs) to pass to the op. Full name consists of the OpInfo's name and variant name +# separated by a period (e.g. special.polygamma.special_polygamma_n_0). These are necessary +# to specify if they cannot be auto-generated for some reason. Try to keep these sorted +# in alphabetical order! +njt_sample_inputs = { + "bmm": sample_inputs_bmm, + "chunk": sample_inputs_chunk, + "clone": sample_inputs_clone, + "count_nonzero": partial(sample_inputs_njt_reduction, supports_keepdim=False), + "fill": sample_inputs_fill, + **{f"mvlgamma.mvlgamma_p_{p}": sample_inputs_mvl_gamma(p=1) for p in (1, 3, 5)}, + "nn.functional.embedding": sample_inputs_nn_functional_embedding, + "nn.functional.embedding_bag": sample_inputs_nn_functional_embedding_bag, + "nn.functional.linear": sample_inputs_nn_functional_linear, + "nn.functional.prelu": sample_inputs_nn_functional_prelu, + "nn.functional.rms_norm": sample_inputs_nn_functional_rms_norm, + "nn.functional.threshold": sample_inputs_nn_functional_threshold, + **{f"polygamma.polygamma_n_{n}": sample_inputs_polygamma_n(n=n) for n in range(5)}, + "special.polygamma.special_polygamma_n_0": sample_inputs_special_polygamma_n(n=0), + "to": sample_inputs_to, + "matmul": sample_inputs_matmul, + "masked_select": sample_inputs_masked_select, + "narrow": sample_inputs_narrow, + "index_put": sample_inputs_index_put, + # these two don't have ReductionOpInfo entries + "max.reduction_with_dim": sample_inputs_njt_reduction, + "min.reduction_with_dim": sample_inputs_njt_reduction, + "select": sample_inputs_select, + "split": sample_inputs_split, + "split_with_sizes": sample_inputs_split_with_sizes, + "squeeze": sample_inputs_squeeze, + "unflatten": sample_inputs_unflatten, + "unsqueeze": sample_inputs_unsqueeze, + "where": sample_inputs_where, +} + +njt_references = { + "bmm": reference_bmm, + "chunk": partial( + unary_dimwise_reference, batchwise_reference=batchwise_reference_chunk + ), + "count_nonzero": reduction_reference, + # these two don't have ReductionOpInfo entries + "max.reduction_with_dim": reduction_reference, + "min.reduction_with_dim": reduction_reference, + "narrow": partial( + unary_dimwise_reference, batchwise_reference=batchwise_reference_narrow + ), + "select": partial( + unary_dimwise_reference, batchwise_reference=batchwise_reference_select + ), + "split": partial( + unary_dimwise_reference, batchwise_reference=batchwise_reference_split + ), + "split_with_sizes": partial( + unary_dimwise_reference, + batchwise_reference=batchwise_reference_split_with_sizes, + ), + "squeeze": unbind_reference, + "nn.functional.embedding_bag": reference_nn_functional_embedding_bag, + "unflatten": partial( + unary_dimwise_reference, batchwise_reference=batchwise_reference_unflatten + ), + "unsqueeze": partial( + unary_dimwise_reference, batchwise_reference=batchwise_reference_unsqueeze + ), +} + + +# Translates an OpInfo entry to one that operates on NJTs. +def translate_opinfo(op): + new_op = copy(op) + new_op.supports_njt = True + # add some extra info for use in generating tests on the right subset of ops + new_op._extra_op_data = extra_op_data.get(op.full_name, ExtraOpData()) + + if op.full_name in njt_sample_inputs: + new_op.sample_inputs_func = njt_sample_inputs[op.full_name] + new_op.ref = njt_references.get(op.full_name, unbind_reference) + elif isinstance(op, UnaryUfuncInfo): + new_op.sample_inputs_func = partial( + sample_inputs_elementwise_njt_unary, op_kwargs=None + ) + new_op.ref = unbind_reference + elif isinstance(op, BinaryUfuncInfo): + new_op.sample_inputs_func = partial( + sample_inputs_elementwise_njt_binary, op_kwargs=None + ) + new_op.ref = unbind_reference + elif isinstance(op, ReductionOpInfo): + new_op.sample_inputs_func = partial(sample_inputs_njt_reduction, op_kwargs=None) + new_op.ref = reduction_reference + # TODO: Translate the rest of the OpInfos + else: + new_op.sample_inputs_func = unsupported_sample_inputs_func(op.full_name) + new_op.ref = unsupported_reference(op.full_name) + new_op.supports_njt = False + + return new_op + + +njt_op_db = [translate_opinfo(op) for op in op_db] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/signal.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/signal.py new file mode 100644 index 0000000000000000000000000000000000000000..f81efd19dbc6c804f066fd89a7068dce8ecf515f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/signal.py @@ -0,0 +1,459 @@ +# mypy: ignore-errors + +import unittest +from collections.abc import Callable +from functools import partial +from itertools import product + +import numpy + +import torch +from torch.testing._internal.common_dtype import floating_types +from torch.testing._internal.common_utils import TEST_SCIPY +from torch.testing._internal.opinfo.core import ( + DecorateInfo, + ErrorInput, + OpInfo, + SampleInput, +) + + +if TEST_SCIPY: + import scipy.signal + + +def sample_inputs_window(op_info, device, dtype, requires_grad, *args, **kwargs): + r"""Base function used to create sample inputs for windows. + + For additional required args you should use *args, as well as **kwargs for + additional keyword arguments. + """ + + # Remove include_conjugated_inputs from kwargs + kwargs.pop("include_conjugated_inputs", None) + # Tests window sizes up to 5 samples. + for size, sym in product(range(6), (True, False)): + yield SampleInput( + size, + *args, + sym=sym, + device=device, + dtype=dtype, + requires_grad=requires_grad, + **kwargs, + ) + + +def reference_inputs_window(op_info, device, dtype, requires_grad, *args, **kwargs): + r"""Reference inputs function to use for windows which have a common signature, i.e., + window size and sym only. + + Implement other special functions for windows that have a specific signature. + See exponential and gaussian windows for instance. + """ + yield from sample_inputs_window( + op_info, device, dtype, requires_grad, *args, **kwargs + ) + + cases = (8, 16, 32, 64, 128, 256) + + for size in cases: + yield SampleInput(size, sym=False) + yield SampleInput(size, sym=True) + + +def reference_inputs_exponential_window( + op_info, device, dtype, requires_grad, **kwargs +): + yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs) + + cases = ( + (8, {"center": 4, "tau": 0.5}), + (16, {"center": 8, "tau": 2.5}), + (32, {"center": 16, "tau": 43.5}), + (64, {"center": 20, "tau": 3.7}), + (128, {"center": 62, "tau": 99}), + (256, {"tau": 10}), + ) + + for size, kw in cases: + yield SampleInput(size, sym=False, **kw) + kw["center"] = None + yield SampleInput(size, sym=True, **kw) + + +def reference_inputs_gaussian_window(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs) + + cases = ( + (8, {"std": 0.1}), + (16, {"std": 1.2}), + (32, {"std": 2.1}), + (64, {"std": 3.9}), + (128, {"std": 4.5}), + (256, {"std": 10}), + ) + + for size, kw in cases: + yield SampleInput(size, sym=False, **kw) + yield SampleInput(size, sym=True, **kw) + + +def reference_inputs_kaiser_window(op_info, device, dtype, requires_grad, **kwargs): + yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs) + + cases = ( + (8, {"beta": 2}), + (16, {"beta": 12}), + (32, {"beta": 30}), + (64, {"beta": 35}), + (128, {"beta": 41.2}), + (256, {"beta": 100}), + ) + + for size, kw in cases: + yield SampleInput(size, sym=False, **kw) + yield SampleInput(size, sym=True, **kw) + + +def reference_inputs_general_cosine_window( + op_info, device, dtype, requires_grad, **kwargs +): + yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs) + + cases = ( + (8, {"a": [0.5, 0.5]}), + (16, {"a": [0.46, 0.54]}), + (32, {"a": [0.46, 0.23, 0.31]}), + (64, {"a": [0.5]}), + (128, {"a": [0.1, 0.8, 0.05, 0.05]}), + (256, {"a": [0.2, 0.2, 0.2, 0.2, 0.2]}), + ) + + for size, kw in cases: + yield SampleInput(size, sym=False, **kw) + yield SampleInput(size, sym=True, **kw) + + +def reference_inputs_general_hamming_window( + op_info, device, dtype, requires_grad, **kwargs +): + yield from sample_inputs_window(op_info, device, dtype, requires_grad, **kwargs) + + cases = ( + (8, {"alpha": 0.54}), + (16, {"alpha": 0.5}), + (32, {"alpha": 0.23}), + (64, {"alpha": 0.8}), + (128, {"alpha": 0.9}), + (256, {"alpha": 0.05}), + ) + + for size, kw in cases: + yield SampleInput(size, sym=False, **kw) + yield SampleInput(size, sym=True, **kw) + + +def error_inputs_window(op_info, device, *args, **kwargs): + # Tests for windows that have a negative size + yield ErrorInput( + SampleInput(-1, *args, dtype=torch.float32, device=device, **kwargs), + error_type=ValueError, + error_regex="requires non-negative window length, got M=-1", + ) + + # Tests for window tensors that are not torch.strided, for instance, torch.sparse_coo. + yield ErrorInput( + SampleInput( + 3, + *args, + layout=torch.sparse_coo, + device=device, + dtype=torch.float32, + **kwargs, + ), + error_type=ValueError, + error_regex="is implemented for strided tensors only, got: torch.sparse_coo", + ) + + # Tests for window tensors that are not floating point dtypes, for instance, torch.long. + yield ErrorInput( + SampleInput(3, *args, dtype=torch.long, device=device, **kwargs), + error_type=ValueError, + error_regex="expects float32 or float64 dtypes, got: torch.int64", + ) + + # Tests for window tensors that are bfloat16 + yield ErrorInput( + SampleInput(3, *args, dtype=torch.bfloat16, device=device, **kwargs), + error_type=ValueError, + error_regex="expects float32 or float64 dtypes, got: torch.bfloat16", + ) + + # Tests for window tensors that are float16 + yield ErrorInput( + SampleInput(3, *args, dtype=torch.float16, device=device, **kwargs), + error_type=ValueError, + error_regex="expects float32 or float64 dtypes, got: torch.float16", + ) + + +def error_inputs_exponential_window(op_info, device, **kwargs): + # Yield common error inputs + yield from error_inputs_window(op_info, device, **kwargs) + + # Tests for negative decay values. + yield ErrorInput( + SampleInput(3, tau=-1, dtype=torch.float32, device=device, **kwargs), + error_type=ValueError, + error_regex="Tau must be positive, got: -1 instead.", + ) + + # Tests for symmetric windows and a given center value. + yield ErrorInput( + SampleInput(3, center=1, sym=True, dtype=torch.float32, device=device), + error_type=ValueError, + error_regex="Center must be None for symmetric windows", + ) + + +def error_inputs_gaussian_window(op_info, device, **kwargs): + # Yield common error inputs + yield from error_inputs_window(op_info, device, std=0.5, **kwargs) + + # Tests for negative standard deviations + yield ErrorInput( + SampleInput(3, std=-1, dtype=torch.float32, device=device, **kwargs), + error_type=ValueError, + error_regex="Standard deviation must be positive, got: -1 instead.", + ) + + +def error_inputs_kaiser_window(op_info, device, **kwargs): + # Yield common error inputs + yield from error_inputs_window(op_info, device, beta=12, **kwargs) + + # Tests for negative beta + yield ErrorInput( + SampleInput(3, beta=-1, dtype=torch.float32, device=device, **kwargs), + error_type=ValueError, + error_regex="beta must be non-negative, got: -1 instead.", + ) + + +def error_inputs_general_cosine_window(op_info, device, **kwargs): + # Yield common error inputs + yield from error_inputs_window(op_info, device, a=[0.54, 0.46], **kwargs) + + # Tests for negative beta + yield ErrorInput( + SampleInput(3, a=None, dtype=torch.float32, device=device, **kwargs), + error_type=TypeError, + error_regex="Coefficients must be a list/tuple", + ) + + yield ErrorInput( + SampleInput(3, a=[], dtype=torch.float32, device=device, **kwargs), + error_type=ValueError, + error_regex="Coefficients cannot be empty", + ) + + +def reference_signal_window(fn: Callable): + r"""Wrapper for scipy signal window references. + + Discards keyword arguments for window reference functions that don't have a matching signature with + torch, e.g., gaussian window. + """ + + def _fn( + *args, + dtype=numpy.float64, + device=None, + layout=torch.strided, + requires_grad=False, + **kwargs, + ): + r"""The unused arguments are defined to disregard those values""" + return fn(*args, **kwargs).astype(dtype) + + return _fn + + +def make_signal_windows_opinfo( + name: str, + ref: Callable, + sample_inputs_func: Callable, + reference_inputs_func: Callable, + error_inputs_func: Callable, + *, + skips: tuple[DecorateInfo, ...] = (), +): + r"""Helper function to create OpInfo objects related to different windows.""" + return OpInfo( + name=name, + ref=ref if TEST_SCIPY else None, + dtypes=floating_types(), + sample_inputs_func=sample_inputs_func, + reference_inputs_func=reference_inputs_func, + error_inputs_func=error_inputs_func, + supports_out=False, + supports_autograd=False, + skips=( + # TODO: same as this? + # https://github.com/pytorch/pytorch/issues/81774 + # also see: arange, new_full + # fails to match any schemas despite working in the interpreter + DecorateInfo( + unittest.expectedFailure, + "TestOperatorSignatures", + "test_get_torch_func_signature_exhaustive", + ), + # fails to match any schemas despite working in the interpreter + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + # skip these tests since we have non tensor input + DecorateInfo( + unittest.skip("Skipped!"), "TestCommon", "test_noncontiguous_samples" + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestCommon", + "test_variant_consistency_eager", + ), + DecorateInfo(unittest.skip("Skipped!"), "TestMathBits", "test_conj_view"), + DecorateInfo( + unittest.skip("Skipped!"), "TestMathBits", "test_neg_conj_view" + ), + DecorateInfo(unittest.skip("Skipped!"), "TestMathBits", "test_neg_view"), + DecorateInfo( + unittest.skip("Skipped!"), + "TestVmapOperatorsOpInfo", + "test_vmap_exhaustive", + ), + DecorateInfo( + unittest.skip("Skipped!"), + "TestVmapOperatorsOpInfo", + "test_op_has_batch_rule", + ), + DecorateInfo( + unittest.skip("Buggy on MPS for now (mistakenly promotes to float64)"), + "TestCommon", + "test_numpy_ref_mps", + ), + *skips, + ), + ) + + +op_db: list[OpInfo] = [ + make_signal_windows_opinfo( + name="signal.windows.hamming", + ref=reference_signal_window(scipy.signal.windows.hamming) + if TEST_SCIPY + else None, + sample_inputs_func=sample_inputs_window, + reference_inputs_func=reference_inputs_window, + error_inputs_func=error_inputs_window, + ), + make_signal_windows_opinfo( + name="signal.windows.hann", + ref=reference_signal_window(scipy.signal.windows.hann) if TEST_SCIPY else None, + sample_inputs_func=sample_inputs_window, + reference_inputs_func=reference_inputs_window, + error_inputs_func=error_inputs_window, + ), + make_signal_windows_opinfo( + name="signal.windows.bartlett", + ref=reference_signal_window(scipy.signal.windows.bartlett) + if TEST_SCIPY + else None, + sample_inputs_func=sample_inputs_window, + reference_inputs_func=reference_inputs_window, + error_inputs_func=error_inputs_window, + ), + make_signal_windows_opinfo( + name="signal.windows.blackman", + ref=reference_signal_window(scipy.signal.windows.blackman) + if TEST_SCIPY + else None, + sample_inputs_func=sample_inputs_window, + reference_inputs_func=reference_inputs_window, + error_inputs_func=error_inputs_window, + ), + make_signal_windows_opinfo( + name="signal.windows.cosine", + ref=reference_signal_window(scipy.signal.windows.cosine) + if TEST_SCIPY + else None, + sample_inputs_func=sample_inputs_window, + reference_inputs_func=reference_inputs_window, + error_inputs_func=error_inputs_window, + ), + make_signal_windows_opinfo( + name="signal.windows.exponential", + ref=reference_signal_window(scipy.signal.windows.exponential) + if TEST_SCIPY + else None, + sample_inputs_func=partial(sample_inputs_window, tau=2.78), + reference_inputs_func=partial(reference_inputs_exponential_window, tau=2.78), + error_inputs_func=error_inputs_exponential_window, + ), + make_signal_windows_opinfo( + name="signal.windows.gaussian", + ref=reference_signal_window(scipy.signal.windows.gaussian) + if TEST_SCIPY + else None, + sample_inputs_func=partial(sample_inputs_window, std=1.92), + reference_inputs_func=partial(reference_inputs_gaussian_window, std=1.92), + error_inputs_func=error_inputs_gaussian_window, + skips=( + DecorateInfo( + unittest.skip("Buggy on MPS for now (mistakenly promotes to float64)"), + "TestCommon", + "test_numpy_ref_mps", + ), + ), + ), + make_signal_windows_opinfo( + name="signal.windows.kaiser", + ref=reference_signal_window(scipy.signal.windows.kaiser) + if TEST_SCIPY + else None, + sample_inputs_func=partial(sample_inputs_window, beta=12.0), + reference_inputs_func=partial(reference_inputs_kaiser_window, beta=12.0), + error_inputs_func=error_inputs_kaiser_window, + ), + make_signal_windows_opinfo( + name="signal.windows.general_cosine", + ref=reference_signal_window(scipy.signal.windows.general_cosine) + if TEST_SCIPY + else None, + sample_inputs_func=partial(sample_inputs_window, a=[0.54, 0.46]), + reference_inputs_func=partial( + reference_inputs_general_cosine_window, a=[0.54, 0.46] + ), + error_inputs_func=error_inputs_general_cosine_window, + ), + make_signal_windows_opinfo( + name="signal.windows.general_hamming", + ref=reference_signal_window(scipy.signal.windows.general_hamming) + if TEST_SCIPY + else None, + sample_inputs_func=partial(sample_inputs_window, alpha=0.54), + reference_inputs_func=partial( + reference_inputs_general_hamming_window, alpha=0.54 + ), + error_inputs_func=error_inputs_window, + ), + make_signal_windows_opinfo( + name="signal.windows.nuttall", + ref=reference_signal_window(scipy.signal.windows.nuttall) + if TEST_SCIPY + else None, + sample_inputs_func=sample_inputs_window, + reference_inputs_func=reference_inputs_window, + error_inputs_func=error_inputs_window, + ), +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/sparse.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/sparse.py new file mode 100644 index 0000000000000000000000000000000000000000..c04009f1f2ac97804079933bca3bab531a96edb1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/sparse.py @@ -0,0 +1,931 @@ +# mypy: ignore-errors + +import os + +import torch +from torch.testing import make_tensor # noqa: F401 +from torch.testing._internal.common_dtype import highest_precision_float +from torch.testing._internal.opinfo.core import ( # noqa: F401 + BinaryUfuncInfo, + ErrorInput, + generate_elementwise_binary_tensors, + ReductionOpInfo, + sample_inputs_reduction, + SampleInput, +) + + +def _check_validate(op_info, sample): + def _check_fail(sample): + try: + op_info( + sample.sample_input.input, + *sample.sample_input.args, + **sample.sample_input.kwargs, + ) + except sample.error_type: + pass + except Exception as msg: + raise AssertionError( # noqa: B904 + f"{op_info.name} on {sample.sample_input=} expected exception " + f"{sample.error_type}: {sample.error_regex}, got {type(msg).__name__}: {msg}" + ) + else: + raise AssertionError( + f"{op_info.name} on {sample.sample_input=} expected exception " + f"{sample.error_type}: {sample.error_regex}, got none." + ) + + def _check_success(sample): + try: + op_info(sample.input, *sample.args, **sample.kwargs) + except Exception as msg: + raise AssertionError( # noqa: B904 + f"{op_info.name} on {sample=} expected to succeed " + f", got {type(msg).__name__}: {msg}" + ) + + if isinstance(sample, ErrorInput): + _check_fail(sample) + else: + _check_success(sample) + + +def _sample_inputs_sparse( + sample_inputs, + maybe_failing_sample_inputs, + validate_sample_input, + op_info, + *args, + **kwargs, +): + check_validate = ( + os.environ.get("PYTORCH_TEST_CHECK_VALIDATE_SPARSE_SAMPLES", "0") == "1" + ) + for sample in sample_inputs(op_info, *args, **kwargs): + sample = validate_sample_input(op_info, sample, check_validate=check_validate) + if isinstance(sample, SampleInput): + yield sample + # Error inputs are handled in error_inputs_sparse + + for sample in maybe_failing_sample_inputs(op_info, *args, **kwargs): + sample = validate_sample_input(op_info, sample, check_validate=check_validate) + if isinstance(sample, SampleInput): + yield sample + + +def _error_inputs_sparse( + maybe_failing_sample_inputs, validate_sample_input, op_info, *args, **kwargs +): + check_validate = ( + os.environ.get("PYTORCH_TEST_CHECK_VALIDATE_SPARSE_SAMPLES", "0") == "1" + ) + for sample in maybe_failing_sample_inputs(op_info, *args, **kwargs): + sample = validate_sample_input(op_info, sample, check_validate=check_validate) + if isinstance(sample, ErrorInput): + yield sample + # Sample inputs are handled in sample_inputs_sparse + + +def _apply_requires_grad_to_samples(sample_inputs): + """Decorator to _maybe_failing_sample_inputs_... generator functions + that clones and sets requires_grad argument to tensors in sample + input arguments. This is needed when the generated samples share + tensor instances. + """ + + def wrapper(op_info, device, dtype, requires_grad, layout, **kwargs): + def apply_requires_grad(x): + if ( + not isinstance(x, torch.Tensor) + or x.requires_grad + or not requires_grad + or not (x.is_floating_point() or x.is_complex()) + ): + return x + return x.detach().clone().requires_grad_(requires_grad) + + if requires_grad: + for sample_input in sample_inputs( + op_info, device, dtype, requires_grad, layout, **kwargs + ): + yield sample_input.transform(apply_requires_grad) + else: + yield from sample_inputs( + op_info, device, dtype, requires_grad, layout, **kwargs + ) + + return wrapper + + +def sample_inputs_sparse_reduction( + op_info, device, dtype, requires_grad, layout, blocksize=None, **kwargs +): + """Sample inputs for reduction operations on sparse tensors.""" + layout_name = str(layout).split(".", 1)[-1].rsplit("_coo", 1)[0] + op_supports_layout = getattr(op_info, "supports_" + layout_name) + if not op_supports_layout: + return + + for sample_input in sample_inputs_reduction( + op_info, device, dtype, requires_grad, **kwargs + ): + if sample_input.input.ndim == 0: + # scalar sparse tensors are not supported + continue + + if layout in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }: + if sample_input.input.ndim < 2: + # conversion to sparse compressed tensors requires at + # least 2 dimensional tensors + continue + if sample_input.input.ndim > 2 and (sample_input.input == 0).any(): + # Skip batched sparse compressed samples that contain + # explicit zeros because to_sparse(layout=..) will + # fail, see gh-98495. + # TODO: remove this if-block after gh-98495 is fixed. + continue + + if layout in {torch.sparse_bsr, torch.sparse_bsc} and blocksize is None: + blocksize = (1, 1) + + yield SampleInput( + sample_input.input.detach() + .to_sparse(layout=layout, blocksize=blocksize) + .requires_grad_(requires_grad), + args=sample_input.args, + kwargs=sample_input.kwargs, + ) + + if layout is torch.sparse_coo and (dtype.is_floating_point or dtype.is_complex): + # uncoalesced samples + inp = sample_input.input.detach().to_sparse(layout=layout) + inp = torch.sparse_coo_tensor( + inp.indices().repeat(1, 2), + inp.values().repeat(2), + inp.shape, + dtype=inp.dtype, + device=inp.device, + ) + if inp.is_coalesced(): + raise AssertionError("Expected inp to not be coalesced") + yield SampleInput( + inp.requires_grad_(requires_grad), + args=sample_input.args, + kwargs=sample_input.kwargs, + ) + + if sample_input.input.ndim > 2: + # hybrid samples + yield SampleInput( + sample_input.input.detach() + .to_sparse( + layout=layout, + blocksize=blocksize, + dense_dim=sample_input.input.ndim - 2, + ) + .requires_grad_(requires_grad), + args=sample_input.args, + kwargs=sample_input.kwargs, + ) + + +def _validate_sample_input_sparse_reduction(op_info, sample, check_validate=False): + """Return the specified sample when it is valid and supported by the + operation. Otherwise, return the sample as ErrorInput instance. + + When check_validate is True, the result is validated against + calling the op on the sample. + """ + UNSPECIFIED = object() + if op_info.name == "sum": + sample = _validate_sample_input_sparse_reduction_sum(sample) + + if op_info.name == "masked.sum": + mask = sample.kwargs.get("mask", UNSPECIFIED) + if ( + mask not in {None, UNSPECIFIED} + and mask.ndim > 2 + and mask.layout is torch.strided + and (mask == 0).any() + ): + # TODO: remove this if-block after gh-98495 is fixed. + sample = ErrorInput( + sample, + error_regex="Expect the same number of specified elements per batch.", + ) + elif not sample.kwargs.get("keepdim"): + sample = ErrorInput( + sample, + error_type=(AssertionError, RuntimeError), + error_regex="reduction operations on (CSR|CSC) tensors with keepdim=False is unsupported", + ) + elif mask is UNSPECIFIED: + sample = ErrorInput( + sample, + error_type=ValueError, + error_regex="masked (.*) expects explicit mask for sparse_csr tensor input", + ) + elif sample.input.ndim > 2: + sample = ErrorInput( + sample, + error_regex="crow_indices is supposed to be a vector, but got 3 dimensional tensor.", + ) + + if op_info.name in {"masked.amax", "masked.amin", "masked.mean", "masked.prod"}: + t_inp = sample.input + mask = sample.kwargs.get("mask") + if ( + mask is not None + and mask.ndim > 2 + and mask.layout is torch.strided + and (mask == 0).any() + ): + # TODO: remove this if-block after gh-98495 is fixed. + sample = ErrorInput( + sample, + error_regex="Expect the same number of specified elements per batch.", + ) + elif mask is None: + sample = ErrorInput( + sample, + error_type=ValueError, + error_regex="masked (.*) expects explicit mask for sparse_csr tensor input", + ) + elif ( + mask.layout is sample.input.layout + and mask.ndim > 2 + and op_info.name == "masked.mean" + ): + sample = ErrorInput( + sample, + error_type=TypeError, + error_regex=( + "where[(][)] received an invalid combination of arguments" + " - got [(]Tensor, Tensor, NoneType[)]" + ), + ) + elif not sample.kwargs.get("keepdim"): + sample = ErrorInput( + sample, + error_type=(AssertionError, RuntimeError), + error_regex="reduction operations on (CSR|CSC) tensors with keepdim=False is unsupported", + ) + elif ( + sample.input.ndim > 2 + and (sample.kwargs.get("dim") not in {0, 1}) + and mask.ndim > 2 + and mask.layout is not torch.strided + ): + if sample.kwargs.get("dim") == (0, -1): + sample = ErrorInput( + sample, + error_regex="tensor dimensionality must be sum of batch, base, and dense dimensionalities", + ) + elif op_info.name == "masked.prod": + sample = ErrorInput( + sample, + error_regex="input_dim == 2 INTERNAL ASSERT FAILED at", + ) + else: + sample = ErrorInput( + sample, + error_type=AssertionError, + error_regex="Sparse CSR tensors are 2D and only support reduction along dim 0 or 1.", + ) + elif sample.input.ndim > 2: + sample = ErrorInput( + sample, + error_regex="crow_indices is supposed to be a vector, but got 3 dimensional tensor.", + ) + elif ( + mask.layout is t_inp.layout + and mask._nnz() != t_inp._nnz() + and t_inp.dense_dim() > 0 + ): + sample = ErrorInput( + sample, + error_regex="Index tensor must have the same number of dimensions as src tensor", + ) + + if check_validate: + _check_validate(op_info, sample) + + return sample + + +def _validate_sample_input_sparse_reduction_sum(sample, check_validate=False): + # NOTE: When fixing a failing sample case, remove the + # corresponding if-block + t_inp, t_kwargs = sample.input, sample.kwargs + dim = t_kwargs.get("dim") + keepdim = t_kwargs.get("keepdim") + layout = t_inp.layout + if isinstance(dim, (int, list, tuple)): + if layout in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }: + if layout in {torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc}: + return ErrorInput( + sample, + error_regex=( + "Currently the only compressed sparse format supported for sum.dim_IntList is CSR, but got layout" + ), + ) + if layout in {torch.sparse_csr, torch.sparse_csc} and not keepdim: + return ErrorInput( + sample, + error_regex=( + "reduction operations on CSR tensors with keepdim=False is unsupported" + ), + ) + if t_inp.dim() != 2: + return ErrorInput( + sample, + error_regex=("input_dim == 2 INTERNAL ASSERT"), + ) + if layout == torch.sparse_csr: + if t_inp.dtype == torch.bool: + return ErrorInput( + sample, + error_regex=("_sparse_csr_sum_cpu not implemented for 'Bool'"), + ) + if t_inp.dtype == torch.complex32: + return ErrorInput( + sample, + error_regex=( + "_sparse_csr_sum_cuda not implemented for 'ComplexHalf'" + ), + ) + return sample + + +def _maybe_failing_sample_inputs_sparse_reduction_sum( + op_info, device, dtype, requires_grad, layout, **kwargs +): + """Generator of samples that are known to fail or that were failing in past.""" + # NOTE: When fixing a failing case, remove the Exception comment + # but keep the `yield sample` statement. + if layout in [ + torch.sparse_csr, + torch.sparse_csc, + ]: + # NotImplementedError: Could not run 'aten::sum.IntList_out' with arguments from the 'SparseCsrCPU' backend. + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype) + .to_sparse(layout=layout) + .requires_grad_(requires_grad), + kwargs=dict(dim=0, keepdim=True), + ) + yield SampleInput( + torch.tensor([[[0, 1]], [[2, 3]]], dtype=dtype) + .to_sparse(layout=layout, dense_dim=1) + .requires_grad_(requires_grad), + kwargs=dict(dim=0), + ) + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype) + .to_sparse(layout=layout) + .requires_grad_(requires_grad), + kwargs=dict(dim=(0,)), + ) + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype) + .to_sparse(layout=layout) + .requires_grad_(requires_grad), + kwargs=dict(dim=(0,), keepdim=True), + ) + yield SampleInput( + torch.tensor([[[0, 1]], [[2, 3]]], dtype=dtype) + .to_sparse(layout=layout, dense_dim=1) + .requires_grad_(requires_grad), + kwargs=dict(dim=(0,)), + ) + + # RuntimeError: torch.empty: Only batched sparse compressed (non-block) tensors are supported, but got size [2] + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype) + .to_sparse(layout=layout) + .requires_grad_(requires_grad), + kwargs=dict(dim=0), + ) + + if layout in [ + torch.sparse_bsr, + torch.sparse_bsc, + ]: + # RuntimeError: empty_sparse_compressed expected sparse compressed (non-block) tensor layout but got SparseBsr + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype) + .to_sparse(layout=layout, blocksize=(2, 2)) + .requires_grad_(requires_grad), + kwargs=dict(dim=0, keepdim=True), + ) + yield SampleInput( + torch.tensor([[[0, 1]], [[2, 3]]], dtype=dtype) + .to_sparse(layout=layout, dense_dim=1, blocksize=(1, 1)) + .requires_grad_(requires_grad), + kwargs=dict(dim=0), + ) + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype) + .to_sparse(layout=layout, blocksize=(1, 1)) + .requires_grad_(requires_grad), + kwargs=dict(dim=(0,)), + ) + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype) + .to_sparse(layout=layout, blocksize=(1, 1)) + .requires_grad_(requires_grad), + kwargs=dict(dim=(0,), keepdim=True), + ) + yield SampleInput( + torch.tensor([[[0, 1]], [[2, 3]]], dtype=dtype) + .to_sparse(layout=layout, blocksize=(1, 1), dense_dim=1) + .requires_grad_(requires_grad), + kwargs=dict(dim=(0,)), + ) + + # RuntimeError: torch.empty: Only batched sparse compressed (non-block) tensors are supported, but got size [2] + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype) + .to_sparse(layout=layout, blocksize=(1, 1)) + .requires_grad_(requires_grad), + kwargs=dict(dim=0), + ) + + +def sample_inputs_sparse_reduction_sum( + op_info, device, dtype, requires_grad, layout, **kwargs +): + """Sample inputs for sum on sparse tensors.""" + yield from _sample_inputs_sparse( + sample_inputs_sparse_reduction, + _maybe_failing_sample_inputs_sparse_reduction_sum, + _validate_sample_input_sparse_reduction, + op_info, + device, + dtype, + requires_grad, + layout, + **kwargs, + ) + + +def error_inputs_sparse_reduction_sum(op_info, device, layout, **kwargs): + """Error inputs for sum on sparse tensors.""" + dtype = torch.float64 + requires_grad = False + yield from _error_inputs_sparse( + _maybe_failing_sample_inputs_sparse_reduction_sum, + _validate_sample_input_sparse_reduction, + op_info, + device, + dtype, + requires_grad, + layout, + **kwargs, + ) + + +def sample_inputs_sparse_elementwise_binary_operation( + op_info, device, dtype, requires_grad, layout, **kwargs +): + """Sample inputs for elementwise binary operations on sparse tensors. + + The samples include regular, zero-sized, batched, and hybrid + sparse tensors as well as rhs scalars. All tensors are full tensors. + """ + + def _to_sparse(tensor, **kwargs): + return tensor.detach().to_sparse(**kwargs).requires_grad_(requires_grad) + + for sample_input in generate_elementwise_binary_tensors( + op_info, + device=device, + dtype=dtype, + requires_grad=requires_grad, + exclude_zero=True, + **kwargs, + ): + lhs, rhs = sample_input.input, sample_input.args[0] + min_dense_dim = 0 + max_dense_dim = lhs.ndim - 1 + if layout in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + }: + if lhs.ndim < 2: + # sparse compressed tensors sparse_dim must be 2 + continue + max_dense_dim = lhs.ndim - 2 + + for dense_dim in range(min_dense_dim, max_dense_dim + 1): + if layout in {torch.sparse_bsr, torch.sparse_bsc}: + blocksizes = [(1, 1)] + if lhs.numel() > 0: + blocksizes.append( + ( + lhs.shape[lhs.ndim - 2 - dense_dim], + lhs.shape[lhs.ndim - 1 - dense_dim], + ) + ) + else: + blocksizes = [None] + for blocksize in blocksizes: + to_sparse_kwargs = dict( + layout=layout, dense_dim=dense_dim, blocksize=blocksize + ) + lhs_sparse = _to_sparse(lhs, **to_sparse_kwargs) + rhs_sparse = _to_sparse(rhs, **to_sparse_kwargs) + # op(sparse, sparse) + yield SampleInput( + lhs_sparse, + args=(rhs_sparse, *sample_input.args[1:]), + kwargs=sample_input.kwargs, + ) + # op(sparse, scalar) + yield SampleInput( + lhs_sparse, + args=( + make_tensor( + (), dtype=dtype, device=device, requires_grad=requires_grad + ), + *sample_input.args[1:], + ), + kwargs=sample_input.kwargs, + ) + + +def _validate_sample_input_elementwise_binary_sparse_mul(sample): + # NOTE: When fixing a failing sample case, remove the + # corresponding if-block + t_inp, t_args = sample.input, sample.args + batch_dim = t_inp.dim() - t_inp.dense_dim() - t_inp.sparse_dim() + layout = t_inp.layout + dtype = t_inp.dtype + if layout is torch.sparse_csr and batch_dim > 0 and t_args[0].ndim > 0: + return ErrorInput( + sample, + error_regex=( + "coo_to_sparse_csr: conversion from Sparse to SparseCsr for input" + " tensors with sparse_dim[(][)]!=2 is not supported" + ), + ) + elif layout is torch.sparse_csc and t_args[0].ndim > 0: + return ErrorInput( + sample, error_regex="Expected result Tensor to be of format CSR" + ) + elif layout is torch.sparse_bsr and t_args[0].ndim > 0: + return ErrorInput( + sample, + error_regex="empty_sparse_compressed expected sparse compressed [(]non-block[)] tensor layout but got SparseBsr", + ) + elif layout is torch.sparse_bsc and t_args[0].ndim > 0: + return ErrorInput( + sample, + error_regex="empty_sparse_compressed expected sparse compressed [(]non-block[)] tensor layout but got SparseBsc", + ) + elif ( + layout is torch.sparse_coo + and dtype is torch.bool + and t_args[0].ndim > 0 + and t_inp.is_cpu + and t_inp.numel() > 0 + and t_inp.dense_dim() > 0 + ): + return ErrorInput( + sample, error_regex="\"addcmul_cpu_out\" not implemented for 'Bool'" + ) + elif ( + layout in {torch.sparse_coo, torch.sparse_csr} + and dtype is torch.bool + and t_inp._nnz() > 0 + and t_args[0].ndim > 0 + and t_inp.is_cpu + and t_inp.numel() > 0 + ): + return ErrorInput( + sample, error_regex="\"mul_out_sparse\" not implemented for 'Bool'" + ) + elif ( + layout is torch.sparse_csr + and t_args[0].layout is torch.strided + and 0 < t_args[0].ndim + and t_args[0].ndim < t_inp.ndim + ): + return ErrorInput( + sample, error_regex="sparse_mask_sparse_csr expects self to be 2D" + ) + elif layout is torch.sparse_csr and ( + (t_args[0].layout is torch.strided and 0 < t_args[0].ndim) + or (t_args[0].layout is layout and t_inp.shape != t_args[0].shape) + ): + return ErrorInput( + sample, + error_regex=( + "expects sparse inputs with equal dimensionality, number of sparse dimensions," + " and shape of sparse dimensions" + ), + ) + elif ( + layout is torch.sparse_csr + and t_inp.dense_dim() > 0 + and t_inp._nnz() > 0 + and t_inp.is_cpu + and dtype is torch.float16 + and t_args[0].ndim > 0 + ): + return ErrorInput( + sample, error_regex="\"addcmul_cpu_out\" not implemented for 'Half'" + ) + return sample + + +@_apply_requires_grad_to_samples +def _maybe_failing_sample_inputs_sparse_elementwise_binary_mul( + op_info, device, dtype, requires_grad, layout, **kwargs +): + """Generator of samples that are known to fail or that were failing in past.""" + # NOTE: When fixing a failing case, remove the Exception comment + # but keep the `yield sample` statement. + + blocksize = (1, 1) if layout in {torch.sparse_bsr, torch.sparse_bsc} else None + regular = torch.tensor([[1, 2], [3, 4]], device=device, dtype=dtype).to_sparse( + layout=layout, dense_dim=0, blocksize=blocksize + ) + batch = torch.tensor( + [[[1, 2], [3, 4]], [[4, 5], [6, 7]]], device=device, dtype=dtype + ).to_sparse(layout=layout, dense_dim=0, blocksize=blocksize) + hybrid = torch.tensor( + [[[1], [2]], [[3], [4]]], device=device, dtype=dtype + ).to_sparse(layout=layout, dense_dim=1, blocksize=blocksize) + + if layout is torch.sparse_csr: + # RuntimeError: crow_indices is supposed to be a vector, but got 2 dimensional tensor + yield SampleInput(batch, args=(batch,)) + # RuntimeError: Only tensors with two sparse dimensions can be + # converted to the SparseCsr layout, got self with 3 sparse + # dimensions. + yield SampleInput( + torch.zeros_like(hybrid).requires_grad_(requires_grad), + args=(torch.zeros_like(hybrid).requires_grad_(requires_grad),), + ) + if dtype is torch.complex32: + # RuntimeError: "mul_out_sparse" not implemented for 'ComplexHalf' + yield SampleInput(regular, args=(regular,)) + if dtype is torch.bool and regular.is_cpu: + # RuntimeError: "mul_out_sparse" not implemented for 'Bool' + yield SampleInput(regular, args=(regular,)) + if layout is torch.sparse_csc: + # RuntimeError: Expected result Tensor to be of format CSR + yield SampleInput(regular, args=(regular,)) + if layout is torch.sparse_bsr: + # RuntimeError: empty_sparse_compressed expected sparse compressed (non-block) tensor layout but got SparseBsr + yield SampleInput(regular, args=(regular,)) + if layout is torch.sparse_bsc: + # RuntimeError: empty_sparse_compressed expected sparse compressed (non-block) tensor layout but got SparseBsc + yield SampleInput(regular, args=(regular,)) + if layout is torch.sparse_coo: + if dtype is torch.complex32: + # RuntimeError: "mul_out_sparse" not implemented for 'ComplexHalf' + yield SampleInput(regular, args=(regular,)) + if dtype is torch.bool and regular.is_cpu: + # RuntimeError: "mul_out_sparse" not implemented for 'Bool' + yield SampleInput(regular, args=(regular,)) + if dtype in {torch.bool, torch.float16} and regular.is_cpu: + # RuntimeError: "addcmul_cpu_out" not implemented for '(Bool|Half)' + yield SampleInput(hybrid, args=(hybrid,)) + + +def _validate_sample_input_sparse_elementwise_binary_operation( + op_info, sample, check_validate=False +): + if op_info.name == "mul": + sample = _validate_sample_input_elementwise_binary_sparse_mul(sample) + + if check_validate: + _check_validate(op_info, sample) + return sample + + +def sample_inputs_sparse_mul(op_info, device, dtype, requires_grad, layout, **kwargs): + """Sample inputs for mul operation on sparse tensors.""" + yield from _sample_inputs_sparse( + sample_inputs_sparse_elementwise_binary_operation, + _maybe_failing_sample_inputs_sparse_elementwise_binary_mul, + _validate_sample_input_sparse_elementwise_binary_operation, + op_info, + device, + dtype, + requires_grad, + layout, + **kwargs, + ) + + +def error_inputs_sparse_mul(op_info, device, layout, **kwargs): + """Error inputs for mul operation on sparse tensors.""" + dtype = torch.float64 + requires_grad = False + yield from _error_inputs_sparse( + _maybe_failing_sample_inputs_sparse_elementwise_binary_mul, + _validate_sample_input_sparse_elementwise_binary_operation, + op_info, + device, + dtype, + requires_grad, + layout, + **kwargs, + ) + + +def _sample_inputs_sparse_like_fns( + op_info, device, dtype, requires_grad, layout, **kwargs +): + from torch.testing._internal.common_utils import TestCase + + for tensor in TestCase().generate_simple_inputs( + layout, + device=device, + dtype=dtype, + enable_batch=True, + enable_hybrid=True, + enable_zero_sized=True, + enable_non_contiguous_indices=False, + enable_non_contiguous_values=False, + ): + yield SampleInput(tensor, args=(), kwargs={}) + yield SampleInput( + tensor, args=(), kwargs=dict(device=device, dtype=dtype, layout=layout) + ) + + hpf = highest_precision_float(device) + if dtype is not hpf: + yield SampleInput(tensor, args=(), kwargs=dict(dtype=hpf)) + + if torch.cuda.is_available(): + other_device = "cuda" if tensor.device.type == "cpu" else "cpu" + yield SampleInput(tensor, args=(), kwargs=dict(device=other_device)) + + if layout is torch.sparse_csr: + other_layout = torch.sparse_csc + elif layout is torch.sparse_csc: + other_layout = torch.sparse_csr + elif layout is torch.sparse_bsr: + other_layout = torch.sparse_bsc + elif layout is torch.sparse_bsc: + other_layout = torch.sparse_bsr + else: + other_layout = torch.strided + yield SampleInput(tensor, args=(), kwargs=dict(layout=other_layout)) + + if layout is not torch.sparse_coo: + yield SampleInput(tensor, args=(), kwargs=dict(layout=torch.sparse_coo)) + + +def _validate_sample_input_sparse_like_fns(op_info, sample, check_validate=False): + if ( + sample.input.layout + in { + torch.sparse_csr, + torch.sparse_csc, + torch.sparse_bsr, + torch.sparse_bsc, + } + and op_info.name != "zeros_like" + ): + if sample.kwargs.get("layout", sample.input.layout) != sample.input.layout: + return ErrorInput( + sample, + error_regex=( + "empty_like with different sparse layout is not supported" + " \\(self is Sparse(Csc|Csr|Bsc|Bsr) but you requested Sparse(Csr|Csc|Bsr|Bsc)\\)" + ), + ) + if sample.input.layout is torch.sparse_coo: + return ErrorInput( + sample, + error_regex=( + "Could not run 'aten::normal_' with arguments from the 'Sparse(CPU|CUDA)' backend." + ), + ) + if check_validate: + _check_validate(op_info, sample) + return sample + + +def _maybe_failing_sample_inputs_sparse_like_fns( + op_info, device, dtype, requires_grad, layout, **kwargs +): + if torch.cuda.is_available() and layout is not torch.sparse_coo: + other_device = "cuda" if torch.device(device).type == "cpu" else "cpu" + if layout is torch.sparse_csr: + other_layout = torch.sparse_csc + elif layout is torch.sparse_csc: + other_layout = torch.sparse_csr + elif layout is torch.sparse_bsr: + other_layout = torch.sparse_bsc + elif layout is torch.sparse_bsc: + other_layout = torch.sparse_bsr + else: + other_layout = torch.strided + + blocksize = (1, 1) if layout in {torch.sparse_bsr, torch.sparse_bsc} else None + + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype, device=device).to_sparse( + layout=layout, blocksize=blocksize + ), + kwargs=dict(device=other_device), + ) + + yield SampleInput( + torch.tensor([[0, 1], [2, 3]], dtype=dtype, device=device).to_sparse( + layout=layout, blocksize=blocksize + ), + kwargs=dict(layout=other_layout), + ) + + +def sample_inputs_sparse_like_fns( + op_info, device, dtype, requires_grad, layout, **kwargs +): + """Sample inputs for like-functions on sparse tensors.""" + yield from _sample_inputs_sparse( + _sample_inputs_sparse_like_fns, + _maybe_failing_sample_inputs_sparse_like_fns, + _validate_sample_input_sparse_like_fns, + op_info, + device, + dtype, + requires_grad, + layout, + **kwargs, + ) + + +def error_inputs_sparse_like_fns(op_info, device, layout, **kwargs): + """Error inputs for like-functions on sparse tensors.""" + dtype = torch.float64 + requires_grad = False + yield from _error_inputs_sparse( + _maybe_failing_sample_inputs_sparse_like_fns, + _validate_sample_input_sparse_like_fns, + op_info, + device, + dtype, + requires_grad, + layout, + **kwargs, + ) + + +def _validate_sample_input_sparse_default(op_info, sample, check_validate=False): + if op_info.name == "to_sparse": + if ( + sample.input.layout + in {torch.sparse_csr, torch.sparse_csc, torch.sparse_bsr, torch.sparse_bsc} + and len(sample.args) == 1 + and isinstance(sample.args[0], int) + and sample.args[0] != 2 + ): + sample = ErrorInput( + sample, + error_regex="sparse dim argument must be 2 for sparse_compressed_to_sparse", + ) + + if check_validate: + _check_validate(op_info, sample) + return sample + + +def validate_sample_input_sparse(op_info, sample, check_validate=False): + """Return the specified sample when it is valid and supported by the + operation. Otherwise, return the sample as ErrorInput instance. + + When check_validate is True, the result is validated against + calling the op on the sample. + """ + if isinstance(op_info, ReductionOpInfo): + return _validate_sample_input_sparse_reduction( + op_info, sample, check_validate=check_validate + ) + elif isinstance(op_info, BinaryUfuncInfo): + return _validate_sample_input_sparse_elementwise_binary_operation( + op_info, sample, check_validate=check_validate + ) + else: + return _validate_sample_input_sparse_default( + op_info, sample, check_validate=check_validate + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/special.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/special.py new file mode 100644 index 0000000000000000000000000000000000000000..1626ead643244354bbbe5f25c144750c196ae44e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/definitions/special.py @@ -0,0 +1,1057 @@ +# mypy: ignore-errors + +import unittest +from functools import partial +from itertools import product + +import numpy as np + +import torch +from torch.testing import make_tensor +from torch.testing._internal.common_device_type import ( + precisionOverride, + tol, + toleranceOverride, +) +from torch.testing._internal.common_dtype import all_types_and, floating_types +from torch.testing._internal.common_utils import TEST_SCIPY, torch_to_numpy_dtype_dict +from torch.testing._internal.opinfo.core import ( + BinaryUfuncInfo, + DecorateInfo, + L, + NumericsFilter, + OpInfo, + S, + SampleInput, + UnaryUfuncInfo, +) +from torch.testing._internal.opinfo.refs import ( + ElementwiseBinaryPythonRefInfo, + ElementwiseUnaryPythonRefInfo, +) +from torch.testing._internal.opinfo.utils import ( + np_unary_ufunc_integer_promotion_wrapper, +) + + +if TEST_SCIPY: + import scipy.special + + +# TODO: Consolidate `i0e` with sample_inputs_unary when `make_tensor`, +# supports `exclude` argument. +# For more context: https://github.com/pytorch/pytorch/pull/56352#discussion_r633277617 +def sample_inputs_i0_i1(op_info, device, dtype, requires_grad, **kwargs): + exclude_zero = requires_grad and op_info.op is torch.special.i0e + make_arg = partial( + make_tensor, + dtype=dtype, + device=device, + requires_grad=requires_grad, + exclude_zero=exclude_zero, + ) + yield SampleInput(make_arg((S,))) + yield SampleInput(make_arg(())) + + if requires_grad and not exclude_zero: + # Special Case for gradient + # Sample with `0` in the input + t = make_arg((S,)) + t[0] = 0 + + yield SampleInput(t) + + +def sample_inputs_polygamma(op_info, device, dtype, requires_grad, **kwargs): + make_arg = partial( + make_tensor, + device=device, + # TODO: eliminate low after gh-106692 is fixed: + low=(1 if dtype in {torch.int32, torch.int64} else None), + dtype=dtype, + requires_grad=requires_grad, + ) + tensor_shapes = ((S, S), ()) + ns = (1, 2, 3, 4, 5) + + for shape, n in product(tensor_shapes, ns): + yield SampleInput(make_arg(shape), args=(n,)) + + +def reference_polygamma(x, n): + # WEIRD `scipy.special.polygamma` behavior + # >>> scipy.special.polygamma(0, np.array(501, dtype=np.float32)).dtype + # dtype('float64') + # >>> scipy.special.polygamma(0, np.array([501], dtype=np.float32)).dtype + # dtype('float32') + # + # Thus we cast output to the default torch dtype or preserve double + result_dtype = torch_to_numpy_dtype_dict[torch.get_default_dtype()] + if x.dtype == np.double: + result_dtype = np.double + return scipy.special.polygamma(n, x).astype(result_dtype) + + +def sample_inputs_entr(op_info, device, dtype, requires_grad, **kwargs): + low, _ = op_info.domain + + if requires_grad: + low = 0 + op_info._domain_eps + + make_arg = partial( + make_tensor, dtype=dtype, device=device, low=low, requires_grad=requires_grad + ) + yield SampleInput(make_arg((L,))) + yield SampleInput(make_arg(())) + + +def sample_inputs_erfcx(op_info, device, dtype, requires_grad, **kwargs): + for shape in ((L,), (1, 0, 3), ()): + yield SampleInput( + make_tensor( + shape, + device=device, + dtype=dtype, + low=-5, + requires_grad=requires_grad, + ), + ) + + +_unsigned_int_types = (torch.uint16, torch.uint32, torch.uint64) + + +op_db: list[OpInfo] = [ + UnaryUfuncInfo( + "special.i0e", + aten_name="special_i0e", + ref=scipy.special.i0e if TEST_SCIPY else None, + decorators=(precisionOverride({torch.bfloat16: 3e-1, torch.float16: 3e-1}),), + dtypes=all_types_and( + torch.bool, torch.half, torch.bfloat16, *_unsigned_int_types + ), + dtypesIfMPS=all_types_and(torch.bool, torch.half, torch.bfloat16), + sample_inputs_func=sample_inputs_i0_i1, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ), + UnaryUfuncInfo( + "special.i1", + aten_name="special_i1", + ref=np_unary_ufunc_integer_promotion_wrapper(scipy.special.i1) + if TEST_SCIPY + else None, + dtypes=all_types_and( + torch.bool, torch.half, torch.bfloat16, *_unsigned_int_types + ), + dtypesIfMPS=all_types_and(torch.bool, torch.half, torch.bfloat16), + backward_dtypes=floating_types(), + sample_inputs_func=sample_inputs_i0_i1, + decorators=( + DecorateInfo( + toleranceOverride( + { + torch.float32: tol(atol=1e-4, rtol=0), + torch.bool: tol(atol=1e-4, rtol=0), + } + ) + ), + ), + skips=( + DecorateInfo( + unittest.skip("Incorrect result!"), + "TestUnaryUfuncs", + "test_reference_numerics_large", + dtypes=(torch.int8,), + ), + ), + supports_fwgrad_bwgrad=True, + supports_forward_ad=True, + ), + UnaryUfuncInfo( + "special.i1e", + aten_name="special_i1e", + ref=scipy.special.i1e if TEST_SCIPY else None, + dtypes=all_types_and( + torch.bool, torch.half, torch.bfloat16, *_unsigned_int_types + ), + dtypesIfMPS=all_types_and(torch.bool, torch.half, torch.bfloat16), + backward_dtypes=floating_types(), + sample_inputs_func=sample_inputs_i0_i1, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + ), + UnaryUfuncInfo( + "special.ndtr", + aten_name="special_ndtr", + decorators=(precisionOverride({torch.bfloat16: 5e-3, torch.float16: 5e-4}),), + ref=scipy.special.ndtr if TEST_SCIPY else None, + dtypes=all_types_and( + torch.bool, torch.half, torch.bfloat16, *_unsigned_int_types + ), + dtypesIfMPS=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # Dispatch stub: unsupported device typemeta + DecorateInfo( + unittest.expectedFailure, + "TestFwdGradients", + "test_fn_fwgrad_bwgrad", + device_type="meta", + ), + ), + ), + # A separate OpInfo entry for special.polygamma is needed to reorder the arguments + # for the alias. See the discussion here: https://github.com/pytorch/pytorch/pull/59691#discussion_r650261939 + UnaryUfuncInfo( + "special.polygamma", + op=lambda x, n, **kwargs: torch.special.polygamma(n, x, **kwargs), + variant_test_name="special_polygamma_n_0", + ref=reference_polygamma if TEST_SCIPY else None, + dtypes=all_types_and(torch.bool, torch.half, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_polygamma, + skips=( + # lambda impl + DecorateInfo( + unittest.expectedFailure, "TestJit", "test_variant_consistency_jit" + ), + DecorateInfo( + unittest.expectedFailure, + "TestNormalizeOperators", + "test_normalize_operator_exhaustive", + ), + ), + sample_kwargs=lambda device, dtype, input: ({"n": 0}, {"n": 0}), + # polygamma functions have multiple singularities at x having non-positive integer value + reference_numerics_filter=NumericsFilter( + condition=lambda x: (x < 0.1) & ((x - x.round()).abs() < 1e-4), safe_val=1 + ), + ), + BinaryUfuncInfo( + "special.xlog1py", + aten_name="special_xlog1py", + dtypes=all_types_and( + torch.bool, torch.half, torch.bfloat16, *_unsigned_int_types + ), + dtypesIfMPS=all_types_and(torch.bool, torch.half, torch.bfloat16), + promotes_int_to_float=True, + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + supports_one_python_scalar=True, + # We don't test -1 as the gradient will be NaN and it'll break + rhs_make_tensor_kwargs=dict(low=-0.99), + ), + BinaryUfuncInfo( + "special.zeta", + aten_name="special_zeta", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + supports_autograd=False, + supports_one_python_scalar=True, + skips=( + # Reference reference_inputs nans and infs on cuda and nan, inf, 0., -inf for cpu + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + ), + # TODO: FIXME + # OpInfo entry to verify the gradient formula of `other`/`q` + # BinaryUfuncInfo('special.zeta', + # op=lambda q, x, **kwargs: torch.special.zeta(x, q, **kwargs), + # aten_name='special_zeta', + # variant_test_name='grad', + # dtypes=all_types_and(torch.bool), + # promotes_int_to_float=True, + # supports_autograd=True, + # supports_rhs_python_scalar=False, + # decorators=[ + # # Derivative wrt first tensor not implemented + # DecorateInfo(unittest.expectedFailure, "TestCommon", + # "test_floating_inputs_are_differentiable") + # ], + # skips=( + # # Lambda doesn't work in JIT test + # # AssertionError: JIT Test does not execute any logic + # DecorateInfo(unittest.skip("Skipped!"), "TestJit", "test_variant_consistency_jit"), + # )), + UnaryUfuncInfo( + "special.entr", + ref=scipy.special.entr if TEST_SCIPY else None, + aten_name="special_entr", + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + decorators=(precisionOverride({torch.float16: 1e-1, torch.bfloat16: 1e-1}),), + dtypes=all_types_and( + torch.bool, torch.half, torch.bfloat16, *_unsigned_int_types + ), + dtypesIfMPS=all_types_and(torch.bool, torch.half, torch.bfloat16), + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestUnaryUfuncs", + "test_reference_numerics_large", + dtypes=[torch.bfloat16, torch.float16], + ), + ), + supports_inplace_autograd=False, + sample_inputs_func=sample_inputs_entr, + ), + UnaryUfuncInfo( + "special.ndtri", + ref=scipy.special.ndtri if TEST_SCIPY else None, + domain=(0, 1), + aten_name="special_ndtri", + dtypes=all_types_and(torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # The operator 'aten::special_ndtri.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + UnaryUfuncInfo( + "special.log_ndtr", + aten_name="special_log_ndtr", + ref=scipy.special.log_ndtr if TEST_SCIPY else None, + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + skips=( + # The operator 'aten::special_log_ndtr.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + UnaryUfuncInfo( + "special.erfcx", + ref=scipy.special.erfcx if TEST_SCIPY else None, + aten_name="special_erfcx", + decorators=( + toleranceOverride( + { + torch.float32: tol(atol=0, rtol=4e-6), + } + ), + ), + skips=( + # Exception: Tensor-likes are not close! + # Greatest absolute difference: inf at index (10,) (up to 1e-05 allowed) + # Greatest relative difference: nan at index (10,) (up to 0.001 allowed) + DecorateInfo( + unittest.expectedFailure, + "TestConsistency", + "test_output_grad_match", + device_type="mps", + dtypes=(torch.float16,), + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + supports_forward_ad=True, + supports_fwgrad_bwgrad=True, + sample_inputs_func=sample_inputs_erfcx, + ), + UnaryUfuncInfo( + "special.airy_ai", + decorators=( + precisionOverride( + { + torch.float32: 1e-03, + torch.float64: 1e-05, + }, + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool), + ref=lambda x: scipy.special.airy(x)[0] if TEST_SCIPY else None, + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestUnaryUfuncs", + "test_reference_numerics_large", + ), + # NotImplementedError: The operator 'aten::special_airy_ai.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.bessel_j0", + decorators=( + precisionOverride( + { + torch.float32: 1e-04, + torch.float64: 1e-05, + }, + ), + # Comparing fp32 CPU grads with fp16 MPS ones leads to high errors + DecorateInfo( + toleranceOverride( + { + torch.float16: tol(atol=1e-2, rtol=1e-4), + torch.bfloat16: tol(atol=1e-2, rtol=1e-4), + } + ), + "TestConsistency", + "test_output_match", + device_type="mps", + ), + DecorateInfo( + toleranceOverride( + { + torch.float16: tol(atol=1e-2, rtol=1e-4), + torch.bfloat16: tol(atol=1e-2, rtol=1e-4), + } + ), + "TestConsistency", + "test_output_grad_match", + device_type="mps", + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=scipy.special.j0 if TEST_SCIPY else None, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.bessel_j1", + decorators=( + precisionOverride( + { + torch.float32: 1e-04, + torch.float64: 1e-05, + }, + ), + # Comparing fp32 CPU grads with fp16 MPS ones leads to high errors + DecorateInfo( + toleranceOverride( + { + torch.float16: tol(atol=1e-2, rtol=1e-4), + torch.bfloat16: tol(atol=1e-2, rtol=1e-4), + } + ), + "TestConsistency", + "test_output_match", + device_type="mps", + ), + DecorateInfo( + toleranceOverride({torch.float16: tol(atol=0.0004, rtol=0.009)}), + "TestConsistency", + "test_output_grad_match", + device_type="mps", + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=scipy.special.j1 if TEST_SCIPY else None, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.bessel_y0", + decorators=( + precisionOverride( + { + torch.float32: 1e-04, + torch.float64: 1e-05, + }, + ), + # Comparing fp32 CPU grads with fp16 MPS ones leads to high errors + DecorateInfo( + toleranceOverride( + { + torch.float16: tol(atol=1e-2, rtol=1e-4), + torch.bfloat16: tol(atol=1e-2, rtol=1e-4), + } + ), + "TestConsistency", + "test_output_match", + device_type="mps", + ), + ), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + ref=scipy.special.y0 if TEST_SCIPY else None, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.bessel_y1", + decorators=( + precisionOverride( + { + torch.float32: 1e-04, + torch.float64: 1e-05, + }, + ), + # Comparing fp32 CPU grads with fp16 MPS ones leads to high errors + DecorateInfo( + toleranceOverride( + { + torch.float16: tol(atol=1e-2, rtol=1e-4), + torch.bfloat16: tol(atol=1e-2, rtol=1e-4), + } + ), + "TestConsistency", + "test_output_match", + device_type="mps", + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=scipy.special.y1 if TEST_SCIPY else None, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.chebyshev_polynomial_t", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.chebyshev_polynomial_u", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.chebyshev_polynomial_v", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.chebyshev_polynomial_w", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.hermite_polynomial_h", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: inf + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + # Too slow + DecorateInfo( + unittest.skip, "TestCommon", "test_compare_cpu", device_type="xpu" + ), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.hermite_polynomial_he", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: inf + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.laguerre_polynomial_l", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + # Too slow + DecorateInfo( + unittest.skip, "TestCommon", "test_compare_cpu", device_type="xpu" + ), + # NotImplementedError: The operator 'aten::special_laguerre_polynomial_l.out' + # is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_promotes_int_to_float", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.legendre_polynomial_p", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + # NotImplementedError: The operator 'aten::special_legendre_polynomial_p.out' + # is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_variant_consistency_eager", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_promotes_int_to_float", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_out_warning", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_out", device_type="mps" + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_noncontiguous_samples", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, "TestCommon", "test_dtypes", device_type="mps" + ), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.modified_bessel_i0", + decorators=( + precisionOverride( + { + torch.float32: 1e-03, + torch.float64: 1e-05, + }, + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=scipy.special.i0 if TEST_SCIPY else None, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.modified_bessel_i1", + decorators=( + precisionOverride( + { + torch.float32: 1e-03, + torch.float64: 1e-05, + }, + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=scipy.special.i1 if TEST_SCIPY else None, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.modified_bessel_k0", + decorators=( + precisionOverride( + { + torch.float32: 1e-03, + torch.float64: 1e-05, + }, + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=scipy.special.k0 if TEST_SCIPY else None, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.modified_bessel_k1", + decorators=( + precisionOverride( + { + torch.float32: 1e-03, + torch.float64: 1e-05, + }, + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=scipy.special.k1 if TEST_SCIPY else None, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.scaled_modified_bessel_k0", + decorators=( + toleranceOverride( + { + torch.float32: tol(atol=1e-03, rtol=1e-03), + torch.float64: tol(atol=1e-05, rtol=1e-03), + } + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=scipy.special.k0e if TEST_SCIPY else None, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.scaled_modified_bessel_k1", + decorators=( + toleranceOverride( + { + torch.float32: tol(atol=1e-03, rtol=1e-03), + torch.float64: tol(atol=1e-05, rtol=1e-03), + } + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=scipy.special.k1e if TEST_SCIPY else None, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.shifted_chebyshev_polynomial_t", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.shifted_chebyshev_polynomial_u", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.shifted_chebyshev_polynomial_v", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + BinaryUfuncInfo( + "special.shifted_chebyshev_polynomial_w", + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + promotes_int_to_float=True, + skips=( + DecorateInfo(unittest.skip("Skipped!"), "TestCudaFuserOpInfo"), + DecorateInfo(unittest.skip("Skipped!"), "TestNNCOpInfo"), + # Greatest absolute difference: nan + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + supports_one_python_scalar=True, + supports_autograd=False, + ), + UnaryUfuncInfo( + "special.spherical_bessel_j0", + decorators=( + toleranceOverride( + { + torch.float16: tol(atol=1e-02, rtol=1e-02), + torch.bfloat16: tol(atol=1e-02, rtol=1e-02), + torch.float32: tol(atol=1e-03, rtol=1e-03), + torch.float64: tol(atol=1e-05, rtol=1e-03), + } + ), + ), + dtypes=all_types_and(torch.bool, *_unsigned_int_types), + dtypesIfMPS=all_types_and(torch.bool, torch.float16, torch.bfloat16), + ref=lambda x: scipy.special.spherical_jn(0, x) if TEST_SCIPY else None, + supports_autograd=False, + skips=( + DecorateInfo( + unittest.skip( + "Scipy doesn't support bool inputs to spherical_bessel_j0" + ), + "TestUnaryUfuncs", + "test_reference_numerics_normal", + dtypes=(torch.bool,), + ), + ), + ), +] + +python_ref_db: list[OpInfo] = [ + # + # Elementwise Unary Special OpInfos + # + ElementwiseUnaryPythonRefInfo( + "_refs.special.bessel_j0", + torch_opinfo_name="special.bessel_j0", + op_db=op_db, + decorators=( + precisionOverride( + { + torch.float32: 1e-04, + torch.float64: 1e-05, + }, + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.bessel_j1", + torch_opinfo_name="special.bessel_j1", + op_db=op_db, + decorators=( + precisionOverride( + { + torch.float32: 1e-04, + torch.float64: 1e-05, + }, + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.entr", + torch_opinfo_name="special.entr", + op_db=op_db, + decorators=(precisionOverride({torch.float16: 1e-1, torch.bfloat16: 1e-1}),), + skips=( + DecorateInfo( + unittest.skip("Skipped!"), + "TestUnaryUfuncs", + "test_reference_numerics_large", + dtypes=[torch.bfloat16, torch.float16], + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.erfcx", + torch_opinfo_name="special.erfcx", + op_db=op_db, + decorators=( + toleranceOverride( + { + torch.float32: tol(atol=0, rtol=4e-6), + } + ), + ), + skips=( + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref", + device_type="mps", + dtypes=(torch.float16,), + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref_torch_fallback", + device_type="mps", + dtypes=(torch.float16,), + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.i0e", + torch_opinfo_name="special.i0e", + op_db=op_db, + decorators=(precisionOverride({torch.bfloat16: 3e-1, torch.float16: 3e-1}),), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.i1", + torch_opinfo_name="special.i1", + op_db=op_db, + decorators=( + DecorateInfo( + toleranceOverride( + { + torch.float32: tol(atol=1e-4, rtol=0), + torch.bool: tol(atol=1e-4, rtol=0), + } + ) + ), + ), + skips=( + DecorateInfo( + unittest.skip("Incorrect result!"), + "TestUnaryUfuncs", + "test_reference_numerics_large", + dtypes=(torch.int8,), + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.i1e", + torch_opinfo_name="special.i1e", + op_db=op_db, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.log_ndtr", + torch_opinfo_name="special.log_ndtr", + op_db=op_db, + skips=( + # The operator 'aten::special_log_ndtr.out' is not currently implemented for the MPS device + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref", + device_type="mps", + ), + DecorateInfo( + unittest.expectedFailure, + "TestCommon", + "test_python_ref_torch_fallback", + device_type="mps", + ), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.ndtr", + torch_opinfo_name="special.ndtr", + op_db=op_db, + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.ndtri", + torch_opinfo_name="special.ndtri", + op_db=op_db, + skips=( + # The operator 'aten::special_ndtri.out' is not currently implemented for the MPS device + DecorateInfo(unittest.expectedFailure, "TestCommon", device_type="mps"), + ), + ), + ElementwiseUnaryPythonRefInfo( + "_refs.special.spherical_bessel_j0", + torch_opinfo_name="special.spherical_bessel_j0", + op_db=op_db, + decorators=( + toleranceOverride( + { + torch.float32: tol(atol=1e-03, rtol=1e-03), + torch.float64: tol(atol=1e-05, rtol=1e-03), + } + ), + ), + skips=( + DecorateInfo( + unittest.skip( + "Scipy doesn't support bool inputs to spherical_bessel_j0" + ), + "TestUnaryUfuncs", + "test_reference_numerics_normal", + dtypes=(torch.bool,), + ), + # Seems to fail on M2 but not M4: + # Exception: Cannot convert a MPS Tensor to float64 dtype as the MPS framework doesn't support float64 + DecorateInfo( + unittest.skip("Platform-specific error"), + "TestCommon", + "test_python_ref", + device_type="mps", + dtypes=(torch.float16,), + ), + DecorateInfo( + unittest.skip("Platform-specific error"), + "TestCommon", + "test_python_ref_torch_fallback", + device_type="mps", + dtypes=(torch.float16,), + ), + ), + ), + # + # Elementwise Binary Special OpInfos + # + ElementwiseBinaryPythonRefInfo( + "_refs.special.zeta", + torch_opinfo_name="special.zeta", + supports_one_python_scalar=True, + op_db=op_db, + skips=( + # Reference reference_inputs nans and infs on cuda and nan, inf, 0., -inf for cpu + DecorateInfo(unittest.expectedFailure, "TestCommon", "test_compare_cpu"), + ), + ), +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/refs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/refs.py new file mode 100644 index 0000000000000000000000000000000000000000..12c672fb282f76c288f44f61c06bd3b058284e54 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/refs.py @@ -0,0 +1,219 @@ +# mypy: ignore-errors + +from torch.testing._internal.opinfo.core import ( + BinaryUfuncInfo, + OpInfo, + ReductionOpInfo, + UnaryUfuncInfo, +) + + +# NOTE [Python References] +# Python References emulate existing PyTorch operations, but can ultimately +# be expressed in terms of "primitive" operations from torch._prims. +# +# These references are experimental. +# See https://dev-discuss.pytorch.org/t/tracing-with-primitives-update-0/577 +# for additional context. +# +# Python Reference OpInfos should be added to the python_ref_db list below. +# Tests can opt-into running on these references by including +# that list in the Sequence they pass to the @ops decorator. +# +# When a Python Reference OpInfo is constructed a pointer to an +# existing OpInfo must be provided using the torch_opinfo_name kwarg. +# The existing OpInfo with that name and no variant will be found +# to inherit from. +# +# Instead of just inheriting the existing OpInfo's metadata, the +# Python Reference OpInfos inherit the existing OpInfo's +# construction arguments. These arguments can be overridden +# by adding kwargs to the constructor. + + +def _find_referenced_opinfo(referenced_name, variant_name, *, op_db=None): + """ + Finds the OpInfo with the given name that has no variant name. + """ + # NOTE: searching the global op_db doesn't work when OpInfos are split into + # different modules, as otherwise the op_db will not be fully constructed + # yet. So, instead the local op_db must be passed in explicitly. + if op_db is None: + from torch.testing._internal.common_methods_invocations import op_db + + for opinfo in op_db: + if opinfo.name == referenced_name and opinfo.variant_test_name == variant_name: + return opinfo + + +def _inherit_constructor_args(name, op, inherited, overrides): + # inherits metadata + common_kwargs = { + "name": name, + "op": op, + "aliases": None, # TODO add a check for alias coverage + "method_variant": None, + "inplace_variant": None, # TODO: add a check for inplace coverage + "supports_scripting": False, + } + + # Acquires inherited kwargs + kwargs = inherited.copy() + + # Fixes metadata + if "kwargs" in kwargs: + kwargs.update(kwargs["kwargs"]) + del kwargs["kwargs"] + if "self" in kwargs: + del kwargs["self"] + if "__class__" in kwargs: + del kwargs["__class__"] + if "skips" in kwargs: + del kwargs["skips"] + if "decorators" in kwargs: + del kwargs["decorators"] + + # Overrides metadata + kwargs.update(common_kwargs) + kwargs.update(overrides) + + # At the moment no prims support autograd, so we must not run autograd + # tests e.g. when testing dtype support. Once we start writing autograd + # formulas for prims this can be removed. + kwargs["supports_autograd"] = False + kwargs["supports_gradgrad"] = False + kwargs["supports_fwgrad_bwgrad"] = False + kwargs["supports_inplace_autograd"] = False + kwargs["supports_forward_ad"] = False + + return kwargs + + +class PythonRefInfo(OpInfo): + """ + An OpInfo for a Python reference of an OpInfo base class operation. + """ + + def __init__( + self, + name, # the stringname of the callable Python reference + *, + op=None, # the function variant of the operation, populated as torch. if None + op_db=None, # The database of opinfos to search for the parent opinfo + torch_opinfo_name, # the string name of the corresponding torch opinfo + torch_opinfo_variant_name="", # the variant name for corresponding torch opinfo + validate_view_consistency=True, + **kwargs, + ): # additional kwargs override kwargs inherited from the torch opinfo + self.torch_opinfo_name = torch_opinfo_name + self.torch_opinfo_variant_name = torch_opinfo_variant_name + self.torch_opinfo = _find_referenced_opinfo( + torch_opinfo_name, torch_opinfo_variant_name, op_db=op_db + ) + self.validate_view_consistency = validate_view_consistency + if not isinstance(self.torch_opinfo, OpInfo): + raise AssertionError( + f"Expected torch_opinfo to be OpInfo, got {type(self.torch_opinfo)}" + ) + + inherited = self.torch_opinfo._original_opinfo_args + ukwargs = _inherit_constructor_args(name, op, inherited, kwargs) + super().__init__(**ukwargs) + + +class ReductionPythonRefInfo(ReductionOpInfo): + """ + An OpInfo for a Python reference of an elementwise unary operation. + """ + + def __init__( + self, + name, # the stringname of the callable Python reference + *, + op=None, # the function variant of the operation, populated as torch. if None + op_db=None, # The database of opinfos to search for the parent opinfo + torch_opinfo_name, # the string name of the corresponding torch opinfo + torch_opinfo_variant_name="", # the variant name for corresponding torch opinfo + **kwargs, + ): # additional kwargs override kwargs inherited from the torch opinfo + self.torch_opinfo_name = torch_opinfo_name + self.torch_opinfo_variant_name = torch_opinfo_variant_name + self.torch_opinfo = _find_referenced_opinfo( + torch_opinfo_name, torch_opinfo_variant_name, op_db=op_db + ) + if not isinstance(self.torch_opinfo, ReductionOpInfo): + raise AssertionError( + f"Expected torch_opinfo to be ReductionOpInfo, got {type(self.torch_opinfo)}" + ) + + inherited = self.torch_opinfo._original_reduction_args + ukwargs = _inherit_constructor_args(name, op, inherited, kwargs) + + # See https://github.com/pytorch/pytorch/issues/77216 + self.validate_view_consistency = False + + super().__init__(**ukwargs) + + +class ElementwiseUnaryPythonRefInfo(UnaryUfuncInfo): + """ + An OpInfo for a Python reference of an elementwise unary operation. + """ + + def __init__( + self, + name, # the stringname of the callable Python reference + *, + op=None, # the function variant of the operation, populated as torch. if None + op_db=None, # The database of opinfos to search for the parent opinfo + torch_opinfo_name, # the string name of the corresponding torch opinfo + torch_opinfo_variant_name="", # the variant name for corresponding torch opinfo + validate_view_consistency=True, + **kwargs, + ): # additional kwargs override kwargs inherited from the torch opinfo + self.torch_opinfo_name = torch_opinfo_name + self.torch_opinfo_variant_name = torch_opinfo_variant_name + self.torch_opinfo = _find_referenced_opinfo( + torch_opinfo_name, torch_opinfo_variant_name, op_db=op_db + ) + self.validate_view_consistency = validate_view_consistency + if not isinstance(self.torch_opinfo, UnaryUfuncInfo): + raise AssertionError( + f"Expected torch_opinfo to be UnaryUfuncInfo, got {type(self.torch_opinfo)}" + ) + + inherited = self.torch_opinfo._original_unary_ufunc_args + ukwargs = _inherit_constructor_args(name, op, inherited, kwargs) + + super().__init__(**ukwargs) + + +class ElementwiseBinaryPythonRefInfo(BinaryUfuncInfo): + """ + An OpInfo for a Python reference of an elementwise binary operation. + """ + + def __init__( + self, + name, # the stringname of the callable Python reference + *, + op=None, # the function variant of the operation, populated as torch. if None + op_db=None, # The database of opinfos to search for the parent opinfo + torch_opinfo_name, # the string name of the corresponding torch opinfo + torch_opinfo_variant_name="", # the variant name for corresponding torch opinfo + **kwargs, + ): # additional kwargs override kwargs inherited from the torch opinfo + self.torch_opinfo_name = torch_opinfo_name + self.torch_opinfo_variant_name = torch_opinfo_variant_name + self.torch_opinfo = _find_referenced_opinfo( + torch_opinfo_name, torch_opinfo_variant_name, op_db=op_db + ) + if not isinstance(self.torch_opinfo, BinaryUfuncInfo): + raise AssertionError( + f"Expected torch_opinfo to be BinaryUfuncInfo, got {type(self.torch_opinfo)}" + ) + + inherited = self.torch_opinfo._original_binary_ufunc_args + ukwargs = _inherit_constructor_args(name, op, inherited, kwargs) + + super().__init__(**ukwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..38e09c8587abd52b314d969ca7f1c53ed676a2de --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/opinfo/utils.py @@ -0,0 +1,282 @@ +# mypy: ignore-errors + +import collections +import warnings +from collections.abc import Sequence +from functools import partial, wraps + +import numpy as np +import numpy.typing as npt + +import torch +from torch.testing._internal.common_cuda import TEST_CUDA +from torch.testing._internal.common_dtype import ( + _dispatch_dtypes, + all_types, + all_types_and, + all_types_and_complex, + all_types_and_complex_and, + all_types_and_half, + complex_types, + floating_and_complex_types, + floating_and_complex_types_and, + floating_types, + floating_types_and, + floating_types_and_half, + integral_types, + integral_types_and, +) +from torch.testing._internal.common_utils import torch_to_numpy_dtype_dict + + +COMPLETE_DTYPES_DISPATCH = ( + all_types, + all_types_and_complex, + all_types_and_half, + floating_types, + floating_and_complex_types, + floating_types_and_half, + integral_types, + complex_types, +) + +EXTENSIBLE_DTYPE_DISPATCH = ( + all_types_and_complex_and, + floating_types_and, + floating_and_complex_types_and, + integral_types_and, + all_types_and, +) + +# Better way to acquire devices? +DEVICES = ["cpu"] + (["cuda"] if TEST_CUDA else []) + + +class _dynamic_dispatch_dtypes(_dispatch_dtypes): + # Class to tag the dynamically generated types. + pass + + +def get_supported_dtypes(op, sample_inputs_fn, device_type): + # Returns the supported dtypes for the given operator and device_type pair. + if device_type not in ["cpu", "cuda"]: + raise AssertionError( + f"Expected device_type in ['cpu', 'cuda'], got {device_type!r}" + ) + if not TEST_CUDA and device_type == "cuda": + warnings.warn( + "WARNING: CUDA is not available, empty_dtypes dispatch will be returned!", + stacklevel=2, + ) + return _dynamic_dispatch_dtypes(()) + + supported_dtypes = set() + for dtype in all_types_and_complex_and(torch.bool, torch.bfloat16, torch.half): + try: + samples = sample_inputs_fn(op, device_type, dtype, False) + except RuntimeError: + # If `sample_inputs_fn` doesn't support sampling for a given + # `dtype`, we assume that the `dtype` is not supported. + # We raise a warning, so that user knows that this was the case + # and can investigate if there was an issue with the `sample_inputs_fn`. + warnings.warn( + f"WARNING: Unable to generate sample for device:{device_type} and dtype:{dtype}", + stacklevel=2, + ) + continue + + # We assume the dtype is supported + # only if all samples pass for the given dtype. + supported = True + for sample in samples: + try: + op(sample.input, *sample.args, **sample.kwargs) + except RuntimeError: + # dtype is not supported + supported = False + break + + if supported: + supported_dtypes.add(dtype) + + return _dynamic_dispatch_dtypes(supported_dtypes) + + +def dtypes_dispatch_hint(dtypes): + # Function returns the appropriate dispatch function (from COMPLETE_DTYPES_DISPATCH and EXTENSIBLE_DTYPE_DISPATCH) + # and its string representation for the passed `dtypes`. + return_type = collections.namedtuple("return_type", "dispatch_fn dispatch_fn_str") + + # CUDA is not available, dtypes will be empty. + if len(dtypes) == 0: + return return_type((), "()") + + set_dtypes = set(dtypes) + for dispatch in COMPLETE_DTYPES_DISPATCH: + # Short circuit if we get an exact match. + if set(dispatch()) == set_dtypes: + return return_type(dispatch, dispatch.__name__ + "()") + + chosen_dispatch = None + chosen_dispatch_score = 0.0 + for dispatch in EXTENSIBLE_DTYPE_DISPATCH: + dispatch_dtypes = set(dispatch()) + if not dispatch_dtypes.issubset(set_dtypes): + continue + + score = len(dispatch_dtypes) + if score > chosen_dispatch_score: + chosen_dispatch_score = score + chosen_dispatch = dispatch + + # If user passed dtypes which are lower than the lowest + # dispatch type available (not likely but possible in code path). + if chosen_dispatch is None: + return return_type((), str(dtypes)) + + return return_type( + partial(dispatch, *tuple(set(dtypes) - set(dispatch()))), + dispatch.__name__ + str(tuple(set(dtypes) - set(dispatch()))), + ) + + +def is_dynamic_dtype_set(op): + # Detect if the OpInfo entry acquired dtypes dynamically + # using `get_supported_dtypes`. + return op.dynamic_dtypes + + +def str_format_dynamic_dtype(op): + fmt_str = f""" + OpInfo({op.name}, + dtypes={dtypes_dispatch_hint(op.dtypes).dispatch_fn_str}, + dtypesIfCUDA={dtypes_dispatch_hint(op.dtypesIfCUDA).dispatch_fn_str}, + ) + """ + + return fmt_str + + +def np_unary_ufunc_integer_promotion_wrapper(fn): + # Wrapper that passes PyTorch's default scalar + # type as an argument to the wrapped NumPy + # unary ufunc when given an integer input. + # This mimics PyTorch's integer->floating point + # type promotion. + # + # This is necessary when NumPy promotes + # integer types to double, since PyTorch promotes + # integer types to the default scalar type. + + # Helper to determine if promotion is needed + def is_integral(dtype): + return dtype in [ + np.bool_, + bool, + np.uint8, + np.int8, + np.int16, + np.int32, + np.int64, + ] + + @wraps(fn) + def wrapped_fn(x): + # As the default dtype can change, acquire it when function is called. + # NOTE: Promotion in PyTorch is from integer types to the default dtype + np_dtype = torch_to_numpy_dtype_dict[torch.get_default_dtype()] + + if is_integral(x.dtype): + return fn(x.astype(np_dtype)) + return fn(x) + + return wrapped_fn + + +def reference_reduction_numpy(f, supports_keepdims=True): + """Wraps a NumPy reduction operator. + + The wrapper function will forward dim, keepdim, mask, and identity + kwargs to the wrapped function as the NumPy equivalent axis, + keepdims, where, and initiak kwargs, respectively. + + Args: + f: NumPy reduction operator to wrap + supports_keepdims (bool, optional): Whether the NumPy operator accepts + keepdims parameter. If it does not, the wrapper will manually unsqueeze + the reduced dimensions if it was called with keepdim=True. Defaults to True. + + Returns: + Wrapped function + + """ + + @wraps(f) + def wrapper(x: npt.NDArray, *args, **kwargs): + # Copy keys into a set + keys = set(kwargs.keys()) + + dim = kwargs.pop("dim", None) + keepdim = kwargs.pop("keepdim", False) + + if "dim" in keys: + dim = tuple(dim) if isinstance(dim, Sequence) else dim + + # NumPy reductions don't accept dim=0 for scalar inputs + # so we convert it to None if and only if dim is equivalent + if x.ndim == 0 and dim in {0, -1, (0,), (-1,)}: + kwargs["axis"] = None + else: + kwargs["axis"] = dim + + if "keepdim" in keys and supports_keepdims: + kwargs["keepdims"] = keepdim + + if "mask" in keys: + mask = kwargs.pop("mask") + if mask is not None: + if mask.layout != torch.strided: + raise AssertionError( + f"Expected mask.layout == torch.strided, got {mask.layout}" + ) + kwargs["where"] = mask.cpu().numpy() + + if "identity" in keys: + identity = kwargs.pop("identity") + if identity is not None: + if identity.dtype is torch.bfloat16: + identity = identity.cpu().to(torch.float32) + else: + identity = identity.cpu() + kwargs["initial"] = identity.numpy() + + result = f(x, *args, **kwargs) + + # Unsqueeze reduced dimensions if NumPy does not support keepdims + if keepdim and not supports_keepdims and x.ndim > 0: + dim = list(range(x.ndim)) if dim is None else dim + result = np.expand_dims(result, dim) + + return result + + return wrapper + + +def prod_numpy(a, *args, **kwargs): + """ + The function will call np.prod with type as np.int64 if the input type + is int or uint64 if is uint. This is necessary because windows np.prod uses by default + int32 while on linux it uses int64. + This is for fixing integer overflow https://github.com/pytorch/pytorch/issues/77320 + + Returns: + np.prod of input + """ + if "dtype" not in kwargs: + if np.issubdtype(a.dtype, np.signedinteger): + a = a.astype(np.int64) + elif np.issubdtype(a.dtype, np.unsignedinteger): + a = a.astype(np.uint64) + + fn = reference_reduction_numpy(np.prod) + return fn(a, *args, **kwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e9125ba0ebe7e0623a12ad1a1cd7eeb7d2749a3a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/__init__.py @@ -0,0 +1,7 @@ +# mypy: ignore-errors + +from .make_fx import make_fx_check +from .aot_autograd import aot_autograd_check, _test_aot_autograd_forwards_backwards_helper +from .fake_tensor import fake_check +from .autograd_registration import autograd_registration_check +from .generate_tests import generate_opcheck_tests, opcheck, OpCheckError, dontGenerateOpCheckTests, is_inside_opcheck_mode diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/aot_autograd.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/aot_autograd.py new file mode 100644 index 0000000000000000000000000000000000000000..b7fe32b81a51e345db356e8e39c71b50e068d4db --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/aot_autograd.py @@ -0,0 +1,176 @@ +# mypy: ignore-errors + +import torch +import torch.utils._pytree as pytree +from torch.testing._utils import wrapper_set_seed +from functorch.compile import compiled_function, min_cut_rematerialization_partition, default_partition, nop +from .make_fx import randomize +import re + + +class assert_raises_regex: + def __init__(self, exception_cls, regex): + self.exception_cls = exception_cls + self.regex = regex + + def __enter__(self): + pass + + def __exit__(self, exc_type, exc_val, traceback): + if exc_type == self.exception_cls: + msg = str(exc_val) + if not re.search(self.regex, msg): + raise AssertionError( + f"Expected exception to match regex. regex: {self.regex}, exception: {msg}") + return True # Squashes the exception + if exc_type is not None: + raise AssertionError( + f"Expected {self.exception_cls} to be raised, instead got exception {exc_type}") + raise AssertionError("Expected exception to be raised but none was") + + +def aot_autograd_check( + func, + args, + kwargs, + dynamic, + assert_raises_regex_fn=assert_raises_regex, + assert_equals_fn=torch.testing.assert_close, + check_gradients=True, + try_check_data_specialization=False, + skip_correctness_check=False, + disable_functionalization=False): + """Compares func(*args, **kwargs) in eager-mode to under AOTAutograd. + + Compares outputs and (if check_gradients=True) gradients produced by + AOTAutograd against eager-mode PyTorch. + + We assume that func(*args, **kwargs) succeeds in eager-mode PyTorch. + + """ + flat_args, args_spec = pytree.tree_flatten((args, kwargs)) + args = [arg for arg in flat_args if isinstance(arg, torch.Tensor)] + + # We construct a new function that only accepts Tensors as inputs + def func_no_tensors(args): + reconstructed_flat_args = [] + args = iter(args) + for v in flat_args: + if isinstance(v, torch.Tensor): + reconstructed_flat_args.append(next(args)) + else: + reconstructed_flat_args.append(v) + + c_args, c_kwargs = pytree.tree_unflatten(reconstructed_flat_args, args_spec) + return func(*c_args, **c_kwargs) + + # cannot use the min cut partitioner without functionalization + if disable_functionalization: + compiled_f = compiled_function( + func_no_tensors, + nop, + nop, + dynamic=dynamic, + partition_fn=default_partition, + keep_inference_input_mutations=True, + disable_functionalization=True + ) + else: + compiled_f = compiled_function( + func_no_tensors, + nop, + nop, + dynamic=dynamic, + partition_fn=min_cut_rematerialization_partition, + keep_inference_input_mutations=True, + disable_functionalization=False + ) + + out = wrapper_set_seed(func_no_tensors, args) + if check_gradients == "auto": + any_tensor_requires_grad = pytree.tree_any_only(torch.Tensor, lambda x: x.requires_grad, args) + any_output_requires_grad = pytree.tree_any_only(torch.Tensor, lambda x: x.requires_grad, out) + check_gradients = any_tensor_requires_grad and any_output_requires_grad + if not check_gradients: + compiled_out = wrapper_set_seed(compiled_f, args) + if not skip_correctness_check: + assert_equals_fn(compiled_out, out, msg=outputs_msg) + return + _test_aot_autograd_forwards_backwards_helper( + func_no_tensors, compiled_f, args, assert_raises_regex_fn, assert_equals_fn, + try_check_data_specialization, skip_correctness_check) + +outputs_msg = ( + "Outputs of the operator are different in eager-mode PyTorch vs " + "AOTDispatcher tracing. This means the operator will have incorrect output " + "underneath torch.compile. This could be because the operator's " + "implementation not traceable." +) + + +def _test_aot_autograd_forwards_backwards_helper( + f, compiled_f, args, assert_raises_regex_fn, assert_equals_fn, + try_check_data_specialization, skip_correctness_check=False): + # Verify grads are equal between compiled and non-compiled versions of f. + + def call_forwards_backwards(f, args): + flat_args = pytree.arg_tree_leaves(*args) + diff_args = [arg for arg in flat_args if isinstance(arg, torch.Tensor) and + arg.requires_grad] + out = wrapper_set_seed(f, args) + flat_out = pytree.tree_leaves(out) + + sm = 0 + for i in flat_out: + if isinstance(i, torch.Tensor): + # We need to call .abs() because it is possible that the output of the + # operator is a complex Tensor and autograd will yell at autograd.grad + # on a complex Tensor unless we manually provide the grad_output flag. + sm += i.sum().abs() + if not isinstance(sm, torch.Tensor): + raise AssertionError(f"Expected sm to be a Tensor, got {type(sm)}") + return out, torch.autograd.grad(sm, diff_args, allow_unused=True) + + def check(args, ignore_failure=False): + try: + orig_out, orig_grad = call_forwards_backwards(f, args) + except Exception: + if ignore_failure: + return + raise + + # See https://github.com/pytorch/pytorch/pull/98960#issuecomment-1505962215 + tensor_args = [x for x in pytree.tree_flatten(args)[0] if isinstance(x, torch.Tensor)] + any_non_leaves = any(x.grad_fn is not None for x in tensor_args) + if all(x is None for x in orig_grad) and any_non_leaves: + with assert_raises_regex_fn(RuntimeError, 'does not require grad and does not have a grad_fn'): + call_forwards_backwards(compiled_f, args) + return + + msg = ( + "Gradients of the operator are different in eager-mode PyTorch vs " + "AOTDispatcher. This means the operator will have incorrect gradients " + "underneath torch.compile. This could be because the operator's " + "backward is incorrectly registered or not traceable." + ) + + compiled_out, compiled_grad = call_forwards_backwards(compiled_f, args) + if not skip_correctness_check: + try: + assert_equals_fn(compiled_out, orig_out) + except Exception as e: + raise type(e)(outputs_msg) from e + try: + assert_equals_fn(compiled_grad, orig_grad) + except Exception as e: + raise type(e)(msg) from e + + check(args, ignore_failure=False) + + # Randomize the data and run the traced graph with it, to catch bugs + # where we may have baked in Tensor data into the trace. + # This is not guaranteed to succeed, because `f` might have preconditions + # on the values of the inputs, so we just ignore if this test fails. + if try_check_data_specialization: + args = randomize(args) + check(args, ignore_failure=True) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/autograd_registration.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/autograd_registration.py new file mode 100644 index 0000000000000000000000000000000000000000..0c94f127b4e5e79d8fb03cac70d6d17dd2a261dc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/autograd_registration.py @@ -0,0 +1,135 @@ +# mypy: ignore-errors + +import contextlib + +import torch +import torch.utils._pytree as pytree + + +@contextlib.contextmanager +def set_autograd_fallback_mode(mode): + prev = torch._C._get_autograd_fallback_mode() + try: + torch._C._set_autograd_fallback_mode(mode) + yield + finally: + torch._C._set_autograd_fallback_mode(prev) + + +def autograd_registration_check(op, args, kwargs): + """Check if autograd was registered correctly (for the operator). + + Operators should have "autograd support" registered directly to an + autograd dispatch key. + An incorrect registration may lead to unexpected silent incorrectness. + Note that this check won't catch all problems but will catch + the most common ones. + + Example usage: + >>> x = torch.randn(3, requires_grad=True) + >>> autograd_registration_check(torch.ops.aten.sin.default, (x,), {}) + + Here are some best practices if you do find your autograd is + registered incorrectly: + - If the operator is composite (i.e. consists of other PyTorch ops) + and you wish the operator to decompose and get autograd support + that way, then please register the implementation to + DispatchKey::CompositeImplicitAutograd + - If you're adding an autograd formula for the operator, the correct + thing to do is to register an autograd.Function to + DispatchKey::Autograd (preferred) or one of the + DispatchKey::Autograd keys. It is NOT OK to register + an autograd.Function to a backend (e.g. CPU/CUDA) key. + - If your operator is non-differentiable, then you should register + an implementation to the Autograd key that uses + AutoDispatchBelowAutograd and re-invokes the operator. + + """ + if not isinstance(op, torch._ops.OpOverload): + raise AssertionError(f"Expected op to be OpOverload, got {type(op)}") + # Implementation details + # ----------------------------------------------- + # If an operator doesn't have an autograd kernel at an autograd key, + # and the operator does not return inputs as-is, then all of + # the outputs should have requires_grad=False before we apply + # special behaviors of our default autograd fallback. + # (The default autograd fallback may set requires_grad=True on output + # tensors in certain modes so that when they are backpropped through, + # they raise an error). + # + # Our strategy for detecting if an operator doesn't have an autograd + # kernel at the autograd key is: + # - set the autograd fallback mode to "nothing" (so it does not change + # the required-gradness of outputs) + # - run the operator + # - Check if any outputs of the operator (that are not inputs) require + # grad. This would only happen if the user calls regular PyTorch + # operations in their backend key (this op should instead be + # CompositeImplicitAutograd or not an op) or if the user invokes + # an autograd.Function in the backend key. + # + # Note that it's already likely a bug if the operator directly returns + # an input as output (because custom ops don't have a good way of + # constructing true in-place or out variants), but we defer that + # responsibility to a different test (schema_check). + + flat_args = pytree.arg_tree_leaves(*args, **kwargs) + all_tensors = [arg for arg in flat_args if isinstance(arg, torch.Tensor)] + if not any(t.requires_grad for t in all_tensors): + raise RuntimeError( + "autograd_registration_check: no inputs have requires_grad=True so " + "we are unable to actually perform this test. Please pass inputs " + "that do require grad." + ) + + # Determine which AutogradBACKEND key to check + all_device_types = {arg.device.type for arg in all_tensors} + if not all_device_types.issubset(["cpu", "cuda", "xpu"]): + # Don't want to support other keys yet + raise NotImplementedError( + f"autograd_registration_check: NYI devices other than CPU/CUDA/XPU, got {all_device_types}" + ) + if "cuda" in all_device_types: + key = "AutogradCUDA" + elif "cpu" in all_device_types: + key = "AutogradCPU" + elif "xpu" in all_device_types: + key = "AutogradXPU" + + if torch._C._dispatch_has_kernel_for_dispatch_key(op.name(), key): + return + if torch._C._dispatch_has_kernel_for_dispatch_key(op.name(), "Autograd"): + return + if torch._C._dispatch_has_kernel_for_dispatch_key( + op.name(), "CompositeImplicitAutograd" + ): + return + + # At this point, we know the operator doesn't have a kernel registered to an + # autograd key. Let's proceed with our test. + with set_autograd_fallback_mode("nothing"): + all_outs = op(*args, **kwargs) + + inp_ids = {id(arg) for arg in flat_args} + + def not_an_input_and_requires_grad(tensor): + if not tensor.requires_grad: + return False + if id(tensor) in inp_ids: + return False + return True + + if not pytree.tree_any_only(torch.Tensor, not_an_input_and_requires_grad, all_outs): + return + + raise AssertionError( + f"{op.name()}: at least one output of this operator has requires_grad=True " + f"but the operator does not have an autograd kernel defined at an autograd " + f"key (e.g. DispatchKey::Autograd). This could mean that you have " + f"incorrectly registered an autograd kernel to a non-Autograd DispatchKey, " + f"which may lead to silently incorrect results. If your operator consists " + f"of regular PyTorch operations, consider not using an operator at all " + f"or registering your operator as CompositeImplicitAutograd. If you have " + f"an autograd.Function registered to a backend (CPU/CUDA/XPU) key, the correct " + f"location for it is the Autograd key." + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/fake_tensor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/fake_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..5e60f50189b5dc3ab43fdd97120d5fa23559a84e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/fake_tensor.py @@ -0,0 +1,12 @@ +# mypy: ignore-errors + +import torch._subclasses + + +def is_builtin(op): + return op.namespace in ('aten', 'prims', 'prim') + + +def fake_check(op, args, kwargs): + with torch._subclasses.CrossRefFakeMode(ignore_op_fn=is_builtin): + op(*args, **kwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/generate_tests.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/generate_tests.py new file mode 100644 index 0000000000000000000000000000000000000000..68ace6688bfeb36ff464199f710bec916fcec479 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/generate_tests.py @@ -0,0 +1,866 @@ +# mypy: ignore-errors + +import datetime +import difflib +import functools +import inspect +import json +import os +import re +import tempfile +import threading +import unittest +from collections.abc import Callable, Sequence +from typing import Any + +import torch +import torch._dynamo +import torch.utils._pytree as pytree +from torch._dynamo.utils import clone_input +from torch._library.custom_ops import CustomOpDef +from torch._subclasses.schema_check_mode import SchemaCheckMode +from torch._utils_internal import get_file_path_2 +from torch.overrides import TorchFunctionMode +from torch.testing._internal.optests import ( + aot_autograd_check, + autograd_registration_check, + fake_check, +) + + +def dontGenerateOpCheckTests(reason: str): + def inner(fun): + fun._torch_dont_generate_opcheck_tests = True + return fun + + return inner + + +def is_abstract(tensor: torch.Tensor) -> bool: + if tensor.is_meta: + return True + if torch._subclasses.fake_tensor.is_fake(tensor): + return True + return False + + +def safe_schema_check( + op: torch._ops.OpOverload, + args: tuple[Any, ...], + kwargs: dict[str, Any], + *, + copy_inputs: bool = True, + rtol: float | None = None, + atol: float | None = None, +) -> Any: + if copy_inputs: + args, kwargs = deepcopy_tensors((args, kwargs)) + if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)): + return None + with SchemaCheckMode(): + result = op(*args, **kwargs) + return result + + +def safe_autograd_registration_check( + op: torch._ops.OpOverload, + args: tuple[Any, ...], + kwargs: dict[str, Any], + *, + copy_inputs: bool = True, + rtol: float | None = None, + atol: float | None = None, +) -> None: + if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)): + return + if copy_inputs: + args, kwargs = deepcopy_tensors((args, kwargs)) + # Don't perform autograd_registration_check if none of the inputs require grad. + if not pytree.tree_any_only( + torch.Tensor, lambda x: x.requires_grad, (args, kwargs) + ): + return + return autograd_registration_check(op, args, kwargs) + + +def safe_fake_check( + op: torch._ops.OpOverload, + args: tuple[Any, ...], + kwargs: dict[str, Any], + *, + copy_inputs: bool = True, + rtol: float | None = None, + atol: float | None = None, +) -> None: + if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)): + return None + if copy_inputs: + args, kwargs = deepcopy_tensors((args, kwargs)) + return fake_check(op, args, kwargs) + + +def safe_aot_autograd_check( + op: torch._ops.OpOverload, + args: tuple[Any, ...], + kwargs: dict[str, Any], + dynamic: bool, + *, + copy_inputs: bool = True, + rtol: float | None = None, + atol: float | None = None, +) -> Any: + # NB: copy_inputs does nothing for aot_autograd_check: it always needs to copy + # inputs. + if pytree.tree_any_only(torch.Tensor, is_abstract, (args, kwargs)): + return None + + def func(*args, **kwargs): + args, kwargs = pytree.tree_map_only(torch.Tensor, torch.clone, (args, kwargs)) + return op(*args, **kwargs) + + # aot_autograd_check runs func(*args, **kwargs) multiple times + # and assumes `func` does not modify its inputs. + if rtol and atol: + assert_equals_fn = functools.partial( + torch.testing.assert_close, rtol=rtol, atol=atol + ) + else: + assert_equals_fn = torch.testing.assert_close + return aot_autograd_check( + func, + args, + kwargs, + dynamic, + check_gradients="auto", + assert_equals_fn=assert_equals_fn, + ) + + +def deepcopy_tensors(inputs: Any) -> Any: + return pytree.tree_map_only(torch.Tensor, clone_input, inputs) + + +# Test util requirements +# - The test util must have signature (op: OpOverload, args, kwargs) +# - The test util must NOT mutate args, kwargs. +# - The test utils in this list must not be prefixes of each other. For example, +# having both "test_schema" and "test_schema_is_functional" is NOT OK. +# - The order of items in this dict matters (for opcheck), we'll run them +# in order. +ALL_TEST_UTILS = { + "test_schema": safe_schema_check, + "test_autograd_registration": safe_autograd_registration_check, + "test_faketensor": safe_fake_check, + "test_aot_dispatch_static": functools.partial( + safe_aot_autograd_check, + dynamic=False, + ), + "test_aot_dispatch_dynamic": functools.partial( + safe_aot_autograd_check, + dynamic=True, + ), +} + +GDOC = "https://docs.google.com/document/d/1Pj5HRZvdOq3xpFpbEjUZp2hBovhy7Wnxw14m6lF2154/edit" + +DEFAULT_TEST_UTILS = [ + "test_schema", + "test_autograd_registration", + "test_faketensor", + "test_aot_dispatch_dynamic", +] + +DEPRECATED_DEFAULT_TEST_UTILS = DEFAULT_TEST_UTILS + [ + "test_aot_dispatch_static", +] + + +def generate_opcheck_tests( + testcase: Any, + namespaces: list[str], + failures_dict_path: str | None = None, + additional_decorators: dict[str, Callable] | None = None, + test_utils: list[str] = DEFAULT_TEST_UTILS, +) -> None: + """Given an existing TestCase, use the existing tests to generate + additional validation tests for custom operators. + + For {all existing tests in the TestCase} x {all test utils}, + we will generate one new test. The new test runs a TorchFunctionMode + that intercepts ``op(*args, **kwargs)`` calls and invokes + ``test_util(op, *args, **kwargs)``, where ``op`` is an operator. + + The test_util that we support are in ALL_TEST_UTILS. They are: + - test_schema: This runs SchemaCheckMode. + - test_autograd_registration: This runs autograd_registration_check. + - test_faketensor: This runs CrossRefFakeMode. + - test_aot_dispatch_static: This runs aot_autograd_check, which: + checks that the outputs (and gradients, if they are computable) + are the same under eager-mode PyTorch and using AOTAutograd. + - test_aot_dispatch_dynamic: Same as aot_dispatch_static, but + runs AOTAutograd using dynamic shapes instead of static shapes. + + The generated test will have name ``{test_util}__{original_name}``. + For example, if there is a method named ``test_cumsum``, then + we will generate a ``test_schema__test_cumsum``, + ``test_faketensor__test_cumsum``, etc. + + For more details, see https://docs.google.com/document/d/1Pj5HRZvdOq3xpFpbEjUZp2hBovhy7Wnxw14m6lF2154/edit + + Args: + testcase: The testcase we will modify and generate additional tests for. + namespaces: We will only intercept calls to custom operators with these + namespaces. + failures_dict_path: See ``validate_failures_dict_structure`` for more details + test_utils: a list of test_utils to generate. Example: ["test_schema", "test_faketensor"] + """ + if additional_decorators is None: + additional_decorators = {} + test_methods = [ + m + for m in dir(testcase) + if m.startswith("test_") and callable(getattr(testcase, m)) + ] + if failures_dict_path is None: + # The default failures_dict_path is failures_dict.json in + # the same directory as the test file. + prev_frame = inspect.currentframe().f_back + filename = inspect.getframeinfo(prev_frame)[0] + failures_dict_path = get_file_path_2( + os.path.dirname(filename), "failures_dict.json" + ) + failures_dict = FailuresDict.load( + failures_dict_path, create_file=should_update_failures_dict() + ) + validate_failures_dict_structure(failures_dict, test_utils, testcase) + validate_failures_dict_formatting(failures_dict_path) + + def construct_method(attr, prefix, tester): + method = getattr(testcase, attr) + if getattr(method, "_torch_dont_generate_opcheck_tests", False): + return + new_method_name = prefix + "__" + attr + + @functools.wraps(method) + def new_method(*args, **kwargs): + with OpCheckMode( + namespaces, + prefix, + tester, + failures_dict, + f"{testcase.__name__}.{new_method_name}", + failures_dict_path, + ): + result = method(*args, **kwargs) + return result + + if pytestmark := new_method.__dict__.get("pytestmark"): + import pytest + + # check if we need to simplify the parametrize marks + # NB: you need to add this mark to your pytest.ini + opcheck_only_one = False + for mark in pytestmark: + if isinstance(mark, pytest.Mark) and mark.name == "opcheck_only_one": + opcheck_only_one = True + + if opcheck_only_one: + new_pytestmark = [] + for mark in pytestmark: + if isinstance(mark, pytest.Mark) and mark.name == "parametrize": + argnames, argvalues = mark.args + if mark.kwargs: + raise AssertionError("NYI: mark.kwargs is not empty") + # Special case for device, we want to run on all + # devices + if argnames != "device": + new_pytestmark.append( + pytest.mark.parametrize( + argnames, (next(iter(argvalues)),) + ) + ) + continue + new_pytestmark.append(mark) + new_method.__dict__["pytestmark"] = new_pytestmark + + if new_method_name in additional_decorators: + for dec in additional_decorators[new_method_name]: + new_method = dec(new_method) + + if hasattr(testcase, new_method_name): + raise RuntimeError( + f"Tried to autogenerate {new_method_name} but {testcase} already " + f"has method named {new_method_name}. Please rename the original " + f"method on the TestCase." + ) + setattr(testcase, new_method_name, new_method) + + test_utils = {name: ALL_TEST_UTILS[name] for name in test_utils} + for attr in test_methods: + for prefix, tester in test_utils.items(): + construct_method(attr, prefix, tester) + + generate_tag_tests(testcase, failures_dict, additional_decorators) + + +def generate_tag_tests(testcase, failures_dict, additional_decorators): + def generate_test(qualname, definitely_not_pt2_compliant, xfailed_tests): + def inner(self): + try: + op = torch._library.utils.lookup_op(qualname) + except AttributeError as e: + # Operator not importable in this test file + raise unittest.SkipTest(f"Can't import operator {qualname}") from e + op_marked_as_compliant = torch.Tag.pt2_compliant_tag in op.tags + if not op_marked_as_compliant: + return + if not definitely_not_pt2_compliant: + return + raise AssertionError( + f"op '{qualname}' was tagged with torch.Tag.pt2_compliant_tag " + f"but it failed some of the generated opcheck tests " + f"({xfailed_tests}). This may lead to silent correctness issues, " + f"please fix this." + ) + + return inner + + for qualname, test_dict in failures_dict.data.items(): + xfailed_tests = [ + test + for test, status_dict in test_dict.items() + # We're about to delete the following test after Ed's PR + # to specialize on C++ .size() calls + if "test_aot_dispatch_static" not in test + and status_dict["status"] == "xfail" + ] + definitely_not_pt2_compliant = len(xfailed_tests) > 0 + generated = generate_test(qualname, definitely_not_pt2_compliant, xfailed_tests) + + # Could result in collisions, but unlikely. We'll raise if we see one below. + mangled_qualname = qualname.replace("::", "_").replace(".", "_") + test_name = "test_pt2_compliant_tag_" + mangled_qualname + + # You can skip this test via the additional_decorators argument + # in generate_opcheck_tests + if test_name in additional_decorators: + for decorator in additional_decorators[test_name]: + generated = decorator(generated) + + if hasattr(testcase, test_name): + raise RuntimeError( + f"Tried to generate a test named {test_name}, but it exists " + f"already. This could be because of a name collision (where " + f"we generated two tests with the same name), or where we " + f"generated a test with the same name as an existing test." + ) + setattr(testcase, test_name, generated) + + +TEST_OPTIONS = ("xfail", "skip", "xsuccess") + + +def validate_failures_dict_formatting(failures_dict_path: str) -> None: + with open(failures_dict_path) as fp: + actual = fp.read() + failures_dict = FailuresDict.load(failures_dict_path) + expected = failures_dict._save(to_str=True) + if actual == expected: + return + if should_update_failures_dict(): + failures_dict = FailuresDict.load(failures_dict_path) + failures_dict.save() + return + expected = expected.splitlines(1) + actual = actual.splitlines(1) + diff = difflib.unified_diff(actual, expected) + diff = "".join(diff) + raise RuntimeError( + f"\n{diff}\n\nExpected the failures dict to be formatted " + f"a certain way. Please see the above diff; you can correct " + f"this either manually or by re-running the test with " + f"PYTORCH_OPCHECK_ACCEPT=1" + ) + + +def validate_failures_dict_structure( + failure_dict: "FailuresDict", test_utils: list[str], testcase: Any +) -> None: + """Validates the failures dict. + + The failure dict looks something like the following. + It maps operator name (qualname) to a list of autogenerated tests. + Each autogenerated test may have a check for the operator (if the operator is + called by the test); the dictionary specifies if we should skip the check, + or if we expect some check to fail. + + { + "fbgemm::split_lengths": { + "test_schema__test_split_lengths": { + "comment": "you can put whatever you want into the comment section", + "status": "xfail", + } + "test_schema__test_split_lengths_empty": { + "comment": "", + "status": "skip", + }, + }, + "fbgemm::gather_lengths": { + "test_schema__test_gather_lengths": { + "comment": "", + "status": "skip", + }, + }, + } + + """ + failure_dict = failure_dict.data + for test_to_option in failure_dict.values(): + for test_name, test_dict in test_to_option.items(): + if set(test_dict.keys()) != set({"comment", "status"}): + raise RuntimeError( + "in failures_dict, expected sub-dict to have keys 'comment' and 'status'" + ) + test_option = test_dict["status"] + if test_option not in TEST_OPTIONS: + raise RuntimeError( + f"In failures_dict, got status={test_option} but it needs to be in {TEST_OPTIONS}" + ) + test_class, actual_test_name = test_name.split(".") + if not any(actual_test_name.startswith(test) for test in test_utils): + raise RuntimeError( + f"In failures_dict, test name '{test_name}' should begin with one of {test_utils}" + ) + for test in test_utils: + if not actual_test_name.startswith(test): + continue + base_test_name = actual_test_name[len(test) + 2 :] + # remove potential pytest parametrization suffix + base_test_name = re.sub(r"\[.*\]", "", base_test_name) + if testcase.__name__ != test_class: + continue + if hasattr(testcase, base_test_name): + continue + raise RuntimeError( + f"In failures dict, got test name '{test_name}'. We parsed this as " + f"running test '{test}' on '{base_test_name}', but " + f"{base_test_name} does not exist on the TestCase '{testcase.__name__}]. " + f"Maybe you need to change the test name?" + ) + + +def should_update_failures_dict() -> bool: + key = "PYTORCH_OPCHECK_ACCEPT" + return key in os.environ and os.environ[key] == "1" + + +_is_inside_opcheck_mode = threading.local() +_is_inside_opcheck_mode.value = False + + +def is_inside_opcheck_mode(): + return _is_inside_opcheck_mode.value + + +class OpCheckMode(TorchFunctionMode): + """ + For a given test, OpCheckMode intercepts calls to operators and runs + test_util(op, args, kwargs) for each intercepted (op, args, kwargs). + """ + + def __init__( + self, + namespaces: list[str], + test_util_name: str, + test_util: Callable, + failures_dict: "FailuresDict", + test_name: str, + failures_dict_path: str, + ): + # We will intercept calls to ops with these namespaces + self.namespaces = namespaces + # The test utility function. Its signature should be (op, args, kwargs) -> None. + # Examples of test utilities are: schema_check, make_fx_check + self.test_util = test_util + self.test_util_name = test_util_name + # The name of the test that is running this OpCheckMode. + self.test_name = test_name + # Maps qualname -> test_name -> skip/xfail + # Tells us if we should skip a test or assert that there is a failure. + self.failures_dict = failures_dict + # Location of the failures dict. Makes it so that the error message is better. + self.failures_dict_path = failures_dict_path + + # OpCheckMode suppresses errors, collects them here, and then raises them on exit. + # Maps qualname -> List[(Exception, func, maybe args, maybe kwargs)] + self.seen_ops_to_errors = {} + + def maybe_raise_errors_on_exit(self) -> None: + # Check expected failures first + for qualname in self.seen_ops_to_errors: + option = self.failures_dict.get_status(qualname, self.test_name) + if len(self.seen_ops_to_errors[qualname]) == 0: + if should_update_failures_dict(): + self.failures_dict.set_status( + qualname, self.test_name, "xsuccess", comment="" + ) + else: + if option == "xfail": + raise OpCheckError( + f"generate_opcheck_tests: Unexpected success for operator " + f"{qualname} on test {self.test_name}. This may mean that " + f"you have fixed this test failure. Please rerun the test with " + f"PYTORCH_OPCHECK_ACCEPT=1 to automatically update the test runner " + f"or manually remove the " + f"expected failure in the failure dict at " + f"{self.failures_dict_path}" + f"For more details, see " + f"{GDOC}" + ) + continue + failed_ops = [] + for qualname in self.seen_ops_to_errors: + option = self.failures_dict.get_status(qualname, self.test_name) + if option != "xsuccess": + continue + if len(self.seen_ops_to_errors[qualname]) == 0: + continue + failed_ops.append(qualname) + if not failed_ops: + return + + if should_update_failures_dict(): + for op in failed_ops: + self.failures_dict.set_status(op, self.test_name, "xfail") + return + + # Raise from the first error but also report about all of them to make + # recording xfails easier. + ex, op, args, kwargs = self.seen_ops_to_errors[failed_ops[0]][0] + repro_command = generate_repro( + self.test_util_name, op, args, kwargs, save_data=should_print_better_repro() + ) + raise OpCheckError( + f"Test generated by `generate_opcheck_tests`, {self.test_name}, " + f"failed on operators {failed_ops}. This usually means that the " + f"operators are not implemented correctly and may lead to silently " + f"incorrect behavior. Set PYTORCH_OPCHECK_PRINT_BETTER_REPRO=1 for a standalone repro, " + f"or please see " + f"{GDOC} " + f"for more recommendations. " + f"To reproduce this problem locally, try to run the following:\n{repro_command}" + ) from ex + + def __enter__(self, *args, **kwargs): + self.prev_is_opcheck_mode = _is_inside_opcheck_mode.value + self.prev_dynamo_disable = os.environ.get("TORCHDYNAMO_DISABLE", "") + _is_inside_opcheck_mode.value = True + os.environ["TORCHDYNAMO_DISABLE"] = "1" + # When running this test mode, we want to disable + # default torch.compile custom op checker + self.prev_functorch_config_for_checking_custom_op = ( + torch._functorch.config.check_custom_op_aliasing + ) + torch._functorch.config.check_custom_op_aliasing = False + return super().__enter__(*args, **kwargs) + + def __exit__(self, *args, **kwargs): + _is_inside_opcheck_mode.value = self.prev_is_opcheck_mode + os.environ["TORCHDYNAMO_DISABLE"] = self.prev_dynamo_disable + torch._functorch.config.check_custom_op_aliasing = ( + self.prev_functorch_config_for_checking_custom_op + ) + try: + self.maybe_raise_errors_on_exit() + if should_update_failures_dict(): + self.failures_dict.save() + finally: + result = super().__exit__(*args, **kwargs) + return result + + def run_test_util(self, op, args, kwargs): + try: + self.test_util(op, args, kwargs, copy_inputs=False) + except torch._subclasses.fake_tensor.UnsupportedFakeTensorException: + # We might get here if the input is already a FakeTensor + # or if we're in a torch.compile block. Just ignore these + # since we can't handle them and reporting them as failures + # is too noisy. + pass + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs if kwargs else {} + + # Only intercept calls to operators + if not isinstance(func, (torch._ops.OpOverloadPacket, torch._ops.OpOverload)): + return func(*args, **kwargs) + if ( + torch.jit.is_tracing() + or torch.jit.is_scripting() + or torch._dynamo.is_compiling() + ): + return func(*args, **kwargs) + # Pre-existing code may not use the .default overload. If we see an + # OpOverloadPacket and we cannot resolve the overload, then we just throw + # and ask the user to clarify. Otherwise, we attempt to resolve the overload. + if isinstance(func, torch._ops.OpOverloadPacket): + func = resolve_unique_overload_or_throw(func) + qualname = func.name() + ns = qualname.split("::")[0] + if ns not in self.namespaces: + return func(*args, **kwargs) + + args_c, kwargs_c = deepcopy_tensors((args, kwargs)) + result = func(*args, **kwargs) + + option = self.failures_dict.get_status(qualname, self.test_name) + if option == "xsuccess" or option == "xfail": + # Suppress all errors during execution. Raise them during __exit__. + try: + if qualname not in self.seen_ops_to_errors: + self.seen_ops_to_errors[qualname] = [] + self.run_test_util(func, args_c, kwargs_c) + except Exception as ex: + if should_print_better_repro(): + self.seen_ops_to_errors[qualname].append((ex, func, args, kwargs)) + else: + self.seen_ops_to_errors[qualname].append((ex, func, None, None)) + elif option == "skip": + pass + return result + + +def should_print_better_repro() -> None: + """If set, the tests generated by `generate_opcheck_tests` will print a + repro command on failure. + + In order to print the repro command, we need to save some tensors to disk. + These will be saved under the following directory: + {tempfile.gettempdir()}/pytorch_opcheck_safe_to_delete/. + + Although this is a temp folder, it will usually not automatically get cleaned + up, so you'll need to manually delete it. + """ + key = "PYTORCH_OPCHECK_PRINT_BETTER_REPRO" + if key not in os.environ: + return False + value = os.environ[key] + return value == "1" or value == 1 + + +def opcheck( + op: torch._ops.OpOverload | torch._ops.OpOverloadPacket | CustomOpDef, + args: tuple[Any, ...], + kwargs: dict[str, Any] | None = None, + *, + test_utils: str | Sequence[str] = DEFAULT_TEST_UTILS, + raise_exception: bool = True, + rtol: float | None = None, + atol: float | None = None, +) -> dict[str, str]: + """See torch.library.opcheck for docstring""" + + if (rtol is None) ^ (atol is None): + raise ValueError( + "opcheck(op, ...): if you specify one of rtol/atol, you must specify both" + ) + + if kwargs is None: + kwargs = {} + if isinstance(op, CustomOpDef): + op = op._opoverload + if isinstance(op, torch._ops.OpOverloadPacket): + op = resolve_unique_overload_or_throw(op) + if not isinstance(op, torch._ops.OpOverload): + raise ValueError( + f"opcheck(op, ...): op must be instance of torch._ops.OpOverload, " + f"e.g. torch.ops.aten.sin.default, got {type(op)}" + ) + if test_utils == "ALL": + test_utils = tuple(ALL_TEST_UTILS.keys()) + if isinstance(test_utils, str): + test_utils = (test_utils,) + if not isinstance(test_utils, (tuple, list)) or not set(test_utils).issubset( + ALL_TEST_UTILS.keys() + ): + raise ValueError( + f"opcheck(op, ..., test_utils={test_utils}), expected test_utils " + f"to be subset of {tuple(ALL_TEST_UTILS.keys())} but it was not" + ) + + results_dict = {} + for test_util in test_utils: + tester = ALL_TEST_UTILS[test_util] + try: + tester(op, args, kwargs, rtol=rtol, atol=atol) + results_dict[test_util] = "SUCCESS" + except Exception as ex: + if raise_exception: + raise OpCheckError( + f"opcheck(op, ...): {test_util} failed with {ex} " + f"(scroll up for stack trace)" + ) from ex + results_dict[test_util] = ex + return results_dict + + +class OpCheckError(Exception): + pass + + +def generate_repro( + test: str, + op: torch._ops.OpOverload, + args: tuple[Any, ...], + kwargs: dict[str, Any], + *, + save_data: bool, + dry_run: bool = False, +) -> str: + if save_data: + now = datetime.datetime.now() + path = os.path.join(tempfile.gettempdir(), "pytorch_opcheck_safe_to_delete") + unix_timestamp = datetime.datetime.timestamp(now) * 100000 + filepath = os.path.join(path, f"repro_{unix_timestamp}.pt") + if not dry_run: + os.makedirs(path, exist_ok=True) + torch.save((args, kwargs), filepath) + args_kwargs = f'args, kwargs = torch.load("{filepath}")' + else: + args_kwargs = ( + "# If you rerun your test with PYTORCH_OPCHECK_PRINT_BETTER_REPRO=1\n" + "# we will fill them in same (args, kwargs) as in your test\n" + "args = () # args to the operator\n" + "kwargs = {} # kwargs to the operator" + ) + + ns, name = op._schema.name.split("::") + overload = op._overloadname + + repro_command = ( + f"# =========================================================\n" + f"# BEGIN REPRO SCRIPT\n" + f"# =========================================================\n" + f"import torch\n" + f"from torch.testing._internal.optests import opcheck\n" + f"\n" + f"# Make sure you have loaded the library that contains the op\n" + f"# via an import or torch.ops.load_library(...)\n" + f"op = torch.ops.{ns}.{name}.{overload}\n" + f"\n" + f"{args_kwargs}\n" + f'opcheck(op, args, kwargs, test_utils="{test}")\n' + f"# =========================================================\n" + f"# END REPRO SCRIPT\n" + f"# =========================================================\n" + ) + return repro_command + + +def resolve_unique_overload_or_throw( + op: torch._ops.OpOverloadPacket, +) -> torch._ops.OpOverload: + all_schemas = torch._C._jit_get_schemas_for_operator(op._qualified_op_name) + if len(all_schemas) != 1: + raise RuntimeError( + f"opcheck can only test operators without overloads. " + f"Got the following overloads for {op._qualified_op_name}: " + f"{[schema.overload_name for schema in all_schemas]}" + ) + + overload_name = all_schemas[0].overload_name + if overload_name == "": + return op.default + return getattr(op, overload_name) + + +DUMP_OPTIONS = {"indent": 2, "sort_keys": True} + + +FailuresDictData = dict[str, dict[str, dict[str, str]]] + + +VERSION = 1 +DESCRIPTION = ( + f"This is a dict containing failures for tests autogenerated by " + f"generate_opcheck_tests. " + f"For more details, please see {GDOC}" +) + + +class FailuresDict: + def __init__(self, path: str, data: FailuresDictData): + self.path = path + self.data = data + + @staticmethod + def load(path, *, create_file=False) -> "FailuresDict": + if create_file and not os.path.exists(path): + result = FailuresDict(path, {}) + FailuresDict.save() + return result + with open(path) as fp: + contents = fp.read() + if contents.strip() == "": + dct = { + "_description": DESCRIPTION, + "data": {}, + "_version": VERSION, + } + else: + dct = json.loads(contents) + if "data" not in dct: + raise AssertionError("Expected 'data' in dct") + if "_version" not in dct or dct["_version"] != VERSION: + raise AssertionError( + f"Expected '_version' in dct with value {VERSION}" + ) + return FailuresDict(path, dct["data"]) + + def _save(self, to_str=False) -> str | None: + to_dump = { + "_description": DESCRIPTION, + "data": self.data, + "_version": VERSION, + } + # json.dumps doesn't end with a newline. Let's add one because files + # should end in newlines. + serialized = json.dumps(to_dump, **DUMP_OPTIONS) + "\n" + if to_str: + return serialized + with open(self.path, "w") as fp: + fp.write(serialized) + return None + + def save(self) -> None: + return self._save() + + def get_status(self, qualname: str, test_name: str) -> str: + if qualname not in self.data: + return "xsuccess" + dct = self.data[qualname] + if test_name not in dct: + return "xsuccess" + return dct[test_name]["status"] + + def set_status( + self, + qualname: str, + test_name: str, + status: str, + *, + comment: str | None = None, + ): + if qualname not in self.data: + self.data[qualname] = {} + dct = self.data[qualname] + if test_name not in dct: + dct[test_name] = {"status": None, "comment": ""} + + if status == "xsuccess": + # The default status is "xsuccess". + del dct[test_name] + else: + dct[test_name]["status"] = status + if comment is not None: + dct[test_name]["comment"] = comment diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/make_fx.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/make_fx.py new file mode 100644 index 0000000000000000000000000000000000000000..970a0be1b36956d3693a5a93d07dbf32027c9773 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/optests/make_fx.py @@ -0,0 +1,89 @@ +# mypy: ignore-errors + +import torch +from torch.fx.experimental.proxy_tensor import make_fx +from torch.testing._utils import wrapper_set_seed +import torch.utils._pytree as pytree + + +def make_fx_check( + func, + args, + kwargs, + tracing_mode, + assert_close=torch.testing.assert_close, + randomize_data=False, +): + f, *new_args = handle_sizes_for_dynamic_shapes(func, args, kwargs) + + def run(f, *args, **kwargs): + return wrapper_set_seed(f, *args, **kwargs) + + traced_f = make_fx(f, tracing_mode=tracing_mode)(*new_args) + + msg = ( + "op(*args, **kwargs) and make_fx(op)(*args, **kwargs) produced different " + "values. This could mean that your abstract impls (meta/FakeTensor impls) " + "are incorrect, that your operator is not completely traceable (e.g., " + "it relies on some global state), or that there is a bug in make_fx. " + "Note that if you passed a python function (and not an operator) to " + "make_fx_check, it is still possible that the python function will still " + "work with torch.compile because it handles capturing pieces of " + "your python code to compile." + ) + + # Randomize the data and run the traced graph with it, to catch bugs + # where we may have baked in Tensor data into the trace. + # This is not guaranteed to succeed, because `f` might have preconditions + # on the values of the inputs, so we just ignore if we used + # random data and it fails. + if randomize_data: + new_args = randomize(new_args) + try: + expected = run(f, *new_args) + except Exception: + if randomize_data: + return + raise + result = run(traced_f, *new_args) + assert_close(result, expected, msg=msg) + + +# Arguably we should make make_fx promote torch.Size() objects to symbolic shapes. +# Absent that, here is our strategy: +# +# If any argument is a torch.Size(), maybe get dynamic shapes for it by: +# - Create a temporary Tensor whose size is the torch.Size() we want. Note that +# we use an expanded Tensor as we cannot pass "meta" Tensors to make_fx. +# - Pass it to make_fx such that it is converted to a proxy Tensor +# - Unpack the size in the wrapper to get a torch.Size with dynamic shapes (in +# symbolic mode, a no-op otherwise) +def handle_sizes_for_dynamic_shapes(func, args, kwargs): + def f(args, kwargs, extra_args, extra_kwargs): + if extra_args: + for i, t in extra_args: + args[i] = t.size() + if extra_kwargs: + for k, t in extra_kwargs.items(): + kwargs[k] = t.size() + + return func(*args, **kwargs) + + extra_args = [] + extra_kwargs = {} + for i, arg in enumerate(args): + if isinstance(arg, torch.Size): + extra_args.append((i, torch.empty(arg, device="cpu"))) + for key, value in kwargs.items(): + if isinstance(value, torch.Size): + extra_kwargs[key] = torch.empty(value, device="cpu") + + return f, args, kwargs, extra_args, extra_kwargs + + +def randomize(args): + def transform(x): + if not x.dtype.is_floating_point: + return x + return x.detach().clone().uniform_(0, 1).requires_grad_(x.requires_grad) + return pytree.tree_map_only(torch.Tensor, transform, args) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/py312_intrinsics.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/py312_intrinsics.py new file mode 100644 index 0000000000000000000000000000000000000000..ae6dc3ddd72446bb0d778b88b0af12dd5a02741f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/py312_intrinsics.py @@ -0,0 +1,29 @@ +import torch + + +# This is a very odd way to define a test case that uses CALL_INTRINSIC_1 4 and +# CALL_INTRINSIC_2 7. The problem is those intrinsics only exists on Python +# 3.12+ and the Python code to produce them is not backwards compatible with +# previous versions. +class Foo: + @classmethod + def _default_update(cls): + def f[T](a: "This is a new annotation"): # noqa: F722 + """This is a test""" + + f.attr = "This is also a test" + f.__wrapped__ = "This is a bald faced lie" + + def wrapper(b: "This is the prior annotation"): # noqa: F722 + pass + + return wrapper, f + + def test_default_update(self): + @torch.compile(backend="eager", fullgraph=True) + def fn(x): + wrapper, f = Foo._default_update() + return x + 1 + + x = torch.randn(2) + fn(x) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/quantization_torch_package_models.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/quantization_torch_package_models.py new file mode 100644 index 0000000000000000000000000000000000000000..abc4ab6f7e4734361ec7ecea3d4755910f9cf2ab --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/quantization_torch_package_models.py @@ -0,0 +1,33 @@ +# mypy: ignore-errors + +import math + +import torch +import torch.nn as nn + + +class LinearReluFunctionalChild(nn.Module): + def __init__(self, N): + super().__init__() + self.w1 = nn.Parameter(torch.empty(N, N)) + self.b1 = nn.Parameter(torch.zeros(N)) + torch.nn.init.kaiming_uniform_(self.w1, a=math.sqrt(5)) + + def forward(self, x): + x = torch.nn.functional.linear(x, self.w1, self.b1) + x = torch.nn.functional.relu(x) + return x + +class LinearReluFunctional(nn.Module): + def __init__(self, N): + super().__init__() + self.child = LinearReluFunctionalChild(N) + self.w1 = nn.Parameter(torch.empty(N, N)) + self.b1 = nn.Parameter(torch.zeros(N)) + torch.nn.init.kaiming_uniform_(self.w1, a=math.sqrt(5)) + + def forward(self, x): + x = self.child(x) + x = torch.nn.functional.linear(x, self.w1, self.b1) + x = torch.nn.functional.relu(x) + return x diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/static_module.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/static_module.py new file mode 100644 index 0000000000000000000000000000000000000000..0a031b0d8f6e685517b7ac51c236e23835501cd9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/static_module.py @@ -0,0 +1,27 @@ +# mypy: allow-untyped-defs +# Owner(s): ["module: unknown"] + +import torch + + +class StaticModule: + def __init__(self, scripted): + # this is an nn.Module + if hasattr(scripted, "_c"): + self.static_module = torch._C._jit_to_static_module(scripted._c) + else: + self.static_module = torch._C._jit_to_static_module(scripted.graph) + + def __call__(self, *args, **kwargs): + return self.static_module(*args, **kwargs) + + def benchmark(self, args, kwargs, warmup_runs, main_runs): + self.static_module.benchmark(args, kwargs, warmup_runs, main_runs) + + def runAsync(self, args, kwargs): + return self.static_module.runAsync(args, kwargs) + + def benchmark_individual_ops(self, args, kwargs, warmup_runs, main_runs): + return self.static_module.benchmark_individual_ops( + args, kwargs, warmup_runs, main_runs + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/subclasses.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/subclasses.py new file mode 100644 index 0000000000000000000000000000000000000000..03326233a8fbbfaad2568c0ca8dbfbc764ef3106 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/subclasses.py @@ -0,0 +1,81 @@ +# mypy: ignore-errors +from typing import Any + +import torch +import torch.utils._pytree as pytree +from torch._subclasses.fake_tensor import is_fake +from torch.testing._internal.two_tensor import TwoTensor +from torch.utils._python_dispatch import return_and_correct_aliasing + + +class WrapperSubclass(torch.Tensor): + @staticmethod + def __new__(cls, a, outer_size=None, outer_stride=None): + if outer_size is None: + outer_size = a.size() + if outer_stride is None: + outer_stride = a.stride() + + kwargs = {} + kwargs["strides"] = outer_stride + kwargs["storage_offset"] = a.storage_offset() + kwargs["device"] = a.device + kwargs["layout"] = a.layout + kwargs["requires_grad"] = a.requires_grad + kwargs["dtype"] = a.dtype + out = torch.Tensor._make_wrapper_subclass(cls, outer_size, **kwargs) + + return out + + def __init__(self, a, outer_size=None, outer_stride=None): + self.a = a + + def __repr__(self): + return f"WrapperSubclass({repr(self.a)})" + + def __tensor_flatten__(self): + return ["a"], None + + @staticmethod + def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride): + if meta is not None: + raise AssertionError("Expected meta to be None") + a = inner_tensors["a"] + if is_fake(a): + if outer_size is None: + raise AssertionError("Expected outer_size to not be None") + if outer_stride is None: + raise AssertionError("Expected outer_stride to not be None") + return WrapperSubclass(a, outer_size, outer_stride) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + if kwargs is None: + kwargs = {} + args_a = pytree.tree_map_only(WrapperSubclass, lambda x: x.a, args) + + kwargs_a = pytree.tree_map_only(WrapperSubclass, lambda x: x.a, kwargs) + + out_a = func(*args_a, **kwargs_a) + out_a_flat, spec = pytree.tree_flatten(out_a) + out_flat = [ + WrapperSubclass(o_a) if isinstance(o_a, torch.Tensor) else o_a + for o_a in out_a_flat + ] + out = pytree.tree_unflatten(out_flat, spec) + from torch._higher_order_ops.cond import cond_op + + if func is cond_op: + return out + else: + return return_and_correct_aliasing(func, args, kwargs, out) + + def __coerce_same_metadata_as_tangent__( + self, expected_metadata: Any, expected_type: type | None = None + ): + if expected_type is type(self.a): + return self.a + elif expected_type is TwoTensor: + return TwoTensor(self.a, self.a.clone()) + + return None diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/test_module/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/test_module/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/test_module/future_div.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/test_module/future_div.py new file mode 100644 index 0000000000000000000000000000000000000000..0a3494f945fad36d84cb8056dcf722d6911f0af2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/test_module/future_div.py @@ -0,0 +1,10 @@ +# mypy: ignore-errors + + + +def div_int_future(): + return 1 / 2 + + +def div_float_future(): + return 3.14 / 0.125 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/test_module/no_future_div.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/test_module/no_future_div.py new file mode 100644 index 0000000000000000000000000000000000000000..164e6d168414a11039f3b63885760ad08b81ae99 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/test_module/no_future_div.py @@ -0,0 +1,11 @@ +# mypy: ignore-errors + +import torch # noqa: F401 + + +def div_int_nofuture(): + return 1 / 2 + + +def div_float_nofuture(): + return 3.14 / 0.125 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/torchbind_impls.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/torchbind_impls.py new file mode 100644 index 0000000000000000000000000000000000000000..02ec2623271c94af9b6910751eec66032f84a146 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/torchbind_impls.py @@ -0,0 +1,193 @@ +# mypy: allow-untyped-defs +import contextlib +from pathlib import Path + +import torch + + +_TORCHBIND_IMPLS_INITIALIZED = False + +_TENSOR_QUEUE_GLOBAL_TEST: torch.ScriptObject | None = None + + +def init_torchbind_implementations(): + global _TORCHBIND_IMPLS_INITIALIZED + global _TENSOR_QUEUE_GLOBAL_TEST + if _TORCHBIND_IMPLS_INITIALIZED: + return + + load_torchbind_test_lib() + register_fake_operators() + register_fake_classes() + _TENSOR_QUEUE_GLOBAL_TEST = _empty_tensor_queue() + _TORCHBIND_IMPLS_INITIALIZED = True + + +def _empty_tensor_queue() -> torch.ScriptObject: + return torch.classes._TorchScriptTesting._TensorQueue( + torch.empty( + 0, + ).fill_(-1) + ) + + +# put these under a function because the corresponding library might not be loaded yet. +def register_fake_operators(): + @torch.library.register_fake("_TorchScriptTesting::takes_foo_python_meta") + def fake_takes_foo(foo, z): + return foo.add_tensor(z) + + @torch.library.register_fake("_TorchScriptTesting::queue_pop") + def fake_queue_pop(tq): + return tq.pop() + + @torch.library.register_fake("_TorchScriptTesting::queue_push") + def fake_queue_push(tq, x): + return tq.push(x) + + torch.library.register_autocast( + "_TorchScriptTesting::queue_push", "cpu", torch.float32 + ) + torch.library.register_autocast( + "_TorchScriptTesting::queue_push", "cuda", torch.float32 + ) + + torch.library.register_autocast( + "_TorchScriptTesting::queue_pop", "cpu", torch.float32 + ) + torch.library.register_autocast( + "_TorchScriptTesting::queue_pop", "cuda", torch.float32 + ) + + @torch.library.register_fake("_TorchScriptTesting::queue_size") + def fake_queue_size(tq): + return tq.size() + + def meta_takes_foo_list_return(foo, x): + a = foo.add_tensor(x) + b = foo.add_tensor(a) + c = foo.add_tensor(b) + return [a, b, c] + + def meta_takes_foo_tuple_return(foo, x): + a = foo.add_tensor(x) + b = foo.add_tensor(a) + return (a, b) + + @torch.library.register_fake("_TorchScriptTesting::takes_foo_tensor_return") + def meta_takes_foo_tensor_return(foo, x): + # This implementation deliberately creates unbacked symint for testing + ctx = torch.library.get_ctx() + fake_shape = [ctx.new_dynamic_size() for _ in range(2)] + return torch.empty(fake_shape, dtype=torch.int, device="cpu") + + torch.ops._TorchScriptTesting.takes_foo_list_return.default.py_impl( + torch._C.DispatchKey.Meta + )(meta_takes_foo_list_return) + + torch.ops._TorchScriptTesting.takes_foo_tuple_return.default.py_impl( + torch._C.DispatchKey.Meta + )(meta_takes_foo_tuple_return) + + torch.ops._TorchScriptTesting.takes_foo.default.py_impl(torch._C.DispatchKey.Meta)( + # make signature match original cpp implementation to support kwargs + lambda foo, x: foo.add_tensor(x) + ) + + +def register_fake_classes(): + # noqa: F841 + @torch._library.register_fake_class("_TorchScriptTesting::_Foo") + class FakeFoo: + def __init__(self, x: int, y: int): + self.x = x + self.y = y + + @classmethod + def __obj_unflatten__(cls, flattend_foo): + return cls(**dict(flattend_foo)) + + def add_tensor(self, z): + return (self.x + self.y) * z + + @torch._library.register_fake_class("_TorchScriptTesting::_ContainsTensor") + class FakeContainsTensor: + def __init__(self, t: torch.Tensor): + self.t = t + + @classmethod + def __obj_unflatten__(cls, flattend_foo): + return cls(**dict(flattend_foo)) + + def get(self): + return self.t + + @torch._library.register_fake_class("_TorchScriptTesting::_TensorQueue") + class FakeTensorQueue: + def __init__(self, queue): + self.queue = queue + + @classmethod + def __obj_unflatten__(cls, flattened_ctx): + return cls(**dict(flattened_ctx)) + + def push(self, x): + self.queue.append(x) + + def pop(self): + if self.is_empty(): + return torch.empty([]) + return self.queue.pop(0) + + def size(self): + return len(self.queue) + + def is_empty(self): + return len(self.queue) == 0 + + def float_size(self): + return float(len(self.queue)) + + @torch._library.register_fake_class("_TorchScriptTesting::_FlattenWithTensorOp") + class FakeFlatten: + def __init__(self, t): + self.t = t + + def get(self): + return self.t + + @classmethod + def __obj_unflatten__(cls, flattened_ctx): + return cls(**dict(flattened_ctx)) + + +def load_torchbind_test_lib(): + import unittest + + from torch.testing._internal.common_utils import ( # type: ignore[attr-defined] + find_library_location, + IS_FBCODE, + IS_MACOS, + IS_SANDCASTLE, + IS_WINDOWS, + ) + + if IS_MACOS: + raise unittest.SkipTest("non-portable load_library call used in test") + elif IS_SANDCASTLE or IS_FBCODE: + lib_file_path = Path("//caffe2/test/cpp/jit:test_custom_class_registrations") + elif IS_WINDOWS: + lib_file_path = find_library_location("torchbind_test.dll") + else: + lib_file_path = find_library_location("libtorchbind_test.so") + torch.ops.load_library(str(lib_file_path)) + + +@contextlib.contextmanager +def _register_py_impl_temporarily(op_overload, key, fn): + try: + op_overload.py_impl(key)(fn) + yield + finally: + del op_overload.py_kernels[key] + op_overload._dispatch_cache.clear() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/triton_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/triton_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..9961330e57bb98b0c2fdab93af180356ce4fd410 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/triton_utils.py @@ -0,0 +1,1109 @@ +# mypy: ignore-errors + +import unittest + +from torch.testing._internal.inductor_utils import ( + HAS_CUDA_AND_TRITON, + HAS_GPU, + HAS_XPU_AND_TRITON, +) +from torch.utils._triton import has_triton + + +requires_cuda_and_triton = unittest.skipUnless( + HAS_CUDA_AND_TRITON, "requires cuda and triton" +) +requires_xpu_and_triton = unittest.skipUnless( + HAS_XPU_AND_TRITON, "requires xpu and triton" +) +requires_gpu_and_triton = unittest.skipUnless( + HAS_XPU_AND_TRITON or HAS_CUDA_AND_TRITON, "requires gpu and triton" +) +requires_gpu = unittest.skipUnless(HAS_GPU, "requires gpu") + +if has_triton(): + import triton + from triton import language as tl + + import torch + + def _get_strange_configs() -> list[triton.Config]: + if torch.version.hip: + configs = [ + triton.Config( + { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 16, + "GROUP_SIZE_M": 4, + "matrix_instr_nonkdim": 16, + "waves_per_eu": 3, + "kpack": 2, + }, + num_stages=4, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 16, + "GROUP_SIZE_M": 4, + "matrix_instr_nonkdim": 16, + "waves_per_eu": 3, + "kpack": 2, + }, + num_stages=4, + num_warps=4, + ), + ] + else: + configs = [ + triton.Config( + { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 16, + "GROUP_SIZE_M": 4, + }, + num_stages=4, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + }, + num_stages=4, + num_warps=4, + ), + ] + return configs + + # Define here so that multiple tests can take advantage of it + @triton.jit + def add_kernel( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x + y + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def sub_kernel( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x - y + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def add_kernel_with_optional_param( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + ARGS_PASSED: "tl.constexpr", + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + if ARGS_PASSED == "two": + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x + y + else: + output = x + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def add_kernel_with_none_param_and_equal_to_1_arg( + in_ptr0, + in_ptr1, # in_ptr1 could be None + out_ptr, + n_elements, + stride, + ARGS_PASSED: "tl.constexpr", + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets * stride, mask=mask) + if ARGS_PASSED == "two": + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x + y + else: + output = x + tl.store(out_ptr + offsets * stride, output, mask=mask) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 128}, num_stages=3, num_warps=8), + triton.Config({"BLOCK_SIZE": 128}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_SIZE": 64}, num_stages=3, num_warps=8), + triton.Config({"BLOCK_SIZE": 64}, num_stages=4, num_warps=4), + ], + key=[], + ) + @triton.jit + def add_kernel_autotuned( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x + y + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 128}, num_stages=3, num_warps=8), + triton.Config({"BLOCK_SIZE": 128}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_SIZE": 64}, num_stages=3, num_warps=8), + triton.Config({"BLOCK_SIZE": 64}, num_stages=4, num_warps=4), + ], + key=[], + ) + @triton.jit + def sub_kernel_autotuned( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x - y + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 16}, num_stages=2, num_warps=2), + ], + key=[], + ) + @triton.jit + def add_kernel_autotuned_weird_param_order( + in_ptr0, + in_ptr1, + n_elements, + BLOCK_SIZE: "tl.constexpr", + out_ptr, + ): + # out_ptr is after an autotuned param that's declared as tl.constexpr. + # This param ordering can create bugs if not handled correctly. + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x + y + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.autotune( + configs=[ + triton.Config( + {"BLOCK_SIZE_X": 128, "BLOCK_SIZE_Y": 128}, num_stages=3, num_warps=8 + ), + triton.Config( + {"BLOCK_SIZE_X": 128, "BLOCK_SIZE_Y": 128}, num_stages=4, num_warps=4 + ), + triton.Config( + {"BLOCK_SIZE_X": 64, "BLOCK_SIZE_Y": 64}, num_stages=3, num_warps=8 + ), + triton.Config( + {"BLOCK_SIZE_X": 64, "BLOCK_SIZE_Y": 64}, num_stages=4, num_warps=4 + ), + ], + key=[], + ) + @triton.jit + def add_kernel_2d_autotuned( + in_ptr0, + in_ptr1, + out_ptr, + x_elements, + y_elements, + BLOCK_SIZE_X: "tl.constexpr", + BLOCK_SIZE_Y: "tl.constexpr", + ): + xoffset = tl.program_id(0) * BLOCK_SIZE_X + xindex = xoffset + tl.arange(0, BLOCK_SIZE_X)[:, None] + xmask = xindex < x_elements + yoffset = tl.program_id(1) * BLOCK_SIZE_Y + yindex = yoffset + tl.arange(0, BLOCK_SIZE_Y)[None, :] + ymask = yindex < y_elements + x1 = xindex + y0 = yindex + tmp0 = tl.load(in_ptr0 + (x1 + (x_elements * y0)), xmask & ymask) + tmp1 = tl.load(in_ptr0 + (y0 + (y_elements * x1)), xmask & ymask) + tmp2 = tmp0 + tmp1 + tl.store(out_ptr + (x1 + (x_elements * y0)), tmp2, xmask & ymask) + + def _dummy_early_config_prune(configs, *_, **__): + return configs + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE": 128}, num_stages=3, num_warps=8), + triton.Config({"BLOCK_SIZE": 64}, num_stages=4, num_warps=4), + ], + key=[], + warmup=10, + rep=20, + prune_configs_by={"early_config_prune": _dummy_early_config_prune}, + ) + @triton.jit + def add_kernel_autotuned_with_unsupported_args( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x + y + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def add_kernel_with_scaling( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + scaling_factor, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + output = (x + y) * scaling_factor + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def add_kernel_with_tma_1d_old_api( + in_desc_ptr0, + in_desc_ptr1, + out_desc_ptr, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + offset = pid * BLOCK_SIZE + + a = tl._experimental_descriptor_load( + in_desc_ptr0, + [offset], + [BLOCK_SIZE], + tl.float32, + ) + b = tl._experimental_descriptor_load( + in_desc_ptr1, + [offset], + [BLOCK_SIZE], + tl.float32, + ) + + output = a + b + + tl._experimental_descriptor_store( + out_desc_ptr, + output, + [offset], + ) + + @triton.jit + def add_kernel_with_tma_2d_old_api( + in_desc_ptr0, + in_desc_ptr1, + out_desc_ptr, + BLOCK_SIZE_X: "tl.constexpr", + BLOCK_SIZE_Y: "tl.constexpr", + ): + pid_x = tl.program_id(axis=0) + pid_y = tl.program_id(axis=1) + offset_x = pid_x * BLOCK_SIZE_X + offset_y = pid_y * BLOCK_SIZE_Y + + x = tl._experimental_descriptor_load( + in_desc_ptr0, + [offset_x, offset_y], + [BLOCK_SIZE_X, BLOCK_SIZE_Y], + tl.float32, + ) + y = tl._experimental_descriptor_load( + in_desc_ptr1, + [offset_x, offset_y], + [BLOCK_SIZE_X, BLOCK_SIZE_Y], + tl.float32, + ) + + output = x + y + + tl._experimental_descriptor_store( + out_desc_ptr, + output, + [offset_x, offset_y], + ) + + @triton.jit + def add_kernel_with_tma_1d_new_api( + in_desc_ptr0, + in_desc_ptr1, + out_desc_ptr, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + offset = pid * BLOCK_SIZE + + a = tl.load_tensor_descriptor( + in_desc_ptr0, + [offset], + ) + b = tl.load_tensor_descriptor( + in_desc_ptr1, + [offset], + ) + + output = a + b + + tl.store_tensor_descriptor( + out_desc_ptr, + [offset], + output, + ) + + @triton.jit + def add_kernel_with_tma_2d_new_api( + in_desc_ptr0, + in_desc_ptr1, + out_desc_ptr, + BLOCK_SIZE_X: "tl.constexpr", + BLOCK_SIZE_Y: "tl.constexpr", + ): + pid_x = tl.program_id(axis=0) + pid_y = tl.program_id(axis=1) + offset_x = pid_x * BLOCK_SIZE_X + offset_y = pid_y * BLOCK_SIZE_Y + + x = tl.load_tensor_descriptor( + in_desc_ptr0, + [offset_x, offset_y], + ) + y = tl.load_tensor_descriptor( + in_desc_ptr1, + [offset_x, offset_y], + ) + + output = x + y + + tl.store_tensor_descriptor( + out_desc_ptr, + [offset_x, offset_y], + output, + ) + + @triton.jit + def add_kernel_on_device_tma_old_api( + a_ptr, + b_ptr, + c_ptr, + m, + n, + workspace, + BLOCK_SIZE: "tl.constexpr", + ): + a_desc_ptr = workspace + b_desc_ptr = workspace + 128 + c_desc_ptr = workspace + 256 + tl.extra.cuda.experimental_device_tensormap_create2d( + desc_ptr=a_desc_ptr, + global_address=a_ptr, + load_size=[BLOCK_SIZE, BLOCK_SIZE], + global_size=[m, n], + element_ty=a_ptr.dtype.element_ty, + ) + tl.extra.cuda.experimental_device_tensormap_create2d( + desc_ptr=b_desc_ptr, + global_address=b_ptr, + load_size=[BLOCK_SIZE, BLOCK_SIZE], + global_size=[m, n], + element_ty=b_ptr.dtype.element_ty, + ) + tl.extra.cuda.experimental_device_tensormap_create2d( + desc_ptr=c_desc_ptr, + global_address=c_ptr, + load_size=[BLOCK_SIZE, BLOCK_SIZE], + global_size=[m, n], + element_ty=c_ptr.dtype.element_ty, + ) + + tl.extra.cuda.experimental_tensormap_fenceproxy_acquire(a_desc_ptr) + tl.extra.cuda.experimental_tensormap_fenceproxy_acquire(b_desc_ptr) + tl.extra.cuda.experimental_tensormap_fenceproxy_acquire(c_desc_ptr) + + pid_x = tl.program_id(axis=0) + pid_y = tl.program_id(axis=1) + offset_x = pid_x * BLOCK_SIZE + offset_y = pid_y * BLOCK_SIZE + + # Load data using the tensor descriptors + a = tl._experimental_descriptor_load( + a_desc_ptr, + [offset_x, offset_y], + [BLOCK_SIZE, BLOCK_SIZE], + tl.float32, + ) + b = tl._experimental_descriptor_load( + b_desc_ptr, + [offset_x, offset_y], + [BLOCK_SIZE, BLOCK_SIZE], + tl.float32, + ) + + # Perform addition + output = a + b + + # Store the result + tl._experimental_descriptor_store( + c_desc_ptr, + output, + [offset_x, offset_y], + ) + + @triton.jit + def add_kernel_on_device_tma_new_api( + a_ptr, + b_ptr, + c_ptr, + m, + n, + workspace, # unused but left here to match the old API kernel + BLOCK_SIZE: "tl.constexpr", + ): + # Create tensor descriptors using the new API + a_desc = tl.make_tensor_descriptor( + base=a_ptr, + shape=[m, n], + strides=[n, 1], + block_shape=[BLOCK_SIZE, BLOCK_SIZE], + ) + b_desc = tl.make_tensor_descriptor( + base=b_ptr, + shape=[m, n], + strides=[n, 1], + block_shape=[BLOCK_SIZE, BLOCK_SIZE], + ) + c_desc = tl.make_tensor_descriptor( + base=c_ptr, + shape=[m, n], + strides=[n, 1], + block_shape=[BLOCK_SIZE, BLOCK_SIZE], + ) + + pid_x = tl.program_id(axis=0) + pid_y = tl.program_id(axis=1) + offset_x = pid_x * BLOCK_SIZE + offset_y = pid_y * BLOCK_SIZE + + # Load data using the tensor descriptors with the new API + a = tl.load_tensor_descriptor( + a_desc, + [offset_x, offset_y], + ) + b = tl.load_tensor_descriptor( + b_desc, + [offset_x, offset_y], + ) + + # Perform addition + output = a + b + + # Store the result with the new API + tl.store_tensor_descriptor( + c_desc, + [offset_x, offset_y], + output, + ) + + @triton.jit + def mul2_kernel( + in_ptr0, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + output = 2 * x + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def mul2_inplace_kernel( + ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(ptr + offsets, mask=mask) + output = 2 * x + tl.store(ptr + offsets, output, mask=mask) + + @triton.jit + def zero_negs(x): + return tl.where(x >= 0, x, 0) + + @triton.jit + def indirection_kernel( + in_ptr0, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ACTIVATION: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + if ACTIVATION == "mul2_inplace_kernel": + mul2_inplace_kernel(in_ptr0, n_elements, BLOCK_SIZE=BLOCK_SIZE) + elif ACTIVATION == "add_kernel": + add_kernel(in_ptr0, in_ptr0, out_ptr, n_elements, BLOCK_SIZE=BLOCK_SIZE) + x = tl.load(in_ptr0 + offsets, mask=mask) + tl.store(out_ptr + offsets, x, mask=mask) + + @triton.jit + def double_strided_kernel( + in_ptr, + out_ptr, + in_y_stride, + out_y_stride, + X_BLOCK_SIZE: "tl.constexpr", + Y_BLOCK_SIZE: "tl.constexpr", + ): + xid = tl.program_id(axis=0) + yid = tl.program_id(axis=1) + x_start = xid * X_BLOCK_SIZE + y_start = yid * Y_BLOCK_SIZE + x_offsets = x_start + tl.arange(0, X_BLOCK_SIZE) + y_offsets = y_start + tl.arange(0, Y_BLOCK_SIZE) + src_offsets = y_offsets[:, None] * in_y_stride + x_offsets[None, :] + dst_offsets = y_offsets[:, None] * out_y_stride + x_offsets[None, :] + src = tl.load(in_ptr + src_offsets) + tl.store(out_ptr + dst_offsets, src * 2.0) + + @triton.jit + def inline_asm_kernel_is_pure_true( + X, Y, Z, n: "tl.constexpr", BLOCK: "tl.constexpr" + ): + x = tl.load(X + tl.arange(0, BLOCK)) + y = tl.load(Y + tl.arange(0, BLOCK)) + s = tl.full([BLOCK], n, tl.int32) + z = tl.inline_asm_elementwise( + "shf.l.wrap.b32 $0, $1, $2, $3;", + "=r,r, r, r", + [x, y, s], + dtype=tl.int32, + is_pure=True, + pack=1, + ) + tl.store(Z + tl.arange(0, BLOCK), z) + + @triton.jit + def inline_asm_kernel_is_pure_false( + X, Y, Z, n: "tl.constexpr", BLOCK: "tl.constexpr" + ): + x = tl.load(X + tl.arange(0, BLOCK)) + y = tl.load(Y + tl.arange(0, BLOCK)) + s = tl.full([BLOCK], n, tl.int32) + z = tl.inline_asm_elementwise( + "shf.l.wrap.b32 $0, $1, $2, $3;", + "=r,r, r, r", + [x, y, s], + dtype=tl.int32, + is_pure=False, + pack=1, + ) + tl.store(Z + tl.arange(0, BLOCK), z) + + @triton.jit + def add_kernel_with_block_ptr( + x_ptr, + y_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + x = tl.load( + tl.make_block_ptr( + base=x_ptr, + shape=[n_elements], + strides=[1], + offsets=[block_start], + block_shape=[BLOCK_SIZE], + order=[0], + ), + boundary_check=[0], + ) + y = tl.load( + tl.make_block_ptr( + base=y_ptr, + shape=[n_elements], + strides=[1], + offsets=[block_start], + block_shape=[BLOCK_SIZE], + order=[0], + ), + boundary_check=[0], + ) + output = x + y + tl.store( + tl.make_block_ptr( + base=output_ptr, + shape=[n_elements], + strides=[1], + offsets=[block_start], + block_shape=[BLOCK_SIZE], + order=[0], + ), + output, + boundary_check=[0], + ) + + @triton.jit + def kernel_with_block_ptr_2d( + x_ptr, + output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + x = tl.load( + tl.make_block_ptr( + base=x_ptr, + shape=[n_elements, 1], + strides=[1, 1], + offsets=[block_start, 0], + block_shape=[BLOCK_SIZE, 1], + order=[1, 0], + ), + boundary_check=[0], + ) + output = x + tl.store( + tl.make_block_ptr( + base=output_ptr, + shape=[n_elements, 1], + strides=[1, 1], + offsets=[block_start, 0], + block_shape=[BLOCK_SIZE, 1], + order=[1, 0], + ), + output, + boundary_check=[0], + ) + + from triton.language import load, store + + @triton.jit + def add_kernel_with_import( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = load(in_ptr0 + offsets, mask=mask) + y = load(in_ptr1 + offsets, mask=mask) + output = x + y + store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def cond_op_kernel( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + if tl.program_id(0) == 0: + output = x + y + else: + output = x * y + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def atomic_add_kernel( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x + y + tl.atomic_add(out_ptr + offsets, output, mask=mask) + + @triton.jit + def add_4_times_kernel( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + for _ in range(2): + output = x + y + tl.store(out_ptr + offsets, output, mask=mask) + i = 2 + while i > 0: + i -= 1 + output = x + y + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def add_kernel_out_of_order_fn2( + in_ptr0, + in_ptr1, + n_elements, + out_ptr, + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x + y + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.autotune( + configs=[ + triton.Config( + { + "BLOCK_SIZE_M": 16, + "BLOCK_SIZE_N": 16, + "BLOCK_SIZE_K": 16, + "GROUP_SIZE_M": 4, + }, + num_stages=4, + num_warps=4, + ), + triton.Config( + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 64, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 8, + }, + num_stages=4, + num_warps=4, + ), + ], + key=["M_ptr", "N", "K"], + ) + @triton.jit + def strange_config_matmul_kernel( + a_ptr, + b_ptr, + c_ptr, + M_ptr, + N, + K, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, + ): + # This is a simplified matmul from Triton tutorial. + pid = tl.program_id(axis=0) + M = tl.load(M_ptr) + if M == 0 and BLOCK_SIZE_M > 32: + # This will run the full matmul if BLOCK_SIZE_M > 32 + M = 4096 + elif M == 0: + # This directly returns, which will cut short the bad config of 16-block size. + return + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] + offs_k[None, :]) + b_ptrs = b_ptr + (offs_k[:, None] + offs_bn[None, :]) + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + accumulator = tl.dot(a, b, accumulator) + a_ptrs += BLOCK_SIZE_K + b_ptrs += BLOCK_SIZE_K + c = accumulator.to(tl.float16) + + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + offs_cm[:, None] + offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) + + @triton.jit + def kernel_with_docstring_double_quotes(out_ptr, numel, BLOCK_SIZE: tl.constexpr): + """ + This kernel contains a triple-quote docstring w/ double quotes. + Make sure that codegen sanitizes the docstring. + """ + pid = tl.program_id(axis=0) + offsets = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + ones = tl.full([BLOCK_SIZE], 1.0, dtype=tl.float32) + tl.store(out_ptr + offsets, ones, mask=offsets < numel) + + @triton.jit + def kernel_with_docstring_single_quotes(out_ptr, numel, BLOCK_SIZE: tl.constexpr): + ''' + This kernel contains a triple-quote docstring w/ single quotes + Make sure that codegen sanitizes the docstring. + To prevent it from being linted to double quotes: """!!!""" + ''' + pid = tl.program_id(axis=0) + offsets = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + ones = tl.full([BLOCK_SIZE], 1.0, dtype=tl.float32) + tl.store(out_ptr + offsets, ones, mask=offsets < numel) + + @triton.jit + def kernel_inline_asm_double_quotes( + in_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr + ): + pid = tl.program_id(axis=0) + offsets = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + data = tl.load(in_ptr + offsets, mask=offsets < numel) + cos_pow = tl.inline_asm_elementwise( + asm=""" + { + cos.approx.f32 $0, $1; + ex2.approx.f32 $0, $0; + } + """, + constraints=("=r, r"), + args=[data], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + tl.store(out_ptr + offsets, cos_pow, mask=offsets < numel) + + @triton.jit + def kernel_inline_asm_single_quotes( + in_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr + ): + pid = tl.program_id(axis=0) + offsets = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + data = tl.load(in_ptr + offsets, mask=offsets < numel) + cos_pow = tl.inline_asm_elementwise( + asm=''' + { + // double quotes to pacify the linter """!!!""" + cos.approx.f32 $0, $1; + ex2.approx.f32 $0, $0; + } + ''', + constraints=("=r, r"), + args=[data], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + tl.store(out_ptr + offsets, cos_pow, mask=offsets < numel) + + @triton.jit + def kernel_inline_asm_rocm_double_quotes( + in_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr + ): + pid = tl.program_id(axis=0) + offsets = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + data = tl.load(in_ptr + offsets, mask=offsets < numel) + cos_pow = tl.inline_asm_elementwise( + asm=""" + v_sin_f32 $0, $1 + v_exp_f32 $0, $0 + """, + constraints=("=v, v"), + args=[data], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + tl.store(out_ptr + offsets, cos_pow, mask=offsets < numel) + + @triton.jit + def kernel_inline_asm_rocm_single_quotes( + in_ptr, out_ptr, numel, BLOCK_SIZE: tl.constexpr + ): + pid = tl.program_id(axis=0) + offsets = tl.arange(0, BLOCK_SIZE) + pid * BLOCK_SIZE + data = tl.load(in_ptr + offsets, mask=offsets < numel) + cos_pow = tl.inline_asm_elementwise( + asm=""" + v_sin_f32 $0, $1 + v_exp_f32 $0, $0 + """, + constraints=("=v, v"), + args=[data], + dtype=tl.float32, + is_pure=True, + pack=1, + ) + tl.store(out_ptr + offsets, cos_pow, mask=offsets < numel) + + @triton.jit + def add_kernel_with_boolean_param( + in_ptr0, + in_ptr1, + out_ptr, + n_elements, + add_xy, # boolean param + BLOCK_SIZE: "tl.constexpr", + ): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=mask) + if add_xy: + y = tl.load(in_ptr1 + offsets, mask=mask) + output = x + y + else: + output = x + tl.store(out_ptr + offsets, output, mask=mask) + + @triton.jit + def masked_add_kernel_with_bool_tensor( + in_ptr0, + in_ptr1, + mask_ptr, + out_ptr, + n_elements, + BLOCK_SIZE: "tl.constexpr", + ): + """Kernel that loads a bool tensor and uses it as a mask.""" + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + valid = offsets < n_elements + x = tl.load(in_ptr0 + offsets, mask=valid) + y = tl.load(in_ptr1 + offsets, mask=valid) + keep = tl.load(mask_ptr + offsets, mask=valid, other=0) != 0 + output = tl.where(keep, x + y, x) + tl.store(out_ptr + offsets, output, mask=valid) + + # support the old (experimental) and new (tensor_descriptor) APIs + def create_tensor_descriptor_shim( + tensor, block_sizes: list[int], new_api: bool = True + ): + if new_api: + return triton.tools.tensor_descriptor.TensorDescriptor.from_tensor( + tensor, block_sizes + ) + else: + if len(block_sizes) == 1: + return triton.tools.experimental_descriptor.create_1d_tma_descriptor( + tensor.data_ptr(), + tensor.size(0), + block_sizes[0], + tensor.element_size(), + ) + else: + if len(block_sizes) != 2: + raise AssertionError( + f"Expected len(block_sizes) == 2, got {len(block_sizes)}" + ) + return triton.tools.experimental_descriptor.create_2d_tma_descriptor( + tensor.data_ptr(), + tensor.size(0), + tensor.size(1), + block_sizes[0], + block_sizes[1], + tensor.element_size(), + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/two_tensor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/two_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..e03f37f561cc5468589e74b4d28d55b92fa137a5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_internal/two_tensor.py @@ -0,0 +1,116 @@ +# mypy: ignore-errors + +import torch +import torch.utils._pytree as pytree +from torch._export.wrappers import mark_subclass_constructor_exportable_experimental +from torch.utils._python_dispatch import return_and_correct_aliasing + + +# A simple tensor subclass that holds two tensors internally, and runs every op on both tensors. +class TwoTensor(torch.Tensor): + @staticmethod + def __new__(cls, a, b, outer_size=None, outer_stride=None, *, requires_grad=None): + if outer_size is None: + outer_size = a.size() + if outer_stride is None: + outer_stride = a.stride() + + if not ( + a.device == b.device + and a.layout == b.layout + and a.requires_grad == b.requires_grad + and a.dtype == b.dtype + ): + raise AssertionError( + "Expected a and b to have same device, layout, requires_grad, and dtype" + ) + # I guess it would be more accurate to represent the shape as torch.cat(a, b).shape + shape = outer_size + kwargs = {} + kwargs["strides"] = outer_stride + kwargs["storage_offset"] = a.storage_offset() + kwargs["device"] = a.device + kwargs["layout"] = a.layout + kwargs["requires_grad"] = requires_grad or a.requires_grad + kwargs["dtype"] = a.dtype + out = torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs) + + if a.shape != b.shape: + raise AssertionError( + f"Expected a.shape == b.shape, got {a.shape} != {b.shape}" + ) + if a.stride() != b.stride(): + raise AssertionError( + f"Expected a.stride() == b.stride(), got {a.stride()} != {b.stride()}" + ) + if a.storage_offset() != b.storage_offset(): + raise AssertionError( + f"Expected a.storage_offset() == b.storage_offset(), " + f"got {a.storage_offset()} != {b.storage_offset()}" + ) + return out + + @torch._disable_dynamo + @mark_subclass_constructor_exportable_experimental + def __init__(self, a, b, outer_size=None, outer_stride=None, *, requires_grad=None): + self.a = a + self.b = b + + def __repr__(self): + a_repr = repr(self.a) + b_repr = repr(self.b) + return f"TwoTensor({a_repr}, {b_repr})" + + def __tensor_flatten__(self): + return ["a", "b"], None + + @staticmethod + def __tensor_unflatten__(inner_tensors, meta, outer_size, outer_stride): + if meta is not None: + raise AssertionError("Expected meta to be None") + a, b = inner_tensors["a"], inner_tensors["b"] + if type(a) is torch.Tensor: + if outer_size is None: + raise AssertionError("Expected outer_size to not be None") + if outer_stride is None: + raise AssertionError("Expected outer_stride to not be None") + return TwoTensor(a, b, outer_size, outer_stride) + + @classmethod + def __torch_dispatch__(cls, func, types, args, kwargs): + if kwargs is None: + kwargs = {} + args_a = pytree.tree_map_only(TwoTensor, lambda x: x.a, args) + args_b = pytree.tree_map_only(TwoTensor, lambda x: x.b, args) + + kwargs_a = pytree.tree_map_only(TwoTensor, lambda x: x.a, kwargs) + kwargs_b = pytree.tree_map_only(TwoTensor, lambda x: x.b, kwargs) + + out_a = func(*args_a, **kwargs_a) + out_b = func(*args_b, **kwargs_b) + out_a_flat, spec = pytree.tree_flatten(out_a) + out_b_flat = pytree.tree_leaves(out_b) + # for aten ops that return non-tensors, just assume that + # our two inner tensors return the same value + out_flat = [ + cls(o_a, o_b) if isinstance(o_a, torch.Tensor) else o_a + for o_a, o_b in zip(out_a_flat, out_b_flat, strict=True) + ] + out = pytree.tree_unflatten(out_flat, spec) + from torch._higher_order_ops.cond import cond_op + + if func is cond_op: + return out + else: + return return_and_correct_aliasing(func, args, kwargs, out) + + def get_elem_a(self): + return self.a + + +class TwoTensorMode(torch.utils._python_dispatch.TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + out = func(*args, **kwargs) + if torch._subclasses.fake_tensor._is_tensor_constructor(func): + out = TwoTensor(out, out.clone()) + return out diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..02eeb1f46f5c8610dd3d227036c7320a8deaeb43 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/testing/_utils.py @@ -0,0 +1,62 @@ +# mypy: allow-untyped-defs +import contextlib + +import torch + + +# Common testing utilities for use in public testing APIs. +# NB: these should all be importable without optional dependencies +# (like numpy and expecttest). + + +def wrapper_set_seed(op, *args, **kwargs): + """Wrapper to set seed manually for some functions like dropout + See: https://github.com/pytorch/pytorch/pull/62315#issuecomment-896143189 for more details. + """ + with freeze_rng_state(): + torch.manual_seed(42) + output = op(*args, **kwargs) + + if isinstance(output, torch.Tensor) and output.device.type == "lazy": + # We need to call mark step inside freeze_rng_state so that numerics + # match eager execution + torch._lazy.mark_step() # type: ignore[attr-defined] + + return output + + +@contextlib.contextmanager +def freeze_rng_state(): + # no_dispatch needed for test_composite_compliance + # Some OpInfos use freeze_rng_state for rng determinism, but + # test_composite_compliance overrides dispatch for all torch functions + # which we need to disable to get and set rng state + with torch.utils._mode_utils.no_dispatch(), torch._C._DisableFuncTorch(): + rng_state = torch.get_rng_state() + if torch.accelerator.is_available(): + accelerator = torch.accelerator.current_accelerator(check_available=True) + if accelerator is not None: + accelerator_rng_state = torch.get_device_module( + accelerator.type + ).get_rng_state() + try: + yield + finally: + # Modes are not happy with torch.cuda.set_rng_state + # because it clones the state (which could produce a Tensor Subclass) + # and then grabs the new tensor's data pointer in generator.set_state. + # + # In the long run torch.cuda.set_rng_state should probably be + # an operator. + # + # NB: Mode disable is to avoid running cross-ref tests on this seeding + with torch.utils._mode_utils.no_dispatch(), torch._C._DisableFuncTorch(): + if torch.accelerator.is_available(): + accelerator = torch.accelerator.current_accelerator( + check_available=True + ) + if accelerator is not None: + torch.get_device_module(accelerator.type).set_rng_state( + accelerator_rng_state # type: ignore[possibly-undefined] + ) + torch.set_rng_state(rng_state) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/torch_version.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/torch_version.py new file mode 100644 index 0000000000000000000000000000000000000000..0496a1b564feefe4a52280e2d7f268516f256a70 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/torch_version.py @@ -0,0 +1,66 @@ +from collections.abc import Iterable +from typing import Any + +from torch._vendor.packaging.version import InvalidVersion, Version +from torch.version import __version__ as internal_version + + +__all__ = ["TorchVersion"] + + +class TorchVersion(str): + """A string with magic powers to compare to both Version and iterables! + Prior to 1.10.0 torch.__version__ was stored as a str and so many did + comparisons against torch.__version__ as if it were a str. In order to not + break them we have TorchVersion which masquerades as a str while also + having the ability to compare against both packaging.version.Version as + well as tuples of values, eg. (1, 2, 1) + Examples: + Comparing a TorchVersion object to a Version object + TorchVersion('1.10.0a') > Version('1.10.0a') + Comparing a TorchVersion object to a Tuple object + TorchVersion('1.10.0a') > (1, 2) # 1.2 + TorchVersion('1.10.0a') > (1, 2, 1) # 1.2.1 + Comparing a TorchVersion object against a string + TorchVersion('1.10.0a') > '1.2' + TorchVersion('1.10.0a') > '1.2.1' + """ + + __slots__ = () + + # fully qualified type names here to appease mypy + def _convert_to_version(self, inp: Any) -> Any: + if isinstance(inp, Version): + return inp + elif isinstance(inp, str): + return Version(inp) + elif isinstance(inp, Iterable): + # Ideally this should work for most cases by attempting to group + # the version tuple, assuming the tuple looks (MAJOR, MINOR, ?PATCH) + # Examples: + # * (1) -> Version("1") + # * (1, 20) -> Version("1.20") + # * (1, 20, 1) -> Version("1.20.1") + return Version(".".join(str(item) for item in inp)) + else: + raise InvalidVersion(inp) + + def _cmp_wrapper(self, cmp: Any, method: str) -> bool: + try: + return getattr(Version(self), method)(self._convert_to_version(cmp)) + except BaseException as e: + if not isinstance(e, InvalidVersion): + raise + # Fall back to regular string comparison if dealing with an invalid + # version like 'parrot' + return getattr(super(), method)(cmp) + + +for cmp_method in ["__gt__", "__lt__", "__eq__", "__ge__", "__le__"]: + setattr( + TorchVersion, + cmp_method, + lambda x, y, method=cmp_method: x._cmp_wrapper(y, method), + ) + +__version__ = TorchVersion(internal_version) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/types.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/types.py new file mode 100644 index 0000000000000000000000000000000000000000..d495a129f40c779e07f55ea8e6acbb0c3734da50 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/types.py @@ -0,0 +1,131 @@ +# In some cases, these basic types are shadowed by corresponding +# top-level values. The underscore variants let us refer to these +# types. See https://github.com/python/mypy/issues/4146 for why these +# workarounds is necessary +import os +from builtins import ( # noqa: F401 + bool as _bool, + bytes as _bytes, + complex as _complex, + float as _float, + int as _int, + str as _str, +) +from collections.abc import Sequence +from typing import Any, IO, TYPE_CHECKING, TypeAlias, Union +from typing_extensions import Self + +# `as` imports have better static analysis support than assignment `ExposedType: TypeAlias = HiddenType` +from torch import ( # noqa: F401 + device as _device, + DispatchKey, + dtype as _dtype, + layout as _layout, + qscheme as _qscheme, + Size, + SymBool, + SymFloat, + SymInt, + Tensor, +) + + +if TYPE_CHECKING: + from torch.autograd.graph import GradientEdge + + +__all__ = ["Number", "Device", "FileLike", "Storage"] + +# Convenience aliases for common composite types that we need +# to talk about in PyTorch +_TensorOrTensors: TypeAlias = Tensor | Sequence[Tensor] # noqa: PYI047 +_TensorOrOptionalTensors: TypeAlias = Tensor | Sequence[Tensor | None] # noqa: PYI047 +_TensorOrTensorsOrGradEdge: TypeAlias = Union[ # noqa: PYI047 + Tensor, + Sequence[Tensor], + "GradientEdge", + Sequence["GradientEdge"], +] + +_size: TypeAlias = Size | list[int] | tuple[int, ...] # noqa: PYI042,PYI047 +_symsize: TypeAlias = Size | Sequence[int | SymInt] # noqa: PYI042,PYI047 +_dispatchkey: TypeAlias = str | DispatchKey # noqa: PYI042,PYI047 + +# int or SymInt +IntLikeType: TypeAlias = int | SymInt +# float or SymFloat +FloatLikeType: TypeAlias = float | SymFloat +# bool or SymBool +BoolLikeType: TypeAlias = bool | SymBool + +py_sym_types = (SymInt, SymFloat, SymBool) # left un-annotated intentionally +PySymType: TypeAlias = SymInt | SymFloat | SymBool + +# Meta-type for "numeric" things; matches our docs +Number: TypeAlias = int | float | bool +# tuple for isinstance(x, Number) checks. +# FIXME: refactor once python 3.9 support is dropped. +_Number = (int, float, bool) + +FileLike: TypeAlias = str | os.PathLike[str] | IO[bytes] + +# Meta-type for "device-like" things. Not to be confused with 'device' (a +# literal device object). This nomenclature is consistent with PythonArgParser. +# None means use the default device (typically CPU) +Device: TypeAlias = _device | str | int | None + + +# Storage protocol implemented by ${Type}StorageBase classes +class Storage: + _cdata: int + device: _device + dtype: _dtype + _torch_load_uninitialized: bool + + def __deepcopy__(self, memo: dict[int, Any]) -> Self: + raise NotImplementedError + + def _new_shared(self, size: int) -> Self: + raise NotImplementedError + + def _write_file( + self, + f: Any, + is_real_file: bool, + save_size: bool, + element_size: int, + ) -> None: + raise NotImplementedError + + def element_size(self) -> int: + raise NotImplementedError + + def is_shared(self) -> bool: + raise NotImplementedError + + def share_memory_(self) -> Self: + raise NotImplementedError + + def nbytes(self) -> int: + raise NotImplementedError + + def cpu(self) -> Self: + raise NotImplementedError + + def data_ptr(self) -> int: + raise NotImplementedError + + def from_file( + self, + filename: str, + shared: bool = False, + nbytes: int = 0, + ) -> Self: + raise NotImplementedError + + def _new_with_file( + self, + f: Any, + element_size: int, + ) -> Self: + raise NotImplementedError diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..1c3ec1579006399a6025e8475470effeccf7cd22 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/__init__.py @@ -0,0 +1,107 @@ +# mypy: allow-untyped-defs + +import copyreg +import os.path as _osp +import weakref + +import torch +from torch.utils import ( + backcompat as backcompat, + collect_env as collect_env, + data as data, + deterministic as deterministic, + hooks as hooks, +) +from torch.utils.backend_registration import ( + generate_methods_for_privateuse1_backend, + rename_privateuse1_backend, +) +from torch.utils.cpp_backtrace import get_cpp_backtrace +from torch.utils.throughput_benchmark import ThroughputBenchmark + + +def set_module(obj, mod): + """ + Set the module attribute on a python object for a given object for nicer printing + """ + if not isinstance(mod, str): + raise TypeError("The mod argument should be a string") + obj.__module__ = mod + + +cmake_prefix_path = _osp.join(_osp.dirname(_osp.dirname(__file__)), "share", "cmake") + + +def swap_tensors(t1, t2): + """ + This function swaps the content of the two Tensor objects. + At a high level, this will make t1 have the content of t2 while preserving + its identity. + + This will not work if t1 and t2 have different slots. + """ + # Ensure there are no weakrefs + if weakref.getweakrefs(t1): + raise RuntimeError("Cannot swap t1 because it has weakref associated with it") + if weakref.getweakrefs(t2): + raise RuntimeError("Cannot swap t2 because it has weakref associated with it") + t1_slots = set(copyreg._slotnames(t1.__class__)) # type: ignore[attr-defined] + t2_slots = set(copyreg._slotnames(t2.__class__)) # type: ignore[attr-defined] + if t1_slots != t2_slots: + raise RuntimeError("Cannot swap t1 and t2 if they have different slots") + + def swap_attr(name): + tmp = getattr(t1, name) + setattr(t1, name, (getattr(t2, name))) + setattr(t2, name, tmp) + + def error_pre_hook(grad_outputs): + raise RuntimeError( + "Trying to execute AccumulateGrad node that was poisoned by swap_tensors " + "this can happen when you try to run backward on a tensor that was swapped. " + "For a module m with `torch.__future__.set_swap_module_params_on_conversion(True)` " + "you should not change the device or dtype of the module (e.g. `m.cpu()` or `m.half()`) " + "between running forward and backward. To resolve this, please only change the " + "device/dtype before running forward (or after both forward and backward)." + ) + + def check_use_count(t, name="t1"): + use_count = t._use_count() + error_str = ( + f"Expected use_count of {name} to be 1 or 2 with an AccumulateGrad node but got {use_count} " + f"make sure you are not holding references to the tensor in other places." + ) + if use_count > 1: + if use_count == 2 and t.is_leaf: + accum_grad_node = torch.autograd.graph.get_gradient_edge(t).node + # Make sure that the accumulate_grad node was not lazy_init-ed by get_gradient_edge + if t._use_count() == 2: + accum_grad_node.register_prehook(error_pre_hook) + else: + raise RuntimeError(error_str) + else: + raise RuntimeError(error_str) + + check_use_count(t1, "t1") + check_use_count(t2, "t2") + + # Swap the types + # Note that this will fail if there are mismatched slots + swap_attr("__class__") + + # Swap the dynamic attributes + swap_attr("__dict__") + + # Swap the slots + for slot in t1_slots: + if hasattr(t1, slot) and hasattr(t2, slot): + swap_attr(slot) + elif hasattr(t1, slot): + setattr(t2, slot, (getattr(t1, slot))) + delattr(t1, slot) + elif hasattr(t2, slot): + setattr(t1, slot, (getattr(t2, slot))) + delattr(t2, slot) + + # Swap the at::Tensor they point to + torch._C._swap_tensor_impl(t1, t2) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_appending_byte_serializer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_appending_byte_serializer.py new file mode 100644 index 0000000000000000000000000000000000000000..82cced0b3dc826484e2d44788e4814b9dbe32693 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_appending_byte_serializer.py @@ -0,0 +1,135 @@ +import base64 +import zlib +from collections.abc import Callable, Iterable +from typing import Generic, TypeVar + + +T = TypeVar("T") + +_ENCODING_VERSION: int = 1 + +__all__ = ["AppendingByteSerializer"] + + +####################################### +# Helper classes +####################################### + +CHECKSUM_DIGEST_SIZE = 4 + + +class BytesWriter: + def __init__(self) -> None: + # Reserve CHECKSUM_DIGEST_SIZE bytes for checksum + self._data = bytearray(CHECKSUM_DIGEST_SIZE) + + def write_uint64(self, i: int) -> None: + self._data.extend(i.to_bytes(8, byteorder="big", signed=False)) + + def write_str(self, s: str) -> None: + payload = base64.b64encode(s.encode("utf-8")) + self.write_bytes(payload) + + def write_bytes(self, b: bytes) -> None: + self.write_uint64(len(b)) + self._data.extend(b) + + def to_bytes(self) -> bytes: + digest = zlib.crc32(self._data[CHECKSUM_DIGEST_SIZE:]).to_bytes( + 4, byteorder="big", signed=False + ) + if len(digest) != CHECKSUM_DIGEST_SIZE: + raise AssertionError("Computed checksum digest has unexpected size") + self._data[0:CHECKSUM_DIGEST_SIZE] = digest + return bytes(self._data) + + +class BytesReader: + def __init__(self, data: bytes) -> None: + # Check for data corruption + if len(data) < CHECKSUM_DIGEST_SIZE: + raise AssertionError("Input data is too short to contain checksum") + digest = zlib.crc32(data[CHECKSUM_DIGEST_SIZE:]).to_bytes( + 4, byteorder="big", signed=False + ) + if len(digest) != CHECKSUM_DIGEST_SIZE: + raise AssertionError("Computed checksum digest has unexpected size") + if data[0:CHECKSUM_DIGEST_SIZE] != digest: + raise RuntimeError( + "Bytes object is corrupted, checksum does not match. " + f"Expected: {data[0:CHECKSUM_DIGEST_SIZE]!r}, Got: {digest!r}" + ) + + self._data = data + self._i = CHECKSUM_DIGEST_SIZE + + def is_finished(self) -> bool: + return len(self._data) == self._i + + def read_uint64(self) -> int: + result = int.from_bytes( + self._data[self._i : self._i + 8], byteorder="big", signed=False + ) + self._i += 8 + return result + + def read_str(self) -> str: + return base64.b64decode(self.read_bytes()).decode("utf-8") + + def read_bytes(self) -> bytes: + size = self.read_uint64() + result = self._data[self._i : self._i + size] + self._i += size + return result + + +####################################### +# AppendingByteSerializer +####################################### + + +class AppendingByteSerializer(Generic[T]): + """ + Provides efficient serialization and deserialization of list of bytes + Note that this does not provide any guarantees around byte order + """ + + _serialize_fn: Callable[[BytesWriter, T], None] + _writer: BytesWriter + + def __init__( + self, + *, + serialize_fn: Callable[[BytesWriter, T], None], + ) -> None: + self._serialize_fn = serialize_fn + self.clear() + + def clear(self) -> None: + self._writer = BytesWriter() + # First 8-bytes are for version + self._writer.write_uint64(_ENCODING_VERSION) + + def append(self, data: T) -> None: + self._serialize_fn(self._writer, data) + + def extend(self, elems: Iterable[T]) -> None: + for elem in elems: + self.append(elem) + + def to_bytes(self) -> bytes: + return self._writer.to_bytes() + + @staticmethod + def to_list(data: bytes, *, deserialize_fn: Callable[[BytesReader], T]) -> list[T]: + reader = BytesReader(data) + if reader.read_uint64() != _ENCODING_VERSION: + raise AssertionError( + f"Encoding version mismatch in AppendingByteSerializer.to_list, \ + got {reader.read_uint64()}" + ) + + result: list[T] = [] + while not reader.is_finished(): + result.append(deserialize_fn(reader)) + return result diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_config_module.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_config_module.py new file mode 100644 index 0000000000000000000000000000000000000000..99e7cf8f9cc77f68c9c8b708ce9a84310ec467eb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_config_module.py @@ -0,0 +1,913 @@ +import contextlib +import copy +import hashlib +import importlib +import inspect +import io +import os +import pickle +import tokenize +import unittest +from collections.abc import Callable +from contextvars import ContextVar +from dataclasses import dataclass +from types import FunctionType, ModuleType +from typing import Any, Generic, NoReturn, Optional, TYPE_CHECKING, TypeVar +from typing_extensions import deprecated + +from torch._utils_internal import justknobs_check + + +# Types saved/loaded in configs +CONFIG_TYPES = (int, float, bool, type(None), str, list, set, tuple, dict) + +# Immutable scalar types that don't need deepcopy when returned from configs. +# Everything else is defensively copied to prevent accidental mutation. +_IMMUTABLE_CONFIG_TYPES = (int, float, bool, type(None), str, tuple) + + +# Duplicated, because mypy needs these types statically +T = TypeVar("T", bound=int | float | bool | str | list | set | tuple | dict | None) + + +_UNSET_SENTINEL = object() + + +@dataclass(kw_only=True) +class _Config(Generic[T]): + """Represents a config with richer behaviour than just a default value. + :: + i.e. + foo = Config(justknob="//foo:bar", default=False) + install_config_module(...) + + This configs must be installed with install_config_module to be used + + Precedence Order: + alias: If set, the directly use the value of the alias. + env_name_force: If set, this environment variable has precedence over + everything after this. + If multiple env variables are given, the precedence order is from + left to right. + user_override: If a user sets a value (i.e. foo.bar=True), that + has precedence over everything after this. User overrides are thread-local. + env_name_default: If set, this environment variable will override everything + after this. + If multiple env variables are given, the precedence order is from + left to right. + justknob: If this pytorch installation supports justknobs, that will + override defaults, but will not override the user_override precedence. + default: This value is the lowest precedence, and will be used if nothing is + set. + + Environment Variables: + These are interpreted to be either "0" or "1" to represent true and false. + + Arguments: + justknob: the name of the feature / JK. In OSS this is unused. + default: is the value to default this knob to in OSS. + alias: The alias config to read instead. + env_name_force: The environment variable, or list of, to read that is a FORCE + environment variable. I.e. it overrides everything except for alias. + env_name_default: The environment variable, or list of, to read that changes the + default behaviour. I.e. user overrides take preference. + """ + + default: T | object + justknob: str | None = None + env_name_default: list[str] | None = None + env_name_force: list[str] | None = None + value_type: type | None = None + alias: str | None = None + # Deprecation support + deprecated: bool = False + deprecation_message: str | None = None + + def __post_init__(self) -> None: + self.env_name_default = _Config.string_or_list_of_string_to_list( + self.env_name_default + ) + self.env_name_force = _Config.string_or_list_of_string_to_list( + self.env_name_force + ) + + if self.alias is not None: + if ( + self.default is not _UNSET_SENTINEL + or self.justknob is not None + or self.env_name_default is not None + or self.env_name_force is not None + ): + raise AssertionError( + "if alias is set, none of {default, justknob, \ + env_name_default and env_name_force} can be set" + ) + + @staticmethod + def string_or_list_of_string_to_list( + val: str | list[str] | None, + ) -> list[str] | None: + if val is None: + return None + if isinstance(val, str): + return [val] + if not isinstance(val, list): + raise AssertionError(f"val is not a list, got {type(val)}") + return val + + +# In runtime, we unbox the Config[T] to a T, but typechecker cannot see this, +# so in order to allow for this dynamic behavior to work correctly with +# typechecking we are going to lie to the typechecker that Config[T] returns +# a T. +if TYPE_CHECKING: + + def Config( + default: T | object = _UNSET_SENTINEL, + justknob: str | None = None, + env_name_default: str | list[str] | None = None, + env_name_force: str | list[str] | None = None, + value_type: type | None = None, + alias: str | None = None, + # Deprecation support + deprecated: bool = False, + deprecation_message: str | None = None, + ) -> T: ... + +else: + + def Config( + default: T | object = _UNSET_SENTINEL, + justknob: str | None = None, + env_name_default: str | list[str] | None = None, + env_name_force: str | list[str] | None = None, + value_type: type | None = None, + alias: str | None = None, + # Deprecation support + deprecated: bool = False, + deprecation_message: str | None = None, + ) -> _Config[T]: + return _Config( + default=default, + justknob=justknob, + env_name_default=env_name_default, + env_name_force=env_name_force, + value_type=value_type, + alias=alias, + # Deprecation support + deprecated=deprecated, + deprecation_message=deprecation_message, + ) + + +def _read_env_variable(name: str) -> bool | str | None: + value = os.environ.get(name) + if value == "1": + return True + if value == "0": + return False + return value + + +def install_config_module(module: ModuleType) -> None: + """ + Converts a module-level config into a `ConfigModule()`. + + See _config_typing.pyi for instructions on how to get the converted module to typecheck. + """ + + class ConfigModuleInstance(ConfigModule): + # __annotations__ is written to by Sphinx autodoc + _bypass_keys = set({"_is_dirty", "_hash_digest", "__annotations__"}) + + def visit( + source: ModuleType | type, + dest: ModuleType | SubConfigProxy, + prefix: str, + ) -> None: + """Walk the module structure and move everything to module._config""" + type_hints = inspect.get_annotations(source) + for key, value in list(source.__dict__.items()): + if ( + key.startswith("__") + or isinstance(value, (ModuleType, FunctionType)) + or ( + hasattr(value, "__module__") + and ( + value.__module__ == "typing" + or value.__module__.startswith("collections.abc") + ) + ) + # Handle from torch.utils._config_module import Config + or (isinstance(value, type) and issubclass(value, _Config)) + ): + continue + + name = f"{prefix}{key}" + annotated_type = type_hints.get(key, None) + if isinstance(value, CONFIG_TYPES): + config[name] = _ConfigEntry( + _Config(default=value, value_type=annotated_type), name + ) + if dest is module: + delattr(module, key) + elif isinstance(value, _Config): + if annotated_type is not None and value.value_type is None: + value.value_type = annotated_type + + config[name] = _ConfigEntry(value, name) + + if dest is module: + delattr(module, key) + elif isinstance(value, type): + if value.__module__ != module.__name__: + raise AssertionError( + f"subconfig class {value} must be defined in module {module.__name__}" + ) + # a subconfig with `class Blah:` syntax + proxy = SubConfigProxy(module, f"{name}.") + visit(value, proxy, f"{name}.") + if dest is module: + setattr(dest, key, proxy) + else: + dest.__dict__[key] = proxy + else: + raise AssertionError(f"Unhandled config {key}={value} ({type(value)})") + + config: dict[str, _ConfigEntry] = {} + + compile_ignored_keys = get_assignments_with_compile_ignored_comments(module) + + visit(module, module, "") + module._config = config # type: ignore[attr-defined] + module._compile_ignored_keys = compile_ignored_keys # type: ignore[attr-defined] + module.__class__ = ConfigModuleInstance + module._is_dirty = True # type: ignore[attr-defined] + module._hash_digest = None # type: ignore[attr-defined] + + +COMPILE_IGNORED_MARKER = "@compile_ignored" + + +# Gets all the keys (i.e. assignments) with a @compile_ignored comment +def get_assignments_with_compile_ignored_comments(module: ModuleType) -> set[str]: + source_code = inspect.getsource(module) + assignments = set() + + # Tokenize the source code to retrieve comments + tokens = tokenize.tokenize(io.BytesIO(source_code.encode("utf-8")).readline) + current_comment = "", -1 + prev_name = "" + + for token in tokens: + if token.type == tokenize.COMMENT: + prev_name = "" + maybe_current = token.string.strip() + if COMPILE_IGNORED_MARKER in maybe_current: + if current_comment != ("", -1): + raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}") + current_comment = maybe_current, token.start[0] + elif token.type == tokenize.NAME: + # Only accept the first name token, to handle if you have + # something like foo: Bar = ... + if not prev_name: + prev_name = token.string + elif token.type == tokenize.OP and token.string == "=": + # Check if the current assignment follows a comment + # with COMPILE_IGNORED_MARKER + if ( + COMPILE_IGNORED_MARKER in current_comment[0] + and current_comment[1] == token.start[0] - 1 + ): + assignments.add(prev_name) + current_comment = "", -1 # reset + prev_name = "" + if current_comment != ("", -1): + raise AssertionError(f"unconsumed {COMPILE_IGNORED_MARKER}") + return assignments + + +@dataclass +class _ConfigEntry: + # The default value specified in the configuration + default: Any + # The type of the configuration value + value_type: type + # The value specified by the user when they overrode the configuration + # _UNSET_SENTINEL indicates the value is not set. + user_override: ContextVar[object] + # The justknob to check for this config + justknob: str | None = None + # environment variables are read at install time + env_value_force: Any = _UNSET_SENTINEL + env_value_default: Any = _UNSET_SENTINEL + # Used to work arounds bad assumptions in unittest.mock.patch + # The code to blame is + # https://github.com/python/cpython/blob/94a7a4e22fb8f567090514785c69e65298acca42/Lib/unittest/mock.py#L1637 + # Essentially, mock.patch requires, that if __dict__ isn't accessible + # (which it isn't), that after delattr is called on the object, the + # object must throw when hasattr is called. Otherwise, it doesn't call + # setattr again. + # Technically we'll have an intermediate state of hiding the config while + # mock.patch is unpatching itself, but it calls setattr after the delete + # call so the final state is correct. It's just very unintuitive. + # upstream bug - python/cpython#126886 + hide: bool = False + alias: str | None = None + # Deprecation support + deprecated: bool = False + deprecation_message: str | None = None + _deprecation_warned: bool = False + + def __init__(self, config: _Config, name: str) -> None: + self.default = config.default + self.value_type = ( + config.value_type if config.value_type is not None else type(self.default) + ) + self.justknob = config.justknob + self.alias = config.alias + # Deprecation fields + self.deprecated = config.deprecated + self.deprecation_message = config.deprecation_message + self._deprecation_warned = False + + self.user_override = ContextVar(name, default=_UNSET_SENTINEL) + if config.env_name_default is not None: + for val in config.env_name_default: + if (env_value := _read_env_variable(val)) is not None: + self.env_value_default = env_value + break + if config.env_name_force is not None: + for val in config.env_name_force: + if (env_value := _read_env_variable(val)) is not None: + self.env_value_force = env_value + break + + # Ensure justknobs and envvars are allowlisted types + if self.justknob is not None and self.default is not None: + if not isinstance(self.default, bool): + raise AssertionError( + f"justknobs only support booleans, {self.default} is not a boolean" + ) + if self.value_type is not None and ( + config.env_name_default is not None or config.env_name_force is not None + ): + if self.value_type not in ( + bool, + str, + Optional[bool], # noqa: UP045 + Optional[str], # noqa: UP045 + ): + raise AssertionError( + f"envvar configs only support (optional) booleans or strings, {self.value_type} is neither" + ) + + +class ConfigModule(ModuleType): + # NOTE: This should be kept in sync with _config_typing.pyi. + + # The actual configuration settings. E.g., torch._dynamo.config.debug + # would live as "debug" in the key, and torch._inductor.config.triton.cudagraphs + # maps as "triton.cudagraphs". See discussion on the class for meaning of various sub items + _config: dict[str, _ConfigEntry] + _bypass_keys: set[str] + _compile_ignored_keys: set[str] + _is_dirty: bool + _hash_digest: bytes | None + + def __init__(self) -> None: + raise NotImplementedError( + f"use {__name__}.install_config_module(sys.modules[__name__])" + ) + + def _warn_if_deprecated(self, name: str, config: _ConfigEntry) -> None: + """Issue deprecation warning for config if not already warned.""" + if config.deprecated and not config._deprecation_warned: + import warnings + + msg = f"{self.__name__}.{name} is deprecated" + if config.deprecation_message: + msg += f" and {config.deprecation_message}" + msg += ". It will be removed in a future version of PyTorch." + warnings.warn(msg, FutureWarning, stacklevel=3) + config._deprecation_warned = True + + def __setattr__(self, name: str, value: object) -> None: + if name in self._bypass_keys: + super().__setattr__(name, value) + elif name not in self._config: + raise AttributeError(f"{self.__name__}.{name} does not exist") + else: + # Issue deprecation warning on write (once per config) + config = self._config[name] + self._warn_if_deprecated(name, config) + + if config.alias is not None: + self._set_alias_val(config, value) + else: + config.user_override.set(value) + self._is_dirty = True + config.hide = False + + def __getattr__(self, name: str) -> Any: + try: + config = self._config[name] + + if config.hide: + raise AttributeError(f"{self.__name__}.{name} does not exist") + + # Issue deprecation warning on read (once per config) + self._warn_if_deprecated(name, config) + + alias_val = self._get_alias_val(config) + if alias_val is not _UNSET_SENTINEL: + return alias_val + + if config.env_value_force is not _UNSET_SENTINEL: + return config.env_value_force + + user_override = config.user_override.get() + if user_override is not _UNSET_SENTINEL: + return user_override + + if config.env_value_default is not _UNSET_SENTINEL: + return config.env_value_default + + if config.justknob is not None: + # JK only supports bools and ints + return justknobs_check(name=config.justknob, default=config.default) + + # Reference types can still be modified, so copy them to + # user_overrides to prevent accidental mutation of defaults. + if not isinstance(config.default, _IMMUTABLE_CONFIG_TYPES): + config.user_override.set(copy.deepcopy(config.default)) + return config.user_override.get() + return config.default + + except KeyError as e: + # make hasattr() work properly + raise AttributeError(f"{self.__name__}.{name} does not exist") from e + + def __delattr__(self, name: str) -> None: + self._is_dirty = True + # must support delete because unittest.mock.patch deletes + # then recreate things + self._config[name].user_override.set(_UNSET_SENTINEL) + self._config[name].hide = True + + def _get_alias_module_and_name( + self, entry: _ConfigEntry + ) -> tuple[ModuleType, str] | None: + alias = entry.alias + if alias is None: + return None + module_name, constant_name = alias.rsplit(".", 1) + try: + module = importlib.import_module(module_name) + except ImportError as e: + raise AttributeError(f"config alias {alias} does not exist") from e + return module, constant_name + + def _get_alias_val(self, entry: _ConfigEntry) -> Any: + data = self._get_alias_module_and_name(entry) + if data is None: + return _UNSET_SENTINEL + module, constant_name = data + constant_value = getattr(module, constant_name) + return constant_value + + def _set_alias_val(self, entry: _ConfigEntry, val: Any) -> None: + data = self._get_alias_module_and_name(entry) + if data is None: + raise AssertionError( + "alias data should not be None when setting alias value" + ) + module, constant_name = data + setattr(module, constant_name, val) + + def _is_default(self, name: str) -> bool: + """ + Returns true if the config is at its default value. + configs overridden by the env are not considered default. + """ + config_val = self._config[name] + # The config is not overridden by the user, and the env_value_default + # is different from the default value (meaning user has set the env to + # change the default value). + not_set_env_default = ( + config_val.env_value_default is _UNSET_SENTINEL + or config_val.env_value_default == config_val.default + ) + not_set_env_force = ( + config_val.env_value_force is _UNSET_SENTINEL + or config_val.env_value_force == config_val.default + ) + + unset = config_val.user_override.get() is _UNSET_SENTINEL + # Handle reference types specially to avoid spammy warnings + if not isinstance(config_val.default, _IMMUTABLE_CONFIG_TYPES): + unset = unset or config_val.user_override.get() == config_val.default + return unset and not_set_env_default and not_set_env_force + + def _get_dict( + self, + ignored_keys: list[str] | None = None, + ignored_prefixes: list[str] | None = None, + skip_default: bool = False, + ) -> dict[str, Any]: + """Export a dictionary of current configuration keys and values. + + This function is design to provide a single point which handles + accessing config options and exporting them into a dictionary. + This is used by a number of different user facing export methods + which all have slightly different semantics re: how and what to + skip. + If a config is aliased, it skips this config. + + Arguments: + ignored_keys are keys that should not be exported. + ignored_prefixes are prefixes that if a key matches should + not be exported + skip_default does two things. One if a key has not been modified + it skips it. + """ + config: dict[str, Any] = {} + for key, entry in self._config.items(): + if entry.alias is not None: + continue + if ignored_keys and key in ignored_keys: + continue + if ignored_prefixes: + if any(key.startswith(prefix) for prefix in ignored_prefixes): + continue + if skip_default and self._is_default(key): + continue + + # Read value directly, bypassing __getattr__ overhead + # (deprecation warnings, alias resolution). + user_override = entry.user_override.get() + if entry.env_value_force is not _UNSET_SENTINEL: + val = entry.env_value_force + elif user_override is not _UNSET_SENTINEL: + val = user_override + elif entry.env_value_default is not _UNSET_SENTINEL: + val = entry.env_value_default + elif entry.justknob is not None: + val = justknobs_check(name=entry.justknob, default=entry.default) + else: + val = entry.default + if not isinstance(val, _IMMUTABLE_CONFIG_TYPES): + val = copy.deepcopy(val) + config[key] = val + + return config + + def get_type(self, config_name: str) -> type: + return self._config[config_name].value_type + + def save_config(self) -> bytes: + """Convert config to a pickled blob""" + ignored_keys = getattr(self, "_save_config_ignore", []) + return pickle.dumps( + self._get_dict(ignored_keys=ignored_keys), + protocol=2, + ) + + def save_config_portable( + self, *, ignore_private_configs: bool = True + ) -> dict[str, Any]: + """Convert config to portable format""" + prefixes = [] + if ignore_private_configs: + prefixes.append("_") + prefixes.extend(getattr(self, "_cache_config_ignore_prefix", [])) + config = self._get_dict(ignored_prefixes=prefixes) + for key in getattr(self, "_cache_config_factory_keys", []): + if key in config and config[key] is not None: + instance = config[key]() + if hasattr(instance, "uuid"): + config[key] = instance.uuid() + else: + raise RuntimeError( + f"Config '{key}' is set to {config[key]} which does not " + f"implement uuid(). Implement uuid() for cache key " + f"participation." + ) + return config + + def codegen_config(self) -> str: + """Convert config to Python statements that replicate current config. + This does NOT include config settings that are at default values. + """ + + # additional imports required + imports = set() + + def get_module_name(func: Callable, add_dot: bool) -> str: + module_name = func.__module__ + if module_name == "builtins": + module_name = "" + if add_dot and module_name != "": + module_name += "." + return module_name + + def add_import(func: Callable) -> None: + module_name = get_module_name(func, False) + if module_name: + imports.add(module_name) + + def list_of_callables_to_string(v: list | set) -> list[str]: + return [f"{get_module_name(item, True)}{item.__name__}" for item in v] + + def importable_callable(v: Any) -> bool: + # functools.partial has no attributes below but is a callable + return callable(v) and hasattr(v, "__module__") and hasattr(v, "__name__") + + def get_config_line(mod, k, v) -> str: # type: ignore[no-untyped-def] + """ + Return a string version of the config line. + Handle v when v is a callable, or a list/dict of callables. Add import statements for callables if necessary. + We assume that the value of a single config won't be a mix of callables and non-callables. + + Example output: + import logging + import _warnings + torch._dynamo.config.reorderable_logging_functions = { _warnings.warn, logging.warn, print } + """ + if importable_callable(v): + add_import(v) + return f"{mod}.{k} = {get_module_name(v, True)}{v.__name__}" + elif isinstance(v, (list, set)) and all( + importable_callable(item) for item in v + ): + for item in v: + add_import(item) + v_list = list_of_callables_to_string(v) + if isinstance(v, list): + return f"{mod}.{k} = {v_list}" + else: + return f"{mod}.{k} = {{ {', '.join(v_list)} }}" + else: + return f"{mod}.{k} = {v!r}" + + lines = [] + mod = self.__name__ + for k, v in self._get_dict( + ignored_keys=getattr(self, "_save_config_ignore", []), skip_default=True + ).items(): + lines.append(get_config_line(mod, k, v)) + for import_name in imports: + lines.insert(0, f"import {import_name}") + return "\n".join(lines) + + def get_hash(self) -> bytes: + """Hashes the configs that are not compile_ignored""" + if self._is_dirty or self._hash_digest is None: + dict_to_hash = self._get_dict(ignored_keys=list(self._compile_ignored_keys)) + string_to_hash = repr(sorted(dict_to_hash.items())) + self._hash_digest = hashlib.md5( + string_to_hash.encode("utf-8"), usedforsecurity=False + ).digest() + self._is_dirty = False + return self._hash_digest + + @deprecated( + "`config.to_dict()` has been deprecated. It no longer changes the underlying config." + " use `config.get_config_copy()` instead if you just want a copy of the config, or " + "config.load_config if you need mutable access", + category=FutureWarning, + ) + def to_dict(self) -> dict[str, Any]: + return self.get_config_copy() + + @deprecated( + "`config.shallow_copy_dict()` has been deprecated. It no longer changes the underlying config." + " use `config.get_config_copy()` instead if you just want a copy of the config, or " + "config.load_config if you need mutable access", + category=FutureWarning, + ) + def shallow_copy_dict(self) -> dict[str, Any]: + return self.get_config_copy() + + def load_config(self, maybe_pickled_config: bytes | dict[str, Any]) -> None: + """Restore from a prior call to save_config() or shallow_copy_dict()""" + if not isinstance(maybe_pickled_config, dict): + config = pickle.loads(maybe_pickled_config) + else: + config = maybe_pickled_config + for k, v in config.items(): + if k in self._config: + setattr(self, k, v) + else: + from torch._dynamo.utils import warn_once + + warn_once(f"key {k} with value {v} is not understood by this config") + + def get_config_copy(self) -> dict[str, Any]: + return self._get_dict() + + def get_serializable_config_copy(self) -> dict[str, Any]: + return self._get_dict(ignored_keys=getattr(self, "_save_config_ignore", [])) + + def patch( + self, + arg1: str | dict[str, Any] | None = None, + arg2: Any = None, + **kwargs: dict[str, Any], + ) -> "ContextDecorator": + """ + Decorator and/or context manager to make temporary changes to a config. Note that patched settings are thread-local. + + As a decorator: + + @config.patch("name", val) + @config.patch(name1=val1, name2=val2) + @config.patch({"name1": val1, "name2", val2}) + def foo(...): + ... + + As a context manager: + + with config.patch("name", val): + ... + """ + changes: dict[str, Any] + if arg1 is not None: + if arg2 is not None: + if not isinstance(arg1, str): + raise AssertionError( + "first argument must be a string when passing 2 positional args to patch" + ) + # patch("key", True) syntax + changes = {arg1: arg2} + else: + if not isinstance(arg1, dict): + raise AssertionError( + "first argument must be a dict when passing a single positional arg to patch" + ) + # patch({"key": True}) syntax + changes = arg1 + if kwargs: + raise AssertionError( + "cannot pass both positional and keyword arguments to patch" + ) + else: + # patch(key=True) syntax + changes = kwargs + if arg2 is not None: + raise AssertionError( + "second positional argument is only valid when first argument is a key string" + ) + if not isinstance(changes, dict): + raise AssertionError(f"expected `dict` got {type(changes)}") + config = self + + class ConfigPatch(ContextDecorator): + def __init__(self) -> None: + self.changes = changes + self._prior: ContextVar[tuple[dict[str, Any], ...]] = ContextVar( + f"{config.__name__}.ConfigPatch[{id(self)}]", + default=(), + ) + + def __enter__(self) -> None: + prior: dict[str, Any] = {} + for key in self.changes: + # KeyError on invalid entry + prior[key] = config.__getattr__(key) + prior_stack = self._prior.get() + self._prior.set((*prior_stack, prior)) + try: + for k, v in self.changes.items(): + config.__setattr__(k, v) + except Exception: + self._prior.set(prior_stack) + raise + + def __exit__(self, exc_type, exc_val, exc_tb): # type: ignore[no-untyped-def] + prior_stack = self._prior.get() + if not prior_stack: + raise AssertionError( + "prior should not be empty when exiting ConfigPatch" + ) + prior = prior_stack[-1] + self._prior.set(prior_stack[:-1]) + for k, v in prior.items(): + config.__setattr__(k, v) + + return ConfigPatch() + + def _make_closure_patcher(self, **changes: dict[str, Any]) -> Any: + """ + A lower-overhead version of patch() for things on the critical path. + + Usage: + + # do this off the critical path + change_fn = config.make_closure_patcher(foo=True) + + ... + + revert = change_fn() + try: + ... + finally: + revert() + + """ + config = self._config + + def change() -> Callable[[], None]: + prior = {k: config[k].user_override.get() for k in changes} + for k, v in changes.items(): + self._config[k].user_override.set(v) + + def revert() -> None: + for k, v in prior.items(): + self._config[k].user_override.set(v) + + return revert + + return change + + +class ContextDecorator(contextlib.ContextDecorator): + """ + Same as contextlib.ContextDecorator, but with support for + `unittest.TestCase` + """ + + def __enter__(self) -> None: + raise NotImplementedError("NYI") + + def __exit__(self, exc_type, exc_val, exc_tb) -> NoReturn: # type: ignore[no-untyped-def] + raise NotImplementedError("NYI") + + def __call__(self, func: Callable[[Any], Any]) -> Any: + if isinstance(func, type) and issubclass(func, unittest.TestCase): + + class _TestCase(func): # type: ignore[valid-type, misc] + @classmethod + def setUpClass(cls) -> None: + self.__enter__() + try: + super().setUpClass() + except Exception: + self.__exit__(None, None, None) + raise + + @classmethod + def tearDownClass(cls) -> None: + try: + super().tearDownClass() + finally: + self.__exit__(None, None, None) + + _TestCase.__name__ = func.__name__ + _TestCase.__qualname__ = func.__qualname__ + _TestCase.__module__ = func.__module__ + + return _TestCase + + return super().__call__(func) + + +class SubConfigProxy: + """ + Shim to redirect to main config. + `config.triton.cudagraphs` maps to _config["triton.cudagraphs"] + """ + + def __init__(self, config: object, prefix: str) -> None: + # `super().__setattr__` to bypass custom `__setattr__` + super().__setattr__("_config", config) + super().__setattr__("_prefix", prefix) + + def __setattr__(self, name: str, value: object) -> None: + return self._config.__setattr__(self._prefix + name, value) + + def __getattr__(self, name: str) -> Any: + return self._config.__getattr__(self._prefix + name) + + def __delattr__(self, name: str) -> None: + return self._config.__delattr__(self._prefix + name) + + +def get_tristate_env(name: str, default: Any = None) -> bool | None: + value = os.environ.get(name) + if value == "1": + return True + if value == "0": + return False + return default + + +def inherit_fields_from(parent_cls): + def wrapper(child_cls): + for k, v in parent_cls.__dict__.items(): + # copy fields that are not private and not overridden + if not k.startswith("_") and k not in child_cls.__dict__: + setattr(child_cls, k, v) + return child_cls + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_config_typing.pyi b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_config_typing.pyi new file mode 100644 index 0000000000000000000000000000000000000000..9cae7368cfa5ee74392cde78942ae1634f5b868e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_config_typing.pyi @@ -0,0 +1,36 @@ +# mypy: allow-untyped-defs +from typing import Any, TYPE_CHECKING + +""" +This was semi-automatically generated by running + + stubgen torch.utils._config_module.py + +And then manually extracting the methods of ConfigModule and converting them into top-level functions. + +This file should be imported into any file that uses install_config_module like so: + + if TYPE_CHECKING: + from torch.utils._config_typing import * # noqa: F401, F403 + + from torch.utils._config_module import install_config_module + + # adds patch, save_config, etc + install_config_module(sys.modules[__name__]) + +Note that the import should happen before the call to install_config_module(), otherwise runtime errors may occur. +""" + +if not TYPE_CHECKING: # noqa: PYI002 + raise AssertionError("Do not use at runtime") # noqa: W291 + +def save_config() -> bytes: ... +def save_config_portable(*, ignore_private_configs: bool = True) -> dict[str, Any]: ... +def codegen_config() -> str: ... +def get_hash() -> bytes: ... +def to_dict() -> dict[str, Any]: ... +def shallow_copy_dict() -> dict[str, Any]: ... +def load_config(config: bytes | dict[str, Any]) -> None: ... +def get_config_copy() -> dict[str, Any]: ... +def get_serializable_config_copy() -> dict[str, Any]: ... +def patch(arg1: str | dict[str, Any] | None = None, arg2: Any = None, **kwargs): ... diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_content_store.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_content_store.py new file mode 100644 index 0000000000000000000000000000000000000000..234355210057ab703b9f0aea3729757787f89eb6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_content_store.py @@ -0,0 +1,241 @@ +# mypy: allow-untyped-defs +# This module provides a FAST (on GPU) content addressable store for storages +# (and tensors on top of them) with VERY WEAK portability guarantees (e.g., +# don't expect CPU/CUDA to address to the same hash, don't expect it to be +# portable across devices) that is NOT cryptographically secure. In return, +# we are able to hash 40G of tensor data on GPU in less than a second, +# compared to running SHA-1 in CPU which would a minute or so. The primary +# use case is for efficiently snapshotting intermediate tensor data for +# offline debugging, but it's been put in this module in case you think of +# another use case for it. The hash function could be replaced with a +# straight reimplementation of SHA-1, which would give us much stronger +# portability guarantees. +# +# WARNING: THERE IS NO BC/FC GUARANTEE FOR THIS FORMAT! If you need to format +# shift the result, consider packing it into a single torch.save object +# with traditional view sharing. +# +# Because of the weak portability guarantees, you can only write to the +# content store from a single process; we don't provide any capability +# of "reopening" a content store to add more things to it. But we don't +# assume that you can keep all of the tensors you want to add to the store +# in memory at once, because you probably can't! Nor do we assume that +# you know a priori whether or not two storages can be deduplicated or not. +# +# Note: only storages are content-addressed; tensors are name addressed +# +# Note: our padding strategy means that [1, 0] and [1] int16 tensors would +# map to the same (padded) storage. We think this will be immaterial for most +# users. + +import ctypes +import functools +import hashlib +import os.path +import struct +from collections import defaultdict + +import torch +import torch._prims as prims +import torch._utils +import torch.nn.functional as F +from torch.multiprocessing.reductions import StorageWeakRef + + +def lazy_compile(**compile_kwargs): + """Lazily wrap a function with torch.compile on the first call + + This avoids eagerly importing dynamo. + """ + + def decorate_fn(fn): + @functools.wraps(fn) + def compile_hook(*args, **kwargs): + compiled_fn = torch.compile(fn, **compile_kwargs) + globals()[fn.__name__] = functools.wraps(fn)(compiled_fn) + return compiled_fn(*args, **kwargs) + + return compile_hook + + return decorate_fn + + +# Use of torch.compile is mandatory for (1) good memory usage +# and (2) xor_sum implementation. This is our first instance of +# using PT2 to implement a kernel in PyTorch; if we get AOT capabilities +# it would be good to apply it here. +@lazy_compile(dynamic=True) +def hash_storage_kernel(x): + # The randint calls are carefully written to hit things we + # have lowerings for in inductor. Lack of unsigned 32-bit integer + # is a pain. + a = torch.randint( + -(2**31), 2**31, x.shape, device=x.device, dtype=torch.int32 + ).abs() + a = ((a % (2**31 - 1)) + 1).long() + b = ( + torch.randint(-(2**31), 2**31, x.shape, device=x.device, dtype=torch.int32) + .abs() + .long() + ) + # This is a standard shift-multiply universal hash family + # plus xor sum hash, using Philox to generate random numbers. + # Our Philox RNG is not deterministic across devices so + # don't use this for stable hashing. + # + # This assumes fixed length so you're also obligated to bucket + # by the length of tensor as well + return prims.xor_sum((a * x + b).int(), [0]) + + +# Returns a hex digest of the data in the storage. Guaranteed to be +# SHA-1 if stable_hash=True, otherwise it will consistent for a single +# process run but not necessarily across processes. +def hash_storage(storage: torch.UntypedStorage, *, stable_hash: bool = False) -> str: + import torch._dynamo + from torch._dynamo.utils import is_compile_supported + + device_type = storage.device.type + if stable_hash or not is_compile_supported(device_type): + cpu_storage = storage.cpu() + # TODO: make storage support buffer protocol so this isn't + # necessary + buf = (ctypes.c_byte * cpu_storage.nbytes()).from_address( + cpu_storage.data_ptr() + ) + sha1 = hashlib.sha1(usedforsecurity=False) + sha1.update(buf) + return sha1.hexdigest() + + # TODO: factor this into a random utility + if device_type == "cpu": + generator = torch._C.default_generator + elif device_type == "cuda": + generator = torch.cuda.default_generators[storage.device.index] + elif device_type == "mps": + generator = torch.mps._get_default_mps_generator() + elif device_type == "xpu": + generator = torch.xpu.default_generators[storage.device.index] + else: + raise AssertionError(f"unhandled device type {device_type}") + state = generator.get_state() + try: + generator.manual_seed(0) + x = torch.empty(0, dtype=torch.uint8, device=storage.device).set_(storage) # type: ignore[call-overload] + # The dtype-casting view cannot be compiled, and so the + # padding/reshaping also needs to be done externally even + # though it could be profitably fused + pad = -x.numel() % 4 + if pad > 0: + x = F.pad(x, (0, pad), "constant", 0) + x = x.view(torch.int32) + # We run the 32-bit hash five times with differing parameters to + # reduce chance of collision + ITER = 5 + cs = [hash_storage_kernel(x).item() for _ in range(ITER)] + return struct.pack(">" + "i" * ITER, *cs).hex() + finally: + generator.set_state(state) + + +class ContentStoreWriter: + # Structure: + # storages/ + # 00/ + # 0000..00 + # tensors/ + # name + def __init__(self, loc: str, stable_hash: bool = False) -> None: + self.loc: str = loc + self.seen_storage_hashes: set[str] = set() + self.stable_hash = stable_hash + + # TODO: offer some sort of non-blocking API to speed things up + def write_storage(self, storage: torch.UntypedStorage) -> str: + h = hash_storage(storage, stable_hash=self.stable_hash) + if h in self.seen_storage_hashes: + return h + # TODO: consider not using torch.save for this; we don't actually + # need any metadata for the storage + subfolder = os.path.join(self.loc, "storages") + os.makedirs(subfolder, exist_ok=True) + target = os.path.join(subfolder, h) + if os.path.exists(target): + return h + torch.save(storage, target) + self.seen_storage_hashes.add(h) + return h + + def compute_tensor_metadata(self, t: torch.Tensor, h=None): + if h is None: + h = hash_storage(t.untyped_storage(), stable_hash=self.stable_hash) + return ( + t.dtype, + h, + t.storage_offset(), + tuple(t.shape), + t.stride(), + torch._utils.get_tensor_metadata(t), + ) + + def write_tensor(self, name: str, t: torch.Tensor) -> None: + storage = t.untyped_storage() + h = self.write_storage(storage) + # TODO: Support more advanced snapshotting of requires_grad/grad/etc + d, f = os.path.split(name) + payload = self.compute_tensor_metadata(t, h=h) + subfolder = os.path.join(self.loc, "tensors", d) + os.makedirs(subfolder, exist_ok=True) + torch.save(payload, os.path.join(subfolder, f)) + + +class ContentStoreReader: + def __init__(self, loc: str, *, cache=True) -> None: + self.loc = loc + self.storage_cache: ( + dict[torch.device | None, dict[str, StorageWeakRef]] | None + ) = None + if cache: + self.storage_cache = defaultdict(dict) + + def read_storage(self, h: str, *, device=None) -> torch.UntypedStorage: + if device is not None: + device = torch.device(device) + ws = ( + self.storage_cache[device].get(h) + if self.storage_cache is not None + else None + ) + s: torch.UntypedStorage | None + if ws is not None: + s = torch.UntypedStorage._new_with_weak_ptr(ws.cdata) + if s is not None: + return s + s = torch.load( + os.path.join(self.loc, "storages", h), + weights_only=True, + map_location=device, + )._untyped_storage + if s is None: + raise AssertionError( + f"expected storage for hash {h} in {os.path.join(self.loc, 'storages')}, got None" + ) + if self.storage_cache is not None: + self.storage_cache[device][h] = StorageWeakRef(s) + return s + + def read_tensor_metadata(self, name: str): + fn = os.path.join(self.loc, "tensors", name) + if not os.path.exists(fn): + raise FileNotFoundError(fn) + return torch.load(fn, weights_only=True) + + def read_tensor(self, name: str, *, device=None) -> torch.Tensor: + dtype, h, storage_offset, size, stride, metadata = self.read_tensor_metadata( + name + ) + storage = self.read_storage(h, device=device) + t = torch.tensor([], dtype=dtype, device=storage.device) + t.set_(storage, storage_offset, size, stride) + torch._utils.set_tensor_metadata(t, metadata) + return t diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_contextlib.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_contextlib.py new file mode 100644 index 0000000000000000000000000000000000000000..408cdfe7d7b778c776ac88a050c3fe8c6a6edbb5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_contextlib.py @@ -0,0 +1,170 @@ +# mypy: allow-untyped-defs +# Extra utilities for working with context managers that should have been +# in the standard library but are not + +import functools +import inspect +import sys +import warnings +from collections.abc import Callable +from typing import Any, cast, overload, TypeVar +from typing_extensions import Self + + +# Used for annotating the decorator usage of _DecoratorContextManager (e.g., +# 'no_grad' and 'enable_grad'). +# See https://mypy.readthedocs.io/en/latest/generics.html#declaring-decorators +FuncType = Callable[..., Any] +F = TypeVar("F", bound=FuncType) + + +def _wrap_generator(ctx_factory, func): + """ + Wrap each generator invocation with the context manager factory. + + The input should be a function that returns a context manager, + not a context manager itself, to handle one-shot context managers. + """ + + @functools.wraps(func) + def generator_context(*args, **kwargs): + gen = func(*args, **kwargs) + + # Generators are suspended and unsuspended at `yield`, hence we + # make sure the grad mode is properly set every time the execution + # flow returns into the wrapped generator and restored when it + # returns through our `yield` to our caller (see PR #49017). + try: + # Issuing `None` to a generator fires it up + with ctx_factory(): + response = gen.send(None) + + while True: + try: + # Forward the response to our caller and get its next request + request = yield response + + except GeneratorExit: + # Inform the still active generator about its imminent closure + with ctx_factory(): + gen.close() + raise + + except BaseException: # noqa: B036 + # Propagate the exception thrown at us by the caller + with ctx_factory(): + response = gen.throw(*sys.exc_info()) + + else: + # Pass the last request to the generator and get its response + with ctx_factory(): + response = gen.send(request) + + # We let the exceptions raised above by the generator's `.throw` or + # `.send` methods bubble up to our caller, except for StopIteration + except StopIteration as e: + # The generator informed us that it is done: take whatever its + # returned value (if any) was and indicate that we're done too + # by returning it (see docs for python's return-statement). + return e.value + + return generator_context + + +def context_decorator(ctx, func): + """ + Like contextlib.ContextDecorator. + + But with the following differences: + 1. Is done by wrapping, rather than inheritance, so it works with context + managers that are implemented from C and thus cannot easily inherit from + Python classes + 2. Wraps generators in the intuitive way (c.f. https://bugs.python.org/issue37743) + 3. Errors out if you try to wrap a class, because it is ambiguous whether + or not you intended to wrap only the constructor + + The input argument can either be a context manager (in which case it must + be a multi-shot context manager that can be directly invoked multiple times) + or a callable that produces a context manager. + """ + if callable(ctx) and hasattr(ctx, "__enter__"): + raise AssertionError( + f"Passed in {ctx} is both callable and also a valid context manager " + "(has __enter__), making it ambiguous which interface to use. If you " + "intended to pass a context manager factory, rewrite your call as " + "context_decorator(lambda: ctx()); if you intended to pass a context " + "manager directly, rewrite your call as context_decorator(lambda: ctx)" + ) + + if not callable(ctx): + + def ctx_factory(): + return ctx + + else: + ctx_factory = ctx + + if inspect.isclass(func): + raise RuntimeError( + "Cannot decorate classes; it is ambiguous whether or not only the " + "constructor or all methods should have the context manager applied; " + "additionally, decorating a class at definition-site will prevent " + "use of the identifier as a conventional type. " + "To specify which methods to decorate, decorate each of them " + "individually." + ) + + if inspect.isgeneratorfunction(func): + return _wrap_generator(ctx_factory, func) + + @functools.wraps(func) + def decorate_context(*args, **kwargs): + # pyrefly: ignore [bad-context-manager] + with ctx_factory(): + return func(*args, **kwargs) + + return decorate_context + + +class _DecoratorContextManager: + """Allow a context manager to be used as a decorator.""" + + def __call__(self, orig_func: F) -> F: + if inspect.isclass(orig_func): + warnings.warn( + "Decorating classes is deprecated and will be disabled in " + "future versions. You should only decorate functions or methods. " + "To preserve the current behavior of class decoration, you can " + "directly decorate the `__init__` method and nothing else.", + FutureWarning, + stacklevel=2, + ) + func = cast(F, lambda *args, **kwargs: orig_func(*args, **kwargs)) + else: + func = orig_func + + return cast(F, context_decorator(self.clone, func)) + + def __enter__(self) -> None: + raise NotImplementedError + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + raise NotImplementedError + + def clone(self): + # override this method if your children class takes __init__ parameters + return self.__class__() + + +class _NoParamDecoratorContextManager(_DecoratorContextManager): + """Allow a context manager to be used as a decorator without parentheses.""" + + @overload + def __new__(cls, orig_func: F) -> F: ... # type: ignore[misc] + @overload + def __new__(cls, orig_func: None = None) -> Self: ... + + def __new__(cls, orig_func: F | None = None) -> Self | F: # type: ignore[misc] + if orig_func is None: + return super().__new__(cls) + return cls()(orig_func) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_cpp_embed_headers.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_cpp_embed_headers.py new file mode 100644 index 0000000000000000000000000000000000000000..88ab41aadffe75dd1bdbe0b3482b414f38b7e093 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_cpp_embed_headers.py @@ -0,0 +1,57 @@ +from collections.abc import Sequence +from pathlib import Path +from re import match as _match + + +def read_file(fname: Path | str) -> list[str]: + with open(fname, encoding="utf-8") as f: + return f.readlines() + + +def _embed_headers( + content: list[str], include_dirs: list[Path], processed_files: set[str] +) -> str: + for line_idx, cur_line in enumerate(content): + # Eliminate warning: `#pragma once in main file` + if cur_line.startswith("#pragma once"): + content[line_idx] = "" + continue + m = _match('^\\s*#include\\s*[<"]([^>"]+)[>"]', cur_line) + if m is None: + continue + for include_dir in include_dirs: + path = include_dir / m[1] + if not path.exists(): + continue + if str(path) in processed_files: + content[line_idx] = "" + continue + processed_files.add(str(path)) + content[line_idx] = _embed_headers( + read_file(path), include_dirs, processed_files + ) + break + return "".join(content) + + +def embed_headers( + fname: str, include_dirs: Sequence[str] | Sequence[Path] | str | None = None +) -> str: + if include_dirs is None: + base_dir = Path(__file__).parent.parent.parent + include_dirs = [base_dir, base_dir / "aten" / "src"] + elif isinstance(include_dirs, str): + include_dirs = [Path(include_dirs)] + else: + include_dirs = [Path(x) for x in include_dirs] + + return _embed_headers(read_file(fname), include_dirs, {fname}) + + +if __name__ == "__main__": + import sys + + if len(sys.argv) < 2: + print(f"Usage:\n {sys.argv[0]} filename") + sys.exit(1) + print(embed_headers(sys.argv[1])) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_cpp_extension_versioner.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_cpp_extension_versioner.py new file mode 100644 index 0000000000000000000000000000000000000000..d1391dd9aaab0e920276c7046282ea1002346a15 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_cpp_extension_versioner.py @@ -0,0 +1,63 @@ +# mypy: allow-untyped-defs +import collections + + +Entry = collections.namedtuple("Entry", "version, hash") + + +def update_hash(seed, value): + # Good old boost::hash_combine + # https://www.boost.org/doc/libs/1_35_0/doc/html/boost/hash_combine_id241013.html + return seed ^ (hash(value) + 0x9E3779B9 + (seed << 6) + (seed >> 2)) + + +def hash_source_files(hash_value, source_files): + for filename in source_files: + with open(filename, "rb") as file: + hash_value = update_hash(hash_value, file.read()) + return hash_value + + +def hash_build_arguments(hash_value, build_arguments): + for group in build_arguments: + if group: + for argument in group: + hash_value = update_hash(hash_value, argument) + return hash_value + + +class ExtensionVersioner: + def __init__(self) -> None: + self.entries = {} + + def get_version(self, name): + entry = self.entries.get(name) + return None if entry is None else entry.version + + def bump_version_if_changed( + self, + name, + source_files, + build_arguments, + build_directory, + with_cuda, + with_sycl, + is_python_module, + is_standalone, + ): + hash_value = 0 + hash_value = hash_source_files(hash_value, source_files) + hash_value = hash_build_arguments(hash_value, build_arguments) + hash_value = update_hash(hash_value, build_directory) + hash_value = update_hash(hash_value, with_cuda) + hash_value = update_hash(hash_value, with_sycl) + hash_value = update_hash(hash_value, is_python_module) + hash_value = update_hash(hash_value, is_standalone) + + entry = self.entries.get(name) + if entry is None: + self.entries[name] = entry = Entry(0, hash_value) + elif hash_value != entry.hash: + self.entries[name] = entry = Entry(entry.version + 1, hash_value) + + return entry.version diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_cxx_pytree.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_cxx_pytree.py new file mode 100644 index 0000000000000000000000000000000000000000..d13cc7eb156f1f414d7ce9c9f82913c14f8780d1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_cxx_pytree.py @@ -0,0 +1,1151 @@ +# Owner(s): ["module: pytree"] + +""" +Contains utility functions for working with nested python data structures. + +A *pytree* is a Python nested data structure. It is a tree in the sense that +nodes are Python collections (e.g., list, tuple, dict) and the leaves are +Python values. Furthermore, a pytree should not contain reference cycles. + +pytrees are useful for working with nested collections of Tensors. For example, +one can use `tree_map` to map a function over all Tensors inside some nested +collection of Tensors and `tree_leaves` to get a flat list of all Tensors +inside some nested collection. pytrees are helpful for implementing nested +collection support for PyTorch APIs. +""" + +import functools +import sys +import types +from collections.abc import Callable, Iterable, Mapping +from typing import Any, overload, TypeAlias, TypeVar +from typing_extensions import deprecated, Self, TypeIs + +import torch.utils._pytree as python_pytree +from torch.torch_version import TorchVersion as _TorchVersion +from torch.utils._pytree import ( + Context, + DumpableContext, + FlattenFn, + FlattenWithKeysFn, + FromDumpableContextFn, + is_namedtuple, + is_namedtuple_class, + is_namedtuple_instance, + is_structseq, + is_structseq_class, + is_structseq_instance, + KeyPath, + PyTree, + ToDumpableContextFn, + UnflattenFn, +) + + +# Do not try to import `optree` package if the static version check already fails. +if not python_pytree._cxx_pytree_dynamo_traceable: + raise ImportError( + f"{__name__} depends on `optree>={python_pytree._optree_minimum_version}`, " + "which is an optional dependency of PyTorch. " + "To use it, please upgrade your optree package via " + "`python3 -m pip install --upgrade optree`" + ) + + +import optree +from optree import PyTreeSpec # direct import for type annotations + + +__all__ = [ + "PyTree", + "Context", + "FlattenFn", + "UnflattenFn", + "DumpableContext", + "ToDumpableContextFn", + "FromDumpableContextFn", + "PyTreeSpec", + "TreeSpec", + "LeafSpec", + "keystr", + "key_get", + "register_pytree_node", + "tree_is_leaf", + "tree_flatten", + "tree_flatten_with_path", + "tree_unflatten", + "tree_iter", + "tree_leaves", + "tree_leaves_with_path", + "tree_structure", + "tree_map", + "tree_map_with_path", + "tree_map_", + "tree_map_only", + "tree_map_only_", + "tree_all", + "tree_any", + "tree_all_only", + "tree_any_only", + "treespec_dumps", + "treespec_loads", + "treespec_pprint", + "is_namedtuple", + "is_namedtuple_class", + "is_namedtuple_instance", + "is_structseq", + "is_structseq_class", + "is_structseq_instance", +] + + +# In-tree installation may have VCS-based versioning. Update the previous static version. +python_pytree._optree_version = _TorchVersion(optree.__version__) # type: ignore[attr-defined] + +__TORCH_DICT_SESSION = optree.dict_insertion_ordered(True, namespace="torch") +__TORCH_DICT_SESSION.__enter__() # enable globally and permanently + + +T = TypeVar("T") +S = TypeVar("S") +U = TypeVar("U") +R = TypeVar("R") + +TreeSpec: TypeAlias = PyTreeSpec +OpTreeUnflattenFn = Callable[[Context, Iterable[Any]], PyTree] + +# Keep deprecated alias for backward compatibility +FlattenFunc = FlattenFn # deprecated +UnflattenFunc = UnflattenFn # deprecated +FlattenWithKeysFunc = FlattenWithKeysFn # deprecated + + +def _reverse_args(func: UnflattenFn) -> OpTreeUnflattenFn: + @functools.wraps(func) + def wrapped(*args: Any, **kwargs: Any) -> Any: + return func(*reversed(args), **kwargs) + + return wrapped + + +def register_pytree_node( + cls: type[Any], + flatten_fn: FlattenFn, + unflatten_fn: UnflattenFn, + *, + serialized_type_name: str | None = None, + to_dumpable_context: ToDumpableContextFn | None = None, + from_dumpable_context: FromDumpableContextFn | None = None, + flatten_with_keys_fn: FlattenWithKeysFn | None = None, +) -> None: + """Register a container-like type as pytree node. + + Args: + cls (type): A Python type to treat as an internal pytree node. + flatten_fn (callable): A function to be used during flattening, taking an instance of + ``cls`` and returning a pair, with (1) an iterable for the children to be flattened + recursively, and (2) some hashable auxiliary data to be stored in the treespec and to be + passed to the ``unflatten_fn``. + unflatten_fn (callable): A function taking two arguments: the auxiliary data that was + returned by ``flatten_fn`` and stored in the treespec, and the unflattened children. + The function should return an instance of ``cls``. + serialized_type_name (str, optional): A keyword argument used to specify the fully + qualified name used when serializing the tree spec. + to_dumpable_context (callable, optional): An optional keyword argument to custom specify how + to convert the context of the pytree to a custom json dumpable representation. This is + used for json serialization, which is being used in :mod:`torch.export` right now. + from_dumpable_context (callable, optional): An optional keyword argument to custom specify + how to convert the custom json dumpable representation of the context back to the + original context. This is used for json deserialization, which is being used in + :mod:`torch.export` right now. + + Example:: + + >>> # xdoctest: +SKIP + >>> # Register a Python type with lambda functions + >>> register_pytree_node( + ... set, + ... lambda s: (sorted(s), None, None), + ... lambda children, _: set(children), + ... ) + """ + if flatten_with_keys_fn is not None: + raise NotImplementedError("KeyPaths are not yet supported in cxx_pytree.") + + _private_register_pytree_node( + cls, + flatten_fn, + unflatten_fn, + serialized_type_name=serialized_type_name, + to_dumpable_context=to_dumpable_context, + from_dumpable_context=from_dumpable_context, + ) + + python_pytree._private_register_pytree_node( + cls, + flatten_fn, + unflatten_fn, + serialized_type_name=serialized_type_name, + to_dumpable_context=to_dumpable_context, + from_dumpable_context=from_dumpable_context, + ) + + +@deprecated( + "`torch.utils._cxx_pytree._register_pytree_node` is deprecated. " + "Please use `torch.utils._cxx_pytree.register_pytree_node` instead.", + category=FutureWarning, +) +def _register_pytree_node( + cls: type[Any], + flatten_fn: FlattenFn, + unflatten_fn: UnflattenFn, + *, + serialized_type_name: str | None = None, + to_dumpable_context: ToDumpableContextFn | None = None, + from_dumpable_context: FromDumpableContextFn | None = None, +) -> None: + """Register a container-like type as pytree node for the C++ pytree only. + + The ``namespace`` argument is used to avoid collisions that occur when different libraries + register the same Python type with different behaviors. It is recommended to add a unique prefix + to the namespace to avoid conflicts with other libraries. Namespaces can also be used to specify + the same class in different namespaces for different use cases. + + .. warning:: + For safety reasons, a ``namespace`` must be specified while registering a custom type. It is + used to isolate the behavior of flattening and unflattening a pytree node type. This is to + prevent accidental collisions between different libraries that may register the same type. + + Args: + cls (type): A Python type to treat as an internal pytree node. + flatten_fn (callable): A function to be used during flattening, taking an instance of + ``cls`` and returning a pair, with (1) an iterable for the children to be flattened + recursively, and (2) some hashable auxiliary data to be stored in the treespec and to be + passed to the ``unflatten_fn``. + unflatten_fn (callable): A function taking two arguments: the auxiliary data that was + returned by ``flatten_fn`` and stored in the treespec, and the unflattened children. + The function should return an instance of ``cls``. + serialized_type_name (str, optional): A keyword argument used to specify the fully + qualified name used when serializing the tree spec. + to_dumpable_context (callable, optional): An optional keyword argument to custom specify how + to convert the context of the pytree to a custom json dumpable representation. This is + used for json serialization, which is being used in :mod:`torch.export` right now. + from_dumpable_context (callable, optional): An optional keyword argument to custom specify + how to convert the custom json dumpable representation of the context back to the + original context. This is used for json deserialization, which is being used in + :mod:`torch.export` right now. + """ + + _private_register_pytree_node( + cls, + flatten_fn, + unflatten_fn, + serialized_type_name=serialized_type_name, + to_dumpable_context=to_dumpable_context, + from_dumpable_context=from_dumpable_context, + ) + + +def _private_register_pytree_node( + cls: type[Any], + flatten_fn: FlattenFn, + unflatten_fn: UnflattenFn, + *, + serialized_type_name: str | None = None, + to_dumpable_context: ToDumpableContextFn | None = None, + from_dumpable_context: FromDumpableContextFn | None = None, +) -> None: + """This is an internal function that is used to register a pytree node type + for the C++ pytree only. End-users should use :func:`register_pytree_node` + instead. + """ + # TODO(XuehaiPan): remove this condition when we make Python pytree out-of-box support + # PyStructSequence types + if not optree.is_structseq_class(cls): + optree.register_pytree_node( + cls, + flatten_fn, + _reverse_args(unflatten_fn), + namespace="torch", + ) + + +def _is_pytreespec_instance( + obj: Any, + /, +) -> TypeIs[TreeSpec | python_pytree.PyTreeSpec]: + if isinstance(obj, (TreeSpec, python_pytree.PyTreeSpec)): + return True + if "torch._dynamo.polyfills.pytree" in sys.modules: + # The PyTorch Dynamo pytree module is not always available, so we check if it is loaded. + # If the PyTorch Dynamo pytree module is loaded, we can check if the treespec + # is an instance of the PyTorch Dynamo TreeSpec class. + import torch._dynamo.polyfills.pytree as dynamo_pytree + + return isinstance(obj, dynamo_pytree.PyTreeSpec) + return False + + +def treespec_leaf() -> TreeSpec: + """Make a treespec representing a leaf node.""" + return optree.treespec_leaf(none_is_leaf=True, namespace="torch") + + +def treespec_tuple(iterable: Iterable[TreeSpec] = (), /) -> TreeSpec: + """Make a tuple treespec from an iterable of child treespecs.""" + return optree.treespec_tuple(iterable, none_is_leaf=True, namespace="torch") + + +def treespec_dict( + mapping: Mapping[Any, TreeSpec] | Iterable[tuple[Any, TreeSpec]] = (), + /, + **kwargs: TreeSpec, +) -> TreeSpec: + """Make a dict treespec from a dict of child treespecs.""" + return optree.treespec_dict( + mapping, + **kwargs, + none_is_leaf=True, + namespace="torch", + ) + + +def tree_is_leaf( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + """Check if a pytree is a leaf. + + >>> tree_is_leaf(1) + True + >>> tree_is_leaf(None) + True + >>> tree_is_leaf([1, 2, 3]) + False + >>> tree_is_leaf((1, 2, 3), is_leaf=lambda x: isinstance(x, tuple)) + True + >>> tree_is_leaf({"a": 1, "b": 2, "c": 3}) + False + >>> tree_is_leaf({"a": 1, "b": 2, "c": None}) + False + + Args: + tree (pytree): A pytree to check if it is a leaf node. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + A boolean indicating if the pytree is a leaf node. + """ + return optree.tree_is_leaf( + tree, + is_leaf=is_leaf, + none_is_leaf=True, + namespace="torch", + ) + + +def tree_flatten( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> tuple[list[Any], TreeSpec]: + """Flatten a pytree. + + See also :func:`tree_unflatten`. + + The flattening order (i.e., the order of elements in the output list) is deterministic, + corresponding to a left-to-right depth-first tree traversal. + + >>> tree = {"b": (2, [3, 4]), "a": 1, "c": None, "d": 5} + >>> tree_flatten(tree) + ([2, 3, 4, 1, None, 5], PyTreeSpec({'b': (*, [*, *]), 'a': *, 'c': *, 'd': *}, NoneIsLeaf, namespace='torch')) + >>> tree_flatten(1) + ([1], PyTreeSpec(*, NoneIsLeaf, namespace='torch')) + >>> tree_flatten(None) + ([None], PyTreeSpec(*, NoneIsLeaf, namespace='torch')) + >>> from collections import OrderedDict + >>> tree = OrderedDict([("b", (2, [3, 4])), ("a", 1), ("c", None), ("d", 5)]) + >>> tree_flatten(tree) + ([2, 3, 4, 1, None, 5], PyTreeSpec(OrderedDict({'b': (*, [*, *]), 'a': *, 'c': *, 'd': *}), NoneIsLeaf, namespace='torch')) + + Args: + tree (pytree): A pytree to flatten. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + A pair ``(leaves, treespec)`` where the first element is a list of leaf values and the + second element is a treespec representing the structure of the pytree. + """ + return optree.tree_flatten( # type: ignore[return-value] + tree, + is_leaf=is_leaf, + none_is_leaf=True, + namespace="torch", + ) + + +def tree_unflatten(leaves: Iterable[Any], treespec: TreeSpec) -> PyTree: + """Reconstruct a pytree from the treespec and the leaves. + + The inverse of :func:`tree_flatten`. + + >>> tree = {"b": (2, [3, 4]), "a": 1, "c": None, "d": 5} + >>> leaves, treespec = tree_flatten(tree) + >>> tree == tree_unflatten(leaves, treespec) + True + + Args: + leaves (iterable): The list of leaves to use for reconstruction. The list must match the + number of leaves of the treespec. + treespec (TreeSpec): The treespec to reconstruct. + + Returns: + The reconstructed pytree, containing the ``leaves`` placed in the structure described by + ``treespec``. + """ + if not _is_pytreespec_instance(treespec): + if not _is_pytreespec_instance(leaves): + raise TypeError( + f"Expected `treespec` to be an instance of " + f"PyTreeSpec but got item of type {type(treespec)}." + ) + # Allow passing the PyTreeSpec instance as the first argument + # pyrefly: ignore [bad-assignment] + leaves, treespec = treespec, leaves + return treespec.unflatten(leaves) + + +def tree_iter( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> Iterable[Any]: + """Get an iterator over the leaves of a pytree. + + See also :func:`tree_flatten`. + + >>> tree = {"b": (2, [3, 4]), "a": 1, "c": None, "d": 5} + >>> list(tree_iter(tree)) + [2, 3, 4, 1, None, 5] + >>> list(tree_iter(1)) + [1] + >>> list(tree_iter(None)) + [None] + + Args: + tree (pytree): A pytree to flatten. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + An iterator over the leaf values. + """ + return optree.tree_iter( + tree, + is_leaf=is_leaf, + none_is_leaf=True, + namespace="torch", + ) + + +def tree_leaves( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> list[Any]: + """Get the leaves of a pytree. + + See also :func:`tree_flatten`. + + >>> tree = {"b": (2, [3, 4]), "a": 1, "c": None, "d": 5} + >>> tree_leaves(tree) + [2, 3, 4, 1, None, 5] + >>> tree_leaves(1) + [1] + >>> tree_leaves(None) + [None] + + Args: + tree (pytree): A pytree to flatten. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + A list of leaf values. + """ + return optree.tree_leaves( + tree, + is_leaf=is_leaf, + none_is_leaf=True, + namespace="torch", + ) + + +def tree_structure( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> TreeSpec: + """Get the treespec for a pytree. + + See also :func:`tree_flatten`. + + >>> tree = {"b": (2, [3, 4]), "a": 1, "c": None, "d": 5} + >>> tree_structure(tree) + PyTreeSpec({'b': (*, [*, *]), 'a': *, 'c': *, 'd': *}, NoneIsLeaf, namespace='torch') + >>> tree_structure(1) + PyTreeSpec(*, NoneIsLeaf, namespace='torch') + >>> tree_structure(None) + PyTreeSpec(*, NoneIsLeaf, namespace='torch') + + Args: + tree (pytree): A pytree to flatten. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + A treespec object representing the structure of the pytree. + """ + return optree.tree_structure( # type: ignore[return-value] + tree, + is_leaf=is_leaf, + none_is_leaf=True, + namespace="torch", + ) + + +def tree_map( + func: Callable[..., Any], + tree: PyTree, + *rests: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + """Map a multi-input function over pytree args to produce a new pytree. + + See also :func:`tree_map_`. + + >>> tree_map(lambda x: x + 1, {"x": 7, "y": (42, 64)}) + {'x': 8, 'y': (43, 65)} + >>> tree_map(lambda x: x is None, {"x": 7, "y": (42, 64), "z": None}) + {'x': False, 'y': (False, False), 'z': True} + + If multiple inputs are given, the structure of the tree is taken from the first input; + subsequent inputs need only have ``tree`` as a prefix: + + >>> tree_map(lambda x, y: [x] + y, [5, 6], [[7, 9], [1, 2]]) + [[5, 7, 9], [6, 1, 2]] + + Args: + func (callable): A function that takes ``1 + len(rests)`` arguments, to be applied at the + corresponding leaves of the pytrees. + tree (pytree): A pytree to be mapped over, with each leaf providing the first positional + argument to function ``func``. + rests (tuple of pytree): A tuple of pytrees, each of which has the same structure as + ``tree`` or has ``tree`` as a prefix. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + A new pytree with the same structure as ``tree`` but with the value at each leaf given by + ``func(x, *xs)`` where ``x`` is the value at the corresponding leaf in ``tree`` and ``xs`` + is the tuple of values at corresponding nodes in ``rests``. + """ + return optree.tree_map( + func, + tree, + *rests, + is_leaf=is_leaf, + none_is_leaf=True, + namespace="torch", + ) + + +def tree_map_( + func: Callable[..., Any], + tree: PyTree, + *rests: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + """Like :func:`tree_map`, but do an inplace call on each leaf and return the original tree. + + See also :func:`tree_map`. + + Args: + func (callable): A function that takes ``1 + len(rests)`` arguments, to be applied at the + corresponding leaves of the pytrees. + tree (pytree): A pytree to be mapped over, with each leaf providing the first positional + argument to function ``func``. + rests (tuple of pytree): A tuple of pytrees, each of which has the same structure as + ``tree`` or has ``tree`` as a prefix. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + The original ``tree`` with the value at each leaf is given by the side-effect of function + ``func(x, *xs)`` (not the return value) where ``x`` is the value at the corresponding leaf + in ``tree`` and ``xs`` is the tuple of values at values at corresponding nodes in ``rests``. + """ + return optree.tree_map_( + func, + tree, + *rests, + is_leaf=is_leaf, + none_is_leaf=True, + namespace="torch", + ) + + +Type2 = tuple[type[T], type[S]] +Type3 = tuple[type[T], type[S], type[U]] +TypeAny = type[Any] | tuple[type[Any], ...] | types.UnionType + +Fn2 = Callable[[T | S], R] +Fn3 = Callable[[T | S | U], R] +Fn = Callable[[T], R] +FnAny = Callable[[Any], R] + +MapOnlyFn = Callable[[T], Callable[[Any], Any]] + + +# These specializations help with type inference on the lambda passed to this +# function +@overload +def map_only(type_or_types_or_pred: type[T], /) -> MapOnlyFn[Fn[T, Any]]: ... + + +@overload +def map_only(type_or_types_or_pred: Type2[T, S], /) -> MapOnlyFn[Fn2[T, S, Any]]: ... + + +@overload +def map_only( + type_or_types_or_pred: Type3[T, S, U], / +) -> MapOnlyFn[Fn3[T, S, U, Any]]: ... + + +# This specialization is needed for the implementations below that call +@overload +def map_only(type_or_types_or_pred: TypeAny, /) -> MapOnlyFn[FnAny[Any]]: ... + + +@overload +def map_only( + type_or_types_or_pred: Callable[[Any], bool], / +) -> MapOnlyFn[FnAny[Any]]: ... + + +def map_only( + type_or_types_or_pred: TypeAny | Callable[[Any], bool], / +) -> MapOnlyFn[FnAny[Any]]: + """ + Suppose you are writing a tree_map over tensors, leaving everything + else unchanged. Ordinarily you would have to write: + + def go(t): + if isinstance(t, Tensor): + return ... + else: + return t + + With this function, you only need to write: + + @map_only(Tensor) + def go(t): + return ... + + You can also directly use 'tree_map_only' + """ + if isinstance(type_or_types_or_pred, (type, tuple, types.UnionType)): + + def pred(x: Any) -> bool: + return isinstance(x, type_or_types_or_pred) # type: ignore[arg-type] + + elif callable(type_or_types_or_pred): + pred = type_or_types_or_pred # type: ignore[assignment] + else: + raise TypeError("Argument must be a type, a tuple of types, or a callable.") + + def wrapper(func: Callable[[T], Any]) -> Callable[[Any], Any]: + @functools.wraps(func) + def wrapped(x: T) -> Any: + if pred(x): + return func(x) + return x + + return wrapped + + return wrapper + + +@overload +def tree_map_only( + type_or_types_or_pred: type[T], + /, + func: Fn[T, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only( + type_or_types_or_pred: Type2[T, S], + /, + func: Fn2[T, S, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only( + type_or_types_or_pred: Type3[T, S, U], + /, + func: Fn3[T, S, U, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only( + type_or_types_or_pred: TypeAny, + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only( + type_or_types_or_pred: Callable[[Any], bool], + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +def tree_map_only( + type_or_types_or_pred: TypeAny | Callable[[Any], bool], + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + return tree_map(map_only(type_or_types_or_pred)(func), tree, is_leaf=is_leaf) + + +@overload +def tree_map_only_( + type_or_types_or_pred: type[T], + /, + func: Fn[T, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only_( + type_or_types_or_pred: Type2[T, S], + /, + func: Fn2[T, S, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only_( + type_or_types_or_pred: Type3[T, S, U], + /, + func: Fn3[T, S, U, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only_( + type_or_types_or_pred: TypeAny, + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only_( + type_or_types_or_pred: Callable[[Any], bool], + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +def tree_map_only_( + type_or_types_or_pred: TypeAny | Callable[[Any], bool], + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + return tree_map_(map_only(type_or_types_or_pred)(func), tree, is_leaf=is_leaf) + + +def tree_all( + pred: Callable[[Any], bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + flat_args = tree_iter(tree, is_leaf=is_leaf) + return all(map(pred, flat_args)) + + +def tree_any( + pred: Callable[[Any], bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + flat_args = tree_iter(tree, is_leaf=is_leaf) + return any(map(pred, flat_args)) + + +@overload +def tree_all_only( + type_or_types: type[T], + /, + pred: Fn[T, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +@overload +def tree_all_only( + type_or_types: Type2[T, S], + /, + pred: Fn2[T, S, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +@overload +def tree_all_only( + type_or_types: Type3[T, S, U], + /, + pred: Fn3[T, S, U, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +def tree_all_only( + type_or_types: TypeAny, + /, + pred: FnAny[bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + flat_args = tree_iter(tree, is_leaf=is_leaf) + return all(pred(x) for x in flat_args if isinstance(x, type_or_types)) + + +@overload +def tree_any_only( + type_or_types: type[T], + /, + pred: Fn[T, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +@overload +def tree_any_only( + type_or_types: Type2[T, S], + /, + pred: Fn2[T, S, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +@overload +def tree_any_only( + type_or_types: Type3[T, S, U], + /, + pred: Fn3[T, S, U, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +def tree_any_only( + type_or_types: TypeAny, + /, + pred: FnAny[bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + flat_args = tree_iter(tree, is_leaf=is_leaf) + return any(pred(x) for x in flat_args if isinstance(x, type_or_types)) + + +def broadcast_prefix( + prefix_tree: PyTree, + full_tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> list[Any]: + """Return a list of broadcasted leaves in ``prefix_tree`` to match the number of leaves in ``full_tree``. + + If a ``prefix_tree`` is a prefix of a ``full_tree``, this means the ``full_tree`` can be + constructed by replacing the leaves of ``prefix_tree`` with appropriate **subtrees**. + + This function returns a list of leaves with the same size as ``full_tree``. The leaves are + replicated from ``prefix_tree``. The number of replicas is determined by the corresponding + subtree in ``full_tree``. + + >>> broadcast_prefix(1, [1, 2, 3]) + [1, 1, 1] + >>> broadcast_prefix([1, 2, 3], [1, 2, 3]) + [1, 2, 3] + >>> broadcast_prefix([1, 2, 3], [1, 2, 3, 4]) + Traceback (most recent call last): + ... + ValueError: list arity mismatch; expected: 3, got: 4; list: [1, 2, 3, 4]. + >>> broadcast_prefix([1, 2, 3], [1, 2, (3, 4)]) + [1, 2, 3, 3] + >>> broadcast_prefix([1, 2, 3], [1, 2, {"a": 3, "b": 4, "c": (None, 5)}]) + [1, 2, 3, 3, 3, 3] + + Args: + prefix_tree (pytree): A pytree with the same structure as a prefix of ``full_tree``. + full_tree (pytree): A pytree with the same structure as a suffix of ``prefix_tree``. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + A list of leaves in ``prefix_tree`` broadcasted to match the number of leaves in ``full_tree``. + """ + result: list[Any] = [] + + def add_leaves(x: Any, subtree: PyTree) -> None: + subtreespec = tree_structure(subtree, is_leaf=is_leaf) + result.extend([x] * subtreespec.num_leaves) + + tree_map_( + add_leaves, + prefix_tree, + full_tree, + is_leaf=is_leaf, + ) + return result + + +# Broadcasts a pytree to the provided TreeSpec and returns the flattened +# values. If this is not possible, then this function returns None. +# +# For example, given pytree=0 and spec=TreeSpec(list, None, [LeafSpec(), LeafSpec()]), +# would return [0, 0]. This is useful for part of the vmap implementation: +# a user can pass in vmap(fn, in_dims)(*inputs). `in_dims` should be +# broadcastable to the tree structure of `inputs` and we use +# _broadcast_to_and_flatten to check this. +def _broadcast_to_and_flatten( + tree: PyTree, + treespec: TreeSpec, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> list[Any] | None: + if not _is_pytreespec_instance(treespec): + raise TypeError( + f"Expected `treespec` to be an instance of " + f"PyTreeSpec but got item of type {type(treespec)}." + ) + full_tree = tree_unflatten([0] * treespec.num_leaves, treespec) + try: + return broadcast_prefix(tree, full_tree, is_leaf=is_leaf) + except ValueError: + return None + + +def treespec_dumps(treespec: TreeSpec, protocol: int | None = None) -> str: + """Serialize a treespec to a JSON string.""" + if not _is_pytreespec_instance(treespec): + raise TypeError( + f"Expected `treespec` to be an instance of " + f"PyTreeSpec but got item of type {type(treespec)}." + ) + + dummy_tree = tree_unflatten([0] * treespec.num_leaves, treespec) + orig_treespec = python_pytree.tree_structure(dummy_tree) + return python_pytree.treespec_dumps(orig_treespec, protocol=protocol) + + +@functools.lru_cache +def treespec_loads(serialized: str) -> TreeSpec: + """Deserialize a treespec from a JSON string.""" + orig_treespec = python_pytree.treespec_loads(serialized) + dummy_tree = python_pytree.tree_unflatten( + [0] * orig_treespec.num_leaves, + orig_treespec, + ) + treespec = tree_structure(dummy_tree) + return treespec + + +class _Asterisk(str): + __slots__ = () + + def __new__(cls) -> Self: + return super().__new__(cls, "*") + + def __repr__(self) -> str: + return "*" # no quotes + + +_asterisk = _Asterisk() +del _Asterisk + + +def treespec_pprint(treespec: TreeSpec) -> str: + dummy_tree = tree_unflatten([_asterisk] * treespec.num_leaves, treespec) + return repr(dummy_tree) + + +class LeafSpecMeta(type(TreeSpec)): # type: ignore[misc] + def __instancecheck__(self, instance: object) -> bool: + return _is_pytreespec_instance(instance) and instance.is_leaf() + + +@deprecated( + "`isinstance(treespec, LeafSpec)` is deprecated, " + "use `isinstance(treespec, TreeSpec)` and `treespec.is_leaf()` instead.", + category=FutureWarning, +) +class LeafSpec(TreeSpec, metaclass=LeafSpecMeta): # type: ignore[misc,final] + def __new__(cls) -> Self: + return treespec_leaf() # type: ignore[return-value] + + +def tree_flatten_with_path( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> tuple[list[tuple[KeyPath, Any]], TreeSpec]: + """Flattens a pytree like :func:`tree_flatten`, but also returns each leaf's key path. + + Args: + tree: a pytree to flatten. If it contains a custom type, that type must be + registered with an appropriate `tree_flatten_with_path_fn` when registered + with :func:`register_pytree_node`. + is_leaf: An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + Returns: + A tuple where the first element is a list of (key path, leaf) pairs, and the + second element is a :class:`TreeSpec` representing the structure of the flattened + tree. + """ + raise NotImplementedError("KeyPaths are not yet supported in cxx_pytree.") + + +def tree_leaves_with_path( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> list[tuple[KeyPath, Any]]: + """Gets the leaves of a pytree like ``tree_leaves`` and returns each leaf's key path. + + Args: + tree: a pytree. If it contains a custom type, that type must be + registered with an appropriate `tree_flatten_with_path_fn` when registered + with :func:`register_pytree_node`. + is_leaf: An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + Returns: + A list of (key path, leaf) pairs. + """ + raise NotImplementedError("KeyPaths are not yet supported in cxx_pytree.") + + +def tree_map_with_path( + func: Callable[..., Any], + tree: PyTree, + *rests: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + """Like :func:`tree_map`, but the provided callable takes an additional key path argument. + + Args: + func: A function that takes ``2 + len(rests)`` arguments, to be applied at the + corresponding leaves of the pytrees. The first positional argument + to ``func`` is the key path of the leaf in question. The second + positional argument is the value of the leaf. + tree: A pytree to be mapped over, with each leaf providing the first positional + argument to function ``func``. + rests: A tuple of pytrees, each of which has the same structure as + ``tree`` or has ``tree`` as a prefix. + is_leaf: An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns + A new pytree with the same structure as ``tree`` but with the value at each leaf given by + ``func(keypath, x, *xs)`` where ``keypath`` is the key path at the + corresponding leaf in ``tree``, ``x`` is the value at that leaf, and + ``xs`` is the tuple of values at corresponding nodes in ``rests``. + """ + raise NotImplementedError("KeyPaths are not yet supported in cxx_pytree.") + + +def keystr(kp: KeyPath) -> str: + """Given a key path, return a pretty-printed representation.""" + raise NotImplementedError("KeyPaths are not yet supported in cxx_pytree.") + + +def key_get(obj: Any, kp: KeyPath) -> Any: + """Given an object and a key path, return the value at the key path.""" + raise NotImplementedError("KeyPaths are not yet supported in cxx_pytree.") + + +with python_pytree._NODE_REGISTRY_LOCK: + # pyrefly: ignore [bad-assignment] + python_pytree._cxx_pytree_imported = True + args, kwargs = (), {} # type: ignore[var-annotated] + for args, kwargs in python_pytree._cxx_pytree_pending_imports: + _private_register_pytree_node(*args, **kwargs) + python_pytree._cxx_pytree_pending_imports.clear() + del args, kwargs diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..fd40af116d1395a1cc03f300feb2cfd9cf55fb39 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/__init__.py @@ -0,0 +1,30 @@ +# mypy: allow-untyped-defs +""" +DebugMode: a debugging TorchDispatchMode that intercepts and logs runtime calls. + +See torch.utils._debug_mode._mode for the full implementation and docstring. +""" + +from torch.utils._debug_mode._calls import ( + _AnnotateCall, + _DebugCall, + _get_call_name, + _OpCall, + _OutputPlacementCall, + _RedistributeCall, + _TritonKernelCall, +) +from torch.utils._debug_mode._mode import ( + DebugInterpreter, + DebugMode, + get_active_debug_mode, +) +from torch.utils._debug_mode._utils import ( + _stringify_shape, + hash_tensor_fn, + norm_hash_fn, + TensorIdTracker, +) + + +__all__ = ["DebugMode", "get_active_debug_mode"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/_calls.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/_calls.py new file mode 100644 index 0000000000000000000000000000000000000000..8da15e03bbb46e753b11ab5d201489669ac0bf54 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/_calls.py @@ -0,0 +1,354 @@ +# mypy: allow-untyped-defs +""" +Call record classes for DebugMode: _DebugCall and its subclasses. +""" + +from typing import Any, TYPE_CHECKING + +import torch + +# Import _utils module for mutable globals +from torch.utils._debug_mode import _utils +from torch.utils._debug_mode._utils import ( + _arg_to_str, + _get_op_name, + _get_stack_trace, + _maybe_get_autograd_trace, + REDISTRIBUTE_FUNC, + TensorIdTracker, +) +from torch.utils._pytree import tree_all, tree_map + + +if TYPE_CHECKING: + from torch._dynamo.device_interface import DeviceInterface + + +class _DebugCall: + """Base class for tracking operator calls in DebugMode""" + + def __init__( + self, + call_depth: int, + record: dict[str, Any] | None = None, + log: dict[str, Any] | None = None, + stack: bool = False, + ) -> None: + self.call_depth = call_depth + if stack: + self.stack_trace = _get_stack_trace() + self.fwd_stack_trace = _maybe_get_autograd_trace() + + # results from dispatch hooks + self.record = record + self.log = log + self.output_str: str | None = None + + def stringify_args( + self, attributes: list[str], tensor_memo: TensorIdTracker | None = None + ) -> None: + """ + To reduce memory consumption, this method stringifies args/kwargs, stores the result, and deletes original args/kwargs. + """ + raise NotImplementedError( + "Subclasses must implement stringify_args(), even if no-op" + ) + + def stringify_output( + self, + output: Any, + attributes: list[str], + tensor_memo: TensorIdTracker | None = None, + ) -> None: + """Store stringified version of call output in self.output_str""" + if tree_all(lambda x: x is None, output): + return + output_str = tree_map(lambda x: _arg_to_str(x, attributes, tensor_memo), output) + self.output_str = f" -> {str(output_str)}" + + def render(self, attributes: list[str]) -> str: + raise NotImplementedError("Subclasses must implement string render()") + + def __repr__(self) -> str: + return self.render([]) + + +class _OpCall(_DebugCall): + """Normal operator call""" + + def __init__( + self, + op, + args: tuple, + kwargs: dict, + call_depth: int, + stack: bool = False, + ) -> None: + super().__init__(call_depth, stack=stack) + self.op = op + self.args = args + self.kwargs = kwargs + + self.args_str: str | None = None + self.kwargs_str: str | None = None + + def stringify_args( + self, attributes: list[str], tensor_memo: TensorIdTracker | None = None + ) -> None: + self.args_str = ", ".join( + _arg_to_str(arg, attributes, tensor_memo) for arg in self.args + ) + if self.kwargs: + self.kwargs_str = ", " + ", ".join( + f"{k}={_arg_to_str(v, attributes, tensor_memo)}" + for k, v in self.kwargs.items() + ) + else: + self.kwargs_str = "" + del self.args + del self.kwargs + + def render(self, attributes: list[str]) -> str: + if self.args_str is not None: + args_str = self.args_str + else: + args_str = ", ".join(_arg_to_str(arg, attributes) for arg in self.args) + + if self.kwargs_str is not None: + kwargs_str = self.kwargs_str + else: + if self.kwargs: + kwargs_str = ", " + ", ".join( + f"{k}={_arg_to_str(v, attributes)}" for k, v in self.kwargs.items() + ) + else: + kwargs_str = "" + + if isinstance(self.op, torch._ops.OpOverload): + op_name = self.op.__qualname__ + elif hasattr(self.op, "__module__") and hasattr(self.op, "__name__"): + op_name = f"{self.op.__module__}.{self.op.__name__}" + else: + op_name = str(self.op) + + base_str = f"{op_name}({args_str}{kwargs_str})" + + if self.output_str: + base_str += self.output_str + if self.log: + base_str += f" # {self.log}" + return base_str + + def __iter__(self): + # for BC; tuple(self) returns (op, args, kwargs, call_depth) + if self.args_str is not None: + yield from [self.op, self.args_str, self.kwargs_str, self.call_depth] + else: + yield from [self.op, self.args, self.kwargs, self.call_depth] + + +class _RedistributeCall(_DebugCall): + def __init__( + self, + arg, + src_placement, + dst_placement, + transform_info_str, + call_depth, + stack=False, + is_explicit=False, + ) -> None: + super().__init__(call_depth, stack=stack) + self.arg = arg + self.src_placement = src_placement + self.dst_placement = dst_placement + self.transform_info_str = transform_info_str + self.is_explicit = is_explicit + self.is_outer_call = isinstance(arg, int) + + self.arg_str: str | None = None + + def stringify_args( + self, attributes: list[str], tensor_memo: TensorIdTracker | None = None + ) -> None: + self.arg_str = f"{_arg_to_str(self.arg, attributes, tensor_memo)}" + del self.arg + + def render(self, attributes: list[str]) -> str: + if self.arg_str is not None: + arg_str = self.arg_str + else: + arg_str = f"{_arg_to_str(self.arg, attributes)}" + + if self.transform_info_str is not None: # prioritize over src/dst placements + placement_str = f"trace: {self.transform_info_str}" + else: + src_placement_str = _arg_to_str(self.src_placement, attributes) + dst_placement_str = _arg_to_str(self.dst_placement, attributes) + placement_str = f"{src_placement_str} -> {dst_placement_str}" + + # DebugMode will add redistribute_input logs at 2 levels, + # once per redistribute decision, and once per redistributed input. + # We only annotate [implicit/explicit] logs on the former (outer-level call). + if self.is_outer_call: + annotation = " [implicit] " + elif self.is_explicit: + annotation = " [explicit] " + else: + annotation = "" + + base_str = f"{REDISTRIBUTE_FUNC}{annotation}({arg_str}, {placement_str})" + if self.output_str: + base_str += self.output_str + return base_str + + def __iter__(self): + # for BC; tuple(self) returns (op, placement info, kwargs, call_depth) + if self.arg_str is not None: + arg = self.arg_str + else: + arg = self.arg + + yield REDISTRIBUTE_FUNC + if self.transform_info_str: + yield [arg, self.transform_info_str] + else: + yield [arg, self.src_placement, self.dst_placement] + yield {} + yield self.call_depth + + +class _OutputPlacementCall(_DebugCall): + """Records output placement for a DTensor op.""" + + def __init__(self, placements_str: str, call_depth: int) -> None: + super().__init__(call_depth) + self.placements_str = placements_str + + def stringify_args( + self, attributes: list[str], tensor_memo: TensorIdTracker | None = None + ) -> None: + pass # Already stringified + + def render(self, attributes: list[str]) -> str: + return f"-> output: {self.placements_str}" + + +class _TritonKernelCall(_DebugCall): + """Triton kernel call from Inductor""" + + def __init__( + self, + kernel_name: str, + kwargs: dict[str, Any], + call_depth: int, + ): + super().__init__(call_depth) + self.kernel_name = kernel_name + self.kwargs = kwargs + self.kwargs_str: str | None = None + + self.pre_hashes: dict[str, Any] | None = None + self.post_hashes: dict[str, Any] | None = None + + def stringify_args( + self, attributes: list[str], tensor_memo: TensorIdTracker | None = None + ) -> None: + # Optionally hash kernel inputs before launch + if hash_fn := _utils._TRITON_INPUT_HASH_FN: + self.pre_hashes = { + k: hash_fn(v) + for k, v in self.kwargs.items() + if isinstance(v, torch.Tensor) + } + + if self.kwargs: + self.kwargs_str = ", ".join( + f"{k}={_arg_to_str(v, attributes, tensor_memo)}" + for k, v in self.kwargs.items() + ) + else: + self.kwargs_str = "" + + def render(self, attributes: list[str]) -> str: + base_str = f"[triton] {self.kernel_name}({self.kwargs_str})" + if self.pre_hashes: + pre_hashes_str = ", ".join(f"{k}: {v}" for k, v in self.pre_hashes.items()) + pre_hashes_str = ( + "\n " + + " " * self.call_depth + + f"# pre-kernel hashes: {{{pre_hashes_str}}}" + ) + else: + pre_hashes_str = "" + if self.post_hashes: + post_hashes_str = ", ".join( + f"{k}: {v}" for k, v in self.post_hashes.items() + ) + post_hashes_str = ( + "\n " + + " " * self.call_depth + + f"# post-kernel hashes: {{{post_hashes_str}}}" + ) + else: + post_hashes_str = "" + return f"{base_str}{pre_hashes_str}{post_hashes_str}\n" + + def finalize(self, device_interface: "DeviceInterface"): + # synchronize -> hash/store kernel results + device_interface.synchronize(device_interface.current_device()) + if _utils._RECORD_TRITON_OUTPUTS: + self.record = { + "output": { + k: v.clone() if isinstance(v, torch.Tensor) else v + for k, v in self.kwargs.items() + } + } + if hash_fn := _utils._TRITON_OUTPUT_HASH_FN: + self.post_hashes = { + k: hash_fn(v) + for k, v in self.kwargs.items() + if isinstance(v, torch.Tensor) + } + + # don't store tensors + del self.kwargs + + def __iter__(self): + yield from [self.kernel_name, (), self.kwargs_str, self.call_depth] + + +class _AnnotateCall(_DebugCall): + """Custom annotation call""" + + def __init__( + self, tag: Any, header: str, call_depth: int, stack: bool = False + ) -> None: + super().__init__(call_depth, stack=stack) + self.tag = tag + self.header = header + + def render(self, attributes: list[str]) -> str: + return f"[{self.header}] {self.tag}" + + def __iter__(self): + yield from [ + f"[{self.header}] {self.tag}", + (), + {}, + self.call_depth, + ] + + +def _get_call_name(call: _DebugCall) -> str: + """String identifying _DebugCall (e.g. func, kernel, module name)""" + if isinstance(call, _OpCall): + return _get_op_name(call.op) + elif isinstance(call, _TritonKernelCall): + return call.kernel_name + elif isinstance(call, _AnnotateCall): + return f"[{call.header}] {call.tag}" + elif isinstance(call, _RedistributeCall): + return REDISTRIBUTE_FUNC + else: + return str(call) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/_mode.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/_mode.py new file mode 100644 index 0000000000000000000000000000000000000000..0d73b0847b8e86d7fe6ff64aaa0f3323127b340f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/_mode.py @@ -0,0 +1,1053 @@ +# mypy: allow-untyped-defs +""" +DebugMode is a debugging TorchDispatchMode that intercepts and logs runtime calls +to a hierarchical string dump. It logs real tensor, DTensor, and optionally FakeTensor +operations, with some additional handling for DTensor internals. + +An example dump from an eager mode DTensor matmul: + + torch.mm(dt$0: f32[8, 8]| S(0), dt$1: f32[8, 32]| S(0)) -> dt$6: f32[8, 32]| S(0) + aten::mm(dt$0: f32[8, 8]| S(0), dt$1: f32[8, 32]| S(0)) + redistribute_input(1, S(0) -> R) + redistribute_input(t$2: f32[1, 32], trace: S(0)->R) + _c10d_functional::all_gather_into_tensor(t$2: f32[1, 32], 8, 0) -> t$3: f32[8, 32] + _c10d_functional::wait_tensor(t$3: f32[8, 32]) -> t$3: f32[8, 32] + aten::mm(t$4: f32[1, 8], t$3: f32[8, 32]) -> t$5: f32[1, 32] + +This mode runs "under" compile, which means it hides itself during compilation, and is re-enabled +at runtime, and DebugMode-related operations won't show up in the compiled region. +DebugMode also provides some visibility into non-torch-dispatch calls (e.g. DTensor redistribute calls, +inductor-generated triton kernels), but requires special handling for these, since dispatch modes +can't intercept them by default. + +The mode also provides some extensions for custom debugging (e.g. adding custom dispatch call hooks +via dispatch_hooks), or numerics debugging (e.g. tensor hashing for bitwise equivalence/closeness, +via log_tensor_hashes). These decorators allow annotating string dumps with additional per-call information, +for any region of runtime code. + +Usage:: + + with DebugMode() as debug_mode: + result = some_pytorch_operation(tensor_input) + print(debug_mode.debug_string()) +""" + +import contextlib +import functools +import logging +from collections.abc import Callable +from typing import Any, TYPE_CHECKING + +import torch +from torch._logging import warning_once +from torch._subclasses.fake_tensor import FakeTensor, FakeTensorMode + +# Import _utils module for mutable globals +from torch.utils._debug_mode import _utils +from torch.utils._debug_mode._calls import ( + _AnnotateCall, + _get_call_name, + _OpCall, + _OutputPlacementCall, + _RedistributeCall, + _TritonKernelCall, +) +from torch.utils._debug_mode._utils import ( + _compute_rel_diff, + _get_user_stack_trace, + _run_dispatch_hooks, + _run_dispatch_pre_log_hooks, + _stringify_dtensor_spec, + hash_tensor_fn, + norm_hash_fn, + TensorIdTracker, +) +from torch.utils._python_dispatch import ( + _get_current_dispatch_mode, + _get_current_dispatch_mode_stack, + TorchDispatchMode, +) +from torch.utils._pytree import ( + keystr, + tree_all, + tree_map, + tree_map_only, + tree_map_with_path, +) + + +if TYPE_CHECKING: + from torch.distributed._tools.mod_tracker import ModTracker + + +log = logging.getLogger(__name__) + + +# Counter for active DebugMode instances (fast path for get_active_debug_mode) +_ACTIVE_DEBUG_MODE_COUNT = 0 + +_annotate_decorated = False + + +def _ensure_annotate_decorated(): + """ + Lazily apply dont_skip_tracing decorator to DebugMode._annotate, to avoid circular import/initialization issues. + """ + global _annotate_decorated + if not _annotate_decorated: + DebugMode._annotate = torch._dynamo.dont_skip_tracing(DebugMode._annotate) # type: ignore[has-type] + + # Mark annotate as side-effectful so aot_eager doesn't DCE it. + from torch.fx.node import _side_effectful_functions + + _side_effectful_functions.add(torch.ops.debug_mode_ops.annotate.default) + + # Register no-op lowering for inductor backend + from torch._inductor.lowering import register_lowering + + @register_lowering(torch.ops.debug_mode_ops.annotate) + def _annotate_lowering(tag: str) -> None: + warning_once(log, 'DebugMode._annotate() is a no-op for backend="inductor"') + return None + + _annotate_decorated = True + + +@torch.library.custom_op("debug_mode_ops::annotate", mutates_args=()) +def _annotate(tag: str) -> None: + # This is special-cased in DebugMode.__torch_dispatch__ + return None + + +@_annotate.register_fake +def _annotate_fake(tag: str) -> None: + return None + + +class DebugInterpreter(torch.fx.Interpreter): + """ + Interpreter class for running aot_eager compiled regions when DebugMode is active, + instead of using the compiled code. This gives us access to fx.Node metadata to decorate + and contextualize DebugMode logs (e.g. nn_module_stack, stack_trace, compiled region boundaries). + + Note: this is currently only enabled with DebugMode(run_compile_with_interpreter=True). + """ + + def __init__(self, module, backend): + super().__init__(module) + self.mode = get_active_debug_mode() + if self.mode is None: + raise RuntimeError("No DebugMode is currently active") + + # for tracking initial nn_module_stack + self.base_nn_module_stack = list(self.mode.current_nn_module_stack) + + # annotate start of region + self.backend = backend + self.mode.operators.append( + _AnnotateCall( + "enter", f"{self.backend} region (compile)", self.mode.call_depth + ) + ) + + def run_node(self, n: torch.fx.Node) -> Any: + if self.mode is None: + raise RuntimeError("No DebugMode is currently active") + + # handling of nn.Module stack + if self.mode.record_nn_module and n.op not in ["placeholder", "output"]: + self.mode._handle_fx_nn_module_stack( + self.base_nn_module_stack, + n.meta.get("nn_module_stack", {}), + n.meta.get("fwd_nn_module_stack", {}), + ) + + # override stack trace with n.meta + if ( + self.mode.record_stack_trace + and n.op not in ["placeholder", "output"] + and (stack_trace := n.meta.get("stack_trace", None)) is not None + ): + with self.mode.set_fx_stack_trace(stack_trace): + return super().run_node(n) + else: + return super().run_node(n) + + def run(self, *args, **kwargs): + if self.mode is None: + raise RuntimeError("No DebugMode is currently active") + result = super().run(*args) + + # reset nn.Module stack to pre-compiled region value + if len(self.mode.current_nn_module_stack) < len(self.base_nn_module_stack): + warning_once( + log, "unexpected handling of nn_module_stack in DebugInterpreter" + ) + while len(self.mode.current_nn_module_stack) > len(self.base_nn_module_stack): + self.mode._exit_nn_module_call() + + # annotate end of region + self.mode.operators.append( + _AnnotateCall( + "exit", f"{self.backend} region (compile)", self.mode.call_depth + ) + ) + + return result + + +class DebugMode(TorchDispatchMode): + def __init__( + self, + *, + record_torchfunction=False, + record_faketensor=False, + record_realtensor=True, + record_tensor_attributes=None, + record_nn_module=False, + store_original_args=False, + record_stack_trace=False, + record_output=True, + record_ids=False, + record_profiler_context=True, + record_localtensor=True, + run_compile_with_interpreter=False, + ) -> None: + super().__init__() + import torch.distributed.tensor # noqa: F401 + + _ensure_annotate_decorated() + self.supports_higher_order_operators = True + + # Pushes DebugMode onto the torchfunction stack, and records __torch_function__ calls as well. + # WARNING: currently incompatible with torch.compile due to dynamo guard failures. + self.record_torchfunction = record_torchfunction + + # Records __torch_dispatch__ calls on FakeTensors. + self.record_faketensor = record_faketensor + + # Records __torch_dispatch__ calls on real tensors. + self.record_realtensor = record_realtensor + + # Records __torch_dispatch__ calls on LocalTensor. + self.record_localtensor = record_localtensor + + # Optional list[str] of tensor attributes, to be annotated in the string dump. + self.record_tensor_attributes = record_tensor_attributes or [] + + # Uses ModTracker to record nn.Module entrances. + # This flag currently has no effect on torch.compiled-regions. + self.record_nn_module = record_nn_module + + self.module_tracker: ModTracker | None = None + if self.record_nn_module: + self.module_tracker_setup() + + # If True, stores call args/kwargs in logs, without immediately stringifying. + # Defaults to False for memory concerns. + self.store_original_args = store_original_args + + # For stack trace recording, stores log call stack traces in .stack_trace. + # For backward graph nodes, will also store the corresponding forward stack traces in .fwd_stack_trace. + # NOTE: this is only available if autograd tracebacks are being set during the forward pass, + # e.g. via DebugMode(record_stack_trace=True), or torch.autograd.set_detect_anomaly(). + self.record_stack_trace = record_stack_trace + + # Records call outputs in logs (e.g. for __torch_dispatch__, __torch_function__, redistribute_input) + self.record_output: bool = record_output + + # Annotates string dumps with graph-style tensor ids, e.g. op($1, $2) -> $3. + self.record_ids: bool = record_ids + + # Annotates string dumps with profiler.record_function contexts from runtime code. + # Currently does not preserve contexts inside torch.compile-d regions. + self.record_profiler_context: bool = record_profiler_context + + # For aot_eager compiled regions, wraps the compiled fx.GraphModule with a DebugInterpreter, + # and uses it at runtime for node metadata visibility. + self.run_compile_with_interpreter: bool = run_compile_with_interpreter + + self.reset() + + def reset(self) -> None: + self.operators = [] + self.call_depth = 0 + self._tensor_memo = TensorIdTracker() + self._output_info: dict[int, object] = {} + self.ignored_record_functions = 0 + self.current_nn_module_stack = [] + self.fx_stack_trace = None + + def _track_op_output(self, op_index, result) -> None: + """Assign IDs to output tensors and store in output_info""" + self._output_info[op_index] = result + + # Without this override, running torch.compile under DebugMode + # will force torch.compile to always use the "eager" backend + # With this, DebugMode will not take effect on torch.compile + @classmethod + def ignore_compile_internals(cls) -> bool: + return True + + def _record_call(self, call) -> None: + if _utils._IN_INDUCTOR_BENCHMARK: + return + + if str(call).startswith("profiler::_record_function"): + return + + if not self.store_original_args: + call.stringify_args( + self.record_tensor_attributes, + self._tensor_memo if self.record_ids else None, + ) + if self.fx_stack_trace: + call.stack_trace = call.fwd_stack_trace = self.fx_stack_trace + self.operators.append(call) + + def _record_call_output(self, call, output) -> None: + if not self.record_output: + return + call.stringify_output( + output, + self.record_tensor_attributes, + self._tensor_memo if self.record_ids else None, + ) + + def __torch_function__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + + call = _OpCall( + func, args, kwargs, self.call_depth, stack=self.record_stack_trace + ) + self._record_call(call) + + try: + self.call_depth += 1 + result = func(*args, **kwargs) + self._record_call_output(call, result) + return result + finally: + self.call_depth -= 1 + + def _maybe_record_function(self, tag): + # filter out tags that appear noisy, or aren't runtime-related + if any( + tag.startswith(prefix) + for prefix in [ + # assuming these are from benchmarking, not the actual runtime call + "CachingAutotuner.", + "InductorBenchmarker.", + # inductor compilation + "compile_fx..", + ] + ): + self.ignored_record_functions += 1 + return + + call = _AnnotateCall( + tag, "record function", self.call_depth, stack=self.record_stack_trace + ) + self.operators.append(call) + self.call_depth += 1 + + def _maybe_exit_record_function(self): + if self.ignored_record_functions < 0: + raise AssertionError( + f"ignored_record_functions is negative: {self.ignored_record_functions}" + ) + if self.ignored_record_functions > 0: + self.ignored_record_functions -= 1 + else: + self.call_depth -= 1 + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + + # Handle record_function entries + if self.record_profiler_context: + if func == torch.ops.profiler._record_function_enter_new.default: + if len(args) != 1: + raise AssertionError(f"expected 1 arg, got {len(args)}") + self._maybe_record_function(args[0]) + elif func == torch.ops.profiler._record_function_exit._RecordFunction: + self._maybe_exit_record_function() + + # Handle DebugMode._annotate() + if func is torch.ops.debug_mode_ops.annotate.default: + if len(args) != 1: + raise AssertionError(f"expected 1 arg, got {len(args)}") + self._handle_annotate(args[0]) + return + + from torch.distributed._functional_collectives import AsyncCollectiveTensor + from torch.distributed._local_tensor import LocalTensor + + # Record the operation with its call depth + call = None + if torch.distributed.tensor.DTensor in types: + call = _OpCall( + func, args, kwargs, self.call_depth, stack=self.record_stack_trace + ) + self._record_call(call) + return NotImplemented + elif FakeTensor in types or isinstance( + _get_current_dispatch_mode(), FakeTensorMode + ): + if self.record_faketensor: + if func != torch.ops.prim.device.default: + call = _OpCall( + func, + args, + kwargs, + self.call_depth + 1, + stack=self.record_stack_trace, + ) + self._record_call(call) + # TODO: check the context manager + elif LocalTensor in types: + if self.record_localtensor: + call = _OpCall( + func, + args, + kwargs, + self.call_depth + 1, + stack=self.record_stack_trace, + ) + self._record_call(call) + elif AsyncCollectiveTensor in types: + # Record AsyncCollectiveTensor operations so debugging/tracing tools can see them + if self.record_realtensor: + call = _OpCall( + func, + args, + kwargs, + self.call_depth + 1, + stack=self.record_stack_trace, + ) + self._record_call(call) + elif len(types) == 0: + if self.record_realtensor: + call = _OpCall( + func, + args, + kwargs, + self.call_depth + 1, + stack=self.record_stack_trace, + ) + self._record_call(call) + + # Run pre-hooks before executing the operation to hash inputs + # We have to run becore the func() call in case there's any + # in-place mutation + if call: + _run_dispatch_pre_log_hooks(call, func, types, args, kwargs) + + result = func(*args, **kwargs) + if call: + self._record_call_output(call, result) + _run_dispatch_hooks(call, func, types, args, kwargs, result) + + return result + + def __enter__(self): + global _ACTIVE_DEBUG_MODE_COUNT + _ACTIVE_DEBUG_MODE_COUNT += 1 + if self.record_torchfunction: + torch._C._push_on_torch_function_stack(self) + + super().__enter__() + if self.record_nn_module: + self.module_tracker.__enter__() # type: ignore[attribute, union-attr] + + if self.record_stack_trace: + self.anomaly_for_traces = torch.autograd.set_detect_anomaly( + True, check_nan=False + ) + self.anomaly_for_traces.__enter__() + return self + + # pyrefly: ignore [bad-override] + def __exit__(self, *args): + global _ACTIVE_DEBUG_MODE_COUNT + _ACTIVE_DEBUG_MODE_COUNT -= 1 + super().__exit__(*args) + if self.record_nn_module: + self.module_tracker.__exit__() # type: ignore[attribute, union-attr] + if self.record_torchfunction: + torch._C._pop_torch_function_stack() + if self.record_stack_trace: + self.anomaly_for_traces.__exit__(*args) + + @contextlib.contextmanager + def set_fx_stack_trace(self, stack_trace): + self.fx_stack_trace = stack_trace + try: + yield + finally: + self.fx_stack_trace = None + + def _enter_nn_module_call(self, fqn, header): + call = _AnnotateCall( + fqn, header, self.call_depth + 1, stack=self.record_stack_trace + ) + self.operators.append(call) + self.current_nn_module_stack.append(fqn) + self.call_depth += 1 + + def _exit_nn_module_call(self): + self.call_depth -= 1 + self.current_nn_module_stack.pop() + + def module_tracker_setup(self) -> None: + from torch.distributed._tools.mod_tracker import ModTracker + + self.module_tracker = ModTracker() + + # module pre-fw hook: record module call + def pre_fw_hook(module, input) -> None: + fqn = self.module_tracker._get_mod_name(module) # type: ignore[attribute, union-attr] + self._enter_nn_module_call(fqn, "nn.Mod") + + # module post-fw hook: decrement call depth + def post_fw_hook(module, input, output) -> None: + self._exit_nn_module_call() + + self.module_tracker.register_user_hooks(pre_fw_hook, post_fw_hook) + + def _handle_fx_nn_module_stack( + self, + base_stack: list[str], + nn_module_stack: dict[str, tuple[str, Any]] | None, + fwd_nn_module_stack: dict[str, tuple[str, Any]] | None, + ) -> None: + """ + Called when DebugInterpreter observes nn_module_stack or fwd_nn_module_stack metadata + from executing the compiled GraphModule. + + If the current module stack is mismatched with what's currently tracked in DebugMode + (current_nn_module_stack), we adjust call depth and add new [nn.Module] log entries accordingly. + """ + + nn_module_stack = nn_module_stack or {} + fwd_nn_module_stack = fwd_nn_module_stack or {} + if nn_module_stack and fwd_nn_module_stack: + raise AssertionError( + "Expecting at most one of nn_module_stack and fwd_nn_module_stack." + ) + + is_fwd = nn_module_stack + stack = nn_module_stack if is_fwd else fwd_nn_module_stack + + # forward stack + current_stack = self.current_nn_module_stack + new_stack = base_stack + [v[0] for v in stack.values()] + + entered = set(new_stack) - set(current_stack) + exited = set(current_stack) - set(new_stack) + + # Decrement depth for exited modules + for _ in exited: + self._exit_nn_module_call() + if self.call_depth < 0: + raise AssertionError("Unexpectedly, DebugMode call_depth is negative") + + # Add [nn.Module] entries for newly entered modules + for fqn in sorted(entered): + self._enter_nn_module_call( + fqn, "nn.Mod (compile)" if is_fwd else "nn.Mod (compile bwd)" + ) + + self.current_nn_module_stack = new_stack + + @contextlib.contextmanager + def record_redistribute_calls( + self, + arg, + src_placement, + dst_placement, + transform_info_str: str | None = None, + is_explicit: bool = False, + ): + try: + self._record_call( + _RedistributeCall( + arg, + src_placement=src_placement, + dst_placement=dst_placement, + transform_info_str=transform_info_str, + call_depth=self.call_depth + 1, + stack=self.record_stack_trace, + is_explicit=is_explicit, + ) + ) + self.call_depth += 1 + yield + finally: + self.call_depth -= 1 + + def record_output_placements(self, output_spec) -> None: + """Record output placements for a DTensor op as a separate line.""" + if not self.record_output: + return + from torch.distributed.tensor._dtensor_spec import DTensorSpec + + placements_str = str( + tree_map_only(DTensorSpec, _stringify_dtensor_spec, output_spec) + ) + call = _OutputPlacementCall(placements_str, self.call_depth + 1) + self._record_call(call) + + def record_triton_kernel( + self, kernel_name: str, kwargs: dict[str, Any] + ) -> _TritonKernelCall: + call = _TritonKernelCall(kernel_name, kwargs, self.call_depth + 1) + call.stringify_args(self.record_tensor_attributes) + self.operators.append(call) + return call + + def debug_string(self, show_stack_trace: bool | None = None) -> str: + """ + show_stack_trace: option to display one-line stack trace summaries above groups + of operations (similar to gm.print_readable() style). + Requires record_stack_trace=True. + if None, uses self.record_stack_trace, otherwise overrides it. + """ + show_stack_trace = ( + self.record_stack_trace if show_stack_trace is None else show_stack_trace + ) + + with torch._C.DisableTorchFunction(): + if not show_stack_trace: + result = "\n".join( + " " + + " " * op.call_depth + + op.render(self.record_tensor_attributes) + for op in self.operators + ) + return result + + # Group operations by stack trace + lines = [] + prev_stack_summary = None + + for op in self.operators: + # Get the stack trace: prefer fwd_stack_trace, fallback to stack_trace + stack_trace = None + if hasattr(op, "fwd_stack_trace") and op.fwd_stack_trace: + stack_trace = op.fwd_stack_trace + elif hasattr(op, "stack_trace") and op.stack_trace: + stack_trace = op.stack_trace + + stack_summary = None + if stack_trace: + stack_summary = _get_user_stack_trace(stack_trace) + + if stack_summary and stack_summary != prev_stack_summary: + # add blank line before stack trace comment for readability + if lines: # don't add blank line at the very start + lines.append("") + indent = " " * (op.call_depth + 1) + lines.append(indent + "# " + stack_summary) + prev_stack_summary = stack_summary + + # Add the operation line + line = ( + " " + + " " * op.call_depth + + op.render(self.record_tensor_attributes) + ) + lines.append(line) + + return "\n".join(lines) + + @staticmethod + @contextlib.contextmanager + def dispatch_hooks( + record_hook: Callable | None = None, + log_hook: Callable | None = None, + pre_log_hook: Callable | None = None, + ): + """ + Allows installing post-hooks on arguments to intercepted __torch_dispatch__ calls; + hook signatures are expected as (func, types, args, kwargs, result), + i.e. __torch_dispatch__ args + return value. + + Logging hook outputs are stored in call.log and annotate calls in debug_string(), + while recording hook outputs are just stored in call.record. + For now hooks are expected to return dictionaries. + + pre_log_hook signature is (func, types, args, kwargs, call) and is executed before + the operation. It allows capturing state before in-place mutations. + """ + if record_hook: + _utils._DISPATCH_RECORD_HOOKS.append(record_hook) + if log_hook: + _utils._DISPATCH_LOG_HOOKS.append(log_hook) + if pre_log_hook: + _utils._DISPATCH_PRE_LOG_HOOKS.append(pre_log_hook) + try: + yield + finally: + if record_hook: + _utils._DISPATCH_RECORD_HOOKS.pop() + if log_hook: + _utils._DISPATCH_LOG_HOOKS.pop() + if pre_log_hook: + _utils._DISPATCH_PRE_LOG_HOOKS.pop() + + @staticmethod + @contextlib.contextmanager + def record_outputs(): + """ + Hook for storing cloned output tensors in .record["output"]. + """ + + def dispatch_hook(func, types, args, kwargs, result): + out = tree_map( + lambda x: x.clone() if isinstance(x, torch.Tensor) else x, result + ) + return {"output": out} + + try: + _old_record_triton = _utils._RECORD_TRITON_OUTPUTS + _utils._RECORD_TRITON_OUTPUTS = True + with DebugMode.dispatch_hooks(record_hook=dispatch_hook): + yield + finally: + _utils._RECORD_TRITON_OUTPUTS = _old_record_triton + + @staticmethod + @contextlib.contextmanager + def log_tensor_hashes( + hash_fn: Callable | str | list[str] = "norm", hash_inputs: bool = False + ): + """ + Installs hook for tensor hash logging. + + hash_fn: One of: + - Custom-defined hash function + - String: one of ("norm", "hash_tensor") + - "norm": uses norm_hash_fn; basically tensor's L1 norm + - "hash_tensor": uses torch.hash_tensor (XOR sum reduction) + - List of strings: returns tuple of hashes from above options + hash_inputs: if True, also hashes tensors in (args, kwargs), storing them in "input_hash". + Input hashes are captured before the operation executes, so they reflect the state before + any in-place mutations. + """ + + def hash_fn_option(hash_type): + if not isinstance(hash_type, str) or hash_type not in [ + "norm", + "hash_tensor", + ]: + raise AssertionError( + f"hash_type must be 'norm' or 'hash_tensor', got {hash_type!r}" + ) + return functools.partial( + norm_hash_fn if hash_type == "norm" else hash_tensor_fn, use_scalar=True + ) + + if callable(hash_fn): + fn = hash_fn + elif isinstance(hash_fn, str): + fn = hash_fn_option(hash_fn) + elif isinstance(hash_fn, list): + fns = [hash_fn_option(fn) for fn in hash_fn] + fn = lambda x: tuple(fn(x) for fn in fns) # noqa: E731 + else: + raise NotImplementedError( + f"log_tensor_hashes() expected hash_fn to be callable, str, or list[str], but found {type(hash_fn)}" + ) + + def _tree_hash(obj): + return tree_map( + lambda x: fn(x) if isinstance(x, torch.Tensor) else None, obj + ) + + def _dispatch_pre_log_hook(func, types, args, kwargs, call): + """Pre-hook to capture input hashes before operation executes""" + if "empty" in str(func) or "profiler" in str(func): + return None + + if hash_inputs: + # Capture input hashes before the operation + input_hash = _tree_hash((args, kwargs)) + if not tree_all(lambda x: x is None, input_hash): + return {"input_hash": input_hash} + return None + + def _dispatch_post_hook(func, types, args, kwargs, result): + """Post-hook to capture output hashes after operation executes""" + if "empty" in str(func) or "profiler" in str(func): + return None + + out = {} + out["hash"] = _tree_hash(result) + + if tree_all(lambda x: x is None, out.values()): + return None + return out + + try: + if hash_inputs: + _old_input_hfn = _utils._TRITON_INPUT_HASH_FN + _utils._TRITON_INPUT_HASH_FN = fn + _old_output_hfn = _utils._TRITON_OUTPUT_HASH_FN + _utils._TRITON_OUTPUT_HASH_FN = fn + with DebugMode.dispatch_hooks( + log_hook=_dispatch_post_hook, + pre_log_hook=_dispatch_pre_log_hook if hash_inputs else None, + ): + yield + finally: + if hash_inputs: + _utils._TRITON_INPUT_HASH_FN = _old_input_hfn # type: ignore[assignment] + _utils._TRITON_OUTPUT_HASH_FN = _old_output_hfn + + @staticmethod + @contextlib.contextmanager + def _benchmarking_inductor(): + """ + Context manager for disabling logging during inductor benchmarking, + so logs don't contain all kernels launched from autotuning. + """ + try: + _utils._IN_INDUCTOR_BENCHMARK = True + yield + finally: + _utils._IN_INDUCTOR_BENCHMARK = False + + @property + def logs(self): + return list(self.operators) + + def _handle_annotate(self, tag): + """Handles DebugMode._annotate()""" + call = _AnnotateCall(tag, "annotate", self.call_depth, self.record_stack_trace) + self.operators.append(call) + + @staticmethod + def _annotate(tag: Any) -> None: + """ + If an active DebugMode exists, adds an "[annotate] " entry to the logs. Useful for contextualizing logs. + Implemented with a custom op. + """ + torch.ops.debug_mode_ops.annotate(tag) + + @staticmethod + def check_hash_mismatches( + logs1: list, logs2: list, compare_inputs: bool = False + ) -> list[dict]: + """ + Compares tensor hashes between two DebugMode runs, for checking run-to-run numerical divergence. + + This first validates the two log sequences have identical structure (same operations, input shapes/dtypes, etc.), + then compares tensor hash values, and returns a list of call outputs where mismatches were found. + Expects input logs to have been run with log_tensor_hashes, and looks for hashes in .log["hash"] & .log["input_hash"] + (or .post_hashes & .pre_hashes for triton kernels). + + note: skips checking log pairs where hashes aren't present, but will raise if present in one & not the other. + + Args: + logs1: logs from the first DebugMode run (from debug_mode.logs) + logs2: logs from the second DebugMode run + compare_inputs: If True, also compare input tensor hashes (default: only output checking) + + Returns: + List of dictionaries describing hash mismatches. Each dict contains: + - call_type: "torch op" or "triton kernel" + - call: Operator/kernel name + - arg_name: For triton kernels, the argument name; None for torch ops + - pytree_path: For torch ops, the pytree path to the differing tensor; None for kernels + - hash1: Hash value from the first run + - hash2: Hash value from the second run + - rel_diff: Relative difference between hash values + - is_input_hash: True if this is an input hash, False for output hash + + Raises: + ValueError: If logs have different lengths, call types, operator names, or call depths + + Usage:: + + # Run model first time + with DebugMode() as debug_mode, DebugMode.log_tensor_hashes(): + model(x) + logs1 = debug_mode.logs + + # Run again, in exactly the same way + with DebugMode() as debug_mode, DebugMode.log_tensor_hashes(): + model(x) + logs2 = debug_mode.logs + + mismatches = DebugMode.check_hash_mismatches(logs1, logs2) + for m in mismatches: + print(f"{m['call']}: hash diff {m['rel_diff']:.2e}") + """ + if len(logs1) != len(logs2): + raise ValueError(f"Log lengths don't match: {len(logs1)} vs {len(logs2)}") + + difference_info = [] + for i, (log1, log2) in enumerate(zip(logs1, logs2)): + # check call type + call1_type = type(log1).__name__ + call2_type = type(log2).__name__ + if call1_type != call2_type: + raise ValueError( + f"Call types don't match at index {i}: {call1_type} vs {call2_type}" + ) + call_type = call1_type + + # check call name + op1_name, op2_name = _get_call_name(log1), _get_call_name(log2) + if op1_name != op2_name: + raise ValueError( + f"Operators don't match at index {i}: {call_type}[{op1_name}] vs {call_type}[{op2_name}]" + ) + op_name = op1_name + + # check call depth + if log1.call_depth != log2.call_depth: + raise ValueError( + f"Call depths for {call_type}[{op_name}] don't match at index {i}: {log1.call_depth} vs {log2.call_depth}" + ) + + # Redistribute: call args should be the same + if isinstance(log1, _RedistributeCall): + if tuple(log1) != tuple(log2): + raise ValueError( + f"Redistribute calls don't match at index {i}: {log1} vs {log2}" + ) + + # Triton kernel: same arg names, arg types + elif isinstance(log1, _TritonKernelCall): + if log1.kwargs_str != log2.kwargs_str: + raise ValueError( + f"Triton kernel call args don't match for {log1.kernel_name} at index {i}:" + f"\n\nlog1: {log1.kwargs_str}\n\nlog2: {log2.kwargs_str}" + ) + + def compare_triton_hashes(hashes1, hashes2, is_input): + if set(hashes1.keys()) != set(hashes2.keys()): # type: ignore[union-attr] + raise AssertionError( + f"hash key mismatch: {set(hashes1.keys())} vs {set(hashes2.keys())}" + ) + for key in hashes1: + if hashes1[key] != hashes2[key]: + difference_info.append( + { + "call_type": "triton kernel", + "call": op_name, + "arg_name": key, + "pytree_path": None, + "hash1": hashes1[key], + "hash2": hashes2[key], + "rel_diff": _compute_rel_diff( + hashes1[key], hashes2[key] + ), + "is_input_hash": is_input, + } + ) + + # check output hashes + has_post_1, has_post_2 = ( + log1.post_hashes is not None, + log2.post_hashes is not None, + ) + if has_post_1 != has_post_2: + raise ValueError( + f"Triton kernel post-hash presence inconsistent for {log1.kernel_name} " + f"at index {i}: log1 has post_hashes={has_post_1}, log2 has post_hashes={has_post_2}" + ) + + if has_post_1: + compare_triton_hashes( + log1.post_hashes, log2.post_hashes, is_input=False + ) + + # maybe check input hashes + if compare_inputs: + has_pre_1, has_pre_2 = ( + log1.pre_hashes is not None, + log2.pre_hashes is not None, + ) + if has_pre_1 != has_pre_2: + raise ValueError( + f"Triton kernel pre-hash presence inconsistent for {log1.kernel_name} " + f"at index {i}: log1 has pre_hashes={has_pre_1}, log2 has pre_hashes={has_pre_2}" + ) + + if has_pre_1: + compare_triton_hashes( + log1.pre_hashes, log2.pre_hashes, is_input=True + ) + + # regular log calls + elif isinstance(log1, _OpCall): + + def compare_op_hashes(hashes1, hashes2, is_input): + def _helper(keypath, hash1, hash2): + if hash1 != hash2: + difference_info.append( + { + "call_type": "torch op", + "call": op_name, + "arg_name": None, + "pytree_path": keystr(keypath), + "hash1": hash1, + "hash2": hash2, + "rel_diff": _compute_rel_diff(hash1, hash2), + "is_input_hash": is_input, + } + ) + + tree_map_with_path(_helper, hashes1, hashes2) + + # check output hashes + has_hash1 = log1.log is not None and "hash" in log1.log + has_hash2 = log2.log is not None and "hash" in log2.log + if has_hash1 != has_hash2: + raise ValueError( + f"Output hash presence inconsistent for triton kernel {call_type}[{op_name}] " + f"at index {i}: log1 has hash={has_hash1}, log2 has hash={has_hash2}" + ) + + if has_hash1: + compare_op_hashes( + log1.log["hash"], # type: ignore[union-attr] + log2.log["hash"], + is_input=False, + ) + + # maybe check input hashes + if compare_inputs: + has_hash1 = log1.log is not None and "input_hash" in log1.log + has_hash2 = log2.log is not None and "input_hash" in log2.log + if has_hash1 != has_hash2: + raise ValueError( + f"Input hash presence inconsistent for triton kernel {call_type}[{op_name}] " + f"at index {i}: log1 has input_hash={has_hash1}, log2 has input_hash={has_hash2}" + ) + + if has_hash1: + compare_op_hashes( + log1.log["input_hash"], # type: ignore[union-attr] + log2.log["input_hash"], + is_input=True, + ) + + return difference_info + + +def get_active_debug_mode() -> DebugMode | None: + # Fast path: if no DebugMode is active, skip the stack walk + if _ACTIVE_DEBUG_MODE_COUNT == 0: + return None + debug_mode = None + for mode in _get_current_dispatch_mode_stack(): + if isinstance(mode, DebugMode): + debug_mode = mode + break + return debug_mode diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0eea2b7cabcd116780c8511ec0ed5d2951fcb6ca --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_debug_mode/_utils.py @@ -0,0 +1,251 @@ +# mypy: allow-untyped-defs +""" +Utility functions for DebugMode: tensor formatting, hashing, stack traces, and hook runners. +""" + +import inspect +import os +import traceback +import weakref +from collections.abc import Callable +from typing import TYPE_CHECKING + +import torch +from torch._subclasses.fake_tensor import FakeTensor +from torch.fx.graph import _parse_stack_trace +from torch.utils._dtype_abbrs import dtype_abbrs +from torch.utils._pytree import tree_map +from torch.utils._traceback import CapturedTraceback +from torch.utils.weak import WeakIdRef + + +if TYPE_CHECKING: + from torch.utils._debug_mode._calls import _DebugCall + + +REDISTRIBUTE_FUNC = "redistribute_input" + +# Tracks if we're in inductor benchmarking, and temporarily disables logging +# (for ignoring autotuning kernel launches which don't affect the user-facing result) +_IN_INDUCTOR_BENCHMARK: bool = False +# For record_outputs, log_tensor_hashes hooks for triton kernels. +# Stores kernel outputs in call.record["output"] +_RECORD_TRITON_OUTPUTS: bool = False +# Annotates kernel output hashes, and stores them in call.post_hashes +_TRITON_OUTPUT_HASH_FN: Callable | None = None +# Annotates kernel input hashes, and stores them in call.pre_hashes +_TRITON_INPUT_HASH_FN: Callable | None = None + +# registered dispatch call hooks +_DISPATCH_RECORD_HOOKS: list[Callable] = [] +_DISPATCH_LOG_HOOKS: list[Callable] = [] +_DISPATCH_PRE_LOG_HOOKS: list[Callable] = [] + + +def _stringify_shape(shape) -> str: + return f"[{', '.join([str(x) for x in shape])}]" + + +def _stringify_device_mesh(mesh) -> str: + return f"DM({', '.join([str(s) for s in mesh.shape])})" + + +def _stringify_placement(placement) -> str: + return f"[{', '.join([str(p) for p in placement])}]" + + +def _stringify_attributes(tensor, attributes) -> str: + pairs = {} + for attr in attributes: + if hasattr(tensor, attr): + pairs[attr] = getattr(tensor, attr) + if len(pairs) == 0: + return "" + return f"{{{', '.join([f'{k}={v}' for k, v in pairs.items()])}}}" + + +def _stringify_dtensor_spec(spec) -> str: + from torch.distributed.tensor._dtensor_spec import DTensorSpec + + return DTensorSpec.format_shard_order_str(spec.placements, spec.shard_order) + + +class TensorIdTracker: + def __init__(self) -> None: + self.tensor_memo: dict[WeakIdRef, int] = {} + self.next_tensor_id = 0 + + def _id(self, tensor) -> int: + with torch._C._DisablePythonDispatcher(): + o = WeakIdRef(tensor) + + def del_memo() -> None: + self.tensor_memo.pop(o, None) + + weakref.finalize(tensor, del_memo) + if o not in self.tensor_memo: + self.tensor_memo[o] = self.next_tensor_id + self.next_tensor_id += 1 + return self.tensor_memo[o] + + +def _tensor_debug_string(tensor, attributes, tensor_memo=None) -> str: + """Convert tensor to debug string representation.""" + + if isinstance(tensor, torch.Tensor): + tensor_debug_str = f"{dtype_abbrs[tensor.dtype]}{_stringify_shape(tensor.shape)}{_stringify_attributes(tensor, attributes)}" + id_str = f"${tensor_memo._id(tensor)}" if tensor_memo is not None else "" + if isinstance(tensor, torch.distributed.tensor.DTensor): + # omitted device mesh + return f"dt{id_str}: {tensor_debug_str}| {_stringify_dtensor_spec(tensor._spec)}" + elif isinstance(tensor, FakeTensor): + return f"ft{id_str}: {tensor_debug_str}" + else: + return f"t{id_str}: {tensor_debug_str}" + else: + raise RuntimeError(f"Unsupported tensor type: {type(tensor)}") + + +def _arg_to_str(arg, attributes, tensor_memo=None) -> str: + from torch.distributed.tensor._dtensor_spec import DTensorSpec + + def to_str(x): + if isinstance(x, torch.Tensor): + return _tensor_debug_string(x, attributes, tensor_memo) + elif isinstance(x, DTensorSpec): + return _stringify_dtensor_spec(x) + return x + + arg = tree_map(to_str, arg) + return str(arg) + + +def norm_hash_fn(t: torch.Tensor, use_scalar: bool = False) -> torch.Tensor | float: + """ + from Observer. Computes a hash for a tensor by converting it to float (if needed), making it contiguous, + replacing NaN/inf values with fixed numbers, and then computing the L1 norm in float64 or complex128. + This is used to generate a deterministic summary value for tensor comparison. + """ + with torch._C._DisablePythonDispatcher(): + if not (t.is_floating_point() or t.is_complex()): + t = t.float() + t = t.contiguous() + + if t.is_complex(): + t_float = t.to(dtype=torch.complex128) + else: + t_float = t.to(dtype=torch.float64) + + out = t_float.norm(p=1) + if use_scalar: + return out.item() + return out + + +def _compute_rel_diff(hash1, hash2): + # Relative difference: |hash1 - hash2| / max(|hash1|, |hash2|, eps) + numerator = abs(hash1 - hash2) + denominator = max(abs(hash1), abs(hash2), 1e-10) + return numerator / denominator + + +def hash_tensor_fn(t: torch.Tensor, use_scalar: bool = False) -> torch.Tensor | int: + """ + wrapper over torch.hash_tensor + """ + if isinstance(t, torch.distributed.tensor.DTensor): + t = t.to_local() + + if t.is_floating_point(): + t_clean = t.to(dtype=torch.float64) + elif t.is_complex(): + t_clean = t.to(dtype=torch.complex128).view(torch.float64) + else: + t_clean = t.to(dtype=torch.int64) + + if t.numel() > 0: + out = torch.hash_tensor(t_clean) + else: + out = torch.zeros((), device=t_clean.device, dtype=torch.uint64) + + if use_scalar: + return out.item() # type: ignore[attribute] + return out + + +def _get_stack_trace() -> str: + from torch.fx.experimental.symbolic_shapes import uninteresting_files + + summary = CapturedTraceback.extract().summary() + summary = summary[:-4] # filter out DebugMode frames + summary = [ + frame for frame in summary if frame.filename not in uninteresting_files() + ] + summary = traceback.StackSummary.from_list(summary) + return "".join(summary.format()) + + +def _get_user_stack_trace(stack_trace_str: str) -> str | None: + # Extract user code stack trace, filtering out torch internals. + torch_dir = os.path.dirname(inspect.getfile(torch)) + filter_fn = lambda file, name, code: not file.startswith(torch_dir + os.path.sep) # noqa: E731 + trace = _parse_stack_trace(stack_trace_str, filter_fn=filter_fn) + if trace: + return f"File: {trace.file}:{trace.lineno} in {trace.name}, code: {trace.code}" + return None + + +def _maybe_get_autograd_trace() -> str | None: + if torch._C._current_autograd_node() is not None: + tb = torch._C._current_autograd_node().metadata.get("traceback_") # type: ignore[attr-defined] + if tb: + return "".join(tb) + return None + + +def _get_op_name(op) -> str: + if isinstance(op, torch._ops.OpOverload): + op_name = op.__qualname__ + elif hasattr(op, "__module__") and hasattr(op, "__name__"): + op_name = f"{op.__module__}.{op.__name__}" + else: + op_name = str(op) + return op_name + + +def _run_hook(hook, *args): + out = hook(*args) + if out is not None and not isinstance(out, dict): + raise AssertionError(f"hook must return None or dict, got {type(out).__name__}") + return out + + +def _run_dispatch_pre_log_hooks(call: "_DebugCall", func, types, args, kwargs) -> None: + if _DISPATCH_PRE_LOG_HOOKS: + for hook in _DISPATCH_PRE_LOG_HOOKS: + hook_out = _run_hook(hook, func, types, args, kwargs, call) + if hook_out is not None: + # Store pre-hook results in call.log + if call.log is None: + call.log = {} + call.log.update(hook_out) + + +def _run_dispatch_hooks(call: "_DebugCall", func, types, args, kwargs, result) -> None: + if _DISPATCH_RECORD_HOOKS: + record = {} + for hook in _DISPATCH_RECORD_HOOKS: + hook_out = _run_hook(hook, func, types, args, kwargs, result) + if hook_out is not None: + record.update(hook_out) + if record: + call.record = record + + if _DISPATCH_LOG_HOOKS: + # Preserve existing log from pre-hooks (e.g., input_hash) + if call.log is None: + call.log = {} + for hook in _DISPATCH_LOG_HOOKS: + hook_out = _run_hook(hook, func, types, args, kwargs, result) + if hook_out is not None: + call.log.update(hook_out) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_device.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_device.py new file mode 100644 index 0000000000000000000000000000000000000000..8b93bbef8193a339c62b8948998a17d28d8796c0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_device.py @@ -0,0 +1,137 @@ +# mypy: allow-untyped-defs +import functools + +import torch +from torch._C import _len_torch_function_stack +from torch.overrides import _pop_mode, _push_mode, TorchFunctionMode +from torch.utils._contextlib import context_decorator + + +CURRENT_DEVICE: torch.device | None = None + + +@functools.lru_cache(1) +def _device_constructors(): + return { + # standard ones + torch.empty, + torch.empty_permuted, + torch.empty_strided, + torch.empty_quantized, + torch.ones, + torch.arange, + torch.bartlett_window, + torch.blackman_window, + torch.eye, + torch.fft.fftfreq, + torch.fft.rfftfreq, + torch.full, + torch.hamming_window, + torch.hann_window, + torch.kaiser_window, + torch.linspace, + torch.logspace, + torch.nested.nested_tensor, + # This function doesn't actually take a device argument + # torch.normal, + torch.rand, + torch.randn, + torch.randint, + torch.randperm, + torch.range, + torch.sparse_coo_tensor, + torch.sparse_compressed_tensor, + torch.sparse_csr_tensor, + torch.sparse_csc_tensor, + torch.sparse_bsr_tensor, + torch.sparse_bsc_tensor, + torch.tril_indices, + torch.triu_indices, + torch.zeros, + torch.asarray, + # weird ones + torch.tensor, + torch.as_tensor, + torch.scalar_tensor, + # *_like may contain device kwarg, but the user implicitly + # expects a specific device even when kwarg unused. + # torch.zeros_like, + # torch.randint_like, + # torch.randn_like, + # torch.ones_like, + # torch.full_like, + # torch.empty_like, + } + + +# NB: This is directly called from C++ in torch/csrc/Device.cpp +class DeviceContext(TorchFunctionMode): + def __init__(self, device) -> None: + self.device = torch.device(device) + self.prev_mode: DeviceContext | None = None + + def __enter__(self): + global CURRENT_DEVICE + self.old_device = CURRENT_DEVICE + CURRENT_DEVICE = self.device + # We need to put the device at the bottom of the stack + # If we set default device within a function mode context + # exiting that context mode will pop the device function mode off + # of the stack incorrectly + cur_stack = [_pop_mode() for _ in range(_len_torch_function_stack())] + + _push_mode(self) + + for mode in reversed(cur_stack): + if isinstance(mode, DeviceContext): + self.prev_mode = mode + else: + _push_mode(mode) + + def __exit__(self, exc_type, exc_val, exc_tb): + global CURRENT_DEVICE + CURRENT_DEVICE = self.old_device + cur_stack = [] + # Invariant: there should only be one DeviceContext on the stack at any time + # (At the bottom), pop all modes until we hit the bottom, assert it's a DeviceContext + # or else someone else has popped it! + for _ in range(_len_torch_function_stack() - 1): + mode = _pop_mode() + if isinstance(mode, DeviceContext): + raise AssertionError( + "Found nested DeviceContext on the mode stack where none expected" + ) + cur_stack.append(mode) + + if _len_torch_function_stack() > 0: + mode = _pop_mode() + if not isinstance(mode, DeviceContext): + raise AssertionError( + "Expected a DeviceContext at the bottom of the mode stack" + ) + if self.prev_mode is not None: + _push_mode(self.prev_mode) + + for mode in reversed(cur_stack): + _push_mode(mode) + + def __torch_function__(self, func, types, args=(), kwargs=None): + kwargs = kwargs or {} + if func in _device_constructors() and kwargs.get("device") is None: + kwargs["device"] = self.device + return func(*args, **kwargs) + + +# NB: This is directly called from C++ in torch/csrc/Device.cpp +def device_decorator(device, func): + return context_decorator(lambda: device, func) + + +def set_device(device): + """ + Set the default device inside of the wrapped function by decorating it with this function. + + If you would like to use this as a context manager, use device as a + context manager directly, e.g., ``with torch.device(device)``. + """ + return lambda func: device_decorator(torch.device(device), func) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_dtype_abbrs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_dtype_abbrs.py new file mode 100644 index 0000000000000000000000000000000000000000..c4eb9c56671dba774aa09d27887330fc350311fd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_dtype_abbrs.py @@ -0,0 +1,30 @@ +import torch + + +# Used for testing and logging +dtype_abbrs = { + torch.bfloat16: "bf16", + torch.float64: "f64", + torch.float32: "f32", + torch.float16: "f16", + torch.float8_e4m3fn: "f8e4m3fn", + torch.float8_e5m2: "f8e5m2", + torch.float8_e4m3fnuz: "f8e4m3fnuz", + torch.float8_e5m2fnuz: "f8e5m2fnuz", + torch.float8_e8m0fnu: "f8e8m0fnu", + torch.float4_e2m1fn_x2: "f4e2m1fnx2", + torch.complex32: "c32", + torch.complex64: "c64", + torch.complex128: "c128", + torch.int8: "i8", + torch.int16: "i16", + torch.int32: "i32", + torch.int64: "i64", + torch.bool: "b8", + torch.uint8: "u8", + torch.uint16: "u16", + torch.uint32: "u32", + torch.uint64: "u64", + torch.bits16: "b16", + torch.bits1x8: "b1x8", +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_exposed_in.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_exposed_in.py new file mode 100644 index 0000000000000000000000000000000000000000..2cca4ce240ad99afd44b8347135300b8c48ae66f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_exposed_in.py @@ -0,0 +1,21 @@ +from collections.abc import Callable +from typing import TypeVar + + +F = TypeVar("F") + + +# Allows one to expose an API in a private submodule publicly as per the definition +# in PyTorch's public api policy. +# +# It is a temporary solution while we figure out if it should be the long-term solution +# or if we should amend PyTorch's public api policy. The concern is that this approach +# may not be very robust because it's not clear what __module__ is used for. +# However, both numpy and jax overwrite the __module__ attribute of their APIs +# without problem, so it seems fine. +def exposed_in(module: str) -> Callable[[F], F]: + def wrapper(fn: F) -> F: + fn.__module__ = module + return fn + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_filelock.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_filelock.py new file mode 100644 index 0000000000000000000000000000000000000000..a291f59b4ba7fc1d9c1ecca703a6ae1fdf8e4cf1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_filelock.py @@ -0,0 +1,41 @@ +from types import TracebackType +from typing_extensions import Self + +from filelock import FileLock as base_FileLock + +from torch.monitor import _WaitCounter + + +class FileLock(base_FileLock): + """ + This behaves like a normal file lock. + + However, it adds waitcounters for acquiring and releasing the filelock + as well as for the critical region within it. + + pytorch.filelock.enter - While we're acquiring the filelock. + pytorch.filelock.region - While we're holding the filelock and doing work. + pytorch.filelock.exit - While we're releasing the filelock. + """ + + def __enter__(self) -> Self: + self.region_counter = _WaitCounter("pytorch.filelock.region").guard() + with _WaitCounter("pytorch.filelock.enter").guard(): + result = super().__enter__() + self.region_counter.__enter__() + return result + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.region_counter.__exit__() + with _WaitCounter("pytorch.filelock.exit").guard(): + # Returns nothing per + # https://github.com/tox-dev/filelock/blob/57f488ff8fdc2193572efe102408fb63cfefe4e4/src/filelock/_api.py#L379 + super().__exit__(exc_type, exc_value, traceback) + # Returns nothing per + # https://github.com/pytorch/pytorch/blob/0f6bfc58a2cfb7a5c052bea618ab62becaf5c912/torch/csrc/monitor/python_init.cpp#L315 + return None diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_foreach_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_foreach_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e88720a93ce3fad081a9ab05c6f7042b603f0e9a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_foreach_utils.py @@ -0,0 +1,60 @@ +from typing import TypeAlias + +import torch +from torch import Tensor +from torch.autograd.grad_mode import no_grad + + +def _get_foreach_kernels_supported_devices() -> list[str]: + r"""Return the device type list that supports foreach kernels.""" + return ["cuda", "xpu", "mtia", torch._C._get_privateuse1_backend_name()] + + +def _get_fused_kernels_supported_devices() -> list[str]: + r"""Return the device type list that supports fused kernels in optimizer.""" + return [ + "mps", + "cuda", + "xpu", + "hpu", + "cpu", + "mtia", + torch._C._get_privateuse1_backend_name(), + ] + + +TensorListList: TypeAlias = list[list[Tensor | None]] +Indices: TypeAlias = list[int] +_foreach_supported_types = [torch.Tensor] + + +# This util function splits tensors into groups by device and dtype, which is useful before sending +# tensors off to a foreach implementation, which requires tensors to be on one device and dtype. +# If tensorlistlist contains more than one tensorlist, the following assumptions are made BUT NOT verified: +# - tensorlists CAN be None +# - all tensors in the first specified list cannot be None +# - given an index i, all specified tensorlist[i]s match in dtype and device +# with_indices (bool, optional): whether to track previous indices as the last list per dictionary entry. +# It comes in handy if there are Nones or literals in the tensorlists that are getting scattered out. +# Whereas mutating a tensor in the resulting split-up tensorlists WILL propagate changes back to the +# original input tensorlists, changing up Nones/literals WILL NOT propagate, and manual propagation +# may be necessary. Check out torch/optim/sgd.py for an example. +@no_grad() +def _group_tensors_by_device_and_dtype( + tensorlistlist: TensorListList, + with_indices: bool = False, +) -> dict[tuple[torch.device, torch.dtype], tuple[TensorListList, Indices]]: + return torch._C._group_tensors_by_device_and_dtype(tensorlistlist, with_indices) + + +def _device_has_foreach_support(device: torch.device) -> bool: + return ( + device.type in (_get_foreach_kernels_supported_devices() + ["cpu"]) + and not torch.jit.is_scripting() + ) + + +def _has_foreach_support(tensors: list[Tensor], device: torch.device) -> bool: + return _device_has_foreach_support(device) and all( + t is None or type(t) in _foreach_supported_types for t in tensors + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_functools.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_functools.py new file mode 100644 index 0000000000000000000000000000000000000000..215aa5103e4450b69e44ad009a5a4cb4a8296852 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_functools.py @@ -0,0 +1,46 @@ +import functools +from collections.abc import Callable +from typing import Concatenate, TypeVar +from typing_extensions import ParamSpec + + +_P = ParamSpec("_P") +_T = TypeVar("_T") +_C = TypeVar("_C") + +# Sentinel used to indicate that cache lookup failed. +_cache_sentinel = object() + + +def cache_method( + f: Callable[Concatenate[_C, _P], _T], +) -> Callable[Concatenate[_C, _P], _T]: + """ + Like `@functools.cache` but for methods. + + `@functools.cache` (and similarly `@functools.lru_cache`) shouldn't be used + on methods because it caches `self`, keeping it alive + forever. `@cache_method` ignores `self` so won't keep `self` alive (assuming + no cycles with `self` in the parameters). + + Footgun warning: This decorator completely ignores self's properties so only + use it when you know that self is frozen or won't change in a meaningful + way (such as the wrapped function being pure). + """ + cache_name = "_cache_method_" + f.__name__ + + @functools.wraps(f) + def wrap(self: _C, *args: _P.args, **kwargs: _P.kwargs) -> _T: + if kwargs: + raise AssertionError("cache_method does not accept keyword arguments") + if not (cache := getattr(self, cache_name, None)): + cache = {} + setattr(self, cache_name, cache) + cached_value = cache.get(args, _cache_sentinel) + if cached_value is not _cache_sentinel: + return cached_value + value = f(self, *args, **kwargs) + cache[args] = value + return value + + return wrap diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_get_clean_triton.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_get_clean_triton.py new file mode 100644 index 0000000000000000000000000000000000000000..a3f0de4f64aefc6ae9524f7d5a18dd19bdb8ac23 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_get_clean_triton.py @@ -0,0 +1,200 @@ +# mypy: allow-untyped-defs +import argparse +import os +import re +import subprocess +import sys +from pathlib import Path + + +def remove_triton_function_declaration(source_code: str) -> str: + remove_head = re.sub(r"(\n.+\s\'\'\'\n)", "\n", source_code) + remove_tail = re.sub(r"(\'\'\'\,.+)", "\n", remove_head) + return remove_tail + + +def remove_async_compile(source_code: str) -> str: + remove_top_level = str.replace(source_code, "async_compile = AsyncCompile()", "") + remove_compile = str.replace(remove_top_level, "async_compile.wait(globals())", "") + remove_del = str.replace(remove_compile, "del async_compile", "") + return remove_del + + +def rename_kernels(source_code: str) -> str: + pattern = r"(\w+)\s*=\s*async_compile\.triton\('triton_',\s" + triton_kernel_decl = "def triton_" + matches = [ + (match.end(), match.group(1)) + for match in re.finditer(pattern, source_code, re.DOTALL) + ] + + # Starting from the last match to avoid issues with shifting indices after replacements + for end_index, captured_string in reversed(matches): + # Find the index of the next "B" after the current match + index_of_B = source_code.find(triton_kernel_decl, end_index) + if index_of_B != -1: + # Replace the triton_kernel_decl with the captured string + source_code = ( + source_code[:index_of_B] + + f"def {captured_string}" + + source_code[index_of_B + len(triton_kernel_decl) :] + ) + else: + # If triton_kernel_decl is not found after the current match, continue to the next + continue + + return source_code + + +def merge_params(original_params: list[str], new_params: list[str]) -> list[str]: + for idx in range(len(new_params)): + if new_params[idx] == "T": + new_params[idx] = original_params[idx] + return new_params + + +def add_launch_params( + original: str, kernel_to_params: dict[str, tuple[str, str]] +) -> str: + # Regex to match the function call in the original string + pattern = r"(\w+)\.run\((.*)\)" + + def replace(match) -> str: + # Extract parts from the regex match + func_name = match.group(1) + params = match.group(2) + new_params, grid = kernel_to_params[func_name] + new_params = merge_params(params.split(", "), new_params.split(", ")) + + # Format the new function call + new_string = f"{func_name}[{grid}]({', '.join(new_params)})" + return new_string + + transformed = re.sub(pattern, replace, original) + + remove_inductor_wrappers = re.sub( + r"@triton_heuristics[^@]*@triton.jit", + r"@triton.jit", + transformed, + flags=re.DOTALL, + ) + + return remove_inductor_wrappers + + +def process_file( + input_filename: str, output_filename: str, auto_generate_params: bool = True +) -> str: + with open(input_filename) as file: + source_code = file.read() + + transformed_code = source_code + if "def triton_(" in source_code: + raise RuntimeError( + "Need to run original Pytorch code generating kernels with TORCHINDUCTOR_UNIQUE_KERNEL_NAMES=1" + ) + # transformed_code = rename_kernels(transformed_code) + transformed_code = remove_triton_function_declaration(transformed_code) + transformed_code = remove_async_compile(transformed_code) + + launch_params_filename = f"{input_filename}.launch_params" + + # Auto-generate launch_params if they don't exist and auto_generate_params is True + if not os.path.exists(launch_params_filename) and auto_generate_params: + print(f"Launch params file {launch_params_filename} not found. Generating...") + try: + # Set environment variable and run the input file + env = os.environ.copy() + env["TORCHINDUCTOR_DUMP_LAUNCH_PARAMS"] = "1" + + result = subprocess.run( + [sys.executable, input_filename], + env=env, + capture_output=True, + text=True, + cwd=os.path.dirname(input_filename) or ".", + ) + + if result.returncode != 0: + print(f"Error running {input_filename}:") + print(f"stdout: {result.stdout}") + print(f"stderr: {result.stderr}") + raise RuntimeError( + f"Failed to generate launch params. Command failed with return code {result.returncode}" + ) + + print(f"Successfully generated {launch_params_filename}") + + except Exception as e: + raise RuntimeError( + f"Failed to generate launch params by running {input_filename}: {str(e)}" + ) from e + + if not os.path.exists(launch_params_filename): + raise RuntimeError( + f"Missing {launch_params_filename}. Run `TORCHINDUCTOR_DUMP_LAUNCH_PARAMS=1 python {input_filename}` first." + ) + + with open(launch_params_filename) as f: + launch_params_meta = f.readlines() + + split_params = [i.split("|") for i in launch_params_meta] + kernel_args_grid = {a.strip(): (b.strip(), c.strip()) for a, b, c in split_params} + transformed_code = add_launch_params(transformed_code, kernel_args_grid) + + with open(output_filename, "w") as file: + file.write(transformed_code) + print(f"Successfully generated {output_filename}") + return transformed_code + + +def get_clean_triton( + input_path: Path, + output_path: Path = Path("triton_only_repro.py"), + auto_generate_params: bool = True, +): + """Run experiments and output results to file + + Args: + input_path (Path): Path to inductor generated output codede + output_path (Path): Path to write out the new python file + auto_generate_params (bool): Whether to automatically generate launch_params if missing + """ + return process_file(str(input_path), str(output_path), auto_generate_params) + + +if __name__ == "__main__": + """Sample usage: + # Running sweep + python _get_clean_triton.py output_code.py + + # To disable auto-generation of launch params: + python _get_clean_triton.py output_code.py --no-auto-generate + """ + parser = argparse.ArgumentParser( + description="Clean Inductor generated code to remove Inductor dependencies" + ) + + # Add the arguments + parser.add_argument( + "input_path", type=Path, help="Path to inductor generated output code" + ) + parser.add_argument( + "--output_path", + type=Path, + default=Path("triton_only_repro.py"), + help="Path to write out the clean triton output", + ) + parser.add_argument( + "--no-auto-generate", + action="store_true", + help="Disable automatic generation of launch_params file", + ) + + # Parse the arguments + args = parser.parse_args() + + # Call the function with parsed arguments + result = get_clean_triton( + args.input_path, args.output_path, not args.no_auto_generate + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_helion.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_helion.py new file mode 100644 index 0000000000000000000000000000000000000000..6d30832cf3f74158267cd82c21f31e5744022161 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_helion.py @@ -0,0 +1,17 @@ +import functools + +from torch.utils._triton import has_triton + + +@functools.cache +def has_helion_package() -> bool: + try: + import helion # type: ignore[import-untyped, import-not-found] # noqa: F401 + except ImportError: + return False + return True + + +@functools.cache +def has_helion() -> bool: + return has_helion_package() and has_triton() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_import_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_import_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..47e48fb7144e5d60f82f746b0dc862677f322552 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_import_utils.py @@ -0,0 +1,37 @@ +import functools +import importlib.util +from types import ModuleType + + +def _check_module_exists(name: str) -> bool: + r"""Returns if a top-level module with :attr:`name` exists *without** + importing it. This is generally safer than try-catch block around a + `import X`. It avoids third party libraries breaking assumptions of some of + our tests, e.g., setting multiprocessing start method when imported + (see librosa/#747, torchvision/#544). + """ + try: + spec = importlib.util.find_spec(name) + return spec is not None + except ImportError: + return False + + +@functools.lru_cache +def dill_available() -> bool: + return _check_module_exists("dill") + + +@functools.lru_cache +def import_dill() -> ModuleType | None: + if not dill_available(): + return None + + import dill + + # XXX: By default, dill writes the Pickler dispatch table to inject its + # own logic there. This globally affects the behavior of the standard library + # pickler for any user who transitively depends on this module! + # Undo this extension to avoid altering the behavior of the pickler globally. + dill.extend(use_dill=False) + return dill diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_inspect.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_inspect.py new file mode 100644 index 0000000000000000000000000000000000000000..67d29b9bc706935bf1c06dbbaff1029178fc137a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_inspect.py @@ -0,0 +1,90 @@ +import inspect +from typing import Any + + +def _signature_metadata( + sig: inspect.Signature, +) -> tuple[tuple[inspect.Parameter, ...], bool, int]: + """ + Returns tuple(sig.parameters.values()), if any has VAR_POSITIONAL or VAR_KEYWORD, and the max_positional + """ + params = tuple(sig.parameters.values()) + has_var_args = False + max_positional = 0 + + for p in params: + kind = p.kind + if kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD): + has_var_args = True + if kind in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ): + max_positional += 1 + + return params, has_var_args, max_positional + + +def _fast_bind( + sig: inspect.Signature, *args: Any, **kwargs: Any +) -> inspect.BoundArguments: + """ + Fast path for inspect.Signature.bind() for signatures without + VAR_POSITIONAL or VAR_KEYWORD parameters. Falls back to sig.bind() + for signatures that contain *args or **kwargs. + """ + params, has_var_args, max_positional = _signature_metadata(sig) + + # fallback for complex signatures + if has_var_args: + return sig.bind(*args, **kwargs) + + len_args = len(args) + + if len_args > max_positional: + raise TypeError( + f"Too many positional arguments: expected max {max_positional}, got {len_args}" + ) + + arguments: dict[str, Any] = {} + arg_i = 0 + + for p in params: + name = p.name + kind = p.kind + + if kind is inspect.Parameter.POSITIONAL_ONLY: + if name in kwargs: + raise TypeError( + f"Got some positional-only arguments passed as keyword arguments: '{name}'" + ) + if arg_i < len_args: + arguments[name] = args[arg_i] + arg_i += 1 + elif p.default is inspect.Parameter.empty: + raise TypeError(f"Missing required argument '{name}'") + + elif kind is inspect.Parameter.POSITIONAL_OR_KEYWORD: + if arg_i < len_args: + if name in kwargs: + raise TypeError(f"Multiple values for argument '{name}'") + arguments[name] = args[arg_i] + arg_i += 1 + elif name in kwargs: + arguments[name] = kwargs[name] + elif p.default is inspect.Parameter.empty: + raise TypeError(f"Missing required argument '{name}'") + + elif kind is inspect.Parameter.KEYWORD_ONLY: + if name in kwargs: + arguments[name] = kwargs[name] + elif p.default is inspect.Parameter.empty: + raise TypeError(f"Missing required argument '{name}'") + + # disallow extra keyword arguments not in the signature + # cause kwargs have been processed by sig.bind at the beginning + for name in kwargs: + if name not in sig.parameters: + raise TypeError(f"Got an unexpected keyword argument '{name}'") + + return inspect.BoundArguments(sig, arguments) # type: ignore[arg-type] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_mode_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_mode_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..b79b52b13449e829b6168bceeb07f254cb6c6180 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_mode_utils.py @@ -0,0 +1,15 @@ +# mypy: allow-untyped-defs +from typing import TypeVar + +import torch + + +T = TypeVar("T") + + +# returns if all are the same mode +def all_same_mode(modes): + return all(tuple(mode == modes[0] for mode in modes)) + + +no_dispatch = torch._C._DisableTorchDispatch diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_ordered_set.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_ordered_set.py new file mode 100644 index 0000000000000000000000000000000000000000..3d218395308f3fb9e4a6156d64d4d3032ed4c0bf --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_ordered_set.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +from collections.abc import ( + Hashable, + Iterable, + Iterator, + MutableSet, + Reversible, + Set as AbstractSet, +) +from typing import Any, cast, TypeVar + + +T = TypeVar("T", bound=Hashable) +T_co = TypeVar("T_co", bound=Hashable, covariant=True) + +__all__ = ["OrderedSet"] + + +class OrderedSet(MutableSet[T], Reversible[T]): + """ + Insertion ordered set, similar to OrderedDict. + """ + + __slots__ = ("_dict",) + + def __init__(self, iterable: Iterable[T] | None = None) -> None: + self._dict = dict.fromkeys(iterable, None) if iterable is not None else {} + + @staticmethod + def _from_dict(dict_inp: dict[T, None]) -> OrderedSet[T]: + s: OrderedSet[T] = OrderedSet() + s._dict = dict_inp + return s + + # + # Required overridden abstract methods + # + def __contains__(self, elem: object) -> bool: + return elem in self._dict + + def __iter__(self) -> Iterator[T]: + return iter(self._dict) + + def __len__(self) -> int: + return len(self._dict) + + def __reversed__(self) -> Iterator[T]: + return reversed(self._dict) + + def add(self, elem: T) -> None: + self._dict[elem] = None + + def discard(self, elem: T) -> None: + self._dict.pop(elem, None) + + def clear(self) -> None: + # overridden because MutableSet impl is slow + self._dict.clear() + + # Unimplemented set() methods in _collections_abc.MutableSet + + @classmethod + def _wrap_iter_in_set(cls, other: Any) -> Any: + """ + Wrap non-Set Iterables in OrderedSets + + Some of the magic methods are more strict on input types than + the public apis, so we need to wrap inputs in sets. + """ + + if not isinstance(other, AbstractSet) and isinstance(other, Iterable): + return cls(other) + else: + return other + + def pop(self) -> T: + if not self: + raise KeyError("pop from an empty set") + return self._dict.popitem()[0] + + def copy(self) -> OrderedSet[T]: + return OrderedSet._from_dict(self._dict.copy()) + + def difference(self, *others: Iterable[T]) -> OrderedSet[T]: + res = self.copy() + res.difference_update(*others) + return res + + def difference_update(self, *others: Iterable[T]) -> None: + for other in others: + self -= other # type: ignore[arg-type] + + def update(self, *others: Iterable[T]) -> None: + for other in others: + self |= other + + def intersection(self, *others: Iterable[T]) -> OrderedSet[T]: + res = self.copy() + for other in others: + if other is not self: + res &= other # type: ignore[arg-type] + return res + + def intersection_update(self, *others: Iterable[T]) -> None: + for other in others: + self &= other # type: ignore[arg-type] + + def issubset(self, other: Iterable[T]) -> bool: + return self <= self._wrap_iter_in_set(other) + + def issuperset(self, other: Iterable[T]) -> bool: + return self >= self._wrap_iter_in_set(other) + + def symmetric_difference(self, other: Iterable[T]) -> OrderedSet[T]: + return self ^ other # type: ignore[operator] + + def symmetric_difference_update(self, other: Iterable[T]) -> None: + self ^= other # type: ignore[arg-type] + + def union(self, *others: Iterable[T]) -> OrderedSet[T]: + res = self.copy() + for other in others: + if other is self: + continue + res |= other + return res + + # Specify here for correct type inference, otherwise would + # return AbstractSet[T] + def __sub__(self, other: AbstractSet[T_co]) -> OrderedSet[T]: + # following cpython set impl optimization + if isinstance(other, OrderedSet) and (len(self) * 4) > len(other): + out = self.copy() + out -= other + return out + return cast(OrderedSet[T], super().__sub__(other)) + + def __ior__(self, other: Iterable[T]) -> OrderedSet[T]: # type: ignore[misc, override] # noqa: PYI034 + if isinstance(other, OrderedSet): + self._dict.update(other._dict) + return self + return super().__ior__(other) # type: ignore[arg-type] + + def __eq__(self, other: object) -> bool: + if isinstance(other, OrderedSet): + return self._dict == other._dict + return super().__eq__(other) + + def __ne__(self, other: object) -> bool: + if isinstance(other, OrderedSet): + return self._dict != other._dict + return super().__ne__(other) + + def __or__(self, other: AbstractSet[T_co]) -> OrderedSet[T]: + return cast(OrderedSet[T], super().__or__(other)) + + def __and__(self, other: AbstractSet[T_co]) -> OrderedSet[T]: + # MutableSet impl will iterate over other, iter over smaller of two sets + if isinstance(other, OrderedSet) and len(self) < len(other): + # pyrefly: ignore [unsupported-operation, bad-return] + return other & self + return cast(OrderedSet[T], super().__and__(other)) + + def __xor__(self, other: AbstractSet[T_co]) -> OrderedSet[T]: + return cast(OrderedSet[T], super().__xor__(other)) + + def __repr__(self) -> str: + return f"{self.__class__.__name__}({list(self)})" + + def __getstate__(self) -> list[T]: + return list(self._dict.keys()) + + def __setstate__(self, state: list[T]) -> None: + self._dict = dict.fromkeys(state, None) + + def __reduce__(self) -> tuple[type[OrderedSet[T]], tuple[list[T]]]: + return (OrderedSet, (list(self),)) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_pallas.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_pallas.py new file mode 100644 index 0000000000000000000000000000000000000000..84e04d2835e7b2beaa1fc5d79530f2df81b76669 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_pallas.py @@ -0,0 +1,125 @@ +import functools + +import torch + + +@functools.cache +def has_jax_package() -> bool: + """Check if JAX is installed.""" + try: + import jax # noqa: F401 # type: ignore[import-not-found] + + return True + except ImportError: + return False + + +@functools.cache +def has_pallas_package() -> bool: + """Check if Pallas (JAX experimental) is available.""" + if not has_jax_package(): + return False + try: + from jax.experimental import ( # noqa: F401 # type: ignore[import-not-found] + pallas as pl, + ) + + return True + except ImportError: + return False + + +@functools.cache +def get_jax_version(fallback: tuple[int, int, int] = (0, 0, 0)) -> tuple[int, int, int]: + """Get JAX version as (major, minor, patch) tuple.""" + try: + import jax # type: ignore[import-not-found] + + version_parts = jax.__version__.split(".") + major, minor, patch = (int(v) for v in version_parts[:3]) + return (major, minor, patch) + except (ImportError, ValueError, AttributeError): + return fallback + + +@functools.cache +def has_jax_cuda_backend() -> bool: + """Check if JAX has CUDA backend support with SM90+ (required by Mosaic GPU).""" + if not has_jax_package(): + return False + try: + import jax # type: ignore[import-not-found] + + # Check if CUDA backend is available + devices = jax.devices("gpu") + if len(devices) == 0: + return False + + # Mosaic GPU requires SM90+ (compute capability 9.0+) + if torch.cuda.is_available(): + major, minor = torch.cuda.get_device_capability() + if major < 9: + return False + + return True + except Exception: + return False + + +@functools.cache +def has_jax_tpu_backend() -> bool: + """Check if JAX has TPU backend support.""" + if not has_jax_package(): + return False + try: + import jax # type: ignore[import-not-found] + + # Check if TPU backend is available + devices = jax.devices("tpu") + return len(devices) > 0 + except Exception: + return False + + +@functools.cache +def has_torch_tpu() -> bool: + """Check if torch_tpu is installed and available.""" + try: + import torch_tpu.api # noqa: F401 # type: ignore[import] + + # Verify hardware/runtime access + torch_tpu.api.tpu_device() + return True + except (ImportError, RuntimeError): + return False + + +@functools.cache +def has_cpu_pallas() -> bool: + """Checks for a full Pallas-on-CPU environment.""" + return has_pallas_package() + + +@functools.cache +def has_cuda_pallas() -> bool: + """Checks for a full Pallas-on-CUDA environment.""" + return has_pallas_package() and torch.cuda.is_available() and has_jax_cuda_backend() + + +@functools.cache +def has_tpu_pallas() -> bool: + """Checks for a full Pallas-on-TPU environment.""" + return has_pallas_package() and has_jax_tpu_backend() and has_torch_tpu() + + +@functools.cache +def has_pallas() -> bool: + """ + Check if Pallas backend is fully available for use. + + Requirements: + - JAX package installed + - Pallas (jax.experimental.pallas) available + - A compatible backend (CUDA or TPU) is available in both PyTorch and JAX. + """ + return has_cpu_pallas() or has_cuda_pallas() or has_tpu_pallas() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_python_dispatch.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_python_dispatch.py new file mode 100644 index 0000000000000000000000000000000000000000..9b07d32c950ddf20ba8a10984e7ed45bc86a1647 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_python_dispatch.py @@ -0,0 +1,965 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import contextlib +import functools +import warnings +from collections import deque +from dataclasses import dataclass +from typing import cast, overload, Protocol, TYPE_CHECKING +from typing_extensions import TypeIs + +import torch +import torchgen +import torchgen.model +from torch._C import ( + _get_dispatch_stack_at, + _len_torch_dispatch_stack, + _pop_torch_dispatch_stack, + _push_on_torch_dispatch_stack, + DispatchKey, +) +from torch._C._dynamo.guards import set_is_in_mode_without_ignore_compile_internals + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# TODO: Limitations and things about enable_torch_dispatch_mode we should fix before exposing it: +# - We need a better user-facing api for _DisableTorchDispatch that +# is able to selectively disable __torch_dispatch__ of a particular class. +# - It doesn't work with the tensor constructors (torch.tensor, torch.Tensor) +# - Better name (see https://github.com/pytorch/pytorch/pull/63496#discussion_r694091694) + +_is_in_torch_dispatch_mode = False +_is_in_non_infra_torch_dispatch_mode = False +# If inside any mode that has ignore_compile_internals() = False +_is_in_any_mode_without_ignore_compile_internals = False + + +def is_in_torch_dispatch_mode(include_infra_modes: bool = True) -> bool: + return ( + _is_in_torch_dispatch_mode + if include_infra_modes + else _is_in_non_infra_torch_dispatch_mode + ) + + +def is_in_any_mode_without_ignore_compile_internals() -> bool: + return _is_in_any_mode_without_ignore_compile_internals + + +def any_torch_dispatch_mode_on_stack() -> bool: + stack_len = torch._C._len_torch_dispatch_stack() + + for idx in range(stack_len): + mode = _get_dispatch_stack_at(idx) + + # Apply filters first + if mode.is_infra_mode(): + continue + + if mode.ignore_compile_internals(): + continue + + return True + return False + + +class TorchDispatchMode: + """ + A ``TorchDispatchMode`` allows you to override the meaning of all + ``__torch_dispatch__`` overridable functions within a dynamic scope, + without having to actually create a tensor subclass or manually + monkey-patch functions in the PyTorch API. Some common situations + where you should use a mode: + + * You want to override the meaning of factory functions, or other + functions that do not otherwise take a tensor as an argument + (these cannot be overridden with tensor subclasses). + + * You want to override the behavior of all functions without needing + to wrap your inputs in tensor subclasses; e.g., if you are just + interested in logging intermediate computations. + + * You want to control the order of execution of various tensor + subclasses explicitly, rather than implicitly via the return of + ``NotImplemented``. + + Independent subclasses of :class:`TorchDispatchMode` are compositional: + modes can be pushed onto a stack using ``with MyMode():``. + When you call functions in the PyTorch API inside your + ``__torch_dispatch__`` implementation, by default, they will forward on to + the next mode on the mode stack. If you want recursively call back into + your current ``__torch_dispatch__`` implementation, either explicitly + invoke ``self.__torch_dispatch__(...)``, or use the context manager + ``self`` to make PyTorch + API self-referential (beware of infinite loops, in this case!) + """ + + # - When False, custom torch dispatch mode will error out explicitly when a hop + # is called under the mode. + # - When True, custom torch dispatch mode's __torch_dispatch__ will be triggered. + # Mode authors can implement how the mode interacts with higher order operators. + supports_higher_order_operators = False + + def __init_subclass__(cls, **kwargs): + super().__init_subclass__(**kwargs) + if cls._should_skip_dynamo(): + if "__torch_dispatch__" in cls.__dict__: + raw = cls.__dict__["__torch_dispatch__"] + if not isinstance(raw, classmethod): + cls.__torch_dispatch__ = torch._disable_dynamo(raw, recursive=True) + + def __init__(self, _dispatch_key=None): + if _dispatch_key is not None: + if not isinstance(_dispatch_key, torch._C.DispatchKey): + raise AssertionError("_dispatch_key must be a torch._C.DispatchKey") + self.__dict__["_dispatch_key"] = _dispatch_key + + self.old_dispatch_mode_flags: deque[bool] = deque() + self.old_non_infra_dispatch_mode_flags: deque[bool] = deque() + self.old_without_ignore_compile_internals_dispatch_mode_flags: deque[bool] = ( + deque() + ) + + def _lazy_init_old_dispatch_mode_flags(self): + if not hasattr(self, "old_dispatch_mode_flags"): + self.old_dispatch_mode_flags: deque[bool] = deque() # type: ignore[no-redef] + + if not hasattr(self, "old_non_infra_dispatch_mode_flags"): + self.old_non_infra_dispatch_mode_flags: deque[bool] = deque() # type: ignore[no-redef] + + if not hasattr( + self, "old_without_ignore_compile_internals_dispatch_mode_flags" + ): + self.old_without_ignore_compile_internals_dispatch_mode_flags: deque[ # type: ignore[no-redef] + bool + ] = deque() + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + raise NotImplementedError + + def __enter__(self): + global _is_in_torch_dispatch_mode + global _is_in_non_infra_torch_dispatch_mode + global _is_in_any_mode_without_ignore_compile_internals + + # Previously, there wasn't any state in this class' constructor + # super calls were added to existing modes, but for any new modes + # this will replicate the previous behavior of not strictly needing + # to call super().__init__() + self._lazy_init_old_dispatch_mode_flags() + self.old_dispatch_mode_flags.append(_is_in_torch_dispatch_mode) + _is_in_torch_dispatch_mode = True + self.old_non_infra_dispatch_mode_flags.append( + _is_in_non_infra_torch_dispatch_mode + ) + _is_in_non_infra_torch_dispatch_mode = ( + _is_in_non_infra_torch_dispatch_mode or not self.is_infra_mode() + ) + self.old_without_ignore_compile_internals_dispatch_mode_flags.append( + _is_in_any_mode_without_ignore_compile_internals + ) + _is_in_any_mode_without_ignore_compile_internals = ( + _is_in_any_mode_without_ignore_compile_internals + or not self.ignore_compile_internals() + ) + set_is_in_mode_without_ignore_compile_internals( + _is_in_any_mode_without_ignore_compile_internals + ) + _push_mode(self) + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + mb_dk_or_mode_key = self.__dict__.get("_dispatch_key", None) + if mb_dk_or_mode_key is None: + # Today, mode keys are not used at all in the per-dispatch-key-mode logic (for pre-dispatch) + # We should probably revisit this. + mb_dk_or_mode_key = self.__dict__.get("_mode_key", None) + global _is_in_torch_dispatch_mode + _is_in_torch_dispatch_mode = self.old_dispatch_mode_flags.pop() + global _is_in_non_infra_torch_dispatch_mode + _is_in_non_infra_torch_dispatch_mode = ( + self.old_non_infra_dispatch_mode_flags.pop() + ) + global _is_in_any_mode_without_ignore_compile_internals + _is_in_any_mode_without_ignore_compile_internals = ( + self.old_without_ignore_compile_internals_dispatch_mode_flags.pop() + ) + set_is_in_mode_without_ignore_compile_internals( + _is_in_any_mode_without_ignore_compile_internals + ) + _pop_mode(mb_dk_or_mode_key) + + @classmethod + def push(cls, *args, **kwargs): + warnings.warn( + "`Mode.push()` is no longer necessary and can be replaced with just `with Mode()`", + stacklevel=2, + ) + instance = cls(*args, **kwargs) + return instance + + @classmethod + def is_infra_mode(cls) -> bool: + return False + + @classmethod + def _should_skip_dynamo(cls) -> bool: + """Skip Dynamo when the flag is set to True + + This is temporary measure to rollout a feature + that skips PT2 compilation inside __torch_dispatch__ + frames. + + If this flag is off, we would expect following: + + class YoloMode(TorchDispatchMode): + @classmethod + def _should_skip_dynamo(cls): + return False + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + return torch.ops.aten.mul.Tensor(args[0], args[1]) + + x = torch.ones(5) + with YoloMode(): + out = torch.compile(torch.add, backend=backend, fullgraph=True)(x, x) + + # instead of recursively disabling, we are compiling into __torch_dispatch__ + assert len(backend.graphs) == 1 + """ + return True + + @classmethod + def ignore_compile_internals(cls) -> bool: + """Ignore operators that are compiled via torch.compile. + + If ``True``, then this TorchDispatchMode ignores operators that + are optimized by :func:`torch.compile`. Mechanically, this involves + turning off the TorchDispatchMode throughout the whole compilation process, + and turning it back on for the runtime of the compiled artifact(s). + For example, + + @torch.compile + def f(x): + return x.sin().cos() + + with LoggingMode(): + f(x) + + The above example will not log anything if + ``LoggingMode.ignore_compile_internals()`` is True. + torch.compile will fuse sin() and cos() into a single operation + and this TorchDispatchMode will not be passed sin and cos. + + If ``False`` (default), :func:`torch.compile` will respect + the eager semantics of passing this TorchDispatchMode all + operators that would have run during eager execution. + The way this will usually happen is that :func:`torch.compile` + will just fallback to eager-mode PyTorch. + """ + if cls.is_infra_mode(): + return True + return False + + +def _get_current_dispatch_mode() -> TorchDispatchMode | None: + """ + Return the top user mode on the stack (the next one that would be + executed) if there are any. + """ + stack_len = _len_torch_dispatch_stack() + if stack_len > 0: + return _get_dispatch_stack_at(stack_len - 1) + return None + + +def _detect_infra_mode(key): + if key not in ( + torch._C._TorchDispatchModeKey.FUNCTIONAL, + torch._C._TorchDispatchModeKey.PROXY, + ): + raise AssertionError( + f"key must be either FUNCTIONAL ({torch._C._TorchDispatchModeKey.FUNCTIONAL}) \ + or PROXY ({torch._C._TorchDispatchModeKey.PROXY}) _TorchDispatchModeKey, \ + got {key}" + ) + from torch._ops import _get_dispatch_mode_pre_dispatch + + pre_dispatch_mode = _get_dispatch_mode_pre_dispatch(key) + post_dispatch_mode = torch._C._get_dispatch_mode(key) + + if pre_dispatch_mode is not None and post_dispatch_mode is not None: + raise AssertionError( + "At most one of pre_dispatch_mode and post_dispatch_mode may be active" + ) + + if pre_dispatch_mode is None: + return post_dispatch_mode + + return pre_dispatch_mode + + +def _unset_infra_mode(key): + from torch._ops import _get_dispatch_mode_pre_dispatch, unset_mode_pre_dispatch + + pre_dispatch_mode = _get_dispatch_mode_pre_dispatch(key) + post_dispatch_mode = torch._C._get_dispatch_mode(key) + if pre_dispatch_mode and post_dispatch_mode: + raise AssertionError( + "Can't have active infra mode on both pre and post dispatch mode stack" + ) + + if pre_dispatch_mode: + mode = unset_mode_pre_dispatch(key) + return mode + if post_dispatch_mode: + return torch._C._unset_dispatch_mode(key) + + +def _disable_infra_mode(key): + if key not in ( + torch._C._TorchDispatchModeKey.FUNCTIONAL, + torch._C._TorchDispatchModeKey.PROXY, + ): + raise AssertionError( + "key must be either FUNCTIONAL or PROXY _TorchDispatchModeKey" + ) + mode_unset = _unset_infra_mode(key) + try: + yield mode_unset + finally: + if mode_unset is not None: + _push_mode(mode_unset) + + +def _get_current_dispatch_mode_stack() -> list[TorchDispatchMode]: + """ + Returns the current stack of dispatch modes, with the most recent + (i.e., the one that will be processed first) at the end of the + list (standard stack convention). + """ + stack_len = _len_torch_dispatch_stack() + return [_get_dispatch_stack_at(i) for i in range(stack_len)] + + +def _push_mode(mode: TorchDispatchMode) -> None: + k = mode._dispatch_key if hasattr(mode, "_dispatch_key") else None + if k is not None and k != torch._C.DispatchKey.PreDispatch: + raise AssertionError( + "mode._dispatch_key must be None or DispatchKey.PreDispatch" + ) + if k is None: + _push_on_torch_dispatch_stack(mode) + return + + from torch._ops import _set_mode_pre_dispatch, get_cached_ops + + # See Note [Not Caching Per-Dispatch-Key Mode Handlers] + # Clear the cache of every op that has been used so far, for this particular key. + ks = torch._C._functionality_to_backend_keys(k) + for op in get_cached_ops(): + for key in ks: + op._uncache_dispatch(key) + _set_mode_pre_dispatch(mode) + + +def _pop_mode(k: DispatchKey | torch._C._TorchDispatchModeKey | None = None): + if k == torch._C.DispatchKey.PreDispatch: # type: ignore[attr-defined] + from torch._ops import _pop_mode_from_pre_dispatch + + return _pop_mode_from_pre_dispatch() + + if k is None or isinstance(k, torch._C._TorchDispatchModeKey): + return _pop_torch_dispatch_stack(k) + + +@contextlib.contextmanager +def _pop_mode_temporarily(k: DispatchKey | None = None): + old = _pop_mode(k) + try: + yield old + finally: + _push_mode(old) + + +@contextlib.contextmanager +def _disable_current_modes(): + from torch._ops import ( + _len_torch_dispatch_stack_pre_dispatch, + _pop_mode_from_pre_dispatch, + ) + from torch._subclasses.functional_tensor import FunctionalTensorMode + from torch._subclasses.schema_check_mode import SchemaCheckMode + from torch.fx.experimental.proxy_tensor import ProxyTorchDispatchMode + + mode_len_pre_dispatch = _len_torch_dispatch_stack_pre_dispatch() + old_pre_dispatch_modes = [ + _pop_mode_from_pre_dispatch() for _ in range(mode_len_pre_dispatch) + ] + + has_proxy_mode_in_pre_dispatch = False + has_functional_mode_in_pre_dispatch = False + has_schema_check_mode_in_pre_dispatch = False + + for i in old_pre_dispatch_modes: + if isinstance(i, ProxyTorchDispatchMode): + has_proxy_mode_in_pre_dispatch = True + if isinstance(i, FunctionalTensorMode): + has_functional_mode_in_pre_dispatch = True + if isinstance(i, SchemaCheckMode): + has_schema_check_mode_in_pre_dispatch = True + + mode_len = _len_torch_dispatch_stack() + old_modes = [_pop_mode() for _ in range(mode_len)] + + for old in old_modes: + if ( + isinstance(old, FunctionalTensorMode) + and has_functional_mode_in_pre_dispatch + ): + raise AssertionError( + "Can't have FunctionalMode available both in PreDispatch and Python Key" + ) + if isinstance(old, ProxyTorchDispatchMode) and has_proxy_mode_in_pre_dispatch: + raise AssertionError( + "Can't have ProxyTorchDispatchMode available both in PreDispatch and Python Key" + ) + if isinstance(old, SchemaCheckMode) and has_schema_check_mode_in_pre_dispatch: + raise AssertionError( + "Can't have SchemaCheckMode available both in PreDispatch and Python Key" + ) + + # Manually disable proxy and fake modes, if any are active + try: + yield old_pre_dispatch_modes + old_modes + finally: + for mode in reversed(old_modes): + _push_mode(mode) + for mode in reversed(old_pre_dispatch_modes): + _push_mode(mode) + + +class BaseTorchDispatchMode(TorchDispatchMode): + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if kwargs is None: + kwargs = {} + return func(*args, **kwargs) + + +# Subtypes which have __tensor_flatten__ and __tensor_unflatten__. +class TensorWithFlatten(Protocol): + def __tensor_flatten__(self) -> tuple[Sequence[str], object]: ... + + @staticmethod + def __tensor_unflatten__( + inner_tensors: int, flatten_spec: int, outer_size: int, outer_stride: int + ) -> torch.Tensor: ... + + # It would be really nice to be able to say that the return of + # is_traceable_wrapper_subclass() is Intersection[torch.Tensor, + # TensorWithFlatten] - but that doesn't exist. + + shape: torch._C.Size + + @overload + def stride(self, dim: None = None) -> tuple[int, ...]: ... + + @overload + def stride(self, dim: int) -> int: ... + + @overload + def size(self, dim: None = None) -> tuple[int, ...]: ... + + @overload + def size(self, dim: int) -> int: ... + + def storage_offset(self) -> int: ... + + def dim(self) -> int: ... + + @overload + def to( + self, + dtype: torch.types._dtype, + non_blocking: bool = False, + copy: bool = False, + *, + memory_format: torch.memory_format | None = None, + ) -> torch.Tensor: ... + + @overload + def to( + self, + device: torch._prims_common.DeviceLikeType | None = None, + dtype: torch.types._dtype | None = None, + non_blocking: bool = False, + copy: bool = False, + *, + memory_format: torch.memory_format | None = None, + ) -> torch.Tensor: ... + + @overload + def to( + self, + other: torch.Tensor, + non_blocking: bool = False, + copy: bool = False, + *, + memory_format: torch.memory_format | None = None, + ) -> torch.Tensor: ... + + +def is_traceable_wrapper_subclass(t: object) -> TypeIs[TensorWithFlatten]: + """ + Returns whether or not a tensor subclass that implements __torch_dispatch__ + is 'traceable' with torch.compile. + In order for a tensor subclass to support TorchDispatchMode-style tracing in PT2, + It must implement two magic methods: __tensor_flatten__ and __tensor_unflatten__. + It is also expected to obey some restrictions around traceability and aliasing: + * The subclass's __torch_dispatch__() implementation should desugar into pytorch + dispatcher operations that can be traced into a graph. + * The subclass should use return_and_correct_aliasing(). This is needed today to make + sure that torch.compile does the right thing in a few cases around input mutation + and output aliasing. + + Expected magic method signatures: + attrs, ctx = t.__tensor_flatten__() + attrs: list of attribute name strings for inner tensors + ctx: dict containing any other subclass-specific metadata needed for unflattening + + t = MySubClass.__tensor_unflatten__(inner_tensors, ctx, outer_size, outer_stride) + inner_tensors: dict mapping attribute name -> tensor for each inner tensor + ctx: dict with subclass metadata in the form that __tensor_flatten__() produces + outer_size: expected (possibly symbolic) size that the returned subclass + instance should have. Note that this arg is useful for certain subclasses + that require the shape info to be constructed. In most cases, this arg can be + safely ignored. + outer_stride: expected (possibly symbolic) stride that the returned subclass + instance should have. Note that this arg is useful for certain subclasses + that require the stride info to be constructed. In most cases, this arg can be + safely ignored. + """ + is_subclass = isinstance(t, torch.Tensor) and type(t) is not torch.Tensor + return ( + is_subclass + and hasattr(t, "__tensor_flatten__") + and hasattr(t, "__tensor_unflatten__") + ) + + +def is_traceable_wrapper_subclass_type(t: type) -> TypeIs[type[TensorWithFlatten]]: + """Same as above, but takes a type argument instead of an instance.""" + return ( + issubclass(t, torch.Tensor) + and t is not torch.Tensor + and hasattr(t, "__tensor_flatten__") + and hasattr(t, "__tensor_unflatten__") + ) + + +def transform_subclass(t, callback, outer_size=None, outer_stride=None): + """ + Given a traceable, wrapper tensor subclass ``t`` that implements + ``__torch_dispatch__`` and holds some inner tensors, + and a callback of type ``Callable[[str, torch.Tensor], torch.Tensor]``, + `transform_subclass` will construct a fresh instance of the wrapper tensor subclass. + It will do so by grabbing each inner tensor attribute from the wrapper, + passing them into ``callback`` to get a transformed tensor, + and putting each transformed tensor into the fresh tensor subclass instance. + + Note: this function will not handle ensuring that the fresh subclass + gets the same (autograd, and aliasing) metadata as the original tensor. + This is generally handled in other subsystems like AOTAutograd. + """ + outer_size = outer_size if outer_size is not None else t.size() + outer_stride = outer_stride if outer_stride is not None else t.stride() + + attrs, ctx = t.__tensor_flatten__() + transformed_tensors_dict = {} + for attr in attrs: + transformed_tensors_dict[attr] = callback(attr, getattr(t, attr)) + sub = type(t).__tensor_unflatten__( + transformed_tensors_dict, ctx, outer_size, outer_stride + ) + + # NB: Purposefully guard here to simplify the inner / outer symbols. + # Using sym_eq() for symbolic comparison can result in an expression that's too + # difficult to guard on, so we use == here. + if sub.shape != outer_size: + raise AssertionError( + f"Expected return value from {type(t)}__tensor_unflatten__() to have " + f"shape equal to {outer_size}, but got: {sub.shape}" + ) + if sub.stride() != outer_stride: + raise AssertionError( + f"Expected return value from {type(t)}__tensor_unflatten__() to have " + f"stride equal to {outer_stride}, but got: {sub.stride()}" + ) + + return sub + + +def _correct_storage_aliasing(func, schema_info, args, outs) -> None: + """ + Given: an OpOverload, a SchemaInfo (cached information from torchgen about schema), + and the inputs/outputs to the OpOverload, + this function checks to see if func is a view operator + (by checking if any of the outputs in the op's schema + are immutable aliases of inputs). + If so, this function manually aliases the storage of the output tensor + with its corresponding input tensor alias. + It does this by unsafely overwriting the storage field of the output tensor + to be the same storage as the input. + """ + if not isinstance(func, torch._ops.OpOverload): + raise AssertionError(f"func must be an OpOverload, got {type(args)}") + if not isinstance(args, tuple): + raise AssertionError(f"args must be a tuple, got {type(args)}") + if not isinstance(outs, (list, tuple)): + raise AssertionError(f"outs must be a list or tuple, got {type(args)}") + + def alias_non_inplace_storage(arg, ret) -> None: + # This is hopefully a reasonable assert: + # subclasses that rely on this API for output aliasing + # should always return wrapper tensor subclasses for us to manually alias. + # in theory if a subclass that needs this API wants to sometimes return + # plain tensors, we could remove the assert and just not perform the aliasing, + # but it seems safer to learn more about this case first. + # + # Performance note: This is all just to assert that the argument and result + # types match, checking that is cheaper than is_traceable_wrapper_subclass_type, + # and multiple returns are relatively unlikely, so just check up front! + arg_type = type(arg) + ret_type = type(ret) + if arg_type is not ret_type and ( + is_traceable_wrapper_subclass_type(arg_type) + or is_traceable_wrapper_subclass_type(ret_type) + ): + ret_list = ret if isinstance(ret, list) else [ret] + for r in ret_list: + if type(arg) is not type(r): + raise AssertionError( + f"Called {str(func)} with input of type {type(arg)}\n" + f"and output of type {type(ret)}. But expected types to match." + ) + # Need to call a non-dispatcher helper, because we explicitly do **not** + # want our subclass to intercept the set_() call. + # instead, our subclass should directly have its storage swapped out. + # we **explicitly** don't want to reset the sizes on ret, if the storage implies a size change. + # Why? + # The purpose of this API is *not* to change the size/strides of our output- we assume it's already correct. + # We just want to "fix up" the storage aliasing, without modifying or output's metadata. + # Example: out = inp.expand(inp.shape[0], inp.shape[0]) + # This requires swapping the storage of out to be the same as inp, + # but we do *not* want it to change the sizes/strides that were compute for out. + + if isinstance(ret, list): + for r in ret: + torch._functionalize_unsafe_set(r, arg) + else: + if not isinstance(ret, torch.Tensor): + raise AssertionError(f"expected torch.Tensor, got {type(ret)}") + torch._functionalize_unsafe_set(ret, arg) + + for arg_idx, return_idx in schema_info.read_only_alias_match_indexes: + alias_non_inplace_storage(args[arg_idx], outs[return_idx]) + + +def _get_write_alias(x) -> str | None: + alias_set = x.alias_set + if not alias_set or not x.is_write: + return None + # torchscript allows for complicated alias sets, but our dispatcher ops only really involve simple aliasing + if len(alias_set) != 1: + raise AssertionError("Expected alias_set to contain exactly one element") + # timeit says next(iter(alias_set)) is faster than list(alias_set)[0] even for + # set of size 1 on Python 3.13. + return next(iter(alias_set)) + + +# This abstracts over the fact that in return_and_correct_aliasing, +# we sometimes use torchgen schema parsing (for aten ops, since torchscript's schema parsing is sometimes buggy), +# and sometimes use torchscript schema parsing (for custom ops, for which torchgen parsing is untested). +@dataclass +class AliasInfo: + alias_set: set[str] + is_write: bool + name: str | None + + +@dataclass +class SchemaInfo: + args: list[AliasInfo] + outs: list[AliasInfo] + + is_inplace_view_op: bool + + # [_get_write_alias(x) for x in outs]. Guaranteed to contain no Nones; we coerce + # all-Nones result to empty list instead, and we don't support + # some-but-not-all-Nones. + outs_write_aliases: list[str] | None + + # List of (arg_idx, return_idx) where args[arg_idx].alias_set & + # outs[out_idx].alias_set is not empty, and not args[arg_idx].is_write. + read_only_alias_match_indexes: list[tuple[int, int]] + + +# Given an OpOverload, returns schema information on it. +# This is cached for efficiency, since it can involve running torchgen +@functools.cache +def get_alias_info(func) -> SchemaInfo: + # For ATen ops: use torchgen (since torchscript parser doesn't handle alias annotations + # properly for some ops that output tensorlists) + if func.namespace == "aten": + torchgen_schema_str = str(func._schema) + if not torchgen_schema_str.startswith("aten::"): + raise AssertionError( + "Expected torchgen schema string to start with 'aten::'" + ) + # remove the aten:: namespace, which is added by the torchscript parser, + # and torchgen doesn't know how to handle + torchgen_schema_str = torchgen_schema_str[6:] + import re + + # the torchscript parser ends up converting int[2]=1 into int[2]=[1, 1], + # which torchgen chokes on. + torchgen_schema_str = re.sub(r"=\[[0, ]+\]", "=0", torchgen_schema_str) + torchgen_schema_str = re.sub(r"=\[[1, ]+\]", "=1", torchgen_schema_str) + # for aten::rot90 / aten:fft_* + torchgen_schema_str = re.sub( + r"=\[(-?[0-9]+), (-?[0-9]+)\]", r"=[\1,\2]", torchgen_schema_str + ) + torchgen_schema = torchgen.model.FunctionSchema.parse(torchgen_schema_str) + arg_schemas = [ + AliasInfo( + alias_set=( + set() if a.annotation is None else set(a.annotation.alias_set) + ), + is_write=a.annotation is not None and a.annotation.is_write, + name=a.name, + ) + for a in torchgen_schema.arguments.flat_all + ] + out_schemas = [ + AliasInfo( + alias_set=( + set() if a.annotation is None else set(a.annotation.alias_set) + ), + is_write=a.annotation is not None and a.annotation.is_write, + name=a.name, + ) + for a in torchgen_schema.returns + ] + else: + # For non-aten ops, torchgen is untested so we rely on torchscript schema parsing + arg_schemas = [ + AliasInfo( + alias_set=( + set() if a.alias_info is None else set(a.alias_info.before_set) + ), + is_write=a.alias_info is not None and a.alias_info.is_write, + name=a.name, + ) + for a in func._schema.arguments + ] + out_schemas = [ + AliasInfo( + alias_set=( + set() if a.alias_info is None else set(a.alias_info.before_set) + ), + is_write=a.alias_info is not None and a.alias_info.is_write, + name=a.name, + ) + for a in func._schema.returns + ] + read_only_alias_match_indexes = [] + for arg_idx, schema_arg in enumerate(arg_schemas): + for return_idx, schema_out in enumerate(out_schemas): + is_read_only_alias_match = ( + schema_arg.alias_set & schema_out.alias_set + ) and not schema_arg.is_write + if is_read_only_alias_match: + read_only_alias_match_indexes.append((arg_idx, return_idx)) + + outs_write_aliases_list: list[str | None] = [ + _get_write_alias(r) for r in out_schemas + ] + non_nones = sum(x is not None for x in outs_write_aliases_list) + if non_nones == 0: + outs_write_aliases: list[str] | None = None + elif non_nones != len(outs_write_aliases_list): + # simplifying assumption: we don't have **any** ops with return types like "-> (Tensor(a!), Tensor)" + raise RuntimeError("Unsupported schema: " + str(func._schema)) + else: + outs_write_aliases = cast(list[str], outs_write_aliases_list) + + schema_info = SchemaInfo( + args=arg_schemas, + outs=out_schemas, + # This check is surprisingly expensive because pybind11 enum_s are + # inefficient. Just cache it. + is_inplace_view_op=torch.Tag.inplace_view in func.tags, + outs_write_aliases=outs_write_aliases, + read_only_alias_match_indexes=read_only_alias_match_indexes, + ) + return schema_info + + +def autograd_would_have_decomposed( + func: torch._ops.OpOverload, flat_args: Sequence[torch.Tensor | object] +) -> bool: + """ + Suppose that an operator has CompositeImplicitAutograd decomp registered. + Would autograd have used this decomposition? It will only use it if there + isn't an explicit backend registration for the device as well. This function + will tell if this would have occurred. + + Why do we need to apply these decompositions later? When inference mode is + on, the autograd key is bypassed entirely, so a lower level mode cannot rely + on the decomposition have been applied. It's easy to accidentally never apply + the decomposition, resulting in an operator showing up in a graph that + is unexpected. + + Why do we need to AVOID applying the decomposition when autograd wouldn't + have decomposed? If autograd doesn't decompose, this means in eager mode + we would have run the fused kernel. It must be possible to trace this + fused kernel directly into the graph for fidelity with eager (NB: a user + has the option of then further decomposing at proxy tensor mode via + decomposition table, but we must preserve it to proxy mode to have the + choice.) + + Why does functionalization need to also perform the test here? This is + because some CompositeImplicitAutograd decompositions are not functional. + If we are eventually going to decompose, we need to do this while we can + still turn functionalization back on, so those decompositions get functionalized. + So an early decomposition in functionalization may still be necessary. Note that + if proxy tensor decomposition process could turn functionalization back on, this + wouldn't be necessary, and maybe that is a useful thing to do anyway because + the decomposition table is user specified and a user could violate the functional + decomp requirement with a bad decomp. If this happened, then you could always + pass through functionalization. + """ + has_backend_registration = False + for a in flat_args: + if isinstance(a, torch.Tensor): + backend_key = torch._C._parse_dispatch_key( + torch._C._dispatch_key_for_device(a.device.type) + ) + if backend_key is None: + raise AssertionError( + f"failed to parse dispatch key for device {a.device.type}" + ) + # TODO: use func.has_kernel_for_dispatch_key(backend_key) + # but this one checks py_impl and CompositeImplicitAutograd + # incorrectly shows up as has backend reg here + has_backend_registration = torch._C._dispatch_has_kernel_for_dispatch_key( + func.name(), backend_key + ) + + # in theory we should take all backend keys and take the highest priority one + # to properly mimic the dispatcher, + # this just grabs the first tensor and takes its device key + break + return not has_backend_registration + + +def return_and_correct_aliasing(func, args, kwargs, out): + """ + This function should be used by wrapper tensor ``__torch_dispatch__`` subclasses + that would like to work with torch.compile. It ensures that the subclass + properly implements the aliasing behavior of every op, + which is needed for correctness in AOTAutograd. + This function will handle: + + * When we see a view op, we will alias the storages of any + input and output tensor subclasses + + * When we see an inplace or out= op, we will directly + return the corresponding input tensor, instead of returning + a (potentially) fresh output tensor. + """ + + # Caching here because torchgen parsing is definitely not fast, and this function is called + # once for every op in the graph during functionalization. + schema_info = get_alias_info(func) + + def get_arg_from_alias(output_alias, schema_info, args, kwargs): + new_args, new_kwargs = torch.fx.operator_schemas.normalize_function( # type: ignore[misc] + func, args=args, kwargs=kwargs + ) + + arg_indices = [ + i for i, a in enumerate(schema_info.args) if output_alias in a.alias_set + ] + # For any dispatcher op with an output alias, we expect it to map to exactly one alias in the schema's input arguments. + if len(arg_indices) != 1: + raise AssertionError( + "Expected exactly one argument index for the given output alias" + ) + idx = arg_indices[0] + arg_info = schema_info.args[idx] + if arg_info.name is not None and arg_info.name in new_kwargs: + return new_kwargs[arg_info.name] + return new_args[idx] + + # Fix up the storages of any outs so that they point to the same storage as the input, + # if func is a view op. + _correct_storage_aliasing( + func, schema_info, args, (out,) if not isinstance(out, tuple) else out + ) + + # For inplace_view ops in particular, we'll try hard to make sure that the wrapper subclass's + # metadata is set correctly. + if schema_info.is_inplace_view_op: + # no_dispatch() to make sure that we secretly change the metadata on the wrapper, + # but don't end up dispatching the op anywhere else. + mutated_args = [ + x + for i, x in enumerate(args) + if _get_write_alias(schema_info.args[i]) is not None + ] + # Assumption: we have a very small number of inplace_view ops that follow a strict schema: + # there is only a single argument that gets its metadata mutated. + if len(mutated_args) != 1: + raise AssertionError( + "expected exactly one mutated arg for inplace_view ops" + ) + # This check exists because we generally *do* want to update the metadata of any wrapper subclasses, + # but FunctionalTensor is special: it overrides all size/stride calls to plumb to the inner tensor. + # so we don't actually need to update the metadata (and attempting to do so causes errors) + from torch._subclasses.functional_tensor import FunctionalTensor + + if not isinstance(mutated_args[0], FunctionalTensor): + with torch.utils._mode_utils.no_dispatch(): + # See Note: [Fake Tensor Dispatch Keys] + # we're borrowing the way it modifies dispatch key TLS. + meta_in_tls = torch._C._meta_in_tls_dispatch_include() + torch._C._set_meta_in_tls_dispatch_include(True) + try: + func(*args, **kwargs) + finally: + torch._C._set_meta_in_tls_dispatch_include(meta_in_tls) + + # Next: we need to make sure to return inputs directly, if the output is a mutable alias (e.g. add_()). + + schema_info_outs_write_aliases = schema_info.outs_write_aliases + # simple case: none of our outputs have mutable aliases, so we can return the output as-is + if schema_info_outs_write_aliases is None: + return out + + if len(schema_info_outs_write_aliases) == 1: + return get_arg_from_alias( + schema_info_outs_write_aliases[0], schema_info, args, kwargs + ) + + # In the multi-return case, all aten ops return a tuple / list, so cast accordingly. + outs_to_return = type(out)( + [ + (get_arg_from_alias(write_alias, schema_info, args, kwargs)) + for write_alias in schema_info_outs_write_aliases + ] + ) + return outs_to_return diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_pytree.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_pytree.py new file mode 100644 index 0000000000000000000000000000000000000000..de90d98f77951c6bbd4b7a3582ecabfa514ce947 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_pytree.py @@ -0,0 +1,2264 @@ +# Owner(s): ["module: pytree"] + +""" +Contains utility functions for working with nested python data structures. + +A *pytree* is a Python nested data structure. It is a tree in the sense that +nodes are Python collections (e.g., list, tuple, dict) and the leaves are +Python values. Furthermore, a pytree should not contain reference cycles. + +pytrees are useful for working with nested collections of Tensors. For example, +one can use `tree_map` to map a function over all Tensors inside some nested +collection of Tensors and `tree_leaves` to get a flat list of all Tensors +inside some nested collection. pytrees are helpful for implementing nested +collection support for PyTorch APIs. + +This pytree implementation is not very performant due to Python overhead +To improve the performance we can move parts of the implementation to C++. +""" + +import dataclasses +import functools +import importlib +import importlib.metadata +import json +import logging +import sys +import threading +import types +import warnings +from collections import defaultdict, deque, namedtuple, OrderedDict +from collections.abc import Callable, Hashable, Iterable, Mapping, Sequence +from enum import Enum +from typing import ( + Any, + cast, + ClassVar, + Final, + Generic, + NoReturn, + overload, + Protocol, + TYPE_CHECKING, + TypeAlias, + TypeVar, +) +from typing_extensions import deprecated, NamedTuple, Self, TypeIs + +from torch.torch_version import TorchVersion as _TorchVersion + + +log = logging.getLogger(__name__) + + +if TYPE_CHECKING: + import torch.utils._cxx_pytree as cxx_pytree + + +__all__ = [ + "PyTree", + "Context", + "FlattenFn", + "UnflattenFn", + "DumpableContext", + "ToDumpableContextFn", + "FromDumpableContextFn", + "PyTreeSpec", + "TreeSpec", + "LeafSpec", + "keystr", + "key_get", + "register_pytree_node", + "tree_is_leaf", + "tree_flatten", + "tree_flatten_with_path", + "tree_unflatten", + "tree_iter", + "tree_leaves", + "tree_leaves_with_path", + "tree_structure", + "tree_map", + "tree_map_with_path", + "tree_map_", + "tree_map_only", + "tree_map_only_", + "tree_all", + "tree_any", + "tree_all_only", + "tree_any_only", + "treespec_dumps", + "treespec_loads", + "treespec_pprint", + "is_namedtuple", + "is_namedtuple_class", + "is_namedtuple_instance", + "is_structseq", + "is_structseq_class", + "is_structseq_instance", +] + + +T = TypeVar("T") +S = TypeVar("S") +U = TypeVar("U") +R = TypeVar("R") + + +DEFAULT_TREESPEC_SERIALIZATION_PROTOCOL = 1 +NO_SERIALIZED_TYPE_NAME_FOUND = "NO_SERIALIZED_TYPE_NAME_FOUND" + + +class KeyEntry(Protocol): + def __hash__(self) -> int: ... + + def __eq__(self, other: object) -> bool: ... + + def __str__(self) -> str: ... + + def get(self, parent: Any) -> Any: ... + + +class EnumEncoder(json.JSONEncoder): + def default(self, obj: object) -> str | dict[str, Any]: + if isinstance(obj, Enum): + return { + "__enum__": True, + "fqn": f"{obj.__class__.__module__}:{obj.__class__.__qualname__}", + "name": obj.name, + } + return cast(str, super().default(obj)) + + +Context = Any +PyTree = Any +FlattenFn = Callable[[PyTree], tuple[list[Any], Context]] +UnflattenFn = Callable[[Iterable[Any], Context], PyTree] +DumpableContext = Any # Any json dumpable text +ToDumpableContextFn = Callable[[Context], DumpableContext] +FromDumpableContextFn = Callable[[DumpableContext], Context] +ToStrFunc = Callable[["TreeSpec", list[str]], str] # deprecated +MaybeFromStrFunc = Callable[[str], tuple[Any, Context, str] | None] # deprecated +KeyPath = tuple[KeyEntry, ...] +FlattenWithKeysFn = Callable[[PyTree], tuple[list[tuple[KeyEntry, Any]], Any]] + +# Keep deprecated alias for backward compatibility +FlattenFunc = FlattenFn # deprecated +UnflattenFunc = UnflattenFn # deprecated +FlattenWithKeysFunc = FlattenWithKeysFn # deprecated + + +# A NodeDef holds two callables: +# - flatten_fn should take the collection and return a flat list of values. +# It can also return some context that is used in reconstructing the +# collection. +# - unflatten_fn should take a flat list of values and some context +# (returned by flatten_fn). It returns the collection by reconstructing +# it from the list and the context. +# - flatten_with_keys_fn, which is a callable that takes a +# pytree and returns a list of (keypath, value) pairs and a context. +class NodeDef(NamedTuple): + type: type[Any] + flatten_fn: FlattenFn + unflatten_fn: UnflattenFn + flatten_with_keys_fn: FlattenWithKeysFn | None + + +_NODE_REGISTRY_LOCK = threading.RLock() +SUPPORTED_NODES: dict[type[Any], NodeDef] = {} + + +# _SerializeNodeDef holds the following: +# - typ: the type of the node (e.g., "Dict", "List", etc) +# - serialized_type_name: the fully qualified name of the type, e.g. "collections.OrderedDict" +# - to_dumpable_context takes a TreeSpec, and returns a serialized string format of the +# context, and the version number +# - from_dumpable_context takes in a string representation of the context, and the +# version, and returns the deserialized context +class _SerializeNodeDef(NamedTuple): + typ: type[Any] + serialized_type_name: str + to_dumpable_context: ToDumpableContextFn | None + from_dumpable_context: FromDumpableContextFn | None + + +SUPPORTED_SERIALIZED_TYPES: dict[type[Any], _SerializeNodeDef] = {} +SERIALIZED_TYPE_TO_PYTHON_TYPE: dict[str, type[Any]] = {} + +# NB: we try really hard to not import _cxx_pytree (which depends on optree) +# as much as possible. This is for isolation: a user who is not using C++ pytree +# shouldn't pay for it, and it helps makes things like cpython upgrades easier. +_optree_minimum_version = _TorchVersion("0.13.0") +try: + _optree_version = importlib.metadata.version("optree") +except importlib.metadata.PackageNotFoundError: + # No optree package found + _cxx_pytree_dynamo_traceable = _cxx_pytree_exists = False + _optree_version = _TorchVersion("0.0.0a0") +else: + _optree_version = _TorchVersion(_optree_version) + if _optree_version < _optree_minimum_version: + # optree package less than our required minimum version. + # Pretend the optree package doesn't exist. + # NB: We will raise ImportError if the user directly tries to + # `import torch.utils._cxx_pytree` (look in that file for the check). + _cxx_pytree_dynamo_traceable = _cxx_pytree_exists = False + else: + _cxx_pytree_dynamo_traceable = _cxx_pytree_exists = True + +_cxx_pytree_imported = False +_cxx_pytree_pending_imports: list[Any] = [] + + +def register_pytree_node( + cls: type[Any], + flatten_fn: FlattenFn, + unflatten_fn: UnflattenFn, + *, + serialized_type_name: str | None = None, + to_dumpable_context: ToDumpableContextFn | None = None, + from_dumpable_context: FromDumpableContextFn | None = None, + flatten_with_keys_fn: FlattenWithKeysFn | None = None, +) -> None: + """Register a container-like type as pytree node. + + Note: + :func:`register_dataclass` is a simpler way of registering a container-like + type as a pytree node. + + Args: + cls: the type to register + flatten_fn: A callable that takes a pytree and returns a flattened + representation of the pytree and additional context to represent the + flattened pytree. + unflatten_fn: A callable that takes a flattened version of the pytree, + additional context, and returns an unflattened pytree. + serialized_type_name: A keyword argument used to specify the fully qualified + name used when serializing the tree spec. + to_dumpable_context: An optional keyword argument to custom specify how + to convert the context of the pytree to a custom json dumpable + representation. This is used for json serialization, which is being + used in torch.export right now. + from_dumpable_context: An optional keyword argument to custom specify how + to convert the custom json dumpable representation of the context + back to the original context. This is used for json deserialization, + which is being used in torch.export right now. + flatten_with_keys_fn: An optional keyword argument to specify how to + access each pytree leaf's keypath when flattening and tree-mapping. + Like ``flatten_fn``, but in place of a List[leaf], it should return + a List[(keypath, leaf)]. + """ + with _NODE_REGISTRY_LOCK: + if cls in SUPPORTED_NODES: + raise ValueError(f"{cls} is already registered as pytree node.") + + _private_register_pytree_node( + cls, + flatten_fn, + unflatten_fn, + serialized_type_name=serialized_type_name, + to_dumpable_context=to_dumpable_context, + from_dumpable_context=from_dumpable_context, + flatten_with_keys_fn=flatten_with_keys_fn, + ) + + if not _cxx_pytree_exists: + return + + if _cxx_pytree_imported: + import torch.utils._cxx_pytree as cxx_pytree + + cxx_pytree._private_register_pytree_node( + cls, + flatten_fn, + unflatten_fn, + serialized_type_name=serialized_type_name, + to_dumpable_context=to_dumpable_context, + from_dumpable_context=from_dumpable_context, + ) + else: + args = (cls, flatten_fn, unflatten_fn) + kwargs = { + "serialized_type_name": serialized_type_name, + "to_dumpable_context": to_dumpable_context, + "from_dumpable_context": from_dumpable_context, + } + _cxx_pytree_pending_imports.append((args, kwargs)) + + +def register_dataclass( + cls: type[Any], + *, + field_names: list[str] | None = None, + drop_field_names: list[str] | None = None, + serialized_type_name: str | None = None, +) -> None: + """ + Registers a type that has the semantics of a ``dataclasses.dataclass`` type + as a pytree node. + + This is a simpler API than :func:`register_pytree_node` for registering + a dataclass or a custom class with the semantics of a dataclass. + + Args: + cls: The python type to register. The class must have the semantics of a + dataclass; in particular, it must be constructed by passing the fields + in. + field_names (Optional[List[str]]): A list of field names that correspond + to the **non-constant data** in this class. This list must contain + all the fields that are used to initialize the class. This argument + is optional if ``cls`` is a dataclass, in which case the fields will + be taken from ``dataclasses.fields()``. + drop_field_names (Optional[List[str]]): A list of field names that + should not be included in the pytree. + serialized_type_name: A keyword argument used to specify the fully + qualified name used when serializing the tree spec. This is only + needed for serializing the treespec in torch.export. + + Example: + + >>> from torch import Tensor + >>> from dataclasses import dataclass + >>> import torch.utils._pytree as pytree + >>> + >>> @dataclass + >>> class Point: + >>> x: Tensor + >>> y: Tensor + >>> + >>> pytree.register_dataclass(Point) + >>> + >>> point = Point(torch.tensor(0), torch.tensor(1)) + >>> point = pytree.tree_map(lambda x: x + 1, point) + >>> assert torch.allclose(point.x, torch.tensor(1)) + >>> assert torch.allclose(point.y, torch.tensor(2)) + + """ + drop_field_names = drop_field_names or [] + + if not dataclasses.is_dataclass(cls): + if field_names is None: + raise ValueError( + "field_names must be specified with a list of all fields used to " + f"initialize {cls}, as it is not a dataclass." + ) + elif field_names is None: + field_names = [f.name for f in dataclasses.fields(cls) if f.init] + else: + dataclass_init_fields = {f.name for f in dataclasses.fields(cls) if f.init} + dataclass_init_fields.difference_update(drop_field_names) + + if dataclass_init_fields != set(field_names): + error_msg = "field_names does not include all dataclass fields.\n" + + if missing := dataclass_init_fields - set(field_names): + error_msg += ( + f"Missing fields in `field_names`: {missing}. If you want " + "to include these fields in the pytree, please add them " + "to `field_names`, otherwise please add them to " + "`drop_field_names`.\n" + ) + + if unexpected := set(field_names) - dataclass_init_fields: + error_msg += ( + f"Unexpected fields in `field_names`: {unexpected}. " + "Please remove these fields, or add them to `drop_field_names`.\n" + ) + + raise ValueError(error_msg) + + def _flatten_fn(obj: Any) -> tuple[list[Any], Context]: + flattened = [] + flat_names = [] + none_names = [] + for name in field_names: + val = getattr(obj, name) + if val is not None: + flattened.append(val) + flat_names.append(name) + else: + none_names.append(name) + return flattened, [flat_names, none_names] + + def _unflatten_fn(values: Iterable[Any], context: Context) -> Any: + flat_names, none_names = context + return cls( + **dict(zip(flat_names, values, strict=True)), **dict.fromkeys(none_names) + ) + + def _flatten_fn_with_keys(obj: Any) -> tuple[list[Any], Context]: + flattened, (flat_names, _none_names) = _flatten_fn(obj) # type: ignore[misc] + return [ + (GetAttrKey(k), v) for k, v in zip(flat_names, flattened, strict=True) + ], flat_names + + _private_register_pytree_node( + cls, + _flatten_fn, + _unflatten_fn, + serialized_type_name=serialized_type_name, + flatten_with_keys_fn=_flatten_fn_with_keys, + ) + + +CONSTANT_NODES: set[type] = set() + + +def register_constant(cls: type[Any]) -> None: + """Registers a type as a pytree node with no leaves. + + In a :func:`torch.compile` region, if instances of these types get passed to + :func:`torch._dynamo.nonstrict_trace`-ed function, they treated as a + constant (sometimes referred to as "static"): + + 1. if the instance object existed before the :func:`torch.compile` region, + we _assume_ no mutation will happen to it inside the :func:`torch.compile` + region, require that it has non-default `__eq__` and `__hash__` methods, and + we guard on the instance based on its `__eq__` method, i.e., if a new + instance fails to match any instances from the previous compilations, + :func:`torch.compile` will recompile the function using the new instance. + + 2. else if the instance object is created inside the :func:`torch.compile` + region, we currently don't support using it in a + :func:`torch._dynamo.nonstrict_trace`-ed function. + + In general, if your class holds Tensors or dynamic int/float/bool (values that + may change from run-to-run of a function being compiled), then you probably + do not want to register it as a constant. + + Otherwise if you want to pass instance of a class to a + :func:`torch._dynamo.nonstrict_trace`-ed function, but you either can't use + :func:`register_pytree_node` on the class, or the class is "constant" enough + that you don't want to bother using :func:`register_pytree_node`, you should + consider using this function. + + Args: + cls: the type to register as a constant. This type must be hashable. + + Example: + + >>> from dataclasses import dataclass + >>> import torch.utils._pytree as pytree + >>> + >>> @dataclass(frozen=True) + >>> class Config: + >>> norm: str + >>> + >>> pytree.register_constant(Config) + >>> + >>> config = Config("l2") + >>> values, spec = pytree.tree_flatten(config) + >>> assert len(values) == 0 + + """ + if cls.__eq__ is object.__eq__: # type: ignore[comparison-overlap] + raise TypeError( + "register_constant(cls) expects `cls` to have a non-default `__eq__` implementation." + ) + + # Class with a custom `__eq__` without `__hash__` won't inherit the default + # `__hash__` from object; see https://stackoverflow.com/a/1608907. + if cls.__hash__ is None: # type: ignore[comparison-overlap] + raise TypeError( + "register_constant(cls) expects `cls` to have a non-default `__hash__` implementation." + ) + + def _flatten(x): # type: ignore[no-untyped-def] + return [], ConstantNode(x) + + def _unflatten(_, context): # type: ignore[no-untyped-def] + return context.value + + def _flatten_with_keys(x): # type: ignore[no-untyped-def] + return [], ConstantNode(x) + + with _NODE_REGISTRY_LOCK: + _private_register_pytree_node( + cls, + _flatten, + _unflatten, + flatten_with_keys_fn=_flatten_with_keys, + ) + CONSTANT_NODES.add(cls) + + +def is_constant_class(cls: type[Any]) -> bool: + return isinstance(cls, type) and cls in CONSTANT_NODES + + +@dataclasses.dataclass(frozen=True, slots=True) +class ConstantNode(Generic[T]): + value: T + + +def _is_constant_holder(spec: "TreeSpec") -> bool: + """Checks if the spec is from a pytree registered with register_constant""" + return isinstance(spec._context, ConstantNode) + + +def _retrieve_constant(spec: "TreeSpec") -> Any: + """Given a spec from a pytree registered with register_constant, retrieves the constant""" + if not _is_constant_holder(spec): + raise AssertionError("spec does not correspond to a registered constant pytree") + return tree_unflatten([], spec) + + +def _register_namedtuple( + cls: type[Any], + *, + serialized_type_name: str, +) -> None: + """ + Registers a namedtuple as a valid pytree node. By default namedtuples are + valid pytree nodes, but they are not serializable. This API provides the + argument `serialized_type_name` which allows these namedtuples to be + serialized. + + Args: + cls: the dataclass type to register + serialized_type_name: The serialized name for the dataclass. This is + required if you want to serialize the pytree TreeSpec containing this + namedtuple. + """ + _private_register_pytree_node( + cls, + _namedtuple_flatten, + _namedtuple_unflatten, + serialized_type_name=serialized_type_name, + to_dumpable_context=_namedtuple_serialize, + from_dumpable_context=_namedtuple_deserialize, + flatten_with_keys_fn=_namedtuple_flatten_with_keys, + ) + + +@deprecated( + "`torch.utils._pytree._register_pytree_node` is deprecated. " + "Please use `torch.utils._pytree.register_pytree_node` instead.", + category=FutureWarning, +) +def _register_pytree_node( + cls: type[Any], + flatten_fn: FlattenFn, + unflatten_fn: UnflattenFn, + to_str_fn: ToStrFunc | None = None, # deprecated + maybe_from_str_fn: MaybeFromStrFunc | None = None, # deprecated + *, + serialized_type_name: str | None = None, + to_dumpable_context: ToDumpableContextFn | None = None, + from_dumpable_context: FromDumpableContextFn | None = None, + flatten_with_keys_fn: FlattenWithKeysFn | None = None, +) -> None: + """Register a container-like type as pytree node for the Python pytree only. + + Args: + cls: the type to register + flatten_fn: A callable that takes a pytree and returns a flattened + representation of the pytree and additional context to represent the + flattened pytree. + unflatten_fn: A callable that takes a flattened version of the pytree, + additional context, and returns an unflattened pytree. + serialized_type_name: A keyword argument used to specify the fully qualified + name used when serializing the tree spec. + to_dumpable_context: An optional keyword argument to custom specify how + to convert the context of the pytree to a custom json dumpable + representation. This is used for json serialization, which is being + used in torch.export right now. + from_dumpable_context: An optional keyword argument to custom specify how + to convert the custom json dumpable representation of the context + back to the original context. This is used for json deserialization, + which is being used in torch.export right now. + flatten_with_keys_fn: An optional keyword argument to specify how to + access each pytree leaf's keypath when flattening and tree-mapping. + Like ``flatten_fn``, but in place of a List[leaf], it should return + a List[(keypath, leaf)]. + """ + if to_str_fn is not None or maybe_from_str_fn is not None: + warnings.warn( + "`to_str_fn` and `maybe_from_str_fn` is deprecated. " + "Please use `to_dumpable_context` and `from_dumpable_context` instead.", + FutureWarning, + stacklevel=2, + ) + + _private_register_pytree_node( + cls, + flatten_fn, + unflatten_fn, + serialized_type_name=serialized_type_name, + to_dumpable_context=to_dumpable_context, + from_dumpable_context=from_dumpable_context, + flatten_with_keys_fn=flatten_with_keys_fn, + ) + + +def _deregister_pytree_node( + cls: type[Any], +) -> None: + """This is an internal function that is used to deregister a pytree node type + for the Python pytree only. This should be only used inside PyTorch. + """ + with _NODE_REGISTRY_LOCK: + del SUPPORTED_NODES[cls] + node_def = SUPPORTED_SERIALIZED_TYPES[cls] + del SERIALIZED_TYPE_TO_PYTHON_TYPE[node_def.serialized_type_name] + del SUPPORTED_SERIALIZED_TYPES[cls] + CONSTANT_NODES.discard(cls) + + +def _private_register_pytree_node( + cls: type[Any], + flatten_fn: FlattenFn, + unflatten_fn: UnflattenFn, + *, + serialized_type_name: str | None = None, + to_dumpable_context: ToDumpableContextFn | None = None, + from_dumpable_context: FromDumpableContextFn | None = None, + flatten_with_keys_fn: FlattenWithKeysFn | None = None, +) -> None: + """This is an internal function that is used to register a pytree node type + for the Python pytree only. End-users should use :func:`register_pytree_node` + instead. + """ + from torch._library.opaque_object import is_opaque_type + + if isinstance(cls, type) and is_opaque_type(cls): + # TODO: remove this allowance once downstream callers stop calling + # register_constant on Enum subclasses. Enums are now natively + # supported as opaque value types and don't need pytree registration. + import enum + + if issubclass(cls, enum.Enum): + log.warning( + "%s is an Enum subclass and is now natively supported by " + "torch.compile as an opaque value type. Calling " + "register_constant() on Enum subclasses is deprecated and " + "will be an error in a future release.", + cls, + ) + else: + raise ValueError( + f"{cls} cannot be registered as a pytree as it has been " + "registered as an opaque object. Opaque objects must be pytree leaves." + ) + + with _NODE_REGISTRY_LOCK: + if cls in SUPPORTED_NODES: + # TODO: change this warning to an error after OSS/internal stabilize + warnings.warn( + f"{cls} is already registered as pytree node. " + "Overwriting the previous registration.", + stacklevel=2, + ) + + node_def = NodeDef(cls, flatten_fn, unflatten_fn, flatten_with_keys_fn) + SUPPORTED_NODES[cls] = node_def + + if (to_dumpable_context is None) ^ (from_dumpable_context is None): + raise ValueError( + f"Both to_dumpable_context and from_dumpable_context for {cls} must " + "be None or registered." + ) + + if serialized_type_name is None: + serialized_type_name = NO_SERIALIZED_TYPE_NAME_FOUND + + serialize_node_def = _SerializeNodeDef( + cls, + serialized_type_name, + to_dumpable_context, + from_dumpable_context, + ) + SUPPORTED_SERIALIZED_TYPES[cls] = serialize_node_def + SERIALIZED_TYPE_TO_PYTHON_TYPE[serialized_type_name] = cls + + +@dataclasses.dataclass(frozen=True, slots=True) +class SequenceKey(Generic[T]): + idx: int + + def __str__(self) -> str: + return f"[{self.idx!r}]" + + def get(self, sequence: Sequence[T]) -> T: + return sequence[self.idx] + + +K = TypeVar("K", bound=Hashable) + + +@dataclasses.dataclass(frozen=True, slots=True) +class MappingKey(Generic[K, T]): + key: K + + def __str__(self) -> str: + return f"[{self.key!r}]" + + def get(self, mapping: Mapping[K, T]) -> T: + return mapping[self.key] + + +@dataclasses.dataclass(frozen=True, slots=True) +class GetAttrKey: + name: str + + def __str__(self) -> str: + return f".{self.name}" + + def get(self, obj: Any) -> Any: + return getattr(obj, self.name) + + +# Reference: https://github.com/metaopt/optree/blob/main/optree/typing.py +def is_namedtuple(obj: object | type) -> bool: + """Return whether the object is an instance of namedtuple or a subclass of namedtuple.""" + cls = obj if isinstance(obj, type) else type(obj) + return is_namedtuple_class(cls) + + +# Reference: https://github.com/metaopt/optree/blob/main/optree/typing.py +def is_namedtuple_class(cls: type) -> bool: + """Return whether the class is a subclass of namedtuple.""" + return ( + isinstance(cls, type) + and issubclass(cls, tuple) + and isinstance(getattr(cls, "_fields", None), tuple) + and all(type(field) is str for field in cls._fields) # type: ignore[attr-defined] + and callable(getattr(cls, "_make", None)) + and callable(getattr(cls, "_asdict", None)) + ) + + +# Reference: https://github.com/metaopt/optree/blob/main/optree/typing.py +def is_namedtuple_instance(obj: object) -> bool: + """Return whether the object is an instance of namedtuple.""" + return is_namedtuple_class(type(obj)) + + +_T_co = TypeVar("_T_co", covariant=True) + + +# Reference: https://github.com/metaopt/optree/blob/main/optree/typing.py +class structseq(tuple[_T_co, ...]): + """A generic type stub for CPython's ``PyStructSequence`` type.""" + + __slots__: ClassVar[tuple[()]] = () + + n_fields: Final[int] # type: ignore[misc] + n_sequence_fields: Final[int] # type: ignore[misc] + n_unnamed_fields: Final[int] # type: ignore[misc] + + def __init_subclass__(cls) -> NoReturn: + """Prohibit subclassing.""" + raise TypeError("type 'structseq' is not an acceptable base type") + + def __new__( + cls: type[Self], + sequence: Iterable[_T_co], + # pyrefly: ignore [bad-function-definition] + dict: dict[str, Any] = ..., + ) -> Self: + raise NotImplementedError + + +# Reference: https://github.com/metaopt/optree/blob/main/optree/typing.py +def is_structseq(obj: object | type) -> bool: + """Return whether the object is an instance of PyStructSequence or a class of PyStructSequence.""" + cls = obj if isinstance(obj, type) else type(obj) + return is_structseq_class(cls) + + +# Set if the type allows subclassing (see CPython's Include/object.h) +Py_TPFLAGS_BASETYPE: int = 1 << 10 + + +# Reference: https://github.com/metaopt/optree/blob/main/optree/typing.py +def is_structseq_class(cls: type) -> bool: + """Return whether the class is a class of PyStructSequence.""" + return ( + isinstance(cls, type) + # Check direct inheritance from `tuple` rather than `issubclass(cls, tuple)` + and cls.__bases__ == (tuple,) + # Check PyStructSequence members + and isinstance(getattr(cls, "n_fields", None), int) + and isinstance(getattr(cls, "n_sequence_fields", None), int) + and isinstance(getattr(cls, "n_unnamed_fields", None), int) + # Check the type does not allow subclassing + and not bool(cls.__flags__ & Py_TPFLAGS_BASETYPE) # only works for CPython + ) + + +# Reference: https://github.com/metaopt/optree/blob/main/optree/typing.py +def is_structseq_instance(obj: object) -> bool: + """Return whether the object is an instance of PyStructSequence.""" + return is_structseq_class(type(obj)) + + +def _tuple_flatten(d: tuple[T, ...]) -> tuple[list[T], Context]: + return list(d), None + + +def _tuple_flatten_with_keys( + d: tuple[T, ...], +) -> tuple[list[tuple[KeyEntry, T]], Context]: + values, context = _tuple_flatten(d) + # pyrefly: ignore [bad-return] + return [(SequenceKey(i), v) for i, v in enumerate(values)], context + + +def _tuple_unflatten(values: Iterable[T], context: Context) -> tuple[T, ...]: + return tuple(values) + + +def _list_flatten(d: list[T]) -> tuple[list[T], Context]: + return d, None + + +def _list_flatten_with_keys(d: list[T]) -> tuple[list[tuple[KeyEntry, T]], Context]: + values, context = _list_flatten(d) + # pyrefly: ignore [bad-return] + return [(SequenceKey(i), v) for i, v in enumerate(values)], context + + +def _list_unflatten(values: Iterable[T], context: Context) -> list[T]: + return list(values) + + +def _dict_flatten(d: dict[Any, T]) -> tuple[list[T], Context]: + return list(d.values()), list(d.keys()) + + +def _dict_flatten_with_keys( + d: dict[Any, T], +) -> tuple[list[tuple[KeyEntry, T]], Context]: + values, context = _dict_flatten(d) + # pyrefly: ignore [bad-return] + return [(MappingKey(k), v) for k, v in zip(context, values, strict=True)], context + + +def _dict_unflatten(values: Iterable[T], context: Context) -> dict[Any, T]: + return dict(zip(context, values, strict=True)) + + +def _namedtuple_flatten(d: NamedTuple) -> tuple[list[Any], Context]: + return list(d), type(d) + + +def _namedtuple_flatten_with_keys( + d: NamedTuple, +) -> tuple[list[tuple[KeyEntry, Any]], Context]: + values, context = _namedtuple_flatten(d) + # pyrefly: ignore [bad-return] + return ( + [ + (GetAttrKey(field), v) + for field, v in zip(context._fields, values, strict=True) + ], + context, + ) + + +def _namedtuple_unflatten(values: Iterable[T], context: Context) -> NamedTuple: + return cast(NamedTuple, context(*values)) + + +def _namedtuple_serialize(context: Context) -> DumpableContext: + if context not in SUPPORTED_SERIALIZED_TYPES: + raise NotImplementedError( + f"Can't serialize TreeSpec of namedtuple class {context} because we " + "didn't register a serializated_type_name. Please register using " + "`_register_namedtuple`." + ) + + serialize_node_def = SUPPORTED_SERIALIZED_TYPES[context] + serialized_type_name = serialize_node_def.serialized_type_name + + if serialized_type_name == NO_SERIALIZED_TYPE_NAME_FOUND: + raise NotImplementedError( + f"Can't serialize TreeSpec of namedtuple class {context} because we " + "couldn't find a serializated_type_name. Please register using " + "`_register_namedtuple`." + ) + return serialized_type_name + + +def _namedtuple_deserialize(dumpable_context: DumpableContext) -> Context: + if dumpable_context not in SERIALIZED_TYPE_TO_PYTHON_TYPE: + raise NotImplementedError( + f"Can't deserialize TreeSpec of namedtuple class {dumpable_context} " + "because we couldn't find a serializated name." + ) + + typ = SERIALIZED_TYPE_TO_PYTHON_TYPE[dumpable_context] + return typ + + +def _ordereddict_flatten(d: OrderedDict[Any, T]) -> tuple[list[T], Context]: + return list(d.values()), list(d.keys()) + + +def _ordereddict_flatten_with_keys( + d: OrderedDict[Any, T], +) -> tuple[list[tuple[KeyEntry, T]], Context]: + values, context = _ordereddict_flatten(d) + # pyrefly: ignore [bad-return] + return [(MappingKey(k), v) for k, v in zip(context, values, strict=True)], context + + +def _ordereddict_unflatten( + values: Iterable[T], + context: Context, +) -> OrderedDict[Any, T]: + return OrderedDict((key, value) for key, value in zip(context, values, strict=True)) + + +_odict_flatten = _ordereddict_flatten +_odict_unflatten = _ordereddict_unflatten + + +def _defaultdict_flatten(d: defaultdict[Any, T]) -> tuple[list[T], Context]: + values, dict_context = _dict_flatten(d) + return values, [d.default_factory, dict_context] + + +def _defaultdict_flatten_with_keys( + d: defaultdict[Any, T], +) -> tuple[list[tuple[KeyEntry, T]], Context]: + values, context = _defaultdict_flatten(d) + _, dict_context = context + # pyrefly: ignore [bad-return] + return [ + (MappingKey(k), v) for k, v in zip(dict_context, values, strict=True) + ], context + + +def _defaultdict_unflatten( + values: Iterable[T], + context: Context, +) -> defaultdict[Any, T]: + default_factory, dict_context = context + return defaultdict(default_factory, _dict_unflatten(values, dict_context)) + + +def _defaultdict_serialize(context: Context) -> DumpableContext: + default_factory, dict_context = context + json_defaultdict = { + "default_factory_module": default_factory.__module__, + "default_factory_name": default_factory.__qualname__, + "dict_context": dict_context, + } + return json_defaultdict + + +def _defaultdict_deserialize(dumpable_context: DumpableContext) -> Context: + if not isinstance(dumpable_context, dict): + raise AssertionError("dumpable_context must be a dict") + + expected_keys = { + "default_factory_module", + "default_factory_name", + "dict_context", + } + if set(dumpable_context) != expected_keys: + raise AssertionError( + f"dumpable_context keys must be {expected_keys}, got {set(dumpable_context)}" + ) + + default_factory_module = dumpable_context["default_factory_module"] + default_factory_name = dumpable_context["default_factory_name"] + if not isinstance(default_factory_module, str): + raise AssertionError("default_factory_module must be a string") + if not isinstance(default_factory_name, str): + raise AssertionError("default_factory_name must be a string") + module = importlib.import_module(default_factory_module) + default_factory = getattr(module, default_factory_name) + + dict_context = dumpable_context["dict_context"] + return [default_factory, dict_context] + + +def _deque_flatten(d: deque[T]) -> tuple[list[T], Context]: + return list(d), d.maxlen + + +def _deque_flatten_with_keys( + d: deque[T], +) -> tuple[list[tuple[KeyEntry, T]], Context]: + values, context = _deque_flatten(d) + # pyrefly: ignore [bad-return] + return [(SequenceKey(i), v) for i, v in enumerate(values)], context + + +def _deque_unflatten(values: Iterable[T], context: Context) -> deque[T]: + return deque(values, maxlen=context) + + +_private_register_pytree_node( + tuple, + _tuple_flatten, + _tuple_unflatten, + serialized_type_name="builtins.tuple", + flatten_with_keys_fn=_tuple_flatten_with_keys, +) +_private_register_pytree_node( + list, + _list_flatten, + _list_unflatten, + serialized_type_name="builtins.list", + flatten_with_keys_fn=_list_flatten_with_keys, +) +_private_register_pytree_node( + dict, + _dict_flatten, + _dict_unflatten, + serialized_type_name="builtins.dict", + flatten_with_keys_fn=_dict_flatten_with_keys, +) +_private_register_pytree_node( + namedtuple, # type: ignore[arg-type] + _namedtuple_flatten, + _namedtuple_unflatten, + serialized_type_name="collections.namedtuple", + to_dumpable_context=_namedtuple_serialize, + from_dumpable_context=_namedtuple_deserialize, + flatten_with_keys_fn=_namedtuple_flatten_with_keys, +) +_private_register_pytree_node( + OrderedDict, + _ordereddict_flatten, + _ordereddict_unflatten, + serialized_type_name="collections.OrderedDict", + flatten_with_keys_fn=_ordereddict_flatten_with_keys, +) +_private_register_pytree_node( + defaultdict, + _defaultdict_flatten, + _defaultdict_unflatten, + serialized_type_name="collections.defaultdict", + to_dumpable_context=_defaultdict_serialize, + from_dumpable_context=_defaultdict_deserialize, + flatten_with_keys_fn=_defaultdict_flatten_with_keys, +) +_private_register_pytree_node( + deque, + _deque_flatten, + _deque_unflatten, + serialized_type_name="collections.deque", + flatten_with_keys_fn=_deque_flatten_with_keys, +) + + +STANDARD_DICT_TYPES: frozenset[type] = frozenset({dict, OrderedDict, defaultdict}) +# pyrefly: ignore [no-matching-overload] +BUILTIN_TYPES: frozenset[type] = frozenset( + # pyrefly: ignore [bad-argument-type] + { + tuple, + list, + dict, + namedtuple, # type: ignore[arg-type] + OrderedDict, + defaultdict, + deque, + }, +) + + +@deprecated( + "torch.utils._pytree._is_namedtuple_instance is private and will be removed in a future release. " + "Please use torch.utils._pytree.is_namedtuple_instance instead.", + category=FutureWarning, +) +def _is_namedtuple_instance(tree: Any) -> bool: + return is_namedtuple_instance(tree) + + +def _get_node_type(tree: Any) -> Any: + node_type = type(tree) + # All namedtuple types are implicitly registered as pytree nodes. + # XXX: Other parts of the codebase expect namedtuple types always return + # `namedtuple` instead of the actual namedtuple type. Even if the type + # is explicitly registered. + if is_namedtuple_class(node_type): + return namedtuple + return node_type + + +# A leaf is defined as anything that is not a Node. +def tree_is_leaf( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + """Check if a pytree is a leaf. + + >>> tree_is_leaf(1) + True + >>> tree_is_leaf(None) + True + >>> tree_is_leaf([1, 2, 3]) + False + >>> tree_is_leaf((1, 2, 3), is_leaf=lambda x: isinstance(x, tuple)) + True + >>> tree_is_leaf({"a": 1, "b": 2, "c": 3}) + False + >>> tree_is_leaf({"a": 1, "b": 2, "c": None}) + False + """ + if is_leaf is not None and is_leaf(tree): + return True + return _get_node_type(tree) not in SUPPORTED_NODES + + +@deprecated( + "torch.utils._pytree._is_leaf is private and will be removed in a future release. " + "Please use torch.utils._pytree.tree_is_leaf instead.", + category=FutureWarning, +) +def _is_leaf(tree: PyTree, is_leaf: Callable[[PyTree], bool] | None = None) -> bool: + return tree_is_leaf(tree, is_leaf=is_leaf) + + +# A TreeSpec represents the structure of a pytree. It holds: +# "type": the type of root Node of the pytree +# context: some context that is useful in unflattening the pytree +# children(): specs for each child of the root Node +# num_nodes: the total number of nodes +# num_leaves: the number of leaves +# num_children: the number of children of the root Node (i.e., len(children())) +# is_leaf(): whether the root Node is a leaf +@dataclasses.dataclass(init=False, frozen=True, eq=True, repr=False, slots=True) +class TreeSpec: + """Representing the structure of the pytree.""" + + type: Any + _context: Context + _children: list[Self] + + num_nodes: int = dataclasses.field(init=False) + num_leaves: int = dataclasses.field(init=False) + num_children: int = dataclasses.field(init=False) + + def __init__( + self, + type: Any, + context: Context, # keep for backward compatibility + children_specs: list[Self], # keep for backward compatibility + ) -> None: + object.__setattr__(self, "type", type) + object.__setattr__(self, "_context", context) + object.__setattr__(self, "_children", children_specs) + self.__post_init__() + + def __post_init__(self) -> None: + if self.type is None: + if self._context is not None: + raise AssertionError("leaf node should not have a context") + if len(self._children) != 0: + raise AssertionError("leaf node should not have children") + num_nodes = 1 + num_leaves = 1 + num_children = 0 + else: + num_nodes = 1 + num_leaves = 0 + for child in self._children: + num_nodes += child.num_nodes + num_leaves += child.num_leaves + num_children = len(self._children) + object.__setattr__(self, "num_nodes", num_nodes) + object.__setattr__(self, "num_leaves", num_leaves) + object.__setattr__(self, "num_children", num_children) + + def __repr__(self, indent: int = 0) -> str: + repr_prefix: str = f"TreeSpec({self.type.__name__}, {self._context}, [" + children_specs_str: str = "" + if self.num_children > 0: + indent += 2 + children_specs_str += self._children[0].__repr__(indent) + children_specs_str += "," if self.num_children > 1 else "" + children_specs_str += ",".join( + [ + "\n" + " " * indent + child.__repr__(indent) + for child in self._children[1:] + ] + ) + repr_suffix: str = f"{children_specs_str}])" + return repr_prefix + repr_suffix + + def __eq__(self, other: PyTree) -> bool: + if self is other: + return True + elif other.__class__ is self.__class__: + if str(self.type) != str(other.type): + return False + if self._context != other._context: + return False + elif self._children != other._children: + return False + return True + return NotImplemented + + @property + def context(self) -> Context: + return self._context + + @property + @deprecated( + "`treespec.children_specs` is deprecated. " + "Use `treespec.child(index)` to access a single child, " + "or `treespec.children()` to get all children.", + category=FutureWarning, + ) + def children_specs(self) -> list[Self]: + return self._children + + def is_leaf(self) -> bool: + """Test whether the treespec represents a leaf.""" + return self.num_nodes == 1 and self.num_leaves == 1 + + def children(self) -> list[Self]: + """Get all the child treespecs.""" + return self._children.copy() + + def child(self, index: int) -> Self: + """Get the child treespec at the given index.""" + return self._children[index] + + def flatten_up_to(self, tree: PyTree) -> list[PyTree]: + """Flatten the subtrees in ``tree`` up to the structure of this treespec and return a list of subtrees.""" + + def helper(treespec: TreeSpec, node: PyTree, subtrees: list[PyTree]) -> None: + if treespec.is_leaf(): + subtrees.append(node) + return + + node_type = _get_node_type(node) + if treespec.type not in BUILTIN_TYPES: + # Always require custom node types to match exactly + if node_type != treespec.type: + raise ValueError( + f"Type mismatch; " + f"expected {treespec.type!r}, but got {node_type!r}.", + ) + flatten_fn = SUPPORTED_NODES[node_type].flatten_fn + children, context = flatten_fn(node) + if len(children) != treespec.num_children: + raise ValueError( + f"Node arity mismatch; " + f"expected {treespec.num_children}, but got {len(children)}.", + ) + if context != treespec._context: + raise ValueError( + f"Node context mismatch for custom node type {treespec.type!r}.", + ) + else: + # For builtin dictionary types, we allow some flexibility + # Otherwise, we require exact matches + both_standard_dict = ( + treespec.type in STANDARD_DICT_TYPES + and node_type in STANDARD_DICT_TYPES + ) + if not both_standard_dict and node_type != treespec.type: + raise ValueError( + f"Node type mismatch; " + f"expected {treespec.type!r}, but got {node_type!r}.", + ) + if len(node) != treespec.num_children: + raise ValueError( + f"Node arity mismatch; " + f"expected {treespec.num_children}, but got {len(node)}.", + ) + + if both_standard_dict: + # dictionary types are compatible with each other + dict_context = ( + treespec._context + if treespec.type is not defaultdict + # ignore mismatch of `default_factory` for defaultdict + else treespec._context[1] + ) + expected_keys = dict_context + got_key_set = set(node) + expected_key_set = set(expected_keys) + if got_key_set != expected_key_set: + missing_keys = expected_key_set.difference(got_key_set) + extra_keys = got_key_set.difference(expected_key_set) + message = "" + if missing_keys: + message += f"; missing key(s): {missing_keys}" + if extra_keys: + message += f"; extra key(s): {extra_keys}" + raise ValueError(f"Node keys mismatch{message}.") + children = [node[key] for key in expected_keys] + else: + # node_type is treespec.type + flatten_fn = SUPPORTED_NODES[node_type].flatten_fn + children, context = flatten_fn(node) + if ( + node_type is not deque # ignore mismatch of `maxlen` for deque + ) and context != treespec._context: + raise ValueError( + f"Node context mismatch for node type {treespec.type!r}; " + f"expected {treespec._context!r}, but got {context!r}.", # namedtuple type mismatch + ) + + for subtree, subspec in zip(children, treespec._children, strict=True): + helper(subspec, subtree, subtrees) + + subtrees: list[PyTree] = [] + helper(self, tree, subtrees) + return subtrees + + def unflatten(self, leaves: Iterable[Any]) -> PyTree: + """Reconstruct a pytree from the leaves.""" + if not isinstance(leaves, (list, tuple)): + leaves = list(leaves) + if len(leaves) != self.num_leaves: + raise ValueError( + f"treespec.unflatten(leaves): `leaves` has length {len(leaves)} " + f"but the spec refers to a pytree that holds {self.num_leaves} " + f"items ({self}).", + ) + if self.is_leaf(): + return leaves[0] + + unflatten_fn = SUPPORTED_NODES[self.type].unflatten_fn + + # Recursively unflatten the children + start = 0 + end = 0 + child_pytrees = [] + for child_spec in self._children: + end += child_spec.num_leaves + child_pytrees.append(child_spec.unflatten(leaves[start:end])) + start = end + + return unflatten_fn(child_pytrees, self._context) + + def __hash__(self) -> int: + node_type = self.type + if node_type is defaultdict: + default_factory, dict_context = self._context + hashable_context = (default_factory, tuple(dict_context)) + elif node_type in (dict, OrderedDict): + hashable_context = tuple(self._context) + elif node_type is None or node_type in BUILTIN_TYPES: + hashable_context = self._context + elif isinstance(self._context, ConstantNode): + hashable_context = self._context.value + else: + # The context for user-defined node types might not be hashable. + # Ignore it for hashing. + # This does not break the correctness that equal objects imply the + # same hash. This might increase the hash collision rate, but we + # don't care about that. + hashable_context = None + return hash((node_type, hashable_context, tuple(self._children))) + + +PyTreeSpec: TypeAlias = TreeSpec + + +# NOTE: subclassing a dataclass is subtle. In order to enable reasoning about +# this class with `dataclasses.fields`, etc., while having a simplified +# constructor that takes no argument, we wrap with `dataclass(init=True, ...)` +# again, with fields that have `init=False`. +@deprecated( + "`isinstance(treespec, LeafSpec)` is deprecated, " + "use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead.", + category=FutureWarning, +) +@dataclasses.dataclass(init=False, frozen=True, eq=False, repr=False, slots=True) +class LeafSpec(TreeSpec): + type: Any = dataclasses.field(default=None, init=False) + _context: Context = dataclasses.field(default=None, init=False) + _children: list[Self] = dataclasses.field(default_factory=list, init=False) + + def __init__(self) -> None: + # On Python 3.10.0, the dataclass-generated __init__ for frozen+slots + # dataclasses doesn't populate init=False fields with simple defaults, + # causing copy.deepcopy to fail with AttributeError. Avoid the bug by + # providing a custom __init__ (init=False above) instead. + object.__setattr__(self, "type", None) + object.__setattr__(self, "_context", None) + object.__setattr__(self, "_children", []) + object.__setattr__(self, "num_nodes", 1) + object.__setattr__(self, "num_leaves", 1) + object.__setattr__(self, "num_children", 0) + + def __repr__(self, indent: int = 0) -> str: + return "*" + + +# All leaves are equivalent, so represent with a single object to save on +# object construction time +with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", category=FutureWarning, module=__name__, append=False + ) + _LEAF_SPEC = LeafSpec() + + +def treespec_leaf() -> LeafSpec: + """Make a treespec representing a leaf node.""" + return _LEAF_SPEC + + +def treespec_tuple(iterable: Iterable[TreeSpec] = (), /) -> TreeSpec: + """Make a tuple treespec from an iterable of child treespecs.""" + children = list(iterable) + if any(not isinstance(child, TreeSpec) for child in children): + raise ValueError(f"Expected a tuple of TreeSpec values, got: {children!r}.") + return TreeSpec(tuple, None, children) + + +def treespec_dict( + mapping: Mapping[Any, TreeSpec] | Iterable[tuple[Any, TreeSpec]] = (), + /, + **kwargs: TreeSpec, +) -> TreeSpec: + """Make a dict treespec from a dict of child treespecs.""" + dct = dict(mapping, **kwargs) + if any(not isinstance(child, TreeSpec) for child in dct.values()): + raise ValueError(f"Expected a dictionary of TreeSpec values, got: {dct!r}.") + return TreeSpec(dict, list(dct.keys()), list(dct.values())) + + +def _is_pytreespec_instance( + obj: Any, +) -> TypeIs["TreeSpec | cxx_pytree.PyTreeSpec"]: + if isinstance(obj, TreeSpec): + return True + if "torch.utils._cxx_pytree" in sys.modules: + # The C++ pytree module is not always available, so we check if it is loaded. + # If the C++ pytree module is loaded, we can check if the treespec + # is an instance of the C++ TreeSpec class. + import torch.utils._cxx_pytree as cxx_pytree + + if isinstance(obj, cxx_pytree.PyTreeSpec): + return True + if "torch._dynamo.polyfills.pytree" in sys.modules: + # The PyTorch Dynamo pytree module is not always available, so we check if it is loaded. + # If the PyTorch Dynamo pytree module is loaded, we can check if the treespec + # is an instance of the PyTorch Dynamo TreeSpec class. + import torch._dynamo.polyfills.pytree as dynamo_pytree + + return isinstance(obj, dynamo_pytree.PyTreeSpec) + return False + + +def _ensure_python_treespec_instance( + treespec: "TreeSpec | cxx_pytree.PyTreeSpec", +) -> TreeSpec: + if isinstance(treespec, TreeSpec): + return treespec + + if not _is_pytreespec_instance(treespec): + raise TypeError( + f"Expected `treespec` to be an instance of " + f"PyTreeSpec but got item of type {type(treespec)}." + ) + dummy_tree = treespec.unflatten([0] * treespec.num_leaves) + return tree_structure(dummy_tree) + + +def tree_flatten( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> tuple[list[Any], TreeSpec]: + """Flattens a pytree into a list of values and a TreeSpec that can be used + to reconstruct the pytree. + """ + + def helper(node: PyTree, leaves: list[Any]) -> TreeSpec: + if tree_is_leaf(node, is_leaf=is_leaf): + leaves.append(node) + return _LEAF_SPEC + + node_type = _get_node_type(node) + flatten_fn = SUPPORTED_NODES[node_type].flatten_fn + children, context = flatten_fn(node) + + # Recursively flatten the children + subspecs = [helper(child, leaves) for child in children] + return TreeSpec(node_type, context, subspecs) + + leaves: list[Any] = [] + treespec = helper(tree, leaves) + return leaves, treespec + + +def tree_unflatten(leaves: Iterable[Any], treespec: TreeSpec) -> PyTree: + """Given a list of values and a TreeSpec, builds a pytree. + This is the inverse operation of `tree_flatten`. + """ + if not _is_pytreespec_instance(treespec): + if not _is_pytreespec_instance(leaves): + raise TypeError( + f"Expected `treespec` to be an instance of " + f"PyTreeSpec but got item of type {type(treespec)}." + ) + # Allow passing the PyTreeSpec instance as the first argument + leaves, treespec = treespec, leaves + return treespec.unflatten(leaves) + + +def tree_iter( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> Iterable[Any]: + """Get an iterator over the leaves of a pytree.""" + if tree_is_leaf(tree, is_leaf=is_leaf): + yield tree + else: + node_type = _get_node_type(tree) + flatten_fn = SUPPORTED_NODES[node_type].flatten_fn + child_pytrees, _ = flatten_fn(tree) + + # Recursively flatten the children + for child in child_pytrees: + yield from tree_iter(child, is_leaf=is_leaf) + + +def tree_leaves( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> list[Any]: + """Get a list of leaves of a pytree.""" + return list(tree_iter(tree, is_leaf=is_leaf)) + + +def tree_structure( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> TreeSpec: + """Get the TreeSpec for a pytree.""" + return tree_flatten(tree, is_leaf=is_leaf)[1] + + +def tree_map( + func: Callable[..., Any], + tree: PyTree, + *rests: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + """Map a multi-input function over pytree args to produce a new pytree. + + See also :func:`tree_map_`. + + >>> tree_map(lambda x: x + 1, {"x": 7, "y": (42, 64)}) + {'x': 8, 'y': (43, 65)} + >>> tree_map(lambda x: x is None, {"x": 7, "y": (42, 64), "z": None}) + {'x': False, 'y': (False, False), 'z': True} + + If multiple inputs are given, the structure of the tree is taken from the first input; + subsequent inputs need only have ``tree`` as a prefix: + + >>> tree_map(lambda x, y: [x] + y, [5, 6], [[7, 9], [1, 2]]) + [[5, 7, 9], [6, 1, 2]] + + Args: + func (callable): A function that takes ``1 + len(rests)`` arguments, to be applied at the + corresponding leaves of the pytrees. + tree (pytree): A pytree to be mapped over, with each leaf providing the first positional + argument to function ``func``. + rests (tuple of pytree): A tuple of pytrees, each of which has the same structure as + ``tree`` or has ``tree`` as a prefix. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + A new pytree with the same structure as ``tree`` but with the value at each leaf given by + ``func(x, *xs)`` where ``x`` is the value at the corresponding leaf in ``tree`` and ``xs`` + is the tuple of values at corresponding nodes in ``rests``. + """ + leaves, treespec = tree_flatten(tree, is_leaf=is_leaf) + flat_args = [leaves] + [treespec.flatten_up_to(r) for r in rests] + return treespec.unflatten(map(func, *flat_args)) + + +def tree_map_( + func: Callable[..., Any], + tree: PyTree, + *rests: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + """Like :func:`tree_map`, but do an inplace call on each leaf and return the original tree. + + See also :func:`tree_map`. + + Args: + func (callable): A function that takes ``1 + len(rests)`` arguments, to be applied at the + corresponding leaves of the pytrees. + tree (pytree): A pytree to be mapped over, with each leaf providing the first positional + argument to function ``func``. + rests (tuple of pytree): A tuple of pytrees, each of which has the same structure as + ``tree`` or has ``tree`` as a prefix. + is_leaf (callable, optional): An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns: + The original ``tree`` with the value at each leaf is given by the side-effect of function + ``func(x, *xs)`` (not the return value) where ``x`` is the value at the corresponding leaf + in ``tree`` and ``xs`` is the tuple of values at values at corresponding nodes in ``rests``. + """ + leaves, treespec = tree_flatten(tree, is_leaf=is_leaf) + flat_args = [leaves] + [treespec.flatten_up_to(r) for r in rests] + deque(map(func, *flat_args), maxlen=0) # consume and exhaust the iterable + return tree + + +Type2 = tuple[type[T], type[S]] +Type3 = tuple[type[T], type[S], type[U]] +TypeAny = type[Any] | tuple[type[Any], ...] | types.UnionType + +Fn2 = Callable[[T | S], R] +Fn3 = Callable[[T | S | U], R] +Fn = Callable[[T], R] +FnAny = Callable[[Any], R] + +MapOnlyFn = Callable[[T], Callable[[Any], Any]] + + +# These specializations help with type inference on the lambda passed to this +# function +@overload +def map_only(type_or_types_or_pred: type[T], /) -> MapOnlyFn[Fn[T, Any]]: ... + + +@overload +def map_only(type_or_types_or_pred: Type2[T, S], /) -> MapOnlyFn[Fn2[T, S, Any]]: ... + + +@overload +def map_only( + type_or_types_or_pred: Type3[T, S, U], / +) -> MapOnlyFn[Fn3[T, S, U, Any]]: ... + + +# This specialization is needed for the implementations below that call +@overload +def map_only(type_or_types_or_pred: TypeAny, /) -> MapOnlyFn[FnAny[Any]]: ... + + +@overload +def map_only( + type_or_types_or_pred: Callable[[Any], bool], / +) -> MapOnlyFn[FnAny[Any]]: ... + + +def map_only( + type_or_types_or_pred: TypeAny | Callable[[Any], bool], / +) -> MapOnlyFn[FnAny[Any]]: + """ + Suppose you are writing a tree_map over tensors, leaving everything + else unchanged. Ordinarily you would have to write: + + def go(t): + if isinstance(t, Tensor): + return ... + else: + return t + + With this function, you only need to write: + + @map_only(Tensor) + def go(t): + return ... + + You can also directly use 'tree_map_only' + """ + if isinstance(type_or_types_or_pred, (type, tuple, types.UnionType)): + + def pred(x: Any) -> bool: + return isinstance(x, type_or_types_or_pred) # type: ignore[arg-type] + + elif callable(type_or_types_or_pred): + pred = type_or_types_or_pred # type: ignore[assignment] + else: + raise TypeError("Argument must be a type, a tuple of types, or a callable.") + + def wrapper(func: Callable[[T], Any]) -> Callable[[Any], Any]: + @functools.wraps(func) + def wrapped(x: T) -> Any: + if pred(x): + return func(x) + return x + + return wrapped + + return wrapper + + +@overload +def tree_map_only( + type_or_types_or_pred: type[T], + /, + func: Fn[T, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only( + type_or_types_or_pred: Type2[T, S], + /, + func: Fn2[T, S, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only( + type_or_types_or_pred: Type3[T, S, U], + /, + func: Fn3[T, S, U, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only( + type_or_types_or_pred: TypeAny, + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only( + type_or_types_or_pred: Callable[[Any], bool], + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +def tree_map_only( + type_or_types_or_pred: TypeAny | Callable[[Any], bool], + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + return tree_map(map_only(type_or_types_or_pred)(func), tree, is_leaf=is_leaf) + + +@overload +def tree_map_only_( + type_or_types_or_pred: type[T], + /, + func: Fn[T, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only_( + type_or_types_or_pred: Type2[T, S], + /, + func: Fn2[T, S, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only_( + type_or_types_or_pred: Type3[T, S, U], + /, + func: Fn3[T, S, U, Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only_( + type_or_types_or_pred: TypeAny, + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +@overload +def tree_map_only_( + type_or_types_or_pred: Callable[[Any], bool], + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: ... + + +def tree_map_only_( + type_or_types_or_pred: TypeAny | Callable[[Any], bool], + /, + func: FnAny[Any], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + return tree_map_(map_only(type_or_types_or_pred)(func), tree, is_leaf=is_leaf) + + +def tree_all( + pred: Callable[[Any], bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + flat_args = tree_iter(tree, is_leaf=is_leaf) + return all(map(pred, flat_args)) + + +def tree_any( + pred: Callable[[Any], bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + flat_args = tree_iter(tree, is_leaf=is_leaf) + return any(map(pred, flat_args)) + + +@overload +def tree_all_only( + type_or_types: type[T], + /, + pred: Fn[T, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +@overload +def tree_all_only( + type_or_types: Type2[T, S], + /, + pred: Fn2[T, S, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +@overload +def tree_all_only( + type_or_types: Type3[T, S, U], + /, + pred: Fn3[T, S, U, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +def tree_all_only( + type_or_types: TypeAny, + /, + pred: FnAny[bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + flat_args = tree_iter(tree, is_leaf=is_leaf) + return all(pred(x) for x in flat_args if isinstance(x, type_or_types)) + + +@overload +def tree_any_only( + type_or_types: type[T], + /, + pred: Fn[T, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +@overload +def tree_any_only( + type_or_types: Type2[T, S], + /, + pred: Fn2[T, S, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +@overload +def tree_any_only( + type_or_types: Type3[T, S, U], + /, + pred: Fn3[T, S, U, bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: ... + + +def tree_any_only( + type_or_types: TypeAny, + /, + pred: FnAny[bool], + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> bool: + flat_args = tree_iter(tree, is_leaf=is_leaf) + return any(pred(x) for x in flat_args if isinstance(x, type_or_types)) + + +# Broadcasts a pytree to the provided TreeSpec and returns the flattened +# values. If this is not possible, then this function returns None. +# +# For example, given pytree=0 and spec=TreeSpec(list, None, [LeafSpec(), LeafSpec()]), +# would return [0, 0]. This is useful for part of the vmap implementation: +# a user can pass in vmap(fn, in_dims)(*inputs). `in_dims` should be +# broadcastable to the tree structure of `inputs` and we use +# _broadcast_to_and_flatten to check this. +def _broadcast_to_and_flatten( + tree: PyTree, + treespec: TreeSpec, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> list[Any] | None: + def broadcast_prefix( + prefix_tree: PyTree, + full_tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, + ) -> list[Any]: + result: list[Any] = [] + + def add_leaves(x: Any, subtree: PyTree) -> None: + subtreespec = tree_structure(subtree, is_leaf=is_leaf) + result.extend([x] * subtreespec.num_leaves) + + tree_map_( + add_leaves, + prefix_tree, + full_tree, + is_leaf=is_leaf, + ) + return result + + full_tree = tree_unflatten([0] * treespec.num_leaves, treespec) + try: + return broadcast_prefix(tree, full_tree, is_leaf=is_leaf) + except ValueError: + return None + + +@dataclasses.dataclass +class _TreeSpecSchema: + """ + _TreeSpecSchema is the schema used to serialize the TreeSpec + It contains the following fields: + - type: A string name of the type. null for the case of a LeafSpec. + - context: Any format which is json dumpable + - children_spec: A list of children serialized specs. + """ + + type: str | None + context: DumpableContext + children_spec: list["_TreeSpecSchema"] + + +class _ProtocolFn(NamedTuple): + treespec_to_json: Callable[[TreeSpec], DumpableContext] + json_to_treespec: Callable[[DumpableContext], TreeSpec] + + +_SUPPORTED_PROTOCOLS: dict[int, _ProtocolFn] = {} + + +def _treespec_to_json(treespec: TreeSpec) -> _TreeSpecSchema: + if treespec.is_leaf(): + return _TreeSpecSchema(None, None, []) + + if treespec.type not in SUPPORTED_SERIALIZED_TYPES: + raise NotImplementedError( + f"Serializing {treespec.type} in pytree is not registered.", + ) + + serialize_node_def = SUPPORTED_SERIALIZED_TYPES[treespec.type] + + serialized_type_name = serialize_node_def.serialized_type_name + + if serialized_type_name == NO_SERIALIZED_TYPE_NAME_FOUND: + raise NotImplementedError( + f"No registered serialization name for {treespec.type} found. " + "Please update your _register_pytree_node call with a `serialized_type_name` kwarg." + ) + + if serialize_node_def.to_dumpable_context is None: + try: + serialized_context = json.dumps(treespec._context, cls=EnumEncoder) + except TypeError as e: + raise TypeError( + "Unable to serialize context. " + "Please make the context json dump-able, or register a " + "custom serializer using _register_pytree_node." + ) from e + else: + serialized_context = serialize_node_def.to_dumpable_context(treespec._context) + + child_schemas = [_treespec_to_json(child) for child in treespec._children] + + return _TreeSpecSchema(serialized_type_name, serialized_context, child_schemas) + + +def enum_object_hook(obj: dict[str, Any]) -> Enum | dict[str, Any]: + if "__enum__" in obj: + modname, _, classname = obj["fqn"].partition(":") + mod = importlib.import_module(modname) + enum_cls = mod + for attr in classname.split("."): + enum_cls = getattr(enum_cls, attr) + enum_cls = cast(type[Enum], enum_cls) + + return enum_cls[obj["name"]] + return obj + + +def _json_to_treespec(json_schema: DumpableContext) -> TreeSpec: + if ( + json_schema["type"] is None + and json_schema["context"] is None + and len(json_schema["children_spec"]) == 0 + ): + return _LEAF_SPEC + + if json_schema["type"] not in SERIALIZED_TYPE_TO_PYTHON_TYPE: + raise NotImplementedError( + f"Deserializing {json_schema['type']} in pytree is not registered.", + ) + + typ = SERIALIZED_TYPE_TO_PYTHON_TYPE[json_schema["type"]] + serialize_node_def = SUPPORTED_SERIALIZED_TYPES[typ] + + if serialize_node_def.from_dumpable_context is None: + try: + context = json.loads(json_schema["context"], object_hook=enum_object_hook) + except TypeError as ex: + raise TypeError( + "Unable to deserialize context. " + "Please make the context json load-able, or register a " + "custom serializer using _register_pytree_node.", + ) from ex + else: + context = serialize_node_def.from_dumpable_context(json_schema["context"]) + + children_specs = [ + _json_to_treespec(child_string) for child_string in json_schema["children_spec"] + ] + + return TreeSpec(typ, context, children_specs) + + +_SUPPORTED_PROTOCOLS[1] = _ProtocolFn(_treespec_to_json, _json_to_treespec) + + +def treespec_dumps(treespec: TreeSpec, protocol: int | None = None) -> str: + treespec = _ensure_python_treespec_instance(treespec) + + if protocol is None: + protocol = DEFAULT_TREESPEC_SERIALIZATION_PROTOCOL + + if protocol in _SUPPORTED_PROTOCOLS: + json_spec = _SUPPORTED_PROTOCOLS[protocol].treespec_to_json(treespec) + else: + raise ValueError( + f"Unknown protocol {protocol}. " + f"Available protocols: {list(_SUPPORTED_PROTOCOLS.keys())}", + ) + + str_spec = json.dumps((protocol, dataclasses.asdict(json_spec)), cls=EnumEncoder) + return str_spec + + +@functools.lru_cache +def treespec_loads(serialized: str) -> TreeSpec: + protocol, json_schema = json.loads(serialized) + + if protocol in _SUPPORTED_PROTOCOLS: + return _SUPPORTED_PROTOCOLS[protocol].json_to_treespec(json_schema) + raise ValueError( + f"Unknown protocol {protocol}. " + f"Available protocols: {list(_SUPPORTED_PROTOCOLS.keys())}", + ) + + +class _Asterisk(str): + __slots__ = () + + def __new__(cls) -> Self: + return super().__new__(cls, "*") + + def __repr__(self) -> str: + return "*" # no quotes + + +_asterisk = _Asterisk() +del _Asterisk + + +def treespec_pprint(treespec: TreeSpec) -> str: + dummy_tree = tree_unflatten([_asterisk] * treespec.num_leaves, treespec) + return repr(dummy_tree) + + +# TODO(angelayi): remove this function after OSS/internal stabilize +@deprecated( + "`pytree_to_str` is deprecated. Please use `treespec_dumps` instead.", + category=FutureWarning, +) +def pytree_to_str(treespec: TreeSpec) -> str: + return treespec_dumps(treespec) + + +# TODO(angelayi): remove this function after OSS/internal stabilize +@deprecated( + "`str_to_pytree` is deprecated. Please use `treespec_loads` instead.", + category=FutureWarning, +) +def str_to_pytree(json: str) -> TreeSpec: + return treespec_loads(json) + + +def arg_tree_leaves(*args: PyTree, **kwargs: PyTree) -> list[Any]: + """Get a flat list of arguments to this function + + A slightly faster version of tree_leaves((args, kwargs)) + """ + leaves: list[Any] = [] + for a in args: + leaves.extend(tree_iter(a)) + for a in kwargs.values(): + leaves.extend(tree_iter(a)) + return leaves + + +def tree_flatten_with_path( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> tuple[list[tuple[KeyPath, Any]], TreeSpec]: + """Flattens a pytree like :func:`tree_flatten`, but also returns each leaf's key path. + + Args: + tree: a pytree to flatten. If it contains a custom type, that type must be + registered with an appropriate `tree_flatten_with_path_fn` when registered + with :func:`register_pytree_node`. + is_leaf: An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + Returns: + A tuple where the first element is a list of (key path, leaf) pairs, and the + second element is a :class:`TreeSpec` representing the structure of the flattened + tree. + """ + _, treespec = tree_flatten(tree, is_leaf) + return list(_generate_key_paths((), tree, is_leaf)), treespec + + +def tree_leaves_with_path( + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> list[tuple[KeyPath, Any]]: + """Gets the leaves of a pytree like ``tree_leaves`` and returns each leaf's key path. + + Args: + tree: a pytree. If it contains a custom type, that type must be + registered with an appropriate `tree_flatten_with_path_fn` when registered + with :func:`register_pytree_node`. + is_leaf: An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + Returns: + A list of (key path, leaf) pairs. + """ + return list(_generate_key_paths((), tree, is_leaf)) + + +def _generate_key_paths( + key_path: KeyPath, + tree: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> Iterable[tuple[KeyPath, Any]]: + if is_leaf and is_leaf(tree): + yield key_path, tree + return + + node_type = _get_node_type(tree) + handler = SUPPORTED_NODES.get(node_type) + if not handler: + # This is a leaf + yield key_path, tree + return + + flatten_with_keys = handler.flatten_with_keys_fn + if flatten_with_keys: + key_children, _ = flatten_with_keys(tree) + for k, c in key_children: + yield from _generate_key_paths((*key_path, k), c, is_leaf) + else: + # We registered this pytree but didn't add a flatten_with_keys_fn, complain. + raise ValueError( + f"Did not find a flatten_with_keys_fn for type: {node_type}. " + "Please pass a flatten_with_keys_fn argument to register_pytree_node." + ) + + +def tree_map_with_path( + func: Callable[..., Any], + tree: PyTree, + *rests: PyTree, + is_leaf: Callable[[PyTree], bool] | None = None, +) -> PyTree: + """Like :func:`tree_map`, but the provided callable takes an additional key path argument. + + Args: + func: A function that takes ``2 + len(rests)`` arguments, to be applied at the + corresponding leaves of the pytrees. The first positional argument + to ``func`` is the key path of the leaf in question. The second + positional argument is the value of the leaf. + tree: A pytree to be mapped over, with each leaf providing the first positional + argument to function ``func``. + rests: A tuple of pytrees, each of which has the same structure as + ``tree`` or has ``tree`` as a prefix. + is_leaf: An extra leaf predicate function that will be called at each + flattening step. The function should have a single argument with signature + ``is_leaf(node) -> bool``. If it returns :data:`True`, the whole subtree being treated + as a leaf. Otherwise, the default pytree registry will be used to determine a node is a + leaf or not. If the function is not specified, the default pytree registry will be used. + + Returns + A new pytree with the same structure as ``tree`` but with the value at each leaf given by + ``func(keypath, x, *xs)`` where ``keypath`` is the key path at the + corresponding leaf in ``tree``, ``x`` is the value at that leaf, and + ``xs`` is the tuple of values at corresponding nodes in ``rests``. + """ + keypath_leaves, treespec = tree_flatten_with_path(tree, is_leaf) + keypath_leaves = list(zip(*keypath_leaves, strict=True)) + all_keypath_leaves = keypath_leaves + [treespec.flatten_up_to(r) for r in rests] + return treespec.unflatten(func(*xs) for xs in zip(*all_keypath_leaves, strict=True)) + + +def keystr(kp: KeyPath) -> str: + """Given a key path, return a pretty-printed representation.""" + return "".join([str(k) for k in kp]) + + +def key_get(obj: Any, kp: KeyPath) -> Any: + """Given an object and a key path, return the value at the key path.""" + for k in kp: + obj = k.get(obj) + return obj diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_runtime_estimation.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_runtime_estimation.py new file mode 100644 index 0000000000000000000000000000000000000000..b2d56dcfc7f79e14c3302d3264b6ec79d623fcf2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_runtime_estimation.py @@ -0,0 +1,151 @@ +import torch +from torch._inductor.utils import get_device_tflops, get_gpu_dram_gbps +from torch.fx.experimental.symbolic_shapes import ( + optimization_hint, + statically_known_true, +) +from torch.utils._ordered_set import OrderedSet + +from .flop_counter import flop_registry + + +aten = torch.ops.aten + +_FLOAT_TYPES = OrderedSet( + [ + torch.float16, + torch.bfloat16, + torch.float32, + torch.float64, + ] +) + +# No fall-back kernel needed/exists for view ops +_VIEW_OPS = OrderedSet( + [ + aten.lift_fresh, + aten.t, + aten.transpose, + aten.view, + aten.detach, + aten._unsafe_view, + aten.split, + aten.adjoint, + aten.as_strided, + aten.diagonal, + aten.expand, + aten.expand_as, + aten.movedim, + aten.permute, + aten.select, + aten.squeeze, + aten.mT, + aten.mH, + aten.real, + aten.imag, + aten.view_as, + aten.unflatten, + aten.unfold, + aten.unbind, + aten.unsqueeze, + aten.vsplit, + aten.hsplit, + aten.split_with_sizes, + aten.swapaxes, + aten.swapdims, + aten.chunk, + ] +) +# We can ignore benchmarking tensor create ops +_CREATE_OPS = OrderedSet( + [ + aten.randint, + aten.randn, + aten.rand, + aten.randn_like, + aten.rand_like, + aten.randint_like, + aten.arange, + aten.ones_like, + aten.zeros_like, + ] +) + +_IGNORE_OPS = _VIEW_OPS | _CREATE_OPS + + +def get_compute_time(func_packet, args, kwargs, out, out_dtypes) -> float: # type: ignore[no-untyped-def] + """ + Estimates the compute time of an aten operator. + + Args: + func_packet: The operator overload packet. + args: The arguments to the operator. + kwargs: The keyword arguments to the operator. + out: The output of the operator. + out_dtypes: The output data types. + + Returns: + float: The estimated compute time in nanoseconds. + """ + if func_packet in flop_registry: + if len(out_dtypes) != 1: + raise AssertionError( + f"Only support single out dtype got {out_dtypes} for {func_packet}" + ) + dtype = out_dtypes.pop() + # This actually gives peta-FLOPs/s hence multiply by 1e15 to get the FLOPs/s + peak_gpu_flops = get_device_tflops(dtype) * 1e15 + # We can expect to achieve 75% of theoretical peak flops + factor = 0.75 + peak_empirical_flops = factor * peak_gpu_flops + flop_count_func = flop_registry[func_packet] + # We divide by a factor of 2 to get the MACs (multiply and accumulate) + flop_count = flop_count_func(*args, **kwargs, out_val=out) / 2 + # We multiply by 1e9 to get the time in nano seconds + compute_time = (flop_count / peak_empirical_flops) * 1e9 + return compute_time + return 0.0 + + +def get_num_bytes(t: torch.Tensor) -> int: + """ + Calculates the memory consumption of a tensor. + + Args: + t (torch.Tensor): The input tensor. + + Returns: + int: The memory consumption of the tensor in bytes. + """ + real_numel = 1 + for size, stride in zip(t.shape, t.stride()): + # For dims with stride=0 (expanded/broadcast), only 1 element accessed + if not statically_known_true(stride == 0): + real_numel *= optimization_hint(size, fallback=0) + + return real_numel * t.element_size() + + +def get_transfer_time(flat_args_kwargs, flat_outs) -> float: # type: ignore[no-untyped-def] + """ + Estimates the memory transfer time of input and output tensors. + + Args: + flat_args_kwargs (List[torch.Tensor]): The flat list of arguments and keyword arguments. + flat_outs (List[torch.Tensor]): The flat list of outputs. + + Returns: + float: The estimated memory transfer time in nanoseconds. + """ + gpu_memory_bandwidth = get_gpu_dram_gbps() + read_bytes = sum( + get_num_bytes(t) for t in flat_args_kwargs if isinstance(t, torch.Tensor) + ) + write_bytes = sum( + get_num_bytes(t) for t in flat_outs if isinstance(t, torch.Tensor) + ) + counted_bytes = read_bytes + write_bytes + # The GPU memory bandwidth is in GB/s so the transfer time is in nanoseconds + transfer_time = counted_bytes / gpu_memory_bandwidth + return transfer_time diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_stats.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_stats.py new file mode 100644 index 0000000000000000000000000000000000000000..b8a2978c3ea7066c56070382d0d4866faf58ab8d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_stats.py @@ -0,0 +1,31 @@ +# NOTE! PLEASE KEEP THIS FILE *FREE* OF TORCH DEPS! IT SHOULD BE IMPORTABLE ANYWHERE. +# IF YOU FEEL AN OVERWHELMING URGE TO ADD A TORCH DEP, MAKE A TRAMPOLINE FILE A LA torch._dynamo.utils +# AND SCRUB AWAY TORCH NOTIONS THERE. +import collections +import functools +from collections import OrderedDict +from collections.abc import Callable +from typing import TypeVar +from typing_extensions import ParamSpec + + +simple_call_counter: OrderedDict[str, int] = collections.OrderedDict() + +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +def count_label(label: str) -> None: + prev = simple_call_counter.setdefault(label, 0) + simple_call_counter[label] = prev + 1 + + +def count(fn: Callable[_P, _R]) -> Callable[_P, _R]: + @functools.wraps(fn) + def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _R: + if fn.__qualname__ not in simple_call_counter: + simple_call_counter[fn.__qualname__] = 0 + simple_call_counter[fn.__qualname__] = simple_call_counter[fn.__qualname__] + 1 + return fn(*args, **kwargs) + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_strobelight/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_strobelight/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_strobelight/cli_function_profiler.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_strobelight/cli_function_profiler.py new file mode 100644 index 0000000000000000000000000000000000000000..d2e1595bf2a1477b33ed00446d86e6cdea267a8f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_strobelight/cli_function_profiler.py @@ -0,0 +1,313 @@ +# mypy: disallow-untyped-defs + +import functools +import logging +import os +import re +import subprocess +import time +from collections.abc import Callable, Sequence +from threading import Lock +from typing import Any, TypeVar +from typing_extensions import ParamSpec + + +logger = logging.getLogger("strobelight_function_profiler") + +console_handler = logging.StreamHandler() +formatter = logging.Formatter( + "%(name)s, line %(lineno)d, %(asctime)s, %(levelname)s: %(message)s" +) +console_handler.setFormatter(formatter) + +logger.addHandler(console_handler) +logger.setLevel(logging.INFO) +logger.propagate = False + +_P = ParamSpec("_P") +_R = TypeVar("_R") + + +class StrobelightCLIProfilerError(Exception): + """ + Raised when an error happens during strobelight profiling + """ + + +def _pid_namespace_link(pid: int | None = None) -> str: + """Returns the link to the process's namespace, example: pid:[4026531836]""" + PID_NAMESPACE_PATH = "/proc/{}/ns/pid" + pid = pid or os.getpid() + return os.readlink(PID_NAMESPACE_PATH.format(pid)) + + +def _pid_namespace(pid: int | None = None) -> int: + """Returns the process's namespace id""" + pid = pid or os.getpid() + link = _pid_namespace_link(pid) + return int(link[link.find("[") + 1 : -1]) + + +def _command_to_string(command: Sequence[str]) -> str: + return " ".join(command) + + +class StrobelightCLIFunctionProfiler: + """ + Note: this is a meta only tool. + + StrobelightCLIFunctionProfiler can be used to profile a python function and + generate a strobelight link with the results. It works on meta servers but + does not requires an fbcode target. + When stop_at_error is false(default), error during profiling does not prevent + the work function from running. + + Check function_profiler_example.py for an example. + """ + + # This lock is used to make sure only one thread is running the profiler at any point. + _lock = Lock() + + def __init__( + self, + *, + stop_at_error: bool = False, + max_profile_duration_sec: int = 60 * 10, + sample_each: float = 1e7, # sample each sample_each cycles. + run_user_name: str = "pytorch-strobelight-ondemand", + timeout_wait_for_running_sec: int = 60, + timeout_wait_for_finished_sec: int = 60, + recorded_env_variables: list[str] | None = None, + sample_tags: list[str] | None = None, + stack_max_len: int = 127, + async_stack_max_len: int = 127, + ) -> None: + self.stop_at_error = stop_at_error + self.max_profile_duration_sec = max_profile_duration_sec + self.sample_each = sample_each + self.run_user_name = run_user_name + self.timeout_wait_for_running_sec = timeout_wait_for_running_sec + self.timeout_wait_for_finished_sec = timeout_wait_for_finished_sec + # Results of the most recent run. + # Tracks the strobelight run id of the most recent run + self.current_run_id: int | None = None + self.sample_tags = sample_tags + + def _run_async(self) -> None: + processId = os.getpid() + namespace = _pid_namespace(processId) + command = [ + "strobeclient", + "run", + "--profiler", + "pyperf", + "--event", + "cycles", + "--async", + "--sample-interval", + f"{int(self.sample_each)}", + "--duration-ms", + f"{int(self.max_profile_duration_sec * 1000)}", + "--pid", + f"{namespace}:{processId}", + ] + + if self.sample_tags: + command.append("--sample-tags") + command.append(",".join(self.sample_tags)) + + logger.debug("running command: %s", _command_to_string(command)) + result = subprocess.run(command, capture_output=True) + output = result.stderr.decode("utf-8") + logger.debug("output:\n{%s}", output) + + if result.returncode != 0: + raise StrobelightCLIProfilerError( + f"failed to start strobelight profiling, error in run_async:{output}" + ) + + if match := re.search(r"INFO Run Id: (-?\d+)", output): + self.current_run_id = int(match.group(1)) + return + + raise StrobelightCLIProfilerError( + f"failed to start strobelight profiling, unexpected result {output}" + ) + + def _wait_for_running(self, counter: int = 0) -> None: + if counter > 20: + raise StrobelightCLIProfilerError( + "wait_for_running called more than 20 times" + ) + + command = ["strobeclient", "getRunStatus", "--run-id", f"{self.current_run_id}"] + logger.debug("running command: %s", _command_to_string(command)) + result = subprocess.run(command, capture_output=True) + output = result.stderr.decode("utf-8") + logger.debug("output:\n{%s}", output) + + if result.returncode != 0: + raise StrobelightCLIProfilerError( + f"failed to start strobelight profiling, error in wait_for_running:{output}" + ) + + if match := re.search("Profile run status: (.*)", output): + current_status = match.group(1) + if current_status == "RUNNING": + return + elif current_status == "PREPARING": + time.sleep(10) + self._wait_for_running(counter + 1) + return + else: + raise StrobelightCLIProfilerError(f"unexpected {current_status} phase") + + raise StrobelightCLIProfilerError(f"unexpected output\n: {output} ") + + def _stop_run(self) -> None: + command = ["strobeclient", "stopRun", "--run-id", str(self.current_run_id)] + logger.debug("running command: %s", _command_to_string(command)) + result = subprocess.run(command, capture_output=True) + output = result.stderr.decode("utf-8") + logger.debug("output:\n{%s}", output) + + if result.returncode != 0: + raise StrobelightCLIProfilerError( + f"failed to stop strobelight profiling, return code is not 0 :{output}" + ) + + if match := re.search("INFO ::1:(.*)", output): + current_status = match.group(1) + if current_status.__contains__("Success!"): + return + else: + raise StrobelightCLIProfilerError( + f"failed to stop strobelight profiling, got {current_status} result" + ) + + raise StrobelightCLIProfilerError(f"unexpected output\n: {output} ") + + def _get_results(self) -> None: + command = ["strobeclient", "getRunStatus", "--run-id", str(self.current_run_id)] + logger.debug("running command: %s", _command_to_string(command)) + result = subprocess.run(command, capture_output=True) + output = result.stderr.decode("utf-8") + logger.debug("output:\n{%s}", output) + + if result.returncode != 0: + raise StrobelightCLIProfilerError( + f"failed to extract profiling results, return code is not 0 : {output}" + ) + + if match := re.search("INFO ::1:(.*)", output): + current_status = match.group(1) + if current_status.__contains__("Profile run status: PROCESSING"): + time.sleep(10) + self._get_results() + return + elif not current_status.__contains__("Profile run finished with SUCCESS"): + raise StrobelightCLIProfilerError( + f"failed to extract profiling results, unexpected response {output}" + ) + + for item in re.findall( + r"(Total samples(.*)|GraphProfiler(.*)|Icicle view \(python stack\)(.*))", + output, + ): + logger.info(item[0]) + + def _stop_strobelight_no_throw( + self, + collect_results: bool, + ) -> None: + try: + # call stop run + self._stop_run() + logger.info("strobelight profiling stopped") + + logger.debug("collection stopped") + + if not collect_results: + return + + self._get_results() + except Exception: + logger.warning("error during stop_strobelight", exc_info=True) + + # Return true if strobelight started and is running. Never throw. + def _start_strobelight(self) -> bool: + strobelight_started = False + try: + self._run_async() + strobelight_started = True + logger.info("strobelight run id is: %s", self.current_run_id) + self._wait_for_running() + logger.info("strobelight profiling running") + return True + + except Exception: + logger.warning("error during start_strobelight:", exc_info=True) + if strobelight_started: + self._stop_strobelight_no_throw(collect_results=False) + return False + + def profile( + self, work_function: Callable[_P, _R], *args: _P.args, **kwargs: _P.kwargs + ) -> _R | None: + self.current_run_id = None + + if locked := StrobelightCLIFunctionProfiler._lock.acquire(False): + if not locked: + if self.stop_at_error: + raise StrobelightCLIProfilerError("concurrent runs not supported") + + logger.warning("concurrent runs not supported") + return work_function(*args, **kwargs) + + started = self._start_strobelight() + if not started: + if self.stop_at_error: + StrobelightCLIFunctionProfiler._lock.release() + raise StrobelightCLIProfilerError( + "failed to start strobelight profiling" + ) + result = work_function(*args, **kwargs) + StrobelightCLIFunctionProfiler._lock.release() + return result + + try: + logger.debug("collection started") + result = work_function(*args, **kwargs) + self._stop_strobelight_no_throw(collect_results=True) + StrobelightCLIFunctionProfiler._lock.release() + return result + except Exception as error: + logger.warning("work function throw exception", exc_info=True) + self._stop_strobelight_no_throw(collect_results=False) + StrobelightCLIFunctionProfiler._lock.release() + raise error + return None + + +# A function decorator that wraps profile, if no profiler is provided one with +# default args is created. A function can be annotated as: +# @strobelight() +# @strobelight(profiler = StrobelightFunctionProfiler(stop_at_error=True,..)) +# @strobelight(stop_at_error=True,...) +def strobelight( + profiler: StrobelightCLIFunctionProfiler | None = None, **kwargs: Any +) -> Callable[[Callable[_P, _R]], Callable[_P, _R | None]]: + if not profiler: + profiler = StrobelightCLIFunctionProfiler(**kwargs) + + def strobelight_inner( + work_function: Callable[_P, _R], + ) -> Callable[_P, _R | None]: + @functools.wraps(work_function) + def wrapper_function(*args: _P.args, **kwargs: _P.kwargs) -> _R | None: + # pyrefly: ignore [bad-argument-type] + return profiler.profile(work_function, *args, **kwargs) + + return wrapper_function + + return strobelight_inner diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/functions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/functions.py new file mode 100644 index 0000000000000000000000000000000000000000..1f95d5337f6a179f45f529168cdf8ca2392fb6fa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/functions.py @@ -0,0 +1,1488 @@ +# mypy: allow-untyped-defs +import functools +import math +import operator +import sys +from collections.abc import Callable +from typing import SupportsFloat, TYPE_CHECKING, TypeVar +from typing_extensions import TypeVarTuple, Unpack + +import sympy +from sympy import S +from sympy.core import sympify +from sympy.core.expr import Expr +from sympy.core.function import Application +from sympy.core.logic import _torf, fuzzy_and, fuzzy_or +from sympy.core.numbers import equal_valued +from sympy.core.operations import LatticeOp, ShortCircuit +from sympy.core.sorting import ordered +from sympy.core.traversal import walk +from sympy.printing.precedence import PRECEDENCE +from sympy.utilities.iterables import sift + +from torch.torch_version import TorchVersion + +from .numbers import int_oo, is_infinite + + +if TYPE_CHECKING: + from collections.abc import Iterable + + +_T = TypeVar("_T", bound=SupportsFloat) +_Ts = TypeVarTuple("_Ts") + +# Portions of this file are adapted from the Sympy codebase, which was +# licensed as follows: +# +# Copyright (c) 2006-2023 SymPy Development Team +# +# All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# +# a. Redistributions of source code must retain the above copyright notice, +# this list of conditions and the following disclaimer. +# b. Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# c. Neither the name of SymPy nor the names of its contributors +# may be used to endorse or promote products derived from this software +# without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL THE REGENTS OR CONTRIBUTORS BE LIABLE FOR +# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT +# LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY +# OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH +# DAMAGE. + +__all__ = [ + "FloorDiv", + "ModularIndexing", + "Where", + "PythonMod", + "Mod", + "CleanDiv", + "CeilToInt", + "FloorToInt", + "CeilDiv", + "IntTrueDiv", + "FloatTrueDiv", + "LShift", + "RShift", + "IsNonOverlappingAndDenseIndicator", + "TruncToFloat", + "TruncToInt", + "RoundToInt", + "RoundDecimal", + "ToFloat", + "FloatPow", + "PowByNatural", + "Identity", +] + + +def _is_symbols_binary_summation(expr: sympy.Expr) -> bool: + # No need to check that two args are not the same, since expr is pr-optimized but we do it anyway. + return ( + isinstance(expr, sympy.Expr) + and expr.is_Add + and len(expr._args) == 2 + and expr._args[0].is_symbol + and expr._args[1].is_symbol + and expr._args[0] is not expr._args[1] + ) + + +def _keep_float( + f: Callable[[Unpack[_Ts]], _T], +) -> Callable[[Unpack[_Ts]], _T | sympy.Float]: + @functools.wraps(f) + def inner(*args: Unpack[_Ts]) -> _T | sympy.Float: + r: _T | sympy.Float = f(*args) + if any(isinstance(a, sympy.Float) for a in args) and not isinstance( + r, sympy.Float + ): + r = sympy.Float(float(r)) + return r + + # pyrefly: ignore [bad-return] + return inner + + +def fuzzy_eq(x: bool | None, y: bool | None) -> bool | None: + if None in (x, y): + return None + return x == y + + +def simple_floordiv_gcd(p: sympy.Basic, q: sympy.Basic) -> sympy.Basic: + """ + Fast path for sympy.gcd, using a simple factoring strategy. + + We try to rewrite p and q in the form n*e*p1 + n*e*p2 and n*e*q0, + where n is the greatest common integer factor and e is the largest + syntactic common factor (i.e., common sub-expression) in p and q. + Then the gcd returned is n*e, cancelling which we would be left with + p1 + p2 and q0. + + Note that further factoring of p1 + p2 and q0 might be possible with + sympy.factor (which uses domain-specific theories). E.g., we are unable + to find that x*y + x + y + 1 is divisible by x + 1. More generally, + when q is of the form q1 + q2 (instead of being already factored) it + might be necessary to fall back on sympy.gcd. + """ + + def integer_coefficient(x: sympy.Basic) -> int: + integer_coefficients: list[int] = [ + abs(int(arg)) + for arg in sympy.Mul.make_args(x) + if isinstance(arg, (int, sympy.Integer)) + ] + return math.prod(integer_coefficients) + + def integer_factor(expr: sympy.Basic) -> int: + integer_factors: Iterable[int] = map( + integer_coefficient, sympy.Add.make_args(expr) + ) + return functools.reduce(math.gcd, integer_factors) + + gcd: int = math.gcd(integer_factor(p), integer_factor(q)) + p, q = p / gcd, q / gcd # type: ignore[operator, assignment] # remove in py3.12 + + base_splits: list[tuple[sympy.Basic, ...]] = list( + map(sympy.Mul.make_args, sympy.Add.make_args(p)) + ) + divisor_split: tuple[sympy.Basic, ...] = sympy.Mul.make_args(q) + for x in divisor_split: + if all(x in base_split for base_split in base_splits): + gcd = gcd * x # type: ignore[operator] # remove in py3.12 + return gcd # type: ignore[return-value] # remove in py3.12 + + +# It would be nice to have assertions on whether or not inputs is_integer +# However, with bugs like https://github.com/sympy/sympy/issues/26620 sympy +# sometimes inconsistently reports floats an integers. +# +# What we can assume from sympy is that if something is an int, it +# definitely is is_integer, but if it is a float it may or may not +# be is_integer. So we are unable to do strong asserts that things +# are NOT integers. + + +# TODO: In Triton, // rounds to zero, but in Python, it is floor division. +# When we can prove both arguments are non-negative, we should just have a +# GenericFloorDiv (name pending) which can codegen efficiently in Python/C, +# and then PythonFloorDiv and CIntDiv which have the appropriate rounding +# semantics. +# +# Right now, FloorDiv de facto changes behavior if arguments are negative or +# not, this can potentially cause correctness issues. +class FloorDiv(sympy.Function): + """ + We maintain this so that: + 1. We can use divisibility guards to simplify FloorDiv(a, b) to a / b. + 2. Printing out the expression is nicer (compared to say, representing a//b as (a - a % b) / b) + + NB: This is Python-style floor division, round to -Inf + """ + + nargs: tuple[int, ...] = (2,) + precedence: int = 35 # lower precedence than add + is_integer: bool = True + + @property + def base(self) -> sympy.Basic: + return self.args[0] + + @property + def divisor(self) -> sympy.Basic: + return self.args[1] + + def _sympystr(self, printer: sympy.printing.StrPrinter) -> str: + base = printer.parenthesize(self.base, PRECEDENCE["Atom"] - 0.5) + divisor = printer.parenthesize(self.divisor, PRECEDENCE["Atom"] - 0.5) + return f"({base}//{divisor})" + + # Automatic evaluation. + # https://docs.sympy.org/latest/guides/custom-functions.html#best-practices-for-eval + @classmethod + def eval(cls, base: sympy.Integer, divisor: sympy.Integer) -> sympy.Basic | None: + # python test/test_dynamic_shapes.py -k TestDimConstraints.test_dim_constraints_solve_full + # Assert triggered by inequality solver + # assert base.is_integer, base + # assert divisor.is_integer, divisor + + # We don't provide the same error message as in Python because SymPy + # makes it difficult to check the types. + if divisor.is_zero: + raise ZeroDivisionError("division by zero") + if is_infinite(base) and is_infinite(divisor): + return sympy.nan + if base is sympy.nan or divisor is sympy.nan: + return sympy.nan + + if base.is_zero: + return sympy.S.Zero + if base.is_integer and equal_valued(divisor, 1): + return base + if base.is_integer and equal_valued(divisor, -1): + return sympy.Mul(base, -1) + if base == divisor: + return sympy.S.One + + if ( + isinstance(base, sympy.Number) + and isinstance(divisor, sympy.Number) + and (is_infinite(base) or is_infinite(divisor)) + ): + r = float(base) / float(divisor) + if r == math.inf: + return int_oo + elif r == -math.inf: + return -int_oo + elif math.isnan(r): + return sympy.nan + else: + return sympy.Integer(math.floor(r)) + if isinstance(base, sympy.Integer) and isinstance(divisor, sympy.Integer): + return sympy.Integer(int(base) // int(divisor)) + if isinstance(base, FloorDiv): + return FloorDiv(base.args[0], base.args[1] * divisor) + + # Expands (x + y) // b into x // b + y // b. + # This only works if floor is an identity, i.e. x / b is an integer. + if isinstance(divisor, sympy.Integer): + quotients = 0 + terms = [] + for term in sympy.Add.make_args(base): + quotient = term / divisor + + # This is a sympy bug fixed in https://github.com/sympy/sympy/pull/28442 + # sympy can generate a quotient with (1/22)*.... such that quotient.is_integer is True + # FloorDiv should not allow that as output. see + quotient_is_integer = None + if isinstance(quotient, sympy.Mul) and TorchVersion( + sympy.__version__ + ) < TorchVersion("1.15.0"): + rationals = quotient.atoms(sympy.Rational) + all_rationals_ints = all(r.q == 1 for r in rationals) + quotient_is_integer = quotient.is_integer and all_rationals_ints + else: + quotient_is_integer = quotient.is_integer + + if quotient_is_integer: + terms.append(term) + quotients += quotient + + if len(terms) != 0: + # Passing evaluate = False since expression will be optimized during the subtraction post its construction. + return ( + FloorDiv(base - sympy.Add(*terms, evaluate=False), divisor) + + quotients + ) + + try: + gcd = simple_floordiv_gcd(base, divisor) + if equal_valued(gcd, 1) and isinstance(divisor, sympy.Add): + gcd = sympy.gcd(base, divisor) + if not equal_valued(gcd, 1): + return FloorDiv( + sympy.simplify(base / gcd), sympy.simplify(divisor / gcd) + ) + except sympy.PolynomialError: + pass # https://github.com/pytorch/pytorch/issues/108276 + + return None + + def _eval_is_nonnegative(self) -> bool | None: + p, q = self.args[:2] + if all([p.is_integer, q.is_integer, p.is_nonnegative, q.is_nonnegative]): + return True + return None + + +class ModularIndexing(sympy.Function): + """ + ModularIndexing(a, b, c) => (a // b) % c where % is the C modulus + """ + + nargs: tuple[int, ...] = (3,) + is_integer: bool = True + precedence: int = 35 # lower precedence than add + + @classmethod + def eval( + cls, base: sympy.Integer, divisor: sympy.Integer, modulus: sympy.Integer + ) -> sympy.Basic | None: + if base == 0 or modulus == 1: + return sympy.S.Zero + if ( + isinstance(base, sympy.Integer) + and isinstance(divisor, sympy.Integer) + and isinstance(modulus, sympy.Integer) + ): + return (base // divisor) % modulus + + try: + if divisor != 1: + gcd = sympy.gcd(base, divisor) + if gcd != 1: + return ModularIndexing( + sympy.simplify(base / gcd), + sympy.simplify(divisor / gcd), + modulus, + ) + except sympy.PolynomialError: + pass # https://github.com/pytorch/pytorch/issues/108276 + + if isinstance(base, sympy.Add): + new_terms: list[sympy.Integer] = [] + all_positive: bool = True + for term in base.args: + if sympy.gcd(term, modulus * divisor) != modulus * divisor: + if (isinstance(term, sympy.Integer) and term < 0) or ( + isinstance(term, sympy.Mul) + and isinstance(term.args[0], sympy.Integer) + and term.args[0] < 0 + ): + # workaround for https://github.com/triton-lang/triton/issues/619, + # if there are negative terms, // produces wrong result + # TODO if https://github.com/triton-lang/triton/issues/619 is fixed + # this optimization would become valid + all_positive = False + break + else: + new_terms.append(term) + + if len(new_terms) != len(base.args) and all_positive: + return ModularIndexing(sum(new_terms), divisor, modulus) + + if isinstance(base, FloorDiv): + return ModularIndexing(base.args[0], base.args[1] * divisor, modulus) + + return None + + def _eval_is_nonnegative(self) -> bool | None: + p, q = self.args[:2] + return fuzzy_eq(p.is_nonnegative, q.is_nonnegative) # type: ignore[attr-defined] + + +class Where(sympy.Function): + """ + Good ol' ternary operator + """ + + nargs: tuple[int, ...] = (3,) + precedence: int = 35 # lower precedence than add + + def _eval_is_integer(self) -> bool | None: + return True if self.args[1].is_integer and self.args[2].is_integer else None # type: ignore[attr-defined] + + def _eval_is_nonnegative(self) -> bool | None: + return ( + True + if self.args[1].is_nonnegative and self.args[2].is_nonnegative # type: ignore[attr-defined] + else None + ) + + def _eval_is_positive(self) -> bool | None: + return True if self.args[1].is_positive and self.args[2].is_positive else None # type: ignore[attr-defined] + + @classmethod + def eval(cls, c: sympy.Basic, p: sympy.Basic, q: sympy.Basic) -> sympy.Basic | None: + if c == sympy.true: + return p + elif c == sympy.false: + return q + return None + + +# Python-style modulus: take sign from RHS +class PythonMod(sympy.Function): + nargs: tuple[int, ...] = (2,) + + precedence: int = 35 # lower precedence than add + is_integer: bool = True + + @classmethod + def eval(cls, p: sympy.Expr, q: sympy.Expr) -> sympy.Expr | None: + # python test/dynamo/test_export.py -k ExportTests.test_trivial_constraint + # Triggered by sympy.solvers.inequalities.reduce_inequalities + # assert p.is_integer, p + # assert q.is_integer, q + + if q.is_zero: + raise ZeroDivisionError("Modulo by zero") + + # Three cases: + # 1. p == 0 + # 2. p is either q or -q + # 3. p is integer and q == 1 + if p is S.Zero or p in (q, -q) or q == 1: + return S.Zero + + # Evaluate if they are both literals. + if q.is_Number and p.is_Number: + return p % q + + # If q == 2, it's a matter of whether p is odd or even. + if q.is_Number and q == 2: + if p.is_even: + return S.Zero + if p.is_odd: + return S.One + + # If p is a multiple of q. + r = p / q + if r.is_integer: + return S.Zero + + # If p < q and its ratio is positive, then: + # - floor(p / q) = 0 + # - p % q = p - floor(p / q) * q = p + less = p < q + # pyrefly: ignore [missing-attribute] + if less.is_Boolean and bool(less) and r.is_positive: + return p + + if sympy.Mod(p, q) == 0: + return S.Zero + + return None + + # NB: args[1] for PythonMod + def _eval_is_nonnegative(self) -> bool | None: + return True if self.args[1].is_positive else None # type: ignore[attr-defined] + + def _eval_is_nonpositive(self) -> bool | None: + return True if self.args[1].is_negative else None # type: ignore[attr-defined] + + def _ccode(self, printer) -> str: + p = printer.parenthesize(self.args[0], PRECEDENCE["Atom"] - 0.5) + q = printer.parenthesize(self.args[1], PRECEDENCE["Atom"] - 0.5) + abs_q = str(q) if self.args[1].is_positive else f"abs({q})" + return f"({p} % {q}) < 0 ? {p} % {q} + {abs_q} : {p} % {q}" + + +# Generic modulus: only defined on non-negative arguments +class Mod(sympy.Function): + nargs = (2,) + precedence: int = 35 # lower precedence than add + + is_integer = True + is_nonnegative = True + + @classmethod + def eval(cls, p, q): + # This was adapted from: sympy/core/mod.py + + # Triggered by + # python test/test_dynamic_shapes.py -k TestDimConstraints.test_dim_constraints_solve_full + # assert p.is_integer, p + # assert q.is_integer, q + + if q.is_zero: + raise ZeroDivisionError("Modulo by zero") + + # Three cases: + # 1. p == 0 + # 2. p is either q or -q + # 3. p is integer and q == 1 + if p is S.Zero or p in (q, -q) or q == 1: + return S.Zero + + # Evaluate if they are both literals. + if q.is_Number and p.is_Number: + if p < 0: + raise AssertionError(p) + if q < 1: + raise AssertionError(q) + return p % q + + # If q == 2, it's a matter of whether p is odd or even. + if q.is_Number and q == 2: + if p.is_even: + return S.Zero + if p.is_odd: + return S.One + + # If p is a multiple of q. + r = p / q + if r.is_integer: + return S.Zero + + # If p < q and its ratio is positive, then: + # - floor(p / q) = 0 + # - p % q = p - floor(p / q) * q = p + less = p < q + if less.is_Boolean and bool(less) and r.is_positive: + return p + + +class CleanDiv(FloorDiv): + """ + Div where we can assume no rounding. + This is to enable future optimizations. + """ + + +# Don't use sympy ceiling/floor as they will attempt simplifications involving +# frac +class CeilToInt(sympy.Function): + is_integer = True + + @classmethod + def eval(cls, number): + # assert number.is_integer is not True, number + if number in (sympy.oo, int_oo): + return int_oo + if number in (-sympy.oo, -int_oo): + return -int_oo + if isinstance(number, sympy.Number): + return sympy.Integer(math.ceil(float(number))) + + def _ccode(self, printer) -> str: + number = printer.parenthesize(self.args[0], self.args[0].precedence - 0.5) + return f"ceil({number})" + + +class FloorToInt(sympy.Function): + is_integer = True + + @classmethod + def eval(cls, number): + if number in (sympy.oo, int_oo): + return int_oo + if number in (-sympy.oo, int_oo): + return -int_oo + if isinstance(number, sympy.Integer): + return number + if isinstance(number, sympy.Number): + return sympy.Integer(math.floor(float(number))) + + +class CeilDiv(sympy.Function): + """ + Div used in indexing that rounds up. + """ + + is_integer = True + + def __new__(cls, base, divisor): + base = sympy.sympify(base) + divisor = sympy.sympify(divisor) + if sympy.gcd(base, divisor) == divisor: + return CleanDiv(base, divisor) + else: + return FloorDiv(base + (divisor - 1), divisor) + + +class LShift(sympy.Function): + is_integer = True + + @classmethod + def eval(cls, base, shift): + if shift < 0: + raise ValueError("negative shift count") + return base * 2**shift + + +class RShift(sympy.Function): + is_integer = True + + @classmethod + def eval(cls, base, shift): + if shift < 0: + raise ValueError("negative shift count") + return FloorDiv(base, 2**shift) + + +class MinMaxBase(Expr, LatticeOp): # type: ignore[misc] + def __new__(cls, *original_args, **assumptions): + from sympy.core.parameters import global_parameters + + evaluate = assumptions.pop("evaluate", global_parameters.evaluate) + args = (sympify(arg) for arg in original_args) + + # See the comment in _satisfy_unique_summations_symbols. + unique_summations_symbols = ( + None + if not evaluate + else cls._satisfy_unique_summations_symbols(original_args) + ) + + if evaluate: + try: + # first standard filter, for cls.zero and cls.identity + # also reshape Max(a, Max(b, c)) to Max(a, b, c) + args = frozenset(cls._new_args_filter(args)) # type: ignore[assignment] + except ShortCircuit: + return cls.zero # type: ignore[attr-defined] + + # No need to run _collapse_arguments and _find_localzeros, see the comment + # in _satisfy_unique_summations_symbols. + if unique_summations_symbols is None: + # remove redundant args that are easily identified + args = cls._collapse_arguments(args, **assumptions) + + # find local zeros + args = cls._find_localzeros(args, **assumptions) + + args = frozenset(args) + + if not args: + return cls.identity # type: ignore[attr-defined] + + if len(args) == 1: + return list(args).pop() + + # base creation + obj = Expr.__new__(cls, *ordered(args), **assumptions) + obj._argset = args + + obj.unique_summations_symbols = unique_summations_symbols + return obj + + @classmethod + def _satisfy_unique_summations_symbols( + cls, args + ) -> set[sympy.core.symbol.Symbol] | None: + """ + One common case in some models is building expressions of the form + max(max(max(a+b...), c+d), e+f) which is simplified to max(a+b, c+d, e+f, ...). + For such expressions, we call the Max constructor X times (once for each nested + max) and the expression gets flattened. + + An expensive cost in constructing those expressions is running _collapse_arguments + and _find_localzeros. However, those two optimizations are unnecessary when the args + to max are all of the form a+b, c+d, ..etc where each term uses a unique set of symbols. + + This function is used to detect such properties of the expressions we are building + and if so inform that we do not need to run those optimizations. To detect those, + we store a property in the expression that tells that this expression is a min/max + operation over terms that use unique symbols "unique_summations_symbols". This property + also memoize the set of symbols used in all the terms to make it faster to detect this + property inductively. + + When we apply max to add a new term, all we need to do is check if the new term uses + unique symbols (with respect to existing terms and itself). + Example: + t = Max(a+b, c+d) ==> satisfies the property + Max(t, h+j) ==> h,j not in [a,b,c,d] => satisfy the property. + + The function returns None if the new expression does not satisfy the unique_summations_symbols + property. Otherwise, it returns a new set of unique symbols. + """ + if len(args) != 2: + return None + + (lhs, rhs) = ( + (args[1], args[0]) + if isinstance(args[1], MinMaxBase) + else (args[0], args[1]) + ) + + if not _is_symbols_binary_summation(rhs): + return None + + # base case max(a+b, c+d) ==> satisfies the property if a+b and c+d use unique symbols. + if _is_symbols_binary_summation(lhs): + return cls._unique_symbols(args) + + # inductive case max(t, h+j) ==> satisfies the property if h, j not in t.unique_summations_symbols + if isinstance(lhs, MinMaxBase): + lhs_unique_summations_symbols = getattr( + lhs, "unique_summations_symbols", None + ) + if lhs_unique_summations_symbols is not None: + return cls._unique_symbols([rhs], lhs_unique_summations_symbols) + + return None + + @classmethod + def _unique_symbols( + cls, args, initial_set: set[sympy.core.symbol.Symbol] | None = None + ) -> set[sympy.core.symbol.Symbol] | None: + """ + Return seen_symbols if all atoms in all args are all unique symbols, + else returns None. initial_set can be used to represent initial value for seen_symbols + """ + seen_symbols = set() if initial_set is None else initial_set.copy() + for arg in args: + for element in arg.atoms(): + if not isinstance(element, sympy.core.symbol.Symbol): + return None + elif element in seen_symbols: + return None + else: + seen_symbols.add(element) + return seen_symbols + + @classmethod + def _collapse_arguments(cls, args, **assumptions): + """Remove redundant args. + + Examples + ======== + + >>> from sympy import Min, Max + >>> from sympy.abc import a, b, c, d, e + + Any arg in parent that appears in any + parent-like function in any of the flat args + of parent can be removed from that sub-arg: + + >>> Min(a, Max(b, Min(a, c, d))) + Min(a, Max(b, Min(c, d))) + + If the arg of parent appears in an opposite-than parent + function in any of the flat args of parent that function + can be replaced with the arg: + + >>> Min(a, Max(b, Min(c, d, Max(a, e)))) + Min(a, Max(b, Min(a, c, d))) + """ + if not args: + return args + args = list(ordered(args)) + if cls is Min: + other = Max + else: + other = Min # type: ignore[assignment] + + # find global comparable max of Max and min of Min if a new + # value is being introduced in these args at position 0 of + # the ordered args + if args[0].is_number: + sifted = mins, maxs = [], [] # type: ignore[var-annotated] + for i in args: + for v in walk(i, Min, Max): + if v.args[0].is_comparable: + sifted[isinstance(v, Max)].append(v) + small = Min.identity + for i in mins: + v = i.args[0] + if v.is_number and (v < small) == True: # noqa: E712 + small = v + big = Max.identity + for i in maxs: + v = i.args[0] + if v.is_number and (v > big) == True: # noqa: E712 + big = v + # at the point when this function is called from __new__, + # there may be more than one numeric arg present since + # local zeros have not been handled yet, so look through + # more than the first arg + if cls is Min: + for arg in args: + if not arg.is_number: + break + if (arg < small) == True: # noqa: E712 + small = arg + elif cls == Max: + for arg in args: + if not arg.is_number: + break + if (arg > big) == True: # noqa: E712 + big = arg + T = None + if cls is Min: + if small != Min.identity: + other = Max + T = small + elif big != Max.identity: + other = Min # type: ignore[assignment] + T = big + if T is not None: + # remove numerical redundancy + for i in range(len(args)): + a = args[i] + if isinstance(a, other): + a0 = a.args[0] + if ( # noqa: E712 + (a0 > T) if other == Max else (a0 < T) # noqa: E712 + ) == True: # noqa: E712 + args[i] = cls.identity # type: ignore[attr-defined] + + # remove redundant symbolic args + def do(ai, a): + if not isinstance(ai, (Min, Max)): + return ai + cond = a in ai.args + if not cond: + return ai.func(*[do(i, a) for i in ai.args], evaluate=False) + if isinstance(ai, cls): + return ai.func(*[do(i, a) for i in ai.args if i != a], evaluate=False) + return a + + for i, a in enumerate(args): + args[i + 1 :] = [do(ai, a) for ai in args[i + 1 :]] + + # factor out common elements as for + # Min(Max(x, y), Max(x, z)) -> Max(x, Min(y, z)) + # and vice versa when swapping Min/Max -- do this only for the + # easy case where all functions contain something in common; + # trying to find some optimal subset of args to modify takes + # too long + + def factor_minmax(args): + is_other = lambda arg: isinstance(arg, other) # noqa: E731 + other_args, remaining_args = sift(args, is_other, binary=True) + if not other_args: + return args + + # Min(Max(x, y, z), Max(x, y, u, v)) -> {x,y}, ({z}, {u,v}) + arg_sets = [set(arg.args) for arg in other_args] + common = set.intersection(*arg_sets) + if not common: + return args + + new_other_args = list(common) + arg_sets_diff = [arg_set - common for arg_set in arg_sets] + + # If any set is empty after removing common then all can be + # discarded e.g. Min(Max(a, b, c), Max(a, b)) -> Max(a, b) + if all(arg_sets_diff): + other_args_diff = [other(*s, evaluate=False) for s in arg_sets_diff] + new_other_args.append(cls(*other_args_diff, evaluate=False)) + + other_args_factored = other(*new_other_args, evaluate=False) + return remaining_args + [other_args_factored] + + if len(args) > 1: + args = factor_minmax(args) + + return args + + @classmethod + def _new_args_filter(cls, arg_sequence): + """ + Generator filtering args. + + first standard filter, for cls.zero and cls.identity. + Also reshape ``Max(a, Max(b, c))`` to ``Max(a, b, c)``, + and check arguments for comparability + """ + for arg in arg_sequence: + # pre-filter, checking comparability of arguments + if ( + not isinstance(arg, Expr) + or arg.is_extended_real is False + or (arg.is_number and not arg.is_comparable) + ): + raise ValueError(f"The argument '{arg}' is not comparable.") + + if arg == cls.zero: # type: ignore[attr-defined] + raise ShortCircuit(arg) + elif arg == cls.identity: # type: ignore[attr-defined] + continue + elif arg.func == cls: + yield from arg.args + else: + yield arg + + @classmethod + def _find_localzeros(cls, values, **options): + """ + Sequentially allocate values to localzeros. + + When a value is identified as being more extreme than another member it + replaces that member; if this is never true, then the value is simply + appended to the localzeros. + + Unlike the sympy implementation, we only look for zero and one, we don't + do generic is connected test pairwise which is slow + """ + + # First, collapse all numeric arguments + other_values = set() + num_value = None + for arg in values: + if arg.is_Number: + if num_value is None: + num_value = arg + else: + if cls is Max: + num_value = max(num_value, arg) + elif cls is Min: + num_value = min(num_value, arg) + else: + raise AssertionError(f"impossible {cls}") + else: + other_values.add(arg) + + # Special cases when there is only one symbolic value + if num_value is None: + return other_values + + if len(other_values) == 0: + return {num_value} + + if len(other_values) == 1: + other_value = next(iter(other_values)) + if num_value in (0.0, 0) and other_value.is_nonnegative: + return other_values if cls is Max else {num_value} + if num_value == 1 and other_value.is_positive: + return other_values if cls is Max else {num_value} + + other_values.add(num_value) + return other_values + + _eval_is_algebraic = lambda s: _torf(i.is_algebraic for i in s.args) # noqa: E731 + _eval_is_antihermitian = lambda s: _torf( # noqa: E731 + i.is_antihermitian + for i in s.args # noqa: E731 + ) # noqa: E731 + _eval_is_commutative = lambda s: _torf( # noqa: E731 + i.is_commutative + for i in s.args # noqa: E731 + ) # noqa: E731 + _eval_is_complex = lambda s: _torf(i.is_complex for i in s.args) # noqa: E731 + _eval_is_composite = lambda s: _torf(i.is_composite for i in s.args) # noqa: E731 + _eval_is_even = lambda s: _torf(i.is_even for i in s.args) # noqa: E731 + _eval_is_finite = lambda s: _torf(i.is_finite for i in s.args) # noqa: E731 + _eval_is_hermitian = lambda s: _torf(i.is_hermitian for i in s.args) # noqa: E731 + _eval_is_imaginary = lambda s: _torf(i.is_imaginary for i in s.args) # noqa: E731 + _eval_is_infinite = lambda s: _torf(i.is_infinite for i in s.args) # noqa: E731 + _eval_is_integer = lambda s: _torf(i.is_integer for i in s.args) # noqa: E731 + _eval_is_irrational = lambda s: _torf(i.is_irrational for i in s.args) # noqa: E731 + _eval_is_negative = lambda s: _torf(i.is_negative for i in s.args) # noqa: E731 + _eval_is_noninteger = lambda s: _torf(i.is_noninteger for i in s.args) # noqa: E731 + _eval_is_nonnegative = lambda s: _torf( # noqa: E731 + i.is_nonnegative + for i in s.args # noqa: E731 + ) # noqa: E731 + _eval_is_nonpositive = lambda s: _torf( # noqa: E731 + i.is_nonpositive + for i in s.args # noqa: E731 + ) # noqa: E731 + _eval_is_nonzero = lambda s: _torf(i.is_nonzero for i in s.args) # noqa: E731 + _eval_is_odd = lambda s: _torf(i.is_odd for i in s.args) # noqa: E731 + _eval_is_polar = lambda s: _torf(i.is_polar for i in s.args) # noqa: E731 + _eval_is_positive = lambda s: _torf(i.is_positive for i in s.args) # noqa: E731 + _eval_is_prime = lambda s: _torf(i.is_prime for i in s.args) # noqa: E731 + _eval_is_rational = lambda s: _torf(i.is_rational for i in s.args) # noqa: E731 + _eval_is_real = lambda s: _torf(i.is_real for i in s.args) # noqa: E731 + _eval_is_extended_real = lambda s: _torf( # noqa: E731 + i.is_extended_real + for i in s.args # noqa: E731 + ) # noqa: E731 + _eval_is_transcendental = lambda s: _torf( # noqa: E731 + i.is_transcendental + for i in s.args # noqa: E731 + ) # noqa: E731 + _eval_is_zero = lambda s: _torf(i.is_zero for i in s.args) # noqa: E731 + + +class Max(MinMaxBase, Application): # type: ignore[misc] + r""" + Return, if possible, the maximum value of the list. + """ + + zero = S.Infinity + identity = S.NegativeInfinity + + def _eval_is_positive(self): # type:ignore[override] + return fuzzy_or(a.is_positive for a in self.args) # type: ignore[attr-defined] + + def _eval_is_nonnegative(self): # type:ignore[override] + return fuzzy_or(a.is_nonnegative for a in self.args) # type: ignore[attr-defined] + + def _eval_is_negative(self): # type:ignore[override] + return fuzzy_and(a.is_negative for a in self.args) + + +class Min(MinMaxBase, Application): # type: ignore[misc] + """ + Return, if possible, the minimum value of the list. + """ + + zero = S.NegativeInfinity + identity = S.Infinity + + def _eval_is_positive(self): # type:ignore[override] + return fuzzy_and(a.is_positive for a in self.args) # type: ignore[attr-defined] + + def _eval_is_nonnegative(self): # type:ignore[override] + return fuzzy_and(a.is_nonnegative for a in self.args) # type: ignore[attr-defined] + + def _eval_is_negative(self): # type:ignore[override] + return fuzzy_or(a.is_negative for a in self.args) + + +def safe_pow(base, exp): + sign = 1 + if base < 0: + base = -base + sign = 1 if exp % 2 == 0 else -1 + return sign * _safe_pow(base, exp) + + +# Prevent people from overflowing pow +def _safe_pow(base, exponent): + if exponent < 0: + raise ValueError("Exponent must be non-negative.") + + if exponent == 0: + return 1 + + half_exp = safe_pow(base, exponent // 2) + if half_exp is int_oo: + return int_oo + + # TODO: microoptimization is to avoid overflowing into arbitrary precision + # and detect overflow prior to doing operations + + result = half_exp * half_exp + if result > sys.maxsize: + return int_oo + + if exponent % 2 == 1: + result *= base + if result > sys.maxsize: + return int_oo + + return result + + +class PowByNatural(sympy.Function): + is_integer = True + + precedence: int = 50 # precedence of mul + + @classmethod + def eval(cls, base, exp): + if isinstance(base, sympy.Integer) and isinstance(exp, sympy.Integer): + r = safe_pow(base, exp) + if r in (-int_oo, int_oo): + return r + return sympy.Integer(r) + if isinstance(exp, sympy.Integer): + # Rely on regular sympy Pow for this (note that iterated + # multiplication turns into a Pow anyway, you can't escape!!) + return sympy.Pow(base, exp) + if exp in (int_oo, sympy.oo): + if base.is_nonnegative: + return int_oo + elif base.is_negative: + return sympy.zoo # this is apparently what (-2)**sympy.oo does + # NB: do NOT translate into sympy.Pow, we will lose knowledge that exp + # is a natural number if we do + + +# base is assumed to be nonnegative, thereby prevent complex numbers from +# occurring +class FloatPow(sympy.Function): + is_real = True + + precedence: int = 60 # precedence of pow + + @classmethod + def eval(cls, base, exp): + # NB: These test sympy.Number, not sympy.Float, because: + # - Sometimes we may have sympy.oo or int_oo, and that's not a Float + # (but coerces to math.Inf) + # - Sometimes Float(0.0) will unpredictably decay to Integer(0), + # but we should still accept it in floatey contexts + if isinstance(base, sympy.Number) and isinstance(exp, sympy.Number): + return sympy.Float(float(base) ** float(exp)) + # NB: do not do any nontrivial reasoning + + +# Overloaded to be compatible with regular Python. +# https://github.com/pytorch/pytorch/issues/90900 +# +# In particular, sympy division is willing to simplify x/x == 1 +# where 1 is an integer, but this must be a float if x was float. +class FloatTrueDiv(sympy.Function): + is_real = True + + precedence: int = 35 # lower precedence than add + + @classmethod + def eval(cls, base, divisor): + # assert base.is_integer is not True, base + # assert divisor.is_integer is not True, divisor + + if divisor.is_zero: + raise ZeroDivisionError("division by zero") + + if isinstance(base, sympy.Number) and isinstance(divisor, sympy.Number): + return sympy.Float(float(base) / float(divisor)) + + +# Overloaded to be compatible with regular Python. We distinguish this from +# FloatTrueDiv, because the code generation has to be different for this case: +# Python has a fancy algorithm for integer true division that isn't just +# "promote both arguments to float and use float division", so you need to +# codegen it differently. While technically you can work it out from the +# types of the input, this is often inconvenient to do in Inductor codegen, +# so just have a different operator +# NB: Right now, Inductor codegen doesn't implement this correctly lol +class IntTrueDiv(sympy.Function): + is_real = True + + precedence: int = 35 # lower precedence than add + + @classmethod + def eval(cls, base, divisor): + if divisor.is_zero: + raise ZeroDivisionError("division by zero") + + if ( + isinstance(base, sympy.Number) + and isinstance(divisor, sympy.Number) + and (is_infinite(base) or is_infinite(divisor)) + ): + # Don't have to worry about precision here, you're getting zero or + # inf from the division + return sympy.Float(float(base) / float(divisor)) + if isinstance(base, sympy.Integer) and isinstance(divisor, sympy.Integer): + return sympy.Float(int(base) / int(divisor)) + + def _ccode(self, printer) -> str: + base = printer.parenthesize(self.args[0], PRECEDENCE["Atom"] - 0.5) + divisor = printer.parenthesize(self.args[1], PRECEDENCE["Atom"] - 0.5) + return f"((int){base}/(int){divisor})" + + +# TODO: As an indicator, this != 0 implies == 1 (and vice versa). +# Because we do not have the ability to guard on the stride permutation +# at the moment, it is hard to make further inferences when this is true, +# as although we know the tensor is contiguous in *some* layout, we don't +# know which one (however, you could, for example, make the inference that +# reshaping this to a 1D tensor can be guard-free.) +class IsNonOverlappingAndDenseIndicator(sympy.Function): + is_integer = True + + @classmethod + def eval(cls, *args): + if len(args) % 2 != 0: + raise AssertionError( + f"expected an even number of arguments, got {len(args)}" + ) + dim = len(args) // 2 + sizes = args[0:dim] + strides = args[dim:] + + # sym_node imported in torch.__init__. Local import to avoid an import cycle + from torch.fx.experimental.symbolic_shapes import ( + eval_is_non_overlapping_and_dense, + ) + + if all(isinstance(a, sympy.Integer) for a in args): + return eval_is_non_overlapping_and_dense( + [int(a) for a in sizes], [int(a) for a in strides] + ) + + if dim == 1: + # Manually implement the rank one short circuit + if strides[0].is_Number and strides[0] == 1: + return 1 + + if sizes[0].is_Number and sizes[0] < 2: + return 1 + + # return 0 case covered by case above + + # TODO: Inability to access size-obliviousness sucks: if we have a + # size oblivious test on a size-like unbacked SymInt, we could + # confidently return zero when we have a size-like u0 stride + # and a size-like u1 size. Maybe a fancy ValueRanges analysis for + # this function could help figure this out. + + if all(isinstance(a, sympy.Integer) for a in strides): + if dim == 0: + raise AssertionError("dim must not be zero") + # When all strides are integral, we can sort, and the size for the + # largest stride doesn't matter and can be arbitrarily symbolic + s_sizes, s_strides = zip( + *sorted(zip(sizes, strides, strict=True), key=operator.itemgetter(1)), + strict=True, + ) + # Put something arbitrary in the max size spot, it'll be ignored + if all(isinstance(a, sympy.Integer) for a in s_sizes[:-1]): + s_sizes = s_sizes[:-1] + (42,) + # We can reuse the regular eval, because it is invariant to + # permutation of dimensions + return eval_is_non_overlapping_and_dense( + [int(a) for a in s_sizes], [int(a) for a in s_strides] + ) + + return None + + +# NB: this is inconsistent with math.trunc in Python +class TruncToFloat(sympy.Function): + is_real = True + + @classmethod + def eval(cls, number): + if number in (sympy.oo, -sympy.oo): + return number + # assert number.is_integer is not True, number + if isinstance(number, sympy.Number): + # NB: It is safe to use truncation to integer, which is what + # math.trunc does, as Python integers are arbitrary precision and + # so we are guaranteed not to lose precision when we do this + return sympy.Float(math.trunc(float(number))) + + +class TruncToInt(sympy.Function): + is_integer = True + + @classmethod + def eval(cls, number): + # assert number.is_integer is not True, number + if number in (sympy.oo, int_oo): + return int_oo + if number in (-sympy.oo, -int_oo): + return -int_oo + if isinstance(number, sympy.Number): + return sympy.Integer(math.trunc(float(number))) + + +# This is float -> int +class RoundToInt(sympy.Function): + is_integer = True + + @classmethod + def eval(cls, number): + # assert number.is_integer is not True, number + + if number is sympy.oo: + return int_oo + if number is -sympy.oo: + return -int_oo + if isinstance(number, sympy.Number): + return sympy.Integer(round(float(number), 0)) + + +# To get float -> int, Python style round semantics. +# +# x = PyFloat_AsDouble(self); +# if (o_ndigits == Py_None) { +# /* single-argument round or with None ndigits: +# * round to nearest integer */ +# rounded = round(x); +# if (fabs(x-rounded) == 0.5) +# /* halfway case: round to even */ +# rounded = 2.0*round(x/2.0); +# return PyLong_FromDouble(rounded); +# } + + +# NB: Like Round, this only ever returns floats. ndigits cannot be None +class RoundDecimal(sympy.Function): + is_real = True + + @classmethod + def eval(cls, number, ndigits): + # assert number.is_integer is not True, number + + if isinstance(number, sympy.Number) and isinstance(ndigits, sympy.Integer): + return sympy.Float(round(float(number), int(ndigits))) + + +class ToFloat(sympy.Function): + is_real = True + + @classmethod + def eval(cls, number): + if number in [sympy.oo, -sympy.oo]: + return number + + if isinstance(number, sympy.Integer): + return sympy.Float(int(number)) + if number is int_oo: + return sympy.oo + if number is -int_oo: + return -sympy.oo + + +class Identity(sympy.Function): + """ + Prevents expansion and other optimizations + """ + + precedence = 10 + + def __repr__(self) -> str: # type: ignore[override] + return f"Identity({self.args[0]})" + + def _sympystr(self, printer) -> str: + """Controls how sympy's StrPrinter prints this""" + return f"({printer.doprint(self.args[0])})" + + def _eval_is_real(self): + return self.args[0].is_real + + def _eval_is_integer(self): + return self.args[0].is_integer # type: ignore[attr-defined] + + @property + def is_number(self): + # Treat Identity as numeric only when the argument is comparable. + # This avoids creating numeric non-comparable Identity(I) terms. + return bool(self.args[0].is_number and self.args[0].is_comparable) + + @property + def is_comparable(self): + # Delegate comparability to the wrapped argument. + return bool(self.args[0].is_comparable) + + def _eval_expand_identity(self, **hints): + # Removes the identity op. + return self.args[0] + + def __int__(self) -> int: + return int(self.args[0]) + + def _identity_atom_compare(self, other, op): + """ + Fast path for comparing wrapped numeric atomics against other numeric atomics. + Keep compound expressions on SymPy's default symbolic path. + """ + arg = self.args[0] + if isinstance(other, int): + other = sympy.Integer(other) + if not isinstance(other, sympy.Expr): + return None + if not (arg.is_Atom and arg.is_number and arg.is_comparable): + return None + if not (other.is_Atom and other.is_number and other.is_comparable): + return None + return sympy.S.true if op(arg, other) else sympy.S.false + + def __ge__(self, other): + out = self._identity_atom_compare(other, lambda a, b: a >= b) + return out if out is not None else super().__ge__(other) + + def __gt__(self, other): + out = self._identity_atom_compare(other, lambda a, b: a > b) + return out if out is not None else super().__gt__(other) + + def __le__(self, other): + out = self._identity_atom_compare(other, lambda a, b: a <= b) + return out if out is not None else super().__le__(other) + + def __lt__(self, other): + out = self._identity_atom_compare(other, lambda a, b: a < b) + return out if out is not None else super().__lt__(other) + + def __float__(self) -> float: + return float(self.args[0]) + + +def make_opaque_unary_fn(name): + class OpaqueUnaryFn(sympy.Function): + """ + Unlike the builtin sympy functions on real numbers like sympy.sqrt, + these equivalents do not do any nontrivial reasoning besides + constant propagation. This helps avoid performing transformations + that are valid for real numbers but are invalid for floating point; + in particular, while we are willing to make optimizations that change + numerics for Tensor compute, we are NOT willing to make optimizations + that change numerics for size compute. + """ + + _torch_handler_name = name + _torch_unpickler = make_opaque_unary_fn + + @classmethod + def eval(cls, a): + if isinstance(a, (sympy.Integer, sympy.Float)): + # Python converts to float64 before computing, c.f. + # >>> math.sin(2**53+1) + # -0.848925964814655 + # >>> math.sin(float(2**53+1)) + # -0.848925964814655 + try: + return sympy.Float(getattr(math, name)(float(a))) + # Just use sympy semantics for infinity/overflow, you might get some + # weird objects but ask silly questions, get silly answers + except OverflowError: + return getattr(sympy, name)(a) + elif a in [sympy.oo, -sympy.oo, sympy.zoo, -sympy.zoo, int_oo, -int_oo]: + if a is int_oo: + a = sympy.oo + if a is -int_oo: + a = -sympy.oo + if name == "log2": + return sympy.log(a, 2) + return getattr(sympy, name)(a) + return None + + nm = "OpaqueUnaryFn_" + name + OpaqueUnaryFn.__name__ = nm + OpaqueUnaryFn.__qualname__ = nm + + return OpaqueUnaryFn + + +# Keep in sync with math_op_names in torch/fx/experimental/sym_node.py +OpaqueUnaryFn_sqrt = make_opaque_unary_fn("sqrt") +OpaqueUnaryFn_cos = make_opaque_unary_fn("cos") +OpaqueUnaryFn_cosh = make_opaque_unary_fn("cosh") +OpaqueUnaryFn_sin = make_opaque_unary_fn("sin") +OpaqueUnaryFn_sinh = make_opaque_unary_fn("sinh") +OpaqueUnaryFn_tan = make_opaque_unary_fn("tan") +OpaqueUnaryFn_tanh = make_opaque_unary_fn("tanh") +OpaqueUnaryFn_asin = make_opaque_unary_fn("asin") +OpaqueUnaryFn_acos = make_opaque_unary_fn("acos") +OpaqueUnaryFn_atan = make_opaque_unary_fn("atan") +OpaqueUnaryFn_exp = make_opaque_unary_fn("exp") +OpaqueUnaryFn_log = make_opaque_unary_fn("log") +OpaqueUnaryFn_asinh = make_opaque_unary_fn("asinh") +OpaqueUnaryFn_log2 = make_opaque_unary_fn("log2") + + +def make_opaque_bitwise_fn(name, real_op_name): + if name == "bitwise_and": + prec = PRECEDENCE["BitwiseAnd"] + elif name == "bitwise_xor": + prec = PRECEDENCE["BitwiseXor"] + elif name == "bitwise_or": + prec = PRECEDENCE["BitwiseOr"] + else: + raise AssertionError(f"unrecognized {name}") + + class BitwiseFn(sympy.Function): + _torch_handler_name = name + precedence: int = prec + _torch_unpickler = functools.partial( + make_opaque_bitwise_fn, real_op_name=real_op_name + ) + + @classmethod + def eval(cls, a, b): + if a.is_Boolean and b.is_Boolean: + return getattr(operator, real_op_name)(a, b) + if a.is_Boolean: + a = sympy.Integer(1 if a else 0) + if b.is_Boolean: + b = sympy.Integer(1 if b else 0) + if isinstance(a, (sympy.Integer, int)) and isinstance( + b, (sympy.Integer, int) + ): + return sympy.Integer(getattr(operator, real_op_name)(int(a), int(b))) + return None + + nm = "BitwiseFn_" + name + BitwiseFn.__name__ = nm + BitwiseFn.__qualname__ = nm + + return BitwiseFn + + +BitwiseFn_bitwise_and = make_opaque_bitwise_fn("bitwise_and", "and_") +BitwiseFn_bitwise_or = make_opaque_bitwise_fn("bitwise_or", "or_") +BitwiseFn_bitwise_xor = make_opaque_bitwise_fn("bitwise_xor", "xor") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/interp.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/interp.py new file mode 100644 index 0000000000000000000000000000000000000000..6eca9e389d85ae452cbf357d01ca9278239a617d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/interp.py @@ -0,0 +1,228 @@ +# mypy: allow-untyped-defs +""" +This is a simple interpreter for Sympy expressions that dispatches to +classes following the torch._inductor.virtualized calling convention. +For directness, the interpreter takes the handler directly rather than +consulting the TLS. It does not use most of the methods on the full +handler; only those with corresponding Sympy expressions. To see an example +of a full handler, see torch.utils._sympy.value_ranges.ValueRangeAnalysis. +""" + +import functools +import logging +from typing import Any + +import sympy +from sympy.logic.boolalg import Boolean as SympyBoolean, BooleanAtom + +import torch + +from .functions import ( + BitwiseFn_bitwise_and, + BitwiseFn_bitwise_or, + BitwiseFn_bitwise_xor, + CeilToInt, + CleanDiv, + FloatPow, + FloatTrueDiv, + FloorDiv, + FloorToInt, + Identity, + IntTrueDiv, + IsNonOverlappingAndDenseIndicator, + Max, + Min, + Mod, + ModularIndexing, + OpaqueUnaryFn_log2, + PowByNatural, + PythonMod, + RoundDecimal, + RoundToInt, + ToFloat, + TruncToFloat, + TruncToInt, + Where, +) + + +log = logging.getLogger(__name__) + + +# TODO: Dedupe this with SYMPY_INTERP + + +@functools.cache +def handlers(): + # TODO add CeilDiv (it doesn't appear in the index_expr) + + # TODO default to some decompositions if the interpreter doesn't have them + # like decomposing ModularIndexing or implementing Le(a,b) as Ge(b, a) + + HANDLERS = { + sympy.Or: "or_", + sympy.And: "and_", + sympy.Eq: "eq", + sympy.Ne: "ne", + sympy.Lt: "lt", + sympy.Gt: "gt", + sympy.Le: "le", + sympy.Ge: "ge", + sympy.Not: "not_", + IntTrueDiv: "int_truediv", + FloatTrueDiv: "truediv", + FloorDiv: "floordiv", + CleanDiv: "floordiv", # TODO: hmm? + TruncToFloat: "trunc", + Where: "where", + sympy.Add: "add", + sympy.Mul: "mul", + FloatPow: "pow", + PowByNatural: "pow_by_natural", + # sympy simplifies x * x into Pow(x, 2), so we need to handle this. + # Do NOT use builtin Pow for floats + # TODO: There is a hazard here, if we have float * float it will + # also get turned into Pow(float, 2) but we don't want this because + # pow_by_natural is assumed to only be integers. Probably the fix is + # to add a FloatMul to impede this optimization + sympy.Pow: "pow_by_natural", + Mod: "mod", + PythonMod: "python_mod", + # TODO: Inductor can generate these, but it's ill-specified which + # semantics were intended here. Needs to be cleaned up along with + # FloorDiv in a bigger cleanup + sympy.Mod: "mod", + sympy.Abs: "abs", + sympy.log: "log", + sympy.exp: "exp", + sympy.Min: "minimum", + sympy.Max: "maximum", + Min: "minimum", + Max: "maximum", + ModularIndexing: "modular_indexing", + sympy.functions.elementary.piecewise.ExprCondPair: "expr_cond_pair", + sympy.Piecewise: "piecewise", + Identity: "identity", + IsNonOverlappingAndDenseIndicator: "is_non_overlapping_and_dense_indicator", + RoundDecimal: "round_decimal", + # TODO: do the rest of the opaque unary functions... + OpaqueUnaryFn_log2: "log2", + BitwiseFn_bitwise_and: "bitwise_and", + BitwiseFn_bitwise_or: "bitwise_or", + BitwiseFn_bitwise_xor: "bitwise_xor", + } + # TODO: This is kind of pointless, we shouldn't be generating sympy.sin + # for these functions, they should be Opaque instead + for name in ["cos", "sin", "tan", "sinh", "cosh", "tanh", "asin", "acos", "atan"]: + HANDLERS[getattr(sympy, name)] = name + + return HANDLERS + + +ASSOCIATIVE_OPS = {"minimum", "maximum", "mul", "add", "and_", "or_"} + + +def _run_sympy_handler(analysis, args, expr, index_dtype=torch.int64): + # Special cases + if isinstance(expr, sympy.Pow) and isinstance( + expr.args[1], sympy.core.numbers.Half + ): + return analysis.sqrt(args[0]) + if isinstance(expr, ToFloat): + return analysis.to_dtype(args[0], torch.float64) + + # These handlers are special because they take an extra dtype argument + # specifying what they should convert to, and we need to appropriately set + # this up when we convert from Sympy. A reasonable default when you + # are translating is to conservatively do int64, and then narrow these + # arguments later when you discover you can narrow the index range. But + # if you already know that 32-bit indexing is OK, you can directly do the + # sympy translation with index_dtype=torch.int32 + INDEX_DTYPE_HANDLERS = { + TruncToInt: "trunc_to_int", + sympy.floor: "floor_to_int", + sympy.ceiling: "ceil_to_int", + FloorToInt: "floor_to_int", + CeilToInt: "ceil_to_int", + RoundToInt: "round_to_int", + } + if (handler_name := INDEX_DTYPE_HANDLERS.get(expr.func)) is not None: + return getattr(analysis, handler_name)(*args, index_dtype) + + # Fastpath for n-ary integral addition + if expr.func is sympy.Add and expr.is_integer and hasattr(analysis, "sym_sum"): + r = analysis.sym_sum(args) + log.debug("sym_sum(%s) -> %s", args, r) + return r + + if hasattr(expr.func, "_torch_handler_name"): + handler_name = expr.func._torch_handler_name + else: + handler_name = handlers()[expr.func] + handler = getattr(analysis, handler_name) + try: + if handler_name in ASSOCIATIVE_OPS: + if len(args) <= 1: + raise AssertionError("associative op needs >1 args") + acc = handler(args[0], args[1]) + for i in range(2, len(args)): + acc = handler(acc, args[i]) + log.debug("%s(%s) -> %s", handler_name, args, acc) + return acc + else: + r = handler(*args) + log.debug("%s(%s) -> %s", handler_name, args, r) + return r + except NotImplementedError: + raise + except Exception: + log.warning("failed while executing %s(%s)", handler_name, args) + raise + + +_nil = object() + + +def sympy_interp( + analysis, + env: dict[sympy.Symbol, Any], + expr: sympy.Expr | SympyBoolean, + *, + index_dtype=torch.int64, + missing_handler=None, +): + # Handle base cases + dtype = None + if isinstance(expr, BooleanAtom): + dtype = torch.bool + elif isinstance(expr, sympy.Integer): + dtype = torch.int64 + elif isinstance(expr, sympy.Number): + dtype = torch.double + + if dtype is not None: + return analysis.constant(expr, dtype) + elif isinstance(expr, sympy.Symbol): + if (r := env.get(expr, _nil)) is not _nil: + return r + elif missing_handler: + return missing_handler(expr) + else: + raise KeyError(expr) + + # Recursive case + return _run_sympy_handler( + analysis, + [ + sympy_interp( + analysis, + env, + arg, + index_dtype=index_dtype, + missing_handler=missing_handler, + ) + for arg in expr.args + ], + expr, + index_dtype=index_dtype, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/numbers.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/numbers.py new file mode 100644 index 0000000000000000000000000000000000000000..792ba9f2e8567bff740b605fd2f6ac32bae5149c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/numbers.py @@ -0,0 +1,417 @@ +# mypy: allow-untyped-defs +import mpmath.libmp as mlib # type: ignore[import-untyped] +import sympy +from sympy import Expr +from sympy.core.decorators import _sympifyit +from sympy.core.expr import AtomicExpr +from sympy.core.numbers import Number +from sympy.core.parameters import global_parameters +from sympy.core.singleton import S, Singleton + + +# pyrefly: ignore [invalid-inheritance] +class IntInfinity(Number, metaclass=Singleton): + r"""Positive integer infinite quantity. + + Integer infinity is a value in an extended integers which + is greater than all other integers. We distinguish it from + sympy's existing notion of infinity in that it reports that + it is_integer. + + Infinity is a singleton, and can be accessed by ``S.IntInfinity``, + or can be imported as ``int_oo``. + """ + + # NB: We can't actually mark this as infinite, as integer and infinite are + # inconsistent assumptions in sympy. We also report that we are complex, + # different from sympy.oo + + is_integer = True + is_commutative = True + is_number = True + is_extended_real = True + is_comparable = True + is_extended_positive = True + is_prime = False + + # Ensure we get dispatched to before plain numbers + _op_priority = 100.0 + + __slots__ = () + + def __new__(cls): + return AtomicExpr.__new__(cls) + + def _sympystr(self, printer) -> str: + return "int_oo" + + def _eval_subs(self, old, new): + if self == old: + return new + + # We could do these, not sure about it + """ + def _eval_evalf(self, prec=None): + return Float('inf') + + def evalf(self, prec=None, **options): + return self._eval_evalf(prec) + """ + + @_sympifyit("other", NotImplemented) + def __add__(self, other): + if isinstance(other, Number) and global_parameters.evaluate: + if other in (S.Infinity, S.NegativeInfinity): + return other + if other in (S.NegativeIntInfinity, S.NaN): + return S.NaN + return self + return Number.__add__(self, other) + + __radd__ = __add__ + + @_sympifyit("other", NotImplemented) + def __sub__(self, other): + if isinstance(other, Number) and global_parameters.evaluate: + if other is S.Infinity: + return S.NegativeInfinity + if other is S.NegativeInfinity: + return S.Infinity + if other in (S.IntInfinity, S.NaN): + return S.NaN + return self + return Number.__sub__(self, other) + + @_sympifyit("other", NotImplemented) + def __rsub__(self, other): + return (-self).__add__(other) + + @_sympifyit("other", NotImplemented) + def __mul__(self, other): + if isinstance(other, Number) and global_parameters.evaluate: + if other.is_zero or other is S.NaN: + return S.NaN + if other.is_extended_positive: + return self + return S.NegativeIntInfinity + return Number.__mul__(self, other) + + __rmul__ = __mul__ + + @_sympifyit("other", NotImplemented) + def __truediv__(self, other): + if isinstance(other, Number) and global_parameters.evaluate: + if other in ( + S.Infinity, + S.IntInfinity, + S.NegativeInfinity, + S.NegativeIntInfinity, + S.NaN, + ): + return S.NaN + if other.is_extended_nonnegative: + return S.Infinity # truediv produces float + return S.NegativeInfinity # truediv produces float + return Number.__truediv__(self, other) + + def __abs__(self): + return S.IntInfinity + + def __neg__(self): + return S.NegativeIntInfinity + + def _eval_power(self, expt): + if expt.is_extended_positive: + return S.IntInfinity + if expt.is_extended_negative: + return S.Zero + if expt is S.NaN: + return S.NaN + if expt is S.ComplexInfinity: + return S.NaN + if expt.is_extended_real is False and expt.is_number: + from sympy.functions.elementary.complexes import re + + expt_real = re(expt) + if expt_real.is_positive: + return S.ComplexInfinity + if expt_real.is_negative: + return S.Zero + if expt_real.is_zero: + return S.NaN + + return self ** expt.evalf() + + def _as_mpf_val(self, prec): + return mlib.finf + + def __hash__(self): + return super().__hash__() + + def __eq__(self, other): + return other is S.IntInfinity + + def __ne__(self, other): + return other is not S.IntInfinity + + def __gt__(self, other): + if other is S.Infinity: + return sympy.false # sympy.oo > int_oo + elif other is S.IntInfinity: + return sympy.false # consistency with sympy.oo + else: + return sympy.true + + def __ge__(self, other): + if other is S.Infinity: + return sympy.false # sympy.oo > int_oo + elif other is S.IntInfinity: + return sympy.true # consistency with sympy.oo + else: + return sympy.true + + def __lt__(self, other): + if other is S.Infinity: + return sympy.true # sympy.oo > int_oo + elif other is S.IntInfinity: + return sympy.false # consistency with sympy.oo + else: + return sympy.false + + def __le__(self, other): + if other is S.Infinity: + return sympy.true # sympy.oo > int_oo + elif other is S.IntInfinity: + return sympy.true # consistency with sympy.oo + else: + return sympy.false + + @_sympifyit("other", NotImplemented) + def __mod__(self, other): + if not isinstance(other, Expr): + return NotImplemented + return S.NaN + + __rmod__ = __mod__ + + def floor(self): + return self + + def ceiling(self): + return self + + +int_oo = S.IntInfinity + + +def is_infinite(expr) -> bool: + """Check if an expression is any type of infinity (positive or negative). + + This handles both sympy's built-in infinities (oo, -oo) and PyTorch's + integer infinities (int_oo, -int_oo). + + Note: We cannot rely on sympy's is_finite property because IntInfinity + and NegativeIntInfinity have is_integer=True, which implies is_finite=True + in sympy's assumption system. + """ + return expr in ( + S.Infinity, + S.NegativeInfinity, + S.IntInfinity, + S.NegativeIntInfinity, + ) + + +# pyrefly: ignore [invalid-inheritance] +class NegativeIntInfinity(Number, metaclass=Singleton): + """Negative integer infinite quantity. + + NegativeInfinity is a singleton, and can be accessed + by ``S.NegativeInfinity``. + + See Also + ======== + + IntInfinity + """ + + # Ensure we get dispatched to before plain numbers + _op_priority = 100.0 + + is_integer = True + is_extended_real = True + is_commutative = True + is_comparable = True + is_extended_negative = True + is_number = True + is_prime = False + + __slots__ = () + + def __new__(cls): + return AtomicExpr.__new__(cls) + + def _eval_subs(self, old, new): + if self == old: + return new + + def _sympystr(self, printer) -> str: + return "-int_oo" + + """ + def _eval_evalf(self, prec=None): + return Float('-inf') + + def evalf(self, prec=None, **options): + return self._eval_evalf(prec) + """ + + @_sympifyit("other", NotImplemented) + def __add__(self, other): + if isinstance(other, Number) and global_parameters.evaluate: + if other is S.Infinity: + return S.Infinity + if other in (S.IntInfinity, S.NaN): + return S.NaN + return self + return Number.__add__(self, other) + + __radd__ = __add__ + + @_sympifyit("other", NotImplemented) + def __sub__(self, other): + if isinstance(other, Number) and global_parameters.evaluate: + if other is S.NegativeInfinity: + return S.Infinity + if other in (S.NegativeIntInfinity, S.NaN): + return S.NaN + return self + return Number.__sub__(self, other) + + @_sympifyit("other", NotImplemented) + def __rsub__(self, other): + return (-self).__add__(other) + + @_sympifyit("other", NotImplemented) + def __mul__(self, other): + if isinstance(other, Number) and global_parameters.evaluate: + if other.is_zero or other is S.NaN: + return S.NaN + if other.is_extended_positive: + return self + return S.IntInfinity + return Number.__mul__(self, other) + + __rmul__ = __mul__ + + @_sympifyit("other", NotImplemented) + def __truediv__(self, other): + if isinstance(other, Number) and global_parameters.evaluate: + if other in ( + S.Infinity, + S.IntInfinity, + S.NegativeInfinity, + S.NegativeIntInfinity, + S.NaN, + ): + return S.NaN + if other.is_extended_nonnegative: + return self + return S.Infinity # truediv returns float + return Number.__truediv__(self, other) + + def __abs__(self): + return S.IntInfinity + + def __neg__(self): + return S.IntInfinity + + def _eval_power(self, expt): + if expt.is_number: + if expt in ( + S.NaN, + S.Infinity, + S.NegativeInfinity, + S.IntInfinity, + S.NegativeIntInfinity, + ): + return S.NaN + + if isinstance(expt, sympy.Integer) and expt.is_extended_positive: + if expt.is_odd: + return S.NegativeIntInfinity + else: + return S.IntInfinity + + inf_part = S.IntInfinity**expt + s_part = S.NegativeOne**expt + if inf_part == 0 and s_part.is_finite: + return inf_part + if ( + inf_part is S.ComplexInfinity + and s_part.is_finite + and not s_part.is_zero + ): + return S.ComplexInfinity + return s_part * inf_part + + def _as_mpf_val(self, prec): + return mlib.fninf + + def __hash__(self): + return super().__hash__() + + def __eq__(self, other): + return other is S.NegativeIntInfinity + + def __ne__(self, other): + return other is not S.NegativeIntInfinity + + def __gt__(self, other): + if other is S.NegativeInfinity: + return sympy.true # -sympy.oo < -int_oo + elif other is S.NegativeIntInfinity: + return sympy.false # consistency with sympy.oo + else: + return sympy.false + + def __ge__(self, other): + if other is S.NegativeInfinity: + return sympy.true # -sympy.oo < -int_oo + elif other is S.NegativeIntInfinity: + return sympy.true # consistency with sympy.oo + else: + return sympy.false + + def __lt__(self, other): + if other is S.NegativeInfinity: + return sympy.false # -sympy.oo < -int_oo + elif other is S.NegativeIntInfinity: + return sympy.false # consistency with sympy.oo + else: + return sympy.true + + def __le__(self, other): + if other is S.NegativeInfinity: + return sympy.false # -sympy.oo < -int_oo + elif other is S.NegativeIntInfinity: + return sympy.true # consistency with sympy.oo + else: + return sympy.true + + @_sympifyit("other", NotImplemented) + def __mod__(self, other): + if not isinstance(other, Expr): + return NotImplemented + return S.NaN + + __rmod__ = __mod__ + + def floor(self): + return self + + def ceiling(self): + return self + + def as_powers_dict(self): + return {S.NegativeOne: 1, S.IntInfinity: 1} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/printers.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/printers.py new file mode 100644 index 0000000000000000000000000000000000000000..e9ca80812b04db0af467c8c70f0133987b585de8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/printers.py @@ -0,0 +1,666 @@ +import sys + +import sympy +from sympy.printing.precedence import PRECEDENCE, precedence +from sympy.printing.str import StrPrinter + + +INDEX_TYPE = "int64_t" +INDEX_TYPE_MAX = (1 << 63) - 1 +INDEX_TYPE_MIN = -1 << 63 + + +# This printer contains rules that are supposed to be generic for both C/C++ and +# Python +class ExprPrinter(StrPrinter): + # override this so that _print_FloorDiv is used + printmethod = "_torch_sympystr" + + def _print_Mul(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, "*", precedence(expr)) + + def _print_Not(self, expr: sympy.Expr) -> str: + # pyrefly: ignore [missing-attribute] + return f"not ({self._print(expr.args[0])})" + + def _print_Add(self, expr: sympy.Expr, order: str | None = None) -> str: + return self.stringify(expr.args, " + ", precedence(expr)) + + def _print_Relational(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, f" {expr.rel_op} ", precedence(expr)) + + def _print_BitwiseFn_bitwise_and(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " & ", PRECEDENCE["BitwiseAnd"]) + + def _print_BitwiseFn_bitwise_or(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " | ", PRECEDENCE["BitwiseOr"]) + + def _print_BitwiseFn_bitwise_xor(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " ^ ", PRECEDENCE["BitwiseXor"]) + + # NB: this is OK to put here, because Mod is only defined for positive + # numbers, and so across C/Python its behavior is consistent + def _print_Mod(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " % ", PRECEDENCE["Atom"] - 0.5) + + def _print_FloatTrueDiv(self, expr: sympy.Expr) -> str: + s = self.stringify(expr.args, " / ", PRECEDENCE["Atom"] - 0.5) + return f"({s})" + + def _print_CleanDiv(self, expr: sympy.Expr) -> str: + return self._print_FloorDiv(expr) + + def _print_Identity(self, expr: sympy.Expr) -> str: + # pyrefly: ignore [missing-attribute] + return self._print(expr.args[0]) + + def _print_Float(self, expr: sympy.Expr) -> str: + if expr._prec == 53: + # IEEE-754 double precision have 53 bits. SymPy prints them with + # 15 digits, but we need 17 for round-trip correctness + return str(sympy.Float(expr, dps=17)) + else: + # We don't use other precisions in pytorch + return str(expr) + + # This must be implemented because sympy will collect x * x into Pow(x, 2), without + # any explicit intervention. We print it just like x * x, notably, we + # never generate sympy.Pow with floats. + # + # NB: this pow by natural, you should never have used builtin sympy.pow + # for FloatPow, and a symbolic exponent should be PowByNatural. These + # means exp is guaranteed to be integer. + def _print_Pow(self, expr: sympy.Expr) -> str: + base, exp = expr.args + if exp != int(exp): + raise AssertionError(exp) + exp = int(exp) + if exp < 0: + raise AssertionError(f"exponent must be non-negative, got {exp}") + if exp > 0: + return self.stringify([base] * exp, "*", PRECEDENCE["Mul"]) + return "1" + + # Explicit NotImplemented functions are to prevent default sympy printing + # behavior, which will just barf out ToFloat(...) to your IR. The error + # message is better here because it tells you which printer class it needs + # to go in. + + def _print_ToFloat(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_ToFloat not implemented for {type(self)}") + + def _print_Infinity(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_Infinity not implemented for {type(self)}") + + def _print_NegativeInfinity(self, expr: sympy.Expr) -> str: + raise NotImplementedError( + f"_print_NegativeInfinity not implemented for {type(self)}" + ) + + def _print_NaN(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_NaN not implemented for {type(self)}") + + def _print_FloorDiv(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_FloorDiv not implemented for {type(self)}") + + def _print_PythonMod(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_PythonMod not implemented for {type(self)}") + + def _print_IntTrueDiv(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_IntTrueDiv not implemented for {type(self)}") + + def _print_PowByNatural(self, expr: sympy.Expr) -> str: + raise NotImplementedError( + f"_print_PowByNatural not implemented for {type(self)}" + ) + + def _print_FloatPow(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_FloatPow not implemented for {type(self)}") + + def _print_TruncToInt(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_TruncToInt not implemented for {type(self)}") + + def _print_RoundToInt(self, expr: sympy.Expr) -> str: + raise NotImplementedError(f"_print_RoundToInt not implemented for {type(self)}") + + def _print_RoundDecimal(self, expr: sympy.Expr) -> str: + raise NotImplementedError( + f"_print_RoundDecimal not implemented for {type(self)}" + ) + + # NB: Some float operations are INTENTIONALLY not implemented for + # printers. You can implement them as a quick unblock, but it is better + # to ask yourself why we haven't done this computation in the Tensor + # universe instead + + def _print_TruncToFloat(self, expr: sympy.Expr) -> str: + raise NotImplementedError( + f"_print_TruncToFloat not implemented for {type(self)}" + ) + + +class PythonPrinter(ExprPrinter): + def _print_ToFloat(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("ToFloat expects exactly one argument") + # NB: We use sym_float here because the printer is used for cache + # serialization, and cache guards get evaluated with SymInt to + # propagate guards to the parent ShapeEnv. However, this comes at a + # runtime cost for guards involving float. If this is unacceptable + # overhead, what you want to do is have two separate printers for + # SymInt, one for when the inputs are guaranteed to be int, and + # another for when they could be SymInt. + # + # NB: sym_min/sym_max also have this problem, but I chose not to fix + # those. + # + # See https://github.com/pytorch/pytorch/issues/142507 for more + # context. + # pyrefly: ignore [missing-attribute] + return f"torch.sym_float({self._print(expr.args[0])})" + + def _print_And(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " and ", precedence(expr)) + + def _print_Or(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " or ", precedence(expr)) + + def _print_ModularIndexing(self, expr: sympy.Expr) -> str: + x, div, mod = ( + self.parenthesize(arg, PRECEDENCE["Atom"] - 0.5) for arg in expr.args + ) + if div != "1": + x = f"({x} // {div})" + return f"({x} % {mod})" + + def _print_Infinity(self, expr: sympy.Expr) -> str: + return "math.inf" + + def _print_NegativeInfinity(self, expr: sympy.Expr) -> str: + return "-math.inf" + + def _print_NaN(self, expr: sympy.Expr) -> str: + return "math.nan" + + # WARNING: this is dangerous for Triton, which has C-style modulus + def _print_PythonMod(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " % ", PRECEDENCE["Atom"] - 0.5) + + # WARNING: this is dangerous for Triton, which has C-style modulus + def _print_FloorDiv(self, expr: sympy.Expr) -> str: + x, div = (self.parenthesize(arg, PRECEDENCE["Atom"] - 0.5) for arg in expr.args) + return f"{x} // {div}" + + # WARNING: this is dangerous for Triton, when lhs, rhs > 2**53, Python + # does a special algorithm + def _print_IntTrueDiv(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " / ", PRECEDENCE["Atom"] - 0.5) + + def _helper_sqrt(self, expr: sympy.Expr) -> str: + # NB: We use torch._sym_sqrt here instead of math.sqrt because the + # guard expression may be evaluated with SymInt/SymFloat inputs (e.g. + # during cache hit re-evaluation in evaluate_guards_expression). + # math.sqrt on a SymFloat triggers evaluate_expr which forces + # concretization/specialization of the symbol, creating spurious + # guards that didn't exist in the original program. + # torch._sym_sqrt properly propagates through the symbolic system + # without forcing specialization. + # See https://github.com/pytorch/pytorch/issues/152435 + # pyrefly: ignore [missing-attribute] + return f"torch._sym_sqrt({self._print(expr)})" + + def _print_OpaqueUnaryFn_sqrt(self, expr: sympy.Expr) -> str: + return self._helper_sqrt(expr.args[0]) + + def _print_FloatPow(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " ** ", PRECEDENCE["Pow"]) + + # TODO: Not sure this works with Triton, even when base/exp are integral + def _print_PowByNatural(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " ** ", PRECEDENCE["Pow"]) + + def _print_floor(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("floor expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.floor({self._print(expr.args[0])})" + + def _print_FloorToInt(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("FloorToInt expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.floor({self._print(expr.args[0])})" + + def _print_TruncToInt(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("TruncToInt expects exactly one argument") + # This also could have been int(), they'll do the same thing for float + # pyrefly: ignore [missing-attribute] + return f"math.trunc({self._print(expr.args[0])})" + + def _print_ceiling(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("ceiling expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.ceil({self._print(expr.args[0])})" + + def _print_CeilToInt(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("CeilToInt expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.ceil({self._print(expr.args[0])})" + + def _print_Abs(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("Abs expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"abs({self._print(expr.args[0])})" + + # NB: It's expected that we've made explicit any promotion in the sympy + # expression, so it doesn't matter that Python max/min doesn't perform + # promotion + def _print_Max(self, expr: sympy.Expr) -> str: + if len(expr.args) < 2: + raise AssertionError("Max expects at least two arguments") + # pyrefly: ignore [missing-attribute] + return f"max({', '.join(map(self._print, expr.args))})" + + def _print_Min(self, expr: sympy.Expr) -> str: + if len(expr.args) < 2: + raise AssertionError("Min expects at least two arguments") + # pyrefly: ignore [missing-attribute] + return f"min({', '.join(map(self._print, expr.args))})" + + def _print_OpaqueUnaryFn_cos(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("cos expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.cos({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_cosh(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("cosh expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.cosh({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_acos(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("acos expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.acos({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_sin(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("sin expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.sin({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_sinh(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("sinh expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.sinh({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_asin(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("asin expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.asin({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_tan(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("tan expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.tan({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_tanh(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("tanh expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.tanh({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_atan(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("atan expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.atan({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_log2(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("log2 expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.log2({self._print(expr.args[0])})" + + def _print_RoundToInt(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("RoundToInt expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"round({self._print(expr.args[0])})" + + def _print_RoundDecimal(self, expr: sympy.Expr) -> str: + if len(expr.args) != 2: + raise AssertionError("RoundDecimal expects exactly two arguments") + number, ndigits = expr.args + if not isinstance(ndigits, sympy.Integer): + raise TypeError("ndigits must be an instance of sympy.Integer") + # pyrefly: ignore [missing-attribute] + return f"round({self._print(number)}, {ndigits})" + + def _print_Piecewise(self, expr: sympy.Expr) -> str: + # Convert Piecewise(expr_cond_pairs) to nested ternary expressions + # Piecewise((e1, c1), (e2, c2), ..., (eN, cN)) + # becomes: e1 if c1 else (e2 if c2 else (... else eN)) + result: str | None = None + for expr_i, cond_i in reversed(expr.args): + # pyrefly: ignore [missing-attribute] + expr_str = self._print(expr_i) + if cond_i == True: # noqa: E712 + # This is the default case + result = expr_str + else: + # pyrefly: ignore [missing-attribute] + cond_str = self._print(cond_i) + if result is None: + result = expr_str + else: + result = f"({expr_str} if {cond_str} else {result})" + return result if result else "0" + + +class CppPrinter(ExprPrinter): + def _print_Integer(self, expr: sympy.Expr) -> str: + suffix = "LL" if sys.platform in ["darwin", "win32"] else "L" + i = int(expr) + if i > INDEX_TYPE_MAX or i < INDEX_TYPE_MIN: + raise OverflowError(f"{i} too big to convert to {INDEX_TYPE}") + elif i == INDEX_TYPE_MIN: + if i != (-1) << 63: + raise AssertionError("unexpected minimum index type value") + # Writing -9223372036854775808L makes the value overflow + # as it is parsed as -(9223372036854775808L) by the C/C++ compiler + return f"(-1{suffix} << 63)" + return f"{i}{suffix}" + + def _print_Where(self, expr: sympy.Expr) -> str: + c, p, q = ( + self.parenthesize(arg, PRECEDENCE["Atom"] - 0.5) for arg in expr.args + ) + return f"{c} ? {p} : {q}" + + def _print_Or(self, expr: sympy.Expr) -> str: + return self.stringify(expr.args, " || ", precedence(expr)) + + def _print_Piecewise(self, expr: sympy.Expr) -> str: + # Convert Piecewise(expr_cond_pairs) to nested ternary operators + # Piecewise((e1, c1), (e2, c2), ..., (eN, cN)) + # becomes: c1 ? e1 : (c2 ? e2 : (... : eN)) + result: str | None = None + for expr_i, cond_i in reversed(expr.args): + expr_str = self.parenthesize(expr_i, PRECEDENCE["Atom"] - 0.5) + if cond_i == True: # noqa: E712 + # This is the default case + result = expr_str + else: + cond_str = self.parenthesize(cond_i, PRECEDENCE["Atom"] - 0.5) + if result is None: + result = expr_str + else: + result = f"{cond_str} ? {expr_str} : {result}" + return f"({result})" if result else "0" + + def _print_ModularIndexing(self, expr: sympy.Expr) -> str: + x, div, mod = expr.args + x = self.doprint(x) + if div != 1: + div = self.doprint(div) + if expr.is_integer: + x = f"c10::div_floor_integer(static_cast({x}), static_cast({div}))" + else: + x = f"c10::div_floor_floating(static_cast({x}), static_cast({div}))" + mod = self.doprint(mod) + return f"(static_cast<{INDEX_TYPE}>({x}) % static_cast<{INDEX_TYPE}>({mod}))" + + def _print_FloorDiv(self, expr: sympy.Expr) -> str: + x, div = expr.args + x = self.doprint(x) + div = self.doprint(div) + if expr.is_integer: + return f"c10::div_floor_integer(static_cast({x}), static_cast({div}))" + return f"c10::div_floor_floating(static_cast({x}), static_cast({div}))" + + def _print_floor(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("floor expects exactly one argument") + # pyrefly: ignore [missing-attribute] + r = f"std::floor({self._print(expr.args[0])})" + return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r + + def _print_FloorToInt(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("FloorToInt expects exactly one argument") + # pyrefly: ignore [missing-attribute] + r = f"std::floor({self._print(expr.args[0])})" + return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r + + def _print_TruncToInt(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("TruncToInt expects exactly one argument") + # pyrefly: ignore [missing-attribute] + r = f"std::trunc({self._print(expr.args[0])})" + return f"static_cast<{INDEX_TYPE}>({r})" + + def _print_TruncToFloat(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("TruncToFloat expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::trunc({self._print(expr.args[0])})" + + def _print_ToFloat(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("ToFloat expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"static_cast({self._print(expr.args[0])})" + + def _print_PythonMod(self, expr: sympy.Expr) -> str: + x, div = expr.args + x = self.doprint(x) + div = self.doprint(div) + return f"c10::div_mod({x}, {div})" + + def _print_IntTrueDiv(self, expr: sympy.Expr) -> str: + lhs, rhs = expr.args + # TODO: This is only accurate up to 2**53 + # pyrefly: ignore [missing-attribute] + return f"static_cast({self._print(lhs)}) / static_cast({self._print(rhs)})" + + # TODO: PowByNatural: we need to implement our own int-int pow. Do NOT + # use std::pow, that operates on floats + def _print_PowByNatural(self, expr: sympy.Expr) -> str: + # Implement the special-case of 2**x for now + base, exp = expr.args + if base == 2: + # pyrefly: ignore [missing-attribute] + return f"(1 << ({self._print(exp)}))" + raise NotImplementedError( + f"_print_PowByNatural not implemented for {type(self)}" + ) + + def _print_FloatPow(self, expr: sympy.Expr) -> str: + base, exp = expr.args + # pyrefly: ignore [missing-attribute] + return f"std::pow({self._print(base)}, {self._print(exp)})" + + def _print_Pow(self, expr: sympy.Expr) -> str: + # Uses float constants to perform FP div + base, exp = expr.args + + if exp == 0.5 or exp == -0.5: + # pyrefly: ignore [missing-attribute] + base = self._print(base) + return f"std::sqrt({base})" if exp == 0.5 else f"1.0/std::sqrt({base})" + if exp.is_integer: + exp = int(exp) + if exp > 0: + r = self.stringify([base] * exp, "*", PRECEDENCE["Mul"]) + elif exp < -1: + r = ( + "1.0/(" + + self.stringify([base] * abs(exp), "*", PRECEDENCE["Mul"]) + + ")" + ) + elif exp == -1: + # pyrefly: ignore [missing-attribute] + r = "1.0/" + self._print(base) + else: # exp == 0 + r = "1.0" + + return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r + else: + # TODO: float vs double + return f"std::pow({base}, {float(exp)})" + + def _print_Rational(self, expr: sympy.Expr) -> str: + # Uses float constants to perform FP div + if expr.q == 1: + r = f"{expr.p}" + else: + r = f"{expr.p}.0/{expr.q}.0" + return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r + + def _print_ceiling(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("ceiling expects exactly one argument") + # pyrefly: ignore [missing-attribute] + r = f"std::ceil({self._print(expr.args[0])})" + return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r + + def _print_CeilToInt(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("CeilToInt expects exactly one argument") + # pyrefly: ignore [missing-attribute] + r = f"std::ceil({self._print(expr.args[0])})" + return f"static_cast<{INDEX_TYPE}>({r})" if expr.is_integer else r + + def _print_Min(self, expr: sympy.Expr) -> str: + # pyrefly: ignore [missing-attribute] + args = [self._print(a) for a in expr.args] + if len(args) == 2: + return f"std::min(static_cast<{INDEX_TYPE}>({args[0]}), static_cast<{INDEX_TYPE}>({args[1]}))" + else: + # Initializer list overload + il = "{" + ", ".join(args) + "}" + return f"std::min<{INDEX_TYPE}>({il})" + + def _print_Max(self, expr: sympy.Expr) -> str: + # pyrefly: ignore [missing-attribute] + args = [self._print(a) for a in expr.args] + if len(args) == 2: + return f"std::max(static_cast<{INDEX_TYPE}>({args[0]}), static_cast<{INDEX_TYPE}>({args[1]}))" + else: + # Initializer list overload + il = "{" + ", ".join(args) + "}" + return f"std::max<{INDEX_TYPE}>({il})" + + def _print_Abs(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("Abs expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::abs({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_cos(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("cos expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::cos({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_cosh(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("cosh expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::cosh({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_acos(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("acos expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::acos({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_sin(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("sin expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"math.sin({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_sinh(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("sinh expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::sinh({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_asin(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("asin expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::asin({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_tan(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("tan expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::tan({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_tanh(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("tanh expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::tanh({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_atan(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("atan expects exactly one argument") + # pyrefly: ignore [missing-attribute] + return f"std::atan({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_sqrt(self, expr: sympy.Expr) -> str: + # pyrefly: ignore [missing-attribute] + return f"std::sqrt({self._print(expr.args[0])})" + + def _print_OpaqueUnaryFn_log2(self, expr: sympy.Expr) -> str: + # pyrefly: ignore [missing-attribute] + return f"std::log2({self._print(expr.args[0])})" + + def _print_RoundToInt(self, expr: sympy.Expr) -> str: + if len(expr.args) != 1: + raise AssertionError("RoundToInt expects exactly one argument") + # TODO: dispatch to llrint depending on index type + # pyrefly: ignore [missing-attribute] + return f"std::lrint({self._print(expr.args[0])})" + + def _print_RoundDecimal(self, expr: sympy.Expr) -> str: + if len(expr.args) != 2: + raise AssertionError("RoundDecimal expects exactly two arguments") + number, ndigits = expr.args + if number.is_integer: + # ndigits < 0 should have been filtered by the sympy function + if ndigits >= 0: + raise AssertionError("ndigits must be negative for integer inputs") + raise ValueError( + f"For integer inputs, only non-negative ndigits are currently supported, but got {ndigits}." + ) + number_str = self.parenthesize(number, PRECEDENCE["Mul"]) + return f"static_cast(std::nearbyint(1e{ndigits} * {number_str}) * 1e{-ndigits})" + + def _print_BooleanTrue(self, expr: sympy.Expr) -> str: + return "true" + + def _print_BooleanFalse(self, expr: sympy.Expr) -> str: + return "false" + + def _print_Infinity(self, expr: sympy.Expr) -> str: + return "std::numeric_limits::infinity()" + + def _print_NegativeInfinity(self, expr: sympy.Expr) -> str: + return f"-{self._print_Infinity(expr)}" + + def _print_NaN(self, expr: sympy.Expr) -> str: + return "std::numeric_limits::quiet_NaN()" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/reference.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/reference.py new file mode 100644 index 0000000000000000000000000000000000000000..3879cef698154e400d38a26167baba496bd978ee --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/reference.py @@ -0,0 +1,615 @@ +# mypy: allow-untyped-defs +import math +import operator +from typing import NoReturn + +import sympy + +import torch +from torch.utils._sympy.functions import ( + _keep_float, + BitwiseFn_bitwise_and, + BitwiseFn_bitwise_or, + BitwiseFn_bitwise_xor, + FloatPow, + FloatTrueDiv, + FloorDiv, + IntTrueDiv, + Max, + Min, + Mod, + OpaqueUnaryFn_exp, + OpaqueUnaryFn_log, + OpaqueUnaryFn_log2, + OpaqueUnaryFn_sqrt, + PowByNatural, + RoundDecimal, + RoundToInt, + ToFloat, + TruncToInt, +) + + +# The sympy interpretation of operators. It will also sometimes work with +# plain int/float, but if you do certain operations you will get out a +# sympy.Basic in the end. If you want the Python/FX traceable interpretation, +# check PythonReferenceAnalysis. +# NB: For magic methods this needs to use normal magic methods +# so that test_magic_methods works +class ReferenceAnalysis: + @staticmethod + def constant(c, dtype): + return sympy.sympify(c) + + @staticmethod + def or_(a, b): + return a | b + + @staticmethod + def and_(a, b): + return a & b + + @staticmethod + def eq(a, b): + if isinstance(a, sympy.Expr) or isinstance(b, sympy.Expr): + return sympy.Eq(a, b) + return a == b + + @classmethod + def ne(cls, a, b): + return cls.not_(cls.eq(a, b)) + + @staticmethod + def lt(a, b): + return a < b + + @staticmethod + def gt(a, b): + return a > b + + @staticmethod + def le(a, b): + return a <= b + + @staticmethod + def ge(a, b): + return a >= b + + @staticmethod + def not_(a): + if isinstance(a, bool): + raise AssertionError("not_ needs sympy expr") + return ~a + + @staticmethod + def reciprocal(x): + return FloatTrueDiv(1.0, x) + + @staticmethod + def square(x): + return PowByNatural(x, 2) + + @staticmethod + def trunc_to_int(x, dtype): + return TruncToInt(x) + + @staticmethod + def ceil_to_int(x, dtype): + return sympy.ceiling(x) + + @staticmethod + def floor_to_int(x, dtype): + return sympy.floor(x) + + @staticmethod + def floor(x): + return _keep_float(sympy.floor)(x) + + @staticmethod + def ceil(x): + return _keep_float(sympy.ceiling)(x) + + @staticmethod + def to_dtype(x, dtype): + if dtype == torch.float64: + return ToFloat(x) + raise NotImplementedError(f"to_dtype {dtype} NYI") + + @staticmethod + def mod(x, y): + return Mod(x, y) + + @staticmethod + def abs(x): + return abs(x) + + @staticmethod + def neg(x): + return -x + + @staticmethod + def truediv(a, b): + return FloatTrueDiv(a, b) + + @staticmethod + def int_truediv(a, b): + return IntTrueDiv(a, b) + + @staticmethod + def floordiv(a, b): + return FloorDiv(a, b) + + @staticmethod + def truncdiv(a, b) -> NoReturn: + raise NotImplementedError("TODO: truncdiv") + + @staticmethod + def add(a, b): + return _keep_float(operator.add)(a, b) + + @classmethod + def sym_sum(cls, args): + return sympy.Add(*args) + + @staticmethod + def mul(a, b): + return _keep_float(operator.mul)(a, b) + + @staticmethod + def sub(a, b): + return _keep_float(operator.sub)(a, b) + + @staticmethod + def exp(x): + return OpaqueUnaryFn_exp(x) + + @staticmethod + def log(x): + return OpaqueUnaryFn_log(x) + + @staticmethod + def log2(x): + return OpaqueUnaryFn_log2(x) + + @staticmethod + def sqrt(x): + return OpaqueUnaryFn_sqrt(x) + + @staticmethod + def pow(a, b): + # pyrefly: ignore [bad-argument-count, bad-argument-type] + return _keep_float(FloatPow)(a, b) + + @staticmethod + def pow_by_natural(a, b): + return PowByNatural(a, b) + + @staticmethod + def minimum(a, b): + return Min(a, b) + + @staticmethod + def maximum(a, b): + return Max(a, b) + + @staticmethod + def round_to_int(a, dtype): + return RoundToInt(a) + + @staticmethod + def round_decimal(a, b): + return RoundDecimal(a, b) + + @staticmethod + def bitwise_and(a, b): + return BitwiseFn_bitwise_and(a, b) + + @staticmethod + def bitwise_or(a, b): + return BitwiseFn_bitwise_or(a, b) + + @staticmethod + def bitwise_xor(a, b): + return BitwiseFn_bitwise_xor(a, b) + + +# Unlike ReferenceAnalysis, does NOT sympyify, instead, works with plain +# Python types and is FX traceable. Inheritance here is purely for code +# sharing (TODO: considering splitting out a BaseReferenceAnalysis). +class PythonReferenceAnalysis(ReferenceAnalysis): + @staticmethod + def constant(c, dtype): + if dtype is torch.int64: + return int(c) + elif dtype is torch.double: + return float(c) + elif dtype is torch.bool: + return bool(c) + else: + raise AssertionError(f"unrecognized dtype {dtype}") + + @staticmethod + def not_(a): + return torch.sym_not(a) + + @classmethod + def sym_sum(cls, args): + if len(args) == 0: + return 0 + if len(args) == 1: + return args[0] + acc = cls.add(args[0], args[1]) + for i in range(2, len(args)): + acc = cls.add(acc, args[i]) + return acc + + @staticmethod + def floordiv(a, b): + return a // b + + @staticmethod + def mod(x, y): + return x % y + + @staticmethod + def python_mod(x, y): + return x % y + + @staticmethod + def truncdiv(a, b): + return a / b + + @staticmethod + def to_dtype(x, dtype): + if dtype == torch.float64: + return torch.sym_float(x) + raise NotImplementedError(f"to_dtype {dtype} NYI") + + @staticmethod + def exp(x) -> NoReturn: + raise AssertionError("exp is not valid shape sympy expr") + + @staticmethod + def log(x) -> NoReturn: + raise AssertionError("log is not valid shape sympy expr") + + @staticmethod + def log2(x): + return torch._sym_log2(x) # type: ignore[attr-defined] + + @staticmethod + def sqrt(x): + return torch._sym_sqrt(x) # type: ignore[attr-defined] + + @staticmethod + def minimum(a, b): + return torch.sym_min(a, b) + + @staticmethod + def maximum(a, b): + return torch.sym_max(a, b) + + @staticmethod + def floor_to_int(x, dtype): + return math.floor(x) + + @staticmethod + def ceil_to_int(x, dtype): + return math.ceil(x) + + @staticmethod + def floor(x): + return float(math.floor(x)) + + @staticmethod + def ceil(x): + return float(math.ceil(x)) + + @staticmethod + def truediv(a, b): + return a / b + + @staticmethod + def pow(a, b): + return a**b + + @staticmethod + def pow_by_natural(a, b): + # Pray that safe_pow is not needed here lol. In particular, this + # never participates in VR low/high ranges, so overflow should be + # unlikely + return a**b + + @staticmethod + def round_to_int(a, dtype): + return round(a) + + @staticmethod + def round_decimal(a, b): + return round(a, ndigits=b) + + @staticmethod + def bitwise_and(a, b): + return a & b + + @staticmethod + def bitwise_or(a, b): + return a | b + + @staticmethod + def bitwise_xor(a, b): + return a ^ b + + @staticmethod + def expr_cond_pair(expr, cond): + return (expr, cond) + + @staticmethod + def piecewise(*pairs): + # Build nested sym_ite from right to left. + # Piecewise((e1, c1), (e2, c2), ..., (en, True)) becomes + # sym_ite(c1, e1, sym_ite(c2, e2, ... en)) + result = pairs[-1][0] + for expr, cond in reversed(pairs[:-1]): + result = torch.sym_ite(cond, expr, result) + return result + + +# Like PythonReferenceAnalysis, but some export-unfriendly choices of +# operators to make things faster +class OptimizedPythonReferenceAnalysis(PythonReferenceAnalysis): + @staticmethod + def sym_sum(args): + return torch.sym_sum(args) + + +def _to_dtype(x: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + return torch.ops.prims.convert_element_type.default(x, dtype) + + +# Suppose we have some int/float arguments. This diagram commutes: +# +# int/float -- PythonReferenceAnalysis.op --> int/float +# | | +# | | +# torch.tensor(..., dtype=torch.int64/torch.float64) +# | | +# V V +# Tensor -- TensorReferenceAnalysis.op --> Tensor +# +# NB: int before and after must be representable in int64 (we will +# insert guards accordingly.) +# +# This is guaranteed to be FX traceable with OpOverloads only. +class TensorReferenceAnalysis: + # NB: This is actually dead, because with Proxy tracing the factory + # function isn't traced correctly. Here for completeness. + @staticmethod + def constant(c, dtype): + d: int | float | bool + if dtype is torch.int64: + d = int(c) + elif dtype is torch.double: + d = float(c) + elif dtype is torch.bool: + d = bool(c) + else: + raise AssertionError(f"unrecognized dtype {dtype}") + return torch.ops.aten.scalar_tensor.default(d, dtype=dtype) + + @staticmethod + def or_(a, b): + return torch.ops.aten.logical_or.default(a, b) + + @staticmethod + def and_(a, b): + return torch.ops.aten.logical_and.default(a, b) + + @staticmethod + def bitwise_and(a, b): + return torch.ops.aten.bitwise_and(a, b) + + @staticmethod + def bitwise_or(a, b): + return torch.ops.aten.bitwise_or(a, b) + + @staticmethod + def bitwise_xor(a, b): + return torch.ops.aten.bitwise_xor(a, b) + + @staticmethod + def eq(a, b): + return torch.ops.aten.eq.Tensor(a, b) + + @classmethod + def ne(cls, a, b): + return torch.ops.aten.ne.Tensor(a, b) + + @staticmethod + def lt(a, b): + return torch.ops.aten.lt.Tensor(a, b) + + @staticmethod + def gt(a, b): + return torch.ops.aten.gt.Tensor(a, b) + + @staticmethod + def le(a, b): + return torch.ops.aten.le.Tensor(a, b) + + @staticmethod + def ge(a, b): + return torch.ops.aten.ge.Tensor(a, b) + + @staticmethod + def not_(a): + return torch.ops.aten.logical_not.default(a) + + @staticmethod + def reciprocal(x): + return torch.ops.aten.reciprocal.default(x) + + @staticmethod + def square(x): + # TODO: maybe composite implicit autograd doesn't work here? + return torch.ops.aten.square.default(x) + + @staticmethod + def trunc_to_int(x, dtype): + return _to_dtype(torch.ops.aten.trunc.default(x), dtype) + + @staticmethod + def ceil_to_int(x, dtype): + return _to_dtype(torch.ops.aten.ceil.default(x), dtype) + + @staticmethod + def floor_to_int(x, dtype): + return _to_dtype(torch.ops.aten.floor.default(x), dtype) + + @staticmethod + def floor(x): + return torch.ops.aten.floor.default(x) + + @staticmethod + def ceil(x): + return torch.ops.aten.ceil.default(x) + + @staticmethod + def to_dtype(x, dtype): + return _to_dtype(x, dtype) + + @staticmethod + def mod(x, y) -> NoReturn: + # TODO: https://github.com/pytorch/pytorch/pull/133654 + raise NotImplementedError( + "no C-style modulus operation available from frontend atm" + ) + + @staticmethod + def abs(x): + return torch.ops.aten.abs.default(x) + + @staticmethod + def neg(x): + return torch.ops.aten.neg.default(x) + + @staticmethod + def truediv(a, b): + return torch.ops.aten.true_divide.Tensor(a, b) + + @staticmethod + def int_truediv(a, b): + raise NotImplementedError( + "Python int truediv difficult to implement in PyTorch atm" + ) + + # TODO: This is wrong, CPython has a custom implementation of true + # division that results in higher precision when the floats are + # sufficiently large. Short term fix: add a guard here + # pyrefly: ignore [unreachable] + return torch.ops.aten.true_divide.default( + _to_dtype(a, torch.float64), _to_dtype(b, torch.float64) + ) + + @staticmethod + def floordiv(a, b): + return torch.ops.aten.div.Tensor_mode(a, b, rounding_mode="floor") + + @staticmethod + def truncdiv(a, b) -> NoReturn: + raise NotImplementedError( + "no C-style truncdiv operation available from frontend atm" + ) + + @staticmethod + def add(a, b): + return torch.ops.aten.add.Tensor(a, b) + + @staticmethod + def mul(a, b): + return torch.ops.aten.mul.Tensor(a, b) + + @staticmethod + def sub(a, b): + return torch.ops.aten.sub.Tensor(a, b) + + @staticmethod + def exp(x): + return torch.ops.aten.exp.default(x) + + @staticmethod + def log(x): + return torch.ops.aten.log.default(x) + + @staticmethod + def log2(x): + return torch.ops.aten.log2.default(x) + + @staticmethod + def sqrt(x): + return torch.ops.aten.sqrt.default(x) + + @staticmethod + def sin(x): + return torch.ops.aten.sin.default(x) + + @staticmethod + def cos(x): + return torch.ops.aten.cos.default(x) + + @staticmethod + def tanh(x): + return torch.ops.aten.tanh.default(x) + + @staticmethod + def sinh(x): + return torch.ops.aten.sinh.default(x) + + @staticmethod + def cosh(x): + return torch.ops.aten.cosh.default(x) + + @staticmethod + def tan(x): + return torch.ops.aten.tan.default(x) + + @staticmethod + def acos(x): + return torch.ops.aten.acos.default(x) + + @staticmethod + def atan(x): + return torch.ops.aten.atan.default(x) + + @staticmethod + def asin(x): + return torch.ops.aten.asin.default(x) + + @staticmethod + def pow(a, b): + return torch.ops.aten.pow.Tensor_Tensor(a, b) + + @staticmethod + def pow_by_natural(a, b): + # NB: pow handles int x int fine + return torch.ops.aten.pow.Tensor_Tensor(a, b) + + @staticmethod + def minimum(a, b): + return torch.ops.aten.minimum.default(a, b) + + @staticmethod + def maximum(a, b): + return torch.ops.aten.maximum.default(a, b) + + @staticmethod + def round_to_int(a, dtype): + return torch.ops.aten.round.default(a) + + @staticmethod + def round_decimal(a, b) -> NoReturn: + raise NotImplementedError( + "round decimal doesn't support Tensor second argument atm" + ) + + # return torch.ops.aten.round.decimals(a, b) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/singleton_int.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/singleton_int.py new file mode 100644 index 0000000000000000000000000000000000000000..57d5615e552711a490e306ebc97a260a55251c21 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/singleton_int.py @@ -0,0 +1,96 @@ +# mypy: allow-untyped-defs +import sympy +from sympy.multipledispatch import dispatch + + +__all__ = ["SingletonInt"] + + +class SingletonInt(sympy.AtomicExpr): + # This is probably not super important unless we are in multiple dispatch + # situations with other more exotic Expr types. + _op_priority = 99999 + + def __new__(cls, *args, coeff=None, **kwargs): + instance = super().__new__(cls, *args, **kwargs) + return instance + + # The semantics of this class should match that of NestedIntSymNodeImpl in + # c10/core/NestedIntSymNodeImpl.h + def __init__(self, val, *, coeff=1) -> None: + self._val = val + self._coeff = coeff + super().__init__() + + # See NOTE [ Inequalities with nested int ] + def _eval_Eq(self, other): + if ( + isinstance(other, SingletonInt) + and other._val == self._val + and self._coeff == other._coeff + ): + return sympy.true + else: + return sympy.false + + # This is necessary so that calling expr.free_symbols on exprs that contain + # this Singleton does not error + @property + def free_symbols(self): + return set() + + def __mul__(self, other): + if isinstance(other, SingletonInt): + raise ValueError( + "SingletonInt cannot be multiplied by another SingletonInt" + ) + return SingletonInt(self._val, coeff=self._coeff * other) + + def __rmul__(self, other): + if isinstance(other, SingletonInt): + raise ValueError( + "SingletonInt cannot be multiplied by another SingletonInt" + ) + return SingletonInt(self._val, coeff=self._coeff * other) + + # Make sure we promptly raise an error instead of falling back to building + # an expression tree. There are probably more ops, how can we be exhaustive? + def __add__(self, other): + raise NotImplementedError("NYI") + + def __sub__(self, other): + raise NotImplementedError("NYI") + + def __truediv__(self, other): + raise NotImplementedError("NYI") + + def __floordiv__(self, other): + raise NotImplementedError("NYI") + + def __mod__(self, other): + raise NotImplementedError("NYI") + + +# See NOTE [ Inequalities with nested int ] +@dispatch(sympy.Integer, SingletonInt) +def _eval_is_ge(a, b): + if a < 2: + return sympy.false + raise ValueError("Symbolic SingletonInt: Relation is indeterminate") + + +@dispatch(SingletonInt, sympy.Integer) # type: ignore[no-redef] +def _eval_is_ge(a, b): # noqa: F811 + if b <= 2: + return sympy.true + raise ValueError("Symbolic SingletonInt: Relation is indeterminate") + + +@dispatch(SingletonInt, SingletonInt) # type: ignore[no-redef] +def _eval_is_ge(a, b): # noqa: F811 + if a._val == b._val: + if a._coeff >= b._coeff: + return sympy.true + else: + return sympy.false + raise ValueError("Symbolic SingletonInt: Relation is indeterminate") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/solve.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/solve.py new file mode 100644 index 0000000000000000000000000000000000000000..3bd5e1484601ffa1c7c2743ffa228c536cd54fb5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/solve.py @@ -0,0 +1,179 @@ +import logging + +import sympy + +from torch.utils._sympy.functions import FloorDiv + + +log = logging.getLogger(__name__) + +_MIRROR_REL_OP: dict[type[sympy.Basic], type[sympy.Rel]] = { + sympy.Eq: sympy.Eq, + sympy.Ne: sympy.Ne, + sympy.Ge: sympy.Le, + sympy.Gt: sympy.Lt, + sympy.Le: sympy.Ge, + sympy.Lt: sympy.Gt, +} + +INEQUALITY_TYPES = (sympy.Gt, sympy.Ge, sympy.Lt, sympy.Le) + + +def mirror_rel_op(type: type) -> type[sympy.Rel] | None: + return _MIRROR_REL_OP.get(type) + + +# Tries to simplify 'expr', so as to leave only 'thing' in the left-hand side. +# +# Returns a tuple of: +# 1. The simplified expression +# 2. The expression on the right-hand side +# +# Returns 'None' if it can't reach a state where the only thing in the left +# hand side is 'thing'. +# +# 'trials': number of times 'try_solve' will try to isolate 'thing' to the +# left-hand side. +# +# 'floordiv_inequality': flag to enable conversion of 'FloorDiv' into +# inequalities. +def try_solve( + expr: sympy.Basic, + thing: sympy.Basic, + trials: int = 5, + floordiv_inequality: bool = True, +) -> tuple[sympy.Rel, sympy.Expr] | None: + mirror = mirror_rel_op(type(expr)) + + # Ignore unsupported expressions: + # - Those that are not relational operations + # - Those that don't have a mirror (just avoiding unexpected classes) + if not isinstance(expr, sympy.Rel) or mirror is None: + log.debug("expression with unsupported type: %s", type(expr)) + return None + + lhs_has_thing = expr.lhs.has(thing) + rhs_has_thing = expr.rhs.has(thing) + + # Give up when 'thing' appears on both sides of the relational expression. + # That is because, as is, we assume the thing we are trying to isolate is + # only on the right-hand side. + if lhs_has_thing and rhs_has_thing: + log.debug("thing (%s) found in both sides of expression: %s", thing, expr) + return None + + # Try considering both LHS and RHS by mirroring the original expression: + # a < b ==> b > a + expressions = [] + + # Add each version of 'expr' if 'thing' is in its left-hand side. + if lhs_has_thing: + expressions.append(expr) + if rhs_has_thing: + expressions.append(mirror(expr.rhs, expr.lhs)) + + for e in expressions: + if e is None: + continue + + if not isinstance(e, sympy.Rel): + raise AssertionError("expected sympy.Rel") + + for _ in range(trials): + trial = _try_isolate_lhs(e, thing, floordiv_inequality=floordiv_inequality) + # Stop if there was no change in this trial. + if trial == e: + break + e = trial # type: ignore[assignment] + + # Return if we were able to isolate 'thing' on the left-hand side. + if isinstance(e, sympy.Rel) and e.lhs == thing: + log.debug("solved: %s ---> %s", expr, e) + return e, e.rhs + + return None + + +def _try_isolate_lhs( + e: sympy.Basic, thing: sympy.Basic, floordiv_inequality: bool +) -> sympy.Basic: + op = type(e) + + if isinstance(e, sympy.Rel): + # Move any constants in the left-hand side to the right-hand side. + lhs_not_thing = ( + sum(a for a in e.lhs.args if not a.has(thing)) + if isinstance(e.lhs, sympy.Add) + else 0 + ) + e = op(e.lhs - lhs_not_thing, e.rhs - lhs_not_thing) # type: ignore[attr-defined] + + # Divide both sides by the factors that don't contain thing. + if isinstance(e, sympy.Rel) and isinstance(e.lhs, sympy.Mul): + lhs, rhs = e.args + other = sympy.Mul(*[a for a in lhs.args if not a.has(thing)]) + + # If we can't tell whether 'other' is negative or positive, we do nothing. + # That is because we don't know whether we have mirror the operation or not. + # We also divide only when we know 'rhs' is not zero. + if not (isinstance(e, INEQUALITY_TYPES) and other.is_negative is None) and not ( + not isinstance(e, INEQUALITY_TYPES) and rhs.is_zero + ): + # Divide both sides by 'other'. + lhs = lhs / other + rhs = rhs / other + + # If 'e' is an inequality and 'other' is negative, we have to + # mirror the expression. + if isinstance(e, INEQUALITY_TYPES) and other.is_negative: + op = mirror_rel_op(op) # type: ignore[assignment] + + if op is None: + raise AssertionError("expected op to be not None") + e = op(lhs, rhs) + + ################################################################################ + # left-hand side is FloorDiv + ################################################################################ + # + # Given the expression: a // b op c + # where 'op' is a relational operation, these rules only work if: + # - b > 0 + # - c is an integer + if ( + floordiv_inequality + and isinstance(e, sympy.Rel) + and isinstance(e.lhs, FloorDiv) + and e.lhs.divisor.is_positive + and e.rhs.is_integer + ): + # a // b == expr + # => a >= (b * expr) and a < (b * (expr + 1)) + if isinstance(e, sympy.Eq): + numerator, denominator = e.lhs.args + return sympy.And( + sympy.Ge(numerator, (e.rhs * denominator)), + sympy.Lt(numerator, ((e.rhs + 1) * denominator)), + ) + # a // b != expr + # => a < (b * expr) or a >= (b * (expr + 1)) + if isinstance(e, sympy.Ne): + numerator, denominator = e.lhs.args + return sympy.Or( + sympy.Lt(numerator, (e.rhs * denominator)), + sympy.Ge(numerator, ((e.rhs + 1) * denominator)), + ) + # The transformations below only work if b is positive. + # Note: we only have this information for constants. + # a // b > expr => a >= b * (expr + 1) + # a // b >= expr => a >= b * expr + if isinstance(e, (sympy.Gt, sympy.Ge)): + quotient = e.rhs if isinstance(e, sympy.Ge) else (e.rhs + 1) + return sympy.Ge(e.lhs.args[0], (quotient * e.lhs.args[1])) + # a // b < expr => a < b * expr + # a // b <= expr => a < b * (expr + 1) + if isinstance(e, (sympy.Lt, sympy.Le)): + quotient = e.rhs if isinstance(e, sympy.Lt) else (e.rhs + 1) + return sympy.Lt(e.lhs.args[0], (quotient * e.lhs.args[1])) + + return e diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/symbol.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/symbol.py new file mode 100644 index 0000000000000000000000000000000000000000..61a7c147458e03e2eaf1704a42335e357eb69be7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/symbol.py @@ -0,0 +1,101 @@ +# mypy: allow-untyped-defs +""" +This file contains canonical definitions for our symbol naming conventions, +across torch.fx.experimental.symbolic_shapes and torch._inductor. The +intention is: + +1. To make it easily greppable where all the sites we use a prefix are +2. Make it possible to easily tell if we can introduce a new prefix without + introducing a conflict + +You can occasionally test if prefixes have been hardcoded by renaming prefixes +in this file and seeing what breaks. +""" + +from collections.abc import Iterable +from enum import auto, Enum + +import sympy + + +class SymT(Enum): + SIZE = auto() + FLOAT = auto() + UNBACKED_INT = auto() + UNBACKED_FLOAT = auto() + # Inductor: The intermediates in inner_fn tmp0, one generated per ops call. + # If one of these shows up in an indexing expression, that means an + # indirect load is happening. + TMP = auto() + # Inductor: Placeholder variable that is later replaced with TMP + INDIRECT = auto() + # Inductor: Some size expressions are replaced with a precomputed size ps0 + # which is computed host side, and then directly reused in the kernel, so + # we don't repeatedly recompute it on device. + PRECOMPUTED_SIZE = auto() + # Inductor: An indexing variable i0 in loops IR which ranges over non-reduced + # dim in the loop + INDEX = auto() + # Inductor: A reduction indexing (r0, r1) variables in loops IR which ranges over + # reduced dim(s) in the loop + R0_INDEX = auto() + R1_INDEX = auto() + # Inductor: In templated kernels torch._inductor.kernel, we have a hook to + # store the final output and append epilogue fusions. To do this, we must + # know what the indexes the outputs range over. NB: These will also + # advertise as INDEX, this is... probably OK? + TEMPLATE_INDEX = auto() + # Inductor: iteration domain for blockIdx.x/blockIdx.y + XBLOCK = auto() + YBLOCK = auto() + ZBLOCK = auto() + # Inductor: this is used solely for dynamic_reshape_indexer + VIEW = auto() + # Alternate (non-modular) indexing used in halide kernels + HALIDE = auto() + + +# Invariant: there must not be a prefix which is a prefix of another string, +# as this introduces ambiguity +prefix_str = { + SymT.SIZE: "s", # integer + SymT.UNBACKED_INT: "u", # integer + # Prefix z here is chosen to avoid false aliasing in symbol_is_type test + # DO NOT add a "z" type. You also need to avoid conflicts on these + # prefixes but this is somewhat easier to manage + SymT.FLOAT: "zf", + SymT.UNBACKED_FLOAT: "zuf", + SymT.TMP: "tmp", + SymT.PRECOMPUTED_SIZE: "ps", + SymT.INDEX: "i", + SymT.R0_INDEX: "r0_", + SymT.R1_INDEX: "r1_", + SymT.TEMPLATE_INDEX: "idx", + SymT.XBLOCK: "x", + SymT.YBLOCK: "y", + SymT.ZBLOCK: "z", + SymT.INDIRECT: "indirect", # false aliasing? + SymT.VIEW: "view", + SymT.HALIDE: "h", +} + + +def make_symbol(prefix: SymT, idx: int, **kwargs) -> sympy.Symbol: + # TODO: maybe put the assumptions here directly + return sympy.Symbol(f"{prefix_str[prefix]}{idx}", **kwargs) + + +# This type is a little wider than it should be, because free_symbols says +# that it contains Basic, rather than Symbol +def symbol_is_type(sym: sympy.Basic, prefix: SymT | Iterable[SymT]) -> bool: + if not isinstance(sym, sympy.Symbol): + raise AssertionError("expected sympy.Symbol") + name_str = sym.name.lower() # Match capitalized names like XBLOCK, RBLOCK + if isinstance(prefix, SymT): + return name_str.startswith(prefix_str[prefix]) + else: + return name_str.startswith(tuple(prefix_str[p] for p in prefix)) + + +def free_symbol_is_type(e: sympy.Expr, prefix: SymT | Iterable[SymT]) -> bool: + return any(symbol_is_type(v, prefix) for v in e.free_symbols) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/value_ranges.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/value_ranges.py new file mode 100644 index 0000000000000000000000000000000000000000..ebd17c467fc39c7ebecc29423751ce45651ef556 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_sympy/value_ranges.py @@ -0,0 +1,1144 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import dataclasses +import functools +import itertools +import logging +import math +import operator +from collections.abc import Callable +from typing import Generic, overload, SupportsFloat, TYPE_CHECKING, TypeGuard, TypeVar +from typing_extensions import TypeIs + +import sympy +from sympy.logic.boolalg import Boolean as SympyBoolean, BooleanAtom + +import torch +from torch._logging import LazyString +from torch._prims_common import dtype_to_type + +from .functions import ( + _keep_float, + FloatTrueDiv, + FloorDiv, + IntTrueDiv, + OpaqueUnaryFn_exp, + OpaqueUnaryFn_log, + OpaqueUnaryFn_log2, + OpaqueUnaryFn_sqrt, + PowByNatural, + RoundDecimal, + RoundToInt, + safe_pow, + ToFloat, + TruncToFloat, + TruncToInt, +) +from .interp import sympy_interp +from .numbers import int_oo, IntInfinity, NegativeIntInfinity + + +log = logging.getLogger(__name__) + +__all__ = ["ValueRanges", "bound_sympy"] + +_T = TypeVar("_T", sympy.Expr, SympyBoolean) + + +class ValueRangeError(RuntimeError): + pass + + +# Like sympify, but supports less stuff, and also ensures that direct +# sympy expressions don't have free variables +def simple_sympify(e): + if isinstance(e, bool): + return sympy.true if e else sympy.false + elif isinstance(e, int): + return sympy.Integer(e) + elif isinstance(e, float): + # infinity is special; we use it to bracket integers as well + if math.isinf(e): + return sympy.oo if e > 0 else -sympy.oo + return sympy.Float(e) + elif isinstance(e, sympy.Expr): + if not getattr(e, "is_number", False): + raise AssertionError(e) + # NaNs can occur when doing things like 0 * sympy.oo, but it is better + # if the operator notices this and takes care of it, because sometimes + # the NaN is inappropriate (for example, for ints, the [-oo, oo] range + # should go to zero when multiplied with [0, 0]) + if e == sympy.nan: + raise AssertionError("sympy expression is NaN") + return e + elif isinstance(e, BooleanAtom): + return e + else: + raise AssertionError(f"not simple sympy type {type(e)}: {e}") + + +# Sympy atomics only. Unlike <=, it also works on Sympy bools. +def sympy_generic_le(lower, upper): + if isinstance(lower, sympy.Expr): + if not isinstance(upper, sympy.Expr): + raise AssertionError( + "upper must be a sympy.Expr when lower is a sympy.Expr" + ) + # instead of lower <= upper, we do upper >= lower since upper is mostly int_oo + # and we have better code paths there. + return upper >= lower + else: + # only negative condition is True > False + if not isinstance(lower, SympyBoolean) or not isinstance(upper, SympyBoolean): + raise AssertionError((lower, upper)) + return not (lower and not upper) + + +def vr_is_bool(vr: ValueRanges[_T]) -> TypeGuard[ValueRanges[SympyBoolean]]: + return vr.is_bool + + +def vr_is_expr(vr: ValueRanges[_T]) -> TypeGuard[ValueRanges[sympy.Expr]]: + return not vr.is_bool + + +def is_sympy_integer(value) -> TypeIs[sympy.Integer]: + return isinstance(value, sympy.Integer) + + +ExprIn = int | float | sympy.Expr +BoolIn = bool | SympyBoolean +AllIn = ExprIn | BoolIn +ExprFn = Callable[[sympy.Expr], sympy.Expr] +ExprFn2 = Callable[[sympy.Expr, sympy.Expr], sympy.Expr] +BoolFn = Callable[[SympyBoolean], SympyBoolean] +BoolFn2 = Callable[[SympyBoolean, SympyBoolean], SympyBoolean] +AllFn = ExprFn | BoolFn +AllFn2 = ExprFn2 | BoolFn2 + + +@dataclasses.dataclass(frozen=True) +class ValueRanges(Generic[_T]): + if TYPE_CHECKING: + # ruff doesn't understand circular references but mypy does + # pyrefly: ignore [unbound-name] + ExprVR = ValueRanges[sympy.Expr] # noqa: F821 + # pyrefly: ignore [unbound-name] + BoolVR = ValueRanges[SympyBoolean] # noqa: F821 + AllVR = ExprVR | BoolVR + + # Although the type signature here suggests you can pass any + # sympy expression, in practice the analysis here only works + # with constant sympy expressions + lower: _T + upper: _T + is_bool: bool + is_int: bool + is_float: bool + + def __repr__(self) -> str: + return f"VR[{self.lower}, {self.upper}]" + + @overload + def __init__( + self: ValueRanges[sympy.Expr], + lower: ExprIn, + upper: ExprIn, + ) -> None: ... + + @overload + def __init__( # type: ignore[misc] + self: ValueRanges[SympyBoolean], + lower: BoolIn, + upper: BoolIn, + ) -> None: ... + + def __init__(self, lower: AllIn, upper: AllIn) -> None: + lower = simple_sympify(lower) + upper = simple_sympify(upper) + # TODO: when the bounds have free variables, this may be + # nontrivial to actually verify + try: + if not sympy_generic_le(lower, upper): + raise ValueRangeError(f"Invalid ranges [{lower}:{upper}]") + except TypeError as e: + raise TypeError(f"Could not compare {lower} <= {upper}") from e + + is_bool_lower = isinstance(lower, SympyBoolean) + is_bool_upper = isinstance(upper, SympyBoolean) + if is_bool_lower != is_bool_upper: + raise AssertionError((lower, upper)) + + # Warning: is_int/is_float is best effort. We do pretty well in + # Dynamo, but in Inductor these attributes are often wrong because we + # are not very rigorous in dtype analysis. This is also why we need + # the flexible analysis for is_int: sometimes a sympy.oo pops in for + # an integer bound. I would /like/ for us not to do this, but it's + # too hard to push the invariant through right now. + if isinstance(lower, sympy.Integer) and upper == sympy.oo: + upper = int_oo + if isinstance(upper, sympy.Integer) and lower == -sympy.oo: + lower = -int_oo + # NB: [-int_oo, -int_oo] and [int_oo, int_oo] are allowed + integer_types = (sympy.Integer, NegativeIntInfinity, IntInfinity) + is_int_lower = isinstance(lower, integer_types) + is_int_upper = isinstance(upper, integer_types) + + # Because this is a frozen class + object.__setattr__(self, "lower", lower) + object.__setattr__(self, "upper", upper) + # Unlike bool/int in Python, we don't report bools are ints + # + # NB: is_bool_lower == is_bool_upper, so we only need to check one + object.__setattr__(self, "is_bool", is_bool_lower) + object.__setattr__( + self, + "is_int", + not self.is_bool and is_int_lower and is_int_upper, + ) + """ + # This assert is just impossible right now, too many sympy bugs + if self.is_int: + # NB: sympy will sometimes randomly lose the float-ness of zero, + # so we also need to account for that in the assertion here. + # See also https://github.com/sympy/sympy/issues/26620 + assert isinstance(lower, sympy.Integer) or lower in [-sympy.oo, 0], ( + lower, + upper, + ) + assert isinstance(upper, sympy.Integer) or upper in [sympy.oo, 0], (lower, upper) + """ + # NB: [-oo, oo] always advertises as float! + object.__setattr__(self, "is_float", not self.is_bool and not self.is_int) + if not self.is_bool and not self.is_int and not self.is_float: + raise AssertionError((lower, upper)) + + def boolify(self) -> ValueRanges[SympyBoolean]: + if vr_is_bool(self): + return self + elif self == ValueRanges.unknown(): + return ValueRanges.unknown_bool() + else: + raise AssertionError(f"not bool like {self}") + + def __contains__(self, x: AllIn) -> bool: + return ValueRanges.wrap(x).issubset(self) + + def issubset(self, other): + if other is self.unknown_int(): + return True + return sympy_generic_le(other.lower, self.lower) and sympy_generic_le( + self.upper, other.upper + ) + + def tighten(self, other) -> ValueRanges: + """Given two ValueRanges, returns their intersection""" + return self & other + + # Intersection + @overload + def __and__( + self: ValueRanges[sympy.Expr], + other: ValueRanges[sympy.Expr], + ) -> ValueRanges[sympy.Expr]: ... + + @overload + def __and__( # type: ignore[misc] + self: ValueRanges[SympyBoolean], + other: ValueRanges[SympyBoolean], + ) -> ValueRanges[SympyBoolean]: ... + + def __and__(self: AllVR, other: AllVR) -> AllVR: + if other in (ValueRanges.unknown(), ValueRanges.unknown_int()): + return self + if self in (ValueRanges.unknown(), ValueRanges.unknown_int()): + return other + if self.is_bool != other.is_bool: + raise AssertionError((self, other)) + if self.is_int != other.is_int: + raise AssertionError((self, other)) + if self.is_float != other.is_float: + raise AssertionError((self, other)) + if self.is_bool: + return ValueRanges( + sympy.Or(self.lower, other.lower), sympy.And(self.upper, other.upper) + ) + else: + return ValueRanges( + sympy.Max(self.lower, other.lower), sympy.Min(self.upper, other.upper) + ) + + # Union + @overload + def __or__( + self: ValueRanges[sympy.Expr], + other: ValueRanges[sympy.Expr], + ) -> ValueRanges[sympy.Expr]: ... + + @overload + def __or__( # type: ignore[misc] + self: ValueRanges[SympyBoolean], + other: ValueRanges[SympyBoolean], + ) -> ValueRanges[SympyBoolean]: ... + + def __or__(self: AllVR, other: AllVR) -> AllVR: + if ValueRanges.unknown() in (self, other): + return ValueRanges.unknown() + if self.is_bool != other.is_bool: + raise AssertionError((self, other)) + if self.is_int != other.is_int: + raise AssertionError((self, other)) + if self.is_float != other.is_float: + raise AssertionError((self, other)) + if self.is_bool: + return ValueRanges( + sympy.And(self.lower, other.lower), sympy.Or(self.upper, other.upper) + ) + else: + return ValueRanges( + sympy.Min(self.lower, other.lower), sympy.Max(self.upper, other.upper) + ) + + def is_singleton(self) -> bool: + return self.lower == self.upper + + @staticmethod + @functools.cache + def unknown() -> ValueRanges[sympy.Expr]: + return ValueRanges(-sympy.oo, sympy.oo) + + @staticmethod + @functools.cache + def unknown_int() -> ValueRanges[sympy.Expr]: + return ValueRanges(-int_oo, int_oo) + + @staticmethod + @functools.cache + def unknown_bool() -> ValueRanges[SympyBoolean]: + return ValueRanges(sympy.false, sympy.true) + + @overload + @staticmethod + # work around the fact that bool and int overlap + def wrap(arg: ExprIn | ExprVR) -> ExprVR: # type: ignore[overload-overlap] + ... + + @overload + @staticmethod + def wrap(arg: BoolIn | BoolVR) -> BoolVR: # type: ignore[misc] + ... + + @staticmethod + def wrap(arg: AllIn | AllVR) -> AllVR: + if isinstance(arg, ValueRanges): + return arg + if isinstance(arg, float) and math.isnan(arg): + return ValueRanges.unknown() + # arg is either ExprIn or BoolIn, but we don't know it here + return ValueRanges(arg, arg) # type: ignore[arg-type] + + @staticmethod + def increasing_map(x: ExprIn | ExprVR, fn: ExprFn) -> ExprVR: + """Increasing: x <= y => f(x) <= f(y).""" + x = ValueRanges.wrap(x) + return ValueRanges(fn(x.lower), fn(x.upper)) + + @overload + @staticmethod + def decreasing_map(x: ExprIn | ExprVR, fn: ExprFn) -> ExprVR: ... + + @overload + @staticmethod + def decreasing_map(x: BoolIn | BoolVR, fn: BoolFn) -> BoolVR: # type: ignore[misc] + ... + + @staticmethod + def decreasing_map(x: AllIn | AllVR, fn: AllFn) -> AllVR: + """Decreasing: x <= y => f(x) >= f(y).""" + x = ValueRanges.wrap(x) + # consistently either Expr or Bool, but we don't know it here + return ValueRanges(fn(x.upper), fn(x.lower)) # type: ignore[arg-type] + + @staticmethod + def monotone_map(x: ExprIn | ExprVR, fn: ExprFn) -> ExprVR: + """It's increasing or decreasing.""" + x = ValueRanges.wrap(x) + l = fn(x.lower) + u = fn(x.upper) + return ValueRanges(min(l, u), max(l, u)) + + @staticmethod + def convex_min_zero_map(x: ExprIn | ExprVR, fn: ExprFn) -> ExprVR: + """Fn is convex and has a minimum at 0.""" + x = ValueRanges.wrap(x) + if 0 in x: + upper = max(fn(x.lower), fn(x.upper)) + upper = simple_sympify(upper) + if isinstance(upper, sympy.Float) or upper == sympy.oo: + return ValueRanges(0.0, upper) + return ValueRanges(0, upper) + return ValueRanges.monotone_map(x, fn) + + @overload + @staticmethod + def coordinatewise_increasing_map( + x: ExprIn | ExprVR, + y: ExprIn | ExprVR, + fn: ExprFn2, + ) -> ExprVR: ... + + @overload + @staticmethod + def coordinatewise_increasing_map( # type: ignore[misc] + x: BoolIn | BoolVR, + y: BoolIn | BoolVR, + fn: BoolFn2, + ) -> BoolVR: ... + + @staticmethod + def coordinatewise_increasing_map( + x: AllIn | AllVR, + y: AllIn | AllVR, + fn: AllFn2, + ) -> AllVR: + """ + It's increasing on each coordinate. + + Mathematically: + For every 1 <= i <= n and x_i <= y_i we have that + f(x1, .., xn) <= f(x1, , yi, ..., xn) + """ + x, y = ValueRanges.wrap(x), ValueRanges.wrap(y) + return ValueRanges( + fn(x.lower, y.lower), # type: ignore[arg-type] + fn(x.upper, y.upper), # type: ignore[arg-type] + ) + + @classmethod + def coordinatewise_monotone_map(cls, x, y, fn): + """It's increasing or decreasing on each coordinate.""" + x, y = cls.wrap(x), cls.wrap(y) + products = [ + fn(a, b) + for a, b in itertools.product([x.lower, x.upper], [y.lower, y.upper]) + ] + return ValueRanges(min(products), max(products)) + + +class SymPyValueRangeAnalysis: + """ + It gives bounds on a SymPy operator given bounds on its arguments + See the function `bound_sympy` for a function that applies this logic to a full SymPy expression + """ + + @staticmethod + def constant(value, dtype): + if isinstance(value, ValueRanges): + if not value.is_singleton(): + raise AssertionError("ValueRanges must be a singleton for constant()") + value = value.lower + # NB: value is NOT a sympy expression, it's a constant! + is_python = isinstance(value, (int, float, bool)) + if not is_python and not isinstance( + value, (BooleanAtom, sympy.Integer, sympy.Number) + ): + raise AssertionError(f"not a supported constant type: {type(value)}") + + # using nan makes subsequent computation throw, and for the purposes of optimization + # returning -math.inf - math.inf is equivalent to giving up + if isinstance(value, SupportsFloat) and math.isnan(value): + if dtype == torch.bool: + return ValueRanges.unknown_bool() + elif dtype.is_floating_point: + return ValueRanges.unknown() + else: + return ValueRanges.unknown_int() + + if is_python: + type_ = dtype_to_type(dtype) + value = type_(value) + else: + # We do a type check on a best-effort basis + # We don't want to force a cast to sympy.Float if the value is Rational to avoid losing precision + if dtype == torch.bool: + if not isinstance(value, BooleanAtom): + raise AssertionError("expected BooleanAtom for bool dtype") + elif dtype.is_floating_point: + # pyrefly: ignore [missing-attribute] + if value.is_finite and not value.is_real: + raise AssertionError( + "expected float-like sympy value for float dtype" + ) + else: + # dtype is intXX + if not getattr(value, "is_integer", False): + raise AssertionError("expected integer sympy value for int dtype") + + r = ValueRanges.wrap(value) + return r + + @staticmethod + def to_dtype(a, dtype, src_dtype=None): + if dtype == torch.float64: + # pyrefly: ignore [bad-argument-type] + return ValueRanges.increasing_map(a, ToFloat) + elif dtype == torch.bool: + return ValueRanges.unknown_bool() + elif not dtype.is_floating_point: + return ValueRanges.unknown_int() + return ValueRanges.unknown() + + @staticmethod + def trunc_to_int(a, dtype): + # pyrefly: ignore [bad-argument-type] + return ValueRanges.increasing_map(a, TruncToInt) + + @staticmethod + def not_(a): + a = ValueRanges.wrap(a) + a = a.boolify() + if not a.is_bool: + raise AssertionError("not_ expects a boolean ValueRanges") + return ValueRanges.decreasing_map(a, sympy.Not) + + @staticmethod + def or_(a, b): + return ValueRanges.coordinatewise_increasing_map(a, b, sympy.Or) + + @staticmethod + def and_(a, b): + return ValueRanges.coordinatewise_increasing_map(a, b, sympy.And) + + @staticmethod + def _bool_to_int(x): + if x.is_singleton(): + return ValueRanges.wrap(sympy.Integer(1 if x.lower else 0)) + else: + return ValueRanges(sympy.Integer(0), sympy.Integer(1)) + + @classmethod + def bitwise_and(cls, a, b): + a, b = ValueRanges.wrap(a), ValueRanges.wrap(b) + if a.is_bool and b.is_bool: + return cls.and_(a, b) + if a.is_bool: + a = cls._bool_to_int(a) + if b.is_bool: + b = cls._bool_to_int(b) + lower = min(a.lower, b.lower) + if lower < 0 and lower != -sympy.oo and lower != -int_oo: + # If both lower bounds are negative, then bits start like + # 1...10..., so the smallest possible value is 1...101...1. + # Thus, we need to find the next smallest power of 2 (inclusive). + try: + lower = -(1 << int(-lower - 1).bit_length()) + except Exception: + lower = -int_oo + else: + lower = 0 + return ValueRanges(lower, max(a.upper, b.upper)) + + @classmethod + def bitwise_or(cls, a, b): + a, b = ValueRanges.wrap(a), ValueRanges.wrap(b) + if a.is_bool and b.is_bool: + return cls.or_(a, b) + if a.is_bool: + a = cls._bool_to_int(a) + if b.is_bool: + b = cls._bool_to_int(b) + upper = max(a.upper, b.upper) + if upper == 0: + upper = 0 + elif upper > 0 and upper != sympy.oo and upper != int_oo: + # If both upper bounds are positive, then the largest + # possible value is 01...1, so we need to find + # next largest power of 2 (exclusive), minus 1 + try: + upper = (1 << int(upper).bit_length()) - 1 + except Exception: + upper = int_oo + elif upper < 0: + upper = -1 + return ValueRanges(min(a.lower, b.lower), upper) + + @classmethod + def bitwise_xor(cls, a, b): + a, b = ValueRanges.wrap(a), ValueRanges.wrap(b) + if a.is_bool and b.is_bool: + bounds = { + a.lower ^ b.lower, + a.lower ^ b.upper, + a.upper ^ b.lower, + a.upper ^ b.upper, + } + + has_false = any(bound == sympy.false for bound in bounds) + has_true = any(bound == sympy.true for bound in bounds) + + if has_false and has_true: + lower, upper = sympy.false, sympy.true + elif has_true: + lower = upper = sympy.true + elif has_false: + lower = upper = sympy.false + else: + raise AssertionError(f"Non-boolean xor result: {bounds}") + + return ValueRanges(lower, upper) + if a.is_bool: + a = cls._bool_to_int(a) + if b.is_bool: + b = cls._bool_to_int(b) + if ( + a.lower == a.upper + and b.lower == b.upper + and is_sympy_integer(a.lower) + and is_sympy_integer(b.lower) + ): + value_range = a.lower ^ b.lower + return ValueRanges(value_range, value_range) + return ValueRanges(-int_oo, int_oo) + + @staticmethod + def eq(a, b): + a = ValueRanges.wrap(a) + b = ValueRanges.wrap(b) + if a.is_singleton() and b.is_singleton() and a.lower == b.lower: + return ValueRanges.wrap(sympy.true) + elif a.lower > b.upper or b.lower > a.upper: # ranges disjoint + return ValueRanges.wrap(sympy.false) + return ValueRanges(sympy.false, sympy.true) + + @classmethod + def ne(cls, a, b): + return cls.not_(cls.eq(a, b)) + + @classmethod + def identity(cls, a): + return ValueRanges.wrap(a) + + @classmethod + def lt(cls, a, b): + a = ValueRanges.wrap(a) + b = ValueRanges.wrap(b) + if a.is_bool != b.is_bool: + raise AssertionError( + "operands must both be boolean ValueRanges or both non-boolean" + ) + if a.is_bool: + return cls.and_(cls.not_(a), b) + else: + if a.upper < b.lower: + return ValueRanges.wrap(sympy.true) + elif a.lower >= b.upper: + return ValueRanges.wrap(sympy.false) + return ValueRanges(sympy.false, sympy.true) + + @classmethod + def gt(cls, a, b): + return cls.lt(b, a) + + @classmethod + def le(cls, a, b): + return cls.not_(cls.gt(a, b)) + + @classmethod + def ge(cls, a, b): + return cls.not_(cls.lt(a, b)) + + @staticmethod + def add(a, b): + return ValueRanges.coordinatewise_increasing_map( + a, b, _keep_float(operator.add) + ) + + @classmethod + def mul(cls, a, b): + a = ValueRanges.wrap(a) + b = ValueRanges.wrap(b) + + if a.is_bool != b.is_bool: + raise AssertionError( + "operands must both be boolean ValueRanges or both non-boolean" + ) + if a.is_bool: + return cls.and_(a, b) + + def safe_mul(a, b): + # Make unknown() * wrap(0.0) == wrap(0.0) + if a == 0.0 or a == 0: + return a + elif b == 0.0 or b == 0: + return b + else: + return a * b + + return ValueRanges.coordinatewise_monotone_map(a, b, _keep_float(safe_mul)) + + @staticmethod + def int_truediv(a, b): + a = ValueRanges.wrap(a) + b = ValueRanges.wrap(b) + if 0 in b or ((-int_oo in a or int_oo in a) and (-int_oo in b or int_oo in b)): + return ValueRanges.unknown() + else: + return ValueRanges.coordinatewise_monotone_map( + a, + b, + _keep_float(IntTrueDiv), + ) + + @staticmethod + def truediv(a, b): + a = ValueRanges.wrap(a) + b = ValueRanges.wrap(b) + if 0 in b or ( + (-sympy.oo in a or sympy.oo in a) and (-sympy.oo in b or sympy.oo in b) + ): + return ValueRanges.unknown() + else: + return ValueRanges.coordinatewise_monotone_map( + a, + b, + _keep_float(FloatTrueDiv), + ) + + @staticmethod + def floordiv(a, b): + a = ValueRanges.wrap(a) + b = ValueRanges.wrap(b) + + # TODO We shall assume division is always valid probably. + if 0 in b: + if b.lower >= 0 and a.lower >= 0: + return ValueRanges(0, int_oo) + if b.upper <= 0 and a.upper <= 0: + return ValueRanges(0, int_oo) + if b.upper <= 0 and a.lower >= 0: + return ValueRanges(-int_oo, 0) + if b.lower >= 0 and a.upper <= 0: + return ValueRanges(-int_oo, 0) + return ValueRanges.unknown_int() + products = [] + for x, y in itertools.product([a.lower, a.upper], [b.lower, b.upper]): + r = FloorDiv(x, y) + if r is sympy.nan: + products.append((sympy.sign(x) * sympy.sign(y)) * int_oo) + else: + products.append(r) + + return ValueRanges(min(products), max(products)) + + @classmethod + def mod(cls, x, y): + x = ValueRanges.wrap(x) + y = ValueRanges.wrap(y) + # nb. We implement C semantics + + def c_mod(a, b): + ret = abs(a) % abs(b) + if a < 0: + ret *= -1 + return ret + + def c_div(a, b): + x = a / b + return sympy.Integer(x) if x.is_finite and x not in (int_oo, -int_oo) else x + + if 0 in y: + return ValueRanges.unknown_int() + elif y.is_singleton(): + y_val = abs(y.lower) + # If it wraps, we need to take the whole interval + + # The function is locally linear if they are in the same class + if c_div(x.lower, y_val) == c_div(x.upper, y_val): + return ValueRanges.increasing_map(x, lambda u: c_mod(u, y_val)) + if x.upper < 0: + # Negative case + return ValueRanges(-y_val + 1, 0) + elif x.lower > 0: + # Positive case + return ValueRanges(0, y_val - 1) + else: + # Mixed case + lower = max(-y_val + 1, x.lower) + upper = min(y_val - 1, x.upper) + return ValueRanges(lower, upper) + else: + # Too difficult, we bail out + upper = cls.abs(y).upper - 1 + return ValueRanges(-upper, upper) + + @classmethod + def python_mod(cls, x, y): + """Python-style modulo: result has same sign as divisor. + + Assumes valid input where y is never 0. + - When y > 0: result is in [0, y - 1] + - When y < 0: result is in [y + 1, 0] + """ + + x = ValueRanges.wrap(x) + y = ValueRanges.wrap(y) + if x.lower >= 0 and y.lower >= 0: + return SymPyValueRangeAnalysis.mod(x, y) + lower = y.lower + 1 if y.lower < 0 else 0 + upper = y.upper - 1 if y.upper > 0 else 0 + return ValueRanges(lower, upper) + + @classmethod + def modular_indexing(cls, a, b, c): + return cls.mod(cls.floordiv(a, b), c) + + @classmethod + def is_non_overlapping_and_dense_indicator(cls, *args): + return ValueRanges.unknown_int() + + @classmethod + def pow_by_natural(cls, a, b): + a = ValueRanges.wrap(a) + b = ValueRanges.wrap(b) + if a.is_singleton() and b.is_singleton(): + return ValueRanges.wrap(safe_pow(a.lower, b.lower)) + # NB: Exclude zero, because zero is special + elif a.lower >= 1: + # We should know that b >= 0 but we may have forgotten this fact due + # to replacements, so don't assert it, but DO clamp it to prevent + # degenerate problems + # pyrefly: ignore [no-matching-overload] + return ValueRanges.coordinatewise_increasing_map( + a, b & ValueRanges(0, int_oo), PowByNatural + ) + elif b.is_singleton(): + if b.lower % 2 == 0: + # x^n where n is even + return ValueRanges.convex_min_zero_map( + a, lambda x: safe_pow(x, b.lower) + ) + else: + # x^n where n is odd + return ValueRanges.increasing_map(a, lambda x: safe_pow(x, b.lower)) + else: + # a is potentially negative, and we don't know if the exponent is + # even or odd. So just conservatively set the upper and lower + # bound based on what the maximum absolute value could be, in both + # directions + max_base = max(a.upper, -a.lower) + return ValueRanges( + -(safe_pow(max_base, b.upper)), safe_pow(max_base, b.upper) + ) + + @classmethod + def pow(cls, a, b): + return ValueRanges.unknown() + + # We could implement all this, but for floating point pow, is there + # really a point? + """ + a = ValueRanges.wrap(a) + b = ValueRanges.wrap(b) + + # Not implemented yet. It's a bit tricky + # If you want to implement it, compute the partial derivatives of a ** b + # and check the ranges where the function is increasing / decreasing + # Another non-tight way of doing this is defaulting to doing noting that for a > 0, a ** b == exp(b * log(a)) + # If this second option is implemented, by carefult about the types and possible infinities here and there. + if not b.is_singleton(): + return ValueRanges.unknown() + + b = b.lower + if a.is_singleton(): + a = a.lower + r = a**b + if not r.is_finite: + return ValueRanges.unknown() + return ValueRanges.wrap(r) + + if b == 0: + if not a.lower.is_finite: + return ValueRanges.unknown() + return ValueRanges.wrap(1.0) + + if b < 0: + a = cls.reciprocal(a) + b = -b + + if a == ValueRanges.unknown(): + return ValueRanges.unknown() + + # If the base is positive, then we're good, otherwise nothing's defined + if a.lower >= 0: + return ValueRanges.increasing_map(a, lambda x: x**b) + else: + return ValueRanges.unknown() + """ + + @staticmethod + def reciprocal(x): + """Needed as it's used in pow, but it won't appear on a SymPy expression""" + x = ValueRanges.wrap(x) + if 0 in x: + return ValueRanges.unknown() + else: + return ValueRanges.decreasing_map(x, lambda y: FloatTrueDiv(1.0, y)) # type: ignore[operator] + + @staticmethod + def abs(x): + return ValueRanges.convex_min_zero_map(x, abs) + + @staticmethod + def exp(x): + return ValueRanges.increasing_map(x, OpaqueUnaryFn_exp) + + @staticmethod + def log(x): + x = ValueRanges.wrap(x) + if x.lower <= 0: + return ValueRanges.unknown() + return ValueRanges.increasing_map(x, OpaqueUnaryFn_log) + + @staticmethod + def log2(x): + x = ValueRanges.wrap(x) + if x.lower <= 0: + return ValueRanges.unknown() + return ValueRanges.increasing_map(x, OpaqueUnaryFn_log2) + + @classmethod + def minimum(cls, a, b): + return cls.min_or_max(a, b, sympy.Min) + + @classmethod + def maximum(cls, a, b): + return cls.min_or_max(a, b, sympy.Max) + + @staticmethod + def min_or_max(a, b, fn): + a = ValueRanges.wrap(a) + b = ValueRanges.wrap(b) + return ValueRanges.coordinatewise_increasing_map(a, b, fn) + + @classmethod + def floor_to_int(cls, x, dtype): + return ValueRanges.increasing_map(x, sympy.functions.elementary.integers.floor) + + @classmethod + def ceil_to_int(cls, x, dtype): + return ValueRanges.increasing_map( + x, sympy.functions.elementary.integers.ceiling + ) + + # I think these implementations are sound. The hazard here is that sympy + # will carry out the floor/ceil at too high precision and then something + # bad will happen when we convert it to float. + # + # For truncation, the implementation is clearly sound, because the desired + # target float is always exactly representable, since you're just chopping + # off bits the mantissa. But what about ceil/floor? + # + # The important constraint here is that we're not defining floor on + # arbitrary real numbers, only representable float numbers. So we can + # take advantage of the fact that before we reach the first + # unrepresentable integer in floating point space, we have the range of + # numbers corresponding to exponent zero: all integers, with no fractional + # amounts. floor/ceil is an identity operation in this case. In the + # range below here, representable floating point numbers are spaced + # exactly 1/2 apart, and notably, both the floor/ceil are defined floating + # point numbers. There is no "gap" as you step up to the next exponent. + + @classmethod + def floor(cls, x): + return ValueRanges.increasing_map( + x, _keep_float(sympy.functions.elementary.integers.floor) + ) + + @classmethod + def ceil(cls, x): + return ValueRanges.increasing_map( + x, _keep_float(sympy.functions.elementary.integers.ceiling) + ) + + @classmethod + def round_decimal(cls, number, ndigits): + if not ndigits.is_singleton(): + return ValueRanges.unknown() + + ndigits = ndigits.lower + # We can't use functools.partial here since sympy doesn't support keyword arguments, but we have to bind + # the second parameter. + fn = lambda number: RoundDecimal(number, ndigits) # type: ignore[misc, assignment] # noqa: E731 + + return ValueRanges.increasing_map(number, fn) + + @classmethod + def round_to_int(cls, number, dtype): + # pyrefly: ignore [bad-argument-type] + return ValueRanges.increasing_map(number, RoundToInt) + + # It's used in some models on symints + @staticmethod + def sqrt(x): + x = ValueRanges.wrap(x) + if x.lower < 0: + return ValueRanges.unknown() + return ValueRanges.increasing_map(x, OpaqueUnaryFn_sqrt) + + @staticmethod + def where(a, b, c): + b = ValueRanges.wrap(b) + c = ValueRanges.wrap(c) + a = a.boolify() + # We sometimes write unknown without specifying the type correctly + # In particular, we do that when initialising the bounds for loads in bounds.py + if b.is_bool != c.is_bool and ValueRanges.unknown() not in (b, c): + raise AssertionError( + "where() requires b and c to have the same boolean-ness or allow unknown()" + ) + if b.is_bool: + return ValueRanges(sympy.And(b.lower, c.lower), sympy.Or(b.upper, c.upper)) + else: + return ValueRanges(sympy.Min(b.lower, c.lower), sympy.Max(b.upper, c.upper)) + + # expr_cond_pair is used to represent a single (expr, condition) pair in piecewise. + # We just return the value range of the expression and its corresponding condition as a tuple + # and defer the analysis to piecewise + @staticmethod + def expr_cond_pair(a, b): + b = b.boolify() + return (a, b) + + # piecewise function can be used to convert a SymBool to SymInt: + # int_expr = Piecewise((1, bool_expr), (0, True)), it evaluates to 1 when sym_bool is True and 0 otherwise. + # + # ranges is a sequence of (expr_range, condition_range) pairs. The range pair is constructed in expr_cond_pair. + # The ValueRange of Piecewise is just the union of all expr ranges whose condition expr can be True. + @staticmethod + def piecewise(*ranges): + init_range = None + for expr_range, cond_range in ranges: + if sympy.true in cond_range: + if init_range is None: + init_range = expr_range + else: + init_range = init_range | expr_range + return init_range + + @staticmethod + def cos(x): + # TODO: We should tighten value ranges + # If input range span is pi + 2*pi*k, then output range is (-1, 1) + # otherwise the minimum of the value of the function on the extremes + return ValueRanges(-1.0, 1.0) + + @staticmethod + def cosh(x): + return ValueRanges(0.0, sympy.oo) + """ + x = ValueRanges.wrap(x) + if x.lower > 0: + return ValueRanges.increasing_map(x, OpaqueUnaryFn_cosh) + elif x.upper < 0: + return ValueRanges.decreasing_map(x, OpaqueUnaryFn_cosh) + return ValueRanges(0.0, sympy.oo) + """ + + @staticmethod + def sin(x): + # TODO: We should tighten value ranges + # See details on cos + return ValueRanges(-1.0, 1.0) + + @staticmethod + def sinh(x): + # return ValueRanges.increasing_map(x, OpaqueUnaryFn_sinh) + return ValueRanges(-sympy.oo, sympy.oo) + + @staticmethod + def tan(x): + return ValueRanges(-sympy.oo, sympy.oo) + + @staticmethod + def tanh(x): + # return ValueRanges.increasing_map(x, OpaqueUnaryFn_tanh) + return ValueRanges(-sympy.oo, sympy.oo) + + @staticmethod + def asin(x): + return ValueRanges(-sympy.oo, sympy.oo) + """ + x = ValueRanges.wrap(x) + if -1 <= x.lower and x.upper <= 1: + return ValueRanges.increasing_map(x, OpaqueUnaryFn_asinh) + return ValueRanges.unknown() + """ + + @staticmethod + def acos(x): + return ValueRanges(-sympy.oo, sympy.oo) + """ + x = ValueRanges.wrap(x) + if -1 <= x.lower and x.upper <= 1: + return ValueRanges.decreasing_map(x, OpaqueUnaryFn_acos) + return ValueRanges.unknown() + """ + + @staticmethod + def atan(x): + return ValueRanges(-sympy.oo, sympy.oo) + # return ValueRanges.increasing_map(x, OpaqueUnaryFn_atan) + + @staticmethod + def trunc(x): + # pyrefly: ignore [bad-argument-type] + return ValueRanges.increasing_map(x, TruncToFloat) + + +def bound_sympy( + expr: sympy.Expr, ranges: dict[sympy.Symbol, ValueRanges] | None = None +) -> ValueRanges: + log.debug( + "bound_sympy(%s)%s", + expr, + LazyString( + lambda: ( + "\n" + + "\n".join( + f" {k}: {r}" for k, r in ranges.items() if k in expr.free_symbols + ) + if ranges + else "" + ) + ), + ) + if isinstance(expr, sympy.Number): + return ValueRanges.wrap(expr) + + ranges = ranges or {} + + # If there's a tracing context, augment available constrained ranges. + context = torch._guards.TracingContext.try_get() + if context and context.fake_mode and context.fake_mode.shape_env: + if ranges: + ranges = {**context.fake_mode.shape_env.var_to_range, **ranges} + else: + ranges = context.fake_mode.shape_env.var_to_range + + def missing_handler(s): + if s.is_integer: # type: ignore[attr-defined] + if s.is_positive: # type: ignore[attr-defined] + vr = ValueRanges(1, int_oo) + elif s.is_nonnegative: # type: ignore[attr-defined] + vr = ValueRanges(0, int_oo) + else: + vr = ValueRanges.unknown_int() + else: + # Don't bother trying very hard here + vr = ValueRanges.unknown() + return vr + + return sympy_interp( + SymPyValueRangeAnalysis, ranges, expr, missing_handler=missing_handler + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_thunk.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_thunk.py new file mode 100644 index 0000000000000000000000000000000000000000..b5ab598077f4e8d3d9de9169a1352918771f07f6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_thunk.py @@ -0,0 +1,29 @@ +from collections.abc import Callable +from typing import Generic, TypeVar + + +R = TypeVar("R") + + +class Thunk(Generic[R]): + """ + A simple lazy evaluation implementation that lets you delay + execution of a function. It properly handles releasing the + function once it is forced. + """ + + f: Callable[[], R] | None + r: R | None + + __slots__ = ["f", "r"] + + def __init__(self, f: Callable[[], R]) -> None: + self.f = f + self.r = None + + def force(self) -> R: + if self.f is None: + return self.r # type: ignore[return-value] + self.r = self.f() + self.f = None + return self.r diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_traceback.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_traceback.py new file mode 100644 index 0000000000000000000000000000000000000000..f5415002092a23350f7d7d3436388d34b1eb9501 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_traceback.py @@ -0,0 +1,260 @@ +# mypy: allow-untyped-defs +import contextlib +import inspect +import os.path +import tempfile +import traceback +from types import TracebackType + + +# This file contains utilities for ensuring dynamically compile()'d +# code fragments display their line numbers in backtraces. +# +# The constraints: +# +# - We don't have control over the user exception printer (in particular, +# we cannot assume the linecache trick will work, c.f. +# https://stackoverflow.com/q/50515651/23845 ) +# +# - We don't want to create temporary files every time we compile() +# some code; file creation should happen lazily only at exception +# time. Arguably, you *should* be willing to write out your +# generated Python code to file system, but in some situations +# (esp. library code) it would violate user expectation to write +# to the file system, so we try to avoid it. In particular, we'd +# like to keep the files around, so users can open up the files +# mentioned in the trace; if the file is invisible, we want to +# avoid clogging up the filesystem. +# +# If this is not a constraint for you, there is a substantially simpler +# way to implement the functionality in this PR: instead of using +# eval/exec directly, just always write a Python file to filesystem +# and compile that. +# +# - You have control over a context where the compiled code will get +# executed, so that we can interpose while the stack is unwinding +# (otherwise, we have no way to interpose on the exception printing +# process.) +# +# There are two things you have to do to make use of the utilities here: +# +# - When you compile your source code, you must save its string source +# in its f_globals under the magic name "__compile_source__" +# +# - Before running the compiled code, enter the +# report_compile_source_on_error() context manager. + + +@contextlib.contextmanager +def report_compile_source_on_error(): + try: + yield + except Exception as exc: + tb = exc.__traceback__ + + # Walk the traceback, looking for frames that have + # source attached + stack = [] + while tb is not None: + filename = tb.tb_frame.f_code.co_filename + source = tb.tb_frame.f_globals.get("__compile_source__") + + if filename == "" and source is not None: + # What black magic are we doing here? Intuitively, what + # we would like to do is overwrite the co_filename on any + # frames that were generated from exec/eval so that they + # point to a temporary file that has the actual line + # information, so Python's default error printer can print + # useful line information on it. + # + # Writing out the temporary file is easy. But overwriting + # co_filename is not! You can't modify the code object + # associated with a frame. You can, however, reconstruct + # a traceback with entirely new frames from scratch, so that's + # what we do. But there's another problem, which is how to + # make the frame? + # + # The black magic is we make a frankenstein frame and code + # object which resembles the original frame/code enough so + # that it will print properly under traceback and the default + # error printer, but IT IS NOT THE ORIGINAL FRAME (you + # couldn't, e.g., execute its code with different variables + # and expect it to work.) + + # Don't delete the temporary file so the user can inspect it + # TODO: This creates a temporary file for every frame, but we + # technically only need one per distinct __compile_source__ + with tempfile.NamedTemporaryFile( + mode="w", delete=False, suffix=".py" + ) as f: + f.write(source) + # Create a frame. Python doesn't let you construct + # FrameType directly, so just make one with compile + frame = tb.tb_frame + code = compile("__inspect_currentframe()", f.name, "eval") + code = code.replace(co_name=frame.f_code.co_name) + # Python 3.11 only + if hasattr(frame.f_code, "co_linetable"): + # We can't copy ALL of the metadata over, because you + # can cause Python to segfault this way. What exactly + # do we need? We need enough information for + # traceback to be able to print the exception + # correctly. Code reading Lib/traceback.py reveals + # that traceback calls code.co_positions() in order to + # get the augmented line/col numbers. Objects/codeobject.c, + # specifically _PyCode_InitAddressRange, reveals that + # this iterator is initialized from co_linetable and + # co_firstfileno. So copy these we must! + code = code.replace( # type: ignore[call-arg] + co_linetable=frame.f_code.co_linetable, # type: ignore[attr-defined] + co_firstlineno=frame.f_code.co_firstlineno, # type: ignore[attr-defined] + ) + fake_frame = eval( + code, + frame.f_globals, + {**frame.f_locals, "__inspect_currentframe": inspect.currentframe}, + ) + fake_tb = TracebackType(None, fake_frame, tb.tb_lasti, tb.tb_lineno) + stack.append(fake_tb) + else: + stack.append(tb) + + tb = tb.tb_next + + # Reconstruct the linked list + tb_next = None + for tb in reversed(stack): + tb.tb_next = tb_next + tb_next = tb + + raise exc.with_traceback(tb_next) # noqa: B904 + + +def shorten_filename(fn, *, base=None): + """Shorten a source filepath, with the assumption that torch/ subdirectories don't need to be shown to user.""" + if base is None: + base = os.path.dirname(os.path.dirname(__file__)) + # Truncate torch/foo.py to foo.py + try: + prefix = os.path.commonpath([fn, base]) + except ValueError: + return fn + else: + return fn[len(prefix) + 1 :] + + +def format_frame(frame, *, base=None, line=False) -> str: + """ + Format a FrameSummary in a short way, without printing full absolute path or code. + + The idea is the result fits on a single line. + """ + extra_line = "" + if line: + extra_line = f"{frame.line} # " + return f"{extra_line}{shorten_filename(frame.filename, base=base)}:{frame.lineno} in {frame.name}" + + +def format_traceback_short(tb): + """Format a TracebackType in a short way, printing only the inner-most frame.""" + return format_frame(traceback.extract_tb(tb)[-1]) + + +class CapturedTraceback: + __slots__ = ["tb", "skip"] + + def __init__(self, tb, skip=0) -> None: + self.tb = tb + self.skip = skip + + def cleanup(self) -> None: + self.tb = None + + def summary(self): + import torch._C._profiler + + if self.tb is None: + # TODO: Maybe indicate that the traceback was elided? + return traceback.StackSummary() + + return _extract_symbolized_tb( + torch._C._profiler.symbolize_tracebacks([self.tb])[0], self.skip + ) + + def __getstate__(self): + return ( + None, + { + "tb": None, # TB is not pickleable + "skip": self.skip, + }, + ) + + @staticmethod + def extract(*, script=False, cpp=False, skip=0): + """ + Like traceback.extract_stack(), but faster (approximately 20x faster); it + is fast enough that you can unconditionally log stacks this way as part of + normal execution. It returns a torch._C._profiler.CapturedTraceback + object that must be formatted specially with format_captured_tb. + + By default, this only reports Python backtraces (like extract_stack). You + can set the script/cpp kwargs to also turn on TorchScript/C++ trace + reporting. + """ + import torch._C._profiler + + if script or cpp: + if skip != 0: + raise AssertionError("skip with script/cpp NYI") + + return CapturedTraceback( + torch._C._profiler.gather_traceback(python=True, script=script, cpp=cpp), + # Elide extract() frame if we don't have script/cpp frames. If + # we do have those frames, it doesn't work so force zero. + 0 if script or cpp else skip + 1, + ) + + def format(self): + """ + Formats a single torch._C._profiler.CapturedTraceback into a list of + strings equivalent to the output of traceback.format_list. Note that if + pass it CapturedTraceback with C++ traces, it is better not to use this + function and use the batch formatting API format_captured_tbs to amortize + the cost of symbolization + """ + return traceback.format_list(self.summary()) + + @staticmethod + def format_all(tbs): + """ + Bulk version of CapturedTraceback.format. Returns a list of list of strings. + """ + import torch._C._profiler + + # Directly populate tracebacks that already have cached summaries + rs: list[list[str] | None] = [] + delayed_idxs = [] + for i, tb in enumerate(tbs): + if tb.tb is None: + rs.append([]) + else: + rs.append(None) + delayed_idxs.append(i) + + torch._C._profiler.symbolize_tracebacks([tbs[i].tb for i in delayed_idxs]) + for i in delayed_idxs: + rs[i] = traceback.format_list(tbs[i].summary()) + + return rs + + +def _extract_symbolized_tb(tb, skip): + """ + Given a symbolized traceback from symbolize_tracebacks, return a StackSummary object of + pre-processed stack trace entries. + """ + stack = traceback.StackSummary() + for f in reversed(tb[skip:]): + stack.append(traceback.FrameSummary(f["filename"], f["line"], f["name"])) + return stack diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_triton.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_triton.py new file mode 100644 index 0000000000000000000000000000000000000000..075aa5a332234ec41cc153ebd34c4ffb586a4013 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_triton.py @@ -0,0 +1,223 @@ +import functools +import hashlib +import os +from typing import Any + + +@functools.cache +def has_triton_package() -> bool: + try: + import triton # noqa: F401 + + return True + except ImportError: + return False + + +@functools.cache +def get_triton_version(fallback: tuple[int, int] = (0, 0)) -> tuple[int, int]: + try: + import triton + + major, minor = tuple(int(v) for v in triton.__version__.split(".")[:2]) + return (major, minor) + except ImportError: + return fallback + + +@functools.cache +def _device_supports_tma() -> bool: + import torch + + return ( + torch.cuda.is_available() + and torch.cuda.get_device_capability() >= (9, 0) + and not torch.version.hip + ) + + +@functools.cache +def has_triton_experimental_host_tma() -> bool: + if has_triton_package(): + if _device_supports_tma(): + try: + from triton.tools.experimental_descriptor import ( # noqa: F401 + create_1d_tma_descriptor, + create_2d_tma_descriptor, + ) + + try: + from triton.tools.experimental_descriptor import enable_in_pytorch + + return enable_in_pytorch() + except ImportError: + return True + except ImportError: + pass + + return False + + +@functools.cache +def has_triton_tensor_descriptor_host_tma() -> bool: + if has_triton_package(): + if _device_supports_tma(): + try: + from triton.tools.tensor_descriptor import ( # noqa: F401 + TensorDescriptor, + ) + + return True + except ImportError: + pass + + return False + + +@functools.cache +def has_triton_tma() -> bool: + return has_triton_tensor_descriptor_host_tma() or has_triton_experimental_host_tma() + + +@functools.cache +def has_triton_tma_device() -> bool: + if has_triton_package(): + import torch + + if ( + torch.cuda.is_available() + and torch.cuda.get_device_capability() >= (9, 0) + and not torch.version.hip + ) or torch.xpu.is_available(): + # old API + try: + from triton.language.extra.cuda import ( # noqa: F401 + experimental_device_tensormap_create1d, + experimental_device_tensormap_create2d, + ) + + return True + except ImportError: + pass + + # new API + try: + from triton.language import make_tensor_descriptor # noqa: F401 + + return True + except ImportError: + pass + + return False + + +@functools.cache +def has_datacenter_blackwell_tma_device() -> bool: + import torch + + if ( + torch.cuda.is_available() + and torch.cuda.get_device_capability() >= (10, 0) + and torch.cuda.get_device_capability() < (11, 0) + and not torch.version.hip + ): + return has_triton_tma_device() and has_triton_tensor_descriptor_host_tma() + + return False + + +@functools.lru_cache(None) +def has_triton_stable_tma_api() -> bool: + if has_triton_package(): + import torch + + if ( + torch.cuda.is_available() + and torch.cuda.get_device_capability() >= (9, 0) + and not torch.version.hip + ) or torch.xpu.is_available(): + try: + from triton.language import make_tensor_descriptor # noqa: F401 + + return True + except ImportError: + pass + return False + + +@functools.cache +def has_triton() -> bool: + if not has_triton_package(): + return False + + from torch._inductor.config import triton_disable_device_detection + + if triton_disable_device_detection: + return False + + from torch._dynamo.device_interface import get_interface_for_device + + def cuda_extra_check(device_interface: Any) -> bool: + return device_interface.Worker.get_device_properties().major >= 7 + + def cpu_extra_check(device_interface: Any) -> bool: + import triton.backends + + return "cpu" in triton.backends.backends + + def _return_true(device_interface: Any) -> bool: + return True + + triton_supported_devices = { + "cuda": cuda_extra_check, + "xpu": _return_true, + "cpu": cpu_extra_check, + "mtia": _return_true, + } + + def is_device_compatible_with_triton() -> bool: + for device, extra_check in triton_supported_devices.items(): + device_interface = get_interface_for_device(device) + if device_interface.is_available() and extra_check(device_interface): + return True + return False + + return is_device_compatible_with_triton() + + +@functools.cache +def triton_backend() -> Any: + from triton.compiler.compiler import make_backend + from triton.runtime.driver import driver + + target = driver.active.get_current_target() + return make_backend(target) + + +def _extern_libs_key(backend: Any) -> str: + """Return a cache key fragment for extern libs (e.g. libdevice.10.bc). + + These files affect codegen but are not covered by triton_key() (Python + sources only) or backend.hash() (ptxas version and arch only). + """ + opts = backend.parse_options({}) + extern_libs = getattr(opts, "extern_libs", None) + if not extern_libs: + return "" + parts = [] + for name, path in sorted(extern_libs): + if os.path.isfile(path): + with open(path, "rb") as f: + parts.append(f"{name}-{hashlib.sha256(f.read()).hexdigest()}") + return "-".join(parts) + + +@functools.cache +def triton_hash_with_backend() -> str: + from torch._inductor.runtime.triton_compat import triton_key + + backend = triton_backend() + key = f"{triton_key()}-{backend.hash()}" + + # Hash is upper case so that it can't contain any Python keywords. + return hashlib.sha256(key.encode("utf-8")).hexdigest().upper() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_typing_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_typing_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..4c59ac453dd81aee9c9660edea5710ee86bbe185 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_typing_utils.py @@ -0,0 +1,76 @@ +"""Miscellaneous utilities to aid with typing.""" + +from collections.abc import Callable +from typing import Any, cast, Concatenate, TypeVar +from typing_extensions import ParamSpec + + +# Helper to turn Optional[T] into T when we know None either isn't +# possible or should trigger an exception. +T = TypeVar("T") + + +def not_none(obj: T | None) -> T: + if obj is None: + raise TypeError("Invariant encountered: value was None when it should not be") + return obj + + +_P = ParamSpec("_P") +_R = TypeVar("_R") +_A1 = TypeVar("_A1") + + +def copy_func_params( + source_func: Callable[_P, Any], +) -> Callable[[Callable[..., _R]], Callable[_P, _R]]: + """Cast the decorated function's call signature to the source_func's. + + Usage: + def upstream_func(a: int, b: float, *, double: bool = False) -> float: ... + @copy_func_params(upstream_func) + def enhanced(a: int, b: float, *args: Any, double: bool = False, **kwargs: Any) -> str: ... + """ + + def return_func(func: Callable[..., _R]) -> Callable[_P, _R]: + return cast(Callable[_P, _R], func) + + return return_func + + +def copy_method_params( + source_method: Callable[Concatenate[Any, _P], Any], +) -> Callable[[Callable[..., _R]], Callable[Concatenate[_A1, _P], _R]]: + """Cast the decorated *method*'s call signature to the source_method's. + Keeps the first argument type (e.g., self/cls). + """ + + def return_func(func: Callable[..., _R]) -> Callable[Concatenate[_A1, _P], _R]: + return cast(Callable[Concatenate[_A1, _P], _R], func) + + return return_func + + +# stricter variants to preserve the origin callers Return Type too. +# TODO: consider folding both these into the above variants with an optional +# parameter to control whether to copy the return type or not. +def copy_func_sig( + source_func: Callable[_P, _R], +) -> Callable[[Callable[..., _R]], Callable[_P, _R]]: + """Cast the decorated function's call signature and return type to the source_func's.""" + + def _return(func: Callable[..., _R]) -> Callable[_P, _R]: + return cast(Callable[_P, _R], func) + + return _return + + +def copy_method_sig( + source_method: Callable[Concatenate[_A1, _P], _R], +) -> Callable[[Callable[..., _R]], Callable[Concatenate[_A1, _P], _R]]: + """Cast the decorated *method*'s call signature to the source_method and return type.""" + + def _return(func: Callable[..., _R]) -> Callable[Concatenate[_A1, _P], _R]: + return cast(Callable[Concatenate[_A1, _P], _R], func) + + return _return diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_zip.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_zip.py new file mode 100644 index 0000000000000000000000000000000000000000..c4bfbcb0b9b637f82fba7c9722cd6fbc5690555c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/_zip.py @@ -0,0 +1,86 @@ +# mypy: allow-untyped-defs +import argparse +import glob +import os +from pathlib import Path +from zipfile import ZipFile + + +# Exclude some standard library modules to: +# 1. Slim down the final zipped file size +# 2. Remove functionality we don't want to support. +DENY_LIST = [ + # Interface to unix databases + "dbm", + # ncurses bindings (terminal interfaces) + "curses", + # Tcl/Tk GUI + "tkinter", + "tkinter", + # Tests for the standard library + "test", + "tests", + "idle_test", + "__phello__.foo.py", + # importlib frozen modules. These are already baked into CPython. + "_bootstrap.py", + "_bootstrap_external.py", +] + +strip_file_dir = "" + + +def remove_prefix(text, prefix): + if text.startswith(prefix): + return text[len(prefix) :] + return text + + +def write_to_zip(file_path, strip_file_path, zf, prepend_str="") -> None: + stripped_file_path = prepend_str + remove_prefix(file_path, strip_file_dir + "/") + path = Path(stripped_file_path) + if path.name in DENY_LIST: + return + zf.write(file_path, stripped_file_path) + + +def main() -> None: + global strip_file_dir + parser = argparse.ArgumentParser(description="Zip py source") + parser.add_argument("paths", nargs="*", help="Paths to zip.") + parser.add_argument( + "--install-dir", "--install_dir", help="Root directory for all output files" + ) + parser.add_argument( + "--strip-dir", + "--strip_dir", + help="The absolute directory we want to remove from zip", + ) + parser.add_argument( + "--prepend-str", + "--prepend_str", + help="A string to prepend onto all paths of a file in the zip", + default="", + ) + parser.add_argument("--zip-name", "--zip_name", help="Output zip name") + + args = parser.parse_args() + + zip_file_name = args.install_dir + "/" + args.zip_name + strip_file_dir = args.strip_dir + prepend_str = args.prepend_str + with ZipFile(zip_file_name, mode="w") as zf: + for p in sorted(args.paths): + if os.path.isdir(p): + files = glob.glob(p + "/**/*.py", recursive=True) + for file_path in sorted(files): + # strip the absolute path + write_to_zip( + file_path, strip_file_dir + "/", zf, prepend_str=prepend_str + ) + else: + write_to_zip(p, strip_file_dir + "/", zf, prepend_str=prepend_str) + + +if __name__ == "__main__": + main() # pragma: no cover diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/backcompat/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/backcompat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f6ec989be1e078ba857d30b06a91e1dc54131e4b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/backcompat/__init__.py @@ -0,0 +1,27 @@ +# mypy: allow-untyped-defs +from torch._C import ( + _get_backcompat_broadcast_warn, + _get_backcompat_keepdim_warn, + _set_backcompat_broadcast_warn, + _set_backcompat_keepdim_warn, +) + + +class Warning: + def __init__(self, setter, getter) -> None: + self.setter = setter + self.getter = getter + + def set_enabled(self, value) -> None: + self.setter(value) + + def get_enabled(self): + return self.getter() + + enabled = property(get_enabled, set_enabled) + + +broadcast_warning = Warning( + _set_backcompat_broadcast_warn, _get_backcompat_broadcast_warn +) +keepdim_warning = Warning(_set_backcompat_keepdim_warn, _get_backcompat_keepdim_warn) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/backend_registration.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/backend_registration.py new file mode 100644 index 0000000000000000000000000000000000000000..f775cc0430d3cfcbe3479937dabe175404289721 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/backend_registration.py @@ -0,0 +1,520 @@ +# mypy: allow-untyped-defs + +import torch +from torch._C import _get_privateuse1_backend_name, _rename_privateuse1_backend +from torch.overrides import handle_torch_function, has_torch_function_unary + + +__all__ = [ + "rename_privateuse1_backend", + "generate_methods_for_privateuse1_backend", +] + +# TODO: Should use `torch._C._get_privateuse1_backend_name()` to get +# renamed-backend name for `privateuse1`, but the func will cause an +# error with torch.jit.script, so we use the global variable named +# `_privateuse1_backend_name`. +_privateuse1_backend_name = "privateuseone" + + +def rename_privateuse1_backend(backend_name: str) -> None: + r""" + Rename the privateuse1 backend device to make it more convenient to use as a device name within PyTorch APIs. + + The steps are: + + (1) (In C++) implement kernels for various torch operations, and register them + to the PrivateUse1 dispatch key. + (2) (In python) call torch.utils.rename_privateuse1_backend("foo") + + You can now use "foo" as an ordinary device string in python. + + Note: this API can only be called once per process. Attempting to change + the external backend after it's already been set will result in an error. + + Note(AMP): If you want to support AMP on your device, you can register a custom backend module. + The backend must register a custom backend module with ``torch._register_device_module("foo", BackendModule)``. + BackendModule needs to have the following API's: + + (1) ``get_amp_supported_dtype() -> List[torch.dtype]`` + get the supported dtypes on your "foo" device in AMP, maybe the "foo" device supports one more dtype. + + Note(random): If you want to support to set seed for your device, BackendModule needs to have the following API's: + + (1) ``_is_in_bad_fork() -> bool`` + Return ``True`` if now it is in bad_fork, else return ``False``. + + (2) ``manual_seed_all(seed int) -> None`` + Sets the seed for generating random numbers for your devices. + + (3) ``device_count() -> int`` + Returns the number of "foo"s available. + + (4) ``get_rng_state(device: Union[int, str, torch.device] = 'foo') -> Tensor`` + Returns a list of ByteTensor representing the random number states of all devices. + + (5) ``set_rng_state(new_state: Tensor, device: Union[int, str, torch.device] = 'foo') -> None`` + Sets the random number generator state of the specified "foo" device. + + And there are some common funcs: + + (1) ``is_available() -> bool`` + Returns a bool indicating if "foo" is currently available. + + (2) ``current_device() -> int`` + Returns the index of a currently selected device. + + For more details, see https://pytorch.org/tutorials/advanced/extend_dispatcher.html#get-a-dispatch-key-for-your-backend + For an existing example, see https://github.com/bdhirsh/pytorch_open_registration_example + + Example:: + + >>> # xdoctest: +SKIP("failing") + >>> torch.utils.rename_privateuse1_backend("foo") + # This will work, assuming that you've implemented the right C++ kernels + # to implement torch.ones. + >>> a = torch.ones(2, device="foo") + + """ + _rename_privateuse1_backend(backend_name) + global _privateuse1_backend_name + _privateuse1_backend_name = backend_name + + +def _check_register_once(module, attr) -> None: + if hasattr(module, attr): + raise RuntimeError( + f"The custom device module of {module} has already been registered with {attr}" + ) + + +def _normalization_device( + custom_backend_name: str, device: int | str | torch.device | None = None +) -> int: + def _get_current_device_index(): + _get_device_index = "current_device" + if hasattr(torch, custom_backend_name) and hasattr( + getattr(torch, custom_backend_name), _get_device_index + ): + return getattr(getattr(torch, custom_backend_name), _get_device_index)() + else: + # The default device index is 0. + return 0 + + if device is None: + return _get_current_device_index() + # if isinstance(device, str), this means that the parameter passed in is in the string format "foo:0" + # convert str object to torch.device object, and then process it uniformly + elif isinstance(device, str): + device = torch.device(device) + + # variable device can only be torch.device type or int type + if isinstance(device, torch.device): + if device.type != custom_backend_name: + raise RuntimeError(f"Invalid device, must be {custom_backend_name} device") + elif device.index is None: + device_idx = _get_current_device_index() + else: + device_idx = device.index + # if isinstance(device, int), we can take the index number directly + else: + device_idx = device + return device_idx + + +def _generate_tensor_methods_for_privateuse1_backend(custom_backend_name: str) -> None: + @property # type: ignore[misc] + def wrap_tensor_backend(self: torch.Tensor) -> bool: + if has_torch_function_unary(self): + # TODO mypy doesn't support @property, see: https://github.com/python/mypy/issues/6185 + return handle_torch_function(wrap_tensor_backend.__get__, (self,), self) # type: ignore[attr-defined] + return self.device.type == custom_backend_name + + _check_register_once(torch.Tensor, f"is_{custom_backend_name}") + wrap_tensor_backend.fget.__name__ = f"is_{custom_backend_name}" # type: ignore[attr-defined] + setattr(torch.Tensor, f"is_{custom_backend_name}", wrap_tensor_backend) + + def wrap_tensor_to( + self: torch.Tensor, + device: int | torch.device | None = None, + non_blocking=False, + **kwargs, + ) -> torch.Tensor: + r"""Perform Tensor device conversion. Call the to operator implementation. + + .. note:: + If the ``self`` Tensor already + has the correct :class:`torch.device`, then ``self`` is returned. + Otherwise, the returned tensor is a copy of ``self`` with the desired :class:`torch.device`. + + Args: + device (int, optional): if specified, all parameters will be copied to that device + non_blocking (bool): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. Otherwise, + the argument has no effect. + **kwargs (dict): For compatibility, may contain the key ``memory_format`` argument. + """ + if has_torch_function_unary(self): + return handle_torch_function( + wrap_tensor_to, + (self,), + self, + device=device, + non_blocking=False, + **kwargs, + ) + device_idx = _normalization_device(custom_backend_name, device) + return self.to( + device=torch.device(f"{custom_backend_name}:{device_idx}"), + non_blocking=non_blocking, + **kwargs, + ) + + _check_register_once(torch.Tensor, custom_backend_name) + wrap_tensor_to.__name__ = custom_backend_name + setattr(torch.Tensor, custom_backend_name, wrap_tensor_to) + + +def _generate_module_methods_for_privateuse1_backend(custom_backend_name: str) -> None: + # Generate Module attributes and methods depends on Tensor methods, + # so we need to check whether Tensor methods is already registered. + if not hasattr(torch.Tensor, custom_backend_name): + raise RuntimeError( + f"Can not automatically generate {custom_backend_name}() method for torch.nn.Module." + f"Because torch.Tensor doesn't has the method {custom_backend_name}()." + f"For this error, you can try setting for_tensor=True." + ) + + def wrap_module_to( + # pyrefly: ignore [invalid-type-var] + self: torch.nn.modules.module.T, + device: int | torch.device | None = None, + ) -> torch.nn.modules.module.T: # pyrefly: ignore [invalid-type-var] + r"""Move all model parameters and buffers to the custom device. + + This also makes associated parameters and buffers different objects. So + it should be called before constructing optimizer if the module will + live on device while being optimized. + + .. note:: + This method modifies the module in-place. + + Args: + device (int, optional): if specified, all parameters will be copied to that device + """ + # pyrefly: ignore [missing-attribute] + return self._apply(lambda t: getattr(t, custom_backend_name)(device)) + + _check_register_once(torch.nn.Module, custom_backend_name) + setattr(torch.nn.Module, custom_backend_name, wrap_module_to) + + +def _generate_packed_sequence_methods_for_privateuse1_backend( + custom_backend_name: str, +) -> None: + # Generate PackedSequence Module attributes and methods depends on Tensor methods, + # so we need to check whether Tensor methods is already registered. + if not hasattr(torch.Tensor, f"is_{custom_backend_name}") or not hasattr( + torch.Tensor, custom_backend_name + ): + raise RuntimeError( + f"Can not automatically generate is_{custom_backend_name}() or " + f"{custom_backend_name}() method for torch.nn.utils.rnn.PackedSequence." + f"Because torch.Tensor doesn't has the method is_{custom_backend_name}()" + f"or {custom_backend_name}()." + f"For this error, you can try setting for_tensor=True." + ) + + @property # type: ignore[misc] + def wrap_tensor_backend(self: torch.nn.utils.rnn.PackedSequence) -> bool: + return self.data.device.type == custom_backend_name + + _check_register_once(torch.nn.utils.rnn.PackedSequence, f"is_{custom_backend_name}") + setattr( + torch.nn.utils.rnn.PackedSequence, + f"is_{custom_backend_name}", + wrap_tensor_backend, + ) + + def wrap_module_to( + self: torch.nn.utils.rnn.PackedSequence, *args, **kwargs + ) -> torch.nn.utils.rnn.PackedSequence: + r"""Move all model parameters and buffers to the custom device. + + This also makes associated parameters and buffers different objects. So + it should be called before constructing optimizer if the module will + live on device while being optimized. + + .. note:: + This method modifies the module in-place. + + Args: + device (int, optional): if specified, all parameters will be copied to that device + """ + ex = torch.tensor((), dtype=self.data.dtype, device=self.data.device).to( + *args, + **kwargs, + ) + if ex.device.type == custom_backend_name: + return self.to(*args, **kwargs) + kwargs.update({"device": custom_backend_name}) + + return self.to(*args, **kwargs) + + _check_register_once(torch.nn.utils.rnn.PackedSequence, custom_backend_name) + setattr(torch.nn.utils.rnn.PackedSequence, custom_backend_name, wrap_module_to) + + +def _generate_storage_methods_for_privateuse1_backend( + custom_backend_name: str, unsupported_dtype: list[torch.dtype] | None = None +) -> None: + # Attribute is registered in the _StorageBase class + # and UntypedStorage obtains through inheritance. + @property # type: ignore[misc] + def wrap_storage_backend(self: torch.storage._StorageBase) -> bool: + r"""Return the internal :class:`torch.UntypedStorage`.""" + return self.device.type == custom_backend_name + + _check_register_once(torch.storage._StorageBase, f"is_{custom_backend_name}") + setattr( + torch.storage._StorageBase, f"is_{custom_backend_name}", wrap_storage_backend + ) + + def wrap_storage_to(self, device=None, non_blocking=False): + r"""Return a copy of this object in custom device memory. + + If this object is already in device memory and on the correct device, then + no copy is performed and the original object is returned. + + Args: + device (int): The destination device id. Defaults to the current device. + non_blocking (bool): If ``True`` and the source is in pinned memory, + the copy will be asynchronous with respect to the host. Otherwise, + the argument has no effect. + """ + # There should be a judgment related to storage device and a judgment related to storage type, + # but it depends on the extended function, so this part is temporarily omitted in the automatic generation. + device_idx = _normalization_device(custom_backend_name, device) + + if getattr(self, f"is_{custom_backend_name}"): + # storage has already on expected device. + if self.get_device() == device_idx: + return self + # For sparse storage, custom need to extend the implementation by themselves. + if self.is_sparse: + raise RuntimeError( + f"Can not support a sparse storage move to {custom_backend_name} backend" + ) + # create untyped_storage and copy data + untyped_storage = torch.UntypedStorage( + self.size(), device=torch.device(f"{custom_backend_name}:{device_idx}") + ) + untyped_storage.copy_(self, non_blocking) + return untyped_storage + + _check_register_once(torch.storage._StorageBase, custom_backend_name) + setattr(torch.storage._StorageBase, custom_backend_name, wrap_storage_to) + + # Register the corresponding attribute for the TypedStorage class. + # When the TypedStorage class is removed, the registration is also removed. + + @property # type: ignore[misc] + def wrap_typed_storage_backend(self: torch.storage.TypedStorage) -> bool: + torch.storage._warn_typed_storage_removal() + return self._untyped_storage.device.type == custom_backend_name + + _check_register_once(torch.TypedStorage, f"is_{custom_backend_name}") + setattr( + torch.storage.TypedStorage, + f"is_{custom_backend_name}", + wrap_typed_storage_backend, + ) + + def wrap_typed_storage_to( + self: torch.storage.TypedStorage, device=None, non_blocking=False, **kwargs + ) -> torch.storage.TypedStorage: + torch.storage._warn_typed_storage_removal() + if unsupported_dtype and self.dtype in unsupported_dtype: + raise RuntimeError( + f"Cannot create {custom_backend_name} storage " + f"as {self.dtype} dtype is not supported by this backend" + ) + custom_backend_storage: torch.UntypedStorage = getattr( + self._untyped_storage, custom_backend_name + )(device, non_blocking, **kwargs) + return self._new_wrapped_storage(custom_backend_storage) + + _check_register_once(torch.TypedStorage, custom_backend_name) + setattr(torch.TypedStorage, custom_backend_name, wrap_typed_storage_to) + + +def generate_methods_for_privateuse1_backend( + for_tensor: bool = True, + for_module: bool = True, + for_packed_sequence: bool = True, + for_storage: bool = False, + unsupported_dtype: list[torch.dtype] | None = None, +) -> None: + r""" + Automatically generate attributes and methods for the custom backend after rename privateuse1 backend. + + In the default scenario, storage-related methods will not be generated automatically. + + When you implement kernels for various torch operations, and register them to the PrivateUse1 dispatch key. + And call the function torch.rename_privateuse1_backend("foo") to rename your backend name. + At this point, you can easily register specific methods and attributes by calling this function. + Just like torch.Tensor.foo(), torch.Tensor.is_foo, torch.Storage.foo(), torch.Storage.is_foo. + + Note: We recommend you use generic functions (check devices are equal or to(device=)). + We provide these methods for convenience only and they will be "monkey patched" onto the objects + and so will not be properly typed. For Storage methods generate, if you need to support sparse data storage, + you need to extend the implementation yourself. + + Args: + for_tensor (bool): whether register related methods for torch.Tensor class. + for_module (bool): whether register related methods for torch.nn.Module class. + for_storage (bool): whether register related methods for torch.Storage class. + unsupported_dtype (List[torch.dtype]): takes effect only when the storage method needs to be generated, + indicating that the storage does not support the torch.dtype type. + + Example:: + + >>> # xdoctest: +SKIP("failing") + >>> torch.utils.rename_privateuse1_backend("foo") + >>> torch.utils.generate_methods_for_privateuse1_backend() + # Then automatically generate backend-related attributes and methods. + >>> a = torch.tensor(2).foo() + >>> a.is_foo + >>> hasattr(torch.nn.Module, 'foo') + """ + custom_backend_name = _get_privateuse1_backend_name() + + if for_tensor: + _generate_tensor_methods_for_privateuse1_backend(custom_backend_name) + + if for_module: + _generate_module_methods_for_privateuse1_backend(custom_backend_name) + + if for_storage: + _generate_storage_methods_for_privateuse1_backend( + custom_backend_name, unsupported_dtype + ) + + if for_packed_sequence: + _generate_packed_sequence_methods_for_privateuse1_backend(custom_backend_name) + + +def _get_custom_mod_func(func_name: str): + r""" + Return the func named `func_name` defined in custom device module. If not defined, + return `None`. And the func is registered with `torch.utils.rename_privateuse1_backend('foo')` + and `torch._register_device_module('foo', BackendModule)`. + If the custom device module or the func is not defined, it will give warning or error message. + Args: + func_name (str): return the callable func named func_name defined in custom device module. + Example:: + class DummyfooModule: + @staticmethod + def is_available(): + return True + @staticmethod + def func_name(*args, **kwargs): + .... + torch.utils.rename_privateuse1_backend("foo") + torch._register_device_module("foo", DummyfooModule) + foo_is_available_func = torch.utils.backend_registration._get_custom_mod_func("is_available") + if foo_is_available_func: + foo_is_available = foo_is_available_func() + func_ = torch.utils.backend_registration._get_custom_mod_func("func_name") + if func_: + result = func_(*args, **kwargs) + Attention: This function is not meant to be used directly by users, which is why + it is marked as private. It is a convenience function for backend implementers to + more easily call the hooks into their backend extensions. + """ + if not isinstance(func_name, str): + raise AssertionError(f"func_name must be `str`, but got `{type(func_name)}`.") + backend_name = _get_privateuse1_backend_name() + custom_device_mod = getattr(torch, backend_name, None) + function = getattr(custom_device_mod, func_name, None) + if custom_device_mod is None or function is None: + message = f"Try to call torch.{backend_name}.{func_name}. The backend must register a custom backend " + message += f"module with `torch._register_device_module('{backend_name}', BackendModule)`. And " + message += f"BackendModule needs to have the following API's:\n `{func_name}(*args, **kwargs)`. \n" + raise RuntimeError(message) + return function + + +class _DummyBackendModule: + def is_initialized(self) -> bool: + return True + + def is_available(self) -> bool: + return True + + def current_device(self) -> int: + return 0 + + def _is_in_bad_fork(self) -> bool: + return False + + def manual_seed_all(self, seed: int) -> None: + pass + + def device_count(self) -> int: + return 1 + + +class _DummyPrivateUse1Hook(torch._C._acc.PrivateUse1Hooks): + def is_available(self) -> bool: + return True + + def has_primary_context(self, dev_id) -> bool: + return True + + def is_built(self) -> bool: + return True + + +class _DummyDeviceGuard(torch._C._acc.DeviceGuard): + def type_(self): + return torch._C._autograd.DeviceType.PrivateUse1 + + +def _setup_privateuseone_for_python_backend( + rename=None, backend_module=None, hook=None, device_guard=None +) -> None: + """This function will prepare the PrivateUse1 dispatch key to be used as a python backend. + + WARNING: this API is experimental and might change without notice. + + Formally, this registers things that Pytorch expects a registered backend + in C++ to have: including device guards, hooks, and backend modules and what not. + + after this call, one can use `torch.library` to write Ops for this dispatch key + and expect it to behave like a backend registered in C++. + + See the unit test at test/test_privateuseone_python_backend.py for more details. + + Args: + rename: str | None, if passed in, we will rename privateuseone backend to + the name given. + backend_module: object | None, if passed in None, we will use DummyBackendModule + hook: object | None, if passed in None, we will use DummyPrivateUse1Hook + device_guard: object | None, if passed in None, we will use DummyDeviceGuard + """ + # NOTE: the ordering of which these functions are called is important. + if rename is not None: + torch.utils.rename_privateuse1_backend(rename) + else: + rename = "privateuseone" + torch.utils.generate_methods_for_privateuse1_backend() + if backend_module is None: + backend_module = _DummyBackendModule() + if hook is None: + hook = _DummyPrivateUse1Hook() + if device_guard is None: + device_guard = _DummyDeviceGuard() + torch._register_device_module(rename, backend_module) + torch._C._acc.register_python_privateuseone_hook(hook) + torch._C._acc.register_python_privateuseone_device_guard(device_guard) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..9e814aaf4671ca35484c43bc38677849d02a81ec --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/__init__.py @@ -0,0 +1,6 @@ +from torch.utils.benchmark.utils.common import * # noqa: F403 +from torch.utils.benchmark.utils.timer import * # noqa: F403 +from torch.utils.benchmark.utils.compare import * # noqa: F403 +from torch.utils.benchmark.utils.fuzzer import * # noqa: F403 +from torch.utils.benchmark.utils.valgrind_wrapper.timer_interface import * # noqa: F403 +from torch.utils.benchmark.utils.sparse_fuzzer import * # noqa: F403 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/compare.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/compare.py new file mode 100644 index 0000000000000000000000000000000000000000..1c266e7cf9a6e604c94dfb28f19f31f1649220f4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/compare.py @@ -0,0 +1,99 @@ +# mypy: allow-untyped-defs +"""Example of Timer and Compare APIs: + +$ python -m examples.compare +""" + +import pickle +import sys +import time + +import torch + +import torch.utils.benchmark as benchmark_utils + + +class FauxTorch: + """Emulate different versions of pytorch. + + In normal circumstances this would be done with multiple processes + writing serialized measurements, but this simplifies that model to + make the example clearer. + """ + def __init__(self, real_torch, extra_ns_per_element) -> None: + self._real_torch = real_torch + self._extra_ns_per_element = extra_ns_per_element + + def extra_overhead(self, result): + # time.sleep has a ~65 us overhead, so only fake a + # per-element overhead if numel is large enough. + numel = int(result.numel()) + if numel > 5000: + time.sleep(numel * self._extra_ns_per_element * 1e-9) + return result + + def add(self, *args, **kwargs): + return self.extra_overhead(self._real_torch.add(*args, **kwargs)) + + def mul(self, *args, **kwargs): + return self.extra_overhead(self._real_torch.mul(*args, **kwargs)) + + def cat(self, *args, **kwargs): + return self.extra_overhead(self._real_torch.cat(*args, **kwargs)) + + def matmul(self, *args, **kwargs): + return self.extra_overhead(self._real_torch.matmul(*args, **kwargs)) + + +def main() -> None: + tasks = [ + ("add", "add", "torch.add(x, y)"), + ("add", "add (extra +0)", "torch.add(x, y + zero)"), + ] + + serialized_results = [] + repeats = 2 + timers = [ + benchmark_utils.Timer( + stmt=stmt, + globals={ + "torch": torch if branch == "master" else FauxTorch(torch, overhead_ns), + "x": torch.ones((size, 4)), + "y": torch.ones((1, 4)), + "zero": torch.zeros(()), + }, + label=label, + sub_label=sub_label, + description=f"size: {size}", + env=branch, + num_threads=num_threads, + ) + for branch, overhead_ns in [("master", None), ("my_branch", 1), ("severe_regression", 5)] + for label, sub_label, stmt in tasks + for size in [1, 10, 100, 1000, 10000, 50000] + for num_threads in [1, 4] + ] + + for i, timer in enumerate(timers * repeats): + serialized_results.append(pickle.dumps( + timer.blocked_autorange(min_run_time=0.05) + )) + print(f"\r{i + 1} / {len(timers) * repeats}", end="") + sys.stdout.flush() + print() + + comparison = benchmark_utils.Compare([ + pickle.loads(i) for i in serialized_results + ]) + + print("== Unformatted " + "=" * 80 + "\n" + "/" * 95 + "\n") + comparison.print() + + print("== Formatted " + "=" * 80 + "\n" + "/" * 93 + "\n") + comparison.trim_significant_figures() + comparison.colorize() + comparison.print() + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/fuzzer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/fuzzer.py new file mode 100644 index 0000000000000000000000000000000000000000..80a4e733928d8b059919d847da1b461d55dd7402 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/fuzzer.py @@ -0,0 +1,86 @@ +# mypy: allow-untyped-defs +"""Example of the Timer and Fuzzer APIs: + +$ python -m examples.fuzzer +""" + +import sys + +import torch.utils.benchmark as benchmark_utils + + +def main() -> None: + add_fuzzer = benchmark_utils.Fuzzer( + parameters=[ + [ + benchmark_utils.FuzzedParameter( + name=f"k{i}", + minval=16, + maxval=16 * 1024, + distribution="loguniform", + ) for i in range(3) + ], + benchmark_utils.FuzzedParameter( + name="d", + distribution={2: 0.6, 3: 0.4}, + ), + ], + tensors=[ + [ + benchmark_utils.FuzzedTensor( + name=name, + size=("k0", "k1", "k2"), + dim_parameter="d", + probability_contiguous=0.75, + min_elements=64 * 1024, + max_elements=128 * 1024, + ) for name in ("x", "y") + ], + ], + seed=0, + ) + + n = 250 + measurements = [] + for i, (tensors, tensor_properties, _) in enumerate(add_fuzzer.take(n=n)): + x, x_order = tensors["x"], str(tensor_properties["x"]["order"]) + y, y_order = tensors["y"], str(tensor_properties["y"]["order"]) + shape = ", ".join(tuple(f'{i:>4}' for i in x.shape)) + + description = "".join([ + f"{x.numel():>7} | {shape:<16} | ", + f"{'contiguous' if x.is_contiguous() else x_order:<12} | ", + f"{'contiguous' if y.is_contiguous() else y_order:<12} | ", + ]) + + timer = benchmark_utils.Timer( + stmt="x + y", + globals=tensors, + description=description, + ) + + measurements.append(timer.blocked_autorange(min_run_time=0.1)) + measurements[-1].metadata = {"numel": x.numel()} + print(f"\r{i + 1} / {n}", end="") + sys.stdout.flush() + print() + + # More string munging to make pretty output. + print(f"Average attempts per valid config: {1. / (1. - add_fuzzer.rejection_rate):.1f}") + + def time_fn(m): + return m.median / m.metadata["numel"] + measurements.sort(key=time_fn) + + template = f"{{:>6}}{' ' * 19}Size Shape{' ' * 13}X order Y order\n{'-' * 80}" + print(template.format("Best:")) + for m in measurements[:15]: + print(f"{time_fn(m) * 1e9:>4.1f} ns / element {m.description}") + + print("\n" + template.format("Worst:")) + for m in measurements[-15:]: + print(f"{time_fn(m) * 1e9:>4.1f} ns / element {m.description}") + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/op_benchmark.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/op_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..f65599ee18a4f2c4a0d35b514c8f87725affae01 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/op_benchmark.py @@ -0,0 +1,107 @@ +# mypy: allow-untyped-defs +"""Example use of Timer and op fuzzers to measure kernel performance. + +$ python -m examples.op_benchmark +""" + +import numpy as np +import torch + +from torch.utils.benchmark import Timer +from torch.utils.benchmark.op_fuzzers.binary import BinaryOpFuzzer +from torch.utils.benchmark.op_fuzzers.unary import UnaryOpFuzzer +import operator + + +_MEASURE_TIME = 1.0 + + +def assert_dicts_equal(dict_0, dict_1) -> None: + """Builtin dict comparison will not compare numpy arrays. + e.g. + x = {"a": np.ones((2, 1))} + x == x # Raises ValueError + """ + if set(dict_0.keys()) != set(dict_0.keys()): + raise AssertionError("dicts must have the same keys") + if all(np.all(v != dict_1[k]) for k, v in dict_0.items() if k != "dtype"): + raise AssertionError("dict values differ for keys other than 'dtype'") + + +def run(n, stmt, fuzzer_cls) -> None: + float_iter = fuzzer_cls(seed=0, dtype=torch.float32).take(n) + int_iter = fuzzer_cls(seed=0, dtype=torch.int32).take(n) + raw_results = [] + for i, (float_values, int_values) in enumerate(zip(float_iter, int_iter, strict=True)): + float_tensors, float_tensor_params, float_params = float_values + int_tensors, int_tensor_params, int_params = int_values + + # This benchmark assumes that the two fuzzers generate identically + # sized and strided Tensors, since the same seed is used. + assert_dicts_equal(float_params, int_params) + assert_dicts_equal(float_tensor_params["x"], int_tensor_params["x"]) + + float_measurement, int_measurement = ( + Timer( + stmt, + globals=tensors, + ).blocked_autorange(min_run_time=_MEASURE_TIME) + for tensors in (float_tensors, int_tensors) + ) + + descriptions = [] + for name in float_tensors: + shape_str = "(" + ", ".join([ + f"2 ** {int(np.log2(i))}" + if 2 ** int(np.log2(i)) == i and i > 1 + else str(i) + for i in float_tensors[name].shape + ]) + ")" + order = float_tensor_params[name]["order"] + order_str = ("" if all(order == np.arange(len(order))) else str(tuple(order))) + steps = float_tensor_params[name]["steps"] + steps_str = str(steps) if sum(steps) > len(steps) else "" + descriptions.append((name, shape_str, order_str, steps_str)) + raw_results.append((float_measurement, int_measurement, descriptions)) + + print(f"\r{i + 1} / {n}", end="") + print() + + parsed_results, name_len, shape_len, order_len, steps_len = [], 0, 0, 0, 0 + for float_measurement, int_measurement, descriptions in raw_results: + t_float = float_measurement.median * 1e6 + t_int = int_measurement.median * 1e6 + rel_diff = abs(t_float - t_int) / (t_float + t_int) * 2 + parsed_results.append((t_float, t_int, rel_diff, descriptions)) + for name, shape, order, steps in descriptions: + name_len = max(name_len, len(name)) + shape_len = max(shape_len, len(shape)) + order_len = max(order_len, len(order)) + steps_len = max(steps_len, len(steps)) + + parsed_results.sort(key=operator.itemgetter(2)) + + print(f"stmt: {stmt}") + print(f" diff faster{'':>17}{' ' * name_len} ", end="") + print(f"{'shape'.ljust(shape_len)}{'':>16}{'order'.ljust(order_len)}", end="") + print(f" steps\n{'-' * 100}") + for results, spacer in [(parsed_results[:10], "..."), (parsed_results[-10:], "")]: + for t_float, t_int, rel_diff, descriptions in results: + time_str = [f"{rel_diff * 100:>4.1f}% {'int' if t_int < t_float else 'float':<20}"] + time_str.extend(["".ljust(len(time_str[0])) for _ in descriptions[:-1]]) + for t_str, (name, shape, order, steps) in zip(time_str, descriptions, strict=True): + name = f"{name}:".ljust(name_len + 1) + shape = shape.ljust(shape_len + 10) + order = order.ljust(order_len) + print(f"{t_str} {name} {shape}| {order} | {steps}") + print(spacer) + + +def main() -> None: + run(n=100, stmt="torch.median(x, dim=0)", fuzzer_cls=UnaryOpFuzzer) + run(n=100, stmt="torch.square(x)", fuzzer_cls=UnaryOpFuzzer) + run(n=100, stmt="x + y", fuzzer_cls=BinaryOpFuzzer) + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/simple_timeit.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/simple_timeit.py new file mode 100644 index 0000000000000000000000000000000000000000..8137d4d8791975b46b1314c2f3a05ed048dbdcd3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/simple_timeit.py @@ -0,0 +1,25 @@ +"""Trivial use of Timer API: + +$ python -m examples.simple_timeit +""" + +import torch + +import torch.utils.benchmark as benchmark_utils + + +def main() -> None: + timer = benchmark_utils.Timer( + stmt="x + y", + globals={"x": torch.ones((4, 8)), "y": torch.ones((1, 8))}, + label="Broadcasting add (4x8)", + ) + + for i in range(3): + print(f"Run: {i}\n{'-' * 40}") + print(f"timeit:\n{timer.timeit(10000)}\n") + print(f"autorange:\n{timer.blocked_autorange()}\n\n") + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/spectral_ops_fuzz_test.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/spectral_ops_fuzz_test.py new file mode 100644 index 0000000000000000000000000000000000000000..81a33c34bc8229a44838ea93c29af34895061c53 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/examples/spectral_ops_fuzz_test.py @@ -0,0 +1,114 @@ +# mypy: allow-untyped-defs +"""Microbenchmarks for the torch.fft module""" +from argparse import ArgumentParser +from collections import namedtuple +from collections.abc import Iterable + +import torch +import torch.fft +from torch.utils import benchmark +from torch.utils.benchmark.op_fuzzers.spectral import SpectralOpFuzzer + + +def _dim_options(ndim): + if ndim == 1: + return [None] + elif ndim == 2: + return [0, 1, None] + elif ndim == 3: + return [0, 1, 2, (0, 1), (0, 2), None] + raise ValueError(f"Expected ndim in range 1-3, got {ndim}") + + +def run_benchmark(name: str, function: object, dtype: torch.dtype, seed: int, device: str, samples: int, + probability_regular: float): + cuda = device == 'cuda' + spectral_fuzzer = SpectralOpFuzzer(seed=seed, dtype=dtype, cuda=cuda, + probability_regular=probability_regular) + results = [] + for tensors, tensor_params, params in spectral_fuzzer.take(samples): + shape = [params['k0'], params['k1'], params['k2']][:params['ndim']] + str_shape = ' x '.join([f"{s:<4}" for s in shape]) + sub_label = f"{str_shape} {'' if tensor_params['x']['is_contiguous'] else '(discontiguous)'}" + for dim in _dim_options(params['ndim']): + for nthreads in (1, 4, 16) if not cuda else (1,): + measurement = benchmark.Timer( + stmt='func(x, dim=dim)', + globals={'func': function, 'x': tensors['x'], 'dim': dim}, + label=f"{name}_{device}", + sub_label=sub_label, + description=f"dim={dim}", + num_threads=nthreads, + ).blocked_autorange(min_run_time=1) + measurement.metadata = { + 'name': name, + 'device': device, + 'dim': dim, + 'shape': shape, + } + measurement.metadata.update(tensor_params['x']) + results.append(measurement) + return results + + +Benchmark = namedtuple('Benchmark', ['name', 'function', 'dtype']) +BENCHMARKS = [ + Benchmark('fft_real', torch.fft.fftn, torch.float32), + Benchmark('fft_complex', torch.fft.fftn, torch.complex64), + Benchmark('ifft', torch.fft.ifftn, torch.complex64), + Benchmark('rfft', torch.fft.rfftn, torch.float32), + Benchmark('irfft', torch.fft.irfftn, torch.complex64), +] +BENCHMARK_MAP = {b.name: b for b in BENCHMARKS} +BENCHMARK_NAMES = [b.name for b in BENCHMARKS] +DEVICE_NAMES = ['cpu', 'cuda'] + +def _output_csv(file, results) -> None: + file.write('benchmark,device,num_threads,numel,shape,contiguous,dim,mean (us),median (us),iqr (us)\n') + for measurement in results: + metadata = measurement.metadata + device, dim, shape, name, numel, contiguous = ( + metadata['device'], metadata['dim'], metadata['shape'], + metadata['name'], metadata['numel'], metadata['is_contiguous']) + + if isinstance(dim, Iterable): + dim_str = '-'.join(str(d) for d in dim) + else: + dim_str = str(dim) + shape_str = 'x'.join(str(s) for s in shape) + + print(name, device, measurement.task_spec.num_threads, numel, shape_str, contiguous, dim_str, # type: ignore[possibly-undefined] + measurement.mean * 1e6, measurement.median * 1e6, measurement.iqr * 1e6, + sep=',', file=file) + + +if __name__ == '__main__': + parser = ArgumentParser(description=__doc__) + parser.add_argument('--device', type=str, choices=DEVICE_NAMES, nargs='+', default=DEVICE_NAMES) + parser.add_argument('--bench', type=str, choices=BENCHMARK_NAMES, nargs='+', default=BENCHMARK_NAMES) + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--samples', type=int, default=10) + parser.add_argument('--probability-regular', '--probability_regular', type=float, default=1.0) + parser.add_argument('-o', '--output', type=str) + args = parser.parse_args() + + num_benchmarks = len(args.device) * len(args.bench) + i = 0 + results = [] + for device in args.device: + for bench in (BENCHMARK_MAP[b] for b in args.bench): + results += run_benchmark( + name=bench.name, function=bench.function, dtype=bench.dtype, + seed=args.seed, device=device, samples=args.samples, + probability_regular=args.probability_regular) + i += 1 + print(f'Completed {bench.name} benchmark on {device} ({i} of {num_benchmarks})') + + if args.output is not None: + with open(args.output, 'w') as f: + _output_csv(f, results) + + compare = benchmark.Compare(results) + compare.trim_significant_figures() + compare.colorize() + compare.print() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/binary.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/binary.py new file mode 100644 index 0000000000000000000000000000000000000000..e53c310111bec8166e6090f351e39153dbe400aa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/binary.py @@ -0,0 +1,107 @@ +# mypy: allow-untyped-defs +import numpy as np +import torch + +from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedTensor + + +_MIN_DIM_SIZE = 16 +_MAX_DIM_SIZE = 16 * 1024 ** 2 +_POW_TWO_SIZES = tuple(2 ** i for i in range( + int(np.log2(_MIN_DIM_SIZE)), + int(np.log2(_MAX_DIM_SIZE)) + 1, +)) + + +class BinaryOpFuzzer(Fuzzer): + def __init__(self, seed, dtype=torch.float32, cuda=False) -> None: + super().__init__( + parameters=[ + # Dimensionality of x and y. (e.g. 1D, 2D, or 3D.) + FuzzedParameter("dim", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True), + + # Shapes for `x` and `y`. + # It is important to test all shapes, however + # powers of two are especially important and therefore + # warrant special attention. This is done by generating + # both a value drawn from all integers between the min and + # max allowed values, and another from only the powers of two + # (both distributions are loguniform) and then randomly + # selecting between the two. + # Moreover, `y` will occasionally have singleton + # dimensions in order to test broadcasting. + [ + FuzzedParameter( + name=f"k_any_{i}", + minval=_MIN_DIM_SIZE, + maxval=_MAX_DIM_SIZE, + distribution="loguniform", + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k_pow2_{i}", + distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES} + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k{i}", + distribution={ + ParameterAlias(f"k_any_{i}"): 0.8, + ParameterAlias(f"k_pow2_{i}"): 0.2, + }, + strict=True, + ) for i in range(3) + ], + + [ + FuzzedParameter( + name=f"y_k{i}", + distribution={ + ParameterAlias(f"k{i}"): 0.8, + 1: 0.2, + }, + strict=True, + ) for i in range(3) + ], + + # Steps for `x` and `y`. (Benchmarks strided memory access.) + [ + FuzzedParameter( + name=f"{name}_step_{i}", + distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04}, + ) + for i in range(3) + for name in ("x", "y") + ], + + # Repeatable entropy for downstream applications. + FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"), + ], + tensors=[ + FuzzedTensor( + name="x", + size=("k0", "k1", "k2"), + steps=("x_step_0", "x_step_1", "x_step_2"), + probability_contiguous=0.75, + min_elements=4 * 1024, + max_elements=32 * 1024 ** 2, + max_allocation_bytes=2 * 1024**3, # 2 GB + dim_parameter="dim", + dtype=dtype, + cuda=cuda, + ), + FuzzedTensor( + name="y", + size=("y_k0", "y_k1", "y_k2"), + steps=("x_step_0", "x_step_1", "x_step_2"), + probability_contiguous=0.75, + max_allocation_bytes=2 * 1024**3, # 2 GB + dim_parameter="dim", + dtype=dtype, + cuda=cuda, + ), + ], + seed=seed, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/sparse_binary.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/sparse_binary.py new file mode 100644 index 0000000000000000000000000000000000000000..8e6269464e0d53d2c3c51ed5406d7c88598fec79 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/sparse_binary.py @@ -0,0 +1,107 @@ +# mypy: allow-untyped-defs +import numpy as np +import torch + +from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedSparseTensor + + +_MIN_DIM_SIZE = 16 +_MAX_DIM_SIZE = 16 * 1024 ** 2 +_POW_TWO_SIZES = tuple(2 ** i for i in range( + int(np.log2(_MIN_DIM_SIZE)), + int(np.log2(_MAX_DIM_SIZE)) + 1, +)) + + +class BinaryOpSparseFuzzer(Fuzzer): + def __init__(self, seed, dtype=torch.float32, cuda=False) -> None: + super().__init__( + parameters=[ + # Dimensionality of x and y. (e.g. 1D, 2D, or 3D.) + FuzzedParameter("dim_parameter", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True), + FuzzedParameter( + name="sparse_dim", + distribution={1: 0.4, 2: 0.4, 3: 0.2}, + strict=True + ), + # Shapes for `x` and `y`. + # It is important to test all shapes, however + # powers of two are especially important and therefore + # warrant special attention. This is done by generating + # both a value drawn from all integers between the min and + # max allowed values, and another from only the powers of two + # (both distributions are loguniform) and then randomly + # selecting between the two. + # Moreover, `y` will occasionally have singleton + # dimensions in order to test broadcasting. + [ + FuzzedParameter( + name=f"k_any_{i}", + minval=_MIN_DIM_SIZE, + maxval=_MAX_DIM_SIZE, + distribution="loguniform", + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k_pow2_{i}", + distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES} + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k{i}", + distribution={ + ParameterAlias(f"k_any_{i}"): 0.8, + ParameterAlias(f"k_pow2_{i}"): 0.2, + }, + strict=True, + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"y_k{i}", + distribution={ + ParameterAlias(f"k{i}"): 1.0}, + strict=True, + ) for i in range(3) + ], + FuzzedParameter( + name="density", + distribution={0.1: 0.4, 0.05: 0.3, 0.01: 0.3}, + ), + FuzzedParameter( + name="coalesced", + distribution={True: 0.5, False: 0.5}, + ), + # Repeatable entropy for downstream applications. + FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"), + ], + tensors=[ + FuzzedSparseTensor( + name="x", + size=("k0", "k1", "k2"), + dim_parameter="dim_parameter", + sparse_dim="sparse_dim", + density="density", + coalesced="coalesced", + min_elements=4 * 1024, + max_elements=32 * 1024 ** 2, + dtype=dtype, + cuda=cuda, + ), + FuzzedSparseTensor( + name="y", + size=("y_k0", "y_k1", "y_k2"), + dim_parameter="dim_parameter", + sparse_dim="sparse_dim", + density="density", + coalesced="coalesced", + min_elements=4 * 1024, + max_elements=32 * 1024 ** 2, + dtype=dtype, + cuda=cuda, + ), + ], + seed=seed, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/sparse_unary.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/sparse_unary.py new file mode 100644 index 0000000000000000000000000000000000000000..18921becd078cb3140a1705078dd57f4a597a2ec --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/sparse_unary.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import torch + +if TYPE_CHECKING: + from torch.types import _dtype + +from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedSparseTensor + +__all__ = ["UnaryOpSparseFuzzer"] + +_MIN_DIM_SIZE = 16 +_MAX_DIM_SIZE = 16 * 1024 ** 2 +_POW_TWO_SIZES = tuple(2 ** i for i in range( + int(np.log2(_MIN_DIM_SIZE)), + int(np.log2(_MAX_DIM_SIZE)) + 1, +)) + +class UnaryOpSparseFuzzer(Fuzzer): + def __init__(self, seed: int | None, dtype: _dtype | None = None, cuda: bool = False) -> None: + if dtype is None: + dtype = getattr(torch, 'float32', None) + super().__init__( + parameters=[ + # Sparse dim parameter of x. (e.g. 1D, 2D, or 3D.) + FuzzedParameter("dim_parameter", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True), + FuzzedParameter( + name="sparse_dim", + distribution={1: 0.4, 2: 0.4, 3: 0.2}, + strict=True + ), + # Shapes for `x`. + # It is important to test all shapes, however + # powers of two are especially important and therefore + # warrant special attention. This is done by generating + # both a value drawn from all integers between the min and + # max allowed values, and another from only the powers of two + # (both distributions are loguniform) and then randomly + # selecting between the two. + [ + FuzzedParameter( + name=f"k_any_{i}", + minval=_MIN_DIM_SIZE, + maxval=_MAX_DIM_SIZE, + distribution="loguniform", + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k_pow2_{i}", + distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES} + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k{i}", + distribution={ + ParameterAlias(f"k_any_{i}"): 0.8, + ParameterAlias(f"k_pow2_{i}"): 0.2, + }, + strict=True, + ) for i in range(3) + ], + FuzzedParameter( + name="density", + distribution={0.1: 0.4, 0.05: 0.3, 0.01: 0.3}, + ), + FuzzedParameter( + name="coalesced", + distribution={True: 0.5, False: 0.5}, + ), + FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"), + ], + tensors=[ + FuzzedSparseTensor( + name="x", + size=("k0", "k1", "k2"), + dim_parameter="dim_parameter", + sparse_dim="sparse_dim", + min_elements=4 * 1024, + max_elements=32 * 1024 ** 2, + density="density", + coalesced="coalesced", + dtype=dtype, + cuda=cuda, + ), + ], + seed=seed, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/spectral.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/spectral.py new file mode 100644 index 0000000000000000000000000000000000000000..c324e338dca5da3d2b8b9a55e7d89f108d6783dd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/spectral.py @@ -0,0 +1,94 @@ +# mypy: allow-untyped-defs +import math + +import torch +from torch.utils import benchmark +from torch.utils.benchmark import FuzzedParameter, FuzzedTensor, ParameterAlias + + +__all__ = ['SpectralOpFuzzer'] + +MIN_DIM_SIZE = 16 +MAX_DIM_SIZE = 16 * 1024 + +def power_range(upper_bound, base): + return (base ** i for i in range(int(math.log(upper_bound, base)) + 1)) + +# List of regular numbers from MIN_DIM_SIZE to MAX_DIM_SIZE +# These numbers factorize into multiples of prime factors 2, 3, and 5 only +# and are usually the fastest in FFT implementations. +REGULAR_SIZES = [] +for i in power_range(MAX_DIM_SIZE, 2): + for j in power_range(MAX_DIM_SIZE // i, 3): + ij = i * j + for k in power_range(MAX_DIM_SIZE // ij, 5): + ijk = ij * k + if ijk > MIN_DIM_SIZE: + REGULAR_SIZES.append(ijk) +REGULAR_SIZES.sort() + +class SpectralOpFuzzer(benchmark.Fuzzer): + def __init__(self, *, seed: int, dtype=torch.float64, + cuda: bool = False, probability_regular: float = 1.0) -> None: + super().__init__( + parameters=[ + # Dimensionality of x. (e.g. 1D, 2D, or 3D.) + FuzzedParameter("ndim", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True), + + # Shapes for `x`. + # It is important to test all shapes, however + # regular sizes are especially important to the FFT and therefore + # warrant special attention. This is done by generating + # both a value drawn from all integers between the min and + # max allowed values, and another from only the regular numbers + # (both distributions are loguniform) and then randomly + # selecting between the two. + [ + FuzzedParameter( + name=f"k_any_{i}", + minval=MIN_DIM_SIZE, + maxval=MAX_DIM_SIZE, + distribution="loguniform", + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k_regular_{i}", + distribution={size: 1. / len(REGULAR_SIZES) for size in REGULAR_SIZES} + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k{i}", + distribution={ + ParameterAlias(f"k_regular_{i}"): probability_regular, + ParameterAlias(f"k_any_{i}"): 1 - probability_regular, + }, + strict=True, + ) for i in range(3) + ], + + # Steps for `x`. (Benchmarks strided memory access.) + [ + FuzzedParameter( + name=f"step_{i}", + distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04}, + ) for i in range(3) + ], + ], + tensors=[ + FuzzedTensor( + name="x", + size=("k0", "k1", "k2"), + steps=("step_0", "step_1", "step_2"), + probability_contiguous=0.75, + min_elements=4 * 1024, + max_elements=32 * 1024 ** 2, + max_allocation_bytes=2 * 1024**3, # 2 GB + dim_parameter="ndim", + dtype=dtype, + cuda=cuda, + ), + ], + seed=seed, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/unary.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/unary.py new file mode 100644 index 0000000000000000000000000000000000000000..6008adfe459218cd0e239efede5a3f1cd35ee61b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/op_fuzzers/unary.py @@ -0,0 +1,82 @@ +# mypy: allow-untyped-defs +import numpy as np +import torch + +from torch.utils.benchmark import Fuzzer, FuzzedParameter, ParameterAlias, FuzzedTensor + + +_MIN_DIM_SIZE = 16 +_MAX_DIM_SIZE = 16 * 1024 ** 2 +_POW_TWO_SIZES = tuple(2 ** i for i in range( + int(np.log2(_MIN_DIM_SIZE)), + int(np.log2(_MAX_DIM_SIZE)) + 1, +)) + + +class UnaryOpFuzzer(Fuzzer): + def __init__(self, seed, dtype=torch.float32, cuda=False) -> None: + super().__init__( + parameters=[ + # Dimensionality of x. (e.g. 1D, 2D, or 3D.) + FuzzedParameter("dim", distribution={1: 0.3, 2: 0.4, 3: 0.3}, strict=True), + + # Shapes for `x`. + # It is important to test all shapes, however + # powers of two are especially important and therefore + # warrant special attention. This is done by generating + # both a value drawn from all integers between the min and + # max allowed values, and another from only the powers of two + # (both distributions are loguniform) and then randomly + # selecting between the two. + [ + FuzzedParameter( + name=f"k_any_{i}", + minval=_MIN_DIM_SIZE, + maxval=_MAX_DIM_SIZE, + distribution="loguniform", + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k_pow2_{i}", + distribution={size: 1. / len(_POW_TWO_SIZES) for size in _POW_TWO_SIZES} + ) for i in range(3) + ], + [ + FuzzedParameter( + name=f"k{i}", + distribution={ + ParameterAlias(f"k_any_{i}"): 0.8, + ParameterAlias(f"k_pow2_{i}"): 0.2, + }, + strict=True, + ) for i in range(3) + ], + + # Steps for `x`. (Benchmarks strided memory access.) + [ + FuzzedParameter( + name=f"x_step_{i}", + distribution={1: 0.8, 2: 0.06, 4: 0.06, 8: 0.04, 16: 0.04}, + ) for i in range(3) + ], + + # Repeatable entropy for downstream applications. + FuzzedParameter(name="random_value", minval=0, maxval=2 ** 32 - 1, distribution="uniform"), + ], + tensors=[ + FuzzedTensor( + name="x", + size=("k0", "k1", "k2"), + steps=("x_step_0", "x_step_1", "x_step_2"), + probability_contiguous=0.75, + min_elements=4 * 1024, + max_elements=32 * 1024 ** 2, + max_allocation_bytes=2 * 1024**3, # 2 GB + dim_parameter="dim", + dtype=dtype, + cuda=cuda, + ), + ], + seed=seed, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/_stubs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/_stubs.py new file mode 100644 index 0000000000000000000000000000000000000000..c91e3d12b29e1c050edbadaebb877d7fc0761e57 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/_stubs.py @@ -0,0 +1,42 @@ +from typing import Any +from collections.abc import Callable +from typing_extensions import Protocol, runtime_checkable + + +class TimerClass(Protocol): + """This is the portion of the `timeit.Timer` API used by benchmark utils.""" + def __init__( + self, + stmt: str, + setup: str, + timer: Callable[[], float], + globals: dict[str, Any], + **kwargs: Any, + ) -> None: + ... + + def timeit(self, number: int) -> float: + ... + + +@runtime_checkable +class TimeitModuleType(Protocol): + """Modules generated from `timeit_template.cpp`.""" + def timeit(self, number: int) -> float: + ... + + +class CallgrindModuleType(Protocol): + """Replicates the valgrind endpoints in `torch._C`. + + These bindings are used to collect Callgrind profiles on earlier versions + of PyTorch and will eventually be removed. + """ + __file__: str + __name__: str + + def _valgrind_supported_platform(self) -> bool: + ... + + def _valgrind_toggle(self) -> None: + ... diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/common.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/common.py new file mode 100644 index 0000000000000000000000000000000000000000..d4f328d19083f0fc92da79e34d70a68b8ef891ff --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/common.py @@ -0,0 +1,359 @@ +"""Base shared classes and utilities.""" + +import collections +import contextlib +import dataclasses +import os +import shutil +import tempfile +import textwrap +import time +from typing import cast, Any +from collections.abc import Iterable, Iterator +import uuid + +import torch + + +__all__ = ["TaskSpec", "Measurement", "select_unit", "unit_to_english", "trim_sigfig", "ordered_unique", "set_torch_threads"] + + +_MAX_SIGNIFICANT_FIGURES = 4 +_MIN_CONFIDENCE_INTERVAL = 25e-9 # 25 ns + +# Measurement will include a warning if the distribution is suspect. All +# runs are expected to have some variation; these parameters set the +# thresholds. +_IQR_WARN_THRESHOLD = 0.1 +_IQR_GROSS_WARN_THRESHOLD = 0.25 + + +@dataclasses.dataclass(init=True, repr=False, eq=True, frozen=True) +class TaskSpec: + """Container for information used to define a Timer. (except globals)""" + stmt: str + setup: str + global_setup: str = "" + label: str | None = None + sub_label: str | None = None + description: str | None = None + env: str | None = None + num_threads: int = 1 + + @property + def title(self) -> str: + """Best effort attempt at a string label for the measurement.""" + if self.label is not None: + return self.label + (f": {self.sub_label}" if self.sub_label else "") + elif "\n" not in self.stmt: + return self.stmt + (f": {self.sub_label}" if self.sub_label else "") + return ( + f"stmt:{f' ({self.sub_label})' if self.sub_label else ''}\n" + f"{textwrap.indent(self.stmt, ' ')}" + ) + + def setup_str(self) -> str: + return ( + "" if (self.setup == "pass" or not self.setup) + else f"setup:\n{textwrap.indent(self.setup, ' ')}" if "\n" in self.setup + else f"setup: {self.setup}" + ) + + def summarize(self) -> str: + """Build TaskSpec portion of repr string for other containers.""" + sections = [ + self.title, + self.description or "", + self.setup_str(), + ] + return "\n".join([f"{i}\n" if "\n" in i else i for i in sections if i]) + +_TASKSPEC_FIELDS = tuple(i.name for i in dataclasses.fields(TaskSpec)) + + +@dataclasses.dataclass(init=True, repr=False) +class Measurement: + """The result of a Timer measurement. + + This class stores one or more measurements of a given statement. It is + serializable and provides several convenience methods + (including a detailed __repr__) for downstream consumers. + """ + number_per_run: int + raw_times: list[float] + task_spec: TaskSpec + metadata: dict[Any, Any] | None = None # Reserved for user payloads. + + def __post_init__(self) -> None: + self._sorted_times: tuple[float, ...] = () + self._warnings: tuple[str, ...] = () + self._median: float = -1.0 + self._mean: float = -1.0 + self._p25: float = -1.0 + self._p75: float = -1.0 + + def __getattr__(self, name: str) -> Any: + # Forward TaskSpec fields for convenience. + if name in _TASKSPEC_FIELDS: + return getattr(self.task_spec, name) + return super().__getattribute__(name) + + # ========================================================================= + # == Convenience methods for statistics =================================== + # ========================================================================= + # + # These methods use raw time divided by number_per_run; this is an + # extrapolation and hides the fact that different number_per_run will + # result in different amortization of overheads, however if Timer has + # selected an appropriate number_per_run then this is a non-issue, and + # forcing users to handle that division would result in a poor experience. + @property + def times(self) -> list[float]: + return [t / self.number_per_run for t in self.raw_times] + + @property + def median(self) -> float: + self._lazy_init() + return self._median + + @property + def mean(self) -> float: + self._lazy_init() + return self._mean + + @property + def iqr(self) -> float: + self._lazy_init() + return self._p75 - self._p25 + + @property + def significant_figures(self) -> int: + """Approximate significant figure estimate. + + This property is intended to give a convenient way to estimate the + precision of a measurement. It only uses the interquartile region to + estimate statistics to try to mitigate skew from the tails, and + uses a static z value of 1.645 since it is not expected to be used + for small values of `n`, so z can approximate `t`. + + The significant figure estimation used in conjunction with the + `trim_sigfig` method to provide a more human interpretable data + summary. __repr__ does not use this method; it simply displays raw + values. Significant figure estimation is intended for `Compare`. + """ + self._lazy_init() + n_total = len(self._sorted_times) + lower_bound = int(n_total // 4) + upper_bound = int(torch.tensor(3 * n_total / 4).ceil()) + interquartile_points: tuple[float, ...] = self._sorted_times[lower_bound:upper_bound] + std = torch.tensor(interquartile_points).std(unbiased=False).item() + sqrt_n = torch.tensor(len(interquartile_points)).sqrt().item() + + # Rough estimates. These are by no means statistically rigorous. + confidence_interval = max(1.645 * std / sqrt_n, _MIN_CONFIDENCE_INTERVAL) + relative_ci = torch.tensor(self._median / confidence_interval).log10().item() + num_significant_figures = int(torch.tensor(relative_ci).floor()) + return min(max(num_significant_figures, 1), _MAX_SIGNIFICANT_FIGURES) + + @property + def has_warnings(self) -> bool: + self._lazy_init() + return bool(self._warnings) + + def _lazy_init(self) -> None: + if self.raw_times and not self._sorted_times: + self._sorted_times = tuple(sorted(self.times)) + _sorted_times = torch.tensor(self._sorted_times, dtype=torch.float64) + self._median = _sorted_times.quantile(.5).item() + self._mean = _sorted_times.mean().item() + self._p25 = _sorted_times.quantile(.25).item() + self._p75 = _sorted_times.quantile(.75).item() + + def add_warning(msg: str) -> None: + rel_iqr = self.iqr / self.median * 100 + self._warnings += ( + f" WARNING: Interquartile range is {rel_iqr:.1f}% " + f"of the median measurement.\n {msg}", + ) + + if not self.meets_confidence(_IQR_GROSS_WARN_THRESHOLD): + add_warning("This suggests significant environmental influence.") + elif not self.meets_confidence(_IQR_WARN_THRESHOLD): + add_warning("This could indicate system fluctuation.") + + + def meets_confidence(self, threshold: float = _IQR_WARN_THRESHOLD) -> bool: + return self.iqr / self.median < threshold + + @property + def title(self) -> str: + return self.task_spec.title + + @property + def env(self) -> str: + return ( + "Unspecified env" if self.taskspec.env is None + else cast(str, self.taskspec.env) + ) + + @property + def as_row_name(self) -> str: + return self.sub_label or self.stmt or "[Unknown]" + + def __repr__(self) -> str: + """ + Example repr: + + Broadcasting add (4x8) + Median: 5.73 us + IQR: 2.25 us (4.01 to 6.26) + 372 measurements, 100 runs per measurement, 1 thread + WARNING: Interquartile range is 39.4% of the median measurement. + This suggests significant environmental influence. + """ + self._lazy_init() + skip_line, newline = "MEASUREMENT_REPR_SKIP_LINE", "\n" + n = len(self._sorted_times) + time_unit, time_scale = select_unit(self._median) + iqr_filter = '' if n >= 4 else skip_line + + repr_str = f""" +{super().__repr__()} +{self.task_spec.summarize()} + {'Median: ' if n > 1 else ''}{self._median / time_scale:.2f} {time_unit} + {iqr_filter}IQR: {self.iqr / time_scale:.2f} {time_unit} ({self._p25 / time_scale:.2f} to {self._p75 / time_scale:.2f}) + {n} measurement{'s' if n > 1 else ''}, {self.number_per_run} runs {'per measurement,' if n > 1 else ','} {self.num_threads} thread{'s' if self.num_threads > 1 else ''} +{newline.join(self._warnings)}""".strip() # noqa: B950 + + return "\n".join(l for l in repr_str.splitlines(keepends=False) if skip_line not in l) + + @staticmethod + def merge(measurements: Iterable["Measurement"]) -> list["Measurement"]: + """Convenience method for merging replicates. + + Merge will extrapolate times to `number_per_run=1` and will not + transfer any metadata. (Since it might differ between replicates) + """ + grouped_measurements: collections.defaultdict[TaskSpec, list[Measurement]] = collections.defaultdict(list) + for m in measurements: + grouped_measurements[m.task_spec].append(m) + + def merge_group(task_spec: TaskSpec, group: list["Measurement"]) -> "Measurement": + times: list[float] = [] + for m in group: + # Different measurements could have different `number_per_run`, + # so we call `.times` which normalizes the results. + times.extend(m.times) + + return Measurement( + number_per_run=1, + raw_times=times, + task_spec=task_spec, + metadata=None, + ) + + return [merge_group(t, g) for t, g in grouped_measurements.items()] + + +def select_unit(t: float) -> tuple[str, float]: + """Determine how to scale times for O(1) magnitude. + + This utility is used to format numbers for human consumption. + """ + time_unit = {-3: "ns", -2: "us", -1: "ms"}.get(int(torch.tensor(t).log10().item() // 3), "s") + time_scale = {"ns": 1e-9, "us": 1e-6, "ms": 1e-3, "s": 1}[time_unit] + return time_unit, time_scale + + +def unit_to_english(u: str) -> str: + return { + "ns": "nanosecond", + "us": "microsecond", + "ms": "millisecond", + "s": "second", + }[u] + + +def trim_sigfig(x: float, n: int) -> float: + """Trim `x` to `n` significant figures. (e.g. 3.14159, 2 -> 3.10000)""" + if n != int(n): + raise AssertionError("Number of significant figures must be an integer") + magnitude = int(torch.tensor(x).abs().log10().ceil().item()) + scale = 10 ** (magnitude - n) + return float(torch.tensor(x / scale).round() * scale) + + +def ordered_unique(elements: Iterable[Any]) -> list[Any]: + return list(collections.OrderedDict(dict.fromkeys(elements)).keys()) + + +@contextlib.contextmanager +def set_torch_threads(n: int) -> Iterator[None]: + prior_num_threads = torch.get_num_threads() + try: + torch.set_num_threads(n) + yield + finally: + torch.set_num_threads(prior_num_threads) + + +def _make_temp_dir(prefix: str | None = None, gc_dev_shm: bool = False) -> str: + """Create a temporary directory. The caller is responsible for cleanup. + + This function is conceptually similar to `tempfile.mkdtemp`, but with + the key additional feature that it will use shared memory if the + `BENCHMARK_USE_DEV_SHM` environment variable is set. This is an + implementation detail, but an important one for cases where many Callgrind + measurements are collected at once. (Such as when collecting + microbenchmarks.) + + This is an internal utility, and is exported solely so that microbenchmarks + can reuse the util. + """ + use_dev_shm: bool = (os.getenv("BENCHMARK_USE_DEV_SHM") or "").lower() in ("1", "true") + if use_dev_shm: + root = "/dev/shm/pytorch_benchmark_utils" + if os.name != "posix": + raise AssertionError(f"tmpfs (/dev/shm) is POSIX only, current platform is {os.name}") + if not os.path.exists("/dev/shm"): + raise AssertionError("This system does not appear to support tmpfs (/dev/shm).") + os.makedirs(root, exist_ok=True) + + # Because we're working in shared memory, it is more important than + # usual to clean up ALL intermediate files. However we don't want every + # worker to walk over all outstanding directories, so instead we only + # check when we are sure that it won't lead to contention. + if gc_dev_shm: + for i in os.listdir(root): + owner_file = os.path.join(root, i, "owner.pid") + if not os.path.exists(owner_file): + continue + + with open(owner_file) as f: + owner_pid = int(f.read()) + + if owner_pid == os.getpid(): + continue + + try: + # https://stackoverflow.com/questions/568271/how-to-check-if-there-exists-a-process-with-a-given-pid-in-python + os.kill(owner_pid, 0) + + except OSError: + print(f"Detected that {os.path.join(root, i)} was orphaned in shared memory. Cleaning up.") + shutil.rmtree(os.path.join(root, i)) + + else: + root = tempfile.gettempdir() + + # We include the time so names sort by creation time, and add a UUID + # to ensure we don't collide. + name = f"{prefix or tempfile.gettempprefix()}__{int(time.time())}__{uuid.uuid4()}" + path = os.path.join(root, name) + os.makedirs(path, exist_ok=False) + + if use_dev_shm: + with open(os.path.join(path, "owner.pid"), "w") as f: + f.write(str(os.getpid())) + + return path diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/compare.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/compare.py new file mode 100644 index 0000000000000000000000000000000000000000..c1e232e6e04260f277254c9b181c63dfeaadee62 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/compare.py @@ -0,0 +1,345 @@ +# mypy: allow-untyped-defs +"""Display class to aggregate and print the results of many measurements.""" +import collections +import enum +import itertools as it + +from torch.utils.benchmark.utils import common +from torch import tensor as _tensor +import operator + +__all__ = ["Colorize", "Compare"] + +BEST = "\033[92m" +GOOD = "\033[34m" +BAD = "\033[2m\033[91m" +VERY_BAD = "\033[31m" +BOLD = "\033[1m" +TERMINATE = "\033[0m" + + +class Colorize(enum.Enum): + NONE = "none" + COLUMNWISE = "columnwise" + ROWWISE = "rowwise" + + +# Classes to separate internal bookkeeping from what is rendered. +class _Column: + def __init__( + self, + grouped_results: list[tuple[common.Measurement | None, ...]], + time_scale: float, + time_unit: str, + trim_significant_figures: bool, + highlight_warnings: bool, + ) -> None: + self._grouped_results = grouped_results + self._flat_results = [*it.chain.from_iterable(grouped_results)] + self._time_scale = time_scale + self._time_unit = time_unit + self._trim_significant_figures = trim_significant_figures + self._highlight_warnings = ( + highlight_warnings + and any(r.has_warnings for r in self._flat_results if r) + ) + leading_digits = [ + int(_tensor(r.median / self._time_scale).log10().ceil()) if r else None + for r in self._flat_results + ] + unit_digits = max(d for d in leading_digits if d is not None) + decimal_digits = min( + max(m.significant_figures - digits, 0) + for digits, m in zip(leading_digits, self._flat_results, strict=True) + if (m is not None) and (digits is not None) + ) if self._trim_significant_figures else 1 + length = unit_digits + decimal_digits + (1 if decimal_digits else 0) + self._template = f"{{:>{length}.{decimal_digits}f}}{{:>{7 if self._highlight_warnings else 0}}}" + + def get_results_for(self, group): + return self._grouped_results[group] + + def num_to_str(self, value: float | None, estimated_sigfigs: int, spread: float | None): + if value is None: + return " " * len(self.num_to_str(1, estimated_sigfigs, None)) + + if self._trim_significant_figures: + value = common.trim_sigfig(value, estimated_sigfigs) + + return self._template.format( + value, + f" (! {spread * 100:.0f}%)" if self._highlight_warnings and spread is not None else "") + + +def optional_min(seq): + l = list(seq) + return None if len(l) == 0 else min(l) + + +class _Row: + def __init__(self, results, row_group, render_env, env_str_len, + row_name_str_len, time_scale, colorize, num_threads=None) -> None: + super().__init__() + self._results = results + self._row_group = row_group + self._render_env = render_env + self._env_str_len = env_str_len + self._row_name_str_len = row_name_str_len + self._time_scale = time_scale + self._colorize = colorize + self._columns: tuple[_Column, ...] = () + self._num_threads = num_threads + + def register_columns(self, columns: tuple[_Column, ...]) -> None: + self._columns = columns + + def as_column_strings(self): + concrete_results = [r for r in self._results if r is not None] + env = f"({concrete_results[0].env})" if self._render_env else "" + env = env.ljust(self._env_str_len + 4) + output = [" " + env + concrete_results[0].as_row_name] + for m, col in zip(self._results, self._columns or (), strict=False): + if m is None: + output.append(col.num_to_str(None, 1, None)) + else: + output.append(col.num_to_str( + m.median / self._time_scale, + m.significant_figures, + m.iqr / m.median if m.has_warnings else None + )) + return output + + @staticmethod + def color_segment(segment, value, best_value): + if value <= best_value * 1.01 or value <= best_value + 100e-9: + return BEST + BOLD + segment + TERMINATE * 2 + if value <= best_value * 1.1: + return GOOD + BOLD + segment + TERMINATE * 2 + if value >= best_value * 5: + return VERY_BAD + BOLD + segment + TERMINATE * 2 + if value >= best_value * 2: + return BAD + segment + TERMINATE * 2 + + return segment + + def row_separator(self, overall_width): + return ( + [f"{self._num_threads} threads: ".ljust(overall_width, "-")] + if self._num_threads is not None else [] + ) + + def finalize_column_strings(self, column_strings, col_widths): + best_values = [-1 for _ in column_strings] + if self._colorize == Colorize.ROWWISE: + row_min = min(r.median for r in self._results if r is not None) + best_values = [row_min for _ in column_strings] + elif self._colorize == Colorize.COLUMNWISE: + best_values = [ + optional_min(r.median for r in column.get_results_for(self._row_group) if r is not None) + for column in (self._columns or ()) + ] + + row_contents = [column_strings[0].ljust(col_widths[0])] + for col_str, width, result, best_value in zip(column_strings[1:], col_widths[1:], self._results, best_values, strict=False): + col_str = col_str.center(width) + if self._colorize != Colorize.NONE and result is not None and best_value is not None: + col_str = self.color_segment(col_str, result.median, best_value) + row_contents.append(col_str) + return row_contents + + +class Table: + def __init__( + self, + results: list[common.Measurement], + colorize: Colorize, + trim_significant_figures: bool, + highlight_warnings: bool + ) -> None: + if len({r.label for r in results}) != 1: + raise AssertionError("All results must share the same label") + + self.results = results + self._colorize = colorize + self._trim_significant_figures = trim_significant_figures + self._highlight_warnings = highlight_warnings + self.label = results[0].label + self.time_unit, self.time_scale = common.select_unit( + min(r.median for r in results) + ) + + self.row_keys = common.ordered_unique([self.row_fn(i) for i in results]) + self.row_keys.sort(key=operator.itemgetter(slice(2))) # preserve stmt order + self.column_keys = common.ordered_unique([self.col_fn(i) for i in results]) + self.rows, self.columns = self.populate_rows_and_columns() + + @staticmethod + def row_fn(m: common.Measurement) -> tuple[int, str | None, str]: + return m.num_threads, m.env, m.as_row_name + + @staticmethod + def col_fn(m: common.Measurement) -> str | None: + return m.description + + def populate_rows_and_columns(self) -> tuple[tuple[_Row, ...], tuple[_Column, ...]]: + rows: list[_Row] = [] + columns: list[_Column] = [] + ordered_results: list[list[common.Measurement | None]] = [ + [None for _ in self.column_keys] + for _ in self.row_keys + ] + row_position = {key: i for i, key in enumerate(self.row_keys)} + col_position = {key: i for i, key in enumerate(self.column_keys)} + for r in self.results: + i = row_position[self.row_fn(r)] + j = col_position[self.col_fn(r)] + ordered_results[i][j] = r + + unique_envs = {r.env for r in self.results} + render_env = len(unique_envs) > 1 + env_str_len = max(len(i) for i in unique_envs) if render_env else 0 + + row_name_str_len = max(len(r.as_row_name) for r in self.results) + + prior_num_threads = -1 + prior_env = "" + row_group = -1 + rows_by_group: list[list[list[common.Measurement | None]]] = [] + for (num_threads, env, _), row in zip(self.row_keys, ordered_results, strict=True): + thread_transition = (num_threads != prior_num_threads) + if thread_transition: + prior_num_threads = num_threads + prior_env = "" + row_group += 1 + rows_by_group.append([]) + rows.append( + _Row( + results=row, + row_group=row_group, + render_env=(render_env and env != prior_env), + env_str_len=env_str_len, + row_name_str_len=row_name_str_len, + time_scale=self.time_scale, + colorize=self._colorize, + num_threads=num_threads if thread_transition else None, + ) + ) + rows_by_group[-1].append(row) + prior_env = env + + for i in range(len(self.column_keys)): + grouped_results = [tuple(row[i] for row in g) for g in rows_by_group] + column = _Column( + grouped_results=grouped_results, + time_scale=self.time_scale, + time_unit=self.time_unit, + trim_significant_figures=self._trim_significant_figures, + highlight_warnings=self._highlight_warnings,) + columns.append(column) + + rows_tuple, columns_tuple = tuple(rows), tuple(columns) + for ri in rows_tuple: + ri.register_columns(columns_tuple) + return rows_tuple, columns_tuple + + def render(self) -> str: + string_rows = [[""] + self.column_keys] + string_rows.extend(r.as_column_strings() for r in self.rows) + num_cols = max(len(i) for i in string_rows) + for sr in string_rows: + sr.extend(["" for _ in range(num_cols - len(sr))]) + + col_widths = [max(len(j) for j in i) for i in zip(*string_rows, strict=True)] + finalized_columns = [" | ".join(i.center(w) for i, w in zip(string_rows[0], col_widths, strict=True))] + overall_width = len(finalized_columns[0]) + for string_row, row in zip(string_rows[1:], self.rows, strict=True): + finalized_columns.extend(row.row_separator(overall_width)) + finalized_columns.append(" | ".join(row.finalize_column_strings(string_row, col_widths))) + + newline = "\n" + has_warnings = self._highlight_warnings and any(ri.has_warnings for ri in self.results) + return f""" +[{(' ' + (self.label or '') + ' ').center(overall_width - 2, '-')}] +{newline.join(finalized_columns)} + +Times are in {common.unit_to_english(self.time_unit)}s ({self.time_unit}). +{'(! XX%) Measurement has high variance, where XX is the IQR / median * 100.' + newline if has_warnings else ""}"""[1:] + + +class Compare: + """Helper class for displaying the results of many measurements in a + formatted table. + + The table format is based on the information fields provided in + :class:`torch.utils.benchmark.Timer` (`description`, `label`, `sub_label`, + `num_threads`, etc). + + The table can be directly printed using :meth:`print` or casted as a `str`. + + For a full tutorial on how to use this class, see: + https://pytorch.org/tutorials/recipes/recipes/benchmark.html + + Args: + results: List of Measurement to display. + """ + def __init__(self, results: list[common.Measurement]) -> None: + self._results: list[common.Measurement] = [] + self.extend_results(results) + self._trim_significant_figures = False + self._colorize = Colorize.NONE + self._highlight_warnings = False + + def __str__(self) -> str: + return "\n".join(self._render()) + + def extend_results(self, results) -> None: + """Append results to already stored ones. + + All added results must be instances of ``Measurement``. + """ + for r in results: + if not isinstance(r, common.Measurement): + raise ValueError( + "Expected an instance of `Measurement`, " f"got {type(r)} instead." + ) + self._results.extend(results) + + def trim_significant_figures(self) -> None: + """Enables trimming of significant figures when building the formatted table.""" + self._trim_significant_figures = True + + def colorize(self, rowwise=False) -> None: + """Colorize formatted table. + + Colorize columnwise by default. + """ + self._colorize = Colorize.ROWWISE if rowwise else Colorize.COLUMNWISE + + def highlight_warnings(self) -> None: + """Enables warning highlighting when building formatted table.""" + self._highlight_warnings = True + + def print(self) -> None: + """Print formatted table""" + print(str(self)) + + def _render(self): + results = common.Measurement.merge(self._results) + grouped_results = self._group_by_label(results) + output = [self._layout(group) for group in grouped_results.values()] + return output + + def _group_by_label(self, results: list[common.Measurement]): + grouped_results: collections.defaultdict[str, list[common.Measurement]] = collections.defaultdict(list) + for r in results: + grouped_results[r.label].append(r) + return grouped_results + + def _layout(self, results: list[common.Measurement]): + table = Table( + results, + self._colorize, + self._trim_significant_figures, + self._highlight_warnings + ) + return table.render() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/compile.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/compile.py new file mode 100644 index 0000000000000000000000000000000000000000..dd15a582a274980bea4aff22f7325ccf562ecb13 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/compile.py @@ -0,0 +1,195 @@ +# mypy: allow-untyped-defs +from typing import Any, cast +from collections.abc import Callable + +import torch +import torch._dynamo +from torch._dynamo.testing import CompileCounterWithBackend +from torch.utils.benchmark import Timer + + +__all__ = ["bench_all", "benchmark_compile"] + + +_warned_tensor_cores = False +_default_float_32_precision = torch.get_float32_matmul_precision() + +try: + + from tabulate import tabulate + + HAS_TABULATE = True +except ModuleNotFoundError: + HAS_TABULATE = False + tabulate = None # type: ignore[assignment] + print("tabulate is not installed, please pip install tabulate to use this utility") + +if HAS_TABULATE: + def _enable_tensor_cores() -> None: + global _warned_tensor_cores + + if torch.cuda.is_available(): + if torch.backends.cuda.matmul.allow_tf32 is False and torch.cuda.get_device_capability() >= (8, 0): + torch.set_float32_matmul_precision("high") + if not _warned_tensor_cores: + print("Your GPU supports tensor cores") + print("we will enable it automatically by setting `torch.set_float32_matmul_precision('high')`") + _warned_tensor_cores = True + + def _disable_tensor_cores() -> None: + torch.set_float32_matmul_precision(_default_float_32_precision) + + def bench_loop( + model: torch.nn.Module | Callable, + sample_input: torch.Tensor | Any, + num_iters: int = 5, + optimizer: torch.optim.Optimizer | None = None, + loss_fn: Callable | None = None, + ): + # Define the statement and setup for the benchmark + if optimizer and loss_fn: + # Training mode + stmt = """ + output = model(sample_input) + loss = loss_fn(output) if loss_fn else output.sum() + loss.backward() + optimizer.step() + optimizer.zero_grad() + """ + else: + # Inference mode + stmt = "model(sample_input)" + + # Create the Timer object + timer = Timer( + stmt=stmt, + globals={"model": model, "sample_input": sample_input, "optimizer": optimizer, "loss_fn": loss_fn}, + ) + + + result = timer.timeit(number=num_iters) + + # Get the average time per iteration in milliseconds + avg_time = result.mean * 1000 + return round(avg_time, 2) + + def benchmark_compile( + model: torch.nn.Module | Callable, + sample_input: torch.Tensor | Any, + num_iters: int = 5, + backend: str | None = None, + mode: str | None = "default", + optimizer: torch.optim.Optimizer | None = None, + loss_fn : torch.nn.Module | Callable | None = None, + ): + """ + Use this utility to benchmark torch.compile + """ + if backend: + try: + torch._dynamo.reset() + compile_counter_with_backend = CompileCounterWithBackend(backend) + opt_model = torch.compile(model, backend=compile_counter_with_backend, mode=mode) + + # Compilation only happens after the first inference + compilation_time = bench_loop(opt_model, sample_input, 1, optimizer, loss_fn) + + running_time = bench_loop(opt_model, sample_input, num_iters, optimizer, loss_fn) + + if compile_counter_with_backend.frame_count == 0: + raise RuntimeError("No compilation occurred during benchmarking.") + + if compile_counter_with_backend.frame_count > 1: + raise RuntimeError("Recompilation occurred during benchmarking.") + + except Exception as e: + print(e) + print(f"Failed to compile {backend} with mode {mode}") + return None, None + else: + opt_model = model + compilation_time = None + running_time = bench_loop(opt_model, sample_input, num_iters, optimizer, loss_fn) + + compilation_time = round(compilation_time, 2) if compilation_time else None + running_time = round(running_time, 2) if running_time else None + + + return compilation_time, running_time + + + def bench_all( + model : torch.nn.Module | Callable, + sample_input: torch.Tensor | Any, + num_iters : int = 5, + optimizer: torch.optim.Optimizer | None = None, + loss_fn : torch.nn.Module | Callable | None = None, + ): + """ + This is a simple utility that can be used to benchmark torch.compile + In particular it ensures that your GPU is setup to use tensor cores if it supports its + It also tries out all the main backends and prints a table of results so you can easily compare them all + Many of the backendds have their own optional dependencies so please pip install them separately + + You will get one table for inference and another for training + If you'd like to leverage this utility for training make sure to pass in a torch.optim.Optimizer + + The important warnings are + Your GPU supports tensor cores + we will enable it automatically by setting `torch.set_float32_matmul_precision('high')` + + If a compilation fails for any reason including the dependency not being included + then we will print Failed to compile {backend} with mode {mode} + """ + field_names = ["Train/Inference", "Backend", "Mode", "Compilation Time", "Average Running Time"] + table = [] + + + eager_time = None + torch._dynamo.reset() + _, eager_time = benchmark_compile(model, sample_input, num_iters, None, None, optimizer) + table.append( + [("Training" if optimizer else "Inference"), "Eager", "-", "-", f"{eager_time} ms"] + ) + + for backend in torch._dynamo.list_backends(): + + if backend == "inductor": + mode_options = cast(list[str | None], list(torch._inductor.list_mode_options().keys())) + [None] + for mode in mode_options: + if mode == "default": + continue + torch._dynamo.reset() + try: + if torch.cuda.is_available(): + _enable_tensor_cores() + compilation_time, running_time = benchmark_compile( + model, sample_input, num_iters, backend, mode, optimizer, loss_fn) + finally: + if torch.cuda.is_available(): + _disable_tensor_cores() + table.append([ + ("Training" if optimizer else "Inference"), + # pyrefly: ignore [redundant-condition] + backend if backend else "-", + mode if mode is not None else "-", + f"{compilation_time} ms " if compilation_time else "-", + f"{running_time} ms " if running_time else "-", + ]) + + else: + torch._dynamo.reset() + compilation_time, running_time = benchmark_compile( + model, sample_input, num_iters, backend, None, optimizer, loss_fn) + + if running_time is not None: + table.append([ + ("Training" if optimizer else "Inference"), + backend, "-", + f"{compilation_time} ms " or "-", + f"{running_time} ms ", + ]) + + + # pyrefly: ignore [not-callable] + return tabulate(table, headers=field_names, tablefmt="github") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/cpp_jit.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/cpp_jit.py new file mode 100644 index 0000000000000000000000000000000000000000..a298146ce17c7ff6f303b4d76c4c96ba786ae774 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/cpp_jit.py @@ -0,0 +1,175 @@ +"""JIT C++ strings into executables.""" +import atexit +import os +import re +import shutil +import textwrap +import threading +from typing import Any + +import torch +from torch.utils.benchmark.utils._stubs import CallgrindModuleType, TimeitModuleType +from torch.utils.benchmark.utils.common import _make_temp_dir +from torch.utils import cpp_extension + + +LOCK = threading.Lock() +SOURCE_ROOT = os.path.split(os.path.abspath(__file__))[0] + +# We calculate uuid once at import time so that separate processes will have +# separate build roots, but threads will share the same build root. +# `cpp_extension` uses build root as part of the cache key, so per-invocation +# uuid's (e.g. different build root per _compile_template call) would lead to +# a 0% cache hit rate and spurious recompilation. Consider the following: +# ``` +# setup = "auto x = torch::ones({1024, 1024});" +# stmt = "torch::mm(x, x);" +# for num_threads in [1, 2, 4, 8]: +# print(Timer(stmt, setup, num_threads=num_threads, language="c++").blocked_autorange()) +# ```` +# `setup` and `stmt` do not change, so we can reuse the executable from the +# first pass through the loop. +_BUILD_ROOT: str | None = None + +def _get_build_root() -> str: + global _BUILD_ROOT + if _BUILD_ROOT is None: + _BUILD_ROOT = _make_temp_dir(prefix="benchmark_utils_jit_build") + # pyrefly: ignore [missing-argument] + atexit.register(shutil.rmtree, _BUILD_ROOT) + return _BUILD_ROOT + + +# BACK_TESTING_NOTE: +# There are two workflows where this code could be used. One is the obvious +# case where someone simply builds or installs PyTorch and uses Timer. +# The other is that the entire `torch/utils/benchmark` folder from a CURRENT +# PyTorch checkout is copy-pasted into a much OLDER version of the PyTorch +# source code. This is what we refer to here as "back testing". The rationale +# is that we might want to use current tooling to study some aspect of an +# earlier version of PyTorch. (e.g. a regression.) +# +# The problem is that Timer relies on several aspects of core PyTorch, namely +# some binding functions for Valgrind symbols in `torch._C` and the +# `torch.__config__._cxx_flags()` method. If we were to naively copy code +# around this wouldn't work as the symbols of interest aren't present in +# earlier versions of PyTorch. In order to work around this, we must add back +# testing shims. These shims will never activate during normal use, but will +# allow Timer to function outside of the "correct" version of PyTorch by +# emulating functionality that was added later. +# +# These shims are temporary, and as Timer becomes more integrated with +# PyTorch the cost and complexity of such shims will increase. Once back +# testing is no longer required (which is to say we have done enough historic +# analysis and the shims no longer justify their maintenance and code +# complexity costs) back testing paths will be removed. + +CXX_FLAGS: list[str] | None +if hasattr(torch.__config__, "_cxx_flags"): + try: + CXX_FLAGS = torch.__config__._cxx_flags().strip().split() + if CXX_FLAGS is not None and "-g" not in CXX_FLAGS: + CXX_FLAGS.append("-g") + # remove "-W" flags to allow build benchmarks + # with a relaxed constraint of compiler versions + if CXX_FLAGS is not None: + CXX_FLAGS = list(filter(lambda x: not x.startswith("-W"), CXX_FLAGS)) + + except RuntimeError: + # We are in FBCode. + CXX_FLAGS = None +else: + # FIXME: Remove when back testing is no longer required. + CXX_FLAGS = ["-O2", "-fPIC", "-g"] + +EXTRA_INCLUDE_PATHS: list[str] = [os.path.join(SOURCE_ROOT, "valgrind_wrapper")] +CONDA_PREFIX = os.getenv("CONDA_PREFIX") +if CONDA_PREFIX is not None: + # Load will automatically search /usr/include, but not conda include. + EXTRA_INCLUDE_PATHS.append(os.path.join(CONDA_PREFIX, "include")) + + +COMPAT_CALLGRIND_BINDINGS: CallgrindModuleType | None = None +def get_compat_bindings() -> CallgrindModuleType: + with LOCK: + global COMPAT_CALLGRIND_BINDINGS + if COMPAT_CALLGRIND_BINDINGS is None: + COMPAT_CALLGRIND_BINDINGS = cpp_extension.load( + name="callgrind_bindings", + sources=[os.path.join( + SOURCE_ROOT, + "valgrind_wrapper", + "compat_bindings.cpp" + )], + extra_cflags=CXX_FLAGS, + extra_include_paths=EXTRA_INCLUDE_PATHS, + ) + return COMPAT_CALLGRIND_BINDINGS + + +def _compile_template( + *, + stmt: str, + setup: str, + global_setup: str, + src: str, + is_standalone: bool +) -> Any: + for before, after, indentation in ( + ("// GLOBAL_SETUP_TEMPLATE_LOCATION", global_setup, 0), + ("// SETUP_TEMPLATE_LOCATION", setup, 4), + ("// STMT_TEMPLATE_LOCATION", stmt, 8) + ): + # C++ doesn't care about indentation so this code isn't load + # bearing the way it is with Python, but this makes the source + # look nicer if a human has to look at it. + src = re.sub( + before, + textwrap.indent(after, " " * indentation)[indentation:], + src + ) + + # We want to isolate different Timers. However `cpp_extension` will + # cache builds which will significantly reduce the cost of repeated + # invocations. + with LOCK: + name = f"timer_cpp_{abs(hash(src))}" + build_dir = os.path.join(_get_build_root(), name) + os.makedirs(build_dir, exist_ok=True) + + src_path = os.path.join(build_dir, "timer_src.cpp") + with open(src_path, "w") as f: + f.write(src) + + # `cpp_extension` has its own locking scheme, so we don't need our lock. + return cpp_extension.load( + name=name, + sources=[src_path], + build_directory=build_dir, + extra_cflags=CXX_FLAGS, + extra_include_paths=EXTRA_INCLUDE_PATHS, + is_python_module=not is_standalone, + is_standalone=is_standalone, + ) + + +def compile_timeit_template(*, stmt: str, setup: str, global_setup: str) -> TimeitModuleType: + template_path: str = os.path.join(SOURCE_ROOT, "timeit_template.cpp") + with open(template_path) as f: + src: str = f.read() + + module = _compile_template(stmt=stmt, setup=setup, global_setup=global_setup, src=src, is_standalone=False) + if not isinstance(module, TimeitModuleType): + raise AssertionError("compiled module is not a TimeitModuleType") + return module + + +def compile_callgrind_template(*, stmt: str, setup: str, global_setup: str) -> str: + template_path: str = os.path.join(SOURCE_ROOT, "valgrind_wrapper", "timer_callgrind_template.cpp") + with open(template_path) as f: + src: str = f.read() + + target = _compile_template(stmt=stmt, setup=setup, global_setup=global_setup, src=src, is_standalone=True) + if not isinstance(target, str): + raise AssertionError("compiled target path is not a string") + return target diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/fuzzer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/fuzzer.py new file mode 100644 index 0000000000000000000000000000000000000000..38f771d23632efd27239e460591d923be3ee59fc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/fuzzer.py @@ -0,0 +1,469 @@ +# mypy: allow-untyped-defs +import functools +import itertools as it +from typing import Any +from collections.abc import Callable + +import torch + + +__all__ = [ + "Fuzzer", + "FuzzedParameter", "ParameterAlias", + "FuzzedTensor", +] + + +_DISTRIBUTIONS = ( + "loguniform", + "uniform", +) + + +class FuzzedParameter: + """Specification for a parameter to be generated during fuzzing.""" + def __init__( + self, + name: str, + minval: int | float | None = None, + maxval: int | float | None = None, + distribution: str | dict[Any, float] | None = None, + strict: bool = False, + ) -> None: + """ + Args: + name: + A string name with which to identify the parameter. + FuzzedTensors can reference this string in their + specifications. + minval: + The lower bound for the generated value. See the description + of `distribution` for type behavior. + maxval: + The upper bound for the generated value. Type behavior is + identical to `minval`. + distribution: + Specifies the distribution from which this parameter should + be drawn. There are three possibilities: + - "loguniform" + Samples between `minval` and `maxval` (inclusive) such + that the probabilities are uniform in log space. As a + concrete example, if minval=1 and maxval=100, a sample + is as likely to fall in [1, 10) as it is [10, 100]. + - "uniform" + Samples are chosen with uniform probability between + `minval` and `maxval` (inclusive). If either `minval` + or `maxval` is a float then the distribution is the + continuous uniform distribution; otherwise samples + are constrained to the integers. + - dict: + If a dict is passed, the keys are taken to be choices + for the variables and the values are interpreted as + probabilities. (And must sum to one.) + If a dict is passed, `minval` and `maxval` must not be set. + Otherwise, they must be set. + strict: + If a parameter is strict, it will not be included in the + iterative resampling process which Fuzzer uses to find a + valid parameter configuration. This allows an author to + prevent skew from resampling for a given parameter (for + instance, a low size limit could inadvertently bias towards + Tensors with fewer dimensions) at the cost of more iterations + when generating parameters. + """ + self._name = name + self._minval = minval + self._maxval = maxval + self._distribution = self._check_distribution(distribution) + self.strict = strict + + @property + def name(self): + return self._name + + def sample(self, state): + if self._distribution == "loguniform": + return self._loguniform(state) + + if self._distribution == "uniform": + return self._uniform(state) + + if isinstance(self._distribution, dict): + return self._custom_distribution(state) + + def _check_distribution(self, distribution): + if not isinstance(distribution, dict): + if distribution not in _DISTRIBUTIONS: + raise AssertionError(f"Unknown distribution: {distribution}") + else: + if any(i < 0 for i in distribution.values()): + raise AssertionError("Probabilities cannot be negative") + if not abs(sum(distribution.values()) - 1) > 1e-5: + raise AssertionError("Distribution is not normalized") + if self._minval is not None: + raise AssertionError("When passing a custom distribution, 'minval' must be None") + if self._maxval is not None: + raise AssertionError("When passing a custom distribution, 'maxval' must be None") + + return distribution + + def _loguniform(self, state): + import numpy as np + output = int(2 ** state.uniform( + low=np.log2(self._minval) if self._minval is not None else None, + high=np.log2(self._maxval) if self._maxval is not None else None, + )) + if self._minval is not None and output < self._minval: + return self._minval + if self._maxval is not None and output > self._maxval: + return self._maxval + return output + + def _uniform(self, state): + if isinstance(self._minval, int) and isinstance(self._maxval, int): + return int(state.randint(low=self._minval, high=self._maxval + 1)) + return state.uniform(low=self._minval, high=self._maxval) + + def _custom_distribution(self, state): + import numpy as np + # If we directly pass the keys to `choice`, numpy will convert + # them to numpy dtypes. + index = state.choice( + np.arange(len(self._distribution)), + p=tuple(self._distribution.values())) + return list(self._distribution.keys())[index] + + +class ParameterAlias: + """Indicates that a parameter should alias the value of another parameter. + + When used in conjunction with a custom distribution, this allows fuzzed + tensors to represent a broader range of behaviors. For example, the + following sometimes produces Tensors which broadcast: + + Fuzzer( + parameters=[ + FuzzedParameter("x_len", 4, 1024, distribution="uniform"), + + # `y` will either be size one, or match the size of `x`. + FuzzedParameter("y_len", distribution={ + 0.5: 1, + 0.5: ParameterAlias("x_len") + }), + ], + tensors=[ + FuzzedTensor("x", size=("x_len",)), + FuzzedTensor("y", size=("y_len",)), + ], + ) + + Chains of alias' are allowed, but may not contain cycles. + """ + def __init__(self, alias_to) -> None: + self.alias_to = alias_to + + def __repr__(self) -> str: + return f"ParameterAlias[alias_to: {self.alias_to}]" + + +def dtype_size(dtype): + if dtype == torch.bool: + return 1 + if dtype.is_floating_point or dtype.is_complex: + return int(torch.finfo(dtype).bits / 8) + return int(torch.iinfo(dtype).bits / 8) + + +def prod(values, base=1): + """np.prod can overflow, so for sizes the product should be done in Python. + + Even though np.prod type promotes to int64, it can still overflow in which + case the negative value will pass the size check and OOM when attempting to + actually allocate the Tensor. + """ + return functools.reduce(lambda x, y: int(x) * int(y), values, base) + + +class FuzzedTensor: + def __init__( + self, + name: str, + size: tuple[str | int, ...], + steps: tuple[str | int, ...] | None = None, + probability_contiguous: float = 0.5, + min_elements: int | None = None, + max_elements: int | None = None, + max_allocation_bytes: int | None = None, + dim_parameter: str | None = None, + roll_parameter: str | None = None, + dtype=torch.float32, + cuda=False, + tensor_constructor: Callable | None = None + ) -> None: + """ + Args: + name: + A string identifier for the generated Tensor. + size: + A tuple of integers or strings specifying the size of the generated + Tensor. String values will replaced with a concrete int during the + generation process, while ints are simply passed as literals. + steps: + An optional tuple with the same length as `size`. This indicates + that a larger Tensor should be allocated, and then sliced to + produce the generated Tensor. For instance, if size is (4, 8) + and steps is (1, 4), then a tensor `t` of size (4, 32) will be + created and then `t[:, ::4]` will be used. (Allowing one to test + Tensors with strided memory.) + probability_contiguous: + A number between zero and one representing the chance that the + generated Tensor has a contiguous memory layout. This is achieved by + randomly permuting the shape of a Tensor, calling `.contiguous()`, + and then permuting back. This is applied before `steps`, which can + also cause a Tensor to be non-contiguous. + min_elements: + The minimum number of parameters that this Tensor must have for a + set of parameters to be valid. (Otherwise they are resampled.) + max_elements: + Like `min_elements`, but setting an upper bound. + max_allocation_bytes: + Like `max_elements`, but for the size of Tensor that must be + allocated prior to slicing for `steps` (if applicable). For + example, a FloatTensor with size (1024, 1024) and steps (4, 4) + would have 1M elements, but would require a 64 MB allocation. + dim_parameter: + The length of `size` and `steps` will be truncated to this value. + This allows Tensors of varying dimensions to be generated by the + Fuzzer. + dtype: + The PyTorch dtype of the generated Tensor. + cuda: + Whether to place the Tensor on a GPU. + tensor_constructor: + Callable which will be used instead of the default Tensor + construction method. This allows the author to enforce properties + of the Tensor (e.g. it can only have certain values). The dtype and + concrete shape of the Tensor to be created will be passed, and + concrete values of all parameters will be passed as kwargs. Note + that transformations to the result (permuting, slicing) will be + performed by the Fuzzer; the tensor_constructor is only responsible + for creating an appropriately sized Tensor. + """ + self._name = name + self._size = size + self._steps = steps + self._probability_contiguous = probability_contiguous + self._min_elements = min_elements + self._max_elements = max_elements + self._max_allocation_bytes = max_allocation_bytes + self._dim_parameter = dim_parameter + self._dtype = dtype + self._cuda = cuda + self._tensor_constructor = tensor_constructor + + @property + def name(self): + return self._name + + @staticmethod + def default_tensor_constructor(size, dtype, **kwargs): + if dtype.is_floating_point or dtype.is_complex: + return torch.rand(size=size, dtype=dtype, device="cpu") + else: + return torch.randint(1, 127, size=size, dtype=dtype, device="cpu") + + def _make_tensor(self, params, state): + import numpy as np + size, steps, allocation_size = self._get_size_and_steps(params) + constructor = ( + self._tensor_constructor or + self.default_tensor_constructor + ) + + raw_tensor = constructor(size=allocation_size, dtype=self._dtype, **params) + if self._cuda: + raw_tensor = raw_tensor.cuda() + + # Randomly permute the Tensor and call `.contiguous()` to force re-ordering + # of the memory, and then permute it back to the original shape. + dim = len(size) + order = np.arange(dim) + if state.rand() > self._probability_contiguous: + while dim > 1 and np.all(order == np.arange(dim)): + order = state.permutation(raw_tensor.dim()) + + raw_tensor = raw_tensor.permute(tuple(order)).contiguous() + raw_tensor = raw_tensor.permute(tuple(np.argsort(order))) + + slices = [slice(0, size * step, step) for size, step in zip(size, steps, strict=True)] + tensor = raw_tensor[tuple(slices)] + + properties = { + "numel": int(tensor.numel()), + "order": order, + "steps": steps, + "is_contiguous": tensor.is_contiguous(), + "dtype": str(self._dtype), + } + + return tensor, properties + + def _get_size_and_steps(self, params): + dim = ( + params[self._dim_parameter] + if self._dim_parameter is not None + else len(self._size) + ) + + def resolve(values, dim): + """Resolve values into concrete integers.""" + values = tuple(params.get(i, i) for i in values) + if len(values) > dim: + values = values[:dim] + if len(values) < dim: + values = values + tuple(1 for _ in range(dim - len(values))) + return values + + size = resolve(self._size, dim) + steps = resolve(self._steps or (), dim) + allocation_size = tuple(size_i * step_i for size_i, step_i in zip(size, steps, strict=True)) + return size, steps, allocation_size + + def satisfies_constraints(self, params) -> bool: + size, _, allocation_size = self._get_size_and_steps(params) + # Product is computed in Python to avoid integer overflow. + num_elements = prod(size) + if num_elements < 0: + raise AssertionError("Computed number of elements is negative") + + allocation_bytes = prod(allocation_size, base=dtype_size(self._dtype)) + + def nullable_greater(left, right): + if left is None or right is None: + return False + return left > right + + return not any(( + nullable_greater(num_elements, self._max_elements), + nullable_greater(self._min_elements, num_elements), + nullable_greater(allocation_bytes, self._max_allocation_bytes), + )) + + +class Fuzzer: + def __init__( + self, + parameters: list[FuzzedParameter | list[FuzzedParameter]], + tensors: list[FuzzedTensor | list[FuzzedTensor]], + constraints: list[Callable] | None = None, + seed: int | None = None + ) -> None: + """ + Args: + parameters: + List of FuzzedParameters which provide specifications + for generated parameters. Iterable elements will be + unpacked, though arbitrary nested structures will not. + tensors: + List of FuzzedTensors which define the Tensors which + will be created each step based on the parameters for + that step. Iterable elements will be unpacked, though + arbitrary nested structures will not. + constraints: + List of callables. They will be called with params + as kwargs, and if any of them return False the current + set of parameters will be rejected. + seed: + Seed for the RandomState used by the Fuzzer. This will + also be used to set the PyTorch random seed so that random + ops will create reproducible Tensors. + """ + import numpy as np + if seed is None: + seed = int(np.random.RandomState().randint(0, 2 ** 32 - 1, dtype=np.int64)) + self._seed = seed + self._parameters = Fuzzer._unpack(parameters, FuzzedParameter) + self._tensors = Fuzzer._unpack(tensors, FuzzedTensor) + self._constraints = constraints or () + + p_names = {p.name for p in self._parameters} + t_names = {t.name for t in self._tensors} + name_overlap = p_names.intersection(t_names) + if name_overlap: + raise ValueError(f"Duplicate names in parameters and tensors: {name_overlap}") + + self._rejections = 0 + self._total_generated = 0 + + @staticmethod + def _unpack(values, cls): + return tuple(it.chain.from_iterable( + [[i] if isinstance(i, cls) else i for i in values] + )) + + def take(self, n): + import numpy as np + state = np.random.RandomState(self._seed) + torch.manual_seed(state.randint(low=0, high=2 ** 63, dtype=np.int64)) + for _ in range(n): + params = self._generate(state) + tensors = {} + tensor_properties = {} + for t in self._tensors: + tensor, properties = t._make_tensor(params, state) + tensors[t.name] = tensor + tensor_properties[t.name] = properties + yield tensors, tensor_properties, params + + @property + def rejection_rate(self): + if not self._total_generated: + return 0. + return self._rejections / self._total_generated + + def _generate(self, state): + strict_params: dict[str, float | int | ParameterAlias] = {} + for _ in range(1000): + candidate_params: dict[str, float | int | ParameterAlias] = {} + for p in self._parameters: + if p.strict: + if p.name in strict_params: + candidate_params[p.name] = strict_params[p.name] + else: + candidate_params[p.name] = p.sample(state) + strict_params[p.name] = candidate_params[p.name] + else: + candidate_params[p.name] = p.sample(state) + + candidate_params = self._resolve_aliases(candidate_params) + + self._total_generated += 1 + if not all(f(candidate_params) for f in self._constraints): + self._rejections += 1 + continue + + if not all(t.satisfies_constraints(candidate_params) for t in self._tensors): + self._rejections += 1 + continue + + return candidate_params + raise ValueError("Failed to generate a set of valid parameters.") + + @staticmethod + def _resolve_aliases(params): + params = dict(params) + alias_count = sum(isinstance(v, ParameterAlias) for v in params.values()) + + keys = list(params.keys()) + while alias_count: + for k in keys: + v = params[k] + if isinstance(v, ParameterAlias): + params[k] = params[v.alias_to] + alias_count_new = sum(isinstance(v, ParameterAlias) for v in params.values()) + if alias_count == alias_count_new: + raise ValueError(f"ParameterAlias cycle detected\n{params}") + + alias_count = alias_count_new + + return params diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/sparse_fuzzer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/sparse_fuzzer.py new file mode 100644 index 0000000000000000000000000000000000000000..cf38c11240b278e07e2b53e4e2306219c4f6a9fb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/sparse_fuzzer.py @@ -0,0 +1,122 @@ +# mypy: allow-untyped-defs +from numbers import Number +import torch +from torch.utils.benchmark import FuzzedTensor +import math + +class FuzzedSparseTensor(FuzzedTensor): + def __init__( + self, + name: str, + size: tuple[str | int, ...], + min_elements: int | None = None, + max_elements: int | None = None, + dim_parameter: str | None = None, + sparse_dim: str | None = None, + nnz: str | None = None, + density: str | None = None, + coalesced: str | None = None, + dtype=torch.float32, + cuda=False + ) -> None: + """ + Args: + name: + A string identifier for the generated Tensor. + size: + A tuple of integers or strings specifying the size of the generated + Tensor. String values will replaced with a concrete int during the + generation process, while ints are simply passed as literals. + min_elements: + The minimum number of parameters that this Tensor must have for a + set of parameters to be valid. (Otherwise they are resampled.) + max_elements: + Like `min_elements`, but setting an upper bound. + dim_parameter: + The length of `size` will be truncated to this value. + This allows Tensors of varying dimensions to be generated by the + Fuzzer. + sparse_dim: + The number of sparse dimensions in a sparse tensor. + density: + This value allows tensors of varying sparsities to be generated by the Fuzzer. + coalesced: + The sparse tensor format permits uncoalesced sparse tensors, + where there may be duplicate coordinates in the indices. + dtype: + The PyTorch dtype of the generated Tensor. + cuda: + Whether to place the Tensor on a GPU. + """ + super().__init__(name=name, size=size, min_elements=min_elements, + max_elements=max_elements, dim_parameter=dim_parameter, dtype=dtype, cuda=cuda) + self._density = density + self._coalesced = coalesced + self._sparse_dim = sparse_dim + + @staticmethod + def sparse_tensor_constructor(size, dtype, sparse_dim, nnz, is_coalesced): + """sparse_tensor_constructor creates a sparse tensor with coo format. + + Note that when `is_coalesced` is False, the number of elements is doubled but the number of indices + represents the same amount of number of non zeros `nnz`, i.e, this is virtually the same tensor + with the same sparsity pattern. Moreover, most of the sparse operation will use coalesce() method + and what we want here is to get a sparse tensor with the same `nnz` even if this is coalesced or not. + + In the other hand when `is_coalesced` is True the number of elements is reduced in the coalescing process + by an unclear amount however the probability to generate duplicates indices are low for most of the cases. + This decision was taken on purpose to maintain the construction cost as low as possible. + """ + if isinstance(size, Number): + size = [size] * sparse_dim + if all(size[d] <= 0 for d in range(sparse_dim)) and nnz != 0: + raise AssertionError('invalid arguments') + v_size = [nnz] + list(size[sparse_dim:]) + if dtype.is_floating_point: + v = torch.rand(size=v_size, dtype=dtype, device="cpu") + else: + v = torch.randint(1, 127, size=v_size, dtype=dtype, device="cpu") + + i = torch.rand(sparse_dim, nnz, device="cpu") + i.mul_(torch.tensor(size[:sparse_dim]).unsqueeze(1).to(i)) + i = i.to(torch.long) + + if not is_coalesced: + v = torch.cat([v, torch.randn_like(v)], 0) + i = torch.cat([i, i], 1) + + x = torch.sparse_coo_tensor(i, v, torch.Size(size)) + if is_coalesced: + x = x.coalesce() + return x + + def _make_tensor(self, params, state): + size, _, _ = self._get_size_and_steps(params) + density = params['density'] + nnz = math.ceil(sum(size) * density) + if nnz > sum(size): + raise AssertionError('nnz cannot exceed total number of elements') + + is_coalesced = params['coalesced'] + sparse_dim = params['sparse_dim'] if self._sparse_dim else len(size) + sparse_dim = min(sparse_dim, len(size)) + tensor = self.sparse_tensor_constructor(size, self._dtype, sparse_dim, nnz, is_coalesced) + + if self._cuda: + tensor = tensor.cuda() + sparse_dim = tensor.sparse_dim() + dense_dim = tensor.dense_dim() + is_hybrid = len(size[sparse_dim:]) > 0 + + properties = { + "numel": int(tensor.numel()), + "shape": tensor.size(), + "is_coalesced": tensor.is_coalesced(), + "density": density, + "sparsity": 1.0 - density, + "sparse_dim": sparse_dim, + "dense_dim": dense_dim, + "is_hybrid": is_hybrid, + "dtype": str(self._dtype), + } + return tensor, properties diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/timeit_template.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/timeit_template.cpp new file mode 100644 index 0000000000000000000000000000000000000000..30b6f79c0b5aebca676035ac0b7c08cfce639b23 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/timeit_template.cpp @@ -0,0 +1,43 @@ +/* C++ template for Timer.timeit + +This template will be consumed by `cpp_jit.py`, and will replace: + `GLOBAL_SETUP_TEMPLATE_LOCATION`, + `SETUP_TEMPLATE_LOCATION` + and + `STMT_TEMPLATE_LOCATION` +sections with user provided statements. +*/ +#include + +#include +#include +#include +#include + +// Global setup. (e.g. #includes) +// GLOBAL_SETUP_TEMPLATE_LOCATION + +double timeit(int n) { + pybind11::gil_scoped_release no_gil; + + // Setup + // SETUP_TEMPLATE_LOCATION + + { + // Warmup + // STMT_TEMPLATE_LOCATION + } + + // Main loop + auto start_time = std::chrono::high_resolution_clock::now(); + for (const auto loop_idx : c10::irange(n)) { + (void)loop_idx; + // STMT_TEMPLATE_LOCATION + } + auto end_time = std::chrono::high_resolution_clock::now(); + return std::chrono::duration(end_time - start_time).count(); +} + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { + m.def("timeit", &timeit); +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/timer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..9686840fa10a1610bdfd4e12b590497af7a4d74d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/timer.py @@ -0,0 +1,533 @@ +"""Timer class based on the timeit.Timer class, but torch aware.""" +import enum +import timeit +import textwrap +from typing import overload, Any, NoReturn +from collections.abc import Callable + +import torch +from torch.utils.benchmark.utils import common, cpp_jit +from torch.utils.benchmark.utils._stubs import TimerClass, TimeitModuleType +from torch.utils.benchmark.utils.valgrind_wrapper import timer_interface as valgrind_timer_interface + + +__all__ = ["Timer", "timer", "Language"] + + +if torch.accelerator.is_available(): + def timer() -> float: + torch.accelerator.synchronize() + return timeit.default_timer() +else: + timer = timeit.default_timer + + +class Language(enum.Enum): + PYTHON = 0 + CPP = 1 + + +class CPPTimer: + def __init__( + self, + stmt: str, + setup: str, + global_setup: str, + timer: Callable[[], float], + globals: dict[str, Any], + ) -> None: + if timer is not timeit.default_timer: + raise NotImplementedError( + "PyTorch was built with accelerators and an accelerator is present; however " + "Timer does not yet support accelerator measurements. If your " + "code is CPU only, pass `timer=timeit.default_timer` to the " + "Timer's constructor to indicate this. (Note that this will " + "produce incorrect results if an accelerator is in fact used, as " + "Timer will not synchronize the accelerator.)" + ) + + if globals: + raise ValueError("C++ timing does not support globals.") + + self._stmt: str = textwrap.dedent(stmt) + self._setup: str = textwrap.dedent(setup) + self._global_setup: str = textwrap.dedent(global_setup) + self._timeit_module: TimeitModuleType | None = None + + def timeit(self, number: int) -> float: + if self._timeit_module is None: + self._timeit_module = cpp_jit.compile_timeit_template( + stmt=self._stmt, + setup=self._setup, + global_setup=self._global_setup, + ) + + return self._timeit_module.timeit(number) + + +class Timer: + """Helper class for measuring execution time of PyTorch statements. + + For a full tutorial on how to use this class, see: + https://pytorch.org/tutorials/recipes/recipes/benchmark.html + + The PyTorch Timer is based on `timeit.Timer` (and in fact uses + `timeit.Timer` internally), but with several key differences: + + 1) Runtime aware: + Timer will perform warmups (important as some elements of PyTorch are + lazily initialized), set threadpool size so that comparisons are + apples-to-apples, and synchronize asynchronous accelerator functions when + necessary. + + 2) Focus on replicates: + When measuring code, and particularly complex kernels / models, + run-to-run variation is a significant confounding factor. It is + expected that all measurements should include replicates to quantify + noise and allow median computation, which is more robust than mean. + To that effect, this class deviates from the `timeit` API by + conceptually merging `timeit.Timer.repeat` and `timeit.Timer.autorange`. + (Exact algorithms are discussed in method docstrings.) The `timeit` + method is replicated for cases where an adaptive strategy is not + desired. + + 3) Optional metadata: + When defining a Timer, one can optionally specify `label`, `sub_label`, + `description`, and `env`. (Defined later) These fields are included in + the representation of result object and by the `Compare` class to group + and display results for comparison. + + 4) Instruction counts + In addition to wall times, Timer can run a statement under Callgrind + and report instructions executed. + + Directly analogous to `timeit.Timer` constructor arguments: + + `stmt`, `setup`, `timer`, `globals` + + PyTorch Timer specific constructor arguments: + + `label`, `sub_label`, `description`, `env`, `num_threads` + + Args: + stmt: Code snippet to be run in a loop and timed. + + setup: Optional setup code. Used to define variables used in `stmt` + + global_setup: (C++ only) + Code which is placed at the top level of the file for things like + `#include` statements. + + timer: + Callable which returns the current time. If PyTorch was built + without accelerators or there is no accelerator present, this defaults to + `timeit.default_timer`; otherwise it will synchronize accelerators before + measuring the time. + + globals: + A dict which defines the global variables when `stmt` is being + executed. This is the other method for providing variables which + `stmt` needs. + + label: + String which summarizes `stmt`. For instance, if `stmt` is + "torch.nn.functional.relu(torch.add(x, 1, out=out))" + one might set label to "ReLU(x + 1)" to improve readability. + + sub_label: + Provide supplemental information to disambiguate measurements + with identical stmt or label. For instance, in our example + above sub_label might be "float" or "int", so that it is easy + to differentiate: + "ReLU(x + 1): (float)" + + "ReLU(x + 1): (int)" + when printing Measurements or summarizing using `Compare`. + + description: + String to distinguish measurements with identical label and + sub_label. The principal use of `description` is to signal to + `Compare` the columns of data. For instance one might set it + based on the input size to create a table of the form: :: + + | n=1 | n=4 | ... + ------------- ... + ReLU(x + 1): (float) | ... | ... | ... + ReLU(x + 1): (int) | ... | ... | ... + + + using `Compare`. It is also included when printing a Measurement. + + env: + This tag indicates that otherwise identical tasks were run in + different environments, and are therefore not equivalent, for + instance when A/B testing a change to a kernel. `Compare` will + treat Measurements with different `env` specification as distinct + when merging replicate runs. + + num_threads: + The size of the PyTorch threadpool when executing `stmt`. Single + threaded performance is important as both a key inference workload + and a good indicator of intrinsic algorithmic efficiency, so the + default is set to one. This is in contrast to the default PyTorch + threadpool size which tries to utilize all cores. + """ + + _timer_cls: type[TimerClass] = timeit.Timer + + def __init__( + self, + stmt: str = "pass", + setup: str = "pass", + global_setup: str = "", + timer: Callable[[], float] = timer, + globals: dict[str, Any] | None = None, + label: str | None = None, + sub_label: str | None = None, + description: str | None = None, + env: str | None = None, + num_threads: int = 1, + language: Language | str = Language.PYTHON, + ) -> None: + if not isinstance(stmt, str): + raise ValueError("Currently only a `str` stmt is supported.") + + # We copy `globals` to prevent mutations from leaking. + # (For instance, `eval` adds the `__builtins__` key) + self._globals = dict(globals or {}) + + timer_kwargs = {} + if language in (Language.PYTHON, "py", "python"): + # Include `torch` if not specified as a convenience feature. + self._globals.setdefault("torch", torch) + self._language: Language = Language.PYTHON + if global_setup: + raise ValueError( + f"global_setup is C++ only, got `{global_setup}`. Most " + "likely this code can simply be moved to `setup`." + ) + + elif language in (Language.CPP, "cpp", "c++"): + if self._timer_cls is not timeit.Timer: + raise AssertionError("_timer_cls has already been swapped.") + self._timer_cls = CPPTimer + setup = ("" if setup == "pass" else setup) + self._language = Language.CPP + timer_kwargs["global_setup"] = global_setup + + else: + raise ValueError(f"Invalid language `{language}`.") + + # Convenience adjustment so that multi-line code snippets defined in + # functions do not IndentationError (Python) or look odd (C++). The + # leading newline removal is for the initial newline that appears when + # defining block strings. For instance: + # textwrap.dedent(""" + # print("This is a stmt") + # """) + # produces '\nprint("This is a stmt")\n'. + # + # Stripping this down to 'print("This is a stmt")' doesn't change + # what gets executed, but it makes __repr__'s nicer. + stmt = textwrap.dedent(stmt) + stmt = (stmt[1:] if stmt and stmt[0] == "\n" else stmt).rstrip() + setup = textwrap.dedent(setup) + setup = (setup[1:] if setup and setup[0] == "\n" else setup).rstrip() + + + self._timer = self._timer_cls( + stmt=stmt, + setup=setup, + timer=timer, + globals=valgrind_timer_interface.CopyIfCallgrind.unwrap_all(self._globals), + **timer_kwargs, + ) + self._task_spec = common.TaskSpec( + stmt=stmt, + setup=setup, + global_setup=global_setup, + label=label, + sub_label=sub_label, + description=description, + env=env, + num_threads=num_threads, + ) + + def _timeit(self, number: int) -> float: + # Even calling a timer in C++ takes ~50 ns, so no real operation should + # take less than 1 ns. (And this prevents divide by zero errors.) + return max(self._timer.timeit(number), 1e-9) + + def timeit(self, number: int = 1000000) -> common.Measurement: + """Mirrors the semantics of timeit.Timer.timeit(). + + Execute the main statement (`stmt`) `number` times. + https://docs.python.org/3/library/timeit.html#timeit.Timer.timeit + """ + with common.set_torch_threads(self._task_spec.num_threads): + # Warmup + self._timeit(number=max(int(number // 100), 2)) + + return common.Measurement( + number_per_run=number, + raw_times=[self._timeit(number=number)], + task_spec=self._task_spec + ) + + def repeat(self, repeat: int = -1, number: int = -1) -> None: + raise NotImplementedError("See `Timer.blocked_autorange.`") + + def autorange(self, callback: Callable[[int, float], NoReturn] | None = None) -> None: + raise NotImplementedError("See `Timer.blocked_autorange.`") + + def _threaded_measurement_loop( + self, + number: int, + time_hook: Callable[[], float], + stop_hook: Callable[[list[float]], bool], + min_run_time: float, + max_run_time: float | None = None, + callback: Callable[[int, float], NoReturn] | None = None + ) -> list[float]: + total_time = 0.0 + can_stop = False + times: list[float] = [] + with common.set_torch_threads(self._task_spec.num_threads): + while (total_time < min_run_time) or (not can_stop): + time_spent = time_hook() + times.append(time_spent) + total_time += time_spent + if callback: + callback(number, time_spent) + can_stop = stop_hook(times) + if max_run_time and total_time > max_run_time: + break + return times + + def _estimate_block_size(self, min_run_time: float) -> int: + with common.set_torch_threads(self._task_spec.num_threads): + # Estimate the block size needed for measurement to be negligible + # compared to the inner loop. This also serves as a warmup. + overhead = torch.tensor([self._timeit(0) for _ in range(5)]).median().item() + number = 1 + while True: + time_taken = self._timeit(number) + relative_overhead = overhead / time_taken + if relative_overhead <= 1e-4 and time_taken >= min_run_time / 1000: + break + if time_taken > min_run_time: + break + # Avoid overflow in C++ pybind11 interface + if number * 10 > 2147483647: + break + number *= 10 + return number + + def blocked_autorange( + self, + callback: Callable[[int, float], NoReturn] | None = None, + min_run_time: float = 0.2, + ) -> common.Measurement: + """Measure many replicates while keeping timer overhead to a minimum. + + At a high level, blocked_autorange executes the following pseudo-code:: + + `setup` + + total_time = 0 + while total_time < min_run_time + start = timer() + for _ in range(block_size): + `stmt` + total_time += (timer() - start) + + Note the variable `block_size` in the inner loop. The choice of block + size is important to measurement quality, and must balance two + competing objectives: + + 1) A small block size results in more replicates and generally + better statistics. + + 2) A large block size better amortizes the cost of `timer` + invocation, and results in a less biased measurement. This is + important because accelerator synchronization time is non-trivial + (order single to low double digit microseconds) and would + otherwise bias the measurement. + + blocked_autorange sets block_size by running a warmup period, + increasing block size until timer overhead is less than 0.1% of + the overall computation. This value is then used for the main + measurement loop. + + Returns: + A `Measurement` object that contains measured runtimes and + repetition counts, and can be used to compute statistics. + (mean, median, etc.) + """ + number = self._estimate_block_size(min_run_time) + + def time_hook() -> float: + return self._timeit(number) + + def stop_hook(times: list[float]) -> bool: + return True + + times = self._threaded_measurement_loop( + number, time_hook, stop_hook, + min_run_time=min_run_time, + callback=callback) + + return common.Measurement( + number_per_run=number, + raw_times=times, + task_spec=self._task_spec + ) + + def adaptive_autorange( + self, + threshold: float = 0.1, + *, + min_run_time: float = 0.01, + max_run_time: float = 10.0, + callback: Callable[[int, float], NoReturn] | None = None, + ) -> common.Measurement: + """Similar to `blocked_autorange` but also checks for variablility in measurements + and repeats until iqr/median is smaller than `threshold` or `max_run_time` is reached. + + + At a high level, adaptive_autorange executes the following pseudo-code:: + + `setup` + + times = [] + while times.sum < max_run_time + start = timer() + for _ in range(block_size): + `stmt` + times.append(timer() - start) + + enough_data = len(times)>3 and times.sum > min_run_time + small_iqr=times.iqr/times.mean float: + return self._timeit(number) + + def stop_hook(times: list[float]) -> bool: + if len(times) > 3: + return common.Measurement( + number_per_run=number, + raw_times=times, + task_spec=self._task_spec + ).meets_confidence(threshold=threshold) + return False + times = self._threaded_measurement_loop( + number, time_hook, stop_hook, min_run_time, max_run_time, callback=callback) + + return common.Measurement( + number_per_run=number, + raw_times=times, + task_spec=self._task_spec + ) + + @overload + def collect_callgrind( + self, + number: int, + *, + repeats: None, + collect_baseline: bool, + retain_out_file: bool, + ) -> valgrind_timer_interface.CallgrindStats: + ... + + @overload + def collect_callgrind( + self, + number: int, + *, + repeats: int, + collect_baseline: bool, + retain_out_file: bool, + ) -> tuple[valgrind_timer_interface.CallgrindStats, ...]: + ... + + def collect_callgrind( + self, + number: int = 100, + *, + repeats: int | None = None, + collect_baseline: bool = True, + retain_out_file: bool = False, + ) -> Any: + """Collect instruction counts using Callgrind. + + Unlike wall times, instruction counts are deterministic + (modulo non-determinism in the program itself and small amounts of + jitter from the Python interpreter.) This makes them ideal for detailed + performance analysis. This method runs `stmt` in a separate process + so that Valgrind can instrument the program. Performance is severely + degraded due to the instrumentation, however this is ameliorated by + the fact that a small number of iterations is generally sufficient to + obtain good measurements. + + In order to use this method `valgrind`, `callgrind_control`, and + `callgrind_annotate` must be installed. + + Because there is a process boundary between the caller (this process) + and the `stmt` execution, `globals` cannot contain arbitrary in-memory + data structures. (Unlike timing methods) Instead, globals are + restricted to builtins, `nn.Modules`'s, and TorchScripted functions/modules + to reduce the surprise factor from serialization and subsequent + deserialization. The `GlobalsBridge` class provides more detail on this + subject. Take particular care with nn.Modules: they rely on pickle and + you may need to add an import to `setup` for them to transfer properly. + + By default, a profile for an empty statement will be collected and + cached to indicate how many instructions are from the Python loop which + drives `stmt`. + + Returns: + A `CallgrindStats` object which provides instruction counts and + some basic facilities for analyzing and manipulating results. + """ + if not isinstance(self._task_spec.stmt, str): + raise ValueError("`collect_callgrind` currently only supports string `stmt`") + + if repeats is not None and repeats < 1: + raise ValueError("If specified, `repeats` must be >= 1") + + # Check that the statement is valid. It doesn't guarantee success, but it's much + # simpler and quicker to raise an exception for a faulty `stmt` or `setup` in + # the parent process rather than the valgrind subprocess. + self._timeit(1) + is_python = (self._language == Language.PYTHON) + if not is_python and self._globals: + raise AssertionError("_timer globals are only supported for Python timers") + result = valgrind_timer_interface.wrapper_singleton().collect_callgrind( + task_spec=self._task_spec, + globals=self._globals, + number=number, + repeats=repeats or 1, + collect_baseline=collect_baseline and is_python, + is_python=is_python, + retain_out_file=retain_out_file, + ) + + return (result[0] if repeats is None else result) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/callgrind.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/callgrind.h new file mode 100644 index 0000000000000000000000000000000000000000..f078cc82b95daf94d2bea51f1e1b1a8c12daea23 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/callgrind.h @@ -0,0 +1,129 @@ + +/* + ---------------------------------------------------------------- + + Notice that the following BSD-style license applies to this one + file (callgrind.h) only. The rest of Valgrind is licensed under the + terms of the GNU General Public License, version 2, unless + otherwise indicated. See the COPYING file in the source + distribution for details. + + ---------------------------------------------------------------- + + This file is part of callgrind, a valgrind tool for cache simulation + and call tree tracing. + + Copyright (C) 2003-2017 Josef Weidendorfer. All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions + are met: + + 1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + 2. The origin of this software must not be misrepresented; you must + not claim that you wrote the original software. If you use this + software in a product, an acknowledgment in the product + documentation would be appreciated but is not required. + + 3. Altered source versions must be plainly marked as such, and must + not be misrepresented as being the original software. + + 4. The name of the author may not be used to endorse or promote + products derived from this software without specific prior written + permission. + + THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS + OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED + WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY + DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE + GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, + WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING + NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + ---------------------------------------------------------------- + + Notice that the above BSD-style license applies to this one file + (callgrind.h) only. The entire rest of Valgrind is licensed under + the terms of the GNU General Public License, version 2. See the + COPYING file in the source distribution for details. + + ---------------------------------------------------------------- +*/ + +#ifndef __CALLGRIND_H +#define __CALLGRIND_H + +#include "valgrind.h" + +/* !! ABIWARNING !! ABIWARNING !! ABIWARNING !! ABIWARNING !! + This enum comprises an ABI exported by Valgrind to programs + which use client requests. DO NOT CHANGE THE ORDER OF THESE + ENTRIES, NOR DELETE ANY -- add new ones at the end. + + The identification ('C','T') for Callgrind has historical + reasons: it was called "Calltree" before. Besides, ('C','G') would + clash with cachegrind. + */ + +typedef + enum { + VG_USERREQ__DUMP_STATS = VG_USERREQ_TOOL_BASE('C','T'), + VG_USERREQ__ZERO_STATS, + VG_USERREQ__TOGGLE_COLLECT, + VG_USERREQ__DUMP_STATS_AT, + VG_USERREQ__START_INSTRUMENTATION, + VG_USERREQ__STOP_INSTRUMENTATION + } Vg_CallgrindClientRequest; + +/* Dump current state of cost centers, and zero them afterwards */ +#define CALLGRIND_DUMP_STATS \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__DUMP_STATS, \ + 0, 0, 0, 0, 0) + +/* Dump current state of cost centers, and zero them afterwards. + The argument is appended to a string stating the reason which triggered + the dump. This string is written as a description field into the + profile data dump. */ +#define CALLGRIND_DUMP_STATS_AT(pos_str) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__DUMP_STATS_AT, \ + pos_str, 0, 0, 0, 0) + +/* Zero cost centers */ +#define CALLGRIND_ZERO_STATS \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__ZERO_STATS, \ + 0, 0, 0, 0, 0) + +/* Toggles collection state. + The collection state specifies whether the happening of events + should be noted or if they are to be ignored. Events are noted + by increment of counters in a cost center */ +#define CALLGRIND_TOGGLE_COLLECT \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__TOGGLE_COLLECT, \ + 0, 0, 0, 0, 0) + +/* Start full callgrind instrumentation if not already switched on. + When cache simulation is done, it will flush the simulated cache; + this will lead to an artificial cache warmup phase afterwards with + cache misses which would not have happened in reality. */ +#define CALLGRIND_START_INSTRUMENTATION \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__START_INSTRUMENTATION, \ + 0, 0, 0, 0, 0) + +/* Stop full callgrind instrumentation if not already switched off. + This flushes Valgrinds translation cache, and does no additional + instrumentation afterwards, which effectivly will run at the same + speed as the "none" tool (ie. at minimal slowdown). + Use this to bypass Callgrind aggregation for uninteresting code parts. + To start Callgrind in this mode to ignore the setup phase, use + the option "--instr-atstart=no". */ +#define CALLGRIND_STOP_INSTRUMENTATION \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__STOP_INSTRUMENTATION, \ + 0, 0, 0, 0, 0) + +#endif /* __CALLGRIND_H */ diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/compat_bindings.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/compat_bindings.cpp new file mode 100644 index 0000000000000000000000000000000000000000..cd41f0de092f0b1488c8945edf2af80c6f9b596c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/compat_bindings.cpp @@ -0,0 +1,35 @@ +/* Used to collect profiles of old versions of PyTorch. */ +#include +#include + +bool _valgrind_supported_platform() { +#if defined(NVALGRIND) + return false; +#else + return true; +#endif +} + +void _valgrind_toggle() { +#if defined(NVALGRIND) + TORCH_CHECK(false, "Valgrind is not supported."); +#else + CALLGRIND_TOGGLE_COLLECT; +#endif +} + +void _valgrind_toggle_and_dump_stats() { +#if defined(NVALGRIND) + TORCH_CHECK(false, "Valgrind is not supported."); +#else + // NB: See note in Module.cpp + CALLGRIND_TOGGLE_COLLECT; + CALLGRIND_DUMP_STATS; +#endif +} + +PYBIND11_MODULE(callgrind_bindings, m) { + m.def("_valgrind_supported_platform", &_valgrind_supported_platform); + m.def("_valgrind_toggle", &_valgrind_toggle); + m.def("_valgrind_toggle_and_dump_stats", &_valgrind_dump_stats); +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/timer_callgrind_template.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/timer_callgrind_template.cpp new file mode 100644 index 0000000000000000000000000000000000000000..587685c7df7445b299c35462307f47cf6012a00d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/timer_callgrind_template.cpp @@ -0,0 +1,68 @@ +/* C++ template for Timer.collect_callgrind + +This template will be consumed by `cpp_jit.py`, and will replace: + `GLOBAL_SETUP_TEMPLATE_LOCATION`, + `SETUP_TEMPLATE_LOCATION` + and + `STMT_TEMPLATE_LOCATION` +sections with user provided statements. +*/ + +#include +#include +#include + +#include + +// Global setup. (e.g. #includes) +// GLOBAL_SETUP_TEMPLATE_LOCATION + +#if defined(NVALGRIND) +static_assert(false); +#endif + +int main(int argc, char* argv[]) { + // This file should only be called inside of `Timer`, so we can adopt a + // very simple and rigid argument parsing scheme. + TORCH_CHECK(argc == 9); + TORCH_CHECK(std::string(argv[1]) == "--number"); + auto number = std::stoi(argv[2]); + + TORCH_CHECK( + std::string(argv[3]) == "--number-warmup" || + std::string(argv[3]) == "--number_warmup"); + auto number_warmup = std::stoi(argv[4]); + + TORCH_CHECK(std::string(argv[5]) == "--repeats"); + auto repeats = std::stoi(argv[6]); + + TORCH_CHECK( + std::string(argv[7]) == "--number-threads" || + std::string(argv[7]) == "--number_threads"); + auto number_threads = std::stoi(argv[8]); + torch::set_num_threads(number_threads); + + // Setup + // SETUP_TEMPLATE_LOCATION + + // Warmup + for (const auto i : c10::irange(number_warmup)) { + (void)i; + // STMT_TEMPLATE_LOCATION + } + + // Main loop + for (const auto repeat : c10::irange(repeats)) { + (void)repeat; + CALLGRIND_TOGGLE_COLLECT; + + for (const auto i : c10::irange(number)) { + (void)i; + // STMT_TEMPLATE_LOCATION + } + + // NB: See note in Module.cpp + CALLGRIND_TOGGLE_COLLECT; + CALLGRIND_DUMP_STATS; + } +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py new file mode 100644 index 0000000000000000000000000000000000000000..02738aa26aa1003b9c2676bc4fd1a3fc04496232 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/timer_interface.py @@ -0,0 +1,919 @@ +"""Intermediate layer between `Timer` and `valgrind`.""" +from __future__ import annotations + +import collections +import enum +import dataclasses +import itertools as it +import os +import pickle +import re +import shutil +import subprocess +import sys +import textwrap +from typing import ( + cast, Any, NamedTuple, + TYPE_CHECKING) + +import torch +from torch.utils.benchmark.utils import common, cpp_jit +import operator + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + from torch.utils.benchmark.utils._stubs import CallgrindModuleType + + +__all__ = ["FunctionCount", "FunctionCounts", "CallgrindStats", "CopyIfCallgrind"] + + +if TYPE_CHECKING: + CompletedProcessType = subprocess.CompletedProcess[str] +else: + CompletedProcessType = subprocess.CompletedProcess + + +class FunctionCount(NamedTuple): + # TODO(#105471): Rename the count field + count: int # type: ignore[assignment] + function: str + + +@dataclasses.dataclass(repr=False, eq=False, frozen=True) +class FunctionCounts: + """Container for manipulating Callgrind results. + + It supports: + 1) Addition and subtraction to combine or diff results. + 2) Tuple-like indexing. + 3) A `denoise` function which strips CPython calls which are known to + be non-deterministic and quite noisy. + 4) Two higher order methods (`filter` and `transform`) for custom + manipulation. + """ + _data: tuple[FunctionCount, ...] + inclusive: bool + truncate_rows: bool = True + + # For normal use, torch._tensor_str.PRINT_OPTS.linewidth determines + # the print settings. This is simply to allow hermetic unit tests. + _linewidth: int | None = None + + def __iter__(self) -> Iterator[FunctionCount]: + yield from self._data + + def __len__(self) -> int: + return len(self._data) + + def __getitem__(self, item: Any) -> FunctionCount | FunctionCounts: + data: FunctionCount | tuple[FunctionCount, ...] = self._data[item] + return ( + FunctionCounts(cast(tuple[FunctionCount, ...], data), self.inclusive, truncate_rows=False) + if isinstance(data, tuple) else data + ) + + def __repr__(self) -> str: + count_len = 0 + for c, _ in self: + # Account for sign in string length. + count_len = max(count_len, len(str(c)) + int(c < 0)) + + lines = [] + linewidth = self._linewidth or torch._tensor_str.PRINT_OPTS.linewidth + fn_str_len = max(linewidth - count_len - 4, 40) + for c, fn in self: + if len(fn) > fn_str_len: + left_len = int((fn_str_len - 5) // 2) + fn = fn[:left_len] + " ... " + fn[-(fn_str_len - left_len - 5):] + lines.append(f" {c:>{count_len}} {fn}") + + if self.truncate_rows and len(lines) > 18: + lines = lines[:9] + ["...".rjust(count_len + 2)] + lines[-9:] + + if not self.inclusive: + lines.extend(["", f"Total: {self.sum()}"]) + + return "\n".join([super().__repr__()] + lines) + + def __add__( + self, + other: FunctionCounts, + ) -> FunctionCounts: + return self._merge(other, lambda c: c) + + def __sub__( + self, + other: FunctionCounts, + ) -> FunctionCounts: + return self._merge(other, operator.neg) + + def __mul__(self, other: int | float) -> FunctionCounts: + return self._from_dict({ + fn: int(c * other) for c, fn in self._data + }, self.inclusive) + + def transform(self, map_fn: Callable[[str], str]) -> FunctionCounts: + """Apply `map_fn` to all of the function names. + + This can be used to regularize function names (e.g. stripping irrelevant + parts of the file path), coalesce entries by mapping multiple functions + to the same name (in which case the counts are added together), etc. + """ + counts: collections.defaultdict[str, int] = collections.defaultdict(int) + for c, fn in self._data: + counts[map_fn(fn)] += c + + return self._from_dict(counts, self.inclusive) + + def filter(self, filter_fn: Callable[[str], bool]) -> FunctionCounts: + """Keep only the elements where `filter_fn` applied to function name returns True.""" + return FunctionCounts(tuple(i for i in self if filter_fn(i.function)), self.inclusive) + + def sum(self) -> int: + return sum(c for c, _ in self) + + def denoise(self) -> FunctionCounts: + """Remove known noisy instructions. + + Several instructions in the CPython interpreter are rather noisy. These + instructions involve unicode to dictionary lookups which Python uses to + map variable names. FunctionCounts is generally a content agnostic + container, however this is sufficiently important for obtaining + reliable results to warrant an exception.""" + return self.filter(lambda fn: "dictobject.c:lookdict_unicode" not in fn) + + def _merge( + self, + second: FunctionCounts, + merge_fn: Callable[[int], int] + ) -> FunctionCounts: + if self.inclusive != second.inclusive: + raise AssertionError("Cannot merge inclusive and exclusive counts.") + counts: collections.defaultdict[str, int] = collections.defaultdict(int) + for c, fn in self: + counts[fn] += c + + for c, fn in second: + counts[fn] += merge_fn(c) + + return self._from_dict(counts, self.inclusive) + + @staticmethod + def _from_dict(counts: dict[str, int], inclusive: bool) -> FunctionCounts: + flat_counts = (FunctionCount(c, fn) for fn, c in counts.items() if c) + return FunctionCounts(tuple(sorted(flat_counts, reverse=True)), inclusive) + + +@dataclasses.dataclass(repr=False, eq=False, frozen=True) +class CallgrindStats: + """Top level container for Callgrind results collected by Timer. + + Manipulation is generally done using the FunctionCounts class, which is + obtained by calling `CallgrindStats.stats(...)`. Several convenience + methods are provided as well; the most significant is + `CallgrindStats.as_standardized()`. + """ + task_spec: common.TaskSpec + number_per_run: int + built_with_debug_symbols: bool + baseline_inclusive_stats: FunctionCounts + baseline_exclusive_stats: FunctionCounts + stmt_inclusive_stats: FunctionCounts + stmt_exclusive_stats: FunctionCounts + stmt_callgrind_out: str | None + + def __repr__(self) -> str: + base_stats = self.baseline_exclusive_stats + output = f""" +{super().__repr__()} +{self.task_spec.summarize()} + {'':>25}All{'':>10}Noisy symbols removed + Instructions: {self.counts(denoise=False):>12}{'':>15}{self.counts(denoise=True):>12} + Baseline: {base_stats.sum():>12}{'':>15}{base_stats.denoise().sum():>12} +{self.number_per_run} runs per measurement, {self.task_spec.num_threads} thread{'s' if self.task_spec.num_threads > 1 else ''} +""".strip() + if not self.built_with_debug_symbols: + output += textwrap.dedent(""" + Warning: PyTorch was not built with debug symbols. + Source information may be limited. Rebuild with + REL_WITH_DEB_INFO=1 for more detailed results.""") + return output + + def stats(self, inclusive: bool = False) -> FunctionCounts: + """Returns detailed function counts. + + Conceptually, the FunctionCounts returned can be thought of as a tuple + of (count, path_and_function_name) tuples. + + `inclusive` matches the semantics of callgrind. If True, the counts + include instructions executed by children. `inclusive=True` is useful + for identifying hot spots in code; `inclusive=False` is useful for + reducing noise when diffing counts from two different runs. (See + CallgrindStats.delta(...) for more details) + """ + return self.stmt_inclusive_stats if inclusive else self.stmt_exclusive_stats + + def counts(self, *, denoise: bool = False) -> int: + """Returns the total number of instructions executed. + + See `FunctionCounts.denoise()` for an explanation of the `denoise` arg. + """ + stats = self.stmt_exclusive_stats + return (stats.denoise() if denoise else stats).sum() + + # FIXME: Once 3.7 is the minimum version, type annotate `other` per PEP 563 + def delta( + self, + other: CallgrindStats, + inclusive: bool = False, + ) -> FunctionCounts: + """Diff two sets of counts. + + One common reason to collect instruction counts is to determine the + the effect that a particular change will have on the number of instructions + needed to perform some unit of work. If a change increases that number, the + next logical question is "why". This generally involves looking at what part + if the code increased in instruction count. This function automates that + process so that one can easily diff counts on both an inclusive and + exclusive basis. + """ + return self.stats(inclusive=inclusive) - other.stats(inclusive=inclusive) + + def as_standardized(self) -> CallgrindStats: + """Strip library names and some prefixes from function strings. + + When comparing two different sets of instruction counts, on stumbling + block can be path prefixes. Callgrind includes the full filepath + when reporting a function (as it should). However, this can cause + issues when diffing profiles. If a key component such as Python + or PyTorch was built in separate locations in the two profiles, which + can result in something resembling:: + + 23234231 /tmp/first_build_dir/thing.c:foo(...) + 9823794 /tmp/first_build_dir/thing.c:bar(...) + ... + 53453 .../aten/src/Aten/...:function_that_actually_changed(...) + ... + -9823794 /tmp/second_build_dir/thing.c:bar(...) + -23234231 /tmp/second_build_dir/thing.c:foo(...) + + Stripping prefixes can ameliorate this issue by regularizing the + strings and causing better cancellation of equivalent call sites + when diffing. + """ + def strip(stats: FunctionCounts) -> FunctionCounts: + transforms = ( + # PyTorch may have been built in different locations. + (r"^.+build/\.\./", "build/../"), + (r"^.+/" + re.escape("build/aten/"), "build/aten/"), + + # "Python" and "Objects" come from CPython. + (r"^.+/" + re.escape("Python/"), "Python/"), + (r"^.+/" + re.escape("Objects/"), "Objects/"), + + # Strip library name. e.g. `libtorch.so` + (r"\s\[.+\]$", ""), + ) + + for before, after in transforms: + stats = stats.transform(lambda fn: re.sub(before, after, fn)) + + return stats + + return CallgrindStats( + task_spec=self.task_spec, + number_per_run=self.number_per_run, + built_with_debug_symbols=self.built_with_debug_symbols, + baseline_inclusive_stats=strip(self.baseline_inclusive_stats), + baseline_exclusive_stats=strip(self.baseline_exclusive_stats), + stmt_inclusive_stats=strip(self.stmt_inclusive_stats), + stmt_exclusive_stats=strip(self.stmt_exclusive_stats), + + # `as_standardized` will change symbol names, so the contents will + # no longer map directly to `callgrind.out` + stmt_callgrind_out=None, + ) + + +class Serialization(enum.Enum): + PICKLE = 0 + TORCH = 1 + TORCH_JIT = 2 + + +_GLOBALS_ALLOWED_TYPES: dict[Serialization, tuple[Any, ...]] = { + Serialization.PICKLE: (str, bytes, bool, int, float, complex), + Serialization.TORCH_JIT: (torch.jit.ScriptFunction, torch.jit.ScriptModule), + Serialization.TORCH: (torch.nn.Module,), +} + + +class CopyIfCallgrind: + """Signal that a global may be replaced with a deserialized copy. + + See `GlobalsBridge` for why this matters. + """ + def __init__(self, value: Any, *, setup: str | None = None) -> None: + for method, supported_types in _GLOBALS_ALLOWED_TYPES.items(): + if any(isinstance(value, t) for t in supported_types): + self._value: Any = value + self._setup: str | None = setup + self._serialization: Serialization = method + break + else: + supported_str = "\n".join([ + getattr(t, "__name__", repr(t)) + for t in it.chain(_GLOBALS_ALLOWED_TYPES.values())]) + + raise ValueError( + f"Unsupported type: {type(value)}\n" + f"`collect_callgrind` restricts globals to the following types:\n" + f"{textwrap.indent(supported_str, ' ')}" + ) + + @property + def value(self) -> Any: + return self._value + + @property + def setup(self) -> str | None: + return self._setup + + @property + def serialization(self) -> Serialization: + return self._serialization + + @staticmethod + def unwrap_all(globals: dict[str, Any]) -> dict[str, Any]: + return { + k: (v.value if isinstance(v, CopyIfCallgrind) else v) + for k, v in globals.items() + } + + +class GlobalsBridge: + """Handle the transfer of (certain) globals when collecting Callgrind statistics. + + Key takeaway: Any globals passed must be wrapped in `CopyIfCallgrind` to + work with `Timer.collect_callgrind`. + + Consider the following code snippet: + ``` + import pickle + import timeit + + class Counter: + value = 0 + + def __call__(self): + self.value += 1 + + counter = Counter() + timeit.Timer("counter()", globals={"counter": counter}).timeit(10) + print(counter.value) # 10 + + timeit.Timer( + "counter()", + globals={"counter": pickle.loads(pickle.dumps(counter))} + ).timeit(20) + print(counter.value) # Still 10 + ``` + + In the first case, `stmt` is executed using the objects in `globals`; + however, the addition of serialization and deserialization changes the + semantics and may meaningfully change behavior. + + This is a practical consideration when collecting Callgrind statistics. + Unlike `exec` based execution (which `timeit` uses under the hood) which + can share in-memory data structures with the caller, Callgrind collection + requires an entirely new process in order to run under Valgrind. This means + that any data structures used for statement execution will have to be + serialized and deserialized in the subprocess. + + In order to avoid surprising semantics from (user invisible) process + boundaries, what can be passed through `globals` is severely restricted + for `Timer.collect_callgrind`. It is expected that most setup should be + achievable (albeit perhaps less ergonomically) by passing a `setup` + string. + + There are, however, exceptions. One such class are TorchScripted functions. + Because they require a concrete file with source code it is not possible + to define them using a `setup` string. Another group are torch.nn.Modules, + whose construction can be complex and prohibitively cumbersome to coerce + into a `setup` string. Finally, most builtin types are sufficiently well + behaved and sufficiently common to warrant allowing as well. (e.g. + `globals={"n": 1}` is very convenient.) + + Fortunately, all have well defined serialization semantics. This class + is responsible for enabling the Valgrind subprocess to use elements in + `globals` so long as they are an allowed type. + + Caveats: + The user is required to acknowledge this serialization by wrapping + elements in `globals` with `CopyIfCallgrind`. + + While ScriptFunction and ScriptModule are expected to save and load + quite robustly, it is up to the user to ensure that an nn.Module can + un-pickle successfully. + + `torch.Tensor` and `np.ndarray` are deliberately excluded. The + serialization/deserialization process perturbs the representation of a + tensor in ways that could result in incorrect measurements. For example, + if a tensor lives in pinned CPU memory, this fact would not be preserved + by a dump, and that will in turn change the performance of certain CUDA + operations. + """ + + def __init__(self, globals: dict[str, Any], data_dir: str) -> None: + self._globals: dict[str, CopyIfCallgrind] = {} + self._data_dir = data_dir + if not os.path.exists(data_dir): + os.mkdir(data_dir) + + if globals.get("torch", torch) is not torch: + raise ValueError("`collect_callgrind` does not support mocking out `torch`.") + + for name, value in globals.items(): + if name in ("torch", "__builtins__"): + # Torch will be imported by the collection script, and + # __builtins__ is added by Timer. + continue + + if not isinstance(value, CopyIfCallgrind): + raise ValueError( + "`collect_callgrind` requires that globals be wrapped in " + "`CopyIfCallgrind` so that serialization is explicit." + ) + + self._globals[name] = value + + def construct(self) -> str: + load_lines = [] + for name, wrapped_value in self._globals.items(): + if wrapped_value.setup is not None: + load_lines.append(textwrap.dedent(wrapped_value.setup)) + + if wrapped_value.serialization == Serialization.PICKLE: + path = os.path.join(self._data_dir, f"{name}.pkl") + load_lines.append( + f"with open({repr(path)}, 'rb') as f:\n {name} = pickle.load(f)") + with open(path, "wb") as f: + pickle.dump(wrapped_value.value, f) + + elif wrapped_value.serialization == Serialization.TORCH: + path = os.path.join(self._data_dir, f"{name}.pt") + # TODO: Figure out if we can use torch.serialization.add_safe_globals here + # Using weights_only=False after the change in + # https://dev-discuss.pytorch.org/t/bc-breaking-change-torch-load-is-being-flipped-to-use-weights-only-true-by-default-in-the-nightlies-after-137602/2573 + load_lines.append(f"{name} = torch.load({repr(path)}, weights_only=False)") + torch.save(wrapped_value.value, path) + + elif wrapped_value.serialization == Serialization.TORCH_JIT: + path = os.path.join(self._data_dir, f"{name}.pt") + load_lines.append(f"{name} = torch.jit.load({repr(path)})") + with open(path, "wb") as f: + torch.jit.save(wrapped_value.value, f) # type: ignore[no-untyped-call] + + else: + raise NotImplementedError( + f"Unknown serialization method: {wrapped_value.serialization}") + + return "\n".join(load_lines) + + +class _ValgrindWrapper: + def __init__(self) -> None: + self._bindings_module: CallgrindModuleType | None = None + valgrind_symbols = ( + "_valgrind_supported_platform", + "_valgrind_toggle", + "_valgrind_toggle_and_dump_stats", + ) + if all(hasattr(torch._C, symbol) for symbol in valgrind_symbols): + self._supported_platform: bool = torch._C._valgrind_supported_platform() + + else: + print("Callgrind bindings are not present in `torch._C`. JIT-ing bindings.") + self._bindings_module = cpp_jit.get_compat_bindings() + if not all(hasattr(self._bindings_module, symbol) for symbol in valgrind_symbols): + raise AssertionError("JIT-compiled callgrind bindings are missing required symbols") + self._supported_platform = self._bindings_module._valgrind_supported_platform() + + self._commands_available: dict[str, bool] = {} + if self._supported_platform: + # Only bother checking on supported platforms. + for cmd in ("valgrind", "callgrind_control", "callgrind_annotate"): + self._commands_available[cmd] = not subprocess.run( + ["which", cmd], + capture_output=True, + check=False, + ).returncode + + self._build_type: str | None = None + build_search = re.search("BUILD_TYPE=(.+),", torch.__config__.show()) # type: ignore[no-untyped-call] + if build_search is not None: + self._build_type = build_search.groups()[0].split(",")[0] + + def _validate(self) -> None: + if not self._supported_platform: + raise OSError("Valgrind is not supported on this platform.") + + missing_cmds = [cmd for cmd, available in self._commands_available.items() if not available] + if missing_cmds: + raise OSError("Missing: " + ", ".join(missing_cmds)) + + def collect_callgrind( + self, + task_spec: common.TaskSpec, + globals: dict[str, Any], + *, + number: int, + repeats: int, + collect_baseline: bool, + is_python: bool, + retain_out_file: bool, + ) -> tuple[CallgrindStats, ...]: + """Collect stats, and attach a reference run which can be used to filter interpreter overhead.""" + self._validate() + if not is_python and collect_baseline: + raise AssertionError("collect_baseline is only supported for Python timers") + + *task_stats, baseline_stats = self._invoke( + task_spec=task_spec, + globals=globals, + number=number, + repeats=repeats, + collect_baseline=collect_baseline, + is_python=is_python, + retain_out_file=retain_out_file, + ) + if len(task_stats) != repeats: + raise AssertionError("Unexpected number of task stats returned from _invoke") + + return tuple( + CallgrindStats( + task_spec=task_spec, + number_per_run=number, + built_with_debug_symbols=self._build_type == "RelWithDebInfo", + baseline_inclusive_stats=baseline_stats[0], + baseline_exclusive_stats=baseline_stats[1], + stmt_inclusive_stats=stmt_inclusive_stats, + stmt_exclusive_stats=stmt_exclusive_stats, + stmt_callgrind_out=out_contents, + ) + for stmt_inclusive_stats, stmt_exclusive_stats, out_contents in task_stats + ) + + def _invoke( + self, + *, + task_spec: common.TaskSpec, + globals: dict[str, Any], + number: int, + repeats: int, + collect_baseline: bool, + is_python: bool, + retain_out_file: bool, + ) -> tuple[tuple[FunctionCounts, FunctionCounts, str | None], ...]: + """Core invocation method for Callgrind collection. + + Valgrind operates by effectively replacing the CPU with an emulated + version which allows it to instrument any code at the cost of severe + performance degradation. This has the practical effect that in order + to collect Callgrind statistics, a new process has to be created + running under `valgrind`. The steps for this process are: + + 1) Create a scratch directory. + 2) Codegen a run script. (_ValgrindWrapper._construct_script) + Inside the run script: + * Validate that Python and torch match the parent process + * Validate that it is indeed running under valgrind + * Execute `setup` and warm up `stmt` + * Begin collecting stats + * Run the `stmt` loop + * Stop collecting stats + 3) Parse the run results. + 4) Cleanup the scratch directory. + """ + working_dir = common._make_temp_dir(prefix="callgrind") + data_dir = os.path.join(working_dir, "data") + script_file = os.path.join(working_dir, "timer_callgrind.py") + callgrind_out = os.path.join(working_dir, "callgrind.out") + error_log = os.path.join(working_dir, "error.txt") + stat_log = os.path.join(working_dir, "callgrind_stat.txt") + stdout_stderr_log = os.path.join(working_dir, "stdout_stderr.log") + + def run(args: list[str], **kwargs: Any) -> tuple[CompletedProcessType, str]: + # https://thraxil.org/users/anders/posts/2008/03/13/Subprocess-Hanging-PIPE-is-your-enemy/ + with open(stdout_stderr_log, "wb") as f_stdout_stderr: + invocation = subprocess.run( + args, + stdout=f_stdout_stderr, + stderr=subprocess.STDOUT, + **kwargs, + ) + with open(stdout_stderr_log) as f: + return invocation, f.read() + + try: + if is_python: + if self._bindings_module is not None: + shutil.copy( + self._bindings_module.__file__, + os.path.join(working_dir, os.path.split(self._bindings_module.__file__)[1]) + ) + + script_file = os.path.join(working_dir, "timer_callgrind.py") + with open(script_file, "w") as f: + f.write(self._construct_script( + task_spec, + globals=GlobalsBridge(globals, data_dir), + number=number, + repeats=repeats, + collect_baseline=collect_baseline, + error_log=error_log, + stat_log=stat_log, + bindings=self._bindings_module)) + + run_loop_cmd = ["python", script_file] + else: + if collect_baseline: + raise AssertionError("collect_baseline must be False for non-Python timers") + run_loop_exec = cpp_jit.compile_callgrind_template( + stmt=task_spec.stmt, + setup=task_spec.setup, + global_setup=task_spec.global_setup, + ) + run_loop_cmd = [ + run_loop_exec, + "--number", str(number), + "--number-warmup", str(min(number, 10)), + "--repeats", str(repeats), + "--number-threads", str(task_spec.num_threads), + ] + + valgrind_invocation, valgrind_invocation_output = run([ + "valgrind", + "--tool=callgrind", + f"--callgrind-out-file={callgrind_out}", + "--dump-line=yes", + "--dump-instr=yes", + "--instr-atstart=yes", + "--collect-atstart=no", + ] + run_loop_cmd) + + if valgrind_invocation.returncode: + error_report = "" + if os.path.exists(error_log): + with open(error_log) as f: + error_report = f.read() + if not error_report: + error_report = "Unknown error.\n" + valgrind_invocation_output + + raise OSError(f"Failed to collect callgrind profile:\n{error_report}") + + def parse_output(fpath: str, inclusive: bool) -> FunctionCounts: + _annotate_invocation, annotate_invocation_output = run([ + "callgrind_annotate", + f"--inclusive={'yes' if inclusive else 'no'}", + "--threshold=100", + "--show-percs=no", + fpath + ], check=True) + + total_pattern = re.compile(r"^([0-9,]+)\s+PROGRAM TOTALS") + begin_pattern = re.compile(r"Ir\s+file:function") + function_pattern = re.compile(r"^\s*([0-9,]+)\s+(.+:.+)$") + + class ScanState(enum.Enum): + SCANNING_FOR_TOTAL = 0 + SCANNING_FOR_START = 1 + PARSING = 2 + + scan_state = ScanState.SCANNING_FOR_TOTAL + fn_counts = [] + for l in annotate_invocation_output.splitlines(keepends=False): + if scan_state == ScanState.SCANNING_FOR_TOTAL: + total_match = total_pattern.match(l) + if total_match: + program_totals = int(total_match.groups()[0].replace(",", "")) + scan_state = ScanState.SCANNING_FOR_START + + elif scan_state == ScanState.SCANNING_FOR_START: + if begin_pattern.match(l): + scan_state = ScanState.PARSING + + else: + if scan_state != ScanState.PARSING: + raise AssertionError("Failed to enter PARSING state while parsing callgrind_annotate output") + fn_match = function_pattern.match(l) + if fn_match: + ir_str, file_function = fn_match.groups() + ir = int(ir_str.replace(",", "")) + if ir == program_totals: # type: ignore[possibly-undefined] + # Callgrind includes some top level red herring symbols when + # a program dumps multiple profiles. + continue + fn_counts.append(FunctionCount(ir, file_function)) + + elif re.match(r"-+", l): + # Ignore heading separator lines. + continue + + else: + break + + if scan_state != ScanState.PARSING: + raise AssertionError(f"Failed to parse {fpath}") + return FunctionCounts(tuple(sorted(fn_counts, reverse=True)), inclusive=inclusive) + + def read_results(i: int) -> tuple[FunctionCounts, FunctionCounts, str | None]: + if i == repeats and not collect_baseline: + # Null baseline. + return ( + FunctionCounts((), inclusive=True), + FunctionCounts((), inclusive=False), + None, + ) + + fpath = f"{callgrind_out}.{i + 1}" # Callgrind one-indexes files. + callgrind_out_contents: str | None = None + if retain_out_file: + with open(fpath) as f: + callgrind_out_contents = f.read() + + return ( + parse_output(fpath, inclusive=True), + parse_output(fpath, inclusive=False), + callgrind_out_contents + ) + + return tuple(read_results(i) for i in range(repeats + 1)) + finally: + shutil.rmtree(working_dir) + + @staticmethod + def _construct_script( + task_spec: common.TaskSpec, + globals: GlobalsBridge, + *, + number: int, + repeats: int, + collect_baseline: bool, + error_log: str, + stat_log: str, + bindings: CallgrindModuleType | None, + ) -> str: + def block_stmt(stmt: str, indent: int = 0) -> str: + """Partially unroll benchmark loop. + + The naive template looks something like: + "for _ in range({number}): {stmt}" + + However a loop in Python is surprisingly expensive, and significantly + increases the number of background Python instructions. So instead we + partially unroll the loops, with a block size of 100 chosen to keep + the instruction overhead from `range` low while also not ballooning + the size of the generated file. + """ + block_size = 100 + loop_count = number // block_size + if loop_count == 1: + # There is no point in having `for _ in range(1): ...` rather + # than just `...`, and this lets us save shave a few background + # instructions. + loop_count = 0 + remainder = number - block_size * loop_count + blocked_stmt = "" + + if loop_count: + unrolled_stmts = textwrap.indent("\n".join([stmt] * block_size), " " * 4) + blocked_stmt += f"for _ in range({loop_count}):\n{unrolled_stmts}\n" + + if remainder: + blocked_stmt += "\n".join([stmt] * remainder) + + return textwrap.indent(blocked_stmt, " " * indent) + + pass_baseline = ( + "callgrind_bindings._valgrind_toggle()\n" + f"{block_stmt('pass')}\n" + "callgrind_bindings._valgrind_toggle_and_dump_stats()" + ) + + return textwrap.dedent(r""" + import gc + import os + import pickle + import subprocess + import sys + import time + + # Mitigate https://github.com/pytorch/pytorch/issues/37377 + # which can sometimes cause the subprocess call to fail. + import numpy as np + + import torch + torch.set_num_threads({num_threads}) + + {bindings_import} + + PID = os.getpid() + + def log_failure(msg): + with open({error_log_repr}, "wt") as f: + f.write(msg) + sys.exit(1) + + def check_result(completed_process): + if completed_process.returncode: + log_failure(f"Command failed: {{' '.join(completed_process.args)}}") + return completed_process + + # ============================================================================= + # == Check that subprocess matches parent ===================================== + # ============================================================================= + if os.path.realpath(sys.executable) != "{parent_interpreter}": + log_failure( + "Interpreter mismatch:\n" + f" {{os.path.realpath(sys.executable)}}\n vs.\n {parent_interpreter}" + ) + + if torch.__file__ != "{torch_file}": + log_failure( + "PyTorch does not match expected file:\n" + f" {{torch.__file__}}\n vs.\n {torch_file}" + ) + + # ============================================================================= + # == User specified setup ===================================================== + # ============================================================================= + # Load serialized globals + {load_globals} + + # User setup str + {setup} + + for _ in range({warmup_number}): + {indented_stmt} + + # ============================================================================= + # == Callgrind management ===================================================== + # ============================================================================= + with open("{stat_log}", "wb") as stat_file: + # If many instances of callgrind are running at once, the output of + # `callgrind_control` may exceed 16kb which would cause `subprocess.PIPE` + # to deadlock. So instead we use a file. + callgrind_stat = check_result(subprocess.run( + ["callgrind_control", "--stat"], + stdout=stat_file, + stderr=subprocess.STDOUT, + )) + + with open("{stat_log}", "rt") as stat_file: + stat_lines = stat_file.read().splitlines() + + if f"PID {{PID}}: python {{__file__}}" not in stat_lines: + log_failure("Process does not appear to be running callgrind.") + + gc.collect() + time.sleep(0.01) + + # ============================================================================= + # == User code block ========================================================== + # ============================================================================= + for _ in range({repeats}): + callgrind_bindings._valgrind_toggle() + {blocked_stmt} + callgrind_bindings._valgrind_toggle_and_dump_stats() + gc.collect() + + {baseline} + """).strip().format( + indented_stmt=textwrap.indent(task_spec.stmt, " " * 4), + blocked_stmt=block_stmt(task_spec.stmt, indent=4), + baseline=(pass_baseline if collect_baseline else ""), + number=number, + repeats=repeats, + load_globals=globals.construct(), + setup=task_spec.setup, + warmup_number=min(number, 10), + num_threads=task_spec.num_threads, + error_log_repr=repr(error_log), + stat_log=stat_log, + parent_interpreter=os.path.realpath(sys.executable), + torch_file=torch.__file__, + bindings_import=( + "import torch._C as callgrind_bindings" if bindings is None + else f"import {bindings.__name__} as callgrind_bindings"), + ) + + +CALLGRIND_SINGLETON: _ValgrindWrapper | None = None +def wrapper_singleton() -> _ValgrindWrapper: + global CALLGRIND_SINGLETON + if CALLGRIND_SINGLETON is None: + CALLGRIND_SINGLETON = _ValgrindWrapper() + return CALLGRIND_SINGLETON diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/valgrind.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/valgrind.h new file mode 100644 index 0000000000000000000000000000000000000000..d33dd30932aa86b8284cb93d0e29ec646e820197 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/benchmark/utils/valgrind_wrapper/valgrind.h @@ -0,0 +1,7157 @@ +/* -*- c -*- + ---------------------------------------------------------------- + + Notice that the following BSD-style license applies to this one + file (valgrind.h) only. The rest of Valgrind is licensed under the + terms of the GNU General Public License, version 2, unless + otherwise indicated. See the COPYING file in the source + distribution for details. + + ---------------------------------------------------------------- + + This file is part of Valgrind, a dynamic binary instrumentation + framework. + + Copyright (C) 2000-2017 Julian Seward. All rights reserved. + + Redistribution and use in source and binary forms, with or without + modification, are permitted provided that the following conditions + are met: + + 1. Redistributions of source code must retain the above copyright + notice, this list of conditions and the following disclaimer. + + 2. The origin of this software must not be misrepresented; you must + not claim that you wrote the original software. If you use this + software in a product, an acknowledgment in the product + documentation would be appreciated but is not required. + + 3. Altered source versions must be plainly marked as such, and must + not be misrepresented as being the original software. + + 4. The name of the author may not be used to endorse or promote + products derived from this software without specific prior written + permission. + + THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS + OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED + WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE + ARE DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY + DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL + DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE + GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS + INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, + WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING + NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS + SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. + + ---------------------------------------------------------------- + + Notice that the above BSD-style license applies to this one file + (valgrind.h) only. The entire rest of Valgrind is licensed under + the terms of the GNU General Public License, version 2. See the + COPYING file in the source distribution for details. + + ---------------------------------------------------------------- +*/ + + +/* This file is for inclusion into client (your!) code. + + You can use these macros to manipulate and query Valgrind's + execution inside your own programs. + + The resulting executables will still run without Valgrind, just a + little bit more slowly than they otherwise would, but otherwise + unchanged. When not running on valgrind, each client request + consumes very few (eg. 7) instructions, so the resulting performance + loss is negligible unless you plan to execute client requests + millions of times per second. Nevertheless, if that is still a + problem, you can compile with the NVALGRIND symbol defined (gcc + -DNVALGRIND) so that client requests are not even compiled in. */ + +#ifndef __VALGRIND_H +#define __VALGRIND_H + + +/* ------------------------------------------------------------------ */ +/* VERSION NUMBER OF VALGRIND */ +/* ------------------------------------------------------------------ */ + +/* Specify Valgrind's version number, so that user code can + conditionally compile based on our version number. Note that these + were introduced at version 3.6 and so do not exist in version 3.5 + or earlier. The recommended way to use them to check for "version + X.Y or later" is (eg) + +#if defined(__VALGRIND_MAJOR__) && defined(__VALGRIND_MINOR__) \ + && (__VALGRIND_MAJOR__ > 3 \ + || (__VALGRIND_MAJOR__ == 3 && __VALGRIND_MINOR__ >= 6)) +*/ +#define __VALGRIND_MAJOR__ 3 +#define __VALGRIND_MINOR__ 17 + + +#include + +/* Nb: this file might be included in a file compiled with -ansi. So + we can't use C++ style "//" comments nor the "asm" keyword (instead + use "__asm__"). */ + +/* Derive some tags indicating what the target platform is. Note + that in this file we're using the compiler's CPP symbols for + identifying architectures, which are different to the ones we use + within the rest of Valgrind. Note, __powerpc__ is active for both + 32 and 64-bit PPC, whereas __powerpc64__ is only active for the + latter (on Linux, that is). + + Misc note: how to find out what's predefined in gcc by default: + gcc -Wp,-dM somefile.c +*/ +#undef PLAT_x86_darwin +#undef PLAT_amd64_darwin +#undef PLAT_x86_win32 +#undef PLAT_amd64_win64 +#undef PLAT_x86_linux +#undef PLAT_amd64_linux +#undef PLAT_ppc32_linux +#undef PLAT_ppc64be_linux +#undef PLAT_ppc64le_linux +#undef PLAT_arm_linux +#undef PLAT_arm64_linux +#undef PLAT_s390x_linux +#undef PLAT_mips32_linux +#undef PLAT_mips64_linux +#undef PLAT_nanomips_linux +#undef PLAT_x86_solaris +#undef PLAT_amd64_solaris + + +#if defined(__APPLE__) && defined(__i386__) +# define PLAT_x86_darwin 1 +#elif defined(__APPLE__) && defined(__x86_64__) +# define PLAT_amd64_darwin 1 +#elif (defined(__MINGW32__) && defined(__i386__)) \ + || defined(__CYGWIN32__) \ + || (defined(_WIN32) && defined(_M_IX86)) +# define PLAT_x86_win32 1 +#elif (defined(__MINGW32__) && defined(__x86_64__)) \ + || (defined(_WIN32) && defined(_M_X64)) +/* __MINGW32__ and _WIN32 are defined in 64 bit mode as well. */ +# define PLAT_amd64_win64 1 +#elif defined(__linux__) && defined(__i386__) +# define PLAT_x86_linux 1 +#elif defined(__linux__) && defined(__x86_64__) && !defined(__ILP32__) +# define PLAT_amd64_linux 1 +#elif defined(__linux__) && defined(__powerpc__) && !defined(__powerpc64__) +# define PLAT_ppc32_linux 1 +#elif defined(__linux__) && defined(__powerpc__) && defined(__powerpc64__) && _CALL_ELF != 2 +/* Big Endian uses ELF version 1 */ +# define PLAT_ppc64be_linux 1 +#elif defined(__linux__) && defined(__powerpc__) && defined(__powerpc64__) && _CALL_ELF == 2 +/* Little Endian uses ELF version 2 */ +# define PLAT_ppc64le_linux 1 +#elif defined(__linux__) && defined(__arm__) && !defined(__aarch64__) +# define PLAT_arm_linux 1 +#elif defined(__linux__) && defined(__aarch64__) && !defined(__arm__) +# define PLAT_arm64_linux 1 +#elif defined(__linux__) && defined(__s390__) && defined(__s390x__) +# define PLAT_s390x_linux 1 +#elif defined(__linux__) && defined(__mips__) && (__mips==64) +# define PLAT_mips64_linux 1 +#elif defined(__linux__) && defined(__mips__) && (__mips==32) +# define PLAT_mips32_linux 1 +#elif defined(__linux__) && defined(__nanomips__) +# define PLAT_nanomips_linux 1 +#elif defined(__sun) && defined(__i386__) +# define PLAT_x86_solaris 1 +#elif defined(__sun) && defined(__x86_64__) +# define PLAT_amd64_solaris 1 +#else +/* If we're not compiling for our target platform, don't generate + any inline asms. */ +# if !defined(NVALGRIND) +# define NVALGRIND 1 +# endif +#endif + + +/* ------------------------------------------------------------------ */ +/* ARCHITECTURE SPECIFICS for SPECIAL INSTRUCTIONS. There is nothing */ +/* in here of use to end-users -- skip to the next section. */ +/* ------------------------------------------------------------------ */ + +/* + * VALGRIND_DO_CLIENT_REQUEST(): a statement that invokes a Valgrind client + * request. Accepts both pointers and integers as arguments. + * + * VALGRIND_DO_CLIENT_REQUEST_STMT(): a statement that invokes a Valgrind + * client request that does not return a value. + + * VALGRIND_DO_CLIENT_REQUEST_EXPR(): a C expression that invokes a Valgrind + * client request and whose value equals the client request result. Accepts + * both pointers and integers as arguments. Note that such calls are not + * necessarily pure functions -- they may have side effects. + */ + +#define VALGRIND_DO_CLIENT_REQUEST(_zzq_rlval, _zzq_default, \ + _zzq_request, _zzq_arg1, _zzq_arg2, \ + _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + do { (_zzq_rlval) = VALGRIND_DO_CLIENT_REQUEST_EXPR((_zzq_default), \ + (_zzq_request), (_zzq_arg1), (_zzq_arg2), \ + (_zzq_arg3), (_zzq_arg4), (_zzq_arg5)); } while (0) + +#define VALGRIND_DO_CLIENT_REQUEST_STMT(_zzq_request, _zzq_arg1, \ + _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + do { (void) VALGRIND_DO_CLIENT_REQUEST_EXPR(0, \ + (_zzq_request), (_zzq_arg1), (_zzq_arg2), \ + (_zzq_arg3), (_zzq_arg4), (_zzq_arg5)); } while (0) + +#if defined(NVALGRIND) + +/* Define NVALGRIND to completely remove the Valgrind magic sequence + from the compiled code (analogous to NDEBUG's effects on + assert()) */ +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + (_zzq_default) + +#else /* ! NVALGRIND */ + +/* The following defines the magic code sequences which the JITter + spots and handles magically. Don't look too closely at them as + they will rot your brain. + + The assembly code sequences for all architectures is in this one + file. This is because this file must be stand-alone, and we don't + want to have multiple files. + + For VALGRIND_DO_CLIENT_REQUEST, we must ensure that the default + value gets put in the return slot, so that everything works when + this is executed not under Valgrind. Args are passed in a memory + block, and so there's no intrinsic limit to the number that could + be passed, but it's currently five. + + The macro args are: + _zzq_rlval result lvalue + _zzq_default default value (result returned when running on real CPU) + _zzq_request request code + _zzq_arg1..5 request params + + The other two macros are used to support function wrapping, and are + a lot simpler. VALGRIND_GET_NR_CONTEXT returns the value of the + guest's NRADDR pseudo-register and whatever other information is + needed to safely run the call original from the wrapper: on + ppc64-linux, the R2 value at the divert point is also needed. This + information is abstracted into a user-visible type, OrigFn. + + VALGRIND_CALL_NOREDIR_* behaves the same as the following on the + guest, but guarantees that the branch instruction will not be + redirected: x86: call *%eax, amd64: call *%rax, ppc32/ppc64: + branch-and-link-to-r11. VALGRIND_CALL_NOREDIR is just text, not a + complete inline asm, since it needs to be combined with more magic + inline asm stuff to be useful. +*/ + +/* ----------------- x86-{linux,darwin,solaris} ---------------- */ + +#if defined(PLAT_x86_linux) || defined(PLAT_x86_darwin) \ + || (defined(PLAT_x86_win32) && defined(__GNUC__)) \ + || defined(PLAT_x86_solaris) + +typedef + struct { + unsigned int nraddr; /* where's the code? */ + } + OrigFn; + +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "roll $3, %%edi ; roll $13, %%edi\n\t" \ + "roll $29, %%edi ; roll $19, %%edi\n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + __extension__ \ + ({volatile unsigned int _zzq_args[6]; \ + volatile unsigned int _zzq_result; \ + _zzq_args[0] = (unsigned int)(_zzq_request); \ + _zzq_args[1] = (unsigned int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned int)(_zzq_arg5); \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %EDX = client_request ( %EAX ) */ \ + "xchgl %%ebx,%%ebx" \ + : "=d" (_zzq_result) \ + : "a" (&_zzq_args[0]), "0" (_zzq_default) \ + : "cc", "memory" \ + ); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + volatile unsigned int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %EAX = guest_NRADDR */ \ + "xchgl %%ecx,%%ecx" \ + : "=a" (__addr) \ + : \ + : "cc", "memory" \ + ); \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_CALL_NOREDIR_EAX \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* call-noredir *%EAX */ \ + "xchgl %%edx,%%edx\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "xchgl %%edi,%%edi\n\t" \ + : : : "cc", "memory" \ + ); \ + } while (0) + +#endif /* PLAT_x86_linux || PLAT_x86_darwin || (PLAT_x86_win32 && __GNUC__) + || PLAT_x86_solaris */ + +/* ------------------------- x86-Win32 ------------------------- */ + +#if defined(PLAT_x86_win32) && !defined(__GNUC__) + +typedef + struct { + unsigned int nraddr; /* where's the code? */ + } + OrigFn; + +#if defined(_MSC_VER) + +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + __asm rol edi, 3 __asm rol edi, 13 \ + __asm rol edi, 29 __asm rol edi, 19 + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + valgrind_do_client_request_expr((uintptr_t)(_zzq_default), \ + (uintptr_t)(_zzq_request), (uintptr_t)(_zzq_arg1), \ + (uintptr_t)(_zzq_arg2), (uintptr_t)(_zzq_arg3), \ + (uintptr_t)(_zzq_arg4), (uintptr_t)(_zzq_arg5)) + +static __inline uintptr_t +valgrind_do_client_request_expr(uintptr_t _zzq_default, uintptr_t _zzq_request, + uintptr_t _zzq_arg1, uintptr_t _zzq_arg2, + uintptr_t _zzq_arg3, uintptr_t _zzq_arg4, + uintptr_t _zzq_arg5) +{ + volatile uintptr_t _zzq_args[6]; + volatile unsigned int _zzq_result; + _zzq_args[0] = (uintptr_t)(_zzq_request); + _zzq_args[1] = (uintptr_t)(_zzq_arg1); + _zzq_args[2] = (uintptr_t)(_zzq_arg2); + _zzq_args[3] = (uintptr_t)(_zzq_arg3); + _zzq_args[4] = (uintptr_t)(_zzq_arg4); + _zzq_args[5] = (uintptr_t)(_zzq_arg5); + __asm { __asm lea eax, _zzq_args __asm mov edx, _zzq_default + __SPECIAL_INSTRUCTION_PREAMBLE + /* %EDX = client_request ( %EAX ) */ + __asm xchg ebx,ebx + __asm mov _zzq_result, edx + } + return _zzq_result; +} + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + volatile unsigned int __addr; \ + __asm { __SPECIAL_INSTRUCTION_PREAMBLE \ + /* %EAX = guest_NRADDR */ \ + __asm xchg ecx,ecx \ + __asm mov __addr, eax \ + } \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_CALL_NOREDIR_EAX ERROR + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm { __SPECIAL_INSTRUCTION_PREAMBLE \ + __asm xchg edi,edi \ + } \ + } while (0) + +#else +#error Unsupported compiler. +#endif + +#endif /* PLAT_x86_win32 */ + +/* ----------------- amd64-{linux,darwin,solaris} --------------- */ + +#if defined(PLAT_amd64_linux) || defined(PLAT_amd64_darwin) \ + || defined(PLAT_amd64_solaris) \ + || (defined(PLAT_amd64_win64) && defined(__GNUC__)) + +typedef + struct { + unsigned long int nraddr; /* where's the code? */ + } + OrigFn; + +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "rolq $3, %%rdi ; rolq $13, %%rdi\n\t" \ + "rolq $61, %%rdi ; rolq $51, %%rdi\n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + __extension__ \ + ({ volatile unsigned long int _zzq_args[6]; \ + volatile unsigned long int _zzq_result; \ + _zzq_args[0] = (unsigned long int)(_zzq_request); \ + _zzq_args[1] = (unsigned long int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned long int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned long int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned long int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned long int)(_zzq_arg5); \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %RDX = client_request ( %RAX ) */ \ + "xchgq %%rbx,%%rbx" \ + : "=d" (_zzq_result) \ + : "a" (&_zzq_args[0]), "0" (_zzq_default) \ + : "cc", "memory" \ + ); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + volatile unsigned long int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %RAX = guest_NRADDR */ \ + "xchgq %%rcx,%%rcx" \ + : "=a" (__addr) \ + : \ + : "cc", "memory" \ + ); \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_CALL_NOREDIR_RAX \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* call-noredir *%RAX */ \ + "xchgq %%rdx,%%rdx\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "xchgq %%rdi,%%rdi\n\t" \ + : : : "cc", "memory" \ + ); \ + } while (0) + +#endif /* PLAT_amd64_linux || PLAT_amd64_darwin || PLAT_amd64_solaris */ + +/* ------------------------- amd64-Win64 ------------------------- */ + +#if defined(PLAT_amd64_win64) && !defined(__GNUC__) + +#error Unsupported compiler. + +#endif /* PLAT_amd64_win64 */ + +/* ------------------------ ppc32-linux ------------------------ */ + +#if defined(PLAT_ppc32_linux) + +typedef + struct { + unsigned int nraddr; /* where's the code? */ + } + OrigFn; + +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "rlwinm 0,0,3,0,31 ; rlwinm 0,0,13,0,31\n\t" \ + "rlwinm 0,0,29,0,31 ; rlwinm 0,0,19,0,31\n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + \ + __extension__ \ + ({ unsigned int _zzq_args[6]; \ + unsigned int _zzq_result; \ + unsigned int* _zzq_ptr; \ + _zzq_args[0] = (unsigned int)(_zzq_request); \ + _zzq_args[1] = (unsigned int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned int)(_zzq_arg5); \ + _zzq_ptr = _zzq_args; \ + __asm__ volatile("mr 3,%1\n\t" /*default*/ \ + "mr 4,%2\n\t" /*ptr*/ \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* %R3 = client_request ( %R4 ) */ \ + "or 1,1,1\n\t" \ + "mr %0,3" /*result*/ \ + : "=b" (_zzq_result) \ + : "b" (_zzq_default), "b" (_zzq_ptr) \ + : "cc", "memory", "r3", "r4"); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + unsigned int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %R3 = guest_NRADDR */ \ + "or 2,2,2\n\t" \ + "mr %0,3" \ + : "=b" (__addr) \ + : \ + : "cc", "memory", "r3" \ + ); \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* branch-and-link-to-noredir *%R11 */ \ + "or 3,3,3\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "or 5,5,5\n\t" \ + ); \ + } while (0) + +#endif /* PLAT_ppc32_linux */ + +/* ------------------------ ppc64-linux ------------------------ */ + +#if defined(PLAT_ppc64be_linux) + +typedef + struct { + unsigned long int nraddr; /* where's the code? */ + unsigned long int r2; /* what tocptr do we need? */ + } + OrigFn; + +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "rotldi 0,0,3 ; rotldi 0,0,13\n\t" \ + "rotldi 0,0,61 ; rotldi 0,0,51\n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + \ + __extension__ \ + ({ unsigned long int _zzq_args[6]; \ + unsigned long int _zzq_result; \ + unsigned long int* _zzq_ptr; \ + _zzq_args[0] = (unsigned long int)(_zzq_request); \ + _zzq_args[1] = (unsigned long int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned long int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned long int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned long int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned long int)(_zzq_arg5); \ + _zzq_ptr = _zzq_args; \ + __asm__ volatile("mr 3,%1\n\t" /*default*/ \ + "mr 4,%2\n\t" /*ptr*/ \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* %R3 = client_request ( %R4 ) */ \ + "or 1,1,1\n\t" \ + "mr %0,3" /*result*/ \ + : "=b" (_zzq_result) \ + : "b" (_zzq_default), "b" (_zzq_ptr) \ + : "cc", "memory", "r3", "r4"); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + unsigned long int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %R3 = guest_NRADDR */ \ + "or 2,2,2\n\t" \ + "mr %0,3" \ + : "=b" (__addr) \ + : \ + : "cc", "memory", "r3" \ + ); \ + _zzq_orig->nraddr = __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %R3 = guest_NRADDR_GPR2 */ \ + "or 4,4,4\n\t" \ + "mr %0,3" \ + : "=b" (__addr) \ + : \ + : "cc", "memory", "r3" \ + ); \ + _zzq_orig->r2 = __addr; \ + } + +#define VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* branch-and-link-to-noredir *%R11 */ \ + "or 3,3,3\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "or 5,5,5\n\t" \ + ); \ + } while (0) + +#endif /* PLAT_ppc64be_linux */ + +#if defined(PLAT_ppc64le_linux) + +typedef + struct { + unsigned long int nraddr; /* where's the code? */ + unsigned long int r2; /* what tocptr do we need? */ + } + OrigFn; + +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "rotldi 0,0,3 ; rotldi 0,0,13\n\t" \ + "rotldi 0,0,61 ; rotldi 0,0,51\n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + \ + __extension__ \ + ({ unsigned long int _zzq_args[6]; \ + unsigned long int _zzq_result; \ + unsigned long int* _zzq_ptr; \ + _zzq_args[0] = (unsigned long int)(_zzq_request); \ + _zzq_args[1] = (unsigned long int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned long int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned long int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned long int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned long int)(_zzq_arg5); \ + _zzq_ptr = _zzq_args; \ + __asm__ volatile("mr 3,%1\n\t" /*default*/ \ + "mr 4,%2\n\t" /*ptr*/ \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* %R3 = client_request ( %R4 ) */ \ + "or 1,1,1\n\t" \ + "mr %0,3" /*result*/ \ + : "=b" (_zzq_result) \ + : "b" (_zzq_default), "b" (_zzq_ptr) \ + : "cc", "memory", "r3", "r4"); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + unsigned long int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %R3 = guest_NRADDR */ \ + "or 2,2,2\n\t" \ + "mr %0,3" \ + : "=b" (__addr) \ + : \ + : "cc", "memory", "r3" \ + ); \ + _zzq_orig->nraddr = __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %R3 = guest_NRADDR_GPR2 */ \ + "or 4,4,4\n\t" \ + "mr %0,3" \ + : "=b" (__addr) \ + : \ + : "cc", "memory", "r3" \ + ); \ + _zzq_orig->r2 = __addr; \ + } + +#define VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* branch-and-link-to-noredir *%R12 */ \ + "or 3,3,3\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "or 5,5,5\n\t" \ + ); \ + } while (0) + +#endif /* PLAT_ppc64le_linux */ + +/* ------------------------- arm-linux ------------------------- */ + +#if defined(PLAT_arm_linux) + +typedef + struct { + unsigned int nraddr; /* where's the code? */ + } + OrigFn; + +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "mov r12, r12, ror #3 ; mov r12, r12, ror #13 \n\t" \ + "mov r12, r12, ror #29 ; mov r12, r12, ror #19 \n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + \ + __extension__ \ + ({volatile unsigned int _zzq_args[6]; \ + volatile unsigned int _zzq_result; \ + _zzq_args[0] = (unsigned int)(_zzq_request); \ + _zzq_args[1] = (unsigned int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned int)(_zzq_arg5); \ + __asm__ volatile("mov r3, %1\n\t" /*default*/ \ + "mov r4, %2\n\t" /*ptr*/ \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* R3 = client_request ( R4 ) */ \ + "orr r10, r10, r10\n\t" \ + "mov %0, r3" /*result*/ \ + : "=r" (_zzq_result) \ + : "r" (_zzq_default), "r" (&_zzq_args[0]) \ + : "cc","memory", "r3", "r4"); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + unsigned int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* R3 = guest_NRADDR */ \ + "orr r11, r11, r11\n\t" \ + "mov %0, r3" \ + : "=r" (__addr) \ + : \ + : "cc", "memory", "r3" \ + ); \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* branch-and-link-to-noredir *%R4 */ \ + "orr r12, r12, r12\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "orr r9, r9, r9\n\t" \ + : : : "cc", "memory" \ + ); \ + } while (0) + +#endif /* PLAT_arm_linux */ + +/* ------------------------ arm64-linux ------------------------- */ + +#if defined(PLAT_arm64_linux) + +typedef + struct { + unsigned long int nraddr; /* where's the code? */ + } + OrigFn; + +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "ror x12, x12, #3 ; ror x12, x12, #13 \n\t" \ + "ror x12, x12, #51 ; ror x12, x12, #61 \n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + \ + __extension__ \ + ({volatile unsigned long int _zzq_args[6]; \ + volatile unsigned long int _zzq_result; \ + _zzq_args[0] = (unsigned long int)(_zzq_request); \ + _zzq_args[1] = (unsigned long int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned long int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned long int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned long int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned long int)(_zzq_arg5); \ + __asm__ volatile("mov x3, %1\n\t" /*default*/ \ + "mov x4, %2\n\t" /*ptr*/ \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* X3 = client_request ( X4 ) */ \ + "orr x10, x10, x10\n\t" \ + "mov %0, x3" /*result*/ \ + : "=r" (_zzq_result) \ + : "r" ((unsigned long int)(_zzq_default)), \ + "r" (&_zzq_args[0]) \ + : "cc","memory", "x3", "x4"); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + unsigned long int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* X3 = guest_NRADDR */ \ + "orr x11, x11, x11\n\t" \ + "mov %0, x3" \ + : "=r" (__addr) \ + : \ + : "cc", "memory", "x3" \ + ); \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* branch-and-link-to-noredir X8 */ \ + "orr x12, x12, x12\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "orr x9, x9, x9\n\t" \ + : : : "cc", "memory" \ + ); \ + } while (0) + +#endif /* PLAT_arm64_linux */ + +/* ------------------------ s390x-linux ------------------------ */ + +#if defined(PLAT_s390x_linux) + +typedef + struct { + unsigned long int nraddr; /* where's the code? */ + } + OrigFn; + +/* __SPECIAL_INSTRUCTION_PREAMBLE will be used to identify Valgrind specific + * code. This detection is implemented in platform specific toIR.c + * (e.g. VEX/priv/guest_s390_decoder.c). + */ +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "lr 15,15\n\t" \ + "lr 1,1\n\t" \ + "lr 2,2\n\t" \ + "lr 3,3\n\t" + +#define __CLIENT_REQUEST_CODE "lr 2,2\n\t" +#define __GET_NR_CONTEXT_CODE "lr 3,3\n\t" +#define __CALL_NO_REDIR_CODE "lr 4,4\n\t" +#define __VEX_INJECT_IR_CODE "lr 5,5\n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + __extension__ \ + ({volatile unsigned long int _zzq_args[6]; \ + volatile unsigned long int _zzq_result; \ + _zzq_args[0] = (unsigned long int)(_zzq_request); \ + _zzq_args[1] = (unsigned long int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned long int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned long int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned long int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned long int)(_zzq_arg5); \ + __asm__ volatile(/* r2 = args */ \ + "lgr 2,%1\n\t" \ + /* r3 = default */ \ + "lgr 3,%2\n\t" \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + __CLIENT_REQUEST_CODE \ + /* results = r3 */ \ + "lgr %0, 3\n\t" \ + : "=d" (_zzq_result) \ + : "a" (&_zzq_args[0]), "0" (_zzq_default) \ + : "cc", "2", "3", "memory" \ + ); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + volatile unsigned long int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + __GET_NR_CONTEXT_CODE \ + "lgr %0, 3\n\t" \ + : "=a" (__addr) \ + : \ + : "cc", "3", "memory" \ + ); \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_CALL_NOREDIR_R1 \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + __CALL_NO_REDIR_CODE + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + __VEX_INJECT_IR_CODE); \ + } while (0) + +#endif /* PLAT_s390x_linux */ + +/* ------------------------- mips32-linux ---------------- */ + +#if defined(PLAT_mips32_linux) + +typedef + struct { + unsigned int nraddr; /* where's the code? */ + } + OrigFn; + +/* .word 0x342 + * .word 0x742 + * .word 0xC2 + * .word 0x4C2*/ +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "srl $0, $0, 13\n\t" \ + "srl $0, $0, 29\n\t" \ + "srl $0, $0, 3\n\t" \ + "srl $0, $0, 19\n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + __extension__ \ + ({ volatile unsigned int _zzq_args[6]; \ + volatile unsigned int _zzq_result; \ + _zzq_args[0] = (unsigned int)(_zzq_request); \ + _zzq_args[1] = (unsigned int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned int)(_zzq_arg5); \ + __asm__ volatile("move $11, %1\n\t" /*default*/ \ + "move $12, %2\n\t" /*ptr*/ \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* T3 = client_request ( T4 ) */ \ + "or $13, $13, $13\n\t" \ + "move %0, $11\n\t" /*result*/ \ + : "=r" (_zzq_result) \ + : "r" (_zzq_default), "r" (&_zzq_args[0]) \ + : "$11", "$12", "memory"); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + volatile unsigned int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* %t9 = guest_NRADDR */ \ + "or $14, $14, $14\n\t" \ + "move %0, $11" /*result*/ \ + : "=r" (__addr) \ + : \ + : "$11" \ + ); \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_CALL_NOREDIR_T9 \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* call-noredir *%t9 */ \ + "or $15, $15, $15\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "or $11, $11, $11\n\t" \ + ); \ + } while (0) + + +#endif /* PLAT_mips32_linux */ + +/* ------------------------- mips64-linux ---------------- */ + +#if defined(PLAT_mips64_linux) + +typedef + struct { + unsigned long nraddr; /* where's the code? */ + } + OrigFn; + +/* dsll $0,$0, 3 + * dsll $0,$0, 13 + * dsll $0,$0, 29 + * dsll $0,$0, 19*/ +#define __SPECIAL_INSTRUCTION_PREAMBLE \ + "dsll $0,$0, 3 ; dsll $0,$0,13\n\t" \ + "dsll $0,$0,29 ; dsll $0,$0,19\n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + __extension__ \ + ({ volatile unsigned long int _zzq_args[6]; \ + volatile unsigned long int _zzq_result; \ + _zzq_args[0] = (unsigned long int)(_zzq_request); \ + _zzq_args[1] = (unsigned long int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned long int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned long int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned long int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned long int)(_zzq_arg5); \ + __asm__ volatile("move $11, %1\n\t" /*default*/ \ + "move $12, %2\n\t" /*ptr*/ \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* $11 = client_request ( $12 ) */ \ + "or $13, $13, $13\n\t" \ + "move %0, $11\n\t" /*result*/ \ + : "=r" (_zzq_result) \ + : "r" (_zzq_default), "r" (&_zzq_args[0]) \ + : "$11", "$12", "memory"); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + volatile unsigned long int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* $11 = guest_NRADDR */ \ + "or $14, $14, $14\n\t" \ + "move %0, $11" /*result*/ \ + : "=r" (__addr) \ + : \ + : "$11"); \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_CALL_NOREDIR_T9 \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* call-noredir $25 */ \ + "or $15, $15, $15\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "or $11, $11, $11\n\t" \ + ); \ + } while (0) + +#endif /* PLAT_mips64_linux */ + +#if defined(PLAT_nanomips_linux) + +typedef + struct { + unsigned int nraddr; /* where's the code? */ + } + OrigFn; +/* + 8000 c04d srl zero, zero, 13 + 8000 c05d srl zero, zero, 29 + 8000 c043 srl zero, zero, 3 + 8000 c053 srl zero, zero, 19 +*/ + +#define __SPECIAL_INSTRUCTION_PREAMBLE "srl[32] $zero, $zero, 13 \n\t" \ + "srl[32] $zero, $zero, 29 \n\t" \ + "srl[32] $zero, $zero, 3 \n\t" \ + "srl[32] $zero, $zero, 19 \n\t" + +#define VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + _zzq_default, _zzq_request, \ + _zzq_arg1, _zzq_arg2, _zzq_arg3, _zzq_arg4, _zzq_arg5) \ + __extension__ \ + ({ volatile unsigned int _zzq_args[6]; \ + volatile unsigned int _zzq_result; \ + _zzq_args[0] = (unsigned int)(_zzq_request); \ + _zzq_args[1] = (unsigned int)(_zzq_arg1); \ + _zzq_args[2] = (unsigned int)(_zzq_arg2); \ + _zzq_args[3] = (unsigned int)(_zzq_arg3); \ + _zzq_args[4] = (unsigned int)(_zzq_arg4); \ + _zzq_args[5] = (unsigned int)(_zzq_arg5); \ + __asm__ volatile("move $a7, %1\n\t" /* default */ \ + "move $t0, %2\n\t" /* ptr */ \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* $a7 = client_request( $t0 ) */ \ + "or[32] $t0, $t0, $t0\n\t" \ + "move %0, $a7\n\t" /* result */ \ + : "=r" (_zzq_result) \ + : "r" (_zzq_default), "r" (&_zzq_args[0]) \ + : "$a7", "$t0", "memory"); \ + _zzq_result; \ + }) + +#define VALGRIND_GET_NR_CONTEXT(_zzq_rlval) \ + { volatile OrigFn* _zzq_orig = &(_zzq_rlval); \ + volatile unsigned long int __addr; \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + /* $a7 = guest_NRADDR */ \ + "or[32] $t1, $t1, $t1\n\t" \ + "move %0, $a7" /*result*/ \ + : "=r" (__addr) \ + : \ + : "$a7"); \ + _zzq_orig->nraddr = __addr; \ + } + +#define VALGRIND_CALL_NOREDIR_T9 \ + __SPECIAL_INSTRUCTION_PREAMBLE \ + /* call-noredir $25 */ \ + "or[32] $t2, $t2, $t2\n\t" + +#define VALGRIND_VEX_INJECT_IR() \ + do { \ + __asm__ volatile(__SPECIAL_INSTRUCTION_PREAMBLE \ + "or[32] $t3, $t3, $t3\n\t" \ + ); \ + } while (0) + +#endif +/* Insert assembly code for other platforms here... */ + +#endif /* NVALGRIND */ + + +/* ------------------------------------------------------------------ */ +/* PLATFORM SPECIFICS for FUNCTION WRAPPING. This is all very */ +/* ugly. It's the least-worst tradeoff I can think of. */ +/* ------------------------------------------------------------------ */ + +/* This section defines magic (a.k.a appalling-hack) macros for doing + guaranteed-no-redirection macros, so as to get from function + wrappers to the functions they are wrapping. The whole point is to + construct standard call sequences, but to do the call itself with a + special no-redirect call pseudo-instruction that the JIT + understands and handles specially. This section is long and + repetitious, and I can't see a way to make it shorter. + + The naming scheme is as follows: + + CALL_FN_{W,v}_{v,W,WW,WWW,WWWW,5W,6W,7W,etc} + + 'W' stands for "word" and 'v' for "void". Hence there are + different macros for calling arity 0, 1, 2, 3, 4, etc, functions, + and for each, the possibility of returning a word-typed result, or + no result. +*/ + +/* Use these to write the name of your wrapper. NOTE: duplicates + VG_WRAP_FUNCTION_Z{U,Z} in pub_tool_redir.h. NOTE also: inserts + the default behaviour equivalance class tag "0000" into the name. + See pub_tool_redir.h for details -- normally you don't need to + think about this, though. */ + +/* Use an extra level of macroisation so as to ensure the soname/fnname + args are fully macro-expanded before pasting them together. */ +#define VG_CONCAT4(_aa,_bb,_cc,_dd) _aa##_bb##_cc##_dd + +#define I_WRAP_SONAME_FNNAME_ZU(soname,fnname) \ + VG_CONCAT4(_vgw00000ZU_,soname,_,fnname) + +#define I_WRAP_SONAME_FNNAME_ZZ(soname,fnname) \ + VG_CONCAT4(_vgw00000ZZ_,soname,_,fnname) + +/* Use this macro from within a wrapper function to collect the + context (address and possibly other info) of the original function. + Once you have that you can then use it in one of the CALL_FN_ + macros. The type of the argument _lval is OrigFn. */ +#define VALGRIND_GET_ORIG_FN(_lval) VALGRIND_GET_NR_CONTEXT(_lval) + +/* Also provide end-user facilities for function replacement, rather + than wrapping. A replacement function differs from a wrapper in + that it has no way to get hold of the original function being + called, and hence no way to call onwards to it. In a replacement + function, VALGRIND_GET_ORIG_FN always returns zero. */ + +#define I_REPLACE_SONAME_FNNAME_ZU(soname,fnname) \ + VG_CONCAT4(_vgr00000ZU_,soname,_,fnname) + +#define I_REPLACE_SONAME_FNNAME_ZZ(soname,fnname) \ + VG_CONCAT4(_vgr00000ZZ_,soname,_,fnname) + +/* Derivatives of the main macros below, for calling functions + returning void. */ + +#define CALL_FN_v_v(fnptr) \ + do { volatile unsigned long _junk; \ + CALL_FN_W_v(_junk,fnptr); } while (0) + +#define CALL_FN_v_W(fnptr, arg1) \ + do { volatile unsigned long _junk; \ + CALL_FN_W_W(_junk,fnptr,arg1); } while (0) + +#define CALL_FN_v_WW(fnptr, arg1,arg2) \ + do { volatile unsigned long _junk; \ + CALL_FN_W_WW(_junk,fnptr,arg1,arg2); } while (0) + +#define CALL_FN_v_WWW(fnptr, arg1,arg2,arg3) \ + do { volatile unsigned long _junk; \ + CALL_FN_W_WWW(_junk,fnptr,arg1,arg2,arg3); } while (0) + +#define CALL_FN_v_WWWW(fnptr, arg1,arg2,arg3,arg4) \ + do { volatile unsigned long _junk; \ + CALL_FN_W_WWWW(_junk,fnptr,arg1,arg2,arg3,arg4); } while (0) + +#define CALL_FN_v_5W(fnptr, arg1,arg2,arg3,arg4,arg5) \ + do { volatile unsigned long _junk; \ + CALL_FN_W_5W(_junk,fnptr,arg1,arg2,arg3,arg4,arg5); } while (0) + +#define CALL_FN_v_6W(fnptr, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { volatile unsigned long _junk; \ + CALL_FN_W_6W(_junk,fnptr,arg1,arg2,arg3,arg4,arg5,arg6); } while (0) + +#define CALL_FN_v_7W(fnptr, arg1,arg2,arg3,arg4,arg5,arg6,arg7) \ + do { volatile unsigned long _junk; \ + CALL_FN_W_7W(_junk,fnptr,arg1,arg2,arg3,arg4,arg5,arg6,arg7); } while (0) + +/* ----------------- x86-{linux,darwin,solaris} ---------------- */ + +#if defined(PLAT_x86_linux) || defined(PLAT_x86_darwin) \ + || defined(PLAT_x86_solaris) + +/* These regs are trashed by the hidden call. No need to mention eax + as gcc can already see that, plus causes gcc to bomb. */ +#define __CALLER_SAVED_REGS /*"eax"*/ "ecx", "edx" + +/* Macros to save and align the stack before making a function + call and restore it afterwards as gcc may not keep the stack + pointer aligned if it doesn't realise calls are being made + to other functions. */ + +#define VALGRIND_ALIGN_STACK \ + "movl %%esp,%%edi\n\t" \ + "andl $0xfffffff0,%%esp\n\t" +#define VALGRIND_RESTORE_STACK \ + "movl %%edi,%%esp\n\t" + +/* These CALL_FN_ macros assume that on x86-linux, sizeof(unsigned + long) == 4. */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[1]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[2]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "subl $12, %%esp\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "subl $8, %%esp\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[4]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "subl $4, %%esp\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[5]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "pushl 16(%%eax)\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[6]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "subl $12, %%esp\n\t" \ + "pushl 20(%%eax)\n\t" \ + "pushl 16(%%eax)\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[7]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "subl $8, %%esp\n\t" \ + "pushl 24(%%eax)\n\t" \ + "pushl 20(%%eax)\n\t" \ + "pushl 16(%%eax)\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[8]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "subl $4, %%esp\n\t" \ + "pushl 28(%%eax)\n\t" \ + "pushl 24(%%eax)\n\t" \ + "pushl 20(%%eax)\n\t" \ + "pushl 16(%%eax)\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[9]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "pushl 32(%%eax)\n\t" \ + "pushl 28(%%eax)\n\t" \ + "pushl 24(%%eax)\n\t" \ + "pushl 20(%%eax)\n\t" \ + "pushl 16(%%eax)\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[10]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "subl $12, %%esp\n\t" \ + "pushl 36(%%eax)\n\t" \ + "pushl 32(%%eax)\n\t" \ + "pushl 28(%%eax)\n\t" \ + "pushl 24(%%eax)\n\t" \ + "pushl 20(%%eax)\n\t" \ + "pushl 16(%%eax)\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[11]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "subl $8, %%esp\n\t" \ + "pushl 40(%%eax)\n\t" \ + "pushl 36(%%eax)\n\t" \ + "pushl 32(%%eax)\n\t" \ + "pushl 28(%%eax)\n\t" \ + "pushl 24(%%eax)\n\t" \ + "pushl 20(%%eax)\n\t" \ + "pushl 16(%%eax)\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[12]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "subl $4, %%esp\n\t" \ + "pushl 44(%%eax)\n\t" \ + "pushl 40(%%eax)\n\t" \ + "pushl 36(%%eax)\n\t" \ + "pushl 32(%%eax)\n\t" \ + "pushl 28(%%eax)\n\t" \ + "pushl 24(%%eax)\n\t" \ + "pushl 20(%%eax)\n\t" \ + "pushl 16(%%eax)\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11,arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[13]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + _argvec[12] = (unsigned long)(arg12); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "pushl 48(%%eax)\n\t" \ + "pushl 44(%%eax)\n\t" \ + "pushl 40(%%eax)\n\t" \ + "pushl 36(%%eax)\n\t" \ + "pushl 32(%%eax)\n\t" \ + "pushl 28(%%eax)\n\t" \ + "pushl 24(%%eax)\n\t" \ + "pushl 20(%%eax)\n\t" \ + "pushl 16(%%eax)\n\t" \ + "pushl 12(%%eax)\n\t" \ + "pushl 8(%%eax)\n\t" \ + "pushl 4(%%eax)\n\t" \ + "movl (%%eax), %%eax\n\t" /* target->%eax */ \ + VALGRIND_CALL_NOREDIR_EAX \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "edi" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#endif /* PLAT_x86_linux || PLAT_x86_darwin || PLAT_x86_solaris */ + +/* ---------------- amd64-{linux,darwin,solaris} --------------- */ + +#if defined(PLAT_amd64_linux) || defined(PLAT_amd64_darwin) \ + || defined(PLAT_amd64_solaris) + +/* ARGREGS: rdi rsi rdx rcx r8 r9 (the rest on stack in R-to-L order) */ + +/* These regs are trashed by the hidden call. */ +#define __CALLER_SAVED_REGS /*"rax",*/ "rcx", "rdx", "rsi", \ + "rdi", "r8", "r9", "r10", "r11" + +/* This is all pretty complex. It's so as to make stack unwinding + work reliably. See bug 243270. The basic problem is the sub and + add of 128 of %rsp in all of the following macros. If gcc believes + the CFA is in %rsp, then unwinding may fail, because what's at the + CFA is not what gcc "expected" when it constructs the CFIs for the + places where the macros are instantiated. + + But we can't just add a CFI annotation to increase the CFA offset + by 128, to match the sub of 128 from %rsp, because we don't know + whether gcc has chosen %rsp as the CFA at that point, or whether it + has chosen some other register (eg, %rbp). In the latter case, + adding a CFI annotation to change the CFA offset is simply wrong. + + So the solution is to get hold of the CFA using + __builtin_dwarf_cfa(), put it in a known register, and add a + CFI annotation to say what the register is. We choose %rbp for + this (perhaps perversely), because: + + (1) %rbp is already subject to unwinding. If a new register was + chosen then the unwinder would have to unwind it in all stack + traces, which is expensive, and + + (2) %rbp is already subject to precise exception updates in the + JIT. If a new register was chosen, we'd have to have precise + exceptions for it too, which reduces performance of the + generated code. + + However .. one extra complication. We can't just whack the result + of __builtin_dwarf_cfa() into %rbp and then add %rbp to the + list of trashed registers at the end of the inline assembly + fragments; gcc won't allow %rbp to appear in that list. Hence + instead we need to stash %rbp in %r15 for the duration of the asm, + and say that %r15 is trashed instead. gcc seems happy to go with + that. + + Oh .. and this all needs to be conditionalised so that it is + unchanged from before this commit, when compiled with older gccs + that don't support __builtin_dwarf_cfa. Furthermore, since + this header file is freestanding, it has to be independent of + config.h, and so the following conditionalisation cannot depend on + configure time checks. + + Although it's not clear from + 'defined(__GNUC__) && defined(__GCC_HAVE_DWARF2_CFI_ASM)', + this expression excludes Darwin. + .cfi directives in Darwin assembly appear to be completely + different and I haven't investigated how they work. + + For even more entertainment value, note we have to use the + completely undocumented __builtin_dwarf_cfa(), which appears to + really compute the CFA, whereas __builtin_frame_address(0) claims + to but actually doesn't. See + https://bugs.kde.org/show_bug.cgi?id=243270#c47 +*/ +#if defined(__GNUC__) && defined(__GCC_HAVE_DWARF2_CFI_ASM) +# define __FRAME_POINTER \ + ,"r"(__builtin_dwarf_cfa()) +# define VALGRIND_CFI_PROLOGUE \ + "movq %%rbp, %%r15\n\t" \ + "movq %2, %%rbp\n\t" \ + ".cfi_remember_state\n\t" \ + ".cfi_def_cfa rbp, 0\n\t" +# define VALGRIND_CFI_EPILOGUE \ + "movq %%r15, %%rbp\n\t" \ + ".cfi_restore_state\n\t" +#else +# define __FRAME_POINTER +# define VALGRIND_CFI_PROLOGUE +# define VALGRIND_CFI_EPILOGUE +#endif + +/* Macros to save and align the stack before making a function + call and restore it afterwards as gcc may not keep the stack + pointer aligned if it doesn't realise calls are being made + to other functions. */ + +#define VALGRIND_ALIGN_STACK \ + "movq %%rsp,%%r14\n\t" \ + "andq $0xfffffffffffffff0,%%rsp\n\t" +#define VALGRIND_RESTORE_STACK \ + "movq %%r14,%%rsp\n\t" + +/* These CALL_FN_ macros assume that on amd64-linux, sizeof(unsigned + long) == 8. */ + +/* NB 9 Sept 07. There is a nasty kludge here in all these CALL_FN_ + macros. In order not to trash the stack redzone, we need to drop + %rsp by 128 before the hidden call, and restore afterwards. The + nastyness is that it is only by luck that the stack still appears + to be unwindable during the hidden call - since then the behaviour + of any routine using this macro does not match what the CFI data + says. Sigh. + + Why is this important? Imagine that a wrapper has a stack + allocated local, and passes to the hidden call, a pointer to it. + Because gcc does not know about the hidden call, it may allocate + that local in the redzone. Unfortunately the hidden call may then + trash it before it comes to use it. So we must step clear of the + redzone, for the duration of the hidden call, to make it safe. + + Probably the same problem afflicts the other redzone-style ABIs too + (ppc64-linux); but for those, the stack is + self describing (none of this CFI nonsense) so at least messing + with the stack pointer doesn't give a danger of non-unwindable + stack. */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[1]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[2]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[4]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[5]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "movq 32(%%rax), %%rcx\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[6]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "movq 40(%%rax), %%r8\n\t" \ + "movq 32(%%rax), %%rcx\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[7]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "movq 48(%%rax), %%r9\n\t" \ + "movq 40(%%rax), %%r8\n\t" \ + "movq 32(%%rax), %%rcx\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[8]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $136,%%rsp\n\t" \ + "pushq 56(%%rax)\n\t" \ + "movq 48(%%rax), %%r9\n\t" \ + "movq 40(%%rax), %%r8\n\t" \ + "movq 32(%%rax), %%rcx\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[9]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "pushq 64(%%rax)\n\t" \ + "pushq 56(%%rax)\n\t" \ + "movq 48(%%rax), %%r9\n\t" \ + "movq 40(%%rax), %%r8\n\t" \ + "movq 32(%%rax), %%rcx\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[10]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $136,%%rsp\n\t" \ + "pushq 72(%%rax)\n\t" \ + "pushq 64(%%rax)\n\t" \ + "pushq 56(%%rax)\n\t" \ + "movq 48(%%rax), %%r9\n\t" \ + "movq 40(%%rax), %%r8\n\t" \ + "movq 32(%%rax), %%rcx\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[11]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "pushq 80(%%rax)\n\t" \ + "pushq 72(%%rax)\n\t" \ + "pushq 64(%%rax)\n\t" \ + "pushq 56(%%rax)\n\t" \ + "movq 48(%%rax), %%r9\n\t" \ + "movq 40(%%rax), %%r8\n\t" \ + "movq 32(%%rax), %%rcx\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[12]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $136,%%rsp\n\t" \ + "pushq 88(%%rax)\n\t" \ + "pushq 80(%%rax)\n\t" \ + "pushq 72(%%rax)\n\t" \ + "pushq 64(%%rax)\n\t" \ + "pushq 56(%%rax)\n\t" \ + "movq 48(%%rax), %%r9\n\t" \ + "movq 40(%%rax), %%r8\n\t" \ + "movq 32(%%rax), %%rcx\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11,arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[13]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + _argvec[12] = (unsigned long)(arg12); \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + VALGRIND_ALIGN_STACK \ + "subq $128,%%rsp\n\t" \ + "pushq 96(%%rax)\n\t" \ + "pushq 88(%%rax)\n\t" \ + "pushq 80(%%rax)\n\t" \ + "pushq 72(%%rax)\n\t" \ + "pushq 64(%%rax)\n\t" \ + "pushq 56(%%rax)\n\t" \ + "movq 48(%%rax), %%r9\n\t" \ + "movq 40(%%rax), %%r8\n\t" \ + "movq 32(%%rax), %%rcx\n\t" \ + "movq 24(%%rax), %%rdx\n\t" \ + "movq 16(%%rax), %%rsi\n\t" \ + "movq 8(%%rax), %%rdi\n\t" \ + "movq (%%rax), %%rax\n\t" /* target->%rax */ \ + VALGRIND_CALL_NOREDIR_RAX \ + VALGRIND_RESTORE_STACK \ + VALGRIND_CFI_EPILOGUE \ + : /*out*/ "=a" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r14", "r15" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#endif /* PLAT_amd64_linux || PLAT_amd64_darwin || PLAT_amd64_solaris */ + +/* ------------------------ ppc32-linux ------------------------ */ + +#if defined(PLAT_ppc32_linux) + +/* This is useful for finding out about the on-stack stuff: + + extern int f9 ( int,int,int,int,int,int,int,int,int ); + extern int f10 ( int,int,int,int,int,int,int,int,int,int ); + extern int f11 ( int,int,int,int,int,int,int,int,int,int,int ); + extern int f12 ( int,int,int,int,int,int,int,int,int,int,int,int ); + + int g9 ( void ) { + return f9(11,22,33,44,55,66,77,88,99); + } + int g10 ( void ) { + return f10(11,22,33,44,55,66,77,88,99,110); + } + int g11 ( void ) { + return f11(11,22,33,44,55,66,77,88,99,110,121); + } + int g12 ( void ) { + return f12(11,22,33,44,55,66,77,88,99,110,121,132); + } +*/ + +/* ARGREGS: r3 r4 r5 r6 r7 r8 r9 r10 (the rest on stack somewhere) */ + +/* These regs are trashed by the hidden call. */ +#define __CALLER_SAVED_REGS \ + "lr", "ctr", "xer", \ + "cr0", "cr1", "cr2", "cr3", "cr4", "cr5", "cr6", "cr7", \ + "r0", "r2", "r3", "r4", "r5", "r6", "r7", "r8", "r9", "r10", \ + "r11", "r12", "r13" + +/* Macros to save and align the stack before making a function + call and restore it afterwards as gcc may not keep the stack + pointer aligned if it doesn't realise calls are being made + to other functions. */ + +#define VALGRIND_ALIGN_STACK \ + "mr 28,1\n\t" \ + "rlwinm 1,1,0,0,27\n\t" +#define VALGRIND_RESTORE_STACK \ + "mr 1,28\n\t" + +/* These CALL_FN_ macros assume that on ppc32-linux, + sizeof(unsigned long) == 4. */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[1]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[2]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[4]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[5]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 6,16(11)\n\t" /* arg4->r6 */ \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[6]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 6,16(11)\n\t" /* arg4->r6 */ \ + "lwz 7,20(11)\n\t" \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[7]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 6,16(11)\n\t" /* arg4->r6 */ \ + "lwz 7,20(11)\n\t" \ + "lwz 8,24(11)\n\t" \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[8]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 6,16(11)\n\t" /* arg4->r6 */ \ + "lwz 7,20(11)\n\t" \ + "lwz 8,24(11)\n\t" \ + "lwz 9,28(11)\n\t" \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[9]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 6,16(11)\n\t" /* arg4->r6 */ \ + "lwz 7,20(11)\n\t" \ + "lwz 8,24(11)\n\t" \ + "lwz 9,28(11)\n\t" \ + "lwz 10,32(11)\n\t" /* arg8->r10 */ \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[10]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + _argvec[9] = (unsigned long)arg9; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "addi 1,1,-16\n\t" \ + /* arg9 */ \ + "lwz 3,36(11)\n\t" \ + "stw 3,8(1)\n\t" \ + /* args1-8 */ \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 6,16(11)\n\t" /* arg4->r6 */ \ + "lwz 7,20(11)\n\t" \ + "lwz 8,24(11)\n\t" \ + "lwz 9,28(11)\n\t" \ + "lwz 10,32(11)\n\t" /* arg8->r10 */ \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[11]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + _argvec[9] = (unsigned long)arg9; \ + _argvec[10] = (unsigned long)arg10; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "addi 1,1,-16\n\t" \ + /* arg10 */ \ + "lwz 3,40(11)\n\t" \ + "stw 3,12(1)\n\t" \ + /* arg9 */ \ + "lwz 3,36(11)\n\t" \ + "stw 3,8(1)\n\t" \ + /* args1-8 */ \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 6,16(11)\n\t" /* arg4->r6 */ \ + "lwz 7,20(11)\n\t" \ + "lwz 8,24(11)\n\t" \ + "lwz 9,28(11)\n\t" \ + "lwz 10,32(11)\n\t" /* arg8->r10 */ \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[12]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + _argvec[9] = (unsigned long)arg9; \ + _argvec[10] = (unsigned long)arg10; \ + _argvec[11] = (unsigned long)arg11; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "addi 1,1,-32\n\t" \ + /* arg11 */ \ + "lwz 3,44(11)\n\t" \ + "stw 3,16(1)\n\t" \ + /* arg10 */ \ + "lwz 3,40(11)\n\t" \ + "stw 3,12(1)\n\t" \ + /* arg9 */ \ + "lwz 3,36(11)\n\t" \ + "stw 3,8(1)\n\t" \ + /* args1-8 */ \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 6,16(11)\n\t" /* arg4->r6 */ \ + "lwz 7,20(11)\n\t" \ + "lwz 8,24(11)\n\t" \ + "lwz 9,28(11)\n\t" \ + "lwz 10,32(11)\n\t" /* arg8->r10 */ \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11,arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[13]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + _argvec[9] = (unsigned long)arg9; \ + _argvec[10] = (unsigned long)arg10; \ + _argvec[11] = (unsigned long)arg11; \ + _argvec[12] = (unsigned long)arg12; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "addi 1,1,-32\n\t" \ + /* arg12 */ \ + "lwz 3,48(11)\n\t" \ + "stw 3,20(1)\n\t" \ + /* arg11 */ \ + "lwz 3,44(11)\n\t" \ + "stw 3,16(1)\n\t" \ + /* arg10 */ \ + "lwz 3,40(11)\n\t" \ + "stw 3,12(1)\n\t" \ + /* arg9 */ \ + "lwz 3,36(11)\n\t" \ + "stw 3,8(1)\n\t" \ + /* args1-8 */ \ + "lwz 3,4(11)\n\t" /* arg1->r3 */ \ + "lwz 4,8(11)\n\t" \ + "lwz 5,12(11)\n\t" \ + "lwz 6,16(11)\n\t" /* arg4->r6 */ \ + "lwz 7,20(11)\n\t" \ + "lwz 8,24(11)\n\t" \ + "lwz 9,28(11)\n\t" \ + "lwz 10,32(11)\n\t" /* arg8->r10 */ \ + "lwz 11,0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + VALGRIND_RESTORE_STACK \ + "mr %0,3" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#endif /* PLAT_ppc32_linux */ + +/* ------------------------ ppc64-linux ------------------------ */ + +#if defined(PLAT_ppc64be_linux) + +/* ARGREGS: r3 r4 r5 r6 r7 r8 r9 r10 (the rest on stack somewhere) */ + +/* These regs are trashed by the hidden call. */ +#define __CALLER_SAVED_REGS \ + "lr", "ctr", "xer", \ + "cr0", "cr1", "cr2", "cr3", "cr4", "cr5", "cr6", "cr7", \ + "r0", "r3", "r4", "r5", "r6", "r7", "r8", "r9", "r10", \ + "r11", "r12", "r13" + +/* Macros to save and align the stack before making a function + call and restore it afterwards as gcc may not keep the stack + pointer aligned if it doesn't realise calls are being made + to other functions. */ + +#define VALGRIND_ALIGN_STACK \ + "mr 28,1\n\t" \ + "rldicr 1,1,0,59\n\t" +#define VALGRIND_RESTORE_STACK \ + "mr 1,28\n\t" + +/* These CALL_FN_ macros assume that on ppc64-linux, sizeof(unsigned + long) == 8. */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+0]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+1]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+2]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+3]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+4]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 6, 32(11)\n\t" /* arg4->r6 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+5]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 6, 32(11)\n\t" /* arg4->r6 */ \ + "ld 7, 40(11)\n\t" /* arg5->r7 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+6]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 6, 32(11)\n\t" /* arg4->r6 */ \ + "ld 7, 40(11)\n\t" /* arg5->r7 */ \ + "ld 8, 48(11)\n\t" /* arg6->r8 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+7]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 6, 32(11)\n\t" /* arg4->r6 */ \ + "ld 7, 40(11)\n\t" /* arg5->r7 */ \ + "ld 8, 48(11)\n\t" /* arg6->r8 */ \ + "ld 9, 56(11)\n\t" /* arg7->r9 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+8]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 6, 32(11)\n\t" /* arg4->r6 */ \ + "ld 7, 40(11)\n\t" /* arg5->r7 */ \ + "ld 8, 48(11)\n\t" /* arg6->r8 */ \ + "ld 9, 56(11)\n\t" /* arg7->r9 */ \ + "ld 10, 64(11)\n\t" /* arg8->r10 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+9]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + _argvec[2+9] = (unsigned long)arg9; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "addi 1,1,-128\n\t" /* expand stack frame */ \ + /* arg9 */ \ + "ld 3,72(11)\n\t" \ + "std 3,112(1)\n\t" \ + /* args1-8 */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 6, 32(11)\n\t" /* arg4->r6 */ \ + "ld 7, 40(11)\n\t" /* arg5->r7 */ \ + "ld 8, 48(11)\n\t" /* arg6->r8 */ \ + "ld 9, 56(11)\n\t" /* arg7->r9 */ \ + "ld 10, 64(11)\n\t" /* arg8->r10 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+10]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + _argvec[2+9] = (unsigned long)arg9; \ + _argvec[2+10] = (unsigned long)arg10; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "addi 1,1,-128\n\t" /* expand stack frame */ \ + /* arg10 */ \ + "ld 3,80(11)\n\t" \ + "std 3,120(1)\n\t" \ + /* arg9 */ \ + "ld 3,72(11)\n\t" \ + "std 3,112(1)\n\t" \ + /* args1-8 */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 6, 32(11)\n\t" /* arg4->r6 */ \ + "ld 7, 40(11)\n\t" /* arg5->r7 */ \ + "ld 8, 48(11)\n\t" /* arg6->r8 */ \ + "ld 9, 56(11)\n\t" /* arg7->r9 */ \ + "ld 10, 64(11)\n\t" /* arg8->r10 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+11]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + _argvec[2+9] = (unsigned long)arg9; \ + _argvec[2+10] = (unsigned long)arg10; \ + _argvec[2+11] = (unsigned long)arg11; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "addi 1,1,-144\n\t" /* expand stack frame */ \ + /* arg11 */ \ + "ld 3,88(11)\n\t" \ + "std 3,128(1)\n\t" \ + /* arg10 */ \ + "ld 3,80(11)\n\t" \ + "std 3,120(1)\n\t" \ + /* arg9 */ \ + "ld 3,72(11)\n\t" \ + "std 3,112(1)\n\t" \ + /* args1-8 */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 6, 32(11)\n\t" /* arg4->r6 */ \ + "ld 7, 40(11)\n\t" /* arg5->r7 */ \ + "ld 8, 48(11)\n\t" /* arg6->r8 */ \ + "ld 9, 56(11)\n\t" /* arg7->r9 */ \ + "ld 10, 64(11)\n\t" /* arg8->r10 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11,arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+12]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + _argvec[2+9] = (unsigned long)arg9; \ + _argvec[2+10] = (unsigned long)arg10; \ + _argvec[2+11] = (unsigned long)arg11; \ + _argvec[2+12] = (unsigned long)arg12; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 11,%1\n\t" \ + "std 2,-16(11)\n\t" /* save tocptr */ \ + "ld 2,-8(11)\n\t" /* use nraddr's tocptr */ \ + "addi 1,1,-144\n\t" /* expand stack frame */ \ + /* arg12 */ \ + "ld 3,96(11)\n\t" \ + "std 3,136(1)\n\t" \ + /* arg11 */ \ + "ld 3,88(11)\n\t" \ + "std 3,128(1)\n\t" \ + /* arg10 */ \ + "ld 3,80(11)\n\t" \ + "std 3,120(1)\n\t" \ + /* arg9 */ \ + "ld 3,72(11)\n\t" \ + "std 3,112(1)\n\t" \ + /* args1-8 */ \ + "ld 3, 8(11)\n\t" /* arg1->r3 */ \ + "ld 4, 16(11)\n\t" /* arg2->r4 */ \ + "ld 5, 24(11)\n\t" /* arg3->r5 */ \ + "ld 6, 32(11)\n\t" /* arg4->r6 */ \ + "ld 7, 40(11)\n\t" /* arg5->r7 */ \ + "ld 8, 48(11)\n\t" /* arg6->r8 */ \ + "ld 9, 56(11)\n\t" /* arg7->r9 */ \ + "ld 10, 64(11)\n\t" /* arg8->r10 */ \ + "ld 11, 0(11)\n\t" /* target->r11 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R11 \ + "mr 11,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(11)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#endif /* PLAT_ppc64be_linux */ + +/* ------------------------- ppc64le-linux ----------------------- */ +#if defined(PLAT_ppc64le_linux) + +/* ARGREGS: r3 r4 r5 r6 r7 r8 r9 r10 (the rest on stack somewhere) */ + +/* These regs are trashed by the hidden call. */ +#define __CALLER_SAVED_REGS \ + "lr", "ctr", "xer", \ + "cr0", "cr1", "cr2", "cr3", "cr4", "cr5", "cr6", "cr7", \ + "r0", "r3", "r4", "r5", "r6", "r7", "r8", "r9", "r10", \ + "r11", "r12", "r13" + +/* Macros to save and align the stack before making a function + call and restore it afterwards as gcc may not keep the stack + pointer aligned if it doesn't realise calls are being made + to other functions. */ + +#define VALGRIND_ALIGN_STACK \ + "mr 28,1\n\t" \ + "rldicr 1,1,0,59\n\t" +#define VALGRIND_RESTORE_STACK \ + "mr 1,28\n\t" + +/* These CALL_FN_ macros assume that on ppc64-linux, sizeof(unsigned + long) == 8. */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+0]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+1]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+2]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+3]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+4]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 6, 32(12)\n\t" /* arg4->r6 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+5]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 6, 32(12)\n\t" /* arg4->r6 */ \ + "ld 7, 40(12)\n\t" /* arg5->r7 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+6]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 6, 32(12)\n\t" /* arg4->r6 */ \ + "ld 7, 40(12)\n\t" /* arg5->r7 */ \ + "ld 8, 48(12)\n\t" /* arg6->r8 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+7]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 6, 32(12)\n\t" /* arg4->r6 */ \ + "ld 7, 40(12)\n\t" /* arg5->r7 */ \ + "ld 8, 48(12)\n\t" /* arg6->r8 */ \ + "ld 9, 56(12)\n\t" /* arg7->r9 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+8]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 6, 32(12)\n\t" /* arg4->r6 */ \ + "ld 7, 40(12)\n\t" /* arg5->r7 */ \ + "ld 8, 48(12)\n\t" /* arg6->r8 */ \ + "ld 9, 56(12)\n\t" /* arg7->r9 */ \ + "ld 10, 64(12)\n\t" /* arg8->r10 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+9]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + _argvec[2+9] = (unsigned long)arg9; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "addi 1,1,-128\n\t" /* expand stack frame */ \ + /* arg9 */ \ + "ld 3,72(12)\n\t" \ + "std 3,96(1)\n\t" \ + /* args1-8 */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 6, 32(12)\n\t" /* arg4->r6 */ \ + "ld 7, 40(12)\n\t" /* arg5->r7 */ \ + "ld 8, 48(12)\n\t" /* arg6->r8 */ \ + "ld 9, 56(12)\n\t" /* arg7->r9 */ \ + "ld 10, 64(12)\n\t" /* arg8->r10 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+10]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + _argvec[2+9] = (unsigned long)arg9; \ + _argvec[2+10] = (unsigned long)arg10; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "addi 1,1,-128\n\t" /* expand stack frame */ \ + /* arg10 */ \ + "ld 3,80(12)\n\t" \ + "std 3,104(1)\n\t" \ + /* arg9 */ \ + "ld 3,72(12)\n\t" \ + "std 3,96(1)\n\t" \ + /* args1-8 */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 6, 32(12)\n\t" /* arg4->r6 */ \ + "ld 7, 40(12)\n\t" /* arg5->r7 */ \ + "ld 8, 48(12)\n\t" /* arg6->r8 */ \ + "ld 9, 56(12)\n\t" /* arg7->r9 */ \ + "ld 10, 64(12)\n\t" /* arg8->r10 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+11]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + _argvec[2+9] = (unsigned long)arg9; \ + _argvec[2+10] = (unsigned long)arg10; \ + _argvec[2+11] = (unsigned long)arg11; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "addi 1,1,-144\n\t" /* expand stack frame */ \ + /* arg11 */ \ + "ld 3,88(12)\n\t" \ + "std 3,112(1)\n\t" \ + /* arg10 */ \ + "ld 3,80(12)\n\t" \ + "std 3,104(1)\n\t" \ + /* arg9 */ \ + "ld 3,72(12)\n\t" \ + "std 3,96(1)\n\t" \ + /* args1-8 */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 6, 32(12)\n\t" /* arg4->r6 */ \ + "ld 7, 40(12)\n\t" /* arg5->r7 */ \ + "ld 8, 48(12)\n\t" /* arg6->r8 */ \ + "ld 9, 56(12)\n\t" /* arg7->r9 */ \ + "ld 10, 64(12)\n\t" /* arg8->r10 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11,arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3+12]; \ + volatile unsigned long _res; \ + /* _argvec[0] holds current r2 across the call */ \ + _argvec[1] = (unsigned long)_orig.r2; \ + _argvec[2] = (unsigned long)_orig.nraddr; \ + _argvec[2+1] = (unsigned long)arg1; \ + _argvec[2+2] = (unsigned long)arg2; \ + _argvec[2+3] = (unsigned long)arg3; \ + _argvec[2+4] = (unsigned long)arg4; \ + _argvec[2+5] = (unsigned long)arg5; \ + _argvec[2+6] = (unsigned long)arg6; \ + _argvec[2+7] = (unsigned long)arg7; \ + _argvec[2+8] = (unsigned long)arg8; \ + _argvec[2+9] = (unsigned long)arg9; \ + _argvec[2+10] = (unsigned long)arg10; \ + _argvec[2+11] = (unsigned long)arg11; \ + _argvec[2+12] = (unsigned long)arg12; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "mr 12,%1\n\t" \ + "std 2,-16(12)\n\t" /* save tocptr */ \ + "ld 2,-8(12)\n\t" /* use nraddr's tocptr */ \ + "addi 1,1,-144\n\t" /* expand stack frame */ \ + /* arg12 */ \ + "ld 3,96(12)\n\t" \ + "std 3,120(1)\n\t" \ + /* arg11 */ \ + "ld 3,88(12)\n\t" \ + "std 3,112(1)\n\t" \ + /* arg10 */ \ + "ld 3,80(12)\n\t" \ + "std 3,104(1)\n\t" \ + /* arg9 */ \ + "ld 3,72(12)\n\t" \ + "std 3,96(1)\n\t" \ + /* args1-8 */ \ + "ld 3, 8(12)\n\t" /* arg1->r3 */ \ + "ld 4, 16(12)\n\t" /* arg2->r4 */ \ + "ld 5, 24(12)\n\t" /* arg3->r5 */ \ + "ld 6, 32(12)\n\t" /* arg4->r6 */ \ + "ld 7, 40(12)\n\t" /* arg5->r7 */ \ + "ld 8, 48(12)\n\t" /* arg6->r8 */ \ + "ld 9, 56(12)\n\t" /* arg7->r9 */ \ + "ld 10, 64(12)\n\t" /* arg8->r10 */ \ + "ld 12, 0(12)\n\t" /* target->r12 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R12 \ + "mr 12,%1\n\t" \ + "mr %0,3\n\t" \ + "ld 2,-16(12)\n\t" /* restore tocptr */ \ + VALGRIND_RESTORE_STACK \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[2]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r28" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#endif /* PLAT_ppc64le_linux */ + +/* ------------------------- arm-linux ------------------------- */ + +#if defined(PLAT_arm_linux) + +/* These regs are trashed by the hidden call. */ +#define __CALLER_SAVED_REGS "r0", "r1", "r2", "r3","r4", "r12", "r14" + +/* Macros to save and align the stack before making a function + call and restore it afterwards as gcc may not keep the stack + pointer aligned if it doesn't realise calls are being made + to other functions. */ + +/* This is a bit tricky. We store the original stack pointer in r10 + as it is callee-saves. gcc doesn't allow the use of r11 for some + reason. Also, we can't directly "bic" the stack pointer in thumb + mode since r13 isn't an allowed register number in that context. + So use r4 as a temporary, since that is about to get trashed + anyway, just after each use of this macro. Side effect is we need + to be very careful about any future changes, since + VALGRIND_ALIGN_STACK simply assumes r4 is usable. */ +#define VALGRIND_ALIGN_STACK \ + "mov r10, sp\n\t" \ + "mov r4, sp\n\t" \ + "bic r4, r4, #7\n\t" \ + "mov sp, r4\n\t" +#define VALGRIND_RESTORE_STACK \ + "mov sp, r10\n\t" + +/* These CALL_FN_ macros assume that on arm-linux, sizeof(unsigned + long) == 4. */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[1]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[2]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[4]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[5]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r3, [%1, #16] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[6]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "sub sp, sp, #4 \n\t" \ + "ldr r0, [%1, #20] \n\t" \ + "push {r0} \n\t" \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r3, [%1, #16] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[7]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr r0, [%1, #20] \n\t" \ + "ldr r1, [%1, #24] \n\t" \ + "push {r0, r1} \n\t" \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r3, [%1, #16] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[8]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "sub sp, sp, #4 \n\t" \ + "ldr r0, [%1, #20] \n\t" \ + "ldr r1, [%1, #24] \n\t" \ + "ldr r2, [%1, #28] \n\t" \ + "push {r0, r1, r2} \n\t" \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r3, [%1, #16] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[9]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr r0, [%1, #20] \n\t" \ + "ldr r1, [%1, #24] \n\t" \ + "ldr r2, [%1, #28] \n\t" \ + "ldr r3, [%1, #32] \n\t" \ + "push {r0, r1, r2, r3} \n\t" \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r3, [%1, #16] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[10]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "sub sp, sp, #4 \n\t" \ + "ldr r0, [%1, #20] \n\t" \ + "ldr r1, [%1, #24] \n\t" \ + "ldr r2, [%1, #28] \n\t" \ + "ldr r3, [%1, #32] \n\t" \ + "ldr r4, [%1, #36] \n\t" \ + "push {r0, r1, r2, r3, r4} \n\t" \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r3, [%1, #16] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[11]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr r0, [%1, #40] \n\t" \ + "push {r0} \n\t" \ + "ldr r0, [%1, #20] \n\t" \ + "ldr r1, [%1, #24] \n\t" \ + "ldr r2, [%1, #28] \n\t" \ + "ldr r3, [%1, #32] \n\t" \ + "ldr r4, [%1, #36] \n\t" \ + "push {r0, r1, r2, r3, r4} \n\t" \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r3, [%1, #16] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[12]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "sub sp, sp, #4 \n\t" \ + "ldr r0, [%1, #40] \n\t" \ + "ldr r1, [%1, #44] \n\t" \ + "push {r0, r1} \n\t" \ + "ldr r0, [%1, #20] \n\t" \ + "ldr r1, [%1, #24] \n\t" \ + "ldr r2, [%1, #28] \n\t" \ + "ldr r3, [%1, #32] \n\t" \ + "ldr r4, [%1, #36] \n\t" \ + "push {r0, r1, r2, r3, r4} \n\t" \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r3, [%1, #16] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11,arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[13]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + _argvec[12] = (unsigned long)(arg12); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr r0, [%1, #40] \n\t" \ + "ldr r1, [%1, #44] \n\t" \ + "ldr r2, [%1, #48] \n\t" \ + "push {r0, r1, r2} \n\t" \ + "ldr r0, [%1, #20] \n\t" \ + "ldr r1, [%1, #24] \n\t" \ + "ldr r2, [%1, #28] \n\t" \ + "ldr r3, [%1, #32] \n\t" \ + "ldr r4, [%1, #36] \n\t" \ + "push {r0, r1, r2, r3, r4} \n\t" \ + "ldr r0, [%1, #4] \n\t" \ + "ldr r1, [%1, #8] \n\t" \ + "ldr r2, [%1, #12] \n\t" \ + "ldr r3, [%1, #16] \n\t" \ + "ldr r4, [%1] \n\t" /* target->r4 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_R4 \ + VALGRIND_RESTORE_STACK \ + "mov %0, r0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "r10" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#endif /* PLAT_arm_linux */ + +/* ------------------------ arm64-linux ------------------------ */ + +#if defined(PLAT_arm64_linux) + +/* These regs are trashed by the hidden call. */ +#define __CALLER_SAVED_REGS \ + "x0", "x1", "x2", "x3","x4", "x5", "x6", "x7", "x8", "x9", \ + "x10", "x11", "x12", "x13", "x14", "x15", "x16", "x17", \ + "x18", "x19", "x20", "x30", \ + "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", "v8", "v9", \ + "v10", "v11", "v12", "v13", "v14", "v15", "v16", "v17", \ + "v18", "v19", "v20", "v21", "v22", "v23", "v24", "v25", \ + "v26", "v27", "v28", "v29", "v30", "v31" + +/* x21 is callee-saved, so we can use it to save and restore SP around + the hidden call. */ +#define VALGRIND_ALIGN_STACK \ + "mov x21, sp\n\t" \ + "bic sp, x21, #15\n\t" +#define VALGRIND_RESTORE_STACK \ + "mov sp, x21\n\t" + +/* These CALL_FN_ macros assume that on arm64-linux, + sizeof(unsigned long) == 8. */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[1]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[2]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[4]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[5]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x3, [%1, #32] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[6]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x3, [%1, #32] \n\t" \ + "ldr x4, [%1, #40] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[7]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x3, [%1, #32] \n\t" \ + "ldr x4, [%1, #40] \n\t" \ + "ldr x5, [%1, #48] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[8]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x3, [%1, #32] \n\t" \ + "ldr x4, [%1, #40] \n\t" \ + "ldr x5, [%1, #48] \n\t" \ + "ldr x6, [%1, #56] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[9]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x3, [%1, #32] \n\t" \ + "ldr x4, [%1, #40] \n\t" \ + "ldr x5, [%1, #48] \n\t" \ + "ldr x6, [%1, #56] \n\t" \ + "ldr x7, [%1, #64] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[10]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "sub sp, sp, #0x20 \n\t" \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x3, [%1, #32] \n\t" \ + "ldr x4, [%1, #40] \n\t" \ + "ldr x5, [%1, #48] \n\t" \ + "ldr x6, [%1, #56] \n\t" \ + "ldr x7, [%1, #64] \n\t" \ + "ldr x8, [%1, #72] \n\t" \ + "str x8, [sp, #0] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[11]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "sub sp, sp, #0x20 \n\t" \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x3, [%1, #32] \n\t" \ + "ldr x4, [%1, #40] \n\t" \ + "ldr x5, [%1, #48] \n\t" \ + "ldr x6, [%1, #56] \n\t" \ + "ldr x7, [%1, #64] \n\t" \ + "ldr x8, [%1, #72] \n\t" \ + "str x8, [sp, #0] \n\t" \ + "ldr x8, [%1, #80] \n\t" \ + "str x8, [sp, #8] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[12]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "sub sp, sp, #0x30 \n\t" \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x3, [%1, #32] \n\t" \ + "ldr x4, [%1, #40] \n\t" \ + "ldr x5, [%1, #48] \n\t" \ + "ldr x6, [%1, #56] \n\t" \ + "ldr x7, [%1, #64] \n\t" \ + "ldr x8, [%1, #72] \n\t" \ + "str x8, [sp, #0] \n\t" \ + "ldr x8, [%1, #80] \n\t" \ + "str x8, [sp, #8] \n\t" \ + "ldr x8, [%1, #88] \n\t" \ + "str x8, [sp, #16] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10,arg11, \ + arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[13]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + _argvec[12] = (unsigned long)(arg12); \ + __asm__ volatile( \ + VALGRIND_ALIGN_STACK \ + "sub sp, sp, #0x30 \n\t" \ + "ldr x0, [%1, #8] \n\t" \ + "ldr x1, [%1, #16] \n\t" \ + "ldr x2, [%1, #24] \n\t" \ + "ldr x3, [%1, #32] \n\t" \ + "ldr x4, [%1, #40] \n\t" \ + "ldr x5, [%1, #48] \n\t" \ + "ldr x6, [%1, #56] \n\t" \ + "ldr x7, [%1, #64] \n\t" \ + "ldr x8, [%1, #72] \n\t" \ + "str x8, [sp, #0] \n\t" \ + "ldr x8, [%1, #80] \n\t" \ + "str x8, [sp, #8] \n\t" \ + "ldr x8, [%1, #88] \n\t" \ + "str x8, [sp, #16] \n\t" \ + "ldr x8, [%1, #96] \n\t" \ + "str x8, [sp, #24] \n\t" \ + "ldr x8, [%1] \n\t" /* target->x8 */ \ + VALGRIND_BRANCH_AND_LINK_TO_NOREDIR_X8 \ + VALGRIND_RESTORE_STACK \ + "mov %0, x0" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS, "x21" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#endif /* PLAT_arm64_linux */ + +/* ------------------------- s390x-linux ------------------------- */ + +#if defined(PLAT_s390x_linux) + +/* Similar workaround as amd64 (see above), but we use r11 as frame + pointer and save the old r11 in r7. r11 might be used for + argvec, therefore we copy argvec in r1 since r1 is clobbered + after the call anyway. */ +#if defined(__GNUC__) && defined(__GCC_HAVE_DWARF2_CFI_ASM) +# define __FRAME_POINTER \ + ,"d"(__builtin_dwarf_cfa()) +# define VALGRIND_CFI_PROLOGUE \ + ".cfi_remember_state\n\t" \ + "lgr 1,%1\n\t" /* copy the argvec pointer in r1 */ \ + "lgr 7,11\n\t" \ + "lgr 11,%2\n\t" \ + ".cfi_def_cfa r11, 0\n\t" +# define VALGRIND_CFI_EPILOGUE \ + "lgr 11, 7\n\t" \ + ".cfi_restore_state\n\t" +#else +# define __FRAME_POINTER +# define VALGRIND_CFI_PROLOGUE \ + "lgr 1,%1\n\t" +# define VALGRIND_CFI_EPILOGUE +#endif + +/* Nb: On s390 the stack pointer is properly aligned *at all times* + according to the s390 GCC maintainer. (The ABI specification is not + precise in this regard.) Therefore, VALGRIND_ALIGN_STACK and + VALGRIND_RESTORE_STACK are not defined here. */ + +/* These regs are trashed by the hidden call. Note that we overwrite + r14 in s390_irgen_noredir (VEX/priv/guest_s390_irgen.c) to give the + function a proper return address. All others are ABI defined call + clobbers. */ +#if defined(__VX__) || defined(__S390_VX__) +#define __CALLER_SAVED_REGS "0", "1", "2", "3", "4", "5", "14", \ + "v0", "v1", "v2", "v3", "v4", "v5", "v6", "v7", \ + "v8", "v9", "v10", "v11", "v12", "v13", "v14", "v15", \ + "v16", "v17", "v18", "v19", "v20", "v21", "v22", "v23", \ + "v24", "v25", "v26", "v27", "v28", "v29", "v30", "v31" +#else +#define __CALLER_SAVED_REGS "0", "1", "2", "3", "4", "5", "14", \ + "f0", "f1", "f2", "f3", "f4", "f5", "f6", "f7" +#endif + +/* Nb: Although r11 is modified in the asm snippets below (inside + VALGRIND_CFI_PROLOGUE) it is not listed in the clobber section, for + two reasons: + (1) r11 is restored in VALGRIND_CFI_EPILOGUE, so effectively it is not + modified + (2) GCC will complain that r11 cannot appear inside a clobber section, + when compiled with -O -fno-omit-frame-pointer + */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[1]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-160\n\t" \ + "lg 1, 0(1)\n\t" /* target->r1 */ \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,160\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "d" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +/* The call abi has the arguments in r2-r6 and stack */ +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[2]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-160\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,160\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1, arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-160\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,160\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1, arg2, arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[4]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-160\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,160\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1, arg2, arg3, arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[5]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-160\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 5,32(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,160\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1, arg2, arg3, arg4, arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[6]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-160\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 5,32(1)\n\t" \ + "lg 6,40(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,160\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"6","7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_6W(lval, orig, arg1, arg2, arg3, arg4, arg5, \ + arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[7]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-168\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 5,32(1)\n\t" \ + "lg 6,40(1)\n\t" \ + "mvc 160(8,15), 48(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,168\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"6","7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1, arg2, arg3, arg4, arg5, \ + arg6, arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[8]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-176\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 5,32(1)\n\t" \ + "lg 6,40(1)\n\t" \ + "mvc 160(8,15), 48(1)\n\t" \ + "mvc 168(8,15), 56(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,176\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"6","7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1, arg2, arg3, arg4, arg5, \ + arg6, arg7 ,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[9]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-184\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 5,32(1)\n\t" \ + "lg 6,40(1)\n\t" \ + "mvc 160(8,15), 48(1)\n\t" \ + "mvc 168(8,15), 56(1)\n\t" \ + "mvc 176(8,15), 64(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,184\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"6","7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1, arg2, arg3, arg4, arg5, \ + arg6, arg7 ,arg8, arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[10]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + _argvec[9] = (unsigned long)arg9; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-192\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 5,32(1)\n\t" \ + "lg 6,40(1)\n\t" \ + "mvc 160(8,15), 48(1)\n\t" \ + "mvc 168(8,15), 56(1)\n\t" \ + "mvc 176(8,15), 64(1)\n\t" \ + "mvc 184(8,15), 72(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,192\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"6","7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1, arg2, arg3, arg4, arg5, \ + arg6, arg7 ,arg8, arg9, arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[11]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + _argvec[9] = (unsigned long)arg9; \ + _argvec[10] = (unsigned long)arg10; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-200\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 5,32(1)\n\t" \ + "lg 6,40(1)\n\t" \ + "mvc 160(8,15), 48(1)\n\t" \ + "mvc 168(8,15), 56(1)\n\t" \ + "mvc 176(8,15), 64(1)\n\t" \ + "mvc 184(8,15), 72(1)\n\t" \ + "mvc 192(8,15), 80(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,200\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"6","7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1, arg2, arg3, arg4, arg5, \ + arg6, arg7 ,arg8, arg9, arg10, arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[12]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + _argvec[9] = (unsigned long)arg9; \ + _argvec[10] = (unsigned long)arg10; \ + _argvec[11] = (unsigned long)arg11; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-208\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 5,32(1)\n\t" \ + "lg 6,40(1)\n\t" \ + "mvc 160(8,15), 48(1)\n\t" \ + "mvc 168(8,15), 56(1)\n\t" \ + "mvc 176(8,15), 64(1)\n\t" \ + "mvc 184(8,15), 72(1)\n\t" \ + "mvc 192(8,15), 80(1)\n\t" \ + "mvc 200(8,15), 88(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,208\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"6","7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1, arg2, arg3, arg4, arg5, \ + arg6, arg7 ,arg8, arg9, arg10, arg11, arg12)\ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[13]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)arg1; \ + _argvec[2] = (unsigned long)arg2; \ + _argvec[3] = (unsigned long)arg3; \ + _argvec[4] = (unsigned long)arg4; \ + _argvec[5] = (unsigned long)arg5; \ + _argvec[6] = (unsigned long)arg6; \ + _argvec[7] = (unsigned long)arg7; \ + _argvec[8] = (unsigned long)arg8; \ + _argvec[9] = (unsigned long)arg9; \ + _argvec[10] = (unsigned long)arg10; \ + _argvec[11] = (unsigned long)arg11; \ + _argvec[12] = (unsigned long)arg12; \ + __asm__ volatile( \ + VALGRIND_CFI_PROLOGUE \ + "aghi 15,-216\n\t" \ + "lg 2, 8(1)\n\t" \ + "lg 3,16(1)\n\t" \ + "lg 4,24(1)\n\t" \ + "lg 5,32(1)\n\t" \ + "lg 6,40(1)\n\t" \ + "mvc 160(8,15), 48(1)\n\t" \ + "mvc 168(8,15), 56(1)\n\t" \ + "mvc 176(8,15), 64(1)\n\t" \ + "mvc 184(8,15), 72(1)\n\t" \ + "mvc 192(8,15), 80(1)\n\t" \ + "mvc 200(8,15), 88(1)\n\t" \ + "mvc 208(8,15), 96(1)\n\t" \ + "lg 1, 0(1)\n\t" \ + VALGRIND_CALL_NOREDIR_R1 \ + "aghi 15,216\n\t" \ + VALGRIND_CFI_EPILOGUE \ + "lgr %0, 2\n\t" \ + : /*out*/ "=d" (_res) \ + : /*in*/ "a" (&_argvec[0]) __FRAME_POINTER \ + : /*trash*/ "cc", "memory", __CALLER_SAVED_REGS,"6","7" \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + + +#endif /* PLAT_s390x_linux */ + +/* ------------------------- mips32-linux ----------------------- */ + +#if defined(PLAT_mips32_linux) + +/* These regs are trashed by the hidden call. */ +#define __CALLER_SAVED_REGS "$2", "$3", "$4", "$5", "$6", \ +"$7", "$8", "$9", "$10", "$11", "$12", "$13", "$14", "$15", "$24", \ +"$25", "$31" + +/* These CALL_FN_ macros assume that on mips-linux, sizeof(unsigned + long) == 4. */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[1]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "subu $29, $29, 16 \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 16\n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[2]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "subu $29, $29, 16 \n\t" \ + "lw $4, 4(%1) \n\t" /* arg1*/ \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 16 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "subu $29, $29, 16 \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 16 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[4]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "subu $29, $29, 16 \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 16 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[5]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "subu $29, $29, 16 \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $7, 16(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 16 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[6]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "lw $4, 20(%1) \n\t" \ + "subu $29, $29, 24\n\t" \ + "sw $4, 16($29) \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $7, 16(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 24 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[7]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "lw $4, 20(%1) \n\t" \ + "subu $29, $29, 32\n\t" \ + "sw $4, 16($29) \n\t" \ + "lw $4, 24(%1) \n\t" \ + "nop\n\t" \ + "sw $4, 20($29) \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $7, 16(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 32 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[8]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "lw $4, 20(%1) \n\t" \ + "subu $29, $29, 32\n\t" \ + "sw $4, 16($29) \n\t" \ + "lw $4, 24(%1) \n\t" \ + "sw $4, 20($29) \n\t" \ + "lw $4, 28(%1) \n\t" \ + "sw $4, 24($29) \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $7, 16(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 32 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[9]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "lw $4, 20(%1) \n\t" \ + "subu $29, $29, 40\n\t" \ + "sw $4, 16($29) \n\t" \ + "lw $4, 24(%1) \n\t" \ + "sw $4, 20($29) \n\t" \ + "lw $4, 28(%1) \n\t" \ + "sw $4, 24($29) \n\t" \ + "lw $4, 32(%1) \n\t" \ + "sw $4, 28($29) \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $7, 16(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 40 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[10]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "lw $4, 20(%1) \n\t" \ + "subu $29, $29, 40\n\t" \ + "sw $4, 16($29) \n\t" \ + "lw $4, 24(%1) \n\t" \ + "sw $4, 20($29) \n\t" \ + "lw $4, 28(%1) \n\t" \ + "sw $4, 24($29) \n\t" \ + "lw $4, 32(%1) \n\t" \ + "sw $4, 28($29) \n\t" \ + "lw $4, 36(%1) \n\t" \ + "sw $4, 32($29) \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $7, 16(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 40 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[11]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "lw $4, 20(%1) \n\t" \ + "subu $29, $29, 48\n\t" \ + "sw $4, 16($29) \n\t" \ + "lw $4, 24(%1) \n\t" \ + "sw $4, 20($29) \n\t" \ + "lw $4, 28(%1) \n\t" \ + "sw $4, 24($29) \n\t" \ + "lw $4, 32(%1) \n\t" \ + "sw $4, 28($29) \n\t" \ + "lw $4, 36(%1) \n\t" \ + "sw $4, 32($29) \n\t" \ + "lw $4, 40(%1) \n\t" \ + "sw $4, 36($29) \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $7, 16(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 48 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[12]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "lw $4, 20(%1) \n\t" \ + "subu $29, $29, 48\n\t" \ + "sw $4, 16($29) \n\t" \ + "lw $4, 24(%1) \n\t" \ + "sw $4, 20($29) \n\t" \ + "lw $4, 28(%1) \n\t" \ + "sw $4, 24($29) \n\t" \ + "lw $4, 32(%1) \n\t" \ + "sw $4, 28($29) \n\t" \ + "lw $4, 36(%1) \n\t" \ + "sw $4, 32($29) \n\t" \ + "lw $4, 40(%1) \n\t" \ + "sw $4, 36($29) \n\t" \ + "lw $4, 44(%1) \n\t" \ + "sw $4, 40($29) \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $7, 16(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 48 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11,arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[13]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + _argvec[12] = (unsigned long)(arg12); \ + __asm__ volatile( \ + "subu $29, $29, 8 \n\t" \ + "sw $28, 0($29) \n\t" \ + "sw $31, 4($29) \n\t" \ + "lw $4, 20(%1) \n\t" \ + "subu $29, $29, 56\n\t" \ + "sw $4, 16($29) \n\t" \ + "lw $4, 24(%1) \n\t" \ + "sw $4, 20($29) \n\t" \ + "lw $4, 28(%1) \n\t" \ + "sw $4, 24($29) \n\t" \ + "lw $4, 32(%1) \n\t" \ + "sw $4, 28($29) \n\t" \ + "lw $4, 36(%1) \n\t" \ + "sw $4, 32($29) \n\t" \ + "lw $4, 40(%1) \n\t" \ + "sw $4, 36($29) \n\t" \ + "lw $4, 44(%1) \n\t" \ + "sw $4, 40($29) \n\t" \ + "lw $4, 48(%1) \n\t" \ + "sw $4, 44($29) \n\t" \ + "lw $4, 4(%1) \n\t" \ + "lw $5, 8(%1) \n\t" \ + "lw $6, 12(%1) \n\t" \ + "lw $7, 16(%1) \n\t" \ + "lw $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "addu $29, $29, 56 \n\t" \ + "lw $28, 0($29) \n\t" \ + "lw $31, 4($29) \n\t" \ + "addu $29, $29, 8 \n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#endif /* PLAT_mips32_linux */ + +/* ------------------------- nanomips-linux -------------------- */ + +#if defined(PLAT_nanomips_linux) + +/* These regs are trashed by the hidden call. */ +#define __CALLER_SAVED_REGS "$t4", "$t5", "$a0", "$a1", "$a2", \ +"$a3", "$a4", "$a5", "$a6", "$a7", "$t0", "$t1", "$t2", "$t3", \ +"$t8","$t9", "$at" + +/* These CALL_FN_ macros assume that on mips-linux, sizeof(unsigned + long) == 4. */ + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[1]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + __asm__ volatile( \ + "lw $t9, 0(%1)\n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[2]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + __asm__ volatile( \ + "lw $t9, 0(%1)\n\t" \ + "lw $a0, 4(%1)\n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[3]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + __asm__ volatile( \ + "lw $t9, 0(%1)\n\t" \ + "lw $a0, 4(%1)\n\t" \ + "lw $a1, 8(%1)\n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[4]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + __asm__ volatile( \ + "lw $t9, 0(%1)\n\t" \ + "lw $a0, 4(%1)\n\t" \ + "lw $a1, 8(%1)\n\t" \ + "lw $a2,12(%1)\n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[5]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + __asm__ volatile( \ + "lw $t9, 0(%1)\n\t" \ + "lw $a0, 4(%1)\n\t" \ + "lw $a1, 8(%1)\n\t" \ + "lw $a2,12(%1)\n\t" \ + "lw $a3,16(%1)\n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[6]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + __asm__ volatile( \ + "lw $t9, 0(%1)\n\t" \ + "lw $a0, 4(%1)\n\t" \ + "lw $a1, 8(%1)\n\t" \ + "lw $a2,12(%1)\n\t" \ + "lw $a3,16(%1)\n\t" \ + "lw $a4,20(%1)\n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[7]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + __asm__ volatile( \ + "lw $t9, 0(%1)\n\t" \ + "lw $a0, 4(%1)\n\t" \ + "lw $a1, 8(%1)\n\t" \ + "lw $a2,12(%1)\n\t" \ + "lw $a3,16(%1)\n\t" \ + "lw $a4,20(%1)\n\t" \ + "lw $a5,24(%1)\n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[8]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + __asm__ volatile( \ + "lw $t9, 0(%1)\n\t" \ + "lw $a0, 4(%1)\n\t" \ + "lw $a1, 8(%1)\n\t" \ + "lw $a2,12(%1)\n\t" \ + "lw $a3,16(%1)\n\t" \ + "lw $a4,20(%1)\n\t" \ + "lw $a5,24(%1)\n\t" \ + "lw $a6,28(%1)\n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[9]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + __asm__ volatile( \ + "lw $t9, 0(%1)\n\t" \ + "lw $a0, 4(%1)\n\t" \ + "lw $a1, 8(%1)\n\t" \ + "lw $a2,12(%1)\n\t" \ + "lw $a3,16(%1)\n\t" \ + "lw $a4,20(%1)\n\t" \ + "lw $a5,24(%1)\n\t" \ + "lw $a6,28(%1)\n\t" \ + "lw $a7,32(%1)\n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[10]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + __asm__ volatile( \ + "addiu $sp, $sp, -16 \n\t" \ + "lw $t9,36(%1) \n\t" \ + "sw $t9, 0($sp) \n\t" \ + "lw $t9, 0(%1) \n\t" \ + "lw $a0, 4(%1) \n\t" \ + "lw $a1, 8(%1) \n\t" \ + "lw $a2,12(%1) \n\t" \ + "lw $a3,16(%1) \n\t" \ + "lw $a4,20(%1) \n\t" \ + "lw $a5,24(%1) \n\t" \ + "lw $a6,28(%1) \n\t" \ + "lw $a7,32(%1) \n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0 \n\t" \ + "addiu $sp, $sp, 16 \n\t" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[11]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + __asm__ volatile( \ + "addiu $sp, $sp, -16 \n\t" \ + "lw $t9,36(%1) \n\t" \ + "sw $t9, 0($sp) \n\t" \ + "lw $t9,40(%1) \n\t" \ + "sw $t9, 4($sp) \n\t" \ + "lw $t9, 0(%1) \n\t" \ + "lw $a0, 4(%1) \n\t" \ + "lw $a1, 8(%1) \n\t" \ + "lw $a2,12(%1) \n\t" \ + "lw $a3,16(%1) \n\t" \ + "lw $a4,20(%1) \n\t" \ + "lw $a5,24(%1) \n\t" \ + "lw $a6,28(%1) \n\t" \ + "lw $a7,32(%1) \n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0 \n\t" \ + "addiu $sp, $sp, 16 \n\t" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[12]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + __asm__ volatile( \ + "addiu $sp, $sp, -16 \n\t" \ + "lw $t9,36(%1) \n\t" \ + "sw $t9, 0($sp) \n\t" \ + "lw $t9,40(%1) \n\t" \ + "sw $t9, 4($sp) \n\t" \ + "lw $t9,44(%1) \n\t" \ + "sw $t9, 8($sp) \n\t" \ + "lw $t9, 0(%1) \n\t" \ + "lw $a0, 4(%1) \n\t" \ + "lw $a1, 8(%1) \n\t" \ + "lw $a2,12(%1) \n\t" \ + "lw $a3,16(%1) \n\t" \ + "lw $a4,20(%1) \n\t" \ + "lw $a5,24(%1) \n\t" \ + "lw $a6,28(%1) \n\t" \ + "lw $a7,32(%1) \n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0 \n\t" \ + "addiu $sp, $sp, 16 \n\t" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11,arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long _argvec[13]; \ + volatile unsigned long _res; \ + _argvec[0] = (unsigned long)_orig.nraddr; \ + _argvec[1] = (unsigned long)(arg1); \ + _argvec[2] = (unsigned long)(arg2); \ + _argvec[3] = (unsigned long)(arg3); \ + _argvec[4] = (unsigned long)(arg4); \ + _argvec[5] = (unsigned long)(arg5); \ + _argvec[6] = (unsigned long)(arg6); \ + _argvec[7] = (unsigned long)(arg7); \ + _argvec[8] = (unsigned long)(arg8); \ + _argvec[9] = (unsigned long)(arg9); \ + _argvec[10] = (unsigned long)(arg10); \ + _argvec[11] = (unsigned long)(arg11); \ + _argvec[12] = (unsigned long)(arg12); \ + __asm__ volatile( \ + "addiu $sp, $sp, -16 \n\t" \ + "lw $t9,36(%1) \n\t" \ + "sw $t9, 0($sp) \n\t" \ + "lw $t9,40(%1) \n\t" \ + "sw $t9, 4($sp) \n\t" \ + "lw $t9,44(%1) \n\t" \ + "sw $t9, 8($sp) \n\t" \ + "lw $t9,48(%1) \n\t" \ + "sw $t9,12($sp) \n\t" \ + "lw $t9, 0(%1) \n\t" \ + "lw $a0, 4(%1) \n\t" \ + "lw $a1, 8(%1) \n\t" \ + "lw $a2,12(%1) \n\t" \ + "lw $a3,16(%1) \n\t" \ + "lw $a4,20(%1) \n\t" \ + "lw $a5,24(%1) \n\t" \ + "lw $a6,28(%1) \n\t" \ + "lw $a7,32(%1) \n\t" \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $a0 \n\t" \ + "addiu $sp, $sp, 16 \n\t" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) _res; \ + } while (0) + +#endif /* PLAT_nanomips_linux */ + +/* ------------------------- mips64-linux ------------------------- */ + +#if defined(PLAT_mips64_linux) + +/* These regs are trashed by the hidden call. */ +#define __CALLER_SAVED_REGS "$2", "$3", "$4", "$5", "$6", \ +"$7", "$8", "$9", "$10", "$11", "$12", "$13", "$14", "$15", "$24", \ +"$25", "$31" + +/* These CALL_FN_ macros assume that on mips64-linux, + sizeof(long long) == 8. */ + +#define MIPS64_LONG2REG_CAST(x) ((long long)(long)x) + +#define CALL_FN_W_v(lval, orig) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[1]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + __asm__ volatile( \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "0" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_W(lval, orig, arg1) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[2]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + __asm__ volatile( \ + "ld $4, 8(%1)\n\t" /* arg1*/ \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_WW(lval, orig, arg1,arg2) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[3]; \ + volatile unsigned long long _res; \ + _argvec[0] = _orig.nraddr; \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + __asm__ volatile( \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + + +#define CALL_FN_W_WWW(lval, orig, arg1,arg2,arg3) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[4]; \ + volatile unsigned long long _res; \ + _argvec[0] = _orig.nraddr; \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + __asm__ volatile( \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_WWWW(lval, orig, arg1,arg2,arg3,arg4) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[5]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + _argvec[4] = MIPS64_LONG2REG_CAST(arg4); \ + __asm__ volatile( \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $7, 32(%1)\n\t" \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_5W(lval, orig, arg1,arg2,arg3,arg4,arg5) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[6]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + _argvec[4] = MIPS64_LONG2REG_CAST(arg4); \ + _argvec[5] = MIPS64_LONG2REG_CAST(arg5); \ + __asm__ volatile( \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $7, 32(%1)\n\t" \ + "ld $8, 40(%1)\n\t" \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_6W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[7]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + _argvec[4] = MIPS64_LONG2REG_CAST(arg4); \ + _argvec[5] = MIPS64_LONG2REG_CAST(arg5); \ + _argvec[6] = MIPS64_LONG2REG_CAST(arg6); \ + __asm__ volatile( \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $7, 32(%1)\n\t" \ + "ld $8, 40(%1)\n\t" \ + "ld $9, 48(%1)\n\t" \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_7W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[8]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + _argvec[4] = MIPS64_LONG2REG_CAST(arg4); \ + _argvec[5] = MIPS64_LONG2REG_CAST(arg5); \ + _argvec[6] = MIPS64_LONG2REG_CAST(arg6); \ + _argvec[7] = MIPS64_LONG2REG_CAST(arg7); \ + __asm__ volatile( \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $7, 32(%1)\n\t" \ + "ld $8, 40(%1)\n\t" \ + "ld $9, 48(%1)\n\t" \ + "ld $10, 56(%1)\n\t" \ + "ld $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_8W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[9]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + _argvec[4] = MIPS64_LONG2REG_CAST(arg4); \ + _argvec[5] = MIPS64_LONG2REG_CAST(arg5); \ + _argvec[6] = MIPS64_LONG2REG_CAST(arg6); \ + _argvec[7] = MIPS64_LONG2REG_CAST(arg7); \ + _argvec[8] = MIPS64_LONG2REG_CAST(arg8); \ + __asm__ volatile( \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $7, 32(%1)\n\t" \ + "ld $8, 40(%1)\n\t" \ + "ld $9, 48(%1)\n\t" \ + "ld $10, 56(%1)\n\t" \ + "ld $11, 64(%1)\n\t" \ + "ld $25, 0(%1) \n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_9W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[10]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + _argvec[4] = MIPS64_LONG2REG_CAST(arg4); \ + _argvec[5] = MIPS64_LONG2REG_CAST(arg5); \ + _argvec[6] = MIPS64_LONG2REG_CAST(arg6); \ + _argvec[7] = MIPS64_LONG2REG_CAST(arg7); \ + _argvec[8] = MIPS64_LONG2REG_CAST(arg8); \ + _argvec[9] = MIPS64_LONG2REG_CAST(arg9); \ + __asm__ volatile( \ + "dsubu $29, $29, 8\n\t" \ + "ld $4, 72(%1)\n\t" \ + "sd $4, 0($29)\n\t" \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $7, 32(%1)\n\t" \ + "ld $8, 40(%1)\n\t" \ + "ld $9, 48(%1)\n\t" \ + "ld $10, 56(%1)\n\t" \ + "ld $11, 64(%1)\n\t" \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "daddu $29, $29, 8\n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_10W(lval, orig, arg1,arg2,arg3,arg4,arg5,arg6, \ + arg7,arg8,arg9,arg10) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[11]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + _argvec[4] = MIPS64_LONG2REG_CAST(arg4); \ + _argvec[5] = MIPS64_LONG2REG_CAST(arg5); \ + _argvec[6] = MIPS64_LONG2REG_CAST(arg6); \ + _argvec[7] = MIPS64_LONG2REG_CAST(arg7); \ + _argvec[8] = MIPS64_LONG2REG_CAST(arg8); \ + _argvec[9] = MIPS64_LONG2REG_CAST(arg9); \ + _argvec[10] = MIPS64_LONG2REG_CAST(arg10); \ + __asm__ volatile( \ + "dsubu $29, $29, 16\n\t" \ + "ld $4, 72(%1)\n\t" \ + "sd $4, 0($29)\n\t" \ + "ld $4, 80(%1)\n\t" \ + "sd $4, 8($29)\n\t" \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $7, 32(%1)\n\t" \ + "ld $8, 40(%1)\n\t" \ + "ld $9, 48(%1)\n\t" \ + "ld $10, 56(%1)\n\t" \ + "ld $11, 64(%1)\n\t" \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "daddu $29, $29, 16\n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_11W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[12]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + _argvec[4] = MIPS64_LONG2REG_CAST(arg4); \ + _argvec[5] = MIPS64_LONG2REG_CAST(arg5); \ + _argvec[6] = MIPS64_LONG2REG_CAST(arg6); \ + _argvec[7] = MIPS64_LONG2REG_CAST(arg7); \ + _argvec[8] = MIPS64_LONG2REG_CAST(arg8); \ + _argvec[9] = MIPS64_LONG2REG_CAST(arg9); \ + _argvec[10] = MIPS64_LONG2REG_CAST(arg10); \ + _argvec[11] = MIPS64_LONG2REG_CAST(arg11); \ + __asm__ volatile( \ + "dsubu $29, $29, 24\n\t" \ + "ld $4, 72(%1)\n\t" \ + "sd $4, 0($29)\n\t" \ + "ld $4, 80(%1)\n\t" \ + "sd $4, 8($29)\n\t" \ + "ld $4, 88(%1)\n\t" \ + "sd $4, 16($29)\n\t" \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $7, 32(%1)\n\t" \ + "ld $8, 40(%1)\n\t" \ + "ld $9, 48(%1)\n\t" \ + "ld $10, 56(%1)\n\t" \ + "ld $11, 64(%1)\n\t" \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "daddu $29, $29, 24\n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#define CALL_FN_W_12W(lval, orig, arg1,arg2,arg3,arg4,arg5, \ + arg6,arg7,arg8,arg9,arg10, \ + arg11,arg12) \ + do { \ + volatile OrigFn _orig = (orig); \ + volatile unsigned long long _argvec[13]; \ + volatile unsigned long long _res; \ + _argvec[0] = MIPS64_LONG2REG_CAST(_orig.nraddr); \ + _argvec[1] = MIPS64_LONG2REG_CAST(arg1); \ + _argvec[2] = MIPS64_LONG2REG_CAST(arg2); \ + _argvec[3] = MIPS64_LONG2REG_CAST(arg3); \ + _argvec[4] = MIPS64_LONG2REG_CAST(arg4); \ + _argvec[5] = MIPS64_LONG2REG_CAST(arg5); \ + _argvec[6] = MIPS64_LONG2REG_CAST(arg6); \ + _argvec[7] = MIPS64_LONG2REG_CAST(arg7); \ + _argvec[8] = MIPS64_LONG2REG_CAST(arg8); \ + _argvec[9] = MIPS64_LONG2REG_CAST(arg9); \ + _argvec[10] = MIPS64_LONG2REG_CAST(arg10); \ + _argvec[11] = MIPS64_LONG2REG_CAST(arg11); \ + _argvec[12] = MIPS64_LONG2REG_CAST(arg12); \ + __asm__ volatile( \ + "dsubu $29, $29, 32\n\t" \ + "ld $4, 72(%1)\n\t" \ + "sd $4, 0($29)\n\t" \ + "ld $4, 80(%1)\n\t" \ + "sd $4, 8($29)\n\t" \ + "ld $4, 88(%1)\n\t" \ + "sd $4, 16($29)\n\t" \ + "ld $4, 96(%1)\n\t" \ + "sd $4, 24($29)\n\t" \ + "ld $4, 8(%1)\n\t" \ + "ld $5, 16(%1)\n\t" \ + "ld $6, 24(%1)\n\t" \ + "ld $7, 32(%1)\n\t" \ + "ld $8, 40(%1)\n\t" \ + "ld $9, 48(%1)\n\t" \ + "ld $10, 56(%1)\n\t" \ + "ld $11, 64(%1)\n\t" \ + "ld $25, 0(%1)\n\t" /* target->t9 */ \ + VALGRIND_CALL_NOREDIR_T9 \ + "daddu $29, $29, 32\n\t" \ + "move %0, $2\n" \ + : /*out*/ "=r" (_res) \ + : /*in*/ "r" (&_argvec[0]) \ + : /*trash*/ "memory", __CALLER_SAVED_REGS \ + ); \ + lval = (__typeof__(lval)) (long)_res; \ + } while (0) + +#endif /* PLAT_mips64_linux */ + +/* ------------------------------------------------------------------ */ +/* ARCHITECTURE INDEPENDENT MACROS for CLIENT REQUESTS. */ +/* */ +/* ------------------------------------------------------------------ */ + +/* Some request codes. There are many more of these, but most are not + exposed to end-user view. These are the public ones, all of the + form 0x1000 + small_number. + + Core ones are in the range 0x00000000--0x0000ffff. The non-public + ones start at 0x2000. +*/ + +/* These macros are used by tools -- they must be public, but don't + embed them into other programs. */ +#define VG_USERREQ_TOOL_BASE(a,b) \ + ((unsigned int)(((a)&0xff) << 24 | ((b)&0xff) << 16)) +#define VG_IS_TOOL_USERREQ(a, b, v) \ + (VG_USERREQ_TOOL_BASE(a,b) == ((v) & 0xffff0000)) + +/* !! ABIWARNING !! ABIWARNING !! ABIWARNING !! ABIWARNING !! + This enum comprises an ABI exported by Valgrind to programs + which use client requests. DO NOT CHANGE THE NUMERIC VALUES OF THESE + ENTRIES, NOR DELETE ANY -- add new ones at the end of the most + relevant group. */ +typedef + enum { VG_USERREQ__RUNNING_ON_VALGRIND = 0x1001, + VG_USERREQ__DISCARD_TRANSLATIONS = 0x1002, + + /* These allow any function to be called from the simulated + CPU but run on the real CPU. Nb: the first arg passed to + the function is always the ThreadId of the running + thread! So CLIENT_CALL0 actually requires a 1 arg + function, etc. */ + VG_USERREQ__CLIENT_CALL0 = 0x1101, + VG_USERREQ__CLIENT_CALL1 = 0x1102, + VG_USERREQ__CLIENT_CALL2 = 0x1103, + VG_USERREQ__CLIENT_CALL3 = 0x1104, + + /* Can be useful in regression testing suites -- eg. can + send Valgrind's output to /dev/null and still count + errors. */ + VG_USERREQ__COUNT_ERRORS = 0x1201, + + /* Allows the client program and/or gdbserver to execute a monitor + command. */ + VG_USERREQ__GDB_MONITOR_COMMAND = 0x1202, + + /* Allows the client program to change a dynamic command line + option. */ + VG_USERREQ__CLO_CHANGE = 0x1203, + + /* These are useful and can be interpreted by any tool that + tracks malloc() et al, by using vg_replace_malloc.c. */ + VG_USERREQ__MALLOCLIKE_BLOCK = 0x1301, + VG_USERREQ__RESIZEINPLACE_BLOCK = 0x130b, + VG_USERREQ__FREELIKE_BLOCK = 0x1302, + /* Memory pool support. */ + VG_USERREQ__CREATE_MEMPOOL = 0x1303, + VG_USERREQ__DESTROY_MEMPOOL = 0x1304, + VG_USERREQ__MEMPOOL_ALLOC = 0x1305, + VG_USERREQ__MEMPOOL_FREE = 0x1306, + VG_USERREQ__MEMPOOL_TRIM = 0x1307, + VG_USERREQ__MOVE_MEMPOOL = 0x1308, + VG_USERREQ__MEMPOOL_CHANGE = 0x1309, + VG_USERREQ__MEMPOOL_EXISTS = 0x130a, + + /* Allow printfs to valgrind log. */ + /* The first two pass the va_list argument by value, which + assumes it is the same size as or smaller than a UWord, + which generally isn't the case. Hence are deprecated. + The second two pass the vargs by reference and so are + immune to this problem. */ + /* both :: char* fmt, va_list vargs (DEPRECATED) */ + VG_USERREQ__PRINTF = 0x1401, + VG_USERREQ__PRINTF_BACKTRACE = 0x1402, + /* both :: char* fmt, va_list* vargs */ + VG_USERREQ__PRINTF_VALIST_BY_REF = 0x1403, + VG_USERREQ__PRINTF_BACKTRACE_VALIST_BY_REF = 0x1404, + + /* Stack support. */ + VG_USERREQ__STACK_REGISTER = 0x1501, + VG_USERREQ__STACK_DEREGISTER = 0x1502, + VG_USERREQ__STACK_CHANGE = 0x1503, + + /* Wine support */ + VG_USERREQ__LOAD_PDB_DEBUGINFO = 0x1601, + + /* Querying of debug info. */ + VG_USERREQ__MAP_IP_TO_SRCLOC = 0x1701, + + /* Disable/enable error reporting level. Takes a single + Word arg which is the delta to this thread's error + disablement indicator. Hence 1 disables or further + disables errors, and -1 moves back towards enablement. + Other values are not allowed. */ + VG_USERREQ__CHANGE_ERR_DISABLEMENT = 0x1801, + + /* Some requests used for Valgrind internal, such as + self-test or self-hosting. */ + /* Initialise IR injection */ + VG_USERREQ__VEX_INIT_FOR_IRI = 0x1901, + /* Used by Inner Valgrind to inform Outer Valgrind where to + find the list of inner guest threads */ + VG_USERREQ__INNER_THREADS = 0x1902 + } Vg_ClientRequest; + +#if !defined(__GNUC__) +# define __extension__ /* */ +#endif + + +/* Returns the number of Valgrinds this code is running under. That + is, 0 if running natively, 1 if running under Valgrind, 2 if + running under Valgrind which is running under another Valgrind, + etc. */ +#define RUNNING_ON_VALGRIND \ + (unsigned)VALGRIND_DO_CLIENT_REQUEST_EXPR(0 /* if not */, \ + VG_USERREQ__RUNNING_ON_VALGRIND, \ + 0, 0, 0, 0, 0) \ + + +/* Discard translation of code in the range [_qzz_addr .. _qzz_addr + + _qzz_len - 1]. Useful if you are debugging a JITter or some such, + since it provides a way to make sure valgrind will retranslate the + invalidated area. Returns no value. */ +#define VALGRIND_DISCARD_TRANSLATIONS(_qzz_addr,_qzz_len) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__DISCARD_TRANSLATIONS, \ + _qzz_addr, _qzz_len, 0, 0, 0) + +#define VALGRIND_INNER_THREADS(_qzz_addr) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__INNER_THREADS, \ + _qzz_addr, 0, 0, 0, 0) + + +/* These requests are for getting Valgrind itself to print something. + Possibly with a backtrace. This is a really ugly hack. The return value + is the number of characters printed, excluding the "**** " part at the + start and the backtrace (if present). */ + +#if defined(__GNUC__) || defined(__INTEL_COMPILER) && !defined(_MSC_VER) +/* Modern GCC will optimize the static routine out if unused, + and unused attribute will shut down warnings about it. */ +static int VALGRIND_PRINTF(const char *format, ...) + __attribute__((format(__printf__, 1, 2), __unused__)); +#endif +static int +#if defined(_MSC_VER) +__inline +#endif +VALGRIND_PRINTF(const char *format, ...) +{ +#if defined(NVALGRIND) + (void)format; + return 0; +#else /* NVALGRIND */ +#if defined(_MSC_VER) || defined(__MINGW64__) + uintptr_t _qzz_res; +#else + unsigned long _qzz_res; +#endif + va_list vargs; + va_start(vargs, format); +#if defined(_MSC_VER) || defined(__MINGW64__) + _qzz_res = VALGRIND_DO_CLIENT_REQUEST_EXPR(0, + VG_USERREQ__PRINTF_VALIST_BY_REF, + (uintptr_t)format, + (uintptr_t)&vargs, + 0, 0, 0); +#else + _qzz_res = VALGRIND_DO_CLIENT_REQUEST_EXPR(0, + VG_USERREQ__PRINTF_VALIST_BY_REF, + (unsigned long)format, + (unsigned long)&vargs, + 0, 0, 0); +#endif + va_end(vargs); + return (int)_qzz_res; +#endif /* NVALGRIND */ +} + +#if defined(__GNUC__) || defined(__INTEL_COMPILER) && !defined(_MSC_VER) +static int VALGRIND_PRINTF_BACKTRACE(const char *format, ...) + __attribute__((format(__printf__, 1, 2), __unused__)); +#endif +static int +#if defined(_MSC_VER) +__inline +#endif +VALGRIND_PRINTF_BACKTRACE(const char *format, ...) +{ +#if defined(NVALGRIND) + (void)format; + return 0; +#else /* NVALGRIND */ +#if defined(_MSC_VER) || defined(__MINGW64__) + uintptr_t _qzz_res; +#else + unsigned long _qzz_res; +#endif + va_list vargs; + va_start(vargs, format); +#if defined(_MSC_VER) || defined(__MINGW64__) + _qzz_res = VALGRIND_DO_CLIENT_REQUEST_EXPR(0, + VG_USERREQ__PRINTF_BACKTRACE_VALIST_BY_REF, + (uintptr_t)format, + (uintptr_t)&vargs, + 0, 0, 0); +#else + _qzz_res = VALGRIND_DO_CLIENT_REQUEST_EXPR(0, + VG_USERREQ__PRINTF_BACKTRACE_VALIST_BY_REF, + (unsigned long)format, + (unsigned long)&vargs, + 0, 0, 0); +#endif + va_end(vargs); + return (int)_qzz_res; +#endif /* NVALGRIND */ +} + + +/* These requests allow control to move from the simulated CPU to the + real CPU, calling an arbitrary function. + + Note that the current ThreadId is inserted as the first argument. + So this call: + + VALGRIND_NON_SIMD_CALL2(f, arg1, arg2) + + requires f to have this signature: + + Word f(Word tid, Word arg1, Word arg2) + + where "Word" is a word-sized type. + + Note that these client requests are not entirely reliable. For example, + if you call a function with them that subsequently calls printf(), + there's a high chance Valgrind will crash. Generally, your prospects of + these working are made higher if the called function does not refer to + any global variables, and does not refer to any libc or other functions + (printf et al). Any kind of entanglement with libc or dynamic linking is + likely to have a bad outcome, for tricky reasons which we've grappled + with a lot in the past. +*/ +#define VALGRIND_NON_SIMD_CALL0(_qyy_fn) \ + VALGRIND_DO_CLIENT_REQUEST_EXPR(0 /* default return */, \ + VG_USERREQ__CLIENT_CALL0, \ + _qyy_fn, \ + 0, 0, 0, 0) + +#define VALGRIND_NON_SIMD_CALL1(_qyy_fn, _qyy_arg1) \ + VALGRIND_DO_CLIENT_REQUEST_EXPR(0 /* default return */, \ + VG_USERREQ__CLIENT_CALL1, \ + _qyy_fn, \ + _qyy_arg1, 0, 0, 0) + +#define VALGRIND_NON_SIMD_CALL2(_qyy_fn, _qyy_arg1, _qyy_arg2) \ + VALGRIND_DO_CLIENT_REQUEST_EXPR(0 /* default return */, \ + VG_USERREQ__CLIENT_CALL2, \ + _qyy_fn, \ + _qyy_arg1, _qyy_arg2, 0, 0) + +#define VALGRIND_NON_SIMD_CALL3(_qyy_fn, _qyy_arg1, _qyy_arg2, _qyy_arg3) \ + VALGRIND_DO_CLIENT_REQUEST_EXPR(0 /* default return */, \ + VG_USERREQ__CLIENT_CALL3, \ + _qyy_fn, \ + _qyy_arg1, _qyy_arg2, \ + _qyy_arg3, 0) + + +/* Counts the number of errors that have been recorded by a tool. Nb: + the tool must record the errors with VG_(maybe_record_error)() or + VG_(unique_error)() for them to be counted. */ +#define VALGRIND_COUNT_ERRORS \ + (unsigned)VALGRIND_DO_CLIENT_REQUEST_EXPR( \ + 0 /* default return */, \ + VG_USERREQ__COUNT_ERRORS, \ + 0, 0, 0, 0, 0) + +/* Several Valgrind tools (Memcheck, Massif, Helgrind, DRD) rely on knowing + when heap blocks are allocated in order to give accurate results. This + happens automatically for the standard allocator functions such as + malloc(), calloc(), realloc(), memalign(), new, new[], free(), delete, + delete[], etc. + + But if your program uses a custom allocator, this doesn't automatically + happen, and Valgrind will not do as well. For example, if you allocate + superblocks with mmap() and then allocates chunks of the superblocks, all + Valgrind's observations will be at the mmap() level and it won't know that + the chunks should be considered separate entities. In Memcheck's case, + that means you probably won't get heap block overrun detection (because + there won't be redzones marked as unaddressable) and you definitely won't + get any leak detection. + + The following client requests allow a custom allocator to be annotated so + that it can be handled accurately by Valgrind. + + VALGRIND_MALLOCLIKE_BLOCK marks a region of memory as having been allocated + by a malloc()-like function. For Memcheck (an illustrative case), this + does two things: + + - It records that the block has been allocated. This means any addresses + within the block mentioned in error messages will be + identified as belonging to the block. It also means that if the block + isn't freed it will be detected by the leak checker. + + - It marks the block as being addressable and undefined (if 'is_zeroed' is + not set), or addressable and defined (if 'is_zeroed' is set). This + controls how accesses to the block by the program are handled. + + 'addr' is the start of the usable block (ie. after any + redzone), 'sizeB' is its size. 'rzB' is the redzone size if the allocator + can apply redzones -- these are blocks of padding at the start and end of + each block. Adding redzones is recommended as it makes it much more likely + Valgrind will spot block overruns. `is_zeroed' indicates if the memory is + zeroed (or filled with another predictable value), as is the case for + calloc(). + + VALGRIND_MALLOCLIKE_BLOCK should be put immediately after the point where a + heap block -- that will be used by the client program -- is allocated. + It's best to put it at the outermost level of the allocator if possible; + for example, if you have a function my_alloc() which calls + internal_alloc(), and the client request is put inside internal_alloc(), + stack traces relating to the heap block will contain entries for both + my_alloc() and internal_alloc(), which is probably not what you want. + + For Memcheck users: if you use VALGRIND_MALLOCLIKE_BLOCK to carve out + custom blocks from within a heap block, B, that has been allocated with + malloc/calloc/new/etc, then block B will be *ignored* during leak-checking + -- the custom blocks will take precedence. + + VALGRIND_FREELIKE_BLOCK is the partner to VALGRIND_MALLOCLIKE_BLOCK. For + Memcheck, it does two things: + + - It records that the block has been deallocated. This assumes that the + block was annotated as having been allocated via + VALGRIND_MALLOCLIKE_BLOCK. Otherwise, an error will be issued. + + - It marks the block as being unaddressable. + + VALGRIND_FREELIKE_BLOCK should be put immediately after the point where a + heap block is deallocated. + + VALGRIND_RESIZEINPLACE_BLOCK informs a tool about reallocation. For + Memcheck, it does four things: + + - It records that the size of a block has been changed. This assumes that + the block was annotated as having been allocated via + VALGRIND_MALLOCLIKE_BLOCK. Otherwise, an error will be issued. + + - If the block shrunk, it marks the freed memory as being unaddressable. + + - If the block grew, it marks the new area as undefined and defines a red + zone past the end of the new block. + + - The V-bits of the overlap between the old and the new block are preserved. + + VALGRIND_RESIZEINPLACE_BLOCK should be put after allocation of the new block + and before deallocation of the old block. + + In many cases, these three client requests will not be enough to get your + allocator working well with Memcheck. More specifically, if your allocator + writes to freed blocks in any way then a VALGRIND_MAKE_MEM_UNDEFINED call + will be necessary to mark the memory as addressable just before the zeroing + occurs, otherwise you'll get a lot of invalid write errors. For example, + you'll need to do this if your allocator recycles freed blocks, but it + zeroes them before handing them back out (via VALGRIND_MALLOCLIKE_BLOCK). + Alternatively, if your allocator reuses freed blocks for allocator-internal + data structures, VALGRIND_MAKE_MEM_UNDEFINED calls will also be necessary. + + Really, what's happening is a blurring of the lines between the client + program and the allocator... after VALGRIND_FREELIKE_BLOCK is called, the + memory should be considered unaddressable to the client program, but the + allocator knows more than the rest of the client program and so may be able + to safely access it. Extra client requests are necessary for Valgrind to + understand the distinction between the allocator and the rest of the + program. + + Ignored if addr == 0. +*/ +#define VALGRIND_MALLOCLIKE_BLOCK(addr, sizeB, rzB, is_zeroed) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__MALLOCLIKE_BLOCK, \ + addr, sizeB, rzB, is_zeroed, 0) + +/* See the comment for VALGRIND_MALLOCLIKE_BLOCK for details. + Ignored if addr == 0. +*/ +#define VALGRIND_RESIZEINPLACE_BLOCK(addr, oldSizeB, newSizeB, rzB) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__RESIZEINPLACE_BLOCK, \ + addr, oldSizeB, newSizeB, rzB, 0) + +/* See the comment for VALGRIND_MALLOCLIKE_BLOCK for details. + Ignored if addr == 0. +*/ +#define VALGRIND_FREELIKE_BLOCK(addr, rzB) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__FREELIKE_BLOCK, \ + addr, rzB, 0, 0, 0) + +/* Create a memory pool. */ +#define VALGRIND_CREATE_MEMPOOL(pool, rzB, is_zeroed) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__CREATE_MEMPOOL, \ + pool, rzB, is_zeroed, 0, 0) + +/* Create a memory pool with some flags specifying extended behaviour. + When flags is zero, the behaviour is identical to VALGRIND_CREATE_MEMPOOL. + + The flag VALGRIND_MEMPOOL_METAPOOL specifies that the pieces of memory + associated with the pool using VALGRIND_MEMPOOL_ALLOC will be used + by the application as superblocks to dole out MALLOC_LIKE blocks using + VALGRIND_MALLOCLIKE_BLOCK. In other words, a meta pool is a "2 levels" + pool : first level is the blocks described by VALGRIND_MEMPOOL_ALLOC. + The second level blocks are described using VALGRIND_MALLOCLIKE_BLOCK. + Note that the association between the pool and the second level blocks + is implicit : second level blocks will be located inside first level + blocks. It is necessary to use the VALGRIND_MEMPOOL_METAPOOL flag + for such 2 levels pools, as otherwise valgrind will detect overlapping + memory blocks, and will abort execution (e.g. during leak search). + + Such a meta pool can also be marked as an 'auto free' pool using the flag + VALGRIND_MEMPOOL_AUTO_FREE, which must be OR-ed together with the + VALGRIND_MEMPOOL_METAPOOL. For an 'auto free' pool, VALGRIND_MEMPOOL_FREE + will automatically free the second level blocks that are contained + inside the first level block freed with VALGRIND_MEMPOOL_FREE. + In other words, calling VALGRIND_MEMPOOL_FREE will cause implicit calls + to VALGRIND_FREELIKE_BLOCK for all the second level blocks included + in the first level block. + Note: it is an error to use the VALGRIND_MEMPOOL_AUTO_FREE flag + without the VALGRIND_MEMPOOL_METAPOOL flag. +*/ +#define VALGRIND_MEMPOOL_AUTO_FREE 1 +#define VALGRIND_MEMPOOL_METAPOOL 2 +#define VALGRIND_CREATE_MEMPOOL_EXT(pool, rzB, is_zeroed, flags) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__CREATE_MEMPOOL, \ + pool, rzB, is_zeroed, flags, 0) + +/* Destroy a memory pool. */ +#define VALGRIND_DESTROY_MEMPOOL(pool) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__DESTROY_MEMPOOL, \ + pool, 0, 0, 0, 0) + +/* Associate a piece of memory with a memory pool. */ +#define VALGRIND_MEMPOOL_ALLOC(pool, addr, size) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__MEMPOOL_ALLOC, \ + pool, addr, size, 0, 0) + +/* Disassociate a piece of memory from a memory pool. */ +#define VALGRIND_MEMPOOL_FREE(pool, addr) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__MEMPOOL_FREE, \ + pool, addr, 0, 0, 0) + +/* Disassociate any pieces outside a particular range. */ +#define VALGRIND_MEMPOOL_TRIM(pool, addr, size) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__MEMPOOL_TRIM, \ + pool, addr, size, 0, 0) + +/* Resize and/or move a piece associated with a memory pool. */ +#define VALGRIND_MOVE_MEMPOOL(poolA, poolB) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__MOVE_MEMPOOL, \ + poolA, poolB, 0, 0, 0) + +/* Resize and/or move a piece associated with a memory pool. */ +#define VALGRIND_MEMPOOL_CHANGE(pool, addrA, addrB, size) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__MEMPOOL_CHANGE, \ + pool, addrA, addrB, size, 0) + +/* Return 1 if a mempool exists, else 0. */ +#define VALGRIND_MEMPOOL_EXISTS(pool) \ + (unsigned)VALGRIND_DO_CLIENT_REQUEST_EXPR(0, \ + VG_USERREQ__MEMPOOL_EXISTS, \ + pool, 0, 0, 0, 0) + +/* Mark a piece of memory as being a stack. Returns a stack id. + start is the lowest addressable stack byte, end is the highest + addressable stack byte. */ +#define VALGRIND_STACK_REGISTER(start, end) \ + (unsigned)VALGRIND_DO_CLIENT_REQUEST_EXPR(0, \ + VG_USERREQ__STACK_REGISTER, \ + start, end, 0, 0, 0) + +/* Unmark the piece of memory associated with a stack id as being a + stack. */ +#define VALGRIND_STACK_DEREGISTER(id) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__STACK_DEREGISTER, \ + id, 0, 0, 0, 0) + +/* Change the start and end address of the stack id. + start is the new lowest addressable stack byte, end is the new highest + addressable stack byte. */ +#define VALGRIND_STACK_CHANGE(id, start, end) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__STACK_CHANGE, \ + id, start, end, 0, 0) + +/* Load PDB debug info for Wine PE image_map. */ +#define VALGRIND_LOAD_PDB_DEBUGINFO(fd, ptr, total_size, delta) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__LOAD_PDB_DEBUGINFO, \ + fd, ptr, total_size, delta, 0) + +/* Map a code address to a source file name and line number. buf64 + must point to a 64-byte buffer in the caller's address space. The + result will be dumped in there and is guaranteed to be zero + terminated. If no info is found, the first byte is set to zero. */ +#define VALGRIND_MAP_IP_TO_SRCLOC(addr, buf64) \ + (unsigned)VALGRIND_DO_CLIENT_REQUEST_EXPR(0, \ + VG_USERREQ__MAP_IP_TO_SRCLOC, \ + addr, buf64, 0, 0, 0) + +/* Disable error reporting for this thread. Behaves in a stack like + way, so you can safely call this multiple times provided that + VALGRIND_ENABLE_ERROR_REPORTING is called the same number of times + to re-enable reporting. The first call of this macro disables + reporting. Subsequent calls have no effect except to increase the + number of VALGRIND_ENABLE_ERROR_REPORTING calls needed to re-enable + reporting. Child threads do not inherit this setting from their + parents -- they are always created with reporting enabled. */ +#define VALGRIND_DISABLE_ERROR_REPORTING \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__CHANGE_ERR_DISABLEMENT, \ + 1, 0, 0, 0, 0) + +/* Re-enable error reporting, as per comments on + VALGRIND_DISABLE_ERROR_REPORTING. */ +#define VALGRIND_ENABLE_ERROR_REPORTING \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__CHANGE_ERR_DISABLEMENT, \ + -1, 0, 0, 0, 0) + +/* Execute a monitor command from the client program. + If a connection is opened with GDB, the output will be sent + according to the output mode set for vgdb. + If no connection is opened, output will go to the log output. + Returns 1 if command not recognised, 0 otherwise. */ +#define VALGRIND_MONITOR_COMMAND(command) \ + VALGRIND_DO_CLIENT_REQUEST_EXPR(0, VG_USERREQ__GDB_MONITOR_COMMAND, \ + command, 0, 0, 0, 0) + + +/* Change the value of a dynamic command line option. + Note that unknown or not dynamically changeable options + will cause a warning message to be output. */ +#define VALGRIND_CLO_CHANGE(option) \ + VALGRIND_DO_CLIENT_REQUEST_STMT(VG_USERREQ__CLO_CHANGE, \ + option, 0, 0, 0, 0) + + +#undef PLAT_x86_darwin +#undef PLAT_amd64_darwin +#undef PLAT_x86_win32 +#undef PLAT_amd64_win64 +#undef PLAT_x86_linux +#undef PLAT_amd64_linux +#undef PLAT_ppc32_linux +#undef PLAT_ppc64be_linux +#undef PLAT_ppc64le_linux +#undef PLAT_arm_linux +#undef PLAT_s390x_linux +#undef PLAT_mips32_linux +#undef PLAT_mips64_linux +#undef PLAT_nanomips_linux +#undef PLAT_x86_solaris +#undef PLAT_amd64_solaris + +#endif /* __VALGRIND_H */ diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/bundled_inputs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/bundled_inputs.py new file mode 100644 index 0000000000000000000000000000000000000000..7fa743f2804693e0096aa3e486f7394d00f49285 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/bundled_inputs.py @@ -0,0 +1,472 @@ +#!/usr/bin/env python3 +# mypy: allow-untyped-defs +from typing import Any, TypeVar, NamedTuple +from collections.abc import Callable, Sequence +import textwrap +import torch +from torch._C import TupleType, ListType +from torch.jit._recursive import wrap_cpp_module + + +T = TypeVar("T") + +MAX_RAW_TENSOR_SIZE = 16 + +class InflatableArg(NamedTuple): + """Helper type for bundled inputs. + + 'value' is the compressed/deflated input that is stored in the model. Value + must be of the same type as the argument to the function that it is a deflated + input for. + + 'fmt' is a formattable code string that is executed to inflate the compressed data into + the appropriate input. It can use 'value' as an input to the format str. It must result + in a value of the same type as 'value'. + + 'fmt_fn' is a formattable function code string that is executed to inflate the compressed + data into the appropriate input. It must result in a value of the same type as 'value'. + The function name should be the formattable part of the string. + + Note: Only top level InflatableArgs can be inflated. i.e. you cannot place + an inflatable arg inside of some other structure. You should instead create + an inflatable arg such that the fmt code string returns the full structure + of your input. + """ + + value: Any + fmt: str = "{}" + fmt_fn: str = "" + + +def bundle_inputs( + model: torch.jit.ScriptModule, + inputs: Sequence[tuple[Any, ...]] | dict[Callable, Sequence[tuple[Any, ...]] | None] | None, + info: list[str] | dict[Callable, list[str]] | None = None, + *, + _receive_inflate_expr: list[str] | None = None, +) -> torch.jit.ScriptModule: + """Create and return a copy of the specified model with inputs attached. + + The original model is not mutated or changed in any way. + + Models with bundled inputs can be invoked in a uniform manner by + benchmarking and code coverage tools. + + If inputs is passed in as a list then the inputs will be bundled for 'forward'. + If inputs is instead passed in as a map then all the methods specified in the map + will have their corresponding inputs bundled. Info should match watchever type is + chosen for the inputs. + + The returned model will support the following methods: + + `get_all_bundled_inputs_for_() -> List[Tuple[Any, ...]]` + Returns a list of tuples suitable for passing to the model like + `for inp in model.get_all_bundled_inputs_for_foo(): model.foo(*inp)` + + `get_bundled_inputs_functions_and_info() -> Dict[str, Dict[str: List[str]]]` + Returns a dictionary mapping function names to a metadata dictionary. + This nested dictionary maps preset strings like: + 'get_inputs_function_name' -> the name of a function attribute in this model that can be + run to get back a list of inputs corresponding to that function. + 'info' -> the user provided extra information about the bundled inputs + + If forward has bundled inputs then these following functions will also be defined on the returned module: + + `get_all_bundled_inputs() -> List[Tuple[Any, ...]]` + Returns a list of tuples suitable for passing to the model like + `for inp in model.get_all_bundled_inputs(): model(*inp)` + + `get_num_bundled_inputs() -> int` + Equivalent to `len(model.get_all_bundled_inputs())`, + but slightly easier to call from C++. + + Inputs can be specified in one of two ways: + + - The model can define `_generate_bundled_inputs_for_`. + If the user chooses this method inputs[] should map to None + + - The `inputs` argument to this function can be a dictionary mapping functions to a + list of inputs, of the same form that will be returned by get_all_bundled_inputs_for_. + Alternatively if only bundling inputs for forward the map can be omitted and a singular list of inputs + can be provided instead. + + The type of the inputs is List[Tuple[Any, ...]]. The outer list corresponds with a + list of inputs, the inner tuple is the list of args that together make up one input. + For inputs of functions that take one arg, this will be a tuple of length one. The Any, ... + is the actual data that makes up the args, e.g. a tensor. + + Info is an optional parameter that maps functions to a list of strings providing extra information about that + function's bundled inputs. Alternatively if only bundling inputs for forward the map can be omitted and + a singular list of information can be provided instead. This could be descriptions, expected outputs, etc. + - Ex: info={model.forward : ['man eating icecream', 'an airplane', 'a dog']} + + This function will attempt to optimize arguments so that (e.g.) + arguments like `torch.zeros(1000)` will be represented compactly. + Only top-level arguments will be optimized. + Tensors in lists or tuples will not. + """ + if not isinstance(model, torch.jit.ScriptModule): + raise Exception("Only ScriptModule is supported.") # noqa: TRY002 + + ignored_methods, ignored_attrs = _get_bundled_inputs_attributes_and_methods(model) + clone = torch._C._hack_do_not_use_clone_module_with_class( # type: ignore[attr-defined] + model._c, + ignored_methods, + ignored_attrs, + ) + + # The above cloning function returns a torch._C.scriptmodule and we need a torch.jit.scriptmodule. + # Fortunately there is a function in _recursive that does exactly that conversion. + cloned_module = wrap_cpp_module(clone) + if isinstance(inputs, dict): + if not isinstance(info, dict) and info is not None: + raise AssertionError("If inputs is a dict, info must be a dict or None") + augment_many_model_functions_with_bundled_inputs(cloned_module, inputs, _receive_inflate_expr, info) + else: + if not isinstance(info, list) and info is not None: + raise AssertionError("If inputs is a list, info must be a list or None") + augment_model_with_bundled_inputs(cloned_module, inputs, _receive_inflate_expr, info) + return cloned_module + +def augment_model_with_bundled_inputs( + model: torch.jit.ScriptModule, + inputs: Sequence[tuple[Any, ...]] | None = None, + _receive_inflate_expr: list[str] | None = None, # For debugging. + info: list[str] | None = None, # Optional argument to provide info about forward or its inputs + skip_size_check=False, +) -> None: + """Add bundled sample inputs to a model for the forward function. + + Models with bundled inputs can be invoked in a uniform manner by + benchmarking and code coverage tools. + + Augmented models will support the following methods: + + `get_all_bundled_inputs() -> List[Tuple[Any, ...]]` + Returns a list of tuples suitable for passing to the model like + `for inp in model.get_all_bundled_inputs(): model(*inp)` + + `get_num_bundled_inputs() -> int` + Equivalent to `len(model.get_all_bundled_inputs())`, + but slightly easier to call from C++. + + `get_bundled_inputs_functions_and_info() -> Dict[str, Dict[str: List[str]]]` + Returns a dictionary mapping function names to a metadata dictionary. + This nested dictionary maps preset strings like: + 'get_inputs_function_name' -> the name of a function attribute in this model that can be + run to get back a list of inputs corresponding to that function. + 'info' -> the user provided extra information about the bundled inputs + + Inputs can be specified in one of two ways: + + - The model can define `_generate_bundled_inputs_for_forward`. + If the user chooses this method inputs should be None + + - `inputs` is a list of inputs of form List[Tuple[Any, ...]]. A list of tuples where the elements + of each tuple are the args that make up one input. + """ + if not isinstance(model, torch.jit.ScriptModule): + raise Exception("Only ScriptModule is supported.") # noqa: TRY002 + + forward: Callable = model.forward + + # Sometimes forward won't have a name attached so just in case + if not hasattr(forward, "__name__"): + forward.__name__ = 'forward' + augment_many_model_functions_with_bundled_inputs( + model, + inputs={forward : inputs}, + _receive_inflate_expr=_receive_inflate_expr, + info={forward : info} if info else None, + skip_size_check=skip_size_check, + ) + + +def augment_many_model_functions_with_bundled_inputs( + model: torch.jit.ScriptModule, + inputs: dict[Callable, Sequence[tuple[Any, ...]] | None], + _receive_inflate_expr: list[str] | None = None, # For debugging. + info: dict[Callable, list[str]] | None = None, # Optional argument to provide info about the function or its inputs + skip_size_check=False, +) -> None: + """Add bundled sample inputs to a model for an arbitrary list of public functions. + + Models with bundled inputs can be invoked in a uniform manner by + benchmarking and code coverage tools. + + Augmented models will support the following methods: + + `get_all_bundled_inputs_for_() -> List[Tuple[Any, ...]]` + Returns a list of tuples suitable for passing to the model like + `for inp in model.get_all_bundled_inputs_for_foo(): model.foo(*inp)` + + `get_bundled_inputs_functions_and_info() -> Dict[str, Dict[str: List[str]]]` + Returns a dictionary mapping function names to a metadata dictionary. + This nested dictionary maps preset strings like: + 'get_inputs_function_name' -> the name of a function attribute in this model that can be + run to get back a list of inputs corresponding to that function. + 'info' -> the user provided extra information about the bundled inputs + + If forward has bundled inputs then these following functions are also defined: + + `get_all_bundled_inputs() -> List[Tuple[Any, ...]]` + Returns a list of tuples suitable for passing to the model like + `for inp in model.get_all_bundled_inputs(): model(*inp)` + + `get_num_bundled_inputs() -> int` + Equivalent to `len(model.get_all_bundled_inputs())`, + but slightly easier to call from C++. + + Inputs can be specified in one of two ways: + + - The model can define `_generate_bundled_inputs_for_`. + If the user chooses this method inputs[] should map to None + + - The `inputs` argument to this function can be a dictionary mapping functions to a + list of inputs, of the same form that will be returned by get_all_bundled_inputs_for_. + The type of the inputs is List[Tuple[Any, ...]]. The outer list corresponds with a + list of inputs, the inner tuple is the list of args that together make up one input. + For inputs of functions that take one arg, this will be a tuple of length one. The Any, ... + is the actual data that makes up the args, e.g. a tensor. + + Info is an optional parameter that maps functions to a list of strings providing extra information about that + function's bundled inputs. This could be descriptions, expected outputs, etc. + - Ex: info={model.forward : ['man eating icecream', 'an airplane', 'a dog']} + + This function will attempt to optimize arguments so that (e.g.) + arguments like `torch.zeros(1000)` will be represented compactly. + Only top-level arguments will be optimized. + Tensors in lists or tuples will not. + """ + if not isinstance(model, torch.jit.ScriptModule): + raise Exception("Only ScriptModule is supported.") # noqa: TRY002 + + if not inputs: + raise Exception("Please provide inputs for at least 1 function") # noqa: TRY002 + + if hasattr(model, "get_all_bundled_inputs") or hasattr(model, "get_bundled_inputs_functions_and_info"): + raise Exception( # noqa: TRY002 + "Models can only be augmented with bundled inputs once. " + "This Model seems to have already been augmented with " + "bundled inputs. Please start afresh with one that " + "doesn't have bundled inputs.", + ) + + get_bundled_inputs_functions_and_info_template = "" + + for function, input_list in inputs.items(): + if hasattr(function, "__name__"): + function_name = function.__name__ + else: + if hasattr(function, "name"): + function_name = function.name # type: ignore[attr-defined] + else: + raise Exception( # noqa: TRY002 + 'At least one of your functions has no attribute name please ensure all have one. m.foo.name = "foo"') + + + if input_list is not None and not isinstance(input_list, Sequence): + raise TypeError(f"Error inputs for function {function_name} is not a Sequence") + + function_arg_types = [arg.type for arg in function.schema.arguments[1:]] # type: ignore[attr-defined] + deflated_inputs_type: ListType = ListType(TupleType(function_arg_types)) + model._c._register_attribute(f"_bundled_inputs_deflated_{function_name}", deflated_inputs_type, []) + + if hasattr(model, "_generate_bundled_inputs_for_" + function_name): + if input_list is not None: + raise Exception( # noqa: TRY002 + f"inputs[{function_name}] is not None, but _generate_bundled_inputs_for_{function_name} is already defined" + ) + # Model author already defined _generate_bundled_inputs_for_. + elif input_list is None or len(input_list) == 0: + raise Exception( # noqa: TRY002 + f"inputs for {function_name} must be specified if " + f"_generate_bundled_inputs_for_{function_name} is not already defined" + ) + else: + # Iterate over the inputs and args in each input. + # Accumulate `deflated_inputs` as (possibly) compressed values + # and `parts` to be joined into the expression that unpacks them. + deflated_inputs = [] + parts = [] + for inp_idx, args in enumerate(input_list): + if not isinstance(args, tuple) and not isinstance(args, list): # type: ignore[arg-type] + raise TypeError( + f"Error bundled input for function {function_name} idx: {inp_idx} is not a Tuple or a List" + ) + deflated_args = [] + parts.append("(") + for arg_idx, arg in enumerate(args): + inflate_helper_fn_name = _get_inflate_helper_fn_name(arg_idx, inp_idx, function_name) + deflated, inflater, helper_definition = _inflate_expr( + arg, + f"deflated[{inp_idx}][{arg_idx}]", + inflate_helper_fn_name, + skip_size_check=skip_size_check, + ) + deflated_args.append(deflated) + parts.append(f" {inflater},") + if helper_definition: + model.define(textwrap.dedent(helper_definition)) + deflated_inputs.append(tuple(deflated_args)) + parts.append("),") + parts.append("") + expr = "\n".join(parts) + + # Back-channel return this expr for debugging. + if _receive_inflate_expr is not None: + _receive_inflate_expr.append(expr) + setattr(model, f"_bundled_inputs_deflated_{function_name}", deflated_inputs) + definition = textwrap.dedent(""" + def _generate_bundled_inputs_for_{name}(self): + deflated = self._bundled_inputs_deflated_{name} + return [ + {expr} + ] + """).format(expr=expr, name=function_name) + model.define(definition) + + # Define get_all_bundled_inputs_for_ that caches the generated inputs. + model.define(textwrap.dedent(""" + def get_all_bundled_inputs_for_{name}(self): + all_inputs = self._generate_bundled_inputs_for_{name}() + assert all_inputs is not None + return all_inputs + """).format(name=function_name)) + + # Add to the high level helper methods + inputs_info = repr(info[function]) if info and function in info else '[]' + get_bundled_inputs_functions_and_info_template += f""" + temp_dict : Dict[str,List[str]] = {{}} + info: List[str] = {inputs_info} + + temp_dict['info'] = info + temp_dict['get_inputs_function_name'] = ['get_all_bundled_inputs_for_{function_name}'] + all_inputs['{function_name}'] = temp_dict + """ + + # To ensure backwards compatibility and a streamlined api for forward these wrappers are provided + if function_name == 'forward': + model.define(textwrap.dedent(""" + def get_all_bundled_inputs(self): + return self.get_all_bundled_inputs_for_forward() + """)) + model.define(textwrap.dedent(""" + def get_num_bundled_inputs(self): + return len(self.get_all_bundled_inputs_for_forward()) + """)) + + # Define some high level helper methods that act on all bundled inputs + model.define(textwrap.dedent(f""" + def get_bundled_inputs_functions_and_info(self): + all_inputs : Dict[str, Dict[str,List[str]]] = {{}} + {get_bundled_inputs_functions_and_info_template} + return all_inputs + """)) + +def _inflate_expr( + arg: T, ref: str, inflate_helper_fn_name: str, skip_size_check: bool = False +) -> tuple[T | torch.Tensor, str, str | None]: + # Allow custom inflation expressions any object. + # For example, calling custom image-decoding ops. + # Or just use "{}" as the format string to ignore size limits. + if isinstance(arg, InflatableArg): + if arg.fmt_fn: + if arg.fmt not in ["{}", ""]: + raise Exception( # noqa: TRY002 + f"Bundled input argument at position '{ref}' has " + f"both arg.fmt_fn => \n{arg.fmt_fn} " + f"\n and arg.fmt => {arg.fmt}. " + "Please choose `arg.fmt` if the deflater is straightforward or " + "`arg.fmt_fn` if you need a function." + ) + + helper_definition = arg.fmt_fn.format(inflate_helper_fn_name) + expr = f"self.{inflate_helper_fn_name}({ref})" + + return arg.value, expr, helper_definition + else: + return arg.value, arg.fmt.format(ref), None + + if isinstance(arg, torch.Tensor): + # Small-storage tensors can just be saved directly. + if arg._typed_storage().size() <= MAX_RAW_TENSOR_SIZE or skip_size_check: + return arg, ref, None + # Small contiguous tensors can be cloned to have small storage. + # TODO: Should we do this even for non-contiguous tensors? + if arg.is_contiguous() and arg.numel() <= MAX_RAW_TENSOR_SIZE: + return arg.clone(), ref, None + # Example inputs commonly come from torch.zeros, torch.ones, or torch.full. + # These can be represented compactly. + for fmt in [torch.contiguous_format, torch.channels_last]: + if arg.is_contiguous(memory_format=fmt) and (arg == arg.flatten()[0]).all().item(): + return (arg.flatten()[0].clone().expand(*arg.size()), + f"{ref}.contiguous(memory_format={fmt})", None) + # Prevent big tensors from being bundled by default. + # TODO: Provide more useful diagnostics. + raise Exception( # noqa: TRY002 + f"Bundled input argument at position '{ref}' is " + f"a tensor with storage size {arg._typed_storage().size()}. " + f"You probably don't want to bundle this as an input. " + ) + else: + return arg, ref, None + +def _get_bundled_inputs_attributes_and_methods(script_module: torch.jit.ScriptModule) -> tuple[list[str], list[str]]: + methods: list[str] = [] + attributes: list[str] = [] + + # Has bundled inputs for forward + if hasattr(script_module, 'get_all_bundled_inputs'): + methods.append('get_all_bundled_inputs') + methods.append('get_num_bundled_inputs') + methods.append('run_on_bundled_input') + + if hasattr(script_module, 'get_bundled_inputs_functions_and_info'): + methods.append('get_bundled_inputs_functions_and_info') + all_info = script_module.get_bundled_inputs_functions_and_info() + for function_name in all_info: + methods.append("get_all_bundled_inputs_for_" + function_name) + methods.append("_generate_bundled_inputs_for_" + function_name) + attributes.append("_bundled_inputs_deflated_" + function_name) + + bundled_inputs_fn = getattr( + script_module, + f"get_all_bundled_inputs_for_{function_name}" + ) + num_bundled_inputs: int = len(bundled_inputs_fn()) + + # Check inflate helper functions for each function, argument and bundled input + func = getattr(script_module, function_name) + for arg_idx in range(len(func.schema.arguments) - 1): + for input_idx in range(num_bundled_inputs): + helper_fn_name = _get_inflate_helper_fn_name( + arg_idx=arg_idx, + input_idx=input_idx, + function_name=function_name + ) + # if the arg has an InflatableArg with fmt_fn, add the helper function name + if hasattr(script_module, helper_fn_name): + methods.append(helper_fn_name) + + return (methods, attributes) + + +def _get_inflate_helper_fn_name( + arg_idx: int, + input_idx: int, + function_name: str, +) -> str: + return f"_inflate_helper_for_{function_name}_input_{input_idx}_arg_{arg_idx}" + + + +def bundle_randn(*size, dtype=None): + """Generate a tensor that will be inflated with torch.randn.""" + stub = torch.zeros(1, dtype=dtype).expand(*size) + return InflatableArg(value=stub, fmt="torch.randn_like({})") + + +def bundle_large_tensor(t): + """Wrap a tensor to allow bundling regardless of size.""" + return InflatableArg(value=t, fmt="{}") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/checkpoint.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..d5f2bb039262b08e156dec471d6a3664b1852f2a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/checkpoint.py @@ -0,0 +1,1711 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import contextlib +import itertools +import platform +import uuid +import warnings +import weakref +from collections import defaultdict +from typing import * # noqa: F403 +from typing_extensions import Self +import enum +from weakref import ReferenceType + +import torch +import torch.fx.traceback as fx_traceback +from torch.utils._pytree import tree_map +from torch.testing._internal.logging_tensor import capture_logs, LoggingTensorMode +from torch.utils._python_dispatch import TorchDispatchMode +from torch._C._autograd import _make_saved_tensor, SavedTensor +from typing import NoReturn + +__all__ = [ + "checkpoint", + "checkpoint_sequential", + "CheckpointError", + "CheckpointFunction", + "check_backward_validity", + "detach_variable", + "get_device_states", + "set_device_states", + "noop_context_fn", + "set_checkpoint_early_stop", + "DefaultDeviceType", + "set_checkpoint_debug_enabled", + "CheckpointPolicy", + "SelectiveCheckpointContext", + "create_selective_checkpoint_contexts", + "SAC_IGNORED_OPS", + "GraphExecGroup", +] + +_DEFAULT_DETERMINISM_MODE = "default" + +_checkpoint_debug_enabled: bool | None = None + + +@contextlib.contextmanager +def set_checkpoint_debug_enabled(enabled: bool | None): + """ + Context manager that sets whether checkpoint should print additional debug + information when running. See the ``debug`` flag for + :func:`~torch.utils.checkpoint.checkpoint` for more information. Note that + when set, this context manager overrides the value of ``debug`` passed to + checkpoint. To defer to the local setting, pass ``None`` to this context. + + Args: + enabled (bool): Whether checkpoint should print debug information. + Default is 'None'. + """ + global _checkpoint_debug_enabled + try: + prev = _checkpoint_debug_enabled + _checkpoint_debug_enabled = enabled + yield + finally: + _checkpoint_debug_enabled = prev + + +def detach_variable(inputs: Tuple[Any, ...]) -> Tuple[torch.Tensor, ...]: + if isinstance(inputs, tuple): + out = [] + for inp in inputs: + if not isinstance(inp, torch.Tensor): + out.append(inp) + continue + + x = inp.detach() + x.requires_grad = inp.requires_grad + out.append(x) + return tuple(out) + else: + raise RuntimeError( + f"Only tuple of tensors is supported. Got Unsupported input type: {type(inputs).__name__}" + ) + + +def check_backward_validity(inputs: Iterable[Any]) -> None: + if not any(inp.requires_grad for inp in inputs if isinstance(inp, torch.Tensor)): + warnings.warn( + "None of the inputs have requires_grad=True. Gradients will be None", stacklevel=2 + ) + + +def _get_device_module(device="cuda"): + if device == "meta": + return torch.device("meta") + device_module = getattr(torch, device) + return device_module + + +class DefaultDeviceType: + r""" + A class that manages the default device type for checkpointing. + + If no non-CPU tensors are present, the default device type will + be used. The default value is 'cuda'. The device type is used in + the checkpointing process when determining which device states + to save and restore for recomputation. + """ + + _default_device_type: str | None = None + + @staticmethod + def set_device_type(device: str = "cuda") -> None: + """ + Set the default device type for checkpointing. + + Args: + device (str): The device type to be set as default. Default is 'cuda'. + """ + DefaultDeviceType._default_device_type = device + + @staticmethod + def get_device_type() -> str: + """ + Get the current default device type for checkpointing. + + Returns: + str: The current default device type. + """ + if not DefaultDeviceType._default_device_type: + DefaultDeviceType._default_device_type = acc.type if (acc := torch.accelerator.current_accelerator(True)) else "cpu" + + return DefaultDeviceType._default_device_type + + +def _infer_device_type(*args): + device_types = [] + + def add_device_types(arg): + nonlocal device_types + if isinstance(arg, torch.Tensor) and arg.device.type != "cpu": + device_types.append(arg.device.type) + return arg + tree_map(add_device_types, args) + + device_types_set = set(device_types) + if len(device_types_set) > 1: + warnings.warn( + "Tensor arguments, excluding CPU tensors, are detected on at least two types of devices. " + "Device state will only be saved for devices of a single device type, and the remaining " + "devices will be ignored. Consequently, if any checkpointed functions involve randomness, " + "this may result in incorrect gradients. (Note that if CUDA devices are among the devices " + "detected, it will be prioritized; otherwise, the first device encountered will be selected.)" + f"\nDevice types: {sorted(device_types_set)} first device type: {device_types[0]}", stacklevel=2 + ) + if len(device_types) == 0: + return DefaultDeviceType.get_device_type() + elif "cuda" in device_types_set: + return "cuda" + else: + return device_types[0] + + +# We can't know if the run_fn will internally move some args to different devices, +# which would require logic to preserve rng states for those devices as well. +# We could paranoically stash and restore ALL the rng states for all visible devices, +# but that seems very wasteful for most cases. Compromise: Stash the RNG state for +# the device of all Tensor args. +# +# To consider: maybe get_device_states and set_device_states should reside in torch/random.py? +def get_device_states(*args) -> Tuple[List[int], List[torch.Tensor]]: + # This will not error out if "arg" is a CPU tensor or a non-tensor type because + # the conditionals short-circuit. + fwd_device_ids = [] + + def add_device_ids(arg): + nonlocal fwd_device_ids + if isinstance(arg, torch.Tensor) and arg.device.type not in {"cpu", "meta"}: + fwd_device_ids.append(arg.get_device()) + return arg + tree_map(add_device_ids, args) + + fwd_device_states = [] + device_module = _get_device_module(_infer_device_type(*args)) + for device_id in fwd_device_ids: + with device_module.device(device_id): + fwd_device_states.append(device_module.get_rng_state()) + + return fwd_device_ids, fwd_device_states + + +def set_device_states(devices, states, *, device_type=None) -> None: + """Sets random number generator states for the specified devices. + + Args: + devices: Device ids to set states for. + states: States to set. + device_type: ``device_type`` of the devices to set states for. Default + is the device returned by a call to ``DefaultDeviceType.get_device_type()``, + which is ``cuda`` if not changed by calling ``DefaultDeviceType::set_device_type()``. + """ + if device_type is None: + device_type = DefaultDeviceType.get_device_type() + if device_type == "meta": + return + device_module = _get_device_module(device_type) + for device, state in zip(devices, states, strict=False): + with device_module.device(device): + device_module.set_rng_state(state) + + +def _get_autocast_kwargs(device_type="cuda"): + if torch.amp.is_autocast_available(device_type): + device_autocast_kwargs = { + "enabled": torch.is_autocast_enabled(device_type), + "dtype": torch.get_autocast_dtype(device_type), + "cache_enabled": torch.is_autocast_cache_enabled(), + } + else: + device_autocast_kwargs = None + + cpu_autocast_kwargs = { + "enabled": torch.is_autocast_enabled('cpu'), + "dtype": torch.get_autocast_dtype('cpu'), + "cache_enabled": torch.is_autocast_cache_enabled(), + } + + return device_autocast_kwargs, cpu_autocast_kwargs + + +class CheckpointFunction(torch.autograd.Function): + @staticmethod + # pyrefly: ignore [bad-override] + def forward(ctx, run_function, preserve_rng_state, *args): + check_backward_validity(args) + ctx.run_function = run_function + ctx.preserve_rng_state = preserve_rng_state + # Accommodates the (remote) possibility that autocast is enabled for cpu AND gpu. + ctx.device_type = _infer_device_type(*args) + ctx.device_autocast_kwargs, ctx.cpu_autocast_kwargs = _get_autocast_kwargs( + ctx.device_type + ) + if preserve_rng_state: + ctx.fwd_cpu_state = torch.get_rng_state() + # Don't eagerly initialize the cuda context by accident. + # (If the user intends that the context is initialized later, within their + # run_function, we SHOULD actually stash the cuda state here. Unfortunately, + # we have no way to anticipate this will happen before we run the function.) + ctx.had_device_in_fwd = False + device_module = _get_device_module(ctx.device_type) + if getattr(device_module, "_initialized", False): + ctx.had_device_in_fwd = True + ctx.fwd_devices, ctx.fwd_device_states = get_device_states(*args) + + # Save non-tensor inputs in ctx, keep a placeholder None for tensors + # to be filled out during the backward. + ctx.inputs = [] + ctx.tensor_indices = [] + tensor_inputs = [] + for i, arg in enumerate(args): + if torch.is_tensor(arg): + tensor_inputs.append(arg) + ctx.tensor_indices.append(i) + ctx.inputs.append(None) + else: + ctx.inputs.append(arg) + + ctx.save_for_backward(*tensor_inputs) + + with torch.no_grad(): + outputs = run_function(*args) + return outputs + + @staticmethod + def backward(ctx, *args): + if not torch.autograd._is_checkpoint_valid(): + raise RuntimeError( + "When use_reentrant=True, torch.utils.checkpoint is incompatible" + " with .grad() or passing an `inputs` parameter to .backward()." + " To resolve this error, you can either set use_reentrant=False," + " or call .backward() without passing the `inputs` argument." + ) + # Copy the list to avoid modifying original list. + inputs = list(ctx.inputs) + tensor_indices = ctx.tensor_indices + tensors = ctx.saved_tensors + + # Fill in inputs with appropriate saved tensors. + for i, idx in enumerate(tensor_indices): + inputs[idx] = tensors[i] + + # Stash the surrounding rng state, and mimic the state that was + # present at this time during forward. Restore the surrounding state + # when we're done. + rng_devices = [] + if ctx.preserve_rng_state and ctx.had_device_in_fwd: + rng_devices = ctx.fwd_devices + with torch.random.fork_rng( + devices=rng_devices, enabled=ctx.preserve_rng_state, device_type=ctx.device_type + ): + if ctx.preserve_rng_state: + torch.set_rng_state(ctx.fwd_cpu_state) + if ctx.had_device_in_fwd: + set_device_states(ctx.fwd_devices, ctx.fwd_device_states, device_type=ctx.device_type) + detached_inputs = detach_variable(tuple(inputs)) + + device_autocast_ctx = torch.amp.autocast( + device_type=ctx.device_type, **ctx.device_autocast_kwargs + ) if torch.amp.is_autocast_available(ctx.device_type) else contextlib.nullcontext() + with torch.enable_grad(), device_autocast_ctx, torch.amp.autocast("cpu", **ctx.cpu_autocast_kwargs): # type: ignore[attr-defined] + outputs = ctx.run_function(*detached_inputs) + + if isinstance(outputs, torch.Tensor): + outputs = (outputs,) + + # run backward() with only tensor that requires grad + outputs_with_grad = [] + args_with_grad = [] + for i in range(len(outputs)): + if torch.is_tensor(outputs[i]) and outputs[i].requires_grad: + outputs_with_grad.append(outputs[i]) + args_with_grad.append(args[i]) + if len(outputs_with_grad) == 0: + raise RuntimeError( + "none of output has requires_grad=True," + " this checkpoint() is not necessary" + ) + torch.autograd.backward(outputs_with_grad, args_with_grad) + grads = tuple( + inp.grad if isinstance(inp, torch.Tensor) else None + for inp in detached_inputs + ) + + return (None, None) + grads + + +def noop_context_fn(): + return contextlib.nullcontext(), contextlib.nullcontext() + +# Note: [torch.compile and checkpoint] +# TorchDynamo does not step inside utils.checkpoint function. The flow +# looks likes this +# 1) TorchDynamo tries to wrap utils.checkpoint in a HigherOrderOp by +# speculatively checking if the forward function is safe to trace. +# 2) If yes, then Dynamo-generated Fx graph has the wrapped higher +# order op. As a result, TorchDynamo does not look inside utils.checkpoint. +# 3) If not, then TorchDynamo falls back to eager by performing a graph +# break. And here, the following disable wrapper ensures that +# TorchDynamo does not trigger again on the frames created by +# utils.checkpoint innards. +@torch._disable_dynamo +def checkpoint( + function, + *args, + use_reentrant: bool | None = None, + context_fn: Callable[[], Tuple[ContextManager, ContextManager]] = noop_context_fn, + determinism_check: str = _DEFAULT_DETERMINISM_MODE, + debug: bool = False, + early_stop: bool = True, + **kwargs +): + r"""Checkpoint a model or part of the model. + + Activation checkpointing is a technique that trades compute for memory. + By default, tensors computed during the forward pass are kept alive until + they are used in gradient computations in the backward pass. To reduce this + memory usage, tensors produced in the passed :attr:`function` are not kept + alive until the backward pass. Instead, any passed tensors in :attr:`args` + are kept alive, and the unsaved tensors are recomputed by re-invoking + :attr:`function` in the backward pass as needed for gradient computation. + Activation checkpointing can be applied to any part of a model -- this is + sometimes described as "checkpointing" that part of the model. + + There are currently two checkpointing implementations available, determined + by the :attr:`use_reentrant` parameter. It is recommended that you use + ``use_reentrant=False``. Please refer the note below for a discussion of + their differences. + + .. warning:: + + If the :attr:`function` invocation during the backward pass differs + from the forward pass, e.g., due to a global variable, the checkpointed + version may not be equivalent, potentially causing an + error being raised or leading to silently incorrect gradients. + + .. warning:: + + The ``use_reentrant`` parameter should be passed explicitly. In version + 2.9 we will raise an exception if ``use_reentrant`` is not passed. + If you are using the ``use_reentrant=True`` variant, please refer to the + note below for important considerations and potential limitations. + + .. note:: + + The reentrant variant of checkpoint (``use_reentrant=True``) and + the non-reentrant variant of checkpoint (``use_reentrant=False``) + differ in the following ways: + + * Non-reentrant checkpoint stops recomputation as soon as all needed + intermediate activations have been recomputed. This feature is enabled + by default, but can be disabled with :func:`set_checkpoint_early_stop`. + Reentrant checkpoint always recomputes :attr:`function` in its + entirety during the backward pass. + + * The reentrant variant does not record the autograd graph during the + forward pass, as it runs with the forward pass under + :func:`torch.no_grad`. The non-reentrant version does record the + autograd graph, allowing one to perform backward on the graph within + checkpointed regions. + + * The reentrant checkpoint only supports the + :func:`torch.autograd.backward` API for the backward pass without its + `inputs` argument, while the non-reentrant version supports all ways + of performing the backward pass. + + * At least one input and output must have ``requires_grad=True`` for the + reentrant variant. If this condition is unmet, the checkpointed part + of the model will not have gradients. The non-reentrant version does + not have this requirement. + + * The reentrant version does not consider tensors in nested structures + (e.g., custom objects, lists, dicts, etc) as participating in + autograd, while the non-reentrant version does. + + * The reentrant checkpoint does not support checkpointed regions with + detached tensors from the computational graph, whereas the + non-reentrant version does. For the reentrant variant, if the + checkpointed segment contains tensors detached using ``detach()`` or + with :func:`torch.no_grad`, the backward pass will raise an error. + This is because ``checkpoint`` makes all the outputs require gradients + and this causes issues when a tensor is defined to have no gradient in + the model. To avoid this, detach the tensors outside of the + ``checkpoint`` function. + + Args: + function: describes what to run in the forward pass of the model or + part of the model. It should also know how to handle the inputs + passed as the tuple. For example, in LSTM, if user passes + ``(activation, hidden)``, :attr:`function` should correctly use the + first input as ``activation`` and the second input as ``hidden`` + args: tuple containing inputs to the :attr:`function` + + Keyword args: + preserve_rng_state(bool, optional): Omit stashing and restoring + the RNG state during each checkpoint. Note that under torch.compile, + this flag doesn't take effect and we always preserve RNG state. + Default: ``True`` + use_reentrant(bool): + specify whether to use the activation checkpoint variant that + requires reentrant autograd. This parameter should be passed + explicitly. In version 2.9 we will raise an exception if + ``use_reentrant`` is not passed. If ``use_reentrant=False``, + ``checkpoint`` will use an implementation that does not require + reentrant autograd. This allows ``checkpoint`` to support additional + functionality, such as working as expected with + ``torch.autograd.grad`` and support for keyword arguments input into + the checkpointed function. + context_fn(Callable, optional): A callable returning a tuple of two + context managers. The function and its recomputation will be run + under the first and second context managers respectively. + This argument is only supported if ``use_reentrant=False``. + determinism_check(str, optional): A string specifying the determinism + check to perform. By default it is set to ``"default"`` which + compares the shapes, dtypes, and devices of the recomputed tensors + against those the saved tensors. To turn off this check, specify + ``"none"``. Currently these are the only two supported values. + Please open an issue if you would like to see more determinism + checks. This argument is only supported if ``use_reentrant=False``, + if ``use_reentrant=True``, the determinism check is always disabled. + debug(bool, optional): If ``True``, error messages will also include + a trace of the operators ran during the original forward computation + as well as the recomputation. This argument is only supported if + ``use_reentrant=False``. + early_stop(bool, optional): If ``True``, non-reentrant checkpoint stops + recomputation as soon as it has computed all needed Tensors. This + argument is ignored if ``use_reentrant=True``. Can be overridden + globally using :func:`set_checkpoint_early_stop` context manager. + Default: ``True``. + + Returns: + Output of running :attr:`function` on :attr:`*args` + """ + if use_reentrant is None: + warnings.warn( + "torch.utils.checkpoint: the use_reentrant parameter should be " + "passed explicitly. Starting in PyTorch 2.9, calling checkpoint " + "without use_reentrant will raise an exception. use_reentrant=False is " + "recommended, but if you need to preserve the current default " + "behavior, you can pass use_reentrant=True. Refer to docs for more " + "details on the differences between the two variants.", + stacklevel=2 + ) + use_reentrant = True + + # Hack to mix *args with **kwargs in a python 2.7-compliant way + preserve = kwargs.pop("preserve_rng_state", True) + if kwargs and use_reentrant: + raise ValueError( + "Unexpected keyword arguments: " + ",".join(arg for arg in kwargs) + ) + + if use_reentrant: + if context_fn is not noop_context_fn or debug is not False: + raise ValueError( + "Passing `context_fn` or `debug` is only supported when " + "use_reentrant=False." + ) + return CheckpointFunction.apply(function, preserve, *args) + else: + gen = _checkpoint_without_reentrant_generator( + function, preserve, context_fn, determinism_check, debug, early_stop, *args, **kwargs + ) + # Runs pre-forward logic + next(gen) + ret = function(*args, **kwargs) + # Runs post-forward logic + try: + next(gen) + except StopIteration: + return ret + + +def checkpoint_sequential(functions, segments, input, use_reentrant=None, **kwargs): + r"""Checkpoint a sequential model to save memory. + + Sequential models execute a list of modules/functions in order + (sequentially). Therefore, we can divide such a model in various segments + and checkpoint each segment. All segments except the last will not store + the intermediate activations. The inputs of each checkpointed segment will + be saved for re-running the segment in the backward pass. + + .. warning:: + The ``use_reentrant`` parameter should be passed explicitly. In version + 2.9 we will raise an exception if ``use_reentrant`` is not passed. + If you are using the ``use_reentrant=True` variant, please see + :func:`~torch.utils.checkpoint.checkpoint` for + the important considerations and limitations of this variant. It is + recommended that you use ``use_reentrant=False``. + + .. warning: + Since PyTorch 1.4, it allows only one Tensor as the input and + intermediate outputs, just like :class:`torch.nn.Sequential`. + + Args: + functions: A :class:`torch.nn.Sequential` or the list of modules or + functions (comprising the model) to run sequentially. + segments: Number of chunks to create in the model + input: A Tensor that is input to :attr:`functions` + preserve_rng_state(bool, optional): Omit stashing and restoring + the RNG state during each checkpoint. + Default: ``True`` + use_reentrant(bool): + specify whether to use the activation checkpoint variant that + requires reentrant autograd. This parameter should be passed + explicitly. In version 2.5 we will raise an exception if + ``use_reentrant`` is not passed. If ``use_reentrant=False``, + ``checkpoint`` will use an implementation that does not require + reentrant autograd. This allows ``checkpoint`` to support additional + functionality, such as working as expected with + ``torch.autograd.grad`` and support for keyword arguments input into + the checkpointed function. + + Returns: + Output of running :attr:`functions` sequentially on :attr:`*inputs` + + Example: + >>> # xdoctest: +SKIP("stub") + >>> model = nn.Sequential(...) + >>> input_var = checkpoint_sequential(model, chunks, input_var) + """ + if use_reentrant is None: + warnings.warn( + "torch.utils.checkpoint.checkpoint_sequential: the use_reentrant " + "parameter should be passed explicitly. " + "In version 2.9 we will raise an exception if use_reentrant " + "is not passed. use_reentrant=False is " + "recommended, but if you need to preserve the current default " + "behavior, you can pass use_reentrant=True. Refer to docs for more " + "details on the differences between the two variants.", stacklevel=2 + ) + use_reentrant = True + + # Hack for keyword-only parameter in a python 2.7-compliant way + preserve = kwargs.pop("preserve_rng_state", True) + if kwargs: + raise ValueError( + "Unexpected keyword arguments: " + ",".join(arg for arg in kwargs) + ) + + def run_function(start, end, functions): + def forward(input): + for j in range(start, end + 1): + input = functions[j](input) + return input + + return forward + + if isinstance(functions, torch.nn.Sequential): + functions = list(functions.children()) + + segment_size = len(functions) // segments + # the last chunk has to be non-volatile + end = -1 + for start in range(0, segment_size * (segments - 1), segment_size): + end = start + segment_size - 1 + input = checkpoint( + run_function(start, end, functions), + input, + use_reentrant=use_reentrant, + preserve_rng_state=preserve, + ) + return run_function(end + 1, len(functions) - 1, functions)(input) + + +def _internal_assert(cond) -> None: + if not cond: + raise AssertionError( + "Something went unexpectedly wrong in activation checkpoint. " + "Please report this bug by filing an issue to PyTorch." + ) + + +# NOTE [ Nestable Checkpoint ] +# +# The semantics of nested checkpoint can be defined by two basic rules. +# Following the two rules leads to an important implication that is central +# to motivating the design. +# +# Rule 1. Saved tensors are managed by inner-most checkpoint only and hidden +# from any outer layers of checkpoint. +# +# Rule 2. The inputs of inner checkpoints are treated as tensors saved to its +# parent checkpoint. +# +# Implication: To recompute any given saved tensor, we need to recompute all of +# the checkpoints wrapping it. +# +# Why is this implied? To unpack a saved tensor X during backward we need to +# recompute the inner-most checkpoint (#1), and in order to recompute that +# checkpoint I need to have its inputs, which are managed by that checkpoint's +# parent (#2), which thus also needs to be recomputed first. Continue this line +# of reasoning and we realize that in order to unpack X, all checkpoints that +# were active at the time X was saved need to be recomputed. (unless we have +# already done so in that backward for some other saved tensor). +# +# In practice, we use a noop autograd Function to save inputs as saved tensors. +# During unpack calling ctx.saved_tensor triggers the parent checkpoint to +# recompute. +# +# Rule 3. We should start recomputation as if there are no checkpoints currently +# active. Checkpoints encountered during recomputation are still +# respected. +# +# When we start recomputation, we push the saved variable hook meant for +# recomputation on the stack. See examples in Rule 6 for more context. +# +# * * * * +# +# Beyond the basic semantics specific to nested checkpoint, we impose several +# more constraints that may apply to checkpointing in general. +# +# Rule 4. Lifetime of recomputed tensors +# +# Recomputed tensors are considered specific to particular invocations +# of backward and are always cleared immediately as they are unpacked +# Particularly, we require this to happen even if retain_graph=True. +# +# [ Implementation details of Rule 4 ] +# +# If we were okay with recomputed tensors staying alive after backward is run +# with retain_graph=True, we would store recomputed variables as the values of a +# WeakKeyDictionary and pack strong references to the keys, so that as we +# backward, those packed keys would be cleared as long as retain_graph=False. +# Clearing the packed key clears the corresponding entry in the WKD. +# +# If we wish recomputed variables to be immediately cleared as we unpack them in +# the retain_graph=True case, we cannot rely on the packed keys to be cleared by +# backward automatically. Instead of packing the strong reference to the key +# directly, we pack a container object, which we manually clear as we unpack. +# +# An important detail is that if a second backward happens, the second +# recomputation needs to reset the container with a newly created key. +# +# Rule 5. Stop recomputation as soon as we've recomputed the saved tensors we +# know we need. +# +# [ Implementation details of Rule 5 ] +# +# During recomputation, raise an exception if the number of recomputed tensors +# matches the number of tensors that we expected to recompute. We wrap the +# recomputation call with a try-catch to catch this specific exception. See +# Rule #6 below for some examples. +# +# Rule 6. We support doing backward inside checkpoint context +# +# [ retain_graph is True] +# +# def fn(x): +# y = x.sin() +# z = y.cos() +# gx, = torch.autograd.grad(z, x, retains_grad=True) +# return gx, z +# +# out = checkpoint(fn)(inp) +# out.backward() +# +# Because z is saved by cos while checkpoint is enabled, it would not be +# actually saved, and so the .grad() call inside must trigger a recomputation. +# +# During recomputation the "inner pack hook" has two responsibilities: +# +# 1) As usual, populating the WeakKeyDictionary storing recomputed tensors +# 2) Pack the actual tensor (detached) so that one may perform backward on the +# recomputed graph. The tensors saved to this graph will live until the end +# of recomputation, or die earlier if someone performs backward with +# retain_graph=False. +# +# More generally performing backward on the recomputed graph occurs in the +# following cases: +# - If backward is performed inside forward, +# - During the original forward IF early-stop is disabled +# - During the original backward +# - If there are multiple .grad()/.backward() calls, we would perform backward +# on the recomputed graph even if early-stop is enabled (see the example below) +# +# [ retain_graph is False ] +# +# The example below shows what happens if during recomputation we find that some +# of the tensors we are trying to recompute have already been cleared. +# +# Spoiler: we don't do anything special, we just skip over them! +# +# def fn(x): +# y = x.sin() # (1) +# z = y.cos() # (2) +# gx, = torch.autograd.grad(z, x) # (3) +# return x.cos() * gx # (4) +# +# out = checkpoint(fn)(inp) +# out.backward() # (5) +# +# 1, 2. Don't save x and y since we are inside a checkpoint. +# 3. Trigger a recompute of fn since x and y weren't saved. +# And depending on whether early stop is enabled, either stop at (2) or +# continue running the function. +# Because we are running backward with retain_graph=False, we clear x and y's +# holders. +# 4. Don't save x since we are inside a checkpoint. +# 5. Calling backward triggers another recompute of fn. During recompute, we see +# that x and y have already been cleared in the original graph as indicated +# by holder=None. We skip over them. We still save x at (4) (since its holder +# is still alive.) + +_enable_checkpoint_early_stop: bool | None = None + + +@contextlib.contextmanager +def set_checkpoint_early_stop(enable: bool): + """Context manager that sets whether checkpoint should stop recomputation early. + + By default, non-reentrant checkpoint stops recomputation as soon as it + has computed all needed Tensors. This context manager can be used to disable + that feature if it is problematic for your specific application. + + This context manager only needs to be active when forward is run. It does + not need to be active during backward. + + Example:: + + >>> # xdoctest: +SKIP(failing) + >>> message = "saved tensors default hooks are disabled" + >>> with set_checkpoint_early_stop(False): + ... # Any checkpoint under this context manager will respect this + ... # context manager, even if its backward is performed outside. + ... out = checkpoint(fn, inputs) + ... + >>> out.backward() + """ + global _enable_checkpoint_early_stop + try: + prev = _enable_checkpoint_early_stop + _enable_checkpoint_early_stop = enable + yield + finally: + _enable_checkpoint_early_stop = prev + + +class _Handle: + pass + + +class _Holder: + def __init__(self) -> None: + self.handles: dict[int, _Handle | None] = {} + + +class _CheckpointFrame: + def __init__(self, recompute_fn, early_stop, unpack_error_cb, metadata_fn) -> None: + self.recompute_fn = recompute_fn + self.saved_args: List[Any] = [] + self.weak_holders: List[ReferenceType] = [] + # We store this as a weakkeydictionary so that in the case of a partial + # backward, the entries in the dict are cleared alongside the Holder + # which will be removed when the SavedVariable is cleared. + self.recomputed: DefaultDict[ + int, weakref.WeakKeyDictionary[_Handle, torch.Tensor] + ] = defaultdict(weakref.WeakKeyDictionary) + # We need both recomp_counter and recomputed since they can diverge + # https://github.com/pytorch/pytorch/pull/90105#discussion_r1135889885 + self.recomp_counter: DefaultDict[int, int] = defaultdict(int) + self.is_recomputed: DefaultDict[int, bool] = defaultdict(bool) + + # See Rule 5 + self.early_stop = early_stop + + # Debugging + self.metadata_fn = metadata_fn + self.unpack_error_cb = unpack_error_cb + self.x_metadatas = [] + self.forward_completed = False + self.ignore_saved_mismatch = False + + def save_inputs(self, *args): + self.saved_args = [ + _make_saved_tensor(arg, is_output=False) + if isinstance(arg, torch.Tensor) else arg + for arg in args + ] + + def get_inputs(self): + return [ + arg.unpack() if isinstance(arg, SavedTensor) else arg + for arg in self.saved_args + ] + + def check_recomputed_tensors_match(self, gid) -> None: + if self.ignore_saved_mismatch: + # TODO: we can probably make this check stricter by checking that + # the metadata of the first tensors still match. + return + # NOTE [ Error handling for checkpoint ] + # + # At a high level, we need to check that the tensors saved + # during original forward matches tensors saved during recompute + # This means handling 3 cases: + # + # 1. During recompute, more tensors were saved. + # + # Usually this is hidden due to the StopRecomputationError + # but if early stop is not enabled, or we would have errored + # anyway because there aren't enough weak_holders. But we + # do want to have a nice error. See the _recomputation_hook + # for details. + if not len(self.weak_holders) == self.recomp_counter[gid]: + # 2. During recompute, fewer tensors were saved + # + # We know that every time we save something do original forward + # we append to weak_holder, and every time we save a tensor + # during recompute we increment recompute_counter. + raise CheckpointError( + "torch.utils.checkpoint: A different number of tensors was saved " + "during the original forward and recomputation.\n" + f"Number of tensors saved during forward: {len(self.weak_holders)}\n" + f"Number of tensors saved during recomputation: {self.recomp_counter[gid]}.\n" + f"{_debug_tip_msg}" + ) + + # 3. During recompute, the same tensors were saved, but they + # have different metadata + nb_meta_different = [] + for idx, weak_holder in enumerate(self.weak_holders): + holder = weak_holder() + if holder is None: + continue + # We've seen all holders since we iterate over them in order + # For every holder that is still alive now, it must've been + # alive when we saw it during recompute, therefore, the + # gid must be set. + _internal_assert(gid in holder.handles) + # We know this is the first unpack, so it couldn't have been set + # to None yet. + _internal_assert(holder.handles[gid] is not None) + # We always set these together in the recomputation hook + _internal_assert(holder.handles[gid] in self.recomputed[gid]) + # see pack hook, x_metadata is 1:1 with weak_holders. + x_meta = self.x_metadatas[idx] + recomputed_x = self.recomputed[gid][holder.handles[gid]] + if x_meta != self.metadata_fn(recomputed_x): + nb_meta_different.append((idx, x_meta, self.metadata_fn(recomputed_x))) + + if len(nb_meta_different) > 0: + mismatched_tensors = "" + for idx, x_meta, recomputed_meta in nb_meta_different: + mismatched_tensors += ( + f"tensor at position {idx}:\n" + f"saved metadata: {x_meta}\n" + f"recomputed metadata: {recomputed_meta}\n" + ) + raise CheckpointError( + "torch.utils.checkpoint: Recomputed values for the following tensors " + "have different metadata than during the forward pass.\n" + f"{mismatched_tensors}.\n" + f"{_debug_tip_msg}" + ) + + +_debug_tip_msg = """ +Tip: To see a more detailed error message, either pass `debug=True` to +`torch.utils.checkpoint.checkpoint(...)` or wrap the code block +with `with torch.utils.checkpoint.set_checkpoint_debug_enabled(True):` to +enable checkpoint‑debug mode globally. + +If this error occurs under torch.compile with automatic_dynamic_shapes enabled, +it may be because the recomputation selected a different compiled graph than the +forward pass (e.g., a dynamic graph instead of the original static graph). +To fix this, either: + - Use torch._dynamo.mark_dynamic() to explicitly mark varying dimensions as + dynamic upfront, avoiding the static-to-dynamic transition. + - Call torch._C._dynamo.eval_frame._set_lru_cache(False) to disable LRU cache + reordering, which can change which graph is checked first between forward + and recompute. +See https://github.com/pytorch/pytorch/issues/166926 for more details and +workaround examples. +""" + + +_checkpoint_error_template = """ \ +An error happened while unpacking tensors; dumping logs of latest computation +because you passed `debug=True` to `torch.utils.checkpoint.checkpoint()`. +Scroll all the way down for guidance on how to navigate these logs. + ++~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~+ +| 1. Stack traces of the operators that ran in the original forward | ++------------------------------------------------------------------------------+ + +{forward_traces} ++~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~+ +| 2. Stack traces of the operators that ran during recomputation | ++------------------------------------------------------------------------------+ + +{recompute_traces} ++~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~+ +| 3. Log of operators in the original forward and recomputation | ++------------------------------------------------------------------------------+ +(Scroll up to correlate stack traces with each operation listed below. This + helps identify their source in the code.) + +IMPORTANT: Differences in "detach" calls between the original forward and the + recomputation are expected. They are introduced by the checkpointing + mechanism and can be ignored. + +Operations executed during the original forward: + +{forward_ops} + +Operations executed during recomputation: + +{recompute_ops} + ++------------------------------------------------------------------------------+ + ERROR: Detected non-determinism while running activation checkpointing + + You are seeing this error because you passed `debug=True` to checkpoint and + tensors to be saved during the original forward and differ between those saved + during recomputation. This can happen if different operators were ran in the + original forward and in the recomputation. + + To identify where the mismatch may be coming from, you can do the following: + + 1) Compare the operators ran during original forward and recomputation to + see where they differ. These operators are printed above in the order they + were executed. + + 2) Review the stack trace for each operator to locate its invocation source. + Each operator's stack trace is printed in their execution order. + + Note that the logs can be quite long. Here's how they are structured: + (Tip: you can Ctrl-f for these headers) + + 1. Stack traces of the operators that ran in the original forward + 2. Stack traces of the operators that ran during recomputation + 3. Log of operators in the original forward and recomputation + 4. Error message <--- You are here +-------------------------------------------------------------------------------- +""" + +class CheckpointError(RuntimeError): + pass + + +def _get_debug_context_and_cb() -> Tuple[Callable[[], Any], Callable[[CheckpointError], None]]: + # This function returns the context_fn and error_cb to be used by the + # checkpointing mechanism. error_cb is invoked when an error is detected + # during unpack. + + cpp_tb = platform.machine() in ('x86_64', 'aarch64', 'arm64') and platform.system() == 'Linux' + + class CaptureLogs: + def __init__(self) -> None: + self.logs = None + self.tbs = None + + def get_context_manager(self): + @contextlib.contextmanager + def logging_mode(): + with LoggingTensorMode(), \ + capture_logs(True, python_tb=True, script_tb=True, cpp_tb=cpp_tb) as logs_and_tb: + self.logs, self.tbs = logs_and_tb + yield logs_and_tb + return logging_mode() + + capture_logs_fwd = CaptureLogs() + capture_logs_recompute = CaptureLogs() + + def unpack_error_cb(e: CheckpointError) -> NoReturn: + def get_str_tb(label, capture_logs): + out = "" + total_len = len(capture_logs.logs) + for i, (log, tb) in enumerate(zip(capture_logs.logs, capture_logs.tbs, strict=False)): + out += f"{log} ({i + 1} of {total_len} in {label})\n\n" + found_torch_dispatch = False + for line in tb: + # Start printing stack trace only after __torch_dispatch__ is found + is_torch_dispatch = line['name'] == '__torch_dispatch__' + if not found_torch_dispatch and not is_torch_dispatch: + continue + elif is_torch_dispatch: + found_torch_dispatch = True + continue + out += f"{line['filename']}:{line['line']}:{line['name']}\n" + out += "\n\n" + return out + if capture_logs_fwd.logs is None: + raise AssertionError("capture_logs_fwd.logs is None") + if capture_logs_recompute.logs is None: + raise AssertionError("capture_logs_recompute.logs is None") + raise CheckpointError( + _checkpoint_error_template.format( + forward_traces=get_str_tb("original", capture_logs_fwd), + recompute_traces=get_str_tb("recompute", capture_logs_recompute), + forward_ops="\n".join(capture_logs_fwd.logs), + recompute_ops="\n".join(capture_logs_recompute.logs) + ) + ) from e + + def context_fn(): + return capture_logs_fwd.get_context_manager(), capture_logs_recompute.get_context_manager() + + return context_fn, unpack_error_cb + +def _default_meta_extractor(x: torch.Tensor) -> Dict[str, Any]: + # These properties are fast to check, easy to understand + return { + "shape": x.shape, + "dtype": x.dtype, + "device": x.device + } + +_allowed_determinism_checks_to_fns: Dict[str, Callable[[torch.Tensor], Any]] = { + _DEFAULT_DETERMINISM_MODE: _default_meta_extractor, + "none": lambda _: None, +} + +# See Rule 5 +class _StopRecomputationError(Exception): + pass + + +class _recomputation_hook(torch.autograd.graph.saved_tensors_hooks): + def __init__(self, target_frame_ref: ReferenceType, gid: GraphExecGroup | int) -> None: + # Dynamo guards on WeakKeyDictionary internals are unstable here + # (dict length/keys change every call), causing recompilation storms. + # with `.compile()` so we disable + @torch._dynamo.disable + def pack_hook(x): + x = x.detach() if x.requires_grad else x + target_frame = target_frame_ref() + if target_frame is None: + raise AssertionError("Internal error: target_frame reference is None") + recomp_idx = target_frame.recomp_counter[gid] + target_frame.recomp_counter[gid] += 1 + + if recomp_idx >= len(target_frame.weak_holders): + if target_frame.early_stop: + raise AssertionError("Unexpected state: target_frame.early_stop is set") + if not target_frame.forward_completed: + # We run into this case when early stop is not enabled and do + # grad within checkpoint. + # We need to set this flag, so we don't error out later when + # we check if the number of tensors saved during forward and + # recomputation match. + target_frame.ignore_saved_mismatch = True + return x + raise CheckpointError( + "torch.utils.checkpoint: trying to save more tensors during " + "recomputation than during the original forward pass.\n" + f"{_debug_tip_msg}" + ) + + holder = target_frame.weak_holders[recomp_idx]() + + # This holder may have been cleared because someone may have called + # backward within forward. If so, we don't need to save. + if holder is not None: + _internal_assert(holder.handles.get(gid, None) is None) + holder.handles[gid] = _Handle() + target_frame.recomputed[gid][holder.handles[gid]] = x + + if target_frame.early_stop and target_frame.recomp_counter[gid] == len( + target_frame.weak_holders + ): + raise _StopRecomputationError + # See Rule 6: [ retain_graph is True ] above + return x + + def unpack_hook(x): + # See Rule 6: [ retain_graph is True ] above for an example of when + # the graph created during recomputation could be backwarded. + return x + + super().__init__(pack_hook, unpack_hook) + + +# torch._disable_dynamo creates a reference cycle with decorated function +# This function is used to ensure that the decorated function does not have +# a closure, so that other objects aren't also kept alive. +# https://github.com/pytorch/pytorch/issues/154642 +# Note: does not work when fn is compiled +@torch._disable_dynamo +def _run_fn_with_dynamo_disabled(fn, *args, **kwargs): + return fn(*args, **kwargs) + + +class _checkpoint_hook(torch.autograd.graph.saved_tensors_hooks): + def __init__(self, frame) -> None: + def pack_hook(x): + # See Rule 4 above + holder = _Holder() + frame.weak_holders.append(weakref.ref(holder)) + # Save metadata to detect non-determinism + if frame.metadata_fn is not None: + with torch.no_grad(): + frame.x_metadatas.append(frame.metadata_fn(x)) + return holder + + def unpack_hook(holder): + # First check if we're inside a GraphExecGroup context + gid: GraphExecGroup | None | int = GraphExecGroup._get_current_group() + if gid is None: + # Fallback to using the current graph task id + gid = torch._C._current_graph_task_id() + if gid == -1: + # generate a temporary id if we trigger unpack outside of a backward call + gid = int(uuid.uuid4()) + + if not frame.is_recomputed[gid]: + args = frame.get_inputs() + + try: + with _recomputation_hook( + weakref.ref(frame), gid + ), torch.autograd.enable_grad(): + # See Note: [compiled autograd and checkpoint unpack hook] + _run_fn_with_dynamo_disabled(frame.recompute_fn, *args) + except _StopRecomputationError: + pass + frame.is_recomputed[gid] = True + frame.check_recomputed_tensors_match(gid) + + _internal_assert(gid in holder.handles) + + if holder.handles[gid] is None: + extra = "" + if torch._C._get_graph_exec_group() is not None: + extra = ( + "Performing two backward calls that overlap (i.e. require the same " + "saved activation in order to compute gradients) is not allowed while " + "under the torch.utils.checkpoint.GraphExecGroup context. " + ) + raise CheckpointError( + "torch.utils.checkpoint: Unpack is being triggered for a tensor that was already " + f"unpacked once. {extra}If you are calling ctx.saved_tensors in backward, make sure " + "to do so only once. Otherwise please open an issue with details on your use case." + ) + _internal_assert(holder.handles[gid] in frame.recomputed[gid]) + ret = frame.recomputed[gid][holder.handles[gid]] + holder.handles[gid] = None + return ret + + if frame.unpack_error_cb is not None: + def unpack_hook_with_error_cb(holder): + try: + return unpack_hook(holder) + except CheckpointError as e: + frame.unpack_error_cb(e) + super().__init__(pack_hook, unpack_hook_with_error_cb) + else: + super().__init__(pack_hook, unpack_hook) + + +def _is_compiling(func, args, kwargs): + # Check if we are under AOTAutograd tracing or export tracing + # Checking that a proxy mode is active should always do what we want + if torch.compiler._is_non_strict_tracing(): + return False + return torch._C._get_dispatch_mode(torch._C._TorchDispatchModeKey.PROXY) is not None + + +class _VersionWrapper: + # Check that cached tensors are not mutated. + def __init__(self, val) -> None: + self.val: torch.Tensor | Any = val + self.version: int | None = val._version if isinstance(val, torch.Tensor) else None + + def get_val(self, allow_cache_entry_mutation): + if self.version is not None and not allow_cache_entry_mutation: + if self.val._version != self.version: + # Can we give user a stack trace of where the mutation happened? + raise RuntimeError( + "Tensor cached during selective activation checkpoint has been mutated" + ) + return self.val + + +def _maybe_detach(x, any_ret_has_alias_info): + # We detach for two separate reasons: + # - For view ops, we need to ensure that when the tensor is returned from + # CachedDispatchMode, as_view sees that the AutogradMeta is nullptr + # - Avoid reference cycles + # For case 1, it is not enough to check whether x has differentiable dtype + # because non-differentiable dtype can have non-nullptr AutogradMeta, e.g. + # when the tensor is a view. + if isinstance(x, torch.Tensor) and (x.is_floating_point() or x.is_complex() or any_ret_has_alias_info): + with torch._C._SetExcludeDispatchKeyGuard(torch._C.DispatchKey.ADInplaceOrView, False): + # Ensure that view performed beneath autograd properly propagates + # version counter. TODO: Use reentrant_dispatch instead of + # manually manipulating dispatch keys. Using reentrant_dispatch + # would respect inference_mode, though that is not relevant for + # this case. + x = x.detach() + return x + + +class SelectiveCheckpointContext: + """ + Context passed to policy function during selective checkpointing. + + This class is used to pass relevant metadata to the policy function during + selective checkpointing. The metadata includes whether the current invocation + of the policy function is during recomputation or not. + + Example: + >>> # xdoctest: +SKIP(stub) + >>> + >>> def policy_fn(ctx, op, *args, **kwargs): + >>> print(ctx.is_recompute) + >>> + >>> context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) + >>> + >>> out = torch.utils.checkpoint.checkpoint( + >>> fn, x, y, + >>> use_reentrant=False, + >>> context_fn=context_fn, + >>> ) + """ + def __init__(self, *, is_recompute, op_output=None) -> None: + self.is_recompute = is_recompute + self.op_output = op_output + + +class CheckpointPolicy(enum.Enum): + """ + Enum for specifying the policy for checkpointing during backpropagation. + + The following policies are supported: + + - ``{MUST,PREFER}_SAVE``: The operation's output will be saved during the forward + pass and will not be recomputed during the backward pass + - ``{MUST,PREFER}_RECOMPUTE``: The operation's output will not be saved during the + forward pass and will be recomputed during the backward pass + - ``{MUST,PREFER}_CPU_OFFLOAD``: The operation's output will be saved during the + forward pass, offloaded to CPU, and reloaded to GPU during the backward pass + + Use ``MUST_*`` over ``PREFER_*`` to indicate that the policy should not be overridden + by other subsystems like `torch.compile`. + + .. note:: + A policy function that always returns ``PREFER_RECOMPUTE`` is + equivalent to vanilla checkpointing. + + A policy function that returns ``PREFER_SAVE`` every op is + NOT equivalent to not using checkpointing. Using such a policy would + save additional tensors not limited to ones that are actually needed for + gradient computation. + """ + MUST_SAVE = 0 + PREFER_SAVE = 1 + MUST_RECOMPUTE = 2 + PREFER_RECOMPUTE = 3 + MUST_CPU_OFFLOAD = 4 + PREFER_CPU_OFFLOAD = 5 + + +def _policy_from_bool(b): + # For backward compatibility + return CheckpointPolicy.MUST_SAVE if b else CheckpointPolicy.PREFER_RECOMPUTE + + +SAC_IGNORED_OPS = { + # AC inserts different number of detach during forward and recompute. + torch.ops.aten.detach.default, + # AC's determinism check invokes additional metadata ops during forward. + # With subclasses involved, these metadata ops become dispatchable, this + # can result in incorrectness if these ops are selected cached. + torch.ops.prim.device.default, +} | set(torch._subclasses.functional_tensor.FunctionalTensor.metadata_fns) # type: ignore[has-type] + + +class _CachingTorchDispatchMode(TorchDispatchMode): + @classmethod + def ignore_compile_internals(cls): + return True + + # Used together with _CachedTorchDispatchMode to implement SAC. + def __init__(self, policy_fn, storage, ac_graph_id=None) -> None: + self.policy_fn = policy_fn + self.storage = storage + self.ac_graph_id = ac_graph_id + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + kwargs = {} if kwargs is None else kwargs + is_compiling = _is_compiling(func, args, kwargs) + + if is_compiling: + fx_traceback.current_meta["ac_graph_id"] = self.ac_graph_id + fx_traceback.current_meta["recompute"] = CheckpointPolicy.PREFER_RECOMPUTE + + if func in SAC_IGNORED_OPS: + return func(*args, **kwargs) + + proxy_mode = None + graph_len_before = 0 + if is_compiling: + from torch.fx.experimental.proxy_tensor import get_proxy_mode + proxy_mode = get_proxy_mode() + if proxy_mode is not None: + graph_len_before = len(proxy_mode.tracer.graph.nodes) + + out = func(*args, **kwargs) + + # HOPs don't support func._schema + # HOPs don't alias -> this is always true today and will be always true for a long time + # TODO HOPs don't mutate -> this is always true today but will not be true forever + if isinstance(func, torch._ops.HigherOrderOperator): + any_ret_has_alias_info = False + else: + any_ret_has_alias_info = any(ret.alias_info is not None for ret in func._schema.returns) + + policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=False, op_output=out), + func, *args, **kwargs) + if isinstance(policy, bool): + policy = _policy_from_bool(policy) + + if is_compiling: + if proxy_mode is not None: + graph = proxy_mode.tracer.graph + num_new = len(graph.nodes) - graph_len_before + for node in itertools.islice(reversed(graph.nodes), num_new): + node.meta["recompute"] = policy + + if policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE) or is_compiling: + self.storage[func].append(tree_map(lambda x: _VersionWrapper(_maybe_detach(x, any_ret_has_alias_info)), out)) + return out + +class _CachedTorchDispatchMode(TorchDispatchMode): + @classmethod + def ignore_compile_internals(cls): + return True + + # Used together with _CachedTorchDispatchMode to implement SAC. + def __init__(self, policy_fn, storage, allow_cache_entry_mutation) -> None: + self.policy_fn = policy_fn + self.storage = storage + self.allow_cache_entry_mutation = allow_cache_entry_mutation + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + if func in SAC_IGNORED_OPS: + return func(*args, **kwargs) + + kwargs = {} if kwargs is None else kwargs + policy = self.policy_fn(SelectiveCheckpointContext(is_recompute=True), + func, *args, **kwargs) + if isinstance(policy, bool): + policy = _policy_from_bool(policy) + + is_compiling = _is_compiling(func, args, kwargs) + + if policy in (CheckpointPolicy.MUST_SAVE, CheckpointPolicy.PREFER_SAVE) or is_compiling: + storage = self.storage.get(func) + if storage is None: + raise RuntimeError(f"{func} encountered during backward, but not found in storage") + if len(storage) == 0: + raise RuntimeError( + "Trying to backward an extra time. You are only allowed to backward once " + "on any region computed under selective activation checkpoint." + ) + out = tree_map(lambda x: x.get_val(self.allow_cache_entry_mutation), storage.pop(0)) + else: + out = func(*args, **kwargs) + return out + + +def create_selective_checkpoint_contexts(policy_fn_or_list, allow_cache_entry_mutation=False): + """ + Helper to avoid recomputing certain ops during activation checkpointing. + + Use this with `torch.utils.checkpoint.checkpoint` to control which + operations are recomputed during the backward pass. + + Args: + policy_fn_or_list (Callable or List): + - If a policy function is provided, it should accept a + :class:`SelectiveCheckpointContext`, the :class:`OpOverload`, args and + kwargs to the op, and return a :class:`CheckpointPolicy` enum value + indicating whether the execution of the op should be recomputed or not. + - If a list of operations is provided, it is equivalent to a policy + returning `CheckpointPolicy.MUST_SAVE` for the specified + operations and `CheckpointPolicy.PREFER_RECOMPUTE` for all other + operations. + allow_cache_entry_mutation (bool, optional): By default, an error is + raised if any tensors cached by selective activation checkpoint are + mutated in order to ensure correctness. If set to `True`, this check + is disabled. + Returns: + A tuple of two context managers. + + Example: + >>> # xdoctest: +REQUIRES(LINUX) + >>> import functools + >>> + >>> x = torch.rand(10, 10, requires_grad=True) + >>> y = torch.rand(10, 10, requires_grad=True) + >>> + >>> ops_to_save = [ + >>> torch.ops.aten.mm.default, + >>> ] + >>> + >>> def policy_fn(ctx, op, *args, **kwargs): + >>> if op in ops_to_save: + >>> return CheckpointPolicy.MUST_SAVE + >>> else: + >>> return CheckpointPolicy.PREFER_RECOMPUTE + >>> + >>> context_fn = functools.partial(create_selective_checkpoint_contexts, policy_fn) + >>> + >>> # or equivalently + >>> context_fn = functools.partial(create_selective_checkpoint_contexts, ops_to_save) + >>> + >>> def fn(x, y): + >>> return torch.sigmoid(torch.matmul(torch.matmul(x, y), y)) * y + >>> + >>> out = torch.utils.checkpoint.checkpoint( + >>> fn, x, y, + >>> use_reentrant=False, + >>> context_fn=context_fn, + >>> ) + """ + # NB: If grad_mode is disabled, checkpoint would not run forward under + # context_fn anyway, so proceed as usual. + if isinstance(policy_fn_or_list, list): + for op in policy_fn_or_list: + if not isinstance(op, (torch._ops.OpOverload, torch._ops.HigherOrderOperator)): + _extra_msg = ( + "Please update the OpOverloadPacket to a specific OpOverload." + "For example, if you have `torch.ops.aten.mm`, change it to `torch.ops.aten.mm.default`." + ) if isinstance(op, torch._ops.OpOverloadPacket) else "" + raise ValueError( + f"Expected op in `op_list` to be an OpOverload but got: {op} " + f"of type {type(op)}. {_extra_msg}" + ) + + def policy_fn(ctx, op, *args, **kwargs): + if op in policy_fn_or_list: + return CheckpointPolicy.MUST_SAVE + else: + return CheckpointPolicy.PREFER_RECOMPUTE + elif callable(policy_fn_or_list): + policy_fn = policy_fn_or_list + else: + raise TypeError("policy_fn_or_list must be either a function or a list of ops.") + + storage: Dict[Any, List[Any]] = defaultdict(list) + return ( + _CachingTorchDispatchMode(policy_fn, storage), + _CachedTorchDispatchMode(policy_fn, storage, allow_cache_entry_mutation), + ) + +# NB: this helper wraps fn before calling checkpoint_impl. kwargs and +# saving/restoring of global state is handled here. + +def _checkpoint_without_reentrant_generator( + fn, + preserve_rng_state=True, + context_fn: Callable[[], Tuple[ContextManager, ContextManager]] = noop_context_fn, + determinism_check: str = _DEFAULT_DETERMINISM_MODE, + debug: bool = False, + early_stop: bool = True, + *args, + **kwargs +): + """Checkpointing without reentrant autograd. + + Args: + fn: describes what to run in the forward pass of the model or + part of the model. It should also know how to handle the inputs + passed as the tuple. For example, in LSTM, if user passes + ``(activation, hidden)``, :attr:`function` should correctly use the + first input as ``activation`` and the second input as ``hidden`` + preserve_rng_state(bool, optional): Omit stashing and restoring + the RNG state during each checkpoint. + Default: ``True`` + context_fn(Callable, optional): A callable returning a tuple of two + context managers. The function and its recomputation will be run + under the first and second context managers respectively. + determinism_check(str, optional): A string specifying the determinism + check to perform. By default it is set to ``"default"`` which + compares the shapes, dtypes, and devices of the recomputed tensors + against those the saved tensors. To turn off this check, specify + ``"none"``. Currently these are the only two supported values. + Please open an issue if you would like to see more determinism + checks. + debug(bool, optional): If ``True``, error messages will also include + a trace of the operators ran during the original forward computation + as well as the recomputation. + early_stop(bool, optional): If ``True``, non-reentrant checkpoint stops + recomputation as soon as it has computed all needed Tensors. Can be + overridden globally using :func:`set_checkpoint_early_stop` context + manager. Default: ``True``. + *args: Arguments to pass in to the given ``function``. + **kwargs: Keyword arguments to pass into the given ``function``. + """ + unpack_error_cb = None + + if _checkpoint_debug_enabled if _checkpoint_debug_enabled is not None else debug: + if context_fn is not noop_context_fn: + raise ValueError( + "debug=True is incompatible with non-default context_fn" + ) + context_fn, unpack_error_cb = _get_debug_context_and_cb() + + if determinism_check in _allowed_determinism_checks_to_fns: + metadata_fn = _allowed_determinism_checks_to_fns[determinism_check] + else: + raise ValueError( + f"determinism_check should be one of {list(_allowed_determinism_checks_to_fns.keys())}, " + f"but got {determinism_check}" + ) + + device_type = _infer_device_type(*args) + device_module = _get_device_module(device_type) + forward_context, recompute_context = context_fn() + if _is_compiling(fn, args, kwargs) and context_fn is not noop_context_fn: + if ( + not isinstance(forward_context, TorchDispatchMode) + or not isinstance(recompute_context, TorchDispatchMode) + ): + raise AssertionError( + "In torch.compile mode, `context_fn` arg passed to `torch.utils.checkpoint` " + "must generate a tuple of two `TorchDispatchMode`s." + ) + # Accommodates the (remote) possibility that autocast is enabled for cpu AND gpu. + device_autocast_kwargs, cpu_autocast_kwargs = _get_autocast_kwargs(device_type=device_type) + + if preserve_rng_state: + fwd_cpu_state = torch.get_rng_state() + # Don't eagerly initialize the cuda context by accident. + # (If the user intends that the context is initialized later, within their + # run_function, we SHOULD actually stash the cuda state here. Unfortunately, + # we have no way to anticipate this will happen before we run the function. + # If they do so, we raise an error.) + had_device_in_fwd = False + if getattr(device_module, "_initialized", False): + had_device_in_fwd = True + fwd_devices, fwd_device_states = get_device_states(*args) + + from torch.overrides import _get_current_function_mode_stack + from torch.utils._device import DeviceContext + + # recompute_fn should respect the device context of the original forward + device_ctx = next( + filter( + lambda mode: isinstance(mode, DeviceContext), + reversed(_get_current_function_mode_stack()), + ), + contextlib.nullcontext(), + ) + error_on_nested_fx_trace = torch._dynamo.config.error_on_nested_fx_trace + is_non_strict_tracing = torch.compiler._is_non_strict_tracing() + + def recompute_fn(*args) -> None: + # This will be called later during recomputation. This wrapping enables + # the necessary global state to be captured. + rng_devices = [] + if preserve_rng_state and had_device_in_fwd: + rng_devices = fwd_devices + with torch.random.fork_rng( + devices=rng_devices, enabled=preserve_rng_state, device_type=device_type + ): + if preserve_rng_state: + torch.set_rng_state(fwd_cpu_state) + if had_device_in_fwd: + set_device_states(fwd_devices, fwd_device_states, device_type=device_type) + + device_autocast_ctx = torch.amp.autocast( + device_type=device_type, **device_autocast_kwargs + ) if torch.amp.is_autocast_available(device_type) else contextlib.nullcontext() + nested_fx_trace_ctx = ( + torch._dynamo.config.patch( + error_on_nested_fx_trace=error_on_nested_fx_trace + ) + if is_non_strict_tracing + else contextlib.nullcontext() + ) + with ( + device_autocast_ctx, + torch.amp.autocast("cpu", **cpu_autocast_kwargs), + recompute_context, + device_ctx, + nested_fx_trace_ctx, + ): # type: ignore[attr-defined] + fn(*args, **kwargs) + + new_frame = _CheckpointFrame( + recompute_fn, + _enable_checkpoint_early_stop if _enable_checkpoint_early_stop is not None else early_stop, + unpack_error_cb, + metadata_fn + ) + + if not torch.is_grad_enabled(): + yield + return + + new_frame.save_inputs(*args) + + with _checkpoint_hook(new_frame), forward_context: + yield + new_frame.forward_completed = True + + if getattr(device_module, "_initialized", False) and \ + preserve_rng_state and not had_device_in_fwd: # type: ignore[possibly-undefined] + # Device was not initialized before running the forward, so we didn't + # stash the device state. + raise RuntimeError( + "PyTorch's device state was initialized in the forward pass " + "of a Checkpoint, which is not allowed. Please open an issue " + "if you need this feature." + ) + + return + + +class GraphExecGroup: + """Any checkpointed regions encountered by backward under the same instance + of this context manager will trigger recompute at most once, even if + there are multiple calls to backward. + + Backward calls under the same instance of this context manager must execute + over non-overlapping regions of the backward graph even if retain_graph=True. + In particular, any two backward call cannot use the same saved activation for + gradient computation. + + .. note:: + This context manager only affects checkpoint with use_reentrant=False, and + is a no-op otherwise. + """ + + def __enter__(self) -> Self: + if torch._C._get_graph_exec_group() is not None: + raise RuntimeError( + "GraphExecGroup contexts cannot be nested. " + f"Already inside group {torch._C._get_graph_exec_group()}" + ) + torch._C._set_graph_exec_group(self) + return self + + def __exit__(self, *args: object) -> None: + torch._C._set_graph_exec_group(None) + + @classmethod + def _get_current_group(cls) -> GraphExecGroup | None: + # Private API to be used by utils like AC + return torch._C._get_graph_exec_group() + + +# Note: [compiled autograd and checkpoint unpack hook] +# When tracing via compiled autograd, this hook will be visible to the +# compiler if the forward of this checkpointed region ran in eager. +# If the forward had ran under compile, it would have been wrapped in a +# higher order op. See Note: [torch.compile and checkpoint]. +# +# Since we run the recomputation hook under a enable_grad context, +# AOTDispatch will trace a joint graph for this hook, and may +# save different activations than in eager. This conflicts with the +# strict activation count checks in `frame.check_recomputed_tensors_match`. +# So, we disable this hook to force it to recompute eager checkpointed regions +# in eager. This could be removed if we can disable the partitioner for this +# graph segment. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/collect_env.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/collect_env.py new file mode 100644 index 0000000000000000000000000000000000000000..b566cd94e5d0314afc50150b24e2af7cb984a390 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/collect_env.py @@ -0,0 +1,944 @@ +# mypy: allow-untyped-defs + +# Unlike the rest of the PyTorch this file must be python2 compliant. +# This script outputs relevant system environment info +# Run it with `python collect_env.py` or `python -m torch.utils.collect_env` +import datetime +import json +import locale +import os +import re +import subprocess +import sys +from collections import namedtuple +from typing import cast as _cast, Dict as _Dict + + +try: + import torch + + TORCH_AVAILABLE = True +except (ImportError, NameError, AttributeError, OSError): + TORCH_AVAILABLE = False + +# System Environment Information +SystemEnv = namedtuple( + "SystemEnv", + [ + "torch_version", + "is_debug_build", + "cuda_compiled_version", + "gcc_version", + "clang_version", + "cmake_version", + "os", + "libc_version", + "python_version", + "python_platform", + "is_cuda_available", + "cuda_runtime_version", + "cuda_module_loading", + "nvidia_driver_version", + "nvidia_gpu_models", + "cudnn_version", + "is_xpu_available", + "pip_version", # 'pip' or 'pip3' + "pip_packages", + "conda_packages", + "hip_compiled_version", + "hip_runtime_version", + "miopen_runtime_version", + "caching_allocator_config", + "is_xnnpack_available", + "cpu_info", + ], +) + +COMMON_PATTERNS = [ + "torch", + "numpy", + "triton", + "optree", +] + +NVIDIA_PATTERNS = [ + "cuda-cudart", + "cuda-cupti", + "cuda-libraries", + "cuda-opencl", + "cuda-nvrtc", + "cuda-runtime", + "cublas", + "cudnn", + "cufft", + "curand", + "cusolver", + "cusparse", + "nccl", + "nvjitlink", + "nvtx", +] + +ONEAPI_PATTERNS = [ + "dpcpp-cpp-rt", + "intel-cmplr-lib-rt", + "intel-cmplr-lib-ur", + "intel-cmplr-lic-rt", + "intel-opencl-rt", + "intel-sycl-rt", + "mkl", + "onemkl-sycl-blas", + "onemkl-sycl-dft", + "onemkl-sycl-lapack", + "onemkl-sycl-rng", + "onemkl-sycl-sparse", + "intel-openmp", + "tbb", + "impi-rt", + "impi-devel", + "oneccl", + "oneccl-devel", + "intel-pti", + "umf", + "tcmlib", +] + +CONDA_PATTERNS = [ + "cudatoolkit", + "soumith", + "mkl", + "magma", +] + +PIP_PATTERNS = [ + "mypy", + "flake8", + "onnx", +] + + +def run(command): + """Return (return-code, stdout, stderr).""" + shell = type(command) is str + p = subprocess.Popen( + command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=shell + ) + raw_output, raw_err = p.communicate() + rc = p.returncode + if get_platform() == "win32": + enc = "oem" + else: + enc = locale.getpreferredencoding() + output = raw_output.decode(enc) + err = raw_err.decode(enc) + return rc, output.strip(), err.strip() + + +def run_and_read_all(run_lambda, command): + """Run command using run_lambda; reads and returns entire output if rc is 0.""" + rc, out, _ = run_lambda(command) + if rc != 0: + return None + return out + + +def run_and_parse_first_match(run_lambda, command, regex): + """Run command using run_lambda, returns the first regex match if it exists.""" + rc, out, _ = run_lambda(command) + if rc != 0: + return None + match = re.search(regex, out) + if match is None: + return None + return match.group(1) + + +def run_and_return_first_line(run_lambda, command): + """Run command using run_lambda and returns first line if output is not empty.""" + rc, out, _ = run_lambda(command) + if rc != 0: + return None + return out.split("\n")[0] + + +def get_conda_packages(run_lambda, patterns=None): + if patterns is None: + patterns = CONDA_PATTERNS + COMMON_PATTERNS + NVIDIA_PATTERNS + ONEAPI_PATTERNS + conda = os.environ.get("CONDA_EXE", "conda") + out = run_and_read_all(run_lambda, "{} list".format(conda)) + if out is None: + return out + + return "\n".join( + line + for line in out.splitlines() + if not line.startswith("#") and any(name in line for name in patterns) + ) + + +def get_gcc_version(run_lambda): + return run_and_parse_first_match(run_lambda, "gcc --version", r"gcc (.*)") + + +def get_clang_version(run_lambda): + return run_and_parse_first_match( + run_lambda, "clang --version", r"clang version (.*)" + ) + + +def get_cmake_version(run_lambda): + return run_and_parse_first_match(run_lambda, "cmake --version", r"cmake (.*)") + + +def get_nvidia_driver_version(run_lambda): + if get_platform() == "darwin": + cmd = "kextstat | grep -i cuda" + return run_and_parse_first_match( + run_lambda, cmd, r"com[.]nvidia[.]CUDA [(](.*?)[)]" + ) + smi = get_nvidia_smi() + return run_and_parse_first_match(run_lambda, smi, r"Driver Version: (.*?) ") + + +def get_gpu_info(run_lambda): + if get_platform() == "darwin" or ( + TORCH_AVAILABLE + and hasattr(torch.version, "hip") + and torch.version.hip is not None + ): + if TORCH_AVAILABLE and torch.cuda.is_available(): + if torch.version.hip is not None: + prop = torch.cuda.get_device_properties(0) + if hasattr(prop, "gcnArchName"): + gcnArch = " ({})".format(prop.gcnArchName) + else: + gcnArch = "NoGCNArchNameOnOldPyTorch" + else: + gcnArch = "" + return torch.cuda.get_device_name(None) + gcnArch + return None + smi = get_nvidia_smi() + uuid_regex = re.compile(r" \(UUID: .+?\)") + rc, out, _ = run_lambda(smi + " -L") + if rc != 0: + return None + # Anonymize GPUs by removing their UUID + return re.sub(uuid_regex, "", out) + + +def get_running_cuda_version(run_lambda): + return run_and_parse_first_match(run_lambda, "nvcc --version", r"release .+ V(.*)") + + +def get_cudnn_version(run_lambda): + """Return a list of libcudnn.so; it's hard to tell which one is being used.""" + if get_platform() == "win32": + system_root = os.environ.get("SYSTEMROOT", "C:\\Windows") + cuda_path = os.environ.get("CUDA_PATH", "%CUDA_PATH%") + where_cmd = os.path.join(system_root, "System32", "where") + cudnn_cmd = '{} /R "{}\\bin" cudnn*.dll'.format(where_cmd, cuda_path) + elif get_platform() == "darwin": + # CUDA libraries and drivers can be found in /usr/local/cuda/. See + # https://docs.nvidia.com/cuda/archive/9.0/cuda-installation-guide-mac-os-x/index.html#installation + # https://docs.nvidia.com/deeplearning/cudnn/installation/latest/ + # Use CUDNN_LIBRARY when cudnn library is installed elsewhere. + cudnn_cmd = "ls /usr/local/cuda/lib/libcudnn*" + else: + cudnn_cmd = 'ldconfig -p | grep libcudnn | rev | cut -d" " -f1 | rev' + rc, out, _ = run_lambda(cudnn_cmd) + # find will return 1 if there are permission errors or if not found + if len(out) == 0 or (rc != 1 and rc != 0): + l = os.environ.get("CUDNN_LIBRARY") + if l is not None and os.path.isfile(l): + return os.path.realpath(l) + return None + files_set = set() + for fn in out.split("\n"): + fn = os.path.realpath(fn) # eliminate symbolic links + if os.path.isfile(fn): + files_set.add(fn) + if not files_set: + return None + # Alphabetize the result because the order is non-deterministic otherwise + files = sorted(files_set) + if len(files) == 1: + return files[0] + result = "\n".join(files) + return "Probably one of the following:\n{}".format(result) + + +def get_nvidia_smi(): + # Note: nvidia-smi is currently available only on Windows and Linux + smi = "nvidia-smi" + if get_platform() == "win32": + system_root = os.environ.get("SYSTEMROOT", "C:\\Windows") + program_files_root = os.environ.get("PROGRAMFILES", "C:\\Program Files") + legacy_path = os.path.join( + program_files_root, "NVIDIA Corporation", "NVSMI", smi + ) + new_path = os.path.join(system_root, "System32", smi) + smis = [new_path, legacy_path] + for candidate_smi in smis: + if os.path.exists(candidate_smi): + smi = '"{}"'.format(candidate_smi) + break + return smi + + +def _detect_linux_pkg_manager(): + if get_platform() != "linux": + return "N/A" + for mgr_name in ["dpkg", "dnf", "yum", "zypper"]: + rc, _, _ = run(f"which {mgr_name}") + if rc == 0: + return mgr_name + return "N/A" + + +def get_linux_pkg_version(run_lambda, pkg_name): + pkg_mgr = _detect_linux_pkg_manager() + if pkg_mgr == "N/A": + return "N/A" + + grep_version = { + "dpkg": { + "field_index": 2, + "command": "dpkg -l | grep {}", + }, + "dnf": { + "field_index": 1, + "command": "dnf list | grep {}", + }, + "yum": { + "field_index": 1, + "command": "yum list | grep {}", + }, + "zypper": { + "field_index": 2, + "command": "zypper info {} | grep Version", + }, + } + + field_index: int = int(_cast(int, grep_version[pkg_mgr]["field_index"])) + cmd: str = str(grep_version[pkg_mgr]["command"]) + cmd = cmd.format(pkg_name) + ret = run_and_read_all(run_lambda, cmd) + if ret is None or ret == "": + return "N/A" + lst = re.sub(" +", " ", ret).split(" ") + if len(lst) <= field_index: + return "N/A" + return lst[field_index] + + +def get_intel_gpu_driver_version(run_lambda): + lst = [] + platform = get_platform() + if platform == "linux": + pkgs = { # type: ignore[var-annotated] + "dpkg": { + "intel-opencl-icd", + "libze1", + "level-zero", + }, + "dnf": { + "intel-opencl", + "level-zero", + }, + "yum": { + "intel-opencl", + "level-zero", + }, + "zypper": { + "intel-opencl", + "level-zero", + }, + }.get(_detect_linux_pkg_manager(), {}) + for pkg in pkgs: + ver = get_linux_pkg_version(run_lambda, pkg) + if ver != "N/A": + lst.append(f"* {pkg}:\t{ver}") + if platform in ["win32", "cygwin"]: + txt = run_and_read_all( + run_lambda, + 'powershell.exe "gwmi -Class Win32_PnpSignedDriver | where{$_.DeviceClass -eq \\"DISPLAY\\"\ + -and $_.Manufacturer -match \\"Intel\\"} | Select-Object -Property DeviceName,DriverVersion,DriverDate\ + | ConvertTo-Json"', + ) + try: + obj = json.loads(txt) + if type(obj) is list: + for o in obj: + lst.append( + f'* {o["DeviceName"]}: {o["DriverVersion"]} ({o["DriverDate"]})' + ) + else: + lst.append(f'* {obj["DriverVersion"]} ({obj["DriverDate"]})') + except ValueError as e: + lst.append(txt) + lst.append(str(e)) + return "\n".join(lst) + + +def get_intel_gpu_onboard(run_lambda): + lst: list[str] = [] + platform = get_platform() + if platform == "linux": + txt = run_and_read_all(run_lambda, "xpu-smi discovery -j") + if txt: + try: + obj = json.loads(txt) + device_list = obj.get("device_list", []) + if isinstance(device_list, list) and device_list: + lst.extend(f'* {device["device_name"]}' for device in device_list) + else: + lst.append("N/A") + except (ValueError, TypeError) as e: + lst.append(txt) + lst.append(str(e)) + else: + lst.append("N/A") + if platform in ["win32", "cygwin"]: + txt = run_and_read_all( + run_lambda, + 'powershell.exe "gwmi -Class Win32_PnpSignedDriver | where{$_.DeviceClass -eq \\"DISPLAY\\"\ + -and $_.Manufacturer -match \\"Intel\\"} | Select-Object -Property DeviceName | ConvertTo-Json"', + ) + if txt: + try: + obj = json.loads(txt) + if isinstance(obj, list) and obj: + lst.extend(f'* {device["DeviceName"]}' for device in obj) + else: + lst.append(f'* {obj.get("DeviceName", "N/A")}') + except ValueError as e: + lst.append(txt) + lst.append(str(e)) + else: + lst.append("N/A") + return "\n".join(lst) + + +def get_intel_gpu_detected(run_lambda): + if not TORCH_AVAILABLE or not hasattr(torch, "xpu"): + return "N/A" + + device_count = torch.xpu.device_count() + if device_count == 0: + return "N/A" + + devices = [ + f"* [{i}] {torch.xpu.get_device_properties(i)}" for i in range(device_count) + ] + return "\n".join(devices) + + +# example outputs of CPU infos +# * linux +# Architecture: x86_64 +# CPU op-mode(s): 32-bit, 64-bit +# Address sizes: 46 bits physical, 48 bits virtual +# Byte Order: Little Endian +# CPU(s): 128 +# On-line CPU(s) list: 0-127 +# Vendor ID: GenuineIntel +# Model name: Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz +# CPU family: 6 +# Model: 106 +# Thread(s) per core: 2 +# Core(s) per socket: 32 +# Socket(s): 2 +# Stepping: 6 +# BogoMIPS: 5799.78 +# Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr +# sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc arch_perfmon rep_good nopl +# xtopology nonstop_tsc cpuid aperfmperf tsc_known_freq pni pclmulqdq monitor ssse3 fma cx16 +# pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand +# hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced +# fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap +# avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 +# xsaves wbnoinvd ida arat avx512vbmi pku ospke avx512_vbmi2 gfni vaes vpclmulqdq +# avx512_vnni avx512_bitalg tme avx512_vpopcntdq rdpid md_clear flush_l1d arch_capabilities +# Virtualization features: +# Hypervisor vendor: KVM +# Virtualization type: full +# Caches (sum of all): +# L1d: 3 MiB (64 instances) +# L1i: 2 MiB (64 instances) +# L2: 80 MiB (64 instances) +# L3: 108 MiB (2 instances) +# NUMA: +# NUMA node(s): 2 +# NUMA node0 CPU(s): 0-31,64-95 +# NUMA node1 CPU(s): 32-63,96-127 +# Vulnerabilities: +# Itlb multihit: Not affected +# L1tf: Not affected +# Mds: Not affected +# Meltdown: Not affected +# Mmio stale data: Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown +# Retbleed: Not affected +# Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp +# Spectre v1: Mitigation; usercopy/swapgs barriers and __user pointer sanitization +# Spectre v2: Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS SW sequence +# Srbds: Not affected +# Tsx async abort: Not affected +# * win32 +# Architecture=9 +# CurrentClockSpeed=2900 +# DeviceID=CPU0 +# Family=179 +# L2CacheSize=40960 +# L2CacheSpeed= +# Manufacturer=GenuineIntel +# MaxClockSpeed=2900 +# Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz +# ProcessorType=3 +# Revision=27142 +# +# Architecture=9 +# CurrentClockSpeed=2900 +# DeviceID=CPU1 +# Family=179 +# L2CacheSize=40960 +# L2CacheSpeed= +# Manufacturer=GenuineIntel +# MaxClockSpeed=2900 +# Name=Intel(R) Xeon(R) Platinum 8375C CPU @ 2.90GHz +# ProcessorType=3 +# Revision=27142 + + +def get_cpu_info(run_lambda): + rc, out, err = 0, "", "" + if get_platform() == "linux": + rc, out, err = run_lambda("lscpu") + elif get_platform() == "win32": + rc, out, err = run_lambda( + 'powershell.exe "gwmi -Class Win32_Processor | Select-Object -Property Name,Manufacturer,Family,\ + Architecture,ProcessorType,DeviceID,CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision\ + | ConvertTo-Json"' + ) + if rc == 0: + lst = [] + try: + obj = json.loads(out) + if type(obj) is list: + for o in obj: + lst.append("----------------------") + lst.extend([f"{k}: {v}" for (k, v) in o.items()]) + else: + lst.extend([f"{k}: {v}" for (k, v) in obj.items()]) + except ValueError as e: + lst.append(out) + lst.append(str(e)) + out = "\n".join(lst) + elif get_platform() == "darwin": + rc, out, err = run_lambda("sysctl -n machdep.cpu.brand_string") + cpu_info = "None" + if rc == 0: + cpu_info = out + else: + cpu_info = err + return cpu_info + + +def get_platform(): + if sys.platform.startswith("linux"): + return "linux" + elif sys.platform.startswith("win32"): + return "win32" + elif sys.platform.startswith("cygwin"): + return "cygwin" + elif sys.platform.startswith("darwin"): + return "darwin" + else: + return sys.platform + + +def get_mac_version(run_lambda): + return run_and_parse_first_match(run_lambda, "sw_vers -productVersion", r"(.*)") + + +def get_windows_version(run_lambda): + ret = run_and_read_all( + run_lambda, + 'powershell.exe "gwmi -Class Win32_OperatingSystem | Select-Object -Property Caption,\ + OSArchitecture,Version | ConvertTo-Json"', + ) + try: + obj = json.loads(ret) + ret = f'{obj["Caption"]} ({obj["Version"]} {obj["OSArchitecture"]})' + except ValueError as e: + ret += f"\n{str(e)}" + return ret + + +def get_lsb_version(run_lambda): + return run_and_parse_first_match( + run_lambda, "lsb_release -a", r"Description:\t(.*)" + ) + + +def check_release_file(run_lambda): + return run_and_parse_first_match( + run_lambda, "cat /etc/*-release", r'PRETTY_NAME="(.*)"' + ) + + +def get_os(run_lambda): + from platform import machine + + platform = get_platform() + + if platform in ["win32", "cygwin"]: + return get_windows_version(run_lambda) + + if platform == "darwin": + version = get_mac_version(run_lambda) + if version is None: + return None + return "macOS {} ({})".format(version, machine()) + + if platform == "linux": + # Ubuntu/Debian based + desc = get_lsb_version(run_lambda) + if desc is not None: + return "{} ({})".format(desc, machine()) + + # Try reading /etc/*-release + desc = check_release_file(run_lambda) + if desc is not None: + return "{} ({})".format(desc, machine()) + + return "{} ({})".format(platform, machine()) + + # Unknown platform + return platform + + +def get_python_platform(): + import platform + + return platform.platform() + + +def get_libc_version(): + import platform + + if get_platform() != "linux": + return "N/A" + return "-".join(platform.libc_ver()) + + +def get_pip_packages(run_lambda, patterns=None): + """Return `pip list` output. Note: will also find conda-installed pytorch and numpy packages.""" + if patterns is None: + patterns = PIP_PATTERNS + COMMON_PATTERNS + NVIDIA_PATTERNS + ONEAPI_PATTERNS + + pip_version = "pip3" if sys.version_info.major == 3 else "pip" + + os.environ["PIP_DISABLE_PIP_VERSION_CHECK"] = "1" + # People generally have pip as `pip` or `pip3` + # But here it is invoked as `python -mpip` + out = run_and_read_all( + run_lambda, [sys.executable, "-mpip", "list", "--format=freeze"] + ) + if out is None: + return pip_version, out + + filtered_out = "\n".join( + line for line in out.splitlines() if any(name in line for name in patterns) + ) + + return pip_version, filtered_out + + +def get_cachingallocator_config() -> _Dict[str, str]: + """Return the caching allocator configuration from environment variables. + """ + # pyrefly: ignore [bad-return] + return { + var: os.environ.get(var) + for var in ( + "PYTORCH_CUDA_ALLOC_CONF", + "PYTORCH_HIP_ALLOC_CONF", + "PYTORCH_ALLOC_CONF", + ) + if os.environ.get(var) + } + + +def get_cuda_module_loading_config(): + if TORCH_AVAILABLE and torch.cuda.is_available(): + torch.cuda.init() + config = os.environ.get("CUDA_MODULE_LOADING", "") + return config + else: + return "N/A" + + +def is_xnnpack_available(): + if TORCH_AVAILABLE: + import torch.backends.xnnpack + + return str(torch.backends.xnnpack.enabled) # type: ignore[attr-defined] + else: + return "N/A" + + +def get_env_info(): + """ + Collects environment information to aid in debugging. + + The returned environment information contains details on torch version, is debug build + or not, cuda compiled version, gcc version, clang version, cmake version, operating + system, libc version, python version, python platform, CUDA availability, CUDA + runtime version, CUDA module loading config, GPU model and configuration, Nvidia + driver version, cuDNN version, pip version and versions of relevant pip and + conda packages, HIP runtime version, MIOpen runtime version, + Caching allocator config, XNNPACK availability and CPU information. + + Returns: + SystemEnv (namedtuple): A tuple containing various environment details + and system information. + """ + run_lambda = run + pip_version, pip_list_output = get_pip_packages(run_lambda) + + if TORCH_AVAILABLE: + version_str = torch.__version__ + debug_mode_str = str(torch.version.debug) + cuda_available_str = str(torch.cuda.is_available()) + cuda_version_str = torch.version.cuda + xpu_available_str = str(torch.xpu.is_available()) + if torch.xpu.is_available(): + xpu_available_str = ( + f"{xpu_available_str}\n" + + f"XPU used to build PyTorch: {torch.version.xpu}\n" + + f"Intel GPU driver version:\n{get_intel_gpu_driver_version(run_lambda)}\n" + + f"Intel GPU models onboard:\n{get_intel_gpu_onboard(run_lambda)}\n" + + f"Intel GPU models detected:\n{get_intel_gpu_detected(run_lambda)}" + ) + if ( + not hasattr(torch.version, "hip") or torch.version.hip is None + ): # cuda version + hip_compiled_version = hip_runtime_version = miopen_runtime_version = "N/A" + else: # HIP version + + def get_version_or_na(cfg, prefix): + _lst = [s.rsplit(None, 1)[-1] for s in cfg if prefix in s] + return _lst[0] if _lst else "N/A" + + cfg = torch._C._show_config().split("\n") + hip_runtime_version = get_version_or_na(cfg, "HIP Runtime") + miopen_runtime_version = get_version_or_na(cfg, "MIOpen") + cuda_version_str = "N/A" + hip_compiled_version = torch.version.hip + else: + version_str = debug_mode_str = cuda_available_str = cuda_version_str = xpu_available_str = "N/A" # type: ignore[assignment] + hip_compiled_version = hip_runtime_version = miopen_runtime_version = "N/A" + + sys_version = sys.version.replace("\n", " ") + + conda_packages = get_conda_packages(run_lambda) + + return SystemEnv( + torch_version=version_str, + is_debug_build=debug_mode_str, + python_version="{} ({}-bit runtime)".format( + sys_version, sys.maxsize.bit_length() + 1 + ), + python_platform=get_python_platform(), + is_cuda_available=cuda_available_str, + cuda_compiled_version=cuda_version_str, + cuda_runtime_version=get_running_cuda_version(run_lambda), + cuda_module_loading=get_cuda_module_loading_config(), + nvidia_gpu_models=get_gpu_info(run_lambda), + nvidia_driver_version=get_nvidia_driver_version(run_lambda), + cudnn_version=get_cudnn_version(run_lambda), + is_xpu_available=xpu_available_str, + hip_compiled_version=hip_compiled_version, + hip_runtime_version=hip_runtime_version, + miopen_runtime_version=miopen_runtime_version, + pip_version=pip_version, + pip_packages=pip_list_output, + conda_packages=conda_packages, + os=get_os(run_lambda), + libc_version=get_libc_version(), + gcc_version=get_gcc_version(run_lambda), + clang_version=get_clang_version(run_lambda), + cmake_version=get_cmake_version(run_lambda), + caching_allocator_config=get_cachingallocator_config(), + is_xnnpack_available=is_xnnpack_available(), + cpu_info=get_cpu_info(run_lambda), + ) + + +env_info_fmt = """ +PyTorch version: {torch_version} +Is debug build: {is_debug_build} +CUDA used to build PyTorch: {cuda_compiled_version} +ROCM used to build PyTorch: {hip_compiled_version} + +OS: {os} +GCC version: {gcc_version} +Clang version: {clang_version} +CMake version: {cmake_version} +Libc version: {libc_version} + +Python version: {python_version} +Python platform: {python_platform} +Is CUDA available: {is_cuda_available} +CUDA runtime version: {cuda_runtime_version} +CUDA_MODULE_LOADING set to: {cuda_module_loading} +GPU models and configuration: {nvidia_gpu_models} +Nvidia driver version: {nvidia_driver_version} +cuDNN version: {cudnn_version} +Is XPU available: {is_xpu_available} +HIP runtime version: {hip_runtime_version} +MIOpen runtime version: {miopen_runtime_version} +Is XNNPACK available: {is_xnnpack_available} +Caching allocator config: {caching_allocator_config} + +CPU: +{cpu_info} + +Versions of relevant libraries: +{pip_packages} +{conda_packages} +""".strip() + + +def pretty_str(envinfo): + def replace_nones(dct, replacement="Could not collect"): + for key in dct: + if dct[key] is not None: + continue + dct[key] = replacement + return dct + + def replace_bools(dct, true="Yes", false="No"): + for key in dct: + if dct[key] is True: + dct[key] = true + elif dct[key] is False: + dct[key] = false + return dct + + def prepend(text, tag="[prepend]"): + lines = text.split("\n") + updated_lines = [tag + line for line in lines] + return "\n".join(updated_lines) + + def replace_if_empty(text, replacement="No relevant packages"): + if text is not None and len(text) == 0: + return replacement + return text + + def maybe_start_on_next_line(string): + # If `string` is multiline, prepend a \n to it. + if string is not None and len(string.split("\n")) > 1: + return "\n{}\n".format(string) + return string + + mutable_dict = envinfo._asdict() + + # If nvidia_gpu_models is multiline, start on the next line + mutable_dict["nvidia_gpu_models"] = maybe_start_on_next_line( + envinfo.nvidia_gpu_models + ) + + # If the machine doesn't have CUDA, report some fields as 'No CUDA' + dynamic_cuda_fields = [ + "cuda_runtime_version", + "nvidia_gpu_models", + "nvidia_driver_version", + ] + all_cuda_fields = dynamic_cuda_fields + ["cudnn_version"] + all_dynamic_cuda_fields_missing = all( + mutable_dict[field] is None for field in dynamic_cuda_fields + ) + if ( + TORCH_AVAILABLE + and not torch.cuda.is_available() + and all_dynamic_cuda_fields_missing + ): + for field in all_cuda_fields: + mutable_dict[field] = "No CUDA" + if envinfo.cuda_compiled_version is None: + mutable_dict["cuda_compiled_version"] = "None" + + # Replace True with Yes, False with No + mutable_dict = replace_bools(mutable_dict) + + # Replace all None objects with 'Could not collect' + mutable_dict = replace_nones(mutable_dict) + + # If either of these are '', replace with 'No relevant packages' + mutable_dict["pip_packages"] = replace_if_empty(mutable_dict["pip_packages"]) + mutable_dict["conda_packages"] = replace_if_empty(mutable_dict["conda_packages"]) + + # Tag conda and pip packages with a prefix + # If they were previously None, they'll show up as ie '[conda] Could not collect' + if mutable_dict["pip_packages"]: + mutable_dict["pip_packages"] = prepend( + mutable_dict["pip_packages"], "[{}] ".format(envinfo.pip_version) + ) + if mutable_dict["conda_packages"]: + mutable_dict["conda_packages"] = prepend( + mutable_dict["conda_packages"], "[conda] " + ) + mutable_dict["cpu_info"] = envinfo.cpu_info + mutable_dict["caching_allocator_config"] = envinfo.caching_allocator_config + if not envinfo.caching_allocator_config: + mutable_dict["caching_allocator_config"] = "N/A" + return env_info_fmt.format(**mutable_dict) + + +def get_pretty_env_info(): + """ + Returns a pretty string of environment information. + + This function retrieves environment information by calling the `get_env_info` function + and then formats the information into a human-readable string. The retrieved environment + information is listed in the document of `get_env_info`. + This function is used in `python collect_env.py` that should be executed when reporting a bug. + + Returns: + str: A pretty string of the environment information. + """ + return pretty_str(get_env_info()) + + +def main() -> None: + print("Collecting environment information...") + output = get_pretty_env_info() + print(output) + + if ( + TORCH_AVAILABLE + and hasattr(torch, "utils") + and hasattr(torch.utils, "_crash_handler") + ): + minidump_dir = torch.utils._crash_handler.DEFAULT_MINIDUMP_DIR + if sys.platform == "linux" and os.path.exists(minidump_dir): + dumps = [ + os.path.join(minidump_dir, dump) for dump in os.listdir(minidump_dir) + ] + latest = max(dumps, key=os.path.getctime) + ctime = os.path.getctime(latest) + creation_time = datetime.datetime.fromtimestamp(ctime).strftime( + "%Y-%m-%d %H:%M:%S" + ) + msg = ( + "\n*** Detected a minidump at {} created on {}, ".format( + latest, creation_time + ) + + "if this is related to your bug please include it when you file a report ***" + ) + print(msg, file=sys.stderr) + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/cpp_backtrace.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/cpp_backtrace.py new file mode 100644 index 0000000000000000000000000000000000000000..af4a7fcb63e263038255359c946a6a0d4a21dbd0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/cpp_backtrace.py @@ -0,0 +1,12 @@ +# mypy: allow-untyped-defs +from torch._C import _get_cpp_backtrace + +def get_cpp_backtrace(frames_to_skip=0, maximum_number_of_frames=64) -> str: + r""" + Return a string containing the C++ stack trace of the current thread. + + Args: + frames_to_skip (int): the number of frames to skip from the top of the stack + maximum_number_of_frames (int): the maximum number of frames to return + """ + return _get_cpp_backtrace(frames_to_skip, maximum_number_of_frames) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/cpp_extension.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/cpp_extension.py new file mode 100644 index 0000000000000000000000000000000000000000..3e9d5d3823cbe00a13bf4e21849edb26b43b6f8a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/cpp_extension.py @@ -0,0 +1,3203 @@ +# mypy: allow-untyped-defs +import copy +import glob +import importlib +import importlib.abc +import os +import re +import shlex +import shutil +import setuptools +import subprocess +import sys +import sysconfig +import types +import collections +from pathlib import Path +import errno +import logging + +logger = logging.getLogger(__name__) + +import torch +import torch._appdirs +from .file_baton import FileBaton +from ._cpp_extension_versioner import ExtensionVersioner +from typing_extensions import deprecated +from torch.torch_version import TorchVersion, Version + + +from setuptools.command.build_ext import build_ext + +IS_WINDOWS = sys.platform == 'win32' +IS_MACOS = sys.platform.startswith('darwin') +IS_LINUX = sys.platform.startswith('linux') +LIB_EXT = '.pyd' if IS_WINDOWS else '.so' +EXEC_EXT = '.exe' if IS_WINDOWS else '' +CLIB_PREFIX = '' if IS_WINDOWS else 'lib' +CLIB_EXT = '.dll' if IS_WINDOWS else '.so' +SHARED_FLAG = '/DLL' if IS_WINDOWS else '-shared' + +_HERE = os.path.abspath(__file__) +_TORCH_PATH = os.path.dirname(os.path.dirname(_HERE)) +TORCH_LIB_PATH = os.path.join(_TORCH_PATH, 'lib') + + +SUBPROCESS_DECODE_ARGS = ('oem',) if IS_WINDOWS else () +MINIMUM_GCC_VERSION = (5, 0, 0) +MINIMUM_MSVC_VERSION = (19, 0, 24215) + +VersionRange = tuple[tuple[int, ...], tuple[int, ...]] +VersionMap = dict[str, VersionRange] +# The following values were taken from the following GitHub gist that +# summarizes the minimum valid major versions of g++/clang++ for each supported +# CUDA version: https://gist.github.com/ax3l/9489132 +# Or from include/crt/host_config.h in the CUDA SDK +# The second value is the exclusive(!) upper bound, i.e. min <= version < max +CUDA_GCC_VERSIONS: VersionMap = { + '11.0': (MINIMUM_GCC_VERSION, (10, 0)), + '11.1': (MINIMUM_GCC_VERSION, (11, 0)), + '11.2': (MINIMUM_GCC_VERSION, (11, 0)), + '11.3': (MINIMUM_GCC_VERSION, (11, 0)), + '11.4': ((6, 0, 0), (12, 0)), + '11.5': ((6, 0, 0), (12, 0)), + '11.6': ((6, 0, 0), (12, 0)), + '11.7': ((6, 0, 0), (12, 0)), + '12.0': ((6, 0, 0), (13, 0)), + '12.1': ((6, 0, 0), (13, 0)), + '12.2': ((6, 0, 0), (13, 0)), + '12.3': ((6, 0, 0), (14, 0)), + '12.4': ((6, 0, 0), (14, 0)), + '12.5': ((6, 0, 0), (14, 0)), + '12.6': ((6, 0, 0), (14, 0)), + '12.7': ((6, 0, 0), (14, 0)), + '12.8': ((6, 0, 0), (15, 0)), + '12.9': ((6, 0, 0), (15, 0)), + '13.0': ((6, 0, 0), (16, 0)), +} + +MINIMUM_CLANG_VERSION = (3, 3, 0) +CUDA_CLANG_VERSIONS: VersionMap = { + '11.1': (MINIMUM_CLANG_VERSION, (11, 0)), + '11.2': (MINIMUM_CLANG_VERSION, (12, 0)), + '11.3': (MINIMUM_CLANG_VERSION, (12, 0)), + '11.4': (MINIMUM_CLANG_VERSION, (13, 0)), + '11.5': (MINIMUM_CLANG_VERSION, (13, 0)), + '11.6': (MINIMUM_CLANG_VERSION, (14, 0)), + '11.7': (MINIMUM_CLANG_VERSION, (14, 0)), + '12.0': ((7, 0), (15, 0)), + '12.1': ((7, 0), (15, 0)), + '12.2': ((7, 0), (16, 0)), + '12.3': ((7, 0), (16, 0)), + '12.4': ((7, 0), (17, 0)), + '12.5': ((7, 0), (18, 0)), + '12.6': ((7, 0), (18, 0)), + '12.7': ((7, 0), (19, 0)), + '12.8': ((7, 0), (19, 0)), + '12.9': ((7, 0), (19, 0)), + '13.0': ((7, 0), (21, 0)), +} + +__all__ = ["get_default_build_root", "check_compiler_ok_for_platform", "get_compiler_abi_compatibility_and_version", "BuildExtension", + "CppExtension", "CUDAExtension", "SyclExtension", "include_paths", "library_paths", "load", "load_inline", "is_ninja_available", + "verify_ninja_availability", "remove_extension_h_precompiler_headers", "get_cxx_compiler", "check_compiler_is_gcc"] +# Taken directly from python stdlib < 3.9 +# See https://github.com/pytorch/pytorch/issues/48617 +def _nt_quote_args(args: list[str] | None) -> list[str]: + """Quote command-line arguments for DOS/Windows conventions. + + Just wraps every argument which contains blanks in double quotes, and + returns a new argument list. + """ + # Cover None-type + if not args: + return [] + return [f'"{arg}"' if ' ' in arg else arg for arg in args] + +def _find_cuda_home() -> str | None: + """Find the CUDA install path.""" + # Guess #1 + cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') + if cuda_home is None: + # Guess #2 + nvcc_path = shutil.which("nvcc") + if nvcc_path is not None: + cuda_home = os.path.dirname(os.path.dirname(nvcc_path)) + else: + # Guess #3 + if IS_WINDOWS: + cuda_homes = glob.glob( + 'C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v*.*') + if len(cuda_homes) == 0: + cuda_home = '' + else: + cuda_home = cuda_homes[0] + else: + cuda_home = '/usr/local/cuda' + if not os.path.exists(cuda_home): + cuda_home = None + if cuda_home and not torch.cuda.is_available(): + logger.warning("No CUDA runtime is found, using CUDA_HOME='%s'", cuda_home) + return cuda_home + +def _find_rocm_home() -> str | None: + """Find the ROCm install path.""" + # Guess #1 + rocm_home = os.environ.get('ROCM_HOME') or os.environ.get('ROCM_PATH') + if rocm_home is None: + # Guess #2 + hipcc_path = shutil.which('hipcc') + if hipcc_path is not None: + rocm_home = os.path.dirname(os.path.dirname( + os.path.realpath(hipcc_path))) + # can be either /hip/bin/hipcc or /bin/hipcc + if os.path.basename(rocm_home) == 'hip': + rocm_home = os.path.dirname(rocm_home) + else: + # Guess #3 + fallback_path = '/opt/rocm' + if os.path.exists(fallback_path): + rocm_home = fallback_path + if rocm_home and torch.version.hip is None: + logger.warning("No ROCm runtime is found, using ROCM_HOME='%s'", rocm_home) + return rocm_home + +def _find_sycl_home() -> str | None: + sycl_home = None + icpx_path = shutil.which('icpx') + # Guess 1: for source code build developer/user, we'll have icpx in PATH, + # which will tell us the SYCL_HOME location. + if icpx_path is not None: + sycl_home = os.path.dirname(os.path.dirname( + os.path.realpath(icpx_path))) + + # Guess 2: for users install Pytorch with XPU support, the sycl runtime is + # inside intel-sycl-rt, which is automatically installed via pip dependency. + else: + try: + files = importlib.metadata.files('intel-sycl-rt') or [] + for f in files: + if f.name == "libsycl.so": + sycl_home = os.path.dirname(Path(f.locate()).parent.resolve()) + break + except importlib.metadata.PackageNotFoundError: + logger.warning("Trying to find SYCL_HOME from intel-sycl-rt package, but it is not installed.") + return sycl_home + +def _join_rocm_home(*paths) -> str: + """ + Join paths with ROCM_HOME, or raises an error if it ROCM_HOME is not set. + + This is basically a lazy way of raising an error for missing $ROCM_HOME + only once we need to get any ROCm-specific path. + """ + if ROCM_HOME is None: + raise OSError('ROCM_HOME environment variable is not set. ' + 'Please set it to your ROCm install root.') + return os.path.join(ROCM_HOME, *paths) + +def _join_sycl_home(*paths) -> str: + """ + Join paths with SYCL_HOME, or raises an error if it SYCL_HOME is not found. + + This is basically a lazy way of raising an error for missing SYCL_HOME + only once we need to get any SYCL-specific path. + """ + if SYCL_HOME is None: + raise OSError('SYCL runtime is not dected. Please setup the pytorch ' + 'prerequisites for Intel GPU following the instruction in ' + 'https://github.com/pytorch/pytorch?tab=readme-ov-file#intel-gpu-support ' + 'or install intel-sycl-rt via pip.') + + return os.path.join(SYCL_HOME, *paths) + + + +ABI_INCOMPATIBILITY_WARNING = ( + " !! WARNING !!" + "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" + "Your compiler (%s) may be ABI-incompatible with PyTorch!" + "Please use a compiler that is ABI-compatible with GCC 5.0 and above." + "See https://gcc.gnu.org/onlinedocs/libstdc++/manual/abi.html." + "See https://gist.github.com/goldsborough/d466f43e8ffc948ff92de7486c5216d6" + "for instructions on how to install GCC 5 or higher." + "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" + " !! WARNING !!" +) +WRONG_COMPILER_WARNING = ( + " !! WARNING !!" + "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" + "Your compiler (%s) is not compatible with the compiler Pytorch was" + "built with for this platform, which is %s on %s. Please" + "use %s to compile your extension. Alternatively, you may" + "compile PyTorch from source using %s, and then you can also use" + "%s to compile your extension." + "See https://github.com/pytorch/pytorch/blob/master/CONTRIBUTING.md for help" + "with compiling PyTorch from source." + "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" + " !! WARNING !!" +) +CUDA_MISMATCH_MESSAGE = ( + "The detected CUDA version (%s) mismatches the version that was used to compile" + "PyTorch (%s). Please make sure to use the same CUDA versions." +) +CUDA_MISMATCH_WARN = ( + "The detected CUDA version (%s) has a minor version mismatch with the version that was used to compile PyTorch (%s). Most likely this shouldn't be a problem." +) +CUDA_NOT_FOUND_MESSAGE = ( + "CUDA was not found on the system, please set the CUDA_HOME or the CUDA_PATH" + "environment variable or add NVCC to your system PATH. The extension compilation will fail." +) +ROCM_HOME = _find_rocm_home() if (torch.cuda._is_compiled() and torch.version.hip) else None +HIP_HOME = _join_rocm_home('hip') if ROCM_HOME else None +IS_HIP_EXTENSION = bool(ROCM_HOME is not None and torch.version.hip is not None) +ROCM_VERSION = None +if torch.version.hip is not None: + ROCM_VERSION = tuple(int(v) for v in torch.version.hip.split('.')[:2]) + +CUDA_HOME = _find_cuda_home() if (torch.cuda._is_compiled() and torch.version.cuda) else None +CUDNN_HOME = os.environ.get('CUDNN_HOME') or os.environ.get('CUDNN_PATH') +SYCL_HOME = _find_sycl_home() if torch.xpu._is_compiled() else None +WINDOWS_CUDA_HOME = os.environ.get('WINDOWS_CUDA_HOME') # used for AOTI cross-compilation + +# PyTorch releases have the version pattern major.minor.patch, whereas when +# PyTorch is built from source, we append the git commit hash, which gives +# it the below pattern. +BUILT_FROM_SOURCE_VERSION_PATTERN = re.compile(r'\d+\.\d+\.\d+\w+\+\w+') + +COMMON_MSVC_FLAGS = ['/MD', '/wd4819', '/wd4251', '/wd4244', '/wd4267', '/wd4275', '/wd4018', '/wd4190', '/wd4624', '/wd4067', '/wd4068', '/EHsc'] + +MSVC_IGNORE_CUDAFE_WARNINGS = [ + 'base_class_has_different_dll_interface', + 'field_without_dll_interface', + 'dll_interface_conflict_none_assumed', + 'dll_interface_conflict_dllexport_assumed' +] + +COMMON_NVCC_FLAGS = [ + '-D__CUDA_NO_HALF_OPERATORS__', + '-D__CUDA_NO_HALF_CONVERSIONS__', + '-D__CUDA_NO_BFLOAT16_CONVERSIONS__', + '-D__CUDA_NO_HALF2_OPERATORS__', + '--expt-relaxed-constexpr' +] + +COMMON_HIP_FLAGS = [ + '-D__HIP_PLATFORM_AMD__=1', + '-DUSE_ROCM=1', + '-DHIPBLAS_V2', +] + +if not IS_WINDOWS: + COMMON_HIP_FLAGS.append('-fPIC') + +COMMON_HIPCC_FLAGS = [ + '-DCUDA_HAS_FP16=1', + '-D__HIP_NO_HALF_OPERATORS__=1', + '-D__HIP_NO_HALF_CONVERSIONS__=1', + '-DHIP_ENABLE_WARP_SYNC_BUILTINS=1' +] + +if IS_WINDOWS: + # Compatibility flags, similar to those set in cmake/Dependencies.cmake. + COMMON_HIPCC_FLAGS.append('-fms-extensions') + # Suppress warnings about dllexport. + COMMON_HIPCC_FLAGS.append('-Wno-ignored-attributes') + + +def _get_icpx_version() -> str: + icpx = 'icx' if IS_WINDOWS else 'icpx' + compiler_info = subprocess.check_output([icpx, '--version']) + match = re.search(r'(\d+)\.(\d+)\.(\d+)', compiler_info.decode().strip()) + version = ['0', '0', '0'] if match is None else list(match.groups()) + version = list(map(int, version)) + if len(version) != 3: + raise AssertionError("Failed to parse DPC++ compiler version") + # Aligning version format with what torch.version.xpu() returns + return f"{version[0]}{version[1]:02}{version[2]:02}" + + +def _get_sycl_arch_list(): + if 'TORCH_XPU_ARCH_LIST' in os.environ: + return os.environ.get('TORCH_XPU_ARCH_LIST') + arch_list = torch.xpu.get_arch_list() + # Dropping dg2* archs since they lack hardware support for fp64 and require + # special consideration from the user. If needed these platforms can + # be requested thru TORCH_XPU_ARCH_LIST environment variable. + arch_list = [x for x in arch_list if not x.startswith('dg2')] + return ','.join(arch_list) + + +# If arch list returned by _get_sycl_arch_list() is empty, then sycl kernels will be compiled +# for default spir64 target and avoid device specific compilations entirely. Further, kernels +# will be JIT compiled at runtime. +def _append_sycl_targets_if_missing(cflags) -> None: + if any(flag.startswith('-fsycl-targets=') for flag in cflags): + # do nothing: user has manually specified sycl targets + return + if _get_sycl_arch_list() != '': + # AOT (spir64_gen) + JIT (spir64) + cflags.append('-fsycl-targets=spir64_gen,spir64') + else: + # JIT (spir64) + cflags.append('-fsycl-targets=spir64') + +def _get_sycl_device_flags(cflags): + # We need last occurrence of -fsycl-targets as it will be the one taking effect. + # So searching in reversed list. + flags = [f for f in reversed(cflags) if f.startswith('-fsycl-targets=')] + if not flags: + raise AssertionError("bug: -fsycl-targets should have been amended to cflags") + + arch_list = _get_sycl_arch_list() + if arch_list != '': + flags += [f'-Xs "-device {arch_list}"'] + return flags + +_COMMON_SYCL_FLAGS = [ + '-fsycl', +] + +_SYCL_DLINK_FLAGS = [ + *_COMMON_SYCL_FLAGS, + '-fsycl-link', + '--offload-compress', +] + +JIT_EXTENSION_VERSIONER = ExtensionVersioner() + +PLAT_TO_VCVARS = { + 'win32' : 'x86', + 'win-amd64' : 'x86_amd64', +} + +min_supported_cpython = "0x030A0000" # Python 3.10 hexcode + +def get_cxx_compiler(): + if IS_WINDOWS: + compiler = os.environ.get('CXX', 'cl') + else: + compiler = os.environ.get('CXX', 'c++') + return compiler + +def _is_binary_build() -> bool: + return not BUILT_FROM_SOURCE_VERSION_PATTERN.match(torch.version.__version__) + + +def _accepted_compilers_for_platform() -> list[str]: + # gnu-c++ and gnu-cc are the conda gcc compilers + return ['clang++', 'clang'] if IS_MACOS else ['g++', 'gcc', 'gnu-c++', 'gnu-cc', 'clang++', 'clang'] + +def _maybe_write(filename, new_content) -> None: + r''' + Equivalent to writing the content into the file but will not touch the file + if it already had the right content (to avoid triggering recompile). + ''' + if os.path.exists(filename): + with open(filename) as f: + content = f.read() + + if content == new_content: + # The file already contains the right thing! + return + + with open(filename, 'w') as source_file: + source_file.write(new_content) + +def get_default_build_root() -> str: + """ + Return the path to the root folder under which extensions will built. + + For each extension module built, there will be one folder underneath the + folder returned by this function. For example, if ``p`` is the path + returned by this function and ``ext`` the name of an extension, the build + folder for the extension will be ``p/ext``. + + This directory is **user-specific** so that multiple users on the same + machine won't meet permission issues. + """ + return os.path.realpath(torch._appdirs.user_cache_dir(appname='torch_extensions')) + + +def check_compiler_ok_for_platform(compiler: str) -> bool: + """ + Verify that the compiler is the expected one for the current platform. + + Args: + compiler (str): The compiler executable to check. + + Returns: + True if the compiler is gcc/g++ on Linux or clang/clang++ on macOS, + and always True for Windows. + """ + if IS_WINDOWS: + return True + compiler_path = shutil.which(compiler) + if compiler_path is None: + return False + # Use os.path.realpath to resolve any symlinks, in particular from 'c++' to e.g. 'g++'. + compiler_path = os.path.realpath(compiler_path) + # Check the compiler name + if any(name in compiler_path for name in _accepted_compilers_for_platform()): + return True + # If compiler wrapper is used try to infer the actual compiler by invoking it with -v flag + env = os.environ.copy() + env['LC_ALL'] = 'C' # Don't localize output + try: + version_string = subprocess.check_output([compiler, '-v'], stderr=subprocess.STDOUT, env=env).decode(*SUBPROCESS_DECODE_ARGS) + except subprocess.CalledProcessError: + # If '-v' fails, try '--version' + version_string = subprocess.check_output([compiler, '--version'], stderr=subprocess.STDOUT, env=env).decode(*SUBPROCESS_DECODE_ARGS) + if IS_LINUX: + # Check for 'gcc' or 'g++' for sccache wrapper + pattern = re.compile("^COLLECT_GCC=(.*)$", re.MULTILINE) + results = re.findall(pattern, version_string) + if len(results) != 1: + # Clang is also a supported compiler on Linux + # Though on Ubuntu it's sometimes called "Ubuntu clang version" + return 'clang version' in version_string + compiler_path = os.path.realpath(results[0].strip()) + # On RHEL/CentOS c++ is a gcc compiler wrapper + if os.path.basename(compiler_path) == 'c++' and 'gcc version' in version_string: + return True + return any(name in compiler_path for name in _accepted_compilers_for_platform()) + if IS_MACOS: + # Check for 'clang' or 'clang++' + return version_string.startswith("Apple clang") + return False + + +def get_compiler_abi_compatibility_and_version(compiler) -> tuple[bool, TorchVersion]: + """ + Determine if the given compiler is ABI-compatible with PyTorch alongside its version. + + Args: + compiler (str): The compiler executable name to check (e.g. ``g++``). + Must be executable in a shell process. + + Returns: + A tuple that contains a boolean that defines if the compiler is (likely) ABI-incompatible with PyTorch, + followed by a `TorchVersion` string that contains the compiler version separated by dots. + """ + if not _is_binary_build(): + return (True, TorchVersion('0.0.0')) + if os.environ.get('TORCH_DONT_CHECK_COMPILER_ABI') in ['ON', '1', 'YES', 'TRUE', 'Y']: + return (True, TorchVersion('0.0.0')) + + # First check if the compiler is one of the expected ones for the particular platform. + if not check_compiler_ok_for_platform(compiler): + logger.warning(WRONG_COMPILER_WARNING, compiler, _accepted_compilers_for_platform()[0], sys.platform, _accepted_compilers_for_platform()[0], compiler, compiler) + return (False, TorchVersion('0.0.0')) + + if IS_MACOS: + # There is no particular minimum version we need for clang, so we're good here. + return (True, TorchVersion('0.0.0')) + try: + if IS_LINUX: + minimum_required_version = MINIMUM_GCC_VERSION + compiler_info = subprocess.check_output([compiler, '-dumpfullversion', '-dumpversion']) + else: + minimum_required_version = MINIMUM_MSVC_VERSION + compiler_info = subprocess.check_output(compiler, stderr=subprocess.STDOUT) + match = re.search(r'(\d+)\.(\d+)\.(\d+)', compiler_info.decode(*SUBPROCESS_DECODE_ARGS).strip()) + version = ['0', '0', '0'] if match is None else list(match.groups()) + except (subprocess.CalledProcessError, OSError): + logger.warning('Error checking compiler version for %s', compiler, exc_info=True) + return (False, TorchVersion('0.0.0')) + + # convert alphanumeric string to numeric string + # amdclang++ returns str like 0.0.0git, others return 0.0.0 + numeric_version = [re.sub(r'\D', '', v) for v in version] + + if tuple(map(int, numeric_version)) >= minimum_required_version: + return (True, TorchVersion('.'.join(numeric_version))) + + compiler = f'{compiler} {".".join(numeric_version)}' + logger.warning(ABI_INCOMPATIBILITY_WARNING, compiler) + + return (False, TorchVersion('.'.join(numeric_version))) + + +def _check_cuda_version(compiler_name: str, compiler_version: TorchVersion) -> None: + if not CUDA_HOME: + raise RuntimeError(CUDA_NOT_FOUND_MESSAGE) + + nvcc = os.path.join(CUDA_HOME, 'bin', 'nvcc.exe' if IS_WINDOWS else 'nvcc') + if not os.path.exists(nvcc): + raise FileNotFoundError(f"nvcc not found at '{nvcc}'. Ensure CUDA path '{CUDA_HOME}' is correct.") + + cuda_version_str = subprocess.check_output([nvcc, '--version']).strip().decode(*SUBPROCESS_DECODE_ARGS) + cuda_version = re.search(r'release (\d+[.]\d+)', cuda_version_str) + if cuda_version is None: + return + + cuda_str_version = cuda_version.group(1) + cuda_ver = Version(cuda_str_version) + if torch.version.cuda is None: + return + + torch_cuda_version = Version(torch.version.cuda) + if cuda_ver != torch_cuda_version: + # major/minor attributes are only available in setuptools>=49.4.0 + if getattr(cuda_ver, "major", None) is None: + raise ValueError("setuptools>=49.4.0 is required") + if cuda_ver.major != torch_cuda_version.major: + raise RuntimeError(CUDA_MISMATCH_MESSAGE, cuda_str_version, torch.version.cuda) + logger.warning(CUDA_MISMATCH_WARN, cuda_str_version, torch.version.cuda) + + if not (sys.platform.startswith('linux') and + os.environ.get('TORCH_DONT_CHECK_COMPILER_ABI') not in ['ON', '1', 'YES', 'TRUE', 'Y'] and + _is_binary_build()): + return + + cuda_compiler_bounds: VersionMap = CUDA_CLANG_VERSIONS if compiler_name.startswith('clang') else CUDA_GCC_VERSIONS + + if cuda_str_version not in cuda_compiler_bounds: + logger.warning('There are no %s version bounds defined for CUDA version %s', compiler_name, cuda_str_version) + else: + min_compiler_version, max_excl_compiler_version = cuda_compiler_bounds[cuda_str_version] + # Special case for 11.4.0, which has lower compiler bounds than 11.4.1 + if "V11.4.48" in cuda_version_str and cuda_compiler_bounds == CUDA_GCC_VERSIONS: + max_excl_compiler_version = (11, 0) + min_compiler_version_str = '.'.join(map(str, min_compiler_version)) + max_excl_compiler_version_str = '.'.join(map(str, max_excl_compiler_version)) + + version_bound_str = f'>={min_compiler_version_str}, <{max_excl_compiler_version_str}' + + if compiler_version < TorchVersion(min_compiler_version_str): + raise RuntimeError( + f'The current installed version of {compiler_name} ({compiler_version}) is less ' + f'than the minimum required version by CUDA {cuda_str_version} ({min_compiler_version_str}). ' + f'Please make sure to use an adequate version of {compiler_name} ({version_bound_str}).' + ) + if compiler_version >= TorchVersion(max_excl_compiler_version_str): + raise RuntimeError( + f'The current installed version of {compiler_name} ({compiler_version}) is greater ' + f'than the maximum required version by CUDA {cuda_str_version}. ' + f'Please make sure to use an adequate version of {compiler_name} ({version_bound_str}).' + ) + + +# Specify Visual Studio C runtime library for hipcc +def _set_hipcc_runtime_lib(is_standalone, debug) -> None: + if is_standalone: + if debug: + COMMON_HIP_FLAGS.append('-fms-runtime-lib=static_dbg') + else: + COMMON_HIP_FLAGS.append('-fms-runtime-lib=static') + else: + if debug: + COMMON_HIP_FLAGS.append('-fms-runtime-lib=dll_dbg') + else: + COMMON_HIP_FLAGS.append('-fms-runtime-lib=dll') + +def _append_sycl_std_if_no_std_present(cflags) -> None: + if not any(flag.startswith('-sycl-std=') for flag in cflags): + cflags.append('-sycl-std=2020') + + +def _wrap_sycl_host_flags(cflags): + host_cflags = [] + host_cxx = get_cxx_compiler() + if IS_WINDOWS: + for flag in cflags: + if flag.startswith("-I"): + flag = flag.replace("\\", "\\\\").replace("-I", "/I") + else: + flag = flag.replace("-D", "/D") + flag = flag.replace('"', '\\"') + host_cflags.append(flag) + joined_host_cflags = ' '.join(host_cflags) + + external_include = _join_sycl_home("include").replace("\\", "\\\\") + + # Some versions of DPC++ compiler pass paths to SYCL headers as user include paths (`-I`) rather + # than system paths (`-isystem`). This makes host compiler to report warnings encountered in the + # SYCL headers, such as deprecated warnings, even if warmed API is not actually used in the program. + # We expect that this issue will be addressed in the later version of DPC++ compiler. To workaround the + # issue now we wrap paths to SYCL headers in `/external:I`. Warning free compilation is especially important + # for Windows build as `/sdl` compilation flag assumes that and we will fail compilation otherwise. + wrapped_host_cflags = [ + f"-fsycl-host-compiler={host_cxx}", + f'-fsycl-host-compiler-options="\\"/external:I{external_include}\\" /external:W0 {joined_host_cflags}"', + ] + else: + joined_host_cflags = ' '.join(cflags) + wrapped_host_cflags = [ + f"-fsycl-host-compiler={host_cxx}", + shlex.quote(f"-fsycl-host-compiler-options={joined_host_cflags}"), + ] + return wrapped_host_cflags + + +class BuildExtension(build_ext): + """ + A custom :mod:`setuptools` build extension . + + This :class:`setuptools.build_ext` subclass takes care of passing the + minimum required compiler flags (e.g. ``-std=c++20``) as well as mixed + C++/CUDA/SYCL compilation (and support for CUDA/SYCL files in general). + + When using :class:`BuildExtension`, it is allowed to supply a dictionary + for ``extra_compile_args`` (rather than the usual list) that maps from + languages/compilers (the only expected values are ``cxx``, ``nvcc`` or + ``sycl``) to a list of additional compiler flags to supply to the compiler. + This makes it possible to supply different flags to the C++, CUDA and SYCL + compiler during mixed compilation. + + ``use_ninja`` (bool): If ``use_ninja`` is ``True`` (default), then we + attempt to build using the Ninja backend. Ninja greatly speeds up + compilation compared to the standard ``setuptools.build_ext``. + Fallbacks to the standard distutils backend if Ninja is not available. + + .. note:: + By default, the Ninja backend uses #CPUS + 2 workers to build the + extension. This may use up too many resources on some systems. One + can control the number of workers by setting the `MAX_JOBS` environment + variable to a non-negative number. + """ + + @classmethod + def with_options(cls, **options): + """Return a subclass with alternative constructor that extends any original keyword arguments to the original constructor with the given options.""" + class cls_with_options(cls): # type: ignore[misc, valid-type] + def __init__(self, *args, **kwargs) -> None: + kwargs.update(options) + super().__init__(*args, **kwargs) + + return cls_with_options + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.no_python_abi_suffix = kwargs.get("no_python_abi_suffix", False) + + self.use_ninja = kwargs.get('use_ninja', True) + if self.use_ninja: + # Test if we can use ninja. Fallback otherwise. + msg = ('Attempted to use ninja as the BuildExtension backend but ' + '%s. Falling back to using the slow distutils backend.') + if not is_ninja_available(): + logger.warning(msg, 'we could not find ninja.') + self.use_ninja = False + + def finalize_options(self) -> None: + super().finalize_options() + if self.use_ninja: + self.force = True + + def build_extensions(self) -> None: + compiler_name, compiler_version = self._check_abi() + + cuda_ext = False + sycl_ext = False + extension_iter = iter(self.extensions) + extension = next(extension_iter, None) + while not (cuda_ext and sycl_ext) and extension: + for source in extension.sources: + _, ext = os.path.splitext(source) + if ext == '.cu': + cuda_ext = True + elif ext == '.sycl': + sycl_ext = True + + # This check accounts on a case when cuda and sycl sources + # are mixed in the same extension. We can stop checking + # sources if both are found or there is no more sources. + if cuda_ext and sycl_ext: + break + + extension = next(extension_iter, None) + + if sycl_ext: + if not self.use_ninja: + raise AssertionError("ninja is required to build sycl extensions.") + + if cuda_ext and not IS_HIP_EXTENSION: + _check_cuda_version(compiler_name, compiler_version) + + for extension in self.extensions: + # Ensure at least an empty list of flags for 'cxx', 'nvcc' and 'sycl' when + # extra_compile_args is a dict. Otherwise, default torch flags do + # not get passed. Necessary when only one of 'cxx', 'nvcc' or 'sycl' is + # passed to extra_compile_args in CUDAExtension or SyclExtension, i.e. + # CUDAExtension(..., extra_compile_args={'cxx': [...]}) + # or + # CUDAExtension(..., extra_compile_args={'nvcc': [...]}) + if isinstance(extension.extra_compile_args, dict): + for ext in ['cxx', 'nvcc', 'sycl']: + if ext not in extension.extra_compile_args: + extension.extra_compile_args[ext] = [] + + self._add_compile_flag(extension, '-DTORCH_API_INCLUDE_EXTENSION_H') + + if IS_HIP_EXTENSION: + self._hipify_compile_flags(extension) + + if extension.py_limited_api: + # compile any extension that has passed in py_limited_api to the + # Extension constructor with the Py_LIMITED_API flag set to our + # min supported CPython version. + # See https://docs.python.org/3/c-api/stable.html#c.Py_LIMITED_API + self._add_compile_flag(extension, f'-DPy_LIMITED_API={min_supported_cpython}') + self._define_torch_extension_name(extension) + + if 'nvcc_dlink' in extension.extra_compile_args: + if not self.use_ninja: + raise AssertionError( + f"With dlink=True, ninja is required to build cuda extension {extension.name}." + ) + + # Register .cu, .cuh, .hip, .mm and .sycl as valid source extensions. + # NOTE: At the moment .sycl is not a standard extension for SYCL supported + # by compiler. Here we introduce a torch level convention that SYCL sources + # should have .sycl file extension. + self.compiler.src_extensions += ['.cu', '.cuh', '.hip', '.sycl'] + if torch.backends.mps.is_built(): + self.compiler.src_extensions += ['.mm'] + # Save the original _compile method for later. + if self.compiler.compiler_type == 'msvc': + self.compiler._cpp_extensions += ['.cu', '.cuh'] + original_compile = self.compiler.compile + original_spawn = self.compiler.spawn + else: + original_compile = self.compiler._compile + + def append_std17_if_no_std_present(cflags) -> None: + # NVCC does not allow multiple -std to be passed, so we avoid + # overriding the option if the user explicitly passed it. + cpp_format_prefix = '/{}:' if self.compiler.compiler_type == 'msvc' else '-{}=' + cpp_flag_prefix = cpp_format_prefix.format('std') + cpp_flag = cpp_flag_prefix + 'c++20' + if not any(flag.startswith(cpp_flag_prefix) for flag in cflags): + cflags.append(cpp_flag) + + def unix_cuda_flags(cflags): + cflags = (COMMON_NVCC_FLAGS + + ['--compiler-options', "'-fPIC'"] + + cflags + _get_cuda_arch_flags(cflags)) + + # NVCC does not allow multiple -ccbin/--compiler-bindir to be passed, so we avoid + # overriding the option if the user explicitly passed it. + _ccbin = os.getenv("CC") + if ( + _ccbin is not None + and not any(flag.startswith(('-ccbin', '--compiler-bindir')) for flag in cflags) + ): + cflags.extend(['-ccbin', _ccbin]) + + return cflags + + def convert_to_absolute_paths_inplace(paths) -> None: + # Helper function. See Note [Absolute include_dirs] + if paths is not None: + for i in range(len(paths)): + if not os.path.isabs(paths[i]): + paths[i] = os.path.abspath(paths[i]) + + def unix_wrap_single_compile(obj, src, ext, cc_args, extra_postargs, pp_opts) -> None: + # Copy before we make any modifications. + cflags = copy.deepcopy(extra_postargs) + try: + original_compiler = self.compiler.compiler_so + if _is_cuda_file(src): + nvcc = [_join_rocm_home('bin', 'hipcc') if IS_HIP_EXTENSION else _join_cuda_home('bin', 'nvcc')] + self.compiler.set_executable('compiler_so', nvcc) + if isinstance(cflags, dict): + cflags = cflags['nvcc'] + if IS_HIP_EXTENSION: + cflags = COMMON_HIPCC_FLAGS + cflags + _get_rocm_arch_flags(cflags) + else: + cflags = unix_cuda_flags(cflags) + elif isinstance(cflags, dict): + cflags = cflags['cxx'] + if IS_HIP_EXTENSION: + cflags = COMMON_HIP_FLAGS + cflags + append_std17_if_no_std_present(cflags) + + original_compile(obj, src, ext, cc_args, cflags, pp_opts) + finally: + # Put the original compiler back in place. + self.compiler.set_executable('compiler_so', original_compiler) + + def unix_wrap_ninja_compile(sources, + output_dir=None, + macros=None, + include_dirs=None, + debug=0, + extra_preargs=None, + extra_postargs=None, + depends=None): + r"""Compiles sources by outputting a ninja file and running it.""" + # NB: I copied some lines from self.compiler (which is an instance + # of distutils.UnixCCompiler). See the following link. + # https://github.com/python/cpython/blob/f03a8f8d5001963ad5b5b28dbd95497e9cc15596/Lib/distutils/ccompiler.py#L564-L567 # codespell:ignore + # This can be fragile, but a lot of other repos also do this + # (see https://github.com/search?q=_setup_compile&type=Code) + # so it is probably OK; we'll also get CI signal if/when + # we update our python version (which is when distutils can be + # upgraded) + + # Use absolute path for output_dir so that the object file paths + # (`objects`) get generated with absolute paths. + # pyrefly: ignore [no-matching-overload] + output_dir = os.path.abspath(output_dir) + + # See Note [Absolute include_dirs] + convert_to_absolute_paths_inplace(self.compiler.include_dirs) + + _, objects, extra_postargs, pp_opts, _ = \ + self.compiler._setup_compile(output_dir, macros, + include_dirs, sources, + depends, extra_postargs) + common_cflags = self.compiler._get_cc_args(pp_opts, debug, extra_preargs) + extra_cc_cflags = self.compiler.compiler_so[1:] + with_cuda = any(map(_is_cuda_file, sources)) + with_sycl = any(map(_is_sycl_file, sources)) + if with_sycl and with_cuda: + raise AssertionError( + "cannot have both SYCL and CUDA files in the same extension" + ) + + # extra_postargs can be either: + # - a dict mapping cxx/nvcc/sycl to extra flags + # - a list of extra flags. + if isinstance(extra_postargs, dict): + post_cflags = extra_postargs['cxx'] + else: + post_cflags = list(extra_postargs) + if IS_HIP_EXTENSION: + post_cflags = COMMON_HIP_FLAGS + post_cflags + append_std17_if_no_std_present(post_cflags) + + cuda_post_cflags = None + cuda_cflags = None + if with_cuda: + cuda_cflags = common_cflags + if isinstance(extra_postargs, dict): + cuda_post_cflags = extra_postargs['nvcc'] + else: + cuda_post_cflags = list(extra_postargs) + if IS_HIP_EXTENSION: + cuda_post_cflags = cuda_post_cflags + _get_rocm_arch_flags(cuda_post_cflags) + cuda_post_cflags = COMMON_HIP_FLAGS + COMMON_HIPCC_FLAGS + cuda_post_cflags + else: + cuda_post_cflags = unix_cuda_flags(cuda_post_cflags) + append_std17_if_no_std_present(cuda_post_cflags) + cuda_cflags = [shlex.quote(f) for f in cuda_cflags] + cuda_post_cflags = [shlex.quote(f) for f in cuda_post_cflags] + + if isinstance(extra_postargs, dict) and 'nvcc_dlink' in extra_postargs: + cuda_dlink_post_cflags = unix_cuda_flags(extra_postargs['nvcc_dlink']) + cuda_dlink_post_cflags = [shlex.quote(f) for f in cuda_dlink_post_cflags] + else: + cuda_dlink_post_cflags = None + + sycl_post_cflags = None + sycl_cflags = None + sycl_dlink_post_cflags = None + if with_sycl: + sycl_cflags = extra_cc_cflags + common_cflags + _COMMON_SYCL_FLAGS + if isinstance(extra_postargs, dict): + sycl_post_cflags = extra_postargs['sycl'] + else: + sycl_post_cflags = list(extra_postargs) + _append_sycl_targets_if_missing(sycl_post_cflags) + append_std17_if_no_std_present(sycl_cflags) + _append_sycl_std_if_no_std_present(sycl_cflags) + host_cflags = extra_cc_cflags + common_cflags + post_cflags + append_std17_if_no_std_present(host_cflags) + # escaping quoted arguments to pass them thru SYCL compiler + icpx_version = _get_icpx_version() + if int(icpx_version) >= 20250200: + host_cflags = [item.replace('"', '\\"') for item in host_cflags] + else: + host_cflags = [item.replace('"', '\\\\"') for item in host_cflags] + # Note the order: shlex.quote sycl_flags first, _wrap_sycl_host_flags + # second. Reason is that sycl host flags are quoted, space containing + # strings passed to SYCL compiler. + sycl_cflags = [shlex.quote(f) for f in sycl_cflags] + sycl_cflags += _wrap_sycl_host_flags(host_cflags) + sycl_dlink_post_cflags = _SYCL_DLINK_FLAGS.copy() + sycl_dlink_post_cflags += _get_sycl_device_flags(sycl_post_cflags) + sycl_post_cflags = [shlex.quote(f) for f in sycl_post_cflags] + + _write_ninja_file_and_compile_objects( + sources=sources, + objects=objects, + cflags=[shlex.quote(f) for f in extra_cc_cflags + common_cflags], + post_cflags=[shlex.quote(f) for f in post_cflags], + cuda_cflags=cuda_cflags, + cuda_post_cflags=cuda_post_cflags, + cuda_dlink_post_cflags=cuda_dlink_post_cflags, + sycl_cflags=sycl_cflags, + sycl_post_cflags=sycl_post_cflags, + sycl_dlink_post_cflags=sycl_dlink_post_cflags, + build_directory=output_dir, + verbose=True, + with_cuda=with_cuda, + with_sycl=with_sycl) + + # Return *all* object filenames, not just the ones we just built. + return objects + + def win_cuda_flags(cflags): + return (COMMON_NVCC_FLAGS + + cflags + _get_cuda_arch_flags(cflags)) + + def win_hip_flags(cflags): + return (COMMON_HIPCC_FLAGS + COMMON_HIP_FLAGS + cflags + _get_rocm_arch_flags(cflags)) + + def win_filter_msvc_include_dirs(pp_opts) -> list[str]: + """Filter out MSVC include dirs from pp_opts for oneAPI 2025.3+.""" + # oneAPI 2025.3+ changed include path ordering to match MSVC behavior. + # Filter out MSVC headers to avoid conflicting declarations with oneAPI's std headers. + icpx_version = int(_get_icpx_version()) + if icpx_version >= 20250300: + vc_tools_dir = os.path.normcase(os.environ.get('VCToolsInstallDir', '')) + if vc_tools_dir: + pp_opts = [ + path for path in pp_opts + if vc_tools_dir not in os.path.normcase(path) + ] + return pp_opts + + def win_wrap_single_compile(sources, + output_dir=None, + macros=None, + include_dirs=None, + debug=0, + extra_preargs=None, + extra_postargs=None, + depends=None): + + self.cflags = copy.deepcopy(extra_postargs) + extra_postargs = None + + def spawn(cmd): + # Using regex to match src, obj and include files + src_regex = re.compile('/T(p|c)(.*)') + src_list = [ + m.group(2) for m in (src_regex.match(elem) for elem in cmd) + if m + ] + + obj_regex = re.compile('/Fo(.*)') # codespell:ignore + obj_list = [ + m.group(1) for m in (obj_regex.match(elem) for elem in cmd) + if m + ] + + include_regex = re.compile(r'((\-|\/)I.*)') + include_list = [ + m.group(1) + for m in (include_regex.match(elem) for elem in cmd) if m + ] + + if len(src_list) >= 1 and len(obj_list) >= 1: + src = src_list[0] + obj = obj_list[0] + if _is_cuda_file(src): + if IS_HIP_EXTENSION: + nvcc = _get_hipcc_path() + else: + nvcc = _join_cuda_home('bin', 'nvcc') + if isinstance(self.cflags, dict): + cflags = self.cflags['nvcc'] + elif isinstance(self.cflags, list): + cflags = self.cflags + else: + cflags = [] + + if IS_HIP_EXTENSION: + cflags = win_hip_flags(cflags) + else: + cflags = win_cuda_flags(cflags) + ['-std=c++20', '--use-local-env'] + for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS: + cflags = ['-Xcudafe', '--diag_suppress=' + ignore_warning] + cflags + for flag in COMMON_MSVC_FLAGS: + cflags = ['-Xcompiler', flag] + cflags + cmd = [nvcc, '-c', src, '-o', obj] + include_list + cflags + elif isinstance(self.cflags, dict): + cflags = COMMON_MSVC_FLAGS + self.cflags['cxx'] + append_std17_if_no_std_present(cflags) + cmd += cflags + elif isinstance(self.cflags, list): + cflags = COMMON_MSVC_FLAGS + self.cflags + append_std17_if_no_std_present(cflags) + cmd += cflags + + return original_spawn(cmd) + + try: + self.compiler.spawn = spawn + return original_compile(sources, output_dir, macros, + include_dirs, debug, extra_preargs, + extra_postargs, depends) + finally: + self.compiler.spawn = original_spawn + + def win_wrap_ninja_compile(sources, + output_dir=None, + macros=None, + include_dirs=None, + debug=0, + extra_preargs=None, + extra_postargs=None, + depends=None, + is_standalone=False): + if not self.compiler.initialized: + self.compiler.initialize() + # pyrefly: ignore [no-matching-overload] + output_dir = os.path.abspath(output_dir) + + # Note [Absolute include_dirs] + # Convert relative path in self.compiler.include_dirs to absolute path if any. + # For ninja build, the build location is not local, but instead, the build happens + # in a script-created build folder. Thus, relative paths lose their correctness. + # To be consistent with jit extension, we allow user to enter relative include_dirs + # in setuptools.setup, and we convert the relative path to absolute path here. + convert_to_absolute_paths_inplace(self.compiler.include_dirs) + + _, objects, extra_postargs, pp_opts, _ = \ + self.compiler._setup_compile(output_dir, macros, + include_dirs, sources, + depends, extra_postargs) + # Replace space with \ when using hipcc (hipcc passes includes to clang without ""s so clang sees space in include paths as new argument) + if IS_HIP_EXTENSION: + pp_opts = ["-I{}".format(s[2:].replace(" ", "\\")) if s.startswith('-I') else s for s in pp_opts] + common_cflags = extra_preargs or [] + cflags = [] + if debug: + cflags.extend(self.compiler.compile_options_debug) + else: + cflags.extend(self.compiler.compile_options) + cflags = cflags + common_cflags + pp_opts + COMMON_MSVC_FLAGS + if IS_HIP_EXTENSION: + _set_hipcc_runtime_lib(is_standalone, debug) + common_cflags.extend(COMMON_HIP_FLAGS) + else: + common_cflags.extend(COMMON_MSVC_FLAGS) + with_cuda = any(map(_is_cuda_file, sources)) + with_sycl = any(map(_is_sycl_file, sources)) + if with_sycl and with_cuda: + raise AssertionError( + "cannot have both SYCL and CUDA files in the same extension" + ) + + # extra_postargs can be either: + # - a dict mapping cxx/nvcc to extra flags + # - a list of extra flags. + if isinstance(extra_postargs, dict): + post_cflags = extra_postargs['cxx'] + else: + post_cflags = list(extra_postargs) + if IS_HIP_EXTENSION: + post_cflags = COMMON_HIP_FLAGS + post_cflags + append_std17_if_no_std_present(post_cflags) + + cuda_post_cflags = None + cuda_cflags = None + if with_cuda: + cuda_cflags = ['-std=c++20'] + for common_cflag in common_cflags: + cuda_cflags.append('-Xcompiler') + cuda_cflags.append(common_cflag) + if not IS_HIP_EXTENSION: + cuda_cflags.append('--use-local-env') + for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS: + cuda_cflags.append('-Xcudafe') + cuda_cflags.append('--diag_suppress=' + ignore_warning) + cuda_cflags.extend(pp_opts) + if isinstance(extra_postargs, dict): + cuda_post_cflags = extra_postargs['nvcc'] + else: + cuda_post_cflags = list(extra_postargs) + if IS_HIP_EXTENSION: + cuda_post_cflags = win_hip_flags(cuda_post_cflags) + else: + cuda_post_cflags = win_cuda_flags(cuda_post_cflags) + cflags = _nt_quote_args(cflags) + post_cflags = _nt_quote_args(post_cflags) + if with_cuda: + cuda_cflags = _nt_quote_args(cuda_cflags) + cuda_post_cflags = _nt_quote_args(cuda_post_cflags) + if isinstance(extra_postargs, dict) and 'nvcc_dlink' in extra_postargs: + cuda_dlink_post_cflags = win_cuda_flags(extra_postargs['nvcc_dlink']) + else: + cuda_dlink_post_cflags = None + + sycl_cflags = None + sycl_post_cflags = None + sycl_dlink_post_cflags = None + if with_sycl: + sycl_cflags = common_cflags + win_filter_msvc_include_dirs(pp_opts) + _COMMON_SYCL_FLAGS + if isinstance(extra_postargs, dict): + sycl_post_cflags = extra_postargs['sycl'] + else: + sycl_post_cflags = list(extra_postargs) + _append_sycl_targets_if_missing(sycl_post_cflags) + append_std17_if_no_std_present(sycl_cflags) + _append_sycl_std_if_no_std_present(sycl_cflags) + host_cflags = common_cflags + pp_opts + post_cflags + append_std17_if_no_std_present(host_cflags) + + sycl_cflags = _nt_quote_args(sycl_cflags) + host_cflags = _nt_quote_args(host_cflags) + + sycl_cflags += _wrap_sycl_host_flags(host_cflags) + sycl_dlink_post_cflags = _SYCL_DLINK_FLAGS.copy() + sycl_dlink_post_cflags += _get_sycl_device_flags(sycl_post_cflags) + sycl_post_cflags = _nt_quote_args(sycl_post_cflags) + + + _write_ninja_file_and_compile_objects( + sources=sources, + objects=objects, + cflags=cflags, + post_cflags=post_cflags, + cuda_cflags=cuda_cflags, + cuda_post_cflags=cuda_post_cflags, + cuda_dlink_post_cflags=cuda_dlink_post_cflags, + sycl_cflags=sycl_cflags, + sycl_post_cflags=sycl_post_cflags, + sycl_dlink_post_cflags=sycl_dlink_post_cflags, + build_directory=output_dir, + verbose=True, + with_cuda=with_cuda, + with_sycl=with_sycl) + + # Return *all* object filenames, not just the ones we just built. + return objects + # Monkey-patch the _compile or compile method. + # https://github.com/python/cpython/blob/dc0284ee8f7a270b6005467f26d8e5773d76e959/Lib/distutils/ccompiler.py#L511 # codespell:ignore + if self.compiler.compiler_type == 'msvc': + if self.use_ninja: + self.compiler.compile = win_wrap_ninja_compile + else: + self.compiler.compile = win_wrap_single_compile + else: + if self.use_ninja: + self.compiler.compile = unix_wrap_ninja_compile + else: + self.compiler._compile = unix_wrap_single_compile + + build_ext.build_extensions(self) + + def get_ext_filename(self, ext_name): + # Get the original shared library name. For Python 3, this name will be + # suffixed with ".so", where will be something like + # cpython-37m-x86_64-linux-gnu. + ext_filename = super().get_ext_filename(ext_name) + # If `no_python_abi_suffix` is `True`, we omit the Python 3 ABI + # component. This makes building shared libraries with setuptools that + # aren't Python modules nicer. + if self.no_python_abi_suffix: + # The parts will be e.g. ["my_extension", "cpython-37m-x86_64-linux-gnu", "so"]. + ext_filename_parts = ext_filename.split('.') + # Remove ABI component only if it actually exists in a file name, see gh-170542. + if len(ext_filename_parts) > 2: + # Omit the second to last element. + without_abi = ext_filename_parts[:-2] + ext_filename_parts[-1:] + ext_filename = '.'.join(without_abi) + return ext_filename + + def get_export_symbols(self, ext): + if IS_WINDOWS: + # Skips exporting the module "PyInit_" function that the + # distutils Extension.get_export_symbols would add to + # ext.export_symbols. Only relevant for Windows builds. + return ext.export_symbols + return super().get_export_symbols(ext) + + def _check_abi(self) -> tuple[str, TorchVersion]: + # On some platforms, like Windows, compiler_cxx is not available. + if hasattr(self.compiler, 'compiler_cxx'): + compiler = self.compiler.compiler_cxx[0] + else: + compiler = get_cxx_compiler() + _, version = get_compiler_abi_compatibility_and_version(compiler) + # Warn user if VC env is activated but `DISTUILS_USE_SDK` is not set. + if IS_WINDOWS and 'VSCMD_ARG_TGT_ARCH' in os.environ and 'DISTUTILS_USE_SDK' not in os.environ: + msg = ('It seems that the VC environment is activated but DISTUTILS_USE_SDK is not set.' + 'This may lead to multiple activations of the VC env.' + 'Please set `DISTUTILS_USE_SDK=1` and try again.') + raise UserWarning(msg) + return compiler, version + + def _add_compile_flag(self, extension, flag) -> None: + extension.extra_compile_args = copy.deepcopy(extension.extra_compile_args) + if isinstance(extension.extra_compile_args, dict): + for args in extension.extra_compile_args.values(): + args.append(flag) + else: + extension.extra_compile_args.append(flag) + + # Simple hipify, replace the first occurrence of CUDA with HIP + # in flags starting with "-" and containing "CUDA", but exclude -I flags + def _hipify_compile_flags(self, extension) -> None: + if isinstance(extension.extra_compile_args, dict) and 'nvcc' in extension.extra_compile_args: + modified_flags = [] + for flag in extension.extra_compile_args['nvcc']: + if flag.startswith("-") and "CUDA" in flag and not flag.startswith("-I"): + # check/split flag into flag and value + parts = flag.split("=", 1) + if len(parts) == 2: + flag_part, value_part = parts + # replace fist instance of "CUDA" with "HIP" only in the flag and not flag value + modified_flag_part = flag_part.replace("CUDA", "HIP", 1) + modified_flag = f"{modified_flag_part}={value_part}" + else: + # replace fist instance of "CUDA" with "HIP" in flag + modified_flag = flag.replace("CUDA", "HIP", 1) + modified_flags.append(modified_flag) + logger.info('Modified flag: %s -> %s', flag, modified_flag) + else: + modified_flags.append(flag) + extension.extra_compile_args['nvcc'] = modified_flags + + def _define_torch_extension_name(self, extension) -> None: + # pybind11 doesn't support dots in the names + # so in order to support extensions in the packages + # like torch._C, we take the last part of the string + # as the library name + names = extension.name.split('.') + name = names[-1] + define = f'-DTORCH_EXTENSION_NAME={name}' + self._add_compile_flag(extension, define) + + +def CppExtension(name, sources, *args, **kwargs): + """ + Create a :class:`setuptools.Extension` for C++. + + Convenience method that creates a :class:`setuptools.Extension` with the + bare minimum (but often sufficient) arguments to build a C++ extension. + + All arguments are forwarded to the :class:`setuptools.Extension` + constructor. Full list arguments can be found at + https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference + + .. warning:: + The PyTorch python API (as provided in libtorch_python) cannot be built + with the flag ``py_limited_api=True``. When this flag is passed, it is + the user's responsibility in their library to not use APIs from + libtorch_python (in particular pytorch/python bindings) and to only use + APIs from libtorch (aten objects, operators and the dispatcher). For + example, to give access to custom ops from python, the library should + register the ops through the dispatcher. + + Contrary to CPython setuptools, who does not define -DPy_LIMITED_API + as a compile flag when py_limited_api is specified as an option for + the "bdist_wheel" command in ``setup``, PyTorch does! We will specify + -DPy_LIMITED_API=min_supported_cpython to best enforce consistency, + safety, and sanity in order to encourage best practices. To target a + different version, set min_supported_cpython to the hexcode of the + CPython version of choice. + + Example: + >>> # xdoctest: +SKIP + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CPP_EXT) + >>> from setuptools import setup + >>> from torch.utils.cpp_extension import BuildExtension, CppExtension + >>> setup( + ... name='extension', + ... ext_modules=[ + ... CppExtension( + ... name='extension', + ... sources=['extension.cpp'], + ... extra_compile_args=['-g'], + ... extra_link_args=['-Wl,--no-as-needed', '-lm']) + ... ], + ... cmdclass={ + ... 'build_ext': BuildExtension + ... }) + """ + include_dirs = kwargs.get('include_dirs', []) + include_dirs += include_paths() + kwargs['include_dirs'] = include_dirs + + library_dirs = kwargs.get('library_dirs', []) + library_dirs += library_paths() + kwargs['library_dirs'] = library_dirs + + libraries = kwargs.get('libraries', []) + libraries.append('c10') + libraries.append('torch') + libraries.append('torch_cpu') + if not kwargs.get('py_limited_api', False): + # torch_python uses more than the python limited api + libraries.append('torch_python') + if IS_WINDOWS: + libraries.append("sleef") + + kwargs['libraries'] = libraries + + kwargs['language'] = 'c++' + return setuptools.Extension(name, sources, *args, **kwargs) + + +def CUDAExtension(name, sources, *args, **kwargs): + """ + Create a :class:`setuptools.Extension` for CUDA/C++. + + Convenience method that creates a :class:`setuptools.Extension` with the + bare minimum (but often sufficient) arguments to build a CUDA/C++ + extension. This includes the CUDA include path, library path and runtime + library. + + All arguments are forwarded to the :class:`setuptools.Extension` + constructor. Full list arguments can be found at + https://setuptools.pypa.io/en/latest/userguide/ext_modules.html#extension-api-reference + + .. warning:: + The PyTorch python API (as provided in libtorch_python) cannot be built + with the flag ``py_limited_api=True``. When this flag is passed, it is + the user's responsibility in their library to not use APIs from + libtorch_python (in particular pytorch/python bindings) and to only use + APIs from libtorch (aten objects, operators and the dispatcher). For + example, to give access to custom ops from python, the library should + register the ops through the dispatcher. + + Contrary to CPython setuptools, who does not define -DPy_LIMITED_API + as a compile flag when py_limited_api is specified as an option for + the "bdist_wheel" command in ``setup``, PyTorch does! We will specify + -DPy_LIMITED_API=min_supported_cpython to best enforce consistency, + safety, and sanity in order to encourage best practices. To target a + different version, set min_supported_cpython to the hexcode of the + CPython version of choice. + + Example: + >>> # xdoctest: +SKIP + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CPP_EXT) + >>> from setuptools import setup + >>> from torch.utils.cpp_extension import BuildExtension, CUDAExtension + >>> setup( + ... name='cuda_extension', + ... ext_modules=[ + ... CUDAExtension( + ... name='cuda_extension', + ... sources=['extension.cpp', 'extension_kernel.cu'], + ... extra_compile_args={'cxx': ['-g'], + ... 'nvcc': ['-O2']}, + ... extra_link_args=['-Wl,--no-as-needed', '-lcuda']) + ... ], + ... cmdclass={ + ... 'build_ext': BuildExtension + ... }) + + Compute capabilities: + + By default the extension will be compiled to run on all archs of the cards visible during the + building process of the extension, plus PTX. If down the road a new card is installed the + extension may need to be recompiled. If a visible card has a compute capability (CC) that's + newer than the newest version for which your nvcc can build fully-compiled binaries, PyTorch + will make nvcc fall back to building kernels with the newest version of PTX your nvcc does + support (see below for details on PTX). + + You can override the default behavior using `TORCH_CUDA_ARCH_LIST` to explicitly specify which + CCs you want the extension to support: + + ``TORCH_CUDA_ARCH_LIST="6.1 8.6" python build_my_extension.py`` + ``TORCH_CUDA_ARCH_LIST="5.2 6.0 6.1 7.0 7.5 8.0 8.6+PTX" python build_my_extension.py`` + + The +PTX option causes extension kernel binaries to include PTX instructions for the specified + CC. PTX is an intermediate representation that allows kernels to runtime-compile for any CC >= + the specified CC (for example, 8.6+PTX generates PTX that can runtime-compile for any GPU with + CC >= 8.6). This improves your binary's forward compatibility. However, relying on older PTX to + provide forward compat by runtime-compiling for newer CCs can modestly reduce performance on + those newer CCs. If you know exact CC(s) of the GPUs you want to target, you're always better + off specifying them individually. For example, if you want your extension to run on 8.0 and 8.6, + "8.0+PTX" would work functionally because it includes PTX that can runtime-compile for 8.6, but + "8.0 8.6" would be better. + + Note that while it's possible to include all supported archs, the more archs get included the + slower the building process will be, as it will build a separate kernel image for each arch. + + Note that CUDA-11.5 nvcc will hit internal compiler error while parsing torch/extension.h on Windows. + To workaround the issue, move python binding logic to pure C++ file. + + Example use: + #include + at::Tensor SigmoidAlphaBlendForwardCuda(....) + + Instead of: + #include + torch::Tensor SigmoidAlphaBlendForwardCuda(...) + + Currently open issue for nvcc bug: https://github.com/pytorch/pytorch/issues/69460 + Complete workaround code example: https://github.com/facebookresearch/pytorch3d/commit/cb170ac024a949f1f9614ffe6af1c38d972f7d48 + + Relocatable device code linking: + + If you want to reference device symbols across compilation units (across object files), + the object files need to be built with `relocatable device code` (-rdc=true or -dc). + An exception to this rule is "dynamic parallelism" (nested kernel launches) which is not used a lot anymore. + `Relocatable device code` is less optimized so it needs to be used only on object files that need it. + Using `-dlto` (Device Link Time Optimization) at the device code compilation step and `dlink` step + helps reduce the protentional perf degradation of `-rdc`. + Note that it needs to be used at both steps to be useful. + + If you have `rdc` objects you need to have an extra `-dlink` (device linking) step before the CPU symbol linking step. + There is also a case where `-dlink` is used without `-rdc`: + when an extension is linked against a static lib containing rdc-compiled objects + like the [NVSHMEM library](https://developer.nvidia.com/nvshmem). + + Note: Ninja is required to build a CUDA Extension with RDC linking. + + Example: + >>> # xdoctest: +SKIP + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CPP_EXT) + >>> CUDAExtension( + ... name='cuda_extension', + ... sources=['extension.cpp', 'extension_kernel.cu'], + ... dlink=True, + ... dlink_libraries=["dlink_lib"], + ... extra_compile_args={'cxx': ['-g'], + ... 'nvcc': ['-O2', '-rdc=true']}) + """ + library_dirs = kwargs.get('library_dirs', []) + library_dirs += library_paths(device_type="cuda") + kwargs['library_dirs'] = library_dirs + + libraries = kwargs.get('libraries', []) + libraries.append('c10') + libraries.append('torch') + libraries.append('torch_cpu') + if not kwargs.get('py_limited_api', False): + # torch_python uses more than the python limited api + libraries.append('torch_python') + if IS_HIP_EXTENSION: + libraries.append('amdhip64') + libraries.append('c10_hip') + libraries.append('torch_hip') + else: + libraries.append('cudart') + libraries.append('c10_cuda') + libraries.append('torch_cuda') + kwargs['libraries'] = libraries + + include_dirs = kwargs.get('include_dirs', []) + + if IS_HIP_EXTENSION: + from .hipify import hipify_python + build_dir = os.getcwd() + hipify_result = hipify_python.hipify( + project_directory=build_dir, + output_directory=build_dir, + header_include_dirs=include_dirs, + includes=[os.path.join(build_dir, '*')], # limit scope to build_dir only + extra_files=[os.path.abspath(s) for s in sources], + show_detailed=True, + is_pytorch_extension=True, + hipify_extra_files_only=True, # don't hipify everything in includes path + ) + + hipified_sources = set() + for source in sources: + s_abs = os.path.abspath(source) + hipified_s_abs = (hipify_result[s_abs].hipified_path if (s_abs in hipify_result and + hipify_result[s_abs].hipified_path is not None) else s_abs) + # setup() arguments must *always* be /-separated paths relative to the setup.py directory, + # *never* absolute paths + hipified_sources.add(os.path.relpath(hipified_s_abs, build_dir)) + + sources = list(hipified_sources) + + include_dirs += include_paths(device_type="cuda") + kwargs['include_dirs'] = include_dirs + + kwargs['language'] = 'c++' + + dlink_libraries = kwargs.get('dlink_libraries', []) + dlink = kwargs.get('dlink', False) or dlink_libraries + if dlink: + extra_compile_args = kwargs.get('extra_compile_args', {}) + + extra_compile_args_dlink = extra_compile_args.get('nvcc_dlink', []) + extra_compile_args_dlink += ['-dlink'] + extra_compile_args_dlink += [f'-L{x}' for x in library_dirs] + extra_compile_args_dlink += [f'-l{x}' for x in dlink_libraries] + + if (torch.version.cuda is not None) and TorchVersion(torch.version.cuda) >= '11.2': + extra_compile_args_dlink += ['-dlto'] # Device Link Time Optimization started from cuda 11.2 + + extra_compile_args['nvcc_dlink'] = extra_compile_args_dlink + + kwargs['extra_compile_args'] = extra_compile_args + + return setuptools.Extension(name, sources, *args, **kwargs) + + +def SyclExtension(name, sources, *args, **kwargs): + r""" + Creates a :class:`setuptools.Extension` for SYCL/C++. + + Convenience method that creates a :class:`setuptools.Extension` with the + bare minimum (but often sufficient) arguments to build a SYCL/C++ + extension. + + All arguments are forwarded to the :class:`setuptools.Extension` + constructor. + + .. warning:: + The PyTorch python API (as provided in libtorch_python) cannot be built + with the flag ``py_limited_api=True``. When this flag is passed, it is + the user's responsibility in their library to not use APIs from + libtorch_python (in particular pytorch/python bindings) and to only use + APIs from libtorch (aten objects, operators and the dispatcher). For + example, to give access to custom ops from python, the library should + register the ops through the dispatcher. + + Contrary to CPython setuptools, who does not define -DPy_LIMITED_API + as a compile flag when py_limited_api is specified as an option for + the "bdist_wheel" command in ``setup``, PyTorch does! We will specify + -DPy_LIMITED_API=min_supported_cpython to best enforce consistency, + safety, and sanity in order to encourage best practices. To target a + different version, set min_supported_cpython to the hexcode of the + CPython version of choice. + + Example: + >>> # xdoctest: +SKIP + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CPP_EXT) + >>> from torch.utils.cpp_extension import BuildExtension, SyclExtension + >>> setup( + ... name='xpu_extension', + ... ext_modules=[ + ... SyclExtension( + ... name='xpu_extension', + ... sources=['extension.cpp', 'extension_kernel.cpp'], + ... extra_compile_args={'cxx': ['-g', '-std=c++20', '-fPIC']}) + ... ], + ... cmdclass={ + ... 'build_ext': BuildExtension + ... }) + + By default the extension will be compiled to run on all archs of the cards visible during the + building process of the extension. If down the road a new card is installed the + extension may need to be recompiled. You can override the default behavior using + `TORCH_XPU_ARCH_LIST` to explicitly specify which device architectures you want the extension + to support: + + ``TORCH_XPU_ARCH_LIST="pvc,xe-lpg" python build_my_extension.py`` + + Note that while it's possible to include all supported archs, the more archs get included the + slower the building process will be, as it will build a separate kernel image for each arch. + + Note: Ninja is required to build SyclExtension. + """ + library_dirs = kwargs.get("library_dirs", []) + library_dirs += library_paths() + kwargs["library_dirs"] = library_dirs + + libraries = kwargs.get("libraries", []) + libraries.append("c10") + libraries.append("c10_xpu") + libraries.append("torch") + libraries.append("torch_cpu") + libraries.append("sycl") + if not kwargs.get('py_limited_api', False): + # torch_python uses more than the python limited api + libraries.append("torch_python") + libraries.append("torch_xpu") + kwargs["libraries"] = libraries + + include_dirs = kwargs.get("include_dirs", []) + include_dirs += include_paths(device_type="xpu") + kwargs["include_dirs"] = include_dirs + + kwargs["language"] = "c++" + + return setuptools.Extension(name, sources, *args, **kwargs) + +def include_paths(device_type: str = "cpu", torch_include_dirs=True) -> list[str]: + """ + Get the include paths required to build a C++ or CUDA or SYCL extension. + + Args: + device_type: Defaults to "cpu". + Returns: + A list of include path strings. + """ + paths = [] + lib_include = os.path.join(_TORCH_PATH, 'include') + if torch_include_dirs: + paths.extend([ + lib_include, + # Remove this once torch/torch.h is officially no longer supported for C++ extensions. + os.path.join(lib_include, 'torch', 'csrc', 'api', 'include'), + ]) + if device_type == "cuda" and IS_HIP_EXTENSION: + paths.append(os.path.join(lib_include, 'THH')) + paths.append(_join_rocm_home('include')) + elif device_type == "cuda": + cuda_home_include = _join_cuda_home('include') + # if we have the Debian/Ubuntu packages for cuda, we get /usr as cuda home. + # but gcc doesn't like having /usr/include passed explicitly + if cuda_home_include != '/usr/include': + paths.append(cuda_home_include) + + # Support CUDA_INC_PATH env variable supported by CMake files + if (cuda_inc_path := os.environ.get("CUDA_INC_PATH", None)) and \ + cuda_inc_path != '/usr/include': + + paths.append(cuda_inc_path) + if CUDNN_HOME is not None: + paths.append(os.path.join(CUDNN_HOME, 'include')) + elif device_type == "xpu": + paths.append(_join_sycl_home('include')) + paths.append(_join_sycl_home('include', 'sycl')) + return paths + + +def library_paths(device_type: str = "cpu", torch_include_dirs: bool = True, cross_target_platform: str | None = None) -> list[str]: + """ + Get the library paths required to build a C++ or CUDA extension. + + Args: + device_type: Defaults to "cpu". + + Returns: + A list of library path strings. + """ + + paths = [] + + if torch_include_dirs: + # We need to link against libtorch.so + paths.extend([TORCH_LIB_PATH]) + + if device_type == "cuda" and IS_HIP_EXTENSION: + lib_dir = 'lib' + paths.append(_join_rocm_home(lib_dir)) + if HIP_HOME is not None: + paths.append(os.path.join(HIP_HOME, 'lib')) + elif device_type == "cuda": + if cross_target_platform == "windows": + lib_dir = os.path.join('lib', 'x64') + if WINDOWS_CUDA_HOME is None: + raise RuntimeError("Need to set WINDOWS_CUDA_HOME for windows cross-compilation") + paths.append(os.path.join(WINDOWS_CUDA_HOME, lib_dir)) + else: + if IS_WINDOWS: + lib_dir = os.path.join('lib', 'x64') + else: + lib_dir = 'lib64' + if (not os.path.exists(_join_cuda_home(lib_dir)) and + os.path.exists(_join_cuda_home('lib'))): + # 64-bit CUDA may be installed in 'lib' (see e.g. gh-16955) + # Note that it's also possible both don't exist (see + # _find_cuda_home) - in that case we stay with 'lib64'. + lib_dir = 'lib' + + paths.append(_join_cuda_home(lib_dir)) + if CUDNN_HOME is not None: + paths.append(os.path.join(CUDNN_HOME, lib_dir)) + elif device_type == "xpu": + if IS_WINDOWS: + lib_dir = os.path.join('lib', 'x64') + else: + lib_dir = 'lib64' + if (not os.path.exists(_join_sycl_home(lib_dir)) and + os.path.exists(_join_sycl_home('lib'))): + lib_dir = 'lib' + + paths.append(_join_sycl_home(lib_dir)) + + return paths + + +def load(name, + sources: str | list[str], + extra_cflags=None, + extra_cuda_cflags=None, + extra_sycl_cflags=None, + extra_ldflags=None, + extra_include_paths=None, + build_directory=None, + verbose=False, + with_cuda: bool | None = None, + with_sycl: bool | None = None, + is_python_module=True, + is_standalone=False, + keep_intermediates=True): + """ + Load a PyTorch C++ extension just-in-time (JIT). + + To load an extension, a Ninja build file is emitted, which is used to + compile the given sources into a dynamic library. This library is + subsequently loaded into the current Python process as a module and + returned from this function, ready for use. + + By default, the directory to which the build file is emitted and the + resulting library compiled to is ``/torch_extensions/``, where + ```` is the temporary folder on the current platform and ```` + the name of the extension. This location can be overridden in two ways. + First, if the ``TORCH_EXTENSIONS_DIR`` environment variable is set, it + replaces ``/torch_extensions`` and all extensions will be compiled + into subfolders of this directory. Second, if the ``build_directory`` + argument to this function is supplied, it overrides the entire path, i.e. + the library will be compiled into that folder directly. + + To compile the sources, the default system compiler (``c++``) is used, + which can be overridden by setting the ``CXX`` environment variable. To pass + additional arguments to the compilation process, ``extra_cflags`` or + ``extra_ldflags`` can be provided. For example, to compile your extension + with optimizations, pass ``extra_cflags=['-O3']``. You can also use + ``extra_cflags`` to pass further include directories. + + CUDA support with mixed compilation is provided. Simply pass CUDA source + files (``.cu`` or ``.cuh``) along with other sources. Such files will be + detected and compiled with nvcc rather than the C++ compiler. This includes + passing the CUDA lib64 directory as a library directory, and linking + ``cudart``. You can pass additional flags to nvcc via + ``extra_cuda_cflags``, just like with ``extra_cflags`` for C++. Various + heuristics for finding the CUDA install directory are used, which usually + work fine. If not, setting the ``CUDA_HOME`` environment variable is the + safest option. + + SYCL support with mixed compilation is provided. Simply pass SYCL source + files (``.sycl``) along with other sources. Such files will be detected + and compiled with SYCL compiler (such as Intel DPC++ Compiler) rather + than the C++ compiler. You can pass additional flags to SYCL compiler + via ``extra_sycl_cflags``, just like with ``extra_cflags`` for C++. + SYCL compiler is expected to be found via system PATH environment + variable. + + Args: + name: The name of the extension to build. This MUST be the same as the + name of the pybind11 module! + sources: A list of relative or absolute paths to C++ source files. + extra_cflags: optional list of compiler flags to forward to the build. + extra_cuda_cflags: optional list of compiler flags to forward to nvcc + when building CUDA sources. + extra_sycl_cflags: optional list of compiler flags to forward to SYCL + compiler when building SYCL sources. + extra_ldflags: optional list of linker flags to forward to the build. + extra_include_paths: optional list of include directories to forward + to the build. + build_directory: optional path to use as build workspace. + verbose: If ``True``, turns on verbose logging of load steps. + with_cuda: Determines whether CUDA headers and libraries are added to + the build. If set to ``None`` (default), this value is + automatically determined based on the existence of ``.cu`` or + ``.cuh`` in ``sources``. Set it to `True`` to force CUDA headers + and libraries to be included. + with_sycl: Determines whether SYCL headers and libraries are added to + the build. If set to ``None`` (default), this value is + automatically determined based on the existence of ``.sycl`` in + ``sources``. Set it to `True`` to force SYCL headers and + libraries to be included. + is_python_module: If ``True`` (default), imports the produced shared + library as a Python module. If ``False``, behavior depends on + ``is_standalone``. + is_standalone: If ``False`` (default) loads the constructed extension + into the process as a plain dynamic library. If ``True``, build a + standalone executable. + + Returns: + If ``is_python_module`` is ``True``: + Returns the loaded PyTorch extension as a Python module. + + If ``is_python_module`` is ``False`` and ``is_standalone`` is ``False``: + Returns nothing. (The shared library is loaded into the process as + a side effect.) + + If ``is_standalone`` is ``True``. + Return the path to the executable. (On Windows, TORCH_LIB_PATH is + added to the PATH environment variable as a side effect.) + + Example: + >>> # xdoctest: +SKIP + >>> from torch.utils.cpp_extension import load + >>> module = load( + ... name='extension', + ... sources=['extension.cpp', 'extension_kernel.cu'], + ... extra_cflags=['-O2'], + ... verbose=True) + """ + return _jit_compile( + name, + [sources] if isinstance(sources, str) else sources, + extra_cflags, + extra_cuda_cflags, + extra_sycl_cflags, + extra_ldflags, + extra_include_paths, + build_directory or _get_build_directory(name, verbose), + verbose, + with_cuda, + with_sycl, + is_python_module, + is_standalone, + keep_intermediates=keep_intermediates) + +@deprecated("PyBind11 ABI handling is internal to PyBind11; this will be removed after PyTorch 2.9.0") +def _get_pybind11_abi_build_flags() -> list[str]: + return [] + +def check_compiler_is_gcc(compiler) -> bool: + if not IS_LINUX: + return False + + env = os.environ.copy() + env['LC_ALL'] = 'C' # Don't localize output + try: + version_string = subprocess.check_output([compiler, '-v'], stderr=subprocess.STDOUT, env=env).decode(*SUBPROCESS_DECODE_ARGS) + except (subprocess.CalledProcessError, OSError): + try: + version_string = subprocess.check_output([compiler, '--version'], stderr=subprocess.STDOUT, env=env).decode(*SUBPROCESS_DECODE_ARGS) + except (subprocess.CalledProcessError, OSError): + return False + # Check for GCC by verifying both COLLECT_GCC and gcc version string are present + # This works for c++, g++, gcc, and versioned variants like g++-13 + pattern = re.compile("^COLLECT_GCC=(.*)$", re.MULTILINE) + has_collect_gcc = pattern.search(version_string) is not None + if has_collect_gcc and 'gcc version' in version_string: + return True + return False + +def _check_and_build_extension_h_precompiler_headers( + extra_cflags, + extra_include_paths, + is_standalone=False) -> None: + r''' + Precompiled Headers(PCH) can pre-build the same headers and reduce build time for pytorch load_inline modules. + GCC official manual: https://gcc.gnu.org/onlinedocs/gcc-4.0.4/gcc/Precompiled-Headers.html + PCH only works when built pch file(header.h.gch) and build target have the same build parameters. So, We need + add a signature file to record PCH file parameters. If the build parameters(signature) changed, it should rebuild + PCH file. + + Note: + 1. Windows and MacOS have different PCH mechanism. We only support Linux currently. + 2. It only works on GCC/G++. + ''' + if not IS_LINUX: + return + + compiler = get_cxx_compiler() + + b_is_gcc = check_compiler_is_gcc(compiler) + if b_is_gcc is False: + return + + head_file = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h') + head_file_pch = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h.gch') + head_file_signature = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h.sign') + + def listToString(s): + # initialize an empty string + string = "" + if s is None: + return string + + # traverse in the string + for element in s: + string += (element + ' ') + # return string + return string + + def format_precompiler_header_cmd(compiler, head_file, head_file_pch, common_cflags, torch_include_dirs, extra_cflags, extra_include_paths): + return re.sub( + r"[ \n]+", + " ", + f""" + {compiler} -x c++-header {head_file} -o {head_file_pch} {torch_include_dirs} {extra_include_paths} {extra_cflags} {common_cflags} + """, + ).strip() + + def command_to_signature(cmd): + signature = cmd.replace(' ', '_') + return signature + + def check_pch_signature_in_file(file_path, signature): + b_exist = os.path.isfile(file_path) + if b_exist is False: + return False + + with open(file_path) as file: + # read all content of a file + content = file.read() + # check if string present in a file + return signature == content + + def _create_if_not_exist(path_dir) -> None: + if not os.path.exists(path_dir): + try: + Path(path_dir).mkdir(parents=True, exist_ok=True) + except OSError as exc: # Guard against race condition + if exc.errno != errno.EEXIST: + raise RuntimeError(f"Fail to create path {path_dir}") from exc + + def write_pch_signature_to_file(file_path, pch_sign) -> None: + _create_if_not_exist(os.path.dirname(file_path)) + with open(file_path, "w") as f: + f.write(pch_sign) + f.close() + + def build_precompile_header(pch_cmd) -> None: + try: + subprocess.check_output(shlex.split(pch_cmd), stderr=subprocess.STDOUT) + except subprocess.CalledProcessError as e: + raise RuntimeError(f"Compile PreCompile Header fail, command: {pch_cmd}") from e + + extra_cflags_str = listToString(extra_cflags) + extra_include_paths_str = " ".join( + [f"-I{include}" for include in extra_include_paths] if extra_include_paths else [] + ) + + lib_include = os.path.join(_TORCH_PATH, 'include') + torch_include_dirs = [ + f"-I {lib_include}", + # Python.h + "-I {}".format(sysconfig.get_path("include")), + # torch/all.h + "-I {}".format(os.path.join(lib_include, 'torch', 'csrc', 'api', 'include')), + ] + + torch_include_dirs_str = listToString(torch_include_dirs) + + common_cflags = [] + if not is_standalone: + common_cflags += ['-DTORCH_API_INCLUDE_EXTENSION_H'] + + common_cflags += ['-std=c++20', '-fPIC'] + common_cflags_str = listToString(common_cflags) + + pch_cmd = format_precompiler_header_cmd(compiler, head_file, head_file_pch, common_cflags_str, torch_include_dirs_str, extra_cflags_str, extra_include_paths_str) + pch_sign = command_to_signature(pch_cmd) + + if os.path.isfile(head_file_pch) is not True: + build_precompile_header(pch_cmd) + write_pch_signature_to_file(head_file_signature, pch_sign) + else: + b_same_sign = check_pch_signature_in_file(head_file_signature, pch_sign) + if b_same_sign is False: + build_precompile_header(pch_cmd) + write_pch_signature_to_file(head_file_signature, pch_sign) + +def remove_extension_h_precompiler_headers() -> None: + def _remove_if_file_exists(path_file) -> None: + if os.path.exists(path_file): + os.remove(path_file) + + head_file_pch = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h.gch') + head_file_signature = os.path.join(_TORCH_PATH, 'include', 'torch', 'extension.h.sign') + + _remove_if_file_exists(head_file_pch) + _remove_if_file_exists(head_file_signature) + +def load_inline(name, + cpp_sources, + cuda_sources=None, + sycl_sources=None, + functions=None, + extra_cflags=None, + extra_cuda_cflags=None, + extra_sycl_cflags=None, + extra_ldflags=None, + extra_include_paths=None, + build_directory=None, + verbose=False, + with_cuda=None, + with_sycl=None, + is_python_module=True, + with_pytorch_error_handling=True, + keep_intermediates=True, + use_pch=False, + no_implicit_headers=False): + r''' + Load a PyTorch C++ extension just-in-time (JIT) from string sources. + + This function behaves exactly like :func:`load`, but takes its sources as + strings rather than filenames. These strings are stored to files in the + build directory, after which the behavior of :func:`load_inline` is + identical to :func:`load`. + + See `the + tests `_ + for good examples of using this function. + + Sources may omit two required parts of a typical non-inline C++ extension: + the necessary header includes, as well as the (pybind11) binding code. More + precisely, strings passed to ``cpp_sources`` are first concatenated into a + single ``.cpp`` file. This file is then prepended with ``#include + `` + + Furthermore, if the ``functions`` argument is supplied, bindings will be + automatically generated for each function specified. ``functions`` can + either be a list of function names, or a dictionary mapping from function + names to docstrings. If a list is given, the name of each function is used + as its docstring. + + The sources in ``cuda_sources`` are concatenated into a separate ``.cu`` + file and prepended with ``torch/types.h``, ``cuda.h`` and + ``cuda_runtime.h`` includes. The ``.cpp`` and ``.cu`` files are compiled + separately, but ultimately linked into a single library. Note that no + bindings are generated for functions in ``cuda_sources`` per se. To bind + to a CUDA kernel, you must create a C++ function that calls it, and either + declare or define this C++ function in one of the ``cpp_sources`` (and + include its name in ``functions``). + + The sources in ``sycl_sources`` are concatenated into a separate ``.sycl`` + file and prepended with ``torch/types.h``, ``sycl/sycl.hpp`` includes. + The ``.cpp`` and ``.sycl`` files are compiled separately, but ultimately + linked into a single library. Note that no bindings are generated for + functions in ``sycl_sources`` per se. To bind to a SYCL kernel, you must + create a C++ function that calls it, and either declare or define this + C++ function in one of the ``cpp_sources`` (and include its name + in ``functions``). + + + + See :func:`load` for a description of arguments omitted below. + + Args: + cpp_sources: A string, or list of strings, containing C++ source code. + cuda_sources: A string, or list of strings, containing CUDA source code. + sycl_sources: A string, or list of strings, containing SYCL source code. + functions: A list of function names for which to generate function + bindings. If a dictionary is given, it should map function names to + docstrings (which are otherwise just the function names). + with_cuda: Determines whether CUDA headers and libraries are added to + the build. If set to ``None`` (default), this value is + automatically determined based on whether ``cuda_sources`` is + provided. Set it to ``True`` to force CUDA headers + and libraries to be included. + with_sycl: Determines whether SYCL headers and libraries are added to + the build. If set to ``None`` (default), this value is + automatically determined based on whether ``sycl_sources`` is + provided. Set it to ``True`` to force SYCL headers + and libraries to be included. + with_pytorch_error_handling: Determines whether pytorch error and + warning macros are handled by pytorch instead of pybind. To do + this, each function ``foo`` is called via an intermediary ``_safe_foo`` + function. This redirection might cause issues in obscure cases + of cpp. This flag should be set to ``False`` when this redirect + causes issues. + no_implicit_headers: If ``True``, skips automatically adding headers, most notably + ``#include `` and ``#include `` lines. + Use this option to improve cold start times when you + already include the necessary headers in your source code. Default: ``False``. + + Example: + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CPP_EXT) + >>> from torch.utils.cpp_extension import load_inline + >>> source = """ + at::Tensor sin_add(at::Tensor x, at::Tensor y) { + return x.sin() + y.sin(); + } + """ + >>> module = load_inline(name='inline_extension', + ... cpp_sources=[source], + ... functions=['sin_add']) + + .. note:: + Since load_inline will just-in-time compile the source code, please ensure + that you have the right toolchains installed in the runtime. For example, + when loading C++, make sure a C++ compiler is available. If you're loading + a CUDA extension, you will need to additionally install the corresponding CUDA + toolkit (nvcc and any other dependencies your code has). Compiling toolchains + are not included when you install torch and must be additionally installed. + + During compiling, by default, the Ninja backend uses #CPUS + 2 workers to build + the extension. This may use up too many resources on some systems. One + can control the number of workers by setting the `MAX_JOBS` environment + variable to a non-negative number. + ''' + build_directory = build_directory or _get_build_directory(name, verbose) + + if isinstance(cpp_sources, str): + cpp_sources = [cpp_sources] + cuda_sources = cuda_sources or [] + if isinstance(cuda_sources, str): + cuda_sources = [cuda_sources] + sycl_sources = sycl_sources or [] + if isinstance(sycl_sources, str): + sycl_sources = [sycl_sources] + + if not no_implicit_headers: + cpp_sources.insert(0, '#include ') + + if use_pch is True: + # Using PreCompile Header('torch/extension.h') to reduce compile time. + _check_and_build_extension_h_precompiler_headers(extra_cflags, extra_include_paths) + else: + remove_extension_h_precompiler_headers() + + # If `functions` is supplied, we create the pybind11 bindings for the user. + # Here, `functions` is (or becomes, after some processing) a map from + # function names to function docstrings. + if functions is not None: + module_def = [] + module_def.append('PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {') + if isinstance(functions, str): + functions = [functions] + if isinstance(functions, list): + # Make the function docstring the same as the function name. + functions = {f: f for f in functions} + elif not isinstance(functions, dict): + raise ValueError(f"Expected 'functions' to be a list or dict, but was {type(functions)}") + for function_name, docstring in functions.items(): + if with_pytorch_error_handling: + module_def.append(f'm.def("{function_name}", torch::wrap_pybind_function({function_name}), "{docstring}");') + else: + module_def.append(f'm.def("{function_name}", {function_name}, "{docstring}");') + module_def.append('}') + cpp_sources += module_def + + cpp_source_path = os.path.join(build_directory, 'main.cpp') + _maybe_write(cpp_source_path, "\n".join(cpp_sources)) + + sources = [cpp_source_path] + + if cuda_sources: + if not no_implicit_headers: + cuda_sources.insert(0, '#include ') + cuda_sources.insert(1, '#include ') + cuda_sources.insert(2, '#include ') + + cuda_source_path = os.path.join(build_directory, 'cuda.cu') + _maybe_write(cuda_source_path, "\n".join(cuda_sources)) + + sources.append(cuda_source_path) + + if sycl_sources: + if not no_implicit_headers: + sycl_sources.insert(0, '#include ') + sycl_sources.insert(1, '#include ') + + sycl_source_path = os.path.join(build_directory, 'sycl.sycl') + _maybe_write(sycl_source_path, "\n".join(sycl_sources)) + + sources.append(sycl_source_path) + + return _jit_compile( + name, + sources, + extra_cflags, + extra_cuda_cflags, + extra_sycl_cflags, + extra_ldflags, + extra_include_paths, + build_directory, + verbose, + with_cuda, + with_sycl, + is_python_module, + is_standalone=False, + keep_intermediates=keep_intermediates) + + +def _jit_compile(name, + sources, + extra_cflags, + extra_cuda_cflags, + extra_sycl_cflags, + extra_ldflags, + extra_include_paths, + build_directory: str, + verbose: bool, + with_cuda: bool | None, + with_sycl: bool | None, + is_python_module, + is_standalone, + keep_intermediates=True) -> types.ModuleType | str: + if is_python_module and is_standalone: + raise ValueError("`is_python_module` and `is_standalone` are mutually exclusive.") + + if with_cuda is None: + with_cuda = any(map(_is_cuda_file, sources)) + with_cudnn = any('cudnn' in f for f in extra_ldflags or []) + if with_sycl is None: + with_sycl = any(map(_is_sycl_file, sources)) + if with_sycl and with_cuda: + raise AssertionError( + "cannot have both SYCL and CUDA files in the same extension" + ) + old_version = JIT_EXTENSION_VERSIONER.get_version(name) + version = JIT_EXTENSION_VERSIONER.bump_version_if_changed( + name, + sources, + build_arguments=[extra_cflags, extra_cuda_cflags, extra_ldflags, extra_include_paths], + build_directory=build_directory, + with_cuda=with_cuda, + with_sycl=with_sycl, + is_python_module=is_python_module, + is_standalone=is_standalone, + ) + if version > 0: + if version != old_version and verbose: + logger.info('The input conditions for extension module %s have changed.', name) + logger.info('Bumping to version %s and re-building as %s_v%s...', version, name, version) + name = f'{name}_v{version}' + + baton = FileBaton(os.path.join(build_directory, 'lock')) + if baton.try_acquire(): + try: + if version != old_version: + if IS_HIP_EXTENSION and (with_cuda or with_cudnn): + from .hipify import hipify_python + from .hipify.hipify_python import GeneratedFileCleaner + clean_ctx_mgr = GeneratedFileCleaner(keep_intermediates=keep_intermediates) + else: + import contextlib + hipify_python = None # type: ignore[assignment] + clean_ctx_mgr = contextlib.nullcontext() + with clean_ctx_mgr as clean_ctx: + if IS_HIP_EXTENSION and (with_cuda or with_cudnn): + assert hipify_python is not None # noqa: S101 + hipify_result = hipify_python.hipify( + project_directory=build_directory, + output_directory=build_directory, + header_include_dirs=(extra_include_paths if extra_include_paths is not None else []), + extra_files=[os.path.abspath(s) for s in sources], + ignores=[_join_rocm_home('*'), os.path.join(_TORCH_PATH, '*')], # no need to hipify ROCm or PyTorch headers + show_detailed=verbose, + show_progress=verbose, + is_pytorch_extension=True, + clean_ctx=clean_ctx + ) + + hipified_sources = set() + for source in sources: + s_abs = os.path.abspath(source) + if s_abs in hipify_result and hipify_result[s_abs].hipified_path is not None: + hipified_s_abs = hipify_result[s_abs].hipified_path + else: + hipified_s_abs = s_abs + hipified_sources.add(hipified_s_abs) + sources = list(hipified_sources) + + _write_ninja_file_and_build_library( + name=name, + sources=sources, + extra_cflags=extra_cflags or [], + extra_cuda_cflags=extra_cuda_cflags or [], + extra_sycl_cflags=extra_sycl_cflags or [], + extra_ldflags=extra_ldflags or [], + extra_include_paths=extra_include_paths or [], + build_directory=build_directory, + verbose=verbose, + with_cuda=with_cuda, + with_sycl=with_sycl, + is_standalone=is_standalone) + elif verbose: + logger.debug('No modifications detected for re-loaded extension module %s, skipping build step...', name) + finally: + baton.release() + else: + baton.wait() + + if verbose: + logger.info('Loading extension module %s...', name) + + if is_standalone: + return _get_exec_path(name, build_directory) + + return _import_module_from_library(name, build_directory, is_python_module) + +def _get_hipcc_path(): + if IS_WINDOWS: + # mypy thinks ROCM_VERSION is None but it will never be None here + hipcc_exe = 'hipcc.exe' if ROCM_VERSION >= (6, 4) else 'hipcc.bat' # type: ignore[operator] + return _join_rocm_home('bin', hipcc_exe) + else: + return _join_rocm_home('bin', 'hipcc') + +def _write_ninja_file_and_compile_objects( + sources: list[str], + objects, + cflags, + post_cflags, + cuda_cflags, + cuda_post_cflags, + cuda_dlink_post_cflags, + sycl_cflags, + sycl_post_cflags, + sycl_dlink_post_cflags, + build_directory: str, + verbose: bool, + with_cuda: bool | None, + with_sycl: bool | None) -> None: + verify_ninja_availability() + + compiler = get_cxx_compiler() + + get_compiler_abi_compatibility_and_version(compiler) + if with_cuda is None: + with_cuda = any(map(_is_cuda_file, sources)) + if with_sycl is None: + with_sycl = any(map(_is_sycl_file, sources)) + if with_sycl and with_cuda: + raise AssertionError( + "cannot have both SYCL and CUDA files in the same extension" + ) + build_file_path = os.path.join(build_directory, 'build.ninja') + if verbose: + logger.debug('Emitting ninja build file %s...', build_file_path) + + # Create build_directory if it does not exist + if not os.path.exists(build_directory): + if verbose: + logger.debug('Creating directory %s...', build_directory) + # This is like mkdir -p, i.e. will also create parent directories. + os.makedirs(build_directory, exist_ok=True) + + _write_ninja_file( + path=build_file_path, + cflags=cflags, + post_cflags=post_cflags, + cuda_cflags=cuda_cflags, + cuda_post_cflags=cuda_post_cflags, + cuda_dlink_post_cflags=cuda_dlink_post_cflags, + sycl_cflags=sycl_cflags, + sycl_post_cflags=sycl_post_cflags, + sycl_dlink_post_cflags=sycl_dlink_post_cflags, + sources=sources, + objects=objects, + ldflags=None, + library_target=None, + with_cuda=with_cuda, + with_sycl=with_sycl) + if verbose: + logger.info('Compiling objects...') + _run_ninja_build( + build_directory, + verbose, + # It would be better if we could tell users the name of the extension + # that failed to build but there isn't a good way to get it here. + error_prefix='Error compiling objects for extension') + + +def _write_ninja_file_and_build_library( + name, + sources: list[str], + extra_cflags, + extra_cuda_cflags, + extra_sycl_cflags, + extra_ldflags, + extra_include_paths, + build_directory: str, + verbose: bool, + with_cuda: bool | None, + with_sycl: bool | None, + is_standalone: bool = False) -> None: + verify_ninja_availability() + + compiler = get_cxx_compiler() + + get_compiler_abi_compatibility_and_version(compiler) + if with_cuda is None: + with_cuda = any(map(_is_cuda_file, sources)) + if with_sycl is None: + with_sycl = any(map(_is_sycl_file, sources)) + if with_sycl and with_cuda: + raise AssertionError( + "cannot have both SYCL and CUDA files in the same extension" + ) + extra_ldflags = _prepare_ldflags( + extra_ldflags or [], + with_cuda, + with_sycl, + verbose, + is_standalone) + build_file_path = os.path.join(build_directory, 'build.ninja') + if verbose: + logger.debug('Emitting ninja build file %s...', build_file_path) + + # Create build_directory if it does not exist + if not os.path.exists(build_directory): + if verbose: + logger.debug('Creating directory %s...', build_directory) + # This is like mkdir -p, i.e. will also create parent directories. + os.makedirs(build_directory, exist_ok=True) + + # NOTE: Emitting a new ninja build file does not cause re-compilation if + # the sources did not change, so it's ok to re-emit (and it's fast). + _write_ninja_file_to_build_library( + path=build_file_path, + name=name, + sources=sources, + extra_cflags=extra_cflags or [], + extra_cuda_cflags=extra_cuda_cflags or [], + extra_sycl_cflags=extra_sycl_cflags or [], + extra_ldflags=extra_ldflags or [], + extra_include_paths=extra_include_paths or [], + with_cuda=with_cuda, + with_sycl=with_sycl, + is_standalone=is_standalone) + + if verbose: + logger.info('Building extension module %s...', name) + _run_ninja_build( + build_directory, + verbose, + error_prefix=f"Error building extension '{name}'") + + +def is_ninja_available() -> bool: + """Return ``True`` if the `ninja `_ build system is available on the system, ``False`` otherwise.""" + try: + subprocess.check_output(['ninja', '--version']) + except Exception: + return False + else: + return True + + +def verify_ninja_availability() -> None: + """Raise ``RuntimeError`` if `ninja `_ build system is not available on the system, does nothing otherwise.""" + if not is_ninja_available(): + raise RuntimeError("Ninja is required to load C++ extensions (pip install ninja to get it)") + + +def _prepare_ldflags(extra_ldflags, with_cuda, with_sycl, verbose, is_standalone): + if IS_WINDOWS: + python_lib_path = os.path.join(sys.base_exec_prefix, 'libs') + + extra_ldflags.append('c10.lib') + if with_cuda: + extra_ldflags.append('c10_hip.lib' if IS_HIP_EXTENSION else 'c10_cuda.lib') + if with_sycl: + extra_ldflags.append('c10_xpu.lib') + extra_ldflags.append('torch_cpu.lib') + if with_cuda: + extra_ldflags.append('torch_hip.lib' if IS_HIP_EXTENSION else 'torch_cuda.lib') + # /INCLUDE is used to ensure torch_cuda is linked against in a project that relies on it. + # Related issue: https://github.com/pytorch/pytorch/issues/31611 + extra_ldflags.append('-INCLUDE:?warp_size@cuda@at@@YAHXZ') + if with_sycl: + extra_ldflags.append('torch_xpu.lib') + extra_ldflags.append('torch.lib') + extra_ldflags.append(f'/LIBPATH:{TORCH_LIB_PATH}') + if not is_standalone: + extra_ldflags.append('torch_python.lib') + extra_ldflags.append(f'/LIBPATH:{python_lib_path}') + + else: + extra_ldflags.append(f'-L{TORCH_LIB_PATH}') + extra_ldflags.append('-lc10') + if with_cuda: + extra_ldflags.append('-lc10_hip' if IS_HIP_EXTENSION else '-lc10_cuda') + if with_sycl: + extra_ldflags.append('-lc10_xpu') + extra_ldflags.append('-ltorch_cpu') + if with_cuda: + extra_ldflags.append('-ltorch_hip' if IS_HIP_EXTENSION else '-ltorch_cuda') + if with_sycl: + extra_ldflags.append('-ltorch_xpu') + extra_ldflags.append('-ltorch') + if not is_standalone: + extra_ldflags.append('-ltorch_python') + + if is_standalone: + extra_ldflags.append(f"-Wl,-rpath,{TORCH_LIB_PATH}") + + if with_cuda: + if verbose: + logger.info('Detected CUDA files, patching ldflags') + if IS_WINDOWS and not IS_HIP_EXTENSION: + extra_ldflags.append(f'/LIBPATH:{_join_cuda_home("lib", "x64")}') + extra_ldflags.append('cudart.lib') + if CUDNN_HOME is not None: + extra_ldflags.append(f'/LIBPATH:{os.path.join(CUDNN_HOME, "lib", "x64")}') + elif not IS_HIP_EXTENSION: + extra_lib_dir = "lib64" + if (not os.path.exists(_join_cuda_home(extra_lib_dir)) and + os.path.exists(_join_cuda_home("lib"))): + # 64-bit CUDA may be installed in "lib" + # Note that it's also possible both don't exist (see _find_cuda_home) - in that case we stay with "lib64" + extra_lib_dir = "lib" + extra_ldflags.append(f'-L{_join_cuda_home(extra_lib_dir)}') + extra_ldflags.append('-lcudart') + if CUDNN_HOME is not None: + extra_ldflags.append(f'-L{os.path.join(CUDNN_HOME, "lib64")}') + elif IS_HIP_EXTENSION: + if IS_WINDOWS: + extra_ldflags.append(f'/LIBPATH:{_join_rocm_home("lib")}') + extra_ldflags.append('amdhip64.lib') + else: + extra_ldflags.append(f'-L{_join_rocm_home("lib")}') + extra_ldflags.append('-lamdhip64') + if with_sycl: + if IS_WINDOWS: + extra_ldflags.append(f'/LIBPATH:{_join_sycl_home("lib")}') + extra_ldflags.append('sycl.lib') + else: + extra_ldflags.append(f'-L{_join_sycl_home("lib")}') + extra_ldflags.append('-lsycl') + return extra_ldflags + + +def _get_cuda_arch_flags(cflags: list[str] | None = None) -> list[str]: + """ + Determine CUDA arch flags to use. + + For an arch, say "6.1", the added compile flag will be + ``-gencode=arch=compute_61,code=sm_61``. + For an added "+PTX", an additional + ``-gencode=arch=compute_xx,code=compute_xx`` is added. + + See select_compute_arch.cmake for corresponding named and supported arches + when building with CMake. + """ + # If cflags is given, there may already be user-provided arch flags in it + # (from `extra_compile_args`) + if cflags is not None: + for flag in cflags: + if 'TORCH_EXTENSION_NAME' in flag: + continue + if 'arch' in flag: + return [] + + # Note: keep combined names ("arch1+arch2") above single names, otherwise + # string replacement may not do the right thing + named_arches = collections.OrderedDict([ + ('Kepler+Tesla', '3.7'), + ('Kepler', '3.5+PTX'), + ('Maxwell+Tegra', '5.3'), + ('Maxwell', '5.0;5.2+PTX'), + ('Pascal', '6.0;6.1+PTX'), + ('Volta+Tegra', '7.2'), + ('Volta', '7.0+PTX'), + ('Turing', '7.5+PTX'), + ('Ampere+Tegra', '8.7'), + ('Ampere', '8.0;8.6+PTX'), + ('Ada', '8.9+PTX'), + ('Hopper', '9.0+PTX'), + ('Blackwell+Tegra', '11.0'), + ('Blackwell', '10.0;10.3;12.0;12.1+PTX'), + ]) + + supported_arches = ['3.5', '3.7', '5.0', '5.2', '5.3', '6.0', '6.1', '6.2', + '7.0', '7.2', '7.5', '8.0', '8.6', '8.7', '8.9', '9.0', '9.0a', + '10.0', '10.0a', '11.0', '11.0a', '10.3', '10.3a', '12.0', + '12.0a', '12.1', '12.1a'] + valid_arch_strings = supported_arches + [s + "+PTX" for s in supported_arches] + + # The default is sm_30 for CUDA 9.x and 10.x + # First check for an env var (same as used by the main setup.py) + # Can be one or more architectures, e.g. "6.1" or "3.5;5.2;6.0;6.1;7.0+PTX" + # See cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake + _arch_list = os.environ.get('TORCH_CUDA_ARCH_LIST', None) + + # If not given or set as native, determine what's best for the GPU / CUDA version that can be found + if not _arch_list or _arch_list == "native": + arch_list = [] + # the assumption is that the extension should run on any of the currently visible cards, + # which could be of different types - therefore all archs for visible cards should be included + for i in range(torch.cuda.device_count()): + capability = torch.cuda.get_device_capability(i) + supported_sm = [int("".join(re.findall(r"\d+", arch.split('_')[1]))) + for arch in torch.cuda.get_arch_list() if 'sm_' in arch] + max_supported_sm = max((sm // 10, sm % 10) for sm in supported_sm) + # Capability of the device may be higher than what's supported by the user's + # NVCC, causing compilation error. User's NVCC is expected to match the one + # used to build pytorch, so we use the maximum supported capability of pytorch + # to clamp the capability. + capability = min(max_supported_sm, capability) + arch = f'{capability[0]}.{capability[1]}' + if arch not in arch_list: + arch_list.append(arch) + arch_list = sorted(arch_list) + arch_list[-1] += '+PTX' + + if not _arch_list: + # Only log on rank 0 in distributed settings to avoid spam + if not torch.distributed.is_available() or not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: + arch_list_str = ';'.join(arch_list) + logger.debug( + "TORCH_CUDA_ARCH_LIST is not set, using TORCH_CUDA_ARCH_LIST='%s' " + "for visible GPU architectures. Set os.environ['TORCH_CUDA_ARCH_LIST'] to override.", + arch_list_str) + else: + # Deal with lists that are ' ' separated (only deal with ';' after) + _arch_list = _arch_list.replace(' ', ';') + # Expand named arches + for named_arch, archival in named_arches.items(): + _arch_list = _arch_list.replace(named_arch, archival) + + arch_list = _arch_list.split(';') + + flags = [] + for arch in arch_list: + if arch not in valid_arch_strings: + raise ValueError(f"Unknown CUDA arch ({arch}) or GPU not supported") + else: + # Handle both single and double-digit architecture versions + version = arch.split('+')[0] # Remove "+PTX" if present + major, minor = version.split('.') + num = f"{major}{minor}" + flags.append(f'-gencode=arch=compute_{num},code=sm_{num}') + if arch.endswith('+PTX'): + flags.append(f'-gencode=arch=compute_{num},code=compute_{num}') + + return sorted(set(flags)) + + +def _get_rocm_arch_flags(cflags: list[str] | None = None) -> list[str]: + # If cflags is given, there may already be user-provided arch flags in it + # (from `extra_compile_args`). If user also specified -fgpu-rdc or -fno-gpu-rdc, we + # assume they know what they're doing. Otherwise, we force -fno-gpu-rdc default. + has_gpu_rdc_flag = False + if cflags is not None: + has_custom_flags = False + for flag in cflags: + if 'amdgpu-target' in flag or 'offload-arch' in flag: + has_custom_flags = True + elif 'gpu-rdc' in flag: + has_gpu_rdc_flag = True + if has_custom_flags: + return [] if has_gpu_rdc_flag else ['-fno-gpu-rdc'] + # Use same defaults as used for building PyTorch + # Allow env var to override, just like during initial cmake build. + _archs = os.environ.get('PYTORCH_ROCM_ARCH', None) + if not _archs: + arch_set = set() + # the assumption is that the extension should run on any of the currently visible cards, + # which could be of different types - therefore all archs for visible cards should be included + for i in range(torch.cuda.device_count()): + device_properties = torch.cuda.get_device_properties(i) + if hasattr(device_properties, "gcnArchName"): + device_arch = (device_properties.gcnArchName).split(":", 1)[0] + arch_set.add(device_arch) + + archs = ";".join(arch_set) + + logger.warning( + "The environment variable `PYTORCH_ROCM_ARCH` is not set, all archs for visible cards are included for compilation (%s).\n" + "If this is not desired, please set the environment variable `PYTORCH_ROCM_ARCH` to specific architectures.", archs) + else: + archs = _archs.replace(' ', ';') + + archs = archs.split(';') + flags = [f'--offload-arch={arch}' for arch in archs] + flags += [] if has_gpu_rdc_flag else ['-fno-gpu-rdc'] + return flags + +def _get_build_directory(name: str, verbose: bool) -> str: + """ + Get the build directory for an extension. + + Args: + name: The name of the extension + verbose: Whether to print verbose information + + Returns: + The path to the build directory + """ + root_extensions_directory = os.environ.get('TORCH_EXTENSIONS_DIR') + if root_extensions_directory is None: + root_extensions_directory = get_default_build_root() + # Determine GPU accelerator prefix based on available accelerators. Fallback to CPU. + # Priority: ROCm/HIP > CUDA > CPU + # Note: torch.backends.cuda.is_built() returns True for both CUDA and ROCm, + # so we need to check torch.version.hip to distinguish them + if torch.version.hip is not None: + accelerator_str = f'rocm{torch.version.hip.replace(".", "")}' + elif torch.version.cuda is not None: + accelerator_str = f'cu{torch.version.cuda.replace(".", "")}' + else: + accelerator_str = 'cpu' + python_version = f'py{sys.version_info.major}{sys.version_info.minor}{getattr(sys, "abiflags", "")}' + build_folder = f'{python_version}_{accelerator_str}' + + root_extensions_directory = os.path.join( + root_extensions_directory, build_folder) + + if verbose: + logger.info('Using %s as PyTorch extensions root...', root_extensions_directory) + + build_directory = os.path.join(root_extensions_directory, name) + if not os.path.exists(build_directory): + if verbose: + logger.debug('Creating extension directory %s...', build_directory) + # This is like mkdir -p, i.e. will also create parent directories. + os.makedirs(build_directory, exist_ok=True) + + return build_directory + + +def _get_num_workers(verbose: bool) -> int | None: + max_jobs = os.environ.get('MAX_JOBS') + if max_jobs is not None and max_jobs.isdigit(): + if verbose: + logger.debug('Using envvar MAX_JOBS (%s) as the number of workers...', max_jobs) + return int(max_jobs) + if verbose: + logger.info( + 'Allowing ninja to set a default number of workers... ' + '(overridable by setting the environment variable MAX_JOBS=N)' + ) + return None + + +def _get_vc_env(vc_arch: str) -> dict[str, str]: + try: + from setuptools import distutils # type: ignore[attr-defined] + + return distutils._msvccompiler._get_vc_env(vc_arch) + except AttributeError: + try: + from setuptools._distutils import _msvccompiler + return _msvccompiler._get_vc_env(vc_arch) # type: ignore[attr-defined] + except AttributeError: + from setuptools._distutils.compilers.C import msvc + return msvc._get_vc_env(vc_arch) # type: ignore[attr-defined] + +def _run_ninja_build(build_directory: str, verbose: bool, error_prefix: str) -> None: + command = ['ninja', '-v'] + num_workers = _get_num_workers(verbose) + if num_workers is not None: + command.extend(['-j', str(num_workers)]) + env = os.environ.copy() + # Try to activate the vc env for the users + if IS_WINDOWS and 'VSCMD_ARG_TGT_ARCH' not in env: + from setuptools import distutils # type: ignore[attr-defined] + + plat_name = distutils.util.get_platform() + plat_spec = PLAT_TO_VCVARS[plat_name] + vc_env = {k.upper(): v for k, v in _get_vc_env(plat_spec).items()} + for k, v in env.items(): + uk = k.upper() + if uk not in vc_env: + vc_env[uk] = v + env = vc_env + try: + sys.stdout.flush() + sys.stderr.flush() + # Warning: don't pass stdout=None to subprocess.run to get output. + # subprocess.run assumes that sys.__stdout__ has not been modified and + # attempts to write to it by default. However, when we call _run_ninja_build + # from ahead-of-time cpp extensions, the following happens: + # 1) If the stdout encoding is not utf-8, setuptools detaches __stdout__. + # https://github.com/pypa/setuptools/blob/7e97def47723303fafabe48b22168bbc11bb4821/setuptools/dist.py#L1110 + # (it probably shouldn't do this) + # 2) subprocess.run (on POSIX, with no stdout override) relies on + # __stdout__ not being detached: + # https://github.com/python/cpython/blob/c352e6c7446c894b13643f538db312092b351789/Lib/subprocess.py#L1214 + # To work around this, we pass in the fileno directly and hope that + # it is valid. + stdout_fileno = 1 + subprocess.run( + command, + shell=IS_WINDOWS and IS_HIP_EXTENSION, + stdout=stdout_fileno if verbose else subprocess.PIPE, + stderr=subprocess.STDOUT, + cwd=build_directory, + check=True, + env=env) + except subprocess.CalledProcessError as e: + # Python 2 and 3 compatible way of getting the error object. + _, error, _ = sys.exc_info() + # error.output contains the stdout and stderr of the build attempt. + message = error_prefix + # `error` is a CalledProcessError (which has an `output`) attribute, but + # mypy thinks it's Optional[BaseException] and doesn't narrow + if hasattr(error, 'output') and error.output: # type: ignore[union-attr] + message += f": {error.output.decode(*SUBPROCESS_DECODE_ARGS)}" # type: ignore[union-attr] + raise RuntimeError(message) from e + + +def _get_exec_path(module_name, path): + if IS_WINDOWS and TORCH_LIB_PATH not in os.getenv('PATH', '').split(';'): + torch_lib_in_path = any( + os.path.exists(p) and os.path.samefile(p, TORCH_LIB_PATH) + for p in os.getenv('PATH', '').split(';') + ) + if not torch_lib_in_path: + os.environ['PATH'] = f"{TORCH_LIB_PATH};{os.getenv('PATH', '')}" + return os.path.join(path, f'{module_name}{EXEC_EXT}') + + +def _import_module_from_library(module_name, path, is_python_module): + filepath = os.path.join(path, f"{module_name}{LIB_EXT}") + if is_python_module: + # https://stackoverflow.com/questions/67631/how-to-import-a-module-given-the-full-path + spec = importlib.util.spec_from_file_location(module_name, filepath) + if spec is None: + raise AssertionError(f"Failed to create spec for module {module_name} at {filepath}") + module = importlib.util.module_from_spec(spec) + if not isinstance(spec.loader, importlib.abc.Loader): + raise AssertionError("spec.loader is not a valid importlib Loader") + spec.loader.exec_module(module) + return module + else: + torch.ops.load_library(filepath) + return filepath + + +def _write_ninja_file_to_build_library(path, + name, + sources, + extra_cflags, + extra_cuda_cflags, + extra_sycl_cflags, + extra_ldflags, + extra_include_paths, + with_cuda, + with_sycl, + is_standalone) -> None: + extra_cflags = [flag.strip() for flag in extra_cflags] + extra_cuda_cflags = [flag.strip() for flag in extra_cuda_cflags] + extra_sycl_cflags = [flag.strip() for flag in extra_sycl_cflags] + extra_ldflags = [flag.strip() for flag in extra_ldflags] + extra_include_paths = [flag.strip() for flag in extra_include_paths] + + # Turn into absolute paths so we can emit them into the ninja build + # file wherever it is. + user_includes = [os.path.abspath(file) for file in extra_include_paths] + + # include_paths() gives us the location of torch/extension.h + # TODO generalize with_cuda as specific device type. + if with_cuda: + system_includes = include_paths("cuda") + elif with_sycl: + system_includes = include_paths("xpu") + else: + system_includes = include_paths("cpu") + # sysconfig.get_path('include') gives us the location of Python.h + # Explicitly specify 'posix_prefix' scheme on non-Windows platforms to workaround error on some MacOS + # installations where default `get_path` points to non-existing `/Library/Python/M.m/include` folder + python_include_path = sysconfig.get_path('include', scheme='nt' if IS_WINDOWS else 'posix_prefix') + if python_include_path is not None: + system_includes.append(python_include_path) + + common_cflags = [] + if not is_standalone: + common_cflags.append(f'-DTORCH_EXTENSION_NAME={name}') + common_cflags.append('-DTORCH_API_INCLUDE_EXTENSION_H') + + # Windows does not understand `-isystem` and quotes flags later. + if IS_WINDOWS: + common_cflags += [f'-I{include}' for include in user_includes + system_includes] + else: + common_cflags += [f'-I{shlex.quote(include)}' for include in user_includes] + common_cflags += [f'-isystem {shlex.quote(include)}' for include in system_includes] + + if IS_WINDOWS: + COMMON_HIP_FLAGS.extend(['-fms-runtime-lib=dll']) + cflags = common_cflags + ['/std:c++20'] + extra_cflags + cflags += COMMON_MSVC_FLAGS + (COMMON_HIP_FLAGS if IS_HIP_EXTENSION else []) + cflags = _nt_quote_args(cflags) + else: + cflags = common_cflags + ['-fPIC', '-std=c++20'] + extra_cflags + + if with_cuda and IS_HIP_EXTENSION: + cuda_flags = ['-DWITH_HIP'] + common_cflags + extra_cflags + COMMON_HIP_FLAGS + COMMON_HIPCC_FLAGS + cuda_flags = cuda_flags + ['-std=c++20'] + cuda_flags += _get_rocm_arch_flags(cuda_flags) + cuda_flags += extra_cuda_cflags + if IS_WINDOWS: + cuda_flags = _nt_quote_args(cuda_flags) + elif with_cuda: + cuda_flags = common_cflags + COMMON_NVCC_FLAGS + _get_cuda_arch_flags(extra_cuda_cflags) + if IS_WINDOWS: + for flag in COMMON_MSVC_FLAGS: + cuda_flags = ['-Xcompiler', flag] + cuda_flags + for ignore_warning in MSVC_IGNORE_CUDAFE_WARNINGS: + cuda_flags = ['-Xcudafe', '--diag_suppress=' + ignore_warning] + cuda_flags + cuda_flags = cuda_flags + ['-std=c++20'] + cuda_flags = _nt_quote_args(cuda_flags) + cuda_flags += _nt_quote_args(extra_cuda_cflags) + else: + cuda_flags += ['--compiler-options', "'-fPIC'"] + cuda_flags += extra_cuda_cflags + if not any(flag.startswith('-std=') for flag in cuda_flags): + cuda_flags.append('-std=c++20') + cc_env = os.getenv("CC") + if cc_env is not None: + cuda_flags = ['-ccbin', cc_env] + cuda_flags + else: + cuda_flags = None + + if with_sycl: + sycl_cflags = cflags + _COMMON_SYCL_FLAGS + sycl_cflags += extra_sycl_cflags + _append_sycl_targets_if_missing(sycl_cflags) + _append_sycl_std_if_no_std_present(sycl_cflags) + host_cflags = cflags + # escaping quoted arguments to pass them thru SYCL compiler + icpx_version = _get_icpx_version() + if int(icpx_version) < 20250200: + host_cflags = [item.replace('\\"', '\\\\"') for item in host_cflags] + + sycl_cflags += _wrap_sycl_host_flags(host_cflags) + sycl_dlink_post_cflags = _SYCL_DLINK_FLAGS.copy() + sycl_dlink_post_cflags += _get_sycl_device_flags(sycl_cflags) + else: + sycl_cflags = None + sycl_dlink_post_cflags = None + + def object_file_path(source_file: str) -> str: + # '/path/to/file.cpp' -> 'file' + file_name = os.path.splitext(os.path.basename(source_file))[0] + if _is_cuda_file(source_file) and with_cuda: + # Use a different object filename in case a C++ and CUDA file have + # the same filename but different extension (.cpp vs. .cu). + target = f'{file_name}.cuda.o' + elif _is_sycl_file(source_file) and with_sycl: + target = f'{file_name}.sycl.o' + else: + target = f'{file_name}.o' + return target + + objects = [object_file_path(src) for src in sources] + ldflags = ([] if is_standalone else [SHARED_FLAG]) + extra_ldflags + + # The darwin linker needs explicit consent to ignore unresolved symbols. + if IS_MACOS: + ldflags.append('-undefined dynamic_lookup') + elif IS_WINDOWS: + ldflags = _nt_quote_args(ldflags) + + ext = EXEC_EXT if is_standalone else LIB_EXT + library_target = f'{name}{ext}' + + _write_ninja_file( + path=path, + cflags=cflags, + post_cflags=None, + cuda_cflags=cuda_flags, + cuda_post_cflags=None, + cuda_dlink_post_cflags=None, + sycl_cflags=sycl_cflags, + sycl_post_cflags=[], + sycl_dlink_post_cflags=sycl_dlink_post_cflags, + sources=sources, + objects=objects, + ldflags=ldflags, + library_target=library_target, + with_cuda=with_cuda, + with_sycl=with_sycl) + + +def _write_ninja_file(path, + cflags, + post_cflags, + cuda_cflags, + cuda_post_cflags, + cuda_dlink_post_cflags, + sycl_cflags, + sycl_post_cflags, + sycl_dlink_post_cflags, + sources, + objects, + ldflags, + library_target, + with_cuda, + with_sycl) -> None: + r"""Write a ninja file that does the desired compiling and linking. + + `path`: Where to write this file + `cflags`: list of flags to pass to $cxx. Can be None. + `post_cflags`: list of flags to append to the $cxx invocation. Can be None. + `cuda_cflags`: list of flags to pass to $nvcc. Can be None. + `cuda_post_cflags`: list of flags to append to the $nvcc invocation. Can be None. + `cuda_dlink_post_cflags`: list of flags to append to the $nvcc device code link invocation. Can be None. + `sycl_cflags`: list of flags to pass to SYCL compiler. Can be None. + `sycl_post_cflags`: list of flags to append to the SYCL compiler invocation. Can be None. + `sycl_dlink_post_cflags`: list of flags to append to the SYCL compiler device code link invocation. Can be None. +e. + `sources`: list of paths to source files + `objects`: list of desired paths to objects, one per source. + `ldflags`: list of flags to pass to linker. Can be None. + `library_target`: Name of the output library. Can be None; in that case, + we do no linking. + `with_cuda`: If we should be compiling with CUDA. + """ + def sanitize_flags(flags): + if flags is None: + return [] + else: + return [flag.strip() for flag in flags] + + cflags = sanitize_flags(cflags) + post_cflags = sanitize_flags(post_cflags) + cuda_cflags = sanitize_flags(cuda_cflags) + cuda_post_cflags = sanitize_flags(cuda_post_cflags) + cuda_dlink_post_cflags = sanitize_flags(cuda_dlink_post_cflags) + sycl_cflags = sanitize_flags(sycl_cflags) + sycl_post_cflags = sanitize_flags(sycl_post_cflags) + sycl_dlink_post_cflags = sanitize_flags(sycl_dlink_post_cflags) + ldflags = sanitize_flags(ldflags) + + # Sanity checks... + if len(sources) != len(objects): + raise AssertionError("sources and objects lists must be the same length") + if len(sources) == 0: + raise AssertionError("At least one source is required to build a library") + + compiler = get_cxx_compiler() + + # Version 1.3 is required for the `deps` directive. + config = ['ninja_required_version = 1.3'] + config.append(f'cxx = {compiler}') + if with_cuda or cuda_dlink_post_cflags: + if "PYTORCH_NVCC" in os.environ: + nvcc = os.getenv("PYTORCH_NVCC") # user can set nvcc compiler with ccache using the environment variable here + else: + if IS_HIP_EXTENSION: + nvcc = _get_hipcc_path() + else: + nvcc = _join_cuda_home('bin', 'nvcc') + config.append(f'nvcc = {nvcc}') + if with_sycl or sycl_dlink_post_cflags: + sycl = 'icx' if IS_WINDOWS else 'icpx' + config.append(f'sycl = {sycl}') + + if IS_HIP_EXTENSION: + post_cflags = COMMON_HIP_FLAGS + post_cflags + flags = [f'cflags = {" ".join(cflags)}'] + flags.append(f'post_cflags = {" ".join(post_cflags)}') + if with_cuda: + flags.append(f'cuda_cflags = {" ".join(cuda_cflags)}') + flags.append(f'cuda_post_cflags = {" ".join(cuda_post_cflags)}') + flags.append(f'cuda_dlink_post_cflags = {" ".join(cuda_dlink_post_cflags)}') + if with_sycl: + flags.append(f'sycl_cflags = {" ".join(sycl_cflags)}') + flags.append(f'sycl_post_cflags = {" ".join(sycl_post_cflags)}') + flags.append(f'sycl_dlink_post_cflags = {" ".join(sycl_dlink_post_cflags)}') + flags.append(f'ldflags = {" ".join(ldflags)}') + + # Turn into absolute paths so we can emit them into the ninja build + # file wherever it is. + sources = [os.path.abspath(file) for file in sources] + + # See https://ninja-build.org/build.ninja.html for reference. + compile_rule = ['rule compile'] + if IS_WINDOWS: + compiler_name = "$cxx" if IS_HIP_EXTENSION else "cl" + compile_rule.append( + f' command = {compiler_name} ' + '/showIncludes $cflags -c $in /Fo$out $post_cflags' # codespell:ignore + ) + if not IS_HIP_EXTENSION: + compile_rule.append(' deps = msvc') + else: + compile_rule.append( + ' command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags') + compile_rule.append(' depfile = $out.d') + compile_rule.append(' deps = gcc') + + if with_cuda: + cuda_compile_rule = ['rule cuda_compile'] + nvcc_gendeps = '' + # -MD is not supported by ROCm + # Nvcc flag `-MD` is not supported by sccache, which may increase build time. + if torch.version.cuda is not None and os.getenv('TORCH_EXTENSION_SKIP_NVCC_GEN_DEPENDENCIES', '0') != '1': + cuda_compile_rule.append(' depfile = $out.d') + cuda_compile_rule.append(' deps = gcc') + # Note: non-system deps with nvcc are only supported + # on Linux so use -MD to make this work on Windows too. + nvcc_gendeps = '-MD -MF $out.d' + cuda_compile_rule.append( + f' command = $nvcc {nvcc_gendeps} $cuda_cflags -c $in -o $out $cuda_post_cflags') + + if with_sycl: + sycl_compile_rule = ['rule sycl_compile'] + # SYCL compiler does not recognize .sycl extension automatically, + # so we pass '-x c++' explicitly notifying compiler of file format + sycl_compile_rule.append( + ' command = $sycl $sycl_cflags -c -x c++ $in -o $out $sycl_post_cflags') + + + # Emit one build rule per source to enable incremental build. + build = [] + for source_file, object_file in zip(sources, objects, strict=True): + is_cuda_source = _is_cuda_file(source_file) and with_cuda + is_sycl_source = _is_sycl_file(source_file) and with_sycl + if is_cuda_source: + rule = 'cuda_compile' + elif is_sycl_source: + rule = 'sycl_compile' + else: + rule = 'compile' + if IS_WINDOWS: + source_file = source_file.replace(':', '$:') + object_file = object_file.replace(':', '$:') + source_file = source_file.replace(" ", "$ ") + object_file = object_file.replace(" ", "$ ") + build.append(f'build {object_file}: {rule} {source_file}') + + if cuda_dlink_post_cflags: + cuda_devlink_out = os.path.join(os.path.dirname(objects[0]), 'dlink.o') + cuda_devlink_rule = ['rule cuda_devlink'] + cuda_devlink_rule.append(' command = $nvcc $in -o $out $cuda_dlink_post_cflags') + cuda_devlink = [f'build {cuda_devlink_out}: cuda_devlink {" ".join(objects)}'] + objects += [cuda_devlink_out] + else: + cuda_devlink_rule, cuda_devlink = [], [] + + if sycl_dlink_post_cflags: + sycl_devlink_out = os.path.join(os.path.dirname(objects[0]), "sycl_dlink.o") + if IS_WINDOWS: + sycl_devlink_objects = [obj.replace(":", "$:") for obj in objects] + objects += [sycl_devlink_out] + sycl_devlink_out = sycl_devlink_out.replace(":", "$:") + else: + sycl_devlink_objects = list(objects) + objects += [sycl_devlink_out] + sycl_devlink_rule = ["rule sycl_devlink"] + sycl_devlink_rule.append( + " command = $sycl $in -o $out $sycl_dlink_post_cflags" + ) + sycl_devlink = [ + f"build {sycl_devlink_out}: sycl_devlink {' '.join(sycl_devlink_objects)}" + ] + else: + sycl_devlink_rule, sycl_devlink = [], [] + + if library_target is not None: + link_rule = ['rule link'] + if IS_WINDOWS: + cl_paths = subprocess.check_output(['where', + 'cl']).decode(*SUBPROCESS_DECODE_ARGS).split('\r\n') + if len(cl_paths) >= 1: + cl_path = os.path.dirname(cl_paths[0]).replace(':', '$:') + else: + raise RuntimeError("MSVC is required to load C++ extensions") + link_rule.append(f' command = "{cl_path}/link.exe" $in /nologo $ldflags /out:$out') + else: + link_rule.append(' command = $cxx $in $ldflags -o $out') + + link = [f'build {library_target}: link {" ".join(objects)}'] + + default = [f'default {library_target}'] + else: + link_rule, link, default = [], [], [] + + # 'Blocks' should be separated by newlines, for visual benefit. + blocks = [config, flags, compile_rule] + if with_cuda: + blocks.append(cuda_compile_rule) # type: ignore[possibly-undefined] + if with_sycl: + blocks.append(sycl_compile_rule) # type: ignore[possibly-undefined] + blocks += [cuda_devlink_rule, sycl_devlink_rule, link_rule, build, cuda_devlink, sycl_devlink, link, default] + content = "\n\n".join("\n".join(b) for b in blocks) + # Ninja requires a new lines at the end of the .ninja file + content += "\n" + _maybe_write(path, content) + +def _join_cuda_home(*paths) -> str: + """ + Join paths with CUDA_HOME, or raises an error if it CUDA_HOME is not set. + + This is basically a lazy way of raising an error for missing $CUDA_HOME + only once we need to get any CUDA-specific path. + """ + if CUDA_HOME is None: + raise OSError('CUDA_HOME environment variable is not set. ' + 'Please set it to your CUDA install root.') + return os.path.join(CUDA_HOME, *paths) + + +def _is_cuda_file(path: str) -> bool: + valid_ext = ['.cu', '.cuh'] + if IS_HIP_EXTENSION: + valid_ext.append('.hip') + return os.path.splitext(path)[1] in valid_ext + +def _is_sycl_file(path: str) -> bool: + valid_ext = ['.sycl'] + return os.path.splitext(path)[1] in valid_ext diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4ab5e7ce7f1c55a7b8bfff0bda646a4635231871 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/__init__.py @@ -0,0 +1,78 @@ +from torch.utils.data.dataloader import ( + _DatasetKind, + DataLoader, + default_collate, + default_convert, + get_worker_info, +) +from torch.utils.data.datapipes._decorator import ( + argument_validation, + functional_datapipe, + guaranteed_datapipes_determinism, + non_deterministic, + runtime_validation, + runtime_validation_disabled, +) +from torch.utils.data.datapipes.datapipe import ( + DataChunk, + DFIterDataPipe, + IterDataPipe, + MapDataPipe, +) +from torch.utils.data.dataset import ( + ChainDataset, + ConcatDataset, + Dataset, + IterableDataset, + random_split, + StackDataset, + Subset, + TensorDataset, +) +from torch.utils.data.distributed import DistributedSampler +from torch.utils.data.sampler import ( + BatchSampler, + RandomSampler, + Sampler, + SequentialSampler, + SubsetRandomSampler, + WeightedRandomSampler, +) + + +__all__ = [ + "BatchSampler", + "ChainDataset", + "ConcatDataset", + "DFIterDataPipe", + "DataChunk", + "DataLoader", + "Dataset", + "DistributedSampler", + "IterDataPipe", + "IterableDataset", + "MapDataPipe", + "RandomSampler", + "Sampler", + "SequentialSampler", + "StackDataset", + "Subset", + "SubsetRandomSampler", + "TensorDataset", + "WeightedRandomSampler", + "_DatasetKind", + "argument_validation", + "default_collate", + "default_convert", + "functional_datapipe", + "get_worker_info", + "guaranteed_datapipes_determinism", + "non_deterministic", + "random_split", + "runtime_validation", + "runtime_validation_disabled", +] + +# Please keep this list sorted +if __all__ != sorted(__all__): + raise AssertionError("__all__ is not sorted") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..44111ef697b7188df38711db1add2b8e0de4a293 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/__init__.py @@ -0,0 +1,53 @@ +r"""Utility classes & functions for data loading. Code in this folder is mostly used by ../dataloder.py. + +A lot of multiprocessing is used in data loading, which only supports running +functions defined in global environment (py2 can't serialize static methods). +Therefore, for code tidiness we put these functions into different files in this +folder. +""" + +import atexit +import sys + +# old private location of the ExceptionWrapper that some users rely on: +from torch._utils import ExceptionWrapper + + +IS_WINDOWS = sys.platform == "win32" + + +MP_STATUS_CHECK_INTERVAL = 5.0 +r"""Interval (in seconds) to check status of processes to avoid hanging in + multiprocessing data loading. This is mainly used in getting data from + another process, in which case we need to periodically check whether the + sender is alive to prevent hanging.""" + + +python_exit_status = False +r"""Whether Python is shutting down. This flag is guaranteed to be set before +the Python core library resources are freed, but Python may already be exiting +for some time when this is set. + +Hook to set this flag is `_set_python_exit_flag`, and is inspired by a similar +hook in Python 3.7 multiprocessing library: +https://github.com/python/cpython/blob/d4d60134b29290049e28df54f23493de4f1824b6/Lib/multiprocessing/util.py#L277-L327 +""" + + +try: + import numpy + + HAS_NUMPY = True +except ModuleNotFoundError: + HAS_NUMPY = False + + +def _set_python_exit_flag() -> None: + global python_exit_status + python_exit_status = True + + +atexit.register(_set_python_exit_flag) + + +from . import collate, fetch, pin_memory, signal_handling, worker diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py new file mode 100644 index 0000000000000000000000000000000000000000..8767fd52383a0604996370e6b2374e998658785b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/collate.py @@ -0,0 +1,401 @@ +# mypy: allow-untyped-defs +r"""Contains definitions of the methods used by the _BaseDataLoaderIter workers. + +These methods are used to collate samples fetched from dataset into Tensor(s). +These **needs** to be in global scope since Py2 doesn't support serializing +static methods. + +`default_collate` and `default_convert` are exposed to users via 'dataloader.py'. +""" + +import collections +import contextlib +import copy +import re +from collections.abc import Callable + +import torch + + +np_str_obj_array_pattern = re.compile(r"[SaUO]") + + +def default_convert(data): + r""" + Convert each NumPy array element into a :class:`torch.Tensor`. + + If the input is a `Sequence`, `Collection`, or `Mapping`, it tries to convert each element inside to a :class:`torch.Tensor`. + If the input is not an NumPy array, it is left unchanged. + This is used as the default function for collation when both `batch_sampler` and `batch_size` + are NOT defined in :class:`~torch.utils.data.DataLoader`. + + The general input type to output type mapping is similar to that + of :func:`~torch.utils.data.default_collate`. See the description there for more details. + + Args: + data: a single data point to be converted + + Examples: + >>> # xdoctest: +SKIP + >>> # Example with `int` + >>> default_convert(0) + 0 + >>> # Example with NumPy array + >>> default_convert(np.array([0, 1])) + tensor([0, 1]) + >>> # Example with NamedTuple + >>> Point = namedtuple("Point", ["x", "y"]) + >>> default_convert(Point(0, 0)) + Point(x=0, y=0) + >>> default_convert(Point(np.array(0), np.array(0))) + Point(x=tensor(0), y=tensor(0)) + >>> # Example with List + >>> default_convert([np.array([0, 1]), np.array([2, 3])]) + [tensor([0, 1]), tensor([2, 3])] + """ + elem_type = type(data) + if isinstance(data, torch.Tensor): + return data + elif ( + elem_type.__module__ == "numpy" + and elem_type.__name__ != "str_" + and elem_type.__name__ != "string_" + ): + # array of string classes and object + if ( + elem_type.__name__ == "ndarray" + and np_str_obj_array_pattern.search(data.dtype.str) is not None + ): + return data + return torch.as_tensor(data) + elif isinstance(data, collections.abc.Mapping): + try: + if isinstance(data, collections.abc.MutableMapping): + # The mapping type may have extra properties, so we can't just + # use `type(data)(...)` to create the new mapping. + # Create a clone and update it if the mapping type is mutable. + clone = copy.copy(data) + clone.update({key: default_convert(data[key]) for key in data}) + return clone + else: + return elem_type({key: default_convert(data[key]) for key in data}) + except TypeError: + # The mapping type may not support `copy()` / `update(mapping)` + # or `__init__(iterable)`. + return {key: default_convert(data[key]) for key in data} + elif isinstance(data, tuple) and hasattr(data, "_fields"): # namedtuple + return elem_type(*(default_convert(d) for d in data)) + elif isinstance(data, tuple): + return [default_convert(d) for d in data] # Backwards compatibility. + elif isinstance(data, collections.abc.Sequence) and not isinstance( + data, (str, bytes) + ): + try: + if isinstance(data, collections.abc.MutableSequence): + # The sequence type may have extra properties, so we can't just + # use `type(data)(...)` to create the new sequence. + # Create a clone and update it if the sequence type is mutable. + clone = copy.copy(data) # type: ignore[arg-type] + for i, d in enumerate(data): + clone[i] = default_convert(d) + return clone + else: + return elem_type([default_convert(d) for d in data]) + except TypeError: + # The sequence type may not support `copy()` / `__setitem__(index, item)` + # or `__init__(iterable)` (e.g., `range`). + return [default_convert(d) for d in data] + else: + return data + + +default_collate_err_msg_format = ( + "default_collate: batch must contain tensors, numpy arrays, numbers, " + "dicts or lists; found {}" +) + + +def collate( + batch, + *, + collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None, +): + r""" + General collate function that handles collection type of element within each batch. + + The function also opens function registry to deal with specific element types. `default_collate_fn_map` + provides default collate functions for tensors, numpy arrays, numbers and strings. + + Args: + batch: a single batch to be collated + collate_fn_map: Optional dictionary mapping from element type to the corresponding collate function. + If the element type isn't present in this dictionary, + this function will go through each key of the dictionary in the insertion order to + invoke the corresponding collate function if the element type is a subclass of the key. + + Examples: + >>> def collate_tensor_fn(batch, *, collate_fn_map): + ... # Extend this function to handle batch of tensors + ... return torch.stack(batch, 0) + >>> def custom_collate(batch): + ... collate_map = {torch.Tensor: collate_tensor_fn} + ... return collate(batch, collate_fn_map=collate_map) + >>> # Extend `default_collate` by in-place modifying `default_collate_fn_map` + >>> default_collate_fn_map.update({torch.Tensor: collate_tensor_fn}) + + Note: + Each collate function requires a positional argument for batch and a keyword argument + for the dictionary of collate functions as `collate_fn_map`. + """ + elem = batch[0] + elem_type = type(elem) + + if collate_fn_map is not None: + if elem_type in collate_fn_map: + return collate_fn_map[elem_type](batch, collate_fn_map=collate_fn_map) + + for collate_type in collate_fn_map: + if isinstance(elem, collate_type): + return collate_fn_map[collate_type]( + batch, collate_fn_map=collate_fn_map + ) + + if isinstance(elem, collections.abc.Mapping): + try: + if isinstance(elem, collections.abc.MutableMapping): + # The mapping type may have extra properties, so we can't just + # use `type(data)(...)` to create the new mapping. + # Create a clone and update it if the mapping type is mutable. + clone = copy.copy(elem) + clone.update( + { + key: collate( + [d[key] for d in batch], collate_fn_map=collate_fn_map + ) + for key in elem + } + ) + return clone + else: + return elem_type( + { + key: collate( + [d[key] for d in batch], collate_fn_map=collate_fn_map + ) + for key in elem + } + ) + except TypeError: + # The mapping type may not support `copy()` / `update(mapping)` + # or `__init__(iterable)`. + return { + key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) + for key in elem + } + elif isinstance(elem, tuple) and hasattr(elem, "_fields"): # namedtuple + return elem_type( + *( + collate(samples, collate_fn_map=collate_fn_map) + for samples in zip(*batch, strict=False) + ) + ) + elif isinstance(elem, collections.abc.Sequence): + # check to make sure that the elements in batch have consistent size + it = iter(batch) + elem_size = len(next(it)) + + if not all(len(elem) == elem_size for elem in it): + raise RuntimeError("each element in list of batch should be of equal size") + transposed = list( + zip(*batch, strict=False) + ) # It may be accessed twice, so we use a list. + + if isinstance(elem, tuple): + return [ + collate(samples, collate_fn_map=collate_fn_map) + for samples in transposed + ] # Backwards compatibility. + else: + try: + if isinstance(elem, collections.abc.MutableSequence): + # The sequence type may have extra properties, so we can't just + # use `type(data)(...)` to create the new sequence. + # Create a clone and update it if the sequence type is mutable. + clone = copy.copy(elem) # type: ignore[arg-type] + for i, samples in enumerate(transposed): + clone[i] = collate(samples, collate_fn_map=collate_fn_map) + return clone + else: + return elem_type( + [ + collate(samples, collate_fn_map=collate_fn_map) + for samples in transposed + ] + ) + except TypeError: + # The sequence type may not support `copy()` / `__setitem__(index, item)` + # or `__init__(iterable)` (e.g., `range`). + return [ + collate(samples, collate_fn_map=collate_fn_map) + for samples in transposed + ] + + raise TypeError(default_collate_err_msg_format.format(elem_type)) + + +def collate_tensor_fn( + batch, + *, + collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None, +): + elem = batch[0] + out = None + if elem.is_nested: + raise RuntimeError( + "Batches of nested tensors are not currently supported by the default collate_fn; " + "please provide a custom collate_fn to handle them appropriately." + ) + if elem.layout in { + torch.sparse_coo, + torch.sparse_csr, + torch.sparse_bsr, + torch.sparse_csc, + torch.sparse_bsc, + }: + raise RuntimeError( + "Batches of sparse tensors are not currently supported by the default collate_fn; " + "please provide a custom collate_fn to handle them appropriately." + ) + if torch.utils.data.get_worker_info() is not None: + # If we're in a background process, concatenate directly into a + # shared memory tensor to avoid an extra copy + numel = sum(x.numel() for x in batch) + storage = elem._typed_storage()._new_shared(numel, device=elem.device) + out = elem.new(storage).resize_(len(batch), *list(elem.size())) + return torch.stack(batch, 0, out=out) + + +def collate_numpy_array_fn( + batch, + *, + collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None, +): + elem = batch[0] + # array of string classes and object + if np_str_obj_array_pattern.search(elem.dtype.str) is not None: + raise TypeError(default_collate_err_msg_format.format(elem.dtype)) + + return collate([torch.as_tensor(b) for b in batch], collate_fn_map=collate_fn_map) + + +def collate_numpy_scalar_fn( + batch, + *, + collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None, +): + return torch.as_tensor(batch) + + +def collate_float_fn( + batch, + *, + collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None, +): + return torch.tensor(batch, dtype=torch.float64) + + +def collate_int_fn( + batch, + *, + collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None, +): + return torch.tensor(batch) + + +def collate_str_fn( + batch, + *, + collate_fn_map: dict[type | tuple[type, ...], Callable] | None = None, +): + return batch + + +default_collate_fn_map: dict[type | tuple[type, ...], Callable] = { + torch.Tensor: collate_tensor_fn +} +with contextlib.suppress(ImportError): + import numpy as np + + # For both ndarray and memmap (subclass of ndarray) + default_collate_fn_map[np.ndarray] = collate_numpy_array_fn + # See scalars hierarchy: https://numpy.org/doc/stable/reference/arrays.scalars.html + # Skip string scalars + default_collate_fn_map[(np.bool_, np.number, np.object_)] = collate_numpy_scalar_fn +default_collate_fn_map[float] = collate_float_fn +default_collate_fn_map[int] = collate_int_fn +default_collate_fn_map[str] = collate_str_fn +default_collate_fn_map[bytes] = collate_str_fn + + +def default_collate(batch): + r""" + Take in a batch of data and put the elements within the batch into a tensor with an additional outer dimension - batch size. + + The exact output type can be a :class:`torch.Tensor`, a `Sequence` of :class:`torch.Tensor`, a + Collection of :class:`torch.Tensor`, or left unchanged, depending on the input type. + This is used as the default function for collation when + `batch_size` or `batch_sampler` is defined in :class:`~torch.utils.data.DataLoader`. + + Here is the general input type (based on the type of the element within the batch) to output type mapping: + + * :class:`torch.Tensor` -> :class:`torch.Tensor` (with an added outer dimension batch size) + * NumPy Arrays -> :class:`torch.Tensor` + * `float` -> :class:`torch.Tensor` + * `int` -> :class:`torch.Tensor` + * `str` -> `str` (unchanged) + * `bytes` -> `bytes` (unchanged) + * `Mapping[K, V_i]` -> `Mapping[K, default_collate([V_1, V_2, ...])]` + * `NamedTuple[V1_i, V2_i, ...]` -> `NamedTuple[default_collate([V1_1, V1_2, ...]), + default_collate([V2_1, V2_2, ...]), ...]` + * `Sequence[V1_i, V2_i, ...]` -> `Sequence[default_collate([V1_1, V1_2, ...]), + default_collate([V2_1, V2_2, ...]), ...]` + + Args: + batch: a single batch to be collated + + Examples: + >>> # xdoctest: +SKIP + >>> # Example with a batch of `int`s: + >>> default_collate([0, 1, 2, 3]) + tensor([0, 1, 2, 3]) + >>> # Example with a batch of `str`s: + >>> default_collate(["a", "b", "c"]) + ['a', 'b', 'c'] + >>> # Example with `Map` inside the batch: + >>> default_collate([{"A": 0, "B": 1}, {"A": 100, "B": 100}]) + {'A': tensor([ 0, 100]), 'B': tensor([ 1, 100])} + >>> # Example with `NamedTuple` inside the batch: + >>> Point = namedtuple("Point", ["x", "y"]) + >>> default_collate([Point(0, 0), Point(1, 1)]) + Point(x=tensor([0, 1]), y=tensor([0, 1])) + >>> # Example with `Tuple` inside the batch: + >>> default_collate([(0, 1), (2, 3)]) + [tensor([0, 2]), tensor([1, 3])] + >>> # Example with `List` inside the batch: + >>> default_collate([[0, 1], [2, 3]]) + [tensor([0, 2]), tensor([1, 3])] + >>> # Two options to extend `default_collate` to handle specific type + >>> # Option 1: Write custom collate function and invoke `default_collate` + >>> def custom_collate(batch): + ... elem = batch[0] + ... if isinstance(elem, CustomType): # Some custom condition + ... return ... + ... else: # Fall back to `default_collate` + ... return default_collate(batch) + >>> # Option 2: In-place modify `default_collate_fn_map` + >>> def collate_customtype_fn(batch, *, collate_fn_map=None): + ... return ... + >>> default_collate_fn_map.update(CustomType, collate_customtype_fn) + >>> default_collate(batch) # Handle `CustomType` automatically + """ + return collate(batch, collate_fn_map=default_collate_fn_map) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/fetch.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/fetch.py new file mode 100644 index 0000000000000000000000000000000000000000..9bcd0ec5b30731269fc304b5ef2e087d94dc3211 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/fetch.py @@ -0,0 +1,57 @@ +# mypy: allow-untyped-defs +r"""Contains definitions of the methods used by the _BaseDataLoaderIter to fetch data from an iterable-style or map-style dataset. + +This logic is shared in both single- and multi-processing data loading. +""" + +from typing import NoReturn + + +class _BaseDatasetFetcher: + def __init__(self, dataset, auto_collation, collate_fn, drop_last) -> None: + self.dataset = dataset + self.auto_collation = auto_collation + self.collate_fn = collate_fn + self.drop_last = drop_last + + def fetch(self, possibly_batched_index) -> NoReturn: + raise NotImplementedError + + +class _IterableDatasetFetcher(_BaseDatasetFetcher): + def __init__(self, dataset, auto_collation, collate_fn, drop_last) -> None: + super().__init__(dataset, auto_collation, collate_fn, drop_last) + self.dataset_iter = iter(dataset) + self.ended = False + + def fetch(self, possibly_batched_index): + if self.ended: + raise StopIteration + + if self.auto_collation: + data = [] + for _ in possibly_batched_index: + try: + data.append(next(self.dataset_iter)) + except StopIteration: + self.ended = True + break + if len(data) == 0 or ( + self.drop_last and len(data) < len(possibly_batched_index) + ): + raise StopIteration + else: + data = next(self.dataset_iter) + return self.collate_fn(data) + + +class _MapDatasetFetcher(_BaseDatasetFetcher): + def fetch(self, possibly_batched_index): + if self.auto_collation: + if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__: + data = self.dataset.__getitems__(possibly_batched_index) + else: + data = [self.dataset[idx] for idx in possibly_batched_index] + else: + data = self.dataset[possibly_batched_index] + return self.collate_fn(data) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/pin_memory.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/pin_memory.py new file mode 100644 index 0000000000000000000000000000000000000000..8d34ad4d959e986e680bdbb5df5d1998c87bd23a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/pin_memory.py @@ -0,0 +1,108 @@ +# mypy: allow-untyped-defs +r"""Contains definitions of the methods used by the _BaseDataLoaderIter to put fetched tensors into pinned memory. + +These **needs** to be in global scope since Py2 doesn't support serializing +static methods. +""" + +import collections +import copy +import queue + +import torch +from torch._utils import ExceptionWrapper + +from . import MP_STATUS_CHECK_INTERVAL + + +def _pin_memory_loop(in_queue, out_queue, device_id, done_event, device) -> None: + # This setting is thread local, and prevents the copy in pin_memory from + # consuming all CPU cores. + torch.set_num_threads(1) + + torch.multiprocessing._set_thread_name("pt_data_pin") + torch.accelerator.set_device_index(device_id) + + def do_one_step() -> None: + try: + r = in_queue.get(timeout=MP_STATUS_CHECK_INTERVAL) + except queue.Empty: + return + idx, data = r + if not done_event.is_set() and not isinstance(data, ExceptionWrapper): + try: + data = pin_memory(data, device) + except Exception: + data = ExceptionWrapper( + where=f"in pin memory thread for device {device_id}" + ) + r = (idx, data) + while not done_event.is_set(): + try: + out_queue.put(r, timeout=MP_STATUS_CHECK_INTERVAL) + break + except queue.Full: + continue + + # See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on the + # logic of this function. + while not done_event.is_set(): + # Make sure that we don't preserve any object from one iteration + # to the next + do_one_step() + + +def pin_memory(data, device=None): + if isinstance(data, torch.Tensor): + return data.pin_memory() + + if hasattr(data, "pin_memory"): + return data.pin_memory() + + if isinstance(data, (str, bytes)): + return data + + if isinstance(data, collections.abc.Mapping): + try: + if isinstance(data, collections.abc.MutableMapping): + # The mapping type may have extra properties, so we can't just + # use `type(data)(...)` to create the new mapping. + # Create a clone and update it if the mapping type is mutable. + clone = copy.copy(data) + clone.update( + {k: pin_memory(sample, device) for k, sample in data.items()} + ) + return clone + else: + # pyrefly: ignore [bad-instantiation] + return type(data)( + # pyrefly: ignore [bad-argument-count] + {k: pin_memory(sample, device) for k, sample in data.items()} + ) # type: ignore[call-arg] + except TypeError: + # The mapping type may not support `copy()` / `update(mapping)` + # or `__init__(iterable)`. + return {k: pin_memory(sample, device) for k, sample in data.items()} + + if isinstance(data, tuple): + if hasattr(data, "_fields"): # namedtuple + return type(data)(*(pin_memory(sample, device) for sample in data)) + return type(data)(pin_memory(sample, device) for sample in data) + + if isinstance(data, collections.abc.Sequence): + try: + if isinstance(data, collections.abc.MutableSequence): + # The sequence type may have extra properties, so we can't just + # use `type(data)(...)` to create the new sequence. + # Create a clone and update it if the sequence type is mutable. + clone = copy.copy(data) # type: ignore[arg-type] + for i, item in enumerate(data): + clone[i] = pin_memory(item, device) + return clone + return type(data)([pin_memory(sample, device) for sample in data]) # type: ignore[call-arg] + except TypeError: + # The sequence type may not support `copy()` / `__setitem__(index, item)` + # or `__init__(iterable)` (e.g., `range`). + return [pin_memory(sample, device) for sample in data] + + return data diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/signal_handling.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/signal_handling.py new file mode 100644 index 0000000000000000000000000000000000000000..abff09bc40819d83420a08e0b90d7ba816f4764f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/signal_handling.py @@ -0,0 +1,80 @@ +# mypy: allow-untyped-defs +r"""Signal handling for multiprocessing data loading. + +NOTE [ Signal handling in multiprocessing data loading ] + +In cases like DataLoader, if a worker process dies due to bus error/segfault +or just hang, the main process will hang waiting for data. This is difficult +to avoid on PyTorch side as it can be caused by limited shm, or other +libraries users call in the workers. In this file and `DataLoader.cpp`, we make +our best effort to provide some error message to users when such unfortunate +events happen. + +When a _BaseDataLoaderIter starts worker processes, their pids are registered in a +defined in `DataLoader.cpp`: id(_BaseDataLoaderIter) => Collection[ Worker pids ] +via `_set_worker_pids`. + +When an error happens in a worker process, the main process received a SIGCHLD, +and Python will eventually call the handler registered below +(in `_set_SIGCHLD_handler`). In the handler, the `_error_if_any_worker_fails` +call checks all registered worker pids and raise proper error message to +prevent main process from hanging waiting for data from worker. + +Additionally, at the beginning of each worker's `_utils.worker._worker_loop`, +`_set_worker_signal_handlers` is called to register critical signal handlers +(e.g., for SIGSEGV, SIGBUS, SIGFPE, SIGTERM) in C, which just prints an error +message to stderr before triggering the default handler. So a message will also +be printed from the worker process when it is killed by such signals. + +See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for the reasoning of +this signal handling design and other mechanism we implement to make our +multiprocessing data loading robust to errors. +""" + +import signal +import threading + +# Some of the following imported functions are not used in this file, but are to +# be used `_utils.signal_handling.XXXXX`. +from torch._C import ( # noqa: F401 + _error_if_any_worker_fails, + _remove_worker_pids, + _set_worker_pids, + _set_worker_signal_handlers, +) + +from . import IS_WINDOWS + + +_SIGCHLD_handler_set = False +r"""Whether SIGCHLD handler is set for DataLoader worker failures. Only one +handler needs to be set for all DataLoaders in a process.""" + + +def _set_SIGCHLD_handler() -> None: + # Windows doesn't support SIGCHLD handler + if IS_WINDOWS: + return + # can't set signal in child threads + if not isinstance(threading.current_thread(), threading._MainThread): # type: ignore[attr-defined] + return + global _SIGCHLD_handler_set + if _SIGCHLD_handler_set: + return + previous_handler = signal.getsignal(signal.SIGCHLD) + if not callable(previous_handler): + # This doesn't catch default handler, but SIGCHLD default handler is a + # no-op. + previous_handler = None + + def handler(signum, frame) -> None: + # This following call uses `waitid` with WNOHANG from C side. Therefore, + # Python can still get and update the process status successfully. + _error_if_any_worker_fails() + if previous_handler is not None: + if not callable(previous_handler): + raise AssertionError("previous_handler is not callable") + previous_handler(signum, frame) + + signal.signal(signal.SIGCHLD, handler) + _SIGCHLD_handler_set = True diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/worker.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/worker.py new file mode 100644 index 0000000000000000000000000000000000000000..850365ee9b19b2136778cc5714610da68ba94e19 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/_utils/worker.py @@ -0,0 +1,399 @@ +# mypy: allow-untyped-defs +r"""Contains definitions of the methods used by the _BaseDataLoaderIter workers. + +These **needs** to be in global scope since Py2 doesn't support serializing +static methods. +""" + +from __future__ import annotations + +import os +import queue +import random +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torch +from torch._utils import ExceptionWrapper + +from . import HAS_NUMPY, IS_WINDOWS, MP_STATUS_CHECK_INTERVAL, signal_handling + + +if TYPE_CHECKING: + from torch.utils.data import Dataset + +if IS_WINDOWS: + import ctypes + from ctypes.wintypes import BOOL, DWORD, HANDLE + + # On Windows, the parent ID of the worker process remains unchanged when the manager process + # is gone, and the only way to check it through OS is to let the worker have a process handle + # of the manager and ask if the process status has changed. + class ManagerWatchdog: + def __init__(self) -> None: + self.manager_pid = os.getppid() + + # mypy cannot detect this code is windows only + self.kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) # type: ignore[attr-defined] + self.kernel32.OpenProcess.argtypes = (DWORD, BOOL, DWORD) + self.kernel32.OpenProcess.restype = HANDLE + self.kernel32.WaitForSingleObject.argtypes = (HANDLE, DWORD) + self.kernel32.WaitForSingleObject.restype = DWORD + + # Value obtained from https://msdn.microsoft.com/en-us/library/ms684880.aspx + SYNCHRONIZE = 0x00100000 + self.manager_handle = self.kernel32.OpenProcess( + SYNCHRONIZE, 0, self.manager_pid + ) + + if not self.manager_handle: + raise ctypes.WinError(ctypes.get_last_error()) # type: ignore[attr-defined] + + self.manager_dead = False + + def is_alive(self) -> bool: + if not self.manager_dead: + # Value obtained from https://msdn.microsoft.com/en-us/library/windows/desktop/ms687032.aspx + self.manager_dead = ( + self.kernel32.WaitForSingleObject(self.manager_handle, 0) == 0 + ) + return not self.manager_dead + +else: + + class ManagerWatchdog: # type: ignore[no-redef] + def __init__(self) -> None: + self.manager_pid = os.getppid() + self.manager_dead = False + + def is_alive(self) -> bool: + if not self.manager_dead: + self.manager_dead = os.getppid() != self.manager_pid + return not self.manager_dead + + +_worker_info: WorkerInfo | None = None + + +@dataclass(frozen=True, slots=True) +class WorkerInfo: + """Information about the current DataLoader worker process or thread. + + Attributes: + id: The current worker id (0 to num_workers - 1) + num_workers: Total number of workers + seed: Random seed set for this worker + dataset: Copy of the dataset object in this worker + rng: Optional RNG state container. Defaults to None. + worker_method: Optional worker method ("multiprocessing" or "thread"). Defaults to "multiprocessing". + """ + + id: int + num_workers: int + seed: int + dataset: Dataset + rng: _RNG | None = None + worker_method: str | None = "multiprocessing" + + +def get_worker_info() -> WorkerInfo | None: + r"""Returns the information about the current + :class:`~torch.utils.data.DataLoader` iterator worker process. + + When called in a worker, this returns an object guaranteed to have the + following attributes: + + * :attr:`id`: the current worker id. + * :attr:`num_workers`: the total number of workers. + * :attr:`seed`: the random seed set for the current worker. This value is + determined by main process RNG and the worker id. See + :class:`~torch.utils.data.DataLoader`'s documentation for more details. + * :attr:`dataset`: the copy of the dataset object in **this** process. Note + that this will be a different object in a different process than the one + in the main process. + + When called in the main process, this returns ``None``. + + .. note:: + When used in a :attr:`worker_init_fn` passed over to + :class:`~torch.utils.data.DataLoader`, this method can be useful to + set up each worker process differently, for instance, using ``worker_id`` + to configure the ``dataset`` object to only read a specific fraction of a + sharded dataset, or use ``seed`` to seed other libraries used in dataset + code. + """ + return _worker_info + + +r"""Dummy class used to signal the end of an IterableDataset""" + + +@dataclass(frozen=True) +class _IterableDatasetStopIteration: + worker_id: int + + +r"""Dummy class used to resume the fetching when worker reuse is enabled""" + + +@dataclass(frozen=True) +class _ResumeIteration: + seed: int | None = None + + +@dataclass(frozen=True, slots=True) +class _RNG: + """Container for thread-local random number generator state. + + Used by thread workers to maintain separate RNG state per worker thread + to avoid race conditions. + + Attributes: + random_generator: Python random.Random generator for this thread + torch_generator: PyTorch Generator for this thread + numpy_generator: NumPy Generator for this thread (None if numpy not available) + """ + + random_generator: random.Random + torch_generator: torch.Generator + numpy_generator: object | None = None + + +# The function `_generate_state` is adapted from `numpy.random.SeedSequence` +# from https://github.com/numpy/numpy/blob/main/numpy/random/bit_generator.pyx +# It's MIT licensed, here is the copyright: + +# Copyright (c) 2015 Melissa E. O'Neill +# Copyright (c) 2019 NumPy Developers +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +# SOFTWARE. + + +# This function generates an array of int32 as the seed for +# `numpy.random`, in order to prevent state collision due to same +# seed and algorithm for `numpy.random` and `random` modules. +# TODO: Implement `SeedSequence` like object for `torch.random` +def _generate_state(base_seed, worker_id): + INIT_A = 0x43B0D7E5 + MULT_A = 0x931E8875 + INIT_B = 0x8B51F9DD + MULT_B = 0x58F38DED + MIX_MULT_L = 0xCA01F9DD + MIX_MULT_R = 0x4973F715 + XSHIFT = 4 * 8 // 2 + MASK32 = 0xFFFFFFFF + + entropy = [worker_id, base_seed & MASK32, base_seed >> 32, 0] + pool = [0] * 4 + + hash_const_A = INIT_A + + def hash(value): + nonlocal hash_const_A + value = (value ^ hash_const_A) & MASK32 + hash_const_A = (hash_const_A * MULT_A) & MASK32 + value = (value * hash_const_A) & MASK32 + value = (value ^ (value >> XSHIFT)) & MASK32 + return value + + def mix(x, y): + result_x = (MIX_MULT_L * x) & MASK32 + result_y = (MIX_MULT_R * y) & MASK32 + result = (result_x - result_y) & MASK32 + result = (result ^ (result >> XSHIFT)) & MASK32 + return result + + # Add in the entropy to the pool. + for i in range(len(pool)): + pool[i] = hash(entropy[i]) + + # Mix all bits together so late bits can affect earlier bits. + for i_src in range(len(pool)): + for i_dst in range(len(pool)): + if i_src != i_dst: + pool[i_dst] = mix(pool[i_dst], hash(pool[i_src])) + + hash_const_B = INIT_B + state = [] + for i_dst in range(4): + data_val = pool[i_dst] + data_val = (data_val ^ hash_const_B) & MASK32 + hash_const_B = (hash_const_B * MULT_B) & MASK32 + data_val = (data_val * hash_const_B) & MASK32 + data_val = (data_val ^ (data_val >> XSHIFT)) & MASK32 + state.append(data_val) + return state + + +def _worker_loop( + dataset_kind, + dataset, + index_queue, + data_queue, + done_event, + auto_collation, + collate_fn, + drop_last, + base_seed, + init_fn, + worker_id, + num_workers, + persistent_workers, + shared_seed, +) -> None: + # See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on the + # logic of this function. + + try: + # Initialize C side signal handlers for SIGBUS and SIGSEGV. Python signal + # module's handlers are executed after Python returns from C low-level + # handlers, likely when the same fatal signal had already happened + # again. + # https://docs.python.org/3/library/signal.html#execution-of-python-signal-handlers + signal_handling._set_worker_signal_handlers() + + torch.multiprocessing._set_thread_name("pt_data_worker") + + torch.set_num_threads(1) + seed = base_seed + worker_id + random.seed(seed) + torch.manual_seed(seed) + if HAS_NUMPY: + np_seed = _generate_state(base_seed, worker_id) + import numpy as np + + np.random.seed(np_seed) + + from torch.utils.data import IterDataPipe + from torch.utils.data.graph_settings import apply_random_seed + + shared_rng = torch.Generator() + if isinstance(dataset, IterDataPipe): + if shared_seed is None: + raise AssertionError( + "shared_seed must be provided for IterDataPipe workers" + ) + shared_rng.manual_seed(shared_seed) + dataset = apply_random_seed(dataset, shared_rng) + + global _worker_info + _worker_info = WorkerInfo( + id=worker_id, num_workers=num_workers, seed=seed, dataset=dataset + ) + + from torch.utils.data import _DatasetKind + + init_exception = None + + try: + if init_fn is not None: + init_fn(worker_id) + + fetcher = _DatasetKind.create_fetcher( + dataset_kind, dataset, auto_collation, collate_fn, drop_last + ) + except Exception: + init_exception = ExceptionWrapper( + where=f"in DataLoader worker process {worker_id}" + ) + + # When using Iterable mode, some worker can exit earlier than others due + # to the IterableDataset behaving differently for different workers. + # When such things happen, an `_IterableDatasetStopIteration` object is + # sent over to the main process with the ID of this worker, so that the + # main process won't send more tasks to this worker, and will send + # `None` to this worker to properly exit it. + # + # Note that we cannot set `done_event` from a worker as it is shared + # among all processes. Instead, we set the `iteration_end` flag to + # signify that the iterator is exhausted. When either `done_event` or + # `iteration_end` is set, we skip all processing step and just wait for + # `None`. + iteration_end = False + + watchdog = ManagerWatchdog() + + while watchdog.is_alive(): + try: + r = index_queue.get(timeout=MP_STATUS_CHECK_INTERVAL) + except queue.Empty: + continue + if isinstance(r, _ResumeIteration): + # Acknowledge the main process + data_queue.put((r, None)) + iteration_end = False + + if isinstance(dataset, IterDataPipe): + if r.seed is None: + raise AssertionError( + "resume iteration seed is None for IterDataPipe" + ) + shared_rng.manual_seed(r.seed) + dataset = apply_random_seed(dataset, shared_rng) + + # Recreate the fetcher for worker-reuse policy + fetcher = _DatasetKind.create_fetcher( + dataset_kind, dataset, auto_collation, collate_fn, drop_last + ) + continue + elif r is None: + # Received the final signal + if not done_event.is_set() and not iteration_end: + raise AssertionError( + "Received final signal but neither done_event nor iteration_end is set" + ) + break + elif done_event.is_set() or iteration_end: + # `done_event` is set. But I haven't received the final signal + # (None) yet. I will keep continuing until get it, and skip the + # processing steps. + continue + idx, index = r + data: _IterableDatasetStopIteration | ExceptionWrapper + if init_exception is not None: + data = init_exception + init_exception = None + else: + try: + data = fetcher.fetch(index) # type: ignore[possibly-undefined] + except Exception as e: + if ( + isinstance(e, StopIteration) + and dataset_kind == _DatasetKind.Iterable + ): + data = _IterableDatasetStopIteration(worker_id) + # Set `iteration_end` + # (1) to save future `next(...)` calls, and + # (2) to avoid sending multiple `_IterableDatasetStopIteration`s. + iteration_end = True + else: + # It is important that we don't store exc_info in a variable. + # `ExceptionWrapper` does the correct thing. + # See NOTE [ Python Traceback Reference Cycle Problem ] + data = ExceptionWrapper( + where=f"in DataLoader worker process {worker_id}" + ) + data_queue.put((idx, data)) + del data, idx, index, r # save memory + except KeyboardInterrupt: + # Main process will raise KeyboardInterrupt anyways. + pass + if done_event.is_set(): + data_queue.cancel_join_thread() + data_queue.close() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/backward_compatibility.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/backward_compatibility.py new file mode 100644 index 0000000000000000000000000000000000000000..5b928aea69fa7a7033a82021c5f41e053ff962fa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/backward_compatibility.py @@ -0,0 +1,11 @@ +# mypy: allow-untyped-defs +from typing_extensions import deprecated as _deprecated + + +@_deprecated( + "Usage of `backward_compatibility.worker_init_fn` is deprecated " + "as `DataLoader` automatically applies sharding in every worker", + category=FutureWarning, +) +def worker_init_fn(worker_id) -> None: + pass diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/dataloader.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..4eaa7559807877edb284fa306e10809de1d559af --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/dataloader.py @@ -0,0 +1,1686 @@ +# mypy: allow-untyped-defs +r"""Definition of the DataLoader and associated iterators that subclass _BaseDataLoaderIter. + +To support these two classes, in `./_utils` we define many utility methods and +functions to be run in multiprocessing. E.g., the data loading worker loop is +in `./_utils/worker.py`. +""" + +from __future__ import annotations + +import contextlib +import functools +import itertools +import logging +import multiprocessing as python_multiprocessing +import os +import queue +import threading +import warnings +from collections.abc import Callable +from typing import Any, Generic, NoReturn, TYPE_CHECKING, TypeVar +from typing_extensions import Self + +import torch +import torch.distributed as dist +import torch.utils.data.graph_settings +from torch._utils import ExceptionWrapper +from torch.utils.data import _utils +from torch.utils.data.datapipes.datapipe import ( + _IterDataPipeSerializationWrapper, + _MapDataPipeSerializationWrapper, + IterDataPipe, + MapDataPipe, +) +from torch.utils.data.dataset import Dataset, IterableDataset +from torch.utils.data.sampler import ( + BatchSampler, + RandomSampler, + Sampler, + SequentialSampler, +) + + +if TYPE_CHECKING: + from collections.abc import Iterable + +__all__ = [ + "DataLoader", + "get_worker_info", + "default_collate", + "default_convert", +] + + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) +_worker_init_fn_t = Callable[[int], None] + +# Ideally we would parameterize `DataLoader` by the return type of `collate_fn`, but there is currently no way to have that +# type parameter set to a default value if the user doesn't pass in a custom 'collate_fn'. +# See https://github.com/python/mypy/issues/3737. +_collate_fn_t = Callable[[list[_T]], Any] + + +# These functions used to be defined in this file. However, it was moved to +# _utils/collate.py. Although it is rather hard to access this from user land +# (one has to explicitly directly `import torch.utils.data.dataloader`), there +# probably is user code out there using it. This aliasing maintains BC in this +# aspect. +default_collate: _collate_fn_t = _utils.collate.default_collate +default_convert = _utils.collate.default_convert + +get_worker_info = _utils.worker.get_worker_info + +logger = logging.getLogger(__name__) + + +class _DatasetKind: + Map = 0 + Iterable = 1 + + @staticmethod + def create_fetcher(kind, dataset, auto_collation, collate_fn, drop_last): + if kind == _DatasetKind.Map: + return _utils.fetch._MapDatasetFetcher( + dataset, auto_collation, collate_fn, drop_last + ) + else: + return _utils.fetch._IterableDatasetFetcher( + dataset, auto_collation, collate_fn, drop_last + ) + + +class _InfiniteConstantSampler(Sampler): + r"""Analogous to ``itertools.repeat(None, None)``. + + Used as sampler for :class:`~torch.utils.data.IterableDataset`. + """ + + def __iter__(self): + while True: + yield None + + +def _get_distributed_settings(): + if dist.is_available() and dist.is_initialized(): + return dist.get_world_size(), dist.get_rank() + else: + return 1, 0 + + +def _sharding_worker_init_fn(worker_init_fn, world_size, rank_id, worker_id) -> None: + global_worker_id = worker_id + info = torch.utils.data.get_worker_info() + if info is None: + raise AssertionError("Worker info is None in sharding worker init function") + total_workers = info.num_workers + datapipe = info.dataset + if not isinstance(datapipe, (IterDataPipe, MapDataPipe)): + raise AssertionError( + "datapipe must be an instance of IterDataPipe or MapDataPipe" + ) + # To distribute elements across distributed process evenly, we should shard data on distributed + # processes first then shard on worker processes + total_workers *= world_size + global_worker_id = global_worker_id * world_size + rank_id + # For BC, use default SHARDING_PRIORITIES + torch.utils.data.graph_settings.apply_sharding( + datapipe, total_workers, global_worker_id + ) + if worker_init_fn is not None: + worker_init_fn(worker_id) + + +def _share_dist_seed(generator, pg): + _shared_seed = torch.empty((), dtype=torch.int64).random_(generator=generator) + if isinstance(pg, dist.ProcessGroup): + dist.broadcast(_shared_seed, src=0, group=pg) + return _shared_seed.item() + + +class DataLoader(Generic[_T_co]): + r""" + Data loader combines a dataset and a sampler, and provides an iterable over the given dataset. + + The :class:`~torch.utils.data.DataLoader` supports both map-style and + iterable-style datasets with single- or multi-process loading, customizing + loading order and optional automatic batching (collation) and memory pinning. + + See :py:mod:`torch.utils.data` documentation page for more details. + + Args: + dataset (Dataset): dataset from which to load the data. + batch_size (int, optional): how many samples per batch to load + (default: ``1``). + shuffle (bool, optional): set to ``True`` to have the data reshuffled + at every epoch (default: ``False``). + sampler (Sampler or Iterable, optional): defines the strategy to draw + samples from the dataset. Can be any ``Iterable`` with ``__len__`` + implemented. If specified, :attr:`shuffle` must not be specified. + batch_sampler (Sampler or Iterable, optional): like :attr:`sampler`, but + returns a batch of indices at a time. Mutually exclusive with + :attr:`batch_size`, :attr:`shuffle`, :attr:`sampler`, + and :attr:`drop_last`. + num_workers (int, optional): how many subprocesses to use for data + loading. ``0`` means that the data will be loaded in the main process. + (default: ``0``) + collate_fn (Callable, optional): merges a list of samples to form a + mini-batch of Tensor(s). Used when using batched loading from a + map-style dataset. + pin_memory (bool, optional): If ``True``, the data loader will copy Tensors + into device/CUDA pinned memory before returning them. If your data elements + are a custom type, or your :attr:`collate_fn` returns a batch that is a custom type, + see the example below. + drop_last (bool, optional): set to ``True`` to drop the last incomplete batch, + if the dataset size is not divisible by the batch size. If ``False`` and + the size of dataset is not divisible by the batch size, then the last batch + will be smaller. (default: ``False``) + timeout (numeric, optional): if positive, the timeout value for collecting a batch + from workers. Should always be non-negative. (default: ``0``) + worker_init_fn (Callable, optional): If not ``None``, this will be called on each + worker subprocess with the worker id (an int in ``[0, num_workers - 1]``) as + input, after seeding and before data loading. (default: ``None``) + multiprocessing_context (str or multiprocessing.context.BaseContext, optional): If + ``None``, the default + `multiprocessing context `_ # noqa: D401 + of your operating system will + be used. (default: ``None``) + generator (torch.Generator, optional): If not ``None``, this RNG will be used + by RandomSampler to generate random indexes and multiprocessing to generate + ``base_seed`` for workers. (default: ``None``) + prefetch_factor (int, optional, keyword-only arg): Number of batches loaded + in advance by each worker. ``2`` means there will be a total of + 2 * num_workers batches prefetched across all workers. (default value depends + on the set value for num_workers. If value of num_workers=0 default is ``None``. + Otherwise, if value of ``num_workers > 0`` default is ``2``). + persistent_workers (bool, optional): If ``True``, the data loader will not shut down + the worker processes after a dataset has been consumed once. This allows to + maintain the workers `Dataset` instances alive. (default: ``False``) + pin_memory_device (str, optional): Deprecated, the current :ref:`accelerator` + will be used as the device if ``pin_memory=True``. + in_order (bool, optional): If ``False``, the data loader will not enforce that batches + are returned in a first-in, first-out order. Only applies when ``num_workers > 0``. (default: ``True``) + + + .. warning:: If the ``spawn`` start method is used, :attr:`worker_init_fn` + cannot be an unpicklable object, e.g., a lambda function. See + :ref:`multiprocessing-best-practices` on more details related + to multiprocessing in PyTorch. + + .. warning:: ``len(dataloader)`` heuristic is based on the length of the sampler used. + When :attr:`dataset` is an :class:`~torch.utils.data.IterableDataset`, + it instead returns an estimate based on ``len(dataset) / batch_size``, with proper + rounding depending on :attr:`drop_last`, regardless of multi-process loading + configurations. This represents the best guess PyTorch can make because PyTorch + trusts user :attr:`dataset` code in correctly handling multi-process + loading to avoid duplicate data. + + However, if sharding results in multiple workers having incomplete last batches, + this estimate can still be inaccurate, because (1) an otherwise complete batch can + be broken into multiple ones and (2) more than one batch worth of samples can be + dropped when :attr:`drop_last` is set. Unfortunately, PyTorch can not detect such + cases in general. + + See `Dataset Types`_ for more details on these two types of datasets and how + :class:`~torch.utils.data.IterableDataset` interacts with + `Multi-process data loading`_. + + .. warning:: See :ref:`reproducibility`, and :ref:`dataloader-workers-random-seed`, and + :ref:`data-loading-randomness` notes for random seed related questions. + + .. warning:: Setting `in_order` to `False` can harm reproducibility and may lead to a skewed data + distribution being fed to the trainer in cases with imbalanced data. + """ + + dataset: Dataset[_T_co] + batch_size: int | None + num_workers: int + pin_memory: bool + drop_last: bool + timeout: float + sampler: Sampler | Iterable + pin_memory_device: str + prefetch_factor: int | None + _iterator: _BaseDataLoaderIter | None + __initialized = False + + def __init__( + self, + dataset: Dataset[_T_co], + batch_size: int | None = 1, + shuffle: bool | None = None, + sampler: Sampler | Iterable | None = None, + batch_sampler: Sampler[list] | Iterable[list] | None = None, + num_workers: int = 0, + collate_fn: _collate_fn_t | None = None, + pin_memory: bool = False, + drop_last: bool = False, + timeout: float = 0, + worker_init_fn: _worker_init_fn_t | None = None, + multiprocessing_context=None, + generator=None, + *, + prefetch_factor: int | None = None, + persistent_workers: bool = False, + pin_memory_device: str = "", + in_order: bool = True, + ) -> None: + torch._C._log_api_usage_once("python.data_loader") + + if num_workers < 0: + raise ValueError( + "num_workers option should be non-negative; " + "use num_workers=0 to disable multiprocessing." + ) + + if timeout < 0: + raise ValueError("timeout option should be non-negative") + + if num_workers == 0 and prefetch_factor is not None: + raise ValueError( + "prefetch_factor option could only be specified in multiprocessing." + "let num_workers > 0 to enable multiprocessing, otherwise set prefetch_factor to None." + ) + elif num_workers > 0 and prefetch_factor is None: + prefetch_factor = 2 + elif prefetch_factor is not None and prefetch_factor < 0: + raise ValueError("prefetch_factor option should be non-negative") + + if persistent_workers and num_workers == 0: + raise ValueError("persistent_workers option needs num_workers > 0") + + self.dataset = dataset + self.num_workers = num_workers + self.prefetch_factor = prefetch_factor + self.pin_memory = pin_memory + self.pin_memory_device = pin_memory_device + self.timeout = timeout + self.worker_init_fn = worker_init_fn + self.multiprocessing_context = multiprocessing_context + self.in_order = in_order + + # Adds forward compatibilities so classic DataLoader can work with DataPipes: + # _DataPipeSerializationWrapper container makes it easier to serialize without redefining pickler + if isinstance(self.dataset, IterDataPipe): + self.dataset = _IterDataPipeSerializationWrapper(self.dataset) + elif isinstance(self.dataset, MapDataPipe): + self.dataset = _MapDataPipeSerializationWrapper(self.dataset) + + # Arg-check dataset related before checking samplers because we want to + # tell users that iterable-style datasets are incompatible with custom + # samplers first, so that they don't learn that this combo doesn't work + # after spending time fixing the custom sampler errors. + if isinstance(dataset, IterableDataset): + self._dataset_kind = _DatasetKind.Iterable + # NOTE [ Custom Samplers and IterableDataset ] + # + # `IterableDataset` does not support custom `batch_sampler` or + # `sampler` since the key is irrelevant (unless we support + # generator-style dataset one day...). + # + # For `sampler`, we always create a dummy sampler. This is an + # infinite sampler even when the dataset may have an implemented + # finite `__len__` because in multi-process data loading, naive + # settings will return duplicated data (which may be desired), and + # thus using a sampler with length matching that of dataset will + # cause data lost (you may have duplicates of the first couple + # batches, but never see anything afterwards). Therefore, + # `Iterabledataset` always uses an infinite sampler, an instance of + # `_InfiniteConstantSampler` defined above. + # + # A custom `batch_sampler` essentially only controls the batch size. + # However, it is unclear how useful it would be since an iterable-style + # dataset can handle that within itself. Moreover, it is pointless + # in multi-process data loading as the assignment order of batches + # to workers is an implementation detail so users can not control + # how to batchify each worker's iterable. Thus, we disable this + # option. If this turns out to be useful in future, we can re-enable + # this, and support custom samplers that specify the assignments to + # specific workers. + if isinstance(dataset, IterDataPipe): + if shuffle is not None: + dataset = torch.utils.data.graph_settings.apply_shuffle_settings( + dataset, shuffle=shuffle + ) + # We cannot check `shuffle is not None` here, since previously `shuffle=False` was the default. + elif shuffle not in {False, None}: + raise ValueError( + f"DataLoader with IterableDataset: expected unspecified shuffle option, but got shuffle={shuffle}" + ) + + if sampler is not None: + # See NOTE [ Custom Samplers and IterableDataset ] + raise ValueError( + f"DataLoader with IterableDataset: expected unspecified sampler option, but got sampler={sampler}" + ) + elif batch_sampler is not None: + # See NOTE [ Custom Samplers and IterableDataset ] + raise ValueError( + "DataLoader with IterableDataset: expected unspecified " + f"batch_sampler option, but got batch_sampler={batch_sampler}" + ) + else: + shuffle = bool(shuffle) + self._dataset_kind = _DatasetKind.Map + + if sampler is not None and shuffle: + raise ValueError("sampler option is mutually exclusive with shuffle") + + if batch_sampler is not None: + # auto_collation with custom batch_sampler + if batch_size != 1 or shuffle or sampler is not None or drop_last: + raise ValueError( + "batch_sampler option is mutually exclusive " + "with batch_size, shuffle, sampler, and " + "drop_last" + ) + batch_size = None + drop_last = False + elif batch_size is None: + # no auto_collation + if drop_last: + raise ValueError( + "batch_size=None option disables auto-batching " + "and is mutually exclusive with drop_last" + ) + + if sampler is None: # give default samplers + if self._dataset_kind == _DatasetKind.Iterable: + # See NOTE [ Custom Samplers and IterableDataset ] + sampler = _InfiniteConstantSampler() + else: # map-style + if shuffle: + sampler = RandomSampler(dataset, generator=generator) # type: ignore[arg-type] + else: + sampler = SequentialSampler(dataset) # type: ignore[arg-type] + + if batch_size is not None and batch_sampler is None: + # auto_collation without custom batch_sampler + batch_sampler = BatchSampler(sampler, batch_size, drop_last) + + self.batch_size = batch_size + self.drop_last = drop_last + self.sampler = sampler + self.batch_sampler = batch_sampler + self.generator = generator + + if collate_fn is None: + if self._auto_collation: + collate_fn = _utils.collate.default_collate + else: + collate_fn = _utils.collate.default_convert + + self.collate_fn = collate_fn + self.persistent_workers = persistent_workers + + self.__initialized = True + self._IterableDataset_len_called = ( + None # See NOTE [ IterableDataset and __len__ ] + ) + + self._iterator = None + + self.check_worker_number_rationality() + + def _get_iterator(self) -> _BaseDataLoaderIter: + if self.num_workers == 0: + return _SingleProcessDataLoaderIter(self) + else: + self.check_worker_number_rationality() + return _MultiProcessingDataLoaderIter(self) + + @property + def multiprocessing_context(self): + return self.__multiprocessing_context + + @multiprocessing_context.setter + def multiprocessing_context(self, multiprocessing_context) -> None: + if multiprocessing_context is not None: + if self.num_workers > 0: + if isinstance(multiprocessing_context, str): + valid_start_methods = torch.multiprocessing.get_all_start_methods() + if multiprocessing_context not in valid_start_methods: + raise ValueError( + "multiprocessing_context option " + f"should specify a valid start method in {valid_start_methods!r}, but got " + f"multiprocessing_context={multiprocessing_context!r}" + ) + multiprocessing_context = torch.multiprocessing.get_context( + multiprocessing_context + ) + + if not isinstance( + multiprocessing_context, python_multiprocessing.context.BaseContext + ): + raise TypeError( + "multiprocessing_context option should be a valid context " + "object or a string specifying the start method, but got " + f"multiprocessing_context={multiprocessing_context}" + ) + else: + raise ValueError( + "multiprocessing_context can only be used with " + "multi-process loading (num_workers > 0), but got " + f"num_workers={self.num_workers}" + ) + + self.__multiprocessing_context = multiprocessing_context + + def __setattr__(self, attr, val) -> None: + if self.__initialized and attr in ( + "batch_size", + "batch_sampler", + "sampler", + "drop_last", + "dataset", + "persistent_workers", + ): + raise ValueError( + f"{attr} attribute should not be set after {self.__class__.__name__} is initialized" + ) + + super().__setattr__(attr, val) + + def __iter__(self) -> _BaseDataLoaderIter: + # When using a single worker the returned iterator should be + # created every time to avoid resetting its state + # However, in the case of a multiple workers iterator + # the iterator is only created once in the lifetime of the + # DataLoader object so that workers can be reused + if self.persistent_workers and self.num_workers > 0: + if self._iterator is None: + self._iterator = self._get_iterator() + else: + self._iterator._reset(self) + return self._iterator + else: + return self._get_iterator() + + @property + def _auto_collation(self): + return self.batch_sampler is not None + + @property + def _index_sampler(self): + # The actual sampler used for generating indices for `_DatasetFetcher` + # (see _utils/fetch.py) to read data at each time. This would be + # `.batch_sampler` if in auto-collation mode, and `.sampler` otherwise. + # We can't change `.sampler` and `.batch_sampler` attributes for BC + # reasons. + if self._auto_collation: + return self.batch_sampler + else: + return self.sampler + + def __len__(self) -> int: + if self._dataset_kind == _DatasetKind.Iterable: + # NOTE [ IterableDataset and __len__ ] + # + # For `IterableDataset`, `__len__` could be inaccurate when one naively + # does multi-processing data loading, since the samples will be duplicated. + # However, no real use case should be actually using that behavior, so + # it should count as a user error. We should generally trust user + # code to do the proper thing (e.g., configure each replica differently + # in `__iter__`), and give us the correct `__len__` if they choose to + # implement it (this will still throw if the dataset does not implement + # a `__len__`). + # + # To provide a further warning, we track if `__len__` was called on the + # `DataLoader`, save the returned value in `self._len_called`, and warn + # if the iterator ends up yielding more than this number of samples. + + # Cannot statically verify that dataset is Sized + length = self._IterableDataset_len_called = len(self.dataset) # type: ignore[assignment, arg-type] + if ( + self.batch_size is not None + ): # IterableDataset doesn't allow custom sampler or batch_sampler + from math import ceil + + if self.drop_last: + length = length // self.batch_size + else: + length = ceil(length / self.batch_size) + return length + else: + return len(self._index_sampler) + + def check_worker_number_rationality(self) -> None: + # This function check whether the dataloader's worker number is rational based on + # current system's resource. Current rule is that if the number of workers this + # Dataloader will create is bigger than the number of logical cpus that is allowed to + # use, than we will pop up a warning to let user pay attention. + # + # eg. If current system has 2 physical CPUs with 16 cores each. And each core support 2 + # threads, then the total logical cpus here is 2 * 16 * 2 = 64. Let's say current + # DataLoader process can use half of them which is 32, then the rational max number of + # worker that initiated from this process is 32. + # Now, let's say the created DataLoader has num_works = 40, which is bigger than 32. + # So the warning message is triggered to notify the user to lower the worker number if + # necessary. + # + # + # [Note] Please note that this function respects `cpuset` only when os.sched_getaffinity is + # available (available in most of Linux system, but not OSX and Windows). + # When os.sched_getaffinity is not available, os.cpu_count() is called instead, but + # it doesn't respect cpuset. + # We don't take threading into account since each worker process is single threaded + # at this time. + # + # We don't set any threading flags (eg. OMP_NUM_THREADS, MKL_NUM_THREADS, etc) + # other than `torch.set_num_threads` to 1 in the worker process, if the passing + # in functions use 3rd party modules that rely on those threading flags to determine + # how many thread to create (eg. numpy, etc), then it is caller's responsibility to + # set those flags correctly. + def _create_warning_msg(num_worker_suggest, num_worker_created, cpuset_checked): + suggested_max_worker_msg = ( + ( + ( + "Our suggested max number of worker in current system is {}{}, which is smaller " + "than what this DataLoader is going to create." + ).format( + num_worker_suggest, + ( + "" + if cpuset_checked + else " (`cpuset` is not taken into account)" + ), + ) + ) + if num_worker_suggest is not None + else ( + "DataLoader is not able to compute a suggested max number of worker in current system." + ) + ) + + warn_msg = ( + f"This DataLoader will create {num_worker_created} worker processes in total. {suggested_max_worker_msg} " + "Please be aware that excessive worker creation might get DataLoader running slow or even freeze, " + "lower the worker number to avoid potential slowness/freeze if necessary." + ) + return warn_msg + + if not self.num_workers or self.num_workers == 0: + return + + # try to compute a suggested max number of worker based on system's resource + max_num_worker_suggest = torch._utils.cpu_count() + cpuset_checked = hasattr(os, "sched_getaffinity") + + if max_num_worker_suggest is None or self.num_workers > max_num_worker_suggest: + warnings.warn( + _create_warning_msg( + max_num_worker_suggest, self.num_workers, cpuset_checked + ), + stacklevel=2, + ) + + +class _BaseDataLoaderIter: + def __init__(self, loader: DataLoader) -> None: + self._dataset = loader.dataset + self._shared_seed = None + self._pg = None + if isinstance(self._dataset, IterDataPipe): + if dist.is_available() and dist.is_initialized(): + self._pg = dist.new_group(backend="gloo") + self._shared_seed = _share_dist_seed(loader.generator, self._pg) + shared_rng = torch.Generator() + shared_rng.manual_seed(self._shared_seed) + self._dataset = torch.utils.data.graph_settings.apply_random_seed( + self._dataset, shared_rng + ) + self._dataset_kind = loader._dataset_kind + self._IterableDataset_len_called = loader._IterableDataset_len_called + self._auto_collation = loader._auto_collation + self._drop_last = loader.drop_last + self._index_sampler = loader._index_sampler + self._num_workers = loader.num_workers + ws, rank = _get_distributed_settings() + self._world_size = ws + self._rank = rank + + if loader.pin_memory and loader.pin_memory_device: + warnings.warn( + "pin_memory_device is deprecated, the current accelerator will be used as the device," + f"ignore pin_memory_device='{loader.pin_memory_device}'.", + stacklevel=2, + ) + if loader.pin_memory and not torch.accelerator.is_available(): + warn_msg = ( + "'pin_memory' argument is set as true but no accelerator is found, " + "then device pinned memory won't be used." + ) + warnings.warn(warn_msg, stacklevel=2) + + # Enabling pin_memory in _BaseDataLoaderIter to support identical + # behavior in forked implementations using _BaseDataLoaderIter. + self._pin_memory = loader.pin_memory and torch.accelerator.is_available() + + # Set pin memory device based on the current accelerator. + self._pin_memory_device = ( + acc.type + if self._pin_memory + and (acc := torch.accelerator.current_accelerator()) is not None + else None + ) + + # Currently, pin_memory would raise error on the MPS backend (see + # https://github.com/pytorch/pytorch/issues/86060), so forcibly + # disable pin_memory on MPS. Remove this restriction once pinned + # memory allocation for MPS is fixed. + if self._pin_memory_device == "mps": + self._pin_memory = False + warn_msg = ( + "'pin_memory' argument is set as true but not supported on MPS now, " + "device pinned memory won't be used." + ) + warnings.warn(warn_msg, stacklevel=2) + + self._timeout = loader.timeout + self._collate_fn = loader.collate_fn + self._sampler_iter = iter(self._index_sampler) + self._base_seed = ( + torch.empty((), dtype=torch.int64) + .random_(generator=loader.generator) + .item() + ) + self._persistent_workers = loader.persistent_workers + self._num_yielded = 0 + self._profile_name = f"enumerate(DataLoader)#{self.__class__.__name__}.__next__" + + def __iter__(self) -> Self: + return self + + def _reset(self, loader, first_iter=False) -> None: + self._sampler_iter = iter(self._index_sampler) + self._num_yielded = 0 + self._IterableDataset_len_called = loader._IterableDataset_len_called + if isinstance(self._dataset, IterDataPipe): + self._shared_seed = _share_dist_seed(loader.generator, self._pg) + shared_rng = torch.Generator() + shared_rng.manual_seed(self._shared_seed) + self._dataset = torch.utils.data.graph_settings.apply_random_seed( + self._dataset, shared_rng + ) + + def _next_index(self): + return next(self._sampler_iter) # may raise StopIteration + + def _next_data(self) -> NoReturn: + raise NotImplementedError + + def __next__(self) -> Any: + with torch.autograd.profiler.record_function(self._profile_name): + if self._sampler_iter is None: + # TODO(https://github.com/pytorch/pytorch/issues/76750) + self._reset() # type: ignore[call-arg] + data = self._next_data() + self._num_yielded += 1 + if ( + self._dataset_kind == _DatasetKind.Iterable + and self._IterableDataset_len_called is not None + and self._num_yielded > self._IterableDataset_len_called + ): + warn_msg = ( + f"Length of IterableDataset {self._dataset} was reported to be {self._IterableDataset_len_called}" + f"(when accessing len(dataloader)), but {self._num_yielded} samples have been fetched. " + ) + if self._num_workers > 0: + warn_msg += ( + "For multiprocessing data-loading, this could be caused by not properly configuring the " + "IterableDataset replica at each worker. Please see " + "https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset for examples." + ) + warnings.warn(warn_msg, stacklevel=2) + return data + + def __len__(self) -> int: + return len(self._index_sampler) + + def __getstate__(self): + # TODO: add limited pickling support for sharing an iterator + # across multiple threads for HOGWILD. + # Probably the best way to do this is by moving the sample pushing + # to a separate thread and then just sharing the data queue + # but signalling the end is tricky without a non-blocking API + raise NotImplementedError("{} cannot be pickled", self.__class__.__name__) + + +class _SingleProcessDataLoaderIter(_BaseDataLoaderIter): + def __init__(self, loader) -> None: + super().__init__(loader) + if self._timeout != 0: + raise AssertionError("_SingleProcessDataLoaderIter requires timeout == 0") + if self._num_workers != 0: + raise AssertionError( + "_SingleProcessDataLoaderIter requires num_workers == 0" + ) + + # Adds forward compatibilities so classic DataLoader can work with DataPipes: + # Taking care of distributed sharding + if isinstance(self._dataset, (IterDataPipe, MapDataPipe)): + # For BC, use default SHARDING_PRIORITIES + torch.utils.data.graph_settings.apply_sharding( + self._dataset, self._world_size, self._rank + ) + + self._dataset_fetcher = _DatasetKind.create_fetcher( + self._dataset_kind, + self._dataset, + self._auto_collation, + self._collate_fn, + self._drop_last, + ) + + def _next_data(self): + index = self._next_index() # may raise StopIteration + data = self._dataset_fetcher.fetch(index) # may raise StopIteration + if self._pin_memory: + data = _utils.pin_memory.pin_memory(data, self._pin_memory_device) + return data + + +class _MultiProcessingDataLoaderIter(_BaseDataLoaderIter): + r"""Iterates once over the DataLoader's dataset, as specified by the sampler.""" + + # NOTE [ Data Loader Multiprocessing Shutdown Logic ] + # + # Preliminary: + # + # Our data model looks like this (queues are indicated with curly brackets): + # + # main process || + # | || + # {index_queue} || + # | || + # worker processes || DATA + # | || + # {worker_result_queue} || FLOW + # | || + # pin_memory_thread of main process || DIRECTION + # | || + # {data_queue} || + # | || + # data output \/ + # + # P.S. `worker_result_queue` and `pin_memory_thread` part may be omitted if + # `pin_memory=False`. + # + # + # Terminating multiprocessing logic requires very careful design. In + # particular, we need to make sure that + # + # 1. The iterator gracefully exits the workers when its last reference is + # gone or it is depleted. + # + # In this case, the workers should be gracefully exited because the + # main process may still need to continue to run, and we want cleaning + # up code in the workers to be executed (e.g., releasing GPU memory). + # Naturally, we implement the shutdown logic in `__del__` of + # DataLoaderIterator. + # + # We delay the discussion on the logic in this case until later. + # + # 2. The iterator exits the workers when the loader process and/or worker + # processes exits normally or with error. + # + # We set all workers and `pin_memory_thread` to have `daemon=True`. + # + # You may ask, why can't we make the workers non-daemonic, and + # gracefully exit using the same logic as we have in `__del__` when the + # iterator gets deleted (see 1 above)? + # + # First of all, `__del__` is **not** guaranteed to be called when + # interpreter exits. Even if it is called, by the time it executes, + # many Python core library resources may already be freed, and even + # simple things like acquiring an internal lock of a queue may hang. + # Therefore, in this case, we actually need to prevent `__del__` from + # being executed, and rely on the automatic termination of daemonic + # children. + # + # Thus, we register an `atexit` hook that sets a global flag + # `_utils.python_exit_status`. Since `atexit` hooks are executed in the + # reverse order of registration, we are guaranteed that this flag is + # set before library resources we use are freed (which, at least in + # CPython, is done via an `atexit` handler defined in + # `multiprocessing/util.py` + # https://github.com/python/cpython/blob/c606624af8d4cb3b4a052fb263bb983b3f87585b/Lib/multiprocessing/util.py#L320-L362 + # registered when an object requiring this mechanism is first + # created, e.g., `mp.Queue` + # https://github.com/python/cpython/blob/c606624af8d4cb3b4a052fb263bb983b3f87585b/Lib/multiprocessing/context.py#L100-L103 + # https://github.com/python/cpython/blob/c606624af8d4cb3b4a052fb263bb983b3f87585b/Lib/multiprocessing/queues.py#L29 + # ) + # + # So in `__del__`, we check if `_utils.python_exit_status` is set or + # `None` (freed), and perform no-op if so. + # + # However, simply letting library clean-up codes run can also be bad, + # because such codes (i.e., `multiprocessing.util._exit_function()`) + # include join putting threads for `mp.Queue`, which can be blocking. + # Hence, the main process putting threads are called with + # `cancel_join_thread` at creation. See later section + # [ 3b. A process won't hang when putting into a queue; ] + # for more details. + # + # Here are two example cases where library clean-up codes can run + # before `__del__` is called: + # + # 1. If we hold onto a reference to the iterator, it more often + # than not tries to do `multiprocessing` library cleaning before + # clearing the alive referenced objects (https://github.com/pytorch/pytorch/issues/48666) + # and thus prevents our cleaning-up code to run first. + # + # 2. A similar issue araises when a `DataLoader` is used in a subprocess. + # When a process ends, it shuts the all its daemonic children + # down with a SIGTERM (instead of joining them without a timeout). + # Similarly for threads, but by a different mechanism. This fact, + # together with a few implementation details of multiprocessing, forces + # us to make workers daemonic. All of our problems arise when a + # DataLoader is used in a subprocess, and are caused by multiprocessing + # code which looks more or less like this: + # + # try: + # your_function_using_a_dataloader() + # finally: + # multiprocessing.util._exit_function() + # + # The joining/termination mentioned above happens inside + # `_exit_function()`. Now, if `your_function_using_a_dataloader()` + # throws, the stack trace stored in the exception will prevent the + # frame which uses `DataLoaderIter` to be freed. If the frame has any + # reference to the `DataLoaderIter` (e.g., in a method of the iter), + # its `__del__`, which starts the shutdown procedure, will not be + # called. That, in turn, means that workers aren't notified. Attempting + # to join in `_exit_function` will then result in a hang. + # + # For context, `_exit_function` is also registered as an `atexit` call. + # So it is unclear to me (@ssnl) why this is needed in a finally block. + # The code dates back to 2008 and there is no comment on the original + # PEP 371 or patch https://bugs.python.org/issue3050 (containing both + # the finally block and the `atexit` registration) that explains this. + # + # + # Finally, another choice is to just shutdown workers with logic in 1 + # above whenever we see an error in `next`. This isn't ideal because + # a. It prevents users from using try-catch to resume data loading. + # b. It doesn't prevent hanging if users have references to the + # iterator. + # + # 3. All processes exit if any of them die unexpectedly by fatal signals. + # + # As shown above, the workers are set as daemonic children of the main + # process. However, automatic cleaning-up of such child processes only + # happens if the parent process exits gracefully (e.g., not via fatal + # signals like SIGKILL). So we must ensure that each process will exit + # even the process that should send/receive data to/from it were + # killed, i.e., + # + # a. A process won't hang when getting from a queue. + # + # Even with carefully designed data dependencies (i.e., a `put()` + # always corresponding to a `get()`), hanging on `get()` can still + # happen when data in queue is corrupted (e.g., due to + # `cancel_join_thread` or unexpected exit). + # + # For child exit, we set a timeout whenever we try to get data + # from `data_queue`, and check the workers' status on each timeout + # and error. + # See `_DataLoaderiter._get_batch()` and + # `_DataLoaderiter._try_get_data()` for details. + # + # Additionally, for child exit on non-Windows platforms, we also + # register a SIGCHLD handler (which is supported on Windows) on + # the main process, which checks if any of the workers fail in the + # (Python) handler. This is more efficient and faster in detecting + # worker failures, compared to only using the above mechanism. + # See `DataLoader.cpp` and `_utils/signal_handling.py` for details. + # + # For `.get()` calls where the sender(s) is not the workers, we + # guard them with timeouts, and check the status of the sender + # when timeout happens: + # + in the workers, the `_utils.worker.ManagerWatchdog` class + # checks the status of the main process. + # + if `pin_memory=True`, when getting from `pin_memory_thread`, + # check `pin_memory_thread` status periodically until `.get()` + # returns or see that `pin_memory_thread` died. + # + # b. A process won't hang when putting into a queue; + # + # We use `mp.Queue` which has a separate background thread to put + # objects from an unbounded buffer array. The background thread is + # daemonic and usually automatically joined when the process + # *exits*. + # + # In case that the receiver has ended abruptly while + # reading from the pipe, the join will hang forever. The usual + # solution for this in Python is calling `q.cancel_join_thread`, + # which prevents automatically joining it when finalizing + # (exiting). + # + # Nonetheless, `cancel_join_thread` must only be called when the + # queue is **not** going to be read from or write into by another + # process, because it may hold onto a lock or leave corrupted data + # in the queue, leading other readers/writers to hang. + # + # Hence, + # + For worker processes, we only do so (for their output + # queues, i.e., `worker_result_queue`) before exiting. + # + For `pin_memory_thread`, its output queue `data_queue` is a + # `queue.Queue` that does blocking `put` if the queue is full. + # So there is no above problem, but as a result, in + # `_pin_memory_loop`, we do need to wrap the `put` in a loop + # that breaks not only upon success, but also when the main + # process stops reading, i.e., is shutting down. + # + For loader process, we `cancel_join_thread()` for all + # `_index_queues` because the whole purpose of workers and + # `pin_memory_thread` is to serve the loader process. If + # loader process is already exiting, we don't really care if + # the queues are corrupted. + # + # + # Now let's get back to 1: + # how we gracefully exit the workers when the last reference to the + # iterator is gone. + # + # To achieve this, we implement the following logic along with the design + # choices mentioned above: + # + # `workers_done_event`: + # A `multiprocessing.Event` shared among the main process and all worker + # processes. This is used to signal the workers that the iterator is + # shutting down. After it is set, they will not send processed data to + # queues anymore, and only wait for the final `None` before exiting. + # `done_event` isn't strictly needed. I.e., we can just check for `None` + # from the input queue, but it allows us to skip wasting resources + # processing data if we are already shutting down. + # + # `pin_memory_thread_done_event`: + # A `threading.Event` for a similar purpose to that of + # `workers_done_event`, but is for the `pin_memory_thread`. The reason + # that separate events are needed is that `pin_memory_thread` reads from + # the output queue of the workers. But the workers, upon seeing that + # `workers_done_event` is set, only wants to see the final `None`, and is + # not required to flush all data in the output queue (e.g., it may call + # `cancel_join_thread` on that queue if its `IterableDataset` iterator + # happens to exhaust coincidentally, which is out of the control of the + # main process). Thus, since we will exit `pin_memory_thread` before the + # workers (see below), two separate events are used. + # + # NOTE: In short, the protocol is that the main process will set these + # `done_event`s and then the corresponding processes/threads a `None`, + # and that they may exit at any time after receiving the `None`. + # + # NOTE: Using `None` as the final signal is valid, since normal data will + # always be a 2-tuple with the 1st element being the index of the data + # transferred (different from dataset index/key), and the 2nd being + # either the dataset key or the data sample (depending on which part + # of the data model the queue is at). + # + # [ worker processes ] + # While loader process is alive: + # Get from `index_queue`. + # If get anything else, + # Check `workers_done_event`. + # If set, continue to next iteration + # i.e., keep getting until see the `None`, then exit. + # Otherwise, process data: + # If is fetching from an `IterableDataset` and the iterator + # is exhausted, send an `_IterableDatasetStopIteration` + # object to signal iteration end. The main process, upon + # receiving such an object, will send `None` to this + # worker and not use the corresponding `index_queue` + # anymore. + # If timed out, + # No matter `workers_done_event` is set (still need to see `None`) + # or not, must continue to next iteration. + # (outside loop) + # If `workers_done_event` is set, (this can be False with `IterableDataset`) + # `data_queue.cancel_join_thread()`. (Everything is ending here: + # main process won't read from it; + # other workers will also call + # `cancel_join_thread`.) + # + # [ pin_memory_thread ] + # # No need to check main thread. If this thread is alive, the main loader + # # thread must be alive, because this thread is set as daemonic. + # While `pin_memory_thread_done_event` is not set: + # Get from `worker_result_queue`. + # If timed out, continue to get in the next iteration. + # Otherwise, process data. + # While `pin_memory_thread_done_event` is not set: + # Put processed data to `data_queue` (a `queue.Queue` with blocking put) + # If timed out, continue to put in the next iteration. + # Otherwise, break, i.e., continuing to the out loop. + # + # NOTE: we don't check the status of the main thread because + # 1. if the process is killed by fatal signal, `pin_memory_thread` + # ends. + # 2. in other cases, either the cleaning-up in __del__ or the + # automatic exit of daemonic thread will take care of it. + # This won't busy-wait either because `.get(timeout)` does not + # busy-wait. + # + # [ main process ] + # In the DataLoader Iter's `__del__` + # b. Exit `pin_memory_thread` + # i. Set `pin_memory_thread_done_event`. + # ii Put `None` in `worker_result_queue`. + # iii. Join the `pin_memory_thread`. + # iv. `worker_result_queue.cancel_join_thread()`. + # + # c. Exit the workers. + # i. Set `workers_done_event`. + # ii. Put `None` in each worker's `index_queue`. + # iii. Join the workers. + # iv. Call `.cancel_join_thread()` on each worker's `index_queue`. + # + # NOTE: (c) is better placed after (b) because it may leave corrupted + # data in `worker_result_queue`, which `pin_memory_thread` + # reads from, in which case the `pin_memory_thread` can only + # happen at timing out, which is slow. Nonetheless, same thing + # happens if a worker is killed by signal at unfortunate times, + # but in other cases, we are better off having a non-corrupted + # `worker_result_queue` for `pin_memory_thread`. + # + # NOTE: If `pin_memory=False`, there is no `pin_memory_thread` and (b) + # can be omitted + # + # NB: `done_event`s isn't strictly needed. E.g., we can just check for + # `None` from `index_queue`, but it allows us to skip wasting resources + # processing indices already in `index_queue` if we are already shutting + # down. + + def __init__(self, loader) -> None: + super().__init__(loader) + + self._prefetch_factor = loader.prefetch_factor + self._in_order = loader.in_order + + if self._num_workers <= 0: + raise AssertionError( + "num_workers must be greater than 0 for MultiProcessingDataLoaderIter" + ) + if self._prefetch_factor <= 0: + raise AssertionError( + "prefetch_factor must be greater than 0 for MultiProcessingDataLoaderIter" + ) + + if loader.multiprocessing_context is None: + multiprocessing_context = torch.multiprocessing + else: + multiprocessing_context = loader.multiprocessing_context + + self._worker_init_fn = loader.worker_init_fn + + # Adds forward compatibilities so classic DataLoader can work with DataPipes: + # Additional worker init function will take care of sharding in MP and Distributed + if isinstance(self._dataset, (IterDataPipe, MapDataPipe)): + self._worker_init_fn = functools.partial( + _sharding_worker_init_fn, + self._worker_init_fn, + self._world_size, + self._rank, + ) + + # No certainty which module multiprocessing_context is + self._worker_result_queue = multiprocessing_context.Queue() # type: ignore[var-annotated] + self._worker_pids_set = False + self._shutdown = False + self._workers_done_event = multiprocessing_context.Event() + + self._index_queues = [] + self._workers = [] + for i in range(self._num_workers): + # No certainty which module multiprocessing_context is + index_queue = multiprocessing_context.Queue() # type: ignore[var-annotated] + # Need to `cancel_join_thread` here! + # See sections (2) and (3b) above. + index_queue.cancel_join_thread() + w = multiprocessing_context.Process( + target=_utils.worker._worker_loop, + args=( + self._dataset_kind, + self._dataset, + index_queue, + self._worker_result_queue, + self._workers_done_event, + self._auto_collation, + self._collate_fn, + self._drop_last, + self._base_seed, + self._worker_init_fn, + i, + self._num_workers, + self._persistent_workers, + self._shared_seed, + ), + ) + w.daemon = True + # NB: Process.start() actually take some time as it needs to + # start a process and pass the arguments over via a pipe. + # Therefore, we only add a worker to self._workers list after + # it started, so that we do not call .join() if program dies + # before it starts, and __del__ tries to join but will get: + # AssertionError: can only join a started process. + from pickle import PicklingError + + try: + w.start() + except (TypeError, AttributeError, PicklingError): + warnings.warn( + "Got pickle error when attempting to start a worker Process. " + "This might be because the worker Process arguments are not picklable. " + "Python 3.14+ changed the multiprocessing start method in non-Mac POSIX platforms " + "to 'forkserver', which requires the worker Process arguments to be picklable. " + "You can also try multiprocessing.set_start_method('fork').", + stacklevel=2, + ) + raise + self._index_queues.append(index_queue) + self._workers.append(w) + + if self._pin_memory: + self._pin_memory_thread_done_event = threading.Event() + + # Queue is not type-annotated + self._data_queue = queue.Queue() # type: ignore[var-annotated] + current_device_id = torch.accelerator.current_device_index() + pin_memory_thread = threading.Thread( + target=_utils.pin_memory._pin_memory_loop, + args=( + self._worker_result_queue, + self._data_queue, + current_device_id, + self._pin_memory_thread_done_event, + self._pin_memory_device, + ), + ) + pin_memory_thread.daemon = True + pin_memory_thread.start() + # Similar to workers (see comment above), we only register + # pin_memory_thread once it is started. + self._pin_memory_thread = pin_memory_thread + else: + self._data_queue = self._worker_result_queue # type: ignore[assignment] + + # In some rare cases, persistent workers (daemonic processes) + # would be terminated before `__del__` of iterator is invoked + # when main process exits + # It would cause failure when pin_memory_thread tries to read + # corrupted data from worker_result_queue + # atexit is used to shutdown thread and child processes in the + # right sequence before main process exits + if self._persistent_workers and self._pin_memory: + import atexit + + for w in self._workers: + atexit.register(_MultiProcessingDataLoaderIter._clean_up_worker, w) + + # .pid can be None only before process is spawned (not the case, so ignore) + _utils.signal_handling._set_worker_pids( + id(self), + tuple(w.pid for w in self._workers), # type: ignore[misc] + ) + _utils.signal_handling._set_SIGCHLD_handler() + self._worker_pids_set = True + self._reset(loader, first_iter=True) + + def _reset(self, loader, first_iter=False) -> None: + super()._reset(loader, first_iter) + self._send_idx = 0 # idx of the next task to be sent to workers + self._rcvd_idx = 0 # idx of the next task to be returned in __next__ + # information about data not yet yielded, i.e., tasks w/ indices in range [rcvd_idx, send_idx). + # map: task idx => - (worker_id,) if data isn't fetched (outstanding) + # \ (worker_id, data) if data is already fetched (out-of-order) + self._task_info = {} + self._tasks_outstanding = ( + 0 # always equal to count(v for v in task_info.values() if len(v) == 1) + ) + # A list of booleans representing whether each worker still has work to + # do, i.e., not having exhausted its iterable dataset object. It always + # contains all `True`s if not using an iterable-style dataset + # (i.e., if kind != Iterable). + # Not that this indicates that a worker still has work to do *for this epoch*. + # It does not mean that a worker is dead. In case of `_persistent_workers`, + # the worker will be reset to available in the next epoch. + self._workers_status = [True for i in range(self._num_workers)] + # A list of integers representing how many tasks are outstanding for each worker + # Incremented when a task is dispatched to the worker + # Decremented when that data has been given to the main thread + # Each worker should have at most self._prefetch_factor tasks outstanding + self._workers_num_tasks = [0 for i in range(self._num_workers)] + # Reset the worker queue cycle so it resumes next epoch at worker 0 + self._worker_queue_idx_cycle = itertools.cycle(range(self._num_workers)) + # We resume the prefetching in case it was enabled + if not first_iter: + for idx in range(self._num_workers): + self._index_queues[idx].put( + _utils.worker._ResumeIteration(self._shared_seed) + ) + resume_iteration_cnt = self._num_workers + while resume_iteration_cnt > 0: + return_idx, return_data = self._get_data() + if isinstance(return_idx, _utils.worker._ResumeIteration): + if return_data is not None: + raise AssertionError( + "Expected return_data to be None when resuming iteration" + ) + resume_iteration_cnt -= 1 + # prime the prefetch loop + for _ in range(self._prefetch_factor * self._num_workers): + self._try_put_index() + + def _try_get_data(self, timeout=_utils.MP_STATUS_CHECK_INTERVAL): + # Tries to fetch data from `self._data_queue` once for a given timeout. + # This can also be used as inner loop of fetching without timeout, with + # the sender status as the loop condition. + # + # This raises a `RuntimeError` if any worker died expectedly. This error + # can come from either the SIGCHLD handler in `_utils/signal_handling.py` + # (only for non-Windows platforms), or the manual check below on errors + # and timeouts. + # + # Returns a 2-tuple: + # (bool: whether successfully get data, any: data if successful else None) + try: + data = self._data_queue.get(timeout=timeout) + return (True, data) + except Exception as e: + # At timeout and error, we manually check whether any worker has + # failed. Note that this is the only mechanism for Windows to detect + # worker failures. + failed_workers = [] + for worker_id, w in enumerate(self._workers): + if self._workers_status[worker_id] and not w.is_alive(): + failed_workers.append(w) + self._mark_worker_as_unavailable(worker_id) + if len(failed_workers) > 0: + pids_str = ", ".join(str(w.pid) for w in failed_workers) + raise RuntimeError( + f"DataLoader worker (pid(s) {pids_str}) exited unexpectedly" + ) from e + if isinstance(e, queue.Empty): + return (False, None) + + import errno + import tempfile + + try: + # Raise an exception if we are this close to the FDs limit. + # Apparently, trying to open only one file is not a sufficient + # test. + # See NOTE [ DataLoader on Linux and open files limit ] + fds_limit_margin = 10 + with contextlib.ExitStack() as stack: + for _ in range(fds_limit_margin): + stack.enter_context( + tempfile.NamedTemporaryFile() # pyrefly: ignore [bad-argument-type] + ) + except OSError as e: + if e.errno == errno.EMFILE: + raise RuntimeError( + "Too many open files. Communication with the" + " workers is no longer possible. Please increase the" + " limit using `ulimit -n` in the shell or change the" + " sharing strategy by calling" + " `torch.multiprocessing.set_sharing_strategy('file_system')`" + " at the beginning of your code" + ) from None + raise + + # NOTE [ DataLoader on Linux and open files limit ] + # + # On Linux when DataLoader is used with multiprocessing we pass the data between + # the root process and the workers through SHM files. We remove those files from + # the filesystem as soon as they are created and keep them alive by + # passing around their file descriptors through AF_UNIX sockets. (See + # docs/source/multiprocessing.rst and 'Multiprocessing Technical Notes` in + # the wiki (https://github.com/pytorch/pytorch/wiki).) + # + # This sometimes leads us to exceeding the open files limit. When that happens, + # and the offending file descriptor is coming over a socket, the `socket` Python + # package silently strips the file descriptor from the message, setting only the + # `MSG_CTRUNC` flag (which might be a bit misleading since the manpage says that + # it _indicates that some control data were discarded due to lack of space in + # the buffer for ancillary data_). This might reflect the C implementation of + # AF_UNIX sockets. + # + # This behaviour can be reproduced with the script and instructions at the + # bottom of this note. + # + # When that happens, the standard Python `multiprocessing` (and not + # `torch.multiprocessing`) raises a `RuntimeError: received 0 items of ancdata` + # + # Sometimes, instead of the FD being stripped, you may get an `OSError: + # Too many open files`, both in the script below and in DataLoader. However, + # this is rare and seems to be nondeterministic. + # + # + # #!/usr/bin/env python3 + # import sys + # import socket + # import os + # import array + # import shutil + # import socket + # + # + # if len(sys.argv) != 4: + # print("Usage: ", sys.argv[0], " tmp_dirname iteration (send|recv)") + # sys.exit(1) + # + # if __name__ == '__main__': + # dirname = sys.argv[1] + # sock_path = dirname + "/sock" + # iterations = int(sys.argv[2]) + # def dummy_path(i): + # return dirname + "/" + str(i) + ".dummy" + # + # + # if sys.argv[3] == 'send': + # while not os.path.exists(sock_path): + # pass + # client = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) + # client.connect(sock_path) + # for i in range(iterations): + # fd = os.open(dummy_path(i), os.O_WRONLY | os.O_CREAT) + # ancdata = array.array('i', [fd]) + # msg = bytes([i % 256]) + # print("Sending fd ", fd, " (iteration #", i, ")") + # client.sendmsg([msg], [(socket.SOL_SOCKET, socket.SCM_RIGHTS, ancdata)]) + # + # + # else: + # assert sys.argv[3] == 'recv' + # + # if os.path.exists(dirname): + # raise Exception("Directory exists") + # + # os.mkdir(dirname) + # + # print("Opening socket...") + # server = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM) + # server.bind(sock_path) + # + # print("Listening...") + # for i in range(iterations): + # a = array.array('i') + # msg, ancdata, flags, addr = server.recvmsg(1, socket.CMSG_SPACE(a.itemsize)) + # assert(len(ancdata) == 1) + # cmsg_level, cmsg_type, cmsg_data = ancdata[0] + # a.frombytes(cmsg_data) + # print("Received fd ", a[0], " (iteration #", i, ")") + # + # shutil.rmtree(dirname) + # + # Steps to reproduce: + # + # 1. Run two shells and set lower file descriptor limit in the receiving one: + # (shell1) ulimit -n 1020 + # (shell2) ulimit -n 1022 + # + # 2. Run the script above with the `recv` option in the first shell + # (shell1) ./test_socket.py sock_tmp 1017 recv + # + # 3. Run the script with the `send` option in the second shell: + # (shell2) ./test_socket.py sock_tmp 1017 send + + def _get_data(self): + # Fetches data from `self._data_queue`. + # + # We check workers' status every `MP_STATUS_CHECK_INTERVAL` seconds, + # which we achieve by running `self._try_get_data(timeout=MP_STATUS_CHECK_INTERVAL)` + # in a loop. This is the only mechanism to detect worker failures for + # Windows. For other platforms, a SIGCHLD handler is also used for + # worker failure detection. + # + # If `pin_memory=True`, we also need check if `pin_memory_thread` had + # died at timeouts. + if self._timeout > 0: + success, data = self._try_get_data(self._timeout) + if success: + return data + else: + raise RuntimeError( + f"DataLoader timed out after {self._timeout} seconds" + ) + elif self._pin_memory: + while self._pin_memory_thread.is_alive(): + success, data = self._try_get_data() + if success: + return data + else: + # while condition is false, i.e., pin_memory_thread died. + raise RuntimeError("Pin memory thread exited unexpectedly") + # In this case, `self._data_queue` is a `queue.Queue`,. But we don't + # need to call `.task_done()` because we don't use `.join()`. + else: + while True: + success, data = self._try_get_data() + if success: + return data + + def _next_data(self): + while True: + # If the worker responsible for `self._rcvd_idx` has already ended + # and was unable to fulfill this task (due to exhausting an `IterableDataset`), + # we try to advance `self._rcvd_idx` to find the next valid index. + # + # This part needs to run in the loop because both the `self._get_data()` + # call and `_IterableDatasetStopIteration` check below can mark + # extra worker(s) as dead. + while self._rcvd_idx < self._send_idx: + info = self._task_info.get(self._rcvd_idx, None) + if info: + worker_id = info[0] + if ( + len(info) == 2 or self._workers_status[worker_id] + ): # has data or is still active + break + del self._task_info[self._rcvd_idx] + self._rcvd_idx += 1 + else: + # no valid `self._rcvd_idx` is found (i.e., didn't break) + if not self._persistent_workers: + self._shutdown_workers() + raise StopIteration + + # Now `self._rcvd_idx` is the batch index we want to fetch + + # Check if the next sample has already been generated + if len(self._task_info[self._rcvd_idx]) == 2: + worker_id, data = self._task_info.pop(self._rcvd_idx) + self._rcvd_idx += 1 + return self._process_data(data, worker_id) + + if self._shutdown or self._tasks_outstanding <= 0: + raise AssertionError( + "Invalid iterator state: shutdown or no outstanding tasks when fetching next data" + ) + idx, data = self._get_data() + self._tasks_outstanding -= 1 + if self._dataset_kind == _DatasetKind.Iterable: + # Check for _IterableDatasetStopIteration + if isinstance(data, _utils.worker._IterableDatasetStopIteration): + if self._persistent_workers: + self._workers_status[data.worker_id] = False + else: + self._mark_worker_as_unavailable(data.worker_id) + self._try_put_index() + continue + + if idx != self._rcvd_idx: + if not self._in_order: + # don't store it for later, process now + # delete from self._task_info immediately + # this keeps the object size manageable + worker_id = self._task_info.pop(idx)[0] + return self._process_data(data, worker_id) + # store out-of-order samples + self._task_info[idx] += (data,) + else: + worker_id = self._task_info.pop(idx)[0] + self._rcvd_idx += 1 + return self._process_data(data, worker_id) + + def _try_put_index(self) -> None: + max_tasks = self._prefetch_factor * self._num_workers + if self._tasks_outstanding >= max_tasks: + raise AssertionError( + "Number of outstanding tasks exceeded maximum allowed tasks" + ) + + try: + index = self._next_index() + except StopIteration: + return + for _ in range(self._num_workers): # find the next active worker, if any + worker_queue_idx = next(self._worker_queue_idx_cycle) + if self._workers_status[worker_queue_idx]: + if self._in_order: + break + elif self._workers_num_tasks[worker_queue_idx] < max_tasks // sum( + self._workers_status + ): + # when self._in_order is False, distribute work to a worker if it has capacity + # _workers_status is updated only in this thread, so the sum is guaranteed > 0 + break + else: + # not found (i.e., didn't break) + return + + self._index_queues[worker_queue_idx].put((self._send_idx, index)) # type: ignore[possibly-undefined] + self._task_info[self._send_idx] = (worker_queue_idx,) + self._workers_num_tasks[worker_queue_idx] += 1 + self._tasks_outstanding += 1 + self._send_idx += 1 + + def _process_data(self, data, worker_idx): + self._workers_num_tasks[worker_idx] -= 1 + self._try_put_index() + if isinstance(data, ExceptionWrapper): + data.reraise() + return data + + def _mark_worker_as_unavailable(self, worker_id, shutdown=False) -> None: + # Mark a worker as having finished its work e.g., due to + # exhausting an `IterableDataset`. This should be used only when this + # `_MultiProcessingDataLoaderIter` is going to continue running. + + if ( + not self._workers_status[worker_id] + and not self._persistent_workers + and not shutdown + ): + raise AssertionError( + "Worker status inconsistent when marking worker as unavailable" + ) + + # Signal termination to that specific worker. + q = self._index_queues[worker_id] + # Indicate that no more data will be put on this queue by the current + # process. + q.put(None) + + # Note that we don't actually join the worker here, nor do we remove the + # worker's pid from C side struct because (1) joining may be slow, and + # (2) since we don't join, the worker may still raise error, and we + # prefer capturing those, rather than ignoring them, even though they + # are raised after the worker has finished its job. + # Joining is deferred to `_shutdown_workers`, which it is called when + # all workers finish their jobs (e.g., `IterableDataset` replicas) or + # when this iterator is garbage collected. + + self._workers_status[worker_id] = False + + if self._workers_done_event.is_set() != shutdown: + raise AssertionError( + "_workers_done_event state does not match shutdown flag" + ) + + def _shutdown_workers(self) -> None: + # Called when shutting down this `_MultiProcessingDataLoaderIter`. + # See NOTE [ Data Loader Multiprocessing Shutdown Logic ] for details on + # the logic of this function. + if ( + _utils is None + # pyrefly: ignore [unnecessary-comparison] + or _utils.python_exit_status is True + # pyrefly: ignore [unnecessary-comparison] + or _utils.python_exit_status is None + ): + # See (2) of the note. If Python is shutting down, do no-op. + return + # Normal exit when last reference is gone / iterator is depleted. + # See (1) and the second half of the note. + if not self._shutdown: + self._shutdown = True + try: + # Normal exit when last reference is gone / iterator is depleted. + # See (1) and the second half of the note. + + # Exit `pin_memory_thread` first because exiting workers may leave + # corrupted data in `worker_result_queue` which `pin_memory_thread` + # reads from. + if hasattr(self, "_pin_memory_thread"): + # Use hasattr in case error happens before we set the attribute. + self._pin_memory_thread_done_event.set() + # Send something to pin_memory_thread in case it is waiting + # so that it can wake up and check `pin_memory_thread_done_event` + self._worker_result_queue.put((None, None)) + self._pin_memory_thread.join() + self._worker_result_queue.cancel_join_thread() + self._worker_result_queue.close() + + # Exit workers now. + self._workers_done_event.set() + for worker_id in range(len(self._workers)): + # Get number of workers from `len(self._workers)` instead of + # `self._num_workers` in case we error before starting all + # workers. + # If we are using workers_status with persistent_workers + # we have to shut it down because the worker is paused + if self._persistent_workers or self._workers_status[worker_id]: + self._mark_worker_as_unavailable(worker_id, shutdown=True) + for w in self._workers: + # We should be able to join here, but in case anything went + # wrong, we set a timeout and if the workers fail to join, + # they are killed in the `finally` block. + w.join(timeout=_utils.MP_STATUS_CHECK_INTERVAL) + for q in self._index_queues: + q.cancel_join_thread() + q.close() + finally: + # Even though all this function does is putting into queues that + # we have called `cancel_join_thread` on, weird things can + # happen when a worker is killed by a signal, e.g., hanging in + # `Event.set()`. So we need to guard this with SIGCHLD handler, + # and remove pids from the C side data structure only at the + # end. + # + # FIXME: Unfortunately, for Windows, we are missing a worker + # error detection mechanism here in this function, as it + # doesn't provide a SIGCHLD handler. + if self._worker_pids_set: + _utils.signal_handling._remove_worker_pids(id(self)) + self._worker_pids_set = False + for w in self._workers: + if w.is_alive(): + # Existing mechanisms try to make the workers exit + # peacefully, but in case that we unfortunately reach + # here, which we shouldn't, (e.g., pytorch/pytorch#39570), + # we kill the worker. + w.terminate() + + # staticmethod is used to remove reference to `_MultiProcessingDataLoaderIter` + @staticmethod + def _clean_up_worker(w) -> None: + try: + w.join(timeout=_utils.MP_STATUS_CHECK_INTERVAL) + finally: + if w.is_alive(): + w.terminate() + + def __del__(self) -> None: + self._shutdown_workers() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ac93de335b2d7379246de9cee658dd9eafe1d303 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/__init__.py @@ -0,0 +1 @@ +from torch.utils.data.datapipes import dataframe as dataframe, iter as iter, map as map diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/_decorator.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/_decorator.py new file mode 100644 index 0000000000000000000000000000000000000000..0289668c03abcfc0a8e37bc9ff62365fea3dd1cf --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/_decorator.py @@ -0,0 +1,213 @@ +# mypy: allow-untyped-defs +import inspect +from collections.abc import Callable +from functools import wraps +from typing import Any, get_type_hints + +from torch.utils.data.datapipes._typing import _DataPipeMeta +from torch.utils.data.datapipes.datapipe import IterDataPipe, MapDataPipe + + +###################################################### +# Functional API +###################################################### +class functional_datapipe: + name: str + + def __init__(self, name: str, enable_df_api_tracing=False) -> None: + """ + Define a functional datapipe. + + Args: + enable_df_api_tracing - if set, any returned DataPipe would accept + DataFrames API in tracing mode. + """ + self.name = name + self.enable_df_api_tracing = enable_df_api_tracing + + def __call__(self, cls): + if issubclass(cls, IterDataPipe): + if isinstance(cls, type): # type: ignore[arg-type] + if not isinstance(cls, _DataPipeMeta): + raise TypeError( + "`functional_datapipe` can only decorate IterDataPipe" + ) + # with non_deterministic decorator + else: + if not isinstance(cls, non_deterministic) and not ( + hasattr(cls, "__self__") + and isinstance(cls.__self__, non_deterministic) + ): + raise TypeError( + "`functional_datapipe` can only decorate IterDataPipe" + ) + IterDataPipe.register_datapipe_as_function( + self.name, cls, enable_df_api_tracing=self.enable_df_api_tracing + ) + elif issubclass(cls, MapDataPipe): + MapDataPipe.register_datapipe_as_function(self.name, cls) + + return cls + + +###################################################### +# Determinism +###################################################### +_determinism: bool = False + + +class guaranteed_datapipes_determinism: + prev: bool + + def __init__(self) -> None: + global _determinism + self.prev = _determinism + _determinism = True + + def __enter__(self) -> None: + pass + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + global _determinism + _determinism = self.prev + + +class non_deterministic: + cls: type[IterDataPipe] | None = None + # TODO: Lambda for picking + deterministic_fn: Callable[..., bool] + + def __init__(self, arg: type[IterDataPipe] | Callable[..., bool]) -> None: + # 1. Decorator doesn't have any argument + if isinstance(arg, type): # type: ignore[arg-type] + if not issubclass(arg, IterDataPipe): # type: ignore[arg-type] + raise TypeError( + "Only `IterDataPipe` can be decorated with `non_deterministic`" + f", but {arg.__name__} is found" + ) + self.cls = arg # type: ignore[assignment] + # 2. Decorator has an argument of a function + # This class should behave differently given different inputs. Use this + # function to verify the determinism for each instance. + # When the function returns True, the instance is non-deterministic. Otherwise, + # the instance is a deterministic DataPipe. + elif isinstance(arg, Callable): # type:ignore[arg-type] + self.deterministic_fn = arg + else: + raise TypeError(f"{arg} can not be decorated by non_deterministic") + + def __call__(self, *args, **kwargs): + global _determinism + # Decorate IterDataPipe + if self.cls is not None: + if _determinism: + raise TypeError( + f"{self.cls.__name__} is non-deterministic, but you set 'guaranteed_datapipes_determinism'. " + "You can turn off determinism for this DataPipe if that is acceptable " + "for your application" + ) + return self.cls(*args, **kwargs) # type: ignore[call-arg] + + # Decorate with a functional argument + if not ( + isinstance(args[0], type) and issubclass(args[0], IterDataPipe) # type: ignore[arg-type] + ): + raise TypeError( + f"Only `IterDataPipe` can be decorated, but {args[0].__name__} is found" + ) + self.cls = args[0] + return self.deterministic_wrapper_fn + + def deterministic_wrapper_fn(self, *args, **kwargs) -> IterDataPipe: + res = self.deterministic_fn(*args, **kwargs) + if not isinstance(res, bool): + raise TypeError( + "deterministic_fn of `non_deterministic` decorator is required " + f"to return a boolean value, but {type(res)} is found" + ) + global _determinism + if _determinism and res: + raise TypeError( + f"{self.cls.__name__} is non-deterministic with the inputs, but you set " # type: ignore[union-attr] + "'guaranteed_datapipes_determinism'. You can turn off determinism " + "for this DataPipe if that is acceptable for your application" + ) + return self.cls(*args, **kwargs) # type: ignore[call-arg, misc] + + +###################################################### +# Type validation +###################################################### +# Validate each argument of DataPipe with hint as a subtype of the hint. +def argument_validation(f): + signature = inspect.signature(f) + hints = get_type_hints(f) + + @wraps(f) + def wrapper(*args, **kwargs): + bound = signature.bind(*args, **kwargs) + for argument_name, value in bound.arguments.items(): + if argument_name in hints and isinstance( + hints[argument_name], _DataPipeMeta + ): + hint = hints[argument_name] + if not isinstance(value, IterDataPipe): + raise TypeError( + f"Expected argument '{argument_name}' as a IterDataPipe, but found {type(value)}" + ) + if not value.type.issubtype(hint.type): + raise TypeError( + f"Expected type of argument '{argument_name}' as a subtype of " + f"hint {hint.type}, but found {value.type}" + ) + + return f(*args, **kwargs) + + return wrapper + + +# Default value is True +_runtime_validation_enabled: bool = True + + +class runtime_validation_disabled: + prev: bool + + def __init__(self) -> None: + global _runtime_validation_enabled + self.prev = _runtime_validation_enabled + _runtime_validation_enabled = False + + def __enter__(self) -> None: + pass + + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: + global _runtime_validation_enabled + _runtime_validation_enabled = self.prev + + +# Runtime checking +# Validate output data is subtype of return hint +def runtime_validation(f): + # TODO: + # Can be extended to validate '__getitem__' and nonblocking + if f.__name__ != "__iter__": + raise TypeError( + f"Can not decorate function {f.__name__} with 'runtime_validation'" + ) + + @wraps(f) + def wrapper(self): + global _runtime_validation_enabled + if not _runtime_validation_enabled: + yield from f(self) + else: + it = f(self) + for d in it: + if not self.type.issubtype_of_instance(d): + raise RuntimeError( + f"Expected an instance as subtype of {self.type}, but found {d}({type(d)})" + ) + yield d + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/_hook_iterator.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/_hook_iterator.py new file mode 100644 index 0000000000000000000000000000000000000000..26836168047497000de0003b0489b19e832015bb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/_hook_iterator.py @@ -0,0 +1,279 @@ +# mypy: allow-untyped-defs +import functools +import inspect +from enum import Enum + +import torch + + +class _SnapshotState(Enum): + r""" + These are the snapshotting-related states that IterDataPipes can be in. + + `NotStarted` - allows you to restore a snapshot and create an iterator with reset + `Restored` - cannot restore again, allows you to create an iterator without resetting the DataPipe + `Iterating` - can restore, will reset if you create a new iterator + """ + + NotStarted = 0 + Restored = 1 + Iterating = 2 + + +def _simplify_obj_name(obj) -> str: + """Simplify the display strings of objects for the purpose of rendering within DataPipe error messages.""" + if inspect.isfunction(obj): + return obj.__name__ + else: + return repr(obj) + + +def _strip_datapipe_from_name(name: str) -> str: + return name.replace("IterDataPipe", "").replace("MapDataPipe", "") + + +def _generate_input_args_string(obj): + """Generate a string for the input arguments of an object.""" + signature = inspect.signature(obj.__class__) + input_param_names = set(signature.parameters.keys()) + result = [] + for name, value in inspect.getmembers(obj): + if name in input_param_names: + result.append((name, _simplify_obj_name(value))) + return ", ".join([f"{name}={value}" for name, value in result]) + + +def _generate_iterdatapipe_msg(datapipe, simplify_dp_name: bool = False): + output_string = ( + f"{datapipe.__class__.__name__}({_generate_input_args_string(datapipe)})" + ) + if simplify_dp_name: + output_string = _strip_datapipe_from_name(output_string) + return output_string + + +def _gen_invalid_iterdatapipe_msg(datapipe) -> str: + return ( + "This iterator has been invalidated because another iterator has been created " + f"from the same IterDataPipe: {_generate_iterdatapipe_msg(datapipe)}\n" + "This may be caused multiple references to the same IterDataPipe. We recommend " + "using `.fork()` if that is necessary." + ) + + +_feedback_msg = ( + "\nFor feedback regarding this single iterator per IterDataPipe constraint, feel free " + "to comment on this issue: https://github.com/pytorch/data/issues/45." +) + + +def _check_iterator_valid(datapipe, iterator_id, next_method_exists=False) -> None: + r""" + Given an instance of a DataPipe and an iterator ID, check if the IDs match, and if not, raises an exception. + + In the case of ChildDataPipe, the ID gets compared to the one stored in `main_datapipe` as well. + """ + if next_method_exists: + # This is the case where `IterDataPipe` has both `__iter__` and `__next__`. + # The `_valid_iterator_id` should either be never set (`None`), or set by at most one + # iterator (`0`). Otherwise, it means there are multiple iterators. + if datapipe._valid_iterator_id is not None and datapipe._valid_iterator_id != 0: + extra_msg = "\nNote that this exception is raised inside your IterDataPipe's a `__next__` method" + raise RuntimeError( + _gen_invalid_iterdatapipe_msg(datapipe) + extra_msg + _feedback_msg + ) + elif ( + hasattr(datapipe, "_is_child_datapipe") and datapipe._is_child_datapipe is True + ): + if hasattr(datapipe, "_check_valid_iterator_id"): + if not datapipe._check_valid_iterator_id(iterator_id): + raise RuntimeError( + "This iterator has been invalidated, because a new iterator has been created " + f"from one of the ChildDataPipes of " + f"{_generate_iterdatapipe_msg(datapipe.main_datapipe)}." + + _feedback_msg + ) + else: + raise RuntimeError( + "ChildDataPipe must have method `_check_valid_iterator_id`." + ) + elif datapipe._valid_iterator_id != iterator_id: + raise RuntimeError(_gen_invalid_iterdatapipe_msg(datapipe) + _feedback_msg) + + +def _set_datapipe_valid_iterator_id(datapipe): + """Given a DataPipe, updates its valid iterator ID and reset the DataPipe.""" + if hasattr(datapipe, "_is_child_datapipe") and datapipe._is_child_datapipe is True: + if hasattr(datapipe, "_set_main_datapipe_valid_iterator_id"): + datapipe._set_main_datapipe_valid_iterator_id() # reset() is called within this method when appropriate + else: + raise RuntimeError( + "ChildDataPipe must have method `_set_main_datapipe_valid_iterator_id`." + ) + else: + if datapipe._valid_iterator_id is None: + datapipe._valid_iterator_id = 0 + else: + datapipe._valid_iterator_id += 1 + datapipe.reset() + return datapipe._valid_iterator_id + + +def hook_iterator(namespace) -> None: + r""" + Define a hook that is applied to all `__iter__` of metaclass `_DataPipeMeta`. + + This is done for the purpose of profiling and checking if an iterator is still valid. + """ + + def profiler_record_fn_context(datapipe): + if not hasattr(datapipe, "_profile_name"): + datapipe._profile_name = _generate_iterdatapipe_msg( + datapipe, simplify_dp_name=True + ) + return torch.autograd.profiler.record_function(datapipe._profile_name) + + class IteratorDecorator: + r""" + Wrap the iterator and modifying its `__next__` method. + + This decorator is applied to DataPipes of which `__iter__` method is NOT a generator function. + Those `__iter__` method commonly returns `self` but not necessarily. + """ + + def __init__(self, iterator, datapipe, iterator_id, has_next_method) -> None: + self.iterator = iterator + self.datapipe = datapipe + self.iterator_id = iterator_id + self._profiler_enabled = torch.autograd._profiler_enabled() + # Check if `__iter__` returns `self` and `DataPipe` has `__next__` + self.self_and_has_next_method = ( + self.iterator is self.datapipe and has_next_method + ) + + def __iter__(self): + return self + + def _get_next(self): + """Return next with logic related to iterator validity, profiler, and incrementation of samples yielded.""" + _check_iterator_valid(self.datapipe, self.iterator_id) + result = next(self.iterator) + if not self.self_and_has_next_method: + self.datapipe._number_of_samples_yielded += 1 + return result + + def __next__(self): + # TODO: Add try-except to in-place reduce traceback from the Exception + # See: https://github.com/pytorch/data/issues/284 + if self._profiler_enabled: + with profiler_record_fn_context(self.datapipe): + return self._get_next() + else: # Decided against using `contextlib.nullcontext` for performance reasons + return self._get_next() + + def __getattr__(self, name): + return getattr(self.iterator, name) + + func = namespace["__iter__"] + + # ``__iter__`` of IterDataPipe is a generator function + if inspect.isgeneratorfunction(func): + + @functools.wraps(func) + def wrap_generator(*args, **kwargs): + gen = func(*args, **kwargs) + datapipe = args[0] + if datapipe._fast_forward_iterator: + it = datapipe._fast_forward_iterator + datapipe._fast_forward_iterator = None + datapipe._snapshot_state = _SnapshotState.Iterating + while True: + try: + yield next(it) + except StopIteration: + return + iterator_id = _set_datapipe_valid_iterator_id( + datapipe + ) # This ID is tied to each created iterator + _profiler_enabled = torch.autograd._profiler_enabled() + try: + if _profiler_enabled: + with profiler_record_fn_context(datapipe): + response = gen.send(None) + else: + response = gen.send(None) + + while True: + datapipe._number_of_samples_yielded += 1 + request = yield response + # Pass through here every time `__next__` is called + if _profiler_enabled: + with profiler_record_fn_context(datapipe): + _check_iterator_valid(datapipe, iterator_id) + response = gen.send(request) + else: # Decided against using `contextlib.nullcontext` for performance reasons + _check_iterator_valid(datapipe, iterator_id) + response = gen.send(request) + except StopIteration: + return + except Exception as e: + # TODO: Simplify the traceback message to skip over `response = gen.send(None)` + # Part of https://github.com/pytorch/data/issues/284 + datapipe = args[0] + msg = "thrown by __iter__ of" + single_iterator_msg = "single iterator per IterDataPipe constraint" + if hasattr(e.args, "__len__"): + full_msg = f"{msg} {datapipe.__class__.__name__}({_generate_input_args_string(datapipe)})" + if len(e.args) == 0 or not isinstance( + e.args[0], str + ): # If an exception message doesn't exist + e.args = (f"\nThis exception is {full_msg}",) + elif msg not in e.args[0] and single_iterator_msg not in e.args[0]: + e.args = ( + e.args[0] + f"\nThis exception is {full_msg}", + ) + e.args[1:] + raise + + namespace["__iter__"] = wrap_generator + else: # ``__iter__`` of IterDataPipe is NOT a generator function + # IterDataPipe is an iterator with both ``__iter__`` and ``__next__`` + # And ``__iter__`` may or may not return `self` + if "__next__" in namespace: # If `__next__` exists, put a wrapper around it + next_func = namespace["__next__"] + + @functools.wraps(next_func) + def wrap_next(*args, **kwargs): + datapipe = args[0] + if torch.autograd._profiler_enabled(): + with profiler_record_fn_context(datapipe): + result = next_func(*args, **kwargs) + else: + result = next_func(*args, **kwargs) + datapipe._number_of_samples_yielded += 1 + return result + + namespace["__next__"] = wrap_next + + # Note that if the `__next__` and `__iter__` do something completely unrelated. It may cause issue but + # the user will be violating the iterator protocol. Potential issue: + # 1. Valid iterator ID may not update or checked properly + # 2. The number of samples yielded will be miscounted + + # Regardless if `__next__` exists or not, `__iter__` needs a wrapper to track the number of valid iterators + @functools.wraps(func) + def wrap_iter(*args, **kwargs): + iter_ret = func(*args, **kwargs) + datapipe = args[0] + datapipe._snapshot_state = _SnapshotState.Iterating + if datapipe._fast_forward_iterator: + iter_ret = datapipe._fast_forward_iterator + datapipe._fast_forward_iterator = None + return iter_ret + iterator_id = _set_datapipe_valid_iterator_id( + datapipe + ) # This ID is tied to each created iterator + return IteratorDecorator( + iter_ret, datapipe, iterator_id, "__next__" in namespace + ) + + namespace["__iter__"] = wrap_iter diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/_typing.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/_typing.py new file mode 100644 index 0000000000000000000000000000000000000000..d44f2f6d67c8302c92d111ab1228c2d6310dab9f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/_typing.py @@ -0,0 +1,483 @@ +# mypy: allow-untyped-defs +# Taking reference from official Python typing +# https://github.com/python/cpython/blob/master/Lib/typing.py + +import collections +import functools +import numbers +import sys + +# Please check [Note: TypeMeta and TypeAlias] +# In case of metaclass conflict due to ABCMeta or _ProtocolMeta +# For Python 3.9, only Protocol in typing uses metaclass +from abc import ABCMeta +from collections.abc import Iterator + +# TODO: Use TypeAlias when Python 3.6 is deprecated +from typing import ( + _eval_type, # pyrefly: ignore [missing-module-attribute] + _GenericAlias, # pyrefly: ignore [missing-module-attribute] + _tp_cache, # pyrefly: ignore [missing-module-attribute] + _type_check, # pyrefly: ignore [missing-module-attribute] + _type_repr, + Any, + ForwardRef, + Generic, + get_type_hints, + TypeVar, + Union, +) + +from torch.utils.data.datapipes._hook_iterator import _SnapshotState, hook_iterator + + +class GenericMeta(ABCMeta): # type: ignore[no-redef] + pass + + +class Integer(numbers.Integral): + pass + + +class Boolean(numbers.Integral): + pass + + +# Python 'type' object is not subscriptable +# Tuple[int, List, dict] -> valid +# tuple[int, list, dict] -> invalid +# Map Python 'type' to abstract base class +TYPE2ABC = { + bool: Boolean, + int: Integer, + float: numbers.Real, + complex: numbers.Complex, + dict: dict, + list: list, + set: set, + tuple: tuple, + None: type(None), +} + + +def issubtype(left, right, recursive=True): + r""" + Check if the left-side type is a subtype of the right-side type. + + If any of type is a composite type like `Union` and `TypeVar` with + bounds, it would be expanded into a list of types and check all + of left-side types are subtypes of either one from right-side types. + """ + left = TYPE2ABC.get(left, left) + right = TYPE2ABC.get(right, right) + + if right is Any or left == right: + return True + + if isinstance(right, _GenericAlias): + if getattr(right, "__origin__", None) is Generic: + return True + + if right is type(None): + return False + + # Right-side type + constraints = _decompose_type(right) + + if len(constraints) == 0 or Any in constraints: + return True + + if left is Any: + return False + + # Left-side type + variants = _decompose_type(left) + + # all() will return True for empty variants + if len(variants) == 0: + return False + + return all( + _issubtype_with_constraints(variant, constraints, recursive) + for variant in variants + ) + + +def _decompose_type(t, to_list=True): + if isinstance(t, TypeVar): + if t.__bound__ is not None: + ts = [t.__bound__] + else: + # For T_co, __constraints__ is () + ts = list(t.__constraints__) + elif hasattr(t, "__origin__") and t.__origin__ == Union: + ts = t.__args__ + else: + if not to_list: + return None + ts = [t] + # Ignored: Generator has incompatible item type "object"; expected "Type[Any]" + ts = [TYPE2ABC.get(_t, _t) for _t in ts] # type: ignore[misc] + return ts + + +def _issubtype_with_constraints(variant, constraints, recursive=True): + r""" + Check if the variant is a subtype of either one from constraints. + + For composite types like `Union` and `TypeVar` with bounds, they + would be expanded for testing. + """ + if variant in constraints: + return True + + # [Note: Subtype for Union and TypeVar] + # Python typing is able to flatten Union[Union[...]] or Union[TypeVar]. + # But it couldn't flatten the following scenarios: + # - Union[int, TypeVar[Union[...]]] + # - TypeVar[TypeVar[...]] + # So, variant and each constraint may be a TypeVar or a Union. + # In these cases, all of inner types from the variant are required to be + # extracted and verified as a subtype of any constraint. And, all of + # inner types from any constraint being a TypeVar or a Union are + # also required to be extracted and verified if the variant belongs to + # any of them. + + # Variant + vs = _decompose_type(variant, to_list=False) + + # Variant is TypeVar or Union + if vs is not None: + return all(_issubtype_with_constraints(v, constraints, recursive) for v in vs) + + # Variant is not TypeVar or Union + if hasattr(variant, "__origin__") and variant.__origin__ is not None: + v_origin = variant.__origin__ + # In Python-3.9 typing library untyped generics do not have args + v_args = getattr(variant, "__args__", None) + else: + v_origin = variant + v_args = None + + # Constraints + for constraint in constraints: + cs = _decompose_type(constraint, to_list=False) + + # Constraint is TypeVar or Union + if cs is not None: + if _issubtype_with_constraints(variant, cs, recursive): + return True + # Constraint is not TypeVar or Union + else: + # __origin__ can be None for plain list, tuple, ... in Python 3.6 + if hasattr(constraint, "__origin__") and constraint.__origin__ is not None: + c_origin = constraint.__origin__ + if v_origin == c_origin: + if not recursive: + return True + # In Python-3.9 typing library untyped generics do not have args + c_args = getattr(constraint, "__args__", None) + if c_args is None or len(c_args) == 0: + return True + if ( + v_args is not None + and len(v_args) == len(c_args) + and all( + issubtype(v_arg, c_arg) + for v_arg, c_arg in zip(v_args, c_args, strict=True) + ) + ): + return True + # Tuple[int] -> Tuple + else: + if v_origin == constraint: + return True + + return False + + +def issubinstance(data, data_type): + if not issubtype(type(data), data_type, recursive=False): + return False + + # In Python-3.9 typing library __args__ attribute is not defined for untyped generics + dt_args = getattr(data_type, "__args__", None) + if isinstance(data, tuple): + if dt_args is None or len(dt_args) == 0: + return True + if len(dt_args) != len(data): + return False + return all(issubinstance(d, t) for d, t in zip(data, dt_args, strict=True)) + elif isinstance(data, (list, set)): + if dt_args is None or len(dt_args) == 0: + return True + t = dt_args[0] + return all(issubinstance(d, t) for d in data) + elif isinstance(data, dict): + if dt_args is None or len(dt_args) == 0: + return True + kt, vt = dt_args + return all( + issubinstance(k, kt) and issubinstance(v, vt) for k, v in data.items() + ) + + return True + + +# [Note: TypeMeta and TypeAlias] +# In order to keep compatibility for Python 3.6, use Meta for the typing. +# TODO: When PyTorch drops the support for Python 3.6, it can be converted +# into the Alias system and using `__class_getitem__` for DataPipe. The +# typing system will gain benefit of performance and resolving metaclass +# conflicts as elaborated in https://www.python.org/dev/peps/pep-0560/ + + +class _DataPipeType: + r"""Save type annotation in `param`.""" + + def __init__(self, param) -> None: + self.param = param + + def __repr__(self) -> str: + return _type_repr(self.param) + + def __eq__(self, other): + if isinstance(other, _DataPipeType): + return self.param == other.param + return NotImplemented + + def __hash__(self): + return hash(self.param) + + def issubtype(self, other): + if isinstance(other.param, _GenericAlias): + if getattr(other.param, "__origin__", None) is Generic: + return True + if isinstance(other, _DataPipeType): + return issubtype(self.param, other.param) + if isinstance(other, type): + return issubtype(self.param, other) + raise TypeError(f"Expected '_DataPipeType' or 'type', but found {type(other)}") + + def issubtype_of_instance(self, other): + return issubinstance(other, self.param) + + +# Default type for DataPipe without annotation +_T_co = TypeVar("_T_co", covariant=True) +# pyrefly: ignore [invalid-annotation] +_DEFAULT_TYPE = _DataPipeType(Generic[_T_co]) + + +class _DataPipeMeta(GenericMeta): + r""" + Metaclass for `DataPipe`. + + Add `type` attribute and `__init_subclass__` based on the type, and validate the return hint of `__iter__`. + + Note that there is subclass `_IterDataPipeMeta` specifically for `IterDataPipe`. + """ + + type: _DataPipeType + + def __new__(cls, name, bases, namespace, **kwargs): + return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] + + # TODO: the statements below are not reachable by design as there is a bug and typing is low priority for now. + cls.__origin__ = None + if "type" in namespace: + return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] + + namespace["__type_class__"] = False + # For plain derived class without annotation + for base in bases: + if isinstance(base, _DataPipeMeta): + return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] + + namespace.update( + {"type": _DEFAULT_TYPE, "__init_subclass__": _dp_init_subclass} + ) + return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] + + def __init__(self, name, bases, namespace, **kwargs) -> None: + super().__init__(name, bases, namespace, **kwargs) # type: ignore[call-overload] + + # TODO: Fix isinstance bug + @_tp_cache + def _getitem_(self, params): + if params is None: + raise TypeError(f"{self.__name__}[t]: t can not be None") + if isinstance(params, str): + params = ForwardRef(params) + if not isinstance(params, tuple): + params = (params,) + + msg = f"{self.__name__}[t]: t must be a type" + params = tuple(_type_check(p, msg) for p in params) + + if isinstance(self.type.param, _GenericAlias): + orig = getattr(self.type.param, "__origin__", None) + if isinstance(orig, type) and orig is not Generic: + p = self.type.param[params] # type: ignore[index] + t = _DataPipeType(p) + l = len(str(self.type)) + 2 + name = self.__name__[:-l] + name = name + "[" + str(t) + "]" + bases = (self,) + self.__bases__ + return self.__class__( + name, + bases, + { + "__init_subclass__": _dp_init_subclass, + "type": t, + "__type_class__": True, + }, + ) + + if len(params) > 1: + raise TypeError( + f"Too many parameters for {self} actual {len(params)}, expected 1" + ) + + t = _DataPipeType(params[0]) + + if not t.issubtype(self.type): + raise TypeError( + f"Can not subclass a DataPipe[{t}] from DataPipe[{self.type}]" + ) + + # Types are equal, fast path for inheritance + if self.type == t: + return self + + name = self.__name__ + "[" + str(t) + "]" + bases = (self,) + self.__bases__ + + return self.__class__( + name, + bases, + {"__init_subclass__": _dp_init_subclass, "__type_class__": True, "type": t}, + ) + + # TODO: Fix isinstance bug + def _eq_(self, other): + if not isinstance(other, _DataPipeMeta): + return NotImplemented + if self.__origin__ is None or other.__origin__ is None: # type: ignore[has-type] + return self is other + return ( + self.__origin__ == other.__origin__ # type: ignore[has-type] + and self.type == other.type + ) + + # TODO: Fix isinstance bug + def _hash_(self): + return hash((self.__name__, self.type)) + + +class _IterDataPipeMeta(_DataPipeMeta): + r""" + Metaclass for `IterDataPipe` and inherits from `_DataPipeMeta`. + + Add various functions for behaviors specific to `IterDataPipe`. + """ + + def __new__(cls, name, bases, namespace, **kwargs): + if "reset" in namespace: + reset_func = namespace["reset"] + + @functools.wraps(reset_func) + def conditional_reset(*args, **kwargs) -> None: + r""" + Only execute DataPipe's `reset()` method if `_SnapshotState` is `Iterating` or `NotStarted`. + + This allows recently restored DataPipe to preserve its restored state during the initial `__iter__` call. + """ + datapipe = args[0] + if datapipe._snapshot_state in ( + _SnapshotState.Iterating, + _SnapshotState.NotStarted, + ): + # Reset `NotStarted` is necessary because the `source_datapipe` of a DataPipe might have + # already begun iterating. + datapipe._number_of_samples_yielded = 0 + datapipe._fast_forward_iterator = None + reset_func(*args, **kwargs) + datapipe._snapshot_state = _SnapshotState.Iterating + + namespace["reset"] = conditional_reset + + if "__iter__" in namespace: + hook_iterator(namespace) + return super().__new__(cls, name, bases, namespace, **kwargs) # type: ignore[call-overload] + + +def _dp_init_subclass(sub_cls, *args, **kwargs) -> None: + # Add function for datapipe instance to reinforce the type + sub_cls.reinforce_type = reinforce_type + + # TODO: + # - add global switch for type checking at compile-time + + # Ignore internal type class + if getattr(sub_cls, "__type_class__", False): + return + + # Check if the string type is valid + if isinstance(sub_cls.type.param, ForwardRef): + base_globals = sys.modules[sub_cls.__module__].__dict__ + try: + param = _eval_type(sub_cls.type.param, base_globals, locals()) + sub_cls.type.param = param + except TypeError as e: + raise TypeError( + f"{sub_cls.type.param.__forward_arg__} is not supported by Python typing" + ) from e + + if "__iter__" in sub_cls.__dict__: + iter_fn = sub_cls.__dict__["__iter__"] + hints = get_type_hints(iter_fn) + if "return" in hints: + return_hint = hints["return"] + # Plain Return Hint for Python 3.6 + if return_hint == Iterator: + return + if not ( + hasattr(return_hint, "__origin__") + and ( + return_hint.__origin__ == Iterator + or return_hint.__origin__ == collections.abc.Iterator + ) + ): + raise TypeError( + "Expected 'Iterator' as the return annotation for `__iter__` of {}" + ", but found {}".format( + sub_cls.__name__, _type_repr(hints["return"]) + ) + ) + data_type = return_hint.__args__[0] + if not issubtype(data_type, sub_cls.type.param): + raise TypeError( + f"Expected return type of '__iter__' as a subtype of {sub_cls.type}," + f" but found {_type_repr(data_type)} for {sub_cls.__name__}" + ) + + +def reinforce_type(self, expected_type): + r""" + Reinforce the type for DataPipe instance. + + And the 'expected_type' is required to be a subtype of the original type + hint to restrict the type requirement of DataPipe instance. + """ + if isinstance(expected_type, tuple): + expected_type = tuple[expected_type] # type: ignore[valid-type] + _type_check(expected_type, msg="'expected_type' must be a type") + + if not issubtype(expected_type, self.type.param): + raise TypeError( + f"Expected 'expected_type' as subtype of {self.type}, but found {_type_repr(expected_type)}" + ) + + self.type = _DataPipeType(expected_type) + return self diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f7f4b7dcb414c205614a694ccaa02961e45e9b3e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/__init__.py @@ -0,0 +1,12 @@ +from torch.utils.data.datapipes.dataframe.dataframes import ( + CaptureDataFrame, + DFIterDataPipe, +) +from torch.utils.data.datapipes.dataframe.datapipes import DataFramesAsTuplesPipe + + +__all__ = ["CaptureDataFrame", "DFIterDataPipe", "DataFramesAsTuplesPipe"] + +# Please keep this list sorted +if __all__ != sorted(__all__): + raise AssertionError("__all__ is not sorted") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/dataframe_wrapper.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/dataframe_wrapper.py new file mode 100644 index 0000000000000000000000000000000000000000..9cfc5c268a17455cb30d981036996a716d0cc668 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/dataframe_wrapper.py @@ -0,0 +1,128 @@ +# mypy: allow-untyped-defs +from typing import Any + + +_pandas: Any = None +_WITH_PANDAS: bool | None = None + + +def _try_import_pandas() -> bool: + try: + import pandas # type: ignore[import] + + global _pandas + _pandas = pandas + return True + except ImportError: + return False + + +# pandas used only for prototyping, will be shortly replaced with TorchArrow +def _with_pandas() -> bool: + global _WITH_PANDAS + if _WITH_PANDAS is None: + _WITH_PANDAS = _try_import_pandas() + return _WITH_PANDAS + + +class PandasWrapper: + @classmethod + def create_dataframe(cls, data, columns): + if not _with_pandas(): + raise RuntimeError("DataFrames prototype requires pandas to function") + return _pandas.DataFrame(data, columns=columns) # type: ignore[union-attr] + + @classmethod + def is_dataframe(cls, data): + if not _with_pandas(): + return False + return isinstance(data, _pandas.core.frame.DataFrame) # type: ignore[union-attr] + + @classmethod + def is_column(cls, data): + if not _with_pandas(): + return False + return isinstance(data, _pandas.core.series.Series) # type: ignore[union-attr] + + @classmethod + def iterate(cls, data): + if not _with_pandas(): + raise RuntimeError("DataFrames prototype requires pandas to function") + yield from data.itertuples(index=False) + + @classmethod + def concat(cls, buffer): + if not _with_pandas(): + raise RuntimeError("DataFrames prototype requires pandas to function") + return _pandas.concat(buffer) # type: ignore[union-attr] + + @classmethod + def get_item(cls, data, idx): + if not _with_pandas(): + raise RuntimeError("DataFrames prototype requires pandas to function") + return data[idx : idx + 1] + + @classmethod + def get_len(cls, df): + if not _with_pandas(): + raise RuntimeError("DataFrames prototype requires pandas to function") + return len(df.index) + + @classmethod + def get_columns(cls, df): + if not _with_pandas(): + raise RuntimeError("DataFrames prototype requires pandas to function") + return list(df.columns.values.tolist()) + + +# When you build own implementation just override it with dataframe_wrapper.set_df_wrapper(new_wrapper_class) +default_wrapper = PandasWrapper + + +def get_df_wrapper(): + return default_wrapper + + +def set_df_wrapper(wrapper) -> None: + global default_wrapper + default_wrapper = wrapper + + +def create_dataframe(data, columns=None): + wrapper = get_df_wrapper() + return wrapper.create_dataframe(data, columns) + + +def is_dataframe(data): + wrapper = get_df_wrapper() + return wrapper.is_dataframe(data) + + +def get_columns(data): + wrapper = get_df_wrapper() + return wrapper.get_columns(data) + + +def is_column(data): + wrapper = get_df_wrapper() + return wrapper.is_column(data) + + +def concat(buffer): + wrapper = get_df_wrapper() + return wrapper.concat(buffer) + + +def iterate(data): + wrapper = get_df_wrapper() + return wrapper.iterate(data) + + +def get_item(data, idx): + wrapper = get_df_wrapper() + return wrapper.get_item(data, idx) + + +def get_len(df): + wrapper = get_df_wrapper() + return wrapper.get_len(df) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/dataframes.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/dataframes.py new file mode 100644 index 0000000000000000000000000000000000000000..d2608796dc75e2ae2c441b77de5bcbc760919459 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/dataframes.py @@ -0,0 +1,468 @@ +# mypy: allow-untyped-defs +from typing import Any, NoReturn + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.dataframe.structures import DataChunkDF +from torch.utils.data.datapipes.datapipe import DFIterDataPipe, IterDataPipe + + +# TODO(VitalyFedyunin): Add error when two different traces get combined + +__all__ = [ + "Capture", + "CaptureA", + "CaptureAdd", + "CaptureCall", + "CaptureControl", + "CaptureDataFrame", + "CaptureDataFrameWithDataPipeOps", + "CaptureF", + "CaptureGetAttr", + "CaptureGetItem", + "CaptureInitial", + "CaptureLikeMock", + "CaptureMul", + "CaptureSetItem", + "CaptureSub", + "CaptureVariable", + "CaptureVariableAssign", + "DataFrameTracer", + "DataFrameTracedOps", + "disable_capture", + "get_val", +] + + +def disable_capture() -> None: + CaptureControl.disabled = True + + +class CaptureControl: + disabled = False + + +class DataFrameTracedOps(DFIterDataPipe): + def __init__(self, source_datapipe, output_var) -> None: + super().__init__() + self.source_datapipe = source_datapipe + self.output_var = output_var + + def __iter__(self): + for item in self.source_datapipe: + yield self.output_var.apply_ops(item) + + +# TODO(VitalyFedyunin): Extract this list from the DFIterDataPipe registered functions +DATAPIPES_OPS = [ + "_dataframes_as_tuples", + "groupby", + "_dataframes_filter", + "map", + "to_datapipe", + "shuffle", + "concat", + "batch", + "_dataframes_per_row", + "_dataframes_concat", + "_dataframes_shuffle", +] + +UNIMPLEMENTED_ATTR = ["__deepcopy__", "__setstate__", "is_shardable", "apply_sharding"] + + +class Capture: + # TODO: All operations are shared across entire InitialCapture, need to figure out what if we join two captures + + def __init__(self, schema_df=None) -> None: + self.ctx = {"operations": [], "variables": [], "schema_df": schema_df} + + def __str__(self) -> str: + return self._ops_str() + + def _ops_str(self): + res = "" + for op in self.ctx["operations"]: + if len(res) > 0: + res += "\n" + res += str(op) + return res + + def __getstate__(self): + # TODO(VitalyFedyunin): Currently can't pickle (why?) + self.ctx["schema_df"] = None + for var in self.ctx["variables"]: + var.calculated_value = None + state = {} + for item in self.__dict__: + state[item] = getattr(self, item) + return state + + def __setstate__(self, state): + for k, v in state.items(): + setattr(self, k, v) + + def __getattr__(self, attrname): + if attrname == "kwarg" or attrname == "kwargs": + raise RuntimeError("no kwargs!") + if attrname == "__deepcopy__": + raise AttributeError + result = CaptureGetAttr(self, attrname, ctx=self.ctx) + return result + + def __getitem__(self, key): + return CaptureGetItem(self, key, ctx=self.ctx) + + def __setitem__(self, key, value) -> None: + self.ctx["operations"].append(CaptureSetItem(self, key, value, ctx=self.ctx)) + + def __add__(self, add_val): + res = CaptureAdd(self, add_val, ctx=self.ctx) + var = CaptureVariable(res, ctx=self.ctx) + self.ctx["operations"].append( + CaptureVariableAssign(variable=var, value=res, ctx=self.ctx) + ) + return var + + def __sub__(self, add_val): + res = CaptureSub(self, add_val, ctx=self.ctx) + var = CaptureVariable(res, ctx=self.ctx) + self.ctx["operations"].append( + CaptureVariableAssign(variable=var, value=res, ctx=self.ctx) + ) + return var + + def __mul__(self, add_val): + res = CaptureMul(self, add_val, ctx=self.ctx) + var = CaptureVariable(res, ctx=self.ctx) + t = CaptureVariableAssign(variable=var, value=res, ctx=self.ctx) + self.ctx["operations"].append(t) + return var + + def _is_context_empty(self): + return len(self.ctx["operations"]) == 0 and len(self.ctx["variables"]) == 0 + + def apply_ops_2(self, dataframe) -> None: + # TODO(VitalyFedyunin): Make this calculation thread safe (as currently it updates pointer) + self.ctx["variables"][0].calculated_value = dataframe + for op in self.ctx["operations"]: + op.execute() + + @property + def columns(self): + self.apply_ops_2(self.ctx["schema_df"]) + value = self.execute() + return value.columns + + # TODO(VitalyFedyunin): Add tests + # TODO(VitalyFedyunin): Need to join context if one of them are empty because we used capture + + def __call__(self, *args, **kwargs): + # TODO: Check if args or kwargs have more than one different context + if self._is_context_empty(): + # TODO: Allow CaptureA to take context from mock + for arg in args: + if isinstance(arg, Capture) and not arg._is_context_empty(): + self.ctx = arg.ctx + break + if self._is_context_empty(): + for k, v in kwargs.items(): + if isinstance(k, Capture) and not k._is_context_empty(): + self.ctx = k.ctx + break + if isinstance(v, Capture) and not v._is_context_empty(): + self.ctx = v.ctx + break + + res = CaptureCall(self, ctx=self.ctx, args=args, kwargs=kwargs) + var = CaptureVariable(None, ctx=self.ctx) + t = CaptureVariableAssign(ctx=self.ctx, variable=var, value=res) + self.ctx["operations"].append(t) + return var + + +class CaptureF(Capture): + def __init__(self, ctx=None, **kwargs) -> None: + super().__init__() + if ctx is None: + self.ctx = {"operations": [], "variables": []} + else: + self.ctx = ctx + self.kwargs = kwargs + + +class CaptureA(CaptureF): + def __str__(self) -> str: + return f"{self.kwargs['name']}" + + def execute(self): + value = self.kwargs["real_attribute"] + return value + + +class CaptureLikeMock: + def __init__(self, name) -> None: + import unittest.mock as mock + + # TODO(VitalyFedyunin): Do not use private function here, copy own implementation instead. + get_target, attribute = mock._get_target(name) # type: ignore[attr-defined] + self.get_target = get_target + self.attribute = attribute + self.name = name + + def __enter__(self): + self.save = getattr(self.get_target(), self.attribute) + capt = CaptureA(name=self.name, real_attribute=self.save) + setattr(self.get_target(), self.attribute, capt) + + def __exit__(self, *exc_info): + setattr(self.get_target(), self.attribute, self.save) + + +class CaptureCall(Capture): + def __init__(self, callable, ctx=None, **kwargs) -> None: + super().__init__() + if ctx is None: + self.ctx = {"operations": [], "variables": []} + else: + self.ctx = ctx + self.kwargs = kwargs + self.callable = callable + + def __str__(self) -> str: + return "{callable}({args},{kwargs})".format( + callable=self.callable, **self.kwargs + ) + + def execute(self): + # TODO: VitalyFedyunin execute kwargs and maybe nested structures + executed_args = [] + for arg in self.kwargs["args"]: + if isinstance(arg, Capture): + executed_args.append(arg.execute()) + else: + executed_args.append(arg) + left = get_val(self.callable) + return left(*executed_args, **self.kwargs["kwargs"]) + + +class CaptureVariableAssign(CaptureF): + def __str__(self) -> str: + variable = self.kwargs["variable"] + value = self.kwargs["value"] + return f"{variable} = {value}" + + def execute(self) -> None: + self.kwargs["variable"].calculated_value = self.kwargs["value"].execute() + + +class CaptureVariable(Capture): + # TODO(VitalyFedyunin): This should be atomic and thread safe + names_idx = 0 + + def __init__(self, value, ctx) -> None: + super().__init__() + if CaptureControl.disabled: + raise RuntimeError("Attempting to create capture variable with capture off") + self.ctx = ctx + self.value = value + self.name = f"var_{CaptureVariable.names_idx}" + CaptureVariable.names_idx += 1 + self.ctx["variables"].append(self) + + def __str__(self) -> str: + return self.name + + def execute(self): + return self.calculated_value + + def apply_ops(self, dataframe): + # TODO(VitalyFedyunin): Make this calculation thread safe (as currently it updates pointer) + self.ctx["variables"][0].calculated_value = dataframe + for op in self.ctx["operations"]: + op.execute() + return self.calculated_value + + +class CaptureGetItem(Capture): + def __init__(self, left, key, ctx) -> None: + super().__init__() + self.ctx = ctx + self.left = left + self.key = key + + def __str__(self) -> str: + return f"{self.left}[{get_val(self.key)}]" + + def execute(self): + left = self.left.execute() + return left[self.key] + + +class CaptureSetItem(Capture): + def __init__(self, left, key, value, ctx) -> None: + super().__init__() + self.ctx = ctx + self.left = left + self.key = key + self.value = value + + def __str__(self) -> str: + return f"{self.left}[{get_val(self.key)}] = {self.value}" + + def execute(self) -> None: + left = self.left.execute() + value = self.value.execute() + left[self.key] = value + + +class CaptureAdd(Capture): + def __init__(self, left, right, ctx) -> None: + super().__init__() + self.ctx = ctx + self.left = left + self.right = right + + def __str__(self) -> str: + return f"{self.left} + {self.right}" + + def execute(self): + return get_val(self.left) + get_val(self.right) + + +class CaptureMul(Capture): + def __init__(self, left, right, ctx) -> None: + super().__init__() + self.ctx = ctx + self.left = left + self.right = right + + def __str__(self) -> str: + return f"{self.left} * {self.right}" + + def execute(self): + return get_val(self.left) * get_val(self.right) + + +class CaptureSub(Capture): + def __init__(self, left, right, ctx) -> None: + super().__init__() + self.ctx = ctx + self.left = left + self.right = right + + def __str__(self) -> str: + return f"{self.left} - {self.right}" + + def execute(self): + return get_val(self.left) - get_val(self.right) + + +class CaptureGetAttr(Capture): + def __init__(self, src, name, ctx) -> None: + super().__init__() + self.ctx = ctx + self.src = src + self.name = name + + def __str__(self) -> str: + return f"{self.src}.{self.name}" + + def execute(self): + val = get_val(self.src) + return getattr(val, self.name) + + +def get_val(capture): + if isinstance(capture, Capture): + return capture.execute() + elif isinstance(capture, str): + return f'"{capture}"' + else: + return capture + + +class CaptureInitial(CaptureVariable): + def __init__(self, schema_df=None) -> None: + # pyrefly: ignore [bad-assignment] + new_ctx: dict[str, list[Any]] = { + "operations": [], + "variables": [], + "schema_df": schema_df, + } + super().__init__(None, new_ctx) + self.name = f"input_{self.name}" + + +class CaptureDataFrame(CaptureInitial): + pass + + +class CaptureDataFrameWithDataPipeOps(CaptureDataFrame): + def as_datapipe(self): + return DataFrameTracedOps(self.ctx["variables"][0].source_datapipe, self) + + def raw_iterator(self): + return self.as_datapipe().__iter__() + + def __iter__(self): + return iter(self._dataframes_as_tuples()) + + def batch(self, batch_size=10, drop_last: bool = False, wrapper_class=DataChunkDF): + dp = self._dataframes_per_row()._dataframes_concat(batch_size) + dp = dp.as_datapipe().batch(1, drop_last=drop_last, wrapper_class=wrapper_class) + dp._dp_contains_dataframe = True + return dp + + def groupby( + self, + group_key_fn, + *, + buffer_size=10000, + group_size=None, + guaranteed_group_size=None, + drop_remaining=False, + ): + dp = self._dataframes_per_row() + dp = dp.as_datapipe().groupby( + group_key_fn, + buffer_size=buffer_size, + group_size=group_size, + guaranteed_group_size=guaranteed_group_size, + drop_remaining=drop_remaining, + ) + return dp + + def shuffle(self, *args, **kwargs): + return self._dataframes_shuffle(*args, **kwargs) + + def filter(self, *args, **kwargs): + return self._dataframes_filter(*args, **kwargs) + + def collate(self, *args, **kwargs) -> NoReturn: + raise RuntimeError("Can't collate unbatched DataFrames stream") + + def __getattr__(self, attrname): # ? + if attrname in UNIMPLEMENTED_ATTR: + raise AttributeError("Attempting to get ", attrname) + if attrname in DATAPIPES_OPS: + return (self.as_datapipe()).__getattr__(attrname) + return super().__getattr__(attrname) + + +@functional_datapipe("trace_as_dataframe") +class DataFrameTracer(CaptureDataFrameWithDataPipeOps, IterDataPipe): # type: ignore[misc] + source_datapipe: Any | None = None + + # TODO(VitalyFedyunin): Must implement all special functions of datapipes + + def set_shuffle_settings(self, *args, **kwargs) -> None: + pass + + def is_shardable(self) -> bool: + return False + + def __init__(self, source_datapipe, schema_df=None) -> None: + self.source_datapipe = source_datapipe + if schema_df is None: + schema_df = next(iter(self.source_datapipe)) + super().__init__(schema_df=schema_df) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/datapipes.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/datapipes.py new file mode 100644 index 0000000000000000000000000000000000000000..2ea8b6fe7b836e8fe8ac6d543e4598b72a3ebb02 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/datapipes.py @@ -0,0 +1,138 @@ +# mypy: allow-untyped-defs +import random +from typing import Any + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper +from torch.utils.data.datapipes.datapipe import DFIterDataPipe, IterDataPipe + + +__all__ = [ + "ConcatDataFramesPipe", + "DataFramesAsTuplesPipe", + "ExampleAggregateAsDataFrames", + "FilterDataFramesPipe", + "PerRowDataFramesPipe", + "ShuffleDataFramesPipe", +] + + +@functional_datapipe("_dataframes_as_tuples") +class DataFramesAsTuplesPipe(IterDataPipe): + def __init__(self, source_datapipe) -> None: + super().__init__() + self.source_datapipe = source_datapipe + + def __iter__(self): + for df in self.source_datapipe: + # for record in df.to_records(index=False): + yield from df_wrapper.iterate(df) + + +@functional_datapipe("_dataframes_per_row", enable_df_api_tracing=True) +class PerRowDataFramesPipe(DFIterDataPipe): + def __init__(self, source_datapipe) -> None: + self.source_datapipe = source_datapipe + + def __iter__(self): + for df in self.source_datapipe: + # TODO(VitalyFedyunin): Replacing with TorchArrow only API, as we are dropping pandas as followup + for i in range(len(df)): + yield df[i : i + 1] + + +@functional_datapipe("_dataframes_concat", enable_df_api_tracing=True) +class ConcatDataFramesPipe(DFIterDataPipe): + def __init__(self, source_datapipe, batch=3) -> None: + self.source_datapipe = source_datapipe + self.n_batch = batch + + def __iter__(self): + buffer = [] + for df in self.source_datapipe: + buffer.append(df) + if len(buffer) == self.n_batch: + yield df_wrapper.concat(buffer) + buffer = [] + if buffer: + yield df_wrapper.concat(buffer) + + +@functional_datapipe("_dataframes_shuffle", enable_df_api_tracing=True) +class ShuffleDataFramesPipe(DFIterDataPipe): + def __init__(self, source_datapipe) -> None: + self.source_datapipe = source_datapipe + + def __iter__(self): + size = None + all_buffer: list[Any] = [] + for df in self.source_datapipe: + if size is None: + size = df_wrapper.get_len(df) + all_buffer.extend( + df_wrapper.get_item(df, i) for i in range(df_wrapper.get_len(df)) + ) + random.shuffle(all_buffer) + buffer = [] + for df in all_buffer: + buffer.append(df) + if len(buffer) == size: + yield df_wrapper.concat(buffer) + buffer = [] + if buffer: + yield df_wrapper.concat(buffer) + + +@functional_datapipe("_dataframes_filter", enable_df_api_tracing=True) +class FilterDataFramesPipe(DFIterDataPipe): + def __init__(self, source_datapipe, filter_fn) -> None: + self.source_datapipe = source_datapipe + self.filter_fn = filter_fn + + def __iter__(self): + size = None + all_buffer = [] + filter_res = [] + # pyrefly: ignore [bad-assignment] + for df in self.source_datapipe: + if size is None: + size = len(df.index) + for i in range(len(df.index)): + all_buffer.append(df[i : i + 1]) + filter_res.append(self.filter_fn(df.iloc[i])) + + buffer = [] + for df, res in zip(all_buffer, filter_res, strict=True): + if res: + buffer.append(df) + if len(buffer) == size: + yield df_wrapper.concat(buffer) + buffer = [] + if buffer: + yield df_wrapper.concat(buffer) + + +@functional_datapipe("_to_dataframes_pipe", enable_df_api_tracing=True) +class ExampleAggregateAsDataFrames(DFIterDataPipe): + def __init__(self, source_datapipe, dataframe_size=10, columns=None) -> None: + self.source_datapipe = source_datapipe + self.columns = columns + self.dataframe_size = dataframe_size + + def _as_list(self, item): + try: + return list(item) + except ( + Exception + ): # TODO(VitalyFedyunin): Replace with better iterable exception + return [item] + + def __iter__(self): + aggregate = [] + for item in self.source_datapipe: + aggregate.append(self._as_list(item)) + if len(aggregate) == self.dataframe_size: + yield df_wrapper.create_dataframe(aggregate, columns=self.columns) + aggregate = [] + if len(aggregate) > 0: + yield df_wrapper.create_dataframe(aggregate, columns=self.columns) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/structures.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/structures.py new file mode 100644 index 0000000000000000000000000000000000000000..26b4c33db03cc584f223444c07730ef67f4495e7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/dataframe/structures.py @@ -0,0 +1,22 @@ +from collections.abc import Iterator +from typing import Any + +from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper +from torch.utils.data.datapipes.datapipe import DataChunk + + +__all__ = ["DataChunkDF"] + + +class DataChunkDF(DataChunk): + """DataChunkDF iterating over individual items inside of DataFrame containers, to access DataFrames user `raw_iterator`.""" + + def __iter__(self) -> Iterator[Any]: + for df in self.items: + yield from df_wrapper.iterate(df) + + def __len__(self) -> int: + total_len = 0 + for df in self.items: + total_len += df_wrapper.get_len(df) + return total_len diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/datapipe.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/datapipe.py new file mode 100644 index 0000000000000000000000000000000000000000..51c1689008530b1ec4e78c9c921fd9aa6629ecfb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/datapipe.py @@ -0,0 +1,427 @@ +import functools +import pickle +from collections.abc import Callable, Iterable, Iterator +from typing import TypeVar + +from torch.utils._import_utils import import_dill +from torch.utils.data.datapipes._hook_iterator import _SnapshotState +from torch.utils.data.datapipes._typing import _DataPipeMeta, _IterDataPipeMeta +from torch.utils.data.datapipes.utils.common import ( + _deprecation_warning, + _iter_deprecated_functional_names, + _map_deprecated_functional_names, +) +from torch.utils.data.dataset import Dataset, IterableDataset + + +dill = import_dill() +HAS_DILL = dill is not None + +__all__ = [ + "DataChunk", + "DFIterDataPipe", + "IterDataPipe", + "MapDataPipe", +] + + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) + +UNTRACABLE_DATAFRAME_PIPES = [ + "batch", # As it returns DataChunks + "groupby", # As it returns DataChunks + "_dataframes_as_tuples", # As it unpacks DF + "trace_as_dataframe", # As it used to mark DF for tracing +] + + +class DataChunk(list[_T]): + def __init__(self, items: Iterable[_T]) -> None: + items = list(items) + super().__init__(items) + self.items = items + + def as_str(self, indent: str = "") -> str: + return indent + "[" + ", ".join(str(i) for i in iter(self)) + "]" + + def __iter__(self) -> Iterator[_T]: + yield from super().__iter__() + + def raw_iterator(self) -> Iterator[_T]: + yield from self.items + + +class IterDataPipe(IterableDataset[_T_co], metaclass=_IterDataPipeMeta): + r""" + Iterable-style DataPipe. + + All DataPipes that represent an iterable of data samples should subclass this. + This style of DataPipes is particularly useful when data come from a stream, or + when the number of samples is too large to fit them all in memory. ``IterDataPipe`` is lazily initialized and its + elements are computed only when ``next()`` is called on the iterator of an ``IterDataPipe``. + + All subclasses should overwrite :meth:`__iter__`, which would return an + iterator of samples in this DataPipe. Calling ``__iter__`` of an ``IterDataPipe`` automatically invokes its + method ``reset()``, which by default performs no operation. When writing a custom ``IterDataPipe``, users should + override ``reset()`` if necessary. The common usages include resetting buffers, pointers, + and various state variables within the custom ``IterDataPipe``. + + Note: + Only `one` iterator can be valid for each ``IterDataPipe`` at a time, + and the creation a second iterator will invalidate the first one. This constraint is necessary because + some ``IterDataPipe`` have internal buffers, whose states can become invalid if there are multiple iterators. + The code example below presents details on how this constraint looks in practice. + If you have any feedback related to this constraint, please see `GitHub IterDataPipe Single Iterator Issue`_. + + These DataPipes can be invoked in two ways, using the class constructor or applying their + functional form onto an existing ``IterDataPipe`` (recommended, available to most but not all DataPipes). + You can chain multiple `IterDataPipe` together to form a pipeline that will perform multiple + operations in succession. + + .. _GitHub IterDataPipe Single Iterator Issue: + https://github.com/pytorch/data/issues/45 + + Note: + When a subclass is used with :class:`~torch.utils.data.DataLoader`, each + item in the DataPipe will be yielded from the :class:`~torch.utils.data.DataLoader` + iterator. When :attr:`num_workers > 0`, each worker process will have a + different copy of the DataPipe object, so it is often desired to configure + each copy independently to avoid having duplicate data returned from the + workers. :func:`~torch.utils.data.get_worker_info`, when called in a worker + process, returns information about the worker. It can be used in either the + dataset's :meth:`__iter__` method or the :class:`~torch.utils.data.DataLoader` 's + :attr:`worker_init_fn` option to modify each copy's behavior. + + Examples: + General Usage: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper, Mapper + >>> dp = IterableWrapper(range(10)) + >>> map_dp_1 = Mapper(dp, lambda x: x + 1) # Using class constructor + >>> map_dp_2 = dp.map( + ... lambda x: x + 1 + ... ) # Using functional form (recommended) + >>> list(map_dp_1) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + >>> list(map_dp_2) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + >>> filter_dp = map_dp_1.filter(lambda x: x % 2 == 0) + >>> list(filter_dp) + [2, 4, 6, 8, 10] + Single Iterator Constraint Example: + >>> from torchdata.datapipes.iter import IterableWrapper, Mapper + >>> source_dp = IterableWrapper(range(10)) + >>> it1 = iter(source_dp) + >>> list(it1) + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + >>> it1 = iter(source_dp) + >>> it2 = iter( + ... source_dp + ... ) # The creation of a new iterator invalidates `it1` + >>> next(it2) + 0 + >>> next(it1) # Further usage of `it1` will raise a `RunTimeError` + """ + + functions: dict[str, Callable] = {} + reduce_ex_hook: Callable | None = None + getstate_hook: Callable | None = None + str_hook: Callable | None = None + repr_hook: Callable | None = None + _valid_iterator_id: int | None = None + _number_of_samples_yielded: int = 0 + _snapshot_state: _SnapshotState = _SnapshotState.NotStarted + _fast_forward_iterator: Iterator | None = None + + def __iter__(self) -> Iterator[_T_co]: + # pyrefly: ignore [bad-return] + return self + + def __getattr__(self, attribute_name): + if attribute_name in IterDataPipe.functions: + if attribute_name in _iter_deprecated_functional_names: + kwargs = _iter_deprecated_functional_names[attribute_name] + _deprecation_warning(**kwargs) + f = IterDataPipe.functions[attribute_name] + function = functools.partial(f, self) + functools.update_wrapper(wrapper=function, wrapped=f, assigned=("__doc__",)) + return function + else: + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{attribute_name}" + ) + + @classmethod + def register_function(cls, function_name, function) -> None: + cls.functions[function_name] = function + + @classmethod + def register_datapipe_as_function( + cls, function_name, cls_to_register, enable_df_api_tracing=False + ) -> None: + if function_name in cls.functions: + raise Exception( # noqa: TRY002 + f"Unable to add DataPipe function name {function_name} as it is already taken" + ) + + def class_function(cls, enable_df_api_tracing, source_dp, *args, **kwargs): + result_pipe = cls(source_dp, *args, **kwargs) + if isinstance(result_pipe, IterDataPipe): + if enable_df_api_tracing or isinstance(source_dp, DFIterDataPipe): + if function_name not in UNTRACABLE_DATAFRAME_PIPES: + result_pipe = result_pipe.trace_as_dataframe() + + return result_pipe + + function = functools.partial( + class_function, cls_to_register, enable_df_api_tracing + ) + functools.update_wrapper( + wrapper=function, wrapped=cls_to_register, assigned=("__doc__",) + ) + cls.functions[function_name] = function + + def __getstate__(self): + """ + Serialize `lambda` functions when `dill` is available. + + If this doesn't cover your custom DataPipe's use case, consider writing custom methods for + `__getstate__` and `__setstate__`, or use `pickle.dumps` for serialization. + """ + state = self.__dict__ + if IterDataPipe.getstate_hook is not None: + return IterDataPipe.getstate_hook(state) + return state + + def __reduce_ex__(self, *args, **kwargs): + if IterDataPipe.reduce_ex_hook is not None: + try: + return IterDataPipe.reduce_ex_hook(self) + except NotImplementedError: + pass + return super().__reduce_ex__(*args, **kwargs) + + @classmethod + def set_getstate_hook(cls, hook_fn) -> None: + if IterDataPipe.getstate_hook is not None and hook_fn is not None: + raise RuntimeError("Attempt to override existing getstate_hook") + IterDataPipe.getstate_hook = hook_fn + + @classmethod + def set_reduce_ex_hook(cls, hook_fn) -> None: + if IterDataPipe.reduce_ex_hook is not None and hook_fn is not None: + raise RuntimeError("Attempt to override existing reduce_ex_hook") + IterDataPipe.reduce_ex_hook = hook_fn + + def __repr__(self) -> str: + if self.repr_hook is not None: + return self.repr_hook(self) + # Instead of showing , return the class name + return str(self.__class__.__qualname__) + + def __str__(self) -> str: + if self.str_hook is not None: + return self.str_hook(self) + # Instead of showing , return the class name + return str(self.__class__.__qualname__) + + def __dir__(self): + # for auto-completion in a REPL (e.g. Jupyter notebook) + return list(super().__dir__()) + list(self.functions.keys()) + + def reset(self) -> None: + r""" + Reset the `IterDataPipe` to the initial state. + + By default, no-op. For subclasses of `IterDataPipe`, depending on their functionalities, + they may want to override this method with implementations that + may clear the buffers and reset pointers of the DataPipe. + The `reset` method is always called when `__iter__` is called as part of `hook_iterator`. + """ + + +class DFIterDataPipe(IterDataPipe): + def _is_dfpipe(self) -> bool: + return True + + +class MapDataPipe(Dataset[_T_co], metaclass=_DataPipeMeta): + r""" + Map-style DataPipe. + + All datasets that represent a map from keys to data samples should subclass this. + Subclasses should overwrite :meth:`__getitem__`, supporting fetching a + data sample for a given, unique key. Subclasses can also optionally overwrite + :meth:`__len__`, which is expected to return the size of the dataset by many + :class:`~torch.utils.data.Sampler` implementations and the default options + of :class:`~torch.utils.data.DataLoader`. + + These DataPipes can be invoked in two ways, using the class constructor or applying their + functional form onto an existing `MapDataPipe` (recommend, available to most but not all DataPipes). + + Note: + :class:`~torch.utils.data.DataLoader` by default constructs an index + sampler that yields integral indices. To make it work with a map-style + DataPipe with non-integral indices/keys, a custom sampler must be provided. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper, Mapper + >>> dp = SequenceWrapper(range(10)) + >>> map_dp_1 = dp.map(lambda x: x + 1) # Using functional form (recommended) + >>> list(map_dp_1) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + >>> map_dp_2 = Mapper(dp, lambda x: x + 1) # Using class constructor + >>> list(map_dp_2) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + >>> batch_dp = map_dp_1.batch(batch_size=2) + >>> list(batch_dp) + [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]] + """ + + functions: dict[str, Callable] = {} + reduce_ex_hook: Callable | None = None + getstate_hook: Callable | None = None + str_hook: Callable | None = None + repr_hook: Callable | None = None + + def __getattr__(self, attribute_name): + if attribute_name in MapDataPipe.functions: + if attribute_name in _map_deprecated_functional_names: + kwargs = _map_deprecated_functional_names[attribute_name] + _deprecation_warning(**kwargs) + f = MapDataPipe.functions[attribute_name] + function = functools.partial(f, self) + functools.update_wrapper(wrapper=function, wrapped=f, assigned=("__doc__",)) + return function + else: + raise AttributeError( + f"'{self.__class__.__name__}' object has no attribute '{attribute_name}" + ) + + @classmethod + def register_function(cls, function_name, function) -> None: + cls.functions[function_name] = function + + @classmethod + def register_datapipe_as_function(cls, function_name, cls_to_register) -> None: + if function_name in cls.functions: + raise Exception( # noqa: TRY002 + f"Unable to add DataPipe function name {function_name} as it is already taken" + ) + + def class_function(cls, source_dp, *args, **kwargs): + result_pipe = cls(source_dp, *args, **kwargs) + return result_pipe + + function = functools.partial(class_function, cls_to_register) + functools.update_wrapper( + wrapper=function, wrapped=cls_to_register, assigned=("__doc__",) + ) + cls.functions[function_name] = function + + def __getstate__(self): + """ + Serialize `lambda` functions when `dill` is available. + + If this doesn't cover your custom DataPipe's use case, consider writing custom methods for + `__getstate__` and `__setstate__`, or use `pickle.dumps` for serialization. + """ + state = self.__dict__ + if MapDataPipe.getstate_hook is not None: + return MapDataPipe.getstate_hook(state) + return state + + def __reduce_ex__(self, *args, **kwargs): + if MapDataPipe.reduce_ex_hook is not None: + try: + return MapDataPipe.reduce_ex_hook(self) + except NotImplementedError: + pass + return super().__reduce_ex__(*args, **kwargs) + + @classmethod + def set_getstate_hook(cls, hook_fn) -> None: + if MapDataPipe.getstate_hook is not None and hook_fn is not None: + raise RuntimeError("Attempt to override existing getstate_hook") + MapDataPipe.getstate_hook = hook_fn + + @classmethod + def set_reduce_ex_hook(cls, hook_fn) -> None: + if MapDataPipe.reduce_ex_hook is not None and hook_fn is not None: + raise RuntimeError("Attempt to override existing reduce_ex_hook") + MapDataPipe.reduce_ex_hook = hook_fn + + def __repr__(self) -> str: + if self.repr_hook is not None: + return self.repr_hook(self) + # Instead of showing , return the class name + return str(self.__class__.__qualname__) + + def __str__(self) -> str: + if self.str_hook is not None: + return self.str_hook(self) + # Instead of showing , return the class name + return str(self.__class__.__qualname__) + + def __dir__(self): + # for auto-completion in a REPL (e.g. Jupyter notebook) + return list(super().__dir__()) + list(self.functions.keys()) + + +class _DataPipeSerializationWrapper: + def __init__(self, datapipe) -> None: + self._datapipe = datapipe + + def __getstate__(self): + use_dill = False + try: + value = pickle.dumps(self._datapipe) + except Exception: + if HAS_DILL: + # pyrefly: ignore [missing-attribute] + value = dill.dumps(self._datapipe) + use_dill = True + else: + raise + return (value, use_dill) + + def __setstate__(self, state): + value, use_dill = state + if use_dill: + # pyrefly: ignore [missing-attribute] + self._datapipe = dill.loads(value) + else: + self._datapipe = pickle.loads(value) + + def __len__(self) -> int: + try: + return len(self._datapipe) + except Exception as e: + raise TypeError( + f"{type(self).__name__} instance doesn't have valid length" + ) from e + + +class _IterDataPipeSerializationWrapper(_DataPipeSerializationWrapper, IterDataPipe): + def __init__(self, datapipe: IterDataPipe[_T_co]) -> None: + super().__init__(datapipe) + # pyrefly: ignore [invalid-type-var] + self._datapipe_iter: Iterator[_T_co] | None = None + + def __iter__(self) -> "_IterDataPipeSerializationWrapper": + self._datapipe_iter = iter(self._datapipe) + return self + + def __next__(self) -> _T_co: # type: ignore[type-var] + if self._datapipe_iter is None: + raise AssertionError( + "Iterator has not been initialized; call __iter__() before __next__()" + ) + return next(self._datapipe_iter) + + +class _MapDataPipeSerializationWrapper(_DataPipeSerializationWrapper, MapDataPipe): + def __getitem__(self, idx): + return self._datapipe[idx] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/datapipe.pyi b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/datapipe.pyi new file mode 100644 index 0000000000000000000000000000000000000000..7f49cc212383b2a635c36e1dc96c040d1d63868d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/datapipe.pyi @@ -0,0 +1,746 @@ +# @generated by torch/utils/data/datapipes/gen_pyi.py from datapipe.pyi.in +# mypy: allow-untyped-defs +# This base template ("datapipe.pyi.in") is generated from mypy stubgen with minimal editing for code injection +# The output file will be "datapipe.pyi". This is executed as part of torch/CMakeLists.txt +# Note that, for mypy, .pyi file takes precedent over .py file, such that we must define the interface for other +# classes/objects here, even though we are not injecting extra code into them at the moment. + +from collections.abc import Callable, Iterable, Iterator +from typing import Any, Literal, TypeVar + +from torch.utils.data import Dataset, default_collate, IterableDataset +from torch.utils.data.datapipes._hook_iterator import _SnapshotState +from torch.utils.data.datapipes._typing import _DataPipeMeta, _IterDataPipeMeta + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) +UNTRACABLE_DATAFRAME_PIPES: Any + +class DataChunk(list[_T]): + items: list[_T] + def __init__(self, items: Iterable[_T]) -> None: ... + def as_str(self, indent: str = "") -> str: ... + def __iter__(self) -> Iterator[_T]: ... + def raw_iterator(self) -> Iterator[_T]: ... + +class MapDataPipe(Dataset[_T_co], metaclass=_DataPipeMeta): + functions: dict[str, Callable] = ... + reduce_ex_hook: Callable | None = ... + getstate_hook: Callable | None = ... + str_hook: Callable | None = ... + repr_hook: Callable | None = ... + def __getattr__(self, attribute_name: Any): ... + @classmethod + def register_function(cls, function_name: Any, function: Any) -> None: ... + @classmethod + def register_datapipe_as_function( + cls, + function_name: Any, + cls_to_register: Any, + ): ... + def __getstate__(self): ... + def __reduce_ex__(self, *args: Any, **kwargs: Any): ... + @classmethod + def set_getstate_hook(cls, hook_fn: Any) -> None: ... + @classmethod + def set_reduce_ex_hook(cls, hook_fn: Any) -> None: ... + # Functional form of 'BatcherMapDataPipe' + def batch( + self, + batch_size: int, + drop_last: bool = False, + wrapper_class: type[DataChunk] = DataChunk, + ) -> MapDataPipe: + r""" + Create mini-batches of data (functional name: ``batch``). + + An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``, + or ``length % batch_size`` for the last batch if ``drop_last`` is set to ``False``. + + Args: + datapipe: Iterable DataPipe being batched + batch_size: The size of each batch + drop_last: Option to drop the last batch if it's not full + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper + >>> dp = SequenceWrapper(range(10)) + >>> batch_dp = dp.batch(batch_size=2) + >>> list(batch_dp) + [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]] + """ + # Functional form of 'ConcaterMapDataPipe' + def concat(self, *datapipes: MapDataPipe) -> MapDataPipe: + r""" + Concatenate multiple Map DataPipes (functional name: ``concat``). + + The new index of is the cumulative sum of source DataPipes. + For example, if there are 2 source DataPipes both with length 5, + index 0 to 4 of the resulting `ConcatMapDataPipe` would refer to + elements of the first DataPipe, and 5 to 9 would refer to elements + of the second DataPipe. + + Args: + datapipes: Map DataPipes being concatenated + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper + >>> dp1 = SequenceWrapper(range(3)) + >>> dp2 = SequenceWrapper(range(3)) + >>> concat_dp = dp1.concat(dp2) + >>> list(concat_dp) + [0, 1, 2, 0, 1, 2] + """ + # Functional form of 'MapperMapDataPipe' + def map(self, fn: Callable = ...) -> MapDataPipe: + r""" + Apply the input function over each item from the source DataPipe (functional name: ``map``). + + The function can be any regular Python function or partial object. Lambda + function is not recommended as it is not supported by pickle. + + Args: + datapipe: Source MapDataPipe + fn: Function being applied to each item + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper, Mapper + >>> def add_one(x): + ... return x + 1 + >>> dp = SequenceWrapper(range(10)) + >>> map_dp_1 = dp.map(add_one) + >>> list(map_dp_1) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + >>> map_dp_2 = Mapper(dp, lambda x: x + 1) + >>> list(map_dp_2) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + """ + # Functional form of 'ShufflerIterDataPipe' + def shuffle(self, *, indices: list | None = None) -> IterDataPipe: + r""" + Shuffle the input MapDataPipe via its indices (functional name: ``shuffle``). + + When it is used with :class:`~torch.utils.data.DataLoader`, the methods to + set up random seed are different based on :attr:`num_workers`. + + For single-process mode (:attr:`num_workers == 0`), the random seed is set before + the :class:`~torch.utils.data.DataLoader` in the main process. For multi-process + mode (:attr:`num_worker > 0`), ``worker_init_fn`` is used to set up a random seed + for each worker process. + + Args: + datapipe: MapDataPipe being shuffled + indices: a list of indices of the MapDataPipe. If not provided, we assume it uses 0-based indexing + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper + >>> dp = SequenceWrapper(range(10)) + >>> shuffle_dp = dp.shuffle().set_seed(0) + >>> list(shuffle_dp) + [7, 8, 1, 5, 3, 4, 2, 0, 9, 6] + >>> list(shuffle_dp) + [6, 1, 9, 5, 2, 4, 7, 3, 8, 0] + >>> # Reset seed for Shuffler + >>> shuffle_dp = shuffle_dp.set_seed(0) + >>> list(shuffle_dp) + [7, 8, 1, 5, 3, 4, 2, 0, 9, 6] + + Note: + Even thought this ``shuffle`` operation takes a ``MapDataPipe`` as the input, it would return an + ``IterDataPipe`` rather than a ``MapDataPipe``, because ``MapDataPipe`` should be non-sensitive to + the order of data order for the sake of random reads, but ``IterDataPipe`` depends on the order + of data during data-processing. + """ + # Functional form of 'ZipperMapDataPipe' + def zip(self, *datapipes: MapDataPipe[_T_co]) -> MapDataPipe: + r""" + Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``). + + This MataPipe is out of bound as soon as the shortest input DataPipe is exhausted. + + Args: + *datapipes: Map DataPipes being aggregated + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper + >>> dp1 = SequenceWrapper(range(3)) + >>> dp2 = SequenceWrapper(range(10, 13)) + >>> zip_dp = dp1.zip(dp2) + >>> list(zip_dp) + [(0, 10), (1, 11), (2, 12)] + """ + +class IterDataPipe(IterableDataset[_T_co], metaclass=_IterDataPipeMeta): + functions: dict[str, Callable] = ... + reduce_ex_hook: Callable | None = ... + getstate_hook: Callable | None = ... + str_hook: Callable | None = ... + repr_hook: Callable | None = ... + _number_of_samples_yielded: int = ... + _snapshot_state: _SnapshotState = _SnapshotState.Iterating # noqa: PYI015 + _fast_forward_iterator: Iterator | None = ... + def __getattr__(self, attribute_name: Any): ... + @classmethod + def register_function(cls, function_name: Any, function: Any) -> None: ... + @classmethod + def register_datapipe_as_function( + cls, + function_name: Any, + cls_to_register: Any, + enable_df_api_tracing: bool = ..., + ): ... + def __getstate__(self): ... + def __reduce_ex__(self, *args: Any, **kwargs: Any): ... + @classmethod + def set_getstate_hook(cls, hook_fn: Any) -> None: ... + @classmethod + def set_reduce_ex_hook(cls, hook_fn: Any) -> None: ... + # Functional form of 'BatcherIterDataPipe' + def batch( + self, + batch_size: int, + drop_last: bool = False, + wrapper_class: type[DataChunk] = DataChunk, + ) -> IterDataPipe: + r""" + Creates mini-batches of data (functional name: ``batch``). + + An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``, or ``length % batch_size`` for the + last batch if ``drop_last`` is set to ``False``. + + Args: + datapipe: Iterable DataPipe being batched + batch_size: The size of each batch + drop_last: Option to drop the last batch if it's not full + wrapper_class: wrapper to apply onto each batch (type ``List``) before yielding, + defaults to ``DataChunk`` + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp = IterableWrapper(range(10)) + >>> dp = dp.batch(batch_size=3, drop_last=True) + >>> list(dp) + [[0, 1, 2], [3, 4, 5], [6, 7, 8]] + """ + # Functional form of 'CollatorIterDataPipe' + def collate( + self, + conversion: Callable[..., Any]| dict[str | Any, Callable | Any]| None = default_collate, + collate_fn: Callable | None = None, + ) -> IterDataPipe: # fmt: skip + r""" + Collates samples from DataPipe to Tensor(s) by a custom collate function (functional name: ``collate``). + + By default, it uses :func:`torch.utils.data.default_collate`. + + .. note:: + While writing a custom collate function, you can import :func:`torch.utils.data.default_collate` for the + default behavior and `functools.partial` to specify any additional arguments. + + Args: + datapipe: Iterable DataPipe being collated + collate_fn: Customized collate function to collect and combine data or a batch of data. + Default function collates to Tensor(s) based on data type. + + Example: + >>> # xdoctest: +SKIP + >>> # Convert integer data to float Tensor + >>> class MyIterDataPipe(torch.utils.data.IterDataPipe): + ... def __init__(self, start, end): + ... super(MyIterDataPipe).__init__() + ... assert end > start, "this example only works with end >= start" + ... self.start = start + ... self.end = end + ... + ... def __iter__(self): + ... return iter(range(self.start, self.end)) + ... + ... def __len__(self): + ... return self.end - self.start + >>> ds = MyIterDataPipe(start=3, end=7) + >>> print(list(ds)) + [3, 4, 5, 6] + >>> def collate_fn(batch): + ... return torch.tensor(batch, dtype=torch.float) + >>> collated_ds = CollateIterDataPipe(ds, collate_fn=collate_fn) + >>> print(list(collated_ds)) + [tensor(3.), tensor(4.), tensor(5.), tensor(6.)] + """ + # Functional form of 'ConcaterIterDataPipe' + def concat(self, *datapipes: IterDataPipe) -> IterDataPipe: + r""" + Concatenates multiple Iterable DataPipes (functional name: ``concat``). + + The resulting DataPipe will yield all the elements from the first input DataPipe, before yielding from the subsequent ones. + + Args: + datapipes: Iterable DataPipes being concatenated + + Example: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> import random + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp1 = IterableWrapper(range(3)) + >>> dp2 = IterableWrapper(range(5)) + >>> list(dp1.concat(dp2)) + [0, 1, 2, 0, 1, 2, 3, 4] + """ + # Functional form of 'DemultiplexerIterDataPipe' + def demux( + self, + num_instances: int, + classifier_fn: Callable[[_T_co], int | None], + drop_none: bool = False, + buffer_size: int = 1000, + ) -> list[IterDataPipe]: + r""" + Splits the input DataPipe into multiple child DataPipes, using the given classification function (functional name: ``demux``). + + A list of the child DataPipes is returned from this operation. + + Args: + datapipe: Iterable DataPipe being filtered + num_instances: number of instances of the DataPipe to create + classifier_fn: a function that maps values to an integer within the range ``[0, num_instances - 1]`` or ``None`` + drop_none: defaults to ``False``, if ``True``, the function will skip over elements classified as ``None`` + buffer_size: this defines the maximum number of inputs that the buffer can hold across all child + DataPipes while waiting for their values to be yielded. + Defaults to ``1000``. Use ``-1`` for the unlimited buffer. + + Examples: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> from torchdata.datapipes.iter import IterableWrapper + >>> def odd_or_even(n): + ... return n % 2 + >>> source_dp = IterableWrapper(range(5)) + >>> dp1, dp2 = source_dp.demux(num_instances=2, classifier_fn=odd_or_even) + >>> list(dp1) + [0, 2, 4] + >>> list(dp2) + [1, 3] + >>> # It can also filter out any element that gets `None` from the `classifier_fn` + >>> def odd_or_even_no_zero(n): + ... return n % 2 if n != 0 else None + >>> dp1, dp2 = source_dp.demux( + ... num_instances=2, classifier_fn=odd_or_even_no_zero, drop_none=True + ... ) + >>> list(dp1) + [2, 4] + >>> list(dp2) + [1, 3] + """ + # Functional form of 'FilterIterDataPipe' + def filter(self, filter_fn: Callable, input_col=None) -> IterDataPipe: + r""" + Filters out elements from the source datapipe according to input ``filter_fn`` (functional name: ``filter``). + + Args: + datapipe: Iterable DataPipe being filtered + filter_fn: Customized function mapping an element to a boolean. + input_col: Index or indices of data which ``filter_fn`` is applied, such as: + + - ``None`` as default to apply ``filter_fn`` to the data directly. + - Integer(s) is used for list/tuple. + - Key(s) is used for dict. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> def is_even(n): + ... return n % 2 == 0 + >>> dp = IterableWrapper(range(5)) + >>> filter_dp = dp.filter(filter_fn=is_even) + >>> list(filter_dp) + [0, 2, 4] + """ + # Functional form of 'ForkerIterDataPipe' + def fork( + self, + num_instances: int, + buffer_size: int = 1000, + copy: Literal["shallow", "deep"] | None = None, + ) -> list[IterDataPipe]: + r""" + Creates multiple instances of the same Iterable DataPipe (functional name: ``fork``). + + Args: + datapipe: Iterable DataPipe being copied + num_instances: number of instances of the datapipe to create + buffer_size: this restricts how far ahead the leading child DataPipe + can read relative to the slowest child DataPipe. + Defaults to ``1000``. Use ``-1`` for the unlimited buffer. + copy: copy strategy to use for items yielded by each branch. Supported + options are ``None`` for no copying, ``"shallow"`` for shallow object + copies, and ``"deep"`` for deep object copies. Defaults to ``None``. + + Note: + All branches of the forked pipeline return the identical object unless + the copy parameter is supplied. If the object is mutable or contains + mutable objects, changing them in one branch will affect all others. + + Example: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> from torchdata.datapipes.iter import IterableWrapper + >>> source_dp = IterableWrapper(range(5)) + >>> dp1, dp2 = source_dp.fork(num_instances=2) + >>> list(dp1) + [0, 1, 2, 3, 4] + >>> list(dp2) + [0, 1, 2, 3, 4] + """ + # Functional form of 'GrouperIterDataPipe' + def groupby( + self, + group_key_fn: Callable[[_T_co], Any], + *, + keep_key: bool = False, + buffer_size: int = 10000, + group_size: int | None = None, + guaranteed_group_size: int | None = None, + drop_remaining: bool = False, + ) -> IterDataPipe: + r""" + Groups data from IterDataPipe by keys from ``group_key_fn``, yielding a ``DataChunk`` with batch size up to ``group_size``. + + (functional name: ``groupby``). + + The samples are read sequentially from the source ``datapipe``, and a batch of samples belonging to the same group + will be yielded as soon as the size of the batch reaches ``group_size``. When the buffer is full, + the DataPipe will yield the largest batch with the same key, provided that its size is larger + than ``guaranteed_group_size``. If its size is smaller, it will be dropped if ``drop_remaining=True``. + + After iterating through the entirety of source ``datapipe``, everything not dropped due to the buffer capacity + will be yielded from the buffer, even if the group sizes are smaller than ``guaranteed_group_size``. + + Args: + datapipe: Iterable datapipe to be grouped + group_key_fn: Function used to generate group key from the data of the source datapipe + keep_key: Option to yield the matching key along with the items in a tuple, + resulting in `(key, [items])` otherwise returning [items] + buffer_size: The size of buffer for ungrouped data + group_size: The max size of each group, a batch is yielded as soon as it reaches this size + guaranteed_group_size: The guaranteed minimum group size to be yielded in case the buffer is full + drop_remaining: Specifies if the group smaller than ``guaranteed_group_size`` will be dropped from buffer + when the buffer is full + + Example: + >>> import os + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> def group_fn(file): + ... return os.path.basename(file).split(".")[0] + >>> source_dp = IterableWrapper( + ... ["a.png", "b.png", "a.json", "b.json", "a.jpg", "c.json"] + ... ) + >>> dp0 = source_dp.groupby(group_key_fn=group_fn) + >>> list(dp0) + [['a.png', 'a.json', 'a.jpg'], ['b.png', 'b.json'], ['c.json']] + >>> # A group is yielded as soon as its size equals to `group_size` + >>> dp1 = source_dp.groupby(group_key_fn=group_fn, group_size=2) + >>> list(dp1) + [['a.png', 'a.json'], ['b.png', 'b.json'], ['a.jpg'], ['c.json']] + >>> # Scenario where `buffer` is full, and group 'a' needs to be yielded since its size > `guaranteed_group_size` + >>> dp2 = source_dp.groupby( + ... group_key_fn=group_fn, + ... buffer_size=3, + ... group_size=3, + ... guaranteed_group_size=2, + ... ) + >>> list(dp2) + [['a.png', 'a.json'], ['b.png', 'b.json'], ['a.jpg'], ['c.json']] + """ + # Functional form of 'FileListerIterDataPipe' + def list_files( + self, + masks: str | list[str] = "", + *, + recursive: bool = False, + abspath: bool = False, + non_deterministic: bool = False, + length: int = -1, + ) -> IterDataPipe: + r""" + Given path(s) to the root directory, yields file pathname(s) (path + filename) of files within the root directory. + + Multiple root directories can be provided (functional name: ``list_files``). + + Args: + root: Root directory or a sequence of root directories + masks: Unix style filter string or string list for filtering file name(s) + recursive: Whether to return pathname from nested directories or not + abspath: Whether to return relative pathname or absolute pathname + non_deterministic: Whether to return pathname in sorted order or not. + If ``False``, the results yielded from each root directory will be sorted + length: Nominal length of the datapipe + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import FileLister + >>> dp = FileLister(root=".", recursive=True) + >>> list(dp) + ['example.py', './data/data.tar'] + """ + # Functional form of 'MapperIterDataPipe' + def map( + self, + fn: Callable, + input_col=None, + output_col=None, + ) -> IterDataPipe: + r""" + Applies a function over each item from the source DataPipe (functional name: ``map``). + + The function can be any regular Python function or partial object. Lambda + function is not recommended as it is not supported by pickle. + + Args: + datapipe: Source Iterable DataPipe + fn: Function being applied over each item + input_col: Index or indices of data which ``fn`` is applied, such as: + + - ``None`` as default to apply ``fn`` to the data directly. + - Integer(s) is used for list/tuple. + - Key(s) is used for dict. + + output_col: Index of data where result of ``fn`` is placed. ``output_col`` can be specified + only when ``input_col`` is not ``None`` + + - ``None`` as default to replace the index that ``input_col`` specified; For ``input_col`` with + multiple indices, the left-most one is used, and other indices will be removed. + - Integer is used for list/tuple. ``-1`` represents to append result at the end. + - Key is used for dict. New key is acceptable. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper, Mapper + >>> def add_one(x): + ... return x + 1 + >>> dp = IterableWrapper(range(10)) + >>> # Invocation via functional form is preferred + ... map_dp_1 = dp.map(add_one) + >>> list(map_dp_1) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + >>> # We discourage the usage of `lambda` functions as they are not serializable with `pickle` + >>> # Use `functools.partial` or explicitly define the function instead + >>> map_dp_2 = Mapper(dp, lambda x: x + 1) + >>> list(map_dp_2) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + """ + # Functional form of 'MultiplexerIterDataPipe' + def mux(self, *datapipes) -> IterDataPipe: + r""" + Yields one element at a time from each of the input Iterable DataPipes (functional name: ``mux``). + + As in, one element from the 1st input DataPipe, then one element from the 2nd DataPipe in the next iteration, + and so on. It ends when the shortest input DataPipe is exhausted. + + Args: + datapipes: Iterable DataPipes that will take turn to yield their elements, until the shortest DataPipe is exhausted + + Example: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp1, dp2, dp3 = ( + ... IterableWrapper(range(3)), + ... IterableWrapper(range(10, 15)), + ... IterableWrapper(range(20, 25)), + ... ) + >>> list(dp1.mux(dp2, dp3)) + [0, 10, 20, 1, 11, 21, 2, 12, 22] + """ + # Functional form of 'FileOpenerIterDataPipe' + def open_files( + self, + mode: str = "r", + encoding: str | None = None, + length: int = -1, + ) -> IterDataPipe: + r""" + Given pathnames, opens files and yield pathname and file stream in a tuple (functional name: ``open_files``). + + Args: + datapipe: Iterable datapipe that provides pathnames + mode: An optional string that specifies the mode in which + the file is opened by ``open()``. It defaults to ``r``, other options are + ``b`` for reading in binary mode and ``t`` for text mode. + encoding: An optional string that specifies the encoding of the + underlying file. It defaults to ``None`` to match the default encoding of ``open``. + length: Nominal length of the datapipe + + Note: + The opened file handles will be closed by Python's GC periodically. Users can choose + to close them explicitly. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import ( + ... FileLister, + ... FileOpener, + ... StreamReader, + ... ) + >>> dp = FileLister(root=".").filter(lambda fname: fname.endswith(".txt")) + >>> dp = FileOpener(dp) + >>> dp = StreamReader(dp) + >>> list(dp) + [('./abc.txt', 'abc')] + """ + # Functional form of 'StreamReaderIterDataPipe' + def read_from_stream(self, chunk: int | None = None) -> IterDataPipe: + r""" + Given IO streams and their label names, yield bytes with label name as tuple. + + (functional name: ``read_from_stream``). + + Args: + datapipe: Iterable DataPipe provides label/URL and byte stream + chunk: Number of bytes to be read from stream per iteration. + If ``None``, all bytes will be read until the EOF. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper, StreamReader + >>> from io import StringIO + >>> dp = IterableWrapper([("alphabet", StringIO("abcde"))]) + >>> list(StreamReader(dp, chunk=1)) + [('alphabet', 'a'), ('alphabet', 'b'), ('alphabet', 'c'), ('alphabet', 'd'), ('alphabet', 'e')] + """ + # Functional form of 'RoutedDecoderIterDataPipe' + def routed_decode( + self, + *handlers: Callable, + key_fn: Callable = ..., + ) -> IterDataPipe: + r""" + Decodes binary streams from input DataPipe, yields pathname and decoded data in a tuple. + + (functional name: ``routed_decode``) + + Args: + datapipe: Iterable datapipe that provides pathname and binary stream in tuples + handlers: Optional user defined decoder handlers. If ``None``, basic and image decoder + handlers will be set as default. If multiple handles are provided, the priority + order follows the order of handlers (the first handler has the top priority) + key_fn: Function for decoder to extract key from pathname to dispatch handlers. + Default is set to extract file extension from pathname + + Note: + When ``key_fn`` is specified returning anything other than extension, the default + handler will not work and users need to specify custom handler. Custom handler + could use regex to determine the eligibility to handle data. + """ + # Functional form of 'ShardingFilterIterDataPipe' + def sharding_filter(self, sharding_group_filter=None) -> IterDataPipe: + r""" + Wrapper that allows DataPipe to be sharded (functional name: ``sharding_filter``). + + After ``apply_sharding`` is called, each instance of the DataPipe (on different workers) will have every `n`-th element of the + original DataPipe, where `n` equals to the number of instances. + + Args: + source_datapipe: Iterable DataPipe that will be sharded + """ + # Functional form of 'ShufflerIterDataPipe' + def shuffle( + self, + *, + buffer_size: int = 10000, + unbatch_level: int = 0, + ) -> IterDataPipe: + r""" + Shuffle the input DataPipe with a buffer (functional name: ``shuffle``). + + The buffer with ``buffer_size`` is filled with elements from the datapipe first. Then, + each item will be yielded from the buffer by reservoir sampling via iterator. + + ``buffer_size`` is required to be larger than ``0``. For ``buffer_size == 1``, the + datapipe is not shuffled. In order to fully shuffle all elements from datapipe, + ``buffer_size`` is required to be greater than or equal to the size of datapipe. + + When it is used with :class:`torch.utils.data.DataLoader`, the methods to + set up random seed are different based on :attr:`num_workers`. + + For single-process mode (:attr:`num_workers == 0`), the random seed is set before + the :class:`~torch.utils.data.DataLoader` in the main process. For multi-process + mode (:attr:`num_worker > 0`), `worker_init_fn` is used to set up a random seed + for each worker process. + + Args: + datapipe: The IterDataPipe being shuffled + buffer_size: The buffer size for shuffling (default to ``10000``) + unbatch_level: Specifies if it is necessary to unbatch source data before + applying the shuffle + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp = IterableWrapper(range(10)) + >>> shuffle_dp = dp.shuffle() + >>> list(shuffle_dp) + [0, 4, 1, 6, 3, 2, 9, 5, 7, 8] + """ + # Functional form of 'UnBatcherIterDataPipe' + def unbatch(self, unbatch_level: int = 1) -> IterDataPipe: + r""" + Undos batching of data (functional name: ``unbatch``). + + In other words, it flattens the data up to the specified level within a batched DataPipe. + + Args: + datapipe: Iterable DataPipe being un-batched + unbatch_level: Defaults to ``1`` (only flattening the top level). If set to ``2``, + it will flatten the top two levels, and ``-1`` will flatten the entire DataPipe. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> source_dp = IterableWrapper([[[0, 1], [2]], [[3, 4], [5]], [[6]]]) + >>> dp1 = source_dp.unbatch() + >>> list(dp1) + [[0, 1], [2], [3, 4], [5], [6]] + >>> dp2 = source_dp.unbatch(unbatch_level=2) + >>> list(dp2) + [0, 1, 2, 3, 4, 5, 6] + """ + # Functional form of 'ZipperIterDataPipe' + def zip(self, *datapipes: IterDataPipe) -> IterDataPipe: + r""" + Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``). + + The output is stopped as soon as the shortest input DataPipe is exhausted. + + Args: + *datapipes: Iterable DataPipes being aggregated + + Example: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp1, dp2, dp3 = ( + ... IterableWrapper(range(5)), + ... IterableWrapper(range(10, 15)), + ... IterableWrapper(range(20, 25)), + ... ) + >>> list(dp1.zip(dp2, dp3)) + [(0, 10, 20), (1, 11, 21), (2, 12, 22), (3, 13, 23), (4, 14, 24)] + """ + +class DFIterDataPipe(IterDataPipe): + def _is_dfpipe(self): ... + def __iter__(self): ... + +class _DataPipeSerializationWrapper: + def __init__(self, datapipe): ... + def __getstate__(self): ... + def __setstate__(self, state): ... + def __len__(self): ... + +class _IterDataPipeSerializationWrapper(_DataPipeSerializationWrapper, IterDataPipe): + def __iter__(self): ... + +class _MapDataPipeSerializationWrapper(_DataPipeSerializationWrapper, MapDataPipe): + def __getitem__(self, idx): ... diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/gen_pyi.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/gen_pyi.py new file mode 100644 index 0000000000000000000000000000000000000000..90f9d80a2e7fef61459d525d32486211415ad3ed --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/gen_pyi.py @@ -0,0 +1,336 @@ +# mypy: allow-untyped-defs +import os +from collections import defaultdict +from pathlib import Path +from typing import Any +from typing_extensions import deprecated + + +try: + from torchgen.api.python import format_function_signature + from torchgen.utils import FileManager +except ImportError: + import sys + + REPO_ROOT = Path(__file__).absolute().parents[4] + sys.path.insert(0, str(REPO_ROOT)) + + from torchgen.api.python import format_function_signature + from torchgen.utils import FileManager + + if len(sys.path) > 0 and sys.path[0] == str(REPO_ROOT): + del sys.path[0] + + +__all__: list[str] = [] # not intended to expose any symbols + + +def __dir__() -> list[str]: + return [] # appease public API test + + +@deprecated( + "`torch.utils.data.datapipes.gen_pyi.materialize_lines` is deprecated and will be removed in the future.", + category=FutureWarning, +) +def materialize_lines(lines: list[str], indentation: int) -> str: + output = "" + new_line_with_indent = "\n" + " " * indentation + for i, line in enumerate(lines): + if i != 0: + output += new_line_with_indent + output += line.replace("\n", new_line_with_indent) + return output + + +@deprecated( + "`torch.utils.data.datapipes.gen_pyi.gen_from_template` is deprecated and will be removed in the future.", + category=FutureWarning, +) +def gen_from_template( + dir: str, + template_name: str, + output_name: str, + replacements: list[tuple[str, Any, int]], +) -> None: + template_path = os.path.join(dir, template_name) + output_path = os.path.join(dir, output_name) + + with open(template_path, encoding="utf-8") as f: + content = f.read() + for placeholder, lines, indentation in replacements: + with open(output_path, "w", encoding="utf-8") as f: + content = content.replace( + placeholder, materialize_lines(lines, indentation) + ) + f.write(content) + + +def find_file_paths(dir_paths: list[str], files_to_exclude: set[str]) -> set[str]: + """ + When given a path to a directory, returns the paths to the relevant files within it. + + This function does NOT recursive traverse to subdirectories. + """ + paths: set[str] = set() + for dir_path in dir_paths: + all_files = os.listdir(dir_path) + python_files = {fname for fname in all_files if ".py" == fname[-3:]} + filter_files = { + fname for fname in python_files if fname not in files_to_exclude + } + paths.update({os.path.join(dir_path, fname) for fname in filter_files}) + return paths + + +def extract_method_name(line: str) -> str: + """Extract method name from decorator in the form of "@functional_datapipe({method_name})".""" + if '("' in line: + start_token, end_token = '("', '")' + elif "('" in line: + start_token, end_token = "('", "')" + else: + raise RuntimeError( + f"Unable to find appropriate method name within line:\n{line}" + ) + start, end = line.find(start_token) + len(start_token), line.find(end_token) + return line[start:end] + + +def extract_class_name(line: str) -> str: + """Extract class name from class definition in the form of "class {CLASS_NAME}({Type}):".""" + start_token = "class " + end_token = "(" + start, end = line.find(start_token) + len(start_token), line.find(end_token) + return line[start:end] + + +def parse_datapipe_file( + file_path: str, +) -> tuple[dict[str, list[str]], dict[str, str], set[str], dict[str, list[str]]]: + """Given a path to file, parses the file and returns a dictionary of method names to function signatures.""" + method_to_signature, method_to_class_name, special_output_type = {}, {}, set() + doc_string_dict = defaultdict(list) + with open(file_path, encoding="utf-8") as f: + open_paren_count = 0 + method_name, class_name, signature = "", "", "" + skip = False + for line in f: + if line.count('"""') % 2 == 1: + skip = not skip + if skip or '"""' in line: # Saving docstrings + doc_string_dict[method_name].append(line) + continue + if "@functional_datapipe" in line: + method_name = extract_method_name(line) + doc_string_dict[method_name] = [] + continue + if method_name and "class " in line: + class_name = extract_class_name(line) + continue + if method_name and ("def __init__(" in line or "def __new__(" in line): + if "def __new__(" in line: + special_output_type.add(method_name) + open_paren_count += 1 + start = line.find("(") + len("(") + line = line[start:] + if open_paren_count > 0: + open_paren_count += line.count("(") + open_paren_count -= line.count(")") + if open_paren_count == 0: + end = line.rfind(")") + signature += line[:end] + method_to_signature[method_name] = process_signature(signature) + method_to_class_name[method_name] = class_name + method_name, class_name, signature = "", "", "" + elif open_paren_count < 0: + raise RuntimeError( + "open parenthesis count < 0. This shouldn't be possible." + ) + else: + signature += line.strip() + return ( + method_to_signature, + method_to_class_name, + special_output_type, + doc_string_dict, + ) + + +def parse_datapipe_files( + file_paths: set[str], +) -> tuple[dict[str, list[str]], dict[str, str], set[str], dict[str, list[str]]]: + methods_and_signatures = {} + methods_and_class_names = {} + methods_with_special_output_types = set() + methods_and_doc_strings = {} + for path in file_paths: + ( + method_to_signature, + method_to_class_name, + methods_needing_special_output_types, + doc_string_dict, + ) = parse_datapipe_file(path) + methods_and_signatures.update(method_to_signature) + methods_and_class_names.update(method_to_class_name) + methods_with_special_output_types.update(methods_needing_special_output_types) + methods_and_doc_strings.update(doc_string_dict) + return ( + methods_and_signatures, + methods_and_class_names, + methods_with_special_output_types, + methods_and_doc_strings, + ) + + +def split_outside_bracket(line: str, delimiter: str = ",") -> list[str]: + """Given a line of text, split it on comma unless the comma is within a bracket '[]'.""" + bracket_count = 0 + curr_token = "" + res = [] + for char in line: + if char == "[": + bracket_count += 1 + elif char == "]": + bracket_count -= 1 + elif char == delimiter and bracket_count == 0: + res.append(curr_token) + curr_token = "" + continue + curr_token += char + res.append(curr_token) + return res + + +def process_signature(line: str) -> list[str]: + """ + Clean up a given raw function signature. + + This includes removing the self-referential datapipe argument, default + arguments of input functions, newlines, and spaces. + """ + tokens: list[str] = split_outside_bracket(line) + for i, token in enumerate(tokens): + tokens[i] = token.strip(" ") + if token == "cls": + tokens[i] = "self" + elif i > 0 and ("self" == tokens[i - 1]) and (tokens[i][0] != "*"): + # Remove the datapipe after 'self' or 'cls' unless it has '*' + tokens[i] = "" + elif "Callable =" in token: # Remove default argument if it is a function + head = token.rpartition("=")[0] + tokens[i] = head.strip(" ") + " = ..." + tokens = [t for t in tokens if t != ""] + return tokens + + +def get_method_definitions( + file_path: str | list[str], + files_to_exclude: set[str], + deprecated_files: set[str], + default_output_type: str, + method_to_special_output_type: dict[str, str], + root: str = "", +) -> list[str]: + """ + #.pyi generation for functional DataPipes Process. + + # 1. Find files that we want to process (exclude the ones who don't) + # 2. Parse method name and signature + # 3. Remove first argument after self (unless it is "*datapipes"), default args, and spaces + """ + if root == "": + root = str(Path(__file__).parent.resolve()) + file_path = [file_path] if isinstance(file_path, str) else file_path + file_path = [os.path.join(root, path) for path in file_path] + file_paths = find_file_paths( + file_path, files_to_exclude=files_to_exclude.union(deprecated_files) + ) + ( + methods_and_signatures, + methods_and_class_names, + methods_w_special_output_types, + methods_and_doc_strings, + ) = parse_datapipe_files(file_paths) + + for fn_name in method_to_special_output_type: + if fn_name not in methods_w_special_output_types: + methods_w_special_output_types.add(fn_name) + + method_definitions = [] + for method_name, arguments in methods_and_signatures.items(): + class_name = methods_and_class_names[method_name] + if method_name in methods_w_special_output_types: + output_type = method_to_special_output_type[method_name] + else: + output_type = default_output_type + doc_string = "".join(methods_and_doc_strings[method_name]) + if doc_string == "": + doc_string = " ..." + else: + doc_string = "\n" + doc_string + definition = format_function_signature(method_name, arguments, output_type) + method_definitions.append( + f"# Functional form of '{class_name}'\n" + + definition.removesuffix("...").rstrip() # remove "..." + + doc_string, + ) + method_definitions.sort( + key=lambda s: s.split("\n")[1] + ) # sorting based on method_name + + return method_definitions + + +# Defined outside of main() so they can be imported by TorchData +iterDP_file_path: str = "iter" +iterDP_files_to_exclude: set[str] = {"__init__.py", "utils.py"} +iterDP_deprecated_files: set[str] = set() +iterDP_method_to_special_output_type: dict[str, str] = { + "demux": "list[IterDataPipe]", + "fork": "list[IterDataPipe]", +} + +mapDP_file_path: str = "map" +mapDP_files_to_exclude: set[str] = {"__init__.py", "utils.py"} +mapDP_deprecated_files: set[str] = set() +mapDP_method_to_special_output_type: dict[str, str] = {"shuffle": "IterDataPipe"} + + +def main() -> None: + """ + # Inject file into template datapipe.pyi.in. + + TODO: The current implementation of this script only generates interfaces for built-in methods. To generate + interface for user-defined DataPipes, consider changing `IterDataPipe.register_datapipe_as_function`. + """ + iter_method_definitions = get_method_definitions( + iterDP_file_path, + iterDP_files_to_exclude, + iterDP_deprecated_files, + "IterDataPipe", + iterDP_method_to_special_output_type, + ) + + map_method_definitions = get_method_definitions( + mapDP_file_path, + mapDP_files_to_exclude, + mapDP_deprecated_files, + "MapDataPipe", + mapDP_method_to_special_output_type, + ) + + path = Path(__file__).absolute().parent + fm = FileManager(install_dir=path, template_dir=path, dry_run=False) + fm.write_with_template( + "datapipe.pyi", + "datapipe.pyi.in", + lambda: { + "IterDataPipeMethods": iter_method_definitions, + "MapDataPipeMethods": map_method_definitions, + }, + ) + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..05831250da468cc76e8c2cc8e4018373e8191951 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/__init__.py @@ -0,0 +1,66 @@ +from torch.utils.data.datapipes.iter.callable import ( + CollatorIterDataPipe as Collator, + MapperIterDataPipe as Mapper, +) +from torch.utils.data.datapipes.iter.combinatorics import ( + SamplerIterDataPipe as Sampler, + ShufflerIterDataPipe as Shuffler, +) +from torch.utils.data.datapipes.iter.combining import ( + ConcaterIterDataPipe as Concater, + DemultiplexerIterDataPipe as Demultiplexer, + ForkerIterDataPipe as Forker, + MultiplexerIterDataPipe as Multiplexer, + ZipperIterDataPipe as Zipper, +) +from torch.utils.data.datapipes.iter.filelister import ( + FileListerIterDataPipe as FileLister, +) +from torch.utils.data.datapipes.iter.fileopener import ( + FileOpenerIterDataPipe as FileOpener, +) +from torch.utils.data.datapipes.iter.grouping import ( + BatcherIterDataPipe as Batcher, + GrouperIterDataPipe as Grouper, + UnBatcherIterDataPipe as UnBatcher, +) +from torch.utils.data.datapipes.iter.routeddecoder import ( + RoutedDecoderIterDataPipe as RoutedDecoder, +) +from torch.utils.data.datapipes.iter.selecting import FilterIterDataPipe as Filter +from torch.utils.data.datapipes.iter.sharding import ( + ShardingFilterIterDataPipe as ShardingFilter, +) +from torch.utils.data.datapipes.iter.streamreader import ( + StreamReaderIterDataPipe as StreamReader, +) +from torch.utils.data.datapipes.iter.utils import ( + IterableWrapperIterDataPipe as IterableWrapper, +) + + +__all__ = [ + "Batcher", + "Collator", + "Concater", + "Demultiplexer", + "FileLister", + "FileOpener", + "Filter", + "Forker", + "Grouper", + "IterableWrapper", + "Mapper", + "Multiplexer", + "RoutedDecoder", + "Sampler", + "ShardingFilter", + "Shuffler", + "StreamReader", + "UnBatcher", + "Zipper", +] + +# Please keep this list sorted +if __all__ != sorted(__all__): + raise AssertionError("__all__ is not sorted") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/callable.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/callable.py new file mode 100644 index 0000000000000000000000000000000000000000..aecaf775929ee846b3508f58be1c0355426b02cd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/callable.py @@ -0,0 +1,244 @@ +# mypy: allow-untyped-defs +import functools +from collections import namedtuple +from collections.abc import Callable, Iterator, Sized +from typing import Any, TypeVar + +import torch +from torch.utils.data._utils.collate import default_collate +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper +from torch.utils.data.datapipes.datapipe import IterDataPipe +from torch.utils.data.datapipes.utils.common import ( + _check_unpickable_fn, + validate_input_col, +) + + +__all__ = [ + "CollatorIterDataPipe", + "MapperIterDataPipe", +] + + +_T_co = TypeVar("_T_co", covariant=True) + + +@functional_datapipe("map") +class MapperIterDataPipe(IterDataPipe[_T_co]): + r""" + Applies a function over each item from the source DataPipe (functional name: ``map``). + + The function can be any regular Python function or partial object. Lambda + function is not recommended as it is not supported by pickle. + + Args: + datapipe: Source Iterable DataPipe + fn: Function being applied over each item + input_col: Index or indices of data which ``fn`` is applied, such as: + + - ``None`` as default to apply ``fn`` to the data directly. + - Integer(s) is used for list/tuple. + - Key(s) is used for dict. + + output_col: Index of data where result of ``fn`` is placed. ``output_col`` can be specified + only when ``input_col`` is not ``None`` + + - ``None`` as default to replace the index that ``input_col`` specified; For ``input_col`` with + multiple indices, the left-most one is used, and other indices will be removed. + - Integer is used for list/tuple. ``-1`` represents to append result at the end. + - Key is used for dict. New key is acceptable. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper, Mapper + >>> def add_one(x): + ... return x + 1 + >>> dp = IterableWrapper(range(10)) + >>> # Invocation via functional form is preferred + ... map_dp_1 = dp.map(add_one) + >>> list(map_dp_1) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + >>> # We discourage the usage of `lambda` functions as they are not serializable with `pickle` + >>> # Use `functools.partial` or explicitly define the function instead + >>> map_dp_2 = Mapper(dp, lambda x: x + 1) + >>> list(map_dp_2) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + """ + + datapipe: IterDataPipe + fn: Callable + + def __init__( + self, + datapipe: IterDataPipe, + fn: Callable, + input_col=None, + output_col=None, + ) -> None: + torch._C._log_api_usage_once("python.data_pipes.map") + super().__init__() + self.datapipe = datapipe + + _check_unpickable_fn(fn) + self.fn = fn # type: ignore[assignment] + + self.input_col = input_col + if input_col is None and output_col is not None: + raise ValueError("`output_col` must be None when `input_col` is None.") + if isinstance(output_col, (list, tuple)): + if len(output_col) > 1: + raise ValueError("`output_col` must be a single-element list or tuple") + output_col = output_col[0] + self.output_col = output_col + validate_input_col(fn, input_col) + + def _apply_fn(self, data): + if self.input_col is None and self.output_col is None: + return self.fn(data) + + if self.input_col is None: + res = self.fn(data) + elif isinstance(self.input_col, (list, tuple)): + args = tuple(data[col] for col in self.input_col) + res = self.fn(*args) + else: + res = self.fn(data[self.input_col]) + + # Copy tuple to list and run in-place modification because tuple is immutable. + if isinstance(data, tuple): + t_flag = True + data = list(data) + else: + t_flag = False + + if self.output_col is None: + if isinstance(self.input_col, (list, tuple)): + data[self.input_col[0]] = res + for idx in sorted(self.input_col[1:], reverse=True): + del data[idx] + else: + data[self.input_col] = res + else: + if self.output_col == -1: + data.append(res) + else: + data[self.output_col] = res + + # Convert list back to tuple + return tuple(data) if t_flag else data + + def __iter__(self) -> Iterator[_T_co]: + for data in self.datapipe: + yield self._apply_fn(data) + + def __len__(self) -> int: + # pyrefly: ignore [unsafe-overlap] + if isinstance(self.datapipe, Sized): + return len(self.datapipe) + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") + + +def _collate_helper(conversion, item): + # TODO(VitalyFedyunin): Verify that item is any sort of batch + if len(item.items) > 1: + # TODO(VitalyFedyunin): Compact all batch dataframes into one + raise RuntimeError("Only supports one DataFrame per batch") + df = item[0] + columns_name = df_wrapper.get_columns(df) + tuple_names: list = [] + tuple_values: list = [] + + for name in conversion: + if name not in columns_name: + raise RuntimeError("Conversion keys mismatch") + + for name in columns_name: + if name in conversion: + if not callable(conversion[name]): + raise RuntimeError( + "Collate (DF)DataPipe requires callable as dict values" + ) + collation_fn = conversion[name] + else: + # TODO(VitalyFedyunin): Add default collation into df_wrapper + try: + import torcharrow.pytorch as tap # type: ignore[import] + + collation_fn = tap.rec.Default() + except Exception as e: + raise RuntimeError( + "unable to import default collation function from the TorchArrow" + ) from e + + tuple_names.append(str(name)) + value = collation_fn(df[name]) + tuple_values.append(value) + + # TODO(VitalyFedyunin): We can dynamically extract types from the tuple_values here + # TODO(VitalyFedyunin): Instead of ignoring mypy error, make sure tuple_names is not empty + tpl_cls = namedtuple("CollateResult", tuple_names) # type: ignore[misc] + tuple = tpl_cls(*tuple_values) + return tuple + + +@functional_datapipe("collate") +class CollatorIterDataPipe(MapperIterDataPipe): + r""" + Collates samples from DataPipe to Tensor(s) by a custom collate function (functional name: ``collate``). + + By default, it uses :func:`torch.utils.data.default_collate`. + + .. note:: + While writing a custom collate function, you can import :func:`torch.utils.data.default_collate` for the + default behavior and `functools.partial` to specify any additional arguments. + + Args: + datapipe: Iterable DataPipe being collated + collate_fn: Customized collate function to collect and combine data or a batch of data. + Default function collates to Tensor(s) based on data type. + + Example: + >>> # xdoctest: +SKIP + >>> # Convert integer data to float Tensor + >>> class MyIterDataPipe(torch.utils.data.IterDataPipe): + ... def __init__(self, start, end): + ... super(MyIterDataPipe).__init__() + ... assert end > start, "this example only works with end >= start" + ... self.start = start + ... self.end = end + ... + ... def __iter__(self): + ... return iter(range(self.start, self.end)) + ... + ... def __len__(self): + ... return self.end - self.start + >>> ds = MyIterDataPipe(start=3, end=7) + >>> print(list(ds)) + [3, 4, 5, 6] + >>> def collate_fn(batch): + ... return torch.tensor(batch, dtype=torch.float) + >>> collated_ds = CollateIterDataPipe(ds, collate_fn=collate_fn) + >>> print(list(collated_ds)) + [tensor(3.), tensor(4.), tensor(5.), tensor(6.)] + """ + + def __init__( + self, + datapipe: IterDataPipe, + conversion: Callable[..., Any] + | dict[str | Any, Callable | Any] + | None = default_collate, + collate_fn: Callable | None = None, + ) -> None: + # TODO(VitalyFedyunin): Replace `Callable[..., Any]` with `Callable[[IColumn], Any]` + # TODO(VitalyFedyunin): Replace with `Dict[Union[str, IColumn], Union[Callable, Enum]]` + if collate_fn is not None: + super().__init__(datapipe, fn=collate_fn) + else: + if callable(conversion): + super().__init__(datapipe, fn=conversion) + else: + # TODO(VitalyFedyunin): Validate passed dictionary + collate_fn = functools.partial(_collate_helper, conversion) + super().__init__(datapipe, fn=collate_fn) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/combinatorics.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/combinatorics.py new file mode 100644 index 0000000000000000000000000000000000000000..f9221e63f12c27f8552c55bafa1ab841beba06af --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/combinatorics.py @@ -0,0 +1,195 @@ +# mypy: allow-untyped-defs +import random +from collections.abc import Iterator, Sized +from typing import TypeVar + +import torch +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import IterDataPipe +from torch.utils.data.sampler import Sampler, SequentialSampler + + +__all__ = [ + "SamplerIterDataPipe", + "ShufflerIterDataPipe", +] + + +_T_co = TypeVar("_T_co", covariant=True) + + +class SamplerIterDataPipe(IterDataPipe[_T_co]): + r""" + Generate sample elements using the provided ``Sampler`` (defaults to :class:`SequentialSampler`). + + Args: + datapipe: IterDataPipe to sample from + sampler: Sampler class to generate sample elements from input DataPipe. + Default is :class:`SequentialSampler` for IterDataPipe + """ + + datapipe: IterDataPipe + sampler: Sampler + + def __init__( + self, + datapipe: IterDataPipe, + sampler: type[Sampler] = SequentialSampler, + sampler_args: tuple | None = None, + sampler_kwargs: dict | None = None, + ) -> None: + # pyrefly: ignore [unsafe-overlap] + if not isinstance(datapipe, Sized): + raise AssertionError( + "Sampler class requires input datapipe implemented `__len__`" + ) + super().__init__() + + self.datapipe = datapipe + self.sampler_args = () if sampler_args is None else sampler_args + self.sampler_kwargs = {} if sampler_kwargs is None else sampler_kwargs + self.sampler_kwargs["data_source"] = self.datapipe + self.sampler = sampler(*self.sampler_args, **self.sampler_kwargs) + + def __iter__(self) -> Iterator[_T_co]: + return iter(self.sampler) + + def __len__(self) -> int: + # Dataset has been tested as `Sized` + if isinstance(self.sampler, Sized): + return len(self.sampler) + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") + + +@functional_datapipe("shuffle") +class ShufflerIterDataPipe(IterDataPipe[_T_co]): + r""" + Shuffle the input DataPipe with a buffer (functional name: ``shuffle``). + + The buffer with ``buffer_size`` is filled with elements from the datapipe first. Then, + each item will be yielded from the buffer by reservoir sampling via iterator. + + ``buffer_size`` is required to be larger than ``0``. For ``buffer_size == 1``, the + datapipe is not shuffled. In order to fully shuffle all elements from datapipe, + ``buffer_size`` is required to be greater than or equal to the size of datapipe. + + When it is used with :class:`torch.utils.data.DataLoader`, the methods to + set up random seed are different based on :attr:`num_workers`. + + For single-process mode (:attr:`num_workers == 0`), the random seed is set before + the :class:`~torch.utils.data.DataLoader` in the main process. For multi-process + mode (:attr:`num_worker > 0`), `worker_init_fn` is used to set up a random seed + for each worker process. + + Args: + datapipe: The IterDataPipe being shuffled + buffer_size: The buffer size for shuffling (default to ``10000``) + unbatch_level: Specifies if it is necessary to unbatch source data before + applying the shuffle + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp = IterableWrapper(range(10)) + >>> shuffle_dp = dp.shuffle() + >>> list(shuffle_dp) + [0, 4, 1, 6, 3, 2, 9, 5, 7, 8] + """ + + datapipe: IterDataPipe[_T_co] + buffer_size: int + _buffer: list[_T_co] + _enabled: bool + _seed: int | None + _rng: random.Random + + def __init__( + self, + datapipe: IterDataPipe[_T_co], + *, + buffer_size: int = 10000, + unbatch_level: int = 0, + ) -> None: + super().__init__() + # TODO: Performance optimization + # buffer can be a fixed size and remove expensive `append()` and `len()` operations + self._buffer: list[_T_co] = [] + if buffer_size <= 0: + raise AssertionError("buffer_size should be larger than 0") + if unbatch_level == 0: + self.datapipe = datapipe + else: + self.datapipe = datapipe.unbatch(unbatch_level=unbatch_level) + self.buffer_size = buffer_size + self._enabled = True + self._seed = None + self._rng = random.Random() + + def set_shuffle(self, shuffle=True): + self._enabled = shuffle + return self + + def set_seed(self, seed: int): + self._seed = seed + return self + + def __iter__(self) -> Iterator[_T_co]: + if not self._enabled: + yield from self.datapipe + else: + for x in self.datapipe: + if len(self._buffer) == self.buffer_size: + idx = self._rng.randint(0, len(self._buffer) - 1) + val, self._buffer[idx] = self._buffer[idx], x + yield val + else: + self._buffer.append(x) + while self._buffer: + idx = self._rng.randint(0, len(self._buffer) - 1) + yield self._buffer.pop(idx) + + def __len__(self) -> int: + # pyrefly: ignore [unsafe-overlap] + if isinstance(self.datapipe, Sized): + return len(self.datapipe) + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") + + def reset(self) -> None: + self._buffer = [] + if self._enabled: + if self._seed is None: + self._seed = int(torch.empty((), dtype=torch.int64).random_().item()) + self._rng.seed(self._seed) + self._seed = None + + def __getstate__(self): + state = ( + self.datapipe, + self.buffer_size, + self._enabled, + self._seed, + self._buffer, + self._rng.getstate(), + self._valid_iterator_id, + self._number_of_samples_yielded, + ) + if IterDataPipe.getstate_hook is not None: + return IterDataPipe.getstate_hook(state) + return state + + def __setstate__(self, state): + ( + self.datapipe, + self.buffer_size, + self._enabled, + self._seed, + self._buffer, + rng_state, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) = state + self._rng = random.Random() + self._rng.setstate(rng_state) + + def __del__(self) -> None: + self._buffer.clear() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/combining.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/combining.py new file mode 100644 index 0000000000000000000000000000000000000000..2e1d04f00c5ffe32dcf909527c9c3d74f29a17f8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/combining.py @@ -0,0 +1,716 @@ +# mypy: allow-untyped-defs +import copy as copymodule +import warnings +from abc import ABC, abstractmethod +from collections import deque +from collections.abc import Callable, Iterator, Sized +from typing import Any, Literal, TypeVar + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes._hook_iterator import _SnapshotState +from torch.utils.data.datapipes.datapipe import IterDataPipe +from torch.utils.data.datapipes.utils.common import _check_unpickable_fn, StreamWrapper + + +__all__ = [ + "ConcaterIterDataPipe", + "DemultiplexerIterDataPipe", + "ForkerIterDataPipe", + "MultiplexerIterDataPipe", + "ZipperIterDataPipe", +] + + +_T_co = TypeVar("_T_co", covariant=True) + + +@functional_datapipe("concat") +class ConcaterIterDataPipe(IterDataPipe): + r""" + Concatenates multiple Iterable DataPipes (functional name: ``concat``). + + The resulting DataPipe will yield all the elements from the first input DataPipe, before yielding from the subsequent ones. + + Args: + datapipes: Iterable DataPipes being concatenated + + Example: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> import random + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp1 = IterableWrapper(range(3)) + >>> dp2 = IterableWrapper(range(5)) + >>> list(dp1.concat(dp2)) + [0, 1, 2, 0, 1, 2, 3, 4] + """ + + datapipes: tuple[IterDataPipe] + + def __init__(self, *datapipes: IterDataPipe) -> None: + if len(datapipes) == 0: + raise ValueError("Expected at least one DataPipe, but got nothing") + if not all(isinstance(dp, IterDataPipe) for dp in datapipes): + raise TypeError("Expected all inputs to be `IterDataPipe`") + self.datapipes = datapipes # type: ignore[assignment] + + def __iter__(self) -> Iterator: + for dp in self.datapipes: + yield from dp + + def __len__(self) -> int: + # pyrefly: ignore [unsafe-overlap] + if all(isinstance(dp, Sized) for dp in self.datapipes): + # pyrefly: ignore [bad-argument-type] + return sum(len(dp) for dp in self.datapipes) + else: + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") + + +@functional_datapipe("fork") +class ForkerIterDataPipe(IterDataPipe): + r""" + Creates multiple instances of the same Iterable DataPipe (functional name: ``fork``). + + Args: + datapipe: Iterable DataPipe being copied + num_instances: number of instances of the datapipe to create + buffer_size: this restricts how far ahead the leading child DataPipe + can read relative to the slowest child DataPipe. + Defaults to ``1000``. Use ``-1`` for the unlimited buffer. + copy: copy strategy to use for items yielded by each branch. Supported + options are ``None`` for no copying, ``"shallow"`` for shallow object + copies, and ``"deep"`` for deep object copies. Defaults to ``None``. + + Note: + All branches of the forked pipeline return the identical object unless + the copy parameter is supplied. If the object is mutable or contains + mutable objects, changing them in one branch will affect all others. + + Example: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> from torchdata.datapipes.iter import IterableWrapper + >>> source_dp = IterableWrapper(range(5)) + >>> dp1, dp2 = source_dp.fork(num_instances=2) + >>> list(dp1) + [0, 1, 2, 3, 4] + >>> list(dp2) + [0, 1, 2, 3, 4] + """ + + def __new__( + cls, + datapipe: IterDataPipe, + num_instances: int, + buffer_size: int = 1000, + copy: Literal["shallow", "deep"] | None = None, + ): + if num_instances < 1: + raise ValueError( + f"Expected `num_instances` larger than 0, but {num_instances} is found" + ) + if num_instances == 1: + return datapipe + container = _ForkerIterDataPipe(datapipe, num_instances, buffer_size, copy) # type: ignore[abstract] + return [_ChildDataPipe(container, i) for i in range(num_instances)] + + +class _ContainerTemplate(ABC): + r"""Abstract class for container ``DataPipes``. The followings are three required methods.""" + + @abstractmethod + def get_next_element_by_instance(self, instance_id: int): ... + + @abstractmethod + def is_every_instance_exhausted(self) -> bool: ... + + @abstractmethod + def reset(self) -> None: ... + + @abstractmethod + def get_length_by_instance(self, instance_id: int): + r"""Raise TypeError if it's not supposed to be implemented to support `list(datapipe)`.""" + + +def _no_op(x): + return x + + +class _ForkerIterDataPipe(IterDataPipe, _ContainerTemplate): + r""" + Container to hold instance-specific information on behalf of ForkerIterDataPipe. + + It tracks the state of its child DataPipes, maintains the buffer, and yields the next value + as requested by the child DataPipes. + """ + + def __init__( + self, + datapipe: IterDataPipe, + num_instances: int, + buffer_size: int = 1000, + copy: Literal["shallow", "deep"] | None = None, + ) -> None: + self.main_datapipe = datapipe + self._datapipe_iterator: Iterator[Any] | None = None + self.num_instances = num_instances + self.buffer: deque = deque() + self.buffer_size = buffer_size + if self.buffer_size < 0: + warnings.warn( + "Unlimited buffer size is set for `fork`, " + "please be aware of OOM at random places", + UserWarning, + stacklevel=2, + ) + if copy is None: + self.copy_fn = _no_op + elif copy == "shallow": + self.copy_fn = copymodule.copy + elif copy == "deep": + self.copy_fn = copymodule.deepcopy + else: + raise ValueError( + f"Unknown copy method `{copy}` requested, choose one of None, `shallow` or `deep`." + ) + + self.child_pointers: list[int] = [ + 0 + ] * num_instances # Indicate the indices of the next element to get + self.slowest_ptr = 0 # The index to read by the slowest child + self.leading_ptr = 0 # The index to read by the fastest child + self.end_ptr: int | None = None # The index to stop child + self._child_stop: list[bool] = [True for _ in range(num_instances)] + + def __len__(self) -> int: + # pyrefly: ignore [bad-argument-type] + return len(self.main_datapipe) + + def get_next_element_by_instance(self, instance_id: int): + if self._datapipe_iterator is None and self._child_stop[instance_id]: + self._datapipe_iterator = iter(self.main_datapipe) + self._snapshot_state = _SnapshotState.Iterating + for i in range(self.num_instances): + self._child_stop[i] = False + try: + while not self._child_stop[instance_id]: + self.child_pointers[instance_id] += 1 + if ( + self.end_ptr is not None + and self.child_pointers[instance_id] == self.end_ptr + ): + self._child_stop[instance_id] = True + break + # Use buffer + if self.buffer and self.child_pointers[instance_id] <= self.leading_ptr: + idx = self.child_pointers[instance_id] - self.slowest_ptr - 1 + return_val = self.buffer[idx] + else: # Retrieve one element from main datapipe + self.leading_ptr = self.child_pointers[instance_id] + try: + return_val = next(self._datapipe_iterator) # type: ignore[arg-type] + self.buffer.append(return_val) + except StopIteration: + self._child_stop[instance_id] = True + self._datapipe_iterator = None + self.end_ptr = self.leading_ptr + continue + if self.child_pointers[instance_id] == self.slowest_ptr + 1: + new_min = min( + self.child_pointers + ) # Can optimize by avoiding the call to min() + if self.slowest_ptr < new_min: + self.slowest_ptr = new_min + self.buffer.popleft() + if ( + self.buffer_size >= 0 + and self.leading_ptr > self.buffer_size + self.slowest_ptr + ): + raise BufferError( + "ForkerIterDataPipe buffer overflow," + + f"buffer size {self.buffer_size} is insufficient." + ) + + yield self.copy_fn(return_val) # type: ignore[possibly-undefined] + finally: + self._child_stop[instance_id] = True + # Cleanup _datapipe_iterator for the case that fork exits earlier + if all(self._child_stop): + self._datapipe_iterator = None + self._cleanup() + + def is_every_instance_exhausted(self) -> bool: + return self.end_ptr is not None and all(self._child_stop) + + def get_length_by_instance(self, instance_id: int) -> int: + # pyrefly: ignore [bad-argument-type] + return len(self.main_datapipe) + + def reset(self) -> None: + self._datapipe_iterator = None + self.buffer = deque() + self.child_pointers = [0] * self.num_instances + self.slowest_ptr = 0 + self.leading_ptr = 0 + self.end_ptr = None + self._child_stop = [True for _ in range(self.num_instances)] + + def __getstate__(self): + state = ( + self.main_datapipe, + self.num_instances, + self.buffer_size, + self.copy_fn, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) + if IterDataPipe.getstate_hook is not None: + return IterDataPipe.getstate_hook(state) + return state + + def __setstate__(self, state): + ( + self.main_datapipe, + self.num_instances, + self.buffer_size, + self.copy_fn, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) = state + self._datapipe_iterator = None + self.buffer = deque() + self.child_pointers = [0] * self.num_instances + self.slowest_ptr = 0 + self.leading_ptr = 0 + self.end_ptr = None + self._child_stop = [True for _ in range(self.num_instances)] + + def _cleanup(self) -> None: + while self.buffer: + d = self.buffer.popleft() + StreamWrapper.close_streams(d) + + def __del__(self) -> None: + self._cleanup() + + +class _ChildDataPipe(IterDataPipe): + r""" + Iterable Datapipe that is a child of a main DataPipe. + + The instance of this class will pass its instance_id to get the next value from its main DataPipe. + + Note: + ChildDataPipe, like all other IterDataPipe, follows the single iterator per IterDataPipe constraint. + Since ChildDataPipes share a common buffer, when an iterator is created for one of the ChildDataPipes, + the previous iterators for all ChildDataPipes must be invalidated, with the exception when a ChildDataPipe + hasn't had an iterator created from it since the last invalidation. See the example below. + + Example: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> # Singler Iterator per IteraDataPipe Invalidation + >>> from torchdata.datapipes.iter import IterableWrapper + >>> source_dp = IterableWrapper(range(10)) + >>> cdp1, cdp2 = source_dp.fork(num_instances=2) + >>> it1, it2 = iter(cdp1), iter(cdp2) + >>> it3 = iter(cdp1) + >>> # The line above invalidates `it1` and `it2`, and resets `ForkerIterDataPipe`. + >>> it4 = iter(cdp2) + >>> # The line above doesn't invalidate `it3`, because an iterator for `cdp2` hasn't been created since + >>> # the last invalidation. + + Args: + main_datapipe: Main DataPipe with a method 'get_next_element_by_instance(instance_id)' + instance_id: integer identifier of this instance + """ + + _is_child_datapipe: bool = True + + def __init__(self, main_datapipe: IterDataPipe, instance_id: int) -> None: + if not isinstance(main_datapipe, _ContainerTemplate): + raise AssertionError("main_datapipe must implement _ContainerTemplate") + + self.main_datapipe: IterDataPipe = main_datapipe + self.instance_id = instance_id + + def __iter__(self): + # Note that the logic behind setting iterator ID and `reset` are handled within `hook_iterator` + # We want to separate the code for reset and yield, so that 'reset' executes before __next__ is called + return self.main_datapipe.get_next_element_by_instance(self.instance_id) + + def __len__(self) -> int: + return self.main_datapipe.get_length_by_instance(self.instance_id) + + # This method is called by `hook_iterator` in `_typing.py`. + def _set_main_datapipe_valid_iterator_id(self) -> int: + r""" + Update the valid iterator ID for both this DataPipe object and `main_datapipe`. + + `main_datapipe.reset()` is called when the ID is incremented to a new generation. + """ + # 1. First time any child iterator is created + if self.main_datapipe._valid_iterator_id is None: + self.main_datapipe._valid_iterator_id = 0 # type: ignore[attr-defined] + # 2. This instance was already in the same generation as `main_datapipe`, + # we need to increment the ID further by 1 + elif self.main_datapipe._valid_iterator_id == self._valid_iterator_id: # type: ignore[has-type] + self.main_datapipe._valid_iterator_id += 1 # type: ignore[attr-defined] + # Whenever a new generation of iterator is created, the `main_datapipe` must reset + if not self.main_datapipe.is_every_instance_exhausted(): + warnings.warn( + "Some child DataPipes are not exhausted when __iter__ is called. We are resetting " + "the buffer and each child DataPipe will read from the start again.", + UserWarning, + stacklevel=2, + ) + self.main_datapipe.reset() + # 3. Otherwise, the iterator is behind the others, so it will just need to catch up by setting + # the instance's iterator to match that of `main_datapipe` + self._valid_iterator_id = self.main_datapipe._valid_iterator_id + return self._valid_iterator_id + + # This method is called by `hook_iterator` in `_typing.py`. + def _check_valid_iterator_id(self, iterator_id) -> bool: + r"""Check the valid iterator ID against that of DataPipe object and that of `main_datapipe`.""" + return ( + iterator_id == self._valid_iterator_id + and iterator_id == self.main_datapipe._valid_iterator_id + ) + + +@functional_datapipe("demux") +class DemultiplexerIterDataPipe(IterDataPipe): + r""" + Splits the input DataPipe into multiple child DataPipes, using the given classification function (functional name: ``demux``). + + A list of the child DataPipes is returned from this operation. + + Args: + datapipe: Iterable DataPipe being filtered + num_instances: number of instances of the DataPipe to create + classifier_fn: a function that maps values to an integer within the range ``[0, num_instances - 1]`` or ``None`` + drop_none: defaults to ``False``, if ``True``, the function will skip over elements classified as ``None`` + buffer_size: this defines the maximum number of inputs that the buffer can hold across all child + DataPipes while waiting for their values to be yielded. + Defaults to ``1000``. Use ``-1`` for the unlimited buffer. + + Examples: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> from torchdata.datapipes.iter import IterableWrapper + >>> def odd_or_even(n): + ... return n % 2 + >>> source_dp = IterableWrapper(range(5)) + >>> dp1, dp2 = source_dp.demux(num_instances=2, classifier_fn=odd_or_even) + >>> list(dp1) + [0, 2, 4] + >>> list(dp2) + [1, 3] + >>> # It can also filter out any element that gets `None` from the `classifier_fn` + >>> def odd_or_even_no_zero(n): + ... return n % 2 if n != 0 else None + >>> dp1, dp2 = source_dp.demux( + ... num_instances=2, classifier_fn=odd_or_even_no_zero, drop_none=True + ... ) + >>> list(dp1) + [2, 4] + >>> list(dp2) + [1, 3] + """ + + def __new__( + cls, + datapipe: IterDataPipe, + num_instances: int, + classifier_fn: Callable[[_T_co], int | None], + drop_none: bool = False, + buffer_size: int = 1000, + ): + if num_instances < 1: + raise ValueError( + f"Expected `num_instances` larger than 0, but {num_instances} is found" + ) + + _check_unpickable_fn(classifier_fn) + + # When num_instances == 1, demux can be replaced by filter, + # but keep it as Demultiplexer for the sake of consistency + # like throwing Error when classification result is out of o range + container = _DemultiplexerIterDataPipe( + datapipe, num_instances, classifier_fn, drop_none, buffer_size + ) # type: ignore[abstract] + return [_ChildDataPipe(container, i) for i in range(num_instances)] + + +class _DemultiplexerIterDataPipe(IterDataPipe, _ContainerTemplate): + r""" + Container to hold instance-specific information on behalf of DemultiplexerIterDataPipe. + + It tracks the state of its child DataPipes, maintains the buffer, classifies and yields the next correct value + as requested by the child DataPipes. + """ + + def __init__( + self, + datapipe: IterDataPipe[_T_co], + num_instances: int, + classifier_fn: Callable[[_T_co], int | None], + drop_none: bool, + buffer_size: int, + ) -> None: + # pyrefly: ignore [invalid-type-var] + self.main_datapipe = datapipe + self._datapipe_iterator: Iterator[Any] | None = None + self.num_instances = num_instances + self.buffer_size = buffer_size + if self.buffer_size < 0: + warnings.warn( + "Unlimited buffer size is set for `demux`, " + "please be aware of OOM at random places", + UserWarning, + stacklevel=2, + ) + self.current_buffer_usage = 0 + # pyrefly: ignore [invalid-type-var] + self.child_buffers: list[deque[_T_co]] = [deque() for _ in range(num_instances)] + # pyrefly: ignore [invalid-type-var] + self.classifier_fn = classifier_fn + self.drop_none = drop_none + self.main_datapipe_exhausted = False + self._child_stop: list[bool] = [True for _ in range(num_instances)] + + def _find_next(self, instance_id: int) -> _T_co: # type: ignore[type-var] + while True: + if self.main_datapipe_exhausted or self._child_stop[instance_id]: + raise StopIteration + if self._datapipe_iterator is None: + raise ValueError( + "_datapipe_iterator has not been set, likely because this private method is called directly " + "without invoking get_next_element_by_instance() first." + ) + value = next(self._datapipe_iterator) + classification = self.classifier_fn(value) + if classification is None and self.drop_none: + StreamWrapper.close_streams(value) + continue + if ( + classification is None + or classification >= self.num_instances + or classification < 0 + ): + raise ValueError( + f"Output of the classification fn should be between 0 and {self.num_instances - 1}. " + + f"{classification} is returned." + ) + if classification == instance_id: + return value + self.child_buffers[classification].append(value) + self.current_buffer_usage += 1 + if self.buffer_size >= 0 and self.current_buffer_usage > self.buffer_size: + raise BufferError( + f"DemultiplexerIterDataPipe buffer overflow, buffer size {self.buffer_size} is insufficient." + ) + + def get_next_element_by_instance(self, instance_id: int): + if self._datapipe_iterator is None and self._child_stop[instance_id]: + self._datapipe_iterator = iter(self.main_datapipe) + self._snapshot_state = ( + _SnapshotState.Iterating + ) # This is necessary for the DataPipe to reset properly. + self.main_datapipe_exhausted = False + for i in range(self.num_instances): + self._child_stop[i] = False + + try: + while not self._child_stop[instance_id]: + if self.child_buffers[instance_id]: + self.current_buffer_usage -= 1 + yield self.child_buffers[instance_id].popleft() + else: + try: + yield self._find_next(instance_id) + except StopIteration: + self._child_stop[instance_id] = True + self.main_datapipe_exhausted = True + self._datapipe_iterator = None + finally: + self._child_stop[instance_id] = True + # Cleanup _datapipe_iterator for the case that demux exits earlier + if all(self._child_stop): + self._datapipe_iterator = None + if self.child_buffers[instance_id]: + self._cleanup(instance_id) + + def is_every_instance_exhausted(self) -> bool: + return self.main_datapipe_exhausted and all(self._child_stop) + + def get_length_by_instance(self, instance_id: int) -> int: + raise TypeError + + def reset(self) -> None: + self._datapipe_iterator = None + self.current_buffer_usage = 0 + self.child_buffers = [deque() for _ in range(self.num_instances)] + self._child_stop = [True for _ in range(self.num_instances)] + self.main_datapipe_exhausted = False + + def __getstate__(self): + state = ( + self.main_datapipe, + self.num_instances, + self.buffer_size, + self.classifier_fn, + self.drop_none, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) + if IterDataPipe.getstate_hook is not None: + return IterDataPipe.getstate_hook(state) + return state + + def __setstate__(self, state): + ( + self.main_datapipe, + self.num_instances, + self.buffer_size, + self.classifier_fn, + self.drop_none, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) = state + self._datapipe_iterator = None + self.current_buffer_usage = 0 + self.child_buffers = [deque() for _ in range(self.num_instances)] + self._child_stop = [True for _ in range(self.num_instances)] + self.main_datapipe_exhausted = False + + def _cleanup(self, instance_id: int | None = None) -> None: + ids = ( + range(self.num_instances) + if instance_id is None + else [ + instance_id, + ] + ) + for i in ids: + q = self.child_buffers[i] + while q: + d = q.popleft() + StreamWrapper.close_streams(d) + + def __del__(self) -> None: + self._cleanup() + + +@functional_datapipe("mux") +class MultiplexerIterDataPipe(IterDataPipe): + r""" + Yields one element at a time from each of the input Iterable DataPipes (functional name: ``mux``). + + As in, one element from the 1st input DataPipe, then one element from the 2nd DataPipe in the next iteration, + and so on. It ends when the shortest input DataPipe is exhausted. + + Args: + datapipes: Iterable DataPipes that will take turn to yield their elements, until the shortest DataPipe is exhausted + + Example: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp1, dp2, dp3 = ( + ... IterableWrapper(range(3)), + ... IterableWrapper(range(10, 15)), + ... IterableWrapper(range(20, 25)), + ... ) + >>> list(dp1.mux(dp2, dp3)) + [0, 10, 20, 1, 11, 21, 2, 12, 22] + """ + + def __init__(self, *datapipes) -> None: + self.datapipes = datapipes + self.buffer: list = [] # Store values to be yielded only when every iterator provides one + + def __iter__(self): + iterators = [iter(x) for x in self.datapipes] + while iterators: + for it in iterators: + try: + value = next(it) + self.buffer.append(value) + except StopIteration: + self.buffer.clear() + return + yield from self.buffer + self.buffer.clear() + + def __len__(self) -> int: + if all(isinstance(dp, Sized) for dp in self.datapipes): + return min(len(dp) for dp in self.datapipes) * len(self.datapipes) + else: + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") + + def reset(self) -> None: + self.buffer = [] + + def __getstate__(self): + state = ( + self.datapipes, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) + if IterDataPipe.getstate_hook is not None: + return IterDataPipe.getstate_hook(state) + return state + + def __setstate__(self, state): + ( + self.datapipes, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) = state + self.buffer = [] + + def __del__(self) -> None: + self.buffer.clear() + + +@functional_datapipe("zip") +class ZipperIterDataPipe(IterDataPipe[tuple[_T_co]]): + r""" + Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``). + + The output is stopped as soon as the shortest input DataPipe is exhausted. + + Args: + *datapipes: Iterable DataPipes being aggregated + + Example: + >>> # xdoctest: +REQUIRES(module:torchdata) + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp1, dp2, dp3 = ( + ... IterableWrapper(range(5)), + ... IterableWrapper(range(10, 15)), + ... IterableWrapper(range(20, 25)), + ... ) + >>> list(dp1.zip(dp2, dp3)) + [(0, 10, 20), (1, 11, 21), (2, 12, 22), (3, 13, 23), (4, 14, 24)] + """ + + datapipes: tuple[IterDataPipe] + + def __init__(self, *datapipes: IterDataPipe) -> None: + if not all(isinstance(dp, IterDataPipe) for dp in datapipes): + raise TypeError( + "All inputs are required to be `IterDataPipe` for `ZipIterDataPipe`." + ) + super().__init__() + self.datapipes = datapipes # type: ignore[assignment] + + def __iter__(self) -> Iterator[tuple[_T_co]]: + iterators = [iter(datapipe) for datapipe in self.datapipes] + yield from zip(*iterators, strict=False) + + def __len__(self) -> int: + # pyrefly: ignore [unsafe-overlap] + if all(isinstance(dp, Sized) for dp in self.datapipes): + # pyrefly: ignore [bad-argument-type] + return min(len(dp) for dp in self.datapipes) + else: + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/filelister.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/filelister.py new file mode 100644 index 0000000000000000000000000000000000000000..352d3c01e12d278cb8e1308ee78feff0610808bc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/filelister.py @@ -0,0 +1,67 @@ +from collections.abc import Iterator, Sequence + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import IterDataPipe +from torch.utils.data.datapipes.iter.utils import IterableWrapperIterDataPipe +from torch.utils.data.datapipes.utils.common import get_file_pathnames_from_root + + +__all__ = ["FileListerIterDataPipe"] + + +@functional_datapipe("list_files") +class FileListerIterDataPipe(IterDataPipe[str]): + r""" + Given path(s) to the root directory, yields file pathname(s) (path + filename) of files within the root directory. + + Multiple root directories can be provided (functional name: ``list_files``). + + Args: + root: Root directory or a sequence of root directories + masks: Unix style filter string or string list for filtering file name(s) + recursive: Whether to return pathname from nested directories or not + abspath: Whether to return relative pathname or absolute pathname + non_deterministic: Whether to return pathname in sorted order or not. + If ``False``, the results yielded from each root directory will be sorted + length: Nominal length of the datapipe + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import FileLister + >>> dp = FileLister(root=".", recursive=True) + >>> list(dp) + ['example.py', './data/data.tar'] + """ + + def __init__( + self, + root: str | Sequence[str] | IterDataPipe = ".", + masks: str | list[str] = "", + *, + recursive: bool = False, + abspath: bool = False, + non_deterministic: bool = False, + length: int = -1, + ) -> None: + super().__init__() + if isinstance(root, str): + root = [root] + if not isinstance(root, IterDataPipe): + root = IterableWrapperIterDataPipe(root) + self.datapipe: IterDataPipe = root + self.masks: str | list[str] = masks + self.recursive: bool = recursive + self.abspath: bool = abspath + self.non_deterministic: bool = non_deterministic + self.length: int = length + + def __iter__(self) -> Iterator[str]: + for path in self.datapipe: + yield from get_file_pathnames_from_root( + path, self.masks, self.recursive, self.abspath, self.non_deterministic + ) + + def __len__(self) -> int: + if self.length == -1: + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") + return self.length diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/fileopener.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/fileopener.py new file mode 100644 index 0000000000000000000000000000000000000000..e77f7a4c8e660ec0e2ff5374fcdc8a30c474ea03 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/fileopener.py @@ -0,0 +1,79 @@ +from collections.abc import Iterable, Iterator +from io import IOBase + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import IterDataPipe +from torch.utils.data.datapipes.utils.common import get_file_binaries_from_pathnames + + +__all__ = [ + "FileOpenerIterDataPipe", +] + + +@functional_datapipe("open_files") +class FileOpenerIterDataPipe(IterDataPipe[tuple[str, IOBase]]): + r""" + Given pathnames, opens files and yield pathname and file stream in a tuple (functional name: ``open_files``). + + Args: + datapipe: Iterable datapipe that provides pathnames + mode: An optional string that specifies the mode in which + the file is opened by ``open()``. It defaults to ``r``, other options are + ``b`` for reading in binary mode and ``t`` for text mode. + encoding: An optional string that specifies the encoding of the + underlying file. It defaults to ``None`` to match the default encoding of ``open``. + length: Nominal length of the datapipe + + Note: + The opened file handles will be closed by Python's GC periodically. Users can choose + to close them explicitly. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import ( + ... FileLister, + ... FileOpener, + ... StreamReader, + ... ) + >>> dp = FileLister(root=".").filter(lambda fname: fname.endswith(".txt")) + >>> dp = FileOpener(dp) + >>> dp = StreamReader(dp) + >>> list(dp) + [('./abc.txt', 'abc')] + """ + + def __init__( + self, + datapipe: Iterable[str], + mode: str = "r", + encoding: str | None = None, + length: int = -1, + ) -> None: + super().__init__() + self.datapipe: Iterable[str] = datapipe + self.mode: str = mode + self.encoding: str | None = encoding + + if self.mode not in ("b", "t", "rb", "rt", "r"): + raise ValueError(f"Invalid mode {mode}") + # TODO: enforce typing for each instance based on mode, otherwise + # `argument_validation` with this DataPipe may be potentially broken + + if "b" in mode and encoding is not None: + raise ValueError("binary mode doesn't take an encoding argument") + + self.length: int = length + + # Remove annotation due to 'IOBase' is a general type and true type + # is determined at runtime based on mode. Some `DataPipe` requiring + # a subtype would cause mypy error. + def __iter__(self) -> Iterator[tuple[str, IOBase]]: + yield from get_file_binaries_from_pathnames( + self.datapipe, self.mode, self.encoding + ) + + def __len__(self) -> int: + if self.length == -1: + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") + return self.length diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/grouping.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/grouping.py new file mode 100644 index 0000000000000000000000000000000000000000..dd3a3775235198361025142dbb03d0abdbe17a6a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/grouping.py @@ -0,0 +1,325 @@ +# mypy: allow-untyped-defs +from collections import defaultdict +from collections.abc import Callable, Iterator, Sized +from typing import Any, NoReturn, TypeVar + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import DataChunk, IterDataPipe +from torch.utils.data.datapipes.utils.common import _check_unpickable_fn + + +__all__ = [ + "BatcherIterDataPipe", + "GrouperIterDataPipe", + "UnBatcherIterDataPipe", +] + + +_T_co = TypeVar("_T_co", covariant=True) + + +def __getattr__(name: str) -> NoReturn: + raise AttributeError(f"module {__name__} has no attribute {name}") + + +@functional_datapipe("batch") +class BatcherIterDataPipe(IterDataPipe[DataChunk]): + r""" + Creates mini-batches of data (functional name: ``batch``). + + An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``, or ``length % batch_size`` for the + last batch if ``drop_last`` is set to ``False``. + + Args: + datapipe: Iterable DataPipe being batched + batch_size: The size of each batch + drop_last: Option to drop the last batch if it's not full + wrapper_class: wrapper to apply onto each batch (type ``List``) before yielding, + defaults to ``DataChunk`` + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp = IterableWrapper(range(10)) + >>> dp = dp.batch(batch_size=3, drop_last=True) + >>> list(dp) + [[0, 1, 2], [3, 4, 5], [6, 7, 8]] + """ + + datapipe: IterDataPipe + batch_size: int + drop_last: bool + + def __init__( + self, + datapipe: IterDataPipe, + batch_size: int, + drop_last: bool = False, + wrapper_class: type[DataChunk] = DataChunk, + ) -> None: + if batch_size <= 0: + raise AssertionError("Batch size is required to be larger than 0!") + super().__init__() + self.datapipe = datapipe + self.batch_size = batch_size + self.drop_last = drop_last + self.wrapper_class = wrapper_class + + def __iter__(self) -> Iterator[DataChunk]: + batch: list = [] + for x in self.datapipe: + batch.append(x) + if len(batch) == self.batch_size: + yield self.wrapper_class(batch) + batch = [] + if len(batch) > 0: + if not self.drop_last: + yield self.wrapper_class(batch) + + def __len__(self) -> int: + # pyrefly: ignore [unsafe-overlap] + if isinstance(self.datapipe, Sized): + if self.drop_last: + return len(self.datapipe) // self.batch_size + else: + return (len(self.datapipe) + self.batch_size - 1) // self.batch_size + else: + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") + + +@functional_datapipe("unbatch") +class UnBatcherIterDataPipe(IterDataPipe): + r""" + Undos batching of data (functional name: ``unbatch``). + + In other words, it flattens the data up to the specified level within a batched DataPipe. + + Args: + datapipe: Iterable DataPipe being un-batched + unbatch_level: Defaults to ``1`` (only flattening the top level). If set to ``2``, + it will flatten the top two levels, and ``-1`` will flatten the entire DataPipe. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> source_dp = IterableWrapper([[[0, 1], [2]], [[3, 4], [5]], [[6]]]) + >>> dp1 = source_dp.unbatch() + >>> list(dp1) + [[0, 1], [2], [3, 4], [5], [6]] + >>> dp2 = source_dp.unbatch(unbatch_level=2) + >>> list(dp2) + [0, 1, 2, 3, 4, 5, 6] + """ + + def __init__(self, datapipe: IterDataPipe, unbatch_level: int = 1) -> None: + self.datapipe = datapipe + self.unbatch_level = unbatch_level + + def __iter__(self): + for element in self.datapipe: + yield from self._dive(element, unbatch_level=self.unbatch_level) + + def _dive(self, element, unbatch_level): + if unbatch_level < -1: + raise ValueError("unbatch_level must be -1 or >= 0") + if unbatch_level == -1: + if isinstance(element, (list, DataChunk)): + for item in element: + yield from self._dive(item, unbatch_level=-1) + else: + yield element + elif unbatch_level == 0: + yield element + else: + if isinstance(element, (list, DataChunk)): + for item in element: + yield from self._dive(item, unbatch_level=unbatch_level - 1) + else: + raise IndexError( + f"unbatch_level {self.unbatch_level} exceeds the depth of the DataPipe" + ) + + +@functional_datapipe("groupby") +class GrouperIterDataPipe(IterDataPipe[DataChunk]): + r""" + Groups data from IterDataPipe by keys from ``group_key_fn``, yielding a ``DataChunk`` with batch size up to ``group_size``. + + (functional name: ``groupby``). + + The samples are read sequentially from the source ``datapipe``, and a batch of samples belonging to the same group + will be yielded as soon as the size of the batch reaches ``group_size``. When the buffer is full, + the DataPipe will yield the largest batch with the same key, provided that its size is larger + than ``guaranteed_group_size``. If its size is smaller, it will be dropped if ``drop_remaining=True``. + + After iterating through the entirety of source ``datapipe``, everything not dropped due to the buffer capacity + will be yielded from the buffer, even if the group sizes are smaller than ``guaranteed_group_size``. + + Args: + datapipe: Iterable datapipe to be grouped + group_key_fn: Function used to generate group key from the data of the source datapipe + keep_key: Option to yield the matching key along with the items in a tuple, + resulting in `(key, [items])` otherwise returning [items] + buffer_size: The size of buffer for ungrouped data + group_size: The max size of each group, a batch is yielded as soon as it reaches this size + guaranteed_group_size: The guaranteed minimum group size to be yielded in case the buffer is full + drop_remaining: Specifies if the group smaller than ``guaranteed_group_size`` will be dropped from buffer + when the buffer is full + + Example: + >>> import os + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> def group_fn(file): + ... return os.path.basename(file).split(".")[0] + >>> source_dp = IterableWrapper( + ... ["a.png", "b.png", "a.json", "b.json", "a.jpg", "c.json"] + ... ) + >>> dp0 = source_dp.groupby(group_key_fn=group_fn) + >>> list(dp0) + [['a.png', 'a.json', 'a.jpg'], ['b.png', 'b.json'], ['c.json']] + >>> # A group is yielded as soon as its size equals to `group_size` + >>> dp1 = source_dp.groupby(group_key_fn=group_fn, group_size=2) + >>> list(dp1) + [['a.png', 'a.json'], ['b.png', 'b.json'], ['a.jpg'], ['c.json']] + >>> # Scenario where `buffer` is full, and group 'a' needs to be yielded since its size > `guaranteed_group_size` + >>> dp2 = source_dp.groupby( + ... group_key_fn=group_fn, + ... buffer_size=3, + ... group_size=3, + ... guaranteed_group_size=2, + ... ) + >>> list(dp2) + [['a.png', 'a.json'], ['b.png', 'b.json'], ['a.jpg'], ['c.json']] + """ + + def __init__( + self, + datapipe: IterDataPipe[_T_co], + group_key_fn: Callable[[_T_co], Any], + *, + keep_key: bool = False, + buffer_size: int = 10000, + group_size: int | None = None, + guaranteed_group_size: int | None = None, + drop_remaining: bool = False, + ) -> None: + _check_unpickable_fn(group_key_fn) + # pyrefly: ignore [invalid-type-var] + self.datapipe = datapipe + # pyrefly: ignore [invalid-type-var] + self.group_key_fn = group_key_fn + + self.keep_key = keep_key + self.max_buffer_size = buffer_size + self.buffer_elements: defaultdict[Any, list] = defaultdict(list) + self.curr_buffer_size = 0 + self.group_size = group_size + self.guaranteed_group_size = None + if group_size is not None and buffer_size is not None: + if not (0 < group_size <= buffer_size): + raise AssertionError("group_size must be > 0 and <= buffer_size") + self.guaranteed_group_size = group_size + if guaranteed_group_size is not None: + if group_size is None or not (0 < guaranteed_group_size <= group_size): + raise AssertionError( + "guaranteed_group_size must be > 0 and <= group_size and group_size must be set" + ) + self.guaranteed_group_size = guaranteed_group_size + self.drop_remaining = drop_remaining + self.wrapper_class = DataChunk + + def _remove_biggest_key(self): + biggest_key = None + biggest_size = 0 + result_to_yield = None + for findkey in self.buffer_elements: + if len(self.buffer_elements[findkey]) > biggest_size: + biggest_size = len(self.buffer_elements[findkey]) + biggest_key = findkey + + if ( + self.guaranteed_group_size is not None + and biggest_size < self.guaranteed_group_size + and not self.drop_remaining + ): + raise RuntimeError( + "Failed to group items", str(self.buffer_elements[biggest_key]) + ) + + if ( + self.guaranteed_group_size is None + or biggest_size >= self.guaranteed_group_size + ): + result_to_yield = self.buffer_elements[biggest_key] + + self.curr_buffer_size -= biggest_size + del self.buffer_elements[biggest_key] + + return result_to_yield + + def __iter__(self): + for x in self.datapipe: + key = self.group_key_fn(x) + + self.buffer_elements[key].append(x) + self.curr_buffer_size += 1 + + if self.group_size is not None and self.group_size == len( + self.buffer_elements[key] + ): + result: DataChunk[Any] = self.wrapper_class(self.buffer_elements[key]) + yield (key, result) if self.keep_key else result + self.curr_buffer_size -= len(self.buffer_elements[key]) + del self.buffer_elements[key] + + if self.curr_buffer_size == self.max_buffer_size: + result_to_yield = self._remove_biggest_key() + if result_to_yield is not None: + result = self.wrapper_class(result_to_yield) + yield (key, result) if self.keep_key else result + + for key in tuple(self.buffer_elements.keys()): + result = self.wrapper_class(self.buffer_elements.pop(key)) + self.curr_buffer_size -= len(result) + yield (key, result) if self.keep_key else result + + def reset(self) -> None: + self.curr_buffer_size = 0 + self.buffer_elements = defaultdict(list) + + def __getstate__(self): + state = ( + self.datapipe, + self.group_key_fn, + self.keep_key, + self.max_buffer_size, + self.group_size, + self.guaranteed_group_size, + self.drop_remaining, + self.wrapper_class, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) + if IterDataPipe.getstate_hook is not None: + return IterDataPipe.getstate_hook(state) + return state + + def __setstate__(self, state): + ( + self.datapipe, + self.group_key_fn, + self.keep_key, + self.max_buffer_size, + self.group_size, + self.guaranteed_group_size, + self.drop_remaining, + self.wrapper_class, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) = state + self.curr_buffer_size = 0 + self.buffer_elements = defaultdict(list) + + def __del__(self) -> None: + self.buffer_elements.clear() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/routeddecoder.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/routeddecoder.py new file mode 100644 index 0000000000000000000000000000000000000000..ba4d708a0a318bd75ab67f456b0a5ef2f24b2c81 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/routeddecoder.py @@ -0,0 +1,70 @@ +from collections.abc import Callable, Iterable, Iterator, Sized +from io import BufferedIOBase +from typing import Any + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import IterDataPipe +from torch.utils.data.datapipes.utils.common import _deprecation_warning +from torch.utils.data.datapipes.utils.decoder import ( + basichandlers as decoder_basichandlers, + Decoder, + extension_extract_fn, + imagehandler as decoder_imagehandler, +) + + +__all__ = ["RoutedDecoderIterDataPipe"] + + +@functional_datapipe("routed_decode") +class RoutedDecoderIterDataPipe(IterDataPipe[tuple[str, Any]]): + r""" + Decodes binary streams from input DataPipe, yields pathname and decoded data in a tuple. + + (functional name: ``routed_decode``) + + Args: + datapipe: Iterable datapipe that provides pathname and binary stream in tuples + handlers: Optional user defined decoder handlers. If ``None``, basic and image decoder + handlers will be set as default. If multiple handles are provided, the priority + order follows the order of handlers (the first handler has the top priority) + key_fn: Function for decoder to extract key from pathname to dispatch handlers. + Default is set to extract file extension from pathname + + Note: + When ``key_fn`` is specified returning anything other than extension, the default + handler will not work and users need to specify custom handler. Custom handler + could use regex to determine the eligibility to handle data. + """ + + def __init__( + self, + datapipe: Iterable[tuple[str, BufferedIOBase]], + *handlers: Callable, + key_fn: Callable = extension_extract_fn, + ) -> None: + super().__init__() + self.datapipe: Iterable[tuple[str, BufferedIOBase]] = datapipe + if not handlers: + handlers = (decoder_basichandlers, decoder_imagehandler("torch")) + self.decoder = Decoder(*handlers, key_fn=key_fn) + _deprecation_warning( + type(self).__name__, + deprecation_version="1.12", + removal_version="1.13", + old_functional_name="routed_decode", + ) + + def add_handler(self, *handler: Callable) -> None: + self.decoder.add_handler(*handler) + + def __iter__(self) -> Iterator[tuple[str, Any]]: + for data in self.datapipe: + pathname = data[0] + result = self.decoder(data) + yield (pathname, result[pathname]) + + def __len__(self) -> int: + if isinstance(self.datapipe, Sized): + return len(self.datapipe) + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/selecting.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/selecting.py new file mode 100644 index 0000000000000000000000000000000000000000..afb0e91d8557911aae6f20d830667c79f7764cc5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/selecting.py @@ -0,0 +1,102 @@ +# mypy: allow-untyped-defs +from collections.abc import Callable, Iterator +from typing import TypeVar + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.dataframe import dataframe_wrapper as df_wrapper +from torch.utils.data.datapipes.datapipe import IterDataPipe +from torch.utils.data.datapipes.utils.common import ( + _check_unpickable_fn, + StreamWrapper, + validate_input_col, +) + + +__all__ = ["FilterIterDataPipe"] + + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) + + +@functional_datapipe("filter") +class FilterIterDataPipe(IterDataPipe[_T_co]): + r""" + Filters out elements from the source datapipe according to input ``filter_fn`` (functional name: ``filter``). + + Args: + datapipe: Iterable DataPipe being filtered + filter_fn: Customized function mapping an element to a boolean. + input_col: Index or indices of data which ``filter_fn`` is applied, such as: + + - ``None`` as default to apply ``filter_fn`` to the data directly. + - Integer(s) is used for list/tuple. + - Key(s) is used for dict. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> def is_even(n): + ... return n % 2 == 0 + >>> dp = IterableWrapper(range(5)) + >>> filter_dp = dp.filter(filter_fn=is_even) + >>> list(filter_dp) + [0, 2, 4] + """ + + datapipe: IterDataPipe[_T_co] + filter_fn: Callable + + def __init__( + self, + datapipe: IterDataPipe[_T_co], + filter_fn: Callable, + input_col=None, + ) -> None: + super().__init__() + self.datapipe = datapipe + + _check_unpickable_fn(filter_fn) + self.filter_fn = filter_fn # type: ignore[assignment] + + self.input_col = input_col + validate_input_col(filter_fn, input_col) + + def _apply_filter_fn(self, data) -> bool: + if self.input_col is None: + return self.filter_fn(data) + elif isinstance(self.input_col, (list, tuple)): + args = tuple(data[col] for col in self.input_col) + return self.filter_fn(*args) + else: + return self.filter_fn(data[self.input_col]) + + def __iter__(self) -> Iterator[_T_co]: + for data in self.datapipe: + condition, filtered = self._returnIfTrue(data) + if condition: + yield filtered + else: + StreamWrapper.close_streams(data) + + def _returnIfTrue(self, data: _T) -> tuple[bool, _T]: + condition = self._apply_filter_fn(data) + + if df_wrapper.is_column(condition): + # We are operating on DataFrames filter here + result = [] + for idx, mask in enumerate(df_wrapper.iterate(condition)): + if mask: + result.append(df_wrapper.get_item(data, idx)) + if result: + return True, df_wrapper.concat(result) + else: + return False, None # type: ignore[return-value] + + if not isinstance(condition, bool): + raise ValueError( + "Boolean output is required for `filter_fn` of FilterIterDataPipe, got", + type(condition), + ) + + return condition, data diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/sharding.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/sharding.py new file mode 100644 index 0000000000000000000000000000000000000000..dd3378016e915fb58c3a6e5d66e1ba7187b267c2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/sharding.py @@ -0,0 +1,106 @@ +# mypy: allow-untyped-defs +from collections.abc import Sized +from enum import IntEnum +from typing import NoReturn + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import IterDataPipe + + +__all__ = [ + "SHARDING_PRIORITIES", + "ShardingFilterIterDataPipe", +] + + +class SHARDING_PRIORITIES(IntEnum): + DEFAULT = 1 + DISTRIBUTED = 2 + MULTIPROCESSING = 3 + + +class _ShardingIterDataPipe(IterDataPipe): + def apply_sharding( + self, + num_of_instances: int, + instance_id: int, + sharding_group: SHARDING_PRIORITIES, + ) -> NoReturn: + raise NotImplementedError + + +@functional_datapipe("sharding_filter") +class ShardingFilterIterDataPipe(_ShardingIterDataPipe): + r""" + Wrapper that allows DataPipe to be sharded (functional name: ``sharding_filter``). + + After ``apply_sharding`` is called, each instance of the DataPipe (on different workers) will have every `n`-th element of the + original DataPipe, where `n` equals to the number of instances. + + Args: + source_datapipe: Iterable DataPipe that will be sharded + """ + + def __init__( + self, source_datapipe: IterDataPipe, sharding_group_filter=None + ) -> None: + super().__init__() + self.source_datapipe = source_datapipe + self.sharding_group_filter = sharding_group_filter + self.groups: dict[int, tuple[int, int]] = {} + self.num_of_instances = 1 + self.instance_id = 0 + self._update_num_of_instances() + + def apply_sharding( + self, num_of_instances, instance_id, sharding_group=SHARDING_PRIORITIES.DEFAULT + ): + if instance_id >= num_of_instances: + raise ValueError( + f"instance_id({instance_id}) should be smaller than num_of_instances({num_of_instances})" + ) + if sharding_group == SHARDING_PRIORITIES.DEFAULT: + if len(self.groups) and SHARDING_PRIORITIES.DEFAULT not in self.groups: + raise RuntimeError( + "ShardingFilter cannot mix DEFAULT and non DEFAULT groups" + ) + else: + if SHARDING_PRIORITIES.DEFAULT in self.groups: + raise RuntimeError( + "ShardingFilter cannot mix DEFAULT and non DEFAULT groups" + ) + self.groups[sharding_group] = (num_of_instances, instance_id) + self._update_num_of_instances() + + def _update_num_of_instances(self) -> None: + sorted_sharding_groups = [ + self.groups[key] + for key in sorted(self.groups.keys()) + if self.sharding_group_filter is None or key == self.sharding_group_filter + ] + + sorted_sharding_groups.reverse() + + self.num_of_instances = 1 + self.instance_id = 0 + + for group_num_of_instances, group_instance_id in sorted_sharding_groups: + self.instance_id += self.num_of_instances * group_instance_id + self.num_of_instances *= group_num_of_instances + + def __iter__(self): + for i, item in enumerate(self.source_datapipe): + if i % self.num_of_instances == self.instance_id: + yield item + + def __len__(self) -> int: + # pyrefly: ignore [unsafe-overlap] + if isinstance(self.source_datapipe, Sized): + return len(self.source_datapipe) // self.num_of_instances + ( + 1 + if ( + self.instance_id < len(self.source_datapipe) % self.num_of_instances + ) + else 0 + ) + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/streamreader.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/streamreader.py new file mode 100644 index 0000000000000000000000000000000000000000..1129c06548e1f406629e25a5a2f558dea3a1475e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/streamreader.py @@ -0,0 +1,45 @@ +from collections.abc import Iterator +from io import IOBase + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import IterDataPipe + + +__all__ = ["StreamReaderIterDataPipe"] + + +@functional_datapipe("read_from_stream") +class StreamReaderIterDataPipe(IterDataPipe[tuple[str, bytes]]): + r""" + Given IO streams and their label names, yield bytes with label name as tuple. + + (functional name: ``read_from_stream``). + + Args: + datapipe: Iterable DataPipe provides label/URL and byte stream + chunk: Number of bytes to be read from stream per iteration. + If ``None``, all bytes will be read until the EOF. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper, StreamReader + >>> from io import StringIO + >>> dp = IterableWrapper([("alphabet", StringIO("abcde"))]) + >>> list(StreamReader(dp, chunk=1)) + [('alphabet', 'a'), ('alphabet', 'b'), ('alphabet', 'c'), ('alphabet', 'd'), ('alphabet', 'e')] + """ + + def __init__( + self, datapipe: IterDataPipe[tuple[str, IOBase]], chunk: int | None = None + ) -> None: + self.datapipe = datapipe + self.chunk = chunk + + def __iter__(self) -> Iterator[tuple[str, bytes]]: + for furl, stream in self.datapipe: + while True: + d = stream.read(self.chunk) + if not d: + stream.close() + break + yield (furl, d) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e45ddab282f7b975732b28ab88339f979792646a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/iter/utils.py @@ -0,0 +1,60 @@ +import copy +import warnings +from collections.abc import Iterable, Iterator, Sized +from typing import TypeVar + +from torch.utils.data.datapipes.datapipe import IterDataPipe + + +_T = TypeVar("_T") + +__all__ = ["IterableWrapperIterDataPipe"] + + +class IterableWrapperIterDataPipe(IterDataPipe[_T]): + r""" + Wraps an iterable object to create an IterDataPipe. + + Args: + iterable: Iterable object to be wrapped into an IterDataPipe + deepcopy: Option to deepcopy input iterable object for each + iterator. The copy is made when the first element is read in ``iter()``. + + .. note:: + If ``deepcopy`` is explicitly set to ``False``, users should ensure + that the data pipeline doesn't contain any in-place operations over + the iterable instance to prevent data inconsistency across iterations. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.iter import IterableWrapper + >>> dp = IterableWrapper(range(10)) + >>> list(dp) + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + """ + + def __init__(self, iterable: Iterable[_T], deepcopy: bool = True) -> None: + self.iterable = iterable + self.deepcopy = deepcopy + + def __iter__(self) -> Iterator[_T]: + source_data = self.iterable + if self.deepcopy: + try: + source_data = copy.deepcopy(self.iterable) + # For the case that data cannot be deep-copied, + # all in-place operations will affect iterable variable. + # When this DataPipe is iterated second time, it will + # yield modified items. + except TypeError: + warnings.warn( + "The input iterable can not be deepcopied, " + "please be aware of in-place modification would affect source data.", + stacklevel=2, + ) + yield from source_data + + def __len__(self) -> int: + if isinstance(self.iterable, Sized): + return len(self.iterable) + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bc555e8fdac26039d36c4c1e1ba8309bfa8b4e5a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/__init__.py @@ -0,0 +1,20 @@ +# Functional DataPipe +from torch.utils.data.datapipes.map.callable import MapperMapDataPipe as Mapper +from torch.utils.data.datapipes.map.combinatorics import ( + ShufflerIterDataPipe as Shuffler, +) +from torch.utils.data.datapipes.map.combining import ( + ConcaterMapDataPipe as Concater, + ZipperMapDataPipe as Zipper, +) +from torch.utils.data.datapipes.map.grouping import BatcherMapDataPipe as Batcher +from torch.utils.data.datapipes.map.utils import ( + SequenceWrapperMapDataPipe as SequenceWrapper, +) + + +__all__ = ["Batcher", "Concater", "Mapper", "SequenceWrapper", "Shuffler", "Zipper"] + +# Please keep this list sorted +if __all__ != sorted(__all__): + raise AssertionError("__all__ is not sorted") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/callable.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/callable.py new file mode 100644 index 0000000000000000000000000000000000000000..3696d34b2a815599709bb09d9b0dfcaca988a6eb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/callable.py @@ -0,0 +1,67 @@ +# mypy: allow-untyped-defs +from collections.abc import Callable +from typing import TypeVar + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import MapDataPipe +from torch.utils.data.datapipes.utils.common import _check_unpickable_fn + + +__all__ = ["MapperMapDataPipe", "default_fn"] + + +_T_co = TypeVar("_T_co", covariant=True) + + +# Default function to return each item directly +# In order to keep datapipe picklable, eliminates the usage +# of python lambda function +def default_fn(data): + return data + + +@functional_datapipe("map") +class MapperMapDataPipe(MapDataPipe[_T_co]): + r""" + Apply the input function over each item from the source DataPipe (functional name: ``map``). + + The function can be any regular Python function or partial object. Lambda + function is not recommended as it is not supported by pickle. + + Args: + datapipe: Source MapDataPipe + fn: Function being applied to each item + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper, Mapper + >>> def add_one(x): + ... return x + 1 + >>> dp = SequenceWrapper(range(10)) + >>> map_dp_1 = dp.map(add_one) + >>> list(map_dp_1) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + >>> map_dp_2 = Mapper(dp, lambda x: x + 1) + >>> list(map_dp_2) + [1, 2, 3, 4, 5, 6, 7, 8, 9, 10] + """ + + datapipe: MapDataPipe + fn: Callable + + def __init__( + self, + datapipe: MapDataPipe, + fn: Callable = default_fn, + ) -> None: + super().__init__() + self.datapipe = datapipe + _check_unpickable_fn(fn) + self.fn = fn # type: ignore[assignment] + + def __len__(self) -> int: + # pyrefly: ignore [bad-argument-type] + return len(self.datapipe) + + def __getitem__(self, index) -> _T_co: + return self.fn(self.datapipe[index]) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/combinatorics.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/combinatorics.py new file mode 100644 index 0000000000000000000000000000000000000000..af4792fc805b824d45a966851e0fae2d853ff99f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/combinatorics.py @@ -0,0 +1,132 @@ +# mypy: allow-untyped-defs +import random +from collections.abc import Iterator +from typing import TypeVar + +import torch +from torch.utils.data.datapipes.datapipe import IterDataPipe, MapDataPipe + + +__all__ = ["ShufflerIterDataPipe"] + + +_T_co = TypeVar("_T_co", covariant=True) + + +# @functional_datapipe('shuffle') +class ShufflerIterDataPipe(IterDataPipe[_T_co]): + r""" + Shuffle the input MapDataPipe via its indices (functional name: ``shuffle``). + + When it is used with :class:`~torch.utils.data.DataLoader`, the methods to + set up random seed are different based on :attr:`num_workers`. + + For single-process mode (:attr:`num_workers == 0`), the random seed is set before + the :class:`~torch.utils.data.DataLoader` in the main process. For multi-process + mode (:attr:`num_worker > 0`), ``worker_init_fn`` is used to set up a random seed + for each worker process. + + Args: + datapipe: MapDataPipe being shuffled + indices: a list of indices of the MapDataPipe. If not provided, we assume it uses 0-based indexing + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper + >>> dp = SequenceWrapper(range(10)) + >>> shuffle_dp = dp.shuffle().set_seed(0) + >>> list(shuffle_dp) + [7, 8, 1, 5, 3, 4, 2, 0, 9, 6] + >>> list(shuffle_dp) + [6, 1, 9, 5, 2, 4, 7, 3, 8, 0] + >>> # Reset seed for Shuffler + >>> shuffle_dp = shuffle_dp.set_seed(0) + >>> list(shuffle_dp) + [7, 8, 1, 5, 3, 4, 2, 0, 9, 6] + + Note: + Even thought this ``shuffle`` operation takes a ``MapDataPipe`` as the input, it would return an + ``IterDataPipe`` rather than a ``MapDataPipe``, because ``MapDataPipe`` should be non-sensitive to + the order of data order for the sake of random reads, but ``IterDataPipe`` depends on the order + of data during data-processing. + """ + + datapipe: MapDataPipe[_T_co] + _enabled: bool + _seed: int | None + _rng: random.Random + + def __init__( + self, + datapipe: MapDataPipe[_T_co], + *, + indices: list | None = None, + ) -> None: + super().__init__() + self.datapipe = datapipe + # pyrefly: ignore [bad-argument-type] + self.indices = list(range(len(datapipe))) if indices is None else indices + self._enabled = True + self._seed = None + self._rng = random.Random() + self._shuffled_indices: list = self.indices + + def set_shuffle(self, shuffle=True): + self._enabled = shuffle + return self + + def set_seed(self, seed: int): + self._seed = seed + return self + + def __iter__(self) -> Iterator[_T_co]: + if not self._enabled: + for idx in self.indices: + yield self.datapipe[idx] + else: + while self._shuffled_indices: + idx = self._shuffled_indices.pop() + yield self.datapipe[idx] + + def reset(self) -> None: + if self._enabled and self._seed is None: + self._seed = int(torch.empty((), dtype=torch.int64).random_().item()) + self._rng.seed(self._seed) + self._seed = None + self._shuffled_indices = self._rng.sample(self.indices, len(self.indices)) + + def __len__(self) -> int: + # pyrefly: ignore [bad-argument-type] + return len(self.datapipe) + + def __getstate__(self): + state = ( + self.datapipe, + self.indices, + self._enabled, + self._seed, + self._rng.getstate(), + self._shuffled_indices, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) + if IterDataPipe.getstate_hook is not None: + return IterDataPipe.getstate_hook(state) + return state + + def __setstate__(self, state): + ( + self.datapipe, + self.indices, + self._enabled, + self._seed, + rng_state, + self._shuffled_indices, + self._valid_iterator_id, + self._number_of_samples_yielded, + ) = state + self._rng = random.Random() + self._rng.setstate(rng_state) + + +MapDataPipe.register_datapipe_as_function("shuffle", ShufflerIterDataPipe) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/combining.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/combining.py new file mode 100644 index 0000000000000000000000000000000000000000..276404de530015f14561832e5b8a0d28383ca57f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/combining.py @@ -0,0 +1,111 @@ +# mypy: allow-untyped-defs +from collections.abc import Sized +from typing import TypeVar + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import MapDataPipe + + +__all__ = ["ConcaterMapDataPipe", "ZipperMapDataPipe"] + +_T_co = TypeVar("_T_co", covariant=True) + + +@functional_datapipe("concat") +class ConcaterMapDataPipe(MapDataPipe): + r""" + Concatenate multiple Map DataPipes (functional name: ``concat``). + + The new index of is the cumulative sum of source DataPipes. + For example, if there are 2 source DataPipes both with length 5, + index 0 to 4 of the resulting `ConcatMapDataPipe` would refer to + elements of the first DataPipe, and 5 to 9 would refer to elements + of the second DataPipe. + + Args: + datapipes: Map DataPipes being concatenated + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper + >>> dp1 = SequenceWrapper(range(3)) + >>> dp2 = SequenceWrapper(range(3)) + >>> concat_dp = dp1.concat(dp2) + >>> list(concat_dp) + [0, 1, 2, 0, 1, 2] + """ + + datapipes: tuple[MapDataPipe] + + def __init__(self, *datapipes: MapDataPipe) -> None: + if len(datapipes) == 0: + raise ValueError("Expected at least one DataPipe, but got nothing") + if not all(isinstance(dp, MapDataPipe) for dp in datapipes): + raise TypeError("Expected all inputs to be `MapDataPipe`") + # pyrefly: ignore [unsafe-overlap] + if not all(isinstance(dp, Sized) for dp in datapipes): + raise TypeError("Expected all inputs to be `Sized`") + self.datapipes = datapipes # type: ignore[assignment] + + def __getitem__(self, index) -> _T_co: # type: ignore[type-var] + offset = 0 + for dp in self.datapipes: + # pyrefly: ignore [bad-argument-type] + if index - offset < len(dp): + return dp[index - offset] + else: + # pyrefly: ignore [bad-argument-type] + offset += len(dp) + raise IndexError(f"Index {index} is out of range.") + + def __len__(self) -> int: + # pyrefly: ignore [bad-argument-type] + return sum(len(dp) for dp in self.datapipes) + + +@functional_datapipe("zip") +class ZipperMapDataPipe(MapDataPipe[tuple[_T_co, ...]]): + r""" + Aggregates elements into a tuple from each of the input DataPipes (functional name: ``zip``). + + This MataPipe is out of bound as soon as the shortest input DataPipe is exhausted. + + Args: + *datapipes: Map DataPipes being aggregated + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper + >>> dp1 = SequenceWrapper(range(3)) + >>> dp2 = SequenceWrapper(range(10, 13)) + >>> zip_dp = dp1.zip(dp2) + >>> list(zip_dp) + [(0, 10), (1, 11), (2, 12)] + """ + + datapipes: tuple[MapDataPipe[_T_co], ...] + + def __init__(self, *datapipes: MapDataPipe[_T_co]) -> None: + if len(datapipes) == 0: + raise ValueError("Expected at least one DataPipe, but got nothing") + if not all(isinstance(dp, MapDataPipe) for dp in datapipes): + raise TypeError("Expected all inputs to be `MapDataPipe`") + # pyrefly: ignore [unsafe-overlap] + if not all(isinstance(dp, Sized) for dp in datapipes): + raise TypeError("Expected all inputs to be `Sized`") + self.datapipes = datapipes + + def __getitem__(self, index) -> tuple[_T_co, ...]: + res = [] + for dp in self.datapipes: + try: + res.append(dp[index]) + except IndexError as e: + raise IndexError( + f"Index {index} is out of range for one of the input MapDataPipes {dp}." + ) from e + return tuple(res) + + def __len__(self) -> int: + # pyrefly: ignore [bad-argument-type] + return min(len(dp) for dp in self.datapipes) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/grouping.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/grouping.py new file mode 100644 index 0000000000000000000000000000000000000000..fdc34d9dc4fdb0d6ce4d973c667bbfc774bd9bcb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/grouping.py @@ -0,0 +1,76 @@ +# mypy: allow-untyped-defs +from collections.abc import Sized +from typing import TypeVar + +from torch.utils.data.datapipes._decorator import functional_datapipe +from torch.utils.data.datapipes.datapipe import DataChunk, MapDataPipe + + +__all__ = ["BatcherMapDataPipe"] + + +_T = TypeVar("_T") + + +@functional_datapipe("batch") +class BatcherMapDataPipe(MapDataPipe[DataChunk]): + r""" + Create mini-batches of data (functional name: ``batch``). + + An outer dimension will be added as ``batch_size`` if ``drop_last`` is set to ``True``, + or ``length % batch_size`` for the last batch if ``drop_last`` is set to ``False``. + + Args: + datapipe: Iterable DataPipe being batched + batch_size: The size of each batch + drop_last: Option to drop the last batch if it's not full + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper + >>> dp = SequenceWrapper(range(10)) + >>> batch_dp = dp.batch(batch_size=2) + >>> list(batch_dp) + [[0, 1], [2, 3], [4, 5], [6, 7], [8, 9]] + """ + + datapipe: MapDataPipe + batch_size: int + drop_last: bool + + def __init__( + self, + datapipe: MapDataPipe[_T], + batch_size: int, + drop_last: bool = False, + wrapper_class: type[DataChunk] = DataChunk, + ) -> None: + if batch_size <= 0: + raise AssertionError("Batch size is required to be larger than 0!") + super().__init__() + self.datapipe = datapipe + self.batch_size = batch_size + self.drop_last = drop_last + self.wrapper_class = wrapper_class + + def __getitem__(self, index) -> DataChunk: + batch: list = [] + indices = range(index * self.batch_size, (index + 1) * self.batch_size) + try: + batch.extend(self.datapipe[i] for i in indices) + return self.wrapper_class(batch) + except IndexError as e: + if not self.drop_last and len(batch) > 0: + return self.wrapper_class(batch) + else: + raise IndexError(f"Index {index} is out of bound.") from e + + def __len__(self) -> int: + # pyrefly: ignore [unsafe-overlap] + if isinstance(self.datapipe, Sized): + if self.drop_last: + return len(self.datapipe) // self.batch_size + else: + return (len(self.datapipe) + self.batch_size - 1) // self.batch_size + else: + raise TypeError(f"{type(self).__name__} instance doesn't have valid length") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a5b9075f1dbbc66a84dfd14d0778cc96ca604da0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/map/utils.py @@ -0,0 +1,61 @@ +import copy +import warnings +from collections.abc import Mapping, Sequence +from typing import Any, TypeVar + +from torch.utils.data.datapipes.datapipe import MapDataPipe + + +_T = TypeVar("_T") + +__all__ = ["SequenceWrapperMapDataPipe"] + + +class SequenceWrapperMapDataPipe(MapDataPipe[_T]): + r""" + Wraps a sequence object into a MapDataPipe. + + Args: + sequence: Sequence object to be wrapped into an MapDataPipe + deepcopy: Option to deepcopy input sequence object + + .. note:: + If ``deepcopy`` is set to False explicitly, users should ensure + that data pipeline doesn't contain any in-place operations over + the iterable instance, in order to prevent data inconsistency + across iterations. + + Example: + >>> # xdoctest: +SKIP + >>> from torchdata.datapipes.map import SequenceWrapper + >>> dp = SequenceWrapper(range(10)) + >>> list(dp) + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9] + >>> dp = SequenceWrapper({"a": 100, "b": 200, "c": 300, "d": 400}) + >>> dp["a"] + 100 + """ + + sequence: Sequence[_T] | Mapping[Any, _T] + + def __init__( + self, sequence: Sequence[_T] | Mapping[Any, _T], deepcopy: bool = True + ) -> None: + if deepcopy: + try: + self.sequence = copy.deepcopy(sequence) + except TypeError: + warnings.warn( + "The input sequence can not be deepcopied, " + "please be aware of in-place modification would affect source data", + stacklevel=2, + ) + self.sequence = sequence + else: + self.sequence = sequence + + def __getitem__(self, index: int) -> _T: + return self.sequence[index] + + def __len__(self) -> int: + return len(self.sequence) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/common.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/common.py new file mode 100644 index 0000000000000000000000000000000000000000..4e78cd5095245247d43e06d69f557e8467214de2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/common.py @@ -0,0 +1,414 @@ +# mypy: allow-untyped-defs +import fnmatch +import functools +import inspect +import os +import warnings +from collections.abc import Callable, Iterable +from io import IOBase +from typing import Any, NoReturn + +from torch.utils._import_utils import dill_available + + +__all__ = [ + "validate_input_col", + "StreamWrapper", + "get_file_binaries_from_pathnames", + "get_file_pathnames_from_root", + "match_masks", + "validate_pathname_binary_tuple", +] + + +# BC for torchdata +DILL_AVAILABLE = dill_available() + + +def validate_input_col(fn: Callable, input_col: int | tuple | list | None) -> None: + """ + Check that function used in a callable datapipe works with the input column. + + This simply ensures that the number of positional arguments matches the size + of the input column. The function must not contain any non-default + keyword-only arguments. + + Examples: + >>> # xdoctest: +SKIP("Failing on some CI machines") + >>> def f(a, b, *, c=1): + >>> return a + b + c + >>> def f_def(a, b=1, *, c=1): + >>> return a + b + c + >>> assert validate_input_col(f, [1, 2]) + >>> assert validate_input_col(f_def, 1) + >>> assert validate_input_col(f_def, [1, 2]) + + Notes: + If the function contains variable positional (`inspect.VAR_POSITIONAL`) arguments, + for example, f(a, *args), the validator will accept any size of input column + greater than or equal to the number of positional arguments. + (in this case, 1). + + Args: + fn: The function to check. + input_col: The input column to check. + + Raises: + ValueError: If the function is not compatible with the input column. + """ + try: + sig = inspect.signature(fn) + except ( + ValueError + ): # Signature cannot be inspected, likely it is a built-in fn or written in C + return + if isinstance(input_col, (list, tuple)): + input_col_size = len(input_col) + else: + input_col_size = 1 + + pos = [] + var_positional = False + non_default_kw_only = [] + + for p in sig.parameters.values(): + if p.kind in ( + inspect.Parameter.POSITIONAL_ONLY, + inspect.Parameter.POSITIONAL_OR_KEYWORD, + ): + pos.append(p) + elif p.kind is inspect.Parameter.VAR_POSITIONAL: + var_positional = True + elif p.kind is inspect.Parameter.KEYWORD_ONLY: + if p.default is p.empty: + non_default_kw_only.append(p) + else: + continue + + if isinstance(fn, functools.partial): + fn_name = getattr(fn.func, "__name__", repr(fn.func)) + else: + fn_name = getattr(fn, "__name__", repr(fn)) + + if len(non_default_kw_only) > 0: + raise ValueError( + f"The function {fn_name} takes {len(non_default_kw_only)} " + f"non-default keyword-only parameters, which is not allowed." + ) + + if len(sig.parameters) < input_col_size: + if not var_positional: + raise ValueError( + f"The function {fn_name} takes {len(sig.parameters)} " + f"parameters, but {input_col_size} are required." + ) + else: + if len(pos) > input_col_size: + if any(p.default is p.empty for p in pos[input_col_size:]): + raise ValueError( + f"The function {fn_name} takes {len(pos)} " + f"positional parameters, but {input_col_size} are required." + ) + elif len(pos) < input_col_size: + if not var_positional: + raise ValueError( + f"The function {fn_name} takes {len(pos)} " + f"positional parameters, but {input_col_size} are required." + ) + + +def _is_local_fn(fn): + # Functions or Methods + if hasattr(fn, "__code__"): + return fn.__code__.co_flags & inspect.CO_NESTED + # Callable Objects + else: + if hasattr(fn, "__qualname__"): + return "" in fn.__qualname__ + fn_type = type(fn) + if hasattr(fn_type, "__qualname__"): + return "" in fn_type.__qualname__ + return False + + +def _check_unpickable_fn(fn: Callable) -> None: + """ + Check function is pickable or not. + + If it is a lambda or local function, a UserWarning will be raised. If it's not a callable function, a TypeError will be raised. + """ + if not callable(fn): + raise TypeError(f"A callable function is expected, but {type(fn)} is provided.") + + # Extract function from partial object + # Nested partial function is automatically expanded as a single partial object + if isinstance(fn, functools.partial): + fn = fn.func + + # Local function + if _is_local_fn(fn) and not dill_available(): + warnings.warn( + "Local function is not supported by pickle, please use " + "regular python function or functools.partial instead.", + stacklevel=2, + ) + return + + # Lambda function + if hasattr(fn, "__name__") and fn.__name__ == "" and not dill_available(): + warnings.warn( + "Lambda function is not supported by pickle, please use " + "regular python function or functools.partial instead.", + stacklevel=2, + ) + return + + +def match_masks(name: str, masks: str | list[str]) -> bool: + # empty mask matches any input name + if not masks: + return True + + if isinstance(masks, str): + return fnmatch.fnmatch(name, masks) + + for mask in masks: + if fnmatch.fnmatch(name, mask): + return True + return False + + +def get_file_pathnames_from_root( + root: str, + masks: str | list[str], + recursive: bool = False, + abspath: bool = False, + non_deterministic: bool = False, +) -> Iterable[str]: + # print out an error message and raise the error out + def onerror(err: OSError) -> NoReturn: + warnings.warn(err.filename + " : " + err.strerror, stacklevel=2) + raise err + + if os.path.isfile(root): + path = root + if abspath: + path = os.path.abspath(path) + fname = os.path.basename(path) + if match_masks(fname, masks): + yield path + else: + for path, dirs, files in os.walk(root, onerror=onerror): + if abspath: + path = os.path.abspath(path) + if not non_deterministic: + files.sort() + for f in files: + if match_masks(f, masks): + yield os.path.join(path, f) + if not recursive: + break + if not non_deterministic: + # Note that this is in-place modifying the internal list from `os.walk` + # This only works because `os.walk` doesn't shallow copy before turn + # https://github.com/python/cpython/blob/f4c03484da59049eb62a9bf7777b963e2267d187/Lib/os.py#L407 + dirs.sort() + + +def get_file_binaries_from_pathnames( + pathnames: Iterable, mode: str, encoding: str | None = None +): + if not isinstance(pathnames, Iterable): + pathnames = [ + pathnames, + ] + + if mode in ("b", "t"): + mode = "r" + mode + + for pathname in pathnames: + if not isinstance(pathname, str): + raise TypeError( + f"Expected string type for pathname, but got {type(pathname)}" + ) + yield pathname, StreamWrapper(open(pathname, mode, encoding=encoding)) # noqa:SIM115 + + +def validate_pathname_binary_tuple(data: tuple[str, IOBase]) -> None: + if not isinstance(data, tuple): + raise TypeError( + f"pathname binary data should be tuple type, but it is type {type(data)}" + ) + if len(data) != 2: + raise TypeError( + f"pathname binary stream tuple length should be 2, but got {len(data)}" + ) + if not isinstance(data[0], str): + raise TypeError( + f"pathname within the tuple should have string type pathname, but it is type {type(data[0])}" + ) + if not isinstance(data[1], IOBase) and not isinstance(data[1], StreamWrapper): + raise TypeError( + f"binary stream within the tuple should have IOBase or" + f"its subclasses as type, but it is type {type(data[1])}" + ) + + +# Deprecated function names and its corresponding DataPipe type and kwargs for the `_deprecation_warning` function +_iter_deprecated_functional_names: dict[str, dict] = {} +_map_deprecated_functional_names: dict[str, dict] = {} + + +def _deprecation_warning( + old_class_name: str, + *, + deprecation_version: str, + removal_version: str, + old_functional_name: str = "", + old_argument_name: str = "", + new_class_name: str = "", + new_functional_name: str = "", + new_argument_name: str = "", + deprecate_functional_name_only: bool = False, +) -> None: + if new_functional_name and not old_functional_name: + raise ValueError( + "Old functional API needs to be specified for the deprecation warning." + ) + if new_argument_name and not old_argument_name: + raise ValueError( + "Old argument name needs to be specified for the deprecation warning." + ) + + if old_functional_name and old_argument_name: + raise ValueError( + "Deprecating warning for functional API and argument should be separated." + ) + + msg = f"`{old_class_name}()`" + if deprecate_functional_name_only and old_functional_name: + msg = f"{msg}'s functional API `.{old_functional_name}()` is" + elif old_functional_name: + msg = f"{msg} and its functional API `.{old_functional_name}()` are" + elif old_argument_name: + msg = f"The argument `{old_argument_name}` of {msg} is" + else: + msg = f"{msg} is" + msg = ( + f"{msg} deprecated since {deprecation_version} and will be removed in {removal_version}." + f"\nSee https://github.com/pytorch/data/issues/163 for details." + ) + + if new_class_name or new_functional_name: + msg = f"{msg}\nPlease use" + if new_class_name: + msg = f"{msg} `{new_class_name}()`" + if new_class_name and new_functional_name: + msg = f"{msg} or" + if new_functional_name: + msg = f"{msg} `.{new_functional_name}()`" + msg = f"{msg} instead." + + if new_argument_name: + msg = f"{msg}\nPlease use `{old_class_name}({new_argument_name}=)` instead." + + warnings.warn(msg, FutureWarning, stacklevel=2) + + +class StreamWrapper: + """ + StreamWrapper is introduced to wrap file handler generated by DataPipe operation like `FileOpener`. + + StreamWrapper would guarantee the wrapped file handler is closed when it's out of scope. + """ + + session_streams: dict[Any, int] = {} + debug_unclosed_streams: bool = False + + def __init__(self, file_obj, parent_stream=None, name=None) -> None: + self.file_obj = file_obj + self.child_counter = 0 + self.parent_stream = parent_stream + self.close_on_last_child = False + self.name = name + self.closed = False + if parent_stream is not None: + if not isinstance(parent_stream, StreamWrapper): + raise RuntimeError( + f"Parent stream should be StreamWrapper, {type(parent_stream)} was given" + ) + parent_stream.child_counter += 1 + self.parent_stream = parent_stream + if StreamWrapper.debug_unclosed_streams: + StreamWrapper.session_streams[self] = 1 + + @classmethod + def close_streams(cls, v, depth=0) -> None: + """Traverse structure and attempts to close all found StreamWrappers on best effort basis.""" + if depth > 10: + return + if isinstance(v, StreamWrapper): + v.close() + else: + # Traverse only simple structures + if isinstance(v, dict): + for vv in v.values(): + cls.close_streams(vv, depth=depth + 1) + elif isinstance(v, (list, tuple)): + for vv in v: + cls.close_streams(vv, depth=depth + 1) + + def __getattr__(self, name): + file_obj = self.__dict__["file_obj"] + return getattr(file_obj, name) + + def close(self, *args, **kwargs) -> None: + if self.closed: + return + if StreamWrapper.debug_unclosed_streams: + del StreamWrapper.session_streams[self] + if hasattr(self, "parent_stream") and self.parent_stream is not None: + self.parent_stream.child_counter -= 1 + if ( + not self.parent_stream.child_counter + and self.parent_stream.close_on_last_child + ): + self.parent_stream.close() + try: + self.file_obj.close(*args, **kwargs) + except AttributeError: + pass + self.closed = True + + def autoclose(self) -> None: + """Automatically close stream when all child streams are closed or if there are none.""" + self.close_on_last_child = True + if self.child_counter == 0: + self.close() + + def __dir__(self): + attrs = list(self.__dict__.keys()) + list(StreamWrapper.__dict__.keys()) + attrs += dir(self.file_obj) + return list(set(attrs)) + + def __del__(self) -> None: + if not self.closed: + self.close() + + def __iter__(self): + yield from self.file_obj + + def __next__(self): + return next(self.file_obj) + + def __repr__(self) -> str: + if self.name is None: + return f"StreamWrapper<{self.file_obj!r}>" + else: + return f"StreamWrapper<{self.name},{self.file_obj!r}>" + + def __getstate__(self): + return self.file_obj + + def __setstate__(self, obj): + self.file_obj = obj diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/decoder.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..2f33c5cc7e0e7312d0746292b8f714c10f9907a2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/decoder.py @@ -0,0 +1,389 @@ +# mypy: allow-untyped-defs +# This file takes partial of the implementation from NVIDIA's webdataset at here: +# https://github.com/tmbdev/webdataset/blob/master/webdataset/autodecode.py + +import io +import json +import os.path +import pickle +import tempfile + +import torch +from torch.utils.data.datapipes.utils.common import StreamWrapper + + +__all__ = [ + "Decoder", + "ImageHandler", + "MatHandler", + "audiohandler", + "basichandlers", + "extension_extract_fn", + "handle_extension", + "imagehandler", + "mathandler", + "videohandler", +] + + +################################################################ +# handle basic datatypes +################################################################ +def basichandlers(extension: str, data): + """Transforms raw data (byte stream) into python objects. + + Looks at the extension and loads the data into a python object supporting + the corresponding extension. + + Args: + extension (str): The file extension + data (byte stream): Data to load into a python object. + + Returns: + object: The data loaded into a corresponding python object + supporting the extension. + + Example: + >>> import pickle + >>> data = pickle.dumps("some data") + >>> new_data = basichandlers("pickle", data) + >>> new_data + some data + + The transformation of data for extensions are: + - txt, text, transcript: utf-8 decoded data of str format + - cls, cls2, class, count, index, inx, id: int + - json, jsn: json loaded data + - pickle, pyd: pickle loaded data + - pt: torch loaded data + """ + + if extension in "txt text transcript": + return data.decode("utf-8") + + if extension in ["cls", "cls2", "class", "count", "index", "inx", "id"]: + try: + return int(data) + except ValueError: + return None + + if extension in "json jsn": + return json.loads(data) + + if extension in ["pyd", "pickle"]: + return pickle.loads(data) + + if extension == "pt": + stream = io.BytesIO(data) + return torch.load(stream) + + # if extension in "ten tb".split(): + # from . import tenbin + # return tenbin.decode_buffer(data) + + # if extension in "mp msgpack msg".split(): + # import msgpack + # return msgpack.unpackb(data) + + return None + + +################################################################ +# handle images +################################################################ +imagespecs = { + "l8": ("numpy", "uint8", "l"), + "rgb8": ("numpy", "uint8", "rgb"), + "rgba8": ("numpy", "uint8", "rgba"), + "l": ("numpy", "float", "l"), + "rgb": ("numpy", "float", "rgb"), + "rgba": ("numpy", "float", "rgba"), + "torchl8": ("torch", "uint8", "l"), + "torchrgb8": ("torch", "uint8", "rgb"), + "torchrgba8": ("torch", "uint8", "rgba"), + "torchl": ("torch", "float", "l"), + "torchrgb": ("torch", "float", "rgb"), + "torch": ("torch", "float", "rgb"), + "torchrgba": ("torch", "float", "rgba"), + "pill": ("pil", None, "l"), + "pil": ("pil", None, "rgb"), + "pilrgb": ("pil", None, "rgb"), + "pilrgba": ("pil", None, "rgba"), +} + + +def handle_extension(extensions, f): + """ + Return a decoder handler function for the list of extensions. + + Extensions can be a space separated list of extensions. + Extensions can contain dots, in which case the corresponding number + of extension components must be present in the key given to f. + Comparisons are case insensitive. + Examples: + handle_extension("jpg jpeg", my_decode_jpg) # invoked for any file.jpg + handle_extension("seg.jpg", special_case_jpg) # invoked only for file.seg.jpg + """ + extensions = extensions.lower().split() + + def g(key, data): + extension = key.lower().split(".") + + for target in extensions: + target = target.split(".") + if len(target) > len(extension): + continue + + if extension[-len(target) :] == target: + return f(data) + return None + + return g + + +class ImageHandler: + """ + Decode image data using the given `imagespec`. + + The `imagespec` specifies whether the image is decoded + to numpy/torch/pi, decoded to uint8/float, and decoded + to l/rgb/rgba: + + - l8: numpy uint8 l + - rgb8: numpy uint8 rgb + - rgba8: numpy uint8 rgba + - l: numpy float l + - rgb: numpy float rgb + - rgba: numpy float rgba + - torchl8: torch uint8 l + - torchrgb8: torch uint8 rgb + - torchrgba8: torch uint8 rgba + - torchl: torch float l + - torchrgb: torch float rgb + - torch: torch float rgb + - torchrgba: torch float rgba + - pill: pil None l + - pil: pil None rgb + - pilrgb: pil None rgb + - pilrgba: pil None rgba + """ + + def __init__(self, imagespec) -> None: + if imagespec not in list(imagespecs.keys()): + raise AssertionError(f"unknown image specification: {imagespec}") + self.imagespec = imagespec.lower() + + def __call__(self, extension, data): + if extension.lower() not in ["jpg", "jpeg", "png", "ppm", "pgm", "pbm", "pnm"]: + return None + + try: + import numpy as np + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "Package `numpy` is required to be installed for default image decoder." + "Please use `pip install numpy` to install the package" + ) from e + + try: + import PIL.Image + except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "Package `PIL` is required to be installed for default image decoder." + "Please use `pip install Pillow` to install the package" + ) from e + + imagespec = self.imagespec + atype, etype, mode = imagespecs[imagespec] + + with io.BytesIO(data) as stream: + img = PIL.Image.open(stream) + img.load() + img = img.convert(mode.upper()) + if atype == "pil": + return img + elif atype == "numpy": + result = np.asarray(img) + if result.dtype != np.uint8: + raise AssertionError( + f"numpy image array should be type uint8, but got {result.dtype}" + ) + if etype == "uint8": + return result + else: + return result.astype("f") / 255.0 + elif atype == "torch": + result = np.asarray(img) + if result.dtype != np.uint8: + raise AssertionError( + f"numpy image array should be type uint8, but got {result.dtype}" + ) + + if etype == "uint8": + result = np.array(result.transpose(2, 0, 1)) + return torch.tensor(result) + else: + result = np.array(result.transpose(2, 0, 1)) + return torch.tensor(result) / 255.0 + return None + + +def imagehandler(imagespec): + return ImageHandler(imagespec) + + +################################################################ +# torch video +################################################################ +def videohandler(extension, data): + if extension not in [ + "mp4", + "ogv", + "mjpeg", + "avi", + "mov", + "h264", + "mpg", + "webm", + "wmv", + ]: + return None + + try: + import torchvision.io + except ImportError as e: + raise ModuleNotFoundError( + "Package `torchvision` is required to be installed for default video file loader." + "Please use `pip install torchvision`" + "to install the package" + ) from e + + with tempfile.TemporaryDirectory() as dirname: + fname = os.path.join(dirname, f"file.{extension}") + with open(fname, "wb") as stream: + stream.write(data) + return torchvision.io.read_video(fname) + + +################################################################ +# torchaudio +################################################################ +def audiohandler(extension, data): + if extension not in ["flac", "mp3", "sox", "wav", "m4a", "ogg", "wma"]: + return None + + try: + import torchaudio # type: ignore[import] + except ImportError as e: + raise ModuleNotFoundError( + "Package `torchaudio` is required to be installed for default audio file loader." + "Please use `pip install torchaudio`" + "to install the package" + ) from e + + with tempfile.TemporaryDirectory() as dirname: + fname = os.path.join(dirname, f"file.{extension}") + with open(fname, "wb") as stream: + stream.write(data) + return torchaudio.load(fname) + + +################################################################ +# mat +################################################################ +class MatHandler: + def __init__(self, **loadmat_kwargs) -> None: + try: + import scipy.io as sio + except ImportError as e: + raise ModuleNotFoundError( + "Package `scipy` is required to be installed for mat file." + "Please use `pip install scipy`" + "to install the package" + ) from e + self.sio = sio + self.loadmat_kwargs = loadmat_kwargs + + def __call__(self, extension, data): + if extension != "mat": + return None + with io.BytesIO(data) as stream: + return self.sio.loadmat(stream, **self.loadmat_kwargs) + + +def mathandler(**loadmat_kwargs): + return MatHandler(**loadmat_kwargs) + + +################################################################ +# a sample decoder +################################################################ +# Extract extension from pathname +def extension_extract_fn(pathname): + ext = os.path.splitext(pathname)[1] + # Remove dot + if ext: + ext = ext[1:] + return ext + + +class Decoder: + """ + Decode key/data sets using a list of handlers. + + For each key/data item, this iterates through the list of + handlers until some handler returns something other than None. + """ + + def __init__(self, *handler, key_fn=extension_extract_fn) -> None: + self.handlers = list(handler) if handler else [] + self.key_fn = key_fn + + # Insert new handler from the beginning of handlers list to make sure the new + # handler having the highest priority + def add_handler(self, *handler) -> None: + if not handler: + return + self.handlers = list(handler) + self.handlers + + @staticmethod + def _is_stream_handle(data): + obj_to_check = data.file_obj if isinstance(data, StreamWrapper) else data + return isinstance(obj_to_check, (io.BufferedIOBase, io.RawIOBase)) + + def decode1(self, key, data): + if not data: + return data + + # if data is a stream handle, we need to read all the content before decoding + if Decoder._is_stream_handle(data): + ds = data + # The behavior of .read can differ between streams (e.g. HTTPResponse), hence this is used instead + data = b"".join(data) + ds.close() + + for f in self.handlers: + result = f(key, data) + if result is not None: + return result + return data + + def decode(self, data): + result = {} + # single data tuple(pathname, data stream) + if isinstance(data, tuple): + data = [data] + + if data is not None: + for k, v in data: + # TODO: xinyu, figure out why Nvidia do this? + if k[0] == "_": + if isinstance(v, bytes): + v = v.decode("utf-8") + result[k] = v + continue + result[k] = self.decode1(self.key_fn(k), v) + return result + + def __call__(self, data): + return self.decode(data) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/snapshot.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/snapshot.py new file mode 100644 index 0000000000000000000000000000000000000000..42aec1aa308a9b21b251de595cddfbe171930bb6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/datapipes/utils/snapshot.py @@ -0,0 +1,65 @@ +# mypy: allow-untyped-defs +from torch.utils.data.datapipes._hook_iterator import _SnapshotState +from torch.utils.data.datapipes.datapipe import IterDataPipe +from torch.utils.data.graph_settings import apply_random_seed + + +# TODO: Caveats +# 1. Caller (either the ReadingService or DataLoader) must pass in the initial RNG +# 2. `in_batch_shuffle` and `bucketbatch` are not compatible with this because they currently +# lack the option to `set_seed`. +def _simple_graph_snapshot_restoration( + datapipe: IterDataPipe, n_iterations: int, rng=None +) -> None: + r""" + Fast-forward the given DataPipe and its parents by ``n_iterations``, re-doing computations to restore a snapshot. + + For instance, applying this function to the final DataPipe of a graph will restore the snapshot + (via fast-forward) every DataPipe within the graph. + + After you deserialize a DataPipe, you can use its `_number_of_samples_yielded` attribute as the input + to this function to forward the DataPipe. + + A DataPipe cannot be restored twice in a row unless there is an iteration started between the restoration + attempts. + + Note: + This is the simplest but least efficient way to fast-forward a DataPipe. Usage of other fast-forwarding + methods (custom ones if necessary) are recommended. + + Args: + datapipe: IterDataPipe to be fast-forwarded + n_iterations: number of iterations to fast-forward + rng: ``Optional[torch.Generator]``. If not ``None``, this RNG will be used for shuffling. The generator + should be in its `initial` state as it was first passed into ``DataLoader`` or ``ReadingService``. + """ + if datapipe._snapshot_state == _SnapshotState.Restored: + raise RuntimeError( + "Snapshot restoration cannot be applied. You can only restore simple snapshot to the graph " + "if your graph has not been restored." + ) + + # For this snapshot restoration function, we want the DataPipe to be at its initial state prior to + # simple fast-forwarding. Therefore, we need to call `reset` twice, because if `SnapshotState` is `Restored`, + # the first reset will not actually reset. + datapipe.reset() # This ensures `SnapshotState` is `Iterating` by this point, even if it was `Restored`. + # pyrefly: ignore [bad-argument-type] + apply_random_seed(datapipe, rng) + + remainder = n_iterations + it = iter(datapipe) # This always reset the DataPipe if it hasn't already. + while remainder > 0: + try: + next(it) + remainder -= 1 + except StopIteration as e: + raise RuntimeError( + f"Fast-forward {datapipe} by {n_iterations} iterations " + "exceeds the number of samples available." + ) from e + datapipe._fast_forward_iterator = it + # While the DataPipe has `_fast_forward_iterator`, `next()` will get result from there instead of elsewhere. + + # This will prevent the DataPipe from resetting in the `iter()` call + # If another DataPipe is consuming it, it won't have to start over again + datapipe._snapshot_state = _SnapshotState.Restored diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/dataset.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..a7c8be6824b1ebeb6b2aaf82a4720870572d2ca9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/dataset.py @@ -0,0 +1,511 @@ +# mypy: allow-untyped-defs +import bisect +import itertools +import math +import warnings +from collections.abc import Sequence + +# UP006 wants 'Iterable' to be imported from collections.abc but it needs to +# stay from typing for now due to BC concerns. In particular several internal +# targets fail to typecheck with: +# TypeError: Cannot create a consistent method resolution order (MRO) for +# bases Iterable, Generic +from typing import cast, Generic, Iterable, TypeVar # noqa: UP035 +from typing_extensions import deprecated + +# No 'default_generator' in torch/__init__.pyi +from torch import default_generator, Generator, randperm, Tensor + + +__all__ = [ + "Dataset", + "IterableDataset", + "TensorDataset", + "StackDataset", + "ConcatDataset", + "ChainDataset", + "Subset", + "random_split", +] + + +_T = TypeVar("_T") +_T_co = TypeVar("_T_co", covariant=True) +_T_dict = dict[str, _T_co] +_T_tuple = tuple[_T_co, ...] +_T_stack = TypeVar("_T_stack", _T_tuple, _T_dict) + + +class Dataset(Generic[_T_co]): + r"""An abstract class representing a :class:`Dataset`. + + All datasets that represent a map from keys to data samples should subclass + it. All subclasses should overwrite :meth:`__getitem__`, supporting fetching a + data sample for a given key. Subclasses could also optionally overwrite + :meth:`__len__`, which is expected to return the size of the dataset by many + :class:`~torch.utils.data.Sampler` implementations and the default options + of :class:`~torch.utils.data.DataLoader`. Subclasses could also + optionally implement :meth:`__getitems__`, for speedup batched samples + loading. This method accepts list of indices of samples of batch and returns + list of samples. + + .. note:: + :class:`~torch.utils.data.DataLoader` by default constructs an index + sampler that yields integral indices. To make it work with a map-style + dataset with non-integral indices/keys, a custom sampler must be provided. + """ + + def __getitem__(self, index) -> _T_co: + raise NotImplementedError("Subclasses of Dataset should implement __getitem__.") + + # def __getitems__(self, indices: List) -> List[_T_co]: + # Not implemented to prevent false-positives in fetcher check in + # torch.utils.data._utils.fetch._MapDatasetFetcher + + def __add__(self, other: "Dataset[_T_co]") -> "ConcatDataset[_T_co]": + return ConcatDataset([self, other]) + + # No `def __len__(self)` default? + # See NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] + # in pytorch/torch/utils/data/sampler.py + + +class IterableDataset(Dataset[_T_co], Iterable[_T_co]): + r"""An iterable Dataset. + + All datasets that represent an iterable of data samples should subclass it. + Such form of datasets is particularly useful when data come from a stream. + + All subclasses should overwrite :meth:`__iter__`, which would return an + iterator of samples in this dataset. + + When a subclass is used with :class:`~torch.utils.data.DataLoader`, each + item in the dataset will be yielded from the :class:`~torch.utils.data.DataLoader` + iterator. When :attr:`num_workers > 0`, each worker process will have a + different copy of the dataset object, so it is often desired to configure + each copy independently to avoid having duplicate data returned from the + workers. :func:`~torch.utils.data.get_worker_info`, when called in a worker + process, returns information about the worker. It can be used in either the + dataset's :meth:`__iter__` method or the :class:`~torch.utils.data.DataLoader` 's + :attr:`worker_init_fn` option to modify each copy's behavior. + + Example 1: splitting workload across all workers in :meth:`__iter__`:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_DATALOADER) + >>> # xdoctest: +SKIP("Fails on MacOS12") + >>> class MyIterableDataset(torch.utils.data.IterableDataset): + ... def __init__(self, start, end): + ... super(MyIterableDataset).__init__() + ... assert end > start, "this example only works with end >= start" + ... self.start = start + ... self.end = end + ... + ... def __iter__(self): + ... worker_info = torch.utils.data.get_worker_info() + ... if worker_info is None: # single-process data loading, return the full iterator + ... iter_start = self.start + ... iter_end = self.end + ... else: # in a worker process + ... # split workload + ... per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers))) + ... worker_id = worker_info.id + ... iter_start = self.start + worker_id * per_worker + ... iter_end = min(iter_start + per_worker, self.end) + ... return iter(range(iter_start, iter_end)) + ... + >>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6]. + >>> ds = MyIterableDataset(start=3, end=7) + + >>> # Single-process loading + >>> print(list(torch.utils.data.DataLoader(ds, num_workers=0))) + [tensor([3]), tensor([4]), tensor([5]), tensor([6])] + + >>> # xdoctest: +REQUIRES(POSIX) + >>> # Multi-process loading with two worker processes + >>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6]. + >>> # xdoctest: +IGNORE_WANT("non deterministic") + >>> print(list(torch.utils.data.DataLoader(ds, num_workers=2))) + [tensor([3]), tensor([5]), tensor([4]), tensor([6])] + + >>> # With even more workers + >>> # xdoctest: +IGNORE_WANT("non deterministic") + >>> print(list(torch.utils.data.DataLoader(ds, num_workers=12))) + [tensor([3]), tensor([5]), tensor([4]), tensor([6])] + + Example 2: splitting workload across all workers using :attr:`worker_init_fn`:: + + >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_DATALOADER) + >>> class MyIterableDataset(torch.utils.data.IterableDataset): + ... def __init__(self, start, end): + ... super(MyIterableDataset).__init__() + ... assert end > start, "this example only works with end >= start" + ... self.start = start + ... self.end = end + ... + ... def __iter__(self): + ... return iter(range(self.start, self.end)) + ... + >>> # should give same set of data as range(3, 7), i.e., [3, 4, 5, 6]. + >>> ds = MyIterableDataset(start=3, end=7) + + >>> # Single-process loading + >>> print(list(torch.utils.data.DataLoader(ds, num_workers=0))) + [3, 4, 5, 6] + >>> + >>> # Directly doing multi-process loading yields duplicate data + >>> print(list(torch.utils.data.DataLoader(ds, num_workers=2))) + [3, 3, 4, 4, 5, 5, 6, 6] + + >>> # Define a `worker_init_fn` that configures each dataset copy differently + >>> def worker_init_fn(worker_id): + ... worker_info = torch.utils.data.get_worker_info() + ... dataset = worker_info.dataset # the dataset copy in this worker process + ... overall_start = dataset.start + ... overall_end = dataset.end + ... # configure the dataset to only process the split workload + ... per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers))) + ... worker_id = worker_info.id + ... dataset.start = overall_start + worker_id * per_worker + ... dataset.end = min(dataset.start + per_worker, overall_end) + ... + + >>> # Mult-process loading with the custom `worker_init_fn` + >>> # Worker 0 fetched [3, 4]. Worker 1 fetched [5, 6]. + >>> print(list(torch.utils.data.DataLoader(ds, num_workers=2, worker_init_fn=worker_init_fn))) + [3, 5, 4, 6] + + >>> # With even more workers + >>> print(list(torch.utils.data.DataLoader(ds, num_workers=12, worker_init_fn=worker_init_fn))) + [3, 4, 5, 6] + """ + + def __add__(self, other: Dataset[_T_co]): + return ChainDataset([self, other]) + + # No `def __len__(self)` default? Subclasses raise `TypeError` when needed. + # See NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] + + +class TensorDataset(Dataset[tuple[Tensor, ...]]): + r"""Dataset wrapping tensors. + + Each sample will be retrieved by indexing tensors along the first dimension. + + Args: + *tensors (Tensor): tensors that have the same size of the first dimension. + """ + + tensors: tuple[Tensor, ...] + + def __init__(self, *tensors: Tensor) -> None: + if any(tensors[0].size(0) != tensor.size(0) for tensor in tensors): + raise AssertionError("Size mismatch between tensors") + self.tensors = tensors + + def __getitem__(self, index): + return tuple(tensor[index] for tensor in self.tensors) + + def __len__(self) -> int: + return self.tensors[0].size(0) + + +class StackDataset(Dataset[_T_stack]): + r"""Dataset as a stacking of multiple datasets. + + This class is useful to assemble different parts of complex input data, given as datasets. + + Example: + >>> # xdoctest: +SKIP + >>> images = ImageDataset() + >>> texts = TextDataset() + >>> tuple_stack = StackDataset(images, texts) + >>> tuple_stack[0] == (images[0], texts[0]) + >>> dict_stack = StackDataset(image=images, text=texts) + >>> dict_stack[0] == {"image": images[0], "text": texts[0]} + + Args: + *args (Dataset): Datasets for stacking returned as tuple. + **kwargs (Dataset): Datasets for stacking returned as dict. + """ + + datasets: tuple | dict + + def __init__(self, *args: Dataset[_T_co], **kwargs: Dataset[_T_co]) -> None: + if args: + if kwargs: + raise ValueError( + "Supported either ``tuple``- (via ``args``) or" + "``dict``- (via ``kwargs``) like input/output, but both types are given." + ) + self._length = len(args[0]) # type: ignore[arg-type] + if any(self._length != len(dataset) for dataset in args): # type: ignore[arg-type] + raise ValueError("Size mismatch between datasets") + self.datasets = args + elif kwargs: + tmp = list(kwargs.values()) + self._length = len(tmp[0]) # type: ignore[arg-type] + if any(self._length != len(dataset) for dataset in tmp): # type: ignore[arg-type] + raise ValueError("Size mismatch between datasets") + self.datasets = kwargs + else: + raise ValueError("At least one dataset should be passed") + + def __getitem__(self, index): + if isinstance(self.datasets, dict): + return {k: dataset[index] for k, dataset in self.datasets.items()} + return tuple(dataset[index] for dataset in self.datasets) + + def __getitems__(self, indices: list): + # add batched sampling support when parent datasets supports it. + if isinstance(self.datasets, dict): + dict_batch: list[_T_dict] = [{} for _ in indices] + for k, dataset in self.datasets.items(): + if callable(getattr(dataset, "__getitems__", None)): + items = dataset.__getitems__(indices) # type: ignore[attr-defined] + if len(items) != len(indices): + raise ValueError( + "Nested dataset's output size mismatch." + f" Expected {len(indices)}, got {len(items)}" + ) + for data, d_sample in zip(items, dict_batch, strict=True): + d_sample[k] = data + else: + for idx, d_sample in zip(indices, dict_batch, strict=True): + d_sample[k] = dataset[idx] + return dict_batch + + # tuple data + list_batch: list[list] = [[] for _ in indices] + for dataset in self.datasets: + if callable(getattr(dataset, "__getitems__", None)): + items = dataset.__getitems__(indices) # type: ignore[attr-defined] + if len(items) != len(indices): + raise ValueError( + "Nested dataset's output size mismatch." + f" Expected {len(indices)}, got {len(items)}" + ) + for data, t_sample in zip(items, list_batch, strict=True): + t_sample.append(data) + else: + for idx, t_sample in zip(indices, list_batch, strict=True): + t_sample.append(dataset[idx]) + tuple_batch: list[_T_tuple] = [tuple(sample) for sample in list_batch] + return tuple_batch + + def __len__(self) -> int: + return self._length + + +class ConcatDataset(Dataset[_T_co]): + r"""Dataset as a concatenation of multiple datasets. + + This class is useful to assemble different existing datasets. + + Args: + datasets (sequence): List of datasets to be concatenated + """ + + datasets: list[Dataset[_T_co]] + cumulative_sizes: list[int] + + @staticmethod + def cumsum(sequence): + r, s = [], 0 + for e in sequence: + l = len(e) + r.append(l + s) + s += l + return r + + def __init__(self, datasets: Iterable[Dataset]) -> None: + super().__init__() + self.datasets = list(datasets) + if len(self.datasets) == 0: + raise AssertionError("datasets should not be an empty iterable") + for d in self.datasets: + if isinstance(d, IterableDataset): + raise AssertionError("ConcatDataset does not support IterableDataset") + self.cumulative_sizes = self.cumsum(self.datasets) + + def __len__(self) -> int: + return self.cumulative_sizes[-1] + + def __getitem__(self, idx): + if idx < 0: + if -idx > len(self): + raise ValueError( + "absolute value of index should not exceed dataset length" + ) + idx = len(self) + idx + dataset_idx = bisect.bisect_right(self.cumulative_sizes, idx) + if dataset_idx == 0: + sample_idx = idx + else: + sample_idx = idx - self.cumulative_sizes[dataset_idx - 1] + return self.datasets[dataset_idx][sample_idx] + + @property + @deprecated( + "`cummulative_sizes` attribute is renamed to `cumulative_sizes`", + category=FutureWarning, + ) + def cummulative_sizes(self): + return self.cumulative_sizes + + +class ChainDataset(IterableDataset): + r"""Dataset for chaining multiple :class:`IterableDataset` s. + + This class is useful to assemble different existing dataset streams. The + chaining operation is done on-the-fly, so concatenating large-scale + datasets with this class will be efficient. + + Args: + datasets (iterable of IterableDataset): datasets to be chained together + """ + + def __init__(self, datasets: Iterable[Dataset]) -> None: + super().__init__() + self.datasets = datasets + + def __iter__(self): + for d in self.datasets: + if not isinstance(d, IterableDataset): + raise AssertionError("ChainDataset only supports IterableDataset") + yield from d + + def __len__(self) -> int: + total = 0 + for d in self.datasets: + if not isinstance(d, IterableDataset): + raise AssertionError("ChainDataset only supports IterableDataset") + total += len(d) # type: ignore[arg-type] + return total + + +class Subset(Dataset[_T_co]): + r""" + Subset of a dataset at specified indices. + + .. note:: + When subclassing `Subset` and overriding `__getitem__`, you **must** also + override `__getitems__` to ensure `DataLoader` works correctly with your + custom logic. If you override only `__getitem__`, a `NotImplementedError` + will be raised when using `DataLoader`. + + A simple implementation of `__getitems__` can delegate to `__getitem__`: + + .. code-block:: python + + def __getitems__(self, indices): + return [self.__getitem__(idx) for idx in indices] + + For better performance, consider implementing batch-aware logic in + `__getitems__` instead of calling `__getitem__` multiple times. + + Args: + dataset (Dataset): The whole Dataset + indices (sequence): Indices in the whole set selected for subset + """ + + dataset: Dataset[_T_co] + indices: Sequence[int] + + def __init__(self, dataset: Dataset[_T_co], indices: Sequence[int]) -> None: + self.dataset = dataset + self.indices = indices + + # Check if __getitem__ is overridden but __getitems__ is not + if ( + type(self).__getitem__ is not Subset.__getitem__ + and type(self).__getitems__ is Subset.__getitems__ + ): + raise NotImplementedError( + f"{type(self).__name__} overrides __getitem__ but not __getitems__. " + "When subclassing Subset and overriding __getitem__, you must also override " + "__getitems__ to ensure DataLoader works correctly with your custom logic. " + "A simple implementation:\n\n" + "def __getitems__(self, indices):\n" + " return [self.__getitem__(idx) for idx in indices]" + ) + + def __getitem__(self, idx): + if isinstance(idx, list): + return self.dataset[[self.indices[i] for i in idx]] + return self.dataset[self.indices[idx]] + + def __getitems__(self, indices: list[int]) -> list[_T_co]: + # add batched sampling support when parent dataset supports it. + # see torch.utils.data._utils.fetch._MapDatasetFetcher + if callable(getattr(self.dataset, "__getitems__", None)): + return self.dataset.__getitems__([self.indices[idx] for idx in indices]) # type: ignore[attr-defined] + else: + return [self.dataset[self.indices[idx]] for idx in indices] + + def __len__(self) -> int: + return len(self.indices) + + +def random_split( + dataset: Dataset[_T], + lengths: Sequence[int | float], + generator: Generator | None = default_generator, +) -> list[Subset[_T]]: + r""" + Randomly split a dataset into non-overlapping new datasets of given lengths. + + If a list of fractions that sum up to 1 is given, + the lengths will be computed automatically as + floor(frac * len(dataset)) for each fraction provided. + + After computing the lengths, if there are any remainders, 1 count will be + distributed in round-robin fashion to the lengths + until there are no remainders left. + + Optionally fix the generator for reproducible results, e.g.: + + Example: + >>> # xdoctest: +SKIP + >>> generator1 = torch.Generator().manual_seed(42) + >>> generator2 = torch.Generator().manual_seed(42) + >>> random_split(range(10), [3, 7], generator=generator1) + >>> random_split(range(30), [0.3, 0.3, 0.4], generator=generator2) + + Args: + dataset (Dataset): Dataset to be split + lengths (sequence): lengths or fractions of splits to be produced + generator (Generator): Generator used for the random permutation. + """ + if math.isclose(sum(lengths), 1) and sum(lengths) <= 1: + subset_lengths: list[int] = [] + for i, frac in enumerate(lengths): + if frac < 0 or frac > 1: + raise ValueError(f"Fraction at index {i} is not between 0 and 1") + n_items_in_split = math.floor(len(dataset) * frac) # type: ignore[arg-type] + subset_lengths.append(n_items_in_split) + remainder = len(dataset) - sum(subset_lengths) # type: ignore[arg-type] + # add 1 to all the lengths in round-robin fashion until the remainder is 0 + for i in range(remainder): + idx_to_add_at = i % len(subset_lengths) + subset_lengths[idx_to_add_at] += 1 + lengths = subset_lengths + for i, length in enumerate(lengths): + if length == 0: + warnings.warn( + f"Length of split at index {i} is 0. " + f"This might result in an empty dataset.", + stacklevel=2, + ) + + # Cannot verify that dataset is Sized + if sum(lengths) != len(dataset): # type: ignore[arg-type] + raise ValueError( + "Sum of input lengths does not equal the length of the input dataset!" + ) + + indices = randperm(sum(lengths), generator=generator).tolist() # type: ignore[arg-type, call-overload] + lengths = cast(Sequence[int], lengths) + return [ + Subset(dataset, indices[offset - length : offset]) + for offset, length in zip(itertools.accumulate(lengths), lengths, strict=True) + ] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/distributed.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/distributed.py new file mode 100644 index 0000000000000000000000000000000000000000..5179d7698ffee0f2acda62a2b2073df176aae794 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/distributed.py @@ -0,0 +1,157 @@ +import math +from collections.abc import Iterator +from typing import TypeVar + +import torch +import torch.distributed as dist +from torch.utils.data.dataset import Dataset +from torch.utils.data.sampler import Sampler + + +__all__ = ["DistributedSampler"] + + +_T_co = TypeVar("_T_co", covariant=True) + + +class DistributedSampler(Sampler[_T_co]): + r"""Sampler that restricts data loading to a subset of the dataset. + + It is especially useful in conjunction with + :class:`torch.nn.parallel.DistributedDataParallel`. In such a case, each + process can pass a :class:`~torch.utils.data.DistributedSampler` instance as a + :class:`~torch.utils.data.DataLoader` sampler, and load a subset of the + original dataset that is exclusive to it. + + .. note:: + Dataset is assumed to be of constant size and that any instance of it always + returns the same elements in the same order. + + Args: + dataset: Dataset used for sampling. + num_replicas (int, optional): Number of processes participating in + distributed training. By default, :attr:`world_size` is retrieved from the + current distributed group. + rank (int, optional): Rank of the current process within :attr:`num_replicas`. + By default, :attr:`rank` is retrieved from the current distributed + group. + shuffle (bool, optional): If ``True`` (default), sampler will shuffle the + indices. + seed (int, optional): random seed used to shuffle the sampler if + :attr:`shuffle=True`. This number should be identical across all + processes in the distributed group. Default: ``0``. + drop_last (bool, optional): if ``True``, then the sampler will drop the + tail of the data to make it evenly divisible across the number of + replicas. If ``False``, the sampler will add extra indices to make + the data evenly divisible across the replicas. Default: ``False``. + + .. warning:: + In distributed mode, calling the :meth:`set_epoch` method at + the beginning of each epoch **before** creating the :class:`DataLoader` iterator + is necessary to make shuffling work properly across multiple epochs. Otherwise, + the same ordering will be always used. + + Example:: + + >>> # xdoctest: +SKIP + >>> sampler = DistributedSampler(dataset) if is_distributed else None + >>> loader = DataLoader(dataset, shuffle=(sampler is None), + ... sampler=sampler) + >>> for epoch in range(start_epoch, n_epochs): + ... if is_distributed: + ... sampler.set_epoch(epoch) + ... train(loader) + """ + + def __init__( + self, + dataset: Dataset, + num_replicas: int | None = None, + rank: int | None = None, + shuffle: bool = True, + seed: int = 0, + drop_last: bool = False, + ) -> None: + if num_replicas is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + num_replicas = dist.get_world_size() + if rank is None: + if not dist.is_available(): + raise RuntimeError("Requires distributed package to be available") + rank = dist.get_rank() + if rank >= num_replicas or rank < 0: + raise ValueError( + f"Invalid rank {rank}, rank should be in the interval [0, {num_replicas - 1}]" + ) + self.dataset = dataset + self.num_replicas = num_replicas + self.rank = rank + self.epoch = 0 + self.drop_last = drop_last + # If the dataset length is evenly divisible by # of replicas, then there + # is no need to drop any data, since the dataset will be split equally. + if self.drop_last and len(self.dataset) % self.num_replicas != 0: # type: ignore[arg-type] + # Split to nearest available length that is evenly divisible. + # This is to ensure each rank receives the same amount of data when + # using this Sampler. + self.num_samples = math.ceil( + (len(self.dataset) - self.num_replicas) / self.num_replicas # type: ignore[arg-type] + ) + else: + self.num_samples = math.ceil(len(self.dataset) / self.num_replicas) # type: ignore[arg-type] + self.total_size = self.num_samples * self.num_replicas + self.shuffle = shuffle + self.seed = seed + + def __iter__(self) -> Iterator[_T_co]: + if self.shuffle: + # deterministically shuffle based on epoch and seed + g = torch.Generator() + g.manual_seed(self.seed + self.epoch) + indices = torch.randperm(len(self.dataset), generator=g).tolist() # type: ignore[arg-type] + else: + indices = list(range(len(self.dataset))) # type: ignore[arg-type] + + if not self.drop_last: + # add extra samples to make it evenly divisible + padding_size = self.total_size - len(indices) + if padding_size <= len(indices): + indices += indices[:padding_size] + else: + indices += (indices * math.ceil(padding_size / len(indices)))[ + :padding_size + ] + else: + # remove tail of data to make it evenly divisible. + indices = indices[: self.total_size] + if len(indices) != self.total_size: + raise AssertionError( + f"Number of indices ({len(indices)}) does not match total_size ({self.total_size})" + ) + + # subsample + indices = indices[self.rank : self.total_size : self.num_replicas] + if len(indices) != self.num_samples: + raise AssertionError( + f"Number of subsampled indices ({len(indices)}) does not match num_samples ({self.num_samples})" + ) + + # pyrefly: ignore [bad-return] + return iter(indices) + + def __len__(self) -> int: + return self.num_samples + + def set_epoch(self, epoch: int) -> None: + r""" + Set the epoch for this sampler. + + When :attr:`shuffle=True`, this ensures all replicas + use a different random ordering for each epoch. Otherwise, the next iteration of this + sampler will yield the same ordering. + + Args: + epoch (int): Epoch number. + """ + self.epoch = epoch diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/graph.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/graph.py new file mode 100644 index 0000000000000000000000000000000000000000..3f8dbb17ceaea9c61f46af7ce29666361f45094b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/graph.py @@ -0,0 +1,163 @@ +# mypy: allow-untyped-defs +import io +import pickle +import warnings +from collections.abc import Collection + +from torch.utils._import_utils import dill_available +from torch.utils.data.datapipes.datapipe import IterDataPipe, MapDataPipe + + +__all__ = ["traverse", "traverse_dps"] + +DataPipe = IterDataPipe | MapDataPipe +# pyrefly: ignore [invalid-type-alias] +DataPipeGraph = dict[int, tuple[DataPipe, "DataPipeGraph"]] + + +def _stub_unpickler() -> str: + return "STUB" + + +# TODO(VitalyFedyunin): Make sure it works without dill module installed +def _list_connected_datapipes( + scan_obj: DataPipe, only_datapipe: bool, cache: set[int] +) -> list[DataPipe]: + f = io.BytesIO() + p = pickle.Pickler( + f + ) # Not going to work for lambdas, but dill infinite loops on typing and can't be used as is + if dill_available(): + from dill import Pickler as dill_Pickler + + d = dill_Pickler(f) + else: + d = None + + captured_connections = [] + + def getstate_hook(ori_state): + state = None + if isinstance(ori_state, dict): + state = {} + for k, v in ori_state.items(): + if isinstance(v, (IterDataPipe, MapDataPipe, Collection)): + state[k] = v + elif isinstance(ori_state, (tuple, list)): + state = [] # type: ignore[assignment] + for v in ori_state: + if isinstance(v, (IterDataPipe, MapDataPipe, Collection)): + state.append(v) # type: ignore[attr-defined] + elif isinstance(ori_state, (IterDataPipe, MapDataPipe, Collection)): + state = ori_state # type: ignore[assignment] + return state + + def reduce_hook(obj): + if obj == scan_obj or id(obj) in cache: + raise NotImplementedError + else: + captured_connections.append(obj) + # Adding id to remove duplicate DataPipe serialized at the same level + cache.add(id(obj)) + return _stub_unpickler, () + + datapipe_classes: tuple[type[DataPipe]] = (IterDataPipe, MapDataPipe) # type: ignore[assignment] + + try: + for cls in datapipe_classes: + cls.set_reduce_ex_hook(reduce_hook) + if only_datapipe: + cls.set_getstate_hook(getstate_hook) + try: + p.dump(scan_obj) + except (pickle.PickleError, AttributeError, TypeError): + if dill_available(): + # pyrefly: ignore [missing-attribute] + d.dump(scan_obj) + else: + raise + finally: + for cls in datapipe_classes: + cls.set_reduce_ex_hook(None) + if only_datapipe: + cls.set_getstate_hook(None) + if dill_available(): + from dill import extend as dill_extend + + dill_extend(False) # Undo change to dispatch table + return captured_connections + + +def traverse_dps(datapipe: DataPipe) -> DataPipeGraph: + r""" + Traverse the DataPipes and their attributes to extract the DataPipe graph. + + This only looks into the attribute from each DataPipe that is either a + DataPipe and a Python collection object such as ``list``, ``tuple``, + ``set`` and ``dict``. + + Args: + datapipe: the end DataPipe of the graph + Returns: + A graph represented as a nested dictionary, where keys are ids of DataPipe instances + and values are tuples of DataPipe instance and the sub-graph + """ + cache: set[int] = set() + return _traverse_helper(datapipe, only_datapipe=True, cache=cache) + + +def traverse(datapipe: DataPipe, only_datapipe: bool | None = None) -> DataPipeGraph: + r""" + Traverse the DataPipes and their attributes to extract the DataPipe graph. + + [Deprecated] + When ``only_dataPipe`` is specified as ``True``, it would only look into the + attribute from each DataPipe that is either a DataPipe and a Python collection object + such as ``list``, ``tuple``, ``set`` and ``dict``. + + Note: + This function is deprecated. Please use `traverse_dps` instead. + + Args: + datapipe: the end DataPipe of the graph + only_datapipe: If ``False`` (default), all attributes of each DataPipe are traversed. + This argument is deprecating and will be removed after the next release. + Returns: + A graph represented as a nested dictionary, where keys are ids of DataPipe instances + and values are tuples of DataPipe instance and the sub-graph + """ + msg = ( + "`traverse` function and will be removed after 1.13. " + "Please use `traverse_dps` instead." + ) + if not only_datapipe: + msg += " And, the behavior will be changed to the equivalent of `only_datapipe=True`." + warnings.warn(msg, FutureWarning, stacklevel=2) + if only_datapipe is None: + only_datapipe = False + cache: set[int] = set() + return _traverse_helper(datapipe, only_datapipe, cache) + + +# Add cache here to prevent infinite recursion on DataPipe +def _traverse_helper( + datapipe: DataPipe, only_datapipe: bool, cache: set[int] +) -> DataPipeGraph: + if not isinstance(datapipe, (IterDataPipe, MapDataPipe)): + raise RuntimeError( + f"Expected `IterDataPipe` or `MapDataPipe`, but {type(datapipe)} is found" + ) + + dp_id = id(datapipe) + if dp_id in cache: + return {} + cache.add(dp_id) + # Using cache.copy() here is to prevent the same DataPipe pollutes the cache on different paths + items = _list_connected_datapipes(datapipe, only_datapipe, cache.copy()) + d: DataPipeGraph = {dp_id: (datapipe, {})} + for item in items: + # Using cache.copy() here is to prevent recursion on a single path rather than global graph + # Single DataPipe can present multiple times in different paths in graph + # pyrefly: ignore [no-matching-overload] + d[dp_id][1].update(_traverse_helper(item, only_datapipe, cache.copy())) + return d diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/graph_settings.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/graph_settings.py new file mode 100644 index 0000000000000000000000000000000000000000..03096398a6738b29c22aad044caaf16e4c45a7d0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/graph_settings.py @@ -0,0 +1,173 @@ +# mypy: allow-untyped-defs +import inspect +import warnings +from typing import Any +from typing_extensions import deprecated + +import torch +from torch.utils.data.datapipes.iter.sharding import ( + _ShardingIterDataPipe, + SHARDING_PRIORITIES, +) +from torch.utils.data.graph import DataPipe, DataPipeGraph, traverse_dps + + +__all__ = [ + "apply_random_seed", + "apply_sharding", + "apply_shuffle_seed", + "apply_shuffle_settings", + "get_all_graph_pipes", +] + + +def get_all_graph_pipes(graph: DataPipeGraph) -> list[DataPipe]: + return _get_all_graph_pipes_helper(graph, set()) + + +def _get_all_graph_pipes_helper( + graph: DataPipeGraph, id_cache: set[int] +) -> list[DataPipe]: + results: list[DataPipe] = [] + for dp_id, (datapipe, sub_graph) in graph.items(): + if dp_id in id_cache: + continue + id_cache.add(dp_id) + results.append(datapipe) + results.extend(_get_all_graph_pipes_helper(sub_graph, id_cache)) + return results + + +def _is_sharding_datapipe(datapipe: DataPipe) -> bool: + return isinstance(datapipe, _ShardingIterDataPipe) or ( + hasattr(datapipe, "apply_sharding") + and inspect.ismethod(datapipe.apply_sharding) + ) + + +def apply_sharding( + datapipe: DataPipe, + num_of_instances: int, + instance_id: int, + sharding_group=SHARDING_PRIORITIES.DEFAULT, +) -> DataPipe: + r""" + Apply dynamic sharding over the ``sharding_filter`` DataPipe that has a method ``apply_sharding``. + + RuntimeError will be raised when multiple ``sharding_filter`` are presented in the same branch. + """ + graph = traverse_dps(datapipe) + + def _helper(graph, prev_applied=None) -> None: + for dp, sub_graph in graph.values(): + applied = None + if _is_sharding_datapipe(dp): + if prev_applied is not None: + raise RuntimeError( + "Sharding twice on a single pipeline is likely unintended and will cause data loss. " + f"Sharding already applied to {prev_applied} while trying to apply to {dp}" + ) + # For BC, only provide sharding_group if accepted + sig = inspect.signature(dp.apply_sharding) + if len(sig.parameters) < 3: + dp.apply_sharding(num_of_instances, instance_id) + else: + dp.apply_sharding( + num_of_instances, instance_id, sharding_group=sharding_group + ) + applied = dp + if applied is None: + applied = prev_applied + _helper(sub_graph, applied) + + _helper(graph) + + return datapipe + + +def _is_shuffle_datapipe(datapipe: DataPipe) -> bool: + return ( + hasattr(datapipe, "set_shuffle") + and hasattr(datapipe, "set_seed") + and inspect.ismethod(datapipe.set_shuffle) + and inspect.ismethod(datapipe.set_seed) + ) + + +def apply_shuffle_settings(datapipe: DataPipe, shuffle: bool | None = None) -> DataPipe: + r""" + Traverse the graph of ``DataPipes`` to find and set shuffle attribute. + + Apply the method to each `DataPipe` that has APIs of ``set_shuffle`` + and ``set_seed``. + + Args: + datapipe: DataPipe that needs to set shuffle attribute + shuffle: Shuffle option (default: ``None`` and no-op to the graph) + """ + if shuffle is None: + return datapipe + + graph = traverse_dps(datapipe) + all_pipes = get_all_graph_pipes(graph) + shufflers = [pipe for pipe in all_pipes if _is_shuffle_datapipe(pipe)] + if not shufflers and shuffle: + warnings.warn( + "`shuffle=True` was set, but the datapipe does not contain a `Shuffler`. Adding one at the end. " + "Be aware that the default buffer size might not be sufficient for your task.", + stacklevel=2, + ) + datapipe = datapipe.shuffle() + shufflers = [ + datapipe, + ] + + for shuffler in shufflers: + shuffler.set_shuffle(shuffle) + + return datapipe + + +@deprecated( + "`apply_shuffle_seed` is deprecated since 1.12 and will be removed in the future releases. " + "Please use `apply_random_seed` instead.", + category=FutureWarning, +) +def apply_shuffle_seed(datapipe: DataPipe, rng: Any) -> DataPipe: + return apply_random_seed(datapipe, rng) + + +def _is_random_datapipe(datapipe: DataPipe) -> bool: + return hasattr(datapipe, "set_seed") and inspect.ismethod(datapipe.set_seed) + + +def apply_random_seed(datapipe: DataPipe, rng: torch.Generator) -> DataPipe: + r""" + Traverse the graph of ``DataPipes`` to find random ``DataPipe`` with an API of ``set_seed``. + + Then set the random seed based on the provided RNG to those ``DataPipe``. + + Args: + datapipe: DataPipe that needs to set randomness + rng: Random number generator to generate random seeds + """ + graph = traverse_dps(datapipe) + all_pipes = get_all_graph_pipes(graph) + # Using a set to track id of DataPipe to prevent setting randomness per DataPipe more than once. + # And, `id` is used in case of unhashable DataPipe + cache = set() + random_datapipes = [] + for pipe in all_pipes: + if id(pipe) in cache: + continue + if _is_random_datapipe(pipe): + random_datapipes.append(pipe) + cache.add(id(pipe)) + + for pipe in random_datapipes: + random_seed = int( + torch.empty((), dtype=torch.int64).random_(generator=rng).item() + ) + pipe.set_seed(random_seed) + + return datapipe diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/sampler.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..aa13bb8e0a3e146bd7bfbc766fdfcb822efa9313 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/data/sampler.py @@ -0,0 +1,354 @@ +# mypy: allow-untyped-defs +import itertools +from collections.abc import Iterable, Iterator, Sequence, Sized +from typing import Generic, TypeVar + +import torch + + +# Note: For benchmarking changes to samplers, see: +# /benchmarks/data/samplers_bench.py +# This benchmark compares the performance of different sampler implementations +# and can be used to evaluate the impact of optimizations. + + +__all__ = [ + "BatchSampler", + "RandomSampler", + "Sampler", + "SequentialSampler", + "SubsetRandomSampler", + "WeightedRandomSampler", +] + + +_T_co = TypeVar("_T_co", covariant=True) + + +class Sampler(Generic[_T_co]): + r"""Base class for all Samplers. + + Every Sampler subclass has to provide an :meth:`__iter__` method, providing a + way to iterate over indices or lists of indices (batches) of dataset elements, + and may provide a :meth:`__len__` method that returns the length of the returned iterators. + + Example: + >>> # xdoctest: +SKIP + >>> class AccedingSequenceLengthSampler(Sampler[int]): + >>> def __init__(self, data: List[str]) -> None: + >>> self.data = data + >>> + >>> def __len__(self) -> int: + >>> return len(self.data) + >>> + >>> def __iter__(self) -> Iterator[int]: + >>> sizes = torch.tensor([len(x) for x in self.data]) + >>> yield from torch.argsort(sizes).tolist() + >>> + >>> class AccedingSequenceLengthBatchSampler(Sampler[List[int]]): + >>> def __init__(self, data: List[str], batch_size: int) -> None: + >>> self.data = data + >>> self.batch_size = batch_size + >>> + >>> def __len__(self) -> int: + >>> return (len(self.data) + self.batch_size - 1) // self.batch_size + >>> + >>> def __iter__(self) -> Iterator[List[int]]: + >>> sizes = torch.tensor([len(x) for x in self.data]) + >>> for batch in torch.chunk(torch.argsort(sizes), len(self)): + >>> yield batch.tolist() + + .. note:: The :meth:`__len__` method isn't strictly required by + :class:`~torch.utils.data.DataLoader`, but is expected in any + calculation involving the length of a :class:`~torch.utils.data.DataLoader`. + """ + + def __iter__(self) -> Iterator[_T_co]: + raise NotImplementedError + + # NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] + # + # Many times we have an abstract class representing a collection/iterable of + # data, e.g., `torch.utils.data.Sampler`, with its subclasses optionally + # implementing a `__len__` method. In such cases, we must make sure to not + # provide a default implementation, because both straightforward default + # implementations have their issues: + # + # + `return NotImplemented`: + # Calling `len(subclass_instance)` raises: + # TypeError: 'NotImplementedType' object cannot be interpreted as an integer + # + # + `raise NotImplementedError`: + # This prevents triggering some fallback behavior. E.g., the built-in + # `list(X)` tries to call `len(X)` first, and executes a different code + # path if the method is not found or `NotImplemented` is returned, while + # raising a `NotImplementedError` will propagate and make the call fail + # where it could have used `__iter__` to complete the call. + # + # Thus, the only two sensible things to do are + # + # + **not** provide a default `__len__`. + # + # + raise a `TypeError` instead, which is what Python uses when users call + # a method that is not defined on an object. + # (@ssnl verifies that this works on at least Python 3.7.) + + +class SequentialSampler(Sampler[int]): + r"""Samples elements sequentially, always in the same order. + + Args: + data_source (Sized): data source to sample from. Must implement __len__. + """ + + data_source: Sized + + def __init__(self, data_source: Sized) -> None: + self.data_source = data_source + + def __iter__(self) -> Iterator[int]: + return iter(range(len(self.data_source))) + + def __len__(self) -> int: + return len(self.data_source) + + +class RandomSampler(Sampler[int]): + r"""Samples elements randomly. If without replacement, then sample from a shuffled dataset. + + If with replacement, then user can specify :attr:`num_samples` to draw. + + Args: + data_source (Sized): data source to sample from. Must implement __len__. + replacement (bool): samples are drawn on-demand with replacement if ``True``, default=``False`` + num_samples (int): number of samples to draw, default=`len(dataset)`. + generator (Generator): Generator used in sampling. + """ + + data_source: Sized + replacement: bool + + def __init__( + self, + data_source: Sized, + replacement: bool = False, + num_samples: int | None = None, + generator=None, + ) -> None: + self.data_source = data_source + self.replacement = replacement + self._num_samples = num_samples + self.generator = generator + + if not isinstance(self.replacement, bool): + raise TypeError( + f"replacement should be a boolean value, but got replacement={self.replacement}" + ) + + if not isinstance(self.num_samples, int) or self.num_samples <= 0: + raise ValueError( + f"num_samples should be a positive integer value, but got num_samples={self.num_samples}" + ) + + @property + def num_samples(self) -> int: + # dataset size might change at runtime + if self._num_samples is None: + return len(self.data_source) + return self._num_samples + + def __iter__(self) -> Iterator[int]: + n = len(self.data_source) + if self.generator is None: + seed = int(torch.empty((), dtype=torch.int64).random_().item()) + generator = torch.Generator() + generator.manual_seed(seed) + else: + generator = self.generator + + if self.replacement: + for _ in range(self.num_samples // 32): + yield from torch.randint( + high=n, size=(32,), dtype=torch.int64, generator=generator + ).tolist() + yield from torch.randint( + high=n, + size=(self.num_samples % 32,), + dtype=torch.int64, + generator=generator, + ).tolist() + else: + for _ in range(self.num_samples // n): + yield from torch.randperm(n, generator=generator).tolist() + yield from torch.randperm(n, generator=generator).tolist()[ + : self.num_samples % n + ] + + def __len__(self) -> int: + return self.num_samples + + +class SubsetRandomSampler(Sampler[int]): + r"""Samples elements randomly from a given list of indices, without replacement. + + Args: + indices (sequence): a sequence of indices + generator (Generator): Generator used in sampling. + """ + + indices: Sequence[int] + + def __init__(self, indices: Sequence[int], generator=None) -> None: + self.indices = indices + self.generator = generator + + def __iter__(self) -> Iterator[int]: + for i in torch.randperm(len(self.indices), generator=self.generator).tolist(): + yield self.indices[i] + + def __len__(self) -> int: + return len(self.indices) + + +class WeightedRandomSampler(Sampler[int]): + r"""Samples elements from ``[0,..,len(weights)-1]`` with given probabilities (weights). + + Args: + weights (sequence) : a sequence of weights, not necessary summing up to one + num_samples (int): number of samples to draw + replacement (bool): if ``True``, samples are drawn with replacement. + If not, they are drawn without replacement, which means that when a + sample index is drawn for a row, it cannot be drawn again for that row. + generator (Generator): Generator used in sampling. + + Example: + >>> # xdoctest: +IGNORE_WANT("non-deterministic") + >>> list( + ... WeightedRandomSampler( + ... [0.1, 0.9, 0.4, 0.7, 3.0, 0.6], 5, replacement=True + ... ) + ... ) + [4, 4, 1, 4, 5] + >>> list( + ... WeightedRandomSampler( + ... [0.9, 0.4, 0.05, 0.2, 0.3, 0.1], 5, replacement=False + ... ) + ... ) + [0, 1, 4, 3, 2] + """ + + weights: torch.Tensor + num_samples: int + replacement: bool + + def __init__( + self, + weights: Sequence[float], + num_samples: int, + replacement: bool = True, + generator=None, + ) -> None: + if ( + not isinstance(num_samples, int) + or isinstance(num_samples, bool) + or num_samples <= 0 + ): + raise ValueError( + f"num_samples should be a positive integer value, but got num_samples={num_samples}" + ) + if not isinstance(replacement, bool): + raise ValueError( + f"replacement should be a boolean value, but got replacement={replacement}" + ) + + weights_tensor = torch.as_tensor(weights, dtype=torch.double) + if len(weights_tensor.shape) != 1: + raise ValueError( + "weights should be a 1d sequence but given " + f"weights have shape {tuple(weights_tensor.shape)}" + ) + + self.weights = weights_tensor + self.num_samples = num_samples + self.replacement = replacement + self.generator = generator + + def __iter__(self) -> Iterator[int]: + rand_tensor = torch.multinomial( + self.weights, self.num_samples, self.replacement, generator=self.generator + ) + yield from iter(rand_tensor.tolist()) + + def __len__(self) -> int: + return self.num_samples + + +class BatchSampler(Sampler[list[int]]): + r"""Wraps another sampler to yield a mini-batch of indices. + + Args: + sampler (Sampler or Iterable): Base sampler. Can be any iterable object + batch_size (int): Size of mini-batch. + drop_last (bool): If ``True``, the sampler will drop the last batch if + its size would be less than ``batch_size`` + + Example: + >>> list( + ... BatchSampler( + ... SequentialSampler(range(10)), batch_size=3, drop_last=False + ... ) + ... ) + [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9]] + >>> list( + ... BatchSampler(SequentialSampler(range(10)), batch_size=3, drop_last=True) + ... ) + [[0, 1, 2], [3, 4, 5], [6, 7, 8]] + """ + + def __init__( + self, + sampler: Sampler[int] | Iterable[int], + batch_size: int, + drop_last: bool, + ) -> None: + # Since collections.abc.Iterable does not check for `__getitem__`, which + # is one way for an object to be an iterable, we don't do an `isinstance` + # check here. + if ( + not isinstance(batch_size, int) + or isinstance(batch_size, bool) + or batch_size <= 0 + ): + raise ValueError( + f"batch_size should be a positive integer value, but got batch_size={batch_size}" + ) + if not isinstance(drop_last, bool): + raise ValueError( + f"drop_last should be a boolean value, but got drop_last={drop_last}" + ) + self.sampler = sampler + self.batch_size = batch_size + self.drop_last = drop_last + + def __iter__(self) -> Iterator[list[int]]: + sampler_iter = iter(self.sampler) + if self.drop_last: + # Create multiple references to the same iterator + args = [sampler_iter] * self.batch_size + for batch_droplast in zip(*args, strict=False): + yield [*batch_droplast] + else: + batch = [*itertools.islice(sampler_iter, self.batch_size)] + while batch: + yield batch + batch = [*itertools.islice(sampler_iter, self.batch_size)] + + def __len__(self) -> int: + # Can only be called if self.sampler has __len__ implemented + # We cannot enforce this condition, so we turn off typechecking for the + # implementation below. + # Somewhat related: see NOTE [ Lack of Default `__len__` in Python Abstract Base Classes ] + if self.drop_last: + return len(self.sampler) // self.batch_size # type: ignore[arg-type] + else: + return (len(self.sampler) + self.batch_size - 1) // self.batch_size # type: ignore[arg-type] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/deterministic.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/deterministic.py new file mode 100644 index 0000000000000000000000000000000000000000..a055c43be531a5c65c4f29f6c8165104e98e5ca0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/deterministic.py @@ -0,0 +1,22 @@ +# mypy: allow-untyped-defs +import sys +import types + +import torch + + +class _Deterministic(types.ModuleType): + @property + def fill_uninitialized_memory(self): + """ + Whether to fill uninitialized memory with a known value when + :meth:`torch.use_deterministic_algorithms()` is set to ``True``. + """ + return torch._C._get_deterministic_fill_uninitialized_memory() + + @fill_uninitialized_memory.setter + def fill_uninitialized_memory(self, mode): + return torch._C._set_deterministic_fill_uninitialized_memory(mode) + + +sys.modules[__name__].__class__ = _Deterministic diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/dlpack.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/dlpack.py new file mode 100644 index 0000000000000000000000000000000000000000..aef32100ee7105d364f0e144dc1cf2e0368f7767 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/dlpack.py @@ -0,0 +1,231 @@ +from typing import Any + +import torch +import enum + +from torch._C import _to_dlpack as to_dlpack +from torch.types import Device as _Device + +__all__ = [ + "DLDeviceType", + "from_dlpack", +] + +class DLDeviceType(enum.IntEnum): + # Enums as in DLPack specification (aten/src/ATen/dlpack.h) + kDLCPU = 1, + kDLCUDA = 2, + kDLCUDAHost = 3, + kDLOpenCL = 4, + kDLVulkan = 7, + kDLMetal = 8, + kDLVPI = 9, + kDLROCM = 10, + kDLROCMHost = 11, + kDLExtDev = 12, + kDLCUDAManaged = 13, + kDLOneAPI = 14, + kDLWebGPU = 15, + kDLHexagon = 16, + kDLMAIA = 17, + + +torch._C._add_docstr(to_dlpack, r"""to_dlpack(tensor) -> PyCapsule + +Returns an opaque object (a "DLPack capsule") representing the tensor. + +.. note:: + ``to_dlpack`` is a legacy DLPack interface. The capsule it returns + cannot be used for anything in Python other than use it as input to + ``from_dlpack``. The more idiomatic use of DLPack is to call + ``from_dlpack`` directly on the tensor object - this works when that + object has a ``__dlpack__`` method, which PyTorch and most other + libraries indeed have now. + +.. warning:: + Only call ``from_dlpack`` once per capsule produced with ``to_dlpack``. + Behavior when a capsule is consumed multiple times is undefined. + +Args: + tensor: a tensor to be exported + +The DLPack capsule shares the tensor's memory. +""") + + +# TODO: add a typing.Protocol to be able to tell Mypy that only objects with +# __dlpack__ and __dlpack_device__ methods are accepted. +def from_dlpack( + ext_tensor: Any, + *, + device: _Device | None = None, + copy: bool | None = None +) -> 'torch.Tensor': + """from_dlpack(ext_tensor) -> Tensor + + Converts a tensor from an external library into a ``torch.Tensor``. + + The returned PyTorch tensor will share the memory with the input tensor + (which may have come from another library). Note that in-place operations + will therefore also affect the data of the input tensor. This may lead to + unexpected issues (e.g., other libraries may have read-only flags or + immutable data structures), so the user should only do this if they know + for sure that this is fine. + + Args: + ext_tensor (object with ``__dlpack__`` attribute, or a DLPack capsule): + The tensor or DLPack capsule to convert. + + If ``ext_tensor`` is a tensor (or ndarray) object, it must support + the ``__dlpack__`` protocol (i.e., have a ``ext_tensor.__dlpack__`` + method). Otherwise ``ext_tensor`` may be a DLPack capsule, which is + an opaque ``PyCapsule`` instance, typically produced by a + ``to_dlpack`` function or method. + + device (torch.device or str or None): An optional PyTorch device + specifying where to place the new tensor. If None (default), the + new tensor will be on the same device as ``ext_tensor``. + + copy (bool or None): An optional boolean indicating whether or not to copy + ``self``. If None, PyTorch will copy only if necessary. + + Examples:: + + >>> import torch.utils.dlpack + >>> t = torch.arange(4) + + # Convert a tensor directly (supported in PyTorch >= 1.10) + >>> t2 = torch.from_dlpack(t) + >>> t2[:2] = -1 # show that memory is shared + >>> t2 + tensor([-1, -1, 2, 3]) + >>> t + tensor([-1, -1, 2, 3]) + + # The old-style DLPack usage, with an intermediate capsule object + >>> capsule = torch.utils.dlpack.to_dlpack(t) + >>> capsule + + >>> t3 = torch.from_dlpack(capsule) + >>> t3 + tensor([-1, -1, 2, 3]) + >>> t3[0] = -9 # now we're sharing memory between 3 tensors + >>> t3 + tensor([-9, -1, 2, 3]) + >>> t2 + tensor([-9, -1, 2, 3]) + >>> t + tensor([-9, -1, 2, 3]) + + """ + + if hasattr(ext_tensor, '__dlpack__'): + # Only populate kwargs if any of the optional arguments are, in fact, not None. Otherwise, + # leave them out, since we might end up falling back to no-extra-kwargs __dlpack__ call. + kwargs: dict[str, Any] = {} + kwargs["max_version"] = (1, 0) + + # Track copy request for potential manual handling + requested_copy = copy + producer_handled_copy = True + cross_device_transfer = False # Will be set to True if device transfer is needed + + if copy is not None: + kwargs["copy"] = copy + + # Parse the device parameter. + # At this moment, it can either be a torch.device or a str representing + # a torch.device, e.g. "cpu", "cuda", etc. + # Get source device first (we need it to detect cross-device transfers) + ext_device = ext_tensor.__dlpack_device__() + + if device is not None: + if isinstance(device, str): + device = torch.device(device) + if not isinstance(device, torch.device): + raise AssertionError(f"from_dlpack: unsupported device type: {type(device)}") + + # Convert target device to DLPack format + target_dl_device = torch._C._torchDeviceToDLDevice(device) + + # Detect cross-device transfer by comparing source and target devices + # E.g. CPU->CUDA, cuda:0->cuda:1, etc. + cross_device_transfer = (ext_device != target_dl_device) + + # Only pass dl_device to producer if NOT cross-device transfer + if not cross_device_transfer: + kwargs["dl_device"] = target_dl_device + + # Cross-device transfer always requires a copy + if cross_device_transfer and copy is False: + raise ValueError( + f"cannot move DLPack tensor from device {ext_device} to {target_dl_device} " + "without copying. Set copy=None or copy=True." + ) + + # ext_device is either CUDA or ROCm, we need to pass the current + # stream + if ext_device[0] in (DLDeviceType.kDLCUDA, DLDeviceType.kDLROCM): + stream = torch.cuda.current_stream(f'cuda:{ext_device[1]}') + # cuda_stream is the pointer to the stream and it is a public + # attribute, but it is not documented + # The array API specify that the default legacy stream must be passed + # with a value of 1 for CUDA + # https://data-apis.org/array-api/latest/API_specification/array_object.html?dlpack-self-stream-none#dlpack-self-stream-none + is_cuda = ext_device[0] == DLDeviceType.kDLCUDA + # Since pytorch is not using PTDS by default, lets directly pass + # the legacy stream + stream_ptr = 1 if is_cuda and stream.cuda_stream == 0 else stream.cuda_stream + kwargs["stream"] = stream_ptr + + # Try different parameter combinations until one works + dlpack = None + + # Attempt 1: Try with all the parameters + try: + dlpack = ext_tensor.__dlpack__(**kwargs) + except TypeError: + pass + + # Attempt 2: Remove max_version + if dlpack is None: + kwargs.pop("max_version", None) + try: + dlpack = ext_tensor.__dlpack__(**kwargs) + except TypeError: + pass + + # Attempt 3: Remove copy + if dlpack is None: + kwargs.pop("copy", None) + producer_handled_copy = False + try: + dlpack = ext_tensor.__dlpack__(**kwargs) + except TypeError: + pass + + # Attempt 4: Remove dl_device + if dlpack is None: + kwargs.pop("dl_device", None) + dlpack = ext_tensor.__dlpack__(**kwargs) + + tensor = torch._C._from_dlpack(dlpack) + + # Manual copy if producer didn't handle it (cross-device already copies via .to()) + if requested_copy is True and not producer_handled_copy and not cross_device_transfer: + tensor = tensor.clone() + + # Handle cross-device transfer by moving tensor to target device + if cross_device_transfer: + tensor = tensor.to(device) + + return tensor + + else: + if device is not None or copy is not None: + raise AssertionError( + "device and copy kwargs not supported when ext_tensor is already a DLPack capsule." + ) + # Old versions just call the converter + dlpack = ext_tensor + return torch._C._from_dlpack(dlpack) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/file_baton.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/file_baton.py new file mode 100644 index 0000000000000000000000000000000000000000..74de1495c8fc572cfc350966f6f0b7997fde7ad1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/file_baton.py @@ -0,0 +1,63 @@ +# mypy: allow-untyped-defs +import os +import time +import warnings + + +class FileBaton: + """A primitive, file-based synchronization utility.""" + + def __init__(self, lock_file_path, wait_seconds=0.1, warn_after_seconds=None) -> None: + """ + Create a new :class:`FileBaton`. + + Args: + lock_file_path: The path to the file used for locking. + wait_seconds: The seconds to periodically sleep (spin) when + calling ``wait()``. + warn_after_seconds: The seconds to wait before showing + lock file path to warn existing lock file. + """ + self.lock_file_path = lock_file_path + self.wait_seconds = wait_seconds + self.fd = None + self.warn_after_seconds = warn_after_seconds + + def try_acquire(self) -> bool | None: + """ + Try to atomically create a file under exclusive access. + + Returns: + True if the file could be created, else False. + """ + try: + self.fd = os.open(self.lock_file_path, os.O_CREAT | os.O_EXCL) + return True + except FileExistsError: + return False + + def wait(self) -> None: + """ + Periodically sleeps for a certain amount until the baton is released. + + The amount of time slept depends on the ``wait_seconds`` parameter + passed to the constructor. + """ + has_warned = False + + start_time = time.time() + while os.path.exists(self.lock_file_path): + time.sleep(self.wait_seconds) + + if self.warn_after_seconds is not None: + if time.time() - start_time > self.warn_after_seconds and not has_warned: + warnings.warn(f'Waited on lock file "{self.lock_file_path}" for ' + f'{self.warn_after_seconds} seconds.', stacklevel=2) + has_warned = True + + def release(self) -> None: + """Release the baton and removes its file.""" + if self.fd is not None: + os.close(self.fd) + + os.remove(self.lock_file_path) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/flop_counter.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/flop_counter.py new file mode 100644 index 0000000000000000000000000000000000000000..5f015a32f9c313c06a7274a7b38c82e5a0b4c362 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/flop_counter.py @@ -0,0 +1,1017 @@ +# mypy: allow-untyped-defs +from types import NoneType +import logging +import torch +from torch.utils._pytree import tree_map, tree_flatten, tree_unflatten +from .module_tracker import ModuleTracker +from typing import Any, TypeVar +from collections.abc import Callable +from collections.abc import Iterator +from typing_extensions import ParamSpec +from collections import defaultdict +from torch.utils._python_dispatch import TorchDispatchMode +from math import prod +from functools import wraps +import warnings + +__all__ = ["FlopCounterMode", "register_flop_formula"] + +_T = TypeVar("_T") +_P = ParamSpec("_P") + +log = logging.getLogger(__name__) + + +try: + from triton.runtime.jit import JITFunction as _JITFunction +except ImportError: + if any(getattr(torch.version, attr, None) is not None for attr in ["cuda", "hip", "xpu"]): + log.warning("triton not found; flop counting will not work for triton kernels") + _JITFunction = NoneType + + +aten = torch.ops.aten + +def get_shape(i): + if isinstance(i, torch.Tensor): + return i.shape + return i + +flop_registry: dict[Any, Any] = {} + +def shape_wrapper(f): + @wraps(f) + def nf(*args, out_val=None, **kwargs): + args, kwargs, out_shape = tree_map(get_shape, (args, kwargs, out_val)) + return f(*args, out_shape=out_shape, **kwargs) + return nf + +def register_flop_formula(targets, get_raw=False) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]: + + def register_fun(flop_formula: Callable[_P, _T]) -> Callable[_P, _T]: + if not get_raw: + flop_formula = shape_wrapper(flop_formula) + + def register(target) -> None: + if not (isinstance(target, (torch._ops.OpOverloadPacket, _JITFunction))): + raise ValueError( + f"register_flop_formula(targets): expected each target to be " + f"OpOverloadPacket (i.e. torch.ops.mylib.foo), or JitFunction" + f", got {target} which is of type {type(target)}") + if target in flop_registry: + raise RuntimeError(f"duplicate registrations for {target}") + flop_registry[target] = flop_formula + + # To handle allowing multiple aten_ops at once + torch.utils._pytree.tree_map_(register, targets) + + return flop_formula + + return register_fun + +@register_flop_formula(aten.mm) +def mm_flop(a_shape, b_shape, *args, out_shape=None, **kwargs) -> int: + """Count flops for matmul.""" + # Inputs should be a list of length 2. + # Inputs contains the shapes of two matrices. + m, k = a_shape + k2, n = b_shape + if k != k2: + raise AssertionError(f"matmul: inner dimensions must match (k == k2), got {k} and {k2}") + # NB(chilli): Should be 2 * k - 1 technically for FLOPs. + return m * n * 2 * k + +@register_flop_formula(aten.addmm) +def addmm_flop(self_shape, a_shape, b_shape, out_shape=None, **kwargs) -> int: + """Count flops for addmm.""" + return mm_flop(a_shape, b_shape) + +@register_flop_formula(aten.bmm) +def bmm_flop(a_shape, b_shape, out_shape=None, **kwargs) -> int: + """Count flops for the bmm operation.""" + # Inputs should be a list of length 2. + # Inputs contains the shapes of two tensor. + b, m, k = a_shape + b2, k2, n = b_shape + if b != b2: + raise AssertionError(f"bmm: batch dimensions must match (b == b2), got {b} and {b2}") + if k != k2: + raise AssertionError(f"bmm: inner dimensions must match (k == k2), got {k} and {k2}") + # NB(chilli): Should be 2 * k - 1 technically for FLOPs. + flop = b * m * n * 2 * k + return flop + +@register_flop_formula(aten.baddbmm) +def baddbmm_flop(self_shape, a_shape, b_shape, out_shape=None, **kwargs) -> int: + """Count flops for the baddbmm operation.""" + # Inputs should be a list of length 3. + # Inputs contains the shapes of three tensors. + return bmm_flop(a_shape, b_shape) + +@register_flop_formula(aten._scaled_mm) +def _scaled_mm_flop( + a_shape, + b_shape, + scale_a_shape, + scale_b_shape, + bias_shape=None, + scale_result_shape=None, + out_dtype=None, + use_fast_accum=False, + out_shape=None, + **kwargs, +) -> int: + """Count flops for _scaled_mm.""" + return mm_flop(a_shape, b_shape) + + +def conv_flop_count( + x_shape: list[int], + w_shape: list[int], + out_shape: list[int], + transposed: bool = False, +) -> int: + """Count flops for convolution. + + Note only multiplication is + counted. Computation for bias are ignored. + Flops for a transposed convolution are calculated as + flops = (x_shape[2:] * prod(w_shape) * batch_size). + Args: + x_shape (list(int)): The input shape before convolution. + w_shape (list(int)): The filter shape. + out_shape (list(int)): The output shape after convolution. + transposed (bool): is the convolution transposed + Returns: + int: the number of flops + """ + batch_size = x_shape[0] + conv_shape = (x_shape if transposed else out_shape)[2:] + c_out, c_in, *filter_size = w_shape + + """ + General idea here is that for a regular conv, for each point in the output + spatial dimension we convolve the filter with something (hence + `prod(conv_shape) * prod(filter_size)` ops). Then, this gets multiplied by + 1. batch_size, 2. the cross product of input and weight channels. + + For the transpose, it's not each point in the *output* spatial dimension but + each point in the *input* spatial dimension. + """ + # NB(chilli): I don't think this properly accounts for padding :think: + # NB(chilli): Should be 2 * c_in - 1 technically for FLOPs. + flop = prod(conv_shape) * prod(filter_size) * batch_size * c_out * c_in * 2 + return flop + +@register_flop_formula([aten.convolution, + aten._convolution, + aten.cudnn_convolution, + aten._slow_conv2d_forward, + aten.convolution_overrideable]) +def conv_flop(x_shape, w_shape, _bias, _stride, _padding, _dilation, transposed, *args, out_shape=None, **kwargs) -> int: + """Count flops for convolution.""" + # pyrefly: ignore [bad-argument-type] + return conv_flop_count(x_shape, w_shape, out_shape, transposed=transposed) + + +@register_flop_formula(aten.convolution_backward) +def conv_backward_flop( + grad_out_shape, + x_shape, + w_shape, + _bias, + _stride, + _padding, + _dilation, + transposed, + _output_padding, + _groups, + output_mask, + out_shape) -> int: + + def t(shape): + return [shape[1], shape[0]] + list(shape[2:]) + flop_count = 0 + + """ + Let's say we have a regular 1D conv + {A, B, C} [inp] + {i, j} [weight] + => (conv) + {Ai + Bj, Bi + Cj} [out] + + And as a reminder, the transposed conv of the above is + => {Ai, Aj + Bi, Bj + Ci, Cj} [transposed conv out] + + For the backwards of conv, we now have + {D, E} [grad_out] + {A, B, C} [inp] + {i, j} [weight] + + # grad_inp as conv_transpose(grad_out, weight) + Let's first compute grad_inp. To do so, we can simply look at all the + multiplications that each element of inp is involved in. For example, A is + only involved in the first element of the output (and thus only depends upon + D in grad_out), and C is only involved in the last element of the output + (and thus only depends upon E in grad_out) + + {Di, Dj + Ei, Ej} [grad_inp] + + Note that this corresponds to the below conv_transpose. This gives us the + output_mask[0] branch, which is grad_inp. + + {D, E} [inp (grad_out)] + {i, j} [weight] + => (conv_transpose) + {Di, Dj + Ei, Ej} [out (grad_inp)] + + I leave the fact that grad_inp for a transposed conv is just conv(grad_out, + weight) as an exercise for the reader. + + # grad_weight as conv(inp, grad_out) + To compute grad_weight, we again look at the terms in the output, which as + a reminder is: + => {Ai + Bj, Bi + Cj} [out] + => {D, E} [grad_out] + If we manually compute the gradient for the weights, we see it's + {AD + BE, BD + CE} [grad_weight] + + This corresponds to the below conv + {A, B, C} [inp] + {D, E} [weight (grad_out)] + => (conv) + {AD + BE, BD + CE} [out (grad_weight)] + + # grad_weight of transposed conv as conv(grad_out, inp) + As a reminder, the terms of the output of a transposed conv are: + => {Ai, Aj + Bi, Bj + Ci, Cj} [transposed conv out] + => {D, E, F, G} [grad_out] + + Manually computing the gradient for the weights, we see it's + {AD + BE + CF, AE + BF + CG} [grad_weight] + + This corresponds to the below conv + {D, E, F, G} [inp (grad_out)] + {A, B, C} [weight (inp)] + => (conv) + {AD + BE + CF, AE + BF + CG} [out (grad_weight)] + + For the full backwards formula, there are also some details involving + transpose of the batch/channel dimensions and groups, but I skip those for + the sake of brevity (and they're pretty similar to matmul backwards) + + Check [conv backwards decomposition as conv forwards] + """ + # grad_inp as conv_transpose(grad_out, weight) + if output_mask[0]: + grad_input_shape = get_shape(out_shape[0]) + flop_count += conv_flop_count(grad_out_shape, w_shape, grad_input_shape, not transposed) + + if output_mask[1]: + grad_weight_shape = get_shape(out_shape[1]) + if transposed: + # grad_weight of transposed conv as conv(grad_out, inp) + flop_count += conv_flop_count(t(grad_out_shape), t(x_shape), t(grad_weight_shape), transposed=False) + else: + # grad_weight as conv(inp, grad_out) + flop_count += conv_flop_count(t(x_shape), t(grad_out_shape), t(grad_weight_shape), transposed=False) + + return flop_count + +def sdpa_flop_count(query_shape, key_shape, value_shape): + """ + Count flops for self-attention. + + Supports GQA (grouped-query attention) where key/value have fewer heads + than the query. The kernel broadcasts KV heads to match query heads. + """ + b, h_q, s_q, d_q = query_shape + _b2, h_kv, s_k, _d2 = key_shape + _b3, _h3, _s3, d_v = value_shape + if not (b == _b2 == _b3 and h_kv == _h3 and d_q == _d2 and s_k == _s3): + raise AssertionError( + f"sdpa_flop_count: query/key/value shapes are incompatible: " + f"q={query_shape}, k={key_shape}, v={value_shape}" + ) + if h_q < h_kv or h_q % h_kv != 0: + raise AssertionError( + f"sdpa_flop_count: query heads ({h_q}) must be a multiple of " + f"key/value heads ({h_kv})" + ) + total_flops = 0 + # q: [b, h_q, s_q, d_q] @ k: [b, h_q, d_q, s_k] -> scores: [b, h_q, s_q, s_k] + total_flops += bmm_flop((b * h_q, s_q, d_q), (b * h_q, d_q, s_k)) + # scores: [b, h_q, s_q, s_k] @ v: [b, h_q, s_k, d_v] -> out: [b, h_q, s_q, d_v] + total_flops += bmm_flop((b * h_q, s_q, s_k), (b * h_q, s_k, d_v)) + return total_flops + + +@register_flop_formula([aten._scaled_dot_product_efficient_attention, + aten._scaled_dot_product_flash_attention, + aten._scaled_dot_product_cudnn_attention]) +def sdpa_flop(query_shape, key_shape, value_shape, *args, out_shape=None, **kwargs) -> int: + """Count flops for self-attention.""" + # NB: We aren't accounting for causal attention here + return sdpa_flop_count(query_shape, key_shape, value_shape) + + +def _offsets_to_lengths(offsets, max_len): + """ + If the offsets tensor is fake, then we don't know the actual lengths. + In that case, we can just assume the worst case; each batch has max length. + """ + from torch._subclasses.fake_tensor import FakeTensor + from torch._subclasses.functional_tensor import FunctionalTensor + if not isinstance(offsets, (FakeTensor, FunctionalTensor)) and offsets.device.type != "meta": + return offsets.diff().tolist() + return [max_len] * (offsets.size(0) - 1) + + +def _unpack_flash_attention_nested_shapes( + *, + query, + key, + value, + grad_out=None, + cum_seq_q, + cum_seq_k, + max_q, + max_k, +) -> Iterator[tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...] | None]]: + """ + Given inputs to a flash_attention_(forward|backward) kernel, this will handle behavior for + NestedTensor inputs by effectively unbinding the NestedTensor and yielding the shapes for + each batch element. + + In the case that this isn't a NestedTensor kernel, then it just yields the original shapes. + """ + if cum_seq_q is not None: + # This means we should be dealing with a Nested Jagged Tensor query. + # The inputs will have shape (sum(sequence len), heads, dimension) + # In comparison, non-Nested inputs have shape (batch, heads, sequence len, dimension) + # To deal with this, we convert to a shape of (batch, heads, max_seq_len, dimension) + # So the flops calculation in this case is an overestimate of the actual flops. + if len(key.shape) != 3: + raise AssertionError("sdpa_flop_count: expected key.shape to be 3-dimensional") + if len(value.shape) != 3: + raise AssertionError("sdpa_flop_count: expected value.shape to be 3-dimensional") + if grad_out is not None and grad_out.shape != query.shape: + raise AssertionError("sdpa_flop_count: grad_out.shape must match query.shape when provided") + _, h_q, d_q = query.shape + _, h_k, d_k = key.shape + _, h_v, d_v = value.shape + if cum_seq_q is None: + raise AssertionError("sdpa_flop_count: cum_seq_q must not be None") + if cum_seq_k is None: + raise AssertionError("sdpa_flop_count: cum_seq_k must not be None") + if cum_seq_q.shape != cum_seq_k.shape: + raise AssertionError("sdpa_flop_count: cum_seq_q and cum_seq_k must have the same shape") + seq_q_lengths = _offsets_to_lengths(cum_seq_q, max_q) + seq_k_lengths = _offsets_to_lengths(cum_seq_k, max_k) + for (seq_q_len, seq_k_len) in zip(seq_q_lengths, seq_k_lengths, strict=True): + new_query_shape = (1, h_q, seq_q_len, d_q) + new_key_shape = (1, h_k, seq_k_len, d_k) + new_value_shape = (1, h_v, seq_k_len, d_v) + new_grad_out_shape = new_query_shape if grad_out is not None else None + yield new_query_shape, new_key_shape, new_value_shape, new_grad_out_shape + return + + yield query.shape, key.shape, value.shape, grad_out.shape if grad_out is not None else None + + +def _unpack_efficient_attention_nested_shapes( + *, + query, + key, + value, + grad_out=None, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, +) -> Iterator[tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...] | None]]: + """ + Given inputs to a efficient_attention_(forward|backward) kernel, this will handle behavior for + NestedTensor inputs by effectively unbinding the NestedTensor and yielding the shapes for + each batch element. + + In the case that this isn't a NestedTensor kernel, then it just yields the original shapes. + """ + if cu_seqlens_q is not None: + # Unlike flash_attention_forward, we get a 4D tensor instead of a 3D tensor for efficient attention. + # + # This means we should be dealing with a Nested Jagged Tensor query. + # The inputs will have shape (sum(sequence len), heads, dimension) + # In comparison, non-Nested inputs have shape (batch, heads, sequence len, dimension) + # To deal with this, we convert to a shape of (batch, heads, max_seq_len, dimension) + # So the flops calculation in this case is an overestimate of the actual flops. + if len(key.shape) != 4: + raise AssertionError("_unpack_efficient_attention_nested_shapes: expected key.shape to be 4-dimensional") + if len(value.shape) != 4: + raise AssertionError("_unpack_efficient_attention_nested_shapes: expected value.shape to be 4-dimensional") + if grad_out is not None and grad_out.shape != query.shape: + raise AssertionError("_unpack_efficient_attention_nested_shapes: grad_out.shape must match query.shape when provided") + _, _, h_q, d_q = query.shape + _, _, h_k, d_k = key.shape + _, _, h_v, d_v = value.shape + if cu_seqlens_q is None: + raise AssertionError("_unpack_efficient_attention_nested_shapes: cu_seqlens_q must not be None") + if cu_seqlens_k is None: + raise AssertionError("_unpack_efficient_attention_nested_shapes: cu_seqlens_k must not be None") + if cu_seqlens_q.shape != cu_seqlens_k.shape: + raise AssertionError("_unpack_efficient_attention_nested_shapes: " + "cu_seqlens_q and cu_seqlens_k must have the same shape") + seqlens_q = _offsets_to_lengths(cu_seqlens_q, max_seqlen_q) + seqlens_k = _offsets_to_lengths(cu_seqlens_k, max_seqlen_k) + for len_q, len_k in zip(seqlens_q, seqlens_k, strict=True): + new_query_shape = (1, h_q, len_q, d_q) + new_key_shape = (1, h_k, len_k, d_k) + new_value_shape = (1, h_v, len_k, d_v) + new_grad_out_shape = new_query_shape if grad_out is not None else None + yield new_query_shape, new_key_shape, new_value_shape, new_grad_out_shape + return + + yield query.shape, key.shape, value.shape, grad_out.shape if grad_out is not None else None + + +@register_flop_formula(aten._flash_attention_forward, get_raw=True) +def _flash_attention_forward_flop( + query, + key, + value, + cum_seq_q, + cum_seq_k, + max_q, + max_k, + *args, + out_shape=None, + **kwargs +) -> int: + """Count flops for self-attention.""" + # NB: We aren't accounting for causal attention here + # in case this is a nested tensor, we unpack the individual batch elements + # and then sum the flops per batch element + sizes = _unpack_flash_attention_nested_shapes( + query=query, + key=key, + value=value, + cum_seq_q=cum_seq_q, + cum_seq_k=cum_seq_k, + max_q=max_q, + max_k=max_k, + ) + return sum( + sdpa_flop_count(query_shape, key_shape, value_shape) + for query_shape, key_shape, value_shape, _ in sizes + ) + + +@register_flop_formula(aten._efficient_attention_forward, get_raw=True) +def _efficient_attention_forward_flop( + query, + key, + value, + bias, + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + *args, + **kwargs +) -> int: + """Count flops for self-attention.""" + # NB: We aren't accounting for causal attention here + # in case this is a nested tensor, we unpack the individual batch elements + # and then sum the flops per batch element + sizes = _unpack_efficient_attention_nested_shapes( + query=query, + key=key, + value=value, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + ) + return sum( + sdpa_flop_count(query_shape, key_shape, value_shape) + for query_shape, key_shape, value_shape, _ in sizes + ) + + +def sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape): + b, h_q, s_q, d_q = query_shape + _b2, h_kv, s_k, _d2 = key_shape + _b3, _h3, _s3, d_v = value_shape + _b4, _h4, _s4, _d4 = grad_out_shape + if not (b == _b2 == _b3 == _b4 and h_kv == _h3 and h_q == _h4): + raise AssertionError( + "sdpa_backward_flop_count: batch/heads mismatch among tensors" + ) + if h_q < h_kv or h_q % h_kv != 0: + raise AssertionError( + f"sdpa_backward_flop_count: query heads ({h_q}) must be a multiple of " + f"key/value heads ({h_kv})" + ) + if not (d_q == _d2 and d_v == _d4 and s_k == _s3 and s_q == _s4): + raise AssertionError( + "sdpa_backward_flop_count: grad_out/value/key/query shapes are incompatible" + ) + total_flops = 0 + # Step 1: We recompute the scores matrix. + # q: [b, h_q, s_q, d_q] @ k: [b, h_q, d_q, s_k] -> scores: [b, h_q, s_q, s_k] + total_flops += bmm_flop((b * h_q, s_q, d_q), (b * h_q, d_q, s_k)) + + # Step 2: We propagate the gradients through the score @ v operation. + # gradOut: [b, h_q, s_q, d_v] @ v: [b, h_q, d_v, s_k] -> gradScores: [b, h_q, s_q, s_k] + total_flops += bmm_flop((b * h_q, s_q, d_v), (b * h_q, d_v, s_k)) + # scores: [b, h_q, s_k, s_q] @ gradOut: [b, h_q, s_q, d_v] -> gradV: [b, h_q, s_k, d_v] + total_flops += bmm_flop((b * h_q, s_k, s_q), (b * h_q, s_q, d_v)) + + # Step 3: We propagate th gradients through the k @ v operation + # gradScores: [b, h_q, s_q, s_k] @ k: [b, h_q, s_k, d_q] -> gradQ: [b, h_q, s_q, d_q] + total_flops += bmm_flop((b * h_q, s_q, s_k), (b * h_q, s_k, d_q)) + # q: [b, h_q, d_q, s_q] @ gradScores: [b, h_q, s_q, s_k] -> gradK: [b, h_q, d_q, s_k] + total_flops += bmm_flop((b * h_q, d_q, s_q), (b * h_q, s_q, s_k)) + return total_flops + + +@register_flop_formula([aten._scaled_dot_product_efficient_attention_backward, + aten._scaled_dot_product_flash_attention_backward, + aten._scaled_dot_product_cudnn_attention_backward]) +def sdpa_backward_flop(grad_out_shape, query_shape, key_shape, value_shape, *args, out_shape=None, **kwargs) -> int: + """Count flops for self-attention backward.""" + return sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape) + +@register_flop_formula(aten._flash_attention_backward, get_raw=True) +def _flash_attention_backward_flop( + grad_out, + query, + key, + value, + out, # named _out_shape to avoid kwarg collision with out_shape created in wrapper + logsumexp, + cum_seq_q, + cum_seq_k, + max_q, + max_k, + *args, + **kwargs, +) -> int: + # in case this is a nested tensor, we unpack the individual batch elements + # and then sum the flops per batch element + shapes = _unpack_flash_attention_nested_shapes( + query=query, + key=key, + value=value, + grad_out=grad_out, + cum_seq_q=cum_seq_q, + cum_seq_k=cum_seq_k, + max_q=max_q, + max_k=max_k, + ) + return sum( + sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape) + for query_shape, key_shape, value_shape, grad_out_shape in shapes + ) + + +@register_flop_formula(aten._efficient_attention_backward, get_raw=True) +def _efficient_attention_backward_flop( + grad_out, + query, + key, + value, + bias, + out, # named _out to avoid kwarg collision with out created in wrapper + cu_seqlens_q, + cu_seqlens_k, + max_seqlen_q, + max_seqlen_k, + *args, + **kwargs, +) -> int: + # in case this is a nested tensor, we unpack the individual batch elements + # and then sum the flops per batch element + shapes = _unpack_efficient_attention_nested_shapes( + query=query, + key=key, + value=value, + grad_out=grad_out, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + ) + return sum( + sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape) + for query_shape, key_shape, value_shape, grad_out_shape in shapes + ) + + +def _varlen_attn_forward_flop( + query, + key, + value, + cu_seq_q, + cu_seq_k, + max_q, + max_k, + *args, + out_val=None, + **kwargs, +) -> int: + """Count flops for varlen_attn forward.""" + sizes = _unpack_flash_attention_nested_shapes( + query=query, + key=key, + value=value, + cum_seq_q=cu_seq_q, + cum_seq_k=cu_seq_k if cu_seq_k is not None else cu_seq_q, + max_q=max_q, + max_k=max_k, + ) + return sum( + sdpa_flop_count(query_shape, key_shape, value_shape) + for query_shape, key_shape, value_shape, _ in sizes + ) + + +def _varlen_attn_out_flop( + out, + query, + key, + value, + cu_seq_q, + cu_seq_k, + max_q, + max_k, + *args, + out_val=None, + **kwargs, +) -> int: + """Count flops for varlen_attn_out forward.""" + return _varlen_attn_forward_flop( + query, key, value, cu_seq_q, cu_seq_k, max_q, max_k, + ) + + +def _varlen_attn_backward_flop( + grad_out, + query, + key, + value, + out, + lse, + cu_seq_q, + cu_seq_k, + max_q, + max_k, + *args, + out_val=None, + **kwargs, +) -> int: + """Count flops for varlen_attn backward.""" + sizes = _unpack_flash_attention_nested_shapes( + query=query, + key=key, + value=value, + grad_out=grad_out, + cum_seq_q=cu_seq_q, + cum_seq_k=cu_seq_k, + max_q=max_q, + max_k=max_k, + ) + return sum( + sdpa_backward_flop_count(grad_out_shape, query_shape, key_shape, value_shape) + for query_shape, key_shape, value_shape, grad_out_shape in sizes + ) + + +flop_registry = { + aten.mm: mm_flop, + aten.addmm: addmm_flop, + aten.bmm: bmm_flop, + aten.baddbmm: baddbmm_flop, + aten._scaled_mm: _scaled_mm_flop, + aten.convolution: conv_flop, + aten._convolution: conv_flop, + aten.cudnn_convolution: conv_flop, + aten.convolution_overrideable: conv_flop, + aten._slow_conv2d_forward: conv_flop, + aten.convolution_backward: conv_backward_flop, + aten._scaled_dot_product_efficient_attention: sdpa_flop, + aten._scaled_dot_product_flash_attention: sdpa_flop, + aten._scaled_dot_product_cudnn_attention: sdpa_flop, + aten._scaled_dot_product_efficient_attention_backward: sdpa_backward_flop, + aten._scaled_dot_product_flash_attention_backward: sdpa_backward_flop, + aten._scaled_dot_product_cudnn_attention_backward: sdpa_backward_flop, + aten._flash_attention_forward: _flash_attention_forward_flop, + aten._efficient_attention_forward: _efficient_attention_forward_flop, + aten._flash_attention_backward: _flash_attention_backward_flop, + aten._efficient_attention_backward: _efficient_attention_backward_flop, +} + +def normalize_tuple(x): + if not isinstance(x, tuple): + return (x,) + return x + + +# Define the suffixes for different orders of magnitude +suffixes = ["", "K", "M", "B", "T"] +# Thanks BingChat! +def get_suffix_str(number): + # Find the index of the appropriate suffix based on the number of digits + # with some additional overflow. + # i.e. 1.01B should be displayed as 1001M, not 1.001B + index = max(0, min(len(suffixes) - 1, (len(str(number)) - 2) // 3)) + return suffixes[index] + +def convert_num_with_suffix(number, suffix): + index = suffixes.index(suffix) + # Divide the number by 1000^index and format it to two decimal places + value = f"{number / 1000 ** index:.3f}" + # Return the value and the suffix as a string + return value + suffixes[index] + +def convert_to_percent_str(num, denom) -> str: + if denom == 0: + return "0%" + return f"{num / denom:.2%}" + +def _pytreeify_preserve_structure(f): + @wraps(f) + def nf(args): + flat_args, spec = tree_flatten(args) + out = f(*flat_args) + return tree_unflatten(out, spec) + + return nf + + +class FlopCounterMode: + """ + ``FlopCounterMode`` is a context manager that counts the number of flops within its context. + + It does this using a ``TorchDispatchMode``. + + It also supports hierarchical output by passing a module (or list of + modules) to FlopCounterMode on construction. If you do not need hierarchical + output, you do not need to use it with a module. + + Example usage + + .. code-block:: python + + mod = ... + with FlopCounterMode(mod) as flop_counter: + mod.sum().backward() + + """ + + def __init__( + self, + mods: torch.nn.Module | list[torch.nn.Module] | None = None, + depth: int = 2, + display: bool = True, + custom_mapping: dict[Any, Any] | None = None) -> None: + super().__init__() + self.flop_counts: dict[str, dict[Any, int]] = defaultdict(lambda: defaultdict(int)) + self.depth = depth + self.display = display + self.mode: _FlopCounterMode | None = None + if custom_mapping is None: + custom_mapping = {} + if mods is not None: + warnings.warn("mods argument is not needed anymore, you can stop passing it", stacklevel=2) + self.flop_registry = { + **flop_registry, + **{k: v if getattr(v, "_get_raw", False) else shape_wrapper(v) for k, v in custom_mapping.items()} + } + self.mod_tracker = ModuleTracker() + + def get_total_flops(self) -> int: + return sum(self.flop_counts['Global'].values()) + + def get_flop_counts(self) -> dict[str, dict[Any, int]]: + """Return the flop counts as a dictionary of dictionaries. + + The outer + dictionary is keyed by module name, and the inner dictionary is keyed by + operation name. + + Returns: + Dict[str, Dict[Any, int]]: The flop counts as a dictionary. + """ + return {k: dict(v) for k, v in self.flop_counts.items()} + + def get_table(self, depth=None): + if depth is None: + depth = self.depth + if depth is None: + depth = 999999 + + + import tabulate + + tabulate.PRESERVE_WHITESPACE = True + header = ["Module", "FLOP", "% Total"] + values = [] + global_flops = self.get_total_flops() + global_suffix = get_suffix_str(global_flops) + is_global_subsumed = False + + def process_mod(mod_name, depth): + nonlocal is_global_subsumed + + total_flops = sum(self.flop_counts[mod_name].values()) + + is_global_subsumed |= total_flops >= global_flops + + padding = " " * depth + values = [] + values.append([ + padding + mod_name, + convert_num_with_suffix(total_flops, global_suffix), + convert_to_percent_str(total_flops, global_flops) + ]) + for k, v in self.flop_counts[mod_name].items(): + values.append([ + padding + " - " + str(k), + convert_num_with_suffix(v, global_suffix), + convert_to_percent_str(v, global_flops) + ]) + return values + + for mod in sorted(self.flop_counts.keys()): + if mod == 'Global': + continue + mod_depth = mod.count(".") + 1 + if mod_depth > depth: + continue + + cur_values = process_mod(mod, mod_depth - 1) + values.extend(cur_values) + + # We do a bit of messing around here to only output the "Global" value + # if there are any FLOPs in there that aren't already fully contained by + # a module. + if 'Global' in self.flop_counts and not is_global_subsumed: + for value in values: + value[0] = " " + value[0] + + values = process_mod('Global', 0) + values + + if len(values) == 0: + values = [["Global", "0", "0%"]] + + return tabulate.tabulate(values, headers=header, colalign=("left", "right", "right")) + + # NB: This context manager is NOT reentrant + def __enter__(self): + self.flop_counts.clear() + self.mod_tracker.__enter__() + self.mode = _FlopCounterMode(self) + self.mode.__enter__() + return self + + def __exit__(self, *args): + if self.mode is None: + raise AssertionError("Internal error: FlopCounter.__exit__ called but mode is None") + b = self.mode.__exit__(*args) + self.mode = None # break cycles + self.mod_tracker.__exit__() + if self.display: + print(self.get_table(self.depth)) + return b + + def _count_flops(self, func_packet, out, args, kwargs): + if func_packet in self.flop_registry: + flop_count_func = self.flop_registry[func_packet] + flop_count = flop_count_func(*args, **kwargs, out_val=out) # type: ignore[operator] + for par in set(self.mod_tracker.parents): + self.flop_counts[par][func_packet] += flop_count + return out + +class _FlopCounterMode(TorchDispatchMode): + supports_higher_order_operators = True + + def __init__(self, counter: FlopCounterMode) -> None: + self.counter = counter + + def _execute_with_isolated_flop_counting(self, branch_fn, operands): + """Execute a branch function and capture its FLOP counts without + affecting self.counter.flop_counts + + Args: + branch_fn: The branch function to execute + operands: Arguments to pass to the branch function + + Returns: + Tuple of (result, flop_counts) where result is the branch output + and flop_counts is a copy of the FLOP counts after execution + """ + import copy + checkpointed_flop_counts = copy.copy(self.counter.flop_counts) + with self: + result = branch_fn(*operands) + flop_counts = copy.copy(self.counter.flop_counts) + self.counter.flop_counts = checkpointed_flop_counts + return result, flop_counts + + def _handle_higher_order_ops(self, func, types, args, kwargs): + is_triton = func in {torch.ops.higher_order.triton_kernel_wrapper_mutation, + torch.ops.higher_order.triton_kernel_wrapper_functional} + if is_triton: + from torch._higher_order_ops.triton_kernel_wrap import get_kernel + # Special case - look in the triton flop registry for the kernel + from triton.runtime.jit import JITFunction + kernel_name = get_kernel(kwargs["kernel_idx"]) + # Unwrap heuristics if they are present + while not isinstance(kernel_name, JITFunction): + if hasattr(kernel_name, "fn"): + kernel_name = kernel_name.fn + else: + break + return self.counter._count_flops(kernel_name, None, args, kwargs) + elif func is torch.ops.higher_order.cond: + # The flop counter for cond counts the upper bound of flops. + # For example, if a matmul is executed 2 times in true branch + # but only 1 time in the false branch, the flop counter will + # record the larger number of flops, i.e. 2 times. + pred, true_branch, false_branch, operands = args + # Step 1: Count flops for true branch and false branch separately + true_out, true_flop_counts = self._execute_with_isolated_flop_counting( + true_branch, operands + ) + if true_out is NotImplemented: + return NotImplemented + + false_out, false_flop_counts = self._execute_with_isolated_flop_counting( + false_branch, operands + ) + if false_out is NotImplemented: + return NotImplemented + + # Step 2: merge flop counts + all_mod_keys = set(true_flop_counts.keys()) | set(false_flop_counts.keys()) + merged_flop_counts = {} + for outer_key in all_mod_keys: + true_func_counts = true_flop_counts[outer_key] + false_func_counts = false_flop_counts[outer_key] + + merged_func_counts = {} + all_func_keys = set(true_func_counts.keys()) | set(false_func_counts.keys()) + + for func_key in all_func_keys: + true_val = true_func_counts.get(func_key, 0) + false_val = false_func_counts.get(func_key, 0) + merged_func_counts[func_key] = max(true_val, false_val) + + merged_flop_counts[outer_key] = merged_func_counts + + # Step 3: update the counter with merged counts + for outer_key, inner_dict in merged_flop_counts.items(): + self.counter.flop_counts[outer_key].update(inner_dict) + + # It doesn't matter which one we return since true_fn and false_fn return + # output with the same structure. + return true_out + else: + return NotImplemented + + def __torch_dispatch__(self, func, types, args=(), kwargs=None): + kwargs = kwargs if kwargs else {} + + # Skip ops from non-standard dispatch_sizes_strides_policy such as NJT + if func in {torch.ops.aten.sym_is_contiguous.default, + torch.ops.aten.is_contiguous.default, + torch.ops.aten.is_contiguous.memory_format, + torch.ops.aten.is_strides_like_format.default, + torch.ops.aten.is_non_overlapping_and_dense.default, + torch.ops.aten.size.default, + torch.ops.aten.sym_size.default, + torch.ops.aten.stride.default, + torch.ops.aten.sym_stride.default, + torch.ops.aten.storage_offset.default, + torch.ops.aten.sym_storage_offset.default, + torch.ops.aten.numel.default, + torch.ops.aten.sym_numel.default, + torch.ops.aten.dim.default, + torch.ops.prim.layout.default}: + + return NotImplemented + + if isinstance(func, torch._ops.HigherOrderOperator): + return self._handle_higher_order_ops(func, types, args, kwargs) + + # If we don't have func in flop_registry, see if it can decompose + if func not in self.counter.flop_registry and func is not torch.ops.prim.device.default: + with self: + r = func.decompose(*args, **kwargs) + if r is not NotImplemented: + return r + + # no further decomposition; execute & count flops + out = func(*args, **kwargs) + return self.counter._count_flops(func._overloadpacket, out, args, kwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..58f3ace6c03d093337c9fa417ccbe8bc267b6c69 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/__init__.py @@ -0,0 +1 @@ +from .version import __version__ diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/constants.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/constants.py new file mode 100644 index 0000000000000000000000000000000000000000..c9e1bc4dc410cce97c6bc0b55391387fd04c349c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/constants.py @@ -0,0 +1,66 @@ +"""Constants for annotations in the mapping. + +The constants defined here are used to annotate the mapping tuples in cuda_to_hip_mappings.py. +They are based on +https://github.com/ROCm/HIPIFY/blob/master/src/Statistics.h +and fall in three categories: 1) type of mapping, 2) API of mapping, 3) unsupported +mapping. +""" + +import warnings +warnings.warn("hipify's constants.py is no longer used as of version 2.0.0", FutureWarning) + +CONV_VERSION = 0, +CONV_INIT = 1 +CONV_DEVICE = 2 +CONV_MEM = 3 +CONV_KERN = 4 +CONV_COORD_FUNC = 5 +CONV_MATH_FUNC = 6 +CONV_DEVICE_FUNC = 7 +CONV_SPECIAL_FUNC = 8 +CONV_STREAM = 9 +CONV_EVENT = 10 +CONV_OCCUPANCY = 11 +CONV_CONTEXT = 12 +CONV_PEER = 13 +CONV_MODULE = 14 +CONV_CACHE = 15 +CONV_EXEC = 16 +CONV_ERROR = 17 +CONV_DEF = 18 +CONV_TEX = 19 +CONV_GL = 20 +CONV_GRAPHICS = 21 +CONV_SURFACE = 22 +CONV_JIT = 23 +CONV_D3D9 = 24 +CONV_D3D10 = 25 +CONV_D3D11 = 26 +CONV_VDPAU = 27 +CONV_EGL = 28 +CONV_THREAD = 29 +CONV_OTHER = 30 +CONV_INCLUDE = 31 +CONV_INCLUDE_CUDA_MAIN_H = 32 +CONV_TYPE = 33 +CONV_LITERAL = 34 +CONV_NUMERIC_LITERAL = 35 +CONV_LAST = 36 + +API_DRIVER = 37 +API_RUNTIME = 38 +API_BLAS = 39 +API_SPECIAL = 40 +API_RAND = 41 +API_LAST = 42 +API_FFT = 43 +API_RTC = 44 +API_ROCTX = 45 +API_PYT_EXT = 46 + +HIP_UNSUPPORTED = 47 +API_PYTORCH = 1337 +API_CAFFE2 = 1338 +API_C10 = 1339 +API_ROCMSMI = 1340 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/cuda_to_hip_mappings.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/cuda_to_hip_mappings.py new file mode 100644 index 0000000000000000000000000000000000000000..9332d0a24f8ed3382b271beb355079bc07152eb5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/cuda_to_hip_mappings.py @@ -0,0 +1,3497 @@ +import collections +import os + +""" Mapping of CUDA functions, include files, constants, and types to ROCm/HIP equivalents """ + +_IS_FBCODE = os.environ.get("IS_FBCODE", "0") == "1" + +# FBCODE compiles against rccl sources instead of an installed rccl package. +# The header location is src/rccl.h versus rccl/rccl.h, respectively. +_RCCL_HEADER = "" if _IS_FBCODE else "" + +# List of math functions that should be replaced inside device code only. +MATH_TRANSPILATIONS = collections.OrderedDict([ + ("std::max", ("::max")), + ("std::min", ("::min")), + ("std::ceil", ("::ceil")), + ("std::floor", ("::floor")), + ("std::exp", ("::exp")), + ("std::log", ("::log")), + ("std::pow", ("::pow")), + ("std::fabs", ("::fabs")), + ("std::fmod", ("::fmod")), + ("std::remainder", ("::remainder")), + ("std::frexp", ("::frexp")), +]) + +CUDA_TYPE_NAME_MAP = collections.OrderedDict([ + ("CUresult", "hipError_t"), + ("cudaError_t", "hipError_t"), + ("cudaError", "hipError_t"), + ("CUDA_ARRAY3D_DESCRIPTOR", "HIP_ARRAY3D_DESCRIPTOR"), + ("CUDA_ARRAY_DESCRIPTOR", "HIP_ARRAY_DESCRIPTOR"), + ("CUDA_MEMCPY2D", "hip_Memcpy2D"), + ("CUDA_MEMCPY3D", "HIP_MEMCPY3D"), + ("CUDA_MEMCPY3D_PEER", "HIP_MEMCPY3D_PEER"), + ("CUDA_POINTER_ATTRIBUTE_P2P_TOKENS", "HIP_POINTER_ATTRIBUTE_P2P_TOKENS"), + ("CUDA_RESOURCE_DESC", "HIP_RESOURCE_DESC"), + ("CUDA_RESOURCE_VIEW_DESC", "HIP_RESOURCE_VIEW_DESC"), + ("CUipcEventHandle", "hipIpcEventHandle"), + ("CUipcMemHandle", "hipIpcMemHandle"), + ("CUaddress_mode", "hipAddress_mode"), + ("CUarray_cubemap_face", "hipArray_cubemap_face"), + ("CUarray_format", "hipArray_format"), + ("CUcomputemode", "hipComputemode"), + ("CUmem_advise", "hipMemAdvise"), + ("CUmem_range_attribute", "hipMemRangeAttribute"), + ("CUctx_flags", "hipCctx_flags"), + ("CUdevice", "hipDevice_t"), + ("CUdevice_attribute_enum", "hipDeviceAttribute_t"), + ("CUdevice_attribute", "hipDeviceAttribute_t"), + ("CUpointer_attribute", "hipPointer_attribute"), + ("CU_POINTER_ATTRIBUTE_DEVICE_ORDINAL", "HIP_POINTER_ATTRIBUTE_DEVICE_ORDINAL"), + ("CU_POINTER_ATTRIBUTE_BUFFER_ID", "HIP_POINTER_ATTRIBUTE_BUFFER_ID"), + ("CUdeviceptr", "hipDeviceptr_t"), + ("CUarray_st", "hipArray"), + ("CUarray", "hipArray *"), + ("CUdevprop_st", "hipDeviceProp_t"), + ("CUdevprop", "hipDeviceProp_t"), + ("CUfunction", "hipFunction_t"), + ("CUgraphicsResource", "hipGraphicsResource_t"), + ("CUmipmappedArray", "hipMipmappedArray_t"), + ("CUfunction_attribute", "hipFuncAttribute_t"), + ("CUfunction_attribute_enum", "hipFuncAttribute_t"), + ("CUgraphicsMapResourceFlags", "hipGraphicsMapFlags"), + ("CUgraphicsMapResourceFlags_enum", "hipGraphicsMapFlags"), + ("CUgraphicsRegisterFlags", "hipGraphicsRegisterFlags"), + ("CUgraphicsRegisterFlags_enum", "hipGraphicsRegisterFlags"), + ("CUoccupancy_flags", "hipOccupancyFlags"), + ("CUoccupancy_flags_enum", "hipOccupancyFlags"), + ("CUfunc_cache_enum", "hipFuncCache"), + ("CUfunc_cache", "hipFuncCache"), + ("CUipcMem_flags", "hipIpcMemFlags"), + ("CUipcMem_flags_enum", "hipIpcMemFlags"), + ("CUjit_cacheMode", "hipJitCacheMode"), + ("CUjit_cacheMode_enum", "hipJitCacheMode"), + ("CUjit_fallback", "hipJitFallback"), + ("CUjit_fallback_enum", "hipJitFallback"), + ("CUjit_option", "hipJitOption"), + ("CUjit_option_enum", "hipJitOption"), + ("CUjit_target", "hipJitTarget"), + ("CUjit_target_enum", "hipJitTarget"), + ("CUjitInputType", "hipJitInputType"), + ("CUjitInputType_enum", "hipJitInputType"), + ("CUlimit", "hipLimit_t"), + ("CUlimit_enum", "hipLimit_t"), + ("CUmemAccessDesc", "hipMemAccessDesc"), + ("CUmemAccessDesc_st", "hipMemAccessDesc"), + ("CUmemAccessDesc_v1", "hipMemAccessDesc"), + ("CUmemAttach_flags", "hipMemAttachFlags_t"), + ("CUmemAttach_flags_enum", "hipMemAttachFlags_t"), + ("CUmemAllocationGranularity_flags", "hipMemAllocationGranularity_flags"), + ("CUmemAllocationGranularity_flags_enum", "hipMemAllocationGranularity_flags"), + ("CUmemAllocationHandleType", "hipMemAllocationHandleType"), + ("CUmemAllocationHandleType_enum", "hipMemAllocationHandleType"), + ("CUmemAllocationProp", "hipMemAllocationProp"), + ("CUmemAllocationProp_st", "hipMemAllocationProp"), + ("CUmemAllocationProp_v1", "hipMemAllocationProp"), + ("CUmemAllocationType", "hipMemAllocationType"), + ("CUmemAllocationType_enum", "hipMemAllocationType"), + ("CUmemGenericAllocationHandle", "hipMemGenericAllocationHandle_t"), + ("CUmemGenericAllocationHandle_v1", "hipMemGenericAllocationHandle_t"), + ("CUmemHandleType", "hipMemHandleType"), + ("CUmemHandleType_enum", "hipMemHandleType"), + ("CUmemLocation", "hipMemLocation"), + ("CUmemLocationType", "hipMemLocationType"), + ("CUmemLocationType_enum", "hipMemLocationType"), + ("CUmemLocation_st", "hipMemLocation"), + ("CUmemLocation_v1", "hipMemLocation"), + ("CUmemOperationType", "hipMemOperationType"), + ("CUmemOperationType_enum", "hipMemOperationType"), + ("CUmemPoolHandle_st", "ihipMemPoolHandle_t"), + ("CUmemPoolProps", "hipMemPoolProps"), + ("CUmemPoolProps_st", "hipMemPoolProps"), + ("CUmemPoolProps_v1", "hipMemPoolProps"), + ("CUmemPoolPtrExportData", "hipMemPoolPtrExportData"), + ("CUmemPoolPtrExportData_st", "hipMemPoolPtrExportData"), + ("CUmemPoolPtrExportData_v1", "hipMemPoolPtrExportData"), + ("CUmemPool_attribute", "hipMemPoolAttr"), + ("CUmemPool_attribute_enum", "hipMemPoolAttr"), + ("CUmem_advise_enum", "hipMemoryAdvise"), + ("CUmem_range_attribute_enum", "hipMemRangeAttribute"), + ("CUmemoryPool", "hipMemPool_t"), + ("CUmemorytype", "hipMemType_t"), + ("CUmemorytype_enum", "hipMemType_t"), + ("CUresourcetype", "hipResourceType"), + ("CUresourcetype_enum", "hipResourceType"), + ("CUresourceViewFormat", "hipResourceViewFormat"), + ("CUresourceViewFormat_enum", "hipResourceViewFormat"), + ("CUsharedconfig", "hipSharedMemConfig"), + ("CUsharedconfig_enum", "hipSharedMemConfig"), + ("CUcontext", "hipCtx_t"), + ("CUmodule", "hipModule_t"), + ("CUstream", "hipStream_t"), + ("CUstream_st", "ihipStream_t"), + ("CUstreamCallback", "hipStreamCallback_t"), + ("CUsurfObject", "hipSurfaceObject"), + ("CUsurfref", "hipSurfaceReference_t"), + ("CUtexObject", "hipTextureObject_t"), + ("CUtexref", "textureReference"), + ("CUstream_flags", "hipStreamFlags"), + ("CUstreamWaitValue_flags", "hipStreamWaitValueFlags"), + ("CUstreamWriteValue_flags", "hipStreamWriteValueFlags"), + ("CUstreamBatchMemOpType", "hipStreamBatchMemOpType"), + ("CUdevice_P2PAttribute", "hipDeviceP2PAttribute"), + ("CUevent", "hipEvent_t"), + ("CUevent_st", "ihipEvent_t"), + ("CUevent_flags", "hipEventFlags"), + ("CUfilter_mode", "hipTextureFilterMode"), + ("CUGLDeviceList", "hipGLDeviceList"), + ("CUGLmap_flags", "hipGLMapFlags"), + ("CUd3d9DeviceList", "hipD3D9DeviceList"), + ("CUd3d9map_flags", "hipD3D9MapFlags"), + ("CUd3d9register_flags", "hipD3D9RegisterFlags"), + ("CUd3d10DeviceList", "hipd3d10DeviceList"), + ("CUd3d10map_flags", "hipD3D10MapFlags"), + ("CUd3d10register_flags", "hipD3D10RegisterFlags"), + ("CUd3d11DeviceList", "hipd3d11DeviceList"), + ("CUeglStreamConnection_st", "hipEglStreamConnection"), + ("CUeglStreamConnection", "hipEglStreamConnection"), + ("libraryPropertyType_t", "hipLibraryPropertyType_t"), + ("libraryPropertyType", "hipLibraryPropertyType_t"), + ("cudaStreamCallback_t", "hipStreamCallback_t"), + ("cudaArray", "hipArray"), + ("cudaArray_t", "hipArray_t"), + ("cudaArray_const_t", "hipArray_const_t"), + ("cudaMipmappedArray_t", "hipMipmappedArray_t"), + ("cudaMipmappedArray_const_t", "hipMipmappedArray_const_t"), + ("cudaArrayDefault", "hipArrayDefault"), + ("cudaArrayLayered", "hipArrayLayered"), + ("cudaArraySurfaceLoadStore", "hipArraySurfaceLoadStore"), + ("cudaArrayCubemap", "hipArrayCubemap"), + ("cudaArrayTextureGather", "hipArrayTextureGather"), + ("cudaMemoryAdvise", "hipMemoryAdvise"), + ("cudaMemRangeAttribute", "hipMemRangeAttribute"), + ("cudaMemcpyKind", "hipMemcpyKind"), + ("cudaMemoryType", "hipMemoryType"), + ("cudaExtent", "hipExtent"), + ("cudaPitchedPtr", "hipPitchedPtr"), + ("cudaPos", "hipPos"), + ("cudaEvent_t", "hipEvent_t"), + ("cudaStream_t", "hipStream_t"), + ("cudaPointerAttributes", "hipPointerAttribute_t"), + ("cudaDeviceAttr", "hipDeviceAttribute_t"), + ("cudaDeviceProp", "hipDeviceProp_t"), + ("cudaDeviceP2PAttr", "hipDeviceP2PAttribute"), + ("cudaComputeMode", "hipComputeMode"), + ("cudaFuncCache", "hipFuncCache_t"), + ("cudaFuncAttributes", "hipFuncAttributes"), + ("cudaSharedMemConfig", "hipSharedMemConfig"), + ("cudaLimit", "hipLimit_t"), + ("cudaOutputMode", "hipOutputMode"), + ("cudaTextureReadMode", "hipTextureReadMode"), + ("cudaTextureFilterMode", "hipTextureFilterMode"), + ("cudaChannelFormatKind", "hipChannelFormatKind"), + ("cudaChannelFormatDesc", "hipChannelFormatDesc"), + ("cudaResourceDesc", "hipResourceDesc"), + ("cudaResourceViewDesc", "hipResourceViewDesc"), + ("cudaTextureDesc", "hipTextureDesc"), + ("surfaceReference", "hipSurfaceReference"), + ("cudaTextureObject_t", "hipTextureObject_t"), + ("cudaResourceType", "hipResourceType"), + ("cudaResourceViewFormat", "hipResourceViewFormat"), + ("cudaTextureAddressMode", "hipTextureAddressMode"), + ("cudaSurfaceBoundaryMode", "hipSurfaceBoundaryMode"), + ("cudaSurfaceFormatMode", "hipSurfaceFormatMode"), + ("cudaSurfaceObject_t", "hipSurfaceObject_t"), + ("cudaTextureType1D", "hipTextureType1D"), + ("cudaTextureType2D", "hipTextureType2D"), + ("cudaTextureType3D", "hipTextureType3D"), + ("cudaTextureTypeCubemap", "hipTextureTypeCubemap"), + ("cudaTextureType1DLayered", "hipTextureType1DLayered"), + ("cudaTextureType2DLayered", "hipTextureType2DLayered"), + ("cudaTextureTypeCubemapLayered", "hipTextureTypeCubemapLayered"), + ("cudaUUID_t", "hipUUID"), + ("cudaIpcEventHandle_t", "hipIpcEventHandle_t"), + ("cudaIpcEventHandle_st", "hipIpcEventHandle_t"), + ("cudaIpcMemHandle_t", "hipIpcMemHandle_t"), + ("cudaIpcMemHandle_st", "hipIpcMemHandle_t"), + ("cudaGraphicsCubeFace", "hipGraphicsCubeFace"), + ("cudaGraphicsMapFlags", "hipGraphicsMapFlags"), + ("cudaGraphicsRegisterFlags", "hipGraphicsRegisterFlags"), + ("cudaGLDeviceList", "hipGLDeviceList"), + ("cudaGLMapFlags", "hipGLMapFlags"), + ("cudaD3D9DeviceList", "hipD3D9DeviceList"), + ("cudaD3D9MapFlags", "hipD3D9MapFlags"), + ("cudaD3D9RegisterFlags", "hipD3D9RegisterFlags"), + ("cudaD3D10DeviceList", "hipd3d10DeviceList"), + ("cudaD3D10MapFlags", "hipD3D10MapFlags"), + ("cudaD3D10RegisterFlags", "hipD3D10RegisterFlags"), + ("cudaD3D11DeviceList", "hipd3d11DeviceList"), + ("cudaEglStreamConnection", "hipEglStreamConnection"), + ("cublasHandle_t", "hipblasHandle_t"), + ("cublasOperation_t", "hipblasOperation_t"), + ("cublasStatus_t", "hipblasStatus_t"), + ("cublasFillMode_t", "hipblasFillMode_t"), + ("cublasDiagType_t", "hipblasDiagType_t"), + ("cublasSideMode_t", "hipblasSideMode_t"), + ("cublasPointerMode_t", "hipblasPointerMode_t"), + ("cublasGemmAlgo_t", "hipblasGemmAlgo_t"), + ("cublasAtomicsMode_t", "hipblasAtomicsMode_t"), + ("cublasDataType_t", "hipblasDatatype_t"), + ("curandStatus", "hiprandStatus_t"), + ("curandStatus_t", "hiprandStatus_t"), + ("curandRngType", "hiprandRngType_t"), + ("curandRngType_t", "hiprandRngType_t"), + ("curandGenerator_st", "hiprandGenerator_st"), + ("curandGenerator_t", "hiprandGenerator_t"), + ("curandDirectionVectorSet", "hiprandDirectionVectorSet_t"), + ("curandDirectionVectorSet_t", "hiprandDirectionVectorSet_t"), + ("curandOrdering", "hiprandOrdering_t"), + ("curandOrdering_t", "hiprandOrdering_t"), + ("curandDistribution_st", "hiprandDistribution_st"), + ("curandHistogramM2V_st", "hiprandDistribution_st"), + ("curandDistribution_t", "hiprandDistribution_t"), + ("curandHistogramM2V_t", "hiprandDistribution_t"), + ("curandDistributionShift_st", "hiprandDistributionShift_st"), + ("curandDistributionShift_t", "hiprandDistributionShift_t"), + ("curandDistributionM2Shift_st", "hiprandDistributionM2Shift_st"), + ("curandDistributionM2Shift_t", "hiprandDistributionM2Shift_t"), + ("curandHistogramM2_st", "hiprandHistogramM2_st"), + ("curandHistogramM2_t", "hiprandHistogramM2_t"), + ("curandHistogramM2K_st", "hiprandHistogramM2K_st"), + ("curandHistogramM2K_t", "hiprandHistogramM2K_t"), + ("curandDiscreteDistribution_st", "hiprandDiscreteDistribution_st"), + ("curandDiscreteDistribution_t", "hiprandDiscreteDistribution_t"), + ("curandMethod", "hiprandMethod_t"), + ("curandMethod_t", "hiprandMethod_t"), + ("curandDirectionVectors32_t", "hiprandDirectionVectors32_t"), + ("curandDirectionVectors64_t", "hiprandDirectionVectors64_t"), + ("curandStateMtgp32_t", "hiprandStateMtgp32_t"), + ("curandStateMtgp32", "hiprandStateMtgp32_t"), + ("curandStateScrambledSobol64_t", "hiprandStateScrambledSobol64_t"), + ("curandStateSobol64_t", "hiprandStateSobol64_t"), + ("curandStateScrambledSobol32_t", "hiprandStateScrambledSobol32_t"), + ("curandStateSobol32_t", "hiprandStateSobol32_t"), + ("curandStateMRG32k3a_t", "hiprandStateMRG32k3a_t"), + ("curandStatePhilox4_32_10_t", "hiprandStatePhilox4_32_10_t"), + ("curandStateXORWOW_t", "hiprandStateXORWOW_t"), + ("curandState_t", "hiprandState_t"), + ("curandState", "hiprandState_t"), + ("CUuuid", "hipUUID"), + ("cudaGraph_t", "hipGraph_t"), + ("cudaGraphExec_t", "hipGraphExec_t"), + ("__nv_bfloat16", "__hip_bfloat16"), + ("__nv_bfloat162", "__hip_bfloat162"), +]) + +CUDA_INCLUDE_MAP = collections.OrderedDict([ + ("include ", _RCCL_HEADER), + ("nvrtc.h", "hip/hiprtc.h"), + ("thrust/system/cuda", "thrust/system/hip"), + ("cub/util_allocator.cuh", "hipcub/hipcub.hpp"), + ("cub/block/block_reduce.cuh", "hipcub/hipcub.hpp"), + ("cub/block/block_raking_layout.cuh", "hipcub/hipcub.hpp"), + ("cub/cub.cuh", "hipcub/hipcub.hpp"), + ("cub/config.cuh", "hipcub/hipcub.hpp"), + ("cub/util_ptx.cuh", "hipcub/hipcub.hpp"), + ("cub/util_type.cuh", "hipcub/hipcub.hpp"), + ("cub/device/device_run_length_encode.cuh", "hipcub/hipcub.hpp"), + ("cub/block/block_load.cuh", "hipcub/hipcub.hpp"), + ("cub/block/block_store.cuh", "hipcub/hipcub.hpp"), + ("cub/block/block_scan.cuh", "hipcub/hipcub.hpp"), + ("cub/device/device_radix_sort.cuh", "hipcub/hipcub.hpp"), + ("cub/device/device_reduce.cuh", "hipcub/hipcub.hpp"), + ("cub/device/device_scan.cuh", "hipcub/hipcub.hpp"), + ("cub/device/device_select.cuh", "hipcub/hipcub.hpp"), + ("nvtx3/nvtx3.hpp", "roctracer/roctx.h"), + ("nvToolsExt.h", "roctracer/roctx.h"), + ("nvml.h", "rocm_smi/rocm_smi.h"), + ("tensorpipe/tensorpipe_cuda.h", "tensorpipe/tensorpipe_hip.h"), +]) + +CUDA_IDENTIFIER_MAP = collections.OrderedDict([ + ("__CUDACC__", "__HIPCC__"), + ("CUDA_ERROR_INVALID_CONTEXT", "hipErrorInvalidContext"), + ("CUDA_ERROR_CONTEXT_ALREADY_CURRENT", "hipErrorContextAlreadyCurrent"), + ("CUDA_ERROR_ARRAY_IS_MAPPED", "hipErrorArrayIsMapped"), + ("CUDA_ERROR_ALREADY_MAPPED", "hipErrorAlreadyMapped"), + ("CUDA_ERROR_ALREADY_ACQUIRED", "hipErrorAlreadyAcquired"), + ("CUDA_ERROR_NOT_MAPPED", "hipErrorNotMapped"), + ("CUDA_ERROR_NOT_MAPPED_AS_ARRAY", "hipErrorNotMappedAsArray"), + ("CUDA_ERROR_NOT_MAPPED_AS_POINTER", "hipErrorNotMappedAsPointer"), + ("CUDA_ERROR_CONTEXT_ALREADY_IN_USE", "hipErrorContextAlreadyInUse"), + ("CUDA_ERROR_INVALID_SOURCE", "hipErrorInvalidSource"), + ("CUDA_ERROR_FILE_NOT_FOUND", "hipErrorFileNotFound"), + ("CUDA_ERROR_NOT_FOUND", "hipErrorNotFound"), + ("CUDA_ERROR_LAUNCH_INCOMPATIBLE_TEXTURING", "hipErrorLaunchIncompatibleTexturing"), + ("CUDA_ERROR_PRIMARY_CONTEXT_ACTIVE", "hipErrorPrimaryContextActive"), + ("CUDA_ERROR_CONTEXT_IS_DESTROYED", "hipErrorContextIsDestroyed"), + ("CUDA_ERROR_NOT_PERMITTED", "hipErrorNotPermitted"), + ("CUDA_ERROR_NOT_SUPPORTED", "hipErrorNotSupported"), + ("cudaErrorMissingConfiguration", "hipErrorMissingConfiguration"), + ("cudaErrorPriorLaunchFailure", "hipErrorPriorLaunchFailure"), + ("cudaErrorInvalidDeviceFunction", "hipErrorInvalidDeviceFunction"), + ("cudaErrorInvalidConfiguration", "hipErrorInvalidConfiguration"), + ("cudaErrorInvalidPitchValue", "hipErrorInvalidPitchValue"), + ("cudaErrorInvalidSymbol", "hipErrorInvalidSymbol"), + ("cudaErrorInvalidHostPointer", "hipErrorInvalidHostPointer"), + ("cudaErrorInvalidDevicePointer", "hipErrorInvalidDevicePointer"), + ("cudaErrorInvalidTexture", "hipErrorInvalidTexture"), + ("cudaErrorInvalidTextureBinding", "hipErrorInvalidTextureBinding"), + ("cudaErrorInvalidChannelDescriptor", "hipErrorInvalidChannelDescriptor"), + ("cudaErrorInvalidMemcpyDirection", "hipErrorInvalidMemcpyDirection"), + ("cudaErrorAddressOfConstant", "hipErrorAddressOfConstant"), + ("cudaErrorTextureFetchFailed", "hipErrorTextureFetchFailed"), + ("cudaErrorTextureNotBound", "hipErrorTextureNotBound"), + ("cudaErrorSynchronizationError", "hipErrorSynchronizationError"), + ("cudaErrorInvalidFilterSetting", "hipErrorInvalidFilterSetting"), + ("cudaErrorInvalidNormSetting", "hipErrorInvalidNormSetting"), + ("cudaErrorMixedDeviceExecution", "hipErrorMixedDeviceExecution"), + ("cudaErrorNotYetImplemented", "hipErrorNotYetImplemented"), + ("cudaErrorMemoryValueTooLarge", "hipErrorMemoryValueTooLarge"), + ("cudaErrorInsufficientDriver", "hipErrorInsufficientDriver"), + ("cudaErrorSetOnActiveProcess", "hipErrorSetOnActiveProcess"), + ("cudaErrorInvalidSurface", "hipErrorInvalidSurface"), + ("cudaErrorDuplicateVariableName", "hipErrorDuplicateVariableName"), + ("cudaErrorDuplicateTextureName", "hipErrorDuplicateTextureName"), + ("cudaErrorDuplicateSurfaceName", "hipErrorDuplicateSurfaceName"), + ("cudaErrorDevicesUnavailable", "hipErrorDevicesUnavailable"), + ("cudaErrorIncompatibleDriverContext", "hipErrorIncompatibleDriverContext"), + ("cudaErrorDeviceAlreadyInUse", "hipErrorDeviceAlreadyInUse"), + ("cudaErrorLaunchMaxDepthExceeded", "hipErrorLaunchMaxDepthExceeded"), + ("cudaErrorLaunchFileScopedTex", "hipErrorLaunchFileScopedTex"), + ("cudaErrorLaunchFileScopedSurf", "hipErrorLaunchFileScopedSurf"), + ("cudaErrorSyncDepthExceeded", "hipErrorSyncDepthExceeded"), + ("cudaErrorLaunchPendingCountExceeded", "hipErrorLaunchPendingCountExceeded"), + ("cudaErrorNotPermitted", "hipErrorNotPermitted"), + ("cudaErrorNotSupported", "hipErrorNotSupported"), + ("cudaErrorStartupFailure", "hipErrorStartupFailure"), + ("cudaErrorApiFailureBase", "hipErrorApiFailureBase"), + ("CUDA_SUCCESS", "hipSuccess"), + ("cudaSuccess", "hipSuccess"), + ("CUDA_ERROR_INVALID_VALUE", "hipErrorInvalidValue"), + ("cudaErrorInvalidValue", "hipErrorInvalidValue"), + ("CUDA_ERROR_OUT_OF_MEMORY", "hipErrorMemoryAllocation"), + ("cudaErrorMemoryAllocation", "hipErrorMemoryAllocation"), + ("CUDA_ERROR_NOT_INITIALIZED", "hipErrorNotInitialized"), + ("cudaErrorInitializationError", "hipErrorInitializationError"), + ("CUDA_ERROR_DEINITIALIZED", "hipErrorDeinitialized"), + ("cudaErrorCudartUnloading", "hipErrorDeinitialized"), + ("CUDA_ERROR_PROFILER_DISABLED", "hipErrorProfilerDisabled"), + ("cudaErrorProfilerDisabled", "hipErrorProfilerDisabled"), + ("CUDA_ERROR_PROFILER_NOT_INITIALIZED", "hipErrorProfilerNotInitialized"), + ("cudaErrorProfilerNotInitialized", "hipErrorProfilerNotInitialized"), + ("CUDA_ERROR_PROFILER_ALREADY_STARTED", "hipErrorProfilerAlreadyStarted"), + ("cudaErrorProfilerAlreadyStarted", "hipErrorProfilerAlreadyStarted"), + ("CUDA_ERROR_PROFILER_ALREADY_STOPPED", "hipErrorProfilerAlreadyStopped"), + ("cudaErrorProfilerAlreadyStopped", "hipErrorProfilerAlreadyStopped"), + ("CUDA_ERROR_NO_DEVICE", "hipErrorNoDevice"), + ("cudaErrorNoDevice", "hipErrorNoDevice"), + ("CUDA_ERROR_INVALID_DEVICE", "hipErrorInvalidDevice"), + ("cudaErrorInvalidDevice", "hipErrorInvalidDevice"), + ("CUDA_ERROR_INVALID_IMAGE", "hipErrorInvalidImage"), + ("cudaErrorInvalidKernelImage", "hipErrorInvalidImage"), + ("CUDA_ERROR_MAP_FAILED", "hipErrorMapFailed"), + ("cudaErrorMapBufferObjectFailed", "hipErrorMapFailed"), + ("CUDA_ERROR_UNMAP_FAILED", "hipErrorUnmapFailed"), + ("cudaErrorUnmapBufferObjectFailed", "hipErrorUnmapFailed"), + ("CUDA_ERROR_NO_BINARY_FOR_GPU", "hipErrorNoBinaryForGpu"), + ("cudaErrorNoKernelImageForDevice", "hipErrorNoBinaryForGpu"), + ("CUDA_ERROR_ECC_UNCORRECTABLE", "hipErrorECCNotCorrectable"), + ("cudaErrorECCUncorrectable", "hipErrorECCNotCorrectable"), + ("CUDA_ERROR_UNSUPPORTED_LIMIT", "hipErrorUnsupportedLimit"), + ("cudaErrorUnsupportedLimit", "hipErrorUnsupportedLimit"), + ("CUDA_ERROR_PEER_ACCESS_UNSUPPORTED", "hipErrorPeerAccessUnsupported"), + ("cudaErrorPeerAccessUnsupported", "hipErrorPeerAccessUnsupported"), + ("CUDA_ERROR_INVALID_PTX", "hipErrorInvalidKernelFile"), + ("cudaErrorInvalidPtx", "hipErrorInvalidKernelFile"), + ("CUDA_ERROR_INVALID_GRAPHICS_CONTEXT", "hipErrorInvalidGraphicsContext"), + ("cudaErrorInvalidGraphicsContext", "hipErrorInvalidGraphicsContext"), + ("CUDA_ERROR_NVLINK_UNCORRECTABLE", "hipErrorNvlinkUncorrectable"), + ("cudaErrorNvlinkUncorrectable", "hipErrorNvlinkUncorrectable"), + ("CUDA_ERROR_SHARED_OBJECT_SYMBOL_NOT_FOUND", "hipErrorSharedObjectSymbolNotFound"), + ("cudaErrorSharedObjectSymbolNotFound", "hipErrorSharedObjectSymbolNotFound"), + ("CUDA_ERROR_SHARED_OBJECT_INIT_FAILED", "hipErrorSharedObjectInitFailed"), + ("cudaErrorSharedObjectInitFailed", "hipErrorSharedObjectInitFailed"), + ("CUDA_ERROR_OPERATING_SYSTEM", "hipErrorOperatingSystem"), + ("cudaErrorOperatingSystem", "hipErrorOperatingSystem"), + ("CUDA_ERROR_INVALID_HANDLE", "hipErrorInvalidResourceHandle"), + ("cudaErrorInvalidResourceHandle", "hipErrorInvalidResourceHandle"), + ("CUDA_ERROR_NOT_READY", "hipErrorNotReady"), + ("cudaErrorNotReady", "hipErrorNotReady"), + ("CUDA_ERROR_ILLEGAL_ADDRESS", "hipErrorIllegalAddress"), + ("cudaErrorIllegalAddress", "hipErrorIllegalAddress"), + ("CUDA_ERROR_LAUNCH_OUT_OF_RESOURCES", "hipErrorLaunchOutOfResources"), + ("cudaErrorLaunchOutOfResources", "hipErrorLaunchOutOfResources"), + ("CUDA_ERROR_LAUNCH_TIMEOUT", "hipErrorLaunchTimeOut"), + ("cudaErrorLaunchTimeout", "hipErrorLaunchTimeOut"), + ("CUDA_ERROR_PEER_ACCESS_ALREADY_ENABLED", "hipErrorPeerAccessAlreadyEnabled"), + ("cudaErrorPeerAccessAlreadyEnabled", "hipErrorPeerAccessAlreadyEnabled"), + ("CUDA_ERROR_PEER_ACCESS_NOT_ENABLED", "hipErrorPeerAccessNotEnabled"), + ("cudaErrorPeerAccessNotEnabled", "hipErrorPeerAccessNotEnabled"), + ("CUDA_ERROR_ASSERT", "hipErrorAssert"), + ("cudaErrorAssert", "hipErrorAssert"), + ("CUDA_ERROR_TOO_MANY_PEERS", "hipErrorTooManyPeers"), + ("cudaErrorTooManyPeers", "hipErrorTooManyPeers"), + ("CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED", "hipErrorHostMemoryAlreadyRegistered"), + ("cudaErrorHostMemoryAlreadyRegistered", "hipErrorHostMemoryAlreadyRegistered"), + ("CUDA_ERROR_HOST_MEMORY_NOT_REGISTERED", "hipErrorHostMemoryNotRegistered"), + ("cudaErrorHostMemoryNotRegistered", "hipErrorHostMemoryNotRegistered"), + ("CUDA_ERROR_HARDWARE_STACK_ERROR", "hipErrorHardwareStackError"), + ("cudaErrorHardwareStackError", "hipErrorHardwareStackError"), + ("CUDA_ERROR_ILLEGAL_INSTRUCTION", "hipErrorIllegalInstruction"), + ("cudaErrorIllegalInstruction", "hipErrorIllegalInstruction"), + ("CUDA_ERROR_MISALIGNED_ADDRESS", "hipErrorMisalignedAddress"), + ("cudaErrorMisalignedAddress", "hipErrorMisalignedAddress"), + ("CUDA_ERROR_INVALID_ADDRESS_SPACE", "hipErrorInvalidAddressSpace"), + ("cudaErrorInvalidAddressSpace", "hipErrorInvalidAddressSpace"), + ("CUDA_ERROR_INVALID_PC", "hipErrorInvalidPc"), + ("cudaErrorInvalidPc", "hipErrorInvalidPc"), + ("CUDA_ERROR_LAUNCH_FAILED", "hipErrorLaunchFailure"), + ("cudaErrorLaunchFailure", "hipErrorLaunchFailure"), + ("CUDA_ERROR_UNKNOWN", "hipErrorUnknown"), + ("cudaErrorUnknown", "hipErrorUnknown"), + ("CU_TR_ADDRESS_MODE_WRAP", "HIP_TR_ADDRESS_MODE_WRAP"), + ("CU_TR_ADDRESS_MODE_CLAMP", "HIP_TR_ADDRESS_MODE_CLAMP"), + ("CU_TR_ADDRESS_MODE_MIRROR", "HIP_TR_ADDRESS_MODE_MIRROR"), + ("CU_TR_ADDRESS_MODE_BORDER", "HIP_TR_ADDRESS_MODE_BORDER"), + ("CU_CUBEMAP_FACE_POSITIVE_X", "HIP_CUBEMAP_FACE_POSITIVE_X"), + ("CU_CUBEMAP_FACE_NEGATIVE_X", "HIP_CUBEMAP_FACE_NEGATIVE_X"), + ("CU_CUBEMAP_FACE_POSITIVE_Y", "HIP_CUBEMAP_FACE_POSITIVE_Y"), + ("CU_CUBEMAP_FACE_NEGATIVE_Y", "HIP_CUBEMAP_FACE_NEGATIVE_Y"), + ("CU_CUBEMAP_FACE_POSITIVE_Z", "HIP_CUBEMAP_FACE_POSITIVE_Z"), + ("CU_CUBEMAP_FACE_NEGATIVE_Z", "HIP_CUBEMAP_FACE_NEGATIVE_Z"), + ("CU_AD_FORMAT_UNSIGNED_INT8", "HIP_AD_FORMAT_UNSIGNED_INT8"), + ("CU_AD_FORMAT_UNSIGNED_INT16", "HIP_AD_FORMAT_UNSIGNED_INT16"), + ("CU_AD_FORMAT_UNSIGNED_INT32", "HIP_AD_FORMAT_UNSIGNED_INT32"), + ("CU_AD_FORMAT_SIGNED_INT8", "HIP_AD_FORMAT_SIGNED_INT8"), + ("CU_AD_FORMAT_SIGNED_INT16", "HIP_AD_FORMAT_SIGNED_INT16"), + ("CU_AD_FORMAT_SIGNED_INT32", "HIP_AD_FORMAT_SIGNED_INT32"), + ("CU_AD_FORMAT_HALF", "HIP_AD_FORMAT_HALF"), + ("CU_AD_FORMAT_FLOAT", "HIP_AD_FORMAT_FLOAT"), + ("CU_COMPUTEMODE_DEFAULT", "hipComputeModeDefault"), + ("CU_COMPUTEMODE_EXCLUSIVE", "hipComputeModeExclusive"), + ("CU_COMPUTEMODE_PROHIBITED", "hipComputeModeProhibited"), + ("CU_COMPUTEMODE_EXCLUSIVE_PROCESS", "hipComputeModeExclusiveProcess"), + ("CU_MEM_ADVISE_SET_READ_MOSTLY", "hipMemAdviseSetReadMostly"), + ("CU_MEM_ADVISE_UNSET_READ_MOSTLY", "hipMemAdviseUnsetReadMostly"), + ("CU_MEM_ADVISE_SET_PREFERRED_LOCATION", "hipMemAdviseSetPreferredLocation"), + ("CU_MEM_ADVISE_UNSET_PREFERRED_LOCATION", "hipMemAdviseUnsetPreferredLocation"), + ("CU_MEM_ADVISE_SET_ACCESSED_BY", "hipMemAdviseSetAccessedBy"), + ("CU_MEM_ADVISE_UNSET_ACCESSED_BY", "hipMemAdviseUnsetAccessedBy"), + ("CU_MEM_RANGE_ATTRIBUTE_READ_MOSTLY", "hipMemRangeAttributeReadMostly"), + ("CU_MEM_RANGE_ATTRIBUTE_PREFERRED_LOCATION", "hipMemRangeAttributePreferredLocation"), + ("CU_MEM_RANGE_ATTRIBUTE_ACCESSED_BY", "hipMemRangeAttributeAccessedBy"), + ("CU_MEM_RANGE_ATTRIBUTE_LAST_PREFETCH_LOCATION", "hipMemRangeAttributeLastPrefetchLocation"), + ("CU_CTX_SCHED_AUTO", "HIP_CTX_SCHED_AUTO"), + ("CU_CTX_SCHED_SPIN", "HIP_CTX_SCHED_SPIN"), + ("CU_CTX_SCHED_YIELD", "HIP_CTX_SCHED_YIELD"), + ("CU_CTX_SCHED_BLOCKING_SYNC", "HIP_CTX_SCHED_BLOCKING_SYNC"), + ("CU_CTX_BLOCKING_SYNC", "HIP_CTX_BLOCKING_SYNC"), + ("CU_CTX_SCHED_MASK", "HIP_CTX_SCHED_MASK"), + ("CU_CTX_MAP_HOST", "HIP_CTX_MAP_HOST"), + ("CU_CTX_LMEM_RESIZE_TO_MAX", "HIP_CTX_LMEM_RESIZE_TO_MAX"), + ("CU_CTX_FLAGS_MASK", "HIP_CTX_FLAGS_MASK"), + ("CU_LAUNCH_PARAM_BUFFER_POINTER", "HIP_LAUNCH_PARAM_BUFFER_POINTER"), + ("CU_LAUNCH_PARAM_BUFFER_SIZE", "HIP_LAUNCH_PARAM_BUFFER_SIZE"), + ("CU_LAUNCH_PARAM_END", "HIP_LAUNCH_PARAM_END"), + ("CU_IPC_HANDLE_SIZE", "HIP_IPC_HANDLE_SIZE"), + ("CU_MEMHOSTALLOC_DEVICEMAP", "HIP_MEMHOSTALLOC_DEVICEMAP"), + ("CU_MEMHOSTALLOC_PORTABLE", "HIP_MEMHOSTALLOC_PORTABLE"), + ("CU_MEMHOSTALLOC_WRITECOMBINED", "HIP_MEMHOSTALLOC_WRITECOMBINED"), + ("CU_MEMHOSTREGISTER_DEVICEMAP", "HIP_MEMHOSTREGISTER_DEVICEMAP"), + ("CU_MEMHOSTREGISTER_IOMEMORY", "HIP_MEMHOSTREGISTER_IOMEMORY"), + ("CU_MEMHOSTREGISTER_PORTABLE", "HIP_MEMHOSTREGISTER_PORTABLE"), + ("CU_PARAM_TR_DEFAULT", "HIP_PARAM_TR_DEFAULT"), + ("CU_STREAM_LEGACY", "HIP_STREAM_LEGACY"), + ("CU_STREAM_PER_THREAD", "HIP_STREAM_PER_THREAD"), + ("CU_TRSA_OVERRIDE_FORMAT", "HIP_TRSA_OVERRIDE_FORMAT"), + ("CU_TRSF_NORMALIZED_COORDINATES", "HIP_TRSF_NORMALIZED_COORDINATES"), + ("CU_TRSF_READ_AS_INTEGER", "HIP_TRSF_READ_AS_INTEGER"), + ("CU_TRSF_SRGB", "HIP_TRSF_SRGB"), + ("CUDA_ARRAY3D_2DARRAY", "HIP_ARRAY3D_LAYERED"), + ("CUDA_ARRAY3D_CUBEMAP", "HIP_ARRAY3D_CUBEMAP"), + ("CUDA_ARRAY3D_DEPTH_TEXTURE", "HIP_ARRAY3D_DEPTH_TEXTURE"), + ("CUDA_ARRAY3D_LAYERED", "HIP_ARRAY3D_LAYERED"), + ("CUDA_ARRAY3D_SURFACE_LDST", "HIP_ARRAY3D_SURFACE_LDST"), + ("CUDA_ARRAY3D_TEXTURE_GATHER", "HIP_ARRAY3D_TEXTURE_GATHER"), + ("CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_BLOCK", "hipDeviceAttributeMaxThreadsPerBlock"), + ("CU_DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_X", "hipDeviceAttributeMaxBlockDimX"), + ("CU_DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_Y", "hipDeviceAttributeMaxBlockDimY"), + ("CU_DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_Z", "hipDeviceAttributeMaxBlockDimZ"), + ("CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X", "hipDeviceAttributeMaxGridDimX"), + ("CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_Y", "hipDeviceAttributeMaxGridDimY"), + ("CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_Z", "hipDeviceAttributeMaxGridDimZ"), + ("CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK", "hipDeviceAttributeMaxSharedMemoryPerBlock"), + ("CU_DEVICE_ATTRIBUTE_SHARED_MEMORY_PER_BLOCK", "hipDeviceAttributeMaxSharedMemoryPerBlock"), + ("CU_DEVICE_ATTRIBUTE_TOTAL_CONSTANT_MEMORY", "hipDeviceAttributeTotalConstantMemory"), + ("CU_DEVICE_ATTRIBUTE_WARP_SIZE", "hipDeviceAttributeWarpSize"), + ("CU_DEVICE_ATTRIBUTE_MAX_PITCH", "hipDeviceAttributeMaxPitch"), + ("CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_BLOCK", "hipDeviceAttributeMaxRegistersPerBlock"), + ("CU_DEVICE_ATTRIBUTE_REGISTERS_PER_BLOCK", "hipDeviceAttributeMaxRegistersPerBlock"), + ("CU_DEVICE_ATTRIBUTE_CLOCK_RATE", "hipDeviceAttributeClockRate"), + ("CU_DEVICE_ATTRIBUTE_TEXTURE_ALIGNMENT", "hipDeviceAttributeTextureAlignment"), + ("CU_DEVICE_ATTRIBUTE_GPU_OVERLAP", "hipDeviceAttributeAsyncEngineCount"), + ("CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT", "hipDeviceAttributeMultiprocessorCount"), + ("CU_DEVICE_ATTRIBUTE_KERNEL_EXEC_TIMEOUT", "hipDeviceAttributeKernelExecTimeout"), + ("CU_DEVICE_ATTRIBUTE_INTEGRATED", "hipDeviceAttributeIntegrated"), + ("CU_DEVICE_ATTRIBUTE_CAN_MAP_HOST_MEMORY", "hipDeviceAttributeCanMapHostMemory"), + ("CU_DEVICE_ATTRIBUTE_COMPUTE_MODE", "hipDeviceAttributeComputeMode"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE1D_WIDTH", "hipDeviceAttributeMaxTexture1DWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_WIDTH", "hipDeviceAttributeMaxTexture2DWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_HEIGHT", "hipDeviceAttributeMaxTexture2DHeight"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE3D_WIDTH", "hipDeviceAttributeMaxTexture3DWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE3D_HEIGHT", "hipDeviceAttributeMaxTexture3DHeight"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE3D_DEPTH", "hipDeviceAttributeMaxTexture3DDepth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_LAYERED_WIDTH", "hipDeviceAttributeMaxTexture2DLayeredWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_LAYERED_HEIGHT", "hipDeviceAttributeMaxTexture2DLayeredHeight"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_LAYERED_LAYERS", "hipDeviceAttributeMaxTexture2DLayeredLayers"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_ARRAY_WIDTH", "hipDeviceAttributeMaxTexture2DLayeredWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_ARRAY_HEIGHT", "hipDeviceAttributeMaxTexture2DLayeredHeight"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_ARRAY_NUMSLICES", "hipDeviceAttributeMaxTexture2DLayeredLayers"), + ("CU_DEVICE_ATTRIBUTE_SURFACE_ALIGNMENT", "hipDeviceAttributeSurfaceAlignment"), + ("CU_DEVICE_ATTRIBUTE_CONCURRENT_KERNELS", "hipDeviceAttributeConcurrentKernels"), + ("CU_DEVICE_ATTRIBUTE_ECC_ENABLED", "hipDeviceAttributeEccEnabled"), + ("CU_DEVICE_ATTRIBUTE_PCI_BUS_ID", "hipDeviceAttributePciBusId"), + ("CU_DEVICE_ATTRIBUTE_PCI_DEVICE_ID", "hipDeviceAttributePciDeviceId"), + ("CU_DEVICE_ATTRIBUTE_TCC_DRIVER", "hipDeviceAttributeTccDriver"), + ("CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE", "hipDeviceAttributeMemoryClockRate"), + ("CU_DEVICE_ATTRIBUTE_GLOBAL_MEMORY_BUS_WIDTH", "hipDeviceAttributeMemoryBusWidth"), + ("CU_DEVICE_ATTRIBUTE_L2_CACHE_SIZE", "hipDeviceAttributeL2CacheSize"), + ("CU_DEVICE_ATTRIBUTE_MAX_THREADS_PER_MULTIPROCESSOR", "hipDeviceAttributeMaxThreadsPerMultiProcessor"), + ("CU_DEVICE_ATTRIBUTE_ASYNC_ENGINE_COUNT", "hipDeviceAttributeAsyncEngineCount"), + ("CU_DEVICE_ATTRIBUTE_UNIFIED_ADDRESSING", "hipDeviceAttributeUnifiedAddressing"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE1D_LAYERED_WIDTH", "hipDeviceAttributeMaxTexture1DLayeredWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE1D_LAYERED_LAYERS", "hipDeviceAttributeMaxTexture1DLayeredLayers"), + ("CU_DEVICE_ATTRIBUTE_CAN_TEX2D_GATHER", "hipDeviceAttributeCanTex2DGather"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_GATHER_WIDTH", "hipDeviceAttributeMaxTexture2DGatherWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_GATHER_HEIGHT", "hipDeviceAttributeMaxTexture2DGatherHeight"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE3D_WIDTH_ALTERNATE", "hipDeviceAttributeMaxTexture3DWidthAlternate"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE3D_HEIGHT_ALTERNATE", "hipDeviceAttributeMaxTexture3DHeightAlternate"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE3D_DEPTH_ALTERNATE", "hipDeviceAttributeMaxTexture3DDepthAlternate"), + ("CU_DEVICE_ATTRIBUTE_PCI_DOMAIN_ID", "hipDeviceAttributePciDomainId"), + ("CU_DEVICE_ATTRIBUTE_TEXTURE_PITCH_ALIGNMENT", "hipDeviceAttributeTexturePitchAlignment"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURECUBEMAP_WIDTH", "hipDeviceAttributeMaxTextureCubemapWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURECUBEMAP_LAYERED_WIDTH", "hipDeviceAttributeMaxTextureCubemapLayeredWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURECUBEMAP_LAYERED_LAYERS", "hipDeviceAttributeMaxTextureCubemapLayeredLayers"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE1D_WIDTH", "hipDeviceAttributeMaxSurface1DWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE2D_WIDTH", "hipDeviceAttributeMaxSurface2DWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE2D_HEIGHT", "hipDeviceAttributeMaxSurface2DHeight"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE3D_WIDTH", "hipDeviceAttributeMaxSurface3DWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE3D_HEIGHT", "hipDeviceAttributeMaxSurface3DHeight"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE3D_DEPTH", "hipDeviceAttributeMaxSurface3DDepth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE1D_LAYERED_WIDTH", "hipDeviceAttributeMaxSurface1DLayeredWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE1D_LAYERED_LAYERS", "hipDeviceAttributeMaxSurface1DLayeredLayers"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE2D_LAYERED_WIDTH", "hipDeviceAttributeMaxSurface2DLayeredWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE2D_LAYERED_HEIGHT", "hipDeviceAttributeMaxSurface2DLayeredHeight"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACE2D_LAYERED_LAYERS", "hipDeviceAttributeMaxSurface2DLayeredLayers"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACECUBEMAP_WIDTH", "hipDeviceAttributeMaxSurfaceCubemapWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACECUBEMAP_LAYERED_WIDTH", "hipDeviceAttributeMaxSurfaceCubemapLayeredWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_SURFACECUBEMAP_LAYERED_LAYERS", "hipDeviceAttributeMaxSurfaceCubemapLayeredLayers"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE1D_LINEAR_WIDTH", "hipDeviceAttributeMaxTexture1DLinearWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_LINEAR_WIDTH", "hipDeviceAttributeMaxTexture2DLinearWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_LINEAR_HEIGHT", "hipDeviceAttributeMaxTexture2DLinearHeight"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_LINEAR_PITCH", "hipDeviceAttributeMaxTexture2DLinearPitch"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_MIPMAPPED_WIDTH", "hipDeviceAttributeMaxTexture2DMipmappedWidth"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE2D_MIPMAPPED_HEIGHT", "hipDeviceAttributeMaxTexture2DMipmappedHeight"), + ("CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR", "hipDeviceAttributeComputeCapabilityMajor"), + ("CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR", "hipDeviceAttributeComputeCapabilityMinor"), + ("CU_DEVICE_ATTRIBUTE_MAXIMUM_TEXTURE1D_MIPMAPPED_WIDTH", "hipDeviceAttributeMaxTexture1DMipmappedWidth"), + ("CU_DEVICE_ATTRIBUTE_STREAM_PRIORITIES_SUPPORTED", "hipDeviceAttributeStreamPrioritiesSupported"), + ("CU_DEVICE_ATTRIBUTE_GLOBAL_L1_CACHE_SUPPORTED", "hipDeviceAttributeGlobalL1CacheSupported"), + ("CU_DEVICE_ATTRIBUTE_LOCAL_L1_CACHE_SUPPORTED", "hipDeviceAttributeLocalL1CacheSupported"), + ("CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_MULTIPROCESSOR", "hipDeviceAttributeMaxSharedMemoryPerMultiprocessor"), + ("CU_DEVICE_ATTRIBUTE_MAX_REGISTERS_PER_MULTIPROCESSOR", "hipDeviceAttributeMaxRegistersPerMultiprocessor"), + ("CU_DEVICE_ATTRIBUTE_MANAGED_MEMORY", "hipDeviceAttributeManagedMemory"), + ("CU_DEVICE_ATTRIBUTE_MULTI_GPU_BOARD", "hipDeviceAttributeIsMultiGpuBoard"), + ("CU_DEVICE_ATTRIBUTE_MULTI_GPU_BOARD_GROUP_ID", "hipDeviceAttributeMultiGpuBoardGroupId"), + ("CU_DEVICE_ATTRIBUTE_HOST_NATIVE_ATOMIC_SUPPORTED", "hipDeviceAttributeHostNativeAtomicSupported"), + ("CU_DEVICE_ATTRIBUTE_SINGLE_TO_DOUBLE_PRECISION_PERF_RATIO", "hipDeviceAttributeSingleToDoublePrecisionPerfRatio"), + ("CU_DEVICE_ATTRIBUTE_PAGEABLE_MEMORY_ACCESS", "hipDeviceAttributePageableMemoryAccess"), + ("CU_DEVICE_ATTRIBUTE_CONCURRENT_MANAGED_ACCESS", "hipDeviceAttributeConcurrentManagedAccess"), + ("CU_DEVICE_ATTRIBUTE_COMPUTE_PREEMPTION_SUPPORTED", "hipDeviceAttributeComputePreemptionSupported"), + ("CU_DEVICE_ATTRIBUTE_CAN_USE_HOST_POINTER_FOR_REGISTERED_MEM", "hipDeviceAttributeCanUseHostPointerForRegisteredMem"), + ("CU_DEVICE_ATTRIBUTE_MAX", "hipDeviceAttributeMax"), + ("CU_POINTER_ATTRIBUTE_CONTEXT", "hipPointerAttributeContext"), + ("CU_POINTER_ATTRIBUTE_MEMORY_TYPE", "hipPointerAttributeMemoryType"), + ("CU_POINTER_ATTRIBUTE_DEVICE_POINTER", "hipPointerAttributeDevicePointer"), + ("CU_POINTER_ATTRIBUTE_HOST_POINTER", "hipPointerAttributeHostPointer"), + ("CU_POINTER_ATTRIBUTE_P2P_TOKENS", "hipPointerAttributeP2pTokens"), + ("CU_POINTER_ATTRIBUTE_SYNC_MEMOPS", "hipPointerAttributeSyncMemops"), + ("CU_POINTER_ATTRIBUTE_BUFFER_ID", "hipPointerAttributeBufferId"), + ("CU_POINTER_ATTRIBUTE_IS_MANAGED", "hipPointerAttributeIsManaged"), + ("CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK", "hipFuncAttributeMaxThreadsPerBlocks"), + ("CU_FUNC_ATTRIBUTE_SHARED_SIZE_BYTES", "hipFuncAttributeSharedSizeBytes"), + ("CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES", "hipFuncAttributeMaxDynamicSharedMemorySize"), + ("CU_FUNC_ATTRIBUTE_CONST_SIZE_BYTES", "hipFuncAttributeConstSizeBytes"), + ("CU_FUNC_ATTRIBUTE_LOCAL_SIZE_BYTES", "hipFuncAttributeLocalSizeBytes"), + ("CU_FUNC_ATTRIBUTE_NUM_REGS", "hipFuncAttributeNumRegs"), + ("CU_FUNC_ATTRIBUTE_PTX_VERSION", "hipFuncAttributePtxVersion"), + ("CU_FUNC_ATTRIBUTE_BINARY_VERSION", "hipFuncAttributeBinaryVersion"), + ("CU_FUNC_ATTRIBUTE_CACHE_MODE_CA", "hipFuncAttributeCacheModeCA"), + ("CU_FUNC_ATTRIBUTE_MAX", "hipFuncAttributeMax"), + ("CU_GRAPHICS_MAP_RESOURCE_FLAGS_NONE", "hipGraphicsMapFlagsNone"), + ("CU_GRAPHICS_MAP_RESOURCE_FLAGS_READ_ONLY", "hipGraphicsMapFlagsReadOnly"), + ("CU_GRAPHICS_MAP_RESOURCE_FLAGS_WRITE_DISCARD", "hipGraphicsMapFlagsWriteDiscard"), + ("CU_GRAPHICS_REGISTER_FLAGS_NONE", "hipGraphicsRegisterFlagsNone"), + ("CU_GRAPHICS_REGISTER_FLAGS_READ_ONLY", "hipGraphicsRegisterFlagsReadOnly"), + ("CU_GRAPHICS_REGISTER_FLAGS_WRITE_DISCARD", "hipGraphicsRegisterFlagsWriteDiscard"), + ("CU_GRAPHICS_REGISTER_FLAGS_SURFACE_LDST", "hipGraphicsRegisterFlagsSurfaceLoadStore"), + ("CU_GRAPHICS_REGISTER_FLAGS_TEXTURE_GATHER", "hipGraphicsRegisterFlagsTextureGather"), + ("CU_OCCUPANCY_DEFAULT", "hipOccupancyDefault"), + ("CU_OCCUPANCY_DISABLE_CACHING_OVERRIDE", "hipOccupancyDisableCachingOverride"), + ("CU_FUNC_CACHE_PREFER_NONE", "hipFuncCachePreferNone"), + ("CU_FUNC_CACHE_PREFER_SHARED", "hipFuncCachePreferShared"), + ("CU_FUNC_CACHE_PREFER_L1", "hipFuncCachePreferL1"), + ("CU_FUNC_CACHE_PREFER_EQUAL", "hipFuncCachePreferEqual"), + ("CU_IPC_MEM_LAZY_ENABLE_PEER_ACCESS", "hipIpcMemLazyEnablePeerAccess"), + ("CUDA_IPC_HANDLE_SIZE", "HIP_IPC_HANDLE_SIZE"), + ("CU_JIT_CACHE_OPTION_NONE", "hipJitCacheModeOptionNone"), + ("CU_JIT_CACHE_OPTION_CG", "hipJitCacheModeOptionCG"), + ("CU_JIT_CACHE_OPTION_CA", "hipJitCacheModeOptionCA"), + ("CU_PREFER_PTX", "hipJitFallbackPreferPtx"), + ("CU_PREFER_BINARY", "hipJitFallbackPreferBinary"), + ("CU_JIT_MAX_REGISTERS", "hipJitOptionMaxRegisters"), + ("CU_JIT_THREADS_PER_BLOCK", "hipJitOptionThreadsPerBlock"), + ("CU_JIT_WALL_TIME", "hipJitOptionWallTime"), + ("CU_JIT_INFO_LOG_BUFFER", "hipJitOptionInfoLogBuffer"), + ("CU_JIT_INFO_LOG_BUFFER_SIZE_BYTES", "hipJitOptionInfoLogBufferSizeBytes"), + ("CU_JIT_ERROR_LOG_BUFFER", "hipJitOptionErrorLogBuffer"), + ("CU_JIT_ERROR_LOG_BUFFER_SIZE_BYTES", "hipJitOptionErrorLogBufferSizeBytes"), + ("CU_JIT_OPTIMIZATION_LEVEL", "hipJitOptionOptimizationLevel"), + ("CU_JIT_TARGET_FROM_CUCONTEXT", "hipJitOptionTargetFromContext"), + ("CU_JIT_TARGET", "hipJitOptionTarget"), + ("CU_JIT_FALLBACK_STRATEGY", "hipJitOptionFallbackStrategy"), + ("CU_JIT_GENERATE_DEBUG_INFO", "hipJitOptionGenerateDebugInfo"), + ("CU_JIT_LOG_VERBOSE", "hipJitOptionLogVerbose"), + ("CU_JIT_GENERATE_LINE_INFO", "hipJitOptionGenerateLineInfo"), + ("CU_JIT_CACHE_MODE", "hipJitOptionCacheMode"), + ("CU_JIT_NEW_SM3X_OPT", "hipJitOptionSm3xOpt"), + ("CU_JIT_FAST_COMPILE", "hipJitOptionFastCompile"), + ("CU_JIT_NUM_OPTIONS", "hipJitOptionNumOptions"), + ("CU_TARGET_COMPUTE_10", "hipJitTargetCompute10"), + ("CU_TARGET_COMPUTE_11", "hipJitTargetCompute11"), + ("CU_TARGET_COMPUTE_12", "hipJitTargetCompute12"), + ("CU_TARGET_COMPUTE_13", "hipJitTargetCompute13"), + ("CU_TARGET_COMPUTE_20", "hipJitTargetCompute20"), + ("CU_TARGET_COMPUTE_21", "hipJitTargetCompute21"), + ("CU_TARGET_COMPUTE_30", "hipJitTargetCompute30"), + ("CU_TARGET_COMPUTE_32", "hipJitTargetCompute32"), + ("CU_TARGET_COMPUTE_35", "hipJitTargetCompute35"), + ("CU_TARGET_COMPUTE_37", "hipJitTargetCompute37"), + ("CU_TARGET_COMPUTE_50", "hipJitTargetCompute50"), + ("CU_TARGET_COMPUTE_52", "hipJitTargetCompute52"), + ("CU_TARGET_COMPUTE_53", "hipJitTargetCompute53"), + ("CU_TARGET_COMPUTE_60", "hipJitTargetCompute60"), + ("CU_TARGET_COMPUTE_61", "hipJitTargetCompute61"), + ("CU_TARGET_COMPUTE_62", "hipJitTargetCompute62"), + ("CU_JIT_INPUT_CUBIN", "hipJitInputTypeBin"), + ("CU_JIT_INPUT_PTX", "hipJitInputTypePtx"), + ("CU_JIT_INPUT_FATBINARY", "hipJitInputTypeFatBinary"), + ("CU_JIT_INPUT_OBJECT", "hipJitInputTypeObject"), + ("CU_JIT_INPUT_LIBRARY", "hipJitInputTypeLibrary"), + ("CU_JIT_NUM_INPUT_TYPES", "hipJitInputTypeNumInputTypes"), + ("CU_LIMIT_STACK_SIZE", "hipLimitStackSize"), + ("CU_LIMIT_PRINTF_FIFO_SIZE", "hipLimitPrintfFifoSize"), + ("CU_LIMIT_MALLOC_HEAP_SIZE", "hipLimitMallocHeapSize"), + ("CU_LIMIT_DEV_RUNTIME_SYNC_DEPTH", "hipLimitDevRuntimeSyncDepth"), + ("CU_LIMIT_DEV_RUNTIME_PENDING_LAUNCH_COUNT", "hipLimitDevRuntimePendingLaunchCount"), + ("CU_MEM_ATTACH_GLOBAL", "hipMemAttachGlobal"), + ("CU_MEM_ATTACH_HOST", "hipMemAttachHost"), + ("CU_MEM_ATTACH_SINGLE", "hipMemAttachSingle"), + ("CU_MEMORYTYPE_HOST", "hipMemTypeHost"), + ("CU_MEMORYTYPE_DEVICE", "hipMemTypeDevice"), + ("CU_MEMORYTYPE_ARRAY", "hipMemTypeArray"), + ("CU_MEMORYTYPE_UNIFIED", "hipMemTypeUnified"), + ("CU_MEMHOSTREGISTER_READ_ONLY", "hipHostRegisterReadOnly"), + ("CU_MEMPOOL_ATTR_RELEASE_THRESHOLD", "hipMemPoolAttrReleaseThreshold"), + ("CU_MEMPOOL_ATTR_RESERVED_MEM_CURRENT", "hipMemPoolAttrReservedMemCurrent"), + ("CU_MEMPOOL_ATTR_RESERVED_MEM_HIGH", "hipMemPoolAttrReservedMemHigh"), + ("CU_MEMPOOL_ATTR_REUSE_ALLOW_INTERNAL_DEPENDENCIES", "hipMemPoolReuseAllowInternalDependencies"), + ("CU_MEMPOOL_ATTR_REUSE_ALLOW_OPPORTUNISTIC", "hipMemPoolReuseAllowOpportunistic"), + ("CU_MEMPOOL_ATTR_REUSE_FOLLOW_EVENT_DEPENDENCIES", "hipMemPoolReuseFollowEventDependencies"), + ("CU_MEMPOOL_ATTR_USED_MEM_CURRENT", "hipMemPoolAttrUsedMemCurrent"), + ("CU_MEMPOOL_ATTR_USED_MEM_HIGH", "hipMemPoolAttrUsedMemHigh"), + ("CU_MEM_ACCESS_FLAGS_PROT_NONE", "hipMemAccessFlagsProtNone"), + ("CU_MEM_ACCESS_FLAGS_PROT_READ", "hipMemAccessFlagsProtRead"), + ("CU_MEM_ACCESS_FLAGS_PROT_READWRITE", "hipMemAccessFlagsProtReadWrite"), + ("CU_MEM_ALLOCATION_TYPE_INVALID", "hipMemAllocationTypeInvalid"), + ("CU_MEM_ALLOCATION_TYPE_MAX", "hipMemAllocationTypeMax"), + ("CU_MEM_ALLOCATION_TYPE_PINNED", "hipMemAllocationTypePinned"), + ("CU_MEM_ALLOC_GRANULARITY_MINIMUM", "hipMemAllocationGranularityMinimum"), + ("CU_MEM_ALLOC_GRANULARITY_RECOMMENDED", "hipMemAllocationGranularityRecommended"), + ("CU_MEM_HANDLE_TYPE_GENERIC", "hipMemHandleTypeGeneric"), + ("CU_MEM_HANDLE_TYPE_NONE", "hipMemHandleTypeNone"), + ("CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR", "hipMemHandleTypePosixFileDescriptor"), + ("CU_MEM_HANDLE_TYPE_WIN32", "hipMemHandleTypeWin32"), + ("CU_MEM_HANDLE_TYPE_WIN32_KMT", "hipMemHandleTypeWin32Kmt"), + ("CU_MEM_LOCATION_TYPE_DEVICE", "hipMemLocationTypeDevice"), + ("CU_MEM_LOCATION_TYPE_HOST", "hipMemLocationTypeHost"), + ("CU_MEM_LOCATION_TYPE_INVALID", "hipMemLocationTypeInvalid"), + ("CU_MEM_OPERATION_TYPE_MAP", "hipMemOperationTypeMap"), + ("CU_MEM_OPERATION_TYPE_UNMAP", "hipMemOperationTypeUnmap"), + ("CU_RESOURCE_TYPE_ARRAY", "hipResourceTypeArray"), + ("CU_RESOURCE_TYPE_MIPMAPPED_ARRAY", "hipResourceTypeMipmappedArray"), + ("CU_RESOURCE_TYPE_LINEAR", "hipResourceTypeLinear"), + ("CU_RESOURCE_TYPE_PITCH2D", "hipResourceTypePitch2D"), + ("CU_RES_VIEW_FORMAT_NONE", "hipResViewFormatNone"), + ("CU_RES_VIEW_FORMAT_UINT_1X8", "hipResViewFormatUnsignedChar1"), + ("CU_RES_VIEW_FORMAT_UINT_2X8", "hipResViewFormatUnsignedChar2"), + ("CU_RES_VIEW_FORMAT_UINT_4X8", "hipResViewFormatUnsignedChar4"), + ("CU_RES_VIEW_FORMAT_SINT_1X8", "hipResViewFormatSignedChar1"), + ("CU_RES_VIEW_FORMAT_SINT_2X8", "hipResViewFormatSignedChar2"), + ("CU_RES_VIEW_FORMAT_SINT_4X8", "hipResViewFormatSignedChar4"), + ("CU_RES_VIEW_FORMAT_UINT_1X16", "hipResViewFormatUnsignedShort1"), + ("CU_RES_VIEW_FORMAT_UINT_2X16", "hipResViewFormatUnsignedShort2"), + ("CU_RES_VIEW_FORMAT_UINT_4X16", "hipResViewFormatUnsignedShort4"), + ("CU_RES_VIEW_FORMAT_SINT_1X16", "hipResViewFormatSignedShort1"), + ("CU_RES_VIEW_FORMAT_SINT_2X16", "hipResViewFormatSignedShort2"), + ("CU_RES_VIEW_FORMAT_SINT_4X16", "hipResViewFormatSignedShort4"), + ("CU_RES_VIEW_FORMAT_UINT_1X32", "hipResViewFormatUnsignedInt1"), + ("CU_RES_VIEW_FORMAT_UINT_2X32", "hipResViewFormatUnsignedInt2"), + ("CU_RES_VIEW_FORMAT_UINT_4X32", "hipResViewFormatUnsignedInt4"), + ("CU_RES_VIEW_FORMAT_SINT_1X32", "hipResViewFormatSignedInt1"), + ("CU_RES_VIEW_FORMAT_SINT_2X32", "hipResViewFormatSignedInt2"), + ("CU_RES_VIEW_FORMAT_SINT_4X32", "hipResViewFormatSignedInt4"), + ("CU_RES_VIEW_FORMAT_FLOAT_1X16", "hipResViewFormatHalf1"), + ("CU_RES_VIEW_FORMAT_FLOAT_2X16", "hipResViewFormatHalf2"), + ("CU_RES_VIEW_FORMAT_FLOAT_4X16", "hipResViewFormatHalf4"), + ("CU_RES_VIEW_FORMAT_FLOAT_1X32", "hipResViewFormatFloat1"), + ("CU_RES_VIEW_FORMAT_FLOAT_2X32", "hipResViewFormatFloat2"), + ("CU_RES_VIEW_FORMAT_FLOAT_4X32", "hipResViewFormatFloat4"), + ("CU_RES_VIEW_FORMAT_UNSIGNED_BC1", "hipResViewFormatUnsignedBlockCompressed1"), + ("CU_RES_VIEW_FORMAT_UNSIGNED_BC2", "hipResViewFormatUnsignedBlockCompressed2"), + ("CU_RES_VIEW_FORMAT_UNSIGNED_BC3", "hipResViewFormatUnsignedBlockCompressed3"), + ("CU_RES_VIEW_FORMAT_UNSIGNED_BC4", "hipResViewFormatUnsignedBlockCompressed4"), + ("CU_RES_VIEW_FORMAT_SIGNED_BC4", "hipResViewFormatSignedBlockCompressed4"), + ("CU_RES_VIEW_FORMAT_UNSIGNED_BC5", "hipResViewFormatUnsignedBlockCompressed5"), + ("CU_RES_VIEW_FORMAT_SIGNED_BC5", "hipResViewFormatSignedBlockCompressed5"), + ("CU_RES_VIEW_FORMAT_UNSIGNED_BC6H", "hipResViewFormatUnsignedBlockCompressed6H"), + ("CU_RES_VIEW_FORMAT_SIGNED_BC6H", "hipResViewFormatSignedBlockCompressed6H"), + ("CU_RES_VIEW_FORMAT_UNSIGNED_BC7", "hipResViewFormatUnsignedBlockCompressed7"), + ("CU_SHARED_MEM_CONFIG_DEFAULT_BANK_SIZE", "hipSharedMemBankSizeDefault"), + ("CU_SHARED_MEM_CONFIG_FOUR_BYTE_BANK_SIZE", "hipSharedMemBankSizeFourByte"), + ("CU_SHARED_MEM_CONFIG_EIGHT_BYTE_BANK_SIZE", "hipSharedMemBankSizeEightByte"), + ("CU_STREAM_DEFAULT", "hipStreamDefault"), + ("CU_STREAM_NON_BLOCKING", "hipStreamNonBlocking"), + ("CU_STREAM_WAIT_VALUE_GEQ", "hipStreamWaitValueGeq"), + ("CU_STREAM_WAIT_VALUE_EQ", "hipStreamWaitValueEq"), + ("CU_STREAM_WAIT_VALUE_AND", "hipStreamWaitValueAnd"), + ("CU_STREAM_WAIT_VALUE_FLUSH", "hipStreamWaitValueFlush"), + ("CU_STREAM_WRITE_VALUE_DEFAULT", "hipStreamWriteValueDefault"), + ("CU_STREAM_WRITE_VALUE_NO_MEMORY_BARRIER", "hipStreamWriteValueNoMemoryBarrier"), + ("CU_STREAM_MEM_OP_WAIT_VALUE_32", "hipStreamBatchMemOpWaitValue32"), + ("CU_STREAM_MEM_OP_WRITE_VALUE_32", "hipStreamBatchMemOpWriteValue32"), + ("CU_STREAM_MEM_OP_FLUSH_REMOTE_WRITES", "hipStreamBatchMemOpFlushRemoteWrites"), + ("cuGetErrorName", "hipGetErrorName"), + ("cuGetErrorString", "hipDrvGetErrorString"), + ("cuInit", "hipInit"), + ("cuDriverGetVersion", "hipDriverGetVersion"), + ("cuCtxCreate", "hipCtxCreate"), + ("cuCtxCreate_v2", "hipCtxCreate"), + ("cuCtxDestroy", "hipCtxDestroy"), + ("cuCtxDestroy_v2", "hipCtxDestroy"), + ("cuCtxGetApiVersion", "hipCtxGetApiVersion"), + ("cuCtxGetCacheConfig", "hipCtxGetCacheConfig"), + ("cuCtxGetCurrent", "hipCtxGetCurrent"), + ("cuCtxGetDevice", "hipCtxGetDevice"), + ("cuCtxGetFlags", "hipCtxGetFlags"), + ("cuDeviceGetUuid", "hipDeviceGetUuid"), + ("cuCtxGetLimit", "hipCtxGetLimit"), + ("cuCtxGetSharedMemConfig", "hipCtxGetSharedMemConfig"), + ("cuCtxGetStreamPriorityRange", "hipCtxGetStreamPriorityRange"), + ("cuCtxPopCurrent_v2", "hipCtxPopCurrent"), + ("cuCtxPushCurrent_v2", "hipCtxPushCurrent"), + ("cuCtxSetCacheConfig", "hipCtxSetCacheConfig"), + ("cuCtxSetCurrent", "hipCtxSetCurrent"), + ("cuCtxSetLimit", "hipCtxSetLimit"), + ("cuCtxSetSharedMemConfig", "hipCtxSetSharedMemConfig"), + ("cuCtxSynchronize", "hipCtxSynchronize"), + ("cuCtxAttach", "hipCtxAttach"), + ("cuCtxDetach", "hipCtxDetach"), + ("cuCtxEnablePeerAccess", "hipCtxEnablePeerAccess"), + ("cuCtxDisablePeerAccess", "hipCtxDisablePeerAccess"), + ("cuDeviceCanAccessPeer", "hipDeviceCanAccessPeer"), + ("cuDeviceGetP2PAttribute", "hipDeviceGetP2PAttribute"), + ("cuDevicePrimaryCtxGetState", "hipDevicePrimaryCtxGetState"), + ("cuDevicePrimaryCtxRelease", "hipDevicePrimaryCtxRelease"), + ("cuDevicePrimaryCtxReset", "hipDevicePrimaryCtxReset"), + ("cuDevicePrimaryCtxRetain", "hipDevicePrimaryCtxRetain"), + ("cuDevicePrimaryCtxSetFlags", "hipDevicePrimaryCtxSetFlags"), + ("cuDeviceGet", "hipDeviceGet"), + ("cuDeviceGetName", "hipDeviceGetName"), + ("cuDeviceGetCount", "hipGetDeviceCount"), + ("cuDeviceGetAttribute", "hipDeviceGetAttribute"), + ("cuDeviceGetPCIBusId", "hipDeviceGetPCIBusId"), + ("cuDeviceGetByPCIBusId", "hipDeviceGetByPCIBusId"), + ("cuDeviceTotalMem_v2", "hipDeviceTotalMem"), + ("cuDeviceComputeCapability", "hipDeviceComputeCapability"), + ("cuDeviceGetProperties", "hipGetDeviceProperties"), + ("cuLinkAddData", "hipLinkAddData"), + ("cuLinkAddFile", "hipLinkAddFile"), + ("cuLinkComplete", "hipLinkComplete"), + ("cuLinkCreate", "hipLinkCreate"), + ("cuLinkDestroy", "hipLinkDestroy"), + ("cuModuleGetFunction", "hipModuleGetFunction"), + ("cuModuleGetGlobal", "hipModuleGetGlobal"), + ("cuModuleGetGlobal_v2", "hipModuleGetGlobal"), + ("cuModuleGetSurfRef", "hipModuleGetSurfRef"), + ("cuModuleGetTexRef", "hipModuleGetTexRef"), + ("cuModuleLoad", "hipModuleLoad"), + ("cuModuleLoadData", "hipModuleLoadData"), + ("cuModuleLoadDataEx", "hipModuleLoadDataEx"), + ("cuModuleLoadFatBinary", "hipModuleLoadFatBinary"), + ("cuModuleUnload", "hipModuleUnload"), + ("CU_DEVICE_P2P_ATTRIBUTE_PERFORMANCE_RANK", "hipDeviceP2PAttributePerformanceRank"), + ("CU_DEVICE_P2P_ATTRIBUTE_ACCESS_SUPPORTED", "hipDeviceP2PAttributeAccessSupported"), + ("CU_DEVICE_P2P_ATTRIBUTE_NATIVE_ATOMIC_SUPPORTED", "hipDeviceP2PAttributeNativeAtomicSupported"), + ("CU_EVENT_DEFAULT", "hipEventDefault"), + ("CU_EVENT_BLOCKING_SYNC", "hipEventBlockingSync"), + ("CU_EVENT_DISABLE_TIMING", "hipEventDisableTiming"), + ("CU_EVENT_INTERPROCESS", "hipEventInterprocess"), + ("cuEventCreate", "hipEventCreate"), + ("cuEventDestroy", "hipEventDestroy"), + ("cuEventDestroy_v2", "hipEventDestroy"), + ("cuEventElapsedTime", "hipEventElapsedTime"), + ("cuEventQuery", "hipEventQuery"), + ("cuEventRecord", "hipEventRecord"), + ("cuEventSynchronize", "hipEventSynchronize"), + ("cuFuncSetAttribute", "hipFuncSetAttribute"), + ("cuFuncGetAttribute", "hipFuncGetAttribute"), + ("cuFuncSetCacheConfig", "hipFuncSetCacheConfig"), + ("cuFuncSetSharedMemConfig", "hipFuncSetSharedMemConfig"), + ("cuLaunchKernel", "hipModuleLaunchKernel"), + ("cuLaunchCooperativeKernel", "hipModuleLaunchCooperativeKernel"), + ("cuFuncSetBlockShape", "hipFuncSetBlockShape"), + ("cuFuncSetSharedSize", "hipFuncSetSharedSize"), + ("cuLaunch", "hipLaunch"), + ("cuLaunchGrid", "hipLaunchGrid"), + ("cuLaunchGridAsync", "hipLaunchGridAsync"), + ("cuParamSetf", "hipParamSetf"), + ("cuParamSeti", "hipParamSeti"), + ("cuParamSetSize", "hipParamSetSize"), + ("cuParamSetv", "hipParamSetv"), + ("cuOccupancyMaxActiveBlocksPerMultiprocessor", "hipModuleOccupancyMaxActiveBlocksPerMultiprocessor"), + ("cuOccupancyMaxActiveBlocksPerMultiprocessorWithFlags", "hipModuleOccupancyMaxActiveBlocksPerMultiprocessorWithFlags"), + ("cuOccupancyMaxPotentialBlockSize", "hipModuleOccupancyMaxPotentialBlockSize"), + ("cuOccupancyMaxPotentialBlockSizeWithFlags", "hipModuleOccupancyMaxPotentialBlockSizeWithFlags"), + ("cuStreamAddCallback", "hipStreamAddCallback"), + ("cuStreamAttachMemAsync", "hipStreamAttachMemAsync"), + ("cuStreamCreate", "hipStreamCreate__"), + ("cuStreamCreateWithPriority", "hipStreamCreateWithPriority"), + ("cuStreamDestroy", "hipStreamDestroy"), + ("cuStreamDestroy_v2", "hipStreamDestroy"), + ("cuStreamGetFlags", "hipStreamGetFlags"), + ("cuStreamGetPriority", "hipStreamGetPriority"), + ("cuStreamQuery", "hipStreamQuery"), + ("cuStreamSynchronize", "hipStreamSynchronize"), + ("cuStreamWaitEvent", "hipStreamWaitEvent"), + ("cuStreamWaitValue32", "hipStreamWaitValue32"), + ("cuStreamWriteValue32", "hipStreamWriteValue32"), + ("cuStreamBatchMemOp", "hipStreamBatchMemOp"), + ("cuArray3DCreate", "hipArray3DCreate"), + ("cuArray3DGetDescriptor", "hipArray3DGetDescriptor"), + ("cuArrayCreate", "hipArrayCreate"), + ("cuArrayDestroy", "hipArrayDestroy"), + ("cuArrayGetDescriptor", "hipArrayGetDescriptor"), + ("cuIpcCloseMemHandle", "hipIpcCloseMemHandle"), + ("cuIpcGetEventHandle", "hipIpcGetEventHandle"), + ("cuIpcGetMemHandle", "hipIpcGetMemHandle"), + ("cuIpcOpenEventHandle", "hipIpcOpenEventHandle"), + ("cuIpcOpenMemHandle", "hipIpcOpenMemHandle"), + ("cuMemAlloc_v2", "hipMalloc"), + ("cuMemAllocHost", "hipMemAllocHost"), + ("cuMemAllocManaged", "hipMemAllocManaged"), + ("cuMemAllocPitch", "hipMemAllocPitch__"), + ("cuMemcpy", "hipMemcpy__"), + ("cuMemcpy2D", "hipMemcpy2D__"), + ("cuMemcpy2DAsync", "hipMemcpy2DAsync__"), + ("cuMemcpy2DUnaligned", "hipMemcpy2DUnaligned"), + ("cuMemcpy3D", "hipMemcpy3D__"), + ("cuMemcpy3DAsync", "hipMemcpy3DAsync__"), + ("cuMemcpy3DPeer", "hipMemcpy3DPeer__"), + ("cuMemcpy3DPeerAsync", "hipMemcpy3DPeerAsync__"), + ("cuMemcpyAsync", "hipMemcpyAsync__"), + ("cuMemcpyAtoA", "hipMemcpyAtoA"), + ("cuMemcpyAtoD", "hipMemcpyAtoD"), + ("cuMemcpyAtoH", "hipMemcpyAtoH"), + ("cuMemcpyAtoHAsync", "hipMemcpyAtoHAsync"), + ("cuMemcpyDtoA", "hipMemcpyDtoA"), + ("cuMemcpyDtoD_v2", "hipMemcpyDtoD"), + ("cuMemcpyDtoDAsync_v2", "hipMemcpyDtoDAsync"), + ("cuMemcpyDtoH_v2", "hipMemcpyDtoH"), + ("cuMemcpyDtoHAsync_v2", "hipMemcpyDtoHAsync"), + ("cuMemcpyHtoA", "hipMemcpyHtoA"), + ("cuMemcpyHtoAAsync", "hipMemcpyHtoAAsync"), + ("cuMemcpyHtoD_v2", "hipMemcpyHtoD"), + ("cuMemcpyHtoDAsync_v2", "hipMemcpyHtoDAsync"), + ("cuMemcpyPeerAsync", "hipMemcpyPeerAsync__"), + ("cuMemcpyPeer", "hipMemcpyPeer__"), + ("cuMemFree", "hipFree"), + ("cuMemFree_v2", "hipFree"), + ("cuMemFreeHost", "hipHostFree"), + ("cuMemGetAddressRange", "hipMemGetAddressRange"), + ("cuMemGetInfo_v2", "hipMemGetInfo"), + ("cuMemHostAlloc", "hipHostMalloc"), + ("cuMemHostGetDevicePointer", "hipMemHostGetDevicePointer"), + ("cuMemHostGetFlags", "hipMemHostGetFlags"), + ("cuMemHostRegister_v2", "hipHostRegister"), + ("cuMemHostUnregister", "hipHostUnregister"), + ("cuMemsetD16_v2", "hipMemsetD16"), + ("cuMemsetD16Async", "hipMemsetD16Async"), + ("cuMemsetD2D16_v2", "hipMemsetD2D16"), + ("cuMemsetD2D16Async", "hipMemsetD2D16Async"), + ("cuMemsetD2D32_v2", "hipMemsetD2D32"), + ("cuMemsetD2D32Async", "hipMemsetD2D32Async"), + ("cuMemsetD2D8_v2", "hipMemsetD2D8"), + ("cuMemsetD2D8Async", "hipMemsetD2D8Async"), + ("cuMemsetD32_v2", "hipMemset"), + ("cuMemsetD32Async", "hipMemsetAsync"), + ("cuMemsetD8_v2", "hipMemsetD8"), + ("cuMemsetD8Async", "hipMemsetD8Async"), + ("cuMipmappedArrayCreate", "hipMipmappedArrayCreate"), + ("cuMipmappedArrayDestroy", "hipMipmappedArrayDestroy"), + ("cuMipmappedArrayGetLevel", "hipMipmappedArrayGetLevel"), + ("cuMemPrefetchAsync", "hipMemPrefetchAsync__"), + ("cuMemAdvise", "hipMemAdvise"), + ("cuMemRangeGetAttribute", "hipMemRangeGetAttribute"), + ("cuMemRangeGetAttributes", "hipMemRangeGetAttributes"), + ("cuPointerGetAttribute", "hipPointerGetAttribute"), + ("cuMemGetAddressRange_v2", "hipMemGetAddressRange"), + ("cuArray3DCreate_v2", "hipArray3DCreate"), + ("cuArray3DGetDescriptor_v2", "hipArray3DGetDescriptor"), + ("cuArrayGetDescriptor_v2", "hipArrayGetDescriptor"), + ("cuMemAlloc", "hipMalloc"), + ("cuMemAllocHost_v2", "hipMemAllocHost"), + ("cuMemAllocPitch_v2", "hipMemAllocPitch"), + ("cuMemGetInfo", "hipMemGetInfo"), + ("cuMemHostGetDevicePointer_v2", "hipHostGetDevicePointer"), + ("cuMemHostRegister", "hipHostRegister"), + ("cuMemcpy2DAsync_v2", "hipMemcpyParam2DAsync"), + ("cuMemcpy2DUnaligned_v2", "hipDrvMemcpy2DUnaligned"), + ("cuMemcpy2D_v2", "hipMemcpyParam2D"), + ("cuMemcpy3DAsync_v2", "hipDrvMemcpy3DAsync"), + ("cuMemcpy3D_v2", "hipDrvMemcpy3D"), + ("cuMemcpyAtoA_v2", "hipMemcpyAtoA"), + ("cuMemcpyAtoD_v2", "hipMemcpyAtoD"), + ("cuMemcpyAtoHAsync_v2", "hipMemcpyAtoHAsync"), + ("cuMemcpyAtoH_v2", "hipMemcpyAtoH"), + ("cuMemcpyDtoA_v2", "hipMemcpyDtoA"), + ("cuMemcpyDtoD", "hipMemcpyDtoD"), + ("cuMemcpyDtoDAsync", "hipMemcpyDtoDAsync"), + ("cuMemcpyDtoH", "hipMemcpyDtoH"), + ("cuMemcpyDtoHAsync", "hipMemcpyDtoHAsync"), + ("cuMemcpyHtoA_v2", "hipMemcpyHtoA"), + ("cuMemcpyHtoD", "hipMemcpyHtoD"), + ("cuMemcpyHtoDAsync", "hipMemcpyHtoDAsync"), + ("cuMemsetD16", "hipMemsetD16"), + ("cuMemsetD32", "hipMemsetD32"), + ("cuMemsetD8", "hipMemsetD8"), + ("cuMemAddressFree", "hipMemAddressFree"), + ("cuMemAddressReserve", "hipMemAddressReserve"), + ("cuMemCreate", "hipMemCreate"), + ("cuMemExportToShareableHandle", "hipMemExportToShareableHandle"), + ("cuMemGetAccess", "hipMemGetAccess"), + ("cuMemGetAllocationGranularity", "hipMemGetAllocationGranularity"), + ("cuMemGetAllocationPropertiesFromHandle", "hipMemGetAllocationPropertiesFromHandle"), + ("cuMemImportFromShareableHandle", "hipMemImportFromShareableHandle"), + ("cuMemMap", "hipMemMap"), + ("cuMemMapArrayAsync", "hipMemMapArrayAsync"), + ("cuMemRelease", "hipMemRelease"), + ("cuMemRetainAllocationHandle", "hipMemRetainAllocationHandle"), + ("cuMemSetAccess", "hipMemSetAccess"), + ("cuMemUnmap", "hipMemUnmap"), + ("cuMemAllocAsync", "hipMallocAsync"), + ("cuMemAllocFromPoolAsync", "hipMallocFromPoolAsync"), + ("cuMemFreeAsync", "hipFreeAsync"), + ("cuMemPoolCreate", "hipMemPoolCreate"), + ("cuMemPoolDestroy", "hipMemPoolDestroy"), + ("cuMemPoolExportPointer", "hipMemPoolExportPointer"), + ("cuMemPoolExportToShareableHandle", "hipMemPoolExportToShareableHandle"), + ("cuMemPoolGetAccess", "hipMemPoolGetAccess"), + ("cuMemPoolGetAttribute", "hipMemPoolGetAttribute"), + ("cuMemPoolImportFromShareableHandle", "hipMemPoolImportFromShareableHandle"), + ("cuMemPoolImportPointer", "hipMemPoolImportPointer"), + ("cuMemPoolSetAccess", "hipMemPoolSetAccess"), + ("cuMemPoolSetAttribute", "hipMemPoolSetAttribute"), + ("cuMemPoolTrimTo", "hipMemPoolTrimTo"), + ("cuPointerGetAttributes", "hipPointerGetAttributes"), + ("cuPointerSetAttribute", "hipPointerSetAttribute"), + ("CU_TR_FILTER_MODE_POINT", "hipFilterModePoint"), + ("CU_TR_FILTER_MODE_LINEAR", "hipFilterModeLinear"), + ("cuTexRefGetAddress", "hipTexRefGetAddress"), + ("cuTexRefGetAddressMode", "hipTexRefGetAddressMode"), + ("cuTexRefGetArray", "hipTexRefGetArray"), + ("cuTexRefGetBorderColor", "hipTexRefGetBorderColor"), + ("cuTexRefGetFilterMode", "hipTexRefGetFilterMode"), + ("cuTexRefGetFlags", "hipTexRefGetFlags"), + ("cuTexRefGetFormat", "hipTexRefGetFormat"), + ("cuTexRefGetMaxAnisotropy", "hipTexRefGetMaxAnisotropy"), + ("cuTexRefGetMipmapFilterMode", "hipTexRefGetMipmapFilterMode"), + ("cuTexRefGetMipmapLevelBias", "hipTexRefGetMipmapLevelBias"), + ("cuTexRefGetMipmapLevelClamp", "hipTexRefGetMipmapLevelClamp"), + ("cuTexRefGetMipmappedArray", "hipTexRefGetMipmappedArray"), + ("cuTexRefSetAddress", "hipTexRefSetAddress"), + ("cuTexRefSetAddress2D", "hipTexRefSetAddress2D"), + ("cuTexRefSetAddressMode", "hipTexRefSetAddressMode"), + ("cuTexRefSetArray", "hipTexRefSetArray"), + ("cuTexRefSetBorderColor", "hipTexRefSetBorderColor"), + ("cuTexRefSetFilterMode", "hipTexRefSetFilterMode"), + ("cuTexRefSetFlags", "hipTexRefSetFlags"), + ("cuTexRefSetFormat", "hipTexRefSetFormat"), + ("cuTexRefSetMaxAnisotropy", "hipTexRefSetMaxAnisotropy"), + ("cuTexRefSetMipmapFilterMode", "hipTexRefSetMipmapFilterMode"), + ("cuTexRefSetMipmapLevelBias", "hipTexRefSetMipmapLevelBias"), + ("cuTexRefSetMipmapLevelClamp", "hipTexRefSetMipmapLevelClamp"), + ("cuTexRefSetMipmappedArray", "hipTexRefSetMipmappedArray"), + ("cuTexRefCreate", "hipTexRefCreate"), + ("cuTexRefDestroy", "hipTexRefDestroy"), + ("cuSurfRefGetArray", "hipSurfRefGetArray"), + ("cuSurfRefSetArray", "hipSurfRefSetArray"), + ("cuTexObjectCreate", "hipTexObjectCreate"), + ("cuTexObjectDestroy", "hipTexObjectDestroy"), + ("cuTexObjectGetResourceDesc", "hipTexObjectGetResourceDesc"), + ("cuTexObjectGetResourceViewDesc", "hipTexObjectGetResourceViewDesc"), + ("cuTexObjectGetTextureDesc", "hipTexObjectGetTextureDesc"), + ("cuSurfObjectCreate", "hipSurfObjectCreate"), + ("cuSurfObjectDestroy", "hipSurfObjectDestroy"), + ("cuSurfObjectGetResourceDesc", "hipSurfObjectGetResourceDesc"), + ("cuGraphicsMapResources", "hipGraphicsMapResources"), + ("cuGraphicsResourceGetMappedMipmappedArray", "hipGraphicsResourceGetMappedMipmappedArray"), + ("cuGraphicsResourceGetMappedPointer", "hipGraphicsResourceGetMappedPointer"), + ("cuGraphicsResourceSetMapFlags", "hipGraphicsResourceSetMapFlags"), + ("cuGraphicsSubResourceGetMappedArray", "hipGraphicsSubResourceGetMappedArray"), + ("cuGraphicsUnmapResources", "hipGraphicsUnmapResources"), + ("cuGraphicsUnregisterResource", "hipGraphicsUnregisterResource"), + ("cuProfilerInitialize", "hipProfilerInitialize"), + ("cuProfilerStart", "hipProfilerStart"), + ("cuProfilerStop", "hipProfilerStop"), + ("CU_GL_DEVICE_LIST_ALL", "HIP_GL_DEVICE_LIST_ALL"), + ("CU_GL_DEVICE_LIST_CURRENT_FRAME", "HIP_GL_DEVICE_LIST_CURRENT_FRAME"), + ("CU_GL_DEVICE_LIST_NEXT_FRAME", "HIP_GL_DEVICE_LIST_NEXT_FRAME"), + ("cuGLGetDevices", "hipGLGetDevices"), + ("cuGraphicsGLRegisterBuffer", "hipGraphicsGLRegisterBuffer"), + ("cuGraphicsGLRegisterImage", "hipGraphicsGLRegisterImage"), + ("cuWGLGetDevice", "hipWGLGetDevice"), + ("CU_GL_MAP_RESOURCE_FLAGS_NONE", "HIP_GL_MAP_RESOURCE_FLAGS_NONE"), + ("CU_GL_MAP_RESOURCE_FLAGS_READ_ONLY", "HIP_GL_MAP_RESOURCE_FLAGS_READ_ONLY"), + ("CU_GL_MAP_RESOURCE_FLAGS_WRITE_DISCARD", "HIP_GL_MAP_RESOURCE_FLAGS_WRITE_DISCARD"), + ("cuGLCtxCreate", "hipGLCtxCreate"), + ("cuGLInit", "hipGLInit"), + ("cuGLMapBufferObject", "hipGLMapBufferObject"), + ("cuGLMapBufferObjectAsync", "hipGLMapBufferObjectAsync"), + ("cuGLRegisterBufferObject", "hipGLRegisterBufferObject"), + ("cuGLSetBufferObjectMapFlags", "hipGLSetBufferObjectMapFlags"), + ("cuGLUnmapBufferObject", "hipGLUnmapBufferObject"), + ("cuGLUnmapBufferObjectAsync", "hipGLUnmapBufferObjectAsync"), + ("cuGLUnregisterBufferObject", "hipGLUnregisterBufferObject"), + ("CU_D3D9_DEVICE_LIST_ALL", "HIP_D3D9_DEVICE_LIST_ALL"), + ("CU_D3D9_DEVICE_LIST_CURRENT_FRAME", "HIP_D3D9_DEVICE_LIST_CURRENT_FRAME"), + ("CU_D3D9_DEVICE_LIST_NEXT_FRAME", "HIP_D3D9_DEVICE_LIST_NEXT_FRAME"), + ("cuD3D9CtxCreate", "hipD3D9CtxCreate"), + ("cuD3D9CtxCreateOnDevice", "hipD3D9CtxCreateOnDevice"), + ("cuD3D9GetDevice", "hipD3D9GetDevice"), + ("cuD3D9GetDevices", "hipD3D9GetDevices"), + ("cuD3D9GetDirect3DDevice", "hipD3D9GetDirect3DDevice"), + ("cuGraphicsD3D9RegisterResource", "hipGraphicsD3D9RegisterResource"), + ("CU_D3D9_MAPRESOURCE_FLAGS_NONE", "HIP_D3D9_MAPRESOURCE_FLAGS_NONE"), + ("CU_D3D9_MAPRESOURCE_FLAGS_READONLY", "HIP_D3D9_MAPRESOURCE_FLAGS_READONLY"), + ("CU_D3D9_MAPRESOURCE_FLAGS_WRITEDISCARD", "HIP_D3D9_MAPRESOURCE_FLAGS_WRITEDISCARD"), + ("CU_D3D9_REGISTER_FLAGS_NONE", "HIP_D3D9_REGISTER_FLAGS_NONE"), + ("CU_D3D9_REGISTER_FLAGS_ARRAY", "HIP_D3D9_REGISTER_FLAGS_ARRAY"), + ("cuD3D9MapResources", "hipD3D9MapResources"), + ("cuD3D9RegisterResource", "hipD3D9RegisterResource"), + ("cuD3D9ResourceGetMappedArray", "hipD3D9ResourceGetMappedArray"), + ("cuD3D9ResourceGetMappedPitch", "hipD3D9ResourceGetMappedPitch"), + ("cuD3D9ResourceGetMappedPointer", "hipD3D9ResourceGetMappedPointer"), + ("cuD3D9ResourceGetMappedSize", "hipD3D9ResourceGetMappedSize"), + ("cuD3D9ResourceGetSurfaceDimensions", "hipD3D9ResourceGetSurfaceDimensions"), + ("cuD3D9ResourceSetMapFlags", "hipD3D9ResourceSetMapFlags"), + ("cuD3D9UnmapResources", "hipD3D9UnmapResources"), + ("cuD3D9UnregisterResource", "hipD3D9UnregisterResource"), + ("CU_D3D10_DEVICE_LIST_ALL", "HIP_D3D10_DEVICE_LIST_ALL"), + ("CU_D3D10_DEVICE_LIST_CURRENT_FRAME", "HIP_D3D10_DEVICE_LIST_CURRENT_FRAME"), + ("CU_D3D10_DEVICE_LIST_NEXT_FRAME", "HIP_D3D10_DEVICE_LIST_NEXT_FRAME"), + ("cuD3D10GetDevice", "hipD3D10GetDevice"), + ("cuD3D10GetDevices", "hipD3D10GetDevices"), + ("cuGraphicsD3D10RegisterResource", "hipGraphicsD3D10RegisterResource"), + ("CU_D3D10_MAPRESOURCE_FLAGS_NONE", "HIP_D3D10_MAPRESOURCE_FLAGS_NONE"), + ("CU_D3D10_MAPRESOURCE_FLAGS_READONLY", "HIP_D3D10_MAPRESOURCE_FLAGS_READONLY"), + ("CU_D3D10_MAPRESOURCE_FLAGS_WRITEDISCARD", "HIP_D3D10_MAPRESOURCE_FLAGS_WRITEDISCARD"), + ("CU_D3D10_REGISTER_FLAGS_NONE", "HIP_D3D10_REGISTER_FLAGS_NONE"), + ("CU_D3D10_REGISTER_FLAGS_ARRAY", "HIP_D3D10_REGISTER_FLAGS_ARRAY"), + ("cuD3D10CtxCreate", "hipD3D10CtxCreate"), + ("cuD3D10CtxCreateOnDevice", "hipD3D10CtxCreateOnDevice"), + ("cuD3D10GetDirect3DDevice", "hipD3D10GetDirect3DDevice"), + ("cuD3D10MapResources", "hipD3D10MapResources"), + ("cuD3D10RegisterResource", "hipD3D10RegisterResource"), + ("cuD3D10ResourceGetMappedArray", "hipD3D10ResourceGetMappedArray"), + ("cuD3D10ResourceGetMappedPitch", "hipD3D10ResourceGetMappedPitch"), + ("cuD3D10ResourceGetMappedPointer", "hipD3D10ResourceGetMappedPointer"), + ("cuD3D10ResourceGetMappedSize", "hipD3D10ResourceGetMappedSize"), + ("cuD3D10ResourceGetSurfaceDimensions", "hipD3D10ResourceGetSurfaceDimensions"), + ("cuD310ResourceSetMapFlags", "hipD3D10ResourceSetMapFlags"), + ("cuD3D10UnmapResources", "hipD3D10UnmapResources"), + ("cuD3D10UnregisterResource", "hipD3D10UnregisterResource"), + ("CU_D3D11_DEVICE_LIST_ALL", "HIP_D3D11_DEVICE_LIST_ALL"), + ("CU_D3D11_DEVICE_LIST_CURRENT_FRAME", "HIP_D3D11_DEVICE_LIST_CURRENT_FRAME"), + ("CU_D3D11_DEVICE_LIST_NEXT_FRAME", "HIP_D3D11_DEVICE_LIST_NEXT_FRAME"), + ("cuD3D11GetDevice", "hipD3D11GetDevice"), + ("cuD3D11GetDevices", "hipD3D11GetDevices"), + ("cuGraphicsD3D11RegisterResource", "hipGraphicsD3D11RegisterResource"), + ("cuD3D11CtxCreate", "hipD3D11CtxCreate"), + ("cuD3D11CtxCreateOnDevice", "hipD3D11CtxCreateOnDevice"), + ("cuD3D11GetDirect3DDevice", "hipD3D11GetDirect3DDevice"), + ("cuGraphicsVDPAURegisterOutputSurface", "hipGraphicsVDPAURegisterOutputSurface"), + ("cuGraphicsVDPAURegisterVideoSurface", "hipGraphicsVDPAURegisterVideoSurface"), + ("cuVDPAUGetDevice", "hipVDPAUGetDevice"), + ("cuVDPAUCtxCreate", "hipVDPAUCtxCreate"), + ("cuEGLStreamConsumerAcquireFrame", "hipEGLStreamConsumerAcquireFrame"), + ("cuEGLStreamConsumerConnect", "hipEGLStreamConsumerConnect"), + ("cuEGLStreamConsumerConnectWithFlags", "hipEGLStreamConsumerConnectWithFlags"), + ("cuEGLStreamConsumerDisconnect", "hipEGLStreamConsumerDisconnect"), + ("cuEGLStreamConsumerReleaseFrame", "hipEGLStreamConsumerReleaseFrame"), + ("cuEGLStreamProducerConnect", "hipEGLStreamProducerConnect"), + ("cuEGLStreamProducerDisconnect", "hipEGLStreamProducerDisconnect"), + ("cuEGLStreamProducerPresentFrame", "hipEGLStreamProducerPresentFrame"), + ("cuEGLStreamProducerReturnFrame", "hipEGLStreamProducerReturnFrame"), + ("cuGraphicsEGLRegisterImage", "hipGraphicsEGLRegisterImage"), + ("cuGraphicsResourceGetMappedEglFrame", "hipGraphicsResourceGetMappedEglFrame"), + ("cudaDataType_t", "hipDataType"), + ("cudaDataType", "hipDataType"), + ("CUDA_R_32F", "HIP_R_32F"), + ("CUDA_R_64F", "HIP_R_64F"), + ("CUDA_R_16F", "HIP_R_16F"), + ("CUDA_R_8I", "HIP_R_8I"), + ("CUDA_C_32F", "HIP_C_32F"), + ("CUDA_C_64F", "HIP_C_64F"), + ("CUDA_C_16F", "HIP_C_16F"), + ("CUDA_C_8I", "HIP_C_8I"), + ("CUDA_R_8U", "HIP_R_8U"), + ("CUDA_C_8U", "HIP_C_8U"), + ("CUDA_R_32I", "HIP_R_32I"), + ("CUDA_C_32I", "HIP_C_32I"), + ("CUDA_R_32U", "HIP_R_32U"), + ("CUDA_C_32U", "HIP_C_32U"), + ("CUDA_R_16BF", "HIP_R_16BF"), + ("CUDA_C_16BF", "HIP_C_16BF"), + ("CUDA_R_4I", "HIP_R_4I"), + ("CUDA_C_4I", "HIP_C_4I"), + ("CUDA_R_4U", "HIP_R_4U"), + ("CUDA_C_4U", "HIP_C_4U"), + ("CUDA_R_16I", "HIP_R_16I"), + ("CUDA_C_16I", "HIP_C_16I"), + ("CUDA_R_16U", "HIP_R_16U"), + ("CUDA_C_16U", "HIP_C_16U"), + ("CUDA_R_64I", "HIP_R_64I"), + ("CUDA_C_64I", "HIP_C_64I"), + ("CUDA_R_64U", "HIP_R_64U"), + ("CUDA_C_64U", "HIP_C_64U"), + ("CUDA_R_8F_E4M3", "HIP_R_8F_E4M3"), + ("CUDA_R_8F_E5M2", "HIP_R_8F_E5M2"), + ("CUDA_R_4F_E2M1", "HIP_R_4F_E2M1"), + ("MAJOR_VERSION", "hipLibraryMajorVersion"), + ("MINOR_VERSION", "hipLibraryMinorVersion"), + ("PATCH_LEVEL", "hipLibraryPatchVersion"), + ("cudaLaunchKernel", "hipLaunchKernel"), + ("cudaMemAllocationHandleType", "hipMemAllocationHandleType"), + ("cudaMemAllocationType", "hipMemAllocationType"), + ("cudaMemLocationType", "hipMemLocationType"), + ("cudaMemAttachGlobal", "hipMemAttachGlobal"), + ("cudaMemAttachHost", "hipMemAttachHost"), + ("cudaMemAttachSingle", "hipMemAttachSingle"), + ("cudaOccupancyDefault", "hipOccupancyDefault"), + ("cudaOccupancyDisableCachingOverride", "hipOccupancyDisableCachingOverride"), + ("cudaGetLastError", "hipGetLastError"), + ("cudaPeekAtLastError", "hipPeekAtLastError"), + ("cudaGetErrorName", "hipGetErrorName"), + ("cudaGetErrorString", "hipGetErrorString"), + ("cudaMemcpy3DParms", "hipMemcpy3DParms"), + ("cudaMemcpy3DPeerParms", "hipMemcpy3DPeerParms"), + ("cudaMemcpy", "hipMemcpy"), + ("cudaMemcpyToArray", "hipMemcpyToArray"), + ("cudaMemcpyToSymbol", "hipMemcpyToSymbol"), + ("cudaMemcpyToSymbolAsync", "hipMemcpyToSymbolAsync"), + ("cudaMemcpyAsync", "hipMemcpyAsync"), + ("cudaMemcpy2D", "hipMemcpy2D"), + ("cudaMemcpy2DAsync", "hipMemcpy2DAsync"), + ("cudaMemcpy2DToArray", "hipMemcpy2DToArray"), + ("cudaMemcpy2DArrayToArray", "hipMemcpy2DArrayToArray"), + ("cudaMemcpy2DFromArray", "hipMemcpy2DFromArray"), + ("cudaMemcpy2DFromArrayAsync", "hipMemcpy2DFromArrayAsync"), + ("cudaMemcpy2DToArrayAsync", "hipMemcpy2DToArrayAsync"), + ("cudaMemcpy3D", "hipMemcpy3D"), + ("cudaMemcpy3DAsync", "hipMemcpy3DAsync"), + ("cudaMemcpy3DPeer", "hipMemcpy3DPeer"), + ("cudaMemcpy3DPeerAsync", "hipMemcpy3DPeerAsync"), + ("cudaMemcpyArrayToArray", "hipMemcpyArrayToArray"), + ("cudaMemcpyFromArrayAsync", "hipMemcpyFromArrayAsync"), + ("cudaMemcpyFromSymbol", "hipMemcpyFromSymbol"), + ("cudaMemcpyFromSymbolAsync", "hipMemcpyFromSymbolAsync"), + ("cudaMemAdvise", "hipMemAdvise"), + ("cudaMemRangeGetAttribute", "hipMemRangeGetAttribute"), + ("cudaMemRangeGetAttributes", "hipMemRangeGetAttributes"), + ("cudaMemAdviseSetReadMostly", "hipMemAdviseSetReadMostly"), + ("cudaMemAdviseUnsetReadMostly", "hipMemAdviseUnsetReadMostly"), + ("cudaMemAdviseSetPreferredLocation", "hipMemAdviseSetPreferredLocation"), + ("cudaMemAdviseUnsetPreferredLocation", "hipMemAdviseUnsetPreferredLocation"), + ("cudaMemAdviseSetAccessedBy", "hipMemAdviseSetAccessedBy"), + ("cudaMemAdviseUnsetAccessedBy", "hipMemAdviseUnsetAccessedBy"), + ("cudaMemRangeAttributeReadMostly", "hipMemRangeAttributeReadMostly"), + ("cudaMemRangeAttributePreferredLocation", "hipMemRangeAttributePreferredLocation"), + ("cudaMemRangeAttributeAccessedBy", "hipMemRangeAttributeAccessedBy"), + ("cudaMemRangeAttributeLastPrefetchLocation", "hipMemRangeAttributeLastPrefetchLocation"), + ("cudaMemcpyHostToHost", "hipMemcpyHostToHost"), + ("cudaMemcpyHostToDevice", "hipMemcpyHostToDevice"), + ("cudaMemcpyDeviceToHost", "hipMemcpyDeviceToHost"), + ("cudaMemcpyDeviceToDevice", "hipMemcpyDeviceToDevice"), + ("cudaMemcpyDefault", "hipMemcpyDefault"), + ("cudaMemset", "hipMemset"), + ("cudaMemsetAsync", "hipMemsetAsync"), + ("cudaMemset2D", "hipMemset2D"), + ("cudaMemset2DAsync", "hipMemset2DAsync"), + ("cudaMemset3D", "hipMemset3D"), + ("cudaMemset3DAsync", "hipMemset3DAsync"), + ("cudaMemGetInfo", "hipMemGetInfo"), + ("cudaDeviceGetDefaultMemPool", "hipDeviceGetDefaultMemPool"), + ("cudaMemAccessDesc", "hipMemAccessDesc"), + ("cudaMemAccessFlagsProtReadWrite", "hipMemAccessFlagsProtReadWrite"), + ("cudaMemLocationTypeDevice", "hipMemLocationTypeDevice"), + ("cudaMemPoolAttrReleaseThreshold", "hipMemPoolAttrReleaseThreshold"), + ("cudaMemPoolAttrReservedMemCurrent", "hipMemPoolAttrReservedMemCurrent"), + ("cudaMemPoolAttrReservedMemHigh", "hipMemPoolAttrReservedMemHigh"), + ("cudaMemPoolAttrUsedMemCurrent", "hipMemPoolAttrUsedMemCurrent"), + ("cudaMemPoolAttrUsedMemHigh", "hipMemPoolAttrUsedMemHigh"), + ("cudaMemPoolGetAttribute", "hipMemPoolGetAttribute"), + ("cudaMemPoolReuseAllowInternalDependencies", "hipMemPoolReuseAllowInternalDependencies"), + ("cudaMemPoolReuseAllowOpportunistic", "hipMemPoolReuseAllowOpportunistic"), + ("cudaMemPoolReuseFollowEventDependencies", "hipMemPoolReuseFollowEventDependencies"), + ("cudaMemPoolSetAccess", "hipMemPoolSetAccess"), + ("cudaMemPoolSetAttribute", "hipMemPoolSetAttribute"), + ("cudaMemPoolTrimTo", "hipMemPoolTrimTo"), + ("cudaMemPool_t", "hipMemPool_t"), + ("cudaArrayGetInfo", "hipArrayGetInfo"), + ("cudaFreeMipmappedArray", "hipFreeMipmappedArray"), + ("cudaGetMipmappedArrayLevel", "hipGetMipmappedArrayLevel"), + ("cudaGetSymbolAddress", "hipGetSymbolAddress"), + ("cudaGetSymbolSize", "hipGetSymbolSize"), + ("cudaMemPrefetchAsync", "hipMemPrefetchAsync"), + ("cudaMallocHost", "hipHostMalloc"), + ("cudaMallocArray", "hipMallocArray"), + ("cudaMalloc", "hipMalloc"), + ("cudaMalloc3D", "hipMalloc3D"), + ("cudaMalloc3DArray", "hipMalloc3DArray"), + ("cudaMallocAsync", "hipMallocAsync"), + ("cudaMallocManaged", "hipMallocManaged"), + ("cudaMallocMipmappedArray", "hipMallocMipmappedArray"), + ("cudaMallocPitch", "hipMallocPitch"), + ("cudaFreeHost", "hipHostFree"), + ("cudaFreeArray", "hipFreeArray"), + ("cudaFree", "hipFree"), + ("cudaFreeAsync", "hipFreeAsync"), + ("cudaHostRegister", "hipHostRegister"), + ("cudaHostUnregister", "hipHostUnregister"), + ("cudaHostAlloc", "hipHostMalloc"), + ("cudaMemoryTypeHost", "hipMemoryTypeHost"), + ("cudaMemoryTypeDevice", "hipMemoryTypeDevice"), + ("cudaMemoryTypeUnregistered", "hipMemoryTypeUnregistered"), + ("cudaMemoryTypeManaged", "hipMemoryTypeManaged"), + ("make_cudaExtent", "make_hipExtent"), + ("make_cudaPitchedPtr", "make_hipPitchedPtr"), + ("make_cudaPos", "make_hipPos"), + ("cudaHostAllocDefault", "hipHostMallocDefault"), + ("cudaHostAllocPortable", "hipHostMallocPortable"), + ("cudaHostAllocMapped", "hipHostMallocMapped"), + ("cudaHostNodeParams", "hipHostNodeParams"), + ("cudaHostAllocWriteCombined", "hipHostMallocWriteCombined"), + ("cudaHostFn_t", "hipHostFn_t"), + ("cudaHostGetFlags", "hipHostGetFlags"), + ("cudaHostRegisterDefault", "hipHostRegisterDefault"), + ("cudaHostRegisterPortable", "hipHostRegisterPortable"), + ("cudaHostRegisterMapped", "hipHostRegisterMapped"), + ("cudaHostRegisterIoMemory", "hipHostRegisterIoMemory"), + ("cudaEventCreate", "hipEventCreate"), + ("cudaEventCreateWithFlags", "hipEventCreateWithFlags"), + ("cudaEventDestroy", "hipEventDestroy"), + ("cudaEventRecord", "hipEventRecord"), + ("cudaEventElapsedTime", "hipEventElapsedTime"), + ("cudaEventSynchronize", "hipEventSynchronize"), + ("cudaEventQuery", "hipEventQuery"), + ("cudaEventDefault", "hipEventDefault"), + ("cudaEventBlockingSync", "hipEventBlockingSync"), + ("cudaEventDisableTiming", "hipEventDisableTiming"), + ("cudaEventInterprocess", "hipEventInterprocess"), + ("cudaEventRecordDefault", "hipEventRecordDefault"), + ("cudaEventRecordExternal", "hipEventRecordExternal"), + ("cudaEventWaitDefault", "hipEventWaitDefault"), + ("cudaEventWaitExternal", "hipEventWaitExternal"), + ("cudaStreamCreate", "hipStreamCreate"), + ("cudaStreamCreateWithFlags", "hipStreamCreateWithFlags"), + ("cudaStreamCreateWithPriority", "hipStreamCreateWithPriority"), + ("cudaStreamDestroy", "hipStreamDestroy"), + ("cudaStreamWaitEvent", "hipStreamWaitEvent"), + ("cudaStreamSynchronize", "hipStreamSynchronize"), + ("cudaStreamGetFlags", "hipStreamGetFlags"), + ("cudaStreamQuery", "hipStreamQuery"), + ("cudaStreamAddCallback", "hipStreamAddCallback"), + ("cudaStreamAttachMemAsync", "hipStreamAttachMemAsync"), + ("cudaStreamGetPriority", "hipStreamGetPriority"), + ("cudaCpuDeviceId", "hipCpuDeviceId"), + ("cudaStreamDefault", "hipStreamDefault"), + ("cudaStreamNonBlocking", "hipStreamNonBlocking"), + ("cudaStreamGetCaptureInfo", "hipStreamGetCaptureInfo"), + ("cudaStreamGetCaptureInfo_v2", "hipStreamGetCaptureInfo_v2"), + ("cudaStreamCaptureStatus", "hipStreamCaptureStatus"), + ("cudaStreamCaptureStatusActive", "hipStreamCaptureStatusActive"), + ("cudaStreamCaptureStatusNone", "hipStreamCaptureStatusNone"), + ("cudaStreamCaptureStatusInvalidated", "hipStreamCaptureStatusInvalidated"), + ("cudaStreamCaptureMode", "hipStreamCaptureMode"), + ("cudaStreamCaptureModeGlobal", "hipStreamCaptureModeGlobal"), + ("cudaStreamCaptureModeRelaxed", "hipStreamCaptureModeRelaxed"), + ("cudaStreamCaptureModeThreadLocal", "hipStreamCaptureModeThreadLocal"), + ("cudaStreamBeginCapture", "hipStreamBeginCapture"), + ("cudaStreamEndCapture", "hipStreamEndCapture"), + ("cudaStreamSetCaptureDependencies", "hipStreamSetCaptureDependencies"), + ("cudaStreamUpdateCaptureDependencies", "hipStreamUpdateCaptureDependencies"), + ("cudaGraphNode_t", "hipGraphNode_t"), + ("cudaGraphInstantiate", "hipGraphInstantiate"), + ("cudaGraphInstantiateWithFlags", "hipGraphInstantiateWithFlags"), + ("cudaGraphInstantiateFlagAutoFreeOnLaunch", "hipGraphInstantiateFlagAutoFreeOnLaunch"), + ("cudaGraphDestroy", "hipGraphDestroy"), + ("cudaGraphExecDestroy", "hipGraphExecDestroy"), + ("cudaGraphLaunch", "hipGraphLaunch"), + ("cudaGraphGetNodes", "hipGraphGetNodes"), + ("cudaGraphDebugDotPrint", "hipGraphDebugDotPrint"), + ("cudaGraphDebugDotFlagsVerbose", "hipGraphDebugDotFlagsVerbose"), + ("cudaGraphRetainUserObject", "hipGraphRetainUserObject"), + ("cudaGraphUserObjectMove", "hipGraphUserObjectMove"), + ("cudaDeviceGetGraphMemAttribute", "hipDeviceGetGraphMemAttribute"), + ("cudaDeviceGraphMemTrim", "hipDeviceGraphMemTrim"), + ("cudaDeviceSetGraphMemAttribute", "hipDeviceSetGraphMemAttribute"), + ("cudaGraphAddChildGraphNode", "hipGraphAddChildGraphNode"), + ("cudaGraphAddDependencies", "hipGraphAddDependencies"), + ("cudaGraphAddEmptyNode", "hipGraphAddEmptyNode"), + ("cudaGraphAddEventRecordNode", "hipGraphAddEventRecordNode"), + ("cudaGraphAddEventWaitNode", "hipGraphAddEventWaitNode"), + ("cudaGraphAddExternalSemaphoresSignalNode", "hipGraphAddExternalSemaphoresSignalNode"), + ("cudaGraphAddExternalSemaphoresWaitNode", "hipGraphAddExternalSemaphoresWaitNode"), + ("cudaGraphAddHostNode", "hipGraphAddHostNode"), + ("cudaGraphAddKernelNode", "hipGraphAddKernelNode"), + ("cudaGraphAddMemAllocNode", "hipGraphAddMemAllocNode"), + ("cudaGraphAddMemFreeNode", "hipGraphAddMemFreeNode"), + ("cudaGraphAddMemcpyNode", "hipGraphAddMemcpyNode"), + ("cudaGraphAddMemcpyNode1D", "hipGraphAddMemcpyNode1D"), + ("cudaGraphAddMemcpyNodeFromSymbol", "hipGraphAddMemcpyNodeFromSymbol"), + ("cudaGraphAddMemcpyNodeToSymbol", "hipGraphAddMemcpyNodeToSymbol"), + ("cudaGraphAddMemsetNode", "hipGraphAddMemsetNode"), + ("cudaGraphAddNode", "hipGraphAddNode"), + ("cudaGraphChildGraphNodeGetGraph", "hipGraphChildGraphNodeGetGraph"), + ("cudaGraphClone", "hipGraphClone"), + ("cudaGraphCreate", "hipGraphCreate"), + ("cudaGraphDestroyNode", "hipGraphDestroyNode"), + ("cudaGraphEventRecordNodeGetEvent", "hipGraphEventRecordNodeGetEvent"), + ("cudaGraphEventRecordNodeSetEvent", "hipGraphEventRecordNodeSetEvent"), + ("cudaGraphEventWaitNodeGetEvent", "hipGraphEventWaitNodeGetEvent"), + ("cudaGraphEventWaitNodeSetEvent", "hipGraphEventWaitNodeSetEvent"), + ("cudaGraphExecChildGraphNodeSetParams", "hipGraphExecChildGraphNodeSetParams"), + ("cudaGraphExecEventRecordNodeSetEvent", "hipGraphExecEventRecordNodeSetEvent"), + ("cudaGraphExecEventWaitNodeSetEvent", "hipGraphExecEventWaitNodeSetEvent"), + ("cudaGraphExecExternalSemaphoresSignalNodeSetParams", "hipGraphExecExternalSemaphoresSignalNodeSetParams"), + ("cudaGraphExecExternalSemaphoresWaitNodeSetParams", "hipGraphExecExternalSemaphoresWaitNodeSetParams"), + ("cudaGraphExecGetFlags", "hipGraphExecGetFlags"), + ("cudaGraphExecHostNodeSetParams", "hipGraphExecHostNodeSetParams"), + ("cudaGraphExecKernelNodeSetParams", "hipGraphExecKernelNodeSetParams"), + ("cudaGraphExecMemcpyNodeSetParams", "hipGraphExecMemcpyNodeSetParams"), + ("cudaGraphExecMemcpyNodeSetParams1D", "hipGraphExecMemcpyNodeSetParams1D"), + ("cudaGraphExecMemcpyNodeSetParamsFromSymbol", "hipGraphExecMemcpyNodeSetParamsFromSymbol"), + ("cudaGraphExecMemcpyNodeSetParamsToSymbol", "hipGraphExecMemcpyNodeSetParamsToSymbol"), + ("cudaGraphExecMemsetNodeSetParams", "hipGraphExecMemsetNodeSetParams"), + ("cudaGraphExecNodeSetParams", "hipGraphExecNodeSetParams"), + ("cudaGraphExecUpdate", "hipGraphExecUpdate"), + ("cudaGraphExternalSemaphoresSignalNodeGetParams", "hipGraphExternalSemaphoresSignalNodeGetParams"), + ("cudaGraphExternalSemaphoresSignalNodeSetParams", "hipGraphExternalSemaphoresSignalNodeSetParams"), + ("cudaGraphExternalSemaphoresWaitNodeGetParams", "hipGraphExternalSemaphoresWaitNodeGetParams"), + ("cudaGraphExternalSemaphoresWaitNodeSetParams", "hipGraphExternalSemaphoresWaitNodeSetParams"), + ("cudaGraphGetEdges", "hipGraphGetEdges"), + ("cudaGraphGetRootNodes", "hipGraphGetRootNodes"), + ("cudaGraphHostNodeGetParams", "hipGraphHostNodeGetParams"), + ("cudaGraphHostNodeSetParams", "hipGraphHostNodeSetParams"), + ("cudaGraphInstantiateWithParams", "hipGraphInstantiateWithParams"), + ("cudaGraphKernelNodeCopyAttributes", "hipGraphKernelNodeCopyAttributes"), + ("cudaGraphKernelNodeGetAttribute", "hipGraphKernelNodeGetAttribute"), + ("cudaGraphKernelNodeGetParams", "hipGraphKernelNodeGetParams"), + ("cudaGraphKernelNodeSetAttribute", "hipGraphKernelNodeSetAttribute"), + ("cudaGraphKernelNodeSetParams", "hipGraphKernelNodeSetParams"), + ("cudaGraphMemAllocNodeGetParams", "hipGraphMemAllocNodeGetParams"), + ("cudaGraphMemFreeNodeGetParams", "hipGraphMemFreeNodeGetParams"), + ("cudaGraphMemcpyNodeGetParams", "hipGraphMemcpyNodeGetParams"), + ("cudaGraphMemcpyNodeSetParams", "hipGraphMemcpyNodeSetParams"), + ("cudaGraphMemcpyNodeSetParams1D", "hipGraphMemcpyNodeSetParams1D"), + ("cudaGraphMemcpyNodeSetParamsFromSymbol", "hipGraphMemcpyNodeSetParamsFromSymbol"), + ("cudaGraphMemcpyNodeSetParamsToSymbol", "hipGraphMemcpyNodeSetParamsToSymbol"), + ("cudaGraphMemsetNodeGetParams", "hipGraphMemsetNodeGetParams"), + ("cudaGraphMemsetNodeSetParams", "hipGraphMemsetNodeSetParams"), + ("cudaGraphNodeFindInClone", "hipGraphNodeFindInClone"), + ("cudaGraphNodeGetDependencies", "hipGraphNodeGetDependencies"), + ("cudaGraphNodeGetDependentNodes", "hipGraphNodeGetDependentNodes"), + ("cudaGraphNodeGetEnabled", "hipGraphNodeGetEnabled"), + ("cudaGraphNodeGetType", "hipGraphNodeGetType"), + ("cudaGraphNodeSetEnabled", "hipGraphNodeSetEnabled"), + ("cudaGraphNodeSetParams", "hipGraphNodeSetParams"), + ("cudaGraphReleaseUserObject", "hipGraphReleaseUserObject"), + ("cudaGraphRemoveDependencies", "hipGraphRemoveDependencies"), + ("cudaGraphUpload", "hipGraphUpload"), + ("cudaUserObjectRelease", "hipUserObjectRelease"), + ("cudaUserObjectRetain", "hipUserObjectRetain"), + ("cudaGraphDebugDotFlags", "hipGraphDebugDotFlags"), + ("cudaGraphDebugDotFlagsEventNodeParams", "hipGraphDebugDotFlagsEventNodeParams"), + ("cudaGraphDebugDotFlagsExtSemasSignalNodeParams", "hipGraphDebugDotFlagsExtSemasSignalNodeParams"), + ("cudaGraphDebugDotFlagsExtSemasWaitNodeParams", "hipGraphDebugDotFlagsExtSemasWaitNodeParams"), + ("cudaGraphDebugDotFlagsHandles", "hipGraphDebugDotFlagsHandles"), + ("cudaGraphDebugDotFlagsHostNodeParams", "hipGraphDebugDotFlagsHostNodeParams"), + ("cudaGraphDebugDotFlagsKernelNodeAttributes", "hipGraphDebugDotFlagsKernelNodeAttributes"), + ("cudaGraphDebugDotFlagsKernelNodeParams", "hipGraphDebugDotFlagsKernelNodeParams"), + ("cudaGraphDebugDotFlagsMemcpyNodeParams", "hipGraphDebugDotFlagsMemcpyNodeParams"), + ("cudaGraphDebugDotFlagsMemsetNodeParams", "hipGraphDebugDotFlagsMemsetNodeParams"), + ("cudaGraphDependencyType", "hipGraphDependencyType"), + ("cudaGraphDependencyTypeDefault", "hipGraphDependencyTypeDefault"), + ("cudaGraphDependencyTypeProgrammatic", "hipGraphDependencyTypeProgrammatic"), + ("cudaGraphDependencyType_enum", "hipGraphDependencyType"), + ("cudaGraphEdgeData", "hipGraphEdgeData"), + ("cudaGraphEdgeData_st", "hipGraphEdgeData"), + ("cudaGraphExecUpdateError", "hipGraphExecUpdateError"), + ("cudaGraphExecUpdateErrorFunctionChanged", "hipGraphExecUpdateErrorFunctionChanged"), + ("cudaGraphExecUpdateErrorNodeTypeChanged", "hipGraphExecUpdateErrorNodeTypeChanged"), + ("cudaGraphExecUpdateErrorNotSupported", "hipGraphExecUpdateErrorNotSupported"), + ("cudaGraphExecUpdateErrorParametersChanged", "hipGraphExecUpdateErrorParametersChanged"), + ("cudaGraphExecUpdateErrorTopologyChanged", "hipGraphExecUpdateErrorTopologyChanged"), + ("cudaGraphExecUpdateErrorUnsupportedFunctionChange", "hipGraphExecUpdateErrorUnsupportedFunctionChange"), + ("cudaGraphExecUpdateResult", "hipGraphExecUpdateResult"), + ("cudaGraphExecUpdateSuccess", "hipGraphExecUpdateSuccess"), + ("cudaGraphInstantiateError", "hipGraphInstantiateError"), + ("cudaGraphInstantiateFlagDeviceLaunch", "hipGraphInstantiateFlagDeviceLaunch"), + ("cudaGraphInstantiateFlagUpload", "hipGraphInstantiateFlagUpload"), + ("cudaGraphInstantiateFlagUseNodePriority", "hipGraphInstantiateFlagUseNodePriority"), + ("cudaGraphInstantiateFlags", "hipGraphInstantiateFlags"), + ("cudaGraphInstantiateInvalidStructure", "hipGraphInstantiateInvalidStructure"), + ("cudaGraphInstantiateMultipleDevicesNotSupported", "hipGraphInstantiateMultipleDevicesNotSupported"), + ("cudaGraphInstantiateNodeOperationNotSupported", "hipGraphInstantiateNodeOperationNotSupported"), + ("cudaGraphInstantiateParams", "hipGraphInstantiateParams"), + ("cudaGraphInstantiateParams_st", "hipGraphInstantiateParams"), + ("cudaGraphInstantiateResult", "hipGraphInstantiateResult"), + ("cudaGraphInstantiateSuccess", "hipGraphInstantiateSuccess"), + ("cudaGraphKernelNodePortDefault", "hipGraphKernelNodePortDefault"), + ("cudaGraphKernelNodePortLaunchCompletion", "hipGraphKernelNodePortLaunchCompletion"), + ("cudaGraphKernelNodePortProgrammatic", "hipGraphKernelNodePortProgrammatic"), + ("cudaGraphMemAttrReservedMemCurrent", "hipGraphMemAttrReservedMemCurrent"), + ("cudaGraphMemAttrReservedMemHigh", "hipGraphMemAttrReservedMemHigh"), + ("cudaGraphMemAttrUsedMemCurrent", "hipGraphMemAttrUsedMemCurrent"), + ("cudaGraphMemAttrUsedMemHigh", "hipGraphMemAttrUsedMemHigh"), + ("cudaGraphMemAttributeType", "hipGraphMemAttributeType"), + ("cudaGraphNodeParams", "hipGraphNodeParams"), + ("cudaGraphNodeType", "hipGraphNodeType"), + ("cudaGraphNodeTypeConditional", "hipGraphNodeTypeConditional"), + ("cudaGraphNodeTypeCount", "hipGraphNodeTypeCount"), + ("cudaGraphNodeTypeEmpty", "hipGraphNodeTypeEmpty"), + ("cudaGraphNodeTypeEventRecord", "hipGraphNodeTypeEventRecord"), + ("cudaGraphNodeTypeExtSemaphoreSignal", "hipGraphNodeTypeExtSemaphoreSignal"), + ("cudaGraphNodeTypeExtSemaphoreWait", "hipGraphNodeTypeExtSemaphoreWait"), + ("cudaGraphNodeTypeGraph", "hipGraphNodeTypeGraph"), + ("cudaGraphNodeTypeHost", "hipGraphNodeTypeHost"), + ("cudaGraphNodeTypeKernel", "hipGraphNodeTypeKernel"), + ("cudaGraphNodeTypeMemAlloc", "hipGraphNodeTypeMemAlloc"), + ("cudaGraphNodeTypeMemFree", "hipGraphNodeTypeMemFree"), + ("cudaGraphNodeTypeMemcpy", "hipGraphNodeTypeMemcpy"), + ("cudaGraphNodeTypeMemset", "hipGraphNodeTypeMemset"), + ("cudaGraphNodeTypeWaitEvent", "hipGraphNodeTypeWaitEvent"), + ("cudaUserObject_t", "hipUserObject_t"), + ("cudaUserObjectCreate", "hipUserObjectCreate"), + ("cudaUserObjectNoDestructorSync", "hipUserObjectNoDestructorSync"), + ("cudaThreadExchangeStreamCaptureMode", "hipThreadExchangeStreamCaptureMode"), + ("cudaStreamIsCapturing", "hipStreamIsCapturing"), + ("cudaDeviceSynchronize", "hipDeviceSynchronize"), + ("cudaDeviceReset", "hipDeviceReset"), + ("cudaSetDevice", "hipSetDevice"), + ("cudaGetDevice", "hipGetDevice"), + ("cudaGetDeviceCount", "hipGetDeviceCount"), + ("cudaChooseDevice", "hipChooseDevice"), + ("cudaThreadExit", "hipDeviceReset"), + ("cudaThreadGetCacheConfig", "hipDeviceGetCacheConfig"), + ("cudaThreadGetLimit", "hipThreadGetLimit"), + ("cudaThreadSetCacheConfig", "hipDeviceSetCacheConfig"), + ("cudaThreadSetLimit", "hipThreadSetLimit"), + ("cudaThreadSynchronize", "hipDeviceSynchronize"), + ("cudaDeviceGetAttribute", "hipDeviceGetAttribute"), + ("cudaDevAttrMaxThreadsPerBlock", "hipDeviceAttributeMaxThreadsPerBlock"), + ("cudaDevAttrMaxBlocksPerMultiprocessor", "hipDeviceAttributeMaxBlocksPerMultiprocessor"), + ("cudaDevAttrMaxBlockDimX", "hipDeviceAttributeMaxBlockDimX"), + ("cudaDevAttrMaxBlockDimY", "hipDeviceAttributeMaxBlockDimY"), + ("cudaDevAttrMaxBlockDimZ", "hipDeviceAttributeMaxBlockDimZ"), + ("cudaDevAttrMaxGridDimX", "hipDeviceAttributeMaxGridDimX"), + ("cudaDevAttrMaxGridDimY", "hipDeviceAttributeMaxGridDimY"), + ("cudaDevAttrMaxGridDimZ", "hipDeviceAttributeMaxGridDimZ"), + ("cudaDevAttrMaxSharedMemoryPerBlock", "hipDeviceAttributeMaxSharedMemoryPerBlock"), + ("cudaDevAttrMaxSharedMemoryPerBlockOptin", "hipDeviceAttributeMaxSharedMemoryPerBlock"), + ("cudaDevAttrTotalConstantMemory", "hipDeviceAttributeTotalConstantMemory"), + ("cudaDevAttrWarpSize", "hipDeviceAttributeWarpSize"), + ("cudaDevAttrMaxPitch", "hipDeviceAttributeMaxPitch"), + ("cudaDevAttrMaxRegistersPerBlock", "hipDeviceAttributeMaxRegistersPerBlock"), + ("cudaDevAttrClockRate", "hipDeviceAttributeClockRate"), + ("cudaDevAttrTextureAlignment", "hipDeviceAttributeTextureAlignment"), + ("cudaDevAttrGpuOverlap", "hipDeviceAttributeGpuOverlap"), + ("cudaDevAttrMultiProcessorCount", "hipDeviceAttributeMultiprocessorCount"), + ("cudaDevAttrKernelExecTimeout", "hipDeviceAttributeKernelExecTimeout"), + ("cudaDevAttrIntegrated", "hipDeviceAttributeIntegrated"), + ("cudaDevAttrReserved94", "hipDeviceAttributeCanUseStreamWaitValue"), + ("cudaDevAttrCooperativeLaunch", "hipDeviceAttributeCooperativeLaunch"), + ("cudaDevAttrCooperativeMultiDeviceLaunch", "hipDeviceAttributeCooperativeMultiDeviceLaunch"), + ("cudaDevAttrCanMapHostMemory", "hipDeviceAttributeCanMapHostMemory"), + ("cudaDevAttrComputeMode", "hipDeviceAttributeComputeMode"), + ("cudaDevAttrMaxTexture1DWidth", "hipDeviceAttributeMaxTexture1DWidth"), + ("cudaDevAttrMaxTexture2DWidth", "hipDeviceAttributeMaxTexture2DWidth"), + ("cudaDevAttrMaxTexture2DHeight", "hipDeviceAttributeMaxTexture2DHeight"), + ("cudaDevAttrMaxTexture3DWidth", "hipDeviceAttributeMaxTexture3DWidth"), + ("cudaDevAttrMaxTexture3DHeight", "hipDeviceAttributeMaxTexture3DHeight"), + ("cudaDevAttrMaxTexture3DDepth", "hipDeviceAttributeMaxTexture3DDepth"), + ("cudaDevAttrMaxTexture2DLayeredWidth", "hipDeviceAttributeMaxTexture2DLayeredWidth"), + ("cudaDevAttrMaxTexture2DLayeredHeight", "hipDeviceAttributeMaxTexture2DLayeredHeight"), + ("cudaDevAttrMaxTexture2DLayeredLayers", "hipDeviceAttributeMaxTexture2DLayeredLayers"), + ("cudaDevAttrSurfaceAlignment", "hipDeviceAttributeSurfaceAlignment"), + ("cudaDevAttrConcurrentKernels", "hipDeviceAttributeConcurrentKernels"), + ("cudaDevAttrEccEnabled", "hipDeviceAttributeEccEnabled"), + ("cudaDevAttrMemoryPoolsSupported", "hipDeviceAttributeMemoryPoolsSupported"), + ("cudaDevAttrPciBusId", "hipDeviceAttributePciBusId"), + ("cudaDevAttrPciDeviceId", "hipDeviceAttributePciDeviceId"), + ("cudaDevAttrTccDriver", "hipDeviceAttributeTccDriver"), + ("cudaDevAttrMemoryClockRate", "hipDeviceAttributeMemoryClockRate"), + ("cudaDevAttrGlobalMemoryBusWidth", "hipDeviceAttributeMemoryBusWidth"), + ("cudaDevAttrL2CacheSize", "hipDeviceAttributeL2CacheSize"), + ("cudaDevAttrMaxThreadsPerMultiProcessor", "hipDeviceAttributeMaxThreadsPerMultiProcessor"), + ("cudaDevAttrAsyncEngineCount", "hipDeviceAttributeAsyncEngineCount"), + ("cudaDevAttrUnifiedAddressing", "hipDeviceAttributeUnifiedAddressing"), + ("cudaDevAttrMaxTexture1DLayeredWidth", "hipDeviceAttributeMaxTexture1DLayeredWidth"), + ("cudaDevAttrMaxTexture1DLayeredLayers", "hipDeviceAttributeMaxTexture1DLayeredLayers"), + ("cudaDevAttrMaxTexture2DGatherWidth", "hipDeviceAttributeMaxTexture2DGatherWidth"), + ("cudaDevAttrMaxTexture2DGatherHeight", "hipDeviceAttributeMaxTexture2DGatherHeight"), + ("cudaDevAttrMaxTexture3DWidthAlt", "hipDeviceAttributeMaxTexture3DWidthAlternate"), + ("cudaDevAttrMaxTexture3DHeightAlt", "hipDeviceAttributeMaxTexture3DHeightAlternate"), + ("cudaDevAttrMaxTexture3DDepthAlt", "hipDeviceAttributeMaxTexture3DDepthAlternate"), + ("cudaDevAttrPciDomainId", "hipDeviceAttributePciDomainId"), + ("cudaDevAttrTexturePitchAlignment", "hipDeviceAttributeTexturePitchAlignment"), + ("cudaDevAttrMaxTextureCubemapWidth", "hipDeviceAttributeMaxTextureCubemapWidth"), + ("cudaDevAttrMaxTextureCubemapLayeredWidth", "hipDeviceAttributeMaxTextureCubemapLayeredWidth"), + ("cudaDevAttrMaxTextureCubemapLayeredLayers", "hipDeviceAttributeMaxTextureCubemapLayeredLayers"), + ("cudaDevAttrMaxSurface1DWidth", "hipDeviceAttributeMaxSurface1DWidth"), + ("cudaDevAttrMaxSurface2DWidth", "hipDeviceAttributeMaxSurface2DWidth"), + ("cudaDevAttrMaxSurface2DHeight", "hipDeviceAttributeMaxSurface2DHeight"), + ("cudaDevAttrMaxSurface3DWidth", "hipDeviceAttributeMaxSurface3DWidth"), + ("cudaDevAttrMaxSurface3DHeight", "hipDeviceAttributeMaxSurface3DHeight"), + ("cudaDevAttrMaxSurface3DDepth", "hipDeviceAttributeMaxSurface3DDepth"), + ("cudaDevAttrMaxSurface1DLayeredWidth", "hipDeviceAttributeMaxSurface1DLayeredWidth"), + ("cudaDevAttrMaxSurface1DLayeredLayers", "hipDeviceAttributeMaxSurface1DLayeredLayers"), + ("cudaDevAttrMaxSurface2DLayeredWidth", "hipDeviceAttributeMaxSurface2DLayeredWidth"), + ("cudaDevAttrMaxSurface2DLayeredHeight", "hipDeviceAttributeMaxSurface2DLayeredHeight"), + ("cudaDevAttrMaxSurface2DLayeredLayers", "hipDeviceAttributeMaxSurface2DLayeredLayers"), + ("cudaDevAttrMaxSurfaceCubemapWidth", "hipDeviceAttributeMaxSurfaceCubemapWidth"), + ("cudaDevAttrMaxSurfaceCubemapLayeredWidth", "hipDeviceAttributeMaxSurfaceCubemapLayeredWidth"), + ("cudaDevAttrMaxSurfaceCubemapLayeredLayers", "hipDeviceAttributeMaxSurfaceCubemapLayeredLayers"), + ("cudaDevAttrMaxTexture1DLinearWidth", "hipDeviceAttributeMaxTexture1DLinearWidth"), + ("cudaDevAttrMaxTexture2DLinearWidth", "hipDeviceAttributeMaxTexture2DLinearWidth"), + ("cudaDevAttrMaxTexture2DLinearHeight", "hipDeviceAttributeMaxTexture2DLinearHeight"), + ("cudaDevAttrMaxTexture2DLinearPitch", "hipDeviceAttributeMaxTexture2DLinearPitch"), + ("cudaDevAttrMaxTexture2DMipmappedWidth", "hipDeviceAttributeMaxTexture2DMipmappedWidth"), + ("cudaDevAttrMaxTexture2DMipmappedHeight", "hipDeviceAttributeMaxTexture2DMipmappedHeight"), + ("cudaDevAttrComputeCapabilityMajor", "hipDeviceAttributeComputeCapabilityMajor"), + ("cudaDevAttrComputeCapabilityMinor", "hipDeviceAttributeComputeCapabilityMinor"), + ("cudaDevAttrMaxTexture1DMipmappedWidth", "hipDeviceAttributeMaxTexture1DMipmappedWidth"), + ("cudaDevAttrStreamPrioritiesSupported", "hipDeviceAttributeStreamPrioritiesSupported"), + ("cudaDevAttrGlobalL1CacheSupported", "hipDeviceAttributeGlobalL1CacheSupported"), + ("cudaDevAttrLocalL1CacheSupported", "hipDeviceAttributeLocalL1CacheSupported"), + ("cudaDevAttrMaxSharedMemoryPerMultiprocessor", "hipDeviceAttributeMaxSharedMemoryPerMultiprocessor"), + ("cudaDevAttrHostRegisterSupported", "hipDeviceAttributeHostRegisterSupported"), + ("cudaDevAttrMaxRegistersPerMultiprocessor", "hipDeviceAttributeMaxRegistersPerMultiprocessor"), + ("cudaDevAttrManagedMemory", "hipDeviceAttributeManagedMemory"), + ("cudaDevAttrDirectManagedMemAccessFromHost", "hipDeviceAttributeDirectManagedMemAccessFromHost"), + ("cudaDevAttrIsMultiGpuBoard", "hipDeviceAttributeIsMultiGpuBoard"), + ("cudaDevAttrMultiGpuBoardGroupID", "hipDeviceAttributeMultiGpuBoardGroupID"), + ("cudaDevAttrHostNativeAtomicSupported", "hipDeviceAttributeHostNativeAtomicSupported"), + ("cudaDevAttrSingleToDoublePrecisionPerfRatio", "hipDeviceAttributeSingleToDoublePrecisionPerfRatio"), + ("cudaDevAttrPageableMemoryAccess", "hipDeviceAttributePageableMemoryAccess"), + ("cudaDevAttrPageableMemoryAccessUsesHostPageTables", "hipDeviceAttributePageableMemoryAccessUsesHostPageTables"), + ("cudaDevAttrConcurrentManagedAccess", "hipDeviceAttributeConcurrentManagedAccess"), + ("cudaDevAttrComputePreemptionSupported", "hipDeviceAttributeComputePreemptionSupported"), + ("cudaDevAttrCanUseHostPointerForRegisteredMem", "hipDeviceAttributeCanUseHostPointerForRegisteredMem"), + ("cudaPointerGetAttributes", "hipPointerGetAttributes"), + ("cudaHostGetDevicePointer", "hipHostGetDevicePointer"), + ("cudaGetDeviceProperties", "hipGetDeviceProperties"), + ("cudaDeviceGetPCIBusId", "hipDeviceGetPCIBusId"), + ("cudaDeviceGetByPCIBusId", "hipDeviceGetByPCIBusId"), + ("cudaDeviceGetStreamPriorityRange", "hipDeviceGetStreamPriorityRange"), + ("cudaSetValidDevices", "hipSetValidDevices"), + ("cudaDevP2PAttrPerformanceRank", "hipDeviceP2PAttributePerformanceRank"), + ("cudaDevP2PAttrAccessSupported", "hipDeviceP2PAttributeAccessSupported"), + ("cudaDevP2PAttrNativeAtomicSupported", "hipDeviceP2PAttributeNativeAtomicSupported"), + ("cudaDeviceGetP2PAttribute", "hipDeviceGetP2PAttribute"), + ("cudaComputeModeDefault", "hipComputeModeDefault"), + ("cudaComputeModeExclusive", "hipComputeModeExclusive"), + ("cudaComputeModeProhibited", "hipComputeModeProhibited"), + ("cudaComputeModeExclusiveProcess", "hipComputeModeExclusiveProcess"), + ("cudaGetDeviceFlags", "hipGetDeviceFlags"), + ("cudaSetDeviceFlags", "hipSetDeviceFlags"), + ("cudaDeviceScheduleAuto", "hipDeviceScheduleAuto"), + ("cudaDeviceScheduleSpin", "hipDeviceScheduleSpin"), + ("cudaDeviceScheduleYield", "hipDeviceScheduleYield"), + ("cudaDeviceBlockingSync", "hipDeviceScheduleBlockingSync"), + ("cudaDeviceScheduleBlockingSync", "hipDeviceScheduleBlockingSync"), + ("cudaDeviceScheduleMask", "hipDeviceScheduleMask"), + ("cudaDeviceMapHost", "hipDeviceMapHost"), + ("cudaDeviceLmemResizeToMax", "hipDeviceLmemResizeToMax"), + ("cudaDeviceMask", "hipDeviceMask"), + ("cudaDeviceSetCacheConfig", "hipDeviceSetCacheConfig"), + ("cudaDeviceGetCacheConfig", "hipDeviceGetCacheConfig"), + ("cudaFuncAttributes", "hipFuncAttributes"), + ("cudaFuncAttributeMaxDynamicSharedMemorySize", "hipFuncAttributeMaxDynamicSharedMemorySize"), + ("cudaFuncAttributePreferredSharedMemoryCarveout", "hipFuncAttributePreferredSharedMemoryCarveout"), + ("cudaFuncSetAttribute", "hipFuncSetAttribute"), + ("cudaFuncSetCacheConfig", "hipFuncSetCacheConfig"), + ("cudaFuncCachePreferNone", "hipFuncCachePreferNone"), + ("cudaFuncCachePreferShared", "hipFuncCachePreferShared"), + ("cudaFuncCachePreferL1", "hipFuncCachePreferL1"), + ("cudaFuncCachePreferEqual", "hipFuncCachePreferEqual"), + ("cudaFuncGetAttributes", "hipFuncGetAttributes"), + ("cudaFuncSetSharedMemConfig", "hipFuncSetSharedMemConfig"), + ("cudaGetParameterBuffer", "hipGetParameterBuffer"), + ("cudaSetDoubleForDevice", "hipSetDoubleForDevice"), + ("cudaSetDoubleForHost", "hipSetDoubleForHost"), + ("cudaConfigureCall", "hipConfigureCall"), + ("cudaLaunch", "hipLaunch"), + ("cudaLaunchCooperativeKernel", "hipLaunchCooperativeKernel"), + ("cudaLaunchHostFunc", "hipLaunchHostFunc"), + ("cudaSetupArgument", "hipSetupArgument"), + ("cudaDriverGetVersion", "hipDriverGetVersion"), + ("cudaRuntimeGetVersion", "hipRuntimeGetVersion"), + ("cudaOccupancyMaxPotentialBlockSize", "hipOccupancyMaxPotentialBlockSize"), + ("cudaErrorContextIsDestroyed", "hipErrorContextIsDestroyed"), + ("cudaDeviceSetLimit", "hipDeviceSetLimit"), + ("cudaMallocFromPoolAsync", "hipMallocFromPoolAsync"), + ("cudaDeviceGetMemPool", "hipDeviceGetMemPool"), + ("cudaDeviceSetMemPool", "hipDeviceSetMemPool"), + ("cudaMemPoolAttr", "hipMemPoolAttr"), + ("cudaMemPoolProps", "hipMemPoolProps"), + ("cudaMemAllocationTypePinned", "hipMemAllocationTypePinned"), + ("cudaMemHandleTypeNone", "hipMemHandleTypeNone"), + ("cudaMemHandleTypePosixFileDescriptor", "hipMemHandleTypePosixFileDescriptor"), + ("cudaMemLocation", "hipMemLocation"), + ("cudaMemPoolCreate", "hipMemPoolCreate"), + ("cudaMemPoolDestroy", "hipMemPoolDestroy"), + ("cudaCreateSurfaceObject", "hipCreateSurfaceObject"), + ("cudaDestroySurfaceObject", "hipDestroySurfaceObject"), + ("cudaGraphUpload", "hipGraphUpload"), + ("cudaOccupancyMaxPotentialBlockSizeWithFlags", "hipOccupancyMaxPotentialBlockSizeWithFlags"), + ("cudaOccupancyMaxActiveBlocksPerMultiprocessor", "hipOccupancyMaxActiveBlocksPerMultiprocessor"), + ("cudaOccupancyMaxActiveBlocksPerMultiprocessorWithFlags", "hipOccupancyMaxActiveBlocksPerMultiprocessorWithFlags"), + ("cudaOccupancyMaxPotentialBlockSizeVariableSMem", "hipOccupancyMaxPotentialBlockSizeVariableSMem"), + ("cudaOccupancyMaxPotentialBlockSizeVariableSMemWithFlags", "hipOccupancyMaxPotentialBlockSizeVariableSMemWithFlags"), + ("cudaDeviceCanAccessPeer", "hipDeviceCanAccessPeer"), + ("cudaDeviceDisablePeerAccess", "hipDeviceDisablePeerAccess"), + ("cudaDeviceEnablePeerAccess", "hipDeviceEnablePeerAccess"), + ("cudaMemcpyPeerAsync", "hipMemcpyPeerAsync"), + ("cudaMemcpyPeer", "hipMemcpyPeer"), + ("cudaIpcMemLazyEnablePeerAccess", "hipIpcMemLazyEnablePeerAccess"), + ("cudaDeviceSetSharedMemConfig", "hipDeviceSetSharedMemConfig"), + ("cudaDeviceGetSharedMemConfig", "hipDeviceGetSharedMemConfig"), + ("cudaSharedMemBankSizeDefault", "hipSharedMemBankSizeDefault"), + ("cudaSharedMemBankSizeFourByte", "hipSharedMemBankSizeFourByte"), + ("cudaSharedMemBankSizeEightByte", "hipSharedMemBankSizeEightByte"), + ("cudaLimitStackSize", "hipLimitStackSize"), + ("cudaLimitPrintfFifoSize", "hipLimitPrintfFifoSize"), + ("cudaLimitMallocHeapSize", "hipLimitMallocHeapSize"), + ("cudaLimitDevRuntimeSyncDepth", "hipLimitDevRuntimeSyncDepth"), + ("cudaLimitDevRuntimePendingLaunchCount", "hipLimitDevRuntimePendingLaunchCount"), + ("cudaDeviceGetLimit", "hipDeviceGetLimit"), + ("cudaProfilerInitialize", "hipProfilerInitialize"), + ("cudaProfilerStart", "hipProfilerStart"), + ("cudaProfilerStop", "hipProfilerStop"), + ("cudaKeyValuePair", "hipKeyValuePair"), + ("cudaCSV", "hipCSV"), + ("cudaReadModeElementType", "hipReadModeElementType"), + ("cudaReadModeNormalizedFloat", "hipReadModeNormalizedFloat"), + ("cudaFilterModePoint", "hipFilterModePoint"), + ("cudaFilterModeLinear", "hipFilterModeLinear"), + ("cudaBindTexture", "hipBindTexture"), + ("cudaUnbindTexture", "hipUnbindTexture"), + ("cudaBindTexture2D", "hipBindTexture2D"), + ("cudaBindTextureToArray", "hipBindTextureToArray"), + ("cudaBindTextureToMipmappedArray", "hipBindTextureToMipmappedArray"), + ("cudaGetTextureAlignmentOffset", "hipGetTextureAlignmentOffset"), + ("cudaGetTextureReference", "hipGetTextureReference"), + ("cudaChannelFormatKindSigned", "hipChannelFormatKindSigned"), + ("cudaChannelFormatKindUnsigned", "hipChannelFormatKindUnsigned"), + ("cudaChannelFormatKindFloat", "hipChannelFormatKindFloat"), + ("cudaChannelFormatKindNone", "hipChannelFormatKindNone"), + ("cudaCreateChannelDesc", "hipCreateChannelDesc"), + ("cudaGetChannelDesc", "hipGetChannelDesc"), + ("cudaResourceTypeArray", "hipResourceTypeArray"), + ("cudaResourceTypeMipmappedArray", "hipResourceTypeMipmappedArray"), + ("cudaResourceTypeLinear", "hipResourceTypeLinear"), + ("cudaResourceTypePitch2D", "hipResourceTypePitch2D"), + ("cudaResViewFormatNone", "hipResViewFormatNone"), + ("cudaResViewFormatUnsignedChar1", "hipResViewFormatUnsignedChar1"), + ("cudaResViewFormatUnsignedChar2", "hipResViewFormatUnsignedChar2"), + ("cudaResViewFormatUnsignedChar4", "hipResViewFormatUnsignedChar4"), + ("cudaResViewFormatSignedChar1", "hipResViewFormatSignedChar1"), + ("cudaResViewFormatSignedChar2", "hipResViewFormatSignedChar2"), + ("cudaResViewFormatSignedChar4", "hipResViewFormatSignedChar4"), + ("cudaResViewFormatUnsignedShort1", "hipResViewFormatUnsignedShort1"), + ("cudaResViewFormatUnsignedShort2", "hipResViewFormatUnsignedShort2"), + ("cudaResViewFormatUnsignedShort4", "hipResViewFormatUnsignedShort4"), + ("cudaResViewFormatSignedShort1", "hipResViewFormatSignedShort1"), + ("cudaResViewFormatSignedShort2", "hipResViewFormatSignedShort2"), + ("cudaResViewFormatSignedShort4", "hipResViewFormatSignedShort4"), + ("cudaResViewFormatUnsignedInt1", "hipResViewFormatUnsignedInt1"), + ("cudaResViewFormatUnsignedInt2", "hipResViewFormatUnsignedInt2"), + ("cudaResViewFormatUnsignedInt4", "hipResViewFormatUnsignedInt4"), + ("cudaResViewFormatSignedInt1", "hipResViewFormatSignedInt1"), + ("cudaResViewFormatSignedInt2", "hipResViewFormatSignedInt2"), + ("cudaResViewFormatSignedInt4", "hipResViewFormatSignedInt4"), + ("cudaResViewFormatHalf1", "hipResViewFormatHalf1"), + ("cudaResViewFormatHalf2", "hipResViewFormatHalf2"), + ("cudaResViewFormatHalf4", "hipResViewFormatHalf4"), + ("cudaResViewFormatFloat1", "hipResViewFormatFloat1"), + ("cudaResViewFormatFloat2", "hipResViewFormatFloat2"), + ("cudaResViewFormatFloat4", "hipResViewFormatFloat4"), + ("cudaResViewFormatUnsignedBlockCompressed1", "hipResViewFormatUnsignedBlockCompressed1"), + ("cudaResViewFormatUnsignedBlockCompressed2", "hipResViewFormatUnsignedBlockCompressed2"), + ("cudaResViewFormatUnsignedBlockCompressed3", "hipResViewFormatUnsignedBlockCompressed3"), + ("cudaResViewFormatUnsignedBlockCompressed4", "hipResViewFormatUnsignedBlockCompressed4"), + ("cudaResViewFormatSignedBlockCompressed4", "hipResViewFormatSignedBlockCompressed4"), + ("cudaResViewFormatUnsignedBlockCompressed5", "hipResViewFormatUnsignedBlockCompressed5"), + ("cudaResViewFormatSignedBlockCompressed5", "hipResViewFormatSignedBlockCompressed5"), + ("cudaResViewFormatUnsignedBlockCompressed6H", "hipResViewFormatUnsignedBlockCompressed6H"), + ("cudaResViewFormatSignedBlockCompressed6H", "hipResViewFormatSignedBlockCompressed6H"), + ("cudaResViewFormatUnsignedBlockCompressed7", "hipResViewFormatUnsignedBlockCompressed7"), + ("cudaAddressModeWrap", "hipAddressModeWrap"), + ("cudaAddressModeClamp", "hipAddressModeClamp"), + ("cudaAddressModeMirror", "hipAddressModeMirror"), + ("cudaAddressModeBorder", "hipAddressModeBorder"), + ("cudaCreateTextureObject", "hipCreateTextureObject"), + ("cudaDestroyTextureObject", "hipDestroyTextureObject"), + ("cudaGetTextureObjectResourceDesc", "hipGetTextureObjectResourceDesc"), + ("cudaGetTextureObjectResourceViewDesc", "hipGetTextureObjectResourceViewDesc"), + ("cudaGetTextureObjectTextureDesc", "hipGetTextureObjectTextureDesc"), + ("cudaBindSurfaceToArray", "hipBindSurfaceToArray"), + ("cudaGetSurfaceReference", "hipGetSurfaceReference"), + ("cudaBoundaryModeZero", "hipBoundaryModeZero"), + ("cudaBoundaryModeClamp", "hipBoundaryModeClamp"), + ("cudaBoundaryModeTrap", "hipBoundaryModeTrap"), + ("cudaFormatModeForced", "hipFormatModeForced"), + ("cudaFormatModeAuto", "hipFormatModeAuto"), + ("cudaGetSurfaceObjectResourceDesc", "hipGetSurfaceObjectResourceDesc"), + ("cudaIpcCloseMemHandle", "hipIpcCloseMemHandle"), + ("cudaIpcGetEventHandle", "hipIpcGetEventHandle"), + ("cudaIpcGetMemHandle", "hipIpcGetMemHandle"), + ("cudaIpcOpenEventHandle", "hipIpcOpenEventHandle"), + ("cudaIpcOpenMemHandle", "hipIpcOpenMemHandle"), + ("cudaGLGetDevices", "hipGLGetDevices"), + ("cudaGraphicsGLRegisterBuffer", "hipGraphicsGLRegisterBuffer"), + ("cudaGraphicsGLRegisterImage", "hipGraphicsGLRegisterImage"), + ("cudaWGLGetDevice", "hipWGLGetDevice"), + ("cudaGraphicsMapResources", "hipGraphicsMapResources"), + ("cudaGraphicsResourceGetMappedMipmappedArray", "hipGraphicsResourceGetMappedMipmappedArray"), + ("cudaGraphicsResourceGetMappedPointer", "hipGraphicsResourceGetMappedPointer"), + ("cudaGraphicsResourceSetMapFlags", "hipGraphicsResourceSetMapFlags"), + ("cudaGraphicsSubResourceGetMappedArray", "hipGraphicsSubResourceGetMappedArray"), + ("cudaGraphicsUnmapResources", "hipGraphicsUnmapResources"), + ("cudaGraphicsUnregisterResource", "hipGraphicsUnregisterResource"), + ("cudaGraphicsCubeFacePositiveX", "hipGraphicsCubeFacePositiveX"), + ("cudaGraphicsCubeFaceNegativeX", "hipGraphicsCubeFaceNegativeX"), + ("cudaGraphicsCubeFacePositiveY", "hipGraphicsCubeFacePositiveY"), + ("cudaGraphicsCubeFaceNegativeY", "hipGraphicsCubeFaceNegativeY"), + ("cudaGraphicsCubeFacePositiveZ", "hipGraphicsCubeFacePositiveZ"), + ("cudaGraphicsCubeFaceNegativeZ", "hipGraphicsCubeFaceNegativeZ"), + ("cudaGraphicsMapFlagsNone", "hipGraphicsMapFlagsNone"), + ("cudaGraphicsMapFlagsReadOnly", "hipGraphicsMapFlagsReadOnly"), + ("cudaGraphicsMapFlagsWriteDiscard", "hipGraphicsMapFlagsWriteDiscard"), + ("cudaGraphicsRegisterFlagsNone", "hipGraphicsRegisterFlagsNone"), + ("cudaGraphicsRegisterFlagsReadOnly", "hipGraphicsRegisterFlagsReadOnly"), + ("cudaGraphicsRegisterFlagsWriteDiscard", "hipGraphicsRegisterFlagsWriteDiscard"), + ("cudaGraphicsRegisterFlagsSurfaceLoadStore", "hipGraphicsRegisterFlagsSurfaceLoadStore"), + ("cudaGraphicsRegisterFlagsTextureGather", "hipGraphicsRegisterFlagsTextureGather"), + ("cudaGLDeviceListAll", "HIP_GL_DEVICE_LIST_ALL"), + ("cudaGLDeviceListCurrentFrame", "HIP_GL_DEVICE_LIST_CURRENT_FRAME"), + ("cudaGLDeviceListNextFrame", "HIP_GL_DEVICE_LIST_NEXT_FRAME"), + ("cudaGLMapFlagsNone", "HIP_GL_MAP_RESOURCE_FLAGS_NONE"), + ("cudaGLMapFlagsReadOnly", "HIP_GL_MAP_RESOURCE_FLAGS_READ_ONLY"), + ("cudaGLMapFlagsWriteDiscard", "HIP_GL_MAP_RESOURCE_FLAGS_WRITE_DISCARD"), + ("cudaGLMapBufferObject", "hipGLMapBufferObject__"), + ("cudaGLMapBufferObjectAsync", "hipGLMapBufferObjectAsync__"), + ("cudaGLRegisterBufferObject", "hipGLRegisterBufferObject"), + ("cudaGLSetBufferObjectMapFlags", "hipGLSetBufferObjectMapFlags"), + ("cudaGLSetGLDevice", "hipGLSetGLDevice"), + ("cudaGLUnmapBufferObject", "hipGLUnmapBufferObject"), + ("cudaGLUnmapBufferObjectAsync", "hipGLUnmapBufferObjectAsync"), + ("cudaGLUnregisterBufferObject", "hipGLUnregisterBufferObject"), + ("cudaD3D9DeviceListAll", "HIP_D3D9_DEVICE_LIST_ALL"), + ("cudaD3D9DeviceListCurrentFrame", "HIP_D3D9_DEVICE_LIST_CURRENT_FRAME"), + ("cudaD3D9DeviceListNextFrame", "HIP_D3D9_DEVICE_LIST_NEXT_FRAME"), + ("cudaD3D9GetDevice", "hipD3D9GetDevice"), + ("cudaD3D9GetDevices", "hipD3D9GetDevices"), + ("cudaD3D9GetDirect3DDevice", "hipD3D9GetDirect3DDevice"), + ("cudaD3D9SetDirect3DDevice", "hipD3D9SetDirect3DDevice"), + ("cudaGraphicsD3D9RegisterResource", "hipGraphicsD3D9RegisterResource"), + ("cudaD3D9MapFlags", "hipD3D9MapFlags"), + ("cudaD3D9MapFlagsNone", "HIP_D3D9_MAPRESOURCE_FLAGS_NONE"), + ("cudaD3D9MapFlagsReadOnly", "HIP_D3D9_MAPRESOURCE_FLAGS_READONLY"), + ("cudaD3D9MapFlagsWriteDiscard", "HIP_D3D9_MAPRESOURCE_FLAGS_WRITEDISCARD"), + ("cudaD3D9RegisterFlagsNone", "HIP_D3D9_REGISTER_FLAGS_NONE"), + ("cudaD3D9RegisterFlagsArray", "HIP_D3D9_REGISTER_FLAGS_ARRAY"), + ("cudaD3D9MapResources", "hipD3D9MapResources"), + ("cudaD3D9RegisterResource", "hipD3D9RegisterResource"), + ("cudaD3D9ResourceGetMappedArray", "hipD3D9ResourceGetMappedArray"), + ("cudaD3D9ResourceGetMappedPitch", "hipD3D9ResourceGetMappedPitch"), + ("cudaD3D9ResourceGetMappedPointer", "hipD3D9ResourceGetMappedPointer"), + ("cudaD3D9ResourceGetMappedSize", "hipD3D9ResourceGetMappedSize"), + ("cudaD3D9ResourceGetSurfaceDimensions", "hipD3D9ResourceGetSurfaceDimensions"), + ("cudaD3D9ResourceSetMapFlags", "hipD3D9ResourceSetMapFlags"), + ("cudaD3D9UnmapResources", "hipD3D9UnmapResources"), + ("cudaD3D9UnregisterResource", "hipD3D9UnregisterResource"), + ("cudaD3D10DeviceListAll", "HIP_D3D10_DEVICE_LIST_ALL"), + ("cudaD3D10DeviceListCurrentFrame", "HIP_D3D10_DEVICE_LIST_CURRENT_FRAME"), + ("cudaD3D10DeviceListNextFrame", "HIP_D3D10_DEVICE_LIST_NEXT_FRAME"), + ("cudaD3D10GetDevice", "hipD3D10GetDevice"), + ("cudaD3D10GetDevices", "hipD3D10GetDevices"), + ("cudaGraphicsD3D10RegisterResource", "hipGraphicsD3D10RegisterResource"), + ("cudaD3D10MapFlagsNone", "HIP_D3D10_MAPRESOURCE_FLAGS_NONE"), + ("cudaD3D10MapFlagsReadOnly", "HIP_D3D10_MAPRESOURCE_FLAGS_READONLY"), + ("cudaD3D10MapFlagsWriteDiscard", "HIP_D3D10_MAPRESOURCE_FLAGS_WRITEDISCARD"), + ("cudaD3D10RegisterFlagsNone", "HIP_D3D10_REGISTER_FLAGS_NONE"), + ("cudaD3D10RegisterFlagsArray", "HIP_D3D10_REGISTER_FLAGS_ARRAY"), + ("cudaD3D10GetDirect3DDevice", "hipD3D10GetDirect3DDevice"), + ("cudaD3D10MapResources", "hipD3D10MapResources"), + ("cudaD3D10RegisterResource", "hipD3D10RegisterResource"), + ("cudaD3D10ResourceGetMappedArray", "hipD3D10ResourceGetMappedArray"), + ("cudaD3D10ResourceGetMappedPitch", "hipD3D10ResourceGetMappedPitch"), + ("cudaD3D10ResourceGetMappedPointer", "hipD3D10ResourceGetMappedPointer"), + ("cudaD3D10ResourceGetMappedSize", "hipD3D10ResourceGetMappedSize"), + ("cudaD3D10ResourceGetSurfaceDimensions", "hipD3D10ResourceGetSurfaceDimensions"), + ("cudaD3D10ResourceSetMapFlags", "hipD3D10ResourceSetMapFlags"), + ("cudaD3D10SetDirect3DDevice", "hipD3D10SetDirect3DDevice"), + ("cudaD3D10UnmapResources", "hipD3D10UnmapResources"), + ("cudaD3D10UnregisterResource", "hipD3D10UnregisterResource"), + ("cudaD3D11DeviceListAll", "HIP_D3D11_DEVICE_LIST_ALL"), + ("cudaD3D11DeviceListCurrentFrame", "HIP_D3D11_DEVICE_LIST_CURRENT_FRAME"), + ("cudaD3D11DeviceListNextFrame", "HIP_D3D11_DEVICE_LIST_NEXT_FRAME"), + ("cudaD3D11GetDevice", "hipD3D11GetDevice"), + ("cudaD3D11GetDevices", "hipD3D11GetDevices"), + ("cudaGraphicsD3D11RegisterResource", "hipGraphicsD3D11RegisterResource"), + ("cudaGraphicsVDPAURegisterOutputSurface", "hipGraphicsVDPAURegisterOutputSurface"), + ("cudaGraphicsVDPAURegisterVideoSurface", "hipGraphicsVDPAURegisterVideoSurface"), + ("cudaVDPAUGetDevice", "hipVDPAUGetDevice"), + ("cudaVDPAUSetVDPAUDevice", "hipVDPAUSetDevice"), + ("cudaEGLStreamConsumerAcquireFrame", "hipEGLStreamConsumerAcquireFrame"), + ("cudaEGLStreamConsumerConnect", "hipEGLStreamConsumerConnect"), + ("cudaEGLStreamConsumerConnectWithFlags", "hipEGLStreamConsumerConnectWithFlags"), + ("cudaEGLStreamConsumerReleaseFrame", "hipEGLStreamConsumerReleaseFrame"), + ("cudaEGLStreamProducerConnect", "hipEGLStreamProducerConnect"), + ("cudaEGLStreamProducerDisconnect", "hipEGLStreamProducerDisconnect"), + ("cudaEGLStreamProducerPresentFrame", "hipEGLStreamProducerPresentFrame"), + ("cudaEGLStreamProducerReturnFrame", "hipEGLStreamProducerReturnFrame"), + ("cudaGraphicsEGLRegisterImage", "hipGraphicsEGLRegisterImage"), + ("cudaGraphicsResourceGetMappedEglFrame", "hipGraphicsResourceGetMappedEglFrame"), + ("cublasInit", "hipblasInit"), + ("cublasShutdown", "hipblasShutdown"), + ("cublasGetVersion", "hipblasGetVersion"), + ("cublasGetError", "hipblasGetError"), + ("cublasAlloc", "hipblasAlloc"), + ("cublasFree", "hipblasFree"), + ("cublasSetKernelStream", "hipblasSetKernelStream"), + ("cublasGetAtomicsMode", "hipblasGetAtomicsMode"), + ("cublasSetAtomicsMode", "hipblasSetAtomicsMode"), + ("cublasGetMathMode", "hipblasGetMathMode"), + ("cublasSetMathMode", "hipblasSetMathMode"), + ("CUBLAS_OP_N", "HIPBLAS_OP_N"), + ("CUBLAS_OP_T", "HIPBLAS_OP_T"), + ("CUBLAS_OP_C", "HIPBLAS_OP_C"), + ("CUBLAS_STATUS_SUCCESS", "HIPBLAS_STATUS_SUCCESS"), + ("CUBLAS_STATUS_NOT_INITIALIZED", "HIPBLAS_STATUS_NOT_INITIALIZED"), + ("CUBLAS_STATUS_ALLOC_FAILED", "HIPBLAS_STATUS_ALLOC_FAILED"), + ("CUBLAS_STATUS_INVALID_VALUE", "HIPBLAS_STATUS_INVALID_VALUE"), + ("CUBLAS_STATUS_MAPPING_ERROR", "HIPBLAS_STATUS_MAPPING_ERROR"), + ("CUBLAS_STATUS_EXECUTION_FAILED", "HIPBLAS_STATUS_EXECUTION_FAILED"), + ("CUBLAS_STATUS_INTERNAL_ERROR", "HIPBLAS_STATUS_INTERNAL_ERROR"), + ("CUBLAS_STATUS_NOT_SUPPORTED", "HIPBLAS_STATUS_NOT_SUPPORTED"), + ("CUBLAS_STATUS_ARCH_MISMATCH", "HIPBLAS_STATUS_ARCH_MISMATCH"), + ("CUBLAS_FILL_MODE_LOWER", "HIPBLAS_FILL_MODE_LOWER"), + ("CUBLAS_FILL_MODE_UPPER", "HIPBLAS_FILL_MODE_UPPER"), + ("CUBLAS_DIAG_NON_UNIT", "HIPBLAS_DIAG_NON_UNIT"), + ("CUBLAS_DIAG_UNIT", "HIPBLAS_DIAG_UNIT"), + ("CUBLAS_SIDE_LEFT", "HIPBLAS_SIDE_LEFT"), + ("CUBLAS_SIDE_RIGHT", "HIPBLAS_SIDE_RIGHT"), + ("CUBLAS_POINTER_MODE_HOST", "HIPBLAS_POINTER_MODE_HOST"), + ("CUBLAS_POINTER_MODE_DEVICE", "HIPBLAS_POINTER_MODE_DEVICE"), + ("CUBLAS_ATOMICS_NOT_ALLOWED", "HIPBLAS_ATOMICS_NOT_ALLOWED"), + ("CUBLAS_ATOMICS_ALLOWED", "HIPBLAS_ATOMICS_ALLOWED"), + ("CUBLAS_DATA_FLOAT", "HIPBLAS_DATA_FLOAT"), + ("CUBLAS_DATA_DOUBLE", "HIPBLAS_DATA_DOUBLE"), + ("CUBLAS_DATA_HALF", "HIPBLAS_DATA_HALF"), + ("CUBLAS_DATA_INT8", "HIPBLAS_DATA_INT8"), + ("CUBLAS_GEMM_DEFAULT", "HIPBLAS_GEMM_DEFAULT"), + ("CUBLAS_GEMM_DEFAULT_TENSOR_OP", "HIPBLAS_GEMM_DEFAULT"), + ("cublasCreate", "hipblasCreate"), + ("cublasDestroy", "hipblasDestroy"), + ("cublasSetVector", "hipblasSetVector"), + ("cublasGetVector", "hipblasGetVector"), + ("cublasSetVectorAsync", "hipblasSetVectorAsync"), + ("cublasGetVectorAsync", "hipblasGetVectorAsync"), + ("cublasSetMatrix", "hipblasSetMatrix"), + ("cublasGetMatrix", "hipblasGetMatrix"), + ("cublasGetMatrixAsync", "hipblasGetMatrixAsync"), + ("cublasSetMatrixAsync", "hipblasSetMatrixAsync"), + ("cublasXerbla", "hipblasXerbla"), + ("cublasSnrm2", "hipblasSnrm2"), + ("cublasDnrm2", "hipblasDnrm2"), + ("cublasScnrm2", "hipblasScnrm2"), + ("cublasDznrm2", "hipblasDznrm2"), + ("cublasNrm2Ex", "hipblasNrm2Ex"), + ("cublasSdot", "hipblasSdot"), + ("cublasSdotBatched", "hipblasSdotBatched"), + ("cublasDdot", "hipblasDdot"), + ("cublasDdotBatched", "hipblasDdotBatched"), + ("cublasCdotu", "hipblasCdotu"), + ("cublasCdotc", "hipblasCdotc"), + ("cublasZdotu", "hipblasZdotu"), + ("cublasZdotc", "hipblasZdotc"), + ("cublasSscal", "hipblasSscal"), + ("cublasSscalBatched", "hipblasSscalBatched"), + ("cublasDscal", "hipblasDscal"), + ("cublasDscalBatched", "hipblasDscalBatched"), + ("cublasCscal", "hipblasCscal"), + ("cublasCsscal", "hipblasCsscal"), + ("cublasZscal", "hipblasZscal"), + ("cublasZdscal", "hipblasZdscal"), + ("cublasSaxpy", "hipblasSaxpy"), + ("cublasSaxpyBatched", "hipblasSaxpyBatched"), + ("cublasDaxpy", "hipblasDaxpy"), + ("cublasCaxpy", "hipblasCaxpy"), + ("cublasZaxpy", "hipblasZaxpy"), + ("cublasScopy", "hipblasScopy"), + ("cublasScopyBatched", "hipblasScopyBatched"), + ("cublasDcopy", "hipblasDcopy"), + ("cublasDcopyBatched", "hipblasDcopyBatched"), + ("cublasCcopy", "hipblasCcopy"), + ("cublasZcopy", "hipblasZcopy"), + ("cublasSswap", "hipblasSswap"), + ("cublasDswap", "hipblasDswap"), + ("cublasCswap", "hipblasCswap"), + ("cublasZswap", "hipblasZswap"), + ("cublasIsamax", "hipblasIsamax"), + ("cublasIdamax", "hipblasIdamax"), + ("cublasIcamax", "hipblasIcamax"), + ("cublasIzamax", "hipblasIzamax"), + ("cublasIsamin", "hipblasIsamin"), + ("cublasIdamin", "hipblasIdamin"), + ("cublasIcamin", "hipblasIcamin"), + ("cublasIzamin", "hipblasIzamin"), + ("cublasSasum", "hipblasSasum"), + ("cublasSasumBatched", "hipblasSasumBatched"), + ("cublasDasum", "hipblasDasum"), + ("cublasDasumBatched", "hipblasDasumBatched"), + ("cublasScasum", "hipblasScasum"), + ("cublasDzasum", "hipblasDzasum"), + ("cublasSrot", "hipblasSrot"), + ("cublasDrot", "hipblasDrot"), + ("cublasCrot", "hipblasCrot"), + ("cublasCsrot", "hipblasCsrot"), + ("cublasZrot", "hipblasZrot"), + ("cublasZdrot", "hipblasZdrot"), + ("cublasSrotg", "hipblasSrotg"), + ("cublasDrotg", "hipblasDrotg"), + ("cublasCrotg", "hipblasCrotg"), + ("cublasZrotg", "hipblasZrotg"), + ("cublasSrotm", "hipblasSrotm"), + ("cublasDrotm", "hipblasDrotm"), + ("cublasSrotmg", "hipblasSrotmg"), + ("cublasDrotmg", "hipblasDrotmg"), + ("cublasSgemv", "hipblasSgemv"), + ("cublasSgemvBatched", "hipblasSgemvBatched"), + ("cublasDgemv", "hipblasDgemv"), + ("cublasCgemv", "hipblasCgemv"), + ("cublasZgemv", "hipblasZgemv"), + ("cublasSgbmv", "hipblasSgbmv"), + ("cublasDgbmv", "hipblasDgbmv"), + ("cublasCgbmv", "hipblasCgbmv"), + ("cublasZgbmv", "hipblasZgbmv"), + ("cublasStrmv", "hipblasStrmv"), + ("cublasDtrmv", "hipblasDtrmv"), + ("cublasCtrmv", "hipblasCtrmv"), + ("cublasZtrmv", "hipblasZtrmv"), + ("cublasStbmv", "hipblasStbmv"), + ("cublasDtbmv", "hipblasDtbmv"), + ("cublasCtbmv", "hipblasCtbmv"), + ("cublasZtbmv", "hipblasZtbmv"), + ("cublasStpmv", "hipblasStpmv"), + ("cublasDtpmv", "hipblasDtpmv"), + ("cublasCtpmv", "hipblasCtpmv"), + ("cublasZtpmv", "hipblasZtpmv"), + ("cublasStrsv", "hipblasStrsv"), + ("cublasDtrsv", "hipblasDtrsv"), + ("cublasCtrsv", "hipblasCtrsv"), + ("cublasZtrsv", "hipblasZtrsv"), + ("cublasStpsv", "hipblasStpsv"), + ("cublasDtpsv", "hipblasDtpsv"), + ("cublasCtpsv", "hipblasCtpsv"), + ("cublasZtpsv", "hipblasZtpsv"), + ("cublasStbsv", "hipblasStbsv"), + ("cublasDtbsv", "hipblasDtbsv"), + ("cublasCtbsv", "hipblasCtbsv"), + ("cublasZtbsv", "hipblasZtbsv"), + ("cublasSsymv", "hipblasSsymv"), + ("cublasDsymv", "hipblasDsymv"), + ("cublasCsymv", "hipblasCsymv"), + ("cublasZsymv", "hipblasZsymv"), + ("cublasChemv", "hipblasChemv"), + ("cublasZhemv", "hipblasZhemv"), + ("cublasSsbmv", "hipblasSsbmv"), + ("cublasDsbmv", "hipblasDsbmv"), + ("cublasChbmv", "hipblasChbmv"), + ("cublasZhbmv", "hipblasZhbmv"), + ("cublasSspmv", "hipblasSspmv"), + ("cublasDspmv", "hipblasDspmv"), + ("cublasChpmv", "hipblasChpmv"), + ("cublasZhpmv", "hipblasZhpmv"), + ("cublasSger", "hipblasSger"), + ("cublasDger", "hipblasDger"), + ("cublasCgeru", "hipblasCgeru"), + ("cublasCgerc", "hipblasCgerc"), + ("cublasZgeru", "hipblasZgeru"), + ("cublasZgerc", "hipblasZgerc"), + ("cublasSsyr", "hipblasSsyr"), + ("cublasDsyr", "hipblasDsyr"), + ("cublasCher", "hipblasCher"), + ("cublasZher", "hipblasZher"), + ("cublasSspr", "hipblasSspr"), + ("cublasDspr", "hipblasDspr"), + ("cublasChpr", "hipblasChpr"), + ("cublasZhpr", "hipblasZhpr"), + ("cublasSsyr2", "hipblasSsyr2"), + ("cublasDsyr2", "hipblasDsyr2"), + ("cublasCher2", "hipblasCher2"), + ("cublasZher2", "hipblasZher2"), + ("cublasSspr2", "hipblasSspr2"), + ("cublasDspr2", "hipblasDspr2"), + ("cublasChpr2", "hipblasChpr2"), + ("cublasZhpr2", "hipblasZhpr2"), + ("cublasSgemmBatched", "hipblasSgemmBatched"), + ("cublasDgemmBatched", "hipblasDgemmBatched"), + ("cublasHgemmBatched", "hipblasHgemmBatched"), + ("cublasSgemmStridedBatched", "hipblasSgemmStridedBatched"), + ("cublasDgemmStridedBatched", "hipblasDgemmStridedBatched"), + ("cublasHgemmStridedBatched", "hipblasHgemmStridedBatched"), + ("cublasCgemmBatched", "hipblasCgemmBatched"), + ("cublasCgemm3mBatched", "hipblasCgemm3mBatched"), + ("cublasZgemmBatched", "hipblasZgemmBatched"), + ("cublasCgemmStridedBatched", "hipblasCgemmStridedBatched"), + ("cublasCgemm3mStridedBatched", "hipblasCgemm3mStridedBatched"), + ("cublasZgemmStridedBatched", "hipblasZgemmStridedBatched"), + ("cublasSgemm", "hipblasSgemm"), + ("cublasDgemm", "hipblasDgemm"), + ("cublasCgemm", "hipblasCgemm"), + ("cublasZgemm", "hipblasZgemm"), + ("cublasHgemm", "hipblasHgemm"), + ("cublasSsyrk", "hipblasSsyrk"), + ("cublasDsyrk", "hipblasDsyrk"), + ("cublasCsyrk", "hipblasCsyrk"), + ("cublasZsyrk", "hipblasZsyrk"), + ("cublasCherk", "hipblasCherk"), + ("cublasZherk", "hipblasZherk"), + ("cublasSsyr2k", "hipblasSsyr2k"), + ("cublasDsyr2k", "hipblasDsyr2k"), + ("cublasCsyr2k", "hipblasCsyr2k"), + ("cublasZsyr2k", "hipblasZsyr2k"), + ("cublasSsyrkx", "hipblasSsyrkx"), + ("cublasDsyrkx", "hipblasDsyrkx"), + ("cublasCsyrkx", "hipblasCsyrkx"), + ("cublasZsyrkx", "hipblasZsyrkx"), + ("cublasCher2k", "hipblasCher2k"), + ("cublasZher2k", "hipblasZher2k"), + ("cublasCherkx", "hipblasCherkx"), + ("cublasZherkx", "hipblasZherkx"), + ("cublasSsymm", "hipblasSsymm"), + ("cublasDsymm", "hipblasDsymm"), + ("cublasCsymm", "hipblasCsymm"), + ("cublasZsymm", "hipblasZsymm"), + ("cublasChemm", "hipblasChemm"), + ("cublasZhemm", "hipblasZhemm"), + ("cublasStrsm", "hipblasStrsm"), + ("cublasDtrsm", "hipblasDtrsm"), + ("cublasCtrsm", "hipblasCtrsm"), + ("cublasZtrsm", "hipblasZtrsm"), + ("cublasStrsmBatched", "hipblasStrsmBatched"), + ("cublasDtrsmBatched", "hipblasDtrsmBatched"), + ("cublasCtrsmBatched", "hipblasCtrsmBatched"), + ("cublasZtrsmBatched", "hipblasZtrsmBatched"), + ("cublasStrmm", "hipblasStrmm"), + ("cublasDtrmm", "hipblasDtrmm"), + ("cublasCtrmm", "hipblasCtrmm"), + ("cublasZtrmm", "hipblasZtrmm"), + ("cublasSgeam", "hipblasSgeam"), + ("cublasDgeam", "hipblasDgeam"), + ("cublasCgeam", "hipblasCgeam"), + ("cublasZgeam", "hipblasZgeam"), + ("cublasSgetrfBatched", "hipblasSgetrfBatched"), + ("cublasDgetrfBatched", "hipblasDgetrfBatched"), + ("cublasCgetrfBatched", "hipblasCgetrfBatched"), + ("cublasZgetrfBatched", "hipblasZgetrfBatched"), + ("cublasSgetriBatched", "hipblasSgetriBatched"), + ("cublasDgetriBatched", "hipblasDgetriBatched"), + ("cublasCgetriBatched", "hipblasCgetriBatched"), + ("cublasZgetriBatched", "hipblasZgetriBatched"), + ("cublasSgetrsBatched", "hipblasSgetrsBatched"), + ("cublasDgetrsBatched", "hipblasDgetrsBatched"), + ("cublasCgetrsBatched", "hipblasCgetrsBatched"), + ("cublasZgetrsBatched", "hipblasZgetrsBatched"), + ("cublasSmatinvBatched", "hipblasSmatinvBatched"), + ("cublasDmatinvBatched", "hipblasDmatinvBatched"), + ("cublasCmatinvBatched", "hipblasCmatinvBatched"), + ("cublasZmatinvBatched", "hipblasZmatinvBatched"), + ("cublasSgeqrfBatched", "hipblasSgeqrfBatched"), + ("cublasDgeqrfBatched", "hipblasDgeqrfBatched"), + ("cublasCgeqrfBatched", "hipblasCgeqrfBatched"), + ("cublasZgeqrfBatched", "hipblasZgeqrfBatched"), + ("cublasSgelsBatched", "hipblasSgelsBatched"), + ("cublasDgelsBatched", "hipblasDgelsBatched"), + ("cublasCgelsBatched", "hipblasCgelsBatched"), + ("cublasZgelsBatched", "hipblasZgelsBatched"), + ("cublasSdgmm", "hipblasSdgmm"), + ("cublasDdgmm", "hipblasDdgmm"), + ("cublasCdgmm", "hipblasCdgmm"), + ("cublasZdgmm", "hipblasZdgmm"), + ("cublasStpttr", "hipblasStpttr"), + ("cublasDtpttr", "hipblasDtpttr"), + ("cublasCtpttr", "hipblasCtpttr"), + ("cublasZtpttr", "hipblasZtpttr"), + ("cublasStrttp", "hipblasStrttp"), + ("cublasDtrttp", "hipblasDtrttp"), + ("cublasCtrttp", "hipblasCtrttp"), + ("cublasZtrttp", "hipblasZtrttp"), + ("cublasCreate_v2", "hipblasCreate_v2"), + ("cublasDestroy_v2", "hipblasDestroy_v2"), + ("cublasGetVersion_v2", "hipblasGetVersion_v2"), + ("cublasSetWorkspace", "hipblasSetWorkspace"), + ("cublasSetStream", "hipblasSetStream"), + ("cublasGetStream", "hipblasGetStream"), + ("cublasSetStream_v2", "hipblasSetStream_v2"), + ("cublasGetStream_v2", "hipblasGetStream_v2"), + ("cublasGetPointerMode", "hipblasGetPointerMode"), + ("cublasSetPointerMode", "hipblasSetPointerMode"), + ("cublasGetPointerMode_v2", "hipblasGetPointerMode_v2"), + ("cublasSetPointerMode_v2", "hipblasSetPointerMode_v2"), + ("cublasSgemv_v2", "hipblasSgemv_v2"), + ("cublasDgemv_v2", "hipblasDgemv_v2"), + ("cublasCgemv_v2", "hipblasCgemv_v2"), + ("cublasZgemv_v2", "hipblasZgemv_v2"), + ("cublasSgbmv_v2", "hipblasSgbmv_v2"), + ("cublasDgbmv_v2", "hipblasDgbmv_v2"), + ("cublasCgbmv_v2", "hipblasCgbmv_v2"), + ("cublasZgbmv_v2", "hipblasZgbmv_v2"), + ("cublasStrmv_v2", "hipblasStrmv_v2"), + ("cublasDtrmv_v2", "hipblasDtrmv_v2"), + ("cublasCtrmv_v2", "hipblasCtrmv_v2"), + ("cublasZtrmv_v2", "hipblasZtrmv_v2"), + ("cublasStbmv_v2", "hipblasStbmv_v2"), + ("cublasDtbmv_v2", "hipblasDtbmv_v2"), + ("cublasCtbmv_v2", "hipblasCtbmv_v2"), + ("cublasZtbmv_v2", "hipblasZtbmv_v2"), + ("cublasStpmv_v2", "hipblasStpmv_v2"), + ("cublasDtpmv_v2", "hipblasDtpmv_v2"), + ("cublasCtpmv_v2", "hipblasCtpmv_v2"), + ("cublasZtpmv_v2", "hipblasZtpmv_v2"), + ("cublasStrsv_v2", "hipblasStrsv_v2"), + ("cublasDtrsv_v2", "hipblasDtrsv_v2"), + ("cublasCtrsv_v2", "hipblasCtrsv_v2"), + ("cublasZtrsv_v2", "hipblasZtrsv_v2"), + ("cublasStpsv_v2", "hipblasStpsv_v2"), + ("cublasDtpsv_v2", "hipblasDtpsv_v2"), + ("cublasCtpsv_v2", "hipblasCtpsv_v2"), + ("cublasZtpsv_v2", "hipblasZtpsv_v2"), + ("cublasStbsv_v2", "hipblasStbsv_v2"), + ("cublasDtbsv_v2", "hipblasDtbsv_v2"), + ("cublasCtbsv_v2", "hipblasCtbsv_v2"), + ("cublasZtbsv_v2", "hipblasZtbsv_v2"), + ("cublasSsymv_v2", "hipblasSsymv_v2"), + ("cublasDsymv_v2", "hipblasDsymv_v2"), + ("cublasCsymv_v2", "hipblasCsymv_v2"), + ("cublasZsymv_v2", "hipblasZsymv_v2"), + ("cublasChemv_v2", "hipblasChemv_v2"), + ("cublasZhemv_v2", "hipblasZhemv_v2"), + ("cublasSsbmv_v2", "hipblasSsbmv_v2"), + ("cublasDsbmv_v2", "hipblasDsbmv_v2"), + ("cublasChbmv_v2", "hipblasChbmv_v2"), + ("cublasZhbmv_v2", "hipblasZhbmv_v2"), + ("cublasSspmv_v2", "hipblasSspmv_v2"), + ("cublasDspmv_v2", "hipblasDspmv_v2"), + ("cublasChpmv_v2", "hipblasChpmv_v2"), + ("cublasZhpmv_v2", "hipblasZhpmv_v2"), + ("cublasSger_v2", "hipblasSger_v2"), + ("cublasDger_v2", "hipblasDger_v2"), + ("cublasCgeru_v2", "hipblasCgeru_v2"), + ("cublasCgerc_v2", "hipblasCergc_v2"), + ("cublasZgeru_v2", "hipblasZgeru_v2"), + ("cublasZgerc_v2", "hipblasZgerc_v2"), + ("cublasSsyr_v2", "hipblasSsyr_v2"), + ("cublasDsyr_v2", "hipblasDsyr_v2"), + ("cublasCsyr_v2", "hipblasCsyr_v2"), + ("cublasZsyr_v2", "hipblasZsyr_v2"), + ("cublasCher_v2", "hipblasCher_v2"), + ("cublasZher_v2", "hipblasZher_v2"), + ("cublasSspr_v2", "hipblasSspr_v2"), + ("cublasDspr_v2", "hipblasDspr_v2"), + ("cublasChpr_v2", "hipblasChpr_v2"), + ("cublasZhpr_v2", "hipblasZhpr_v2"), + ("cublasSsyr2_v2", "hipblasSsyr2_v2"), + ("cublasDsyr2_v2", "hipblasDsyr2_v2"), + ("cublasCsyr2_v2", "hipblasCsyr2_v2"), + ("cublasZsyr2_v2", "hipblasZsyr2_v2"), + ("cublasCher2_v2", "hipblasCher2_v2"), + ("cublasZher2_v2", "hipblasZher2_v2"), + ("cublasSspr2_v2", "hipblasSspr2_v2"), + ("cublasDspr2_v2", "hipblasDspr2_v2"), + ("cublasChpr2_v2", "hipblasChpr2_v2"), + ("cublasZhpr2_v2", "hipblasZhpr2_v2"), + ("cublasSgemm_v2", "hipblasSgemm_v2"), + ("cublasDgemm_v2", "hipblasDgemm_v2"), + ("cublasCgemm_v2", "hipblasCgemm_v2"), + ("cublasCgemm3m", "hipblasCgemm3m"), + ("cublasCgemm3mEx", "hipblasCgemm3mEx"), + ("cublasZgemm_v2", "hipblasZgemm_v2"), + ("cublasZgemm3m", "hipblasZgemm3m"), + ("cublasSgemmEx", "hipblasSgemmEx"), + ("cublasGemmEx", "hipblasGemmEx"), + ("cublasGemmBatchedEx", "hipblasGemmBatchedEx"), + ("cublasGemmStridedBatchedEx", "hipblasGemmStridedBatchedEx"), + ("cublasCgemmEx", "hipblasCgemmEx"), + ("cublasUint8gemmBias", "hipblasUint8gemmBias"), + ("cublasSsyrk_v2", "hipblasSsyrk_v2"), + ("cublasDsyrk_v2", "hipblasDsyrk_v2"), + ("cublasCsyrk_v2", "hipblasCsyrk_v2"), + ("cublasZsyrk_v2", "hipblasZsyrk_v2"), + ("cublasCsyrkEx", "hipblasCsyrkEx"), + ("cublasCsyrk3mEx", "hipblasCsyrk3mEx"), + ("cublasCherk_v2", "hipblasCherk_v2"), + ("cublasCherkEx", "hipblasCherkEx"), + ("cublasCherk3mEx", "hipblasCherk3mEx"), + ("cublasZherk_v2", "hipblasZherk_v2"), + ("cublasSsyr2k_v2", "hipblasSsyr2k_v2"), + ("cublasDsyr2k_v2", "hipblasDsyr2k_v2"), + ("cublasCsyr2k_v2", "hipblasCsyr2k_v2"), + ("cublasZsyr2k_v2", "hipblasZsyr2k_v2"), + ("cublasCher2k_v2", "hipblasCher2k_v2"), + ("cublasZher2k_v2", "hipblasZher2k_v2"), + ("cublasSsymm_v2", "hipblasSsymm_v2"), + ("cublasDsymm_v2", "hipblasDsymm_v2"), + ("cublasCsymm_v2", "hipblasCsymm_v2"), + ("cublasZsymm_v2", "hipblasZsymm_v2"), + ("cublasChemm_v2", "hipblasChemm_v2"), + ("cublasZhemm_v2", "hipblasZhemm_v2"), + ("cublasStrsm_v2", "hipblasStrsm_v2"), + ("cublasDtrsm_v2", "hipblasDtrsm_v2"), + ("cublasCtrsm_v2", "hipblasCtrsm_v2"), + ("cublasZtrsm_v2", "hipblasZtrsm_v2"), + ("cublasStrmm_v2", "hipblasStrmm_v2"), + ("cublasDtrmm_v2", "hipblasDtrmm_v2"), + ("cublasCtrmm_v2", "hipblasCtrmm_v2"), + ("cublasZtrmm_v2", "hipblasZtrmm_v2"), + ("cublasSnrm2_v2", "hipblasSnrm2_v2"), + ("cublasDnrm2_v2", "hipblasDnrm2_v2"), + ("cublasScnrm2_v2", "hipblasScnrm2_v2"), + ("cublasDznrm2_v2", "hipblasDznrm2_v2"), + ("cublasDotEx", "hipblasDotEx"), + ("cublasDotcEx", "hipblasDotcEx"), + ("cublasSdot_v2", "hipblasSdot_v2"), + ("cublasDdot_v2", "hipblasDdot_v2"), + ("cublasCdotu_v2", "hipblasCdotu_v2"), + ("cublasCdotc_v2", "hipblasCdotc_v2"), + ("cublasZdotu_v2", "hipblasZdotu_v2"), + ("cublasZdotc_v2", "hipblasZdotc_v2"), + ("cublasScalEx", "hipblasScalEx"), + ("cublasSscal_v2", "hipblasSscal_v2"), + ("cublasDscal_v2", "hipblasDscal_v2"), + ("cublasCscal_v2", "hipblasCscal_v2"), + ("cublasCsscal_v2", "hipblasCsscal_v2"), + ("cublasZscal_v2", "hipblasZcsal_v2"), + ("cublasZdscal_v2", "hipblasZdscal_v2"), + ("cublasAxpyEx", "hipblasAxpyEx"), + ("cublasSaxpy_v2", "hipblasSaxpy_v2"), + ("cublasDaxpy_v2", "hipblasDaxpy_v2"), + ("cublasCaxpy_v2", "hipblasCaxpy_v2"), + ("cublasZaxpy_v2", "hipblasZaxpy_v2"), + ("cublasScopy_v2", "hipblasScopy_v2"), + ("cublasDcopy_v2", "hipblasDcopy_v2"), + ("cublasCcopy_v2", "hipblasCcopy_v2"), + ("cublasZcopy_v2", "hipblasZcopy_v2"), + ("cublasSswap_v2", "hipblasSswap_v2"), + ("cublasDswap_v2", "hipblasDswap_v2"), + ("cublasCswap_v2", "hipblasCswap_v2"), + ("cublasZswap_v2", "hipblasZswap_v2"), + ("cublasIsamax_v2", "hipblasIsamax_v2"), + ("cublasIdamax_v2", "hipblasIdamax_v2"), + ("cublasIcamax_v2", "hipblasIcamax_v2"), + ("cublasIzamax_v2", "hipblasIzamax_v2"), + ("cublasIsamin_v2", "hipblasIsamin_v2"), + ("cublasIdamin_v2", "hipblasIdamin_v2"), + ("cublasIcamin_v2", "hipblasIcamin_v2"), + ("cublasIzamin_v2", "hipblasIzamin_v2"), + ("cublasSasum_v2", "hipblasSasum_v2"), + ("cublasDasum_v2", "hipblasDasum_v2"), + ("cublasScasum_v2", "hipblasScasum_v2"), + ("cublasDzasum_v2", "hipblasDzasum_v2"), + ("cublasSrot_v2", "hipblasSrot_v2"), + ("cublasDrot_v2", "hipblasDrot_v2"), + ("cublasCrot_v2", "hipblasCrot_v2"), + ("cublasCsrot_v2", "hipblasCsrot_v2"), + ("cublasZrot_v2", "hipblasZrot_v2"), + ("cublasZdrot_v2", "hipblasZdrot_v2"), + ("cublasSrotg_v2", "hipblasSrotg_v2"), + ("cublasDrotg_v2", "hipblasDrotg_v2"), + ("cublasCrotg_v2", "hipblasCrotg_v2"), + ("cublasZrotg_v2", "hipblasZrotg_v2"), + ("cublasSrotm_v2", "hipblasSrotm_v2"), + ("cublasDrotm_v2", "hipblasDrotm_v2"), + ("cublasSrotmg_v2", "hipblasSrotmg_v2"), + ("cublasDrotmg_v2", "hipblasDrotmg_v2"), + ("cublasComputeType_t", "hipblasComputeType_t"), + ("CUBLAS_COMPUTE_32I", "HIPBLAS_COMPUTE_32I"), + ("CUBLAS_COMPUTE_32F", "HIPBLAS_COMPUTE_32F"), + ("CUBLAS_COMPUTE_32F_FAST_TF32", "HIPBLAS_COMPUTE_32F_FAST_TF32"), + ("CUBLAS_COMPUTE_64F", "HIPBLAS_COMPUTE_64F"), + ("cublasLtEpilogue_t", "hipblasLtEpilogue_t"), + ("CUBLASLT_EPILOGUE_DEFAULT", "HIPBLASLT_EPILOGUE_DEFAULT"), + ("CUBLASLT_EPILOGUE_RELU", "HIPBLASLT_EPILOGUE_RELU"), + ("CUBLASLT_EPILOGUE_BIAS", "HIPBLASLT_EPILOGUE_BIAS"), + ("CUBLASLT_EPILOGUE_RELU_BIAS", "HIPBLASLT_EPILOGUE_RELU_BIAS"), + ("CUBLASLT_EPILOGUE_GELU", "HIPBLASLT_EPILOGUE_GELU"), + ("CUBLASLT_EPILOGUE_GELU_BIAS", "HIPBLASLT_EPILOGUE_GELU_BIAS"), + ("cublasLtHandle_t", "hipblasLtHandle_t"), + ("cublasLtMatmulDesc_t", "hipblasLtMatmulDesc_t"), + ("cublasLtMatmulDescOpaque_t", "hipblasLtMatmulDescOpaque_t"), + ("cublasLtMatmulDescAttributes_t", "hipblasLtMatmulDescAttributes_t"), + ("CUBLASLT_MATMUL_DESC_TRANSA", "HIPBLASLT_MATMUL_DESC_TRANSA"), + ("CUBLASLT_MATMUL_DESC_TRANSB", "HIPBLASLT_MATMUL_DESC_TRANSB"), + ("CUBLASLT_MATMUL_DESC_EPILOGUE", "HIPBLASLT_MATMUL_DESC_EPILOGUE"), + ("CUBLASLT_MATMUL_DESC_BIAS_POINTER", "HIPBLASLT_MATMUL_DESC_BIAS_POINTER"), + ("CUBLASLT_MATMUL_DESC_A_SCALE_MODE", "HIPBLASLT_MATMUL_DESC_A_SCALE_MODE"), + ("CUBLASLT_MATMUL_DESC_B_SCALE_MODE", "HIPBLASLT_MATMUL_DESC_B_SCALE_MODE"), + ("CUBLASLT_MATMUL_DESC_A_SCALE_POINTER", "HIPBLASLT_MATMUL_DESC_A_SCALE_POINTER"), + ("CUBLASLT_MATMUL_DESC_B_SCALE_POINTER", "HIPBLASLT_MATMUL_DESC_B_SCALE_POINTER"), + ("CUBLASLT_MATMUL_DESC_D_SCALE_POINTER", "HIPBLASLT_MATMUL_DESC_D_SCALE_POINTER"), + ("CUBLASLT_MATMUL_DESC_AMAX_D_POINTER", "HIPBLASLT_MATMUL_DESC_AMAX_D_POINTER"), + ("CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE", "HIPBLASLT_MATMUL_DESC_BIAS_DATA_TYPE"), + ("CUBLASLT_MATMUL_DESC_POINTER_MODE", "HIPBLASLT_MATMUL_DESC_POINTER_MODE"), + ("CUBLASLT_MATMUL_MATRIX_SCALE_OUTER_VEC_32F", "HIPBLASLT_MATMUL_MATRIX_SCALE_OUTER_VEC_32F"), + ("CUBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0", "HIPBLASLT_MATMUL_MATRIX_SCALE_VEC32_UE8M0"), + ("CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3", "HIPBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3"), + ("CUBLASLT_POINTER_MODE_DEVICE", "HIPBLASLT_POINTER_MODE_DEVICE"), + ("CUBLASLT_POINTER_MODE_HOST", "HIPBLASLT_POINTER_MODE_HOST"), + ("cublasLtMatrixLayout_t", "hipblasLtMatrixLayout_t"), + ("cublasLtMatrixLayoutOpaque_t", "hipblasLtMatrixLayoutOpaque_t"), + ("cublasLtMatrixLayoutAttribute_t", "hipblasLtMatrixLayoutAttribute_t"), + ("cublasLtMatrixLayoutCreate", "hipblasLtMatrixLayoutCreate"), + ("cublasLtMatrixLayoutDestroy", "hipblasLtMatrixLayoutDestroy"), + ("cublasLtMatrixLayoutSetAttribute", "hipblasLtMatrixLayoutSetAttribute"), + ("CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT", "HIPBLASLT_MATRIX_LAYOUT_BATCH_COUNT"), + ("CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET", "HIPBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET"), + ("cublasLtMatmulPreference_t", "hipblasLtMatmulPreference_t"), + ("cublasLtMatmulPreferenceOpaque_t", "hipblasLtMatmulPreferenceOpaque_t"), + ("cublasLtMatmulPreferenceAttributes_t", "hipblasLtMatmulPreferenceAttributes_t"), + ("CUBLASLT_MATMUL_PREF_SEARCH_MODE", "HIPBLASLT_MATMUL_PREF_SEARCH_MODE"), + ("CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES", "HIPBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES"), + ("cublasLtMatmulAlgo_t", "hipblasLtMatmulAlgo_t"), + ("cublasLtMatmulHeuristicResult_t", "hipblasLtMatmulHeuristicResult_t"), + ("cublasLtCreate", "hipblasLtCreate"), + ("cublasLtDestroy", "hipblasLtDestroy"), + ("cublasLtMatmulDescCreate", "hipblasLtMatmulDescCreate"), + ("cublasLtMatmulDescDestroy", "hipblasLtMatmulDescDestroy"), + ("cublasLtMatmulDescSetAttribute", "hipblasLtMatmulDescSetAttribute"), + ("cublasLtMatmulPreferenceCreate", "hipblasLtMatmulPreferenceCreate"), + ("cublasLtMatmulPreferenceDestroy", "hipblasLtMatmulPreferenceDestroy"), + ("cublasLtMatmulPreferenceSetAttribute", "hipblasLtMatmulPreferenceSetAttribute"), + ("cublasLtMatmulAlgoGetHeuristic", "hipblasLtMatmulAlgoGetHeuristic"), + ("cublasLtMatmul", "hipblasLtMatmul"), + ("CURAND_STATUS_SUCCESS", "HIPRAND_STATUS_SUCCESS"), + ("CURAND_STATUS_VERSION_MISMATCH", "HIPRAND_STATUS_VERSION_MISMATCH"), + ("CURAND_STATUS_NOT_INITIALIZED", "HIPRAND_STATUS_NOT_INITIALIZED"), + ("CURAND_STATUS_ALLOCATION_FAILED", "HIPRAND_STATUS_ALLOCATION_FAILED"), + ("CURAND_STATUS_TYPE_ERROR", "HIPRAND_STATUS_TYPE_ERROR"), + ("CURAND_STATUS_OUT_OF_RANGE", "HIPRAND_STATUS_OUT_OF_RANGE"), + ("CURAND_STATUS_LENGTH_NOT_MULTIPLE", "HIPRAND_STATUS_LENGTH_NOT_MULTIPLE"), + ("CURAND_STATUS_DOUBLE_PRECISION_REQUIRED", "HIPRAND_STATUS_DOUBLE_PRECISION_REQUIRED"), + ("CURAND_STATUS_LAUNCH_FAILURE", "HIPRAND_STATUS_LAUNCH_FAILURE"), + ("CURAND_STATUS_PREEXISTING_FAILURE", "HIPRAND_STATUS_PREEXISTING_FAILURE"), + ("CURAND_STATUS_INITIALIZATION_FAILED", "HIPRAND_STATUS_INITIALIZATION_FAILED"), + ("CURAND_STATUS_ARCH_MISMATCH", "HIPRAND_STATUS_ARCH_MISMATCH"), + ("CURAND_STATUS_INTERNAL_ERROR", "HIPRAND_STATUS_INTERNAL_ERROR"), + ("CURAND_RNG_TEST", "HIPRAND_RNG_TEST"), + ("mtgp32dc_params_fast_11213", "mtgp32dc_params_fast_11213"), + ("CURAND_RNG_PSEUDO_DEFAULT", "HIPRAND_RNG_PSEUDO_DEFAULT"), + ("CURAND_RNG_PSEUDO_XORWOW", "HIPRAND_RNG_PSEUDO_XORWOW"), + ("CURAND_RNG_PSEUDO_MRG32K3A", "HIPRAND_RNG_PSEUDO_MRG32K3A"), + ("CURAND_RNG_PSEUDO_MTGP32", "HIPRAND_RNG_PSEUDO_MTGP32"), + ("CURAND_RNG_PSEUDO_MT19937", "HIPRAND_RNG_PSEUDO_MT19937"), + ("CURAND_RNG_PSEUDO_PHILOX4_32_10", "HIPRAND_RNG_PSEUDO_PHILOX4_32_10"), + ("CURAND_RNG_QUASI_DEFAULT", "HIPRAND_RNG_QUASI_DEFAULT"), + ("CURAND_RNG_QUASI_SOBOL32", "HIPRAND_RNG_QUASI_SOBOL32"), + ("CURAND_RNG_QUASI_SCRAMBLED_SOBOL32", "HIPRAND_RNG_QUASI_SCRAMBLED_SOBOL32"), + ("CURAND_RNG_QUASI_SOBOL64", "HIPRAND_RNG_QUASI_SOBOL64"), + ("CURAND_RNG_QUASI_SCRAMBLED_SOBOL64", "HIPRAND_RNG_QUASI_SCRAMBLED_SOBOL64"), + ("curand_ORDERING_PSEUDO_BEST", "HIPRAND_ORDERING_PSEUDO_BEST"), + ("curand_ORDERING_PSEUDO_DEFAULT", "HIPRAND_ORDERING_PSEUDO_DEFAULT"), + ("curand_ORDERING_PSEUDO_SEEDED", "HIPRAND_ORDERING_PSEUDO_SEEDED"), + ("curand_ORDERING_QUASI_DEFAULT", "HIPRAND_ORDERING_QUASI_DEFAULT"), + ("curand_DIRECTION_VECTORS_32_JOEKUO6", "HIPRAND_DIRECTION_VECTORS_32_JOEKUO6"), + ("curand_SCRAMBLED_DIRECTION_VECTORS_32_JOEKUO6", "HIPRAND_SCRAMBLED_DIRECTION_VECTORS_32_JOEKUO6"), + ("curand_DIRECTION_VECTORS_64_JOEKUO6", "HIPRAND_DIRECTION_VECTORS_64_JOEKUO6"), + ("curand_SCRAMBLED_DIRECTION_VECTORS_64_JOEKUO6", "HIPRAND_SCRAMBLED_DIRECTION_VECTORS_64_JOEKUO6"), + ("curand_CHOOSE_BEST", "HIPRAND_CHOOSE_BEST"), + ("curand_ITR", "HIPRAND_ITR"), + ("curand_KNUTH", "HIPRAND_KNUTH"), + ("curand_HITR", "HIPRAND_HITR"), + ("curand_M1", "HIPRAND_M1"), + ("curand_M2", "HIPRAND_M2"), + ("curand_BINARY_SEARCH", "HIPRAND_BINARY_SEARCH"), + ("curand_DISCRETE_GAUSS", "HIPRAND_DISCRETE_GAUSS"), + ("curand_REJECTION", "HIPRAND_REJECTION"), + ("curand_DEVICE_API", "HIPRAND_DEVICE_API"), + ("curand_FAST_REJECTION", "HIPRAND_FAST_REJECTION"), + ("curand_3RD", "HIPRAND_3RD"), + ("curand_DEFINITION", "HIPRAND_DEFINITION"), + ("curand_POISSON", "HIPRAND_POISSON"), + ("curandCreateGenerator", "hiprandCreateGenerator"), + ("curandCreateGeneratorHost", "hiprandCreateGeneratorHost"), + ("curandCreatePoissonDistribution", "hiprandCreatePoissonDistribution"), + ("curandDestroyDistribution", "hiprandDestroyDistribution"), + ("curandDestroyGenerator", "hiprandDestroyGenerator"), + ("curandGenerate", "hiprandGenerate"), + ("curandGenerateLogNormal", "hiprandGenerateLogNormal"), + ("curandGenerateLogNormalDouble", "hiprandGenerateLogNormalDouble"), + ("curandGenerateLongLong", "hiprandGenerateLongLong"), + ("curandGenerateNormal", "hiprandGenerateNormal"), + ("curandGenerateNormalDouble", "hiprandGenerateNormalDouble"), + ("curandGeneratePoisson", "hiprandGeneratePoisson"), + ("curandGenerateSeeds", "hiprandGenerateSeeds"), + ("curandGenerateUniform", "hiprandGenerateUniform"), + ("curandGenerateUniformDouble", "hiprandGenerateUniformDouble"), + ("curandGetDirectionVectors32", "hiprandGetDirectionVectors32"), + ("curandGetDirectionVectors64", "hiprandGetDirectionVectors64"), + ("curandGetProperty", "hiprandGetProperty"), + ("curandGetScrambleConstants32", "hiprandGetScrambleConstants32"), + ("curandGetScrambleConstants64", "hiprandGetScrambleConstants64"), + ("curandGetVersion", "hiprandGetVersion"), + ("curandSetGeneratorOffset", "hiprandSetGeneratorOffset"), + ("curandSetGeneratorOrdering", "hiprandSetGeneratorOrdering"), + ("curandSetPseudoRandomGeneratorSeed", "hiprandSetPseudoRandomGeneratorSeed"), + ("curandSetQuasiRandomGeneratorDimensions", "hiprandSetQuasiRandomGeneratorDimensions"), + ("curandSetStream", "hiprandSetStream"), + ("curand", "hiprand"), + ("curand4", "hiprand4"), + ("curand_init", "hiprand_init"), + ("curand_log_normal", "hiprand_log_normal"), + ("curand_log_normal_double", "hiprand_log_normal_double"), + ("curand_log_normal2", "hiprand_log_normal2"), + ("curand_log_normal2_double", "hiprand_log_normal2_double"), + ("curand_log_normal4", "hiprand_log_normal4"), + ("curand_log_normal4_double", "hiprand_log_normal4_double"), + ("curand_mtgp32_single", "hiprand_mtgp32_single"), + ("curand_mtgp32_single_specific", "hiprand_mtgp32_single_specific"), + ("curand_mtgp32_specific", "hiprand_mtgp32_specific"), + ("curand_normal", "hiprand_normal"), + ("curandMakeMTGP32Constants", "hiprandMakeMTGP32Constants"), + ("curandMakeMTGP32KernelState", "hiprandMakeMTGP32KernelState"), + ("curand_normal_double", "hiprand_normal_double"), + ("curand_normal2", "hiprand_normal2"), + ("curand_normal2_double", "hiprand_normal2_double"), + ("curand_normal4", "hiprand_normal4"), + ("curand_normal4_double", "hiprand_normal4_double"), + ("curand_uniform", "hiprand_uniform"), + ("curand_uniform_double", "hiprand_uniform_double"), + ("curand_uniform2_double", "hiprand_uniform2_double"), + ("curand_uniform4", "hiprand_uniform4"), + ("curand_uniform4_double", "hiprand_uniform4_double"), + ("curand_discrete", "hiprand_discrete"), + ("curand_discrete4", "hiprand_discrete4"), + ("curand_poisson", "hiprand_poisson"), + ("curand_poisson4", "hiprand_poisson4"), + ("curand_Philox4x32_10", "hiprand_Philox4x32_10"), + ("mtgp32_kernel_params", "mtgp32_kernel_params_t"), + ("CUFFT_FORWARD", "HIPFFT_FORWARD"), + ("CUFFT_INVERSE", "HIPFFT_BACKWARD"), + ("CUFFT_COMPATIBILITY_DEFAULT", "HIPFFT_COMPATIBILITY_DEFAULT"), + ("cuComplex", "hipComplex"), + ("cuDoubleComplex", "hipDoubleComplex"), + ("cufftResult_t", "hipfftResult_t"), + ("cufftResult", "hipfftResult"), + ("CUFFT_SUCCESS", "HIPFFT_SUCCESS"), + ("CUFFT_INVALID_PLAN", "HIPFFT_INVALID_PLAN"), + ("CUFFT_ALLOC_FAILED", "HIPFFT_ALLOC_FAILED"), + ("CUFFT_INVALID_TYPE", "HIPFFT_INVALID_TYPE"), + ("CUFFT_INVALID_VALUE", "HIPFFT_INVALID_VALUE"), + ("CUFFT_INTERNAL_ERROR", "HIPFFT_INTERNAL_ERROR"), + ("CUFFT_EXEC_FAILED", "HIPFFT_EXEC_FAILED"), + ("CUFFT_SETUP_FAILED", "HIPFFT_SETUP_FAILED"), + ("CUFFT_INVALID_SIZE", "HIPFFT_INVALID_SIZE"), + ("CUFFT_UNALIGNED_DATA", "HIPFFT_UNALIGNED_DATA"), + ("CUFFT_INCOMPLETE_PARAMETER_LIST", "HIPFFT_INCOMPLETE_PARAMETER_LIST"), + ("CUFFT_INVALID_DEVICE", "HIPFFT_INVALID_DEVICE"), + ("CUFFT_PARSE_ERROR", "HIPFFT_PARSE_ERROR"), + ("CUFFT_NO_WORKSPACE", "HIPFFT_NO_WORKSPACE"), + ("CUFFT_NOT_IMPLEMENTED", "HIPFFT_NOT_IMPLEMENTED"), + ("CUFFT_LICENSE_ERROR", "HIPFFT_LICENSE_ERROR"), + ("CUFFT_NOT_SUPPORTED", "HIPFFT_NOT_SUPPORTED"), + ("cufftType_t", "hipfftType_t"), + ("cufftType", "hipfftType"), + ("CUFFT_R2C", "HIPFFT_R2C"), + ("CUFFT_C2R", "HIPFFT_C2R"), + ("CUFFT_C2C", "HIPFFT_C2C"), + ("CUFFT_D2Z", "HIPFFT_D2Z"), + ("CUFFT_Z2D", "HIPFFT_Z2D"), + ("CUFFT_Z2Z", "HIPFFT_Z2Z"), + ("cufftCompatibility_t", "hipfftCompatibility_t"), + ("cufftCompatibility", "hipfftCompatibility"), + ("CUFFT_COMPATIBILITY_FFTW_PADDING", "HIPFFT_COMPATIBILITY_FFTW_PADDING"), + ("cufftReal", "hipfftReal"), + ("cufftDoubleReal", "hipfftDoubleReal"), + ("cufftComplex", "hipfftComplex"), + ("cufftDoubleComplex", "hipfftDoubleComplex"), + ("cufftHandle", "hipfftHandle"), + ("cufftPlan1d", "hipfftPlan1d"), + ("cufftPlan2d", "hipfftPlan2d"), + ("cufftPlan3d", "hipfftPlan3d"), + ("cufftPlanMany", "hipfftPlanMany"), + ("cufftMakePlan1d", "hipfftMakePlan1d"), + ("cufftMakePlan2d", "hipfftMakePlan2d"), + ("cufftMakePlan3d", "hipfftMakePlan3d"), + ("cufftMakePlanMany", "hipfftMakePlanMany"), + ("cufftMakePlanMany64", "hipfftMakePlanMany64"), + ("cufftGetSizeMany64", "hipfftGetSizeMany64"), + ("cufftEstimate1d", "hipfftEstimate1d"), + ("cufftEstimate2d", "hipfftEstimate2d"), + ("cufftEstimate3d", "hipfftEstimate3d"), + ("cufftEstimateMany", "hipfftEstimateMany"), + ("cufftCreate", "hipfftCreate"), + ("cufftGetSize1d", "hipfftGetSize1d"), + ("cufftGetSize2d", "hipfftGetSize2d"), + ("cufftGetSize3d", "hipfftGetSize3d"), + ("cufftGetSizeMany", "hipfftGetSizeMany"), + ("cufftGetSize", "hipfftGetSize"), + ("cufftSetWorkArea", "hipfftSetWorkArea"), + ("cufftSetAutoAllocation", "hipfftSetAutoAllocation"), + ("cufftXtExec", "hipfftXtExec"), + ("cufftXtMakePlanMany", "hipfftXtMakePlanMany"), + ("cufftExecC2C", "hipfftExecC2C"), + ("cufftExecR2C", "hipfftExecR2C"), + ("cufftExecC2R", "hipfftExecC2R"), + ("cufftExecZ2Z", "hipfftExecZ2Z"), + ("cufftExecD2Z", "hipfftExecD2Z"), + ("cufftExecZ2D", "hipfftExecZ2D"), + ("cufftSetStream", "hipfftSetStream"), + ("cufftDestroy", "hipfftDestroy"), + ("cufftGetVersion", "hipfftGetVersion"), + ("cufftGetProperty", "hipfftGetProperty"), + ("nvrtcResult", "hiprtcResult"), + ("NVRTC_SUCCESS", "HIPRTC_SUCCESS"), + ("NVRTC_ERROR_OUT_OF_MEMORY", "HIPRTC_ERROR_OUT_OF_MEMORY"), + ("NVRTC_ERROR_PROGRAM_CREATION_FAILURE", "HIPRTC_ERROR_PROGRAM_CREATION_FAILURE"), + ("NVRTC_ERROR_INVALID_INPUT", "HIPRTC_ERROR_INVALID_INPUT"), + ("NVRTC_ERROR_INVALID_PROGRAM", "HIPRTC_ERROR_INVALID_PROGRAM"), + ("NVRTC_ERROR_COMPILATION", "HIPRTC_ERROR_COMPILATION"), + ("NVRTC_ERROR_BUILTIN_OPERATION_FAILURE", "HIPRTC_ERROR_BUILTIN_OPERATION_FAILURE"), + ("NVRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION", "HIPRTC_ERROR_NO_NAME_EXPRESSIONS_AFTER_COMPILATION"), + ("NVRTC_ERROR_NAME_EXPRESSION_NOT_VALID", "HIPRTC_ERROR_NAME_EXPRESSION_NOT_VALID"), + ("NVRTC_ERROR_INTERNAL_ERROR", "HIPRTC_ERROR_INTERNAL_ERROR"), + ("nvrtcGetErrorString", "hiprtcGetErrorString"), + ("nvrtcVersion", "hiprtcVersion"), + ("nvrtcProgram", "hiprtcProgram"), + ("nvrtcAddNameExpression", "hiprtcAddNameExpression"), + ("nvrtcCompileProgram", "hiprtcCompileProgram"), + ("nvrtcCreateProgram", "hiprtcCreateProgram"), + ("nvrtcDestroyProgram", "hiprtcDestroyProgram"), + ("nvrtcGetLoweredName", "hiprtcGetLoweredName"), + ("nvrtcGetProgramLog", "hiprtcGetProgramLog"), + ("nvrtcGetProgramLogSize", "hiprtcGetProgramLogSize"), + ("nvrtcGetPTX", "hiprtcGetCode"), + ("nvrtcGetPTXSize", "hiprtcGetCodeSize"), + ("nvrtcGetCUBIN", "hiprtcGetBitcode"), + ("nvrtcGetCUBINSize", "hiprtcGetBitcodeSize"), + ("thrust::cuda", "thrust::hip"), + ("cub::", "hipcub::"), + ("cub::ArgMax", "hipcub::ArgMax"), + ("cub::ArgMin", "hipcub::ArgMin"), + ("cub::BLOCK_SCAN_WARP_SCANS", "hipcub::BLOCK_SCAN_WARP_SCANS"), + ("cub::BLOCK_REDUCE_WARP_REDUCTIONS", "hipcub::BLOCK_REDUCE_WARP_REDUCTIONS"), + ("cub::BLOCK_STORE_WARP_TRANSPOSE", "hipcub::BLOCK_STORE_WARP_TRANSPOSE"), + ("cub::BLOCK_LOAD_DIRECT", "hipcub::BLOCK_LOAD_DIRECT"), + ("cub::BLOCK_STORE_DIRECT", "hipcub::BLOCK_STORE_DIRECT"), + ("cub::BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY", "hipcub::BLOCK_REDUCE_RAKING_COMMUTATIVE_ONLY"), + ("cub::BlockReduce", "hipcub::BlockReduce"), + ("cub::BlockScan", "hipcub::BlockScan"), + ("cub::BlockLoad", "hipcub::BlockLoad"), + ("cub::BlockStore", "hipcub::BlockStore"), + ("cub::BlockRakingLayout", "hipcub::BlockRakingLayout"), + ("cub::BlockRadixSort", "hipcub::BlockRadixSort"), + ("cub::Uninitialized", "hipcub::Uninitialized"), + ("cub::RowMajorTid", "hipcub::RowMajorTid"), + ("cub::CachingDeviceAllocator", "hipcub::CachingDeviceAllocator"), + ("cub::CountingInputIterator", "hipcub::CountingInputIterator"), + ("cub::DeviceRadixSort", "hipcub::DeviceRadixSort"), + ("cub::DeviceReduce", "hipcub::DeviceReduce"), + ("cub::DeviceRunLengthEncode", "hipcub::DeviceRunLengthEncode"), + ("cub::DeviceScan", "hipcub::DeviceScan"), + ("cub::DeviceSegmentedRadixSort", "hipcub::DeviceSegmentedRadixSort"), + ("cub::DeviceSegmentedReduce", "hipcub::DeviceSegmentedReduce"), + ("cub::DeviceSelect", "hipcub::DeviceSelect"), + ("cub::FpLimits", "hipcub::FpLimits"), + ("cub::KeyValuePair", "hipcub::KeyValuePair"), + ("cub::Max", "hipcub::Max"), + ("cub::Min", "hipcub::Min"), + ("cub::Sum", "hipcub::Sum"), + ("cub::Log2", "hipcub::Log2"), + ("cub::LaneId", "hipcub::LaneId"), + ("cub::WarpMask", "hipcub::WarpMask"), + ("cub::ShuffleIndex", "hipcub::ShuffleIndex"), + ("cub::ShuffleDown", "hipcub::ShuffleDown"), + ("cub::ArgIndexInputIterator", "hipcub::ArgIndexInputIterator"), + ("cub::TransformInputIterator", "hipcub::TransformInputIterator"), + ("cub::WarpReduce", "hipcub::WarpReduce"), + ("cub::CTA_SYNC", "hipcub::CTA_SYNC"), + ("nvtxMark", "roctxMark"), + ("nvtxMarkA", "roctxMarkA"), + ("nvtxRangePushA", "roctxRangePushA"), + ("nvtxRangePush", "roctxRangePush"), + ("nvtxRangePop", "roctxRangePop"), + ("nvtxRangeStartA", "roctxRangeStartA"), + ("nvtxRangeStart", "roctxRangeStart"), + ("nvtxRangeEnd", "roctxRangeStop"), + ("nvtxRangeId_t", "int"), + ("nvmlReturn_t", "rsmi_status_t"), + ("NVML_SUCCESS", "RSMI_STATUS_SUCCESS"), + ("NVML_P2P_CAPS_INDEX_READ", "RSMI_STATUS_SUCCESS"), + ("NVML_P2P_STATUS_OK", "RSMI_STATUS_SUCCESS"), + ("NVML_ERROR_INSUFFICIENT_SIZE", "RSMI_STATUS_INSUFFICIENT_SIZE"), + ("nvmlDevice_t", "uint32_t"), + ("nvmlGpuP2PStatus_t", "bool"), + ("nvmlProcessInfo_t", "rsmi_process_info_t"), + ("nvmlGpuP2PCapsIndex_t", "uint32_t"), +]) + +CUDA_SPECIAL_MAP = collections.OrderedDict([ + ("cusparseStatus_t", "hipsparseStatus_t"), + ("cusparseHandle_t", "hipsparseHandle_t"), + ("cuComplex", "hipComplex"), + ("cuDoubleComplex", "hipDoubleComplex"), + ("CUSPARSE_POINTER_MODE_HOST", "HIPSPARSE_POINTER_MODE_HOST"), + ("cusparseOperation_t", "hipsparseOperation_t"), + ("cusparseCreateMatDescr", "hipsparseCreateMatDescr"), + ("cusparseCreate", "hipsparseCreate"), + ("cusparseDestroyMatDescr", "hipsparseDestroyMatDescr"), + ("cusparseDestroy", "hipsparseDestroy"), + ("cusparseGetVersion", "hipsparseGetVersion"), + ("cusparseXcoo2csr", "hipsparseXcoo2csr"), + ("cusparseMatDescr_t", "hipsparseMatDescr_t"), + ("cusparseDiagType_t", "hipsparseDiagType_t"), + ("CUSPARSE_DIAG_TYPE_UNIT", "HIPSPARSE_DIAG_TYPE_UNIT"), + ("CUSPARSE_DIAG_TYPE_NON_UNIT", "HIPSPARSE_DIAG_TYPE_NON_UNIT"), + ("cusparseSetMatDiagType", "hipsparseSetMatDiagType"), + ("cusparseFillMode_t", "hipsparseFillMode_t"), + ("CUSPARSE_FILL_MODE_UPPER", "HIPSPARSE_FILL_MODE_UPPER"), + ("CUSPARSE_FILL_MODE_LOWER", "HIPSPARSE_FILL_MODE_LOWER"), + ("cusparseSetMatFillMode", "hipsparseSetMatFillMode"), + ("cusparseDirection_t", "hipsparseDirection_t"), + ("CUSPARSE_DIRECTION_ROW", "HIPSPARSE_DIRECTION_ROW"), + ("CUSPARSE_DIRECTION_COLUMN", "HIPSPARSE_DIRECTION_COLUMN"), + ("cusparseSolvePolicy_t", "hipsparseSolvePolicy_t"), + ("CUSPARSE_SOLVE_POLICY_NO_LEVEL", "HIPSPARSE_SOLVE_POLICY_NO_LEVEL"), + ("CUSPARSE_SOLVE_POLICY_USE_LEVEL", "HIPSPARSE_SOLVE_POLICY_USE_LEVEL"), + ("cusparseCreateBsrsv2Info", "hipsparseCreateBsrsv2Info"), + ("cusparseCreateBsrsm2Info", "hipsparseCreateBsrsm2Info"), + ("cusparseDestroyBsrsv2Info", "hipsparseDestroyBsrsv2Info"), + ("cusparseDestroyBsrsm2Info", "hipsparseDestroyBsrsm2Info"), + ("cusparseSbsrmm", "hipsparseSbsrmm"), + ("cusparseDbsrmm", "hipsparseDbsrmm"), + ("cusparseCbsrmm", "hipsparseCbsrmm"), + ("cusparseZbsrmm", "hipsparseZbsrmm"), + ("cusparseSbsrmv", "hipsparseSbsrmv"), + ("cusparseDbsrmv", "hipsparseDbsrmv"), + ("cusparseCbsrmv", "hipsparseCbsrmv"), + ("cusparseZbsrmv", "hipsparseZbsrmv"), + ("cusparseSbsrsv2_bufferSize", "hipsparseSbsrsv2_bufferSize"), + ("cusparseDbsrsv2_bufferSize", "hipsparseDbsrsv2_bufferSize"), + ("cusparseCbsrsv2_bufferSize", "hipsparseCbsrsv2_bufferSize"), + ("cusparseZbsrsv2_bufferSize", "hipsparseZbsrsv2_bufferSize"), + ("cusparseSbsrsv2_analysis", "hipsparseSbsrsv2_analysis"), + ("cusparseDbsrsv2_analysis", "hipsparseDbsrsv2_analysis"), + ("cusparseCbsrsv2_analysis", "hipsparseCbsrsv2_analysis"), + ("cusparseZbsrsv2_analysis", "hipsparseZbsrsv2_analysis"), + ("cusparseSbsrsv2_solve", "hipsparseSbsrsv2_solve"), + ("cusparseDbsrsv2_solve", "hipsparseDbsrsv2_solve"), + ("cusparseCbsrsv2_solve", "hipsparseCbsrsv2_solve"), + ("cusparseZbsrsv2_solve", "hipsparseZbsrsv2_solve"), + ("cusparseSbsrsm2_bufferSize", "hipsparseSbsrsm2_bufferSize"), + ("cusparseDbsrsm2_bufferSize", "hipsparseDbsrsm2_bufferSize"), + ("cusparseCbsrsm2_bufferSize", "hipsparseCbsrsm2_bufferSize"), + ("cusparseZbsrsm2_bufferSize", "hipsparseZbsrsm2_bufferSize"), + ("cusparseSbsrsm2_analysis", "hipsparseSbsrsm2_analysis"), + ("cusparseDbsrsm2_analysis", "hipsparseDbsrsm2_analysis"), + ("cusparseCbsrsm2_analysis", "hipsparseCbsrsm2_analysis"), + ("cusparseZbsrsm2_analysis", "hipsparseZbsrsm2_analysis"), + ("cusparseSbsrsm2_solve", "hipsparseSbsrsm2_solve"), + ("cusparseDbsrsm2_solve", "hipsparseDbsrsm2_solve"), + ("cusparseCbsrsm2_solve", "hipsparseCbsrsm2_solve"), + ("cusparseZbsrsm2_solve", "hipsparseZbsrsm2_solve"), + ("cusparseScsrmm2", "hipsparseScsrmm2"), + ("cusparseDcsrmm2", "hipsparseDcsrmm2"), + ("cusparseCcsrmm2", "hipsparseCcsrmm2"), + ("cusparseZcsrmm2", "hipsparseZcsrmm2"), + ("cusparseScsrmm", "hipsparseScsrmm"), + ("cusparseDcsrmm", "hipsparseDcsrmm"), + ("cusparseCcsrmm", "hipsparseCcsrmm"), + ("cusparseZcsrmm", "hipsparseZcsrmm"), + ("cusparseXcsrsort_bufferSizeExt", "hipsparseXcsrsort_bufferSizeExt"), + ("cusparseCreateCsrgemm2Info", "hipsparseCreateCsrgemm2Info"), + ("cusparseDestroyCsrgemm2Info", "hipsparseDestroyCsrgemm2Info"), + ("cusparseXcsrgemm2Nnz", "hipsparseXcsrgemm2Nnz"), + ("cusparseDcsrgemm2_bufferSizeExt", "hipsparseDcsrgemm2_bufferSizeExt"), + ("cusparseScsrgemm2_bufferSizeExt", "hipsparseScsrgemm2_bufferSizeExt"), + ("cusparseDcsrgemm2", "hipsparseDcsrgemm2"), + ("cusparseScsrgemm2", "hipsparseScsrgemm2"), + ("cusparseScsrgemm", "hipsparseScsrgemm"), + ("cusparseDcsrgemm", "hipsparseDcsrgemm"), + ("cusparseCcsrgemm", "hipsparseCcsrgemm"), + ("cusparseZcsrgemm", "hipsparseZcsrgemm"), + ("cusparseSetPointerMode", "hipsparseSetPointerMode"), + ("cusparseXcsrgeam2Nnz", "hipsparseXcsrgeam2Nnz"), + ("cusparseScsrgeam", "hipsparseScsrgeam"), + ("cusparseDcsrgeam", "hipsparseDcsrgeam"), + ("cusparseCcsrgeam", "hipsparseCcsrgeam"), + ("cusparseZcsrgeam", "hipsparseZcsrgeam"), + ("cusparseScsrgeam2_bufferSizeExt", "hipsparseScsrgeam2_bufferSizeExt"), + ("cusparseDcsrgeam2_bufferSizeExt", "hipsparseDcsrgeam2_bufferSizeExt"), + ("cusparseCcsrgeam2_bufferSizeExt", "hipsparseCcsrgeam2_bufferSizeExt"), + ("cusparseZcsrgeam2_bufferSizeExt", "hipsparseZcsrgeam2_bufferSizeExt"), + ("cusparseScsrgeam2", "hipsparseScsrgeam2"), + ("cusparseDcsrgeam2", "hipsparseDcsrgeam2"), + ("cusparseCcsrgeam2", "hipsparseCcsrgeam2"), + ("cusparseZcsrgeam2", "hipsparseZcsrgeam2"), + ("cusparseXcsrsort", "hipsparseXcsrsort"), + ("cusparseXbsrsm2_zeroPivot", "hipsparseXbsrsm2_zeroPivot"), + ("cusparseXbsrsv2_zeroPivot", "hipsparseXbsrsv2_zeroPivot"), + ("cusparseXcoosort_bufferSizeExt", "hipsparseXcoosort_bufferSizeExt"), + ("cusparseXcoosortByRow", "hipsparseXcoosortByRow"), + ("cusparseSetStream", "hipsparseSetStream"), + ("cusparseGetStream", "hipsparseGetStream"), + ("cusparseCreateIdentityPermutation", "hipsparseCreateIdentityPermutation"), + ("cusparseSetMatIndexBase", "hipsparseSetMatIndexBase"), + ("cusparseSetMatType", "hipsparseSetMatType"), + ("cusparseSgthr", "hipsparseSgthr"), + ("cusparseDgthr", "hipsparseDgthr"), + ("cusparseCgthr", "hipsparseCgthr"), + ("cusparseZgthr", "hipsparseZgthr"), + ("cusparseScsrmv", "hipsparseScsrmv"), + ("cusparseDcsrmv", "hipsparseDcsrmv"), + ("cusparseCcsrmv", "hipsparseCcsrmv"), + ("cusparseZcsrmv", "hipsparseZcsrmv"), + ("cusparseSpMV", "hipsparseSpMV"), + ("cusparseSpMV_bufferSize", "hipsparseSpMV_bufferSize"), + ("cusparseSpMM", "hipsparseSpMM"), + ("cusparseSpMM_bufferSize", "hipsparseSpMM_bufferSize"), + ("cusparseCreateCsrsv2Info", "hipsparseCreateCsrsv2Info"), + ("cusparseDestroyCsrsv2Info", "hipsparseDestroyCsrsv2Info"), + ("cusparseScsrsv2_bufferSize", "hipsparseScsrsv2_bufferSize"), + ("cusparseDcsrsv2_bufferSize", "hipsparseDcsrsv2_bufferSize"), + ("cusparseCcsrsv2_bufferSize", "hipsparseCcsrsv2_bufferSize"), + ("cusparseZcsrsv2_bufferSize", "hipsparseZcsrsv2_bufferSize"), + ("cusparseScsrsv2_analysis", "hipsparseScsrsv2_analysis"), + ("cusparseDcsrsv2_analysis", "hipsparseDcsrsv2_analysis"), + ("cusparseCcsrsv2_analysis", "hipsparseCcsrsv2_analysis"), + ("cusparseZcsrsv2_analysis", "hipsparseZcsrsv2_analysis"), + ("cusparseScsrsv2_solve", "hipsparseScsrsv2_solve"), + ("cusparseDcsrsv2_solve", "hipsparseDcsrsv2_solve"), + ("cusparseCcsrsv2_solve", "hipsparseCcsrsv2_solve"), + ("cusparseZcsrsv2_solve", "hipsparseZcsrsv2_solve"), + ("cusparseXcsrsv2_zeroPivot", "hipsparseXcsrsv2_zeroPivot"), + ("cusparseCreateCsrsm2Info", "hipsparseCreateCsrsm2Info"), + ("cusparseDestroyCsrsm2Info", "hipsparseDestroyCsrsm2Info"), + ("cusparseScsrsm2_bufferSizeExt", "hipsparseScsrsm2_bufferSizeExt"), + ("cusparseDcsrsm2_bufferSizeExt", "hipsparseDcsrsm2_bufferSizeExt"), + ("cusparseCcsrsm2_bufferSizeExt", "hipsparseCcsrsm2_bufferSizeExt"), + ("cusparseZcsrsm2_bufferSizeExt", "hipsparseZcsrsm2_bufferSizeExt"), + ("cusparseScsrsm2_analysis", "hipsparseScsrsm2_analysis"), + ("cusparseDcsrsm2_analysis", "hipsparseDcsrsm2_analysis"), + ("cusparseCcsrsm2_analysis", "hipsparseCcsrsm2_analysis"), + ("cusparseZcsrsm2_analysis", "hipsparseZcsrsm2_analysis"), + ("cusparseScsrsm2_solve", "hipsparseScsrsm2_solve"), + ("cusparseDcsrsm2_solve", "hipsparseDcsrsm2_solve"), + ("cusparseCcsrsm2_solve", "hipsparseCcsrsm2_solve"), + ("cusparseZcsrsm2_solve", "hipsparseZcsrsm2_solve"), + ("cusparseXcsrsm2_zeroPivot", "hipsparseXcsrsm2_zeroPivot"), + ("cusparseXcsrgeamNnz", "hipsparseXcsrgeamNnz"), + ("cusparseXcsrgemmNnz", "hipsparseXcsrgemmNnz"), + ("cusparseCcsrgemm2_bufferSizeExt", "hipsparseCcsrgemm2_bufferSizeExt"), + ("cusparseZcsrgemm2_bufferSizeExt", "hipsparseZcsrgemm2_bufferSizeExt"), + ("cusparseCcsrgemm2", "hipsparseCcsrgemm2"), + ("cusparseZcsrgemm2", "hipsparseZcsrgemm2"), + ("cusparseScsc2dense", "hipsparseScsc2dense"), + ("cusparseDcsc2dense", "hipsparseDcsc2dense"), + ("cusparseCcsc2dense", "hipsparseCcsc2dense"), + ("cusparseZcsc2dense", "hipsparseZcsc2dense"), + ("cusparseXcsr2coo", "hipsparseXcsr2coo"), + ("cusparseScsr2csc", "hipsparseScsr2csc"), + ("cusparseDcsr2csc", "hipsparseDcsr2csc"), + ("cusparseCcsr2csc", "hipsparseCcsr2csc"), + ("cusparseZcsr2csc", "hipsparseZcsr2csc"), + ("cusparseScsr2dense", "hipsparseScsr2dense"), + ("cusparseDcsr2dense", "hipsparseDcsr2dense"), + ("cusparseCcsr2dense", "hipsparseCcsr2dense"), + ("cusparseZcsr2dense", "hipsparseZcsr2dense"), + ("cusparseSnnz_compress", "hipsparseSnnz_compress"), + ("cusparseDnnz_compress", "hipsparseDnnz_compress"), + ("cusparseCnnz_compress", "hipsparseCnnz_compress"), + ("cusparseZnnz_compress", "hipsparseZnnz_compress"), + ("cusparseScsr2csr_compress", "hipsparseScsr2csr_compress"), + ("cusparseDcsr2csr_compress", "hipsparseDcsr2csr_compress"), + ("cusparseCcsr2csr_compress", "hipsparseCcsr2csr_compress"), + ("cusparseZcsr2csr_compress", "hipsparseZcsr2csr_compress"), + ("cusparseSdense2csc", "hipsparseSdense2csc"), + ("cusparseDdense2csc", "hipsparseDdense2csc"), + ("cusparseCdense2csc", "hipsparseCdense2csc"), + ("cusparseZdense2csc", "hipsparseZdense2csc"), + ("cusparseSdense2csr", "hipsparseSdense2csr"), + ("cusparseDdense2csr", "hipsparseDdense2csr"), + ("cusparseCdense2csr", "hipsparseCdense2csr"), + ("cusparseZdense2csr", "hipsparseZdense2csr"), + ("cusparseSnnz", "hipsparseSnnz"), + ("cusparseDnnz", "hipsparseDnnz"), + ("cusparseCnnz", "hipsparseCnnz"), + ("cusparseZnnz", "hipsparseZnnz"), + ("cusparseXcoosortByColumn", "hipsparseXcoosortByColumn"), + ("cusparseXcscsort_bufferSizeExt", "hipsparseXcscsort_bufferSizeExt"), + ("cusparseXcscsort", "hipsparseXcscsort"), + ("cusparseCreateCsrilu02Info", "hipsparseCreateCsrilu02Info"), + ("cusparseDestroyCsrilu02Info", "hipsparseDestroyCsrilu02Info"), + ("cusparseCreateBsrilu02Info", "hipsparseCreateBsrilu02Info"), + ("cusparseDestroyBsrilu02Info", "hipsparseDestroyBsrilu02Info"), + ("cusparseCreateCsric02Info", "hipsparseCreateCsric02Info"), + ("cusparseDestroyCsric02Info", "hipsparseDestroyCsric02Info"), + ("cusparseCreateBsric02Info", "hipsparseCreateBsric02Info"), + ("cusparseDestroyBsric02Info", "hipsparseDestroyBsric02Info"), + ("cusparseScsrilu02_numericBoost", "hipsparseScsrilu02_numericBoost"), + ("cusparseDcsrilu02_numericBoost", "hipsparseDcsrilu02_numericBoost"), + ("cusparseCcsrilu02_numericBoost", "hipsparseCcsrilu02_numericBoost"), + ("cusparseZcsrilu02_numericBoost", "hipsparseZcsrilu02_numericBoost"), + ("cusparseXcsrilu02_zeroPivot", "hipsparseXcsrilu02_zeroPivot"), + ("cusparseScsrilu02_bufferSize", "hipsparseScsrilu02_bufferSize"), + ("cusparseDcsrilu02_bufferSize", "hipsparseDcsrilu02_bufferSize"), + ("cusparseCcsrilu02_bufferSize", "hipsparseCcsrilu02_bufferSize"), + ("cusparseZcsrilu02_bufferSize", "hipsparseZcsrilu02_bufferSize"), + ("cusparseScsrilu02_analysis", "hipsparseScsrilu02_analysis"), + ("cusparseDcsrilu02_analysis", "hipsparseDcsrilu02_analysis"), + ("cusparseCcsrilu02_analysis", "hipsparseCcsrilu02_analysis"), + ("cusparseZcsrilu02_analysis", "hipsparseZcsrilu02_analysis"), + ("cusparseScsrilu02", "hipsparseScsrilu02"), + ("cusparseDcsrilu02", "hipsparseDcsrilu02"), + ("cusparseCcsrilu02", "hipsparseCcsrilu02"), + ("cusparseZcsrilu02", "hipsparseZcsrilu02"), + ("cusparseSbsrilu02_numericBoost", "hipsparseSbsrilu02_numericBoost"), + ("cusparseDbsrilu02_numericBoost", "hipsparseDbsrilu02_numericBoost"), + ("cusparseCbsrilu02_numericBoost", "hipsparseCbsrilu02_numericBoost"), + ("cusparseZbsrilu02_numericBoost", "hipsparseZbsrilu02_numericBoost"), + ("cusparseXbsrilu02_zeroPivot", "hipsparseXbsrilu02_zeroPivot"), + ("cusparseSbsrilu02_bufferSize", "hipsparseSbsrilu02_bufferSize"), + ("cusparseDbsrilu02_bufferSize", "hipsparseDbsrilu02_bufferSize"), + ("cusparseCbsrilu02_bufferSize", "hipsparseCbsrilu02_bufferSize"), + ("cusparseZbsrilu02_bufferSize", "hipsparseZbsrilu02_bufferSize"), + ("cusparseSbsrilu02_analysis", "hipsparseSbsrilu02_analysis"), + ("cusparseDbsrilu02_analysis", "hipsparseDbsrilu02_analysis"), + ("cusparseCbsrilu02_analysis", "hipsparseCbsrilu02_analysis"), + ("cusparseZbsrilu02_analysis", "hipsparseZbsrilu02_analysis"), + ("cusparseSbsrilu02", "hipsparseSbsrilu02"), + ("cusparseDbsrilu02", "hipsparseDbsrilu02"), + ("cusparseCbsrilu02", "hipsparseCbsrilu02"), + ("cusparseZbsrilu02", "hipsparseZbsrilu02"), + ("cusparseXcsric02_zeroPivot", "hipsparseXcsric02_zeroPivot"), + ("cusparseScsric02_bufferSize", "hipsparseScsric02_bufferSize"), + ("cusparseDcsric02_bufferSize", "hipsparseDcsric02_bufferSize"), + ("cusparseCcsric02_bufferSize", "hipsparseCcsric02_bufferSize"), + ("cusparseZcsric02_bufferSize", "hipsparseZcsric02_bufferSize"), + ("cusparseScsric02_analysis", "hipsparseScsric02_analysis"), + ("cusparseDcsric02_analysis", "hipsparseDcsric02_analysis"), + ("cusparseCcsric02_analysis", "hipsparseCcsric02_analysis"), + ("cusparseZcsric02_analysis", "hipsparseZcsric02_analysis"), + ("cusparseScsric02", "hipsparseScsric02"), + ("cusparseDcsric02", "hipsparseDcsric02"), + ("cusparseCcsric02", "hipsparseCcsric02"), + ("cusparseZcsric02", "hipsparseZcsric02"), + ("cusparseXbsric02_zeroPivot", "hipsparseXbsric02_zeroPivot"), + ("cusparseSbsric02_bufferSize", "hipsparseSbsric02_bufferSize"), + ("cusparseDbsric02_bufferSize", "hipsparseDbsric02_bufferSize"), + ("cusparseCbsric02_bufferSize", "hipsparseCbsric02_bufferSize"), + ("cusparseZbsric02_bufferSize", "hipsparseZbsric02_bufferSize"), + ("cusparseSbsric02_analysis", "hipsparseSbsric02_analysis"), + ("cusparseDbsric02_analysis", "hipsparseDbsric02_analysis"), + ("cusparseCbsric02_analysis", "hipsparseCbsric02_analysis"), + ("cusparseZbsric02_analysis", "hipsparseZbsric02_analysis"), + ("cusparseSbsric02", "hipsparseSbsric02"), + ("cusparseDbsric02", "hipsparseDbsric02"), + ("cusparseCbsric02", "hipsparseCbsric02"), + ("cusparseZbsric02", "hipsparseZbsric02"), + ("cusparseSgtsv2_bufferSizeExt", "hipsparseSgtsv2_bufferSizeExt"), + ("cusparseDgtsv2_bufferSizeExt", "hipsparseDgtsv2_bufferSizeExt"), + ("cusparseCgtsv2_bufferSizeExt", "hipsparseCgtsv2_bufferSizeExt"), + ("cusparseZgtsv2_bufferSizeExt", "hipsparseZgtsv2_bufferSizeExt"), + ("cusparseSgtsv2", "hipsparseSgtsv2"), + ("cusparseDgtsv2", "hipsparseDgtsv2"), + ("cusparseCgtsv2", "hipsparseCgtsv2"), + ("cusparseZgtsv2", "hipsparseZgtsv2"), + ("cusparseSgtsv2_nopivot_bufferSizeExt", "hipsparseSgtsv2_nopivot_bufferSizeExt"), + ("cusparseDgtsv2_nopivot_bufferSizeExt", "hipsparseDgtsv2_nopivot_bufferSizeExt"), + ("cusparseCgtsv2_nopivot_bufferSizeExt", "hipsparseCgtsv2_nopivot_bufferSizeExt"), + ("cusparseZgtsv2_nopivot_bufferSizeExt", "hipsparseZgtsv2_nopivot_bufferSizeExt"), + ("cusparseSgtsv2_nopivot", "hipsparseSgtsv2_nopivot"), + ("cusparseDgtsv2_nopivot", "hipsparseDgtsv2_nopivot"), + ("cusparseCgtsv2_nopivot", "hipsparseCgtsv2_nopivot"), + ("cusparseZgtsv2_nopivot", "hipsparseZgtsv2_nopivot"), + ("cusparseSgtsv2StridedBatch_bufferSizeExt", "hipsparseSgtsv2StridedBatch_bufferSizeExt"), + ("cusparseDgtsv2StridedBatch_bufferSizeExt", "hipsparseDgtsv2StridedBatch_bufferSizeExt"), + ("cusparseCgtsv2StridedBatch_bufferSizeExt", "hipsparseCgtsv2StridedBatch_bufferSizeExt"), + ("cusparseZgtsv2StridedBatch_bufferSizeExt", "hipsparseZgtsv2StridedBatch_bufferSizeExt"), + ("cusparseSgtsv2StridedBatch", "hipsparseSgtsv2StridedBatch"), + ("cusparseDgtsv2StridedBatch", "hipsparseDgtsv2StridedBatch"), + ("cusparseCgtsv2StridedBatch", "hipsparseCgtsv2StridedBatch"), + ("cusparseZgtsv2StridedBatch", "hipsparseZgtsv2StridedBatch"), + ("cusparseSgtsvInterleavedBatch_bufferSizeExt", "hipsparseSgtsvInterleavedBatch_bufferSizeExt"), + ("cusparseDgtsvInterleavedBatch_bufferSizeExt", "hipsparseDgtsvInterleavedBatch_bufferSizeExt"), + ("cusparseCgtsvInterleavedBatch_bufferSizeExt", "hipsparseCgtsvInterleavedBatch_bufferSizeExt"), + ("cusparseZgtsvInterleavedBatch_bufferSizeExt", "hipsparseZgtsvInterleavedBatch_bufferSizeExt"), + ("cusparseSgtsvInterleavedBatch", "hipsparseSgtsvInterleavedBatch"), + ("cusparseDgtsvInterleavedBatch", "hipsparseDgtsvInterleavedBatch"), + ("cusparseCgtsvInterleavedBatch", "hipsparseCgtsvInterleavedBatch"), + ("cusparseZgtsvInterleavedBatch", "hipsparseZgtsvInterleavedBatch"), + ("cusparseSgpsvInterleavedBatch_bufferSizeExt", "hipsparseSgpsvInterleavedBatch_bufferSizeExt"), + ("cusparseDgpsvInterleavedBatch_bufferSizeExt", "hipsparseDgpsvInterleavedBatch_bufferSizeExt"), + ("cusparseCgpsvInterleavedBatch_bufferSizeExt", "hipsparseCgpsvInterleavedBatch_bufferSizeExt"), + ("cusparseZgpsvInterleavedBatch_bufferSizeExt", "hipsparseZgpsvInterleavedBatch_bufferSizeExt"), + ("cusparseSgpsvInterleavedBatch", "hipsparseSgpsvInterleavedBatch"), + ("cusparseDgpsvInterleavedBatch", "hipsparseDgpsvInterleavedBatch"), + ("cusparseCgpsvInterleavedBatch", "hipsparseCgpsvInterleavedBatch"), + ("cusparseZgpsvInterleavedBatch", "hipsparseZgpsvInterleavedBatch"), + ("cusparseCreateSpVec", "hipsparseCreateSpVec"), + ("cusparseDestroySpVec", "hipsparseDestroySpVec"), + ("cusparseSpVecGet", "hipsparseSpVecGet"), + ("cusparseSpVecGetIndexBase", "hipsparseSpVecGetIndexBase"), + ("cusparseSpVecGetValues", "hipsparseSpVecGetValues"), + ("cusparseSpVecSetValues", "hipsparseSpVecSetValues"), + ("cusparseCreateCooAoS", "hipsparseCreateCooAoS"), + ("cusparseCooGet", "hipsparseCooGet"), + ("cusparseCooAoSGet", "hipsparseCooAoSGet"), + ("cusparseCsrGet", "hipsparseCsrGet"), + ("cusparseSpMatGetFormat", "hipsparseSpMatGetFormat"), + ("cusparseSpMatGetIndexBase", "hipsparseSpMatGetIndexBase"), + ("cusparseSpMatGetValues", "hipsparseSpMatGetValues"), + ("cusparseSpMatGetStridedBatch", "hipsparseSpMatGetStridedBatch"), + ("cusparseSpMatSetStridedBatch", "hipsparseSpMatSetStridedBatch"), + ("cusparseDnVecGet", "hipsparseDnVecGet"), + ("cusparseDnVecGetValues", "hipsparseDnVecGetValues"), + ("cusparseDnVecSetValues", "hipsparseDnVecSetValues"), + ("cusparseDnMatGet", "hipsparseDnMatGet"), + ("cusparseDnMatGetValues", "hipsparseDnMatGetValues"), + ("cusparseDnMatSetValues", "hipsparseDnMatSetValues"), + ("cusparseDnMatGetStridedBatch", "hipsparseDnMatGetStridedBatch"), + ("cusparseSpVV_bufferSize", "hipsparseSpVV_bufferSize"), + ("cusparseSpVV", "hipsparseSpVV"), + ("cusparseCsr2cscEx2_bufferSize", "hipsparseCsr2cscEx2_bufferSize"), + ("cusparseCsr2cscEx2", "hipsparseCsr2cscEx2"), + ("cusparseCreateDnMat", "hipsparseCreateDnMat"), + ("cusparseDnMatSetStridedBatch", "hipsparseDnMatSetStridedBatch"), + ("cusparseCsrSetStridedBatch", "hipsparseCsrSetStridedBatch"), + ("cusparseCreateDnVec", "hipsparseCreateDnVec"), + ("cusparseCreateCsr", "hipsparseCreateCsr"), + ("cusparseDestroyDnMat", "hipsparseDestroyDnMat"), + ("cusparseDestroyDnVec", "hipsparseDestroyDnVec"), + ("cusparseDestroySpMat", "hipsparseDestroySpMat"), + ("cusparseSpGEMM_destroyDescr", "hipsparseSpGEMM_destroyDescr"), + ("cusparseCreateCoo", "hipsparseCreateCoo"), + ("cusparseSpGEMM_createDescr", "hipsparseSpGEMM_createDescr"), + ("cusparseSpGEMM_copy", "hipsparseSpGEMM_copy"), + ("cusparseSDDMM_bufferSize", "hipsparseSDDMM_bufferSize"), + ("cusparseSDDMM_preprocess", "hipsparseSDDMM_preprocess"), + ("cusparseSDDMM", "hipsparseSDDMM"), + ("cusparseSpGEMM_compute", "hipsparseSpGEMM_compute"), + ("cusparseSpGEMM_workEstimation", "hipsparseSpGEMM_workEstimation"), + ("cusparseSpMatGetSize", "hipsparseSpMatGetSize"), + ("cusparseCsrSetPointers", "hipsparseCsrSetPointers"), + ("cusparseCreateCsc", "hipsparseCreateCsc"), + ("cusparseSpMatSetValues", "hipsparseSpMatSetValues"), + ("cusparseSpMatSetAttribute", "hipsparseSpMatSetAttribute"), + ("cusparseSpSM_createDescr", "hipsparseSpSM_createDescr"), + ("cusparseSpSM_destroyDescr", "hipsparseSpSM_destroyDescr"), + ("cusparseSpSM_bufferSize", "hipsparseSpSM_bufferSize"), + ("cusparseSpSM_analysis", "hipsparseSpSM_analysis"), + ("cusparseSpSM_solve", "hipsparseSpSM_solve"), + ("cusparseSpSV_createDescr", "hipsparseSpSV_createDescr"), + ("cusparseSpSV_destroyDescr", "hipsparseSpSV_destroyDescr"), + ("cusparseSpSV_bufferSize", "hipsparseSpSV_bufferSize"), + ("cusparseSpSV_analysis", "hipsparseSpSV_analysis"), + ("cusparseSpSV_solve", "hipsparseSpSV_solve"), + ("cusparseSparseToDense_bufferSize", "hipsparseSparseToDense_bufferSize"), + ("cusparseSparseToDense", "hipsparseSparseToDense"), + ("cusparseDenseToSparse_bufferSize", "hipsparseDenseToSparse_bufferSize"), + ("cusparseDenseToSparse_analysis", "hipsparseDenseToSparse_analysis"), + ("cusparseDenseToSparse_convert", "hipsparseDenseToSparse_convert"), + ("cusparseSpMVAlg_t", "hipsparseSpMVAlg_t"), + ("cusparseSpMMAlg_t", "hipsparseSpMMAlg_t"), + ("cusparseIndexType_t", "hipsparseIndexType_t"), + ("cusparseDnMatDescr_t", "hipsparseDnMatDescr_t"), + ("cusparseDnVecDescr_t", "hipsparseDnVecDescr_t"), + ("cusparseSpMatDescr_t", "hipsparseSpMatDescr_t"), + ("cusparseSpGEMMDescr_t", "hipsparseSpGEMMDescr_t"), + ("CUSPARSE_INDEX_32I", "HIPSPARSE_INDEX_32I"), + ("CUSPARSE_INDEX_64I", "HIPSPARSE_INDEX_64I"), + ("CUSPARSE_ORDER_COL", "HIPSPARSE_ORDER_COL"), + ("CUSPARSE_MV_ALG_DEFAULT", "HIPSPARSE_MV_ALG_DEFAULT"), + ("CUSPARSE_MM_ALG_DEFAULT", "HIPSPARSE_MM_ALG_DEFAULT"), + ("CUSPARSE_SPMM_COO_ALG1", "HIPSPARSE_SPMM_COO_ALG1"), + ("CUSPARSE_SPMM_COO_ALG2", "HIPSPARSE_SPMM_COO_ALG2"), + ("CUSPARSE_COOMM_ALG1", "HIPSPARSE_COOMM_ALG1"), + ("CUSPARSE_COOMM_ALG2", "HIPSPARSE_COOMM_ALG2"), + ("CUSPARSE_COOMM_ALG3", "HIPSPARSE_COOMM_ALG3"), + ("CUSPARSE_COOMV_ALG", "HIPSPARSE_COOMV_ALG"), + ("CUSPARSE_CSRMM_ALG1", "HIPSPARSE_CSRMM_ALG1"), + ("CUSPARSE_SPMM_CSR_ALG1", "HIPSPARSE_SPMM_CSR_ALG1"), + ("CUSPARSE_SPMM_CSR_ALG2", "HIPSPARSE_SPMM_CSR_ALG2"), + ("CUSPARSE_SPMM_CSR_ALG3", "HIPSPARSE_SPMM_CSR_ALG3"), + ("CUSPARSE_SPGEMM_DEFAULT", "HIPSPARSE_SPGEMM_DEFAULT"), + ("CUSPARSE_SDDMM_ALG_DEFAULT", "HIPSPARSE_SDDMM_ALG_DEFAULT"), + ("CUSPARSE_STATUS_SUCCESS", "HIPSPARSE_STATUS_SUCCESS"), + ("CUSPARSE_STATUS_NOT_INITIALIZED", "HIPSPARSE_STATUS_NOT_INITIALIZED"), + ("CUSPARSE_STATUS_ALLOC_FAILED", "HIPSPARSE_STATUS_ALLOC_FAILED"), + ("CUSPARSE_STATUS_INVALID_VALUE", "HIPSPARSE_STATUS_INVALID_VALUE"), + ("CUSPARSE_STATUS_MAPPING_ERROR", "HIPSPARSE_STATUS_MAPPING_ERROR"), + ("CUSPARSE_STATUS_EXECUTION_FAILED", "HIPSPARSE_STATUS_EXECUTION_FAILED"), + ("CUSPARSE_STATUS_INTERNAL_ERROR", "HIPSPARSE_STATUS_INTERNAL_ERROR"), + ("CUSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED", "HIPSPARSE_STATUS_MATRIX_TYPE_NOT_SUPPORTED"), + ("CUSPARSE_STATUS_ARCH_MISMATCH", "HIPSPARSE_STATUS_ARCH_MISMATCH"), + ("CUSPARSE_STATUS_ZERO_PIVOT", "HIPSPARSE_STATUS_ZERO_PIVOT"), + ("CUSPARSE_OPERATION_TRANSPOSE", "HIPSPARSE_OPERATION_TRANSPOSE"), + ("CUSPARSE_OPERATION_NON_TRANSPOSE", "HIPSPARSE_OPERATION_NON_TRANSPOSE"), + ("CUSPARSE_OPERATION_CONJUGATE_TRANSPOSE", "HIPSPARSE_OPERATION_CONJUGATE_TRANSPOSE"), + ("CUSPARSE_INDEX_BASE_ZERO", "HIPSPARSE_INDEX_BASE_ZERO"), + ("CUSPARSE_INDEX_BASE_ONE", "HIPSPARSE_INDEX_BASE_ONE"), + ("CUSPARSE_MATRIX_TYPE_GENERAL", "HIPSPARSE_MATRIX_TYPE_GENERAL"), + ("cusparseGetErrorName", "hipsparseGetErrorName"), + ("cusparseOrder_t", "hipsparseOrder_t"), + ("cusparseSpGEMMAlg_t", "hipsparseSpGEMMAlg_t"), + ("cusparseCsr2CscAlg_t", "hipsparseCsr2CscAlg_t"), + ("cusparseGetErrorString", "hipsparseGetErrorString"), + ("cusparseGather", "hipsparseGather"), + ("cusparseSparseToDenseAlg_t", "hipsparseSparseToDenseAlg_t"), + ("cusparseDenseToSparseAlg_t", "hipsparseDenseToSparseAlg_t"), + ("cusparseIndexBase_t", "hipsparseIndexBase_t"), + ("cusparseMatrixType_t", "hipsparseMatrixType_t"), + ("cusparsePointerMode_t", "hipsparsePointerMode_t"), + ("cusparseAction_t", "hipsparseAction_t"), + ("cusparseFormat_t", "hipsparseFormat_t"), + ("cusparseSpSMAlg_t", "hipsparseSpSMAlg_t"), + ("cusparseSpSVAlg_t", "hipsparseSpSVAlg_t"), + ("cusparseSpMatAttribute_t", "hipsparseSpMatAttribute_t"), + ("cusparseSpVecDescr_t", "hipsparseSpVecDescr_t"), + ("cusparseSpSMDescr_t", "hipsparseSpSMDescr_t"), + ("cusparseSpSVDescr_t", "hipsparseSpSVDescr_t"), + ("CUSPARSE_POINTER_MODE_DEVICE", "HIPSPARSE_POINTER_MODE_DEVICE"), + ("CUSPARSE_ACTION_SYMBOLIC", "HIPSPARSE_ACTION_SYMBOLIC"), + ("CUSPARSE_ACTION_NUMERIC", "HIPSPARSE_ACTION_NUMERIC"), + ("CUSPARSE_MATRIX_TYPE_SYMMETRIC", "HIPSPARSE_MATRIX_TYPE_SYMMETRIC"), + ("CUSPARSE_MATRIX_TYPE_HERMITIAN", "HIPSPARSE_MATRIX_TYPE_HERMITIAN"), + ("CUSPARSE_MATRIX_TYPE_TRIANGULAR", "HIPSPARSE_MATRIX_TYPE_TRIANGULAR"), + ("CUSPARSE_FORMAT_CSR", "HIPSPARSE_FORMAT_CSR"), + ("CUSPARSE_FORMAT_CSC", "HIPSPARSE_FORMAT_CSC"), + ("CUSPARSE_FORMAT_COO", "HIPSPARSE_FORMAT_COO"), + ("CUSPARSE_FORMAT_COO_AOS", "HIPSPARSE_FORMAT_COO_AOS"), + ("CUSPARSE_ORDER_ROW", "HIPSPARSE_ORDER_ROW"), + ("CUSPARSE_CSRMV_ALG1", "HIPSPARSE_CSRMV_ALG1"), + ("CUSPARSE_CSRMV_ALG2", "HIPSPARSE_CSRMV_ALG2"), + ("CUSPARSE_INDEX_16U", "HIPSPARSE_INDEX_16U"), + ("CUSPARSE_SPMAT_FILL_MODE", "HIPSPARSE_SPMAT_FILL_MODE"), + ("CUSPARSE_SPMAT_DIAG_TYPE", "HIPSPARSE_SPMAT_DIAG_TYPE"), + ("CUSPARSE_CSR2CSC_ALG1", "HIPSPARSE_CSR2CSC_ALG1"), + ("CUSPARSE_CSR2CSC_ALG2", "HIPSPARSE_CSR2CSC_ALG2"), + ("CUSPARSE_SPSM_ALG_DEFAULT", "HIPSPARSE_SPSM_ALG_DEFAULT"), + ("CUSPARSE_SPSV_ALG_DEFAULT", "HIPSPARSE_SPSV_ALG_DEFAULT"), + ("CUSPARSE_SPARSETODENSE_ALG_DEFAULT", "HIPSPARSE_SPARSETODENSE_ALG_DEFAULT"), + ("CUSPARSE_DENSETOSPARSE_ALG_DEFAULT", "HIPSPARSE_DENSETOSPARSE_ALG_DEFAULT"), + ("cuSPARSELt", "hipSPARSELt"), + ("AT_CUSPARSELT_ENABLED", "AT_HIPSPARSELT_ENABLED"), + ("CUSPARSELT_SPARSITY_50_PERCENT", "HIPSPARSELT_SPARSITY_50_PERCENT"), + ("cusparseComputeType", "hipsparseLtComputetype_t"), + ("CUSPARSE_COMPUTE_32F", "HIPSPARSELT_COMPUTE_32F"), + ("CUSPARSE_COMPUTE_16F", "HIPSPARSELT_COMPUTE_16F"), + ("CUSPARSE_COMPUTE_32I", "HIPSPARSELT_COMPUTE_32I"), + ("CUSPARSE_COMPUTE_TF32", "HIPSPARSELT_COMPUTE_TF32"), + ("CUSPARSELT_MATMUL_BIAS_POINTER", "HIPSPARSELT_MATMUL_BIAS_POINTER"), + ("CUSPARSELT_MATMUL_ALG_DEFAULT", "HIPSPARSELT_MATMUL_ALG_DEFAULT"), + ("CUSPARSELT_MATMUL_ALG_CONFIG_ID", "HIPSPARSELT_MATMUL_ALG_CONFIG_ID"), + ("CUSPARSELT_MATMUL_ALG_CONFIG_MAX_ID", "HIPSPARSELT_MATMUL_ALG_CONFIG_MAX_ID"), + ("CUSPARSELT_MATMUL_ALPHA_VECTOR_SCALING", "HIPSPARSELT_MATMUL_ALPHA_VECTOR_SCALING"), + ("CUSPARSELT_MATMUL_SPLIT_K", "HIPSPARSELT_MATMUL_SPLIT_K"), + ("CUSPARSELT_MATMUL_SPLIT_K_MODE", "HIPSPARSELT_MATMUL_SPLIT_K_MODE"), + ("cusparseLtHandle_t", "hipsparseLtHandle_t"), + ("cusparseLtMatDescriptor_t", "hipsparseLtMatDescriptor_t"), + ("cusparseLtInit", "hipsparseLtInit"), + ("cusparseLtStructuredDescriptorInit", "hipsparseLtStructuredDescriptorInit"), + ("cusparseLtSplitKMode_t", "hipsparseLtSplitKMode_t"), + ("cusparseLtSpMMACompressedSize2", "hipsparseLtSpMMACompressedSize2"), + ("cusparseLtSpMMACompress2", "hipsparseLtSpMMACompress2"), + ("cusparseLtMatmulDescriptor_t", "hipsparseLtMatmulDescriptor_t"), + ("cusparseLtMatmulPlan_t", "hipsparseLtMatmulPlan_t"), + ("cusparseLtMatmulAlgSelection_t", "hipsparseLtMatmulAlgSelection_t"), + ("cusparseLtDenseDescriptorInit", "hipsparseLtDenseDescriptorInit"), + ("cusparseLtMatmulDescriptorInit", "hipsparseLtMatmulDescriptorInit"), + ("cusparseLtMatmulDescSetAttribute", "hipsparseLtMatmulDescSetAttribute"), + ("cusparseLtMatmulAlgSelectionInit", "hipsparseLtMatmulAlgSelectionInit"), + ("cusparseLtMatmulAlgSetAttribute", "hipsparseLtMatmulAlgSetAttribute"), + ("cusparseLtMatmulPlanInit", "hipsparseLtMatmulPlanInit"), + ("cusparseLtMatmulGetWorkspace", "hipsparseLtMatmulGetWorkspace"), + ("cusparseLtMatmulSearch", "hipsparseLtMatmulSearch"), + ("cusparseLtMatmulAlgGetAttribute", "hipsparseLtMatmulAlgGetAttribute"), + ("cusparseLtMatmul", "hipsparseLtMatmul"), + ("cusparseLtMatDescriptorDestroy", "hipsparseLtMatDescriptorDestroy"), + ("cusparseLtMatmulPlanDestroy", "hipsparseLtMatmulPlanDestroy"), + ("cusolverEigMode_t", "hipsolverEigMode_t"), + ("cusolverEigType_t", "hipsolverEigType_t"), + ("CUSOLVER_EIG_MODE_VECTOR", "HIPSOLVER_EIG_MODE_VECTOR"), + ("CUSOLVER_EIG_MODE_NOVECTOR", "HIPSOLVER_EIG_MODE_NOVECTOR"), + ("CUSOLVER_EIG_TYPE_1", "HIPSOLVER_EIG_TYPE_1"), + ("CUSOLVER_EIG_TYPE_2", "HIPSOLVER_EIG_TYPE_2"), + ("CUSOLVER_EIG_TYPE_3", "HIPSOLVER_EIG_TYPE_3"), + ("syevjInfo_t", "hipsolverSyevjInfo_t"), + ("cusolverDnCreateSyevjInfo", "hipsolverDnCreateSyevjInfo"), + ("cusolverDnXsyevjSetSortEig", "hipsolverDnXsyevjSetSortEig"), + ("cusolverDnDestroySyevjInfo", "hipsolverDnDestroySyevjInfo"), + ("gesvdjInfo_t", "hipsolverGesvdjInfo_t"), + ("cusolverDnCreateGesvdjInfo", "hipsolverDnCreateGesvdjInfo"), + ("cusolverDnXgesvdjSetSortEig", "hipsolverDnXgesvdjSetSortEig"), + ("cusolverDnDestroyGesvdjInfo", "hipsolverDnDestroyGesvdjInfo"), + ("cusolverDnHandle_t", "hipsolverDnHandle_t"), + ("cusolverDnCreate", "hipsolverDnCreate"), + ("cusolverDnSetStream", "hipsolverDnSetStream"), + ("cusolverDnGetStream", "hipsolverDnGetStream"), + ("cusolverDnDestroy", "hipsolverDnDestroy"), + ("cusolverDnParams_t", "hipsolverDnParams_t"), + ("cusolverDnCgeqrf", "hipsolverDnCgeqrf"), + ("cusolverDnCgeqrf_bufferSize", "hipsolverDnCgeqrf_bufferSize"), + ("cusolverDnCgesvd", "hipsolverDnCgesvd"), + ("cusolverDnCgesvd_bufferSize", "hipsolverDnCgesvd_bufferSize"), + ("cusolverDnCgesvdj", "hipsolverDnCgesvdj"), + ("cusolverDnCgesvdjBatched", "hipsolverDnCgesvdjBatched"), + ("cusolverDnCgesvdjBatched_bufferSize", "hipsolverDnCgesvdjBatched_bufferSize"), + ("cusolverDnCgesvdj_bufferSize", "hipsolverDnCgesvdj_bufferSize"), + ("cusolverDnCgetrf", "hipsolverDnCgetrf"), + ("cusolverDnCgetrf_bufferSize", "hipsolverDnCgetrf_bufferSize"), + ("cusolverDnCgetrs", "hipsolverDnCgetrs"), + ("cusolverDnCheevd", "hipsolverDnCheevd"), + ("cusolverDnCheevd_bufferSize", "hipsolverDnCheevd_bufferSize"), + ("cusolverDnCheevj", "hipsolverDnCheevj"), + ("cusolverDnCheevjBatched", "hipsolverDnCheevjBatched"), + ("cusolverDnCheevjBatched_bufferSize", "hipsolverDnCheevjBatched_bufferSize"), + ("cusolverDnCheevj_bufferSize", "hipsolverDnCheevj_bufferSize"), + ("cusolverDnCpotrf", "hipsolverDnCpotrf"), + ("cusolverDnCpotrfBatched", "hipsolverDnCpotrfBatched"), + ("cusolverDnCpotrf_bufferSize", "hipsolverDnCpotrf_bufferSize"), + ("cusolverDnCpotrs", "hipsolverDnCpotrs"), + ("cusolverDnCpotrsBatched", "hipsolverDnCpotrsBatched"), + ("cusolverDnCungqr", "hipsolverDnCungqr"), + ("cusolverDnCungqr_bufferSize", "hipsolverDnCungqr_bufferSize"), + ("cusolverDnCunmqr", "hipsolverDnCunmqr"), + ("cusolverDnCunmqr_bufferSize", "hipsolverDnCunmqr_bufferSize"), + ("cusolverDnDgeqrf", "hipsolverDnDgeqrf"), + ("cusolverDnDgeqrf_bufferSize", "hipsolverDnDgeqrf_bufferSize"), + ("cusolverDnDgesvd", "hipsolverDnDgesvd"), + ("cusolverDnDgesvd_bufferSize", "hipsolverDnDgesvd_bufferSize"), + ("cusolverDnDgesvdj", "hipsolverDnDgesvdj"), + ("cusolverDnDgesvdjBatched", "hipsolverDnDgesvdjBatched"), + ("cusolverDnDgesvdjBatched_bufferSize", "hipsolverDnDgesvdjBatched_bufferSize"), + ("cusolverDnDgesvdj_bufferSize", "hipsolverDnDgesvdj_bufferSize"), + ("cusolverDnDgetrf", "hipsolverDnDgetrf"), + ("cusolverDnDgetrf_bufferSize", "hipsolverDnDgetrf_bufferSize"), + ("cusolverDnDgetrs", "hipsolverDnDgetrs"), + ("cusolverDnDorgqr", "hipsolverDnDorgqr"), + ("cusolverDnDorgqr_bufferSize", "hipsolverDnDorgqr_bufferSize"), + ("cusolverDnDormqr", "hipsolverDnDormqr"), + ("cusolverDnDormqr_bufferSize", "hipsolverDnDormqr_bufferSize"), + ("cusolverDnDpotrf", "hipsolverDnDpotrf"), + ("cusolverDnDpotrfBatched", "hipsolverDnDpotrfBatched"), + ("cusolverDnDpotrf_bufferSize", "hipsolverDnDpotrf_bufferSize"), + ("cusolverDnDpotrs", "hipsolverDnDpotrs"), + ("cusolverDnDpotrsBatched", "hipsolverDnDpotrsBatched"), + ("cusolverDnDsyevd", "hipsolverDnDsyevd"), + ("cusolverDnDsyevd_bufferSize", "hipsolverDnDsyevd_bufferSize"), + ("cusolverDnDsyevj", "hipsolverDnDsyevj"), + ("cusolverDnDsyevjBatched", "hipsolverDnDsyevjBatched"), + ("cusolverDnDsyevjBatched_bufferSize", "hipsolverDnDsyevjBatched_bufferSize"), + ("cusolverDnDsyevj_bufferSize", "hipsolverDnDsyevj_bufferSize"), + ("cusolverDnSgeqrf", "hipsolverDnSgeqrf"), + ("cusolverDnSgeqrf_bufferSize", "hipsolverDnSgeqrf_bufferSize"), + ("cusolverDnSgesvd", "hipsolverDnSgesvd"), + ("cusolverDnSgesvd_bufferSize", "hipsolverDnSgesvd_bufferSize"), + ("cusolverDnSgesvdj", "hipsolverDnSgesvdj"), + ("cusolverDnSgesvdjBatched", "hipsolverDnSgesvdjBatched"), + ("cusolverDnSgesvdjBatched_bufferSize", "hipsolverDnSgesvdjBatched_bufferSize"), + ("cusolverDnSgesvdj_bufferSize", "hipsolverDnSgesvdj_bufferSize"), + ("cusolverDnSgetrf", "hipsolverDnSgetrf"), + ("cusolverDnSgetrf_bufferSize", "hipsolverDnSgetrf_bufferSize"), + ("cusolverDnSgetrs", "hipsolverDnSgetrs"), + ("cusolverDnSorgqr", "hipsolverDnSorgqr"), + ("cusolverDnSorgqr_bufferSize", "hipsolverDnSorgqr_bufferSize"), + ("cusolverDnSormqr", "hipsolverDnSormqr"), + ("cusolverDnSormqr_bufferSize", "hipsolverDnSormqr_bufferSize"), + ("cusolverDnSpotrf", "hipsolverDnSpotrf"), + ("cusolverDnSpotrfBatched", "hipsolverDnSpotrfBatched"), + ("cusolverDnSpotrf_bufferSize", "hipsolverDnSpotrf_bufferSize"), + ("cusolverDnSpotrs", "hipsolverDnSpotrs"), + ("cusolverDnSpotrsBatched", "hipsolverDnSpotrsBatched"), + ("cusolverDnSsyevd", "hipsolverDnSsyevd"), + ("cusolverDnSsyevd_bufferSize", "hipsolverDnSsyevd_bufferSize"), + ("cusolverDnSsyevj", "hipsolverDnSsyevj"), + ("cusolverDnSsyevjBatched", "hipsolverDnSsyevjBatched"), + ("cusolverDnSsyevjBatched_bufferSize", "hipsolverDnSsyevjBatched_bufferSize"), + ("cusolverDnSsyevj_bufferSize", "hipsolverDnSsyevj_bufferSize"), + ("cusolverDnXgeqrf", "hipsolverDnXgeqrf"), + ("cusolverDnXgeqrf_bufferSize", "hipsolverDnXgeqrf_bufferSize"), + ("cusolverDnXpotrf", "hipsolverDnXpotrf"), + ("cusolverDnXpotrf_bufferSize", "hipsolverDnXpotrf_bufferSize"), + ("cusolverDnXpotrs", "hipsolverDnXpotrs"), + ("cusolverDnXsyevd", "hipsolverDnXsyevd"), + ("cusolverDnXsyevd_bufferSize", "hipsolverDnXsyevd_bufferSize"), + ("cusolverDnZgeqrf", "hipsolverDnZgeqrf"), + ("cusolverDnZgeqrf_bufferSize", "hipsolverDnZgeqrf_bufferSize"), + ("cusolverDnZgesvd", "hipsolverDnZgesvd"), + ("cusolverDnZgesvd_bufferSize", "hipsolverDnZgesvd_bufferSize"), + ("cusolverDnZgesvdj", "hipsolverDnZgesvdj"), + ("cusolverDnZgesvdjBatched", "hipsolverDnZgesvdjBatched"), + ("cusolverDnZgesvdjBatched_bufferSize", "hipsolverDnZgesvdjBatched_bufferSize"), + ("cusolverDnZgesvdj_bufferSize", "hipsolverDnZgesvdj_bufferSize"), + ("cusolverDnZgetrf", "hipsolverDnZgetrf"), + ("cusolverDnZgetrf_bufferSize", "hipsolverDnZgetrf_bufferSize"), + ("cusolverDnZgetrs", "hipsolverDnZgetrs"), + ("cusolverDnZheevd", "hipsolverDnZheevd"), + ("cusolverDnZheevd_bufferSize", "hipsolverDnZheevd_bufferSize"), + ("cusolverDnZheevj", "hipsolverDnZheevj"), + ("cusolverDnZheevjBatched", "hipsolverDnZheevjBatched"), + ("cusolverDnZheevjBatched_bufferSize", "hipsolverDnZheevjBatched_bufferSize"), + ("cusolverDnZheevj_bufferSize", "hipsolverDnZheevj_bufferSize"), + ("cusolverDnZpotrf", "hipsolverDnZpotrf"), + ("cusolverDnZpotrfBatched", "hipsolverDnZpotrfBatched"), + ("cusolverDnZpotrf_bufferSize", "hipsolverDnZpotrf_bufferSize"), + ("cusolverDnZpotrs", "hipsolverDnZpotrs"), + ("cusolverDnZpotrsBatched", "hipsolverDnZpotrsBatched"), + ("cusolverDnZungqr", "hipsolverDnZungqr"), + ("cusolverDnZungqr_bufferSize", "hipsolverDnZungqr_bufferSize"), + ("cusolverDnZunmqr", "hipsolverDnZunmqr"), + ("cusolverDnZunmqr_bufferSize", "hipsolverDnZunmqr_bufferSize"), + ("cusolverDnDsytrf_bufferSize", "hipsolverDnDsytrf_bufferSize"), + ("cusolverDnSsytrf_bufferSize", "hipsolverDnSsytrf_bufferSize"), + ("cusolverDnZsytrf_bufferSize", "hipsolverDnZsytrf_bufferSize"), + ("cusolverDnCsytrf_bufferSize", "hipsolverDnCsytrf_bufferSize"), + ("cusolverDnDsytrf", "hipsolverDnDsytrf"), + ("cusolverDnSsytrf", "hipsolverDnSsytrf"), + ("cusolverDnZsytrf", "hipsolverDnZsytrf"), + ("cusolverDnCsytrf", "hipsolverDnCsytrf"), + ("cusolverDnSgesvdaStridedBatched_bufferSize", "hipsolverDnSgesvdaStridedBatched_bufferSize"), + ("cusolverDnDgesvdaStridedBatched_bufferSize", "hipsolverDnDgesvdaStridedBatched_bufferSize"), + ("cusolverDnCgesvdaStridedBatched_bufferSize", "hipsolverDnCgesvdaStridedBatched_bufferSize"), + ("cusolverDnZgesvdaStridedBatched_bufferSize", "hipsolverDnZgesvdaStridedBatched_bufferSize"), + ("cusolverDnSgesvdaStridedBatched", "hipsolverDnSgesvdaStridedBatched"), + ("cusolverDnDgesvdaStridedBatched", "hipsolverDnDgesvdaStridedBatched"), + ("cusolverDnCgesvdaStridedBatched", "hipsolverDnCgesvdaStridedBatched"), + ("cusolverDnZgesvdaStridedBatched", "hipsolverDnZgesvdaStridedBatched"), + ("cusolverDnXgesvdjSetTolerance", "hipsolverDnXgesvdjSetTolerance"), + ("cusolverDnXgesvdjSetMaxSweeps", "hipsolverDnXgesvdjSetMaxSweeps"), + ("cusolverDnSgebrd_bufferSize", "hipsolverDnSgebrd_bufferSize"), + ("cusolverDnDgebrd_bufferSize", "hipsolverDnDgebrd_bufferSize"), + ("cusolverDnCgebrd_bufferSize", "hipsolverDnCgebrd_bufferSize"), + ("cusolverDnZgebrd_bufferSize", "hipsolverDnZgebrd_bufferSize"), + ("cusolverDnSgebrd", "hipsolverDnSgebrd"), + ("cusolverDnDgebrd", "hipsolverDnDgebrd"), + ("cusolverDnCgebrd", "hipsolverDnCgebrd"), + ("cusolverDnZgebrd", "hipsolverDnZgebrd"), + ("cusolverDnXgesvdjGetSweeps", "hipsolverDnXgesvdjGetSweeps"), + ("cusolverDnXsyevjSetTolerance", "hipsolverDnXsyevjSetTolerance"), + ("cusolverDnXsyevjSetMaxSweeps", "hipsolverDnXsyevjSetMaxSweeps"), + ("cusolverDnXsyevjGetResidual", "hipsolverDnXsyevjGetResidual"), + ("cusolverDnXgesvdjGetResidual", "hipsolverDnXgesvdjGetResidual"), + ("cusolverDnXsyevjGetSweeps", "hipsolverDnXsyevjGetSweeps"), +]) + +PYTORCH_SPECIFIC_MAPPINGS = collections.OrderedDict([ + ("USE_CUDA", "USE_ROCM"), + ("CUDA_VERSION", "TORCH_HIP_VERSION"), + ("gloo/cuda.h", "gloo/hip.h"), + ("gloo/cuda_allreduce_halving_doubling.h", "gloo/hip_allreduce_halving_doubling.h"), + ("gloo/cuda_allreduce_halving_doubling_pipelined.h", "gloo/hip_allreduce_halving_doubling_pipelined.h"), + ("gloo/cuda_allreduce_ring.h", "gloo/hip_allreduce_ring.h"), + ("gloo/cuda_allreduce_ring_chunked.h", "gloo/hip_allreduce_ring_chunked.h"), + ("gloo/cuda_broadcast_one_to_all.h", "gloo/hip_broadcast_one_to_all.h"), + ("gloo::CudaAllreduceHalvingDoublingPipelined", "gloo::HipAllreduceHalvingDoublingPipelined"), + ("gloo::CudaAllreduceRingChunked", "gloo::HipAllreduceRingChunked"), + ("gloo::CudaBroadcastOneToAll", "gloo::HipBroadcastOneToAll"), + ("gloo::CudaHostWorkspace", "gloo::HipHostWorkspace"), + ("gloo::CudaDeviceWorkspace", "gloo::HipDeviceWorkspace"), + ("CUDNN_RNN_RELU", "miopenRNNRELU"), + ("CUDNN_RNN_TANH", "miopenRNNTANH"), + ("CUDNN_LSTM", "miopenLSTM"), + ("CUDNN_GRU", "miopenGRU"), + ("cudnnRNNMode_t", "miopenRNNMode_t"), + ("magma_queue_create_from_cuda", "magma_queue_create_from_hip"), + # TODO: Remove these. They were necessary for Meta-internal builds. + ("cudnnHandle_t", "miopenHandle_t"), + ("cudnnCreate", "miopenCreate"), + ("cudnnDestroy", "miopenDestroy"), + ("cudnnSetStream", "miopenSetStream"), + ("cudnnTensorDescriptor_t ", "miopenTensorDescriptor_t "), + ("CUDNN_ENFORCE", "MIOPEN_ENFORCE"), + ("CUDNN_CHECK", "MIOPEN_CHECK"), + # NVSHMEM → rocSHMEM mappings (only symbols used in hipified files: + # NVSHMEMSymmetricMemory.cpp and nvshmem_team_manager.hpp). + ("NVSHMEM_TEAM_INVALID", "rocshmem::ROCSHMEM_TEAM_INVALID"), + ("NVSHMEM_TEAM_WORLD", "rocshmem::ROCSHMEM_TEAM_WORLD"), + ("NVSHMEMX_INIT_WITH_UNIQUEID", "rocshmem::ROCSHMEM_INIT_WITH_UNIQUEID"), + + ("nvshmem_malloc", "rocshmem::rocshmem_malloc"), + ("nvshmem_free", "rocshmem::rocshmem_free"), + ("nvshmem_ptr", "rocshmem::rocshmem_ptr"), + ("nvshmem_team_t", "rocshmem::rocshmem_team_t"), + ("nvshmem_team_split_strided", "rocshmem::rocshmem_team_split_strided"), + + ("nvshmemx_uniqueid_t", "rocshmem::rocshmem_uniqueid_t"), + ("nvshmemx_get_uniqueid", "rocshmem::rocshmem_get_uniqueid"), + ("nvshmemx_init_attr", "rocshmem::rocshmem_init_attr"), + ("nvshmemx_init_attr_t", "rocshmem::rocshmem_init_attr_t"), + ("nvshmemx_set_attr_uniqueid_args", "rocshmem::rocshmem_set_attr_uniqueid_args"), +]) + +C10_MAPPINGS = collections.OrderedDict([ + ("CUDA_VERSION", "TORCH_HIP_VERSION"), + ("CUDA_LAUNCH_BLOCKING=1", "AMD_SERIALIZE_KERNEL=3"), + ("CUDA_LAUNCH_BLOCKING", "AMD_SERIALIZE_KERNEL"), + ("c10/cuda/CUDAAlgorithm.h", "c10/hip/HIPAlgorithm.h"), + ("c10/cuda/CUDAAllocatorConfig.h", "c10/hip/HIPAllocatorConfig.h"), + ("c10/cuda/CUDACachingAllocator.h", "c10/hip/HIPCachingAllocator.h"), + ("c10/cuda/CUDADeviceAssertion.h", "c10/hip/HIPDeviceAssertion.h"), + ("c10/cuda/CUDADeviceAssertionHost.h", "c10/hip/HIPDeviceAssertionHost.h"), + ("c10/cuda/CUDAException.h", "c10/hip/HIPException.h"), + ("c10/cuda/CUDAFunctions.h", "c10/hip/HIPFunctions.h"), + ("c10/cuda/CUDAGraphsC10Utils.h", "c10/hip/HIPGraphsC10Utils.h"), + ("c10/cuda/CUDAGuard.h", "c10/hip/HIPGuard.h"), + ("c10/cuda/CUDAMacros.h", "c10/hip/HIPMacros.h"), + ("c10/cuda/CUDAMathCompat.h", "c10/hip/HIPMathCompat.h"), + ("c10/cuda/CUDAMiscFunctions.h", "c10/hip/HIPMiscFunctions.h"), + ("c10/cuda/CUDAStream.h", "c10/hip/HIPStream.h"), + ("c10/cuda/PeerToPeerAccess.h", "c10/hip/PeerToPeerAccess.h"), + ("c10/cuda/CUDAEvent.h", "c10/hip/HIPEvent.h"), + ("c10/cuda/impl/CUDAGuardImpl.h", "c10/hip/impl/HIPGuardImpl.h"), + ("c10/cuda/impl/CUDATest.h", "c10/hip/impl/HIPTest.h"), + ("CUDATest.hpp", "HIPTest.hpp"), + ("c10/cuda/impl/cuda_cmake_macros.h", "c10/hip/impl/hip_cmake_macros.h"), + # TODO: Remove these. They were necessary for Meta-internal builds. + ("c10::hip::c10_hip_check_implementation", "c10::cuda::c10_cuda_check_implementation"), +]) + +# TODO: Remove CAFFE2_SPECIFIC_MAPPINGS. They were necessary for Meta-internal builds. +# CAFFE2 mappings for simple filename patterns (no path separators) +# These work with the word-boundary regex approach +# NOTE: Removed broad "context_gpu" mapping to prevent double transformation. +# Use CAFFE2_PATH_MAPPINGS for specific path transformations instead. +CAFFE2_SPECIFIC_MAPPINGS = collections.OrderedDict([ + ("cuda_nccl_gpu", "hip/hip_nccl_gpu"), + ("mixed_utils", "hip/mixed_utils"), + ("operator_fallback_gpu", "hip/operator_fallback_gpu"), + ("spatial_batch_norm_op_impl", "hip/spatial_batch_norm_op_impl"), + ("recurrent_network_executor_gpu", "hip/recurrent_network_executor_gpu"), + ("generate_proposals_op_util_nms_gpu", "hip/generate_proposals_op_util_nms_gpu"), + ("max_pool_with_index_gpu", "hip/max_pool_with_index_gpu"), + ("THCCachingAllocator_gpu", "hip/THCCachingAllocator_gpu"), + ("top_k_heap_selection", "hip/top_k_heap_selection"), + ("top_k_radix_selection", "hip/top_k_radix_selection"), + ("GpuAtomics", "hip/GpuAtomics"), + ("GpuDefs", "hip/GpuDefs"), + ("GpuScanUtils", "hip/GpuScanUtils"), + ("GpuBitonicSort", "hip/GpuBitonicSort"), + # ("gather_op", "hip/gather_op"), # gather_op.h is device-agnostic, no path transform needed + ("HIP_CHECK", "CUDA_CHECK"), + # ("HIPContext", "CUDAContext"), + ("CUBLAS_ENFORCE", "HIPBLAS_ENFORCE"), + ("CaffeHipGetDevice", "CaffeCudaGetDevice"), +]) + +# TODO: Remove CAFFE2_PATH_MAPPINGS. They were necessary for Meta-internal builds. +# CAFFE2 path mappings (contain slashes - need special handling in hipify_python.py) +CAFFE2_PATH_MAPPINGS = collections.OrderedDict([ + ("math/reduce.cuh", "math/hip/reduce.cuh"), + ("operators/gather_op.cuh", "operators/hip/gather_op.cuh"), + ("sgd/adagrad_fused_op_gpu.cuh", "sgd/hip/adagrad_fused_op_gpu.cuh"), + ("operators/segment_reduction_op_gpu.cuh", "operators/hip/segment_reduction_op_gpu.cuh"), + ("caffe2/core/common_cudnn.h", "caffe2/core/hip/common_miopen.h"), + # Add specific mappings for common_gpu and context_gpu to prevent double transformation + ("caffe2/core/common_gpu.h", "caffe2/core/hip/common_gpu.h"), + ("caffe2/core/context_gpu.h", "caffe2/core/hip/context_gpu.h"), + ("hip/hip/", "hip/"), +]) + +CUDA_TO_HIP_MAPPINGS = [ + CUDA_IDENTIFIER_MAP, + CUDA_TYPE_NAME_MAP, + CUDA_INCLUDE_MAP, + CUDA_SPECIAL_MAP, + PYTORCH_SPECIFIC_MAPPINGS, + C10_MAPPINGS, + # TODO: Remove CAFFE2_SPECIFIC_MAPPINGS and CAFFE2_PATH_MAPPINGS. See above. + CAFFE2_SPECIFIC_MAPPINGS, + CAFFE2_PATH_MAPPINGS +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/hipify_python.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/hipify_python.py new file mode 100644 index 0000000000000000000000000000000000000000..3919d61bbaa0cd15740d1efbb33c1e2e7434fc56 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/hipify_python.py @@ -0,0 +1,1175 @@ +#!/usr/bin/env python3 +# mypy: allow-untyped-defs +""" The Python Hipify script. +## +# Copyright (c) 2015-2016 Advanced Micro Devices, Inc. All rights reserved. +# 2017-2018 Advanced Micro Devices, Inc. and +# Facebook Inc. All rights reserved. +# +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: +# +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. +# +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +""" +import argparse +import fnmatch +import re +import shutil +import sys +import os +import warnings + +from .cuda_to_hip_mappings import CUDA_TO_HIP_MAPPINGS +from .cuda_to_hip_mappings import MATH_TRANSPILATIONS +from .cuda_to_hip_mappings import CAFFE2_PATH_MAPPINGS + +from collections.abc import Iterator +from collections.abc import Mapping, Iterable +from enum import Enum +import functools +import hashlib + +def _deprecated(name): + warnings.warn(f"hipify version 2.0.0 no longer uses function {name}", FutureWarning, stacklevel=2) + +class CurrentState(Enum): + INITIALIZED = 1 + DONE = 2 + +class HipifyResult: + def __init__(self, current_state, hipified_path) -> None: + self.current_state = current_state + self.hipified_path = hipified_path + self.status = "" + + def __str__(self) -> str: + return (f"HipifyResult:: current_state: {self.current_state}, hipified_path : {self.hipified_path}, status: {self.status}") + +HipifyFinalResult = dict[str, HipifyResult] +HIPIFY_C_BREADCRUMB = "// !!! This is a file automatically generated by hipify!!!\n" +HIPIFY_FINAL_RESULT: HipifyFinalResult = {} + +# Hardcode the PyTorch template map +"""This dictionary provides the mapping from PyTorch kernel template types +to their actual types.""" +PYTORCH_TEMPLATE_MAP = {"Dtype": "scalar_t", "T": "scalar_t"} + +__all__ = ['InputError', 'openf', 'bcolors', 'GeneratedFileCleaner', 'match_extensions', 'matched_files_iter', + 'preprocess_file_and_save_result', 'compute_stats', 'add_dim3', 'processKernelLaunches', 'find_closure_group', + 'find_bracket_group', 'find_parentheses_group', 'replace_math_functions', 'hip_header_magic', 'replace_extern_shared', + 'get_hip_file_path', 'is_out_of_place', 'is_pytorch_file', 'is_cusparse_file', 'is_special_file', 'is_caffe2_gpu_file', + 'Trie', 'preprocessor', 'file_specific_replacement', 'file_add_header', + 'fix_static_global_kernels', 'extract_arguments', 'str2bool', 'CurrentState', 'HipifyResult', 'hipify'] + + +class InputError(Exception): + # Exception raised for errors in the input. + + def __init__(self, message) -> None: + super().__init__(message) + self.message = message + + def __str__(self) -> str: + return f"Input error: {self.message}" + + +def openf(filename, mode): + return open(filename, mode, errors='ignore') + + +# Color coding for printing +class bcolors: + HEADER = '\033[95m' + OKBLUE = '\033[94m' + OKGREEN = '\033[92m' + WARNING = '\033[93m' + FAIL = '\033[91m' + ENDC = '\033[0m' + BOLD = '\033[1m' + UNDERLINE = '\033[4m' + + +# To the programmer, the output of hipify most likely are intermediates. +# This class allows users of hipify to ask for a cleanup by running the +# hipify and compilation in a with instantiating this context manager class +# with keep_intermediates=False. +# The main usecase is the cpp_extensions, specifically the load method. +# It is a good idea to keep intermediates (in case of errors or to +# not recompile unchanged files), but in cases where you don't want to +# keep them (e.g. in the CI), this can be used to remove files. +class GeneratedFileCleaner: + """Context Manager to clean up generated files""" + def __init__(self, keep_intermediates=False) -> None: + self.keep_intermediates = keep_intermediates + self.files_to_clean = set() + self.dirs_to_clean = [] + + def __enter__(self): + return self + + def open(self, fn, *args, **kwargs): + if not os.path.exists(fn): + self.files_to_clean.add(os.path.abspath(fn)) + + return open(fn, *args, **kwargs) + + def makedirs(self, dn, exist_ok=False) -> None: + parent, n = os.path.split(dn) + if not n: + parent, n = os.path.split(parent) + if parent and n and not os.path.exists(parent): + self.makedirs(parent, exist_ok=True) + if not os.path.isdir(dn) or not exist_ok: + os.mkdir(dn) + self.dirs_to_clean.append(os.path.abspath(dn)) + + def __exit__(self, type, value, traceback): + if not self.keep_intermediates: + for f in self.files_to_clean: + os.unlink(f) + for d in self.dirs_to_clean[::-1]: + os.rmdir(d) + +# Follow UNIX convention for paths to use '/' instead of '\\' on Windows +def _to_unix_path(path: str) -> str: + return path.replace(os.sep, '/') + +def match_extensions(filename: str, extensions: Iterable) -> bool: + """Helper method to see if filename ends with certain extension""" + return any(filename.endswith(e) for e in extensions) + + +def _fnmatch(filepath, patterns): + return any(fnmatch.fnmatch(filepath, pattern) for pattern in patterns) + + +def matched_files_iter( + root_path: str, + includes: Iterable = (), + ignores: Iterable = (), + extensions: Iterable = (), + out_of_place_only: bool = False, + is_pytorch_extension: bool = False) -> Iterator[str]: + + exact_matches = set(includes) + + # This is a very rough heuristic; really, we want to avoid scanning + # any file which is not checked into source control, but this script + # needs to work even if you're in a Git or Hg checkout, so easier to + # just block the biggest time sinks that won't matter in the + # end. + for (abs_dirpath, dirs, filenames) in os.walk(root_path, topdown=True): + rel_dirpath = os.path.relpath(abs_dirpath, root_path) + if rel_dirpath == '.': + # Blah blah blah O(n) blah blah + if ".git" in dirs: + dirs.remove(".git") + if "build" in dirs: + dirs.remove("build") + if "third_party" in dirs: + dirs.remove("third_party") + dirs.append("third_party/nvfuser") + for filename in filenames: + filepath = _to_unix_path(os.path.join(abs_dirpath, filename)) + # We respect extensions, UNLESS you wrote the entire + # filename verbatim, in which case we always accept it + if ( + _fnmatch(filepath, includes) + and (not _fnmatch(filepath, ignores)) + and (match_extensions(filepath, extensions) or filepath in exact_matches) + ): + yield filepath + + +def preprocess_file_and_save_result( + output_directory: str, + filepath: str, + all_files: Iterable, + header_include_dirs: Iterable, + stats: dict[str, list], + hip_clang_launch: bool, + is_pytorch_extension: bool, + clean_ctx: GeneratedFileCleaner, + show_progress: bool) -> None: + fin_path = os.path.abspath(os.path.join(output_directory, filepath)) + hipify_result = HipifyResult(current_state=CurrentState.INITIALIZED, hipified_path=fin_path) + HIPIFY_FINAL_RESULT[fin_path] = hipify_result + result = preprocessor(output_directory, filepath, all_files, header_include_dirs, stats, + hip_clang_launch, is_pytorch_extension, clean_ctx, show_progress) + + # Show what happened + if show_progress and "ignored" not in result.status: + print( + fin_path, "->", + result.hipified_path, result.status, flush=True) + + HIPIFY_FINAL_RESULT[fin_path] = result + + +def compute_stats(stats) -> None: + unsupported_calls = {cuda_call for (cuda_call, _filepath) in stats["unsupported_calls"]} + + # Print the number of unsupported calls + print(f"Total number of unsupported CUDA function calls: {len(unsupported_calls):d}") + + # Print the list of unsupported calls + print(", ".join(unsupported_calls)) + + # Print the number of kernel launches + print(f"\nTotal number of replaced kernel launches: {len(stats['kernel_launches']):d}") + + +def add_dim3(kernel_string, cuda_kernel): + '''adds dim3() to the second and third arguments in the kernel launch''' + count = 0 + closure = 0 + kernel_string = kernel_string.replace("<<<", "").replace(">>>", "") + arg_locs: list[dict[str, int]] = [{} for _ in range(2)] + arg_locs[count]['start'] = 0 + for ind, c in enumerate(kernel_string): + if count > 1: + break + if c == "(": + closure += 1 + elif c == ")": + closure -= 1 + if (c == "," or ind == len(kernel_string) - 1) and closure == 0: + arg_locs[count]['end'] = ind + (c != ",") + count += 1 + if count < 2: + arg_locs[count]['start'] = ind + 1 + + first_arg_raw = kernel_string[arg_locs[0]['start']:arg_locs[0]['end'] + 1] + second_arg_raw = kernel_string[arg_locs[1]['start']:arg_locs[1]['end']] + + first_arg_clean = kernel_string[arg_locs[0]['start']:arg_locs[0]['end']].replace("\n", "").strip(" ") + second_arg_clean = kernel_string[arg_locs[1]['start']:arg_locs[1]['end']].replace("\n", "").strip(" ") + + first_arg_dim3 = f"dim3({first_arg_clean})" + second_arg_dim3 = f"dim3({second_arg_clean})" + + first_arg_raw_dim3 = first_arg_raw.replace(first_arg_clean, first_arg_dim3) + second_arg_raw_dim3 = second_arg_raw.replace(second_arg_clean, second_arg_dim3) + cuda_kernel = cuda_kernel.replace(first_arg_raw + second_arg_raw, first_arg_raw_dim3 + second_arg_raw_dim3) + return cuda_kernel + + +RE_KERNEL_LAUNCH = re.compile(r'([ ]+)(detail?)::[ ]+\\\n[ ]+') + + +def processKernelLaunches(string, stats): + """ Replace the CUDA style Kernel launches with the HIP style kernel launches.""" + # Concat the namespace with the kernel names. (Find cleaner way of doing this later). + string = RE_KERNEL_LAUNCH.sub(lambda inp: f"{inp.group(1)}{inp.group(2)}::", string) + + def grab_method_and_template(in_kernel): + # The positions for relevant kernel components. + pos = { + "kernel_launch": {"start": in_kernel["start"], "end": in_kernel["end"]}, + "kernel_name": {"start": -1, "end": -1}, + "template": {"start": -1, "end": -1} + } + + # Count for balancing template + count = {"<>": 0} + + # Status for whether we are parsing a certain item. + START = 0 + AT_TEMPLATE = 1 + AFTER_TEMPLATE = 2 + AT_KERNEL_NAME = 3 + + status = START + + # Parse the string character by character + for i in range(pos["kernel_launch"]["start"] - 1, -1, -1): + char = string[i] + + # Handle Templating Arguments + if status in (START, AT_TEMPLATE): + if char == ">": + if status == START: + status = AT_TEMPLATE + pos["template"]["end"] = i + count["<>"] += 1 + + if char == "<": + count["<>"] -= 1 + if count["<>"] == 0 and (status == AT_TEMPLATE): + pos["template"]["start"] = i + status = AFTER_TEMPLATE + + # Handle Kernel Name + if status != AT_TEMPLATE: + if string[i].isalnum() or string[i] in {'(', ')', '_', ':', '#'}: + if status != AT_KERNEL_NAME: + status = AT_KERNEL_NAME + pos["kernel_name"]["end"] = i + + # Case: Kernel name starts the string. + if i == 0: + pos["kernel_name"]["start"] = 0 + + # Finished + return [(pos["kernel_name"]), (pos["template"]), (pos["kernel_launch"])] + + else: + # Potential ending point if we're already traversing a kernel's name. + if status == AT_KERNEL_NAME: + pos["kernel_name"]["start"] = i + + # Finished + return [(pos["kernel_name"]), (pos["template"]), (pos["kernel_launch"])] + + def find_kernel_bounds(string): + """Finds the starting and ending points for all kernel launches in the string.""" + kernel_end = 0 + kernel_positions = [] + + # Continue until we cannot find any more kernels anymore. + while string.find("<<<", kernel_end) != -1: + # Get kernel starting position (starting from the previous ending point) + kernel_start = string.find("<<<", kernel_end) + + # Get kernel ending position (adjust end point past the >>>) + kernel_end = string.find(">>>", kernel_start) + 3 + if kernel_end <= 0: + raise InputError("no kernel end found") + + # Add to list of traversed kernels + kernel_positions.append({"start": kernel_start, "end": kernel_end, + "group": string[kernel_start: kernel_end]}) + + return kernel_positions + + # Replace comments and string literals from the code so that find_kernel_bounds does not + # wrongly capture kernels in comments and string literals. + # This function replaces them with "x" to keep positions. + def mask_comments(string): + in_comment = '' + prev_c = '' + new_string = '' + for c in string: + if in_comment == '': + # Outside comments + if c == '/' and prev_c == '/': + in_comment = '//' + elif c == '*' and prev_c == '/': + in_comment = '/*' + elif c == '"' and prev_c != '\\' and prev_c != "'": + in_comment = '"' + elif in_comment == '//': + # In // xxx + if c == '\r' or c == '\n': + in_comment = '' + elif in_comment == '/*': + # In /* xxx */ + if c == '/' and prev_c == '*': + in_comment = '' + elif in_comment == '"': + # In "" + if c == '"' and prev_c != '\\': + in_comment = '' + prev_c = c + if in_comment == '': + new_string += c + else: + new_string += 'x' + return new_string + + # Grab positional ranges of all kernel launches + get_kernel_positions = list(find_kernel_bounds(mask_comments(string))) + output_string = string + + # Replace each CUDA kernel with a HIP kernel. + for kernel in get_kernel_positions: + # Get kernel components + params = grab_method_and_template(kernel) + + # Find parenthesis after kernel launch + parenthesis = string.find("(", kernel["end"]) + + # Extract cuda kernel + cuda_kernel = string[params[0]["start"]:parenthesis + 1] + kernel_string = string[kernel['start']:kernel['end']] + end_param_index = 0 if params[1]['end'] == -1 else 1 + kernel_name_with_template = string[params[0]['start']:params[end_param_index]['end'] + 1] + cuda_kernel_dim3 = add_dim3(kernel_string, cuda_kernel) + # Keep number of kernel launch params consistent (grid dims, group dims, stream, dynamic shared size) + num_klp = len(extract_arguments(0, kernel["group"].replace("<<<", "(").replace(">>>", ")"))) + + hip_kernel = "hipLaunchKernelGGL(" + cuda_kernel_dim3[0:-1].replace( + ">>>", ", 0" * (4 - num_klp) + ">>>").replace("<<<", ", ").replace( + ">>>", ", ").replace(kernel_name_with_template, "(" + kernel_name_with_template + ")") + + # Replace cuda kernel with hip kernel + output_string = output_string.replace(cuda_kernel, hip_kernel) + + # Update the statistics + stats["kernel_launches"].append(hip_kernel) + + return output_string + + +def find_closure_group(input_string, start, group): + """Generalization for finding a balancing closure group + + if group = ["(", ")"], then finds the first balanced parentheses. + if group = ["{", "}"], then finds the first balanced bracket. + + Given an input string, a starting position in the input string, and the group type, + find_closure_group returns the positions of group[0] and group[1] as a tuple. + + Example: + >>> find_closure_group("(hi)", 0, ["(", ")"]) + (0, 3) + """ + + inside_parenthesis = False + parens = 0 + pos = start + p_start, p_end = -1, -1 + + while pos < len(input_string): + if input_string[pos] == group[0]: + if inside_parenthesis is False: + inside_parenthesis = True + parens = 1 + p_start = pos + else: + parens += 1 + elif input_string[pos] == group[1] and inside_parenthesis: + parens -= 1 + + if parens == 0: + p_end = pos + return p_start, p_end + + pos += 1 + return None, None + + +def find_bracket_group(input_string, start): + """Finds the first balanced parentheses.""" + return find_closure_group(input_string, start, group=["{", "}"]) + + +def find_parentheses_group(input_string, start): + """Finds the first balanced bracket.""" + return find_closure_group(input_string, start, group=["(", ")"]) + + +RE_ASSERT = re.compile(r"\bassert[ ]*\(") + + +def replace_math_functions(input_string): + """FIXME: Temporarily replace std:: invocations of math functions + with non-std:: versions to prevent linker errors NOTE: This + can lead to correctness issues when running tests, since the + correct version of the math function (exp/expf) might not get + called. Plan is to remove this function once HIP supports + std:: math function calls inside device code + + """ + output_string = input_string + for func in MATH_TRANSPILATIONS: + output_string = output_string.replace(fr'{func}(', f'{MATH_TRANSPILATIONS[func]}(') + + return output_string + + +RE_SYNCTHREADS = re.compile(r":?:?\b(__syncthreads)\b(\w*\()") + + +def hip_header_magic(input_string): + """If the file makes kernel builtin calls and does not include the cuda_runtime.h header, + then automatically add an #include to match the "magic" includes provided by NVCC. + TODO: + Update logic to ignore cases where the cuda_runtime.h is included by another file. + """ + + # Copy the input. + output_string = input_string + + # Check if one of the following headers is already included. + headers = ["hip/hip_runtime.h", "hip/hip_runtime_api.h"] + if any(re.search(fr'#include ("{ext}"|<{ext}>)', output_string) for ext in headers): + return output_string + + # Rough logic to detect if we're inside device code + hasDeviceLogic: int + hasDeviceLogic = "hipLaunchKernelGGL" in output_string + hasDeviceLogic += "__global__" in output_string + hasDeviceLogic += "__shared__" in output_string + hasDeviceLogic += RE_SYNCTHREADS.search(output_string) is not None + + # If device logic found, provide the necessary header. + if hasDeviceLogic: + output_string = '#include "hip/hip_runtime.h"\n' + input_string + + return output_string + + +RE_EXTERN_SHARED = re.compile(r"extern\s+([\w\(\)]+)?\s*__shared__\s+([\w:<>\s]+)\s+(\w+)\s*\[\s*\]\s*;") + + +def replace_extern_shared(input_string): + """ + Match 'extern __shared__ type foo[];' syntax and use HIP_DYNAMIC_SHARED() MACRO instead. + See: https://github.com/ROCm/hip/blob/master/docs/markdown/hip_kernel_language.md#__shared__ + Examples: + "extern __shared__ char smemChar[];" + => "HIP_DYNAMIC_SHARED( char, smemChar)" + "extern __shared__ unsigned char smem[];" + => "HIP_DYNAMIC_SHARED( unsigned char, my_smem)" + """ + output_string = input_string + output_string = RE_EXTERN_SHARED.sub( + lambda inp: f"HIP_DYNAMIC_SHARED({inp.group(1) or ''} {inp.group(2)}, {inp.group(3)})", output_string) + + return output_string + + +def get_hip_file_path(rel_filepath, is_pytorch_extension=False): + """ + Returns the new name of the hipified file + """ + # At the moment, some PyTorch source files are HIPified in place. The predicate + # is_out_of_place tells us if this is the case or not. + if os.path.isabs(rel_filepath): + raise AssertionError("rel_filepath must be a relative path") + if not is_pytorch_extension and not is_out_of_place(rel_filepath): + return rel_filepath + + dirpath, filename = os.path.split(rel_filepath) + root, ext = os.path.splitext(filename) + + # Here's the plan: + # + # In general, we need to disambiguate the HIPified filename so that + # it gets a different name from the original filename, so + # that we don't overwrite the original file + # + # There's a lot of different naming conventions across PyTorch, + # but the general recipe is to convert occurrences + # of cuda/gpu to hip, and add hip if there are no occurrences + # of cuda/gpu anywhere. + # + # Concretely, we do the following: + # + # - If there is a directory component named "cuda", replace + # it with "hip", AND + # + # - If the file name contains "CUDA", replace it with "HIP", AND + # + # - ALWAYS replace '.cu' with '.hip', because those files + # contain CUDA kernels that needs to be hipified and processed with + # hip compiler + # + # - If we are not hipifying a PyTorch extension, and the parent + # directory name did not change as a result of the above + # transformations, insert "hip" in the file path + # as the direct parent folder of the file + # + # - If we are hipifying a PyTorch extension, and the parent directory + # name as well as the filename (incl. extension) did not change as + # a result of the above transformations, insert "_hip" in the filename + # + # This isn't set in stone; we might adjust this to support other + # naming conventions. + + if ext == '.cu': + ext = '.hip' + + orig_filename = filename + orig_dirpath = dirpath + + dirpath = dirpath.replace('cuda', 'hip') + dirpath = dirpath.replace('CUDA', 'HIP') + dirpath = dirpath.replace('THC', 'THH') + + root = root.replace('cuda', 'hip') + root = root.replace('CUDA', 'HIP') + # Special case to handle caffe2/core/THCCachingAllocator + if dirpath != "caffe2/core": + root = root.replace('THC', 'THH') + + if not is_pytorch_extension and dirpath == orig_dirpath: + dirpath = os.path.join(dirpath, 'hip') + + if is_pytorch_extension and dirpath == orig_dirpath and (root + ext) == orig_filename: + root = root + "_hip" + + return os.path.join(dirpath, root + ext) + + +def is_out_of_place(rel_filepath) -> bool: + if os.path.isabs(rel_filepath): + raise AssertionError("rel_filepath must be a relative path") + if rel_filepath.startswith("torch/"): + return False + if rel_filepath.startswith("third_party/nvfuser/"): + return False + if rel_filepath.startswith("tools/autograd/templates/"): + return False + return True + + +# Keep this synchronized with includes/ignores in build_amd.py +def is_pytorch_file(rel_filepath) -> bool: + _deprecated("is_pytorch_file") + if os.path.isabs(rel_filepath): + raise AssertionError("rel_filepath must be a relative path") + if rel_filepath.startswith("aten/"): + if rel_filepath.startswith("aten/src/ATen/core/"): + return False + return True + if rel_filepath.startswith("torch/"): + return True + if rel_filepath.startswith("third_party/nvfuser/"): + return True + if rel_filepath.startswith("third_party/fbgemm/"): + return True + if rel_filepath.startswith("third_party/mslk/"): + return True + if rel_filepath.startswith("tools/autograd/templates/"): + return True + if rel_filepath.startswith("test/cpp/c10d/"): + return True + return False + + +def is_cusparse_file(rel_filepath): + _deprecated("is_cusparse_file") + if is_pytorch_file(rel_filepath): + return "sparse" in rel_filepath.lower() + return False + + +def is_special_file(rel_filepath) -> bool: + _deprecated("is_special_file") + if is_pytorch_file(rel_filepath): + if "sparse" in rel_filepath.lower(): + return True + elif "linalg" in rel_filepath.lower(): + if "batchlinearalgebralibblas" in rel_filepath.lower(): + return False # don't use "special" mappings for this specific linalg cublas file + return True + return False + + +def is_caffe2_gpu_file(rel_filepath): + _deprecated("is_caffe2_gpu_file") + if os.path.isabs(rel_filepath): + raise AssertionError("rel_filepath must be a relative path") + if rel_filepath.startswith("c10/cuda"): + return True + filename = os.path.basename(rel_filepath) + _, ext = os.path.splitext(filename) + + return ('gpu' in filename or ext in ['.cu', '.cuh']) and ('cudnn' not in filename) + + +class TrieNode: + """A Trie node whose children are represented as a directory of char: TrieNode. + A special char '' represents end of word + """ + + def __init__(self) -> None: + self.children = {} + + +class Trie: + """Creates a Trie out of a list of words. The trie can be exported to a Regex pattern. + The corresponding Regex should match much faster than a simple Regex union.""" + + def __init__(self) -> None: + """Initialize the trie with an empty root node.""" + self.root = TrieNode() + self._hash = hashlib.md5(usedforsecurity=False) + self._digest = self._hash.digest() + + def add(self, word) -> None: + """Add a word to the Trie. """ + self._hash.update(word.encode()) + self._digest = self._hash.digest() + node = self.root + + for char in word: + node.children.setdefault(char, TrieNode()) + node = node.children[char] + node.children[''] = True # Mark the end of the word + + def dump(self): + """Return the root node of Trie. """ + return self.root + + def quote(self, char): + """ Escape a char for regex. """ + return re.escape(char) + + def search(self, word): + """Search whether word is present in the Trie. + Returns True if yes, else return False""" + node = self.root + for char in word: + if char in node.children: + node = node.children[char] + else: + return False + + # make sure to check the end-of-word marker present + return '' in node.children + + @functools.lru_cache # noqa: B019 + def _pattern(self, root, digest): + """Convert a Trie into a regular expression pattern + + Memoized on the hash digest of the trie, which is built incrementally + during add(). + """ + node = root + + if "" in node.children and len(node.children.keys()) == 1: + return None + + alt = [] # store alternative patterns + cc = [] # store char to char classes + q = 0 # for node representing the end of word + for char in sorted(node.children.keys()): + if isinstance(node.children[char], TrieNode): + try: + recurse = self._pattern(node.children[char], self._digest) + alt.append(self.quote(char) + recurse) + except Exception: + cc.append(self.quote(char)) + else: + q = 1 + cconly = not len(alt) > 0 + + if len(cc) > 0: + if len(cc) == 1: + alt.append(cc[0]) + else: + alt.append('[' + ''.join(cc) + ']') + + if len(alt) == 1: + result = alt[0] + else: + result = "(?:" + "|".join(alt) + ")" + + if q: + if cconly: + result += "?" + else: + result = f"(?:{result})?" + return result + + def pattern(self): + """Export the Trie to a regex pattern.""" + return self._pattern(self.root, self._digest) + + def export_to_regex(self): + """Export the Trie to a regex pattern.""" + return self._pattern(self.root, self._digest) + +PYTORCH_TRIE = Trie() +PYTORCH_MAP: dict[str, object] = {} + +for mapping in CUDA_TO_HIP_MAPPINGS: + if not isinstance(mapping, Mapping): + raise TypeError("Expected each mapping in CUDA_TO_HIP_MAPPINGS to be a Mapping") + for src, dst in mapping.items(): + PYTORCH_TRIE.add(src) + PYTORCH_MAP[src] = dst + +RE_PYTORCH_PREPROCESSOR = re.compile(fr'(?<=\W)({PYTORCH_TRIE.export_to_regex()})(?=\W)') + +RE_QUOTE_HEADER = re.compile(r'#include "([^"]+)"') +RE_ANGLE_HEADER = re.compile(r'#include <([^>]+)>') +RE_THC_GENERIC_FILE = re.compile(r'#define THC_GENERIC_FILE "([^"]+)"') +RE_CU_SUFFIX = re.compile(r'\.cu\b') # be careful not to pick up .cuh + +""" +Returns a HipifyResult object with the following details: + "hipified_path" : absolute path of hipified source file + "status" : "ok" if hipified file was written out + "skipped" if an identical hipified file already existed or hipified file couldn't be written out + "ignored" if the source file was a hipified file itself or not meant to be hipified + "current_state" : CurrentState.INITIALIZED if source file is first ready to be hipified + CurrentState.DONE if source file is done with hipification process +""" + + +def preprocessor( + output_directory: str, + filepath: str, + all_files: Iterable, + header_include_dirs: Iterable, + stats: dict[str, list], + hip_clang_launch: bool, + is_pytorch_extension: bool, + clean_ctx: GeneratedFileCleaner, + show_progress: bool) -> HipifyResult: + """ Executes the CUDA -> HIP conversion on the specified file. """ + fin_path = os.path.abspath(os.path.join(output_directory, filepath)) + filepath = _to_unix_path(filepath) + hipify_result = HIPIFY_FINAL_RESULT[fin_path] + if filepath not in all_files: + hipify_result.hipified_path = None + hipify_result.status = "[ignored, not to be hipified]" + hipify_result.current_state = CurrentState.DONE + return hipify_result + + rel_filepath = _to_unix_path(os.path.relpath(filepath, output_directory)) + + with open(fin_path, encoding='utf-8') as fin: + if fin.readline() == HIPIFY_C_BREADCRUMB: + hipify_result.hipified_path = None + hipify_result.status = "[ignored, input is hipified output]" + hipify_result.current_state = CurrentState.DONE + return hipify_result + fin.seek(0) + output_source = fin.read() + + orig_output_source = output_source + + # get_hip_file_path needs a relative path to work correctly + fout_path = os.path.abspath(os.path.join(output_directory, get_hip_file_path(rel_filepath, is_pytorch_extension))) + if not os.path.exists(os.path.dirname(fout_path)): + clean_ctx.makedirs(os.path.dirname(fout_path)) + + # unsupported_calls statistics reporting is broken atm + def pt_repl(m): + return PYTORCH_MAP[m.group(0)] + + output_source = RE_PYTORCH_PREPROCESSOR.sub(pt_repl, output_source) + + # TODO: Remove CAFFE2_PATH_MAPPINGS. They were necessary for Meta-internal builds. + # Apply CAFFE2 path mappings (simple string replacement for paths containing slashes) + # Need to be careful to avoid double-transformations when source file has #ifdef blocks + # with HIP-specific paths already in them (e.g., caffe2/core/hip/context_gpu.h) + for cuda_path, hip_path in CAFFE2_PATH_MAPPINGS.items(): + # Use regex to ensure we don't match paths that already have been hipified + # We need to avoid transforming "caffe2/core/hip/context_gpu.h" when looking for "caffe2/core/context_gpu.h" + # The key insight: if hip_path contains /hip/ and cuda_path doesn't, we need to be careful + if "/hip/" in hip_path and "/hip/" not in cuda_path: + # Only replace cuda_path if it's not preceded by "/hip/" + # Use negative lookbehind to prevent matching already-hipified paths + # The pattern checks that the cuda_path is not immediately preceded by "/hip/" + pattern = r'(?', False), output_source) + output_source = RE_THC_GENERIC_FILE.sub(mk_repl('#define THC_GENERIC_FILE "{0}"'), output_source) + + # CMakeLists.txt rewrites + if filepath.endswith('CMakeLists.txt'): + output_source = output_source.replace('CUDA', 'HIP') + output_source = output_source.replace('THC', 'THH') + output_source = RE_CU_SUFFIX.sub('.hip', output_source) + + # Perform Kernel Launch Replacements + if not hip_clang_launch: + output_source = processKernelLaunches(output_source, stats) + + # Replace std:: with non-std:: versions + if (filepath.endswith((".cu", ".cuh"))) and "PowKernel" not in filepath: + output_source = replace_math_functions(output_source) + + # Include header if device code is contained. + output_source = hip_header_magic(output_source) + + # Replace the extern __shared__ + # NOTE: No longer needed after transition from hcc to hipclang. + # output_source = replace_extern_shared(output_source) + + # Don't write out identical hipified files for extensions if dirpath has not changed + if ( + is_pytorch_extension + and orig_output_source == output_source + and os.path.dirname(fin_path) == os.path.dirname(fout_path) + ): + hipify_result.hipified_path = fin_path + hipify_result.status = "[skipped, no changes]" + hipify_result.current_state = CurrentState.DONE + return hipify_result + + # Add hipify breadcrumb for C-style files to avoid re-hipification + if fin_path != fout_path and match_extensions(fin_path, (".cu", ".cuh", ".c", ".cc", ".cpp", ".h", ".hpp")): + output_source = HIPIFY_C_BREADCRUMB + output_source + + do_write = True + if os.path.exists(fout_path): + with open(fout_path, encoding='utf-8') as fout_old: + do_write = fout_old.read() != output_source + if do_write: + try: + with clean_ctx.open(fout_path, 'w', encoding='utf-8') as fout: + fout.write(output_source) + hipify_result.hipified_path = fout_path + hipify_result.status = "[ok]" + hipify_result.current_state = CurrentState.DONE + return hipify_result + except OSError as e: + print(f'{bcolors.WARNING}Failed to save {fout_path} with "{e.strerror}", leaving {fin_path} unchanged.{bcolors.ENDC}', + file=sys.stderr) + hipify_result.hipified_path = fin_path + hipify_result.status = "[skipped, no permissions]" + hipify_result.current_state = CurrentState.DONE + return hipify_result + else: + hipify_result.hipified_path = fout_path + hipify_result.status = "[skipped, already hipified]" + hipify_result.current_state = CurrentState.DONE + return hipify_result + +def file_specific_replacement(filepath, search_string, replace_string, strict=False) -> None: + with openf(filepath, "r+") as f: + contents = f.read() + if strict: + contents = re.sub(fr'\b({re.escape(search_string)})\b', lambda x: replace_string, contents) + else: + contents = contents.replace(search_string, replace_string) + f.seek(0) + f.write(contents) + f.truncate() + + +def file_add_header(filepath, header) -> None: + with openf(filepath, "r+") as f: + contents = f.read() + if header[0] != "<" and header[-1] != ">": + header = f'"{header}"' + contents = (f'#include {header} \n') + contents + f.seek(0) + f.write(contents) + f.truncate() + + +def fix_static_global_kernels(in_txt): + """Static global kernels in HIP results in a compilation error.""" + in_txt = in_txt.replace(" __global__ static", "__global__") + return in_txt + + +RE_INCLUDE = re.compile(r"#include .*\n") + + +def extract_arguments(start, string): + """ + Return the list of arguments in the upcoming function parameter closure. + Example: + string (input): '(blocks, threads, 0, THCState_getCurrentStream(state))' + arguments (output): [{'start': 1, 'end': 7}, {'start': 8, 'end': 16}, \ + {'start': 17, 'end': 19}, {'start': 20, 'end': 53}] + """ + + arguments = [] + closures = { + "<": 0, + "(": 0 + } + current_position = start + argument_start_pos = current_position + 1 + + # Search for final parenthesis + while current_position < len(string): + if string[current_position] == "(": + closures["("] += 1 + elif string[current_position] == ")": + closures["("] -= 1 + elif string[current_position] == "<": + closures["<"] += 1 + elif string[current_position] == ">" and string[current_position - 1] != "-" and closures["<"] > 0: + closures["<"] -= 1 + + # Finished all arguments + if closures["("] == 0 and closures["<"] == 0: + # Add final argument + arguments.append({"start": argument_start_pos, "end": current_position}) + break + + # Finished current argument + if closures["("] == 1 and closures["<"] == 0 and string[current_position] == ",": + arguments.append({"start": argument_start_pos, "end": current_position}) + argument_start_pos = current_position + 1 + + current_position += 1 + + return arguments + + +def str2bool(v : str) -> bool: + """ArgumentParser doesn't support type=bool. Thus, this helper method will convert + from possible string types to True / False.""" + if v.lower() in ('yes', 'true', 't', 'y', '1'): + return True + elif v.lower() in ('no', 'false', 'f', 'n', '0'): + return False + else: + raise argparse.ArgumentTypeError('Boolean value expected.') + + +def hipify( + project_directory: str, + show_detailed: bool = False, + extensions: Iterable = (".cu", ".cuh", ".c", ".cc", ".cpp", ".h", ".in", ".hpp"), + header_extensions: Iterable = (".cuh", ".h", ".hpp"), + output_directory: str = "", + header_include_dirs: Iterable = (), + includes: Iterable = ('*',), + extra_files: Iterable = (), + out_of_place_only: bool = False, + ignores: Iterable = (), + show_progress: bool = True, + hip_clang_launch: bool = False, + is_pytorch_extension: bool = False, + hipify_extra_files_only: bool = False, + clean_ctx: GeneratedFileCleaner | None = None +) -> HipifyFinalResult: + if project_directory == "": + project_directory = os.getcwd() + + # Verify the project directory exists. + if not os.path.exists(project_directory): + print("The project folder specified does not exist.") + sys.exit(1) + + # If no output directory, provide a default one. + if not output_directory: + project_directory.rstrip("/") + output_directory = project_directory + "_amd" + + if project_directory != output_directory: + includes = [include.replace(project_directory, output_directory) for include in includes] + ignores = [ignore.replace(project_directory, output_directory) for ignore in ignores] + + # Copy from project directory to output directory if not done already. + if not os.path.exists(output_directory): + shutil.copytree(project_directory, output_directory) + + includes = list(map(_to_unix_path, includes)) + ignores = list(map(_to_unix_path, ignores)) + + all_files = list(matched_files_iter(output_directory, includes=includes, + ignores=ignores, extensions=extensions, + out_of_place_only=out_of_place_only, + is_pytorch_extension=is_pytorch_extension)) + all_files_set = set(all_files) + + for f in extra_files: + if not os.path.isabs(f): + f = os.path.join(output_directory, f) + if f not in all_files_set: + all_files.append(f) + + # List all files in header_include_paths to ensure they are hipified + from pathlib import Path + for header_include_dir in header_include_dirs: + if os.path.isabs(header_include_dir): + header_include_dir_path = Path(header_include_dir) + else: + header_include_dir_path = Path(os.path.join(output_directory, header_include_dir)) + all_files.extend( + str(path) for path in header_include_dir_path.rglob('*') if path.is_file() + and _fnmatch(str(path), includes) + and (not _fnmatch(str(path), ignores)) + and match_extensions(path.name, header_extensions) + ) + + if clean_ctx is None: + clean_ctx = GeneratedFileCleaner(keep_intermediates=True) + + # Preprocessing statistics. + stats: dict[str, list] = {"unsupported_calls": [], "kernel_launches": []} + + for filepath in (all_files if not hipify_extra_files_only else extra_files): + preprocess_file_and_save_result(output_directory, filepath, all_files, header_include_dirs, + stats, hip_clang_launch, is_pytorch_extension, clean_ctx, show_progress) + + print(bcolors.OKGREEN + "Successfully preprocessed all matching files." + bcolors.ENDC, file=sys.stderr) + + # Show detailed summary + if show_detailed: + compute_stats(stats) + + return HIPIFY_FINAL_RESULT diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/version.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/version.py new file mode 100644 index 0000000000000000000000000000000000000000..afced14728f75ce3cd253465517ccc1032a62309 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hipify/version.py @@ -0,0 +1 @@ +__version__ = '2.0.0' diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hooks.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hooks.py new file mode 100644 index 0000000000000000000000000000000000000000..1e3f6fb9ab09d755e9ed2154f91bb1184395d61c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/hooks.py @@ -0,0 +1,257 @@ +# mypy: allow-untyped-defs +import torch +from collections import OrderedDict +import weakref +import warnings +from typing import Any + +__all__ = ["RemovableHandle", "unserializable_hook", "warn_if_has_hooks", "BackwardHook"] + +class RemovableHandle: + r""" + A handle which provides the capability to remove a hook. + + Args: + hooks_dict (dict): A dictionary of hooks, indexed by hook ``id``. + extra_dict (Union[dict, List[dict]]): An additional dictionary or list of + dictionaries whose keys will be deleted when the same keys are + removed from ``hooks_dict``. + """ + + id: int + next_id: int = 0 + + def __init__(self, hooks_dict: Any, *, extra_dict: Any = None) -> None: + self.hooks_dict_ref = weakref.ref(hooks_dict) + self.id = RemovableHandle.next_id + RemovableHandle.next_id += 1 + + self.extra_dict_ref: tuple = () + if isinstance(extra_dict, dict): + self.extra_dict_ref = (weakref.ref(extra_dict),) + elif isinstance(extra_dict, list): + self.extra_dict_ref = tuple(weakref.ref(d) for d in extra_dict) + + def remove(self) -> None: + hooks_dict = self.hooks_dict_ref() + if hooks_dict is not None and self.id in hooks_dict: + del hooks_dict[self.id] + + for ref in self.extra_dict_ref: + extra_dict = ref() + if extra_dict is not None and self.id in extra_dict: + del extra_dict[self.id] + + def __getstate__(self): + if self.extra_dict_ref is None: + return (self.hooks_dict_ref(), self.id) + else: + return (self.hooks_dict_ref(), self.id, tuple(ref() for ref in self.extra_dict_ref)) + + def __setstate__(self, state) -> None: + if state[0] is None: + # create a dead reference + self.hooks_dict_ref = weakref.ref(OrderedDict()) + else: + self.hooks_dict_ref = weakref.ref(state[0]) + self.id = state[1] + RemovableHandle.next_id = max(RemovableHandle.next_id, self.id + 1) + + if len(state) < 3 or state[2] is None: + self.extra_dict_ref = () + else: + self.extra_dict_ref = tuple(weakref.ref(d) for d in state[2]) + + def __enter__(self) -> "RemovableHandle": + return self + + def __exit__(self, type: Any, value: Any, tb: Any) -> None: + self.remove() + + +def unserializable_hook(f): + """ + Mark a function as an unserializable hook with this decorator. + + This suppresses warnings that would otherwise arise if you attempt + to serialize a tensor that has a hook. + """ + f.__torch_unserializable__ = True + return f + + +def warn_if_has_hooks(tensor) -> None: + if tensor._backward_hooks: + for k in tensor._backward_hooks: + hook = tensor._backward_hooks[k] + if not hasattr(hook, "__torch_unserializable__"): + warnings.warn(f"backward hook {repr(hook)} on tensor will not be " + "serialized. If this is expected, you can " + "decorate the function with @torch.utils.hooks.unserializable_hook " + "to suppress this warning", stacklevel=2) + +class BackwardHook: + """ + A wrapper class to implement nn.Module backward hooks. + + It handles: + - Ignoring non-Tensor inputs and replacing them by None before calling the user hook + - Generating the proper Node to capture a set of Tensor's gradients + - Linking the gradients captures for the outputs with the gradients captured for the input + - Calling the user hook once both output and input gradients are available + """ + + def __init__(self, module, user_hooks, user_pre_hooks) -> None: + self.user_hooks = user_hooks + self.user_pre_hooks = user_pre_hooks + self.module = module + + self.grad_outputs = None + self.n_outputs = -1 + self.output_tensors_index = None + self.n_inputs = -1 + self.input_tensors_index = None + + def _pack_with_none(self, indices, values, size): + res = [None] * size + for idx, val in zip(indices, values, strict=True): + res[idx] = val + + return tuple(res) + + def _unpack_none(self, indices, values): + res = [values[idx] for idx in indices] + + return tuple(res) + + def _set_user_hook(self, grad_fn) -> None: + def hook(grad_input, _): + if self.grad_outputs is None: + # This happens because the gradient in your nn.Module flows to + # the Module's input without " passing through the Module's + # output, e.g. when you're doing double backward. + return + res = self._pack_with_none(self.input_tensors_index, grad_input, self.n_inputs) + + for hook in self.user_hooks: + out = hook(self.module, res, self.grad_outputs) + + if out is None: + continue + + if len(out) != len(res): + raise RuntimeError("Backward hook returned an invalid number of grad_input, " + f"got {len(out)}, but expected {len(res)}") + + res = out + + self.grad_outputs = None + + return self._unpack_none(self.input_tensors_index, res) + + grad_fn.register_hook(hook) + + def _apply_on_tensors(self, fn, args): + # Can be used to apply the given function to the tensors contained in the + # args. Will return updated args and the tensors indices + tensors_idx = [] + tensors = [] + + requires_grad = False + for i, arg in enumerate(args): + if isinstance(arg, torch.Tensor): + tensors_idx.append(i) + tensors.append(arg) + requires_grad |= arg.requires_grad + + if not (requires_grad and torch.is_grad_enabled()): + return args, None + + new_tensors = torch.nn.modules._functions.BackwardHookFunction.apply(*tensors) + if len(new_tensors) == 0: + raise RuntimeError("Cannot set Module backward hook for a Module with no input Tensors.") + + grad_fns = [t.grad_fn for t in new_tensors if t.grad_fn is not None and t.grad_fn.name() == "BackwardHookFunctionBackward"] + if len(grad_fns) == 0: + raise RuntimeError("Error while setting up backward hooks. Please open " + "an issue with a code sample to reproduce this.") + + fn(grad_fns[0]) + + arg_list = list(args) + for idx, val in zip(tensors_idx, new_tensors, strict=True): + arg_list[idx] = val + + if type(args) is tuple: + out = tuple(arg_list) + else: + out = type(args)(*arg_list) + return out, tensors_idx + + def setup_input_hook(self, args): + def fn(grad_fn) -> None: + self._set_user_hook(grad_fn) + + res, input_idx = self._apply_on_tensors(fn, args) + self.n_inputs = len(args) + self.input_tensors_index = input_idx + return res + + def setup_output_hook(self, args): + def fn(grad_fn) -> None: + def hook(_, grad_output): + self.grad_outputs = self._pack_with_none(self.output_tensors_index, + grad_output, + self.n_outputs) + + if self.user_pre_hooks: + expected_len = len(self.grad_outputs) + for user_pre_hook in self.user_pre_hooks: + hook_grad_outputs = user_pre_hook(self.module, self.grad_outputs) + if hook_grad_outputs is None: + continue + + actual_len = len(hook_grad_outputs) + if actual_len != expected_len: + raise RuntimeError("Backward pre hook returned an invalid number of grad_output, " + f"got {actual_len}, but expected {expected_len}") + self.grad_outputs = hook_grad_outputs + + # We need to be able to clear self.grad_outputs but also return it + local_grad_outputs = self.grad_outputs + + # Special case if no input required gradients, this hook should call the user + # hook directly + if self.input_tensors_index is None: + warnings.warn("Full backward hook is firing when gradients are computed " + "with respect to module outputs since no inputs require gradients. See " + "https://docs.pytorch.org/docs/main/generated/torch.nn.Module.html#torch.nn.Module.register_full_backward_hook " # noqa: B950 + "for more details.", + stacklevel=5) + grad_inputs = self._pack_with_none([], [], self.n_inputs) + for user_hook in self.user_hooks: + res = user_hook(self.module, grad_inputs, self.grad_outputs) + if res is not None and not (isinstance(res, tuple) and all(el is None for el in res)): + raise RuntimeError("Backward hook for Modules where no input requires " + "gradient should always return None or None for all gradients.") + self.grad_outputs = None + + if local_grad_outputs is not None: + if self.output_tensors_index is None: + raise AssertionError("output_tensors_index should not be None when grad_outputs is not None") + return tuple(local_grad_outputs[i] for i in self.output_tensors_index) + + grad_fn.register_hook(hook) + + is_tuple = True + if not isinstance(args, tuple): + args = (args,) + is_tuple = False + + res, output_idx = self._apply_on_tensors(fn, args) + self.n_outputs = len(args) + self.output_tensors_index = output_idx + + if not is_tuple: + res = res[0] + return res diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/jit/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/jit/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/jit/log_extract.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/jit/log_extract.py new file mode 100644 index 0000000000000000000000000000000000000000..9e018457802f4aafd05ba6a8d10ef1c4953b1047 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/jit/log_extract.py @@ -0,0 +1,118 @@ +# mypy: allow-untyped-defs +from contextlib import contextmanager +from typing import Any, cast +import random +import torch +import time +from torch.utils.benchmark import Timer + +def extract_ir(filename: str) -> list[str]: + BEGIN = "" + END = "" + pfx = None + graphs = [] + with open(filename) as f: + split_strs = f.read().split(BEGIN) + for i, split_str in enumerate(split_strs): + if i == 0: + continue + end_loc = split_str.find(END) + if end_loc == -1: + continue + s = split_str[:end_loc] + pfx = split_strs[i - 1].splitlines()[-1] + lines = [x[len(pfx):] for x in s.splitlines(keepends=True)] + graphs.append(''.join(lines)) + + return graphs + + +def make_tensor_from_type(inp_type: torch._C.TensorType): + size = inp_type.sizes() + stride = inp_type.strides() + device = inp_type.device() + dtype = inp_type.dtype() + if size is None: + raise AssertionError("make_tensor_from_type: 'size' is None (inp_type.sizes() returned None)") + if stride is None: + raise AssertionError("make_tensor_from_type: 'stride' is None (inp_type.strides() returned None)") + if device is None: + raise AssertionError("make_tensor_from_type: 'device' is None (inp_type.device() returned None)") + if dtype is None: + raise AssertionError("make_tensor_from_type: 'dtype' is None (inp_type.dtype() returned None)") + return torch.empty_strided(size=size, stride=stride, device=device, dtype=dtype) + +def load_graph_and_inputs(ir: str) -> tuple[Any, list[Any]]: + graph = torch._C.parse_ir(ir, parse_tensor_constants=True) + graph.makeMultiOutputIntoTuple() + inputs = [] + for inp in graph.inputs(): + if isinstance(inp.type(), torch._C.FloatType): + inputs.append(random.uniform(.1, 100)) + elif isinstance(inp.type(), torch._C.IntType): + inputs.append(random.randint(1, 100)) + elif isinstance(inp.type(), torch._C.TensorType): + tensorType = cast(torch._C.TensorType, inp.type()) + inputs.append(make_tensor_from_type(tensorType)) + elif isinstance(inp.type(), torch._C.BoolType): + inputs.append(random.randint(0, 1) == 1) + else: + raise NotImplementedError(f"A default value is not implemented for type {inp.type()}") + + func = torch._C._create_function_from_graph("forward", graph) + torch._C._jit_pass_erase_shape_information(func.graph) + return (func, inputs) + +def time_cuda(fn, inputs, test_runs): + t = Timer(stmt="fn(*inputs)", globals={"fn": fn, "inputs" : inputs}) + times = t.blocked_autorange() + return times.median * 1000 # time in ms + +def time_cpu(fn, inputs, test_runs): + s = time.perf_counter() + for _ in range(test_runs): + fn(*inputs) + e = time.perf_counter() + return (e - s) / test_runs * 1000 # time in ms + +def run_test(ir, inputs, *, warmup_runs=10, test_runs=20) -> float: + graph, _ = load_graph_and_inputs(ir) + for _ in range(warmup_runs): + graph(*inputs) + + is_cpu = None + for input in inputs: + if isinstance(input, torch.Tensor): + is_cpu = input.device.type == "cpu" + break + if is_cpu is None: + raise AssertionError("No tensor found in inputs") + + out = time_cpu(graph, inputs, test_runs) if is_cpu else time_cuda(graph, inputs, test_runs) + return out + +@contextmanager +def no_fuser(*args, **kwargs): + old_optimize = torch._C._get_graph_executor_optimize(False) + try: + yield + finally: + torch._C._get_graph_executor_optimize(old_optimize) + +def run_baseline_no_fusion(ir, inputs) -> float: + with no_fuser(): + return run_test(ir, inputs) + + +def run_nnc(ir, inputs, dynamic) -> float: + try: + strat = [("DYNAMIC", 10)] if dynamic else [("STATIC", 10)] + old_strat = torch.jit.set_fusion_strategy(strat) + with torch.jit.fuser("fuser1"): + return run_test(ir, inputs) + finally: + torch.jit.set_fusion_strategy(old_strat) + +def run_nvfuser(ir, inputs) -> float: + with torch.jit.fuser("fuser2"): + return run_test(ir, inputs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/mkldnn.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/mkldnn.py new file mode 100644 index 0000000000000000000000000000000000000000..11bb4e442b2960a601c3c6c66c5e326ac3c9c166 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/mkldnn.py @@ -0,0 +1,238 @@ +# mypy: allow-untyped-defs +import torch + + +class MkldnnLinear(torch.jit.ScriptModule): + def __init__(self, dense_module, dtype) -> None: + super().__init__() + self.register_buffer('weight', dense_module.weight.to_mkldnn(dtype)) + if dense_module.bias is not None: + # Bias can be fp32 or bf16 for OneDNN bf16 path, but for good accuracy, + # we use fp32 dtype. + self.register_buffer('bias', dense_module.bias.to_mkldnn()) + else: + # TODO: Remove this once ScriptModule supports registering None buffer + self.register_buffer( + 'bias', + torch.zeros([dense_module.weight.size(0)], dtype=torch.float).to_mkldnn()) + + @torch.jit.script_method + def __getstate__(self): + return (self.weight.to_dense(), self.bias.to_dense(), self.training) + + @torch.jit.script_method + def __setstate__(self, state): + self.weight = state[0].to_mkldnn() + self.bias = state[1].to_mkldnn() + self.training = state[2] + + @torch.jit.script_method + def forward(self, x): + x_mkldnn = x if x.is_mkldnn else x.to_mkldnn() + y_mkldnn = torch._C._nn.mkldnn_linear(x_mkldnn, self.weight, self.bias) + y = y_mkldnn if x.is_mkldnn else y_mkldnn.to_dense() + return y + + +class _MkldnnConvNd(torch.jit.ScriptModule): + """Common base of MkldnnConv1d and MkldnnConv2d.""" + + __constants__ = ['stride', 'padding', 'dilation', 'groups'] + + def __init__(self, dense_module) -> None: + super().__init__() + + self.stride = dense_module.stride + self.padding = dense_module.padding + self.dilation = dense_module.dilation + self.groups = dense_module.groups + + if dense_module.bias is not None: + self.register_buffer('bias', dense_module.bias.to_mkldnn()) + else: + # Bias can be fp32 or bf16 for OneDNN bf16 path, but for good accuracy, + # we use fp32 dtype. + # TODO: Remove this once ScriptModule supports registering None buffer + self.register_buffer( + 'bias', + torch.zeros([dense_module.weight.size(0)], dtype=torch.float).to_mkldnn()) + + @torch.jit.script_method + def __getstate__(self): + return (self.weight.to_dense(), self.bias.to_dense(), self.training) + + @torch.jit.script_method + def forward(self, x): + return torch.mkldnn_convolution( + x, + self.weight, + self.bias, + self.padding, + self.stride, + self.dilation, + self.groups) + + +class MkldnnConv1d(_MkldnnConvNd): + def __init__(self, dense_module, dtype) -> None: + super().__init__(dense_module) + + self.register_buffer('weight', dense_module.weight.to_mkldnn(dtype)) + + @torch.jit.script_method + def __setstate__(self, state): + self.weight = state[0].to_mkldnn() + self.bias = state[1].to_mkldnn() + self.training = state[2] + + +class MkldnnConv2d(_MkldnnConvNd): + def __init__(self, dense_module, dtype) -> None: + super().__init__(dense_module) + + self.register_buffer('weight', torch._C._nn.mkldnn_reorder_conv2d_weight( + dense_module.weight.to_mkldnn(dtype), + self.padding, + self.stride, + self.dilation, + self.groups)) + + @torch.jit.script_method + def __setstate__(self, state): + self.weight = torch._C._nn.mkldnn_reorder_conv2d_weight( + state[0].to_mkldnn(), + self.padding, + self.stride, + self.dilation, + self.groups) + self.bias = state[1].to_mkldnn() + self.training = state[2] + +class MkldnnConv3d(_MkldnnConvNd): + def __init__(self, dense_module, dtype) -> None: + super().__init__(dense_module) + + self.register_buffer('weight', torch._C._nn.mkldnn_reorder_conv3d_weight( + dense_module.weight.to_mkldnn(dtype), + self.padding, + self.stride, + self.dilation, + self.groups)) + + @torch.jit.script_method + def __setstate__(self, state): + self.weight = torch._C._nn.mkldnn_reorder_conv3d_weight( + state[0].to_mkldnn(), + self.padding, + self.stride, + self.dilation, + self.groups) + self.bias = state[1].to_mkldnn() + self.training = state[2] + + +class MkldnnBatchNorm(torch.jit.ScriptModule): + __constants__ = ['exponential_average_factor', 'eps'] + + def __init__(self, dense_module) -> None: + super().__init__() + + if dense_module.training: + raise AssertionError("Only support eval mode batchnorm for mkldnn path now") + if not dense_module.track_running_stats: + raise AssertionError("Only support track_running_stats=True for mkldnn path now") + if not dense_module.affine: + raise AssertionError("Only support affine=True for mkldnn path now") + + if dense_module.momentum is None: + self.exponential_average_factor = 0.0 + else: + self.exponential_average_factor = dense_module.momentum + self.eps = dense_module.eps + + self.register_buffer('weight', dense_module.weight.to_mkldnn()) + self.register_buffer('bias', dense_module.bias.to_mkldnn()) + self.register_buffer('running_mean', dense_module.running_mean.to_mkldnn()) + self.register_buffer('running_var', dense_module.running_var.to_mkldnn()) + + @torch.jit.script_method + def __getstate__(self): + weight = self.weight.to_dense() + bias = self.bias.to_dense() + running_mean = self.running_mean.to_dense() + running_var = self.running_var.to_dense() + return (weight, bias, running_mean, running_var, self.training) + + @torch.jit.script_method + def __setstate__(self, state): + self.weight = state[0].to_mkldnn() + self.bias = state[1].to_mkldnn() + self.running_mean = state[2].to_mkldnn() + self.running_var = state[3].to_mkldnn() + self.training = state[4] + + @torch.jit.script_method + def forward(self, x): + return torch.batch_norm( + x, + self.weight, + self.bias, + self.running_mean, + self.running_var, + False, # training + self.exponential_average_factor, + self.eps, + False, # cuda_enabled + ) + +class MkldnnPrelu(torch.jit.ScriptModule): + def __init__(self, dense_module, dtype) -> None: + super().__init__() + self.register_buffer('weight', dense_module.weight.to_mkldnn(dtype)) + + @torch.jit.script_method + def __getstate__(self): + return (self.weight.to_dense(), self.training) + + @torch.jit.script_method + def __setstate__(self, state): + self.weight = state[0].to_mkldnn() + self.training = state[1] + + @torch.jit.script_method + def forward(self, x): + x_mkldnn = x if x.is_mkldnn else x.to_mkldnn() + y_mkldnn = torch.prelu(x_mkldnn, self.weight) + y = y_mkldnn if x.is_mkldnn else y_mkldnn.to_dense() + return y + +def to_mkldnn(module, dtype=torch.float): + if dtype not in (torch.float, torch.bfloat16, torch.half): + raise AssertionError("MKLDNN only support float, bfloat16, and half path now") + + + def m_fn(m, d): + if isinstance(m, torch.nn.Linear): + return MkldnnLinear(m, d) + elif isinstance(m, torch.nn.Conv1d): + return MkldnnConv1d(m, d) + elif isinstance(m, torch.nn.Conv2d): + return MkldnnConv2d(m, d) + elif isinstance(m, torch.nn.Conv3d): + return MkldnnConv3d(m, d) + elif isinstance(m, (torch.nn.BatchNorm2d, torch.nn.BatchNorm3d)): + # For batchnorm bf16 path, OneDNN requires weight and bias need fp32 dtype. + # so it doesn't need dtype argument. + return MkldnnBatchNorm(m) + elif isinstance(m, torch.nn.PReLU): + return MkldnnPrelu(m, d) + else: + return m + + def m_fn_rec(m, d): + new_m = m_fn(m, d) + for name, sub_m in m.named_children(): + setattr(new_m, name, m_fn_rec(sub_m, d)) + return new_m + + return m_fn_rec(module, dtype) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/mobile_optimizer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/mobile_optimizer.py new file mode 100644 index 0000000000000000000000000000000000000000..1ad0a65204a4733323e7ed29a51403aa47556bbd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/mobile_optimizer.py @@ -0,0 +1,135 @@ +# mypy: allow-untyped-defs +"""This module contains utility method for mobile model optimization and lint.""" + +import torch +from enum import Enum +from torch._C import _MobileOptimizerType as MobileOptimizerType +from typing import AnyStr + +class LintCode(Enum): + BUNDLED_INPUT = 1 + REQUIRES_GRAD = 2 + DROPOUT = 3 + BATCHNORM = 4 + +def optimize_for_mobile( + script_module: torch.jit.ScriptModule, + optimization_blocklist: set[MobileOptimizerType] | None = None, + preserved_methods: list[AnyStr] | None = None, + backend: str = 'CPU') -> torch.jit.RecursiveScriptModule: + """ + Optimize a torch script module for mobile deployment. + + Args: + script_module: An instance of torch script module with type of ScriptModule. + optimization_blocklist: A set with type of MobileOptimizerType. When set is not passed, + optimization method will run all the optimizer pass; otherwise, optimizer + method will run the optimization pass that is not included inside optimization_blocklist. + preserved_methods: A list of methods that needed to be preserved when freeze_module pass is invoked + backend: Device type to use for running the result model ('CPU'(default), 'Vulkan' or 'Metal'). + Returns: + A new optimized torch script module + """ + if not isinstance(script_module, torch.jit.ScriptModule): + raise TypeError( + f'Got {type(script_module)}, but ScriptModule is expected.') + + if optimization_blocklist is None: + optimization_blocklist = set() + + if preserved_methods is None: + preserved_methods = [] + + # Convert potential byte arrays into strings (if there is any) to pass type checking + # Here we use a new name as assigning it back to preserved_methods will invoke + # mypy errors (i.e. List[AnyStr] = List[str]) + preserved_methods_str: list[str] = [str(method) for method in preserved_methods] + + bundled_inputs_attributes = _get_bundled_inputs_preserved_attributes(script_module, preserved_methods_str) + if all(hasattr(script_module, method) for method in bundled_inputs_attributes): + preserved_methods_str = list(set(preserved_methods_str + bundled_inputs_attributes)) + + non_exist_methods = [method for method in preserved_methods_str if not hasattr(script_module, method)] + if non_exist_methods: + raise AttributeError( + f"The following methods to preserve do not exist in script_module: {', '.join(non_exist_methods)}") + + backend = backend.lower() + if backend == 'cpu': + optimized_cpp_module = torch._C._jit_pass_optimize_for_mobile( + script_module._c, + optimization_blocklist, + preserved_methods_str) + elif backend == 'vulkan': + optimized_cpp_module = torch._C._jit_pass_vulkan_optimize_for_mobile( + script_module._c, + optimization_blocklist, + preserved_methods_str) + elif backend == 'metal': + optimized_cpp_module = torch._C._jit_pass_metal_optimize_for_mobile(script_module._c, preserved_methods_str) + else: + raise TypeError("Unknown backend, must be one of 'CPU', 'Vulkan' or 'Metal'") + + return torch.jit._recursive.wrap_cpp_module(optimized_cpp_module) + + +def generate_mobile_module_lints(script_module: torch.jit.ScriptModule): + """ + Generate a list of lints for a given torch script module. + + Args: + script_module: An instance of torch script module with type of ScriptModule. + + Returns: + lint_map: A list of dictionary that contains modules lints + """ + if not isinstance(script_module, torch.jit.ScriptModule): + raise TypeError( + f'Got {type(script_module)}, but ScriptModule is expected.') + + lint_list = [] + + if not hasattr(script_module, "_generate_bundled_inputs_for_forward"): + lint_list.append({"name": LintCode.BUNDLED_INPUT.name, "message": "No bundled input for forward, please add bundled inputs " + "before saving the module using torch.utils.bundled_inputs.augment_model_with_bundled_inputs."}) + + for name, param in script_module.named_parameters(): + if param.requires_grad: + lint_list.append({"name": LintCode.REQUIRES_GRAD.name, "message": f"Param {name} requires grad, " + "please set torch.no_grad() to reduce memory usage and improve computation speed during " + "inference phase."}) + + op_names = torch.jit.export_opnames(script_module) + for op_name in op_names: + if "dropout" in op_name: + lint_list.append({"name": LintCode.DROPOUT.name, + "message": f"Operator {op_name} exists, remember to call eval() before " + "saving the module.and call torch.utils.mobile_optimizer.optimize_for_mobile to drop dropout " + "operator."}) + if "batch_norm" in op_name: + lint_list.append({"name": LintCode.BATCHNORM.name, + "message": f"Operator {op_name} exists, remember to call eval() before " + "saving the module and call torch.utils.mobile_optimizer.optimize_for_mobile to drop batch_norm " + "operator."}) + + return lint_list + +def _get_bundled_inputs_preserved_attributes(script_module: torch.jit.ScriptModule, preserved_methods: list[str]) -> list[str]: + + bundled_inputs_attributes = [] + # Has bundled inputs for forward + if hasattr(script_module, 'get_all_bundled_inputs'): + bundled_inputs_attributes.append('get_all_bundled_inputs') + bundled_inputs_attributes.append('get_num_bundled_inputs') + + # Bundled inputs in module after the change that introduced bundled inputs for multiple functions + if hasattr(script_module, 'get_bundled_inputs_functions_and_info'): + bundled_inputs_attributes.append('get_bundled_inputs_functions_and_info') + all_info = script_module.get_bundled_inputs_functions_and_info() + for function_name in all_info: + if function_name not in preserved_methods: + bundled_inputs_attributes.append(function_name) + bundled_inputs_attributes.append("get_all_bundled_inputs_for_" + function_name) + bundled_inputs_attributes.append("_bundled_inputs_deflated_" + function_name) + + return bundled_inputs_attributes diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..16d1ab1c6dd1a2d422cae74eaa5b5888dd2fa175 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/__init__.py @@ -0,0 +1,450 @@ +#!/usr/bin/env python3 +# mypy: allow-untyped-defs +""" +model_dump: a one-stop shop for TorchScript model inspection. + +The goal of this tool is to provide a simple way to extract lots of +useful information from a TorchScript model and make it easy for humans +to consume. It (mostly) replaces zipinfo, common uses of show_pickle, +and various ad-hoc analysis notebooks. + +The tool extracts information from the model and serializes it as JSON. +That JSON can then be rendered by an HTML+JS page, either by +loading the JSON over HTTP or producing a fully self-contained page +with all of the code and data burned-in. +""" + +# Maintainer notes follow. +""" +The implementation strategy has tension between 3 goals: +- Small file size. +- Fully self-contained. +- Easy, modern JS environment. +Using Preact and HTM achieves 1 and 2 with a decent result for 3. +However, the models I tested with result in ~1MB JSON output, +so even using something heavier like full React might be tolerable +if the build process can be worked out. + +One principle I have followed that I think is very beneficial +is to keep the JSON data as close as possible to the model +and do most of the rendering logic on the client. +This makes for easier development (just refresh, usually), +allows for more laziness and dynamism, and lets us add more +views of the same data without bloating the HTML file. + +Currently, this code doesn't actually load the model or even +depend on any part of PyTorch. I don't know if that's an important +feature to maintain, but it's probably worth preserving the ability +to run at least basic analysis on models that cannot be loaded. + +I think the easiest way to develop this code is to cd into model_dump and +run "python -m http.server", then load http://localhost:8000/skeleton.html +in the browser. In another terminal, run +"python -m torch.utils.model_dump --style=json FILE > \ + torch/utils/model_dump/model_info.json" +every time you update the Python code or model. +When you update JS, just refresh. + +Possible improvements: + - Fix various TODO comments in this file and the JS. + - Make the HTML much less janky, especially the auxiliary data panel. + - Make the auxiliary data panel start small, expand when + data is available, and have a button to clear/contract. + - Clean up the JS. There's a lot of copypasta because + I don't really know how to use Preact. + - Make the HTML render and work nicely inside a Jupyter notebook. + - Add the ability for JS to choose the URL to load the JSON based + on the page URL (query or hash). That way we could publish the + inlined skeleton once and have it load various JSON blobs. + - Add a button to expand all expandable sections so ctrl-F works well. + - Add hyperlinking from data to code, and code to code. + - Add hyperlinking from debug info to Diffusion. + - Make small tensor contents available. + - Do something nice for quantized models + (they probably don't work at all right now). +""" + +import argparse +import io +import itertools +import json +import os +import pickle +import pprint +import re +import sys +import urllib.parse +import zipfile +from pathlib import Path +import warnings + +import torch.utils.show_pickle + + +DEFAULT_EXTRA_FILE_SIZE_LIMIT = 16 * 1024 + +__all__ = ['get_storage_info', 'hierarchical_pickle', 'get_model_info', 'get_inline_skeleton', + 'burn_in_info', 'get_info_and_burn_skeleton'] + +def get_storage_info(storage): + if not isinstance(storage, torch.utils.show_pickle.FakeObject): + raise AssertionError(f"storage is not FakeObject: {type(storage)}") + if storage.module != "pers": + raise AssertionError(f"storage.module is not 'pers': {storage.module!r}") + if storage.name != "obj": + raise AssertionError(f"storage.name is not 'obj': {storage.name!r}") + if storage.state is not None: + raise AssertionError(f"storage.state is not None: {storage.state!r}") + if not isinstance(storage.args, tuple): + raise AssertionError(f"storage.args is not a tuple: {type(storage.args)}") + if len(storage.args) != 1: + raise AssertionError(f"len(storage.args) is not 1: {len(storage.args)}") + sa = storage.args[0] + if not isinstance(sa, tuple): + raise AssertionError(f"sa is not a tuple: {type(sa)}") + if len(sa) != 5: + raise AssertionError(f"len(sa) is not 5: {len(sa)}") + if sa[0] != "storage": + raise AssertionError(f"sa[0] is not 'storage': {sa[0]!r}") + if not isinstance(sa[1], torch.utils.show_pickle.FakeClass): + raise AssertionError(f"sa[1] is not FakeClass: {type(sa[1])}") + if sa[1].module != "torch": + raise AssertionError(f"sa[1].module is not 'torch': {sa[1].module!r}") + if not sa[1].name.endswith("Storage"): + raise AssertionError(f"sa[1].name does not end with 'Storage': {sa[1].name!r}") + storage_info = [sa[1].name.replace("Storage", "")] + list(sa[2:]) + return storage_info + + +def hierarchical_pickle(data): + if isinstance(data, (bool, int, float, str, type(None))): + return data + if isinstance(data, list): + return [hierarchical_pickle(d) for d in data] + if isinstance(data, tuple): + return { + "__tuple_values__": hierarchical_pickle(list(data)), + } + if isinstance(data, dict): + return { + "__is_dict__": True, + "keys": hierarchical_pickle(list(data.keys())), + "values": hierarchical_pickle(list(data.values())), + } + if isinstance(data, torch.utils.show_pickle.FakeObject): + typename = f"{data.module}.{data.name}" + if ( + typename.startswith(('__torch__.', 'torch.jit.LoweredWrapper.', 'torch.jit.LoweredModule.')) + ): + if data.args != (): + raise AssertionError("data.args is not ()") + return { + "__module_type__": typename, + "state": hierarchical_pickle(data.state), + } + if typename == "torch._utils._rebuild_tensor_v2": + if data.state is not None: + raise AssertionError("data.state is not None") + storage, offset, size, stride, requires_grad, *_ = data.args + storage_info = get_storage_info(storage) + return {"__tensor_v2__": [storage_info, offset, size, stride, requires_grad]} + if typename == "torch._utils._rebuild_qtensor": + if data.state is not None: + raise AssertionError("data.state is not None") + storage, offset, size, stride, quantizer, requires_grad, *_ = data.args + storage_info = get_storage_info(storage) + if not isinstance(quantizer, tuple): + raise AssertionError("quantizer is not a tuple") + if not isinstance(quantizer[0], torch.utils.show_pickle.FakeClass): + raise AssertionError("quantizer[0] is not a FakeClass") + if quantizer[0].module != "torch": + raise AssertionError("quantizer[0].module is not torch") + if quantizer[0].name == "per_tensor_affine": + if len(quantizer) != 3: + raise AssertionError("len(quantizer) is not 3") + if not isinstance(quantizer[1], float): + raise AssertionError("quantizer[1] is not a float") + if not isinstance(quantizer[2], int): + raise AssertionError("quantizer[2] is not an int") + quantizer_extra = list(quantizer[1:3]) + else: + quantizer_extra = [] + quantizer_json = [quantizer[0].name] + quantizer_extra + return {"__qtensor__": [storage_info, offset, size, stride, quantizer_json, requires_grad]} + if typename == "torch.jit._pickle.restore_type_tag": + if data.state is not None: + raise AssertionError("data.state is not None") + obj, typ = data.args + if not isinstance(typ, str): + raise AssertionError("typ is not a string") + return hierarchical_pickle(obj) + if re.fullmatch(r"torch\.jit\._pickle\.build_[a-z]+list", typename): + if data.state is not None: + raise AssertionError("data.state is not None") + ls, = data.args + if not isinstance(ls, list): + raise AssertionError("ls is not a list") + return hierarchical_pickle(ls) + if typename == "torch.device": + if data.state is not None: + raise AssertionError("data.state is not None") + name, = data.args + if not isinstance(name, str): + raise AssertionError("name is not a string") + # Just forget that it was a device and return the name. + return name + if typename == "builtin.UnicodeDecodeError": + if data.state is not None: + raise AssertionError("data.state is not None") + msg, = data.args + if not isinstance(msg, str): + raise AssertionError("msg is not a string") + # Hack: Pretend this is a module so we don't need custom serialization. + # Hack: Wrap the message in a tuple so it looks like a nice state object. + # TODO: Undo at least that second hack. We should support string states. + return { + "__module_type__": typename, + "state": hierarchical_pickle((msg,)), + } + raise Exception(f"Can't prepare fake object of type for JS: {typename}") # noqa: TRY002 + raise Exception(f"Can't prepare data of type for JS: {type(data)}") # noqa: TRY002 + + +def get_model_info( + path_or_file, + title=None, + extra_file_size_limit=DEFAULT_EXTRA_FILE_SIZE_LIMIT): + """Get JSON-friendly information about a model. + + The result is suitable for being saved as model_info.json, + or passed to burn_in_info. + """ + + if isinstance(path_or_file, os.PathLike): + default_title = os.fspath(path_or_file) + file_size = path_or_file.stat().st_size # type: ignore[attr-defined] + elif isinstance(path_or_file, str): + default_title = path_or_file + file_size = Path(path_or_file).stat().st_size + else: + default_title = "buffer" + path_or_file.seek(0, io.SEEK_END) + file_size = path_or_file.tell() + path_or_file.seek(0) + + title = title or default_title + + with zipfile.ZipFile(path_or_file) as zf: + path_prefix = None + zip_files = [] + # pyrefly: ignore [bad-assignment] + for zi in zf.infolist(): + prefix = re.sub("/.*", "", zi.filename) + if path_prefix is None: + path_prefix = prefix + elif prefix != path_prefix: + raise Exception(f"Mismatched prefixes: {path_prefix} != {prefix}") # noqa: TRY002 + zip_files.append( + { + "filename": zi.filename, + "compression": zi.compress_type, + "compressed_size": zi.compress_size, + "file_size": zi.file_size, + } + ) + if path_prefix is None: + raise AssertionError("path_prefix is None") + version = zf.read(path_prefix + "/version").decode("utf-8").strip() + + def get_pickle(name): + if path_prefix is None: + raise AssertionError("path_prefix is None") + with zf.open(path_prefix + f"/{name}.pkl") as handle: + raw = torch.utils.show_pickle.DumpUnpickler(handle, catch_invalid_utf8=True).load() + return hierarchical_pickle(raw) + + model_data = get_pickle("data") + constants = get_pickle("constants") + + # Intern strings that are likely to be reused. + # Pickle automatically detects shared structure, + # so reused strings are stored efficiently. + # However, JSON has no way of representing this, + # so we have to do it manually. + interned_strings : dict[str, int] = {} + + def intern(s): + if s not in interned_strings: + interned_strings[s] = len(interned_strings) + return interned_strings[s] + + code_files = {} + for zi in zf.infolist(): + if not zi.filename.endswith(".py"): + continue + with zf.open(zi) as handle: + raw_code = handle.read() + with zf.open(zi.filename + ".debug_pkl") as handle: + raw_debug = handle.read() + + # Parse debug info and add begin/end markers if not present + # to ensure that we cover the entire source code. + debug_info_t = pickle.loads(raw_debug) + text_table = None + + if (len(debug_info_t) == 3 and + isinstance(debug_info_t[0], str) and + debug_info_t[0] == 'FORMAT_WITH_STRING_TABLE'): + _, text_table, content = debug_info_t + + def parse_new_format(line): + # (0, (('', '', 0), 0, 0)) + num, ((text_indexes, fname_idx, offset), start, end), tag = line + text = ''.join(text_table[x] for x in text_indexes) # type: ignore[index] + fname = text_table[fname_idx] # type: ignore[index] + return num, ((text, fname, offset), start, end), tag + + debug_info_t = map(parse_new_format, content) + + debug_info = list(debug_info_t) + if not debug_info: + debug_info.append((0, (('', '', 0), 0, 0))) + if debug_info[-1][0] != len(raw_code): + debug_info.append((len(raw_code), (('', '', 0), 0, 0))) + + code_parts = [] + for di, di_next in itertools.pairwise(debug_info): + start, source_range, *_ = di + end = di_next[0] + if end <= start: + raise AssertionError("end is not greater than start") + source, s_start, s_end = source_range + s_text, s_file, s_line = source + # TODO: Handle this case better. TorchScript ranges are in bytes, + # but JS doesn't really handle byte strings. + # if bytes and chars are not equivalent for this string, + # zero out the ranges so we don't highlight the wrong thing. + if len(s_text) != len(s_text.encode("utf-8")): + s_start = 0 + s_end = 0 + text = raw_code[start:end] + code_parts.append([text.decode("utf-8"), intern(s_file), s_line, intern(s_text), s_start, s_end]) + code_files[zi.filename] = code_parts + + extra_files_json_pattern = re.compile(re.escape(path_prefix) + "/extra/.*\\.json") + extra_files_jsons = {} + for zi in zf.infolist(): + if not extra_files_json_pattern.fullmatch(zi.filename): + continue + if zi.file_size > extra_file_size_limit: + continue + with zf.open(zi) as handle: + try: + json_content = json.load(handle) + extra_files_jsons[zi.filename] = json_content + except json.JSONDecodeError: + extra_files_jsons[zi.filename] = "INVALID JSON" + + always_render_pickles = { + "bytecode.pkl", + } + extra_pickles = {} + for zi in zf.infolist(): + if not zi.filename.endswith(".pkl"): + continue + with zf.open(zi) as handle: + # TODO: handle errors here and just ignore the file? + # NOTE: For a lot of these files (like bytecode), + # we could get away with just unpickling, but this should be safer. + obj = torch.utils.show_pickle.DumpUnpickler(handle, catch_invalid_utf8=True).load() + buf = io.StringIO() + pprint.pprint(obj, buf) + contents = buf.getvalue() + # Checked the rendered length instead of the file size + # because pickles with shared structure can explode in size during rendering. + if os.path.basename(zi.filename) not in always_render_pickles and \ + len(contents) > extra_file_size_limit: + continue + extra_pickles[zi.filename] = contents + + return { + "model": { + "title": title, + "file_size": file_size, + "version": version, + "zip_files": zip_files, + "interned_strings": list(interned_strings), + "code_files": code_files, + "model_data": model_data, + "constants": constants, + "extra_files_jsons": extra_files_jsons, + "extra_pickles": extra_pickles, + } + } + + +def get_inline_skeleton(): + """Get a fully-inlined skeleton of the frontend. + + The returned HTML page has no external network dependencies for code. + It can load model_info.json over HTTP, or be passed to burn_in_info. + """ + + import importlib.resources + + # pyrefly: ignore [bad-argument-type] + skeleton = importlib.resources.read_text(__package__, "skeleton.html") + # pyrefly: ignore [bad-argument-type] + js_code = importlib.resources.read_text(__package__, "code.js") + for js_module in ["preact", "htm"]: + # pyrefly: ignore [bad-argument-type] + js_lib = importlib.resources.read_binary(__package__, f"{js_module}.mjs") + js_url = "data:application/javascript," + urllib.parse.quote(js_lib) + js_code = js_code.replace(f"https://unpkg.com/{js_module}?module", js_url) + skeleton = skeleton.replace(' src="./code.js">', ">\n" + js_code) + return skeleton + + +def burn_in_info(skeleton, info): + """Burn model info into the HTML skeleton. + + The result will render the hard-coded model info and + have no external network dependencies for code or data. + """ + + # Note that Python's json serializer does not escape slashes in strings. + # Since we're inlining this JSON directly into a script tag, a string + # containing "" would end the script prematurely and + # mess up our page. Unconditionally escape fixes that. + return skeleton.replace( + "BURNED_IN_MODEL_INFO = null", + "BURNED_IN_MODEL_INFO = " + json.dumps(info, sort_keys=True).replace("/", "\\/")) + + +def get_info_and_burn_skeleton(path_or_bytesio, **kwargs): + model_info = get_model_info(path_or_bytesio, **kwargs) + skeleton = get_inline_skeleton() + page = burn_in_info(skeleton, model_info) + return page + + +def main(argv, *, stdout=None) -> None: + warnings.warn("torch.utils.model_dump is deprecated and will be removed in a future PyTorch release.", stacklevel=2) + parser = argparse.ArgumentParser() + parser.add_argument("--style", choices=["json", "html"]) + parser.add_argument("--title") + parser.add_argument("model") + args = parser.parse_args(argv[1:]) + + info = get_model_info(args.model, title=args.title) + + output = stdout or sys.stdout + + if args.style == "json": + output.write(json.dumps(info, sort_keys=True) + "\n") + elif args.style == "html": + skeleton = get_inline_skeleton() + page = burn_in_info(skeleton, info) + output.write(page) + else: + raise Exception("Invalid style") # noqa: TRY002 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/__main__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..5d4bdac389bb1f270d74efb6c876258d46077110 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/__main__.py @@ -0,0 +1,5 @@ +#!/usr/bin/env python3 +import sys +from . import main + +sys.exit(main(sys.argv)) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/code.js b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/code.js new file mode 100644 index 0000000000000000000000000000000000000000..173ddfb639d847159ee4fdf46691404bf1bbb7a3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/code.js @@ -0,0 +1,689 @@ +import { h, Component, render } from 'https://unpkg.com/preact?module'; +import htm from 'https://unpkg.com/htm?module'; + +const html = htm.bind(h); + +const BURNED_IN_MODEL_INFO = null; + +// https://stackoverflow.com/a/20732091 +function humanFileSize(size) { + if (size == 0) { return "0 B"; } + var i = Math.floor( Math.log(size) / Math.log(1024) ); + return (size / Math.pow(1024, i)).toFixed(2) * 1 + ' ' + ['B', 'kB', 'MB', 'GB', 'TB'][i]; +} + +function caret(down) { + return down ? "\u25BE" : "\u25B8"; +} + +class Blamer { + constructor() { + this.blame_on_click = false; + this.aux_content_pane = null; + } + + setAuxContentPane(pane) { + this.aux_content_pane = pane; + } + + readyBlame() { + this.blame_on_click = true; + } + + maybeBlame(arg) { + if (!this.blame_on_click) { + return; + } + this.blame_on_click = false; + if (!this.aux_content_pane) { + return; + } + this.aux_content_pane.doBlame(arg); + } +} + +let blame = new Blamer(); + +class Hider extends Component { + constructor() { + super(); + this.state = { shown: null }; + } + + componentDidMount() { + this.setState({ shown: this.props.shown === "true" }); + } + + render({name, children}, {shown}) { + let my_caret = html` this.click()} >${caret(shown)}`; + return html`
+

${my_caret} ${name}

+
${shown ? this.props.children : []}
`; + } + + click() { + this.setState({shown: !this.state.shown}); + } +} + +function ModelSizeSection({model: {file_size, zip_files}}) { + let store_size = 0; + let compr_size = 0; + for (const zi of zip_files) { + if (zi.compression === 0) { + // TODO: Maybe check that compressed_size === file_size. + store_size += zi.compressed_size; + } else { + compr_size += zi.compressed_size; + } + } + let zip_overhead = file_size - store_size - compr_size; + // TODO: Better formatting. Right-align this. + return html` + <${Hider} name="Model Size" shown=true> + 
.
+      Model size: ${file_size} (${humanFileSize(file_size)})
+      Stored files: ${store_size} (${humanFileSize(store_size)})
+      Compressed files: ${compr_size} (${humanFileSize(compr_size)})
+      Zip overhead: ${zip_overhead} (${humanFileSize(zip_overhead)})
+
`; +} + +function StructuredDataSection({name, data, shown}) { + return html` + <${Hider} name=${name} shown=${shown}> +
+ <${StructuredData} data=${data} indent="" prefix=""/> +
`; +} + +class StructuredData extends Component { + constructor() { + super(); + this.state = { shown: false }; + + this.INLINE_TYPES = new Set(["boolean", "number", "string"]) + this.IGNORED_STATE_KEYS = new Set(["training", "_is_full_backward_hook"]) + } + + click() { + this.setState({shown: !this.state.shown}); + } + + expando(data) { + if (data === null || this.INLINE_TYPES.has(typeof(data))) { + return false; + } + if (typeof(data) != "object") { + throw new Error("Not an object"); + } + if (Array.isArray(data)) { + // TODO: Maybe show simple lists and tuples on one line. + return true; + } + if (data.__tuple_values__) { + // TODO: Maybe show simple lists and tuples on one line. + return true; + } + if (data.__is_dict__) { + // TODO: Maybe show simple (empty?) dicts on one line. + return true; + } + if (data.__module_type__) { + return true; + } + if (data.__tensor_v2__) { + return false; + } + if (data.__qtensor__) { + return false; + } + throw new Error("Can't handle data type.", data); + } + + renderHeadline(data) { + if (data === null) { + return "None"; + } + if (typeof(data) == "boolean") { + const sd = String(data); + return sd.charAt(0).toUpperCase() + sd.slice(1); + } + if (typeof(data) == "number") { + return JSON.stringify(data); + } + if (typeof(data) == "string") { + return JSON.stringify(data); + } + if (typeof(data) != "object") { + throw new Error("Not an object"); + } + if (Array.isArray(data)) { + return "list(["; + } + if (data.__tuple_values__) { + return "tuple(("; + } + if (data.__is_dict__) { + return "dict({"; + } + if (data.__module_type__) { + return data.__module_type__ + "()"; + } + if (data.__tensor_v2__) { + const [storage, offset, size, stride, grad] = data.__tensor_v2__; + const [dtype, key, device, numel] = storage; + return this.renderTensor( + "tensor", dtype, key, device, numel, offset, size, stride, grad, []); + } + if (data.__qtensor__) { + const [storage, offset, size, stride, quantizer, grad] = data.__qtensor__; + const [dtype, key, device, numel] = storage; + let extra_parts = []; + if (quantizer[0] == "per_tensor_affine") { + extra_parts.push(`scale=${quantizer[1]}`); + extra_parts.push(`zero_point=${quantizer[2]}`); + } else { + extra_parts.push(`quantizer=${quantizer[0]}`); + } + return this.renderTensor( + "qtensor", dtype, key, device, numel, offset, size, stride, grad, extra_parts); + } + throw new Error("Can't handle data type.", data); + } + + renderTensor( + prefix, + dtype, + storage_key, + device, + storage_numel, + offset, + size, + stride, + grad, + extra_parts) { + let parts = [ + "(" + size.join(",") + ")", + dtype, + ]; + parts.push(...extra_parts); + if (device != "cpu") { + parts.push(device); + } + if (grad) { + parts.push("grad"); + } + // TODO: Check stride and indicate if the tensor is channels-last or non-contiguous + // TODO: Check size, stride, offset, and numel and indicate if + // the tensor doesn't use all data in storage. + // TODO: Maybe show key? + void(offset); + void(stride); + void(storage_key); + void(storage_numel); + return prefix + "(" + parts.join(", ") + ")"; + } + + renderBody(indent, data) { + if (data === null || this.INLINE_TYPES.has(typeof(data))) { + throw "Should not reach here." + } + if (typeof(data) != "object") { + throw new Error("Not an object"); + } + if (Array.isArray(data)) { + let new_indent = indent + "\u00A0\u00A0"; + let parts = []; + for (let idx = 0; idx < data.length; idx++) { + // Does it make sense to put explicit index numbers here? + parts.push(html`
<${StructuredData} prefix=${idx + ": "} indent=${new_indent} data=${data[idx]} />`); + } + return parts; + } + if (data.__tuple_values__) { + // Handled the same as lists. + return this.renderBody(indent, data.__tuple_values__); + } + if (data.__is_dict__) { + let new_indent = indent + "\u00A0\u00A0"; + let parts = []; + for (let idx = 0; idx < data.keys.length; idx++) { + if (typeof(data.keys[idx]) != "string") { + parts.push(html`
${new_indent}Non-string key`); + } else { + parts.push(html`
<${StructuredData} prefix=${data.keys[idx] + ": "} indent=${new_indent} data=${data.values[idx]} />`); + } + } + return parts; + } + if (data.__module_type__) { + const mstate = data.state; + if (mstate === null || typeof(mstate) != "object") { + throw new Error("Bad module state"); + } + let new_indent = indent + "\u00A0\u00A0"; + let parts = []; + if (mstate.__is_dict__) { + // TODO: Less copy/paste between this and normal dicts. + for (let idx = 0; idx < mstate.keys.length; idx++) { + if (typeof(mstate.keys[idx]) != "string") { + parts.push(html`
${new_indent}Non-string key`); + } else if (this.IGNORED_STATE_KEYS.has(mstate.keys[idx])) { + // Do nothing. + } else { + parts.push(html`
<${StructuredData} prefix=${mstate.keys[idx] + ": "} indent=${new_indent} data=${mstate.values[idx]} />`); + } + } + } else if (mstate.__tuple_values__) { + parts.push(html`
<${StructuredData} prefix="" indent=${new_indent} data=${mstate} />`); + } else if (mstate.__module_type__) { + // We normally wouldn't have the state of a module be another module, + // but we use "modules" to encode special values (like Unicode decode + // errors) that might be valid states. Just go with it. + parts.push(html`
<${StructuredData} prefix="" indent=${new_indent} data=${mstate} />`); + } else { + throw new Error("Bad module state"); + } + return parts; + } + if (data.__tensor_v2__) { + throw "Should not reach here." + } + if (data.__qtensor__) { + throw "Should not reach here." + } + throw new Error("Can't handle data type.", data); + } + + render({data, indent, prefix}, {shown}) { + const exp = this.expando(data) ? html` this.click()} >${caret(shown)} ` : ""; + const headline = this.renderHeadline(data); + const body = shown ? this.renderBody(indent, data) : ""; + return html`${indent}${exp}${prefix}${headline}${body}`; + } +} + +function ZipContentsSection({model: {zip_files}}) { + // TODO: Add human-readable sizes? + // TODO: Add sorting options? + // TODO: Add hierarchical collapsible tree? + return html` + <${Hider} name="Zip Contents" shown=false> + + + + + + + + + + + ${zip_files.map(zf => html` + + + + + `)} + +
ModeSizeCompressedName
${{0: "store", 8: "deflate"}[zf.compression] || zf.compression}${zf.file_size}${zf.compressed_size}${zf.filename}
`; +} + +function CodeSection({model: {code_files}}) { + return html` + <${Hider} name="Code" shown=false> +
+ ${Object.entries(code_files).map(([fn, code]) => html`<${OneCodeSection} + filename=${fn} code=${code} />`)} +
`; +} + +class OneCodeSection extends Component { + constructor() { + super(); + this.state = { shown: false }; + } + + click() { + const shown = !this.state.shown; + this.setState({shown: shown}); + } + + render({filename, code}, {shown}) { + const header = html` +

+ this.click()} >${caret(shown)} + ${filename}

+ `; + if (!shown) { + return header; + } + return html` + ${header} + 
${code.map(c => this.renderBlock(c))}
+ `; + } + + renderBlock([text, ist_file, line, ist_s_text, s_start, s_end]) { + return html` blame.maybeBlame({ist_file, line, ist_s_text, s_start, s_end})} + >${text}`; + } +} + +function ExtraJsonSection({files}) { + return html` + <${Hider} name="Extra files (JSON)" shown=false> +
+

Use "Log Raw Model Info" for hierarchical view in browser console.

+ ${Object.entries(files).map(([fn, json]) => html`<${OneJsonSection} + filename=${fn} json=${json} />`)} +
`; +} + +class OneJsonSection extends Component { + constructor() { + super(); + this.state = { shown: false }; + } + + click() { + const shown = !this.state.shown; + this.setState({shown: shown}); + } + + render({filename, json}, {shown}) { + const header = html` +

+ this.click()} >${caret(shown)} + ${filename}

+ `; + if (!shown) { + return header; + } + return html` + ${header} + 
${JSON.stringify(json, null, 2)}
+ `; + } +} + +function ExtraPicklesSection({files}) { + return html` + <${Hider} name="Extra Pickles" shown=false> +
+ ${Object.entries(files).map(([fn, content]) => html`<${OnePickleSection} + filename=${fn} content=${content} />`)} +
`; +} + +class OnePickleSection extends Component { + constructor() { + super(); + this.state = { shown: false }; + } + + click() { + const shown = !this.state.shown; + this.setState({shown: shown}); + } + + render({filename, content}, {shown}) { + const header = html` +

+ this.click()} >${caret(shown)} + ${filename}

+ `; + if (!shown) { + return header; + } + return html` + ${header} + 
${content}
+ `; + } +} + +function assertStorageAreEqual(key, lhs, rhs) { + if (lhs.length !== rhs.length || + !lhs.every((val, idx) => val === rhs[idx])) { + throw new Error("Storage mismatch for key '" + key + "'"); + } +} + +function computeTensorMemory(numel, dtype) { + const sizes = { + "Byte": 1, + "Char": 1, + "Short": 2, + "Int": 4, + "Long": 8, + "Half": 2, + "Float": 4, + "Double": 8, + "ComplexHalf": 4, + "ComplexFloat": 8, + "ComplexDouble": 16, + "Bool": 1, + "QInt8": 1, + "QUInt8": 1, + "QInt32": 4, + "BFloat16": 2, + }; + let dtsize = sizes[dtype]; + if (!dtsize) { + throw new Error("Unrecognized dtype: " + dtype); + } + return numel * dtsize; +} + +// TODO: Maybe track by dtype as well. +// TODO: Maybe distinguish between visible size and storage size. +function getTensorStorages(data) { + if (data === null) { + return new Map(); + } + if (typeof(data) == "boolean") { + return new Map(); + } + if (typeof(data) == "number") { + return new Map(); + } + if (typeof(data) == "string") { + return new Map(); + } + if (typeof(data) != "object") { + throw new Error("Not an object"); + } + if (Array.isArray(data)) { + let result = new Map(); + for (const item of data) { + const tensors = getTensorStorages(item); + for (const [key, storage] of tensors.entries()) { + if (!result.has(key)) { + result.set(key, storage); + } else { + const old_storage = result.get(key); + assertStorageAreEqual(key, old_storage, storage); + } + } + } + return result; + } + if (data.__tuple_values__) { + return getTensorStorages(data.__tuple_values__); + } + if (data.__is_dict__) { + return getTensorStorages(data.values); + } + if (data.__module_type__) { + return getTensorStorages(data.state); + } + if (data.__tensor_v2__) { + const [storage, offset, size, stride, grad] = data.__tensor_v2__; + const [dtype, key, device, numel] = storage; + return new Map([[key, storage]]); + } + if (data.__qtensor__) { + const [storage, offset, size, stride, quantizer, grad] = data.__qtensor__; + const [dtype, key, device, numel] = storage; + return new Map([[key, storage]]); + } + throw new Error("Can't handle data type.", data); +} + +function getTensorMemoryByDevice(pickles) { + let all_tensors = []; + for (const [name, pickle] of pickles) { + const tensors = getTensorStorages(pickle); + all_tensors.push(...tensors.values()); + } + let result = {}; + for (const storage of all_tensors.values()) { + const [dtype, key, device, numel] = storage; + const size = computeTensorMemory(numel, dtype); + result[device] = (result[device] || 0) + size; + } + return result; +} + +// Make this a separate component so it is rendered lazily. +class OpenTensorMemorySection extends Component { + render({model: {model_data, constants}}) { + let sizes = getTensorMemoryByDevice(new Map([ + ["data", model_data], + ["constants", constants], + ])); + return html` + + + + + + + + + + ${Object.entries(sizes).map(([dev, size]) => html` + + + + `)} + +
DeviceBytesHuman
${dev}${size}${humanFileSize(size)}
`; + } +} + +function TensorMemorySection({model}) { + return html` + <${Hider} name="Tensor Memory" shown=false> + <${OpenTensorMemorySection} model=${model} />`; +} + +class AuxContentPane extends Component { + constructor() { + super(); + this.state = { + blame_info: null, + }; + } + + doBlame(arg) { + this.setState({...this.state, blame_info: arg}); + } + + render({model: {interned_strings}}, {blame_info}) { + let blame_content = ""; + if (blame_info) { + const {ist_file, line, ist_s_text, s_start, s_end} = blame_info; + let s_text = interned_strings[ist_s_text]; + if (s_start != 0 || s_end != s_text.length) { + let prefix = s_text.slice(0, s_start); + let main = s_text.slice(s_start, s_end); + let suffix = s_text.slice(s_end); + s_text = html`${prefix}${main}${suffix}`; + } + blame_content = html` +

${interned_strings[ist_file]}:${line}

+ 
${s_start}:${s_end}
+ 
${s_text}

+ `; + } + return html` + +
+ ${blame_content} + `; + } +} + +class App extends Component { + constructor() { + super(); + this.state = { + err: false, + model: null, + }; + } + + componentDidMount() { + const app = this; + if (BURNED_IN_MODEL_INFO !== null) { + app.setState({model: BURNED_IN_MODEL_INFO}); + } else { + fetch("./model_info.json").then(function(response) { + if (!response.ok) { + throw new Error("Response not ok."); + } + return response.json(); + }).then(function(body) { + app.setState({model: body}); + }).catch(function(error) { + console.log("Top-level error: ", error); + }); + } + } + + componentDidCatch(error) { + void(error); + this.setState({...this.state, err: true}); + } + + render(_, {err}) { + if (this.state.model === null) { + return html`

Loading...

`; + } + + const model = this.state.model.model; + + let error_msg = ""; + if (err) { + error_msg = html`

An error occurred. Check console

`; + } + + return html` + ${error_msg} +
+

TorchScript Model (version ${model.version}): ${model.title}

+ + <${ModelSizeSection} model=${model}/> + <${StructuredDataSection} name="Model Data" data=${model.model_data} shown=true/> + <${StructuredDataSection} name="Constants" data=${model.constants} shown=false/> + <${ZipContentsSection} model=${model}/> + <${CodeSection} model=${model}/> + <${ExtraJsonSection} files=${model.extra_files_jsons}/> + <${ExtraPicklesSection} files=${model.extra_pickles}/> + <${TensorMemorySection} model=${model}/> +
+
+ <${AuxContentPane} + err=${this.state.error} + model=${model} + ref=${(p) => blame.setAuxContentPane(p)}/> +
+ `; + } +} + +render(h(App), document.body); diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/htm.mjs b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/htm.mjs new file mode 100644 index 0000000000000000000000000000000000000000..06f25a13d8021ff4f43de442bbf0279f24735d6c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/htm.mjs @@ -0,0 +1,2 @@ +// HTM, Apache License +var n=function(t,s,r,e){var u;s[0]=0;for(var h=1;h=5&&((e||!n&&5===r)&&(h.push(r,0,e,s),r=6),n&&(h.push(r,n,0,s),r=6)),e=""},a=0;a"===t?(r=1,e=""):e=t+e[0]:u?t===u?u="":e+=t:'"'===t||"'"===t?u=t:">"===t?(p(),r=1):r&&("="===t?(r=5,s=e,e=""):"/"===t&&(r<5||">"===n[a][l+1])?(p(),3===r&&(h=h[0]),r=h,(h=h[0]).push(2,0,r),r=0):" "===t||"\t"===t||"\n"===t||"\r"===t?(p(),r=2):e+=t),3===r&&"!--"===e&&(r=4,h=h[0])}return p(),h}(s)),r),arguments,[])).length>1?r:r[0]} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/preact.mjs b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/preact.mjs new file mode 100644 index 0000000000000000000000000000000000000000..8c85bd948c6772ca8d40fc8d6fab6a220d55a1ef --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_dump/preact.mjs @@ -0,0 +1,2 @@ +// Preact, MIT License +var n,l,u,i,t,o,r={},f=[],e=/acit|ex(?:s|g|n|p|$)|rph|grid|ows|mnc|ntw|ine[ch]|zoo|^ord|itera/i;function c(e,n){for(var t in n)e[t]=n[t];return e}function s(e){var n=e.parentNode;n&&n.removeChild(e)}function a(e,n,t){var _,l,o,r=arguments,i={};for(o in n)"key"==o?_=n[o]:"ref"==o?l=n[o]:i[o]=n[o];if(arguments.length>3)for(t=[t],o=3;o0?v(m.type,m.props,m.key,null,m.__v):m)){if(m.__=t,m.__b=t.__b+1,null===(h=P[p])||h&&m.key==h.key&&m.type===h.type)P[p]=void 0;else for(a=0;a3)for(t=[t],o=3;o + + + TorchScript Model + + + + + + + + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_zoo.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_zoo.py new file mode 100644 index 0000000000000000000000000000000000000000..e0c6004e23ea806a2c83e12cd2998e0279e0b16f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/model_zoo.py @@ -0,0 +1,2 @@ +# torchvision imports tqdm from here. +from torch.hub import tqdm, load_state_dict_from_url as load_url # noqa: F401 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/module_tracker.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/module_tracker.py new file mode 100644 index 0000000000000000000000000000000000000000..7b5a8aad4dda9880c860850e42071a55ee7442e5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/module_tracker.py @@ -0,0 +1,160 @@ +# mypy: allow-untyped-defs +import logging +import weakref +from typing import TYPE_CHECKING + +import torch +from torch.autograd.graph import register_multi_grad_hook +from torch.nn.modules.module import ( + register_module_forward_hook, + register_module_forward_pre_hook, +) +from torch.utils._pytree import tree_flatten + + +if TYPE_CHECKING: + from torch.utils.hooks import RemovableHandle + + +logger = logging.getLogger(__name__) + + +__all__ = ["ModuleTracker"] + + +class ModuleTracker: + """ + ``ModuleTracker`` is a context manager that tracks the nn.Module hierarchy during execution + so that other system can query which Module is currently being executed (or its backward is being + executed). + + You can access the ``parents`` attribute on this context manager to get the set of all the + Modules currently being executed via their fqn (fully qualified name, also used as the key within + the state_dict). + You can access the ``is_bw`` attribute to know if you are currently running in backward or not. + + Note that ``parents`` is never empty and always contains the "Global" key. The ``is_bw`` flag + will remain ``True`` after the forward until another Module is executed. If you need it to be + more accurate, please submit an issue requesting this. Adding a map from fqn to the module instance + is possible but not done yet, please submit an issue requesting this if you need it. + + Example usage + + .. code-block:: python + + mod = torch.nn.Linear(2, 2) + + with ModuleTracker() as tracker: + # Access anything during the forward pass + def my_linear(m1, m2, bias): + print(f"Current modules: {tracker.parents}") + return torch.mm(m1, m2.t()) + bias + + torch.nn.functional.linear = my_linear + + mod(torch.rand(2, 2)) + + """ + + parents: set[str] + """ + A Set containing the fqn for each module currently running their forward + """ + + def __init__(self) -> None: + self.parents = {"Global"} + self._known_modules: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + self._seen_modules: weakref.WeakSet = weakref.WeakSet() + self._has_callback = False + self._hooks: list[RemovableHandle] = [] + + def _maybe_set_engine_callback(self) -> None: + # This assumes no concurrent calls to backward + if self._has_callback: + return + + def callback() -> None: + self.parents = {"Global"} + self._has_callback = False + + torch.autograd.Variable._execution_engine.queue_callback(callback) + self._has_callback = True + + @property + def is_bw(self): + """ + A boolean marking if this is currently running during the backward pass or not + """ + return torch._C._current_graph_task_id() != -1 + + def _get_mod_name(self, mod): + if mod not in self._known_modules: + self._known_modules[mod] = type(mod).__name__ + mod_name = self._known_modules[mod] + if mod not in self._seen_modules: + for name, submod in mod.named_children(): + self._known_modules[submod] = f"{mod_name}.{name}" + self._get_mod_name(submod) + self._seen_modules.add(mod) + return mod_name + + def _get_append_fn(self, name, is_bw): + def fn(*args) -> None: + if is_bw: + self._maybe_set_engine_callback() + if name in self.parents: + logger.info( + "The module hierarchy tracking seems to be broken as this Module was already entered. %s during %s", + name, + "backward" if is_bw else "forward", + ) + self.parents.add(name) + + return fn + + def _get_pop_fn(self, name, is_bw): + def fn(*args) -> None: + if name in self.parents: + self.parents.remove(name) + else: + logger.info( + "The Module hierarchy tracking is confused as we're exiting a Module that was never entered. %s during %s", + name, + "backward" if is_bw else "forward", + ) + + return fn + + def _fw_pre_hook(self, mod, input) -> None: + name = self._get_mod_name(mod) + self._get_append_fn(name, False)() + + args, _ = tree_flatten(input) + tensors = [a for a in args if isinstance(a, torch.Tensor) and a.requires_grad] + if tensors: + self._hooks.append( + register_multi_grad_hook(tensors, self._get_pop_fn(name, True)) + ) + + def _fw_post_hook(self, mod, input, output) -> None: + name = self._get_mod_name(mod) + self._get_pop_fn(name, False)() + + args, _ = tree_flatten(output) + tensors = [a for a in args if isinstance(a, torch.Tensor) and a.requires_grad] + if tensors: + self._hooks.append( + register_multi_grad_hook(tensors, self._get_append_fn(name, True)) + ) + + def __enter__(self): + self._fw_pre_handle = register_module_forward_pre_hook(self._fw_pre_hook) + self._fw_post_handle = register_module_forward_hook(self._fw_post_hook) + return self + + def __exit__(self, *args): + self._fw_pre_handle.remove() + self._fw_post_handle.remove() + for hook in self._hooks: + hook.remove() + self._hooks.clear() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/serialization/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/serialization/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d63bc18b69b138a026622de599aed656cc868c8e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/serialization/__init__.py @@ -0,0 +1 @@ +from . import config diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/serialization/config.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/serialization/config.py new file mode 100644 index 0000000000000000000000000000000000000000..c3e6729c68583f7206d07df7bfa2666007a6bd67 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/serialization/config.py @@ -0,0 +1,25 @@ +import sys +from typing import Optional as _Optional, TYPE_CHECKING as _TYPE_CHECKING + + +if _TYPE_CHECKING: + from torch.serialization import LoadEndianness as _LoadEndianess + +from torch.utils._config_module import install_config_module as _install_config_module + + +class load: + mmap: bool = False + endianness: _Optional["_LoadEndianess"] = None + # MAP_PRIVATE = 2 + mmap_flags: int | None = None if sys.platform == "win32" else 2 + calculate_storage_offsets: bool = False + + +class save: + compute_crc32: bool = True + use_pinned_memory_for_d2h: bool = False + storage_alignment: int = 64 + + +_install_config_module(sys.modules[__name__]) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/show_pickle.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/show_pickle.py new file mode 100644 index 0000000000000000000000000000000000000000..269ba3fbda4230c71ff14fe0e97872f6a8c57e6d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/show_pickle.py @@ -0,0 +1,151 @@ +#!/usr/bin/env python3 +# mypy: allow-untyped-defs +import sys +import pickle +import struct +import pprint +import zipfile +import fnmatch +from typing import Any, IO + +__all__ = ["FakeObject", "FakeClass", "DumpUnpickler", "main"] + +class FakeObject: + def __init__(self, module, name, args) -> None: + self.module = module + self.name = name + self.args = args + # NOTE: We don't distinguish between state never set and state set to None. + self.state = None + + def __repr__(self) -> str: + state_str = "" if self.state is None else f"(state={self.state!r})" + return f"{self.module}.{self.name}{self.args!r}{state_str}" + + def __setstate__(self, state): + self.state = state + + @staticmethod + def pp_format(printer, obj, stream, indent, allowance, context, level) -> None: + if not obj.args and obj.state is None: + stream.write(repr(obj)) + return + if obj.state is None: + stream.write(f"{obj.module}.{obj.name}") + printer._format(obj.args, stream, indent + 1, allowance + 1, context, level) + return + if not obj.args: + stream.write(f"{obj.module}.{obj.name}()(state=\n") + indent += printer._indent_per_level + stream.write(" " * indent) + printer._format(obj.state, stream, indent, allowance + 1, context, level + 1) + stream.write(")") + return + raise Exception("Need to implement") # noqa: TRY002 + + +class FakeClass: + def __init__(self, module, name) -> None: + self.module = module + self.name = name + self.__new__ = self.fake_new # type: ignore[assignment] + + def __repr__(self) -> str: + return f"{self.module}.{self.name}" + + def __call__(self, *args): + return FakeObject(self.module, self.name, args) + + def fake_new(self, *args): + return FakeObject(self.module, self.name, args[1:]) + + +class DumpUnpickler(pickle._Unpickler): # type: ignore[name-defined] + def __init__( + self, + file, + *, + catch_invalid_utf8=False, + **kwargs) -> None: + super().__init__(file, **kwargs) + self.catch_invalid_utf8 = catch_invalid_utf8 + + def find_class(self, module, name): + return FakeClass(module, name) + + def persistent_load(self, pid): + return FakeObject("pers", "obj", (pid,)) + + dispatch = dict(pickle._Unpickler.dispatch) # type: ignore[attr-defined] + + # Custom objects in TorchScript are able to return invalid UTF-8 strings + # from their pickle (__getstate__) functions. Install a custom loader + # for strings that catches the decode exception and replaces it with + # a sentinel object. + def load_binunicode(self) -> None: + strlen, = struct.unpack(" sys.maxsize: + raise Exception("String too long.") # noqa: TRY002 + str_bytes = self.read(strlen) # type: ignore[attr-defined] + obj: Any + try: + obj = str(str_bytes, "utf-8", "surrogatepass") + except UnicodeDecodeError as exn: + if not self.catch_invalid_utf8: + raise + obj = FakeObject("builtin", "UnicodeDecodeError", (str(exn),)) + self.append(obj) # type: ignore[attr-defined] + dispatch[pickle.BINUNICODE[0]] = load_binunicode # type: ignore[assignment] + + @classmethod + def dump(cls, in_stream, out_stream): + value = cls(in_stream).load() + pprint.pprint(value, stream=out_stream) + return value + + +def main(argv, output_stream=None) -> int | None: + if len(argv) != 2: + # Don't spam stderr if not using stdout. + if output_stream is not None: + raise Exception("Pass argv of length 2.") # noqa: TRY002 + sys.stderr.write("usage: show_pickle PICKLE_FILE\n") + sys.stderr.write(" PICKLE_FILE can be any of:\n") + sys.stderr.write(" path to a pickle file\n") + sys.stderr.write(" file.zip@member.pkl\n") + sys.stderr.write(" file.zip@*/pattern.*\n") + sys.stderr.write(" (shell glob pattern for members)\n") + sys.stderr.write(" (only first match will be shown)\n") + return 2 + + fname = argv[1] + handle: IO[bytes] + if "@" not in fname: + with open(fname, "rb") as handle: + DumpUnpickler.dump(handle, output_stream) + else: + zfname, mname = fname.split("@", 1) + with zipfile.ZipFile(zfname) as zf: + if "*" not in mname: + with zf.open(mname) as handle: + DumpUnpickler.dump(handle, output_stream) + else: + found = False + for info in zf.infolist(): + if fnmatch.fnmatch(info.filename, mname): + with zf.open(info) as handle: + DumpUnpickler.dump(handle, output_stream) + found = True + break + if not found: + raise Exception(f"Could not find member matching {mname} in {zfname}") # noqa: TRY002 + + +if __name__ == "__main__": + # This hack works on every version of Python I've tested. + # I've tested on the following versions: + # 3.7.4 + if True: + pprint.PrettyPrinter._dispatch[FakeObject.__repr__] = FakeObject.pp_format # type: ignore[attr-defined] + + sys.exit(main(sys.argv)) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a9b2ac5edd05e16ef51e75f2ca68864b65da5d58 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/__init__.py @@ -0,0 +1,19 @@ +import tensorboard +from torch._vendor.packaging.version import Version + +if not hasattr(tensorboard, "__version__") or Version( + tensorboard.__version__ +) < Version("1.15"): + raise ImportError("TensorBoard logging requires TensorBoard version 1.15 or above") + +del Version +del tensorboard + +from .writer import FileWriter, SummaryWriter +from tensorboard.summary.writer.record_writer import RecordWriter + +__all__ = [ + "FileWriter", + "RecordWriter", + "SummaryWriter", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_convert_np.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_convert_np.py new file mode 100644 index 0000000000000000000000000000000000000000..f0e8910580de16d9e7cf90f10d6327556e9a37a6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_convert_np.py @@ -0,0 +1,37 @@ +"""This module converts objects into numpy array.""" + +import numpy as np + +import torch + + +def make_np(x: torch.Tensor) -> np.ndarray: + """ + Convert an object into numpy array. + + Args: + x: An instance of torch tensor + + Returns: + numpy.array: Numpy array + """ + if isinstance(x, np.ndarray): + return x + if np.isscalar(x): + return np.array([x]) + if isinstance(x, torch.Tensor): + if x.device.type == "meta": + return np.random.randn(1) + return _prepare_pytorch(x) + raise NotImplementedError( + f"Got {type(x)}, but numpy array or torch tensor are expected." + ) + + +def _prepare_pytorch(x: torch.Tensor) -> np.ndarray: + if x.dtype == torch.bfloat16: + x = x.to(torch.float16) + # pyrefly: ignore [bad-assignment] + x = x.detach().cpu().numpy() + # pyrefly: ignore [bad-return] + return x diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_embedding.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..73413e219d0efbabe7d66747bd108ee5e8be4319 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_embedding.py @@ -0,0 +1,87 @@ +# mypy: allow-untyped-defs +import math +import numpy as np +from ._convert_np import make_np +from ._utils import make_grid +from tensorboard.compat import tf +from tensorboard.plugins.projector.projector_config_pb2 import EmbeddingInfo + + +_HAS_GFILE_JOIN = hasattr(tf.io.gfile, "join") + + +def _gfile_join(a, b): + # The join API is different between tensorboard's TF stub and TF: + # https://github.com/tensorflow/tensorboard/issues/6080 + # We need to try both because `tf` may point to either the stub or the real TF. + if _HAS_GFILE_JOIN: + return tf.io.gfile.join(a, b) + else: + fs = tf.io.gfile.get_filesystem(a) + return fs.join(a, b) + + +def make_tsv(metadata, save_path, metadata_header=None) -> None: + if not metadata_header: + metadata = [str(x) for x in metadata] + else: + if len(metadata_header) != len( + metadata[0] + ): + raise AssertionError("len of header must be equal to the number of columns in metadata") + metadata = ["\t".join(str(e) for e in l) for l in [metadata_header] + metadata] + + metadata_bytes = tf.compat.as_bytes("\n".join(metadata) + "\n") + with tf.io.gfile.GFile(_gfile_join(save_path, "metadata.tsv"), "wb") as f: + f.write(metadata_bytes) + + +# https://github.com/tensorflow/tensorboard/issues/44 image label will be squared +def make_sprite(label_img, save_path) -> None: + from PIL import Image + from io import BytesIO + + # this ensures the sprite image has correct dimension as described in + # https://www.tensorflow.org/get_started/embedding_viz + nrow = math.ceil((label_img.size(0)) ** 0.5) + arranged_img_CHW = make_grid(make_np(label_img), ncols=nrow) + + # augment images so that #images equals nrow*nrow + arranged_augment_square_HWC = np.zeros( + (arranged_img_CHW.shape[2], arranged_img_CHW.shape[2], 3) + ) + arranged_img_HWC = arranged_img_CHW.transpose(1, 2, 0) # chw -> hwc + arranged_augment_square_HWC[: arranged_img_HWC.shape[0], :, :] = arranged_img_HWC + im = Image.fromarray(np.uint8((arranged_augment_square_HWC * 255).clip(0, 255))) + + with BytesIO() as buf: + im.save(buf, format="PNG") + im_bytes = buf.getvalue() + + with tf.io.gfile.GFile(_gfile_join(save_path, "sprite.png"), "wb") as f: + f.write(im_bytes) + + +def get_embedding_info(metadata, label_img, subdir, global_step, tag): + info = EmbeddingInfo() + info.tensor_name = f"{tag}:{str(global_step).zfill(5)}" + info.tensor_path = _gfile_join(subdir, "tensors.tsv") + if metadata is not None: + info.metadata_path = _gfile_join(subdir, "metadata.tsv") + if label_img is not None: + info.sprite.image_path = _gfile_join(subdir, "sprite.png") + info.sprite.single_image_dim.extend([label_img.size(3), label_img.size(2)]) + return info + + +def write_pbtxt(save_path, contents) -> None: + config_path = _gfile_join(save_path, "projector_config.pbtxt") + with tf.io.gfile.GFile(config_path, "wb") as f: + f.write(tf.compat.as_bytes(contents)) + + +def make_mat(matlist, save_path) -> None: + with tf.io.gfile.GFile(_gfile_join(save_path, "tensors.tsv"), "wb") as f: + for x in matlist: + x = [str(i.item()) for i in x] + f.write(tf.compat.as_bytes("\t".join(x) + "\n")) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_onnx_graph.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_onnx_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..59d1d5b5f0a5cc0bf1d3694184b9727661d06093 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_onnx_graph.py @@ -0,0 +1,60 @@ +# mypy: allow-untyped-defs +from tensorboard.compat.proto.graph_pb2 import GraphDef +from tensorboard.compat.proto.node_def_pb2 import NodeDef +from tensorboard.compat.proto.versions_pb2 import VersionDef +from tensorboard.compat.proto.attr_value_pb2 import AttrValue +from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto + + +def load_onnx_graph(fname): + import onnx + m = onnx.load(fname) # type: ignore[attr-defined] + g = m.graph + return parse(g) + + +def parse(graph): + nodes = [] + import itertools + + nodes_proto = list(itertools.chain(graph.input, graph.output)) + + for node in nodes_proto: + print(node.name) + shapeproto = TensorShapeProto( + dim=[ + TensorShapeProto.Dim(size=d.dim_value) + for d in node.type.tensor_type.shape.dim + ] + ) + nodes.append( + NodeDef( + name=node.name.encode(encoding="utf_8"), + op="Variable", + input=[], + attr={ + "dtype": AttrValue(type=node.type.tensor_type.elem_type), + "shape": AttrValue(shape=shapeproto), + }, + ) + ) + + for node in graph.node: + _attr = [" = ".join([str(f[1]) for f in s.ListFields()]) for s in node.attribute] + attr = ", ".join(_attr).encode(encoding="utf_8") + print(node.output[0]) + nodes.append( + NodeDef( + name=node.output[0].encode(encoding="utf_8"), + op=node.op_type, + input=node.input, + attr={"parameters": AttrValue(s=attr)}, + ) + ) + + # two pass token replacement, appends opname to object id + mapping = {} + for node in nodes: + mapping[node.name] = node.op + "_" + node.name + + return GraphDef(node=nodes, versions=VersionDef(producer=22)) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_proto_graph.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_proto_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..8a9dc8536755c173f64bd5186af680d8438067ab --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_proto_graph.py @@ -0,0 +1,55 @@ +import torch + +from collections.abc import Sequence +from tensorboard.compat.proto.node_def_pb2 import NodeDef +from tensorboard.compat.proto.attr_value_pb2 import AttrValue +from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto + + +def attr_value_proto(dtype: object, shape: Sequence[int] | None, s: str | None) -> dict[str, AttrValue]: + """Create a dict of objects matching a NodeDef's attr field. + + Follows https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/attr_value.proto + specifically designed for a NodeDef. The values have been reverse engineered from + standard TensorBoard logged data. + """ + attr = {} + if s is not None: + attr["attr"] = AttrValue(s=s.encode(encoding="utf_8")) + if shape is not None: + shapeproto = tensor_shape_proto(shape) + attr["_output_shapes"] = AttrValue(list=AttrValue.ListValue(shape=[shapeproto])) + return attr + + +def tensor_shape_proto(outputsize: Sequence[int]) -> TensorShapeProto: + """Create an object matching a tensor_shape field. + + Follows https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/tensor_shape.proto . + """ + return TensorShapeProto(dim=[TensorShapeProto.Dim(size=d) for d in outputsize]) + + +def node_proto( + name: str, + op: str = "UnSpecified", + input: list[str] | str | None = None, + dtype: torch.dtype | None = None, + shape: tuple[int, ...] | None = None, + outputsize: Sequence[int] | None = None, + attributes: str = "", +) -> NodeDef: + """Create an object matching a NodeDef. + + Follows https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/node_def.proto . + """ + if input is None: + input = [] + if not isinstance(input, list): + input = [input] + return NodeDef( + name=name.encode(encoding="utf_8"), + op=op, + input=input, + attr=attr_value_proto(dtype, outputsize, attributes), + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_pytorch_graph.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_pytorch_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..e8f9fa11a8e7aca80af73a356768e8f1e71610fb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_pytorch_graph.py @@ -0,0 +1,378 @@ +# mypy: allow-untyped-defs +from collections import OrderedDict +import contextlib +from typing import Any + +from tensorboard.compat.proto.config_pb2 import RunMetadata +from tensorboard.compat.proto.graph_pb2 import GraphDef +from tensorboard.compat.proto.step_stats_pb2 import StepStats, DeviceStepStats +from tensorboard.compat.proto.versions_pb2 import VersionDef + +import torch +from ._proto_graph import node_proto + +methods_OP = [ + "attributeNames", + "hasMultipleOutputs", + "hasUses", + "inputs", + "kind", + "outputs", + "outputsSize", + "scopeName", +] +# Some additional methods to explure for methods_IO are +# +# 'unique' (type int) +# 'type' (type >) +# +# But the below are sufficient for now. +methods_IO = ["node", "offset", "debugName"] + +GETATTR_KIND = "prim::GetAttr" +CLASSTYPE_KIND = "ClassType" + + +class NodeBase: + def __init__( + self, + debugName=None, + inputs=None, + scope=None, + tensor_size=None, + op_type="UnSpecified", + attributes="", + ) -> None: + # TODO; Specify a __slots__ for this class or potentially + # used namedtuple instead + self.debugName = debugName + self.inputs = inputs + self.tensor_size = tensor_size + self.kind = op_type + self.attributes = attributes + self.scope = scope + + def __repr__(self) -> str: + repr = [] + repr.append(str(type(self))) + repr.extend( + m + ": " + str(getattr(self, m)) + str(type(getattr(self, m))) + for m in dir(self) + if "__" not in m + ) + return "\n".join(repr) + "\n\n" + + +class NodePy(NodeBase): + def __init__(self, node_cpp, valid_methods) -> None: + super().__init__(node_cpp) + valid_methods = valid_methods[:] + self.inputs = [] + + for m in valid_methods: + if m == "inputs" or m == "outputs": + list_of_node = list(getattr(node_cpp, m)()) + io_unique_names = [] + io_tensor_sizes = [] + for n in list_of_node: + io_unique_names.append(n.debugName()) + if n.isCompleteTensor(): + io_tensor_sizes.append(n.type().sizes()) + else: + io_tensor_sizes.append(None) + + setattr(self, m, io_unique_names) + setattr(self, m + "tensor_size", io_tensor_sizes) + + else: + setattr(self, m, getattr(node_cpp, m)()) + + +class NodePyIO(NodePy): + def __init__(self, node_cpp, input_or_output=None) -> None: + super().__init__(node_cpp, methods_IO) + try: + tensor_size = node_cpp.type().sizes() + except RuntimeError: + tensor_size = [ + 1, + ] # fail when constant model is used. + self.tensor_size = tensor_size + # Kind attribute string is purely descriptive and will be shown + # in detailed information for the node in TensorBoard's graph plugin. + # + # NodePyOP nodes get this from their kind() method. + self.kind = "Parameter" + if input_or_output: + self.input_or_output = input_or_output + self.kind = "IO Node" + + +class NodePyOP(NodePy): + def __init__(self, node_cpp) -> None: + super().__init__(node_cpp, methods_OP) + # Replace single quote which causes strange behavior in TensorBoard + # TODO: See if we can remove this in the future + self.attributes = str( + {k: _node_get(node_cpp, k) for k in node_cpp.attributeNames()} + ).replace("'", " ") + self.kind = node_cpp.kind() + + +class GraphPy: + """Helper class to convert torch.nn.Module to GraphDef proto and visualization with TensorBoard. + + GraphDef generation operates in two passes: + + In the first pass, all nodes are read and saved to two lists. + One list is for input/output nodes (nodes_io), which only have inbound + or outbound connections, but not both. Another list is for internal + operator nodes (nodes_op). The first pass also saves all scope name + appeared in the nodes in scope_name_appeared list for later processing. + + In the second pass, scope names are fully applied to all nodes. + debugNameToScopedName is a mapping from a node's ID to its fully qualified + scope name. e.g. Net1/Linear[0]/1. Unfortunately torch.jit doesn't have + totally correct scope output, so this is nontrivial. The function + populate_namespace_from_OP_to_IO and find_common_root are used to + assign scope name to a node based on the connection between nodes + in a heuristic kind of way. Bookkeeping is done with shallowest_scope_name + and scope_name_appeared. + """ + + def __init__(self) -> None: + self.nodes_op = [] + self.nodes_io = OrderedDict() + self.unique_name_to_scoped_name = {} + self.shallowest_scope_name = "default" + self.scope_name_appeared = [] + + def append(self, x) -> None: + if isinstance(x, NodePyIO): + self.nodes_io[x.debugName] = x + if isinstance(x, NodePyOP): + self.nodes_op.append(x) + + def printall(self) -> None: + print("all nodes") + for node in self.nodes_op: + print(node) + for key in self.nodes_io: + print(self.nodes_io[key]) + + def find_common_root(self) -> None: + for fullscope in self.scope_name_appeared: + if fullscope: + self.shallowest_scope_name = fullscope.split("/")[0] + + def populate_namespace_from_OP_to_IO(self) -> None: + for node in self.nodes_op: + for node_output, outputSize in zip(node.outputs, node.outputstensor_size, strict=True): + self.scope_name_appeared.append(node.scopeName) + self.nodes_io[node_output] = NodeBase( + node_output, + node.inputs, + node.scopeName, + outputSize, + op_type=node.kind, + attributes=node.attributes, + ) + + self.find_common_root() + + for node in self.nodes_op: + for input_node_id in node.inputs: + self.unique_name_to_scoped_name[input_node_id] = ( + node.scopeName + "/" + input_node_id + ) + + for key, node in self.nodes_io.items(): + if type(node) is NodeBase: + # pyrefly: ignore [unsupported-operation] + self.unique_name_to_scoped_name[key] = node.scope + "/" + node.debugName + if hasattr(node, "input_or_output"): + self.unique_name_to_scoped_name[key] = ( + node.input_or_output + "/" + node.debugName + ) + + if hasattr(node, "scope") and node.scope is not None: + self.unique_name_to_scoped_name[key] = node.scope + "/" + node.debugName + if node.scope == "" and self.shallowest_scope_name: + self.unique_name_to_scoped_name[node.debugName] = ( + + self.shallowest_scope_name + "/" + node.debugName + ) + + # replace name + for key, node in self.nodes_io.items(): + self.nodes_io[key].inputs = [ + self.unique_name_to_scoped_name[node_input_id] + for node_input_id in node.inputs + ] + if node.debugName in self.unique_name_to_scoped_name: + self.nodes_io[key].debugName = self.unique_name_to_scoped_name[ + node.debugName + ] + + def to_proto(self): + """Convert graph representation of GraphPy object to TensorBoard required format.""" + # TODO: compute correct memory usage and CPU time once + # PyTorch supports it + nodes = [ + node_proto( + v.debugName, + input=v.inputs, + outputsize=v.tensor_size, + op=v.kind, + attributes=v.attributes, + ) + for v in self.nodes_io.values() + ] + return nodes + + +def parse(graph, trace, args=None, omit_useless_nodes=True): + """Parse an optimized PyTorch model graph and produces a list of nodes and node stats. + + Useful for eventual conversion to TensorBoard protobuf format. + + Args: + graph (PyTorch module): The model graph to be parsed. + trace (PyTorch JIT TracedModule): The model trace to be parsed. + args (tuple): input tensor[s] for the model. + omit_useless_nodes (boolean): Whether to remove nodes from the graph. + """ + nodes_py = GraphPy() + for node in graph.inputs(): + if omit_useless_nodes: + if ( + len(node.uses()) == 0 + ): # number of user of the node (= number of outputs/ fanout) + continue + + if node.type().kind() != CLASSTYPE_KIND: + nodes_py.append(NodePyIO(node, "input")) + + attr_to_scope: dict[Any, str] = {} + for node in graph.nodes(): + if node.kind() == GETATTR_KIND: + attr_name = node.s("name") + attr_key = node.output().debugName() + parent = node.input().node() + if ( + parent.kind() == GETATTR_KIND + ): # If the parent node is not the top-level "self" node + parent_attr_key = parent.output().debugName() + parent_scope = attr_to_scope[parent_attr_key] + attr_scope = parent_scope.split("/")[-1] + attr_to_scope[attr_key] = f"{parent_scope}/{attr_scope}.{attr_name}" + else: + attr_to_scope[attr_key] = f"__module.{attr_name}" + # We don't need classtype nodes; scope will provide this information + if node.output().type().kind() != CLASSTYPE_KIND: + node_py = NodePyOP(node) + node_py.scopeName = attr_to_scope[attr_key] # type: ignore[attr-defined] + nodes_py.append(node_py) + else: + nodes_py.append(NodePyOP(node)) + + for i, node in enumerate(graph.outputs()): # Create sink nodes for output ops + node_pyio = NodePyIO(node, "output") + node_pyio.debugName = f"output.{i + 1}" + node_pyio.inputs = [node.debugName()] + nodes_py.append(node_pyio) + + def parse_traced_name(module): + if isinstance(module, torch.jit.TracedModule): + module_name = module._name + else: + module_name = getattr(module, "original_name", "Module") + return module_name + + alias_to_name = {} + base_name = parse_traced_name(trace) + for name, module in trace.named_modules(prefix="__module"): + mod_name = parse_traced_name(module) + attr_name = name.split(".")[-1] + alias_to_name[name] = f"{mod_name}[{attr_name}]" + + for node in nodes_py.nodes_op: + module_aliases = node.scopeName.split("/") + replacements = [ + alias_to_name[alias] if alias in alias_to_name else alias.split(".")[-1] + for alias in module_aliases + ] + node.scopeName = base_name + if any(replacements): + node.scopeName += "/" + "/".join(replacements) + + nodes_py.populate_namespace_from_OP_to_IO() + return nodes_py.to_proto() + + +def graph(model, args, verbose=False, use_strict_trace=True): + """ + Process a PyTorch model and produces a `GraphDef` proto that can be logged to TensorBoard. + + Args: + model (PyTorch module): The model to be parsed. + args (tuple): input tensor[s] for the model. + verbose (bool): Whether to print out verbose information while + processing. + use_strict_trace (bool): Whether to pass keyword argument `strict` to + `torch.jit.trace`. Pass False when you want the tracer to + record your mutable container types (list, dict) + """ + with _set_model_to_eval(model): + try: + trace = torch.jit.trace(model, args, strict=use_strict_trace) + graph = trace.graph + torch._C._jit_pass_inline(graph) + except RuntimeError as e: + print(e) + print("Error occurs, No graph saved") + raise e + + if verbose: + print(graph) + list_of_nodes = parse(graph, trace, args) + # We are hardcoding that this was run on CPU even though it might have actually + # run on GPU. Note this is what is shown in TensorBoard and has no bearing + # on actual execution. + # TODO: See if we can extract GPU vs CPU information from the PyTorch model + # and pass it correctly to TensorBoard. + # + # Definition of StepStats and DeviceStepStats can be found at + # https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/graph/tf_graph_common/proto.ts + # and + # https://github.com/tensorflow/tensorboard/blob/master/tensorboard/compat/proto/step_stats.proto + stepstats = RunMetadata( + step_stats=StepStats(dev_stats=[DeviceStepStats(device="/device:CPU:0")]) + ) + return GraphDef(node=list_of_nodes, versions=VersionDef(producer=22)), stepstats + # The producer version has been reverse engineered from standard + # TensorBoard logged data. + + +@contextlib.contextmanager +def _set_model_to_eval(model): + """Context manager to temporarily set the training mode of ``model`` to eval.""" + if not isinstance(model, torch.jit.ScriptFunction): + originally_training = model.training + model.train(False) + try: + yield + finally: + model.train(originally_training) + else: + # Do nothing for ScriptFunction + try: + yield + finally: + pass + + +def _node_get(node: torch._C.Node, key: str): + """Get attributes of a node which is polymorphic over return type.""" + sel = node.kindOf(key) + return getattr(node, sel)(key) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..0d1a0289f4230c9b131733695da082651b09a1c9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/_utils.py @@ -0,0 +1,131 @@ +# mypy: allow-untyped-defs +import numpy as np +import numpy.typing as npt + + +# Functions for converting +def figure_to_image(figures, close=True): + """Render matplotlib figure to numpy format. + + Note that this requires the ``matplotlib`` package. + + Args: + figures (matplotlib.pyplot.figure or list of figures): figure or a list of figures + close (bool): Flag to automatically close the figure + + Returns: + numpy.array: image in [CHW] order + """ + import matplotlib.pyplot as plt + import matplotlib.backends.backend_agg as plt_backend_agg + + def render_to_rgb(figure): + canvas = plt_backend_agg.FigureCanvasAgg(figure) + canvas.draw() + data: npt.NDArray = np.frombuffer(canvas.buffer_rgba(), dtype=np.uint8) + w, h = figure.canvas.get_width_height() + image_hwc = data.reshape([h, w, 4])[:, :, 0:3] + image_chw = np.moveaxis(image_hwc, source=2, destination=0) + if close: + plt.close(figure) + return image_chw + + if isinstance(figures, list): + images = [render_to_rgb(figure) for figure in figures] + return np.stack(images) + else: + image = render_to_rgb(figures) + return image + + +def _prepare_video(V): + """ + Convert a 5D tensor into 4D tensor. + + Convesrion is done from [batchsize, time(frame), channel(color), height, width] (5D tensor) + to [time(frame), new_width, new_height, channel] (4D tensor). + + A batch of images are spread to a grid, which forms a frame. + e.g. Video with batchsize 16 will have a 4x4 grid. + """ + b, t, c, h, w = V.shape + + if V.dtype == np.uint8: + V = np.float32(V) / 255.0 + + def is_power2(num): + return num != 0 and ((num & (num - 1)) == 0) + + # pad to nearest power of 2, all at once + + if not is_power2(V.shape[0]): + + len_addition = int(2 ** V.shape[0].bit_length() - V.shape[0]) + V = np.concatenate((V, np.zeros(shape=(len_addition, t, c, h, w))), axis=0) + + n_rows = 2 ** ((b.bit_length() - 1) // 2) + + n_cols = V.shape[0] // n_rows + + V = np.reshape(V, (n_rows, n_cols, t, c, h, w)) + V = np.transpose(V, axes=(2, 0, 4, 1, 5, 3)) + V = np.reshape(V, (t, n_rows * h, n_cols * w, c)) + + return V + + +def make_grid(I, ncols=8): + # I: N1HW or N3HW + if not isinstance(I, np.ndarray): + raise AssertionError("plugin error, should pass numpy array here") + if I.shape[1] == 1: + I = np.concatenate([I, I, I], 1) + if I.ndim != 4 or I.shape[1] != 3: + raise AssertionError("Input should be a 4D numpy array with 3 channels") + nimg = I.shape[0] + H = I.shape[2] + W = I.shape[3] + ncols = min(nimg, ncols) + nrows = int(np.ceil(float(nimg) / ncols)) + canvas = np.zeros((3, H * nrows, W * ncols), dtype=I.dtype) + i = 0 + for y in range(nrows): + for x in range(ncols): + if i >= nimg: + break + canvas[:, y * H : (y + 1) * H, x * W : (x + 1) * W] = I[i] + i = i + 1 + return canvas + + # if modality == 'IMG': + # if x.dtype == np.uint8: + # x = x.astype(np.float32) / 255.0 + + +def convert_to_HWC(tensor, input_format): # tensor: numpy array + if len(set(input_format)) != len(input_format): + raise AssertionError(f"You can not use the same dimension shordhand twice. \ + input_format: {input_format}") + if len(tensor.shape) != len(input_format): + raise AssertionError(f"size of input tensor and input format are different. \ + tensor shape: {tensor.shape}, input_format: {input_format}") + input_format = input_format.upper() + + if len(input_format) == 4: + index = [input_format.find(c) for c in "NCHW"] + tensor_NCHW = tensor.transpose(index) + tensor_CHW = make_grid(tensor_NCHW) + return tensor_CHW.transpose(1, 2, 0) + + if len(input_format) == 3: + index = [input_format.find(c) for c in "HWC"] + tensor_HWC = tensor.transpose(index) + if tensor_HWC.shape[2] == 1: + tensor_HWC = np.concatenate([tensor_HWC, tensor_HWC, tensor_HWC], 2) + return tensor_HWC + + if len(input_format) == 2: + index = [input_format.find(c) for c in "HW"] + tensor = tensor.transpose(index) + tensor = np.stack([tensor, tensor, tensor], 2) + return tensor diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/summary.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/summary.py new file mode 100644 index 0000000000000000000000000000000000000000..f5d86875387bfe8032a5440208512fea801b79e4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/summary.py @@ -0,0 +1,986 @@ +# mypy: allow-untyped-defs +import json +import logging +import struct + +from typing import Any + +import torch +import numpy as np + + +from google.protobuf import struct_pb2 + +from tensorboard.compat.proto.summary_pb2 import ( + HistogramProto, + Summary, + SummaryMetadata, +) +from tensorboard.compat.proto.tensor_pb2 import TensorProto +from tensorboard.compat.proto.tensor_shape_pb2 import TensorShapeProto +from tensorboard.plugins.custom_scalar import layout_pb2 +from tensorboard.plugins.pr_curve.plugin_data_pb2 import PrCurvePluginData +from tensorboard.plugins.text.plugin_data_pb2 import TextPluginData + +from ._convert_np import make_np +from ._utils import _prepare_video, convert_to_HWC + +__all__ = [ + "half_to_int", + "int_to_half", + "hparams", + "scalar", + "histogram_raw", + "histogram", + "make_histogram", + "image", + "image_boxes", + "draw_boxes", + "make_image", + "video", + "make_video", + "audio", + "custom_scalars", + "text", + "tensor_proto", + "pr_curve_raw", + "pr_curve", + "compute_curve", + "mesh", +] + +logger = logging.getLogger(__name__) + +def half_to_int(f: float) -> int: + """Casts a half-precision float value into an integer. + + Converts a half precision floating point value, such as `torch.half` or + `torch.bfloat16`, into an integer value which can be written into the + half_val field of a TensorProto for storage. + + To undo the effects of this conversion, use int_to_half(). + + """ + buf = struct.pack("f", f) + return struct.unpack("i", buf)[0] + +def int_to_half(i: int) -> float: + """Casts an integer value to a half-precision float. + + Converts an integer value obtained from half_to_int back into a floating + point value. + + """ + buf = struct.pack("i", i) + return struct.unpack("f", buf)[0] + +def _tensor_to_half_val(t: torch.Tensor) -> list[int]: + return [half_to_int(x) for x in t.flatten().tolist()] + +def _tensor_to_complex_val(t: torch.Tensor) -> list[float]: + return torch.view_as_real(t).flatten().tolist() + +def _tensor_to_list(t: torch.Tensor) -> list[Any]: + return t.flatten().tolist() + +# type maps: torch.Tensor type -> (protobuf type, protobuf val field) +_TENSOR_TYPE_MAP = { + torch.half: ("DT_HALF", "half_val", _tensor_to_half_val), + torch.float16: ("DT_HALF", "half_val", _tensor_to_half_val), + torch.bfloat16: ("DT_BFLOAT16", "half_val", _tensor_to_half_val), + torch.float32: ("DT_FLOAT", "float_val", _tensor_to_list), + torch.float: ("DT_FLOAT", "float_val", _tensor_to_list), + torch.float64: ("DT_DOUBLE", "double_val", _tensor_to_list), + torch.double: ("DT_DOUBLE", "double_val", _tensor_to_list), + torch.int8: ("DT_INT8", "int_val", _tensor_to_list), + torch.uint8: ("DT_UINT8", "int_val", _tensor_to_list), + torch.qint8: ("DT_UINT8", "int_val", _tensor_to_list), + torch.int16: ("DT_INT16", "int_val", _tensor_to_list), + torch.short: ("DT_INT16", "int_val", _tensor_to_list), + torch.int: ("DT_INT32", "int_val", _tensor_to_list), + torch.int32: ("DT_INT32", "int_val", _tensor_to_list), + torch.qint32: ("DT_INT32", "int_val", _tensor_to_list), + torch.int64: ("DT_INT64", "int64_val", _tensor_to_list), + torch.complex32: ("DT_COMPLEX32", "scomplex_val", _tensor_to_complex_val), + torch.chalf: ("DT_COMPLEX32", "scomplex_val", _tensor_to_complex_val), + torch.complex64: ("DT_COMPLEX64", "scomplex_val", _tensor_to_complex_val), + torch.cfloat: ("DT_COMPLEX64", "scomplex_val", _tensor_to_complex_val), + torch.bool: ("DT_BOOL", "bool_val", _tensor_to_list), + torch.complex128: ("DT_COMPLEX128", "dcomplex_val", _tensor_to_complex_val), + torch.cdouble: ("DT_COMPLEX128", "dcomplex_val", _tensor_to_complex_val), + torch.uint8: ("DT_UINT8", "uint32_val", _tensor_to_list), + torch.quint8: ("DT_UINT8", "uint32_val", _tensor_to_list), + torch.quint4x2: ("DT_UINT8", "uint32_val", _tensor_to_list), +} + + +def _calc_scale_factor(tensor) -> int: + converted = tensor.numpy() if not isinstance(tensor, np.ndarray) else tensor + return 1 if converted.dtype == np.uint8 else 255 + + +def _draw_single_box( + image, + xmin, + ymin, + xmax, + ymax, + display_str, + color="black", + color_text="black", + thickness=2, +): + from PIL import ImageDraw, ImageFont + + font = ImageFont.load_default() + draw = ImageDraw.Draw(image) + (left, right, top, bottom) = (xmin, xmax, ymin, ymax) + draw.line( + [(left, top), (left, bottom), (right, bottom), (right, top), (left, top)], + width=thickness, + fill=color, + ) + if display_str: + text_bottom = bottom + # Reverse list and print from bottom to top. + _left, _top, _right, _bottom = font.getbbox(display_str) + text_width, text_height = _right - _left, _bottom - _top + margin = np.ceil(0.05 * text_height) + draw.rectangle( + [ + (left, text_bottom - text_height - 2 * margin), + (left + text_width, text_bottom), + ], + fill=color, + ) + draw.text( + (left + margin, text_bottom - text_height - margin), + display_str, + fill=color_text, + font=font, + ) + return image + + +def hparams(hparam_dict=None, metric_dict=None, hparam_domain_discrete=None): + """Output three `Summary` protocol buffers needed by hparams plugin. + + `Experiment` keeps the metadata of an experiment, such as the name of the + hyperparameters and the name of the metrics. + `SessionStartInfo` keeps key-value pairs of the hyperparameters + `SessionEndInfo` describes status of the experiment e.g. STATUS_SUCCESS + + Args: + hparam_dict: A dictionary that contains names of the hyperparameters + and their values. + metric_dict: A dictionary that contains names of the metrics + and their values. + hparam_domain_discrete: (Optional[Dict[str, List[Any]]]) A dictionary that + contains names of the hyperparameters and all discrete values they can hold + + Returns: + The `Summary` protobufs for Experiment, SessionStartInfo and + SessionEndInfo + """ + import torch + from tensorboard.plugins.hparams.api_pb2 import ( + DataType, + Experiment, + HParamInfo, + MetricInfo, + MetricName, + Status, + ) + from tensorboard.plugins.hparams.metadata import ( + EXPERIMENT_TAG, + PLUGIN_DATA_VERSION, + PLUGIN_NAME, + SESSION_END_INFO_TAG, + SESSION_START_INFO_TAG, + ) + from tensorboard.plugins.hparams.plugin_data_pb2 import ( + HParamsPluginData, + SessionEndInfo, + SessionStartInfo, + ) + + # TODO: expose other parameters in the future. + # hp = HParamInfo(name='lr',display_name='learning rate', + # type=DataType.DATA_TYPE_FLOAT64, domain_interval=Interval(min_value=10, + # max_value=100)) + # mt = MetricInfo(name=MetricName(tag='accuracy'), display_name='accuracy', + # description='', dataset_type=DatasetType.DATASET_VALIDATION) + # exp = Experiment(name='123', description='456', time_created_secs=100.0, + # hparam_infos=[hp], metric_infos=[mt], user='tw') + + if not isinstance(hparam_dict, dict): + logger.warning("parameter: hparam_dict should be a dictionary, nothing logged.") + raise TypeError( + "parameter: hparam_dict should be a dictionary, nothing logged." + ) + if not isinstance(metric_dict, dict): + logger.warning("parameter: metric_dict should be a dictionary, nothing logged.") + raise TypeError( + "parameter: metric_dict should be a dictionary, nothing logged." + ) + + hparam_domain_discrete = hparam_domain_discrete or {} + if not isinstance(hparam_domain_discrete, dict): + raise TypeError( + "parameter: hparam_domain_discrete should be a dictionary, nothing logged." + ) + for k, v in hparam_domain_discrete.items(): + if ( + k not in hparam_dict + or not isinstance(v, list) + or not all(isinstance(d, type(hparam_dict[k])) for d in v) + ): + raise TypeError( + f"parameter: hparam_domain_discrete[{k}] should be a list of same type as hparam_dict[{k}]." + ) + hps = [] + + ssi = SessionStartInfo() + for k, v in hparam_dict.items(): + if v is None: + continue + if isinstance(v, (int, float)): + ssi.hparams[k].number_value = v + + if k in hparam_domain_discrete: + domain_discrete: struct_pb2.ListValue | None = struct_pb2.ListValue( + values=[ + struct_pb2.Value(number_value=d) + for d in hparam_domain_discrete[k] + ] + ) + else: + domain_discrete = None + + hps.append( + HParamInfo( + name=k, + type=DataType.Value("DATA_TYPE_FLOAT64"), + domain_discrete=domain_discrete, + ) + ) + continue + + if isinstance(v, str): + ssi.hparams[k].string_value = v + + if k in hparam_domain_discrete: + domain_discrete = struct_pb2.ListValue( + values=[ + struct_pb2.Value(string_value=d) + for d in hparam_domain_discrete[k] + ] + ) + else: + domain_discrete = None + + hps.append( + HParamInfo( + name=k, + type=DataType.Value("DATA_TYPE_STRING"), + domain_discrete=domain_discrete, + ) + ) + continue + + if isinstance(v, bool): + ssi.hparams[k].bool_value = v + + if k in hparam_domain_discrete: + domain_discrete = struct_pb2.ListValue( + values=[ + struct_pb2.Value(bool_value=d) + for d in hparam_domain_discrete[k] + ] + ) + else: + domain_discrete = None + + hps.append( + HParamInfo( + name=k, + type=DataType.Value("DATA_TYPE_BOOL"), + domain_discrete=domain_discrete, + ) + ) + continue + + if isinstance(v, torch.Tensor): + v = make_np(v)[0] + ssi.hparams[k].number_value = v + hps.append(HParamInfo(name=k, type=DataType.Value("DATA_TYPE_FLOAT64"))) + continue + raise ValueError( + "value should be one of int, float, str, bool, or torch.Tensor" + ) + + content = HParamsPluginData(session_start_info=ssi, version=PLUGIN_DATA_VERSION) + smd = SummaryMetadata( + plugin_data=SummaryMetadata.PluginData( + plugin_name=PLUGIN_NAME, content=content.SerializeToString() + ) + ) + ssi = Summary(value=[Summary.Value(tag=SESSION_START_INFO_TAG, metadata=smd)]) + + mts = [MetricInfo(name=MetricName(tag=k)) for k in metric_dict] + + exp = Experiment(hparam_infos=hps, metric_infos=mts) + + content = HParamsPluginData(experiment=exp, version=PLUGIN_DATA_VERSION) + smd = SummaryMetadata( + plugin_data=SummaryMetadata.PluginData( + plugin_name=PLUGIN_NAME, content=content.SerializeToString() + ) + ) + exp = Summary(value=[Summary.Value(tag=EXPERIMENT_TAG, metadata=smd)]) + + sei = SessionEndInfo(status=Status.Value("STATUS_SUCCESS")) + content = HParamsPluginData(session_end_info=sei, version=PLUGIN_DATA_VERSION) + smd = SummaryMetadata( + plugin_data=SummaryMetadata.PluginData( + plugin_name=PLUGIN_NAME, content=content.SerializeToString() + ) + ) + sei = Summary(value=[Summary.Value(tag=SESSION_END_INFO_TAG, metadata=smd)]) + + return exp, ssi, sei + + +def scalar(name, tensor, collections=None, new_style=False, double_precision=False): + """Output a `Summary` protocol buffer containing a single scalar value. + + The generated Summary has a Tensor.proto containing the input Tensor. + Args: + name: A name for the generated node. Will also serve as the series name in + TensorBoard. + tensor: A real numeric Tensor containing a single value. + collections: Optional list of graph collections keys. The new summary op is + added to these collections. Defaults to `[GraphKeys.SUMMARIES]`. + new_style: Whether to use new style (tensor field) or old style (simple_value + field). New style could lead to faster data loading. + Returns: + A scalar `Tensor` of type `string`. Which contains a `Summary` protobuf. + Raises: + ValueError: If tensor has the wrong shape or type. + """ + tensor = make_np(tensor).squeeze() + if tensor.ndim != 0: + raise AssertionError(f"Tensor should contain one element (0 dimensions). \ + Was given size: {tensor.size} and {tensor.ndim} dimensions.") + # python float is double precision in numpy + scalar = float(tensor) + if new_style: + tensor_proto = TensorProto(float_val=[scalar], dtype="DT_FLOAT") + if double_precision: + tensor_proto = TensorProto(double_val=[scalar], dtype="DT_DOUBLE") + + plugin_data = SummaryMetadata.PluginData(plugin_name="scalars") + smd = SummaryMetadata(plugin_data=plugin_data) + return Summary( + value=[ + Summary.Value( + tag=name, + tensor=tensor_proto, + metadata=smd, + ) + ] + ) + else: + return Summary(value=[Summary.Value(tag=name, simple_value=scalar)]) + + +def tensor_proto(tag, tensor): + """Outputs a `Summary` protocol buffer containing the full tensor. + The generated Summary has a Tensor.proto containing the input Tensor. + Args: + tag: A name for the generated node. Will also serve as the series name in + TensorBoard. + tensor: Tensor to be converted to protobuf + Returns: + A tensor protobuf in a `Summary` protobuf. + Raises: + ValueError: If tensor is too big to be converted to protobuf, or + tensor data type is not supported + """ + if tensor.numel() * tensor.itemsize >= (1 << 31): + raise ValueError( + "tensor is bigger than protocol buffer's hard limit of 2GB in size" + ) + + if tensor.dtype in _TENSOR_TYPE_MAP: + dtype, field_name, conversion_fn = _TENSOR_TYPE_MAP[tensor.dtype] + tensor_proto = TensorProto( + **{ + "dtype": dtype, + "tensor_shape": TensorShapeProto( + dim=[TensorShapeProto.Dim(size=x) for x in tensor.shape] + ), + field_name: conversion_fn(tensor), + }, + ) + else: + raise ValueError(f"{tag} has unsupported tensor dtype {tensor.dtype}") + + plugin_data = SummaryMetadata.PluginData(plugin_name="tensor") + smd = SummaryMetadata(plugin_data=plugin_data) + return Summary(value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor_proto)]) + + +def histogram_raw(name, min, max, num, sum, sum_squares, bucket_limits, bucket_counts): + # pylint: disable=line-too-long + """Output a `Summary` protocol buffer with a histogram. + + The generated + [`Summary`](https://www.tensorflow.org/code/tensorflow/core/framework/summary.proto) + has one summary value containing a histogram for `values`. + Args: + name: A name for the generated node. Will also serve as a series name in + TensorBoard. + min: A float or int min value + max: A float or int max value + num: Int number of values + sum: Float or int sum of all values + sum_squares: Float or int sum of squares for all values + bucket_limits: A numeric `Tensor` with upper value per bucket + bucket_counts: A numeric `Tensor` with number of values per bucket + Returns: + A scalar `Tensor` of type `string`. The serialized `Summary` protocol + buffer. + """ + hist = HistogramProto( + min=min, + max=max, + num=num, + sum=sum, + sum_squares=sum_squares, + bucket_limit=bucket_limits, + bucket=bucket_counts, + ) + return Summary(value=[Summary.Value(tag=name, histo=hist)]) + + +def histogram(name, values, bins, max_bins=None): + # pylint: disable=line-too-long + """Output a `Summary` protocol buffer with a histogram. + + The generated + [`Summary`](https://www.tensorflow.org/code/tensorflow/core/framework/summary.proto) + has one summary value containing a histogram for `values`. + This op reports an `InvalidArgument` error if any value is not finite. + Args: + name: A name for the generated node. Will also serve as a series name in + TensorBoard. + values: A real numeric `Tensor`. Any shape. Values to use to + build the histogram. + Returns: + A scalar `Tensor` of type `string`. The serialized `Summary` protocol + buffer. + """ + values = make_np(values) + hist = make_histogram(values.astype(float), bins, max_bins) + return Summary(value=[Summary.Value(tag=name, histo=hist)]) + + +def make_histogram(values, bins, max_bins=None): + """Convert values into a histogram proto using logic from histogram.cc.""" + if values.size == 0: + raise ValueError("The input has no element.") + values = values.reshape(-1) + counts, limits = np.histogram(values, bins=bins) + num_bins = len(counts) + if max_bins is not None and num_bins > max_bins: + subsampling = num_bins // max_bins + subsampling_remainder = num_bins % subsampling + if subsampling_remainder != 0: + # pyrefly: ignore [no-matching-overload] + counts = np.pad( + counts, + pad_width=[[0, subsampling - subsampling_remainder]], + mode="constant", + constant_values=0, + ) + counts = counts.reshape(-1, subsampling).sum(axis=-1) + new_limits = np.empty((counts.size + 1,), limits.dtype) + new_limits[:-1] = limits[:-1:subsampling] + new_limits[-1] = limits[-1] + limits = new_limits + + # Find the first and the last bin defining the support of the histogram: + + cum_counts = np.cumsum(np.greater(counts, 0)) + start, end = np.searchsorted(cum_counts, [0, cum_counts[-1] - 1], side="right") + start = int(start) + end = int(end) + 1 + del cum_counts + + # TensorBoard only includes the right bin limits. To still have the leftmost limit + # included, we include an empty bin left. + # If start == 0, we need to add an empty one left, otherwise we can just include the bin left to the + # first nonzero-count bin: + counts = ( + counts[start - 1 : end] if start > 0 else np.concatenate([[0], counts[:end]]) + ) + limits = limits[start : end + 1] + + if counts.size == 0 or limits.size == 0: + raise ValueError("The histogram is empty, please file a bug report.") + + sum_sq = values.dot(values) + return HistogramProto( + min=values.min(), + max=values.max(), + num=len(values), + sum=values.sum(), + sum_squares=sum_sq, + bucket_limit=limits.tolist(), + bucket=counts.tolist(), + ) + + +def image(tag, tensor, rescale=1, dataformats="NCHW"): + """Output a `Summary` protocol buffer with images. + + The summary has up to `max_images` summary values containing images. The + images are built from `tensor` which must be 3-D with shape `[height, width, + channels]` and where `channels` can be: + * 1: `tensor` is interpreted as Grayscale. + * 3: `tensor` is interpreted as RGB. + * 4: `tensor` is interpreted as RGBA. + The `name` in the outputted Summary.Value protobufs is generated based on the + name, with a suffix depending on the max_outputs setting: + * If `max_outputs` is 1, the summary value tag is '*name*/image'. + * If `max_outputs` is greater than 1, the summary value tags are + generated sequentially as '*name*/image/0', '*name*/image/1', etc. + Args: + tag: A name for the generated node. Will also serve as a series name in + TensorBoard. + tensor: A 3-D `uint8` or `float32` `Tensor` of shape `[height, width, + channels]` where `channels` is 1, 3, or 4. + 'tensor' can either have values in [0, 1] (float32) or [0, 255] (uint8). + The image() function will scale the image values to [0, 255] by applying + a scale factor of either 1 (uint8) or 255 (float32). Out-of-range values + will be clipped. + Returns: + A scalar `Tensor` of type `string`. The serialized `Summary` protocol + buffer. + """ + tensor = make_np(tensor) + tensor = convert_to_HWC(tensor, dataformats) + # Do not assume that user passes in values in [0, 255], use data type to detect + scale_factor = _calc_scale_factor(tensor) + tensor = tensor.astype(np.float32) + tensor = (tensor * scale_factor).clip(0, 255).astype(np.uint8) + image = make_image(tensor, rescale=rescale) + return Summary(value=[Summary.Value(tag=tag, image=image)]) + + +def image_boxes( + tag, tensor_image, tensor_boxes, rescale=1, dataformats="CHW", labels=None +): + """Output a `Summary` protocol buffer with images.""" + tensor_image = make_np(tensor_image) + tensor_image = convert_to_HWC(tensor_image, dataformats) + tensor_boxes = make_np(tensor_boxes) + tensor_image = tensor_image.astype(np.float32) * _calc_scale_factor(tensor_image) + image = make_image( + tensor_image.clip(0, 255).astype(np.uint8), + rescale=rescale, + rois=tensor_boxes, + labels=labels, + ) + return Summary(value=[Summary.Value(tag=tag, image=image)]) + + +def draw_boxes(disp_image, boxes, labels=None): + # xyxy format + num_boxes = boxes.shape[0] + list_gt = range(num_boxes) + for i in list_gt: + disp_image = _draw_single_box( + disp_image, + boxes[i, 0], + boxes[i, 1], + boxes[i, 2], + boxes[i, 3], + display_str=None if labels is None else labels[i], + color="Red", + ) + return disp_image + + +def make_image(tensor, rescale=1, rois=None, labels=None): + """Convert a numpy representation of an image to Image protobuf.""" + from PIL import Image + + height, width, channel = tensor.shape + scaled_height = int(height * rescale) + scaled_width = int(width * rescale) + image = Image.fromarray(tensor) + if rois is not None: + image = draw_boxes(image, rois, labels=labels) + ANTIALIAS = Image.Resampling.LANCZOS + image = image.resize((scaled_width, scaled_height), ANTIALIAS) + import io + + output = io.BytesIO() + image.save(output, format="PNG") + image_string = output.getvalue() + output.close() + return Summary.Image( + height=height, + width=width, + colorspace=channel, + encoded_image_string=image_string, + ) + + +def video(tag, tensor, fps=4): + tensor = make_np(tensor) + tensor = _prepare_video(tensor) + # If user passes in uint8, then we don't need to rescale by 255 + scale_factor = _calc_scale_factor(tensor) + tensor = tensor.astype(np.float32) + tensor = (tensor * scale_factor).clip(0, 255).astype(np.uint8) + video = make_video(tensor, fps) + return Summary(value=[Summary.Value(tag=tag, image=video)]) + + +def make_video(tensor, fps): + try: + import moviepy # noqa: F401 + except ImportError: + print("add_video needs package moviepy") + return + try: + from moviepy import editor as mpy + except ImportError: + print( + "moviepy is installed, but can't import moviepy.editor.", + "Some packages could be missing [imageio, requests]", + ) + return + import tempfile + + _t, h, w, c = tensor.shape + + # encode sequence of images into gif string + clip = mpy.ImageSequenceClip(list(tensor), fps=fps) + + with tempfile.NamedTemporaryFile(suffix=".gif") as f: + filename = f.name + try: # newer version of moviepy use logger instead of progress_bar argument. + clip.write_gif(filename, verbose=False, logger=None) + except TypeError: + try: # older version of moviepy does not support progress_bar argument. + clip.write_gif(filename, verbose=False, progress_bar=False) + except TypeError: + clip.write_gif(filename, verbose=False) + + f.seek(0) + tensor_string = f.read() + + return Summary.Image( + height=h, width=w, colorspace=c, encoded_image_string=tensor_string + ) + + +def audio(tag, tensor, sample_rate=44100): + array = make_np(tensor) + array = array.squeeze() + if abs(array).max() > 1: + print("warning: audio amplitude out of range, auto clipped.") + array = array.clip(-1, 1) + if array.ndim != 1: + raise AssertionError("input tensor should be 1 dimensional.") + # pyrefly: ignore [no-matching-overload] + array = (array * np.iinfo(np.int16).max).astype(" 127: # weird, value > 127 breaks protobuf + num_thresholds = 127 + data = np.stack((tp, fp, tn, fn, precision, recall)) + pr_curve_plugin_data = PrCurvePluginData( + version=0, num_thresholds=num_thresholds + ).SerializeToString() + plugin_data = SummaryMetadata.PluginData( + plugin_name="pr_curves", content=pr_curve_plugin_data + ) + smd = SummaryMetadata(plugin_data=plugin_data) + tensor = TensorProto( + dtype="DT_FLOAT", + float_val=data.reshape(-1).tolist(), + tensor_shape=TensorShapeProto( + dim=[ + TensorShapeProto.Dim(size=data.shape[0]), + TensorShapeProto.Dim(size=data.shape[1]), + ] + ), + ) + return Summary(value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)]) + + +def pr_curve(tag, labels, predictions, num_thresholds=127, weights=None): + # weird, value > 127 breaks protobuf + num_thresholds = min(num_thresholds, 127) + data = compute_curve( + labels, predictions, num_thresholds=num_thresholds, weights=weights + ) + pr_curve_plugin_data = PrCurvePluginData( + version=0, num_thresholds=num_thresholds + ).SerializeToString() + plugin_data = SummaryMetadata.PluginData( + plugin_name="pr_curves", content=pr_curve_plugin_data + ) + smd = SummaryMetadata(plugin_data=plugin_data) + tensor = TensorProto( + dtype="DT_FLOAT", + float_val=data.reshape(-1).tolist(), + tensor_shape=TensorShapeProto( + dim=[ + TensorShapeProto.Dim(size=data.shape[0]), + TensorShapeProto.Dim(size=data.shape[1]), + ] + ), + ) + return Summary(value=[Summary.Value(tag=tag, metadata=smd, tensor=tensor)]) + + +# https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/pr_curve/summary.py +def compute_curve(labels, predictions, num_thresholds=None, weights=None): + _MINIMUM_COUNT = 1e-7 + + if weights is None: + weights = 1.0 + + # Compute bins of true positives and false positives. + # pyrefly: ignore [unsupported-operation] + bucket_indices = np.int32(np.floor(predictions * (num_thresholds - 1))) + float_labels = labels.astype(np.float64) + # pyrefly: ignore [unsupported-operation] + histogram_range = (0, num_thresholds - 1) + tp_buckets, _ = np.histogram( + bucket_indices, + # pyrefly: ignore [bad-argument-type] + bins=num_thresholds, + range=histogram_range, + weights=float_labels * weights, + ) + fp_buckets, _ = np.histogram( + bucket_indices, + # pyrefly: ignore [bad-argument-type] + bins=num_thresholds, + range=histogram_range, + weights=(1.0 - float_labels) * weights, + ) + + # Obtain the reverse cumulative sum. + tp = np.cumsum(tp_buckets[::-1])[::-1] + fp = np.cumsum(fp_buckets[::-1])[::-1] + tn = fp[0] - fp + fn = tp[0] - tp + precision = tp / np.maximum(_MINIMUM_COUNT, tp + fp) + recall = tp / np.maximum(_MINIMUM_COUNT, tp + fn) + return np.stack((tp, fp, tn, fn, precision, recall)) + + +def _get_tensor_summary( + name, display_name, description, tensor, content_type, components, json_config +): + """Create a tensor summary with summary metadata. + + Args: + name: Uniquely identifiable name of the summary op. Could be replaced by + combination of name and type to make it unique even outside of this + summary. + display_name: Will be used as the display name in TensorBoard. + Defaults to `name`. + description: A longform readable description of the summary data. Markdown + is supported. + tensor: Tensor to display in summary. + content_type: Type of content inside the Tensor. + components: Bitmask representing present parts (vertices, colors, etc.) that + belong to the summary. + json_config: A string, JSON-serialized dictionary of ThreeJS classes + configuration. + + Returns: + Tensor summary with metadata. + """ + import torch + from tensorboard.plugins.mesh import metadata + + tensor = torch.as_tensor(tensor) + + tensor_metadata = metadata.create_summary_metadata( + name, + display_name, + content_type, + components, + tensor.shape, + description, + json_config=json_config, + ) + + tensor = TensorProto( + dtype="DT_FLOAT", + float_val=tensor.reshape(-1).tolist(), + tensor_shape=TensorShapeProto( + dim=[ + TensorShapeProto.Dim(size=tensor.shape[0]), + TensorShapeProto.Dim(size=tensor.shape[1]), + TensorShapeProto.Dim(size=tensor.shape[2]), + ] + ), + ) + + tensor_summary = Summary.Value( + tag=metadata.get_instance_name(name, content_type), + tensor=tensor, + metadata=tensor_metadata, + ) + + return tensor_summary + + +def _get_json_config(config_dict): + """Parse and returns JSON string from python dictionary.""" + json_config = "{}" + if config_dict is not None: + json_config = json.dumps(config_dict, sort_keys=True) + return json_config + + +# https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/mesh/summary.py +def mesh( + tag, vertices, colors, faces, config_dict, display_name=None, description=None +): + """Output a merged `Summary` protocol buffer with a mesh/point cloud. + + Args: + tag: A name for this summary operation. + vertices: Tensor of shape `[dim_1, ..., dim_n, 3]` representing the 3D + coordinates of vertices. + faces: Tensor of shape `[dim_1, ..., dim_n, 3]` containing indices of + vertices within each triangle. + colors: Tensor of shape `[dim_1, ..., dim_n, 3]` containing colors for each + vertex. + display_name: If set, will be used as the display name in TensorBoard. + Defaults to `name`. + description: A longform readable description of the summary data. Markdown + is supported. + config_dict: Dictionary with ThreeJS classes names and configuration. + + Returns: + Merged summary for mesh/point cloud representation. + """ + from tensorboard.plugins.mesh import metadata + from tensorboard.plugins.mesh.plugin_data_pb2 import MeshPluginData + + json_config = _get_json_config(config_dict) + + summaries = [] + tensors = [ + (vertices, MeshPluginData.VERTEX), + (faces, MeshPluginData.FACE), + (colors, MeshPluginData.COLOR), + ] + tensors = [tensor for tensor in tensors if tensor[0] is not None] + components = metadata.get_components_bitmask( + [content_type for (tensor, content_type) in tensors] + ) + + for tensor, content_type in tensors: + summaries.append( + _get_tensor_summary( + tag, + display_name, + description, + tensor, + content_type, + components, + json_config, + ) + ) + + return Summary(value=summaries) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/writer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/writer.py new file mode 100644 index 0000000000000000000000000000000000000000..2f1ccce77cf7211492898cbfe6bd48ab610d6ccb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/tensorboard/writer.py @@ -0,0 +1,1216 @@ +# mypy: allow-untyped-defs +"""Provide an API for writing protocol buffers to event files to be consumed by TensorBoard for visualization.""" +from __future__ import annotations + +import os +import time +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from matplotlib.figure import Figure +from tensorboard.compat import tf +from tensorboard.compat.proto import event_pb2 +from tensorboard.compat.proto.event_pb2 import Event, SessionLog +from tensorboard.plugins.projector.projector_config_pb2 import ProjectorConfig +from tensorboard.summary.writer.event_file_writer import EventFileWriter + +from ._convert_np import make_np +from ._embedding import get_embedding_info, make_mat, make_sprite, make_tsv, write_pbtxt +from ._onnx_graph import load_onnx_graph +from ._pytorch_graph import graph +from ._utils import figure_to_image +from .summary import ( + audio, + custom_scalars, + histogram, + histogram_raw, + hparams, + image, + image_boxes, + mesh, + pr_curve, + pr_curve_raw, + scalar, + tensor_proto, + text, + video, +) + +__all__ = ["FileWriter", "SummaryWriter"] + + +class FileWriter: + """Writes protocol buffers to event files to be consumed by TensorBoard. + + The `FileWriter` class provides a mechanism to create an event file in a + given directory and add summaries and events to it. The class updates the + file contents asynchronously. This allows a training program to call methods + to add data to the file directly from the training loop, without slowing down + training. + """ + + def __init__(self, log_dir, max_queue=10, flush_secs=120, filename_suffix="") -> None: + """Create a `FileWriter` and an event file. + + On construction the writer creates a new event file in `log_dir`. + The other arguments to the constructor control the asynchronous writes to + the event file. + + Args: + log_dir: A string. Directory where event file will be written. + max_queue: Integer. Size of the queue for pending events and + summaries before one of the 'add' calls forces a flush to disk. + Default is ten items. + flush_secs: Number. How often, in seconds, to flush the + pending events and summaries to disk. Default is every two minutes. + filename_suffix: A string. Suffix added to all event filenames + in the log_dir directory. More details on filename construction in + tensorboard.summary.writer.event_file_writer.EventFileWriter. + """ + # Sometimes PosixPath is passed in and we need to coerce it to + # a string in all cases + # TODO: See if we can remove this in the future if we are + # actually the ones passing in a PosixPath + log_dir = str(log_dir) + self.event_writer = EventFileWriter( + log_dir, max_queue, flush_secs, filename_suffix + ) + + def get_logdir(self): + """Return the directory where event file will be written.""" + return self.event_writer.get_logdir() + + def add_event(self, event, step=None, walltime=None) -> None: + """Add an event to the event file. + + Args: + event: An `Event` protocol buffer. + step: Number. Optional global step value for training process + to record with the event. + walltime: float. Optional walltime to override the default (current) + walltime (from time.time()) seconds after epoch + """ + event.wall_time = time.time() if walltime is None else walltime + if step is not None: + # Make sure step is converted from numpy or other formats + # since protobuf might not convert depending on version + event.step = int(step) + self.event_writer.add_event(event) + + def add_summary(self, summary, global_step=None, walltime=None) -> None: + """Add a `Summary` protocol buffer to the event file. + + This method wraps the provided summary in an `Event` protocol buffer + and adds it to the event file. + + Args: + summary: A `Summary` protocol buffer. + global_step: Number. Optional global step value for training process + to record with the summary. + walltime: float. Optional walltime to override the default (current) + walltime (from time.time()) seconds after epoch + """ + event = event_pb2.Event(summary=summary) + self.add_event(event, global_step, walltime) + + def add_graph(self, graph_profile, walltime=None) -> None: + """Add a `Graph` and step stats protocol buffer to the event file. + + Args: + graph_profile: A `Graph` and step stats protocol buffer. + walltime: float. Optional walltime to override the default (current) + walltime (from time.time()) seconds after epoch + """ + graph = graph_profile[0] + stepstats = graph_profile[1] + event = event_pb2.Event(graph_def=graph.SerializeToString()) + self.add_event(event, None, walltime) + + trm = event_pb2.TaggedRunMetadata( + tag="step1", run_metadata=stepstats.SerializeToString() + ) + event = event_pb2.Event(tagged_run_metadata=trm) + self.add_event(event, None, walltime) + + def add_onnx_graph(self, graph, walltime=None) -> None: + """Add a `Graph` protocol buffer to the event file. + + Args: + graph: A `Graph` protocol buffer. + walltime: float. Optional walltime to override the default (current) + _get_file_writerfrom time.time()) + """ + event = event_pb2.Event(graph_def=graph.SerializeToString()) + self.add_event(event, None, walltime) + + def flush(self) -> None: + """Flushes the event file to disk. + + Call this method to make sure that all pending events have been written to + disk. + """ + self.event_writer.flush() + + def close(self) -> None: + """Flushes the event file to disk and close the file. + + Call this method when you do not need the summary writer anymore. + """ + self.event_writer.close() + + def reopen(self) -> None: + """Reopens the EventFileWriter. + + Can be called after `close()` to add more events in the same directory. + The events will go into a new events file. + Does nothing if the EventFileWriter was not closed. + """ + self.event_writer.reopen() + + +class SummaryWriter: + """Writes entries directly to event files in the log_dir to be consumed by TensorBoard. + + The `SummaryWriter` class provides a high-level API to create an event file + in a given directory and add summaries and events to it. The class updates the + file contents asynchronously. This allows a training program to call methods + to add data to the file directly from the training loop, without slowing down + training. + """ + + def __init__( + self, + log_dir=None, + comment="", + purge_step=None, + max_queue=10, + flush_secs=120, + filename_suffix="", + ) -> None: + """Create a `SummaryWriter` that will write out events and summaries to the event file. + + Args: + log_dir (str): Save directory location. Default is + runs/**CURRENT_DATETIME_HOSTNAME**, which changes after each run. + Use hierarchical folder structure to compare + between runs easily. e.g. pass in 'runs/exp1', 'runs/exp2', etc. + for each new experiment to compare across them. + comment (str): Comment log_dir suffix appended to the default + ``log_dir``. If ``log_dir`` is assigned, this argument has no effect. + purge_step (int): + When logging crashes at step :math:`T+X` and restarts at step :math:`T`, + any events whose global_step larger or equal to :math:`T` will be + purged and hidden from TensorBoard. + Note that crashed and resumed experiments should have the same ``log_dir``. + max_queue (int): Size of the queue for pending events and + summaries before one of the 'add' calls forces a flush to disk. + Default is ten items. + flush_secs (int): How often, in seconds, to flush the + pending events and summaries to disk. Default is every two minutes. + filename_suffix (str): Suffix added to all event filenames in + the log_dir directory. More details on filename construction in + tensorboard.summary.writer.event_file_writer.EventFileWriter. + + Examples:: + + from torch.utils.tensorboard import SummaryWriter + + # create a summary writer with automatically generated folder name. + writer = SummaryWriter() + # folder location: runs/May04_22-14-54_s-MacBook-Pro.local/ + + # create a summary writer using the specified folder name. + writer = SummaryWriter("my_experiment") + # folder location: my_experiment + + # create a summary writer with comment appended. + writer = SummaryWriter(comment="LR_0.1_BATCH_16") + # folder location: runs/May04_22-14-54_s-MacBook-Pro.localLR_0.1_BATCH_16/ + + """ + torch._C._log_api_usage_once("tensorboard.create.summarywriter") + if not log_dir: + import socket + from datetime import datetime + + current_time = datetime.now().strftime("%b%d_%H-%M-%S") + log_dir = os.path.join( + "runs", current_time + "_" + socket.gethostname() + comment + ) + self.log_dir = log_dir + self.purge_step = purge_step + self.max_queue = max_queue + self.flush_secs = flush_secs + self.filename_suffix = filename_suffix + + # Initialize the file writers, but they can be cleared out on close + # and recreated later as needed. + self.file_writer = self.all_writers = None + self._get_file_writer() + + # Create default bins for histograms, see generate_testdata.py in tensorflow/tensorboard + v = 1e-12 + buckets = [] + neg_buckets = [] + while v < 1e20: + buckets.append(v) + neg_buckets.append(-v) + v *= 1.1 + self.default_bins = neg_buckets[::-1] + [0] + buckets + + def _get_file_writer(self): + """Return the default FileWriter instance. Recreates it if closed.""" + if self.all_writers is None or self.file_writer is None: + # pyrefly: ignore [bad-assignment] + self.file_writer = FileWriter( + self.log_dir, self.max_queue, self.flush_secs, self.filename_suffix + ) + # pyrefly: ignore [bad-assignment, missing-attribute] + self.all_writers = {self.file_writer.get_logdir(): self.file_writer} + if self.purge_step is not None: + most_recent_step = self.purge_step + # pyrefly: ignore [missing-attribute] + self.file_writer.add_event( + Event(step=most_recent_step, file_version="brain.Event:2") + ) + # pyrefly: ignore [missing-attribute] + self.file_writer.add_event( + Event( + step=most_recent_step, + session_log=SessionLog(status=SessionLog.START), + ) + ) + self.purge_step = None + return self.file_writer + + def get_logdir(self): + """Return the directory where event files will be written.""" + return self.log_dir + + def add_hparams( + self, + hparam_dict, + metric_dict, + hparam_domain_discrete=None, + run_name=None, + global_step=None, + ) -> None: + """Add a set of hyperparameters to be compared in TensorBoard. + + Args: + hparam_dict (dict): Each key-value pair in the dictionary is the + name of the hyper parameter and it's corresponding value. + The type of the value can be one of `bool`, `string`, `float`, + `int`, or `None`. + metric_dict (dict): Each key-value pair in the dictionary is the + name of the metric and it's corresponding value. Note that the key used + here should be unique in the tensorboard record. Otherwise the value + you added by ``add_scalar`` will be displayed in hparam plugin. In most + cases, this is unwanted. + hparam_domain_discrete: (Optional[Dict[str, List[Any]]]) A dictionary that + contains names of the hyperparameters and all discrete values they can hold + run_name (str): Name of the run, to be included as part of the logdir. + If unspecified, will use current timestamp. + global_step (int): Global step value to record + + Examples:: + + from torch.utils.tensorboard import SummaryWriter + with SummaryWriter() as w: + for i in range(5): + w.add_hparams({'lr': 0.1*i, 'bsize': i}, + {'hparam/accuracy': 10*i, 'hparam/loss': 10*i}) + + Expected result: + + .. image:: _static/img/tensorboard/add_hparam.png + :scale: 50 % + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_hparams") + if type(hparam_dict) is not dict or type(metric_dict) is not dict: + raise TypeError("hparam_dict and metric_dict should be dictionary.") + exp, ssi, sei = hparams(hparam_dict, metric_dict, hparam_domain_discrete) + + if not run_name: + run_name = str(time.time()) + logdir = os.path.join(self._get_file_writer().get_logdir(), run_name) + with SummaryWriter(log_dir=logdir) as w_hp: + w_hp.file_writer.add_summary(exp, global_step) + w_hp.file_writer.add_summary(ssi, global_step) + w_hp.file_writer.add_summary(sei, global_step) + for k, v in metric_dict.items(): + w_hp.add_scalar(k, v, global_step) + + def add_scalar( + self, + tag, + scalar_value, + global_step=None, + walltime=None, + new_style=False, + double_precision=False, + ) -> None: + """Add scalar data to summary. + + Args: + tag (str): Data identifier + scalar_value (float or string/blobname): Value to save + global_step (int): Global step value to record + walltime (float): Optional override default walltime (time.time()) + with seconds after epoch of event + new_style (boolean): Whether to use new style (tensor field) or old + style (simple_value field). New style could lead to faster data loading. + Examples:: + + from torch.utils.tensorboard import SummaryWriter + writer = SummaryWriter() + x = range(100) + for i in x: + writer.add_scalar('y=2x', i * 2, i) + writer.close() + + Expected result: + + .. image:: _static/img/tensorboard/add_scalar.png + :scale: 50 % + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_scalar") + + summary = scalar( + tag, scalar_value, new_style=new_style, double_precision=double_precision + ) + self._get_file_writer().add_summary(summary, global_step, walltime) + + def add_scalars(self, main_tag, tag_scalar_dict, global_step=None, walltime=None) -> None: + """Add many scalar data to summary. + + Args: + main_tag (str): The parent name for the tags + tag_scalar_dict (dict): Key-value pair storing the tag and corresponding values + global_step (int): Global step value to record + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + + Examples:: + + from torch.utils.tensorboard import SummaryWriter + writer = SummaryWriter() + r = 5 + for i in range(100): + writer.add_scalars('run_14h', {'xsinx':i*np.sin(i/r), + 'xcosx':i*np.cos(i/r), + 'tanx': np.tan(i/r)}, i) + writer.close() + # This call adds three values to the same scalar plot with the tag + # 'run_14h' in TensorBoard's scalar section. + + Expected result: + + .. image:: _static/img/tensorboard/add_scalars.png + :scale: 50 % + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_scalars") + walltime = time.time() if walltime is None else walltime + fw_logdir = self._get_file_writer().get_logdir() + for tag, scalar_value in tag_scalar_dict.items(): + fw_tag = fw_logdir + "/" + main_tag.replace("/", "_") + "_" + tag + if self.all_writers is None: + raise AssertionError("self.all_writers is None") + if fw_tag in self.all_writers: + fw = self.all_writers[fw_tag] + else: + fw = FileWriter( + fw_tag, self.max_queue, self.flush_secs, self.filename_suffix + ) + self.all_writers[fw_tag] = fw + fw.add_summary(scalar(main_tag, scalar_value), global_step, walltime) + + def add_tensor( + self, + tag, + tensor, + global_step=None, + walltime=None, + ) -> None: + """Add tensor data to summary. + + Args: + tag (str): Data identifier + tensor (torch.Tensor): tensor to save + global_step (int): Global step value to record + Examples:: + + from torch.utils.tensorboard import SummaryWriter + writer = SummaryWriter() + x = torch.tensor([1,2,3]) + writer.add_scalar('x', x) + writer.close() + + Expected result: + Summary::tensor::float_val [1,2,3] + ::tensor::shape [3] + ::tag 'x' + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_tensor") + + summary = tensor_proto(tag, tensor) + self._get_file_writer().add_summary(summary, global_step, walltime) + + def add_histogram( + self, + tag, + values, + global_step=None, + bins="tensorflow", + walltime=None, + max_bins=None, + ) -> None: + """Add histogram to summary. + + Args: + tag (str): Data identifier + values (torch.Tensor, numpy.ndarray, or string/blobname): Values to build histogram + global_step (int): Global step value to record + bins (str): One of {'tensorflow','auto', 'fd', ...}. This determines how the bins are made. You can find + other options in: https://numpy.org/doc/stable/reference/generated/numpy.histogram.html + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + + Examples:: + + from torch.utils.tensorboard import SummaryWriter + import numpy as np + writer = SummaryWriter() + for i in range(10): + x = np.random.random(1000) + writer.add_histogram('distribution centers', x + i, i) + writer.close() + + Expected result: + + .. image:: _static/img/tensorboard/add_histogram.png + :scale: 50 % + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_histogram") + if isinstance(bins, str) and bins == "tensorflow": + bins = self.default_bins + self._get_file_writer().add_summary( + histogram(tag, values, bins, max_bins=max_bins), global_step, walltime + ) + + def add_histogram_raw( + self, + tag, + min, + max, + num, + sum, + sum_squares, + bucket_limits, + bucket_counts, + global_step=None, + walltime=None, + ) -> None: + """Add histogram with raw data. + + Args: + tag (str): Data identifier + min (float or int): Min value + max (float or int): Max value + num (int): Number of values + sum (float or int): Sum of all values + sum_squares (float or int): Sum of squares for all values + bucket_limits (torch.Tensor, numpy.ndarray): Upper value per bucket. + The number of elements of it should be the same as `bucket_counts`. + bucket_counts (torch.Tensor, numpy.ndarray): Number of values per bucket + global_step (int): Global step value to record + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + see: https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/histogram/README.md + + Examples:: + + from torch.utils.tensorboard import SummaryWriter + import numpy as np + writer = SummaryWriter() + dummy_data = [] + for idx, value in enumerate(range(50)): + dummy_data += [idx + 0.001] * value + + bins = list(range(50+2)) + bins = np.array(bins) + values = np.array(dummy_data).astype(float).reshape(-1) + counts, limits = np.histogram(values, bins=bins) + sum_sq = values.dot(values) + writer.add_histogram_raw( + tag='histogram_with_raw_data', + min=values.min(), + max=values.max(), + num=len(values), + sum=values.sum(), + sum_squares=sum_sq, + bucket_limits=limits[1:].tolist(), + bucket_counts=counts.tolist(), + global_step=0) + writer.close() + + Expected result: + + .. image:: _static/img/tensorboard/add_histogram_raw.png + :scale: 50 % + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_histogram_raw") + if len(bucket_limits) != len(bucket_counts): + raise ValueError( + "len(bucket_limits) != len(bucket_counts), see the document." + ) + self._get_file_writer().add_summary( + histogram_raw( + tag, min, max, num, sum, sum_squares, bucket_limits, bucket_counts + ), + global_step, + walltime, + ) + + def add_image( + self, tag, img_tensor, global_step=None, walltime=None, dataformats="CHW" + ) -> None: + """Add image data to summary. + + Note that this requires the ``pillow`` package. + + Args: + tag (str): Data identifier + img_tensor (torch.Tensor, numpy.ndarray, or string/blobname): Image data + global_step (int): Global step value to record + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + dataformats (str): Image data format specification of the form + CHW, HWC, HW, WH, etc. + Shape: + img_tensor: Default is :math:`(3, H, W)`. You can use ``torchvision.utils.make_grid()`` to + convert a batch of tensor into 3xHxW format or call ``add_images`` and let us do the job. + Tensor with :math:`(1, H, W)`, :math:`(H, W)`, :math:`(H, W, 3)` is also suitable as long as + corresponding ``dataformats`` argument is passed, e.g. ``CHW``, ``HWC``, ``HW``. + + Examples:: + + from torch.utils.tensorboard import SummaryWriter + import numpy as np + img = np.zeros((3, 100, 100)) + img[0] = np.arange(0, 10000).reshape(100, 100) / 10000 + img[1] = 1 - np.arange(0, 10000).reshape(100, 100) / 10000 + + img_HWC = np.zeros((100, 100, 3)) + img_HWC[:, :, 0] = np.arange(0, 10000).reshape(100, 100) / 10000 + img_HWC[:, :, 1] = 1 - np.arange(0, 10000).reshape(100, 100) / 10000 + + writer = SummaryWriter() + writer.add_image('my_image', img, 0) + + # If you have non-default dimension setting, set the dataformats argument. + writer.add_image('my_image_HWC', img_HWC, 0, dataformats='HWC') + writer.close() + + Expected result: + + .. image:: _static/img/tensorboard/add_image.png + :scale: 50 % + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_image") + self._get_file_writer().add_summary( + image(tag, img_tensor, dataformats=dataformats), global_step, walltime + ) + + def add_images( + self, tag, img_tensor, global_step=None, walltime=None, dataformats="NCHW" + ) -> None: + """Add batched image data to summary. + + Note that this requires the ``pillow`` package. + + Args: + tag (str): Data identifier + img_tensor (torch.Tensor, numpy.ndarray, or string/blobname): Image data + global_step (int): Global step value to record + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + dataformats (str): Image data format specification of the form + NCHW, NHWC, CHW, HWC, HW, WH, etc. + Shape: + img_tensor: Default is :math:`(N, 3, H, W)`. If ``dataformats`` is specified, other shape will be + accepted. e.g. NCHW or NHWC. + + Examples:: + + from torch.utils.tensorboard import SummaryWriter + import numpy as np + + img_batch = np.zeros((16, 3, 100, 100)) + for i in range(16): + img_batch[i, 0] = np.arange(0, 10000).reshape(100, 100) / 10000 / 16 * i + img_batch[i, 1] = (1 - np.arange(0, 10000).reshape(100, 100) / 10000) / 16 * i + + writer = SummaryWriter() + writer.add_images('my_image_batch', img_batch, 0) + writer.close() + + Expected result: + + .. image:: _static/img/tensorboard/add_images.png + :scale: 30 % + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_images") + self._get_file_writer().add_summary( + image(tag, img_tensor, dataformats=dataformats), global_step, walltime + ) + + def add_image_with_boxes( + self, + tag, + img_tensor, + box_tensor, + global_step=None, + walltime=None, + rescale=1, + dataformats="CHW", + labels=None, + ) -> None: + """Add image and draw bounding boxes on the image. + + Args: + tag (str): Data identifier + img_tensor (torch.Tensor, numpy.ndarray, or string/blobname): Image data + box_tensor (torch.Tensor, numpy.ndarray, or string/blobname): Box data (for detected objects) + box should be represented as [x1, y1, x2, y2]. + global_step (int): Global step value to record + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + rescale (float): Optional scale override + dataformats (str): Image data format specification of the form + NCHW, NHWC, CHW, HWC, HW, WH, etc. + labels (list of string): The label to be shown for each bounding box. + Shape: + img_tensor: Default is :math:`(3, H, W)`. It can be specified with ``dataformats`` argument. + e.g. CHW or HWC + + box_tensor: (torch.Tensor, numpy.ndarray, or string/blobname): NX4, where N is the number of + boxes and each 4 elements in a row represents (xmin, ymin, xmax, ymax). + """ + torch._C._log_api_usage_once("tensorboard.logging.add_image_with_boxes") + if labels is not None: + if isinstance(labels, str): + labels = [labels] + if len(labels) != box_tensor.shape[0]: + labels = None + self._get_file_writer().add_summary( + image_boxes( + tag, + img_tensor, + box_tensor, + rescale=rescale, + dataformats=dataformats, + labels=labels, + ), + global_step, + walltime, + ) + + def add_figure( + self, + tag: str, + figure: Figure | list[Figure], + global_step: int | None = None, + close: bool = True, + walltime: float | None = None, + ) -> None: + """Render matplotlib figure into an image and add it to summary. + + Note that this requires the ``matplotlib`` package. + + Args: + tag: Data identifier + figure: Figure or a list of figures + global_step: Global step value to record + close: Flag to automatically close the figure + walltime: Optional override default walltime (time.time()) + seconds after epoch of event + """ + torch._C._log_api_usage_once("tensorboard.logging.add_figure") + if isinstance(figure, list): + self.add_image( + tag, + figure_to_image(figure, close), + global_step, + walltime, + dataformats="NCHW", + ) + else: + self.add_image( + tag, + figure_to_image(figure, close), + global_step, + walltime, + dataformats="CHW", + ) + + def add_video(self, tag, vid_tensor, global_step=None, fps=4, walltime=None) -> None: + """Add video data to summary. + + Note that this requires the ``moviepy`` package. + + Args: + tag (str): Data identifier + vid_tensor (torch.Tensor): Video data + global_step (int): Global step value to record + fps (float or int): Frames per second + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + Shape: + vid_tensor: :math:`(N, T, C, H, W)`. The values should lie in [0, 255] for type `uint8` or [0, 1] for type `float`. + """ + torch._C._log_api_usage_once("tensorboard.logging.add_video") + self._get_file_writer().add_summary( + video(tag, vid_tensor, fps), global_step, walltime + ) + + def add_audio( + self, tag, snd_tensor, global_step=None, sample_rate=44100, walltime=None + ) -> None: + """Add audio data to summary. + + Args: + tag (str): Data identifier + snd_tensor (torch.Tensor): Sound data + global_step (int): Global step value to record + sample_rate (int): sample rate in Hz + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + Shape: + snd_tensor: :math:`(1, L)`. The values should lie between [-1, 1]. + """ + torch._C._log_api_usage_once("tensorboard.logging.add_audio") + self._get_file_writer().add_summary( + audio(tag, snd_tensor, sample_rate=sample_rate), global_step, walltime + ) + + def add_text(self, tag, text_string, global_step=None, walltime=None) -> None: + """Add text data to summary. + + Args: + tag (str): Data identifier + text_string (str): String to save + global_step (int): Global step value to record + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + Examples:: + + writer.add_text('lstm', 'This is an lstm', 0) + writer.add_text('rnn', 'This is an rnn', 10) + """ + torch._C._log_api_usage_once("tensorboard.logging.add_text") + self._get_file_writer().add_summary( + text(tag, text_string), global_step, walltime + ) + + def add_onnx_graph(self, prototxt) -> None: + torch._C._log_api_usage_once("tensorboard.logging.add_onnx_graph") + self._get_file_writer().add_onnx_graph(load_onnx_graph(prototxt)) + + def add_graph( + self, model, input_to_model=None, verbose=False, use_strict_trace=True + ) -> None: + """Add graph data to summary. + + Args: + model (torch.nn.Module): Model to draw. + input_to_model (torch.Tensor or list of torch.Tensor): A variable or a tuple of + variables to be fed. + verbose (bool): Whether to print graph structure in console. + use_strict_trace (bool): Whether to pass keyword argument `strict` to + `torch.jit.trace`. Pass False when you want the tracer to + record your mutable container types (list, dict) + """ + torch._C._log_api_usage_once("tensorboard.logging.add_graph") + # A valid PyTorch model should have a 'forward' method + self._get_file_writer().add_graph( + graph(model, input_to_model, verbose, use_strict_trace) + ) + + @staticmethod + def _encode(rawstr): + # I'd use urllib but, I'm unsure about the differences from python3 to python2, etc. + retval = rawstr + retval = retval.replace("%", f"%{ord('%'):02x}") + retval = retval.replace("/", f"%{ord('/'):02x}") + retval = retval.replace("\\", "%%%02x" % (ord("\\"))) # noqa: UP031 + return retval + + def add_embedding( + self, + mat, + metadata=None, + label_img=None, + global_step=None, + tag="default", + metadata_header=None, + ) -> None: + """Add embedding projector data to summary. + + Args: + mat (torch.Tensor or numpy.ndarray): A matrix which each row is the feature vector of the data point + metadata (list): A list of labels, each element will be converted to string + label_img (torch.Tensor): Images correspond to each data point + global_step (int): Global step value to record + tag (str): Name for the embedding + metadata_header (list): A list of headers for multi-column metadata. If given, each metadata must be + a list with values corresponding to headers. + Shape: + mat: :math:`(N, D)`, where N is number of data and D is feature dimension + + label_img: :math:`(N, C, H, W)` + + Examples:: + + import keyword + import torch + meta = [] + while len(meta)<100: + meta = meta+keyword.kwlist # get some strings + meta = meta[:100] + + for i, v in enumerate(meta): + meta[i] = v+str(i) + + label_img = torch.rand(100, 3, 10, 32) + for i in range(100): + label_img[i]*=i/100.0 + + writer.add_embedding(torch.randn(100, 5), metadata=meta, label_img=label_img) + writer.add_embedding(torch.randn(100, 5), label_img=label_img) + writer.add_embedding(torch.randn(100, 5), metadata=meta) + + .. note:: + Categorical (i.e. non-numeric) metadata cannot have more than 50 unique values if they are to be used for + coloring in the embedding projector. + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_embedding") + mat = make_np(mat) + if global_step is None: + global_step = 0 + # clear pbtxt? + + # Maybe we should encode the tag so slashes don't trip us up? + # I don't think this will mess us up, but better safe than sorry. + subdir = f"{str(global_step).zfill(5)}/{self._encode(tag)}" + save_path = os.path.join(self._get_file_writer().get_logdir(), subdir) + + fs = tf.io.gfile + if fs.exists(save_path): + if fs.isdir(save_path): + print( + "warning: Embedding dir exists, did you set global_step for add_embedding()?" + ) + else: + raise NotADirectoryError( + f"Path: `{save_path}` exists, but is a file. Cannot proceed." + ) + else: + fs.makedirs(save_path) + + if metadata is not None: + if mat.shape[0] != len( + metadata + ): + raise AssertionError("#labels should equal with #data points") + make_tsv(metadata, save_path, metadata_header=metadata_header) + + if label_img is not None: + if mat.shape[0] != label_img.shape[0]: + raise AssertionError("#images should equal with #data points") + make_sprite(label_img, save_path) + + if mat.ndim != 2: + raise AssertionError("mat should be 2D, where mat.size(0) is the number of data points") + make_mat(mat, save_path) + + # Filesystem doesn't necessarily have append semantics, so we store an + # internal buffer to append to and re-write whole file after each + # embedding is added + if not hasattr(self, "_projector_config"): + self._projector_config = ProjectorConfig() + embedding_info = get_embedding_info( + metadata, label_img, subdir, global_step, tag + ) + self._projector_config.embeddings.extend([embedding_info]) + + + from google.protobuf import text_format + + config_pbtxt = text_format.MessageToString(self._projector_config) + write_pbtxt(self._get_file_writer().get_logdir(), config_pbtxt) + + def add_pr_curve( + self, + tag, + labels, + predictions, + global_step=None, + num_thresholds=127, + weights=None, + walltime=None, + ) -> None: + """Add precision recall curve. + + Plotting a precision-recall curve lets you understand your model's + performance under different threshold settings. With this function, + you provide the ground truth labeling (T/F) and prediction confidence + (usually the output of your model) for each target. The TensorBoard UI + will let you choose the threshold interactively. + + Args: + tag (str): Data identifier + labels (torch.Tensor, numpy.ndarray, or string/blobname): + Ground truth data. Binary label for each element. + predictions (torch.Tensor, numpy.ndarray, or string/blobname): + The probability that an element be classified as true. + Value should be in [0, 1] + global_step (int): Global step value to record + num_thresholds (int): Number of thresholds used to draw the curve. + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + + Examples:: + + from torch.utils.tensorboard import SummaryWriter + import numpy as np + labels = np.random.randint(2, size=100) # binary label + predictions = np.random.rand(100) + writer = SummaryWriter() + writer.add_pr_curve('pr_curve', labels, predictions, 0) + writer.close() + + """ + torch._C._log_api_usage_once("tensorboard.logging.add_pr_curve") + labels, predictions = make_np(labels), make_np(predictions) + self._get_file_writer().add_summary( + pr_curve(tag, labels, predictions, num_thresholds, weights), + global_step, + walltime, + ) + + def add_pr_curve_raw( + self, + tag, + true_positive_counts, + false_positive_counts, + true_negative_counts, + false_negative_counts, + precision, + recall, + global_step=None, + num_thresholds=127, + weights=None, + walltime=None, + ) -> None: + """Add precision recall curve with raw data. + + Args: + tag (str): Data identifier + true_positive_counts (torch.Tensor, numpy.ndarray, or string/blobname): true positive counts + false_positive_counts (torch.Tensor, numpy.ndarray, or string/blobname): false positive counts + true_negative_counts (torch.Tensor, numpy.ndarray, or string/blobname): true negative counts + false_negative_counts (torch.Tensor, numpy.ndarray, or string/blobname): false negative counts + precision (torch.Tensor, numpy.ndarray, or string/blobname): precision + recall (torch.Tensor, numpy.ndarray, or string/blobname): recall + global_step (int): Global step value to record + num_thresholds (int): Number of thresholds used to draw the curve. + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + see: https://github.com/tensorflow/tensorboard/blob/master/tensorboard/plugins/pr_curve/README.md + """ + torch._C._log_api_usage_once("tensorboard.logging.add_pr_curve_raw") + self._get_file_writer().add_summary( + pr_curve_raw( + tag, + true_positive_counts, + false_positive_counts, + true_negative_counts, + false_negative_counts, + precision, + recall, + num_thresholds, + weights, + ), + global_step, + walltime, + ) + + def add_custom_scalars_multilinechart( + self, tags, category="default", title="untitled" + ) -> None: + """Shorthand for creating multilinechart. Similar to ``add_custom_scalars()``, but the only necessary argument is *tags*. + + Args: + tags (list): list of tags that have been used in ``add_scalar()`` + + Examples:: + + writer.add_custom_scalars_multilinechart(['twse/0050', 'twse/2330']) + """ + torch._C._log_api_usage_once( + "tensorboard.logging.add_custom_scalars_multilinechart" + ) + layout = {category: {title: ["Multiline", tags]}} + self._get_file_writer().add_summary(custom_scalars(layout)) + + def add_custom_scalars_marginchart( + self, tags, category="default", title="untitled" + ) -> None: + """Shorthand for creating marginchart. + + Similar to ``add_custom_scalars()``, but the only necessary argument is *tags*, + which should have exactly 3 elements. + + Args: + tags (list): list of tags that have been used in ``add_scalar()`` + + Examples:: + + writer.add_custom_scalars_marginchart(['twse/0050', 'twse/2330', 'twse/2006']) + """ + torch._C._log_api_usage_once( + "tensorboard.logging.add_custom_scalars_marginchart" + ) + if len(tags) != 3: + raise AssertionError(f"Expected 3 tags, got {len(tags)}.") + layout = {category: {title: ["Margin", tags]}} + self._get_file_writer().add_summary(custom_scalars(layout)) + + def add_custom_scalars(self, layout) -> None: + """Create special chart by collecting charts tags in 'scalars'. + + NOTE: This function can only be called once for each SummaryWriter() object. + + Because it only provides metadata to tensorboard, the function can be called before or after the training loop. + + Args: + layout (dict): {categoryName: *charts*}, where *charts* is also a dictionary + {chartName: *ListOfProperties*}. The first element in *ListOfProperties* is the chart's type + (one of **Multiline** or **Margin**) and the second element should be a list containing the tags + you have used in add_scalar function, which will be collected into the new chart. + + Examples:: + + layout = {'Taiwan':{'twse':['Multiline',['twse/0050', 'twse/2330']]}, + 'USA':{ 'dow':['Margin', ['dow/aaa', 'dow/bbb', 'dow/ccc']], + 'nasdaq':['Margin', ['nasdaq/aaa', 'nasdaq/bbb', 'nasdaq/ccc']]}} + + writer.add_custom_scalars(layout) + """ + torch._C._log_api_usage_once("tensorboard.logging.add_custom_scalars") + self._get_file_writer().add_summary(custom_scalars(layout)) + + def add_mesh( + self, + tag, + vertices, + colors=None, + faces=None, + config_dict=None, + global_step=None, + walltime=None, + ) -> None: + """Add meshes or 3D point clouds to TensorBoard. + + The visualization is based on Three.js, + so it allows users to interact with the rendered object. Besides the basic definitions + such as vertices, faces, users can further provide camera parameter, lighting condition, etc. + Please check https://threejs.org/docs/index.html#manual/en/introduction/Creating-a-scene for + advanced usage. + + Args: + tag (str): Data identifier + vertices (torch.Tensor): List of the 3D coordinates of vertices. + colors (torch.Tensor): Colors for each vertex + faces (torch.Tensor): Indices of vertices within each triangle. (Optional) + config_dict: Dictionary with ThreeJS classes names and configuration. + global_step (int): Global step value to record + walltime (float): Optional override default walltime (time.time()) + seconds after epoch of event + + Shape: + vertices: :math:`(B, N, 3)`. (batch, number_of_vertices, channels) + + colors: :math:`(B, N, 3)`. The values should lie in [0, 255] for type `uint8` or [0, 1] for type `float`. + + faces: :math:`(B, N, 3)`. The values should lie in [0, number_of_vertices] for type `uint8`. + + Examples:: + + from torch.utils.tensorboard import SummaryWriter + vertices_tensor = torch.as_tensor([ + [1, 1, 1], + [-1, -1, 1], + [1, -1, -1], + [-1, 1, -1], + ], dtype=torch.float).unsqueeze(0) + colors_tensor = torch.as_tensor([ + [255, 0, 0], + [0, 255, 0], + [0, 0, 255], + [255, 0, 255], + ], dtype=torch.int).unsqueeze(0) + faces_tensor = torch.as_tensor([ + [0, 2, 3], + [0, 3, 1], + [0, 1, 2], + [1, 3, 2], + ], dtype=torch.int).unsqueeze(0) + + writer = SummaryWriter() + writer.add_mesh('my_mesh', vertices=vertices_tensor, colors=colors_tensor, faces=faces_tensor) + + writer.close() + """ + torch._C._log_api_usage_once("tensorboard.logging.add_mesh") + self._get_file_writer().add_summary( + mesh(tag, vertices, colors, faces, config_dict), global_step, walltime + ) + + def flush(self) -> None: + """Flushes the event file to disk. + + Call this method to make sure that all pending events have been written to + disk. + """ + if self.all_writers is None: + return + for writer in self.all_writers.values(): + writer.flush() + + def close(self) -> None: + if self.all_writers is None: + return # ignore double close + for writer in self.all_writers.values(): + writer.flush() + writer.close() + # pyrefly: ignore [bad-assignment] + self.file_writer = self.all_writers = None + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + self.close() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/throughput_benchmark.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/throughput_benchmark.py new file mode 100644 index 0000000000000000000000000000000000000000..d4b94e0b13a39fbe192f89e791c663b2ecf45ea2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/throughput_benchmark.py @@ -0,0 +1,161 @@ +# mypy: allow-untyped-defs + +import torch._C + + +def format_time(time_us=None, time_ms=None, time_s=None) -> str: + """Define time formatting.""" + if sum([time_us is not None, time_ms is not None, time_s is not None]) != 1: + raise AssertionError("Expected only one of time_us, time_ms, time_s is given.") + + US_IN_SECOND = 1e6 + US_IN_MS = 1e3 + + if time_us is None: + if time_ms is not None: + time_us = time_ms * US_IN_MS + elif time_s is not None: + time_us = time_s * US_IN_SECOND + else: + raise AssertionError("Shouldn't reach here :)") + + if time_us >= US_IN_SECOND: + return f'{time_us / US_IN_SECOND:.3f}s' + if time_us >= US_IN_MS: + return f'{time_us / US_IN_MS:.3f}ms' + return f'{time_us:.3f}us' + + +class ExecutionStats: + def __init__(self, c_stats, benchmark_config) -> None: + self._c_stats = c_stats + self.benchmark_config = benchmark_config + + @property + def latency_avg_ms(self): + return self._c_stats.latency_avg_ms + + @property + def num_iters(self): + return self._c_stats.num_iters + + @property + def iters_per_second(self): + """Return total number of iterations per second across all calling threads.""" + return self.num_iters / self.total_time_seconds + + @property + def total_time_seconds(self): + return self.num_iters * ( + self.latency_avg_ms / 1000.0) / self.benchmark_config.num_calling_threads + + def __str__(self) -> str: + return '\n'.join([ + "Average latency per example: " + format_time(time_ms=self.latency_avg_ms), + f"Total number of iterations: {self.num_iters}", + f"Total number of iterations per second (across all threads): {self.iters_per_second:.2f}", + "Total time: " + format_time(time_s=self.total_time_seconds) + ]) + + +class ThroughputBenchmark: + """ + This class is a wrapper around a c++ component throughput_benchmark::ThroughputBenchmark. + + This wrapper on the throughput_benchmark::ThroughputBenchmark component is responsible + for executing a PyTorch module (nn.Module or ScriptModule) under an inference + server like load. It can emulate multiple calling threads to a single module + provided. In the future we plan to enhance this component to support inter and + intra-op parallelism as well as multiple models running in a single process. + + Please note that even though nn.Module is supported, it might incur an overhead + from the need to hold GIL every time we execute Python code or pass around + inputs as Python objects. As soon as you have a ScriptModule version of your + model for inference deployment it is better to switch to using it in this + benchmark. + + Example:: + + >>> # xdoctest: +SKIP("undefined vars") + >>> from torch.utils import ThroughputBenchmark + >>> bench = ThroughputBenchmark(my_module) + >>> # Pre-populate benchmark's data set with the inputs + >>> for input in inputs: + ... # Both args and kwargs work, same as any PyTorch Module / ScriptModule + ... bench.add_input(input[0], x2=input[1]) + >>> # Inputs supplied above are randomly used during the execution + >>> stats = bench.benchmark( + ... num_calling_threads=4, + ... num_warmup_iters = 100, + ... num_iters = 1000, + ... ) + >>> print("Avg latency (ms): {}".format(stats.latency_avg_ms)) + >>> print("Number of iterations: {}".format(stats.num_iters)) + """ + + def __init__(self, module) -> None: + if isinstance(module, torch.jit.ScriptModule): + self._benchmark = torch._C.ThroughputBenchmark(module._c) + else: + self._benchmark = torch._C.ThroughputBenchmark(module) + + def run_once(self, *args, **kwargs): + """ + Given input id (input_idx) run benchmark once and return prediction. + + This is useful for testing that benchmark actually runs the module you + want it to run. input_idx here is an index into inputs array populated + by calling add_input() method. + """ + return self._benchmark.run_once(*args, **kwargs) + + def add_input(self, *args, **kwargs) -> None: + """ + Store a single input to a module into the benchmark memory and keep it there. + + During the benchmark execution every thread is going to pick up a + random input from the all the inputs ever supplied to the benchmark via + this function. + """ + self._benchmark.add_input(*args, **kwargs) + + def benchmark( + self, + num_calling_threads=1, + num_warmup_iters=10, + num_iters=100, + profiler_output_path=""): + """ + Run a benchmark on the module. + + Args: + num_warmup_iters (int): Warmup iters are used to make sure we run a module + a few times before actually measuring things. This way we avoid cold + caches and any other similar problems. This is the number of warmup + iterations for each of the thread in separate + + num_iters (int): Number of iterations the benchmark should run with. + This number is separate from the warmup iterations. Also the number is + shared across all the threads. Once the num_iters iterations across all + the threads is reached, we will stop execution. Though total number of + iterations might be slightly larger. Which is reported as + stats.num_iters where stats is the result of this function + + profiler_output_path (str): Location to save Autograd Profiler trace. + If not empty, Autograd Profiler will be enabled for the main benchmark + execution (but not the warmup phase). The full trace will be saved + into the file path provided by this argument + + + This function returns BenchmarkExecutionStats object which is defined via pybind11. + It currently has two fields: + - num_iters - number of actual iterations the benchmark have made + - avg_latency_ms - average time it took to infer on one input example in milliseconds + """ + config = torch._C.BenchmarkConfig() + config.num_calling_threads = num_calling_threads + config.num_warmup_iters = num_warmup_iters + config.num_iters = num_iters + config.profiler_output_path = profiler_output_path + c_stats = self._benchmark.benchmark(config) + return ExecutionStats(c_stats, config) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/viz/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/viz/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/viz/_cycles.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/viz/_cycles.py new file mode 100644 index 0000000000000000000000000000000000000000..2f68cf75dee4d37bb8d3dead4bcecccd7048159b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/viz/_cycles.py @@ -0,0 +1,505 @@ +# mypy: allow-untyped-defs +import gc +import sys +from typing import Any, NamedTuple +import types +import weakref +import json +from tempfile import NamedTemporaryFile +import torch +from torch.cuda._memory_viz import _frames_fmt, _block_extra +import atexit +import logging +logger = logging.getLogger(__name__) + +def observe_garbage(observer): + enabled = True + + def disable() -> None: + # when GC runs during exit, things like `sys` will already be unloaded + # so we have to disable the callback to avoid hitting errors. + nonlocal enabled + enabled = False + atexit.register(disable) + + def gc_callback(phase, info) -> None: + nonlocal enabled + if not enabled: + return + if phase == "start": + gc.set_debug(gc.DEBUG_SAVEALL) + elif phase == "stop": + orig_trace = sys.getprofile() + self_return = [False] + + def do_collect(*args, **kwargs): + nonlocal enabled + if not self_return[0]: + self_return[0] = True + else: + sys.setprofile(orig_trace) + enabled = False + try: + # things in gc.garbage have survived a collection + # so to free them we have to collect a generation greater than them + # but that might _also_ free other stuff and we don't want to miss + # that stuff. So we have to now force gc at the highest level here, + # report all of what we found, _then_ we can free it up. + if info['generation'] != 2: + gc.collect() + observer(gc.garbage) + gc.garbage.clear() + # we have to re-run GC to clean up the cycles + # we saved from before. + gc.set_debug(0) + before = torch.cuda.memory_allocated() + gc.collect() + after = torch.cuda.memory_allocated() + if before != after: + logger.warning("CUDA Memory changed during GC, %d bytes freed.", before - after) + finally: + enabled = True + if orig_trace is not None: + return orig_trace(*args, **kwargs) + sys.setprofile(do_collect) + + gc.callbacks.append(gc_callback) + + # provide a way to disarm the callback + def remove() -> None: + gc.callbacks.remove(gc_callback) + return remove + +# Function to visualize cycles adapted from refcycle: +# Copyright 2013 Mark Dickinson +# +# 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 +# +# http://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. + +def _get_cell_type(): + def f(x=None): + return lambda: x + return type(f().__closure__[0]) + +CellType = _get_cell_type() + +def annotated_references(obj): + """ + Return known information about references held by the given object. + + Returns a mapping from referents to lists of descriptions. Note that there + may be more than one edge leading to any particular referent; hence the + need for a list. Descriptions are currently strings. + + """ + references: dict[int, list[str]] = {} + + def add_reference(name, obj) -> None: + references.setdefault(id(obj), []).append(name) + + def add_attrs(*attrs) -> None: + for attr in attrs: + if hasattr(obj, attr): + add_reference(attr, getattr(obj, attr)) + + def add_cell_references() -> None: + try: + add_attrs("cell_contents") + except ValueError: + # if cell_contents is empty, + # accessing it raises ValueError + # in this case there is no object to + # annotate + pass + + def add_function_references() -> None: + add_attrs("__defaults__", + "__closure__", + "__globals__", + "__code__", + "__name__", + "__module__", + "__doc__" + "__qualname__", + "__annotations__", + "__kwdefaults__") + + + def add_sequence_references() -> None: + for position, item in enumerate(obj): + add_reference(f"[{position}]", item) + + def add_dict_references() -> None: + for key, value in obj.items(): + add_reference("key", key) + add_reference(f"[{repr(key)}]", value) + + def add_set_references() -> None: + for elt in obj: + add_reference("element", elt) + + def add_bound_method_references() -> None: + add_attrs("__self__", "__func__", "im_class") + + def add_weakref_references() -> None: + # For subclasses of weakref, we can't reliably distinguish the + # callback (if any) from other attributes. + if type(obj) is weakref.ref: + referents = gc.get_referents(obj) + if len(referents) == 1: + target = referents[0] + add_reference("__callback__", target) + + + def add_frame_references() -> None: + f_locals = obj.f_locals + add_attrs("f_back", "f_code", "f_builtins", "f_globals", "f_trace", "f_locals") + # Some badly-behaved code replaces the f_locals dict with + # something that doesn't support the full dict interface. So we + # only continue with the annotation if f_locals is a Python dict. + if type(f_locals) is dict: + for name, local in obj.f_locals.items(): + add_reference(f"local {name}", local) + + def add_getset_descriptor_references() -> None: + add_attrs("__objclass__", "__name__", "__doc__") + + type_based_references = { + tuple: add_sequence_references, + list: add_sequence_references, + dict: add_dict_references, + set: add_set_references, + frozenset: add_set_references, + types.FunctionType: add_function_references, + types.FrameType: add_frame_references, + CellType: add_cell_references, + types.MethodType: add_bound_method_references, + weakref.ref: add_weakref_references, + types.GetSetDescriptorType: add_getset_descriptor_references, + } + + for type_ in type(obj).__mro__: + if type_ in type_based_references: + type_based_references[type_]() + + add_attrs("__dict__", "__class__") + if isinstance(obj, type): + add_attrs("__mro__") + + return references + +############################################################################### +# Object annotations. + + +BASE_TYPES = (int, float, complex, type(None), str, bytes) +FRAME_FILENAME_LIMIT = 32 + +def object_annotation(obj): + """ + Return a string to be used for Graphviz nodes. + + The string should be short but as informative as possible. + """ + + def format_sequence(obj): + body = ','.join(repr(x) if isinstance(x, BASE_TYPES) else type(x).__name__ for x in obj[:8]) + if len(obj) > 8: + body = f'{body}, ...{len(obj) - 8}' + return body + + # For basic types, use the repr. + if isinstance(obj, BASE_TYPES): + return repr(obj) + if type(obj).__name__ == 'function': + return f"function\n{obj.__name__}" + elif isinstance(obj, types.MethodType): + try: + func_name = obj.__func__.__qualname__ + except AttributeError: + func_name = "" + return f"instancemethod\n{func_name}" + elif isinstance(obj, list): + return f"[{format_sequence(obj)}]" + elif isinstance(obj, tuple): + return f"({format_sequence(obj)})" + elif isinstance(obj, dict): + return f"dict[{len(obj)}]" + elif isinstance(obj, types.ModuleType): + return f"module\n{obj.__name__}" + elif isinstance(obj, type): + return f"type\n{obj.__name__}" + elif isinstance(obj, weakref.ref): + referent = obj() + if referent is None: + return "weakref (dead referent)" + else: + return f"weakref to id 0x{id(referent):x}" + elif isinstance(obj, types.FrameType): + filename = obj.f_code.co_filename + if len(filename) > FRAME_FILENAME_LIMIT: + filename = "..." + filename[-(FRAME_FILENAME_LIMIT - 3):] + return f"frame\n{filename}:{obj.f_lineno}" + elif is_cuda_tensor(obj): + return f"object\n{type(obj).__module__}.{type(obj).__name__} ({obj.shape})" + else: + return f"object\n{type(obj).__module__}.{type(obj).__name__}" + + + +class Node(NamedTuple): + label: str + context: str | None + root: bool + referrents: list[tuple[str, int]] + +def create_graph(objects, *, context=None, filter=None): + if context is None: + context = cuda_allocation_context() + if filter is None: + filter = is_cuda_tensor + + objects = [obj for obj in objects if not isinstance(obj, weakref.ProxyTypes)] + nodes = [Node(object_annotation(obj), context(obj), filter(obj), []) for obj in objects] + node_referrers: list[list[int]] = [[] for obj in objects] + + id_to_node = {id(obj): i for i, obj in enumerate(objects)} + for obj in objects: + fidx = id_to_node[id(obj)] + f = nodes[fidx] + references = annotated_references(obj) + for referrent in gc.get_referents(obj): + rid = id(referrent) + tidx = id_to_node.get(rid) + if tidx is None: + continue + labels = references.get(rid, ["?"]) + node_referrers[tidx].append(fidx) + for label in labels: + f.referrents.append((label, tidx)) + + to_search = [i for i, n in enumerate(nodes) if n.root] + to_keep = set() + while to_search: + idx = to_search.pop() + if idx in to_keep: + continue + to_keep.add(idx) + referrers = node_referrers[idx] + to_search.extend(referrers) + id_to_filtered_id: dict[int, int] = {} + filtered: list[Any] = [] + for i, n in enumerate(nodes): + if i in to_keep: + id_to_filtered_id[i] = len(id_to_filtered_id) + filtered.append(n) + for n in filtered: + n.referrents[:] = [(label, id_to_filtered_id[idx]) + for (label, idx) in n.referrents + if idx in id_to_filtered_id] + return filtered + +def escape(n): + return json.dumps(n) + + +def is_cuda_tensor(obj): + return ( + isinstance(obj, torch.Tensor) and + obj.device.type == "cuda" and + not isinstance(obj, torch._subclasses.FakeTensor) + ) + +def cuda_allocation_context(): + snapshot = torch.cuda.memory._snapshot() + addr_to_frame = {} + for seg in snapshot['segments']: + addr = seg['address'] + for blk in seg['blocks']: + if blk['state'] == 'active_allocated': + frames, _real_size = _block_extra(blk) + addr_to_frame[addr] = frames + addr += blk['size'] + + def object_context(obj): + if is_cuda_tensor(obj): + addr = obj.untyped_storage().data_ptr() + frames = addr_to_frame.get(addr) + if frames is not None: + return '\n'.join(_frames_fmt(frames, full_filename=True)) + return None + return object_context + +def to_dot(nodes): + lines = ["digraph GraphName {", "node [shape=rect];", 'rankdir=LR;'] + for i, n in enumerate(nodes): + lines.append(f'{i} [label={escape(n.label)}, color={"red" if n.root else "black"}];') + + for i, f in enumerate(nodes): + for label, j in f.referrents: + lines.append(f'{i} -> {j} [label = {escape(label)}]') + lines.append("}\n") + return '\n'.join(lines) + +_template = """ + + + + + + +
+
+
+
+ 
Mouse over tensor objects to see where they were allocated.
+
+
+ + + + +""" +_listener_template = """ +document.getElementById('node{id}').addEventListener('mouseover', function(event) {{ + document.getElementById("stacktrace").textContent = {stack} +}}) +""" +def to_html(nodes): + listeners = [] + for i, n in enumerate(nodes): + if n.context is None: + continue + s = _listener_template.format(id=str(i + 1), stack=escape(f'{n.label}:\n{n.context}')) + listeners.append(s) + dot = to_dot(nodes) + return _template.replace('$DOT', repr(dot)).replace('$LISTENERS', '\n'.join(listeners)) + +def observe_tensor_cycles(callback): + torch.cuda.memory._record_memory_history(max_entries=100000) + + def observer(garbage) -> None: + if garbage: + if not any(is_cuda_tensor(obj) for obj in garbage): + logger.info("No CUDA Tensors found in garbage") + return + callback(to_html(create_graph(garbage))) + return observe_garbage(observer) + + +def warn_tensor_cycles(): + """ + Install a warning that reports whenever a cycle that is holding CUDA memory is observed. + + The warning produces an .html file that visualizes the cycle, + and links it to the stack frame that allocated the CUDA tensor. + + Reference cycles are freed by the cycle collector rather than being cleaned up + when the objects in the cycle first become unreachable. If a cycle points to a tensor, + the CUDA memory for that tensor will not be freed until garbage collection runs. + Accumulation of CUDA allocations can lead to out of memory errors (OOMs), as well as + non-deterministic allocation behavior which is harder to debug. + """ + logger.info("Watching Python reference cycles for CUDA Tensors.") + + def write_and_log(html) -> None: + with NamedTemporaryFile('w', suffix='.html') as f: + f.write(html) + logger.warning('Reference cycle includes a CUDA Tensor see visualization of cycle %s', f.name) + return observe_tensor_cycles(write_and_log) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/weak.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/weak.py new file mode 100644 index 0000000000000000000000000000000000000000..fa22b2b3765e9f68dd37b3af32e764ce3c074e9d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/utils/weak.py @@ -0,0 +1,367 @@ +# mypy: allow-untyped-defs +from __future__ import annotations + +import collections.abc as _collections_abc +import weakref +from collections.abc import Mapping, MutableMapping +from weakref import ref + +from torch import Tensor + + +WeakRef = ref + + +__all__ = [ + "TensorWeakRef", + "WeakIdRef", + "WeakIdKeyDictionary", + "WeakTensorKeyDictionary", +] + + +# TODO: make weakref properly thread safe following +# https://github.com/python/cpython/pull/125325 +class _IterationGuard: + # This context manager registers itself in the current iterators of the + # weak container, such as to delay all removals until the context manager + # exits. + # This technique should be relatively thread-safe (since sets are). + + def __init__(self, weakcontainer) -> None: + # Don't create cycles + self.weakcontainer = ref(weakcontainer) + + def __enter__(self): + w = self.weakcontainer() + if w is not None: + w._iterating.add(self) + return self + + def __exit__(self, e, t, b): + w = self.weakcontainer() + if w is not None: + s = w._iterating + s.remove(self) + if not s: + w._commit_removals() + + +# This file defines a variant of WeakKeyDictionary that overrides the hashing +# behavior of the key to use object identity, rather than the builtin +# __eq__/__hash__ functions. This is useful for Tensor weak keys, as their +# __eq__ implementation return a Tensor (elementwise equality), which means +# you can't use them directly with the WeakKeyDictionary in standard library. +# +# Our implementation strategy is to create a wrapper weak key object, which we +# use as a key in a stock Python dictionary. This is similar to how weakref +# implements WeakKeyDictionary, but instead of using weakref.ref as the +# wrapper, we use a custom wrapper that has different __eq__ and __hash__ +# behavior. Note that we subsequently store this weak key directly in an +# ORDINARY dictionary, since the newly constructed WeakIdKey's only use would +# be a dictionary so it would have no strong references. Ensuring that +# only live WeakIdKeys are in the map is handled by putting finalizers on the +# original key object. + + +# It is simpler to implement this with composition, but if we want to +# directly reuse the callback mechanism on weakref, we need the weakref +# and the key to be exactly the same object. Reusing the callback mechanism +# minimizes the divergence between our implementation and Lib/weakref.py +# +# NB: Prefer using this when working with weakrefs of Tensors; e.g., do +# WeakIdRef(tensor) rather than weakref.ref(tensor); it handles a number of +# easy to get wrong cases transparently for you. +class WeakIdRef(weakref.ref): + __slots__ = ["_id"] + + def __init__(self, key, callback=None) -> None: + # Unlike stock weakref, which preserves hash semantics of the + # original object but lazily defers hash calls until the first + # time the user attempts to hash the weakref, we can eagerly + # cache the id of the key as we know this is definitely the hash + # method + self._id = id(key) + super().__init__(key, callback) # type: ignore[call-arg] + + def __call__(self): + r = super().__call__() + # Special logic for Tensor PyObject resurrection + if hasattr(r, "_fix_weakref"): + r._fix_weakref() # type: ignore[union-attr] + return r + + def __hash__(self): + return self._id + + def __eq__(self, other): + # An attractive but wrong alternate implementation is to only test if + # the stored _ids match. This can lead to an ABA problem if you have: + # + # a1 = A() + # w1 = WeakIdRef(a1) + # del a1 + # a2 = A() # suppose it gets the same ID as a1 + # w2 = WeakIdRef(a2) + # print(w1 == w2) + # + # This should be False, as a1 and a2 are unrelated (and a1 is + # dead anyway) + a = self() + b = other() + if a is not None and b is not None: + return a is b + return self is other + + +# This is the same as WeakIdRef but equality is checked using hash() rather than id. +# This will be equivalent to the one above except for classes where hash is not their id. +class _WeakHashRef(weakref.ref): + __slots__ = ["_id"] + + def __init__(self, key, callback=None) -> None: + # Unlike stock weakref, which preserves hash semantics of the + # original object but lazily defers hash calls until the first + # time the user attempts to hash the weakref, we can eagerly + # cache the id of the key as we know this is definitely the hash + # method + self._id = hash(key) + super().__init__(key, callback) # type: ignore[call-arg] + + def __call__(self): + r = super().__call__() + # Special logic for Tensor PyObject resurrection + if hasattr(r, "_fix_weakref"): + r._fix_weakref() # type: ignore[union-attr] + return r + + def __hash__(self): + return self._id + + def __eq__(self, other): + # Use hash equality to determine ref equality. + # ScriptObject implements __hash__ to return the wrapped IValue's id, so + # this is equivalent to doing an identity comparison. + a = self() + b = other() + if a is not None and b is not None: + return hash(a) == hash(b) + return self is other + + +# This is directly adapted from cpython/Lib/weakref.py +class WeakIdKeyDictionary(MutableMapping): + def __init__(self, dict=None, ref_type=WeakIdRef) -> None: # CHANGED + self.data = {} + + self.ref_type = ref_type # CHANGED + + def remove(k, selfref=ref(self)) -> None: + self = selfref() + if self is not None: + if self._iterating: + self._pending_removals.append(k) + else: + try: + del self.data[k] + except KeyError: + pass + + self._remove = remove + # A list of dead weakrefs (keys to be removed) + self._pending_removals = [] + self._iterating = set() + self._dirty_len = False + if dict is not None: + self.update(dict) + + def _commit_removals(self) -> None: + # NOTE: We don't need to call this method before mutating the dict, + # because a dead weakref never compares equal to a live weakref, + # even if they happened to refer to equal objects. + # However, it means keys may already have been removed. + pop = self._pending_removals.pop + d = self.data + while True: + try: + key = pop() + except IndexError: + return + + try: + del d[key] + except KeyError: + pass + + def _scrub_removals(self) -> None: + d = self.data + self._pending_removals = [k for k in self._pending_removals if k in d] + self._dirty_len = False + + def __delitem__(self, key) -> None: + self._dirty_len = True + del self.data[self.ref_type(key)] # CHANGED + + def __getitem__(self, key): + return self.data[self.ref_type(key)] # CHANGED + + def __len__(self) -> int: + if self._dirty_len and self._pending_removals: + # self._pending_removals may still contain keys which were + # explicitly removed, we have to scrub them (see issue #21173). + self._scrub_removals() + return len(self.data) - len(self._pending_removals) + + def __repr__(self) -> str: + return f"<{self.__class__.__name__} at {id(self):#x}>" + + def __setitem__(self, key, value) -> None: + self.data[self.ref_type(key, self._remove)] = value # CHANGED + + def copy(self): + new = WeakIdKeyDictionary() + with _IterationGuard(self): + for key, value in self.data.items(): + o = key() + if o is not None: + new[o] = value + return new + + __copy__ = copy + + def __deepcopy__(self, memo): + from copy import deepcopy + + new = self.__class__() + with _IterationGuard(self): + for key, value in self.data.items(): + o = key() + if o is not None: + new[o] = deepcopy(value, memo) + return new + + def get(self, key, default=None): + return self.data.get(self.ref_type(key), default) # CHANGED + + def __contains__(self, key) -> bool: + try: + wr = self.ref_type(key) # CHANGED + except TypeError: + return False + return wr in self.data + + def items(self): + with _IterationGuard(self): + for wr, value in self.data.items(): + key = wr() + if key is not None: + yield key, value + + def keys(self): + with _IterationGuard(self): + for wr in self.data: + obj = wr() + if obj is not None: + yield obj + + __iter__ = keys + + def values(self): + with _IterationGuard(self): + for wr, value in self.data.items(): + if wr() is not None: + yield value + + def keyrefs(self): + """Return a list of weak references to the keys. + + The references are not guaranteed to be 'live' at the time + they are used, so the result of calling the references needs + to be checked before being used. This can be used to avoid + creating references that will cause the garbage collector to + keep the keys around longer than needed. + + """ + return list(self.data) + + def popitem(self): + self._dirty_len = True + while True: + key, value = self.data.popitem() + o = key() + if o is not None: + return o, value + + # pyrefly: ignore [bad-override] + def pop(self, key, *args): + self._dirty_len = True + + return self.data.pop(self.ref_type(key), *args) # CHANGED + + def setdefault(self, key, default=None): + return self.data.setdefault( + self.ref_type(key, self._remove), default + ) # CHANGED + + def update(self, dict=None, **kwargs) -> None: # type: ignore[override] + d = self.data + if dict is not None: + if not hasattr(dict, "items"): + dict = type({})(dict) + for key, value in dict.items(): + d[self.ref_type(key, self._remove)] = value # CHANGED + if kwargs: + self.update(kwargs) + + def __ior__(self, other): + self.update(other) + return self + + def __or__(self, other): + if isinstance(other, _collections_abc.Mapping): + c = self.copy() + c.update(other) + return c + return NotImplemented + + def __ror__(self, other): + if isinstance(other, _collections_abc.Mapping): + c = self.__class__() + c.update(other) + c.update(self) + return c + return NotImplemented + + # Default Mapping equality will tests keys for equality, but + # we want to test ids for equality + def __eq__(self, other): + if not isinstance(other, Mapping): + return NotImplemented + return {id(k): v for k, v in self.items()} == { + id(k): v for k, v in other.items() + } + + +# Convenience alias +WeakTensorKeyDictionary = WeakIdKeyDictionary + + +class TensorWeakRef: + """Wrapper around a weak ref of a Tensor that handles the _fix_weakref() call required when unwrapping a Tensor weakref.""" + + ref: WeakRef[Tensor] + + def __init__(self, tensor: Tensor) -> None: + if not isinstance(tensor, Tensor): + raise AssertionError(f"expected torch.Tensor, got {type(tensor)}.") + self.ref = weakref.ref(tensor) + + def __call__(self): + out = self.ref() + if out is None: + return out + if not isinstance(out, Tensor): + raise AssertionError(f"expected torch.Tensor, got {type(out)}.") + # TODO, add _fix_weakref type binding + out._fix_weakref() # type: ignore[attr-defined] + return out diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/version.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/version.py new file mode 100644 index 0000000000000000000000000000000000000000..0e30481414e080486881ca642c43937529522293 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/version.py @@ -0,0 +1,10 @@ +from typing import Optional + +__all__ = ['__version__', 'debug', 'cuda', 'git_version', 'hip', 'rocm', 'xpu'] +__version__ = '2.12.1' +debug = False +cuda: Optional[str] = '13.2' +git_version = '7269437d655783a26cba32aa88195b741ff496aa' +hip: Optional[str] = None +rocm: Optional[str] = None +xpu: Optional[str] = None diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0c77d45201c09ebc980321327c581d808982a957 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/__init__.py @@ -0,0 +1,650 @@ +# mypy: allow-untyped-defs +r""" +This package introduces support for the XPU backend, specifically tailored for +Intel GPU optimization. + +This package is lazily initialized, so you can always import it, and use +:func:`is_available()` to determine if your system supports XPU. +""" + +from __future__ import annotations + +import threading +import traceback +from functools import lru_cache +from typing import Any, NewType, TYPE_CHECKING + +import torch +import torch._C +from torch._utils import _dummy_type, _LazySeedTracker + + +if TYPE_CHECKING: + from collections.abc import Callable + + from torch.types import Device + +from ._utils import _get_device_index +from .graphs import ( + graph, + graph_pool_handle, + is_current_stream_capturing, + make_graphed_callables, + XPUGraph, +) +from .streams import Event, Stream + + +_initialized = False +_tls = threading.local() +_initialization_lock = threading.Lock() +_queued_calls: list[ + tuple[Callable[[], None], list[str]] +] = [] # don't invoke these until initialization occurs +_is_in_bad_fork = getattr(torch._C, "_xpu_isInBadFork", lambda: False) +_lazy_seed_tracker = _LazySeedTracker() +default_generators: tuple[torch._C.Generator] = () # type: ignore[assignment] + + +def _is_compiled() -> bool: + r"""Return true if compile with XPU support.""" + return torch._C._has_xpu + + +if _is_compiled(): + _XpuDeviceProperties = torch._C._XpuDeviceProperties + _exchange_device = torch._C._xpu_exchangeDevice + _maybe_exchange_device = torch._C._xpu_maybeExchangeDevice +else: + # Define dummy if PyTorch was compiled without XPU + _XpuDeviceProperties = _dummy_type("_XpuDeviceProperties") # type: ignore[assignment, misc] + + def _exchange_device(device: int) -> int: + raise NotImplementedError("PyTorch was compiled without XPU support") + + def _maybe_exchange_device(device: int) -> int: + raise NotImplementedError("PyTorch was compiled without XPU support") + + +@lru_cache(maxsize=1) +def device_count() -> int: + r"""Return the number of XPU device available.""" + if not _is_compiled(): + return 0 + return torch._C._xpu_getDeviceCount() + + +def is_available() -> bool: + r"""Return a bool indicating if XPU is currently available.""" + # This function never throws. + return device_count() > 0 + + +def is_bf16_supported(including_emulation: bool = True) -> bool: + r"""Return a bool indicating if the current XPU device supports dtype bfloat16.""" + if not is_available(): + return False + return ( + including_emulation + or torch.xpu.get_device_properties().has_bfloat16_conversions + ) + + +def is_tf32_supported() -> bool: + r"""Return a bool indicating if the current XPU device supports dtype tf32.""" + if not is_available(): + return False + # On Intel Xe architecture and newer, TF32 operations can be accelerated + # through DPAS (Dot Product Accumulate Systolic) instructions. Therefore, + # TF32 support can be determined by checking whether the device supports + # subgroup matrix multiply-accumulate operations. + return torch.xpu.get_device_properties().has_subgroup_matrix_multiply_accumulate + + +def is_initialized(): + r"""Return whether PyTorch's XPU state has been initialized.""" + return _initialized and not _is_in_bad_fork() + + +def _lazy_call(callable, **kwargs) -> None: + if is_initialized(): + callable() + else: + global _lazy_seed_tracker + if kwargs.get("seed_all", False): + _lazy_seed_tracker.queue_seed_all(callable, traceback.format_stack()) + elif kwargs.get("seed", False): + _lazy_seed_tracker.queue_seed(callable, traceback.format_stack()) + else: + # Don't store the actual traceback to avoid memory cycle + _queued_calls.append((callable, traceback.format_stack())) + + +def init() -> None: + r"""Initialize PyTorch's XPU state. + This is a Python API about lazy initialization that avoids initializing + XPU until the first time it is accessed. Does nothing if the XPU state is + already initialized. + """ + _lazy_init() + + +def _lazy_init() -> None: + global _initialized, _queued_calls + if is_initialized() or hasattr(_tls, "is_initializing"): + return + with _initialization_lock: + # This test was was protected via GIL. Double-check whether XPU has + # already been initialized. + if is_initialized(): + return + # Stop promptly upon encountering a bad fork error. + if _is_in_bad_fork(): + raise RuntimeError( + "Cannot re-initialize XPU in forked subprocess. To use XPU with " + "multiprocessing, you must use the 'spawn' start method" + ) + if not _is_compiled(): + raise AssertionError("Torch not compiled with XPU enabled") + # This function inits XPU backend and detects bad fork processing. + torch._C._xpu_init() + # Some of the queued calls may reentrantly call _lazy_init(); We need to + # just return without initializing in that case. + _tls.is_initializing = True + + _queued_calls.extend(calls for calls in _lazy_seed_tracker.get_calls() if calls) + + try: + for queued_call, orig_traceback in _queued_calls: + try: + queued_call() + except Exception as e: + msg = ( + f"XPU call failed lazily at initialization with error: {str(e)}\n\n" + f"XPU call was originally invoked at:\n\n{''.join(orig_traceback)}" + ) + raise Exception(msg) from e # noqa: TRY002 + finally: + delattr(_tls, "is_initializing") + _initialized = True + + +class _DeviceGuard: + def __init__(self, index: int) -> None: + self.idx = index + self.prev_idx = -1 + + def __enter__(self): + self.prev_idx = torch.xpu._exchange_device(self.idx) + + def __exit__(self, type: Any, value: Any, traceback: Any): + self.idx = torch.xpu._maybe_exchange_device(self.prev_idx) + return False + + +class device: + r"""Context-manager that changes the selected device. + + Args: + device (torch.device or int or str): device index to select. It's a no-op if + this argument is a negative integer or ``None``. + """ + + def __init__(self, device: Any) -> None: + self.idx = _get_device_index(device, optional=True) + self.prev_idx = -1 + + def __enter__(self): + self.prev_idx = torch.xpu._exchange_device(self.idx) + + def __exit__(self, type: Any, value: Any, traceback: Any): + self.idx = torch.xpu._maybe_exchange_device(self.prev_idx) + return False + + +class device_of(device): + r"""Context-manager that changes the current device to that of given object. + + You can use both tensors and storages as arguments. If a given object is + not allocated on a XPU, this is a no-op. + + Args: + obj (Tensor or Storage): object allocated on the selected device. + """ + + def __init__(self, obj) -> None: + idx = obj.get_device() if obj.is_xpu else -1 + super().__init__(idx) + + +def set_device(device: Device) -> None: + r"""Set the current device. + + Args: + device (torch.device or int or str): selected device. This function is a + no-op if this argument is negative. + """ + _lazy_init() + device = _get_device_index(device) + if device >= 0: + torch._C._xpu_setDevice(device) + + +def get_device_name(device: Device = None) -> str: + r"""Get the name of a device. + + Args: + device (torch.device or int or str, optional): device for which to + return the name. This function is a no-op if this argument is a + negative integer. It uses the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + + Returns: + str: the name of the device + """ + return get_device_properties(device).name + + +@lru_cache(None) +def get_device_capability(device: Device = None) -> dict[str, Any]: + r"""Get the xpu capability of a device. + + Args: + device (torch.device or int or str, optional): device for which to + return the device capability. This function is a no-op if this + argument is a negative integer. It uses the current device, given by + :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` + (default). + + Returns: + dict[str, Any]: the xpu capability dictionary of the device + """ + props = get_device_properties(device) + # Only keep attributes that are safe for dictionary serialization. + serializable_types = (int, float, bool, str, type(None), list, tuple, dict) + return { + key: value + for key in dir(props) + if not key.startswith("__") + and isinstance((value := getattr(props, key)), serializable_types) + } + + +def get_device_properties( + device: Device = None, +) -> _XpuDeviceProperties: + r"""Get the properties of a device. Returns _XpuDeviceProperties containing the following device properties: + + - ``name`` (str): device name. + - ``platform_name`` (str): SYCL platform name. + - ``vendor`` (str): device vendor. + - ``device_id`` (int): device identifier (product ID). + - ``driver_version`` (str): driver version. + - ``version`` (str): runtime version. + - ``max_compute_units`` (int): number of parallel compute units. + - ``gpu_eu_count`` (int): number of EUs (Execution Unit). + - ``max_work_group_size``: (int): maximum number of work-items permitted in a work-group. + - ``max_num_sub_groups`` (int): maximum number of sub-groups supported in a work-group. + - ``memory_clock_rate`` (int) maximum clock rate of device's global memory in MHz. + - ``memory_bus_width`` (int) maximum bus width between device and memory in bits. + - ``sub_group_sizes``: (list[int]): a list of supported sub-group sizes. + - ``local_mem_size`` (int): device local memory capacity that can be allocated per work-group in bytes. + - ``has_fp16`` (bool): whether float16 dtype is supported. + - ``has_fp64`` (bool): whether float64 dtype is supported. + - ``has_atomic64`` (bool): whether 64-bit atomic operations are supported. + - ``has_bfloat16_conversions`` (bool): whether bfloat16 conversions are supported. + - ``has_subgroup_matrix_multiply_accumulate`` (bool): whether DPAS (Dot Product Accumulate Systolic) is supported. + - ``has_subgroup_matrix_multiply_accumulate_tensor_float32`` (bool): whether DPAS with tf32 inputs is supported. + - ``has_subgroup_2d_block_io`` (bool): whether 2D block I/O for efficient matrix multiplication is supported. + - ``total_memory`` (int): device global memory in bytes. + - ``gpu_subslice_count`` (int): number of subslice. + - ``architecture`` (int): device architecture identifier (experimental). + - ``type`` (str): device type, e.g. 'cpu', 'gpu', accelerator', 'host', 'unknown'. + - ``uuid`` (Any): device UUID (Universal Unique ID), 16 bytes. + + Args: + device (torch.device or int or str): device for which to return the + properties of the device. + + Returns: + _XpuDeviceProperties: the properties of the device + """ + _lazy_init() + device = _get_device_index(device, optional=True) + return _get_device_properties(device) # type: ignore[name-defined] # noqa: F821 + + +def current_device() -> int: + r"""Return the index of a currently selected device.""" + _lazy_init() + return torch._C._xpu_getDevice() + + +def _get_device(device: int | str | torch.device) -> torch.device: + r"""Return the torch.device type object from the passed in device. + + Args: + device (torch.device or int or str): selected device. + """ + if isinstance(device, str): + device = torch.device(device) + elif isinstance(device, int): + device = torch.device("xpu", device) + return device + + +def can_device_access_peer(device: Device, peer: Device) -> bool: + r"""Query whether a device can access a peer device's memory. + + Args: + device (torch.device or int or str): selected device. + peer (torch.device or int or str): peer device to query access to. + + Returns: + bool: ``True`` if ``device`` can access ``peer``, ``False`` otherwise. + """ + _lazy_init() + device = _get_device_index(device, optional=True) + peer = _get_device_index(peer, optional=True) + return torch._C._xpu_canDeviceAccessPeer(device, peer) + + +class StreamContext: + r"""Context-manager that selects a given stream. + + All XPU kernels queued within its context will be enqueued on a selected + stream. + + Args: + Stream (Stream): selected stream. This manager is a no-op if it's + ``None``. + .. note:: Streams are per-device. + """ + + cur_stream: torch.xpu.Stream | None + + def __init__(self, stream: torch.xpu.Stream | None) -> None: + self.stream = stream + self.idx = _get_device_index(None, True) + if self.idx is None: + self.idx = -1 # pyrefly: ignore [bad-assignment] + + def __enter__(self): + cur_stream = self.stream + if cur_stream is None or self.idx == -1: + return + self.src_prev_stream = torch.xpu.current_stream(None) + + # If the stream is not on the current device, then set the current stream on the device + if self.src_prev_stream.device != cur_stream.device: + with device(cur_stream.device): + self.dst_prev_stream = torch.xpu.current_stream(cur_stream.device) + torch.xpu.set_stream(cur_stream) + + def __exit__(self, type: Any, value: Any, traceback: Any): + cur_stream = self.stream + if cur_stream is None or self.idx == -1: + return + + # Reset the stream on the original device and destination device + if self.src_prev_stream.device != cur_stream.device: + torch.xpu.set_stream(self.dst_prev_stream) + torch.xpu.set_stream(self.src_prev_stream) + + +def stream(stream: torch.xpu.Stream | None) -> StreamContext: + r"""Wrap around the Context-manager StreamContext that selects a given stream. + + Arguments: + stream (Stream): selected stream. This manager is a no-op if it's ``None``. + """ + return StreamContext(stream) + + +def _set_stream_by_id(stream_id, device_index, device_type) -> None: + r"""set stream specified by the stream id, device index and device type + + Args: stream_id (int): not visible to the user, used to assigned to the specific stream. + device_index (int): selected device index. + device_type (int): selected device type. + """ + torch._C._xpu_setStream( + stream_id=stream_id, + device_index=device_index, + device_type=device_type, + ) + + +def set_stream(stream: Stream) -> None: + r"""Set the current stream. This is a wrapper API to set the stream. + Usage of this function is discouraged in favor of the ``stream`` + context manager. + + Args: + stream (Stream): selected stream. This function is a no-op + if this argument is ``None``. + """ + if stream is None: + return + _lazy_init() + _set_stream_by_id( + stream_id=stream.stream_id, + device_index=stream.device_index, + device_type=stream.device_type, + ) + + +def current_stream(device: Device = None) -> Stream: + r"""Return the currently selected :class:`Stream` for a given device. + + Args: + device (torch.device or int, optional): selected device. Returns + the currently selected :class:`Stream` for the current device, given + by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` + (default). + """ + _lazy_init() + streamdata = torch._C._xpu_getCurrentStream( + _get_device_index(device, optional=True) + ) + return Stream( + stream_id=streamdata[0], device_index=streamdata[1], device_type=streamdata[2] + ) + + +def get_stream_from_external(data_ptr: int, device: Device = None) -> Stream: + r"""Return a :class:`Stream` from an external SYCL queue. + + This function is used to wrap SYCL queue created in other libraries in order + to facilitate data exchange and multi-library interactions. + + .. note:: This function doesn't manage the queue life-cycle, it is the user + responsibility to keep the referenced queue alive while this returned stream is + being used. The different SYCL queue pointers will result in distinct + :class:`Stream` objects, even if the SYCL queues they dereference are equivalent. + + Args: + data_ptr(int): Integer representation of the `sycl::queue*` value passed externally. + device(torch.device or int, optional): the device where the queue was originally created. + It is the user responsibility to ensure the device is specified correctly. + """ + _lazy_init() + streamdata = torch._C._xpu_getStreamFromExternal( + data_ptr, _get_device_index(device, optional=True) + ) + return Stream( + stream_id=streamdata[0], device_index=streamdata[1], device_type=streamdata[2] + ) + + +def synchronize(device: Device = None) -> None: + r"""Wait for all kernels in all streams on a XPU device to complete. + + Args: + device (torch.device or int, optional): device for which to synchronize. + It uses the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + _lazy_init() + device = _get_device_index(device, optional=True) + return torch._C._xpu_synchronize(device) + + +def get_arch_list() -> list[str]: + r"""Return list XPU architectures this library was compiled for.""" + if not _is_compiled(): + return [] + arch_flags = torch._C._xpu_getArchFlags() + if arch_flags is None: + return [] + return arch_flags.split() + + +def get_gencode_flags() -> str: + r"""Return XPU AOT(ahead-of-time) build flags this library was compiled with.""" + arch_list = get_arch_list() + if len(arch_list) == 0: + return "" + return f"-device {','.join(arch for arch in arch_list)}" + + +def _get_generator(device: torch.device) -> torch._C.Generator: + r"""Return the XPU Generator object for the given device. + + Args: + device (torch.device): selected device. + """ + idx = device.index + if idx is None: + idx = current_device() + return torch.xpu.default_generators[idx] + + +def _set_rng_state_offset( + offset: int, device: int | str | torch.device = "xpu" +) -> None: + r"""Set the random number generator state offset of the specified GPU. + + Args: + offset (int): The desired offset + device (torch.device or int, optional): The device to set the RNG state. + Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). + """ + final_device = _get_device(device) + + def cb() -> None: + default_generator = _get_generator(final_device) + default_generator.set_offset(offset) + + _lazy_call(cb) + + +def _get_rng_state_offset(device: int | str | torch.device = "xpu") -> int: + r"""Return the random number generator state offset of the specified GPU. + + Args: + device (torch.device or int, optional): The device to return the RNG state offset of. + Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). + + .. warning:: + This function eagerly initializes XPU. + """ + _lazy_init() + final_device = _get_device(device) + default_generator = _get_generator(final_device) + return default_generator.get_offset() + + +# import here to avoid circular import +from .memory import ( + change_current_allocator, + empty_cache, + get_per_process_memory_fraction, + max_memory_allocated, + max_memory_reserved, + mem_get_info, + memory_allocated, + memory_reserved, + memory_snapshot, + memory_stats, + memory_stats_as_nested_dict, + MemPool, + reset_accumulated_memory_stats, + reset_peak_memory_stats, + set_per_process_memory_fraction, + use_mem_pool, + XPUPluggableAllocator, +) +from .random import ( + get_rng_state, + get_rng_state_all, + initial_seed, + manual_seed, + manual_seed_all, + seed, + seed_all, + set_rng_state, + set_rng_state_all, +) + + +_POOL_HANDLE = NewType("_POOL_HANDLE", tuple[int, int]) +__all__ = [ + "Event", + "Stream", + "StreamContext", + "XPUPluggableAllocator", + "XPUGraph", + "can_device_access_peer", + "change_current_allocator", + "current_device", + "current_stream", + "default_generators", + "device", + "device_of", + "device_count", + "empty_cache", + "get_arch_list", + "get_device_capability", + "get_device_name", + "get_device_properties", + "get_gencode_flags", + "get_per_process_memory_fraction", + "get_rng_state", + "get_rng_state_all", + "get_stream_from_external", + "graph", + "graph_pool_handle", + "init", + "initial_seed", + "is_available", + "is_bf16_supported", + "is_current_stream_capturing", + "is_initialized", + "is_tf32_supported", + "make_graphed_callables", + "manual_seed", + "manual_seed_all", + "max_memory_allocated", + "max_memory_reserved", + "mem_get_info", + "memory_allocated", + "memory_reserved", + "memory_snapshot", + "memory_stats", + "memory_stats_as_nested_dict", + "MemPool", + "use_mem_pool", + "reset_accumulated_memory_stats", + "reset_peak_memory_stats", + "seed", + "seed_all", + "set_device", + "set_per_process_memory_fraction", + "set_rng_state", + "set_rng_state_all", + "set_stream", + "stream", + "streams", + "synchronize", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/_gpu_trace.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/_gpu_trace.py new file mode 100644 index 0000000000000000000000000000000000000000..7c3a8b9bf785bee0d46f657d0ea1754dea3c7dcc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/_gpu_trace.py @@ -0,0 +1,69 @@ +from collections.abc import Callable + +from torch._utils import CallbackRegistry + + +EventCreationCallbacks: "CallbackRegistry[int]" = CallbackRegistry("XPU event creation") +EventDeletionCallbacks: "CallbackRegistry[int]" = CallbackRegistry("XPU event deletion") +EventRecordCallbacks: "CallbackRegistry[int, int]" = CallbackRegistry( + "XPU event record" +) +EventWaitCallbacks: "CallbackRegistry[int, int]" = CallbackRegistry("XPU event wait") +MemoryAllocationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU memory allocation" +) +MemoryDeallocationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU memory deallocation" +) +StreamCreationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU stream creation" +) +DeviceSynchronizationCallbacks: "CallbackRegistry[[]]" = CallbackRegistry( + "XPU device synchronization" +) +StreamSynchronizationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU stream synchronization" +) +EventSynchronizationCallbacks: "CallbackRegistry[int]" = CallbackRegistry( + "XPU event synchronization" +) + + +def register_callback_for_event_creation(cb: Callable[[int], None]) -> None: + EventCreationCallbacks.add_callback(cb) + + +def register_callback_for_event_deletion(cb: Callable[[int], None]) -> None: + EventDeletionCallbacks.add_callback(cb) + + +def register_callback_for_event_record(cb: Callable[[int, int], None]) -> None: + EventRecordCallbacks.add_callback(cb) + + +def register_callback_for_event_wait(cb: Callable[[int, int], None]) -> None: + EventWaitCallbacks.add_callback(cb) + + +def register_callback_for_memory_allocation(cb: Callable[[int], None]) -> None: + MemoryAllocationCallbacks.add_callback(cb) + + +def register_callback_for_memory_deallocation(cb: Callable[[int], None]) -> None: + MemoryDeallocationCallbacks.add_callback(cb) + + +def register_callback_for_stream_creation(cb: Callable[[int], None]) -> None: + StreamCreationCallbacks.add_callback(cb) + + +def register_callback_for_device_synchronization(cb: Callable[[], None]) -> None: + DeviceSynchronizationCallbacks.add_callback(cb) + + +def register_callback_for_stream_synchronization(cb: Callable[[int], None]) -> None: + StreamSynchronizationCallbacks.add_callback(cb) + + +def register_callback_for_event_synchronization(cb: Callable[[int], None]) -> None: + EventSynchronizationCallbacks.add_callback(cb) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..8f738267459a2791a4a33ca4bec74800a58f0b9a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/_utils.py @@ -0,0 +1,39 @@ +from typing import Any + +import torch + +# The _get_device_index has been moved to torch.utils._get_device_index +from torch._utils import _get_device_index as _torch_get_device_index + + +def _get_device_index( + device: Any, optional: bool = False, allow_cpu: bool = False +) -> int: + r"""Get the device index from :attr:`device`, which can be a torch.device + object, a Python integer, or ``None``. + + If :attr:`device` is a torch.device object, returns the device index if it + is a XPU device. Note that for a XPU device without a specified index, + i.e., ``torch.device('xpu')``, this will return the current default XPU + device if :attr:`optional` is ``True``. If :attr:`allow_cpu` is ``True``, + CPU devices will be accepted and ``-1`` will be returned in this case. + + If :attr:`device` is a Python integer, it is returned as is. + + If :attr:`device` is ``None``, this will return the current default XPU + device if :attr:`optional` is ``True``. + """ + if isinstance(device, int): + return device + if isinstance(device, str): + device = torch.device(device) + if isinstance(device, torch.device): + if allow_cpu: + if device.type not in ["xpu", "cpu"]: + raise ValueError(f"Expected a xpu or cpu device, but got: {device}") + elif device.type != "xpu": + raise ValueError(f"Expected a xpu device, but got: {device}") + if not torch.jit.is_scripting(): + if isinstance(device, torch.xpu.device): + return device.idx + return _torch_get_device_index(device, optional, allow_cpu) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/graphs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/graphs.py new file mode 100644 index 0000000000000000000000000000000000000000..51780050f59370ec172b1ee226f24e5dda6d104f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/graphs.py @@ -0,0 +1,529 @@ +from __future__ import annotations + +import typing +from collections.abc import Callable +from typing import overload, TYPE_CHECKING, TypeAlias +from typing_extensions import ParamSpec, Self, TypeVar + +import torch +from torch import Tensor + + +if TYPE_CHECKING: + from torch.xpu import _POOL_HANDLE + +from .._utils import _dummy_type + + +__all__ = [ + "is_current_stream_capturing", + "graph_pool_handle", + "XPUGraph", + "graph", + "make_graphed_callables", +] + +_R = TypeVar("_R") +_P = ParamSpec("_P") + +if not hasattr(torch._C, "_XpuStreamBase"): + # Define dummy base classes + torch._C.__dict__["_XPUGraph"] = _dummy_type("_XPUGraph") + torch._C.__dict__["_xpu_graph_pool_handle"] = _dummy_type("_xpu_graph_pool_handle") + torch._C.__dict__["_xpu_isCurrentStreamCapturing"] = _dummy_type( + "_xpu_isCurrentStreamCapturing" + ) + +from torch._C import _xpu_graph_pool_handle, _xpu_isCurrentStreamCapturing, _XPUGraph + + +def is_current_stream_capturing() -> bool: + r"""Return True if XPU graph capture is underway on the current XPU stream, False otherwise. + + If a XPU context does not exist on the current device, returns False without initializing the context. + """ + return _xpu_isCurrentStreamCapturing() + + +def graph_pool_handle() -> _POOL_HANDLE: + r"""Return an opaque token representing the id of a graph memory pool.""" + return torch.xpu._POOL_HANDLE(_xpu_graph_pool_handle()) + + +class XPUGraph(_XPUGraph): + r"""Wrapper around a XPU graph. + + Arguments: + keep_graph (bool, optional): If ``keep_graph=False``, the + executable command graph will be instantiated on GPU at the end of + ``capture_end`` and the underlying modifiable command graph will be + destroyed. Note that the executable command graph will not be + instantiated at the end of ``capture_end`` in this + case. Instead, it will be instantiated via an explicit called + to ``instantiate`` or automatically on the first call to + ``replay`` if ``instantiate`` was not already called. Calling + ``instantiate`` manually before ``replay`` is recommended to + prevent increased latency on the first call to ``replay``. + + """ + + def __new__(cls, keep_graph: bool = False) -> Self: + return super().__new__(cls, keep_graph) + + def capture_begin(self, pool: _POOL_HANDLE | None = None) -> None: + r"""Begin capturing XPU work on the current xpu stream. + + Typically, you shouldn't call ``capture_begin`` yourself. + Use :class:`~torch.xpu.graph`, which call ``capture_begin`` internally. + + Arguments: + pool (optional): Token (returned by :func:`~torch.xpu.graph_pool_handle` or + :meth:`other_Graph_instance.pool()`) that hints this graph may share memory + with the indicated pool. + """ + super().capture_begin(pool=pool) + + def capture_end(self) -> None: + r"""End XPU graph capture on the current stream. + + After ``capture_end``, ``replay`` may be called on this instance. + + Typically, you shouldn't call ``capture_end`` yourself. + Use :class:`~torch.xpu.graph`, which call ``capture_end`` internally. + """ + super().capture_end() + + def instantiate(self) -> None: + r"""Instantiate the XPU graph. Will be called by + ``capture_end`` if ``keep_graph=False``, or by ``replay`` if + ``keep_graph=True`` and ``instantiate`` has not already been + explicitly called. Does not destroy the xpu modify command graph returned + by ``raw_xpu_graph``. + """ + super().instantiate() + + def replay(self) -> None: + r"""Replay the XPU work captured by this graph.""" + super().replay() + + def reset(self) -> None: + r"""Delete the graph currently held by this instance.""" + super().reset() + + def pool(self) -> _POOL_HANDLE: + r"""Return an opaque token representing the id of this graph's memory pool. + + This id can optionally be passed to another graph's ``capture_begin``, + which hints the other graph may share the same memory pool. + """ + return super().pool() + + def enable_debug_mode(self) -> None: + r"""Enable debugging mode for XPUGraph.debug_dump.""" + return super().enable_debug_mode() + + def debug_dump(self, debug_path: str) -> None: + r""" + Arguments: + debug_path (required): Path to dump the graph to. + + Calls a debugging function to dump the graph if the debugging is + enabled via XPUGraph.enable_debug_mode() + """ + return super().debug_dump(debug_path) + + def raw_xpu_graph(self) -> int: + r"""Returns the underlying xpuGraph_t. ``keep_graph`` must be True. + + XPU doesn't provide APIs to manipulate this object. + """ # noqa: B950 + return super().raw_xpu_graph() + + def raw_xpu_graph_exec(self) -> int: + r"""Returns the underlying xpuGraphExec_t. ``instantiate`` must have been called if ``keep_graph`` is True, or ``capture_end`` must have been called if ``keep_graph`` is False. If you call ``instantiate()`` after ``raw_xpu_graph_exec()``, the previously returned xpuGraphExec_t will be destroyed. It is your responsibility not to use this object after destruction. + + XPU doesn't provide APIs to manipulate this object. + """ # noqa: B950 + return super().raw_xpu_graph_exec() + + +class graph: + r"""Context-manager that captures XPU work into a :class:`torch.xpu.XPUGraph` object for later replay. + + Arguments: + xpu_graph (torch.xpu.XPUGraph): Graph object used for capture. + pool (optional): Opaque token (returned by a call to :func:`~torch.xpu.graph_pool_handle()` or + :meth:`other_Graph_instance.pool()`) hinting this graph's capture + may share memory from the specified pool. + stream (torch.xpu.Stream, optional): If supplied, will be set as the current stream in the context. + If not supplied, ``graph`` sets its own internal side stream as the current stream in the context. + + .. note:: + For effective memory sharing, if you pass a ``pool`` used by a previous capture and the previous capture + used an explicit ``stream`` argument, you should pass the same ``stream`` argument to this capture. + + """ # noqa: B950 + + default_capture_stream: torch.xpu.Stream | None = None + + def __init__( + self, + xpu_graph: XPUGraph, + pool: _POOL_HANDLE | None = None, + stream: torch.xpu.Stream | None = None, + ): + # Lazy-init of default_capture_stream helps avoid circular-import errors. + # Not thread safe, but graphs already have the general (explicitly documented) + # restriction that only one capture may be underway at a time in the process. + if self.__class__.default_capture_stream is None: + self.__class__.default_capture_stream = torch.xpu.Stream() + + self.pool: tuple[()] | tuple[_POOL_HANDLE] = () if pool is None else (pool,) + self.capture_stream = ( + stream if stream is not None else self.__class__.default_capture_stream + ) + if self.capture_stream is None: + raise AssertionError("capture_stream must not be None") + self.stream_ctx = self.capture_stream + self.xpu_graph = xpu_graph + + def __enter__(self) -> None: + # Free as much memory as we can for the graph + torch.xpu.synchronize() + + torch.xpu.empty_cache() + self.stream_ctx.__enter__() + + self.xpu_graph.capture_begin(*self.pool) + + def __exit__(self, *args: object) -> None: + self.xpu_graph.capture_end() + self.stream_ctx.__exit__(*args) + + +_ModuleOrCallable: TypeAlias = torch.nn.Module | Callable[..., object] + + +@overload +def make_graphed_callables( + callables: _ModuleOrCallable, + sample_args: tuple[Tensor, ...], + num_warmup_iters: int = 3, + allow_unused_input: bool = False, + pool: _POOL_HANDLE | None = None, +) -> _ModuleOrCallable: ... + + +@overload +def make_graphed_callables( + callables: tuple[_ModuleOrCallable, ...], + sample_args: tuple[tuple[Tensor, ...], ...], + num_warmup_iters: int = 3, + allow_unused_input: bool = False, + pool: _POOL_HANDLE | None = None, +) -> tuple[_ModuleOrCallable, ...]: ... + + +def make_graphed_callables( + callables: _ModuleOrCallable | tuple[_ModuleOrCallable, ...], + sample_args: tuple[Tensor, ...] | tuple[tuple[Tensor, ...], ...], + num_warmup_iters: int = 3, + allow_unused_input: bool = False, + pool: _POOL_HANDLE | None = None, +) -> _ModuleOrCallable | tuple[_ModuleOrCallable, ...]: + r"""Accept callables (functions or :class:`nn.Module`\ s) and returns graphed versions. + + Each graphed callable's forward pass runs its source callable's + forward XPU work as a XPU graph inside a single autograd node. + + The graphed callable's forward pass also appends + a backward node to the autograd graph. During backward, this node runs the + callable's backward work as a XPU graph. + + Therefore, each graphed callable should be a drop-in replacement for its source callable + in an autograd-enabled training loop. + + See :ref:`Partial-network capture` for detailed use and constraints. + + If you pass a tuple of several callables, their captures will use the same memory pool. + + Arguments: + callables (torch.nn.Module or Python function, or tuple of these): Callable or callables to graph. + If you pass a tuple of callables, their order in the tuple must be the same order they'll run + in the live workload. + sample_args (tuple of Tensors, or tuple of tuples of Tensors): Samples args for each callable. + If a single callable was passed, ``sample_args`` must be a single tuple of argument Tensors. + If a tuple of callables was passed, ``sample_args`` must be tuple of tuples of argument Tensors. + num_warmup_iters (int): The number of warmup iterations. Currently, ``DataDistributedParallel`` needs + 11 iterations for warm up. Default: ``3``. + allow_unused_input (bool): If False, specifying inputs that were not used when computing outputs + (and therefore their grad is always zero) is an error. Defaults to False. + pool (optional): Token (returned by :func:`~torch.xpu.graph_pool_handle` or + :meth:`other_Graph_instance.pool()`) that hints this graph may share memory + with the indicated pool. + .. note:: + The ``requires_grad`` state of each Tensor in ``sample_args`` must match the state + that's expected for the corresponding real input in the training loop. + + .. warning:: + This API is in beta and may change in future releases. + + .. warning:: + ``sample_args`` for each callable must contain only Tensors. Other types are not allowed. + + .. warning:: + Returned callables do not support higher order differentiation (e.g., double backward). + + .. warning:: + In any :class:`~torch.nn.Module` passed to :func:`~make_graphed_callables`, only parameters + may be trainable. Buffers must have ``requires_grad=False``. + + .. warning:: + After you pass a :class:`torch.nn.Module` through :func:`~make_graphed_callables`, + you may not add or remove any of that Module's parameters or buffers. + + .. warning:: + :class:`torch.nn.Module`\s passed to :func:`~torch.xpu.make_graphed_callables` must not have module hooks + registered on them at the time they are passed. However, registering hooks on modules *after* passing them + through :func:`~torch.xpu.make_graphed_callables` is allowed. + + .. warning:: + When running a graphed callable, you must pass its arguments in the same order and format + they appeared in that callable's ``sample_args``. + + .. warning:: + The automatic mixed precision is supported in :func:`~torch.xpu.make_graphed_callables` only with disabled + caching. The context manager `torch.amp.autocast()` must have `cache_enabled=False`. + """ + if torch.is_autocast_enabled() and torch.is_autocast_cache_enabled(): + raise RuntimeError( + "make_graphed_callables does not support the autocast caching. Please set `cache_enabled=False`." + ) + + just_one_callable = False + + _sample_args: tuple[tuple[Tensor, ...], ...] + if not isinstance(callables, tuple): + just_one_callable = True + callables = (callables,) + _sample_args = (typing.cast(tuple[Tensor, ...], sample_args),) + else: + _sample_args = typing.cast(tuple[tuple[Tensor, ...], ...], sample_args) + + flatten_sample_args = [] + + for c, args in zip(callables, _sample_args): + if isinstance(c, torch.nn.Module): + if not ( + len(c._backward_hooks) == 0 + and len(c._forward_hooks) == 0 + and len(c._forward_pre_hooks) == 0 + ): + raise RuntimeError( + "Modules must not have hooks registered at the time they are passed. However, registering hooks " + + "on modules after passing them through make_graphed_callables is allowed." + ) + if not all(b.requires_grad is False for b in c.buffers()): + raise RuntimeError( + "In any :class:`~torch.nn.Module` passed to " + + ":func:`~make_graphed_callables`, only parameters may be trainable. All buffers must have " + + "``requires_grad=False``." + ) + flatten_arg = torch.utils._pytree.arg_tree_leaves(*args) + flatten_sample_args.append(tuple(flatten_arg)) + if not all(isinstance(arg, torch.Tensor) for arg in flatten_arg): + raise TypeError( + "In the beta API, sample_args " + + "for each callable must contain only Tensors. Other types are not allowed." + ) + + # If a callable is an nn.Module, its graph's full input surface is the args the user explicitly + # passes to forward (ie, its sample_args) AND the module's parameter attributes. + per_callable_len_user_args = [len(args) for args in flatten_sample_args] + per_callable_module_params = [ + tuple(c.parameters()) if isinstance(c, torch.nn.Module) else () + for c in callables + ] + per_callable_static_input_surfaces = [ + flatten_sample_args[i] + per_callable_module_params[i] + for i in range(len(callables)) + ] + + fwd_graphs = [torch.xpu.XPUGraph() for _ in range(len(callables))] + bwd_graphs = [torch.xpu.XPUGraph() for _ in range(len(callables))] + + mempool = graph_pool_handle() if pool is None else pool + + # Warmup + torch.xpu.synchronize() + with torch.xpu.stream(torch.xpu.Stream()): + for func, args, static_input_surface in zip( + callables, _sample_args, per_callable_static_input_surfaces + ): + grad_inputs, outputs, outputs_grad = None, None, None + for _ in range(num_warmup_iters): + outputs = torch.utils._pytree.tree_leaves(func(*args)) + outputs_grad = tuple(o for o in outputs if o.requires_grad) + if len(outputs_grad) > 0: + grad_inputs = torch.autograd.grad( + outputs=outputs_grad, + inputs=tuple( + i for i in static_input_surface if i.requires_grad + ), + grad_outputs=tuple( + torch.empty_like(o) for o in outputs if o.requires_grad + ), + only_inputs=True, + allow_unused=allow_unused_input, + ) + for v in [outputs, outputs_grad, grad_inputs]: + del v + + torch.xpu.synchronize() + + # Capture forward graphs + per_callable_static_outputs = [] + per_callable_output_unflatten_spec = [] + for func, args, fwd_graph in zip(callables, _sample_args, fwd_graphs): + # each graph uses the same mempool + with torch.xpu.graph(fwd_graph, pool=mempool): + func_outputs = func(*args) + + flatten_outputs, spec = torch.utils._pytree.tree_flatten(func_outputs) + per_callable_static_outputs.append(tuple(flatten_outputs)) + per_callable_output_unflatten_spec.append(spec) + + # Capture backward graphs in reverse order + per_callable_static_grad_outputs = [] + per_callable_static_grad_inputs = [] + for static_input_surface, static_outputs, bwd_graph in zip( + reversed(per_callable_static_input_surfaces), + reversed(per_callable_static_outputs), + reversed(bwd_graphs), + ): + static_grad_outputs = tuple( + torch.empty_like(o) if o.requires_grad else None for o in static_outputs + ) + + outputs_grad = tuple(o for o in static_outputs if o.requires_grad) + grad_inputs = None + if len(outputs_grad) > 0: + with torch.xpu.graph(bwd_graph, pool=mempool): + grad_inputs = torch.autograd.grad( + outputs=outputs_grad, + inputs=tuple(i for i in static_input_surface if i.requires_grad), + grad_outputs=tuple(o for o in static_grad_outputs if o is not None), + only_inputs=True, + allow_unused=allow_unused_input, + ) + + static_grad_inputs = [] + grad_idx = 0 + for arg in static_input_surface: + if arg.requires_grad and grad_inputs is not None: + static_grad_inputs.append(grad_inputs[grad_idx]) + grad_idx += 1 + else: + static_grad_inputs.append(None) # type: ignore[arg-type] + static_grad_inputs = tuple(static_grad_inputs) # type: ignore[assignment] + + per_callable_static_grad_outputs.append(static_grad_outputs) + per_callable_static_grad_inputs.append(static_grad_inputs) + + # Reverses the most recent two lists + per_callable_static_grad_outputs.reverse() + per_callable_static_grad_inputs.reverse() + + def make_graphed_autograd_function( + fwd_graph: XPUGraph, + bwd_graph: XPUGraph, + module_params: tuple[torch.nn.Parameter, ...], + len_user_args: int, + output_unflatten_spec: torch.utils._pytree.TreeSpec, + static_input_surface: tuple[Tensor, ...], + static_outputs: tuple[Tensor, ...], + static_grad_outputs: tuple[Tensor | None, ...], + static_grad_inputs: tuple[Tensor, ...], + ) -> Callable[..., object]: + class Graphed(torch.autograd.Function): + @staticmethod + # pyrefly: ignore [bad-override] + def forward(ctx: object, *inputs: Tensor) -> tuple[Tensor, ...]: + # At this stage, only the user args may (potentially) be new tensors. + for i in range(len_user_args): + if static_input_surface[i].data_ptr() != inputs[i].data_ptr(): + static_input_surface[i].copy_(inputs[i]) + fwd_graph.replay() + if not isinstance(static_outputs, tuple): + raise RuntimeError("static_outputs must be a tuple") + return tuple(o.detach() for o in static_outputs) + + @staticmethod + @torch.autograd.function.once_differentiable + # pyrefly: ignore [bad-override] + def backward(ctx: object, *grads: Tensor) -> tuple[Tensor, ...]: + if len(grads) != len(static_grad_outputs): + raise RuntimeError( + f"Expected {len(static_grad_outputs)} gradients but got {len(grads)}" + ) + for g, grad in zip(static_grad_outputs, grads): + if g is not None: + if g.data_ptr() != grad.data_ptr(): + g.copy_(grad) + bwd_graph.replay() + + if not isinstance(static_grad_inputs, tuple): + raise RuntimeError("static_grad_inputs must be a tuple") + return tuple( + b.detach() if b is not None else b for b in static_grad_inputs + ) + + def functionalized(*user_args: object) -> object: + # Runs the new autograd function which replays the XPU graphs + flatten_user_args = torch.utils._pytree.arg_tree_leaves(*user_args) + out = Graphed.apply(*(tuple(flatten_user_args) + module_params)) + return torch.utils._pytree.tree_unflatten(out, output_unflatten_spec) + + return functionalized + + ret: list[_ModuleOrCallable] = [] + for i, func in enumerate(callables): + graphed = make_graphed_autograd_function( + fwd_graphs[i], + bwd_graphs[i], + per_callable_module_params[i], + per_callable_len_user_args[i], + per_callable_output_unflatten_spec[i], + per_callable_static_input_surfaces[i], + per_callable_static_outputs[i], + per_callable_static_grad_outputs[i], + per_callable_static_grad_inputs[i], + ) + + if isinstance(func, torch.nn.Module): + + def make_graphed_forward( + func: torch.nn.Module, + graph_training_state: bool, + graphed: Callable[_P, _R], + orig_fwd: Callable[_P, _R], + ) -> Callable[_P, _R]: + def new_fwd(*user_args: _P.args, **user_kwargs: _P.kwargs) -> _R: + if func.training == graph_training_state: + return graphed(*user_args, **user_kwargs) + else: + return orig_fwd(*user_args, **user_kwargs) + + return new_fwd + + func.forward = make_graphed_forward( + func, func.training, graphed, func.forward + ) + ret.append(func) + else: + ret.append(graphed) + + if just_one_callable: + return ret[0] + + return tuple(ret) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/memory.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/memory.py new file mode 100644 index 0000000000000000000000000000000000000000..04ca0dd9fc3977adcdc766260f555866ac1feef6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/memory.py @@ -0,0 +1,643 @@ +import collections +import contextlib +import ctypes +import pickle +import sys +from typing import Any, Literal + +import torch +from torch._utils import _augment_memory_snapshot_stack_traces, _dummy_type +from torch.types import Device + +from . import _get_device_index, _is_compiled, _lazy_init, is_initialized + + +if not _is_compiled(): + # Define dummy base classes + torch._C.__dict__["_xpu_XPUAllocator"] = _dummy_type("_xpu_XPUAllocator") + torch._C.__dict__["_XPUMemPool"] = _dummy_type("_XPUMemPool") + torch._C.__dict__["_xpu_beginAllocateCurrentThreadToPool"] = _dummy_type( + "_xpu_beginAllocateCurrentThreadToPool" + ) + torch._C.__dict__["_xpu_endAllocateToPool"] = _dummy_type("_xpu_endAllocateToPool") + torch._C.__dict__["_xpu_releasePool"] = _dummy_type("_xpu_releasePool") + + +def empty_cache() -> None: + r"""Release all unoccupied cached memory currently held by the caching + allocator so that those can be used in other XPU application. + + .. note:: + :func:`~torch.xpu.empty_cache` doesn't increase the amount of XPU + memory available for PyTorch. However, it may help reduce fragmentation + of XPU memory in certain cases. + """ + if is_initialized(): + torch._C._xpu_emptyCache() + + +def reset_peak_memory_stats(device: Device = None) -> None: + r"""Reset the "peak" stats tracked by the XPU memory allocator. + + See :func:`~torch.xpu.memory_stats` for details. Peak stats correspond to the + `"peak"` key in each individual stat dict. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + device = _get_device_index(device, optional=True) + return torch._C._xpu_resetPeakMemoryStats(device) + + +def reset_accumulated_memory_stats(device: Device = None) -> None: + r"""Reset the "accumulated" (historical) stats tracked by the XPU memory allocator. + + See :func:`~torch.xpu.memory_stats` for details. Accumulated stats correspond to + the `"allocated"` and `"freed"` keys in each individual stat dict. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + device = _get_device_index(device, optional=True) + return torch._C._xpu_resetAccumulatedMemoryStats(device) + + +def memory_stats_as_nested_dict(device: Device = None) -> dict[str, Any]: + r"""Return the result of :func:`~torch.xpu.memory_stats` as a nested dictionary.""" + if not is_initialized(): + return {} + device = _get_device_index(device, optional=True) + return torch._C._xpu_memoryStats(device) + + +def memory_stats(device: Device = None) -> dict[str, Any]: + r"""Return a dictionary of XPU memory allocator statistics for a given device. + + The return value of this function is a dictionary of statistics, each of + which is a non-negative integer. + + Core statistics: + + - ``"allocated_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``: + amount of allocated memory. + - ``"reserved_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``: + amount of reserved memory. + - ``"active_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``: + amount of active memory. + - ``"requested_bytes.{all,large_pool,small_pool}.{current,peak,allocated,freed}"``: + memory requested by client code, compare this with allocated_bytes to check if + allocation rounding adds too much overhead. + + For these core statistics, values are broken down as follows. + + Pool type: + + - ``all``: combined statistics across all memory pools. + - ``large_pool``: statistics for the large allocation pool (for size >= 1MB allocations). + - ``small_pool``: statistics for the small allocation pool (for size < 1MB allocations). + + Metric type: + + - ``current``: current value of this metric. + - ``peak``: maximum value of this metric. + - ``allocated``: historical total increase in this metric. + - ``freed``: historical total decrease in this metric. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistics for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + result = [] + + def _recurse_add_to_result(prefix: str, obj: Any) -> None: + if isinstance(obj, dict): + if len(prefix) > 0: + prefix += "." + for k, v in obj.items(): + _recurse_add_to_result(prefix + k, v) + else: + result.append((prefix, obj)) + + stats = memory_stats_as_nested_dict(device=device) + _recurse_add_to_result("", stats) + result.sort() + + return collections.OrderedDict(result) + + +def memory_allocated(device: Device = None) -> int: + r"""Return the current GPU memory occupied by tensors in bytes for a given device. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + + .. note:: + This is likely less than the amount shown in `xpu-smi` since some + unused memory can be held by the caching allocator and some context + needs to be created on GPU. + """ + return memory_stats(device=device).get("allocated_bytes.all.current", 0) + + +def max_memory_allocated(device: Device = None) -> int: + r"""Return the maximum GPU memory occupied by tensors in bytes for a given device. + + By default, this returns the peak allocated memory since the beginning of + this program. :func:`~torch.xpu.reset_peak_memory_stats` can be used to + reset the starting point in tracking this metric. For example, these two + functions can measure the peak allocated memory usage of each iteration in a + training loop. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + return memory_stats(device=device).get("allocated_bytes.all.peak", 0) + + +def memory_reserved(device: Device = None) -> int: + r"""Return the current GPU memory managed by the caching allocator in bytes for a given device. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + return memory_stats(device=device).get("reserved_bytes.all.current", 0) + + +def max_memory_reserved(device: Device = None) -> int: + r"""Return the maximum GPU memory managed by the caching allocator in bytes for a given device. + + By default, this returns the peak cached memory since the beginning of this + program. :func:`~torch.xpu.reset_peak_memory_stats` can be used to reset + the starting point in tracking this metric. For example, these two functions + can measure the peak cached memory amount of each iteration in a training + loop. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + """ + return memory_stats(device=device).get("reserved_bytes.all.peak", 0) + + +def mem_get_info(device: Device = None) -> tuple[int, int]: + r"""Return the global free and total GPU memory for a given device. + + Args: + device (torch.device or int or str, optional): selected device. Returns + statistic for the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + + Returns: + tuple[int, int]: a tuple of two integers (free_memory, total_memory) in bytes. + The first value is the free memory on the device (available across all processes and applications), + The second value is the device's total hardware memory capacity. + """ + _lazy_init() + device = _get_device_index(device, optional=True) + return torch._C._xpu_getMemoryInfo(device) + + +def get_per_process_memory_fraction(device: Device = None) -> float: + r""" + Retrieve the memory fraction currently set for a process on a given XPU device. + This fraction represents the portion of the total device memory that + the caching allocator is allowed to use. The allowed memory is calculated as: + + .. math:: \text{allowed\_memory} = \text{total\_memory} \times \text{fraction} + + Args: + device (torch.device or int or str, optional): selected device. It uses the current device, + given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default). + + Returns: + float: The memory fraction in the range 0.0 to 1.0. + """ + _lazy_init() + device = _get_device_index(device, optional=True) + return torch._C._xpu_getMemoryFraction(device) + + +def set_per_process_memory_fraction(fraction: float, device: Device = None) -> None: + r""" + Set the memory fraction for a single process on XPU device. + This function limits the amount of memory that the caching allocator can allocate + on the specified XPU device. The allowed memory is computed as: + + .. math:: \text{allowed\_memory} = \text{total\_memory} \times \text{fraction} + + If the process attempts to allocate more than this allowed memory, + an out-of-memory error will be raised by the allocator. + + Arguments: + fraction (float): Range: 0~1. Allowed memory equals total_memory * fraction. + device (torch.device or int or str, optional): selected device. It uses the current device, + given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default). + + .. note:: In general, the total available free memory is less than the total capacity. + """ + _lazy_init() + device = _get_device_index(device, optional=True) + if not isinstance(fraction, float): + raise TypeError("Invalid type for fraction argument, must be `float`") + torch._C._xpu_setMemoryFraction(fraction, device) + + +def memory_snapshot( + mempool_id: tuple[int, int] | None = None, +) -> list[dict[str, Any]]: + r""" + Return a snapshot of the XPU memory allocator state across all devices. + Provides detailed information for each memory segment managed by the allocator + including its size, owning pool, associated stream, call stack traces, and other relevant attributes. + + Arguments: + mempool_id (tuple[int, int] or None, optional): The memory pool id. If None, the default memory pool is used. + + Returns: + list[dict[str, Any]]: List of memory segments and their attributes. + """ + if not is_initialized(): + return [] + return torch._C._xpu_memorySnapshot(mempool_id)["segments"] + + +def _snapshot(device: Device = None, augment_with_fx_traces: bool = False): + """ + Capture a snapshot of the XPU memory state at the time this function is called. + + The returned snapshot is a dictionary with the following structure. + + .. code-block:: python + + class Snapshot(TypedDict): + segments: List[Segment] + device_traces: List[List[TraceEntry]] + + + class Segment(TypedDict): + # A Segment represents a contiguous memory region returned by the SYCL runtime. + # + # All reserved memory is composed of these segments. Segments are + # cached and reused by the allocator. When allocations are smaller + # than the segment, the segment may be split into multiple Blocks. + # + # Calling :func:`~torch.xpu.memory.empty_cache` releases segments that are entirely inactive. + address: int + total_size: int # total size of segment + stream: int + segment_type: Literal["small", "large"] # 'large' (>1MB) + allocated_size: int # size of memory in use + active_size: int # size of memory in use or in active_awaiting_free state + blocks: List[Block] + + + class Block(TypedDict): + # A sub-region of a Segment, either currently allocated or cached for reuse. + size: int + requested_size: int # Original requested size (may be smaller than `size`) + address: int + state: Literal[ + "active_allocated", # used by a tensor + "active_awaiting_free", # waiting for another stream synchronization, then become free + "inactive", # free for reuse + ] + frames: List[Frame] # stack trace from where the allocation occurred + + + class Frame(TypedDict): + filename: str + line: int + name: str + # Optional fields when `augment_with_fx_traces=True` and the frame + # corresponds to FX-generated code. + fx_node_op: str # FX node operation type (e.g., 'call_function', 'output') + fx_node_name: str # FX node name (e.g., 'linear', 'relu_1') + fx_original_trace: str # Original model source code stack trace + + + class TraceEntry(TypedDict): + # Trace entries are recorded only when :func:`~torch.xpu.memory._record_memory_history` is enabled. + action: Literal[ + "alloc" # memory allocated + "free_requested", # received a call to free memory + "free_completed", # memory reclaimed and reusable + "segment_alloc", # ask SYCL runtime for more memory + "segment_free", # called SYCL runtime to return memory to XPU + "segment_map", # ask SYCL runtime to map memory + "segment_unmap", # called SYCL runtime to unmap memory + "snapshot", # snapshot taken + "oom", # threw an OOM exception + ] + addr: int # not present for OOM + frames: List[Frame] + size: int + stream: int + device_free: int # only present for OOM, the amount of free memory reported by the device + + Arguments: + device (torch.device or int or str, optional): selected device. It uses the current device, + given by :func:`~torch.xpu.current_device`, if :attr:`device` is ``None`` (default). + augment_with_fx_traces (bool, optional): If True, augment stack trace frames with FX debug information + that maps generated FX code back to original model source code. This adds the FX-related + fields (fx_node_op, fx_node_name, fx_original_trace) to Frame objects. Default is ``False``. + + Returns: + The Snapshot dictionary object + """ + s = torch._C._xpu_memorySnapshot(None) + if augment_with_fx_traces: + s = _augment_memory_snapshot_stack_traces(s) # type: ignore[assignment, arg-type] + return s + + +def _dump_snapshot( + filename: str = "dump_snapshot.pickle", augment_with_fx_traces: bool = False +) -> None: + """ + Save a pickled version of the `torch.memory._snapshot()` dictionary to a file. + + This file can be opened by the interactive snapshot viewer at pytorch.org/memory_viz + + Snapshot file sizes scale with `max_entries` and stack trace depth per entry, + with several KB per entry. These can easily be in the GB range for longer running + workflows with large `max_entries`. + + Arguments: + filename (str, optional): Name of the file to create. Defaults to "dump_snapshot.pickle". + augment_with_fx_traces (bool, optional): If True, augment the snapshot with FX debug information + before dumping. This maps generated FX code stack traces back to original model + source code. Defaults to ``False``. + """ + s = _snapshot(augment_with_fx_traces=augment_with_fx_traces) + + with open(filename, "wb") as f: + pickle.dump(s, f) + + +def _record_memory_history( + enabled: Literal["state", "all"] | None = "all", + context: Literal["state", "alloc", "all"] | None = "all", + stacks: Literal["python", "all"] = "all", + max_entries: int = sys.maxsize, + clear_history: bool = False, + skip_actions: list[str] | None = None, +) -> None: + """ + Enable recording of stack traces associated with memory allocations, so you can + tell what allocated any piece of memory in :func:`~torch.xpu.memory._snapshot()`. + + In addition to keeping stack traces with each current allocation and free, + this will also enable recording of a history of all alloc/free events. + + Use :func:`~torch.xpu.memory._snapshot()` to retrieve this information, + and the tools in `_memory_viz.py` to visualize snapshots. + + Buffer behavior + --------------- + + This will store up to `max_entries` instances of `TraceEntry` when enabled. + Python trace collection defaults to `sys.maxsize`, meaning long-running + or indefinitely running jobs should set a reasonable limit to avoid excessive + memory use. Expect each entry to be several KB. + + Longer running workflows or those with smaller `max_entries` values will only + store the last accumulated `max_entries` entries, meaning new entries overwrite + older entries, reference to ring buffer behavior. + + Latency impact + -------------- + + The Python trace collection is fast (2us per trace), so you may consider + enabling this on production jobs if you anticipate ever having to debug + memory issues. + + C++ trace collection is also fast (~50ns/frame), which for many typical programs + works out to ~2us per trace, but can vary depending on stack depth. + + Arguments: + enabled (Literal["state", "all"], optional): + `None`, disable recording memory history. + `"state"`, keep information for currently allocated memory. + `"all"`, additionally keep a history of all alloc/free calls. + Defaults to "all". + context (Literal["state", "alloc", "all"], optional): + `None`, Do not record any tracebacks. + `"state"`, Record tracebacks for currently allocated memory. + `"alloc"`, additionally keep tracebacks for alloc calls. + `"all"`, additionally keep tracebacks for free calls. + Defaults to "all". + stacks (Literal["python", "all"], optional): + `"python"`, include Python, TorchScript, and inductor frames in tracebacks. + `"all"`, additionally include C++ frames. + Defaults to "all". + max_entries (int, optional): Keep a maximum of `max_entries` + alloc/free events in the recorded history recorded. + clear_history (bool, optional): Clear history when enabling, defaults to ``False``. + skip_actions (list[str], optional): List of action types to skip when recording + memory history. This can be used to reduce memory overhead by excluding + certain types of events from being recorded. Valid action types are: + + - `"alloc"`: Memory allocation events + - `"free_requested"`: Free requests (memory marked for freeing) + - `"free_completed"`: Completed free operations (memory actually freed) + - `"segment_alloc"`: Segment allocation from SYCL runtime + - `"segment_free"`: Segment freed back to XPU via SYCL runtime + - `"segment_map"`: Segment map events + - `"segment_unmap"`: Segment unmap events + - `"snapshot"`: Memory snapshot generation events + - `"oom"`: Out-of-memory exceptions + + For example, to skip recording free_requested events: + `skip_actions=["free_requested"]` + + Defaults to ``None`` (record all actions). + """ + torch._C._xpu_recordMemoryHistory( + enabled, + context, + stacks, + max_entries, + clear_history, + skip_actions if skip_actions is not None else [], + ) + + +class _XPUAllocator: + r"""Wrapper over internal XPU memory allocators.""" + + def __init__(self, allocator: torch._C._xpu_XPUAllocator): + self._allocator = allocator + + def allocator(self): + return self._allocator + + +class XPUPluggableAllocator(_XPUAllocator): + r""" + XPU memory allocator loaded dynamically from a shared library. + + This lets users provide custom allocation and free functions implemented + in a separate shared library. The allocator is registered and could become + available for use via :func:`~torch.xpu.memory.change_current_allocator`. + + Arguments: + path_to_lib_file (str): + Filesystem path to the shared library file containing the allocation + and free functions. + alloc_fn_name (str): + Name of the allocation function exported from the shared library. + The function must have the signature: + + ``void* alloc_fn(size_t size, int device, sycl::queue* queue);`` + + free_fn_name (str): + Name of the free function exported from the shared library. + The function must have the signature: + + ``void free_fn(void* ptr, size_t size, int device, sycl::queue* queue);`` + """ + + def __init__(self, path_to_lib_file: str, alloc_fn_name: str, free_fn_name: str): + allocator_lib = ctypes.CDLL(path_to_lib_file) + + alloc_fn_ptr = getattr(allocator_lib, alloc_fn_name) + free_fn_ptr = getattr(allocator_lib, free_fn_name) + + alloc_fn_addr = ctypes.cast(alloc_fn_ptr, ctypes.c_void_p).value + free_fn_addr = ctypes.cast(free_fn_ptr, ctypes.c_void_p).value + + if alloc_fn_addr is None or free_fn_addr is None: + raise RuntimeError( + "Failed to load allocator symbols from the shared library." + ) + + self._allocator = torch._C._xpu_customAllocator(alloc_fn_addr, free_fn_addr) + + +def change_current_allocator(allocator: _XPUAllocator) -> None: + r"""Change the currently used memory allocator to be the one provided. + + .. note:: + If the current allocator has already been used/initialized, this function will error. + + Arguments: + allocator (torch.xpu.memory._XPUAllocator): allocator to be set as the active one. + """ + torch._C._xpu_changeCurrentAllocator(allocator.allocator()) + + +def _get_current_allocator() -> _XPUAllocator: + r"""Return the allocator being currently used. + + Returns: + _XPUAllocator: the allocator being currently used. + """ + return _XPUAllocator(torch._C._xpu_getAllocator()) + + +class MemPool(torch._C._XPUMemPool): + r"""MemPool represents a pool of memory in a caching allocator. Currently, + it's just the ID of the pool object maintained in the XPUCachingAllocator. + + Args: + allocator(torch._C._xpu_XPUAllocator, optional): a + torch._C._xpu_XPUAllocator object that can be used to + define how memory gets allocated in the pool. If :attr:`allocator` + is ``None`` (default), memory allocation follows the default/ + current configuration of the XPUCachingAllocator. + use_on_oom(bool): a bool that indicates if this pool can be used + as a last resort if a memory allocation outside of the pool fails due + to Out Of Memory. This is ``False`` by default. + """ + + def __init__( + self, + allocator: torch._C._xpu_XPUAllocator | None = None, + use_on_oom: bool = False, + ): + super().__init__(allocator, True, use_on_oom) + + @property + def id(self) -> tuple[int, int]: + r"""Returns the ID of this pool as a tuple of two ints.""" + return super().id + + @property + def allocator(self) -> torch._C._xpu_XPUAllocator | None: + r"""Returns the allocator this MemPool routes allocations to.""" + return super().allocator + + def use_count(self) -> int: + r"""Returns the reference count of this pool.""" + return super().use_count() + + def snapshot(self): + r"""Return a snapshot of the XPU memory allocator pool state across all + devices. + + Interpreting the output of this function requires familiarity with the + memory allocator internals. + """ + snapshot = torch.xpu.memory_snapshot(self.id) + return snapshot + + +@contextlib.contextmanager +def use_mem_pool(pool: MemPool, device: "Device" = None): + r"""A context manager that routes allocations to a given pool. + + Args: + pool(torch.xpu.MemPool): a :class:`MemPool` object to be made active so that + allocations route to this pool. + device (torch.device or int, optional): selected device. Uses :class:`MemPool on + the current device, given by :func:`~torch.xpu.current_device`, + if :attr:`device` is ``None`` (default). + + .. note:: + This context manager makes only current thread's allocations route to + the given pool. If a new thread is spawned inside the context manager + (e.g. by calling backward) the allocations in that thread will not + route to the given pool. + """ + device_index = ( + torch.xpu.current_device() if device is None else _get_device_index(device) + ) + torch._C._xpu_beginAllocateCurrentThreadToPool(device_index, pool.id) + try: + yield + finally: + torch._C._xpu_endAllocateToPool(device_index, pool.id) + torch._C._xpu_releasePool(device_index, pool.id) + + +__all__ = [ + "MemPool", + "XPUPluggableAllocator", + "change_current_allocator", + "empty_cache", + "get_per_process_memory_fraction", + "max_memory_allocated", + "max_memory_reserved", + "mem_get_info", + "memory_allocated", + "memory_reserved", + "memory_snapshot", + "memory_stats", + "memory_stats_as_nested_dict", + "reset_accumulated_memory_stats", + "reset_peak_memory_stats", + "set_per_process_memory_fraction", + "use_mem_pool", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/random.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/random.py new file mode 100644 index 0000000000000000000000000000000000000000..f58e49e29d1a93954353f6f24cb0696a37e17d23 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/random.py @@ -0,0 +1,176 @@ +# mypy: allow-untyped-defs +from collections.abc import Iterable + +import torch +from torch import Tensor + +from . import _lazy_call, _lazy_init, current_device, device_count, is_initialized + + +def get_rng_state(device: int | str | torch.device = "xpu") -> Tensor: + r"""Return the random number generator state of the specified GPU as a ByteTensor. + + Args: + device (torch.device or int, optional): The device to return the RNG state of. + Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). + + .. warning:: + This function eagerly initializes XPU. + """ + _lazy_init() + if isinstance(device, str): + device = torch.device(device) + elif isinstance(device, int): + device = torch.device("xpu", device) + idx = device.index + if idx is None: + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + return default_generator.get_state() + + +def get_rng_state_all() -> list[Tensor]: + r"""Return a list of ByteTensor representing the random number states of all devices.""" + results = [get_rng_state(i) for i in range(device_count())] + return results + + +def set_rng_state(new_state: Tensor, device: int | str | torch.device = "xpu") -> None: + r"""Set the random number generator state of the specified GPU. + + Args: + new_state (torch.ByteTensor): The desired state + device (torch.device or int, optional): The device to set the RNG state. + Default: ``'xpu'`` (i.e., ``torch.device('xpu')``, the current XPU device). + """ + if not is_initialized(): + with torch._C._DisableFuncTorch(): + new_state = new_state.clone(memory_format=torch.contiguous_format) + + if isinstance(device, str): + device = torch.device(device) + elif isinstance(device, int): + device = torch.device("xpu", device) + + def cb() -> None: + idx = device.index + if idx is None: + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + default_generator.set_state(new_state) + + _lazy_call(cb) + + +def set_rng_state_all(new_states: Iterable[Tensor]) -> None: + r"""Set the random number generator state of all devices. + + Args: + new_states (Iterable of torch.ByteTensor): The desired state for each device. + """ + for i, state in enumerate(new_states): + set_rng_state(state, i) + + +def manual_seed(seed: int) -> None: + r"""Set the seed for generating random numbers for the current GPU. + + It's safe to call this function if XPU is not available; in that case, it is silently ignored. + + Args: + seed (int): The desired seed. + + .. warning:: + If you are working with a multi-GPU model, this function is insufficient + to get determinism. To seed all GPUs, use :func:`manual_seed_all`. + """ + seed = int(seed) + + def cb() -> None: + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + default_generator.manual_seed(seed) + + _lazy_call(cb, seed=True) + + +def manual_seed_all(seed: int) -> None: + r"""Set the seed for generating random numbers on all GPUs. + + It's safe to call this function if XPU is not available; in that case, it is silently ignored. + + Args: + seed (int): The desired seed. + """ + seed = int(seed) + + def cb() -> None: + for i in range(device_count()): + default_generator = torch.xpu.default_generators[i] + default_generator.manual_seed(seed) + + _lazy_call(cb, seed_all=True) + + +def seed() -> None: + r"""Set the seed for generating random numbers to a random number for the current GPU. + + It's safe to call this function if XPU is not available; in that case, it is silently ignored. + + .. warning:: + If you are working with a multi-GPU model, this function will only initialize + the seed on one GPU. To initialize all GPUs, use :func:`seed_all`. + """ + + def cb() -> None: + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + default_generator.seed() + + _lazy_call(cb) + + +def seed_all() -> None: + r"""Set the seed for generating random numbers to a random number on all GPUs. + + It's safe to call this function if XPU is not available; in that case, it is silently ignored. + """ + + def cb() -> None: + random_seed = 0 + seeded = False + for i in range(device_count()): + default_generator = torch.xpu.default_generators[i] + if not seeded: + default_generator.seed() + random_seed = default_generator.initial_seed() + seeded = True + else: + default_generator.manual_seed(random_seed) + + _lazy_call(cb) + + +def initial_seed() -> int: + r"""Return the current random seed of the current GPU. + + .. warning:: + This function eagerly initializes XPU. + """ + _lazy_init() + idx = current_device() + default_generator = torch.xpu.default_generators[idx] + return default_generator.initial_seed() + + +__all__ = [ + "get_rng_state", + "get_rng_state_all", + "set_rng_state", + "set_rng_state_all", + "manual_seed", + "manual_seed_all", + "seed", + "seed_all", + "initial_seed", +] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/streams.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/streams.py new file mode 100644 index 0000000000000000000000000000000000000000..3e5068ee3b44b4427c21c15ac82cbdbfcb456372 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torch/xpu/streams.py @@ -0,0 +1,181 @@ +# mypy: allow-untyped-defs +# pylint: disable=useless-parent-delegation +from __future__ import annotations + +import ctypes + +import torch +from torch._utils import _dummy_type + + +if not hasattr(torch._C, "_XpuStreamBase"): + # Define dummy base classes + torch._C.__dict__["_XpuStreamBase"] = _dummy_type("_XpuStreamBase") + torch._C.__dict__["_XpuEventBase"] = _dummy_type("_XpuEventBase") + + +class Stream(torch._C._XpuStreamBase): + r"""Wrapper around a XPU stream. + + A XPU stream is a linear sequence of execution that belongs to a specific + device, independent from other streams. It supports with statement as a + context manager to ensure the operators within the with block are running + on the corresponding stream. + + Args: + device(torch.device or int, optional): a device on which to allocate + the stream. If :attr:`device` is ``None`` (default) or a negative + integer, this will use the current device. + priority(int, optional): priority of the stream, which can be positive, 0, or negative. + A lower number indicates a higher priority. By default, the priority is set to 0. + If the value falls outside of the allowed priority range, it will automatically be + mapped to the nearest valid priority (lowest for large positive numbers or + highest for large negative numbers). + """ + + def __new__(cls, device=None, priority=0, **kwargs): + # setting device manager is expensive, so we avoid it unless necessary + if device is None or ("stream_id" in kwargs and "device_index" in kwargs): + return super().__new__(cls, priority=priority, **kwargs) + else: + with torch.xpu.device(device): + return super().__new__(cls, priority=priority, **kwargs) + + def wait_event(self, event: Event | torch.Event) -> None: + r"""Make all future work submitted to the stream wait for an event. + + Args: + event (Event, torch.Event): an event to wait for. + """ + event.wait(self) + + def wait_stream(self, stream: Stream | torch.Stream) -> None: + r"""Synchronize with another stream. + + All future work submitted to this stream will wait until all kernels + submitted to a given stream at the time of call complete. + + Args: + stream (Stream, torch.Stream): a stream to synchronize. + """ + self.wait_event(stream.record_event()) + + def record_event(self, event: Event | torch.Event | None = None): + r"""Record an event. + + Args: + event (Event, torch.Event, optional): event to record. If not given, a new one + will be allocated. + + Returns: + Recorded event. + """ + if event is None: + event = Event() + event.record(self) + return event + + def query(self) -> bool: + r"""Check if all the work submitted has been completed. + + Returns: + A boolean indicating if all kernels in this stream are completed. + """ + return super().query() + + def synchronize(self) -> None: + r"""Wait for all the kernels in this stream to complete.""" + super().synchronize() + + @property + def _as_parameter_(self): + return ctypes.c_void_p(self.sycl_queue) + + def __eq__(self, o): + if isinstance(o, Stream): + return super().__eq__(o) + return False + + def __hash__(self): + return hash((self.sycl_queue, self.device)) + + def __repr__(self) -> str: + return f"torch.xpu.Stream(device={self.device} sycl_queue={self.sycl_queue:#x})" + + +class Event(torch._C._XpuEventBase): + r"""Wrapper around a XPU event. + + XPU events are synchronization markers that can be used to monitor the + device's progress, and to synchronize XPU streams. + + The underlying XPU events are lazily initialized when the event is first + recorded. After creation, only streams on the same device may record the + event. However, streams on any device can wait on the event. + + Args: + enable_timing (bool, optional): indicates if the event should measure time + (default: ``False``) + """ + + def __new__(cls, enable_timing=False): + return super().__new__(cls, enable_timing=enable_timing) + + def record(self, stream: Stream | torch.Stream | None = None) -> None: + r"""Record the event in a given stream. + + Args: + stream (Stream, torch.Stream, optional): Uses ``torch.xpu.current_stream()`` if no stream is specified. + The stream's device must match the event's device. + """ + if stream is None: + stream = torch.xpu.current_stream() + super().record(stream) + + def wait(self, stream: Stream | torch.Stream | None = None) -> None: + r"""Make all future work submitted to the given stream wait for this event. + + Args: + stream (Stream, torch.Stream, optional): Uses ``torch.xpu.current_stream()`` if no stream is specified. + """ + if stream is None: + stream = torch.xpu.current_stream() + super().wait(stream) + + def query(self) -> bool: + r"""Check if all work currently captured by event has completed. + + Returns: + A boolean indicating if all work currently captured by event has + completed. + """ + return super().query() + + def elapsed_time(self, end_event: Event): + r"""Return the time elapsed. + + Time reported in milliseconds after the event was recorded and + before the end_event was recorded. + + Args: + end_event (Event): the end event. + """ + return super().elapsed_time(end_event) + + def synchronize(self) -> None: + r"""Wait for the event to complete. + + Waits until the completion of all work currently captured in this event. + This prevents the CPU thread from proceeding until the event completes. + """ + super().synchronize() + + @property + def _as_parameter_(self): + return ctypes.c_void_p(self.sycl_event) + + def __repr__(self) -> str: + if self.sycl_event: + return f"torch.xpu.Event(sycl_event={self.sycl_event:#x})" + else: + return "torch.xpu.Event(uninitialized)" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2d5dbf0667a022caa07ec30bb10db5b4f83159dd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/__init__.py @@ -0,0 +1,10 @@ +"""torchgen + +This module contains codegeneration utilities for PyTorch. It is used to +build PyTorch from source, but may also be used for out-of-tree projects +that extend PyTorch. + +Note well that we provide no BC guarantees for torchgen. If you're interested +in using torchgen and want the PyTorch team to be aware, please reach out +on GitHub. +""" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/aoti/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/aoti/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/aoti/fallback_ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/aoti/fallback_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..84d413d426d198e6751d63a54084f361838d9d52 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/aoti/fallback_ops.py @@ -0,0 +1,206 @@ +# Be extra careful when you edit this file, because it affects AOTInductor ABI compatibility. See +# https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436 +# for details. +# +# The inductor_fallback_ops list is based on the fallback ops from torch/_inductor/lowering.py. +# +# Generally speaking, it is ok to add a new op to the list, but you need to run +# `python torchgen/gen.py --update-aoti-c-shim` in order to regenerate C shim header files. +# But it is NOT ok to remove an existing fallback op from the list, since that will break +# some existing AOTInductor-compiled models. +# +# A fallback op version defaults to 1. If you want to extend an existing fallback op by adding +# a new argument with a default value, while it is fine in the Python world, it will be BC-breaking +# when generating C shim. Thus you need to bump up the version number of that fallback op by +# updating the entry in the inductor_fallback_ops list, adding a new version number with a list +# of new arguments, and then run `python torchgen/gen.py --update-aoti-c-shim` to regenerate. + +inductor_fallback_ops: dict[str, dict[str, list[str]]] = { + "aten._adaptive_avg_pool2d_backward.default": {}, + "aten._adaptive_avg_pool2d.default": {}, + "aten._adaptive_avg_pool3d_backward.default": {}, + "aten._adaptive_avg_pool3d.default": {}, + "aten._addmm_activation.default": {}, + "aten._cdist_backward.default": {}, + "aten._cdist_forward.default": {}, + "aten._cudnn_rnn.default": {}, + "aten._dyn_quant_matmul_4bit.default": {}, + "aten._dyn_quant_pack_4bit_weight.default": {}, + "aten._efficient_attention_backward.default": {}, + "aten._efficient_attention_forward.default": {}, + "aten._efficientzerotensor.default": {}, + "aten._embedding_bag_dense_backward.default": {}, + "aten._embedding_bag_forward_only.default": {}, + "aten._embedding_bag_per_sample_weights_backward.default": {}, + "aten._embedding_bag.default": {}, + "aten._fft_c2c.default": {}, + "aten._fft_r2c.default": {}, + "aten._flash_attention_backward.default": {}, + "aten._flash_attention_forward.default": {"v2": ["block_table", "num_splits"]}, + "aten._flash_attention_forward_no_dropout_inplace.default": {"v2": ["num_splits"]}, + "aten._flash_attention_forward.quantized": {}, + "aten._fused_moving_avg_obs_fq_helper_functional.default": {}, + "aten._fused_moving_avg_obs_fq_helper.default": {}, + "aten._fused_rms_norm.default": {}, + "aten._grouped_mm.default": {}, + "aten._histogramdd_from_bin_cts.default": {}, + "aten._int_mm.out": {}, + "aten._pdist_backward.default": {}, + "aten._pdist_forward.default": {}, + "aten._scaled_dot_product_attention_math_for_mps.default": {"v2": ["enable_gqa"]}, + "aten._scaled_dot_product_cudnn_attention_backward.default": {}, + "aten._scaled_dot_product_cudnn_attention.default": {}, + "aten._scaled_dot_product_efficient_attention_backward.default": {}, + "aten._scaled_dot_product_efficient_attention.default": {}, + "aten._scaled_dot_product_flash_attention_backward.default": {}, + "aten._scaled_dot_product_flash_attention_for_cpu_backward.default": {}, + "aten._scaled_dot_product_flash_attention_for_cpu.default": {}, + "aten._scaled_dot_product_flash_attention.default": {}, + "aten._scaled_dot_product_flash_attention.quantized": {}, + "aten._scaled_dot_product_fused_attention_overrideable_backward.default": {}, + "aten._scaled_dot_product_fused_attention_overrideable.default": {}, + "aten._scaled_mm.default": {}, + "aten._scaled_grouped_mm.default": {}, + "aten._scaled_mm.out": {}, + "aten._segment_reduce_backward.default": {}, + "aten._thnn_fused_lstm_cell.default": {}, + "aten._to_sparse.default": {}, + "aten._trilinear.default": {}, + "aten._weight_int4pack_mm.default": {}, + "aten._weight_int8pack_mm.default": {}, + "aten.abs.default": {}, + "aten.adaptive_max_pool2d_backward.default": {}, + "aten.adaptive_max_pool2d.default": {}, + "aten.adaptive_max_pool3d_backward.default": {}, + "aten.adaptive_max_pool3d.default": {}, + "aten.add.Scalar": {}, + "aten.add.Tensor": {}, + "aten.addbmm.default": {}, + "aten.addmm.out": {}, + "aten.addmv.default": {}, + "aten.angle.default": {}, + "aten.avg_pool2d_backward.default": {}, + "aten.avg_pool2d.default": {}, + "aten.avg_pool3d_backward.default": {}, + "aten.avg_pool3d.default": {}, + "aten.baddbmm.out": {}, + "aten.bernoulli_.float": {}, + "aten.bernoulli_.Tensor": {}, + "aten.bmm.out": {}, + "aten.bucketize.Tensor": {}, + "aten.cat.default": {}, + "aten.cholesky_inverse.default": {}, + "aten.cholesky_solve.default": {}, + "aten.convolution_backward.default": {}, + "aten.convolution.default": {}, + "aten.cummax.default": {}, + "aten.cummin.default": {}, + "aten.cumprod.default": {}, + "aten.cumsum.default": {}, + "aten.exponential.default": {}, + "aten.fill_.Scalar": {}, + "aten.fractional_max_pool2d_backward.default": {}, + "aten.fractional_max_pool2d.default": {}, + "aten.fractional_max_pool3d_backward.default": {}, + "aten.fractional_max_pool3d.default": {}, + "aten.gcd.default": {}, + "aten.geqrf.default": {}, + "aten.grid_sampler_2d_backward.default": {}, + "aten.hann_window.default": {}, + "aten.histc.default": {}, + "aten.histogram.bin_ct": {}, + "aten.index_put.default": {}, + "aten.index_reduce.default": {}, + "aten.index.Tensor": {}, + "aten.kthvalue.default": {}, + "aten.logcumsumexp.default": {}, + "aten.lu_unpack.default": {}, + "aten.masked_scatter_backward.default": {}, + "aten.masked_scatter.default": {}, + "aten.masked_select.default": {}, + "aten.max_pool2d_with_indices_backward.default": {}, + "aten.max_pool2d_with_indices.default": {}, + "aten.max_pool3d_with_indices_backward.default": {}, + "aten.max_pool3d_with_indices.default": {}, + "aten.max_unpool2d.default": {}, + "aten.max_unpool3d.default": {}, + "aten.median.default": {}, + "aten.mm.out": {}, + "aten.mm.dtype_out": {}, + "aten.mode.default": {}, + "aten.mul.Scalar": {}, + "aten.mul.Tensor": {}, + "aten.nanmedian.default": {}, + "aten.narrow.default": {}, + "aten.native_dropout.default": {}, + "aten.nonzero.default": {}, + "aten.nonzero_static.default": {}, + "aten.normal_functional.default": {}, + "aten.ormqr.default": {}, + "aten.pad.default": {}, + "aten.permute.default": {}, + "aten.polar.default": {}, + "aten.pow.Scalar": {}, + "aten.pow.Tensor_Scalar": {}, + "aten.pow.Tensor_Tensor": {}, + "aten.rand.default": {}, + "aten.rand.generator": {}, + "aten.randint.default": {}, + "aten.randint.generator": {}, + "aten.randint.low_out": {}, + "aten.randint.low": {}, + "aten.randn.default": {}, + "aten.randn.generator": {}, + "aten.randperm.default": {}, + "aten.rand_like.default": {}, + "aten.rand_like.generator": {}, + "aten.randint_like.default": {}, + "aten.randint_like.low_dtype": {}, + "aten.randn_like.default": {}, + "aten.randn_like.generator": {}, + "aten.repeat_interleave.Tensor": {}, + "aten.replication_pad1d_backward.default": {}, + "aten.replication_pad2d_backward.default": {}, + "aten.reshape.default": {}, + "aten.resize_.default": {}, + "aten.resize_as_.default": {}, + "aten.scatter_reduce.two_out": {}, + "aten.scatter.src_out": {}, + "aten.scatter.value_out": {}, + "aten.searchsorted.Scalar": {}, + "aten.searchsorted.Tensor": {}, + "aten.segment_reduce.default": {}, + "aten.set_.source_Tensor": {}, + "aten.slice.Tensor": {}, + "aten.soft_margin_loss_backward.default": {}, + "aten.sort.default": {}, + "aten.sort.stable": {}, + "aten.squeeze.dim": {}, + "aten.to_sparse.default": {}, + "aten.topk.default": {}, + "aten.triangular_solve.default": {}, + "aten.uniform.default": {}, + "aten.upsample_bicubic2d_backward.default": {}, + "aten.upsample_linear1d_backward.default": {}, + "aten.upsample_trilinear3d_backward.default": {}, + "aten.view_as_complex.default": {}, + "aten.view_as_real.default": {}, + "aten.view.dtype": {}, + "aten._weight_int4pack_mm_with_scales_and_zeros.default": {}, +} + +# `python torchgen/gen.py --update-aoti-c-shim` will automatically generate +# c_shim_aten.{h/cpp} based on the list below. +# Operators in this list are intended to be used in torch/csrc/stable/ops.h +# Unlike other c_shims, operators in this file do not bypass the dispatcher. +# The same BC rules apply as inductor_fallback_ops. +aten_shimified_ops: dict[str, dict[str, list[str]]] = { + "aten.fill_.Scalar": {}, + "aten.pad.default": {}, + "aten.narrow.default": {}, + "aten.amax.default": {}, + "aten.new_empty.default": {}, + "aten.new_zeros.default": {}, + "aten.full.default": {}, + "aten.subtract.Tensor": {}, +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/autograd.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/autograd.py new file mode 100644 index 0000000000000000000000000000000000000000..23a5cd6b1b61e7c8f7ecd1abca310216583deea0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/autograd.py @@ -0,0 +1,892 @@ +from __future__ import annotations + +import re +from dataclasses import dataclass +from typing import cast, TYPE_CHECKING + +from torchgen import local +from torchgen.api import cpp +from torchgen.api.types import BaseCType, Binding, NamedCType, tensorListT +from torchgen.model import ( + BaseTy, + BaseType, + FunctionSchema, + ListType, + NativeFunction, + NativeFunctionsViewGroup, + SchemaKind, + Type, +) +from torchgen.utils import IDENT_REGEX + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# Represents a saved attribute involved in backward calculation. +# Note that it can be a derived property of an input argument, e.g.: +# we could save `other.scalar_type()` instead of the entire `other` tensor. +@dataclass(frozen=True) +class SavedAttribute: + # The NamedCType holds the updated name and cpp type of the attribute + # for the name, Suffix is appended if it's derived property, e.g.: `other_scalar_type` + nctype: NamedCType + + # The expression to read the derived property at save time, e.g.: + # `other.scalar_type()`. + expr: str + + +# Represents a backward formula that calculates derivatives for one +# or more tensors. +@dataclass(frozen=True) +class Derivative: + # The formula string (legit C++ expression). + # Note that expressions against input arguments have been replaced with the + # corresponding saved attributes. + # E.g.: + # raw formula: `mul_tensor_backward(grad, self, other.scalar_type())` + # here: `mul_tensor_backward(grad, self, other_scalar_type)` + formula: str + + # The formula string before input argument replacement + original_formula: str + + # Names of the arguments for which this formula calculates derivatives. + var_names: tuple[str, ...] + + # Saved inputs that are referenced by the formula. + saved_inputs: tuple[SavedAttribute, ...] + + # Saved outputs that are referenced by the formula. + saved_outputs: tuple[SavedAttribute, ...] + + # Gradients that are referenced by name in the formula. + named_gradients: set[str] + + +# Represents a forward formula that calculates forward derivatives +# for one tensor. +@dataclass(frozen=True) +class ForwardDerivative: + # The formula string (legit C++ expression). + # Note that special keywords such as "linear" or "element_wise" have been + # replaced by the automatically generated formula. + formula: str + + # Name of the output arguments for which this formula calculates forward + # derivatives + var_names: tuple[str, ...] + + # Type of the output arguments for which this formula calculates forward + # derivatives + var_types: tuple[Type, ...] + + # Inputs for which the forward derivatives are required for this formula + required_inputs_fw_grad: tuple[str, ...] | None + + # Inputs for which the primal is required for this formula + required_inputs_primal: tuple[str, ...] | None + + # Flag to specify if this formula requires the original value of self + # This is only used by inplace operations + required_original_self_value: bool + + # If this formula is specified in derivatives.yaml or if we are reusing the + # out of place formula for inplace + is_reusing_outplace_formula: bool + + +# Represents differentiability info for a NativeFunction. +@dataclass(frozen=True) +class DifferentiabilityInfo: + # The base name read from derivatives.yaml. + name: str + + # The matching native function. + # + # There can be multiple NativeFunction having the same base name: + # - different overloads with different types of input arguments; + # - in-place/out/functional variants of the same function; + # + # We first use the schema string (under the 'name' key) in derivatives.yaml + # to find the NativeFunction having the same schema string. + # Then we find the in-place/out/functional variants of the matching function. + # Among these variants, we choose the one having the same name as the + # derivatives.yaml entry. If there is no exact match, then we choose the + # in-place variant. + # TODO: maybe the logic to search for all variants is no longer necessary? + func: NativeFunction + + # The name of the generated autograd function. + # It's set only if we will calculate a derivative, i.e. + # 'args_with_derivatives' is not empty. + op: str | None + + # The derivatives formulae for this function. + # Note that the length of this sequence is the number of differentiable inputs + derivatives: Sequence[Derivative] + + # The forward derivatives formulae for this function. + # Note that the length of this sequence is the number of differentiable outputs + forward_derivatives: Sequence[ForwardDerivative] + + # The union of 'saved_inputs' of all 'derivatives'. + all_saved_inputs: Sequence[SavedAttribute] + + # The union of 'saved_outputs' of all 'derivatives'. + all_saved_outputs: Sequence[SavedAttribute] + + # All named gradients that are available for use, in the same + # order as in the grads vector. + available_named_gradients: Sequence[str] + + # The named gradients that are used in any of the derivatives. + # Invariant: all(name in available_named_gradients for name in used_named_gradients) + used_named_gradients: set[str] + + # The function's input arguments for which it calculates derivatives. + # It's the union of 'var_names' of all 'derivatives', sorted by the + # argument order in the function schema. + args_with_derivatives: Sequence[Binding] + + # Names of arguments whose derivative formula is 'non_differentiable'. + non_differentiable_arg_names: Sequence[str] + + # Raw data read from derivatives.yaml. + output_differentiability: list[bool] | None + + # output_differentiability in derivatives.yaml can be a list of + # conditions that express if the output is differentiable. In this case, + # the number of conditions must match the number of outputs + # (NB: we only support one condition right now). + # output_differentiability gets populated with True for each condition, + # while output_differentiability_conditions gets populated with the conditions + output_differentiability_conditions: list[str] | None + + @property + def has_derivatives(self) -> bool: + return len(self.args_with_derivatives) > 0 + + # Generates a new DifferentiabilityInfo using the exact same set of derivative information, + # but with a new operator name. + # This is used when generating "copy" variants of view ops, + # which are able to use the exact same derivative formula as the original view op + # See Note [Codegen'd {view}_copy Operators] + def create_view_copy_from_view_derivative( + self, g: NativeFunctionsViewGroup + ) -> DifferentiabilityInfo | None: + if g.view_copy is None: + return None + f = g.view_copy + + name_split_by_period = self.name.split(".", maxsplit=2) + # Append a "_copy" to the base name of the operator (but keep the overload name the same) + view_copy_name = f"{name_split_by_period[0]}_copy." + ".".join( + name_split_by_period[1:] + ) + view_copy_op_name = None if self.op is None else f"{self.op}_copy" + + return DifferentiabilityInfo( + # Use the "_copy" version of name/func/op + name=view_copy_name, + func=f, + op=view_copy_op_name, + # But keep all derivative info the same + derivatives=self.derivatives, + forward_derivatives=self.forward_derivatives, + all_saved_inputs=self.all_saved_inputs, + all_saved_outputs=self.all_saved_outputs, + available_named_gradients=self.available_named_gradients, + used_named_gradients=self.used_named_gradients, + args_with_derivatives=self.args_with_derivatives, + non_differentiable_arg_names=self.non_differentiable_arg_names, + output_differentiability=self.output_differentiability, + output_differentiability_conditions=self.output_differentiability_conditions, + ) + + +def uses_ident(info: DifferentiabilityInfo | None, ident: str) -> bool: + if info is None: + return False + for derivative in info.derivatives: + formula = derivative.formula + if re.search(IDENT_REGEX.format(ident), formula): + return True + return False + + +def uses_retain_variables(info: DifferentiabilityInfo | None) -> bool: + return uses_ident(info, "retain_variables") + + +def uses_single_grad(info: DifferentiabilityInfo | None) -> bool: + return uses_ident(info, "grad") + + +# Represents a differentiable `Argument`. +# How is it different from the `Argument` type? +# - It's processed Arguments which are differentiable and only used in the +# context of the autograd codegen; +# - It can represent SelfArgument or regular Argument but not TensorOptionsArgument; +@dataclass(frozen=True) +class DifferentiableInput: + name: str + type: Type + + # TODO: only to keep it byte-for-byte compatible with the old codegen, should remove. + cpp_type: str + + +# Represents a differentiable `Return`. +# How it it different from the `Return` type? +# - The name in `Return` is optional. Here it is always populated using the same +# `cpp.return_names()` method. +# TODO: some cpp naming logic (e.g. resolving name conflict) might be irrelevant? +# - It's processed Returns which are differentiable, in compliance with the +# `output_differentiability` field defined in derivatives.yaml (if specified), +# and are only used in the context of the autograd codegen; +@dataclass(frozen=True) +class DifferentiableOutput: + name: str + type: Type + + # TODO: only to keep it byte-for-byte compatible with the old codegen, should remove. + cpp_type: str + + +@dataclass(frozen=True) +class NativeFunctionWithDifferentiabilityInfo: + func: NativeFunction + info: dict[str, DifferentiabilityInfo] | None + fw_derivatives: dict[str, Sequence[ForwardDerivative]] | None + + +# TODO: Update comment below since it is out of date. +def dispatch_strategy(fn: NativeFunctionWithDifferentiabilityInfo) -> str: + """How are we going to call the underlying implementation of a + declaration? There are two strategies: + - use_derived: we want to call the implementation on CPUDoubleType + (or a similar, derived Type instance). Because these derived + instances deal in Tensors, not Variables (it's a completely different + object, so it doesn't dispatch back to VariableType), code on + this dispatch path needs to wrap/unwrap tensors. If the + derived implementation takes and returns tensors, the + implementation is usually differentiable (although we also use + the derived dispatch path for non-differentiable functions + that we still want to dispatch on the derived Type instance; + e.g., size()) + - use_type: we want to call the implementation on Type, because + it is implemented concretely, and the functions it invokes will + get dispatched back to VariableType (which will ensure that they + are differentiable.) + """ + # fn is derived as long as any of its per-key differentiability infos + # has_derivatives. dispatch_strategy() is used to guard generation of fns in VariableType + # and ADInplaceOrViewType. We want to generate these functions as long as a + # derivative is defined for ANY dispatch key. + if fn.func.is_abstract or ( + fn.info is not None and any(info.has_derivatives for info in fn.info.values()) + ): + # If the function is abstract (not implemented on at::Type), we must + # call the implementation on the derived type with unpacked tensors. + + # If the function has a derivative specified and is concrete, we could + # call either implementation. We prefer the calling the derived + # type's implementation with unpacked tensors because it is more + # performant in some cases: any internal calls to other ATen functions + # won't have the history tracked. + + # If the function has a type dispatched argument (i.e. is a factory), + # we prefer calling the derived type's implementation both because it is + # more performant and to ensure factory functions return tensors with _version + # of 0 (probably not strictly necessary, but nice to have to keeps versions simple + # to understand. + + return "use_derived" + else: + # If the function is concrete (we don't have to override it) and we + # didn't declare it in derivatives.yaml, we'll assume that it is + # actually implemented out of differentiable functions. (This + # assumption might not hold, but then you'll see gradcheck fail.) + return "use_type" + + +def is_foreach_func(f: NativeFunction) -> bool: + return f.func.name.name.base.startswith("_foreach_") + + +# note(crcrpar): Most foreach functions can reference an out-place `torch` function whose schema kind +# is functional for their backward derivatives (and forward derivatives in the future), i.e., +# they would find such one in `functional_info_by_signature`. There however are some exceptions: +_foreach_with_inplace_ref = {"_foreach_zero_"} +_foreach_with_tensor_overload = { + "_foreach_add.Tensor", + "_foreach_mul.Tensor", + "_foreach_div.Tensor", +} +# The following do not support the alpha kwarg, which the nonforeach versions support. +_skip_argument_len_check = { + "_foreach_add.Scalar", + "_foreach_add_.Scalar", + "_foreach_add.ScalarList", + "_foreach_add_.ScalarList", + "_foreach_sub.Scalar", + "_foreach_sub_.Scalar", + "_foreach_sub.ScalarList", + "_foreach_sub_.ScalarList", +} + + +# Checks if `function_schema` is a native, non-foreach function which `f`, a foreach function +# reference to generate derivatives. +def is_reference_for_foreach( + f: NativeFunction, + function_schema: FunctionSchema, +) -> bool: + return ( + f.func.name.name.base.split("_foreach_")[-1] == function_schema.name.name.base + and ( + not function_schema.name.name.inplace + or str(f.func.name) in _foreach_with_inplace_ref + ) + and ( + str(f.func.name) in _skip_argument_len_check + or len(f.func.arguments.flat_non_out) + == len(function_schema.arguments.flat_non_out) + ) + and all( + ref_arg.type in (arg.type, getattr(arg.type, "elem", None)) + for arg, ref_arg in zip( + f.func.arguments.flat_non_out, + function_schema.arguments.flat_non_out, + ) + ) + ) + + +# TODO(crcrpar): Avoid hard coding "Default" ideally. +def gen_foreach_derivativeinfo( + foreach_function: NativeFunction, + functional_info_by_signature: dict[ + FunctionSchema, dict[str, DifferentiabilityInfo] + ], + non_functional_info_by_signature: dict[ + FunctionSchema, dict[str, DifferentiabilityInfo] + ], + dispatch_key: str = "Default", +) -> tuple[DifferentiabilityInfo | None, bool]: + """Generate DifferentiabilityInfo for out-place foreach function, return the existing one for in-place. + + The second return value indicates whether the info is generated in this function. + """ + ref_diff_info: DifferentiabilityInfo | None = None + + for function_schema, diff_info in functional_info_by_signature.items(): + if not is_reference_for_foreach(foreach_function, function_schema): + continue + ref_diff_info = diff_info[dispatch_key] + if ref_diff_info is not None: + break + # note(crcrpar): It seems like `zero`'s info isn't available in functional_info_by_signature + # while the info of `zero_` is in non_functional_info_by_signature + if ( + ref_diff_info is None + and foreach_function.func.kind() == SchemaKind.inplace + and str(foreach_function.func.name) in _foreach_with_inplace_ref + ): + for function_schema, diff_info in non_functional_info_by_signature.items(): + if not is_reference_for_foreach(foreach_function, function_schema): + continue + ref_diff_info = diff_info[dispatch_key] + if ref_diff_info is not None: + break + if ref_diff_info is None: + return None, False + + # non out-place uses the existing Derivative. + if foreach_function.func.kind() == SchemaKind.inplace: + return ref_diff_info, False + + map_refarg2foreacharg, map_name2arg = {}, {} + for i, (arg, ref_arg) in enumerate( + zip( + foreach_function.func.arguments.flat_non_out, + function_schema.arguments.flat_non_out, + ) + ): + map_refarg2foreacharg[ref_arg.name] = arg.name + map_name2arg[arg.name] = arg + + all_saved_inputs, all_saved_outputs, all_var_names = [], [], [] + modified_derivative_formulas = [] + for i, derivative in enumerate(ref_diff_info.derivatives): + modified_formula = derivative.formula.replace("grad", "grads[i]").replace( + "result", "result[i]" + ) + saved_inputs, saved_outputs = [], [] + # note(crcrpar): This context seems necessary to call `cpp.argument_type` + with local.parametrize( + use_const_ref_for_mutable_tensors=foreach_function.use_const_ref_for_mutable_tensors, + use_ilistref_for_tensor_lists=foreach_function.part_of_structured_group, + ): + for ref_input in derivative.saved_inputs: + ref_input_jit_name = ref_input.expr.split(".")[0] + mapped_name = map_refarg2foreacharg[ref_input_jit_name] + if isinstance(map_name2arg[mapped_name].type, ListType): + mapped_expr = mapped_name + "[i]" + else: + mapped_expr = mapped_name + new_expr = ref_input.expr.replace(ref_input_jit_name, mapped_expr) + modified_formula = modified_formula.replace( + cast(str, ref_input.nctype.name), new_expr + ) + + nctype = cpp.argument_type(map_name2arg[mapped_name], binds=mapped_name) + canonical_nctype = NamedCType( + nctype.name, nctype.type.remove_const_ref() + ) + saved_inputs.append( + SavedAttribute(nctype=canonical_nctype, expr=mapped_name) + ) + for ref_output in derivative.saved_outputs: + if ref_output.nctype.name == "result": + saved_outputs.append( + SavedAttribute( + nctype=NamedCType( + name="result", type=BaseCType(tensorListT) + ), + expr="result", + ) + ) + else: + raise RuntimeError("") + var_names = [map_refarg2foreacharg[var] for var in derivative.var_names] + all_var_names.extend(var_names) + all_saved_inputs.extend(saved_inputs) + all_saved_outputs.extend(saved_outputs) + modified_derivative = Derivative( + formula=modified_formula, + original_formula=derivative.formula, + var_names=tuple(var_names), + saved_inputs=tuple(saved_inputs), + saved_outputs=tuple(saved_outputs), + named_gradients=set(), + ) + modified_derivative_formulas.append(modified_derivative) + + with local.parametrize( + use_const_ref_for_mutable_tensors=foreach_function.use_const_ref_for_mutable_tensors, + use_ilistref_for_tensor_lists=foreach_function.part_of_structured_group, + ): + args_with_derivatives = [ + Binding( + name=arg.name, + nctype=cpp.argument_type(arg, binds=arg.name), + argument=arg, + default=None, + ) + for arg in foreach_function.func.arguments.flat_non_out + if arg.name in all_var_names + ] + + forward_derivatives: list[ForwardDerivative] = [] + fw_derivative: ForwardDerivative + for fw_derivative in ref_diff_info.forward_derivatives: + var_names: list[str] = list(fw_derivative.var_names) # type: ignore[no-redef] + var_types: list[Type] = list(fw_derivative.var_types) + required_inputs_fw_grad: list[str] = [] + required_inputs_primal: list[str] = [] + if fw_derivative.required_inputs_fw_grad is not None: + required_inputs_fw_grad = list(fw_derivative.required_inputs_fw_grad) + if fw_derivative.required_inputs_primal: + required_inputs_primal = list(fw_derivative.required_inputs_primal) + modified_formula = fw_derivative.formula + + # Foreach's result is TensorList + if "result" in modified_formula: + modified_formula = fw_derivative.formula.replace("result", "result[i]") + + for foreach_arg, ref_arg in zip( + foreach_function.func.arguments.flat_non_out, + ref_diff_info.func.func.arguments.flat_non_out, + ): + # Modify reference forward formula + if ( + isinstance(foreach_arg.type, ListType) + and not foreach_arg.type.is_tensor_like() + ): + # Assuming ScalarList + modified_formula = modified_formula.replace( + ref_arg.name, foreach_arg.name + "[i]" + ) + elif foreach_arg.type.is_tensor_like(): + # Assuming TensorList / Tensor + if not ( + isinstance(foreach_arg.type, ListType) + or ( + foreach_arg.type == BaseType(BaseTy.Tensor) + and str(foreach_function.func.name) + in _foreach_with_tensor_overload + ) + ): + raise AssertionError( + f"{foreach_function.func.name}, {foreach_arg.type}" + ) + for suffix in ("_p", "_t"): + curr_expr = ref_arg.name + suffix + if curr_expr in modified_formula: + new_expr = foreach_arg.name + suffix + modified_formula = modified_formula.replace(curr_expr, new_expr) + else: + # Assuming Scalar + if foreach_arg.name != ref_arg.name: + modified_formula = modified_formula.replace( + ref_arg.name, foreach_arg.name + ) + + # note(crcrpar): there should exist a cooler way... + for i, name in enumerate(var_names): + if name == ref_arg.name: + var_names[i] = foreach_arg.name + var_types[i] = foreach_arg.type + for i, name in enumerate(required_inputs_fw_grad): + if name == ref_arg.name: + required_inputs_fw_grad[i] = foreach_arg.name + for i, name in enumerate(required_inputs_primal): + if name == ref_arg.name: + required_inputs_primal[i] = foreach_arg.name + forward_derivatives.append( + ForwardDerivative( + formula=modified_formula, + var_names=tuple(var_names), + var_types=tuple(var_types), + required_inputs_fw_grad=tuple(required_inputs_fw_grad), + required_inputs_primal=tuple(required_inputs_primal), + required_original_self_value=fw_derivative.required_original_self_value, + is_reusing_outplace_formula=fw_derivative.is_reusing_outplace_formula, + ) + ) + + return ( + DifferentiabilityInfo( + name=foreach_function.func.name.name.base, + func=foreach_function, + op=f"Foreach{ref_diff_info.op}{foreach_function.func.name.overload_name}", + derivatives=modified_derivative_formulas, + forward_derivatives=forward_derivatives, + all_saved_inputs=tuple(set(all_saved_inputs)), + all_saved_outputs=tuple(set(all_saved_outputs)), + available_named_gradients=(), + used_named_gradients=set(), + args_with_derivatives=args_with_derivatives, + non_differentiable_arg_names=[], + output_differentiability=None, + output_differentiability_conditions=None, + ), + True, + ) + + +def match_differentiability_info( + native_functions: list[NativeFunction], + differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]], +) -> list[NativeFunctionWithDifferentiabilityInfo]: + """Sets the "derivative" key on declarations to matching autograd function + In-place functions will use the out-of-place derivative definition if there + is no in-place specific derivative. + """ + + functional_info_by_signature = { + schema.signature(strip_default=True): info_dict + for schema, info_dict in differentiability_infos.items() + if schema.kind() == SchemaKind.functional + } + non_functional_info_by_signature = { + schema.signature(strip_default=True): info_dict + for schema, info_dict in differentiability_infos.items() + if schema.kind() != SchemaKind.functional + } + + def find_info( + f: NativeFunction, + ) -> tuple[dict[str, DifferentiabilityInfo] | None, bool]: + # Don't bother matching info to generated out= variants + if "generated" in f.tags and f.func.kind() == SchemaKind.out: + return None, False + + # (1) Check for an exact match + if f.func in differentiability_infos: + return differentiability_infos[f.func], True + + # (2) If no exact match, check if the out-of-place variant + # of this operator has a match. + # i.e mul() for mul_() or mul_out() + # note(crcrpar): Check foreach or not because in-place foreach functions use backward defined for the existing + # native functions instead of the out-place counterparts. + f_sig = f.func.signature(strip_default=True) + if f_sig in functional_info_by_signature and not is_foreach_func(f): + return functional_info_by_signature[f_sig], False + + # (3) Some operators have a derivative explicitly defined for the mutable + # variant, but get a code-generated out-of-place variant which does *not* + # come with a derivative formula. + # For the generated out-of-place variant, use the mutable variant's formula + # if it exists. + if "generated" in f.tags and f_sig in non_functional_info_by_signature: + info_dict = non_functional_info_by_signature[f_sig] + # See https://github.com/pytorch/pytorch/pull/76320/files#r874816389 + if any( + any("self" in str(input.nctype.name) for input in info.all_saved_inputs) + for info in info_dict.values() + ): + raise AssertionError( + f"Attempted to convert a derivative formula for a mutable operator " + f'to be used automatically by its functional variant ("{str(f.func)}"). ' + "This is not currently supported (we'd need to fix up the formula in the codegen)." + ) + return info_dict, False + + # (4) Generate derivative information of foreach functions if none is defined in `derivatives.yaml` + if is_foreach_func(f): + if f.func in differentiability_infos: + raise AssertionError( + f"Foreach function {f.func.name} already has differentiability info" + ) + diff_info, is_generated = gen_foreach_derivativeinfo( + f, + functional_info_by_signature, + non_functional_info_by_signature, + ) + if diff_info is None: + return None, False + # TODO(crcrpar): Avoid hard coding "Default" ideally. + diff_info_dict = {"Default": diff_info} + if is_generated: + differentiability_infos[f.func] = diff_info_dict + functional_info_by_signature[f.func] = diff_info_dict + return diff_info_dict, is_generated + + return None, False + + result: list[NativeFunctionWithDifferentiabilityInfo] = [] + for f in native_functions: + info_dict, is_exact_match = find_info(f) + + # Currently, the '.strides()' to 'strides_or_error' replacement does not support + # 'self' derivatives of an inplace function, so we must check for this case. + if f.func.kind() == SchemaKind.inplace and (info_dict is not None): + for info in info_dict.values(): + for derivative in info.derivatives: + if "self" in derivative.var_names: + for saved_input in derivative.saved_inputs: + if "strides_or_error" in saved_input.expr: + raise AssertionError( + "Calling '.strides()' in the 'self' derivative formula of an " + f"in-place function is not supported: {f.func}" + ) + + if not info_dict: + result.append( + NativeFunctionWithDifferentiabilityInfo( + func=f, info=None, fw_derivatives=None + ) + ) + continue + + fw_derivative_dict: dict[str, Sequence[ForwardDerivative]] = {} + for key, info in info_dict.items(): + if not info.forward_derivatives: + fw_derivative_dict[key] = [] + continue + + forward_derivatives = info.forward_derivatives + + # For functions that have a single def for out-of-place and inplace (like abs()) + if f.func.kind() == SchemaKind.inplace: + # For inplace functions there is a little bit of work to do: + # 1) Validate the formula and make sure the input that is modified in not used: + # - If there is a formula for the inplace variant of the function (is_exact_match == True) then + # we make sure that the original value of the input that is being modified inplace (self_p) is + # not used in the formula. Note that the formula can use "original_self_p" here and that would + # trigger a clone of the original input. + # - If we are reusing the out of place formula (is_exact_match == False) then we replace every + # occurrence of self_p and self_t by original_self_p and original_self_t. These will be + # populated by cloned version of the original input (either the clone done by the backward AD + # logic if self is also used in a backward formula or a special clone that we add). + # 2) At this point, there cannot be a self_p in the formula. + # 3) Change "result" into "self_p" as by design, in the inplace function codegen, the result is + # simply called self (as it is modified inplace). + # 4) Update the required primals data in case it used to contain "result" but should now contain + # "self" + # 5) If it is not an exact match, the user formula is not modifying the existing forward grad + # inplace as it should. So add some code that makes sure that we do so if the forward grad + # already exists. + + if len(info.forward_derivatives) != 1: + raise AssertionError( + "Only single output inplace should exist, " + f"got {len(info.forward_derivatives)}" + ) + fw_info = info.forward_derivatives[0] + formula = fw_info.formula + + def replace_self_with_original_self(formula: str, postfix: str) -> str: + def repl(m: re.Match[str]) -> str: + return f"{m.group(1)}original_self{postfix}{m.group(2)}" + + return re.sub(IDENT_REGEX.format(f"self{postfix}"), repl, formula) + + if re.search(IDENT_REGEX.format("self_p"), formula): + if is_exact_match: + # For manually defined formulas, don't allow the original value to be used + raise RuntimeError( + f'The formula for "{f.func.name}" is using the original value of self ' + "that is being modified inplace. This would lead to wrong forward gradients. " + 'Please use "result" in the formula only.' + ) + else: + # When the original formula is out of place, we save a clone of the primal + # value to be able to access this value if needed + # replace "self_p"/"self_t" from the formula by "original_self_p"/"original_self_t" + formula = replace_self_with_original_self(formula, "_p") + formula = replace_self_with_original_self(formula, "_t") + + # replace "result" from the formula by "self_p" + def repl(m: re.Match[str]) -> str: + return f"{m.group(1)}self_p{m.group(2)}" + + formula = re.sub(IDENT_REGEX.format("result"), repl, formula) + + required_primals = fw_info.required_inputs_primal + if re.search(IDENT_REGEX.format("self_p"), formula): + required_primals = ( + required_primals + ("self",) if required_primals else ("self",) + ) + + if not is_exact_match: + # NOTE [In-place forward AD formula Optimization] + # + # This optimization transforms the formula to directly do inplace, i.e. + # instead of self_t.copy_(self_t.op()) we do self_t.op_() when the following are met: + # + # 1) the formula satisfies the pattern: "self_t.op(*args)" + # 2) "op" in (1) needs to be the same as the op the derivative is for + # + # (2) may seem too strict, but currently the only ops that satisfy (1) also satisfy (2) + # If there is a need, we can relax (2) to allow any op that has an in-place variant + is_single_method_on_self_t = False + directly_do_inplace = False + op_name: str | None = None + between_parens: str | None = None + match = re.fullmatch(r"self_t.([\w]*)\((.*)\)", formula) + if match: + op_name, between_parens = match.group(1), match.group(2) + + # We want to... + # Match: self_t.op1(other_p.op2(arg)) + # Avoid: self_t.op1(args) + self_t.op2(args) + # Avoid: self_t.op1(other_p.op2(arg)) + self_t.op2(args) + def check_parens_nest_level_gt_zero(s: str) -> bool: + level = 1 + for ch in s: + if ch == ")": + level -= 1 + if level == 0: + return False + if ch == "(": + level += 1 + return True + + is_single_method_on_self_t = check_parens_nest_level_gt_zero( + between_parens + ) + directly_do_inplace = ( + is_single_method_on_self_t and op_name == info.name + ) + + if directly_do_inplace: + if op_name is None: + raise AssertionError("op_name must be non-None for inplace") + if between_parens is None: + raise AssertionError( + "between_parens must be non-None for inplace" + ) + formula = f"self_t_raw.defined() ? self_t_raw.{op_name}_({between_parens}) : {formula}" + else: + # Make sure that the forward grad is modified inplace when the original formula + # is out of place + formula = f"self_t_raw.defined() ? self_t_raw.copy_({formula}) : {formula}" + + required_original_self_value = bool( + re.search(IDENT_REGEX.format("original_self_p"), formula) + ) or bool(re.search(IDENT_REGEX.format("original_self_t"), formula)) + + forward_derivatives = [ + ForwardDerivative( + formula=formula, + var_names=("self",), + var_types=fw_info.var_types, + required_inputs_fw_grad=fw_info.required_inputs_fw_grad, + required_inputs_primal=required_primals, + required_original_self_value=required_original_self_value, + is_reusing_outplace_formula=not is_exact_match, + ), + ] + + fw_derivative_dict[key] = forward_derivatives + + result.append( + NativeFunctionWithDifferentiabilityInfo( + func=f, info=info_dict, fw_derivatives=fw_derivative_dict + ) + ) + + return result + + +def is_differentiable( + name: str, type: Type, info: DifferentiabilityInfo | None +) -> bool: + return type.is_tensor_like() and ( + info is None or name not in info.non_differentiable_arg_names + ) + + +def gen_differentiable_outputs( + fn: NativeFunctionWithDifferentiabilityInfo, key: str = "Default" +) -> list[DifferentiableOutput]: + f = fn.func + info = fn.info[key] if fn.info else None + outputs: list[DifferentiableOutput] = [ + DifferentiableOutput( + name=name, + type=ret.type, + cpp_type=cpp.return_type(ret, symint=True).cpp_type(), + ) + for name, ret in zip(cpp.return_names(f), f.func.returns) + ] + output_differentiability = info.output_differentiability if info else None + if output_differentiability is not None: + if len(output_differentiability) != len(outputs): + raise RuntimeError( + f"The length of output_differentiability ({len(output_differentiability)}), " + f"does not match the number of outputs ({len(outputs)})." + ) + differentiable_outputs: list[DifferentiableOutput] = [] + if False in output_differentiability and f.func.kind() == SchemaKind.inplace: + raise RuntimeError( + "output_differentiability=False for inplace operation (version_counter won't get updated)" + ) + for differentiable, output in zip(output_differentiability, outputs): + if differentiable: + differentiable_outputs.append(output) + return differentiable_outputs + candidate_differentiable_outputs = list( + filter(lambda r: is_differentiable(r.name, r.type, info), outputs) + ) + if uses_single_grad(info): + return candidate_differentiable_outputs[:1] + else: + return candidate_differentiable_outputs diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/cpp.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/cpp.py new file mode 100644 index 0000000000000000000000000000000000000000..f2ac560246f304fa0b5dfad0bea6fb2d1c37a7ff --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/cpp.py @@ -0,0 +1,475 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing_extensions import assert_never + +from torchgen import local +from torchgen.api.types import ( + ArgName, + ArrayCType, + ArrayRefCType, + BaseCType, + BaseTypeToCppMapping, + Binding, + boolT, + ConstRefCType, + CType, + dimnameListT, + intArrayRefT, + iTensorListRefT, + ListCType, + longT, + MutRefCType, + NamedCType, + OptionalCType, + optionalIntArrayRefT, + optionalSymIntArrayRefT, + scalarT, + SpecialArgName, + symIntArrayRefT, + SymIntT, + tensorListT, + tensorOptionsT, + tensorT, + TupleCType, + VectorCType, + voidT, +) +from torchgen.model import ( + Argument, + Arguments, + BaseTy, + BaseType, + FunctionSchema, + ListType, + NativeFunction, + OptionalType, + Return, + SelfArgument, + TensorOptionsArguments, + Type, +) + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# This file describes the translation of JIT schema to the public C++ +# API, which is what people use when they call functions like at::add. +# +# Prominent characteristics of the C++ API: +# +# - dtype, layout, device and pin_memory are collected into +# a single C++ type TensorOptions (the native functions API +# also has this, but tensor options is really most relevant +# for the C++ API; it makes calling kwarg factory functions +# pleasant) +# +# - defaulting lives here (in fact, the dispatcher is completely +# oblivious of defaults!) +# +# BTW: policy on name collisions: we try not to have types with +# collisions, but functions are fair game to collide + + +def name( + func: FunctionSchema, + *, + faithful_name_for_out_overloads: bool = False, + symint_overload: bool = False, +) -> str: + name = str(func.name.name) + if symint_overload: + name += "_symint" + if func.is_out_fn(): + if faithful_name_for_out_overloads: + name += "_outf" + else: + name += "_out" + + return name + + +# Translation of "value types" in JIT schema to C++ API type. Value +# types look the same no matter if they are argument types or return +# types. Returns None if the type in question is not a value type. +def valuetype_type( + t: Type, + *, + binds: ArgName, + mutable: bool = True, + symint: bool = False, +) -> NamedCType | None: + if isinstance(t, BaseType): + if t.name in (BaseTy.Tensor, BaseTy.Scalar): + return None + elif str(t) == "SymInt": + if symint: + return NamedCType(binds, BaseCType(SymIntT)) + else: + return NamedCType(binds, BaseCType(longT)) + # All other BaseType currently map directly to BaseCppTypes. + return NamedCType(binds, BaseCType(BaseTypeToCppMapping[t.name])) + elif isinstance(t, OptionalType): + elem = valuetype_type(t.elem, binds=binds, mutable=mutable, symint=symint) + if elem is None: + return None + return NamedCType(binds, OptionalCType(elem.type)) + elif isinstance(t, ListType): + if str(t.elem) == "bool": + if t.size is None: + raise AssertionError("bool ListType must have a size") + return NamedCType(binds, ArrayCType(BaseCType(boolT), t.size)) + else: + return None + else: + raise AssertionError(f"unrecognized type {repr(t)}") + + +# Translation of types occurring in JIT arguments to a C++ argument type. +# If remove_non_owning_ref_types is set, we'll guarantee that the output CType is not a non-owning reference type. +# For example, we'll return std::vector instead of IntArrayRef. +# See Note [translation from C++ reference to value types] +def argumenttype_type( + t: Type, + *, + mutable: bool, + binds: ArgName, + remove_non_owning_ref_types: bool = False, + symint: bool = False, +) -> NamedCType: + # If it's a value type, do the value type translation + r = valuetype_type( + t, + binds=binds, + mutable=mutable, + symint=symint, + ) + if r is not None: + return r + + if isinstance(t, BaseType): + if t.name == BaseTy.Tensor: + if mutable and not local.use_const_ref_for_mutable_tensors(): + return NamedCType(binds, MutRefCType(BaseCType(tensorT))) + else: + return NamedCType(binds, ConstRefCType(BaseCType(tensorT))) + elif t.name == BaseTy.Scalar: + return NamedCType(binds, ConstRefCType(BaseCType(scalarT))) + else: + raise AssertionError(f"base type should have been value type {t}") + elif isinstance(t, OptionalType): + if str(t.elem) == "Tensor": + if mutable and not local.use_const_ref_for_mutable_tensors(): + return NamedCType( + binds, MutRefCType(BaseCType(tensorT)) + ) # TODO: fix this discrepancy + else: + return NamedCType( + binds, ConstRefCType(OptionalCType(BaseCType(tensorT))) + ) + elif str(t.elem) == "Scalar": + return NamedCType(binds, ConstRefCType(OptionalCType(BaseCType(scalarT)))) + elif isinstance(t.elem, ListType) and str(t.elem.elem) == "int": + return NamedCType(binds, BaseCType(optionalIntArrayRefT)) + elif isinstance(t.elem, ListType) and str(t.elem.elem) == "SymInt": + if symint: + return NamedCType(binds, BaseCType(optionalSymIntArrayRefT)) + else: + return NamedCType(binds, BaseCType(optionalIntArrayRefT)) + elem = argumenttype_type(t.elem, mutable=mutable, binds=binds, symint=symint) + return NamedCType(binds, OptionalCType(elem.type)) + elif isinstance(t, ListType): + # TODO: remove these special cases, ArrayRef fallthrough works fine + if str(t.elem) == "int": + if remove_non_owning_ref_types: + return NamedCType(binds, VectorCType(BaseCType(longT))) + else: + return NamedCType(binds, BaseCType(intArrayRefT)) + if str(t.elem) == "SymInt": + if remove_non_owning_ref_types: + if symint: + return NamedCType(binds, VectorCType(BaseCType(SymIntT))) + else: + return NamedCType(binds, VectorCType(BaseCType(longT))) + else: + if symint: + return NamedCType(binds, BaseCType(symIntArrayRefT)) + else: + return NamedCType(binds, BaseCType(intArrayRefT)) + if str(t.elem) == "Tensor": + if local.use_ilistref_for_tensor_lists(): + return NamedCType(binds, ConstRefCType(BaseCType(iTensorListRefT))) + else: + return NamedCType(binds, BaseCType(tensorListT)) + elif str(t.elem) == "Scalar": + return NamedCType(binds, ArrayRefCType(BaseCType(scalarT))) + elif str(t.elem) == "Dimname": + return NamedCType(binds, BaseCType(dimnameListT)) + elif str(t.elem) == "Tensor?": + return NamedCType( + binds, ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT)))) + ) + elem = argumenttype_type(t.elem, mutable=mutable, binds=binds, symint=symint) + return NamedCType(binds, ArrayRefCType(elem.type)) + else: + raise AssertionError(f"unrecognized type {repr(t)}") + + +# Translate a JIT argument into its C++ type +def argument_type(a: Argument, *, binds: ArgName, symint: bool = False) -> NamedCType: + return argumenttype_type(a.type, mutable=a.is_write, symint=symint, binds=binds) + + +# Translation of a (non-multi) return type from JIT to C++ +# N.B: returntype_type returns a CType, not a NamedCType. +# This is mostly because of the mismatch between return types and return names. +# e.g. a function with a return type of 'void' has 0 return names, +# and a function with a return type of 'std::tuple' has >1 return name. +def returntype_type(t: Type, *, mutable: bool, symint: bool = False) -> CType: + # placeholder is ignored + # NB: symint is ALWAYS respected for return types. So symint argument + # here is IGNORED + r = valuetype_type(t, binds="__placeholder__", mutable=mutable, symint=True) + if r is not None: + return r.type + + if isinstance(t, BaseType): + if t.name == BaseTy.Tensor: + if mutable: + if local.use_const_ref_for_mutable_tensors(): + return ConstRefCType(BaseCType(tensorT)) + else: + return MutRefCType(BaseCType(tensorT)) + else: + # Note [Tensor Copy Returns] + # Currently, we use "Argument.is_write" to determine + # whether or not Tensor return types should be copies or references. + # If that ever changes, take a look at other locations of this note! + return BaseCType(tensorT) + elif t.name == BaseTy.Scalar: + return BaseCType(scalarT) + elif isinstance(t, ListType): + if mutable: + raise AssertionError( + "Native functions should never return a mutable tensor list. " + "They should return void." + ) + elem = returntype_type(t.elem, mutable=False) + if t.size is not None: + raise AssertionError(f"fixed size list returns not supported: {t}") + return VectorCType(elem) + elif isinstance(t, OptionalType): + elem = returntype_type(t.elem, mutable=mutable) + if str(t.elem) == "Tensor": + return OptionalCType(elem) + + raise AssertionError(f"unrecognized return type {t}") + + +# Translation of a single return to its C++ type +def return_type(r: Return, *, symint: bool = False) -> CType: + return returntype_type(r.type, mutable=r.is_write, symint=symint) + + +# Translation of a full (possibly multi) return from JIT to its C++ type +def returns_type(rs: Sequence[Return], *, symint: bool = False) -> CType: + if len(rs) == 0: + return BaseCType(voidT) + elif len(rs) == 1: + return return_type(rs[0], symint=symint) + else: + return TupleCType([return_type(r, symint=symint) for r in rs]) + + +def return_names(f: NativeFunction, *, fallback_name: str = "result") -> Sequence[str]: + returns: list[str] = [] + for i, r in enumerate(f.func.returns): + # If we have an inplace function, the return argument is + # implicitly named self. + # TODO: Consider incorporating this into the data model + if f.func.name.name.inplace: + if i != 0: + raise AssertionError("illegal inplace function with multiple returns") + name = "self" + # If we are out function, the name is the name of the + # corresponding output function (r.name will get recorded + # in field_name later.) + elif f.func.is_out_fn(): + name = f.func.arguments.out[i].name + # If the return argument is explicitly named... + elif r.name: + name_conflict = any( + r.name == a.name for a in f.func.schema_order_arguments() + ) + if name_conflict and not f.func.is_out_fn(): + name = f"{r.name}_return" + else: + name = r.name + # If there is no explicit name and no fallback name was passed in, we just name the output result, + # unless it's a multi-return, in which case it's result0, + # result1, etc (zero-indexed) + else: + name = fallback_name if len(f.func.returns) == 1 else f"{fallback_name}{i}" + returns.append(name) + return returns + + +JIT_TO_CPP_DEFAULT = { + "False": "false", + "True": "true", + "None": "::std::nullopt", # UGH this one is type directed + "Mean": "at::Reduction::Mean", + "[]": "{}", + "contiguous_format": "c10::MemoryFormat::Contiguous", + "long": "at::kLong", +} + + +# Convert a JIT default into C++ expression representing the default +def default_expr(d: str, t: Type, *, symint: bool) -> str: + if d == "None" and str(t) == "Tensor?": + return "{}" + if isinstance(t, BaseType) and t.name is BaseTy.str: + # Schema allows single quotes but C++ needs double + if len(d) >= 2 and d[0] == "'" and d[-1] == "'": + s = "" + i = 1 + while i + 1 < len(d): + if d[i] != "\\": + if d[i] == '"': + s += '\\"' + else: + s += d[i] + i += 1 + else: + if d[i + 1] == "'": + s += "'" + else: + s += d[i : i + 2] + i += 2 + + return f'"{s}"' + + if isinstance(t, OptionalType): + if d == "None": + return "::std::nullopt" + + return default_expr(d, t.elem, symint=symint) + + if isinstance(t, ListType): + if d.startswith("[") and d.endswith("]"): + return "{" + d[1:-1] + "}" + elif symint and d.isdigit() and str(t.elem) == "SymInt": + return f"c10::SymInt({d})" + elif t.size is None: + # NOTE: Sized lists can have scalar defaults + raise ValueError(f"Expected a list default '[...]' but found: '{d}'") + + return JIT_TO_CPP_DEFAULT.get(d, d) + + +# Convert an argument into its C++ API form + + +def argument( + a: Argument | TensorOptionsArguments | SelfArgument, + *, + cpp_no_default_args: set[str], + method: bool, + faithful: bool, + symint: bool = False, + has_tensor_options: bool, +) -> list[Binding]: + def sub_argument( + a: Argument | TensorOptionsArguments | SelfArgument, + ) -> list[Binding]: + return argument( + a, + cpp_no_default_args=cpp_no_default_args, + method=method, + faithful=faithful, + symint=symint, + has_tensor_options=has_tensor_options, + ) + + if isinstance(a, Argument): + binds: ArgName + if a.name == "memory_format" and has_tensor_options: + binds = SpecialArgName.possibly_redundant_memory_format + else: + binds = a.name + default: str | None = None + if a.name not in cpp_no_default_args and a.default is not None: + default = default_expr(a.default, a.type, symint=symint) + return [ + Binding( + nctype=argument_type(a, binds=binds, symint=symint), + name=a.name, + default=default, + argument=a, + ) + ] + elif isinstance(a, TensorOptionsArguments): + if faithful: + return ( + sub_argument(a.dtype) + + sub_argument(a.layout) + + sub_argument(a.device) + + sub_argument(a.pin_memory) + ) + else: + default = None + # Enforced by NativeFunction.__post_init__ + if "options" in cpp_no_default_args: + raise AssertionError("'options' should not be in cpp_no_default_args") + if all(x.default == "None" for x in a.all()): + default = "{}" + elif a.dtype.default == "long": + default = "at::kLong" # TODO: this is wrong + return [ + Binding( + nctype=NamedCType("options", BaseCType(tensorOptionsT)), + name="options", + default=default, + argument=a, + ) + ] + elif isinstance(a, SelfArgument): + if method: + # Caller is responsible for installing implicit this in context! + return [] + else: + return sub_argument(a.argument) + else: + assert_never(a) + + +def arguments( + arguments: Arguments, + *, + faithful: bool, + symint: bool = False, + method: bool, + cpp_no_default_args: set[str], +) -> list[Binding]: + args: list[Argument | TensorOptionsArguments | SelfArgument] = [] + if faithful: + args.extend(arguments.non_out) + args.extend(arguments.out) + else: + args.extend(arguments.out) + args.extend(arguments.non_out) + return [ + r.no_default() if faithful else r + for a in args + for r in argument( + a, + faithful=faithful, + symint=symint, + method=method, + has_tensor_options=arguments.tensor_options is not None, + cpp_no_default_args=cpp_no_default_args, + ) + ] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/dispatcher.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/dispatcher.py new file mode 100644 index 0000000000000000000000000000000000000000..fcca7a60fec1829c5783197055733467fcdd63fe --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/dispatcher.py @@ -0,0 +1,125 @@ +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING +from typing_extensions import assert_never + +from torchgen.api import cpp +from torchgen.api.types import ArgName, Binding, CType, NamedCType +from torchgen.model import ( + Argument, + FunctionSchema, + Return, + SelfArgument, + TensorOptionsArguments, + Type, +) +from torchgen.utils import concatMap + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# This file describes the translation of JIT schema to the dispatcher +# API, the *unboxed* calling convention by which invocations through +# the dispatcher are made. Historically, the dispatcher API matched +# the C++ API, but with the establishment of the boxed API, we've +# made changes to the dispatcher API to so that the unboxed API +# better aligns with the boxed API. The dispatcher API hooks heavily +# into our template based boxing/unboxing machinery, so changes +# to this convention will usually need template updates too. +# +# Prominent characteristics of the dispatcher API: +# +# - dtype, layout, device and pin_memory are represented as separate +# arguments. +# + + +def name(func: FunctionSchema) -> str: + return cpp.name(func) + + +def argumenttype_type( + t: Type, + *, + mutable: bool, + binds: ArgName, + remove_non_owning_ref_types: bool = False, + symint: bool = True, +) -> NamedCType: + # This is a faux amis. If it makes sense in the future to add + # more special cases here, or invert things so cpp.argument_type + # calls this, or just completely inline the function, please do + # it. + return cpp.argumenttype_type( + t, + mutable=mutable, + binds=binds, + symint=symint, + remove_non_owning_ref_types=remove_non_owning_ref_types, + ) + + +def argument_type( + a: Argument, + *, + binds: ArgName, + remove_non_owning_ref_types: bool = False, + symint: bool = True, +) -> NamedCType: + return argumenttype_type( + a.type, + mutable=a.is_write, + binds=binds, + remove_non_owning_ref_types=remove_non_owning_ref_types, + symint=symint, + ) + + +def returns_type(rs: Sequence[Return], *, symint: bool = True) -> CType: + # At present, there is no difference. But there could be! + return cpp.returns_type(rs, symint=symint) + + +def jit_arguments(func: FunctionSchema) -> list[Argument]: + def to_argument( + a: Argument | TensorOptionsArguments | SelfArgument, + ) -> list[Argument]: + if isinstance(a, Argument): + return [a] + elif isinstance(a, SelfArgument): + return [a.argument] + elif isinstance(a, TensorOptionsArguments): + return [a.dtype, a.layout, a.device, a.pin_memory] + else: + assert_never(a) + + return list( + concatMap( + to_argument, + itertools.chain( + func.arguments.positional, func.arguments.kwarg_only, func.arguments.out + ), + ) + ) + + +def argument( + a: Argument, *, remove_non_owning_ref_types: bool = False, symint: bool = True +) -> Binding: + return Binding( + nctype=argument_type( + a, + binds=a.name, + remove_non_owning_ref_types=remove_non_owning_ref_types, + symint=symint, + ), + name=a.name, + argument=a, + ) + + +def arguments(func: FunctionSchema, *, symint: bool = True) -> list[Binding]: + return [argument(a, symint=symint) for a in jit_arguments(func)] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/functionalization.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/functionalization.py new file mode 100644 index 0000000000000000000000000000000000000000..0d097e28a04a93df5b6ea9cef0c5abe41ef55509 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/functionalization.py @@ -0,0 +1,222 @@ +from __future__ import annotations + +from torchgen.api import dispatcher +from torchgen.api.types import ( + BaseCppType, + BaseCType, + Binding, + boolT, + ConstRefCType, + CType, + longT, + NamedCType, + tensorT, +) +from torchgen.model import ( + Argument, + BaseTy, + BaseType, + FunctionSchema, + NativeFunction, + NativeFunctionsViewGroup, +) + + +# This file describes the translation of JIT schema to API's used +# when creating `ViewMeta` specializations that are used by the functionalization pass. +# These API's mostly follow the dispatcher API, with one difference: +# - While the forward function just directly calls into the at::_ops API +# (following the dispatcher convention), the logic here for the reverse function +# is responsible for generating both the call-site, and the declarations +# (which are implemented manually in the at::functionalization::impl namespace). + +# Define some specific lambda input arguments. +base_binding = Binding( + name="base", + nctype=NamedCType(name="base", type=ConstRefCType(BaseCType(tensorT))), + argument=Argument( + name="base", type=BaseType(BaseTy.Tensor), default=None, annotation=None + ), + default=None, +) + +has_symbolic_inputs_binding = Binding( + name="has_symbolic_inputs", + nctype=NamedCType(name="has_symbolic_inputs", type=BaseCType(boolT)), + argument=Argument( + name="has_symbolic_inputs", + type=BaseType(BaseTy.bool), + default=None, + annotation=None, + ), + default=None, +) +mutated_view_binding = Binding( + name="mutated_view", + nctype=NamedCType(name="mutated_view", type=ConstRefCType(BaseCType(tensorT))), + argument=Argument( + name="base", type=BaseType(BaseTy.Tensor), default=None, annotation=None + ), + default=None, +) +out_index_binding = Binding( + name="out_index", + nctype=NamedCType(name="out_index", type=BaseCType(longT)), + argument=Argument( + name="out_index", type=BaseType(BaseTy.int), default=None, annotation=None + ), + default=None, +) +reapply_views_binding = Binding( + name="reapply_views", + nctype=NamedCType(name="reapply_views", type=BaseCType(boolT)), + argument=Argument( + name="reapply_views", type=BaseType(BaseTy.bool), default=None, annotation=None + ), + default=None, +) + +InverseReturnModeT = BaseCppType("at::functionalization", "InverseReturnMode") +inverse_return_mode_binding = Binding( + name="inverse_return_mode", + nctype=NamedCType(name="inverse_return_mode", type=BaseCType(InverseReturnModeT)), + argument=Argument( + name="inverse_return_mode", + # NB: not actually a bool but it doesn't matter because this isn't used + type=BaseType(BaseTy.bool), + default=None, + annotation=None, + ), + default=None, +) + + +# Name of the `ViewMeta` specialization class created. +def classname(func: FunctionSchema, with_namespace: bool = False) -> str: + namespace = "at::functionalization::" if with_namespace else "" + return f"{namespace}{func.name.unambiguous_name()}_ViewMeta" + + +# Name of the operation called inside the `forward`/`reverse` implementations. +def name( + g: NativeFunctionsViewGroup, + *, + is_reverse: bool, + include_namespace: bool, + reapply_views: bool | None = None, +) -> str: + if reapply_views is None: + # reapply_views is only important for the fwd lambda, + # since we always plumb the runtime "reapply_views" argument into the reverse function. + if not is_reverse: + raise AssertionError("reapply_views can only be None for reverse") + if is_reverse: + return reverse_name(g.view, include_namespace) + # in the forward case, we just directly call into the at::_ops API (so we always need the namespace) + if not include_namespace: + raise AssertionError("include_namespace must be True for forward") + if g.view_copy is None: + raise AssertionError("view_copy must be non-None for forward") + api_name = ( + g.view.func.name.unambiguous_name() + if reapply_views + else g.view_copy.func.name.unambiguous_name() + ) + return f"at::_ops::{api_name}::call" + + +def reverse_name(f: NativeFunction, include_namespace: bool) -> str: + # for the reverse: we plumb the "reapply_views" flag into that function and support + # both copy and non-copy variants. (We could avoid doing that, but that would require + # writing out twice as many view inverse functions). + api_name = f.func.name.unambiguous_name() + # in the reverse case, we codegen both the call-sites (which need the full namespace) and the declarations (which don't) + if include_namespace: + return f"at::functionalization::FunctionalInverses::{api_name}_inverse" + else: + return f"{api_name}_inverse" + + +def returns_type(func: FunctionSchema) -> CType: + # Assertion: all view ops return tensor-like outputs + if len(func.returns) < 1: + raise AssertionError("Expected at least one return value") + for ret in func.returns: + if not ret.type.is_tensor_like(): + raise AssertionError(f"Expected tensor-like return type, got {ret.type}") + # However, the return type of the lambda is always an individual tensor. + # For multi-tensor outputs, each tensor needs to be tracked individually. + return BaseCType(tensorT) + + +# Checks whether `func` might return more than one value. +def is_multi_output(func: FunctionSchema) -> bool: + return len(func.returns) > 1 or ( + len(func.returns) == 1 and func.returns[0].type.is_list_like() is not None + ) + + +# `ViewMeta` specialization constructor parameters. +def base_ctor_arguments(func: FunctionSchema) -> list[Binding]: + # All specializations are parematerized by `has_symbolic_inputs` flag. + arguments = [has_symbolic_inputs_binding] + + # If `func` might return more than 1 value, we also parameterize this specialization + # with the output index. + if is_multi_output(func): + arguments.append(out_index_binding) + + return arguments + + +# `ViewMeta` specialized class' constructor arguments. +# +# Values needed specifically by this specialization, that the base class does not need. +# Same as the class' attributes, but non-owning. +def extra_ctor_arguments(func: FunctionSchema) -> list[Binding]: + return attributes(func, owning=False) + + +# `ViewMeta` specialized class' non-static member data. +# +# Essential data for calling the instance's `forward` and `reverse functions. You can +# think of them as values that should be captured from the functionalization kernel. +def attributes(func: FunctionSchema, owning: bool = True) -> list[Binding]: + args = func.arguments.flat_all + if args[0].type != BaseType(BaseTy.Tensor): + raise AssertionError(f"Expected first arg to be Tensor, got {args[0].type}") + return [ + reapply_views_binding, + inverse_return_mode_binding, + *[dispatcher.argument(a, remove_non_owning_ref_types=owning) for a in args[1:]], + ] + + +def op_arguments(func: FunctionSchema, is_reverse: bool) -> list[Binding]: + args = func.arguments.flat_all + if args[0].type != BaseType(BaseTy.Tensor): + raise AssertionError(f"Expected first arg to be Tensor, got {args[0].type}") + non_self_args = args[1:] + # The forward lambda calls the at::_ops API, while the reverse lambda calls the view inverse API. + # Both of these follow the dispatcher API. + non_self_bindings = [dispatcher.argument(a) for a in non_self_args] + if not is_reverse: + # the forward lambda swaps out the original tensor argument with the lambd arg "base" + return [base_binding] + non_self_bindings + else: + # the reverse lambda does the same, but with an additional "mutated_view" arg + # additionally, we have a calling convention: for view ops that return multiple tensor outputs + # their corresponding view_inverse function takes in an additional index argument. + if is_multi_output(func): + return [ + base_binding, + mutated_view_binding, + inverse_return_mode_binding, + out_index_binding, + ] + non_self_bindings + else: + return [ + base_binding, + mutated_view_binding, + inverse_return_mode_binding, + ] + non_self_bindings diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/lazy.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/lazy.py new file mode 100644 index 0000000000000000000000000000000000000000..eff973ae11de99dfe2dea98bcfa7a85d25c99682 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/lazy.py @@ -0,0 +1,470 @@ +from __future__ import annotations + +from typing import Any + +from torchgen.api.types import ( + BaseCppType, + BaseCType, + boolT, + CType, + deviceT, + doubleT, + generatorT, + layoutT, + ListCType, + longT, + memoryFormatT, + NamedCType, + OptionalCType, + scalarT, + scalarTypeT, + stringT, + SymIntT, + VectorCType, +) +from torchgen.model import ( + Argument, + BaseTy, + BaseType, + FunctionSchema, + ListType, + OperatorName, + OptionalType, + Return, + TensorOptionsArguments, + Type, +) + + +_valueT: BaseCppType | None = None + + +# A ValueT is an IR type which represents the computation of a Tensor. In other +# words, a PyTorch user will do operations on lazy tensors, and each output lazy +# tensor internally tracks a ValueT representing the IR node that would have +# actually produced the value of this tensor for real. +# +# This is configurable because different lazy tensor backends (LTC vs XLA) will +# have different IR representations. (Though, arguably, after unification they +# shouldn't!) +def getValueT() -> BaseCppType: + global _valueT + if not _valueT: + raise NotImplementedError( + "The value type needs to be set with setValueT() in run_gen_lazy_tensor()" + ) + + return _valueT + + +def setValueT(val: BaseCppType) -> None: + global _valueT + _valueT = val + + +# this is a bad hack. I need to refactor the data model to represent each arg in the schema as an object, +# making it easier to represent special properties of an arg. +tensorListValueT = BaseCppType("torch::lazy", "Value") + + +def process_ir_type( + typ: Type, properties: LazyIrProperties, *, symint: bool +) -> BaseCType | VectorCType | OptionalCType | ListCType: + """ + This function takes a type from NativeFunctions and converts it for use with + lazy tensor codegen. + + Type conversion for lazy currently consists of + (1) changing at::Tensors into lazy::Values + (2) wrapping everything in a BaseCType + (3) making cpp-reference types into cpp-value types (e.g. vector instead of IntArrayRef) + + (1) converts at::Tensors to lazy::Values (which wrap lazy::Nodes, with which Lazy IR represents tensors.) + There is special handling for Optional[Tensor] or list[Tensor], etc- hence 'tensor-like' + + This is incomplete- there are assertions in places that it's expected to need to add + more types as the codegen is used with more operators. + """ + if isinstance(typ, BaseType): + if typ.name == BaseTy.Tensor: + return BaseCType(getValueT()) + elif typ.name == BaseTy.Scalar: + if properties.TreatScalarsAsConstants: + return BaseCType(scalarT) + # at::scalar has special handling, + # and is wrapped in an lazy::Value just like at::tensor + return BaseCType(getValueT()) + elif typ.name == BaseTy.ScalarType: + return BaseCType(scalarTypeT) + elif typ.name == BaseTy.int: + return BaseCType(longT) + elif typ.name == BaseTy.SymInt: + if symint: + return BaseCType(getValueT()) + else: + return BaseCType(longT) + elif typ.name == BaseTy.bool: + return BaseCType(boolT) + elif typ.name == BaseTy.float: + return BaseCType(doubleT) + elif typ.name == BaseTy.str: + return BaseCType(stringT) + elif typ.name == BaseTy.Device: + return BaseCType(deviceT) + elif typ.name == BaseTy.Generator: + return BaseCType(generatorT) + elif typ.name == BaseTy.Layout: + return BaseCType(layoutT) + elif typ.name == BaseTy.MemoryFormat: + return BaseCType(memoryFormatT) + else: + raise AssertionError(f"TODO add support for type {repr(typ)}") + elif isinstance(typ, OptionalType): + return OptionalCType(process_ir_type(typ.elem, properties, symint=symint)) + elif isinstance(typ, ListType): + if str(typ.elem) == "Tensor?": + # TODO(whc) is this actually correct? or should it use a Vector like above + return ListCType(OptionalCType(BaseCType(getValueT()))) + elif str(typ.elem) == "Tensor": + # this is a TensorList which comes in from GetTensorList as a Value + return BaseCType(tensorListValueT) + elif typ.elem == BaseType(BaseTy.SymInt): + # TODO: return a value type. The problem here is analogous to + # the problem with tensorListValueT: if you have SymInt[] you + # cannot conveniently save the list of Value directly, as nodes + # expect to save values as a vector for ALL arguments. So you + # need a separate IR node that represents all of the size nodes + # assembled into a list. I'm not an LTC dev so I don't want to + # figure it out right now. Y'all figure it out... + return VectorCType(BaseCType(longT)) + + else: + return VectorCType(process_ir_type(typ.elem, properties, symint=symint)) + else: + raise AssertionError(f"unrecognized type {repr(typ)}") + + +# TODO: Determining this based off of CType is bad; this should be computed +# from Type directly; then the same logic as process_ir_type can be used +# +# Invariant: passed typ should be an *owning* CType (e.g., we will report +# that ArrayRef is NOT a value type) +def isValueType(typ: CType, properties: LazyIrProperties | None = None) -> bool: + """ + Given a type, determine if it is a Value-like type. This is equivalent to + being Tensor-like, but assumes the type has already been transformed. + """ + if isinstance(typ, BaseCType): + # I am regretting my naming conventions, but now we are wrapping at::scalar in + # lazy value, while preserving other 'scalar' types as scalars in the IR + treat_scalars_as_constants = properties and properties.TreatScalarsAsConstants + return ( + typ.type == getValueT() + or (typ.type == scalarT and not treat_scalars_as_constants) + or typ.type == SymIntT + ) + elif typ == VectorCType(BaseCType(SymIntT)): + # TODO: report True for this + return False + elif isinstance(typ, (OptionalCType, ListCType, VectorCType)): + return isValueType(typ.elem, properties) + return False + + +def isSymIntType(typ: Type) -> bool: + return isinstance(typ, BaseType) and typ.name == BaseTy.SymInt + + +def isWrappedScalarType(typ: Type) -> bool: + """ + Given a type, determine if it is a c10::scalar which we will wrap in a lazy Value. + Since we literally change the type from scalarT to valueT, information is lost. + This function helps build a list of wrapped scalars to save that information + """ + if isinstance(typ, BaseType): + # I am regretting my naming conventions, but now we are wrapping at::scalar in + # lazy value, while preserving other 'scalar' types as scalars in the IR + return typ.name == BaseTy.Scalar + elif isinstance(typ, (OptionalType, ListType)): + return isWrappedScalarType(typ.elem) + return False + + +# TODO: dedupe with Type.is_generator_like +def isGeneratorType(typ: Type) -> bool: + if isinstance(typ, BaseType): + return typ.name == BaseTy.Generator + elif isinstance(typ, (OptionalType)): + return isGeneratorType(typ.elem) + return False + + +# This class caches a few derived properties computed from an Argument +# and LazyIrProperties +class LazyArgument: + name: str + orig_type: Type + lazy_type_: CType | None + is_wrapped_scalar: bool + is_generator: bool + # TODO: this is lies, it is false for symint list + is_symint_or_list: bool + + # Whether or not we are treating this as symint or not + symint: bool + + # true if this argument is or contains a lazy IR value + is_lazy_value: bool + + def __init__( + self, arg: Argument, properties: LazyIrProperties, *, symint: bool + ) -> None: + self.name = arg.name + self.orig_type = arg.type + self.symint = symint + self.is_optional = isinstance(arg.type, OptionalType) + self.is_generator = isGeneratorType(arg.type) + self.lazy_type_ = process_ir_type(arg.type, properties, symint=symint) + self.is_wrapped_scalar = isWrappedScalarType(arg.type) + self.is_symint_or_list = symint and ( + isSymIntType(arg.type) + or (isinstance(arg.type, OptionalType) and isSymIntType(arg.type.elem)) + # TODO: lists of symints are not currently treated as value types + # or (isinstance(arg.type, ListType) and isSymIntType(arg.type.elem)) + ) + + self.is_lazy_value = isValueType(self.lazy_type, properties) + + @property + def lazy_type(self) -> CType: + if self.lazy_type_ is None: + raise AssertionError( + f"Attempted to access lazy_type for invalid argument {self.name}" + ) + return self.lazy_type_ + + +class LazyIrProperties: + """Collection of properties for an IR node + + The property groups are listed below. Each group is mutually + exclusive, meaning that only one property from each group can be True + at any one time. The properties can be accessed as if they were normal + attributes. The mutual exclusivity is automatically handled. + """ + + Properties: tuple[tuple[str, ...], ...] = ( + ( + "ShapePrecompute", # Assume shape has been precomputed + "ShapeCompute", # Need to compute the shape on construction + "ShapeCache", # Utilize the shape cache to defer computation + ), + ( + "Lower", # Codegen full lower function + "LowerDeclOnly", # Codegen only lower function declaration + ), + ( + "CanBeReused", # Codegen full reuse function + "CanBeReusedDeclOnly", # Codegen only reuse function declaration + ), + ( + "CreateFn", # Codegen full create function + "CreateFnDeclOnly", # Codegen only create function declaration + ), + ( + "TreatScalarsAsConstants", # Treat Scalars as constants instead of handling like values + ), + ) + + def __init__(self, *default_properties: str) -> None: + properties: dict[tuple[str, ...], str | None] = dict.fromkeys( + LazyIrProperties.Properties + ) + self.__dict__["properties"] = properties + for p in default_properties: + setattr(self, p, True) + + def __getattr__(self, key: str) -> Any: + properties = self.__dict__["properties"] + for values in LazyIrProperties.Properties: + if key in values: + return properties[values] == key + + return self.__getattribute__(key) + + def __setattr__(self, key: str, value: Any) -> Any: + properties = self.__dict__["properties"] + for values in LazyIrProperties.Properties: + if key in values: + properties[values] = key if value else None + return value + + raise KeyError(f"Invalid property: {key}") + + +# Inspired by a FunctionSchema object, a LazyIrSchema holds the schema of a Lazy IR node. +# Unlike a FunctionSchema, it has no round-trippable string form (relating to the YAML), +# but carries type information from a native FunctionSchema modified for use with IR nodes, +# and preserving original argument names. +# +# TODO: This is not idiomatic with how other torchgen APIs transform on schema. +class LazyIrSchema: + # The name of the operator this function schema describes. + name: OperatorName + + positional_args: tuple[LazyArgument, ...] + keyword_args: tuple[LazyArgument, ...] + + # TODO: Need to handle collisions with argument names at some point + returns: tuple[Return, ...] + + # if this schema has a Generator arg, list its orig ctype/name but don't + # build a LazyArgument since lazy IR doesn't support it + generator_arg: NamedCType | None = None + + # original function schema + func: FunctionSchema + + # Whether or not we are code-genning for SymInt or not + symint: bool + + properties: LazyIrProperties = LazyIrProperties( + # default properties + "ShapePrecompute", + "Lower", + "CanBeReused", + ) + opkind: str | None = None + + def __init__( + self, + func: FunctionSchema, + properties: LazyIrProperties | None = None, + *, + symint: bool, + ) -> None: + if properties: + self.properties = properties + + self.func = func + self.symint = symint + positional_args: list[LazyArgument] = [] + for arg_field in ["pre_self_positional", "self_arg", "post_self_positional"]: + if arg_field == "self_arg" and func.arguments.self_arg is not None: + arg = func.arguments.self_arg.argument + positional_args.append( + LazyArgument(arg, self.properties, symint=symint) + ) + elif getattr(func.arguments, arg_field) is not None: + positional_args.extend( + LazyArgument(arg, self.properties, symint=symint) + for arg in getattr(func.arguments, arg_field) + ) + self.positional_args = tuple(positional_args) + + keyword_args: list[LazyArgument] = [] + for arg_field in [ + "pre_tensor_options_kwarg_only", + "tensor_options", + "post_tensor_options_kwarg_only", + "out", + ]: + curr_args = getattr(func.arguments, arg_field) + if curr_args is not None: + if isinstance(curr_args, TensorOptionsArguments): + curr_args = curr_args.all() + for arg in curr_args: + if isGeneratorType(arg.type): + if self.generator_arg is not None: + raise AssertionError( + "We expect there is only one generator arg" + ) + self.generator_arg = NamedCType( + arg.name, + arg.type, # type:ignore[arg-type] + ) + keyword_args.extend( + LazyArgument(arg, self.properties, symint=symint) + for arg in curr_args + ) + self.keyword_args = tuple(keyword_args) + self.name = func.name + self.returns = func.returns + + @property + def node_name(self) -> str: + """ + Return camel-case version of op in node. + + Note: This function also appends any `overload_name` in the operation. + For example, if the op is `bitwise_and.Tensor`, the returned name + will be `BitwiseAndTensor`. + """ + op_name = f"{self.name.name}_{self.name.overload_name}".lower() + return "".join(word.capitalize() or "" for word in op_name.split("_")) + + @property + def aten_name(self) -> str: + return str(self.name.name) + + @property + def base_name(self) -> str: + return f"{self.name.name.base}" + + def filtered_args( + self, + positional: bool = True, + keyword: bool = True, + values: bool = True, + scalars: bool = True, + generator: bool = True, + ) -> list[LazyArgument]: + # This function maintains the sorted order of arguments but provides different filtered views. + # Some parts of the code care about kwargs vs args (TS lowerings), + # other parts care about whether they need to wrap the arg in a lazy value or leave it alone. + # Generators are special cased, as they are needed for fallback/shape-inference but not supported + # in TS lowerings and therefore also omitted from lazy IR. + args: list[LazyArgument] = [] + if positional: + args.extend(self.positional_args) + if keyword: + args.extend(self.keyword_args) + + if values and scalars and generator: + return args + elif values and scalars: + return [a for a in args if not a.is_generator] + elif values: + return [a for a in args if a.is_lazy_value] + elif scalars: + return [ + a + for a in args + if not a.is_lazy_value and (generator or not a.is_generator) + ] + + return [] + + @property + def positional_values(self) -> list[LazyArgument]: + return self.filtered_args( + positional=True, keyword=False, values=True, scalars=False + ) + + @property + def positional_scalars(self) -> list[LazyArgument]: + return self.filtered_args( + positional=True, keyword=False, values=False, scalars=True + ) + + @property + def keyword_values(self) -> list[LazyArgument]: + return self.filtered_args( + positional=False, keyword=True, values=True, scalars=False + ) + + @property + def keyword_scalars(self) -> list[LazyArgument]: + return self.filtered_args( + positional=False, keyword=True, values=False, scalars=True + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/meta.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/meta.py new file mode 100644 index 0000000000000000000000000000000000000000..2e99d151faeaccea7ca47f372fd26f9985ce7249 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/meta.py @@ -0,0 +1,13 @@ +from torchgen.model import NativeFunctionsGroup + + +# Follows dispatcher calling convention, but: +# - Mutable arguments not allowed. Meta functions are always +# written in functional form. Look at FunctionSchema.signature() +# - No tensor returns; instead we return a TensorMeta describing +# the tensor in question + + +def name(g: NativeFunctionsGroup) -> str: + # use the overload name from the functional version + return str(g.functional.func.name).replace(".", "_") diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/native.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/native.py new file mode 100644 index 0000000000000000000000000000000000000000..632216704d2d47606b977d487335ca196e2e1842 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/native.py @@ -0,0 +1,159 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING +from typing_extensions import assert_never + +from torchgen import local +from torchgen.api import cpp +from torchgen.api.types import ( + ArgName, + BaseCType, + Binding, + boolT, + ConstRefCType, + CType, + deviceT, + layoutT, + ListCType, + MutRefCType, + NamedCType, + OptionalCType, + scalarT, + scalarTypeT, + tensorT, +) +from torchgen.model import ( + Argument, + FunctionSchema, + Return, + SelfArgument, + TensorOptionsArguments, + Type, +) + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# This file describes the translation of JIT schema to the native functions API. +# This looks a lot like the C++ API (which makes historical sense, because the +# idea was you wrote native functions to implement functions in the C++ API), +# but over time we have evolved the C++ API without actually changing our +# native:: kernels. The intention is to make native API and dispatcher API +# line up as closely as possible, since this results in the least overhead +# (no translation is needed from dispatcher API to native API). +# +# NB: this is symint aware, you will get the non-SymInt variant for some +# dispatch entries and SymInt for others. + + +def name(func: FunctionSchema) -> str: + name = str(func.name.name) + # TODO: delete this! + if func.is_out_fn(): + name += "_out" + if func.name.overload_name: + name += f"_{func.name.overload_name}" + return name + + +def argumenttype_type( + t: Type, *, mutable: bool, binds: ArgName, symint: bool +) -> NamedCType: + if str(t) == "Tensor?": + tensor_type: OptionalCType = OptionalCType(BaseCType(tensorT)) + if mutable and not local.use_const_ref_for_mutable_tensors(): + return NamedCType(binds, MutRefCType(tensor_type)) + else: + return NamedCType(binds, ConstRefCType(tensor_type)) + elif str(t) == "Tensor?[]": + return NamedCType( + binds, ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT)))) + ) + elif str(t) == "Scalar": + return NamedCType(binds, ConstRefCType(BaseCType(scalarT))) + elif str(t) == "Scalar?": + return NamedCType(binds, ConstRefCType(OptionalCType(BaseCType(scalarT)))) + return cpp.argumenttype_type(t, mutable=mutable, binds=binds, symint=symint) + + +def returns_type(rs: Sequence[Return], *, symint: bool) -> CType: + return cpp.returns_type(rs, symint=symint) + + +def argument_type(a: Argument, *, binds: ArgName, symint: bool) -> NamedCType: + return argumenttype_type(a.type, mutable=a.is_write, binds=binds, symint=symint) + + +def argument( + a: Argument | SelfArgument | TensorOptionsArguments, + *, + is_out: bool, + symint: bool, +) -> list[Binding]: + # Ideally, we NEVER default native functions. However, there are a number + # of functions that call native:: directly and rely on the defaulting + # existing. So for BC, we generate defaults for non-out variants (but not + # for out variants, where it is impossible to generate an appropriate + # default) + should_default = not is_out + if isinstance(a, Argument): + default: str | None = None + if should_default and a.default is not None: + default = cpp.default_expr(a.default, a.type, symint=symint) + return [ + Binding( + nctype=argument_type(a, binds=a.name, symint=symint), + name=a.name, + default=default, + argument=a, + ) + ] + elif isinstance(a, SelfArgument): + # Erase SelfArgument from the distinction + return argument(a.argument, is_out=is_out, symint=symint) + elif isinstance(a, TensorOptionsArguments): + default = None + if should_default: + default = "{}" + # TODO: Not sure why the arguments assigned here are for + # TensorOptionsArguments and not the constituent pieces. It seems + # to matter + return [ + Binding( + nctype=NamedCType("dtype", OptionalCType(BaseCType(scalarTypeT))), + name="dtype", + default=default, + argument=a, + ), + Binding( + nctype=NamedCType("layout", OptionalCType(BaseCType(layoutT))), + name="layout", + default=default, + argument=a, + ), + Binding( + nctype=NamedCType("device", OptionalCType(BaseCType(deviceT))), + name="device", + default=default, + argument=a, + ), + Binding( + nctype=NamedCType("pin_memory", OptionalCType(BaseCType(boolT))), + name="pin_memory", + default=default, + argument=a, + ), + ] + else: + assert_never(a) + + +def arguments(func: FunctionSchema, *, symint: bool) -> list[Binding]: + args: list[Argument | TensorOptionsArguments | SelfArgument] = [] + args.extend(func.arguments.non_out) + args.extend(func.arguments.out) + return [ + r for arg in args for r in argument(arg, symint=symint, is_out=func.is_out_fn()) + ] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/python.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/python.py new file mode 100644 index 0000000000000000000000000000000000000000..254d7c1ee9b438758aa8ab02af5632da6eee0d5b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/python.py @@ -0,0 +1,1553 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from torchgen.api import cpp +from torchgen.api.types import Binding, CppSignature, CppSignatureGroup +from torchgen.gen import pythonify_default +from torchgen.model import ( + Argument, + BaseTy, + BaseType, + FunctionSchema, + ListType, + NativeFunction, + OptionalType, + Return, + Type, + Variant, +) + + +if TYPE_CHECKING: + from collections.abc import Iterable, Sequence + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Data Models +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# [Notes] python binding codegen +# +# The Python binding codegen produces code that takes the input list of +# PyObjects, finds the matching ATen C++ function using PythonArgParser, +# converts the PyObjects into C++ types and calls the ATen C++ function: +# +# +--------+ parsing +------------------------+ binding +-----------------------+ +# | PyObjs | ---------> | PythonArgParser Output | ---------> | Cpp Function Dispatch | +# +--------+ +------------------------+ +-----------------------+ +# +# The following examples demonstrate the data models the Python binding +# codegen needs to deal with and the tasks it needs to accomplish. It +# helps understand the purpose of the new data types we introduced below. +# +# - Function Schema (source of truth) +# +# aten::empty.names(int[] size, *, Dimname[]? names, +# ScalarType? dtype=None, Layout? layout=None, +# Device? device=None, bool? pin_memory=None, +# MemoryFormat? memory_format=None) -> Tensor +# +# - Python Signature +# +# It's used to generate input schema string for PythonArgParser. +# Note: TensorOptions fields are reordered and the additional +# 'requires_grad' field is added: +# +# empty(IntArrayRef size, *, DimnameList? names, +# MemoryFormat? memory_format=None, ScalarType dtype=None, +# Layout layout=torch.strided, Device device=None, +# bool pin_memory=False, bool requires_grad=False) +# +# - C++ Signature +# +# It's used to generate C++ lambda formals & dispatch call. +# Note: the scattered TensorOptions fields are packed into 'options'. +# +# auto dispatch_empty = +# [](IntArrayRef size, std::optional names, +# const TensorOptions & options, +# std::optional memory_format) -> Tensor { +# pybind11::gil_scoped_release no_gil; +# return torch::empty(size, names, options, memory_format); +# }; +# +# - Binding between Python Arguments and C++ Arguments +# +# Given a set of Python Arguments in scope, we need produce the +# binding expressions that translate the Python API into C++ API: +# +# Python Args Cpp Args Binding Exprs +# ----------------------------------------------------------------- +# 0: size size '_r.intlist(0)' +# 1: names names 'names' [special init] +# 2: memory_format -------+ +# 3: dtype -----+-|--> options 'options' [special packing] +# 4: layout / | +# 5: device / +--> memory_format '_r.memoryformatOptional(2)' +# 6: pin_memory / +# 7: requires_grad -+ +# +# So the full dispatch expression would look like: +# +# dispatch_empty(_r.intlist(0), names, options, +# _r.memoryformatOptional(2)) +# +# Where does 'names' come from? It involves special local init: +# +# auto __names = _r.toDimnameListOptional(1); +# std::optional names = +# __names ? std::make_optional(DimnameList(__names.value())) +# : std::nullopt; +# +# Where does 'options' come from? It involves special local init +# for TensorOptions. Note that Python side has the additional +# 'requires_grad' field: +# +# const auto options = TensorOptions() +# .dtype(_r.scalartype(3)) +# .device(_r.device(5)) +# .layout(_r.layoutOptional(4)) +# .requires_grad(_r.toBool(7)) +# .pinned_memory(_r.toBool(6)); +# +# In some other cases one Python Argument can map to multiple C++ +# Arguments. For example: +# +# aten::max.names_dim(Tensor self, Dimname dim, bool keepdim=False) +# -> (Tensor values, Tensor indices) +# +# Python Args Cpp Args Binding Exprs +# --------------------------------------------------------------------- +# +----> max 'out[0]' +# /-----> max_values 'out[1] +# 0: input / self '_r.tensor(0)' +# 1: dim / dim '_r.dimname(1)' +# 2: keepdim / keepdim '_r.toBool(2)' +# 3: out -----+ [local init] out '_r.tensorlist_n<2>(3)' +# +# As demonstrated above, the binding can involve reordering, +# packing, unpacking and special local inits. +# +# +# Let's look at a concrete example: +# +# static PythonArgParser parser({ +# "abs(Tensor input, *, Tensor out=None)", +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# ^ +# +--- Python Schema, represented by PythonSignature and PythonArgument +# +# }, /*traceable=*/true); +# +# ParsedArgs<2> parsed_args; +# auto _r = parser.parse(nullptr, args, kwargs, parsed_args); +# +# ... +# +# if (_r.isNone(1)) { +# ~~~~~~~~~~~~ <--- Scattered PythonArgParser output (arg name = 'out') +# represented by PythonArgParserOutputExpr +# +# // aten::abs(Tensor self) -> Tensor +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# ^ +# +--- NativeFunction schema, base version +# +# auto dispatch_abs = [](const Tensor & self) -> Tensor { +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# ^ +# +--- dispatch_lambda_args / dispatch_lambda_return_str +# generated from NativeFunction / CppSignature +# (deprecated PythonSignature is special) +# arguments are represented by DispatchLambdaArgument +# +# pybind11::gil_scoped_release no_gil; +# return self.abs(); +# ~~~~~~~~~~~ <--- cpp_dispatch_target / cpp_dispatch_exprs +# generated from NativeFunction / CppSignature +# }; +# return wrap(dispatch_abs(_r.tensor(0))); +# ~~~~~~~~~~~~~ +# ^ +# +--- dispatch_lambda_exprs +# binding PythonArgParserOutputExpr (python args) +# and DispatchLambdaArgument (c++ args) +# +# } else { +# // aten::abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# ^ +# +--- NativeFunction schema, out-variant +# +# auto dispatch_abs_out = [](Tensor out, const Tensor & self) -> Tensor { +# pybind11::gil_scoped_release no_gil; +# return at::abs_out(out, self); +# }; +# return wrap(dispatch_abs_out(_r.tensor(1), _r.tensor(0))); +# } +# +# +# [Notes] python interface codegen +# The python dataclasses below are used used to generate both python binding code +# and pyi type hint signatures. +# In theory these two should look very similar, but there are number of differences +# in how pyi signatures vs. python_arg_parser signatures are generated. +# These differences have been encapsulated in signature_str() vs. signature_str_pyi() +# to display the full signatures, and argument_str() vs argument_str_pyi() to display arguments. +# For examples, only pyi signatures include return types. + + +def format_function_signature( + name: str, arguments: Iterable[str] = (), return_type: str | None = None +) -> str: + if not isinstance(arguments, (list, tuple)): + arguments = tuple(arguments) + return_type = f" -> {return_type}" if return_type is not None else "" + + sig = f"def {name}({', '.join(arguments)}){return_type}: ..." + if len(sig) <= 80 or len(arguments) == 0 or tuple(arguments) == ("self",): + return sig + + lines = [ + f"def {name}(", + *(f" {arg}," for arg in arguments), + f"){return_type}: ...", + ] + sig = "\n".join(lines) + if all(len(line) <= 80 for line in lines): + return sig + # ruff format bug for compound statements: https://github.com/astral-sh/ruff/issues/18658 + # use `skip` instead of `on` + `off` + return sig.removesuffix(" ...") + " # fmt: skip\n ..." + + +@dataclass(frozen=True) +class PythonReturns: + returns: tuple[Return, ...] + + +@dataclass(frozen=True) +class PythonArgument: + name: str + type: Type + default: str | None + + # Used to generate the default init expr for some PythonArgParser outputs, e.g.: + # + # _r.layoutWithDefault(3, layout_from_backend(self.options().backend()))) + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # ^ + # +--- default_init str + default_init: str | None + + # Compute argument formal for python argument parsing. + # Needs to be consistent with torch/csrc/utils/python_arg_parser.h. + def argument_str(self, *, method: bool = False, symint: bool = True) -> str: + type_str = ( + argument_type_str(self.type, symint=symint) + .replace("const ", "") + .replace(" &", "") + ) + + name = self.name + # s/self/input/ outside method bindings + # [old codegen] TODO: remove this? doesn't rename in codegen, it's just + # for the parse string + if name == "self" and type_str in ["Tensor", "Number"] and not method: + name = "input" + + # add default + if self.default is not None: + default = { + "nullptr": "None", + "::std::nullopt": "None", + "std::nullopt": "None", + "{}": "None", + }.get(self.default, self.default) + return f"{type_str} {name}={default}" + else: + return f"{type_str} {name}" + + def argument_str_pyi( + self, *, method: bool = False, deprecated: bool = False + ) -> str: + type_str = argument_type_str_pyi(self.type) + + name = self.name + # s/self/input/ outside method bindings + # [old codegen] TODO: remove this? doesn't rename in codegen, it's just + # for the parse string + if name == "self" and type_str == "Tensor" and not method and not deprecated: + name = "input" + + if name == "from": # from is a Python keyword... + name += "_" + + # pyi merges the _out and functional variants into the same signature, with an optional out arg + if name == "out" and not deprecated: + type_str = f"{type_str} | None".replace(" | None | None", " | None") + + # pyi deprecated signatures don't get defaults for their out arg + treat_as_no_default = ( + deprecated + and isinstance(self, PythonOutArgument) + and self.default == "None" + ) + + # add default + if self.default is not None and not treat_as_no_default: + if ( + isinstance(self.type, ListType) + and self.type.elem == BaseType(BaseTy.int) + and self.default.startswith("{") + and self.default.endswith("}") + ): + default = ( + "(" + ", ".join(map(str.strip, self.default[1:-1].split(","))) + ")" + ) + else: + default = { + "nullptr": "None", + "::std::nullopt": "None", + "std::nullopt": "None", + "{}": "None", + "c10::MemoryFormat::Contiguous": "contiguous_format", + "QScheme::PER_TENSOR_AFFINE": "per_tensor_affine", + }.get(self.default, self.default) + return f"{name}: {type_str} = {default}" + else: + return f"{name}: {type_str}" + + +@dataclass(frozen=True) +class PythonOutArgument(PythonArgument): + # In Python signature multiple output fields are packed into one 'out' argument. + # When binding to C++, it's first binded to a local 'out' variable: + # 'auto out = _r.tensorlist_n<2>(2);', + # then binded to scattered C++ output arguments as 'out[0]', 'out[1]', and etc. + # TODO: maybe don't need keep scattered out fields for python signature? + outputs: tuple[PythonArgument, ...] + + @staticmethod + def from_outputs(outputs: tuple[PythonArgument, ...]) -> PythonOutArgument | None: + if not outputs: + return None + + size = len(outputs) + if size == 1: + return PythonOutArgument( + name=outputs[0].name, + type=outputs[0].type, + default="None", + default_init=None, + outputs=outputs, + ) + elif size > 1: + if any(not a.type.is_tensor_like() for a in outputs): + raise RuntimeError(f"Unsupported output type: {outputs}") + return PythonOutArgument( + name="out", + # TODO: shouldn't this be OptionalType[ListType[...]], since it defaults to None? + type=ListType(BaseType(BaseTy.Tensor), size), + default="None", + default_init=None, + outputs=outputs, + ) + raise AssertionError(r"Unexpected PythonOutArgument size") + + +@dataclass(frozen=True) +class PythonSignature: + # Base operator name, without inplace/outplace suffix. + name: str + + # Positional arguments. + # TODO: create a dedicated SelfArgument type for 'self'? + input_args: tuple[PythonArgument, ...] + + # Keyword arguments excluding the 'out' argument and scattered kwargs belonging + # to TensorOptions (dtype, layout, device, pin_memory, requires_grad, etc). + input_kwargs: tuple[PythonArgument, ...] + + output_args: PythonOutArgument | None + + # Return types, which are only used by pyi + returns: PythonReturns + + # These are scattered kwargs arguments belonging to TensorOptions. + # When binding to C++, they are packed into a TensorOptions object 'options'. + # It's possible that the C++ signature doesn't take TensorOptions object (e.g. + # for out variant), in which case they will be used as scattered fields without + # being packed into 'options'. + # TODO: maybe create a PythonTensorOptionsArgument? + tensor_options_args: tuple[PythonArgument, ...] + + # method or function signature? + method: bool + + @property + def deprecated(self) -> bool: + return False + + def arguments( + self, *, skip_outputs: bool = False, skip_tensor_options: bool = False + ) -> tuple[PythonArgument | PythonOutArgument, ...]: + result: list[PythonArgument | PythonOutArgument] = [] + result.extend(self.input_args) + result.extend(self.input_kwargs) + if self.output_args is not None and not skip_outputs: + result.append(self.output_args) + if not skip_tensor_options: + result.extend(self.tensor_options_args) + return tuple(result) + + def arguments_count(self) -> int: + return len(self.arguments()) + + def output_idx(self) -> int: + return len(self.input_args) + len(self.input_kwargs) + + # [old codegen] Compute the Python function signature for argument parsing, + # as specified in torch/csrc/utils/python_arg_parser.h. WARNING: + # this is NOT the same type signature as specified by PEP 484 + # as understood by mypy; our format was independently developed + # and has some quirks to make it more suitable specifically + # for error parsing. + # + # For a translation to mypy-valid type signatures, see + # signature_str_pyi(). + def signature_str(self, *, skip_outputs: bool = False, symint: bool = True) -> str: + args = self.arguments(skip_outputs=skip_outputs) + schema_formals: list[str] = [ + a.argument_str(method=self.method, symint=symint) for a in args + ] + positional_argc = len(self.input_args) + if len(schema_formals) > positional_argc: + schema_formals.insert(positional_argc, "*") + + return f"{self.name}({', '.join(schema_formals)})" + + def signature_str_pyi(self, *, skip_outputs: bool = False) -> str: + args = self.arguments(skip_outputs=skip_outputs) + schema_formals: list[str] = [ + a.argument_str_pyi(method=self.method) for a in args + ] + positional_argc = len(self.input_args) + if len(schema_formals) > positional_argc: + schema_formals.insert(positional_argc, "*") + + # only pyi signatures include returns + returns_str = returns_str_pyi(self) + # pyi also includes self (with no typing/defaults) for methods + if self.method: + schema_formals.insert(0, "self") + return format_function_signature(self.name, schema_formals, returns_str) + + def signature_str_pyi_vararg(self, *, skip_outputs: bool = False) -> str | None: + # only pyi uses vararg signatures + args = self.arguments(skip_outputs=skip_outputs) + schema_formals: list[str] = [ + a.argument_str_pyi(method=self.method) for a in args + ] + # vararg only applies to pyi signatures. vararg variants are not generated for all signatures + num_args = self.arguments_count() + if num_args == 0: + return None + + num_positionalargs = len(self.input_args) + + vararg_type = args[0].type + if not ( + isinstance(vararg_type, ListType) + and str(vararg_type.elem) in ["int", "SymInt"] + and num_positionalargs == 1 + ): + return None + + # Below are the major changes in vararg vs. regular pyi signatures + # vararg signatures also omit the asterix + if not isinstance(vararg_type, ListType): + raise AssertionError(f"Expected ListType, got {type(vararg_type)}") + schema_formals[0] = ( + "*" + args[0].name + ": " + argument_type_str_pyi(vararg_type.elem) + ) + + returns_str = returns_str_pyi(self) + # pyi also includes self (with no typing/defaults) for methods + if self.method: + schema_formals.insert(0, "self") + return format_function_signature(self.name, schema_formals, returns_str) + + +# The deprecated python signature involves some special logic, so create a +# dedicated data model to store these extra properties. +@dataclass(frozen=True) +class PythonSignatureDeprecated(PythonSignature): + # Schema for the deprecated function + deprecated_schema: FunctionSchema + + # The deprecated signature might miss some arguments that the corresponding + # C++ signature expects. We need store the constant default values to pass in. + # For example: + # [deprecate signature]: addmm(Scalar beta, Tensor self, Tensor mat1, Tensor mat2) + # [func schema]: aten::addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + # [func call]: self.addmm(mat1, mat2, beta, 1) + # We store ['self', 'mat1', 'mat2', 'beta', '1'] in this case. + deprecated_args_exprs: tuple[str, ...] + + @property + def deprecated(self) -> bool: + return True + + def signature_str(self, *, skip_outputs: bool = False, symint: bool = True) -> str: + return ( + PythonSignature.signature_str( + self, skip_outputs=skip_outputs, symint=symint + ) + + "|deprecated" + ) + + def signature_str_pyi(self, *, skip_outputs: bool = False) -> str: + args = self.arguments(skip_outputs=skip_outputs) + schema_formals: list[str] = [ + a.argument_str_pyi(method=self.method, deprecated=True) for a in args + ] + positional_argc = len(self.input_args) + if len(schema_formals) > positional_argc: + schema_formals.insert(positional_argc, "*") + + returns_str = returns_str_pyi(self) + return format_function_signature(self.name, schema_formals, returns_str) + + def signature_str_pyi_vararg(self, *, skip_outputs: bool = False) -> str | None: + # the codegen doesn't include vararg variants for deprecated signatures + return None + + +# This struct is used to hold the PythonSignature and its corresponding +# NativeFunction BEFORE grouping base and out-variant functions. +# Why not store NativeFunction in PythonSignature or construct PythonSignature +# from NativeFunction? Because they are not 1-1 mapped. +# One native function could have both deprecated and non-deprecated python +# signatures - NativeFunction doesn't contain information to construct the +# deprecated python signature. +# One python signature is used to handle both the base and the out-variant +# function - see 'PythonSignatureGroup'. +@dataclass(frozen=True) +class PythonSignatureNativeFunctionPair: + signature: PythonSignature + function: NativeFunction + + +# We merge pairs of functions with signatures that are equivalent mod +# output arguments, and use a single entry in the python_arg_parser sig +# list for both (output arguments become optional). +@dataclass(frozen=True) +class PythonSignatureGroup: + # The signature used for Python argument parsing. The outplace signature + # is preferred if exists, because it can be used to parse inputs for both + # the out-place variant and the base version (with output omitted). + signature: PythonSignature + + # The regular ATen declaration (e.g. conv2d) + base: NativeFunction + + # The out variant (e.g. conv2d_out) + outplace: NativeFunction | None + + @classmethod + def from_pairs( + cls, + functional: PythonSignatureNativeFunctionPair, + out: PythonSignatureNativeFunctionPair | None, + ) -> PythonSignatureGroup: + if out is None: + return PythonSignatureGroup( + signature=functional.signature, + base=functional.function, + outplace=None, + ) + + # prefer the signature with optional out=... arguments because it's the + # superset that can be used to parse input for both base and outplace. + signature_kwargs = out.signature.__dict__.copy() + + # Out overloads in C++ don't have TensorOptions arguments, + # so take these from the functional variant + signature_kwargs["tensor_options_args"] = ( + functional.signature.tensor_options_args + ) + + return PythonSignatureGroup( + signature=type(out.signature)(**signature_kwargs), + base=functional.function, + outplace=out.function, + ) + + +# C++ function dispatch is wrapped in a lambda function. The lambda function +# has almost the same signature as the C++ function, only with some small +# variants - see details below. +# This data model is used to represent arguments of the lambda function +# signature. +@dataclass(frozen=True) +class DispatchLambdaArgument: + name: str + type_str: str + is_out_arg: bool + + +# To pass PyObjects arguments to C++ function (via the lambda wrapper), +# we need first convert PyObjects into simple C++ objects. This work +# is done by PythonArgParser. +# This data model is used to represent the output of PythonArgParser. +# It has 1-1 mapping with PythonArgument in PythonSignature. +@dataclass(frozen=True) +class PythonArgParserOutputExpr: + # argument name + name: str + + # RHS expression to reference PythonArgParser output. + expr: str + + # In some special cases we need create different expr, e.g.: + # '_r.isNone(1)' instead of '_r.tensor(1)'. + index: int + + # The python argument it maps to. + argument: PythonArgument + + @property + def is_none_expr(self) -> str: + return f"_r.isNone({self.index})" + + +# To pass PythonArgParser output to the lambda wrapper, we need bind +# PythonArgParserOutputExpr to DispatchLambdaArgument. +# They are not always 1-1 mapped, e.g. scattered TensorOptions fields +# need be packed into a TensorOptions object, which is the argument +# that the lambda function wrapper takes. +@dataclass(frozen=True) +class DispatchLambdaArgumentExprs: + # The exprs that provide the binding for lambda arguments, e.g.: + # + # 'self' -> '_r.tensor(0)' + # 'min' -> 'out[0]' / 'min_indices' -> 'out[1]' + # 'options' -> 'options' + # + # It has 1-1 mapping with DispatchLambdaArgument. + exprs: Sequence[str] + + # Special local inits, which might introduce new variables that + # the 'exprs' above reference, e.g.: + # + # 'auto out = _r.tensorlist_n<2>(2);' + # + inits: Sequence[str] + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Helper Functions +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +def _cpp_signature(f: NativeFunction, *, method: bool = False) -> CppSignature: + return CppSignatureGroup.from_native_function(f, method=method).signature + + +def has_tensor_options(f: NativeFunction) -> bool: + return f.func.arguments.tensor_options is not None + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Python Signature +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +# 'simple_type' was introduced by the old codegen, which is slightly +# different from the python schema type, e.g.: doesn't have '?' suffix +# for optional Tensor/TensorList; doesn't have '[size]' suffix for list type. +def argument_type_str( + t: Type, *, simple_type: bool = False, symint: bool = True +) -> str: + if isinstance(t, BaseType): + if t.name == BaseTy.int: + return "int64_t" + elif t.name == BaseTy.float: + return "double" + elif t.name == BaseTy.str: + return "c10::string_view" + elif t.name in [ + BaseTy.Tensor, + BaseTy.bool, + BaseTy.QScheme, + BaseTy.Scalar, + BaseTy.ScalarType, + BaseTy.Generator, + BaseTy.Storage, + BaseTy.Layout, + BaseTy.Device, + BaseTy.DeviceIndex, + BaseTy.MemoryFormat, + BaseTy.Dimname, + BaseTy.Stream, + BaseTy.SymInt, + ]: + # These python schema type names line up with their function schema names + return t.name.name + + elif isinstance(t, OptionalType): + elem = argument_type_str(t.elem, simple_type=simple_type, symint=symint) + return f"{elem}?" + elif isinstance(t, ListType): + size = t.size if not simple_type else None + if str(t.elem) == "bool": + if t.size is None: + raise AssertionError("bool ListType must have a size") + return f"::std::array" + elif str(t.elem) == "int": + return f"IntArrayRef[{size}]" if size is not None else "IntArrayRef" + elif str(t.elem) == "SymInt": + if symint: + return ( + f"SymIntArrayRef[{size}]" if size is not None else "SymIntArrayRef" + ) + else: + return f"IntArrayRef[{size}]" if size is not None else "IntArrayRef" + elif str(t.elem) == "Tensor": + return f"TensorList[{size}]" if size is not None else "TensorList" + elif str(t.elem) == "Scalar": + return f"ScalarList[{size}]" if size is not None else "ScalarList" + elif str(t.elem) == "Tensor?": + if simple_type: + return "c10::List<::std::optional>" + else: + return "const c10::List<::std::optional> &" + elif str(t.elem) == "Dimname": + return f"DimnameList[{size}]" if size is not None else "DimnameList" + elem = argument_type_str(t.elem, simple_type=simple_type, symint=symint) + return f"ArrayRef<{elem}>" + + raise RuntimeError(f"unrecognized type {repr(t)}") + + +def argument_type_size(t: Type) -> int | None: + l = t.is_list_like() + if l is not None and str(l.elem) != "bool": + return l.size + else: + return None + + +def argument(a: Argument) -> PythonArgument: + return PythonArgument( + name=a.name, + type=a.type, + # TODO: directly translate a.default to python default + default=( + str(pythonify_default(cpp.default_expr(a.default, a.type, symint=False))) + if a.default is not None + else None + ), + default_init=None, + ) + + +# Generates a PythonSignature that can be used for either .pyi or PythonArgParser codegen +def signature( + f: NativeFunction, *, method: bool = False, pyi: bool = False +) -> PythonSignature: + return signature_from_schema( + f.func, category_override=f.category_override, method=method, pyi=pyi + ) + + +def signature_from_schema( + func: FunctionSchema, + *, + category_override: str | None, + method: bool = False, + pyi: bool = False, +) -> PythonSignature: + args: list[Argument] = [] + args.extend(func.arguments.pre_self_positional) + # Skip SelfArgument if this is method. + if not method and func.arguments.self_arg is not None: + args.append(func.arguments.self_arg.argument) + args.extend(func.arguments.post_self_positional) + args.extend(func.arguments.pre_tensor_options_kwarg_only) + # Skip TensorOptionsArguments. Python side TensorOptions + # arguments are created based on different rules - see below. + args.extend(func.arguments.post_tensor_options_kwarg_only) + args.extend(func.arguments.out) + + input_arg_set = {a.name for a in func.arguments.flat_positional} + kwarg_only_set = {a.name for a in func.arguments.flat_kwarg_only} + out_arg_set = {a.name for a in func.arguments.out} + + input_args = tuple(map(argument, filter(lambda a: a.name in input_arg_set, args))) + input_kwargs = tuple( + map(argument, filter(lambda a: a.name in kwarg_only_set, args)) + ) + outputs = tuple(map(argument, filter(lambda a: a.name in out_arg_set, args))) + + # Reintroduce the scattered fields of TensorOptions for Python. + # Compared to the cpp counterpart, the python arguments have new property + # (default_init) and a new argument 'requires_grad', which require some + # special handlings. + # [old codegen] TODO: because these aren't guaranteed to be 100% faithful + # to the original versions in the yaml, this recreation is a potential + # source of drift between eager and JIT. Pull this logic out to a shared place. + + has_tensor_input_arg = any( + a.type.is_tensor_like() for a in func.arguments.flat_non_out + ) + if any(a.name == "requires_grad" for a in func.schema_order_arguments()): + raise ValueError( + "argument named requires_grad is reserved, should not explicitly add it in the schema" + ) + + # [old codegen] this probably won't work if one of the returns is not a tensor, + # but it will produce a compile-time error that is obvious. + has_tensor_return = any(r.type.is_tensor_like() for r in func.returns) + + name: str = cpp.name(func) + is_factory_function = category_override == "factory" or ( + has_tensor_return and not has_tensor_input_arg + ) + is_like_or_new_function = ( + category_override in ("new", "like") + or name.startswith("new_") + or name.endswith("_like") + ) + is_dummy_function = category_override == "dummy" + + tensor_options_args: list[PythonArgument] = [] + if (is_factory_function or is_like_or_new_function) and not is_dummy_function: + + def topt_default_init(name: str) -> str | None: + topt_args = func.arguments.tensor_options + if topt_args is None: + return None + a = getattr(topt_args, name) + if a.default is None or a.default == "None": + return None + return cpp.default_expr(a.default, a.type, symint=False) + + tensor_options_args.append( + PythonArgument( + name="dtype", + type=OptionalType(BaseType(BaseTy.ScalarType)), + default="None", + default_init=( + None if is_like_or_new_function else topt_default_init("dtype") + ), + ) + ) + tensor_options_args.append( + PythonArgument( + name="layout", + type=OptionalType(BaseType(BaseTy.Layout)), + default="None", + default_init=( + None if is_like_or_new_function else topt_default_init("layout") + ), + ) + ) + tensor_options_args.append( + PythonArgument( + name="device", + type=OptionalType(BaseType(BaseTy.Device)), + default="None", + default_init=( + None + if is_like_or_new_function + else ( + topt_default_init("device") + or "torch::tensors::get_default_device()" + ) + ), + ) + ) + tensor_options_args.append( + PythonArgument( + name="pin_memory", + type=OptionalType(BaseType(BaseTy.bool)), + default="False", + default_init=None, + ) + ) + tensor_options_args.append( + PythonArgument( + name="requires_grad", + type=OptionalType(BaseType(BaseTy.bool)), + default="False", + default_init=None, + ) + ) + + returns = PythonReturns(returns=func.returns) + + return PythonSignature( + name=str(func.name.name), + input_args=input_args, + input_kwargs=input_kwargs, + output_args=PythonOutArgument.from_outputs(outputs), + tensor_options_args=tuple(tensor_options_args), + returns=returns, + method=method, + ) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Python Interface +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +def structseq_fieldnames(returns: tuple[Return, ...]) -> list[str]: + if len(returns) <= 1 or all(r.name is None for r in returns): + return [] + else: + if any(r.name is None for r in returns): + # When building on Windows, `PyStructSequence_UnnamedField` could not be + # resolved by the linker for some reason, which cause error in building: + # + # python_nn_functions.cpp.obj : error LNK2001: unresolved external symbol + # PyStructSequence_UnnamedField + # + # Thus, at this point in time, we do not support unnamed + # fields in structseq; you must either name all fields, + # or none of them. + raise ValueError("Unnamed field is not supported by codegen") + + return [str(r.name) for r in returns] + + +def argument_type_str_pyi(t: Type) -> str: + add_optional = False + if isinstance(t, OptionalType): + t = t.elem + add_optional = True + + ret = "" + if isinstance(t, BaseType): + if t.name in [BaseTy.int, BaseTy.DeviceIndex]: + ret = "_int" + if t.name == BaseTy.SymInt: + ret = "_int | SymInt" + elif t.name == BaseTy.float: + ret = "_float" + elif t.name == BaseTy.str: + ret = "str" + elif t.name == BaseTy.Scalar: + ret = "Number | _complex" + elif t.name == BaseTy.ScalarType: + ret = "_dtype" + elif t.name == BaseTy.bool: + ret = "_bool" + elif t.name == BaseTy.QScheme: + ret = "_qscheme" + elif t.name == BaseTy.Layout: + ret = "_layout" + elif t.name == BaseTy.Device: + ret = "DeviceLikeType | None" + elif t.name == BaseTy.MemoryFormat: + ret = "memory_format" + elif t.name == BaseTy.Dimname: + ret = "str | EllipsisType | None" + elif t.name == BaseTy.Storage: + ret = "Storage | UntypedStorage" + elif t.name in [BaseTy.Tensor, BaseTy.Generator, BaseTy.Stream]: + # These python schema type names line up with their function schema names + ret = t.name.name + + elif isinstance(t, ListType): + if str(t.elem) == "int": + ret = "_int | _size" if t.size is not None else "_size" + elif t.is_tensor_like(): + # Tensor?[] translates to tuple[Tensor | None, ...] | list[Tensor | None] | None + # Tensor[] translates to tuple[Tensor, ...] | list[Tensor] + if isinstance(t.elem, OptionalType): + add_optional = True + elem_str = "Tensor | None" + else: + elem_str = "Tensor" + ret = ( + f"Tensor | tuple[{elem_str}, ...] | list[{elem_str}]" + if t.size is not None + else f"tuple[{elem_str}, ...] | list[{elem_str}]" + ) + elif str(t.elem) == "float": + ret = "Sequence[_float]" + elif str(t.elem) == "SymInt" and t.size is not None: + elem = argument_type_str_pyi(t.elem) + ret = f"{elem} | Sequence[{elem}]" + else: + elem = argument_type_str_pyi(t.elem) + ret = f"Sequence[{elem}]" + + else: + raise RuntimeError(f"unrecognized type {repr(t)}") + + if add_optional: + ret = f"{ret} | None".replace(" | None | None", " | None") + + return ret + + +def return_type_str_pyi(t: Type) -> str: + # Where arguments are open to accepting Union, return types should return + # concrete types + + if isinstance(t, OptionalType): + inner = return_type_str_pyi(t.elem) + return f"{inner} | None".replace(" | None | None", " | None") + + if isinstance(t, BaseType): + if t.name == BaseTy.Device: + return "_device" + elif t.name == BaseTy.Dimname: + return "str | None" + else: + return argument_type_str_pyi(t) + + if isinstance(t, ListType): + inner = return_type_str_pyi(t.elem) + return f"tuple[{inner}, ...]" + + return argument_type_str_pyi(t) + + +def returns_structseq_pyi(signature: PythonSignature) -> tuple[str, str] | None: + python_returns = [return_type_str_pyi(r.type) for r in signature.returns.returns] + structseq_name = signature.name + field_names = structseq_fieldnames(signature.returns.returns) + if field_names: + # These types are structseq objects which act like named NamedTuples, but + # the constructor acts like the constructor of tuple. Using typing.NamedTuple + # does not allow us to override __init__. + seq_type = f"tuple[{', '.join(python_returns)}]" + structseq_def_lines = [ + f"class {structseq_name}({seq_type}): # fmt: skip", + ] + for name, ret_type in zip(field_names, python_returns): + structseq_def_lines.extend( + [ + " @property", + f" def {name}(self) -> {ret_type}: ...", + ] + ) + structseq_def_lines.extend( + [ + " def __new__(", + " cls,", + f" sequence: {seq_type},", + " ) -> Self: # fmt: skip", + " ...", + f" n_fields: Final[_int] = {len(field_names)}", + f" n_sequence_fields: Final[_int] = {len(field_names)}", + " n_unnamed_fields: Final[_int] = 0", + " def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing", + "", # add an extra newline + ] + ) + structseq_def = "\n".join(structseq_def_lines) + # Example: + # structseq_def = ( + # "class max(tuple[Tensor, Tensor]): # fmt: skip\n" + # " @property\n" + # " def values(self) -> Tensor: ...\n" + # " @property\n" + # " def indices(self) -> Tensor: ...\n" + # " def __new__(\n" + # " cls,\n" + # " sequence: tuple[Tensor, Tensor],\n" + # " ) -> Self: # fmt: skip\n" + # " ...\n" + # " n_fields: Final[_int] = 2", + # " n_sequence_fields: Final[_int] = 2", + # " n_unnamed_fields: Final[_int] = 0", + # " def __init_subclass__(cls) -> NoReturn: ... # prohibit subclassing", + # ) + return structseq_name, structseq_def + return None + + +def returns_str_pyi(signature: PythonSignature) -> str: + field_names = structseq_fieldnames(signature.returns.returns) + if field_names: + return f"torch.return_types.{signature.name}" + + python_returns = [return_type_str_pyi(r.type) for r in signature.returns.returns] + if len(python_returns) > 1: + return "tuple[" + ", ".join(python_returns) + "]" + if len(python_returns) == 1: + return python_returns[0] + return "None" + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# C++ Function Dispatch +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# This section provides APIs to generate the code that does C++ function +# dispatch. The C++ function call is wrapped by a lambda function. +# For example: +# +# // aten::selu_(Tensor(a!) self) -> Tensor(a!) +# auto dispatch_selu_ = [](Tensor self) -> Tensor { +# pybind11::gil_scoped_release no_gil; +# return at::selu_(self); +# }; +# +# The lambda function's signature follows the C++ signature in common +# cases, e.g.: +# +# // aten::add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor +# [](const Tensor & self, const Tensor & other, Scalar alpha) -> Tensor +# +# For out variant the 'out' argument's type is changed from 'Tensor &' +# to 'Tensor'. It's because when calling the lambda it passes in the +# PythonArgParser output '_r.tensor(3)', which is stack allocated object +# and needs to pass by value. Also see comments in 'dispatch_lambda_return_str()'. +# +# // aten::add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) +# [](Tensor out, const Tensor & self, const Tensor & other, Scalar alpha) -> Tensor +# +# For multi-output case it can keep using reference type because the +# PythonArgParser output has been unpacked to local variables, e.g.: +# +# // aten::max.names_dim_max(Tensor self, Dimname dim, bool keepdim=False, *, +# // Tensor(a!) max, Tensor(b!) max_values) -> (Tensor(a!) values, Tensor(b!) indices) +# [](Tensor & max, Tensor & max_values, const Tensor & self, Dimname dim, bool keepdim) -> std::tuple +# +# For deprecated python signature, it should follow deprecated python arg order. +# TODO: This is to keep same byte-for-byte result as the old codegen - maybe unnecessary? + + +def dispatch_lambda_args( + ps: PythonSignature, f: NativeFunction, symint: bool = True +) -> tuple[DispatchLambdaArgument, ...]: + if isinstance(ps, PythonSignatureDeprecated): + schema = ps.deprecated_schema + else: + schema = f.func + + # Start with cpp arguments - dispatch lambda signature always include 'self' + cpp_args = cpp.arguments( + arguments=schema.arguments, + faithful=False, + symint=symint, + method=False, + cpp_no_default_args=f.cpp_no_default_args, + ) + out_args: set[str] = {a.name for a in schema.arguments.out} + + # Convert from cpp argument to lambda argument + def dispatch_lambda_arg(cpp_arg: Binding) -> DispatchLambdaArgument: + type_str = cpp_arg.type + is_out_arg = cpp_arg.name in out_args + if ps.method and cpp_arg.name == "self": + # For method's 'self', we can use 'const Tensor &' and simply ignore mutability! + type_str = "const at::Tensor &" + else: + # For other cases we need prevent dangling refs to temps (unless it's + # unpacked scattered output) + # The reason is explained in the comments above and in 'dispatch_lambda_return_str()'. + # TODO: avoid this special handling? + ensure_temp_safe = len(out_args) <= 1 or not is_out_arg + if ensure_temp_safe: + type_str = { + "at::Tensor &": "at::Tensor", + }.get(type_str, type_str) + return DispatchLambdaArgument( + name=cpp_arg.name, + type_str=type_str, + is_out_arg=is_out_arg, + ) + + return tuple(map(dispatch_lambda_arg, cpp_args)) + + +# [old codegen] XXX: if you got here because of an assertion failure, it doesn't mean +# it's enough to just extend the list here. Before you do this, make sure +# to add an appropriate wrap() overload in torch/csrc/autograd/utils/wrap_outputs.h. +SUPPORTED_RETURN_TYPES = { + "at::Tensor", + "::std::tuple", + "::std::tuple", + "::std::tuple", + "::std::tuple", + "::std::tuple", + "::std::tuple", + "::std::tuple", + "::std::tuple", + "::std::tuple", + "::std::tuple", + "::std::tuple>", + "::std::vector", + # Needed for flash attention forw/backward + "::std::tuple", + "at::Scalar", + "bool", + "int64_t", + "void*", + "void", + "at::QScheme", + "double", + "at::IntArrayRef", + "at::ScalarType", + "at::Stream", +} + + +def dispatch_lambda_return_str(f: NativeFunction) -> str: + # [old codegen] Remove type annotation (e.g. 'Tensor' rather than 'Tensor &') + # because the dispatch lambdas take mutable arguments *by value*, not + # by reference. If you then return a reference to such an argument, you + # will now have a pointer to a dangling stack entry. Not good. + # + # You want: + # + # auto dispatch_selu_ = [](Tensor self) -> Tensor { ...; return at::selu_(self); }; + # ^^^^^^ + # + # *not* + # + # auto dispatch_selu_ = [](Tensor self) -> Tensor& { ...; return at::selu_(self); }; + # ^^^^^^^ + # + # (NB: We can't make dispatch_selu_ take Tensor&, because the enclosing + # codegen looks like dispatch_selu_(_r.tensor(0)), and you can't take a + # mutable reference to temporary. Maybe we could assign it to a + # variable itself.) + returns_without_annotation = tuple( + Return(r.name, r.type, None) for r in f.func.returns + ) + return_str = cpp.returns_type(returns_without_annotation, symint=True).cpp_type() + if return_str not in SUPPORTED_RETURN_TYPES: + raise RuntimeError(f"{f.func.name} returns unsupported type {return_str}") + return return_str + + +def cpp_dispatch_target(f: NativeFunction) -> str: + symint = f.func.has_symint() + name = cpp.name(f.func, symint_overload=symint) + if Variant.method in f.variants: + return f"self.{name}" + if Variant.function in f.variants: + if has_tensor_options(f) or f.func.name.name.base.endswith("_like"): + namespace = "torch" + else: + namespace = "at" + return f"{namespace}::{name}" + raise RuntimeError(f"could not dispatch, neither function nor method: {f.func}") + + +def cpp_dispatch_exprs( + f: NativeFunction, + *, + python_signature: PythonSignature | None = None, +) -> tuple[str, ...]: + cpp_args: Sequence[Binding] = _cpp_signature(f, method=False).arguments() + + exprs: tuple[str, ...] = () + if not isinstance(python_signature, PythonSignatureDeprecated): + # By default the exprs are consistent with the C++ signature. + exprs = tuple(a.name for a in cpp_args) + else: + # For deprecated python signature we may need fill in some constants. + exprs = tuple( + filter( + lambda n: n != "out" or f.func.is_out_fn(), + python_signature.deprecated_args_exprs, + ) + ) + + if Variant.method in f.variants: + exprs = tuple(filter("self".__ne__, exprs)) + + return exprs + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Python / C++ Args Binding +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +# We explicitly enumerate the PythonArgParser unpacking methods for all +# supported types. This might be more verbose than necessary, partially +# because of the irregularity of unpacking method naming, partially +# because we want to mimic the old codegen behavior - to reject +# unexpected and/or unsupported cases which the old codegen rejects. +# For certain cases it is intentionally more restrictive than necessary, +# e.g.: it doesn't accepts doublelist with definite size. +def arg_parser_unpack_method( + t: Type, default: str | None, default_init: str | None, *, symint: bool = True +) -> str: + has_default_init = default_init is not None + if has_default_init and str(t) not in ( + "ScalarType?", + "ScalarType", + "Device", + "Device?", + "Layout", + "Layout?", + "bool", + "bool?", + ): + raise RuntimeError(f"type '{t}' does not supported unpacking with default") + + if isinstance(t, BaseType): + if t.name in [ + BaseTy.Tensor, + BaseTy.Stream, + BaseTy.Storage, + BaseTy.Scalar, + BaseTy.Dimname, + ]: + # These unpack methods line up with their schema names + return t.name.name.lower() + elif t.name == BaseTy.ScalarType: + return "scalartypeWithDefault" if has_default_init else "scalartype" + elif t.name == BaseTy.Device: + return "deviceWithDefault" if has_default_init else "device" + elif t.name == BaseTy.DeviceIndex: + return "toInt64" + elif t.name == BaseTy.int: + return "toInt64" + elif t.name == BaseTy.SymInt: + return "toSymInt" if symint else "toInt64" + elif t.name == BaseTy.bool: + return "toBoolWithDefault" if has_default_init else "toBool" + elif t.name == BaseTy.float: + return "toDouble" + elif t.name == BaseTy.str: + return "stringView" + elif t.name == BaseTy.Layout: + return "layoutWithDefault" if has_default_init else "layout" + elif t.name == BaseTy.MemoryFormat: + return "memoryformat" + + elif isinstance(t, OptionalType): + if str(t.elem) == "Tensor": + return "optionalTensor" + elif str(t.elem) == "Generator": + return "generator" + elif str(t.elem) == "Dimname[]": + return "toDimnameListOptional" + elif not has_default_init and default in ( + None, + "None", + "::std::nullopt", + "std::nullopt", + ): + # If default is None: append 'Optional' to elem's unpacking method + return ( + arg_parser_unpack_method(t.elem, None, None, symint=symint) + "Optional" + ) + else: + # Otherwise, load as underlying type with default + return arg_parser_unpack_method( + t.elem, default, default_init, symint=symint + ) + + elif isinstance(t, ListType): + if str(t.elem) == "Tensor": + # accept and use definite size + return f"tensorlist_n<{t.size}>" if t.size is not None else "tensorlist" + elif str(t.elem) == "Tensor?": + return "list_of_optional_tensors" + elif str(t.elem) == "Dimname": + # accept definite size + return "dimnamelist" + elif str(t.elem) == "int": + # accept definite size + return "intlist" + elif str(t.elem) == "float": + return "doublelist" + elif str(t.elem) == "SymInt": + # accept definite size + return "symintlist" if symint else "intlist" + elif str(t.elem) == "Scalar": + return "scalarlist" + raise RuntimeError(f"type '{t}' is not supported by PythonArgParser") + + +# Return RHS expression for python argument using PythonArgParser output. +# e.g. for arg name 'foo', arg type 'bool', arg_index = 2, returns '_r.toBool(2)' +def arg_parser_output_expr( + arg_index: int, a: PythonArgument, *, symint: bool = True +) -> PythonArgParserOutputExpr: + has_default = a.default_init is not None + unpack_method = arg_parser_unpack_method( + t=a.type, default=a.default, default_init=a.default_init, symint=symint + ) + default = f", {a.default_init}" if has_default else "" + expr = f"_r.{unpack_method}({arg_index}{default})" + + return PythonArgParserOutputExpr( + name=a.name, + expr=expr, + index=arg_index, + argument=a, + ) + + +# Returns a map with key = arg_name and value = PythonArgParserOutputExpr. +def arg_parser_output_exprs( + ps: PythonSignature, f: NativeFunction, *, symint: bool = True +) -> dict[str, PythonArgParserOutputExpr]: + return { + e.name: e + for i, a in enumerate(ps.arguments()) + for e in (arg_parser_output_expr(i, a, symint=symint),) + } + + +# argument name to type for scattered tensor options fields +TENSOR_OPTIONS_FIELDS = { + "dtype": "ScalarType?", + "device": "Device?", + "layout": "Layout?", + "pin_memory": "bool?", + "requires_grad": "bool?", +} + + +# bind arg parser outputs (python args) with dispatch lambda arguments (c++ args). +def dispatch_lambda_exprs( + ps: PythonSignature, f: NativeFunction, *, symint: bool = True +) -> DispatchLambdaArgumentExprs: + # This method is to bind 'arg_parser_outputs' and 'lambda_args' by producing + # 'inits' and 'lambda_args_exprs' for each lambda argument using arg parser + # outputs. + arg_parser_outputs = arg_parser_output_exprs(ps, f, symint=symint) + lambda_args = dispatch_lambda_args(ps, f, symint=symint) + inits: list[str] = [] + lambda_args_exprs: dict[str, str] = {} + + has_toptions = has_tensor_options(f) + + # 1. special inits/unpacking to provide binding exprs for lambda arguments. + for a in ps.arguments(skip_tensor_options=True): + name = a.name + arg_parser_expr = arg_parser_outputs[a.name].expr + + if has_toptions and name == "self": + # TODO: why this needs to be special case? + inits.extend( + [ + f"auto self = {arg_parser_expr};", + ] + ) + lambda_args_exprs[name] = name + elif ( + isinstance(a, PythonOutArgument) + and len(a.outputs) > 1 + and f.func.is_out_fn() + ): + inits.extend( + [ + f"auto out = {arg_parser_expr};", + ] + ) + for i, out_arg in enumerate(a.outputs): + lambda_args_exprs[out_arg.name] = f"out[{i}]" + elif str(a.type) == "Dimname[]?": + # [old codegen] + # TODO: make this part of something more general, or get rid of it. + # optional> are special. The PythonArgParser returns an + # optional>, which cannot be implicitly converted to + # optional>. One needs to unwrap the optional and rewrap. + inits.extend( + [ + f"auto __{name} = {arg_parser_expr};", + f"::std::optional {name} = __{name} ? ::std::make_optional(DimnameList(__{name}.value())) : ::std::nullopt;", # noqa: B950 + ] + ) + lambda_args_exprs[name] = name + else: + # default case - directly using PythonArgParser output expr + lambda_args_exprs[name] = arg_parser_expr + + # method's self is passed directly to python binding, rather than parsed + if ps.method: + lambda_args_exprs["self"] = "self" + + # 2. special packing/checking for TensorOptions. + tensor_options_args_names = [a.name for a in ps.tensor_options_args] + if has_toptions: + if f.func.is_out_fn(): + raise RuntimeError(f"{f.func}: tensor options with output arg") + for a in ps.tensor_options_args: + if a.name not in TENSOR_OPTIONS_FIELDS: + raise RuntimeError( + f"{f.func}: unrecognized tensor options field '{a.name}' in python binding arguments" + ) + if str(a.type) != TENSOR_OPTIONS_FIELDS.get(a.name): + raise RuntimeError( + f"{f.func}: unrecognized type '{str(a.type)}' for tensor options field '{a.name}'" + ) + if not all(a in tensor_options_args_names for a in TENSOR_OPTIONS_FIELDS): + raise RuntimeError( + f"{f.func}: incomplete tensor options args: {tensor_options_args_names}" + ) + + inits.append( + f"""\ +const auto options = TensorOptions() + .dtype({arg_parser_outputs["dtype"].expr}) + .device({arg_parser_outputs["device"].expr}) + .layout({arg_parser_outputs["layout"].expr}) + .requires_grad({arg_parser_outputs["requires_grad"].expr}) + .pinned_memory({arg_parser_outputs["pin_memory"].expr}); +torch::utils::maybe_initialize_device(options); +""" + ) + lambda_args_exprs["options"] = "options" + + # 3. special case - access scattered TensorOptions fields without packing + # TODO: maybe move to the generator side as it's not related to binding. + if not has_toptions and tensor_options_args_names: + if "dtype" in tensor_options_args_names: + # we're an output-arg variant, check these args against output tensor + if not f.func.is_out_fn(): + raise RuntimeError( + f"{f.func}: dtype in tensor_options_args without output arg, {ps} {ps.arguments}" + ) + if not all(a in tensor_options_args_names for a in ("layout", "device")): + raise RuntimeError( + f"{f.func}: incomplete tensor options for output check" + ) + + inits.append( + f"""\ +check_out_type_matches({arg_parser_outputs["out"].expr}, {arg_parser_outputs["dtype"].expr}, + {arg_parser_outputs["dtype"].is_none_expr}, {arg_parser_outputs["layout"].expr}, + {arg_parser_outputs["device"].expr}, {arg_parser_outputs["device"].is_none_expr}); +""" + ) + # we'll set requires_grad on outgoing tensor + if "requires_grad" not in tensor_options_args_names: + raise RuntimeError( + f'{f.func}: expected "requires_grad" in tensor_options_args absent, but found [{tensor_options_args_names}]' + ) + + return DispatchLambdaArgumentExprs( + exprs=tuple(lambda_args_exprs[a.name] for a in lambda_args), + inits=inits, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/structured.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/structured.py new file mode 100644 index 0000000000000000000000000000000000000000..a0e14e5b69e6421fce5ddd247958876061d72b2c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/structured.py @@ -0,0 +1,158 @@ +from __future__ import annotations + +from typing_extensions import assert_never + +from torchgen.api import cpp +from torchgen.api.types import ( + ArgName, + ArrayRefCType, + BaseCType, + Binding, + ConstRefCType, + dimnameListT, + intArrayRefT, + iOptTensorListRefT, + iTensorListRefT, + NamedCType, + OptionalCType, + optionalIntArrayRefT, + optionalScalarRefT, + optionalTensorRefT, + scalarT, + tensorT, +) +from torchgen.model import ( + Argument, + BaseTy, + BaseType, + ListType, + NativeFunctionsGroup, + OptionalType, + SelfArgument, + TensorOptionsArguments, + Type, +) + + +# This file describes the translation of JIT schema to the structured functions API. +# This is similar to native API, but a number of historical problems with native +# API have been fixed. + + +# Translation of types occurring in JIT arguments to a C++ argument type. +# NB: For now, mutable doesn't do anything; but it could if we make +# some more nominal types +def argumenttype_type(t: Type, *, mutable: bool, binds: ArgName) -> NamedCType: + # If it's a value type, do the value type translation + # NB: structured kernels ALWAYS have symint off, since they involve actual + # kernels that require real ints. The one exception is the + # CompositeExplicitAutograd and the meta function (which could + # hypothetically be SymInt), but for simplicity we plan for these to just + # be handled in Python + r = cpp.valuetype_type(t, symint=False, binds=binds, mutable=mutable) + if r is not None: + return r + + if isinstance(t, BaseType): + if t.name == BaseTy.Tensor: + return NamedCType(binds, ConstRefCType(BaseCType(tensorT))) + elif t.name == BaseTy.Scalar: + return NamedCType(binds, ConstRefCType(BaseCType(scalarT))) + else: + raise AssertionError(f"base type should have been value type {t}") + elif isinstance(t, OptionalType): + if t.elem == BaseType(BaseTy.Tensor): + return NamedCType(binds, BaseCType(optionalTensorRefT)) + elif t.elem == BaseType(BaseTy.Scalar): + return NamedCType(binds, BaseCType(optionalScalarRefT)) + elif isinstance(t.elem, ListType) and str(t.elem.elem) == "int": + return NamedCType(binds, BaseCType(optionalIntArrayRefT)) + elem = argumenttype_type(t.elem, mutable=mutable, binds=binds) + return NamedCType(binds, OptionalCType(elem.type)) + elif isinstance(t, ListType): + if t.elem == BaseType(BaseTy.Tensor): + return NamedCType(binds, ConstRefCType(BaseCType(iTensorListRefT))) + elif t.elem == OptionalType(BaseType(BaseTy.Tensor)): + return NamedCType(binds, BaseCType(iOptTensorListRefT)) + # TODO: delete these special cases; see torchgen.api.cpp--these + # must be changed in tandem, but there are problems; see + # https://github.com/pytorch/pytorch/pull/51485 + elif str(t.elem) == "int": + return NamedCType(binds, BaseCType(intArrayRefT)) + elif str(t.elem) == "Dimname": + return NamedCType(binds, BaseCType(dimnameListT)) + elem = argumenttype_type(t.elem, mutable=mutable, binds=binds) + return NamedCType(binds, ArrayRefCType(elem.type)) + else: + raise AssertionError(f"unrecognized type {repr(t)}") + + +def argument_type(a: Argument, *, binds: ArgName) -> NamedCType: + return argumenttype_type(a.type, mutable=a.is_write, binds=binds) + + +# returns_type intentionally omitted, because structured kernels never "return"; +# instead, they always indirectly report their outputs (in the case of a meta +# function, by calling set_output; in the case of an impl function, by writing +# directly into the provided out argument). + + +# Structured kernels are never defaulted +def argument(a: Argument | SelfArgument | TensorOptionsArguments) -> list[Binding]: + if isinstance(a, Argument): + return [ + Binding( + nctype=argument_type(a, binds=a.name), + name=a.name, + default=None, + argument=a, + ) + ] + elif isinstance(a, SelfArgument): + return argument(a.argument) + elif isinstance(a, TensorOptionsArguments): + raise AssertionError("structured kernels don't support TensorOptions yet") + else: + assert_never(a) + + +def impl_arguments(g: NativeFunctionsGroup) -> list[Binding]: + args: list[Argument | TensorOptionsArguments | SelfArgument] = [] + + if g.out.precomputed: + # A list of parameters for the impl function with + # certain parameters replaced with precomputed counterparts + # as specified in native_functions.yaml. + non_out_args_replaced: list[ + Argument | TensorOptionsArguments | SelfArgument + ] = [] + for a in g.out.func.arguments.non_out: + if isinstance(a, Argument) and a.name in g.out.precomputed.replace: + # If a is in precompute.replace, append the parameters + # that should replace it onto non_out_args_replaced. + non_out_args_replaced.extend(g.out.precomputed.replace[a.name]) + else: + # If not, push a as it is. + non_out_args_replaced.append(a) + + args.extend(non_out_args_replaced) + # g.out.precomputed.add is the list of parameters that are added + # without replacement after the non out args and just before the out args + args.extend(g.out.precomputed.add) + else: + args.extend(g.out.func.arguments.non_out) + + args.extend(g.out.func.arguments.out) + return [r for arg in args for r in argument(arg)] + + +def meta_arguments(g: NativeFunctionsGroup) -> list[Binding]: + args: list[Argument | TensorOptionsArguments | SelfArgument] = [] + args.extend(g.functional.func.arguments.non_out) + return [r for arg in args for r in argument(arg)] + + +def out_arguments(g: NativeFunctionsGroup) -> list[Binding]: + args: list[Argument | TensorOptionsArguments | SelfArgument] = [] + args.extend(g.out.func.arguments.out) + return [r for arg in args for r in argument(arg)] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/translate.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/translate.py new file mode 100644 index 0000000000000000000000000000000000000000..f98ce09bbfafb875a619ea01eae7b6f82d76ef71 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/translate.py @@ -0,0 +1,437 @@ +from __future__ import annotations + +from typing import NoReturn, TYPE_CHECKING + +from torchgen.api.types import ( + ArrayRefCType, + BaseCType, + Binding, + boolT, + ConstRefCType, + deviceT, + Expr, + intArrayRefT, + iOptTensorListRefT, + layoutT, + ListCType, + longT, + memoryFormatT, + MutRefCType, + NamedCType, + opmath_t, + OptionalCType, + optionalIntArrayRefT, + optionalScalarRefT, + optionalSymIntArrayRefT, + optionalTensorRefT, + scalar_t, + scalarT, + scalarTypeT, + SpecialArgName, + symIntArrayRefT, + SymIntT, + tensorOptionsT, + tensorT, + VectorCType, +) + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# This file implements a small program synthesis engine that implements +# conversions between one API to another. +# +# The key data type in this file in NamedCType, short for Named C++ semantic type. A NamedCType +# represents a C++ type, plus semantic information about what it represents. +# For example, consider the argument "bool pin_memory"; its normal C++ type is +# "bool", but its C++ semantic type also keeps track that this represents a +# "pin_memory"; you can't just use a random other boolean in a context where you +# need a "pin_memory"! +# +# The translator takes a list of needed NamedCTypes, and then figures out how +# to construct expressions with these NamedCTypes from the given bindings. Many +# of these expressions are trivial (I need a Tensor other; there's a Tensor +# other scope); others are more nontrivial and may require packing/unpacking. +# Some examples of non-trivial action: +# +# - Need the "dtype" binding? Well, maybe "dtype" isn't available +# in the context, instead, "options" is, and you need to extract +# it from there. (Gather) +# +# - Need the "context" binding? Well, maybe "context" isn't available +# in the context, and you need to construct it from "dtype", "device", +# etc. (Scatter) +# +# - Need the "memory_format" binding? Well, actually, it's available +# from both "memory_format" and "options", so you had better make sure +# they are consistent. (Join) + +options_ctype = NamedCType("options", ConstRefCType(BaseCType(tensorOptionsT))) + +out_tensor_ctype = NamedCType("out", ConstRefCType(BaseCType(tensorT))) + +longVec_ctype = VectorCType(BaseCType(longT)) +longSymVec_ctype = VectorCType(BaseCType(SymIntT)) +optionalLongVec_ctype = OptionalCType(VectorCType(BaseCType(longT))) +optionalScalar_ctype = OptionalCType(BaseCType(scalarT)) +optionalTensor_ctype = OptionalCType(BaseCType(tensorT)) + + +class UnsatError(RuntimeError): + pass + + +# Given a set of in-scope bindings and a set of target bindings, synthesize +# a list of expressions that uses only the in-scope bindings (bindings) that +# have all of the types of goals. You may want to use this function if +# you're generating code for a function like: +# +# void f({args}) { +# g({exprs}); // g is a different API +# } +# +# and you need to generate "exprs". +# +# Typically, a list of Bindings is convenient to get (you usually call something +# like arguments() to get them); but technically you only need less information: +# for 'bindings' an (un-ordered) list of Exprs is sufficient; similarly, for +# 'goals', an (ordered) list of NamedCType goals is sufficient. If you are doing +# something more complicated, e.g., tracking the set of bindings in a context, +# you may find using these smaller types more convenient. +def translate( + bindings: Sequence[Expr | Binding], + goals: Sequence[NamedCType | Binding], + *, + method: bool = False, + allow_expensive_conversions: bool = False, +) -> list[Expr]: + binding_exprs: list[Expr] = [] + for b in bindings: + if isinstance(b, Binding): + binding_exprs.append( + Expr( + expr=b.name, + type=b.nctype, + ) + ) + else: + binding_exprs.append(b) + + goal_ctypes: list[NamedCType] = [] + for g in goals: + if isinstance(g, Binding): + goal_ctypes.append(g.nctype) + else: + goal_ctypes.append(g) + + # Add all the bindings to the context + ctx: dict[NamedCType, str] = {} + for b in binding_exprs: + ctx[b.type] = b.expr + + # While we're at it, do some simple forward inference, looking through + # constructors. + # + # NB: When should you do forward inference versus backward inference? + # The general idea: + # + # - Backward inference WHEN the goal gets smaller + # - Forward inference WHEN the hypothesis gets smaller + # + # This helps ensure termination: backward inference starts with a goal + # and tries to make it simpler and simpler until it's trivial; if the + # goal can grow in size, we blow up to a really huge goal size. + # Similarly, with forward inference we take hypotheses and decompose + # them into simpler hypotheses; if hypotheses could expand in size, + # we also have potential nontermination. (In the code below, forward + # inference is only ever carried out at a single step, but you could + # imagine repeated application of forward inference being profitable.) + # + # A good starting point in the literature for exploring more about proof + # search are these lecture notes + # https://www.cs.cmu.edu/~fp/courses/oregon-m10/04-focusing.pdf + # + # TODO: My kingdom for a pattern matcher + # https://www.python.org/dev/peps/pep-0634/ + # + # TODO: This could get us in recomputation trouble if b.expr is nontrivial. + # Fix this by implementing some sort of sharing so that if multiple + # goals share the same expression, we only compute it once. This seems + # to matter in practice as compiler is often unwilling to CSE nontrivial + # expressions like scalar.to() + t = b.type + if ( + isinstance(t, ConstRefCType) + and isinstance(t.elem, OptionalCType) + and isinstance(t.elem.elem, BaseCType) + and str(t.elem.elem.type) == "at::Tensor" + ): + ctx[NamedCType(t.elem.elem.name, ConstRefCType(BaseCType(tensorT)))] = ( + f"({b.expr}.has_value() ? *{b.expr} : at::Tensor())" + ) + + if t.type == ConstRefCType(OptionalCType(BaseCType(tensorT))): + ctx[NamedCType(t.name, BaseCType(optionalTensorRefT))] = ( + f"(({b.expr}.has_value() && (*{b.expr}).defined()) ? at::OptionalTensorRef(*{b.expr}) : at::OptionalTensorRef())" + ) + + if t.type == ConstRefCType(BaseCType(scalarT)): + ctx[NamedCType(t.name, BaseCType(opmath_t))] = f"({b.expr}).to()" + + if t.type == ConstRefCType(OptionalCType(BaseCType(scalarT))): + ctx[NamedCType(t.name, BaseCType(optionalScalarRefT))] = ( + f"({b.expr}.has_value() ? at::OptionalScalarRef(&({b.expr}.value())) : at::OptionalScalarRef())" + ) + + if t.type == BaseCType(scalar_t): + ctx[NamedCType(t.name, BaseCType(opmath_t))] = ( + f"static_cast({b.expr})" + ) + + # [Note: IOptTensorListRef] + if t.type == ConstRefCType(ListCType(OptionalCType(BaseCType(tensorT)))): + ctx[NamedCType(t.name, BaseCType(iOptTensorListRefT))] = ( + f"at::IOptTensorListRef({b.expr})" + ) + + # Add implicit bindings if the generated code is inside a Tensor method + if method: + ctx[NamedCType("self", MutRefCType(BaseCType(tensorT)))] = ( + "const_cast(*this)" + ) + ctx[NamedCType("self", ConstRefCType(BaseCType(tensorT)))] = ( + "const_cast(*this)" + ) + # This is better! Byte-for-byte compat + # ctx[NamedCType("self", ConstRefCType(BaseCType(tensorT)))] = "*this" + + def unsat(goal: NamedCType) -> NoReturn: + ctx_desc = "\n".join( + f" {t.cpp_type()} {t.name}; // {e}" for t, e in ctx.items() + ) + raise UnsatError( + f""" +Failed to synthesize the expression "{goal.cpp_type()} {goal.name}". +When I failed, the following bindings were available in the context: + +{ctx_desc} + +This probably means there is a missing rule in the rules of torchgen.api.translate. +Check this module for more information. +""" + ) + + # A shitty backtracking search implementation. It's shitty because it + # does backtracking via stack (bad idea!) and for the most part tries to + # avoid backtracking. In particular, if + # direct=True, we won't try to do any fancy synthesis, just trivial + # conversions (e.g., "T a" is OK for "const T& a"). So all of the + # existing rules in this function simply try to solve immediately, + # and bail if things don't work out. + def solve(goal: NamedCType, *, direct: bool) -> str: + def direct_solve(goal: NamedCType) -> str: + return solve(goal, direct=True) + + if goal in ctx: + # Trivial + return ctx[goal] + + # const & is satisfied with mutable & + if isinstance(goal.type, ConstRefCType): + try: + # WARNING: not strictly decreasing; be careful not + # to add a direct conversion that goes satisfies + # mutable& with const& + return solve( + NamedCType(goal.name, MutRefCType(goal.type.elem)), direct=direct + ) + except UnsatError: + pass + + # mutable & is satisfied with value + if isinstance(goal.type, MutRefCType): + try: + return solve(NamedCType(goal.name, goal.type.elem), direct=direct) + except UnsatError: + pass + + # TODO: These are referentially equal, shouldn't have to do this; + # ensuring we don't use type synonym IntArrayRef in codegen would + # help + if goal.type == ArrayRefCType(BaseCType(longT)): + return solve(NamedCType(goal.name, BaseCType(intArrayRefT)), direct=direct) + + if direct: + unsat(goal) + + # For now, all of these rules are mutually exclusive. + if goal == NamedCType("memory_format", OptionalCType(BaseCType(memoryFormatT))): + memory_format = direct_solve( + NamedCType( + SpecialArgName.possibly_redundant_memory_format, + OptionalCType(BaseCType(memoryFormatT)), + ) + ) + # No need to join "memory_format" and "options" if the target API takes "options" directly. + # Otherwise it will cause the redundant memory_format error. + if options_ctype in goal_ctypes: + return memory_format + try: + options = direct_solve(options_ctype) + return f"c10::impl::check_tensor_options_and_extract_memory_format({options}, {memory_format})" + except UnsatError: + return memory_format + elif goal == NamedCType("options", BaseCType(tensorOptionsT)): + dtype = direct_solve( + NamedCType("dtype", OptionalCType(BaseCType(scalarTypeT))) + ) + pin_memory = direct_solve( + NamedCType("pin_memory", OptionalCType(BaseCType(boolT))) + ) + device = direct_solve( + NamedCType("device", OptionalCType(BaseCType(deviceT))) + ) + layout = direct_solve( + NamedCType("layout", OptionalCType(BaseCType(layoutT))) + ) + return f"TensorOptions().dtype({dtype}).layout({layout}).device({device}).pinned_memory({pin_memory})" + + elif goal == NamedCType("dtype", OptionalCType(BaseCType(scalarTypeT))): + try: + options = direct_solve(options_ctype) + return f"c10::optTypeMetaToScalarType({options}.dtype_opt())" + except UnsatError: + out_tensor = direct_solve(out_tensor_ctype) + return f"{out_tensor}.scalar_type()" + + elif goal == NamedCType("layout", OptionalCType(BaseCType(layoutT))): + try: + options = direct_solve(options_ctype) + return f"{options}.layout_opt()" + except UnsatError: + out_tensor = direct_solve(out_tensor_ctype) + return f"{out_tensor}.layout()" + + elif goal == NamedCType("device", OptionalCType(BaseCType(deviceT))): + try: + options = direct_solve(options_ctype) + return f"{options}.device_opt()" + except UnsatError: + out_tensor = direct_solve(out_tensor_ctype) + return f"{out_tensor}.device()" + + elif goal == NamedCType("pin_memory", OptionalCType(BaseCType(boolT))): + try: + options = direct_solve(options_ctype) + return f"{options}.pinned_memory_opt()" + except UnsatError: + # If we're calling a factory op from its out= variant, + # We don't actually care about the value of pin_memory. + out_tensor = direct_solve(out_tensor_ctype) + return "::std::nullopt" + + # We can always do translations from value types to reference types, like vector -> IntArrayRef + elif goal.type == BaseCType(intArrayRefT): + try: + return direct_solve(NamedCType(goal.name, longVec_ctype)) + except UnsatError: + # We can also go SymIntArrayRef -> IntArrayRef + symIntArrayRef_type = direct_solve( + NamedCType(goal.name, BaseCType(symIntArrayRefT)) + ) + return f"C10_AS_INTARRAYREF_SLOW({symIntArrayRef_type})" + elif goal.type == BaseCType(symIntArrayRefT): + try: + r = direct_solve(NamedCType(goal.name, BaseCType(intArrayRefT))) + return f"c10::fromIntArrayRefSlow({r})" + except UnsatError: + return direct_solve(NamedCType(goal.name, longSymVec_ctype)) + elif goal.type == BaseCType(SymIntT): + return direct_solve(NamedCType(goal.name, BaseCType(longT))) + elif goal.type == OptionalCType(BaseCType(SymIntT)): + argname = direct_solve( + NamedCType(goal.name, OptionalCType(BaseCType(longT))) + ) + return f"{argname}.has_value() ? ::std::make_optional(c10::SymInt(*{argname})) : ::std::nullopt" + elif goal.type == BaseCType(longT): + symInt_type = direct_solve(NamedCType(goal.name, BaseCType(SymIntT))) + return f"{symInt_type}.guard_int(__FILE__, __LINE__)" + elif goal.type == OptionalCType(BaseCType(longT)): + argname = direct_solve( + NamedCType(goal.name, OptionalCType(BaseCType(SymIntT))) + ) + return f"{argname}.has_value() ? ::std::make_optional({argname}->guard_int(__FILE__, __LINE__)) : ::std::nullopt" + elif goal.type == BaseCType(optionalIntArrayRefT): + try: + return direct_solve(NamedCType(goal.name, optionalLongVec_ctype)) + except UnsatError: + argname = direct_solve( + NamedCType(goal.name, BaseCType(optionalSymIntArrayRefT)) + ) + return f"{argname}.has_value() ? ::std::make_optional(C10_AS_INTARRAYREF_SLOW(*{argname})) : ::std::nullopt" + elif goal.type == BaseCType(optionalSymIntArrayRefT): + # TODO: You might also want to solve this from longSymVec_ctype or + # an optional version of it + argname = direct_solve( + NamedCType(goal.name, BaseCType(optionalIntArrayRefT)) + ) + return f"{argname}.has_value() ? ::std::make_optional(c10::fromIntArrayRefSlow(*{argname})) : ::std::nullopt" + elif goal.type == BaseCType(optionalScalarRefT): + return direct_solve(NamedCType(goal.name, optionalScalar_ctype)) + elif goal.type == BaseCType(optionalTensorRefT): + return direct_solve(NamedCType(goal.name, optionalTensor_ctype)) + + # Note [translation from C++ reference to value types] + # The below cases are all for when we have an argument with a reference type, + # and a corresponding goal with a value type. + # These are needed when we populate the inputs to a lambda capture and we need + # to guarantee the lifetime of each captured argument. + # We guard it with an explicit kwarg because converting to a value type is expensive + # (O(n)) to convert from IntArrayRef to vector), + # so the caller of translate() should be explicit that they need it. + if allow_expensive_conversions: + if goal.type == VectorCType(BaseCType(longT)): + intArrayRef_ctype = NamedCType(goal.name, BaseCType(intArrayRefT)) + argname = direct_solve(intArrayRef_ctype) + return f"{argname}.vec()" + if goal.type == VectorCType(BaseCType(SymIntT)): + symIntArrayRef_ctype = NamedCType(goal.name, BaseCType(symIntArrayRefT)) + argname = direct_solve(symIntArrayRef_ctype) + return f"{argname}.vec()" + elif goal.type == OptionalCType(VectorCType(BaseCType(longT))): + optionalIntArrayRef_ctype = NamedCType( + goal.name, BaseCType(optionalIntArrayRefT) + ) + argname = direct_solve(optionalIntArrayRef_ctype) + return f"{argname}.has_value() ? ::std::make_optional({argname}->vec()) : ::std::nullopt" + elif goal.type == OptionalCType(BaseCType(scalarT)): + optionalScalarRef_ctype = NamedCType( + goal.name, BaseCType(optionalScalarRefT) + ) + argname = direct_solve(optionalScalarRef_ctype) + return f"{argname}.has_value() ? ::std::make_optional({argname}) : ::std::nullopt" + elif goal.type == OptionalCType(BaseCType(scalarT)): + optionalTensorRef_ctype = NamedCType( + goal.name, BaseCType(optionalTensorRefT) + ) + argname = direct_solve(optionalTensorRef_ctype) + return f"{argname}.has_value() ? ::std::make_optional({argname}) : ::std::nullopt" + # Technically, we also need to handle cases of C++ containers holding reference types. + # But there currently aren't any ops that require lambda capture codegen + # With arguments like ::std::vector. + # If that changes, we'll have to add the translation here. + + # We allow const casting on tensors, since const-correctness is a bit broken for at::Tensor. + # We could probably generalize this to non-tensor types too. + if goal.type == MutRefCType(BaseCType(tensorT)): + const_ref_tensor_ctype = NamedCType( + goal.name, ConstRefCType(BaseCType(tensorT)) + ) + argname = direct_solve(const_ref_tensor_ctype) + return f"const_cast({argname})" + + unsat(goal) + + return [Expr(solve(g, direct=False), g) for g in goal_ctypes] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4e98bb8df493f2375b514e6c6aeb897cebe8ec7d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/__init__.py @@ -0,0 +1,5 @@ +from torchgen.api.types.types import * +from torchgen.api.types.types_base import * + + +from torchgen.api.types.signatures import * # usort: skip diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/signatures.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/signatures.py new file mode 100644 index 0000000000000000000000000000000000000000..2eb6e926fc3a327287b73793e223b6d6f0b8d39c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/signatures.py @@ -0,0 +1,358 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from torchgen.api.types.types_base import Binding, CType, Expr + + +if TYPE_CHECKING: + from collections.abc import Iterator, Sequence + + from torchgen.model import ( + BackendIndex, + FunctionSchema, + NativeFunction, + NativeFunctionsGroup, + NativeFunctionsViewGroup, + ) + + +@dataclass(frozen=True) +class CppSignature: + """ + A CppSignature represents a single overload in the C++ API. For + any given function schema, there may be multiple CppSignatures + corresponding to it, based on how we desugar to C++. See also + CppSignatureGroup. + """ + + # The schema this signature is derived from + func: FunctionSchema + + # Is this a C++ signature for a method, i.e. Tensor::my_op(...)? + method: bool + + # Is this a faithful C++ signature (i.e. following the JIT schema) or a convenience API + # (i.e. with a potential TensorOptions argument and out arguments in the front) + faithful: bool + + # Is this a symint C++ signature. For BC reasons, functions that take + # SymInts still present as int64_t in C++, and the SymInt variant is + # offered at a different overload name + # + # NB: If a function RETURNS a SymInt, this is ALWAYS false + symint: bool + + # The set of C++ arguments which should not have defaults applied to them + cpp_no_default_args: set[str] + + # Is this a fallback C++ binding? Fallback bindings are enabled by + # manual_cpp_binding: True and are alternate, non-public API that + # lets manual C++ binding implementers access the binding that would + # have been automatically generated + fallback_binding: bool = False + + # Return the unpacked argument structure of this signature, + # discarding information about which arguments are semantically + # related to each other. + def arguments(self) -> Sequence[Binding]: + return cpp.arguments( + self.func.arguments, + faithful=self.faithful, + symint=self.symint, + method=self.method, + cpp_no_default_args=self.cpp_no_default_args, + ) + + def name(self, *, suppress_symint_suffix: bool = False) -> str: + n = cpp.name( + self.func, + faithful_name_for_out_overloads=self.faithful, + symint_overload=False if suppress_symint_suffix else self.symint, + ) + if self.fallback_binding: + n = f"__dispatch_{n}" + return n + + # Render the C++ declaration for this signature + def decl( + self, + *, + name: str | None = None, + prefix: str = "", + is_redispatching_fn: bool = False, + suppress_symint_suffix: bool = False, + ) -> str: + returns_type = cpp.returns_type( + self.func.returns, symint=self.symint + ).cpp_type() + cpp_args = [a.decl() for a in self.arguments()] + if is_redispatching_fn: + cpp_args = ["c10::DispatchKeySet dispatchKeySet"] + cpp_args + cpp_args_str = ", ".join(cpp_args) + if name is None: + name = prefix + self.name(suppress_symint_suffix=suppress_symint_suffix) + return f"{returns_type} {name}({cpp_args_str})" + + # Render the C++ definition for this signature, not including + # the body (with curly braces) + def defn( + self, + *, + name: str | None = None, + prefix: str = "", + is_redispatching_fn: bool = False, + ) -> str: + returns_type = cpp.returns_type( + self.func.returns, symint=self.symint + ).cpp_type() + cpp_args = [a.defn() for a in self.arguments()] + if is_redispatching_fn: + cpp_args = ["c10::DispatchKeySet dispatchKeySet"] + cpp_args + cpp_args_str = ", ".join(cpp_args) + if name is None: + name = prefix + self.name() + return f"{returns_type} {name}({cpp_args_str})" + + def ptr_type(self) -> str: + args_types_str = ", ".join(a.type for a in self.arguments()) + return f"{cpp.returns_type(self.func.returns, symint=self.symint).cpp_type()} (*)({args_types_str})" + + # Return the C++ function type, e.g., something like int(bool) + def type(self) -> str: + args_types_str = ", ".join(a.type for a in self.arguments()) + return f"{cpp.returns_type(self.func.returns, symint=self.symint).cpp_type()} ({args_types_str})" + + +# Represents group of all CppSignatures associated with a +# FunctionSchema. Right now, that's the regular, user-visible +# signature, as well as a "faithful" signature which doesn't +# have grouping. +@dataclass(frozen=True) +class CppSignatureGroup: + func: FunctionSchema + signature: CppSignature + faithful_signature: CppSignature | None + symint_signature: CppSignature | None + symint_faithful_signature: CppSignature | None + + def most_faithful_signature(self) -> CppSignature: + if self.faithful_signature: + return self.faithful_signature + else: + return self.signature + + def signatures(self, *, symint: bool = True) -> Iterator[CppSignature]: + yield self.signature + if self.faithful_signature: + yield self.faithful_signature + if symint: + if self.symint_signature: + yield self.symint_signature + if self.symint_faithful_signature: + yield self.symint_faithful_signature + + @staticmethod + def from_native_function( + f: NativeFunction, *, method: bool, fallback_binding: bool = False + ) -> CppSignatureGroup: + func = f.func + + def make_sig(*, faithful: bool, symint: bool) -> CppSignature: + return CppSignature( + func=func, + faithful=faithful, + symint=symint, + method=method, + fallback_binding=fallback_binding, + cpp_no_default_args=f.cpp_no_default_args, + ) + + def make_sigs(*, symint: bool) -> tuple[CppSignature, CppSignature | None]: + faithful_signature: CppSignature | None = None + if func.arguments.tensor_options is not None or len(func.arguments.out) > 0: + faithful_signature = make_sig(faithful=True, symint=symint) + signature = make_sig(faithful=False, symint=symint) + return signature, faithful_signature + + signature, faithful_signature = make_sigs(symint=False) + symint_signature: CppSignature | None = None + symint_faithful_signature: CppSignature | None = None + if func.has_symint(): + symint_signature, symint_faithful_signature = make_sigs(symint=True) + + return CppSignatureGroup( + func=func, + signature=signature, + faithful_signature=faithful_signature, + symint_signature=symint_signature, + symint_faithful_signature=symint_faithful_signature, + ) + + +@dataclass(frozen=True) +class DispatcherSignature: + # The schema this signature is derived from + func: FunctionSchema + + # Allows you to prepend an arbitrary prefix to the signature name. + # This is useful for parts of the codegen that generate wrappers around kernels, + # and need to avoid naming collisions. + prefix: str = "" + + symint: bool = True + + def arguments(self) -> list[Binding]: + return dispatcher.arguments(self.func, symint=self.symint) + + def name(self) -> str: + return self.prefix + dispatcher.name(self.func) + + def decl(self, name: str | None = None) -> str: + args_str = ", ".join(a.decl() for a in self.arguments()) + if name is None: + name = self.name() + return f"{self.returns_type().cpp_type()} {name}({args_str})" + + def defn( + self, name: str | None = None, *, is_redispatching_fn: bool = False + ) -> str: + args = [a.defn() for a in self.arguments()] + if is_redispatching_fn: + args = ["c10::DispatchKeySet dispatchKeySet"] + args + args_str = ", ".join(args) + if name is None: + name = self.name() + return f"{self.returns_type().cpp_type()} {name}({args_str})" + + def exprs(self) -> list[Expr]: + return [Expr(a.name, a.nctype) for a in self.arguments()] + + def returns_type(self) -> CType: + return dispatcher.returns_type(self.func.returns, symint=self.symint) + + def ptr_type(self) -> str: + dispatcher_args_types_str = ", ".join(a.type for a in self.arguments()) + return f"{self.returns_type().cpp_type()} (*)({dispatcher_args_types_str})" + + # Return the C++ function type, e.g., something like int(bool) + def type(self) -> str: + dispatcher_args_types_str = ", ".join(a.type for a in self.arguments()) + return f"{self.returns_type().cpp_type()} ({dispatcher_args_types_str})" + + @staticmethod + def from_schema( + func: FunctionSchema, *, prefix: str = "", symint: bool = True + ) -> DispatcherSignature: + return DispatcherSignature(func, prefix, symint) + + +@dataclass(frozen=True) +class NativeSignature: + # The schema this signature is derived from + func: FunctionSchema + + symint: bool + + prefix: str = "" + + def name(self) -> str: + return self.prefix + native.name(self.func) + + def decl(self, name: str | None = None) -> str: + args_str = ", ".join(a.decl() for a in self.arguments()) + if name is None: + name = self.name() + return f"{native.returns_type(self.func.returns, symint=self.symint).cpp_type()} {name}({args_str})" + + def defn(self, name: str | None = None) -> str: + args_str = ", ".join(a.defn() for a in self.arguments()) + if name is None: + name = self.name() + return f"{native.returns_type(self.func.returns, symint=self.symint).cpp_type()} {name}({args_str})" + + def ptr_type(self) -> str: + # don't include defaults in type signature! + args_str = ", ".join(a.defn() for a in self.arguments()) + return f"{native.returns_type(self.func.returns, symint=self.symint).cpp_type()} (*)({args_str})" + + def arguments(self) -> list[Binding]: + return native.arguments(self.func, symint=self.symint) + + def returns_type(self) -> CType: + return native.returns_type(self.func.returns, symint=self.symint) + + def dispatcher_exprs(self) -> list[Expr]: + return translate.translate( + self.arguments(), dispatcher.arguments(self.func), method=False + ) + + +@dataclass(frozen=True) +class ViewInverseSignature: + g: NativeFunctionsViewGroup + + def name(self) -> str: + return functionalization.reverse_name(self.g.view, include_namespace=False) + + def decl(self) -> str: + return_type = functionalization.returns_type(self.g.view.func) + decls = [ + a.decl() + for a in functionalization.op_arguments(self.g.view.func, is_reverse=True) + ] + return f"static {return_type.cpp_type()} {self.name()}({', '.join(decls)});" + + +@dataclass(frozen=True) +class StructuredImplSignature: + g: NativeFunctionsGroup + name: str + + def defn(self, name: str | None = None) -> str: + args_str = ", ".join(a.defn() for a in self.arguments()) + return f"TORCH_IMPL_FUNC({self.name})({args_str})" + + def arguments(self) -> list[Binding]: + return structured.impl_arguments(self.g) + + +# Helper functions + + +def kernel_signature( + f: NativeFunction, backend_index: BackendIndex, *, prefix: str = "" +) -> NativeSignature | DispatcherSignature: + # Note [External Backends Follow Dispatcher API] + # Kernel signatures for in-tree backends follow the "native" API, + # while kernels for out-of-tree backends follow the dispatcher API. + # See the comments in `native.py` for details, but historically there have been + # some small differences in schema convention between them and the Dispatcher API. + # Any differences that require translating between the two will results in a runtime cost, + # so we'd like to keep the differences as small as possible. + # With external backends, we'd like to enforce that they write their kernels with schemas + # that match the Dispatcher API directly, if they can. + meta = backend_index.get_kernel(f) + symint = meta is not None and meta.supports_symint() + if symint: + if not f.func.has_symint(): + raise AssertionError( + f"attempted to define symint kernel for {backend_index.dispatch_key} " + "without SymInt in schema" + ) + if backend_index.external: + return DispatcherSignature.from_schema(f.func, prefix=prefix, symint=symint) + else: + return NativeSignature(f.func, prefix=prefix, symint=symint) + + +# Functions only, no types +from torchgen.api import ( + cpp, + dispatcher, + functionalization, + native, + structured, + translate, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/types.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/types.py new file mode 100644 index 0000000000000000000000000000000000000000..41c05653fffdf3d04fc7078e7df142124ed96e00 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/types.py @@ -0,0 +1,183 @@ +""" +Where should I add a new type? `types_base.py` vs `types.py` + +This file defines data model classes for torchgen typing system, as well as some base types such as int32_t. + +`types.py` defines ATen Tensor type and some c10 types, along with signatures that use these types. + +The difference between these two files, is `types_base.py` should be implementation-agnostic, meaning it shouldn't +contain any type definition that is tight to a specific C++ library (e.g., ATen), so that it can be easily reused +if we want to generate code for another C++ library. + +Add new types to `types.py` if these types are ATen/c10 related. +Add new types to `types_base.py` if they are basic and not attached to ATen/c10. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from torchgen.api.types.types_base import ( + BaseCppType, + BaseCType, + boolT, + byteT, + charT, + CType, + doubleT, + floatT, + int32T, + longT, + shortT, +) +from torchgen.model import BaseTy, ScalarType + + +TENSOR_LIST_LIKE_CTYPES = [ + "at::TensorList", + "const c10::List<::std::optional> &", + "const at::ITensorListRef &", +] + + +halfT = BaseCppType("at", "Half") +complexHalfT = BaseCppType( + "c10", "complex" +) # stuffing template param here is an abuse +complexFloatT = BaseCppType("c10", "complex") +complexDoubleT = BaseCppType("c10", "complex") +bfloat16T = BaseCppType("at", "BFloat16") +float8_e5m2T = BaseCppType("at", "Float8_e5m2") +float8_e5m2fnuzT = BaseCppType("at", "Float8_e5m2fnuz") +float8_e4m3fnT = BaseCppType("at", "Float8_e4m3fn") +float8_e4m3fnuzT = BaseCppType("at", "Float8_e4m3fnuz") +float8_e8m0fnuT = BaseCppType("at", "Float8_e8m0fnu") +stringT = BaseCppType("c10", "string_view") +generatorT = BaseCppType("at", "Generator") +scalarTypeT = BaseCppType("at", "ScalarType") +tensorT = BaseCppType("at", "Tensor") +optionalTensorRefT = BaseCppType("at", "OptionalTensorRef") +tensorListT = BaseCppType("at", "TensorList") +iTensorListRefT = BaseCppType("at", "ITensorListRef") +iOptTensorListRefT = BaseCppType("at", "IOptTensorListRef") +dimnameT = BaseCppType("at", "Dimname") +dimnameListT = BaseCppType("at", "DimnameList") +dimVectorT = BaseCppType("at", "DimVector") +layoutT = BaseCppType("at", "Layout") +deviceT = BaseCppType("at", "Device") +deviceIndexT = BaseCppType("at", "DeviceIndex") +scalarT = BaseCppType("at", "Scalar") +optionalScalarRefT = BaseCppType("at", "OptionalScalarRef") +memoryFormatT = BaseCppType("at", "MemoryFormat") +qschemeT = BaseCppType("at", "QScheme") +storageT = BaseCppType("at", "Storage") +streamT = BaseCppType("at", "Stream") +intArrayRefT = BaseCppType("at", "IntArrayRef") +optionalIntArrayRefT = BaseCppType("at", "OptionalIntArrayRef") +optionalSymIntArrayRefT = BaseCppType("at", "OptionalSymIntArrayRef") +tensorOptionsT = BaseCppType("at", "TensorOptions") +typeAndSizeT = BaseCppType("torch::autograd::generated", "TypeAndSize") +tensorGeometryT = BaseCppType("at", "TensorGeometry") +SymIntT = BaseCppType("c10", "SymInt") +SymBoolT = BaseCppType("c10", "SymBool") +symIntArrayRefT = BaseCppType("c10", "SymIntArrayRef") + +# Types representing template parameters. Technically, we probably shouldn't +# represent them this way in codegen, but it was pretty convenient. +scalar_t = BaseCppType("", "scalar_t") +opmath_t = BaseCppType("", "opmath_t") + +ScalarTypeToCppMapping: dict[ScalarType, BaseCppType] = { + ScalarType.Byte: byteT, + ScalarType.Char: charT, + ScalarType.Short: shortT, + ScalarType.Int: int32T, + ScalarType.Long: longT, + ScalarType.Half: halfT, + ScalarType.Float: floatT, + ScalarType.Double: doubleT, + ScalarType.ComplexHalf: complexHalfT, + ScalarType.ComplexFloat: complexFloatT, + ScalarType.ComplexDouble: complexDoubleT, + ScalarType.Bool: boolT, + ScalarType.Float8_e5m2: float8_e5m2T, + ScalarType.Float8_e5m2fnuz: float8_e5m2fnuzT, + ScalarType.Float8_e4m3fn: float8_e4m3fnT, + ScalarType.Float8_e4m3fnuz: float8_e4m3fnuzT, + ScalarType.Float8_e8m0fnu: float8_e8m0fnuT, +} + +BaseTypeToCppMapping: dict[BaseTy, BaseCppType] = { + BaseTy.int: longT, + BaseTy.float: doubleT, + BaseTy.bool: boolT, + BaseTy.str: stringT, + BaseTy.Generator: generatorT, + BaseTy.ScalarType: scalarTypeT, + BaseTy.Tensor: tensorT, + BaseTy.Dimname: dimnameT, + BaseTy.DimVector: dimVectorT, + BaseTy.Layout: layoutT, + BaseTy.Device: deviceT, + BaseTy.DeviceIndex: deviceIndexT, + BaseTy.Scalar: scalarT, + BaseTy.MemoryFormat: memoryFormatT, + BaseTy.QScheme: qschemeT, + BaseTy.Storage: storageT, + BaseTy.Stream: streamT, + BaseTy.SymInt: SymIntT, + BaseTy.SymBool: SymBoolT, +} + +# CTypes encode C++ type structure as needed for translation. + + +@dataclass(frozen=True) +class OptionalCType(CType): + elem: CType + + def cpp_type(self, *, strip_ref: bool = False) -> str: + # Do not pass `strip_ref` recursively. + return f"::std::optional<{self.elem.cpp_type()}>" + + def remove_const_ref(self) -> CType: + return OptionalCType(self.elem.remove_const_ref()) + + +@dataclass(frozen=True) +class ListCType(CType): + elem: CType + + def cpp_type(self, *, strip_ref: bool = False) -> str: + # Do not pass `strip_ref` recursively. + return f"c10::List<{self.elem.cpp_type()}>" + + def remove_const_ref(self) -> CType: + return ListCType(self.elem.remove_const_ref()) + + +@dataclass(frozen=True) +class ArrayRefCType(CType): + elem: CType + + def cpp_type(self, *, strip_ref: bool = False) -> str: + # Do not pass `strip_ref` recursively. + return f"at::ArrayRef<{self.elem.cpp_type()}>" + + def remove_const_ref(self) -> CType: + return ArrayRefCType(self.elem.remove_const_ref()) + + +@dataclass(frozen=True) +class VectorizedCType(CType): + # This template is explicitly specialized, so the only valid + # elems are those we have specializations for (e.g., float, double, ...) + # scalar_t is also a common argument here (when we are codegen in + # a templated context) + elem: BaseCType + + def cpp_type(self, *, strip_ref: bool = False) -> str: + return f"at::vec::Vectorized<{self.elem.cpp_type()}>" + + def remove_const_ref(self) -> CType: + return self diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/types_base.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/types_base.py new file mode 100644 index 0000000000000000000000000000000000000000..322ae1c39c1ede4a5080b93dbf01ed30589b427c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/types/types_base.py @@ -0,0 +1,238 @@ +""" +Where should I add a new type? `types_base.py` vs `types.py` + +This file defines data model classes for torchgen typing system, as well as some base types such as int32_t. + +`types.py` defines ATen Tensor type and some c10 types, along with signatures that use these types. + +The difference between these two files, is `types_base.py` should be implementation-agnostic, meaning it shouldn't +contain any type definition that is tight to a specific C++ library (e.g., ATen), so that it can be easily reused +if we want to generate code for another C++ library. + +Add new types to `types.py` if these types are ATen/c10 related. +Add new types to `types_base.py` if they are basic and not attached to ATen/c10. +""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from enum import auto, Enum +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from torchgen.model import Argument, SelfArgument, TensorOptionsArguments + + +# An ArgName is just the str name of the argument in schema; +# but in some special circumstances, we may add a little extra +# context. The Enum SpecialArgName covers all of these cases; +# grep for their construction sites to see when they can occur. + + +class SpecialArgName(Enum): + possibly_redundant_memory_format = auto() + + +ArgName = str | SpecialArgName + + +# This class shouldn't be created directly; instead, use/create one of the singletons below. +@dataclass(frozen=True) +class BaseCppType: + ns: str | None + name: str + + def __str__(self) -> str: + if self.ns is None or self.ns == "": + return self.name + return f"{self.ns}::{self.name}" + + +# The set of all non-templated, valid, fully-qualified names of C++ types that are used in the codegen. +# Templated types get their own dataclass, mainly to make namespace parsing easier. +byteT = BaseCppType("", "uint8_t") +charT = BaseCppType("", "int8_t") +shortT = BaseCppType("", "int16_t") +# It would be more symmetric for this to be called intT, but it easy to mix +# this up with JIT int (which is int64_t in C++), so we intentionally don't +# define intT to make it obvious when you've stuffed it up +int32T = BaseCppType("", "int32_t") +longT = BaseCppType("", "int64_t") +doubleT = BaseCppType("", "double") +floatT = BaseCppType("", "float") +boolT = BaseCppType("", "bool") +voidT = BaseCppType("", "void") + + +class CType(ABC): + @abstractmethod + def cpp_type(self, *, strip_ref: bool = False) -> str: + raise NotImplementedError + + @abstractmethod + def remove_const_ref(self) -> CType: + return self + + +@dataclass(frozen=True) +class BaseCType(CType): + type: BaseCppType + + def cpp_type(self, *, strip_ref: bool = False) -> str: + return str(self.type) + + def remove_const_ref(self) -> CType: + return self + + +@dataclass(frozen=True) +class ConstRefCType(CType): + elem: CType + + def cpp_type(self, *, strip_ref: bool = False) -> str: + if strip_ref: + return self.elem.cpp_type(strip_ref=strip_ref) + return f"const {self.elem.cpp_type()} &" + + def remove_const_ref(self) -> CType: + return self.elem.remove_const_ref() + + +@dataclass(frozen=True) +class VectorCType(CType): + elem: CType + + def cpp_type(self, *, strip_ref: bool = False) -> str: + # Do not pass `strip_ref` recursively. + return f"::std::vector<{self.elem.cpp_type()}>" + + def remove_const_ref(self) -> CType: + return VectorCType(self.elem.remove_const_ref()) + + +@dataclass(frozen=True) +class ArrayCType(CType): + elem: CType + size: int + + def cpp_type(self, *, strip_ref: bool = False) -> str: + # Do not pass `strip_ref` recursively. + return f"::std::array<{self.elem.cpp_type()},{self.size}>" + + def remove_const_ref(self) -> CType: + return ArrayCType(self.elem.remove_const_ref(), self.size) + + +@dataclass(frozen=True) +class TupleCType(CType): + elems: list[CType] + + def cpp_type(self, *, strip_ref: bool = False) -> str: + # Do not pass `strip_ref` recursively. + return f"::std::tuple<{','.join([e.cpp_type() for e in self.elems])}>" + + def remove_const_ref(self) -> CType: + return TupleCType([e.remove_const_ref() for e in self.elems]) + + +@dataclass(frozen=True) +class MutRefCType(CType): + elem: CType + + def cpp_type(self, *, strip_ref: bool = False) -> str: + if strip_ref: + return self.elem.cpp_type(strip_ref=strip_ref) + return f"{self.elem.cpp_type()} &" + + def remove_const_ref(self) -> CType: + return self.elem.remove_const_ref() + + +# A NamedCType is short for Named C++ semantic type. A NamedCType represents a C++ type, plus +# semantic information about what it represents. For example, consider the +# argument "bool pin_memory"; its normal C++ type is "bool", but its C++ +# semantic type also keeps track that this represents a "pin_memory"; you can't +# just use a random other boolean in a context where you need a "pin_memory"! +# + + +@dataclass(frozen=True) +class NamedCType: + name: ArgName + type: CType + + def cpp_type(self, *, strip_ref: bool = False) -> str: + return self.type.cpp_type(strip_ref=strip_ref) + + def remove_const_ref(self) -> NamedCType: + return NamedCType(self.name, self.type.remove_const_ref()) + + def with_name(self, name: str) -> NamedCType: + return NamedCType(name, self.type) + + +# A binding represents any C++ binding site for a formal parameter. +# We don't distinguish between binding sites for different APIs; +# instead, all of the important distinctions are encoded in CType, +# which you can use to figure out if a given Binding is appropriate +# for use in another context. (See torchgen.api.translate) + + +@dataclass(frozen=True) +class Binding: + name: str + nctype: NamedCType + argument: Argument | TensorOptionsArguments | SelfArgument + # TODO: maybe don't represent default here + default: str | None = None + + def rename(self, name: str) -> Binding: + return Binding( + name=name, + nctype=self.nctype, + argument=self.argument, + default=self.default, + ) + + @property + def type(self) -> str: + return self.nctype.cpp_type() + + def no_default(self) -> Binding: + return Binding( + name=self.name, + nctype=self.nctype, + default=None, + argument=self.argument, + ) + + def decl(self, *, func_ptr_cast: bool = False) -> str: + mb_default = "" + if self.default is not None: + mb_default = f"={self.default}" + + # casting only needs to know the type + if func_ptr_cast: + return f"{self.type}" + else: + return f"{self.type} {self.name}{mb_default}" + + def defn(self) -> str: + return f"{self.type} {self.name}" + + def with_name(self, name: str) -> Binding: + return Binding( + name=name, nctype=self.nctype, argument=self.argument, default=self.default + ) + + +# An Expr is a C++ expression. It has a C++ string representing its syntax, +# as well as a CType saying what it provides. + + +@dataclass(frozen=True) +class Expr: + expr: str + type: NamedCType diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/ufunc.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/ufunc.py new file mode 100644 index 0000000000000000000000000000000000000000..0ced02fc6d11a5923ef6d0e0086142c34ee31b17 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/ufunc.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torchgen.api.types as api_types +from torchgen.api import cpp, structured +from torchgen.api.types import ( + ArgName, + BaseCppType, + BaseCType, + Binding, + ConstRefCType, + CType, + NamedCType, + scalarT, +) +from torchgen.model import ( + Argument, + BaseTy, + BaseType, + DispatchKey, + FunctionSchema, + NativeFunctionsGroup, + Type, +) + + +def schema_kernel_name(func: FunctionSchema, dispatch_key: DispatchKey) -> str: + if not func.is_out_fn(): + raise AssertionError("ufunc.kernel_name should only be invoked on out schemas") + return f"ufunc_{func.name.name}_{dispatch_key}" + + +def kernel_name(g: NativeFunctionsGroup, dispatch_key: DispatchKey) -> str: + return schema_kernel_name(g.out.func, dispatch_key) + + +# Tensors are omitted (as they are stored in TensorIterator), everything else is +# passed along (technically, we can pass tensors along too, it just wastes +# argument registers) +# +# NB: used for CPU only +def dispatchstub_type(t: Type, *, binds: ArgName) -> NamedCType | None: + # Dispatch stubs are always plain ints + r = cpp.valuetype_type(t, binds=binds, symint=False) + if r is not None: + return r + + if t == BaseType(BaseTy.Scalar): + return NamedCType(binds, ConstRefCType(BaseCType(scalarT))) + elif t == BaseType(BaseTy.Tensor): + return None + else: + raise AssertionError(f"unrecognized type {repr(t)}") + + +def opmath_type(scalar_t: BaseCppType) -> BaseCppType: + if scalar_t == api_types.scalar_t: + return api_types.opmath_t + raise NotImplementedError + + +# NB: Tensors in constructor are stored in opmath_t, not scalar_t +# because Tensor in constructor = its a scalar tensor partially applied = +# it can be higher precision and we want to compute in that higher precision +# +# NB: CUDA only +def ufunctor_ctor_type(t: Type, *, binds: ArgName, scalar_t: BaseCppType) -> NamedCType: + r = cpp.valuetype_type(t, binds=binds, symint=False) + if r is not None: + return r + + if t == BaseType(BaseTy.Scalar): + return NamedCType(binds, BaseCType(opmath_type(scalar_t))) + elif t == BaseType(BaseTy.Tensor): + return NamedCType(binds, BaseCType(opmath_type(scalar_t))) + else: + raise AssertionError(f"unrecognized type {repr(t)}") + + +# Only Tensors ever get passed directly to operator() +# +# NB: CUDA only +# (Actually, this works for CPU too) +def ufunctor_apply_type( + t: Type, *, binds: ArgName, scalar_t: BaseCppType +) -> NamedCType: + if t == BaseType(BaseTy.Tensor): + return NamedCType(binds, BaseCType(scalar_t)) + else: + raise AssertionError(f"unrecognized type {repr(t)}") + + +# The actual ufunc template function the user writes. Everything here +# is done in the computation type. compute_t is opmath_t in CUDA and scalar_t +# in CPU +def ufunc_type(t: Type, *, binds: ArgName, compute_t: CType) -> NamedCType: + r = cpp.valuetype_type(t, binds=binds, symint=False) + if r is not None: + return r + + if t == BaseType(BaseTy.Scalar): + return NamedCType(binds, compute_t) + elif t == BaseType(BaseTy.Tensor): + return NamedCType(binds, compute_t) + else: + raise AssertionError(f"unrecognized type {repr(t)}") + + +def ufunctor_ctor_argument(a: Argument, scalar_t: BaseCppType) -> Binding: + return Binding( + nctype=ufunctor_ctor_type(a.type, binds=a.name, scalar_t=scalar_t), + name=a.name, + default=None, + argument=a, + ) + + +def ufunctor_apply_argument(a: Argument, scalar_t: BaseCppType) -> Binding: + return Binding( + nctype=ufunctor_apply_type(a.type, binds=a.name, scalar_t=scalar_t), + name=a.name, + default=None, + argument=a, + ) + + +def ufunc_argument(a: Argument, compute_t: CType) -> Binding: + return Binding( + nctype=ufunc_type(a.type, binds=a.name, compute_t=compute_t), + name=a.name, + default=None, + argument=a, + ) + + +@dataclass(frozen=True) +class UfunctorBindings: + ctor: list[Binding] + apply: list[Binding] + + +# ufunctors are a CUDA-only concept representing functors that take some of +# their arguments on a host-side constructor, and the rest in the device-side +# apply. E.g., +# +# template +# struct CUDAFunctorOnSelf_add { +# using opmath_t = at::opmath_type; +# opmath_t other_; +# opmath_t alpha_; +# CUDAFunctorOnSelf_add(opmath_t other, opmath_t alpha) : other_(other), alpha_(alpha) {} +# __device__ scalar_t operator()(scalar_t self) { +# return ufunc::add(static_cast(self), other_, alpha_); +# } +# }; +# +# The ctor refers to the constructor CUDAFunctorOnSelf_add, while apply refers +# to the operator() definition +def ufunctor_arguments( + g: NativeFunctionsGroup, *, scalar_tensor_idx: int | None, scalar_t: BaseCppType +) -> UfunctorBindings: + ctor = [] + apply = [] + for a in g.functional.func.arguments.flat_non_out: + if a.type.is_tensor_like(): + if scalar_tensor_idx == 0: + # put it in the ctor anyway + ctor.append(ufunctor_ctor_argument(a, scalar_t=scalar_t)) + scalar_tensor_idx = None + else: + if scalar_tensor_idx is not None: + scalar_tensor_idx -= 1 + apply.append(ufunctor_apply_argument(a, scalar_t=scalar_t)) + else: + ctor.append(ufunctor_ctor_argument(a, scalar_t=scalar_t)) + if scalar_tensor_idx is not None: + raise AssertionError("scalar_tensor_idx should be None at end of processing") + return UfunctorBindings(ctor=ctor, apply=apply) + + +# ufuncs are the inner loop template functions that you wrote in ufunc/add.h +# which do the actual computation in question. E.g., +# +# template +# C10_HOST_DEVICE T add(T self, T other, T alpha) __ubsan_ignore_undefined__ { +# return self + alpha * other; +# } +# +# In this file, we refer to T as compute_t which is bound by caller +def ufunc_arguments(g: NativeFunctionsGroup, *, compute_t: CType) -> list[Binding]: + return [ + ufunc_argument(a, compute_t=compute_t) + for a in g.functional.func.arguments.flat_non_out + ] + + +# Stubs are the DispatchStub trampolines that CPU kernels use to get to their +# vectorized versions. E.g., +# +# using structured_binary_fn_alpha = void(*)(TensorIteratorBase&, const Scalar& alpha); +# DECLARE_DISPATCH(structured_binary_fn_alpha, add_stub); +def stub_arguments(g: NativeFunctionsGroup) -> list[Binding]: + # stubs drop all tensor arguments (they are implicit in the TensorIterator + # argument and keep everything else) + return [ + r + for a in g.out.func.arguments.flat_non_out + if not a.type.is_tensor_like() + for r in structured.argument(a) + ] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/unboxing.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/unboxing.py new file mode 100644 index 0000000000000000000000000000000000000000..edb48ec5d172a7063b4003536506ed33f0f293fa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/api/unboxing.py @@ -0,0 +1,241 @@ +from __future__ import annotations + +from torchgen.api import cpp +from torchgen.api.types import Binding, CppSignatureGroup, CType +from torchgen.model import ( + Argument, + BaseTy, + BaseType, + ListType, + NativeFunction, + OptionalType, + Type, +) + + +# This file generates the code for unboxing wrappers, i.e., the glue logic to unbox a boxed operator and convert the +# ivalues from stack to correct arguments to the unboxed kernel, based on corresponding JIT schema. This codegen is +# an alternative way to generate unboxing wrappers similar to the existing C++ metaprogramming approach but gets the +# job done statically. These generated unboxing wrappers will be useful under the scenario where we need to register +# a fixed set of operators known at compile time and thus can save some time in runtime initialization phase. +# +# Here's an example on how the codegen works: +# +# - Function Schema (source of truth) +# +# aten::empty.names(int[] size, *, Dimname[]? names, +# ScalarType? dtype=None, Layout? layout=None, +# Device? device=None, bool? pin_memory=None, +# MemoryFormat? memory_format=None) -> Tensor +# - Argument Conversion +# Generates C++ code to convert an ivalue (from stack) to its underlying C++ type. +# - int[] size +# ```cpp +# const c10::List size_list_in = (std::move(peek(stack, 0, 7))).toList(); +# +# std::vector size_vec; +# for (c10::IValue size_elem: size_list_in) { +# int64_t size_base = size_elem.to(); +# size_vec.push_back(size_base); +# } +# at::ArrayRef size_list_out(size_vec); +# ~~~~~~~~~~~~~ <-- The converted argument from ivalues in the stack. +# Will be passed to unboxed kernel. +# ``` +# - Dimname[]? names +# ```cpp +# ::std::optional names_opt = (std::move(peek(stack, 1, 7))).toOptional(); +# ::std::optional> names_opt_out; +# if (names_opt.has_value()) { +# ~~~~~~~~~~~ <-- Unwrapping optional shell +# const c10::IValue names_opt_in = names_opt.value(); +# const c10::List names_list_in = names_opt_in.toList(); +# +# std::vector names_vec; +# for (c10::IValue names_elem: names_list_in) { +# ~~~~~~~~~~~~~~~~~~~~~~~~~ <-- Unrolling list, then convert elements one by one. +# at::Dimname names_base = names_elem.to(); +# names_vec.push_back(names_base); +# } +# at::ArrayRef names_list_out(names_vec); +# +# names_opt_out = ::std::optional>(names_list_out); +# } else { +# names_opt_out = ::std::optional>(); +# } +# ``` +# - ScalarType? dtype (similarly for the rest of the arguments) +# ```cpp +# ::std::optional dtype_opt = (std::move(peek(stack, 2, 7))).toOptional(); +# ::std::optional dtype_opt_out; +# if (dtype_opt.has_value()) { +# const c10::IValue dtype_opt_in = dtype_opt.value(); +# at::ScalarType dtype_base = dtype_opt_in.to(); +# ~~~~~~~~~~~~~~~~~~~~ <-- For base types, convert ivalue to it +# directly using ".to()" API. +# dtype_opt_out = ::std::optional(dtype_base); +# } else { +# dtype_opt_out = ::std::optional(); +# } +# ``` +# +# - Unboxed Kernel Call +# ```cpp +# auto result_ = torch::empty( +# size_list_out, +# names_opt_out, +# options, +# memory_format_opt_out +# ); +# ``` +# +# - Push Result Back to Stack +# ```cpp +# drop(stack, 7); +# pack(stack, std::move(result_)); +# ``` +connector = "\n\t" + + +# Return unboxing function name for a NativeFunction +def name(f: NativeFunction) -> str: + return f.func.name.unambiguous_name() + + +# Convert all the arguments in a NativeFunction to C++ code +def convert_arguments(f: NativeFunction) -> tuple[list[Binding], list[str]]: + # we need the 'self' argument so method needs to be False + args = ( + CppSignatureGroup.from_native_function(f, method=False) + .most_faithful_signature() + .arguments() + ) + code_list = [ + f"c10::IValue {args[i].name} = std::move(peek(stack, {i}, {len(args)}));" + for i in range(len(args)) + ] + [""] + binding_list = [] + for arg in args: + # expecting only Argument + if not isinstance(arg.argument, Argument): + raise Exception( # noqa: TRY002 + f"Unexpected argument type, expecting `Argument` but got {arg}" + ) + argument: Argument = arg.argument + unboxed_name, _, code, decl = argumenttype_ivalue_convert( + argument.type, + argument.name, + mutable=argument.is_write, + ) + code_list.extend(decl) + code_list.extend(code) + binding_list.append(arg.with_name(unboxed_name)) + return binding_list, code_list + + +# Takes in the type, name and mutability corresponding to an argument, and generates a tuple of: +# (1) the C++ code necessary to unbox the argument +# (2) A Binding corresponding to the newly created unboxed variable, including variable name and its CType +def argumenttype_ivalue_convert( + t: Type, arg_name: str, *, mutable: bool = False +) -> tuple[str, CType, list[str], list[str]]: + # Unboxing is for mobile, which doesn't care about SymInts + ctype = cpp.argumenttype_type( + t=t, mutable=mutable, binds=arg_name, symint=False + ).type + + if isinstance(t, BaseType): + out_name = f"{arg_name}_base" + code, decl = _gen_code_base_type( + arg_name=arg_name, out_name=out_name, ctype=ctype + ) + elif isinstance(t, OptionalType): + out_name = f"{arg_name}_opt_out" + code, decl = _gen_code_optional_type( + arg_name=arg_name, + out_name=out_name, + t=t, + ctype=ctype, + ) + elif isinstance(t, ListType): + out_name = f"{arg_name}_list_out" + code, decl = _gen_code_list_type( + arg_name=arg_name, + out_name=out_name, + t=t, + ctype=ctype, + ) + else: + raise Exception(f"Cannot handle type {t}. arg_name: {arg_name}") # noqa: TRY002 + return out_name, ctype, code, decl + + +def _gen_code_base_type( + arg_name: str, out_name: str, ctype: CType +) -> tuple[list[str], list[str]]: + return [ + f"{ctype.cpp_type(strip_ref=True)} {out_name} = {arg_name}.to<{ctype.cpp_type(strip_ref=True)}>();" + ], [] + + +def _gen_code_optional_type( + arg_name: str, out_name: str, t: OptionalType, ctype: CType +) -> tuple[list[str], list[str]]: + in_name = f"{arg_name}_opt_in" + res_name, _, res_code, decl = argumenttype_ivalue_convert(t.elem, in_name) + return ( + f""" +auto {arg_name}_opt = {arg_name}.toOptional(); +{ctype.cpp_type(strip_ref=True)} {out_name}; +if ({arg_name}_opt.has_value()) {{ + const c10::IValue {in_name} = {arg_name}_opt.value(); + {connector.join(res_code)} + {out_name} = {ctype.cpp_type(strip_ref=True)}({res_name}); +}} else {{ + {out_name} = {ctype.cpp_type(strip_ref=True)}(); +}} + """.split("\n"), + decl, + ) + + +def _gen_code_list_type( + arg_name: str, out_name: str, t: ListType, ctype: CType +) -> tuple[list[str], list[str]]: + in_name = f"{arg_name}_list_in" + elem_name = f"{arg_name}_elem" + code = [f"const c10::List {in_name} = {arg_name}.toList();"] + res_name, res_ctype, res_code, decl = argumenttype_ivalue_convert(t.elem, elem_name) + # handle list type with size, e.g., bool[4] + if isinstance(t.elem, BaseType) and t.elem.name == BaseTy.bool and t.size: + code.extend( + f""" +{ctype.cpp_type(strip_ref=True)} {out_name} = as_array<{res_ctype.cpp_type(strip_ref=True)}, {t.size}>({in_name}); + """.split("\n") + ) + # we have to use c10::List for optional element. e.g., Tensor?[] -> c10::List<::std::optional> + elif isinstance(t.elem, OptionalType): + code.extend( + f""" +{ctype.cpp_type(strip_ref=True)} {out_name}; +for (c10::IValue {elem_name}: {in_name}) {{ + {connector.join(res_code)} + {out_name}.push_back({res_name}); +}} + """.split("\n") + ) + else: + # use ArrayRef as default. + vec_name = arg_name + "_vec" + # need to bring vector instantiation out of scope so that ArrayRef has valid data + decl.append(f"std::vector<{res_ctype.cpp_type(strip_ref=True)}> {vec_name};") + code.extend( + f""" +for (c10::IValue {elem_name}: {in_name}) {{ + {connector.join(res_code)} + {vec_name}.push_back({res_name}); +}} +{ctype.cpp_type(strip_ref=True)} {out_name}({vec_name}); + """.split("\n") + ) + return code, decl diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/code_template.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/code_template.py new file mode 100644 index 0000000000000000000000000000000000000000..8fcae32c79f53a2c4536987f7236cae00d87312e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/code_template.py @@ -0,0 +1,109 @@ +from __future__ import annotations + +import itertools +import re +import textwrap +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + +# match $identifier or ${identifier} and replace with value in env +# If this identifier is at the beginning of whitespace on a line +# and its value is a list then it is treated as +# block substitution by indenting to that depth and putting each element +# of the list on its own line +# if the identifier is on a line starting with non-whitespace and a list +# then it is comma separated ${,foo} will insert a comma before the list +# if this list is not empty and ${foo,} will insert one after. + + +class CodeTemplate: + substitution_str = r"(^[^\n\S]*)?\$([^\d\W]\w*|\{,?[^\d\W]\w*\,?})" + substitution = re.compile(substitution_str, re.MULTILINE) + + pattern: str + filename: str + + @staticmethod + def from_file(filename: str) -> CodeTemplate: + with open(filename) as f: + return CodeTemplate(f.read(), filename) + + def __init__(self, pattern: str, filename: str = "") -> None: + self.pattern = pattern + self.filename = filename + + def substitute( + self, env: Mapping[str, object] | None = None, **kwargs: object + ) -> str: + if env is None: + env = {} + + def lookup(v: str) -> object: + if env is None: + raise AssertionError("env must be non-None") + return kwargs[v] if v in kwargs else env[v] + + def indent_lines(indent: str, v: Sequence[object]) -> str: + content = "\n".join( + itertools.chain.from_iterable(str(e).splitlines() for e in v) + ) + content = textwrap.indent(content, prefix=indent) + # Remove trailing whitespace on each line + return "\n".join(map(str.rstrip, content.splitlines())).rstrip() + + def replace(match: re.Match[str]) -> str: + indent = match.group(1) + key = match.group(2) + comma_before = "" + comma_after = "" + if key[0] == "{": + key = key[1:-1] + if key[0] == ",": + comma_before = ", " + key = key[1:] + if key[-1] == ",": + comma_after = ", " + key = key[:-1] + v = lookup(key) + if indent is not None: + if not isinstance(v, list): + v = [v] + return indent_lines(indent, v) + elif isinstance(v, list): + middle = ", ".join([str(x) for x in v]) + if len(v) == 0: + return middle + return comma_before + middle + comma_after + else: + return str(v) + + return self.substitution.sub(replace, self.pattern) + + +if __name__ == "__main__": + c = CodeTemplate( + """\ + int foo($args) { + + $bar + $bar + $a+$b + } + int commatest(int a${,stuff}) + int notest(int a${,empty,}) + """ + ) + print( + c.substitute( + args=["hi", 8], + bar=["what", 7], + a=3, + b=4, + stuff=["things...", "others"], + empty=[], + ) + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/context.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/context.py new file mode 100644 index 0000000000000000000000000000000000000000..a99d7119c656f27fabe4accdd2096d997416f4b6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/context.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import contextlib +import functools +from typing import Any, TYPE_CHECKING, TypeVar + +import torchgen.local as local +from torchgen.model import ( + BackendIndex, + DispatchKey, + NativeFunction, + NativeFunctionsGroup, + NativeFunctionsViewGroup, +) +from torchgen.utils import context, S, T + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator + + +# Helper functions for defining generators on things in the model + +F = TypeVar( + "F", + NativeFunction, + NativeFunctionsGroup, + NativeFunctionsViewGroup, + NativeFunction | NativeFunctionsGroup, + NativeFunction | NativeFunctionsViewGroup, +) + +F2 = TypeVar( + "F2", + NativeFunction, + NativeFunctionsGroup, + NativeFunction | None, + bool, + str, +) + +F3 = TypeVar("F3", tuple[NativeFunction, Any], list[NativeFunction]) + + +@contextlib.contextmanager +def native_function_manager( + g: NativeFunctionsGroup | NativeFunctionsViewGroup | NativeFunction, +) -> Iterator[None]: + if isinstance(g, NativeFunctionsGroup): + # By default, we associate all errors with structured native functions + # with the out variant. In some cases, it might be better to have + # a more specific place to hang things; if so, use + # native_function_manager again on the inside + f = g.out + elif isinstance(g, NativeFunctionsViewGroup): + # We associate errors with the view operator + f = g.view + else: + f = g + with context(lambda: f"in native_functions.yaml line {f.loc}:\n {f.func}"): + with local.parametrize( + use_const_ref_for_mutable_tensors=f.use_const_ref_for_mutable_tensors, + use_ilistref_for_tensor_lists=f.part_of_structured_group, + ): + yield + + +# Given a function that operates on NativeFunction, wrap it into a new function +# that sets some appropriate context managers for that native function. +# YOU MUST WRAP FUNCTIONS IN THIS for calls to api modules to be sound +# (you will get an error if we try to access the local variables without having +# set them). +def with_native_function(func: Callable[[F], T]) -> Callable[[F], T]: + @functools.wraps(func) + def wrapper(f: F) -> T: + with native_function_manager(f): + return func(f) + + return wrapper + + +def with_native_function_and(func: Callable[[F, F2], T]) -> Callable[[F, F2], T]: + @functools.wraps(func) + def wrapper(f: F, f2: F2) -> T: + # The first native_function is assumed to be the one with the appropriate context. + with native_function_manager(f): + return func(f, f2) + + return wrapper + + +def method_with_native_function(func: Callable[[S, F], T]) -> Callable[[S, F], T]: + @functools.wraps(func) + def wrapper(slf: S, f: F) -> T: + with native_function_manager(f): + return func(slf, f) + + return wrapper + + +def method_with_nested_native_function( + func: Callable[[S, F3], T], +) -> Callable[[S, F3], T]: + @functools.wraps(func) + def wrapper(slf: S, f: F3) -> T: + with native_function_manager(f[0]): + return func(slf, f) + + return wrapper + + +# Convenience decorator for functions that explicitly take in a BackendIndex, +# instead of indirectly taking one in as a closure +def with_native_function_and_index( + func: Callable[[F, BackendIndex], T], +) -> Callable[[F, BackendIndex], T]: + @functools.wraps(func) + def wrapper(f: F, backend_index: BackendIndex) -> T: + with native_function_manager(f): + return func(f, backend_index) + + return wrapper + + +# Convenience decorator for functions that explicitly take in a Dict of BackendIndices +def with_native_function_and_indices( + func: Callable[[F, dict[DispatchKey, BackendIndex]], T], +) -> Callable[[F, dict[DispatchKey, BackendIndex]], T]: + @functools.wraps(func) + def wrapper(f: F, backend_indices: dict[DispatchKey, BackendIndex]) -> T: + with native_function_manager(f): + return func(f, backend_indices) + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8f08a743ae2dc766530fd8f93be9ebb8b7733f21 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/__init__.py @@ -0,0 +1,19 @@ +from torchgen.dest.lazy_ir import ( + generate_non_native_lazy_ir_nodes as generate_non_native_lazy_ir_nodes, + GenLazyIR as GenLazyIR, + GenLazyNativeFuncDefinition as GenLazyNativeFuncDefinition, + GenLazyShapeInferenceDefinition as GenLazyShapeInferenceDefinition, +) +from torchgen.dest.native_functions import ( + compute_native_function_declaration as compute_native_function_declaration, +) +from torchgen.dest.register_dispatch_key import ( + gen_registration_headers as gen_registration_headers, + gen_registration_helpers as gen_registration_helpers, + RegisterDispatchKey as RegisterDispatchKey, +) +from torchgen.dest.ufunc import ( + compute_ufunc_cpu as compute_ufunc_cpu, + compute_ufunc_cpu_kernel as compute_ufunc_cpu_kernel, + compute_ufunc_cuda as compute_ufunc_cuda, +) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/lazy_ir.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/lazy_ir.py new file mode 100644 index 0000000000000000000000000000000000000000..b0e41a0deadcb7d411c29eb00bfdd61c3f6b2695 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/lazy_ir.py @@ -0,0 +1,718 @@ +from __future__ import annotations + +import itertools +from abc import ABC +from dataclasses import dataclass +from typing import Any + +import torchgen.api.dispatcher as dispatcher +from torchgen.api.lazy import ( + getValueT, + isValueType, + LazyArgument, + LazyIrProperties, + LazyIrSchema, + tensorListValueT, +) +from torchgen.api.translate import translate +from torchgen.api.types import ( + BaseCType, + Binding, + deviceT, + DispatcherSignature, + kernel_signature, + NativeSignature, + OptionalCType, + VectorCType, +) +from torchgen.context import method_with_native_function +from torchgen.dest.lazy_ts_lowering import ts_lowering_body +from torchgen.model import ( + Argument, + BackendIndex, + BackendMetadata, + BaseTy, + BaseType, + FunctionSchema, + ListType, + NativeFunction, + NativeFunctionsGroup, +) + + +def node_ctor_arg_rvalue_string(arg: LazyArgument) -> str: + """ + Given a LazyArgument, + generate a c++ string for materializing an rvalue of that arg for passing into + a lazy Node constructor. + """ + + # TODO: Matching on CType seems wrong; should be matching on Type + if isValueType(arg.lazy_type): + if isinstance(arg.lazy_type, BaseCType): + if arg.is_wrapped_scalar: + return f"node_{arg.name}" + elif arg.lazy_type.type is tensorListValueT: + return f"lazy_{arg.name}_tensorlist" + elif arg.is_symint_or_list: + return f"GetSymIntValue({arg.name})" + return f"lazy_{arg.name}->GetIrValue()" + elif isinstance(arg.lazy_type, OptionalCType): + if arg.is_symint_or_list: + # TODO: I don't understand when you should put lazy_ in the name + # or not + return f"{arg.name} ? std::make_optional(GetSymIntValue(*{arg.name})) : ::std::nullopt" + elif arg.is_wrapped_scalar: + return f"node_{arg.name}" + return ( + f"lazy_{arg.name} ? " + f"std::make_optional(lazy_{arg.name}->GetIrValue()) : " + "::std::nullopt" + ) + else: + raise AssertionError( + f"TODO not sure if there are other valid types to handle here ({arg.lazy_type})" + ) + else: + # NB: this is here because right now we aren't treating SymInt[] as a + # value type; when we do this needs to move above + # NB: we cannot test arg.lazy_type as we've already specified it is an + # int64_t and so we cannot distinguish between SymInt and int64_t + if isinstance(arg.orig_type, ListType) and arg.orig_type.elem == BaseType( + BaseTy.SymInt + ): + if arg.symint: + return f"GetSymIntArrayRefValue({arg.name})" + else: + return f"std::vector({arg.name}.begin(), {arg.name}.end())" + elif isinstance(arg.lazy_type, VectorCType) and isinstance( + arg.lazy_type.elem, BaseCType + ): + return f"std::vector<{arg.lazy_type.elem.type}>({arg.name}.begin(), {arg.name}.end())" + elif ( + isinstance(arg.lazy_type, OptionalCType) + and isinstance(arg.lazy_type.elem, VectorCType) + and isinstance(arg.lazy_type.elem.elem, BaseCType) + ): + return f"torch::lazy::ToOptionalVector<{arg.lazy_type.elem.elem.type}>({arg.name})" + else: + return f"{arg.name}" + + +def node_ctor_inputs(schema: LazyIrSchema) -> str: + """ + Produce a formatted string with the arguments as passed into the constructor of a node class. + """ + node_ctor_values = [ + node_ctor_arg_rvalue_string(arg) for arg in schema.filtered_args() + ] + return ", ".join(node_ctor_values) + + +def gen_fallback_code( + schema: LazyIrSchema, + sig: DispatcherSignature | NativeSignature, + overload_name: str, +) -> str: + """ + Generate code that falls back to eager conditioned on a predicate + """ + dispatcher_sig = DispatcherSignature.from_schema(schema.func) + exprs = translate(sig.arguments(), dispatcher_sig.arguments()) + fallback_args = ",\n ".join([a.expr for a in exprs]) + if len(overload_name): + aten_op_str = f"ATEN_OP2({schema.aten_name}, {overload_name})" + else: + aten_op_str = f"ATEN_OP({schema.aten_name})" + return f""" + if (force_eager_fallback({aten_symbol(schema)})) {{ + return at::native::call_fallback_fn_symint<<c_eager_fallback, {aten_op_str}>::call( + {fallback_args} + ); + }} +""" + + +def aten_symbol(schema: LazyIrSchema) -> str: + missing_interned_strings = { + "sigmoid_backward", + } + if schema.aten_name in missing_interned_strings: + return f'c10::Symbol::fromQualString("aten::{schema.aten_name}")' + + if not schema.aten_name.startswith("at::"): + return f"at::aten::{schema.aten_name}" + else: + return schema.aten_name + + +# converts all tensor-like arguments to meta tensors. Returns: +# (1) a string containing all of the logic that does the conversions. +# (2) a context, to be used by translate(), with all of the relevant bindings. +def convert_to_meta_tensors(sig: DispatcherSignature) -> tuple[str, list[Binding]]: + context: list[Binding] = [] + unwrapped_tensor_args: list[str] = [] + for arg in sig.arguments(): + if isinstance(arg.argument, Argument) and arg.argument.type.is_tensor_like(): + unwrapped_name = f"{arg.name}_meta" + unwrapped_tensor_args.append( + f"auto {unwrapped_name} = to_meta({arg.name});" + ) + context.append(arg.with_name(unwrapped_name)) + else: + context.append(arg) + unwrap_tensor_args_str = "\n ".join(unwrapped_tensor_args) + return unwrap_tensor_args_str, context + + +@dataclass(frozen=True) +class GenLazyIR(ABC): + backend_index: BackendIndex + backend_name: str + node_base: str + use_lazy_shape: bool + + @method_with_native_function + def __call__(self, f: NativeFunctionsGroup | NativeFunction) -> list[str]: + func = f.functional.func if isinstance(f, NativeFunctionsGroup) else f.func + metadata = self.backend_index.get_kernel( + f.functional if isinstance(f, NativeFunctionsGroup) else f + ) + schema = LazyIrSchema( + func, symint=metadata is not None and metadata.supports_symint() + ) + return self.gen(schema) + + # there is no lowering functionality generated unless this IR base class is subclassed and + # implemented as a backend-specific node + def lowering_function(self, schema: LazyIrSchema) -> str: + return "" + + def create_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str: + return "" + + def can_be_reused_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str: + return f"""bool CanBeReused({node_ctor_args}) const {{ + return false; + }}""" + + def node_base_ctor_call(self, schema: LazyIrSchema) -> str: + value_args = schema.filtered_args(values=True, scalars=False) + # backends can customize the way the node base class constructor is called, + # as long as all of its arguments can be generated from information available from the schema + base_ctor_value_args_list = [] + for arg in value_args: + if isinstance(arg.lazy_type, (BaseCType, VectorCType)): + base_ctor_value_args_list.append(f"{arg.name}") + elif isinstance(arg.lazy_type, OptionalCType): + base_ctor_value_args_list.append(f"{arg.name}.value_or(kNullValue)") + else: + raise AssertionError( + f"Unsupported type ({arg.lazy_type}) - add support if necessary" + ) + base_ctor_value_args = ", ".join(base_ctor_value_args_list) + + scalar_args = schema.filtered_args(values=False, scalars=True) + + # Shape construction. + # Conditionally build shape depending on specified shape property + if schema.properties.ShapePrecompute: + shape_ctor_arg = "std::move(shapes)," + elif schema.properties.ShapeCompute: + shape_args = [a.name for a in value_args] + shape_args.extend(a.name for a in scalar_args) + shape_ctor_arg = f"compute_shape_{schema.name}({', '.join(shape_args)})," + elif schema.properties.ShapeCache: + shape_args = [f"operand({i})" for i in range(len(value_args))] + shape_args.extend(a.name for a in scalar_args) + shape_ctor_arg = f"[&](){{ return compute_shape_{schema.name}({', '.join(shape_args)})[0]; }}," + else: + shape_ctor_arg = "" + + scalar_hashes = ", ".join(f"{a.name}" for a in scalar_args) + + return f"""{self.node_base}( + {schema.node_name}::ClassOpKind(), + OpList{{{base_ctor_value_args}}}, + {shape_ctor_arg} + /* num_outputs */ {len(schema.returns)}, + torch::lazy::MHash({scalar_hashes}))""" + + def gen(self, schema: LazyIrSchema) -> list[str]: + opkind = schema.opkind or aten_symbol(schema) + + # for now, we just want one IR class decl and soon after also the method defs + # and we use the functional version not out/inplace. + all_args = schema.filtered_args() + scalar_args = schema.filtered_args(values=False, scalars=True) + + ctor_args = [f"const {i.lazy_type.cpp_type()}& {i.name}" for i in all_args] + reuse_ctor_args = ", ".join(ctor_args) + if self.use_lazy_shape and schema.properties.ShapePrecompute: + ctor_args.append("std::vector&& shapes") + node_ctor_args = ", ".join(ctor_args) + + scalar_initializers = ",\n ".join( + [ + # This code is just special casing the mapping from string_view -> strings + f"{a.name}({a.name}.has_value() ? ::std::make_optional(std::string(*{a.name})) : ::std::nullopt)" + if a.lazy_type.cpp_type() == "::std::optional" + else f"{a.name}({a.name})" + for a in scalar_args + ] + ) + if len(scalar_initializers): + scalar_initializers = f",\n {scalar_initializers}" + scalar_decls = "\n ".join( + [ + f"std::string {a.name};" + if a.lazy_type.cpp_type() == "c10::string_view" + else f"::std::optional {a.name};" + if a.lazy_type.cpp_type() == "::std::optional" + else f"{a.lazy_type.cpp_type()} {a.name};" + for a in scalar_args + ] + ) + optional_values = [ + arg.name + for arg in schema.filtered_args(values=True, scalars=False) + if isinstance(arg.lazy_type, OptionalCType) + ] + has_optional_decls = "\n ".join( + [f"bool has_{value}: 1;" for value in optional_values] + ) + has_optional_defs = "\n ".join( + [f"has_{value} = !!{value};" for value in optional_values] + ) + members_to_string = [] + for arg in scalar_args: + if isinstance(arg.lazy_type, OptionalCType): + value = f"{arg.name}.value()" + if arg.is_generator: + value = '"torch.Generator()"' + members_to_string.append( + f"""if ({arg.name}.has_value()) {{ + ss << ", {arg.name}=" << {value}; + }} else {{ + ss << ", {arg.name}=null"; + }}""" + ) + else: + members_to_string.append(f'ss << ", {arg.name}=" << {arg.name};') + members_to_string_str = "\n ".join(members_to_string) + + return [ + f"""\ +class {schema.node_name} : public {self.node_base} {{ + public: + static torch::lazy::OpKind ClassOpKind() {{ + return torch::lazy::OpKind({opkind}); + }} + + {schema.node_name}({node_ctor_args}) + : {self.node_base_ctor_call(schema)}{scalar_initializers} + {{ + {has_optional_defs} + }} + + std::string ToString() const override {{ + std::stringstream ss; + ss << {self.node_base}::ToString(); + {members_to_string_str} + return ss.str(); + }} + + {self.create_function(schema, reuse_ctor_args)} + + {self.can_be_reused_function(schema, reuse_ctor_args)} + + {self.lowering_function(schema)} + + {scalar_decls} + {has_optional_decls} + +}}; + +""", + ] + + +@dataclass(frozen=True) +class GenTSLazyIR(GenLazyIR): + def lowering_function(self, schema: LazyIrSchema) -> str: + signature = """ + torch::lazy::TSOpVector Lower( + std::shared_ptr function, + torch::lazy::TSLoweringContext* loctx) const override""" + + if schema.properties.LowerDeclOnly: + return f"{signature};" + elif schema.properties.Lower: + return f"""{signature} {{ + {ts_lowering_body(schema)} + }} + """ + else: + return "" + + def create_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str: + signature = f"static NodePtr Create({node_ctor_args})" + if schema.properties.CreateFnDeclOnly: + return f"{signature};" + elif not schema.properties.CreateFn: + return "" + return f"""{signature} {{ + return ReuseOrMakeNode<{schema.node_name}>(data); + }}""" + + def can_be_reused_function(self, schema: LazyIrSchema, node_ctor_args: str) -> str: + signature = f"bool CanBeReused({node_ctor_args}) const" + if schema.properties.CanBeReusedDeclOnly: + return f"{signature};" + elif not schema.properties.CanBeReused: + return "" + value_comparison = [] + for arg in itertools.chain(schema.positional_values, schema.keyword_values): + if isinstance(arg.lazy_type, OptionalCType): + value_comparison.append( + f"nullable_operand(i++) == {arg.name}.value_or(kNullValue)" + ) + else: + value_comparison.append(f"operand(i++) == {arg.name}") + for arg in itertools.chain(schema.positional_scalars, schema.keyword_scalars): + if isinstance(arg.lazy_type, OptionalCType): + value_comparison.append( + f"((!this->{arg.name}&&!{arg.name}) || (this->{arg.name}&&{arg.name} && *(this->{arg.name}) == *{arg.name}))" + ) + else: + value_comparison.append(f"this->{arg.name} == {arg.name}") + value_comparison_str = " &&\n ".join(value_comparison) + + return f"""{signature} {{ + size_t i = 0; + return ({value_comparison_str}); + }}""" + + +@dataclass(frozen=True) +class GenLazyNativeFuncDefinition: + class_method_name: str + backend_index: BackendIndex + tensor_class: str + gen_forced_fallback_code: bool + backend_namespace: str + get_tensorlist: str + get_tensor_or_wrap_number: str + try_get_tensor: str + metrics_counter: str + create_tensor: str + create_from_first_tensor: bool + create_aten_from_ltc_tensor: str + tuple_aten_from_ltc_tensors: str + lazy_tensor_ptr: str + get_device_fn: str + + def lazy_tensor_decls(self, func: NativeFunction, schema: LazyIrSchema) -> str: + value_args = schema.filtered_args(values=True, scalars=False) + # Generates lazy_{name} variables for LazyTensors wrapping input tensors + lazy_tensor_decls: list[str] = [] + for arg in value_args: + if arg.is_wrapped_scalar: + if isinstance(arg.lazy_type, OptionalCType): + lazy_tensor_decls.append( + f"""auto node_{arg.name} = {arg.name} ? + std::make_optional(torch::lazy::LazyGraphExecutor::Get()-> + GetIrValueForScalarFromCodegen(*{arg.name}, *common_device)): + ::std::nullopt;""" + ) + else: + lazy_tensor_decls.append( + f"""auto node_{arg.name} = torch::lazy::LazyGraphExecutor::Get()-> + GetIrValueForScalarFromCodegen({arg.name}, *common_device);""" + ) + elif arg.is_symint_or_list: + continue # values are extracted in isValueType + elif isinstance(arg.lazy_type, BaseCType): + if arg.lazy_type.type is tensorListValueT: + lazy_tensor_decls.append( + f"auto lazy_{arg.name}_tensorlist = " + f"{self.backend_namespace}::{self.get_tensorlist}({arg.name});" + ) + else: + lazy_tensor_decls.append( + f"{self.lazy_tensor_ptr} lazy_{arg.name} = " + f"{self.backend_namespace}::{self.get_tensor_or_wrap_number}({arg.name}, *common_device);" + ) + elif isinstance(arg.lazy_type, OptionalCType): + if arg.lazy_type.elem != BaseCType(getValueT()): + raise AssertionError( + f"Expected OptionalCType elem to be {BaseCType(getValueT())}, " + f"got {arg.lazy_type.elem}" + ) + # TODO(alanwaketan): Maybe we want to apply GetLtcTensorOrCreateForWrappedNumber here, but hold it + # until we encounter a real world example. + lazy_tensor_decls.append( + f"{self.lazy_tensor_ptr} lazy_{arg.name} = " + f"{self.backend_namespace}::{self.try_get_tensor}({arg.name}.value_or(at::Tensor()));" + ) + else: + raise AssertionError( + f"TODO not sure if there are other valid types to handle here ({arg.lazy_type})" + ) + return ("\n ").join(lazy_tensor_decls) + + def force_eager_fallback( + self, + func: NativeFunction, + schema: LazyIrSchema, + metadata: BackendMetadata, + sig: DispatcherSignature | NativeSignature, + ) -> str: + if self.gen_forced_fallback_code: + return gen_fallback_code( + schema, sig, overload_name=func.func.name.overload_name + ) + return "" + + def metrics(self, func: NativeFunction, schema: LazyIrSchema) -> str: + return f"{self.metrics_counter};" + + def get_device(self, func: NativeFunction, schema: LazyIrSchema) -> str: + value_args = schema.filtered_args(values=True, scalars=False) + scalar_args = schema.filtered_args(values=False, scalars=True) + value_types_names = [f"{a.name}" for a in value_args if not a.is_wrapped_scalar] + optional_device = OptionalCType(BaseCType(deviceT)) + optional_devices = [ + a.name for a in scalar_args if a.lazy_type == optional_device + ] + if len(value_types_names) == 0 and len(optional_devices) == 0: + raise AssertionError("Expected at least one Value or Device type") + get_device_str = ( + f"{self.get_device_fn}({', '.join(value_types_names + optional_devices)})" + ) + return f"""auto common_device = {get_device_str}; + TORCH_INTERNAL_ASSERT(common_device); + """ + + def shape_inference(self, func: NativeFunction, schema: LazyIrSchema) -> str: + metadata = self.backend_index.get_kernel(func) + if metadata is None: + raise AssertionError(f"No kernel metadata found for {func.func.name}") + all_args = schema.filtered_args() + returns_length = len(schema.returns) + # call the meta kernel if it exists, to compute output shape/dtype for our IR + # Note [Generated LTC Shape Functions] + # LTC uses meta tensors from core to do shape inference when possible, and otherwise + # we generate a shape function declaration that needs to be manually implemented. + # How do we detect which ops are eligible to use meta tensors? + # In general we should be able to use meta tensors not just on structured operators, + # but also on composite operators that are implemented in terms of structured kernels. + # We don't currently have a way of knowing at codegen time which ops are implemented that way. + # This is the case for all view and view_copy operators however, so we're going to + # use them specifically for all of the view_copy ops (instead of manually writing shape rules for all of them). + is_view_copy_op = "view_copy" in func.tags + is_structured = func.structured or func.structured_delegate is not None + if is_structured or is_view_copy_op: + meta_out = """ +std::vector shapes{torch::lazy::Shape(out_meta.scalar_type(), out_meta.sizes().vec())};""" + if returns_length > 1: + + def this_shape(i: int) -> str: + return f"torch::lazy::Shape(std::get<{i}>(out_meta).scalar_type(), std::get<{i}>(out_meta).sizes().vec())" + + shapes_str = ",".join([this_shape(i) for i in range(returns_length)]) + meta_out = "std::vector shapes{" + shapes_str + "};" + + # Convert tensor args to the meta device and call it. + # (We can't pass in the input tensors directly, because they are "functional wrappers". + # If any of the meta kernels call a tensor op and redispatch, we don't want to hit the functionalize kernels.) + # Even at::meta:: functions might redispatch, e.g. if they call into view ops. + dispatcher_sig = DispatcherSignature.from_schema(func.func) + meta_conversion_str, meta_call_ctx = convert_to_meta_tensors(dispatcher_sig) + meta_call_args = [ + e.expr + for e in translate( + meta_call_ctx, dispatcher_sig.arguments(), method=False + ) + ] + if is_view_copy_op: + # view_copy ops always have a CompositeExplicitAutogradNonFunctional kernel + if not func.has_composite_explicit_autograd_non_functional_kernel: + raise AssertionError( + f"view_copy op {func.func.name} must have " + "CompositeExplicitAutogradNonFunctional kernel" + ) + dispatch_ns = "compositeexplicitautogradnonfunctional" + else: + dispatch_ns = "meta" + aten_name = schema.aten_name + # TODO: this is trolling + if func.func.has_symint() and metadata.supports_symint(): + aten_name += "_symint" + shape_str = f"""\ + {meta_conversion_str} + auto out_meta = at::{dispatch_ns}::{aten_name}({", ".join(meta_call_args)}); + {meta_out}""" + else: + shape_sig = ComputeShapeSignature( + metadata.kernel, func, symint=metadata.supports_symint() + ) + shape_str = f""" + auto shapes = {shape_sig.shape_call};""" + + shape_str += f""" + TORCH_INTERNAL_ASSERT(shapes.size() == {returns_length});""" + + # Calculating which dimensions are symbolic + func_schema_str = "aten::" + str(func.func) + shape_str += f""" + if(torch::lazy::symbolicShapeEnabled()){{ + std::vector inputs = {{ {", ".join(str(a.name) for a in all_args)} }}; + const char* schema_str = "{func_schema_str}"; + applySymbolicShapesOnLT(schema_str, inputs, shapes); + }} + """ + return shape_str + + def build_ir_node(self, func: NativeFunction, schema: LazyIrSchema) -> str: + node_ctor_input_str = node_ctor_inputs(schema) + return f"""torch::lazy::NodePtr node = torch::lazy::ReuseNode<{schema.node_name}>({node_ctor_input_str}); + if (!node) {{ + {self.shape_inference(func, schema)} + node = torch::lazy::MakeNode<{schema.node_name}>({node_ctor_input_str}, std::move(shapes)); + CacheNode(node); + }} + """ + + def create_lazy_tensor(self, first_tensor_name: str | None = None) -> str: + # xla uses an instance method for tensor creation, for the time being + if self.create_from_first_tensor: + # TODO(whc) remove this if XLA switches to using static method for creation + if first_tensor_name is None: + raise AssertionError("Requires first tensor to create lazy tensor") + return f"{first_tensor_name}.{self.create_tensor}" + return f"{self.backend_namespace}::{self.create_tensor}" + + def return_aten_tensor(self, func: NativeFunction, schema: LazyIrSchema) -> str: + returns_length = len(schema.returns) + value_args = schema.filtered_args(values=True, scalars=False) + value_types_names = [f"{a.name}" for a in value_args if not a.is_wrapped_scalar] + first_tensor_name = value_types_names[0] if len(value_types_names) > 0 else None + bridge_str = f"""auto result = {self.create_aten_from_ltc_tensor}( + {self.create_lazy_tensor(first_tensor_name)}(std::move(node), *common_device));""" + + if returns_length > 1: + if len(value_types_names) == 0: + raise AssertionError( + "Code below assumes there is at least one tensor arg" + ) + bridge_str = f"""std::vector<{self.lazy_tensor_ptr}> lazy_tensors; + for (int i = 0; i < {returns_length}; i++) {{ + lazy_tensors.push_back({self.create_lazy_tensor(first_tensor_name)}({getValueT()}(node, i), *common_device)); + }} + auto result = {self.tuple_aten_from_ltc_tensors}<{returns_length}>(lazy_tensors);""" + + if schema.name.name.inplace or func.func.is_out_fn(): + if returns_length != 1: + raise AssertionError( + "We assumed there was no such case where an op is an in-place variant " + f"and has tuple outputs, but got tuple of len {returns_length}." + ) + bridge_str = f"""lazy_{first_tensor_name}->SetInPlaceIrValue(node); + auto& result = {first_tensor_name};""" + + bridge_str += """ + return result;""" + return bridge_str + + @method_with_native_function + def __call__(self, func: NativeFunction) -> list[str]: + sig = kernel_signature(func, self.backend_index) + metadata = self.backend_index.get_kernel(func) + if metadata is None: + raise AssertionError(f"No kernel metadata found for {func.func.name}") + schema = LazyIrSchema(func.func, symint=metadata.supports_symint()) + return [ + f"""\ + {sig.decl(name=f"{self.class_method_name}::{metadata.kernel}")} {{ + {self.force_eager_fallback(func, schema, metadata, sig)} + {self.metrics(func, schema)} + {self.get_device(func, schema)} + {self.lazy_tensor_decls(func, schema)} + {self.build_ir_node(func, schema)} + {self.return_aten_tensor(func, schema)} + }}\n + """ + ] + + +class ComputeShapeSignature: + """ + Here we use the base name as the suffix of the signature to avoid generating for in-place variants. + """ + + def __init__(self, kernel_name: str, f: NativeFunction, *, symint: bool) -> None: + self.__schema = LazyIrSchema(f.func, symint=symint) + self.__dispatch_args = ", ".join( + [a.decl() for a in dispatcher.arguments(f.func, symint=symint)] + ) + self.__call_args = ", ".join( + [f"{arg.name}" for arg in self.__schema.filtered_args(generator=True)] + ) + self.__kernel_name = kernel_name + + def __decl_suffix(self) -> str: + return f"{self.__kernel_name}({self.__dispatch_args})" + + def __call_suffix(self) -> str: + return f"{self.__kernel_name}({self.__call_args})" + + @property + def shape_decl(self) -> str: + return f"TORCH_API std::vector compute_shape_{self.__decl_suffix()}" + + @property + def shape_call(self) -> str: + return f"torch::lazy::compute_shape_{self.__call_suffix()}" + + +@dataclass(frozen=True) +class GenLazyShapeInferenceDefinition: + backend_index: BackendIndex + tensor_class: str + + @method_with_native_function + def __call__(self, f: NativeFunction) -> list[str]: + metadata = self.backend_index.get_kernel(f) + if metadata is None: + raise AssertionError(f"No kernel metadata found for {f.func.name}") + + # See Note [Generated LTC Shape Functions] + is_view_copy_op = "view_copy" in f.tags + is_structured = f.structured or f.structured_delegate is not None + if is_structured or is_view_copy_op: + return [] + else: + shape_sig = ComputeShapeSignature( + metadata.kernel, f, symint=metadata.supports_symint() + ) + return ["\n".join([f"{shape_sig.shape_decl};"])] + + +def generate_non_native_lazy_ir_nodes( + non_native: list[dict[str, Any]], gen_lazy_ir: GenLazyIR +) -> list[str]: + """Generate the non-native lazy IR node classes""" + nodes = [] + for op in non_native: + # Set default properties for Non-Native IRs + properties = LazyIrProperties("ShapeCache", "CanBeReused", "LowerDeclOnly") + for p in op.get("properties", []): + setattr(properties, p, True) + + # non-native is assumed to want symint bindings if you wrote symint + schema = LazyIrSchema(FunctionSchema.parse(op["func"]), properties, symint=True) + schema.opkind = op.get("opkind") + nodes.append(gen_lazy_ir.gen(schema)[0]) + + return nodes diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/lazy_ts_lowering.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/lazy_ts_lowering.py new file mode 100644 index 0000000000000000000000000000000000000000..70161216d8e7c95e194b0d89b345e0da886ef989 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/lazy_ts_lowering.py @@ -0,0 +1,48 @@ +from torchgen.api.lazy import LazyArgument, LazyIrSchema +from torchgen.api.types import OptionalCType + + +def ts_lowering_body(schema: LazyIrSchema) -> str: + # for now, we just want one IR class decl and soon after also the method defs + # and we use the functional version not out/inplace. + emplace_arguments = [] + + def get_value(arg: LazyArgument) -> str: + if isinstance(arg.lazy_type, OptionalCType): + return f"has_{arg.name} ? loctx->GetOutputOp(operand(i++)) : nullptr" + return "loctx->GetOutputOp(operand(i++))" + + for arg in schema.positional_args: + if arg.is_lazy_value: + emplace_arguments.append(get_value(arg)) + continue + emplace_arguments.append(f'"{arg.name}", {arg.name}') + + emplace_arguments_str = "\n ".join( + [f"arguments.emplace_back({a});" for a in emplace_arguments] + ) + emplace_kwarg_values = [ + f'"{arg.name}", {get_value(arg)}' for arg in schema.keyword_values + ] + emplace_kwarg_scalars = [ + f'"{arg.name}", {arg.name}' for arg in schema.keyword_scalars + ] + emplace_kwarguments = "\n ".join( + [ + f"kwarguments.emplace_back({a});" + for a in emplace_kwarg_values + emplace_kwarg_scalars + ] + ) + return f"""\ + std::vector arguments; + std::vector kwarguments; + arguments.reserve({len(emplace_arguments)}); + kwarguments.reserve({len(emplace_kwarg_values + emplace_kwarg_scalars)}); + size_t i = 0; + {emplace_arguments_str} + {emplace_kwarguments} + torch::lazy::TSOpVector {schema.aten_name}_out = torch::lazy::LowerTSBuiltin(function, op().op, arguments, kwarguments); + TORCH_CHECK_EQ({schema.aten_name}_out.size(), {len(schema.returns)}); + + return {schema.aten_name}_out; +""" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/native_functions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/native_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..05e252d09f9c16888dec66045a92b8aefa19b667 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/native_functions.py @@ -0,0 +1,84 @@ +from __future__ import annotations + +import torchgen.api.meta as meta +import torchgen.api.structured as structured +from torchgen.api.types import kernel_signature +from torchgen.context import with_native_function_and_index +from torchgen.model import BackendIndex, NativeFunction, NativeFunctionsGroup +from torchgen.utils import mapMaybe + + +def torch_api_key_word_prefix(bankend_index: BackendIndex) -> str: + if bankend_index.external: + return "" + + # Although Intel GPU ATen library is out-of-tree, it still utilizes torchgen to produce structured + # kernels. Regarding these produced structured kernels, they should be visible for the Intel GPU ATen + # library. Therefore, we need to add "TORCH_XPU_API" prefix to these structured kernels, + # rather than "TORCH_API". Because the semantic of "TORCH_API" is "hidden" for out-of-tree backends. + # For other in-tree backends like cpu and cuda, they still use "TORCH_API" prefix with "visible" semantic. + device_torch_api_key_word_mapping = { + "XPU": "TORCH_XPU_API", + } + + return ( + device_torch_api_key_word_mapping.get( + bankend_index.dispatch_key.name, "TORCH_API" + ) + + " " + ) + + +@with_native_function_and_index +def gen_unstructured(f: NativeFunction, backend_index: BackendIndex) -> str | None: + sig = kernel_signature(f, backend_index) + metadata = backend_index.get_kernel(f) + if metadata is None: + return None + if "legacy::" in metadata.kernel: + return None + else: + prefix = "static" if backend_index.external else "TORCH_API" + return f"{prefix} {sig.decl(name=metadata.kernel)};" + + +@with_native_function_and_index +def gen_structured(g: NativeFunctionsGroup, backend_index: BackendIndex) -> list[str]: + meta_name = meta.name(g) + out_args = structured.impl_arguments(g) + metadata = backend_index.get_kernel(g) + if metadata is None: + return [] + prefix = torch_api_key_word_prefix(backend_index) + return [ + f"""\ +struct {prefix}structured_{metadata.kernel} : public at::meta::structured_{meta_name} {{ +void impl({", ".join(a.decl() for a in out_args)}); +}}; +""" + ] + + +# Generates NativeFunctions.h, a list of forward declarations of all +# actual kernel definitions we keep in aten/src/ATen/native/ +@with_native_function_and_index +def compute_native_function_declaration( + g: NativeFunctionsGroup | NativeFunction, backend_index: BackendIndex +) -> list[str]: + metadata = backend_index.get_kernel(g) + if isinstance(g, NativeFunctionsGroup): + if metadata is not None and metadata.structured: + if backend_index.external: + # Structured hasn't been tested with external backends yet. + raise AssertionError( + "Structured external backend functions are not implemented yet." + ) + else: + return gen_structured(g, backend_index) + else: + return list( + mapMaybe(lambda f: gen_unstructured(f, backend_index), g.functions()) + ) + else: + x = gen_unstructured(g, backend_index) + return [] if x is None else [x] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/register_dispatch_key.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/register_dispatch_key.py new file mode 100644 index 0000000000000000000000000000000000000000..55694fb602f51e50e66ede5002e79ab72560bdfa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/register_dispatch_key.py @@ -0,0 +1,1035 @@ +from __future__ import annotations + +import itertools +import textwrap +from dataclasses import dataclass +from typing import Literal, TYPE_CHECKING +from typing_extensions import assert_never + +import torchgen.api.cpp as cpp +import torchgen.api.meta as meta +import torchgen.api.structured as structured +from torchgen.api.translate import translate +from torchgen.api.types import ( + BaseCType, + Binding, + ConstRefCType, + CppSignature, + CppSignatureGroup, + DispatcherSignature, + Expr, + kernel_signature, + MutRefCType, + NamedCType, + NativeSignature, + tensorT, +) +from torchgen.context import method_with_native_function, native_function_manager +from torchgen.model import ( + Argument, + BackendIndex, + DeviceCheckType, + DispatchKey, + gets_generated_out_inplace_wrapper, + is_cuda_dispatch_key, + NativeFunction, + NativeFunctionsGroup, + SchemaKind, + TensorOptionsArguments, +) +from torchgen.utils import mapMaybe, Target + + +if TYPE_CHECKING: + from torchgen.selective_build.selector import SelectiveBuilder + + +def gen_registration_headers( + backend_index: BackendIndex, + per_operator_headers: bool, + rocm: bool, +) -> list[str]: + if per_operator_headers: + headers = ["#include "] + else: + headers = ["#include "] + + if backend_index.dispatch_key in (DispatchKey.CPU, DispatchKey.Meta): + headers.append("#include ") + elif backend_index.dispatch_key == DispatchKey.CUDA: + if rocm: + headers.append("#include ") + else: + headers.append("#include ") + elif backend_index.dispatch_key == DispatchKey.MPS: + headers.append("#include ") + elif backend_index.dispatch_key == DispatchKey.XPU: + # XPU specific, this header resides in third_party/torch-xpu-ops + headers.append("#include ") + elif backend_index.dispatch_key == DispatchKey.MTIA: + headers.append("#include ") + elif per_operator_headers: + headers += [ + "#include ", + "#include ", + "#include ", + "#include ", + ] + else: + headers.append("#include ") + + headers.append("#include ") + return headers + + +def gen_empty_impl_names( + backend_index: BackendIndex, +) -> tuple[str | None, str | None]: + empty_impl = None + empty_strided_impl = None + + if backend_index.dispatch_key in ( + DispatchKey.Meta, + DispatchKey.CPU, + DispatchKey.CUDA, + DispatchKey.MPS, + DispatchKey.XPU, + DispatchKey.MTIA, + ): + dispatch = str(backend_index.dispatch_key).lower() + empty_impl = f"at::detail::empty_{dispatch}" + empty_strided_impl = f"at::detail::empty_strided_{dispatch}" + elif backend_index.dispatch_key in ( + DispatchKey.CompositeExplicitAutogradNonFunctional, + DispatchKey.QuantizedCPU, + DispatchKey.QuantizedCUDA, + DispatchKey.XPU, + ): + empty_impl = "at::empty" + empty_strided_impl = "at::empty_strided" + + return empty_impl, empty_strided_impl + + +def gen_create_out_helper(backend_index: BackendIndex) -> list[str]: + if backend_index.dispatch_key == DispatchKey.Meta: + empty_options = "options.device(at::kMeta)" + else: + empty_options = "options" + + empty_impl, empty_strided_impl = gen_empty_impl_names(backend_index) + if empty_impl is None: + return [] + + return [ + f""" +Tensor create_out(IntArrayRef sizes, IntArrayRef strides, const TensorOptions &options) {{ + if (strides.empty()) {{ + return {empty_impl}(sizes, {empty_options}); + }} else {{ + return {empty_strided_impl}(sizes, strides, {empty_options}); + }} +}} +""" + ] + + +def gen_maybe_create_proxy_helper(backend_index: BackendIndex) -> list[str]: + _, empty_strided_impl = gen_empty_impl_names(backend_index) + return ( + [] + if empty_strided_impl is None + else [ + f""" +std::optional maybe_create_proxy(const Tensor &out, IntArrayRef sizes, IntArrayRef strides, const TensorOptions &options) {{ + if (out.strides() != strides) {{ + return {empty_strided_impl}(sizes, strides, options); + }} + return std::nullopt; +}} +""" + ] + ) + + +def gen_resize_out_helper(backend_index: BackendIndex) -> list[str]: + if backend_index.dispatch_key == DispatchKey.CompositeExplicitAutogradNonFunctional: + # The function isn't used by this key (since only functional ops have a kernel for this key), + # so we need to not include it to avoid a defined-but-not-used error. + return [] + return [ + """ +void resize_out(const Tensor &out, IntArrayRef sizes, IntArrayRef strides, const TensorOptions &options) { + TORCH_CHECK(options.dtype() == out.dtype(), + "Expected out tensor to have dtype ", options.dtype(), ", but got ", out.dtype(), " instead"); + TORCH_CHECK(options.device() == out.device(), + "Expected out tensor to have device ", options.device(), ", but got ", out.device(), " instead"); + const bool resized = at::native::resize_output(out, sizes); + // Only restride if a resize occurred; otherwise we ignore the (advisory) + // strides from the meta function and directly use the output tensor's + // preexisting strides + if (resized) { + if (!strides.empty()) { + TORCH_INTERNAL_ASSERT(!options.memory_format_opt().has_value()); + // TODO: avoid the redispatch here + out.as_strided_(sizes, strides); + } else if (options.memory_format_opt().has_value()) { + out.unsafeGetTensorImpl()->empty_tensor_restride(*options.memory_format_opt()); + } + } +} +""" + ] + + +def gen_check_inplace_helper(backend_index: BackendIndex) -> list[str]: + return [ + """ +void check_inplace(const Tensor &self, IntArrayRef sizes, const TensorOptions &options) { + // These checks are needed on those operators that: + // 1) don't use 'TensorIterator' (e.g. 'addmm' and 'baddbmm') + // 2) have particular typing rules (e.g. 'cumsum' and 'cumprod') + // For other operators (e.g. 'add'), 'TensorIterator' already checks + // these things separately. + TORCH_CHECK(options.dtype() == self.dtype(), + "Bad in-place call: ", + "input tensor dtype ", self.dtype(), " and output tensor dtype ", options.dtype(), " should match"); + TORCH_CHECK(options.device() == self.device(), + "Bad in-place call: ", + "input tensor device ", self.device(), " and output tensor device ", options.device(), " should match"); + TORCH_CHECK(sizes == self.sizes(), + "Bad in-place call: ", + "input tensor size ", self.sizes(), " and output tensor size ", sizes, " should match"); +} +""" + ] + + +def gen_registration_helpers(backend_index: BackendIndex) -> list[str]: + return [ + 'C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wunused-function")', + *gen_create_out_helper(backend_index), + *gen_resize_out_helper(backend_index), + *gen_check_inplace_helper(backend_index), + *gen_maybe_create_proxy_helper(backend_index), + "C10_DIAGNOSTIC_POP()", + ] + + +# Generates Register{dispatch}.cpp (e.g., RegisterCPU.cpp). +# +# - The primary function of this file is to register all of the +# implementations for the given dispatch key to the dispatcher, +# so they are available for use in PyTorch. If dispatch is +# None, we generate schema (def) registrations and catchall +# registrations. +# - The secondary function of this file is to generate a wrapper +# around functions. In CPUType these wrappers do nothing +# (and should be removed), but in other cases they handle +# DeviceGuard. A small extra benefit of wrappers is they +# are not overloaded, so they can be used in the registration +# API without having to disambiguate which overload you want +# (as would be the case if you directly registered native:: +# functions). +# - The tertiary function of this file is to generate *static* +# cpp API bindings which can be used to bypass dispatcher +# directly to kernels, but with user-friendly cpp-style API +@dataclass(frozen=True) +class RegisterDispatchKey: + backend_index: BackendIndex + + target: Literal[ + Target.ANONYMOUS_DEFINITION, + Target.NAMESPACED_DEFINITION, + Target.NAMESPACED_DECLARATION, + Target.REGISTRATION, + ] + + # Selector object to determine which operators to generate + # registration code for. + selector: SelectiveBuilder + + # Whether or not we are actually code-genning for ROCm + rocm: bool + + # Whether or not to generate symint registrations or not. External users + # of codegen who don't care about symints can set this to false to get + # non-SymInt codegen + symint: bool + + # The class that all unstructured native functions live under. This is used to improve + # compiler error messages when a kernel writer adds a native function with the wrong signature. + # This is only used in unstructured kernels, since structured kernels already live in a class. + # Finally, this field is currently Optional because it is only used by external backends. + # It would be nice if we can add the same logic to in-tree kernels too, but that requires updating + # all of the existing kernel signatures scattered across aten/src/ATen/native. + class_method_name: str | None + + # Only set to true in lightweight dispatch. If lightweight dispatch is enabled we are registering + # operators into JIT op registry, thus we need to avoid generating code to register into the dispatcher. + skip_dispatcher_op_registration: bool + + @staticmethod + def gen_device_check( + type: DeviceCheckType, args: list[Argument], method_name: str + ) -> str: + if type == DeviceCheckType.NoCheck: + return " // No device check\n" + + device_check = "std::optional common_device = std::nullopt;\n" + device_check += "(void)common_device; // Suppress unused variable warning\n" + for arg in args: + # Only tensor like arguments are eligible + if arg.type.is_tensor_like(): + device_check += f""" + c10::impl::check_and_update_common_device(common_device, {arg.name}, "{method_name}", "{arg.name}");""" + return device_check + + @method_with_native_function + def __call__(self, f: NativeFunctionsGroup | NativeFunction) -> list[str]: + if isinstance(f, NativeFunctionsGroup): + g: NativeFunctionsGroup = f + # Note: We call gen_structured() if the operator is marked structured, regardless of the backend. + # gen_structured() has special logic to handle auto-generated kernels. + if g.structured: + return self.gen_structured(g) + else: + return list( + mapMaybe(lambda f: self.gen_unstructured(f, g), g.functions()) + ) + elif isinstance(f, NativeFunction): + r = self.gen_unstructured(f) + return [] if r is None else [r] + else: + assert_never(f) + + def wrapper_kernel_sig( + self, f: NativeFunction + ) -> NativeSignature | DispatcherSignature: + # The prefix is just to ensure uniqueness. The Dispatcher API doesn't guarantee unique kernel names. + return DispatcherSignature.from_schema( + f.func, + prefix=f"wrapper_{self.backend_index.dispatch_key}_{f.func.name.overload_name}_", + symint=self.symint, + ) + + def gen_out_inplace_wrapper( + self, f: NativeFunction, g: NativeFunctionsGroup | None + ) -> str | None: + if g is None: + return None + k = f.func.kind() + if k is SchemaKind.inplace: + copy_op = "at::_copy_from" + elif k is SchemaKind.out: + copy_op = "at::_copy_from_and_resize" + else: + raise AssertionError("gen_out_inplace_wrapper called on a functional op") + + sig = self.wrapper_kernel_sig(f) + name = sig.name() + + func_res = f"{name}_tmp" + return_names = cpp.return_names(f) + if len(return_names) > 1: + updates = "\n ".join( + f"{copy_op}(std::get<{i}>({func_res}), {ret_name});" + for i, ret_name in enumerate(return_names) + ) + returns = f"{sig.returns_type().cpp_type()}({', '.join(return_names)})" + elif len(return_names) == 1: + ret_name = return_names[0] + updates = f"{copy_op}({func_res}, {ret_name});" + returns = ret_name + else: + if len(f.func.arguments.out) != 1: + raise AssertionError( + f"Expected exactly 1 out argument, got {len(f.func.arguments.out)}" + ) + returns = "" + out_arg = f.func.arguments.out[0] + if out_arg.type.is_list_like(): + updates = f"""\ + for (int64_t i = 0; i < {func_res}.size(); ++i) {{ + {copy_op}({func_res}[i], {out_arg.name}[i]); + }}""" + else: + updates = f"{copy_op}({func_res}, {out_arg.name});" + + functional_sig = self.wrapper_kernel_sig(g.functional) + wrapper_name = sig.name() + + return f"""\ +{sig.defn(name=wrapper_name)} {{ + auto {func_res} = {functional_sig.name()}({", ".join(e.expr for e in translate(sig.arguments(), functional_sig.arguments()))}); + {updates} + return {returns}; +}} +""" + + def gen_structured(self, g: NativeFunctionsGroup) -> list[str]: + metadata = self.backend_index.get_kernel(g) + if self.backend_index.dispatch_key == DispatchKey.Meta: + if self.backend_index.has_kernel(g.out): + raise AssertionError( + "Do not explicitly specify Meta dispatch key on structured " + "functions, they will be automatically generated for you" + ) + elif ( + self.backend_index.dispatch_key + == DispatchKey.CompositeExplicitAutogradNonFunctional + ): + if self.backend_index.has_kernel(g.out): + raise AssertionError( + "Do not explicitly specify CompositeExplicitAutograd dispatch key on " + "structured functions, they will be automatically generated for you" + ) + elif metadata is None or not metadata.structured: + return list(mapMaybe(lambda f: self.gen_unstructured(f, g), g.functions())) + structured_gen = StructuredRegisterDispatchKey( + self.backend_index, + self.target, + self.selector, + self.rocm, + self.symint, + self.class_method_name, + self.skip_dispatcher_op_registration, + g, + ) + return list(mapMaybe(structured_gen.gen_one, g.functions())) + + def gen_unstructured( + self, f: NativeFunction, g: NativeFunctionsGroup | None = None + ) -> str | None: + with native_function_manager(f): + inplace_meta = False + gets_out_inplace_wrapper = False + if not self.backend_index.has_kernel(f): + if ( + self.backend_index.dispatch_key == DispatchKey.Meta + and f.func.kind() is SchemaKind.inplace + and + # Defer to composites for meta implementation + not f.has_composite_kernel + and + # Inplace list operations are not supported + len(f.func.returns) == 1 + ): + inplace_meta = True + elif ( + not self.backend_index.use_out_as_primary + and g is not None + and gets_generated_out_inplace_wrapper(f, g, self.backend_index) + ): + # We want to generate inplace/out wrappers, that don't have a kernel for the backend. + gets_out_inplace_wrapper = True + else: + return None + if f.manual_kernel_registration: + return None + + if ( + self.target is Target.REGISTRATION + and not self.selector.is_native_function_selected(f) + ): + return None + + sig = self.wrapper_kernel_sig(f) + + name = sig.name() + returns_type = sig.returns_type().cpp_type() + args = sig.arguments() + args_str = ", ".join(a.defn() for a in args) + + # See Note [Direct dispatch bindings] + cpp_sig_group = CppSignatureGroup.from_native_function( + f, method=False, fallback_binding=False + ) + + # TODO: dedupe this with the structured codegen + if self.target is Target.NAMESPACED_DECLARATION: + result = "" + for cpp_sig in cpp_sig_group.signatures(symint=self.symint): + result += f"TORCH_API {cpp_sig.decl()};\n" + return result + elif self.target is Target.NAMESPACED_DEFINITION: + + def generate_defn(cpp_sig: CppSignature) -> str: + return f""" +{cpp_sig.defn()} {{ +return {sig.name()}({", ".join(e.expr for e in translate(cpp_sig.arguments(), sig.arguments()))}); +}} +""" + + result = "" + for cpp_sig in cpp_sig_group.signatures(symint=self.symint): + result += generate_defn(cpp_sig) + return result + + elif self.target is Target.ANONYMOUS_DEFINITION: + # short circuit for inplace_meta + if inplace_meta: + if f.func.arguments.self_arg is None: + raise AssertionError( + "Expected self_arg to be non-None for inplace_meta" + ) + self_arg_name = f.func.arguments.self_arg.argument.name + # TODO: handle in place on tensor list + return f""" +{returns_type} {name}({args_str}) {{ + TORCH_CHECK_NOT_IMPLEMENTED({self_arg_name}.is_meta(), + "Cannot inplace into non-meta tensor with meta tensor argument"); + return {self_arg_name}; +}} +""" + + # short circuit for generated inplace/out wrappers + if gets_out_inplace_wrapper: + return self.gen_out_inplace_wrapper(f, g) + + metadata = self.backend_index.get_kernel(f) + if metadata is None: + return None + if self.class_method_name is None: + impl_name = f"{metadata.cpp_namespace}::{metadata.kernel}" + else: + impl_name = f"{metadata.cpp_namespace}::{self.class_method_name}::{metadata.kernel}" + + kernel_sig = kernel_signature(f, self.backend_index) + + args_exprs_str = ", ".join( + e.expr + for e in translate( + sig.arguments(), kernel_sig.arguments(), method=False + ) + ) + + device_check = " // No device check\n" + # Backends that require device guards presumably also require device checks. + if self.backend_index.device_guard: + device_check_args = itertools.chain( + f.func.arguments.out, f.func.arguments.flat_positional + ) + device_check = RegisterDispatchKey.gen_device_check( + f.device_check, list(device_check_args), name + ) + + device_guard = "// DeviceGuard omitted" # default + if f.device_guard and self.backend_index.device_guard: + has_tensor_options = any( + isinstance(a, TensorOptionsArguments) + for a in f.func.arguments.non_out + ) + if has_tensor_options: + # kernel is creating a tensor + device_guard = """ + const DeviceGuard device_guard(device_or_default(device));""" + + # CUDA requires special handling + if is_cuda_dispatch_key(self.backend_index.dispatch_key): + device_guard = f"globalContext().lazyInitDevice(c10::DeviceType::CUDA);\n{device_guard}" + else: + # kernel is operating on existing tensors + + # There is precedence for which argument we use to do + # device guard. This describes the precedence order. + self_arg = ( + [f.func.arguments.self_arg.argument] + if f.func.arguments.self_arg is not None + else [] + ) + candidate_args = itertools.chain( + self_arg, + f.func.arguments.out, + f.func.arguments.flat_positional, + ) + + # Only tensor like arguments are eligible + device_of = next( + ( + f"{a.name}" + for a in candidate_args + if a.type.is_tensor_like() + ), + None, + ) + if device_of is not None: + device_guard = f"const OptionalDeviceGuard device_guard(device_of({device_of}));" + + return f"""\ +namespace {{ + +{returns_type} {name}({args_str}) {{ + {device_check} + + {device_guard} + return {impl_name}({args_exprs_str}); +}} + +}} // anonymous namespace +""" + + elif self.target is Target.REGISTRATION: + if f.manual_kernel_registration or self.skip_dispatcher_op_registration: + return None + else: + payload = f"TORCH_FN({name})" + return f'm.impl("{f.func.name}",\n{payload});\n' + else: + assert_never(self.target) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# STRUCTURED +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +@dataclass(frozen=True) +class StructuredRegisterDispatchKey(RegisterDispatchKey): + g: NativeFunctionsGroup + + def gen_class_set_output_functions( + self, k: SchemaKind, parent_class: str, generate_super: bool + ) -> str: + if generate_super: + set_output_super = f"{parent_class}::set_output_raw_strided(output_idx, sizes, strides, options, names);" + else: + set_output_super = "" + + def gen_set_output_function(name: str, maybe_create_proxy: bool) -> str: + return f""" +void set_output_{name}( + int64_t output_idx, IntArrayRef sizes, IntArrayRef strides, + TensorOptions options, DimnameList names +) override {{ +{textwrap.indent(self.gen_class_set_output_body(k, maybe_create_proxy), " ")} + if (!names.empty()) {{ + namedinference::propagate_names(outputs_[output_idx], names); + }} + // super must happen after, so that downstream can use maybe_get_output + // to retrieve the output +{textwrap.indent(set_output_super, " ")} +}} +""" + + return f""" +{gen_set_output_function("strided", maybe_create_proxy=True)} +{gen_set_output_function("raw_strided", maybe_create_proxy=False)} +""" + + def gen_class_set_output_body(self, k: SchemaKind, maybe_create_proxy: bool) -> str: + if self.backend_index.dispatch_key in [ + DispatchKey.CUDA, + DispatchKey.MPS, + DispatchKey.XPU, + DispatchKey.CompositeExplicitAutogradNonFunctional, + ]: + maybe_set_guard = """ +auto current_device = guard_.current_device(); +if (C10_UNLIKELY(current_device.has_value())) { + TORCH_INTERNAL_ASSERT(*current_device == options.device(), + "structured kernels don't support multi-device outputs"); +} else { + guard_.reset_device(options.device()); +} +""" + maybe_set_guard_line = maybe_set_guard + "\n" + else: + maybe_set_guard_line = maybe_set_guard = "" + + if maybe_create_proxy: + create_proxy = """ +auto maybe_proxy = maybe_create_proxy(out, sizes, strides, options); +if (C10_UNLIKELY(maybe_proxy.has_value())) { + proxy_outputs_[output_idx] = std::move(maybe_proxy).value(); +} +""" + else: + create_proxy = "" + + if k is SchemaKind.functional: + if self.backend_index.dispatch_key not in ( + DispatchKey.Meta, + DispatchKey.CPU, + DispatchKey.CUDA, + DispatchKey.MPS, + DispatchKey.XPU, + DispatchKey.MTIA, + DispatchKey.CompositeExplicitAutogradNonFunctional, + ): + raise AssertionError( + f"Unexpected dispatch key {self.backend_index.dispatch_key} " + "for functional schema" + ) + return f"""{maybe_set_guard_line} +outputs_[output_idx] = create_out(sizes, strides, options);""" + elif k is SchemaKind.inplace: + return f"""{maybe_set_guard_line} +const auto& out = outputs_[output_idx].get(); +check_inplace(out, sizes, options); +{create_proxy}""" + elif k is SchemaKind.out: + return f"""{maybe_set_guard_line} +const auto& out = outputs_[output_idx].get(); +resize_out(out, sizes, strides, options); +{create_proxy}""" + elif k is SchemaKind.mutable or k is SchemaKind.scratch: + raise AssertionError( + f"{k} structured operators are currently not supported" + ) + else: + assert_never(k) + + # returns the definition of a ctor, as well as how to construct + # this class to a variable named op + def gen_class_ctor(self, k: SchemaKind, class_name: str, returns: int) -> str: + if k is SchemaKind.functional: + return "" + elif k is SchemaKind.inplace: + # TODO: Make sure out argument is guaranteed to be self + return f"{class_name}(Tensor& self) : outputs_{{std::ref(self)}} {{}}" + elif k is SchemaKind.out: + out_args = ", ".join(f"Tensor& out{i}" for i in range(returns)) + out_refs = ", ".join(f"std::ref(out{i})" for i in range(returns)) + return f"{class_name}({out_args}) : outputs_{{ {out_refs} }} {{}}" + elif k is SchemaKind.mutable or k is SchemaKind.scratch: + raise AssertionError( + f"{k} structured operators are currently not supported" + ) + else: + assert_never(k) + + def gen_class( + self, + f: NativeFunction, + k: SchemaKind, + *, + class_name: str, + parent_class: str, + generate_super: bool, + ) -> str: + if k is SchemaKind.functional: + output_type = "Tensor" + output_value = "outputs_[output_idx]" + proxy_field = "" + elif k is SchemaKind.inplace: + output_type = "std::reference_wrapper" + output_value = "proxy_outputs_[output_idx].has_value() ? *proxy_outputs_[output_idx] : outputs_[output_idx].get()" + proxy_field = f"std::array<::std::optional, {len(f.func.returns)}> proxy_outputs_;" + elif k is SchemaKind.out: + output_type = "std::reference_wrapper" + output_value = "proxy_outputs_[output_idx].has_value() ? *proxy_outputs_[output_idx] : outputs_[output_idx].get()" + proxy_field = f"std::array<::std::optional, {len(f.func.returns)}> proxy_outputs_;" + else: + raise RuntimeError(f"Unsupported SchemaKind {k}") + + if self.backend_index.dispatch_key == DispatchKey.CUDA: + guard_field = "c10::cuda::OptionalCUDAGuard guard_;" + elif ( + self.backend_index.dispatch_key + == DispatchKey.CompositeExplicitAutogradNonFunctional + ): + guard_field = "c10::OptionalDeviceGuard guard_;" + elif self.backend_index.dispatch_key == DispatchKey.MPS: + # TODO: Move to OptionalMPSGuard. + guard_field = "c10::OptionalDeviceGuard guard_;" + elif self.backend_index.dispatch_key == DispatchKey.XPU: + guard_field = "c10::OptionalDeviceGuard guard_;" + elif self.backend_index.dispatch_key == DispatchKey.MTIA: + guard_field = "c10::OptionalDeviceGuard guard_;" + else: + guard_field = "" + + indent = " " * 4 + class_ctor_str = self.gen_class_ctor(k, class_name, len(f.func.returns)) + lines = ( + f"struct {class_name} final : public {parent_class} {{", + f"{textwrap.indent(class_ctor_str, indent)}", + f"{textwrap.indent(self.gen_class_set_output_functions(k, parent_class, generate_super), indent)}", + " const Tensor& maybe_get_output(int64_t output_idx) override {", + f" return {output_value};\n", # type: ignore[possibly-undefined] # TODO: audit + " }", + # type: ignore[possibly-undefined] # TODO: audit + f" std::array<{output_type}, {len(f.func.returns)}> outputs_;", + f"{textwrap.indent(proxy_field, indent)}", # type: ignore[possibly-undefined] # TODO: audit + f"{textwrap.indent(guard_field, indent)}", + "};", + ) + return "\n".join(line for line in lines if line) + + @method_with_native_function + def gen_one(self, f: NativeFunction) -> str | None: + if f.manual_kernel_registration: + raise AssertionError( + f"Function {f.func.name} has manual_kernel_registration=True" + ) + + if ( + self.target is Target.REGISTRATION + and not self.selector.is_native_function_selected(f) + ): + return None + + # TODO: Now, there is something interesting going on here. In the code below, + # we generate CompositeExplicitAutogradNonFunctional implementations of functional and inplace + # based on the out implementation. But in fact, out is definable by + # functional too (just not very efficiently), and this is honestly the + # MORE likely situation for a backend implementer. How do we pick? + # Well, taking a page from Haskell type classes and default methods, + # we could conceivably register a circular definition (out in terms + # of functional, and functional in terms of out) and just require + # someone to implement one or the other. We'd have to do a little bit + # of work to not register one of these "weak" definitions unless there + # is a strong definition somewhere in the DAG! So it's not implemented yet. + if ( + self.backend_index.dispatch_key + == DispatchKey.CompositeExplicitAutogradNonFunctional + and f.func.kind() is SchemaKind.out + ): + # Never generate a default implementation for out, that's what you + # have to define as a backend implementer + return None + + # Note [Direct dispatch bindings] + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Signature of the non-dispatched function we'll expose in a header + # (e.g., at::cpu::add). We don't generate methods (TODO: do this + # when CPUTensor class is a thing); nor do we generate fallback + # bindings for manual_cpp_binding functions. + cpp_sig_group = CppSignatureGroup.from_native_function( + f, method=False, fallback_binding=False + ) + + # Signature of the wrapper function we'll register to the dispatcher + kern = self.backend_index.get_kernel(f) + sig = NativeSignature( + f.func, + prefix=f"wrapper_{self.backend_index.dispatch_key}_", + symint=kern is not None and kern.supports_symint(), + ) + + if self.target is Target.NAMESPACED_DECLARATION: + result = "" + for cpp_sig in cpp_sig_group.signatures(symint=self.symint): + result += f"TORCH_API {cpp_sig.decl()};\n" + return result + + elif self.target is Target.NAMESPACED_DEFINITION: + + def generate_defn(cpp_sig: CppSignature) -> str: + return f""" +{cpp_sig.defn()} {{ +return {sig.name()}({", ".join(e.expr for e in translate(cpp_sig.arguments(), sig.arguments()))}); +}} +""" + + result = "" + for cpp_sig in cpp_sig_group.signatures(symint=self.symint): + result += generate_defn(cpp_sig) + return result + + elif self.target is Target.ANONYMOUS_DEFINITION: + k = f.func.kind() + + # Construct the body of the wrapper function with signature sig + sig_body = [] + # We'll use context to keep track of any variables we've brought + # into scope while generating code + context: list[Binding | Expr] = list(sig.arguments()) + + # Initialize the class corresponding to this structured + # operator; feeding it the output argument(s) if it is known + if self.backend_index.dispatch_key is DispatchKey.Meta: + class_name = f"structured_{meta.name(self.g)}_meta_{k.name}" + parent_class = f"at::meta::structured_{meta.name(self.g)}" + elif ( + self.backend_index.dispatch_key + is DispatchKey.CompositeExplicitAutogradNonFunctional + ): + # TODO: dedup this branch + class_name = f"structured_{meta.name(self.g)}_default_backend_{k.name}" + parent_class = f"at::meta::structured_{meta.name(self.g)}" + else: + metadata = self.backend_index.get_kernel(self.g) + if metadata is None: + raise AssertionError( + f"No kernel metadata found for {self.g.functional.func.name}" + ) + class_name = f"structured_{metadata.kernel}_{k.name}" + parent_class = f"{metadata.cpp_namespace}::structured_{metadata.kernel}" + + if self.backend_index.device_guard: + device_check_args = itertools.chain( + f.func.arguments.out, f.func.arguments.flat_positional + ) + sig_body.append( + RegisterDispatchKey.gen_device_check( + f.device_check, list(device_check_args), sig.name() + ) + ) + + if k is SchemaKind.functional: + sig_body.append(f"{class_name} op;") + elif k is SchemaKind.inplace: + sig_body.append(f"{class_name} op(self);") + elif k is SchemaKind.out: + out_args_str = ", ".join(a.name for a in f.func.arguments.out) + sig_body.append(f"{class_name} op({out_args_str});") + + # Translate the input native arguments into structured + # arguments for the meta call + meta_exprs = ", ".join( + e.expr + for e in translate( + context, structured.meta_arguments(self.g), method=False + ) + ) + + if self.g.out.precomputed: + # If this function group has precomputed elements, the meta function + # returns a struct containing them which must be saved so that it + # can be unpacked when generating code to call the impl. + sig_body.append(f"auto precompute = op.meta({meta_exprs});") + + # Put all of the contents of the precompute struct into the context + # so that translate will be able to return the correct args for the + # call to the impl. + precomputed_values = [ + *self.g.out.precomputed.replace.values(), + self.g.out.precomputed.add, + ] + for precomputed_elems in precomputed_values: + context.extend( + Expr( + expr=f"precompute.{arg.name}", + type=structured.argument_type(arg, binds=arg.name), + ) + for arg in precomputed_elems + ) + + # Add a use of the precompute struct so FB internal compilers don't + # complain that there is an unused variable. + sig_body.append("(void)precompute;") + else: + sig_body.append(f"op.meta({meta_exprs});") + + # After running meta, op.outputs_ is guaranteed to be valid; + # add it to the context + out_args = structured.out_arguments(self.g) + for i, out_arg in enumerate(out_args): + if ConstRefCType(BaseCType(tensorT)) != out_arg.nctype.type: + raise AssertionError( + f"Expected out_arg type to be ConstRefCType(BaseCType(tensorT)), " + f"got {out_arg.nctype.type}" + ) + + if k is SchemaKind.out: + expr = f"op.maybe_get_output({i})" + else: + expr = f"op.outputs_[{i}]" + + context.append( + Expr( + expr=expr, + # TODO: Stop hardcoding that the output type is a Tensor. Note + # that for the codegen here this is fine because outputs_ is + # hardcoded to be tensor already + type=NamedCType( + out_arg.nctype.name, MutRefCType(BaseCType(tensorT)) + ), + ) + ) + + # With the expanded context, do the impl call (if not a meta + # function) + if ( + self.backend_index.dispatch_key + == DispatchKey.CompositeExplicitAutogradNonFunctional + ): + # TODO: https://github.com/pytorch/pytorch/issues/53023 + out_sig_group = CppSignatureGroup.from_native_function( + self.g.out, method=False, fallback_binding=f.manual_cpp_binding + ) + out_sig = out_sig_group.most_faithful_signature() + api_name = out_sig.name() + out_exprs = ", ".join( + e.expr + for e in translate(context, out_sig.arguments(), method=False) + ) + # TODO: I think this means structured won't work with method + # only functions (but maybe you're saved by faithful? iunno.) + # NB: Originally I wrote this as an at::redispatch call, but + # I got in trouble because that meant I needed a DispatchKeySet + # in the wrapper function, which meant I needed a DispatchKeySet + # in the DispatchKeyFunctions declarations, but the defined API + # there does NOT permit a dispatch key set. I think you can + # probably unwind this by calling some function to do the TLS + # fetch and get the DispatchKeySet when you don't have it, but + # I didn't do it for this version + sig_body.append(f"at::{api_name}({out_exprs});") + elif self.backend_index.dispatch_key != DispatchKey.Meta: + impl_exprs = ", ".join( + e.expr + for e in translate( + context, structured.impl_arguments(self.g), method=False + ) + ) + sig_body.append(f"op.impl({impl_exprs});") + + # Go over each output, and check if there is a proxy created for it. + # If so, copy it over to the original output. + if k is SchemaKind.out or k is SchemaKind.inplace: + for i in range(len(f.func.returns)): + sig_body.append( + f"if (op.proxy_outputs_[{i}].has_value()) op.outputs_[{i}].get().copy_(*op.proxy_outputs_[{i}]);" + ) + + # Destructively return the final tensors + # TODO: Do this in translate instead + if k is SchemaKind.functional: + if len(f.func.returns) == 1: + ret_expr = "std::move(op.outputs_[0])" # small optimization + else: + moved = ", ".join( + f"std::move(op.outputs_[{i}])" + for i in range(len(f.func.returns)) + ) + ret_expr = f"std::make_tuple({moved})" + elif k is SchemaKind.inplace: + ret_expr = "self" + elif k is SchemaKind.out: + if len(f.func.returns) == 1: + ret_expr = f.func.arguments.out[0].name + else: + refs = ", ".join(a.name for a in f.func.arguments.out) + ret_expr = f"std::forward_as_tuple({refs})" + sig_body.append(f"return {ret_expr};") # type: ignore[possibly-undefined] # TODO: audit + + sig_body_str = "\n".join(sig_body) + + # For an overview of what this template code looks like, see + # https://github.com/pytorch/rfcs/pull/9 + return f"""\ +{ + self.gen_class( + f, + k, + class_name=class_name, + parent_class=parent_class, + generate_super=self.g.out.structured_inherits is not None, + ) + } + +{sig.defn()} {{ +{sig_body_str} +}} +""" + + elif self.target is Target.REGISTRATION: + return f'm.impl("{f.func.name}", TORCH_FN({sig.name()}));' + else: + assert_never(self.target) + # Silence mypy's "Missing return statement" error + return None diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/ufunc.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/ufunc.py new file mode 100644 index 0000000000000000000000000000000000000000..82ba5352e586d509cd8a26fc6c63ed249e5e35c8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/dest/ufunc.py @@ -0,0 +1,561 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import torchgen.api.ufunc as ufunc +from torchgen.api.translate import translate +from torchgen.api.types import ( + BaseCType, + Binding, + CType, + Expr, + NamedCType, + opmath_t, + scalar_t, + StructuredImplSignature, + VectorizedCType, +) +from torchgen.context import with_native_function +from torchgen.model import ( + Argument, + BaseTy, + BaseType, + DispatchKey, + NativeFunctionsGroup, + ScalarType, + UfuncKey, +) +from torchgen.utils import OrderedSet + + +if TYPE_CHECKING: + from collections.abc import Sequence + + from torchgen.api.ufunc import UfunctorBindings + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# CUDA STUFF +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + +# NB: not bothering to generate dispatch stub forward declaration in header, +# we can just paste it wherever necessary + +# TODO: use BackendIndex +# dispatch_key: DispatchKey # only CPU/CUDA right now + + +# Represents functors for implementing CUDA ufuncs. +# Functors are templated by scalar_t because when USERS instantiate functors +# they are templated. A functor looks something like this: +# +# template +# struct CUDAFunctorOnSelf_add { +# using opmath_t = at::opmath_type; +# opmath_t other_; +# opmath_t alpha_; +# CUDAFunctorOnSelf_add(opmath_t other, opmath_t alpha) +# : other_(other), alpha_(alpha) {} +# __device__ scalar_t operator()(scalar_t self) { +# return ufunc::add(static_cast(self), other_, alpha_); +# } +# }; +# +@dataclass(frozen=True) +class UfunctorSignature: + g: NativeFunctionsGroup + scalar_tensor_idx: int | None + name: str + + def arguments(self) -> UfunctorBindings: + return ufunc.ufunctor_arguments( + self.g, scalar_tensor_idx=self.scalar_tensor_idx, scalar_t=scalar_t + ) + + def fields(self) -> list[Binding]: + # fields are renamed to have a trailing underscore, as is conventional + return [b.rename(f"{b.name}_") for b in self.arguments().ctor] + + def returns_type(self) -> CType: + # TODO: don't hardcode; return type will be inferred based on tags on + # the native function + return BaseCType(scalar_t) + + def decl_fields(self) -> str: + return "\n".join(f"{f.type} {f.name};" for f in self.fields()) + + def inline_defn_ctor(self) -> str: + args_str = ", ".join(a.decl() for a in self.arguments().ctor) + # NB: hypothetically could do this with translate but the + # transition here is very regular + init_str = ", ".join(f"{a.name}_({a.name})" for a in self.arguments().ctor) + return f"{self.name}({args_str}) : {init_str} {{}}" + + def decl_apply(self) -> str: + args_str = ", ".join(a.decl() for a in self.arguments().apply) + return f"{self.returns_type().cpp_type()} operator()({args_str}) const" + + +@dataclass(frozen=True) +class UfuncSignature: + g: NativeFunctionsGroup + name: str + compute_t: CType + + def arguments(self) -> list[Binding]: + return ufunc.ufunc_arguments(self.g, compute_t=self.compute_t) + + def call(self, ctx: Sequence[Binding | Expr]) -> str: + return f"{self.name}({', '.join(a.expr for a in translate(ctx, self.arguments()))})" + + +# steps: +# 1. take the functional signature +# 2. use api.ufunc to convert it to template signature. this establishes +# the type of the template function +# 3. use api.ufunc (II) to generate a split struct / operator() signature. +# this establish context in which we call the template signature +# +# StructuredImplSignature context +# ~> functor constructor sig +# +# Functor constructor context +# ~> functor fields sig +# +# Functor apply context (functor fields + functor apply sig) +# ~> template sig +# + + +def eligible_for_binary_scalar_specialization(g: NativeFunctionsGroup) -> bool: + num_tensors = sum( + 1 for a in g.functional.func.arguments.flat_non_out if a.type.is_tensor_like() + ) + return num_tensors == 2 + + +def compute_ufunc_cuda_functors( + g: NativeFunctionsGroup, +) -> tuple[dict[ScalarType, dict[UfuncKey, UfunctorSignature]], str]: + # First, build the functors. + ufunctor_sigs: dict[ScalarType, dict[UfuncKey, UfunctorSignature]] = {} + ufunctors: list[str] = [] + loops = g.out.ufunc_inner_loop + scalar_tensor_idx_lookup = { + UfuncKey.CUDAFunctorOnSelf: 1, + UfuncKey.CUDAFunctorOnOther: 0, + UfuncKey.CUDAFunctor: None, + } + if eligible_for_binary_scalar_specialization(g): + keys = [ + UfuncKey.CUDAFunctorOnSelf, + UfuncKey.CUDAFunctorOnOther, + UfuncKey.CUDAFunctor, + ] + else: + keys = [UfuncKey.CUDAFunctor] + for k in [UfuncKey.CUDAFunctorOnSelf, UfuncKey.CUDAFunctorOnOther]: + if k in loops: + raise AssertionError(f"cannot use {k} on non-binary function") + for k in keys: + # If the key was directly defined, skip functor codegen; we assume the + # user already done it for us + if k in loops: + ufunctor_sig = UfunctorSignature( + g, scalar_tensor_idx=scalar_tensor_idx_lookup[k], name=loops[k].name + ) + for dtype in loops[k].supported_dtypes: + ufunctor_sigs.setdefault(dtype, {})[k] = ufunctor_sig + continue + + # Note [ScalarOnly and Generic must match names for CUDA] + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # Otherwise, look in ANY of the generic entries. For simplicity of + # codegen, both ScalarOnly and Generic are defined, the ufunc name + # must match (if they didn't match, we'd have to generate distinct + # functors per dtype, which is awful, so we're not going to do it unless + # someone really forces us to) + ufunc_name = None + supported_dtypes: OrderedSet[ScalarType] = OrderedSet() + for lk in [UfuncKey.ScalarOnly, UfuncKey.Generic]: + if lk not in loops: + continue + if ufunc_name is None: + ufunc_name = loops[lk].name + else: + # See Note [ScalarOnly and Generic must match names for CUDA] + if ufunc_name != loops[lk].name: + raise AssertionError( + "ScalarOnly and Generic must have same ufunc name" + ) + supported_dtypes |= loops[lk].supported_dtypes + if ufunc_name is None: + raise AssertionError("ufunc_name must be non-None") + + name = f"{k}_{ufunc_name}" + ufunctor_sig = UfunctorSignature( + g, scalar_tensor_idx=scalar_tensor_idx_lookup[k], name=name + ) + for dtype in supported_dtypes: + ufunctor_sigs.setdefault(dtype, {})[k] = ufunctor_sig + + ufunc_sig = UfuncSignature( + g, name=f"ufunc::{ufunc_name}", compute_t=BaseCType(opmath_t) + ) + apply_ctx = ufunctor_sig.fields() + ufunctor_sig.arguments().apply + ufunctors.append( + f""" +template +struct {ufunctor_sig.name} {{ + using opmath_t = at::opmath_type; + {ufunctor_sig.decl_fields()} + {ufunctor_sig.inline_defn_ctor()} + __device__ {ufunctor_sig.decl_apply()} {{ + return {ufunc_sig.call(apply_ctx)}; + }} +}}; +""" + ) + + return ufunctor_sigs, "\n".join(ufunctors) + + +@dataclass(frozen=True) +class BinaryScalarSpecializationConfig: + scalar_idx: int + ctor_tensor: str + ufunc_key: UfuncKey + + +BinaryScalarSpecializationConfigs = [ + BinaryScalarSpecializationConfig( + scalar_idx=0, + ctor_tensor="self", + ufunc_key=UfuncKey.CUDAFunctorOnOther, + ), + BinaryScalarSpecializationConfig( + scalar_idx=1, + ctor_tensor="other", + ufunc_key=UfuncKey.CUDAFunctorOnSelf, + ), +] + + +def compute_ufunc_cuda_dtype_body( + g: NativeFunctionsGroup, + dtype: ScalarType, + inner_loops: dict[UfuncKey, UfunctorSignature], + parent_ctx: Sequence[Binding], +) -> str: + body = "using opmath_t = at::opmath_type;" + body += "if (false) {}\n" # for ease of codegen + for config in BinaryScalarSpecializationConfigs: + if config.ufunc_key not in inner_loops: + continue + ufunctor_sig = inner_loops[config.ufunc_key] + scalar_idx = config.scalar_idx + 1 + # Make a copy and at the same time widen the type (not permissible + # without copy; we don't want to mutate the input argument anyway) + ctx: list[Expr | Binding] = list(parent_ctx) + ctx.append( + Expr( + expr=f"iter.scalar_value({scalar_idx})", + type=NamedCType(config.ctor_tensor, BaseCType(opmath_t)), + ) + ) + ufunctor_ctor_exprs_str = ", ".join( + a.expr for a in translate(ctx, ufunctor_sig.arguments().ctor) + ) + + # NB: ufunctor must be allocated before iter.remove_operand is called, + # as it relies on iter + body += f"""\ +else if (iter.is_cpu_scalar({scalar_idx})) {{ + {ufunctor_sig.name} ufunctor({ufunctor_ctor_exprs_str}); + iter.remove_operand({scalar_idx}); + gpu_kernel(iter, ufunctor); +}}""" + + ufunctor_sig = inner_loops[UfuncKey.CUDAFunctor] + ufunctor_ctor_exprs_str = ", ".join( + a.expr for a in translate(parent_ctx, ufunctor_sig.arguments().ctor) + ) + body += f""" +else {{ + gpu_kernel(iter, {ufunctor_sig.name}({ufunctor_ctor_exprs_str})); +}} + """ + return body + + +@with_native_function +def compute_ufunc_cuda(g: NativeFunctionsGroup) -> str: + # First, build the functors, indexing them by dtype + ufunctor_sigs, ufunctors = compute_ufunc_cuda_functors(g) + + # Next, build the conditionals + sig = StructuredImplSignature(g, ufunc.kernel_name(g, DispatchKey.CUDA)) + dtype_cases = [] + for dtype, inner_ufunc_sigs in ufunctor_sigs.items(): + dtype_cases.append( + f""" +AT_DISPATCH_CASE(at::ScalarType::{dtype}, + [&]() {{ + {compute_ufunc_cuda_dtype_body(g, dtype, inner_ufunc_sigs, sig.arguments())} + }} +) +""" + ) + + dtype_cases_str = "\n".join(dtype_cases) + + stub_sig = StubSignature(g) + + return f""" +{ufunctors} + +{stub_sig.type_defn()}; +{stub_sig.dispatch_decl()} + +{stub_sig.kernel_defn()} {{ + AT_DISPATCH_SWITCH(iter.common_dtype(), "{sig.name}", + {dtype_cases_str} + ); +}} +REGISTER_DISPATCH({stub_sig.name}, &{stub_sig.kernel_name}) + +{sig.defn()} {{ + {stub_sig.direct_call(sig.arguments())}; +}} +""" + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# CPU STUFF +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +@dataclass(frozen=True) +class StubSignature: + g: NativeFunctionsGroup + + @property + def name(self) -> str: + return f"{str(self.g.functional.func.name.name)}_stub" + + @property + def kernel_name(self) -> str: + return f"{str(self.g.functional.func.name.name)}_kernel" + + @property + def type_name(self) -> str: + return f"{str(self.g.functional.func.name.name)}_fn" + + def arguments(self) -> list[Binding]: + return ufunc.stub_arguments(self.g) + + def type(self) -> str: + cpp_args = self.arguments() + return f"void(*)(TensorIteratorBase&, {', '.join(a.type for a in cpp_args)})" + + def dispatch_decl(self) -> str: + return f"DECLARE_DISPATCH({self.type_name}, {self.name})" + + def dispatch_defn(self) -> str: + return f"DEFINE_DISPATCH({self.name})" + + def kernel_defn(self) -> str: + return f"void {self.kernel_name}(TensorIteratorBase& iter, {', '.join(a.defn() for a in self.arguments())})" + + def type_defn(self) -> str: + return f"using {self.type_name} = {self.type()}" + + # must be called from context where this is TensorIteratorBase* + def call(self, ctx: Sequence[Binding]) -> str: + return f"{self.name}(device_type(), *this, {', '.join(a.expr for a in translate(ctx, self.arguments()))})" + + # used in CUDA to skip the unnecessary dynamic dispatch + def direct_call(self, ctx: Sequence[Binding]) -> str: + return f"{self.kernel_name}(*this, {', '.join(a.expr for a in translate(ctx, self.arguments()))})" + + +@with_native_function +def compute_ufunc_cpu(g: NativeFunctionsGroup) -> str: + stub_sig = StubSignature(g) + sig = StructuredImplSignature(g, ufunc.kernel_name(g, DispatchKey.CPU)) + + return f""" +{stub_sig.type_defn()}; +{stub_sig.dispatch_decl()} +{stub_sig.dispatch_defn()}; + +{sig.defn()} {{ + {stub_sig.call(sig.arguments())}; +}} +""" + + +def compute_ufunc_cpu_dtype_body( + g: NativeFunctionsGroup, + dtype: ScalarType, + inner_loops: dict[UfuncKey, UfuncSignature], + parent_ctx: Sequence[Binding], +) -> str: + if UfuncKey.CPUScalar not in inner_loops: + raise AssertionError(f"{dtype}, {inner_loops.keys()}") + if not inner_loops.keys() <= {UfuncKey.CPUScalar, UfuncKey.CPUVector}: + raise AssertionError( + f"inner_loops keys must be subset of CPUScalar/CPUVector, got {inner_loops.keys()}" + ) + scalar_loop = inner_loops[UfuncKey.CPUScalar] + vec_loop = None + if UfuncKey.CPUVector in inner_loops: + vec_loop = inner_loops[UfuncKey.CPUVector] + + # NB: We DON'T use translate here, because translate is + # incapable of CSE'ing the scalar accesses in case it is also + # used by Vectorized; also, the unpacking here is very simple + # and only affects Scalar; everything else is implicitly captured + # by the lambda + + # Setup scalar in scope + body = [] + ctx = [] + for b in parent_ctx: + if isinstance(b.argument, Argument) and b.argument.type != BaseType( + BaseTy.Scalar + ): + continue + body.append(f"auto _s_{b.name} = {b.name}.to();") + ctx.append(Expr(f"_s_{b.name}", NamedCType(b.nctype.name, BaseCType(scalar_t)))) + if vec_loop is not None: + for b in parent_ctx: + if isinstance(b.argument, Argument) and b.argument.type != BaseType( + BaseTy.Scalar + ): + continue + body.append( + f"auto _v_{b.name} = at::vec::Vectorized(_s_{b.name});" + ) + ctx.append( + Expr( + f"_v_{b.name}", + NamedCType(b.nctype.name, VectorizedCType(BaseCType(scalar_t))), + ) + ) + + # Setup lambda signature + # NB: simplified version of ufunctor_arguments + scalar_bindings = [] + vec_bindings = [] + for a in g.functional.func.arguments.flat_non_out: + if not a.type.is_tensor_like(): + continue + if a.type != BaseType(BaseTy.Tensor): + raise AssertionError(f"Expected Tensor type, got {a.type}") + scalar_bindings.append( + Binding( + name=a.name, + nctype=NamedCType(a.name, BaseCType(scalar_t)), + argument=a, + ) + ) + if vec_loop is not None: + vec_bindings.append( + Binding( + name=a.name, + nctype=NamedCType(a.name, VectorizedCType(BaseCType(scalar_t))), + argument=a, + ) + ) + + def with_ctx(b: Sequence[Binding]) -> list[Expr | Binding]: + r: list[Expr | Binding] = [] + r.extend(ctx) + r.extend(b) + return r + + body_str = "\n".join(body) + if vec_loop is not None: + return f""" +{body_str} +cpu_kernel_vec(iter, + [=]({", ".join(b.decl() for b in scalar_bindings)}) {{ return {scalar_loop.call(with_ctx(scalar_bindings))}; }}, + [=]({", ".join(b.decl() for b in vec_bindings)}) {{ return {vec_loop.call(with_ctx(vec_bindings))}; }} +); +""" + else: + return f""" +{body_str} +cpu_kernel(iter, + [=]({", ".join(b.decl() for b in scalar_bindings)}) {{ return {scalar_loop.call(with_ctx(scalar_bindings))}; }} +); +""" + + +@with_native_function +def compute_ufunc_cpu_kernel(g: NativeFunctionsGroup) -> str: + stub_sig = StubSignature(g) + + # Reindex the ufunc by dtypes; processing generic/scalaronly as well + loops = g.out.ufunc_inner_loop + ufunc_sigs: dict[ScalarType, dict[UfuncKey, UfuncSignature]] = {} + for k in [UfuncKey.CPUScalar, UfuncKey.CPUVector]: + lks = [] + # ORDER MATTERS: this specifies overriding precedence + if k in loops: # should happen rarely + lks.append(k) + if UfuncKey.ScalarOnly in loops and k is UfuncKey.CPUScalar: + lks.append(UfuncKey.ScalarOnly) + if UfuncKey.Generic in loops: + lks.append(UfuncKey.Generic) + # TODO: don't hardcode ufunc:: namespace here, should be centralized smh + for lk in lks: + for dtype in loops[lk].supported_dtypes: + compute_t: CType + if k is UfuncKey.CPUScalar: + compute_t = BaseCType(scalar_t) + elif k is UfuncKey.CPUVector: + compute_t = VectorizedCType(BaseCType(scalar_t)) + else: + raise AssertionError + inner_ufunc_sigs = ufunc_sigs.setdefault(dtype, {}) + if k not in inner_ufunc_sigs: + inner_ufunc_sigs[k] = UfuncSignature( + g, name=f"ufunc::{loops[lk].name}", compute_t=compute_t + ) + + # Build the conditionals + dtype_cases = [] + for dtype, inner_ufunc_sigs in ufunc_sigs.items(): + dtype_cases.append( + f""" +AT_DISPATCH_CASE(at::ScalarType::{dtype}, + [&]() {{ + {compute_ufunc_cpu_dtype_body(g, dtype, inner_ufunc_sigs, stub_sig.arguments())} + }} +) +""" + ) + + dtype_cases_str = "\n".join(dtype_cases) + return f""" +namespace {{ + +{stub_sig.kernel_defn()} {{ + AT_DISPATCH_SWITCH(iter.common_dtype(), "{stub_sig.name}", + {dtype_cases_str} + ); +}} + +}} // anonymous namespace + +{stub_sig.type_defn()}; +{stub_sig.dispatch_decl()} +REGISTER_DISPATCH({stub_sig.name}, &{stub_sig.kernel_name}) +""" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen.py new file mode 100644 index 0000000000000000000000000000000000000000..10f727b9dfba0c14f02b2cce430ff3f3ba9a260b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen.py @@ -0,0 +1,3088 @@ +from __future__ import annotations + +import argparse +import functools +import json +import keyword +import os +from collections import defaultdict, namedtuple, OrderedDict +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Literal, TYPE_CHECKING, TypeVar +from typing_extensions import assert_never + +import yaml + +import torchgen.api.dispatcher as dispatcher +import torchgen.api.meta as meta +import torchgen.api.native as native +import torchgen.api.structured as structured +import torchgen.dest as dest +from torchgen.api import cpp +from torchgen.api.translate import translate +from torchgen.api.types import ( + Binding, + CppSignature, + CppSignatureGroup, + DispatcherSignature, + NamedCType, + NativeSignature, + SpecialArgName, +) +from torchgen.context import ( + method_with_native_function, + native_function_manager, + with_native_function, + with_native_function_and_indices, +) +from torchgen.gen_aoti_c_shim import ( + gen_aoti_c_shim_files, + gen_static_dispatch_backend_call_signature, +) +from torchgen.gen_functionalization_type import ( + gen_functionalization_definition, + gen_functionalization_registration, + gen_functionalization_view_inverse_declaration, + gen_functionalization_view_meta_classes_decl, + gen_functionalization_view_meta_classes_impl, + GenCompositeViewCopyKernel, +) +from torchgen.gen_vmap_plumbing import gen_all_vmap_plumbing +from torchgen.model import ( + Argument, + BackendIndex, + BackendMetadata, + BaseOperatorName, + DEFAULT_KERNEL_NAMESPACE, + dispatch_device_map, + DispatchKey, + FRAGMENT_NAMESPACES, + FunctionSchema, + is_cuda_dispatch_key, + is_generic_dispatch_key, + is_ufunc_dispatch_key, + is_xpu_dispatch_key, + Location, + NativeFunction, + NativeFunctionsGroup, + NativeFunctionsViewGroup, + OperatorName, + OptionalType, + SchemaKind, + SelfArgument, + STRUCTURED_DISPATCH_KEYS, + TensorOptionsArguments, + Type, + Variant, + ViewSchemaKind, +) +from torchgen.native_function_generation import ( + add_generated_native_functions, + gen_composite_functional_kernel, + gen_composite_out_kernel, + pre_group_native_functions, +) +from torchgen.selective_build.selector import SelectiveBuilder +from torchgen.utils import ( + concatMap, + context, + FileManager, + make_file_manager, + mapMaybe, + NamespaceHelper, + Target, +) +from torchgen.yaml_utils import YamlDumper, YamlLoader + + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + +T = TypeVar("T") + +# Welcome to the ATen code generator v2! The ATen code generator is +# responsible for parsing native_functions.yaml and then generating +# various generated files (e.g., TypeDefault.cpp) based on the operators +# defined in this file. This means that the code generator knows how to +# parse function schema, and then translate this into various C++ types +# and boilerplate code. +# +# Some things to know about this file when you modify it: +# +# - This file has STRICT mypy typechecking. Typecheck it with +# `mypy --config mypy-strict.ini` in the root source directory +# +# - Most of the heavy lifting lives in external modules: +# - 'model' has the data model for native_functions.yaml. The classes +# in those file represent what you see when you look at +# a native_functions.yaml +# - 'api' has conversions for how to translate JIT schema into +# the various C++ APIs that the codegen interacts with. There +# are in fact THREE different C++ APIs: the public C++ API, +# the dispatcher API, and the legacy dispatcher API. See each +# of these respective files for more information + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# HELPER FUNCTIONS +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +# A custom loader for YAML to let us also keep track of line numbers +# of each entry in the YAML file +class LineLoader(YamlLoader): + def construct_mapping(self, node, deep=False): # type: ignore[no-untyped-def] + mapping = super().construct_mapping(node, deep=deep) # type: ignore[no-untyped-call] + # Add 1 so line numbering starts at 1 + mapping["__line__"] = node.start_mark.line + 1 + return mapping + + +# Parse native_functions.yaml into a sequence of NativeFunctions and Backend Indices. +ParsedYaml = namedtuple("ParsedYaml", ["native_functions", "backend_indices"]) + + +_GLOBAL_PARSE_NATIVE_YAML_CACHE: dict[str, ParsedYaml] = {} +_GLOBAL_PARSE_TAGS_YAML_CACHE: dict[str, set[str]] = {} + + +def file_manager_from_dispatch_key( + dispatch_key: DispatchKey, + device_fms: dict[str, FileManager], + default_fm: FileManager, +) -> FileManager: + fm = device_fms.get( + next( + ( + device + for check, device in dispatch_device_map.items() + if check(dispatch_key) + ), + "", + ), + default_fm, + ) + return fm + + +def parse_native_yaml_struct( + es: object, + valid_tags: set[str], + ignore_keys: set[DispatchKey] | None = None, + path: str = "", + skip_native_fns_gen: bool = False, +) -> ParsedYaml: + if not isinstance(es, list): + raise AssertionError(f"Expected 'es' to be a list, but got {type(es)}") + rs: list[NativeFunction] = [] + bs: dict[DispatchKey, dict[OperatorName, BackendMetadata]] = defaultdict(dict) + for e in es: + if not isinstance(e, dict): + raise AssertionError(f"Expected to be dict: {e}") + if not isinstance(e.get("__line__"), int): + raise AssertionError(f"Expected '__line__' to be int: {e}") + loc = Location(path, e["__line__"]) + funcs = e.get("func") + if funcs is None: + raise AssertionError(f"Missed 'func' in {e}") + with context(lambda: f"in {loc}:\n {funcs}"): + func, m = NativeFunction.from_yaml(e, loc, valid_tags, ignore_keys) + rs.append(func) + BackendIndex.grow_index(bs, m) + error_check_native_functions(rs) + # Default dict is to prevent the codegen from barfing when we have a dispatch key that has no kernels yet. + indices: dict[DispatchKey, BackendIndex] = defaultdict( + lambda: BackendIndex( + dispatch_key=DispatchKey.Undefined, + use_out_as_primary=True, + external=False, + device_guard=False, + # I'm actually not sure about this; undefined could be hit on + # empty TensorList, hypothetically that could have sizes in it + index={}, + ) + ) + if not skip_native_fns_gen: + add_generated_native_functions(rs, bs) + for k, v in bs.items(): + # All structured in-tree operators are implemented in terms of their out operator. + indices[k] = BackendIndex( + dispatch_key=k, + use_out_as_primary=True, + external=False, + # Only cuda-like devices in tree require device guards + device_guard=is_cuda_dispatch_key(k) or is_xpu_dispatch_key(k), + index=v, + ) + return ParsedYaml(rs, indices) + + +def parse_tags_yaml_struct(es: object, path: str = "") -> set[str]: + if not isinstance(es, list): + raise AssertionError(f"Expected 'es' to be a list, but got {type(es)}") + rs: set[str] = set() + for e in es: + if not isinstance(e.get("__line__"), int): + raise AssertionError(f"Expected '__line__' to be int: {e}") + loc = Location(path, e["__line__"]) + tags = e.get("tag") + with context(lambda: f"in {loc}:\n {tags}"): + e_i = e.copy() + name = e_i.pop("tag") + desc = e_i.pop("desc", "") + # ensure that each tag has a non-empty description + if desc == "": + raise AssertionError(f"Tag '{name}' must have a non-empty description") + rs.add(name) + return rs + + +@functools.cache +def parse_tags_yaml(path: str) -> set[str]: + global _GLOBAL_PARSE_TAGS_YAML_CACHE + if path not in _GLOBAL_PARSE_TAGS_YAML_CACHE: + with open(path) as f: + es = yaml.load(f, Loader=LineLoader) + _GLOBAL_PARSE_TAGS_YAML_CACHE[path] = parse_tags_yaml_struct(es, path=path) + + return _GLOBAL_PARSE_TAGS_YAML_CACHE[path] + + +def parse_native_yaml( + path: str, + tags_yaml_path: str, + ignore_keys: set[DispatchKey] | None = None, + *, + skip_native_fns_gen: bool = False, + loaded_yaml: object | None = None, +) -> ParsedYaml: + global _GLOBAL_PARSE_NATIVE_YAML_CACHE + if path not in _GLOBAL_PARSE_NATIVE_YAML_CACHE: + valid_tags = parse_tags_yaml(tags_yaml_path) + + # if a loaded yaml is provided, use that instead of reading from path + if loaded_yaml is None: + with open(path) as f: + es = yaml.load(f, Loader=LineLoader) + else: + es = loaded_yaml + + _GLOBAL_PARSE_NATIVE_YAML_CACHE[path] = parse_native_yaml_struct( + es, + valid_tags, + ignore_keys, + path=path, + skip_native_fns_gen=skip_native_fns_gen, + ) + + return _GLOBAL_PARSE_NATIVE_YAML_CACHE[path] + + +# Some assertions are already performed during parsing, but those are only within a single NativeFunction. +# Assertions here are meant to be performed across NativeFunctions. +def error_check_native_functions(funcs: Sequence[NativeFunction]) -> None: + func_map: dict[OperatorName, NativeFunction] = {} + base_func_map: dict[BaseOperatorName, list[NativeFunction]] = defaultdict(list) + for f in funcs: + func_map[f.func.name] = f + base_func_map[f.func.name.name].append(f) + for f in funcs: + if f.structured_delegate is not None: + delegate_func = func_map.get(f.structured_delegate) + if delegate_func is None: + raise AssertionError( + f"{f.func.name} is marked as a structured_delegate pointing to " + f"{f.structured_delegate}, but {f.structured_delegate} is missing." + ) + if not delegate_func.structured: + raise AssertionError( + f"{f.func.name} is marked as a structured_delegate pointing to " + f"{f.structured_delegate}, but {f.structured_delegate} is not marked as structured. " + f"Consider adding 'structured=True' to the delegated operator" + ) + + # Check for reserved Python keywords + PYTHON_RESERVED_KEYWORDS = set(keyword.kwlist) + # List of pre-existing operators that are known to have reserved keywords + # Exclusion list is used to suppress the assertion for these operators + EXCLUSION_LIST = { + ("_has_compatible_shallow_copy_type", "from"), + ("random_.from", "from"), + ("uniform_", "from"), + } + + for arg in f.func.arguments.flat_all: + if arg.name in PYTHON_RESERVED_KEYWORDS: + if (str(f.func.name), arg.name) not in EXCLUSION_LIST: + raise AssertionError( + f"Argument name '{arg.name}' in function '{f.func.name}' is a reserved Python keyword." + ) + # See Note [resize_ in Functionalization] + # resize_() is technically an inplace view op (and therefore needs the tag), + # but it would be overkill to add a true "view" variant of resize. + # Instead, resize_() gets special treatment in functionalization, + # and we have a resize() op that is non-aliasing + functional. + if ( + "inplace_view" in f.tags + and str(f.func.name) != "resize_" + and str(f.func.name) != "resize_as_" + and str(f.func.name.name) != "set_" + ): + base_name = f.func.name.name + if not base_name.inplace: + raise AssertionError( + f"{f.func.name} is marked with tag: inplace_view, but it doesn't follow the naming " + "convention for inplace ops - the codegen expects the base name to have a trailing underscore." + ) + out_of_place_base_name = BaseOperatorName( + base_name.base, False, base_name.dunder_method + ) + if len(base_func_map[out_of_place_base_name]) == 0: + raise AssertionError( + f"{f.func.name} is marked with tag: inplace_view. The codegen expects there to be a corresponding " + f"out-of-place view op with the name '{base_name}' and matching schema, but it didn't find one." + ) + + +def cpp_string(s: str) -> str: + """Convert a python string into a c++ string literal""" + s = s.replace("\\", "\\\\") + s = s.replace('"', '\\"') + s = s.replace("\a", "\\a") + s = s.replace("\b", "\\b") + s = s.replace("\f", "\\f") + s = s.replace("\n", "\\n") + s = s.replace("\v", "\\v") + s = s.replace("\t", "\\t") + return f'"{s}"' + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# C++ CODE GENERATION +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + +# Most functions in this section are curried: they consist of a function +# that takes some parameters (e.g., what is to be generated) which itself +# returns a function that actually maps NativeFunction to the code +# to be generated. This pattern makes it convenient to use map, concatMap +# and similar functional combinators. + + +def static_dispatch_keys(backends: list[BackendIndex]) -> list[DispatchKey]: + if len(backends) == 0: + return [] + else: + return [backend.dispatch_key for backend in backends] + [ + DispatchKey.CompositeImplicitAutograd, + DispatchKey.CompositeImplicitAutogradNestedTensor, + DispatchKey.CompositeExplicitAutograd, + DispatchKey.CompositeExplicitAutogradNonFunctional, + ] + + +def get_static_dispatch_backend( + f: NativeFunction, backend_index: BackendIndex +) -> DispatchKey | None: + if f.structured_delegate is not None or backend_index.has_kernel(f): + # TODO: for ops with structured_delegate it should check the dispatch table of + # the out variant instead. For now, these structured ops all have CPU/CUDA kernels + # so we always dispatch to the `backend`, but this could be wrong when we + # migrate math/default_backend ops to use structured delegate. + return backend_index.dispatch_key + elif f.has_composite_explicit_autograd_kernel: + return DispatchKey.CompositeExplicitAutograd + elif f.has_composite_explicit_autograd_non_functional_kernel: + return DispatchKey.CompositeExplicitAutogradNonFunctional + elif f.has_composite_implicit_autograd_kernel: + return DispatchKey.CompositeImplicitAutograd + elif f.has_composite_implicit_autograd_nested_tensor_kernel: + return DispatchKey.CompositeImplicitAutogradNestedTensor + return None + + +def static_dispatch_ops_header( + f: NativeFunction, backend_index: list[BackendIndex] +) -> str | None: + if backend_index is None or f.manual_kernel_registration: + return None + + output = [] + for index in backend_index: + dispatch_key = get_static_dispatch_backend(f, index) + if dispatch_key is not None: + output.append( + f"#include " + ) + return "\n".join(output) + + +def static_dispatch_extra_headers(backends: list[BackendIndex]) -> list[str]: + return [ + f"#include " + for dispatch_key in static_dispatch_keys(backends) + ] + + +# Translates arguments of `sig` to CppSignature bindings. +# Note that we have a special case for `memory_format` argument and this case is not covered by +# tools.codegen.api.translate() yet as its application is limited to static dispatch. +def translate_args( + sig: CppSignature | DispatcherSignature, + cpp_sig: CppSignature, +) -> str: + # Adds SpecialArgName.possibly_redundant_memory_format NamedCType for memory_format bindings + def add_spl_memory_format_binding(input_bindings: list[Binding]) -> list[Binding]: + output_bindings: list[Binding] = [] + for binding in input_bindings: + if binding.name == "memory_format": + spl_mem_format_binding = Binding( + nctype=NamedCType( + SpecialArgName.possibly_redundant_memory_format, + binding.nctype.type, + ), + name=binding.name, + default=binding.default, + argument=binding.argument, + ) + output_bindings.append(spl_mem_format_binding) + else: + output_bindings.append(binding) + return output_bindings + + src_bindings = list(sig.arguments()) + goal_bindings = list(cpp_sig.arguments()) + # When last argument of CPP signature has SpecialArgName.possibly_redundant_memory_format NCType, + # get memory_format bindings of dispatcher signature to have the same NCType as well + for arg in goal_bindings: + if arg.nctype.name == SpecialArgName.possibly_redundant_memory_format: + src_bindings = add_spl_memory_format_binding(src_bindings) + break + exprs = translate(src_bindings, goal_bindings) + return ", ".join(a.expr for a in exprs) + + +def generate_static_dispatch_backend_call( + sig: CppSignature | DispatcherSignature, + f: NativeFunction, + backend_index: BackendIndex, +) -> str: + cpp_sig = gen_static_dispatch_backend_call_signature(sig, f) + name = cpp_sig.name() + exprs = translate_args(sig, cpp_sig) + backend_metadata = backend_index.get_kernel(f) + kernel_ns = ( + backend_metadata.cpp_namespace + if backend_metadata and backend_metadata.cpp_namespace + else DEFAULT_KERNEL_NAMESPACE + ) + ns = kernel_ns.replace("::native", "") + return f"return {ns}::{backend_index.dispatch_key.lower()}::{name}({exprs});" + + +def generate_static_dispatch_fallback_call( + sig: CppSignature | DispatcherSignature, + f: NativeFunction, + backend_indices: list[BackendIndex], +) -> str: + cpp_sigs = CppSignatureGroup.from_native_function( + f, method=False, fallback_binding=False + ) + if sig.symint and f.func.has_symint(): + cpp_sig = cpp_sigs.symint_signature + else: + cpp_sig = cpp_sigs.signature + if cpp_sig is None: + raise AssertionError("Expected cpp_sig to be non-None") + name = cpp_sig.name() + exprs = translate_args(sig, cpp_sig) + ns = DEFAULT_KERNEL_NAMESPACE.replace("::native", "") + if f.has_composite_explicit_autograd_kernel: + return f"return {ns}::{DispatchKey.CompositeExplicitAutograd.lower()}::{name}({exprs});" + elif f.has_composite_explicit_autograd_non_functional_kernel: + return f"return {ns}::{DispatchKey.CompositeExplicitAutogradNonFunctional.lower()}::{name}({exprs});" + elif f.has_composite_implicit_autograd_kernel: + return f"return {ns}::{DispatchKey.CompositeImplicitAutograd.lower()}::{name}({exprs});" + elif f.has_composite_implicit_autograd_nested_tensor_kernel: + return f"return {ns}::{DispatchKey.CompositeImplicitAutogradNestedTensor.lower()}::{name}({exprs});" + else: + return f"""TORCH_CHECK(false, "Static dispatch does not support {name} for\ +{", ".join([str(index.dispatch_key) for index in backend_indices])} ");""" + + +def static_dispatch( + sig: CppSignature | DispatcherSignature, + f: NativeFunction, + backend_indices: list[BackendIndex], +) -> str: + """ + For a given `NativeFunction`, find out the corresponding backend and dispatch to it. If more than one + backends exist, fallback to static dispatch by determining dispatch key from inputs. + Arguments: + sig: A CppSignature or DispatcherSignature for this native function we want to use. + f: NativeFunction to generate static dispatch. + backend_indices: All available backends. + Return: + C++ code to call backend-specific functions, e.g., "return at::cpu::add(self, other, scale);" + """ + if len(backend_indices) == 0 or f.manual_kernel_registration: + return "" + + keys = [ + b + for b in backend_indices + if b.has_kernel(f) + or ( + f.structured_delegate is not None + and b.dispatch_key in STRUCTURED_DISPATCH_KEYS + ) + ] + if len(keys) == 1: + return generate_static_dispatch_backend_call(sig, f, keys[0]) + elif len(keys) == 0: + return generate_static_dispatch_fallback_call(sig, f, backend_indices) + + native_tensor_args = [ + a.name + for a in sig.arguments() + if isinstance(a.argument, SelfArgument) + or isinstance(a.argument, Argument) + and a.argument.type.is_tensor_like() + ] + tensor_args = ", ".join(native_tensor_args) + tensor_opts = f.func.arguments.tensor_options + + stmts = [] + subexprs: list[str] = [] + if tensor_opts is not None: + subexprs.append( + "DispatchKeySet(c10::computeDispatchKey(dtype, layout, device))" + ) + if tensor_args != "": + subexprs.append(f"c10::detail::multi_dispatch_key_set({tensor_args})") + stmts.append(f"""DispatchKeySet _dk_set = {" | ".join(subexprs)};""") + stmts.append("DispatchKey _dk = c10::highestPriorityBackendTypeId(_dk_set);") + + dispatch_code = [] + for index in keys: + dispatch_code.append(f"""case DispatchKey::{index.dispatch_key}:""") + dispatch_code.append( + f"""\t{generate_static_dispatch_backend_call(sig, f, index)};""" + ) + + fallback = generate_static_dispatch_fallback_call(sig, f, backend_indices) + connector = "\n\t\t" + + return f""" + {connector.join(stmts)} + switch (_dk) {{ + {connector.join(dispatch_code)} + default: + {fallback} + }} + """ + + +# Generates RegisterSchema.cpp. Depending on the selector, either +# all schemas are registered, or only some are (in the case of +# selective build) +@dataclass(frozen=True) +class RegisterSchema: + selector: SelectiveBuilder + known_tags: dict[str, int] = field(default_factory=dict) + + @method_with_native_function + def __call__(self, f: NativeFunction) -> str | None: + if not self.selector.is_native_function_selected(f): + return None + tags = "{" + ", ".join(f"at::Tag::{tag}" for tag in sorted(f.tags)) + "}" + if tags == "{}": + return f"m.def({cpp_string(str(f.func))}, {{}});\n" + maybe_tags = "" + if tags not in self.known_tags: + idx = len(self.known_tags) + self.known_tags[tags] = idx + maybe_tags = f"const std::vector tags_{idx} = {tags};\n" + return f"{maybe_tags}m.def({cpp_string(str(f.func))}, tags_{self.known_tags[tags]});\n" + + +# Generates Operators.h and Operators.cpp. +# These provide macros that, given an operator and overload name, allow users +# to access an "un-overloaded" function version of the operator. This +# is useful for extension writers who want to (1) want to decltype the operator +# and (2) don't want to worry about method-only operators. +@dataclass(frozen=True) +class ComputeOperators: + target: Literal[Target.DECLARATION, Target.DEFINITION] + static_dispatch_backend_indices: list[BackendIndex] + + @method_with_native_function + def __call__(self, f: NativeFunction) -> str: + sig = DispatcherSignature.from_schema(f.func) + name = f.func.name.unambiguous_name() + + if self.target is Target.DECLARATION: + # Note [The ATen Operators API] + # The ATen Operators API lives in the at::_ops namespace, and contains compile-time + # metadata about each operator + entry points into the Dispatcher. + # The C++ function, method, and redispatch API's are all implemented as wrappers + # into various bits of the structs defined here. + # + # Important characteristics about the Operators API: + # (1) It follows the Dispatcher API. + # This is kind of necessary to avoid overhead. + # For example: if it followed the C++ API, then all of the faithful C++ factory functions + # would need to wrap their arguments into TensorOptions only to unwrap them again. + # (2) Overload names are disambiguated. + # This is helpful for pytorch extenders who would like to decltype() an aten operator, + # that has overloads, e.g. decltype(at::_ops::mul_Tensor::call) + # (3) No argument defaulting is allowed. + # This is more of an implementation detail to avoid #include cycles, + # since TensorBody.h (which defines the Tensor class) needs to include this file. + # (4) manual_cpp_bindings and faithful names are not included in the API. + # This applies to stuff like __dispatch__is_complex(), and add_outf(). + # These aren't "real aten ops", they're just additional functions provided by the C++ API. + # They're implemented as wrappers in Functions.h that call into the actual operators + # defined here, i.e. at::_ops::is_complex::call() and at::_ops::add_out::call(). + # This means that ATEN_OP(is_complex) will not fastpath, and will go through the dispatcher. + return f""" +struct TORCH_API {name} {{ + using schema = {sig.type()}; + using ptr_schema = schema*; + // See Note [static constexpr char* members for windows NVCC] + static constexpr const char* name = "aten::{f.func.name.name}"; + static constexpr const char* overload_name = "{f.func.name.overload_name}"; + static constexpr const char* schema_str = {cpp_string(str(f.func))}; + static {sig.defn(name="call", is_redispatching_fn=False)}; + static {sig.defn(name="redispatch", is_redispatching_fn=True)}; +}};""" + + elif self.target is Target.DEFINITION: + defns = f""" +// aten::{f.func} +static C10_NOINLINE c10::TypedOperatorHandle<{name}::schema> create_{name}_typed_handle() {{ + return c10::Dispatcher::singleton() + .findSchemaOrThrow({name}::name, {name}::overload_name) + .typed<{name}::schema>(); +}} +""" + for is_redispatching_fn in [False, True]: + if is_redispatching_fn: + dispatcher_exprs_str = ", ".join( + ["dispatchKeySet"] + [a.name for a in sig.arguments()] + ) + method_base = "redispatch" + else: + dispatcher_exprs_str = ", ".join([a.name for a in sig.arguments()]) + method_base = "call" + + dispatcher_call = method_base + method_name = f"{name}::{method_base}" + + fn_body = f""" + static auto op = create_{name}_typed_handle(); + return op.{dispatcher_call}({dispatcher_exprs_str});""" + + if ( + not is_redispatching_fn + and len(self.static_dispatch_backend_indices) > 0 + ): + # call() should go through static dispatch + fn_body = static_dispatch( + sig, f, backend_indices=self.static_dispatch_backend_indices + ) + defns += f""" +// aten::{f.func} +{sig.defn(name=method_name, is_redispatching_fn=is_redispatching_fn)} {{ + {fn_body} +}} +""" + return defns + else: + assert_never(self.target) + + +# Generates Functions.h, which provides the functional public C++ API, +# and the scaffolding to call into the dispatcher from these functions. +@dataclass(frozen=True) +class ComputeFunction: + @method_with_native_function + def __call__(self, f: NativeFunction) -> str | None: + sig_group = CppSignatureGroup.from_native_function( + f, method=False, fallback_binding=f.manual_cpp_binding + ) + has_symint = f.func.has_symint() + + result = "" + for sig in sig_group.signatures(): + # See Note [The ATen Operators API] + target_sig = DispatcherSignature.from_schema(f.func) + exprs = translate(sig.arguments(), target_sig.arguments()) + exprs_str = ", ".join([e.expr for e in exprs]) + + if sig.symint: + intlike_t = "c10::SymInt" + else: + intlike_t = "int64_t" + + if Variant.function in f.variants: + result += f""" +// aten::{f.func} +inline {sig.decl()} {{ + return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str}); +}}""" + + # The template function can be used from template situations + # where you want to switch between the symint or not version + # depending on a template argument + # + # NB: we ALWAYS generate this even for methods. But we put it in + # this header so it can take advantage of per-op headers + if has_symint: + result += f""" +namespace symint {{ + template >> + {sig.decl(suppress_symint_suffix=True)} {{ + return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str}); + }} +}} +""" + return result + + +# Generates TensorBody.h. This file provides the object-oriented (method-based) +# public C++ API, and the scaffolding to call into the dispatcher from these functions. +@dataclass(frozen=True) +class ComputeTensorMethod: + target: Literal[Target.DECLARATION, Target.DEFINITION] + static_dispatch_backend_indices: list[BackendIndex] + + @method_with_native_function + def __call__(self, f: NativeFunction) -> str | None: + if Variant.method not in f.variants: + return None + + if f.func.is_out_fn(): + raise AssertionError(f"Method variant cannot be an out function: {f.func}") + if f.func.arguments.self_arg is None: + raise AssertionError(f"Method variant must have self_arg: {f.func}") + + sig_group = CppSignatureGroup.from_native_function( + f, method=True, fallback_binding=f.manual_cpp_binding + ) + + if self.target is Target.DECLARATION: + result = "" + for sig in sig_group.signatures(): + result += f"{sig.decl()} const;\n" + return result + + if self.target is not Target.DEFINITION: + assert_never(self.target) + + result = "" + + for sig in sig_group.signatures(): + target_sig = DispatcherSignature.from_schema(f.func) + exprs = translate(sig.arguments(), target_sig.arguments(), method=True) + exprs_str = ", ".join([e.expr for e in exprs]) + + result += f""" +// aten::{f.func} +inline {sig.defn(prefix="Tensor::")} const {{ + return at::_ops::{f.func.name.unambiguous_name()}::call({exprs_str}); +}} +""" + + return result + + +# Generates RedispatchFunctions.h. +# This is similar to the C++ API defined in Functions.h, but provides access +# to the dispatcher's redispatch API. +@dataclass(frozen=True) +class ComputeRedispatchFunction: + @method_with_native_function + def __call__(self, f: NativeFunction) -> str | None: + # We unconditionally generate function variants of the redispatch API. + # This is mainly because we can namespace functions separately, but not methods, + sig_group = CppSignatureGroup.from_native_function( + f, method=False, fallback_binding=f.manual_cpp_binding + ) + + result = "" + for sig in sig_group.signatures(): + target_sig = DispatcherSignature.from_schema(f.func) + exprs = translate(sig.arguments(), target_sig.arguments()) + exprs_str = ", ".join(["dispatchKeySet"] + [a.expr for a in exprs]) + + result += f""" +// aten::{f.func} +inline {sig.decl(is_redispatching_fn=True)} {{ + return at::_ops::{f.func.name.unambiguous_name()}::redispatch({exprs_str}); +}} +""" + + return result + + +# Generates ATenOpList.cpp, a runtime accessible list of all aten +# operators. +# TODO: This was historically used to help some JIT interop code +# figure out whether or not to treat aten namespace'd operators +# one way or another, we should reevaluate if this is actually needed. +@with_native_function +def compute_aten_op(f: NativeFunction) -> str: + return f'{{"aten::{f.func.name.name}", "{f.func.name.overload_name}"}},' + + +# Generates MetaFunctions.h +def compute_meta_function_declaration(g: NativeFunctionsGroup) -> str | None: + if not g.structured: + return None + with native_function_manager(g.out): + name = meta.name(g) + args = structured.meta_arguments(g) + args_str = ", ".join(a.decl() for a in args) + parent_class = g.out.structured_inherits + if parent_class is None: + parent_class = "at::impl::MetaBase" + meta_return = "void" + precomputed = g.out.precomputed if g.structured else None + + if precomputed: + # Generate the template declaration with one bool parameter for each + # precomputed element. Each parameter is true if the corresponding (in + # terms of position) precomputed element has been set. + precomputed_values = [*precomputed.replace.values(), precomputed.add] + precomputed_elements = [ + elem for replace_list in precomputed_values for elem in replace_list + ] + precomputed_template_parameters = [ + elem.name.upper() for elem in precomputed_elements + ] + precomputed_template_params_str = ", ".join( + f"bool {param} = false" for param in precomputed_template_parameters + ) + precompute_template_decl = f"template <{precomputed_template_params_str}>" + + # Generate a string containing declarations of all precomputed elements. + precomputed_elements_with_cpp_types = [ + structured.argument_type(elem, binds=elem.name) + for elem in precomputed_elements + ] + + precomputed_elements_decl = ";\n".join( + f"{elem.cpp_type(strip_ref=True)} {elem.name}" + for elem in precomputed_elements_with_cpp_types + ) + + # Generate "setter" methods for each precomputed element. Each method will return + # a new instance of precompute_out with the template parameter that corresponds to + # the member set by the method to true (to indicate that it has been set). + setter_methods = [] + for i, elem in enumerate(precomputed_elements): + # Generate the signature. The return type will be the same + # as the type of `this` but with the template parameter + # corresponding to the element set by this method set to true. + # The assert generated below will ensure that this template + # parameter is false on the type of `this`. + return_ty_templates = ", ".join( + precomputed_template_parameters[:i] + + ["true"] + + precomputed_template_parameters[i + 1 :] + ) + return_ty = f"precompute_out<{return_ty_templates}>" + elem_cpp_ty = precomputed_elements_with_cpp_types[i].cpp_type( + strip_ref=True + ) + signature = f"{return_ty} set_{elem.name}({elem_cpp_ty} value)" + + # Generate an assert which checks that the + # template parameter corresponding to the precomputed + # element that is set by this method is false on the + # class corresponding to the object that `this` points to. + # This ensures that each element can be set only once. + assert_msg = f'"{elem.name} already set"' + assert_stmt = f"static_assert({precomputed_template_parameters[i]} == false, {assert_msg});" + + # Generate the new object construction block. All state + # except the element that this method sets is copied from the + # object that `this` points to. The value for the element that + # the method sets is taken from a method parameter. + construction_stmts = [] + construction_stmts.append(f"{return_ty} ret;") + + for j, elem in enumerate(precomputed_elements): + if i == j: + construction_stmts.append(f"ret.{elem.name} = value;") + else: + construction_stmts.append( + f"ret.{elem.name} = this->{elem.name};" + ) + + construction_stmts.append("return ret;") + construction_block = "\n".join(construction_stmts) + + setter_methods.append( + f""" + {signature} {{ + {assert_stmt} + {construction_block} + }} + """ + ) + setter_methods_decl = "\n".join(setter_methods) + + # Meta should return an instance of the struct containing the precomputed elements. + meta_return_template_params = ", ".join( + ["true"] * len(precomputed_template_parameters) + ) + # This typedef (actually a using statement) is needed so that TORCH_META_FUNC can reuse the return + # type (which has a variable number of template parameters). + meta_return_typedef = f"using meta_return_ty = precompute_out <{meta_return_template_params}>;" + meta_return = "meta_return_ty" + precomputed_decl = f""" + {precompute_template_decl} + struct TORCH_API precompute_out {{ + {setter_methods_decl} + {precomputed_elements_decl}; + }};""" + else: + meta_return_typedef = "" + precomputed_decl = "" + + return f"""\ +struct TORCH_API structured_{name} : public {parent_class} {{ + {precomputed_decl} + {meta_return_typedef} + {meta_return} meta({args_str}); +}}; +""" + + +def needs_backend_select(f: NativeFunction, selector: SelectiveBuilder) -> bool: + name = str(f.func.name.name) + if name.endswith("_like") or name.startswith("new_"): + return False + if f.func.arguments.tensor_options is None: + return False + return selector.is_native_function_selected(f) + + +# Generates RegisterBackendSelect.cpp, a series of kernels which provide +# specialized computation of dispatch key for operator signatures which cannot +# be easily done automatically using templating. +@dataclass(frozen=True) +class ComputeBackendSelect: + target: Literal[Target.DEFINITION, Target.REGISTRATION] + + # Selector object to determine which operators to generate + # registration code for. + selector: SelectiveBuilder + + @method_with_native_function + def __call__(self, f: NativeFunction) -> str | None: + if not needs_backend_select(f, self.selector): + return None + + name = native.name(f.func) + # BackendSelect can go to Meta, so it must preserve symints + native_sig = NativeSignature(f.func, symint=True) + + native_tensor_args = [ + a + for a in native_sig.arguments() + if isinstance(a.argument, Argument) and a.argument.type.is_tensor_like() + ] + + dispatcher_sig = DispatcherSignature.from_schema(f.func) + + sig: NativeSignature | DispatcherSignature + sig = dispatcher_sig + dispatcher_exprs = dispatcher_sig.exprs() + dispatch_key = "c10::computeDispatchKey(dtype, layout, device)" + + if self.target is Target.DEFINITION: + # I don't think there's actually a good reason to generate + # these two cases differently + # The first case could probably be improved though- it calls computeDispatchKeySet(), + # which looks at TLS dispatch keys- there should not be any by the time we reach backend select. + if native_tensor_args: + if not f.func.arguments.has_tensor_arg(): + raise AssertionError( + f"Expected function to have tensor args: {f.func}" + ) + tensor_args = ", ".join(a.name for a in native_tensor_args) + compute_dk = f"""\ +DispatchKeySet _dk_set = c10::DispatchKeySet({dispatch_key}) | c10::detail::multi_dispatch_key_set({tensor_args}); +DispatchKeySet _dk_mask = c10::DispatchKeySet(DispatchKeySet::FULL_AFTER, DispatchKey::BackendSelect); +DispatchKeySet _dk = c10::impl::computeDispatchKeySet(_dk_set, _dk_mask);""" + else: + if f.func.arguments.has_tensor_arg(): + raise AssertionError( + f"Expected function to not have tensor args: {f.func}" + ) + compute_dk = ( + f"DispatchKeySet _dk = c10::DispatchKeySet({dispatch_key});" + ) + return f"""\ +// aten::{f.func} +C10_ALWAYS_INLINE +{sig.defn(name)} {{ + {compute_dk} + return at::_ops::{f.func.name.unambiguous_name()}::redispatch( + _dk, {", ".join(a.expr for a in dispatcher_exprs)}); +}} +""" + elif self.target is Target.REGISTRATION: + return f"""m.impl("aten::{f.func.name}", TORCH_FN({name}));""" + else: + assert_never(self.target) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# YAML CODE GENERATION +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +def format_yaml(data: object) -> str: + # Ignore alias in Dumper + YamlDumper.ignore_aliases = lambda self, data: True # type: ignore[assignment] + + # Support serializing OrderedDict + def dict_representer(dumper: Any, data: Any) -> Any: + return dumper.represent_dict(data.items()) + + YamlDumper.add_representer(OrderedDict, dict_representer) # type: ignore[no-untyped-call] + # Some yaml parsers (e.g. Haskell's) don't understand line breaks. + # width=1e9 turns off optional line breaks and improves + # the portability of the outputted yaml. + return yaml.dump(data, default_flow_style=False, Dumper=YamlDumper, width=1e9) # type: ignore[no-any-return, call-overload] + + +# For some reason, some defaults we write to YAML are written as native +# YAML objects, rather than doing them uniformly as strings. This +# function detects those cases and converts them into native Python +# objects. +def pythonify_default(s: str) -> object: + if s == "true": + return True + elif s == "false": + return False + + try: + return int(s) + except ValueError: + try: + return float(s) + except ValueError: + return s + + +# What is a dynamic type? Over time, the semantic meaning of +# dynamic type has degraded to meaninglessness (in the old days, +# it captured dtype-ness of types, but that has gone away with +# the removal of TH). These days, it's mostly the same thing as +# the C++ API argument type, except that Tensor and Tensor? +# arguments simply present as Tensor. +# +# TODO: Get rid of dynamic_type, after getting tools/autograd +# to use the new codegen framework +def dynamic_type(t: Type) -> str: + if isinstance(t, OptionalType): + return dynamic_type(t.elem) + # Note we don't use t.is_tensor_like() here because it would + # also include Tensor[] + if str(t) == "Tensor": + return "at::Tensor" + # This is a legacy concept, so never report SymInt + return cpp.argumenttype_type( + t, mutable=False, binds="__placeholder__", symint=False + ).cpp_type() + + +def compute_method_of_yaml(variants: set[Variant]) -> list[str]: + # This is written out explicitly to ensure that Tensor and + # namespace are put into the list in the right order + method_of = ["Type"] + if Variant.method in variants: + method_of.append("Tensor") + if Variant.function in variants: + method_of.append("namespace") + return method_of + + +def compute_returns_yaml( + f: NativeFunction, +) -> tuple[list[dict[str, str]], dict[str, str]]: + # Note [name and field_name] + # ~~~~~~~~~~~~~~~~~~~~~~~~~~ + # To understand name_to_field_name, we must first talk about this + # schema: + # + # lstsq.X(Tensor self, Tensor A, *, Tensor(a!) X, Tensor(b!) qr) -> (Tensor(a!) solution, Tensor(b!) QR) + # + # There is something very odd about this schema: it is an out + # variant of the function (that is to say, it will convert into + # at::lstsq_out() in the C++ API), but the names of the output + # return arguments don't match the keyword argument names of + # the inputs. It TURNS OUT that in this situation, the historical + # Declarations.yaml we want to output is this (abbreviated to + # only show relevant fields): + # + # arguments: + # ... + # - field_name: solution + # name: X + # - field_name: QR + # name: qr + # ... + # + # returns: + # - field_name: solution + # name: X + # - field_name: QR + # name: qr + # + # The name of the return fields is stored in 'field_name', and the + # name of the arguments is stored in 'name'. So when we process + # arguments, we need a way to get at the corresponding return. At + # the moment, this is most conveniently done by constructing a + # mapping from name (the argument concept) to field_name (the + # return concept) while processing return arguments, since we don't + # directly maintain this correspondence in the modeling of function + # schema itself. + # + # See also https://github.com/pytorch/pytorch/issues/43114 + name_to_field_name: dict[str, str] = {} + + # Compute the returns field of the YAML entry + names = cpp.return_names(f) + returns = [] + for i, (r, name) in enumerate(zip(f.func.returns, names)): + ret = { + "dynamic_type": dynamic_type(r.type), + "name": name, + # legacy, report ints + "type": cpp.return_type(r, symint=False).cpp_type(), + } + + if r.name: + # See Note [name and field_name] + ret["field_name"] = r.name + if f.func.is_out_fn(): + name_to_field_name[f.func.arguments.out[i].name] = r.name + + returns.append(ret) + + return returns, name_to_field_name + + +# arguments in yaml roughly corresponds to the public C++ API +def compute_cpp_argument_yaml( + cpp_a: Binding, + *, + schema_order: bool, + kwarg_only_set: set[str], + out_arg_set: set[str], + name_to_field_name: dict[str, str], +) -> object: + if isinstance(cpp_a.argument, TensorOptionsArguments): + arg: dict[str, object] = { + "annotation": None, + "dynamic_type": "at::TensorOptions", + "is_nullable": False, + "name": cpp_a.name, + "type": cpp_a.type, + "kwarg_only": True, + } + if cpp_a.default is not None: + arg["default"] = cpp_a.default + return arg + elif isinstance(cpp_a.argument, SelfArgument): + raise AssertionError + elif isinstance(cpp_a.argument, Argument): + return compute_argument_yaml( + cpp_a.argument, + schema_order=schema_order, + kwarg_only_set=kwarg_only_set, + out_arg_set=out_arg_set, + name_to_field_name=name_to_field_name, + ) + + +def compute_argument_yaml( + a: Argument, + *, + schema_order: bool, + kwarg_only_set: set[str], + out_arg_set: set[str], + name_to_field_name: dict[str, str], +) -> object: + arg: dict[str, object] = { + "annotation": str(a.annotation) if a.annotation else None, + "dynamic_type": dynamic_type(a.type), + "is_nullable": a.type.is_nullable(), + "name": a.name, + # legacy, report ints + "type": cpp.argument_type(a, binds="__placeholder__", symint=False).cpp_type(), + } + if a.default is not None: + arg["default"] = pythonify_default( + cpp.default_expr(a.default, a.type, symint=False) + ) + if a.name in kwarg_only_set: + arg["kwarg_only"] = True + if a.name in out_arg_set: + arg["output"] = True + arg["allocate"] = True + # See Note [name and field_name] + if a.name in name_to_field_name: + arg["field_name"] = name_to_field_name[a.name] + # Historically, booleans don't get their size recorded, because it + # is already built into the cpp type (e.g., std::array) + l = a.type.is_list_like() + if l is not None and l.size is not None and str(l.elem) != "bool": + arg["size"] = l.size + return arg + + +@with_native_function +def compute_declaration_yaml(f: NativeFunction) -> object: + returns, name_to_field_name = compute_returns_yaml(f) + + # These sets are used to conveniently test if an argument is a + # kwarg-only or out argument + kwarg_only_set = {a.name for a in f.func.arguments.flat_kwarg_only} + out_arg_set = {a.name for a in f.func.arguments.out} + + sig_group = CppSignatureGroup.from_native_function( + f, method=False, fallback_binding=False + ) + cpp_args = sig_group.signature.arguments() + arguments = [ + compute_cpp_argument_yaml( + cpp_a, + schema_order=False, + kwarg_only_set=kwarg_only_set, + out_arg_set=out_arg_set, + name_to_field_name=name_to_field_name, + ) + for cpp_a in cpp_args + ] + + schema_order_jit_arguments = list(f.func.schema_order_arguments()) + + schema_order_arguments = [ + compute_argument_yaml( + a, + schema_order=True, + kwarg_only_set=kwarg_only_set, + out_arg_set=out_arg_set, + name_to_field_name=name_to_field_name, + ) + for a in schema_order_jit_arguments + ] + + cpp_schema_order_types = [ + # NB: method here doesn't matter + r.type + for a in schema_order_jit_arguments + for r in cpp.argument( + a, + method=False, + cpp_no_default_args=set(), + faithful=False, + symint=False, + has_tensor_options=False, + ) + ] + + # legacy, report ints + cpp_returns = cpp.returns_type(f.func.returns, symint=False).cpp_type() + schema_order_cpp_signature = f"{cpp_returns} ({', '.join(cpp_schema_order_types)})" + + is_factory_method = ( + any(isinstance(a.argument, TensorOptionsArguments) for a in cpp_args) + and Variant.method not in f.variants + ) + + return OrderedDict( + [ + ("name", cpp.name(f.func)), + ("operator_name", str(f.func.name.name)), + ("overload_name", str(f.func.name.overload_name)), + ("manual_kernel_registration", f.manual_kernel_registration), + ( + "category_override", + f.category_override if f.category_override is not None else "", + ), + ("schema_string", f"aten::{f.func}"), + ("arguments", arguments), + ("schema_order_cpp_signature", schema_order_cpp_signature), + ("schema_order_arguments", schema_order_arguments), + ("method_of", compute_method_of_yaml(f.variants)), + ("mode", "native"), + ("python_module", "" if f.python_module is None else f.python_module), + ("returns", returns), + ("inplace", f.func.name.name.inplace), + ("is_factory_method", is_factory_method), + ("abstract", f.is_abstract), + ("device_guard", f.device_guard), + ("with_gil", False), + ("deprecated", False), + ("has_math_kernel", f.has_composite_implicit_autograd_kernel), + ] + ) + + +# See Note [Auto generated composite kernels] +def has_autogenerated_composite_kernel(f: NativeFunction) -> bool: + return (f.structured or f.structured_delegate is not None) and ( + f.func.kind() == SchemaKind.functional or f.func.kind() == SchemaKind.inplace + ) + + +@with_native_function_and_indices +def compute_registration_declarations( + f: NativeFunction, backend_indices: dict[DispatchKey, BackendIndex] +) -> str: + name = dispatcher.name(f.func) + returns_type = dispatcher.returns_type(f.func.returns).cpp_type() + args = dispatcher.arguments(f.func) + args_str = ", ".join(a.no_default().decl() for a in args) + comment_data: dict[str, str] = { + "schema": f"aten::{f.func}", + # TODO: What exactly is the semantics of the 'dispatch' field? + "dispatch": str( + {k for k, v in backend_indices.items() if v.has_kernel(f)} + != {DispatchKey.CompositeImplicitAutograd} + and {k for k, v in backend_indices.items() if v.has_kernel(f)} + != { + DispatchKey.CompositeImplicitAutograd, + DispatchKey.CompositeImplicitAutogradNestedTensor, + } + ), + "default": str(f.has_composite_kernel or has_autogenerated_composite_kernel(f)), + } + return f"""{returns_type} {name}({args_str}); // {json.dumps(comment_data)} +""" + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# RUN IT ALL +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +def get_custom_build_selector( + provided_op_registration_allowlist: list[str] | None, + op_selection_yaml_path: str | None, +) -> SelectiveBuilder: + if ( + provided_op_registration_allowlist is not None + and op_selection_yaml_path is not None + ): + raise AssertionError( + "Both provided_op_registration_allowlist and op_selection_yaml_path " + "can NOT be provided at the same time." + ) + + op_registration_allowlist: set[str] | None = None + if provided_op_registration_allowlist is not None: + op_registration_allowlist = set(provided_op_registration_allowlist) + + if op_registration_allowlist is not None: + selector = SelectiveBuilder.from_legacy_op_registration_allow_list( + op_registration_allowlist, + True, + False, + ) + elif op_selection_yaml_path is not None: + selector = SelectiveBuilder.from_yaml_path(op_selection_yaml_path) + else: + selector = SelectiveBuilder.get_nop_selector() + + return selector + + +def get_grouped_by_view_native_functions( + native_functions: Sequence[NativeFunction], +) -> Sequence[NativeFunction | NativeFunctionsViewGroup]: + def maybe_create_view_group( + d: dict[ViewSchemaKind | SchemaKind, NativeFunction], + ) -> list[NativeFunction | NativeFunctionsViewGroup]: + funcs: list[NativeFunction | NativeFunctionsViewGroup] = [] + if ViewSchemaKind.aliasing in d: + view = d.pop(ViewSchemaKind.aliasing) + view_inplace = d.pop(ViewSchemaKind.aliasing_inplace, None) + view_copy = d.pop(SchemaKind.functional, None) + + funcs.append( + NativeFunctionsViewGroup( + view=view, + view_copy=view_copy, + view_inplace=view_inplace, + ) + ) + # Take the remaining functions that weren't part of the view group + # and emit them separately + funcs.extend(d.values()) + return funcs + + grouped_by_views: dict[ + FunctionSchema, dict[SchemaKind | ViewSchemaKind, NativeFunction] + ] = defaultdict(dict) + for f in native_functions: + schema = f.func.view_signature() + view_kind: ViewSchemaKind = f.view_schema_kind + # We need to group up ops relevant to the same "view", consisting of: + # view op (ViewSchemaKind.aliasing) + # view_inplace op (ViewSchemaKind.aliasing_inplace) + # view_copy op (SchemaKind.functional) + if view_kind == ViewSchemaKind.non_aliasing: + kind = f.func.kind() + if kind in grouped_by_views[schema]: + raise AssertionError( + f"Duplicate schema kind {kind} in {grouped_by_views[schema].keys()}" + ) + grouped_by_views[schema][kind] = f + else: + if view_kind in grouped_by_views[schema]: + raise AssertionError( + f"{view_kind} already in {grouped_by_views[schema].keys()}" + ) + grouped_by_views[schema][view_kind] = f + + return list(concatMap(maybe_create_view_group, grouped_by_views.values())) + + +def get_grouped_native_functions( + native_functions: Sequence[NativeFunction], +) -> Sequence[NativeFunction | NativeFunctionsGroup]: + def flatten_pre_group( + d: dict[SchemaKind, NativeFunction], + ) -> Sequence[NativeFunction | NativeFunctionsGroup]: + r = NativeFunctionsGroup.from_dict(d) + if r is None: + # Invariant: any NativeFunctions that are code-generated + # should have been grouped into NativeFunctionsGroup objects + if any("generated" in f.tags for f in d.values()): + raise AssertionError( + "Generated NativeFunctions should have been grouped into " + f"NativeFunctionsGroup objects: {list(d.values())}" + ) + return list(d.values()) + else: + return [r] + + # TODO: how come ValuesView isn't a Sequence lol + pre_grouped_native_functions = pre_group_native_functions(native_functions) + return list( + concatMap(flatten_pre_group, list(pre_grouped_native_functions.values())) + ) + + +def get_ns_grouped_kernels( + *, + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + backend_indices: dict[DispatchKey, BackendIndex], + native_function_decl_gen: Callable[ + [NativeFunctionsGroup | NativeFunction, BackendIndex], list[str] + ] = dest.compute_native_function_declaration, +) -> dict[str, list[str]]: + ns_grouped_kernels: dict[str, list[str]] = defaultdict(list) + for f in grouped_native_functions: + native_function_namespaces = set() + dispatch_keys = set() + for dispatch_key, backend_idx in backend_indices.items(): + backend_metadata = backend_idx.get_kernel(f) + if backend_metadata: + namespace = backend_metadata.cpp_namespace + dispatch_keys.add(dispatch_key) + native_function_namespaces.add(namespace) + else: + namespace = DEFAULT_KERNEL_NAMESPACE + if len(native_function_namespaces) > 1: + raise AssertionError( + f"Codegen only supports one namespace per operator, " + f"got {native_function_namespaces} from {dispatch_keys}" + ) + ns_grouped_kernels[namespace].extend( + native_function_decl_gen(f, backend_idx) + ) + return ns_grouped_kernels + + +def get_native_function_declarations_from_ns_grouped_kernels( + *, + ns_grouped_kernels: dict[str, list[str]], +) -> list[str]: + declarations: list[str] = [] + newline = "\n" + for namespace, kernels in ns_grouped_kernels.items(): + ns_helper = NamespaceHelper( + namespace_str=namespace, + entity_name="", + max_level=4, + ) + # Convert to a set first to remove duplicate kernel names. Backends are + # allowed to repeat kernel names; only generate the declaration once! + ordered_kernels = list(OrderedDict.fromkeys(kernels)) + declarations.extend( + f""" +{ns_helper.prologue} +{newline.join(ordered_kernels)} +{ns_helper.epilogue} + """.split(newline) + ) + return declarations + + +# Return native function declarations grouped by their namespaces. +def get_native_function_declarations( + *, + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + backend_indices: dict[DispatchKey, BackendIndex], + native_function_decl_gen: Callable[ + [NativeFunctionsGroup | NativeFunction, BackendIndex], list[str] + ] = dest.compute_native_function_declaration, +) -> list[str]: + """ + Generate kernel declarations, in `NativeFunction(s).h`. + :param grouped_native_functions: a sequence of `NativeFunction` or `NativeFunctionGroup`. + :param backend_indices: kernel collections grouped by dispatch key. + :param native_function_decl_gen: callable to generate kernel declaration for each `NativeFunction`. + :return: a list of string, from the string with all declarations, grouped by namespaces, split by newline. + """ + + ns_grouped_kernels = get_ns_grouped_kernels( + grouped_native_functions=grouped_native_functions, + backend_indices=backend_indices, + native_function_decl_gen=native_function_decl_gen, + ) + return get_native_function_declarations_from_ns_grouped_kernels( + ns_grouped_kernels=ns_grouped_kernels + ) + + +def get_kernel_namespace( + *, f: NativeFunction | NativeFunctionsGroup, backend_idx: BackendIndex +) -> str: + backend_metadata = backend_idx.get_kernel(f) + if backend_metadata and "::native" not in backend_metadata.cpp_namespace: + func_name = ( + f.func.name if isinstance(f, NativeFunction) else f.functional.func.name + ) + raise AssertionError( + f"The kernel for function {func_name} " + f"with dispatch key {backend_idx.dispatch_key} " + f"has a namespace {backend_metadata.cpp_namespace} and it's not ending with '::native'." + ) + return ( + backend_metadata.cpp_namespace if backend_metadata else DEFAULT_KERNEL_NAMESPACE + ) + + +# Return native function definitions grouped by dispatch key and custom namespace. +# Used in RegisterDispatchKey.cpp and etc. +def get_native_function_definitions( + *, + fm: FileManager, + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + dispatch_key: DispatchKey, + backend_idx: BackendIndex, + selector: SelectiveBuilder, + rocm: bool, + symint: bool, + skip_dispatcher_op_registration: bool, + gen_dispatch_helpers: bool, +) -> list[str]: + definitions: list[str] = [] + ns_definitions: dict[str, list[str]] = defaultdict(list) + anonymous_definitions: dict[str, list[str]] = defaultdict(list) + registrations: dict[str, dict[str, list[str]]] = defaultdict(dict) + newline = "\n" + ns_gen = dest.RegisterDispatchKey( + backend_idx, + Target.NAMESPACED_DEFINITION, + selector, + rocm=rocm, + symint=symint, + class_method_name=None, + skip_dispatcher_op_registration=skip_dispatcher_op_registration, + ) + anonymous_gen = dest.RegisterDispatchKey( + backend_idx, + Target.ANONYMOUS_DEFINITION, + selector, + rocm=rocm, + symint=symint, + class_method_name=None, + skip_dispatcher_op_registration=skip_dispatcher_op_registration, + ) + reg_gen = dest.RegisterDispatchKey( + backend_idx, + Target.REGISTRATION, + selector, + rocm=rocm, + symint=symint, + class_method_name=None, + skip_dispatcher_op_registration=skip_dispatcher_op_registration, + ) + for f in grouped_native_functions: + kernel_namespace = get_kernel_namespace(f=f, backend_idx=backend_idx).replace( + "::native", "" + ) + + ns_definitions[kernel_namespace].extend( + ns_gen(f), + ) + anonymous_definitions[kernel_namespace].extend( + anonymous_gen(f), + ) + namespace = ( + f.namespace if isinstance(f, NativeFunction) else f.functional.namespace + ) + if namespace not in registrations[kernel_namespace]: + registrations[kernel_namespace] = defaultdict(list) + registrations[kernel_namespace][namespace].extend( + reg_gen(f), + ) + + for kernel_namespace in ns_definitions: + if len(ns_definitions[kernel_namespace]) == 0: + continue + ns_helper = NamespaceHelper(namespace_str=kernel_namespace) + registration_body = "" + for namespace in registrations[kernel_namespace]: + if not registrations[kernel_namespace][namespace]: + continue + registration_body += f""" +TORCH_LIBRARY_IMPL({namespace}, {dispatch_key}, m) {{ + {newline.join(registrations[kernel_namespace][namespace])} +}}""" + definitions.extend( + fm.substitute_with_template( + "RegisterDispatchDefinitions.ini", + lambda: { + "ns_prologue": ns_helper.prologue, + "ns_epilogue": ns_helper.epilogue, + "dispatch_anonymous_definitions": anonymous_definitions[ + kernel_namespace + ], + "static_init_dispatch_registrations": "" + if skip_dispatcher_op_registration + else registration_body, + "deferred_dispatch_registrations": "", + "dispatch_namespace": dispatch_key.lower(), + "dispatch_namespaced_definitions": ns_definitions[kernel_namespace], + }, + ).split(newline) + ) + + return definitions + + +# Return native function declarations grouped by dispatch key and custom namespace. +# Used in CPUFunctions_inl.h and etc. +def get_namespaced_declaration( + *, + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + dispatch_key: DispatchKey, + backend_idx: BackendIndex, + selector: SelectiveBuilder, + rocm: bool, + symint: bool, +) -> list[str]: + declarations: list[str] = [] + ns_grouped_kernels: dict[str, list[str]] = defaultdict(list) + newline = "\n" + func = dest.RegisterDispatchKey( + backend_idx, + Target.NAMESPACED_DECLARATION, + selector, + rocm=rocm, + class_method_name=None, + skip_dispatcher_op_registration=False, + symint=symint, + ) + for f in grouped_native_functions: + namespace = get_kernel_namespace(f=f, backend_idx=backend_idx).replace( + "native", dispatch_key.lower() + ) + + ns_grouped_kernels[namespace].extend( + func(f), + ) + + for namespace, kernels in ns_grouped_kernels.items(): + if len(kernels) == 0: + continue + ns_helper = NamespaceHelper( + namespace_str=namespace, entity_name="", max_level=3 + ) + ordered_kernels = list(OrderedDict.fromkeys(kernels)) + declarations.extend( + f""" +{ns_helper.prologue} +{newline.join(ordered_kernels)} +{ns_helper.epilogue} + """.split(newline) + ) + return declarations + + +# Return native function schema registration code for aten and other namespaces. +def get_native_function_schema_registrations( + *, + native_functions: Sequence[NativeFunction], + schema_selector: SelectiveBuilder, +) -> tuple[list[str], str]: + ns_native_functions: dict[str, list[NativeFunction]] = defaultdict(list) + for native_function in native_functions: + ns_native_functions[native_function.namespace].append(native_function) + schema_registrations = "" + aten_schema_registrations = [] + custom_namespace = None + for namespace, funcs in ns_native_functions.items(): + schema_registrations_body = list( + mapMaybe(RegisterSchema(schema_selector), funcs) + ) + # NB: we have to separate aten namespace registration from other namespaces, + # because in the template we hardcoded an operator for ATen already. + if namespace == "aten": + aten_schema_registrations = schema_registrations_body + else: + custom_namespace = namespace + tab = "\t" + # if the namespace is predefined, we should use define a library fragment + # instead of a new library + torch_library_macro = ( + "TORCH_LIBRARY_FRAGMENT" + if namespace in FRAGMENT_NAMESPACES + else "TORCH_LIBRARY" + ) + schema_registrations += f""" +{torch_library_macro}({custom_namespace}, m) {{ + {tab.join(schema_registrations_body)} +}};""" + return (aten_schema_registrations, schema_registrations) + + +def gen_aggregated_headers( + *, + native_functions: Sequence[NativeFunction], + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + structured_native_functions: Sequence[NativeFunctionsGroup], + static_dispatch_idx: list[BackendIndex], + selector: SelectiveBuilder, + backend_indices: dict[DispatchKey, BackendIndex], + cpu_fm: FileManager, + device_fms: dict[str, FileManager], + functions_keys: set[DispatchKey], + dispatch_keys: Sequence[DispatchKey], + rocm: bool, +) -> None: + # Buck doesn't support dynamic output files, so we aggregate all operator + # headers into a single file + cpu_fm.write( + "NativeMetaFunctions.h", + lambda: { + "NativeMetaFunctions_includes": [], + "NativeMetaFunctions_declarations": list( + mapMaybe(compute_meta_function_declaration, structured_native_functions) + ), + }, + ) + method_native_functions = [ + fn for fn in native_functions if Variant.method in fn.variants + ] + non_method_native_functions = [ + fn for fn in native_functions if fn not in method_native_functions + ] + cpu_fm.write( + "MethodOperators.h", + lambda: { + "MethodOperators_includes": [], + "MethodOperators_declarations": list( + mapMaybe( + ComputeOperators( + Target.DECLARATION, + static_dispatch_backend_indices=static_dispatch_idx, + ), + method_native_functions, + ) + ), + }, + ) + cpu_fm.write( + "Operators.h", + lambda: { + "Operators_includes": ["#include "], + "Operators_declarations": list( + mapMaybe( + ComputeOperators( + Target.DECLARATION, + static_dispatch_backend_indices=static_dispatch_idx, + ), + non_method_native_functions, + ) + ), + }, + ) + cpu_fm.write( + "Functions.h", + lambda: { + "static_dispatch_extra_headers": static_dispatch_extra_headers( + static_dispatch_idx + ), + "Functions_includes": ["#include "], + "Functions_declarations": list( + mapMaybe( + ComputeFunction(), + native_functions, + ) + ), + }, + ) + declarations = get_native_function_declarations( + grouped_native_functions=grouped_native_functions, + backend_indices=backend_indices, + ) + cpu_fm.write( + "NativeFunctions.h", + lambda: { + "NativeFunctions_includes": ["#include "], + "NativeFunctions_declarations": declarations, + }, + ) + + for dispatch_key in dispatch_keys: + fm = file_manager_from_dispatch_key(dispatch_key, device_fms, cpu_fm) + if dispatch_key in functions_keys: + inl_headers = f"#include " + + fm.write_with_template( + f"{dispatch_key}Functions.h", + "DispatchKeyFunctions.h", + lambda: { + "dispatch_key": str(dispatch_key), + "inline_headers": inl_headers, + }, + ) + fm.write_with_template( + f"{dispatch_key}Functions_inl.h", + "DispatchKeyFunctions_inl.h", + lambda: { + "DispatchKeyFunctions_inl_includes": [], + "dispatch_namespace": dispatch_key.lower(), + "dispatch_namespaced_declarations": get_namespaced_declaration( + grouped_native_functions=grouped_native_functions, + dispatch_key=dispatch_key, + backend_idx=backend_indices[dispatch_key], + selector=selector, + rocm=rocm, + symint=True, + ), + }, + ) + + del fm + + +def gen_per_operator_headers( + *, + native_functions: Sequence[NativeFunction], + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + static_dispatch_idx: list[BackendIndex], + selector: SelectiveBuilder, + backend_indices: dict[DispatchKey, BackendIndex], + cpu_fm: FileManager, + device_fms: dict[str, FileManager], + ops_fm: FileManager, + functions_keys: set[DispatchKey], + dispatch_keys: Sequence[DispatchKey], + rocm: bool, +) -> None: + # For CMake builds, split operator declarations into separate headers in + # the ATen/ops folder to split up header dependencies + functions_by_root_name: dict[str, list[NativeFunction]] = defaultdict(list) + for fn in native_functions: + functions_by_root_name[fn.root_name].append(fn) + + grouped_functions_by_root_name: dict[ + str, list[NativeFunction | NativeFunctionsGroup] + ] = defaultdict(list) + for group in grouped_native_functions: + name = group.root_name + grouped_functions_by_root_name[name].append(group) + + for name, functions in functions_by_root_name.items(): + ops_fm.write_with_template( + f"{name}_ops.h", + "Operator.h", + lambda: { + "declarations": list( + mapMaybe( + ComputeOperators( + Target.DECLARATION, + static_dispatch_backend_indices=static_dispatch_idx, + ), + functions, + ) + ), + }, + ) + + ops_fm.write_with_template( + f"{name}.h", + "Function.h", + lambda: { + "static_dispatch_ops_headers": list( + mapMaybe( + lambda fn: static_dispatch_ops_header( + fn, backend_index=static_dispatch_idx + ), + functions, + ) + ), + "operator_includes": f"#include ", + "function_definitions": list( + mapMaybe( + ComputeFunction(), + functions, + ) + ), + }, + ) + + grouped_functions = grouped_functions_by_root_name.get(name, []) + structured_functions = [ + fn + for fn in grouped_functions + if isinstance(fn, NativeFunctionsGroup) and fn.structured + ] + is_structured = len(structured_functions) > 0 + + if is_structured: + ops_fm.write_with_template( + f"{name}_meta.h", + "NativeMetaFunction.h", + lambda: { + "meta_function_declarations": list( + mapMaybe( + compute_meta_function_declaration, structured_functions + ) + ), + }, + ) + declarations = get_native_function_declarations( + grouped_native_functions=grouped_functions, + backend_indices=backend_indices, + native_function_decl_gen=dest.compute_native_function_declaration, + ) + ops_fm.write_with_template( + f"{name}_native.h", + "NativeFunction.h", + lambda: { + "extra_includes": ( + f"#include " if is_structured else [] + ), + "native_function_declarations": declarations, + }, + ) + + for category, suffix in [ + ("Functions", ""), + ("Operators", "_ops"), + ("NativeMetaFunctions", "_meta"), + ("NativeFunctions", "_native"), + ]: + cpu_fm.write( + f"{category}.h", + lambda: { + f"{category}_includes": [ + f"#include " + for name in sorted(functions_by_root_name.keys()) + ], + f"{category}_declarations": [], + }, + ) + + for dispatch_key in dispatch_keys: + if dispatch_key not in functions_keys: + continue + + dispatch_namespace = dispatch_key.lower() + dispatch_names = [] + + for name, functions in functions_by_root_name.items(): + grouped_functions = grouped_functions_by_root_name.get(name, []) + declarations = list( + concatMap( + dest.RegisterDispatchKey( + backend_indices[dispatch_key], + Target.NAMESPACED_DECLARATION, + selector, + rocm=rocm, + symint=True, + class_method_name=None, + skip_dispatcher_op_registration=False, + ), + grouped_functions, + ) + ) + + if len(declarations) == 0: + continue + + dispatch_names.append(name) + ops_fm.write_with_template( + f"{name}_{dispatch_namespace}_dispatch.h", + "DispatchKeyFunction.h", + lambda: { + "dispatch_namespace": dispatch_namespace, + "dispatch_namespaced_declarations": declarations, + }, + ) + + fm = file_manager_from_dispatch_key(dispatch_key, device_fms, cpu_fm) + inl_headers = f"#include " + + fm.write_with_template( + f"{dispatch_key}Functions.h", + "DispatchKeyFunctions.h", + lambda: { + "dispatch_key": str(dispatch_key), + "inline_headers": inl_headers, + }, + ) + fm.write_with_template( + f"{dispatch_key}Functions_inl.h", + "DispatchKeyFunctions_inl.h", + lambda: { + "dispatch_namespace": dispatch_namespace, + "DispatchKeyFunctions_inl_includes": [ + f"#include " + for name in sorted(dispatch_names) + ], + "dispatch_namespaced_declarations": [], + }, + ) + del fm + + cpu_fm.write( + "MethodOperators.h", + lambda: { + "MethodOperators_includes": sorted( + f"#include " + for name, functions in functions_by_root_name.items() + if any(Variant.method in fn.variants for fn in functions) + ), + "MethodOperators_declarations": [], + }, + ) + + +def gen_headers( + *, + native_functions: Sequence[NativeFunction], + valid_tags: set[str], + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + structured_native_functions: Sequence[NativeFunctionsGroup], + static_dispatch_idx: list[BackendIndex], + selector: SelectiveBuilder, + backend_indices: dict[DispatchKey, BackendIndex], + headeronly_fm: FileManager, + core_fm: FileManager, + cpu_fm: FileManager, + device_fms: dict[str, FileManager], + ops_fm: FileManager, + dispatch_keys: Sequence[DispatchKey], + functions_keys: set[DispatchKey], + rocm: bool, + per_operator_headers: bool, +) -> None: + if per_operator_headers: + gen_per_operator_headers( + native_functions=native_functions, + grouped_native_functions=grouped_native_functions, + static_dispatch_idx=static_dispatch_idx, + selector=selector, + backend_indices=backend_indices, + cpu_fm=cpu_fm, + device_fms=device_fms, + ops_fm=ops_fm, + dispatch_keys=dispatch_keys, + functions_keys=functions_keys, + rocm=rocm, + ) + else: + gen_aggregated_headers( + native_functions=native_functions, + grouped_native_functions=grouped_native_functions, + structured_native_functions=structured_native_functions, + static_dispatch_idx=static_dispatch_idx, + selector=selector, + backend_indices=backend_indices, + cpu_fm=cpu_fm, + device_fms=device_fms, + dispatch_keys=dispatch_keys, + functions_keys=functions_keys, + rocm=rocm, + ) + + core_fm.write( + "TensorBody.h", + lambda: { + "tensor_method_declarations": list( + mapMaybe( + ComputeTensorMethod( + target=Target.DECLARATION, + static_dispatch_backend_indices=static_dispatch_idx, + ), + native_functions, + ) + ), + "tensor_method_definitions": list( + mapMaybe( + ComputeTensorMethod( + target=Target.DEFINITION, + static_dispatch_backend_indices=static_dispatch_idx, + ), + native_functions, + ) + ), + }, + ) + + cpu_fm.write( + "RedispatchFunctions.h", + lambda: { + "function_redispatch_definitions": list( + mapMaybe(ComputeRedispatchFunction(), native_functions) + ), + }, + ) + + cpu_fm.write( + "RegistrationDeclarations.h", + lambda: { + "registration_declarations": [ + compute_registration_declarations(f, backend_indices) + for f in native_functions + ], + }, + ) + + cpu_fm.write( + "VmapGeneratedPlumbing.h", lambda: gen_all_vmap_plumbing(native_functions) + ) + + def gen_aten_interned_strings() -> dict[str, str]: + attrs: set[str] = set() # All function argument names + names = set() # All ATen function names + for func in native_functions: + names.add(str(func.func.name.name)) + # Some operators don't have a functional variant but we still create a + # symbol without the underscore + names.add(func.func.name.name.base) + + attrs.update(arg.name for arg in func.func.schema_order_arguments()) + + # These are keywords in C++, so aren't valid symbol names + # https://en.cppreference.com/w/cpp/language/operator_alternative + names -= { + "and", + "and_eq", + "bitand", + "bitor", + "compl", + "not", + "not_eq", + "or", + "or_eq", + "xor", + "xor_eq", + } + + return { + "aten_symbols": " \\\n".join( + [f"_(aten, {name})" for name in sorted(names)] + ), + "attr_symbols": " \\\n".join( + [f"_(attr, {name})" for name in sorted(attrs)] + ), + } + + core_fm.write("aten_interned_strings.h", gen_aten_interned_strings) + + def gen_tags_enum() -> dict[str, str]: + return {"enum_of_valid_tags": (",\n".join(sorted(valid_tags)))} + + headeronly_fm.write("enum_tag.h", gen_tags_enum) + + +def gen_source_files( + *, + native_functions: Sequence[NativeFunction], + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + structured_native_functions: Sequence[NativeFunctionsGroup], + view_groups: Sequence[NativeFunctionsViewGroup], + selector: SelectiveBuilder, + static_dispatch_idx: list[BackendIndex], + backend_indices: dict[DispatchKey, BackendIndex], + aoti_fm: FileManager, + core_fm: FileManager, + cpu_vec_fm: FileManager, + cpu_fm: FileManager, + device_fms: dict[str, FileManager], + dispatch_keys: Sequence[DispatchKey], + functions_keys: set[DispatchKey], + rocm: bool, + force_schema_registration: bool, + per_operator_headers: bool, + skip_dispatcher_op_registration: bool, + update_aoti_c_shim: bool, + aoti_backends: set[DispatchKey | None], + extend_aoti_c_shim: bool, +) -> None: + extra_cuda_headers = """\ +#include +#include +#include +#include """ + if rocm: + extra_cuda_headers = """\ +#include +#include +#include +#include """ + + for dispatch_key in dispatch_keys: + fm = file_manager_from_dispatch_key(dispatch_key, device_fms, cpu_fm) + if per_operator_headers: + + def operator_headers() -> list[str]: + headers = [] + for g in grouped_native_functions: + is_registered = False + if backend_index.has_kernel(g): + is_registered = True + # The above has_kernel test on a group will only test for + # the existence of out dispatch, because that's how + # structured kernels work. But sometimes functions can be + # grouped but not be structured, and then you need to check + # each individual piece, as they may have manual dispatch + # entries. + elif isinstance(g, NativeFunctionsGroup) and any( + backend_index.has_kernel(fn) for fn in g.functions() + ): + is_registered = True + # TODO: this condition is a bit questionable + # (It has to do with the fact that structured kernels get generated kernels + # to the Meta + CompositeExplicitAutogradNonFunctional keys). + elif g.structured and dispatch_key in ( + DispatchKey.Meta, + DispatchKey.CompositeExplicitAutogradNonFunctional, + ): + is_registered = True + if not is_registered: + continue + + headers.append(f"#include ") + if ( + dispatch_key + == DispatchKey.CompositeExplicitAutogradNonFunctional + ): + headers.append(f"#include ") + if dispatch_key in functions_keys: + headers.append( + f"#include " + ) + + return sorted(set(headers)) + + else: + + def operator_headers() -> list[str]: + headers = ["#include "] + if dispatch_key == DispatchKey.CompositeExplicitAutogradNonFunctional: + headers.append("#include ") + if dispatch_key in functions_keys: + headers.append(f"#include ") + return headers + + backend_index = backend_indices[dispatch_key] + ns_grouped_native_functions = defaultdict(list) + for grouped_native_function in grouped_native_functions: + namespace = ( + grouped_native_function.namespace + if isinstance(grouped_native_function, NativeFunction) + else grouped_native_function.functional.namespace + ) + ns_grouped_native_functions[namespace].append(grouped_native_function) + + dispatch_namespace = str(dispatch_key).lower() + + # CompositeImplicitAutogradNestdTensor does not currently user the helpers generated + # compilation will fail when `-Werror=unused-function` flag is set + gen_dispatch_helpers: bool = ( + dispatch_key != DispatchKey.CompositeImplicitAutogradNestedTensor + ) + + register_dispatch_key_base_env = { + "extra_cuda_headers": extra_cuda_headers + if is_cuda_dispatch_key(dispatch_key) + else "", + "external_backend_headers": "", + "dispatch_headers": dest.gen_registration_headers( + backend_index, per_operator_headers, rocm + ), + # ops_headers *could* be sharded, but doesn't seem necessary? + "ops_headers": operator_headers(), + "dispatch_helpers": ( + dest.gen_registration_helpers(backend_index) + if gen_dispatch_helpers + else [] + ), + } + + def register_dispatch_key_env_callable( + gnf: NativeFunction | NativeFunctionsGroup, + ) -> dict[str, list[str]]: + return { + "dispatch_definitions": get_native_function_definitions( + fm=fm, # noqa: F821 + grouped_native_functions=[gnf], + dispatch_key=dispatch_key, + backend_idx=backend_index, + selector=selector, + rocm=rocm, + symint=True, + skip_dispatcher_op_registration=skip_dispatcher_op_registration, + gen_dispatch_helpers=gen_dispatch_helpers, + ) + } + + fm.write_sharded_with_template( + f"Register{dispatch_key}.cpp", + "RegisterDispatchKey.cpp", + grouped_native_functions, + key_fn=lambda x: x.root_name, + env_callable=register_dispatch_key_env_callable, + num_shards=4 if dispatch_key == DispatchKey.CPU else 1, + base_env=register_dispatch_key_base_env, + sharded_keys={"dispatch_definitions"}, + ) + + for g in structured_native_functions: + if not g.out.ufunc_inner_loop or not is_ufunc_dispatch_key(dispatch_key): + continue + name = g.functional.func.name.name + if dispatch_key is DispatchKey.CPU: + if fm is not cpu_fm: + raise AssertionError("Expected fm to be cpu_fm for DispatchKey.CPU") + fm.write_with_template( + f"UfuncCPU_{name}.cpp", + "UfuncCPU.cpp", + lambda: { + "meta_declaration": compute_meta_function_declaration(g), + "native_declaration": dest.compute_native_function_declaration( + g, backend_indices[dispatch_key] + ), + "native_definitions": dest.compute_ufunc_cpu(g), + }, + ) + cpu_vec_fm.write_with_template( + f"UfuncCPUKernel_{name}.cpp", + "UfuncCPUKernel.cpp", + lambda: { + "name": name, + "native_definitions": dest.compute_ufunc_cpu_kernel(g), + }, + ) + elif dispatch_key is DispatchKey.CUDA: + cuda_headers = "#include " + if rocm: + cuda_headers = "#include " + fm.write_with_template( + f"UfuncCUDA_{name}.cu", + "UfuncCUDA.cu", + lambda: { + "name": name, + "cuda_headers": cuda_headers, + "meta_declaration": compute_meta_function_declaration(g), + "native_declaration": dest.compute_native_function_declaration( + g, backend_indices[dispatch_key] + ), + "native_definitions": dest.compute_ufunc_cuda(g), + }, + ) + else: + raise AssertionError(f"unrecognized {dispatch_key} for ufunc") + + del fm + + gen_aoti_c_shim_files( + aoti_fm=aoti_fm, + aoti_backends=aoti_backends, + native_functions=native_functions, + backend_indices=backend_indices, + structured_native_functions=structured_native_functions, + extra_cuda_headers=extra_cuda_headers, + update_aoti_c_shim=update_aoti_c_shim, + extend_aoti_c_shim=extend_aoti_c_shim, + ) + + # BackendSelect is generated specially + def gen_backend_select() -> dict[str, list[str]]: + relevant_fns = [ + fn for fn in native_functions if needs_backend_select(fn, selector) + ] + return { + "ops_headers": [ + f"#include " for fn in relevant_fns + ], + "backend_select_method_definitions": list( + mapMaybe( + ComputeBackendSelect(Target.DEFINITION, selector), relevant_fns + ) + ), + "backend_select_function_registrations": list( + mapMaybe( + ComputeBackendSelect(Target.REGISTRATION, selector), relevant_fns + ) + ), + } + + cpu_fm.write("RegisterBackendSelect.cpp", gen_backend_select) + + schema_selector = selector + if force_schema_registration: + schema_selector = SelectiveBuilder.get_nop_selector() + + ( + aten_schema_registrations, + schema_registrations, + ) = get_native_function_schema_registrations( + native_functions=native_functions, schema_selector=schema_selector + ) + cpu_fm.write( + "RegisterSchema.cpp", + lambda: { + "aten_schema_registrations": [] + if skip_dispatcher_op_registration + else aten_schema_registrations, + "schema_registrations": [] + if skip_dispatcher_op_registration + else schema_registrations, + }, + ) + + def key_func( + fn: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup, + ) -> str: + return fn.root_name + + cpu_fm.write_sharded( + "Operators.cpp", + native_functions, + key_fn=key_func, + env_callable=lambda fn: { + "operator_headers": [f"#include "], + "definitions": [ + ComputeOperators( + Target.DEFINITION, + static_dispatch_backend_indices=static_dispatch_idx, + )(fn) + ], + }, + base_env={ + "static_dispatch_extra_headers": static_dispatch_extra_headers( + static_dispatch_idx + ), + }, + num_shards=5, + sharded_keys={ + "operator_headers", + "definitions", + "static_dispatch_extra_headers", + }, + ) + + cpu_fm.write("Functions.cpp", dict) + + core_fm.write("TensorMethods.cpp", dict) + + core_fm.write( + "ATenOpList.cpp", + lambda: { + "aten_ops": list(mapMaybe(compute_aten_op, native_functions)), + }, + ) + + def gen_op_headers( + g: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup, + ) -> list[str]: + if isinstance(g, NativeFunctionsViewGroup): + # view ops always get a functionalization kernel + headers = [ + f"#include ", + f"#include ", + ] + if g.view_copy is not None: + headers += [ + f"#include ", + f"#include ", + ] + return headers + elif isinstance(g, NativeFunctionsGroup): + headers = [ + f"#include ", + f"#include ", + f"#include ", + f"#include ", + ] + if g.inplace is not None: + headers += [ + f"#include ", + f"#include ", + ] + if g.mutable is not None: + headers += [ + f"#include ", + f"#include ", + ] + return headers + else: + return [ + f"#include ", + f"#include ", + ] + + def functionalization_env_callable( + g: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup, + ) -> dict[str, list[str]]: + return { + "ops_headers": gen_op_headers(g), + "func_definitions": gen_functionalization_definition( + selector, + g, + ), + "func_registrations": gen_functionalization_registration( + selector, + g, + backend_indices[DispatchKey.CompositeImplicitAutograd], + ), + } + + all_groups: list[ + NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup + ] = list(structured_native_functions) + list( + view_groups # type: ignore[assignment, arg-type, operator] + ) + # Note: all operators that functionalization needs to handle (mutable and aliasing ops) should be grouped properly. + # The only reason we really need to deal with direct NativeFunctions here (instead of the groups) is because: + # (1) We can provide better error checking (error out if someone introduces a mutable op that doesn't obey the grouping logic) + # (2) functionalization needs to manually register CompositeImplicitAutograd kernels, which might not be grouped. + # Although this could go away long-term if we add a dedicated dispatch key for decompositions. + structured_map: dict[OperatorName, NativeFunction] = { + f.func.name: f + for f in concatMap(lambda g: list(g.functions()), structured_native_functions) + } + view_map: dict[OperatorName, NativeFunction] = { + f.func.name: f for f in concatMap(lambda g: list(g.functions()), view_groups) + } + all_groups.extend( + f + for f in native_functions + if f.func.name not in structured_map and f.func.name not in view_map + ) + + cpu_fm.write_sharded( + "RegisterFunctionalization.cpp", + all_groups, + key_fn=key_func, + env_callable=functionalization_env_callable, + num_shards=4, + sharded_keys={ + "ops_headers", + "func_definitions", + "func_registrations", + "func_add_back_views_definitions", + "func_add_back_views_registrations", + }, + ) + + cpu_fm.write( + "FunctionalInverses.h", + lambda: { + "view_inverse_declarations": list( + mapMaybe( + lambda g: gen_functionalization_view_inverse_declaration( + selector, g + ), + view_groups, + ) + ) + }, + ) + + cpu_fm.write( + "ViewMetaClasses.h", + lambda: { + "view_meta_declarations": list( + concatMap( + lambda g: gen_functionalization_view_meta_classes_decl(selector, g), + view_groups, + ) + ) + }, + ) + + cpu_fm.write( + "ViewMetaClasses.cpp", + lambda: { + "view_meta_implementations": list( + concatMap( + lambda g: gen_functionalization_view_meta_classes_impl(selector, g), + view_groups, + ) + ), + "op_headers": list(concatMap(gen_op_headers, view_groups)), + }, + ) + + # Note [view_copy NativeFunctions] + # Every view operator in native_functions.yaml that is not CompositeImplicitAutograd + # needs to have a corresponding non-aliasing {view}_copy variant. + # Backends that use functionalization and don't know how to handle aliasing ops + # are expected to implement kernels for these {view}_copy kernels instead. + # The code for {view}_copy operators in core is pretty boilerplate-heavy however, + # so we codegen the following: + # (1) A CompositeExplicitAutogradNonFunctional kernel for every {view}_copy operator. + # These are never explicitly invoked by the functionalization pass, + # but they could theoretically be called from user code (I added these kernels for completeness, + # since the ops are part of the public API). + # (2) A derivative formula for every {view}_copy operator + # {view}_copy operators can reuse the same derivative formulas as their {view} op counterparts, + # so rather than stamping all of the entries out in derivatives.yaml, + # we codegen them in. + # This is similar to how autograd codegen doesn't require inplace ops to have a derivatives.yaml entry. + cpu_fm.write( + "CompositeViewCopyKernels.cpp", + lambda: { + "ops_headers": [ + "\n".join( + f"#include \n" + # NB: this include is important as it ensures we + # set the visibility on generated view_copy kernels + # correctly + f"#include " + for f in ( + [g.view] if g.view_copy is None else [g.view, g.view_copy] + ) + ) + for g in view_groups + ] + + [ + "\n".join( + f"#include \n" + # NB: this include is also important for correct visibility + f"#include " + for f in [g.inplace, g.mutable, g.functional] + if f is not None and "generated" not in f.tags + ) + for g in structured_native_functions + ], + "CompositeViewCopyKernel_Definitions": list( + mapMaybe( + GenCompositeViewCopyKernel( + backend_indices[ + DispatchKey.CompositeExplicitAutogradNonFunctional + ] + ), + view_groups, + ) + ), + "GeneratedCompositeFunctional_Definitions": list( + mapMaybe( + gen_composite_functional_kernel, + structured_native_functions, + ) + ), + "GeneratedCompositeOut_Definitions": list( + mapMaybe( + gen_composite_out_kernel, + structured_native_functions, + ) + ), + }, + ) + + +def gen_declarations_yaml( + cpu_fm: FileManager, native_functions: Sequence[NativeFunction] +) -> None: + cpu_fm.write( + "Declarations.yaml", + lambda: format_yaml([compute_declaration_yaml(f) for f in native_functions]), + ) + + +def get_torchgen_root() -> Path: + """ + If you're depending on torchgen out-of-tree, you can use the root to figure + out the path to native_functions.yaml + """ + return Path(__file__).parent.resolve() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate ATen source files") + parser.add_argument( + "-s", + "--source-path", + help="path to source directory for ATen", + default="aten/src/ATen", + ) + parser.add_argument( + "-o", + "--output-dependencies", + help="output a list of dependencies into the given file and exit", + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="run without writing any files (still updates outputs)", + ) + parser.add_argument( + "--per-operator-headers", + action="store_true", + help="generate separate headers per operator in ATen/ops", + ) + parser.add_argument( + "-d", + "--install-dir", + "--install_dir", + help="output directory", + default="build/aten/src/ATen", + ) + parser.add_argument( + "--aoti-install-dir", + "--aoti_install_dir", + help="output directory for AOTInductor shim", + default="torch/csrc/inductor/aoti_torch/generated", + ) + parser.add_argument( + "--headeronly-install-dir", + "--headeronly_install_dir", + help="output directory for header-only generated files (e.g. enum_tag.h). " + "Defaults to `/core` when --install-dir is set, otherwise " + "`build/torch/headeronly/core`.", + default=None, + ) + parser.add_argument( + "--rocm", + action="store_true", + help="reinterpret CUDA as ROCm/HIP and adjust filepaths accordingly", + ) + parser.add_argument( + "--mps", + action="store_true", + help="Generate MPS registration code when set", + ) + parser.add_argument( + "--xpu", + action="store_true", + help="Generate XPU registration code when set", + ) + parser.add_argument( + "--mtia", + action="store_true", + help="Generate MTIA registration code when set", + ) + + # TODO: --op-registration-whitelist will be removed when all call-sites + # for gen.py are moved over to using the operator YAML file for mobile + # custom build. + parser.add_argument( + "--op-registration-whitelist", + "--op_registration_whitelist", + nargs="*", + help="filter op registrations by the whitelist (if set); " + "each item is `namespace`::`operator name` without overload name; " + "e.g.: aten::empty aten::conv2d ...", + ) + parser.add_argument( + "--op-selection-yaml-path", + "--op_selection_yaml_path", + help="Provide a path to the operator selection (for custom build) YAML " + "that contains the information about the set of selected operators " + "and their categories (training, ...). Each operator is either a " + "full operator name with overload or just a bare operator name. " + "The operator names also contain the namespace prefix (e.g. aten::)", + ) + parser.add_argument( + "--backend-whitelist", + "--backend_whitelist", + nargs="*", + help="filter dispatch backend by the whitelist (if set), " + "e.g.: CPU CUDA QuantizedCPU ...", + ) + parser.add_argument( + "--static-dispatch-backend", + "--static_dispatch_backend", + nargs="*", + help="generate static dispatch code for the specific backend (if set)", + ) + parser.add_argument( + "--skip-dispatcher-op-registration", + "--skip_dispatcher_op_registration", + action="store_true", + help="Avoid registering operators into the dispatcher.", + ) + parser.add_argument( + "--force-schema-registration", + "--force_schema_registration", + action="store_true", + help="force it to generate schema-only registrations for all ops, including" + "those that are not listed on --op-registration-whitelist", + ) + parser.add_argument( + "--generate", + type=str, + nargs="*", + choices=["headers", "sources", "declarations_yaml"], + default=["headers", "sources", "declarations_yaml"], + help="Generate only a subset of files", + ) + parser.add_argument( + "--update-aoti-c-shim", + action="store_true", + help="Update AOTInductor C shim after adding an entry to inductor_fallback_ops in torchgen/aoti/fallback_ops.py. " + "WARNING: Do not use this unless you are sure what you are doing!!!", + ) + parser.add_argument( + "--extend-aoti-c-shim", + action="store_true", + help="This Flag indicates the generation of c shims for out-of-tree ATen ops," + "which is an extension to the In-tree ATen op c shims. This flag needs to be combined with" + "---source-path=" + "--aoti-install-dir=/extend" + " default is torch/csrc/inductor/aoti_torch/generated/extend" + "WARNING: Do not use this unless you are sure what you are doing!!!", + ) + + options = parser.parse_args() + + selector = get_custom_build_selector( + options.op_registration_whitelist, + options.op_selection_yaml_path, + ) + + native_yaml_path = os.path.join(options.source_path, "native/native_functions.yaml") + tags_yaml_path = os.path.join(options.source_path, "native/tags.yaml") + + from torchgen.model import dispatch_keys + + # Only a limited set of dispatch keys get CPUFunctions.h headers generated + # for them; this is the set + functions_keys = { + DispatchKey.CPU, + DispatchKey.CUDA, + DispatchKey.CompositeImplicitAutograd, + DispatchKey.CompositeImplicitAutogradNestedTensor, + DispatchKey.CompositeExplicitAutograd, + DispatchKey.CompositeExplicitAutogradNonFunctional, + DispatchKey.Meta, + DispatchKey.MTIA, + } + + aoti_backends = { + DispatchKey.CPU, + DispatchKey.CUDA, + # None will generate the aten shim based on aten_shimified_ops + # which does not bypass the dispatcher + None, + } + + # TODO: stop generating CUDA kernels for non-CUDA builds + ignore_keys = set() + + MPS_KEYS = {DispatchKey.MPS, DispatchKey.SparseMPS, DispatchKey.SparseCsrMPS} + if options.mps or options.update_aoti_c_shim: + functions_keys.update(MPS_KEYS) + aoti_backends.add(DispatchKey.MPS) + else: + ignore_keys.update(MPS_KEYS) + dispatch_keys[:] = [k for k in dispatch_keys if k not in MPS_KEYS] + + if options.xpu or options.update_aoti_c_shim: + functions_keys.add(DispatchKey.XPU) + aoti_backends.add(DispatchKey.XPU) + else: + ignore_keys.add(DispatchKey.XPU) + + if DispatchKey.XPU in dispatch_keys: + del dispatch_keys[dispatch_keys.index(DispatchKey.XPU)] + + if not options.mtia: + ignore_keys.add(DispatchKey.MTIA) + + if DispatchKey.MTIA in dispatch_keys: + del dispatch_keys[dispatch_keys.index(DispatchKey.MTIA)] + + if options.backend_whitelist: + dispatch_keys = [ + k + for k in dispatch_keys + if is_generic_dispatch_key(k) or str(k) in options.backend_whitelist + ] + + parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path, ignore_keys) + valid_tags = _GLOBAL_PARSE_TAGS_YAML_CACHE[tags_yaml_path] + native_functions, backend_indices = ( + parsed_yaml.native_functions, + parsed_yaml.backend_indices, + ) + + grouped_native_functions = get_grouped_native_functions(native_functions) + + structured_native_functions = [ + g for g in grouped_native_functions if isinstance(g, NativeFunctionsGroup) + ] + native_functions_with_view_groups = get_grouped_by_view_native_functions( + native_functions + ) + view_groups = [ + g + for g in native_functions_with_view_groups + if isinstance(g, NativeFunctionsViewGroup) + ] + + # NB: It is mandatory to NOT use os.path.join here, as the install directory + # will eventually be ingested by cmake, which does not respect Windows style + # path slashes. If you switch this to use os.path.join, you'll get an error + # like: + # + # Syntax error in cmake code when parsing string + # + # C:/Jenkins/workspace/pytorch-builds/pytorch-win-ws2016-cuda9-cudnn7-py3-build/build/aten/src/ATen\core/TensorMethods.h + # + # Invalid character escape '\c'. + core_install_dir = f"{options.install_dir}/core" + Path(core_install_dir).mkdir(parents=True, exist_ok=True) + ops_install_dir = f"{options.install_dir}/ops" + Path(ops_install_dir).mkdir(parents=True, exist_ok=True) + + aoti_install_dir = f"{options.aoti_install_dir}" + Path(aoti_install_dir).mkdir(parents=True, exist_ok=True) + + if options.headeronly_install_dir is not None: + headeronly_install_dir = options.headeronly_install_dir + elif options.install_dir is not None: + headeronly_install_dir = f"{options.install_dir}/core" + else: + headeronly_install_dir = "build/torch/headeronly/core" + Path(headeronly_install_dir).mkdir(parents=True, exist_ok=True) + + core_fm = make_file_manager(options=options, install_dir=core_install_dir) + cpu_fm = make_file_manager(options=options) + cpu_vec_fm = make_file_manager(options=options) + cuda_fm = make_file_manager(options=options) + ops_fm = make_file_manager(options=options, install_dir=ops_install_dir) + aoti_fm = make_file_manager(options=options, install_dir=aoti_install_dir) + headeronly_fm = make_file_manager( + options=options, install_dir=headeronly_install_dir + ) + device_fms = {"cuda": cuda_fm} + if options.xpu: + device_fms["xpu"] = make_file_manager(options=options) + + static_dispatch_idx: list[BackendIndex] = [] + if options.static_dispatch_backend: + static_dispatch_idx = [ + backend_indices[DispatchKey.parse(key)] + for key in options.static_dispatch_backend + ] + for key in options.static_dispatch_backend: + dp_key = DispatchKey.parse(key) + if dp_key not in functions_keys: + functions_keys.add(dp_key) + + if "sources" in options.generate: + gen_source_files( + native_functions=native_functions, + grouped_native_functions=grouped_native_functions, + structured_native_functions=structured_native_functions, + view_groups=view_groups, + selector=selector, + static_dispatch_idx=static_dispatch_idx, + backend_indices=backend_indices, + aoti_fm=aoti_fm, + core_fm=core_fm, + cpu_vec_fm=cpu_vec_fm, + cpu_fm=cpu_fm, + device_fms=device_fms, + dispatch_keys=dispatch_keys, + functions_keys=functions_keys, + rocm=options.rocm, + force_schema_registration=options.force_schema_registration, + per_operator_headers=options.per_operator_headers, + skip_dispatcher_op_registration=options.skip_dispatcher_op_registration, + update_aoti_c_shim=options.update_aoti_c_shim, + aoti_backends=aoti_backends, + extend_aoti_c_shim=options.extend_aoti_c_shim, + ) + + if "headers" in options.generate: + gen_headers( + native_functions=native_functions, + valid_tags=valid_tags, + grouped_native_functions=grouped_native_functions, + structured_native_functions=structured_native_functions, + static_dispatch_idx=static_dispatch_idx, + selector=selector, + backend_indices=backend_indices, + headeronly_fm=headeronly_fm, + core_fm=core_fm, + cpu_fm=cpu_fm, + device_fms=device_fms, + ops_fm=ops_fm, + dispatch_keys=dispatch_keys, + functions_keys=functions_keys, + rocm=options.rocm, + per_operator_headers=options.per_operator_headers, + ) + + if "declarations_yaml" in options.generate: + gen_declarations_yaml(native_functions=native_functions, cpu_fm=cpu_fm) + + if options.output_dependencies: + depfile_path = Path(options.output_dependencies).resolve() + depfile_name = depfile_path.name + depfile_stem = depfile_path.stem + + for fm, prefix in [ + (cpu_fm, ""), + (cpu_vec_fm, "cpu_vec_"), + (core_fm, "core_"), + (ops_fm, "ops_"), + ] + [(device_fm, f"{device}_") for device, device_fm in device_fms.items()]: + varname = prefix + depfile_stem + path = depfile_path.parent / (prefix + depfile_name) + fm.write_outputs(varname, str(path)) + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_aoti_c_shim.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_aoti_c_shim.py new file mode 100644 index 0000000000000000000000000000000000000000..dde6a6c8eda987cdc32c3c486ee5ffe59f89c8dc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_aoti_c_shim.py @@ -0,0 +1,783 @@ +from __future__ import annotations + +import difflib +import os +import textwrap +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from torchgen.aoti.fallback_ops import aten_shimified_ops, inductor_fallback_ops +from torchgen.api.types import DispatcherSignature +from torchgen.api.types.signatures import CppSignature, CppSignatureGroup +from torchgen.context import method_with_native_function +from torchgen.model import ( + Argument, + BackendIndex, + BaseTy, + BaseType, + DispatchKey, + FunctionSchema, + is_cuda_dispatch_key, + ListType, + NativeFunction, + NativeFunctionsGroup, + OperatorName, + OptionalType, + Type, + Variant, +) +from torchgen.utils import FileManager, mapMaybe + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +base_type_to_c_type = { + BaseTy.Tensor: "AtenTensorHandle", + BaseTy.bool: "int32_t", # Use int to pass bool + BaseTy.int: "int64_t", + BaseTy.SymInt: "int64_t", # Inductor-generated code won't see a SymInt + BaseTy.Scalar: "double", # Use double to pass both integer and floating point + BaseTy.float: "double", # TODO: how about other floating point types? + BaseTy.str: "const char*", + BaseTy.DeviceIndex: "int32_t", + BaseTy.Layout: "int32_t", # Represent enum as int + BaseTy.MemoryFormat: "int32_t", # Represent enum as int + BaseTy.ScalarType: "int32_t", # Represent enum as int + BaseTy.Generator: "AtenGeneratorHandle", +} + +base_type_to_aten_type = { + BaseTy.Tensor: "at::Tensor", + BaseTy.bool: "bool", + BaseTy.int: "int64_t", + BaseTy.SymInt: "c10::SymInt", + BaseTy.Scalar: "c10::Scalar", + BaseTy.float: "double", + BaseTy.str: "::std::string_view", + BaseTy.DeviceIndex: "c10::DeviceIndex", + BaseTy.Layout: "c10::Layout", + BaseTy.MemoryFormat: "c10::MemoryFormat", + BaseTy.ScalarType: "c10::ScalarType", + BaseTy.Generator: "at::Generator", +} + +base_type_to_callsite_expr = { + BaseTy.Tensor: "resolve_tensor_dispatch_flags", + BaseTy.bool: "", + BaseTy.int: "", + BaseTy.SymInt: "", + BaseTy.Scalar: "", + BaseTy.float: "", + BaseTy.str: "", + BaseTy.DeviceIndex: "static_cast", + BaseTy.Layout: "static_cast", + BaseTy.MemoryFormat: "static_cast", + BaseTy.ScalarType: "static_cast", + BaseTy.Generator: "*generator_handle_to_generator_pointer", +} + + +# convert args to C types, names in declarations, and expressions in function bodies +def convert_arg_type_and_name( + typ: Type, + name: str, + is_write: bool = False, +) -> tuple[list[str], list[str], list[str], list[str]]: + if isinstance(typ, BaseType): + if typ.name in base_type_to_c_type: + if typ.name == BaseTy.Tensor and is_write: + # For output tensors, our normal call to resolve_tensor_dispatch_flags + # results in an rvalue tensor, which can't be passed to at::Tensor&. + # Override this case specifically. + callsite_expr = [f"*tensor_handle_to_tensor_pointer({name})"] + else: + callsite_expr = [ + f"{base_type_to_callsite_expr[typ.name]}({name})" + if base_type_to_callsite_expr[typ.name] + else name + ] + + return ( + [base_type_to_c_type[typ.name]], + [name], + [base_type_to_aten_type[typ.name]], + callsite_expr, + ) + elif typ.name == BaseTy.Device: + return ( + ["int32_t", "int32_t"], + [name, name + "_index_"], + ["c10::Device"], + [ + f"c10::Device(static_cast({name}), static_cast({name}_index_))" + ], + ) + else: + # TODO: BaseTy.Dimname, etc. + raise NotImplementedError(f"TODO: add support for arg type {repr(typ)}") + elif isinstance(typ, OptionalType): + c_types, names, aten_types, callsite_exprs = convert_arg_type_and_name( + typ.elem, name + ) + j = 0 # index for names + new_aten_types = [] + new_callsite_exprs = [] + for aten_type in aten_types: + # Use pointer to denote optional type + c_types[j] = c_types[j] + "*" + if aten_type.startswith("c10::ArrayRef<"): + # ArrayRef is passed as pointer + size, but no need to add "*" to the size argument + new_aten_types.append(f"::std::optional<{aten_type}>") + base_type = aten_type[len("c10::ArrayRef<") : -1] + new_callsite_exprs.append( + f"pointer_to_optional_list<{base_type}>({names[j]}, {names[j + 1]})" + ) + j += 2 + elif aten_type == "c10::Device": + # Device is passed as device_type + device_index + new_aten_types.append("::std::optional") + new_callsite_exprs.append( + f"pointer_to_optional_device({names[j]}, {names[j + 1]})" + ) + j += 2 + elif aten_type == "at::Tensor": + new_aten_types.append(f"::std::optional<{aten_type}>") + new_callsite_exprs.append(f"resolve_tensor_dispatch_flags({names[j]})") + j += 1 + else: + new_aten_types.append(f"::std::optional<{aten_type}>") + new_callsite_exprs.append( + f"pointer_to_optional<{aten_type}>({names[j]})" + ) + j += 1 + + return ( + c_types, + names, + new_aten_types, + new_callsite_exprs, + ) + elif isinstance(typ, ListType): + # Need to explicitly pass the list as pointer + length + c_types, names, aten_types, _ = convert_arg_type_and_name(typ.elem, name) + if len(c_types) != 1: + raise AssertionError(f"ListType with unsupported element type {repr(typ)}") + + # The list content should never be modified + c_types[0] = f"const {c_types[0]}*" + c_types.append("int64_t") + name = names[0] + names.append(name + "_len_") + + atype = aten_types[0] + callsite_exprs = [] + if atype == "bool": + # no converter from std::vector to c10::ArrayRef + # construct std::array instead + if typ.size is None: + raise AssertionError("bool ListType must have a size") + callsite_exprs.append(f"pointer_to_list<{typ.size}>({name})") + elif atype == "at::Tensor" and not is_write: + callsite_exprs.append( + f"resolve_tensor_list_dispatch_flags({name}, {name}_len_)" + ) + elif atype == "::std::optional": + # convert from std::vector<::std::optional> to c10::List<::std::optional> + callsite_exprs.append( + f"c10::List<{atype}>(c10::ArrayRef<{atype}>(resolve_tensor_list_dispatch_flags({name}, {name}_len_)))" + ) + else: + callsite_exprs.append(f"pointer_to_list<{atype}>({name}, {name}_len_)") + + aten_types = [f"c10::ArrayRef<{t}>" for t in aten_types] + return ( + c_types, + names, + aten_types, + callsite_exprs, + ) + raise NotImplementedError(f"Argument type {repr(typ)} not supported!") + + +def zip_type_and_name(types: list[str], names: list[str]) -> list[str]: + return [typ + " " + name for typ, name in zip(types, names)] + + +# Generate argument declarations and callsite expressions +def gen_arguments( + flat_arguments: Sequence[Argument], skipped_args: set[str] +) -> tuple[list[str], list[str]]: + types: list[str] = [] + new_names: list[str] = [] + callsite_exprs: list[str] = [] + for arg in flat_arguments: + if arg.name in skipped_args: + # Pass the arg's schema default when available (e.g. "false" for + # a bool arg with default=False), so non-optional args with defaults + # can be versioned too. Fall back to std::nullopt for optional args + # with no default (matches historical behavior). + if arg.default is not None: + from torchgen.api.cpp import default_expr + + callsite_exprs.append(default_expr(arg.default, arg.type, symint=False)) + else: + callsite_exprs.append("std::nullopt") + continue + new_types, names, _, new_callsite_exprs = convert_arg_type_and_name( + arg.type, arg.name, arg.is_write + ) + types.extend(new_types) + new_names.extend(names) + callsite_exprs.extend(new_callsite_exprs) + return zip_type_and_name(types, new_names), callsite_exprs + + +# Return values are passed out as pointer arguments because all the C shim functions +# are expected to return AOTITorchError. +# Generate returns as declarations and callsite expressions +def gen_returns(schema: FunctionSchema) -> tuple[list[str], list[str]]: + types = [] + names = [] + for idx, ret in enumerate(schema.returns): + names.append(f"ret{idx}") + if isinstance(ret.type, BaseType) and ret.type.name in base_type_to_c_type: + types.append(base_type_to_c_type[ret.type.name] + "*") + else: + raise NotImplementedError( + f"TODO: add support for return type {repr(ret.type)}" + ) + + def convert_return(typ: BaseType, val: str) -> str: + if typ.name == BaseTy.Tensor: + return f"new_tensor_handle(std::move({val}))" + elif typ.name == BaseTy.SymInt: + return f"{val}.expect_int()" + elif typ.name == BaseTy.Scalar: + return f"{val}.toDouble()" + else: + return val + + ret_pointer_can_be_null = False + unambiguous_name = schema.name.unambiguous_name() + for name in ( + "_functional_sym_constrain_range", + "_scaled_dot_product_cudnn_attention", + "_scaled_dot_product_efficient_attention_backward", + "_scaled_dot_product_efficient_attention", + "_scaled_dot_product_flash_attention", + "_scaled_dot_product_fused_attention_overrideable", + "_thhn_fused_lstm_cell_backward_impl", + "convolution_backward", + "grid_sampler_2d_backward", + "grid_sampler_3d_backward", + "linear_backward", + ): + if name in unambiguous_name: + ret_pointer_can_be_null = True + break + + callsite_exprs: list[str] = [] + for idx, ret in enumerate(schema.returns): + tmp = "tmp_result" if len(names) == 1 else f"std::get<{idx}>(tmp_result)" + if not isinstance(ret.type, BaseType): + raise AssertionError(f"Expected BaseType for return, got {type(ret.type)}") + rval = convert_return(ret.type, tmp) + if ret_pointer_can_be_null: + callsite_exprs.append(f"if ({names[idx]}) {{ *{names[idx]} = {rval}; }}") + else: + callsite_exprs.append(f"*{names[idx]} = {rval};") + + return zip_type_and_name(types, names), callsite_exprs + + +# gen.py generates header first and then src, so caching the result here to avoid duplicate work +declaration_definition_cache: dict[tuple[str, str, str], tuple[str, str]] = {} + + +def gen_declaration_and_definition( + schema: FunctionSchema, + device: str, + backend_call: str, + version_info: dict[str, list[str]], +) -> tuple[str, str]: + base_name = schema.name.unambiguous_name() + + global declaration_definition_cache + if (base_name, device, backend_call) in declaration_definition_cache: + return declaration_definition_cache[(base_name, device, backend_call)] + + # Check the validity of version_info. The format should look like + # {"v2" : ["new_arg1"], "v3": ["new_arg2, new_arg3"]}. + indexed_version_info: dict[int, list[str]] = {1: []} + for ver_str, new_args in sorted(version_info.items()): + if not ver_str.startswith("v"): + raise AssertionError( + f"Version number for {base_name} is {ver_str}, not starting with 'v'" + ) + try: + ver_id = int(ver_str[1:]) + except ValueError as e: + raise AssertionError( + f"Version number for {base_name} is {ver_str}, not a valid integer after 'v'" + ) from e + if ver_id in indexed_version_info: + raise AssertionError(f"{ver_str} for {base_name} has already been defined") + indexed_version_info[ver_id] = new_args + + declarations: list[str] = [] + definitions: list[str] = [] + skipped_args: set[str] = set() + + for ver_id, new_args in sorted(indexed_version_info.items(), reverse=True): + # Iterate in the reverse order, so the latest version of an op will get generated first + # with all the arguments included, while a set of to-be-trimmed args is carried down + # to generate earlier version of the op. + func_name = base_name if ver_id == 1 else f"{base_name}_v{ver_id}" + if schema.is_out_fn(): + # out_variant has out arguments in the front, and it's ok to ignore return values + # because C shim functions only return AOTITorchError + args, callsite_exprs = gen_arguments( + [*schema.arguments.out, *schema.arguments.flat_non_out], skipped_args + ) + ret_assignments: list[str] = [] + else: + args, callsite_exprs = gen_arguments( + schema.arguments.flat_all, skipped_args + ) + # ignore return values for inplace ops + ret_declarations, ret_assignments = ( + ([], []) if schema.name.name.inplace else gen_returns(schema) + ) + args.extend(ret_declarations) + + declaration = textwrap.dedent( + f"AOTITorchError aoti_torch_{device}_{func_name}({', '.join(args)})" + ) + + tmp_result = "auto tmp_result = " if ret_assignments else "" + indent = "\t\t" + ret_assignments_str = ( + "\n".join(indent + r for r in ret_assignments) if ret_assignments else "" + ) + definition = ( + textwrap.dedent(f""" + {declaration} {{ + AOTI_TORCH_CONVERT_EXCEPTION_TO_ERROR_CODE({{ + {tmp_result}{backend_call}( + {", ".join(callsite_exprs)} + ); + """) + + ret_assignments_str + + textwrap.dedent(""" + }); + } + """) + ) + skipped_args.update(new_args) + declarations.append(f"AOTI_TORCH_EXPORT {declaration};") + definitions.append(definition) + + declaration_definition_cache[(base_name, device, backend_call)] = ( + "\n".join(declarations), + "\n".join(definitions), + ) + return declaration_definition_cache[(base_name, device, backend_call)] + + +def gen_static_dispatch_backend_call_signature( + sig: CppSignature | DispatcherSignature, + f: NativeFunction, +) -> CppSignature: + sig = DispatcherSignature.from_schema(f.func) + cpp_sigs = CppSignatureGroup.from_native_function( + f, method=False, fallback_binding=False + ) + if sig.symint and f.func.has_symint(): + cpp_sig = cpp_sigs.symint_signature + else: + cpp_sig = cpp_sigs.signature + if cpp_sig is None: + raise AssertionError(f"No cpp signature found for {f.func.name}") + return cpp_sig + + +def gen_static_dispatch_backend_call( + f: NativeFunction, + backend_index: BackendIndex | None = None, +) -> str: + sig = DispatcherSignature.from_schema(f.func) + cpp_sig = gen_static_dispatch_backend_call_signature(sig, f) + + if backend_index is None: + # Check if this is a symint function and if the function only has method variants + if sig.symint and f.func.has_symint(): + has_function_variant = Variant.function in f.variants + + if not has_function_variant: + # Functions with both function and method variants can use the at::{*}_symint version + # (e.g., narrow -> at::narrow_symint), BUT + # Method-only functions with symint parameters should use at::symint:: namespace + # Remove the _symint suffix since at::symint:: namespace uses the base name + # (e.g., new_empty -> at::symint::new_empty) + base_name = cpp_sig.name() + base_name = base_name.removesuffix("_symint") # Remove "_symint" suffix + return f"at::symint::{base_name}" + + return f"at::{cpp_sig.name()}" + else: + return f"at::{backend_index.dispatch_key.lower()}::{cpp_sig.name()}" + + +def get_backend_index_for_aoti( + func: NativeFunction, + func_group_mapping: dict[OperatorName, NativeFunctionsGroup], + dispatch_key: DispatchKey | None, + backend_indices: dict[DispatchKey, BackendIndex], + extend_aoti_c_shim: bool, +) -> BackendIndex | None: + backend_index = None + + if dispatch_key is None: + return backend_index + + if backend_indices[dispatch_key].has_kernel(func) or ( + func.structured_delegate is not None + and func.structured_delegate in func_group_mapping + and backend_indices[dispatch_key].has_kernel( + func_group_mapping[func.structured_delegate] + ) + ): + backend_index = backend_indices[dispatch_key] + else: + # for the extend out-of-tree kernels, we don't need to + # duplicatly create C shim wrappers for other dispatch keys + if extend_aoti_c_shim: + return backend_index + + elif backend_indices[DispatchKey.CompositeExplicitAutograd].has_kernel(func): + # We need to create C shim wrappers for CompositeExplicitAutograd kernels + backend_index = backend_indices[DispatchKey.CompositeExplicitAutograd] + elif backend_indices[ + DispatchKey.CompositeExplicitAutogradNonFunctional + ].has_kernel(func): + # We need to create C shim wrappers for CompositeExplicitAutogradNonFunctional kernels + backend_index = backend_indices[ + DispatchKey.CompositeExplicitAutogradNonFunctional + ] + elif backend_indices[DispatchKey.CompositeImplicitAutograd].has_kernel(func): + backend_index = backend_indices[DispatchKey.CompositeImplicitAutograd] + + return backend_index + + +def get_header_for_aoti( + func: NativeFunction, + func_group_mapping: dict[OperatorName, NativeFunctionsGroup], + dispatch_key: DispatchKey | None, + backend_indices: dict[DispatchKey, BackendIndex], + extend_aoti_c_shim: bool, +) -> str | None: + backend_index = get_backend_index_for_aoti( + func, func_group_mapping, dispatch_key, backend_indices, extend_aoti_c_shim + ) + if backend_index is None: + if dispatch_key is None: + return f"#include " + return None + + return f"#include " + + +def get_fallback_op_name(func: NativeFunction) -> str: + return ( + f"{func.namespace}.{func.func.name.name}.{func.func.name.overload_name}" + if func.func.name.overload_name + else f"{func.namespace}.{func.func.name.name}.default" + ) + + +def gen_c_shim( + func: NativeFunction, + version_info: dict[str, list[str]], + func_group_mapping: dict[OperatorName, NativeFunctionsGroup], + dispatch_key: DispatchKey | None, + backend_indices: dict[DispatchKey, BackendIndex], + header: bool, + extend_aoti_c_shim: bool, +) -> str | None: + backend_index = get_backend_index_for_aoti( + func, func_group_mapping, dispatch_key, backend_indices, extend_aoti_c_shim + ) + if backend_index is None and dispatch_key is not None: + return None + + schema = func.func + device = "aten" if dispatch_key is None else dispatch_key.lower() + backend_call = gen_static_dispatch_backend_call( + func, + backend_index, + ) + + try: + if header: + declaration, _ = gen_declaration_and_definition( + schema, device, backend_call, version_info + ) + return declaration + else: + _, definition = gen_declaration_and_definition( + schema, device, backend_call, version_info + ) + return definition + + except NotImplementedError: + return None + + +@dataclass(frozen=True) +class ShimGenerator: + inductor_fallback_ops: dict[str, dict[str, list[str]]] + func_group_mapping: dict[OperatorName, NativeFunctionsGroup] + dispatch_key: DispatchKey | None + backend_indices: dict[DispatchKey, BackendIndex] + header: bool # True to generate .h and False to generate .cpp + extend_aoti_c_shim: bool + + @method_with_native_function + def __call__( + self, + func: NativeFunction, + ) -> str | None: + version_info = self.inductor_fallback_ops[get_fallback_op_name(func)] + result = gen_c_shim( + func, + version_info, + self.func_group_mapping, + self.dispatch_key, + self.backend_indices, + self.header, + self.extend_aoti_c_shim, + ) + return result + + +def gen_aoti_c_shim( + native_functions: Sequence[NativeFunction], + inductor_fallback_ops: dict[str, dict[str, list[str]]], + func_group_mapping: dict[OperatorName, NativeFunctionsGroup], + dispatch_key: DispatchKey | None, + backend_indices: dict[DispatchKey, BackendIndex], + header: bool, + extend_aoti_c_shim: bool, + includes: str = "", +) -> str: + body = "\n".join( + list( + mapMaybe( + ShimGenerator( + inductor_fallback_ops, + func_group_mapping, + dispatch_key, + backend_indices, + header, + extend_aoti_c_shim, + ), + native_functions, + ) + ) + ) + device = "aten" if dispatch_key is None else dispatch_key.lower() + include_device_functions = ( + "#include " + if dispatch_key is None + else f"#include " + ) + aten_warning = ( + ( + "\n\n// This file corresponds to the aten_shimified_ops list in torchgen/aoti/fallback_ops.py\n" + ) + if dispatch_key is None + else "" + ) + warning = """ + +// WARNING: THIS FILE IS AUTOGENERATED BY torchgen. DO NOT MODIFY BY HAND. +// See https://github.com/pytorch/pytorch/blob/7e86a7c0155295539996e0cf422883571126073e/torchgen/gen.py#L2424-L2436 for details""" + + if header: + return ( + warning + + aten_warning + + textwrap.dedent(""" + + #pragma once + + #include + + #ifdef __cplusplus + extern "C" { + #endif + + """) + + body + + textwrap.dedent(""" + + #ifdef __cplusplus + } // extern "C" + #endif + """) + ) + else: + return ( + warning + + aten_warning + + textwrap.dedent(f""" + + #include + #include + + #ifndef AT_PER_OPERATOR_HEADERS + {include_device_functions} + #include + #include + #include + #else + """) + + includes + + textwrap.dedent(""" + #endif // AT_PER_OPERATOR_HEADERS + + using namespace torch::aot_inductor; + + """) + + body + ) + + +def gen_aoti_c_shim_files( + aoti_fm: FileManager, + aoti_backends: set[DispatchKey | None], + native_functions: Sequence[NativeFunction], + backend_indices: dict[DispatchKey, BackendIndex], + structured_native_functions: Sequence[NativeFunctionsGroup], + extra_cuda_headers: str, + extend_aoti_c_shim: bool, + update_aoti_c_shim: bool, +) -> None: + structured_func_group_dict = {} + for func_group in structured_native_functions: + for func in func_group.functions(): + if func.structured_delegate is not None: + structured_func_group_dict[func.structured_delegate] = func_group + break + + for dispatch_key in aoti_backends: + # Use aten_shimified_ops for the aten backend, inductor_fallback_ops for others + fallback_ops_dict = ( + aten_shimified_ops if dispatch_key is None else inductor_fallback_ops + ) + fallbacks = {} + for func in native_functions: + op_name = get_fallback_op_name(func) + if op_name in fallback_ops_dict: + fallbacks[op_name] = func + fallback_native_functions = tuple( + value for _, value in sorted(fallbacks.items()) + ) + + # Use "aten" as the device name when dispatch_key is Generic + device_name = "aten" if dispatch_key is None else dispatch_key.lower() + + # header files were checked in for ABI-compatibility checking + header_file_name = f"c_shim_{device_name}.h" + new_header = gen_aoti_c_shim( + fallback_native_functions, + fallback_ops_dict, + structured_func_group_dict, + dispatch_key, + backend_indices, + header=True, + extend_aoti_c_shim=extend_aoti_c_shim, + includes="", + ) + if update_aoti_c_shim: + aoti_fm.write( + header_file_name, + lambda: new_header, + ) + else: + try: + with open( + os.path.join(aoti_fm.install_dir, header_file_name) + ) as old_file: + old_header = old_file.read() + + if old_header != new_header: + diff = "\n".join( + difflib.unified_diff( + old_header.splitlines(), + new_header.splitlines(), + fromfile="expected", + tofile="actual", + lineterm="", + ) + ) + + raise RuntimeError(f""" +The generated AOTInductor C shim header files have unexpectedly changed. This +indicates an AOTInductor fallback operator ABI backward compatibility breakage!!! +Only in a limited number of situations, this is allowed: + +1. You added a fallback op to the inductor_fallback_ops list in torchgen/aoti/fallback_ops.py. +If that's the case, run `python torchgen/gen.py --update-aoti-c-shim` to add a new entry to +existing C shim header files. + +2. You added a new default argument to an existing fallback op. This is clearly a BC breaking +change in the AOTInductor land. You need to annotate the new default argument in +torchgen/aoti/fallback_ops.py, and then run `python torchgen/gen.py --update-aoti-c-shim` to +update the C shim header files by creating different versions of the fallback op. See +https://github.com/pytorch/pytorch/pull/154848 as an example. + +{diff} + """) + except FileNotFoundError: + print( + f"{os.path.join(aoti_fm.install_dir, header_file_name)} not found" + ) + + # cpp files are always generated on-the-fly + def headers_for_aoti() -> str: + headers = [] + for func in fallback_native_functions: + header = get_header_for_aoti( + func, + structured_func_group_dict, + dispatch_key, + backend_indices, + extend_aoti_c_shim=extend_aoti_c_shim, + ) + if header is not None: + headers.append(header) + return "\n".join(sorted(set(headers))) + + extra_headers = ( + extra_cuda_headers + if dispatch_key is not None and is_cuda_dispatch_key(dispatch_key) + else "" + ) + + aoti_fm.write( + f"c_shim_{device_name}.cpp", + lambda: gen_aoti_c_shim( + fallback_native_functions, + fallback_ops_dict, + structured_func_group_dict, + dispatch_key, + backend_indices, + header=False, + extend_aoti_c_shim=extend_aoti_c_shim, + includes=headers_for_aoti() + "\n" + extra_headers, + ), + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_backend_stubs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_backend_stubs.py new file mode 100644 index 0000000000000000000000000000000000000000..efe63a80249eb4226001a686c65e59e47740c1aa --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_backend_stubs.py @@ -0,0 +1,635 @@ +from __future__ import annotations + +import argparse +import os +import re +from collections import Counter, defaultdict, namedtuple +from pathlib import Path +from typing import TYPE_CHECKING + +import yaml + +import torchgen.api.dispatcher as dispatcher +import torchgen.dest as dest +from torchgen.api.types import DispatcherSignature +from torchgen.code_template import CodeTemplate +from torchgen.context import native_function_manager +from torchgen.gen import get_grouped_native_functions, parse_native_yaml +from torchgen.model import ( + BackendIndex, + BackendMetadata, + DispatchKey, + NativeFunction, + NativeFunctionsGroup, + OperatorName, +) +from torchgen.selective_build.selector import SelectiveBuilder +from torchgen.utils import concatMap, context, FileManager, NamespaceHelper, Target +from torchgen.yaml_utils import YamlLoader + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# Parses the external backend's yaml, and adds a new BackendIndex for the backend's dispatch key. +# Returns a Tuple of (backend_key, autograd_key, cpp_namespace, updated BackendIndex mapping) +ParsedExternalYaml = namedtuple( + "ParsedExternalYaml", + ["backend_key", "autograd_key", "class_name", "cpp_namespace", "backend_indices"], +) + + +def parse_backend_yaml( + backend_yaml_path: str, + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + backend_indices: dict[DispatchKey, BackendIndex], +) -> ParsedExternalYaml: + native_functions_map: dict[OperatorName, NativeFunction] = { + f.func.name: f + for f in concatMap( + lambda f: [f] if isinstance(f, NativeFunction) else list(f.functions()), + grouped_native_functions, + ) + } + + with open(backend_yaml_path) as f: + yaml_values = yaml.load(f, Loader=YamlLoader) + if not isinstance(yaml_values, dict): + raise AssertionError( + f"Expected yaml_values to be a dict, got {type(yaml_values)}" + ) + + valid_keys = [ + "backend", + "class_name", + "cpp_namespace", + "extra_headers", + "supported", + "autograd", + "full_codegen", + "non_native", + "ir_gen", + "symint", + ] + + backend = yaml_values.pop("backend", None) + if backend is None: + raise AssertionError('You must provide a value for "backend"') + + class_name = yaml_values.pop("class_name", None) + + cpp_namespace = yaml_values.pop("cpp_namespace", None) + if cpp_namespace is None: + raise AssertionError('You must provide a value for "cpp_namespace"') + + # Mostly just defaulting to false to stick with LazyTensor convention. + use_out_as_primary = yaml_values.pop("use_out_as_primary", False) + if not isinstance(use_out_as_primary, bool): + raise AssertionError( + f"You must provide either True or False for use_out_as_primary. Provided: {use_out_as_primary}" + ) + + use_device_guard = yaml_values.pop("device_guard", False) + if not isinstance(use_device_guard, bool): + raise AssertionError( + f"You must provide either True or False for device_guard. Provided: {use_device_guard}" + ) + + supported = yaml_values.pop("supported", []) + if supported is None: + supported = [] # Allow an empty list of supported ops + if not isinstance(supported, list): + raise AssertionError( + f'expected "supported" to be a list, but got: {supported} (of type {type(supported)})' + ) + + symint = yaml_values.pop("symint", []) + if symint is None: + symint = [] # Allow an empty list of symint ops + if not isinstance(symint, list): + raise AssertionError( + f'expected "symint" to be a list, but got: {symint} (of type {type(symint)})' + ) + symint_set = set(symint) + + supported_autograd = yaml_values.pop("autograd", []) + if not isinstance(supported_autograd, list): + raise AssertionError( + f'expected "autograd" to be a list, but got: {supported_autograd}' + ) + + # full_codegen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py + full_codegen = yaml_values.pop("full_codegen", []) + supported.extend(full_codegen) + + # non_native is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py + yaml_values.pop("non_native", {}) + + # ir_gen is ignored by parse_backend_yaml, and re-parsed in gen_lazy_tensor.py + yaml_values.pop("ir_gen", {}) + + if len(yaml_values.keys()) != 0: + raise AssertionError( + f"{backend_yaml_path} contains unexpected keys: {', '.join(yaml_values.keys())}. " + f"Only the following keys are supported: {', '.join(valid_keys)}" + ) + + def create_backend_index( + backend_ops: list[str], + symint_ops: set[str], + dispatch_key: DispatchKey, + *, + use_out_as_primary: bool, + use_device_guard: bool, + ) -> BackendIndex: + metadata: dict[OperatorName, BackendMetadata] = {} + for op in backend_ops: + op_name = OperatorName.parse(op) + if op_name not in native_functions_map: + raise AssertionError(f"Found an invalid operator name: {op_name}") + # See Note [External Backends Follow Dispatcher API] + kernel_name = dispatcher.name(native_functions_map[op_name].func) + if op in symint_ops: + kernel_name += "_symint" + # TODO: allow structured external backends later. + m = BackendMetadata( + kernel=kernel_name, structured=False, cpp_namespace=cpp_namespace + ) + metadata[op_name] = m + return BackendIndex( + dispatch_key=dispatch_key, + use_out_as_primary=use_out_as_primary, + external=True, + device_guard=use_device_guard, + index=metadata, + ) + + backend_key: DispatchKey | None = None + if len(supported) > 0: + with context( + lambda: f'The provided value for "backend" must be a valid DispatchKey, but got {backend}.' + ): + backend_key = DispatchKey.parse(backend) + + backend_idx = create_backend_index( + supported, + symint_set, + backend_key, + use_out_as_primary=use_out_as_primary, + use_device_guard=use_device_guard, + ) + if backend_key in backend_indices: + raise AssertionError(f"Duplicate backend key: {backend_key}") + backend_indices[backend_key] = backend_idx + + autograd_key: DispatchKey | None = None + if len(supported_autograd) > 0: + with context( + lambda: f'The "autograd" key was specified, which indicates that you would like to override \ +the behavior of autograd for some operators on your backend. However "Autograd{backend}" is not a valid DispatchKey.' + ): + autograd_key = DispatchKey.parse(f"Autograd{backend}") + + autograd_idx = create_backend_index( + supported_autograd, + symint_set, + autograd_key, + use_out_as_primary=use_out_as_primary, + use_device_guard=use_device_guard, + ) + if autograd_key in backend_indices: + raise AssertionError(f"Duplicate autograd key: {autograd_key}") + backend_indices[autograd_key] = autograd_idx + + for g in grouped_native_functions: + if isinstance(g, NativeFunction): + forward_kernels = ( + [] + if backend_key is None + else [ + m + for m in [backend_indices[backend_key].get_kernel(g)] + if m is not None + ] + ) + backward_kernels = ( + [] + if autograd_key is None + else [ + m + for m in [backend_indices[autograd_key].get_kernel(g)] + if m is not None + ] + ) + else: + forward_kernels = ( + [] + if backend_key is None + else [ + m + for m in [ + backend_indices[backend_key].get_kernel(f) + for f in g.functions() + ] + if m is not None + ] + ) + backward_kernels = ( + [] + if autograd_key is None + else [ + m + for m in [ + backend_indices[autograd_key].get_kernel(f) + for f in g.functions() + ] + if m is not None + ] + ) + + forward_kernels = [f for f in forward_kernels if f is not None] + backward_kernels = [f for f in backward_kernels if f is not None] + if not (len(forward_kernels) == 0 or len(backward_kernels) == 0): + raise AssertionError( + f"Currently, all variants of an op must either be registered to a backend key, " + f"or to a backend's autograd key. They cannot be mix and matched. " + f"If this is something you need, feel free to create an issue! " + f'{forward_kernels[0].kernel} is listed under "supported", ' + f'but {backward_kernels[0].kernel} is listed under "autograd".' + ) + + return ParsedExternalYaml( + backend_key, autograd_key, class_name, cpp_namespace, backend_indices + ) + + +def error_on_missing_kernels( + native_functions: Sequence[NativeFunction], + backend_indices: dict[DispatchKey, BackendIndex], + backend_key: DispatchKey, + autograd_key: DispatchKey | None, + class_name: str, + kernel_defn_file_path: str, + full_codegen: list[OperatorName] | None = None, +) -> None: + try: + with open(kernel_defn_file_path) as f: + backend_defns = f.read() + except OSError as e: + raise AssertionError( + f"Unable to read from the specified impl_path file: {kernel_defn_file_path}" + ) from e + + if full_codegen is None: + full_codegen = [] + + indices = [backend_indices[backend_key].index] + ( + [] if autograd_key is None else [backend_indices[autograd_key].index] + ) + # Quick mapping from each OperatorName used by the external backend + # to its backend kernel name + expected_backend_op_names: dict[OperatorName, str] = dict( + list( + concatMap( + lambda index: [ + (op_name, metadata.kernel) for op_name, metadata in index.items() + ], + indices, + ) + ) + ) + expected_backend_native_funcs: list[NativeFunction] = [ + f + for f in native_functions + if f.func.name in expected_backend_op_names and f.func.name not in full_codegen + ] + expected_backend_kernel_name_counts: dict[str, list[NativeFunction]] = defaultdict( + list + ) + for native_f in expected_backend_native_funcs: + expected_backend_kernel_name_counts[ + expected_backend_op_names[native_f.func.name] + ].append(native_f) + + # This just looks for lines containing "foo(", and assumes that the kernel foo has been implemented. + # It might cause false negatives (we won't catch all cases), but that's ok - if we catch a missing kernel + # here, then we get a nicer error message. If we miss it, you get a linker error. + kernel_defn_regex = rf"(.*){class_name}::\s*([\w\d]*)\(" + actual_backend_kernel_name_counts = Counter( + # A bit unwieldy (this could probably be moved into regex), + # but we don't want to include kernel names that come from function calls, + # like "return torch_xla::XLANativeFunctions::empty_strided_symint(...)". + # Easy check is to ignore any lines with colons before the class name. + [ + y + for (x, y) in re.findall(kernel_defn_regex, backend_defns) + if not x.endswith(":") + ] + ) + + missing_kernels_err_msg = "" + for expected_name, funcs in expected_backend_kernel_name_counts.items(): + expected_overload_count = len(funcs) + actual_overload_count = actual_backend_kernel_name_counts[expected_name] + if expected_overload_count != actual_overload_count: + + def create_decl(f: NativeFunction) -> str: + with native_function_manager(f): + return DispatcherSignature.from_schema(f.func).decl() + + expected_schemas_str = "\n".join([create_decl(f) for f in funcs]) + missing_kernels_err_msg += f""" +{class_name} is missing a kernel definition for {expected_name}. We found {actual_overload_count} kernel(s) with that name, +but expected {expected_overload_count} kernel(s). The expected function schemas for the missing operator are: +{expected_schemas_str} + +""" + if missing_kernels_err_msg != "": + raise AssertionError(missing_kernels_err_msg) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate backend stub files") + parser.add_argument( + "-s", + "--source-yaml", + "--source_yaml", + help="path to source yaml file containing operator external definitions", + ) + parser.add_argument("-o", "--output-dir", "--output_dir", help="output directory") + parser.add_argument( + "--dry-run", "--dry_run", type=bool, default=False, help="output directory" + ) + parser.add_argument( + "--impl-path", + "--impl_path", + type=str, + default=None, + help="path to the source C++ file containing kernel definitions", + ) + options = parser.parse_args() + + run(options.source_yaml, options.output_dir, options.dry_run, options.impl_path) + + +def gen_dispatchkey_nativefunc_headers( + fm: FileManager, + class_name: str, + cpp_namespace: str, + backend_indices: dict[DispatchKey, BackendIndex], + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + backend_dispatch_key: DispatchKey, + autograd_dispatch_key: DispatchKey | None, + backend_name: str = "", +) -> None: + if class_name is None: + raise AssertionError("class_name must not be None") + generated_comment = ( + "Autogenerated file by gen_backend_stubs.py. Do not edit directly!" + ) + + # Convert to a set first to remove duplicate kernel names. + # Backends are allowed to repeat kernel names; only generate the declaration once! + # Sort for deterministic output. + backend_declarations = sorted( + set( + concatMap( + lambda f: dest.compute_native_function_declaration( + f, backend_indices[backend_dispatch_key] + ), + grouped_native_functions, + ) + ) + ) + autograd_declarations = sorted( + set( + concatMap( + lambda f: [] + if autograd_dispatch_key is None + else dest.compute_native_function_declaration( + f, backend_indices[autograd_dispatch_key] + ), + grouped_native_functions, + ) + ) + ) + + ns_helper = NamespaceHelper(cpp_namespace) + fm.write_with_template( + f"{backend_dispatch_key}NativeFunctions.h", + "DispatchKeyNativeFunctions.h", + lambda: { + "generated_comment": generated_comment, + "namespace_prologue": ns_helper.prologue, + "class_name": class_name, + "namespace_epilogue": ns_helper.epilogue, + "dispatch_declarations": backend_declarations + autograd_declarations, + "BackendName": backend_name, + "DispatchKey": backend_dispatch_key, + }, + ) + + +def gen_dispatcher_registrations( + fm: FileManager, + output_dir: str, + class_name: str, + backend_indices: dict[DispatchKey, BackendIndex], + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], + backend_dispatch_key: DispatchKey, + dispatch_key: DispatchKey, + selector: SelectiveBuilder, + # build_in_tree is true for lazy TS backend and affects include paths, not used for external backends + build_in_tree: bool = False, + per_operator_headers: bool = False, + backend_name: str = "", + eager_registration: bool = True, +) -> None: + headers = [ + f"{output_dir}/{backend_dispatch_key}NativeFunctions.h", + ] + if build_in_tree: + external_backend_headers_str = "\n".join(f"#include <{h}>" for h in headers) + else: + external_backend_headers_str = "\n".join(f'#include "{h}"' for h in headers) + + if class_name is None: + raise AssertionError("class_name must not be None") + backend_index = backend_indices[dispatch_key] + + dispatch_registrations_body = list( + concatMap( + dest.RegisterDispatchKey( + backend_index, + Target.REGISTRATION, + selector, + rocm=False, + symint=True, + class_method_name=f"{class_name}", + skip_dispatcher_op_registration=False, + ), + grouped_native_functions, + ) + ) + newline = "\n" + ns_helper = NamespaceHelper(namespace_str="at") + deferred_dispatch_registrations = "" + static_init_dispatch_registrations = "" + if eager_registration: + static_template = CodeTemplate( + """\ +TORCH_LIBRARY_IMPL(aten, $dispatch_key, m) { + $dispatch_registrations_body +}""" + ) + static_init_dispatch_registrations = static_template.substitute( + dispatch_key=dispatch_key, + dispatch_registrations_body=dispatch_registrations_body, + ) + else: + deferred_template = CodeTemplate( + """\ +TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions(); +TORCH_API void Register${backend_name}${dispatch_key}NativeFunctions() { + static auto m = MAKE_TORCH_LIBRARY_IMPL(aten, $dispatch_key); + $dispatch_registrations_body +}""" + ) + deferred_dispatch_registrations = deferred_template.substitute( + backend_name=backend_name, + dispatch_key=dispatch_key, + dispatch_registrations_body=dispatch_registrations_body, + ) + + fm.write_with_template( + f"Register{dispatch_key}.cpp", + "RegisterDispatchKey.cpp", + lambda: { + "extra_cuda_headers": "", + "external_backend_headers": external_backend_headers_str, + "ops_headers": "#include " + if not per_operator_headers + else "", + "DispatchKey": dispatch_key, + "dispatch_namespace": dispatch_key.lower(), + "dispatch_headers": dest.gen_registration_headers( + backend_index, per_operator_headers=per_operator_headers, rocm=False + ), + "dispatch_helpers": dest.gen_registration_helpers(backend_index), + "dispatch_definitions": fm.substitute_with_template( + "RegisterDispatchDefinitions.ini", + lambda: { + "ns_prologue": ns_helper.prologue, + "ns_epilogue": ns_helper.epilogue, + "static_init_dispatch_registrations": static_init_dispatch_registrations, + "deferred_dispatch_registrations": deferred_dispatch_registrations, + "dispatch_namespace": dispatch_key.lower(), + "dispatch_namespaced_definitions": "", + "dispatch_anonymous_definitions": list( + concatMap( + dest.RegisterDispatchKey( + backend_index, + Target.ANONYMOUS_DEFINITION, + selector, + rocm=False, + symint=True, + class_method_name=f"{class_name}", + skip_dispatcher_op_registration=False, + ), + grouped_native_functions, + ) + ), + }, + ).split(newline), + }, + ) + + +def run( + source_yaml: str, output_dir: str, dry_run: bool, impl_path: str | None = None +) -> None: + # Assumes that this file lives at torchgen/gen_backend_stubs.py + root = Path(__file__).absolute().parent.parent + common_dir = os.path.join(root, "aten/src") # Assumes root is pytorch_root + if not os.path.exists(common_dir): # This file is out-of-tree. + common_dir = os.path.join(root, "torchgen/packaged") + + template_dir = os.path.join(common_dir, "ATen/templates") + + def make_file_manager(install_dir: str) -> FileManager: + return FileManager( + install_dir=install_dir, template_dir=template_dir, dry_run=dry_run + ) + + fm = make_file_manager(output_dir) + + native_yaml_path = os.path.join(common_dir, "ATen/native/native_functions.yaml") + tags_yaml_path = os.path.join(common_dir, "ATen/native/tags.yaml") + parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path) + native_functions, backend_indices = ( + parsed_yaml.native_functions, + parsed_yaml.backend_indices, + ) + grouped_native_functions = get_grouped_native_functions(native_functions) + parsed_backend_yaml = parse_backend_yaml( + source_yaml, grouped_native_functions, backend_indices + ) + backend_key = parsed_backend_yaml.backend_key + autograd_key = parsed_backend_yaml.autograd_key + cpp_namespace = parsed_backend_yaml.cpp_namespace + class_name = parsed_backend_yaml.class_name + backend_indices = parsed_backend_yaml.backend_indices + + selector = SelectiveBuilder.get_nop_selector() + + if backend_key is None: + # This could be useful if a backend wants to quickly set up a noop yaml file but doesn't have any kernels ready yet. + return + + if class_name is None: + # class_name is an optional argument to backend yaml file. + # if specified it allows an external backend to override + # the name of the class that all generated kernel definitions live under. + # if not specified, its value is given as native_function_class_name. + class_name = backend_indices[backend_key].native_function_class_name() + if class_name is None: + raise AssertionError("class_name must not be None") + + if impl_path is not None: + error_on_missing_kernels( + native_functions, + backend_indices, + backend_key, + autograd_key, + class_name, + impl_path, + ) + + gen_dispatchkey_nativefunc_headers( + fm, + class_name, + cpp_namespace, + backend_indices, + grouped_native_functions, + backend_key, + autograd_key, + ) + + for dispatch_key in ( + [backend_key] if autograd_key is None else [backend_key, autograd_key] + ): + gen_dispatcher_registrations( + fm, + output_dir, + class_name, + backend_indices, + grouped_native_functions, + backend_key, + dispatch_key, + selector, + ) + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_functionalization_type.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_functionalization_type.py new file mode 100644 index 0000000000000000000000000000000000000000..251ba64248a3c2120a286f8553528518b472aa1b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_functionalization_type.py @@ -0,0 +1,1227 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from torchgen.api import cpp, dispatcher, functionalization +from torchgen.api.translate import translate +from torchgen.api.types import ( + BaseCType, + Binding, + CType, + DispatcherSignature, + iTensorListRefT, + NativeSignature, + OptionalCType, + optionalSymIntArrayRefT, + symIntArrayRefT, + SymIntT, + tensorListT, + tensorT, + VectorCType, + ViewInverseSignature, +) +from torchgen.context import ( + method_with_native_function, + native_function_manager, + with_native_function, + with_native_function_and, +) +from torchgen.model import ( + Argument, + BackendIndex, + BaseTy, + BaseType, + FunctionSchema, + ListType, + NativeFunction, + NativeFunctionsGroup, + NativeFunctionsViewGroup, + OperatorName, + Return, + SchemaKind, + SelfArgument, + TensorOptionsArguments, +) +from torchgen.native_function_generation import ( + INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY, + MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT, + OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY, +) +from torchgen.utils import concatMap, dataclass_repr, FileManager + + +if TYPE_CHECKING: + from collections.abc import Callable + + from torchgen.selective_build.selector import SelectiveBuilder + + +# Note: [Mutable Ops Not Using Functionalization] +# Ops in this list currently do not work with functionalization and should be fixed. +MUTABLE_OPS_NOT_USING_FUNCTIONALIZATION = ( + OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY + + MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT + + INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY + + [ + # It will be BC-breaking, but we should fix their schemas. + # should be inplace? + "record_stream", + # See Note [resize_ in Functionalization] + "resize_", + "resize_as_", + # This function is used as for testing purposes only. + "_fill_mem_eff_dropout_mask_", + # Inference-only op called behind a custom op graph break. + "_flash_attention_forward_no_dropout_inplace", + ] +) + +# Eager cumulative out variants compute in the out dtype when dtype is omitted. +# Functionalization normally lowers mutable ops through their functional variants, +# so these need to thread the out dtype explicitly to preserve eager semantics. +CUMULATIVE_OUT_OPS_PRESERVING_OUT_DTYPE = { + OperatorName.parse("cumsum.out"), + OperatorName.parse("cumprod.out"), + OperatorName.parse("cumsum.dimname_out"), + OperatorName.parse("cumprod.dimname_out"), +} + +# This file contains codegen that relates to the functionalization pass. +# It includes: +# - gen_functionalization_definition +# Generates dispatcher kernel definitions for the functionalization pass. +# - gen_functionalization_registration +# Generates dispatcher kernel registrations for the functionalization pass. +# - gen_functionalization_view_inverse_declaration +# Generates a declaration for an "inverse view", for every view op +# that is needed in functionalization. We manually implement their definitions. +# - gen_composite_view_copy_kernel +# Generates view_copy() composite kernels for all view_copy operators. + + +# Generates the body of the default composite C++ kernel for a {view}_copy NativeFunction +# See Note [view_copy NativeFunctions] +@dataclass(frozen=True) +class GenCompositeViewCopyKernel: + backend_index: BackendIndex + + @method_with_native_function + def __call__(self, g: NativeFunctionsViewGroup) -> str | None: + if g.view_copy is None: + return None + elif g.view_copy.func.name.name.base != f"{g.view.func.name.name}_copy": + # If the view_copy doesn't match the standard naming scheme of _copy, + # assume it already exists and doesn't need to be generated. + # Example: slice_inverse() with the copy variant named slice_scatter() + # instead of slice_inverse_copy() + return None + + metadata = self.backend_index.get_kernel(g.view_copy) + if metadata is None: + raise AssertionError( + f"Expected metadata for view_copy kernel: {g.view_copy}" + ) + + # We can make view_copy work in more cases by using reshape() + # when a normal view call would ordinarily fail. + # This also makes LTC more efficient, because they don't need to include + # clone() calls in their graph (which is normally needed by reshape). + if str(g.view_copy.func.name) == "view_copy": + if metadata.kernel != "view_copy_symint": + raise AssertionError( + f"Expected kernel 'view_copy_symint', got '{metadata.kernel}'" + ) + return """\ +at::Tensor view_copy_symint(const at::Tensor & self, at::SymIntArrayRef size) { + c10::SymDimVector shape = infer_size_dv(size, self.sym_numel()); + if (!at::detail::computeStride(self.sym_sizes(), self.sym_strides(), shape).has_value()) { + return self.reshape_symint(size); + } else { + auto output = at::_ops::view::call(self, size); + return output.clone(/*memory_format=*/at::MemoryFormat::Contiguous); + } +} +""" + # view_copy is a native signature, since we're generating an at::native:: kernel + # Functionalization always operates on symints though + view_copy_sig = NativeSignature( + g.view_copy.func, symint=metadata.supports_symint() + ) + + # view is a dispatcher signature, since we're calling into the at::_ops API + view_sig = DispatcherSignature(g.view.func) + + view_api_name = g.view.func.name.unambiguous_name() + exprs = ", ".join( + [e.expr for e in translate(view_copy_sig.arguments(), view_sig.arguments())] + ) + + # view ops today always return either a Tensor or a list of Tensors + if len(g.view.func.returns) != 1: + raise AssertionError(f"Expected 1 return, got {len(g.view.func.returns)}") + if not ( + g.view.func.returns[0].type == BaseType(BaseTy.Tensor) + or g.view.func.returns[0].type == ListType(BaseType(BaseTy.Tensor), None) + ): + raise AssertionError( + f"Expected Tensor or Tensor[] return type, got {g.view.func.returns[0].type}" + ) + + if g.view.func.returns[0].type == BaseType(BaseTy.Tensor): + return_cloned_output = """\ + return output.clone(/*memory_format=*/at::MemoryFormat::Contiguous);""" + else: + # If the return type is a list, we need to clone each tensor in the list. + return_cloned_output = f"""\ + {view_copy_sig.returns_type().cpp_type()} out_clone; + for (const auto i : c10::irange(output.size())) {{ + out_clone.push_back(output[i].clone(/*memory_format=*/at::MemoryFormat::Contiguous)); + }} + return out_clone;""" + + # The default generated composite kernel for {view}_copy() operators just clones + # the input tensor, and runs the underlying view on the clone. + return f""" +{view_copy_sig.defn(name=metadata.kernel)} {{ + auto output = at::_ops::{view_api_name}::call({exprs}); + {return_cloned_output} +}} +""" + + +def return_str(rets: tuple[Return, ...], names: list[str]) -> str: + if len(rets) != len(names): + raise AssertionError(f"Expected {len(rets)} names, got {len(names)}") + if len(rets) == 0: + return "" + elif len(rets) == 1: + return f"return {names[0]};" + else: + return f"return {dispatcher.returns_type(rets).cpp_type()}({', '.join(names)});" + + +def modifies_arguments(f: NativeFunction) -> bool: + return any( + a.annotation is not None and a.annotation.is_write + for a in f.func.arguments.flat_all + ) + + +def wrapper_name(func: FunctionSchema) -> str: + if func.name.overload_name: + return f"{cpp.name(func)}_{func.name.overload_name}" + else: + return cpp.name(func) + + +def is_tensor_like(a: Argument | TensorOptionsArguments | SelfArgument) -> bool: + return isinstance(a, SelfArgument) or ( + isinstance(a, Argument) and a.type.is_tensor_like() + ) + + +# We need to wrap / unwrap various arguments from the op in the functionalization kernels. +# Some op schemas include non-owning types though (like TensorList), +# and when we unwrap them we expect to get out an owning type!. +# We also return a lambda that tells you how to convert the non-owning type argument into the owning type. +def get_owning_type(t: CType) -> tuple[CType, Callable[[str], str]]: + if t == BaseCType(tensorListT): + return VectorCType(BaseCType(tensorT)), lambda x: f"{x}.vec()" + if t == BaseCType(iTensorListRefT): + return VectorCType(BaseCType(tensorT)), lambda x: f"{{{x}.begin(), {x}.end()}}" + # There are technically other non-owning types out there (like IntArrayRef), + # but functionalization only actually cares about the ones involving tensors. + return t, lambda x: x + + +# unwraps all tensor-like arguments, returning: +# (1) a string containing all of the logic that does the unwrapping +# (2) a context, to be used by translate(), with all of the relevant bindings. +def unwrap_tensor_args( + sig: DispatcherSignature, *, is_view_op: bool +) -> tuple[str, list[Binding]]: + context: list[Binding] = [] + unwrapped_tensor_args: list[str] = [] + for arg in sig.arguments(): + if is_tensor_like(arg.argument): + # for tensor inputs, we want to unwrap them before passing them into the redispatch calls. + unwrapped_name = f"{arg.name}_" + # For most ops, the functionalization needs to sync any pending updates on the input tensors + # before calling the operator, since otherwise the operator will act on stale data. + # For view ops though, we can continue to defer syncing until the tensor is used by + # a non-view operator. + maybe_sync_input = ( + "" if is_view_op else f"at::functionalization::impl::sync({arg.name});" + ) + unwrapped_type, conversion_fn = get_owning_type( + arg.nctype.remove_const_ref().type + ) + unwrapped_tensor_args.append( + f""" + {unwrapped_type.cpp_type()} {unwrapped_name}; + if (at::functionalization::impl::isFunctionalTensor({arg.name})) {{ + {maybe_sync_input} + {unwrapped_name} = at::functionalization::impl::from_functional_tensor({arg.name}); + }} else {{ + {unwrapped_name} = {conversion_fn(arg.name)}; + }}""" + ) + context.append(arg.with_name(unwrapped_name)) + else: + # for non-tensor inputs, we want to pass them directly into the redispatch calls. + context.append(arg) + unwrap_tensor_args_str = "\n ".join(unwrapped_tensor_args) + return unwrap_tensor_args_str, context + + +# converts all tensor-like arguments to meta tensors, which are used to compute stride info. Returns: +# (1) a string containing all of the logic that does the conversions. +# (2) a context, to be used by translate(), with all of the relevant bindings. +def convert_to_meta_tensors(sig: DispatcherSignature) -> tuple[str, list[Binding]]: + context: list[Binding] = [] + unwrapped_tensor_args: list[str] = [] + for arg in sig.arguments(): + if is_tensor_like(arg.argument): + # for tensor inputs, we want to unwrap them before passing them into the redispatch calls. + a_ = arg.name + unwrapped_name = f"{arg.name}_meta" + unwrapped_tensor_args.append(f"auto {unwrapped_name} = to_meta({a_});") + context.append(arg.with_name(unwrapped_name)) + else: + # for non-tensor inputs, we want to pass them directly into the redispatch calls. + context.append(arg) + unwrap_tensor_args_str = "\n ".join(unwrapped_tensor_args) + return unwrap_tensor_args_str, context + + +# The functionalization codegen currently expects view op schemas to have this form: +# foo(Tensor(a), ...) -> Tensor(a) (e.g. transpose) +# foo(Tensor(a!), ...) -> Tensor(a!) (e.g. transpose_) +def assert_view_op_properties(func: FunctionSchema) -> None: + def is_alias(a: Argument) -> bool: + return a.annotation is not None + + args = func.arguments.flat_non_out + # The first argument is a tensor with an alias semantics (annotations) + if not (len(args) > 0 and args[0].type == BaseType(BaseTy.Tensor)): + raise AssertionError( + f"In the functionalization codegen, we expect the first argument of every view operator to be a tensor, " + f"but found an argument of type {str(args[0].type)} for operator: {str(func.name)}." + ) + # No other arguments have aliasing semantics + if not (is_alias(args[0]) and not any(is_alias(a) for a in args[1:])): + raise AssertionError( + "In the functionalization codegen, we expect the first argument of every view " + "operator to alias the output. View operators with multiple aliasing inputs " + "aren't supported yet. Found an operator that doesn't satisfy this constraint" + ) + + +# One-liner expression for checking if an expression expr of type type has any +# symbolic values. +def emit_expr_has_symbolic_values(expr: str, type: CType) -> str: + if type == BaseCType(SymIntT): + return f"{expr}.is_symbolic()" + + if isinstance(type, OptionalCType): + innerexpr = f"(*{expr})" + return f"{expr}.has_value() ? {emit_expr_has_symbolic_values(innerexpr, type.elem)} : false" + + if type == BaseCType(optionalSymIntArrayRefT): + return emit_expr_has_symbolic_values( + expr, OptionalCType(BaseCType(symIntArrayRefT)) + ) + + if type in (BaseCType(symIntArrayRefT), VectorCType(BaseCType(SymIntT))): + argname = "arg" + lambda_check = emit_expr_has_symbolic_values(argname, BaseCType(SymIntT)) + return ( + "std::any_of(" + f"{expr}.begin(), {expr}.end(), " + f"[=](auto& {argname}) {{ return {lambda_check}; }})" + ) + + raise ValueError( + "unsupported type for has_symbolic_values check. " + "It should be a SymInt or a collection of those. " + f"Got: {type.cpp_type()}" + ) + + +# Detects whether any of the SymInt arguments are, in fact, symbolic values. +# This is used in the constructor of ViewMeta. +def emit_has_symbolic_inputs(sig: DispatcherSignature) -> tuple[str, str]: + name = "has_symbolic_inputs" + statements = [ + f"{name} = {name} | ({emit_expr_has_symbolic_values(binding.name, binding.nctype.type)});" + for binding in sig.arguments() + if ( + isinstance(binding.argument, Argument) + and binding.argument.type.is_symint_like() + ) + ] + body = "\n ".join(statements) + return ( + name, + f""" + bool {name} = false; + {body}""", + ) + + +# Generates the Functionalization kernel for: +# - ops that create aliases (e.g. transpose()) +# - ops that are views AND mutations (e.g. transpose_()) +def emit_view_functionalization_body( + g: NativeFunctionsViewGroup, *, view_inplace: bool +) -> str: + if view_inplace: + # This op is both an inplace op AND a view op. + # See Note [Functionalization Pass - Inplace View Ops] for details. + # I currently have the view meta call into the out-of-place variant of the view, to avoid + # having to define an extra ~20 inplace {view}_inverse_ functions. + # Most view ops don't have NativeFunctionGroup's both, because we don't define out= variants for view ops. + # I'm assuming that every inplace-view op has a corresponding out-of-place view op, + # with the same name but the trailing underscore removed. + # This is currently asserted at parse time in gen.py (see error_check_native_functions). + if g.view_inplace is None: + raise AssertionError( + "Expected view_inplace to be non-None for inplace view" + ) + f = g.view_inplace + else: + f = g.view + + if g.view_copy is None: + raise AssertionError("Expected view_copy to be non-None") + with native_function_manager(f): + call_sig = DispatcherSignature.from_schema(g.view_copy.func) + + spec = ViewMetaSpecialization(g, f=f) + + # the "view_copy" op name that the functionalization kernels need to call + api_name = g.view_copy.func.name.unambiguous_name() + # Sometimes the functionalization pass needs to no-op (e.g. if it was passed non-functional tensors) + # "no-op"ing in this context is just redispatching to the original op. + noop_api_name = f.func.name.unambiguous_name() + + dispatcher_sig = DispatcherSignature.from_schema(f.func) + assert_view_op_properties(f.func) + view_tensor_name = dispatcher_sig.arguments()[0].name + + return_type = dispatcher_sig.returns_type().remove_const_ref().cpp_type() + + unwrap_tensor_args_str, unwrapped_args_ctx = unwrap_tensor_args( + dispatcher_sig, is_view_op=True + ) + view_redispatch_args = [ + e.expr + for e in translate(unwrapped_args_ctx, call_sig.arguments(), method=False) + ] + + # The meta API call should use the same arguments, but convert all tensors to meta tensors first. + meta_conversion_str, meta_call_ctx = convert_to_meta_tensors(dispatcher_sig) + meta_call_args = [ + e.expr for e in translate(meta_call_ctx, call_sig.arguments(), method=False) + ] + + ( + symbolic_inputs_varname, + symbolic_inputs_check, + ) = emit_has_symbolic_inputs(call_sig) + + if "inplace_view" in f.tags: + # See Note [Functionalization Pass - Inplace View Ops] for more details + return f""" + {dispatcher_sig.defn(name=wrapper_name(f.func), is_redispatching_fn=True)} {{ + if (!at::functionalization::impl::isFunctionalTensor({view_tensor_name})) {{ + // functionalization is re-entrant, but will no-op if it wasn't passed a FunctionalTensorWrapper. + {unwrap_tensor_args_str} + at::AutoDispatchSkipFunctionalize guard; + return at::_ops::{noop_api_name}::call({", ".join(view_redispatch_args)}); + }} + auto reapply_views = at::functionalization::impl::getFunctionalizationReapplyViewsTLS(); + auto inverse_return_mode = ( + reapply_views ? at::functionalization::InverseReturnMode::ViewOrScatterInverse + : at::functionalization::InverseReturnMode::NeverView + ); + {symbolic_inputs_check} + auto view_meta = {spec.new()}; + auto compute_reference_meta = + {view_tensor_name}.key_set().has_backend(c10::BackendComponent::XLABit) || + {view_tensor_name}.key_set().has_backend(c10::BackendComponent::LazyBit); + {return_type} reference_tensor_output; + if (compute_reference_meta && !disable_meta_reference()) {{ + {meta_conversion_str} + at::AutoDispatchSkipFunctionalize func_guard; + c10::impl::ExcludeDispatchKeyGuard guard(exclude_keys_for_meta_dispatch); + reference_tensor_output = at::_ops::{noop_api_name}::call({", ".join(meta_call_args)}); + }} + // This function adds the above view meta to the current tensor and replays them off the base, + // mutating the size/stride info of the current FunctionalTensorWrapper. + // Because of this, we need to make sure to run the reference shape function above, + // BEFORE doing this (otherwise we'll end up running the reference function using the wrong sizes/strides) + at::functionalization::impl::mutate_view_meta({view_tensor_name}, view_meta); + // See Note [Propagating strides in the functionalization pass] + // XLA/LTC don't implement the logic to propagate strides correctly, so we need to rely + // on a reference implementation here (instead of relying on the output from the forward lambda + // having the correct stride info) + if (compute_reference_meta && !disable_meta_reference()) {{ + at::functionalization::impl::set_sizes_strides_offset({view_tensor_name}, reference_tensor_output); + }} + return {view_tensor_name}; + }} +""" + + else: + return f""" + {dispatcher_sig.defn(name=wrapper_name(f.func), is_redispatching_fn=True)} {{ + {unwrap_tensor_args_str} + if (!at::functionalization::impl::isFunctionalTensor({view_tensor_name})) {{ + // functionalization is re-entrant, but will no-op if it wasn't passed a FunctionalTensorWrapper. + at::AutoDispatchSkipFunctionalize guard; + return at::_ops::{noop_api_name}::call({", ".join(view_redispatch_args)}); + }} + auto reapply_views = at::functionalization::impl::getFunctionalizationReapplyViewsTLS(); + auto inverse_return_mode = ( + reapply_views ? at::functionalization::InverseReturnMode::ViewOrScatterInverse + : at::functionalization::InverseReturnMode::NeverView + ); + auto compute_reference_meta = + {view_tensor_name}.key_set().has_backend(c10::BackendComponent::XLABit) || + {view_tensor_name}.key_set().has_backend(c10::BackendComponent::LazyBit); + {return_type} reference_tensor_output; + if (compute_reference_meta && !disable_meta_reference()) {{ + {meta_conversion_str} + at::AutoDispatchSkipFunctionalize func_guard; + c10::impl::ExcludeDispatchKeyGuard guard(exclude_keys_for_meta_dispatch); + reference_tensor_output = at::_ops::{noop_api_name}::call({", ".join(meta_call_args)}); + }} + {return_type} tmp_output; + {{ + at::AutoDispatchSkipFunctionalize guard; + if (reapply_views) {{ + tmp_output = at::_ops::{noop_api_name}::call({", ".join(view_redispatch_args)}); + }} else {{ + tmp_output = at::_ops::{api_name}::call({", ".join(view_redispatch_args)}); + }} + }} + {symbolic_inputs_check} + auto view_meta = {spec.new()}; + auto out = at::functionalization::impl::create_functional_tensor_with_view_meta(tmp_output, {view_tensor_name}, view_meta); + // See Note [Propagating strides in the functionalization pass] + if (compute_reference_meta && !disable_meta_reference()) {{ + at::functionalization::impl::set_sizes_strides_offset(out, reference_tensor_output); + }} + return out; + }} +""" + + +def maybe_create_output(f: NativeFunction, var_name: str) -> str: + if len(f.func.returns) == 0: + return "" + return_type = dispatcher.returns_type(f.func.returns).remove_const_ref().cpp_type() + return f"{return_type} {var_name} = " + + +# Given a NativeFunction, and a variable name corresponding to the output of redispatching on the function, +# this returns two lists of names, consisting of: +# - the names of returns corresponding to the original (mutable) inputs of the outer function +# - the names of returns corresponding to the (immutable) outputs of the inner redispatched function +def get_mutable_redispatch_return_names( + f: NativeFunction, inner_return_var: str +) -> tuple[list[str], list[str]]: + aliased_returns = [] + non_aliased_returns = [] + for i, name in enumerate(f.func.aliased_return_names()): + if name is not None: + aliased_returns.append(name) + else: + non_aliased_returns.append( + inner_return_var + if len(f.func.returns) == 1 + else f"std::get<{i}>({inner_return_var})" + ) + return aliased_returns, non_aliased_returns + + +# When functionalization "no-op's" and redispatches on a mutable operator, we need to take care so that: +# - For fresh outputs, we return the result of the redispatch (without wrapping outputs) +# - For outputs that were aliased to inputs, we return the inputs directly (since some of them might have been wrapped) +def return_from_mutable_noop_redispatch( + f: NativeFunction, inner_return_var: str +) -> str: + aliased, non_aliased = get_mutable_redispatch_return_names(f, inner_return_var) + # Just get all of the return names, and immediately return them + return return_str(f.func.returns, aliased + non_aliased) + + +def wrap_propagate_mutations_and_return( + f: NativeFunction, functional_op: NativeFunction, inner_return_var: str +) -> str: + mutable_arg_names = f.func.arguments.mutable_arg_names() + ( + aliased_outer_rets, + non_aliased_outer_rets, + ) = get_mutable_redispatch_return_names(f, inner_return_var) + _, non_aliased_inner_rets = get_mutable_redispatch_return_names( + functional_op, inner_return_var + ) + # The outer function may have a mix of aliased and non-aliased outputs, + # But the inner functional op that we're transforming to should only have non-aliased outputs + if len(mutable_arg_names) + len(non_aliased_outer_rets) != len( + non_aliased_inner_rets + ): + raise AssertionError( + f"Expected {len(mutable_arg_names)} + {len(non_aliased_outer_rets)} == {len(non_aliased_inner_rets)}" + ) + + # First, take all of the newly created outputs from the inner call and wrap them into functional tensors + updates = [] + non_aliased_wrapped_ret_names = [] + for i, inner_ret in enumerate( + non_aliased_inner_rets[: len(non_aliased_outer_rets)] + ): + ret_name = f"output_{i}" + updates.append( + f"""\ + auto output_{i} = at::functionalization::impl::to_functional_tensor({inner_ret});""" + ) + non_aliased_wrapped_ret_names.append(ret_name) + + # Next, take all of the mutated outputs from the inner call corresponding to mutated inputs, + # and propagate the mutations + for outer_arg, inner_ret in zip( + mutable_arg_names, non_aliased_inner_rets[len(non_aliased_outer_rets) :] + ): + updates.append( + f"""\ + auto {outer_arg}_inner = at::functionalization::impl::from_functional_tensor({outer_arg}); + at::functionalization::impl::replace_({outer_arg}, {inner_ret}); + at::functionalization::impl::commit_update({outer_arg}); + at::functionalization::impl::sync({outer_arg}); + auto {outer_arg}_inner_updated = at::functionalization::impl::from_functional_tensor({outer_arg}); + at::functionalization::impl::propagate_xla_data_direct({outer_arg}_inner, {outer_arg}_inner_updated);""" + ) + + # Finally, we return: + # - Any mutable arguments that also returns + # - Any immutable returns that were created wrapping the output from the inner call + returns_str = return_str( + f.func.returns, aliased_outer_rets + non_aliased_wrapped_ret_names + ) + updates_str = "\n".join(updates) + return f"""\ +{updates_str} + {returns_str}""" + + +def maybe_replace_cumulative_out_dtype_exprs( + f: NativeFunction, + functional_sig: DispatcherSignature, + functional_exprs: list[str], +) -> list[str]: + if ( + f.func.kind() != SchemaKind.out + or f.func.name not in CUMULATIVE_OUT_OPS_PRESERVING_OUT_DTYPE + ): + return functional_exprs + + if len(f.func.arguments.out) != 1: + raise AssertionError( + f"Expected a single out argument for cumulative out op: {f.func.name}" + ) + + dtype_arg_idx = next( + (i for i, arg in enumerate(functional_sig.arguments()) if arg.name == "dtype"), + None, + ) + if dtype_arg_idx is None: + raise AssertionError( + f"Expected dtype argument for cumulative out op: {f.func.name}" + ) + + adjusted_exprs = functional_exprs.copy() + dtype_expr = adjusted_exprs[dtype_arg_idx] + adjusted_exprs[dtype_arg_idx] = ( + f"{dtype_expr}.has_value() ? {dtype_expr} : " + f"std::optional({f.func.arguments.out[0].name}_.scalar_type())" + ) + return adjusted_exprs + + +# Generates the Functionalization kernel for: +# - mutation ops (inplace and out= ops) +@with_native_function_and +def emit_inplace_functionalization_body( + f: NativeFunction, g: NativeFunctionsGroup +) -> str: + # mutation case + if not modifies_arguments(f): + raise AssertionError(f"Expected function to modify arguments: {f.func}") + + dispatcher_sig = DispatcherSignature.from_schema(f.func) + + unwrap_tensor_args_str, unwrapped_args_ctx = unwrap_tensor_args( + dispatcher_sig, is_view_op=False + ) + + mutated_names = [ + a.name + for a in f.func.arguments.flat_all + if a.type.is_tensor_like() and a.annotation is not None + ] + non_mutated_names = [ + a.name + for a in f.func.arguments.flat_all + if a.type.is_tensor_like() and a.annotation is None + ] + non_mutated_tensor_names = [ + a.name + for a in f.func.arguments.flat_all + if a.type == BaseType(BaseTy.Tensor) and a.annotation is None + ] + # all mutable inputs must be functional tensors in order to participate in functionalization + check_all_mutated_args_are_functional = " && ".join( + ["true"] + + [ + f"at::functionalization::impl::isFunctionalTensor({a})" + for a in mutated_names + ] + ) + check_any_non_mutated_args_are_functional = " || ".join( + ["false"] + + [ + f"at::functionalization::impl::isFunctionalTensor({a})" + for a in non_mutated_names + ] + ) + + check_any_non_mutated_tensors_are_xla = " || ".join( + ["false"] + + [ + f"{a}.device().type() == c10::DeviceType::XLA" + for a in non_mutated_tensor_names + ] + ) + # These are used in the cases where we don't functionalize and redispatch to the inplace op + # case 1: we hit an inplace op that doesn't have an out-of-place equivalent + # case 2: we hit an inplace ops but our inputs are not functional tensors (in which case our kernel just no-ops) + inplace_exprs = [ + e.expr + for e in translate(unwrapped_args_ctx, dispatcher_sig.arguments(), method=False) + ] + + # call the out-of-place variant of the op + return_type = ( + dispatcher.returns_type(g.functional.func.returns).remove_const_ref().cpp_type() + ) + functional_sig = DispatcherSignature.from_schema(g.functional.func) + functional_exprs = [ + e.expr + for e in translate(unwrapped_args_ctx, functional_sig.arguments(), method=False) + ] + functional_exprs = maybe_replace_cumulative_out_dtype_exprs( + f, functional_sig, functional_exprs + ) + + meta_conversion_str, meta_call_ctx = convert_to_meta_tensors(dispatcher_sig) + # We don't want to run the inplace meta func for ops like .set_(), because: + # (1) they're unnecessary: inplace meta checks are only useful for ops like add_(), + # where broadcasting will work for the out-of-place case but should fail on the inplace call + # (2) They'll also fail without adding extra infra: we'd need to convert the input storage argument + # into a meta storage + any_storage_args = any( + a.type == BaseType(BaseTy.Storage) for a in f.func.arguments.flat_all + ) + + return f""" + {dispatcher_sig.defn(name=wrapper_name(f.func), is_redispatching_fn=True)} {{ + if ({str(not any_storage_args and f.func.kind() == SchemaKind.inplace).lower()} && !disable_meta_reference()) {{ + // Before converting the mutable op to its functional variant, run meta tensors through the original op. + // This will help us catch shape errors that apply to inplace ops that wouldn't apply to their functional variants. + // (We can only do this for inplace ops today though, because they technically all support meta tensors). + {meta_conversion_str} + at::AutoDispatchSkipFunctionalize func_guard; + c10::impl::ExcludeDispatchKeyGuard guard(exclude_keys_for_meta_dispatch); + at::_ops::{f.func.name.unambiguous_name()}::call({", ".join(a.name for a in meta_call_ctx)}); + }} + {unwrap_tensor_args_str} + if (!({check_all_mutated_args_are_functional})) {{ + // We want to disable this check if there are any XLA tensors. + // cpu_tensor.copy_(xla_tensor) is valid code. + if (!({check_any_non_mutated_tensors_are_xla}) && ({check_any_non_mutated_args_are_functional})) {{ + // case 1: trying to mutate a non functional tensor with a functional tensor is an error + TORCH_INTERNAL_ASSERT(false, + "mutating a non-functional tensor with a functional tensor is not allowed.", + " Please ensure that all of your inputs are wrapped inside of a functionalize() call."); + }} else {{ + // case 2: arguments are not functional tensors, so we no-op and redispatch. + at::AutoDispatchSkipFunctionalize guard; + {maybe_create_output(f, "tmp_output")}at::_ops::{f.func.name.unambiguous_name()}::call({", ".join(inplace_exprs)}); + {return_from_mutable_noop_redispatch(f, "tmp_output")} + }} + }} else {{ + {return_type} tmp_output; + {{ + at::AutoDispatchSkipFunctionalize guard; + tmp_output = at::_ops::{g.functional.func.name.unambiguous_name()}::call({", ".join(functional_exprs)}); + }} + {wrap_propagate_mutations_and_return(f, g.functional, "tmp_output")} + }} + }}""" + + +# The below functions generate RegisterFunctionalization.cpp +# These files provide the kernels that run the functionalization pass, which can be opted into +# per backend (e.g. XLA or Vulkan), or as a composable transform (functionalize() in functorch). + + +# See Note [Functionalization Pass: View Inverses]. +def gen_functionalization_view_inverse_declaration( + selector: SelectiveBuilder, g: NativeFunctionsViewGroup +) -> str | None: + # For every (non-composite) view op, we need a corresponding "inverse view" function. + # This generates the declarations so we get a good compiler error when someone adds a new view. + @with_native_function + def emit_decl_helper(g: NativeFunctionsViewGroup) -> str | None: + if g.view.has_composite_implicit_autograd_kernel: + return None + view_inverse_sig = ViewInverseSignature(g) + return view_inverse_sig.decl() + + return emit_decl_helper(g) + + +# Helper class for generating `ViewMeta` specializations. +@dataclass +class ViewMetaSpecialization: + g: NativeFunctionsViewGroup + f: NativeFunction + + @property + def is_multi_output(self) -> bool: + return functionalization.is_multi_output(self.f.func) + + @property + def is_as_strided(self) -> bool: + return str(self.f.func.name) == "as_strided" + + @property + def out_index(self) -> str: + if self.is_multi_output: + return functionalization.out_index_binding.name + return "0" + + @property + def classname(self) -> str: + return functionalization.classname(self.f.func) + + def decl(self) -> list[str]: + base_ctor_arguments = functionalization.base_ctor_arguments(self.f.func) + extra_ctor_arguments = functionalization.extra_ctor_arguments(self.f.func) + attributes = functionalization.attributes(self.f.func) + + # List of types for declaring the `SerializableTuple` type. + serializable_tuple_args = ",\n".join( + f" {binding.type} /* {binding.name} */" + for binding in (base_ctor_arguments + attributes) + ) + + # Arguments used for forwarding the tuple elements to the constructor. + destructure_tuple_args = ", ".join( + f"std::get<{i}>(tpl)" + for i in range(len(base_ctor_arguments) + len(extra_ctor_arguments)) + ) + + # List of constructor parameters + ctor_parameters = ", ".join( + binding.decl() for binding in (base_ctor_arguments + extra_ctor_arguments) + ) + + # Call the base class `ViewMeta` constructor. + # + # Both of `is_multi_output` and `is_as_strided` are known values, given the + # operation schema. + is_multi_output_str = str(self.is_multi_output).lower() + is_as_strided_str = str(self.is_as_strided).lower() + + base_ctor_bindings = ", ".join( + [ + # `has_symbolic_inputs` is always taken as parameter. + functionalization.has_symbolic_inputs_binding.name, + f"/*is_multi_output=*/{is_multi_output_str}", + f"/*is_as_strided=*/{is_as_strided_str}", + # `out_index` is know if the operation returns only one value. Otherwise, + # we also take it as parameter. + f"/*out_index=*/{self.out_index}", + ] + ) + + # Assignments of `extra_ctor_arguments` to their corresponding fields. + # These are extra fields to-be-declared in this specialization. + # + # We need to set `allow_expensive_conversions`, since we are storing owned versions + # of the non-owning arguments. + ctor_assignments = ",\n".join( + f" {e.type.name}({e.expr})" + for e in translate( + extra_ctor_arguments, + attributes, + method=False, + allow_expensive_conversions=True, + ) + ) + + # List of arguments for constructing the `SerializableTuple` from an instance. + tuple_arguments = ", ".join( + binding.name for binding in (base_ctor_arguments + attributes) + ) + + # List of field declarations. + attr_declarations = "\n".join(f" {binding.decl()};" for binding in attributes) + + # Override `to_out_index` if this operation returns more than 1 value. + to_out_index_decl = "" + if self.is_multi_output: + to_out_index_decl = ( + " std::shared_ptr to_out_index(int64_t out_idx) override;" + ) + + return [ + f""" +struct TORCH_API {self.classname} : public ViewMeta {{ + FUNCTIONALIZATION_VIEWMETA_NAME({self.classname}) + FUNCTIONALIZATION_VIEWMETA_SERIALIZABLE_TUPLE(\n{serializable_tuple_args}); + + {self.classname}(const SerializableTuple& tpl) + : {self.classname}({destructure_tuple_args}) {{}} + + {self.classname}({ctor_parameters}) + : at::functionalization::ViewMeta({base_ctor_bindings}), +{ctor_assignments} {{}} + + Tensor forward(const Tensor& base) override; + Tensor reverse(const Tensor& base, const Tensor& mutated_view) override; +{to_out_index_decl} + + SerializableTuple to_serializable_tuple() {{ + return std::make_tuple({tuple_arguments}); + }} + +{attr_declarations} +}}; +""" + ] + + # Generate a call to the actual operation. + def opcall(self, is_reverse: bool, reapply_views: bool) -> str: + opname = functionalization.name( + self.g, + is_reverse=is_reverse, + include_namespace=True, + reapply_views=reapply_views, + ) + + # Expected arguments for the operation. + if self.g.view_copy is None: + raise AssertionError("Expected view_copy to be non-None") + op_arguments = functionalization.op_arguments(self.g.view_copy.func, is_reverse) + + # The context is composed by the constructor arguments (which are also + # the field variables stored in the instance), and the `base` tensor. + context = [functionalization.base_binding] + context += functionalization.base_ctor_arguments(self.f.func) + context += functionalization.attributes(self.f.func) + + # If we are generating the call for the reverse function, we also have + # access to `mutated_view` argument. + if is_reverse: + context.append(functionalization.mutated_view_binding) + + arguments = ", ".join( + [e.expr for e in translate(context, op_arguments, method=False)] + ) + + # Index the result if this operation returns multiple values. + maybe_index = "" + if not is_reverse and self.is_multi_output: + maybe_index = f"[{self.out_index}]" + + return f"{opname}({arguments}){maybe_index}" + + def impl(self) -> list[str]: + functions = [ + f""" +at::Tensor {self.classname}::forward(const at::Tensor& base) {{ + if (reapply_views) {{ + return {self.opcall(is_reverse=False, reapply_views=True)}; + }} else {{ + return {self.opcall(is_reverse=False, reapply_views=False)}; + }} +}}""", + f""" +at::Tensor {self.classname}::reverse(const at::Tensor& base, const Tensor& mutated_view) {{ + return {self.opcall(is_reverse=True, reapply_views=True)}; +}}""", + ] + + # If this operation returns multiple values, also generate a `to_out_index` + # implementation. + if self.is_multi_output: + functions.append(f""" +std::shared_ptr {self.classname}::to_out_index(int64_t out_index) {{ + return {self.new("out_index")}; +}} +""") + + return functions + + # Create the Python binding for this specialized class. + def binding(self) -> list[str]: + name = functionalization.classname(self.f.func, with_namespace=True) + return [f" create_binding_with_pickle<{name}>(functionalization);"] + + # Generate an instantiation of this specialized class. + def new(self, out_index: str = "0") -> str: + name = functionalization.classname(self.f.func, with_namespace=True) + ctor_arguments = functionalization.base_ctor_arguments( + self.f.func + ) + functionalization.extra_ctor_arguments(self.f.func) + # Replace the `out_index` parameter with the given `out_index`. + arguments = ", ".join( + binding.name if binding.name != "out_index" else out_index + for binding in ctor_arguments + ) + return f"std::make_shared<{name}>({arguments})" + + # Run the function `run` for both: `view` and `view_inplace` functions. + @staticmethod + def map( + g: NativeFunctionsViewGroup, run: Callable[[ViewMetaSpecialization], list[str]] + ) -> list[str]: + def maybe_run(f: NativeFunction | None) -> list[str]: + if f is None: + return [] + with native_function_manager(f): + return run(ViewMetaSpecialization(g, f)) + + return list(concatMap(maybe_run, (g.view, g.view_inplace))) + + +def gen_functionalization_view_meta_classes_base( + selector: SelectiveBuilder, + g: NativeFunctionsViewGroup, + run: Callable[[ViewMetaSpecialization], list[str]], +) -> list[str]: + if not selector.include_all_operators: + return [] + + if g.composite: + return [] + + return ViewMetaSpecialization.map(g, run) + + +def gen_functionalization_view_meta_classes_decl( + selector: SelectiveBuilder, g: NativeFunctionsViewGroup +) -> list[str]: + return gen_functionalization_view_meta_classes_base( + selector, g, ViewMetaSpecialization.decl + ) + + +def gen_functionalization_view_meta_classes_impl( + selector: SelectiveBuilder, g: NativeFunctionsViewGroup +) -> list[str]: + return gen_functionalization_view_meta_classes_base( + selector, g, ViewMetaSpecialization.impl + ) + + +def gen_functionalization_view_meta_classes_binding( + selector: SelectiveBuilder, g: NativeFunctionsViewGroup +) -> list[str]: + return gen_functionalization_view_meta_classes_base( + selector, g, ViewMetaSpecialization.binding + ) + + +# Generates the Python bindings for the `ViewMeta` specialized classes. +def gen_functionalization_view_meta_classes( + native_functions_path: str, + tags_path: str, + selector: SelectiveBuilder, + install_dir: str, + template_dir: str, +) -> None: + from torchgen.gen import get_grouped_by_view_native_functions, parse_native_yaml + + # Parse the native_functions.yaml. + # Then, group them into `NativeFunctionsViewGroup`. + # + # This is the same steps we do in gen.py (ATen codegen). + native_functions = parse_native_yaml( + native_functions_path, tags_path + ).native_functions + native_functions_with_view_groups = get_grouped_by_view_native_functions( + native_functions + ) + view_groups = [ + g + for g in native_functions_with_view_groups + if isinstance(g, NativeFunctionsViewGroup) + ] + + fm = FileManager(install_dir=install_dir, template_dir=template_dir, dry_run=False) + fm.write( + "ViewMetaClassesPythonBinding.cpp", + lambda: { + "view_meta_bindings": list( + concatMap( + lambda g: gen_functionalization_view_meta_classes_binding( + selector, g + ), + view_groups, + ) + ), + }, + ) + + +def gen_functionalization_registration( + selector: SelectiveBuilder, + g: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup, + composite_implicit_autograd_index: BackendIndex, +) -> list[str]: + @with_native_function + def emit_registration_helper(f: NativeFunction) -> str: + if f.has_composite_implicit_autograd_kernel: + metadata = composite_implicit_autograd_index.get_kernel(f) + if metadata is None: + raise AssertionError( + f"Expected metadata for composite implicit autograd kernel: {f.func}" + ) + native_api_name = metadata.kernel + sig = NativeSignature(f.func, symint=metadata.supports_symint()) + # Note [Composite view ops in the functionalization pass] + # We don't need to worry about implemententing functionalization kernels for views with + # CompositeImplicitAutograd kernels, because we can just decompose them into their base operators. + # We can't just opt the entire Functionalization dispatch key into the composite keyset though, + # because we don't want to decompose non-view ops that are composite, like `at::ones`. + registration_str = ( + f"static_cast<{sig.ptr_type()}>(at::native::{native_api_name})" + ) + else: + # non-composite view ops (and inplace ops) get a normal registration. + registration_str = f"TORCH_FN(functionalization::{wrapper_name(f.func)})" + return f'm.impl("{f.func.name}", {registration_str});' + + # Don't generate kernels in mobile build + if not selector.include_all_operators: + return [] + + if isinstance(g, NativeFunctionsViewGroup): + # functionalization needs to register kernels for view + view_inplace ops + # See Note [Functionalization <> torch.Tensor constructor] + if str(g.view.func.name) == "lift_fresh": + return [] + view_str = [] + view_str.append(emit_registration_helper(g.view)) + if g.view_inplace is not None: + if not g.view_inplace.is_view_op: + raise AssertionError( + f"Expected view_inplace to be a view op: {g.view_inplace.func}" + ) + view_str.append(emit_registration_helper(g.view_inplace)) + return view_str + + elif isinstance(g, NativeFunctionsGroup): + # Gets a hand-written functionalization kernel + if g.inplace is not None and str(g.inplace.func.name) == "set_.source_Tensor": + fns = [] + else: + fns = list(g.functions()) + else: + if str(g.func.name) in MUTABLE_OPS_NOT_USING_FUNCTIONALIZATION: + return [] + fns = [g] + + registrations = [] + for f in fns: + if f.has_composite_implicit_autograd_kernel: + continue + if str(f.func.name) == "lift": + # See Note [Functionalization <> torch.Tensor constructor] + return [] + if str(f.func.name) == "resize_": + # See Note [resize_ in Functionalization] + return [] + if str(f.func.name.name) != "set_": + if f.is_view_op: + raise AssertionError(f"Unexpected view op: {f.func}") + # functionalization needs to generate and register kernels for inplace ops. + # We *also* need to directly register CompositeImplicitAUtograd kernels + # so that they decompose properly before functioanlization. + if modifies_arguments(f): + registrations.append(emit_registration_helper(f)) + return registrations + + +def gen_functionalization_definition( + selector: SelectiveBuilder, + # Note: Ideally this code should never have to look at NativeFunction + # (and instead only need to operate on grouped NativeFunctions). + # The only reason currently is because we need to emit direct dispatch registrations + # For CompositeImplicitAutograd operators, which are potentially ungrouped. + g: NativeFunction | NativeFunctionsGroup | NativeFunctionsViewGroup, +) -> list[str]: + # Don't generate kernels in mobile build + if not selector.include_all_operators: + return [] + + if isinstance(g, NativeFunctionsViewGroup): + # Case 1: emit view -> view_copy kernels for the functionalization pass + view_defs = [] + if not g.composite: + # invariant: NativeFunctionsViewGroup's always have a view_copy operator + # if the view is not composite (implicit autograd) + if g.view_copy is None: + raise AssertionError( + f"Expected view_copy to be non-None: {dataclass_repr(g, indent=1)}" + ) + view_defs.append(emit_view_functionalization_body(g, view_inplace=False)) + if g.view_inplace is not None: + view_defs.append(emit_view_functionalization_body(g, view_inplace=True)) + return view_defs + elif isinstance(g, NativeFunction): + # Invariant: all mutable operators that we need to handle in functionalization + # should have been properly grouped up. + # TODO: The below ops all have "problematic" schemas that prevent them from + # getting functionalized. Instead of bending over backwards to get things to work, + # I think we should either: + # (1) fix their schemas (BC-breaking) + # (2) hand-write their functionalization kernels + if ( + str(g.func.name) not in MUTABLE_OPS_NOT_USING_FUNCTIONALIZATION + and str(g.func.name.name) not in MUTABLE_OPS_NOT_USING_FUNCTIONALIZATION + ): + if not ( + g.has_composite_implicit_autograd_kernel or not modifies_arguments(g) + ): + raise AssertionError( + f"Expected composite implicit autograd kernel or non-modifying function: {g.func}" + ) + return [] + else: + # Case 2: emit inplace -> out-of-place kernels for the functionalization pass + mutation_defs = [] + mutation_defs.append(emit_inplace_functionalization_body(g.out, g)) + if g.inplace is not None: + mutation_defs.append(emit_inplace_functionalization_body(g.inplace, g)) + if g.mutable is not None: + mutation_defs.append(emit_inplace_functionalization_body(g.mutable, g)) + return mutation_defs + return [] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_lazy_tensor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_lazy_tensor.py new file mode 100644 index 0000000000000000000000000000000000000000..5cacf4002dba9036023795ebe88c417fdb60e4fd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_lazy_tensor.py @@ -0,0 +1,593 @@ +from __future__ import annotations + +import argparse +import os +from collections import namedtuple +from pathlib import Path +from typing import Any, TYPE_CHECKING + +import yaml + +import torchgen.dest as dest +from torchgen.api.lazy import setValueT +from torchgen.api.types import BaseCppType +from torchgen.dest.lazy_ir import GenLazyIR, GenLazyNativeFuncDefinition, GenTSLazyIR +from torchgen.gen import get_grouped_native_functions, parse_native_yaml +from torchgen.gen_backend_stubs import ( + error_on_missing_kernels, + gen_dispatcher_registrations, + gen_dispatchkey_nativefunc_headers, + parse_backend_yaml, +) +from torchgen.model import NativeFunction, NativeFunctionsGroup, OperatorName +from torchgen.selective_build.selector import SelectiveBuilder +from torchgen.utils import FileManager, NamespaceHelper +from torchgen.yaml_utils import YamlLoader + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable, Iterator, Sequence + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Lazy Tensor Codegen +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# Overview +# ~~~~~~~~ +# +# This codegen script builds on existing data models and helpers used +# by all ATen backends, and adds new functionality specific to lazy +# tensor backends. +# +# Inputs: +# - _native_functions.yaml: controls which operators are +# supported by the backend. +# +# Outputs: +# (for all backends) +# Ir.h defines Lazy IR classes to be constructed during tracing +# - opt-in: also generate 'lowering' methods for the TorchScript backend only +# NativeFunctions.cpp defines implementations of native functions which perform lazy tracing +# - opt-in: 'full_codegen' section of backend yaml; 'supported' section omits these implementations +# NativeFunctions.h declares implementations of native functions for both 'supported' and 'full_codegen' +# ops +# +# Register.cpp registers all op implementations with the dispatcher +# RegisterAutograd.cpp registers all autograd implementations with the dispatcher +# +# Validation Helpers: +# - Shape Inference: errs if any ops in backend yaml require shape inference not provided by meta kernels or +# implementations in torch/csrc/lazy/core/shape_inference.* +# - native function impls: errs if any 'supported' ops do not have an implementation defined in the backend +# (non-codegen) implementation file +# +# +# About the Data Model +# ~~~~~~~~~~~~~~~~~~~~ +# +# Modeled after ATen codegen, the first step is to parse yaml and build a data model for the operators +# we care about. In this case, the _native_functions yaml defines a subset of the core operators +# (defined in more detail in the main native_functions.yaml), which will be supported by your backend. +# Backends can list ops in two categories: +# - `supported` ops require hand-implementations but still get codegenned declarations and registrations +# - `full_codegen` ops get implementations (and IR classes) generated too +# +# Each native function is modeled as an object with a schema, and each schema has objects representing their +# arguments. Much of the codegen is manipulation of the arguments and their types. For example, lazy tensor +# backends need to transform 'at::Tensor' arguments into 'lazy::Value' objects, as well as replacing reference +# types (stringref) with actual string objects, and this is done by manipulating the data model objects. +# - see api/lazy.py for the lazy data model +# +# Once the data model is set up, the rest of this script processes a number of templates for output CPP file +# and fills in the template values using helpers in `dest/lazy_ir.py` and `dest/lazy_ts_lowering.py`. These +# helpers mostly iterate over functions and their arguments, outputting different c++ snippets. +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +# Parses the external backend's yaml, and adds a new BackendIndex for the backend's dispatch key. +# Returns a Tuple of (backend_key, autograd_key, cpp_namespace, updated BackendIndex mapping, full_codegen) +ParsedExternalYaml = namedtuple( + "ParsedExternalYaml", + ["backend_key", "autograd_key", "cpp_namespace", "backend_indices", "full_codegen"], +) + + +def parse_native_functions_keys( + backend_yaml_path: str, + grouped_native_functions: Sequence[NativeFunction | NativeFunctionsGroup], +) -> tuple[list[OperatorName], list[Any], list[OperatorName]]: + with open(backend_yaml_path) as f: + yaml_values = yaml.load(f, Loader=YamlLoader) + if not isinstance(yaml_values, dict): + raise AssertionError(f"Expected dict from YAML, got {type(yaml_values)}") + + full_codegen = yaml_values.pop("full_codegen", []) + non_native = yaml_values.pop("non_native", []) + ir_gen = yaml_values.pop("ir_gen", []) + if not isinstance(full_codegen, list): + raise AssertionError( + f"Expected full_codegen to be list, got {type(full_codegen)}" + ) + if not isinstance(non_native, list): + raise AssertionError(f"Expected non_native to be list, got {type(non_native)}") + if not isinstance(ir_gen, list): + raise AssertionError(f"Expected ir_gen to be list, got {type(ir_gen)}") + full_codegen_opnames = [OperatorName.parse(name) for name in full_codegen] + ir_gen_opnames = [OperatorName.parse(name) for name in ir_gen] + return full_codegen_opnames, non_native, ir_gen_opnames + + +def validate_shape_inference_header( + shape_inference_hdr: str, expected_shape_infr_decls: list[str] +) -> None: + try: + with open(shape_inference_hdr) as f: + shape_infr_decls = f.read() + shape_infr_decl_lines = set(shape_infr_decls.split("\n")) + except OSError as e: + raise AssertionError( + f"Unable to read from the specified shape_inference_hdr file: {shape_inference_hdr}" + ) from e + + # TODO(whc) add a check for shape inference functions that have meta kernels implement and should be retired. + + missing_decls = [ + decl for decl in expected_shape_infr_decls if decl not in shape_infr_decl_lines + ] + if missing_decls: + raise Exception( # noqa: TRY002 + f"""Missing shape inference function.\n +Please add declare this function in {shape_inference_hdr}:\n +and implement it in the corresponding shape_inference.cpp file.\n +{os.linesep.join(missing_decls)}""" + ) + + +# Some helper functions for the codegen. +def get_ltc_helper_fns() -> str: + return """\ +at::Tensor to_meta(const at::Tensor& tensor) { + // undefined tensors can't be converted to the meta device, since they don't have sizes/strides + if (!tensor.defined()) return tensor; + auto out = at::native::empty_strided_meta_symint(tensor.sym_sizes(), tensor.sym_strides(), \ +/*dtype=*/tensor.scalar_type(), /*layout=*/tensor.layout(), \ +/*device=*/c10::Device(c10::kMeta), /*pin_memory=*/std::nullopt); + // needs to handle wrapped numbers, so dtype promotion works properly. + if (tensor.unsafeGetTensorImpl()->is_wrapped_number()) { + out.unsafeGetTensorImpl()->set_wrapped_number(true); + } + return out; +} +std::optional to_meta(const std::optional& tensor) { + if (tensor.has_value()) { + return to_meta(*tensor); + } + return std::nullopt; +} + +std::vector to_meta(at::ITensorListRef t_list) { + std::vector outs; + outs.reserve(t_list.size()); + for (const auto& tensor : t_list) { + outs.push_back(to_meta(tensor)); + } + return outs; +} +""" + + +class default_args: + node_base: str = "Node" + node_base_hdr: str | None = None + shape_inference_hdr: str = "torch/csrc/lazy/core/shape_inference.h" + tensor_class: str = "torch::lazy::LazyTensor" + tensor_class_hdr: str = "torch/csrc/lazy/core/tensor.h" + lazy_ir_generator: type[GenLazyIR] = GenLazyIR + native_func_definition_generator: type[GenLazyNativeFuncDefinition] = ( + GenLazyNativeFuncDefinition + ) + backend_name: str = "TorchScript" + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate Lazy Tensor backend files") + parser.add_argument( + "-s", + "--source-yaml", + "--source_yaml", + help="path to source yaml file containing operator external definitions", + ) + parser.add_argument("-o", "--output-dir", "--output_dir", help="output directory") + parser.add_argument( + "--dry-run", "--dry_run", type=bool, default=False, help="output directory" + ) + parser.add_argument( + "--impl-path", + "--impl_path", + type=str, + default=None, + help="path to the source C++ file containing kernel definitions", + ) + parser.add_argument( + "--gen-ts-lowerings", + "--gen_ts_lowerings", + action="store_true", + help="Generate TorchScript lowerings in addition to Lazy IR and NativeFunctions", + ) + parser.add_argument( + "--node-base", + "--node_base", + type=str, + default=default_args.node_base, + help="Name of backend specific custom Lazy IR Node base class", + ) + parser.add_argument( + "--node-base-hdr", + "--node_base_hdr", + type=str, + default=default_args.node_base_hdr, + help="Path to header file defining custom Lazy IR Node base class", + ) + parser.add_argument( + "--shape-inference-hdr", + "--shape_inference_hdr", + type=str, + default=default_args.shape_inference_hdr, + help="Path to header file defining custom Lazy shape inference functions", + ) + parser.add_argument( + "--tensor-class", + "--tensor_class", + type=str, + default=default_args.tensor_class, + help="Name of backend specific custom Lazy Tensor class", + ) + parser.add_argument( + "--tensor-class-hdr", + "--tensor_class_hdr", + type=str, + default=default_args.tensor_class_hdr, + help="Path to header file defining custom Lazy Tensor class", + ) + parser.add_argument( + "--backend-name", + "--backend_name", + type=str, + default=default_args.backend_name, + help="Name of the backend to generate", + ) + options = parser.parse_args() + + # Assumes that this file lives at PYTORCH_ROOT/torchgen/gen_backend_stubs.py + torch_root = Path(__file__).absolute().parents[2] + aten_path = str(torch_root / "aten" / "src" / "ATen") + lazy_ir_generator: type[GenLazyIR] = default_args.lazy_ir_generator + if options.gen_ts_lowerings: + lazy_ir_generator = GenTSLazyIR + native_func_definition_generator: type[GenLazyNativeFuncDefinition] = ( + default_args.native_func_definition_generator + ) + + run_gen_lazy_tensor( + aten_path, + options.source_yaml, + options.output_dir, + options.dry_run, + options.impl_path, + options.node_base, + options.node_base_hdr, + options.tensor_class, + options.tensor_class_hdr, + options.shape_inference_hdr, + lazy_ir_generator, + native_func_definition_generator, + options.backend_name, + ) + + +def run_gen_lazy_tensor( + aten_path: str, + source_yaml: str, + output_dir: str, + dry_run: bool, + impl_path: str | None, + node_base: str = default_args.node_base, + node_base_hdr: str | None = default_args.node_base_hdr, + tensor_class: str = default_args.tensor_class, + tensor_class_hdr: str = default_args.tensor_class_hdr, + shape_inference_hdr: str = default_args.shape_inference_hdr, + lazy_ir_generator: type[GenLazyIR] = default_args.lazy_ir_generator, + native_func_definition_generator: type[ + GenLazyNativeFuncDefinition + ] = default_args.native_func_definition_generator, + # build_in_tree is true for TS backend and affects include paths + build_in_tree: bool = False, + # per_operator_headers changes whether ATen/Functions.h or individual operator headers are used + # it must match how ATen was built + per_operator_headers: bool = False, + backend_name: str = default_args.backend_name, + gen_forced_fallback_code: bool = False, + use_lazy_shape: bool = True, + # the following arguments are temporary customization points for xla backend migration. + # do not rely on them otherwise, they should be removed once migration is complete + backend_namespace: str = "torch::lazy", + get_tensorlist: str = "GetTensorList", + get_tensor_or_wrap_number: str = "GetLtcTensorOrCreateForWrappedNumber", + try_get_tensor: str = "TryGetLtcTensor", + metrics_counter: str = 'TORCH_LAZY_FN_COUNTER("lazy::")', + create_tensor: str = "LazyTensor::Create", + create_from_first_tensor: bool = False, + create_aten_from_ltc_tensor: str = "torch::lazy::CreateAtenFromLtcTensor", + tuple_aten_from_ltc_tensors: str = "torch::lazy::TupleAtenFromLtcTensors", + lazy_value_class: str = "torch::lazy::Value", + lazy_tensor_ptr: str = "LazyTensorPtr", + get_device_fn: str = "torch::lazy::GetBackendDevice", +) -> None: + lv_tokens = lazy_value_class.split("::") + lv_class = lv_tokens[-1] + lv_ns = "::".join(lv_tokens[:-1]) + setValueT(BaseCppType(lv_ns, lv_class)) + template_dir = os.path.join(aten_path, "templates") + + def make_file_manager(install_dir: str) -> FileManager: + return FileManager( + install_dir=install_dir, template_dir=template_dir, dry_run=dry_run + ) + + fm = make_file_manager(output_dir) + + native_yaml_path = os.path.join(aten_path, "native/native_functions.yaml") + tags_yaml_path = os.path.join(aten_path, "native/tags.yaml") + parsed_yaml = parse_native_yaml(native_yaml_path, tags_yaml_path) + native_functions, backend_indices = ( + parsed_yaml.native_functions, + parsed_yaml.backend_indices, + ) + grouped_native_functions = get_grouped_native_functions(native_functions) + + def sort_native_function(f: NativeFunctionsGroup | NativeFunction) -> str: + """ + We sort the native function because of the note in concat_map_codegen. + TODO(alanwaketan): Remove this sorting hack once all ops are grouped properly. + """ + func = f.functional.func if isinstance(f, NativeFunctionsGroup) else f.func + return str(func.name.name) + + grouped_native_functions = sorted( + grouped_native_functions, key=sort_native_function + ) + + parsed_backend_yaml = parse_backend_yaml( + source_yaml, grouped_native_functions, backend_indices + ) + backend_key = parsed_backend_yaml.backend_key + autograd_key = parsed_backend_yaml.autograd_key + cpp_namespace = parsed_backend_yaml.cpp_namespace + backend_indices = parsed_backend_yaml.backend_indices + # the following 3 keys are all processed differently + # for full_codegen, we generate IR, kernels, etc + # for ir_gen, we generate only IR + # non_native is used to register kernels not declared in + # native_functions.yaml + full_codegen, non_native, ir_gen = parse_native_functions_keys( + source_yaml, grouped_native_functions + ) + + def concat_map_codegen( + func: Callable[[NativeFunction], Sequence[str]], + xs: Iterable[NativeFunctionsGroup | NativeFunction], + ops_list: list[OperatorName] = full_codegen, + ) -> Iterator[str]: + """ + We code-gen for the functional variant, which is all we need for IR classes/lowerings/shape inferences, but we + only code-gen additional entries for the inplace variant for the native functions. + """ + + for x in xs: + fs = list(x.functions()) if isinstance(x, NativeFunctionsGroup) else [x] + for f in fs: + if f.func.name in ops_list: + yield from func(f) + + selector = SelectiveBuilder.get_nop_selector() + + if backend_key is None: + raise AssertionError("backend_key must be non-None") + class_name = backend_indices[backend_key].native_function_class_name() + + if impl_path is not None: + error_on_missing_kernels( + native_functions, + backend_indices, + backend_key, + autograd_key, + class_name, + impl_path, + full_codegen, + ) + + """ Validate Shape Inference Definitions + + Generated lazy native functions all perform shape inference, by first using a meta:: kernel + if available for that op, and otherwise using a 'compute_shape_{op}' function instead. The generator + knows the call signature for compute_shape_{op} because it matches the nativefunction (and meta::) signature, + so it just has to check whether the op is structured and generate a call for one or the other. It's up to the dev + to supply the missing compute_shape_{op} function, but the codegen at least warns you about this and provides + the expected signature which can be copy-pasted into shape_inference.h. + + compute_shape_{op} functions are handwritten and should be replaced over time as ops get ported + to structured kernels. + + See torch/csrc/lazy/core/shape_inference.cpp #READ THIS! for more information. + """ + if shape_inference_hdr is not None: + expected_shape_infr_decls = list( + concat_map_codegen( + dest.GenLazyShapeInferenceDefinition( + backend_indices[backend_key], tensor_class + ), + grouped_native_functions, + ) + ) + + validate_shape_inference_header(shape_inference_hdr, expected_shape_infr_decls) + if class_name is None: + raise AssertionError("class_name must be non-None") + + # Generate nativefunction declarations + # Note, eager registrations is set to False for the lazy TS backend as another LTC backend + # may want to register their own lazy kernels instead of registering the TS ones. + # The registration will lazily happen when init_ts_backend is called. + gen_dispatchkey_nativefunc_headers( + fm, + class_name, + cpp_namespace, + backend_indices, + grouped_native_functions, + backend_key, + autograd_key, + backend_name, + ) + + # Generate Dispatcher registrations which hook up the nativefunctions + for dispatch_key in ( + [backend_key] if autograd_key is None else [backend_key, autograd_key] + ): + gen_dispatcher_registrations( + fm, + output_dir, + class_name, + backend_indices, + grouped_native_functions, + backend_key, + dispatch_key, + selector, + build_in_tree=build_in_tree, + per_operator_headers=per_operator_headers, + backend_name=backend_name, + eager_registration=False, + ) + + # Generate native function impls that build IR nodes + ns_helper = NamespaceHelper(cpp_namespace) + fm.write_with_template( + f"{backend_key}NativeFunctions.cpp", + "DispatchKeyNativeFunctions.cpp", + lambda: { + "includes": [ + f"#include <{path}>" + for path in [ + tensor_class_hdr, + shape_inference_hdr, + "ATen/Functions.h", + "ATen/native/TensorConversions.h", + "ATen/NativeFunctions.h", + "ATen/CompositeExplicitAutogradNonFunctionalFunctions.h", + "ATen/MetaFunctions.h", + "ATen/Operators.h", + "ATen/native/CPUFallback.h", + "torch/csrc/lazy/core/ir_builder.h", + "torch/csrc/lazy/core/lazy_graph_executor.h", + "torch/csrc/lazy/core/metrics.h", + "torch/csrc/lazy/core/shape.h", + f"{output_dir}/{backend_key}NativeFunctions.h", + f"{output_dir}/LazyIr.h", + ] + + ( + ["torch/csrc/lazy/ts_backend/ts_eager_fallback.h"] + if gen_forced_fallback_code + else [] + ) + ], + "helper_fns": get_ltc_helper_fns(), + "native_functions_include": "", + "namespace_prologue": ns_helper.prologue, + "namespace_epilogue": ns_helper.epilogue, + "native_function_definitions": list( + concat_map_codegen( + native_func_definition_generator( + f"{backend_key}NativeFunctions", + backend_indices[backend_key], + tensor_class, + gen_forced_fallback_code, + backend_namespace, + get_tensorlist, + get_tensor_or_wrap_number, + try_get_tensor, + metrics_counter, + create_tensor, + create_from_first_tensor, + create_aten_from_ltc_tensor, + tuple_aten_from_ltc_tensors, + lazy_tensor_ptr, + get_device_fn, + ), + grouped_native_functions, + ) + ), + }, + ) + # Generate IR node classes + lazy_ir_obj = lazy_ir_generator( + backend_indices[backend_key], backend_name, node_base, use_lazy_shape + ) + + fm.write_with_template( + "LazyIr.h", + "LazyIr.h", + lambda: { + "lazy_ir_sysinc": [ + f"#include <{path}>" + for path in [ + "ATen/core/Formatting.h", + "c10/core/ScalarType.h", + "torch/csrc/lazy/core/hash.h", + "torch/csrc/lazy/core/ir.h", + "torch/csrc/lazy/core/shape.h", + "optional", + "vector", + ] + ], + "lazy_ir_inc": [f'#include "{node_base_hdr}"'] + if node_base_hdr is not None + else [], + "ir_declarations": list( + concat_map_codegen( + lazy_ir_obj, grouped_native_functions, full_codegen + ir_gen + ) + ), + "namespace_prologue": ns_helper.prologue, + "namespace_epilogue": ns_helper.epilogue, + }, + ) + + # Generate Non Native IR Node classes + fm.write_with_template( + "LazyNonNativeIr.h", + "LazyNonNativeIr.h", + lambda: { + "lazy_non_native_ir_inc": [ + f"#include <{path}>" + for path in [ + "torch/csrc/lazy/core/ir.h", + "torch/csrc/lazy/core/ir_builder.h", + "torch/csrc/lazy/core/internal_ops/ltc_ops.h", + "torch/csrc/lazy/core/shape_inference.h", + ] + + ([node_base_hdr] if node_base_hdr else []) + if path + ], + "non_native_ir_nodes": dest.generate_non_native_lazy_ir_nodes( + non_native, lazy_ir_obj + ), + "namespace_prologue": ns_helper.prologue, + "namespace_epilogue": ns_helper.epilogue, + }, + ) + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_schema_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_schema_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..45ffffba2e07b8f6eaf6d6a549c05bcf4bce2da0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_schema_utils.py @@ -0,0 +1,98 @@ +from typing import Any + +from torchgen.model import ( + Annotation, + Argument, + Arguments, + BaseOperatorName, + BaseTy, + BaseType, + CustomClassType, + FunctionSchema, + ListType, + OperatorName, + Return, +) + + +# Note: These aren't actually used in torchgen, they're some utilities for generating a schema +# from real arguments. For example, this is used to generate HigherOrderOperators' schema since +# their schemas can vary for different instances of the same HOP. + + +class TypeGen: + convert_to_base_ty = { + int: BaseTy.int, + float: BaseTy.float, + str: BaseTy.str, + bool: BaseTy.bool, + } + + @staticmethod + def from_example(obj: Any) -> BaseType | ListType | CustomClassType: + import torch + + if isinstance(obj, torch.fx.GraphModule): + return BaseType(BaseTy.GraphModule) + elif isinstance(obj, torch.Tensor): + return BaseType(BaseTy.Tensor) + elif isinstance(obj, torch.SymInt): + return BaseType(BaseTy.SymInt) + elif isinstance(obj, torch.SymBool): + return BaseType(BaseTy.SymBool) + elif isinstance(obj, torch.ScriptObject): + return CustomClassType(obj._type().name()) # type: ignore[attr-defined] + elif isinstance(obj, (list, tuple)): + if len(obj) == 0: + raise AssertionError("list/tuple must be non-empty") + all_base_tys = [TypeGen.from_example(x) for x in obj] + if len(set(all_base_tys)) > 1: + raise RuntimeError( + f"Cannot generate schema for a sequence of args of heterogeneous types: {all_base_tys}. " + "Consider unpacking the argument and give proper names to them if possible " + "instead of using *args." + ) + return ListType(all_base_tys[0], len(obj)) + tp = type(obj) + if tp not in TypeGen.convert_to_base_ty: + raise RuntimeError(f"unsupported type {tp}") + return BaseType(TypeGen.convert_to_base_ty[tp]) + + +class ReturnGen: + @staticmethod + def from_example( + name: str | None, obj: Any, annotation: Annotation | None + ) -> Return: + return Return(name, TypeGen.from_example(obj), annotation) + + +class ArgumentGen: + @staticmethod + def from_example( + name: str, obj: Any, default: str | None, annotation: Annotation | None + ) -> Argument: + return Argument( + name, TypeGen.from_example(obj), default=default, annotation=annotation + ) + + +class FunctionSchemaGen: + @staticmethod + def from_example( + op_name: str, + example_inputs: tuple[tuple[str, Any], ...], + example_outputs: tuple[Any, ...], + ) -> FunctionSchema: + args = [] + for name, inp in example_inputs: + args.append(ArgumentGen.from_example(name, inp, None, None)) + # ignore the annotations and other attributes for now, we could add more when needed. + arguments = Arguments( + tuple(), None, tuple(args), tuple(), None, tuple(), tuple() + ) + returns = tuple( + ReturnGen.from_example(None, out, None) for out in example_outputs + ) + op_name = OperatorName(BaseOperatorName(op_name, False, False, False), "") + return FunctionSchema(op_name, arguments, returns) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_vmap_plumbing.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_vmap_plumbing.py new file mode 100644 index 0000000000000000000000000000000000000000..16582db2db877dd264b4beeb2da45621becdea2c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/gen_vmap_plumbing.py @@ -0,0 +1,277 @@ +from __future__ import annotations + +import textwrap +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from torchgen.api.translate import translate +from torchgen.api.types import DispatcherSignature +from torchgen.context import method_with_native_function +from torchgen.model import ( + Argument, + BaseTy, + BaseType, + FunctionSchema, + ListType, + NativeFunction, + OptionalType, + Return, + SchemaKind, + Type, +) +from torchgen.utils import mapMaybe + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +def is_tensor(typ: Type) -> bool: + return isinstance(typ, BaseType) and typ.name == BaseTy.Tensor + + +def is_optional_tensor(typ: Type) -> bool: + return isinstance(typ, OptionalType) and is_tensor(typ.elem) + + +def is_tensor_list(typ: Type) -> bool: + return isinstance(typ, ListType) and is_tensor(typ.elem) + + +def unwrap_tensor(name: str, cur_level_var: str) -> list[str]: + result = f"""\ + auto [{name}_value, {name}_bdim] = unwrapTensorAtLevel({name}, {cur_level_var});""" + return textwrap.dedent(result).split("\n") + + +def unwrap_optional_tensor(name: str, cur_level_var: str) -> list[str]: + result = f"""\ + std::optional {name}_value; + std::optional {name}_bdim; + if ({name}) {{ + std::tie({name}_value, {name}_bdim) = unwrapTensorAtLevel({name}.value(), {cur_level_var}); + }}""" + return textwrap.dedent(result).split("\n") + + +def gen_unwraps( + flat_arguments: Sequence[Argument], cur_level_var: str +) -> tuple[str, list[str]]: + arg_names = [a.name for a in flat_arguments] + arg_types = [a.type for a in flat_arguments] + + tensors = [name for typ, name in zip(arg_types, arg_names) if is_tensor(typ)] + optional_tensors = [ + name for typ, name in zip(arg_types, arg_names) if is_optional_tensor(typ) + ] + + unwraps = [] + for tensor in tensors: + unwraps += unwrap_tensor(tensor, cur_level_var) + + for opt_tensor in optional_tensors: + unwraps += unwrap_optional_tensor(opt_tensor, cur_level_var) + unwrap_code = "\n".join(unwraps) + + unwrapped_arg_list = [] + for arg in arg_names: + if arg in tensors or arg in optional_tensors: + unwrapped_arg_list += [f"{arg}_value", f"{arg}_bdim"] + else: + unwrapped_arg_list.append(arg) + return unwrap_code, unwrapped_arg_list + + +def gen_case_where_all_bdims_are_none( + outer_sig: DispatcherSignature, schema: FunctionSchema, cur_level_var: str +) -> str: + conditions = [] + flat_args = schema.arguments.flat_all + for arg in flat_args: + if not arg.type.is_tensor_like(): + continue + conditions.append(f"!isBatchedAtLevel({arg.name}, {cur_level_var})") + + sig = DispatcherSignature.from_schema(schema) + translated_args = ", ".join( + e.expr for e in translate(outer_sig.arguments(), sig.arguments()) + ) + return f"""\ +if ({" && ".join(conditions)}) {{ + return at::_ops::{sig.func.name.unambiguous_name()}::call({translated_args}); +}}""" + + +def gen_returns( + returns: tuple[Return, ...], cur_level_var: str, results_var: str +) -> str: + idx = 0 + wrapped_returns = [] + for ret in returns: + if is_tensor(ret.type): + wrapped_returns.append( + f"makeBatched(std::get<{idx}>({results_var}), std::get<{idx + 1}>({results_var}), {cur_level_var})" + ) + idx += 2 + elif is_tensor_list(ret.type): + wrapped_returns.append( + f"makeBatchedVector(std::get<{idx}>({results_var}), std::get<{idx + 1}>({results_var}), {cur_level_var})" + ) + idx += 2 + else: + wrapped_returns.append(f"std::get<{idx}>({results_var})") + idx += 1 + if len(wrapped_returns) == 1: + result = f"return {wrapped_returns[0]};" + else: + result = f"return std::make_tuple({', '.join(wrapped_returns)});" + return result + + +def accepts_at_least_one_tensor_input(schema: FunctionSchema) -> bool: + return any(a.type.is_tensor_like() for a in schema.arguments.flat_all) + + +def is_mutated_arg(argument: Argument) -> bool: + return argument.annotation is not None and argument.annotation.is_write + + +def gen_vmap_inplace_plumbing(native_function: NativeFunction) -> str | None: + # Assumptions: + # - only one argument is being modified in-place + # - the argument that is being modified in-place is the first argument + # - all returns are either Tensor, tuple of Tensor, or TensorList + schema = native_function.func + sig = DispatcherSignature.from_schema(schema) + returns = schema.returns + + # Check assumptions. If these are invalid we return None + # and punt the work to handle them to the future. + if schema.kind() != SchemaKind.inplace: + raise AssertionError(f"Expected inplace schema, got {schema.kind()}") + if not is_mutated_arg(schema.arguments.flat_all[0]): + return None + if len([arg for arg in schema.arguments.flat_all if is_mutated_arg(arg)]) != 1: + return None + + # Only support cases where all returns are Tensors or vector + if len(returns) == 0: + return None + if not all(is_tensor(ret.type) or is_tensor_list(ret.type) for ret in returns): + return None + if not accepts_at_least_one_tensor_input(schema): + return None + + cur_level_var = "cur_level" + + unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var) + bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var) + + return f"""\ +template +{sig.decl(name=schema.name.unambiguous_name() + "_generated_plumbing")} {{ + c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched); + auto maybe_layer = maybeCurrentDynamicLayer(); + vmap_check_escaped(maybe_layer, "gen_vmap_inplace_plumbing"); + int64_t {cur_level_var} = maybe_layer->layerId(); +{textwrap.indent(bdims_all_none_case, " ")} +{textwrap.indent(unwraps, " ")} + batch_rule({", ".join(unwrapped_arg_list)}); + return {schema.arguments.flat_all[0].name}; +}}""" + + +def gen_vmap_plumbing_no_returns(native_function: NativeFunction) -> str: + schema = native_function.func + sig = DispatcherSignature.from_schema(schema) + cur_level_var = "cur_level" + + unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var) + bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var) + + return f"""\ +template +{sig.decl(name=schema.name.unambiguous_name() + "_generated_plumbing")} {{ + c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched); + auto maybe_layer = maybeCurrentDynamicLayer(); + vmap_check_escaped(maybe_layer, "gen_vmap_plumbing_no_returns"); + int64_t {cur_level_var} = maybe_layer->layerId(); +{textwrap.indent(bdims_all_none_case, " ")} +{textwrap.indent(unwraps, " ")} + batch_rule({", ".join(unwrapped_arg_list)}); +}}""" + + +def gen_vmap_plumbing(native_function: NativeFunction) -> str | None: + schema = native_function.func + sig = DispatcherSignature.from_schema(schema) + returns = schema.returns + + # Only support cases where all returns are Tensors or vector + if not accepts_at_least_one_tensor_input(schema): + return None + if len(returns) == 0: + return gen_vmap_plumbing_no_returns(native_function) + return_symint_overrides = [ + "_scaled_dot_product_flash_attention", + "_scaled_dot_product_cudnn_attention", + "_scaled_dot_product_flash_attention_quantized", + ] + if ( + not all(ret.type.is_tensor_like() for ret in returns) + and schema.name.unambiguous_name() not in return_symint_overrides + ): + return None + # in-place views need special handling + if "inplace_view" in native_function.tags: + return None + + if schema.kind() == SchemaKind.inplace: + return gen_vmap_inplace_plumbing(native_function) + + # Don't support these (mutable, out, scratch) + if schema.kind() != SchemaKind.functional: + return None + + results_var = "results" + cur_level_var = "cur_level" + + unwraps, unwrapped_arg_list = gen_unwraps(schema.arguments.flat_all, cur_level_var) + bdims_all_none_case = gen_case_where_all_bdims_are_none(sig, schema, cur_level_var) + + wrapped_returns = gen_returns(returns, cur_level_var, results_var) + return f"""\ +template +{sig.decl(name=schema.name.unambiguous_name() + "_generated_plumbing")} {{ + c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchBatched); + auto maybe_layer = maybeCurrentDynamicLayer(); + vmap_check_escaped(maybe_layer, "gen_vmap_plumbing"); + int64_t {cur_level_var} = maybe_layer->layerId(); +{textwrap.indent(bdims_all_none_case, " ")} +{textwrap.indent(unwraps, " ")} + auto {results_var} = batch_rule({", ".join(unwrapped_arg_list)}); + {wrapped_returns} +}}""" + + +@dataclass(frozen=True) +class ComputeBatchRulePlumbing: + @method_with_native_function + def __call__(self, f: NativeFunction) -> str | None: + result = gen_vmap_plumbing(f) + return result + + +def gen_all_vmap_plumbing(native_functions: Sequence[NativeFunction]) -> str: + body = "\n".join(list(mapMaybe(ComputeBatchRulePlumbing(), native_functions))) + return f""" +#pragma once +#include +#include + +namespace at {{ namespace functorch {{ + +{body} + +}}}} // namespace at::functorch +""" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/local.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/local.py new file mode 100644 index 0000000000000000000000000000000000000000..d2f9965970009b10977bdb69849dea20d483b007 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/local.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import threading +from contextlib import contextmanager +from typing import TYPE_CHECKING + + +if TYPE_CHECKING: + from collections.abc import Iterator + + +# Simple dynamic scoping implementation. The name "parametrize" comes +# from Racket. +# +# WARNING WARNING: LOOKING TO EDIT THIS FILE? Think carefully about +# why you need to add a toggle to the global behavior of code +# generation. The parameters here should really only be used +# for "temporary" situations, where we need to temporarily change +# the codegen in some cases because we cannot conveniently update +# all call sites, and are slated to be eliminated once all call +# sites are eliminated. If you don't have a plan for how to get there, +# DON'T add a new entry here. + + +class Locals(threading.local): + use_const_ref_for_mutable_tensors: bool | None = None + use_ilistref_for_tensor_lists: bool | None = None + + +_locals = Locals() + + +def use_const_ref_for_mutable_tensors() -> bool: + if _locals.use_const_ref_for_mutable_tensors is None: + raise AssertionError( + "need to initialize local.use_const_ref_for_mutable_tensors with " + "local.parametrize" + ) + return _locals.use_const_ref_for_mutable_tensors + + +def use_ilistref_for_tensor_lists() -> bool: + if _locals.use_ilistref_for_tensor_lists is None: + raise AssertionError( + "need to initialize local.use_ilistref_for_tensor_lists with local.parametrize" + ) + return _locals.use_ilistref_for_tensor_lists + + +@contextmanager +def parametrize( + *, use_const_ref_for_mutable_tensors: bool, use_ilistref_for_tensor_lists: bool +) -> Iterator[None]: + old_use_const_ref_for_mutable_tensors = _locals.use_const_ref_for_mutable_tensors + old_use_ilistref_for_tensor_lists = _locals.use_ilistref_for_tensor_lists + try: + _locals.use_const_ref_for_mutable_tensors = use_const_ref_for_mutable_tensors + _locals.use_ilistref_for_tensor_lists = use_ilistref_for_tensor_lists + yield + finally: + _locals.use_const_ref_for_mutable_tensors = ( + old_use_const_ref_for_mutable_tensors + ) + _locals.use_ilistref_for_tensor_lists = old_use_ilistref_for_tensor_lists diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/model.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/model.py new file mode 100644 index 0000000000000000000000000000000000000000..54aeaab9fb9d5bfd8036677d040847e97e1ec60a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/model.py @@ -0,0 +1,3086 @@ +from __future__ import annotations + +import dataclasses +import itertools +import re +from dataclasses import dataclass +from enum import auto, Enum +from typing import TYPE_CHECKING +from typing_extensions import assert_never + +from torchgen.utils import NamespaceHelper, OrderedSet + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterator, Sequence + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# DATA MODEL +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Some general principles for our data model. +# +# - Stop using C++ data types as the internal data representation +# format. Instead, the internal data structures are centered +# around JIT schema representation. This avoid a big problem +# with the old codegen where we read in all the types from +# native_functions.yaml and then immediately had to retranslate +# them into C++ types. +# +# - More semantic data representation. Instead of representing +# everything as dicts and strings, we define dataclasses for +# every interesting entity the code generation has to deal with. +# These dataclasses have strong semantic invariants: for example, +# we generally require them to roundtrip losslessly into the +# form they were parsed from. These structures are immutable +# and you're expected to populate information once during +# construction. + + +# Represent a source location; used for better error reporting +@dataclass(frozen=True) +class Location: + file: str + line: int + + def __str__(self) -> str: + return f"{self.file}:{self.line}" + + +# Valid values of the 'variants' field in native_functions.yaml +class Variant(Enum): + function = auto() + method = auto() + + +# Default kernel namespace +DEFAULT_KERNEL_NAMESPACE = "at::native" + +# NOTE: Keep the list in sync with `DispatchKey` in c10/core/DispatchKey.h +BACKEND_COMPONENTS = [ + "CPU", + "CUDA", + "HIP", + "XLA", + "MTIA", + "MPS", + "IPU", + "XPU", + "HPU", + "VE", + "Lazy", + "Meta", + "PrivateUse1", + "PrivateUse2", + "PrivateUse3", +] +FUNCTIONALITY_KEYS = [ + "", + "Quantized", + "Sparse", + "SparseCsr", + "NestedTensor", + "Autograd", +] + +# This list guards dispatches that can be used in derivatives.yaml +# For now we omit AutogradFunctionality and AutogradOther +AUTOGRAD_KEYS = ["AutogradNestedTensor"] + [ + "Autograd" + component for component in BACKEND_COMPONENTS +] + +FRAGMENT_NAMESPACES = {"quantized", "quantized_decomposed"} + + +# This doesn't have to be in sync with the header, it only needs to contain +# entries that we actually use in the codegen or want pyi entries for +class DispatchKey(Enum): + Undefined = 0 + CatchAll = Undefined + + FPGA = auto() + MAIA = auto() + Vulkan = auto() + Metal = auto() + MKLDNN = auto() + OpenGL = auto() + OpenCL = auto() + IDEEP = auto() + CustomRNGKeyId = auto() + MkldnnCPU = auto() + Sparse = auto() + SparseCsr = auto() + NestedTensor = auto() + Dense = auto() + + PythonTLSSnapshot = auto() + PreDispatch = auto() + PythonDispatcher = auto() + Python = auto() + FuncTorchDynamicLayerBackMode = auto() + ZeroTensor = auto() + Conjugate = auto() + Negative = auto() + BackendSelect = auto() + Named = auto() + AutogradOther = auto() + AutogradFunctionality = auto() + AutogradNestedTensor = auto() + Tracer = auto() + Autocast = auto() + AutocastCPU = auto() + AutocastCUDA = auto() + Batched = auto() + VmapMode = auto() + FuncTorchGradWrapper = auto() + FuncTorchBatched = auto() + BatchedNestedTensor = auto() + FuncTorchVmapMode = auto() + FuncTorchDynamicLayerFrontMode = auto() + Functionalize = auto() + TESTING_ONLY_GenericWrapper = auto() + TESTING_ONLY_GenericMode = auto() + + ADInplaceOrView = auto() + Autograd = auto() + CompositeImplicitAutograd = auto() + CompositeImplicitAutogradNestedTensor = auto() + CompositeExplicitAutograd = auto() + CompositeExplicitAutogradNonFunctional = auto() + FuncTorchBatchedDecomposition = auto() + + # BEGIN autogenerated + CPU = auto() + CUDA = auto() + HIP = auto() + XLA = auto() + MTIA = auto() + MPS = auto() + IPU = auto() + XPU = auto() + HPU = auto() + VE = auto() + Lazy = auto() + Meta = auto() + PrivateUse1 = auto() + PrivateUse2 = auto() + PrivateUse3 = auto() + QuantizedCPU = auto() + QuantizedCUDA = auto() + QuantizedHIP = auto() + QuantizedXLA = auto() + QuantizedMTIA = auto() + QuantizedMPS = auto() + QuantizedIPU = auto() + QuantizedXPU = auto() + QuantizedHPU = auto() + QuantizedVE = auto() + QuantizedLazy = auto() + QuantizedMeta = auto() + QuantizedPrivateUse1 = auto() + QuantizedPrivateUse2 = auto() + QuantizedPrivateUse3 = auto() + SparseCPU = auto() + SparseCUDA = auto() + SparseHIP = auto() + SparseXLA = auto() + SparseMTIA = auto() + SparseMPS = auto() + SparseIPU = auto() + SparseXPU = auto() + SparseHPU = auto() + SparseVE = auto() + SparseLazy = auto() + SparseMeta = auto() + SparsePrivateUse1 = auto() + SparsePrivateUse2 = auto() + SparsePrivateUse3 = auto() + SparseCsrCPU = auto() + SparseCsrCUDA = auto() + SparseCsrHIP = auto() + SparseCsrXLA = auto() + SparseCsrMTIA = auto() + SparseCsrMPS = auto() + SparseCsrIPU = auto() + SparseCsrXPU = auto() + SparseCsrHPU = auto() + SparseCsrVE = auto() + SparseCsrLazy = auto() + SparseCsrMeta = auto() + SparseCsrPrivateUse1 = auto() + SparseCsrPrivateUse2 = auto() + SparseCsrPrivateUse3 = auto() + NestedTensorCPU = auto() + NestedTensorCUDA = auto() + NestedTensorHIP = auto() + NestedTensorXLA = auto() + NestedTensorMTIA = auto() + NestedTensorMPS = auto() + NestedTensorIPU = auto() + NestedTensorXPU = auto() + NestedTensorHPU = auto() + NestedTensorVE = auto() + NestedTensorLazy = auto() + NestedTensorMeta = auto() + NestedTensorPrivateUse1 = auto() + NestedTensorPrivateUse2 = auto() + NestedTensorPrivateUse3 = auto() + AutogradCPU = auto() + AutogradCUDA = auto() + AutogradHIP = auto() + AutogradXLA = auto() + AutogradMTIA = auto() + AutogradMPS = auto() + AutogradIPU = auto() + AutogradXPU = auto() + AutogradHPU = auto() + AutogradVE = auto() + AutogradLazy = auto() + AutogradMeta = auto() + AutogradPrivateUse1 = auto() + AutogradPrivateUse2 = auto() + AutogradPrivateUse3 = auto() + # END autogenerated + + def __str__(self) -> str: + return self.name + + def lower(self) -> str: + return str(self).lower() + + @staticmethod + def parse(value: str) -> DispatchKey: + for k, v in DispatchKey.__members__.items(): + if k == value: + return v + raise AssertionError(f"unknown dispatch key {value}") + + +class _TorchDispatchModeKey(Enum): + FAKE = auto() + PROXY = auto() + FUNCTIONAL = auto() + + +def codegen_per_backend_entries() -> str: + r: list[str] = [] + for fk in FUNCTIONALITY_KEYS: + r.extend(f" {fk}{bc} = auto()" for bc in BACKEND_COMPONENTS) + return "\n".join(r) + + +for fk in FUNCTIONALITY_KEYS: + for bc in BACKEND_COMPONENTS: + if not hasattr(DispatchKey, fk + bc): + r = codegen_per_backend_entries() + print(r) + raise RuntimeError( + f"Missing {fk}{bc} from DispatchKey enum. Here is the autogenerated list we expect to have:\n\n{r}" + ) + + +STRUCTURED_DISPATCH_KEYS = { + DispatchKey.MPS, + DispatchKey.CUDA, + DispatchKey.CPU, + DispatchKey.XPU, + DispatchKey.MTIA, +} +UFUNC_DISPATCH_KEYS = {DispatchKey.CUDA, DispatchKey.CPU} + +# Set of supported dispatch keys +dispatch_keys = [ + DispatchKey.CPU, + DispatchKey.SparseCPU, + DispatchKey.SparseCsrCPU, + DispatchKey.MkldnnCPU, + DispatchKey.CUDA, + DispatchKey.MPS, + DispatchKey.XPU, + DispatchKey.SparseXPU, + DispatchKey.SparseCsrXPU, + DispatchKey.SparseCUDA, + DispatchKey.SparseCsrCUDA, + DispatchKey.SparseMPS, + DispatchKey.SparseCsrMPS, + DispatchKey.QuantizedCPU, + DispatchKey.QuantizedCUDA, + DispatchKey.CompositeImplicitAutograd, + DispatchKey.CompositeImplicitAutogradNestedTensor, + DispatchKey.CompositeExplicitAutograd, + DispatchKey.CompositeExplicitAutogradNonFunctional, + DispatchKey.NestedTensorCPU, + DispatchKey.NestedTensorCUDA, + DispatchKey.NestedTensorXPU, + DispatchKey.NestedTensorHPU, + # Meta is a magic key: it is automatically generated for structured + # kernels + DispatchKey.Meta, + DispatchKey.SparseMeta, + DispatchKey.SparseCsrMeta, + DispatchKey.QuantizedMeta, + DispatchKey.NestedTensorMeta, + DispatchKey.ZeroTensor, + DispatchKey.MTIA, +] + + +# Dispatch keys that "support all backends". These codegen slightly differently +# then backend specific keys. +def is_generic_dispatch_key(dk: DispatchKey) -> bool: + return dk in { + DispatchKey.CompositeExplicitAutograd, + DispatchKey.CompositeExplicitAutogradNonFunctional, + DispatchKey.CompositeImplicitAutograd, + DispatchKey.CompositeImplicitAutogradNestedTensor, + } + + +# CUDA specific dispatch keys +def is_cuda_dispatch_key(dk: DispatchKey) -> bool: + return dk in { + DispatchKey.CUDA, + DispatchKey.QuantizedCUDA, + DispatchKey.SparseCUDA, + DispatchKey.SparseCsrCUDA, + DispatchKey.NestedTensorCUDA, + DispatchKey.AutogradCUDA, + } + + +# XPU specific dispatcy keys +def is_xpu_dispatch_key(dk: DispatchKey) -> bool: + return dk in { + DispatchKey.XPU, + DispatchKey.QuantizedXPU, + DispatchKey.SparseXPU, + DispatchKey.SparseCsrXPU, + DispatchKey.NestedTensorXPU, + DispatchKey.AutogradXPU, + } + + +# Structured kernel generation is only supported for certain key types; +# otherwise use old-style +def is_structured_dispatch_key(dk: DispatchKey) -> bool: + return dk in STRUCTURED_DISPATCH_KEYS + + +def is_ufunc_dispatch_key(dk: DispatchKey) -> bool: + # For now, ufunc dispatch keys coincide with structured keys + return dk in UFUNC_DISPATCH_KEYS + + +dispatch_device_map = {is_cuda_dispatch_key: "cuda", is_xpu_dispatch_key: "xpu"} + + +# This is oddly named ScalarType and not DType for symmetry with C++ +class ScalarType(Enum): + Byte = auto() + Char = auto() + Short = auto() + Int = auto() + Long = auto() + Half = auto() + Float = auto() + Double = auto() + ComplexHalf = auto() + ComplexFloat = auto() + ComplexDouble = auto() + Bool = auto() + BFloat16 = auto() + Float8_e5m2 = auto() + Float8_e5m2fnuz = auto() + Float8_e4m3fn = auto() + Float8_e4m3fnuz = auto() + Float8_e8m0fnu = auto() + + def __str__(self) -> str: + return self.name + + @staticmethod + def maybe_parse(value: str) -> ScalarType | None: + for k, v in ScalarType.__members__.items(): + if k == value: + return v + return None + + @staticmethod + def parse(value: str) -> ScalarType: + mb_r = ScalarType.maybe_parse(value) + if mb_r is None: + raise AssertionError(f"unknown dtype {value}") + return mb_r + + @staticmethod + def parse_set(values: str) -> OrderedSet[ScalarType]: + dtypes: OrderedSet[ScalarType] = OrderedSet() + for value in values.split(", "): + if value in DTYPE_CLASSES: + dtypes.update(DTYPE_CLASSES[value]) + else: + dtypes.add(ScalarType.parse(value)) + return dtypes + + +DTYPE_CLASSES: dict[str, OrderedSet[ScalarType]] = {} +# NB: Integral doesn't include boolean +DTYPE_CLASSES["Integral"] = OrderedSet( + [ + ScalarType.Byte, + ScalarType.Char, + ScalarType.Int, + ScalarType.Long, + ScalarType.Short, + ] +) +# NB: Floating doesn't include low precision types +DTYPE_CLASSES["Floating"] = OrderedSet([ScalarType.Float, ScalarType.Double]) +DTYPE_CLASSES["Complex"] = OrderedSet( + [ScalarType.ComplexFloat, ScalarType.ComplexDouble] +) +DTYPE_CLASSES["All"] = DTYPE_CLASSES["Integral"] | DTYPE_CLASSES["Floating"] +DTYPE_CLASSES["AllAndComplex"] = DTYPE_CLASSES["All"] | DTYPE_CLASSES["Complex"] +DTYPE_CLASSES["FloatingAndComplex"] = ( + DTYPE_CLASSES["Floating"] | DTYPE_CLASSES["Complex"] +) + + +# Represents the valid entries for ufunc_inner_loop in native_functions.yaml. +# NB: if you add a new UfuncKey, you will teach torchgen.dest.ufunc how +# to process it. Most logic will ignore keys they don't understand, so your +# new key will get silently ignored until you hook in logic to deal with it. +class UfuncKey(Enum): + # These are low level keys that represent exactly one particular + # instantiation of the kernel produced by codegen + CUDAFunctor = auto() + CUDAFunctorOnOther = auto() + CUDAFunctorOnSelf = auto() + + CPUScalar = auto() + CPUVector = auto() + + # These are the ones users will usually specify, and + # implicitly "fill in" the low level keys + ScalarOnly = auto() # CUDA*, CPUScalar + Generic = auto() # CUDA*, CPU* + + def __str__(self) -> str: + return self.name + + @staticmethod + def parse(value: str) -> UfuncKey: + for k, v in UfuncKey.__members__.items(): + if k == value: + return v + raise AssertionError(f"unknown ufunc key {value}") + + +class DeviceCheckType(Enum): + NoCheck = 0 + ExactSame = 1 + + +class ViewSchemaKind(Enum): + aliasing = auto() + aliasing_inplace = auto() + non_aliasing = auto() + + +# The basic input to the code generation is native_functions.yaml. +# The name "native", BTW, comes from the distinction between native +# functions and legacy TH functions. The legacy TH functions are gone, +# but the "native" descriptor has stuck. +# +# NativeFunction models a single entry in native_functions.yaml. Its +# fields roughly correspond to what you would see in the YAML itself, +# but after canonicalization and parsing has occurred. +# +# You can see some of the overall design patterns for how we setup +# dataclasses in this class, but we will defer a complete discussion +# of this at FunctionSchema. +@dataclass(frozen=True) +class NativeFunction: + # The namespace for this operator. For example, if we have "at::add" + # then the namespace would be "at". This enables ops to be registered + # through the same DSL with a custom namespace. If not specified, the + # default namespace would be "at". + namespace: str + + # The function schema of the operator in question. This schema + # has been parsed; see FunctionSchema for more about its structure. + # (This type is quoted as we are forward referencing a type + # defined later in the file. I opted for this ordering of the + # classes for expository clarity.) + func: FunctionSchema + + # Whether or not to generate mutable tensor arguments like regular + # ones + use_const_ref_for_mutable_tensors: bool + + # Whether or not to omit automatic generation of a DeviceGuard + device_guard: bool + + # How to emit automatic generation of device check + device_check: DeviceCheckType + + # What python module to put the function in + python_module: str | None + + # TODO: figure out what this does + category_override: str | None + + # If no variants are specified in native_functions.yaml, this is + # assumed to be {'function'}. + variants: set[Variant] + + # Whether or not we should skip generating registrations for + # this kernel. This is a bit of a double-edged sword, as manual + # registrations don't participate in codegen-based selective build! + manual_kernel_registration: bool + + # Whether or not to skip generating TensorMethod/Functions bindings + # for this kernel. Technically, this doesn't actually skip generating + # the binding; instead, the binding gets generated to __dispatch_{funcname} + # so you can make use of the normal binding if you need it. + manual_cpp_binding: bool + + # The location in the YAML file were this native function entry was + # defined. This is for conveniently reporting error messages! + loc: Location + + # A list of operators that are expected to be auto-generated for this NativeFunction. + # Note: This list isn't actually directly used by the codegen to generate anything. + # Instead, the codegen figures out what operators to generate purely based off of + # function schema, and uses the autogen declarations to error check. + # We expect every NativeFunction that gets auto-generated be explicitly called out + # in native_functions.yaml + autogen: list[OperatorName] + + # If non-empty, this kernel is subject to ufunc codegen. + # Sorted by ufunc_key + ufunc_inner_loop: dict[UfuncKey, UfuncInnerLoop] + + # Whether or not this out functions is a "structured kernel". Structured + # kernels are defined a little differently from normal kernels; in + # particular, their shape checking logic is defined separately from + # the kernel. Only out functions can be structured; other functions + # delegate to the out function using the structured_delegate keyword. + # Every structured kernel must have at least an out and a functional + # variant. + structured: bool + + # Whether or not this non-out function is a structured kernel, defined + # in terms of the out kernel referenced by the string here. + structured_delegate: OperatorName | None + + # Only valid for structured kernels. Specifies alternative of what + # to inherit from when defining the meta class for the structured + # operator. This will usually be TensorIteratorBase. This also + # changes the semantics of set_output to call the parent class. + structured_inherits: str | None + + # Structured kernels can declare elements as "precomputed". These elements + # are returned by the meta function in one struct and passed to the impl + # function in lieu of certain kernel arguments that these precomputed + # elements supersede. Information about the names and types of these + # precomputed elements and how they correspond to kernel arguments is stored + # in this member, if applicable. + precomputed: Precompute | None + + # Argument names whose default should be excluded from the C++ interface. + # Intended for resolving overload ambiguities between signatures. + cpp_no_default_args: set[str] + + # Note [Abstract ATen methods] + # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + # An abstract ATen method is one whose dispatch differs between + # types. These are implemented in derived types (with a + # standard (throwing) definition in Type). A concrete ATen + # method is one which has the same dispatch for all types; + # we just implement it in the base Type. This is exposed + # in Declarations.yaml via a field named 'abstract'. + is_abstract: bool + + # Whether or not the NativeFunction contains a backend-agnostic kernel + has_composite_implicit_autograd_kernel: bool + has_composite_implicit_autograd_nested_tensor_kernel: bool + has_composite_explicit_autograd_kernel: bool + has_composite_explicit_autograd_non_functional_kernel: bool + + # Tags are used to describe semantic information about (groups of) operators, + # That aren't easily inferable directly from the operator's schema. + tags: set[str] + + # NB: The benefit of defining a dataclass is that we automatically get + # a constructor defined for all the fields we specify. No need + # to explicitly write it out. + + # We parse both the NativeFunction + backend-specific information about it, which it stored in a corresponding BackendIndex. + @staticmethod + def from_yaml( + ei: dict[str, object], + loc: Location, + valid_tags: set[str], + ignore_keys: set[DispatchKey] | None = None, + ) -> tuple[NativeFunction, dict[DispatchKey, dict[OperatorName, BackendMetadata]]]: + """ + Parse a NativeFunction from a dictionary as directly parsed + from native_functions.yaml + """ + e = ei.copy() + + funcs = e.pop("func") + if not isinstance(funcs, str): + raise AssertionError(f"not a str: {funcs}") + # only support one level of namespace. E.g., aten::add + namespace_helper = NamespaceHelper.from_namespaced_entity( + namespaced_entity=funcs, max_level=1 + ) + namespace = namespace_helper.get_cpp_namespace(default="aten") + func = FunctionSchema.parse(namespace_helper.entity_name) + + cpp_no_default_args_list = e.pop("cpp_no_default_args", []) + if not isinstance(cpp_no_default_args_list, list): + raise AssertionError( + f"cpp_no_default_args is not a list: {cpp_no_default_args_list}" + ) + cpp_no_default_args = set(cpp_no_default_args_list) + + use_const_ref_for_mutable_tensors = e.pop( + "use_const_ref_for_mutable_tensors", False + ) + if not isinstance(use_const_ref_for_mutable_tensors, bool): + raise AssertionError( + f"use_const_ref_for_mutable_tensors is not a bool: {use_const_ref_for_mutable_tensors}" + ) + + if use_const_ref_for_mutable_tensors: + if func.arguments.out: + raise AssertionError( + "see https://github.com/pytorch/pytorch/issues/145522" + ) + + variants_s = e.pop("variants", "function") + if not isinstance(variants_s, str): + raise AssertionError(f"variants is not a str: {variants_s}") + variants: set[Variant] = set() + for v in variants_s.split(", "): + if v == "function": + variants.add(Variant.function) + elif v == "method": + variants.add(Variant.method) + else: + raise AssertionError(f"illegal variant {v}") + + manual_kernel_registration = e.pop("manual_kernel_registration", False) + if not isinstance(manual_kernel_registration, bool): + raise AssertionError(f"not a bool: {manual_kernel_registration}") + + manual_cpp_binding = e.pop("manual_cpp_binding", False) + if not isinstance(manual_cpp_binding, bool): + raise AssertionError(f"not a bool: {manual_cpp_binding}") + + device_guard = e.pop("device_guard", True) + if not isinstance(device_guard, bool): + raise AssertionError(f"not a bool: {device_guard}") + + device_check_s = e.pop("device_check", None) + if not (device_check_s is None or isinstance(device_check_s, str)): + raise AssertionError(f"not a str: {device_check_s}") + if not ( + device_check_s is None or device_check_s in DeviceCheckType.__members__ + ): + raise AssertionError(f"illegal device_check: {device_check_s}") + device_check: DeviceCheckType + if device_check_s is None: + device_check = DeviceCheckType.ExactSame + else: + device_check = DeviceCheckType[device_check_s] + + structured = e.pop("structured", False) + if not isinstance(structured, bool): + raise AssertionError(f"not a bool: {structured}") + + structured_delegate_s = e.pop("structured_delegate", None) + if not ( + structured_delegate_s is None or isinstance(structured_delegate_s, str) + ): + raise AssertionError(f"not a str: {structured_delegate_s}") + if structured_delegate_s is not None and "::" in structured_delegate_s: + raise AssertionError( + "namespace is not supported in structured delegate," + " using the same namespace as the native function" + ) + structured_delegate: OperatorName | None = None + if structured_delegate_s is not None: + structured_delegate = OperatorName.parse(structured_delegate_s) + + structured_inherits = e.pop("structured_inherits", None) + if not (structured_inherits is None or isinstance(structured_inherits, str)): + raise AssertionError(f"not a str: {structured_inherits}") + if structured_inherits is not None and "::" in structured_inherits: + raise AssertionError( + "namespace is not supported in structured inherits," + " using the same namespace as the native function" + ) + + python_module = e.pop("python_module", None) + if not (python_module is None or isinstance(python_module, str)): + raise AssertionError(f"not a str: {python_module}") + if python_module is not None and Variant.method in variants: + raise AssertionError("functions in modules cannot be methods") + + category_override = e.pop("category_override", None) + if not (category_override is None or isinstance(category_override, str)): + raise AssertionError(f"not a str: {category_override}") + + precomputed_dict = e.pop("precomputed", None) + if precomputed_dict is not None and structured is not True: + raise AssertionError( + f"precomputed requires structured=True, got structured={structured}" + ) + precomputed = Precompute.parse(precomputed_dict) if precomputed_dict else None + + tags_inp = e.pop("tags", []) + if isinstance(tags_inp, str): + tags_inp = [tags_inp] + if not isinstance(tags_inp, list): + raise AssertionError(f"tags is not a list: {tags_inp}") + + # All aten ops generated by torchgen receive the pt2_compliant tag. + if namespace == "aten" and "pt2_compliant_tag" in valid_tags: + tags_inp.append("pt2_compliant_tag") + + tags: set[str] = set() + for t in tags_inp: + if len(valid_tags) == 0: + raise AssertionError("valid_tags is empty") + # TODO: verify that the tag is valid and has an entry in tags.yaml + if t in valid_tags: + tags.add(t) + else: + raise AssertionError(f"illegal tag {t}") + + from torchgen.api import cpp + + raw_dispatch = e.pop("dispatch", None) + if not (raw_dispatch is None or isinstance(raw_dispatch, dict)): + raise AssertionError(f"dispatch is not a dict: {e}") + dispatch: dict[DispatchKey, BackendMetadata] = {} + num_dispatch_keys: int = 0 + if raw_dispatch is not None: + if manual_kernel_registration: + raise AssertionError( + "cannot specify both manual_kernel_registration and dispatch; with " + "manual registration, dispatch has no effect!" + ) + redundant_composite_implicit_autograd = False + for ks, v in raw_dispatch.items(): + if ks == "__line__": + continue # not worth tracking line numbers for dispatch entries + if not isinstance(ks, str): + raise AssertionError( + f"illegal dispatch key '{ks}' in {raw_dispatch}" + ) + if not isinstance(v, str): + raise AssertionError( + f"illegal dispatch value '{v}' in {raw_dispatch}" + ) + for k in ks.split(","): + dispatch_key = DispatchKey.parse(k.strip()) + num_dispatch_keys += 1 + + if ignore_keys and dispatch_key in ignore_keys: + continue + if dispatch_key not in dispatch_keys: + raise AssertionError( + f"Dispatch key {dispatch_key} of kernel {v} " + "is not a supported dispatch key." + ) + # We only allow at most 3 levels of namespace for kernels. + # We will append "native" to a custom kernel namespace. + namespace_helper = NamespaceHelper.from_namespaced_entity( + v, max_level=3 + ) + kernel_namespace = namespace_helper.get_cpp_namespace(default="at") + # Why is 'structured' included? External backends (e.g. + # XLA) opt into which ops are structured independently + # of which in-tree ops are structured + dispatch[dispatch_key] = BackendMetadata( + kernel=namespace_helper.entity_name, + structured=structured + and is_structured_dispatch_key(dispatch_key), + cpp_namespace=(kernel_namespace + "::native"), + ) + if ( + dispatch_key is DispatchKey.CompositeImplicitAutograd + and v == cpp.name(func) + ): + redundant_composite_implicit_autograd = True + + # We count the number of dispatch keys which have not been ignored to prevent a dispatch table + # in which all backend keys are ignored but necessarily kept, remaining compositeimplicit, + # from being treated as redundant. + if num_dispatch_keys == 1 and redundant_composite_implicit_autograd: + raise AssertionError( + "unnecessary dispatch table for this function; just delete the dispatch " + "key entirely" + ) + # if a function is a structured delegate, deleting the dispatch + # table is NOT semantics preserving + if not ( + structured_delegate + or dispatch.keys() != {DispatchKey.CompositeImplicitAutograd} + or dispatch[DispatchKey.CompositeImplicitAutograd].supports_symint() + or num_dispatch_keys != 1 + ): + raise AssertionError( + f"unexpected name for singleton CompositeImplicitAutograd dispatch entry: expected {cpp.name(func)} " + f"but got {dispatch[DispatchKey.CompositeImplicitAutograd]}. Rename your implementation to the expected " + "name, then delete the dispatch table" + ) + elif not structured and structured_delegate is None: + name = str(func.name.name) + if ( + name.startswith("new_") + or name.endswith("_like") + # TODO: maybe it's better to test the return + or ( + func.arguments.tensor_options + and not func.arguments.has_tensor_arg() + ) + ): + raise AssertionError( + f"expected {name} to have a CompositeExplicitAutograd " + "dispatch entry, but there was no dispatch table. Factory functions " + "should not have implicit dispatch as they should not be decomposed " + "for __torch_dispatch__" + ) + dispatch[DispatchKey.CompositeImplicitAutograd] = BackendMetadata( + cpp.name(func), structured=False, cpp_namespace=DEFAULT_KERNEL_NAMESPACE + ) + + composites_in_dispatch = [ + d + for d in dispatch + if d == DispatchKey.CompositeExplicitAutograd + or d == DispatchKey.CompositeExplicitAutogradNonFunctional + or d == DispatchKey.CompositeImplicitAutograd + or d == DispatchKey.CompositeImplicitAutogradNestedTensor + ] + + if not ( + len(composites_in_dispatch) <= 1 + or ( + len(composites_in_dispatch) == 2 + and ( + DispatchKey.CompositeExplicitAutogradNonFunctional + not in composites_in_dispatch + ) + and ( + DispatchKey.CompositeImplicitAutogradNestedTensor + in composites_in_dispatch + ) + ) + ): + raise AssertionError( + "cannot specify more than one of CompositeExplicitAutograd, CompositeExplicitAutogradNonFunctional, " + "or CompositeImplicitAutograd on a single kernel; each " + "strictly subsumes the other. If you wanted to provide an explicit autograd " + "implementation, specify CompositeExplicitAutograd; otherwise specify CompositeImplicitAutograd only" + ) + + autogen_str = e.pop("autogen", "") + if not isinstance(autogen_str, str): + raise AssertionError(f"autogen is not a str: {autogen_str}") + autogen = ( + [] + if autogen_str == "" + else [OperatorName.parse(x) for x in autogen_str.split(", ")] + ) + + raw_ufunc_inner_loop = e.pop("ufunc_inner_loop", {}) + ufunc_inner_loop = {} + if isinstance(raw_ufunc_inner_loop, str): + ufunc_inner_loop[UfuncKey.Generic] = UfuncInnerLoop.parse( + raw_ufunc_inner_loop, UfuncKey.Generic + ) + elif isinstance(raw_ufunc_inner_loop, dict): + for k, vo in raw_ufunc_inner_loop.items(): + if k == "__line__": + continue + if not isinstance(k, str): + raise AssertionError(f"ufunc_inner_loop key is not a str: {k}") + if not isinstance(vo, str): + raise AssertionError(f"ufunc_inner_loop value is not a str: {vo}") + ufunc_key = UfuncKey.parse(k) + ufunc_inner_loop[ufunc_key] = UfuncInnerLoop.parse(vo, ufunc_key) + else: + raise AssertionError( + f"ufunc_inner_loop not str or dict: {raw_ufunc_inner_loop}" + ) + # Program the BackendIndex for the implicit dispatch entry from ufunc + if ufunc_inner_loop: + if not structured: + raise AssertionError("ufunc must be structured") + + # Delay import ufunc here to avoid circular import issue + # See: https://github.com/pytorch/pytorch/issues/81294 + import torchgen.api.ufunc as ufunc + + for dispatch_key in UFUNC_DISPATCH_KEYS: + if dispatch_key in dispatch: + raise AssertionError( + f"ufunc should not have explicit dispatch entry for {dispatch_key}" + ) + dispatch[dispatch_key] = BackendMetadata( + kernel=ufunc.schema_kernel_name(func, dispatch_key), + structured=True, + cpp_namespace=DEFAULT_KERNEL_NAMESPACE, + ) + + if structured_delegate: + # Structured functions MUST have a dispatch table + is_abstract = True + else: + is_abstract = ( + dispatch.keys() != {DispatchKey.CompositeImplicitAutograd} + and dispatch.keys() + != {DispatchKey.CompositeImplicitAutogradNestedTensor} + and dispatch.keys() + != { + DispatchKey.CompositeImplicitAutograd, + DispatchKey.CompositeImplicitAutogradNestedTensor, + } + ) + + has_composite_implicit_autograd_kernel = ( + DispatchKey.CompositeImplicitAutograd in dispatch + ) + has_composite_implicit_autograd_nested_tensor_kernel = ( + DispatchKey.CompositeImplicitAutogradNestedTensor in dispatch + ) + has_composite_explicit_autograd_kernel = ( + DispatchKey.CompositeExplicitAutograd in dispatch + ) + has_composite_explicit_autograd_non_functional_kernel = ( + DispatchKey.CompositeExplicitAutogradNonFunctional in dispatch + ) + + # We aren't going to store dispatch metadata inline in NativeFunctions; + # instead it is separately indexed by backend (so other backends can + # add more dispatch entries after the fact). Reindex the individual + # metadata by OperatorName! + backend_metadata = {k: {func.name: v} for k, v in dispatch.items()} + + # don't care if it exists or not; make it easier to use this function + # with other yaml parsers that aren't setting __line__ in the dict + e.pop("__line__", None) + if e: + raise AssertionError(f"leftover entries: {e}") + + # Asserts that we can't do in post_init, because they rely on backend-specific info + if structured_delegate is not None: + for key in STRUCTURED_DISPATCH_KEYS: + if key in dispatch: + raise AssertionError( + f"if structured_delegate, then must not have {key} in dispatch dictionary " + "(it is delegated!)" + ) + + return ( + NativeFunction( + func=func, + use_const_ref_for_mutable_tensors=use_const_ref_for_mutable_tensors, + variants=variants, + structured=structured, + structured_delegate=structured_delegate, + structured_inherits=structured_inherits, + precomputed=precomputed, + autogen=autogen, + ufunc_inner_loop=ufunc_inner_loop, + manual_kernel_registration=manual_kernel_registration, + manual_cpp_binding=manual_cpp_binding, + python_module=python_module, + category_override=category_override, + device_guard=device_guard, + device_check=device_check, + loc=loc, + cpp_no_default_args=cpp_no_default_args, + is_abstract=is_abstract, + has_composite_implicit_autograd_kernel=has_composite_implicit_autograd_kernel, + has_composite_implicit_autograd_nested_tensor_kernel=has_composite_implicit_autograd_nested_tensor_kernel, + has_composite_explicit_autograd_kernel=has_composite_explicit_autograd_kernel, + has_composite_explicit_autograd_non_functional_kernel=has_composite_explicit_autograd_non_functional_kernel, + tags=tags, + namespace=namespace, + ), + backend_metadata, + ) + + def validate_unstructured(self) -> None: + # TODO: probably better to accumulate these errors and report them all + # at once + if self.structured: + raise AssertionError( + "This function is structured, but there was " + "no valid functional variant of it." + ) + if not self.structured_delegate: + raise AssertionError( + "This function delegates to another structured out function, " + "but no valid function was found (the delegate may not exist, or it has the wrong type)" + ) + + # __post_init__ functions in dataclasses can be used to do extra + # validation after construction. + # + # Notice that we don't do any type validation here. In fact, we + # rely exclusively on mypy to check if you've done types correctly! + # Validation is for nontrivial invariants that cannot be (conveniently) + # encoded in the type system. + def __post_init__(self) -> None: + if self.func.arguments.out: + if self.variants != {Variant.function}: + raise AssertionError( + "Native functions with out arguments MUST " + "be declared with only function variant; e.g., variants: function; " + "otherwise you will tickle a Python argument binding bug " + "(which usually manifests itself as the result variable being undefined.)" + ) + if self.structured: + if self.func.kind() != SchemaKind.out: + raise AssertionError( + "Put structured field on the out= " + "variant of a function; did you mean structured_delegate?" + ) + if not self.device_guard: + raise AssertionError( + "device_guard: False is not respected by structured kernels" + ) + if self.structured_delegate: + if self.func.kind() == SchemaKind.out: + raise AssertionError( + "structured_delegate field not allowed " + "on out= functions; did you mean structured?" + ) + if not self.device_guard: + raise AssertionError( + "device_guard: False is not respected by structured kernels" + ) + # Technically, with the asserts above, this assert is impossible to + # happen + if self.structured and self.structured_delegate: + raise AssertionError( + "Cannot have both structured and structured_delegate on function" + ) + defaulted_arguments = { + a.name for a in self.func.schema_order_arguments() if a.default is not None + } + invalid_args = set.difference(self.cpp_no_default_args, defaulted_arguments) + if len(invalid_args) != 0: + raise AssertionError(f"Invalid cpp_no_default_args: {invalid_args}") + if self.structured_inherits is not None: + if not self.structured: + raise AssertionError( + "structured_inherits must also imply structured: True" + ) + if str(self.func.name).startswith("_foreach"): + if self.device_check != DeviceCheckType.NoCheck: + raise AssertionError( + "foreach kernels fall back to slow path when tensor are on different devices, " + "device_check not allowed to be enabled" + ) + + # NB: if your function accidentally has rand/dropout/... in its name + # but is not actually random, feel free to amend this to special case + if ( + "rand" in str(self.func.name) + or ( + ( + "dropout" in str(self.func.name) + or any( + "dropout" in arg.name for arg in self.func.arguments.flat_all + ) + ) + # Backwards of dropout is typically deterministic + and "backward" not in str(self.func.name) + and str(self.func.name.name) != "_cudnn_init_dropout_state" + ) + or self.func.arguments.has_generator_arg() + ): + if "nondeterministic_seeded" not in self.tags: + raise AssertionError( + f"nondeterministic_seeded tag missing for {self.func.name}" + ) + + @property + def has_composite_kernel(self) -> bool: + return ( + self.has_composite_implicit_autograd_kernel + or self.has_composite_explicit_autograd_kernel + or self.has_composite_explicit_autograd_non_functional_kernel + ) or ( + self.has_composite_implicit_autograd_kernel + and self.has_composite_implicit_autograd_nested_tensor_kernel + ) + + @property + def is_view_op(self) -> bool: + rets = self.func.returns + is_non_mutating_view = len(rets) > 0 and any( + r.annotation is not None and not r.annotation.is_write for r in rets + ) + # See Note [resize_ in Functionalization] for more dtails + is_inplace_view = ( + "inplace_view" in self.tags + and str(self.func.name) != "resize_" + and str(self.func.name) != "resize_as_" + ) + is_wildcard_view = any( + inp.annotation is not None and "*" in inp.annotation.alias_set_after + for inp in self.func.schema_order_arguments() + ) + return is_non_mutating_view or is_inplace_view or is_wildcard_view + + @property + def view_schema_kind(self) -> ViewSchemaKind: + if self.is_view_op and self.func.name.name.inplace: + if "inplace_view" not in self.tags: + raise AssertionError(f"inplace_view tag missing for {self.func.name}") + return ViewSchemaKind.aliasing_inplace + if self.is_view_op: + return ViewSchemaKind.aliasing + else: + return ViewSchemaKind.non_aliasing + + @property + def root_name(self) -> str: + return self.func.name.name.base + + @property + def part_of_structured_group(self) -> bool: + return self.structured or self.structured_delegate is not None + + +class SchemaKind(Enum): + functional = auto() + inplace = auto() + out = auto() + mutable = auto() + scratch = auto() + + +# A structured kernel is guaranteed to have a functional and out variant, and +# optionally an inplace variant. +# +# NB: we create NativeFunctionsGroup *even if* the function is not +# actually annotated structured. Test the structured boolean to see if it +# actually is structured or not. +@dataclass(frozen=True) +class NativeFunctionsGroup: + functional: NativeFunction + inplace: NativeFunction | None + mutable: NativeFunction | None + out: NativeFunction + + @property + def structured(self) -> bool: + # Whether or not the operator has a meta() function. This information is backend-agnostic. + return self.out.structured + + def __post_init__(self) -> None: + test_sig: FunctionSchema = self.functional.func.signature() + for f in self.functions(): + if test_sig != f.func.signature(): + raise AssertionError( + "NativeFunctionsGroup constructed from two NativeFunctions " + f"that don't have matching signatures: {test_sig} != {f.func.signature()}" + ) + + if self.structured != f.part_of_structured_group: + raise AssertionError( + "NativeFunctionsGroup constructed from structured and unstructured " + f"functions: {self.out.func.name} and {f.func.name}" + ) + if self.functional.func.kind() != SchemaKind.functional: + raise AssertionError( + f"functional.func.kind() is {self.functional.func.kind()}, expected SchemaKind.functional" + ) + if self.out.func.kind() != SchemaKind.out: + raise AssertionError( + f"out.func.kind() is {self.out.func.kind()}, expected SchemaKind.out" + ) + if self.functional.namespace != self.out.namespace: + raise AssertionError( + f"functional.namespace ({self.functional.namespace}) != out.namespace ({self.out.namespace})" + ) + if self.inplace is not None: + if self.inplace.func.kind() != SchemaKind.inplace: + raise AssertionError( + f"inplace.func.kind() is {self.inplace.func.kind()}, expected SchemaKind.inplace" + ) + if self.inplace.namespace != self.functional.namespace: + raise AssertionError( + f"inplace.namespace ({self.inplace.namespace}) != functional.namespace ({self.functional.namespace})" + ) + + if self.mutable is not None: + if self.mutable.func.kind() != SchemaKind.mutable: + raise AssertionError( + f"mutable.func.kind() is {self.mutable.func.kind()}, expected SchemaKind.mutable" + ) + if self.mutable.namespace != self.functional.namespace: + raise AssertionError( + f"mutable.namespace ({self.mutable.namespace}) != functional.namespace ({self.functional.namespace})" + ) + # See Note [Overload Ambiguity With Functional Variants] + if not self.functional.func.name.name.functional_overload: + raise AssertionError( + "functional.func.name.name.functional_overload must be True when mutable is not None" + ) + + if self.structured: + # For now, structured composite kernels are not supported (need some + # design work to figure out how to make the composite case work) + if ( + self.out.has_composite_implicit_autograd_kernel + or self.out.has_composite_implicit_autograd_nested_tensor_kernel + ): + raise AssertionError("structured composite kernels are not supported") + + if self.functional.structured_delegate != self.out.func.name: + raise AssertionError( + f"{self.functional.func.name} delegates to {self.functional.structured_delegate} " + f"but its actual delegate is {self.out.func.name}" + ) + if self.inplace is not None: + if self.inplace.structured_delegate != self.out.func.name: + raise AssertionError( + f"{self.inplace.func.name} delegates to {self.inplace.structured_delegate} " + f"but its actual delegate is {self.out.func.name}" + ) + + generated_fns = sorted( + [str(f.func.name) for f in self.functions() if "generated" in f.tags] + ) + generated_fns_str = ", ".join(str(x) for x in generated_fns) + expected_generated_fns: set[str] = set() + for f in self.functions(): + expected_generated_fns.update(str(op) for op in f.autogen) + expected_generated_fns_str = ", ".join( + str(x) for x in sorted(expected_generated_fns) + ) + if len(expected_generated_fns) == 0 and len(generated_fns) > 0: + raise RuntimeError( + f"The codegen expects to be able to generate '{generated_fns_str}'." + " In order to generate them however, we expect them to be called out explicitly in the yaml." + f" Please add an 'autogen: {generated_fns_str}' line to the entry for {str(f.func.name)}" + ) + if expected_generated_fns_str != generated_fns_str: + raise RuntimeError( + f"The codegen expects to be able to generate '{generated_fns_str}'." + f" To do so, it expects a line: 'autogen: {generated_fns_str}'." + f" Instead, it found 'autogen: {expected_generated_fns_str}'" + ) + + def signature(self) -> FunctionSchema: + return self.out.func.signature() + + def functions(self) -> Iterator[NativeFunction]: + yield self.functional + yield self.out + if self.inplace is not None: + yield self.inplace + if self.mutable is not None: + yield self.mutable + + @property + def root_name(self) -> str: + return self.functional.root_name + + @staticmethod + def from_dict(d: dict[SchemaKind, NativeFunction]) -> NativeFunctionsGroup | None: + if not d: + raise AssertionError("from_dict called with empty dict") + if len(d) == 1: + return None + d = dict(d) # non-destructive updates please + functional = d.pop(SchemaKind.functional, None) + inplace = d.pop(SchemaKind.inplace, None) + mutable = d.pop(SchemaKind.mutable, None) + out = d.pop(SchemaKind.out, None) + if d: + raise AssertionError(f"unexpected keys in dict: {d}") + if functional is None: + raise AssertionError("functional variant is required") + # There are a few operators which only have functional/inplace variants; + # these don't count as structured for our purposes here + if out is None: + return None + # assuming all variants have the same namespace + return NativeFunctionsGroup( + functional=functional, + inplace=inplace, + mutable=mutable, + out=out, + ) + + +@dataclass(frozen=True) +class BackendMetadata: + # The name of the backend kernel, for a given operator + # for in-tree backends. These names come directly from the 'dispatch" field + # in native_functions.yaml. The dispatch entry is optional; in that + # case, that is equivalent to having written: + # + # dispatch: + # CompositeImplicitAutograd: $operator_name + kernel: str + # Whether or not the operator has a structured kernel implemented, for this particular backend. + # For in-tree backends, they all have the same value for structured- this is listed + # in native_functions.yaml. + # However, external backends like XLA can indendently toggle which ops are structured. + structured: bool + + # The namespace for kernels, default value: DEFAULT_KERNEL_NAMESPACE + cpp_namespace: str + + def supports_symint(self) -> bool: + return "_symint" in self.kernel + + +@dataclass(frozen=True) +class UfuncInnerLoop: + name: str + supported_dtypes: OrderedSet[ScalarType] + # key is stored here because it affects the semantics of name, + # so its helpful to have them together for further processing + ufunc_key: UfuncKey + + @staticmethod + def parse(value: str, ufunc_key: UfuncKey) -> UfuncInnerLoop: + name, supported_dtypes_str = value.split(" ", 1) + if supported_dtypes_str[0] != "(": + raise AssertionError( + f"expected '(' at start of supported_dtypes, got: {supported_dtypes_str}" + ) + if supported_dtypes_str[-1] != ")": + raise AssertionError( + f"expected ')' at end of supported_dtypes, got: {supported_dtypes_str}" + ) + supported_dtypes: OrderedSet[ScalarType] = OrderedSet() + for k in supported_dtypes_str[1:-1].split(", "): + supported_dtypes |= ScalarType.parse_set(k) + return UfuncInnerLoop( + name=name, supported_dtypes=supported_dtypes, ufunc_key=ufunc_key + ) + + +# BackendIndex represents a backend. +# The BackendIndex encodes per-operator information that is potentially different +# for each backend. The most obvious example is the name of the kernel +# (the 'dispatch' entry in native_functions.yaml). +# However, there can be other examples of different backends having different information. +# External backends can choose to opt their kernels to be structured independently from in-tree backends, +# which means that this information isn't inherently tied to a NativeFunction- it's different per backend. +@dataclass(frozen=True) +class BackendIndex: + dispatch_key: DispatchKey + # Mainly important for structured kernels, this determines which variant in the operator group is used to implement the others. + # All in-tree ops use out kernels, while XLA uses functional kernels. + use_out_as_primary: bool + # Whether the backend requires a device guard, and device checks. + # For in-tree backends, this is currently just CUDA/HIP + # For out-of-tree backends, this is currently just Intel XPU + device_guard: bool + # Whether the backend is in-tree (CPU/CUDA) or out-of-tree (XLA) + external: bool + # Other backend-specific information that is on a per-operator basis + index: dict[OperatorName, BackendMetadata] + + @staticmethod + def grow_index( + parent_index: dict[DispatchKey, dict[OperatorName, BackendMetadata]], + child_index: dict[DispatchKey, dict[OperatorName, BackendMetadata]], + ) -> None: + for k, v in child_index.items(): + for op_name, metadata in v.items(): + if op_name in parent_index[k]: + raise AssertionError( + f"duplicate operator {op_name} for dispatch key {k}" + ) + parent_index[k][op_name] = metadata + + def primary(self, g: NativeFunctionsGroup) -> NativeFunction: + if self.use_out_as_primary: + return g.out + else: + return g.functional + + def has_kernel(self, g: NativeFunction | NativeFunctionsGroup) -> bool: + m = self.get_kernel(g) + return m is not None + + def get_kernel( + self, g: NativeFunction | NativeFunctionsGroup + ) -> BackendMetadata | None: + if isinstance(g, NativeFunction): + f = g + elif isinstance(g, NativeFunctionsGroup): + f = self.primary(g) + else: + assert_never(g) + if f.func.name not in self.index: + return None + return self.index[f.func.name] + + def native_function_class_name(self) -> str | None: + if self.external: + return f"{str(self.dispatch_key)}NativeFunctions" + else: + # TODO: This discrepancy isn't required; we could also generated + # a class for in-tree kernels. It'll just require carefully + # updating every kernel definition + callsite of every in-tree aten kernel. + return None + + +# The function schema is undoubtedly the most important data structure +# in all of the codegen, as it defines the type signature for operators, +# and most of the code generation we do is type directed (e.g., look at +# the types, decide what to do. Think about how we code generate +# C++ function stubs!) +# +# We will also see in this class the general structure for how we model +# data in this code generation. A few notable properties to point out +# ahead of time: +# +# - These dataclasses are a *lossless* representation of the strings +# they are parsed from. In fact, we assert that given the +# information stored in the dataclass, we can exactly reconstruct +# the string we parsed from (and assert this inside the parse +# definition). There are a few reasons for this: +# +# - If you find that it is difficult to reconstruct the string +# given a dataclass, that is a clue that you are data +# representation is wrong. +# +# - It helps ensure that all relevant information is present +# in the dataclass, so that downstream users aren't tempted +# to reparse the original string to get some information +# that was omitted. +# +# - It forces you to represent the data in-memory in the same way +# it is recorded textually, which makes the dataclasses easier +# to understand for someone who is familiar with the +# textual format. (As a tradeoff, it means you have to model +# the syntax, even when it is inconvenient. But maybe that means +# the syntax is bad!) If you don't understand the internal +# representation, go look at the printing code to see how +# it maps onto the surface syntax! +# +# - It makes it easy to test the parsing code, as parsing code +# that is inconsistent with the string code will fail early +# and loudly. (As a tradeoff, it makes the parsing code a bit +# brittle (in particular, with trivial whitespace changes you +# are likely to trigger an assert error). +# +# In general, try to make the __str__ code as simple as possible +# (even at the cost of more complex parsing logic.) Additionally, +# try to minimize redundancy in data representation. (Precomputed +# fields are OK though: they are defined as a simple function on +# the canonical representation in question.) +# +# - These dataclasses are all frozen; once constructed their +# values never change. This makes it easy to tell where any +# given data came from: just look to the constructor. As a +# tradeoff, you can't easily "decorate" a schema with extra +# information from a post-facto analysis. We impose this +# restriction to make these structures more understandable. +# +@dataclass(frozen=True) +class FunctionSchema: + # The name of the operator this function schema describes. + name: OperatorName + + arguments: Arguments + + # TODO: Need to handle collisions with argument names at some point + returns: tuple[Return, ...] + + @property + def is_mutable(self) -> bool: + def is_write(arg: Argument) -> bool: + if arg.annotation is None: + return False + return arg.annotation.is_write + + # Corresponds to torch._C._FunctionSchema.is_mutable + # See aten/src/ATen/core/function_schema.h (keep these in sync) + return any(is_write(a) for a in self.arguments.flat_all) + + def schema_order_arguments(self) -> Iterator[Argument]: + return itertools.chain( + self.arguments.flat_positional, + self.arguments.flat_kwarg_only, + self.arguments.out, + ) + + decl_re = re.compile(r"(?P[^\(]+)\((?P.*)\) -> (?P.*)") + + @staticmethod + def parse(func: str) -> FunctionSchema: + # We should probably get a proper parser here + decls = FunctionSchema.decl_re.findall(func) + if len(decls) != 1: + raise AssertionError(f"Invalid function schema: {func}") + ops, args, return_decl = decls[0] + name = OperatorName.parse(ops) + arguments = Arguments.parse(args) + returns = parse_returns(return_decl) + r = FunctionSchema(name=name, arguments=arguments, returns=returns) + if str(r) != func: + raise AssertionError(f"{str(r)} != {func}") + return r + + def returns_are_aliased(self) -> bool: + # We assert earlier that schemas can't have a mix of aliased and non-aliased returns + return any( + r + for r in self.returns + if r.annotation is not None and r.annotation.is_write + ) + + def __post_init__(self) -> None: + for arg, ret in zip(self.arguments.out, self.returns): + if arg.annotation != ret.annotation: + raise AssertionError( + "Out arguments must have matching return Tensor; furthermore, " + f"the ith-argument needs to correspond to the ith return. " + f"arg.annotation={arg.annotation}, ret.annotation={ret.annotation}" + ) + # We also enforce that if you have any mutable, positional args, then they are not returned. + # This makes it easier to group these functions properly with their functional/out= counterparts. + for a in self.arguments.post_self_positional_mutable: + if any(a.annotation == r.annotation for r in self.returns): + raise AssertionError( + f"If you have a schema with mutable positional args, we expect them to not be returned. schema: {str(self)}" + ) + # Invariant: we expect out arguments to appear as keyword arguments in the schema. + # This means that all mutable returns should be aliased to a keyword argument + # (except for "self", which we explicitly don't treat as an out argument because of its use in methods) + # See Note [is_out_fn] + out_and_self = list(self.arguments.out) + [ + arg for arg in self.arguments.flat_positional if arg.name == "self" + ] + mutable_returns = [ + ret + for ret in self.returns + if ret.annotation is not None and ret.annotation.is_write + ] + immutable_returns = [ + ret + for ret in self.returns + if ret.annotation is None or not ret.annotation.is_write + ] + # Some assertions: We don't want any functions with a return type of "-> (Tensor(a!), Tensor)", + # because: + # (1) It's more annoying to handle properly + # (2) It's unnecessary - you can't method-chain on the first (mutated) output because it's part of a tuple. + # Instead, we expect the (a!) argument to not be returned. + if not (len(mutable_returns) == 0 or len(immutable_returns) == 0): + raise AssertionError( + f"NativeFunctions must have either only mutable returns, or only immutable returns. Found: {str(self)}" + ) + for ret in mutable_returns: + if not any(ret.annotation == arg.annotation for arg in out_and_self): + raise AssertionError( + 'All mutable returns must be aliased either to a keyword argument, or to "self". ' + "Did you forget to mark an out argument as keyword-only?" + ) + if self.arguments.out: + # out= ops that return their mutable inputs are only really useful for method chaining. + # And method chaining is only really useful if the thing you're returning is a plain Tensor. + # So ideally, we'd enforce that out= ops with a single plain mutable tensor should return the tensor, + # and all other types of out= op schemas should return void. + # There are a bunch of existing out= ops that return tuples of tensors though, so we're stuck with allowing that. + if any(a.type != BaseType(BaseTy.Tensor) for a in self.arguments.out): + if len(self.returns) != 0: + raise AssertionError( + "out= ops that accept tensor lists as out arguments " + "are expected to have no return type (since you can't do method chaining on them)" + ) + else: + # mutable keyword arguments whose name has _scratch_ prefix are + # scratch tensors for memory planning and should not be returned + non_scratch_out_args = len( + [ + arg + for arg in self.arguments.out + if not arg.name.startswith("_scratch_") + ] + ) + if non_scratch_out_args != len(self.returns): + raise AssertionError( + f"Must return as many arguments as there are out arguments, or no return at all. " + f"Got {non_scratch_out_args} non-scratch out args and {len(self.returns)} returns" + ) + + if self.name.name.inplace: + self_a = self.arguments.self_arg + if not ( + self_a + and self_a.argument.annotation + and self_a.argument.annotation.is_write + ): + raise AssertionError( + f"Inplace op {self.name} must have a self argument with a mutable annotation" + ) + if self_a.argument.type == BaseType(BaseTy.Tensor): + # All inplace ops with an ordinary `Tensor self` argument should return self, + # to allow for method chaining. + if not ( + len(self.returns) == 1 + and self.returns[0].annotation == self_a.argument.annotation + ): + raise AssertionError( + f"Inplace op {self.name} with Tensor self must return self" + ) + else: + # You can't method chain on non-tensor self arguments though (like a list[Tensor]) + # so in all other cases we expect the return type to be none. + if len(self.returns) != 0: + raise AssertionError( + f"Inplace op {self.name} with non-Tensor self must have no returns" + ) + + if self.arguments.tensor_options is not None: + if self.kind() != SchemaKind.functional: + raise AssertionError( + "Found an operator that is not functional or out variant, but has tensor options arguments." + "This is not allowed- tensor options arguments are only allowed for factory functions." + f"schema: {str(self)}" + ) + if self.is_functional_fn(): + if self.kind() != SchemaKind.functional: + raise AssertionError( + "Found an operator that is not functional, but its overload contains the string 'functional'." + "This is a special keyword in the codegen, please use a different overload name." + f"schema: {str(self)}" + ) + + def is_functional_fn(self) -> bool: + return "functional" in self.name.overload_name + + def is_out_fn(self) -> bool: + # Note [is_out_fn] + # + # out functions are the variants which take an explicit out= argument + # to populate into. We need to know if a schema corresponds to an + # out function for several reasons: + # + # - They codegen differently in C++ API + # - codegen to at::add_out rather than at::add + # - out argument is moved to front of C++ argument list + # + # out functions are DEFINED to be any function with a keyword-only + # argument that is mutable. In principle, this could lead to a + # false positive if you define a function that mutates a + # kwarg only argument, but this isn't the "true" output of this + # function. A more robust definition that would work in this + # case would also look at: + # + # - The output types. Out functions take in the arguments + # they mutate and then return them again; this is sort + # of "definitionally" what makes something an out function. + # Historically, we DO check this for consistency. + # - Correspondence with pure variant. An out function + # should have a signature equivalent to its pure variant, + # but just with extra kwargs for the output elements. This + # is difficult to actually check for and historically + # we only do this check in tools/ + return bool(self.arguments.out) + + def kind(self) -> SchemaKind: + """ + What kind of schema is this? A functional schema is one + that returns a newly allocated output; an inplace schema + modifies the self argument inplace; an out schema writes + the result into an explicitly provided out argument. + """ + is_out = bool(self.arguments.out) + is_scratch = bool( + [arg for arg in self.arguments.out if arg.name.startswith("_scratch_")] + ) + is_inplace = self.name.name.inplace + is_mutable = any( + a.annotation is not None and a.annotation.is_write + for a in self.arguments.post_self_positional + ) + if is_out and is_inplace: + raise AssertionError("A schema cannot be both out= and inplace") + # out= and inplace schemas can also have post_self_positional mutable args, + # but we give precedence to out= and inplace when deciding the schema kind. + # Tradeoff: we probably don't want to have to teach codegen that looks at inplace ops + # to also worry about mutable post_self_positional arguments, + # but it seems like a much bigger lift to classify them has having a new schema kind. + # The number of ops that fit in this strange category is small enough that + # we can probably manually write code for them instead of forcing the codegen to handle them. + if is_inplace: + return SchemaKind.inplace + elif is_scratch: + if not is_out: + raise AssertionError( + "invariant: all scratch operators are expected to be out= operators too" + ) + return SchemaKind.scratch + elif is_out: + if is_scratch: + raise AssertionError( + "We should not categorize a scratch op as an out variant. Check if the order of if statements are expected!" + ) + return SchemaKind.out + elif is_mutable: + return SchemaKind.mutable + else: + return SchemaKind.functional + + # For every return: + # - If the return aliases an input, we return the input name + # - Otherwise, we return None. + # If return names were enforced to be consistent with aliasing information, then we wouldn't need this. + def aliased_return_names(self) -> list[str | None]: + outs: list[str | None] = [] + for r in self.returns: + aliased_args = [ + a + for a in self.arguments.flat_all + if a.annotation is not None and a.annotation == r.annotation + ] + if len(aliased_args) == 0: + outs.append(None) + elif len(aliased_args) == 1: + outs.append(aliased_args[0].name) + else: + aliased_names = ", ".join(a.name for a in aliased_args) + raise AssertionError( + f"Found a return ({r.name})that aliases multiple inputs ({aliased_names})" + ) + return outs + + def signature( + self, + *, + strip_default: bool = False, + strip_view_copy_name: bool = False, + keep_return_names: bool = False, + ) -> FunctionSchema: + """ + Certain schemas are 'related', in that they are simply + inplace/out/functional versions of the same function. This method + factors these schemas into the "core" functional signature which + is equal across all versions. + + Here is what normalization happens to the schema to convert + it to a signature: + - The overload name is stripped (name is retained, since + it expresses semantic content about what the function does) + - Inplace is set False + - Out arguments are stripped + - Mutable post_self_positional args are converted to returns + - Mutability annotations are stripped (this is sound + because you cannot overload on mutability annotation) + - Return names are stripped since they are not overloadable and + some variants have return names but some not + - TensorOptions are dropped + because out= variants of factory functions don't include them + (and we want to be able to pair up factory functions with their out variants) + + Finally, we want to be able to pair up related "view" and their + corresponding "view_copy" operators. We do this by optionally + stripping the trailing "_copy" from the base name. + + Example of a mutable op before and after: + + f.func (Mutable operator): + _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask) # noqa: B950 + + f.func (Corresponding functional operator): + _fused_moving_avg_obs_fq_helper.functional(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor running_min, Tensor running_max, Tensor scale, Tensor zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask, Tensor running_min_out, Tensor running_max_out, Tensor scale_out, Tensor zero_point_out) # noqa: B950 + + f.func.signature() output: + _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor running_min, Tensor running_max, Tensor scale, Tensor zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) # noqa: B950 + """ + + def strip_ret_annotation(r: Return) -> Return: + return Return( + name=r.name if keep_return_names else None, + type=r.type, + annotation=None, + ) + + base_name = self.name.name.base + if strip_view_copy_name: + if base_name.endswith("_copy"): + base_name = base_name.replace("_copy", "") + elif base_name.endswith("_scatter"): + base_name = base_name.replace("scatter", "inverse") + + # find mutable inputs that are not originally returned, and convert them to returns + returns_from_mutable_inputs = tuple( + # When we're grouping functions we strip the return names, + # but when we're generating the actual functional variants then we follow + # a convention for what to name the returns + Return( + name=f"{a.name}_out" if keep_return_names else None, + type=a.type, + annotation=None, + ) + for a in itertools.chain( + # Order is important here (otherwise e.g. inplace with mutable args + # and out= with mutable args won't have the same signature) + ( + [self.arguments.self_arg.argument] + if self.arguments.self_arg is not None + else [] + ), + self.arguments.out, + self.arguments.post_self_positional, + ) + if a.annotation is not None + and a.annotation.is_write + and not any(a.annotation == r.annotation for r in self.returns) + ) + original_returns = tuple(map(strip_ret_annotation, self.returns)) + # Ordering is important here. We expect the "mutable input" returns to come last. + returns = original_returns + returns_from_mutable_inputs + + args_sig = self.arguments.signature(strip_default=strip_default) + # See Note [bernoulli.p schema] + if str(self.name) == "bernoulli.p": + args_sig = Arguments.parse(str(args_sig).replace("float p", "float p=0.5")) + + return FunctionSchema( + name=OperatorName( + name=BaseOperatorName( + base=base_name, + inplace=False, + dunder_method=self.name.name.dunder_method, + ), + overload_name="", # stripped + ), + arguments=args_sig, + returns=returns, + ) + + def view_signature(self) -> FunctionSchema: + return self.signature(strip_view_copy_name=True) + + def with_name(self, name: OperatorName) -> FunctionSchema: + return FunctionSchema( + name=name, + arguments=self.arguments, + returns=self.returns, + ) + + @property + def modifies_arguments(self) -> bool: + return self.kind() in [SchemaKind.inplace, SchemaKind.out, SchemaKind.mutable] + + def has_symint(self) -> bool: + return self.arguments.has_symint_arg() + + def __str__(self) -> str: + all_arguments_str = str(self.arguments) + if len(self.returns) == 1: + returns = str(self.returns[0]) # omit parentheses + else: + returns = "(" + ", ".join(map(str, self.returns)) + ")" + return f"{self.name}({all_arguments_str}) -> {returns}" + + +# Here is the rest of the data model, described more briefly. + + +# Simplified version for what actually shows up in built-ins. +# Look at alias_info.h for expanded syntax. If you need the structure, +# you also need to make this structure recursive so it can be lined +# up with the type components too. For primitives this isn't really +# necessary +@dataclass(frozen=True) +class Annotation: + # Typically only has one element. Not actually a set so + # we can conveniently assume it is canonically ordered + alias_set: tuple[str, ...] + is_write: bool + alias_set_after: tuple[str, ...] + + @staticmethod + def parse(ann: str) -> Annotation: + # TODO: implement a proper parser if this gets more ugly + # Regex Explanation: + # Example: "a! -> a|b" + # Group #1: alias before optional '|', required. Matches the first + # character 'a' in the example + # Group #2: optional alias set after optional '|', matches empty string + # in the example + # Group #3: optional "is write" flag, matches '!' in the example. + # Group #4: optional section containing arrow, matches " -> a|b" in the + # example. + # Group #5: optional alias after set, supports wildcard, matches "a|b" + # in the example. + # Group #6: optional sub-section of alias after set, matches "|b" in the + # example. + m = re.match(r"^([a-z])(\|[a-z])*(!?)( -> (\*|[a-z](\|[a-z])*))?$", ann) + + if m is None: + raise AssertionError(f"unrecognized alias annotation {ann}") + before_alias = m.group(1) + (m.group(2) if m.group(2) else "") + alias_set = tuple(before_alias.split("|")) + is_write = m.group(3) == "!" + if is_write and len(alias_set) > 1: + raise AssertionError( + f"alias set larger than 1 is not mutable, got {ann} instead." + ) + after_set = tuple(m.group(5).split("|")) if m.group(5) else () + if len(before_alias) > 1 and len(after_set) > 1: + raise AssertionError( + f"before alias set and after alias set cannot be larger than 1 at the same time, got {ann} instead." + ) + r = Annotation( + alias_set=alias_set, is_write=is_write, alias_set_after=after_set + ) + if str(r) != ann: + raise AssertionError(f"{r} != {ann}") + return r + + def __str__(self) -> str: + alias_set = "|".join(self.alias_set) + if self.is_write: + alias_set = f"{alias_set}!" + alias_set_after = "|".join(self.alias_set_after) + if alias_set_after: + alias_set = f"{alias_set} -> {alias_set_after}" + return alias_set + + +# The base class for the type system. This is also loosely modeled +# off of jit_type.h, but we've simplified the hierarchy to focus +# in on the aspects of the type system that matter for code generation +# (for example, there's no SingleElementType subclass anymore). +# You never actually construct a Type; usually it's going to be one +# of the subclasses. If Python had ADTs this would be one! +@dataclass(frozen=True) +class Type: + @staticmethod + def parse(t: str) -> Type: + r = Type._parse(t) + if str(r) != t: + raise AssertionError(f"{r} != {t}") + return r + + @staticmethod + def _parse(t: str) -> Type: + m = re.match(r"^(.+)\?$", t) + if m is not None: + return OptionalType(Type.parse(m.group(1))) + m = re.match(r"^(.+)\[([0-9]+)?\]$", t) + if m is not None: + size = int(m.group(2)) if m.group(2) is not None else None + return ListType(elem=Type.parse(m.group(1)), size=size) + + # '__torch__.torch.classes.' is the prefix for custom class + m = re.match(r"^__torch__\.torch\.classes\.([a-zA-Z0-9_.]+)$", t) + if m is not None: + return CustomClassType(m.group(1)) + try: + return BaseType(BaseTy[t]) + except KeyError as e: + raise RuntimeError(f"unrecognized type {t}") from e + + def __str__(self) -> str: + raise NotImplementedError + + # WARNING: These concepts are not very well-defined. For example, + # is "int?" nullable? How about "int?[]". They are defined + # so we can conveniently generate legacy Declarations.yaml but + # really we should probably just remove these at some point + + def is_base_ty_like(self, base_ty: BaseTy) -> bool: + raise NotImplementedError + + def is_tensor_like(self) -> bool: + return self.is_base_ty_like(BaseTy.Tensor) + + def is_generator_like(self) -> bool: + return self.is_base_ty_like(BaseTy.Generator) + + def is_symint_like(self) -> bool: + return self.is_base_ty_like(BaseTy.SymInt) + + def is_nullable(self) -> bool: + raise NotImplementedError + + def is_list_like(self) -> ListType | None: + raise NotImplementedError + + +# Base types are simple, atomic types with no further structure +class BaseTy(Enum): + Generator = auto() + ScalarType = auto() + Tensor = auto() + int = auto() + Dimname = auto() + DimVector = auto() + float = auto() + str = auto() + bool = auto() + Layout = auto() + Device = auto() + DeviceIndex = auto() + Scalar = auto() + MemoryFormat = auto() + QScheme = auto() + Storage = auto() + Stream = auto() + SymInt = auto() + SymBool = auto() + GraphModule = auto() + + +@dataclass(frozen=True) +class BaseType(Type): + name: BaseTy + + def __str__(self) -> str: + return f"{self.name.name}" + + def is_base_ty_like(self, base_ty: BaseTy) -> bool: + return self.name == base_ty + + def is_nullable(self) -> bool: + return False + + def is_list_like(self) -> ListType | None: + return None + + def is_symint_like(self) -> bool: + return self.name == BaseTy.SymInt + + +# Optional types may be specified, or may also be validly given None +@dataclass(frozen=True) +class OptionalType(Type): + elem: Type + + def __str__(self) -> str: + return f"{self.elem}?" + + def is_base_ty_like(self, base_ty: BaseTy) -> bool: + return self.elem.is_base_ty_like(base_ty) + + def is_symint_like(self) -> bool: + return self.elem.is_symint_like() + + def is_nullable(self) -> bool: + return True + + def is_list_like(self) -> ListType | None: + return self.elem.is_list_like() + + +# A type representing a PyTorch custom class +@dataclass(frozen=True) +class CustomClassType(Type): + class_name: str + + def __str__(self) -> str: + """ + Return the class name will prefix __torch__.torch.classes + """ + return f"__torch__.torch.classes.{self.class_name}" + + def is_base_ty_like(self, base_ty: BaseTy) -> bool: + return False + + def is_symint_like(self) -> bool: + return False + + def is_nullable(self) -> bool: + """ + Assume a custom class is not nullable. + """ + return False + + def is_list_like(self) -> ListType | None: + return None + + +# List types specify that we may have multiples of an element. We +# also support explicit sizes on list types, but these have +# some nontrivial semantics! (However, for C++ API purposes, explicit +# sizes are mostly erased from the type system.) +# +# DANGER WILL ROBINSON: C++ elaboration depends on elem type; e.g., +# int[] elaborates differently than bool[3]! +@dataclass(frozen=True) +class ListType(Type): + elem: Type + size: int | None + + def __str__(self) -> str: + size = f"{self.size}" if self.size else "" + return f"{self.elem}[{size}]" + + def is_base_ty_like(self, base_ty: BaseTy) -> bool: + return self.elem.is_base_ty_like(base_ty) + + def is_symint_like(self) -> bool: + return self.elem.is_symint_like() + + def is_nullable(self) -> bool: + return self.elem.is_nullable() + + def is_list_like(self) -> ListType | None: + return self + + +@dataclass(frozen=True) +class Argument: + # NB: I didn't put kwarg_only as a boolean field here, unlike + # c10::Argument, so that printing works correctly + + name: str + type: Type + default: str | None + + # The semantics of the annotation field are a little strange. + # + # Alias annotations parametrize Tensors (since Tensors are the only things + # that can alias.) This motivates why I write Tensor(a!)? (and not, for + # example, Tensor?(a!)), because the (a!) describes aliasing on the tensor, + # which may be optional (i.e., the alias annotation should bind first to + # Tensor, before the optional postfix annotation). + # + # However, despite being a property of Tensor, we (and c10::Argument) + # store the annotation at the top level of the Argument, rather than + # inside the embedded Tensor type. In the C++ version of this + # class, we then go through great lengths to mimic the type + # structure in the annotation structure so we can correlate + # annotations with types. + # + # Now, it turns out, in all applications in code generation, the + # structure of annotated types is very simple. So we just hard + # code it here. But if we ever do get anything more complex, this + # model will have to change! + annotation: Annotation | None + + @property + def alias_info(self) -> Annotation | None: + return self.annotation + + @staticmethod + def parse(arg: str) -> Argument: + name: str + default: str | None + if " " not in arg: + raise AssertionError(f"illegal argument '{arg}'") + if "=" in arg: + if arg.count("=") != 1: + raise AssertionError(f"illegal argument with default value: '{arg}'") + type_and_annot_and_name, default = arg.split("=") + type_and_annot, name = type_and_annot_and_name.rsplit(" ", 1) + name_and_default = f"{name}={default}" + else: + type_and_annot, name_and_default = arg.rsplit(" ", 1) + name = name_and_default + default = None + # TODO: deduplicate annotation matching with Return + match = re.match(r"Tensor\((.+)\)(.*)", type_and_annot) + annotation: Annotation | None + if match: + # If you update this, make sure the __str__ still works too + if match.group(2) not in ["", "?", "[]"]: + raise AssertionError( + f"unrecognized alias analysis form with Tensor: {match.group(2)}" + ) + type_s = "Tensor" + match.group(2) + annotation = Annotation.parse(match.group(1)) + else: + type_s = type_and_annot + annotation = None + type = Type.parse(type_s) + r = Argument( + name=name, + type=type, + default=default, + annotation=annotation, + ) + if str(r) != arg: + raise AssertionError(f"{str(r)} != {arg}") + return r + + @property + def is_write(self) -> bool: + return self.annotation is not None and self.annotation.is_write + + def __str__(self) -> str: + type = f"{self.type}" + if self.annotation: + if type not in ["Tensor", "Tensor?", "Tensor[]"]: + raise AssertionError(f"annotation on non-Tensor type: {type}") + type = type.replace("Tensor", f"Tensor({self.annotation})") + if self.name is None: + return type + else: + mb_default = "" + if self.default: + mb_default = f"={self.default}" + return f"{type} {self.name}{mb_default}" + + +@dataclass(frozen=True) +class Return: + name: str | None + type: Type + annotation: Annotation | None + + @property + def alias_info(self) -> Annotation | None: + return self.annotation + + @staticmethod + def parse(arg: str) -> Return: + name: str | None + if " " in arg: + type_and_annot, name = arg.rsplit(" ", 1) + else: + type_and_annot = arg + name = None + match = re.match(r"Tensor\((.+)\)(.*)", type_and_annot) + annotation: Annotation | None + if match: + # If you update this, make sure the __str__ still works too + if match.group(2) not in ["", "?", "[]"]: + raise AssertionError( + f"unrecognized alias analysis form with Tensor: {match.group(2)}" + ) + type_s = "Tensor" + match.group(2) + annotation = Annotation.parse(match.group(1)) + else: + type_s = type_and_annot + annotation = None + type = Type.parse(type_s) + r = Return( + name=name, + type=type, + annotation=annotation, + ) + if str(r) != arg: + raise AssertionError(f"{str(r)} != {arg}") + return r + + @property + def is_write(self) -> bool: + return self.annotation is not None and self.annotation.is_write + + def __str__(self) -> str: + type = f"{self.type}" + if self.annotation: + if type not in ["Tensor", "Tensor?", "Tensor[]"]: + raise AssertionError(f"annotation on non-Tensor type: {type}") + type = type.replace("Tensor", f"Tensor({self.annotation})") + if self.name is None: + return type + else: + return f"{type} {self.name}" + + +# Represents the self argument for functions that may be methods +@dataclass(frozen=True) +class SelfArgument: + argument: Argument + + +# Bundle of arguments that represent a TensorOptions. This is mostly +# relevant for the public C++ API but we bake it into the core data +# model because other APIs often have to interact with it +@dataclass(frozen=True) +class TensorOptionsArguments: + dtype: Argument + layout: Argument + device: Argument + pin_memory: Argument + + def all(self) -> Sequence[Argument]: + return [self.dtype, self.layout, self.device, self.pin_memory] + + +@dataclass(frozen=True) +class Arguments: + # pre_self_positional is usually empty, but is notably non-empty + # for where.self, where the condition argument comes before the + # self argument + pre_self_positional: tuple[Argument, ...] + self_arg: SelfArgument | None + post_self_positional: tuple[Argument, ...] + + pre_tensor_options_kwarg_only: tuple[Argument, ...] + tensor_options: TensorOptionsArguments | None + # post_tensor_options is typically memory format, which should be + # part of tensor options but isn't right now, and is usually + # placed after the tensor options arguments + post_tensor_options_kwarg_only: tuple[Argument, ...] + + # Unlike in the previous codegen, we have factored out 'out' arguments + # in the canonical representation, removing them from kwarg + # arguments. This choice is justified by numerous downstream + # transformations which treat out arguments specially; additionally, + # you can see that canonicity is not violated! + out: tuple[Argument, ...] # these are also kwarg-only + + @property + def flat_non_out(self) -> Sequence[Argument]: + ret: list[Argument] = [] + ret.extend(self.flat_positional) + ret.extend(self.flat_kwarg_only) + return ret + + @property + def flat_positional(self) -> Sequence[Argument]: + ret: list[Argument] = [] + ret.extend(self.pre_self_positional) + if self.self_arg is not None: + ret.append(self.self_arg.argument) + ret.extend(self.post_self_positional) + return ret + + @property + def post_self_positional_mutable(self) -> Sequence[Argument]: + return [a for a in self.post_self_positional if a.is_write] + + # NB: doesn't contain out arguments + @property + def flat_kwarg_only(self) -> Sequence[Argument]: + ret: list[Argument] = [] + ret.extend(self.pre_tensor_options_kwarg_only) + if self.tensor_options is not None: + ret.extend(self.tensor_options.all()) + ret.extend(self.post_tensor_options_kwarg_only) + return ret + + @property + def flat_all(self) -> Sequence[Argument]: + ret: list[Argument] = [] + ret.extend(self.flat_positional) + ret.extend(self.flat_kwarg_only) + ret.extend(self.out) + return ret + + @property + def non_out( + self, + ) -> Sequence[Argument | SelfArgument | TensorOptionsArguments]: + ret: list[Argument | SelfArgument | TensorOptionsArguments] = [] + ret.extend(self.positional) + ret.extend(self.kwarg_only) + return ret + + @property + def positional(self) -> Sequence[Argument | SelfArgument]: + ret: list[Argument | SelfArgument] = [] + ret.extend(self.pre_self_positional) + if self.self_arg is not None: + ret.append(self.self_arg) + ret.extend(self.post_self_positional) + return ret + + @property + def kwarg_only(self) -> Sequence[Argument | TensorOptionsArguments]: + ret: list[Argument | TensorOptionsArguments] = [] + ret.extend(self.pre_tensor_options_kwarg_only) + if self.tensor_options is not None: + ret.append(self.tensor_options) + ret.extend(self.post_tensor_options_kwarg_only) + return ret + + @property + def all(self) -> Sequence[Argument | SelfArgument | TensorOptionsArguments]: + ret: list[Argument | SelfArgument | TensorOptionsArguments] = [] + ret.extend(self.positional) + ret.extend(self.kwarg_only) + ret.extend(self.out) + return ret + + def mutable_arg_names(self) -> list[str]: + return [ + a.name + for a in self.flat_all + if a.annotation is not None and a.annotation.is_write + ] + + def has_tensor_arg(self) -> bool: + return any(a.type.is_tensor_like() for a in self.flat_non_out) + + def has_symint_arg(self) -> bool: + return any(a.type.is_symint_like() for a in self.flat_non_out) + + def has_generator_arg(self) -> bool: + return any(a.type.is_generator_like() for a in self.flat_non_out) + + def signature(self, *, strip_default: bool = False) -> Arguments: + # dataclasses.replace could be used here, but it is less + # type safe so for now I've opted to type everything out + def strip_arg_annotation(a: Argument) -> Argument: + return Argument( + name=a.name, + type=a.type, + default=a.default if not strip_default else None, + annotation=None, + ) + + return Arguments( + pre_self_positional=tuple( + map(strip_arg_annotation, self.pre_self_positional) + ), + self_arg=( + SelfArgument(strip_arg_annotation(self.self_arg.argument)) + if self.self_arg is not None + else None + ), + post_self_positional=tuple( + map(strip_arg_annotation, self.post_self_positional) + ), + # Since TensorOptions are dropped, the post_tensor_options_kwargs are + # converted to pre_tensor_options_kwargs + pre_tensor_options_kwarg_only=tuple( + map(strip_arg_annotation, self.pre_tensor_options_kwarg_only) + ) + + tuple(map(strip_arg_annotation, self.post_tensor_options_kwarg_only)), + # TensorOptions are dropped in signature, + # so we can pair factory functions with their out= variants. + tensor_options=None, + post_tensor_options_kwarg_only=(), + # out arguments are dropped in signature + out=(), + ) + + def remove_self_annotation(self) -> Arguments: + if self.self_arg is None: + raise AssertionError("remove_self_annotation called but self_arg is None") + return dataclasses.replace( + self, + self_arg=SelfArgument( + dataclasses.replace(self.self_arg.argument, annotation=None) + ), + ) + + def with_out_args(self, outs: list[Argument]) -> Arguments: + if len(self.out) != 0: + raise AssertionError( + f"with_out_args called but self.out is not empty: {self.out}" + ) + return dataclasses.replace( + self, + out=tuple(outs), + ) + + @staticmethod + def _preparse(args: str) -> tuple[list[Argument], list[Argument], list[Argument]]: + positional: list[Argument] = [] + kwarg_only: list[Argument] = [] + out: list[Argument] = [] + arguments_acc = positional + + # TODO: Use a real parser here; this will get bamboozled + # by signatures that contain things like std::array (note the space) + for arg in args.split(", "): + if not arg: + continue + if arg == "*": + if arguments_acc is not positional: + raise AssertionError( + "invalid syntax: kwarg-only specifier * can only occur once" + ) + arguments_acc = kwarg_only + continue + parg = Argument.parse(arg) + # Currently, we rely directly on the invariant that there are NO + # kwarg-only mutating arguments. If you want to relax this, + # we will need a more semantic way of matching that takes + # into account return arguments. In that case, you will have + # to manage out computation a level up, in FunctionSchema. See Note + # [is_out_fn] + if parg.annotation is not None and parg.annotation.is_write: + if arguments_acc is positional: + pass # do nothing + elif arguments_acc is kwarg_only: + arguments_acc = out + else: + if arguments_acc is out: + raise AssertionError( + f"non-mutable argument '{parg.name}' cannot follow mutable out arguments" + ) + arguments_acc.append(parg) + + return positional, kwarg_only, out + + @staticmethod + def parse(args: str) -> Arguments: + """ + Input: 'int x, int y, int z' + """ + + # We do this in two phases. First we parse into three + # main categories: positional, kwarg_only, out. + # Then, we reparse positional and kwarg_only to separate + # out the self argument and tensor options arguments. + + positional, kwarg_only, out = Arguments._preparse(args) + + # Split self argument + self_ix = None + for i, a in enumerate(positional): + if a.name == "self": + self_ix = i + break + pre_self_positional: list[Argument] + self_arg: SelfArgument | None + post_self_positional: list[Argument] + if self_ix is not None: + pre_self_positional = positional[:self_ix] + self_arg = SelfArgument(positional[self_ix]) + post_self_positional = positional[self_ix + 1 :] + else: + pre_self_positional = [] + self_arg = None + post_self_positional = positional + + # Group tensor options arguments + pre_tensor_options_kwarg_only: list[Argument] = [] + tensor_options: TensorOptionsArguments | None = None + post_tensor_options_kwarg_only: list[Argument] = [] + kwarg_only_acc = pre_tensor_options_kwarg_only + + def pred(name: str, ty: Type) -> Callable[[Argument], bool]: + return lambda a: a.name == name and a.type in [ty, OptionalType(ty)] + + predicates = [ # order matters + pred("dtype", Type.parse("ScalarType")), + pred("layout", Type.parse("Layout")), + pred("device", Type.parse("Device")), + pred("pin_memory", Type.parse("bool")), + ] + + i = 0 + while i < len(kwarg_only): + # If there is enough space... + if i <= len(kwarg_only) - len(predicates): + # And the next len(predicates) arguments look like TensorOptions arguments + if all( + p(a) + for p, a in zip(predicates, kwarg_only[i : i + len(predicates)]) + ): + if kwarg_only_acc is not pre_tensor_options_kwarg_only: + raise AssertionError( + "tensor options arguments can only appear once" + ) + # Group them together as one argument + tensor_options = TensorOptionsArguments( + dtype=kwarg_only[i], + layout=kwarg_only[i + 1], + device=kwarg_only[i + 2], + pin_memory=kwarg_only[i + 3], + ) + i += len(predicates) + kwarg_only_acc = post_tensor_options_kwarg_only + continue + kwarg_only_acc.append(kwarg_only[i]) + i += 1 + + return Arguments( + pre_self_positional=tuple(pre_self_positional), + self_arg=self_arg, + post_self_positional=tuple(post_self_positional), + pre_tensor_options_kwarg_only=tuple(pre_tensor_options_kwarg_only), + tensor_options=tensor_options, + post_tensor_options_kwarg_only=tuple(post_tensor_options_kwarg_only), + out=tuple(out), + ) + + def __str__(self) -> str: + all_arguments: list[str] = [] + all_arguments.extend(map(str, self.flat_positional)) + if self.flat_kwarg_only or self.out: + all_arguments.append("*") + all_arguments.extend(map(str, self.flat_kwarg_only)) + all_arguments.extend(map(str, self.out)) + return ", ".join(all_arguments) + + def __post_init__(self) -> None: + # TODO: These invariants are weirdly asymmetric? + # TODO: Fancier types? + if self.self_arg is None: + if self.pre_self_positional: + raise AssertionError( + "pre_self_positional is non-empty but self_arg is None" + ) + if self.tensor_options is None: + if self.post_tensor_options_kwarg_only: + raise AssertionError( + "post_tensor_options_kwarg_only is non-empty but tensor_options is None" + ) + + # We don't allow any of the following to have argument annotations, + # to keep things simple. + mutable_pre_self_positionals = [ + a + for a in self.pre_self_positional + if a.annotation is not None and a.annotation.is_write + ] + if len(mutable_pre_self_positionals) != 0: + raise AssertionError( + f"mutable pre_self_positional arguments are not currently supported in the schema: {mutable_pre_self_positionals}" + ) + + +# Names that validly are __iXXX__ indicating inplace operations. +# Taken from https://www.python.org/dev/peps/pep-0203/#new-methods +# NB: PyTorch hasn't actually implemented all of these +AUGMENTED_ASSIGNMENT_NAMES = [ + "add", + "sub", + "mul", + "div", + "mod", + "pow", + "lshift", + "rshift", + "and", + "xor", + "or", +] + + +# A BaseOperatorName is what we think of the operator name, without +# the overload name. Unusually, we don't represent this as just a +# string; instead, we directly represent a few important semantic +# bits of information we derive from the string: namely whether +# or not it's inplace (add_) and whether or not it's a double-underscore +# method (__add__) +@dataclass(frozen=True) +class BaseOperatorName: + base: str + inplace: bool + dunder_method: bool + # Note [Overload Ambiguity With Functional Variants] + # A handful of operators have both a "mutable" and a "functional" variant. + # (native_batch_norm is a good example, although this isn't the case today). + # For those operators, the mutable and functional variant take in the same set of + # arguments, but have different alias annotations. + # this makes it ambiguous when you try to resolve an OverloadPacket into an overload, + # given a set of input arguments. + # + # So instead of making the "functional" variant in this case a real overload, e.g: + # native_batch_norm (mutable variant) + # native_batch_norm.functional (functional variant) + # we make it a new base operator, + # native_batch_norm_functional (functional variant) + # + # In an ideal world, we would probably invert this so the operators were: + # native_batch_norm.mutable (mutable variant) + # native_batch_norm (functional variant) + # + # Doing that is BC-breaking though, so we're stuck with the above modeling. + functional_overload: bool = False + + # NB: We don't officially support namespace in FunctionSchema, we treat this prefix + # as part of the base operator name, for __str__() to consume. + # The canonical input (from the rest of the infra) will not contain namespace, but + # we have a usecase in ExecuTorch where we want to support BaseOperatorName with namespace. + namespace: str | None = None + + @staticmethod + def parse(op: str) -> BaseOperatorName: + if op == "": + raise AssertionError("operator name cannot be empty") + if op.endswith("_out"): + raise AssertionError( + "_out suffix is reserved and not permitted for operator names; " + "did you mean to specify an out overload name instead?" + ) + # Extract namespace out. Base operator name may or may not contain namespace. + # E.g., aten::__lshift__ is a valid base operator name, __lshift__ is also valid. + # We want to split the namespace out from the base operator name. + match = re.match(r"^(?:(.*)::)?(.*)$", op) + namespace = match.group(1) if match else "" + op_without_ns = match.group(2) if match else op + m = re.match(r"^__([^_]+)__$", op_without_ns) + if m is not None: + dunder_method = True + base = m.group(1) + if any(base == f"i{n}" for n in AUGMENTED_ASSIGNMENT_NAMES): + inplace = True + base = base[1:] + else: + inplace = False + # temporary, this is not intrinsically true but + # has been historically true for dunder methods + # we support (but, if we ever got, say, __int__, this would + # be wrong!) + if base[0] == "i": + raise AssertionError( + f"unexpected dunder method starting with 'i': {op}" + ) + else: + dunder_method = False + base = op_without_ns + if base[-1] == "_": + inplace = True + base = base[:-1] + else: + inplace = False + + # See Note [Overload Ambiguity With Functional Variants] + functional_suffix = "_functional" + if base.endswith(functional_suffix): + functional_overload = True + base = base[: -len(functional_suffix)] + # This seems complicated and unnecessary, so banning dunder methods + # for now on ops that have a functional + mutable variant (like native_batch_norm). + if dunder_method or inplace: + raise AssertionError( + f"functional overload cannot be a dunder method or inplace: {op}" + ) + else: + functional_overload = False + + r = BaseOperatorName( + base=base, + inplace=inplace, + dunder_method=dunder_method, + functional_overload=functional_overload, + namespace=namespace, + ) + if str(r) != op: + raise AssertionError(f"{str(r)} != {op}") + return r + + def __str__(self) -> str: + namespace_prefix = f"{self.namespace}::" if self.namespace else "" + if self.dunder_method: + i = "i" if self.inplace else "" + return f"{namespace_prefix}__{i}{self.base}__" + else: + i = ( + "_" + if self.inplace + else "_functional" + if self.functional_overload + else "" + ) + return f"{namespace_prefix}{self.base}{i}" + + +# Operator name is the base operator name along with the (typically not +# user visible) overload string. +@dataclass(frozen=True) +class OperatorName: + name: BaseOperatorName + overload_name: str + + @staticmethod + def parse(op_name: str) -> OperatorName: + if "." in op_name: + name, overload_name = op_name.split(".", 1) + else: + name = op_name + overload_name = "" + r = OperatorName(name=BaseOperatorName.parse(name), overload_name=overload_name) + if str(r) != op_name: + raise AssertionError(f"{str(r)} != {op_name}") + return r + + def __str__(self) -> str: + if self.overload_name: + return f"{self.name}.{self.overload_name}" + else: + return f"{self.name}" + + # NB: This must be synchronized with the naming scheme in + # aten/src/ATen/templates/Operators.h + # Given a function schema "aten::op.overload(...)", + # If there is no overload name, this returns f"{op}" + # If there is an overload name, this returns f"{op}_{overload}" + def unambiguous_name(self) -> str: + if self.overload_name: + return f"{self.name}_{self.overload_name}" + else: + return f"{self.name}" + + def remove_inplace(self) -> OperatorName: + return OperatorName( + name=BaseOperatorName( + base=self.name.base, + inplace=False, + dunder_method=self.name.dunder_method, + ), + overload_name=self.overload_name, + ) + + def with_overload(self, overload: str) -> OperatorName: + return OperatorName( + name=BaseOperatorName( + base=self.name.base, + inplace=False, + dunder_method=self.name.dunder_method, + ), + overload_name=overload, + ) + + +def gets_generated_out_inplace_wrapper( + f: NativeFunction, g: NativeFunctionsGroup, b: BackendIndex +) -> bool: + return ( + f.func.kind() is not SchemaKind.functional + and not b.has_kernel(f) + and b.has_kernel(g.functional) + ) + + +# NativeFunction objects that are views (f.is_view_op returns True) +# are added into a `NativeFunctionsViewGroup`, which we can use to +# easily access the generated (optional) view_copy NativeFunction. +# It's convenient to group them together, so we pair them up in NativeFunctionsViewGroup. +# See Note [Codegen'd {view}_copy Operators] +# +# One property of this representation is that in order for a view-like op to be part of +# a NativeFunctionsViewGroup, the "aliasing" version of that view op must exist. +# There's one case where that doesn't happen: we have a non-aliasing `narrow_copy.out` op, +# but don't have corresponding aliasing `narrow.out` op. +# This means that `narrow_copy.out` won't appear as a NativeFunctionsViewGroup. +@dataclass(frozen=True) +class NativeFunctionsViewGroup: + view: NativeFunction + # Note: the {view}_copy operator is optional because we currently don't generate copy variants + # for all view ops. Notably, we don't generate them for CompositeImplicitAutograd views + # (we already get them "for free" through decomposition) + view_copy: NativeFunction | None + # view_inplace ops are also optional, but every view_inplace op should have out-of-place variant. + view_inplace: NativeFunction | None + + def __post_init__(self) -> None: + if not self.view.is_view_op: + raise AssertionError(f"view is not a view op: {self.view.func.name}") + if self.view_copy is None: + if gets_generated_view_copy(self.view): + raise AssertionError( + f"{str(self.view.func.name)} appears to be a new operator that aliases its inputs." + " The codegen expects you to add a corresponding operator to native_functions.yaml:" + f" {get_view_copy_name(self.view)!s}." + " See Note [view_copy NativeFunctions] for details." + ) + else: + if not self.view_copy.func.name.name.base.endswith(("_copy", "_scatter")): + raise AssertionError( + f"view_copy name must end with '_copy' or '_scatter': {self.view_copy.func.name}" + ) + if self.view.func.signature() != self.view_copy.func.signature( + strip_view_copy_name=True, + ): + view_sig = self.view.func.signature() + view_copy_sig = self.view_copy.func.signature(strip_view_copy_name=True) + raise AssertionError( + f"view and view_copy signatures don't match: {view_sig} != {view_copy_sig}" + ) + if "view_copy" not in self.view_copy.tags: + raise AssertionError( + f"{str(self.view_copy.func.name), str(self.view.tags)} appears to be a view_copy operator. The codegen expects" + " view_copy operators to be annotated with the 'view_copy' tag in native_functions.yaml." + " See Note [view_copy NativeFunction] for details." + ) + if self.view_inplace is not None: + if self.view.func.signature() != self.view_inplace.func.signature(): + view_sig = self.view.func.signature() + view_inplace_sig = self.view_inplace.func.signature() + raise AssertionError( + f"view and view_inplace signatures don't match: {view_sig} != {view_inplace_sig}" + ) + + if self.view.has_composite_implicit_autograd_kernel: + if self.view_inplace is not None: + if not self.view_inplace.has_composite_implicit_autograd_kernel: + raise AssertionError( + f"{str(self.view.func.name)} and {str(self.view_inplace.func.name)} must either" + " both have CompositeImplicitAutograd kernels, or both not have composite kernels." + ) + if self.view.has_composite_implicit_autograd_nested_tensor_kernel: + if self.view_inplace is not None: + if not self.view_inplace.has_composite_implicit_autograd_nested_tensor_kernel: + raise AssertionError( + f"{str(self.view.func.name)} and {str(self.view_inplace.func.name)} must either" + " both have CompositeImplicitAutogradNestedTensor kernels, or both not have composite kernels." + ) + + def functions(self, *, include_copy: bool = True) -> Iterator[NativeFunction]: + yield self.view + if self.view_inplace is not None: + yield self.view_inplace + if self.view_copy is not None and include_copy: + yield self.view_copy + + @property + def root_name(self) -> str: + return self.view.root_name + + @property + def composite(self) -> bool: + # We currently assert that the "group" is consistent. + # If the view op is composite, then its view_inplace op is too. + return self.view.has_composite_implicit_autograd_kernel + + +def gets_generated_view_copy(f: NativeFunction) -> bool: + # Only aliasing (view) operators get a copy variant. + if not f.is_view_op: + return False + # We don't need to bother generating copy variants for CompositeImplicitAutograd ops, + # because we can let them decompose into base view ops. + if f.has_composite_implicit_autograd_kernel: + return False + # We also don't need to generate copy variants for inplace views. + if "inplace_view" in f.tags: + return False + # Assume ops ending in _inverse have manually-defined copy variants + # (e.g. slice_inverse() has the copy variant slice_scatter()). + # We -could- probably generate these as well, but the codegen will be + # slightly different, and hand-writing these few kernels keeps codegen + # complexity lower. + if f.func.name.name.base.endswith("_inverse"): + return False + return True + + +# Given a NativeFunction that corresponds to a view op, +# returns the OperatorName of the corresponding "copy" variant of the op. +def get_view_copy_name(f: NativeFunction) -> OperatorName: + # Right now, when asking for a view op's corresponding "view_copy" name + # we assert for sanity that the op is allowed to have a generated view_copy variant. + # (We can do this because "gets_generated_view_copy()" tell us which ops get a generated view_copy op). + # However, narrow_copy() already exists as an op directly in native_functions.yaml. + # I'm hardcoding narrow_copy here for now to maintain the assert, + # But we could also just get rid of the assert. + list_of_ops_with_explicit_view_copy_operators = ["narrow"] + if str(f.func.name) not in list_of_ops_with_explicit_view_copy_operators: + if not gets_generated_view_copy(f): + raise AssertionError( + f"{f.func.name} does not have a generated view_copy variant" + ) + + base_name = f"{f.func.name.name.base}_copy" + view_copy_name = OperatorName( + name=BaseOperatorName( + base=base_name, inplace=False, dunder_method=f.func.name.name.dunder_method + ), + overload_name=f.func.name.overload_name, + ) + return view_copy_name + + +# Helper functions for parsing argument lists (both inputs and returns) + + +def parse_returns(return_decl: str) -> tuple[Return, ...]: + """ + Input: '()' + Output: [] + """ + if return_decl == "()": + return () + if return_decl[0] == "(" and return_decl[-1] == ")": + return_decl = return_decl[1:-1] + return tuple(Return.parse(arg) for arg in return_decl.split(", ")) + + +# A Precompute instance consists of a map from kernel argument name +# to the list of Argument instances that should replace that +# kernel argument in the impl function. +@dataclass(frozen=True) +class Precompute: + # A map from kernel argument name -> a list of precomputed + # elements that replaces/supersedes it. + replace: dict[str, list[Argument]] + # List of precomputed args added without replacement + add: list[Argument] + + @staticmethod + def parse(src: object) -> Precompute: + if not isinstance(src, list): + raise AssertionError(f"precomputed must be a list, got {type(src)}") + + # src is a list of strings of the format: + # {kernel param name} -> {replacement decl}[, {replacement decl}, ...] + # [{add decl}[, {add decl}, ...]] + # The last line is optional and contains the precomputed parameters that are + # added without replacement. + # The other lines are parsed to get the names of which precomputed elements + # should replace which kernel arguments. + add_args = [] + if " -> " not in src[-1]: + add_list = src[-1].split(",") + add_args = [Argument.parse(name.strip()) for name in add_list] + src = src[:-1] + + replace = {} + for raw_replace_item in src: + if not isinstance(raw_replace_item, str): + raise AssertionError( + f"precomputed item must be a str, got {type(raw_replace_item)}" + ) + if " -> " not in raw_replace_item: + raise AssertionError( + f"precomputed parameters without replacement are allowed only in the last line, got: {raw_replace_item}" + ) + + arg, with_list_raw = raw_replace_item.split(" -> ") + if " " in arg: + raise AssertionError( + f"illegal kernel param name '{arg}' in precomputed parameters" + ) + with_list = with_list_raw.split(",") + with_list_args = [Argument.parse(name.strip()) for name in with_list] + replace[arg] = with_list_args + + r = Precompute(replace=replace, add=add_args) + if r.to_list() != src: + raise AssertionError(f"r.to_list() != src: {r.to_list()} != {src}") + return r + + def __post_init__(self) -> None: + # the template parameters are upper so if these are the + # same then it is ambiguous + for a in self.add: + if a.name.upper() == a.name: + raise AssertionError( + f"precomputed argument name must not be all uppercase: {a.name}" + ) + for args in self.replace.values(): + for a in args: + if a.name.upper() == a.name: + raise AssertionError( + f"precomputed argument name must not be all uppercase: {a.name}" + ) + + def to_list(self) -> list[str]: + replace_list = [] + for kernel_param, replacement_params in self.replace.items(): + replacements = ", ".join(str(param) for param in replacement_params) + replace_list.append(f"{kernel_param} -> {replacements}") + + return replace_list diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/native_function_generation.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/native_function_generation.py new file mode 100644 index 0000000000000000000000000000000000000000..107802f60051e92e82031d75f56b2ad6d5db674a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/native_function_generation.py @@ -0,0 +1,662 @@ +from __future__ import annotations + +import string +from collections import defaultdict +from typing import TYPE_CHECKING + +import torchgen.api.dispatcher as dispatcher +from torchgen.api.translate import translate +from torchgen.api.types import Binding, DispatcherSignature, Expr +from torchgen.context import with_native_function +from torchgen.model import ( + Annotation, + Argument, + BackendIndex, + BackendMetadata, + BaseOperatorName, + BaseTy, + BaseType, + DEFAULT_KERNEL_NAMESPACE, + DeviceCheckType, + DispatchKey, + FunctionSchema, + NativeFunction, + NativeFunctionsGroup, + OperatorName, + Return, + SchemaKind, + Variant, +) +from torchgen.utils import concatMap + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# See Note: [Out ops with functional variants that don't get grouped properly] +OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY = [ + # This has a functional variant, but it's currently marked private. + # This function should be marked private as well (*_backward ops aren't exposed to python anyway). + "adaptive_avg_pool3d_backward.grad_input", + # There's a functional variant, _slow_conv2d_backward.output_mask, that isn't grouped properly. + # Maybe we can kill this operator in favor of convolution_backward? + "_slow_conv2d_backward.grad_input", +] + + +# See Note: [Mutable ops that cannot get an out variant] +MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT = [ + # should be out=? + "_cummax_helper", + # should be out=? + "_cummin_helper", +] + +# All of these operators don't have any tensor like returns +FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT = [ + "_assert_async", # no return + "_assert_async.msg", # no return + "_assert_tensor_metadata", # no return + "_cslt_sparse_mm_search", # returns an int + "_assert_scalar", # no return + "_dimI", # returns an int + "_dimV", # returns an int + "_has_same_storage_numel", # returns a boolean + "_linalg_check_errors", # no return + "_local_scalar_dense", # returns a Scalar + "_nested_tensor_from_mask_left_aligned", # returns a boolean + "_nnz", # returns an int + "_use_cudnn_ctc_loss", # returns a boolean + "_use_cudnn_ctc_loss.Tensor", # returns a boolean + "_use_miopen_ctc_loss", # returns a boolean + "_use_miopen_ctc_loss.Tensor", # returns a boolean + "_validate_compressed_sparse_indices", # no return + "allclose", # returns a boolean + "dense_dim", # returns an int + "equal", # returns a boolean + "is_coalesced", # returns an boolean + "is_pinned", # returns a boolean + "is_same_size", # returns a boolean + "is_set_to", # returns a boolean + "q_per_channel_axis", # returns an int + "q_scale", # returns a float + "q_zero_point", # returns an int + "qscheme", # returns a QScheme + "record_stream", # no return + "sparse_dim", # returns an int + "sym_constrain_range", # no return + "sym_constrain_range_for_size", # no return + "_nested_tensor_storage_offsets", # returns a vector of ints + "_chunk_grad_outputs_efficient_attention", # returns a bool + "_fused_sdp_choice", # returns an int + "_print", # no return + "_sink_tokens", # no return + "_nested_get_ragged_idx", # returns an int +] + +INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY = [ + # polygamma and polygamma.out both exist, but have a + # pre-self arg (while polygamma_ does not) + # We should either fix this schema so it can be grouped properly, + # or allow the codegen to generate new functional/out= NativeFunctions for this op + # (which would require changing its overload name to prevent overload ambiguity). + "polygamma_" +] + + +# Groups "similar" NativeFunctions together +# example add.Tensor, add_.Tensor, add.out +# "similar" NativeFunctions are all expected to have an identical `signature()`, +# But have differing SchemaKinds. +def pre_group_native_functions( + native_functions: Sequence[NativeFunction], +) -> dict[FunctionSchema, dict[SchemaKind, NativeFunction]]: + pre_grouped_native_functions: dict[ + FunctionSchema, dict[SchemaKind, NativeFunction] + ] = defaultdict(dict) + for f in native_functions: + d = pre_grouped_native_functions[f.func.signature()] + if f.func.kind() in d: + raise AssertionError(f"Duplicate schema kind {f.func.kind()} for {f.func}") + d[f.func.kind()] = f + return pre_grouped_native_functions + + +# Returns the out variant overload name given a base function overload name +def get_expected_out_variant_overload_name(overload_name: str | None) -> str: + return "out" if not overload_name else f"{overload_name}_out" + + +# Helper function: given an inplace FunctionSchema, generate its corresponding out= variant +# Example before: +# _add_relu_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!) +# Example after: +# _add_relu.Scalar_out(Tensor self, Scalar other, Scalar alpha=1, *, Tensor(a!) out) +def self_to_out_signature(func: FunctionSchema) -> FunctionSchema: + # Generating an out= schema from an inplace schema. + if func.kind() != SchemaKind.inplace: + raise AssertionError(f"Expected inplace schema, got {func.kind()}") + if func.arguments.self_arg is None: + raise AssertionError("Expected self_arg to be non-None") + # The new out= schema has: + # - a new out argument with the same type as "func" (but with a mutable annotation) + # - The returns (if any) now alias the out= argument instead of "func" + # - an "out" overload name + return FunctionSchema( + name=func.name.remove_inplace().with_overload( + get_expected_out_variant_overload_name(func.name.overload_name) + ), + arguments=func.arguments.remove_self_annotation().with_out_args( + [ + Argument( + name="out", + type=func.arguments.self_arg.argument.type, + default=None, + annotation=func.arguments.self_arg.argument.annotation, + ) + ] + ), + returns=func.returns, + ) + + +# Helper function: given a functional FunctionSchema, generate its corresponding out= variant +# Example before: +# _to_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, +# bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor +# Example after: +# _to_copy._out(Tensor self, *, bool non_blocking=False, MemoryFormat? memory_format=None, +# Tensor(a!) out) -> Tensor(a!) +def functional_to_out_signature(func: FunctionSchema) -> FunctionSchema: + # Generating an out= schema from a functional schema. + if func.kind() != SchemaKind.functional: + raise AssertionError(f"Expected functional schema, got {func.kind()}") + + new_returns, new_out_args = generate_out_args_from_schema(func) + # The new out= schema has: + # - one or more new out argument(s) with the same type as returns (but with a mutable annotation) + # - The returns now alias the out= arguments + # - an "_out" overload name + return FunctionSchema( + name=func.name.with_overload( + get_expected_out_variant_overload_name(func.name.overload_name) + ), + arguments=func.arguments.signature().with_out_args( + new_out_args, + ), + returns=tuple(new_returns), + ) + + +# Helper function: given a function schema, generate corresponding out arguments, also the updated return annotations. +def generate_out_args_from_schema( + func: FunctionSchema, +) -> tuple[list[Return], list[Argument]]: + # More of a sanity check - our existing restrictions on schemas should enforce that + # mutable schema kinds never return their mutable arguments. + if any(r.annotation is not None and r.annotation.is_write for r in func.returns): + raise AssertionError("Mutable schema kinds should not return mutable arguments") + + tensorlike_rets = [r for r in func.returns if r.type.is_tensor_like()] + if len(tensorlike_rets) == 0: + raise AssertionError("Expected at least one tensor-like return") + + used_annotations = concatMap( + lambda a: [] if a.annotation is None else a.annotation.alias_set, + func.arguments.flat_all, + ) + valid_annotations = [x for x in string.ascii_lowercase if x not in used_annotations] + + all_rets_are_tensors = all(r.type == BaseType(BaseTy.Tensor) for r in func.returns) + + new_out_args: list[Argument] = [] + # The end result of new_returns is that: + # - If every return is a plain tensor, then the new returns == the old returns, but with the out= alias annotations added. + # - Otherwise, none of the out arguments show up in the returns (and we're only left with non-tensor-like returns, if any). + new_returns: list[Return] = [] + for i, r in enumerate(func.returns): + if r.type.is_tensor_like(): + new_out = Argument( + name="out" if len(func.returns) == 1 else f"out{i}", + type=r.type, + default=None, + annotation=Annotation.parse(f"{valid_annotations[i]}!"), + ) + new_out_args.append(new_out) + if all_rets_are_tensors: + # The convention for out= schemas is that they only return their out arguments + # if the return is a plain Tensor (or if it's a tuple of plain Tensors) + new_ret = Return( + name=None, type=new_out.type, annotation=new_out.annotation + ) + new_returns.append(new_ret) + else: + new_returns.append(r) + return new_returns, new_out_args + + +# Helper function: given a mutable FunctionSchema, generate its corresponding out= variant +# Example before: +# _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask) # noqa: B950 +# Example after: +# _fused_moving_avg_obs_fq_helper._out(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False, *, Tensor(e!) out0, Tensor(f!) out1) -> (Tensor(e!), Tensor(f!)) # noqa: B950 +def mutable_to_out_signature(func: FunctionSchema) -> FunctionSchema: + # Generating an out= schema from a mutable schema. + if func.kind() != SchemaKind.mutable: + raise AssertionError(f"Expected mutable schema, got {func.kind()}") + # The new out= schema has: + # - Any non-aliased tensor-like returns are converted to mutable, aliased out= arguments + # (if the argument is a tensor then we also return it for method chaining, + # otherwise we return nothing) + # - an "out" overload name + # + # Note that: + # (1) This also means that we can *only* generate an out= variant from a mutable schema + # if the mutable schema has at least one tensor-like non-aliasing return. + # (2) The generated out= variant still has mutable positional arguments, + # but if necessary we could probably add another out= variant that also + # functionalizes the mutable arguments (a functional_out variant) + + new_returns, new_out_args = generate_out_args_from_schema(func) + + return FunctionSchema( + name=func.name.remove_inplace().with_overload( + get_expected_out_variant_overload_name(func.name.overload_name) + ), + arguments=func.arguments.with_out_args(new_out_args), + returns=tuple(new_returns), + ) + + +# This function, given function of one SchemaKind, as well as a target SchemaKind, +# generates a new NativeFunction with the same properties, but using the target SchemaKind. +# We only actually generate functions for either functional or out= SchemaKinds. +# This function returns a tuple, with: +# - The generated NativeFunction +# - a dictionary of `BackendIndex` objects, describing which dispatch keys +# we will generate kernels for, for the new NativeFunction. +# Details are in the function, but we only generate composite kernels (in some cases) today. +def generate_function( + f: NativeFunction, k: SchemaKind +) -> tuple[NativeFunction, dict[DispatchKey, dict[OperatorName, BackendMetadata]]]: + from torchgen.api import cpp + + if k == SchemaKind.functional: + if f.func.kind() == SchemaKind.functional: + raise AssertionError("Cannot generate functional from functional schema") + # The new "functional" NativeFunction has: + # - any mutable arguments have been converted into (immutable) returns. + # (if a mutable argument was not also a return, it gets converted to one) + # - "_functional" appended to the base name, ONLY IF this op has a mutable variant. + # See Note [Overload Ambiguity With Functional Variants] + # The default grouping logic in signature() actually already does this, + # so we can piggy-back off it (but we still want return names) + func = f.func.signature(keep_return_names=True).with_name( + OperatorName( + name=BaseOperatorName( + base=f.func.name.name.base, + inplace=False, + dunder_method=f.func.name.name.dunder_method, + # See Note [Overload Ambiguity With Functional Variants] + functional_overload=f.func.kind() == SchemaKind.mutable, + ), + overload_name=f.func.name.overload_name, + ) + ) + elif k == SchemaKind.out: + # We generate out= ops mostly just so that we can pair up NativeFunctions into groups easily, + # but at least today, there is no good reason to actually use them. + # we'll generate a dispatcher entry for them, but won't actually register any kernels for them. + if f.func.kind() == SchemaKind.inplace: + func = self_to_out_signature(f.func) + elif f.func.kind() == SchemaKind.mutable: + func = mutable_to_out_signature(f.func) + elif f.func.kind() == SchemaKind.functional: + func = functional_to_out_signature(f.func) + else: + raise AssertionError( + "We only bother generating out= functions from either inplace or mutable or functional variants" + ) + else: + raise AssertionError( + "We currently only generate either functional or out= NativeFunctions" + ) + + # Generated kernel naming convention for out: _. The reason for this is to + # disambiguate operator with the same name but different overload name, e.g., `randn.names_out` and + # `randn.generator_with_names_out`. + kernel_name = ( + func.name.unambiguous_name() + if func.kind() == SchemaKind.out + else cpp.name(func) + ) + if f.func.has_symint(): + kernel_name += "_symint" + backend_metadata = { + DispatchKey.CompositeExplicitAutograd: { + func.name: BackendMetadata( + kernel=kernel_name, + structured=False, + cpp_namespace=DEFAULT_KERNEL_NAMESPACE, + ) + } + } + tags = {"generated"} | set( + f.tags & {"nondeterministic_seeded", "view_copy", "pt2_compliant_tag"} + ) + + return ( + NativeFunction( + func=func, + use_const_ref_for_mutable_tensors=f.use_const_ref_for_mutable_tensors, + # These generated fn's aren't meant to be user friendly- don't generate methods. + variants={Variant.function}, + structured=False, + structured_delegate=None, + structured_inherits=None, + precomputed=None, + autogen=[], + ufunc_inner_loop={}, + manual_kernel_registration=False, + manual_cpp_binding=False, + python_module=None, + category_override=None, + device_guard=False, + device_check=DeviceCheckType.NoCheck, + loc=f.loc, + cpp_no_default_args=set(), + is_abstract=f.is_abstract, + has_composite_implicit_autograd_kernel=False, + has_composite_implicit_autograd_nested_tensor_kernel=False, + has_composite_explicit_autograd_kernel=True, + has_composite_explicit_autograd_non_functional_kernel=False, + # Every generated NativeFunction gets a "generated" tag, so it's easy to tell + # which NativeFunction objects did not come directly from native_functions.yaml. + tags=tags, + namespace=f.namespace, + ), + backend_metadata, + ) + + +# This function is responsible for adding generated NativeFunctions which don't appear +# explicitly in the codegen. +# You can inspect the full list of NativeFunctions yourself with the torchgen package, by running +# torchgen.parse_native_yaml("aten/src/ATen/native/native_functions.yaml", "aten/src/ATen/native/tags.yaml") +# (Maybe we should make a friendly API for this) +# +# Note: this function *mutates* its two inputs, +# adding the new NativeFunctions / BackendMetadata to them +def add_generated_native_functions( + rs: list[NativeFunction], + indices: dict[DispatchKey, dict[OperatorName, BackendMetadata]], +) -> None: + # The main code for generating new NativeFunctions + # First we group of NativeFunctions by schema kind, + # then we detect which ones are missing and generate them. + pre_grouped_native_functions = pre_group_native_functions(rs) + for d in pre_grouped_native_functions.values(): + has_functional = SchemaKind.functional in d + has_inplace = SchemaKind.inplace in d + has_mutable = SchemaKind.mutable in d + has_out = SchemaKind.out in d + is_core = any("core" in variant.tags for variant in d.values()) + + # We automatically generate a few native functions that don't exist in the yaml, for a few reasons: + # (1) If an operator has an inplace/out= variant but no functional variant, we can generate + # a simple functional variant that the functionalization pass can consume. + # (2) If an operator has an inplace or functional but no out= variant, we generate an out= + # variant, mostly so we can easily pair up functions into NativeFunctionsGroup, + # while maintaining the constraint that the out= variant is "required". + if has_mutable or has_inplace or has_out or has_functional: + # Don't bother generating functions trio's for native functions that bypass the dispatcher. + are_manual = all(f.manual_cpp_binding for f in d.values()) + # Don't bother generating functional + out= variants for view operators + # set_ is technically an inplace_view, but for now it is treated + # as a normal inplace op in the codegen + has_view_ops = any( + f.is_view_op and str(f.func.name.name) != "set_" for f in d.values() + ) + # Don't generate the other variants for non-core CompositeImplicitAutograd operators. + # We could probably do this, but the main benefit of generating the function triplets + # is for transforms that need them, and transforms don't need to act directly + # on CompositeImplicitAutograd operators (since we let them decompose). + are_composite_implicit = all( + f.has_composite_implicit_autograd_kernel for f in d.values() + ) + if are_manual or has_view_ops or are_composite_implicit and not is_core: + continue + if has_out and len(d.values()) == 1: + # Note: [Out ops with functional variants that don't get grouped properly] + # In theory we could validly have an out= operator in native_functions.yaml + # that has no other variants. + # But today, all of the operators where that's the case actually do have + # functional variants, that we are just unable to pair up properly. + # I think banning this all together is probably safer + # (you can always add a functional variant yourself if you want to add a new out= operator). + # + # We should probably fix the existing cases; this check is to prevent us from adding more over time. + if ( + str(d[SchemaKind.out].func.name) + not in OUT_OPS_THAT_DONT_GET_GROUPED_PROPERLY + ): + raise AssertionError( + f"Found an out= operator that we could not find any other variants of: {str(d[SchemaKind.out].func)}" + ) + continue + + # Some inplace ops that have problematic schemas (that we should fix), which prevent us + # from generating out= and functional variants + if ( + has_inplace + and str(d[SchemaKind.inplace].func.name) + in INPLACE_OPS_THAT_DONT_GET_GROUPED_PROPERLY + ): + continue + + base_fn = ( + d[SchemaKind.mutable] + if has_mutable + else d[SchemaKind.inplace] + if has_inplace + else d[SchemaKind.out] + if has_out + else d[SchemaKind.functional] + ) + + # Note: [Mutable ops that cannot get an out variant] + # We can only generate an out= variant if either: + # - the original function has tensor-like returns (since we can convert them to out kwargs) + # - or it's inplace (since we can convert `self` to an out kwarg) + # There are only two functions that don't fit this criteria today though, + # and they both look like they should be fixed to be out= variants, + # so if feels safer to ban this schema all-together + base_fn_valid = base_fn.func.kind() == SchemaKind.inplace or any( + r.type.is_tensor_like() for r in base_fn.func.returns + ) + # Note: [Loosen the assertion that all functional should have out variant] + # By design all functional operators should have our variants. The needs_out check + # is loosening this requirement, changing it to only generate out variant if there's + # an `autogen` block in the native function, in the long run it should be removed. + # FIXME: Remove this after figuring out CI job failures related to min, max, mean + needs_out = any("out" in str(op_name) for op_name in base_fn.autogen) + gets_out_variant = not has_out and base_fn_valid and needs_out + if not has_out and not base_fn_valid: + if ( + str(base_fn.func.name) + not in MUTABLE_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT + and str(base_fn.func.name) + not in FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT + ): + raise AssertionError( + f"""Found an operator that we could not generate an out= variant for: {str(base_fn.func)}. +This type of operators don't have tensor-like return, making it difficult to generate a proper out= variant. If +out= variant is not needed, please add the function name into FUNCTIONAL_OPS_THAT_CANNOT_GET_AN_OUT_VARIANT list.""" + ) + + # Generate an out= variant + if gets_out_variant: + fn, metadata = generate_function(base_fn, SchemaKind.out) + d[SchemaKind.out] = fn + BackendIndex.grow_index(indices, metadata) + rs.append(fn) + + # Generate a functional variant, but only do it if the operator got an out= variant + # (Functional variants are only useful if we can group up the variants, + # which we can only do if they have an out= variant) + if not has_functional and (has_out or gets_out_variant): + fn, metadata = generate_function(base_fn, SchemaKind.functional) + d[SchemaKind.functional] = fn + BackendIndex.grow_index(indices, metadata) + rs.append(fn) + + +def return_str(rets: tuple[Return, ...], names: list[str]) -> str: + if len(rets) != len(names): + raise AssertionError( + f"Returns and names length mismatch: {len(rets)} vs {len(names)}" + ) + if len(rets) == 0: + return "" + elif len(rets) == 1: + return f"return {names[0]};" + else: + return f"return {dispatcher.returns_type(rets).cpp_type()}({', '.join(names)});" + + +# Given a function, and the name of a variable corresponding to the output of that function, +# gather up all of the individual returns that are not aliased +def gather_nonaliased_inner_rets(func: FunctionSchema, out_var: str) -> list[str]: + aliased_rets = func.aliased_return_names() + non_aliased_names = [] + is_out_var_a_tuple = len(func.returns) > 1 + for i, r in enumerate(aliased_rets): + if r is None: + non_aliased_names.append( + f"std::get<{i}>({out_var})" if is_out_var_a_tuple else out_var + ) + return non_aliased_names + + +# Generates functional kernels in terms of their inplace.mutable counterparts. +# We only do this for "generated" NativeFunctions +@with_native_function +def gen_composite_functional_kernel(g: NativeFunctionsGroup) -> str | None: + # We should only be generating these for code-generated NativeFunctions + if "generated" not in g.functional.tags: + return None + # And we always write the kernel for a generated op in terms of a non-generated op. + if g.inplace is not None and "generated" not in g.inplace.tags: + target_f = g.inplace + elif g.mutable is not None and "generated" not in g.mutable.tags: + target_f = g.mutable + else: + # We should be guaranteed to have a valid inplace/mutable variant to call into. + # See Note: [Mutable Ops Not Using Functionalization] + raise AssertionError(str(g.functional.func)) + + sig = DispatcherSignature(g.functional.func) + target_sig = DispatcherSignature(target_f.func) + + context: list[Binding | Expr] = [] + clone_mutable_inputs = [] + cloned_return_names = [] + # We can't just directly pass all of the arguments from the functional op into the mutating op. + # We need to check for which inputs to the mutating operator are mutable, + # and clone those inputs first. + for a_curr, a_tgt in zip( + dispatcher.jit_arguments(g.functional.func), + dispatcher.jit_arguments(target_f.func), + ): + if a_tgt.annotation is not None and a_tgt.annotation.is_write: + clone_mutable_inputs.append( + f"auto {a_curr.name}_clone = clone_arg({a_curr.name});" + ) + context.append( + Expr( + expr=f"{a_curr.name}_clone", + type=dispatcher.argument_type(a_curr, binds=a_curr.name), + ) + ) + # Invariant: mutable arguments on the inner mutable op are always returns on the functional op. + cloned_return_names.append(f"{a_curr.name}_clone") + else: + context.append(dispatcher.argument(a_curr)) + exprs = ", ".join([e.expr for e in translate(context, target_sig.arguments())]) + + out_name = "output" + maybe_assign = f"auto {out_name} = " if len(target_f.func.returns) > 0 else "" + inner_return_names = gather_nonaliased_inner_rets(target_f.func, out_name) + ret_str = return_str( + g.functional.func.returns, inner_return_names + cloned_return_names + ) + + clone_mutable_inputs_str = "\n".join(clone_mutable_inputs) + return f""" +{sig.defn(name=sig.name() + ("_symint" if g.out.func.has_symint() else ""))} {{ + {clone_mutable_inputs_str} + {maybe_assign}at::_ops::{target_f.func.name.unambiguous_name()}::call({exprs}); + {ret_str} +}} +""" + + +# Generates out= kernels in terms of their functional counterparts. +# We only do this for "generated" NativeFunctions +@with_native_function +def gen_composite_out_kernel(g: NativeFunctionsGroup) -> str | None: + # We should only be generating these for code-generated NativeFunctions + if "generated" not in g.out.tags: + return None + # And we always write the kernel for the out= op in terms of the functional. + # Note that the functional op might have also been generated, but we don't have to + # worry about cycles, because the generated functional kernels are always implemented + # in terms of non-generated kernels (see gen_composite_functional_kernel). + + sig = DispatcherSignature(g.out.func) + target_sig = DispatcherSignature(g.functional.func) + + exprs = ", ".join( + [e.expr for e in translate(sig.arguments(), target_sig.arguments())] + ) + + copy_outs = [] + out_name = "tmp_output" + for i, out_arg in enumerate(g.out.func.arguments.out): + functional_return_name = ( + out_name + if len(g.functional.func.returns) == 1 + else f"std::get<{i}>({out_name})" + ) + copy_outs.append( + f"""\ + resize_out_helper({out_arg.name}, {functional_return_name}); + copy_arg({out_arg.name}, {functional_return_name});""" + ) + + rets = [] + # For each return arg in the calling (out=) operator, + # If it corresponds to an aliased input, return the input. + # Otherwise, return the corresponding output from calling the functional operator. + for i, ret_name in enumerate(g.out.func.aliased_return_names()): + if ret_name is not None: + rets.append(ret_name) + else: + functional_return_name = ( + out_name + if len(g.functional.func.returns) == 1 + else f"std::get<{i}>({out_name})" + ) + rets.append(functional_return_name) + + copy_outs_str = "\n".join(copy_outs) + + # Kernel name needs to follow the naming convention defined in `generate_function()` + return f""" +{sig.defn(name=g.out.func.name.unambiguous_name() + ("_symint" if g.out.func.has_symint() else ""))} {{ + auto {out_name} = at::_ops::{g.functional.func.name.unambiguous_name()}::call({exprs}); + {copy_outs_str} + {return_str(g.out.func.returns, rets)} +}} +""" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/operator_versions/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/operator_versions/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/operator_versions/gen_mobile_upgraders.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/operator_versions/gen_mobile_upgraders.py new file mode 100644 index 0000000000000000000000000000000000000000..0521e49f772973eaf7fb926b55a9c63cb1108327 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/operator_versions/gen_mobile_upgraders.py @@ -0,0 +1,390 @@ +#!/usr/bin/env python3 + +from __future__ import annotations + +import os +from enum import Enum +from operator import itemgetter +from pathlib import Path +from typing import Any + +import torch +from torch.jit.generate_bytecode import generate_upgraders_bytecode +from torchgen.code_template import CodeTemplate +from torchgen.operator_versions.gen_mobile_upgraders_constant import ( + MOBILE_UPGRADERS_HEADER_DESCRIPTION, +) + + +class ByteCode(Enum): + instructions = 1 + constants = 2 + types = 3 + operators = 4 + register_size = 5 + + +EXCLUDED_OP_SET = [ + "aten::full.names", + "aten::full.out", + "aten::full", +] + +EXCLUE_UPGRADER_SET = ["full_0_4", "full_out_0_4"] + +ONE_INSTRUCTION = CodeTemplate( + """ + Instruction{OpCode::${operator_name}, ${X}, ${N}},""" +) + +INSTRUCTION_LIST = CodeTemplate( + """std::vector({ + ${instruction_list} + }), // instructions list""" +) + +ONE_CONSTANT = CodeTemplate( + """ + c10::IValue(${constant}),""" +) + +CONSTANT_LIST = CodeTemplate( + """std::vector({ + ${constant_list} + }), // constants list""" +) + +CONSTANTS_LIST_EMPTY = """std::vector(), // constants list""" + +ONE_TYPE = CodeTemplate("""c10::parseType("${type_str}"),""") + +TYPE_LIST = CodeTemplate( + """std::vector({ + ${type_list} + }), // types list""" +) + +TYPE_LIST_EMPTY = """std::vector(), // types list""" + +ONE_OPERATOTR_STRING = CodeTemplate( + """ + OperatorString({"${operator_name}", "${overload_name}", ${num_of_args}}),""" +) + +OPERATOR_STRING_LIST = CodeTemplate( + """ + std::vector({ + ${operator_string_list} + }), // operators list""" +) + +ONE_UPGRADER_FUNCTION = CodeTemplate( + """ + mobile::Function::registerFunc( + "${upgrader_name}", + ${instruction_list}, + ${constant_list}, + ${type_list}, + ${register_size} + )""" +) + +ONE_UPGRADER_SRC = CodeTemplate( + """ + ByteCodeFunctionWithOperator({ + ${bytecode_function}, + ${operator_string_list} + }),""" +) + + +ONE_UPGRADER_IN_VERSION_MAP = CodeTemplate( + """Upgrader({${upgrader_min_version}, ${upgrader_max_version}, "${upgrader_name}", ${bytecode_func_index}})""" +) # noqa: E501 + +ONE_OPERATOR_IN_VERSION_MAP = CodeTemplate( + """ + {std::string("${operator_name}"), + std::vector({ + ${upgrader_list_in_version_map} + })},""" +) + + +OPERATOR_VERSION_MAP = CodeTemplate( + """ +const std::unordered_map> +getOperatorVersionMapForMobile() { + static std::unordered_map> + operatorVersionMapForMobile({ + ${operator_list_in_version_map} + }); + return operatorVersionMapForMobile; +} +""" +) + + +UPGRADER_CPP_SRC = CodeTemplate( + MOBILE_UPGRADERS_HEADER_DESCRIPTION + + """ +#include +#include +#include + +namespace torch { +namespace jit { + +// clang-format off + +// From operator_versions_map +${operator_version_map} + +const std::vector& getUpgraderBytecodeList() { + auto generate_upgrader_bytecode_list = []() { + std::vector upgrader_function_list({ + ${upgrader_bytecode} + }); + for (const auto& upgrader_function : upgrader_function_list) { + for (const auto& op : upgrader_function.operators) { + upgrader_function.function.append_operator( + op.name, + op.overload_name, + op.num_specified_args); + } + } + return upgrader_function_list; + }; + static std::vector upgraderBytecodeList = + generate_upgrader_bytecode_list(); + return upgraderBytecodeList; +} + +// clang-format on + +} // namespace jit +} // namespace torch +""" +) + +UPGRADER_MOBILE_FILE_NAME = "upgrader_mobile.cpp" + +UPGRADER_ELEMENT = CodeTemplate( + """\ +Upgrader({${min_version}, ${max_version}, ${operator_name}, ${index}}), +""" +) + +PER_OPERATOR_UPGRADER_LIST = CodeTemplate( + """\ +{ + std::string(${operator_name}), + std::vector({${upgrader_list}}); +} +""" +) + + +def construct_instruction(instruction_list_from_yaml: list[Any]) -> str: + instruction_list_part = [ + ONE_INSTRUCTION.substitute( + operator_name=instruction[0], + X=instruction[1], + N=instruction[2], + ) + for instruction in instruction_list_from_yaml + ] + return INSTRUCTION_LIST.substitute( + instruction_list="".join(instruction_list_part).lstrip("\n") + ) + + +def construct_constants(constants_list_from_yaml: list[Any]) -> str: + constants_list_part = [] + for constant_from_yaml in constants_list_from_yaml: + convert_constant = None + if isinstance(constant_from_yaml, str): + # Add quotes if it's string + convert_constant = f'"{constant_from_yaml}"' + elif isinstance(constant_from_yaml, bool): + convert_constant = "true" if constant_from_yaml else "false" + elif constant_from_yaml is None: + convert_constant = "" + elif isinstance(constant_from_yaml, int): + convert_constant = str(constant_from_yaml) + else: + raise ValueError( + f"The type of {constant_from_yaml} is {type(constant_from_yaml)}. " + "Please add change in construct_constants function in gen_mobile_upgraders.py." + ) + constants_list_part.append(ONE_CONSTANT.substitute(constant=convert_constant)) + if len(constants_list_part) == 0: + return CONSTANTS_LIST_EMPTY + return CONSTANT_LIST.substitute( + constant_list="".join(constants_list_part).lstrip("\n") + ) + + +def construct_operators(operator_list_from_yaml: list[Any]) -> str: + operator_list_part = [ + ONE_OPERATOTR_STRING.substitute( + operator_name=operator[0], + overload_name=operator[1], + num_of_args=operator[2], + ) + for operator in operator_list_from_yaml + ] + return OPERATOR_STRING_LIST.substitute( + operator_string_list="".join(operator_list_part).lstrip("\n") + ) + + +def construct_types(types_tr_list_from_yaml: list[Any]) -> str: + types_tr_list_part = [ + ONE_TYPE.substitute(type_str=types_tr) for types_tr in types_tr_list_from_yaml + ] + if len(types_tr_list_part) == 0: + return TYPE_LIST_EMPTY + return TYPE_LIST.substitute(type_list="".join(types_tr_list_part).lstrip("\n")) + + +def construct_register_size(register_size_from_yaml: int) -> str: + if not isinstance(register_size_from_yaml, int): + raise ValueError( + f"Input register size is {register_size_from_yaml} and" + "it's type is {type(register_size_from_yaml)}. An int type is expected." + ) + return str(register_size_from_yaml) + + +def construct_version_maps( + upgrader_bytecode_function_to_index_map: dict[str, Any], +) -> str: + version_map = torch._C._get_operator_version_map() + sorted_version_map_ = sorted(version_map.items(), key=itemgetter(0)) # type: ignore[no-any-return] + sorted_version_map = dict(sorted_version_map_) + + operator_list_in_version_map_part = [] + for op_name in sorted_version_map: + upgraders_in_version_map_part = [] + # TODO: remove the skip after these two operators schemas are fixed + if op_name in EXCLUDED_OP_SET: + continue + upgrader_ranges = torch._C._get_upgrader_ranges(op_name) + upgrader_entries = sorted_version_map[op_name] + if len(upgrader_ranges) != len(upgrader_entries): + raise AssertionError( + f"upgrader_ranges and upgrader_entries length mismatch for {op_name}: " + f"{len(upgrader_ranges)} != {len(upgrader_entries)}" + ) + for idx, upgrader_entry in enumerate(upgrader_entries): + upgrader_name = upgrader_entry.upgrader_name + bytecode_function_index = upgrader_bytecode_function_to_index_map[ + upgrader_name + ] + upgraders_in_version_map_part.append( + ONE_UPGRADER_IN_VERSION_MAP.substitute( + upgrader_min_version=upgrader_ranges[idx].min_version, + upgrader_max_version=upgrader_ranges[idx].max_version, + upgrader_name=upgrader_name, + bytecode_func_index=bytecode_function_index, + ) + ) + operator_list_in_version_map_part.append( + ONE_OPERATOR_IN_VERSION_MAP.substitute( + operator_name=op_name, + upgrader_list_in_version_map="".join(upgraders_in_version_map_part), + ) + ) + return OPERATOR_VERSION_MAP.substitute( + operator_list_in_version_map="".join(operator_list_in_version_map_part).lstrip( + "\n" + ) + ) + + +def get_upgrader_bytecode_function_to_index_map( + upgrader_dict: list[dict[str, Any]], +) -> dict[str, Any]: + upgrader_bytecode_function_to_index_map = {} + index = 0 + for upgrader_bytecode in upgrader_dict: + for upgrader_name in upgrader_bytecode: + if upgrader_name in EXCLUE_UPGRADER_SET: + continue + upgrader_bytecode_function_to_index_map[upgrader_name] = index + index += 1 + return upgrader_bytecode_function_to_index_map + + +def write_cpp(cpp_path: str, upgrader_dict: list[dict[str, Any]]) -> None: + upgrader_bytecode_function_to_index_map = ( + get_upgrader_bytecode_function_to_index_map(upgrader_dict) + ) + version_map_src = construct_version_maps(upgrader_bytecode_function_to_index_map) + all_upgrader_src_string = [] + for upgrader_bytecode in upgrader_dict: + for upgrader_name, bytecode in upgrader_bytecode.items(): + # TODO: remove the skip after these two operators schemas are fixed + if upgrader_name in EXCLUE_UPGRADER_SET: + continue + instruction_list_str = "" + constant_list_str = "" + type_list_str = "" + register_size_str = "" + operator_list_str = "" + for table_name, contents in bytecode.items(): + element = ByteCode[table_name] + if element is ByteCode.instructions: + instruction_list_str = construct_instruction(contents) + elif element is ByteCode.constants: + constant_list_str = construct_constants(contents) + elif element is ByteCode.operators: + operator_list_str = construct_operators(contents) + elif element is ByteCode.types: + type_list_str = construct_types(contents) + elif element is ByteCode.register_size: + register_size_str = construct_register_size(contents) + + one_upgrader_function_string = ONE_UPGRADER_FUNCTION.substitute( + upgrader_name=upgrader_name, + instruction_list=instruction_list_str, + constant_list=constant_list_str, + type_list=type_list_str, + register_size=register_size_str, + ) + one_upgrader_src_string = ONE_UPGRADER_SRC.substitute( + bytecode_function=one_upgrader_function_string.lstrip("\n"), + operator_string_list=operator_list_str.lstrip("\n"), + ) + all_upgrader_src_string.append(one_upgrader_src_string) + + upgrader_file_content = UPGRADER_CPP_SRC.substitute( + operator_version_map=version_map_src, + upgrader_bytecode="".join(all_upgrader_src_string).lstrip("\n"), + ) + print("writing file to : ", cpp_path + "/" + UPGRADER_MOBILE_FILE_NAME) + with open(os.path.join(cpp_path, UPGRADER_MOBILE_FILE_NAME), "wb") as out_file: + out_file.write(upgrader_file_content.encode("utf-8")) + + +def sort_upgrader(upgrader_list: list[dict[str, Any]]) -> list[dict[str, Any]]: + sorted_upgrader_list = sorted( + upgrader_list, key=lambda one_upgrader: next(iter(one_upgrader)) + ) + return sorted_upgrader_list + + +def main() -> None: + upgrader_list = generate_upgraders_bytecode() + sorted_upgrader_list = sort_upgrader(upgrader_list) + for up in sorted_upgrader_list: + print("after sort upgrader : ", next(iter(up))) + + pytorch_dir = Path(__file__).resolve().parents[2] + upgrader_path = pytorch_dir / "torch" / "csrc" / "jit" / "mobile" + write_cpp(str(upgrader_path), sorted_upgrader_list) + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/operator_versions/gen_mobile_upgraders_constant.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/operator_versions/gen_mobile_upgraders_constant.py new file mode 100644 index 0000000000000000000000000000000000000000..04b5ad887e54153115eeca7b6686d7c2de8dfc06 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/operator_versions/gen_mobile_upgraders_constant.py @@ -0,0 +1,7 @@ +MOBILE_UPGRADERS_HEADER_DESCRIPTION = """/** + * @generated + * This is an auto-generated file. Please do not modify it by hand. + * To re-generate, please run: + * cd ~/pytorch && python torchgen/operator_versions/gen_mobile_upgraders.py + */ +""" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/native/native_functions.yaml b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/native/native_functions.yaml new file mode 100644 index 0000000000000000000000000000000000000000..9e39eda0368fe01b747f61a33454881d72ec3430 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/native/native_functions.yaml @@ -0,0 +1,16269 @@ +# See README.md in this directory for more guidance + +# *********NB: _cast_* operators are DEPRECATED and will be removed +# eventually. These were previously used before TorchScript IR supported +# representing ScalarType's. They are now superseded by usage of +# `aten::to()`. The ops remain here for backward compatibility purposes. + +# DEPRECATED. DO NOT USE +- func: _cast_Byte(Tensor self, bool non_blocking=False) -> Tensor + variants: function + +# DEPRECATED. DO NOT USE +- func: _cast_Char(Tensor self, bool non_blocking=False) -> Tensor + variants: function + +# DEPRECATED. DO NOT USE +- func: _cast_Double(Tensor self, bool non_blocking=False) -> Tensor + variants: function + +# DEPRECATED. DO NOT USE +- func: _cast_Float(Tensor self, bool non_blocking=False) -> Tensor + variants: function + +# DEPRECATED. DO NOT USE +- func: _cast_Int(Tensor self, bool non_blocking=False) -> Tensor + variants: function + +# DEPRECATED. DO NOT USE +- func: _cast_Long(Tensor self, bool non_blocking=False) -> Tensor + variants: function + +# DEPRECATED. DO NOT USE +- func: _cast_Short(Tensor self, bool non_blocking=False) -> Tensor + variants: function + +# DEPRECATED. DO NOT USE +- func: _cast_Half(Tensor self, bool non_blocking=False) -> Tensor + variants: function + +# Computes the gradient of current tensor w.r.t. graph leaves. +- func: _backward(Tensor self, Tensor[] inputs, Tensor? gradient=None, bool? retain_graph=None, bool create_graph=False) -> () + manual_cpp_binding: True + variants: method + +# DEPRECATED. Sets the tensor data held by this `Variable` to be the same as +# `new_data`. It requires that `new_data` and `Variable` have compatible tensor +# type, by checking `_has_compatible_shallow_copy_type(this, new_data)`. +# +# This function is deprecated because it doesn't really make sense in a world +# where Variables *are* Tensors (as opposed to them containing tensors, which +# is what the previous interpretation was.) +- func: set_data(Tensor(a!) self, Tensor new_data) -> () + manual_cpp_binding: True + variants: method + +- func: data(Tensor self) -> Tensor + manual_cpp_binding: True + variants: method + +# True if this `Variable` is a leaf and thus does not have a `grad_fn`. +- func: is_leaf(Tensor self) -> bool + manual_cpp_binding: True + variants: method + +# Returns the output index of this variable from the forward operation that +# produced it. Conversely, it returns the input index of the gradient `Node` to +# which this `Variable` is connected (because in the gradient computation, +# inputs and outputs switch meaning). For example: +# +# y0, y1, y2 = f(x) +# assert y0.output_nr == 0 +# assert y1.output_nr == 1 +# assert y2.output_nr == 2 +# +- func: output_nr(Tensor self) -> int + manual_cpp_binding: True + variants: method + +- func: _version(Tensor self) -> int + manual_cpp_binding: True + variants: method + +- func: requires_grad_(Tensor(a!) self, bool requires_grad=True) -> Tensor(a!) + manual_cpp_binding: True + variants: method + +# Enables .grad attribute for non-leaf Tensors. +- func: retain_grad(Tensor(a!) self) -> () + manual_cpp_binding: True + variants: method + +- func: retains_grad(Tensor self) -> bool + manual_cpp_binding: True + variants: method + +- func: _fw_primal(Tensor(a) self, int level) -> Tensor(a) + variants: method + dispatch: + CompositeExplicitAutograd: _fw_primal + +- func: _make_dual(Tensor(a) primal, Tensor tangent, int level) -> Tensor(a) + variants: function + dispatch: + CompositeExplicitAutograd: _make_dual + +- func: _unpack_dual(Tensor(a) dual, int level) -> (Tensor(a) primal, Tensor tangent) + variants: function + +# NOTE: [_new_zeros_with_same_feature_meta] +# This function creates a new tensor with the layout and TensorOptions +# of `other` but also takes into account the batch dimensions of `self` +# +# This function has a couple extra constraints because it is also used for `jvp` +# in functorch. +# - is used for forward AD because there is the restriction +# that the primal and tangent must have the same layout +# - We cannot assume that `self` and `other` have the same sizes or even dim +# because in the inplace over view case, `other` is the base tensor, and +# `self` is the forward grad with respect to the view, which can have an +# entirely different shape +# - takes the number of batch dims for `self` because we also handle +# some batching logic. We handle that here instead of a batching rule because +# we'd like to avoid calling as_strided in the batching rule (as to enable +# nested vmap in functorch). +# - needs to be CompositeExplicitAutograd for jvp support in functorch. +# functorch currently relies on TensorWrapper which does not have storage +# CompositeExplicitAutograd makes sure the TensorWrapper is unwrapped. +# - this function may eventually take on another int argument to store the +# the number of batch dims for other once we support that use case +- func: _new_zeros_with_same_feature_meta(Tensor self, Tensor other, *, int self_num_batch_dims=0) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: _new_zeros_with_same_feature_meta + autogen: _new_zeros_with_same_feature_meta.out + +# This function compares the storage numel of self with that of other, where +# storage numel is computed as: `other.storage().nbytes() / other.itemsize()`. +# We create this function for composite compliance purposes. The batching rule +# always returns true because vmapped as_strided does not support accessing +# storage locations not indexable by the input tensor. +# See the note above for more information. +- func: _has_same_storage_numel(Tensor self, Tensor other) -> bool + variants: function + dispatch: + CompositeExplicitAutograd: _has_same_storage_numel + +- func: rename_(Tensor(a!) self, Dimname[]? names) -> Tensor(a!) + variants: method + tags: inplace_view + +- func: rename(Tensor(a) self, Dimname[]? names) -> Tensor(a) + variants: method + +- func: align_to(Tensor(a) self, Dimname[] names) -> Tensor(a) + variants: method + +- func: align_to.ellipsis_idx(Tensor(a) self, Dimname[] order, int ellipsis_idx) -> Tensor(a) + variants: method + +- func: align_as(Tensor self, Tensor other) -> Tensor + variants: method + +- func: align_tensors(Tensor[] tensors) -> Tensor[] + +# Not assert because it's a keyword; not Assert because FX already +# took that syntax +# TODO: need to specify this is side-effectful somehow +- func: _assert_async(Tensor self) -> () + dispatch: + CPU: _assert_async_cpu + CUDA: _assert_async_cuda + +- func: _assert_async.msg(Tensor self, str assert_msg) -> () + dispatch: + CPU: _assert_async_msg_cpu + CUDA: _assert_async_msg_cuda + +- func: _assert_scalar(Scalar self, str assert_msg) -> () + dispatch: + CompositeExplicitAutograd: _assert_scalar + +- func: _functional_assert_scalar(Scalar self, str assert_msg, Tensor dep_token) -> Tensor + dispatch: + CompositeExplicitAutograd: _functional_assert_scalar + +- func: _functional_assert_async.msg(Tensor self, str assert_msg, Tensor dep_token) -> Tensor + dispatch: + CPU: _functional_assert_async_msg_cpu + +- func: _assert_tensor_metadata(Tensor a, SymInt[]? size=None, SymInt[]? stride=None, ScalarType? dtype=None, *, Device? device=None, Layout? layout=None) -> () + dispatch: + CompositeExplicitAutograd: _assert_tensor_metadata + Meta: _assert_tensor_metadata_meta_symint + +- func: _print(str s) -> () + dispatch: + CompositeExplicitAutograd: _print + +- func: sym_constrain_range(Scalar size, *, int? min=None, int? max=None) -> () + dispatch: + CompositeExplicitAutograd: sym_constrain_range + +- func: sym_constrain_range_for_size(Scalar size, *, int? min=None, int? max=None) -> () + dispatch: + CompositeExplicitAutograd: sym_constrain_range_for_size + +- func: _functional_sym_constrain_range(Scalar size, int? min, int? max, Tensor dep_token) -> Tensor + dispatch: + CompositeExplicitAutograd: _functional_sym_constrain_range + +- func: _functional_sym_constrain_range_for_size(Scalar size, int? min, int? max, Tensor dep_token) -> Tensor + dispatch: + CompositeExplicitAutograd: _functional_sym_constrain_range_for_size + +- func: _make_dep_token(*, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + dispatch: + CPU: _make_dep_token_cpu + +- func: refine_names(Tensor(a) self, Dimname[] names) -> Tensor(a) + variants: method + +- func: _use_cudnn_ctc_loss(Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, int blank) -> bool + device_check: NoCheck # Tensor arguments allowed to be on different devices, see also _cudnn_ctc_loss + dispatch: + CUDA: _use_cudnn_ctc_loss + +- func: _use_cudnn_ctc_loss.Tensor(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank) -> bool + device_check: NoCheck # Tensor arguments allowed to be on different devices, see also _cudnn_ctc_loss + dispatch: + CUDA: _use_cudnn_ctc_loss_tensor + +- func: _cudnn_ctc_loss(Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, int blank, bool deterministic, bool zero_infinity) -> (Tensor, Tensor) + device_check: NoCheck # log_probs is expected to be on CUDA while targets is expected to be on CPU + dispatch: + CUDA: _cudnn_ctc_loss + autogen: _cudnn_ctc_loss.out + +- func: _cudnn_ctc_loss.Tensor(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank, bool deterministic, bool zero_infinity) -> (Tensor, Tensor) + device_check: NoCheck # log_probs is expected to be on CUDA while targets is expected to be on CPU + dispatch: + CUDA: _cudnn_ctc_loss_tensor + +- func: _use_cudnn_rnn_flatten_weight() -> bool + +- func: _cudnn_rnn_flatten_weight(Tensor[] weight_arr, int weight_stride0, SymInt input_size, int mode, SymInt hidden_size, SymInt proj_size, int num_layers, bool batch_first, bool bidirectional) -> Tensor + dispatch: + CUDA: _cudnn_rnn_flatten_weight + autogen: _cudnn_rnn_flatten_weight.out + +- func: _cudnn_rnn(Tensor input, Tensor[] weight, int weight_stride0, Tensor? weight_buf, Tensor hx, Tensor? cx, int mode, SymInt hidden_size, SymInt proj_size, int num_layers, bool batch_first, float dropout, bool train, bool bidirectional, SymInt[] batch_sizes, Tensor? dropout_state) -> (Tensor, Tensor, Tensor, Tensor, Tensor) + # rnn_tanh may or may not redispatch to _cudnn_rnn based on algorithm and build. Thus it might hit dispatch or kernel device check. + # Disable dispatch time device check for consistent behavior. + device_check: NoCheck + dispatch: + CUDA: _cudnn_rnn + autogen: _cudnn_rnn.out + tags: nondeterministic_seeded + +- func: _cudnn_rnn_backward(Tensor input, Tensor[] weight, int weight_stride0, Tensor weight_buf, Tensor hx, Tensor? cx, Tensor output, Tensor? grad_output, Tensor? grad_hy, Tensor? grad_cy, int mode, SymInt hidden_size, SymInt proj_size, int num_layers, bool batch_first, float dropout, bool train, bool bidirectional, SymInt[] batch_sizes, Tensor? dropout_state, Tensor reserve, bool[4] output_mask) -> (Tensor, Tensor, Tensor, Tensor[]) + dispatch: + CUDA: _cudnn_rnn_backward + autogen: _cudnn_rnn_backward.out + +- func: _cudnn_init_dropout_state(float dropout, bool train, int dropout_seed, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor + dispatch: + CUDA: _cudnn_init_dropout_state + autogen: _cudnn_init_dropout_state.out + tags: nondeterministic_seeded + +- func: _debug_has_internal_overlap(Tensor self) -> int + variants: function + +- func: _fused_dropout(Tensor self, float p, Generator? generator=None) -> (Tensor, Tensor) + variants: function + dispatch: + CUDA: fused_dropout_cuda + tags: nondeterministic_seeded + autogen: _fused_dropout.out + +- func: _masked_scale(Tensor self, Tensor mask, float scale) -> Tensor + variants: function + dispatch: + CUDA: masked_scale_cuda + autogen: _masked_scale.out + +- func: native_dropout(Tensor input, float p, bool? train) -> (Tensor, Tensor) + variants: function + dispatch: + CPU: native_dropout_cpu + CUDA: native_dropout_cuda + MPS: native_dropout_mps + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: native_dropout_nested + tags: [nondeterministic_seeded, core] + autogen: native_dropout.out + +- func: native_dropout_backward(Tensor grad_output, Tensor mask, float scale) -> Tensor + dispatch: + CPU, NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: native_dropout_backward + CUDA: native_dropout_backward_cuda + MPS: native_dropout_backward_mps + autogen: native_dropout_backward.out + tags: pointwise + +- func: _sobol_engine_draw(Tensor quasi, int n, Tensor sobolstate, int dimension, int num_generated, ScalarType? dtype) -> (Tensor, Tensor) + +- func: _sobol_engine_ff_(Tensor(a!) self, int n, Tensor sobolstate, int dimension, int num_generated) -> Tensor(a!) + +- func: _sobol_engine_scramble_(Tensor(a!) self, Tensor ltm, int dimension) -> Tensor(a!) + +- func: _sobol_engine_initialize_state_(Tensor(a!) self, int dimension) -> Tensor(a!) + +- func: _reshape_from_tensor(Tensor self, Tensor shape) -> Tensor + +- func: _shape_as_tensor(Tensor self) -> Tensor + +- func: dropout(Tensor input, float p, bool train) -> Tensor + tags: [nondeterministic_seeded, maybe_aliasing_or_mutating] + +- func: dropout_(Tensor(a!) self, float p, bool train) -> Tensor(a!) + tags: nondeterministic_seeded + +- func: feature_dropout(Tensor input, float p, bool train) -> Tensor + tags: [nondeterministic_seeded, maybe_aliasing_or_mutating] + +- func: feature_dropout_(Tensor(a!) self, float p, bool train) -> Tensor(a!) + tags: nondeterministic_seeded + +- func: alpha_dropout(Tensor input, float p, bool train) -> Tensor + tags: [nondeterministic_seeded, maybe_aliasing_or_mutating] + +- func: alpha_dropout_(Tensor(a!) self, float p, bool train) -> Tensor(a!) + tags: nondeterministic_seeded + +- func: feature_alpha_dropout(Tensor input, float p, bool train) -> Tensor + tags: [nondeterministic_seeded, maybe_aliasing_or_mutating] + +- func: feature_alpha_dropout_(Tensor(a!) self, float p, bool train) -> Tensor(a!) + tags: nondeterministic_seeded + +- func: abs(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: abs + SparseCPU, SparseCUDA, SparseMPS: abs_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: abs_sparse_csr + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_abs + tags: [core, pointwise] + +- func: abs_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: abs_ + SparseCPU, SparseCUDA, SparseMPS: abs_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: abs_sparse_csr_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_abs_ + +- func: abs.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS, MTIA: abs_out + SparseCPU, SparseCUDA, SparseMPS: abs_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: abs_sparse_csr_out + tags: pointwise + +# Note [Adding an alias] +# To add an alias do the following: +# +# 1) Copy the original functions native_functions.yaml entry, but replace the +# original function's name with their own and delete any dispatch +# keys for the aliases. Specifying a dispatch key will prevent +# autograd from recording the operations the alias performs, which +# will stop it from "inheriting" the original operation's autograd behavior. +# 2) Implement the corresponding functions and have them redispatch to the +# original function. +# 3) Add docstrings to the new function that reference the original function, +# and document the method as usual (if it exists.) +# (See torch/_torch_docs.py and docs/source/torch.rst if adding a function, +# torch/_tensor_docs.py and docs/source/tensors.rst if adding a method, +# or module-specific doc bindings (like torch/linalg/__init__.py) if +# adding an alias in a namespace.) +# 4) Update torch/overrides.py consistent with the original function. +# 5) Update the alias_map in torch/csrc/jit/passes/normalize_ops.cpp. +# 6) Add aliases argument to existing OpInfo/UnaryUfuncInfo or create new OpInfo/UnaryUfuncInfo entry +# in op_db list in torch/testing/_internal/common_methods_invocations.py +# +# See torch.absolute, an alias for torch.abs, as an example. +# Absolute, alias for abs + +- func: absolute(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + +- func: absolute_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + +- func: absolute.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + +- func: angle(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA, MPS: angle + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: angle_sparse_csr + tags: pointwise + +- func: angle.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS: angle_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: angle_sparse_csr_out + tags: pointwise + +- func: view_as_real(Tensor(a) self) -> Tensor(a) + variants: function + dispatch: + CPU, CUDA, MPS, Meta: view_as_real + SparseCPU, SparseCUDA, SparseMPS: view_as_real_sparse + +- func: view_as_complex(Tensor(a) self) -> Tensor(a) + variants: function + dispatch: + CPU, CUDA, MPS, Meta: view_as_complex + SparseCPU, SparseCUDA, SparseMPS: view_as_complex_sparse + +- func: sgn(Tensor self) -> Tensor + variants: function, method + structured_delegate: sgn.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sgn_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sgn_sparse_csr + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_sgn + tags: pointwise + +- func: sgn_(Tensor(a!) self) -> Tensor(a!) + variants: method + structured_delegate: sgn.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sgn_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sgn_sparse_csr_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_sgn_ + tags: pointwise + +- func: sgn.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: sgn_out + MPS: sgn_out_mps + SparseCPU, SparseCUDA, SparseMPS: sgn_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sgn_sparse_csr_out + tags: pointwise + +- func: chalf(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor + variants: method + +- func: real(Tensor(a) self) -> Tensor(a) + device_check: NoCheck # TensorIterator + variants: function + +- func: imag(Tensor(a) self) -> Tensor(a) + device_check: NoCheck # TensorIterator + variants: function + +- func: _conj(Tensor(a) self) -> Tensor(a) + variants: function, method + dispatch: + CompositeExplicitAutograd: _conj + +- func: conj(Tensor(a) self) -> Tensor(a) + variants: function, method + manual_cpp_binding: True + +- func: _conj_physical(Tensor self) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: _conj_physical + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: conj_physical_sparse_csr + autogen: _conj_physical.out + tags: pointwise + +- func: conj_physical(Tensor self) -> Tensor + variants: function, method + tags: [pointwise, maybe_aliasing_or_mutating] + +- func: conj_physical.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: conj_physical_out + MPS: conj_physical_out_mps + SparseCPU, SparseCUDA, SparseMPS: conj_physical_out_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: conj_physical_sparse_csr_out + tags: pointwise + +- func: conj_physical_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + dispatch: + CompositeExplicitAutograd: conj_physical_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: conj_physical_sparse_csr_ + tags: pointwise + +- func: resolve_conj(Tensor(a) self) -> Tensor(a) + variants: function, method + +- func: resolve_neg(Tensor(a) self) -> Tensor(a) + variants: function, method + +- func: _neg_view(Tensor(a) self) -> Tensor(a) + variants: function, method + dispatch: + CompositeExplicitAutograd: _neg_view + +- func: acos(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: acos.out + tags: [core, pointwise] + +- func: acos_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: acos.out + tags: pointwise + +- func: acos.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: acos_out + tags: pointwise + +# arccos, alias of acos +- func: arccos(Tensor self) -> Tensor + variants: function, method + +- func: arccos_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + +- func: arccos.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: avg_pool1d(Tensor self, int[1] kernel_size, int[1] stride=[], int[1] padding=0, bool ceil_mode=False, bool count_include_pad=True) -> Tensor + tags: core + autogen: avg_pool1d.out + +- func: adaptive_avg_pool1d(Tensor self, int[1] output_size) -> Tensor + tags: core + autogen: adaptive_avg_pool1d.out + +# Return: (Tensor output, Tensor indices) +- func: adaptive_max_pool1d(Tensor self, int[1] output_size) -> (Tensor, Tensor) + +- func: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: add.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: add_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: add_sparse_csr + MkldnnCPU: mkldnn_add + ZeroTensor: add_zerotensor + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_add_Tensor + tags: [core, pointwise] + +- func: add_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: add.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: add_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: add_sparse_csr_ + MkldnnCPU: mkldnn_add_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_add__Tensor + tags: pointwise + +- func: add.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + ufunc_inner_loop: + Generic: add (AllAndComplex, BFloat16, Half, ComplexHalf) + ScalarOnly: add (Bool) + dispatch: + SparseCPU, SparseMeta: add_out_sparse_cpu + SparseCUDA: add_out_sparse_cuda + SparseMPS: add_out_sparse_mps + SparseCsrCPU, SparseCsrMeta: add_out_sparse_compressed_cpu + SparseCsrCUDA: add_out_sparse_compressed_cuda + MkldnnCPU: mkldnn_add_out + MPS: add_out_mps + MTIA: add_out_mtia + tags: pointwise + +- func: _add_relu.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor + variants: function + dispatch: + CPU: add_relu + +- func: _add_relu_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!) + variants: function + dispatch: + CPU: add_relu_ + +- func: _add_relu.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CPU: add_relu_out + +- func: _add_relu.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor + variants: function + dispatch: + CPU: add_relu + +- func: _add_relu_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!) + variants: function + dispatch: + CPU: add_relu_ + autogen: _add_relu.Scalar_out + +# For C++ only, until we have conversion from C++ numbers to Tensor +- func: add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: add + tags: [core, pointwise] + +- func: add_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: add_ + autogen: add.Scalar_out + tags: pointwise + +- func: addmv(Tensor self, Tensor mat, Tensor vec, *, Scalar beta=1, Scalar alpha=1) -> Tensor + structured_delegate: addmv.out + variants: function, method + +- func: addmv_(Tensor(a!) self, Tensor mat, Tensor vec, *, Scalar beta=1, Scalar alpha=1) -> Tensor(a!) + structured_delegate: addmv.out + variants: function, method + +- func: addmv.out(Tensor self, Tensor mat, Tensor vec, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: addmv_out_cpu + CUDA: addmv_out_cuda + MPS: addmv_out_mps + XPU: addmv_out_xpu + SparseCsrCPU: addmv_out_sparse_compressed + SparseCsrCUDA: addmv_out_sparse_compressed_cuda + +- func: addr(Tensor self, Tensor vec1, Tensor vec2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + variants: function, method + dispatch: + CPU, CUDA: addr + MPS: addr_mps + CompositeExplicitAutograd: math_addr + +- func: addr_(Tensor(a!) self, Tensor vec1, Tensor vec2, *, Scalar beta=1, Scalar alpha=1) -> Tensor(a!) + variants: method + dispatch: + CompositeExplicitAutograd: addr_ + +- func: addr.out(Tensor self, Tensor vec1, Tensor vec2, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: addr_out + MPS: addr_out_mps + CompositeExplicitAutograd: math_addr_out + +- func: affine_grid_generator(Tensor theta, SymInt[] size, bool align_corners) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: affine_grid_generator + autogen: affine_grid_generator.out + +- func: affine_grid_generator_backward(Tensor grad, SymInt[] size, bool align_corners) -> Tensor + variants: function + +- func: _is_all_true(Tensor self) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: _is_all_true + +- func: _is_any_true(Tensor self) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: _is_any_true + +# Note: this function is only for testing. +- func: _test_check_tensor(Tensor self) -> Tensor + variants: function + +# Note; this function is only for testing +- func: _test_functorch_fallback(Tensor self, Tensor other) -> Tensor + variants: function + dispatch: + CPU: _test_functorch_fallback + autogen: _test_functorch_fallback.out + +- func: all.dim(Tensor self, int dim, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: all.out + variants: function, method + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_all + tags: reduction + + +- func: all.dims(Tensor self, int[]? dim=None, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: all.dims_out + variants: function, method + cpp_no_default_args: ['dim'] + dispatch: + CompositeExplicitAutograd: all_dims_default + tags: reduction + +- func: all.out(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + dispatch: + CPU, CUDA: all_out + MPS: all_out_mps + MTIA: all_out_mtia + tags: reduction + +- func: all.dims_out(Tensor self, int[]? dim=None, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + dispatch: + CPU, CUDA: all_dims_out + CompositeExplicitAutograd: all_dims_out_default + cpp_no_default_args: ['dim'] + tags: reduction + +- func: all.dimname(Tensor self, Dimname dim, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: all.dimname_out(Tensor self, Dimname dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: reduction + +- func: allclose(Tensor self, Tensor other, float rtol=1e-05, float atol=1e-08, bool equal_nan=False) -> bool + variants: function, method + tags: data_dependent_output + dispatch: + CompositeExplicitAutograd: allclose + +- func: any.dim(Tensor self, int dim, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: any.out + variants: function, method + tags: [core, reduction] + +- func: any.dims(Tensor self, int[]? dim=None, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: any.dims_out + variants: function, method + cpp_no_default_args: ['dim'] + tags: [core, reduction] + dispatch: + CompositeExplicitAutograd: any_dims_default + +- func: any.out(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + dispatch: + CPU, CUDA: any_out + MPS: any_out_mps + tags: reduction + +- func: any.dims_out(Tensor self, int[]? dim=None, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + dispatch: + CPU, CUDA: any_dims_out + CompositeExplicitAutograd: any_dims_out_default + cpp_no_default_args: ['dim'] + tags: reduction + +- func: any.dimname(Tensor self, Dimname dim, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: any.dimname_out(Tensor self, Dimname dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: reduction + +- func: arange(Scalar end, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: arange + +- func: arange.start(Scalar start, Scalar end, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: arange + +# This operator should be named `arange.start_out` if following the naming convention. However that +# name is already taken. Disabled because of CI job failures. +# FIXME: enable this +#- func: arange.start_out_(Scalar start, Scalar end, *, Tensor(a!) out) -> Tensor(a!) +# dispatch: +# CompositeExplicitAutograd: arange_start_out + +- func: arange.start_step(Scalar start, Scalar end, Scalar step=1, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: arange + cpp_no_default_args: ['step'] + tags: core + +- func: arange.out(Scalar end, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: arange_out + +- func: arange.start_out(Scalar start, Scalar end, Scalar step=1, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, Meta: arange_out + CUDA: arange_cuda_out + MPS: arange_mps_out + MTIA: arange_mtia_out + cpp_no_default_args: ['step'] + +# This function is a temporary hack to allow tracing of arange like constructs with dynamic +# bounds on arange. Normal arange is not traceable because it does not take any tensor inputs; +# if the range you need is based on another tensor, calling this function directly will +# preserve tracing. Get rid of this when arange can directly take tensors for bounds +# (so that it can be traced directly). +- func: _dim_arange(Tensor like, int dim) -> Tensor + +- func: argmax(Tensor self, int? dim=None, bool keepdim=False) -> Tensor + structured_delegate: argmax.out + device_check: NoCheck # TensorIterator + variants: function, method + tags: [core, reduction] + +- func: argmax.out(Tensor self, int? dim=None, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU, CUDA: argmax_out + MPS: argmax_out_mps + tags: reduction + +- func: argmin(Tensor self, int? dim=None, bool keepdim=False) -> Tensor + structured_delegate: argmin.out + device_check: NoCheck # TensorIterator + variants: function, method + tags: [core, reduction] + +- func: argmin.out(Tensor self, int? dim=None, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU, CUDA: argmin_out + MPS: argmin_out_mps + tags: reduction + +- func: acosh(Tensor self) -> Tensor + variants: function, method + structured_delegate: acosh.out + tags: [core, pointwise] + +- func: acosh_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + structured_delegate: acosh.out + tags: pointwise + +- func: acosh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: acosh_out + MPS: acosh_out_mps + tags: pointwise +# arccosh, alias for acosh + +- func: arccosh(Tensor self) -> Tensor + variants: function, method + +- func: arccosh_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + +- func: arccosh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: asinh(Tensor self) -> Tensor + variants: function, method + structured_delegate: asinh.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: asinh_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: asinh_sparse_csr + tags: [core, pointwise] + +- func: asinh_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + structured_delegate: asinh.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: asinh_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: asinh_sparse_csr_ + tags: pointwise + +- func: asinh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: asinh_out + MPS: asinh_out_mps + SparseCPU, SparseCUDA, SparseMPS: asinh_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: asinh_sparse_csr_out + tags: pointwise + +# arcsinh, alias for asinh +- func: arcsinh(Tensor self) -> Tensor + variants: function, method + +- func: arcsinh_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + +- func: arcsinh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: atanh(Tensor self) -> Tensor + structured_delegate: atanh.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: atanh_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: atanh_sparse_csr + tags: [core, pointwise] + +- func: atanh_(Tensor(a!) self) -> Tensor(a!) + structured_delegate: atanh.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: atanh_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: atanh_sparse_csr_ + tags: pointwise + +- func: atanh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: atanh_out + MPS: atanh_out_mps + SparseCPU, SparseCUDA, SparseMPS: atanh_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: atanh_sparse_csr_out + tags: pointwise +# arctanh, alias for atanh + +- func: arctanh(Tensor self) -> Tensor + variants: function, method + +- func: arctanh_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + +- func: arctanh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: as_strided(Tensor(a) self, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor(a) + variants: function, method + dispatch: + ZeroTensor, CPU, CUDA, MTIA, MPS: as_strided_tensorimpl + Meta: as_strided_tensorimpl_meta_symint + QuantizedCPU, QuantizedCUDA: as_strided_qtensorimpl + device_check: NoCheck + device_guard: False + tags: core + +- func: as_strided_(Tensor(a!) self, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor(a!) + use_const_ref_for_mutable_tensors: True + variants: function, method + device_check: NoCheck + device_guard: False + tags: inplace_view + dispatch: + CompositeExplicitAutogradNonFunctional: as_strided__symint + +- func: asin(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: asin.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: asin_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: asin_sparse_csr + tags: [core, pointwise] + +- func: asin_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: asin.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: asin_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: asin_sparse_csr_ + tags: pointwise + +- func: asin.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: asin_out + SparseCPU, SparseCUDA, SparseMPS: asin_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: asin_sparse_csr_out + tags: pointwise + +# arcsin, alias of asin +- func: arcsin(Tensor self) -> Tensor + variants: function, method + +- func: arcsin_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + +- func: arcsin.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: atan(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: atan.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: atan_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: atan_sparse_csr + tags: [core, pointwise] + +- func: atan_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: atan.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: atan_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: atan_sparse_csr_ + tags: pointwise + +- func: atan.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: atan_out + SparseCPU, SparseCUDA, SparseMPS: atan_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: atan_sparse_csr_out + tags: pointwise + +# arctan, alias of atan +- func: arctan(Tensor self) -> Tensor + variants: function, method + +- func: arctan_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + +- func: arctan.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: atleast_1d(Tensor self) -> Tensor + variants: function + tags: maybe_aliasing_or_mutating + +- func: atleast_1d.Sequence(Tensor[] tensors) -> Tensor[] + +- func: atleast_2d(Tensor self) -> Tensor + variants: function + tags: maybe_aliasing_or_mutating + +- func: atleast_2d.Sequence(Tensor[] tensors) -> Tensor[] + variants: function + +- func: atleast_3d(Tensor self) -> Tensor + variants: function + tags: maybe_aliasing_or_mutating + +- func: atleast_3d.Sequence(Tensor[] tensors) -> Tensor[] + variants: function + +- func: baddbmm(Tensor self, Tensor batch1, Tensor batch2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + variants: function, method + structured_delegate: baddbmm.out + +- func: baddbmm_(Tensor(a!) self, Tensor batch1, Tensor batch2, *, Scalar beta=1, Scalar alpha=1) -> Tensor(a!) + variants: method + structured_delegate: baddbmm.out + +- func: baddbmm.out(Tensor self, Tensor batch1, Tensor batch2, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + dispatch: + CPU: baddbmm_out_cpu + CUDA: baddbmm_out_cuda + MPS: baddbmm_out_mps + XPU: baddbmm_out_xpu + MTIA: baddbmm_out_mtia + SparseCsrCUDA: baddbmm_out_sparse_csr_cuda + +- func: baddbmm.dtype(Tensor self, Tensor batch1, Tensor batch2, ScalarType out_dtype, *, Scalar beta=1, Scalar alpha=1) -> Tensor + variants: function + dispatch: + CUDA: _baddbmm_dtype_cuda + XPU: _baddbmm_dtype_xpu + +- func: baddbmm.dtype_out(Tensor self, Tensor batch1, Tensor batch2, ScalarType out_dtype, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CUDA: _baddbmm_out_dtype_cuda + XPU: _baddbmm_out_dtype_xpu + +- func: bartlett_window(int window_length, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: bartlett_window + autogen: bartlett_window.out + +- func: bartlett_window.periodic(int window_length, bool periodic, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: bartlett_window + autogen: bartlett_window.periodic_out + +- func: batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps, bool cudnn_enabled) -> Tensor + tags: maybe_aliasing_or_mutating + +- func: quantized_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor mean, Tensor var, float eps, float output_scale, int output_zero_point) -> Tensor + dispatch: + QuantizedCPU: quantized_batch_norm + autogen: quantized_batch_norm.out + +- func: _batch_norm_impl_index(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps, bool cudnn_enabled) -> (Tensor, Tensor, Tensor, Tensor, int) + tags: maybe_aliasing_or_mutating + +- func: _batch_norm_impl_index_backward(int impl_index, Tensor input, Tensor grad_output, Tensor? weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_var_transform, bool train, float eps, bool[3] output_mask, Tensor reservedSpace) -> (Tensor, Tensor, Tensor) + +# Sample bernoulli with values in `self` as probability. +- func: bernoulli(Tensor self, *, Generator? generator=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: bernoulli + tags: nondeterministic_seeded + +- func: bernoulli.out(Tensor self, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + tags: nondeterministic_seeded + dispatch: + CPU, CUDA: bernoulli_out + MPS: bernoulli_out_mps + +- func: bernoulli_.Tensor(Tensor(a!) self, Tensor p, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + tags: nondeterministic_seeded + dispatch: + CPU, CUDA: bernoulli_ + MPS: bernoulli_mps_ + autogen: bernoulli.Tensor, bernoulli.Tensor_out + +- func: bernoulli_.float(Tensor(a!) self, float p=0.5, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + tags: nondeterministic_seeded + dispatch: + CPU, CUDA: bernoulli_ + MPS: bernoulli_mps_ + autogen: bernoulli.float_out + +# Note [bernoulli.p schema] +# We should probably just fix the overload ambiguity by appending a _functional to the C++ API name (BC breaking) +# This out-of-place version isn't used explicitly, but needed by jit. +# There is no default valid on `p` here because it would introduce ambiguity +# with `bernoulli(Tensor self, *, Generator? generator=None)` declaration. +- func: bernoulli.p(Tensor self, float p, *, Generator? generator=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutogradNonFunctional: bernoulli + +- func: bilinear(Tensor input1, Tensor input2, Tensor weight, Tensor? bias=None) -> Tensor + +- func: binary_cross_entropy(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean) -> Tensor + device_check: NoCheck # TensorIterator + python_module: nn + variants: function + dispatch: + CPU: binary_cross_entropy_cpu + CUDA: binary_cross_entropy_cuda + MPS: binary_cross_entropy_mps + +- func: binary_cross_entropy.out(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: nn + variants: function + dispatch: + CPU: binary_cross_entropy_out_cpu + CUDA: binary_cross_entropy_out_cuda + MPS: binary_cross_entropy_out_mps + +- func: binary_cross_entropy_backward(Tensor grad_output, Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean) -> Tensor + python_module: nn + variants: function + dispatch: + CPU: binary_cross_entropy_backward_cpu + CUDA: binary_cross_entropy_backward_cuda + MPS: binary_cross_entropy_backward_mps + +- func: binary_cross_entropy_backward.grad_input(Tensor grad_output, Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + variants: function + dispatch: + CPU: binary_cross_entropy_backward_out_cpu + CUDA: binary_cross_entropy_backward_out_cuda + MPS: binary_cross_entropy_backward_out_mps + +- func: binary_cross_entropy_with_logits(Tensor self, Tensor target, Tensor? weight=None, Tensor? pos_weight=None, int reduction=Mean) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: binary_cross_entropy_with_logits + autogen: binary_cross_entropy_with_logits.out + +- func: bincount(Tensor self, Tensor? weights=None, SymInt minlength=0) -> Tensor + variants: function, method + dispatch: + CPU: _bincount_cpu + CUDA: _bincount_cuda + MPS: _bincount_mps + tags: dynamic_output_shape + autogen: bincount.out + +- func: bitwise_not(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: bitwise_not.out + variants: function, method + tags: [core, pointwise] + +- func: bitwise_not_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: bitwise_not.out + variants: method + tags: pointwise + +- func: bitwise_not.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: bitwise_not_out + tags: pointwise + +- func: copysign.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: copysign_out + tags: pointwise + +- func: copysign.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: copysign.out + tags: pointwise + +- func: copysign_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: copysign.out + +- func: copysign.Scalar(Tensor self, Scalar other) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: copysign + tags: pointwise + +- func: copysign_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + variants: method + dispatch: + CompositeExplicitAutograd: copysign_ + +- func: copysign.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: copysign_out + tags: pointwise + +- func: _lazy_clone(Tensor self) -> Tensor + # Like clone, but the copy takes place lazily, only if either the + # input or the output are written. + variants: function, method + dispatch: + CompositeExplicitAutograd: _lazy_clone + +- func: logical_not(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: logical_not + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_logical_not + tags: [core, pointwise] + +- func: logical_not_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: logical_not_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_logical_not_ + tags: pointwise + +- func: logical_not.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: logical_not_out + MPS: logical_not_out_mps + tags: pointwise + +- func: logical_xor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: logical_xor + tags: [core, pointwise] + +- func: logical_xor_(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: logical_xor_ + tags: pointwise + +- func: logical_xor.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: logical_xor_out + MPS: logical_xor_out_mps + tags: pointwise + +- func: logical_and(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: logical_and + tags: [core, pointwise] + +- func: logical_and_(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: logical_and_ + tags: pointwise + +- func: logical_and.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: logical_and_out + MPS: logical_and_out_mps + tags: pointwise + +- func: logical_or(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: logical_or + tags: [core, pointwise] + +- func: logical_or_(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: logical_or_ + tags: pointwise + +- func: logical_or.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: logical_or_out + MPS: logical_or_out_mps + tags: pointwise + +- func: blackman_window(int window_length, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: blackman_window + autogen: blackman_window.out + +- func: blackman_window.periodic(int window_length, bool periodic, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: blackman_window + autogen: blackman_window.periodic_out + +- func: bmm(Tensor self, Tensor mat2) -> Tensor + structured_delegate: bmm.out + variants: function, method + dispatch: + SparseCPU: bmm_sparse_cpu + SparseCUDA: bmm_sparse_cuda + SparseMPS: bmm_sparse_mps + NestedTensorCPU: bmm_nested + NestedTensorCUDA: bmm_nested_cuda + tags: core + +- func: bmm.out(Tensor self, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + dispatch: + CPU: bmm_out_cpu + CUDA: bmm_out_cuda + MPS: bmm_out_mps + XPU: bmm_out_xpu + MTIA: bmm_out_mtia + SparseCPU: bmm_out_sparse_cpu + SparseCUDA: bmm_out_sparse_cuda + SparseMPS: bmm_out_sparse_mps + SparseCsrCUDA: bmm_out_sparse_csr_cuda + +- func: bmm.dtype(Tensor self, Tensor mat2, ScalarType out_dtype) -> Tensor + variants: function + dispatch: + CUDA: _bmm_dtype_cuda + XPU: _bmm_dtype_xpu + +- func: bmm.dtype_out(Tensor self, Tensor mat2, ScalarType out_dtype, *, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CUDA: _bmm_out_dtype_cuda + XPU: _bmm_out_dtype_xpu + +- func: broadcast_tensors(Tensor[] tensors) -> Tensor[] + device_check: NoCheck + device_guard: False + +- func: broadcast_to(Tensor(a) self, SymInt[] size) -> Tensor(a) + variants: function, method + dispatch: + CompositeImplicitAutograd: broadcast_to_symint + +- func: _sparse_broadcast_to(Tensor(a) self, int[] size) -> Tensor(a) + variants: function + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sparse_broadcast_to + +- func: cat(Tensor[] tensors, int dim=0) -> Tensor + structured_delegate: cat.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: cat_sparse + QuantizedCPU: cat_quantized_cpu + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: cat_nested + tags: core + +- func: cat.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!) + structured: True + precomputed: + - dim -> int dim, int valid, bool all_contiguous, bool all_same_dtype, bool all_same_sizes_and_stride, MemoryFormat memory_format + dispatch: + CPU: cat_out_cpu + CUDA: cat_out_cuda + MPS: cat_out_mps + QuantizedCPU: cat_out_quantized_cpu + +- func: cat.names(Tensor[] tensors, Dimname dim) -> Tensor + +- func: cat.names_out(Tensor[] tensors, Dimname dim, *, Tensor(a!) out) -> Tensor(a!) + +# alias for torch.cat +- func: concat(Tensor[] tensors, int dim=0) -> Tensor + +- func: concat.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!) + +- func: concat.names(Tensor[] tensors, Dimname dim) -> Tensor + +- func: concat.names_out(Tensor[] tensors, Dimname dim, *, Tensor(a!) out) -> Tensor(a!) + +# alias for torch.cat +- func: concatenate(Tensor[] tensors, int dim=0) -> Tensor + +- func: concatenate.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!) + +- func: concatenate.names(Tensor[] tensors, Dimname dim) -> Tensor + +- func: concatenate.names_out(Tensor[] tensors, Dimname dim, *, Tensor(a!) out) -> Tensor(a!) + +- func: block_diag(Tensor[] tensors) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: block_diag + autogen: block_diag.out + +- func: ceil(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: ceil.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: ceil_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: ceil_sparse_csr + tags: [core, pointwise] + +- func: ceil_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: ceil.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: ceil_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: ceil_sparse_csr_ + tags: pointwise + +- func: ceil.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: ceil_out + SparseCPU, SparseCUDA, SparseMPS: ceil_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: ceil_sparse_csr_out + tags: pointwise + +# alias for torch.linalg.multi_dot +- func: chain_matmul(Tensor[] matrices) -> Tensor + variants: function + +# alias for torch.linalg.multi_dot +- func: chain_matmul.out(Tensor[] matrices, *, Tensor(a!) out) -> Tensor(a!) + +- func: unsafe_chunk(Tensor self, int chunks, int dim=0) -> Tensor[] + variants: function, method + device_check: NoCheck + device_guard: False + tags: maybe_aliasing_or_mutating + +- func: chunk(Tensor(a -> *) self, int chunks, int dim=0) -> Tensor(a)[] + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: chunk + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: chunk_nested_tensor + +- func: tensor_split.sections(Tensor(a -> *) self, SymInt sections, int dim=0) -> Tensor(a)[] + variants: function, method + dispatch: + CompositeImplicitAutograd: tensor_split_sections_symint + +- func: tensor_split.indices(Tensor(a -> *) self, SymInt[] indices, int dim=0) -> Tensor(a)[] + variants: function, method + dispatch: + CompositeImplicitAutograd: tensor_split_indices_symint + +- func: tensor_split.tensor_indices_or_sections(Tensor(a -> *) self, Tensor tensor_indices_or_sections, int dim=0) -> Tensor(a)[] + variants: function, method + +- func: clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + cpp_no_default_args: ['min'] + structured_delegate: clamp.out + dispatch: + QuantizedCPU: clamp_quantized_cpu + tags: [core, pointwise] + +- func: clamp.Tensor(Tensor self, Tensor? min=None, Tensor? max=None) -> Tensor + variants: function, method + structured_delegate: clamp.Tensor_out + tags: [core, pointwise] + +- func: clamp_(Tensor(a!) self, Scalar? min=None, Scalar? max=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + cpp_no_default_args: ['min'] + structured_delegate: clamp.out + tags: pointwise + +- func: clamp_.Tensor(Tensor(a!) self, Tensor? min=None, Tensor? max=None) -> Tensor(a!) + variants: function, method + structured_delegate: clamp.Tensor_out + tags: pointwise + +- func: clamp.out(Tensor self, Scalar? min=None, Scalar? max=None, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + cpp_no_default_args: ['min'] + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MTIA, MPS: clamp_out + tags: pointwise + +- func: clamp.Tensor_out(Tensor self, Tensor? min=None, Tensor? max=None, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: clamp_Tensor_out + tags: pointwise + +- func: clamp_max(Tensor self, Scalar max) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: clamp_max.out + tags: pointwise + +- func: clamp_max.Tensor(Tensor self, Tensor max) -> Tensor + variants: function, method + structured_delegate: clamp_max.Tensor_out + tags: pointwise + +- func: clamp_max_(Tensor(a!) self, Scalar max) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: clamp_max.out + tags: pointwise + +- func: clamp_max_.Tensor(Tensor(a!) self, Tensor max) -> Tensor(a!) + variants: function, method + structured_delegate: clamp_max.Tensor_out + tags: pointwise + +- func: clamp_max.out(Tensor self, Scalar max, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MTIA, MPS: clamp_max_out + tags: pointwise + +- func: clamp_max.Tensor_out(Tensor self, Tensor max, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: clamp_max_Tensor_out + tags: pointwise + +- func: clamp_min(Tensor self, Scalar min) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: clamp_min.out + tags: pointwise + +- func: clamp_min.Tensor(Tensor self, Tensor min) -> Tensor + variants: function, method + structured_delegate: clamp_min.Tensor_out + tags: pointwise + +- func: clamp_min_(Tensor(a!) self, Scalar min) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: clamp_min.out + tags: pointwise + +- func: clamp_min_.Tensor(Tensor(a!) self, Tensor min) -> Tensor(a!) + variants: function, method + structured_delegate: clamp_min.Tensor_out + tags: pointwise + +- func: clamp_min.out(Tensor self, Scalar min, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MTIA, MPS: clamp_min_out + tags: pointwise + +- func: clamp_min.Tensor_out(Tensor self, Tensor min, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: clamp_min_Tensor_out + tags: pointwise + +# clip is an alias for clamp +- func: clip(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor + cpp_no_default_args: ['min'] + variants: function, method + tags: pointwise + +- func: clip.Tensor(Tensor self, Tensor? min=None, Tensor? max=None) -> Tensor + variants: function, method + tags: pointwise + +- func: clip_(Tensor(a!) self, Scalar? min=None, Scalar? max=None) -> Tensor(a!) + cpp_no_default_args: ['min'] + variants: function, method + tags: pointwise + +- func: clip_.Tensor(Tensor(a!) self, Tensor? min=None, Tensor? max=None) -> Tensor(a!) + variants: function, method + tags: pointwise + +- func: clip.out(Tensor self, Scalar? min=None, Scalar? max=None, *, Tensor(a!) out) -> Tensor(a!) + cpp_no_default_args: ['min'] + tags: pointwise + +- func: clip.Tensor_out(Tensor self, Tensor? min=None, Tensor? max=None, *, Tensor(a!) out) -> Tensor(a!) + +- func: cudnn_is_acceptable(Tensor self) -> bool + device_check: NoCheck + device_guard: False + +- func: complex(Tensor real, Tensor imag) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: complex + +- func: complex.out(Tensor real, Tensor imag, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: complex_out + +- func: polar(Tensor abs, Tensor angle) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: polar + +- func: polar.out(Tensor abs, Tensor angle, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: polar_out + +- func: constant_pad_nd(Tensor self, SymInt[] pad, Scalar value=0) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: constant_pad_nd + MPS: constant_pad_nd_mps + autogen: constant_pad_nd.out + tags: core + +- func: contiguous(Tensor(a) self, *, MemoryFormat memory_format=contiguous_format) -> Tensor(a) + variants: method + manual_cpp_binding: True + +- func: convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups) -> Tensor + dispatch: + CompositeExplicitAutograd: convolution + autogen: convolution.out + tags: core + +- func: convolution_backward(Tensor grad_output, Tensor input, Tensor weight, SymInt[]? bias_sizes, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + dispatch: + CompositeExplicitAutograd, CUDA: convolution_backward + autogen: convolution_backward.out + tags: core + +- func: convolution_overrideable(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups) -> Tensor + dispatch: + CompositeExplicitAutograd: convolution_overrideable + autogen: convolution_overrideable.out + +- func: convolution_backward_overrideable(Tensor grad_output, Tensor input, Tensor weight, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups, bool[3] output_mask) -> (Tensor grad_input, Tensor grad_weight, Tensor grad_bias) + dispatch: + CompositeExplicitAutograd: convolution_backward_overrideable + autogen: convolution_backward_overrideable.out + +- func: _convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups, bool benchmark, bool deterministic, bool cudnn_enabled, bool allow_tf32) -> Tensor + dispatch: + CompositeExplicitAutograd: _convolution + autogen: _convolution.out + +- func: _convolution.deprecated(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, int[] output_padding, SymInt groups, bool benchmark, bool deterministic, bool cudnn_enabled) -> Tensor + +- func: _convolution_mode(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, str padding, SymInt[] dilation, SymInt groups) -> Tensor + dispatch: + CompositeImplicitAutograd: _convolution_mode_symint + +- func: _convolution_double_backward(Tensor? ggI, Tensor? ggW, Tensor? ggb, Tensor gO, Tensor weight, Tensor self, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + +- func: conv1d(Tensor input, Tensor weight, Tensor? bias=None, SymInt[1] stride=1, SymInt[1] padding=0, SymInt[1] dilation=1, SymInt groups=1) -> Tensor + dispatch: + CompositeImplicitAutograd: conv1d_symint + +- func: conv2d(Tensor input, Tensor weight, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0, SymInt[2] dilation=1, SymInt groups=1) -> Tensor + dispatch: + CompositeImplicitAutograd: conv2d_symint + +- func: conv3d(Tensor input, Tensor weight, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0, SymInt[3] dilation=1, SymInt groups=1) -> Tensor + dispatch: + CompositeImplicitAutograd: conv3d_symint + +- func: conv1d.padding(Tensor input, Tensor weight, Tensor? bias=None, SymInt[1] stride=1, str padding="valid", SymInt[1] dilation=1, SymInt groups=1) -> Tensor + cpp_no_default_args: ['bias', 'stride', 'padding'] + dispatch: + CompositeImplicitAutograd: conv1d_padding_symint + +- func: conv2d.padding(Tensor input, Tensor weight, Tensor? bias=None, SymInt[2] stride=1, str padding="valid", SymInt[2] dilation=1, SymInt groups=1) -> Tensor + cpp_no_default_args: ['bias', 'stride', 'padding'] + dispatch: + CompositeImplicitAutograd: conv2d_padding_symint + +- func: conv3d.padding(Tensor input, Tensor weight, Tensor? bias=None, SymInt[3] stride=1, str padding="valid", SymInt[3] dilation=1, SymInt groups=1) -> Tensor + cpp_no_default_args: ['bias', 'stride', 'padding'] + dispatch: + CompositeImplicitAutograd: conv3d_padding_symint + +- func: conv_tbc(Tensor self, Tensor weight, Tensor bias, int pad=0) -> Tensor + dispatch: + CompositeExplicitAutograd: conv_tbc + autogen: conv_tbc.out + +- func: conv_tbc_backward(Tensor self, Tensor input, Tensor weight, Tensor bias, int pad) -> (Tensor, Tensor, Tensor) + +# NB: we inherit the goofy argument order from PyTorch torch.nn.functional +- func: conv_transpose1d(Tensor input, Tensor weight, Tensor? bias=None, SymInt[1] stride=1, SymInt[1] padding=0, SymInt[1] output_padding=0, SymInt groups=1, SymInt[1] dilation=1) -> Tensor + dispatch: + CompositeImplicitAutograd: conv_transpose1d_symint + +- func: conv_transpose2d.input(Tensor input, Tensor weight, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0, SymInt[2] output_padding=0, SymInt groups=1, SymInt[2] dilation=1) -> Tensor + dispatch: + CompositeImplicitAutograd: conv_transpose2d_symint + +- func: conv_transpose3d.input(Tensor input, Tensor weight, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0, SymInt[3] output_padding=0, SymInt groups=1, SymInt[3] dilation=1) -> Tensor + dispatch: + CompositeImplicitAutograd: conv_transpose3d_symint + +- func: copy(Tensor self, Tensor src, bool non_blocking=False) -> Tensor + variants: function + dispatch: + Meta: copy_meta + CompositeExplicitAutogradNonFunctional: copy + tags: core + +- func: copy_(Tensor(a!) self, Tensor src, bool non_blocking=False) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + dispatch: + MkldnnCPU: copy_mkldnn_ + SparseCPU, SparseCUDA, SparseMPS: copy_sparse_wrapper_ + CompositeExplicitAutograd: copy_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: copy_sparse_compressed_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: copy_nested_ + autogen: copy.out + +- func: _copy_from(Tensor self, Tensor dst, bool non_blocking=False) -> Tensor + dispatch: + MPS: _copy_from_mps + autogen: _copy_from.out + +# We need this to be able to properly copy from a CPU to an XLA tensor with different sizes. +# See https://github.com/pytorch/xla/issues/2881 +- func: _copy_from_and_resize(Tensor self, Tensor dst) -> Tensor + dispatch: + MPS: _copy_from_and_resize_mps + autogen: _copy_from_and_resize.out + +- func: cos(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: cos.out + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_cos + tags: [core, pointwise] + +- func: cos_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: cos.out + tags: pointwise + +- func: cos.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: cos_out + tags: pointwise + +- func: cosh(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: cosh.out + tags: [core, pointwise] + +- func: cosh_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: cosh.out + tags: pointwise + +- func: cosh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: cosh_out + tags: pointwise + +- func: cosine_embedding_loss(Tensor input1, Tensor input2, Tensor target, float margin=0.0, int reduction=Mean) -> Tensor + +- func: count_nonzero.dim_IntList(Tensor self, int[] dim) -> Tensor + variants: function, method + dispatch: + CPU: count_nonzero_cpu + CUDA: count_nonzero_cuda + MPS: count_nonzero_mps + autogen: count_nonzero.dim_IntList_out + tags: reduction + +- func: count_nonzero(Tensor self, int? dim=None) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: count_nonzero + autogen: count_nonzero.out + tags: reduction + +- func: cov(Tensor self, *, int correction=1, Tensor? fweights=None, Tensor? aweights=None) -> Tensor + variants: function, method + +- func: corrcoef(Tensor self) -> Tensor + variants: function, method + +- func: cudnn_affine_grid_generator(Tensor theta, int N, int C, int H, int W) -> Tensor grid + dispatch: + CUDA: cudnn_affine_grid_generator_forward + autogen: cudnn_affine_grid_generator.out + +# TODO: Why do I have to call this grad?! +- func: cudnn_affine_grid_generator_backward(Tensor grad, int N, int C, int H, int W) -> Tensor grad_theta + dispatch: + CUDA: cudnn_affine_grid_generator_backward + autogen: cudnn_affine_grid_generator_backward.out + +- func: cudnn_batch_norm(Tensor input, Tensor weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float exponential_average_factor, float epsilon) -> (Tensor, Tensor, Tensor, Tensor) + dispatch: + CUDA: cudnn_batch_norm + +- func: cudnn_batch_norm.out(Tensor input, Tensor weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float exponential_average_factor, float epsilon, *, Tensor(a!) out0, Tensor(b!) out1, Tensor(c!) out2, Tensor(d!) out3) -> (Tensor(a!), Tensor(b!), Tensor(c!), Tensor(d!)) + dispatch: + CUDA: cudnn_batch_norm_out + +# NB: You can only use this if you used cudnn_batch_norm training=True +- func: cudnn_batch_norm_backward(Tensor input, Tensor grad_output, Tensor weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_var, float epsilon, Tensor reserveSpace) -> (Tensor, Tensor, Tensor) + dispatch: + CUDA: cudnn_batch_norm_backward + autogen: cudnn_batch_norm_backward.out + +- func: cudnn_convolution(Tensor self, Tensor weight, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic, bool allow_tf32) -> Tensor + dispatch: + CUDA: cudnn_convolution + +- func: cudnn_convolution.out(Tensor self, Tensor weight, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic, bool allow_tf32, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CUDA: cudnn_convolution_out + +- func: cudnn_convolution_transpose(Tensor self, Tensor weight, SymInt[] padding, SymInt[] output_padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic, bool allow_tf32) -> Tensor + dispatch: + CUDA: cudnn_convolution_transpose + autogen: cudnn_convolution_transpose.out + +- func: _mps_convolution_transpose(Tensor self, Tensor weight, SymInt[] padding, SymInt[] output_padding, SymInt[] stride, SymInt[] dilation, SymInt groups) -> Tensor + dispatch: + MPS: _mps_convolution_transpose + autogen: _mps_convolution_transpose.out + +- func: mps_convolution_transpose_backward(Tensor self, Tensor grad_output, Tensor weight, SymInt[] padding, SymInt[] output_padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool[2] output_mask) -> (Tensor, Tensor) + dispatch: + MPS: mps_convolution_transpose_backward + autogen: mps_convolution_transpose_backward.out + +- func: cudnn_convolution_relu(Tensor self, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, SymInt groups) -> Tensor + dispatch: + CUDA: cudnn_convolution_relu + autogen: cudnn_convolution_relu.out + +- func: cudnn_convolution_add_relu(Tensor self, Tensor weight, Tensor z, Scalar? alpha, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, SymInt groups) -> Tensor + dispatch: + CUDA: cudnn_convolution_add_relu + autogen: cudnn_convolution_add_relu.out + +# NB: input is special cased in a way I don't quite understand +- func: cudnn_grid_sampler(Tensor self, Tensor grid) -> Tensor output + dispatch: + CUDA: cudnn_grid_sampler_forward + autogen: cudnn_grid_sampler.out + +- func: cudnn_grid_sampler_backward(Tensor self, Tensor grid, Tensor grad_output) -> (Tensor grad_self, Tensor grad_grid) + dispatch: + CUDA: cudnn_grid_sampler_backward + autogen: cudnn_grid_sampler_backward.out + +- func: cummax(Tensor self, int dim) -> (Tensor values, Tensor indices) + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: cummax + +- func: cummax.out(Tensor self, int dim, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + device_check: NoCheck # TensorIterator + dispatch: + CompositeExplicitAutograd: cummax_out + +- func: cummax.dimname(Tensor self, Dimname dim) -> (Tensor values, Tensor indices) + device_check: NoCheck # TensorIterator + variants: function, method + +- func: cummax.dimname_out(Tensor self, Dimname dim, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + device_check: NoCheck # TensorIterator + +- func: _cummax_helper(Tensor self, Tensor(a!) values, Tensor(b!) indices, int dim) -> () + variants: function + dispatch: + CPU: cummax_helper_cpu + CUDA: cummax_helper_cuda + MPS: cummax_helper_mps + +- func: cummin(Tensor self, int dim) -> (Tensor values, Tensor indices) + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: cummin + +- func: cummin.out(Tensor self, int dim, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + device_check: NoCheck # TensorIterator + dispatch: + CompositeExplicitAutograd: cummin_out + +- func: cummin.dimname(Tensor self, Dimname dim) -> (Tensor values, Tensor indices) + device_check: NoCheck # TensorIterator + variants: function, method + +- func: cummin.dimname_out(Tensor self, Dimname dim, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + device_check: NoCheck # TensorIterator + +- func: _cummin_helper(Tensor self, Tensor(a!) values, Tensor(b!) indices, int dim) -> () + variants: function + dispatch: + CPU: cummin_helper_cpu + CUDA: cummin_helper_cuda + MPS: cummin_helper_mps + +- func: cummaxmin_backward(Tensor grad, Tensor input, Tensor indices, int dim) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + +- func: cumprod(Tensor self, int dim, *, ScalarType? dtype=None) -> Tensor + structured_delegate: cumprod.out + device_check: NoCheck # TensorIterator + variants: function, method + +- func: cumprod_(Tensor(a!) self, int dim, *, ScalarType? dtype=None) -> Tensor(a!) + structured_delegate: cumprod.out + variants: method + +- func: cumprod.out(Tensor self, int dim, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + structured: True + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: cumprod_out + MPS: cumprod_out_mps + +- func: cumprod.dimname(Tensor self, Dimname dim, *, ScalarType? dtype=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + +- func: cumprod_.dimname(Tensor(a!) self, Dimname dim, *, ScalarType? dtype=None) -> Tensor(a!) + variants: method + +- func: cumprod.dimname_out(Tensor self, Dimname dim, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + +- func: cumprod_backward(Tensor grad, Tensor input, int dim, Tensor output) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + +- func: cumsum(Tensor self, int dim, *, ScalarType? dtype=None) -> Tensor + structured_delegate: cumsum.out + device_check: NoCheck # TensorIterator + variants: function, method + tags: core + +- func: cumsum_(Tensor(a!) self, int dim, *, ScalarType? dtype=None) -> Tensor(a!) + structured_delegate: cumsum.out + variants: method + +- func: cumsum.out(Tensor self, int dim, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + structured: True + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: cumsum_out + MPS: cumsum_out_mps + +- func: cumsum.dimname(Tensor self, Dimname dim, *, ScalarType? dtype=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + +- func: cumsum_.dimname(Tensor(a!) self, Dimname dim, *, ScalarType? dtype=None) -> Tensor(a!) + variants: method + +- func: cumsum.dimname_out(Tensor self, Dimname dim, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + +- func: cumulative_trapezoid.x(Tensor y, Tensor x, *, int dim=-1) -> Tensor + +- func: cumulative_trapezoid.dx(Tensor y, *, Scalar dx=1, int dim=-1) -> Tensor + +- func: ctc_loss.IntList(Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, int blank=0, int reduction=Mean, bool zero_infinity=False) -> Tensor + +# convenience function that converts to intlists for you +- func: ctc_loss.Tensor(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank=0, int reduction=Mean, bool zero_infinity=False) -> Tensor + +- func: _ctc_loss(Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, int blank=0, bool zero_infinity=False) -> (Tensor, Tensor) + dispatch: + CPU: ctc_loss_cpu + CUDA: ctc_loss_gpu + Meta: ctc_loss_meta + autogen: _ctc_loss.out + tags: dynamic_output_shape # the shape of second output is data dependent + +- func: _ctc_loss.Tensor(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank=0, bool zero_infinity=False) -> (Tensor, Tensor) + dispatch: + CPU, CUDA: ctc_loss_tensor + autogen: _ctc_loss.Tensor_out + tags: dynamic_output_shape # the shape of second output is data dependent + +- func: _ctc_loss_backward(Tensor grad, Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, Tensor neg_log_likelihood, Tensor log_alpha, int blank, bool zero_infinity=False) -> Tensor + dispatch: + CPU: ctc_loss_backward_cpu + CUDA: ctc_loss_backward_gpu + autogen: _ctc_loss_backward.out + +- func: _ctc_loss_backward.Tensor(Tensor grad, Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, Tensor neg_log_likelihood, Tensor log_alpha, int blank, bool zero_infinity=False) -> Tensor + dispatch: + CPU, CUDA: ctc_loss_backward_tensor + +- func: diag_embed(Tensor self, int offset=0, int dim1=-2, int dim2=-1) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutogradNonFunctional: diag_embed + autogen: diag_embed.out + +- func: diagflat(Tensor self, int offset=0) -> Tensor + variants: function, method + +- func: diagonal(Tensor(a) self, int offset=0, int dim1=0, int dim2=1) -> Tensor(a) + variants: function, method + dispatch: + CompositeExplicitAutograd: diagonal + tags: core + +- func: linalg_diagonal(Tensor(a) A, *, int offset=0, int dim1=-2, int dim2=-1) -> Tensor(a) + python_module: linalg + variants: function + +- func: diagonal.Dimname(Tensor(a) self, *, Dimname outdim, Dimname dim1, Dimname dim2, int offset=0) -> Tensor(a) + variants: function, method + +- func: diagonal_backward(Tensor grad_output, SymInt[] input_sizes, int offset, int dim1, int dim2) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: diagonal_backward_symint + autogen: diagonal_backward.out + +- func: fill_diagonal_(Tensor(a!) self, Scalar fill_value, bool wrap=False) -> Tensor(a!) + variants: method + +- func: diff(Tensor self, int n=1, int dim=-1, Tensor? prepend=None, Tensor? append=None) -> Tensor + variants: function, method + +- func: diff.out(Tensor self, int n=1, int dim=-1, Tensor? prepend=None, Tensor? append=None, *, Tensor(a!) out) -> Tensor(a!) + variants: function + +- func: gradient.scalarint(Tensor self, *, Scalar? spacing=None, int? dim=None, int edge_order=1) -> Tensor[] + variants: function + +- func: gradient.scalararray(Tensor self, *, Scalar spacing, int[] dim, int edge_order=1) -> Tensor[] + variants: function + +- func: gradient.array(Tensor self, *, int[] dim, int edge_order=1) -> Tensor[] + variants: function + +- func: gradient.scalarrayint(Tensor self, *, Scalar[] spacing, int? dim=None, int edge_order=1) -> Tensor[] + variants: function + +- func: gradient.scalarrayarray(Tensor self, *, Scalar[] spacing, int[] dim, int edge_order=1) -> Tensor[] + variants: function + +- func: gradient.tensorarrayint(Tensor self, *, Tensor[] spacing, int? dim=None, int edge_order=1) -> Tensor[] + variants: function + +- func: gradient.tensorarray(Tensor self, *, Tensor[] spacing, int[] dim, int edge_order=1) -> Tensor[] + variants: function + +- func: div.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: div.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: div_sparse + ZeroTensor: div_zerotensor + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_div_Tensor + tags: [core, pointwise] + +- func: div_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: div.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: div_sparse_ + tags: pointwise + +- func: div.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: div_out + SparseCPU, SparseCUDA, SparseMPS: div_out_sparse_zerodim + tags: pointwise + +- func: div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: div.out_mode + dispatch: + SparseCPU, SparseCUDA, SparseMPS: div_sparse + tags: [core, pointwise] + +- func: div_.Tensor_mode(Tensor(a!) self, Tensor other, *, str? rounding_mode) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: div.out_mode + dispatch: + SparseCPU, SparseCUDA, SparseMPS: div_sparse_ + tags: pointwise + +- func: div.out_mode(Tensor self, Tensor other, *, str? rounding_mode, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: div_out_mode + SparseCPU, SparseCUDA, SparseMPS: div_out_sparse_zerodim + tags: pointwise + +# For C++ only, until we have conversion from C++ numbers to Tensor +- func: div.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: div + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_div_Scalar + tags: [core, pointwise] + +- func: div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: div_ + autogen: div.Scalar_out + tags: pointwise + +- func: div.Scalar_mode(Tensor self, Scalar other, *, str? rounding_mode) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: div + tags: [core, pointwise] + +- func: div_.Scalar_mode(Tensor(a!) self, Scalar other, *, str? rounding_mode) -> Tensor(a!) + variants: method + dispatch: + CompositeExplicitAutograd: div_ + autogen: div.Scalar_mode_out + tags: pointwise + +# divide, alias for div +- func: divide.Tensor(Tensor self, Tensor other) -> Tensor + variants: function, method + +- func: divide_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + variants: method + +- func: divide.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + +- func: divide.Scalar(Tensor self, Scalar other) -> Tensor + variants: function, method + +- func: divide_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + variants: method + +- func: divide.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> Tensor + variants: function, method + +- func: divide_.Tensor_mode(Tensor(a!) self, Tensor other, *, str? rounding_mode) -> Tensor(a!) + variants: method + +- func: divide.out_mode(Tensor self, Tensor other, *, str? rounding_mode, Tensor(a!) out) -> Tensor(a!) + +- func: divide.Scalar_mode(Tensor self, Scalar other, *, str? rounding_mode) -> Tensor + variants: function, method + +- func: divide_.Scalar_mode(Tensor(a!) self, Scalar other, *, str? rounding_mode) -> Tensor(a!) + variants: method + + # true_divide, an alias for div +- func: true_divide.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: pointwise + +- func: true_divide_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + +- func: true_divide.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + +- func: true_divide.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + +- func: true_divide_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + +- func: dot(Tensor self, Tensor tensor) -> Tensor + variants: function, method + dispatch: + CPU: dot + CUDA: dot_cuda + MPS: dot_mps + +- func: dot.out(Tensor self, Tensor tensor, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: dot_out + +- func: vdot(Tensor self, Tensor other) -> Tensor + variants: function, method + dispatch: + CPU: vdot + CUDA: vdot_cuda + MPS: vdot_mps + +- func: vdot.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: vdot_out + +- func: einsum(str equation, Tensor[] tensors, *, int[]? path=None) -> Tensor + +- func: embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor + dispatch: + CompositeExplicitAutograd: embedding_symint + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_embedding + autogen: embedding.out + tags: core + +- func: embedding_backward(Tensor grad, Tensor indices, SymInt num_weights, SymInt padding_idx, bool scale_grad_by_freq, bool sparse) -> Tensor + dispatch: + CompositeImplicitAutograd: embedding_backward_symint + +- func: embedding_dense_backward(Tensor grad_output, Tensor indices, SymInt num_weights, SymInt padding_idx, bool scale_grad_by_freq) -> Tensor + dispatch: + CPU: embedding_dense_backward_cpu + CUDA: embedding_dense_backward_cuda + MPS: embedding_dense_backward_mps + autogen: embedding_dense_backward.out + tags: core + +- func: embedding_renorm_(Tensor(a!) self, Tensor indices, float max_norm, float norm_type) -> Tensor(a!) + dispatch: + CPU: embedding_renorm_cpu_ + CUDA: embedding_renorm_cuda_ + autogen: embedding_renorm, embedding_renorm.out + +- func: embedding_sparse_backward(Tensor grad, Tensor indices, int num_weights, int padding_idx, bool scale_grad_by_freq) -> Tensor + +# NOTE [ embedding_bag Native Functions ] +# The `_embedding_bag.*` variants assume that input tensors except for `weight`, +# e.g. `indices` and `offsets` (and `offset2bag`), are contiguous. +# We really only need to enforce this for `_embedding_bag` (the forward) because +# the backward inputs are the same as forward ones. +# The above `embedding_bag` wrapper is created to achieve this, e.g., +# applying indices = indices.contiguous(). +# The backward functions apply a check that these input tensors are contiguous. + + +- func: _embedding_bag_forward_only(Tensor weight, Tensor indices, Tensor offsets, bool scale_grad_by_freq=False, int mode=0, bool sparse=False, Tensor? per_sample_weights=None, bool include_last_offset=False, int padding_idx=-1) -> (Tensor, Tensor, Tensor, Tensor) + dispatch: + CPU: _embedding_bag_forward_only_cpu + CUDA: _embedding_bag_forward_only_cuda + MPS: _embedding_bag_forward_only_mps + autogen: _embedding_bag_forward_only.out + +- func: _rowwise_prune(Tensor weight, Tensor mask, ScalarType compressed_indices_dtype) -> (Tensor, Tensor) + +# row_stack is the alias of vstack +- func: row_stack(Tensor[] tensors) -> Tensor + +- func: row_stack.out(Tensor[] tensors, *, Tensor(a!) out) -> Tensor(a!) + +- func: embedding_bag(Tensor weight, Tensor indices, Tensor offsets, bool scale_grad_by_freq=False, int mode=0, bool sparse=False, Tensor? per_sample_weights=None, bool include_last_offset=False) -> (Tensor, Tensor, Tensor, Tensor) + +# To keep backward and forward compatibility, and to avoid ambiguity with the +# original signature above, scale_grad_by_freq, mode, sparse, +# per_sample_weights, and include_last_offset parameters do not have default +# values. Once the original signature is removed, default values can be added. +- func: embedding_bag.padding_idx(Tensor weight, Tensor indices, Tensor offsets, bool scale_grad_by_freq, int mode, bool sparse, Tensor? per_sample_weights, bool include_last_offset, int? padding_idx) -> (Tensor, Tensor, Tensor, Tensor) + +- func: _embedding_bag(Tensor weight, Tensor indices, Tensor offsets, bool scale_grad_by_freq=False, int mode=0, bool sparse=False, Tensor? per_sample_weights=None, bool include_last_offset=False, int padding_idx=-1) -> (Tensor, Tensor, Tensor, Tensor) + dispatch: + CPU: _embedding_bag_cpu + CUDA: _embedding_bag_cuda + MPS: _embedding_bag_mps + autogen: _embedding_bag.out + tags: core + +- func: _embedding_bag_backward(Tensor grad, Tensor indices, Tensor offsets, Tensor offset2bag, Tensor bag_size, Tensor maximum_indices, SymInt num_weights, bool scale_grad_by_freq, int mode, bool sparse, Tensor? per_sample_weights, int padding_idx=-1) -> Tensor + dispatch: + CPU, CUDA, MPS: _embedding_bag_backward_symint + +- func: _embedding_bag_sparse_backward(Tensor grad, Tensor indices, Tensor offsets, Tensor offset2bag, Tensor bag_size, SymInt num_weights, bool scale_grad_by_freq, int mode, Tensor? per_sample_weights, int padding_idx=-1) -> Tensor + dispatch: + CompositeImplicitAutograd: _embedding_bag_sparse_backward_symint + +- func: _embedding_bag_dense_backward(Tensor grad, Tensor indices, Tensor offset2bag, Tensor bag_size, Tensor maximum_indices, SymInt num_weights, bool scale_grad_by_freq, int mode, Tensor? per_sample_weights, int padding_idx=-1) -> Tensor + dispatch: + CPU: _embedding_bag_dense_backward_cpu + CUDA: _embedding_bag_dense_backward_cuda + MPS: _embedding_bag_dense_backward_mps + autogen: _embedding_bag_dense_backward.out + +- func: _embedding_bag_per_sample_weights_backward(Tensor grad, Tensor weight, Tensor indices, Tensor offsets, Tensor offset2bag, int mode, int padding_idx=-1) -> Tensor + dispatch: + CPU: _embedding_bag_per_sample_weights_backward_cpu + CUDA: _embedding_bag_per_sample_weights_backward_cuda + MPS: _embedding_bag_per_sample_weights_backward_mps + autogen: _embedding_bag_per_sample_weights_backward.out + +- func: empty.names(int[] size, *, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: empty_names + autogen: empty.names_out + +- func: empty.memory_format(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + dispatch: + CPU: empty_cpu + CUDA: empty_cuda + MPS: empty_mps + Meta: empty_meta_symint + MkldnnCPU: empty_mkldnn + SparseCPU, SparseCUDA, SparseMPS: empty_sparse + SparseMeta: empty_sparse_symint + SparseCsrCPU, SparseCsrCUDA: empty_sparse_compressed + SparseCsrMeta: empty_sparse_compressed_symint + QuantizedCPU, QuantizedCUDA, QuantizedMeta: empty_unknown_quantized + tags: core + +- func: empty_permuted(SymInt[] size, int[] physical_layout, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: empty_permuted_symint + autogen: empty_permuted.out + +# We do not make new_empty a composite that calls into new_empty_strided, as the strided version +# is significantly more difficult to implement by different backends +- func: new_empty(Tensor self, SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + variants: method + dispatch: + CompositeExplicitAutograd: new_empty_symint + autogen: new_empty.out + +- func: new_empty_strided(Tensor self, SymInt[] size, SymInt[] stride, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + variants: method + dispatch: + CompositeExplicitAutogradNonFunctional: new_empty_strided_symint + autogen: new_empty_strided.out + +- func: new_full(Tensor self, SymInt[] size, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + variants: method + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: new_full + autogen: new_full.out + +- func: new_zeros(Tensor self, SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + variants: method + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: new_zeros + autogen: new_zeros.out + +- func: new_ones(Tensor self, SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + variants: method + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: new_ones + autogen: new_ones.out + +# other overrides are to provide a more helpful error message that dtype is required +- func: _empty_affine_quantized(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, float scale=1, int zero_point=0, MemoryFormat? memory_format=contiguous_format) -> Tensor + dispatch: + CPU: empty_affine_quantized_other_backends_stub + QuantizedCPU, QuantizedCUDA: empty_affine_quantized + autogen: _empty_affine_quantized.out + +# it's a factory function receiving a tensor argument, thus overriding explicitly +# other overrides are to provide a more helpful error message that dtype is required +- func: _empty_per_channel_affine_quantized(SymInt[] size, *, Tensor scales, Tensor zero_points, int axis, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=contiguous_format) -> Tensor + category_override: factory + dispatch: + CPU: empty_per_channel_affine_quantized_other_backends_stub + QuantizedCPU, QuantizedCUDA: empty_per_channel_affine_quantized + autogen: _empty_per_channel_affine_quantized.out + +- func: resize_(Tensor(a!) self, SymInt[] size, *, MemoryFormat? memory_format=None) -> Tensor(a!) + use_const_ref_for_mutable_tensors: True + variants: method + device_check: NoCheck + device_guard: False + tags: [core, inplace_view] + dispatch: + Meta: resize__symint + CPU: resize_ + CUDA: resize_cuda_ + MPS: resize_mps_ + QuantizedCPU: quantized_resize_cpu_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: resize_sparse_csr_ + autogen: resize, resize.out + +# This is a utility function to enable users to resize out tensor while registering kernels for out variants. +# Eventually, we can consider exposing `resize_output` as a public API to ship it with python op registration +# to make it easy to register out variants for ops. +- func: _resize_output_(Tensor(a!) self, SymInt[] size, Device device) -> Tensor(a!) + use_const_ref_for_mutable_tensors: True + variants: function + dispatch: + Meta: _resize_output_ + autogen: _resize_output, _resize_output.out + +- func: empty_quantized(int[] size, Tensor qtensor, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + category_override: factory + variants: function + dispatch: + QuantizedCPU, QuantizedCUDA: empty_quantized + autogen: empty_quantized.out + +- func: empty.out(SymInt[] size, *, MemoryFormat? memory_format=None, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + device_guard: False + +- func: empty_like(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: empty_like + QuantizedCPU, QuantizedCUDA: empty_like_quantized + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: empty_like_sparse_coo + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: empty_like_sparse_csr + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: empty_like_nested + autogen: empty_like.out + +- func: empty_strided(SymInt[] size, SymInt[] stride, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CPU: empty_strided_cpu + CUDA: empty_strided_cuda + MPS: empty_strided_mps + Meta: empty_strided_meta_symint + QuantizedCPU, QuantizedCUDA: empty_strided_unknown_quantized + autogen: empty_strided.out + tags: core + +- func: erf(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: erf.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: erf_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: erf_sparse_csr + tags: [core, pointwise] + +- func: erf_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: erf.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: erf_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: erf_sparse_csr_ + tags: pointwise + +- func: erf.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: erf_out + SparseCPU, SparseCUDA, SparseMPS: erf_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: erf_sparse_csr_out + tags: pointwise + +- func: erfc(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: erfc.out + variants: function, method + tags: pointwise + +- func: erfc_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: erfc.out + variants: function, method + tags: pointwise + +- func: erfc.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: erfc_out + tags: pointwise + +- func: exp(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: exp.out + variants: function, method + tags: [core, pointwise] + +- func: exp_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: exp.out + variants: function, method + tags: pointwise + +- func: exp.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: exp_out + tags: pointwise + +- func: exp2(Tensor self) -> Tensor + structured_delegate: exp2.out + variants: function, method + tags: pointwise + +- func: exp2_(Tensor(a!) self) -> Tensor(a!) + structured_delegate: exp2.out + variants: function, method + tags: pointwise + +- func: exp2.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: exp2_out + tags: pointwise + +- func: expm1(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: expm1.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: expm1_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: expm1_sparse_csr + tags: [core, pointwise] + +- func: expm1_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: expm1.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: expm1_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: expm1_sparse_csr_ + tags: pointwise + +- func: expm1.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: expm1_out + SparseCPU, SparseCUDA, SparseMPS: expm1_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: expm1_sparse_csr_out + tags: pointwise + +- func: expand(Tensor(a) self, SymInt[] size, *, bool implicit=False) -> Tensor(a) + variants: method # This is method-only to match the previous tensor API. In the future we could make this a function too. + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: expand + tags: core + +- func: expand_as(Tensor(a) self, Tensor other) -> Tensor(a) + variants: method # This is method-only to match the previous tensor API. In the future we could make this a function too. + device_check: NoCheck + device_guard: False + +# decomposes to eye.m +- func: eye(SymInt n, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: eye + +- func: eye.m(SymInt n, SymInt m, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: eye + +- func: eye.out(SymInt n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, Meta: eye_out_cpu + CUDA: eye_out_cuda + MPS: eye_out_mps + +- func: eye.m_out(SymInt n, SymInt m, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, Meta: eye_out_cpu + CUDA: eye_out_cuda + MPS: eye_out_mps + +- func: flatten.using_ints(Tensor(a) self, int start_dim=0, int end_dim=-1) -> Tensor(a) + variants: function, method + +- func: flatten.named_out_dim(Tensor(a) self, int start_dim, int end_dim, Dimname out_dim) -> Tensor(a) + variants: function, method + +- func: flatten.using_names(Tensor(a) self, Dimname start_dim, Dimname end_dim, Dimname out_dim) -> Tensor(a) + variants: function, method + +- func: flatten.DimnameList(Tensor(a) self, Dimname[] dims, Dimname out_dim) -> Tensor(a) + variants: function, method + +- func: unflatten.int(Tensor(a) self, int dim, SymInt[] sizes) -> Tensor(a) + variants: function, method + dispatch: + CompositeImplicitAutograd: unflatten_symint + +- func: unflatten.Dimname(Tensor(a) self, Dimname dim, SymInt[] sizes, Dimname[] names) -> Tensor(a) + variants: function, method + dispatch: + CompositeImplicitAutograd: unflatten_dimname_symint + +- func: fill.Scalar(Tensor self, Scalar value) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: fill + tags: core + +- func: fill.Tensor(Tensor self, Tensor value) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: fill + +- func: fill_.Scalar(Tensor(a!) self, Scalar value) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA, MPS: fill_ + QuantizedCPU, QuantizedCUDA: fill_quantized_ + Meta: fill_meta_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: fill_sparse_csr_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: fill_nested_ + autogen: fill.Scalar_out + +- func: fill_.Tensor(Tensor(a!) self, Tensor value) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA, MPS: fill_ + QuantizedCPU, QuantizedCUDA: fill_quantized_ + Meta: fill_meta_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: fill_nested_ + autogen: fill.Tensor_out + +- func: floor(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: floor.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: floor_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: floor_sparse_csr + tags: [core, pointwise] + +- func: floor_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: floor.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: floor_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: floor_sparse_csr_ + tags: pointwise + +- func: floor.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: floor_out + SparseCPU, SparseCUDA, SparseMPS: floor_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: floor_sparse_csr_out + tags: pointwise + +- func: floor_divide(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA, MPS, MTIA: floor_divide + SparseCPU, SparseCUDA, SparseMPS: floor_divide_sparse + +- func: floor_divide_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CPU, CUDA, MPS: floor_divide_ + SparseCPU, SparseCUDA, SparseMPS: floor_divide_sparse_ + +- func: floor_divide.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS, MTIA: floor_divide_out + SparseCPU, SparseCUDA, SparseMPS: floor_divide_out_sparse_zerodim + +- func: floor_divide.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: floor_divide + +- func: floor_divide_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: floor_divide_ + autogen: floor_divide.Scalar_out + +- func: frac(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: frac.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: frac_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: frac_sparse_csr + tags: pointwise + +- func: frac_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: frac.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: frac_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: frac_sparse_csr_ + tags: pointwise + +- func: frac.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: frac_out + MPS: frac_out_mps + SparseCPU, SparseCUDA, SparseMPS: frac_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: frac_sparse_csr_out + tags: pointwise + +- func: full.names(int[] size, Scalar fill_value, *, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: full + autogen: full.names_out + +- func: full(SymInt[] size, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: full + tags: core + +- func: full.out(SymInt[] size, Scalar fill_value, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: full_out + +- func: full_like(Tensor self, Scalar fill_value, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: full_like + autogen: full_like.out + tags: core + +- func: from_file(str filename, bool? shared=None, int? size=0, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CPU: from_file + autogen: from_file.out + +- func: gcd.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: gcd_out + tags: pointwise + +- func: gcd(Tensor self, Tensor other) -> Tensor + structured_delegate: gcd.out + variants: function, method + tags: pointwise + +- func: gcd_(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: gcd.out + variants: function, method + +- func: lcm.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: lcm_out + tags: pointwise + +- func: lcm(Tensor self, Tensor other) -> Tensor + structured_delegate: lcm.out + variants: function, method + tags: pointwise + +- func: lcm_(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: lcm.out + variants: function, method + +# NOTE [ grid_sampler Native Functions ] +# `grid_sampler` is _supposed to_ do all the shape checking and then dispatch to +# one of `cudnn_grid_sampler`, `grid_sampler_2d`, or `grid_sampler_3d`, each of +# which has the corresponding backward defined as native functions as well. +# However, we do shape checking everywhere for now since each of the mentioned +# functions can be called directly, which will lead to crashes otherwise. +# See https://github.com/pytorch/pytorch/issues/73187 for more information. +# +# There is also _grid_sampler_2d_backward_cpu_fallback which is an +# implementation detail of grid_sampler_2d and is only exposed here for testing +# purposes. +# +# Additionally, arguments `padding_mode` and `interpolation_mode` are cast to +# enums defined in `native/GridSampler.h`. `cudnn_grid_sampler` doesn't take in +# `interpolation_mode` because it only supports Bilinear interpolation mode. +# Nor does it take in `align_corners` because it only supports the mode +# `align_corners = True`. +- func: grid_sampler(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor + +- func: grid_sampler_2d(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor + dispatch: + CPU, QuantizedCPU: grid_sampler_2d_cpu + CUDA: grid_sampler_2d_cuda + MPS: grid_sampler_2d_mps + autogen: grid_sampler_2d.out + tags: core + +# `grid_sampler_2d_backward` takes in `output_mask` to optimize performance for +# the case where `input` doesn't require gradient. Gradient for `grid` is always +# computed (only `output_mask[0]` is checked by the implementations). +- func: grid_sampler_2d_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners, bool[2] output_mask) -> (Tensor, Tensor) + dispatch: + CPU: grid_sampler_2d_backward_cpu + CUDA: grid_sampler_2d_backward_cuda + autogen: grid_sampler_2d_backward.out + +# See NOTE [ grid_sample CPU fallback ] +- func: _grid_sampler_2d_cpu_fallback(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor + dispatch: + CompositeExplicitAutograd: _grid_sampler_2d_cpu_fallback + autogen: _grid_sampler_2d_cpu_fallback.out + +- func: _grid_sampler_2d_cpu_fallback_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> (Tensor, Tensor) + dispatch: + CompositeExplicitAutograd: _grid_sampler_2d_cpu_fallback_backward + +- func: grid_sampler_3d(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor + dispatch: + CPU: grid_sampler_3d_cpu + CUDA: grid_sampler_3d_cuda + MPS: grid_sampler_3d_mps + autogen: grid_sampler_3d.out + +# `grid_sampler_3d_backward` takes in `output_mask` to optimize performance for +# the case where `input` doesn't require gradient. Gradient for `grid` is always +# computed (only `output_mask[0]` is checked by the implementations). +- func: grid_sampler_3d_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners, bool[2] output_mask) -> (Tensor, Tensor) + dispatch: + CPU: grid_sampler_3d_backward_cpu + CUDA: grid_sampler_3d_backward_cuda + autogen: grid_sampler_3d_backward.out + +- func: hann_window(int window_length, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: hann_window + autogen: hann_window.out + +- func: hann_window.periodic(int window_length, bool periodic, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: hann_window + autogen: hann_window.periodic_out + +- func: hamming_window(int window_length, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: hamming_window + autogen: hamming_window.out + +- func: hamming_window.periodic(int window_length, bool periodic, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: hamming_window + autogen: hamming_window.periodic_out + +- func: hamming_window.periodic_alpha(int window_length, bool periodic, float alpha, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: hamming_window + autogen: hamming_window.periodic_alpha_out + +- func: hamming_window.periodic_alpha_beta(int window_length, bool periodic, float alpha, float beta, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: hamming_window + autogen: hamming_window.periodic_alpha_beta_out + +- func: kaiser_window(int window_length, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: kaiser_window + autogen: kaiser_window.out + +- func: kaiser_window.periodic(int window_length, bool periodic, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: kaiser_window + autogen: kaiser_window.periodic_out + +- func: kaiser_window.beta(int window_length, bool periodic, float beta, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: kaiser_window + autogen: kaiser_window.beta_out + +- func: hinge_embedding_loss(Tensor self, Tensor target, float margin=1.0, int reduction=Mean) -> Tensor + +- func: group_norm(Tensor input, int num_groups, Tensor? weight=None, Tensor? bias=None, float eps=1e-05, bool cudnn_enabled=True) -> Tensor + +- func: native_group_norm(Tensor input, Tensor? weight, Tensor? bias, SymInt N, SymInt C, SymInt HxW, int group, float eps) -> (Tensor, Tensor, Tensor) + dispatch: + CPU, CUDA: native_group_norm + CompositeExplicitAutograd: math_group_norm + autogen: native_group_norm.out + tags: core + +- func: native_group_norm_backward(Tensor grad_out, Tensor input, Tensor mean, Tensor rstd, Tensor? weight, SymInt N, SymInt C, SymInt HxW, int group, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + dispatch: + CPU, CUDA: native_group_norm_backward + autogen: native_group_norm_backward.out + tags: core + +# Real to complex forward FFT +- func: _fft_r2c(Tensor self, int[] dim, int normalization, bool onesided) -> Tensor + variants: function + dispatch: + CPU: _fft_r2c_mkl + CUDA: _fft_r2c_cufft + MPS: _fft_r2c_mps + tags: core + +- func: _fft_r2c.out(Tensor self, int[] dim, int normalization, bool onesided, *, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CPU: _fft_r2c_mkl_out + CUDA: _fft_r2c_cufft_out + MPS: _fft_r2c_mps_out + +# Complex to real inverse FFT +- func: _fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor + variants: function + dispatch: + CPU: _fft_c2r_mkl + CUDA: _fft_c2r_cufft + MPS: _fft_c2r_mps + tags: core + +- func: _fft_c2r.out(Tensor self, int[] dim, int normalization, SymInt last_dim_size, *, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CPU: _fft_c2r_mkl_out + CUDA: _fft_c2r_cufft_out + MPS: _fft_c2r_mps_out + +# Standard complex to complex FFT (forward or backward) +- func: _fft_c2c(Tensor self, SymInt[] dim, int normalization, bool forward) -> Tensor + variants: function + dispatch: + CPU: _fft_c2c_mkl + CUDA: _fft_c2c_cufft + MPS: _fft_c2c_mps + +- func: _fft_c2c.out(Tensor self, SymInt[] dim, int normalization, bool forward, *, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CPU: _fft_c2c_mkl_out + CUDA: _fft_c2c_cufft_out + MPS: _fft_c2c_mps_out + +- func: _validate_compressed_sparse_indices(bool is_crow, Tensor compressed_idx, Tensor plain_idx, int cdim, int dim, int nnz) -> () + device_check: NoCheck + variants: function + dispatch: + CPU: _validate_compressed_sparse_indices_cpu + CUDA: _validate_compressed_sparse_indices_cuda + +- func: _cufft_get_plan_cache_size(DeviceIndex device_index) -> int + +- func: _cufft_get_plan_cache_max_size(DeviceIndex device_index) -> int + +- func: _cufft_set_plan_cache_max_size(DeviceIndex device_index, int max_size) -> () + +- func: _cufft_clear_plan_cache(DeviceIndex device_index) -> () + +- func: index.Tensor(Tensor self, Tensor?[] indices) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: index.Tensor_out + variants: function, method + dispatch: + QuantizedCPU: quantized_index + tags: [core, dynamic_output_shape] + # NB: This function is special-cased in tools/autograd/gen_variable_type.py + # NB: The following functions are declared in aten/src/ATen/templates/TensorBody.h and defined in aten/src/ATen/TensorIndexing.cpp: + # - Tensor Tensor::index(ArrayRef indices) + # - Tensor Tensor::index(std::initializer_list indices) + +- func: index.Tensor_out(Tensor self, Tensor?[] indices, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + structured: True + structured_inherits: TensorIteratorBase + precomputed: + - indices -> DimVector sizes, DimVector strides + dispatch: + CPU, CUDA, MPS: index_out + +# Used by inductor to signal indexing without bounds checks +# Note that we don't support boolean indexing, to avoid dynamic output shapes +- func: _unsafe_index.Tensor(Tensor self, Tensor?[] indices) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: _unsafe_index + +# Used by inductor to generate masked loads +# Note that we don't support boolean indexing, to avoid dynamic output shapes +- func: _unsafe_masked_index(Tensor self, Tensor mask, Tensor?[] indices, Scalar fill) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: _unsafe_masked_index + +- func: _unsafe_masked_index_put_accumulate(Tensor self, Tensor mask, Tensor?[] indices, Tensor values) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: _unsafe_masked_index_put_accumulate + +- func: index_copy.out(Tensor self, int dim, Tensor index, Tensor source, *, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + precomputed: + - dim -> int dim + dispatch: + CPU, CUDA: index_copy_out + MPS: index_copy_out_mps + +- func: index_copy_(Tensor(a!) self, int dim, Tensor index, Tensor source) -> Tensor(a!) + variants: method + structured_delegate: index_copy.out + +- func: index_copy(Tensor self, int dim, Tensor index, Tensor source) -> Tensor + variants: function, method + structured_delegate: index_copy.out + +- func: index_copy_.dimname(Tensor(a!) self, Dimname dim, Tensor index, Tensor source) -> Tensor(a!) + variants: method + +- func: index_copy.dimname(Tensor self, Dimname dim, Tensor index, Tensor source) -> Tensor + variants: function, method + +- func: index_put_(Tensor(a!) self, Tensor?[] indices, Tensor values, bool accumulate=False) -> Tensor(a!) + device_check: NoCheck # delegate to _index_put_impl_, which leverages TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: index_put_ + autogen: index_put.out + # NB: The following functions are declared in aten/src/ATen/templates/TensorBody.h and defined in aten/src/ATen/TensorIndexing.cpp: + # - Tensor & Tensor::index_put_(ArrayRef indices, Tensor const & rhs) + # - Tensor & Tensor::index_put_(ArrayRef indices, Scalar v) + # - Tensor & Tensor::index_put_(std::initializer_list indices, Tensor const & rhs) + # - Tensor & Tensor::index_put_(std::initializer_list indices, Scalar v) + +- func: index_put(Tensor self, Tensor?[] indices, Tensor values, bool accumulate=False) -> Tensor + device_check: NoCheck # delegate to _index_put_impl_ after clone, which leverages TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: index_put + tags: core + +- func: _unsafe_index_put(Tensor self, Tensor?[] indices, Tensor values, bool accumulate=False) -> Tensor + device_check: NoCheck # delegate to _index_put_impl_ after clone, which leverages TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: _unsafe_index_put + +- func: _index_put_impl_(Tensor(a!) self, Tensor?[] indices, Tensor values, bool accumulate=False, bool unsafe=False) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CPU, CUDA, MPS: _index_put_impl_ + QuantizedCPU: _index_put_impl_quantized_cpu_ + QuantizedCUDA: _index_put_impl_quantized_cuda_ + autogen: _index_put_impl, _index_put_impl.out + +- func: instance_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool use_input_stats, float momentum, float eps, bool cudnn_enabled) -> Tensor + variants: function + +- func: isclose(Tensor self, Tensor other, float rtol=1e-05, float atol=1e-08, bool equal_nan=False) -> Tensor + variants: function, method + +- func: isin.Tensor_Tensor_out(Tensor elements, Tensor test_elements, *, bool assume_unique=False, bool invert=False, Tensor(a!) out) -> Tensor(a!) + variants: function + structured: True + dispatch: + CPU, CUDA: isin_Tensor_Tensor_out + MPS: isin_Tensor_Tensor_out_mps + +- func: isin.Tensor_Tensor(Tensor elements, Tensor test_elements, *, bool assume_unique=False, bool invert=False) -> Tensor + variants: function + structured_delegate: isin.Tensor_Tensor_out + +- func: isin.Tensor_Scalar_out(Tensor elements, Scalar test_element, *, bool assume_unique=False, bool invert=False, Tensor(a!) out) -> Tensor(a!) + variants: function + structured: True + dispatch: + CPU, CUDA, MPS: isin_Tensor_Scalar_out + +- func: isin.Tensor_Scalar(Tensor elements, Scalar test_element, *, bool assume_unique=False, bool invert=False) -> Tensor + variants: function + structured_delegate: isin.Tensor_Scalar_out + +- func: isin.Scalar_Tensor_out(Scalar element, Tensor test_elements, *, bool assume_unique=False, bool invert=False, Tensor(a!) out) -> Tensor(a!) + variants: function + structured: True + dispatch: + CPU, CUDA: isin_Scalar_Tensor_out + MPS: isin_Scalar_Tensor_out_mps + +- func: isin.Scalar_Tensor(Scalar element, Tensor test_elements, *, bool assume_unique=False, bool invert=False) -> Tensor + variants: function + structured_delegate: isin.Scalar_Tensor_out + +- func: isnan(Tensor self) -> Tensor + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CPU, CUDA, MPS, MTIA: isnan + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_isnan + SparseCPU, SparseCUDA, SparseMPS: isnan_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: isnan_sparse_csr + autogen: isnan.out + tags: [core, pointwise] + +- func: is_distributed(Tensor self) -> bool + variants: function, method + device_check: NoCheck + device_guard: False + +- func: is_floating_point(Tensor self) -> bool + variants: function, method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: is_complex(Tensor self) -> bool + variants: function, method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: is_conj(Tensor self) -> bool + variants: function, method + device_guard: False + manual_cpp_binding: True + +- func: _is_zerotensor(Tensor self) -> bool + variants: function, method + device_guard: False + manual_cpp_binding: True + +- func: is_neg(Tensor self) -> bool + variants: function, method + device_guard: False + manual_cpp_binding: True + +- func: isreal(Tensor self) -> Tensor + variants: function, method + +- func: is_nonzero(Tensor self) -> bool + variants: function, method + device_check: NoCheck + device_guard: False + +- func: is_same_size(Tensor self, Tensor other) -> bool + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: nested_is_same_size + CompositeExplicitAutograd: is_same_size + +- func: is_signed(Tensor self) -> bool + variants: function, method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: is_inference(Tensor self) -> bool + variants: function, method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: numel(Tensor self) -> int + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: dim(Tensor self) -> int + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: get_device(Tensor self) -> int + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: storage_offset(Tensor self) -> int + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: is_contiguous(Tensor self) -> bool + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: is_contiguous.memory_format(Tensor self, MemoryFormat memory_format) -> bool + variants: method + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: kl_div(Tensor self, Tensor target, int reduction=Mean, *, bool log_target=False) -> Tensor + +- func: kron(Tensor self, Tensor other) -> Tensor + variants: function, method + +- func: kron.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + +- func: kthvalue(Tensor self, SymInt k, int dim=-1, bool keepdim=False) -> (Tensor values, Tensor indices) + variants: function, method + dispatch: + CompositeExplicitAutograd: kthvalue + +- func: kthvalue.values(Tensor self, SymInt k, int dim=-1, bool keepdim=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + dispatch: + CPU: kthvalue_out_cpu + CUDA: kthvalue_out_cuda + MPS: kthvalue_out_mps + +- func: kthvalue.dimname(Tensor self, SymInt k, Dimname dim, bool keepdim=False) -> (Tensor values, Tensor indices) + variants: function, method + +- func: kthvalue.dimname_out(Tensor self, SymInt k, Dimname dim, bool keepdim=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + +- func: layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight=None, Tensor? bias=None, float eps=1e-05, bool cudnn_enable=True) -> Tensor + dispatch: + CompositeImplicitAutograd: layer_norm_symint + +- func: native_layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight, Tensor? bias, float eps) -> (Tensor, Tensor, Tensor) + dispatch: + CPU: layer_norm_cpu + CUDA: layer_norm_cuda + MPS: layer_norm_mps + CompositeExplicitAutograd: math_native_layer_norm + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: nested_layer_norm + autogen: native_layer_norm.out + tags: core + +- func: native_layer_norm_backward(Tensor grad_out, Tensor input, SymInt[] normalized_shape, Tensor mean, Tensor rstd, Tensor? weight, Tensor? bias, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + dispatch: + CPU: layer_norm_backward_cpu + CUDA: layer_norm_backward_cuda + MPS: layer_norm_backward_mps + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: layer_norm_backward_nested + autogen: native_layer_norm_backward.out + tags: core + +- func: rms_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight=None, float? eps=None) -> Tensor + dispatch: + CompositeImplicitAutograd: rms_norm_symint + +- func: _fused_rms_norm(Tensor input, int[] normalized_shape, Tensor? weight, float? eps) -> (Tensor, Tensor) + dispatch: + CUDA: _fused_rms_norm_cuda + MPS: _fused_rms_norm_mps + XPU: _fused_rms_norm_xpu + CompositeImplicitAutograd: rms_norm_composite + +- func: _fused_rms_norm_backward(Tensor grad_out, Tensor input, int[] normalized_shape, Tensor rstd, Tensor? weight, bool[2] output_mask) -> (Tensor, Tensor) + dispatch: + CUDA: _fused_rms_norm_backward_cuda + XPU: _fused_rms_norm_backward_xpu + +- func: nan_to_num(Tensor self, float? nan=None, float? posinf=None, float? neginf=None) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: nan_to_num + SparseCPU, SparseCUDA, SparseMPS: nan_to_num_sparse + tags: pointwise + +- func: nan_to_num_(Tensor(a!) self, float? nan=None, float? posinf=None, float? neginf=None) -> Tensor(a!) + variants: function, method + dispatch: + CompositeExplicitAutograd: nan_to_num_ + SparseCPU, SparseCUDA, SparseMPS: nan_to_num_sparse_ + tags: pointwise + +- func: nan_to_num.out(Tensor self, float? nan=None, float? posinf=None, float? neginf=None, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MTIA: nan_to_num_out + MPS: nan_to_num_out_mps + SparseCPU, SparseCUDA, SparseMPS: nan_to_num_sparse_out + tags: pointwise + +- func: linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: linear + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: nested_linear + MPS: _mps_linear + +- func: linear_backward(Tensor self, Tensor grad_output, Tensor weight, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: nested_linear_backward + MPS: mps_linear_backward + autogen: linear_backward.out + +- func: linear.out(Tensor input, Tensor weight, Tensor? bias=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CompositeExplicitAutograd: linear_out + +- func: mkldnn_linear(Tensor self, Tensor weight, Tensor? bias=None) -> Tensor + python_module: nn + dispatch: + MkldnnCPU: mkldnn_linear + autogen: mkldnn_linear.out + +- func: mkldnn_linear_backward_input(int[] input_size, Tensor grad_output, Tensor weight) -> Tensor + dispatch: + MkldnnCPU: mkldnn_linear_backward_input + autogen: mkldnn_linear_backward_input.out + +- func: mkldnn_linear_backward_weights(Tensor grad_output, Tensor input, Tensor weight, bool bias_defined) -> (Tensor, Tensor) + dispatch: + MkldnnCPU: mkldnn_linear_backward_weights + autogen: mkldnn_linear_backward_weights.out + +- func: mkldnn_linear_backward(Tensor self, Tensor grad_output, Tensor weight, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + dispatch: + MkldnnCPU: mkldnn_linear_backward + autogen: mkldnn_linear_backward.out + +- func: _cslt_compress(Tensor input) -> Tensor + dispatch: + CUDA: _cslt_compress + +- func: _cslt_sparse_mm(Tensor compressed_A, Tensor dense_B, Tensor? bias=None, Tensor? alpha=None, ScalarType? out_dtype=None, bool transpose_result=False, int alg_id=0, int split_k=1, int split_k_mode=-1) -> Tensor + dispatch: + CUDA: _cslt_sparse_mm + tags: needs_fixed_stride_order + +- func: _cslt_sparse_mm_search(Tensor compressed_A, Tensor dense_B, Tensor? bias=None, Tensor? alpha=None, ScalarType? out_dtype=None, bool transpose_result=False) -> int + dispatch: + CUDA: _cslt_sparse_mm_search + +- func: _sparse_semi_structured_tile(Tensor input, str algorithm="", bool use_cutlass=True) -> (Tensor, Tensor, Tensor, Tensor, Tensor) + dispatch: + CUDA: _sparse_semi_structured_tile + +- func: _sparse_semi_structured_apply(Tensor input, Tensor thread_masks) -> (Tensor, Tensor) + dispatch: + CUDA: _sparse_semi_structured_apply + +- func: _sparse_semi_structured_apply_dense(Tensor input, Tensor thread_masks) -> Tensor + dispatch: + CUDA: _sparse_semi_structured_apply_dense + +# DEPRECATED: Use torch.__sparse_semi_structured_mm/torch._sparse_semi_structured_addmm instead +- func: _sparse_semi_structured_linear(Tensor input, Tensor weight, Tensor meta, *, Tensor? bias=None, str? activation=None, ScalarType? out_dtype=None) -> Tensor + dispatch: + CUDA: _sparse_semi_structured_linear + +- func: _sparse_semi_structured_mm(Tensor mat1, Tensor mat1_meta, Tensor mat2, *, ScalarType? out_dtype=None) -> Tensor + dispatch: + CUDA: _sparse_semi_structured_mm + +- func: _sparse_semi_structured_addmm(Tensor input, Tensor mat1, Tensor mat1_meta, Tensor mat2, *, Scalar alpha=1, Scalar beta=1, ScalarType? out_dtype=None) -> Tensor + dispatch: + CUDA: _sparse_semi_structured_addmm + +- func: _mixed_dtypes_linear(Tensor input, Tensor weight, Tensor scale, *, Tensor? bias=None, str? activation=None) -> Tensor + dispatch: + CUDA: _mixed_dtypes_linear + +- func: fbgemm_linear_int8_weight_fp32_activation(Tensor input, Tensor weight, Tensor packed, Tensor col_offsets, Scalar weight_scale, Scalar weight_zero_point, Tensor bias) -> Tensor + +- func: fbgemm_linear_int8_weight(Tensor input, Tensor weight, Tensor packed, Tensor col_offsets, Scalar weight_scale, Scalar weight_zero_point, Tensor bias) -> Tensor + +- func: fbgemm_linear_quantize_weight(Tensor input) -> (Tensor, Tensor, float, int) + +- func: fbgemm_pack_gemm_matrix_fp16(Tensor input) -> Tensor + +- func: _wrapped_linear_prepack(Tensor weight, Tensor weight_scale, Tensor weight_zero_point, Tensor bias) -> Tensor + +- func: _wrapped_quantized_linear_prepacked(Tensor input, Tensor input_scale, Tensor input_zero_point, Tensor packed_weight, Tensor output_scale, Tensor output_zero_point, int out_channel) -> Tensor + +- func: fbgemm_linear_fp16_weight_fp32_activation(Tensor input, Tensor packed_weight, Tensor? bias) -> Tensor + +- func: fbgemm_linear_fp16_weight_fp32_activation.out(Tensor input, Tensor packed_weight, Tensor? bias, Tensor(a!) output) -> Tensor + +- func: fbgemm_linear_fp16_weight(Tensor input, Tensor packed_weight, Tensor bias) -> Tensor + +- func: fbgemm_linear_fp16_weight.out(Tensor input, Tensor packed_weight, Tensor bias, Tensor(a!) output) -> Tensor + +- func: fbgemm_pack_quantized_matrix(Tensor input) -> Tensor + +- func: fbgemm_pack_quantized_matrix.KN(Tensor input, int K, int N) -> Tensor + +- func: ldexp.Tensor(Tensor self, Tensor other) -> Tensor + variants: function, method + tags: pointwise + dispatch: + CompositeExplicitAutograd: ldexp + +- func: ldexp_(Tensor(a!) self, Tensor other) -> Tensor(a!) + variants: function, method + tags: pointwise + dispatch: + CompositeExplicitAutograd: ldexp_ + +- func: ldexp.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + tags: pointwise + dispatch: + CompositeExplicitAutograd: ldexp_out + +- func: linspace(Scalar start, Scalar end, int steps, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: linspace + +- func: linspace.Tensor_Tensor(Tensor start, Tensor end, int steps, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + category_override: factory + dispatch: + CompositeExplicitAutograd: linspace + +- func: linspace.Tensor_Scalar(Tensor start, Scalar end, int steps, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + category_override: factory + dispatch: + CompositeExplicitAutograd: linspace + +- func: linspace.Scalar_Tensor(Scalar start, Tensor end, int steps, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + category_override: factory + dispatch: + CompositeExplicitAutograd: linspace + +- func: linspace.out(Scalar start, Scalar end, int steps, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, Meta: linspace_out + CUDA: linspace_cuda_out + MPS: linspace_out_mps + +- func: linspace.Tensor_Tensor_out(Tensor start, Tensor end, int steps, *, Tensor(a!) out) -> Tensor(a!) + category_override: factory + dispatch: + CompositeExplicitAutograd: linspace_out + +- func: linspace.Tensor_Scalar_out(Tensor start, Scalar end, int steps, *, Tensor(a!) out) -> Tensor(a!) + category_override: factory + dispatch: + CompositeExplicitAutograd: linspace_out + +- func: linspace.Scalar_Tensor_out(Scalar start, Tensor end, int steps, *, Tensor(a!) out) -> Tensor(a!) + category_override: factory + dispatch: + CompositeExplicitAutograd: linspace_out + +- func: log(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: log.out + variants: function, method + tags: [core, pointwise] + +- func: log_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: log.out + variants: function, method + tags: pointwise + +- func: log.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: log_out + tags: pointwise + +- func: log10(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: log10.out + variants: function, method + tags: [core, pointwise] + +- func: log10_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: log10.out + variants: function, method + tags: pointwise + +- func: log10.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: log10_out + tags: pointwise + +- func: log1p(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: log1p.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: log1p_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: log1p_sparse_csr + tags: [core, pointwise] + +- func: log1p_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: log1p.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: log1p_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: log1p_sparse_csr_ + tags: pointwise + +- func: log1p.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: log1p_out + SparseCPU, SparseCUDA, SparseMPS: log1p_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: log1p_sparse_csr_out + tags: pointwise + +- func: log2(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: log2.out + variants: function, method + tags: [core, pointwise] + +- func: log2_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: log2.out + variants: function, method + tags: pointwise + +- func: log2.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: log2_out + tags: pointwise + +- func: logaddexp.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: logaddexp_out + tags: pointwise + +- func: logaddexp(Tensor self, Tensor other) -> Tensor + variants: method, function + structured_delegate: logaddexp.out + tags: pointwise + +- func: logaddexp2.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: logaddexp2_out + tags: pointwise + +- func: logaddexp2(Tensor self, Tensor other) -> Tensor + variants: method, function + structured_delegate: logaddexp2.out + tags: pointwise + +- func: xlogy.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: xlogy.OutTensor + variants: function, method + tags: pointwise + +- func: xlogy.Scalar_Self(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: xlogy + tags: pointwise + +- func: xlogy.Scalar_Other(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: xlogy + tags: pointwise + +# xlogy: inplace variant +- func: xlogy_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: xlogy.OutTensor + tags: pointwise + +- func: xlogy_.Scalar_Other(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: xlogy_ + +# xlogy: out variant +- func: xlogy.OutTensor(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + variants: function + dispatch: + CPU, CUDA, MPS: xlogy_out + tags: pointwise + +- func: xlogy.OutScalar_Self(Scalar self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: xlogy_out + tags: pointwise + +- func: xlogy.OutScalar_Other(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: xlogy_out + tags: pointwise + +- func: logspace(Scalar start, Scalar end, int steps, float base=10.0, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: logspace + +- func: logspace.Tensor_Tensor(Tensor start, Tensor end, int steps, float base=10.0, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + category_override: factory + dispatch: + CompositeExplicitAutograd: logspace + +- func: logspace.Tensor_Scalar(Tensor start, Scalar end, int steps, float base=10.0, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + category_override: factory + dispatch: + CompositeExplicitAutograd: logspace + +- func: logspace.Scalar_Tensor(Scalar start, Tensor end, int steps, float base=10.0, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + category_override: factory + dispatch: + CompositeExplicitAutograd: logspace + +- func: logspace.out(Scalar start, Scalar end, int steps, float base=10.0, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, Meta: logspace_out + CUDA: logspace_cuda_out + +- func: logspace.Tensor_Tensor_out(Tensor start, Tensor end, int steps, float base=10.0, *, Tensor(a!) out) -> Tensor(a!) + category_override: factory + dispatch: + CompositeExplicitAutograd: logspace_out + +- func: logspace.Tensor_Scalar_out(Tensor start, Scalar end, int steps, float base=10.0, *, Tensor(a!) out) -> Tensor(a!) + category_override: factory + dispatch: + CompositeExplicitAutograd: logspace_out + +- func: logspace.Scalar_Tensor_out(Scalar start, Tensor end, int steps, float base=10.0, *, Tensor(a!) out) -> Tensor(a!) + category_override: factory + dispatch: + CompositeExplicitAutograd: logspace_out + +# log_softmax allows positional dtype, unlike most operators, because kwonly is BC-breaking when loading jit models. +- func: log_softmax.int(Tensor self, int dim, ScalarType? dtype=None) -> Tensor + variants: function, method + +- func: log_softmax.int_out(Tensor self, int dim, ScalarType? dtype=None, *, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CompositeExplicitAutograd: log_softmax_out + +- func: log_softmax.Dimname(Tensor self, Dimname dim, *, ScalarType? dtype=None) -> Tensor + variants: function, method + +- func: _log_softmax(Tensor self, int dim, bool half_to_float) -> Tensor + structured_delegate: _log_softmax.out + tags: core + +- func: _log_softmax.out(Tensor self, int dim, bool half_to_float, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: log_softmax_cpu_out + CUDA: log_softmax_cuda_out + MTIA: log_softmax_mtia_out + MPS: log_softmax_mps_out + +- func: _log_softmax_backward_data(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype) -> Tensor + structured_delegate: _log_softmax_backward_data.out + +- func: _log_softmax_backward_data.out(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: log_softmax_backward_cpu_out + CUDA: log_softmax_backward_cuda_out + MTIA: log_softmax_backward_mtia_out + MPS: log_softmax_backward_mps_out + +- func: _logcumsumexp(Tensor self, int dim) -> Tensor + dispatch: + CPU: _logcumsumexp_cpu + CUDA: _logcumsumexp_cuda + MPS: _logcumsumexp_mps + +- func: _logcumsumexp.out(Tensor self, int dim, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: _logcumsumexp_out_cpu + CUDA: _logcumsumexp_out_cuda + MPS: _logcumsumexp_out_mps + +- func: logcumsumexp(Tensor self, int dim) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: logcumsumexp + +- func: logcumsumexp.out(Tensor self, int dim, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: logcumsumexp_out + +- func: logcumsumexp.dimname(Tensor self, Dimname dim) -> Tensor + variants: function, method + +- func: logcumsumexp.dimname_out(Tensor self, Dimname dim, *, Tensor(a!) out) -> Tensor(a!) + +- func: logsumexp(Tensor self, int[1] dim, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: logsumexp + tags: reduction + +- func: logsumexp.out(Tensor self, int[1] dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + # calls squeeze + CompositeExplicitAutogradNonFunctional: logsumexp_out + tags: reduction + +- func: logsumexp.names(Tensor self, Dimname[1] dim, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: logsumexp.names_out(Tensor self, Dimname[1] dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: reduction + +- func: margin_ranking_loss(Tensor input1, Tensor input2, Tensor target, float margin=0.0, int reduction=Mean) -> Tensor + +- func: matmul(Tensor self, Tensor other) -> Tensor + variants: function, method + dispatch: + CompositeImplicitAutograd: matmul + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: matmul_nested + +- func: matmul_backward(Tensor grad, Tensor self, Tensor other, bool[2] mask) -> (Tensor, Tensor) + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: matmul_backward_nested + autogen: matmul_backward.out + +- func: matmul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeImplicitAutograd: matmul_out + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: matmul_out_nested + +# Alias to linalg.matrix_power +- func: matrix_power(Tensor self, int n) -> Tensor + variants: function, method + +# Alias to linalg.matrix_power +- func: matrix_power.out(Tensor self, int n, *, Tensor(a!) out) -> Tensor(a!) + +# Alias to linalg.matrix_exp +- func: matrix_exp(Tensor self) -> Tensor + variants: function, method + +# This function should be deprecated in favor of differential_analytic_matrix_function in FunctionsManual.cpp +- func: matrix_exp_backward(Tensor self, Tensor grad) -> Tensor + +# DEPRECATED: Use torch.aminmax instead +- func: _aminmax(Tensor self) -> (Tensor, Tensor) + dispatch: + CPU, CUDA: _aminmax_all + autogen: _aminmax.out + +# DEPRECATED: Use torch.aminmax instead +- func: _aminmax.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor, Tensor) + dispatch: + CPU, CUDA: _aminmax + autogen: _aminmax.dim_out + +- func: aminmax(Tensor self, *, int? dim=None, bool keepdim=False) -> (Tensor min, Tensor max) + device_check: NoCheck # TensorIterator + structured_delegate: aminmax.out + variants: function, method + tags: reduction + +- func: aminmax.out(Tensor self, *, int? dim=None, bool keepdim=False, Tensor(a!) min, Tensor(b!) max) -> (Tensor(a!) min, Tensor(b!) max) + device_check: NoCheck # TensorIterator + structured: True + dispatch: + CPU, CUDA: aminmax_out + MPS: aminmax_out_mps + tags: reduction + +- func: _compute_linear_combination(Tensor input, Tensor coefficients) -> Tensor + dispatch: + CPU, CUDA: _compute_linear_combination + +- func: _compute_linear_combination.out(Tensor input, Tensor coefficients, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: _compute_linear_combination_out + +- func: max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) + device_check: NoCheck # TensorIterator + structured_delegate: max.dim_max + variants: function, method + dispatch: + QuantizedCPU, QuantizedCUDA: qmax + tags: [core, reduction] + +- func: max.dim_max(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) max, Tensor(b!) max_values) -> (Tensor(a!) values, Tensor(b!) indices) + device_check: NoCheck # TensorIterator + structured: True + precomputed: + - dim -> int dim + dispatch: + CPU, CUDA, MTIA: max_out + MPS: max_out_mps + tags: reduction + +- func: max.names_dim(Tensor self, Dimname dim, bool keepdim=False) -> (Tensor values, Tensor indices) + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: max.names_dim_max(Tensor self, Dimname dim, bool keepdim=False, *, Tensor(a!) max, Tensor(b!) max_values) -> (Tensor(a!) values, Tensor(b!) indices) + device_check: NoCheck # TensorIterator + tags: reduction + +- func: value_selecting_reduction_backward(Tensor grad, int dim, Tensor indices, SymInt[] sizes, bool keepdim) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: value_selecting_reduction_backward_symint + NestedTensorCPU, NestedTensorCUDA: value_selecting_reduction_backward_nested_symint + +- func: amax(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor + variants: function, method + structured_delegate: amax.out + tags: [core, reduction] + +- func: amax.out(Tensor self, int[1] dim=[], bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU, CUDA: amax_out + MPS: amax_out_mps + tags: reduction + +# Return: (Tensor output, Tensor indices) +- func: max_pool1d_with_indices(Tensor self, int[1] kernel_size, int[1] stride=[], int[1] padding=0, int[1] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor) + +- func: max_pool1d(Tensor self, int[1] kernel_size, int[1] stride=[], int[1] padding=0, int[1] dilation=1, bool ceil_mode=False) -> Tensor + +- func: max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor + dispatch: + CompositeImplicitAutograd: max_pool2d + MPS: mps_max_pool2d + +- func: max_pool2d_backward(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor + dispatch: + MPS: mps_max_pool2d_backward + autogen: max_pool2d_backward.out + +- func: mkldnn_max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor + dispatch: + MkldnnCPU: mkldnn_max_pool2d + autogen: mkldnn_max_pool2d.out + +- func: mkldnn_max_pool2d_backward(Tensor grad_output, Tensor output, Tensor input, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor + dispatch: + MkldnnCPU: mkldnn_max_pool2d_backward + autogen: mkldnn_max_pool2d_backward.out + +- func: mkldnn_max_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> Tensor + dispatch: + MkldnnCPU: mkldnn_max_pool3d + autogen: mkldnn_max_pool3d.out + +- func: mkldnn_max_pool3d_backward(Tensor grad_output, Tensor output, Tensor input, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> Tensor + dispatch: + MkldnnCPU: mkldnn_max_pool3d_backward + autogen: mkldnn_max_pool3d_backward.out + +- func: quantized_max_pool1d(Tensor self, int[1] kernel_size, int[1] stride=[], int[1] padding=0, int[1] dilation=1, bool ceil_mode=False) -> Tensor + dispatch: + QuantizedCPU: quantized_max_pool1d + autogen: quantized_max_pool1d.out + +- func: quantized_max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor + dispatch: + QuantizedCPU: quantized_max_pool2d + QuantizedCUDA: quantized_max_pool2d_cudnn + autogen: quantized_max_pool2d.out + +- func: quantized_max_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> Tensor + dispatch: + QuantizedCPU: quantized_max_pool3d + autogen: quantized_max_pool3d.out + +- func: max_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> Tensor + +# The CPU and GPU dispatch variants are named weirdly here because otherwise there +# are namespacing issues in C++ +- func: mean(Tensor self, *, ScalarType? dtype=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: mean + tags: [core, reduction] + +# For normal naming convention this should be `mean.out`. However since we already have `mean.out` we have to rename this. +- func: mean.dtype_out(Tensor self, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CompositeExplicitAutograd: mean_dtype_out + tags: reduction + +- func: mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + structured_delegate: mean.out + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + QuantizedCPU: mean_quantized_cpu + tags: [core, reduction] + +- func: mean.out(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + structured: True + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: mean_out + MPS: mean_out_mps + QuantizedCPU: mean_out_quantized_cpu + tags: reduction + +- func: mean.names_dim(Tensor self, Dimname[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: mean.names_out(Tensor self, Dimname[1] dim, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: reduction + +- func: nanmean(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + device_check: NoCheck # Composite + variants: function, method + +- func: nanmean.out(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # Composite + +- func: median(Tensor self) -> Tensor + variants: function, method + dispatch: + CPU: median_cpu + CUDA: median_cuda + MPS: median_mps + autogen: median.out + +- func: median.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) + variants: function, method + dispatch: + CompositeExplicitAutograd: median + +- func: median.dim_values(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + dispatch: + CPU: median_out_cpu + CUDA: median_out_cuda + MPS: median_out_mps + +- func: median.names_dim(Tensor self, Dimname dim, bool keepdim=False) -> (Tensor values, Tensor indices) + variants: function, method + +- func: median.names_dim_values(Tensor self, Dimname dim, bool keepdim=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + +- func: nanmedian(Tensor self) -> Tensor + variants: function, method + dispatch: + CPU: nanmedian_cpu + CUDA: nanmedian_cuda + MPS: nanmedian_mps + autogen: nanmedian.out + +- func: nanmedian.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) + variants: function, method + dispatch: + CompositeExplicitAutograd: nanmedian + +- func: nanmedian.dim_values(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + dispatch: + CPU: nanmedian_out_cpu + CUDA: nanmedian_out_cuda + MPS: nanmedian_out_mps + +- func: nanmedian.names_dim(Tensor self, Dimname dim, bool keepdim=False) -> (Tensor values, Tensor indices) + variants: function, method + +- func: nanmedian.names_dim_values(Tensor self, Dimname dim, bool keepdim=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + +- func: min.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) + device_check: NoCheck # TensorIterator + structured_delegate: min.dim_min + variants: function, method + dispatch: + QuantizedCPU, QuantizedCUDA: qmin + tags: [core, reduction] + +- func: min.dim_min(Tensor self, int dim, bool keepdim=False, *, Tensor(a!) min, Tensor(b!) min_indices) -> (Tensor(a!) values, Tensor(b!) indices) + device_check: NoCheck # TensorIterator + structured: True + precomputed: + - dim -> int dim + dispatch: + CPU, CUDA, MTIA: min_out + MPS: min_out_mps + tags: reduction + +- func: min.names_dim(Tensor self, Dimname dim, bool keepdim=False) -> (Tensor values, Tensor indices) + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: min.names_dim_min(Tensor self, Dimname dim, bool keepdim=False, *, Tensor(a!) min, Tensor(b!) min_indices) -> (Tensor(a!) values, Tensor(b!) indices) + device_check: NoCheck # TensorIterator + tags: reduction + +- func: amin(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor + variants: function, method + structured_delegate: amin.out + tags: [core, reduction] + +- func: amin.out(Tensor self, int[1] dim=[], bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU, CUDA: amin_out + MPS: amin_out_mps + tags: reduction + +# TODO: Add this function to MPS dispatch key so that we avoid declaring it in +# native_functions.yaml +# https://github.com/pytorch/pytorch/issues/77394 +- func: _mps_convolution(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups) -> Tensor + dispatch: + MPS: _mps_convolution + autogen: _mps_convolution.out + +- func: mps_convolution_backward(Tensor self, Tensor grad_output, Tensor weight, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + dispatch: + MPS: mps_convolution_backward + autogen: mps_convolution_backward.out + +- func: mkldnn_convolution(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups) -> Tensor + dispatch: + CompositeExplicitAutograd: mkldnn_convolution + autogen: mkldnn_convolution.out + +- func: mkldnn_rnn_layer(Tensor input, Tensor weight0, Tensor weight1, Tensor weight2, Tensor weight3, Tensor hx_, Tensor cx_, bool reverse, int[] batch_sizes, int mode, int hidden_size, int num_layers, bool has_biases, bool bidirectional, bool batch_first, bool train) -> (Tensor, Tensor, Tensor, Tensor) + dispatch: + CPU: mkldnn_rnn_layer + MkldnnCPU: mkldnn_rnn_layer + autogen: mkldnn_rnn_layer.out + +- func: mkldnn_rnn_layer_backward(Tensor input, Tensor weight1, Tensor weight2, Tensor weight3, Tensor weight4, Tensor hx_, Tensor cx_tmp, Tensor output, Tensor hy_, Tensor cy_, Tensor? grad_output, Tensor? grad_hy, Tensor? grad_cy, bool reverse, int mode, int hidden_size, int num_layers, bool has_biases, bool train, bool bidirectional, int[] batch_sizes, bool batch_first, Tensor workspace) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) + dispatch: + CPU: mkldnn_rnn_layer_backward + autogen: mkldnn_rnn_layer_backward.out + +- func: miopen_batch_norm(Tensor input, Tensor weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float exponential_average_factor, float epsilon) -> (Tensor, Tensor, Tensor) + dispatch: + CUDA: miopen_batch_norm + autogen: miopen_batch_norm.out + +- func: miopen_batch_norm_backward(Tensor input, Tensor grad_output, Tensor weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_var, float epsilon) -> (Tensor, Tensor, Tensor) + dispatch: + CUDA: miopen_batch_norm_backward + autogen: miopen_batch_norm_backward.out + +- func: miopen_convolution(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic) -> Tensor + dispatch: + CUDA: miopen_convolution + autogen: miopen_convolution.out + +- func: miopen_convolution_transpose(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] output_padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic) -> Tensor + dispatch: + CUDA: miopen_convolution_transpose + autogen: miopen_convolution_transpose.out + +- func: miopen_depthwise_convolution(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic) -> Tensor + dispatch: + CUDA: miopen_depthwise_convolution + autogen: miopen_depthwise_convolution.out + +- func: miopen_convolution_relu(Tensor self, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, SymInt groups) -> Tensor + dispatch: + CUDA: miopen_convolution_relu + +- func: miopen_convolution_add_relu(Tensor self, Tensor weight, Tensor z, Scalar? alpha, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, SymInt groups) -> Tensor + dispatch: + CUDA: miopen_convolution_add_relu + +- func: miopen_rnn(Tensor input, Tensor[] weight, int weight_stride0, Tensor hx, Tensor? cx, int mode, int hidden_size, int num_layers, bool batch_first, float dropout, bool train, bool bidirectional, int[] batch_sizes, Tensor? dropout_state) -> (Tensor, Tensor, Tensor, Tensor, Tensor) + dispatch: + CUDA: miopen_rnn + autogen: miopen_rnn.out + tags: nondeterministic_seeded + + +- func: miopen_rnn_backward(Tensor input, Tensor[] weight, int weight_stride0, Tensor weight_buf, Tensor hx, Tensor? cx, Tensor output, Tensor? grad_output, Tensor? grad_hy, Tensor? grad_cy, int mode, int hidden_size, int num_layers, bool batch_first, float dropout, bool train, bool bidirectional, int[] batch_sizes, Tensor? dropout_state, Tensor reserve, bool[4] output_mask) -> (Tensor, Tensor, Tensor, Tensor[]) + dispatch: + CUDA: miopen_rnn_backward + autogen: miopen_rnn_backward.out + +- func: _use_miopen_ctc_loss(Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, int blank) -> bool + device_check: NoCheck # Tensor arguments allowed to be on different devices, see also miopen_ctc_loss + dispatch: + CUDA: _use_miopen_ctc_loss + +- func: _use_miopen_ctc_loss.Tensor(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank) -> bool + device_check: NoCheck # Tensor arguments allowed to be on different devices, see also miopen_ctc_loss + dispatch: + CUDA: _use_miopen_ctc_loss_tensor + +- func: miopen_ctc_loss(Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, int blank, bool deterministic, bool zero_infinity) -> (Tensor, Tensor) + device_check: NoCheck # log_probs is expected to be on CUDA while targets is expected to be on CPU + dispatch: + CUDA: miopen_ctc_loss + autogen: miopen_ctc_loss.out + +- func: miopen_ctc_loss.Tensor(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank, bool deterministic, bool zero_infinity) -> (Tensor, Tensor) + device_check: NoCheck # log_probs is expected to be on CUDA while targets is expected to be on CPU + dispatch: + CUDA: miopen_ctc_loss_tensor + +- func: mm(Tensor self, Tensor mat2) -> Tensor + structured_delegate: mm.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: _sparse_mm + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: _sparse_csr_mm + tags: core + +- func: mm.out(Tensor self, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: mm_out_cpu + CUDA: mm_out_cuda + MTIA: mm_out_mtia + MPS: mm_out_mps + XPU: mm_out_xpu + SparseCPU, SparseCUDA, SparseMPS: _sparse_mm_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: _sparse_csr_mm_out + +- func: mm.dtype(Tensor self, Tensor mat2, ScalarType out_dtype) -> Tensor + dispatch: + CUDA: _mm_dtype_cuda + XPU: _mm_dtype_xpu + +- func: mm.dtype_out(Tensor self, Tensor mat2, ScalarType out_dtype, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CUDA: _mm_dtype_out_cuda + XPU: _mm_dtype_out_xpu + +- func: _int_mm(Tensor self, Tensor mat2) -> Tensor + dispatch: + CPU: _int_mm_cpu + CUDA: _int_mm_cuda + XPU: _int_mm_xpu + +- func: _int_mm.out(Tensor self, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: _int_mm_out_cpu + CUDA: _int_mm_out_cuda + XPU: _int_mm_out_xpu + +- func: _convert_weight_to_int4pack(Tensor self, int innerKTiles) -> Tensor + dispatch: + CUDA: _convert_weight_to_int4pack_cuda + MPS: _convert_weight_to_int4pack_mps + +- func: _weight_int4pack_mm(Tensor self, Tensor mat2, int qGroupSize, Tensor qScaleAndZeros) -> Tensor + dispatch: + MPS: _weight_int4pack_mm_mps + CUDA: _weight_int4pack_mm_cuda + +- func: _weight_int4pack_mm_with_scales_and_zeros(Tensor self, Tensor mat2, int qGroupSize, Tensor qScale, Tensor qZeros) -> Tensor + dispatch: + XPU: _weight_int4pack_mm_xpu + +# Split int4 pack weight between cpu and other devices due to +# https://github.com/pytorch/ao/issues/1117#issuecomment-2451252756. +- func: _convert_weight_to_int4pack_for_cpu(Tensor self, int innerKTiles) -> Tensor + dispatch: + CPU: _convert_weight_to_int4pack_cpu + +- func: _weight_int4pack_mm_for_cpu(Tensor self, Tensor mat2, int qGroupSize, Tensor qScaleAndZeros) -> Tensor + dispatch: + CPU: _weight_int4pack_mm_cpu + +- func: _dyn_quant_pack_4bit_weight(Tensor weights, Tensor scales_zeros, Tensor? bias, int block_size, int in_features, int out_features) -> Tensor + dispatch: + CPU: _dyn_quant_pack_4bit_weight_cpu + +- func: _dyn_quant_matmul_4bit(Tensor inp, Tensor packed_weights, int block_size, int in_features, int out_features) -> Tensor + dispatch: + CPU: _dyn_quant_matmul_4bit_cpu + +- func: _weight_int8pack_mm(Tensor self, Tensor mat2, Tensor scales) -> Tensor + dispatch: + CPU: _weight_int8pack_mm_cpu + CUDA: _weight_int8pack_mm_cuda + MPS: _weight_int8pack_mm_mps + XPU: _weight_int8pack_mm_xpu + +- func: _sparse_mm(Tensor sparse, Tensor dense) -> Tensor + python_module: sparse + +- func: _sparse_mm.reduce(Tensor sparse, Tensor dense, str reduce) -> Tensor + python_module: sparse + +- func: _sparse_sparse_matmul(Tensor self, Tensor other) -> Tensor + dispatch: + SparseCPU: sparse_sparse_matmul_cpu + SparseCUDA: sparse_sparse_matmul_cuda + SparseMPS: sparse_sparse_matmul_mps + autogen: _sparse_sparse_matmul.out + +- func: mode(Tensor self, int dim=-1, bool keepdim=False) -> (Tensor values, Tensor indices) + variants: function, method + dispatch: + CPU, CUDA: mode + +- func: mode.values(Tensor self, int dim=-1, bool keepdim=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + dispatch: + CompositeExplicitAutograd: mode_out + +- func: mode.dimname(Tensor self, Dimname dim, bool keepdim=False) -> (Tensor values, Tensor indices) + variants: function, method + +- func: mode.dimname_out(Tensor self, Dimname dim, bool keepdim=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + +- func: mul.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: mul.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: mul_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: mul_sparse_csr + MkldnnCPU: mkldnn_mul + ZeroTensor: mul_zerotensor + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_mul_Tensor + tags: [core, pointwise] + +- func: mul_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: mul.out + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: mul_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: mul_sparse_csr_ + MkldnnCPU: mkldnn_mul_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_mul__Tensor + tags: pointwise + +- func: mul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: mul_out + SparseCPU: mul_out_sparse_cpu + SparseCUDA: mul_out_sparse_cuda + SparseMPS: mul_out_sparse_mps + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: mul_out_sparse_csr + MkldnnCPU: mkldnn_mul_out + tags: pointwise + # For C++ only, until we have conversion from C++ numbers to Tensor + +- func: mul.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: mul + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: mul_scalar_sparse_csr + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_mul_Scalar + tags: [core, pointwise] + +- func: mul_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: mul_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: mul__scalar_sparse_csr + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_mul__Scalar + autogen: mul.Scalar_out + tags: pointwise +# multiply, alias for mul + +- func: multiply.Tensor(Tensor self, Tensor other) -> Tensor + variants: function, method + +- func: multiply_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + variants: method + +- func: multiply.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + +- func: multiply.Scalar(Tensor self, Scalar other) -> Tensor + variants: function, method + +- func: multiply_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + variants: method + +- func: mv(Tensor self, Tensor vec) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: mv + SparseCPU, SparseCUDA, SparseMPS: mv_sparse + +- func: mv.out(Tensor self, Tensor vec, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: mv_out + +- func: mvlgamma.out(Tensor self, int p, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: mvlgamma_out + tags: pointwise + +- func: mvlgamma(Tensor self, int p) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: mvlgamma + tags: pointwise + +- func: mvlgamma_(Tensor(a!) self, int p) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: mvlgamma_ + tags: pointwise + +- func: narrow_copy(Tensor self, int dim, SymInt start, SymInt length) -> Tensor + variants: function, method + dispatch: + CPU: narrow_copy_dense_cpu + SparseCPU, SparseCUDA, SparseMPS: narrow_copy_sparse + CompositeExplicitAutogradNonFunctional: narrow_copy_dense_symint + tags: view_copy + +- func: narrow_copy.out(Tensor self, int dim, SymInt start, SymInt length, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: narrow_copy_dense_cpu_out + +- func: narrow(Tensor(a) self, int dim, SymInt start, SymInt length) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: narrow_symint + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: narrow_nested_symint + +- func: narrow.Tensor(Tensor(a) self, int dim, Tensor start, SymInt length) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: narrow_tensor_symint + +- func: native_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor) + dispatch: + CPU: batch_norm_cpu + CUDA: batch_norm_cuda + MPS: batch_norm_mps + MkldnnCPU: mkldnn_batch_norm + +- func: native_batch_norm.out(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps, *, Tensor(a!) out, Tensor(b!) save_mean, Tensor(c!) save_invstd) -> (Tensor(a!), Tensor(b!), Tensor(c!)) + dispatch: + CUDA: batch_norm_cuda_out + MPS: batch_norm_mps_out + CPU: batch_norm_cpu_out + +# TODO: In 2 weeks, we should make native_batch_norm composite implicit so that this correct schema percolates correctly through our dispatching +- func: _native_batch_norm_legit(Tensor input, Tensor? weight, Tensor? bias, Tensor(a!) running_mean, Tensor(b!) running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor) + dispatch: + CPU: _batch_norm_legit_cpu + CUDA: _batch_norm_legit_cuda + MPS: _batch_norm_legit_mps + MkldnnCPU: _mkldnn_batch_norm_legit + autogen: _native_batch_norm_legit_functional + tags: core + +# HACK: identical to _native_batch_norm_legit, but training is known to be False, +# So we known that running stats will not be mutated. +# The real fix here is batch norm consolidation. +- func: _native_batch_norm_legit_no_training(Tensor input, Tensor? weight, Tensor? bias, Tensor running_mean, Tensor running_var, float momentum, float eps) -> (Tensor, Tensor, Tensor) + dispatch: + CompositeExplicitAutograd: _batch_norm_legit_no_training + autogen: _native_batch_norm_legit_no_training.out + tags: core + +- func: _native_batch_norm_legit.out(Tensor input, Tensor? weight, Tensor? bias, Tensor(a!) running_mean, Tensor(b!) running_var, bool training, float momentum, float eps, *, Tensor(d!) out, Tensor(e!) save_mean, Tensor(f!) save_invstd) -> (Tensor(d!), Tensor(e!), Tensor(f!)) + dispatch: + CPU: _batch_norm_legit_cpu_out + CUDA: _batch_norm_legit_cuda_out + MPS: _batch_norm_legit_mps_out + +- func: _native_batch_norm_legit.no_stats(Tensor input, Tensor? weight, Tensor? bias, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor) + dispatch: + CPU: _batch_norm_legit_no_stats_cpu + CUDA: _batch_norm_legit_no_stats_cuda + MPS: _batch_norm_legit_no_stats_mps + MkldnnCPU: _mkldnn_batch_norm_legit_no_stats + tags: core + +- func: _native_batch_norm_legit.no_stats_out(Tensor input, Tensor? weight, Tensor? bias, bool training, float momentum, float eps, *, Tensor(a!) out, Tensor(b!) save_mean, Tensor(c!) save_invstd) -> (Tensor(a!), Tensor(b!), Tensor(c!)) + dispatch: + CPU: _batch_norm_legit_no_stats_cpu_out + CUDA: _batch_norm_legit_no_stats_cuda_out + MPS: _batch_norm_legit_no_stats_mps_out + +- func: batch_norm_stats(Tensor input, float eps) -> (Tensor, Tensor) + dispatch: + CUDA: batch_norm_stats_cuda + autogen: batch_norm_stats.out + +- func: batch_norm_elemt(Tensor input, Tensor? weight, Tensor? bias, Tensor mean, Tensor invstd, float eps) -> Tensor + dispatch: + CUDA: batch_norm_elemt_cuda + +- func: batch_norm_elemt.out(Tensor input, Tensor? weight, Tensor? bias, Tensor mean, Tensor invstd, float eps, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CUDA: batch_norm_elemt_cuda_out + +# for backward compatibility +- func: batch_norm_gather_stats(Tensor input, Tensor mean, Tensor invstd, Tensor? running_mean, Tensor? running_var, float momentum, float eps, int count) -> (Tensor, Tensor) + dispatch: + CUDA: batch_norm_gather_stats_cuda + autogen: batch_norm_gather_stats.out + +- func: batch_norm_gather_stats_with_counts(Tensor input, Tensor mean, Tensor invstd, Tensor? running_mean, Tensor? running_var, float momentum, float eps, Tensor counts) -> (Tensor, Tensor) + dispatch: + CUDA: batch_norm_gather_stats_with_counts_cuda + autogen: batch_norm_gather_stats_with_counts.out + +- func: native_batch_norm_backward(Tensor grad_out, Tensor input, Tensor? weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_invstd, bool train, float eps, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + dispatch: + CPU: batch_norm_backward_cpu + CUDA: batch_norm_backward_cuda + MPS: batch_norm_backward_mps + MkldnnCPU: mkldnn_batch_norm_backward + autogen: native_batch_norm_backward.out + +- func: batch_norm_backward_reduce(Tensor grad_out, Tensor input, Tensor mean, Tensor invstd, Tensor? weight, bool input_g, bool weight_g, bool bias_g) -> (Tensor, Tensor, Tensor, Tensor) + dispatch: + CUDA: batch_norm_backward_reduce_cuda + autogen: batch_norm_backward_reduce.out + +- func: batch_norm_backward_elemt(Tensor grad_out, Tensor input, Tensor mean, Tensor invstd, Tensor? weight, Tensor sum_dy, Tensor sum_dy_xmu, Tensor count) -> Tensor + dispatch: + CUDA: batch_norm_backward_elemt_cuda + autogen: batch_norm_backward_elemt.out + +- func: batch_norm_update_stats(Tensor input, Tensor? running_mean, Tensor? running_var, float momentum) -> (Tensor, Tensor) + dispatch: + CPU: batch_norm_update_stats_cpu + CUDA: batch_norm_update_stats_cuda + autogen: batch_norm_update_stats.out + +- func: is_vulkan_available() -> bool + +- func: _nnpack_available() -> bool + +- func: _nnpack_spatial_convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[2] padding, SymInt[2] stride=1) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: _nnpack_spatial_convolution + autogen: _nnpack_spatial_convolution.out + +- func: ones.names(int[] size, *, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: ones + autogen: ones.names_out + +- func: ones(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: ones + +- func: ones.out(SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: ones_out + +- func: ones_like(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: ones_like + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: ones_like + autogen: ones_like.out + +- func: pairwise_distance(Tensor x1, Tensor x2, float p=2, float eps=1e-06, bool keepdim=False) -> Tensor + +- func: cdist(Tensor x1, Tensor x2, float p=2, int? compute_mode=None) -> Tensor + +- func: _euclidean_dist(Tensor x1, Tensor x2) -> Tensor + dispatch: + CompositeExplicitAutograd: _euclidean_dist + autogen: _euclidean_dist.out + +- func: _cdist_forward(Tensor x1, Tensor x2, float p, int? compute_mode) -> Tensor + dispatch: + CPU, CUDA: _cdist_forward + MTIA: _cdist_forward_mtia + MPS: _cdist_forward_mps + autogen: _cdist_forward.out + tags: core + +- func: _cdist_backward(Tensor grad, Tensor x1, Tensor x2, float p, Tensor cdist) -> Tensor + dispatch: + CPU, CUDA: _cdist_backward + autogen: _cdist_backward.out + +- func: pdist(Tensor self, float p=2) -> Tensor + +- func: _pdist_forward(Tensor self, float p=2) -> Tensor + dispatch: + CPU, CUDA: _pdist_forward + autogen: _pdist_forward.out + tags: core + +- func: _pdist_backward(Tensor grad, Tensor self, float p, Tensor pdist) -> Tensor + dispatch: + CPU, CUDA: _pdist_backward + autogen: _pdist_backward.out + +- func: cosine_similarity(Tensor x1, Tensor x2, int dim=1, float eps=1e-08) -> Tensor + variants: function + +- func: permute(Tensor(a) self, int[] dims) -> Tensor(a) + variants: function, method + dispatch: + CompositeExplicitAutograd: permute + MPS: permute_mps + SparseCPU, SparseCUDA, SparseMPS: permute_sparse_coo + tags: core + +- func: movedim.intlist(Tensor(a) self, int[] source, int[] destination) -> Tensor(a) + variants: function, method + +- func: movedim.int(Tensor(a) self, int source, int destination) -> Tensor(a) + variants: function, method + +# moveaxis, alias for movedim +- func: moveaxis.intlist(Tensor(a) self, int[] source, int[] destination) -> Tensor(a) + variants: function, method + +- func: moveaxis.int(Tensor(a) self, int source, int destination) -> Tensor(a) + variants: function, method + +# Only exposed from C++ -- in Python, +# we expose it as an attribute `T`, not a function. +# +# I'd like to name this "T" in C++ too, but +# calling a native function "T" causes undefined +# behavior on Windows, for reasons I don't understand +# (maybe related to capital letter collation somehow...) +- func: numpy_T(Tensor(a) self) -> Tensor(a) + variants: method + +# Exposed on Python as an attribute 'H' +- func: matrix_H(Tensor(a) self) -> Tensor(a) + variants: method + +# Exposed on Python as an attribute 'mT' +- func: mT(Tensor(a) self) -> Tensor(a) + variants: method + +# Exposed on Python as an attribute 'mH' +- func: mH(Tensor(a) self) -> Tensor(a) + variants: method + +- func: adjoint(Tensor(a) self) -> Tensor(a) + variants: function, method + +- func: pixel_shuffle(Tensor self, int upscale_factor) -> Tensor + dispatch: + CPU: pixel_shuffle_cpu + MPS: pixel_shuffle_mps + CompositeExplicitAutogradNonFunctional: math_pixel_shuffle + autogen: pixel_shuffle.out + +- func: pixel_unshuffle(Tensor self, int downscale_factor) -> Tensor + dispatch: + CPU: pixel_unshuffle_cpu + MPS: pixel_unshuffle_mps + CompositeExplicitAutogradNonFunctional: math_pixel_unshuffle + autogen: pixel_unshuffle.out + +- func: channel_shuffle(Tensor self, SymInt groups) -> Tensor + dispatch: + CPU, CUDA: channel_shuffle + QuantizedCPU: channel_shuffle_quantized_cpu + autogen: channel_shuffle.out + +- func: native_channel_shuffle(Tensor self, SymInt groups) -> Tensor + dispatch: + CPU: channel_shuffle_cpu + CompositeImplicitAutograd: math_channel_shuffle + +- func: is_pinned(Tensor self, Device? device=None) -> bool + variants: method + dispatch: + # the NestedTensor keys are necessary because NestedTensor has been removed + # from the CompositeExplicitAutograd keyset see Note [NestedTensor Not Included in Backend Keys] + CompositeExplicitAutograd, NestedTensorCPU: is_pinned + SparseCsrCPU: is_pinned_sparse_compressed + SparseCPU: is_pinned_sparse_coo + +# TODO: add a copy kwarg that guarantees that the tensor is put into fresh +# pinned memory +- func: pin_memory(Tensor(a) self, Device? device=None) -> Tensor(a) + variants: method + +# Unlike pin_memory, this is guaranteed to give a new non-aliasing tensor +- func: _pin_memory(Tensor self, Device? device=None) -> Tensor + dispatch: + CompositeExplicitAutograd: _pin_memory + NestedTensorCPU: _pin_memory_nested + SparseCPU: _pin_memory_sparse_coo + SparseCsrCPU: _pin_memory_sparse_compressed + autogen: _pin_memory.out + +- func: pinverse(Tensor self, float rcond=1e-15) -> Tensor + variants: function, method + +- func: poisson_nll_loss(Tensor input, Tensor target, bool log_input, bool full, float eps, int reduction) -> Tensor + variants: function + +- func: rad2deg(Tensor self) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: rad2deg + SparseCPU, SparseCUDA, SparseMPS: rad2deg_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: rad2deg_sparse_csr + tags: pointwise + +- func: rad2deg_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + dispatch: + CompositeExplicitAutograd: rad2deg_ + SparseCPU, SparseCUDA, SparseMPS: rad2deg_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: rad2deg_sparse_csr_ + tags: pointwise + +- func: rad2deg.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: rad2deg_out + SparseCPU, SparseCUDA, SparseMPS: rad2deg_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: rad2deg_sparse_csr_out + tags: pointwise + +- func: deg2rad(Tensor self) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: deg2rad + SparseCPU, SparseCUDA, SparseMPS: deg2rad_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: deg2rad_sparse_csr + tags: pointwise + +- func: deg2rad_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + dispatch: + CompositeExplicitAutograd: deg2rad_ + SparseCPU, SparseCUDA, SparseMPS: deg2rad_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: deg2rad_sparse_csr_ + tags: pointwise + +- func: deg2rad.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: deg2rad_out + SparseCPU, SparseCUDA, SparseMPS: deg2rad_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: deg2rad_sparse_csr_out + tags: pointwise + +- func: scalar_tensor(Scalar s, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: scalar_tensor + autogen: scalar_tensor.out + tags: core + +- func: rand.names(SymInt[] size, *, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: rand + autogen: rand.names_out + tags: nondeterministic_seeded + +- func: rand.generator_with_names(SymInt[] size, *, Generator? generator, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + device_check: NoCheck + device_guard: False + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: rand + autogen: rand.generator_with_names_out + +- func: rand(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: [core, nondeterministic_seeded] + dispatch: + CompositeExplicitAutograd: rand + +- func: rand.generator(SymInt[] size, *, Generator? generator, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: rand + +- func: rand.out(SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: rand_out + +- func: rand.generator_out(SymInt[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + +- func: rand_like(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: rand_like + autogen: rand_like.out + +- func: rand_like.generator(Tensor self, *, Generator? generator, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: rand_like + autogen: rand_like.generator_out + +- func: randint(SymInt high, SymInt[] size, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randint + +- func: randint.generator(SymInt high, SymInt[] size, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randint + +- func: randint.low(SymInt low, SymInt high, SymInt[] size, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randint + +- func: randint.low_generator(SymInt low, SymInt high, SymInt[] size, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randint + +- func: randint.out(SymInt high, SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randint_out + +- func: randint.generator_out(SymInt high, SymInt[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randint_out + +- func: randint.low_out(SymInt low, SymInt high, SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randint_out + +- func: randint.low_generator_out(SymInt low, SymInt high, SymInt[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randint_out + +- func: randint_like(Tensor self, SymInt high, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: randint_like + autogen: randint_like.out + +- func: randint_like.generator(Tensor self, SymInt high, *, Generator? generator, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: randint_like + autogen: randint_like.generator_out + +- func: randint_like.Tensor(Tensor self, Tensor high, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: randint_like + autogen: randint_like.Tensor_out + +- func: randint_like.Tensor_generator(Tensor self, Tensor high, *, Generator? generator, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: randint_like + autogen: randint_like.Tensor_generator_out + +- func: randint_like.low_dtype(Tensor self, SymInt low, SymInt high, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: randint_like + autogen: randint_like.low_dtype_out + +- func: randint_like.low_generator_dtype(Tensor self, SymInt low, SymInt high, *, Generator? generator, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd: randint_like + autogen: randint_like.low_generator_dtype_out + +- func: randn(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: [core, nondeterministic_seeded] + dispatch: + CompositeExplicitAutograd: randn + +- func: randn.generator(SymInt[] size, *, Generator? generator, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randn + +- func: randn.names(SymInt[] size, *, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: nondeterministic_seeded + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: randn + autogen: randn.names_out + +- func: randn.generator_with_names(SymInt[] size, *, Generator? generator, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: nondeterministic_seeded + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: randn + autogen: randn.generator_with_names_out + +- func: randn.out(SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + +- func: randn.generator_out(SymInt[] size, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + +- func: randn_like(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd, CompositeImplicitAutogradNestedTensor: randn_like + autogen: randn_like.out + +- func: randn_like.generator(Tensor self, *, Generator? generator, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd, CompositeImplicitAutogradNestedTensor: randn_like + autogen: randn_like.generator_out + +- func: randperm(SymInt n, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: [core, nondeterministic_seeded] + dispatch: + CompositeExplicitAutograd: randperm + +- func: randperm.generator(SymInt n, *, Generator? generator, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randperm + +- func: randperm.out(SymInt n, *, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: randperm_out + +- func: randperm.generator_out(SymInt n, *, Generator? generator, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + dispatch: + CPU: randperm_out_cpu + CUDA: randperm_out_cuda + MPS: randperm_out_mps + +- func: range.step(Scalar start, Scalar end, Scalar step=1, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: range + +- func: range(Scalar start, Scalar end, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: range + +- func: range.out_(Scalar start, Scalar end, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: range_out_no_step + +- func: range.out(Scalar start, Scalar end, Scalar step=1, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, Meta: range_out + CUDA: range_cuda_out + MPS: range_mps_out + cpp_no_default_args: ['step'] + +- func: ravel(Tensor(a) self) -> Tensor(a) + variants: function, method + +- func: reciprocal(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: reciprocal.out + variants: function, method + tags: [core, pointwise] + +- func: reciprocal_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: reciprocal.out + variants: function, method + tags: pointwise + +- func: reciprocal.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: reciprocal_out + tags: pointwise + +- func: neg(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: neg.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: neg_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: neg_sparse_csr + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_neg + tags: [core, pointwise] + +- func: neg_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: neg.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: neg_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: neg_sparse_csr_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_neg_ + tags: pointwise + +- func: neg.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: neg_out + SparseCPU, SparseCUDA, SparseMPS: neg_out_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: neg_sparse_csr_out + tags: pointwise +# Alias for neg + +- func: negative(Tensor self) -> Tensor + variants: function, method + +- func: negative_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + +- func: negative.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: repeat(Tensor self, SymInt[] repeats) -> Tensor + variants: method # This is method-only to match the previous tensor API. In the future we could make this a function too. + dispatch: + CompositeExplicitAutograd: repeat + MPS: repeat_mps + autogen: repeat.out + tags: core + +- func: repeat_interleave.Tensor(Tensor repeats, *, SymInt? output_size=None) -> Tensor + variants: function + dispatch: + CPU: repeat_interleave_cpu + CUDA: repeat_interleave_cuda + MPS: repeat_interleave_mps + tags: dynamic_output_shape + autogen: repeat_interleave.Tensor_out + +- func: repeat_interleave.self_Tensor(Tensor self, Tensor repeats, int? dim=None, *, SymInt? output_size=None) -> Tensor + variants: function, method + dispatch: + CompositeImplicitAutograd: repeat_interleave_symint + +- func: repeat_interleave.self_int(Tensor self, SymInt repeats, int? dim=None, *, SymInt? output_size=None) -> Tensor + variants: function, method + dispatch: + CompositeImplicitAutograd: repeat_interleave_symint + +- func: reshape(Tensor(a) self, SymInt[] shape) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: reshape_symint + CompositeImplicitAutogradNestedTensor: reshape_nested_symint + +- func: _reshape_copy(Tensor self, SymInt[] size) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: _reshape_copy_symint + +# NOTE [ _reshape_alias ] is meant to be used in the implementation of reshape. +# They are not user-facing, hence the leading underscore. Please don't use it +# anywhere else. +- func: _reshape_alias(Tensor(a) self, SymInt[] size, SymInt[] stride) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CPU, CUDA, Meta, QuantizedCPU, QuantizedCUDA, ZeroTensor, MPS, MTIA: _reshape_alias + # We don't need to support mkldnn since this is handled explicitly by the reshape operator. + +- func: _mkldnn_reshape(Tensor self, int[] shape) -> Tensor + device_check: NoCheck + device_guard: False + dispatch: + MkldnnCPU: mkldnn_reshape + autogen: _mkldnn_reshape.out + +- func: reshape_as(Tensor(a) self, Tensor other) -> Tensor(a) + variants: method + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: reshape_as + CompositeImplicitAutogradNestedTensor: reshape_as_nested + +- func: round(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: round.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: round_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: round_sparse_csr + tags: [core, pointwise] + +- func: round_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: round.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: round_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: round_sparse_csr_ + tags: pointwise + +- func: round.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: round_out + SparseCPU, SparseCUDA, SparseMPS: round_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: round_sparse_csr_out + tags: pointwise + +- func: round.decimals(Tensor self, *, int decimals) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: round.decimals_out + variants: function, method + tags: pointwise + +- func: round_.decimals(Tensor(a!) self, *, int decimals) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: round.decimals_out + variants: function, method + tags: pointwise + +- func: round.decimals_out(Tensor self, *, int decimals, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: round_decimals_out + tags: pointwise + +- func: rrelu(Tensor self, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor + device_check: NoCheck # TensorIterator + tags: [pointwise, nondeterministic_seeded] + +- func: rrelu_(Tensor(a!) self, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor(a!) + tags: nondeterministic_seeded + device_check: NoCheck # TensorIterator + +- func: relu(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA: relu + MPS: relu_mps + MTIA: relu_mtia + MkldnnCPU: mkldnn_relu + QuantizedCPU: relu_quantized_cpu + QuantizedCUDA: relu_quantized_cuda + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_relu + SparseCPU, SparseCUDA, SparseMPS: relu_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: relu_sparse_csr + tags: [core, pointwise] + +- func: relu_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA: relu_ + MPS: relu_mps_ + MTIA: relu_mtia_ + MkldnnCPU: mkldnn_relu_ + QuantizedCPU: relu_quantized_cpu_ + QuantizedCUDA: relu_quantized_cuda_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_relu_ + SparseCPU, SparseCUDA, SparseMPS: relu_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: relu_sparse_csr_ + autogen: relu.out + tags: pointwise + +- func: relu6(Tensor self) -> Tensor + python_module: nn + tags: pointwise + +- func: relu6_(Tensor(a!) self) -> Tensor(a!) + python_module: nn + +- func: prelu(Tensor self, Tensor weight) -> Tensor + variants: function, method + autogen: prelu.out + +- func: _prelu_kernel(Tensor self, Tensor weight) -> Tensor + dispatch: + CPU, CUDA: _prelu_kernel + QuantizedCPU: _prelu_kernel_quantized_cpu + MkldnnCPU: mkldnn_prelu + MPS: prelu_mps + +- func: _prelu_kernel_backward(Tensor grad_output, Tensor self, Tensor weight) -> (Tensor, Tensor) + dispatch: + CPU, CUDA: _prelu_kernel_backward + MkldnnCPU: mkldnn_prelu_backward + MPS: prelu_backward_mps + +- func: gelu.out(Tensor self, *, str approximate='none', Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU: gelu_out_cpu + CUDA: gelu_out_cuda + MPS: gelu_out_mps + +- func: gelu_(Tensor(a!) self, *, str approximate='none') -> Tensor(a!) + structured_delegate: gelu.out + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + QuantizedCPU: gelu_quantized_cpu_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_gelu_ + +- func: gelu(Tensor self, *, str approximate='none') -> Tensor + structured_delegate: gelu.out + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + MkldnnCPU: mkldnn_gelu + QuantizedCPU: gelu_quantized_cpu + QuantizedCUDA: gelu_quantized_cuda + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_gelu + tags: [core, pointwise] + +- func: gelu_backward.grad_input(Tensor grad_output, Tensor self, *, str approximate='none', Tensor(a!) grad_input) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU: gelu_backward_out_cpu + CUDA: gelu_backward_out_cuda + MPS: gelu_backward_out_mps + +- func: gelu_backward(Tensor grad_output, Tensor self, *, str approximate='none') -> Tensor + structured_delegate: gelu_backward.grad_input + python_module: nn + dispatch: + MkldnnCPU: mkldnn_gelu_backward + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: gelu_backwards_nested + tags: pointwise + +- func: infinitely_differentiable_gelu_backward(Tensor grad, Tensor self) -> Tensor + variants: function + python_module: nn + device_check: NoCheck + device_guard: False + +- func: hardshrink.out(Tensor self, Scalar lambd=0.5, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS: hardshrink_out + +- func: hardshrink(Tensor self, Scalar lambd=0.5) -> Tensor + structured_delegate: hardshrink.out + device_check: NoCheck # TensorIterator + variants: function, method + tags: pointwise + +- func: hardshrink_backward.grad_input(Tensor grad_out, Tensor self, Scalar lambd, *, Tensor(a!) grad_input) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: hardshrink_backward_out + +- func: hardshrink_backward(Tensor grad_out, Tensor self, Scalar lambd) -> Tensor + structured_delegate: hardshrink_backward.grad_input + variants: function, method + +- func: rsqrt(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: rsqrt.out + variants: function, method + tags: [core, pointwise] + +- func: rsqrt_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: rsqrt.out + variants: function, method + tags: pointwise + +- func: rsqrt.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: rsqrt_out + tags: pointwise + +- func: select.Dimname(Tensor(a) self, Dimname dim, int index) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + +- func: select.int(Tensor(a) self, int dim, SymInt index) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: select_symint + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: select_sparse_csr + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: select_nested + tags: core + +- func: select_backward(Tensor grad_output, SymInt[] input_sizes, int dim, SymInt index) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutogradNonFunctional: select_backward_symint + autogen: select_backward.out + +- func: _nested_select_backward(Tensor grad_output, Tensor self, int dim, SymInt index) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: _nested_select_backward_symint + +- func: selu(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + tags: pointwise + +- func: selu_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: pointwise + +- func: celu(Tensor self, Scalar alpha=1.0) -> Tensor + device_check: NoCheck # TensorIterator + dispatch: + CompositeExplicitAutograd: celu + tags: pointwise + +- func: celu_(Tensor(a!) self, Scalar alpha=1.0) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CompositeExplicitAutograd: celu_ + autogen: celu.out + tags: pointwise + +- func: silu(Tensor self) -> Tensor + structured_delegate: silu.out + python_module: nn + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_silu + tags: pointwise + +- func: silu_(Tensor(a!) self) -> Tensor(a!) + structured_delegate: silu.out + python_module: nn + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_silu_ + tags: pointwise + +- func: silu.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA, MPS, MTIA: silu_out + tags: pointwise + +- func: silu_backward.grad_input(Tensor grad_output, Tensor self, *, Tensor(a!) grad_input) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA, MPS: silu_backward_out + tags: pointwise + +- func: silu_backward(Tensor grad_output, Tensor self) -> Tensor + structured_delegate: silu_backward.grad_input + python_module: nn + dispatch: + CompositeImplicitAutograd: math_silu_backward + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: silu_backward_nested + tags: pointwise + +- func: mish(Tensor self) -> Tensor + structured_delegate: mish.out + python_module: nn + tags: pointwise + +- func: mish_(Tensor(a!) self) -> Tensor(a!) + structured_delegate: mish.out + python_module: nn + +- func: mish.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA: mish_out + MPS: mish_out_mps + +- func: mish_backward(Tensor grad_output, Tensor self) -> Tensor + python_module: nn + dispatch: + CPU, CUDA: mish_backward + MPS: mish_backward_mps + CompositeImplicitAutograd: math_mish_backward + +- func: sigmoid(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: sigmoid.out + variants: function, method + dispatch: + QuantizedCPU: sigmoid_quantized_cpu + MkldnnCPU: mkldnn_sigmoid + tags: [core, pointwise] + +- func: sigmoid_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: sigmoid.out + variants: function, method + dispatch: + MkldnnCPU: mkldnn_sigmoid_ + tags: pointwise + +- func: sigmoid.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: sigmoid_out + tags: pointwise + +- func: logit(Tensor self, float? eps=None) -> Tensor + variants: function, method + dispatch: + CPU, CUDA, MTIA: logit + MPS: logit_mps + tags: pointwise + +- func: logit_(Tensor(a!) self, float? eps=None) -> Tensor(a!) + variants: function, method + dispatch: + CPU, CUDA: logit_ + tags: pointwise + +- func: logit.out(Tensor self, float? eps=None, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: logit_out + MPS: logit_out_mps + tags: pointwise + +- func: sin(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: sin.out + variants: function, method + dispatch: + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sin_sparse_csr + SparseCPU, SparseCUDA, SparseMPS: sin_sparse + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_sin + tags: [core, pointwise] + +- func: sin_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: sin.out + variants: function, method + dispatch: + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sin_sparse_csr_ + SparseCPU, SparseCUDA, SparseMPS: sin_sparse_ + tags: pointwise + +- func: sin.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: sin_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sin_sparse_csr_out + SparseCPU, SparseCUDA, SparseMPS: sin_sparse_out + tags: pointwise + +- func: sinc(Tensor self) -> Tensor + structured_delegate: sinc.out + variants: function, method + tags: pointwise + +- func: sinc_(Tensor(a!) self) -> Tensor(a!) + structured_delegate: sinc.out + variants: function, method + tags: pointwise + +- func: sinc.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: sinc_out + tags: pointwise + +- func: sinh(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: sinh.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sinh_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sinh_sparse_csr + tags: [core, pointwise] + +- func: sinh_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: sinh.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sinh_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sinh_sparse_csr_ + tags: pointwise + +- func: sinh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: sinh_out + SparseCPU, SparseCUDA, SparseMPS: sinh_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sinh_sparse_csr_out + +# Returns a copy of this `Variable` that is detached from its autograd graph. +# This method is OK to call if the `Variable` is a view. +# +# NOTE: Previously, if we change the tensor metadata (e.g. sizes / strides / +# storage / storage_offset) of a tensor created from `detach()`, those metadata +# in the original tensor will also be updated. However, the new behavior is that +# those metadata changes to the detached tensor will not update the original tensor +# anymore, and in the `detach()` function we need to set `allow_tensor_metadata_change_` +# to false to make such changes explicitly illegal, in order to prevent users from +# changing metadata of the detached tensor and expecting the original tensor to also +# be updated. + tags: pointwise +- func: detach(Tensor(a) self) -> Tensor(a) + variants: function, method + dispatch: + CompositeExplicitAutograd: detach + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: detach + +# Like `detach()`, but modifies this `Variable` in-place. This method may +# only be called on non-view `Variable`s. You can use `is_view()` to check +# this. If this `Variable` is a view, throws an `std::runtime_error()`. +- func: detach_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + tags: inplace_view + dispatch: + CompositeExplicitAutograd: detach_ + +- func: size.int(Tensor self, int dim) -> int + variants: function + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: size.Dimname(Tensor self, Dimname dim) -> int + variants: function, method + device_check: NoCheck + device_guard: False + +- func: sym_size.int(Tensor self, int dim) -> SymInt + variants: function + device_check: NoCheck + device_guard: False + tags: core + manual_cpp_binding: True + +- func: sym_is_contiguous(Tensor self, MemoryFormat memory_format=contiguous_format) -> SymBool + variants: function + device_check: NoCheck + device_guard: False + tags: core + manual_cpp_binding: True + +- func: sym_numel(Tensor self) -> SymInt + variants: function + device_check: NoCheck + device_guard: False + tags: core + manual_cpp_binding: True + +- func: sym_storage_offset(Tensor self) -> SymInt + variants: function + device_check: NoCheck + device_guard: False + tags: core + manual_cpp_binding: True + +- func: slice.Tensor(Tensor(a) self, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: slice + tags: core + +# NOTE: The implementation of split_with_sizes bypasses the dispatcher to call this; undo +# that if adding specific implementations here! + +- func: slice_backward(Tensor grad_output, SymInt[] input_sizes, int dim, SymInt start, SymInt end, SymInt step) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: slice_backward + autogen: slice_backward.out + +# NB: This op exists to back the implementation of reverse view_funcs for various views (chunk, +# slice.Tensor, split_with_sizes, et al.). Currently, these are only used during fake-ification +# of PT2 graph input subclass instances that are views. This means: +# * This op shouldn't really show up in eager mode (so e.g. XLA shouldn't have to implement it) +# * This op shouldn't show up in a PT2 graph (so a PT2 backend shouldn't have to implement it) +# * A subclass will have to implement this to work in PT2 if a subclass view is used as a graph +# input AND the view utilizes this op in its inverse. The idea is that slice_inverse() is +# easier to implement for a subclass than as_strided() +- func: slice_inverse(Tensor(a) self, Tensor src, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: slice_inverse_symint + +- func: slice_scatter(Tensor self, Tensor src, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutogradNonFunctional: slice_scatter + autogen: slice_scatter.out + tags: [core, view_copy] + +- func: select_scatter(Tensor self, Tensor src, int dim, SymInt index) -> Tensor + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutogradNonFunctional: select_scatter_symint + autogen: select_scatter.out + tags: core + +- func: diagonal_scatter(Tensor self, Tensor src, int offset=0, int dim1=0, int dim2=1) -> Tensor + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutogradNonFunctional: diagonal_scatter + autogen: diagonal_scatter.out + +- func: as_strided_scatter(Tensor self, Tensor src, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutogradNonFunctional: as_strided_scatter_symint + autogen: as_strided_scatter.out + +- func: smm(Tensor self, Tensor mat2) -> Tensor + variants: function, method + +# softmax allows positional dtype, unlike most operators, because kwonly is BC-breaking when loading jit models. +- func: softmax.int(Tensor self, int dim, ScalarType? dtype=None) -> Tensor + variants: function, method + +- func: softmax.int_out(Tensor self, int dim, ScalarType? dtype=None, *, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CompositeExplicitAutograd: softmax_out + +- func: softmax.Dimname(Tensor self, Dimname dim, *, ScalarType? dtype=None) -> Tensor + variants: function, method + +- func: _softmax(Tensor self, int dim, bool half_to_float) -> Tensor + structured_delegate: _softmax.out + dispatch: + MkldnnCPU: mkldnn_softmax + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: softmax_nested + tags: core + +- func: _softmax.out(Tensor self, int dim, bool half_to_float, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: softmax_cpu_out + CUDA: softmax_cuda_out + MPS: softmax_mps_out + +- func: _softmax_backward_data(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype) -> Tensor + structured_delegate: _softmax_backward_data.out + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: nested_softmax_backward + +- func: _softmax_backward_data.out(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype, *, Tensor(a!) grad_input) -> Tensor(a!) + structured: True + dispatch: + CPU: softmax_backward_cpu_out + CUDA: softmax_backward_cuda_out + MPS: softmax_backward_mps_out + +- func: unsafe_split.Tensor(Tensor self, SymInt split_size, int dim=0) -> Tensor[] + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: unsafe_split + autogen: unsafe_split.Tensor_out + +- func: split.Tensor(Tensor(a -> *) self, SymInt split_size, int dim=0) -> Tensor(a)[] + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: split + +- func: split.sizes(Tensor(a -> *) self, SymInt[] split_size, int dim=0) -> Tensor(a)[] + variants: function, method + device_guard: False + dispatch: + CompositeImplicitAutograd: split_symint + +- func: unsafe_split_with_sizes(Tensor self, SymInt[] split_sizes, int dim=0) -> Tensor[] + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: unsafe_split_with_sizes + autogen: unsafe_split_with_sizes.out + +- func: split_with_sizes(Tensor(a -> *) self, SymInt[] split_sizes, int dim=0) -> Tensor(a)[] + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: split_with_sizes + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: split_with_sizes_nested + tags: core + +- func: hsplit.int(Tensor(a -> *) self, int sections) -> Tensor(a)[] + variants: function, method + +- func: hsplit.array(Tensor(a -> *) self, int[] indices) -> Tensor(a)[] + variants: function, method + +- func: vsplit.int(Tensor(a -> *) self, int sections) -> Tensor(a)[] + variants: function, method + +- func: vsplit.array(Tensor(a -> *) self, int[] indices) -> Tensor(a)[] + variants: function, method + +- func: dsplit.int(Tensor(a -> *) self, int sections) -> Tensor(a)[] + variants: function, method + +- func: dsplit.array(Tensor(a -> *) self, int[] indices) -> Tensor(a)[] + variants: function, method + +- func: squeeze(Tensor(a) self) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: squeeze + QuantizedCPU, QuantizedCUDA: squeeze_quantized + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: squeeze_nested + +- func: squeeze.dim(Tensor(a) self, int dim) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: squeeze + QuantizedCPU, QuantizedCUDA: squeeze_quantized + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: squeeze_dim_nested + tags: core + +- func: squeeze.dimname(Tensor(a) self, Dimname dim) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + + +- func: squeeze.dims(Tensor(a) self, int[] dim) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: squeeze + QuantizedCPU, QuantizedCUDA: squeeze_quantized + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: squeeze_dim_nested + tags: core + +- func: squeeze_(Tensor(a!) self) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + tags: inplace_view + dispatch: + CompositeExplicitAutograd: squeeze_ + +- func: squeeze_.dim(Tensor(a!) self, int dim) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + tags: inplace_view + dispatch: + CompositeExplicitAutograd: squeeze_ + +- func: squeeze_.dims(Tensor(a!) self, int[] dim) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + tags: inplace_view + dispatch: + CompositeExplicitAutograd: squeeze_ + +- func: squeeze_.dimname(Tensor(a!) self, Dimname dim) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + tags: inplace_view + +- func: sspaddmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + variants: function, method + +- func: sspaddmm.out(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: _sspaddmm_out_only_sparse + CUDA: _sspaddmm_out_only_sparse_cuda + SparseCPU: _sspaddmm_out_cpu + SparseCUDA: _sspaddmm_out_cuda + +- func: _chunk_cat(Tensor[] tensors, int dim, int num_chunks) -> Tensor + dispatch: + CompositeExplicitAutograd: _chunk_cat + CUDA: _chunk_cat_cuda + +- func: _chunk_cat.out(Tensor[] tensors, int dim, int num_chunks, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: _chunk_cat_out + CUDA: _chunk_cat_out_cuda + +- func: stack(Tensor[] tensors, int dim=0) -> Tensor + dispatch: + CompositeExplicitAutograd: stack + +- func: stack.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: stack_out + +- func: _stack(Tensor[] tensors, int dim=0) -> Tensor + dispatch: # match the backends supported by _cat + CPU: _stack_cpu + CompositeExplicitAutograd: _stack + +- func: _stack.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!) + dispatch: # match the backends supported by _cat_out + CPU: _stack_out_cpu + CompositeExplicitAutograd: _stack_out + +- func: hstack(Tensor[] tensors) -> Tensor + +- func: hstack.out(Tensor[] tensors, *, Tensor(a!) out) -> Tensor(a!) + +- func: vstack(Tensor[] tensors) -> Tensor + +- func: vstack.out(Tensor[] tensors, *, Tensor(a!) out) -> Tensor(a!) + +- func: dstack(Tensor[] tensors) -> Tensor + +- func: dstack.out(Tensor[] tensors, *, Tensor(a!) out) -> Tensor(a!) + +# Overload without center & pad mode, needed for forward-compatibility +- func: stft(Tensor self, int n_fft, int? hop_length=None, int? win_length=None, Tensor? window=None, bool normalized=False, bool? onesided=None, bool? return_complex=None, bool? align_to_window=None) -> Tensor + variants: function, method + cpp_no_default_args: ['hop_length', 'win_length', 'window', 'normalized'] + +- func: stft.center(Tensor self, int n_fft, int? hop_length=None, int? win_length=None, Tensor? window=None, bool center=True, str pad_mode="reflect", bool normalized=False, bool? onesided=None, bool? return_complex=None, bool? align_to_window=None) -> Tensor + variants: function, method + +- func: istft(Tensor self, int n_fft, int? hop_length=None, int? win_length=None, Tensor? window=None, bool center=True, bool normalized=False, bool? onesided=None, int? length=None, bool return_complex=False) -> Tensor + variants: function, method + +- func: stride.int(Tensor self, int dim) -> int + variants: function + device_check: NoCheck + device_guard: False + manual_cpp_binding: True + +- func: stride.Dimname(Tensor self, Dimname dim) -> int + variants: function, method + device_check: NoCheck + device_guard: False + +- func: sym_stride.int(Tensor self, int dim) -> SymInt + variants: function + device_check: NoCheck + device_guard: False + tags: core + manual_cpp_binding: True + +- func: sum(Tensor self, *, ScalarType? dtype=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: sum + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: sum_coo + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sum_csr + autogen: sum.out + tags: reduction + +- func: sum.dim_IntList(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + # TODO: Align the signature of sum.dim_IntList and _sparse_csr_sum.dim_dtype + structured_delegate: sum.IntList_out + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + NestedTensorCPU: NestedTensor_sum_dim_CPU + SparseCPU, SparseCUDA, SparseMPS: sum_sparse_coo + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sum_sparse_compressed + tags: [core, reduction] + +- func: sum.dim_DimnameList(Tensor self, Dimname[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: sum.IntList_out(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + structured: True + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: sum_out + MPS: sum_out_mps + tags: reduction + +- func: sum.DimnameList_out(Tensor self, Dimname[1] dim, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: reduction + +# TODO: this function will be replaced once nested expand semantics have been settled on +- func: _nested_sum_backward(Tensor grad, Tensor self, int[1]? dim, bool keepdim=False) -> Tensor + dispatch: + NestedTensorCPU: _nested_sum_backward_cpu + +- func: nansum(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + variants: function, method + dispatch: + CPU, CUDA: nansum + MPS: nansum_mps + tags: reduction + +- func: nansum.out(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: nansum_out + MPS: nansum_out_mps + tags: reduction + +- func: hash_tensor(Tensor self, int[1] dim=[], *, bool keepdim=False, int mode=0) -> Tensor + variants: function, method + structured_delegate: hash_tensor.out + +- func: hash_tensor.out(Tensor self, int[1] dim=[], *, bool keepdim=False, int mode=0, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU, CUDA: hash_tensor_out + +- func: sum_to_size(Tensor self, SymInt[] size) -> Tensor + variants: method + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: sum_to_size_symint + +- func: sqrt(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: sqrt.out + variants: function, method + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_sqrt + SparseCPU, SparseCUDA, SparseMPS: sqrt_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sqrt_sparse_csr + tags: [core, pointwise] + +- func: sqrt_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: sqrt.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sqrt_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sqrt_sparse_csr_ + tags: pointwise + +- func: sqrt.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: sqrt_out + SparseCPU, SparseCUDA, SparseMPS: sqrt_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sqrt_sparse_csr_out + tags: pointwise + +- func: square(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: pointwise + +- func: square_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function, method + tags: pointwise + +- func: square.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + tags: pointwise + +- func: std(Tensor self, bool unbiased=True) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: std.dim(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: std.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA: std + MPS: std_mps + QuantizedCPU: std_quantized_cpu + tags: reduction + +- func: std_mean(Tensor self, bool unbiased=True) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: std_mean.dim(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: std_mean.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CPU, CUDA: std_mean + MPS: std_mean_mps + autogen: std_mean.correction_out + tags: reduction + +- func: std_mean.names_dim(Tensor self, Dimname[1] dim, bool unbiased=True, bool keepdim=False) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: std_mean.correction_names(Tensor self, Dimname[1] dim, *, Scalar? correction=None, bool keepdim=False) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + tags: reduction + +- func: std.out(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: std.correction_out(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: std_out + QuantizedCPU: std_out_quantized_cpu + tags: reduction + +- func: std.names_dim(Tensor self, Dimname[1] dim, bool unbiased=True, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: std.names_out(Tensor self, Dimname[1] dim, bool unbiased=True, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: std.correction_names(Tensor self, Dimname[1] dim, *, Scalar? correction=None, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: std.correction_names_out(Tensor self, Dimname[1] dim, *, Scalar? correction=None, bool keepdim=False, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + tags: reduction + +- func: prod(Tensor self, *, ScalarType? dtype=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA: prod + MPS: prod_mps + autogen: prod.out + tags: [core, reduction] + +- func: prod.dim_int(Tensor self, int dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + structured_delegate: prod.int_out + device_check: NoCheck # TensorIterator + variants: function, method + tags: [core, reduction] + +- func: prod.int_out(Tensor self, int dim, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + structured: True + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: prod_out + MPS: prod_out_mps + tags: reduction + +- func: prod.dim_Dimname(Tensor self, Dimname dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: prod.Dimname_out(Tensor self, Dimname dim, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: reduction + +- func: t(Tensor(a) self) -> Tensor(a) + device_check: NoCheck + device_guard: False + variants: function, method + dispatch: + CompositeExplicitAutograd: t + +- func: t_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck + device_guard: False + variants: method + tags: inplace_view + dispatch: + CompositeExplicitAutograd: t_ + +- func: tan(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: tan.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: tan_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: tan_sparse_csr + tags: [core, pointwise] + +- func: tan_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: tan.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: tan_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: tan_sparse_csr_ + tags: pointwise + +- func: tan.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: tan_out + SparseCPU, SparseCUDA, SparseMPS: tan_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: tan_sparse_csr_out + tags: pointwise + +- func: tanh(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: tanh.out + variants: function, method + dispatch: + QuantizedCPU: tanh_quantized_cpu + MkldnnCPU: mkldnn_tanh + SparseCPU, SparseCUDA, SparseMPS: tanh_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: tanh_sparse_csr + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_tanh + tags: [core, pointwise] + +- func: tanh_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: tanh.out + variants: function, method + dispatch: + MkldnnCPU: mkldnn_tanh_ + SparseCPU, SparseCUDA, SparseMPS: tanh_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: tanh_sparse_csr_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_tanh_ + tags: pointwise + +- func: tanh.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: tanh_out + SparseCPU, SparseCUDA, SparseMPS: tanh_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: tanh_sparse_csr_out + tags: pointwise + +- func: tensordot(Tensor self, Tensor other, int[] dims_self, int[] dims_other) -> Tensor + variants: function + +- func: tensordot.out(Tensor self, Tensor other, int[] dims_self, int[] dims_other, *, Tensor(a!) out) -> Tensor(a!) + variants: function + +# TODO: namespace threshold in 'nn' +- func: threshold(Tensor self, Scalar threshold, Scalar value) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + structured_delegate: threshold.out + dispatch: + QuantizedCPU: threshold_quantized_cpu + tags: pointwise + +- func: threshold_(Tensor(a!) self, Scalar threshold, Scalar value) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + structured_delegate: threshold.out + +- func: threshold.out(Tensor self, Scalar threshold, Scalar value, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: threshold_out + MPS: threshold_out_mps + +- func: threshold_backward.grad_input(Tensor grad_output, Tensor self, Scalar threshold, *, Tensor(a!) grad_input) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: threshold_backward_out + MPS: threshold_backward_out_mps + SparseCPU, SparseCUDA: threshold_backward_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: threshold_backward_sparse_compressed_out + +- func: threshold_backward(Tensor grad_output, Tensor self, Scalar threshold) -> Tensor + variants: function + structured_delegate: threshold_backward.grad_input + dispatch: + MkldnnCPU: mkldnn_relu_backward + SparseCPU, SparseCUDA: threshold_backward_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: threshold_backward_sparse_compressed + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: threshold_backwards_nested + tags: pointwise + +- func: tile(Tensor self, SymInt[] dims) -> Tensor + variants: function, method + dispatch: + CompositeImplicitAutograd: tile_symint + +- func: transpose.int(Tensor(a) self, int dim0, int dim1) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: transpose + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: transpose_nested + +- func: transpose.Dimname(Tensor(a) self, Dimname dim0, Dimname dim1) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + +- func: _mkldnn_transpose(Tensor self, int dim0, int dim1) -> Tensor + device_check: NoCheck + device_guard: False + dispatch: + MkldnnCPU: mkldnn_transpose + +- func: transpose_(Tensor(a!) self, int dim0, int dim1) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + tags: inplace_view + dispatch: + CompositeExplicitAutograd: transpose_ + +- func: _mkldnn_transpose_(Tensor(a!) self, int dim0, int dim1) -> Tensor(a!) + device_check: NoCheck + device_guard: False + dispatch: + MkldnnCPU: mkldnn_transpose_ + autogen: _mkldnn_transpose.out + +- func: one_hot(Tensor self, int num_classes=-1) -> Tensor + python_module: nn + variants: function + tags: dynamic_output_shape + +- func: flip(Tensor self, int[] dims) -> Tensor + variants: function, method + dispatch: + CPU, QuantizedCPU, CUDA, QuantizedCUDA: flip + MPS: flip_mps + autogen: flip.out + tags: core + +- func: fliplr(Tensor self) -> Tensor + variants: function, method + +- func: flipud(Tensor self) -> Tensor + variants: function, method + +- func: roll(Tensor self, SymInt[1] shifts, int[1] dims=[]) -> Tensor + variants: function, method + dispatch: + CPU, MPS: roll + CUDA: roll_cuda + autogen: roll.out + +# default int[] value [0,1] should not add space after comma, since codegen parser uses ', ' to split args + +- func: rot90(Tensor self, int k=1, int[] dims=[0,1]) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: rot90 + autogen: rot90.out + +- func: trapezoid.x(Tensor y, Tensor x, *, int dim=-1) -> Tensor + +- func: trapezoid.dx(Tensor y, *, Scalar dx=1, int dim=-1) -> Tensor + +- func: trapz.x(Tensor y, Tensor x, *, int dim=-1) -> Tensor + +- func: trapz.dx(Tensor y, *, float dx=1, int dim=-1) -> Tensor + +# Fused implementation detail for transformers. Adds in-projection bias to QKV and divides Q by sqrt(D/num_heads). +- func: _transform_bias_rescale_qkv(Tensor qkv, Tensor qkv_bias, int num_heads) -> (Tensor, Tensor, Tensor) + dispatch: + CPU, NestedTensorCPU: transform_bias_rescale_qkv_cpu + CUDA, NestedTensorCUDA: transform_bias_rescale_qkv_cuda + autogen: _transform_bias_rescale_qkv.out + +- func: _nested_tensor_from_mask(Tensor t, Tensor mask, bool mask_check=True) -> Tensor + dispatch: + CPU, CUDA: NestedTensor_nested_tensor_from_mask + autogen: _nested_tensor_from_mask.out + +- func: _nested_tensor_from_mask_left_aligned(Tensor t, Tensor mask) -> bool + dispatch: + CPU, CUDA: NestedTensor_nested_tensor_from_mask_left_aligned + +- func: _nested_from_padded(Tensor padded, Tensor cpu_nested_shape_example, bool fuse_transform_0213=False) -> Tensor + device_check: NoCheck # cpu_nested_shape_example will always be on CPU + dispatch: + CPU: nested_from_padded_generic + CUDA: nested_from_padded_cuda + autogen: _nested_from_padded.out + +# These private functions are temporary. They will be updated/deleted when nested tensors switch to using SymInts for their metadata representation +- func: _nested_tensor_size(Tensor self) -> Tensor + variants: method + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: _nested_tensor_size + autogen: _nested_tensor_size.out + +- func: _nested_tensor_strides(Tensor self) -> Tensor + variants: method + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: _nested_tensor_strides + autogen: _nested_tensor_strides.out + +- func: _nested_tensor_storage_offsets(Tensor self) -> Tensor + variants: method + dispatch: + NestedTensorCPU, NestedTensorCUDA, NestedTensorMeta: _nested_tensor_storage_offsets + autogen: _nested_tensor_storage_offsets.out + +# _nested_from_padded is not usable from Python, so +# _nested_from_padded_and_nested_example is available for testing. +- func: _nested_from_padded_and_nested_example(Tensor padded, Tensor nt_example) -> Tensor + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_from_padded_and_nested_example + autogen: _nested_from_padded_and_nested_example.out + +# The input arguments' types to this functions are temporary. When nested tensors switch to using SymInts for their metadata representation +# this will need to be updated +- func: _nested_view_from_buffer(Tensor(a) self, Tensor nested_size, Tensor nested_strides, Tensor offsets) -> Tensor(a) + variants: function + device_check: NoCheck + dispatch: + CPU, CUDA: _nested_view_from_buffer + +- func: _nested_view_from_buffer_copy(Tensor self, Tensor nested_size, Tensor nested_strides, Tensor offsets) -> Tensor + variants: function + device_check: NoCheck + tags: view_copy + dispatch: + CompositeExplicitAutogradNonFunctional: _nested_view_from_buffer_copy + autogen: _nested_view_from_buffer_copy.out + +- func: _nested_view_from_jagged(Tensor(a) self, Tensor offsets, Tensor dummy, Tensor? lengths=None, int ragged_idx=1, Tensor? min_seqlen=None, Tensor? max_seqlen=None) -> Tensor(a) + variants: function + device_check: NoCheck + dispatch: {} + +- func: _nested_view_from_jagged_copy(Tensor self, Tensor offsets, Tensor dummy, Tensor? lengths=None, int ragged_idx=1, Tensor? min_seqlen=None, Tensor? max_seqlen=None) -> Tensor + variants: function + device_check: NoCheck + tags: view_copy + dispatch: + CompositeExplicitAutogradNonFunctional: _nested_view_from_jagged_copy + autogen: _nested_view_from_jagged_copy.out + +- func: _nested_get_values(Tensor(a) self) -> Tensor(a) + variants: function + device_check: NoCheck + dispatch: {} + +- func: _nested_get_values_copy(Tensor self) -> Tensor + variants: function + device_check: NoCheck + tags: view_copy + dispatch: + CompositeExplicitAutogradNonFunctional: _nested_get_values_copy + autogen: _nested_get_values_copy.out + +- func: _nested_get_offsets(Tensor self) -> Tensor + variants: function + device_check: NoCheck + dispatch: {} + +# returns undefined Tensor if no lengths present +- func: _nested_get_lengths(Tensor self) -> Tensor + variants: function + device_check: NoCheck + dispatch: {} + +- func: _nested_get_ragged_idx(Tensor self) -> int + variants: function + device_check: NoCheck + dispatch: {} + +- func: _nested_get_min_seqlen(Tensor self) -> Tensor + variants: function + device_check: NoCheck + dispatch: {} + +- func: _nested_get_max_seqlen(Tensor self) -> Tensor + variants: function + device_check: NoCheck + dispatch: {} + +- func: _nested_get_jagged_dummy(Tensor any) -> Tensor + category_override: dummy + dispatch: {} + +- func: _nested_compute_contiguous_strides_offsets(Tensor nested_size) -> (Tensor, Tensor) + variants: function + device_check: NoCheck + dispatch: + CPU, CUDA: _nested_compute_contiguous_strides_offsets + +- func: _trilinear(Tensor i1, Tensor i2, Tensor i3, int[] expand1, int[] expand2, int[] expand3, int[] sumdim, int unroll_dim=1) -> Tensor + dispatch: + # calls unsqueeze + CompositeExplicitAutogradNonFunctional: _trilinear + autogen: _trilinear.out + +- func: triplet_margin_loss(Tensor anchor, Tensor positive, Tensor negative, float margin=1.0, float p=2, float eps=1e-06, bool swap=False, int reduction=Mean) -> Tensor + +- func: trunc(Tensor self) -> Tensor + structured_delegate: trunc.out + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: trunc_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: trunc_sparse_csr + tags: [core, pointwise] + +- func: trunc_(Tensor(a!) self) -> Tensor(a!) + structured_delegate: trunc.out + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: trunc_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: trunc_sparse_csr_ + tags: pointwise + +- func: trunc.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS: trunc_out + SparseCPU, SparseCUDA, SparseMPS: trunc_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: trunc_sparse_csr_out + tags: pointwise +# Alias for trunc + +- func: fix(Tensor self) -> Tensor + variants: function, method + +- func: fix_(Tensor(a!) self) -> Tensor(a!) + variants: function, method + +- func: fix.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: type_as(Tensor self, Tensor other) -> Tensor + variants: method + +- func: _has_compatible_shallow_copy_type(Tensor self, Tensor from) -> bool + variants: function + +- func: _unique(Tensor self, bool sorted=True, bool return_inverse=False) -> (Tensor, Tensor) + variants: function + dispatch: + CPU: _unique_cpu + CUDA: _unique_cuda + MPS: _unique_mps + autogen: _unique.out + +- func: unique_dim(Tensor self, int dim, bool sorted=True, bool return_inverse=False, bool return_counts=False) -> (Tensor, Tensor, Tensor) + variants: function + dispatch: + CPU: unique_dim_cpu + CUDA: unique_dim_cuda + MPS: unique_dim_mps + tags: dynamic_output_shape + autogen: unique_dim.out + +- func: unique_consecutive(Tensor self, bool return_inverse=False, bool return_counts=False, int? dim=None) -> (Tensor, Tensor, Tensor) + variants: function + dispatch: + CPU: unique_consecutive_cpu + CUDA: unique_consecutive_cuda + MPS: unique_consecutive_mps + tags: dynamic_output_shape + autogen: unique_consecutive.out + +- func: unique_dim_consecutive(Tensor self, int dim, bool return_inverse=False, bool return_counts=False) -> (Tensor, Tensor, Tensor) + variants: function + dispatch: + CPU: unique_dim_consecutive_cpu + CUDA: unique_dim_consecutive_cuda + MPS: unique_dim_consecutive_mps + tags: dynamic_output_shape + autogen: unique_dim_consecutive.out + +# _unique and _unique_dim are fragile and modifying them easily cause internal break +# the below operator is a temporary hack for adding return_counts support +# Please don't rely on these two operators, they will be removed soon + +- func: _unique2(Tensor self, bool sorted=True, bool return_inverse=False, bool return_counts=False) -> (Tensor, Tensor, Tensor) + variants: function + dispatch: + CPU: _unique2_cpu + CUDA: _unique2_cuda + MPS: _unique2_mps + tags: dynamic_output_shape + autogen: _unique2.out + +- func: _unsafe_view(Tensor self, SymInt[] size) -> Tensor + dispatch: + CompositeExplicitAutograd: _unsafe_view + autogen: _unsafe_view.out + +- func: unsqueeze(Tensor(a) self, int dim) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: unsqueeze + SparseCPU, SparseCUDA, SparseMPS: unsqueeze_sparse + QuantizedCPU, QuantizedCUDA: unsqueeze_quantized + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: unsqueeze_nested + tags: core + +- func: unsqueeze_(Tensor(a!) self, int dim) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + tags: inplace_view + dispatch: + CompositeExplicitAutograd: unsqueeze_ + +- func: vander(Tensor x, int? N=None, bool increasing=False) -> Tensor + +- func: var(Tensor self, bool unbiased=True) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: var.dim(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: [core, reduction] + cpp_no_default_args: ["unbiased"] + +- func: var.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA: var + MPS: var_mps + tags: [core, reduction] + +- func: var.out(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: var.correction_out(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: var_out + tags: reduction + +- func: var.names_dim(Tensor self, Dimname[1] dim, bool unbiased=True, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: var.names_out(Tensor self, Dimname[1] dim, bool unbiased=True, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: var.correction_names(Tensor self, Dimname[1] dim, *, Scalar? correction=None, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: var.correction_names_out(Tensor self, Dimname[1] dim, *, Scalar? correction=None, bool keepdim=False, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + tags: reduction + +- func: var_mean(Tensor self, bool unbiased=True) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: var_mean.dim(Tensor self, int[1]? dim, bool unbiased=True, bool keepdim=False) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: var_mean.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CPU, CUDA: var_mean + MPS: var_mean_mps + autogen: var_mean.correction_out + tags: reduction + +- func: var_mean.names_dim(Tensor self, Dimname[1] dim, bool unbiased=True, bool keepdim=False) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + cpp_no_default_args: ["unbiased"] + tags: reduction + +- func: var_mean.correction_names(Tensor self, Dimname[1] dim, *, Scalar? correction=None, bool keepdim=False) -> (Tensor, Tensor) + device_check: NoCheck # TensorIterator + variants: function + tags: reduction + +- func: view_as(Tensor(a) self, Tensor other) -> Tensor(a) + variants: method + device_check: NoCheck + device_guard: False + +- func: where.self(Tensor condition, Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CPU, CUDA, MPS, MTIA: where + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_where + tags: [core, pointwise] + +- func: where.self_out(Tensor condition, Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS, MTIA: where_self_out + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_where_out + +- func: where.ScalarSelf(Tensor condition, Scalar self, Tensor other) -> Tensor + variants: function + +- func: where.ScalarOther(Tensor condition, Tensor self, Scalar other) -> Tensor + variants: function, method + +- func: where.Scalar(Tensor condition, Scalar self, Scalar other) -> Tensor + variants: function + +- func: where(Tensor condition) -> Tensor[] + device_check: NoCheck # TensorIterator + variants: function + +- func: norm_except_dim(Tensor v, int pow=2, int dim=0) -> Tensor + variants: function + +# VariableType::_weight_norm does not want to be given a gap in the autograd graph, +# so we don't define "dispatch" variants for it. +- func: _weight_norm(Tensor v, Tensor g, int dim=0) -> Tensor + variants: function + +- func: _weight_norm_interface(Tensor v, Tensor g, int dim=0) -> (Tensor, Tensor) + variants: function + dispatch: + CPU: weight_norm_cpu + CUDA: weight_norm_cuda + MPS: weight_norm_mps + autogen: _weight_norm_interface.out + +- func: _weight_norm_interface_backward(Tensor grad_w, Tensor saved_v, Tensor saved_g, Tensor saved_norms, int dim) -> (Tensor, Tensor) + variants: function + dispatch: + CPU: weight_norm_backward_cpu + CUDA: weight_norm_backward_cuda + MPS: weight_norm_backward_mps + autogen: _weight_norm_interface_backward.out + +- func: _weight_norm_differentiable_backward(Tensor grad_w, Tensor saved_v, Tensor saved_g, Tensor saved_norms, int dim) -> (Tensor, Tensor) + variants: function + +- func: zeros.names(int[] size, *, Dimname[]? names, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: zeros + autogen: zeros.names_out + +- func: _efficientzerotensor(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CPU: _efficientzerotensor + CUDA: _efficientzerotensor_cuda + MPS: _efficientzerotensor_mps + Meta: _efficientzerotensor_meta_symint + autogen: _efficientzerotensor.out + +- func: zeros(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: zeros_symint + +- func: zeros.out(SymInt[] size, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: zeros_out + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: zeros_sparse_out + +- func: zeros_like(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, MemoryFormat? memory_format=None) -> Tensor + dispatch: + # NB: Although this composite mutates on the inside, it is + # non-differentiable so NonFunctional doesn't apply + CompositeExplicitAutograd, CompositeImplicitAutogradNestedTensor: zeros_like + autogen: zeros_like.out + +- func: _standard_gamma_grad(Tensor self, Tensor output) -> Tensor + variants: function + dispatch: + CPU: _standard_gamma_grad_cpu + CUDA: _standard_gamma_grad_cuda + MPS: _standard_gamma_grad_mps + autogen: _standard_gamma_grad.out + +- func: _standard_gamma(Tensor self, Generator? generator=None) -> Tensor + variants: function + dispatch: + CPU: _s_gamma_cpu + CUDA: _s_gamma_cuda + MPS: _s_gamma_mps + tags: nondeterministic_seeded + autogen: _standard_gamma.out + +- func: _philox_key_split(Tensor key, int num_splits) -> Tensor + variants: function + dispatch: + CUDA: _philox_key_split_cuda + +- func: _philox_key_fold_in(Tensor key, int data) -> Tensor + variants: function + dispatch: + CUDA: _philox_key_fold_in_cuda + +- func: _philox_normal_(Tensor(a!) self, Tensor key, float mean=0, float std=1) -> Tensor(a!) + variants: function, method + dispatch: + CUDA: _philox_normal_cuda_ + autogen: _philox_normal, _philox_normal.out + +- func: _philox_uniform_(Tensor(a!) self, Tensor key, float low=0, float high=1) -> Tensor(a!) + variants: function, method + dispatch: + CUDA: _philox_uniform_cuda_ + autogen: _philox_uniform, _philox_uniform.out + +- func: _dirichlet_grad(Tensor x, Tensor alpha, Tensor total) -> Tensor + dispatch: + CPU: _dirichlet_grad_cpu + CUDA: _dirichlet_grad_cuda + autogen: _dirichlet_grad.out + +- func: _sample_dirichlet(Tensor self, Generator? generator=None) -> Tensor + tags: nondeterministic_seeded + variants: function + dispatch: + CPU: _s_dirichlet_cpu + CUDA: _s_dirichlet_cuda + autogen: _sample_dirichlet.out + +- func: poisson(Tensor self, Generator? generator=None) -> Tensor + device_check: NoCheck # TensorIterator + dispatch: + CPU: _s_poisson_cpu + CUDA: _s_poisson_cuda + tags: nondeterministic_seeded + autogen: poisson.out + +- func: binomial(Tensor count, Tensor prob, Generator? generator=None) -> Tensor + device_check: NoCheck # TensorIterator + dispatch: + CPU: _s_binomial_cpu + CUDA: _s_binomial_cuda + tags: nondeterministic_seeded + autogen: binomial.out + +# When more variants get ported to native, this dispatch will get more +# complicated + +- func: native_norm(Tensor self, Scalar p=2) -> Tensor + dispatch: + SparseCPU, SparseCUDA, SparseMPS: norm_sparse + autogen: native_norm.out + +- func: native_norm.ScalarOpt_dim_dtype(Tensor self, Scalar? p, int[1] dim, bool keepdim, ScalarType? dtype) -> Tensor + dispatch: + SparseCPU, SparseCUDA, SparseMPS: norm_sparse + autogen: native_norm.ScalarOpt_dim_dtype_out + +- func: _batch_norm_with_update(Tensor input, Tensor? weight, Tensor? bias, Tensor(a!) running_mean, Tensor(b!) running_var, float momentum, float eps) -> (Tensor, Tensor, Tensor, Tensor) + dispatch: + CPU: _batch_norm_with_update_cpu + CUDA: _batch_norm_with_update_cuda + MPS: _batch_norm_with_update_mps + MkldnnCPU: _batch_norm_with_update_mkldnn + autogen: _batch_norm_with_update_functional + +- func: _batch_norm_with_update.out(Tensor input, Tensor? weight, Tensor? bias, Tensor(a!) running_mean, Tensor(b!) running_var, float momentum, float eps, *, Tensor(d!) out, Tensor(e!) save_mean, Tensor(f!) save_invstd, Tensor(g!) reserve) -> (Tensor(d!), Tensor(e!), Tensor(f!), Tensor(g!)) + dispatch: + CPU: _batch_norm_with_update_cpu_out + CUDA: _batch_norm_with_update_cuda_out + MPS: _batch_norm_with_update_mps_out + +- func: _batch_norm_no_update(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, float momentum, float eps) -> (Tensor, Tensor, Tensor, Tensor) + dispatch: + CompositeExplicitAutograd: _batch_norm_no_update + autogen: _batch_norm_no_update.out + +- func: batch_norm_backward(Tensor grad_out, Tensor input, Tensor weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_var, bool update, float eps, bool[3] output_mask, Tensor reserve) -> (Tensor, Tensor, Tensor) + dispatch: + CPU: _new_batch_norm_backward_cpu + CUDA: _new_batch_norm_backward_cuda + MPS: _new_batch_norm_backward_mps + MkldnnCPU: _new_batch_norm_backward_mkldnn + +# TODO: reduce signatures down to one when optional args is available +- func: _sparse_sum(Tensor self) -> Tensor + +- func: _sparse_sum.dtype(Tensor self, *, ScalarType dtype) -> Tensor + +- func: _sparse_sum.dim(Tensor self, int[1] dim) -> Tensor + dispatch: + CompositeExplicitAutograd: _sparse_sum + autogen: _sparse_sum.dim_out + +- func: _sparse_sum.dim_dtype(Tensor self, int[1] dim, *, ScalarType dtype) -> Tensor + +- func: _sparse_sum_backward(Tensor grad, Tensor self, int[] dim) -> Tensor + dispatch: + SparseCPU: _sparse_sum_backward_cpu + SparseCUDA: _sparse_sum_backward_cuda + SparseMPS: _sparse_sum_backward_mps + autogen: _sparse_sum_backward.out + +- func: _sparse_csr_sum.dim_dtype(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + dispatch: + SparseCsrCPU: _sparse_csr_sum_cpu + SparseCsrCUDA: _sparse_csr_sum_cuda + autogen: _sparse_csr_sum.dim_dtype_out + +- func: _sparse_csr_prod.dim_dtype(Tensor self, int[1] dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + dispatch: + SparseCsrCPU: _sparse_csr_prod_cpu + SparseCsrCUDA: _sparse_csr_prod_cuda + autogen: _sparse_csr_prod.dim_dtype_out + +- func: _sparse_softmax.int(Tensor self, int dim, ScalarType? dtype=None) -> Tensor + python_module: sparse + variants: function + +- func: _sparse_softmax.Dimname(Tensor self, Dimname dim, *, ScalarType? dtype=None) -> Tensor + python_module: sparse + variants: function + +- func: _sparse_softmax(Tensor self, int dim, bool half_to_float) -> Tensor + python_module: sparse + dispatch: + SparseCPU: softmax_sparse_cpu + SparseCUDA: softmax_sparse_cuda + SparseMPS: softmax_sparse_mps + autogen: _sparse_softmax.out + +- func: _sparse_softmax_backward_data(Tensor grad_output, Tensor output, int dim, Tensor self) -> Tensor + dispatch: + SparseCPU: softmax_backward_sparse_cpu + SparseCUDA: softmax_backward_sparse_cuda + SparseMPS: softmax_backward_sparse_mps + autogen: _sparse_softmax_backward_data.out + +- func: _sparse_log_softmax.int(Tensor self, int dim, ScalarType? dtype=None) -> Tensor + python_module: sparse + variants: function + +- func: _sparse_log_softmax.Dimname(Tensor self, Dimname dim, *, ScalarType? dtype=None) -> Tensor + python_module: sparse + variants: function + +- func: _sparse_log_softmax(Tensor self, int dim, bool half_to_float) -> Tensor + python_module: sparse + dispatch: + SparseCPU: log_softmax_sparse_cpu + SparseCUDA: log_softmax_sparse_cuda + SparseMPS: log_softmax_sparse_mps + autogen: _sparse_log_softmax.out + +- func: _sparse_log_softmax_backward_data(Tensor grad_output, Tensor output, int dim, Tensor self) -> Tensor + dispatch: + SparseCPU: log_softmax_backward_sparse_cpu + SparseCUDA: log_softmax_backward_sparse_cuda + SparseMPS: log_softmax_backward_sparse_mps + autogen: _sparse_log_softmax_backward_data.out + +- func: _spdiags(Tensor diagonals, Tensor offsets, int[] shape, Layout? layout=None) -> Tensor + python_module: sparse + dispatch: + CPU: spdiags + autogen: _spdiags.out + +- func: norm.ScalarOpt_dtype(Tensor self, Scalar? p, *, ScalarType dtype) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: norm + autogen: norm.ScalarOpt_dtype_out + tags: reduction + +- func: norm.Scalar(Tensor self, Scalar p=2) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: norm + autogen: norm.Scalar_out + tags: reduction + +- func: norm.ScalarOpt_dim_dtype(Tensor self, Scalar? p, int[1] dim, bool keepdim, *, ScalarType dtype) -> Tensor + structured_delegate: norm.dtype_out + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sparse_dtype_norm + tags: reduction + +- func: norm.ScalarOpt_dim(Tensor self, Scalar? p, int[1] dim, bool keepdim=False) -> Tensor + structured_delegate: norm.out + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sparse_norm + tags: reduction + +- func: norm.dtype_out(Tensor self, Scalar? p, int[1] dim, bool keepdim, *, ScalarType dtype, Tensor(a!) out) -> Tensor(a!) + structured: True + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS: norm_dtype_out + tags: reduction + +- func: norm.out(Tensor self, Scalar? p, int[1] dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + structured: True + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS: norm_out + tags: reduction + +# These four redispatch in their implementation, so OK to be CompositeImplicitAutograd +- func: norm.names_ScalarOpt_dim_dtype(Tensor self, Scalar? p, Dimname[1] dim, bool keepdim, *, ScalarType dtype) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: norm.names_ScalarOpt_dim(Tensor self, Scalar? p, Dimname[1] dim, bool keepdim=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + tags: reduction + +- func: norm.names_dtype_out(Tensor self, Scalar? p, Dimname[1] dim, bool keepdim, *, ScalarType dtype, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: reduction + +- func: norm.names_out(Tensor self, Scalar? p, Dimname[1] dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: reduction + +- func: frexp.Tensor(Tensor self) -> (Tensor mantissa, Tensor exponent) + variants: method, function + dispatch: + CompositeExplicitAutograd: frexp + tags: pointwise + +- func: frexp.Tensor_out(Tensor self, *, Tensor(a!) mantissa, Tensor(b!) exponent) -> (Tensor(a!) mantissa, Tensor(b!) exponent) + dispatch: + CPU, CUDA: frexp_out + tags: pointwise + +# Deprecated (v.1.12) +- func: frobenius_norm.dim(Tensor self, int[1] dim, bool keepdim=False) -> Tensor + variants: function + +# Deprecated (v.1.12) +- func: frobenius_norm.out(Tensor self, int[1] dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + variants: function + +# Deprecated (v.1.12) +- func: nuclear_norm(Tensor self, bool keepdim=False) -> Tensor + variants: function + +# Deprecated (v.1.12) +- func: nuclear_norm.out(Tensor self, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + variants: function + +# Deprecated (v.1.12) +- func: nuclear_norm.dim(Tensor self, int[2] dim, bool keepdim=False) -> Tensor + variants: function + +# Deprecated (v.1.12) +- func: nuclear_norm.dim_out(Tensor self, int[2] dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + variants: function + +- func: clone(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: clone + SparseCPU, SparseCUDA, SparseMPS: clone_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: clone_sparse_compressed + MkldnnCPU: mkldnn_clone + QuantizedCPU, QuantizedCUDA: quantized_clone + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: clone_nested + autogen: clone.out + tags: [core, pointwise] + +- func: positive(Tensor(a) self) -> Tensor(a) + variants: function, method + tags: pointwise + +- func: resize_as_(Tensor(a!) self, Tensor the_template, *, MemoryFormat? memory_format=None) -> Tensor(a!) + use_const_ref_for_mutable_tensors: True + variants: function, method + dispatch: + CompositeExplicitAutograd: resize_as_ + autogen: resize_as, resize_as.out + tags: inplace_view + +- func: resize_as_sparse_(Tensor(a!) self, Tensor the_template) -> Tensor(a!) + use_const_ref_for_mutable_tensors: True + variants: function, method + dispatch: + SparseCPU, SparseCUDA: resize_as_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: resize_as_sparse_compressed_ + autogen: resize_as_sparse, resize_as_sparse.out + +- func: zero_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CPU, CUDA, MPS: zero_ + Meta: zero_meta_ + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: zero_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: zero_sparse_csr_ + MkldnnCPU: mkldnn_zero_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: zero_nested_ + autogen: zero, zero.out + +- func: sub.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: sub_out + MPS: sub_out_mps + MTIA: sub_out_mtia + SparseCPU, SparseCUDA, SparseMPS: sub_out_sparse + tags: pointwise + +- func: sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: sub.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sub_sparse + ZeroTensor: sub_zerotensor + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_sub_Tensor + tags: [core, pointwise] + +- func: sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: sub.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sub_sparse_ + tags: pointwise +# For C++ only, until we have conversion from C++ numbers to Tensor + +- func: sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: sub + tags: [core, pointwise] + +- func: sub_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: sub_ + autogen: sub.Scalar_out + tags: pointwise +# subtract, alias for sub + +- func: subtract.out(Tensor self, Tensor other, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + +- func: subtract.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor + variants: function, method + +- func: subtract_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> Tensor(a!) + variants: method + +# For C++ only, until we have conversion from C++ numbers to Tensor +- func: subtract.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor + variants: function, method + +- func: subtract_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!) + variants: method + +- func: rsub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CPU, CUDA, MPS, MTIA: rsub + autogen: rsub.Tensor_out + +- func: heaviside.out(Tensor self, Tensor values, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: heaviside_out + tags: pointwise + +- func: heaviside(Tensor self, Tensor values) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: heaviside.out + tags: pointwise + +- func: heaviside_(Tensor(a!) self, Tensor values) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: heaviside.out + +# For C++ only, until we have conversion from C++ numbers to Tensor +- func: rsub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: rsub + autogen: rsub.Scalar_out + +# Functionally the same as addmm, but we give it a different derivative formula +# that doesn't propagate gradients to non-present entries on sparse. + tags: pointwise +- func: _sparse_addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + python_module: sparse + dispatch: + CompositeExplicitAutograd: _sparse_addmm + autogen: _sparse_addmm.out + +- func: sparse_sampled_addmm.out(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + python_module: sparse + dispatch: + SparseCsrCUDA: sparse_sampled_addmm_out_sparse_csr_cuda + SparseCsrCPU: sparse_sampled_addmm_out_sparse_csr_cpu + +- func: sparse_sampled_addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + python_module: sparse + dispatch: + SparseCsrCUDA: sparse_sampled_addmm_sparse_csr_cuda + SparseCsrCPU: sparse_sampled_addmm_sparse_csr_cpu + +- func: _sparse_mm_reduce_impl(Tensor self, Tensor other, str reduce) -> (Tensor, Tensor) + python_module: sparse + dispatch: + SparseCsrCPU: _sparse_mm_reduce_impl_sparse_csr_cpu + +- func: _sparse_mm_reduce_impl_backward(Tensor self, Tensor grad_out, Tensor weight, str reduce, Tensor arg_out, bool[2] output_mask) -> (Tensor, Tensor) + python_module: sparse + dispatch: + SparseCsrCPU: _sparse_mm_reduce_impl_backward_sparse_csr_cpu + +- func: addmm.out(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: addmm_out_cpu + CUDA: addmm_out_cuda + MPS: addmm_out_mps + XPU: addmm_out_xpu + MTIA: addmm_out_mtia + SparseCPU: addmm_out_sparse_dense_cpu + SparseCUDA: addmm_out_sparse_dense_cuda + SparseMPS: addmm_out_sparse_dense_mps + SparseCsrCPU: addmm_out_sparse_compressed_cpu + SparseCsrCUDA: addmm_out_sparse_compressed_cuda + +- func: addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + structured_delegate: addmm.out + variants: function, method + dispatch: + SparseCPU: addmm_sparse_dense_cpu + SparseCUDA: addmm_sparse_dense_cuda + SparseMPS: addmm_sparse_dense_mps + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: addmm_sparse_compressed_dense + tags: core + +- func: addmm.dtype(Tensor self, Tensor mat1, Tensor mat2, ScalarType out_dtype, *, Scalar beta=1, Scalar alpha=1) -> Tensor + dispatch: + CUDA: _addmm_dtype_cuda + XPU: _addmm_dtype_xpu + +- func: addmm.dtype_out(Tensor self, Tensor mat1, Tensor mat2, ScalarType out_dtype, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + dispatch: + CUDA: _addmm_dtype_out_cuda + XPU: _addmm_dtype_out_xpu + +- func: addmm_(Tensor(a!) self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor(a!) + structured_delegate: addmm.out + variants: method + dispatch: + # Warning! For whatever reason, the inplace sparse addmm is NON + # broadcasting + SparseCPU: s_addmm_sparse_dense_cpu_ + SparseCUDA: s_addmm_sparse_dense_cuda_ + +- func: _addmm_activation.out(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1, bool use_gelu=False, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: addmm_activation_out_cpu + CUDA: addmm_activation_out_cuda + XPU: addmm_activation_out_xpu + +- func: _addmm_activation(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1, bool use_gelu=False) -> Tensor + structured_delegate: _addmm_activation.out + variants: function, method + +- func: _scaled_mm(Tensor self, Tensor mat2, Tensor scale_a, Tensor scale_b, Tensor? bias=None, Tensor? scale_result=None, ScalarType? out_dtype=None, bool use_fast_accum=False) -> Tensor + variants: function + dispatch: + CPU: _scaled_mm_cpu + CUDA: _scaled_mm_cuda + XPU: _scaled_mm_xpu + tags: needs_exact_strides + + +- func: _scaled_mm.out(Tensor self, Tensor mat2, Tensor scale_a, Tensor scale_b, Tensor? bias=None, Tensor? scale_result=None, ScalarType? out_dtype=None, bool use_fast_accum=False, *, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CPU: _scaled_mm_out_cpu + CUDA: _scaled_mm_out_cuda + XPU: _scaled_mm_out_xpu + tags: needs_exact_strides + +- func: _scaled_mm_v2(Tensor self, Tensor mat2, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? bias, ScalarType? out_dtype, int[] contraction_dim=[], bool use_fast_accum=False) -> Tensor + variants: function + dispatch: + CPU: _scaled_mm_cpu_v2 + CUDA: _scaled_mm_cuda_v2 + XPU: _scaled_mm_xpu_v2 + tags: needs_exact_strides + +- func: _scaled_mm_v2.out(Tensor self, Tensor mat2, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? bias, ScalarType? out_dtype, int[] contraction_dim=[], bool use_fast_accum=False, *, Tensor(a!) out) -> Tensor(a!) + variants: function + dispatch: + CPU: _scaled_mm_cpu_v2_out + CUDA: _scaled_mm_cuda_v2_out + XPU: _scaled_mm_xpu_v2_out + tags: needs_exact_strides + + +- func: _scaled_grouped_mm(Tensor self, Tensor mat2, Tensor scale_a, Tensor scale_b, Tensor? offs=None, Tensor? bias=None, Tensor? scale_result=None, ScalarType? out_dtype=None, bool use_fast_accum=False) -> Tensor + variants: function + dispatch: + CUDA: _scaled_grouped_mm_cuda + tags: needs_exact_strides + +- func: _scaled_grouped_mm_v2(Tensor self, Tensor mat2, Tensor[] scale_a, int[] recipe_a, int[] swizzle_a, Tensor[] scale_b, int[] recipe_b, int[] swizzle_b, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None, int[] contraction_dim=[], bool use_fast_accum=False) -> Tensor + variants: function + dispatch: + CUDA: _scaled_grouped_mm_cuda_v2 + tags: needs_exact_strides + +- func: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: _grouped_mm + CUDA: _grouped_mm_cuda + +# NOTE [ Sparse: autograd and API ] +# +# +# Sparse Tensor Constructors +# ~~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# The API entry points to sparse tensor construction should be +# `sparse_coo tensor` and `_sparse_coo_tensor_unsafe`. Depending on whether the +# indices and values tensors are given, they eventually dispatch to either +# `sparse_coo_tensor_with_dims` or `sparse_coo_tensor_with_dims_and_tensors`. +# +# The autograd support for ctor is implement on `sparse_coo_tensor_with_dims_and_tensors`. +# +# The API methods `sparse_coo tensor` and `_sparse_coo_tensor_unsafe` +# **must not** have specific type dispatches because otherwise codegen will +# consider them as abstract methods (see Note [Abstract ATen methods]), dispatch +# using **Tensor** type, and thus lose autograd tracking on the actual method +# they dispatch to, e.g., `sparse_coo_tensor_with_dims_and_tensors`. +# +# +# Sparse Methods API Design +# ~~~~~~~~~~~~~~~~~~~~~~~~~ +# +# Goals: 1. Flexible API for users to write custom sparse ops +# 2. ctor and member accessor with autograd support +# +# To achieve 1, we need to provide a set of *dangerous* APIs (dangerous in the +# sense that misusing them will break sparse tensor invariant and may out in +# unexpected behavior, e.g., crash). These methods are all prefixed with +# underscore "_" to indicate that they should be used with care. We provide: +# +# + `_indices()`: returns the *raw* indices within the sparse tensor (not just +# sharing storage). Any inplace operation will change the +# actual indices, including t_, set_, as_strided_, resize_, +# etc. +# + `_values()`: returns the *raw* values within the sparse tensor. Similar +# semantics as `_indices()` +# + `_nnz()`: returns the number of non-zero entries. This will always be +# determined by the shapes of indices and values. +# + `_coalesced_(bool)`: inplace sets whether the tensor is coalesced, and +# returns itself. +# +# These methods are very useful in writing new operations, e.g., a custom +# autograd Function. +# +# We also provide other public *safe* APIs: +# + `indices()`: returns a **view** of the indices tensor if the sparse tensor +# is **coalesced**. +# + `values()`: returns a **view** of the values tensor if the containing +# sparse tensor is **coalesced**. +# + `sparse_dim()`: number of sparse dimensions +# + `dense_dim()`: number of dense dimensions +# + `is_coalesced()`: whether the sparse tensor is coalesced +# +# `_indices()` and `_values()` should returns the raw indices and values dense +# tensors within a sparse tensor. They can be quite unsafe with inplace +# operations like `t_()`, and exposes uncoalesced indices and values. The public +# recommended API is `indices()` and `values()`, both of which first check that +# the tensor is coalesced and return views on those tensors. +# +# +# Autograd Support +# ~~~~~~~~~~~~~~~~ +# +# Autograd is supported on `values()` and sparse tensor ctor with indices and +# values tensors. E.g., `torch.sparse_coo_tensor(i, v).values().sum()` is +# differentiable w.r.t. `v`. +# +# NB: The `values()` and `_values()` operators are special in that they are +# layout-aware, i.e., the output depends not just on the data it represents, but +# also on the input layout details (in this case, the `indices` tensor). See +# NOTE [ as_strided Backward and layout-aware/agnostic autograd ] in Functions.cpp +# for discussion on layout-aware vs layout-agnostic autograd. Since PyTorch ops +# operate in the layout-agnostic mode, similar to `as_strided`, backward of +# these two operators need to consider them in a layout-agnostic way: +# + `values()`: +# Input is coalesced. +# We just pretend having `input.indices()` as an additional argument +# `input_indices`, then forward is similar to +# `input.to(kStrided).index_select(input_indices)` regardless of the layout. +# Note that `values()` normally is layout-aware even if we constrain +# ourselves on sparse inputs since it may include all zeros values entries +# as "present" entries. +# + `_values()`: +# Input may be uncoalesced. +# It is not straightforward to construct a layout-agnostic version because +# duplicate indices entries may exist and additional parameterization is +# needed to distribute the value into different values entries. Furthermore, +# this op is intended to provide ways to write custom sparse ops, rather +# than being used in autograd graph, so it is marked as *non-differentiable* +# in derivatives.yaml. +# +# Before reading the following, see NOTE [ Autograd Variable Views ] in +# variable.h for details on views that are tracked by autograd, and views that +# are not. +# +# Moreover, these methods return tensors that share storage with inputs, so we +# mark these methods as view ops to support autograd history tracking. +# The sparse tensor ctor output should technically be view of both input indices +# and values tensors, but currently we only support setting as view of a single +# Variable, so it is only view of the values tensor. +# TODO: clone indices in sparse tensor ctor. +# +# For other methods that return outputs that share storage with inputs, i.e., +# `indices()` and `_indices()`. We mark their outputs as non-differentiable, so +# the view relation is not tracked by autograd, but the version counter is still +# shared. In other words, their outputs are non-differentiable views of the +# sparse tensor. +# FIXME: would be nicer if TensorOptions was optional based; not adding default arguments for options given +# the default would never make sense. + +- func: _sparse_compressed_tensor_with_dims(int nnz, int dense_dim, int[] size, int[] blocksize, ScalarType index_dtype, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor + dispatch: + CompositeExplicitAutograd: sparse_compressed_tensor_with_dims + +- func: sparse_compressed_tensor.comp_plain_value_size(Tensor compressed_indices, Tensor plain_indices, Tensor values, SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor + dispatch: + CompositeExplicitAutograd: sparse_compressed_tensor + +- func: sparse_csr_tensor.crow_col_value_size(Tensor crow_indices, Tensor col_indices, Tensor values, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor +- func: sparse_csc_tensor.ccol_row_value_size(Tensor ccol_indices, Tensor row_indices, Tensor values, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor +- func: sparse_bsr_tensor.crow_col_value_size(Tensor crow_indices, Tensor col_indices, Tensor values, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor +- func: sparse_bsc_tensor.ccol_row_value_size(Tensor ccol_indices, Tensor row_indices, Tensor values, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor + +- func: sparse_compressed_tensor.comp_plain_value(Tensor compressed_indices, Tensor plain_indices, Tensor values, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor + dispatch: + CompositeExplicitAutograd: sparse_compressed_tensor +- func: sparse_csr_tensor.crow_col_value(Tensor crow_indices, Tensor col_indices, Tensor values, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor +- func: sparse_csc_tensor.ccol_row_value(Tensor ccol_indices, Tensor row_indices, Tensor values, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor +- func: sparse_bsr_tensor.crow_col_value(Tensor crow_indices, Tensor col_indices, Tensor values, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor +- func: sparse_bsc_tensor.ccol_row_value(Tensor ccol_indices, Tensor row_indices, Tensor values, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor + +- func: _sparse_compressed_tensor_unsafe(Tensor compressed_indices, Tensor plain_indices, Tensor values, SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeImplicitAutograd: _sparse_compressed_tensor_unsafe_symint + +- func: _sparse_csr_tensor_unsafe(Tensor crow_indices, Tensor col_indices, Tensor values, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor +- func: _sparse_csc_tensor_unsafe(Tensor ccol_indices, Tensor row_indices, Tensor values, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor +- func: _sparse_bsr_tensor_unsafe(Tensor crow_indices, Tensor col_indices, Tensor values, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor +- func: _sparse_bsc_tensor_unsafe(Tensor ccol_indices, Tensor row_indices, Tensor values, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + +- func: sparse_coo_tensor.size(int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor + dispatch: + CompositeExplicitAutograd: sparse_coo_tensor + autogen: sparse_coo_tensor.size_out + +- func: sparse_coo_tensor.indices(Tensor indices, Tensor values, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool? is_coalesced=None) -> Tensor + +- func: sparse_coo_tensor.indices_size(Tensor indices, Tensor values, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool? is_coalesced=None) -> Tensor + +- func: _sparse_coo_tensor_unsafe(Tensor indices, Tensor values, SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool? is_coalesced=None) -> Tensor + dispatch: + CompositeImplicitAutograd: _sparse_coo_tensor_unsafe_symint + +- func: _validate_sparse_coo_tensor_args(Tensor indices, Tensor values, int[] size, bool? is_coalesced=None, bool? check_pinning=None) -> () + +- func: _validate_sparse_compressed_tensor_args(Tensor compressed_indices, Tensor plain_indices, Tensor values, int[] size, Layout layout, bool? check_pinning=None) -> () +- func: _validate_sparse_csr_tensor_args(Tensor crow_indices, Tensor col_indices, Tensor values, int[] size, bool? check_pinning=None) -> () +- func: _validate_sparse_csc_tensor_args(Tensor ccol_indices, Tensor row_indices, Tensor values, int[] size, bool? check_pinning=None) -> () +- func: _validate_sparse_bsr_tensor_args(Tensor crow_indices, Tensor col_indices, Tensor values, int[] size, bool? check_pinning=None) -> () +- func: _validate_sparse_bsc_tensor_args(Tensor ccol_indices, Tensor row_indices, Tensor values, int[] size, bool? check_pinning=None) -> () + +- func: _sparse_coo_tensor_with_dims(int sparse_dim, int dense_dim, int[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor + dispatch: + SparseCPU, SparseCUDA, SparseMeta, SparseMPS, Meta: new_with_dims_sparse + autogen: _sparse_coo_tensor_with_dims.out + +- func: _sparse_coo_tensor_with_dims_and_tensors(int sparse_dim, int dense_dim, SymInt[] size, Tensor indices, Tensor values, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False, bool? is_coalesced=None) -> Tensor + dispatch: + SparseCPU, SparseCUDA, SparseMeta, SparseMPS, Meta: new_with_dims_and_tensor_sparse_symint + autogen: _sparse_coo_tensor_with_dims_and_tensors.out + +- func: sparse_resize_(Tensor(a!) self, int[] size, int sparse_dim, int dense_dim) -> Tensor(a!) + use_const_ref_for_mutable_tensors: True + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: sparse_resize_ + autogen: sparse_resize, sparse_resize.out + +- func: sparse_resize_and_clear_(Tensor(a!) self, int[] size, int sparse_dim, int dense_dim) -> Tensor(a!) + use_const_ref_for_mutable_tensors: True + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: sparse_resize_and_clear_ + autogen: sparse_resize_and_clear, sparse_resize_and_clear.out + +- func: sparse_mask(Tensor self, Tensor mask) -> Tensor + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sparse_mask + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sparse_mask_sparse_compressed + autogen: sparse_mask.out + +- func: _sparse_mask_projection(Tensor self, Tensor mask, bool accumulate_matches=False) -> Tensor + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sparse_mask_projection + autogen: _sparse_mask_projection.out + +- func: _to_cpu(Tensor[] tensors) -> Tensor[] + variants: function + +- func: to_dense(Tensor self, ScalarType? dtype=None, *, bool? masked_grad=None) -> Tensor + variants: method + +# Special case of to_dense with custom derivative +- func: _to_dense(Tensor self, ScalarType? dtype=None, bool? masked_grad=None) -> Tensor + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sparse_to_dense + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: sparse_compressed_to_dense + MkldnnCPU: mkldnn_to_dense + autogen: _to_dense.out + +- func: to_dense_backward(Tensor grad, Tensor input, bool? masked_grad=None) -> Tensor + +- func: sparse_dim(Tensor self) -> int + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: sparse_dim_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: sparse_dim_sparse_csr + CompositeExplicitAutograd: sparse_dim_default + device_check: NoCheck + device_guard: False + +# legacy method +- func: _dimI(Tensor self) -> int + variants: method + dispatch: + SparseCPU, SparseCUDA: sparse_dim_sparse + device_check: NoCheck + device_guard: False + +- func: dense_dim(Tensor self) -> int + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: dense_dim_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: dense_dim_sparse_csr + CompositeExplicitAutograd: dense_dim_default + device_check: NoCheck + device_guard: False + +# legacy method +- func: _dimV(Tensor self) -> int + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMeta: dense_dim_sparse + device_check: NoCheck + device_guard: False + +- func: _nnz(Tensor self) -> int + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: _nnz_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMPS, SparseCsrMeta: _nnz_sparse_csr + device_check: NoCheck + device_guard: False + +# NOTE: [ coalesce autograd ] +# coalesce returns self directly for already coalesced sparse tensors. +# This means coalesce cannot have a derivative registered, otherwise it creates +# circular references in the autograd graph (see gh-52874). +# Instead, the derivative is registered on the slow-path "_coalesce" +- func: coalesce(Tensor(a) self) -> Tensor(a) + variants: method + +- func: _coalesce(Tensor self) -> Tensor + dispatch: + SparseCPU: _coalesce_sparse_cpu + SparseCUDA: _coalesce_sparse_cuda + SparseMPS: _coalesce_sparse_mps + autogen: _coalesce.out + +- func: is_coalesced(Tensor self) -> bool + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: is_coalesced_sparse + CompositeExplicitAutograd: is_coalesced_default + device_check: NoCheck + device_guard: False + +- func: _indices(Tensor(a) self) -> Tensor(a) + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: _indices_sparse + device_check: NoCheck + device_guard: False + +- func: _values(Tensor(a) self) -> Tensor(a) + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: _values_sparse + device_check: NoCheck + device_guard: False + +# This method doesn't do any check but only directly sets the flag. So it can be +# a bit unsafe. Similar to _indices and _values, this is useful for implementing +# custom sparse operations in Python/C++ extension. +- func: _coalesced_(Tensor(a!) self, bool coalesced) -> Tensor(a!) + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: _coalesced_sparse_ + device_check: NoCheck + device_guard: False + autogen: _coalesced, _coalesced.out + +- func: indices(Tensor(a) self) -> Tensor(a) + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: indices_sparse + CompositeExplicitAutograd: indices_default + device_check: NoCheck + device_guard: False + +- func: values(Tensor(a) self) -> Tensor(a) + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: values_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: values_sparse_csr + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: values_nested + CompositeExplicitAutograd: values_default + device_check: NoCheck + device_guard: False + +- func: crow_indices(Tensor(a) self) -> Tensor(a) + variants: method + dispatch: + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: crow_indices_sparse_csr + CompositeExplicitAutograd: crow_indices_default + device_check: NoCheck + device_guard: False + +- func: col_indices(Tensor(a) self) -> Tensor(a) + variants: method + dispatch: + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: col_indices_sparse_csr + CompositeExplicitAutograd: col_indices_default + device_check: NoCheck + device_guard: False + +- func: ccol_indices(Tensor(a) self) -> Tensor(a) + variants: method + dispatch: + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: ccol_indices_sparse_csr + CompositeExplicitAutograd: ccol_indices_default + device_check: NoCheck + device_guard: False + +- func: row_indices(Tensor(a) self) -> Tensor(a) + variants: method + dispatch: + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: row_indices_sparse_csr + CompositeExplicitAutograd: row_indices_default + device_check: NoCheck + device_guard: False + +- func: hspmm.out(Tensor mat1, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + SparseCPU: hspmm_out_sparse_cpu + SparseCUDA: hspmm_out_sparse_cuda + +- func: hspmm(Tensor mat1, Tensor mat2) -> Tensor + dispatch: + SparseCPU: hspmm_sparse_cpu + SparseCUDA: hspmm_sparse_cuda + +- func: copy_sparse_to_sparse_(Tensor(a!) self, Tensor src, bool non_blocking=False) -> Tensor(a!) + device_check: NoCheck # Allows copy into different device + variants: function + dispatch: + SparseCPU, SparseCUDA, SparseMPS, SparseMeta: copy_sparse_ + autogen: copy_sparse_to_sparse, copy_sparse_to_sparse.out + +# By adding the AutogradNestedTensor this makes this function CompositeImplicit-like for nested tensors +- func: unbind.int(Tensor(a -> *) self, int dim=0) -> Tensor(a)[] + variants: function, method + dispatch: + CompositeExplicitAutograd: unbind + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_unbind + +- func: unbind.Dimname(Tensor(a -> *) self, Dimname dim) -> Tensor(a)[] + variants: function, method + +- func: to_sparse.sparse_dim(Tensor self, int sparse_dim) -> Tensor + variants: method + +# Special case of to_sparse.sparse_dim with custom derivative +- func: _to_sparse.sparse_dim(Tensor self, int sparse_dim) -> Tensor + variants: method + dispatch: + CPU, CUDA, MPS: dense_to_sparse + SparseCPU, SparseCUDA, SparseMPS: sparse_coo_to_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta, SparseCsrMPS: sparse_compressed_to_sparse + autogen: _to_sparse.sparse_dim_out + +- func: to_sparse(Tensor self, *, Layout? layout=None, int[2]? blocksize=None, int? dense_dim=None) -> Tensor + variants: method + +# Special case of to_sparse with custom derivative +- func: _to_sparse(Tensor self, *, Layout? layout=None, int[2]? blocksize=None, int? dense_dim=None) -> Tensor + variants: method + dispatch: + CPU, CUDA, MPS: dense_to_sparse + SparseCPU, SparseCUDA, SparseMPS: sparse_coo_to_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sparse_compressed_to_sparse + autogen: _to_sparse.out + +- func: to_sparse_csr(Tensor self, int? dense_dim=None) -> Tensor + variants: method + +# Special case of to_sparse_csr with custom derivative +- func: _to_sparse_csr(Tensor self, int? dense_dim=None) -> Tensor + variants: method + dispatch: + CPU, CUDA: dense_to_sparse_csr + SparseCPU, SparseCUDA: coo_to_sparse_csr + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sparse_compressed_to_sparse_csr + autogen: _to_sparse_csr.out + +- func: to_sparse_csc(Tensor self, int? dense_dim=None) -> Tensor + variants: method + +# Special case of to_sparse_csc with custom derivative +- func: _to_sparse_csc(Tensor self, int? dense_dim=None) -> Tensor + variants: method + dispatch: + CPU, CUDA: dense_to_sparse_csc + SparseCPU, SparseCUDA: coo_to_sparse_csc + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sparse_compressed_to_sparse_csc + autogen: _to_sparse_csc.out + +- func: to_sparse_bsr(Tensor self, int[2] blocksize, int? dense_dim=None) -> Tensor + variants: method + +# Special case of to_sparse_bsr with custom derivative +- func: _to_sparse_bsr(Tensor self, int[2] blocksize, int? dense_dim=None) -> Tensor + variants: method + dispatch: + CPU, CUDA: dense_to_sparse_bsr + SparseCPU, SparseCUDA: coo_to_sparse_bsr + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sparse_compressed_to_sparse_bsr + autogen: _to_sparse_bsr.out + +- func: to_sparse_bsc(Tensor self, int[2] blocksize, int? dense_dim=None) -> Tensor + variants: method + +# Special case of to_sparse_bsc with custom derivative +- func: _to_sparse_bsc(Tensor self, int[2] blocksize, int? dense_dim=None) -> Tensor + variants: method + dispatch: + CPU, CUDA: dense_to_sparse_bsc + SparseCPU, SparseCUDA: coo_to_sparse_bsc + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sparse_compressed_to_sparse_bsc + autogen: _to_sparse_bsc.out + +- func: _to_sparse_semi_structured(Tensor dense) -> (Tensor, Tensor) + variants: function + dispatch: + CUDA: _to_sparse_semi_structured + +- func: to_mkldnn(Tensor self, ScalarType? dtype=None) -> Tensor + variants: method + dispatch: + CPU: dense_to_mkldnn + autogen: to_mkldnn.out + +- func: mkldnn_reorder_conv2d_weight(Tensor self, SymInt[2] padding=0, SymInt[2] stride=1, SymInt[2] dilation=1, SymInt groups=1, SymInt[]? input_size=None) -> Tensor + variants: function + python_module: nn + dispatch: + MkldnnCPU: mkldnn_reorder_conv2d_weight + autogen: mkldnn_reorder_conv2d_weight.out + +- func: mkldnn_reorder_conv3d_weight(Tensor self, SymInt[3] padding=0, SymInt[3] stride=1, SymInt[3] dilation=1, SymInt groups=1, SymInt[]? input_size=None) -> Tensor + variants: function + python_module: nn + dispatch: + MkldnnCPU: mkldnn_reorder_conv3d_weight + autogen: mkldnn_reorder_conv3d_weight.out + +- func: to_mkldnn_backward(Tensor grad, Tensor input) -> Tensor + +- func: quantize_per_tensor_dynamic(Tensor self, ScalarType dtype, bool reduce_range) -> Tensor + variants: function + dispatch: + CPU, CUDA: quantize_per_tensor_dynamic + autogen: quantize_per_tensor_dynamic.out + +- func: quantize_per_tensor(Tensor self, float scale, int zero_point, ScalarType dtype) -> Tensor + variants: function + dispatch: + CPU, CUDA: quantize_per_tensor + autogen: quantize_per_tensor.out + +- func: quantize_per_tensor.tensor_qparams(Tensor self, Tensor scale, Tensor zero_point, ScalarType dtype) -> Tensor + variants: function + dispatch: + CPU, CUDA: quantize_per_tensor_tensor_qparams + autogen: quantize_per_tensor.tensor_qparams_out + +- func: quantize_per_tensor.tensors(Tensor[] tensors, Tensor scales, Tensor zero_points, ScalarType dtype) -> Tensor[] + variants: function + dispatch: + CPU: quantize_per_tensor_list_cpu + autogen: quantize_per_tensor.tensors_out + +- func: quantize_per_channel(Tensor self, Tensor scales, Tensor zero_points, int axis, ScalarType dtype) -> Tensor + variants: function + dispatch: + CPU, CUDA: quantize_per_channel + autogen: quantize_per_channel.out + +- func: dequantize.self(Tensor self) -> Tensor + variants: function, method + dispatch: + CPU, CUDA: dequantize_cpu_or_cuda + QuantizedCPU, QuantizedCUDA: dequantize_quantized + autogen: dequantize.self_out + +- func: dequantize.tensors(Tensor[] tensors) -> Tensor[] + variants: function + dispatch: + QuantizedCPU: dequantize_tensors_quantized_cpu + autogen: dequantize.tensors_out + +- func: q_scale(Tensor self) -> float + variants: function, method + dispatch: + QuantizedCPU, QuantizedCUDA: q_scale_quant + +- func: q_zero_point(Tensor self) -> int + variants: function, method + dispatch: + QuantizedCPU, QuantizedCUDA: q_zero_point_quant + +- func: q_per_channel_scales(Tensor self) -> Tensor + variants: function, method + dispatch: + QuantizedCPU, QuantizedCUDA: q_per_channel_scales + autogen: q_per_channel_scales.out + +- func: q_per_channel_zero_points(Tensor self) -> Tensor + variants: function, method + dispatch: + QuantizedCPU, QuantizedCUDA: q_per_channel_zero_points + autogen: q_per_channel_zero_points.out + +- func: q_per_channel_axis(Tensor self) -> int + variants: function, method + dispatch: + QuantizedCPU, QuantizedCUDA: q_per_channel_axis + +- func: int_repr(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + QuantizedCPU: int_repr_quantized_cpu + QuantizedCUDA: int_repr_quantized_cuda + autogen: int_repr.out + +- func: _make_per_tensor_quantized_tensor(Tensor self, float scale, int zero_point) -> Tensor + dispatch: + CPU: make_per_tensor_quantized_tensor_cpu + CUDA: make_per_tensor_quantized_tensor_cuda + autogen: _make_per_tensor_quantized_tensor.out + +- func: _make_per_channel_quantized_tensor(Tensor self, Tensor scale, Tensor zero_point, int axis) -> Tensor + dispatch: + CPU: make_per_channel_quantized_tensor_cpu + CUDA: make_per_channel_quantized_tensor_cuda + autogen: _make_per_channel_quantized_tensor.out + +- func: qscheme(Tensor self) -> QScheme + variants: method + dispatch: + QuantizedCPU, QuantizedCUDA: qscheme_quant + +- func: fake_quantize_per_tensor_affine(Tensor self, float scale, int zero_point, int quant_min, int quant_max) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + +- func: fake_quantize_per_tensor_affine.tensor_qparams(Tensor self, Tensor scale, Tensor zero_point, int quant_min, int quant_max) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + +- func: fake_quantize_per_tensor_affine_cachemask(Tensor self, float scale, int zero_point, int quant_min, int quant_max) -> (Tensor output, Tensor mask) + variants: function + dispatch: + CPU, CUDA: fake_quantize_per_tensor_affine_cachemask + autogen: fake_quantize_per_tensor_affine_cachemask.out + +- func: _fake_quantize_per_tensor_affine_cachemask_tensor_qparams(Tensor self, Tensor scale, Tensor zero_point, Tensor fake_quant_enabled, int quant_min, int quant_max) -> (Tensor output, Tensor mask) + variants: function + dispatch: + CPU, CUDA: _fake_quantize_per_tensor_affine_cachemask_tensor_qparams + autogen: _fake_quantize_per_tensor_affine_cachemask_tensor_qparams.out + +- func: fake_quantize_per_tensor_affine_cachemask_backward(Tensor grad, Tensor mask) -> Tensor + variants: function + +- func: _fake_quantize_learnable_per_tensor_affine(Tensor self, Tensor scale, Tensor zero_point, int quant_min, int quant_max, float grad_factor=1.0) -> Tensor + variants: function + dispatch: + CPU, CUDA: _fake_quantize_learnable_per_tensor_affine + autogen: _fake_quantize_learnable_per_tensor_affine.out + +- func: _fake_quantize_learnable_per_tensor_affine_backward(Tensor grad, Tensor self, Tensor scale, Tensor zero_point, int quant_min, int quant_max, float grad_factor=1.0) -> (Tensor, Tensor, Tensor) + variants: function + dispatch: + CPU, CUDA: _fake_quantize_learnable_per_tensor_affine_backward + +- func: fake_quantize_per_channel_affine(Tensor self, Tensor scale, Tensor zero_point, int axis, int quant_min, int quant_max) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + +- func: fake_quantize_per_channel_affine_cachemask(Tensor self, Tensor scale, Tensor zero_point, int axis, int quant_min, int quant_max) -> (Tensor output, Tensor mask) + variants: function + dispatch: + CPU, CUDA: fake_quantize_per_channel_affine_cachemask + autogen: fake_quantize_per_channel_affine_cachemask.out + +- func: fake_quantize_per_channel_affine_cachemask_backward(Tensor grad, Tensor mask) -> Tensor + variants: function + +- func: _fake_quantize_learnable_per_channel_affine(Tensor self, Tensor scale, Tensor zero_point, int axis, int quant_min, int quant_max, float grad_factor=1.0) -> Tensor + variants: function + dispatch: + CPU, CUDA: _fake_quantize_learnable_per_channel_affine + autogen: _fake_quantize_learnable_per_channel_affine.out + +- func: _fake_quantize_learnable_per_channel_affine_backward(Tensor grad, Tensor self, Tensor scale, Tensor zero_point, int axis, int quant_min, int quant_max, float grad_factor=1.0) -> (Tensor, Tensor, Tensor) + variants: function + dispatch: + CPU, CUDA: _fake_quantize_learnable_per_channel_affine_backward + +- func: fused_moving_avg_obs_fake_quant(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> Tensor + variants: function + +- func: _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask) + dispatch: + CPU: fused_moving_avg_obs_fake_quant_cpu + CUDA: fused_moving_avg_obs_fake_quant_cuda + autogen: _fused_moving_avg_obs_fq_helper_functional, _fused_moving_avg_obs_fq_helper.out + +- func: _choose_qparams_per_tensor(Tensor self, bool reduce_range=False) -> (float, int) + variants: function + +- func: _saturate_weight_to_fp16(Tensor weight) -> Tensor + variants: function + +- func: choose_qparams_optimized(Tensor input, int numel, int n_bins, float ratio, int bit_width) -> (Tensor, Tensor) + variants: function + +- func: _autocast_to_reduced_precision(Tensor(a) self, bool cuda_enabled, bool cpu_enabled, ScalarType cuda_dtype, ScalarType cpu_dtype) -> Tensor(a) + variants: method + device_guard: False + +- func: _autocast_to_full_precision(Tensor(a) self, bool cuda_enabled, bool cpu_enabled) -> Tensor(a) + variants: method + device_guard: False + +- func: _to_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: _to_copy + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: _to_copy_nested + autogen: _to_copy.out + tags: core + +# to(Device) must not exist because all constructors of Device also works for +# TensorOptions. Otherwise, an ambiguity error is thrown. +# See NOTE [ TensorOptions Constructors ]. +- func: to.dtype_layout(Tensor(a) self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a) + variants: method + device_check: NoCheck + device_guard: False + +- func: to.device(Tensor(a) self, Device device, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a) + variants: method + device_check: NoCheck + device_guard: False + +- func: to.dtype(Tensor(a) self, ScalarType dtype, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a) + variants: method + device_check: NoCheck + device_guard: False + +- func: to.other(Tensor(a) self, Tensor other, bool non_blocking=False, bool copy=False, MemoryFormat? memory_format=None) -> Tensor(a) + variants: method + device_check: NoCheck + device_guard: False + +- func: meshgrid(Tensor[] tensors) -> Tensor[] + +# TODO: Two weeks after this lands, combine these two overloads, +# making "indexing" optional. These are temporarily distinct for +# forward-compatibility reasons. +- func: meshgrid.indexing(Tensor[] tensors, *, str indexing) -> Tensor[] + +- func: cartesian_prod(Tensor[] tensors) -> Tensor + variants: function + tags: maybe_aliasing_or_mutating + +- func: combinations(Tensor self, int r=2, bool with_replacement=False) -> Tensor + variants: function + +- func: item(Tensor self) -> Scalar + tags: data_dependent_output + variants: method + +- func: result_type.Tensor(Tensor tensor, Tensor other) -> ScalarType + variants: function + +- func: result_type.Scalar(Tensor tensor, Scalar other) -> ScalarType + variants: function + +- func: result_type.Scalar_Tensor(Scalar scalar, Tensor tensor) -> ScalarType + variants: function + +- func: result_type.Scalar_Scalar(Scalar scalar1, Scalar scalar2) -> ScalarType + +- func: can_cast(ScalarType from_, ScalarType to) -> bool + variants: function + +- func: promote_types(ScalarType type1, ScalarType type2) -> ScalarType + variants: function + +# NB: Does NOT check precondition that numel == 1 +- func: _local_scalar_dense(Tensor self) -> Scalar + tags: [core, data_dependent_output] + dispatch: + CPU: _local_scalar_dense_cpu + CUDA: _local_scalar_dense_cuda + MPS: _local_scalar_dense_mps + variants: function + +# MPS LSTM implementation + +- func: _lstm_mps(Tensor input, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) + dispatch: + MPS: _lstm_mps + autogen: _lstm_mps.out + tags: nondeterministic_seeded + +- func: lstm_mps_backward(Tensor? grad_y, Tensor? grad_hy, Tensor? grad_cy, Tensor z_state, Tensor cell_state_fwd, Tensor input, Tensor layersOutputs, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor[], Tensor[]) + dispatch: + MPS: lstm_mps_backward + autogen: lstm_mps_backward.out + + +# Fused RNN kernels +- func: _thnn_fused_lstm_cell(Tensor input_gates, Tensor hidden_gates, Tensor cx, Tensor? input_bias=None, Tensor? hidden_bias=None) -> (Tensor, Tensor, Tensor) + dispatch: + CUDA: _thnn_fused_lstm_cell_cuda + autogen: _thnn_fused_lstm_cell.out + +# NB: The composite version of this function below is a simple wrapper that duplicates some of the outputs +# It is necessary to avoid triggering TensorImpl use count checks in debug mode +# NB: this is function is NOT differentiable +- func: _thnn_fused_lstm_cell_backward_impl(Tensor? grad_hy, Tensor? grad_cy, Tensor cx, Tensor cy, Tensor workspace, bool has_bias) -> (Tensor, Tensor, Tensor) + dispatch: + CUDA: _thnn_fused_lstm_cell_backward_impl_cuda + autogen: _thnn_fused_lstm_cell_backward_impl.out + +- func: _thnn_fused_lstm_cell_backward(Tensor? grad_hy, Tensor? grad_cy, Tensor cx, Tensor cy, Tensor workspace, bool has_bias) -> (Tensor, Tensor, Tensor, Tensor, Tensor) + +- func: _thnn_differentiable_lstm_cell_backward(Tensor? grad_hy, Tensor? grad_cy, Tensor input_gates, Tensor hidden_gates, Tensor? input_bias, Tensor? hidden_bias, Tensor cx, Tensor cy) -> (Tensor, Tensor, Tensor, Tensor, Tensor) + +- func: _thnn_fused_gru_cell(Tensor input_gates, Tensor hidden_gates, Tensor hx, Tensor? input_bias=None, Tensor? hidden_bias=None) -> (Tensor, Tensor) + dispatch: + CUDA: _thnn_fused_gru_cell_cuda + autogen: _thnn_fused_gru_cell.out + +- func: _thnn_fused_gru_cell_backward(Tensor grad_hy, Tensor workspace, bool has_bias) -> (Tensor, Tensor, Tensor, Tensor, Tensor) + dispatch: + CUDA: _thnn_fused_gru_cell_backward_cuda + autogen: _thnn_fused_gru_cell_backward.out + +- func: _thnn_differentiable_gru_cell_backward(Tensor grad_hy, Tensor input_gates, Tensor hidden_gates, Tensor hx, Tensor? input_bias, Tensor? hidden_bias) -> (Tensor, Tensor, Tensor, Tensor, Tensor) + +# RNN cells and layers +- func: lstm.input(Tensor input, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor, Tensor) + tags: nondeterministic_seeded + +- func: lstm.data(Tensor data, Tensor batch_sizes, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional) -> (Tensor, Tensor, Tensor) + tags: nondeterministic_seeded + +- func: gru.input(Tensor input, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor) + tags: nondeterministic_seeded + +- func: gru.data(Tensor data, Tensor batch_sizes, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional) -> (Tensor, Tensor) + tags: nondeterministic_seeded + +- func: rnn_tanh.input(Tensor input, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor) + tags: nondeterministic_seeded + +- func: rnn_tanh.data(Tensor data, Tensor batch_sizes, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional) -> (Tensor, Tensor) + tags: nondeterministic_seeded + +- func: rnn_relu.input(Tensor input, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor) + tags: nondeterministic_seeded + +- func: rnn_relu.data(Tensor data, Tensor batch_sizes, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional) -> (Tensor, Tensor) + tags: nondeterministic_seeded + +- func: lstm_cell(Tensor input, Tensor[] hx, Tensor w_ih, Tensor w_hh, Tensor? b_ih=None, Tensor? b_hh=None) -> (Tensor, Tensor) + +- func: gru_cell(Tensor input, Tensor hx, Tensor w_ih, Tensor w_hh, Tensor? b_ih=None, Tensor? b_hh=None) -> Tensor + +- func: rnn_tanh_cell(Tensor input, Tensor hx, Tensor w_ih, Tensor w_hh, Tensor? b_ih=None, Tensor? b_hh=None) -> Tensor + +- func: rnn_relu_cell(Tensor input, Tensor hx, Tensor w_ih, Tensor w_hh, Tensor? b_ih=None, Tensor? b_hh=None) -> Tensor + +# Quantized RNN layer registration has been moved to C10 dispatch in `RNN.cpp` + +# Quantized RNN layers +# - func: quantized_lstm(Tensor input, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first, *, ScalarType? dtype=None, bool use_dynamic=False) -> (Tensor, Tensor, Tensor) + + +# - func: quantized_lstm.data(Tensor data, Tensor batch_sizes, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, *, ScalarType? dtype=None, bool use_dynamic=False) -> (Tensor, Tensor, Tensor) + + +# Quantized GRU layers + +# - func: quantized_gru.input(Tensor input, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor) +# + +# - func: quantized_gru.data(Tensor data, Tensor batch_sizes, Tensor hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional) -> (Tensor, Tensor) +# + +# Quantized RNN cells +- func: quantized_lstm_cell(Tensor input, Tensor[] hx, Tensor w_ih, Tensor w_hh, Tensor b_ih, Tensor b_hh, Tensor packed_ih, Tensor packed_hh, Tensor col_offsets_ih, Tensor col_offsets_hh, Scalar scale_ih, Scalar scale_hh, Scalar zero_point_ih, Scalar zero_point_hh) -> (Tensor, Tensor) + +- func: quantized_gru_cell(Tensor input, Tensor hx, Tensor w_ih, Tensor w_hh, Tensor b_ih, Tensor b_hh, Tensor packed_ih, Tensor packed_hh, Tensor col_offsets_ih, Tensor col_offsets_hh, Scalar scale_ih, Scalar scale_hh, Scalar zero_point_ih, Scalar zero_point_hh) -> Tensor + +- func: quantized_rnn_relu_cell(Tensor input, Tensor hx, Tensor w_ih, Tensor w_hh, Tensor b_ih, Tensor b_hh, Tensor packed_ih, Tensor packed_hh, Tensor col_offsets_ih, Tensor col_offsets_hh, Scalar scale_ih, Scalar scale_hh, Scalar zero_point_ih, Scalar zero_point_hh) -> Tensor + +- func: quantized_rnn_tanh_cell(Tensor input, Tensor hx, Tensor w_ih, Tensor w_hh, Tensor b_ih, Tensor b_hh, Tensor packed_ih, Tensor packed_hh, Tensor col_offsets_ih, Tensor col_offsets_hh, Scalar scale_ih, Scalar scale_hh, Scalar zero_point_ih, Scalar zero_point_hh) -> Tensor + +# PackedSequence utilities +- func: _pack_padded_sequence(Tensor input, Tensor lengths, bool batch_first) -> (Tensor, Tensor) + dispatch: + CompositeExplicitAutograd: _pack_padded_sequence + autogen: _pack_padded_sequence.out + +- func: _pack_padded_sequence_backward(Tensor grad, SymInt[] input_size, Tensor batch_sizes, bool batch_first) -> Tensor + dispatch: + CompositeImplicitAutograd: _pack_padded_sequence_backward_symint + +- func: _pad_packed_sequence(Tensor data, Tensor batch_sizes, bool batch_first, Scalar padding_value, int total_length) -> (Tensor, Tensor) + +# wrappers for legacy TH methods + +- func: set_.source_Storage(Tensor(a!) self, Storage source) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + dispatch: + CPU, CUDA, Meta, MPS: set_ + autogen: set.source_Storage, set.source_Storage_out + tags: inplace_view + +- func: set_.source_Storage_storage_offset(Tensor(a!) self, Storage source, SymInt storage_offset, SymInt[] size, SymInt[] stride=[]) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + dispatch: + CPU: set_storage_cpu_ + Meta: set_storage_meta__symint + CUDA: set_storage_cuda_ + MPS: set_storage_mps_ + QuantizedCPU, QuantizedCUDA: set_storage_quantized_ + autogen: set.source_Storage_storage_offset, set.source_Storage_storage_offset_out + tags: inplace_view + +- func: set_.source_Tensor_storage_offset(Tensor(a!) self, Tensor source, SymInt storage_offset, SymInt[] size, SymInt[] stride=[]) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: set__symint + tags: inplace_view + +- func: set_.source_Tensor(Tensor(a!) self, Tensor source) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + dispatch: + CPU, CUDA, Meta, MPS: set_tensor_ + autogen: set.source_Tensor, set.source_Tensor_out + tags: inplace_view + +- func: set_(Tensor(a!) self) -> Tensor(a!) + variants: method + dispatch: + CPU: set_cpu_ + CUDA: set_cuda_ + Meta: set_meta_ + MPS: set_mps_ + autogen: set, set.out + tags: inplace_view + +# Not making it CompositeImplicitAutograd because lift +# should be a primitive w.r.t. functorch + +# TODO: this should have a view annotation +# TODO: shouldn't be a method +- func: lift(Tensor self) -> Tensor + dispatch: + CompositeExplicitAutograd: lift + autogen: lift.out + +# lift_fresh is called with an argument that is guaranteed to be +# fresh (i.e., newly allocated). This is ONLY called from a +# torch.tensor call; if you FX trace a lift_fresh, you are obligated +# to convert this into a lift_fresh_copy (because FX will violate the +# freshness invariant when tracing). +- func: lift_fresh(Tensor(a) self) -> Tensor(a) + dispatch: + CompositeExplicitAutograd: lift_fresh + +# Like lift, but it clones the input. +- func: lift_fresh_copy(Tensor self) -> Tensor + tags: view_copy + dispatch: + CompositeExplicitAutogradNonFunctional: lift_fresh_copy + autogen: lift_fresh_copy.out + +- func: is_set_to(Tensor self, Tensor tensor) -> bool + variants: method + device_check: NoCheck + device_guard: False + dispatch: + CPU, CUDA, MPS: is_set_to + +- func: masked_fill_.Scalar(Tensor(a!) self, Tensor mask, Scalar value) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CPU: masked_fill__cpu + CUDA: masked_fill__cuda + QuantizedCPU: masked_fill__quantized_cpu + QuantizedCUDA: masked_fill__quantized_cuda + MPS: masked_fill__mps + autogen: masked_fill.Scalar_out + +- func: masked_fill.Scalar(Tensor self, Tensor mask, Scalar value) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: masked_fill + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_masked_fill + tags: pointwise + +- func: masked_fill_.Tensor(Tensor(a!) self, Tensor mask, Tensor value) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CPU: masked_fill__cpu + CUDA: masked_fill__cuda + QuantizedCPU: masked_fill__quantized_cpu + QuantizedCUDA: masked_fill__quantized_cuda + MPS: masked_fill__mps + autogen: masked_fill.Tensor_out + +- func: masked_fill.Tensor(Tensor self, Tensor mask, Tensor value) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: masked_fill + +- func: masked_scatter_(Tensor(a!) self, Tensor mask, Tensor source) -> Tensor(a!) + variants: method + dispatch: + CPU: masked_scatter__cpu + CUDA: masked_scatter__cuda + MPS: masked_scatter__mps + autogen: masked_scatter.out + +- func: masked_scatter(Tensor self, Tensor mask, Tensor source) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: masked_scatter + tags: core + +- func: masked_scatter_backward(Tensor grad_output, Tensor mask, SymInt[] sizes) -> Tensor + dispatch: + CompositeExplicitAutograd: masked_scatter_backward_symint + +- func: _masked_softmax(Tensor self, Tensor mask, int? dim=None, int? mask_type=None) -> Tensor + dispatch: + CUDA: masked_softmax_cuda + CPU: masked_softmax_cpu + autogen: _masked_softmax.out + +- func: _masked_softmax_backward(Tensor grad_output, Tensor output, Tensor mask, int? dim=None) -> Tensor + dispatch: + CUDA: masked_softmax_backward_cuda + CPU: masked_softmax_backward_cpu + autogen: _masked_softmax_backward.out + +- func: view(Tensor(a) self, SymInt[] size) -> Tensor(a) + variants: method + device_check: NoCheck + device_guard: False + dispatch: + ZeroTensor, Meta, CPU, CUDA, QuantizedCPU, QuantizedCUDA, MPS, MTIA: view + MkldnnCPU: mkldnn_view + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: view_nested + tags: core + +# Warning: If you want to change the name or overload name of this +# operator, you might also want to change the `isBlockListedSchema` +# function in `torch/csrc/jit/frontend/schema_catching.cpp`. +# The name and overload name of this operator is hardcoded in that +# function in order to workaround a bug: +# https://github.com/pytorch/pytorch/issues/47964 +- func: view.dtype(Tensor(a) self, ScalarType dtype) -> Tensor(a) + variants: method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: view_dtype + +- func: put_(Tensor(a!) self, Tensor index, Tensor source, bool accumulate=False) -> Tensor(a!) + variants: method + dispatch: + CPU, CUDA: put_ + autogen: put.out + +- func: put(Tensor self, Tensor index, Tensor source, bool accumulate=False) -> Tensor + variants: function, method + dispatch: + CompositeExplicitAutograd: put + +- func: index_add.out(Tensor self, int dim, Tensor index, Tensor source, *, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + precomputed: + - dim -> int dim + dispatch: + CPU: index_add_cpu_out + CUDA: index_add_cuda_out + MPS: index_add_mps_out + +- func: index_add_(Tensor(a!) self, int dim, Tensor index, Tensor source, *, Scalar alpha=1) -> Tensor(a!) + structured_delegate: index_add.out + variants: method + +- func: index_add(Tensor self, int dim, Tensor index, Tensor source, *, Scalar alpha=1) -> Tensor + structured_delegate: index_add.out + variants: function, method + +- func: index_add.dimname(Tensor self, Dimname dim, Tensor index, Tensor source, *, Scalar alpha=1) -> Tensor + variants: function, method + +- func: index_reduce.out(Tensor self, int dim, Tensor index, Tensor source, str reduce, *, bool include_self=True, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + precomputed: + - dim -> int dim + dispatch: + CPU: index_reduce_cpu_out + CUDA: index_reduce_cuda_out + MPS: index_reduce_mps_out + +- func: index_reduce_(Tensor(a!) self, int dim, Tensor index, Tensor source, str reduce, *, bool include_self=True) -> Tensor(a!) + structured_delegate: index_reduce.out + variants: method + +- func: index_reduce(Tensor self, int dim, Tensor index, Tensor source, str reduce, *, bool include_self=True) -> Tensor + structured_delegate: index_reduce.out + variants: function, method + +- func: index_fill_.int_Scalar(Tensor(a!) self, int dim, Tensor index, Scalar value) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CPU, CUDA, MPS: index_fill_ + autogen: index_fill.int_Scalar_out + +- func: index_fill.int_Scalar(Tensor self, int dim, Tensor index, Scalar value) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: index_fill + +- func: index_fill_.int_Tensor(Tensor(a!) self, int dim, Tensor index, Tensor value) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CPU, CUDA, MPS: index_fill_ + autogen: index_fill.int_Tensor_out + +- func: index_fill.int_Tensor(Tensor self, int dim, Tensor index, Tensor value) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + dispatch: + CompositeExplicitAutograd: index_fill + +- func: index_fill_.Dimname_Scalar(Tensor(a!) self, Dimname dim, Tensor index, Scalar value) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + +- func: index_fill_.Dimname_Tensor(Tensor(a!) self, Dimname dim, Tensor index, Tensor value) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + +- func: index_fill.Dimname_Scalar(Tensor self, Dimname dim, Tensor index, Scalar value) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + +- func: index_fill.Dimname_Tensor(Tensor self, Dimname dim, Tensor index, Tensor value) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + +- func: scatter.src(Tensor self, int dim, Tensor index, Tensor src) -> Tensor + structured_delegate: scatter.src_out + variants: function, method + tags: core + +- func: scatter_.src(Tensor(a!) self, int dim, Tensor index, Tensor src) -> Tensor(a!) + structured_delegate: scatter.src_out + variants: method + +- func: scatter.src_out(Tensor self, int dim, Tensor index, Tensor src, *, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + dispatch: + CPU, CUDA: scatter_src_out + MPS: scatter_src_out_mps + +- func: scatter.value(Tensor self, int dim, Tensor index, Scalar value) -> Tensor + structured_delegate: scatter.value_out + variants: function, method + tags: core + +- func: scatter_.value(Tensor(a!) self, int dim, Tensor index, Scalar value) -> Tensor(a!) + structured_delegate: scatter.value_out + variants: method + +- func: scatter.value_out(Tensor self, int dim, Tensor index, Scalar value, *, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + dispatch: + CPU, CUDA: scatter_value_out + MPS: scatter_value_out_mps + +- func: scatter.reduce(Tensor self, int dim, Tensor index, Tensor src, *, str reduce) -> Tensor + structured_delegate: scatter.reduce_out + variants: function, method + +- func: scatter_.reduce(Tensor(a!) self, int dim, Tensor index, Tensor src, *, str reduce) -> Tensor(a!) + structured_delegate: scatter.reduce_out + variants: method + +- func: scatter.reduce_out(Tensor self, int dim, Tensor index, Tensor src, *, str reduce, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + dispatch: + CPU, CUDA: scatter_reduce_out + MPS: scatter_reduce_out_mps + +- func: scatter.value_reduce(Tensor self, int dim, Tensor index, Scalar value, *, str reduce) -> Tensor + structured_delegate: scatter.value_reduce_out + variants: function, method + +- func: scatter_.value_reduce(Tensor(a!) self, int dim, Tensor index, Scalar value, *, str reduce) -> Tensor(a!) + structured_delegate: scatter.value_reduce_out + variants: method + +- func: scatter.value_reduce_out(Tensor self, int dim, Tensor index, Scalar value, *, str reduce, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + dispatch: + CPU, CUDA: scatter_value_reduce_out + MPS: scatter_value_reduce_out_mps + +- func: scatter.dimname_src(Tensor self, Dimname dim, Tensor index, Tensor src) -> Tensor + variants: function, method + +- func: scatter.dimname_value(Tensor self, Dimname dim, Tensor index, Scalar value) -> Tensor + variants: function, method + +- func: scatter_add(Tensor self, int dim, Tensor index, Tensor src) -> Tensor + structured_delegate: scatter_add.out + variants: function, method + tags: core + +- func: scatter_add_(Tensor(a!) self, int dim, Tensor index, Tensor src) -> Tensor(a!) + structured_delegate: scatter_add.out + variants: method + +- func: scatter_add.out(Tensor self, int dim, Tensor index, Tensor src, *, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + dispatch: + CPU, CUDA: scatter_add + MPS: scatter_add_mps_out + +- func: scatter_add.dimname(Tensor self, Dimname dim, Tensor index, Tensor src) -> Tensor + variants: function, method + +- func: scatter_reduce.two(Tensor self, int dim, Tensor index, Tensor src, str reduce, *, bool include_self=True) -> Tensor + structured_delegate: scatter_reduce.two_out + variants: function, method + tags: core + +- func: scatter_reduce_.two(Tensor(a!) self, int dim, Tensor index, Tensor src, str reduce, *, bool include_self=True) -> Tensor(a!) + structured_delegate: scatter_reduce.two_out + variants: method + +- func: scatter_reduce.two_out(Tensor self, int dim, Tensor index, Tensor src, str reduce, *, bool include_self=True, Tensor(a!) out) -> Tensor(a!) + structured: True + variants: function + dispatch: + CPU, CUDA, MPS: scatter_reduce_two + +- func: eq_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + structured_delegate: eq.Scalar_out + device_check: NoCheck # TensorIterator + variants: method + +- func: eq_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: eq.Tensor_out + device_check: NoCheck # TensorIterator + variants: method + +- func: bitwise_and.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + variants: function + dispatch: + CPU, CUDA, MTIA: bitwise_and_out + MPS: bitwise_and_out_mps + tags: pointwise + +- func: bitwise_and.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_and_out + tags: pointwise + +- func: bitwise_and.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CompositeExplicitAutograd: bitwise_and + tags: [core, pointwise] + +- func: bitwise_and.Scalar_Tensor(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_and + autogen: bitwise_and.Scalar_Tensor_out + tags: pointwise + +- func: bitwise_and.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + structured_delegate: bitwise_and.Tensor_out + tags: [core, pointwise] + +- func: bitwise_and_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: bitwise_and_ + tags: pointwise + +- func: bitwise_and_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: bitwise_and.Tensor_out + tags: pointwise + +- func: __and__.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + +- func: __and__.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + +- func: __iand__.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + +- func: __iand__.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + +- func: bitwise_or.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + variants: function + dispatch: + CPU, CUDA, MTIA: bitwise_or_out + MPS: bitwise_or_out_mps + tags: pointwise + +- func: bitwise_or.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_or_out + tags: pointwise + +- func: bitwise_or.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CompositeExplicitAutograd: bitwise_or + tags: [core, pointwise] + +- func: bitwise_or.Scalar_Tensor(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_or + autogen: bitwise_or.Scalar_Tensor_out + tags: pointwise + +- func: bitwise_or.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + structured_delegate: bitwise_or.Tensor_out + tags: [core, pointwise] + +- func: bitwise_or_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: bitwise_or_ + tags: pointwise + +- func: bitwise_or_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: bitwise_or.Tensor_out + tags: pointwise + +- func: __or__.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + +- func: __or__.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + +- func: __ior__.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + +- func: __ior__.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + +- func: bitwise_xor.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + variants: function + dispatch: + CPU, CUDA: bitwise_xor_out + MPS: bitwise_xor_out_mps + tags: pointwise + +- func: bitwise_xor.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_xor_out + tags: pointwise + +- func: bitwise_xor.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CompositeExplicitAutograd: bitwise_xor + tags: [core, pointwise] + +- func: bitwise_xor.Scalar_Tensor(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_xor + autogen: bitwise_xor.Scalar_Tensor_out + tags: pointwise + +- func: bitwise_xor.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + structured_delegate: bitwise_xor.Tensor_out + tags: [core, pointwise] + +- func: bitwise_xor_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: bitwise_xor_ + tags: pointwise + +- func: bitwise_xor_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: bitwise_xor.Tensor_out + tags: pointwise + +- func: __xor__.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + tags: pointwise + +- func: __xor__.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + tags: pointwise + +- func: __ixor__.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + tags: pointwise + +- func: __ixor__.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + tags: pointwise + +- func: __lshift__.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CPU, CUDA, MPS, MTIA: __lshift__ + tags: pointwise + +- func: __lshift__.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CPU, CUDA, MPS, MTIA: __lshift__ + tags: pointwise + +- func: __ilshift__.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CPU, CUDA, MPS: __ilshift__ + autogen: __lshift__.Scalar_out + tags: pointwise + +- func: __ilshift__.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CPU, CUDA, MPS: __ilshift__ + autogen: __lshift__.Tensor_out + tags: pointwise + +- func: bitwise_left_shift.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: bitwise_left_shift.Tensor_out + tags: pointwise + +- func: bitwise_left_shift_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: bitwise_left_shift.Tensor_out + tags: pointwise + +- func: bitwise_left_shift.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: bitwise_left_shift_out + tags: pointwise + +- func: bitwise_left_shift.Tensor_Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CompositeExplicitAutograd: bitwise_left_shift + tags: pointwise + +- func: bitwise_left_shift_.Tensor_Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: bitwise_left_shift_ + tags: pointwise + +- func: bitwise_left_shift.Tensor_Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_left_shift_out + tags: pointwise + +- func: bitwise_left_shift.Scalar_Tensor(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_left_shift + autogen: bitwise_left_shift.Scalar_Tensor_out + tags: pointwise + +- func: __rshift__.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CPU, CUDA, MPS, MTIA: __rshift__ + tags: pointwise + +- func: __rshift__.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CPU, CUDA, MPS, MTIA: __rshift__ + tags: pointwise + +- func: __irshift__.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CPU, CUDA, MPS: __irshift__ + autogen: __rshift__.Scalar_out + +- func: __irshift__.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CPU, CUDA, MPS: __irshift__ + autogen: __rshift__.Tensor_out + +- func: bitwise_right_shift.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function, method + structured_delegate: bitwise_right_shift.Tensor_out + tags: pointwise + +- func: bitwise_right_shift_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: bitwise_right_shift.Tensor_out + tags: pointwise + +- func: bitwise_right_shift.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: bitwise_right_shift_out + tags: pointwise + +- func: bitwise_right_shift.Tensor_Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CompositeExplicitAutograd: bitwise_right_shift + tags: pointwise + +- func: bitwise_right_shift_.Tensor_Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: bitwise_right_shift_ + tags: pointwise + +- func: bitwise_right_shift.Tensor_Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_right_shift_out + tags: pointwise + +- func: bitwise_right_shift.Scalar_Tensor(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CompositeExplicitAutograd: bitwise_right_shift + autogen: bitwise_right_shift.Scalar_Tensor_out + tags: pointwise + +- func: tril_(Tensor(a!) self, SymInt diagonal=0) -> Tensor(a!) + structured_delegate: tril.out + variants: method + +- func: triu_(Tensor(a!) self, SymInt diagonal=0) -> Tensor(a!) + structured_delegate: triu.out + variants: method + +- func: digamma_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: digamma.out + variants: method + tags: pointwise + +- func: lerp_.Scalar(Tensor(a!) self, Tensor end, Scalar weight) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: lerp.Scalar_out + tags: pointwise + +- func: lerp_.Tensor(Tensor(a!) self, Tensor end, Tensor weight) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: lerp.Tensor_out + tags: pointwise + +- func: addbmm_(Tensor(a!) self, Tensor batch1, Tensor batch2, *, Scalar beta=1, Scalar alpha=1) -> Tensor(a!) + variants: method + dispatch: + CPU, CUDA, XPU: addbmm_ + MPS: addbmm_mps_ + +- func: addbmm.out(Tensor self, Tensor batch1, Tensor batch2, *, Scalar beta=1, Scalar alpha=1, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, XPU: addbmm_out + MPS: addbmm_out_mps + +- func: addbmm(Tensor self, Tensor batch1, Tensor batch2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + variants: method, function + dispatch: + CPU, CUDA, XPU: addbmm + MPS: addbmm_mps + +- func: random_.from(Tensor(a!) self, int from, int? to, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + tags: nondeterministic_seeded + dispatch: + CPU, CUDA: random_ + Meta: random_meta_ + MPS: random_mps_ + autogen: random.from, random.from_out + +- func: random_.to(Tensor(a!) self, int to, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: nondeterministic_seeded + variants: method + dispatch: + CPU, CUDA: random_ + Meta: random_meta_ + MPS: random_mps_ + autogen: random.to, random.to_out + +- func: random_(Tensor(a!) self, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: nondeterministic_seeded + variants: method + dispatch: + CPU, CUDA: random_ + MPS: random_mps_ + Meta: random_meta_ + autogen: random, random.out + +- func: uniform_(Tensor(a!) self, float from=0, float to=1, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: nondeterministic_seeded + variants: method + dispatch: + CPU, CUDA: uniform_ + MPS: uniform_mps_ + Meta: uniform_meta_ + autogen: uniform, uniform.out + +- func: cauchy_(Tensor(a!) self, float median=0, float sigma=1, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + tags: nondeterministic_seeded + dispatch: + CPU, CUDA: cauchy_ + MPS: cauchy_mps_ + autogen: cauchy, cauchy.out + +- func: log_normal_(Tensor(a!) self, float mean=1, float std=2, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: nondeterministic_seeded + variants: method + dispatch: + CPU, CUDA: log_normal_ + MPS: log_normal_mps_ + autogen: log_normal, log_normal.out + +- func: exponential_(Tensor(a!) self, float lambd=1, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: nondeterministic_seeded + variants: method + dispatch: + CPU, CUDA: exponential_ + MPS: exponential_mps_ + autogen: exponential, exponential.out + +- func: geometric_(Tensor(a!) self, float p, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: nondeterministic_seeded + variants: method + dispatch: + CPU, CUDA: geometric_ + MPS: geometric_mps_ + + # wrappers for TH functions + autogen: geometric, geometric.out + +- func: diag.out(Tensor self, int diagonal=0, *, Tensor(a!) out) -> Tensor(a!) + +- func: diag(Tensor self, int diagonal=0) -> Tensor + variants: method, function + +- func: cross.out(Tensor self, Tensor other, int? dim=None, *, Tensor(a!) out) -> Tensor(a!) + +- func: cross(Tensor self, Tensor other, int? dim=None) -> Tensor + variants: method, function + +- func: triu.out(Tensor self, SymInt diagonal=0, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: triu_cpu + CUDA: triu_cuda + MPS: triu_mps_out + +- func: triu(Tensor self, SymInt diagonal=0) -> Tensor + structured_delegate: triu.out + variants: method, function + +- func: tril.out(Tensor self, SymInt diagonal=0, *, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: tril_cpu + CUDA: tril_cuda + MPS: tril_mps_out + +- func: tril(Tensor self, SymInt diagonal=0) -> Tensor + structured_delegate: tril.out + variants: method, function + +- func: tril_indices(int row, int col, int offset=0, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CPU: tril_indices_cpu + CUDA: tril_indices_cuda + MPS: tril_indices_mps + autogen: tril_indices.out + +- func: triu_indices(int row, int col, int offset=0, *, ScalarType? dtype=long, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CPU: triu_indices_cpu + CUDA: triu_indices_cuda + MPS: triu_indices_mps + autogen: triu_indices.out + +- func: trace(Tensor self) -> Tensor + variants: method, function + dispatch: + CPU: trace_cpu + CUDA: trace_cuda + MPS: trace_mps + autogen: trace.out + +- func: trace_backward(Tensor grad, SymInt[] sizes) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: trace_backward_symint + +- func: ne.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: ne_Scalar_out + MPS: ne_scalar_out_mps + QuantizedCPU: ne_out_quantized_cpu + tags: pointwise + +- func: ne.Scalar(Tensor self, Scalar other) -> Tensor + structured_delegate: ne.Scalar_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: ne_quantized_cpu + tags: [core, pointwise] + +- func: ne.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: ne_Tensor_out + MPS: ne_tensor_out_mps + QuantizedCPU: ne_out_quantized_cpu + tags: pointwise + +- func: ne.Tensor(Tensor self, Tensor other) -> Tensor + structured_delegate: ne.Tensor_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: ne_quantized_cpu + tags: [core, pointwise] + +- func: ne_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + structured_delegate: ne.Scalar_out + device_check: NoCheck # TensorIterator + variants: method + +- func: ne_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: ne.Tensor_out + device_check: NoCheck # TensorIterator + variants: method + +# not_equal, alias for torch.ne +- func: not_equal.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + +- func: not_equal.Scalar(Tensor self, Scalar other) -> Tensor + variants: method, function + +- func: not_equal.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + +- func: not_equal.Tensor(Tensor self, Tensor other) -> Tensor + variants: method, function + +- func: not_equal_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + variants: method + +- func: not_equal_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + variants: method + +- func: eq.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: eq_Scalar_out + MPS: eq_scalar_out_mps + QuantizedCPU: eq_out_quantized_cpu + tags: pointwise + +- func: eq.Scalar(Tensor self, Scalar other) -> Tensor + structured_delegate: eq.Scalar_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: eq_quantized_cpu + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: eq_scalar_nested + tags: [core, pointwise] + +- func: eq.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: eq_Tensor_out + MPS: eq_tensor_out_mps + QuantizedCPU: eq_out_quantized_cpu + tags: pointwise + +- func: eq.Tensor(Tensor self, Tensor other) -> Tensor + structured_delegate: eq.Tensor_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: eq_quantized_cpu + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: eq_tensor_nested + tags: [core, pointwise] + +- func: ge.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: ge_Scalar_out + MPS: ge_scalar_out_mps + QuantizedCPU: ge_out_quantized_cpu + tags: pointwise + +- func: ge.Scalar(Tensor self, Scalar other) -> Tensor + structured_delegate: ge.Scalar_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: ge_quantized_cpu + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: ge_scalar_nested + tags: [core, pointwise] + +- func: ge.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: ge_Tensor_out + MPS: ge_tensor_out_mps + QuantizedCPU: ge_out_quantized_cpu + tags: pointwise + +- func: ge.Tensor(Tensor self, Tensor other) -> Tensor + structured_delegate: ge.Tensor_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: ge_quantized_cpu + tags: [core, pointwise] + +- func: ge_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + structured_delegate: ge.Scalar_out + device_check: NoCheck # TensorIterator + variants: method + +- func: ge_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: ge.Tensor_out + device_check: NoCheck # TensorIterator + variants: method + +# greater_equal, alias for torch.ge +- func: greater_equal.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + +- func: greater_equal.Scalar(Tensor self, Scalar other) -> Tensor + variants: method, function + +- func: greater_equal.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + +- func: greater_equal.Tensor(Tensor self, Tensor other) -> Tensor + variants: method, function + +- func: greater_equal_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + variants: method + +- func: greater_equal_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + variants: method + +- func: le.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: le_Scalar_out + MPS: le_scalar_out_mps + QuantizedCPU: le_out_quantized_cpu + tags: pointwise + +- func: le.Scalar(Tensor self, Scalar other) -> Tensor + structured_delegate: le.Scalar_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: le_quantized_cpu + tags: [core, pointwise] + +- func: le.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: le_Tensor_out + MPS: le_tensor_out_mps + QuantizedCPU: le_out_quantized_cpu + tags: pointwise + +- func: le.Tensor(Tensor self, Tensor other) -> Tensor + structured_delegate: le.Tensor_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: le_quantized_cpu + tags: [core, pointwise] + +- func: le_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + structured_delegate: le.Scalar_out + device_check: NoCheck # TensorIterator + variants: method + +- func: le_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: le.Tensor_out + device_check: NoCheck # TensorIterator + variants: method + +# less_equal, alias for torch.le +- func: less_equal.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + +- func: less_equal.Scalar(Tensor self, Scalar other) -> Tensor + variants: method, function + +- func: less_equal.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + +- func: less_equal.Tensor(Tensor self, Tensor other) -> Tensor + variants: method, function + +- func: less_equal_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + variants: method + +- func: less_equal_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + variants: method + +- func: gt.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA,MTIA: gt_Scalar_out + MPS: gt_scalar_out_mps + QuantizedCPU: gt_out_quantized_cpu + tags: pointwise + +- func: gt.Scalar(Tensor self, Scalar other) -> Tensor + structured_delegate: gt.Scalar_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: gt_quantized_cpu + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: gt_scalar_nested + tags: [core, pointwise] + +- func: gt.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: gt_Tensor_out + MPS: gt_tensor_out_mps + QuantizedCPU: gt_out_quantized_cpu + tags: pointwise + +- func: gt.Tensor(Tensor self, Tensor other) -> Tensor + structured_delegate: gt.Tensor_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: gt_quantized_cpu + tags: [core, pointwise] + +- func: gt_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + structured_delegate: gt.Scalar_out + device_check: NoCheck # TensorIterator + variants: method + +- func: gt_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: gt.Tensor_out + device_check: NoCheck # TensorIterator + variants: method + +# greater, alias for torch.gt +- func: greater.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + +- func: greater.Scalar(Tensor self, Scalar other) -> Tensor + variants: method, function + +- func: greater.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + +- func: greater.Tensor(Tensor self, Tensor other) -> Tensor + variants: method, function + +- func: greater_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + variants: method + +- func: greater_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + variants: method + +- func: lt.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: lt_Scalar_out + MPS: lt_scalar_out_mps + QuantizedCPU: lt_out_quantized_cpu + tags: pointwise + +- func: lt.Scalar(Tensor self, Scalar other) -> Tensor + structured_delegate: lt.Scalar_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: lt_quantized_cpu + tags: [core, pointwise] + +- func: lt.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: lt_Tensor_out + MPS: lt_tensor_out_mps + QuantizedCPU: lt_out_quantized_cpu + tags: pointwise + +- func: lt.Tensor(Tensor self, Tensor other) -> Tensor + structured_delegate: lt.Tensor_out + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + QuantizedCPU: lt_quantized_cpu + tags: [core, pointwise] + +- func: lt_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + structured_delegate: lt.Scalar_out + device_check: NoCheck # TensorIterator + variants: method + +- func: lt_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: lt.Tensor_out + device_check: NoCheck # TensorIterator + variants: method + +# less, alias for torch.lt +- func: less.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + +- func: less.Scalar(Tensor self, Scalar other) -> Tensor + variants: method, function + +- func: less.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + +- func: less.Tensor(Tensor self, Tensor other) -> Tensor + variants: method, function + +- func: less_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + variants: method + +- func: less_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + variants: method + +- func: take.out(Tensor self, Tensor index, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: take_out + +- func: take(Tensor self, Tensor index) -> Tensor + variants: method, function + dispatch: + CPU, CUDA: take + +- func: take_along_dim.out(Tensor self, Tensor indices, int? dim=None, *, Tensor(a!) out) -> Tensor(a!) + +- func: take_along_dim(Tensor self, Tensor indices, int? dim=None) -> Tensor + variants: method, function + +- func: index_select.out(Tensor self, int dim, Tensor index, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, QuantizedCPU: index_select_out_cpu_ + CUDA, QuantizedCUDA: index_select_out_cuda + MPS: index_select_out_mps + +- func: index_select(Tensor self, int dim, Tensor index) -> Tensor + variants: method, function + dispatch: + CPU: index_select_cpu_ + QuantizedCPU: index_select_quantized_cpu_ + CUDA: index_select_cuda + QuantizedCUDA: index_select_quantized_cuda + SparseCPU: index_select_sparse_cpu + SparseCUDA: index_select_sparse_cuda + SparseMPS: index_select_sparse_mps + MPS: index_select_mps + tags: core + +- func: index_select.dimname_out(Tensor self, Dimname dim, Tensor index, *, Tensor(a!) out) -> Tensor(a!) + +- func: index_select.dimname(Tensor self, Dimname dim, Tensor index) -> Tensor + variants: method, function + +- func: index_select_backward(Tensor grad, SymInt[] self_sizes, int dim, Tensor index) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + dispatch: + CompositeImplicitAutograd: index_select_backward_symint + +- func: masked_select.out(Tensor self, Tensor mask, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: masked_select_out_cpu + CUDA: masked_select_out_cuda + MPS: masked_select_out_mps + tags: dynamic_output_shape + +- func: masked_select(Tensor self, Tensor mask) -> Tensor + variants: method, function + dispatch: + CPU: masked_select_cpu + CUDA: masked_select_cuda + MPS: masked_select_mps + tags: dynamic_output_shape + +- func: masked_select_backward(Tensor grad, Tensor input, Tensor mask) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + +- func: nonzero.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: nonzero_out_cpu + CUDA: nonzero_out_cuda + MPS: nonzero_out_mps + tags: dynamic_output_shape + +- func: nonzero(Tensor self) -> Tensor + variants: method, function + dispatch: + CPU: nonzero_cpu + CUDA: nonzero_cuda + MPS: nonzero_mps + tags: [dynamic_output_shape, core] + +- func: nonzero_static.out(Tensor self, *, SymInt size, int fill_value=-1, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: nonzero_static_out_cpu + CUDA: nonzero_static_out_cuda + MPS: nonzero_static_out_mps + +- func: nonzero_static(Tensor self, *, SymInt size, int fill_value=-1) -> Tensor + variants: method, function + dispatch: + CPU: nonzero_static_cpu + CUDA: nonzero_static_cuda + MPS: nonzero_static_mps + +- func: nonzero_numpy(Tensor self) -> Tensor[] + variants: method, function + +- func: argwhere(Tensor self) -> Tensor + variants: method, function + tags: dynamic_output_shape + +- func: gather.out(Tensor self, int dim, Tensor index, *, bool sparse_grad=False, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU, CUDA: gather_out + MPS: gather_out_mps + +- func: gather(Tensor self, int dim, Tensor index, *, bool sparse_grad=False) -> Tensor + variants: method, function + structured_delegate: gather.out + tags: core + +- func: gather_backward(Tensor grad, Tensor self, int dim, Tensor index, bool sparse_grad) -> Tensor + variants: function + device_check: NoCheck + device_guard: False + +- func: gather.dimname_out(Tensor self, Dimname dim, Tensor index, *, bool sparse_grad=False, Tensor(a!) out) -> Tensor(a!) + +- func: gather.dimname(Tensor self, Dimname dim, Tensor index, *, bool sparse_grad=False) -> Tensor + variants: method, function + +- func: _gather_sparse_backward(Tensor self, int dim, Tensor index, Tensor grad) -> Tensor + +- func: addcmul.out(Tensor self, Tensor tensor1, Tensor tensor2, *, Scalar value=1, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: addcmul_out + MPS: addcmul_out_mps + tags: pointwise + +- func: addcmul(Tensor self, Tensor tensor1, Tensor tensor2, *, Scalar value=1) -> Tensor + structured_delegate: addcmul.out + device_check: NoCheck # TensorIterator + variants: method, function + tags: pointwise + +- func: addcmul_(Tensor(a!) self, Tensor tensor1, Tensor tensor2, *, Scalar value=1) -> Tensor(a!) + structured_delegate: addcmul.out + device_check: NoCheck # TensorIterator + variants: method + tags: pointwise + +- func: addcdiv.out(Tensor self, Tensor tensor1, Tensor tensor2, *, Scalar value=1, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA: addcdiv_out + MPS: addcdiv_out_mps + tags: pointwise + +- func: addcdiv(Tensor self, Tensor tensor1, Tensor tensor2, *, Scalar value=1) -> Tensor + structured_delegate: addcdiv.out + device_check: NoCheck # TensorIterator + variants: method, function + tags: pointwise + +- func: addcdiv_(Tensor(a!) self, Tensor tensor1, Tensor tensor2, *, Scalar value=1) -> Tensor(a!) + structured_delegate: addcdiv.out + device_check: NoCheck # TensorIterator + variants: method + tags: pointwise + +- func: cross_entropy_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100, float label_smoothing=0.0) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: cross_entropy_loss_symint + +- func: triangular_solve.X(Tensor self, Tensor A, bool upper=True, bool transpose=False, bool unitriangular=False, *, Tensor(a!) X, Tensor(b!) M) -> (Tensor(a!) solution, Tensor(b!) cloned_coefficient) + structured: True + dispatch: + CPU, CUDA: triangular_solve_out + MPS: triangular_solve_mps_out + SparseCsrCPU: triangular_solve_out_sparse_csr_cpu + SparseCsrCUDA: triangular_solve_out_sparse_csr_cuda + +- func: triangular_solve(Tensor self, Tensor A, bool upper=True, bool transpose=False, bool unitriangular=False) -> (Tensor solution, Tensor cloned_coefficient) + structured_delegate: triangular_solve.X + variants: method, function + +- func: _linalg_check_errors(Tensor info, str api_name, *, bool is_matrix) -> () + dispatch: + CompositeExplicitAutograd: _linalg_check_errors + +- func: linalg_solve_triangular.out(Tensor self, Tensor B, *, bool upper, bool left=True, bool unitriangular=False, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + dispatch: + CPU, CUDA: linalg_solve_triangular_out + MPS: linalg_solve_triangular_mps_out + +- func: linalg_solve_triangular(Tensor self, Tensor B, *, bool upper, bool left=True, bool unitriangular=False) -> Tensor + python_module: linalg + variants: function + dispatch: + CPU, CUDA: linalg_solve_triangular + MPS: linalg_solve_triangular_mps + +- func: linalg_vander(Tensor x, *, SymInt? N=None) -> Tensor + python_module: linalg + dispatch: + CompositeImplicitAutograd: linalg_vander_symint + +- func: svd.U(Tensor self, bool some=True, bool compute_uv=True, *, Tensor(a!) U, Tensor(b!) S, Tensor(c!) V) -> (Tensor(a!) U, Tensor(b!) S, Tensor(c!) V) + +- func: svd(Tensor self, bool some=True, bool compute_uv=True) -> (Tensor U, Tensor S, Tensor V) + variants: method, function + +# swapaxes, alias for transpose +- func: swapaxes(Tensor(a) self, int axis0, int axis1) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + +- func: swapaxes_(Tensor(a!) self, int axis0, int axis1) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + tags: inplace_view + +# swapdims, alias for transpose +- func: swapdims(Tensor(a) self, int dim0, int dim1) -> Tensor(a) + variants: function, method + device_check: NoCheck + device_guard: False + +- func: swapdims_(Tensor(a!) self, int dim0, int dim1) -> Tensor(a!) + variants: method + device_check: NoCheck + device_guard: False + tags: inplace_view + +- func: cholesky.out(Tensor self, bool upper=False, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: cholesky_out + +- func: cholesky(Tensor self, bool upper=False) -> Tensor + variants: method, function + dispatch: + CPU, CUDA, MPS: cholesky + +- func: cholesky_solve.out(Tensor self, Tensor input2, bool upper=False, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: cholesky_solve_out + +- func: cholesky_solve(Tensor self, Tensor input2, bool upper=False) -> Tensor + variants: method, function + dispatch: + CompositeExplicitAutograd: cholesky_solve + +- func: _cholesky_solve_helper(Tensor self, Tensor A, bool upper) -> Tensor + variants: function + dispatch: + CPU: _cholesky_solve_helper_cpu + CUDA: _cholesky_solve_helper_cuda + MPS: _cholesky_solve_helper_mps + autogen: _cholesky_solve_helper.out + +- func: cholesky_inverse(Tensor self, bool upper=False) -> Tensor + variants: method, function + dispatch: + CPU, CUDA, MPS: cholesky_inverse + +- func: cholesky_inverse.out(Tensor self, bool upper=False, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: cholesky_inverse_out + +- func: qr.Q(Tensor self, bool some=True, *, Tensor(a!) Q, Tensor(b!) R) -> (Tensor(a!) Q, Tensor(b!) R) + +- func: qr(Tensor self, bool some=True) -> (Tensor Q, Tensor R) + variants: method, function + +- func: geqrf.a(Tensor self, *, Tensor(a!) a, Tensor(b!) tau) -> (Tensor(a!) a, Tensor(b!) tau) + dispatch: + CPU, CUDA: geqrf_out + +- func: geqrf(Tensor self) -> (Tensor a, Tensor tau) + variants: method, function + dispatch: + CPU, CUDA: geqrf + +# orgqr, alias for linalg_householder_product +- func: orgqr(Tensor self, Tensor input2) -> Tensor + variants: method, function + +- func: orgqr.out(Tensor self, Tensor input2, *, Tensor(a!) out) -> Tensor(a!) + +- func: ormqr.out(Tensor self, Tensor input2, Tensor input3, bool left=True, bool transpose=False, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: ormqr_out + +- func: ormqr(Tensor self, Tensor input2, Tensor input3, bool left=True, bool transpose=False) -> Tensor + variants: method, function + dispatch: + CPU, CUDA: ormqr + +- func: _lu_with_info(Tensor self, bool pivot=True, bool check_errors=True) -> (Tensor LU, Tensor pivots, Tensor info) + variants: function + +- func: lu_solve.out(Tensor self, Tensor LU_data, Tensor LU_pivots, *, Tensor(a!) out) -> Tensor(a!) + +- func: lu_solve(Tensor self, Tensor LU_data, Tensor LU_pivots) -> Tensor + variants: method, function + +# lu_unpack +- func: lu_unpack(Tensor LU_data, Tensor LU_pivots, bool unpack_data=True, bool unpack_pivots=True) -> (Tensor P, Tensor L, Tensor U) + structured_delegate: lu_unpack.out + variants: function + +- func: lu_unpack.out(Tensor LU_data, Tensor LU_pivots, bool unpack_data=True, bool unpack_pivots=True, *, Tensor(a!) P, Tensor(b!) L, Tensor(c!) U) -> (Tensor(a!) P, Tensor(b!) L, Tensor(c!) U) + variants: function + structured: True + dispatch: + CPU, CUDA, MPS: lu_unpack_out + +# TODO: remove dispatch section when porting TH CUDA to ATen +- func: multinomial.out(Tensor self, SymInt num_samples, bool replacement=False, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + dispatch: + CPU, CUDA: multinomial_out + MPS: multinomial_out_mps + +- func: multinomial(Tensor self, SymInt num_samples, bool replacement=False, *, Generator? generator=None) -> Tensor + variants: method, function + dispatch: + CPU, CUDA: multinomial + MPS: multinomial_mps + tags: nondeterministic_seeded + +- func: lgamma.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: lgamma_out + tags: pointwise + +- func: lgamma_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: lgamma.out + variants: method + tags: pointwise + +- func: lgamma(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: lgamma.out + variants: method, function + tags: pointwise + +- func: digamma.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: digamma_out + tags: pointwise + +- func: digamma(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: digamma.out + variants: method, function + tags: pointwise + +- func: polygamma.out(int n, Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: polygamma_out + tags: pointwise + +- func: polygamma(int n, Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: polygamma.out + variants: method, function + tags: pointwise + +- func: polygamma_(Tensor(a!) self, int n) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: polygamma_ + tags: pointwise + +- func: erfinv(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: erfinv.out + variants: method, function + dispatch: + SparseCPU, SparseCUDA, SparseMPS: erfinv_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: erfinv_sparse_csr + tags: pointwise + +- func: erfinv_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: erfinv.out + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: erfinv_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: erfinv_sparse_csr_ + tags: pointwise + +- func: erfinv.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: erfinv_out + SparseCPU, SparseCUDA, SparseMPS: erfinv_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: erfinv_sparse_csr_out + tags: pointwise + +- func: i0(Tensor self) -> Tensor + structured_delegate: i0.out + variants: function, method + tags: pointwise + +- func: i0_(Tensor(a!) self) -> Tensor(a!) + structured_delegate: i0.out + variants: function, method + tags: pointwise + +- func: i0.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: i0_out + tags: pointwise + +- func: sign(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: sign.out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sign_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sign_sparse_csr + tags: [core, pointwise] + +- func: sign_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: sign.out + variants: method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: sign_sparse_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sign_sparse_csr_ + tags: pointwise + +- func: sign.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: sign_out + MPS: sign_out_mps + SparseCPU, SparseCUDA, SparseMPS: sign_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: sign_sparse_csr_out + tags: pointwise + +- func: signbit(Tensor self) -> Tensor + variants: function, method + structured_delegate: signbit.out + dispatch: + SparseCPU, SparseCUDA, SparseMPS: signbit_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: signbit_sparse_csr + tags: pointwise + +- func: signbit.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU: signbit_out + CUDA: signbit_out + MPS: signbit_out_mps + SparseCPU, SparseCUDA, SparseMPS: signbit_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: signbit_sparse_csr_out + tags: pointwise + +- func: dist(Tensor self, Tensor other, Scalar p=2) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CompositeExplicitAutograd: dist + autogen: dist.out + +- func: atan2.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: atan2_out + tags: [core, pointwise] + +- func: atan2_(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: atan2.out + variants: method + tags: pointwise + +- func: atan2(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: atan2.out + variants: method, function + tags: [core, pointwise] +# arctan2, alias of atan2 + +- func: arctan2(Tensor self, Tensor other) -> Tensor + variants: method, function + +- func: arctan2.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + +- func: arctan2_(Tensor(a!) self, Tensor other) -> Tensor(a!) + variants: method + +- func: lerp.Scalar_out(Tensor self, Tensor end, Scalar weight, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: lerp_Scalar + tags: pointwise + +- func: lerp.Tensor_out(Tensor self, Tensor end, Tensor weight, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: lerp_Tensor + tags: pointwise + +- func: lerp.Scalar(Tensor self, Tensor end, Scalar weight) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + structured_delegate: lerp.Scalar_out + tags: pointwise + +- func: lerp.Tensor(Tensor self, Tensor end, Tensor weight) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + structured_delegate: lerp.Tensor_out + tags: pointwise + +- func: histc.out(Tensor self, int bins=100, Scalar min=0, Scalar max=0, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, MPS: histogram_histc_out + CUDA: _histc_out_cuda + +- func: histc(Tensor self, int bins=100, Scalar min=0, Scalar max=0) -> Tensor + variants: method, function + dispatch: + CPU, MPS: histogram_histc + CUDA: _histc_cuda + +- func: histogram.bins_tensor_out(Tensor self, Tensor bins, *, Tensor? weight=None, bool density=False, Tensor(a!) hist, Tensor(b!) bin_edges) -> (Tensor(a!) hist, Tensor(b!) bin_edges) + dispatch: + CPU, MPS: histogram_out + +- func: histogram.bins_tensor(Tensor self, Tensor bins, *, Tensor? weight=None, bool density=False) -> (Tensor hist, Tensor bin_edges) + variants: method, function + dispatch: + CPU, MPS: histogram + +- func: histogram.bin_ct_out(Tensor self, int bins=100, *, float[]? range=None, Tensor? weight=None, bool density=False, Tensor(a!) hist, Tensor(b!) bin_edges) -> (Tensor(a!) hist, Tensor(b!) bin_edges) + dispatch: + CPU, MPS: histogram_out + +- func: histogram.bin_ct(Tensor self, int bins=100, *, float[]? range=None, Tensor? weight=None, bool density=False) -> (Tensor hist, Tensor bin_edges) + variants: method, function + dispatch: + CPU, MPS: histogram + +- func: _histogramdd_bin_edges(Tensor self, int[] bins, *, float[]? range=None, Tensor? weight=None, bool density=False) -> Tensor[] + dispatch: + CPU, MPS: histogramdd_bin_edges + autogen: _histogramdd_bin_edges.out + +- func: _histogramdd_from_bin_cts(Tensor self, int[] bins, *, float[]? range=None, Tensor? weight=None, bool density=False) -> Tensor + dispatch: + CPU, MPS: _histogramdd + autogen: _histogramdd_from_bin_cts.out + +- func: _histogramdd_from_bin_tensors(Tensor self, Tensor[] bins, *, Tensor? weight=None, bool density=False) -> Tensor + dispatch: + CPU, MPS: _histogramdd + autogen: _histogramdd_from_bin_tensors.out + +- func: histogramdd(Tensor self, int[] bins, float[]? range=None, Tensor? weight=None, bool density=False) -> (Tensor hist, Tensor[] bin_edges) + +- func: histogramdd.int_bins(Tensor self, int bins, float[]? range=None, Tensor? weight=None, bool density=False) -> (Tensor hist, Tensor[] bin_edges) + +- func: histogramdd.TensorList_bins(Tensor self, Tensor[] bins, float[]? range=None, Tensor? weight=None, bool density=False) -> (Tensor hist, Tensor[] bin_edges) + +- func: fmod.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CompositeExplicitAutograd: fmod_out + tags: pointwise + +- func: fmod.Scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CompositeExplicitAutograd: fmod + tags: [core, pointwise] + +- func: fmod_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + dispatch: + CompositeExplicitAutograd: fmod_ + tags: pointwise + +- func: fmod.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: fmod_out + tags: pointwise + +- func: fmod.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: fmod.Tensor_out + variants: method, function + tags: [core, pointwise] + +- func: fmod_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: fmod.Tensor_out + tags: pointwise + +- func: hypot.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: hypot_out + tags: pointwise + +- func: hypot(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: hypot.out + variants: method, function + tags: pointwise + +- func: hypot_(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: hypot.out + variants: method + tags: pointwise + +- func: igamma.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: igamma_out + tags: pointwise + +- func: igamma(Tensor self, Tensor other) -> Tensor + structured_delegate: igamma.out + variants: method, function + tags: pointwise + +- func: igamma_(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: igamma.out + variants: method + tags: pointwise + +- func: igammac.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: igammac_out + tags: pointwise + +- func: igammac(Tensor self, Tensor other) -> Tensor + structured_delegate: igammac.out + variants: method, function + tags: pointwise + +- func: igammac_(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: igammac.out + variants: method + tags: pointwise + +- func: nextafter.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: nextafter_out + tags: pointwise + +- func: nextafter(Tensor self, Tensor other) -> Tensor + structured_delegate: nextafter.out + variants: method, function + tags: pointwise + +- func: nextafter_(Tensor(a!) self, Tensor other) -> Tensor(a!) + structured_delegate: nextafter.out + variants: method + tags: pointwise + +- func: remainder.Scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: remainder_out + tags: pointwise + +- func: remainder.Scalar(Tensor self, Scalar other) -> Tensor + variants: method, function + dispatch: + CompositeExplicitAutograd: remainder + tags: [core, pointwise] + +- func: remainder_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + variants: method + dispatch: + CompositeExplicitAutograd: remainder_ + tags: pointwise + +- func: remainder.Tensor_out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS, MTIA: remainder_out + tags: pointwise + +- func: remainder.Tensor(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: remainder.Tensor_out + variants: method, function + tags: [core, pointwise] + +- func: remainder_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: remainder.Tensor_out + variants: method + tags: pointwise + +- func: remainder.Scalar_Tensor(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: function + dispatch: + CPU, CUDA, MPS: remainder + autogen: remainder.Scalar_Tensor_out + tags: pointwise + +- func: min(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CPU, CUDA: min + MPS: min_mps + QuantizedCPU: min_quantized_cpu + tags: [reduction] + +- func: min.unary_out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: min_unary_out + QuantizedCPU: min_quantized_unary_out + tags: [reduction] + +- func: fmin(Tensor self, Tensor other) -> Tensor + structured_delegate: fmin.out + device_check: NoCheck # TensorIterator + variants: method, function + tags: pointwise + +- func: fmin.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS: fmin_out + tags: pointwise + +- func: max(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CPU, CUDA: max + MPS: max_mps + QuantizedCPU: max_quantized_cpu + tags: [reduction] + +- func: fmax(Tensor self, Tensor other) -> Tensor + structured_delegate: fmax.out + device_check: NoCheck # TensorIterator + variants: method, function + tags: pointwise + +- func: fmax.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MPS: fmax_out + tags: pointwise + +- func: maximum(Tensor self, Tensor other) -> Tensor + structured_delegate: maximum.out + device_check: NoCheck # TensorIterator + variants: method, function + tags: [core, pointwise] + +- func: maximum.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA, MPS: maximum_out + tags: pointwise + +# binary max, alias of maximum +# NOTE: max is not an alias for maximum, since there is also unary max +- func: max.other(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + tags: pointwise + +- func: max.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: pointwise + +- func: max.unary_out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA: max_unary_out + QuantizedCPU: max_quantized_unary_out + tags: [reduction] + +- func: minimum(Tensor self, Tensor other) -> Tensor + structured_delegate: minimum.out + device_check: NoCheck # TensorIterator + variants: method, function + tags: [core, pointwise] + +- func: minimum.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + dispatch: + CPU, CUDA, MTIA, MPS: minimum_out + tags: pointwise + +# binary min, alias for minimum +# NOTE: min is not an alias for minimum, since there is also unary min +- func: min.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: pointwise + +- func: min.other(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + tags: pointwise + +- func: quantile(Tensor self, Tensor q, int? dim=None, bool keepdim=False, *, str interpolation='linear') -> Tensor + variants: method, function + +- func: quantile.out(Tensor self, Tensor q, int? dim=None, bool keepdim=False, *, str interpolation='linear', Tensor(a!) out) -> Tensor(a!) + +- func: quantile.scalar(Tensor self, float q, int? dim=None, bool keepdim=False, *, str interpolation='linear') -> Tensor + variants: method, function + +- func: quantile.scalar_out(Tensor self, float q, int? dim=None, bool keepdim=False, *, str interpolation='linear', Tensor(a!) out) -> Tensor(a!) + +- func: nanquantile(Tensor self, Tensor q, int? dim=None, bool keepdim=False, *, str interpolation='linear') -> Tensor + variants: method, function + +- func: nanquantile.out(Tensor self, Tensor q, int? dim=None, bool keepdim=False, *, str interpolation='linear', Tensor(a!) out) -> Tensor(a!) + +- func: nanquantile.scalar(Tensor self, float q, int? dim=None, bool keepdim=False, *, str interpolation='linear') -> Tensor + variants: method, function + +- func: nanquantile.scalar_out(Tensor self, float q, int? dim=None, bool keepdim=False, *, str interpolation='linear', Tensor(a!) out) -> Tensor(a!) + +- func: sort.values(Tensor self, int dim=-1, bool descending=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + device_check: NoCheck # TensorIterator + dispatch: + CompositeExplicitAutograd: sort_out + +- func: sort.values_stable(Tensor self, *, bool? stable, int dim=-1, bool descending=False, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + structured: True + dispatch: + CPU, CUDA: sort_stable_out + MPS: sort_stable_out_mps + +- func: sort(Tensor self, int dim=-1, bool descending=False) -> (Tensor values, Tensor indices) + device_check: NoCheck # TensorIterator + variants: method, function + dispatch: + CompositeExplicitAutograd: sort + tags: core + +- func: sort.stable(Tensor self, *, bool? stable, int dim=-1, bool descending=False) -> (Tensor values, Tensor indices) + structured_delegate: sort.values_stable + variants: method, function + dispatch: + QuantizedCPU: sort_quantized_cpu_stable + +- func: sort.dimname_values(Tensor self, Dimname dim, bool descending=False, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + +- func: sort.dimname_values_stable(Tensor self, *, bool? stable, Dimname dim, bool descending=False, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + +- func: sort.dimname(Tensor self, Dimname dim, bool descending=False) -> (Tensor values, Tensor indices) + variants: method, function + +- func: sort.dimname_stable(Tensor self, *, bool? stable, Dimname dim, bool descending=False) -> (Tensor values, Tensor indices) + variants: method, function + +- func: msort.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: msort(Tensor self) -> Tensor + variants: method, function + +- func: argsort(Tensor self, int dim=-1, bool descending=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + +- func: argsort.stable(Tensor self, *, bool stable, int dim=-1, bool descending=False) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + +- func: argsort.stable_out(Tensor self, *, bool stable, int dim=-1, bool descending=False, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: function + +- func: argsort.dimname(Tensor self, Dimname dim, bool descending=False) -> Tensor + variants: method, function + +- func: topk.values(Tensor self, SymInt k, int dim=-1, bool largest=True, bool sorted=True, *, Tensor(a!) values, Tensor(b!) indices) -> (Tensor(a!) values, Tensor(b!) indices) + structured: True + dispatch: + CPU: topk_out_cpu + CUDA: topk_out_cuda + MPS: topk_out_mps + +- func: topk(Tensor self, SymInt k, int dim=-1, bool largest=True, bool sorted=True) -> (Tensor values, Tensor indices) + variants: method, function + structured_delegate: topk.values + dispatch: + QuantizedCPU: topk_quantized_cpu + tags: core + +- func: all(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: all.all_out + variants: method, function + tags: reduction + +- func: all.all_out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + structured: True + dispatch: + CPU, CUDA: all_all_out + MTIA: all_all_out_mtia + MPS: all_all_out_mps + tags: reduction + +- func: any(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: any.all_out + variants: method, function + dispatch: + SparseCPU, SparseCUDA, SparseMPS: any_sparse + tags: [core, reduction] + +- func: any.all_out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + structured: True + dispatch: + CPU, CUDA: any_all_out + MPS: any_all_out_mps + tags: reduction + +- func: renorm.out(Tensor self, Scalar p, int dim, Scalar maxnorm, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + dispatch: + CPU, CUDA: renorm_out + MPS: renorm_out_mps + +- func: renorm(Tensor self, Scalar p, int dim, Scalar maxnorm) -> Tensor + device_check: NoCheck # TensorIterator + variants: method, function + structured_delegate: renorm.out + +- func: renorm_(Tensor(a!) self, Scalar p, int dim, Scalar maxnorm) -> Tensor(a!) + device_check: NoCheck # TensorIterator + variants: method + structured_delegate: renorm.out + +- func: unfold(Tensor(a) self, int dimension, int size, int step) -> Tensor(a) + variants: method + device_check: NoCheck + device_guard: False + dispatch: + CPU, CUDA, Meta, MPS, MTIA: unfold + QuantizedCPU, QuantizedCUDA: unfold + +- func: unfold_backward(Tensor grad_in, SymInt[] input_sizes, int dim, int size, int step) -> Tensor + variants: function + dispatch: + CPU, CUDA, MPS: unfold_backward + autogen: unfold_backward.out + +- func: equal(Tensor self, Tensor other) -> bool + tags: [data_dependent_output, pointwise] + variants: method, function + dispatch: + CPU: cpu_equal + CUDA: cuda_equal + MPS: mps_equal + QuantizedCPU: equal_quantized_cpu + +- func: pow.Tensor_Tensor_out(Tensor self, Tensor exponent, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: pow_Tensor_Tensor_out + MPS: pow_tensor_tensor_out_mps + tags: pointwise + +- func: pow.Tensor_Tensor(Tensor self, Tensor exponent) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: pow.Tensor_Tensor_out + variants: method, function + tags: [core, pointwise] + +- func: pow.Scalar_out(Scalar self, Tensor exponent, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + dispatch: + CPU, CUDA: pow_Scalar_out + MPS: pow_Scalar_out_mps + tags: pointwise + +- func: pow.Scalar(Scalar self, Tensor exponent) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: pow.Scalar_out + tags: [core, pointwise] + +- func: pow.Tensor_Scalar_out(Tensor self, Scalar exponent, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: pow_Tensor_Scalar_out + SparseCPU, SparseCUDA, SparseMPS: pow_out_sparse_scalar + tags: pointwise + +- func: pow.Tensor_Scalar(Tensor self, Scalar exponent) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: pow.Tensor_Scalar_out + variants: function, method + dispatch: + SparseCPU, SparseCUDA, SparseMPS: pow_sparse_scalar + tags: [core, pointwise] + +- func: pow_.Scalar(Tensor(a!) self, Scalar exponent) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: pow.Tensor_Scalar_out + variants: method + tags: pointwise + +- func: pow_.Tensor(Tensor(a!) self, Tensor exponent) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured_delegate: pow.Tensor_Tensor_out + variants: method + tags: pointwise + +- func: float_power.Tensor_Tensor_out(Tensor self, Tensor exponent, *, Tensor(a!) out) -> Tensor(a!) + tags: pointwise + +- func: float_power.Tensor_Tensor(Tensor self, Tensor exponent) -> Tensor + variants: function, method + tags: pointwise + +- func: float_power.Scalar_out(Scalar self, Tensor exponent, *, Tensor(a!) out) -> Tensor(a!) + tags: pointwise + +- func: float_power.Scalar(Scalar self, Tensor exponent) -> Tensor + tags: pointwise + +- func: float_power.Tensor_Scalar_out(Tensor self, Scalar exponent, *, Tensor(a!) out) -> Tensor(a!) + tags: pointwise + +- func: float_power.Tensor_Scalar(Tensor self, Scalar exponent) -> Tensor + variants: function, method + tags: pointwise + +- func: float_power_.Scalar(Tensor(a!) self, Scalar exponent) -> Tensor(a!) + variants: method + tags: pointwise + +- func: float_power_.Tensor(Tensor(a!) self, Tensor exponent) -> Tensor(a!) + variants: method + tags: pointwise + +- func: normal_(Tensor(a!) self, float mean=0, float std=1, *, Generator? generator=None) -> Tensor(a!) + device_check: NoCheck # TensorIterator + tags: nondeterministic_seeded + variants: method + dispatch: + CPU, CUDA: normal_ + MPS: normal_mps_ + Meta: normal_meta_ + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: normal_sparse_csr_ + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: normal_nested_ + autogen: normal.out + +# Only used by the functionalization pass. +# Normally, the codegen would be able to generate a normal() NativeFunction, +# but we can't due to overload ambiguity with normal.Tensor_float. +- func: normal_functional(Tensor self, float mean=0, float std=1, *, Generator? generator=None) -> Tensor + device_check: NoCheck # TensorIterator + tags: nondeterministic_seeded + dispatch: + CompositeExplicitAutograd: normal_functional + +- func: normal.Tensor_float_out(Tensor mean, float std=1, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!) + tags: nondeterministic_seeded + dispatch: + CPU, CUDA: normal_out + MPS: normal_mps_out + Meta: normal_out_meta + +- func: normal.Tensor_float(Tensor mean, float std=1, *, Generator? generator=None) -> Tensor + dispatch: + CPU, CUDA: normal + MPS: normal_mps + Meta: normal_meta + tags: nondeterministic_seeded + +- func: normal.float_Tensor_out(float mean, Tensor std, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: normal_out + Meta: normal_out_meta + MPS: normal_mps_out + tags: nondeterministic_seeded + +- func: normal.float_Tensor(float mean, Tensor std, *, Generator? generator=None) -> Tensor + dispatch: + CPU, CUDA: normal + MPS: normal_mps + Meta: normal_meta + tags: nondeterministic_seeded + +- func: normal.Tensor_Tensor_out(Tensor mean, Tensor std, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: normal_out + Meta: normal_out_meta + MPS: normal_mps_out + tags: nondeterministic_seeded + +- func: normal.Tensor_Tensor(Tensor mean, Tensor std, *, Generator? generator=None) -> Tensor + dispatch: + CPU, CUDA: normal + MPS: normal_mps + Meta: normal_meta + tags: nondeterministic_seeded + +- func: normal.float_float(float mean, float std, SymInt[] size, *, Generator? generator=None, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + dispatch: + CompositeExplicitAutograd: normal + tags: nondeterministic_seeded + +- func: normal.float_float_out(float mean, float std, SymInt[] size, *, Generator? generator=None, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: normal_out + tags: nondeterministic_seeded + +- func: alias(Tensor(a) self) -> Tensor(a) + variants: method, function + dispatch: + CompositeExplicitAutograd: alias + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: alias_nested + tags: core + +- func: _amp_foreach_non_finite_check_and_unscale_(Tensor(a!)[] self, Tensor(b!) found_inf, Tensor inv_scale) -> () + variants: function + dispatch: + CUDA: _amp_foreach_non_finite_check_and_unscale_cuda_ + CPU: _amp_foreach_non_finite_check_and_unscale_cpu_ + MPS: _amp_foreach_non_finite_check_and_unscale_mps_ + autogen: _amp_foreach_non_finite_check_and_unscale, _amp_foreach_non_finite_check_and_unscale.out + +- func: _amp_update_scale_(Tensor(a!) self, Tensor(b!) growth_tracker, Tensor found_inf, float scale_growth_factor, float scale_backoff_factor, int growth_interval) -> Tensor(a!) + variants: function + dispatch: + CUDA: _amp_update_scale_cuda_ + CPU: _amp_update_scale_cpu_ + MPS: _amp_update_scale_mps_ + autogen: _amp_update_scale, _amp_update_scale.out + + #- func: _cat(Tensor[] tensors, int dim=0) -> Tensor + #dispatch: + #CPU: _cat_cpu + #CUDA: cat_cuda + #MPS: cat_mps + #QuantizedCPU: cat_quantized_cpu + + #- func: _cat.out(Tensor[] tensors, int dim=0, *, Tensor(a!) out) -> Tensor(a!) + #dispatch: + #CPU: _cat_out_cpu + #CUDA: cat_out_cuda + #QuantizedCPU: cat_out_quantized_cpu + +- func: _foreach_add.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_add_scalar_kernel_slow + CUDA: foreach_tensor_add_scalar_kernel_cuda + +- func: _foreach_add_.Scalar(Tensor(a!)[] self, Scalar scalar) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_add_scalar_kernel_slow_ + CUDA: foreach_tensor_add_scalar_kernel_cuda_ + MTIA: foreach_tensor_add_scalar_kernel_mtia_ + autogen: _foreach_add.Scalar_out + +- func: _foreach_add.List(Tensor[] self, Tensor[] other, *, Scalar alpha=1) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_add_list_kernel_slow + CUDA: foreach_tensor_add_list_kernel_cuda + MTIA: foreach_tensor_add_list_kernel_mtia + +- func: _foreach_add_.List(Tensor(a!)[] self, Tensor[] other, *, Scalar alpha=1) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_add_list_kernel_slow_ + CUDA: foreach_tensor_add_list_kernel_cuda_ + MTIA: foreach_tensor_add_list_kernel_mtia_ + autogen: _foreach_add.List_out + +- func: _foreach_add.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_add_scalarlist_kernel_slow + CUDA: foreach_tensor_add_scalarlist_kernel_cuda + +- func: _foreach_add_.ScalarList(Tensor(a!)[] self, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_add_scalarlist_kernel_slow_ + CUDA: foreach_tensor_add_scalarlist_kernel_cuda_ + autogen: _foreach_add.ScalarList_out + +- func: _foreach_add.Tensor(Tensor[] self, Tensor other, *, Scalar alpha=1) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_add_tensor_kernel_slow + CUDA: foreach_tensor_add_tensor_kernel_cuda + +- func: _foreach_add_.Tensor(Tensor(a!)[] self, Tensor other, *, Scalar alpha=1) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_add_tensor_kernel_slow_ + CUDA: foreach_tensor_add_tensor_kernel_cuda_ + MTIA: foreach_tensor_add_tensor_kernel_mtia_ + autogen: _foreach_add.Tensor_out + +- func: _foreach_sub.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sub_scalar_kernel_slow + CUDA: foreach_tensor_sub_scalar_kernel_cuda + +- func: _foreach_sub_.Scalar(Tensor(a!)[] self, Scalar scalar) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sub_scalar_kernel_slow_ + CUDA: foreach_tensor_sub_scalar_kernel_cuda_ + autogen: _foreach_sub.Scalar_out + +- func: _foreach_sub.List(Tensor[] self, Tensor[] other, *, Scalar alpha=1) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sub_list_kernel_slow + CUDA: foreach_tensor_sub_list_kernel_cuda + +- func: _foreach_sub_.List(Tensor(a!)[] self, Tensor[] other, *, Scalar alpha=1) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sub_list_kernel_slow_ + CUDA: foreach_tensor_sub_list_kernel_cuda_ + autogen: _foreach_sub.List_out + +- func: _foreach_sub.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sub_scalarlist_kernel_slow + CUDA: foreach_tensor_sub_scalarlist_kernel_cuda + +- func: _foreach_sub_.ScalarList(Tensor(a!)[] self, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sub_scalarlist_kernel_slow_ + CUDA: foreach_tensor_sub_scalarlist_kernel_cuda_ + autogen: _foreach_sub.ScalarList_out + +- func: _foreach_mul.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_mul_scalar_kernel_slow + CUDA: foreach_tensor_mul_scalar_kernel_cuda + +- func: _foreach_mul_.Scalar(Tensor(a!)[] self, Scalar scalar) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_mul_scalar_kernel_slow_ + CUDA: foreach_tensor_mul_scalar_kernel_cuda_ + MTIA: foreach_tensor_mul_scalar_kernel_mtia_ + autogen: _foreach_mul.Scalar_out + +- func: _foreach_mul.List(Tensor[] self, Tensor[] other) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_mul_list_kernel_slow + CUDA: foreach_tensor_mul_list_kernel_cuda + MTIA: foreach_tensor_mul_list_kernel_mtia + +- func: _foreach_mul_.List(Tensor(a!)[] self, Tensor[] other) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_mul_list_kernel_slow_ + CUDA: foreach_tensor_mul_list_kernel_cuda_ + MTIA: foreach_tensor_mul_list_kernel_mtia_ + autogen: _foreach_mul.List_out + +- func: _foreach_mul.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_mul_scalarlist_kernel_slow + CUDA: foreach_tensor_mul_scalarlist_kernel_cuda + +- func: _foreach_mul_.ScalarList(Tensor(a!)[] self, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_mul_scalarlist_kernel_slow_ + CUDA: foreach_tensor_mul_scalarlist_kernel_cuda_ + autogen: _foreach_mul.ScalarList_out + +- func: _foreach_mul.Tensor(Tensor[] self, Tensor other) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_mul_tensor_kernel_slow + CUDA: foreach_tensor_mul_tensor_kernel_cuda + MTIA: foreach_tensor_mul_tensor_kernel_mtia + +- func: _foreach_mul_.Tensor(Tensor(a!)[] self, Tensor other) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_mul_tensor_kernel_slow_ + CUDA: foreach_tensor_mul_tensor_kernel_cuda_ + MTIA: foreach_tensor_mul_tensor_kernel_mtia_ + autogen: _foreach_mul.Tensor_out + +- func: _foreach_div.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_div_scalar_kernel_slow + CUDA: foreach_tensor_div_scalar_kernel_cuda + +- func: _foreach_div_.Scalar(Tensor(a!)[] self, Scalar scalar) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_div_scalar_kernel_slow_ + CUDA: foreach_tensor_div_scalar_kernel_cuda_ + autogen: _foreach_div.Scalar_out + +- func: _foreach_div.List(Tensor[] self, Tensor[] other) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_div_list_kernel_slow + CUDA: foreach_tensor_div_list_kernel_cuda + MTIA: foreach_tensor_div_list_kernel_mtia + +- func: _foreach_div_.List(Tensor(a!)[] self, Tensor[] other) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_div_list_kernel_slow_ + CUDA: foreach_tensor_div_list_kernel_cuda_ + MTIA: foreach_tensor_div_list_kernel_mtia_ + autogen: _foreach_div.List_out + +- func: _foreach_div.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_div_scalarlist_kernel_slow + CUDA: foreach_tensor_div_scalarlist_kernel_cuda + +- func: _foreach_div_.ScalarList(Tensor(a!)[] self, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_div_scalarlist_kernel_slow_ + CUDA: foreach_tensor_div_scalarlist_kernel_cuda_ + autogen: _foreach_div.ScalarList_out + +- func: _foreach_div.Tensor(Tensor[] self, Tensor other) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_div_tensor_kernel_slow + CUDA: foreach_tensor_div_tensor_kernel_cuda + MTIA: foreach_tensor_div_tensor_kernel_mtia + +- func: _foreach_div_.Tensor(Tensor(a!)[] self, Tensor other) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_div_tensor_kernel_slow_ + CUDA: foreach_tensor_div_tensor_kernel_cuda_ + MTIA: foreach_tensor_div_tensor_kernel_mtia_ + autogen: _foreach_div.Tensor_out + +- func: _foreach_clamp_max.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_scalar_kernel_slow + CUDA: foreach_tensor_clamp_max_scalar_kernel_cuda + +- func: _foreach_clamp_max_.Scalar(Tensor(a!)[] self, Scalar scalar) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_scalar_kernel_slow_ + CUDA: foreach_tensor_clamp_max_scalar_kernel_cuda_ + autogen: _foreach_clamp_max.Scalar_out + +- func: _foreach_clamp_max.List(Tensor[] self, Tensor[] other) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_list_kernel_slow + CUDA: foreach_tensor_clamp_max_list_kernel_cuda + +- func: _foreach_clamp_max_.List(Tensor(a!)[] self, Tensor[] other) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_list_kernel_slow_ + CUDA: foreach_tensor_clamp_max_list_kernel_cuda_ + autogen: _foreach_clamp_max.List_out + +- func: _foreach_clamp_max.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_scalarlist_kernel_slow + CUDA: foreach_tensor_clamp_max_scalarlist_kernel_cuda + +- func: _foreach_clamp_max_.ScalarList(Tensor(a!)[] self, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_scalarlist_kernel_slow_ + CUDA: foreach_tensor_clamp_max_scalarlist_kernel_cuda_ + autogen: _foreach_clamp_max.ScalarList_out + +- func: _foreach_clamp_min.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_scalar_kernel_slow + CUDA: foreach_tensor_clamp_min_scalar_kernel_cuda + +- func: _foreach_clamp_min_.Scalar(Tensor(a!)[] self, Scalar scalar) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_scalar_kernel_slow_ + CUDA: foreach_tensor_clamp_min_scalar_kernel_cuda_ + autogen: _foreach_clamp_min.Scalar_out + +- func: _foreach_clamp_min.List(Tensor[] self, Tensor[] other) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_list_kernel_slow + CUDA: foreach_tensor_clamp_min_list_kernel_cuda + +- func: _foreach_clamp_min_.List(Tensor(a!)[] self, Tensor[] other) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_list_kernel_slow_ + CUDA: foreach_tensor_clamp_min_list_kernel_cuda_ + autogen: _foreach_clamp_min.List_out + +- func: _foreach_clamp_min.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_scalarlist_kernel_slow + CUDA: foreach_tensor_clamp_min_scalarlist_kernel_cuda + +- func: _foreach_clamp_min_.ScalarList(Tensor(a!)[] self, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_scalarlist_kernel_slow_ + CUDA: foreach_tensor_clamp_min_scalarlist_kernel_cuda_ + autogen: _foreach_clamp_min.ScalarList_out + +# foreach_minimum/maximum dispatches to clamp_max/min +- func: _foreach_maximum.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_scalar_kernel_slow + CUDA: foreach_tensor_clamp_min_scalar_kernel_cuda + +- func: _foreach_maximum_.Scalar(Tensor(a!)[] self, Scalar scalar) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_scalar_kernel_slow_ + CUDA: foreach_tensor_clamp_min_scalar_kernel_cuda_ + MTIA: foreach_tensor_maximum_scalar_kernel_mtia_ + autogen: _foreach_maximum.Scalar_out + +# foreach_minimum/maximum dispatches to clamp_max/min +- func: _foreach_maximum.List(Tensor[] self, Tensor[] other) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_list_kernel_slow + CUDA: foreach_tensor_clamp_min_list_kernel_cuda + +- func: _foreach_maximum_.List(Tensor(a!)[] self, Tensor[] other) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_list_kernel_slow_ + CUDA: foreach_tensor_clamp_min_list_kernel_cuda_ + autogen: _foreach_maximum.List_out + +# foreach_minimum/maximum dispatches to clamp_max/min +- func: _foreach_maximum.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_scalarlist_kernel_slow + CUDA: foreach_tensor_clamp_min_scalarlist_kernel_cuda + +- func: _foreach_maximum_.ScalarList(Tensor(a!)[] self, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_min_scalarlist_kernel_slow_ + CUDA: foreach_tensor_clamp_min_scalarlist_kernel_cuda_ + autogen: _foreach_maximum.ScalarList_out + +- func: _foreach_minimum.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_scalar_kernel_slow + CUDA: foreach_tensor_clamp_max_scalar_kernel_cuda + +- func: _foreach_minimum_.Scalar(Tensor(a!)[] self, Scalar scalar) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_scalar_kernel_slow_ + CUDA: foreach_tensor_clamp_max_scalar_kernel_cuda_ + autogen: _foreach_minimum.Scalar_out + +- func: _foreach_minimum.List(Tensor[] self, Tensor[] other) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_list_kernel_slow + CUDA: foreach_tensor_clamp_max_list_kernel_cuda + +- func: _foreach_minimum_.List(Tensor(a!)[] self, Tensor[] other) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_list_kernel_slow_ + CUDA: foreach_tensor_clamp_max_list_kernel_cuda_ + autogen: _foreach_minimum.List_out + +- func: _foreach_minimum.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_scalarlist_kernel_slow + CUDA: foreach_tensor_clamp_max_scalarlist_kernel_cuda + +- func: _foreach_minimum_.ScalarList(Tensor(a!)[] self, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clamp_max_scalarlist_kernel_slow_ + CUDA: foreach_tensor_clamp_max_scalarlist_kernel_cuda_ + autogen: _foreach_minimum.ScalarList_out + +- func: _foreach_addcdiv.Scalar(Tensor[] self, Tensor[] tensor1, Tensor[] tensor2, Scalar value=1) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcdiv_scalar_slow + CUDA: foreach_tensor_addcdiv_scalar_cuda + +- func: _foreach_addcdiv.ScalarList(Tensor[] self, Tensor[] tensor1, Tensor[] tensor2, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcdiv_scalarlist_slow + CUDA: foreach_tensor_addcdiv_scalarlist_cuda + +- func: _foreach_addcdiv.Tensor(Tensor[] self, Tensor[] tensor1, Tensor[] tensor2, Tensor scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcdiv_tensor_slow + CUDA: foreach_tensor_addcdiv_tensor_cuda + +- func: _foreach_addcdiv_.Scalar(Tensor(a!)[] self, Tensor[] tensor1, Tensor[] tensor2, Scalar value=1) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcdiv_scalar_slow_ + CUDA: foreach_tensor_addcdiv_scalar_cuda_ + autogen: _foreach_addcdiv.Scalar_out + +- func: _foreach_addcdiv_.ScalarList(Tensor(a!)[] self, Tensor[] tensor1, Tensor[] tensor2, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcdiv_scalarlist_slow_ + CUDA: foreach_tensor_addcdiv_scalarlist_cuda_ + autogen: _foreach_addcdiv.ScalarList_out + +- func: _foreach_addcdiv_.Tensor(Tensor(a!)[] self, Tensor[] tensor1, Tensor[] tensor2, Tensor scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcdiv_tensor_slow_ + CUDA: foreach_tensor_addcdiv_tensor_cuda_ + autogen: _foreach_addcdiv.Tensor_out + +- func: _foreach_addcmul.Scalar(Tensor[] self, Tensor[] tensor1, Tensor[] tensor2, Scalar value=1) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcmul_scalar_slow + CUDA: foreach_tensor_addcmul_scalar_cuda + MTIA: foreach_tensor_addcmul_scalar_mtia + +- func: _foreach_addcmul.ScalarList(Tensor[] self, Tensor[] tensor1, Tensor[] tensor2, Scalar[] scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcmul_scalarlist_slow + CUDA: foreach_tensor_addcmul_scalarlist_cuda + +- func: _foreach_addcmul.Tensor(Tensor[] self, Tensor[] tensor1, Tensor[] tensor2, Tensor scalars) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcmul_tensor_slow + CUDA: foreach_tensor_addcmul_tensor_cuda + +- func: _foreach_addcmul_.Scalar(Tensor(a!)[] self, Tensor[] tensor1, Tensor[] tensor2, Scalar value=1) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcmul_scalar_slow_ + CUDA: foreach_tensor_addcmul_scalar_cuda_ + MTIA: foreach_tensor_addcmul_scalar_mtia_ + autogen: _foreach_addcmul.Scalar_out + +- func: _foreach_addcmul_.ScalarList(Tensor(a!)[] self, Tensor[] tensor1, Tensor[] tensor2, Scalar[] scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcmul_scalarlist_slow_ + CUDA: foreach_tensor_addcmul_scalarlist_cuda_ + autogen: _foreach_addcmul.ScalarList_out + +- func: _foreach_addcmul_.Tensor(Tensor(a!)[] self, Tensor[] tensor1, Tensor[] tensor2, Tensor scalars) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_addcmul_tensor_slow_ + CUDA: foreach_tensor_addcmul_tensor_cuda_ + autogen: _foreach_addcmul.Tensor_out + +- func: _foreach_abs(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_abs_slow + CUDA: foreach_tensor_abs_cuda + MTIA: foreach_tensor_abs_mtia + +- func: _foreach_abs_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_abs_slow_ + CUDA: foreach_tensor_abs_cuda_ + MTIA: foreach_tensor_abs_mtia_ + autogen: _foreach_abs.out + +- func: _foreach_acos(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_acos_slow + CUDA: foreach_tensor_acos_cuda + +- func: _foreach_acos_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_acos_slow_ + CUDA: foreach_tensor_acos_cuda_ + autogen: _foreach_acos.out + +- func: _foreach_asin(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_asin_slow + CUDA: foreach_tensor_asin_cuda + +- func: _foreach_asin_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_asin_slow_ + CUDA: foreach_tensor_asin_cuda_ + autogen: _foreach_asin.out + +- func: _foreach_atan(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_atan_slow + CUDA: foreach_tensor_atan_cuda + +- func: _foreach_atan_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_atan_slow_ + CUDA: foreach_tensor_atan_cuda_ + autogen: _foreach_atan.out + +- func: _foreach_ceil(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_ceil_slow + CUDA: foreach_tensor_ceil_cuda + +- func: _foreach_ceil_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_ceil_slow_ + CUDA: foreach_tensor_ceil_cuda_ + autogen: _foreach_ceil.out + +- func: _foreach_cos(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_cos_slow + CUDA: foreach_tensor_cos_cuda + +- func: _foreach_cos_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_cos_slow_ + CUDA: foreach_tensor_cos_cuda_ + autogen: _foreach_cos.out + +- func: _foreach_cosh(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_cosh_slow + CUDA: foreach_tensor_cosh_cuda + +- func: _foreach_cosh_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_cosh_slow_ + CUDA: foreach_tensor_cosh_cuda_ + autogen: _foreach_cosh.out + +- func: _foreach_erf(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_erf_slow + CUDA: foreach_tensor_erf_cuda + +- func: _foreach_erf_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_erf_slow_ + CUDA: foreach_tensor_erf_cuda_ + autogen: _foreach_erf.out + +- func: _foreach_erfc(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_erfc_slow + CUDA: foreach_tensor_erfc_cuda + +- func: _foreach_erfc_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_erfc_slow_ + CUDA: foreach_tensor_erfc_cuda_ + autogen: _foreach_erfc.out + +- func: _foreach_exp(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_exp_slow + CUDA: foreach_tensor_exp_cuda + +- func: _foreach_exp_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_exp_slow_ + CUDA: foreach_tensor_exp_cuda_ + autogen: _foreach_exp.out + +- func: _foreach_expm1(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_expm1_slow + CUDA: foreach_tensor_expm1_cuda + +- func: _foreach_expm1_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_expm1_slow_ + CUDA: foreach_tensor_expm1_cuda_ + autogen: _foreach_expm1.out + +- func: _foreach_floor(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_floor_slow + CUDA: foreach_tensor_floor_cuda + +- func: _foreach_floor_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_floor_slow_ + CUDA: foreach_tensor_floor_cuda_ + autogen: _foreach_floor.out + +- func: _foreach_frac(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_frac_slow + CUDA: foreach_tensor_frac_cuda + +- func: _foreach_frac_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_frac_slow_ + CUDA: foreach_tensor_frac_cuda_ + autogen: _foreach_frac.out + +- func: _foreach_lerp.List(Tensor[] self, Tensor[] tensors1, Tensor[] weights) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensors are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_ternary_lerp_slow + CUDA: foreach_tensor_lerp_ternary_cuda + autogen: _foreach_lerp.List_out + +- func: _foreach_lerp_.List(Tensor(a!)[] self, Tensor[] tensors1, Tensor[] weights) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensors are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_ternary_lerp_slow_ + CUDA: foreach_tensor_lerp_ternary_cuda_ + autogen: _foreach_lerp.List_out + +- func: _foreach_lerp.Scalar(Tensor[] self, Tensor[] tensors1, Scalar weight) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensors are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_lerp_list_kernel_slow + CUDA: foreach_tensor_lerp_list_cuda + autogen: _foreach_lerp.Scalar_out + +- func: _foreach_lerp_.Scalar(Tensor(a!)[] self, Tensor[] tensors1, Scalar weight) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensors are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_lerp_list_kernel_slow_ + CUDA: foreach_tensor_lerp_list_cuda_ + autogen: _foreach_lerp.Scalar_out + +- func: _foreach_lerp.ScalarList(Tensor[] self, Tensor[] tensors1, Scalar[] weight) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensors are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_lerp_scalarlist_kernel_slow + CUDA: foreach_tensor_lerp_scalarlist_cuda + autogen: _foreach_lerp.ScalarList_out + +- func: _foreach_lerp_.ScalarList(Tensor(a!)[] self, Tensor[] tensors1, Scalar[] weight) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensors are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_lerp_scalarlist_kernel_slow_ + CUDA: foreach_tensor_lerp_scalarlist_cuda_ + autogen: _foreach_lerp.ScalarList_out + +- func: _foreach_lgamma(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_lgamma_slow + CUDA: foreach_tensor_lgamma_cuda + +- func: _foreach_lgamma_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_lgamma_slow_ + CUDA: foreach_tensor_lgamma_cuda_ + autogen: _foreach_lgamma.out + +- func: _foreach_log(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_log_slow + CUDA: foreach_tensor_log_cuda + +- func: _foreach_log_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_log_slow_ + CUDA: foreach_tensor_log_cuda_ + autogen: _foreach_log.out + +- func: _foreach_log10(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_log10_slow + CUDA: foreach_tensor_log10_cuda + +- func: _foreach_log10_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_log10_slow_ + CUDA: foreach_tensor_log10_cuda_ + autogen: _foreach_log10.out + +- func: _foreach_log1p(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_log1p_slow + CUDA: foreach_tensor_log1p_cuda + +- func: _foreach_log1p_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_log1p_slow_ + CUDA: foreach_tensor_log1p_cuda_ + autogen: _foreach_log1p.out + +- func: _foreach_log2(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_log2_slow + CUDA: foreach_tensor_log2_cuda + +- func: _foreach_log2_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_log2_slow_ + CUDA: foreach_tensor_log2_cuda_ + autogen: _foreach_log2.out + +- func: _foreach_max(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_max_slow + CUDA: foreach_tensor_max_cuda + autogen: _foreach_max.out + +- func: _foreach_neg(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_neg_slow + CUDA: foreach_tensor_neg_cuda + +- func: _foreach_neg_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_neg_slow_ + CUDA: foreach_tensor_neg_cuda_ + autogen: _foreach_neg.out + +- func: _foreach_norm.Scalar(Tensor[] self, Scalar ord=2, ScalarType? dtype=None) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_norm_slow + CUDA: foreach_tensor_norm_cuda + MTIA: foreach_tensor_norm_mtia + autogen: _foreach_norm.Scalar_out + +# Like _foreach_norm but returns sum(|x|^ord) without the final root +- func: _foreach_powsum.Scalar(Tensor[] self, Scalar ord=2, ScalarType? dtype=None) -> Tensor[] + device_check: NoCheck + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_powsum_slow + CUDA: foreach_tensor_powsum_cuda + autogen: _foreach_powsum.Scalar_out + +- func: _foreach_pow.List(Tensor[] self, Tensor[] exponent) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_pow_list_kernel_slow + CUDA: foreach_tensor_pow_list_kernel_cuda + +- func: _foreach_pow.Scalar(Tensor[] self, Scalar exponent) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_pow_scalar_kernel_slow + CUDA: foreach_tensor_pow_scalar_kernel_cuda + +- func: _foreach_pow.ScalarList(Tensor[] self, Scalar[] exponent) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_pow_scalarlist_kernel_slow + CUDA: foreach_tensor_pow_scalarlist_kernel_cuda + +- func: _foreach_pow.ScalarAndTensor(Scalar self, Tensor[] exponent) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_scalar_pow_list_kernel_slow + CUDA: foreach_scalar_pow_list_kernel_cuda + +- func: _foreach_pow_.List(Tensor(a!)[] self, Tensor[] exponent) -> () + device_check: NoCheck + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_pow_list_kernel_slow_ + CUDA: foreach_tensor_pow_list_kernel_cuda_ + autogen: _foreach_pow.List_out + +- func: _foreach_pow_.Scalar(Tensor(a!)[] self, Scalar exponent) -> () + device_check: NoCheck + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_pow_scalar_kernel_slow_ + CUDA: foreach_tensor_pow_scalar_kernel_cuda_ + autogen: _foreach_pow.Scalar_out + +- func: _foreach_pow_.ScalarList(Tensor(a!)[] self, Scalar[] exponent) -> () + device_check: NoCheck + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_pow_scalarlist_kernel_slow_ + CUDA: foreach_tensor_pow_scalarlist_kernel_cuda_ + autogen: _foreach_pow.ScalarList_out + +- func: _foreach_reciprocal(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_reciprocal_slow + CUDA: foreach_tensor_reciprocal_cuda + +- func: _foreach_reciprocal_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_reciprocal_slow_ + CUDA: foreach_tensor_reciprocal_cuda_ + autogen: _foreach_reciprocal.out + +- func: _foreach_round(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_round_slow + CUDA: foreach_tensor_round_cuda + +- func: _foreach_round_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_round_slow_ + CUDA: foreach_tensor_round_cuda_ + autogen: _foreach_round.out + +- func: _foreach_rsqrt(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_rsqrt_slow + CUDA: foreach_tensor_rsqrt_cuda + +- func: _foreach_rsqrt_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_rsqrt_slow_ + CUDA: foreach_tensor_rsqrt_cuda_ + autogen: _foreach_rsqrt.out + +- func: _foreach_sigmoid(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sigmoid_slow + CUDA: foreach_tensor_sigmoid_cuda + +- func: _foreach_sigmoid_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sigmoid_slow_ + CUDA: foreach_tensor_sigmoid_cuda_ + autogen: _foreach_sigmoid.out + +- func: _foreach_sign(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sign_slow + CUDA: foreach_tensor_sign_cuda + +- func: _foreach_sign_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sign_slow_ + CUDA: foreach_tensor_sign_cuda_ + autogen: _foreach_sign.out + +- func: _foreach_sin(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sin_slow + CUDA: foreach_tensor_sin_cuda + +- func: _foreach_sin_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sin_slow_ + CUDA: foreach_tensor_sin_cuda_ + autogen: _foreach_sin.out + +- func: _foreach_sinh(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sinh_slow + CUDA: foreach_tensor_sinh_cuda + +- func: _foreach_sinh_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sinh_slow_ + CUDA: foreach_tensor_sinh_cuda_ + autogen: _foreach_sinh.out + +- func: _foreach_sqrt(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sqrt_slow + CUDA: foreach_tensor_sqrt_cuda + +- func: _foreach_sqrt_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_sqrt_slow_ + CUDA: foreach_tensor_sqrt_cuda_ + MTIA: foreach_tensor_sqrt_mtia_ + autogen: _foreach_sqrt.out + +- func: _foreach_tan(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_tan_slow + CUDA: foreach_tensor_tan_cuda + +- func: _foreach_tan_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_tan_slow_ + CUDA: foreach_tensor_tan_cuda_ + autogen: _foreach_tan.out + +- func: _foreach_tanh(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_tanh_slow + CUDA: foreach_tensor_tanh_cuda + +- func: _foreach_tanh_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_tanh_slow_ + CUDA: foreach_tensor_tanh_cuda_ + autogen: _foreach_tanh.out + +- func: _foreach_trunc(Tensor[] self) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_trunc_slow + CUDA: foreach_tensor_trunc_cuda + +- func: _foreach_trunc_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_trunc_slow_ + CUDA: foreach_tensor_trunc_cuda_ + autogen: _foreach_trunc.out + +- func: _foreach_zero_(Tensor(a!)[] self) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_zero_slow_ + CUDA: foreach_tensor_zero_cuda_ + autogen: _foreach_zero, _foreach_zero.out + +- func: _foreach_clone(Tensor[] self, *, MemoryFormat? memory_format=None) -> Tensor[] + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_clone_slow + CUDA: foreach_tensor_clone_cuda + autogen: _foreach_clone.out + +- func: _foreach_copy_(Tensor(a!)[] self, Tensor[] src, bool non_blocking=False) -> () + device_check: NoCheck # foreach kernels fall back to slow path when tensor are on different devices + variants: function + dispatch: + CompositeExplicitAutograd: foreach_tensor_copy_list_kernel_slow_ + CUDA: foreach_tensor_copy_list_kernel_cuda_ + MTIA: foreach_tensor_copy_list_kernel_mtia_ + autogen: _foreach_copy.out + +- func: _foreach_copy(Tensor[] self, Tensor[] src, bool non_blocking=False) -> Tensor[] self_out + device_check: NoCheck + variants: function + dispatch: + CompositeExplicitAutograd: _foreach_copy + MTIA: foreach_tensor_copy_list_kernel_mtia + +- func: bucketize.Tensor(Tensor self, Tensor boundaries, *, bool out_int32=False, bool right=False) -> Tensor + dispatch: + CPU: bucketize_cpu + CUDA: bucketize_cuda + MPS: bucketize_mps + +- func: bucketize.Tensor_out(Tensor self, Tensor boundaries, *, bool out_int32=False, bool right=False, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: bucketize_out_cpu + CUDA: bucketize_out_cuda + MPS: bucketize_out_mps + +- func: bucketize.Scalar(Scalar self, Tensor boundaries, *, bool out_int32=False, bool right=False) -> Tensor + dispatch: + CPU: bucketize_cpu + CUDA: bucketize_cuda + MPS: bucketize_mps + autogen: bucketize.Scalar_out + +- func: searchsorted.Tensor(Tensor sorted_sequence, Tensor self, *, bool out_int32=False, bool right=False, str? side=None, Tensor? sorter=None) -> Tensor + dispatch: + CPU: searchsorted_cpu + CUDA: searchsorted_cuda + MPS: searchsorted_mps + +- func: searchsorted.Tensor_out(Tensor sorted_sequence, Tensor self, *, bool out_int32=False, bool right=False, str? side=None, Tensor? sorter=None, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: searchsorted_out_cpu + CUDA: searchsorted_out_cuda + MPS: searchsorted_out_mps + +- func: searchsorted.Scalar(Tensor sorted_sequence, Scalar self, *, bool out_int32=False, bool right=False, str? side=None, Tensor? sorter=None) -> Tensor + dispatch: + CPU: searchsorted_cpu + CUDA: searchsorted_cuda + MPS: searchsorted_mps + +- func: searchsorted.Scalar_out(Tensor sorted_sequence, Scalar self, *, bool out_int32=False, bool right=False, str? side=None, Tensor? sorter=None, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU: searchsorted_out_cpu + CUDA: searchsorted_out_cuda + MPS: searchsorted_out_mps + +- func: _convert_indices_from_coo_to_csr(Tensor self, int size, *, bool out_int32=False) -> Tensor + structured_delegate: _convert_indices_from_coo_to_csr.out + +- func: _convert_indices_from_coo_to_csr.out(Tensor self, int size, *, bool out_int32=False, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: _convert_indices_from_coo_to_csr_structured_cpu + CUDA: _convert_indices_from_coo_to_csr_structured_cuda + +- func: _convert_indices_from_csr_to_coo(Tensor crow_indices, Tensor col_indices, *, bool out_int32=False, bool transpose=False) -> Tensor + structured_delegate: _convert_indices_from_csr_to_coo.out + +- func: _convert_indices_from_csr_to_coo.out(Tensor crow_indices, Tensor col_indices, *, bool out_int32=False, bool transpose=False, Tensor(a!) out) -> Tensor(a!) + structured: True + dispatch: + CPU: _convert_indices_from_csr_to_coo_structured_cpu + CUDA: _convert_indices_from_csr_to_coo_structured_cuda + +## NN wrappers + +- func: mse_loss.out(Tensor self, Tensor target, int reduction=Mean, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA: mse_loss_out + MPS: mse_loss_out_mps + +- func: mse_loss(Tensor self, Tensor target, int reduction=Mean) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: mse_loss.out + python_module: nn + +- func: mse_loss_backward.grad_input(Tensor grad_output, Tensor self, Tensor target, int reduction, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU, CUDA: mse_loss_backward_out + MPS: mse_loss_backward_out_mps + +- func: mse_loss_backward(Tensor grad_output, Tensor self, Tensor target, int reduction) -> Tensor + python_module: nn + dispatch: + CPU, CUDA: mse_loss_backward + MPS: mse_loss_backward_mps + +- func: l1_loss(Tensor self, Tensor target, int reduction=Mean) -> Tensor + python_module: nn + +- func: multi_margin_loss.out(Tensor self, Tensor target, Scalar p=1, Scalar margin=1, Tensor? weight=None, int reduction=Mean, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU: multi_margin_loss_cpu_out + CUDA: multi_margin_loss_cuda_out + +- func: multi_margin_loss(Tensor self, Tensor target, Scalar p=1, Scalar margin=1, Tensor? weight=None, int reduction=Mean) -> Tensor + python_module: nn + dispatch: + CPU: multi_margin_loss_cpu + CUDA: multi_margin_loss_cuda + +- func: multi_margin_loss_backward.grad_input(Tensor grad_output, Tensor self, Tensor target, Scalar p, Scalar margin, Tensor? weight=None, int reduction=Mean, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: multi_margin_loss_cpu_backward_out + CUDA: multi_margin_loss_cuda_backward_out + +- func: multi_margin_loss_backward(Tensor grad_output, Tensor self, Tensor target, Scalar p, Scalar margin, Tensor? weight=None, int reduction=Mean) -> Tensor + python_module: nn + dispatch: + CPU: multi_margin_loss_cpu_backward + CUDA: multi_margin_loss_cuda_backward + +- func: multilabel_margin_loss.out(Tensor self, Tensor target, int reduction=Mean, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + +- func: multilabel_margin_loss(Tensor self, Tensor target, int reduction=Mean) -> Tensor + python_module: nn + +- func: multilabel_margin_loss_forward.output(Tensor self, Tensor target, int reduction, *, Tensor(a!) output, Tensor(b!) is_target) -> (Tensor(a!), Tensor(b!)) + python_module: nn + dispatch: + CPU: multilabel_margin_loss_forward_out_cpu + CUDA: multilabel_margin_loss_forward_out_cuda + +- func: multilabel_margin_loss_forward(Tensor self, Tensor target, int reduction) -> (Tensor output, Tensor is_target) + python_module: nn + dispatch: + CPU: multilabel_margin_loss_forward_cpu + CUDA: multilabel_margin_loss_forward_cuda + +- func: multilabel_margin_loss_backward.grad_input(Tensor grad_output, Tensor self, Tensor target, int reduction, Tensor is_target, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: multilabel_margin_loss_backward_cpu_out + CUDA: multilabel_margin_loss_backward_cuda_out + +- func: multilabel_margin_loss_backward(Tensor grad_output, Tensor self, Tensor target, int reduction, Tensor is_target) -> Tensor + python_module: nn + dispatch: + CPU: multilabel_margin_loss_backward_cpu + CUDA: multilabel_margin_loss_backward_cuda + +- func: nll_loss.out(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + +- func: nll_loss_nd(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: nll_loss_nd_symint + +- func: nll_loss(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: nll_loss_symint + +- func: nll_loss_forward.output(Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index, *, Tensor(a!) output, Tensor(b!) total_weight) -> (Tensor(a!), Tensor(b!)) + python_module: nn + structured: True + dispatch: + CPU: nll_loss_forward_out_cpu + CUDA: nll_loss_forward_out_cuda + MPS: nll_loss_forward_out_mps + +- func: nll_loss_forward(Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index) -> (Tensor output, Tensor total_weight) + python_module: nn + structured_delegate: nll_loss_forward.output + +- func: nll_loss_backward.grad_input(Tensor grad_output, Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index, Tensor total_weight, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: nll_loss_backward_out_cpu + CUDA: nll_loss_backward_out_cuda + MPS: nll_loss_backward_out_mps + +- func: nll_loss_backward(Tensor grad_output, Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index, Tensor total_weight) -> Tensor + python_module: nn + structured_delegate: nll_loss_backward.grad_input + +- func: nll_loss2d.out(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + +- func: nll_loss2d(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean, SymInt ignore_index=-100) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: nll_loss2d_symint + +- func: nll_loss2d_forward.output(Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index, *, Tensor(a!) output, Tensor(b!) total_weight) -> (Tensor(a!), Tensor(b!)) + python_module: nn + dispatch: + CPU: nll_loss2d_forward_out_cpu + CUDA: nll_loss2d_forward_out_cuda + MPS: nll_loss2d_forward_out_mps + +- func: nll_loss2d_forward(Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index) -> (Tensor output, Tensor total_weight) + python_module: nn + dispatch: + CPU: nll_loss2d_forward_cpu + CUDA: nll_loss2d_forward_cuda + MPS: nll_loss2d_forward_mps + +- func: nll_loss2d_backward.grad_input(Tensor grad_output, Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index, Tensor total_weight, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: nll_loss2d_backward_out_cpu + CUDA: nll_loss2d_backward_out_cuda + MPS: nll_loss2d_backward_out_mps + +- func: nll_loss2d_backward(Tensor grad_output, Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index, Tensor total_weight) -> Tensor + python_module: nn + dispatch: + CPU: nll_loss2d_backward_cpu + CUDA: nll_loss2d_backward_cuda + MPS: nll_loss2d_backward_mps + +- func: smooth_l1_loss.out(Tensor self, Tensor target, int reduction=Mean, float beta=1.0, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA: smooth_l1_loss_out + MPS: smooth_l1_loss_out_mps + +- func: smooth_l1_loss(Tensor self, Tensor target, int reduction=Mean, float beta=1.0) -> Tensor + device_check: NoCheck # TensorIterator + structured_delegate: smooth_l1_loss.out + python_module: nn + +- func: smooth_l1_loss_backward.grad_input(Tensor grad_output, Tensor self, Tensor target, int reduction, float beta, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: smooth_l1_loss_backward_out + CUDA: smooth_l1_loss_backward_out + MPS: smooth_l1_loss_backward_out_mps + +- func: smooth_l1_loss_backward(Tensor grad_output, Tensor self, Tensor target, int reduction, float beta) -> Tensor + python_module: nn + dispatch: + CompositeExplicitAutograd: smooth_l1_loss_backward + +- func: huber_loss.out(Tensor self, Tensor target, int reduction=Mean, float delta=1.0, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU, CUDA: huber_loss_out + MPS: huber_loss_out_mps + +- func: huber_loss(Tensor self, Tensor target, int reduction=Mean, float delta=1.0) -> Tensor + python_module: nn + dispatch: + CPU, CUDA: huber_loss + MPS: huber_loss_mps + +- func: huber_loss_backward.out(Tensor grad_output, Tensor self, Tensor target, int reduction, float delta, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU, CUDA: huber_loss_backward_out + MPS: huber_loss_backward_out_mps + +- func: huber_loss_backward(Tensor grad_output, Tensor self, Tensor target, int reduction, float delta) -> Tensor + python_module: nn + dispatch: + CompositeExplicitAutograd: huber_loss_backward + +- func: soft_margin_loss.out(Tensor self, Tensor target, int reduction=Mean, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CompositeExplicitAutograd: soft_margin_loss_out + +- func: soft_margin_loss(Tensor self, Tensor target, int reduction=Mean) -> Tensor + python_module: nn + dispatch: + CompositeExplicitAutograd: soft_margin_loss + +- func: soft_margin_loss_backward.grad_input(Tensor grad_output, Tensor self, Tensor target, int reduction, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CompositeExplicitAutograd: soft_margin_loss_backward_out + +- func: soft_margin_loss_backward(Tensor grad_output, Tensor self, Tensor target, int reduction) -> Tensor + python_module: nn + dispatch: + CompositeExplicitAutograd: soft_margin_loss_backward + +- func: elu.out(Tensor self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: elu_out + +- func: elu(Tensor self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> Tensor + structured_delegate: elu.out + device_check: NoCheck # TensorIterator + python_module: nn + tags: [core, pointwise] + +- func: elu_backward.grad_input(Tensor grad_output, Scalar alpha, Scalar scale, Scalar input_scale, bool is_result, Tensor self_or_result, *, Tensor(a!) grad_input) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA, MPS: elu_backward_out + +- func: elu_backward(Tensor grad_output, Scalar alpha, Scalar scale, Scalar input_scale, bool is_result, Tensor self_or_result) -> Tensor + structured_delegate: elu_backward.grad_input + python_module: nn + +- func: elu_(Tensor(a!) self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> Tensor(a!) + structured_delegate: elu.out + device_check: NoCheck # TensorIterator + python_module: nn + tags: pointwise + +- func: glu.out(Tensor self, int dim=-1, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA: glu_out + MPS: glu_out_mps + +- func: glu(Tensor self, int dim=-1) -> Tensor + structured_delegate: glu.out + device_check: NoCheck # TensorIterator + python_module: nn + +- func: glu_backward.grad_input(Tensor grad_output, Tensor self, int dim, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: glu_backward_cpu_out + CUDA: glu_backward_cuda_out + MPS: glu_backward_mps_out + +- func: glu_backward(Tensor grad_output, Tensor self, int dim) -> Tensor + python_module: nn + dispatch: + CPU: glu_backward_cpu + CUDA: glu_backward_cuda + MPS: glu_backward_mps + +- func: glu_jvp(Tensor glu, Tensor x, Tensor dx, int dim) -> Tensor + python_module: nn + dispatch: + CPU, CUDA: glu_jvp + autogen: glu_jvp.out + +- func: glu_backward_jvp(Tensor grad_x, Tensor grad_glu, Tensor x, Tensor dgrad_glu, Tensor dx, int dim) -> Tensor + python_module: nn + dispatch: + CPU, CUDA: glu_backward_jvp + autogen: glu_backward_jvp.out + +- func: hardsigmoid.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: hardsigmoid_out + QuantizedCPU: hardsigmoid_out_quantized_cpu + +- func: hardsigmoid(Tensor self) -> Tensor + structured_delegate: hardsigmoid.out + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + QuantizedCPU: hardsigmoid_quantized_cpu + tags: pointwise + +- func: hardsigmoid_(Tensor(a!) self) -> Tensor(a!) + structured_delegate: hardsigmoid.out + device_check: NoCheck # TensorIterator + python_module: nn + tags: pointwise + +- func: hardsigmoid_backward.grad_input(Tensor grad_output, Tensor self, *, Tensor(a!) grad_input) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA, MPS: hardsigmoid_backward_out + +- func: hardsigmoid_backward(Tensor grad_output, Tensor self) -> Tensor + structured_delegate: hardsigmoid_backward.grad_input + python_module: nn + +- func: hardtanh.out(Tensor self, Scalar min_val=-1, Scalar max_val=1, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: hardtanh_out + QuantizedCPU: hardtanh_out_quantized_cpu + +- func: hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> Tensor + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: hardtanh + QuantizedCPU: hardtanh_quantized_cpu + tags: [pointwise, core] + +- func: hardtanh_backward.grad_input(Tensor grad_output, Tensor self, Scalar min_val, Scalar max_val, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU, CUDA: hardtanh_backward_out + MPS: hardtanh_backward_out_mps + +- func: hardtanh_backward(Tensor grad_output, Tensor self, Scalar min_val, Scalar max_val) -> Tensor + python_module: nn + dispatch: + CPU, CUDA: hardtanh_backward + MPS: hardtanh_backward_mps + +- func: hardtanh_(Tensor(a!) self, Scalar min_val=-1, Scalar max_val=1) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: hardtanh_ + QuantizedCPU: hardtanh_quantized_cpu_ + tags: pointwise + +- func: hardswish.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: hardswish_out + +- func: hardswish(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: hardswish + +- func: hardswish_(Tensor(a!) self) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: hardswish_ + +- func: hardswish_backward(Tensor grad_output, Tensor self) -> Tensor + python_module: nn + dispatch: + CPU, CUDA, MPS: hardswish_backward + autogen: hardswish_backward.out + +- func: leaky_relu.out(Tensor self, Scalar negative_slope=0.01, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: leaky_relu_out + QuantizedCPU: leaky_relu_out_quantized_cpu + +- func: leaky_relu(Tensor self, Scalar negative_slope=0.01) -> Tensor + structured_delegate: leaky_relu.out + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + QuantizedCPU: leaky_relu_quantized_cpu + tags: [core, pointwise] + +- func: leaky_relu_backward.grad_input(Tensor grad_output, Tensor self, Scalar negative_slope, bool self_is_result, *, Tensor(a!) grad_input) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA, MPS: leaky_relu_backward_out + +- func: leaky_relu_backward(Tensor grad_output, Tensor self, Scalar negative_slope, bool self_is_result) -> Tensor + structured_delegate: leaky_relu_backward.grad_input + python_module: nn + +- func: leaky_relu_(Tensor(a!) self, Scalar negative_slope=0.01) -> Tensor(a!) + structured_delegate: leaky_relu.out + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + QuantizedCPU: leaky_relu_quantized_cpu_ + tags: pointwise + +- func: log_sigmoid.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: nn + +- func: log_sigmoid(Tensor self) -> Tensor + device_check: NoCheck # TensorIterator + python_module: nn + +- func: log_sigmoid_forward.output(Tensor self, *, Tensor(a!) output, Tensor(b!) buffer) -> (Tensor(a!), Tensor(b!)) + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU: log_sigmoid_forward_out_cpu + CUDA: log_sigmoid_forward_out_cuda + MPS: log_sigmoid_forward_out_mps + +- func: log_sigmoid_forward(Tensor self) -> (Tensor output, Tensor buffer) + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU: log_sigmoid_forward_cpu + CUDA: log_sigmoid_forward_cuda + MPS: log_sigmoid_forward_mps + +- func: log_sigmoid_backward.grad_input(Tensor grad_output, Tensor self, Tensor buffer, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: log_sigmoid_backward_cpu_out + CUDA: log_sigmoid_backward_cuda_out + MPS: log_sigmoid_backward_mps_out + +- func: log_sigmoid_backward(Tensor grad_output, Tensor self, Tensor buffer) -> Tensor + python_module: nn + dispatch: + CPU: log_sigmoid_backward_cpu + CUDA: log_sigmoid_backward_cuda + MPS: log_sigmoid_backward_mps + +- func: rrelu_with_noise.out(Tensor self, Tensor(b!) noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + tags: nondeterministic_seeded + dispatch: + CPU: rrelu_with_noise_out_cpu + CUDA: rrelu_with_noise_out_cuda + +- func: rrelu_with_noise(Tensor self, Tensor(b!) noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor + python_module: nn + dispatch: + CPU: rrelu_with_noise_cpu + CUDA: rrelu_with_noise_cuda + tags: nondeterministic_seeded + autogen: rrelu_with_noise_functional + +- func: rrelu_with_noise_backward(Tensor grad_output, Tensor self, Tensor noise, Scalar lower, Scalar upper, bool training, bool self_is_result) -> Tensor + python_module: nn + dispatch: + CompositeExplicitAutograd: rrelu_with_noise_backward + autogen: rrelu_with_noise_backward.out + +- func: rrelu_with_noise_(Tensor(a!) self, Tensor(b!) noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor(a!) + python_module: nn + tags: nondeterministic_seeded + dispatch: + CPU: rrelu_with_noise_cpu_ + CUDA: rrelu_with_noise_cuda_ + +- func: softplus.out(Tensor self, Scalar beta=1, Scalar threshold=20, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA: softplus_out + MPS: softplus_out_mps + +- func: softplus(Tensor self, Scalar beta=1, Scalar threshold=20) -> Tensor + structured_delegate: softplus.out + device_check: NoCheck # TensorIterator + python_module: nn + tags: pointwise + +- func: softplus_backward.grad_input(Tensor grad_output, Tensor self, Scalar beta, Scalar threshold, *, Tensor(a!) grad_input) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA: softplus_backward_out + MPS: softplus_backward_out_mps + +- func: softplus_backward(Tensor grad_output, Tensor self, Scalar beta, Scalar threshold) -> Tensor + structured_delegate: softplus_backward.grad_input + python_module: nn + +- func: softshrink.out(Tensor self, Scalar lambd=0.5, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + device_check: NoCheck # TensorIterator + python_module: nn + dispatch: + CPU, CUDA, MPS: softshrink_out + +- func: softshrink(Tensor self, Scalar lambd=0.5) -> Tensor + structured_delegate: softshrink.out + device_check: NoCheck # TensorIterator + python_module: nn + tags: pointwise + +- func: softshrink_backward.grad_input(Tensor grad_output, Tensor self, Scalar lambd, *, Tensor(a!) grad_input) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: nn + dispatch: + CPU, CUDA, MPS: softshrink_backward_out + +- func: softshrink_backward(Tensor grad_output, Tensor self, Scalar lambd) -> Tensor + structured_delegate: softshrink_backward.grad_input + python_module: nn + +- func: adaptive_avg_pool2d.out(Tensor self, SymInt[2] output_size, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU: adaptive_avg_pool2d_out_cpu + CUDA: adaptive_avg_pool2d_out_cuda + MPS: adaptive_avg_pool2d_out_mps + MkldnnCPU: mkldnn_adaptive_avg_pool2d_out_stub + +- func: adaptive_avg_pool2d(Tensor self, SymInt[2] output_size) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: adaptive_avg_pool2d_symint + +- func: mkldnn_adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor + dispatch: + MkldnnCPU: mkldnn_adaptive_avg_pool2d + +- func: mkldnn_adaptive_avg_pool2d.out(Tensor self, int[2] output_size, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + MkldnnCPU: mkldnn_adaptive_avg_pool2d_out + +- func: mkldnn_adaptive_avg_pool2d_backward(Tensor grad_output, Tensor self) -> Tensor + dispatch: + MkldnnCPU: mkldnn_adaptive_avg_pool2d_backward + autogen: mkldnn_adaptive_avg_pool2d_backward.out + +- func: _adaptive_avg_pool2d(Tensor self, SymInt[2] output_size) -> Tensor + dispatch: + CPU: adaptive_avg_pool2d_cpu + CUDA: adaptive_avg_pool2d_cuda + MPS: adaptive_avg_pool2d_mps + QuantizedCPU: adaptive_avg_pool2d_quantized_cpu + QuantizedCUDA: adaptive_avg_pool2d_quantized_cuda + autogen: _adaptive_avg_pool2d.out + tags: core + +- func: _adaptive_avg_pool2d_backward(Tensor grad_output, Tensor self) -> Tensor + python_module: nn + dispatch: + CPU: adaptive_avg_pool2d_backward_cpu + CUDA: adaptive_avg_pool2d_backward_cuda + MPS: adaptive_avg_pool2d_backward_mps + autogen: _adaptive_avg_pool2d_backward.out + tags: core + +- func: adaptive_avg_pool3d.out(Tensor self, SymInt[3] output_size, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU: adaptive_avg_pool3d_out_cpu + CUDA: adaptive_avg_pool3d_out_cuda + QuantizedCPU: adaptive_avg_pool3d_out_quantized_cpu + +- func: adaptive_avg_pool3d(Tensor self, SymInt[3] output_size) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: adaptive_avg_pool3d_symint + +- func: _adaptive_avg_pool3d(Tensor self, SymInt[3] output_size) -> Tensor + dispatch: + CPU: adaptive_avg_pool3d_cpu + CUDA: adaptive_avg_pool3d_cuda + QuantizedCPU: adaptive_avg_pool3d_quantized_cpu + autogen: _adaptive_avg_pool3d.out + tags: core + +- func: adaptive_avg_pool3d_backward.grad_input(Tensor grad_output, Tensor self, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: adaptive_avg_pool3d_backward_out_cpu + CUDA: adaptive_avg_pool3d_backward_out_cuda + +- func: _adaptive_avg_pool3d_backward(Tensor grad_output, Tensor self) -> Tensor + python_module: nn + dispatch: + CPU: adaptive_avg_pool3d_backward_cpu + CUDA: adaptive_avg_pool3d_backward_cuda + autogen: _adaptive_avg_pool3d_backward.out + +# Return: (Tensor output, Tensor indices) +- func: adaptive_max_pool2d.out(Tensor self, int[2] output_size, *, Tensor(a!) out, Tensor(b!) indices) -> (Tensor(a!), Tensor(b!)) + python_module: nn + structured: True + dispatch: + CPU: adaptive_max_pool2d_out_cpu + CUDA: adaptive_max_pool2d_out_cuda + MPS: adaptive_max_pool2d_out_mps + +# Return: (Tensor output, Tensor indices) +- func: adaptive_max_pool2d(Tensor self, int[2] output_size) -> (Tensor, Tensor) + python_module: nn + structured_delegate: adaptive_max_pool2d.out + +- func: adaptive_max_pool2d_backward.grad_input(Tensor grad_output, Tensor self, Tensor indices, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: adaptive_max_pool2d_backward_out_cpu + CUDA: adaptive_max_pool2d_backward_out_cuda + MPS: adaptive_max_pool2d_backward_out_mps + +- func: adaptive_max_pool2d_backward(Tensor grad_output, Tensor self, Tensor indices) -> Tensor + python_module: nn + structured_delegate: adaptive_max_pool2d_backward.grad_input + +# Return: (Tensor output, Tensor indices) +- func: adaptive_max_pool3d.out(Tensor self, int[3] output_size, *, Tensor(a!) out, Tensor(b!) indices) -> (Tensor(a!), Tensor(b!)) + python_module: nn + structured: True + dispatch: + CPU: adaptive_max_pool3d_out_cpu + CUDA: adaptive_max_pool3d_out_cuda + +# Return: (Tensor output, Tensor indices) +- func: adaptive_max_pool3d(Tensor self, int[3] output_size) -> (Tensor, Tensor) + python_module: nn + structured_delegate: adaptive_max_pool3d.out + +- func: adaptive_max_pool3d_backward.grad_input(Tensor grad_output, Tensor self, Tensor indices, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: adaptive_max_pool3d_backward_out_cpu + CUDA: adaptive_max_pool3d_backward_out_cuda + +- func: adaptive_max_pool3d_backward(Tensor grad_output, Tensor self, Tensor indices) -> Tensor + python_module: nn + structured_delegate: adaptive_max_pool3d_backward.grad_input + +- func: avg_pool2d.out(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + precomputed: + - kernel_size -> int kH, int kW + - stride -> int dH, int dW + - padding -> int padH, int padW + dispatch: + CPU: avg_pool2d_out_cpu + CUDA: avg_pool2d_out_cuda + MPS: avg_pool2d_out_mps + MkldnnCPU: mkldnn_avg_pool2d_out + +- func: avg_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> Tensor + python_module: nn + structured_delegate: avg_pool2d.out + dispatch: + MkldnnCPU: mkldnn_avg_pool2d + QuantizedCPU: avg_pool2d_quantized_cpu + tags: core + +- func: avg_pool2d_backward.grad_input(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] stride, int[2] padding, bool ceil_mode, bool count_include_pad, int? divisor_override, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: avg_pool2d_backward_out_cpu + CUDA: avg_pool2d_backward_out_cuda + MPS: avg_pool2d_backward_out_mps + MkldnnCPU: mkldnn_avg_pool2d_backward_out + +- func: avg_pool2d_backward(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] stride, int[2] padding, bool ceil_mode, bool count_include_pad, int? divisor_override) -> Tensor + python_module: nn + structured_delegate: avg_pool2d_backward.grad_input + dispatch: + MkldnnCPU: mkldnn_avg_pool2d_backward + tags: core + +- func: avg_pool3d.out(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: avg_pool3d_out_cpu + CUDA: avg_pool3d_out_cuda + MPS: avg_pool3d_out_mps + MkldnnCPU: mkldnn_avg_pool3d_out + +- func: avg_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> Tensor + python_module: nn + structured_delegate: avg_pool3d.out + dispatch: + MkldnnCPU: mkldnn_avg_pool3d + QuantizedCPU: avg_pool3d_quantized_cpu + tags: core + +- func: avg_pool3d_backward.grad_input(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] stride, int[3] padding, bool ceil_mode, bool count_include_pad, int? divisor_override, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: avg_pool3d_backward_out_cpu + CUDA: avg_pool3d_backward_out_cuda + MPS: avg_pool3d_backward_out_mps + MkldnnCPU: mkldnn_avg_pool3d_backward_out + +- func: avg_pool3d_backward(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] stride, int[3] padding, bool ceil_mode, bool count_include_pad, int? divisor_override) -> Tensor + python_module: nn + structured_delegate: avg_pool3d_backward.grad_input + dispatch: + MkldnnCPU: mkldnn_avg_pool3d_backward + +# Return: (Tensor output, Tensor indices) +- func: fractional_max_pool2d.output(Tensor self, int[2] kernel_size, int[2] output_size, Tensor random_samples, *, Tensor(a!) output, Tensor(b!) indices) -> (Tensor(a!), Tensor(b!)) + python_module: nn + structured: True + dispatch: + CPU: fractional_max_pool2d_out_cpu + CUDA: fractional_max_pool2d_out_cuda + +# Return: (Tensor output, Tensor indices) +- func: fractional_max_pool2d(Tensor self, int[2] kernel_size, int[2] output_size, Tensor random_samples) -> (Tensor, Tensor) + python_module: nn + structured_delegate: fractional_max_pool2d.output + +- func: fractional_max_pool2d_backward.grad_input(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] output_size, Tensor indices, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: fractional_max_pool2d_backward_cpu + CUDA: fractional_max_pool2d_backward_cuda + +- func: fractional_max_pool2d_backward(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] output_size, Tensor indices) -> Tensor + python_module: nn + structured_delegate: fractional_max_pool2d_backward.grad_input + +# Return: (Tensor output, Tensor indices) +- func: fractional_max_pool3d.output(Tensor self, int[3] kernel_size, int[3] output_size, Tensor random_samples, *, Tensor(a!) output, Tensor(b!) indices) -> (Tensor(a!), Tensor(b!)) + python_module: nn + structured: True + precomputed: + - kernel_size -> int poolSizeT, int poolSizeH, int poolSizeW + - output_size -> int outputT, int outputH, int outputW + - int numBatch, int numPlanes, int inputT, int inputH, int inputW + dispatch: + CPU: fractional_max_pool3d_out_cpu + CUDA: fractional_max_pool3d_out_cuda + +# Return: (Tensor output, Tensor indices) +- func: fractional_max_pool3d(Tensor self, int[3] kernel_size, int[3] output_size, Tensor random_samples) -> (Tensor, Tensor) + python_module: nn + structured_delegate: fractional_max_pool3d.output + +- func: fractional_max_pool3d_backward.grad_input(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] output_size, Tensor indices, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: fractional_max_pool3d_backward_out_cpu + CUDA: fractional_max_pool3d_backward_out_cuda + +- func: fractional_max_pool3d_backward(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] output_size, Tensor indices) -> Tensor + python_module: nn + dispatch: + CPU: fractional_max_pool3d_backward_cpu + CUDA: fractional_max_pool3d_backward_cuda + +# Return: (Tensor output, Tensor indices) +- func: max_pool2d_with_indices.out(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False, *, Tensor(a!) out, Tensor(b!) indices) -> (Tensor(a!), Tensor(b!)) + python_module: nn + structured: True + dispatch: + CPU: max_pool2d_with_indices_out_cpu + CUDA: max_pool2d_with_indices_out_cuda + MPS: max_pool2d_with_indices_out_mps + +# Return: (Tensor output, Tensor indices) +- func: max_pool2d_with_indices(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor) + python_module: nn + structured_delegate: max_pool2d_with_indices.out + tags: core + +- func: max_pool2d_with_indices_backward.grad_input(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] stride, int[2] padding, int[2] dilation, bool ceil_mode, Tensor indices, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: max_pool2d_with_indices_backward_out_cpu + CUDA: max_pool2d_with_indices_backward_out_cuda + MPS: max_pool2d_with_indices_backward_out_mps + +- func: max_pool2d_with_indices_backward(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] stride, int[2] padding, int[2] dilation, bool ceil_mode, Tensor indices) -> Tensor + python_module: nn + structured_delegate: max_pool2d_with_indices_backward.grad_input + tags: core + +# Return: (Tensor output, Tensor indices) +- func: max_pool3d_with_indices.out(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False, *, Tensor(a!) out, Tensor(b!) indices) -> (Tensor(a!), Tensor(b!)) + python_module: nn + dispatch: + CPU: max_pool3d_with_indices_out_cpu + CUDA: max_pool3d_with_indices_out_cuda + MPS: max_pool3d_with_indices_out_mps + +# Return: (Tensor output, Tensor indices) +- func: max_pool3d_with_indices(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor) + python_module: nn + dispatch: + CPU: max_pool3d_with_indices_cpu + CUDA: max_pool3d_with_indices_cuda + MPS: max_pool3d_with_indices_mps + tags: core + +- func: max_pool3d_with_indices_backward.grad_input(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] stride, int[3] padding, int[3] dilation, bool ceil_mode, Tensor indices, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: max_pool3d_with_indices_backward_out_cpu + CUDA: max_pool3d_with_indices_backward_out_cuda + MPS: max_pool3d_with_indices_backward_out_mps + +- func: max_pool3d_with_indices_backward(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] stride, int[3] padding, int[3] dilation, bool ceil_mode, Tensor indices) -> Tensor + python_module: nn + dispatch: + CPU: max_pool3d_with_indices_backward_cpu + CUDA: max_pool3d_with_indices_backward_cuda + MPS: max_pool3d_with_indices_backward_mps + +- func: max_unpool2d.out(Tensor self, Tensor indices, SymInt[2] output_size, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU: max_unpooling2d_forward_out_cpu + CUDA: max_unpooling2d_forward_out_cuda + MPS: max_unpooling2d_forward_out_mps + +- func: max_unpool2d(Tensor self, Tensor indices, SymInt[2] output_size) -> Tensor + python_module: nn + dispatch: + CPU: max_unpooling2d_forward_cpu + CUDA: max_unpooling2d_forward_cuda + MPS: max_unpooling2d_forward_mps + +- func: max_unpool3d.out(Tensor self, Tensor indices, SymInt[3] output_size, int[3] stride, int[3] padding, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU: max_unpooling3d_forward_out_cpu + CUDA: max_unpooling3d_forward_out_cuda + MPS: max_unpooling3d_forward_out_mps + +- func: max_unpool3d(Tensor self, Tensor indices, SymInt[3] output_size, int[3] stride, int[3] padding) -> Tensor + python_module: nn + dispatch: + CPU: max_unpooling3d_forward_cpu + CUDA: max_unpooling3d_forward_cuda + MPS: max_unpooling3d_forward_mps + +- func: reflection_pad1d.out(Tensor self, SymInt[2] padding, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: reflection_pad1d_out_cpu + QuantizedCPU: reflection_pad1d_out_quantized_cpu + CUDA: reflection_pad1d_out_cuda + MPS: reflection_pad1d_out_mps + +- func: reflection_pad1d(Tensor self, SymInt[2] padding) -> Tensor + python_module: nn + structured_delegate: reflection_pad1d.out + tags: core + +- func: reflection_pad1d_backward.grad_input(Tensor grad_output, Tensor self, SymInt[2] padding, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: reflection_pad1d_backward_out_cpu + CUDA: reflection_pad1d_backward_out_cuda + MPS: reflection_pad1d_backward_out_mps + +- func: reflection_pad1d_backward(Tensor grad_output, Tensor self, SymInt[2] padding) -> Tensor + python_module: nn + structured_delegate: reflection_pad1d_backward.grad_input + +- func: reflection_pad2d.out(Tensor self, SymInt[4] padding, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU, QuantizedCPU: reflection_pad2d_out_cpu + CUDA: reflection_pad2d_out_cuda + MPS: reflection_pad2d_out_mps + +- func: reflection_pad2d(Tensor self, SymInt[4] padding) -> Tensor + python_module: nn + dispatch: + CPU: reflection_pad2d_cpu + QuantizedCPU: reflection_pad2d_quantized_cpu + CUDA: reflection_pad2d_cuda + MPS: reflection_pad2d_mps + tags: core + +- func: reflection_pad2d_backward.grad_input(Tensor grad_output, Tensor self, SymInt[4] padding, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: reflection_pad2d_backward_out_cpu + CUDA: reflection_pad2d_backward_out_cuda + MPS: reflection_pad2d_backward_out_mps + +- func: reflection_pad2d_backward(Tensor grad_output, Tensor self, SymInt[4] padding) -> Tensor + python_module: nn + dispatch: + CPU: reflection_pad2d_backward_cpu + CUDA: reflection_pad2d_backward_cuda + MPS: reflection_pad2d_backward_mps + +- func: reflection_pad3d.out(Tensor self, SymInt[6] padding, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: reflection_pad3d_out_cpu + CUDA: reflection_pad3d_out_cuda + MPS: reflection_pad3d_out_mps + +- func: reflection_pad3d(Tensor self, SymInt[6] padding) -> Tensor + python_module: nn + structured_delegate: reflection_pad3d.out + tags: core + +- func: reflection_pad3d_backward.grad_input(Tensor grad_output, Tensor self, SymInt[6] padding, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: reflection_pad3d_backward_out_cpu + CUDA: reflection_pad3d_backward_out_cuda + MPS: reflection_pad3d_backward_out_mps + +- func: reflection_pad3d_backward(Tensor grad_output, Tensor self, SymInt[6] padding) -> Tensor + python_module: nn + structured_delegate: reflection_pad3d_backward.grad_input + +- func: replication_pad1d.out(Tensor self, SymInt[2] padding, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: replication_pad1d_out_cpu + CUDA: replication_pad1d_out_cuda + MPS: replication_pad1d_out_mps + +- func: replication_pad1d(Tensor self, SymInt[2] padding) -> Tensor + python_module: nn + structured_delegate: replication_pad1d.out + +- func: replication_pad1d_backward.grad_input(Tensor grad_output, Tensor self, SymInt[2] padding, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: replication_pad1d_backward_out_cpu + CUDA: replication_pad1d_backward_out_cuda + MPS: replication_pad1d_backward_out_mps + +- func: replication_pad1d_backward(Tensor grad_output, Tensor self, SymInt[2] padding) -> Tensor + python_module: nn + structured_delegate: replication_pad1d_backward.grad_input + +- func: replication_pad2d.out(Tensor self, SymInt[4] padding, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: replication_pad2d_out_cpu + CUDA: replication_pad2d_out_cuda + MPS: replication_pad2d_out_mps + +- func: replication_pad2d(Tensor self, SymInt[4] padding) -> Tensor + python_module: nn + structured_delegate: replication_pad2d.out + tags: core + +- func: replication_pad2d_backward.grad_input(Tensor grad_output, Tensor self, SymInt[4] padding, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: replication_pad2d_backward_out_cpu + CUDA: replication_pad2d_backward_out_cuda + MPS: replication_pad2d_backward_out_mps + +- func: replication_pad2d_backward(Tensor grad_output, Tensor self, SymInt[4] padding) -> Tensor + python_module: nn + dispatch: + CPU: replication_pad2d_backward_cpu + CUDA: replication_pad2d_backward_cuda + MPS: replication_pad2d_backward_mps + +- func: replication_pad3d.out(Tensor self, SymInt[6] padding, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: replication_pad3d_out_cpu + CUDA: replication_pad3d_out_cuda + MPS: replication_pad3d_out_mps + +- func: replication_pad3d(Tensor self, SymInt[6] padding) -> Tensor + python_module: nn + structured_delegate: replication_pad3d.out + tags: core + + +- func: replication_pad3d_backward.grad_input(Tensor grad_output, Tensor self, SymInt[6] padding, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + dispatch: + CPU: replication_pad3d_backward_out_cpu + CUDA: replication_pad3d_backward_out_cuda + MPS: replication_pad3d_backward_out_mps + +- func: replication_pad3d_backward(Tensor grad_output, Tensor self, SymInt[6] padding) -> Tensor + python_module: nn + dispatch: + CPU: replication_pad3d_backward_cpu + CUDA: replication_pad3d_backward_cuda + MPS: replication_pad3d_backward_mps + +- func: _pad_circular(Tensor self, SymInt[] pad) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: _pad_circular_symint + +- func: _pad_enum(Tensor self, SymInt[] pad, int mode, float? value=None) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: _pad_enum_symint + +- func: pad(Tensor self, SymInt[] pad, str mode="constant", float? value=None) -> Tensor + python_module: nn + dispatch: + CompositeImplicitAutograd: pad_symint + +- func: upsample_linear1d.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor + python_module: nn + autogen: upsample_linear1d.vec_out + +- func: upsample_bilinear2d.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor + python_module: nn + autogen: upsample_bilinear2d.vec_out + tags: core + +- func: _upsample_bilinear2d_aa.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor + python_module: nn + autogen: _upsample_bilinear2d_aa.vec_out + +- func: upsample_trilinear3d.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor + python_module: nn + autogen: upsample_trilinear3d.vec_out + +- func: upsample_bicubic2d.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor + python_module: nn + autogen: upsample_bicubic2d.vec_out + +- func: _upsample_bicubic2d_aa.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor + python_module: nn + autogen: _upsample_bicubic2d_aa.vec_out + +- func: _upsample_lanczos2d_aa.vec(Tensor input, SymInt[]? output_size, bool align_corners, float[]? scale_factors) -> Tensor + python_module: nn + autogen: _upsample_lanczos2d_aa.vec_out + +- func: upsample_nearest1d.vec(Tensor input, SymInt[]? output_size, float[]? scale_factors) -> Tensor + python_module: nn + autogen: upsample_nearest1d.vec_out + +- func: _upsample_nearest_exact1d.vec(Tensor input, SymInt[]? output_size, float[]? scale_factors) -> Tensor + python_module: nn + autogen: _upsample_nearest_exact1d.vec_out + +- func: upsample_nearest2d.vec(Tensor input, SymInt[]? output_size, float[]? scale_factors) -> Tensor + python_module: nn + autogen: upsample_nearest2d.vec_out + tags: core + +- func: _upsample_nearest_exact2d.vec(Tensor input, SymInt[]? output_size, float[]? scale_factors) -> Tensor + python_module: nn + autogen: _upsample_nearest_exact2d.vec_out + +- func: upsample_nearest3d.vec(Tensor input, SymInt[]? output_size, float[]? scale_factors) -> Tensor + python_module: nn + autogen: upsample_nearest3d.vec_out + +- func: _upsample_nearest_exact3d.vec(Tensor input, SymInt[]? output_size, float[]? scale_factors) -> Tensor + python_module: nn + autogen: _upsample_nearest_exact3d.vec_out + +# NOTE: all of the non-"vec" upsample overloads are only kept for backward compatibility. +- func: upsample_linear1d.out(Tensor self, SymInt[1] output_size, bool align_corners, float? scales=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_linear1d_out_cpu + CUDA: upsample_linear1d_out_cuda + MPS: upsample_linear1d_out_mps + +- func: upsample_linear1d(Tensor self, SymInt[1] output_size, bool align_corners, float? scales=None) -> Tensor + python_module: nn + structured_delegate: upsample_linear1d.out + +- func: upsample_linear1d_backward.grad_input(Tensor grad_output, SymInt[1] output_size, SymInt[3] input_size, bool align_corners, float? scales=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_linear1d_backward_out_cpu + CUDA: upsample_linear1d_backward_out_cuda + MPS: upsample_linear1d_backward_out_mps + +- func: upsample_linear1d_backward(Tensor grad_output, SymInt[1] output_size, SymInt[3] input_size, bool align_corners, float? scales=None) -> Tensor + python_module: nn + structured_delegate: upsample_linear1d_backward.grad_input + +- func: upsample_bilinear2d.out(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_bilinear2d_out_cpu + CUDA: upsample_bilinear2d_out_cuda + MPS: upsample_bilinear2d_out_mps + +- func: upsample_bilinear2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_bilinear2d.out + dispatch: + QuantizedCPU: upsample_bilinear2d_quantized_cpu + +- func: upsample_bilinear2d_backward.grad_input(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_bilinear2d_backward_out_cpu + CUDA: upsample_bilinear2d_backward_out_cuda + MPS: upsample_bilinear2d_backward_out_mps + +- func: upsample_bilinear2d_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_bilinear2d_backward.grad_input + +- func: _upsample_bilinear2d_aa.out(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_bilinear2d_aa_out_cpu + CUDA: _upsample_bilinear2d_aa_out_cuda + MPS: _upsample_bilinear2d_aa_out_mps + +- func: _upsample_bilinear2d_aa(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_bilinear2d_aa.out + +- func: _upsample_bilinear2d_aa_backward.grad_input(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_bilinear2d_aa_backward_out_cpu + CUDA: _upsample_bilinear2d_aa_backward_out_cuda + +- func: _upsample_bilinear2d_aa_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_bilinear2d_aa_backward.grad_input + +- func: upsample_bicubic2d.out(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_bicubic2d_out_cpu + CUDA: upsample_bicubic2d_out_cuda + MPS: upsample_bicubic2d_out_mps + +- func: upsample_bicubic2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_bicubic2d.out + +- func: upsample_bicubic2d_backward.grad_input(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_bicubic2d_backward_out_cpu + CUDA: upsample_bicubic2d_backward_out_cuda + MPS: upsample_bicubic2d_backward_out_mps + +- func: upsample_bicubic2d_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_bicubic2d_backward.grad_input + +- func: _upsample_bicubic2d_aa.out(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_bicubic2d_aa_out_cpu + CUDA: _upsample_bicubic2d_aa_out_cuda + MPS: _upsample_bicubic2d_aa_out_mps + +- func: _upsample_bicubic2d_aa(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_bicubic2d_aa.out + +- func: _upsample_bicubic2d_aa_backward.grad_input(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_bicubic2d_aa_backward_out_cpu + CUDA: _upsample_bicubic2d_aa_backward_out_cuda + +- func: _upsample_bicubic2d_aa_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_bicubic2d_aa_backward.grad_input + +- func: _upsample_lanczos2d_aa.out(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_lanczos2d_aa_out_cpu + +- func: _upsample_lanczos2d_aa(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_lanczos2d_aa.out + +- func: _upsample_lanczos2d_aa_backward.grad_input(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_lanczos2d_aa_backward_out_cpu + +- func: _upsample_lanczos2d_aa_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_lanczos2d_aa_backward.grad_input + +- func: upsample_trilinear3d.out(Tensor self, SymInt[3] output_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_trilinear3d_out_cpu + CUDA: upsample_trilinear3d_out_cuda + MPS: upsample_trilinear3d_out_mps + +- func: upsample_trilinear3d(Tensor self, SymInt[3] output_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_trilinear3d.out + +- func: upsample_trilinear3d_backward.grad_input(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_trilinear3d_backward_out_cpu + CUDA: upsample_trilinear3d_backward_out_cuda + MPS: upsample_trilinear3d_backward_out_mps + +- func: upsample_trilinear3d_backward(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_trilinear3d_backward.grad_input + +- func: upsample_nearest1d.out(Tensor self, SymInt[1] output_size, float? scales=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_nearest1d_out_cpu + CUDA: upsample_nearest1d_out_cuda + MPS: upsample_nearest1d_out_mps + +- func: _upsample_nearest_exact1d.out(Tensor self, SymInt[1] output_size, float? scales=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_nearest_exact1d_out_cpu + CUDA: _upsample_nearest_exact1d_out_cuda + MPS: _upsample_nearest_exact1d_out_mps + +- func: upsample_nearest1d(Tensor self, SymInt[1] output_size, float? scales=None) -> Tensor + python_module: nn + structured_delegate: upsample_nearest1d.out + +- func: _upsample_nearest_exact1d(Tensor self, SymInt[1] output_size, float? scales=None) -> Tensor + python_module: nn + structured_delegate: _upsample_nearest_exact1d.out + +- func: upsample_nearest1d_backward.grad_input(Tensor grad_output, SymInt[1] output_size, SymInt[3] input_size, float? scales=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_nearest1d_backward_out_cpu + CUDA: upsample_nearest1d_backward_out_cuda + MPS: upsample_nearest1d_backward_out_mps + +- func: _upsample_nearest_exact1d_backward.grad_input(Tensor grad_output, SymInt[1] output_size, SymInt[3] input_size, float? scales=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_nearest_exact1d_backward_out_cpu + CUDA: _upsample_nearest_exact1d_backward_out_cuda + MPS: _upsample_nearest_exact1d_backward_out_mps + +- func: upsample_nearest1d_backward(Tensor grad_output, SymInt[1] output_size, SymInt[3] input_size, float? scales=None) -> Tensor + python_module: nn + structured_delegate: upsample_nearest1d_backward.grad_input + +- func: _upsample_nearest_exact1d_backward(Tensor grad_output, SymInt[1] output_size, SymInt[3] input_size, float? scales=None) -> Tensor + python_module: nn + structured_delegate: _upsample_nearest_exact1d_backward.grad_input + +- func: upsample_nearest2d.out(Tensor self, SymInt[2] output_size, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_nearest2d_out_cpu + CUDA: upsample_nearest2d_out_cuda + MPS: upsample_nearest2d_out_mps + +- func: _upsample_nearest_exact2d.out(Tensor self, SymInt[2] output_size, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_nearest_exact2d_out_cpu + CUDA: _upsample_nearest_exact2d_out_cuda + MPS: _upsample_nearest_exact2d_out_mps + +- func: upsample_nearest2d(Tensor self, SymInt[2] output_size, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_nearest2d.out + dispatch: + QuantizedCPU: upsample_nearest2d_quantized_cpu + +- func: _upsample_nearest_exact2d(Tensor self, SymInt[2] output_size, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_nearest_exact2d.out + dispatch: + QuantizedCPU: _upsample_nearest_exact2d_quantized_cpu + +- func: upsample_nearest2d_backward.grad_input(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_nearest2d_backward_out_cpu + CUDA: upsample_nearest2d_backward_out_cuda + MPS: upsample_nearest2d_backward_out_mps + +- func: _upsample_nearest_exact2d_backward.grad_input(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_nearest_exact2d_backward_out_cpu + CUDA: _upsample_nearest_exact2d_backward_out_cuda + MPS: _upsample_nearest_exact2d_backward_out_mps + +- func: upsample_nearest2d_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_nearest2d_backward.grad_input + +- func: _upsample_nearest_exact2d_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_nearest_exact2d_backward.grad_input + +- func: upsample_nearest3d.out(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_nearest3d_out_cpu + CUDA: upsample_nearest3d_out_cuda + MPS: upsample_nearest3d_out_mps + +- func: _upsample_nearest_exact3d.out(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_nearest_exact3d_out_cpu + CUDA: _upsample_nearest_exact3d_out_cuda + MPS: _upsample_nearest_exact3d_out_mps + +- func: upsample_nearest3d(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_nearest3d.out + dispatch: + QuantizedCPU: upsample_nearest3d_quantized_cpu + +- func: _upsample_nearest_exact3d(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_nearest_exact3d.out + dispatch: + QuantizedCPU: _upsample_nearest_exact3d_quantized_cpu + +- func: upsample_nearest3d_backward.grad_input(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, float? scales_d=None, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: upsample_nearest3d_backward_out_cpu + CUDA: upsample_nearest3d_backward_out_cuda + MPS: upsample_nearest3d_backward_out_mps + +- func: _upsample_nearest_exact3d_backward.grad_input(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, float? scales_d=None, float? scales_h=None, float? scales_w=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: _upsample_nearest_exact3d_backward_out_cpu + CUDA: _upsample_nearest_exact3d_backward_out_cuda + MPS: _upsample_nearest_exact3d_backward_out_mps + +- func: upsample_nearest3d_backward(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: upsample_nearest3d_backward.grad_input + +- func: _upsample_nearest_exact3d_backward(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + python_module: nn + structured_delegate: _upsample_nearest_exact3d_backward.grad_input + +- func: sigmoid_backward.grad_input(Tensor grad_output, Tensor output, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: sigmoid_backward_out + MPS: sigmoid_backward_out_mps + tags: pointwise + +- func: sigmoid_backward(Tensor grad_output, Tensor output) -> Tensor + python_module: nn + structured_delegate: sigmoid_backward.grad_input + tags: pointwise + +- func: logit_backward.grad_input(Tensor grad_output, Tensor self, float? eps=None, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA: logit_backward_out + MPS: logit_backward_out_mps + tags: pointwise + +- func: logit_backward(Tensor grad_output, Tensor self, float? eps=None) -> Tensor + python_module: nn + structured_delegate: logit_backward.grad_input + tags: pointwise + +- func: tanh_backward.grad_input(Tensor grad_output, Tensor output, *, Tensor(a!) grad_input) -> Tensor(a!) + python_module: nn + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MTIA: tanh_backward_out + MPS: tanh_backward_out_mps + tags: pointwise + +- func: tanh_backward(Tensor grad_output, Tensor output) -> Tensor + python_module: nn + structured_delegate: tanh_backward.grad_input + +# What's a thnn_conv_ versus a slow_conv_? +# +# Historically, we have inefficient implementations of convolutions +# coming from the THNN/THCUNN library. These convolutions typically +# operated by computing the Toeplitz matrix and then doing a matrix +# multiply with the input; this is very memory inefficient! However, +# occasionally, we really don't have anything better, so it's helpful +# to have these fallbacks when there is no more optimized implementation +# in cudnn or mkldnn, etc. Both thnn_ and slow_ convolutions fall +# into this bucket. +# +# The difference between these two designations, is that thnn_ refers +# to a convolution that is still written in the "legacy" style; that is, +# C code in the THNN/ or THCUNN/ directory. A slow_ convolution is +# one that is written in the native style: modern C++. Algorithmically, +# these are the same thing, but we give them different prefixes to +# make the operational distinction clear. + tags: pointwise + +- func: slow_conv_transpose2d.out(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0, SymInt[2] output_padding=0, SymInt[2] dilation=1, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + structured: True + dispatch: + CPU: slow_conv_transpose2d_structured_cpu + CUDA: slow_conv_transpose2d_structured_cuda + +- func: slow_conv_transpose2d(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0, SymInt[2] output_padding=0, SymInt[2] dilation=1) -> Tensor + python_module: nn + structured_delegate: slow_conv_transpose2d.out + +- func: slow_conv_transpose3d.out(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0, SymInt[3] output_padding=0, SymInt[3] dilation=1, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU: slow_conv_transpose3d_out_cpu + CUDA: slow_conv_transpose3d_out_cuda + +- func: slow_conv_transpose3d(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0, SymInt[3] output_padding=0, SymInt[3] dilation=1) -> Tensor + python_module: nn + dispatch: + CPU: slow_conv_transpose3d_cpu + CUDA: slow_conv_transpose3d_cuda + +- func: thnn_conv2d.out(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + +- func: thnn_conv2d(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0) -> Tensor + python_module: nn + +- func: _slow_conv2d_forward.output(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias, SymInt[2] stride, SymInt[2] padding, *, Tensor(a!) output) -> Tensor(a!) + python_module: nn + dispatch: + CPU: slow_conv2d_forward_out_cpu + CUDA: slow_conv2d_forward_out_cuda + +- func: _slow_conv2d_forward(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias, SymInt[2] stride, SymInt[2] padding) -> Tensor + python_module: nn + dispatch: + CPU: slow_conv2d_forward_cpu + CUDA: slow_conv2d_forward_cuda + +- func: _slow_conv2d_backward.grad_input(Tensor grad_output, Tensor self, Tensor weight, SymInt[2] kernel_size, SymInt[2] stride, SymInt[2] padding, *, Tensor(a!) grad_input, Tensor(b!) grad_weight, Tensor(c!) grad_bias) -> (Tensor(a!), Tensor(b!), Tensor(c!)) + python_module: nn + dispatch: + CPU: slow_conv2d_backward_out_cpu + CUDA: slow_conv2d_backward_out_cuda + +- func: _slow_conv2d_backward.output_mask(Tensor grad_output, Tensor self, Tensor weight, SymInt[2] kernel_size, SymInt[2] stride, SymInt[2] padding, bool[3] output_mask) -> (Tensor grad_input, Tensor grad_weight, Tensor grad_bias) + python_module: nn + dispatch: + CPU: slow_conv2d_backward_cpu + CUDA: slow_conv2d_backward_cuda + autogen: _slow_conv2d_backward.output_mask_out + +- func: _conv_depthwise2d.out(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias, SymInt[2] stride, SymInt[2] padding, SymInt[2] dilation, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CUDA: conv_depthwise2d_cuda_out + +- func: _conv_depthwise2d(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias, SymInt[2] stride, SymInt[2] padding, SymInt[2] dilation) -> Tensor + python_module: nn + dispatch: + CUDA: conv_depthwise2d_cuda + +- func: conv_depthwise3d(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias, SymInt[3] stride, SymInt[3] padding, SymInt[3] dilation) -> Tensor + python_module: nn + dispatch: + CUDA: conv_depthwise3d_cuda + autogen: conv_depthwise3d.out + +- func: slow_conv3d.out(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + +- func: slow_conv3d(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0) -> Tensor + python_module: nn + +- func: slow_conv3d_forward.output(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias, SymInt[3] stride, SymInt[3] padding, *, Tensor(a!) output) -> Tensor(a!) + python_module: nn + dispatch: + CPU: slow_conv3d_forward_out_cpu + +- func: slow_conv3d_forward(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias, SymInt[3] stride, SymInt[3] padding) -> Tensor + python_module: nn + dispatch: + CPU: slow_conv3d_forward_cpu + +- func: slow_conv_dilated2d(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0, SymInt[2] dilation=1) -> Tensor + python_module: nn + dispatch: + CPU: slow_conv_dilated2d_cpu + CUDA: slow_conv_dilated2d_cuda + autogen: slow_conv_dilated2d.out + +- func: slow_conv_dilated3d(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0, SymInt[3] dilation=1) -> Tensor + python_module: nn + dispatch: + CPU: slow_conv_dilated3d_cpu + CUDA: slow_conv_dilated3d_cuda + autogen: slow_conv_dilated3d.out + +- func: col2im.out(Tensor self, SymInt[2] output_size, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU: col2im_out_cpu + CUDA: col2im_out_cuda + MPS: col2im_out_mps + +- func: col2im(Tensor self, SymInt[2] output_size, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride) -> Tensor + python_module: nn + dispatch: + CPU: col2im_cpu + CUDA: col2im_cuda + MPS: col2im_mps + tags: core + +- func: column_stack(Tensor[] tensors) -> Tensor + +- func: column_stack.out(Tensor[] tensors, *, Tensor(a!) out) -> Tensor(a!) + +- func: im2col.out(Tensor self, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride, *, Tensor(a!) out) -> Tensor(a!) + python_module: nn + dispatch: + CPU: im2col_out_cpu + CUDA: im2col_out_cuda + MPS: im2col_out_mps + +- func: im2col(Tensor self, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride) -> Tensor + python_module: nn + dispatch: + CPU: im2col_cpu + CUDA: im2col_cuda + MPS: im2col_mps + +- func: isfinite(Tensor self) -> Tensor + variants: function, method + device_check: NoCheck + device_guard: False + tags: pointwise + +- func: isinf(Tensor self) -> Tensor + variants: function, method + device_check: NoCheck + device_guard: False + dispatch: + CompositeExplicitAutograd: isinf + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_isinf + SparseCPU, SparseCUDA, SparseMPS: isinf_sparse + SparseMeta: isinf_sparse_meta + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: isinf_sparse_csr + autogen: isinf.out + tags: [core, pointwise] + +- func: record_stream(Tensor(a!) self, Stream s) -> () + variants: method + dispatch: + CUDA: record_stream_cuda + +- func: isposinf(Tensor self) -> Tensor + variants: function, method + structured_delegate: isposinf.out + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_isposinf + SparseCPU, SparseCUDA, SparseMPS: isposinf_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: isposinf_sparse_csr + tags: pointwise + +- func: isposinf.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: isposinf_out + SparseCPU, SparseCUDA, SparseMPS: isposinf_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: isposinf_sparse_csr_out + tags: pointwise + +- func: isneginf(Tensor self) -> Tensor + variants: function, method + structured_delegate: isneginf.out + dispatch: + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: NestedTensor_isneginf + SparseCPU, SparseCUDA, SparseMPS: isneginf_sparse + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: isneginf_sparse_csr + tags: pointwise + +- func: isneginf.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: isneginf_out + SparseCPU, SparseCUDA, SparseMPS: isneginf_sparse_out + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: isneginf_sparse_csr_out + tags: pointwise + +# NOTE [_add_batch_dim and _remove_batch_dim] +# _add_batch_dim and _remove_batch_dim are meant to be used in the implementation +# of the vmap frontend API (see torch/_vmap_internals.py). They are not +# user-facing, hence the leading underscore. Please don't use them them anywhere else. +- func: _add_batch_dim(Tensor self, int batch_dim, int level) -> Tensor + variants: function + +# See NOTE [_add_batch_dim and _remove_batch_dim] +- func: _remove_batch_dim(Tensor self, int level, SymInt batch_size, int out_dim) -> Tensor + variants: function + +## Functions related to the `torch.special` namespace +# Note [special namespace binding] +# Functions in the special python module should have their names start with +# "special_" underscore and be bound to the desired Python name in +# torch/special/__init__.py, and the desired C++ name in torch/csrc/api/include/torch/special.h. +# The "special_" names should be hidden from the user and not documented. + +- func: special_entr(Tensor self) -> Tensor + structured_delegate: special_entr.out + python_module: special + variants: function + tags: pointwise + +- func: special_entr.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: special + variants: function + dispatch: + CPU, CUDA, MPS: special_entr_out + tags: pointwise + +- func: special_ndtri(Tensor self) -> Tensor + structured_delegate: special_ndtri.out + python_module: special + variants: function + tags: pointwise + +- func: special_ndtri.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: special + variants: function + dispatch: + CPU, CUDA: special_ndtri_out + tags: pointwise + +- func: special_log_ndtr(Tensor self) -> Tensor + structured_delegate: special_log_ndtr.out + python_module: special + variants: function + tags: pointwise + +- func: special_log_ndtr.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + structured: True + structured_inherits: TensorIteratorBase + python_module: special + variants: function + dispatch: + CPU, CUDA: special_log_ndtr_out + tags: pointwise + +- func: special_expm1(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_expm1.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_exp2(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_exp2.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_psi(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_psi.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_digamma(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_digamma.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_gammaln(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_gammaln.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_erf(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_erf.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_erfc(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_erfc.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + +- func: special_erfcx(Tensor self) -> Tensor + python_module: special + variants: function + structured_delegate: special_erfcx.out + tags: pointwise + +- func: special_erfcx.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: special_erfcx_out + tags: pointwise + +- func: special_erfinv(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_erfinv.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + +- func: special_ndtr(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_ndtr.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_xlog1py(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + python_module: special + variants: function + structured_delegate: special_xlog1py.out + tags: pointwise + +- func: special_xlog1py.self_scalar(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + python_module: special + variants: function + dispatch: + CompositeExplicitAutograd: special_xlog1py + tags: pointwise + +- func: special_xlog1py.other_scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + python_module: special + variants: function + dispatch: + CompositeExplicitAutograd: special_xlog1py + tags: pointwise + +- func: special_xlog1py.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + python_module: special + variants: function + dispatch: + CPU, CUDA, MPS: special_xlog1py_out + tags: pointwise + +- func: special_xlog1py.self_scalar_out(Scalar self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: special + variants: function + dispatch: + CompositeExplicitAutograd: special_xlog1py_out + tags: pointwise + +- func: special_xlog1py.other_scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: special + variants: function + dispatch: + CompositeExplicitAutograd: special_xlog1py_out + tags: pointwise + +- func: special_xlogy(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + python_module: special + variants: function + +- func: special_xlogy.self_scalar(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + python_module: special + variants: function + +- func: special_xlogy.other_scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + python_module: special + variants: function + +- func: special_xlogy.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: special + variants: function + +- func: special_xlogy.self_scalar_out(Scalar self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: special + variants: function + +- func: special_xlogy.other_scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: special + variants: function + +- func: special_zeta(Tensor self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + python_module: special + variants: function + structured_delegate: special_zeta.out + tags: pointwise + +- func: special_zeta.self_scalar(Scalar self, Tensor other) -> Tensor + device_check: NoCheck # TensorIterator + python_module: special + variants: function + dispatch: + CompositeExplicitAutograd: special_zeta + tags: pointwise + +- func: special_zeta.other_scalar(Tensor self, Scalar other) -> Tensor + device_check: NoCheck # TensorIterator + python_module: special + variants: function + dispatch: + CompositeExplicitAutograd: special_zeta + tags: pointwise + +- func: special_zeta.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + structured: True + structured_inherits: TensorIteratorBase + python_module: special + variants: function + dispatch: + CPU, CUDA, MPS: special_zeta_out + tags: pointwise + +- func: special_zeta.self_scalar_out(Scalar self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: special + variants: function + dispatch: + CompositeExplicitAutograd: special_zeta_out + tags: pointwise + +- func: special_zeta.other_scalar_out(Tensor self, Scalar other, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck # TensorIterator + python_module: special + variants: function + dispatch: + CompositeExplicitAutograd: special_zeta_out + tags: pointwise + +- func: special_i0(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_i0.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_i0e(Tensor self) -> Tensor + python_module: special + variants: function + structured_delegate: special_i0e.out + tags: pointwise + +- func: special_i0e.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: special_i0e_out + tags: pointwise + +- func: special_i1(Tensor self) -> Tensor + python_module: special + variants: function + structured_delegate: special_i1.out + tags: pointwise + +- func: special_i1.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: special_i1_out + tags: pointwise + +- func: special_i1e(Tensor self) -> Tensor + python_module: special + variants: function + structured_delegate: special_i1e.out + tags: pointwise + +- func: special_i1e.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + structured: True + structured_inherits: TensorIteratorBase + dispatch: + CPU, CUDA, MPS: special_i1e_out + tags: pointwise + +- func: special_logit(Tensor self, float? eps=None) -> Tensor + python_module: special + variants: function + +- func: special_logit.out(Tensor self, float? eps=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + +- func: special_polygamma(int n, Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_polygamma.out(int n, Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + +- func: special_logsumexp(Tensor self, int[1] dim, bool keepdim=False) -> Tensor + python_module: special + variants: function + +- func: special_logsumexp.out(Tensor self, int[1] dim, bool keepdim=False, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + +- func: special_expit(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_expit.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_sinc(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_sinc.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_round(Tensor self, *, int decimals=0) -> Tensor + python_module: special + variants: function + +- func: special_round.out(Tensor self, *, int decimals=0, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_log1p(Tensor self) -> Tensor + python_module: special + variants: function + +- func: special_log1p.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_log_softmax(Tensor self, int dim, *, ScalarType? dtype=None) -> Tensor + python_module: special + variants: function + +- func: special_gammainc.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_gammainc(Tensor self, Tensor other) -> Tensor + python_module: special + variants: function + +- func: special_gammaincc.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_gammaincc(Tensor self, Tensor other) -> Tensor + python_module: special + variants: function + +- func: special_multigammaln(Tensor self, int p) -> Tensor + python_module: special + variants: function + +- func: special_multigammaln.out(Tensor self, int p, *, Tensor(a!) out) -> Tensor(a!) + python_module: special + variants: function + +- func: special_softmax(Tensor self, int dim, ScalarType? dtype=None) -> Tensor + python_module: special + variants: function + +## Functions related to the fast Fourier transform and the torch.fft namespace +# Note [FFT namespace binding] +# Functions in the fft python module should have their names start with +# "fft_" underscore and be bound to the desired Python name in +# torch/fft/__init__.py, and the desired C++ name in torch/csrc/api/include/torch/fft.h. +# The "fft_" names should be hidden from the user and not documented. +# +# See fft_fft as an example. + +# torch.fft.fft +# NOTE: NOT an alias for torch.fft, which has different semantics +- func: fft_fft(Tensor self, SymInt? n=None, int dim=-1, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_fft_symint + +- func: fft_fft.out(Tensor self, SymInt? n=None, int dim=-1, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_fft_symint_out + +- func: fft_ifft(Tensor self, SymInt? n=None, int dim=-1, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ifft_symint + +- func: fft_ifft.out(Tensor self, SymInt? n=None, int dim=-1, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ifft_symint_out + +- func: fft_rfft(Tensor self, SymInt? n=None, int dim=-1, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_rfft_symint + +- func: fft_rfft.out(Tensor self, SymInt? n=None, int dim=-1, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_rfft_symint_out + +- func: fft_irfft(Tensor self, SymInt? n=None, int dim=-1, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_irfft_symint + +- func: fft_irfft.out(Tensor self, SymInt? n=None, int dim=-1, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_irfft_symint_out + +- func: fft_hfft(Tensor self, SymInt? n=None, int dim=-1, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_hfft_symint + +- func: fft_hfft.out(Tensor self, SymInt? n=None, int dim=-1, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_hfft_symint_out + +- func: fft_ihfft(Tensor self, SymInt? n=None, int dim=-1, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ihfft_symint + +- func: fft_ihfft.out(Tensor self, SymInt? n=None, int dim=-1, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ihfft_symint_out + +- func: fft_fft2(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_fft2_symint + +- func: fft_fft2.out(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_fft2_symint_out + +- func: fft_ifft2(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ifft2_symint + +- func: fft_ifft2.out(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ifft2_symint_out + +- func: fft_rfft2(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_rfft2_symint + +- func: fft_rfft2.out(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_rfft2_symint_out + +- func: fft_irfft2(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_irfft2_symint + +- func: fft_irfft2.out(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_irfft2_symint_out + +- func: fft_hfft2(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None) -> Tensor + use_const_ref_for_mutable_tensors: True + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_hfft2_symint + +- func: fft_hfft2.out(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_hfft2_symint_out + +- func: fft_ihfft2(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None) -> Tensor + use_const_ref_for_mutable_tensors: True + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ihfft2_symint + +- func: fft_ihfft2.out(Tensor self, SymInt[1]? s=None, int[1] dim=[-2,-1], str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ihfft2_symint_out + +- func: fft_fftn(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_fftn_symint + +- func: fft_fftn.out(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_fftn_symint_out + +- func: fft_ifftn(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ifftn_symint + +- func: fft_ifftn.out(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ifftn_symint_out + +- func: fft_rfftn(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_rfftn_symint + +- func: fft_rfftn.out(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_rfftn_symint_out + +- func: fft_irfftn(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_irfftn_symint + +- func: fft_irfftn.out(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_irfftn_symint_out + +- func: fft_hfftn(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor + use_const_ref_for_mutable_tensors: True + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_hfftn_symint + +- func: fft_hfftn.out(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_hfftn_symint_out + +- func: fft_ihfftn(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None) -> Tensor + use_const_ref_for_mutable_tensors: True + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ihfftn_symint + +- func: fft_ihfftn.out(Tensor self, SymInt[1]? s=None, int[1]? dim=None, str? norm=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeImplicitAutograd: fft_ihfftn_symint_out + +- func: fft_fftfreq(int n, float d=1.0, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeExplicitAutograd: fft_fftfreq + +- func: fft_fftfreq.out(int n, float d=1.0, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeExplicitAutograd: fft_fftfreq_out + +- func: fft_rfftfreq(int n, float d=1.0, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + python_module: fft + variants: function + dispatch: + CompositeExplicitAutograd: fft_rfftfreq + +- func: fft_rfftfreq.out(int n, float d=1.0, *, Tensor(a!) out) -> Tensor(a!) + python_module: fft + variants: function + dispatch: + CompositeExplicitAutograd: fft_rfftfreq_out + +- func: fft_fftshift(Tensor self, int[1]? dim=None) -> Tensor + python_module: fft + variants: function + +- func: fft_ifftshift(Tensor self, int[1]? dim=None) -> Tensor + python_module: fft + variants: function + +## Functions for linear algebra and the torch.linalg namespace +# Note [linalg namespace binding] +# Functions in the linalg python module should have their names start with +# "linalg_" and be bound to the desired Python name in +# torch/linalg/__init__.py, and the desired C++ name in torch/csrc/api/include/torch/linalg.h. +# The "linalg_" names should be hidden from the user and not documented. +# +# See linalg_det as an example. + +# "_ex" stands for experimental +- func: linalg_cholesky_ex(Tensor self, *, bool upper=False, bool check_errors=False) -> (Tensor L, Tensor info) + python_module: linalg + structured_delegate: linalg_cholesky_ex.L + +- func: linalg_cholesky_ex.L(Tensor self, *, bool upper=False, bool check_errors=False, Tensor(a!) L, Tensor(b!) info) -> (Tensor(a!) L, Tensor(b!) info) + python_module: linalg + structured: True + dispatch: + CPU, CUDA, MPS: linalg_cholesky_ex_out + +- func: linalg_cholesky(Tensor self, *, bool upper=False) -> Tensor + python_module: linalg + +- func: linalg_cholesky.out(Tensor self, *, bool upper=False, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +- func: linalg_cross(Tensor self, Tensor other, *, int dim=-1) -> Tensor + python_module: linalg + variants: function + structured_delegate: linalg_cross.out + dispatch: + ZeroTensor: linalg_cross_zerotensor + +- func: linalg_cross.out(Tensor self, Tensor other, *, int dim=-1, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + structured: True + dispatch: + CPU, CUDA, MPS: linalg_cross_out + +# linalg.lu_factor +- func: linalg_lu_factor(Tensor A, *, bool pivot=True) -> (Tensor LU, Tensor pivots) + python_module: linalg + variants: function + +- func: linalg_lu_factor.out(Tensor A, *, bool pivot=True, Tensor(a!) LU, Tensor(b!) pivots) -> (Tensor(a!) LU, Tensor(b!) pivots) + python_module: linalg + variants: function + +- func: linalg_lu_factor_ex(Tensor A, *, bool pivot=True, bool check_errors=False) -> (Tensor LU, Tensor pivots, Tensor info) + python_module: linalg + structured_delegate: linalg_lu_factor_ex.out + variants: function + +- func: linalg_lu_factor_ex.out(Tensor A, *, bool pivot=True, bool check_errors=False, Tensor(a!) LU, Tensor(b!) pivots, Tensor(c!) info) -> (Tensor(a!) LU, Tensor(b!) pivots, Tensor(c!) info) + python_module: linalg + variants: function + structured: True + dispatch: + CPU, CUDA: linalg_lu_factor_ex_out + MPS: linalg_lu_factor_ex_out_mps + +# linalg.lu +- func: linalg_lu(Tensor A, *, bool pivot=True) -> (Tensor P, Tensor L, Tensor U) + python_module: linalg + structured_delegate: linalg_lu.out + variants: function + +- func: linalg_lu.out(Tensor A, *, bool pivot=True, Tensor(a!) P, Tensor(b!) L, Tensor(c!) U) -> (Tensor(a!) P, Tensor(b!) L, Tensor(c!) U) + python_module: linalg + variants: function + structured: True + dispatch: + CPU, CUDA, MPS: linalg_lu_out + +# linalg.lu_solve +- func: linalg_lu_solve(Tensor LU, Tensor pivots, Tensor B, *, bool left=True, bool adjoint=False) -> Tensor + python_module: linalg + structured_delegate: linalg_lu_solve.out + variants: function + +- func: linalg_lu_solve.out(Tensor LU, Tensor pivots, Tensor B, *, bool left=True, bool adjoint=False, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + structured: True + dispatch: + CPU, CUDA: linalg_lu_solve_out + +# linalg.det +- func: _linalg_det(Tensor A) -> (Tensor result, Tensor LU, Tensor pivots) + structured_delegate: _linalg_det.result + +- func: _linalg_det.result(Tensor A, *, Tensor(a!) result, Tensor(b!) LU, Tensor(c!) pivots) -> (Tensor(a!) result, Tensor(b!) LU, Tensor(c!) pivots) + structured: True + dispatch: + CPU, CUDA, MPS: _linalg_det_out + +- func: linalg_det(Tensor A) -> Tensor + python_module: linalg + variants: function + +- func: linalg_det.out(Tensor A, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +# torch.det, alias for torch.linalg.det +- func: det(Tensor self) -> Tensor + variants: function, method + +- func: linalg_ldl_factor_ex(Tensor self, *, bool hermitian=False, bool check_errors=False) -> (Tensor LD, Tensor pivots, Tensor info) + structured_delegate: linalg_ldl_factor_ex.out + python_module: linalg + variants: function + +- func: linalg_ldl_factor_ex.out(Tensor self, *, bool hermitian=False, bool check_errors=False, Tensor(a!) LD, Tensor(b!) pivots, Tensor(c!) info) -> (Tensor(a!) LD, Tensor(b!) pivots, Tensor(c!) info) + structured: True + python_module: linalg + variants: function + dispatch: + CPU, CUDA: linalg_ldl_factor_ex_out + +- func: linalg_ldl_factor(Tensor self, *, bool hermitian=False) -> (Tensor LD, Tensor pivots) + python_module: linalg + variants: function + +- func: linalg_ldl_factor.out(Tensor self, *, bool hermitian=False, Tensor(a!) LD, Tensor(b!) pivots) -> (Tensor(a!) LD, Tensor(b!) pivots) + python_module: linalg + variants: function + +- func: linalg_ldl_solve(Tensor LD, Tensor pivots, Tensor B, *, bool hermitian=False) -> Tensor + structured_delegate: linalg_ldl_solve.out + python_module: linalg + variants: function + +- func: linalg_ldl_solve.out(Tensor LD, Tensor pivots, Tensor B, *, bool hermitian=False, Tensor(a!) out) -> Tensor(a!) + structured: True + python_module: linalg + variants: function + dispatch: + CPU, CUDA: linalg_ldl_solve_out + +- func: linalg_lstsq(Tensor self, Tensor b, float? rcond=None, *, str? driver=None) -> (Tensor solution, Tensor residuals, Tensor rank, Tensor singular_values) + python_module: linalg + variants: function + dispatch: + CompositeExplicitAutograd: linalg_lstsq + tags: dynamic_output_shape + +- func: linalg_lstsq.out(Tensor self, Tensor b, float? rcond=None, *, str? driver=None, Tensor(a!) solution, Tensor(b!) residuals, Tensor(c!) rank, Tensor(d!) singular_values) -> (Tensor(a!) solution, Tensor(b!) residuals, Tensor(c!) rank, Tensor(d!) singular_values) + python_module: linalg + variants: function + dispatch: + CPU, CUDA: linalg_lstsq_out + tags: dynamic_output_shape + +# torch.linalg.matmul, alias for torch.matmul +- func: linalg_matmul(Tensor self, Tensor other) -> Tensor + python_module: linalg + variants: function + +- func: linalg_matmul.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +- func: linalg_vecdot(Tensor x, Tensor y, *, int dim=-1) -> Tensor + python_module: linalg + variants: function + +- func: linalg_vecdot.out(Tensor x, Tensor y, *, int dim=-1, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +- func: linalg_matrix_exp(Tensor self) -> Tensor + python_module: linalg + variants: function + dispatch: + CPU, CUDA: linalg_matrix_exp + autogen: linalg_matrix_exp.out + +- func: _linalg_slogdet(Tensor A) -> (Tensor sign, Tensor logabsdet, Tensor LU, Tensor pivots) + structured_delegate: _linalg_slogdet.sign + +- func: _linalg_slogdet.sign(Tensor A, *, Tensor(a!) sign, Tensor(b!) logabsdet, Tensor(c!) LU, Tensor(d!) pivots) -> (Tensor(a!) sign, Tensor(b!) logabsdet, Tensor(c!) LU, Tensor(d!) pivots) + structured: True + dispatch: + CPU, CUDA, MPS: _linalg_slogdet_out + +- func: linalg_slogdet(Tensor A) -> (Tensor sign, Tensor logabsdet) + python_module: linalg + +- func: linalg_slogdet.out(Tensor A, *, Tensor(a!) sign, Tensor(b!) logabsdet) -> (Tensor(a!) sign, Tensor(b!) logabsdet) + python_module: linalg + +- func: slogdet(Tensor self) -> (Tensor sign, Tensor logabsdet) + variants: function, method + +- func: slogdet.out(Tensor self, *, Tensor(a!) sign, Tensor(b!) logabsdet) -> (Tensor(a!) sign, Tensor(b!) logabsdet) + variants: function + +- func: logdet(Tensor self) -> Tensor + variants: function, method + +- func: linalg_eig(Tensor self) -> (Tensor eigenvalues, Tensor eigenvectors) + python_module: linalg + variants: function + dispatch: + CPU, CUDA: linalg_eig + +- func: linalg_eig.out(Tensor self, *, Tensor(a!) eigenvalues, Tensor(b!) eigenvectors) -> (Tensor(a!) eigenvalues, Tensor(b!) eigenvectors) + python_module: linalg + dispatch: + CPU, CUDA: linalg_eig_out + +- func: _linalg_eigvals(Tensor self) -> Tensor + python_module: linalg + dispatch: + CPU, CUDA: _linalg_eigvals + +- func: linalg_eigvals(Tensor self) -> Tensor + python_module: linalg + +- func: linalg_eigvals.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + dispatch: + CPU, CUDA: linalg_eigvals_out + +# This function is exposes the `compute_v` flag, which is then used to implement `linalg.eigh` and +# `linalg.eigvalsh` as composite functions that call this one +- func: _linalg_eigh(Tensor A, str UPLO="L", bool compute_v=True) -> (Tensor eigenvalues, Tensor eigenvectors) + structured_delegate: _linalg_eigh.eigenvalues + +- func: _linalg_eigh.eigenvalues(Tensor A, str UPLO="L", bool compute_v=True, *, Tensor(a!) eigenvalues, Tensor(b!) eigenvectors) -> (Tensor(a!) eigenvalues, Tensor(b!) eigenvectors) + structured: True + dispatch: + CPU, CUDA: _linalg_eigh_out + +- func: linalg_eigh(Tensor self, str UPLO="L") -> (Tensor eigenvalues, Tensor eigenvectors) + python_module: linalg + +- func: linalg_eigh.eigvals(Tensor self, str UPLO="L", *, Tensor(a!) eigvals, Tensor(b!) eigvecs) -> (Tensor(a!) eigenvalues, Tensor(b!) eigenvectors) + python_module: linalg + +- func: linalg_eigvalsh(Tensor self, str UPLO="L") -> Tensor + python_module: linalg + +- func: linalg_eigvalsh.out(Tensor self, str UPLO="L", *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +- func: linalg_householder_product(Tensor input, Tensor tau) -> Tensor + python_module: linalg + variants: function + dispatch: + CPU, CUDA, MPS: linalg_householder_product + +- func: linalg_householder_product.out(Tensor input, Tensor tau, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + dispatch: + CPU, CUDA, MPS: linalg_householder_product_out + +- func: linalg_inv_ex(Tensor A, *, bool check_errors=False) -> (Tensor inverse, Tensor info) + python_module: linalg + structured_delegate: linalg_inv_ex.inverse + +- func: linalg_inv_ex.inverse(Tensor A, *, bool check_errors=False, Tensor(a!) inverse, Tensor(b!) info) -> (Tensor(a!) inverse, Tensor(b!) info) + python_module: linalg + structured: True + dispatch: + CPU, CUDA: linalg_inv_ex_out + MPS: linalg_inv_ex_out_mps + +- func: linalg_inv(Tensor A) -> Tensor + python_module: linalg + +- func: linalg_inv.out(Tensor A, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +- func: inverse(Tensor self) -> Tensor + variants: function, method + +- func: inverse.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + +- func: inner(Tensor self, Tensor other) -> Tensor + variants: function, method + +- func: inner.out(Tensor self, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + +- func: outer(Tensor self, Tensor vec2) -> Tensor + variants: function, method + +- func: outer.out(Tensor self, Tensor vec2, *, Tensor(a!) out) -> Tensor(a!) + +# torch.ger, alias for torch.outer +- func: ger(Tensor self, Tensor vec2) -> Tensor + variants: function, method + +- func: ger.out(Tensor self, Tensor vec2, *, Tensor(a!) out) -> Tensor(a!) + +- func: linalg_norm(Tensor self, Scalar? ord=None, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + python_module: linalg + variants: function + +- func: linalg_norm.ord_str(Tensor self, str ord, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + python_module: linalg + variants: function + +- func: linalg_norm.out(Tensor self, Scalar? ord=None, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_norm.ord_str_out(Tensor self, str ord, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_vector_norm(Tensor self, Scalar ord=2, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + python_module: linalg + variants: function + structured_delegate: linalg_vector_norm.out + tags: reduction + +- func: linalg_vector_norm.out(Tensor self, Scalar ord=2, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + structured: True + dispatch: + CPU, CUDA, MPS: linalg_vector_norm_out + tags: reduction + +# Computes sum(|x|^ord) - the "power sum" without the final root. +# This is useful for distributed computing where partial power sums +# can be reduced across shards before taking the final root. +- func: linalg__powsum(Tensor self, Scalar ord=2, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + python_module: linalg + variants: function + dispatch: + CompositeExplicitAutograd: linalg__powsum_slow + CPU, CUDA: linalg__powsum + tags: reduction + +- func: linalg_matrix_norm(Tensor self, Scalar ord, int[] dim=[-2,-1], bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + python_module: linalg + +- func: linalg_matrix_norm.out(Tensor self, Scalar ord, int[] dim=[-2,-1], bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +- func: linalg_matrix_norm.str_ord(Tensor self, str ord='fro', int[] dim=[-2,-1], bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + python_module: linalg + +- func: linalg_matrix_norm.str_ord_out(Tensor self, str ord='fro', int[] dim=[-2,-1], bool keepdim=False, *, ScalarType? dtype=None, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +# This function is exposes the `compute_uv` flag, which is then used to implement `linalg.svd` and +# `linalg.svdvals` as composite functions that call this one +- func: _linalg_svd(Tensor A, bool full_matrices=False, bool compute_uv=True, *, str? driver=None) -> (Tensor U, Tensor S, Tensor Vh) + variants: function + structured_delegate: _linalg_svd.U + +- func: _linalg_svd.U(Tensor A, bool full_matrices=False, bool compute_uv=True, *, str? driver=None, Tensor(a!) U, Tensor(b!) S, Tensor(c!) Vh) -> (Tensor(a!) U, Tensor(b!) S, Tensor(c!) Vh) + structured: True + dispatch: + CPU, CUDA: _linalg_svd_out + +- func: linalg_svd(Tensor A, bool full_matrices=True, *, str? driver=None) -> (Tensor U, Tensor S, Tensor Vh) + python_module: linalg + variants: function + +- func: linalg_svd.U(Tensor A, bool full_matrices=True, *, str? driver=None, Tensor(a!) U, Tensor(b!) S, Tensor(c!) Vh) -> (Tensor(a!) U, Tensor(b!) S, Tensor(c!) Vh) + python_module: linalg + variants: function + +- func: linalg_svdvals(Tensor A, *, str? driver=None) -> Tensor + python_module: linalg + variants: function + +- func: linalg_svdvals.out(Tensor A, *, str? driver=None, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_cond(Tensor self, Scalar? p=None) -> Tensor + python_module: linalg + variants: function + +- func: linalg_cond.out(Tensor self, Scalar? p=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_cond.p_str(Tensor self, str p) -> Tensor + python_module: linalg + variants: function + +- func: linalg_cond.p_str_out(Tensor self, str p, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_pinv.atol_rtol_tensor(Tensor self, *, Tensor? atol=None, Tensor? rtol=None, bool hermitian=False) -> Tensor + python_module: linalg + variants: function + dispatch: + # calls svd, which calls mH() (view op) + # also calls narrow() + CompositeExplicitAutogradNonFunctional: linalg_pinv + +- func: linalg_pinv.atol_rtol_tensor_out(Tensor self, *, Tensor? atol=None, Tensor? rtol=None, bool hermitian=False, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + dispatch: + CompositeExplicitAutograd: linalg_pinv_out + +- func: linalg_pinv.atol_rtol_float(Tensor self, *, float? atol=None, float? rtol=None, bool hermitian=False) -> Tensor + cpp_no_default_args: ['atol', 'rtol'] + python_module: linalg + variants: function + +- func: linalg_pinv.atol_rtol_float_out(Tensor self, *, float? atol=None, float? rtol=None, bool hermitian=False, Tensor(a!) out) -> Tensor(a!) + cpp_no_default_args: ['atol', 'rtol'] + python_module: linalg + variants: function + +- func: linalg_pinv(Tensor self, float rcond, bool hermitian=False) -> Tensor + python_module: linalg + variants: function + +- func: linalg_pinv.rcond_tensor(Tensor self, Tensor rcond, bool hermitian=False) -> Tensor + python_module: linalg + variants: function + +- func: linalg_pinv.out(Tensor self, float rcond, bool hermitian=False, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_pinv.out_rcond_tensor(Tensor self, Tensor rcond, bool hermitian=False, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: _linalg_solve_ex(Tensor A, Tensor B, *, bool left=True, bool check_errors=False) -> (Tensor result, Tensor LU, Tensor pivots, Tensor info) + structured_delegate: _linalg_solve_ex.result + +- func: _linalg_solve_ex.result(Tensor A, Tensor B, *, bool left=True, bool check_errors=False, Tensor(a!) result, Tensor(b!) LU, Tensor(c!) pivots, Tensor(d!) info) -> (Tensor(a!) result, Tensor(b!) LU, Tensor(c!) pivots, Tensor(d!) info) + structured: True + dispatch: + CPU, CUDA: _linalg_solve_ex_out + MPS: _linalg_solve_ex_out_mps + +- func: linalg_solve_ex(Tensor A, Tensor B, *, bool left=True, bool check_errors=False) -> (Tensor result, Tensor info) + python_module: linalg + +- func: linalg_solve_ex.out(Tensor A, Tensor B, *, bool left=True, bool check_errors=False, Tensor(a!) result, Tensor(b!) info) -> (Tensor(a!) result, Tensor(b!) info) + python_module: linalg + +- func: linalg_solve(Tensor A, Tensor B, *, bool left=True) -> Tensor + python_module: linalg + +- func: _spsolve(Tensor A, Tensor B, *, bool left=True) -> Tensor + python_module: sparse + dispatch: + SparseCsrCUDA: _sparse_csr_linear_solve + +- func: linalg_solve.out(Tensor A, Tensor B, *, bool left=True, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +- func: linalg_tensorinv(Tensor self, int ind=2) -> Tensor + python_module: linalg + variants: function + +- func: linalg_tensorinv.out(Tensor self, int ind=2, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_tensorsolve(Tensor self, Tensor other, int[]? dims=None) -> Tensor + python_module: linalg + variants: function + +- func: linalg_tensorsolve.out(Tensor self, Tensor other, int[]? dims=None, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_qr(Tensor A, str mode='reduced') -> (Tensor Q, Tensor R) + python_module: linalg + variants: function + structured_delegate: linalg_qr.out + +- func: linalg_qr.out(Tensor A, str mode='reduced', *, Tensor(a!) Q, Tensor(b!) R) -> (Tensor(a!) Q, Tensor(b!) R) + python_module: linalg + structured: True + dispatch: + CPU, CUDA: linalg_qr_out + MPS: linalg_qr_out_mps + +- func: linalg_matrix_power(Tensor self, int n) -> Tensor + python_module: linalg + +- func: linalg_matrix_power.out(Tensor self, int n, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +- func: linalg_matrix_rank.atol_rtol_tensor(Tensor input, *, Tensor? atol=None, Tensor? rtol=None, bool hermitian=False) -> Tensor + python_module: linalg + variants: function + +- func: linalg_matrix_rank.atol_rtol_tensor_out(Tensor input, *, Tensor? atol=None, Tensor? rtol=None, bool hermitian=False, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_matrix_rank.atol_rtol_float(Tensor self, *, float? atol=None, float? rtol=None, bool hermitian=False) -> Tensor + cpp_no_default_args: ['atol', 'rtol'] + python_module: linalg + variants: function + +- func: linalg_matrix_rank.atol_rtol_float_out(Tensor self, *, float? atol=None, float? rtol=None, bool hermitian=False, Tensor(a!) out) -> Tensor(a!) + cpp_no_default_args: ['atol', 'rtol'] + python_module: linalg + variants: function + +- func: linalg_matrix_rank(Tensor self, float tol, bool hermitian=False) -> Tensor + python_module: linalg + variants: function + +- func: linalg_matrix_rank.out(Tensor self, float tol, bool hermitian=False, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_matrix_rank.tol_tensor(Tensor input, Tensor tol, bool hermitian=False) -> Tensor + python_module: linalg + variants: function + +- func: linalg_matrix_rank.out_tol_tensor(Tensor input, Tensor tol, bool hermitian=False, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + variants: function + +- func: linalg_multi_dot(Tensor[] tensors) -> Tensor + python_module: linalg + +- func: linalg_multi_dot.out(Tensor[] tensors, *, Tensor(a!) out) -> Tensor(a!) + python_module: linalg + +## Functions related to the `torch.nested` namespace +# Note [nested namespace binding] +# Functions in the nested python module should have their names start with +# "nested_" underscore and be bound to the desired Python name in +# torch/nested/__init__.py, and the desired C++ name in torch/csrc/api/include/torch/nested.h. +# The "nested_" names should be hidden from the user and not documented. + +- func: nested_to_padded_tensor(Tensor self, float padding, int[]? output_size=None) -> Tensor + python_module: nested + variants: function + +## Functions that are only for testing +# It is undocumented and should not be used outside of tests. +- func: _test_serialization_subcmul(Tensor self, Tensor other, Scalar alpha=1) -> Tensor + +# Note: for testing COW materialization within `at::parallel_for` loop function +- func: _test_parallel_materialize(Tensor self, int num_parallel, bool skip_first=False) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: _test_parallel_materialize + +# Note: this function is only for testing. +- func: _test_optional_intlist(Tensor values, int[]? addends) -> Tensor + python_module: nn + dispatch: + CPU: _test_optional_intlist + autogen: _test_optional_intlist.out + +# Note: this function is only for testing. +- func: _test_optional_filled_intlist(Tensor values, int[2]? addends) -> Tensor + python_module: nn + dispatch: + CPU: _test_optional_intlist + autogen: _test_optional_filled_intlist.out + +# Note: this function is only for testing. +- func: _test_optional_floatlist(Tensor values, float[]? addends) -> Tensor + python_module: nn + dispatch: + CPU: _test_optional_floatlist + autogen: _test_optional_floatlist.out + +# Note: this function is only for testing. +- func: _test_string_default(Tensor dummy, str a="\"'\\", str b='"\'\\') -> Tensor + python_module: nn + +# Note: this function is only for testing. +- func: _test_ambiguous_defaults.a(Tensor dummy, int a=1, int b=1) -> Tensor + python_module: nn + +# Note: this function is only for testing. +- func: _test_ambiguous_defaults.b(Tensor dummy, int a=2, str b="2") -> Tensor + cpp_no_default_args: ['a', 'b'] + python_module: nn + +# Note: this function is only for testing. +- func: _test_warn_in_autograd(Tensor self) -> Tensor + python_module: nn + dispatch: + CompositeExplicitAutograd: _test_warn_in_autograd + autogen: _test_warn_in_autograd.out + +# Note: this function is only for testing. +- func: _test_autograd_multiple_dispatch.fullcoverage(Tensor self) -> Tensor + dispatch: + # the NestedTensor keys are necessary because NestedTensor has been removed + # from the CompositeExplicitAutograd keyset see Note [NestedTensor Not Included in Backend Keys] + CompositeExplicitAutograd, NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: _test_autograd_multiple_dispatch_fullcoverage + autogen: _test_autograd_multiple_dispatch.fullcoverage_out + +# Note: this function is only for testing. +- func: _test_autograd_multiple_dispatch.ntonly(Tensor self, bool b) -> Tensor + dispatch: + CompositeImplicitAutograd, NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: _test_autograd_multiple_dispatch_ntonly + +# Note: this function is only for testing. +- func: _test_autograd_multiple_dispatch_view(Tensor(a) self) -> Tensor(a) + dispatch: + CompositeExplicitAutograd: _test_autograd_multiple_dispatch_view + +# Note: this function is only for testing. +- func: _test_autograd_multiple_dispatch_view_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: _test_autograd_multiple_dispatch_view_copy + tags: view_copy + autogen: _test_autograd_multiple_dispatch_view_copy.out + +- func: segment_reduce(Tensor data, str reduce, *, Tensor? lengths=None, Tensor? indices=None, Tensor? offsets=None, int axis=0, bool unsafe=False, Scalar? initial=None) -> Tensor + variants: function + dispatch: + CPU, CUDA: segment_reduce_kernel + autogen: segment_reduce.out + +- func: _segment_reduce_backward(Tensor grad, Tensor output, Tensor data, str reduce, *, Tensor? lengths=None, Tensor? offsets=None, int axis=0, Scalar? initial=None) -> Tensor + variants: function + dispatch: + CPU, CUDA: _segment_reduce_backward_kernel + autogen: _segment_reduce_backward.out + +- func: pad_sequence(Tensor[] sequences, bool batch_first=False, float padding_value=0.0, str padding_side="right") -> Tensor + python_module: nn + variants: function + +- func: flatten_dense_tensors(Tensor[] tensors) -> Tensor + variants: function + python_module: nn + +- func: unflatten_dense_tensors(Tensor flat, Tensor[] tensors) -> Tensor[] + variants: function + python_module: nn + +- func: _nested_tensor_from_tensor_list(Tensor[] list, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + variants: function + dispatch: + CompositeExplicitAutograd: _nested_tensor_from_tensor_list + autogen: _nested_tensor_from_tensor_list.out + +- func: _fw_primal_copy(Tensor self, int level) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: _fw_primal_copy + tags: view_copy + autogen: _fw_primal_copy.out + +- func: _make_dual_copy(Tensor primal, Tensor tangent, int level) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: _make_dual_copy + tags: view_copy + autogen: _make_dual_copy.out + +- func: view_as_real_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: view_as_real_copy + tags: view_copy + autogen: view_as_real_copy.out + +- func: view_as_complex_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: view_as_complex_copy + tags: view_copy + autogen: view_as_complex_copy.out + +- func: _conj_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: _conj_copy + tags: view_copy + autogen: _conj_copy.out + +- func: _neg_view_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: _neg_view_copy + tags: view_copy + autogen: _neg_view_copy.out + +- func: as_strided_copy(Tensor self, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: as_strided_copy_symint + tags: view_copy + autogen: as_strided_copy.out + +- func: _sparse_broadcast_to_copy(Tensor self, int[] size) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: _sparse_broadcast_to_copy + tags: view_copy + autogen: _sparse_broadcast_to_copy.out + +- func: diagonal_copy(Tensor self, int offset=0, int dim1=0, int dim2=1) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: diagonal_copy + tags: view_copy + autogen: diagonal_copy.out + +- func: expand_copy(Tensor self, SymInt[] size, *, bool implicit=False) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: expand_copy_symint + tags: view_copy + autogen: expand_copy.out + +- func: permute_copy(Tensor self, int[] dims) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: permute_copy + tags: view_copy + autogen: permute_copy.out + +- func: _reshape_alias_copy(Tensor self, SymInt[] size, SymInt[] stride) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: _reshape_alias_copy_symint + tags: view_copy + autogen: _reshape_alias_copy.out + +- func: select_copy.int(Tensor self, int dim, SymInt index) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: select_copy_symint + SparseCsrCPU, SparseCsrCUDA, SparseCsrMeta: select_copy_sparse_csr + tags: view_copy + autogen: select_copy.int_out + +- func: detach_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: detach_copy + tags: view_copy + autogen: detach_copy.out + +- func: slice_copy.Tensor(Tensor self, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: slice_copy_Tensor_symint + tags: view_copy + autogen: slice_copy.Tensor_out + +- func: split_copy.Tensor(Tensor self, SymInt split_size, int dim=0) -> Tensor[] + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: split_copy_Tensor_symint + tags: view_copy + +- func: split_with_sizes_copy(Tensor self, SymInt[] split_sizes, int dim=0) -> Tensor[] + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: split_with_sizes_copy_symint + tags: view_copy + +- func: squeeze_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: squeeze_copy + tags: view_copy + autogen: squeeze_copy.out + +- func: squeeze_copy.dim(Tensor self, int dim) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: squeeze_copy_dim + tags: view_copy + autogen: squeeze_copy.dim_out + +- func: squeeze_copy.dims(Tensor self, int[] dim) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: squeeze_copy_dims + tags: view_copy + autogen: squeeze_copy.dims_out + +- func: t_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: t_copy + tags: view_copy + autogen: t_copy.out + +- func: transpose_copy.int(Tensor self, int dim0, int dim1) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: transpose_copy_int + tags: view_copy + autogen: transpose_copy.int_out + +- func: unsqueeze_copy(Tensor self, int dim) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: unsqueeze_copy + tags: view_copy + autogen: unsqueeze_copy.out + +- func: _indices_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: _indices_copy + tags: view_copy + autogen: _indices_copy.out + +- func: _values_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: _values_copy + tags: view_copy + autogen: _values_copy.out + +- func: indices_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: indices_copy + tags: view_copy + autogen: indices_copy.out + +- func: values_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: values_copy + tags: view_copy + autogen: values_copy.out + +- func: crow_indices_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: crow_indices_copy + tags: view_copy + autogen: crow_indices_copy.out + +- func: col_indices_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: col_indices_copy + tags: view_copy + autogen: col_indices_copy.out + +- func: ccol_indices_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: ccol_indices_copy + tags: view_copy + autogen: ccol_indices_copy.out + +- func: row_indices_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: row_indices_copy + tags: view_copy + autogen: row_indices_copy.out + +- func: unbind_copy.int(Tensor self, int dim=0) -> Tensor[] + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: unbind_copy_int + tags: view_copy + +- func: unbind_copy.int_out(Tensor self, int dim=0, *, Tensor(a!)[] out) -> () + variants: function + dispatch: + CompositeExplicitAutograd: unbind_copy_int_out + +- func: split_copy.Tensor_out(Tensor self, SymInt split_size, int dim=0, *, Tensor(a!)[] out) -> () + variants: function + dispatch: + CompositeExplicitAutograd: split_copy_Tensor_out + + +- func: split_with_sizes_copy.out(Tensor self, SymInt[] split_sizes, int dim=0, *, Tensor(a!)[] out) -> () + variants: function + dispatch: + CompositeExplicitAutograd: split_with_sizes_copy_out + CUDA: split_with_sizes_copy_out_cuda + +- func: view_copy(Tensor self, SymInt[] size) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: view_copy_symint + tags: view_copy + autogen: view_copy.out + +- func: view_copy.dtype(Tensor self, ScalarType dtype) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: view_copy_dtype + tags: view_copy + autogen: view_copy.dtype_out + +- func: unfold_copy(Tensor self, int dimension, int size, int step) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: unfold_copy + tags: view_copy + autogen: unfold_copy.out + +- func: alias_copy(Tensor self) -> Tensor + variants: function + dispatch: + CompositeExplicitAutogradNonFunctional: alias_copy + tags: view_copy + autogen: alias_copy.out + +- func: to_padded_tensor(Tensor self, float padding, SymInt[]? output_size=None) -> Tensor + variants: method + dispatch: + NestedTensorCPU: NestedTensor_to_padded_tensor_generic + NestedTensorCUDA: NestedTensor_to_padded_tensor_cuda + autogen: to_padded_tensor.out + +- func: _jagged_to_padded_dense_forward(Tensor values, Tensor[] offsets, SymInt[] max_lengths, float padding_value=0.0) -> Tensor + variants: function + dispatch: + CUDA: _fbgemm_jagged_to_padded_dense_forward + CPU: _jagged_to_padded_dense_forward_cpu + +- func: _padded_dense_to_jagged_forward(Tensor dense, Tensor[] offsets, SymInt? total_L=None) -> Tensor + variants: function + dispatch: + CUDA: _fbgemm_dense_to_jagged_forward_symint + CPU: _padded_dense_to_jagged_forward_cpu + +- func: _nested_from_padded_tensor(Tensor padded, Tensor offsets, Tensor dummy, int ragged_idx=1, Tensor? min_seqlen=None, Tensor? max_seqlen=None, SymInt? sum_S=None) -> Tensor + variants: function + device_check: NoCheck + dispatch: {} + +- func: _nested_tensor_softmax_with_shape(Tensor self, Tensor query) -> Tensor + dispatch: + NestedTensorCPU: NestedTensor_softmax_dropout + NestedTensorCUDA: NestedTensor_softmax_dropout_cuda + tags: nondeterministic_seeded + +- func: _safe_softmax(Tensor self, int dim, ScalarType? dtype=None) -> Tensor + dispatch: + CompositeExplicitAutograd: _safe_softmax + NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: _safe_softmax + +# Apparently, putting "forward" in the name will cause Python bindings to be skipped, so "fwd" it is. +- func: _transformer_encoder_layer_fwd(Tensor src, int embed_dim, int num_heads, Tensor qkv_weight, Tensor qkv_bias, Tensor proj_weight, Tensor proj_bias, bool use_gelu, bool norm_first, float eps, Tensor norm_weight_1, Tensor norm_bias_1, Tensor norm_weight_2, Tensor norm_bias_2, Tensor ffn_weight_1, Tensor ffn_bias_1, Tensor ffn_weight_2, Tensor ffn_bias_2, Tensor? mask=None, int? mask_type=None) -> Tensor + variants: function + dispatch: + CPU, CUDA, NestedTensorCPU, NestedTensorHPU, NestedTensorCUDA: transformer_encoder_layer_forward + autogen: _transformer_encoder_layer_fwd.out + +- func: _native_multi_head_attention(Tensor query, Tensor key, Tensor value, int embed_dim, int num_head, Tensor qkv_weight, Tensor qkv_bias, Tensor proj_weight, Tensor proj_bias, Tensor? mask=None, bool need_weights=True, bool average_attn_weights=True, int? mask_type=None) -> (Tensor, Tensor) + variants: function + dispatch: + CPU, NestedTensorCPU: native_multi_head_attention_cpu + CUDA, NestedTensorCUDA: native_multi_head_attention_cuda + autogen: _native_multi_head_attention.out + +- func: scaled_dot_product_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None, bool enable_gqa=False) -> Tensor + python_module: nn + variants: function + autogen: scaled_dot_product_attention.out + tags: nondeterministic_seeded + +# This aten function is kept so that we can test the choice function from Python +- func: _fused_sdp_choice(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, *, float? scale=None, bool enable_gqa=False) -> int + dispatch: + Meta: _fused_sdp_choice_meta + CPU, NestedTensorCPU: _fused_sdp_choice_cpp + CUDA, NestedTensorCUDA: _fused_sdp_choice_cuda + XPU: _fused_sdp_choice_xpu + tags: nondeterministic_seeded + +- func: _scaled_dot_product_attention_math(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, Tensor? dropout_mask=None, *, float? scale=None, bool enable_gqa=False) -> (Tensor, Tensor) + variants: function + tags: nondeterministic_seeded + +- func: _scaled_dot_product_attention_math_for_mps(Tensor query, Tensor key, Tensor value, Tensor? attn_mask=None, float dropout_p=0.0, bool is_causal=False, Tensor? dropout_mask=None, *, float? scale=None, bool enable_gqa=False) -> (Tensor, Tensor) + dispatch: + MPS: _scaled_dot_product_attention_math_mps + tags: nondeterministic_seeded + +- func: _scaled_dot_product_flash_attention(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor rng_state, Tensor unused, Tensor debug_attn_mask) + dispatch: + CUDA: _scaled_dot_product_flash_attention_cuda + XPU: _scaled_dot_product_flash_attention_xpu + NestedTensorCUDA: _scaled_dot_product_flash_attention_nestedtensor_cuda + tags: nondeterministic_seeded + +- func: _scaled_dot_product_flash_attention.quantized(Tensor query, Tensor key, Tensor value, Tensor? q_descale, Tensor? k_descale, Tensor? v_descale, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor rng_state, Tensor unused, Tensor debug_attn_mask) + dispatch: + CUDA: _scaled_dot_product_flash_attention_cuda_quantized + tags: nondeterministic_seeded + +- func: _scaled_dot_product_flash_attention_for_cpu(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, *, Tensor? attn_mask=None, float? scale=None) -> (Tensor output, Tensor logsumexp) + dispatch: + CPU: _scaled_dot_product_flash_attention_cpu + tags: nondeterministic_seeded + +- func: _scaled_dot_product_fused_attention_overrideable(Tensor query, Tensor key, Tensor value, Tensor? attn_bias=None, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask) + dispatch: + CompositeExplicitAutograd: _scaled_dot_product_fused_attention_overrideable + XPU: _scaled_dot_product_fused_attention_overrideable_xpu + tags: nondeterministic_seeded + +- func: _scaled_dot_product_flash_attention_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, Tensor philox_seed, Tensor philox_offset, *, float? scale=None) -> (Tensor grad_query, Tensor grad_key, Tensor grad_value) + device_check: NoCheck + variants: function + dispatch: + CUDA: _scaled_dot_product_flash_attention_backward_cuda + XPU: _scaled_dot_product_flash_attention_backward_xpu + NestedTensorCUDA: _scaled_dot_product_flash_attention_backward_nested + +- func: _scaled_dot_product_flash_attention_for_cpu_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor logsumexp, float dropout_p, bool is_causal, *, Tensor? attn_mask=None, float? scale=None) -> (Tensor grad_query, Tensor grad_key, Tensor grad_value) + device_check: NoCheck + variants: function + dispatch: + CPU: _scaled_dot_product_flash_attention_cpu_backward + +- func: _scaled_dot_product_fused_attention_overrideable_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor attn_bias, bool[4] grad_input_mask, Tensor out, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, Tensor philox_seed, Tensor philox_offset, *, float? scale=None) -> (Tensor grad_query, Tensor grad_key, Tensor grad_value, Tensor grad_attn_bias) + device_check: NoCheck + variants: function + dispatch: + CompositeExplicitAutograd: _scaled_dot_product_fused_attention_overrideable_backward + +- func: _scaled_dot_product_efficient_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_bias, bool compute_log_sumexp, float dropout_p=0.0, bool is_causal=False, *, float? scale=None) -> (Tensor output, Tensor log_sumexp, Tensor philox_seed, Tensor philox_offset) + dispatch: + CUDA: _scaled_dot_product_efficient_attention_cuda + NestedTensorCUDA: _scaled_dot_product_efficient_attention_nestedtensor_cuda + tags: nondeterministic_seeded + +- func: _scaled_dot_product_efficient_attention_backward(Tensor grad_out_, Tensor query, Tensor key, Tensor value, Tensor attn_bias, Tensor out, Tensor logsumexp, Tensor philox_seed, Tensor philox_offset, float dropout_p, bool[4] grad_input_mask, bool is_causal=False, *, float? scale=None) -> (Tensor, Tensor, Tensor, Tensor) + device_check: NoCheck + dispatch: + CUDA: _scaled_dot_product_efficient_attention_backward_cuda + tags: nondeterministic_seeded + +- func: _scaled_dot_product_cudnn_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_bias, bool compute_log_sumexp, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask) + dispatch: + CUDA: _scaled_dot_product_cudnn_attention_cuda + NestedTensorCUDA: _scaled_dot_product_cudnn_attention_nestedtensor_cuda + tags: nondeterministic_seeded + +- func: _scaled_dot_product_cudnn_attention_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor logsumexp, Tensor philox_seed, Tensor philox_offset, Tensor attn_bias, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, *, float? scale=None) -> (Tensor, Tensor, Tensor) + dispatch: + CUDA: _scaled_dot_product_cudnn_attention_backward_cuda + NestedTensorCUDA: _scaled_dot_product_cudnn_attention_nestedtensor_backward_cuda + tags: nondeterministic_seeded + +- func: _flash_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None, Tensor? block_table=None, int? num_splits=None) -> (Tensor output, Tensor softmax_logsumexp, Tensor rng_state, Tensor unused, Tensor debug_attn_mask) + variants: function + dispatch: + CUDA: _flash_attention_forward + tags: nondeterministic_seeded + +- func: _flash_attention_forward_no_dropout_inplace(Tensor(a!) out, Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None, Tensor? block_table=None, int? num_splits=None) -> Tensor softmax_logsumexp + variants: function + dispatch: + CUDA: _flash_attention_forward_no_dropout_inplace + tags: nondeterministic_seeded + +- func: _flash_attention_forward.quantized(Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, Tensor? q_descale, Tensor? k_descale, Tensor? v_descale, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None) -> (Tensor output, Tensor softmax_logsumexp, Tensor rng_state, Tensor unused, Tensor debug_attn_mask) + variants: function + dispatch: + CUDA: _flash_attention_forward_quantized + tags: nondeterministic_seeded + +- func: _flash_attention_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, Tensor rng_state, Tensor unused, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None) -> (Tensor, Tensor, Tensor) + device_check: NoCheck + variants: function + dispatch: + CUDA: _flash_attention_backward + +# Returns output, logsumexp if compute_logsumexp +- func: _efficient_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? bias, Tensor? cu_seqlens_q, Tensor? cu_seqlens_k, SymInt? max_seqlen_q, SymInt? max_seqlen_k, float dropout_p, int custom_mask_type, bool compute_log_sumexp=False, *, float? scale=None, Tensor? seqlen_k=None, int? window_size=None) -> (Tensor output, Tensor logsumexp, Tensor philox_seed, Tensor philox_offset, SymInt max_seqlen_batch_q, SymInt max_seqlen_batch_k) + variants: function + dispatch: + CUDA: _efficient_attention_forward + tags: nondeterministic_seeded + +- func: _efficient_attention_backward(Tensor grad_out_, Tensor query, Tensor key, Tensor value, Tensor? bias, Tensor out, Tensor? cu_seqlens_q, Tensor? cu_seqlens_k, SymInt max_seqlen_q, SymInt max_seqlen_k, Tensor logsumexp, float dropout_p, Tensor philox_seed, Tensor philox_offset, int custom_mask_type, bool bias_requires_grad, *, float? scale=None, int? num_splits_key=None, int? window_size=None, bool shared_storage_dqdkdv=False) -> (Tensor, Tensor, Tensor, Tensor) + device_check: NoCheck + variants: function + dispatch: + CUDA: _efficient_attention_backward + +- func: _cudnn_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? attn_bias, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, bool compute_log_sumexp, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask) + dispatch: + CUDA: _cudnn_attention_forward + tags: nondeterministic_seeded + +- func: _cudnn_attention_backward(Tensor grad_out, Tensor query, Tensor key, Tensor value, Tensor out, Tensor logsumexp, Tensor philox_seed, Tensor philox_offset, Tensor attn_bias, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, *, float? scale=None) -> (Tensor, Tensor, Tensor) + dispatch: + CUDA: _cudnn_attention_backward + tags: nondeterministic_seeded + +- func: _triton_scaled_dot_attention(Tensor q, Tensor k, Tensor v, float dropout_p=0.0) -> Tensor + variants: function + dispatch: + CUDA: triton_scaled_dot_attention + tags: nondeterministic_seeded + autogen: _triton_scaled_dot_attention.out + +- func: _fill_mem_eff_dropout_mask_(Tensor(a!) self, float dropout_p, int seed, int offset) -> Tensor(a!) + variants: function + dispatch: + CUDA: _fill_mem_eff_dropout_mask_ + tags: nondeterministic_seeded + +- func: _triton_multi_head_attention(Tensor query, Tensor key, Tensor value, int embed_dim, int num_head, Tensor qkv_weight, Tensor qkv_bias, Tensor proj_weight, Tensor proj_bias, Tensor? mask=None) -> Tensor + variants: function + dispatch: + CUDA: triton_multi_head_attention + autogen: _triton_multi_head_attention.out + +- func: special_airy_ai(Tensor x) -> Tensor + python_module: special + structured_delegate: special_airy_ai.out + variants: function + tags: pointwise + +- func: special_airy_ai.out(Tensor x, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA: special_airy_ai_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_bessel_j0(Tensor self) -> Tensor + python_module: special + structured_delegate: special_bessel_j0.out + variants: function + tags: pointwise + +- func: special_bessel_j0.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_bessel_j0_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_bessel_j1(Tensor self) -> Tensor + python_module: special + structured_delegate: special_bessel_j1.out + variants: function + tags: pointwise + +- func: special_bessel_j1.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_bessel_j1_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_bessel_y0(Tensor self) -> Tensor + python_module: special + structured_delegate: special_bessel_y0.out + variants: function + tags: pointwise + +- func: special_bessel_y0.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_bessel_y0_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_bessel_y1(Tensor self) -> Tensor + python_module: special + structured_delegate: special_bessel_y1.out + variants: function + tags: pointwise + +- func: special_bessel_y1.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_bessel_y1_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_t(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_chebyshev_polynomial_t.out + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_t.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_t + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_t.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_t + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_t.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_chebyshev_polynomial_t_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_t.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_t_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_t.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_t_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_u(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_chebyshev_polynomial_u.out + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_u.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_u + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_u.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_u + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_u.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_chebyshev_polynomial_u_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_u.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_u_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_u.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_u_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_v(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_chebyshev_polynomial_v.out + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_v.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_v + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_v.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_v + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_v.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_chebyshev_polynomial_v_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_v.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_v_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_v.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_v_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_w(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_chebyshev_polynomial_w.out + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_w.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_w + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_w.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_w + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_w.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_chebyshev_polynomial_w_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_w.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_w_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_chebyshev_polynomial_w.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_chebyshev_polynomial_w_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_hermite_polynomial_h(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_hermite_polynomial_h.out + variants: function + tags: pointwise + +- func: special_hermite_polynomial_h.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_hermite_polynomial_h + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_hermite_polynomial_h.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_hermite_polynomial_h + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_hermite_polynomial_h.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_hermite_polynomial_h_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_hermite_polynomial_h.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_hermite_polynomial_h_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_hermite_polynomial_h.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_hermite_polynomial_h_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_hermite_polynomial_he(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_hermite_polynomial_he.out + variants: function + tags: pointwise + +- func: special_hermite_polynomial_he.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_hermite_polynomial_he + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_hermite_polynomial_he.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_hermite_polynomial_he + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_hermite_polynomial_he.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_hermite_polynomial_he_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_hermite_polynomial_he.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_hermite_polynomial_he_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_hermite_polynomial_he.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_hermite_polynomial_he_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_laguerre_polynomial_l(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_laguerre_polynomial_l.out + variants: function + tags: pointwise + +- func: special_laguerre_polynomial_l.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_laguerre_polynomial_l + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_laguerre_polynomial_l.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_laguerre_polynomial_l + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_laguerre_polynomial_l.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA: special_laguerre_polynomial_l_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_laguerre_polynomial_l.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_laguerre_polynomial_l_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_laguerre_polynomial_l.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_laguerre_polynomial_l_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_legendre_polynomial_p(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_legendre_polynomial_p.out + variants: function + tags: pointwise + +- func: special_legendre_polynomial_p.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_legendre_polynomial_p + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_legendre_polynomial_p.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_legendre_polynomial_p + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_legendre_polynomial_p.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA: special_legendre_polynomial_p_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_legendre_polynomial_p.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_legendre_polynomial_p_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_legendre_polynomial_p.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_legendre_polynomial_p_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_modified_bessel_i0(Tensor self) -> Tensor + python_module: special + structured_delegate: special_modified_bessel_i0.out + variants: function + tags: pointwise + +- func: special_modified_bessel_i0.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_modified_bessel_i0_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_modified_bessel_i1(Tensor self) -> Tensor + python_module: special + structured_delegate: special_modified_bessel_i1.out + variants: function + tags: pointwise + +- func: special_modified_bessel_i1.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_modified_bessel_i1_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_modified_bessel_k0(Tensor self) -> Tensor + python_module: special + structured_delegate: special_modified_bessel_k0.out + variants: function + tags: pointwise + +- func: special_modified_bessel_k0.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_modified_bessel_k0_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_modified_bessel_k1(Tensor self) -> Tensor + python_module: special + structured_delegate: special_modified_bessel_k1.out + variants: function + tags: pointwise + +- func: special_modified_bessel_k1.out(Tensor self, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_modified_bessel_k1_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_scaled_modified_bessel_k0(Tensor x) -> Tensor + python_module: special + structured_delegate: special_scaled_modified_bessel_k0.out + variants: function + tags: pointwise + +- func: special_scaled_modified_bessel_k0.out(Tensor x, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_scaled_modified_bessel_k0_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_scaled_modified_bessel_k1(Tensor x) -> Tensor + python_module: special + structured_delegate: special_scaled_modified_bessel_k1.out + variants: function + tags: pointwise + +- func: special_scaled_modified_bessel_k1.out(Tensor x, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_scaled_modified_bessel_k1_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_t(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_shifted_chebyshev_polynomial_t.out + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_t.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_t + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_t.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_t + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_t.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_shifted_chebyshev_polynomial_t_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_t.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_t_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_t.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_t_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_u(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_shifted_chebyshev_polynomial_u.out + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_u.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_u + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_u.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_u + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_u.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_shifted_chebyshev_polynomial_u_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_u.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_u_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_u.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_u_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_v(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_shifted_chebyshev_polynomial_v.out + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_v.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_v + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_v.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_v + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_v.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_shifted_chebyshev_polynomial_v_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_v.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_v_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_v.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_v_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_w(Tensor x, Tensor n) -> Tensor + device_check: NoCheck + python_module: special + structured_delegate: special_shifted_chebyshev_polynomial_w.out + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_w.x_scalar(Scalar x, Tensor n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_w + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_w.n_scalar(Tensor x, Scalar n) -> Tensor + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_w + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_w.out(Tensor x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + device_check: NoCheck + dispatch: + CPU, CUDA, MPS: special_shifted_chebyshev_polynomial_w_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_w.x_scalar_out(Scalar x, Tensor n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_w_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_shifted_chebyshev_polynomial_w.n_scalar_out(Tensor x, Scalar n, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CompositeExplicitAutograd: special_shifted_chebyshev_polynomial_w_out + device_check: NoCheck + python_module: special + variants: function + tags: pointwise + +- func: special_spherical_bessel_j0(Tensor x) -> Tensor + python_module: special + structured_delegate: special_spherical_bessel_j0.out + variants: function + tags: pointwise + +- func: special_spherical_bessel_j0.out(Tensor x, *, Tensor(a!) out) -> Tensor(a!) + dispatch: + CPU, CUDA, MPS: special_spherical_bessel_j0_out + python_module: special + structured_inherits: TensorIteratorBase + structured: True + variants: function + tags: pointwise + +# Aux function used in the test TestPythonDispatch.test_kwarg_only_and_positional_default +# within test/test_python_dispatch.py +- func: _foobar(Tensor self, bool arg1=True, bool arg2=True, *, bool arg3=True) -> Tensor + dispatch: + CPU: foobar + autogen: _foobar.out + +- func: _fused_adam_(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] exp_avgs, Tensor(d!)[] exp_avg_sqs, Tensor(e!)[] max_exp_avg_sqs, Tensor[] state_steps, *, float lr, float beta1, float beta2, float weight_decay, float eps, bool amsgrad, bool maximize, Tensor? grad_scale=None, Tensor? found_inf=None) -> () + # Unlike "foreach" functions, lists of tensors should be guaranteed to be on the same device (for now). + variants: function + dispatch: + CPU: _fused_adam_kernel_cpu_ + CUDA: _fused_adam_kernel_cuda_ + MPS: _fused_adam_kernel_mps_ + autogen: _fused_adam, _fused_adam.out + +- func: _fused_adam_.tensor_lr(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] exp_avgs, Tensor(d!)[] exp_avg_sqs, Tensor(e!)[] max_exp_avg_sqs, Tensor[] state_steps, *, Tensor lr, float beta1, float beta2, float weight_decay, float eps, bool amsgrad, bool maximize, Tensor? grad_scale=None, Tensor? found_inf=None) -> () + # Unlike "foreach" functions, lists of tensors should be guaranteed to be on the same device (for now), + # but still skip the device check as the Tensor LR can be on CPU + device_check: NoCheck + variants: function + dispatch: + CPU: _fused_adam_kernel_cpu_ + CUDA: _fused_adam_kernel_cuda_ + MPS: _fused_adam_kernel_mps_ + autogen: _fused_adam.tensor_lr, _fused_adam.tensor_lr_out + +- func: _fused_adamw_(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] exp_avgs, Tensor(d!)[] exp_avg_sqs, Tensor(e!)[] max_exp_avg_sqs, Tensor[] state_steps, *, float lr, float beta1, float beta2, float weight_decay, float eps, bool amsgrad, bool maximize, Tensor? grad_scale=None, Tensor? found_inf=None) -> () + # Unlike "foreach" functions, lists of tensors should be guaranteed to be on the same device (for now). + variants: function + dispatch: + CPU: _fused_adamw_kernel_cpu_ + CUDA: _fused_adamw_kernel_cuda_ + MPS: _fused_adamw_kernel_mps_ + autogen: _fused_adamw, _fused_adamw.out + +- func: _fused_adamw_.tensor_lr(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] exp_avgs, Tensor(d!)[] exp_avg_sqs, Tensor(e!)[] max_exp_avg_sqs, Tensor[] state_steps, *, Tensor lr, float beta1, float beta2, float weight_decay, float eps, bool amsgrad, bool maximize, Tensor? grad_scale=None, Tensor? found_inf=None) -> () + # Unlike "foreach" functions, lists of tensors should be guaranteed to be on the same device (for now), + # but still skip the device check as the Tensor LR can be on CPU + device_check: NoCheck + variants: function + dispatch: + CPU: _fused_adamw_kernel_cpu_ + CUDA: _fused_adamw_kernel_cuda_ + MPS: _fused_adamw_kernel_mps_ + autogen: _fused_adamw.tensor_lr, _fused_adamw.tensor_lr_out + +- func: _fused_sgd_(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] momentum_buffer_list, *, float weight_decay, float momentum, float lr, float dampening, bool nesterov, bool maximize, bool is_first_step, Tensor? grad_scale=None, Tensor? found_inf=None) -> () + # Unlike "foreach" functions, lists of tensors should be guaranteed to be on the same device (for now). + variants: function + dispatch: + CPU: _fused_sgd_kernel_cpu_ + CUDA: _fused_sgd_kernel_cuda_ + MPS: _fused_sgd_kernel_mps_ + autogen: _fused_sgd, _fused_sgd.out + +- func: _fused_sgd_.tensor_lr(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] momentum_buffer_list, *, float weight_decay, float momentum, Tensor lr, float dampening, bool nesterov, bool maximize, bool is_first_step, Tensor? grad_scale=None, Tensor? found_inf=None) -> () + # Unlike "foreach" functions, lists of tensors should be guaranteed to be on the same device (for now). + # but still skip the device check as the Tensor LR can be on CPU + device_check: NoCheck + variants: function + dispatch: + CPU: _fused_sgd_kernel_cpu_ + CUDA: _fused_sgd_kernel_cuda_ + MPS: _fused_sgd_kernel_mps_ + autogen: _fused_sgd.tensor_lr, _fused_sgd.tensor_lr_out + +- func: _fused_adagrad_(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] state_sums, Tensor(d!)[] state_steps, *, float lr, float lr_decay, float weight_decay, float eps, bool maximize, Tensor? grad_scale=None, Tensor? found_inf=None) -> () + variants: function + dispatch: + CPU: _fused_adagrad_kernel_cpu_ + CUDA: _fused_adagrad_kernel_cuda_ + autogen: _fused_adagrad, _fused_adagrad.out + +- func: _fused_adagrad_.tensor_lr(Tensor(a!)[] self, Tensor(b!)[] grads, Tensor(c!)[] state_sums, Tensor[] state_steps, *, Tensor lr, float lr_decay, float weight_decay, float eps, bool maximize, Tensor? grad_scale=None, Tensor? found_inf=None) -> () + device_check: NoCheck + variants: function + dispatch: + CPU: _fused_adagrad_kernel_cpu_ + CUDA: _fused_adagrad_kernel_cuda_ + autogen: _fused_adagrad.tensor_lr, _fused_adagrad.tensor_lr_out + +# This op is ONLY used by pytorch/XLA in functionalization, and should never show up in vanilla eager mode or in any pytorch tracing contexts. +- func: _propagate_xla_data(Tensor input, Tensor output) -> () + variants: function diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/native/tags.yaml b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/native/tags.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d76e38dae183d06348d3015ecaca9542e47a7da5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/native/tags.yaml @@ -0,0 +1,105 @@ +# This yaml file contains all the possible tags that can be defined in `tags` in `native_functions.yaml` + +- tag: inplace_view + desc: | + This tag indicates if an operator *only* modifies the tensor metadata +- tag: pt2_compliant_tag + desc: | + This tag indicates if the operator is guaranteed to + work with the PT2 compilation APIs (torch.compile, + torch.export, etc). If you add this tag to an + operator, please use + `torch.testing._internal.optest.opcheck` to test that + the operator has been registered correctly and + works with torch.compile +- tag: view_copy + desc: | + This tag indicates operators that are *_copy* variants + of view/aliasing operators. If an operator has a view_copy tag, + then it should have the name {op}_copy, where {op} is a view operator. +- tag: dynamic_output_shape + desc: | + This tag indicates if an operator's output's shape depends on input Tensor + data. +- tag: data_dependent_output + desc: | + Operator has a non-Tensor output whose value is dependent on the data + of Tensor inputs. Among other things, this implies that this operator + cannot be run with meta tensor (since data is not available), nor + can it be symbolically traced. +- tag: generated + desc: | + This tag indicates that the operator doesn't have an explicit entry in + native_functions.yaml, and instead was generated automatically by the codegen. +- tag: nondeterministic_seeded + desc: | + This tag indicates if an operator is nondeterministically seeded + (i.e., is random) such that the operator intentionally produces + different results when run twice on the same inputs, but this randomness + is controlled by a Generator which, if reseeded would give you the + same result. +- tag: nondeterministic_bitwise + desc: | + This tag indicates if an operator doesn't guarantee bitwise equivalence + across different runs of an operator with identical inputs. +- tag: needs_exact_strides + desc: | + This tag indicates that the operator should be passed Tensors following + the same strides as observed in eager when compiled in inductor. + Only one of {needs_exact_strides, needs_contiguous_strides, needs_fixed_stride_order, flexible_layout} + can apply; if multiple are assigned then we assume the most restrictive one. +- tag: needs_contiguous_strides + desc: | + This tag indicates that the operator should be passed contiguous Tensors. + Failure to do so will result in undefined behavior. +- tag: needs_fixed_stride_order + desc: | + This tag indicates that the operator should be passed Tensors following + the same stride permutation as observed in eager when compiled in inductor. + Only one of {needs_exact_strides, needs_contiguous_strides, needs_fixed_stride_order, flexible_layout} + can apply; if multiple are assigned then we assume the most restrictive one. +- tag: flexible_layout + desc: | + This tag indicates that the custom operator can accept inputs with varying + strides/storage_offset and that when compiled, Inductor is allowed to change + the strides/storage_offset of inputs to the custom operator. + Only one of {needs_exact_strides, needs_contiguous_strides, needs_fixed_stride_order, flexible_layout} + can apply; if multiple are assigned then we assume the most restrictive one. + +# NOTE [Core ATen Ops] +- tag: core + desc: | + Core aten ops is a subset of aten ops that remains after aten-to-aten decomposition and + functionalization pass. Core aten ops are fully functional and adhere to single static + assignment (SSA): this implies there will be no `inplace` or `_out` variants in this opset. + This opset is designed to serve as the functional IR to interface with compiler backends. + In contrast to primTorch, core aten opset doesn't decompose ops into explicit + type promotion and broadcasting ops. + Core aten ops is also effectively the opset produced by torchdynamo.export(aten_graph=True), + and thus can be used as an opset for export purpose. +- tag: pointwise + desc: | + Pointwise operators are operators where each element of the output is computed only by accessing + the corresponding element of all the broadcasted inputs. The output shape will be the broadcasted + shape of the inputs. +- tag: maybe_aliasing_or_mutating + desc: | + For some ops, we can't statically determine whether the op is functional or not. Note that this is only + relevant to CIA ops that decompose before functionalization/autograd. It is useful to + know this information for export as we would want to decompose these ops as they are unsafe to be + preserved. +- tag: cudagraph_unsafe + desc: | + This operator does not support cudagraphs. The presence of this tag on an operator will cause + Inductor to split the graph around this operator. Note that operators without this tag may still + not support CUDAGraphs. Inductor may have other hardcoded lists around that. +- tag: reduction + desc: | + This tag indicates that an operator performs a reduction operation, computing aggregate values + (sum, mean, max, min, etc.) across one or more dimensions of the input tensor(s). +- tag: out_variant + desc: | + This tag indicates that the operator is an out variant of a functional + operator. This tag only applies to custom ops. Out variant operators + write their results to pre-allocated output tensors (the out args) + rather than allocating new tensors. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ATenOpList.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ATenOpList.cpp new file mode 100644 index 0000000000000000000000000000000000000000..5de3424857e236917eb68940e7904446de59f586 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ATenOpList.cpp @@ -0,0 +1,36 @@ +#include + +#include +#include +#include +#include +#include + +// ${generated_comment} + +namespace at { + +namespace { +struct OpNameEquals final { + bool operator()(const std::pair& lhs, const std::pair& rhs) const { + return 0 == strcmp(lhs.first, rhs.first) && 0 == strcmp(lhs.second, rhs.second); + } +}; + +struct OpNameHash final { + size_t operator()(const std::pair& p) const { + // use std::hash because std::hash would hash pointers and not pointed-to strings + return std::hash()(p.first) ^ (~ std::hash()(p.second)); + } +}; +} + +bool is_custom_op(const c10::OperatorName& opName) { + static std::unordered_set, OpNameHash, OpNameEquals> ops { + ${aten_ops} + {"", ""} + }; + return ops.count(std::make_pair( + opName.name.c_str(), opName.overload_name.c_str())) == 0; +} +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/CompositeViewCopyKernels.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/CompositeViewCopyKernels.cpp new file mode 100644 index 0000000000000000000000000000000000000000..47097d7aa4320674bec4bddbb5ac861309334f0c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/CompositeViewCopyKernels.cpp @@ -0,0 +1,73 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +#include +$ops_headers +#endif + +namespace at { +namespace native { + +// This file contains a number of kernels for aten functions that are fully code-generated. +// TODO: rename this file to something more generic. + +namespace { +at::Tensor clone_arg(const at::Tensor& t) { + return t.clone(); +} + +std::vector clone_arg(const at::TensorList& t_list) { + std::vector out(t_list.size()); + for (const auto& i : c10::irange(t_list.size())) { + out[i] = t_list[i].clone(); + } + return out; +} + +// duped with gen_resize_out_helper from structured kernels +void copy_arg(const at::Tensor& dst, const at::Tensor& src) { + TORCH_CHECK(src.dtype() == dst.dtype(), + "Expected out tensor to have dtype ", src.dtype(), ", but got ", dst.dtype(), " instead"); + TORCH_CHECK(src.device() == dst.device(), + "Expected out tensor to have device ", src.device(), ", but got ", dst.device(), " instead"); + dst.copy_(src); +} + +void copy_arg(const at::TensorList& dst, const at::TensorList& src) { + TORCH_INTERNAL_ASSERT(dst.size() == src.size()); + for (const auto& i : c10::irange(dst.size())) { + copy_arg(dst[i], src[i]); + } +} + +// TODO: this doesn't handle restriding empty tensors correctly; see +// gen_resize_out_helper for the correct algorithm + +void resize_out_helper(const at::Tensor& dst, const at::Tensor& src) { + at::native::resize_output(dst, src.sizes()); +} + +void resize_out_helper(const at::TensorList& dst, const at::TensorList& src) { + TORCH_INTERNAL_ASSERT(dst.size() == src.size()); + for (const auto& i : c10::irange(dst.size())) { + at::native::resize_output(dst[i], src[i].sizes()); + } +} +} + + +${CompositeViewCopyKernel_Definitions} + +${GeneratedCompositeFunctional_Definitions} + +${GeneratedCompositeOut_Definitions} + +} // namespace native +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyFunction.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyFunction.h new file mode 100644 index 0000000000000000000000000000000000000000..c92d5eb3898ecea0fb9e1f79c2725d1bc6dfa7fb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyFunction.h @@ -0,0 +1,23 @@ +#pragma once +// ${generated_comment} + +// NB: The implementing C++ file is RegisterDispatchKey.cpp + +// The only #includes we need are for custom classes that have defaults in the C++ API +#include +#include +#include + +// Forward declarations of any types needed in the operator signatures. +// We can't directly include these classes because it will cause circular include dependencies. +// This file is included by TensorBody.h, which defines the Tensor class. +#include + +namespace at { + +namespace ${dispatch_namespace} { + +${dispatch_namespaced_declarations} + +} // namespace ${dispatch_namespace} +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyFunctions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyFunctions.h new file mode 100644 index 0000000000000000000000000000000000000000..35f43297fdd9ca9f932c8c53b5b773f1b9b8a427 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyFunctions.h @@ -0,0 +1,29 @@ +#include + +// TODO Undo all logic introduced for Note [Avoiding Include Cycles In Static Dispatch] +// Code introduced to avoid cyclic dependency in static dispatch is no longer +// needed as static dispatch logic is moved from TensorBody.h, which caused cycles in the first place, +// to Operators.cpp for supporting multiple backends with multiple kernels. +// +// Note [Avoiding Include Cycles In Static Dispatch] +// In order to avoid #include cycles in the static dispatch build, we've carefully split out +// the static function definition files into {DispatchKey}Functions.h and {DispatchKey}Functions_inl.h. +// +// Without this split, the include cycle looks like TensorBody.h -> CPUFunctions.h -> TensorBody.h. +// - TensorBody.h #includes CPUFunctions.h in the static dispatch build, because the tensor methods +// all need to call into the fastpath C++ API defined in CPUFunctions.h. The methods are also all +// directly inlined into TensorBody.h. +// - CPUFunctions.h #includes TensorBody.h because it contains function declarations for the entire C++ API, +// which include functions that have defaultable std::optional arguments. +// That requires knowing the full Tensor class definition. +// +// We break the cycle by doing the following: +// - Split out CPUFunction.h into two files: CPUFunctions.h and CPUFunctions_inl.h +// - CPUFunction.h is a dummy file that just includes the Tensor class and includes CPUFunctions_inl., +// - CPUFunctions_inl.h includes everything else +// - (only in the static dispatch build) TensorBody.h makes sure to finish defining the Tensor class, +// and then it includes CPUFunctions_inl.h. +// - All other files that want the cpu fastpath functions can include CPUFunctions.h directly. +// - This also means that static dispatch build, CPUFunctions.h only needs to +// #include TensorBody.h, and it will automatically bring in CPUFunctions_inl.h. +${inline_headers} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyFunctions_inl.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyFunctions_inl.h new file mode 100644 index 0000000000000000000000000000000000000000..fbb71c2cb123cb21fb57ec32341d86bff06f6a17 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyFunctions_inl.h @@ -0,0 +1,22 @@ +#pragma once +// ${generated_comment} + +// NB: The implementing C++ file is RegisterDispatchKey.cpp + +// The only #includes we need are for custom classes that have defaults in the C++ API +#include +#include +#include + +#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS) +#error This change adds a dependency on all pytorch operators, meaning the \ + file will need to be re-compiled every time an operator is changed or added. \ + Consider including a specific operator from \ + . \ + See NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS]. +#endif + +${DispatchKeyFunctions_inl_includes} + + +${dispatch_namespaced_declarations} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyNativeFunctions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyNativeFunctions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..7647f459a744b2eacfac6aaea4f49b86babbb234 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyNativeFunctions.cpp @@ -0,0 +1,13 @@ +// ${generated_comment} +${includes} +${native_functions_include} + +namespace { +${helper_fns} +} // namespace + +${namespace_prologue} + +${native_function_definitions} + +${namespace_epilogue} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyNativeFunctions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyNativeFunctions.h new file mode 100644 index 0000000000000000000000000000000000000000..b45a17b5922f8a0b76e0237616914ce9969efca5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/DispatchKeyNativeFunctions.h @@ -0,0 +1,19 @@ +#pragma once + +// an external backend might generate file within its code tree +// and check all the source files within the tree with clang-format. +// so, disable it since the backend might have a different config. +// clang-format off + +// ${generated_comment} + +#include + +${namespace_prologue} + +struct ${class_name} { + +${dispatch_declarations} + +}; +${namespace_epilogue} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Function.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Function.h new file mode 100644 index 0000000000000000000000000000000000000000..73096afbf11571cbe4147bb63f035a054ca842db --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Function.h @@ -0,0 +1,27 @@ +#pragma once + +// ${generated_comment} + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +${static_dispatch_ops_headers} + +${operator_includes} + +namespace at { + +${function_definitions} + +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/FunctionalInverses.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/FunctionalInverses.h new file mode 100644 index 0000000000000000000000000000000000000000..b15cd09a6c65da3127be8245b87bff2f8c795a3d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/FunctionalInverses.h @@ -0,0 +1,23 @@ +#pragma once + +// ${generated_comment} + +#include +#include + +namespace at { +namespace functionalization { + +struct FunctionalInverses { + +${view_inverse_declarations} + +// NB: These are not generated! They're manually implemented in the template. +// TODO: Change codegen to generate these. See the following link: +// https://github.com/pytorch/pytorch/blob/main/torchgen/model.py#L2583-L2585 +static at::Tensor chunk_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int64_t mutated_view_idx, int chunks, int dim); +static at::Tensor narrow_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int dim, c10::SymInt start, c10::SymInt length); + +}; +} +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..f210402e543aa2de27ea0f510bb869e0c7010e22 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Functions.cpp @@ -0,0 +1,105 @@ +#include + +#include +#include +#include + +namespace at { + +Tensor TensorMaker::make_tensor() { + AutoDispatchBelowADInplaceOrView guard{}; // TODO: Remove. + tracer::impl::NoTracerDispatchMode tracer_guard{}; + + check_size_nonnegative(sizes_); + + TORCH_CHECK_VALUE( + !deleter_ || !ctx_, + "The deleter and context arguments are mutually exclusive."); + + if (device_ == std::nullopt) { + device_ = globalContext().getDeviceFromPtr(data_, opts_.device().type()); + } + + if (opts_.device().has_index()) { + // clang-format off + TORCH_CHECK_VALUE( + opts_.device() == *device_, + "Specified device ", opts_.device(), " does not match device of data ", *device_); + // clang-format on + } + + std::size_t size_bytes = computeStorageSize(); + + DataPtr data_ptr{}; + if (deleter_) { + data_ptr = makeDataPtrFromDeleter(); + } else { + data_ptr = makeDataPtrFromContext(); + } + + TORCH_CHECK(!resizeable_ || allocator_ != nullptr, "Must specify an allocator with allocator() if you want to use resizeable_storage()"); + Storage storage{Storage::use_byte_size_t{}, size_bytes, std::move(data_ptr), /*allocator=*/allocator_, /*resizable=*/resizeable_}; + + Tensor tensor = detail::make_tensor( + std::move(storage), opts_.computeDispatchKey(), opts_.dtype()); + + TensorImpl* tensor_impl = tensor.unsafeGetTensorImpl(); + if (strides_) { + tensor_impl->set_sizes_and_strides(sizes_, *strides_); + } else { + tensor_impl->set_sizes_contiguous(sizes_); + } + if (storage_offset_) { + tensor_impl->set_storage_offset(*storage_offset_); + } + + tensor_impl->set_requires_grad(opts_.requires_grad()); + + return tensor; + } + + std::size_t TensorMaker::computeStorageSize() const noexcept { + std::size_t itemsize = opts_.dtype().itemsize(); + + if (strides_) { + auto storage_size = detail::computeStorageNbytes(sizes_, *strides_, itemsize); + if (storage_offset_) { + storage_size += storage_offset_.value() * itemsize; + } + return storage_size; + } + + std::size_t size = 1; + for (std::int64_t s : sizes_) { + size *= static_cast(s); + } + auto storage_size = size * itemsize; + if (storage_offset_) { + storage_size += storage_offset_.value() * itemsize; + } + return storage_size; + } + + inline DataPtr TensorMaker::makeDataPtrFromDeleter() noexcept { + return InefficientStdFunctionContext::makeDataPtr(data_, std::move(deleter_), *device_); + } + + inline DataPtr TensorMaker::makeDataPtrFromContext() noexcept { + return DataPtr{data_, ctx_.release(), ctx_.get_deleter(), *device_}; + } + + IntArrayRef TensorMaker::makeTempSizes() const noexcept { + static std::int64_t zeros[5] = {0, 0, 0, 0, 0}; + if (opts_.has_memory_format()) { + MemoryFormat format = *opts_.memory_format_opt(); + if (format == MemoryFormat::ChannelsLast) { + return IntArrayRef(zeros, 4); + } + if (format == MemoryFormat::ChannelsLast3d) { + return IntArrayRef(zeros, 5); + } + } + return IntArrayRef(zeros, 1); + } + +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Functions.h new file mode 100644 index 0000000000000000000000000000000000000000..b1feaf9d4daa9786359c97434e4c59d3c75778c7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Functions.h @@ -0,0 +1,143 @@ +#pragma once + +// ${generated_comment} + +#ifdef TORCH_ASSERT_NO_OPERATORS +#error This change adds a dependency on native_functions.yaml, \ + meaning the file will need to be re-compiled every time an operator \ + is changed or added. Consider if your change would be better placed in \ + another file, or if a more specific header might achieve the same goal. \ + See NOTE: [Tensor vs. TensorBase] +#endif + +#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS) +#error This change adds a dependency on all pytorch operators, meaning the \ + file will need to be re-compiled every time an operator is changed or added. \ + Consider including a specific operator from and \ + see NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS]. +#endif + +// NOTE: [TORCH_ASSERT_ONLY_METHOD_OPERATORS] +// +// In ATen, certain generated headers files include the definitions of +// every single operator in PyTorch. Unfortunately this means every +// time an operator signature is updated or changed in +// native_functions.yaml, you (and every other PyTorch developer) need +// to recompile every source file that includes any of these headers. +// +// To break up these header dependencies, and improve incremental +// build times for all PyTorch developers. These headers are split +// into per-operator headers in the `ATen/ops` folder. This limits +// incremental builds to only changes to methods of `Tensor`, or files +// that use the specific operator being changed. With `at::sum` as an +// example, you should include +// +// // instead of ATen/Functions.h +// // instead of ATen/NativeFunctions.h +// // instead of ATen/Operators.h +// // instead of ATen/CPUFunctions.h +// +// However, even if you're careful to use this in your own code. +// `Functions.h` might be included indirectly through another header +// without you realising. To avoid this, you can add +// +// #define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// +// to the top of your source file. This way any time the non-specific +// headers are included, the compiler will error out. +// +// Also, be aware that `ops` are not available in all build +// configurations (namely fb-internal) so you must guard these +// includes with `#ifdef AT_PER_OPERATOR_HEADERS`. e.g. +// +// #ifndef AT_PER_OPERATOR_HEADERS +// #include +// #else +// #include +// #endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +${Functions_includes} + +namespace at { + +${Functions_declarations} + +// Special C++ only overloads for std()-like functions (See gh-40287) +// These are needed because int -> bool conversion takes precedence over int -> IntArrayRef +// So, for example std(0) would select the std(unbiased=False) overload +inline Tensor var(const Tensor& self, int dim) { + return at::var(self, IntArrayRef{dim}); +} +inline std::tuple var_mean(const Tensor& self, int dim) { + return at::var_mean(self, IntArrayRef{dim}); +} +inline Tensor std(const Tensor& self, int dim) { + return at::std(self, IntArrayRef{dim}); +} +inline std::tuple std_mean(const Tensor& self, int dim) { + return at::std_mean(self, IntArrayRef{dim}); +} + +inline int64_t numel(const Tensor& tensor) { + return tensor.numel(); +} + +inline int64_t size(const Tensor& tensor, int64_t dim) { + return tensor.size(dim); +} + +inline int64_t stride(const Tensor& tensor, int64_t dim) { + return tensor.stride(dim); +} + +inline bool is_complex(const Tensor& tensor) { + return tensor.is_complex(); +} + +inline bool is_floating_point(const Tensor& tensor) { + return tensor.is_floating_point(); +} + +inline bool is_signed(const Tensor& tensor) { + return tensor.is_signed(); +} + +inline bool is_inference(const Tensor& tensor) { + return tensor.is_inference(); +} + +inline bool _is_zerotensor(const Tensor& tensor) { + return tensor._is_zerotensor(); +} + +inline bool is_conj(const Tensor& tensor) { + return tensor.is_conj(); +} + +inline Tensor conj(const Tensor& tensor) { + return tensor.conj(); +} + +inline bool is_neg(const Tensor& tensor) { + return tensor.is_neg(); +} + +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/LazyIr.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/LazyIr.h new file mode 100644 index 0000000000000000000000000000000000000000..9190ff8243d316fd2bd472bb3f0603701761bdb7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/LazyIr.h @@ -0,0 +1,19 @@ +#pragma once + +// This file contains autogenerated LazyTensor IR nodes +${lazy_ir_sysinc} +${lazy_ir_inc} + +${namespace_prologue} +using at::operator<<; + +// kNullValue is used to contribute a static hash value any time +// a node has an Optional input that is nullopt. It is important +// to differentiate between HASH(std::nullopt, something) and HASH(something, std::nullopt), +// and using kNullValue in the hash function in the order of arguments +// serves this purpose. +static const torch::lazy::Value kNullValue = torch::lazy::Value(); + +${ir_declarations} + +${namespace_epilogue} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/LazyNonNativeIr.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/LazyNonNativeIr.h new file mode 100644 index 0000000000000000000000000000000000000000..18eaf6da52e4b3654becac6cc89849bc0806ae09 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/LazyNonNativeIr.h @@ -0,0 +1,11 @@ +#pragma once + +${lazy_non_native_ir_inc} + +// This file contains autogenerated LazyTensor Non Native IR nodes + +${namespace_prologue} + +${non_native_ir_nodes} + +${namespace_epilogue} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/MethodOperators.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/MethodOperators.h new file mode 100644 index 0000000000000000000000000000000000000000..0e192cd05ef3c78fa74848c93de32150c1e3fd8b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/MethodOperators.h @@ -0,0 +1,24 @@ +#pragma once + +// ${generated_comment} + +#ifdef TORCH_ASSERT_NO_OPERATORS +#error This change adds a dependency on native_functions.yaml, \ + meaning the file will need to be re-compiled every time an operator \ + is changed or added. Consider if your change would be better placed in \ + another file, or if a more specific header might achieve the same goal. \ + See NOTE: [Tensor vs. TensorBase] +#endif + +// Forward declarations of any types needed in the operator signatures. +// We can't directly include these classes because it will cause circular include dependencies. +// This file is included by TensorBody.h, which defines the Tensor class. +#include + +${MethodOperators_includes} + +namespace at { +namespace _ops { +${MethodOperators_declarations} +} // namespace _ops +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeFunction.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeFunction.h new file mode 100644 index 0000000000000000000000000000000000000000..a5441ad85d1d5e28c4e31dd3f0dc7f66dfbff9e7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeFunction.h @@ -0,0 +1,17 @@ +#pragma once + +// ${generated_comment} + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +${extra_includes} + +${native_function_declarations} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeFunctions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeFunctions.h new file mode 100644 index 0000000000000000000000000000000000000000..9dc972495ca038bddb7b887c39c2e0507e487213 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeFunctions.h @@ -0,0 +1,33 @@ +#pragma once + +// ${generated_comment} + +#ifdef TORCH_ASSERT_NO_OPERATORS +#error This change adds a dependency on native_functions.yaml, \ + meaning the file will need to be re-compiled every time an operator \ + is changed or added. Consider if your change would be better placed in \ + another file, or if a more specific header might achieve the same goal. \ + See NOTE: [Tensor vs. TensorBase] +#endif + +#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS) +#error This change adds a dependency on all pytorch operators, meaning the \ + file will need to be re-compiled every time an operator is changed or added. \ + Consider including a specific operator from \ + and see NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS]. +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +${NativeFunctions_includes} + +${NativeFunctions_declarations} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeMetaFunction.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeMetaFunction.h new file mode 100644 index 0000000000000000000000000000000000000000..6522c97546d0498e4b3825fb4eafefbb34c71911 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeMetaFunction.h @@ -0,0 +1,23 @@ +#pragma once + +// ${generated_comment} + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace at { +namespace meta { + +${meta_function_declarations} + +} // namespace native +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeMetaFunctions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeMetaFunctions.h new file mode 100644 index 0000000000000000000000000000000000000000..89989e2121c9aa34a4583205c3541a04edd36700 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/NativeMetaFunctions.h @@ -0,0 +1,19 @@ +#pragma once + +// ${generated_comment} + +#include +#include +#include +#include + +${NativeMetaFunctions_includes} + +namespace at { + +namespace meta { + +${NativeMetaFunctions_declarations} + +} // namespace meta +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Operator.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Operator.h new file mode 100644 index 0000000000000000000000000000000000000000..ed220f917290c2062481eb53dca232b47d180e2d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Operator.h @@ -0,0 +1,19 @@ +#pragma once + +// ${generated_comment} + +#include +#include +#include + +// Forward declarations of any types needed in the operator signatures. +// We can't directly include these classes because it will cause circular include dependencies. +// This file is included by TensorBody.h, which defines the Tensor class. +#include + +namespace at { +namespace _ops { + +${declarations} + +}} // namespace at::_ops diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Operators.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Operators.cpp new file mode 100644 index 0000000000000000000000000000000000000000..082bb67c3e2043f2c36b29345f57048ec2e9eea7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Operators.cpp @@ -0,0 +1,19 @@ +#include +#include + +// ${generated_comment} +// NOTE See [Sharded File] comment in VariableType + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +${operator_headers} +#endif + +${static_dispatch_extra_headers} + +namespace at { namespace _ops { + +${definitions} + +}} // namespace at::_ops diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Operators.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Operators.h new file mode 100644 index 0000000000000000000000000000000000000000..e74b96ef3d5c6b6d50fe63eac4dca51f0655daa5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/Operators.h @@ -0,0 +1,74 @@ +#pragma once + +// ${generated_comment} + +#ifdef TORCH_ASSERT_NO_OPERATORS +#error This change adds a dependency on native_functions.yaml, \ + meaning the file will need to be re-compiled every time an operator \ + is changed or added. Consider if your change would be better placed in \ + another file, or if a more specific header might achieve the same goal. \ + See NOTE: [Tensor vs. TensorBase] +#endif + +#if defined(AT_PER_OPERATOR_HEADERS) && defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS) +#error This change adds a dependency on all pytorch operators, meaning the \ + file will need to be re-compiled every time an operator is changed or added. \ + Consider including a specific operator from \ + and see NOTE [TORCH_ASSERT_ONLY_METHOD_OPERATORS]. +#endif + +#include +#include +#include +#include +#include +#include +#include +#include + +${Operators_includes} + +// Extension writers: do you write wrapper functions? Are you frustrated with +// resolving overloads of operators? Are you frustrated with dealing with +// pointer-to-methods and resolving overloads of pointer-to-methods?? Look no +// further, this is the utility for you. +// +// Given an operator schema: aten::op.overload(... +// +// Use ATEN_FN2(op, overload) to get a *function* version of the operator +// that is guaranteed to not be overloaded. This means that you can safely +// decltype(&ATEN_FN2(op, overload)) it. NB: the 2 means this macro takes 2 args. +// +// Given an operator schema without an overload name: aten::op(... +// +// Use ATEN_FN(op) to get an unambiguous *function* version of the operator. +// +// There is some interesting behavior for out= operations. +// ATEN_FN2(sin, out) gives a function that is *faithful* to the schema; +// that is, the order of arguments is exactly what it looks like in the schema. + +#define ATEN_FN2(op_name, overload) at::_ops::op_name##_##overload::call +#define ATEN_FN(op_name) at::_ops::op_name::call + +// Separately, ATEN_OP(op) and ATEN_OP2(op, overload) define a class containing compile-time +// metadata about a given aten operator. +// Notable data on the class includes: +// - ATEN_OP2(add, Tensor)::name // returns the string name: "add" +// - ATEN_OP2(add, Tensor)::overload_name // returns the string overload name: "Tensor" +// - ATEN_OP2(add, Tensor)::schema // returns the C++ schema type: at::Tensor (const at::Tensor &, const at::Tensor &, const at::Scalar &) +// - ATEN_OP2(add, Tensor)::schema_str // returns the string jit type: "add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor" + +#define ATEN_OP2(op_name, overload) at::_ops::op_name##_##overload +#define ATEN_OP(op_name) at::_ops::op_name + +// WARNING: Please do not call any of the ops in the _ops namespace directly. +// Use the ATEN_FN macros. We do not guarantee stability of the naming +// scheme for the functions in at::_ops + +// See Note [The ATen Operators API] for details of the at::_ops namespace + +namespace at { +namespace _ops { +${Operators_declarations} +} // namespace _ops +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RedispatchFunctions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RedispatchFunctions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..58102bd97fca4eaef477818b0b0a92b7995e38b1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RedispatchFunctions.cpp @@ -0,0 +1,15 @@ +// ${generated_comment} + +#include +#include + +#include +#include + +namespace at { + +namespace redispatch { + ${function_redispatch_definitions} +} // namespace redispatch + +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RedispatchFunctions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RedispatchFunctions.h new file mode 100644 index 0000000000000000000000000000000000000000..2422cdd409cfdd59c2a05df27d28bb25ee610463 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RedispatchFunctions.h @@ -0,0 +1,32 @@ +#pragma once + +// ${generated_comment} + +#ifdef TORCH_ASSERT_ONLY_METHOD_OPERATORS +#error This change adds a dependency on all pytorch operators, meaning the \ + file will need to be re-compiled every time an operator is changed or added. \ + Consider using the at::_ops::{name}::redispatch() interface by including \ + the specific operator from +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace at { + +namespace redispatch { + ${function_redispatch_definitions} +} // namespace redispatch + +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterBackendSelect.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterBackendSelect.cpp new file mode 100644 index 0000000000000000000000000000000000000000..018cf358f11237d5bdc9bca01aa8d09d1462f574 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterBackendSelect.cpp @@ -0,0 +1,29 @@ +// We register ops with a higher priority dispatch key (BackendSelect) than the usual backend-specific keys (e.g. CPU) +// which makes calls to the factory functions dispatch to here. +// We then 'manually' compute a lower-priority to re-dispatch to (e.g. CPU) to get to the eventually correct backend. +// ${generated_comment} + +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +#include +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else + +${ops_headers} +#endif + +namespace at { + +namespace { + +${backend_select_method_definitions} + +TORCH_LIBRARY_IMPL(aten, BackendSelect, m) { + ${backend_select_function_registrations}; +} + +} // namespace +} // at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterCodegenUnboxedKernels.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterCodegenUnboxedKernels.cpp new file mode 100644 index 0000000000000000000000000000000000000000..279f987c66a26c2eb5d11c664c85b3604b67684b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterCodegenUnboxedKernels.cpp @@ -0,0 +1,41 @@ +#include +#include +#include + +#include + +// ${generated_comment} + +// NOTE [Sharded File]: This file is generated in a sharded fashion to speed up +// incremental rebuilds. See the comment at the top of +// templates/VariableType.cpp for an analogous, in-depth discussion. +// +// Generated by tools/jit/gen_unboxing.py. This file registers all ATen ops into JIT op registry instead of c10 +// dispatcher. JIT op registry only takes boxed kernels, so we are calling unboxing functions in UnboxingFunctions.h +// to cast arguments into C++ types (instead of IValue) and delegate to unboxed kernels. + +namespace torch { namespace jit { + +using autograd::Variable; +using autograd::variable_list; +using at::Scalar; +using at::ScalarType; +using at::Tensor; +using at::TensorOptions; +using at::DeviceGuard; + +using ::c10::fmap; +using ::c10::filter; + +namespace { + +RegisterOperators reg({ + + // Generated operators + ${unboxed_ops} +}); + +} // anon namespace + + +}} // namespace torch::jit diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterDispatchDefinitions.ini b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterDispatchDefinitions.ini new file mode 100644 index 0000000000000000000000000000000000000000..97c921de18f62832d1ca09c245f2466541fe908d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterDispatchDefinitions.ini @@ -0,0 +1,22 @@ +${ns_prologue} + +// NB: TORCH_LIBRARY_IMPL must be in an anonymous namespace to avoid +// ambiguity with conflicting identifiers that may have been defined in +// at namespace already. +namespace { + +${dispatch_anonymous_definitions} + +${static_init_dispatch_registrations} + +} // anonymous namespace + +${deferred_dispatch_registrations} + +namespace ${dispatch_namespace} { + +${dispatch_namespaced_definitions} + +} // namespace ${dispatch_namespace} + +${ns_epilogue} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterDispatchKey.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterDispatchKey.cpp new file mode 100644 index 0000000000000000000000000000000000000000..bd4131b1ad41d0c8eb49b408309e00e7062f6b53 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterDispatchKey.cpp @@ -0,0 +1,51 @@ +// an external backend might generate file within its code tree +// and check all the source files within the tree with clang-format. +// so, disable it since the backend might have a different config. +// clang-format off + +// NOTE: This condition is true for all PyTorch internal libraries, it +// just excludes external projects such as torch_xla which +// reuse some of the PyTorch codegen machinery. +#if defined(CAFFE2_BUILD_MAIN_LIB) || \ + defined(TORCH_CUDA_BUILD_MAIN_LIB) || \ + defined(TORCH_XPU_BUILD_MAIN_LIB) +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +#endif + +// ${generated_comment} + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include +$extra_cuda_headers +$external_backend_headers +$dispatch_headers +$ops_headers + +namespace at { +namespace { +$dispatch_helpers +} // namespace +} // namespace at + +// See template file RegisterDispatchDefinitions.ini +$dispatch_definitions diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterFunctionalization.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterFunctionalization.cpp new file mode 100644 index 0000000000000000000000000000000000000000..408aff0cdab40461a7ba731bab216a7b7435331e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterFunctionalization.cpp @@ -0,0 +1,116 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include +#include +#include +#include +#include +#include + +#include +#ifndef AT_PER_OPERATOR_HEADERS +#include +#include +#else +// needed for the meta tensor calls to get stride info in functionalization +#include +// needed for special handling of copy_(). +// See Note [functionalizating copy_() and not preserving strides] +#include +#include + +$ops_headers +#endif + +namespace at { +namespace functionalization { + +// This keyset is used by functionalization when it calls into meta kernels +// to accurately propagate stride metadata. +// Exclude any modes: the purpose of calling into meta kernels is only as an implementation +// detail to perform shape inference, and we don't want any modal keys to run. +// Specifically, we want to prevent functionalization and Python modes from running. +constexpr auto exclude_keys_for_meta_dispatch = + c10::functorch_transforms_ks | + c10::DispatchKeySet({ + c10::DispatchKey::FuncTorchDynamicLayerBackMode, + c10::DispatchKey::FuncTorchDynamicLayerFrontMode, + c10::DispatchKey::Python, + c10::DispatchKey::PreDispatch, + + }); + +// Helper around at::has_internal_overlap. +// The ATen util is used in hot-path eager mode: it's always fast, +// but might return TOO_HARD sometimes. +// During functionalization, we're ok taking a bit longer +// to detect memory overlap. +inline bool has_internal_overlap_helper(const at::Tensor t) { + auto has_overlap = at::has_internal_overlap(t); + if (has_overlap == at::MemOverlap::Yes) return true; + if (has_overlap == at::MemOverlap::No) return false; + return false; +} + + +inline Tensor to_meta(const Tensor& t) { + if (!t.defined()) return t; + return at::native::empty_strided_meta_symint(t.sym_sizes(), t.sym_strides(), +/*dtype=*/t.scalar_type(), /*layout=*/t.layout(), +/*device=*/c10::Device(kMeta), /*pin_memory=*/std::nullopt); +} + +inline std::optional to_meta(const std::optional& t) { + if (t.has_value()) { + return to_meta(*t); + } + return std::nullopt; +} + +inline std::vector to_meta(at::ITensorListRef t_list) { + std::vector outputs; + outputs.reserve(t_list.size()); + for (const auto& tensor : t_list) { + outputs.push_back(to_meta(tensor)); + } + return outputs; +} + +inline c10::List to_meta(const c10::List& t_list) { + c10::List outputs; + outputs.reserve(t_list.size()); + for (const auto i : c10::irange(t_list.size())) { + outputs.push_back(to_meta(t_list[i])); + } + return outputs; +} + +inline c10::List<::std::optional> to_meta(const c10::List<::std::optional>& t_list) { + c10::List<::std::optional> outputs; + outputs.reserve(t_list.size()); + for (const auto i : c10::irange(t_list.size())) { + outputs.push_back(to_meta(t_list[i])); + } + return outputs; +} + +static bool disable_meta_reference() { + static auto env = c10::utils::get_env("TORCH_DISABLE_FUNCTIONALIZATION_META_REFERENCE"); + return env == "1"; +} + + +${func_definitions} + +} // namespace functionalization + +namespace { + +TORCH_LIBRARY_IMPL(aten, Functionalize, m) { + ${func_registrations}; +} + +} // namespace + +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterSchema.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterSchema.cpp new file mode 100644 index 0000000000000000000000000000000000000000..029796d3e575b2bde85cfd44af9e6fcbb56466cd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegisterSchema.cpp @@ -0,0 +1,13 @@ +// ${generated_comment} +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +#include + +namespace at { +TORCH_LIBRARY(aten, m) { + ${aten_schema_registrations}; + // Distributed Ops + // Implementations located in torch/csrc/jit/runtime/register_distributed_ops.cpp + m.def("get_gradients(int context_id) -> Dict(Tensor, Tensor)"); +} +${schema_registrations} +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegistrationDeclarations.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegistrationDeclarations.h new file mode 100644 index 0000000000000000000000000000000000000000..5a0f0d0c7b44dabb60061d32ced243fe607069d8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/RegistrationDeclarations.h @@ -0,0 +1,4 @@ +// This file contains all native_functions that can be registered to +// and the schema string that they should be registered with + +${registration_declarations} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/TensorBody.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/TensorBody.h new file mode 100644 index 0000000000000000000000000000000000000000..5ca62cc1f7a5171ab204a36db1959dc4f14f4892 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/TensorBody.h @@ -0,0 +1,758 @@ +#pragma once + +#ifdef TORCH_ASSERT_NO_OPERATORS +#error This change adds a dependency on native_functions.yaml, \ + meaning the file will need to be re-compiled every time an operator \ + is changed or added. Consider if your change would be better placed in \ + another file, or if a more specific header might achieve the same goal. \ + See NOTE: [Tensor vs. TensorBase] +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + + +#include + +namespace c10{ +template class List; +template class IListRef; +} +namespace at { +struct Generator; +struct Type; +class DeprecatedTypeProperties; +class Tensor; +} // namespace at +namespace at { +namespace indexing { +struct TensorIndex; +} // namespace indexing +} // namespace at + +namespace torch { namespace autograd { + +struct Node; + +}} // namespace torch::autograd + +namespace at { + +class OptionalTensorRef; +class TensorRef; +class Tensor; +using TensorList = ArrayRef; +using ITensorList = c10::IListRef; + +using Stream = c10::Stream; + +// Tensor is a "generic" object holding a pointer to the underlying TensorImpl object, which +// has an embedded reference count. In this way, Tensor is similar to boost::intrusive_ptr. +// +// For example: +// +// void func(Tensor a) { +// Tensor b = a; +// ... +// } +// +// In this example, when we say Tensor b = a, we are creating a new object that points to the +// same underlying TensorImpl, and bumps its reference count. When b goes out of scope, the +// destructor decrements the reference count by calling release() on the TensorImpl it points to. +// The existing constructors, operator overloads, etc. take care to implement the correct semantics. +// +// Note that Tensor can also be NULL, i.e. it is not associated with any underlying TensorImpl, and +// special care must be taken to handle this. +class TORCH_API Tensor: public TensorBase { + protected: + // Create a Tensor with a +0 reference count. Special care must be + // taken to avoid decrementing this reference count at destruction + // time. Intended to support MaybeOwnedTraits. + explicit Tensor(unsafe_borrow_t, const TensorBase& rhs): TensorBase(unsafe_borrow_t{}, rhs) {} + friend MaybeOwnedTraits; + friend OptionalTensorRef; + friend TensorRef; + + public: + Tensor() = default; + // This constructor should not be used by end users and is an implementation + // detail invoked by autogenerated code. + explicit Tensor( + c10::intrusive_ptr tensor_impl) + : TensorBase(std::move(tensor_impl)) {} + Tensor(const Tensor &tensor) = default; + Tensor(Tensor &&tensor) = default; + + // Implicitly move-constructible from TensorBase, but must be explicit to increase refcount + explicit Tensor(const TensorBase &base): TensorBase(base) {} + /*implicit*/ Tensor(TensorBase &&base): TensorBase(std::move(base)) {} + + // Creates a new wrapper from TensorImpl. Intentionally a free method because + // it should be used with care. Checks necessary invariants + static Tensor wrap_tensor_impl( + c10::intrusive_ptr tensor_impl) { + return TensorBase::wrap_tensor_impl(std::move(tensor_impl)); + } + + Tensor contiguous(MemoryFormat memory_format=MemoryFormat::Contiguous) const { + return TensorBase::contiguous(memory_format); + } + + Tensor conj() const { + if (!this->is_complex()) { + return *this; + } + + C10_DIAGNOSTIC_PUSH_AND_IGNORED_IF_DEFINED("-Wswitch-enum") + switch (this->layout()) { + case at::kSparse: + case at::kSparseCsr: + case at::kSparseCsc: + case at::kSparseBsr: + case at::kSparseBsc: + return this->conj_physical(); + default: + return this->_conj(); + } + C10_DIAGNOSTIC_POP() + } + + // Aliased by Dimname overloads, so need explicit using + using TensorBase::size; + using TensorBase::sym_size; + using TensorBase::stride; + + /// Should be used if *this can reasonably be expected to be contiguous and + /// performance is important. + /// Compared to contiguous, it saves a reference count + /// increment/decrement if *this is already contiguous, at the cost + /// in all cases of an extra pointer of stack usage, an extra branch + /// to access, and an extra branch at destruction time. + c10::MaybeOwned expect_contiguous(MemoryFormat memory_format=MemoryFormat::Contiguous) const &; + + // Use .contiguous() instead. Trying to borrow from a prvalue Tensor + // will only lead to trouble and dangling references. + c10::MaybeOwned expect_contiguous(MemoryFormat memory_format=MemoryFormat::Contiguous) && = delete; + + // The following overloads are very intriguing. Consider the following + // program: + // + // x[1] = 3; + // + // We would expect that the first entry of x is written to 3. But how can we + // actually achieve this? x[1] evaluates to a tensor... + // + // The answer is, using a ref-qualifier. x[1] is an rvalue, which cannot be + // (profitably) assigned to in the traditional sense, so we overload + // assignment to mean, "Actually, copy 3 into the tensor data." This is done + // with an rvalue-reference ref-qualified overload (the methods with && at the + // end of their type.) + // + // There's one more fly in the ointment: We also want + // + // Tensor x = y; + // + // to work, and we want it NOT to copy. So we need a traditional operator= + // overload. But we MUST specify a mutable lvalue ref-qualifier, to + // disambiguate the traditional overload from the rvalue-reference + // ref-qualified overload. Otherwise, it will be ambiguous, because + // a non ref-qualified method is eligible for all situations. + + // Unfortunately, we have to write these constructors out manually + // to work around an MSVC bug: + // error C2580: 'at::Tensor &at::Tensor::operator =(const at::Tensor &) &': + // multiple versions of a defaulted special member functions are not allowed + // Tensor& operator=(const Tensor&) & = default; + // Tensor& operator=(Tensor&&) & = default; + + // Also MSVC will wrongly issue the following warning with the aforementioned fix + // warning C4522: 'at::Tensor': multiple assignment operators specified + // Let's just skip the warning. + // + // TODO: temporarily disabled + + Tensor& operator=(const TensorBase& x) & noexcept { + impl_ = x.getIntrusivePtr(); + return *this; + } + Tensor& operator=(TensorBase&& x) & noexcept { + impl_ = x.unsafeReleaseIntrusivePtr(); + return *this; + } + + Tensor& operator=(const Tensor &x) & noexcept { + return operator=(static_cast(x)); + } + Tensor& operator=(Tensor &&x) & noexcept { + return operator=(static_cast(x)); + } + + Tensor& operator=(const Scalar &v) && { + return fill_(v); + } + Tensor& operator=(const Tensor &rhs) && { + return copy_(rhs); + } + + // NOLINTNEXTLINE(performance-noexcept-move-constructor) + Tensor& operator=(Tensor&& rhs) && { + return copy_(rhs); + } + + C10_DEPRECATED_MESSAGE("Tensor.type() is deprecated. Instead use Tensor.options(), which in many cases (e.g. in a constructor) is a drop-in replacement. If you were using data from type(), that is now available from Tensor itself, so instead of tensor.type().scalar_type(), use tensor.scalar_type() instead and instead of tensor.type().backend() use tensor.device().") + DeprecatedTypeProperties & type() const { + return globalDeprecatedTypePropertiesRegistry().getDeprecatedTypeProperties( + dispatchKeyToBackend(legacyExtractDispatchKey(key_set())), + scalar_type()); + } + + Tensor toType(ScalarType t) const { + return to(options().dtype(t), /*non_blocking*/ false, /*copy*/ false); + } + + // TODO: Deprecate me + Tensor toBackend(Backend b) const { + return to(options().device(backendToDeviceType(b)).layout(layout_from_backend(b)), /*non_blocking*/ false, /*copy*/ false); + } + + C10_DEPRECATED_MESSAGE("Tensor.is_variable() is deprecated; everything is a variable now. (If you want to assert that variable has been appropriately handled already, use at::impl::variable_excluded_from_dispatch())") + bool is_variable() const noexcept { + return !at::impl::variable_excluded_from_dispatch(); + } + + template + C10_DEPRECATED_MESSAGE("Tensor.data() is deprecated. Please use Tensor.data_ptr() instead.") + T * data() const { + return data_ptr(); + } + + template + T item() const; + + template class PtrTraits = DefaultPtrTraits, typename index_t = int64_t> + C10_DEPRECATED_MESSAGE("packed_accessor is deprecated, use packed_accessor32 or packed_accessor64 instead") + GenericPackedTensorAccessor packed_accessor() const & { + return generic_packed_accessor(); + } + template class PtrTraits = DefaultPtrTraits, typename index_t = int64_t> + C10_DEPRECATED_MESSAGE("packed_accessor is deprecated, use packed_accessor32 or packed_accessor64 instead") + GenericPackedTensorAccessor packed_accessor() && = delete; + + Tensor operator~() const { + return bitwise_not(); + } + Tensor operator-() const { + return neg(); + } + Tensor& operator+=(const Tensor & other) { + return add_(other); + } + Tensor& operator+=(const Scalar & other) { + return add_(other); + } + Tensor& operator-=(const Tensor & other) { + return sub_(other); + } + Tensor& operator-=(const Scalar & other) { + return sub_(other); + } + Tensor& operator*=(const Tensor & other) { + return mul_(other); + } + Tensor& operator*=(const Scalar & other) { + return mul_(other); + } + Tensor& operator/=(const Tensor & other) { + return div_(other); + } + Tensor& operator/=(const Scalar & other) { + return div_(other); + } + Tensor& operator&=(const Tensor & other) { + return bitwise_and_(other); + } + Tensor& operator|=(const Tensor & other) { + return bitwise_or_(other); + } + Tensor& operator^=(const Tensor & other) { + return bitwise_xor_(other); + } + Tensor operator[](const Scalar & index) const { + if (!index.isIntegral(false)) { + TORCH_CHECK_INDEX(false, "Can only index tensors with integral scalars"); + } + return this->operator[](index.toLong()); + } + Tensor operator[](const Tensor & index) const { + // These properties are checked in the Scalar constructor, but we already + // check them here to provide more useful diagnostics for the user. + if (!index.defined()) { + TORCH_CHECK_INDEX(false, "Can only index with tensors that are defined"); + } + if (index.dim() != 0) { + TORCH_CHECK_INDEX(false, + "Can only index with tensors that are scalars (zero-dim)"); + } + // The Scalar(Tensor) constructor is explicit, so we need to call it. + return this->operator[](index.item()); + } + Tensor operator[](int64_t index) const { + return select(0, index); + } + + Tensor index(ArrayRef indices) const; + Tensor index(std::initializer_list indices) const; + + Tensor & index_put_(ArrayRef indices, Tensor const & rhs); + Tensor & index_put_(ArrayRef indices, const Scalar& v); + Tensor & index_put_(std::initializer_list indices, Tensor const & rhs); + Tensor & index_put_(std::initializer_list indices, const Scalar& v); + + Tensor cpu() const { + return to(options().device(c10::DeviceType::CPU), /*non_blocking*/ false, /*copy*/ false); + } + + // TODO: The Python version also accepts arguments + Tensor cuda() const { + return to(options().device(c10::DeviceType::CUDA), /*non_blocking*/ false, /*copy*/ false); + } + + Tensor hip() const { + return to(options().device(c10::DeviceType::HIP), /*non_blocking*/ false, /*copy*/ false); + } + + Tensor ve() const { + return to(options().device(c10::DeviceType::VE), /*non_blocking*/ false, /*copy*/ false); + } + + Tensor vulkan() const { + return to(options().device(c10::DeviceType::Vulkan), /*non_blocking*/ false, /*copy*/ false); + } + + Tensor metal() const { + return to(options().device(c10::DeviceType::Metal), /*non_blocking*/ false, /*copy*/ false); + } + + Tensor meta() const { + return to(options().device(c10::DeviceType::Meta), /*non_blocking*/ false, /*copy*/ false); + } + + // ~~~~~ Autograd API ~~~~~ + + /// \fn bool is_leaf() const; + /// + /// All Tensors that have `requires_grad()` which is ``false`` will be leaf Tensors by convention. + /// + /// For Tensors that have `requires_grad()` which is ``true``, they will be leaf Tensors if they were + /// created by the user. This means that they are not the result of an operation and so + /// `grad_fn()` is `nullptr`. + /// + /// Only leaf Tensors will have their `grad()` populated during a call to `backward()`. + /// To get `grad()` populated for non-leaf Tensors, you can use `retain_grad()`. + /// + /// Example: + /// @code + /// auto a = torch::rand(10, torch::requires_grad()); + /// std::cout << a.is_leaf() << std::endl; // prints `true` + /// + /// auto b = torch::rand(10, torch::requires_grad()).to(torch::kCUDA); + /// std::cout << b.is_leaf() << std::endl; // prints `false` + /// // b was created by the operation that cast a cpu Tensor into a cuda Tensor + /// + /// auto c = torch::rand(10, torch::requires_grad()) + 2; + /// std::cout << c.is_leaf() << std::endl; // prints `false` + /// // c was created by the addition operation + /// + /// auto d = torch::rand(10).cuda(); + /// std::cout << d.is_leaf() << std::endl; // prints `true` + /// // d does not require gradients and so has no operation creating it (that is tracked by the autograd engine) + /// + /// auto e = torch::rand(10).cuda().requires_grad_(); + /// std::cout << e.is_leaf() << std::endl; // prints `true` + /// // e requires gradients and has no operations creating it + /// + /// auto f = torch::rand(10, torch::device(torch::kCUDA).requires_grad(true)); + /// std::cout << f.is_leaf() << std::endl; // prints `true` + /// // f requires grad, has no operation creating it + /// @endcode + + /// \fn void backward(const Tensor & gradient={}, std::optional retain_graph=std::nullopt, bool create_graph=false, std::optional inputs=std::nullopt) const; + /// + /// Computes the gradient of current tensor with respect to graph leaves. + /// + /// The graph is differentiated using the chain rule. If the tensor is + /// non-scalar (i.e. its data has more than one element) and requires + /// gradient, the function additionally requires specifying ``gradient``. + /// It should be a tensor of matching type and location, that contains + /// the gradient of the differentiated function w.r.t. this Tensor. + /// + /// This function accumulates gradients in the leaves - you might need to + /// zero them before calling it. + /// + /// \param gradient Gradient w.r.t. the + /// tensor. If it is a tensor, it will be automatically converted + /// to a Tensor that does not require grad unless ``create_graph`` is True. + /// None values can be specified for scalar Tensors or ones that + /// don't require grad. If a None value would be acceptable then + /// this argument is optional. + /// \param retain_graph If ``false``, the graph used to compute + /// the grads will be freed. Note that in nearly all cases setting + /// this option to True is not needed and often can be worked around + /// in a much more efficient way. Defaults to the value of + /// ``create_graph``. + /// \param create_graph If ``true``, graph of the derivative will + /// be constructed, allowing to compute higher order derivative + /// products. Defaults to ``false``. + /// \param inputs Inputs w.r.t. which the gradient will be accumulated into + /// ``at::Tensor::grad``. All other Tensors will be ignored. If not + /// provided, the gradient is accumulated into all the leaf Tensors + /// that were used to compute the current tensor. + /// When inputs are provided and a given input is not a leaf, + /// the current implementation will call its grad_fn (even though it is not strictly needed to get this gradients). + /// It is an implementation detail on which the user should not rely. + /// See https://github.com/pytorch/pytorch/pull/60521#issuecomment-867061780 for more details. + void backward(const Tensor & gradient={}, std::optional retain_graph=std::nullopt, bool create_graph=false, std::optional inputs=std::nullopt) const { + // NB: Adding this wrapper to _backward here because we'd like our + // 'backwards' api to accept the 'inputs' argument optionally. Since code gen + // currently does not support optional of TensorList our approach is to replace + // backward in native_functions.yaml with _backward and call it here instead. + if (inputs.has_value()) { + TORCH_CHECK(inputs.value().size() > 0, "'inputs' argument to backward cannot be empty") + this->_backward(inputs.value(), gradient, retain_graph, create_graph); + } else { + this->_backward({}, gradient, retain_graph, create_graph); + } + } + + /// \fn Tensor detach() const; + /// + /// Returns a new Tensor, detached from the current graph. + /// The result will never require gradient. + + /// \fn Tensor & detach_() const; + /// + /// Detaches the Tensor from the graph that created it, making it a leaf. + /// Views cannot be detached in-place. + + /// \fn void retain_grad() const; + /// + /// Enables this Tensor to have their :attr:`grad` populated during + /// :func:`backward`. This is a no-op for leaf tensors. + + /// \fn bool retains_grad() const; + /// + /// Is ``true`` if this Tensor is non-leaf and its :attr:`grad` is enabled to be + /// populated during :func:`backward`, ``false`` otherwise. + + const Tensor& set_requires_grad(bool requires_grad) const { + TensorBase::set_requires_grad(requires_grad); + return *this; + } + + /// Return a mutable reference to the gradient. This is conventionally + /// used as `t.grad() = x` to set a gradient to a completely new tensor. + /// Note that this function work with a non-const Tensor and is not + /// thread safe. + Tensor& mutable_grad() const { + return impl_->mutable_grad(); + } + + /// This function returns an undefined tensor by default and returns a defined tensor + /// the first time a call to `backward()` computes gradients for this Tensor. + /// The attribute will then contain the gradients computed and future calls + /// to `backward()` will accumulate (add) gradients into it. + const Tensor& grad() const { + const Tensor& maybe_grad = impl_->grad(); + if (!is_leaf() && !retains_grad() && !maybe_grad.defined()) { + TORCH_WARN( + "The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad " + "attribute won't be populated during autograd.backward(). If you indeed want the .grad " + "field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. " + "If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor " + "instead. See github.com/pytorch/pytorch/pull/30531 for more information."); + } + return maybe_grad; + } + + // The Forward AD API functions below are low level and are not to be used by end + // users who should use the API provided in torch/csrc/autograd.h + + /// This function returns the forward gradient for this Tensor at the given level. + const Tensor& _fw_grad(uint64_t level) const { + return impl_->_fw_grad(level, *this); + } + + /// This function can be used to set the value of the forward grad. + /// Note that the given new_grad might not be used directly if it has different + /// metadata (size/stride/storage offset) compared to this Tensor. In that case, + /// new_grad content will be copied into a new Tensor + void _set_fw_grad(const TensorBase& new_grad, uint64_t level, bool is_inplace_op) const { + impl_->_set_fw_grad(new_grad, *this, level, is_inplace_op); + } + + + // STOP. Thinking of adding a method here, which only makes use + // of other ATen methods? Define it in native_functions.yaml. + + //example + //Tensor * add(Tensor & b); + ${tensor_method_declarations} + + // Special C++ only overloads for std()-like functions (See gh-40287) + // These are needed because int -> bool conversion takes precedence over int -> IntArrayRef + // So, for example std(0) would select the std(unbiased=False) overload + + Tensor var(int dim) const { + return var(IntArrayRef{dim}); + } + + Tensor std(int dim) const { + return std(IntArrayRef{dim}); + } + + // We changed .dtype() to return a TypeMeta in #12766. Ideally, we want the + // at::kDouble and its friends to be TypeMeta's, but that hasn't happened yet. + // Before that change, we make this method to maintain BC for C++ usage like + // `x.to(y.dtype)`. + // TODO: remove following two after at::kDouble and its friends are TypeMeta's. + inline Tensor to(caffe2::TypeMeta type_meta, bool non_blocking=false, bool copy=false) const { + return this->to(/*scalar_type=*/typeMetaToScalarType(type_meta), non_blocking, copy); + } + inline Tensor to(Device device, caffe2::TypeMeta type_meta, bool non_blocking=false, bool copy=false) const { + return this->to(device, /*scalar_type=*/typeMetaToScalarType(type_meta), non_blocking, copy); + } + + template + decltype(auto) m(F func, Args&&... params) const { + return func(*this, std::forward(params)...); + } + + /// NOTE: This is similar to the legacy `.data()` function on `Variable`, and is intended + /// to be used from functions that need to access the `Variable`'s equivalent `Tensor` + /// (i.e. `Tensor` that shares the same storage and tensor metadata with the `Variable`). + /// + /// One notable difference with the legacy `.data()` function is that changes to the + /// returned `Tensor`'s tensor metadata (e.g. sizes / strides / storage / storage_offset) + /// will not update the original `Variable`, due to the fact that this function + /// shallow-copies the `Variable`'s underlying TensorImpl. + at::Tensor tensor_data() const { + return TensorBase::tensor_data(); + } + + /// NOTE: `var.variable_data()` in C++ has the same semantics as `tensor.data` + /// in Python, which create a new `Variable` that shares the same storage and + /// tensor metadata with the original `Variable`, but with a completely new + /// autograd history. + /// + /// NOTE: If we change the tensor metadata (e.g. sizes / strides / + /// storage / storage_offset) of a variable created from `var.variable_data()`, those + /// changes will not update the original variable `var`. In `.variable_data()`, we set + /// `allow_tensor_metadata_change_` to false to make such changes explicitly illegal, + /// in order to prevent users from changing metadata of `var.variable_data()` + /// and expecting the original variable `var` to also be updated. + at::Tensor variable_data() const { + return TensorBase::variable_data(); + } + + // Hooks + //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + template + using hook_return_void_t = std::enable_if_t>::value, unsigned>; + template + using hook_return_var_t = std::enable_if_t, Tensor>, unsigned>; + + /// Registers a backward hook. + /// + /// The hook will be called every time a gradient with respect to the Tensor is computed. + /// The hook should have one of the following signature: + /// ``` + /// hook(Tensor grad) -> Tensor + /// ``` + /// ``` + /// hook(Tensor grad) -> void + /// ``` + /// The hook should not modify its argument, but it can optionally return a new gradient + /// which will be used in place of `grad`. + /// + /// This function returns the index of the hook in the list which can be used to remove hook. + /// + /// Example: + /// @code + /// auto v = torch::tensor({0., 0., 0.}, torch::requires_grad()); + /// auto h = v.register_hook([](torch::Tensor grad){ return grad * 2; }); // double the gradient + /// v.backward(torch::tensor({1., 2., 3.})); + /// // This prints: + /// // ``` + /// // 2 + /// // 4 + /// // 6 + /// // [ CPUFloatType{3} ] + /// // ``` + /// std::cout << v.grad() << std::endl; + /// v.remove_hook(h); // removes the hook + /// @endcode + template + hook_return_void_t register_hook(T&& hook) const; + template + hook_return_var_t register_hook(T&& hook) const; + + // Variable methods + //~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + + Tensor data() const { + return TensorBase::data(); + } + + void _backward(TensorList inputs, const std::optional& gradient, std::optional keep_graph, bool create_graph) const; + + const Tensor& requires_grad_(bool _requires_grad=true) const { + TensorBase::requires_grad_(_requires_grad); + return *this; + } +}; + +namespace detail { +// Helper creator for Tensor class which doesn't requires the users to pass +// in an intrusive_ptr instead it just converts the argument passed to +// requested intrusive_ptr type. +template +Tensor make_tensor(Args&&... args) { + return Tensor(c10::make_intrusive(std::forward(args)...)); +} + +} // namespace detail + +} // namespace at + + +namespace at { +${tensor_method_definitions} +} // namespace at + + +namespace c10 { +template <> +struct MaybeOwnedTraits { + using owned_type = at::Tensor; + using borrow_type = at::Tensor; + + static borrow_type createBorrow(const owned_type& from) { + // NOTE: this can be implemented without the special + // unsafe_borrow_t Tensor constructor as + // + // return borrow_type(c10::intrusive_ptr::reclaim(from.unsafeGetTensorImpl())); + // + // but that hurts inlining due to the nullptr check in the + // Tensor(c10::intrusive_ptr<...>) constructor. We already know + // that from.impl_ isn't null because from is a valid Tensor, so + // we needn't do the check again. (using __builtin_assume can + // avoid this, but wouldn't be portable to MSVC.) + return borrow_type(borrow_type::unsafe_borrow_t{}, from); + } + + static void assignBorrow(borrow_type& lhs, const borrow_type& rhs) { + lhs.unsafeReleaseTensorImpl(); + // See above note: this can be implemented with public API + // similarly to createBorrow(), but that would hurt inlining. + lhs = borrow_type(borrow_type::unsafe_borrow_t{}, rhs); + } + + static void destroyBorrow(borrow_type& toDestroy) { + toDestroy.unsafeReleaseTensorImpl(); // "leak" it, but it was already +0. + } + + static const owned_type& referenceFromBorrow(const borrow_type& borrow) { + return borrow; + } + + static const owned_type* pointerFromBorrow(const borrow_type& borrow) { + return &borrow; + } + + static bool debugBorrowIsValid(const borrow_type& /*borrow*/) { + return true; + } +}; + +template <> +struct ExclusivelyOwnedTraits { + using repr_type = at::Tensor; + using pointer_type = at::Tensor*; + using const_pointer_type = const at::Tensor*; + + static repr_type nullRepr() { + return at::Tensor(); + } + + template + static repr_type createInPlace(Args&&... args) { + return at::Tensor(std::forward(args)...); + } + + static repr_type moveToRepr(at::Tensor&& x) { + return std::move(x); + } + + static void destroyOwned(at::Tensor& x) { + return ExclusivelyOwnedTraits::destroyOwned(x); + } + + static at::Tensor take(at::Tensor& x) { + return std::move(x); + } + + static pointer_type getImpl(repr_type& x) { + return &x; + } + + static const_pointer_type getImpl(const repr_type& x) { + return &x; + } +}; +} // namespace c10 + +namespace at { + +inline c10::MaybeOwned borrow_from_optional_tensor( + const std::optional& opt) { + return opt.has_value() + ? c10::MaybeOwned::borrowed(*opt) + : c10::MaybeOwned::owned(std::in_place); +} + +inline c10::MaybeOwned Tensor::expect_contiguous(MemoryFormat memory_format) const & { + if (is_contiguous(memory_format)) { + return c10::MaybeOwned::borrowed(*this); + } else { + return c10::MaybeOwned::owned(__dispatch_contiguous(memory_format)); + } +} +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/TensorMethods.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/TensorMethods.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c1440abc4cb0e547112393b2bdd1812feb84d0c8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/TensorMethods.cpp @@ -0,0 +1,61 @@ +#include +#include + +#include + +namespace at { + +namespace { + +// Verifies the requested type is the same as the Tensor's type. +void check_type(const TensorBase& tensor, ScalarType type) { + TORCH_CHECK( + tensor.scalar_type() == type + || (isQIntType(tensor.scalar_type()) + && toUnderlying(tensor.scalar_type()) == type), + "expected scalar type ", type, " but found ", tensor.scalar_type()); +} + +} // namespace + +template +const T* TensorBase::const_data_ptr() const { + using NonConstT = std::remove_const_t; + check_type(*this, c10::CppTypeToScalarType()); + return this->unsafeGetTensorImpl()->data_ptr_impl(); +} + +template +T* TensorBase::mutable_data_ptr() const { + check_type(*this, c10::CppTypeToScalarType()); + return this->unsafeGetTensorImpl()->mutable_data_ptr_impl(); +} + +template +T* TensorBase::data_ptr() const { + return this->mutable_data_ptr(); +} + +#define DEFINE_CAST(T, name) \ + template TORCH_API const T* TensorBase::const_data_ptr() const; \ + template TORCH_API const T* TensorBase::const_data_ptr() const; \ + template TORCH_API T* TensorBase::mutable_data_ptr() const; \ + template TORCH_API T* TensorBase::data_ptr() const; + + AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_CAST) + AT_FORALL_QINT_TYPES(DEFINE_CAST) + DEFINE_CAST(uint16_t, UInt16) + DEFINE_CAST(uint32_t, UInt32) + DEFINE_CAST(uint64_t, UInt64) + #undef DEFINE_CAST + + #define DEFINE_ITEM(T, name) \ + template <> \ + TORCH_API T Tensor::item() const { \ + return item().to##name(); \ + } + + AT_FORALL_SCALAR_TYPES_WITH_COMPLEX(DEFINE_ITEM) + #undef DEFINE_ITEM + + } //namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UfuncCPU.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UfuncCPU.cpp new file mode 100644 index 0000000000000000000000000000000000000000..6b363a508907cc064e41794720657541fc28c301 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UfuncCPU.cpp @@ -0,0 +1,19 @@ +#define TORCH_ASSERT_NO_OPERATORS + +#include +#include +#include + +namespace at { + +// NB: this is explicitly copied here (via codegen) rather than +// included via NativeFunctions.h to avoid recompiling this file when +// NativeFunctions.h changes +namespace meta { +${meta_declaration} +} + +namespace native { +${native_declaration} +${native_definitions} +}} // namespace at::native diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UfuncCPUKernel.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UfuncCPUKernel.cpp new file mode 100644 index 0000000000000000000000000000000000000000..0cac55664d6125287bdee0bd94c150462b81d5b9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UfuncCPUKernel.cpp @@ -0,0 +1,14 @@ +#define TORCH_ASSERT_NO_OPERATORS + +#include +#include +#include +#include +#include +#include +#include + +namespace at { +namespace native { +${native_definitions} +}} // namespace at::native diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UfuncCUDA.cu b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UfuncCUDA.cu new file mode 100644 index 0000000000000000000000000000000000000000..e75d82d9cc84bd8fddfd303f610412e5d0a98729 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UfuncCUDA.cu @@ -0,0 +1,21 @@ +#define TORCH_ASSERT_NO_OPERATORS + +#include +#include +#include +#include +${cuda_headers} + +namespace at { + +// NB: this is explicitly copied here (via codegen) rather than +// included via NativeFunctions.h to avoid recompiling this file when +// NativeFunctions.h changes +namespace meta { +${meta_declaration} +} + +namespace native { +${native_declaration} +${native_definitions} +}} // namespace at::native diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UnboxingFunctions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UnboxingFunctions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..86c13235d8623964d734e743f5f15cf68a8df63c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UnboxingFunctions.cpp @@ -0,0 +1,35 @@ +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +namespace at { +namespace unboxing { + +using ::c10::fmap; +using ::c10::filter; +using torch::jit::peek; +using torch::jit::drop; +using torch::jit::pack; +using torch::jit::pop; + +// Generated function declaration +${definitions} + +} // namespace unboxing +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UnboxingFunctions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UnboxingFunctions.h new file mode 100644 index 0000000000000000000000000000000000000000..a65469a9b0123cbfd4075ff3c263276aa47f137f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/UnboxingFunctions.h @@ -0,0 +1,32 @@ +// ${generated_comment} + +// Generated by tools/jit/gen_unboxing.py. This file declares code generated boxed C++ functions for operators, +// base off of native_functions.yaml (or similar yaml file with the same syntax). The definition of such a boxed +// function will pop out IValues from the stack then convert them into the correct C++ types based on given schema. This +// unboxing logic is an alternative to template-based metaprogramming unboxing. + +#pragma once + +#include +namespace at { +namespace unboxing { +namespace { + +template +std::array as_array(const c10::List& list) { + std::array res; + AT_ASSERT(list.size() == N); + std::vector vec; + for (c10::IValue elem : list) { + vec.push_back(elem.to()); + } + std::copy(vec.begin(), vec.end(), res.begin()); + return res; +} +} // namespace +using Stack = std::vector; +// Generated function declaration +${declarations} + +} // namespace unboxing +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ViewMetaClasses.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ViewMetaClasses.cpp new file mode 100644 index 0000000000000000000000000000000000000000..0fd53171935f9147ba54bcd39a886e2f4dda6b2f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ViewMetaClasses.cpp @@ -0,0 +1,19 @@ +// ${generated_comment} + +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#include +#else +${op_headers} +#endif + +namespace at { +namespace functionalization { + +${view_meta_implementations} + +} // namespace functionalization +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ViewMetaClasses.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ViewMetaClasses.h new file mode 100644 index 0000000000000000000000000000000000000000..be2dee2a871b35258864377fbac83e3037108b2b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ViewMetaClasses.h @@ -0,0 +1,12 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include + +namespace at { +namespace functionalization { + +${view_meta_declarations} + +} // namespace functionalization +} // namespace at diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ViewMetaClassesPythonBinding.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ViewMetaClassesPythonBinding.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c784e5abe5c88dfb5bc418e60d48b28391274718 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/ViewMetaClassesPythonBinding.cpp @@ -0,0 +1,11 @@ +#include +#include + +namespace torch::functionalization { + +void initGenerated(PyObject* module) { + auto functionalization = py::handle(module).cast(); + $view_meta_bindings +} + +} // namespace torch::functionalization diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/aten_interned_strings.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/aten_interned_strings.h new file mode 100644 index 0000000000000000000000000000000000000000..326d4622334a776f4f1f94fb49a70f2c53c7e6eb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/aten_interned_strings.h @@ -0,0 +1,22 @@ +#pragma once + +// ${generated_comment} + +#if defined(TORCH_ASSERT_NO_OPERATORS) || defined(TORCH_ASSERT_ONLY_METHOD_OPERATORS) +#error This change adds a dependency on native_functions.yaml, \ + meaning the file will need to be re-compiled every time an operator \ + is changed or added. Consider if including for \ + the c10::Symbol class would be sufficient, or if your change would be \ + better placed in another file. +#endif + +// ATen symbols correspond exactly to operators defined in ATen. Every +// symbol here corresponds exactly to an ATen operation defined in +// native_functions.yaml; attributes are in one-to-one correspondence +// with their ATen name. + +#define FORALL_ATEN_BASE_SYMBOLS(_) \ +${aten_symbols} + +#define FORALL_ATTR_BASE_SYMBOLS(_) \ +${attr_symbols} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/enum_tag.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/enum_tag.h new file mode 100644 index 0000000000000000000000000000000000000000..49c174b721153eeba19f98601bd742a0eb5b64ab --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/ATen/templates/enum_tag.h @@ -0,0 +1,19 @@ +#pragma once + +// ${generated_comment} + +#include + +HIDDEN_NAMESPACE_BEGIN(torch, headeronly) + +// Enum of valid tags obtained from the entries in tags.yaml +enum class Tag { + ${enum_of_valid_tags} +}; + +HIDDEN_NAMESPACE_END(torch, headeronly) + +// Re-expose in the at:: namespace for backward compatibility +namespace at { + using torch::headeronly::Tag; +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/BUILD.bazel b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/BUILD.bazel new file mode 100644 index 0000000000000000000000000000000000000000..d1a0db360d230fe0f027c19869c6307f17010503 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/BUILD.bazel @@ -0,0 +1,4 @@ +load("//:tools/bazel.bzl", "rules") +load(":build.bzl", "define_targets") + +define_targets(rules = rules) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/README.md b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/README.md new file mode 100644 index 0000000000000000000000000000000000000000..bfa43899cc590959c2bfd74e38662ec03aaee3d6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/README.md @@ -0,0 +1,3 @@ +If you add a file to this directory, you **MUST** update +`torch/CMakeLists.txt` and add the file as a dependency to +the `add_custom_command` call. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/build.bzl b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/build.bzl new file mode 100644 index 0000000000000000000000000000000000000000..c5ddf7a20b800a714431fdc9feb57679783410f4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/build.bzl @@ -0,0 +1,20 @@ +def define_targets(rules): + rules.py_library( + name = "autograd", + srcs = rules.glob(["*.py"]), + data = rules.glob([ + "*.yaml", + "templates/*", + ]), + visibility = ["//:__subpackages__"], + deps = [ + rules.requirement("PyYAML"), + "//torchgen", + ], + ) + + rules.filegroup( + name = "deprecated_yaml", + srcs = ["deprecated.yaml"], + visibility = ["//:__subpackages__"], + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/context.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/context.py new file mode 100644 index 0000000000000000000000000000000000000000..0ed4b2ee4d014be3dca01c3f2293b36b03b7880b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/context.py @@ -0,0 +1,31 @@ +import functools +from collections.abc import Callable + +from torchgen.api.autograd import NativeFunctionWithDifferentiabilityInfo as NFWDI +from torchgen.context import native_function_manager +from torchgen.utils import T + + +# Like tools.api.context.with_native_function, but for +# NativeFunctionWithDifferentiabilityInfo. +def with_native_function_with_differentiability_info( + func: Callable[[NFWDI], T], +) -> Callable[[NFWDI], T]: + @functools.wraps(func) + def wrapper(f: NFWDI) -> T: + with native_function_manager(f.func): + return func(f) + + return wrapper + + +# Like the above but with an additional dispatch key string argument +def with_native_function_with_differentiability_info_and_key( + func: Callable[[NFWDI, str], T], +) -> Callable[[NFWDI, str], T]: + @functools.wraps(func) + def wrapper(f: NFWDI, key: str) -> T: + with native_function_manager(f.func): + return func(f, key) + + return wrapper diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/deprecated.yaml b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/deprecated.yaml new file mode 100644 index 0000000000000000000000000000000000000000..52f7ec50b6ea15dae1c3308358997950d295c924 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/deprecated.yaml @@ -0,0 +1,134 @@ +# Deprecated function signatures. These are exposed in Python, but not included +# in the error message suggestions. + +- name: add(Tensor self, Scalar alpha, Tensor other) -> Tensor + aten: add(self, other, alpha) + +- name: add_(Tensor(a!) self, Scalar alpha, Tensor other) -> Tensor(a!) + aten: add_(self, other, alpha) + +- name: add(Tensor self, Scalar alpha, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + aten: add_out(out, self, other, alpha) + +- name: addbmm(Scalar beta, Tensor self, Scalar alpha, Tensor batch1, Tensor batch2) -> Tensor + aten: addbmm(self, batch1, batch2, beta, alpha) + +- name: addbmm_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor batch1, Tensor batch2) -> Tensor(a!) + aten: addbmm_(self, batch1, batch2, beta, alpha) + +- name: addbmm(Scalar beta, Tensor self, Scalar alpha, Tensor batch1, Tensor batch2, *, Tensor(a!) out) -> Tensor(a!) + aten: addbmm_out(out, self, batch1, batch2, beta, alpha) + +- name: addbmm(Scalar beta, Tensor self, Tensor batch1, Tensor batch2) -> Tensor + aten: addbmm(self, batch1, batch2, beta, 1) + +- name: addbmm_(Scalar beta, Tensor(a!) self, Tensor batch1, Tensor batch2) -> Tensor(a!) + aten: addbmm_(self, batch1, batch2, beta, 1) + +- name: addbmm(Scalar beta, Tensor self, Tensor batch1, Tensor batch2, *, Tensor(a!) out) -> Tensor(a!) + aten: addbmm_out(out, self, batch1, batch2, beta, 1) + +- name: addcdiv(Tensor self, Scalar value, Tensor tensor1, Tensor tensor2) -> Tensor + aten: addcdiv(self, tensor1, tensor2, value) + +- name: addcdiv_(Tensor(a!) self, Scalar value, Tensor tensor1, Tensor tensor2) -> Tensor(a!) + aten: addcdiv_(self, tensor1, tensor2, value) + +- name: addcdiv(Tensor self, Scalar value, Tensor tensor1, Tensor tensor2, *, Tensor(a!) out) -> Tensor(a!) + aten: addcdiv_out(out, self, tensor1, tensor2, value) + +- name: addcmul(Tensor self, Scalar value, Tensor tensor1, Tensor tensor2) -> Tensor + aten: addcmul(self, tensor1, tensor2, value) + +- name: addcmul_(Tensor(a!) self, Scalar value, Tensor tensor1, Tensor tensor2) -> Tensor(a!) + aten: addcmul_(self, tensor1, tensor2, value) + +- name: addcmul(Tensor self, Scalar value, Tensor tensor1, Tensor tensor2, *, Tensor(a!) out) -> Tensor(a!) + aten: addcmul_out(out, self, tensor1, tensor2, value) + +- name: addmm(Scalar beta, Tensor self, Scalar alpha, Tensor mat1, Tensor mat2) -> Tensor + aten: addmm(self, mat1, mat2, beta, alpha) + +- name: addmm_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor mat1, Tensor mat2) -> Tensor(a!) + aten: addmm_(self, mat1, mat2, beta, alpha) + +- name: addmm(Scalar beta, Tensor self, Scalar alpha, Tensor mat1, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!) + aten: addmm_out(out, self, mat1, mat2, beta, alpha) + +- name: addmm(Scalar beta, Tensor self, Tensor mat1, Tensor mat2) -> Tensor + aten: addmm(self, mat1, mat2, beta, 1) + +- name: addmm_(Scalar beta, Tensor(a!) self, Tensor mat1, Tensor mat2) -> Tensor(a!) + aten: addmm_(self, mat1, mat2, beta, 1) + +- name: addmm(Scalar beta, Tensor self, Tensor mat1, Tensor mat2, *, Tensor(a!) out) -> Tensor(a!) + aten: addmm_out(out, self, mat1, mat2, beta, 1) + +- name: sspaddmm(Scalar beta, Tensor self, Scalar alpha, Tensor mat1, Tensor mat2) -> Tensor + aten: sspaddmm(self, mat1, mat2, beta, alpha) + +- name: sspaddmm(Scalar beta, Tensor self, Tensor mat1, Tensor mat2) -> Tensor + aten: sspaddmm(self, mat1, mat2, beta, 1) + +- name: addmv(Scalar beta, Tensor self, Scalar alpha, Tensor mat, Tensor vec) -> Tensor + aten: addmv(self, mat, vec, beta, alpha) + +- name: addmv_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor mat, Tensor vec) -> Tensor(a!) + aten: addmv_(self, mat, vec, beta, alpha) + +- name: addmv(Scalar beta, Tensor self, Scalar alpha, Tensor mat, Tensor vec, *, Tensor(a!) out) -> Tensor(a!) + aten: addmv_out(out, self, mat, vec, beta, alpha) + +- name: addmv(Scalar beta, Tensor self, Tensor mat, Tensor vec) -> Tensor + aten: addmv(self, mat, vec, beta, 1) + +- name: addmv_(Scalar beta, Tensor(a!) self, Tensor mat, Tensor vec) -> Tensor(a!) + aten: addmv_(self, mat, vec, beta, 1) + +- name: addmv(Scalar beta, Tensor self, Tensor mat, Tensor vec, *, Tensor(a!) out) -> Tensor(a!) + aten: addmv_out(out, self, mat, vec, beta, 1) + +- name: addr(Scalar beta, Tensor self, Scalar alpha, Tensor vec1, Tensor vec2) -> Tensor + aten: addr(self, vec1, vec2, beta, alpha) + +- name: addr_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor vec1, Tensor vec2) -> Tensor(a!) + aten: addr_(self, vec1, vec2, beta, alpha) + +- name: addr(Scalar beta, Tensor self, Scalar alpha, Tensor vec1, Tensor vec2, *, Tensor(a!) out) -> Tensor(a!) + aten: addr_out(out, self, vec1, vec2, beta, alpha) + +- name: addr(Scalar beta, Tensor self, Tensor vec1, Tensor vec2) -> Tensor + aten: addr(self, vec1, vec2, beta, 1) + +- name: addr_(Scalar beta, Tensor(a!) self, Tensor vec1, Tensor vec2) -> Tensor(a!) + aten: addr_(self, vec1, vec2, beta, 1) + +- name: addr(Scalar beta, Tensor self, Tensor vec1, Tensor vec2, *, Tensor(a!) out) -> Tensor(a!) + aten: addr_out(out, self, vec1, vec2, beta, 1) + +- name: baddbmm(Scalar beta, Tensor self, Scalar alpha, Tensor batch1, Tensor batch2) -> Tensor + aten: baddbmm(self, batch1, batch2, beta, alpha) + +- name: baddbmm_(Scalar beta, Tensor(a!) self, Scalar alpha, Tensor batch1, Tensor batch2) -> Tensor(a!) + aten: baddbmm_(self, batch1, batch2, beta, alpha) + +- name: baddbmm(Scalar beta, Tensor self, Scalar alpha, Tensor batch1, Tensor batch2, *, Tensor(a!) out) -> Tensor(a!) + aten: baddbmm_out(out, self, batch1, batch2, beta, alpha) + +- name: baddbmm(Scalar beta, Tensor self, Tensor batch1, Tensor batch2) -> Tensor + aten: baddbmm(self, batch1, batch2, beta, 1) + +- name: baddbmm_(Scalar beta, Tensor(a!) self, Tensor batch1, Tensor batch2) -> Tensor(a!) + aten: baddbmm_(self, batch1, batch2, beta, 1) + +- name: baddbmm(Scalar beta, Tensor self, Tensor batch1, Tensor batch2, *, Tensor(a!) out) -> Tensor(a!) + aten: baddbmm_out(out, self, batch1, batch2, beta, 1) + +- name: sub(Tensor self, Scalar alpha, Tensor other) -> Tensor + aten: sub(self, other, alpha) + +- name: sub_(Tensor(a!) self, Scalar alpha, Tensor other) -> Tensor(a!) + aten: sub_(self, other, alpha) + +- name: sub(Tensor self, Scalar alpha, Tensor other, *, Tensor(a!) out) -> Tensor(a!) + aten: sub_out(out, self, other, alpha) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/derivatives.yaml b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/derivatives.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fa611a88889b0518e9f7f6bef17c317631aa6a65 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/derivatives.yaml @@ -0,0 +1,3282 @@ +# Defines derivative formulas and Python signatures of methods on Variable +# +# Note about possibly confusing nomenclature: An 'output gradient' is the +# gradient of an output of a forward function. Output gradients are used as +# the inputs to backward functions. `grads` is a vector of output gradients, +# and `grad == grads[0]`, in all the derivative formulas in this file. +# An 'input gradient' is the gradient of an input to a forward function. +# Input gradients are the outputs of backward functions, corresponding to the +# input names included in the derivative formulas defined in this file. +# Also, every time we talk computing "gradient" we actually mean computing +# the vector jacobian product using the given 'output gradient' as the vector. +# +# Each entry consists of: +# - A 'name', which specifies the ATen name of the function you +# are defining derivatives for, and an argument specification. +# - An optional 'dispatch' entry which can be used to specify +# per-autograd dispatch key derivatives. If this entry is not +# specified, then the gradient entries will be taken as the +# default gradients (i.e. registered for every backward dispatch +# key). (see _test_autograd_multiple_dispatch for an example +# of how to register separate derivates for different dispatch keys). +# The list of allowed dispatch keys (in addition to 'Default' which +# represents the Autograd alias key) is torchgen/model.py:AUTOGRAD_KEYS. +# - One or more gradients entries, mapping differentiable input +# names to a formula specifying how to compute its gradient. +# Note that a single gradient entry can specify the gradient +# formula for multiple input names, by specifying a key +# "input1, input2" (see atan2 for an example). +# - An argument can be flagged as 'non_differentiable'. +# - Optional entry with key 'output_differentiability' and value a list of the +# same length as the number of outputs from the forward function. The list +# should contain only booleans, specifying whether each of the output Tensor +# is differentiable. +# If it is not specified for a function that returns multiple elements but +# uses `grad` instead of `grads[idx]`, then all but the first output will +# be marked as non-differentiable. +# If None of the output is differentiable, you can also add the function +# name to `gen_variable_type.py`'s `DONT_REQUIRE_DERIVATIVE` list. +# +# There are two cases for Tensor and TensorList arguments here: +# - If that argument is differentiable, in the sense that a gradient with respect +# to that argument could exist. You should either: +# - Specify the formula for that gradient +# - Specify not_implemented("function_name") as a formula to say that this is not +# implemented yet (but might be in the future and the user can request that on an issue) +# - If that argument is not differentiable, because it is not a floating point dtype or the +# function is not differentiable with respect to that argument for +# example. You should either: +# - Do not specify any formula for this argument +# - Specify explicitly that this argument is "non_differentiable". Note that in this case, +# we trust you that this argument will never have requires_grad=True and it will be silently +# ignored if it does. +# +# If a function has out-of-place and in-place variants, then the derivative +# definition for the in-place variant is optional. It will default to the +# definition for the out-of-place variant. Note that _out variants are never +# differentiable. +# +# Gradient expressions are standard C++ expressions operating on ATen +# variables. In a gradient expression, the following variables/functions +# are in scope: +# +# - 'grad', the gradient of the output (often spelled grad_output +# in Python) which we are going to left-multiply. +# +# When a function returns multiple *differentiable* outputs, +# you can refer to the gradients of each outputs using 'grads', +# e.g., 'grads[0]', 'grads[1]'. +# +# When a function returns multiple *differentiable* outputs that +# are named, you can refer to the gradients of each outputs using +# 'grad_{name}', e.g., 'grad_x', 'grad_y'. +# +# When a function returns *one* differentiable output (the +# first output) and some more nondifferentiable outputs, +# you MUST refer to the gradient of the differentiable output with +# 'grad' (this case is special-cased in our code generation). +# +# Note that the number of differentiable outputs can be modified by the +# 'output_differentiability' entry (see above). +# +# Across a differentiable function's derivatives set, it is not +# permitted to mix the use of "grad", "grads", and +# "grad_{name}". You must be consistent for that differentiable +# function. +# +# - Any of the input arguments, tensor or non-tensor, including +# argument names that only appear in Declarations.yaml, e.g. 'output'. +# +# - 'result', representing the result of evaluating the forward +# expression for ATen native function declarations. If the forward +# expression outputs a tuple, use 'resultX' instead to access the +# X-th entry +# +# - 'grad_input_mask', a std::array, specifies which input +# gradients are actually needed. For example, in the entry +# `input0, input1: foo(grad_input_mask)`, `grad_input_mask` is a size +# two array, where `grad_input_mask[0]` is true if `input0` requires +# grad, and `grad_input_mask[1]` is true if `input1` requires grad. +# +# (NB: if your function computes gradient for a list of tensors, +# the `grad_input_mask` will only have a single entry for the list +# specifying if either zero or at least one tensor from the list requires +# grad. If we want to support more fine-grained signalling, +# we'll need some alternate variable which is not a std::array) +# +# - 'retain_variables', a bool which is true if a user has specified +# that saved variables should be retained in case the backwards is +# run again later. This allows an optimization where we can +# destroy saved buffers if we know variables are not going to be retained, +# e.g., it is used by _cudnn_rnn +# +# - `wrap_opt_if`, is a 2-argument function that accepts a tensor +# variable and a boolean condition that dictates whether to save that +# variable in a graph. The result of this function is `std::optional`, +# and it is `::std::nullopt` when the condition evaluates to `false`, +# otherwise it is the variable wrapped in `std::optional`. +# For example, wrap_opt_if(var_0, grad_input_mask[1] || grad_input_mask[2]) +# would mean that `var_0` is saved as long as the second (grad_input_mask[1]) +# or the third (grad_input_mask[2]) argument requires gradients. +# Another interpretation of this expression would read as `var_0` is needed +# in the backward computation of the second or the third argument. +# NOTE: the usage of `var_i.requires_grad()` in the conditional expression +# is not supported, use `grad_input_mask[i]` instead. +# NOTE: `wrap_opt_if` could be used to prevent saving redundant variables +# with multi-output backward formulas. +# See https://github.com/pytorch/pytorch/issues/97575 for more details +# on the issue. +# +# If you need a complex expression, e.g., with local variables, +# write a _backward function in torch/csrc/autograd/FunctionsManual.cpp +# and invoke it from here. By the way, go read +# https://github.com/zdevito/ATen/issues/163; this describes an +# important hazard that occurs when porting backwards from Python to C++ +# +# Double backwards gradient expressions can be somewhat confusing; +# the most important thing to remember is: (1) you need to define a +# derivative formula for every input, including inputs named things +# like 'grad_output', and (2) the gradient to multiply with is always +# called 'grad' (even though it really is a grad-grad). +# +# You can also add forward derivative definition by defining a formula for +# a returned value (in general "result" if the name is not specified). This +# formula works the same way as the backward one and advanced implementations +# should also be placed in the FunctionsManual file. +# This formula should compute a single Jacobian vector product using the (primal) +# value of the argument "foo_p", its forward grad "foo_t" and the result of the +# function as "result". +# Note that the forward derivative can be automatically generated in two cases: +# - if your function is linear (NOT affine or multi-linear), then you can +# specify so by just using the string "auto_linear" for the formula. +# - if your function is applied element wise (and has a single input), you +# can specify so by just using the string "auto_element_wise" for the formula. +# +# Note that to avoid unpacking overhead, functions taking TensorList as inputs +# will always have their forward grad formula called. This function is responsible +# to check if any computation is needed and should return an undefined Tensor when +# there is nothing to do. You can check "cat_forward" for a full example. +# +# NB: There are a number of gradient definitions in here which are bogus +# (implemented using zeros_like). These gradients are (hopefully) not +# used by our frontend. You MUST check the frontend code; search for +# OpName.apply to see if it's still using a legacy Python style API. +# +# Note: Returning views. +# The following cases exist: +# - If a function returns no view, it can have arbitrary outputs. +# - If a function return at least one Tensor that is a differentiable view +# of one of its input: +# - If there is only one differentiable output, this Tensor is marked as a +# differentiable view. (alias or transpose for example) +# - If there are more than one differentiable output, by default all the views are +# marked as differentiable views and created with allow_rebase_history=false. +# Meaning that any inplace operation on it will raise an error. (unbind for example) +# +# Notes about undefined output gradients: +# All backward functions must support all combinations of undefined output +# gradient Tensors, where `grad[i].defined() == false`. Depending on the +# number of input and output grads your derivative formula uses, code +# generation may automatically add some level of undefined grad support, +# according to these three cases: +# +# * 1 input grad and 1 output grad: +# Complete undefined grad support is automatically added, so you +# shouldn't have to think about it, unless there is a bug in the code +# generation. +# +# * 1 input grad and multiple output grads: +# Undefined grad support is automatically added ONLY in the case where +# all output grads are undefined. You will have to add explicit support +# for cases where a subset of output grads is undefined. +# +# * multiple input grads: +# No automatic support, so you will need to add it. +# +# If your derivative formula uses more than one output grad, it is usually +# preferable to add undefined grad support in the backward function itself +# (if you're using one), rather than in the derivative formula in this file. +# +# Undefined Tensors are created with the default constructor `at::Tensor()`. +# It is an efficient way to represent a Tensor filled with zeros because +# the Tensor holds no sizing information and no Storage data is allocated. +# But consequently, Tensor operations cannot be performed on them. +# Therefore, your backward function should treat an undefined output grad as +# a zero, and it needs to be a special case. +# +# If all output grads are undefined, then it should be correct for the +# backward function to return undefined input grads. Since we use the chain +# rule, output grads equal to zero should result in input grads equal to zero, +# unless there is some rare special case. +# +# If a subset of output grads is undefined, then it may be acceptable for +# the backward function to return undefined input grads--it depends on the +# specific function, so you'll have to determine that yourself. If returning +# an undefined Tensor is correct for a given input grad, it is also logically +# correct to return a defined grad full of zeros, but that would not be +# preferable since it would be less efficient. +# +# NB: The parameter names here MUST be consistent with the parameter names +# in native_functions.yaml +- name: abs(Tensor self) -> Tensor + self: grad * self.sgn() + result: handle_r_to_c(result.scalar_type(), self_t.conj() * self_p.sgn()) + +- name: acos(Tensor self) -> Tensor + self: grad * -((-self * self + 1).rsqrt()).conj() + result: auto_element_wise + +- name: add.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor + self: handle_r_to_c(self.scalar_type(), grad) + other: handle_r_to_c(other.scalar_type(), maybe_multiply(grad, alpha.conj())) + result: self_t + maybe_multiply(other_t, alpha) + +- name: add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor + self: handle_r_to_c(self.scalar_type(), grad) + result: self_t.clone() + +- name: addbmm(Tensor self, Tensor batch1, Tensor batch2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + self: maybe_multiply(grad, beta.conj()) + batch1: maybe_multiply(grad.unsqueeze(0).expand_symint({ batch1.sym_size(0), batch1.sym_size(1), batch2.sym_size(2) }).bmm(batch2.transpose(1, 2).conj()), alpha.conj()) + batch2: maybe_multiply(batch1.transpose(1, 2).conj().bmm(grad.unsqueeze(0).expand_symint({ batch1.sym_size(0), batch1.sym_size(1), batch2.sym_size(2) })), alpha.conj()) + result: maybe_multiply(self_t, beta) + maybe_multiply(batch1_t.bmm(batch2_p).sum(0), alpha) + maybe_multiply(batch1_p.bmm(batch2_t).sum(0), alpha) + +- name: addcdiv(Tensor self, Tensor tensor1, Tensor tensor2, *, Scalar value=1) -> Tensor + self: handle_r_to_c(self.scalar_type(), grad) + tensor1: handle_r_to_c(tensor1.scalar_type(), grad * (value / tensor2).conj()) + tensor2: handle_r_to_c(tensor2.scalar_type(), -grad * (value * tensor1 / (tensor2 * tensor2)).conj()) + result: self_t + maybe_multiply(tensor1_t / tensor2_p, value) - maybe_multiply(tensor2_t * (tensor1_p / tensor2_p) / tensor2_p, value) + +- name: addcmul(Tensor self, Tensor tensor1, Tensor tensor2, *, Scalar value=1) -> Tensor + self: handle_r_to_c(self.scalar_type(), grad) + tensor1: handle_r_to_c(tensor1.scalar_type(), grad * (tensor2 * value).conj()) + tensor2: handle_r_to_c(tensor2.scalar_type(), grad * (tensor1 * value).conj()) + result: self_t + maybe_multiply(tensor1_t * tensor2_p, value) + maybe_multiply(tensor2_t * tensor1_p, value) + +- name: addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + self: maybe_multiply(grad, beta.conj()) + mat1: mm_mat1_backward(grad, mat2, mat1.sym_sizes(), mat1.sym_strides(), mat1.layout(), alpha) + mat2: mm_mat2_backward(grad, mat1, mat2.sym_sizes(), mat2.sym_strides(), mat2.layout(), alpha) + result: maybe_multiply(self_t, beta) + maybe_multiply(mat1_t.mm(mat2_p), alpha) + maybe_multiply(mat1_p.mm(mat2_t), alpha) + +- name: _sparse_addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + self: maybe_multiply(grad, beta) + mat1: mm_mat1_sparse_backward(grad, mat1, mat2, alpha) + mat2: mm_mat2_backward(grad, mat1, mat2.sym_sizes(), mat2.sym_strides(), mat2.layout(), alpha) + +- name: addmv(Tensor self, Tensor mat, Tensor vec, *, Scalar beta=1, Scalar alpha=1) -> Tensor + self: maybe_multiply(grad, beta.conj()) + mat: maybe_multiply(grad.ger(vec.conj()), alpha.conj()) + vec: maybe_multiply(mat.t().conj().mv(grad), alpha.conj()) + result: maybe_multiply(self_t, beta) + maybe_multiply(mat_t.mv(vec_p), alpha) + maybe_multiply(mat_p.mv(vec_t), alpha) + +- name: addr(Tensor self, Tensor vec1, Tensor vec2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + self: maybe_multiply(grad, beta.conj()) + vec1: maybe_multiply(grad.mv(vec2.conj()), alpha.conj()) + vec2: maybe_multiply(grad.t().mv(vec1.conj()), alpha.conj()) + result: maybe_multiply(self_t, beta) + maybe_multiply(vec1_t.outer(vec2_p), alpha) + maybe_multiply(vec1_p.outer(vec2_t), alpha) + +- name: affine_grid_generator(Tensor theta, SymInt[] size, bool align_corners) -> Tensor + theta: affine_grid_generator_backward_symint(grad, size, align_corners) + result: auto_linear + +- name: alias(Tensor(a) self) -> Tensor(a) + self: grad + result: self_t + +- name: angle(Tensor self) -> Tensor + self: angle_backward(grad, self) + result: handle_r_to_c(result.scalar_type(), angle_backward(self_t.conj(), self_p).conj()) + +# The four items below are necessary because TensorIterator doesn't work on +# Variables (codegen does not unwrap the input Tensor for all() and any() ). +- name: any(Tensor self) -> Tensor + output_differentiability: [False] + +- name: any.dim(Tensor self, int dim, bool keepdim=False) -> Tensor + output_differentiability: [False] + +- name: any.dims(Tensor self, int[]? dim=None, bool keepdim=False) -> Tensor + output_differentiability: [False] + +- name: _is_all_true(Tensor self) -> Tensor + self: non_differentiable + +- name: _is_any_true(Tensor self) -> Tensor + self: non_differentiable + +- name: all(Tensor self) -> Tensor + output_differentiability: [False] + +- name: all.dim(Tensor self, int dim, bool keepdim=False) -> Tensor + output_differentiability: [False] + +- name: all.dims(Tensor self, int[]? dim=None, bool keepdim=False) -> Tensor + output_differentiability: [False] + +- name: acosh(Tensor self) -> Tensor +# Save one rsqrt in the real case by using that for x real and positive sqrt(x*y) = sqrt(x)*sqrt(y) (not true in the complex case) + self: "self.is_complex() ? grad * ((self + 1).rsqrt() * (self - 1).rsqrt()).conj() : grad * (self * self - 1).rsqrt()" + result: auto_element_wise + +- name: acosh_(Tensor(a!) self) -> Tensor(a!) + self: not_implemented("inplace version of acosh") + +- name: asinh(Tensor self) -> Tensor + self: grad * (self.pow(2) + 1).rsqrt().conj() + result: auto_element_wise + +- name: asinh_(Tensor(a!) self) -> Tensor(a!) + self: not_implemented("inplace version of asinh") + +- name: atanh(Tensor self) -> Tensor + self: grad * 1 / (1 - self.pow(2)).conj() + result: auto_element_wise + +- name: atanh_(Tensor(a!) self) -> Tensor(a!) + self: not_implemented("inplace version of atanh") + +- name: as_strided(Tensor(a) self, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor(a) + self: as_strided_backward(grad, TensorGeometry(self), size, stride, storage_offset) + result: auto_linear + +- name: as_strided_(Tensor(a!) self, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor(a!) + self: as_strided_backward(grad, TensorGeometry(self), size, stride, storage_offset) + result: auto_linear + +- name: asin(Tensor self) -> Tensor + self: grad * (-self * self + 1).rsqrt().conj() + result: auto_element_wise + +- name: atan(Tensor self) -> Tensor + self: grad / (self * self + 1).conj() + result: auto_element_wise + +- name: atan2(Tensor self, Tensor other) -> Tensor + self, other: atan2_backward(grad, self, other, grad_input_mask) + result: (-self_p * other_t + other_p * self_t) / (self_p.pow(2) + other_p.pow(2)) + +- name: baddbmm(Tensor self, Tensor batch1, Tensor batch2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + self: maybe_multiply(grad, beta.conj()) + batch1: maybe_multiply(grad.bmm(batch2.transpose(1, 2).conj()), alpha.conj()) + batch2: maybe_multiply(batch1.transpose(1, 2).conj().bmm(grad), alpha.conj()) + result: maybe_multiply(self_t, beta) + maybe_multiply(batch1_t.bmm(batch2_p), alpha) + maybe_multiply(batch1_p.bmm(batch2_t), alpha) + +- name: bernoulli(Tensor self, *, Generator? generator=None) -> Tensor + self: zeros_like(grad) + result: auto_element_wise + +- name: bernoulli_.Tensor(Tensor(a!) self, Tensor p, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + p: zeros_like(p) + result: self_t.zero_() + +- name: bernoulli_.float(Tensor(a!) self, float p=0.5, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: bmm(Tensor self, Tensor mat2) -> Tensor + self: grad.bmm(mat2.transpose(1, 2).conj()) + mat2: self.transpose(1, 2).conj().bmm(grad) + result: self_t.bmm(mat2_p) + self_p.bmm(mat2_t) + +- name: matmul(Tensor self, Tensor other) -> Tensor + self, other: matmul_backward(grad, self, other, grad_input_mask) + +- name: cat(Tensor[] tensors, int dim=0) -> Tensor + tensors: cat_tensors_backward(grad, to_args_sizes_symint(tensors), to_args_scalartypes(tensors), dim) + result: cat_jvp(tensors, dim) + +- name: cauchy_(Tensor(a!) self, float median=0, float sigma=1, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: ceil(Tensor self) -> Tensor + self: zeros_like(grad) + result: auto_element_wise + +- name: cholesky(Tensor self, bool upper=False) -> Tensor + self: cholesky_backward(grad, upper, result) + +- name: chunk(Tensor(a -> *) self, int chunks, int dim=0) -> Tensor(a)[] + dispatch: + Default: + # the default case will use the CompositeImplicitAutograd + self: not_implemented("chunk") + AutogradNestedTensor: + self: chunk_backward_nested(grads, self, chunks, dim) + +- name: linalg_cholesky_ex(Tensor self, *, bool upper=False, bool check_errors=False) -> (Tensor L, Tensor info) + self: cholesky_backward(grad, upper, L) + L: cholesky_jvp(self_t, L, upper) + +- name: cholesky_solve(Tensor self, Tensor input2, bool upper=False) -> Tensor + self, input2: cholesky_solve_backward(grad, self, input2, result, upper, grad_input_mask) + result: cholesky_solve_jvp(result, input2_p, input2_t, self_t, upper) + +- name: cholesky_inverse(Tensor self, bool upper=False) -> Tensor + self: cholesky_inverse_backward(grad, self, upper, result) + result: cholesky_inverse_jvp(self_p, self_t, result, upper) + +# For clamp, gradient is not defined at the boundaries. But empirically it's helpful +# to be able to get gradient on min and max, so we return the subgradient 1 for these cases. +- name: clamp.Tensor(Tensor self, Tensor? min=None, Tensor? max=None) -> Tensor + self: clamp_backward(grad, self, min, max) + min, max: clamp_backward_min_max(grad, self, min, max, grad_input_mask) + result: clamp_jvp(self_p, self_t, min_p, min_t, max_p, max_t) + +- name: clamp(Tensor self, Scalar? min=None, Scalar? max=None) -> Tensor + self: clamp_backward(grad, self, min, max) + result: auto_element_wise + +- name: clamp_min(Tensor self, Scalar min) -> Tensor + self: where(self >= min, grad, at::scalar_tensor(0., grad.options())) + result: auto_element_wise + +- name: clamp_min.Tensor(Tensor self, Tensor min) -> Tensor + self: where(self >= min, grad, at::scalar_tensor(0., grad.options())) + min: where(self < min, grad, at::scalar_tensor(0., grad.options())) + result: where(self_p >= min_p, self_t, min_t) + +- name: clamp_max(Tensor self, Scalar max) -> Tensor + self: where(self <= max, grad, at::scalar_tensor(0., grad.options())) + result: auto_element_wise + +- name: clamp_max.Tensor(Tensor self, Tensor max) -> Tensor + self: where(self <= max, grad, at::scalar_tensor(0., grad.options())) + max: where(self > max, grad, at::scalar_tensor(0., grad.options())) + result: where(self_p <= max_p, self_t, max_t) + +- name: clone(Tensor self, *, MemoryFormat? memory_format=None) -> Tensor + self: grad + result: auto_linear + +- name: _lazy_clone(Tensor self) -> Tensor + self: grad + result: auto_linear + +- name: _to_copy(Tensor self, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None, bool non_blocking=False, MemoryFormat? memory_format=None) -> Tensor + self: _to_copy_backward(grad, self.options()) + result: _to_copy(self_t, dtype, layout, device, pin_memory, non_blocking, memory_format) + # The condition is: if dtype is not nullopt, then isDifferentiableType(*dtype) + # (If dtype IS nullopt, we rely on the regular check that any input requires grad). + output_differentiability: ["!dtype || isDifferentiableType(*dtype)"] + +- name: _coalesce(Tensor self) -> Tensor + self: grad + +- name: complex(Tensor real, Tensor imag) -> Tensor + real: at::real(grad) + imag: at::imag(grad) + result: at::complex(real_t, imag_t) + +- name: polar(Tensor abs, Tensor angle) -> Tensor + abs, angle: polar_backward(grad, result) + result: at::complex(abs_t*angle_p.cos() - angle_t*abs_p*angle_p.sin(), abs_t*angle_p.sin() + angle_t*abs_p*angle_p.cos()) + +- name: _conj(Tensor(a) self) -> Tensor(a) + self: grad.conj() + result: self_t.conj() + +- name: _neg_view(Tensor(a) self) -> Tensor(a) + self: grad.neg() + result: self_t._neg_view() + +- name: _conj_physical(Tensor self) -> Tensor + self: grad.conj_physical() + result: self_t.conj_physical() + +- name: conj_physical_(Tensor(a!) self) -> Tensor(a!) + self: grad.conj_physical() + result: self_t.conj_physical_() + +- name: copysign.Tensor(Tensor self, Tensor other) -> Tensor + self: copysign_tensor_self_backward(grad, self, result) + other: zeros_like(other) + result: copysign_tensor_self_backward(self_t, self_p, result) + +- name: copysign.Scalar(Tensor self, Scalar other) -> Tensor + self: copysign_tensor_self_backward(grad, self, result) + result: auto_element_wise + +- name: cos(Tensor self) -> Tensor + self: grad * -self.sin().conj() + result: auto_element_wise + +- name: cosh(Tensor self) -> Tensor + self: grad * self.sinh().conj() + result: auto_element_wise + +- name: count_nonzero.dim_IntList(Tensor self, int[] dim) -> Tensor + output_differentiability: [False] + +- name: count_nonzero(Tensor self, int? dim=None) -> Tensor + output_differentiability: [False] + +- name: linalg_cross(Tensor self, Tensor other, *, int dim=-1) -> Tensor + self: at::linalg_cross(other.conj(), grad, dim) + other: at::linalg_cross(grad, self.conj(), dim) + result: "at::linalg_cross(self_t, other_p, dim) + at::linalg_cross(self_p, other_t, dim)" + +- name: logcumsumexp(Tensor self, int dim) -> Tensor + self: logcumsumexp_backward(grad, self, result, dim) + result: logcumsumexp_jvp(self_p, self_t, dim) + +- name: cumprod(Tensor self, int dim, *, ScalarType? dtype=None) -> Tensor + self: cumprod_backward(grad.to(self.scalar_type()), self, dim, result) + result: "cumprod_jvp(self_t, self_p, result, dim).to(dtype.has_value() ? *dtype : self_p.scalar_type())" + +- name: cumsum(Tensor self, int dim, *, ScalarType? dtype=None) -> Tensor + self: cumsum_backward(grad.to(self.scalar_type()), dim) + result: auto_linear + +- name: cummax(Tensor self, int dim) -> (Tensor values, Tensor indices) + self: cummaxmin_backward(grad, self, indices, dim) + values: self_t.gather(dim, indices) + +- name: cummin(Tensor self, int dim) -> (Tensor values, Tensor indices) + self: cummaxmin_backward(grad, self, indices, dim) + values: self_t.gather(dim, indices) + +- name: conv_tbc(Tensor self, Tensor weight, Tensor bias, int pad=0) -> Tensor + self, weight, bias: "grad.defined() ? conv_tbc_backward(grad, self, weight, bias, pad) : std::tuple()" + +- name: _ctc_loss(Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, int blank=0, bool zero_infinity=False) -> (Tensor, Tensor) + log_probs: _ctc_loss_backward(grad, log_probs, targets, input_lengths, target_lengths, result0, result1, blank, zero_infinity) + +- name: _ctc_loss.Tensor(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank=0, bool zero_infinity=False) -> (Tensor, Tensor) + log_probs: _ctc_loss_backward(grad, log_probs, targets, input_lengths, target_lengths, result0, result1, blank, zero_infinity) + +- name: deg2rad(Tensor self) -> Tensor + self: deg2rad_backward(grad) + result: auto_element_wise + +- name: _linalg_det(Tensor A) -> (Tensor result, Tensor LU, Tensor pivots) + A: linalg_det_backward(grad, result, A, LU, pivots) + result: linalg_det_jvp(A_t, result, LU, pivots, A_p.is_contiguous() && !A_p.is_complex()) + output_differentiability: [True, False, False] + +- name: _linalg_slogdet(Tensor A) -> (Tensor sign, Tensor logabsdet, Tensor LU, Tensor pivots) + A: slogdet_backward(grad_sign, grad_logabsdet, A, sign, LU, pivots) + sign, logabsdet: slogdet_jvp(LU, pivots, A_t, sign, A_p.is_contiguous() && !A_p.is_complex()) + output_differentiability: [True, True, False, False] + +- name: block_diag(Tensor[] tensors) -> Tensor + tensors: block_diag_backward(grad, to_args_sizes(tensors), to_args_scalartypes(tensors)) + result: block_diag_jvp(tensors) + +- name: diag_embed(Tensor self, int offset=0, int dim1=-2, int dim2=-1) -> Tensor + self: grad.diagonal(offset, dim1, dim2) + result: auto_linear + +- name: diagonal(Tensor(a) self, int offset=0, int dim1=0, int dim2=1) -> Tensor(a) + self: diagonal_backward_symint(grad, self.sym_sizes(), offset, dim1, dim2) + result: auto_linear + +- name: diagonal_backward(Tensor grad_output, SymInt[] input_sizes, int offset, int dim1, int dim2) -> Tensor + grad_output: grad.diagonal(offset, dim1, dim2) + result: auto_linear + +- name: dist(Tensor self, Tensor other, Scalar p=2) -> Tensor + self: norm_backward(grad, self - other, p, result) + other: -norm_backward(grad, self - other, p, result) + result: norm_jvp(self_p - other_p, self_t - other_t, p, result, {}, false) + +# The backward formula is done in this order to improve numerical stability +# of the higher order derivatives, see https://github.com/pytorch/pytorch/issues/43414 +# Note that we don't use "result" because saving it would be BC-breaking when it is used in an inplace operation later +- name: div.Tensor(Tensor self, Tensor other) -> Tensor + self: div_tensor_self_backward(grad, other, self.scalar_type()) + other: div_tensor_other_backward(grad, self, other) + result: (self_t - other_t * result) / other_p + +- name: div.Scalar(Tensor self, Scalar other) -> Tensor + self: div_tensor_self_backward(grad, other, self.scalar_type()) + result: self_t / other + +- name: div.Tensor_mode(Tensor self, Tensor other, *, str? rounding_mode) -> Tensor + self: div_tensor_self_backward(grad, other, self.scalar_type(), rounding_mode) + other: div_tensor_other_backward(grad, self, other, rounding_mode) + result: "rounding_mode.has_value() ? result.new_zeros_symint(result.sym_sizes()) : self_t / other_p - other_t * (self_p / other_p) / other_p" + +- name: div.Scalar_mode(Tensor self, Scalar other, *, str? rounding_mode) -> Tensor + self: div_tensor_self_backward(grad, other, self.scalar_type(), rounding_mode) + result: "rounding_mode.has_value() ? result.new_zeros_symint(result.sym_sizes()) : self_t / other" + +- name: dot(Tensor self, Tensor tensor) -> Tensor + self: grad * tensor.conj() + tensor: grad * self.conj() + result: at::dot(self_t, tensor_p) + at::dot(self_p, tensor_t) + +- name: vdot(Tensor self, Tensor other) -> Tensor + self: grad.conj() * other + other: grad * self + result: at::vdot(self_t, other_p) + at::vdot(self_p, other_t) + +- name: _fused_dropout(Tensor self, float p, Generator? generator=None) -> (Tensor, Tensor) + self: _fused_dropout_backward(grad, result1, p) + +- name: native_dropout(Tensor input, float p, bool? train) -> (Tensor, Tensor) + input: "GradMode::is_enabled() ? infinitely_differentiable_native_dropout_backward(grad, result1, (!train.has_value() || !train.value() ? 1 : (p == 1 ? 0.0 : 1.0 / (1.0 - p)))) : native_dropout_backward(grad, result1, (!train.has_value() || !train.value() ? 1 : (p == 1 ? 0.0 : 1.0 / (1.0 - p))))" + result0: "(!train.has_value() || train.value()) ? (p == 1 ? 0.0 : 1.0 / (1.0 - p)) * input_t * result1 : input_t" + +- name: native_dropout_backward(Tensor grad_output, Tensor mask, float scale) -> Tensor + grad_output: "native_dropout_double_backward(grad, grad_output, mask, scale)" + mask: 'not_implemented("native_dropout_backward: mask")' + +- name: eq_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + self: zeros_like(self) + result: self_t.zero_() + +- name: eq_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + self: zeros_like(self) + other: zeros_like(other) + result: self_t.zero_() + +- name: erf(Tensor self) -> Tensor + self: 2.0 / sqrt(M_PI) * exp(-(self.pow(2))) * grad + result: auto_element_wise + +- name: erfc(Tensor self) -> Tensor + self: -2.0 / sqrt(M_PI) * exp(-(self.pow(2))) * grad + result: auto_element_wise + +- name: special_erfcx(Tensor self) -> Tensor + self: (2.0 * self * result - 2.0 / sqrt(M_PI)) * grad + result: auto_element_wise + +- name: erfinv(Tensor self) -> Tensor + self: 0.5 * sqrt(M_PI) * exp(self.erfinv().pow(2)) * grad + result: auto_element_wise + +- name: exp(Tensor self) -> Tensor + self: grad * result.conj() + result: auto_element_wise + +- name: exp2(Tensor self) -> Tensor + self: grad * result.conj() * M_LN2 + result: auto_element_wise + +- name: expm1(Tensor self) -> Tensor + self: grad * (result.conj() + 1) + result: auto_element_wise + +# TODO: this derivative is not SymInt safe, need sum_to support +- name: expand(Tensor(a) self, SymInt[] size, *, bool implicit=False) -> Tensor(a) + self: at::sum_to(grad, self.sym_sizes()) + result: auto_linear + +- name: exponential_(Tensor(a!) self, float lambd=1, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: fake_quantize_per_tensor_affine_cachemask(Tensor self, float scale, int zero_point, int quant_min, int quant_max) -> (Tensor output, Tensor mask) + self: fake_quantize_per_tensor_affine_cachemask_backward(grad, mask) + +- name: _fake_quantize_per_tensor_affine_cachemask_tensor_qparams(Tensor self, Tensor scale, Tensor zero_point, Tensor fake_quant_enabled, int quant_min, int quant_max) -> (Tensor output, Tensor mask) + self: fake_quantize_per_tensor_affine_cachemask_backward(grad, mask) + +- name: _fake_quantize_learnable_per_tensor_affine(Tensor self, Tensor scale, Tensor zero_point, int quant_min, int quant_max, float grad_factor=1.0) -> Tensor + self, scale, zero_point: "grad.defined() ? _fake_quantize_learnable_per_tensor_affine_backward(grad, self, scale, zero_point, quant_min, quant_max, grad_factor) : std::tuple()" + +- name: fake_quantize_per_channel_affine_cachemask(Tensor self, Tensor scale, Tensor zero_point, int axis, int quant_min, int quant_max) -> (Tensor output, Tensor mask) + self: fake_quantize_per_channel_affine_cachemask_backward(grad, mask) + +- name: _fake_quantize_learnable_per_channel_affine(Tensor self, Tensor scale, Tensor zero_point, int axis, int quant_min, int quant_max, float grad_factor=1.0) -> Tensor + self, scale, zero_point: "grad.defined() ? _fake_quantize_learnable_per_channel_affine_backward(grad, self, scale, zero_point, axis, quant_min, quant_max, grad_factor) : std::tuple()" + +- name: _fused_moving_avg_obs_fq_helper(Tensor self, Tensor observer_on, Tensor fake_quant_on, Tensor(a!) running_min, Tensor(b!) running_max, Tensor(c!) scale, Tensor(d!) zero_point, float averaging_const, int quant_min, int quant_max, int ch_axis, bool per_row_fake_quant=False, bool symmetric_quant=False) -> (Tensor output, Tensor mask) + self: fake_quantize_per_tensor_affine_cachemask_backward(grad, mask) + +- name: fill.Scalar(Tensor self, Scalar value) -> Tensor + self: zeros_like(grad) + result: at::fill(self_t, 0) + +- name: fill.Tensor(Tensor self, Tensor value) -> Tensor + self: zeros_like(grad) + value: grad.sum() + result: at::fill(self_t, value_t) + +- name: fill_.Scalar(Tensor(a!) self, Scalar value) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.fill_(0) + +- name: fill_.Tensor(Tensor(a!) self, Tensor value) -> Tensor(a!) + self: zeros_like(grad) + value: grad.sum() + result: self_t.fill_(value_t) + +- name: floor(Tensor self) -> Tensor + self: zeros_like(grad) + result: auto_element_wise + +- name: fmod.Scalar(Tensor self, Scalar other) -> Tensor + self: grad + result: auto_element_wise + +- name: fmod.Tensor(Tensor self, Tensor other) -> Tensor + self: grad + other: -grad * self.div(other, /*rounding_mode=*/"trunc") + result: self_t - other_t * self_p.div(other_p, /*rounding_mode=*/"trunc") + +- name: frac(Tensor self) -> Tensor + self: grad + result: self_t + +- name: frexp.Tensor(Tensor self) -> (Tensor mantissa, Tensor exponent) + self: grad / exponent.exp2() + mantissa: self_t / exponent.exp2() + +- name: gather(Tensor self, int dim, Tensor index, *, bool sparse_grad=False) -> Tensor + self: gather_backward(grad, self, dim, index, sparse_grad) + index: non_differentiable + result: auto_linear + +- name: ge_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + self: zeros_like(self) + result: self_t.zero_() + +- name: ge_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + self: zeros_like(self) + other: zeros_like(other) + result: self_t.zero_() + +- name: geometric_(Tensor(a!) self, float p, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: geqrf(Tensor self) -> (Tensor a, Tensor tau) + self: not_implemented("geqrf") + +- name: indices(Tensor(a) self) -> Tensor(a) + output_differentiability: [False] + +- name: _indices(Tensor(a) self) -> Tensor(a) + output_differentiability: [False] + +- name: crow_indices(Tensor(a) self) -> Tensor(a) + output_differentiability: [False] + +- name: col_indices(Tensor(a) self) -> Tensor(a) + output_differentiability: [False] + +- name: ccol_indices(Tensor(a) self) -> Tensor(a) + output_differentiability: [False] + +- name: row_indices(Tensor(a) self) -> Tensor(a) + output_differentiability: [False] + +- name: grid_sampler_2d(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor + input, grid: "grad.defined() ? grid_sampler_2d_backward(grad, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask) : std::tuple()" + +- name: grid_sampler_2d_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners, bool[2] output_mask) -> (Tensor, Tensor) + grad_output, input, grid: grid_sampler_2d_double_backward(grads[0], grads[1], grad_output, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask) + +- name: grid_sampler_3d(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor + input, grid: "grad.defined() ? grid_sampler_3d_backward(grad, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask) : std::tuple()" + +- name: grid_sampler_3d_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners, bool[2] output_mask) -> (Tensor, Tensor) + grad_output, input, grid: grid_sampler_3d_double_backward(grads[0], grads[1], grad_output, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask) + +# See NOTE [ grid_sample CPU fallback ] +- name: _grid_sampler_2d_cpu_fallback(Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> Tensor + input, grid: "grad.defined() ? _grid_sampler_2d_cpu_fallback_backward(grad, input, grid, interpolation_mode, padding_mode, align_corners) : std::tuple()" + +- name: _grid_sampler_2d_cpu_fallback_backward(Tensor grad_output, Tensor input, Tensor grid, int interpolation_mode, int padding_mode, bool align_corners) -> (Tensor, Tensor) + grad_output, input, grid: grid_sampler_2d_double_backward(grads[0], grads[1], grad_output, input, grid, interpolation_mode, padding_mode, align_corners, grad_input_mask) + +- name: gt_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + self: zeros_like(self) + result: self_t.zero_() + +- name: gt_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + self: zeros_like(self) + other: zeros_like(other) + result: self_t.zero_() + +- name: hardsigmoid(Tensor self) -> Tensor + self: hardsigmoid_backward(grad, self) + result: auto_element_wise + +- name: histc(Tensor self, int bins=100, Scalar min=0, Scalar max=0) -> Tensor + output_differentiability: [False] + +- name: hardswish(Tensor self) -> Tensor + self: hardswish_backward(grad, self) + result: auto_element_wise + +- name: hardswish_backward(Tensor grad_output, Tensor self) -> Tensor + grad_output: hardswish_backward(grad, self) + self: at::where(at::logical_and(-3.0 < self, self < 3.0), grad * grad_output / 3.0, at::zeros({}, self.options())) + result: "hardswish_backward(grad_output_t, self_p) + + at::where(at::logical_and(-3.0 < self_p, self_p < 3.0), self_t * grad_output_p / 3.0, at::zeros({}, self_p.options()))" + +- name: hypot(Tensor self, Tensor other) -> Tensor + self: grad * self / result + other: grad * other / result + result: self_t * self_p / result + other_t * other_p / result + +- name: i0(Tensor self) -> Tensor + self: grad * at::special_i1(self) + result: auto_element_wise + +- name: special_i0e(Tensor self) -> Tensor + self: grad * (at::special_i1e(self) - self.sgn() * result) + result: auto_element_wise + +- name: special_i1(Tensor self) -> Tensor + self: i1_backward(grad, self, result) + result: auto_element_wise + +- name: special_i1e(Tensor self) -> Tensor + self: i1e_backward(grad, self, result) + result: auto_element_wise + +- name: igamma(Tensor self, Tensor other) -> Tensor + self: 'not_implemented("igamma: input")' + other: grad * exp((self - 1) * log(other) - other - lgamma(self)) + +- name: igammac(Tensor self, Tensor other) -> Tensor + self: 'not_implemented("igammac: input")' + other: -grad * exp((self - 1) * log(other) - other - lgamma(self)) + +- name: index.Tensor(Tensor self, Tensor?[] indices) -> Tensor + self: index_backward(grad.new_zeros_symint(self.sym_sizes(), self.options()), indices, grad) + result: auto_linear + +- name: _unsafe_index.Tensor(Tensor self, Tensor?[] indices) -> Tensor + self: at::_unsafe_index_put(grad.new_zeros_symint(self.sym_sizes(), self.options()), indices, grad, true) + result: auto_linear + +- name: _unsafe_masked_index(Tensor self, Tensor mask, Tensor?[] indices, Scalar fill) -> Tensor + self: at::_unsafe_masked_index_put_accumulate(grad.new_zeros_symint(self.sym_sizes(), self.options()), mask, indices, grad) + mask: non_differentiable + result: _unsafe_masked_index(self_t, mask, indices, 0) + +- name: _unsafe_masked_index_put_accumulate(Tensor self, Tensor mask, Tensor?[] indices, Tensor values) -> Tensor + self: grad + mask: non_differentiable + values: at::_unsafe_masked_index(grad, mask, indices, 0) + result: at::_unsafe_masked_index_put_accumulate(self_t, mask, indices, values_t) + +- name: index_add(Tensor self, int dim, Tensor index, Tensor source, *, Scalar alpha=1) -> Tensor + self: grad + # The case source.dim() == 0 is necessary to support scalar tensors of the form + # source.dim() == 0 and index.dim() == 1 and index.size() == (1,), + # This is because source is not broadcastable to index, as source.dim() < index.dim() + source: "maybe_multiply(source.dim() > 0 ? grad.index_select(dim, index).expand_as(source) : grad.index_select(dim, index.squeeze(0)), alpha)" + index: non_differentiable + result: at::index_add(self_t, dim, index, maybe_multiply(source_t, alpha)) + +- name: index_reduce(Tensor self, int dim, Tensor index, Tensor source, str reduce, *, bool include_self=True) -> Tensor + self, source: index_reduce_backward(grad, self, dim, index, source, reduce, include_self, result) + index: non_differentiable + +- name: index_copy(Tensor self, int dim, Tensor index, Tensor source) -> Tensor + self: grad.index_fill(dim, index, 0) + # The case source.dim() == 0 is necessary to support scalar tensors of the form + # source.dim() == 0 and index.dim() == 1 and index.size() == (1,), + # This is because source is not broadcastable to index, as source.dim() < index.dim() + source: "source.dim() > 0 ? grad.index_select(dim, index).expand_as(source) : grad.index_select(dim, index.squeeze(0))" + index: non_differentiable + result: self_t.index_copy(dim, index, source_t) + +- name: index_fill.int_Scalar(Tensor self, int dim, Tensor index, Scalar value) -> Tensor + self: grad.index_fill(dim, index, 0) + index: non_differentiable + result: self_t.index_fill(dim, index, 0) + +- name: index_fill.int_Tensor(Tensor self, int dim, Tensor index, Tensor value) -> Tensor + self: grad.index_fill(dim, index, 0) + value: grad.index_select(dim, std::get<0>(at::_unique(index, /*sorted=*/false))).sum() + index: non_differentiable + result: self_t.index_fill(dim, index, value_t) + +- name: index_put(Tensor self, Tensor?[] indices, Tensor values, bool accumulate=False) -> Tensor + self: "accumulate ? grad : grad.index_put(indices, zeros_like(values), false)" + values: grad.index(indices) + result: self_t.index_put(indices, values_t, accumulate) + +- name: _unsafe_index_put(Tensor self, Tensor?[] indices, Tensor values, bool accumulate=False) -> Tensor + self: "accumulate ? grad : at::_unsafe_index_put(grad, indices, zeros_like(values), false)" + values: at::_unsafe_index(grad, indices) + result: at::_unsafe_index_put(self_t, indices, values_t, accumulate) + +- name: _index_put_impl_(Tensor(a!) self, Tensor?[] indices, Tensor values, bool accumulate=False, bool unsafe=False) -> Tensor(a!) + self: "accumulate ? grad : grad.index_put(indices, zeros_like(values), false)" + values: grad.index(indices) + result: at::_index_put_impl_(self_t, indices, values_t, accumulate, unsafe) + +- name: index_select(Tensor self, int dim, Tensor index) -> Tensor + self: index_select_backward_symint(grad, self.sym_sizes(), dim, index) + index: non_differentiable + result: auto_linear + +- name: linalg_inv_ex(Tensor A, *, bool check_errors=False) -> (Tensor inverse, Tensor info) + A: -at::matmul(inverse.mH(), at::matmul(grad, inverse.mH())) + inverse: -at::matmul(at::matmul(inverse, A_t), inverse) + output_differentiability: [True, False] + +- name: linalg_pinv.atol_rtol_tensor(Tensor self, *, Tensor? atol=None, Tensor? rtol=None, bool hermitian=False) -> Tensor + self: pinv_backward(grad, result, self) + result: pinv_jvp(self_p, result, self_t) + +- name: isnan(Tensor self) -> Tensor + self: non_differentiable + +- name: kthvalue(Tensor self, SymInt k, int dim=-1, bool keepdim=False) -> (Tensor values, Tensor indices) + self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) + +- name: ldexp.Tensor(Tensor self, Tensor other) -> Tensor + self: grad * at::pow(2, other).conj() + other: grad * result.conj() * M_LN2 + result: self_t * at::pow(2, other_p) + other_t * result * M_LN2 + +- name: le_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + self: zeros_like(self) + result: self_t.zero_() + +- name: le_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + self: zeros_like(self) + other: zeros_like(other) + result: self_t.zero_() + +- name: lerp.Scalar(Tensor self, Tensor end, Scalar weight) -> Tensor + self: "weight.isComplex() ? grad * (1 - weight.conj().toComplexDouble()) : grad * (1 - weight.toDouble())" + end: grad * weight.conj() + result: at::lerp(self_t, end_t, weight) + +- name: lerp.Tensor(Tensor self, Tensor end, Tensor weight) -> Tensor + self: grad * (1 - weight).conj() + end: grad * weight.conj() + weight: grad * (end - self).conj() + result: at::lerp(self_t, end_t, weight_p) + weight_t * (end_p - self_p) + +- name: lgamma(Tensor self) -> Tensor + self: grad * digamma(self) + result: auto_element_wise + +- name: digamma(Tensor self) -> Tensor + self: grad * polygamma(1, self) + result: auto_element_wise + +- name: polygamma(int n, Tensor self) -> Tensor + self: grad * polygamma(n + 1, self) + result: auto_element_wise + +- name: polygamma_(Tensor(a!) self, int n) -> Tensor(a!) + self: grad * polygamma(n + 1, self) + result: self_t.mul_(polygamma(n + 1, original_self_p)) + +- name: log(Tensor self) -> Tensor + self: grad.div(self.conj()) + result: auto_element_wise + +- name: log10(Tensor self) -> Tensor + self: grad / (self.conj() * 2.3025850929940456) + result: auto_element_wise + +- name: log1p(Tensor self) -> Tensor + self: log1p_backward(grad, self) + result: auto_element_wise + +- name: log2(Tensor self) -> Tensor + self: grad / (self.conj() * 0.6931471805599453) + result: auto_element_wise + +- name: logaddexp(Tensor self, Tensor other) -> Tensor + self: grad / (1 + exp(other - self)).conj() + other: grad / (1 + exp(self - other)).conj() + result: self_t / (1 + exp(other_p - self_p)) + other_t / (1 + exp(self_p - other_p)) + +- name: logaddexp2(Tensor self, Tensor other) -> Tensor + self: grad / (1 + pow(2, other - self)) + other: grad / (1 + pow(2, self - other)) + result: self_t / (1 + pow(2, other_p - self_p)) + other_t / (1 + pow(2, self_p - other_p)) + +# Note [Gradient formula for xlogy at x = 0, y <= 0] +# x * log(y) is not defined at y <= 0, so we cannot even talk about differentiability +# Now, xlogy(0, y) = 0 by definition. +# This does not make it differentiable as it's not defined in a neighbourhood of a point +# (0, y) when y <= 0. +# Now, when a function is non-differentiable, sometimes we return "a relatively sensible value" +# In this case, as per the discussion in https://github.com/pytorch/pytorch/issues/80770, we choose +# this value to be zero, which is the directional derivative along the line {x = 0}. +- name: xlogy.Tensor(Tensor self, Tensor other) -> Tensor + self: at::xlogy(grad, other).masked_fill((self == 0.) & (other <= 0.), 0.) + other: grad * self / other + result: at::xlogy(self_t, other_p).masked_fill((self_p == 0.) & (other_p <= 0.), 0.) + other_t * self_p / other_p + +- name: xlogy.Scalar_Self(Scalar self, Tensor other) -> Tensor + other: grad * self / other + result: auto_element_wise + +- name: xlogy.Scalar_Other(Tensor self, Scalar other) -> Tensor + self: "other.toDouble() > 0. + ? at::xlogy(grad, other) + : at::xlogy(grad, other).masked_fill(self == 0., 0.)" + result: auto_element_wise + +# See Note [Gradient formula for xlogy at x = 0, y <= 0] +# Same here but with y <= -1 +- name: special_xlog1py(Tensor self, Tensor other) -> Tensor + self: at::special_xlog1py(grad, other).masked_fill((self == 0.) & (other <= -1.), 0.) + other: grad * self / (other + 1) + result: at::special_xlog1py(self_t, other_p).masked_fill((self_p == 0.) & (other_p <= -1.), 0.) + other_t * self_p / (other_p + 1) + +- name: special_xlog1py.self_scalar(Scalar self, Tensor other) -> Tensor + other: grad * self / (other + 1) + result: auto_element_wise + +- name: special_xlog1py.other_scalar(Tensor self, Scalar other) -> Tensor + self: "other.toDouble() > -1. + ? at::special_xlog1py(grad, other) + : at::special_xlog1py(grad, other).masked_fill(self == 0., 0.)" + result: auto_element_wise + +- name: special_zeta(Tensor self, Tensor other) -> Tensor + self: not_implemented("zeta") + other: grad * -self * special_zeta(self + 1., other) + +- name: special_zeta.self_scalar(Scalar self, Tensor other) -> Tensor + other: grad * -self * special_zeta(self.toDouble() + 1., other) + +- name: special_zeta.other_scalar(Tensor self, Scalar other) -> Tensor + self: not_implemented("zeta") + +- name: log_normal_(Tensor(a!) self, float mean=1, float std=2, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: logsumexp(Tensor self, int[1] dim, bool keepdim=False) -> Tensor + self: logsumexp_backward(grad, self, result, dim, keepdim) + result: logsumexp_jvp(self_p, self_t, dim, keepdim) + +- name: linalg_lstsq(Tensor self, Tensor b, float? rcond=None, *, str? driver=None) -> (Tensor solution, Tensor residuals, Tensor rank, Tensor singular_values) + self, b: linalg_lstsq_backward(grads[0], grads[1], self, b, solution, grad_input_mask) + solution: linalg_lstsq_solution_jvp(self_p, b_p, self_t, b_t) + residuals: linalg_lstsq_residuals_jvp(self_p, b_p, self_t, b_t, solution, residuals) + output_differentiability: [True, True, False, False] + +- name: lt_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + self: zeros_like(self) + result: self_t.zero_() + +- name: lt_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + self: zeros_like(self) + other: zeros_like(other) + result: self_t.zero_() + +- name: linalg_lu_factor_ex(Tensor A, *, bool pivot=True, bool check_errors=False) -> (Tensor LU, Tensor pivots, Tensor info) + A: lu_factor_ex_backward(grad, LU, pivots, pivot) + LU: lu_factor_ex_jvp(A_t, LU, pivots, pivot) + output_differentiability: [True, False, False] + +- name: linalg_lu(Tensor A, *, bool pivot=True) -> (Tensor P, Tensor L, Tensor U) + A: linalg_lu_backward(grad_L, grad_U, P, L, U, pivot) + L: std::get<0>(linalg_lu_jvp(A_t, P, L, U, pivot)) + U: std::get<1>(linalg_lu_jvp(A_t, P, L, U, pivot)) + output_differentiability: [False, True, True] + +- name: linalg_lu_solve(Tensor LU, Tensor pivots, Tensor B, *, bool left=True, bool adjoint=False) -> Tensor + LU: linalg_lu_solve_LU(grad, LU, pivots, result, left, adjoint) + B: "at::linalg_lu_solve(LU, pivots, grad, left, !adjoint)" + result: linalg_lu_solve_jvp(result, LU_p, pivots, LU_t, B_t, left, adjoint) + +- name: lu_unpack(Tensor LU_data, Tensor LU_pivots, bool unpack_data=True, bool unpack_pivots=True) -> (Tensor P, Tensor L, Tensor U) + LU_data: lu_unpack_backward(grad_L, grad_U, LU_data.sym_size(-2), LU_data.sym_size(-1)) + LU_pivots: non_differentiable + L: "LU_data_t.sym_size(-2) >= LU_data_t.sym_size(-1) ? LU_data_t.tril_symint(-1) : LU_data_t.narrow_symint(-1, 0, LU_data_t.sym_size(-2)).tril_symint(-1)" + U: "LU_data_t.sym_size(-1) >= LU_data_t.sym_size(-2) ? LU_data_t.triu_symint() : LU_data_t.narrow_symint(-2, 0, LU_data_t.sym_size(-1)).triu_symint()" + output_differentiability: [False, True, True] + +- name: masked_fill.Scalar(Tensor self, Tensor mask, Scalar value) -> Tensor + self: grad.masked_fill(mask, 0) + mask: non_differentiable + result: self_t.masked_fill(mask, 0) + +- name: masked_fill.Tensor(Tensor self, Tensor mask, Tensor value) -> Tensor + self: grad.masked_fill(mask, 0) + value: masked_fill_backward(grad, mask) + mask: non_differentiable + result: self_t.masked_fill(mask, value_t) + +- name: masked_scatter(Tensor self, Tensor mask, Tensor source) -> Tensor + self: grad.masked_fill(mask, 0) + source: masked_scatter_backward_symint(grad, mask, source.sym_sizes()) + mask: non_differentiable + result: self_t.masked_scatter(mask, source_t) + +- name: masked_scatter_backward(Tensor grad_output, Tensor mask, SymInt[] sizes) -> Tensor + grad_output: zeros_like(grad_output).masked_scatter(mask, grad) + mask: non_differentiable + result: masked_scatter_backward(grad_output_t, mask, grad_output_t.sizes()) + +- name: masked_select(Tensor self, Tensor mask) -> Tensor + self: masked_select_backward(grad, self, mask) + mask: non_differentiable + result: auto_linear + +- name: linalg_matrix_exp(Tensor self) -> Tensor + self: linalg_matrix_exp_differential(self, grad, /*adjoint*/ true) + result: linalg_matrix_exp_differential(self_p, self_t, /*adjoint*/ false) + +- name: max.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) + self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) + +- name: max(Tensor self) -> Tensor + self: evenly_distribute_backward(grad, self, result) + result: evenly_read_jvp(self_t, self_p, result) + +- name: maximum(Tensor self, Tensor other) -> Tensor + self: at::where(self == other, grad / 2, grad).masked_fill_(self < other, 0) + other: at::where(self == other, grad / 2, grad).masked_fill_(self > other, 0) + result: other_t + at::where(self_p == other_p, at::scalar_tensor(0.5, result.options()), (self_p > other_p).to(result.scalar_type())) * (self_t - other_t) + +- name: fmax(Tensor self, Tensor other) -> Tensor + self: grad.masked_fill((self >= other).logical_or_(other.isnan()).logical_not_(), 0) + other: grad.masked_fill((self >= other).logical_or_(other.isnan()), 0) + result: other_t + (self_p > other_p).logical_or_(other_p.isnan()) * (self_t - other_t) + +- name: mean(Tensor self, *, ScalarType? dtype=None) -> Tensor + dispatch: + Default: + self: grad.expand_symint(self.sym_sizes()) / self.sym_numel() + result: auto_linear + AutogradNestedTensor: + # TODO: replace this with grad.expand_as(self) / self.sym_numel() when that is supported + self: (ones_like(self) * grad) / self.sym_numel() + result: auto_linear + +- name: mean.dim(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + self: mean_backward(grad, self.sym_sizes(), dim, self.sym_numel(), keepdim) + result: auto_linear + +- name: median(Tensor self) -> Tensor + self: evenly_distribute_backward(grad, self, result) + result: evenly_read_jvp(self_t, self_p, result) + +- name: nanmedian(Tensor self) -> Tensor + self: evenly_distribute_backward(grad, self, result) + result: evenly_read_jvp(self_t, self_p, result) + +# This is in theory incorrect in the following case: +# sorted list: [..., a, b, b, ..., b, b, c, ...] with median = b and the value +# | at middle position of the +# | list between two `b`s. E.g., +# | +# ^the middle position +# The gradient exists and is essentially 0 in this case. +# +# In case where the middle position is at the boundary of `b` range, e.g., +# sorted list: [..., a, b, b, ..., b, b, c, ...] +# | +# ^the middle position +# The backward implementation is correct in the sense that it returns the +# subgradient on one side. +- name: median.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) + self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) + +- name: nanmedian.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) + self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) + +- name: min.dim(Tensor self, int dim, bool keepdim=False) -> (Tensor values, Tensor indices) + self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) + +- name: min(Tensor self) -> Tensor + self: evenly_distribute_backward(grad, self, result) + result: evenly_read_jvp(self_t, self_p, result) + +- name: minimum(Tensor self, Tensor other) -> Tensor + self: at::where(self == other, grad / 2, grad).masked_fill_(self > other, 0) + other: at::where(self == other, grad / 2, grad).masked_fill_(self < other, 0) + result: other_t + at::where(self_p == other_p, at::scalar_tensor(0.5, result.options()), (self_p < other_p).to(result.scalar_type())) * (self_t - other_t) + +- name: fmin(Tensor self, Tensor other) -> Tensor + self: grad.masked_fill((self <= other).logical_or_(other.isnan()).logical_not_(), 0) + other: grad.masked_fill((self <= other).logical_or_(other.isnan()), 0) + result: other_t + (self_p <= other_p).logical_or_(other_p.isnan()) * (self_t - other_t) + +- name: amax(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor + self: scale_grad_by_count(restore_reduced_dims(grad, dim, keepdim), restore_reduced_dims(result, dim, keepdim) == self, dim) + result: amaxamin_jvp(self_p, self_t, result, dim, keepdim) + +- name: amin(Tensor self, int[1] dim=[], bool keepdim=False) -> Tensor + self: scale_grad_by_count(restore_reduced_dims(grad, dim, keepdim), restore_reduced_dims(result, dim, keepdim) == self, dim) + result: amaxamin_jvp(self_p, self_t, result, dim, keepdim) + +- name: aminmax(Tensor self, *, int? dim=None, bool keepdim=False) -> (Tensor min, Tensor max) + self: aminmax_backward(self, dim, keepdim, grad_min, grad_max, min, max) + min: aminmax_jvp(self_p, self_t, min, dim, keepdim) + max: aminmax_jvp(self_p, self_t, max, dim, keepdim) + +- name: mm(Tensor self, Tensor mat2) -> Tensor + self: mm_mat1_backward(grad, mat2, self.sym_sizes(), self.sym_strides(), self.layout(), 1) + mat2: mm_mat2_backward(grad, self, mat2.sym_sizes(), mat2.sym_strides(), mat2.layout(), 1) + result: at::mm(self_t, mat2_p) + at::mm(self_p, mat2_t) + +- name: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None, Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor + self: _grouped_mm_mat1_backward(grad, mat2, self.sym_sizes(), self.sym_strides(), self.layout(), offs, 1) + mat2: _grouped_mm_mat2_backward(grad, self, mat2.sym_sizes(), mat2.sym_strides(), mat2.layout(), offs, 1) + +- name: mode(Tensor self, int dim=-1, bool keepdim=False) -> (Tensor values, Tensor indices) + self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), keepdim) + values: gather_with_keepdimed_indices(self_t, dim, indices, keepdim) + +- name: mul.Tensor(Tensor self, Tensor other) -> Tensor + self: mul_tensor_backward(grad, other, self.scalar_type()) + other: mul_tensor_backward(grad, self, other.scalar_type()) + result: other_t * self_p + self_t * other_p + +- name: mul.Scalar(Tensor self, Scalar other) -> Tensor + self: mul_tensor_backward(grad, other, self.scalar_type()) + result: self_t * other + +- name: mv(Tensor self, Tensor vec) -> Tensor + self: grad.ger(vec.conj()) + vec: self.conj().t().mv(grad) + result: mv(self_t, vec_p) + mv(self_p, vec_t) + +- name: mvlgamma(Tensor self, int p) -> Tensor + self: mvlgamma_backward(grad, self, p) + result: auto_element_wise + +- name: nan_to_num(Tensor self, float? nan=None, float? posinf=None, float? neginf=None) -> Tensor + self: grad * at::isfinite(self) + result: auto_element_wise + +- name: native_batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor) + input, weight, bias: "grad.defined() ? native_batch_norm_backward(grad, input, weight, running_mean, running_var, result1, result2, training, eps, grad_input_mask) : std::tuple()" + result0: batch_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, running_mean, running_var, result1, result2, training, eps) + +- name: _native_batch_norm_legit(Tensor input, Tensor? weight, Tensor? bias, Tensor(a!) running_mean, Tensor(b!) running_var, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor) + input, weight, bias: "grad.defined() ? native_batch_norm_backward(grad, input, weight, running_mean, running_var, result1, result2, training, eps, grad_input_mask) : std::tuple()" + result0: batch_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, running_mean, running_var, result1, result2, training, eps) + +- name: _native_batch_norm_legit_no_training(Tensor input, Tensor? weight, Tensor? bias, Tensor running_mean, Tensor running_var, float momentum, float eps) -> (Tensor, Tensor, Tensor) + input, weight, bias: "grad.defined() ? native_batch_norm_backward(grad, input, weight, running_mean, running_var, result1, result2, /*training=*/false, eps, grad_input_mask) : std::tuple()" + result0: batch_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, running_mean, running_var, result1, result2, /*training=*/false, eps) + +- name: _native_batch_norm_legit.no_stats(Tensor input, Tensor? weight, Tensor? bias, bool training, float momentum, float eps) -> (Tensor, Tensor, Tensor) + input, weight, bias: "grad.defined() ? native_batch_norm_backward(grad, input, weight, Tensor(), Tensor(), result1, result2, training, eps, grad_input_mask) : std::tuple()" + result0: batch_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, Tensor(), Tensor(), result1, result2, training, eps) + +- name: native_batch_norm_backward(Tensor grad_out, Tensor input, Tensor? weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_invstd, bool train, float eps, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + input, weight, grad_out: batchnorm_double_backward(input, weight, grads[0], grads[1], grads[2], grad_out, running_mean, running_var, train, eps, save_mean, save_invstd, grad_input_mask) + save_mean: not_implemented("native_batch_norm_backward save_mean") + save_invstd: not_implemented("native_batch_norm_backward save_invstd") + +- name: native_layer_norm(Tensor input, SymInt[] normalized_shape, Tensor? weight, Tensor? bias, float eps) -> (Tensor, Tensor, Tensor) + input, weight, bias: "grad.defined() ? native_layer_norm_backward_symint(grad, input, normalized_shape, result1, result2, weight, bias, grad_input_mask) : std::tuple()" + result0: layer_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, result1, result2, normalized_shape) + +- name: native_layer_norm_backward(Tensor grad_out, Tensor input, SymInt[] normalized_shape, Tensor mean, Tensor rstd, Tensor? weight, Tensor? bias, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + input, weight, grad_out: layer_norm_double_backward(input, weight, grads[0], grads[1], grads[2], grad_out, mean, rstd, normalized_shape, grad_input_mask) + bias: Tensor() + mean: not_implemented("native_layer_norm_backward mean") + rstd: not_implemented("native_layer_norm_backward rstd") + +- name: _fused_rms_norm(Tensor input, int[] normalized_shape, Tensor? weight, float? eps) -> (Tensor, Tensor) + input, weight: "GradMode::is_enabled() || grads[1].defined() ? infinitely_differentiable_native_rms_norm_backward(grads[0], grads[1], input, normalized_shape, result1, weight, grad_input_mask) : (grads[0].defined() ? _fused_rms_norm_backward(grads[0], input, normalized_shape, result1, weight, grad_input_mask) : std::tuple())" + result0: rms_norm_jvp(input_p, input_t, weight_p, weight_t, result1, normalized_shape) + result1: rms_norm_rstd_jvp(input_p, input_t, result1, normalized_shape) + +- name: native_group_norm(Tensor input, Tensor? weight, Tensor? bias, SymInt N, SymInt C, SymInt HxW, int group, float eps) -> (Tensor, Tensor, Tensor) + input, weight, bias: "GradMode::is_enabled() || grads[1].defined() || grads[2].defined() ? infinitely_differentiable_native_group_norm_backward(grads[0], grads[1], grads[2], input, result1, result2, weight, N, C, HxW, group, eps, grad_input_mask) : (grads[0].defined() ? native_group_norm_backward_symint(grads[0].device().is_xpu() ? grads[0] : grads[0].contiguous(grads[0].device().is_cpu() ? input.suggest_memory_format() : c10::MemoryFormat::Contiguous), input.device().is_xpu() ? input : input.contiguous(input.device().is_cpu() ? input.suggest_memory_format() : c10::MemoryFormat::Contiguous), result1, result2, weight, N, C, HxW, group, grad_input_mask) : std::tuple())" + result0: group_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, result1, result2, group) + result1: group_norm_mean_jvp(input_t, result1, group) + result2: group_norm_invstd_jvp(input_p, input_t, result1, result2, group) + +- name: ne_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!) + self: zeros_like(self) + result: self_t.zero_() + +- name: ne_.Tensor(Tensor(a!) self, Tensor other) -> Tensor(a!) + self: zeros_like(self) + other: zeros_like(other) + result: self_t.zero_() + +- name: neg(Tensor self) -> Tensor + self: grad.neg() + result: auto_element_wise + +- name: _batch_norm_with_update(Tensor input, Tensor? weight, Tensor? bias, Tensor(a!) running_mean, Tensor(b!) running_var, float momentum, float eps) -> (Tensor, Tensor, Tensor, Tensor) + input, weight, bias: "grad.defined() ? batch_norm_backward(grad, input, weight, running_mean, running_var, result1, result2, /*update*/true, eps, grad_input_mask, retain_variables ? result3.clone() : result3) : std::tuple()" + result0: batch_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, running_mean, running_var, result1, result2, true, eps) + +- name: _batch_norm_no_update(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, float momentum, float eps) -> (Tensor, Tensor, Tensor, Tensor) + input, weight, bias: "grad.defined() ? batch_norm_backward(grad, input, weight, running_mean, running_var, result1, result2, /*update*/false, eps, grad_input_mask, retain_variables ? result3.clone() : result3) : std::tuple()" + result0: batch_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, running_mean, running_var, result1, result2, false, eps) + +- name: batch_norm_backward(Tensor grad_out, Tensor input, Tensor weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_var, bool update, float eps, bool[3] output_mask, Tensor reserve) -> (Tensor, Tensor, Tensor) + input, weight, grad_out: batchnorm_double_backward(input, weight, grads[0], grads[1], grads[2], grad_out, running_mean, running_var, update, eps, save_mean, save_var, grad_input_mask) + save_mean: not_implemented("batch_norm_backward save_mean") + save_var: not_implemented("batch_norm_backward save_var") + reserve: not_implemented("batch_norm_backward reserve") + +- name: nextafter(Tensor self, Tensor other) -> Tensor + self: not_implemented("nextafter") + other: not_implemented("nextafter") + +- name: norm.Scalar(Tensor self, Scalar p=2) -> Tensor + self: norm_backward(grad, self, p, result) + result: norm_jvp(self_p, self_t, p, result) + +- name: norm.ScalarOpt_dim(Tensor self, Scalar? p, int[1] dim, bool keepdim=False) -> Tensor + self: norm_backward(grad, self, p, result, dim, keepdim) + result: norm_jvp(self_p, self_t, p, result, dim, keepdim) + +- name: norm.ScalarOpt_dtype(Tensor self, Scalar? p, *, ScalarType dtype) -> Tensor + self: norm_backward(grad, self.to(grad.scalar_type()), p, result) + result: norm_jvp(self_p, self_t, p, result) + +- name: norm.ScalarOpt_dim_dtype(Tensor self, Scalar? p, int[1] dim, bool keepdim, *, ScalarType dtype) -> Tensor + self: norm_backward(grad, self.to(grad.scalar_type()), p, result, dim, keepdim) + result: norm_jvp(self_p, self_t, p, result, dim, keepdim) + +- name: linalg_vector_norm(Tensor self, Scalar ord=2, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + self: linalg_vector_norm_backward(grad, self, ord, result, dim, keepdim) + result: linalg_vector_norm_jvp(self_p, self_t, ord, result, dim, keepdim) + +- name: _pdist_forward(Tensor self, float p=2) -> Tensor + self: _pdist_backward(grad, self, p, result) + +- name: _pdist_backward(Tensor grad, Tensor self, float p, Tensor pdist) -> Tensor + grad: not_implemented("_pdist_backward") + self: not_implemented("_pdist_backward") + pdist: not_implemented("_pdist_backward") + +- name: _euclidean_dist(Tensor x1, Tensor x2) -> Tensor + x1, x2: _euclidean_dist_backward(grad, x1, x2, result) + +- name: _cdist_forward(Tensor x1, Tensor x2, float p, int? compute_mode) -> Tensor + x1: _cdist_backward(grad.contiguous(), x1, x2, p, result) + x2: _cdist_backward(grad.mT().contiguous(), x2, x1, p, result.mT().contiguous()) + +- name: _cdist_backward(Tensor grad, Tensor x1, Tensor x2, float p, Tensor cdist) -> Tensor + grad: not_implemented("_cdist_backward") + x1: not_implemented("_cdist_backward") + x2: not_implemented("_cdist_backward") + cdist: not_implemented("_cdist_backward") + +- name: normal_(Tensor(a!) self, float mean=0, float std=1, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: normal.Tensor_float(Tensor mean, float std=1, *, Generator? generator=None) -> Tensor + mean: at::zeros_symint(mean.sym_sizes(), grad.options()) + result: auto_element_wise + +- name: normal.float_Tensor(float mean, Tensor std, *, Generator? generator=None) -> Tensor + std: at::zeros_symint(std.sym_sizes(), grad.options()) + result: auto_element_wise + +- name: normal.Tensor_Tensor(Tensor mean, Tensor std, *, Generator? generator=None) -> Tensor + mean: at::zeros_symint(mean.sym_sizes(), grad.options()) + std: at::zeros_symint(std.sym_sizes(), grad.options()) + result: zeros_like(mean_t) + +- name: linalg_householder_product(Tensor input, Tensor tau) -> Tensor + input, tau: householder_product_backward(grad, result, input, tau) + result: householder_product_jvp(input_t, tau_t, result, input_p, tau_p) + +- name: ormqr(Tensor self, Tensor input2, Tensor input3, bool left=True, bool transpose=False) -> Tensor + self, input2, input3: ormqr_backward(grad, result, self, input2, input3, left, transpose, grad_input_mask) + +- name: permute(Tensor(a) self, int[] dims) -> Tensor(a) + self: permute_backwards(grad, dims) + result: auto_linear + +- name: poisson(Tensor self, Generator? generator=None) -> Tensor + self: zeros_like(self) + result: auto_element_wise + +- name: pow.Tensor_Scalar(Tensor self, Scalar exponent) -> Tensor + self: pow_backward(grad, self, exponent) + result: auto_element_wise + +- name: pow.Tensor_Tensor(Tensor self, Tensor exponent) -> Tensor + self: pow_backward_self(grad, self, exponent) + exponent: pow_backward_exponent(grad, self, exponent, result) + result: (pow_backward_self(self_t.conj(), self_p, exponent_p) + pow_backward_exponent(exponent_t.conj(), self_p, exponent_p, result)).conj() + +- name: pow.Scalar(Scalar self, Tensor exponent) -> Tensor + exponent: pow_backward_exponent(grad, self, exponent, result) + result: auto_element_wise + +- name: prod(Tensor self, *, ScalarType? dtype=None) -> Tensor + self: prod_backward(grad, self.to(grad.scalar_type()), result) + result: (prod_backward(at::ones({}, result.options()).expand_as(result), self_p.to(result.scalar_type()), result) * self_t.conj()).sum().conj() + +- name: prod.dim_int(Tensor self, int dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + self: prod_backward(grad, self.to(grad.scalar_type()), result, dim, keepdim) + result: (prod_backward(at::ones({}, result.options()).expand_as(result), self_p.to(result.scalar_type()), result, dim, keepdim) * self_t.conj()).sum(dim, keepdim).conj() + +- name: put(Tensor self, Tensor index, Tensor source, bool accumulate=False) -> Tensor + self: "accumulate ? grad : grad.put(index, zeros_like(source), false)" + index: non_differentiable + source: grad.take(index).reshape_as(source) + result: self_t.put(index, source_t, accumulate) + +- name: linalg_qr(Tensor A, str mode='reduced') -> (Tensor Q, Tensor R) + A: linalg_qr_backward(grad_Q, grad_R, Q, R, mode) + Q, R: linalg_qr_jvp(A_t, Q, R, mode) + +- name: rad2deg(Tensor self) -> Tensor + self: rad2deg_backward(grad) + result: auto_element_wise + +- name: random_.from(Tensor(a!) self, int from, int? to, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: random_.to(Tensor(a!) self, int to, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: random_(Tensor(a!) self, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: reciprocal(Tensor self) -> Tensor + self: -grad * (result * result).conj() + result: auto_element_wise + +- name: remainder.Scalar(Tensor self, Scalar other) -> Tensor + self: grad + result: auto_element_wise + +- name: remainder.Tensor(Tensor self, Tensor other) -> Tensor + self: grad + other: -grad * self.div(other, /*rounding_mode=*/"floor") + result: self_t - other_t * self_p.div(other_p, /*rounding_mode=*/"floor") + +- name: renorm(Tensor self, Scalar p, int dim, Scalar maxnorm) -> Tensor + self: renorm_backward(grad, self, p, dim, maxnorm) + result: renorm_jvp(self_p, self_t, p, dim, maxnorm) + +- name: repeat(Tensor self, SymInt[] repeats) -> Tensor + self: repeat_backward(grad, repeats, self.sym_sizes()) + result: auto_linear + +- name: special_entr(Tensor self) -> Tensor + self: grad * (-(1 + self.log())) + result: auto_element_wise + +- name: special_ndtri(Tensor self) -> Tensor + self: grad * std::sqrt(2 * M_PI) * (result.square() / 2).exp() + result: auto_element_wise + +- name: special_log_ndtr(Tensor self) -> Tensor + self: grad / std::sqrt(2 * M_PI) * (result + self.pow(2) / 2).neg().exp() + result: auto_element_wise + +# [Note: Sometimes view derivatives] +# The following situation applies to other operations as well. +# TODO: This note is only referenced by to_dense and to_sparse*. Make +# this more generic if it's been referenced more than once. +# +# DO NOT define a backward for reshape! +# reshape is special in that it sometimes returns a view, and sometimes not. +# Defining a backward will make codegen spit out the forward call as +# as_variable(baseType->reshape(self)), +# making it impossible (hard) to detect when it is actually a view. +# - name: reshape(Tensor self, IntArrayRef shape) + +- name: _reshape_alias(Tensor(a) self, SymInt[] size, SymInt[] stride) -> Tensor(a) + self: grad.reshape_symint(self.sym_sizes()) + result: auto_linear + +- name: round(Tensor self) -> Tensor + self: zeros_like(grad) + result: auto_element_wise + +- name: round.decimals(Tensor self, *, int decimals) -> Tensor + self: zeros_like(grad) + result: auto_element_wise + +- name: rsqrt(Tensor self) -> Tensor + self: -0.5 * grad * result.pow(3).conj() + result: auto_element_wise + +- name: scatter.src(Tensor self, int dim, Tensor index, Tensor src) -> Tensor + self: grad.scatter(dim, index, 0) + index: non_differentiable + src: grad.gather(dim, index) + result: self_t.scatter(dim, index, src_t) + +- name: scatter.value(Tensor self, int dim, Tensor index, Scalar value) -> Tensor + self: grad.scatter(dim, index, 0) + index: non_differentiable + result: self_t.scatter(dim, index, 0) + +- name: scatter_add(Tensor self, int dim, Tensor index, Tensor src) -> Tensor + self: grad + index: non_differentiable + src: grad.gather(dim, index) + result: scatter_add(self_t, dim, index, src_t) + +- name: select.int(Tensor(a) self, int dim, SymInt index) -> Tensor(a) + dispatch: + Default: + self: select_backward_symint(grad, self.sym_sizes(), dim, index) + result: auto_linear + AutogradNestedTensor: + self: _nested_select_backward_symint(grad, self, dim, index) + +- name: select_backward(Tensor grad_output, SymInt[] input_sizes, int dim, SymInt index) -> Tensor + grad_output: grad.select_symint(dim, index) + result: auto_linear + +- name: sigmoid(Tensor self) -> Tensor + self: sigmoid_backward(grad, result) + result: auto_element_wise + +- name: logit(Tensor self, float? eps=None) -> Tensor + self: "GradMode::is_enabled() ? infinitely_differentiable_logit_backward(grad, self, eps) : logit_backward(grad, self, eps)" + result: auto_element_wise + +- name: sign(Tensor self) -> Tensor + self: zeros_like(grad) + result: auto_element_wise + +- name: sgn(Tensor self) -> Tensor + self: sgn_backward(self, grad, result) + # Cannot use auto_element_wise here because the Jacobian is *not* Hermitian (in fact, it is symmetric) + # The function is not holomorphic, so there's no reason for its Jacobian to be Hermitian + # auto_element_wise has a name that's a bit deceiving in the complex case + result: sgn_backward(self_p, self_t, result) + +- name: sin(Tensor self) -> Tensor + self: grad * self.cos().conj() + result: auto_element_wise + +- name: sinc(Tensor self) -> Tensor + self: sinc_backward(grad, self) + result: auto_element_wise + +- name: sinh(Tensor self) -> Tensor + self: grad * self.cosh().conj() + result: auto_element_wise + +- name: slice.Tensor(Tensor(a) self, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor(a) + self: slice_backward_wrapper(grad, self.sym_sizes(), dim, start, end, step) + result: auto_linear + +- name: slice_backward(Tensor grad_output, SymInt[] input_sizes, int dim, SymInt start, SymInt end, SymInt step) -> Tensor + grad_output: grad.slice_symint(dim, start, end, step) + result: auto_linear + +- name: slice_inverse(Tensor(a) self, Tensor src, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor(a) + self: grad.slice_symint(dim, start, end, step) + src: slice_scatter_symint(grad, zeros_like(self), dim, start, end, step) + result: auto_linear + +- name: slice_scatter(Tensor self, Tensor src, int dim=0, SymInt? start=None, SymInt? end=None, SymInt step=1) -> Tensor + self: slice_scatter_symint(grad, zeros_like(src), dim, start, end, step) + src: grad.slice_symint(dim, start, end, step) + result: auto_linear + +- name: select_scatter(Tensor self, Tensor src, int dim, SymInt index) -> Tensor + self: select_scatter_symint(grad, zeros_like(src), dim, index) + src: grad.select_symint(dim, index) + result: auto_linear + +- name: diagonal_scatter(Tensor self, Tensor src, int offset=0, int dim1=0, int dim2=1) -> Tensor + self: diagonal_scatter(grad, zeros_like(src), offset, dim1, dim2) + src: grad.diagonal(offset, dim1, dim2) + result: auto_linear + +- name: as_strided_scatter(Tensor self, Tensor src, SymInt[] size, SymInt[] stride, SymInt? storage_offset=None) -> Tensor + self: as_strided_scatter_backward(grad, TensorGeometry(self), TensorGeometry(src), size, stride, storage_offset) + # See Note [as_strided_scatter backward support] + src: grad.contiguous().as_strided_symint(size, stride, storage_offset) + result: auto_linear + +- name: _linalg_solve_ex(Tensor A, Tensor B, *, bool left=True, bool check_errors=False) -> (Tensor result, Tensor LU, Tensor pivots, Tensor info) + A, B: linalg_solve_backward(grad, result, A, LU, pivots, left, grad_input_mask[1]) + result: "linalg_solve_jvp(A_t, B_t, result, LU, pivots, left)" + output_differentiability: [True, False, False, False] # LU is an auxiliary tensor not exposed to the user + +- name: sort(Tensor self, int dim=-1, bool descending=False) -> (Tensor values, Tensor indices) + self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), true) + output_differentiability: [True, False] + values: gather_with_keepdimed_indices(self_t, dim, indices, true) + +- name: sort.stable(Tensor self, *, bool? stable, int dim=-1, bool descending=False) -> (Tensor values, Tensor indices) + self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), true) + output_differentiability: [True, False] + values: gather_with_keepdimed_indices(self_t, dim, indices, true) + +- name: split.Tensor(Tensor(a -> *) self, SymInt split_size, int dim=0) -> Tensor(a)[] + self: split_backward(grads, split_size, dim, self.sym_sizes(), self.options()) + result: auto_linear + +- name: unsafe_split.Tensor(Tensor self, SymInt split_size, int dim=0) -> Tensor[] + self: split_backward(grads, split_size, dim, self.sym_sizes(), self.options()) + result: auto_linear + +- name: split_with_sizes(Tensor(a -> *) self, SymInt[] split_sizes, int dim=0) -> Tensor(a)[] + dispatch: + Default: + self: split_with_sizes_backward(grads, split_sizes, dim, self.sym_sizes(), self.options()) + result: auto_linear + AutogradNestedTensor: + self: _nested_split_with_sizes_backward(grads, split_sizes, dim, at::native::get_nested_tensor_impl(self)->get_nested_sizes(), self.options()) + +- name: unsafe_split_with_sizes(Tensor self, SymInt[] split_sizes, int dim=0) -> Tensor[] + self: split_with_sizes_backward(grads, split_sizes, dim, self.sym_sizes(), self.options()) + result: auto_linear + +- name: sqrt(Tensor self) -> Tensor + self: grad / (2 * result.conj()) + result: auto_element_wise + +- name: squeeze(Tensor(a) self) -> Tensor(a) + self: unsqueeze_to(grad, self.sym_sizes()) + result: auto_linear + +- name: squeeze.dim(Tensor(a) self, int dim) -> Tensor(a) + dispatch: + Default: + self: unsqueeze_to(grad, dim, self.sym_sizes()) + result: auto_linear + AutogradNestedTensor: + self: grad.unsqueeze(dim) + +- name: squeeze.dims(Tensor(a) self, int[] dim) -> Tensor(a) + dispatch: + Default: + self: unsqueeze_to(grad, dim, self.sym_sizes()) + result: auto_linear + AutogradNestedTensor: + self: unsqueeze_multiple(grad, dim, self.dim()) + +- name: squeeze_(Tensor(a!) self) -> Tensor(a!) + self: unsqueeze_to(grad, self.sym_sizes()) + result: auto_linear + +- name: squeeze_.dim(Tensor(a!) self, int dim) -> Tensor(a!) + self: unsqueeze_to(grad, dim, self.sym_sizes()) + result: auto_linear + +- name: squeeze_.dims(Tensor(a!) self, int[] dim) -> Tensor(a!) + self: unsqueeze_to(grad, dim, self.sym_sizes()) + result: auto_linear + +- name: std.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> Tensor + self: std_backward(result, grad, self, dim, correction, keepdim) + # pointwise (variance) + sum + sqrt + result: (at::real(var_backward(self_t.conj(), self_p, dim, correction, true).sum(dim.value_or(IntArrayRef({})), keepdim)) / (2. * result)).masked_fill_(result == 0, 0) + +- name: std_mean.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> (Tensor, Tensor) + self: std_mean_backward(grads[0], grads[1], self, result0, dim, correction, keepdim) + result0: (at::real(var_backward(self_t.conj(), self_p, dim, correction, true).sum(dim.value_or(IntArrayRef({})), keepdim)) / (2. * result0)).masked_fill_(result0 == 0, 0) + # linear + result1: mean(self_t, dim.value_or(IntArrayRef({})), keepdim) + +- name: sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor + self: handle_r_to_c(self.scalar_type(), grad) + other: handle_r_to_c(other.scalar_type(), maybe_multiply(-grad, alpha.conj())) + result: self_t - maybe_multiply(other_t, alpha) + +- name: sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor + self: handle_r_to_c(self.scalar_type(), grad) + result: auto_element_wise + +- name: rsub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor + self: handle_r_to_c(self.scalar_type(), maybe_multiply(-grad, alpha.conj())) + other: handle_r_to_c(other.scalar_type(), grad) + result: -maybe_multiply(self_t, alpha) + other_t + +- name: rsub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor + self: handle_r_to_c(self.scalar_type(), maybe_multiply(-grad, alpha.conj())) + result: auto_element_wise + +- name: sum(Tensor self, *, ScalarType? dtype=None) -> Tensor + dispatch: + Default: + self: grad.expand_symint(self.sym_sizes()) + result: auto_linear + AutogradNestedTensor: + # TODO: replace this with grad.expand_as(self) when that is supported + self: ones_like(self) * grad + result: auto_linear + +- name: sum.dim_IntList(Tensor self, int[1]? dim, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + dispatch: + Default: + self: sum_backward(grad, self.sym_sizes(), dim, keepdim) + result: auto_linear + AutogradNestedTensor: + # TODO: replace this function once semantics for nested tensor expand have been settled on + self: _nested_sum_backward(grad, self, dim, keepdim) + +- name: nansum(Tensor self, int[1]? dim=None, bool keepdim=False, *, ScalarType? dtype=None) -> Tensor + self: nansum_backward(grad.to(self.scalar_type()), self, dim, keepdim) + result: at::where(self_p.isnan(), 0, self_t).sum(dim, keepdim, dtype) + +# We never call _linalg_svd with compute_uv=False in an autograd context, so we don't even consider it here +- name: _linalg_svd(Tensor A, bool full_matrices=False, bool compute_uv=True, *, str? driver=None) -> (Tensor U, Tensor S, Tensor Vh) + A: "svd_backward(full_matrices && grad_U.defined() ? grad_U.narrow_symint(-1, 0, S.sym_size(-1)) : grad_U, + grad_S, + full_matrices && grad_Vh.defined() ? grad_Vh.narrow_symint(-2, 0, S.sym_size(-1)) : grad_Vh, + full_matrices ? U.narrow_symint(-1, 0, S.sym_size(-1)) : U, + S, + full_matrices ? Vh.narrow_symint(-2, 0, S.sym_size(-1)) : Vh)" + U, S, Vh: linalg_svd_jvp(A_t, U, S, Vh, full_matrices) + +- name: _linalg_eigh(Tensor A, str UPLO="L", bool compute_v=True) -> (Tensor eigenvalues, Tensor eigenvectors) + A: linalg_eig_backward(grads[0], grads[1], eigenvalues, eigenvectors, /*is_hermitian=*/true) + eigenvalues, eigenvectors: linalg_eig_jvp(A_t, eigenvalues, eigenvectors, /*is_hermitian=*/true) + +- name: linalg_eig(Tensor self) -> (Tensor eigenvalues, Tensor eigenvectors) + self: handle_r_to_c(self.scalar_type(), linalg_eig_backward(grads[0], grads[1], eigenvalues, eigenvectors, /*is_hermitian=*/false)) + eigenvalues, eigenvectors: linalg_eig_jvp(self_t, eigenvalues, eigenvectors, /*is_hermitian=*/false) + +- name: t(Tensor(a) self) -> Tensor(a) + self: grad.t() + result: auto_linear + +- name: t_(Tensor(a!) self) -> Tensor(a!) + self: grad.t() + result: auto_linear + +- name: one_hot(Tensor self, int num_classes=-1) -> Tensor + self: non_differentiable + +- name: flip(Tensor self, int[] dims) -> Tensor + self: grad.flip(dims) + result: auto_linear + +- name: roll(Tensor self, SymInt[1] shifts, int[1] dims=[]) -> Tensor + self: grad.roll_symint(fmap(reverse_list_symint(shifts), [](c10::SymInt i){return -i;}), reverse_list(dims)) + result: auto_linear + +- name: rot90(Tensor self, int k=1, int[] dims=[0,1]) -> Tensor + self: grad.rot90(-k, dims) + result: auto_linear + +- name: take(Tensor self, Tensor index) -> Tensor + self: take_backward(grad, self, index) + index: non_differentiable + result: auto_linear + +- name: tan(Tensor self) -> Tensor + self: grad * (1 + result.pow(2)).conj() + result: auto_element_wise + +- name: tanh(Tensor self) -> Tensor + self: tanh_backward(grad, result) + result: auto_element_wise + +- name: topk(Tensor self, SymInt k, int dim=-1, bool largest=True, bool sorted=True) -> (Tensor values, Tensor indices) + self: value_selecting_reduction_backward_symint(grad, dim, indices, self.sym_sizes(), true) + output_differentiability: [True, False] + values: gather(self_t, dim, indices) + +- name: trace(Tensor self) -> Tensor + self: trace_backward_symint(grad, self.sym_sizes()) + result: auto_linear + +- name: transpose.int(Tensor(a) self, int dim0, int dim1) -> Tensor(a) + self: grad.transpose(dim0, dim1) + result: auto_linear + +- name: transpose_(Tensor(a!) self, int dim0, int dim1) -> Tensor(a!) + self: grad.transpose(dim0, dim1) + result: auto_linear + +- name: triangular_solve(Tensor self, Tensor A, bool upper=True, bool transpose=False, bool unitriangular=False) -> (Tensor solution, Tensor cloned_coefficient) + self, A: triangular_solve_backward(grad_solution, grad_cloned_coefficient, self, A, solution, upper, transpose, unitriangular, grad_input_mask) + solution: triangular_solve_jvp(solution, A_p, A_t, self_t, upper, transpose, unitriangular) + cloned_coefficient: A_t + +- name: linalg_solve_triangular(Tensor self, Tensor B, *, bool upper, bool left=True, bool unitriangular=False) -> Tensor + self, B: linalg_solve_triangular_backward(grad, self, result, upper, left, unitriangular, grad_input_mask) + result: linalg_solve_triangular_forward_AD(self_t, B_t, self_p, result, upper, left, unitriangular) + +- name: tril(Tensor self, SymInt diagonal=0) -> Tensor + self: grad.tril_symint(diagonal) + result: auto_linear + +- name: triu(Tensor self, SymInt diagonal=0) -> Tensor + self: grad.triu_symint(diagonal) + result: auto_linear + +- name: trunc(Tensor self) -> Tensor + self: zeros_like(grad) + result: auto_element_wise + +- name: hash_tensor(Tensor self, int[1] dim=[], *, bool keepdim=False, int mode=0) -> Tensor + output_differentiability: [False] + +# DO NOT define a backward for to_dense +# See [Note: Sometimes view derivatives] +# - name: to_dense(Tensor self, ScalarType? dtype=None, *, bool? masked_grad=None) -> Tensor +# +- name: _to_dense(Tensor self, ScalarType? dtype=None, bool? masked_grad=None) -> Tensor + self: to_dense_backward(grad, self, masked_grad) + +# DO NOT define a backward for to_sparse.sparse_dim +# See [Note: Sometimes view derivatives] +# - name: to_sparse.sparse_dim(Tensor self, int sparse_dim) -> Tensor +# +- name: _to_sparse.sparse_dim(Tensor self, int sparse_dim) -> Tensor + self: to_sparse_backward(grad, self.layout(), self.sym_blocksize()) + +# DO NOT define a backward for to_sparse +# See [Note: Sometimes view derivatives] +# - name: to_sparse(Tensor self, *, Layout? layout=None, int[2]? blocksize=None, int? dense_dim=None) -> Tensor +# +- name: _to_sparse(Tensor self, *, Layout? layout=None, int[2]? blocksize=None, int? dense_dim=None) -> Tensor + self: to_sparse_backward(grad, self.layout(), self.sym_blocksize()) + +# DO NOT define a backward for to_sparse_csr +# See [Note: Sometimes view derivatives] +# - name: to_sparse_csr(Tensor self, int? dense_dim=None) -> Tensor +# +- name: _to_sparse_csr(Tensor self, int? dense_dim=None) -> Tensor + self: to_sparse_backward(grad, self.layout(), self.sym_blocksize()) + +# DO NOT define a backward for to_sparse_csc +# See [Note: Sometimes view derivatives] +# - name: to_sparse_csc(Tensor self, int? dense_dim=None) -> Tensor +# +- name: _to_sparse_csc(Tensor self, int? dense_dim=None) -> Tensor + self: to_sparse_backward(grad, self.layout(), self.sym_blocksize()) + +# DO NOT define a backward for to_sparse_bsr +# See [Note: Sometimes view derivatives] +# - name: to_sparse_bsr(Tensor self, int[2] blocksize, int? dense_dim=None) -> Tensor +# +- name: _to_sparse_bsr(Tensor self, int[2] blocksize, int? dense_dim=None) -> Tensor + self: to_sparse_backward(grad, self.layout(), self.sym_blocksize()) + +# DO NOT define a backward for to_sparse_bsc +# See [Note: Sometimes view derivatives] +# - name: to_sparse_bsc(Tensor self, int[2] blocksize, int? dense_dim=None) -> Tensor +# +- name: _to_sparse_bsc(Tensor self, int[2] blocksize, int? dense_dim=None) -> Tensor + self: to_sparse_backward(grad, self.layout(), self.sym_blocksize()) + +- name: to_mkldnn(Tensor self, ScalarType? dtype=None) -> Tensor + self: to_mkldnn_backward(grad, self) + +- name: unfold(Tensor(a) self, int dimension, int size, int step) -> Tensor(a) + self: unfold_backward_symint(grad, self.sym_sizes(), dimension, size, step) + result: auto_linear + +- name: unfold_backward(Tensor grad_in, SymInt[] input_sizes, int dim, int size, int step) -> Tensor + grad_in: grad.unfold(dim, size, step) + result: auto_linear + +- name: uniform_(Tensor(a!) self, float from=0, float to=1, *, Generator? generator=None) -> Tensor(a!) + self: zeros_like(grad) + result: self_t.zero_() + +- name: _unique(Tensor self, bool sorted=True, bool return_inverse=False) -> (Tensor, Tensor) + output_differentiability: [True, False] + self: not_implemented("_unique") + +- name: unique_dim(Tensor self, int dim, bool sorted=True, bool return_inverse=False, bool return_counts=False) -> (Tensor, Tensor, Tensor) + output_differentiability: [True, False, False] + self: not_implemented("unique_dim") + +- name: unique_consecutive(Tensor self, bool return_inverse=False, bool return_counts=False, int? dim=None) -> (Tensor, Tensor, Tensor) + output_differentiability: [True, False, False] + self: not_implemented("unique_consecutive") + +- name: unique_dim_consecutive(Tensor self, int dim, bool return_inverse=False, bool return_counts=False) -> (Tensor, Tensor, Tensor) + output_differentiability: [True, False, False] + self: not_implemented("unique_dim_consecutive") + +- name: _unique2(Tensor self, bool sorted=True, bool return_inverse=False, bool return_counts=False) -> (Tensor, Tensor, Tensor) + output_differentiability: [True, False, False] + self: not_implemented("_unique2") + +- name: _unsafe_view(Tensor self, SymInt[] size) -> Tensor + self: grad.reshape_symint(self.sym_sizes()) + result: auto_linear + +- name: lift(Tensor self) -> Tensor + self: grad + result: auto_linear + +- name: lift_fresh(Tensor(a) self) -> Tensor(a) + self: grad + result: auto_linear + +- name: unsqueeze(Tensor(a) self, int dim) -> Tensor(a) + self: grad.squeeze(dim) + result: auto_linear + +- name: unsqueeze_(Tensor(a!) self, int dim) -> Tensor(a!) + self: grad.squeeze(dim) + result: auto_linear + +- name: var.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> Tensor + self: var_backward(grad, self, dim, correction, keepdim) + # pointwise + sum + result: at::real(var_backward(self_t.conj(), self_p, dim, correction, true).sum(dim.value_or(IntArrayRef({})), keepdim)) + +- name: var_mean.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False) -> (Tensor, Tensor) + self: var_mean_backward(grads[0], grads[1], self, dim, correction, keepdim) + result0: at::real(var_backward(self_t.conj(), self_p, dim, correction, true).sum(dim.value_or(IntArrayRef({})), keepdim)) + # linear + result1: mean(self_t, dim.value_or(IntArrayRef({})), keepdim) + +- name: view(Tensor(a) self, SymInt[] size) -> Tensor(a) + dispatch: + Default: + self: grad.reshape_symint(self.sym_sizes()) + result: auto_linear + AutogradNestedTensor: + self: grad.reshape_as(self) + result: auto_linear + +- name: view.dtype(Tensor(a) self, ScalarType dtype) -> Tensor(a) + output_differentiability: [False] + +- name: view_as_real(Tensor(a) self) -> Tensor(a) + self: at::view_as_complex(grad.contiguous()) # gx0 + 1j * gx1 + result: at::view_as_real(self_t) + +- name: view_as_complex(Tensor(a) self) -> Tensor(a) + self: at::view_as_real(grad.contiguous().resolve_conj()) # [gx, gy] + result: at::view_as_complex(self_t) + +- name: where.self(Tensor condition, Tensor self, Tensor other) -> Tensor + condition: non_differentiable + self: where(condition, grad, 0) + other: where(condition, 0, grad) + result: where(condition, self_t, other_t) + +# weight_norm_cuda_interface_backward does not have an explicitly defined derivative, so if we do happen +# to be running backward with create_graph=True, fall back to a backward function that uses +# differentiable ops. +- name: _weight_norm_interface(Tensor v, Tensor g, int dim=0) -> (Tensor, Tensor) + v, g: "grad.defined() ? (GradMode::is_enabled() ? _weight_norm_differentiable_backward(grad.contiguous(), v, g, result1, dim) : _weight_norm_interface_backward(grad.contiguous(), v, g, result1, dim)) : std::tuple()" + +- name: zero_(Tensor(a!) self) -> Tensor(a!) + self: zeros_like(grad) + result: auto_linear + +- name: sparse_mask(Tensor self, Tensor mask) -> Tensor + self: sparse_mask_backward(grad, mask, self.layout()) + mask: non_differentiable + +- name: _sparse_coo_tensor_with_dims_and_tensors(int sparse_dim, int dense_dim, SymInt[] size, Tensor indices, Tensor values, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False, bool? is_coalesced=None) -> Tensor + indices: non_differentiable + values: grad.sparse_mask(result)._values() + +- name: sparse_compressed_tensor.comp_plain_value_size(Tensor compressed_indices, Tensor plain_indices, Tensor values, SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=False) -> Tensor + compressed_indices: non_differentiable + plain_indices: non_differentiable + # TODO: remove to_dense after gh-107381 is fixed + values: grad.to_dense().sparse_mask(result).values() + +- name: _sparse_sum.dim(Tensor self, int[1] dim) -> Tensor + self: at::_sparse_sum_backward(grad, self, dim) + +- name: _standard_gamma(Tensor self, Generator? generator=None) -> Tensor + self: grad * _standard_gamma_grad(self, result) + +- name: _standard_gamma_grad(Tensor self, Tensor output) -> Tensor + self: not_implemented("_standard_gamma_grad") + +- name: values(Tensor(a) self) -> Tensor(a) + dispatch: + Default: + self: values_backward(grad, self) + AutogradNestedTensor: + self: at::_nested_view_from_buffer(grad.contiguous(), self._nested_tensor_size(), self._nested_tensor_strides(), self._nested_tensor_storage_offsets()) + +# Why is _values() not differentiable? +# See NOTE [ Sparse: autograd and API ] +- name: _values(Tensor(a) self) -> Tensor(a) + output_differentiability: [False] + +# NN +- name: _trilinear(Tensor i1, Tensor i2, Tensor i3, int[] expand1, int[] expand2, int[] expand3, int[] sumdim, int unroll_dim=1) -> Tensor + i1, i2, i3: "_trilinear_backward(grad, + wrap_opt_if(i1, grad_input_mask[1] || grad_input_mask[2]), + wrap_opt_if(i2, grad_input_mask[0] || grad_input_mask[2]), + wrap_opt_if(i3, grad_input_mask[0] || grad_input_mask[1]), + expand1, expand2, expand3, sumdim, grad_input_mask)" + result: "_trilinear(i1_t, i2_p, i3_p, expand1, expand2, expand3, sumdim, unroll_dim) + + _trilinear(i1_p, i2_t, i3_p, expand1, expand2, expand3, sumdim, unroll_dim) + + _trilinear(i1_p, i2_p, i3_t, expand1, expand2, expand3, sumdim, unroll_dim)" + +- name: constant_pad_nd(Tensor self, SymInt[] pad, Scalar value=0) -> Tensor + self: constant_pad_nd_backward(grad, pad) + result: constant_pad_nd_symint(self_t, pad, 0) + +- name: binary_cross_entropy(Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean) -> Tensor + self: binary_cross_entropy_backward(grad, self, target, weight, reduction) + target: binary_cross_entropy_target_backward(grad, self, target, weight, reduction) + result: "apply_loss_reduction( + binary_cross_entropy_backward(self_t, self_p, target_p, weight, at::Reduction::None) + + binary_cross_entropy_target_backward(target_t, self_p, target_p, weight, at::Reduction::None), + reduction)" + +- name: binary_cross_entropy_backward(Tensor grad_output, Tensor self, Tensor target, Tensor? weight=None, int reduction=Mean) -> Tensor + self: binary_cross_entropy_double_backward(grad_output, grad, self, target, weight, reduction) + target: binary_cross_entropy_double_backward_target(grad, grad_output, self, target, weight, reduction) + grad_output: binary_cross_entropy_double_backward_grad_output(grad, self, target, weight, reduction) + result: " binary_cross_entropy_double_backward(grad_output_p, self_t, self_p, target_p, weight, reduction) + + binary_cross_entropy_double_backward_target(target_t, grad_output_p, self_p, target_p, weight, reduction) + + binary_cross_entropy_double_backward_grad_output(grad_output_t, self_p, target_p, weight, reduction)" + +- name: binary_cross_entropy_with_logits(Tensor self, Tensor target, Tensor? weight=None, Tensor? pos_weight=None, int reduction=Mean) -> Tensor + self: binary_cross_entropy_with_logits_backward(grad, self, target, weight, pos_weight, reduction) + target: binary_cross_entropy_with_logits_target_backward(grad, self, target, weight, pos_weight, reduction) + result: "apply_loss_reduction( + binary_cross_entropy_with_logits_backward(self_t, self_p, target_p, weight, pos_weight, at::Reduction::None) + + binary_cross_entropy_with_logits_target_backward(target_t, self_p, target_p, weight, pos_weight, at::Reduction::None), + reduction)" + +- name: embedding(Tensor weight, Tensor indices, SymInt padding_idx=-1, bool scale_grad_by_freq=False, bool sparse=False) -> Tensor + indices: non_differentiable + weight: embedding_backward_symint(grad, indices, weight.sym_size(0), padding_idx, scale_grad_by_freq, sparse) + result: auto_linear + +- name: embedding_dense_backward(Tensor grad_output, Tensor indices, SymInt num_weights, SymInt padding_idx, bool scale_grad_by_freq) -> Tensor + grad_output: embedding_dense_double_backward_symint(grad, indices, padding_idx) + indices: non_differentiable + result: auto_linear + +- name: _embedding_bag(Tensor weight, Tensor indices, Tensor offsets, bool scale_grad_by_freq=False, int mode=0, bool sparse=False, Tensor? per_sample_weights=None, bool include_last_offset=False, int padding_idx=-1) -> (Tensor, Tensor, Tensor, Tensor) + indices: non_differentiable + offsets: non_differentiable + weight: _embedding_bag_backward_symint(grad, indices, offsets, result1, result2, result3, weight.sym_size(0), scale_grad_by_freq, mode, sparse, per_sample_weights, padding_idx) + per_sample_weights: _embedding_bag_per_sample_weights_backward(grad, weight, indices, offsets, result1, mode, padding_idx) + +- name: _embedding_bag_backward(Tensor grad, Tensor indices, Tensor offsets, Tensor offset2bag, Tensor bag_size, Tensor maximum_indices, SymInt num_weights, bool scale_grad_by_freq, int mode, bool sparse, Tensor? per_sample_weights, int padding_idx=-1) -> Tensor + grad: not_implemented("_embedding_bag_backward") + indices: non_differentiable + offsets: non_differentiable + offset2bag: non_differentiable + bag_size: non_differentiable + maximum_indices: non_differentiable + per_sample_weights: not_implemented("_embedding_bag_backward") + +- name: _embedding_bag_dense_backward(Tensor grad, Tensor indices, Tensor offset2bag, Tensor bag_size, Tensor maximum_indices, SymInt num_weights, bool scale_grad_by_freq, int mode, Tensor? per_sample_weights, int padding_idx=-1) -> Tensor + grad: not_implemented("_embedding_bag_dense_backward") + indices: non_differentiable + offset2bag: non_differentiable + bag_size: non_differentiable + maximum_indices: non_differentiable + per_sample_weights: not_implemented("_embedding_bag_dense_backward") + +- name: embedding_renorm_(Tensor(a!) self, Tensor indices, float max_norm, float norm_type) -> Tensor(a!) + indices: non_differentiable + self: not_implemented("embedding_renorm") + +- name: mse_loss(Tensor self, Tensor target, int reduction=Mean) -> Tensor + self: mse_loss_backward(grad, self, target, reduction) + target: mse_loss_backward(grad, target, self, reduction) + result: apply_loss_reduction(mse_loss_backward(self_t.conj(), self_p, target_p, at::Reduction::None).conj() + mse_loss_backward(target_t.conj(), target_p, self_p, at::Reduction::None).conj(), reduction) + +- name: multi_margin_loss(Tensor self, Tensor target, Scalar p=1, Scalar margin=1, Tensor? weight=None, int reduction=Mean) -> Tensor + self: multi_margin_loss_backward(grad, self, target, p, margin, weight, reduction) + target: non_differentiable + +- name: multilabel_margin_loss_forward(Tensor self, Tensor target, int reduction) -> (Tensor output, Tensor is_target) + self: multilabel_margin_loss_backward(grad, self, target, reduction, is_target) + target: non_differentiable + +- name: nll_loss_forward(Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index) -> (Tensor output, Tensor total_weight) + self: nll_loss_backward_symint(grad, self, target, weight, reduction, ignore_index, total_weight) + target: non_differentiable + output: std::get<0>(nll_loss_forward_symint(self_t, target, weight, reduction, ignore_index)) + +- name: nll_loss2d_forward(Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index) -> (Tensor output, Tensor total_weight) + self: nll_loss2d_backward_symint(grad, self, target, weight, reduction, ignore_index, total_weight) + target: non_differentiable + output: std::get<0>(nll_loss2d_forward_symint(self_t, target, weight, reduction, ignore_index)) + +- name: smooth_l1_loss(Tensor self, Tensor target, int reduction=Mean, float beta=1.0) -> Tensor + self: smooth_l1_loss_backward(grad, self, target, reduction, beta) + target: smooth_l1_loss_backward(grad, target, self, reduction, beta) + result: apply_loss_reduction(smooth_l1_loss_backward(self_t.conj(), self_p, target_p, at::Reduction::None, beta).conj() + smooth_l1_loss_backward(target_t.conj(), target_p, self_p, at::Reduction::None, beta).conj(), reduction) + +- name: huber_loss(Tensor self, Tensor target, int reduction=Mean, float delta=1.0) -> Tensor + self: huber_loss_backward(grad, self, target, reduction, delta) + target: huber_loss_backward(grad, target, self, reduction, delta) + result: apply_loss_reduction(huber_loss_backward(self_t.conj(), self_p, target_p, at::Reduction::None, delta).conj() + huber_loss_backward(target_t.conj(), target_p, self_p, at::Reduction::None, delta).conj(), reduction) + +- name: soft_margin_loss(Tensor self, Tensor target, int reduction=Mean) -> Tensor + self: soft_margin_loss_backward(grad, self, target, reduction) + result: apply_loss_reduction(soft_margin_loss_backward(self_t.conj(), self_p, target, at::Reduction::None).conj(), reduction) + +- name: relu(Tensor self) -> Tensor + self: threshold_backward(grad, result, 0) + result: auto_element_wise + +- name: silu(Tensor self) -> Tensor + self: "GradMode::is_enabled() ? infinitely_differentiable_silu_backward(grad, self) : silu_backward(grad, self)" + result: auto_element_wise + +- name: mish(Tensor self) -> Tensor + self: "GradMode::is_enabled() ? infinitely_differentiable_mish_backward(grad, self) : mish_backward(grad, self)" + result: auto_element_wise + +- name: elu(Tensor self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> Tensor + self: elu_backward(grad, alpha, scale, input_scale, /* is_result */ false, self) + result: auto_element_wise + +- name: elu_(Tensor(a!) self, Scalar alpha=1, Scalar scale=1, Scalar input_scale=1) -> Tensor(a!) + self: elu_backward(grad, alpha, scale, input_scale, /* is_result */ true, result) + result: self_t.copy_(elu_backward(original_self_t, alpha, scale, input_scale, /* is_result */ true, result)) + +- name: celu(Tensor self, Scalar alpha=1.0) -> Tensor + self: elu_backward(grad, alpha, 1, 1.0/alpha.toFloat(), /* is_result */ false, self) + result: auto_element_wise + +- name: celu_(Tensor(a!) self, Scalar alpha=1.0) -> Tensor(a!) + self: elu_backward(grad, alpha, 1, 1.0/alpha.toFloat(), /* is_result */ true, result) + result: self_t.copy_(elu_backward(original_self_t, alpha, 1, 1.0/alpha.toFloat(), /* is_result */ true, result)) + +- name: gelu(Tensor self, *, str approximate='none') -> Tensor + self: gelu_backward(grad, self, approximate) + result: auto_element_wise + +- name: gelu_backward(Tensor grad_output, Tensor self, *, str approximate='none') -> Tensor + grad_output: gelu_backward(grad, self, approximate) + self: gelu_double_backward(grad, grad_output, self, approximate) + result: gelu_backward(grad_output_t, self_p, approximate) + gelu_double_backward(self_t, grad_output_p, self_p, approximate) + +- name: glu(Tensor self, int dim=-1) -> Tensor + # TODO: glu_backward can benefit from forward result, + # and forward ad/forward over reverse ad for that matter + self: glu_backward(grad, self, dim) + result: glu_jvp(result, self_p, self_t, dim) + +- name: hardshrink(Tensor self, Scalar lambd=0.5) -> Tensor + self: hardshrink_backward(grad, self, lambd) + result: auto_element_wise + +- name: hardshrink_backward(Tensor grad_out, Tensor self, Scalar lambd) -> Tensor + grad_out: hardshrink_backward(grad, self, lambd) + self: zeros_like(grad) + result: at::where((self_p > lambd).logical_or(self_p < -lambd), grad_out_t, at::zeros({}, result.options()).expand_as(result)) + +- name: hardtanh(Tensor self, Scalar min_val=-1, Scalar max_val=1) -> Tensor + self: hardtanh_backward(grad, self, min_val, max_val) + result: auto_element_wise + +- name: leaky_relu(Tensor self, Scalar negative_slope=0.01) -> Tensor + self: leaky_relu_backward(grad, self, negative_slope, false) + result: auto_element_wise + +- name: leaky_relu_(Tensor(a!) self, Scalar negative_slope=0.01) -> Tensor(a!) + self: leaky_relu_backward(grad, result, negative_slope, true) + result: self_t.copy_(leaky_relu_backward(original_self_t.conj(), result, negative_slope, true).conj()) + +- name: log_sigmoid_forward(Tensor self) -> (Tensor output, Tensor buffer) + self: log_sigmoid_backward(grad, self, buffer) + output: log_sigmoid_backward(self_t.conj(), self_p, buffer).conj() + output_differentiability: [True, False] + +- name: _log_softmax(Tensor self, int dim, bool half_to_float) -> Tensor + self: _log_softmax_backward_data(grad, result, dim, self.scalar_type()) + result: self_t - logsumexp_jvp(self_p, self_t, {dim}, true) + +- name: _sparse_log_softmax(Tensor self, int dim, bool half_to_float) -> Tensor + self: _sparse_log_softmax_backward_data(grad, result, dim, self) + +- name: _masked_softmax(Tensor self, Tensor mask, int? dim=None, int? mask_type=None) -> Tensor + self: _masked_softmax_backward(grad, result, mask, dim) + mask: non_differentiable + +- name: _prelu_kernel(Tensor self, Tensor weight) -> Tensor + self, weight: "grad.defined() ? _prelu_kernel_backward(grad, self, weight) : std::tuple()" + result: at::where(self_p >= 0, self_t, weight_p * self_t + weight_t * self_p) + +- name: _prelu_kernel_backward(Tensor grad_output, Tensor self, Tensor weight) -> (Tensor, Tensor) + grad_output: "grads[0].defined() ? + (grads[1].defined() ? at::where(self >= 0, grads[0], grads[0] * weight + grads[1] * self) + : at::where(self >= 0, grads[0], grads[0] * weight)) + : at::where(self >= 0, at::zeros({}, grad_output.options()), grads[1] * self)" + self: "grads[1].defined() ? at::where(self >= 0, at::zeros({}, self.options()), grad_output * grads[1]) : zeros_like(self)" + weight: "grads[0].defined() ? at::where(self >= 0, at::zeros({}, weight.options()), grad_output * grads[0]) : zeros_like(self)" + result0: at::where(self_p >= 0, grad_output_t, grad_output_t * weight_p + grad_output_p * weight_t) + result1: at::where(self_p >= 0, at::zeros({}, self_p.options()), grad_output_p * self_t + grad_output_t * self_p) + +- name: rrelu_with_noise(Tensor self, Tensor(b!) noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor + self: rrelu_with_noise_backward(grad, self, noise, lower, upper, training, false) + result: auto_element_wise + +- name: rrelu_with_noise_(Tensor(a!) self, Tensor(b!) noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> Tensor(a!) + self: rrelu_with_noise_backward(grad, result, noise, lower, upper, training, true) + +- name: rrelu_with_noise_functional(Tensor self, Tensor noise, Scalar lower=0.125, Scalar upper=0.3333333333333333, bool training=False, Generator? generator=None) -> (Tensor, Tensor noise_out) + noise: non_differentiable + self: rrelu_with_noise_backward(grad, self, noise, lower, upper, training, false) + +- name: _softmax(Tensor self, int dim, bool half_to_float) -> Tensor + self: _softmax_backward_data(grad, result, dim, self.scalar_type()) + result: result * (self_t - logsumexp_jvp(self_p, self_t, {dim}, true)) + +- name: _sparse_softmax(Tensor self, int dim, bool half_to_float) -> Tensor + self: _sparse_softmax_backward_data(grad, result, dim, self) + +- name: _sparse_sparse_matmul(Tensor self, Tensor other) -> Tensor + self: sparse_sparse_matmul_backward(grad, self, other, 0) + other: sparse_sparse_matmul_backward(grad, self, other, 1) + +- name: softplus(Tensor self, Scalar beta=1, Scalar threshold=20) -> Tensor + self: softplus_backward(grad, self, beta, threshold) + result: auto_element_wise + +- name: softshrink(Tensor self, Scalar lambd=0.5) -> Tensor + self: softshrink_backward(grad, self, lambd) + result: auto_element_wise + +- name: threshold(Tensor self, Scalar threshold, Scalar value) -> Tensor + self: threshold_backward(grad, self, threshold) + result: auto_element_wise + +- name: threshold_(Tensor(a!) self, Scalar threshold, Scalar value) -> Tensor(a!) + self: threshold_backward(grad, self, threshold) + result: self_t.copy_(threshold_backward(self_t.conj(), original_self_p, threshold).conj()) + +- name: reflection_pad1d(Tensor self, SymInt[2] padding) -> Tensor + self: reflection_pad1d_backward_symint(grad, self, padding) + result: auto_linear + +- name: reflection_pad2d(Tensor self, SymInt[4] padding) -> Tensor + self: reflection_pad2d_backward_symint(grad, self, padding) + result: auto_linear + +- name: reflection_pad3d(Tensor self, SymInt[6] padding) -> Tensor + self: reflection_pad3d_backward_symint(grad, self, padding) + result: auto_linear + +- name: replication_pad1d(Tensor self, SymInt[2] padding) -> Tensor + self: replication_pad1d_backward_symint(grad, self, padding) + result: auto_linear + +- name: replication_pad2d(Tensor self, SymInt[4] padding) -> Tensor + self: replication_pad2d_backward_symint(grad, self, padding) + result: auto_linear + +- name: replication_pad3d(Tensor self, SymInt[6] padding) -> Tensor + self: replication_pad3d_backward_symint(grad, self, padding) + result: auto_linear + +- name: upsample_linear1d(Tensor self, SymInt[1] output_size, bool align_corners, float? scales=None) -> Tensor + self: upsample_linear1d_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales) + result: auto_linear + +- name: upsample_bilinear2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + self: upsample_bilinear2d_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales_h, scales_w) + result: auto_linear + +- name: _upsample_bilinear2d_aa(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + self: _upsample_bilinear2d_aa_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales_h, scales_w) + result: auto_linear + +- name: upsample_bicubic2d(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + self: upsample_bicubic2d_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales_h, scales_w) + result: auto_linear + +- name: _upsample_bicubic2d_aa(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + self: _upsample_bicubic2d_aa_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales_h, scales_w) + result: auto_linear + +- name: _upsample_lanczos2d_aa(Tensor self, SymInt[2] output_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + self: _upsample_lanczos2d_aa_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales_h, scales_w) + result: auto_linear + +- name: upsample_trilinear3d(Tensor self, SymInt[3] output_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + self: upsample_trilinear3d_backward_symint(grad, output_size, self.sym_sizes(), align_corners, scales_d, scales_h, scales_w) + result: auto_linear + +- name: upsample_nearest1d(Tensor self, SymInt[1] output_size, float? scales=None) -> Tensor + self: upsample_nearest1d_backward_symint(grad, output_size, self.sym_sizes(), scales) + result: auto_linear + +- name: _upsample_nearest_exact1d(Tensor self, SymInt[1] output_size, float? scales=None) -> Tensor + self: _upsample_nearest_exact1d_backward_symint(grad, output_size, self.sym_sizes(), scales) + result: auto_linear + +- name: upsample_nearest2d(Tensor self, SymInt[2] output_size, float? scales_h=None, float? scales_w=None) -> Tensor + self: upsample_nearest2d_backward_symint(grad, output_size, self.sym_sizes(), scales_h, scales_w) + result: auto_linear + +- name: _upsample_nearest_exact2d(Tensor self, SymInt[2] output_size, float? scales_h=None, float? scales_w=None) -> Tensor + self: _upsample_nearest_exact2d_backward_symint(grad, output_size, self.sym_sizes(), scales_h, scales_w) + result: auto_linear + +- name: upsample_nearest3d(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + self: upsample_nearest3d_backward_symint(grad, output_size, self.sym_sizes(), scales_d, scales_h, scales_w) + result: auto_linear + +- name: _upsample_nearest_exact3d(Tensor self, SymInt[3] output_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + self: _upsample_nearest_exact3d_backward_symint(grad, output_size, self.sym_sizes(), scales_d, scales_h, scales_w) + result: auto_linear + +- name: pixel_shuffle(Tensor self, int upscale_factor) -> Tensor + self: pixel_unshuffle(grad, upscale_factor) + result: auto_linear + +- name: pixel_unshuffle(Tensor self, int downscale_factor) -> Tensor + self: pixel_shuffle(grad, downscale_factor) + result: auto_linear + +- name: channel_shuffle(Tensor self, SymInt groups) -> Tensor + self: channel_shuffle_symint(grad, grad.sym_size(1) / groups) + result: auto_linear + +- name: _adaptive_avg_pool2d(Tensor self, SymInt[2] output_size) -> Tensor + self: _adaptive_avg_pool2d_backward(grad, self) + result: auto_linear + +- name: _adaptive_avg_pool3d(Tensor self, SymInt[3] output_size) -> Tensor + self: _adaptive_avg_pool3d_backward(grad, self) + result: auto_linear + +- name: adaptive_max_pool2d(Tensor self, int[2] output_size) -> (Tensor, Tensor) + self: adaptive_max_pool2d_backward(grad, self, result1) + result0: gather(self_t.flatten(-2), -1, result1.flatten(-2)).view_as(result1) + output_differentiability: [True, False] + +- name: adaptive_max_pool3d(Tensor self, int[3] output_size) -> (Tensor, Tensor) + self: adaptive_max_pool3d_backward(grad, self, result1) + result0: gather(self_t.flatten(-3), -1, result1.flatten(-3)).view_as(result1) + output_differentiability: [True, False] + +- name: avg_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> Tensor + self: avg_pool2d_backward(grad, self, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override) + result: auto_linear + +- name: avg_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, bool ceil_mode=False, bool count_include_pad=True, int? divisor_override=None) -> Tensor + self: avg_pool3d_backward(grad, self, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override) + result: auto_linear + +- name: fractional_max_pool2d(Tensor self, int[2] kernel_size, int[2] output_size, Tensor random_samples) -> (Tensor, Tensor) + self: fractional_max_pool2d_backward(grad, self, kernel_size, output_size, result1) + result0: gather(self_t.flatten(-2), -1, result1.flatten(-2)).view_as(result1) + output_differentiability: [True, False] + +- name: fractional_max_pool3d(Tensor self, int[3] kernel_size, int[3] output_size, Tensor random_samples) -> (Tensor, Tensor) + self: fractional_max_pool3d_backward(grad, self, kernel_size, output_size, result1) + result0: gather(self_t.flatten(-3), -1, result1.flatten(-3)).view_as(result1) + output_differentiability: [True, False] + +- name: linear(Tensor input, Tensor weight, Tensor? bias=None) -> Tensor + input, weight, bias: "grad.defined() ? linear_backward(input, grad, weight, grad_input_mask) : std::tuple()" + +- name: linear_backward(Tensor self, Tensor grad_output, Tensor weight, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + self, grad_output, weight: linear_double_backward(grads, self, grad_output, weight) + +#mps +- name: max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor + self: max_pool2d_backward(grad, self, kernel_size, stride, padding, dilation, ceil_mode) + +- name: _mps_convolution(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups) -> Tensor + self, weight, bias: "grad.defined() ? mps_convolution_backward_symint(self, grad, weight, padding, stride, dilation, groups, grad_input_mask) : std::tuple()" + +- name: mps_convolution_backward(Tensor self, Tensor grad_output, Tensor weight, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + grad_output, self, weight: _convolution_double_backward_symint(grads[0], grads[1], grads[2], grad_output, weight, self, stride, padding, dilation, false, std::vector(padding.size(), 0), groups, grad_input_mask) + +- name: max_pool2d_with_indices(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor) + self: max_pool2d_with_indices_backward(grad, self, kernel_size, stride, padding, dilation, ceil_mode, result1) + result0: gather(self_t.flatten(-2), -1, result1.flatten(-2)).view_as(result1) + output_differentiability: [True, False] + +- name: max_pool3d_with_indices(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> (Tensor, Tensor) + self: max_pool3d_with_indices_backward(grad, self, kernel_size, stride, padding, dilation, ceil_mode, result1) + result0: gather(self_t.flatten(-3), -1, result1.flatten(-3)).view_as(result1) + output_differentiability: [True, False] + +- name: max_unpool2d(Tensor self, Tensor indices, SymInt[2] output_size) -> Tensor + self: max_pool_double_backward(grad, indices, 2) + indices: non_differentiable + result: auto_linear + +- name: max_unpool3d(Tensor self, Tensor indices, SymInt[3] output_size, int[3] stride, int[3] padding) -> Tensor + self: max_pool_double_backward(grad, indices, 3) + indices: non_differentiable + result: auto_linear + +- name: convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups) -> Tensor + input, weight, bias: "grad.defined() ? convolution_backward_symint(grad, input, weight, bias->sym_sizes(), stride, padding, dilation, transposed, output_padding, groups, grad_input_mask) : std::tuple()" + result: convolution_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, stride, padding, dilation, transposed, output_padding, groups) + +# TorchScript serializes calls to _convolution so this entry is present until that is changed to use convolution. +# Note that the benchmark, deterministic, cudnn_enabled, and allow_tf32 flags are queried from the global context +# by convolution_backward instead of being passed along from the forward pass. +- name: _convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups, bool benchmark, bool deterministic, bool cudnn_enabled, bool allow_tf32) -> Tensor + input, weight, bias: "grad.defined() ? convolution_backward_symint(grad, input, weight, bias->sym_sizes(), stride, padding, dilation, transposed, output_padding, groups, grad_input_mask) : std::tuple()" + result: _convolution_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, stride, padding, dilation, transposed, output_padding, groups, benchmark, deterministic, cudnn_enabled, allow_tf32) + +- name: convolution_backward(Tensor grad_output, Tensor input, Tensor weight, SymInt[]? bias_sizes, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups, bool[3] output_mask) -> (Tensor, Tensor, Tensor) + grad_output, input, weight: _convolution_double_backward_symint(grads[0], grads[1], grads[2], grad_output, weight, input, stride, padding, dilation, transposed, output_padding, groups, grad_input_mask) + result0: std::get<0>(convolution_backward_symint(grad_output_p, input_p, weight_t, bias_sizes, stride, padding, dilation, transposed, output_padding, groups, {true, false, false})) + std::get<0>(convolution_backward_symint(grad_output_t, input_p, weight_p, bias_sizes, stride, padding, dilation, transposed, output_padding, groups, {true, false, false})) + result1: std::get<1>(convolution_backward_symint(grad_output_p, input_t, weight_p, bias_sizes, stride, padding, dilation, transposed, output_padding, groups, {false, true, false})) + std::get<1>(convolution_backward_symint(grad_output_t, input_p, weight_p, bias_sizes, stride, padding, dilation, transposed, output_padding, groups, {false, true, false})) + result2: convolution_backward_jvp_grad_bias(grad_output_t, result2) + +- name: convolution_overrideable(Tensor input, Tensor weight, Tensor? bias, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups) -> Tensor + input, weight, bias: "grad.defined() ? convolution_backward_overrideable_symint(grad, input, weight, stride, padding, dilation, transposed, output_padding, groups, grad_input_mask) : std::tuple()" + +- name: convolution_backward_overrideable(Tensor grad_output, Tensor input, Tensor weight, SymInt[] stride, SymInt[] padding, SymInt[] dilation, bool transposed, SymInt[] output_padding, SymInt groups, bool[3] output_mask) -> (Tensor grad_input, Tensor grad_weight, Tensor grad_bias) + grad_output, input, weight: _convolution_double_backward_symint(grads[0], grads[1], grads[2], grad_output, weight, input, stride, padding, dilation, transposed, output_padding, groups, grad_input_mask) + +- name: slow_conv_transpose2d(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0, SymInt[2] output_padding=0, SymInt[2] dilation=1) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad, self, weight, bias->sym_sizes(), stride, padding, dilation, true, output_padding, 1, grad_input_mask) : std::tuple()" + +- name: slow_conv_transpose3d(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0, SymInt[3] output_padding=0, SymInt[3] dilation=1) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad, self, weight, bias->sym_sizes(), stride, padding, dilation, true, output_padding, 1, grad_input_mask) : std::tuple()" + +- name: _slow_conv2d_forward(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias, SymInt[2] stride, SymInt[2] padding) -> Tensor + self, weight, bias: "grad.defined() ? _slow_conv2d_backward_symint(grad, self, weight, kernel_size, stride, padding, grad_input_mask) : std::tuple()" + +- name: _slow_conv2d_backward.output_mask(Tensor grad_output, Tensor self, Tensor weight, SymInt[2] kernel_size, SymInt[2] stride, SymInt[2] padding, bool[3] output_mask) -> (Tensor grad_input, Tensor grad_weight, Tensor grad_bias) + grad_output, self, weight: _convolution_double_backward_symint(grads[0], grads[1], grads[2], grad_output, weight, self, stride, padding, {{1, 1}}, false, {{0, 0}}, 1, grad_input_mask) + +- name: _conv_depthwise2d(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias, SymInt[2] stride, SymInt[2] padding, SymInt[2] dilation) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad.contiguous(), self, weight, bias->sym_sizes(), stride, padding, dilation, /*transposed=*/ false, /*output_padding=*/ {{0, 0}}, /*groups=*/ 1, grad_input_mask) : std::tuple()" + +- name: conv_depthwise3d(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias, SymInt[3] stride, SymInt[3] padding, SymInt[3] dilation) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad.contiguous(), self, weight, bias->sym_sizes(), stride, padding, dilation, /*transposed=*/ false, /*output_padding=*/ {{0, 0, 0}}, /*groups=*/ 1, grad_input_mask) : std::tuple()" + +- name: slow_conv3d_forward(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias, SymInt[3] stride, SymInt[3] padding) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad, self, weight, bias->sym_sizes(), stride, padding, /*dilation=*/ {{1, 1, 1}}, false, /*output_padding=*/ {{0, 0, 0}}, 1, grad_input_mask) : std::tuple()" + +- name: slow_conv_dilated2d(Tensor self, Tensor weight, SymInt[2] kernel_size, Tensor? bias=None, SymInt[2] stride=1, SymInt[2] padding=0, SymInt[2] dilation=1) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad, self, weight, bias->sym_sizes(), stride, padding, dilation, false, std::vector(padding.size(), 0), 1, grad_input_mask) : std::tuple()" + +- name: slow_conv_dilated3d(Tensor self, Tensor weight, SymInt[3] kernel_size, Tensor? bias=None, SymInt[3] stride=1, SymInt[3] padding=0, SymInt[3] dilation=1) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad, self, weight, bias->sym_sizes(), stride, padding, dilation, false, std::vector(padding.size(), 0), 1, grad_input_mask) : std::tuple()" + +- name: col2im(Tensor self, SymInt[2] output_size, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride) -> Tensor + self: im2col(grad, kernel_size, dilation, padding, stride) + result: auto_linear + +- name: im2col(Tensor self, int[2] kernel_size, int[2] dilation, int[2] padding, int[2] stride) -> Tensor + self: col2im_symint(grad, {self.sym_size(-2), self.sym_size(-1)}, kernel_size, dilation, padding, stride) + result: auto_linear + +- name: _adaptive_avg_pool2d_backward(Tensor grad_output, Tensor self) -> Tensor + grad_output: _adaptive_avg_pool2d_symint(grad, {grad_output.sym_size(-2), grad_output.sym_size(-1)}) + self: zeros_like(self) + result: _adaptive_avg_pool2d_backward(grad_output_t, self_p) + +- name: _adaptive_avg_pool3d_backward(Tensor grad_output, Tensor self) -> Tensor + grad_output: _adaptive_avg_pool3d_symint(grad, { grad_output.sym_size(-3), grad_output.sym_size(-2), grad_output.sym_size(-1) }) + self: zeros_like(self) + result: _adaptive_avg_pool3d_backward(grad_output_t, self_p) + +- name: adaptive_max_pool2d_backward(Tensor grad_output, Tensor self, Tensor indices) -> Tensor + grad_output: max_pool_double_backward(grad, indices, 2) + self: zeros_like(self) + result: auto_linear + +- name: adaptive_max_pool3d_backward(Tensor grad_output, Tensor self, Tensor indices) -> Tensor + grad_output: max_pool_double_backward(grad, indices, 3) + self: zeros_like(self) + result: auto_linear + +- name: avg_pool2d_backward(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] stride, int[2] padding, bool ceil_mode, bool count_include_pad, int? divisor_override) -> Tensor + grad_output: avg_pool2d(grad, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override) + self: zeros_like(self) + result: avg_pool2d_backward(grad_output_t, self_p, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override) + +- name: avg_pool3d_backward(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] stride, int[3] padding, bool ceil_mode, bool count_include_pad, int? divisor_override) -> Tensor + grad_output: avg_pool3d(grad, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override) + self: zeros_like(self) + result: avg_pool3d_backward(grad_output_t, self_p, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override) + +- name: elu_backward(Tensor grad_output, Scalar alpha, Scalar scale, Scalar input_scale, bool is_result, Tensor self_or_result) -> Tensor + grad_output: elu_backward(grad, alpha, scale, input_scale, is_result, self_or_result) + self_or_result: elu_double_backward(grad, grad_output, alpha, scale, input_scale, is_result, self_or_result) + result: elu_backward(grad_output_t, alpha, scale, input_scale, is_result, self_or_result_p) + elu_double_backward(self_or_result_t, grad_output_p, alpha, scale, input_scale, is_result, self_or_result_p) + +- name: fractional_max_pool2d_backward(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] output_size, Tensor indices) -> Tensor + grad_output: max_pool_double_backward(grad, indices, 2) + self: zeros_like(self) + result: auto_linear + +- name: fractional_max_pool3d_backward(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] output_size, Tensor indices) -> Tensor + grad_output: max_pool_double_backward(grad, indices, 3) + self: zeros_like(self) + result: auto_linear + +- name: glu_backward(Tensor grad_output, Tensor self, int dim) -> Tensor + grad_output: glu_double_backward_grad_output(grad, self, dim) + self: glu_double_backward(grad, grad_output, self, dim) + result: glu_backward_jvp(result, grad_output_p, self_p, grad_output_t, self_t, dim) + +- name: hardtanh_backward(Tensor grad_output, Tensor self, Scalar min_val, Scalar max_val) -> Tensor + grad_output: hardtanh_backward(grad, self, min_val, max_val) + self: zeros_like(grad) + result: at::where((self_p > min_val).logical_and(self_p < max_val), grad_output_t, at::zeros({}, result.options()).expand_as(result)) + +- name: log_sigmoid_backward(Tensor grad_output, Tensor self, Tensor buffer) -> Tensor + grad_output: log_sigmoid_backward(grad, self, buffer) + self: log_sigmoid_double_backward(grad * grad_output, self) + result: log_sigmoid_backward(grad_output_t, self_p, buffer) + log_sigmoid_double_backward(self_t * grad_output_p, self_p) + +- name: _log_softmax_backward_data(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype) -> Tensor + grad_output: grad.to(output.dtype()) - (grad.to(output.dtype()) * output.exp()).sum(dim, true) + output: (-grad_output.sum(dim, true) * output.exp() * grad.to(output.dtype())).to(output.dtype()) + +- name: leaky_relu_backward(Tensor grad_output, Tensor self, Scalar negative_slope, bool self_is_result) -> Tensor + # self_is_result is always false here since double backward call is an out-of-place call, self is input itself + grad_output: leaky_relu_backward(grad, self, negative_slope, false) + self: zeros_like(grad) + # leaky_relu_backward(grad_output, self, negative_slope, false) + # computes grad_output * at::where(self_p > 0, 1, negative_slope) + # so the jvp formula is the following: + # grad_output_t * at::where(self_p > 0, self_p.new_ones([]), negative_slope); + # + # leaky_relu_backward(grad_output, result, negative_slope, true) + # computes grad_output * at::where(result > 0, 1, negative_slope) + # under the assumption that `negative_slope` is positive (otherwise, + # it is not possible to compute the gradient). + # + # so the jvp formula is the following: + # grad_output_t * at::where(result_p > 0, result_p.new_ones([]), negative_slope); + # with the assumption that negative_slope is positive. + # + # Combined together that results in the following optimized kernel which + # also checks the assumption that negative_slope is positive when self_is_result + # is True: + result: leaky_relu_backward(grad_output_t, self_p, negative_slope, self_is_result) + +# This derivative is mps-only, and `error_for_max_pool2d_double_backward` just raises an error. +- name: max_pool2d_backward(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor + grad_output: error_for_max_pool2d_double_backward() + self: zeros_like(self) + result: auto_linear + +- name: max_pool2d_with_indices_backward(Tensor grad_output, Tensor self, int[2] kernel_size, int[2] stride, int[2] padding, int[2] dilation, bool ceil_mode, Tensor indices) -> Tensor + grad_output: max_pool_double_backward(grad, indices, 2) + self: zeros_like(self) + indices: non_differentiable + result: auto_linear + +- name: max_pool3d_with_indices_backward(Tensor grad_output, Tensor self, int[3] kernel_size, int[3] stride, int[3] padding, int[3] dilation, bool ceil_mode, Tensor indices) -> Tensor + grad_output: max_pool_double_backward(grad, indices, 3) + self: zeros_like(self) + indices: non_differentiable + result: auto_linear + +- name: mse_loss_backward(Tensor grad_output, Tensor self, Tensor target, int reduction) -> Tensor + grad_output: mse_loss_backward(grad, self, target, reduction) + self: mse_loss_double_backward(grad * grad_output, self, reduction) + target: -mse_loss_double_backward(grad * grad_output, target, reduction) + result: " mse_loss_double_backward(self_t * grad_output_p, self_p, reduction) + - mse_loss_double_backward(target_t * grad_output_p, target_p, reduction) + + mse_loss_backward(grad_output_t, self_p, target_p, reduction) + " + +- name: nll_loss_backward(Tensor grad_output, Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index, Tensor total_weight) -> Tensor + grad_output: nll_loss_symint(grad, target, weight, reduction, ignore_index) + self: zeros_like(grad) + target: non_differentiable + +- name: nll_loss2d_backward(Tensor grad_output, Tensor self, Tensor target, Tensor? weight, int reduction, SymInt ignore_index, Tensor total_weight) -> Tensor + grad_output: nll_loss2d_symint(grad, target, weight, reduction, ignore_index) + self: zeros_like(grad) + target: non_differentiable + +- name: rrelu_with_noise_backward(Tensor grad_output, Tensor self, Tensor noise, Scalar lower, Scalar upper, bool training, bool self_is_result) -> Tensor + # self_is_result is always false here since double backward call is an out-of-place call, self is input itself + grad_output: rrelu_with_noise_backward(grad, self, noise, lower, upper, training, false) + self: zeros_like(grad) + result: rrelu_with_noise_backward(grad_output_t, self_p, noise, lower, upper, training, false) + +- name: reflection_pad1d_backward(Tensor grad_output, Tensor self, SymInt[2] padding) -> Tensor + grad_output: reflection_pad1d_symint(grad, padding) + self: zeros_like(self) + result: reflection_pad1d_backward_symint(grad_output_t, self_p, padding) + +- name: reflection_pad2d_backward(Tensor grad_output, Tensor self, SymInt[4] padding) -> Tensor + grad_output: reflection_pad2d_symint(grad, padding) + self: zeros_like(self) + result: reflection_pad2d_backward_symint(grad_output_t, self_p, padding) + +- name: reflection_pad3d_backward(Tensor grad_output, Tensor self, SymInt[6] padding) -> Tensor + grad_output: reflection_pad3d_symint(grad, padding) + self: zeros_like(self) + result: reflection_pad3d_backward_symint(grad_output_t, self_p, padding) + +- name: replication_pad1d_backward(Tensor grad_output, Tensor self, SymInt[2] padding) -> Tensor + grad_output: replication_pad1d_symint(grad, padding) + self: zeros_like(self) + result: replication_pad1d_backward_symint(grad_output_t, self_p, padding) + +- name: replication_pad2d_backward(Tensor grad_output, Tensor self, SymInt[4] padding) -> Tensor + grad_output: replication_pad2d_symint(grad, padding) + self: zeros_like(self) + result: replication_pad2d_backward_symint(grad_output_t, self_p, padding) + +- name: replication_pad3d_backward(Tensor grad_output, Tensor self, SymInt[6] padding) -> Tensor + grad_output: replication_pad3d_symint(grad, padding) + self: zeros_like(self) + result: replication_pad3d_backward_symint(grad_output_t, self_p, padding) + +- name: sparse_sampled_addmm(Tensor self, Tensor mat1, Tensor mat2, *, Scalar beta=1, Scalar alpha=1) -> Tensor + self, mat1, mat2: "sparse_sampled_addmm_backward(grad, + self, + wrap_opt_if(mat1, grad_input_mask[2]), + wrap_opt_if(mat2, grad_input_mask[1]), + alpha, beta, grad_input_mask)" + +- name: _sparse_mm_reduce_impl(Tensor self, Tensor other, str reduce) -> (Tensor, Tensor) + output_differentiability: [True, False] + self, other: "grad.defined() ? _sparse_mm_reduce_impl_backward(self, grad, other, reduce, result1, grad_input_mask) : std::tuple()" + +- name: smooth_l1_loss_backward(Tensor grad_output, Tensor self, Tensor target, int reduction, float beta) -> Tensor + grad_output: smooth_l1_loss_backward(grad, self, target, reduction, beta) + self: smooth_l1_loss_double_backward(grad * grad_output, self, target, reduction, beta) + target: -smooth_l1_loss_double_backward(grad * grad_output, self, target, reduction, beta) + result: " smooth_l1_loss_double_backward(self_t * grad_output_p, self_p, target_p, reduction, beta) + - smooth_l1_loss_double_backward(target_t * grad_output_p, self_p, target_p, reduction, beta) + + smooth_l1_loss_backward(grad_output_t, self_p, target_p, reduction, beta) + " + +- name: huber_loss_backward(Tensor grad_output, Tensor self, Tensor target, int reduction, float delta) -> Tensor + grad_output: huber_loss_double_backward_grad_output(grad, grad_output, self, target, reduction, delta) + self: huber_loss_double_backward(grad * grad_output, self, target, reduction, delta) + target: -huber_loss_double_backward(grad * grad_output, self, target, reduction, delta) + +- name: softplus_backward(Tensor grad_output, Tensor self, Scalar beta, Scalar threshold) -> Tensor + grad_output: softplus_backward(grad, self, beta, threshold) + self: softplus_double_backward(grad * grad_output, self, beta, threshold) + result: "softplus_backward(grad_output_t, self_p, beta, threshold) + + softplus_double_backward(self_t * grad_output_p, self_p, beta, threshold)" + +- name: _softmax_backward_data(Tensor grad_output, Tensor output, int dim, ScalarType input_dtype) -> Tensor + grad_output: _softmax_backward_data(grad.to(output.dtype()), output, dim, input_dtype) + output: softmax_double_backward(grad.to(output.dtype()), grad_output, dim, output).to(output.dtype()) + +- name: soft_margin_loss_backward(Tensor grad_output, Tensor self, Tensor target, int reduction) -> Tensor + grad_output: soft_margin_loss_double_backward_grad_output(grad, grad_output, self, target, reduction) + self: soft_margin_loss_double_backward(grad * grad_output, self, target, reduction) + +- name: softshrink_backward(Tensor grad_output, Tensor self, Scalar lambd) -> Tensor + grad_output: softshrink_backward(grad, self, lambd) + self: zeros_like(grad) + result: at::where((self_p > lambd).logical_or(self_p < -lambd), grad_output_t, at::zeros({}, result.options()).expand_as(result)) + +- name: threshold_backward(Tensor grad_output, Tensor self, Scalar threshold) -> Tensor + grad_output: threshold_backward(grad, self, threshold) + self: zeros_like(grad) + result: zeros_like(self_t) + threshold_backward(grad_output_t, self_p, threshold) + +- name: upsample_linear1d_backward(Tensor grad_output, SymInt[1] output_size, SymInt[3] input_size, bool align_corners, float? scales=None) -> Tensor + grad_output: upsample_linear1d_symint(grad, output_size, align_corners, scales) + result: auto_linear + +- name: upsample_bilinear2d_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: upsample_bilinear2d_symint(grad, output_size, align_corners, scales_h, scales_w) + result: auto_linear + +- name: _upsample_bilinear2d_aa_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: _upsample_bilinear2d_aa_symint(grad, output_size, align_corners, scales_h, scales_w) + result: auto_linear + +- name: upsample_bicubic2d_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: upsample_bicubic2d_symint(grad, output_size, align_corners, scales_h, scales_w) + result: auto_linear + +- name: _upsample_bicubic2d_aa_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: _upsample_bicubic2d_aa_symint(grad, output_size, align_corners, scales_h, scales_w) + result: auto_linear + +- name: _upsample_lanczos2d_aa_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, bool align_corners, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: _upsample_lanczos2d_aa_symint(grad, output_size, align_corners, scales_h, scales_w) + result: auto_linear + +- name: upsample_trilinear3d_backward(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, bool align_corners, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: upsample_trilinear3d_symint(grad, output_size, align_corners, scales_d, scales_h, scales_w) + result: auto_linear + +- name: upsample_nearest1d_backward(Tensor grad_output, SymInt[1] output_size, SymInt[3] input_size, float? scales=None) -> Tensor + grad_output: upsample_nearest1d_symint(grad, output_size, scales) + result: auto_linear + +- name: _upsample_nearest_exact1d_backward(Tensor grad_output, SymInt[1] output_size, SymInt[3] input_size, float? scales=None) -> Tensor + grad_output: _upsample_nearest_exact1d_symint(grad, output_size, scales) + result: auto_linear + +- name: upsample_nearest2d_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: upsample_nearest2d_symint(grad, output_size, scales_h, scales_w) + result: auto_linear + +- name: _upsample_nearest_exact2d_backward(Tensor grad_output, SymInt[2] output_size, SymInt[4] input_size, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: _upsample_nearest_exact2d_symint(grad, output_size, scales_h, scales_w) + result: auto_linear + +- name: upsample_nearest3d_backward(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: upsample_nearest3d_symint(grad, output_size, scales_d, scales_h, scales_w) + result: auto_linear + +- name: _upsample_nearest_exact3d_backward(Tensor grad_output, SymInt[3] output_size, SymInt[5] input_size, float? scales_d=None, float? scales_h=None, float? scales_w=None) -> Tensor + grad_output: _upsample_nearest_exact3d_symint(grad, output_size, scales_d, scales_h, scales_w) + result: auto_linear + +- name: sigmoid_backward(Tensor grad_output, Tensor output) -> Tensor + grad_output: sigmoid_backward(grad, output.conj()) + output: grad.conj() * grad_output * (-2 * output.conj() + 1) + result: sigmoid_backward(grad_output_t, output_p) + output_t.conj() * grad_output_p * (-2 * output_p.conj() + 1) + +- name: tanh_backward(Tensor grad_output, Tensor output) -> Tensor + grad_output: tanh_backward(grad, output.conj()) + output: grad.conj() * (-2 * output.conj() * grad_output) + result: tanh_backward(grad_output_t, output_p) + output_t.conj() * (-2 * output_p.conj() * grad_output_p) + +# cudnn +- name: _cudnn_ctc_loss(Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, int blank, bool deterministic, bool zero_infinity) -> (Tensor, Tensor) + log_probs: _cudnn_ctc_loss_backward(grad, result0, result1, zero_infinity) + +- name: _cudnn_ctc_loss.Tensor(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank, bool deterministic, bool zero_infinity) -> (Tensor, Tensor) + log_probs: _cudnn_ctc_loss_backward(grad, result0, result1, zero_infinity) + +- name: cudnn_convolution_transpose(Tensor self, Tensor weight, SymInt[] padding, SymInt[] output_padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic, bool allow_tf32) -> Tensor + self, weight: "_cudnn_convolution_backward(self, grad, weight, padding, output_padding, stride, dilation, true, groups, {grad_input_mask[0], grad_input_mask[1]})" + +- name: _mps_convolution_transpose(Tensor self, Tensor weight, SymInt[] padding, SymInt[] output_padding, SymInt[] stride, SymInt[] dilation, SymInt groups) -> Tensor + self, weight: "grad.defined() ? mps_convolution_transpose_backward_symint(self, grad, weight, padding, output_padding, stride, dilation, groups, grad_input_mask) : std::tuple()" + +- name: cudnn_convolution(Tensor self, Tensor weight, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic, bool allow_tf32) -> Tensor + self, weight: "_cudnn_convolution_backward(self, grad, weight, padding, std::vector(padding.size(), 0), stride, dilation, false, groups, {grad_input_mask[0], grad_input_mask[1]})" + +- name: cudnn_grid_sampler(Tensor self, Tensor grid) -> Tensor output + self, grid: "grad.defined() ? cudnn_grid_sampler_backward(self, grid, grad) : std::tuple()" + +- name: cudnn_grid_sampler_backward(Tensor self, Tensor grid, Tensor grad_output) -> (Tensor grad_self, Tensor grad_grid) + grad_output, self, grid: grid_sampler_2d_double_backward(grads[0], grads[1], grad_output, self, grid, 0, 0, true, grad_input_mask) + +- name: cudnn_affine_grid_generator(Tensor theta, int N, int C, int H, int W) -> Tensor grid + theta: cudnn_affine_grid_generator_backward(grad, N, C, H, W) + +# NB: Why is the backwards here so complicated? CuDNN cannot be used to compute +# backward in evaluation mode, because the math for backward in evaluation mode +# is different (since the forward math is different), and CuDNN does not support +# it. And in any case, you shouldn't be using this bn in evaluation mode, +# because it should be merged into the previous convolution (left for future +# work.) +# NB2: The quotes around the gradient are needed to appease YAML parsing rules. +- name: cudnn_batch_norm(Tensor input, Tensor weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float exponential_average_factor, float epsilon) -> (Tensor, Tensor, Tensor, Tensor) + input, weight, bias: "grad.defined() ? (training ? cudnn_batch_norm_backward(input, grad.contiguous(input.suggest_memory_format()), weight, running_mean, running_var, result1, result2, epsilon, retain_variables ? result3.clone() : result3) : native_batch_norm_backward(grad, input, weight, running_mean, running_var, result1, result2, training, epsilon, grad_input_mask)) : std::tuple()" + result0: batch_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, running_mean, running_var, result1, result2, training, epsilon) + +# HACK: save_mean and save_var are going to be passed in as +# requires_grad variables (even though we'll never backprop through +# them) so we need to prevent the unpacking from triggering an error. +- name: cudnn_batch_norm_backward(Tensor input, Tensor grad_output, Tensor weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_var, float epsilon, Tensor reserveSpace) -> (Tensor, Tensor, Tensor) + save_mean: not_implemented("cudnn_batch_norm_backward save_mean") + save_var: not_implemented("cudnn_batch_norm_backward save_var") + reserveSpace: not_implemented("cudnn_batch_norm_backward reserveSpace") + input, weight, grad_output: batchnorm_double_backward(input, weight, grads[0], grads[1], grads[2], grad_output, running_mean, running_var, true, epsilon, save_mean, save_var, grad_input_mask) + +# nnpack + +- name: _nnpack_spatial_convolution(Tensor input, Tensor weight, Tensor? bias, SymInt[2] padding, SymInt[2] stride=1) -> Tensor + # NNPACK does not support strided convolutions in the backwards path, which is the reason why we are using the closest available function that does here. + input, weight, bias: "grad.defined() ? convolution_backward_symint(grad, input, weight, bias->sym_sizes(), stride, padding, std::vector(padding.size(), 1), false, std::vector(padding.size(), 0), 1, grad_input_mask) : std::tuple()" + +#LSTM MPS +- name: _lstm_mps(Tensor input, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) + output_differentiability: [True, True, True, False, False, False] + input, hx, params: "lstm_mps_backward(grads[0], grads[1], grads[2], result3, result4, input, result5, hx, params, has_biases, num_layers, dropout, train, bidirectional, batch_first)" + +- name: lstm_mps_backward(Tensor? grad_y, Tensor? grad_hy, Tensor? grad_cy, Tensor z_state, Tensor cell_state_fwd, Tensor input, Tensor layersOutputs, Tensor[] hx, Tensor[] params, bool has_biases, int num_layers, float dropout, bool train, bool bidirectional, bool batch_first) -> (Tensor, Tensor[], Tensor[]) + + + +# Only frst three of _cudnn_rnn outputs can have gradients. +# _cudnn_rnn outputs: (output, hy, cy, reserve, weight_buf) +- name: _cudnn_rnn(Tensor input, Tensor[] weight, int weight_stride0, Tensor? weight_buf, Tensor hx, Tensor? cx, int mode, SymInt hidden_size, SymInt proj_size, int num_layers, bool batch_first, float dropout, bool train, bool bidirectional, SymInt[] batch_sizes, Tensor? dropout_state) -> (Tensor, Tensor, Tensor, Tensor, Tensor) + dropout_state: non_differentiable + output_differentiability: [True, True, True, False, False] + input, hx, cx, weight: "_cudnn_rnn_backward_symint(input, weight, weight_stride0, result4, hx, cx, result0, grads[0], grads[1], grads[2], mode, hidden_size, proj_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state, retain_variables ? result3.clone() : result3, grad_input_mask)" + +- name: _cudnn_rnn_backward(Tensor input, Tensor[] weight, int weight_stride0, Tensor weight_buf, Tensor hx, Tensor? cx, Tensor output, Tensor? grad_output, Tensor? grad_hy, Tensor? grad_cy, int mode, SymInt hidden_size, SymInt proj_size, int num_layers, bool batch_first, float dropout, bool train, bool bidirectional, SymInt[] batch_sizes, Tensor? dropout_state, Tensor reserve, bool[4] output_mask) -> (Tensor, Tensor, Tensor, Tensor[]) + dropout_state: non_differentiable + input: not_implemented("_cudnn_rnn_backward", kCudnnDoubleBackwardMsg) + weight: not_implemented_list("_cudnn_rnn_backward", kCudnnDoubleBackwardMsg) + hx: not_implemented("_cudnn_rnn_backward", kCudnnDoubleBackwardMsg) + cx: not_implemented("_cudnn_rnn_backward", kCudnnDoubleBackwardMsg) + output: not_implemented("_cudnn_rnn_backward", kCudnnDoubleBackwardMsg) + grad_output: not_implemented("_cudnn_rnn_backward", kCudnnDoubleBackwardMsg) + grad_hy: not_implemented("_cudnn_rnn_backward", kCudnnDoubleBackwardMsg) + grad_cy: not_implemented("_cudnn_rnn_backward", kCudnnDoubleBackwardMsg) + +# miopen + +- name: miopen_convolution_transpose(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] output_padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad, self, weight, bias->sym_sizes(), stride, padding, dilation, true, output_padding, groups, grad_input_mask) : std::tuple()" + +- name: miopen_convolution(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad, self, weight, bias->sym_sizes(), stride, padding, dilation, false, std::vector(padding.size(), 0), groups, grad_input_mask) : std::tuple()" + +- name: miopen_depthwise_convolution(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups, bool benchmark, bool deterministic) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad, self, weight, bias->sym_sizes(), stride, padding, dilation, false, std::vector(padding.size(), 0), groups, grad_input_mask) : std::tuple()" + +- name: miopen_batch_norm(Tensor input, Tensor weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float exponential_average_factor, float epsilon) -> (Tensor, Tensor, Tensor) + input, weight, bias: "grad.defined() ? (training ? miopen_batch_norm_backward(input, grad.contiguous(input.suggest_memory_format()), weight, running_mean, running_var, result1, result2, epsilon) : native_batch_norm_backward(grad, input, weight, running_mean, running_var, result1, result2, training, epsilon, grad_input_mask)) : std::tuple()" + result0: batch_norm_jvp(input_p, input_t, weight_p, weight_t, bias_p, bias_t, running_mean, running_var, result1, result2, training, epsilon) + +- name: miopen_batch_norm_backward(Tensor input, Tensor grad_output, Tensor weight, Tensor? running_mean, Tensor? running_var, Tensor? save_mean, Tensor? save_var, float epsilon) -> (Tensor, Tensor, Tensor) + save_mean: not_implemented("miopen_batch_norm_backward save_mean") + save_var: not_implemented("miopen_batch_norm_backward save_var") + input, weight, grad_output: batchnorm_double_backward(input, weight, grads[0], grads[1], grads[2], grad_output, running_mean, running_var, true, epsilon, save_mean, save_var, grad_input_mask) + +- name: miopen_rnn(Tensor input, Tensor[] weight, int weight_stride0, Tensor hx, Tensor? cx, int mode, int hidden_size, int num_layers, bool batch_first, float dropout, bool train, bool bidirectional, int[] batch_sizes, Tensor? dropout_state) -> (Tensor, Tensor, Tensor, Tensor, Tensor) + dropout_state: non_differentiable + output_differentiability: [True, True, True, False, False] + input, hx, cx, weight: "miopen_rnn_backward(input, weight, weight_stride0, result4, hx, cx, result0, grads[0], grads[1], grads[2], mode, hidden_size, num_layers, batch_first, dropout, train, bidirectional, batch_sizes, dropout_state, retain_variables ? result3.clone() : result3, grad_input_mask)" + +- name: miopen_rnn_backward(Tensor input, Tensor[] weight, int weight_stride0, Tensor weight_buf, Tensor hx, Tensor? cx, Tensor output, Tensor? grad_output, Tensor? grad_hy, Tensor? grad_cy, int mode, int hidden_size, int num_layers, bool batch_first, float dropout, bool train, bool bidirectional, int[] batch_sizes, Tensor? dropout_state, Tensor reserve, bool[4] output_mask) -> (Tensor, Tensor, Tensor, Tensor[]) + dropout_state: non_differentiable + +- name: miopen_ctc_loss(Tensor log_probs, Tensor targets, int[] input_lengths, int[] target_lengths, int blank, bool deterministic, bool zero_infinity) -> (Tensor, Tensor) + log_probs: _miopen_ctc_loss_backward(grad, result0, result1, zero_infinity) + +- name: miopen_ctc_loss.Tensor(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank, bool deterministic, bool zero_infinity) -> (Tensor, Tensor) + log_probs: _miopen_ctc_loss_backward(grad, result0, result1, zero_infinity) + +- name: mkldnn_rnn_layer(Tensor input, Tensor weight0, Tensor weight1, Tensor weight2, Tensor weight3, Tensor hx_, Tensor cx_, bool reverse, int[] batch_sizes, int mode, int hidden_size, int num_layers, bool has_biases, bool bidirectional, bool batch_first, bool train) -> (Tensor, Tensor, Tensor, Tensor) + output_differentiability: [True, True, True, False] + input, weight0, weight1, weight2, weight3, hx_, cx_: "GradMode::is_enabled() ? mkldnn_rnn_layer_differentiable_backward(input, weight0, weight1, weight2, weight3, hx_, cx_, result0, result1, result2, grads[0], grads[1], grads[2], reverse, mode, hidden_size, num_layers, has_biases, train, bidirectional, batch_sizes, batch_first, result3) : mkldnn_rnn_layer_backward(input, weight0, weight1, weight2, weight3, hx_, cx_, result0, result1, result2, grads[0], grads[1], grads[2], reverse, mode, hidden_size, num_layers, has_biases, train, bidirectional, batch_sizes, batch_first, result3)" + +- name: mkldnn_rnn_layer_backward(Tensor input, Tensor weight1, Tensor weight2, Tensor weight3, Tensor weight4, Tensor hx_, Tensor cx_tmp, Tensor output, Tensor hy_, Tensor cy_, Tensor? grad_output, Tensor? grad_hy, Tensor? grad_cy, bool reverse, int mode, int hidden_size, int num_layers, bool has_biases, bool train, bool bidirectional, int[] batch_sizes, bool batch_first, Tensor workspace) -> (Tensor, Tensor, Tensor, Tensor, Tensor, Tensor, Tensor) + +# mkldnn +- name: mkldnn_convolution(Tensor self, Tensor weight, Tensor? bias, SymInt[] padding, SymInt[] stride, SymInt[] dilation, SymInt groups) -> Tensor + self, weight, bias: "grad.defined() ? convolution_backward_symint(grad, self, weight, bias->sym_sizes(), stride, padding, dilation, /*transposed=*/ false, /*output_padding=*/ std::vector(padding.size(), 0), groups, grad_input_mask) : std::tuple()" + +- name: mkldnn_linear(Tensor self, Tensor weight, Tensor? bias=None) -> Tensor + self, weight, bias: mkldnn_linear_backward(self, grad, weight, grad_input_mask) + +- name: mkldnn_max_pool2d(Tensor self, int[2] kernel_size, int[2] stride=[], int[2] padding=0, int[2] dilation=1, bool ceil_mode=False) -> Tensor + self: mkldnn_max_pool2d_backward(grad, result, self, kernel_size, stride, padding, dilation, ceil_mode) + +- name: mkldnn_max_pool3d(Tensor self, int[3] kernel_size, int[3] stride=[], int[3] padding=0, int[3] dilation=1, bool ceil_mode=False) -> Tensor + self: mkldnn_max_pool3d_backward(grad, result, self, kernel_size, stride, padding, dilation, ceil_mode) + +- name: mkldnn_adaptive_avg_pool2d(Tensor self, int[2] output_size) -> Tensor + self: mkldnn_adaptive_avg_pool2d_backward(grad, self) + +- name: _mkldnn_reshape(Tensor self, int[] shape) -> Tensor + self: grad.reshape_symint(self.sym_sizes()) + +# NestedTensor +- name: _nested_tensor_from_tensor_list(Tensor[] list, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + list: "grad.defined()? at::unbind(grad) : std::vector(list.size())" + +- name: _nested_tensor_from_mask(Tensor t, Tensor mask, bool mask_check=True) -> Tensor + t: grad.to_padded_tensor_symint(0, t.sym_sizes()) + mask: non_differentiable + +- name: _nested_from_padded(Tensor padded, Tensor cpu_nested_shape_example, bool fuse_transform_0213=False) -> Tensor + padded: _nested_from_padded_backward(grad, padded, fuse_transform_0213) + cpu_nested_shape_example: non_differentiable + +- name: to_padded_tensor(Tensor self, float padding, SymInt[]? output_size=None) -> Tensor + self: "self.layout() == c10::kJagged ? at::_nested_from_padded_tensor_symint(grad, at::_nested_get_offsets(self), at::_nested_get_jagged_dummy(self), at::_nested_get_ragged_idx(self), at::_nested_get_min_seqlen(self).defined() ? std::optional(at::_nested_get_min_seqlen(self)) : ::std::nullopt, at::_nested_get_max_seqlen(self).defined() ? std::optional(at::_nested_get_max_seqlen(self)) : ::std::nullopt, std::optional(at::_nested_get_values(self).sym_size(0))) : at::_nested_from_padded(grad, self._nested_tensor_size())" + padding: non_differentiable + +- name: _nested_from_padded_tensor(Tensor padded, Tensor offsets, Tensor dummy, int ragged_idx=1, Tensor? min_seqlen=None, Tensor? max_seqlen=None, SymInt? sum_S=None) -> Tensor + padded: grad.to_padded_tensor_symint(0.0, at::OptionalArrayRef(padded.sym_sizes())) + offsets: non_differentiable + dummy: non_differentiable + +- name: _nested_view_from_buffer(Tensor(a) self, Tensor nested_size, Tensor nested_strides, Tensor offsets) -> Tensor(a) + self: grad.values() + nested_size: non_differentiable + nested_strides: non_differentiable + +- name: _nested_view_from_jagged(Tensor(a) self, Tensor offsets, Tensor dummy, Tensor? lengths=None, int ragged_idx=1, Tensor? min_seqlen=None, Tensor? max_seqlen=None) -> Tensor(a) + self: grad.values() + offsets: non_differentiable + lengths: non_differentiable + dummy: non_differentiable + min_seqlen: non_differentiable + max_seqlen: non_differentiable + +- name: _nested_get_values(Tensor(a) self) -> Tensor(a) + self: "_nested_view_from_jagged(grad, at::_nested_get_offsets(self), at::_nested_get_jagged_dummy(self), at::_nested_get_lengths(self), at::_nested_get_ragged_idx(self), at::_nested_get_min_seqlen(self).defined() ? std::optional(at::_nested_get_min_seqlen(self)) : ::std::nullopt, at::_nested_get_max_seqlen(self).defined() ? std::optional(at::_nested_get_max_seqlen(self)) : ::std::nullopt)" + +# Transformer +- name: _safe_softmax(Tensor self, int dim, ScalarType? dtype=None) -> Tensor + self: _softmax_backward_data(grad, result, dim, self.scalar_type()) + result: result * (self_t - safe_logsumexp_jvp(self_p, self_t, {dim}, true)) + +- name: _scaled_dot_product_efficient_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_bias, bool compute_log_sumexp, float dropout_p=0.0, bool is_causal=False, *, float? scale=None) -> (Tensor output, Tensor log_sumexp, Tensor philox_seed, Tensor philox_offset) + output_differentiability: [True, False, False, False] + query, key, value, attn_bias: _scaled_dot_product_efficient_attention_backward(grad, query, key, value, attn_bias, output, log_sumexp, philox_seed, philox_offset, dropout_p, grad_input_mask, is_causal, scale) + +- name: _scaled_dot_product_flash_attention(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor rng_state, Tensor unused, Tensor debug_attn_mask) + output_differentiability: [True, False, False, False, False, False, False, False, False] + query, key, value: _scaled_dot_product_flash_attention_backward_symint(grad, query, key, value, output, logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, rng_state, unused, scale) + +- name: _scaled_dot_product_flash_attention_for_cpu(Tensor query, Tensor key, Tensor value, float dropout_p=0.0, bool is_causal=False, *, Tensor? attn_mask=None, float? scale=None) -> (Tensor output, Tensor logsumexp) + output_differentiability: [True, False] + query, key, value: _scaled_dot_product_flash_attention_for_cpu_backward(grad, query, key, value, output, logsumexp, dropout_p, is_causal, attn_mask, scale) + +- name: _flash_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, float dropout_p, bool is_causal, bool return_debug_mask, *, float? scale=None, SymInt? window_size_left=None, SymInt? window_size_right=None, Tensor? seqused_k=None, Tensor? alibi_slopes=None, Tensor? block_table=None, int? num_splits=None) -> (Tensor output, Tensor softmax_logsumexp, Tensor rng_state, Tensor unused, Tensor debug_attn_mask) + output_differentiability: [True, False, False, False, False] + query, key, value: _flash_attention_backward_symint(grad, query, key, value, output, softmax_logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, rng_state, unused, scale, window_size_left, window_size_right) + +- name: _efficient_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? bias, Tensor? cu_seqlens_q, Tensor? cu_seqlens_k, SymInt? max_seqlen_q, SymInt? max_seqlen_k, float dropout_p, int custom_mask_type, bool compute_log_sumexp=False, *, float? scale=None, Tensor? seqlen_k=None, int? window_size=None) -> (Tensor output, Tensor logsumexp, Tensor philox_seed, Tensor philox_offset, SymInt max_seqlen_batch_q, SymInt max_seqlen_batch_k) + output_differentiability: [True, False, False, False, False, False] + query, key, value, bias: _efficient_attention_backward_symint(grad, query, key, value, bias, output, cu_seqlens_q, cu_seqlens_k, max_seqlen_batch_q, max_seqlen_batch_k, logsumexp, dropout_p, philox_seed, philox_offset, custom_mask_type, bias.requires_grad(), scale) + +- name: _cudnn_attention_forward(Tensor query, Tensor key, Tensor value, Tensor? attn_bias, Tensor? cum_seq_q, Tensor? cum_seq_k, SymInt max_q, SymInt max_k, bool compute_log_sumexp, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask) + output_differentiability: [True, False, False, False, False, False, False, False, False] + query, key, value: _cudnn_attention_backward_symint(grad, query, key, value, output, logsumexp, philox_seed, philox_offset, attn_bias, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, scale) + +- name: _scaled_dot_product_cudnn_attention(Tensor query, Tensor key, Tensor value, Tensor? attn_bias, bool compute_log_sumexp, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask) + output_differentiability: [True, False, False, False, False, False, False, False, False] + query, key, value: _scaled_dot_product_cudnn_attention_backward_symint(grad, query, key, value, output, logsumexp, philox_seed, philox_offset, attn_bias, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, scale) + +- name: _scaled_dot_product_fused_attention_overrideable(Tensor query, Tensor key, Tensor value, Tensor? attn_bias=None, float dropout_p=0.0, bool is_causal=False, bool return_debug_mask=False, *, float? scale=None) -> (Tensor output, Tensor logsumexp, Tensor cum_seq_q, Tensor cum_seq_k, SymInt max_q, SymInt max_k, Tensor philox_seed, Tensor philox_offset, Tensor debug_attn_mask) + output_differentiability: [True, False, False, False, False, False, False, False, False] + query, key, value, attn_bias: _scaled_dot_product_fused_attention_overrideable_backward_symint(grad, query, key, value, attn_bias, grad_input_mask, output, logsumexp, cum_seq_q, cum_seq_k, max_q, max_k, dropout_p, is_causal, philox_seed, philox_offset, scale) + +# fft +- name: _fft_r2c(Tensor self, int[] dim, int normalization, bool onesided) -> Tensor + self: fft_r2c_backward(grad, dim, normalization, onesided, self.sym_size(dim.back())) + result: auto_linear + +- name: _fft_c2r(Tensor self, int[] dim, int normalization, SymInt last_dim_size) -> Tensor + self: fft_c2r_backward(grad, dim, normalization) + result: auto_linear + +- name: _fft_c2c(Tensor self, SymInt[] dim, int normalization, bool forward) -> Tensor + self: _fft_c2c_symint(grad, dim, normalization, !forward) + result: auto_linear + +- name: unbind.int(Tensor(a -> *) self, int dim=0) -> Tensor(a)[] + dispatch: + Default: + self: unbind_backward(grads, dim) + result: auto_linear + AutogradNestedTensor: + self: "self.layout() == c10::kJagged ? unbind_backward_nested_jagged(grads, self, dim) : unbind_backward_nested(grads, at::native::get_nested_tensor_impl(self)->get_nested_sizes(), dim, self.options())" + result: auto_linear + +- name: stack(Tensor[] tensors, int dim=0) -> Tensor + tensors: stack_tensors_backward(grad, dim, to_args_scalartypes(tensors)) + result: stack_jvp(tensors, dim) + +# fused RNN kernels + +# Only frst two of _thnn_fused_lstm_cell outputs can have gradients. +# _thnn_fused_lstm_cell outputs: (hy, cy, workspace) +- name: _thnn_fused_lstm_cell(Tensor input_gates, Tensor hidden_gates, Tensor cx, Tensor? input_bias=None, Tensor? hidden_bias=None) -> (Tensor, Tensor, Tensor) + output_differentiability: [True, True, False] + input_gates, hidden_gates, cx, input_bias, hidden_bias: "GradMode::is_enabled() ? _thnn_differentiable_lstm_cell_backward(grads[0], grads[1], input_gates, hidden_gates, input_bias, hidden_bias, cx, result1) : _thnn_fused_lstm_cell_backward(grads[0], grads[1], cx, result1, result2, input_bias.defined())" + +- name: _thnn_fused_gru_cell(Tensor input_gates, Tensor hidden_gates, Tensor hx, Tensor? input_bias=None, Tensor? hidden_bias=None) -> (Tensor, Tensor) + input_gates, hidden_gates, hx, input_bias, hidden_bias: "grad.defined() ? (GradMode::is_enabled() ? _thnn_differentiable_gru_cell_backward(grad, input_gates, hidden_gates, hx, input_bias, hidden_bias) : _thnn_fused_gru_cell_backward(grad, result1, input_bias.defined())) : std::tuple()" + +# PackedSequence helpers +- name: _pack_padded_sequence(Tensor input, Tensor lengths, bool batch_first) -> (Tensor, Tensor) + input: _pack_padded_sequence_backward_symint(grad, input.sym_sizes(), result1, batch_first) + +# TH wrappers +- name: eq.Scalar(Tensor self, Scalar other) -> Tensor + output_differentiability: [False] + +- name: eq.Tensor(Tensor self, Tensor other) -> Tensor + output_differentiability: [False] + +- name: ge.Scalar(Tensor self, Scalar other) -> Tensor + output_differentiability: [False] + +- name: ge.Tensor(Tensor self, Tensor other) -> Tensor + output_differentiability: [False] + +- name: gt.Scalar(Tensor self, Scalar other) -> Tensor + output_differentiability: [False] + +- name: gt.Tensor(Tensor self, Tensor other) -> Tensor + output_differentiability: [False] + +- name: le.Scalar(Tensor self, Scalar other) -> Tensor + output_differentiability: [False] + +- name: le.Tensor(Tensor self, Tensor other) -> Tensor + output_differentiability: [False] + +- name: lt.Scalar(Tensor self, Scalar other) -> Tensor + output_differentiability: [False] + +- name: lt.Tensor(Tensor self, Tensor other) -> Tensor + output_differentiability: [False] + +- name: ne.Scalar(Tensor self, Scalar other) -> Tensor + output_differentiability: [False] + +- name: ne.Tensor(Tensor self, Tensor other) -> Tensor + output_differentiability: [False] + +- name: multinomial(Tensor self, SymInt num_samples, bool replacement=False, *, Generator? generator=None) -> Tensor + output_differentiability: [False] + +- name: nonzero(Tensor self) -> Tensor + output_differentiability: [False] + +- name: segment_reduce(Tensor data, str reduce, *, Tensor? lengths=None, Tensor? indices=None, Tensor? offsets=None, int axis=0, bool unsafe=False, Scalar? initial=None) -> Tensor + data: _segment_reduce_backward(grad, result, data, reduce, lengths, offsets, axis, initial) + +- name: _pin_memory(Tensor self, Device? device=None) -> Tensor + self: grad + +- name: _new_zeros_with_same_feature_meta(Tensor self, Tensor other, *, int self_num_batch_dims=0) -> Tensor + self: non_differentiable + other: non_differentiable + output_differentiability: [False] + +- name: _test_warn_in_autograd(Tensor self) -> Tensor + self: warn_backwards(grad) + +- name: _test_autograd_multiple_dispatch.fullcoverage(Tensor self) -> Tensor + dispatch: + Default: + self: grad.expand_symint(self.sym_sizes()) + 1 + result: auto_linear + AutogradNestedTensor: + self: grad.mul(grad) + AutogradCUDA: + self: grad.expand_symint(self.sym_sizes()) * 2 + +- name: _test_autograd_multiple_dispatch.ntonly(Tensor self, bool b) -> Tensor + dispatch: + AutogradNestedTensor: + self: grad.mul(grad).add(grad) + +- name: _test_autograd_multiple_dispatch_view(Tensor(a) self) -> Tensor(a) + dispatch: + Default: + self: grad.reshape_as(self) + AutogradCUDA: + self: grad.reshape_as(self) + 1 + +- name: _efficientzerotensor(SymInt[] size, *, ScalarType? dtype=None, Layout? layout=None, Device? device=None, bool? pin_memory=None) -> Tensor + output_differentiability: [False] + +- name: scatter_reduce.two(Tensor self, int dim, Tensor index, Tensor src, str reduce, *, bool include_self=True) -> Tensor + self, src: scatter_reduce_backward(grad, self, dim, index, src, reduce, include_self, result) + index: non_differentiable + result: scatter_reduce_jvp(self_p, self_t, dim, index, src_p, src_t, reduce, include_self, result) + +- name: special_airy_ai(Tensor x) -> Tensor + x: non_differentiable + +- name: special_bessel_j0(Tensor self) -> Tensor + self: non_differentiable + +- name: special_bessel_j1(Tensor self) -> Tensor + self: non_differentiable + +- name: special_bessel_y0(Tensor self) -> Tensor + self: non_differentiable + +- name: special_bessel_y1(Tensor self) -> Tensor + self: non_differentiable + +- name: special_chebyshev_polynomial_t(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_chebyshev_polynomial_t.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_chebyshev_polynomial_t.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_chebyshev_polynomial_u(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_chebyshev_polynomial_u.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_chebyshev_polynomial_u.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_chebyshev_polynomial_v(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_chebyshev_polynomial_v.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_chebyshev_polynomial_v.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_chebyshev_polynomial_w(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_chebyshev_polynomial_w.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_chebyshev_polynomial_w.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_hermite_polynomial_h(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_hermite_polynomial_h.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_hermite_polynomial_h.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_hermite_polynomial_he(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_hermite_polynomial_he.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_hermite_polynomial_he.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_laguerre_polynomial_l(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_laguerre_polynomial_l.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_laguerre_polynomial_l.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_legendre_polynomial_p(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_legendre_polynomial_p.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_legendre_polynomial_p.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_modified_bessel_i0(Tensor self) -> Tensor + self: non_differentiable + +- name: special_modified_bessel_i1(Tensor self) -> Tensor + self: non_differentiable + +- name: special_modified_bessel_k0(Tensor self) -> Tensor + self: non_differentiable + +- name: special_modified_bessel_k1(Tensor self) -> Tensor + self: non_differentiable + +- name: special_scaled_modified_bessel_k0(Tensor x) -> Tensor + x: non_differentiable + +- name: special_scaled_modified_bessel_k1(Tensor x) -> Tensor + x: non_differentiable + +- name: special_shifted_chebyshev_polynomial_t(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_shifted_chebyshev_polynomial_t.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_shifted_chebyshev_polynomial_t.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_shifted_chebyshev_polynomial_u(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_shifted_chebyshev_polynomial_u.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_shifted_chebyshev_polynomial_u.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_shifted_chebyshev_polynomial_v(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_shifted_chebyshev_polynomial_v.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_shifted_chebyshev_polynomial_v.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_shifted_chebyshev_polynomial_w(Tensor x, Tensor n) -> Tensor + x: non_differentiable + n: non_differentiable + +- name: special_shifted_chebyshev_polynomial_w.x_scalar(Scalar x, Tensor n) -> Tensor + n: non_differentiable + +- name: special_shifted_chebyshev_polynomial_w.n_scalar(Tensor x, Scalar n) -> Tensor + x: non_differentiable + +- name: special_spherical_bessel_j0(Tensor x) -> Tensor + x: non_differentiable + +- name: _reshape_copy(Tensor self, SymInt[] size) -> Tensor + self: grad.reshape_symint(self.sym_sizes()) + result: auto_linear + +- name: narrow_copy(Tensor self, int dim, SymInt start, SymInt length) -> Tensor + self: slice_backward_wrapper(grad, self.sym_sizes(), dim, start, start + length, 1) + result: auto_linear + +# note(crcrpar): `torchgen/api/autograd` logic would unwantedly replace substrings of `self` and `other` of function names. +- name: _foreach_div.List(Tensor[] self, Tensor[] other) -> Tensor[] + self: div_tensor_self_backward(grads[i], other[i], self[i].scalar_type()) + other: div_tensor_other_backward(grads[i], self[i], other[i]) + result: (self_t - other_t * result[i]) / other_p + +- name: _foreach_pow.List(Tensor[] self, Tensor[] exponent) -> Tensor[] + self: pow_backward_self(grads[i], self[i], exponent[i]) + exponent: pow_backward_exponent(grads[i], self[i], exponent[i], result[i]) + result: (pow_backward_self(self_t.conj(), self_p, exponent_p) + pow_backward_exponent(exponent_t.conj(), self_p, exponent_p, result[i])).conj() + +- name: _foreach_pow.ScalarList(Tensor[] self, Scalar[] exponent) -> Tensor[] + self: pow_backward(grads[i], self[i], exponent[i]) + result: pow_backward(self_t.conj(), self_p, exponent[i]).conj() + +- name: _foreach_pow.ScalarAndTensor(Scalar self, Tensor[] exponent) -> Tensor[] + exponent: pow_backward_exponent(grads[i], self, exponent[i], result[i]) + +# note(crcrpar): following definitions seem necessary because the reference native functions +# of `maximum` and `minimum` don't have the overload def with Scalar as their second argument. +- name: _foreach_minimum.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + self: at::where(self[i] == scalar, grads[i] / 2, grads[i]).masked_fill_(self[i] > scalar, 0) + result: scalar + at::where(self_p == scalar, at::scalar_tensor(0.5, result[i].options()), (self_p < scalar).to(result[i].scalar_type())) * (self_t - scalar) + +- name: _foreach_minimum.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + self: at::where(self[i] == scalars[i], grads[i] / 2, grads[i]).masked_fill_(self[i] > scalars[i], 0) + result: scalars[i] + at::where(self_p == scalars[i], at::scalar_tensor(0.5, result[i].options()), (self_p < scalars[i]).to(result[i].scalar_type())) * (self_t - scalars[i]) + +- name: _foreach_maximum.Scalar(Tensor[] self, Scalar scalar) -> Tensor[] + self: at::where(self[i] == scalar, grads[i] / 2, grads[i]).masked_fill_(self[i] < scalar, 0) + result: scalar + at::where(self_p == scalar, at::scalar_tensor(0.5, result[i].options()), (self_p > scalar).to(result[i].scalar_type())) * (self_t - scalar) + +- name: _foreach_maximum.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[] + self: at::where(self[i] == scalars[i], grads[i] / 2, grads[i]).masked_fill_(self[i] < scalars[i], 0) + result: scalars[i] + at::where(self_p == scalars[i], at::scalar_tensor(0.5, result[i].options()), (self_p > scalars[i]).to(result[i].scalar_type())) * (self_t - scalars[i]) + +# note(crcrpar): forward-mode AD is tricky for a simple string replace to handle: +# formula.replace("p", "ord") produces `norm_jvord(self_ord, self_t, ord, result)` +- name: _foreach_norm.Scalar(Tensor[] self, Scalar ord=2, ScalarType? dtype=None) -> Tensor[] + self: norm_backward(grads[i], self[i], ord, result[i]) + result: norm_jvp(self_p, self_t, ord, result[i]) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_annotated_fn_args.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_annotated_fn_args.py new file mode 100644 index 0000000000000000000000000000000000000000..2f61209fa6fd0041b732f1400e1162d2f124ad34 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_annotated_fn_args.py @@ -0,0 +1,134 @@ +""" +For procedural tests needed for __torch_function__, we use this function +to export method names and signatures as needed by the tests in +test/test_overrides.py. + +python -m tools.autograd.gen_annotated_fn_args \ + aten/src/ATen/native/native_functions.yaml \ + aten/src/ATen/native/tags.yaml \ + $OUTPUT_DIR \ + tools/autograd + +Where $OUTPUT_DIR is where you would like the files to be +generated. In the full build system, OUTPUT_DIR is +torch/testing/_internal/generated +""" + +from __future__ import annotations + +import argparse +import os +import textwrap +from collections import defaultdict +from typing import Any, TYPE_CHECKING + +import torchgen.api.python as python +from torchgen.context import with_native_function +from torchgen.gen import parse_native_yaml +from torchgen.utils import FileManager + +from .gen_python_functions import ( + is_py_fft_function, + is_py_linalg_function, + is_py_nn_function, + is_py_special_function, + is_py_torch_function, + is_py_variable_method, + should_generate_py_binding, +) + + +if TYPE_CHECKING: + from collections.abc import Sequence + + from torchgen.model import Argument, BaseOperatorName, NativeFunction + + +def gen_annotated( + native_yaml_path: str, tags_yaml_path: str, out: str, autograd_dir: str +) -> None: + native_functions = parse_native_yaml( + native_yaml_path, tags_yaml_path + ).native_functions + mappings = ( + (is_py_torch_function, "torch._C._VariableFunctions"), + (is_py_nn_function, "torch._C._nn"), + (is_py_linalg_function, "torch._C._linalg"), + (is_py_special_function, "torch._C._special"), + (is_py_fft_function, "torch._C._fft"), + (is_py_variable_method, "torch.Tensor"), + ) + annotated_args: list[str] = [] + for pred, namespace in mappings: + groups: dict[BaseOperatorName, list[NativeFunction]] = defaultdict(list) + for f in native_functions: + if not should_generate_py_binding(f) or not pred(f): + continue + groups[f.func.name.name].append(f) + for group in groups.values(): + for f in group: + annotated_args.append(f"{namespace}.{gen_annotated_args(f)}") + + template_path = os.path.join(autograd_dir, "templates") + fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + fm.write_with_template( + "annotated_fn_args.py", + "annotated_fn_args.py.in", + lambda: { + "annotated_args": textwrap.indent("\n".join(annotated_args), " "), + }, + ) + + +@with_native_function +def gen_annotated_args(f: NativeFunction) -> str: + def _get_kwargs_func_exclusion_list() -> list[str]: + # functions that currently don't work with kwargs in test_overrides.py + return [ + "diagonal", + "round_", + "round", + "scatter_", + ] + + def _add_out_arg( + out_args: list[dict[str, Any]], args: Sequence[Argument], *, is_kwarg_only: bool + ) -> None: + for arg in args: + if arg.default is not None: + continue + out_arg: dict[str, Any] = {} + out_arg["is_kwarg_only"] = str(is_kwarg_only) + out_arg["name"] = arg.name + out_arg["simple_type"] = python.argument_type_str( + arg.type, simple_type=True + ) + size_t = python.argument_type_size(arg.type) + if size_t: + out_arg["size"] = size_t + out_args.append(out_arg) + + out_args: list[dict[str, Any]] = [] + _add_out_arg(out_args, f.func.arguments.flat_positional, is_kwarg_only=False) + if f"{f.func.name.name}" not in _get_kwargs_func_exclusion_list(): + _add_out_arg(out_args, f.func.arguments.flat_kwarg_only, is_kwarg_only=True) + + return f"{f.func.name.name}: {repr(out_args)}," + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate annotated_fn_args script") + parser.add_argument( + "native_functions", metavar="NATIVE", help="path to native_functions.yaml" + ) + parser.add_argument("tags", metavar="TAGS", help="path to tags.yaml") + parser.add_argument("out", metavar="OUT", help="path to output directory") + parser.add_argument( + "autograd", metavar="AUTOGRAD", help="path to template directory" + ) + args = parser.parse_args() + gen_annotated(args.native_functions, args.tags, args.out, args.autograd) + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_autograd.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_autograd.py new file mode 100644 index 0000000000000000000000000000000000000000..d93d3f4cab4a6f37c0c81c548b4da3b6c5b9dc95 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_autograd.py @@ -0,0 +1,147 @@ +""" +To run this file by hand from the root of the PyTorch +repository, run: + +python -m tools.autograd.gen_autograd \ + aten/src/ATen/native/native_functions.yaml \ + aten/src/ATen/native/tags.yaml \ + $OUTPUT_DIR \ + tools/autograd + +Where $OUTPUT_DIR is where you would like the files to be +generated. In the full build system, OUTPUT_DIR is +torch/csrc/autograd/generated/ +""" + +# gen_autograd.py generates C++ autograd functions and Python bindings. +# +# It delegates to the following scripts: +# +# gen_autograd_functions.py: generates subclasses of torch::autograd::Node +# gen_variable_type.py: generates VariableType.h which contains all tensor methods +# gen_python_functions.py: generates Python bindings to THPVariable +# + +from __future__ import annotations + +import argparse +import os + +from torchgen.api import cpp +from torchgen.api.autograd import ( + match_differentiability_info, + NativeFunctionWithDifferentiabilityInfo, +) +from torchgen.gen import parse_native_yaml +from torchgen.selective_build.selector import SelectiveBuilder + +from . import gen_python_functions +from .gen_autograd_functions import ( + gen_autograd_functions_lib, + gen_autograd_functions_python, +) +from .gen_inplace_or_view_type import gen_inplace_or_view_type +from .gen_trace_type import gen_trace_type +from .gen_variable_factories import gen_variable_factories +from .gen_variable_type import gen_variable_type +from .gen_view_funcs import gen_view_funcs +from .load_derivatives import load_derivatives + + +def gen_autograd( + native_functions_path: str, + tags_path: str, + out: str, + autograd_dir: str, + operator_selector: SelectiveBuilder, + disable_autograd: bool = False, +) -> None: + # Parse and load derivatives.yaml + differentiability_infos, used_dispatch_keys = load_derivatives( + os.path.join(autograd_dir, "derivatives.yaml"), native_functions_path, tags_path + ) + + template_path = os.path.join(autograd_dir, "templates") + + native_funcs = parse_native_yaml(native_functions_path, tags_path).native_functions + fns = sorted( + filter( + operator_selector.is_native_function_selected_for_training, native_funcs + ), + key=lambda f: cpp.name(f.func), + ) + fns_with_diff_infos: list[NativeFunctionWithDifferentiabilityInfo] = ( + match_differentiability_info(fns, differentiability_infos) + ) + + # Generate VariableType.h/cpp + if not disable_autograd: + gen_variable_type( + out, + native_functions_path, + tags_path, + fns_with_diff_infos, + template_path, + used_dispatch_keys, + ) + + gen_inplace_or_view_type( + out, native_functions_path, tags_path, fns_with_diff_infos, template_path + ) + + # operator filter not applied as tracing sources are excluded in selective build + gen_trace_type(out, native_funcs, template_path) + # Generate Functions.h/cpp + gen_autograd_functions_lib(out, differentiability_infos, template_path) + + # Generate variable_factories.h + gen_variable_factories(out, native_functions_path, tags_path, template_path) + + # Generate ViewFuncs.h/cpp + gen_view_funcs(out, fns_with_diff_infos, template_path) + + +def gen_autograd_python( + native_functions_path: str, + tags_path: str, + out: str, + autograd_dir: str, +) -> None: + differentiability_infos, _ = load_derivatives( + os.path.join(autograd_dir, "derivatives.yaml"), native_functions_path, tags_path + ) + + template_path = os.path.join(autograd_dir, "templates") + + # Generate Functions.h/cpp + gen_autograd_functions_python(out, differentiability_infos, template_path) + + # Generate Python bindings + deprecated_path = os.path.join(autograd_dir, "deprecated.yaml") + gen_python_functions.gen( + out, native_functions_path, tags_path, deprecated_path, template_path + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate autograd C++ files script") + parser.add_argument( + "native_functions", metavar="NATIVE", help="path to native_functions.yaml" + ) + parser.add_argument("tags", metavar="NATIVE", help="path to tags.yaml") + parser.add_argument("out", metavar="OUT", help="path to output directory") + parser.add_argument( + "autograd", metavar="AUTOGRAD", help="path to autograd directory" + ) + args = parser.parse_args() + gen_autograd( + args.native_functions, + args.tags, + args.out, + args.autograd, + SelectiveBuilder.get_nop_selector(), + ) + + +if __name__ == "__main__": + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_autograd_functions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_autograd_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..56e622d38d65d61aa193cb2a22771b7ffbc49901 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_autograd_functions.py @@ -0,0 +1,1082 @@ +# Generates C++ autograd functions for the derivatives of ATen operations +# +# This writes two files: +# Functions.h/cpp: subclasses of autograd::Node +# python_functions.h/cpp: Python bindings for the above classes +# + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from torchgen.api.autograd import ( + Derivative, + DifferentiabilityInfo, + SavedAttribute, + uses_retain_variables, + uses_single_grad, +) +from torchgen.api.types import ( + ArrayRefCType, + BaseCppType, + BaseCType, + Binding, + boolT, + doubleT, + intArrayRefT, + iTensorListRefT, + ListCType, + longT, + MutRefCType, + OptionalCType, + optionalIntArrayRefT, + optionalSymIntArrayRefT, + scalarT, + stringT, + symIntArrayRefT, + SymIntT, + TENSOR_LIST_LIKE_CTYPES, + tensorListT, + tensorT, + VectorCType, +) +from torchgen.code_template import CodeTemplate +from torchgen.model import Argument, FunctionSchema +from torchgen.utils import FileManager + +from .gen_inplace_or_view_type import VIEW_FUNCTIONS + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +FUNCTION_DECLARATION = CodeTemplate( + """\ +#ifdef _WIN32 +struct ${op} : public ${superclass} { + TORCH_API ${op}() = default; +#else +struct TORCH_API ${op} : public ${superclass} { +#endif + using ${superclass}::${superclass}; + variable_list apply(variable_list&& grads) override; + std::string name() const override { return "${op}"; } + void release_variables() override { + ${thread_lock} + ${release_variables} + } + ${will_release_variables} + void compiled_args(CompiledNodeArgs& args) const override; + variable_list apply_with_saved(const variable_list& inputs, SwapSavedVariables& saved) override; + ${saved_variables} + ${saved_list_sizes} +}; +""" +) + +WILL_RELEASE_VARIABLES = CodeTemplate( + """\ +bool retain_variables = true; +void will_release_variables() override { + retain_variables = false; +} +""" +) + +# We generate e.g. MulBackward0::apply and have that call into +# MulBackward0_apply_functional. The apply_functional is a pure function, +# that is, it does not rely on global state. MulBackward0::apply +# is responsible for querying the autograd engine for which outputs should +# be computed (needs_input_grad), applying locks, +# and unpacking saved variables to pass to MulBackward0_apply_functional. +# +# needs_input_grad is a mapping from input index to if that input needs +# gradients computed. For operators that take in List[Tensor], the List[Tensor] +# is one element in the needs_input_grad that specifies if *any* of the +# List[Tensor] needs input grad. In theory this could be optimized. +FUNCTION_DEFINITION = CodeTemplate( + """\ +static variable_list ${op}_apply_functional( + variable_list&& grads, + std::array needs_input_grad${,apply_functional_args_signature}) +{ + IndexRangeGenerator gen; + ${compute_index_ranges} + variable_list grad_inputs(gen.size()); + ${body} + return grad_inputs; +} +inline variable_list ${op}_apply_functional_ivalue(const variable_list& grads, const ivalue_list& args) +{ +#ifdef C10_MOBILE + TORCH_INTERNAL_ASSERT(false, "compiled autograd doesn't work on mobile"); +#else + auto packed_args = PackedArgs(args); + auto needs_input_grad = packed_args.unpack>(); + ${unpack_ivalues} + return ${op}_apply_functional(variable_list(grads), needs_input_grad${,apply_functional_args}); +#endif +} + +variable_list ${op}::apply(variable_list&& grads) { + ${thread_lock} + ${asserts} + ${unpacks} + ${compute_needs_input_grad} + return ${op}_apply_functional(std::move(grads), needs_input_grad${,apply_functional_args}); +} + +void ${op}::compiled_args(CompiledNodeArgs& args) const { + ${compiled_args} +} +variable_list ${op}::apply_with_saved(const variable_list& grads, SwapSavedVariables& saved) { +#ifdef C10_MOBILE + TORCH_INTERNAL_ASSERT(false, "compiled autograd doesn't work on mobile"); +#else + ${apply_with_saved_before} + + static bool called = false; + if (!called) { + called = true; + ${compute_schema} + const auto& pyinterface = torch::dynamo::autograd::getPyCompilerInterface(); + pyinterface->bind_function(saved.get_py_compiler(), name(), ${op}_apply_functional_ivalue, schema); + } + + variable_list output_result; + + PackedArgs packed_args; + ${asserts} + ${unpacks} + ${compute_needs_input_grad} + packed_args.pack(needs_input_grad); + ${get_packed_args} + + output_result = compiled_autograd_apply_functional(packed_args, next_edges(), saved, grads, name()); + + ${apply_with_saved_after} + return output_result; +#endif +} + +""" +) + +GRAD_INPUT_MASK = CodeTemplate( + """\ + auto grad_input_mask = std::array{ + ${masks} + }; +""" +) + +COMPUTE_NEEDS_INPUT_GRAD = CodeTemplate( + """\ +IndexRangeGenerator gen; +${compute_index_ranges} +auto needs_input_grad = std::array{ + ${masks} +};\ +""" +) + + +DERIVATIVE_SINGLE = CodeTemplate( + """\ +if (needs_input_grad[/*${name}*/${idx}]) { + auto grad_result = ${derivative}; + copy_range(grad_inputs, ${name}_ix, grad_result); +} +""" +) + +# note(crcrpar): `self` argument and other optional positional argument +# of foreach functions are basically a list of n `Tensor`s thus iterating over +# `grads` in order to utilize and apply the existing derivative definitions +# to each `Tensor`(s) of `self`, and the others. +DERIVATIVE_SINGLE_FOREACH = CodeTemplate( + """\ +if (needs_input_grad[/*${name}*/${idx}]) { // ${name} + std::vector grad_result; + grad_result.reserve(grads.size()); + for (const auto & i : c10::irange(grads.size())) { + if (grads[i].defined()) { + grad_result.emplace_back(${derivative}); + } else { + grad_result.emplace_back(Tensor()); + } + } + copy_range(grad_inputs, ${name}_ix, grad_result); +} +""" +) + +DERIVATIVE_MULTI_COPY_RANGE = CodeTemplate( + """\ + if (needs_input_grad[/*${name}*/${idx}]) { + copy_range(grad_inputs, ${name}_ix, std::get<${i}>(grad_result)); + } +""" +) + +DERIVATIVE_MULTI = CodeTemplate( + """\ +if (${needs_input_grad}) { + ${grad_input_mask} + auto grad_result = ${derivative}; + ${copy_ranges} +} +""" +) + +# Generates python bindings +# +# This generates the definitions for: +# (1) The PyTypeObject for each backward grad_fn subclassing Node +# (2) The entry for PyTypeObject's tp_getset slot (an array of PyGetSetDef structs) +# We generate one PyGetSetDef struct for each of grad_fn's saved inputs and outputs +# Each PyGetSetDef has a function ptr to a getter, also defined here (3). +# (3) Getters for each of grad_fn's saved inputs and outputs. +# +PY_FUNCTION_DEFINITION = CodeTemplate( + """\ +static PyTypeObject ${op}Class; +addClass<${op}>(module, ${op}Class, "${op}", ${op}_properties); +""" +) + +PY_FUNCTION_PROPS_AND_GETTERS = CodeTemplate( + """\ +${all_getter_definitions} + +static struct PyGetSetDef ${op}_properties[] = { + THP_FUNCTION_DEFAULT_PROPERTIES, + ${all_getsetdef_structs} + {nullptr} /* sentinel */ +}; + +""" +) + +PY_GETSETDEF_STRUCT = CodeTemplate( + """\ +{(char*)"_saved_${name}", (getter)THP${op}_${name}_getter, nullptr, nullptr, nullptr}""" +) + +PY_RAW_GETSETDEF_STRUCT = CodeTemplate( + """\ +{(char*)"_raw_saved_${name}", (getter)THP${op}_${name}_raw_getter, nullptr, nullptr, nullptr}""" +) + +# Getter templates +GETTER_DEFINITION = CodeTemplate( + """\ +static PyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) { + HANDLE_TH_ERRORS + auto prop = static_cast<${op}*>(self->cdata.get())->${name}; + ${body} + END_HANDLE_TH_ERRORS +} +""" +) + +GETTER_DEFINITION_SAVEDVAR = CodeTemplate( + """\ +static PyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) { + HANDLE_TH_ERRORS + const auto& prop = static_cast<${op}*>(self->cdata.get())->${name}_; + ${body} + END_HANDLE_TH_ERRORS +} +""" +) + +GETTER_DEFINITION_RAW_SAVEDVAR = CodeTemplate( + """\ +static PyObject* THP${op}_${name}_raw_getter(THPCppFunction *self, void *_unused) { + HANDLE_TH_ERRORS + const auto& prop = static_cast<${op}*>(self->cdata.get())->${name}_; + ${body} + END_HANDLE_TH_ERRORS +} +""" +) + +GETTER_DEFINITION_VEC_SAVEDVAR = CodeTemplate( + """\ +static PyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) { + HANDLE_TH_ERRORS + const auto *node = static_cast<${op}*>(self->cdata.get()); + const auto& prop = node->${name}_; + if (node->${name}_released_) { + PyErr_SetString(PyExc_RuntimeError, ERR_BACKWARD_TWICE); + return nullptr; + } + ${body} + END_HANDLE_TH_ERRORS +} +""" +) + +GETTER_DEFINITION_RAW_VEC_SAVEDVAR = CodeTemplate( + """\ +static PyObject* THP${op}_${name}_raw_getter(THPCppFunction *self, void *_unused) { + HANDLE_TH_ERRORS + const auto *node = static_cast<${op}*>(self->cdata.get()); + const auto& prop = node->${name}_; + if (node->${name}_released_) { + PyErr_SetString(PyExc_RuntimeError, ERR_BACKWARD_TWICE); + return nullptr; + } + ${body} + END_HANDLE_TH_ERRORS +} +""" +) + +GETTER_DEFINITION_OPT = CodeTemplate( + """\ +static PyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) { + HANDLE_TH_ERRORS + auto opt_prop = static_cast<${op}*>(self->cdata.get())->${name}; + if (!opt_prop.has_value()) { + Py_RETURN_NONE; + } + auto prop = opt_prop.value(); + ${body} + END_HANDLE_TH_ERRORS +} +""" +) + +GETTER_DEFINITION_OPT_ARRAYREF = CodeTemplate( + """\ +static PyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) { + HANDLE_TH_ERRORS + auto opt_prop = static_cast<${op}*>(self->cdata.get())->${name}; + if (!opt_prop.list.has_value()) { + Py_RETURN_NONE; + } + auto prop = opt_prop.list.value(); + ${body} + END_HANDLE_TH_ERRORS +} +""" +) + +# Getter body +GETTER_BODY_SAVEDVAR = """\ +return THPVariable_Wrap(prop.unpack(self->cdata)); +""" + +GETTER_BODY_RAW_SAVEDVAR = """\ +pybind11::object obj = pybind11::cast(prop, pybind11::return_value_policy::reference); +return obj.release().ptr(); +""" + +GETTER_BODY_VEC_SAVEDVAR = """\ +PyObject* tup = PyTuple_New((Py_ssize_t) prop.size()); +for (auto i: c10::irange(prop.size())) { + PyTuple_SetItem(tup, (Py_ssize_t) i, THPVariable_Wrap(prop[i].unpack(self->cdata))); +} +return tup; +""" + +GETTER_BODY_RAW_VEC_SAVEDVAR = """\ +PyObject* tup = PyTuple_New((Py_ssize_t) prop.size()); +for (auto i : c10::irange(prop.size())) { + pybind11::object obj = pybind11::cast(prop[i], pybind11::return_value_policy::reference); + PyTuple_SetItem(tup, (Py_ssize_t) i, obj.release().ptr()); +} +return tup; +""" + +GETTER_BODY_ARRAYREF_LONG = """\ +PyObject* tup = PyTuple_New((Py_ssize_t) prop.size()); +for (auto i : c10::irange(prop.size())) { + PyTuple_SetItem(tup, (Py_ssize_t) i, PyLong_FromUnsignedLong((uint64_t) prop[i])); +} +return tup; +""" + +GETTER_BODY_ARRAYREF_SYMINT = """\ +PyObject* tup = PyTuple_New((Py_ssize_t) prop.size()); +for (auto i : c10::irange(prop.size())) { + auto si = prop[i]; + if (auto m = si.maybe_as_int()) { + PyTuple_SetItem(tup, (Py_ssize_t) i, PyLong_FromUnsignedLong(*m)); + } else { + auto py_symint = py::cast(si).release().ptr(); + PyTuple_SetItem(tup, (Py_ssize_t) i, py_symint); + } +} +return tup; +""" + +GETTER_BODY_ARRAYREF_DOUBLE = """\ +PyObject* tup = PyTuple_New((Py_ssize_t) prop.size()); +for (auto i : c10::irange(prop.size())) { + PyTuple_SetItem(tup, (Py_ssize_t) i, PyFloat_FromDouble((double) prop[i])); +} +return tup; +""" + +GETTER_BODY_INT64_T = """\ +return PyLong_FromUnsignedLong((int64_t) prop); +""" + +GETTER_BODY_SYMINT = """\ +if (auto m = prop.maybe_as_int()) { + return PyLong_FromUnsignedLong(*m); +} else { + return py::cast(prop).release().ptr(); +} +""" + +GETTER_BODY_DOUBLE = """\ +return PyFloat_FromDouble((double) prop); +""" + +GETTER_BODY_BOOL = """\ +if (prop) { + Py_RETURN_TRUE; +} else { + Py_RETURN_FALSE; +} +""" + +GETTER_BODY_STRING = """\ +return PyUnicode_FromStringAndSize(prop.data(), prop.size()); +""" + +GETTER_BODY_SCALAR = """\ +if (prop.isComplex()) { + auto cprop = prop.to>(); + return PyComplex_FromDoubles(cprop.real(), cprop.imag()); +} else if (prop.isFloatingPoint()) { + return PyFloat_FromDouble(prop.to()); +} else if (prop.isIntegral(/*includeBool=*/false)) { + return PyLong_FromLong(prop.to()); +} else if (prop.isBoolean()) { + if (prop.to()) { + Py_RETURN_TRUE; + } else { + Py_RETURN_FALSE; + } +} else { + PyErr_SetString(PyExc_RuntimeError, "Unknown scalar type"); + return nullptr; +} +""" + + +GETTER_BODY_VEC_SCALAR = """\ +PyObject* tup = PyTuple_New((Py_ssize_t) prop.size()); +for (auto i: c10::irange(prop.size())) { + if (prop[i].isComplex()) { + auto cprop = prop[i].to>(); + PyTuple_SetItem(tup, (Py_ssize_t) i, PyComplex_FromDoubles(cprop.real(), cprop.imag())); + } else if (prop[i].isFloatingPoint()) { + auto double_prop = prop[i].to(); + PyTuple_SetItem(tup, (Py_ssize_t) i, PyFloat_FromDouble(double_prop)); + } else if (prop[i].isIntegral(/*includeBool=*/false)) { + auto long_prop = prop[i].to(); + PyTuple_SetItem(tup, (Py_ssize_t) i, PyLong_FromLong(long_prop)); + } else if (prop[i].isBoolean()) { + if (prop[i].to()) { + PyTuple_SetItem(tup, (Py_ssize_t) i, Py_True); + } else { + PyTuple_SetItem(tup, (Py_ssize_t) i, Py_False); + } + } else { + PyErr_SetString(PyExc_RuntimeError, "Unknown scalar type"); + return nullptr; + } +} +return tup; +""" + + +MISC_GETTER_DEFS = { + OptionalCType(BaseCType(longT)): (GETTER_DEFINITION_OPT, GETTER_BODY_INT64_T), + OptionalCType(BaseCType(SymIntT)): (GETTER_DEFINITION_OPT, GETTER_BODY_SYMINT), + BaseCType(doubleT): (GETTER_DEFINITION, GETTER_BODY_DOUBLE), + OptionalCType(BaseCType(doubleT)): (GETTER_DEFINITION_OPT, GETTER_BODY_DOUBLE), + BaseCType(boolT): (GETTER_DEFINITION, GETTER_BODY_BOOL), + BaseCType(scalarT): (GETTER_DEFINITION, GETTER_BODY_SCALAR), + OptionalCType(BaseCType(scalarT)): (GETTER_DEFINITION_OPT, GETTER_BODY_SCALAR), +} + +# These functions have backwards which cannot be traced, and so must have +# their backward functions traced opaquely. +# VIEW_FUNCTIONS are not traceable because they use as_strided, which +# has an untraceable backwards, see +# https://github.com/pytorch/pytorch/issues/4250 +# TODO: This is probably not exhaustive, but it's a start +UNTRACEABLE_FUNCTIONS = VIEW_FUNCTIONS + + +def get_infos_with_derivatives_list( + differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]], +) -> list[DifferentiabilityInfo]: + diff_info_list = [ + info + for diffinfo_dict in differentiability_infos.values() + for info in diffinfo_dict.values() + ] + + return list(filter(lambda info: info.args_with_derivatives, diff_info_list)) + + +def gen_autograd_functions_lib( + out: str, + differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]], + template_path: str, +) -> None: + """Functions.h and Functions.cpp body + + These contain the auto-generated subclasses of torch::autograd::Node + for each every differentiable torch function. + """ + + # get a 1D list of diffinfos, we do not need them to be per FunctionSchema/DispatchKey here + # infos with the diff dispatchkeys but the same name will still be in the same shard. + infos = get_infos_with_derivatives_list(differentiability_infos) + declarations = [process_function(f, FUNCTION_DECLARATION) for f in infos] + definitions = [process_function(f, FUNCTION_DEFINITION) for f in infos] + + file_basename = "Functions" + fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + for suffix in [".h", ".cpp"]: + fname = file_basename + suffix + fm.write_with_template( + fname, + fname, + lambda: { + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/{fname}", + "autograd_function_declarations": declarations, + "autograd_function_definitions": definitions, + }, + ) + + +def gen_autograd_functions_python( + out: str, + differentiability_infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]], + template_path: str, +) -> None: + fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + num_shards = 5 + fm.write( + "python_functions.h", + lambda: { + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/python_functions.h", + "shard_forward_declare": [ + f"void initialize_autogenerated_functions_{i}(PyObject* module);" + for i in range(num_shards) + ], + "shard_call": [ + f"initialize_autogenerated_functions_{i}(module);" + for i in range(num_shards) + ], + }, + ) + + # get a 1D list of diffinfos, we do not need them to be per FunctionSchema/DispatchKey here + # infos with the diff dispatchkeys but the same name will still be in the same shard. + infos = get_infos_with_derivatives_list(differentiability_infos) + fm.write_sharded( + "python_functions.cpp", + infos, + key_fn=lambda info: info.name, + base_env={ + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/python_functions.cpp", + }, + env_callable=lambda info: { + "py_function_initializers": [ + process_function(info, PY_FUNCTION_DEFINITION) + ], + "py_function_props_and_getters": [ + process_function(info, PY_FUNCTION_PROPS_AND_GETTERS) + ], + }, + num_shards=num_shards, + sharded_keys={"py_function_initializers", "py_function_props_and_getters"}, + ) + + +def process_function(info: DifferentiabilityInfo, template: CodeTemplate) -> str: + saved_variables: list[str] = [] + release_variables: list[str] = [] + saved_list_sizes: list[str] = [] + unpack: list[str] = [] + asserts: list[str] = [] + compute_index_ranges: list[str] = [] + getter_definitions: list[str] = [] + py_getsetdef_structs: list[str] = [] + compiled_args: list[str] = [] + apply_with_saved_before: list[str] = [] + apply_with_saved_after: list[str] = [] + apply_functional_args: list[str] = [] + apply_functional_args_ref_types: list[str] = [] + # Maps the name of an input (to the original forward operator; + # examples are "self", "other") to the order in which they appear in the + # operator. + # For example; if the operator is foo(Tensor self, int64_t k, Tensor other), + # the mapping is: {"self": 0, "other": 1}. + # We use this mapping to populate needs_input_grad in some order and then grab + # values from it. + input_name_to_idx: dict[str, int] = {} + + for idx, arg in enumerate(info.args_with_derivatives): + if arg.type in TENSOR_LIST_LIKE_CTYPES: + size = f"{arg.name}_size_" + saved_list_sizes.append(f"size_t {arg.name}_size_;") + apply_functional_args.append(f"{arg.name}_size_") + apply_functional_args_ref_types.append("size_t") + else: + size = "1" + compute_index_ranges.append(f"auto {arg.name}_ix = gen.range({size});") + input_name_to_idx[arg.name] = idx + + def save_var(var: SavedAttribute, is_output: bool) -> None: + name = var.nctype.name + type = var.nctype.type + should_append_getsetdef = True + should_append_raw_getsetdef = False + visit_name = name + uses_cpp_saved_variable_cls = False + unpacked_ref_type = None + + if ( + type == BaseCType(tensorT) + or type == OptionalCType(BaseCType(tensorT)) + or type == MutRefCType(OptionalCType(BaseCType(tensorT))) + or (type == BaseCType(scalarT) and is_output) + ): + uses_cpp_saved_variable_cls = True + saved_variables.append(f"SavedVariable {name}_;") + release_variables.append(f"{name}_.reset_data();") + ptr = "shared_from_this()" if is_output else "" + unpack.append(f"auto {name} = {name}_.unpack({ptr});") + getter_definitions.append( + GETTER_DEFINITION_SAVEDVAR.substitute( + op=info.op, name=name, body=GETTER_BODY_SAVEDVAR + ) + ) + getter_definitions.append( + GETTER_DEFINITION_RAW_SAVEDVAR.substitute( + op=info.op, name=name, body=GETTER_BODY_RAW_SAVEDVAR + ) + ) + should_append_raw_getsetdef = True + visit_name = f"{name}_" + unpacked_ref_type = "Tensor&" + elif ( + type == BaseCType(tensorListT) + or type == BaseCType(iTensorListRefT) + or type == VectorCType(BaseCType(tensorT)) + ): + # note(crcrpar): [nuanced return type of out-of-place foreach functions] + # When an out-of-place foreach function whose return signature is `Tensor[]` + # spells out its backward definitions in `derivatives.yaml`, and some of them depend on + # `result`, `result`'s type is interpreted and treated as `std::vector`. + # An out-of-place foreach whose backwards rely on their output doesn't suffer from this + # difference if the definitions are codegen'ed. + # This special case is needed for `_foreach_pow.List` and `_foreach_pow.ScalarAndTensor` + # as of https://github.com/pytorch/pytorch/pull/105504. + if type == VectorCType(BaseCType(tensorT)): + if not ( + info.func.func.name.name.base.startswith("_foreach") and is_output + ): + raise AssertionError( + "VectorCType(BaseCType(tensorT)) requires foreach function and is_output" + ) + uses_cpp_saved_variable_cls = True + saved_variables.append(f"std::vector {name}_;") + saved_variables.append(f"bool {name}_released_ = false;") + # Just clear() is sufficient, we don't need to loop and clear each variable. + # Because the SavedVariable owns a tensor and a grad_fn, removing the SavedVariable makes them go away as well. + release_variables.append(f"{name}_.clear();") + release_variables.append(f"{name}_released_ = true;") + ptr = "shared_from_this()" if is_output else "nullptr" + unpack.append(f"auto {name} = unpack_list({name}_, {ptr});") + asserts.append(f"TORCH_CHECK(!{name}_released_, ERR_BACKWARD_TWICE);") + getter_definitions.append( + GETTER_DEFINITION_VEC_SAVEDVAR.substitute( + op=info.op, name=name, body=GETTER_BODY_VEC_SAVEDVAR + ) + ) + getter_definitions.append( + GETTER_DEFINITION_RAW_VEC_SAVEDVAR.substitute( + op=info.op, name=name, body=GETTER_BODY_RAW_VEC_SAVEDVAR + ) + ) + should_append_raw_getsetdef = True + visit_name = f"{name}_" + unpacked_ref_type = "std::vector&" + elif type == ListCType(OptionalCType(BaseCType(tensorT))): + uses_cpp_saved_variable_cls = True + saved_variables.append(f"std::vector {name}_;") + saved_variables.append(f"bool {name}_released_ = false;") + # Just clear() is sufficient, we don't need to loop and clear each variable. + # Because the SavedVariable owns a tensor and a grad_fn, removing the SavedVariable makes them go away as well. + release_variables.append(f"{name}_.clear();") + release_variables.append(f"{name}_released_ = true;") + unpack.append(f"auto {name} = unpack_opt_list({name}_);") + asserts.append(f"TORCH_CHECK(!{name}_released_, ERR_BACKWARD_TWICE);") + getter_definitions.append( + GETTER_DEFINITION_VEC_SAVEDVAR.substitute( + op=info.op, name=name, body=GETTER_BODY_VEC_SAVEDVAR + ) + ) + getter_definitions.append( + GETTER_DEFINITION_RAW_VEC_SAVEDVAR.substitute( + op=info.op, name=name, body=GETTER_BODY_RAW_VEC_SAVEDVAR + ) + ) + should_append_raw_getsetdef = True + visit_name = f"{name}_" + unpacked_ref_type = "torch::List>&" + elif type == BaseCType(intArrayRefT): + saved_variables.append(f"std::vector {name};") + getter_definitions.append( + GETTER_DEFINITION.substitute( + op=info.op, name=name, body=GETTER_BODY_ARRAYREF_LONG + ) + ) + elif type == BaseCType(symIntArrayRefT): + saved_variables.append(f"std::vector {name};") + getter_definitions.append( + GETTER_DEFINITION.substitute( + op=info.op, name=name, body=GETTER_BODY_ARRAYREF_SYMINT + ) + ) + elif type == BaseCType(optionalIntArrayRefT): + saved_variables.append(f"c10::OptionalArray {name};") + getter_definitions.append( + GETTER_DEFINITION_OPT_ARRAYREF.substitute( + op=info.op, name=name, body=GETTER_BODY_ARRAYREF_LONG + ) + ) + elif type == BaseCType(optionalSymIntArrayRefT): + saved_variables.append(f"c10::OptionalArray {name};") + getter_definitions.append( + GETTER_DEFINITION_OPT_ARRAYREF.substitute( + op=info.op, name=name, body=GETTER_BODY_ARRAYREF_SYMINT + ) + ) + elif type == OptionalCType(BaseCType(intArrayRefT)): + saved_variables.append(f"c10::OptionalArray {name};") + getter_definitions.append( + GETTER_DEFINITION_OPT_ARRAYREF.substitute( + op=info.op, name=name, body=GETTER_BODY_ARRAYREF_LONG + ) + ) + elif type == OptionalCType(BaseCType(symIntArrayRefT)): + saved_variables.append(f"c10::OptionalArray {name};") + getter_definitions.append( + GETTER_DEFINITION_OPT_ARRAYREF.substitute( + op=info.op, name=name, body=GETTER_BODY_ARRAYREF_SYMINT + ) + ) + elif type == OptionalCType(ArrayRefCType(BaseCType(doubleT))): + saved_variables.append(f"c10::OptionalArray {name};") + getter_definitions.append( + GETTER_DEFINITION_OPT_ARRAYREF.substitute( + op=info.op, name=name, body=GETTER_BODY_ARRAYREF_DOUBLE + ) + ) + elif type == BaseCType(longT): + saved_variables.append(f"{type.cpp_type()} {name} = 0;") + getter_definitions.append( + GETTER_DEFINITION.substitute( + op=info.op, name=name, body=GETTER_BODY_INT64_T + ) + ) + elif type == BaseCType(SymIntT): + saved_variables.append(f"c10::SymInt {name};") + getter_definitions.append( + GETTER_DEFINITION.substitute( + op=info.op, name=name, body=GETTER_BODY_SYMINT + ) + ) + elif type == BaseCType(stringT): + saved_variables.append(f"std::string {name};") + getter_definitions.append( + GETTER_DEFINITION.substitute( + op=info.op, name=name, body=GETTER_BODY_STRING + ) + ) + elif type == OptionalCType(BaseCType(stringT)): + saved_variables.append(f"std::optional {name};") + getter_definitions.append( + GETTER_DEFINITION_OPT.substitute( + op=info.op, name=name, body=GETTER_BODY_STRING + ) + ) + elif type == ArrayRefCType( + elem=BaseCType(type=BaseCppType(ns="at", name="Scalar")) + ): + saved_variables.append(f"std::vector {name};") + unpacked_ref_type = "std::vector&" + saved_variables.append(f"bool {name}_released_ = false;") + # Just clear() is sufficient, we don't need to loop and clear each variable. + # Because the SavedVariable owns a tensor and a grad_fn, removing the SavedVariable makes them go away as well. + release_variables.append(f"{name}.clear();") + # release_variables.append(f"{name}_released_ = true;") + # unpack.append(f"auto {name} = unpack_list({name}_);") + # asserts.append(f"TORCH_CHECK(!{name}_released_, ERR_BACKWARD_TWICE);") + getter_definitions.append( + CodeTemplate( + """\ +static PyObject* THP${op}_${name}_getter(THPCppFunction *self, void *_unused) { + HANDLE_TH_ERRORS + const auto *node = static_cast<${op}*>(self->cdata.get()); + const auto& prop = node->${name}; + if (node->${name}_released_) { + PyErr_SetString(PyExc_RuntimeError, ERR_BACKWARD_TWICE); + return nullptr; + } + ${body} + END_HANDLE_TH_ERRORS +} + """ + ).substitute( + op=info.op, + name=name, + body=GETTER_BODY_VEC_SCALAR, + ) + ) + else: + # Check for indicators that you're putting a non-owning reference + # into the saved variable field. If this is spuriously firing, + # edit this field. Otherwise, you probably need to add a case + # above. + if not ( + "ref" not in type.cpp_type().lower() + and "view" not in type.cpp_type().lower() + and "*" not in type.cpp_type() + and "&" not in type.cpp_type() + ): + raise AssertionError( + f"{type.cpp_type()} looks like it contains a non-owning reference" + ) + saved_variables.append(f"{type.cpp_type()} {name};") + + if type in MISC_GETTER_DEFS: + # pyrefly: ignore [bad-index, index-error] + getter_def, body = MISC_GETTER_DEFS[type] + getter_definitions.append( + getter_def.substitute(op=info.op, name=name, body=body) + ) + else: + # Types we don't expose python bindings to yet: + # TypeAndSize, at::ScalarType, TensorOptions, TensorGeometry, + # std::vector>, std::vector + should_append_getsetdef = False + + if should_append_getsetdef: + py_getsetdef_structs.append( + PY_GETSETDEF_STRUCT.substitute(op=info.op, name=name) + ) + if should_append_raw_getsetdef: + py_getsetdef_structs.append( + PY_RAW_GETSETDEF_STRUCT.substitute(op=info.op, name=name) + ) + + if uses_cpp_saved_variable_cls: + compiled_args.append( + f"args.collect({visit_name}, {'true' if is_output else 'false'});" + ) + else: + compiled_args.append(f"args.collect({visit_name});") + apply_with_saved_before.append(f"saved.before({visit_name});") + apply_with_saved_after.append(f"saved.after({visit_name});") + + if unpacked_ref_type is None: + unpacked_ref_type = f"{saved_variables[-1].split(' ')[0]}&" + apply_functional_args.append(str(name)) + apply_functional_args_ref_types.append(unpacked_ref_type) + + for var in sorted(info.all_saved_inputs, key=lambda sa: str(sa.nctype.name)): + save_var(var, is_output=False) + for var in sorted(info.all_saved_outputs, key=lambda sa: str(sa.nctype.name)): + save_var(var, is_output=True) + + # lock the mutex when we release variables and in Node::apply to protect thread safety + # see Note [Thread Safety on Autograd Node] + if len(release_variables) > 0: + thread_lock = "std::lock_guard lock(mutex_);" + else: + thread_lock = "" + + if uses_retain_variables(info): + apply_functional_args.append("retain_variables") + apply_functional_args_ref_types.append("bool") + will_release_variables = WILL_RELEASE_VARIABLES.substitute() + else: + will_release_variables = "" + + body: list[str] = [] + + if uses_single_grad(info): + body.append("const auto& grad = grads[0];") + else: + # Generate aliases for gradients named for returned values. + body.extend( + f"const auto& {name} = grads[{info.available_named_gradients.index(name)}];" + for name in sorted(info.used_named_gradients) + ) + + def emit_derivative( + derivative: Derivative, + args_with_derivatives: Sequence[Binding], + ) -> tuple[bool, str]: + formula = derivative.formula + var_names = derivative.var_names + + if len(var_names) == 1: + checks_any_grad_defined = False + if "not_implemented" not in formula: + matching_args = [ + arg for arg in args_with_derivatives if arg.name == var_names[0] + ] + if len(matching_args) == 1: + # We can add undefined grad support if the input variable is a Tensor + arg = matching_args[0] + if isinstance(arg.argument, Argument) and str( + arg.argument.type + ) in ("Tensor", "Tensor?"): + formula = "any_grad_defined ? (" + formula + ") : Tensor()" + checks_any_grad_defined = True + if info.name.startswith("_foreach_"): + derivative_template = DERIVATIVE_SINGLE_FOREACH + else: + derivative_template = DERIVATIVE_SINGLE + return ( + checks_any_grad_defined, + derivative_template.substitute( + name=var_names[0], + derivative=formula, + idx=input_name_to_idx[var_names[0]], + ), + ) + + else: + if "grad_input_mask" in formula: + masks = [ + f"needs_input_grad[{input_name_to_idx[name]}]," + for name in var_names + ] + grad_input_mask = GRAD_INPUT_MASK.substitute( + n=len(var_names), masks=masks + ) + else: + grad_input_mask = "" + needs_input_grad = [ + f"needs_input_grad[{input_name_to_idx[name]}]" for name in var_names + ] + needs_input_grad = " || ".join(needs_input_grad) + copy_ranges: list[str] = [] + for i, n in enumerate(var_names): + copy_ranges.append( + DERIVATIVE_MULTI_COPY_RANGE.substitute( + name=n, i=i, idx=input_name_to_idx[n] + ) + ) + return False, DERIVATIVE_MULTI.substitute( + needs_input_grad=needs_input_grad, + copy_ranges=copy_ranges, + derivative=formula, + grad_input_mask=grad_input_mask, + ) + + masks = [] + + need_any_grad_defined_var = False + for derivative in info.derivatives: + checks_any_grad_defined, derivative_text = emit_derivative( + derivative, info.args_with_derivatives + ) + body.append(derivative_text) + need_any_grad_defined_var |= checks_any_grad_defined + + for name in input_name_to_idx: + masks.append(f"task_should_compute_output({{ {name}_ix }}),") + + # Since single-output derivative formulas need to check if grads are + # defined, only perform the check once, before all the formulas + if need_any_grad_defined_var: + body.insert( + -len(info.derivatives), + "bool any_grad_defined = any_variable_defined(grads);", + ) + + if info.name in UNTRACEABLE_FUNCTIONS: + superclass = "Node" + else: + superclass = "TraceableFunction" + + all_getsetdef_structs = ( + ",\n".join(py_getsetdef_structs) + "," if len(py_getsetdef_structs) != 0 else "" + ) + all_getter_definitions = "\n".join(getter_definitions) + + compute_needs_input_grad = COMPUTE_NEEDS_INPUT_GRAD.substitute( + n=len(masks), compute_index_ranges=compute_index_ranges, masks=masks + ) + apply_functional_args_signature = [ + f"{T} {x}" + for T, x in zip(apply_functional_args_ref_types, apply_functional_args) + ] + get_packed_args = "\n".join( + f"packed_args.pack({name});" for name in apply_functional_args + ) + unpack_ivalues = [] + for typ, name in zip(apply_functional_args_ref_types, apply_functional_args): + typ = typ.removesuffix("&") + # pyrefly: ignore [bad-argument-type] + unpack_ivalues.append(f"auto {name} = packed_args.unpack<{typ}>();") + + schema_args = [f"std::array"] + for typ in apply_functional_args_ref_types: + typ = typ.removesuffix("&") + typ = typ.removeprefix("const") + schema_args.append(typ.strip()) + compute_schema = ["std::vector schema = {"] + for schema_arg in schema_args: + compute_schema.append( + f" torch::dynamo::autograd::IValuePacker<{schema_arg}>::packed_type()," + ) + compute_schema.append("};") + + return template.substitute( + unpacks="\n".join(unpack), + op=info.op, + compute_schema="\n".join(compute_schema), + apply_functional_args=apply_functional_args, + apply_functional_args_signature=apply_functional_args_signature, + compute_needs_input_grad=compute_needs_input_grad, + num_inputs=len(input_name_to_idx), + unpack_ivalues="\n".join(unpack_ivalues), + compute_index_ranges=compute_index_ranges, + saved_variables=saved_variables, + release_variables=release_variables, + saved_list_sizes=saved_list_sizes, + asserts=asserts, + thread_lock=thread_lock, + will_release_variables=will_release_variables, + body=body, + superclass=superclass, + all_getter_definitions=all_getter_definitions, + all_getsetdef_structs=all_getsetdef_structs, + compiled_args=compiled_args, + apply_with_saved_before=apply_with_saved_before, + apply_with_saved_after=apply_with_saved_after, + get_packed_args=get_packed_args, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_inplace_or_view_type.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_inplace_or_view_type.py new file mode 100644 index 0000000000000000000000000000000000000000..bbdb833143c47225e9ebbea6de94778c6b298db8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_inplace_or_view_type.py @@ -0,0 +1,676 @@ +# Generates ADInplaceOrViewType.h/cpp +# +# NOTE: If any changes are being made to the ADInplaceOrView codegen please also check +# if updates are needed in torch/csrc/autograd/autograd_not_implemented_fallback.cpp +# The fallback is expected to mimic this codegen, so we should keep the two in sync. + +from __future__ import annotations + +from torchgen.api import cpp +from torchgen.api.autograd import ( + dispatch_strategy, + gen_differentiable_outputs, + NativeFunctionWithDifferentiabilityInfo, +) +from torchgen.api.types import ( + BaseCType, + Binding, + boolT, + ConstRefCType, + CType, + DispatcherSignature, + intArrayRefT, + longT, + OptionalCType, + symIntArrayRefT, + SymIntT, + tensorT, +) +from torchgen.code_template import CodeTemplate +from torchgen.context import with_native_function +from torchgen.model import ( + NativeFunction, + SchemaKind, + SelfArgument, + TensorOptionsArguments, + Type, +) +from torchgen.utils import FileManager + +from .context import with_native_function_with_differentiability_info +from .gen_trace_type import ( + get_return_value, + MANUAL_AUTOGRAD, + tie_return_values, + type_wrapper_name, +) + + +# See NOTE [ Autograd View Variables ] in variable.h for details. +# If you update list VIEW_FUNCTIONS or RETURNS_VIEWS_OF_INPUT, +# you **MUST** also update the public list of view ops accordingly in +# docs/source/tensor_view.rst. Note not all ATen functions are exposed to public, +# e.g alias & sparse_coo_tensor_with_dims_and_tensors. +# +# A map: function name => name of the argument that all outputs are view of + +VIEW_FUNCTIONS_WITH_METADATA_CHANGE = [ + "view_as_complex", + "view_as_real", + "_conj", + "_neg_view", + "_nested_get_values", + "_nested_view_from_buffer", + "_nested_view_from_jagged", +] + +VIEW_FUNCTIONS = { + "numpy_T": "self", + "alias": "self", + "as_strided": "self", + "diagonal": "self", + "expand": "self", + "permute": "self", + "select": "self", + "slice": "self", + "slice_inverse": "self", + "split": "self", + "split_with_sizes": "self", + "squeeze": "self", + "t": "self", + "transpose": "self", + "unfold": "self", + "unsqueeze": "self", + "flatten": "self", + "view": "self", + "unbind": "self", + "_indices": "self", + "_values": "self", + "indices": "self", + "values": "self", + "crow_indices": "self", + "col_indices": "self", + "ccol_indices": "self", + "row_indices": "self", + # sparse_coo ctor output should really be views of both indices and values, + # but we only supports making as view of a single variable, and indices is + # discrete anyways. + # FIXME: clone indices on construction. + "sparse_coo_tensor_with_dims_and_tensors": "values", + "_reshape_alias": "self", + "_test_autograd_multiple_dispatch_view": "self", +} + +for key in VIEW_FUNCTIONS_WITH_METADATA_CHANGE: + VIEW_FUNCTIONS[key] = "self" + +# note: some VIEW_FUNCTIONS are just compositions of the view functions above +# this list contains both the root view functions and any that are purely composed +# of viewing functions, and is used by the JIT to determine when an operator +# may return a view of its inputs; however they may sometimes return a copy. +# (e.g. `contiguous`) +RETURNS_VIEWS_OF_INPUT = set(VIEW_FUNCTIONS.keys()).union( + { + "chunk", + "detach", + "contiguous", + "reshape", + "reshape_as", + "expand_as", + "view_as", + "real", + "imag", + "narrow", + "movedim", + "tensor_split", + "swapdims", + "swapaxes", + "mT", + "mH", + "adjoint", + "matrix_H", + } +) + +# These are the functions we consider views for the purposes of validating +# StorageImpl and TensorImpl in gen_variable_type. +# `_unsafe_view` is not included in VIEW_FUNCTIONS above because it is not a +# view for the purposes of ADInplaceOrView kernel, we do not want to call as_view +# See NOTE [Unsafe View] for more info. +ALL_VIEW_FUNCTIONS = { + **VIEW_FUNCTIONS, + "_unsafe_view": "self", +} + +ARRAYREF_TO_VEC = CodeTemplate( + """\ +auto ${vec} = ${arg}.vec(); +""" +) + +OPTIONAL_TO_VAL = CodeTemplate( + """\ +auto ${val} = ${arg}.value_or(${default}); +""" +) + +CALL_DISPATCH = CodeTemplate( + """\ +at::_ops::${unambiguous_name}::call(${unpacked_args})""" +) + +REVERSE_VIEW_DISPATCH = CodeTemplate( + """\ +${reverse_name}(${unpacked_args})""" +) + +MULTI_OUTPUT_VIEW_ITERATION = CodeTemplate( + """\ +for (auto ${view_idx} : c10::irange(${var}.size())) { + ${body} +} +""" +) + +SETUP_REPLAY_VIEW_IF_NOT_SUPPORT_AS_STRIDED_OR_VIEW_WITH_METADATA_CHANGE = CodeTemplate( + """\ +std::unique_ptr func(nullptr); +std::function rev_func=nullptr; +if (${is_view_with_metadata_change} || + !self.unsafeGetTensorImpl()->support_as_strided() || + self.unsafeGetTensorImpl()->is_python_dispatch() || + c10::AutogradState::get_tls_state().get_view_replay_enabled()) { + ${replay_view_func} + ${reverse_replay_view_func} +} +""" +) + +REPLAY_VIEW_FUNC = CodeTemplate( + """\ +func = std::make_unique<${view_func_name}>(${view_func_args}); +""" +) + +REVERSE_REPLAY_VIEW_LAMBDA_FUNC = CodeTemplate( + """\ +rev_func = [=](const at::Tensor& ${input_view}) { + return ${reverse_replay_view_call}; +}; +""" +) + +METHOD_DEFINITION = CodeTemplate( + """\ +${return_type} ${type_wrapper_name}(${formals}) { + ${type_definition_body} +} +""" +) + +WRAPPER_REGISTRATION = CodeTemplate( + """\ +m.impl("${unqual_operator_name_with_overload}", + TORCH_FN(${class_type}::${type_wrapper_name}) +); +""" +) + +AUTOGRAD_NOT_IMPLEMENTED_REGISTRATION = CodeTemplate( + """\ +m.impl("${unqual_operator_name_with_overload}", torch::autograd::autogradNotImplementedFallback()); +""" +) + +INPLACE_REDISPATCH = CodeTemplate( + """\ +{ + at::AutoDispatchBelowADInplaceOrView guard; + at::_ops::${unambiguous_name}::redispatch(${unpacked_args}); +} +""" +) + +ASSIGN_RETURN_VALUE = CodeTemplate( + """\ +${return_values} = ${rhs_value}; +""" +) + +VIEW_REDISPATCH = CodeTemplate( + """\ +${assign_return_values} ([&]() { + at::AutoDispatchBelowADInplaceOrView guard; + return at::_ops::${unambiguous_name}::redispatch(${unpacked_args}); +})(); +""" +) + +TMP_VAR = "_tmp" + + +# FIXME: Ideally these functions should be methods on Type class, but we have a +# comment in codegen/model.py there saying these concepts are not well defined. +# Thus we put a version that commonly used by autograd codegen here. +def is_tensor_type(t: Type) -> bool: + # TODO: Should handle optional here? + return t.is_tensor_like() and t.is_list_like() is None + + +def is_tensor_list_type(t: Type) -> bool: + # TODO: Should handle optional here? + return t.is_tensor_like() and t.is_list_like() is not None + + +UNPACK_TENSOR = CodeTemplate( + """\ +auto${ref} ${arg_name}_ = unpack${suffix}(${arg_name}, "${arg_name}", ${arg_pos});""" +) + + +def unpacked_name(arg_name: str) -> str: + return arg_name + "_" + + +# e.g. select.int -> select_copy_int_inverse() +def inverse_view_name(f: NativeFunction) -> str: + copy_variant = f"{f.root_name}_copy" + overload = f"{f.func.name.overload_name}" + if overload != "": + overload = "_" + overload + return f"{copy_variant}{overload}_inverse" + + +def extract_bindings(f: NativeFunction) -> list[Binding]: + return [ + r + for a in f.func.schema_order_arguments() + for r in cpp.argument( + a, + method=False, + symint=True, + cpp_no_default_args=set(), + faithful=False, + has_tensor_options=False, + ) + ] + + +@with_native_function +def unpack_args(f: NativeFunction) -> tuple[list[str], list[Binding]]: + body: list[str] = [] + unpacked_bindings: list[Binding] = [] + + for i, binding in enumerate(extract_bindings(f)): + if isinstance(binding.argument, SelfArgument): + raise AssertionError("Binding argument should not be SelfArgument") + if isinstance(binding.argument, TensorOptionsArguments): + raise RuntimeError("VariableKernel shouldn't take TensorOptions") + + is_nullable = binding.argument.type.is_nullable() + if not binding.argument.type.is_tensor_like() or is_nullable: + unpacked_bindings.append(binding) + continue + + is_tensor_list = is_tensor_list_type(binding.argument.type) + ref = (not is_nullable) and not is_tensor_list + suffix = "_opt" if is_nullable and not is_tensor_list else "" + body.append( + UNPACK_TENSOR.substitute( + arg_name=binding.name, + arg_pos=i, + suffix=suffix, + ref="&" if ref else "", + ) + ) + unpacked_bindings.append( + Binding( + name=unpacked_name(binding.name), + nctype=binding.nctype, + argument=binding.argument, + default=binding.default, + ) + ) + + return body, unpacked_bindings + + +def get_base_name(f: NativeFunction) -> str: + return f.func.name.name.base # TODO: should be str(f.func.name.name)? + + +def get_view_info(f: NativeFunction) -> str | None: + base_name = get_base_name(f) + view_info = VIEW_FUNCTIONS.get(base_name) + if view_info is None and base_name in RETURNS_VIEWS_OF_INPUT: + view_info = "self" + return view_info + + +def emit_view_func( + f: NativeFunction, bindings: list[Binding], view_idx: str | None = None +) -> str: + """Generate an additional lambda function to recover views in backward when as_strided is not supported. + See Note [View + Inplace update for base tensor] and [View + Inplace update for view tensor] for more details. + """ + # TODO: Clean this logic up if we get rid of reverse view funcs or reify them. + input_base = "input_base" + replay_view_func = "" + updated_args: list[str] = [] + known_view_arg_simple_types: list[CType] = [ + BaseCType(longT), + OptionalCType(BaseCType(longT)), + BaseCType(SymIntT), + OptionalCType(BaseCType(SymIntT)), + BaseCType(boolT), + BaseCType(intArrayRefT), + BaseCType(symIntArrayRefT), + ConstRefCType(BaseCType(tensorT)), + ConstRefCType(OptionalCType(BaseCType(tensorT))), + ] + for binding in bindings: + arg, arg_type = binding.name, binding.nctype.type + if arg == "self": + updated_args.append(input_base) + continue + if arg_type not in known_view_arg_simple_types: + known_types_str = ", ".join([str(t) for t in known_view_arg_simple_types]) + raise TypeError( + f"You are adding an {arg_type} {arg} argument to op {cpp.name(f.func)} in addition to known types: " + f"{known_types_str}. Please update the list or materialize it so that it can be closed " + "over by value, also add a test in pytorch/xla/test/test_operations.py where this code " + "is exercised." + ) + if arg_type == BaseCType(intArrayRefT) or arg_type == BaseCType( + symIntArrayRefT + ): + # It's not safe to close over IntArrayRef by value, since this is a + # reference type, so materialize a vector to close over by value + arg_vec = arg + "_vec" + replay_view_func += ARRAYREF_TO_VEC.substitute(arg=arg, vec=arg_vec) + updated_args.append(arg_vec) + elif arg_type == OptionalCType(BaseCType(longT)): + # Materialize int64_t? to int64_t + arg_value = arg + "_val" + replay_view_func += OPTIONAL_TO_VAL.substitute( + arg=arg, val=arg_value, default="0" + ) + updated_args.append(arg_value) + elif arg_type == ConstRefCType(BaseCType(tensorT)) or arg_type == ConstRefCType( + OptionalCType(BaseCType(tensorT)) + ): + # NB: Closing over a tensor. If a user modifies this tensor, this will be silently + # incorrect. The proper thing to do is to store the version counter and copy on write. + updated_args.append(arg) + else: + updated_args.append(arg) + + from .gen_view_funcs import view_func_name + + view_func_args = [b.name for b in bindings if b.name != "self"] + if view_idx is not None: + view_func_args.append(f"{view_idx}") + replay_view_func += REPLAY_VIEW_FUNC.substitute( + view_func_name=view_func_name(f, include_namespace=True), + view_func_args=view_func_args, + ) + + input_view = "input_view" + reverse_unpacked_args = [ + "self", + f"{input_view}", + # inverse_return_mode= + "at::functionalization::InverseReturnMode::AlwaysView", + *(() if view_idx is None else (f"{view_idx}",)), + # skip input_base arg + *updated_args[1:], + ] + + from torchgen.api.functionalization import reverse_name + + reverse_replay_view_call = REVERSE_VIEW_DISPATCH.substitute( + reverse_name=reverse_name(f, include_namespace=True), + unpacked_args=reverse_unpacked_args, + ) + reverse_replay_view_func = REVERSE_REPLAY_VIEW_LAMBDA_FUNC.substitute( + input_view=input_view, reverse_replay_view_call=reverse_replay_view_call + ) + + is_view_with_metadata_change = ( + "true" if cpp.name(f.func) in VIEW_FUNCTIONS_WITH_METADATA_CHANGE else "false" + ) + + return SETUP_REPLAY_VIEW_IF_NOT_SUPPORT_AS_STRIDED_OR_VIEW_WITH_METADATA_CHANGE.substitute( + is_view_with_metadata_change=is_view_with_metadata_change, + replay_view_func=replay_view_func, + reverse_replay_view_func=reverse_replay_view_func, + ) + + +def emit_view_body( + fn: NativeFunctionWithDifferentiabilityInfo, var: str +) -> tuple[str, str]: + # See NOTE [ Autograd View Variables ] in variable.h for details. + f = fn.func + base_name = get_base_name(f) + view_info = get_view_info(f) + call = "" + differentiable_outputs = gen_differentiable_outputs(fn) + differentiable_output_vars = {r.name for r in differentiable_outputs} + if not isinstance(view_info, str): + raise TypeError( + f"The view info should be a string for {base_name}, but it is: {view_info}" + ) + if len(differentiable_output_vars) == 0: + # no output is differentiable (.indices() for SparseTensors for example) + rhs_value = ( + f"as_view({view_info}, {var}, " + f"/* is_bw_differentiable */ false, /* is_fw_differentiable */ false)" + ) + elif len(differentiable_output_vars) == 1: + # Single differentiable output (Tensor or Tensor[]) + return_info = differentiable_outputs[0] + # We only support simple Tensor or a TensorList for functions that return views + if not is_tensor_type(return_info.type) and not is_tensor_list_type( + return_info.type + ): + raise RuntimeError( + f"{base_name} that return differentiable views can only return Tensor or Tensor[]" + ) + + # See Note [ View + Inplace detection] + def get_creation_meta_in_mode(original: str) -> str: + creation_meta_with_grad_mode = f"(at::GradMode::is_enabled() ? {original} : CreationMeta::NO_GRAD_MODE)" + return f"InferenceMode::is_enabled() ? CreationMeta::INFERENCE_MODE : {creation_meta_with_grad_mode}" + + # Only allow rebasing of the history if we return a single Tensor + # If we are in a no grad block, raise a warning + # See NOTE [ View + Inplace detection ] for more details about this logic + if is_tensor_list_type(return_info.type): + creation_meta = get_creation_meta_in_mode("CreationMeta::MULTI_OUTPUT_NODE") + view_idx = "view_idx" + view_func = emit_view_func( + f, extract_bindings(f), view_idx=view_idx + ).strip() + as_view_call = ( + f"as_view(/* base */ {view_info}, /* output */ {var}[{view_idx}], " + "/* is_bw_differentiable */ true, /* is_fw_differentiable */ true, " + "/* view_func */ std::move(func), /* rev_view_func */ rev_func, " + f"/* creation_meta */ {creation_meta});" + ) + call += MULTI_OUTPUT_VIEW_ITERATION.substitute( + var=var, view_idx=view_idx, body=f"{view_func}\n{as_view_call}" + ) + rhs_value = f"std::move({var})" + else: + call += emit_view_func(f, extract_bindings(f), view_idx=None) + creation_meta = get_creation_meta_in_mode("CreationMeta::DEFAULT") + rhs_value = ( + f"as_view(/* base */ {view_info}, /* output */ {var}, /* is_bw_differentiable */ true, " + "/* is_fw_differentiable */ true, " + f"/* view_func */ std::move(func), /* rev_view_func */ rev_func, /* creation_meta */ {creation_meta})" + ) + else: + # This could be supported but we don't need it at the moment, so keeping things simple. + raise RuntimeError( + "Function that return multiple differentiable output " + "when at least one of them is view is not supported." + ) + return call, rhs_value + + +def modifies_arguments(f: NativeFunction) -> bool: + return f.func.kind() in [SchemaKind.inplace, SchemaKind.out] + + +@with_native_function_with_differentiability_info +def emit_inplace_or_view_body(fn: NativeFunctionWithDifferentiabilityInfo) -> list[str]: + f = fn.func + inplace_view_body: list[str] = [] + + dispatcher_sig = DispatcherSignature.from_schema(f.func) + dispatcher_exprs = dispatcher_sig.exprs() + + # code-generated ADInplaceOrView kernels plumb and recompute dispatch keys directly through the kernel for performance. + # See Note [Plumbing Keys Through The Dispatcher] for details. + dispatch_key_set = "ks & c10::after_ADInplaceOrView_keyset" + redispatch_args = ", ".join([dispatch_key_set] + [a.expr for a in dispatcher_exprs]) + + # Note that this calls the slow, dispatching variants of manual_cpp_binding ops. + # We could probably work harder to ensure that the fast variants are called instead, but the perf benefit would be minimal. + if modifies_arguments(f): # inplace op + inplace_view_body.append( + INPLACE_REDISPATCH.substitute( + unambiguous_name=f.func.name.unambiguous_name(), + unpacked_args=redispatch_args, + ) + ) + for r in cpp.return_names(f): + inplace_view_body.append(f"increment_version({r});") + else: + if get_view_info(f) is None: + raise AssertionError("Expected view info to be non-None") + inplace_view_body.append( + VIEW_REDISPATCH.substitute( + assign_return_values="auto " + TMP_VAR + " = ", + unambiguous_name=f.func.name.unambiguous_name(), + unpacked_args=redispatch_args, + ) + ) + call, rhs_value = emit_view_body(fn, TMP_VAR) + inplace_view_body.append(call) + if rhs_value is None: + raise AssertionError("Expected rhs_value to be non-None") + inplace_view_body.append( + ASSIGN_RETURN_VALUE.substitute( + return_values=tie_return_values(f), rhs_value=rhs_value + ) + ) + if f.func.returns: + inplace_view_body.append(f"return {get_return_value(f)};") + return inplace_view_body + + +@with_native_function +def gen_formals(f: NativeFunction) -> str: + return ", ".join( + # code-generated autograd kernels plumb and recompute dispatch keys directly through the kernel for performance. + # See Note [Plumbing Keys Through The Dispatcher] for details. + ["c10::DispatchKeySet ks"] + + [ + f"{cpp.argument_type(a, binds='__placeholder__', symint=True).cpp_type()} {a.name}" + for a in f.func.schema_order_arguments() + ] + ) + + +@with_native_function_with_differentiability_info +def inplace_or_view_method_definition( + fn: NativeFunctionWithDifferentiabilityInfo, +) -> str | None: + f = fn.func + if get_view_info(f) is None and ( + # For functions that modify their inputs but don't return them, + # we can't give them autograd support. + # See https://github.com/pytorch/pytorch/issues/53796 + not modifies_arguments(f) or len(f.func.returns) == 0 + ): + return None + return METHOD_DEFINITION.substitute( + return_type=cpp.returns_type(f.func.returns, symint=True).cpp_type(), + type_wrapper_name=type_wrapper_name(f), + formals=gen_formals(f), + type_definition_body=emit_inplace_or_view_body(fn), + ) + + +@with_native_function_with_differentiability_info +def inplace_or_view_method_registration( + fn: NativeFunctionWithDifferentiabilityInfo, +) -> str | None: + f = fn.func + if get_view_info(f) is None and ( + not modifies_arguments(f) or len(f.func.returns) == 0 + ): + return None + return WRAPPER_REGISTRATION.substitute( + unqual_operator_name_with_overload=f.func.name, + type_wrapper_name=type_wrapper_name(f), + class_type="ADInplaceOrView", + ) + + +def use_derived(fn: NativeFunctionWithDifferentiabilityInfo) -> bool: + f = fn.func + name = cpp.name(f.func) + return name not in MANUAL_AUTOGRAD and dispatch_strategy(fn) == "use_derived" + + +def gen_inplace_or_view_type_env( + fn: NativeFunctionWithDifferentiabilityInfo, +) -> dict[str, list[str]]: + definition = inplace_or_view_method_definition(fn) + registration = inplace_or_view_method_registration(fn) + + return { + "ops_headers": ( + [f"#include "] + if definition is not None + else [] + ), + "inplace_or_view_method_definitions": [definition] + if definition is not None + else [], + "inplace_or_view_wrapper_registrations": [registration] + if registration is not None + else [], + } + + +def gen_inplace_or_view_type( + out: str, + native_yaml_path: str, + tags_yaml_path: str, + fns_with_infos: list[NativeFunctionWithDifferentiabilityInfo], + template_path: str, +) -> None: + # NOTE: see Note [Sharded File] at the top of the VariableType.cpp + # template regarding sharding of the generated files. + + fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + fm.write_sharded( + "ADInplaceOrViewType.cpp", + [fn for fn in fns_with_infos if use_derived(fn)], + key_fn=lambda fn: fn.func.root_name, + base_env={ + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/ADInplaceOrViewType.cpp", + }, + env_callable=gen_inplace_or_view_type_env, + num_shards=2, + sharded_keys={ + "ops_headers", + "inplace_or_view_method_definitions", + "inplace_or_view_wrapper_registrations", + }, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_python_functions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_python_functions.py new file mode 100644 index 0000000000000000000000000000000000000000..09688e561e52582feff50eeda1beb28fc82f3112 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_python_functions.py @@ -0,0 +1,1413 @@ +# Generates Python bindings for ATen functions +# +# The bindings are generated as methods on python_variable or functions on the +# torch._C._nn. torch._C._fft, torch._C._linalg, torch._C._nested, torch._C._sparse +# or torch._C._special objects. +# + +# Code tries to stick to the following rules: +# +# - templates should be colocated with the functions that use them. +# no templates are currently shared between functions, but if that +# happens, maybe put the template with the first one +# +# - don't use environment dictionaries when calling template.substitute(). +# pass named arguments directly for everything, otherwise it's much too +# hard to track what's actually being used and by who +# +# - colocate any new hacks/adjustments with existing ones of the same kind. +# ideally in a data structure rather than code if possible. See e.g. +# SCHEMA_DEFAULT_CONVERSION_HACKS, etc. +# +# - similarly, conversions from one format to another should ideally happen +# all at once in a single place. +# +# - no nontrivial nested functions. couple-liners are ok but please no more. +# especially avoid functions that read/write outer variables defined far away. +# +# - raise RuntimeError instead of asserting, and put as much +# information as is available into the message. I.e. no need to +# plumb in new params whose only purpose is to fill out an error +# message, but use what's there +# + +from __future__ import annotations + +import itertools +import re +from collections import defaultdict +from typing import TYPE_CHECKING + +import yaml + +from torchgen.api import cpp +from torchgen.api.python import ( + arg_parser_output_exprs, + cpp_dispatch_exprs, + cpp_dispatch_target, + dispatch_lambda_args, + dispatch_lambda_exprs, + dispatch_lambda_return_str, + has_tensor_options, + PythonSignature, + PythonSignatureDeprecated, + PythonSignatureGroup, + PythonSignatureNativeFunctionPair, + signature, + signature_from_schema, + structseq_fieldnames, +) +from torchgen.code_template import CodeTemplate +from torchgen.context import with_native_function +from torchgen.gen import cpp_string, parse_native_yaml, parse_tags_yaml +from torchgen.model import ( + Argument, + BaseOperatorName, + FunctionSchema, + NativeFunction, + SchemaKind, + Type, + Variant, +) +from torchgen.utils import FileManager, split_name_params +from torchgen.yaml_utils import YamlLoader + +from .gen_inplace_or_view_type import is_tensor_list_type +from .gen_trace_type import should_trace + + +if TYPE_CHECKING: + from collections.abc import Callable, Iterable, Sequence + + +# +# declarations blocklist +# We skip codegen for these functions, for various reasons. +# Future PRs will categorize this list and eliminate or hoist +# them out of eager-only codegen. +# See https://github.com/pytorch/pytorch/issues/30788 +# + +# These functions require manual Python bindings or are not exposed to Python +_SKIP_PYTHON_BINDINGS = [ + "alias", + "contiguous", + "dim", + "get_device", + "is_contiguous", + "is_cuda", + "is_sparse", + "is_sparse_csr", + "numel", + "size", + "storage_offset", + "stride", + "sym_is_contiguous", + "sym_size", + "sym_stride", + "sym_storage_offset", + "sym_numel", + ".*_backward", + ".*_backward_(out|input|weight|bias)", + ".*_forward", + ".*_forward_out", + ".*_jvp", + "_unsafe_view", + "tensor", + "_?sparse_(coo|compressed|csr|csc|bsr|bsc)_tensor.*", + "_range.*", + "_sparse_add_out", + "_sparse_div.*", + "_sparse_mul.*", + "_sparse_sub.*", + "_sparse_dense_add_out", + "index", + "index_out", + "unique_dim_consecutive", + "_cumsum.*", + "_cumprod.*", + "_sum.*", + "_prod.*", + "_th_.*", + "_thnn_.*", + "range.*", + "_solve.*", + "_inverse.*", + "_cholesky.*", + "_triangular_solve.*", + "_qr.*", + "_svd.*", + "slice", + "item", + "_local_scalar_dense", + "to", + "_to_copy", + "_to_copy_out", + "_reshape_copy", + "_reshape_copy_out", + "copy_sparse_to_sparse_", + "copy_", + "_foreach_copy", + "numpy_T", + "matrix_H", + "mT", + "mH", # these need to be an attributes in Python, not functions + "nonzero(_(out|numpy))?", + "set_data", + ".*_overrideable", # overridable functions for backend extension + "data", + "is_leaf", + "output_nr", + "_version", + "requires_grad_", + "retains_grad", + "set_", + "_fw_primal", + "fake_quantize_per_tensor_affine_cachemask", + "fake_quantize_per_channel_affine_cachemask", + "_new_zeros_with_same_feature_meta", + "_has_same_storage_numel", # used for forward AD internals + "_reshape_alias", + "replace_", # only used by the functionalization pass, doesn't need to be exposed to python + "copy", # only used by the functionalization pass + "fill.Tensor", # only used by the functionalization pass + "fill.Scalar", # only used by the functionalization pass + "lift.*", + "normal_functional", # only used by the functionalization pass + "nbytes", + "itemsize", + "_batch_norm_with_update", + "_batch_norm_with_update_out", + "_batch_norm_no_update", +] + +SKIP_PYTHON_BINDINGS = [ + re.compile(rf"^{pattern}$") for pattern in _SKIP_PYTHON_BINDINGS +] + +# These function signatures are not exposed to Python. Note that this signature +# list does not support regex. +SKIP_PYTHON_BINDINGS_SIGNATURES = [ + "add.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor", + "add_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)", + "sub.Scalar(Tensor self, Scalar other, Scalar alpha=1) -> Tensor", + "sub_.Scalar(Tensor(a!) self, Scalar other, Scalar alpha=1) -> Tensor(a!)", + "mul.Scalar(Tensor self, Scalar other) -> Tensor", + "mul_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)", + "div.Scalar(Tensor self, Scalar other) -> Tensor", + "div_.Scalar(Tensor(a!) self, Scalar other) -> Tensor(a!)", +] + + +@with_native_function +def should_generate_py_binding(f: NativeFunction) -> bool: + # NativeFunctions that are entirely code-generated should not get python bindings + # because these codegen implementations are often inefficient. A handful of + # view_copy style ops were exposed accidentally when they were handwritten and now + # that we are moving them to codegen for bc reasons we need to keep them exposed in + # python. + if "generated" in f.tags and "view_copy" not in f.tags: + return False + + name = cpp.name(f.func) + for skip_regex in SKIP_PYTHON_BINDINGS: + if skip_regex.match(name): + return False + + signature = str(f.func) + for pattern in SKIP_PYTHON_BINDINGS_SIGNATURES: + if pattern == signature: + return False + return True + + +def get_pycname(name: BaseOperatorName) -> str: + return f"THPVariable_{name}" + + +def is_noarg(overloads: Sequence[PythonSignatureNativeFunctionPair]) -> bool: + return len(overloads) == 1 and overloads[0].signature.arguments_count() == 0 + + +def is_py_variable_method(f: NativeFunction) -> bool: + return f.python_module is None and Variant.method in f.variants + + +def is_py_torch_function(f: NativeFunction) -> bool: + return f.python_module is None and Variant.function in f.variants + + +def is_py_nn_function(f: NativeFunction) -> bool: + return f.python_module == "nn" + + +def is_py_fft_function(f: NativeFunction) -> bool: + return f.python_module == "fft" + + +def is_py_linalg_function(f: NativeFunction) -> bool: + return f.python_module == "linalg" + + +def is_py_nested_function(f: NativeFunction) -> bool: + return f.python_module == "nested" + + +def is_py_sparse_function(f: NativeFunction) -> bool: + return f.python_module == "sparse" + + +def is_py_special_function(f: NativeFunction) -> bool: + return f.python_module == "special" + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Main Function +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +def gen( + out: str, + native_yaml_path: str, + tags_yaml_path: str, + deprecated_yaml_path: str, + template_path: str, + *, + symint: bool = True, +) -> None: + fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + native_functions = parse_native_yaml( + native_yaml_path, tags_yaml_path + ).native_functions + native_functions = list(filter(should_generate_py_binding, native_functions)) + + methods = load_signatures(native_functions, deprecated_yaml_path, method=True) + create_python_bindings( + fm, + methods, + is_py_variable_method, + None, + "python_variable_methods.cpp", + method=True, + symint=symint, + ) + + # NOTE: num_shards here must be synced with gatherTorchFunctions in + # torch/csrc/autograd/python_torch_functions_manual.cpp + functions = load_signatures(native_functions, deprecated_yaml_path, method=False) + create_python_bindings_sharded( + fm, + functions, + is_py_torch_function, + "torch", + "python_torch_functions.cpp", + method=False, + num_shards=3, + symint=symint, + ) + + create_python_bindings( + fm, + functions, + is_py_nn_function, + "torch.nn", + "python_nn_functions.cpp", + method=False, + symint=symint, + ) + + create_python_bindings( + fm, + functions, + is_py_fft_function, + "torch.fft", + "python_fft_functions.cpp", + method=False, + symint=symint, + ) + + create_python_bindings( + fm, + functions, + is_py_linalg_function, + "torch.linalg", + "python_linalg_functions.cpp", + method=False, + symint=symint, + ) + + create_python_bindings( + fm, + functions, + is_py_nested_function, + "torch.nested", + "python_nested_functions.cpp", + method=False, + ) + + create_python_bindings( + fm, + functions, + is_py_sparse_function, + "torch.sparse", + "python_sparse_functions.cpp", + method=False, + symint=symint, + ) + + create_python_bindings( + fm, + functions, + is_py_special_function, + "torch.special", + "python_special_functions.cpp", + method=False, + symint=symint, + ) + + # Currently, we only use `functions` to generate `return_types` bindings. + # All methods which return structseq have function variant at this point. + # If any method only operator with structseq is added in the future, + # we will have to address that. + create_python_return_type_bindings( + fm, functions, lambda fn: True, "python_return_types.cpp" + ) + create_python_return_type_bindings_header( + fm, functions, lambda fn: True, "python_return_types.h" + ) + + valid_tags = parse_tags_yaml(tags_yaml_path) + + def gen_tags_enum() -> dict[str, str]: + return { + "enum_of_valid_tags": ( + "".join( + [f'\n.value("{tag}", at::Tag::{tag})' for tag in sorted(valid_tags)] + ) + ) + } + + fm.write("python_enum_tag.cpp", gen_tags_enum) + + +def group_filter_overloads( + pairs: Sequence[PythonSignatureNativeFunctionPair], + pred: Callable[[NativeFunction], bool], +) -> dict[BaseOperatorName, list[PythonSignatureNativeFunctionPair]]: + grouped: dict[BaseOperatorName, list[PythonSignatureNativeFunctionPair]] = ( + defaultdict(list) + ) + for pair in pairs: + if pred(pair.function): + grouped[pair.function.func.name.name].append(pair) + return grouped + + +def create_python_bindings( + fm: FileManager, + pairs: Sequence[PythonSignatureNativeFunctionPair], + pred: Callable[[NativeFunction], bool], + module: str | None, + filename: str, + *, + method: bool, + symint: bool = True, +) -> None: + """Generates Python bindings to ATen functions""" + py_methods: list[str] = [] + ops_headers: list[str] = [] + py_method_defs: list[str] = [] + py_forwards: list[str] = [] + + grouped = group_filter_overloads(pairs, pred) + + for name in sorted(grouped.keys(), key=str): + overloads = grouped[name] + py_methods.append( + method_impl(name, module, overloads, method=method, symint=symint) + ) + py_method_defs.append(method_def(name, module, overloads, method=method)) + py_forwards.extend(forward_decls(name, overloads, method=method)) + ops_headers.append(f"#include ") + + fm.write_with_template( + filename, + filename, + lambda: { + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/{filename}", + "ops_headers": ops_headers, + "py_forwards": py_forwards, + "py_methods": py_methods, + "py_method_defs": py_method_defs, + }, + ) + + +def create_python_return_type_bindings( + fm: FileManager, + pairs: Sequence[PythonSignatureNativeFunctionPair], + pred: Callable[[NativeFunction], bool], + filename: str, +) -> None: + """ + Generate function to initialize and return named tuple for native functions + which returns named tuple and registration invocations in `python_return_types.cpp`. + """ + py_return_types_definition: list[str] = [] + py_return_types_registrations: list[str] = [] + + grouped = group_filter_overloads(pairs, pred) + + for name in sorted(grouped.keys(), key=str): + overloads = grouped[name] + definitions, registrations = generate_return_type_definition_and_registrations( + overloads + ) + py_return_types_definition.append( + "" if not definitions else "\n".join(definitions) + ) + py_return_types_registrations.append( + "" if not registrations else "\n".join(registrations) + ) + + fm.write_with_template( + filename, + filename, + lambda: { + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/{filename}", + "py_return_types": py_return_types_definition, + "py_return_types_registrations": py_return_types_registrations, + }, + ) + + +def create_python_return_type_bindings_header( + fm: FileManager, + pairs: Sequence[PythonSignatureNativeFunctionPair], + pred: Callable[[NativeFunction], bool], + filename: str, +) -> None: + """ + Generate function to initialize and return named tuple for native functions + which returns named tuple and relevant entry for the map in `python_return_types.cpp`. + """ + py_return_types_declarations: list[str] = [] + + grouped = group_filter_overloads(pairs, pred) + + for name in sorted(grouped.keys(), key=str): + overloads = grouped[name] + declarations = generate_return_type_declarations(overloads) + py_return_types_declarations.append( + "" if not declarations else "\n".join(declarations) + ) + + fm.write_with_template( + filename, + filename, + lambda: { + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/{filename}", + "py_return_types_declarations": py_return_types_declarations, + }, + ) + + +def create_python_bindings_sharded( + fm: FileManager, + pairs: Sequence[PythonSignatureNativeFunctionPair], + pred: Callable[[NativeFunction], bool], + module: str | None, + filename: str, + *, + method: bool, + num_shards: int, + symint: bool = True, +) -> None: + """Generates Python bindings to ATen functions""" + grouped = group_filter_overloads(pairs, pred) + + def key_func( + kv: tuple[BaseOperatorName, list[PythonSignatureNativeFunctionPair]], + ) -> str: + return kv[0].base + + def env_func( + kv: tuple[BaseOperatorName, list[PythonSignatureNativeFunctionPair]], + ) -> dict[str, list[str]]: + name, fn_pairs = kv + return { + "ops_headers": [f"#include "], + "py_forwards": list(forward_decls(name, fn_pairs, method=method)), + "py_methods": [ + method_impl(name, module, fn_pairs, method=method, symint=symint) + ], + "py_method_defs": [method_def(name, module, fn_pairs, method=method)], + } + + fm.write_sharded( + filename, + grouped.items(), + base_env={ + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/{filename}", + }, + key_fn=key_func, + env_callable=env_func, + num_shards=num_shards, + sharded_keys={"ops_headers", "py_forwards", "py_methods", "py_method_defs"}, + ) + + +def load_signatures( + native_functions: list[NativeFunction], + deprecated_yaml_path: str, + *, + method: bool, + skip_deprecated: bool = False, + pyi: bool = False, +) -> Sequence[PythonSignatureNativeFunctionPair]: + @with_native_function + def gen_signature_pairs(f: NativeFunction) -> PythonSignatureNativeFunctionPair: + return PythonSignatureNativeFunctionPair( + signature=signature(f, method=method, pyi=pyi), + function=f, + ) + + pairs = list(map(gen_signature_pairs, native_functions)) + deprecated = load_deprecated_signatures( + pairs, deprecated_yaml_path, method=method, pyi=pyi + ) + return pairs if skip_deprecated else pairs + deprecated + + +def load_deprecated_signatures( + pairs: Sequence[PythonSignatureNativeFunctionPair], + deprecated_yaml_path: str, + *, + method: bool, + pyi: bool, +) -> list[PythonSignatureNativeFunctionPair]: + # The deprecated.yaml doesn't have complete type information, we need + # find and leverage the original ATen signature (to which it delegates + # the call) to generate the full python signature. + # We join the deprecated and the original signatures using type-only form. + + # group the original ATen signatures by name + grouped: dict[str, list[PythonSignatureNativeFunctionPair]] = defaultdict(list) + for pair in pairs: + grouped[pair.signature.name].append(pair) + + # find matching original signatures for each deprecated signature + results: list[PythonSignatureNativeFunctionPair] = [] + + with open(deprecated_yaml_path) as f: + deprecated_defs = yaml.load(f, Loader=YamlLoader) + + for deprecated in deprecated_defs: + schema = FunctionSchema.parse(deprecated["name"]) + aten_name, call_args = split_name_params(deprecated["aten"]) + is_out = aten_name.endswith("_out") + if is_out: + aten_name = aten_name.replace("_out", "") + + # HACK: these are fixed constants used to pass the aten function. + # The type must be known ahead of time + known_constants = { + "1": Type.parse("Scalar"), + } + schema_args_by_name = {a.name: a for a in schema.arguments.flat_all} + for name in call_args: + if name not in schema_args_by_name and name not in known_constants: + raise AssertionError( + f"deprecation definition: Unrecognized value {name}" + ) + + # Map deprecated signature arguments to their aten signature and test + # if the types and alias annotation match. + def is_schema_compatible( + aten_schema: FunctionSchema, + ) -> bool: + arguments: Iterable[Argument] + if is_out: + arguments = itertools.chain( + aten_schema.arguments.out, aten_schema.arguments.flat_non_out + ) + else: + arguments = aten_schema.arguments.flat_all + + for i, arg in enumerate(arguments): + if i < len(call_args): + arg_name = call_args[i] + if arg_name in known_constants: + schema_type = known_constants[arg_name] + schema_annotation = None + else: + schema_arg = schema_args_by_name[arg_name] + schema_type = schema_arg.type + schema_annotation = schema_arg.annotation + + if schema_type != arg.type or schema_annotation != arg.annotation: + return False + else: + if arg.default is None: + return False + + return len(schema.returns) == len(aten_schema.returns) and all( + a == b for a, b in zip(schema.returns, aten_schema.returns) + ) + + any_schema_found = False + for pair in grouped[aten_name]: + if not is_schema_compatible(pair.function.func): + continue + any_schema_found = True + + python_sig = signature_from_schema( + schema, + category_override=pair.function.category_override, + method=method, + pyi=pyi, + ) + + results.append( + PythonSignatureNativeFunctionPair( + signature=PythonSignatureDeprecated( + name=python_sig.name, + input_args=python_sig.input_args, + input_kwargs=python_sig.input_kwargs, + output_args=python_sig.output_args, + tensor_options_args=python_sig.tensor_options_args, + method=python_sig.method, + deprecated_schema=schema, + deprecated_args_exprs=tuple(call_args), + returns=python_sig.returns, + ), + function=pair.function, + ) + ) + if not any_schema_found: + raise AssertionError( + f"No native function with name {aten_name} matched signature:\n {str(schema)}" + ) + + return results + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Named Tuple Codegen +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +@with_native_function +def gen_structseq_typename_key(f: NativeFunction) -> str: + name = cpp.name(f.func) + fieldnames = structseq_fieldnames(f.func.returns) + return "_".join([name] + fieldnames) + + +def emit_structseq_call( + overloads: Sequence[PythonSignatureNativeFunctionPair], +) -> tuple[list[str], dict[str, str]]: + """ + Generate block of named tuple type def inits, and add typeref snippets + to declarations that use them + """ + typenames: dict[ + str, str + ] = {} # map from unique name + field name lists to typedef name + typedefs: list[str] = [] # typedef declarations and init code + + for overload in overloads: + fieldnames = structseq_fieldnames(overload.function.func.returns) + if not fieldnames: + continue + + name = cpp.name(overload.function.func) # use @with_native_function? + tn_key = gen_structseq_typename_key(overload.function) + typename = typenames.get(tn_key) + if typename is None: + typename = f"NamedTuple{'' if not typedefs else len(typedefs)}" + typenames[tn_key] = typename + typedefs.append( + f"""\ +static PyTypeObject* {typename} = generated::get_{name}_structseq();""" + ) + + return typedefs, typenames + + +def generate_return_type_definition_and_registrations( + overloads: Sequence[PythonSignatureNativeFunctionPair], +) -> tuple[list[str], list[str]]: + """ + Generate block of function in `python_return_types.cpp` to initialize + and return named tuple for a native function which returns named tuple + and registration invocations in same file. + """ + typenames: dict[ + str, str + ] = {} # map from unique name + field name lists to typedef name + definitions: list[str] = [] # function definition to register the typedef + registrations: list[str] = [] # register call for the typedef + + for overload in overloads: + fieldnames = structseq_fieldnames(overload.function.func.returns) + if not fieldnames: + continue + + fields = ", ".join(f'{{"{fn}", ""}}' for fn in fieldnames) + + name = cpp.name(overload.function.func) # use @with_native_function? + tn_key = gen_structseq_typename_key(overload.function) + typename = typenames.get(tn_key) + + if typename is None: + typename = f"{name}NamedTuple{'' if not definitions else len(definitions)}" + typenames[tn_key] = typename + definitions.append( + f"""\ +PyTypeObject* get_{name}_structseq() {{ + static PyStructSequence_Field NamedTuple_fields[] = {{ {fields}, {{nullptr}} }}; + static PyTypeObject {typename}; + static bool is_initialized = false; + static PyStructSequence_Desc desc = {{ "torch.return_types.{name}", nullptr, NamedTuple_fields, {len(fieldnames)} }}; + if (!is_initialized) {{ + PyStructSequence_InitType(&{typename}, &desc); + {typename}.tp_repr = (reprfunc)torch::utils::returned_structseq_repr; + is_initialized = true; + }} + return &{typename}; +}} +""" + ) + registrations.append( + f'addReturnType(return_types_module, "{name}", generated::get_{name}_structseq());' + ) + + return definitions, registrations + + +def generate_return_type_declarations( + overloads: Sequence[PythonSignatureNativeFunctionPair], +) -> list[str]: + """ + Generate block of function declarations in `python_return_types.h` to initialize + and return named tuple for a native function. + """ + typenames: dict[ + str, str + ] = {} # map from unique name + field name lists to typedef name + declarations: list[str] = [] # function declaration to register the typedef + + for overload in overloads: + fieldnames = structseq_fieldnames(overload.function.func.returns) + if not fieldnames: + continue + + name = cpp.name(overload.function.func) # use @with_native_function? + tn_key = gen_structseq_typename_key(overload.function) + typename = typenames.get(tn_key) + + if typename is None: + typename = ( + f"{name}NamedTuple{'' if not declarations else len(declarations)}" + ) + typenames[tn_key] = typename + declarations.append(f"PyTypeObject* get_{name}_structseq();") + + return declarations + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Method Impl Codegen +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + +# python binding for all overloads of a particular function/method +PY_VARIABLE_METHOD_VARARGS = CodeTemplate( + r"""\ +// ${name} +static PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs) +{ + ${method_header} + static PythonArgParser parser({ + ${signatures} + }, /*traceable=*/${traceable}); + + ParsedArgs<${max_args}> parsed_args; + auto _r = parser.parse(${self_}, args, kwargs, parsed_args); + ${check_has_torch_function} + switch (_r.idx) { + ${dispatch} + } + ${method_footer} +} + +""" +) + +# handler for a single parsed signature - may be a single overload or +# a pair of overloads that whose signatures only differ in output params +# (plugged into PY_VARIABLE_METHOD_VARARGS as an item in ${dispatch}) +PY_VARIABLE_CASE = CodeTemplate( + """\ +case ${overload_index}: { + ${body} +} +""" +) + +# python binding for single-overload function/method +PY_VARIABLE_METHOD_VARARGS_SINGLETON = CodeTemplate( + """\ +// ${name} +static PyObject * ${pycname}(PyObject* self_, PyObject* args, PyObject* kwargs) +{ + ${method_header} + static PythonArgParser parser({ + ${signatures} + }, /*traceable=*/${traceable}); + + ParsedArgs<${max_args}> parsed_args; + auto _r = parser.parse(${self_}, args, kwargs, parsed_args); + ${check_has_torch_function} + ${dispatch} + ${method_footer} +} + +""" +) + +# python binding for a method with no args, shortcuts parsing +PY_VARIABLE_METHOD_NOARGS = CodeTemplate( + """\ +// ${name} +static PyObject * ${pycname}(PyObject* self_, PyObject* args) +{ + ${method_header} + ${check_has_torch_function} + ${dispatch} + ${method_footer} +} + +""" +) + + +def method_impl( + name: BaseOperatorName, + module: str | None, + overloads: Sequence[PythonSignatureNativeFunctionPair], + *, + method: bool, + symint: bool = True, +) -> str: + """ + Generate a python binding for all overloads of an op. + """ + pycname = get_pycname(name) + noarg = is_noarg(overloads) + structseq_inits, structseq_typenames = emit_structseq_call(overloads) + + method_header = ["HANDLE_TH_ERRORS"] + method_header += structseq_inits + method_header += ( + ["const Tensor& self = THPVariable_Unpack(self_);"] if method else [] + ) + + method_footer = ([] if noarg else ["Py_RETURN_NONE;"]) + ["END_HANDLE_TH_ERRORS"] + + traceable = "true" if all(should_trace(o.function) for o in overloads) else "false" + + grouped_overloads: Sequence[PythonSignatureGroup] = group_overloads( + overloads, symint=symint + ) + is_singleton = len(grouped_overloads) == 1 + signatures: list[str] = [] + dispatch: list[str] = [] + for overload_index, overload in enumerate(grouped_overloads): + signature = overload.signature.signature_str(symint=symint) + signatures.append(f"{cpp_string(str(signature))},") + dispatch_body = emit_dispatch_case(overload, structseq_typenames, symint=symint) + dispatch.append( + PY_VARIABLE_CASE.substitute( + overload_index=overload_index, body=dispatch_body + ) + if not is_singleton + else dispatch_body + ) + + if noarg: + template = PY_VARIABLE_METHOD_NOARGS + elif is_singleton: + template = PY_VARIABLE_METHOD_VARARGS_SINGLETON + else: + template = PY_VARIABLE_METHOD_VARARGS + + return template.substitute( + name=name, + pycname=pycname, + method_header=method_header, + max_args=max(o.signature.arguments_count() for o in overloads), + signatures=signatures, + traceable=traceable, + check_has_torch_function=gen_has_torch_function_check( + name=name, + module=module, + noarg=noarg, + method=method, + ), + dispatch=dispatch, + method_footer=method_footer, + self_="self_" if method else "nullptr", + ) + + +def gen_has_torch_function_check( + name: BaseOperatorName, module: str | None, *, noarg: bool, method: bool +) -> str: + if noarg: + if method: + return f"""\ +if(check_has_torch_function(self_)) {{ + return handle_torch_function(self_, "{name}"); +}} +""" + else: + return "" + + self_ = "self_" if method else "nullptr" + namespace = ( + { + "torch": "THPVariableFunctionsModule", + "torch.nn": "THPNNVariableFunctionsModule", + "torch.fft": "THPFFTVariableFunctionsModule", + "torch.linalg": "THPLinalgVariableFunctionsModule", + "torch.nested": "THPNestedVariableFunctionsModule", + "torch.sparse": "THPSparseVariableFunctionsModule", + "torch.special": "THPSpecialVariableFunctionsModule", + }[module] + if module + else "THPVariableClass" + ) + + return f"""\ +if(_r.has_torch_function()) {{ + return handle_torch_function(_r, {self_}, args, kwargs, {namespace}, "{module or "torch.Tensor"}"); +}} +""" + + +# handler for output/no-output overload pair +PY_VARIABLE_OUT = CodeTemplate( + """\ +if (_r.isNone(${out_idx})) { + ${call_dispatch} +} else { + ${call_dispatch_out} +} +""" +) + + +def emit_dispatch_case( + overload: PythonSignatureGroup, + structseq_typenames: dict[str, str], + *, + symint: bool = True, +) -> str: + """ + Emit dispatch code for a single parsed signature. This corresponds to either + a single native function, or a pair that differ only in output params. In the + latter case, a single python signature is used for both and dispatching + switches on the presence/absence of passed output args. + """ + if overload.outplace is not None: + # dispatch output and no-output variants, branch on _r.isNone() + return PY_VARIABLE_OUT.substitute( + out_idx=overload.signature.output_idx(), + call_dispatch=emit_single_dispatch( + overload.signature, overload.base, structseq_typenames, symint=symint + ), + call_dispatch_out=emit_single_dispatch( + overload.signature, + overload.outplace, + structseq_typenames, + symint=symint, + ), + ) + else: + # no-output version only + return emit_single_dispatch( + overload.signature, overload.base, structseq_typenames, symint=symint + ) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Forward Declarations Codegen +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +def forward_decls( + name: BaseOperatorName, + overloads: Sequence[PythonSignatureNativeFunctionPair], + *, + method: bool, +) -> tuple[str, ...]: + if method: + return () + + pycname = get_pycname(name) + if is_noarg(overloads): + return ( + f"""\ +static PyObject * {pycname}(PyObject* self_, PyObject* args); +""", + ) + else: + return ( + f"""\ +static PyObject * {pycname}(PyObject* self_, PyObject* args, PyObject* kwargs); +""", + ) + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Method Def (Binding Table Entry) Codegen +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +def method_def( + name: BaseOperatorName, + module: str | None, + overloads: Sequence[PythonSignatureNativeFunctionPair], + *, + method: bool, +) -> str: + """ + Generate method def entry. + """ + pycname = get_pycname(name) + + if name.dunder_method: + # PyMethodDef entry for binary op, throws not implemented error + pycname = f"TypeError_to_NotImplemented_<{pycname}>" + + if is_noarg(overloads): + flags = "METH_NOARGS" if method else "METH_VARARGS | METH_KEYWORDS" + else: + pycname = f"castPyCFunctionWithKeywords({pycname})" + flags = "METH_VARARGS | METH_KEYWORDS" + + if module == "torch": + flags += " | METH_STATIC" + + return f'{{"{name}", {pycname}, {flags}, nullptr}},' + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Overload Sorting and Grouping +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +def group_overloads( + overloads: Sequence[PythonSignatureNativeFunctionPair], *, symint: bool = True +) -> Sequence[PythonSignatureGroup]: + bases: dict[str, PythonSignatureNativeFunctionPair] = {} + outplaces: dict[str, PythonSignatureNativeFunctionPair] = {} + + # first group by signature ignoring out arguments + for overload in overloads: + sig = overload.signature.signature_str(skip_outputs=True, symint=symint) + if overload.function.func.is_out_fn(): + if sig in outplaces: + raise RuntimeError( + f"Found duplicated function definition:\n- {overload.function.func}.\n" + f"Existing definition:\n- {outplaces[sig].function.func}." + ) + outplaces[sig] = overload + else: + if sig in bases: + raise RuntimeError( + f"Found duplicated function definition:\n- {overload.function.func}.\n" + f"Existing definition:\n- {bases[sig].function.func}." + ) + bases[sig] = overload + + for sig, out in outplaces.items(): + if sig not in bases: + candidates: list[str] = [] + for overload in overloads: + if ( + str(overload.function.func.name.name) + == str(out.function.func.name.name) + and not overload.function.func.is_out_fn() + and not overload.signature.deprecated + ): + candidates.append( + overload.signature.signature_str( + skip_outputs=True, symint=symint + ) + ) + out_sig = out.signature.signature_str(symint=symint) + raise RuntimeError( + f"While identifying overloads, we found an out schema {out_sig} without a corresponding non-out variant. " + f"We expected the non-out variant to have schema: \n- {sig}\nPlease check that you spelled the schema " + "correctly in native_functions.yaml. We discovered the following candidate(s): \n" + + "\n".join(f"- {candidate}" for candidate in candidates) + ) + + grouped = [ + PythonSignatureGroup.from_pairs( + functional=base, + out=outplaces.get(sig), + ) + for sig, base in bases.items() + ] + return sort_overloads(grouped, symint=symint) + + +# This function declares a partial order on declarations, and sorts them according +# to its linear extension. This is necessary, because there's some ambiguity in the +# choice of overload, and we want a different order. +# +# See Note[Order of overloads matters] +# +# A few examples of ambiguous python signature pairs. +# +# All parameters have the same type, except one taking Tensor the other taking +# Scalar. A numeric PyObject can be casted into Tensor, and a zero-dim Tensor +# object can be accepted as Scalar type parameter (see python_arg_parser.cpp). +# Therefore, same input arguments might be accepted by either python signature. +# We want to always parse the one taking Tensor first. +# +# bitwise_and(Tensor input, Tensor other, *, Tensor out=None) +# bitwise_and(Tensor input, Scalar other, *, Tensor out=None) +# +# If they have different number of parameters then they are not ambiguous - but +# the difference on output param can be ignored as it's optional. +# +# multiply(Tensor input, Tensor other, *, Tensor out=None) +# multiply(Tensor input, Scalar other) +# +# Both positional args and keyword-only args are considered together. +# +# subtract(Tensor other, *, Scalar alpha=1) +# subtract(Scalar other, Scalar alpha=1) +# +# A few ambiguous cases which it does NOT handle yet. +# +# If there is any difference in other parameters besides the Tensor/Scalar +# difference, then they are not considered ambiguous by this method anymore. +# However, the difference could be too trivial to disambiguate. +# +# foo(Tensor input, Scalar other, Scalar bar) +# foo(Tensor input, Tensor other, double bar) +# +# If they are taking different number of parameters then they are not considered +# ambiguous anymore, even if the difference is only on optional kwargs. +# +# foo(Scalar other, Scalar alpha=1) +# foo(Tensor other, *, Scalar alpha=1, Scalar beta=1) +# + + +def sort_overloads( + grouped_overloads: Sequence[PythonSignatureGroup], *, symint: bool = True +) -> Sequence[PythonSignatureGroup]: + # NB: Smaller here means lower priority + + def is_arg_smaller(t1: Type, t2: Type) -> bool: + return ( + str(t1) == "Scalar" + and str(t2) == "Tensor" + or str(t1) == "Scalar?" + and str(t2) == "Tensor?" + or "Dimname" in str(t1) + and "Dimname" not in str(t2) + or + # In the discussion https://github.com/pytorch/pytorch/issues/54555 it has been + # discussed why it is important to prioritize int/int? over int[] + str(t1) == "int[]" + and (str(t2) == "int" or str(t2) == "int?") + or + # TensorList currently throws an error during argument parsing, that's why it needs to be + # last in signature ordering. See discussion: https://github.com/pytorch/pytorch/issues/58087 + str(t1) == "Tensor[]" + and str(t2).find("[]") != -1 + or + # Prioritize IntArrayRef overload over SymIntArrayRef + str(t1) == "SymInt[]" + and str(t2) == "int[]" + or + # Make sure both in, SymInt are sorted consistently w.r.t. Tensor since Tensor can be implicitly + # converted to either int or SymInt. Prioritize the Tensor overload since it otherwise gets shadowed. + (str(t1) == "SymInt" or str(t1) == "int") + and str(t2) == "Tensor" + ) + + def is_smaller(s1: PythonSignature, s2: PythonSignature) -> bool: + """Returns True if s1 < s2 in the partial order.""" + args1, args2 = s1.arguments(skip_outputs=True), s2.arguments(skip_outputs=True) + if len(args1) != len(args2): + return False + # TODO: should use some canonical form instead of 'str(arg.type)' - see comments + # above. The old codegen used the deprecated 'dynamic_type(arg.type)', which + # ignores the optional annotation, i.e. 'Scalar' and 'Scalar?'. + equal = all(arg1.type == arg2.type for arg1, arg2 in zip(args1, args2)) + smaller_or_equal = all( + str(arg1.type) == str(arg2.type) or is_arg_smaller(arg1.type, arg2.type) + for arg1, arg2 in zip(args1, args2) + ) + return smaller_or_equal and not equal + + # First sort by signature + grouped_overloads = sorted( + grouped_overloads, key=lambda x: x.signature.signature_str(symint=symint) + ) + + # Construct the relation graph + larger_than: dict[int, set[int]] = defaultdict(set) + for i1, overload1 in enumerate(grouped_overloads): + for i2, overload2 in enumerate(grouped_overloads): + if is_smaller(overload1.signature, overload2.signature): + larger_than[i1].add(i2) + + if not larger_than: + return list(grouped_overloads) + + # Use a topological sort to sort overloads according to the partial order. + N = len(grouped_overloads) + sorted_ids: list[int] = list(filter(lambda x: x not in larger_than, range(N))) + + for idx in range(N): + # The size of sorted_ids will grow to N eventually. + i = sorted_ids[idx] + for j in sorted(larger_than.keys()): + larger = larger_than[j] + larger.discard(i) + if not larger: + del larger_than[j] + sorted_ids.append(j) + + return [grouped_overloads[x] for x in sorted_ids] + + +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # +# +# Codegen API Integration +# +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ # + + +def emit_single_dispatch( + ps: PythonSignature, + f: NativeFunction, + structseq_typenames: dict[str, str], + *, + symint: bool = True, +) -> str: + """ + Emit dispatch code for a single native function. + """ + + @with_native_function + def go(f: NativeFunction) -> str: + # header comments + if isinstance(ps, PythonSignatureDeprecated): + schema_comment = f"// [deprecated] aten::{ps.deprecated_schema}" + else: + schema_comment = f"// aten::{f.func}" + + # dispatch lambda signature + name = cpp.name(f.func) + lambda_formals = ", ".join( + f"{a.type_str} {a.name}" for a in dispatch_lambda_args(ps, f, symint=symint) + ) + lambda_return = dispatch_lambda_return_str(f) + + # dispatch lambda body + dispatch_callee = cpp_dispatch_target(f) + dispatch_args = ", ".join(cpp_dispatch_exprs(f, python_signature=ps)) + + # from arg parser outputs to dispatch lambda arguments + parser_outputs = arg_parser_output_exprs(ps, f, symint=symint) + lambda_arg_exprs = dispatch_lambda_exprs(ps, f, symint=symint) + inits = "\n".join(lambda_arg_exprs.inits) + lambda_args = ", ".join(lambda_arg_exprs.exprs) + + # scatter fields + # TODO: Checking `ps.method and ('requires_grad' in parser_outputs)` is a hacky + # solution for enabling the 'requires_grad' argument for tensor methods + # new_full, new_empty, and new_zeros. A much better but more difficult to + # implement solution involves refactoring according to Ed's description here: + # https://github.com/pytorch/pytorch/issues/36455#issuecomment-614767589 + need_set_requires_grad = ps.tensor_options_args and ( + not has_tensor_options(f) + or (ps.method and ("requires_grad" in parser_outputs)) + ) + set_requires_grad = ( + f".set_requires_grad({parser_outputs['requires_grad'].expr})" + if need_set_requires_grad + else "" + ) + + if lambda_return == "void": + # Make in-place foreach return `self` at python-binding level. + # ref: https://github.com/pytorch/pytorch/pull/118622#pullrequestreview-1904804954 + self_arg = f.func.arguments.self_arg + return_stmt: str + if ( + str(f.func.name).startswith("_foreach_") + and f.func.kind() == SchemaKind.inplace + ): + # note(crcrpar): `_foreach_pow.ScalarAndTensor` does NOT have its in-place + # variant and it unlikely to have it in the future. Thus it's safe to have the following check. + if self_arg is None or not is_tensor_list_type(self_arg.argument.type): + raise AssertionError( + "Expected self_arg to be a tensor list type for inplace foreach" + ) + return_stmt = """PyObject* self_tensorlist = _r.args[0]; +Py_INCREF(self_tensorlist); +return self_tensorlist; +""" + else: + return_stmt = "Py_RETURN_NONE;" + return f"""\ +{schema_comment} +{inits} +auto dispatch_{name} = []({lambda_formals}) -> {lambda_return} {{ + pybind11::gil_scoped_release no_gil; + {dispatch_callee}({dispatch_args}); +}}; +dispatch_{name}({lambda_args}){set_requires_grad}; +{return_stmt} +""" + else: + typename = structseq_typenames.get(gen_structseq_typename_key(f)) + structseq_typeref = f"{typename}, " if typename is not None else "" + return f"""\ +{schema_comment} +{inits} +auto dispatch_{name} = []({lambda_formals}) -> {lambda_return} {{ + pybind11::gil_scoped_release no_gil; + return {dispatch_callee}({dispatch_args}); +}}; +return wrap({structseq_typeref}dispatch_{name}({lambda_args}){set_requires_grad}); +""" + + return go(f) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_trace_type.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_trace_type.py new file mode 100644 index 0000000000000000000000000000000000000000..ba75c6e84df0396e4668b5628f8fca87adeaf93b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_trace_type.py @@ -0,0 +1,542 @@ +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +from torchgen.api import cpp +from torchgen.api.types import DispatcherSignature +from torchgen.code_template import CodeTemplate +from torchgen.context import with_native_function +from torchgen.model import Argument, NativeFunction, SchemaKind, TensorOptionsArguments +from torchgen.utils import FileManager + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# Note [Manual Backend kernels] +# For these ops, we want to manually register to dispatch key Backend and +# skip codegen-ed registration to all keys before Backend. +# For codegen this means: +# - op set below must match ops with manual_kernel_registration=True in native_functions.yaml +# where we skip codegen backend kernels +# - all ops below are part of MANUAL_AUTOGRAD to skip codegen Autograd kernel registration +# - all ops below are part of MANUAL_TRACER to skip codegen Tracer kernel registration +# Note: we still register to dispatch key Profiler for these ops, keeping it untouched for now. +# You can find the manual registration in torch/csrc/autograd/VariableTypeManual.cpp +MANUAL_BACKEND = { + "options", + "data", + "set_data", + "is_leaf", + "output_nr", + "_version", + "retain_grad", + "_backward", + "requires_grad_", +} + +# For these ops we want to skip the codegen-ed registration to both Autograd and Tracer keys. +# You can find the manual registration in torch/csrc/autograd/VariableTypeManual.cpp +MANUAL_AUTOGRAD_AND_TRACER = { + "resize_", + "resize_as_", + "detach", + "detach_", + "copy_", + "_fw_primal", + "_make_dual", +} + +# Currently MANUAL_AUTOGRAD and MANUAL_TRACER share the same set of ops: +# union(MANUAL_BACKEND, MANUAL_AUTOGRAD_AND_TRACER) +# You can find the manual registration in torch/csrc/autograd/VariableTypeManual.cpp +MANUAL_AUTOGRAD = MANUAL_TRACER = MANUAL_BACKEND | MANUAL_AUTOGRAD_AND_TRACER + +# These functions we don't want to record for tracing, because we always want +# to trace their constituent parts. This is a temporary hack in lieue +# of proper scopes, where subsequent compilation passes can ask for the unfolding +# on demand. Only concrete ATen methods can be disabled this way; it will have +# NO EFFECT otherwise. +DONT_RECORD_TRACE = { + "convolution", + "conv1d", + "conv2d", + "conv3d", + "conv_transpose1d", + "conv_transpose2d", + "conv_transpose3d", + "lstm_cell", + "gru_cell", + "rnn_tanh_cell", + "rnn_relu_cell", + # FIXME: figure out a better way when we support sparse tensors in jit + "_coalesced", +} + + +def should_trace(f: NativeFunction) -> bool: + # Operations involving Storage or Type are not traceable at the moment + if any( + str(arg.type) in {"Storage", "Type"} for arg in f.func.schema_order_arguments() + ): + return False + # We can't trace functions which don't have any Tensor or TensorList returns + if not any(r.type.is_tensor_like() for r in f.func.returns): + return False + return f.func.name.name.base not in DONT_RECORD_TRACE + + +SELECT = CodeTemplate( + """\ + +if (${cond}) { + ${true} +} else { + ${false} +} +""" +) + +OP_NAME = CodeTemplate( + """\ +op_name = c10::Symbol::fromQualString("aten::${trace_name}"); +""" +) + +# These functions have their names recorded under trace renamed, +RENAME_TRACE = { + "zero": "zeros_like", # replacing aten::zero_ with aten::zeros_like + "fill": "full_like", # replacing aten::fill_ with aten::full_like +} + + +def format_trace_op_name(f: NativeFunction) -> str: + # TODO: byte-for-byte compatible with old codegen behavior - should clean up + if ( + f.func.kind() in (SchemaKind.functional, SchemaKind.out) + or f.func.name.name.dunder_method + ): + # special case for *_out functions: the in-place and out-of-place ops + # are overloaded with the same name in the JIT + trace_name = str(f.func.name.name) + trace_name = RENAME_TRACE.get(trace_name, trace_name) + return OP_NAME.substitute(trace_name=trace_name) + + # otherwise, this is an in-place op and we need to emit both in- and + # out-of-place versions + outplace_trace_name = f.func.name.name.base + inplace_trace_name = cpp.name(f.func) + outplace_trace_name = RENAME_TRACE.get(outplace_trace_name, outplace_trace_name) + inplace_trace_name = RENAME_TRACE.get(inplace_trace_name, inplace_trace_name) + + return SELECT.substitute( + cond="tracer_state->force_outplace", + true=OP_NAME.substitute(trace_name=outplace_trace_name), + false=OP_NAME.substitute(trace_name=inplace_trace_name), + ) + + +ADD_TRACE_INPUT = CodeTemplate("""jit::tracer::addInputs(node, "${name}", ${input});""") + + +def format_trace_inputs(f: NativeFunction) -> str: + def dispatch_trace_input(arg: Argument | TensorOptionsArguments) -> Sequence[str]: + if isinstance(arg, TensorOptionsArguments): + name = "options" + return [ + ADD_TRACE_INPUT.substitute( + name=name, input="c10::optTypeMetaToScalarType(options.dtype_opt())" + ), + ADD_TRACE_INPUT.substitute(name=name, input="options.layout()"), + ADD_TRACE_INPUT.substitute(name=name, input="options.device()"), + ADD_TRACE_INPUT.substitute(name=name, input="options.pinned_memory()"), + ] + else: + name = arg.name + if str(arg.type) == "Tensor?[]": + return [f'jit::tracer::addInputs(node, "{name}", {name});'] + else: + return [ADD_TRACE_INPUT.substitute(name=name, input=name)] + + args: list[Argument | TensorOptionsArguments] = list( + f.func.schema_order_arguments() + ) + + if f.func.is_out_fn(): + # *_out functions take the result as a separate argument, but we don't want to + # trace that argument directly. Instead, we trace its TensorOptions. + # So first, we need to remove the out argument from the list of arguments to trace. + num_out_args = len(f.func.arguments.out) + args = args[:-num_out_args] + + trace_inputs = itertools.chain.from_iterable( + dispatch_trace_input(arg) for arg in args + ) + + if f.func.is_out_fn(): + # for *_out functions, handle the result argument differently for inplace/outplace. + # For inplace: just add the input to the end to confirm with the JIT schema + inplace = [ + ADD_TRACE_INPUT.substitute( + name=f.func.arguments.out[i].name, input=f.func.arguments.out[i].name + ) + # pyrefly: ignore [unbound-name] + for i in range(num_out_args) + ] + + # for outplace: do nothing, except if the function is a factory. + # Factories are a bit special because their out-of-place overloads + # take an extra TensorOptions argument, which is missing in the _out function + has_tensor_return = any(r.type.is_tensor_like() for r in f.func.returns) + has_tensor_input_arg = any( + a.type.is_tensor_like() for a in f.func.arguments.flat_non_out + ) + is_factory_method = f.category_override == "factory" or ( + has_tensor_return and not has_tensor_input_arg + ) + + # HACK: preserve old codegen behavior - the old codegen set the `is_factory_method` + # flag for the whole family of ops with the same basename if any of them is a + # factory method. For most cases the whole family of ops are indeed all factory + # method - 'normal' is the only exception. So we handle it specially here to avoid + # cloning the old logic. + if f.func.name.name.base == "normal": + is_factory_method = True + + if is_factory_method: + outplace = [ + ADD_TRACE_INPUT.substitute( + name="out", + input="c10::optTypeMetaToScalarType(out.options().dtype_opt())", + ), + ADD_TRACE_INPUT.substitute(name="out", input="out.options().layout()"), + ADD_TRACE_INPUT.substitute(name="out", input="out.options().device()"), + ADD_TRACE_INPUT.substitute( + name="out", input="out.options().pinned_memory()" + ), + ] + else: + outplace = [] + + trace_inputs = itertools.chain( + trace_inputs, + [ + SELECT.substitute( + cond="tracer_state->force_outplace", + true="\n".join(outplace), + false="\n".join(inplace), + ) + ], + ) + + return "\n".join(trace_inputs) + + +# `torch.jit.trace` have undocumented keyword argument `_force_outplace`, +# which force jit to replace functions with outplace variants (for +# example `aten::add_` becomes `aten::add`). +# +# This replacement implemented in-place with minimum modifications of +# arguments stack (as it assumes that outplace call has the same arguments +# as inplace version). +# +# However there are no such substitutions available for `aten::fill_` +# and `aten::zero_` operators, as we never implemented `aten::fill` +# and `aten::zero`. So jit tracing hack replacing `aten::zero_` with +# `aten::zeros_like` and replacing `aten::fill_` with `aten::full_like`. +# +# But as they potentially can have different arguments, we also have +# to hack into the stack and add missing ones. +# +# A possible alternative would be: +# +# - Add `aten::fill` and `aten::zero` +# +# - Or keep `aten::zeros_like` arguments aligned with `aten::zero_` +# arguments (inside of the `native_functions.yaml`) +RENAME_TRACE_ADD_ARGS = { + "fill": """\ + jit::tracer::addInputs(node, "options", ::std::optional()); + jit::tracer::addInputs(node, "options", layout_or_default(::std::nullopt)); + jit::tracer::addInputs(node, "options", device_or_default(::std::nullopt)); + jit::tracer::addInputs(node, "options", pinned_memory_or_default(::std::nullopt)); + ::std::optional memory_format = c10::MemoryFormat::Preserve; + jit::tracer::addInputs(node, "memory_format", memory_format); +""", + "zero": """\ + jit::tracer::addInputs(node, "options", ::std::optional()); + jit::tracer::addInputs(node, "options", layout_or_default(::std::nullopt)); + jit::tracer::addInputs(node, "options", device_or_default(::std::nullopt)); + jit::tracer::addInputs(node, "options", pinned_memory_or_default(::std::nullopt)); + ::std::optional memory_format = c10::MemoryFormat::Preserve; + jit::tracer::addInputs(node, "memory_format", memory_format); +""", +} + +INPLACE_GUARD = CodeTemplate( + """\ +jit::tracer::ensureUniqueIfOutOfPlaced("${name}", ${mutable_input}); +""" +) + +PRE_RECORD_TRACE = CodeTemplate( + """\ +torch::jit::Node* node = nullptr; +std::shared_ptr tracer_state; +if (jit::tracer::isTracing()) { + tracer_state = jit::tracer::getTracingState(); + at::Symbol op_name; + ${set_op_name} + node = tracer_state->createNode(op_name, /*num_outputs=*/0); + jit::tracer::recordSourceLocation(node); + ${add_trace_inputs} + tracer_state->insertNode(node); + ${inplace_guard} + jit::tracer::setTracingState(nullptr); +} +""" +) + + +def format_prerecord_trace(f: NativeFunction) -> str: + if not should_trace(f): + return "" + + # TODO: clean up old codegen behavior + is_inplace = ( + f.func.kind() in (SchemaKind.inplace, SchemaKind.out) + and not f.func.name.name.dunder_method + ) + add_args = ( + RENAME_TRACE_ADD_ARGS.get(f.func.name.name.base, "") if is_inplace else "" + ) + additional_inputs = ( + SELECT.substitute( + cond="tracer_state->force_outplace", + true=add_args, + false="", + ) + if add_args + else "" + ) + + return PRE_RECORD_TRACE.substitute( + set_op_name=format_trace_op_name(f), + add_trace_inputs=format_trace_inputs(f) + additional_inputs, + inplace_guard=INPLACE_GUARD.substitute( + name=cpp.name(f.func), + mutable_input=f.func.arguments.out[0].name + if f.func.arguments.out + else "self", + ) + if is_inplace + else "", + ) + + +POST_RECORD_TRACE = CodeTemplate( + """\ +if (tracer_state) { + jit::tracer::setTracingState(std::move(tracer_state)); + ${add_trace_outputs} +} +""" +) + + +def format_postrecord_trace(f: NativeFunction) -> str: + if not should_trace(f): + return "" + + # For outplacing ops, *_out overloads require special handling to move the + # output *argument* to a return value + if f.func.is_out_fn(): + output_names_outplace = [arg.name for arg in f.func.arguments.out] + output_names_inplace = cpp.return_names(f) + + # Code size optimization: the common case is that the return value is + # the same for both variants + if output_names_outplace == output_names_inplace: + outputs = [ + f"jit::tracer::addOutput(node, {n});" for n in output_names_outplace + ] + return POST_RECORD_TRACE.substitute(add_trace_outputs=outputs) + + selection = SELECT.substitute( + cond="force_outplace", + true="\n".join( + f"jit::tracer::addOutput(node, {n});" for n in output_names_outplace + ), + false="\n".join( + f"jit::tracer::addOutput(node, {n});" for n in output_names_inplace + ), + ) + return POST_RECORD_TRACE.substitute(add_trace_outputs=selection) + else: + output_names = cpp.return_names(f) + outputs = [f"jit::tracer::addOutput(node, {n});" for n in output_names] + return POST_RECORD_TRACE.substitute(add_trace_outputs=outputs) + + +def tie_return_values(f: NativeFunction) -> str: + if len(f.func.returns) == 1: + return f"auto {f.func.returns[0].name or 'result'}" + names = cpp.return_names(f) + return f"auto [{', '.join(names)}]" + + +def get_return_value(f: NativeFunction) -> str: + names = cpp.return_names(f) + if len(f.func.returns) == 1: + return names[0] + if f.func.kind() == SchemaKind.out: + return f"std::forward_as_tuple({', '.join(names)})" + else: + moved = ", ".join(f"std::move({name})" for name in names) + return f"std::make_tuple({moved})" + + +TRACE_DISPATCH = CodeTemplate( + """\ +${assign_return_values}at::_ops::${unambiguous_name}::redispatch(${unpacked_args});""" +) + + +def emit_trace_body(f: NativeFunction) -> list[str]: + trace_body: list[str] = [] + + trace_body.append(format_prerecord_trace(f)) + + dispatcher_sig = DispatcherSignature.from_schema(f.func) + dispatcher_exprs = dispatcher_sig.exprs() + + # code-generated tracing kernels plumb and recompute dispatch keys directly through the kernel for performance. + # See Note [Plumbing Keys Through The Dispatcher] for details. + dispatch_key_set = "ks & c10::DispatchKeySet(c10::DispatchKeySet::FULL_AFTER, c10::DispatchKey::Tracer)" + redispatch_args = ", ".join([dispatch_key_set] + [a.expr for a in dispatcher_exprs]) + + assign_return_values = ( + f"{tie_return_values(f)} = " + if f.func.kind() in [SchemaKind.functional, SchemaKind.mutable] + and f.func.returns + else "" + ) + + # Note that this calls the slow, dispatching variants of manual_cpp_binding ops. + # We could probably work harder to ensure that the fast variants are + # called instead, but the perf benefit would be minimal. + trace_body.append( + TRACE_DISPATCH.substitute( + assign_return_values=assign_return_values, + unambiguous_name=f.func.name.unambiguous_name(), + unpacked_args=redispatch_args, + ) + ) + + trace_body.append(format_postrecord_trace(f)) + if f.func.returns: + trace_body.append(f"return {get_return_value(f)};") + return trace_body + + +METHOD_DEFINITION = CodeTemplate( + """\ +${return_type} ${type_wrapper_name}(${formals}) { + ${type_definition_body} +} +""" +) + + +def type_wrapper_name(f: NativeFunction, key: str = "Default") -> str: + if f.func.name.overload_name: + name = f"{cpp.name(f.func)}_{f.func.name.overload_name}" + else: + name = cpp.name(f.func) + + # The key argument is only used in gen_variable_type where we need fns per autograd dispatch key. + # In gen_trace_type and gen_inplace_view_type where only one fn per native_fn must be generated, + # the key argument should not be passed. + # We do not append key if it is Default so that generated functions from + # before per-dispatch-key derivatives were added retain the same names. + if key != "Default": + name = name + f"_{key}" + return name + + +@with_native_function +def method_definition(f: NativeFunction) -> str: + if cpp.name(f.func) in MANUAL_TRACER: + raise AssertionError(f"Function {cpp.name(f.func)} is in MANUAL_TRACER") + + formals = ", ".join( + # code-generated tracing kernels plumb and recompute dispatch keys directly through the kernel for performance. + # See Note [Plumbing Keys Through The Dispatcher] for details. + ["c10::DispatchKeySet ks"] + + [ + f"{cpp.argument_type(a, binds='__placeholder__', symint=True).cpp_type()} {a.name}" + for a in f.func.schema_order_arguments() + ] + ) + + return METHOD_DEFINITION.substitute( + return_type=cpp.returns_type(f.func.returns, symint=True).cpp_type(), + type_wrapper_name=type_wrapper_name(f), + formals=formals, + type_definition_body=emit_trace_body(f), + ) + + +WRAPPER_REGISTRATION = CodeTemplate( + """\ +m.impl("${name}", + TORCH_FN(${class_type}::${type_wrapper_name}) +); +""" +) + + +@with_native_function +def method_registration(f: NativeFunction) -> str: + if cpp.name(f.func) in MANUAL_TRACER: + raise AssertionError(f"Function {cpp.name(f.func)} is in MANUAL_TRACER") + + return WRAPPER_REGISTRATION.substitute( + name=f.func.name, + type_wrapper_name=type_wrapper_name(f), + class_type="TraceType", + ) + + +def gen_trace_type_func(fn: NativeFunction) -> dict[str, list[str]]: + return { + "ops_headers": [f"#include "], + "trace_method_definitions": [method_definition(fn)], + "trace_wrapper_registrations": [method_registration(fn)], + } + + +def gen_trace_type( + out: str, native_functions: list[NativeFunction], template_path: str +) -> None: + # NOTE: see Note [Sharded File] at the top of the VariableType.cpp + # template regarding sharding of the generated files. + fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + fm.write_sharded( + "TraceType.cpp", + [fn for fn in native_functions if cpp.name(fn.func) not in MANUAL_TRACER], + key_fn=lambda fn: fn.root_name, + base_env={ + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/TraceType.cpp", + }, + env_callable=gen_trace_type_func, + num_shards=5, + sharded_keys={ + "ops_headers", + "trace_method_definitions", + "trace_wrapper_registrations", + }, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_variable_factories.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_variable_factories.py new file mode 100644 index 0000000000000000000000000000000000000000..9916a77385d38f01e83416d4303cb17ac17de700 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_variable_factories.py @@ -0,0 +1,116 @@ +# Generates C++ functions that wrap ATen tensor factory methods to turn them into Variables. +# +# This writes one file: variable_factories.h + +from __future__ import annotations + +import re + +import torchgen.api.python as python +from torchgen.api import cpp +from torchgen.api.types import CppSignatureGroup +from torchgen.context import with_native_function +from torchgen.gen import parse_native_yaml +from torchgen.model import NativeFunction, TensorOptionsArguments, Variant +from torchgen.utils import FileManager, mapMaybe + + +OPTIONAL_TYPE_PATTERN = re.compile(r"std::optional<(.+)>") +TYPE_PATTERN = re.compile(r"(?:const\s+)?([A-Z]\w+)") + + +# Add 'at::' to types defined in ATen namespace, e.g. Tensor, TensorList, IntArrayRef and etc. +# TODO: maybe update the cpp argument API to take optional namespace argument? +def fully_qualified_type(argument_type: str) -> str: + def maybe_optional_type(type: str, is_opt: bool) -> str: + return f"std::optional<{type}>" if is_opt else type + + opt_match = OPTIONAL_TYPE_PATTERN.match(argument_type) + is_opt = opt_match is not None + if opt_match: + argument_type = argument_type[opt_match.start(1) : opt_match.end(1)] + match = TYPE_PATTERN.match(argument_type) + if match is None: + return maybe_optional_type(argument_type, is_opt) + index = match.start(1) + qualified_type = f"{argument_type[:index]}at::{argument_type[index:]}" + return maybe_optional_type(qualified_type, is_opt) + + +def gen_variable_factories( + out: str, native_yaml_path: str, tags_yaml_path: str, template_path: str +) -> None: + native_functions = parse_native_yaml( + native_yaml_path, tags_yaml_path + ).native_functions + factory_functions = [fn for fn in native_functions if is_factory_function(fn)] + fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + fm.write_with_template( + "variable_factories.h", + "variable_factories.h", + lambda: { + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/variable_factories.h", + "ops_headers": [ + f"#include " for fn in factory_functions + ], + "function_definitions": list(mapMaybe(process_function, factory_functions)), + }, + ) + + +@with_native_function +def is_factory_function(f: NativeFunction) -> bool: + if Variant.function not in f.variants: + return False + + name = cpp.name(f.func) + has_tensor_options = python.has_tensor_options(f) + return has_tensor_options or name.endswith("_like") + + +@with_native_function +def process_function(f: NativeFunction) -> str | None: + name = cpp.name(f.func) + has_tensor_options = python.has_tensor_options(f) + is_factory = has_tensor_options or name.endswith("_like") + + if Variant.function not in f.variants or not is_factory: + return None + + cpp_sigs = CppSignatureGroup.from_native_function(f, method=False) + sigs = [cpp_sigs.signature] + if cpp_sigs.symint_signature is not None: + sigs.append(cpp_sigs.symint_signature) + r = "" + for sig in sigs: + formals: list[str] = [] + exprs: list[str] = [] + requires_grad = "false" + for arg in sig.arguments(): + qualified_type = fully_qualified_type(arg.type) + if arg.default: + formals.append(f"{qualified_type} {arg.name} = {arg.default}") + else: + formals.append(f"{qualified_type} {arg.name}") + + if isinstance(arg.argument, TensorOptionsArguments): + # note: we remove the requires_grad setting from the TensorOptions because + # it is ignored anyways (and we actually have an assertion that it isn't set + # which would fail otherwise). We handle requires_grad explicitly here + # instead of passing it through to the kernel. + exprs.append( + f"at::TensorOptions({arg.name}).requires_grad(::std::nullopt)" + ) + # Manually set the requires_grad bit on the result tensor. + requires_grad = f"{arg.name}.requires_grad()" + else: + exprs.append(arg.name) + + r += f"""\ +inline at::Tensor {sig.name()}({", ".join(formals)}) {{ + at::AutoDispatchBelowADInplaceOrView guard; + return autograd::make_variable(at::{sig.name()}({", ".join(exprs)}), /*requires_grad=*/{requires_grad}); +}} +""" + return r diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_variable_type.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_variable_type.py new file mode 100644 index 0000000000000000000000000000000000000000..6e45f29a8232a451d6f96e17b31b613d14b3fd90 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_variable_type.py @@ -0,0 +1,2266 @@ +# Generates VariableType.h/cpp +# +# **If any changes are being made to the VariableType codegen please also check +# if updates are needed in torch/csrc/autograd/autograd_not_implemented_fallback.cpp +# +# VariableType is a subclass of at::Type that provides the binding code +# necessary to provide a differentiable version of ATen operators. There are a +# number of different things we could mean: +# +# - Given a non-differentiable forward implementation, we might +# directly associate it with a backward implementation to make +# it differentiable. This is the common case. +# +# - Some functions don't need a backwards implementation, because +# backpropagation will never propagate beyond them. There are a +# number of different reasons why this may be the case: +# +# - The function has no differentiable inputs +# - The function's output is not differentiable +# - The function has no data dependency on its input +# +# - Some function don't need a backwards implementation because they +# are implemented as a composition of other (differentiable) ATen +# functions. These are dispatched directly to the Type superclass, +# which will in turn dispatch back to VariableType for its +# differentiable subcomponents. +# + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING + +from torchgen.api import cpp +from torchgen.api.autograd import ( + DifferentiableInput, + dispatch_strategy, + ForwardDerivative, + gen_differentiable_outputs, + is_differentiable, + NativeFunctionWithDifferentiabilityInfo, + SavedAttribute, +) +from torchgen.api.types import ( + ArrayRefCType, + BaseCppType, + BaseCType, + Binding, + intArrayRefT, + iTensorListRefT, + ListCType, + MutRefCType, + OptionalCType, + scalarT, + SpecialArgName, + stringT, + symIntArrayRefT, + TENSOR_LIST_LIKE_CTYPES, + tensorListT, + tensorT, + TupleCType, + VectorCType, +) +from torchgen.code_template import CodeTemplate +from torchgen.context import ( + native_function_manager, + with_native_function, + with_native_function_and, +) +from torchgen.model import ( + Argument, + BaseType, + ListType, + NativeFunction, + SchemaKind, + SelfArgument, + TensorOptionsArguments, +) +from torchgen.utils import FileManager, mapMaybe + +from .context import with_native_function_with_differentiability_info_and_key +from .gen_inplace_or_view_type import ( + ALL_VIEW_FUNCTIONS, + ASSIGN_RETURN_VALUE, + AUTOGRAD_NOT_IMPLEMENTED_REGISTRATION, + gen_formals, + get_base_name, + get_view_info, + is_tensor_list_type, + is_tensor_type, + METHOD_DEFINITION, + modifies_arguments, + TMP_VAR, + unpack_args, + unpacked_name, + use_derived, + WRAPPER_REGISTRATION, +) +from .gen_trace_type import ( + get_return_value, + MANUAL_AUTOGRAD_AND_TRACER, + MANUAL_BACKEND, + tie_return_values, + type_wrapper_name, +) + + +if TYPE_CHECKING: + from collections.abc import Callable, Sequence + + +# We don't set or modify grad_fn on these methods. Generally, they return +# tensors that have requires_grad=False. In-place functions listed here will +# not examine or modify requires_grad or grad_fn. +# NB: this does NOT include overload name +DONT_REQUIRE_DERIVATIVE = { + # These only depend on the input Tensor's shape and device, not the data + "empty_like", + "ones_like", + "full_like", + "zeros_like", + "rand_like", + "randn_like", + "new_empty", + "new_empty_strided", + "new_full", + "new_zeros", + "new_ones", + # These are only implemented on integral types + "__and__", + "__iand__", + "__ilshift__", + "__ior__", + "__irshift__", + "__ixor__", + "__lshift__", + "__or__", + "__rshift__", + "__xor__", + # These work on integral data types, and hence don't require derivative + "_sobol_engine_draw", + "_sobol_engine_ff", + "_sobol_engine_scramble_", + "_sobol_engine_initialize_state_", + # This is an unsafe method that is meant to be out of reach of autograd. + "_coalesced_", + # Quantize functions should not record gradients + "quantize_per_tensor", + "quantize_per_channel", + # Functions that return integers should not have output that require gradients + "argmax", + "argmin", + "argsort", + "searchsorted", + "bucketize", + # Functions that return booleans are not differentiable + "isnan", + "isposinf", + "isneginf", + "isinf", + "signbit", + "isin", + "allclose", + # Functions return none are not differentiable + "record_stream", + # These functions are not differentiable + "logical_and", + "logical_xor", + "logical_not", + "logical_or", + # This function returns nested_tensor shape as a tensor that is non-differentiable + "_nested_tensor_size", + "_nested_tensor_strides", + "_nested_tensor_storage_offsets", +} + +# The C -> R functions at the time of adding this are still being audited and tested +# but will not error out. +# C -> C, R -> C functions for which backward is correctly implemented and tested +GRADIENT_IMPLEMENTED_FOR_COMPLEX = { + "fill", + "t", + "t_copy", + "view", + "reshape", + "reshape_as", + "view_as", + "view_copy", + "roll", + "clone", + "block_diag", + "diag_embed", + "repeat", + "expand", + "expand_copy", + "flip", + "fliplr", + "flipud", + "rot90", + "nanmean", + "nansum", + "transpose", + "transpose_copy", + "permute", + "permute_copy", + "squeeze", + "squeeze_copy", + "unsqueeze", + "unsqueeze_copy", + "resize", + "resize_as", + "tril", + "triu", + "chunk", + "zero_", + "eq_", + "ne_", + "add", + "__radd__", + "sum", + "_conj", + "sin", + "cos", + "mul", + "sinc", + "sinh", + "cosh", + "__rmul__", + "sgn", + "asin", + "acos", + "sub", + "div", + "cat", + "view_as_complex", + "index_put", + "neg", + "complex", + "select", + "where", + "as_strided", + "as_strided_copy", + "as_strided_scatter", + "slice", + "constant_pad_nd", + "unbind", + "unbind_copy", + "split", + "split_with_sizes", + "unsafe_split", + "split_with_sizes_backward", + "dot", + "vdot", + "cholesky", + "triangular_solve", + "mm", + "_unsafe_view", + "mv", + "outer", + "bmm", + "diagonal", + "alias", + "atan", + "ldexp", + "linear", + "log", + "log10", + "log1p", + "log2", + "logaddexp", + "logsumexp", + "logcumsumexp", + "reciprocal", + "tan", + "pow", + "rsqrt", + "tanh", + "tanh_backward", + "asinh", + "acosh", + "atanh", + "take", + "fill_", + "exp", + "exp2", + "expm1", + "nonzero", + "mean", + "std_mean", + "var_mean", + "inverse", + "solve", + "linalg_cholesky", + "addcmul", + "addcdiv", + "matrix_exp", + "linalg_matrix_exp", + "_linalg_eigh", + "cholesky_solve", + "linalg_qr", + "_linalg_svd", + "_fft_c2c", + "_fft_r2c", + "linalg_solve", + "sqrt", + "stack", + "gather", + "index_select", + "index_add_", + "linalg_inv", + "linalg_inv_ex", + "baddbmm", + "addbmm", + "addmm", + "addmv", + "addr", + "linalg_householder_product", + "ormqr", + "reflection_pad1d", + "reflection_pad2d", + "reflection_pad3d", + "linalg_cholesky_ex", + "linalg_eig", + "diagonal_copy", + "diagonal_scatter", + "alias_copy", + "select_backward", + "diagonal_backward", + "slice_backward", + "reflection_pad1d_backward", + "reflection_pad2d_backward", + "reflection_pad3d_backward", + "_sparse_sparse_matmul", + "replication_pad1d", + "replication_pad2d", + "replication_pad3d", + "put", + "put_", + "_to_copy", + "replication_pad1d_backward", + "replication_pad2d_backward", + "replication_pad3d_backward", + "diag", + "masked_scatter", + "masked_select", + "index_add", + "index_fill", + "trace", + "polar", + "cumsum", + "rsub", + "eig", + "lerp", + "linalg_vector_norm", + "cumprod", + "prod", + "index_copy", + "lu", + "unfold", + "unfold_backward", + "index", + "masked_fill", + "masked_scatter_backward", + "linalg_cross", + "lu_unpack", + "renorm", + "_conj_physical", + "linalg_lu_factor_ex", + "scatter", + "scatter_add", + "sigmoid", + "sigmoid_backward", + "sparse_mask", + "trapezoid", + "cumulative_trapezoid", + "conj_physical_", + "_neg_view", + "_reshape_alias", + "_reshape_copy", + "narrow_copy", + "_linalg_det", + "lu_solve", + "linalg_solve_triangular", + "linalg_pinv", + "linalg_lstsq", + "unfold_copy", + "col2im", + "im2col", + "cholesky_inverse", + "to_sparse", + "sparse_sampled_addmm", + "linalg_lu", + "pixel_shuffle", + "pixel_unshuffle", + "channel_shuffle", + "linalg_lu_solve", + "_linalg_slogdet", + "_linalg_solve_ex", + "_unsafe_index", + "_unsafe_index_put", + "_unsafe_masked_index", + "_unsafe_masked_index_put_accumulate", +} + +GRADIENT_IMPLEMENTED_FOR_SPARSE_COMPLEX = { + "_to_dense", + "_coalesce", + "coalesce", + "values", + "_sparse_coo_tensor_with_dims_and_tensors", + "_sparse_addmm", +} + +GRADIENT_IMPLEMENTED_FOR_COMPLEX.update(GRADIENT_IMPLEMENTED_FOR_SPARSE_COMPLEX) + +# Some operators invalidate the grad_accumulator. Let's reset it. +RESET_GRAD_ACCUMULATOR = {"set_", "resize_"} + +# NOTE [ TensorImpl and Storage Pointer Sanity Checks ] +# +# We check the following properties: +# 1) A function should never change the input tensors' underlying c10::TensorImpl +# pointers or c10::Storage pointers, even if it modifies its input tensors (via +# inplace or out-variants) +# If the function does not modify its arguments, we also check the following properties +# pertaining to its output: +# 2) Its TensorImpl has use_count of 1 (or 2 if it has a PyObject) +# 3) If the function is a view function, it has the same StorageImpl as that of +# the input it is aliased with. Otherwise, its StorageImpl has use_count of 1 +# +# The following code templates implement the checks for this invariant: +SAVE_TENSOR_STORAGE = CodeTemplate( + """\ +auto ${tensor_name}_storage_saved = + ${tensor_name}.has_storage() ? ::std::optional(${tensor_name}.storage()) : ::std::nullopt; +""" +) + + +# If tensor_name == out_tensor_name, used to enforce (1), otherwise used for (2) +ENFORCE_SAME_TENSOR_STORAGE = CodeTemplate( + """\ +if (${tensor_name}_storage_saved.has_value() && + !at::impl::dispatch_mode_enabled() && + !at::impl::tensor_has_dispatch(${tensor_name}) && + !at::impl::tensor_has_dispatch(${out_tensor_name})) + TORCH_INTERNAL_ASSERT(${tensor_name}_storage_saved.value().is_alias_of(${out_tensor_name}.storage())); +""" +) + +SAVE_TENSORLIST_STORAGE = CodeTemplate( + """\ +std::vector<::std::optional> ${tensorlist_name}_storage_saved(${tensorlist_name}.size()); +for (const Tensor& tensor : ${tensorlist_name}) + ${tensorlist_name}_storage_saved.push_back( + tensor.has_storage() ? ::std::optional(tensor.storage()) : ::std::nullopt); +""" +) + +ENFORCE_SAME_TENSORLIST_STORAGE = CodeTemplate( + """\ +for (size_t i=0; i<${tensorlist_name}.size() && !at::impl::dispatch_mode_enabled(); i++) { + if (${tensorlist_name}_storage_saved[i].has_value() && !at::impl::tensorlist_has_dispatch(${tensorlist_name})) + TORCH_INTERNAL_ASSERT(${tensorlist_name}_storage_saved[i].value().is_alias_of(${tensorlist_name}[i].storage())); +} +""" +) + +SAVE_OPTIONALTENSORLIST_STORAGE = CodeTemplate( + """\ +std::vector<::std::optional> ${tensorlist_name}_storage_saved(${tensorlist_name}.size()); +for (const ::std::optional& tensor : ${tensorlist_name}) + ${tensorlist_name}_storage_saved.push_back( + tensor.has_value() && tensor->has_storage() ? ::std::optional(tensor->storage()) : ::std::nullopt); +""" +) + +ENFORCE_SAME_OPTIONALTENSORLIST_STORAGE = CodeTemplate( + """\ +for (size_t i=0; i<${tensorlist_name}.size() && !at::impl::dispatch_mode_enabled(); i++) { + if (${tensorlist_name}_storage_saved[i].has_value() && !at::impl::tensorlist_has_dispatch(${tensorlist_name})) + TORCH_INTERNAL_ASSERT(${tensorlist_name}_storage_saved[i].value().is_alias_of( + static_cast<::std::optional>(${tensorlist_name}[i])->storage())); +} +""" +) + +SAVE_TENSOR_IMPL = CodeTemplate( + """\ +c10::intrusive_ptr ${tensor_name}_impl_saved; +if (${tensor_name}.defined()) ${tensor_name}_impl_saved = ${tensor_name}.getIntrusivePtr(); +""" +) + +ENFORCE_SAME_TENSOR_IMPL = CodeTemplate( + """\ +if (${tensor_name}_impl_saved && !at::impl::dispatch_mode_enabled() && !at::impl::tensor_has_dispatch(${tensor_name})) + TORCH_INTERNAL_ASSERT(${tensor_name}_impl_saved == ${tensor_name}.getIntrusivePtr()); +""" +) + +ENFORCE_TENSOR_IMPL_USE_COUNT = CodeTemplate( + """\ +if (!at::impl::dispatch_mode_enabled() && !at::impl::tensor_has_dispatch(${tensor_name})) + TORCH_INTERNAL_ASSERT(${tensor_name}.use_count() == expected_fresh_use_count(${tensor_name}), "function: ${fn_name}"); +""" +) + +ENFORCE_TENSOR_STORAGE_USE_COUNT_EQUALS_ONE = CodeTemplate( + """\ +if (${tensor_name}.has_storage() && !at::impl::dispatch_mode_enabled() && !at::impl::tensor_has_dispatch(${tensor_name})) { + TORCH_INTERNAL_ASSERT(${tensor_name}.storage().use_count() == 1, "function: ${fn_name}"); +} +""" +) + +SAVE_TENSORLIST_IMPL = CodeTemplate( + """\ +std::vector> ${tensorlist_name}_impl_saved(${tensorlist_name}.size()); +for (size_t i=0; i<${tensorlist_name}.size(); i++) + if (${tensorlist_name}[i].defined()) ${tensorlist_name}_impl_saved[i] = ${tensorlist_name}[i].getIntrusivePtr(); +""" +) + +ENFORCE_SAME_TENSORLIST_IMPL = CodeTemplate( + """\ +for (size_t i=0; i<${tensorlist_name}.size() && !at::impl::dispatch_mode_enabled(); i++) { + if (${tensorlist_name}_impl_saved[i] && !at::impl::tensorlist_has_dispatch(${tensorlist_name})) + TORCH_INTERNAL_ASSERT(${tensorlist_name}_impl_saved[i] == ${tensorlist_name}[i].getIntrusivePtr()); +} +""" +) + +SAVE_OPTIONALTENSORLIST_IMPL = CodeTemplate( + """\ +std::vector> ${tensorlist_name}_impl_saved(${tensorlist_name}.size()); +for (size_t i=0; i<${tensorlist_name}.size(); i++) { + ::std::optional t = ${tensorlist_name}[i]; + if (t.has_value() && t->defined()) ${tensorlist_name}_impl_saved[i] = t->getIntrusivePtr(); +} +""" +) + +ENFORCE_SAME_OPTIONALTENSORLIST_IMPL = CodeTemplate( + """\ +for (size_t i=0; i<${tensorlist_name}.size() && !at::impl::dispatch_mode_enabled(); i++) { + if (${tensorlist_name}_impl_saved[i]) + TORCH_INTERNAL_ASSERT( + ${tensorlist_name}_impl_saved[i] == static_cast<::std::optional>(${tensorlist_name}[i])->getIntrusivePtr()); +} +""" +) + +# The following list contains functions that we don't enforce the invariant on. +DONT_ENFORCE_SAME_TENSOR_IMPL_OR_STORAGE = { + # These functions are expected to change impl or storage of input tensors + "set_", + "_cudnn_rnn_flatten_weight", + "_unsafe_masked_index", + "_unsafe_masked_index_put_accumulate", +} +DONT_ENFORCE_TENSOR_IMPL_USE_COUNT = { + # These non-inplace, non-out functions return tensors with use_count > 1 + # Therefore, they MAY (but not necessarily) return one of its inputs as-is + # See https://github.com/pytorch/pytorch/issues/60426 for more information + "_embedding_bag", + "_embedding_bag_forward_only", + "q_per_channel_scales", + "q_per_channel_zero_points", + "lu_unpack", + "_cudnn_rnn_backward", + # The below failed StorageImpl use_count check but we skip tensor_impl check + # just in case + "_cudnn_rnn", + "dequantize_self", + # lift() should never actually be called with a requires_grad=True tensor, + "lift", + "lift_fresh", + "lift_fresh_copy", + # Nested Tensors related functions + # _nested_tensor_size() should never actually be called with requires_grad=True tensor + "_nested_tensor_size", + "_nested_tensor_strides", + "_nested_tensor_storage_offsets", +} + +DONT_ENFORCE_STORAGE_IMPL_USE_COUNT = { + # These non-view functions return tensors with storage use_count != 1 + "_slow_conv2d_forward", + "slow_conv3d_forward", + "channel_shuffle", + # If an input is returned as-is in output, we cannot guarantee its storage_impl + # use count to be 1 either. + *DONT_ENFORCE_TENSOR_IMPL_USE_COUNT, +} +# END CHECKS FOR [ TensorImpl and Storage Pointer Sanity Checks ] + +DECLARE_GRAD_FN = CodeTemplate( + """\ +std::shared_ptr<${op}> grad_fn; +""" +) + +DECLARE_VECTOR_OF_GRAD_FN = CodeTemplate( + """\ +std::vector> grad_fns; +""" +) + +SETUP_ANY_REQUIRES_GRAD = CodeTemplate( + """\ +[[maybe_unused]] auto _any_requires_grad = compute_requires_grad( ${args_with_derivatives} ); +${extra_differentiability_conditions} +""" +) + +SETUP_DERIVATIVE = CodeTemplate( + """\ +if (_any_requires_grad) { + ${setup} +} +""" +) + +SETUP_NONE_REQUIRES_GRAD = CodeTemplate( + """\ +if (compute_requires_grad( ${args_to_check} )) { + throw_error_out_requires_grad("${base_name}"); +} +""" +) + +ASSIGN_GRAD_FN = CodeTemplate( + """\ +grad_fn = std::shared_ptr<${op}>(new ${op}(${op_ctor}), deleteNode); +grad_fn->set_next_edges(collect_next_edges( ${args_with_derivatives} )); +""" +) + +# note(crcrpar): `compute_requires_grad` in the template below is supplied with arguments indexed with `i` +# while the `SETUP_ANY_REQUIRES_GRAD` above takes whole tensors and scalars. +ASSIGN_VECTOR_OF_GRAD_FN = CodeTemplate( + """\ +for (const auto& i : c10::irange( ${irange} )) { + const auto ith_requires_grad = compute_requires_grad(${args_with_derivatives}); + check_inplace(self[i], ith_requires_grad); + grad_fns.push_back([&]() -> std::shared_ptr<${op}> { + if (!ith_requires_grad) { + return nullptr; + } else { + auto grad_fn = std::shared_ptr<${op}>(new ${op}(${op_ctor}), deleteNode); + grad_fn->set_next_edges(collect_next_edges( ${args_with_derivatives} )); + return grad_fn; + } + }()); +} +""" +) + +CALL_REDISPATCH = CodeTemplate( + """\ +at::redispatch::${api_name}(${unpacked_args})""" +) +# If the non-variable operation has return values, we use the `tmp` variable to hold the +# values temporarily and pass the values to the return variables outside of the +# `at::AutoDispatchBelowAutograd` guard block. +DISPATCH_TO_NON_VAR_TYPE_WITH_TMP_RETURN_VALUES_JVP_DECOMP = CodeTemplate( + """\ +auto ${tmp_var} = ([&]() { + if (${any_has_forward_grad}) { + static c10::OperatorName full_name("aten::${op_name}", "${op_overload}"); + static ::std::optional opt_op = c10::Dispatcher::singleton().findSchema(full_name); + return impl::run_jit_decomposition_with_args_for_jvp<${return_types}>("${op_name}", *opt_op, ks, ${arg_names}); + } else { + ${guard} + return ${base_type_call}; + } +})(); +""" +) + +DISPATCH_TO_NON_VAR_TYPE_WITH_TMP_RETURN_VALUES = CodeTemplate( + """\ +auto ${tmp_var} = ([&]() { + ${guard} + return ${base_type_call}; +})(); +""" +) + +DISPATCH_TO_NON_VAR_TYPE_WITHOUT_RETURN_VALUES = CodeTemplate( + """\ +{ + ${guard} + ${base_type_call}; +} +""" +) + +SET_HISTORY = CodeTemplate( + """\ +if (grad_fn) { + ${fn}_history(${differentiable_outputs}, grad_fn); +} +""" +) + +LOOP_OVER_VECTOR_OF_GRAD_FNS = CodeTemplate( + """\ +if (!grad_fns.empty()) { + ${preamble} + for (const auto& i : c10::irange(grad_fns.size())) { + auto grad_fn = grad_fns[i]; + if (grad_fn != nullptr) { + ${statements} + } + } +} +""" +) + +CONDITIONAL = CodeTemplate( + """\ +if (${cond}) { + ${statements} +} +""" +) + +RUN_ONLY_IN_DEBUG_MODE = CodeTemplate( + """\ +#ifndef NDEBUG +${statements} +#endif +""" +) + +FW_DERIVATIVE_CHECK_TEMPLATE = CodeTemplate( + """\ +isFwGradDefined(${req_inp})\ +""" +) +FW_DERIVATIVE_SIZE_CHECK_TEMPLATE = CodeTemplate( + """\ +TORCH_CHECK( + self.size() == ${inp_name}.size(), + "Tensor lists must have the same number of tensors, got ", + self.size(), + " and ", + ${inp_name}.size()); +""" +) + +FW_DERIVATIVE_TENSORLIST_CHECK_TEMPLATE = CodeTemplate( + """\ +isFwGradDefinedTensorList(${req_inp})\ +""" +) + +FW_DERIVATIVE_DEFINED_GRAD_TEMPLATE = CodeTemplate( + """\ +auto ${inp_name}_t_raw = toNonOptFwGrad(${inp}); +auto ${inp_name}_tensor = toNonOptTensor(${inp}); +auto ${inp_name}_t = (${inp_name}_t_raw.defined() || !${inp_name}_tensor.defined()) + ? ${inp_name}_t_raw : at::${zeros_fn}(${inp_name}_tensor.sym_sizes(), ${inp_name}_tensor.options()); +""" +) + +FW_DERIVATIVE_UPDATE_WRAPPED_NUM_TEMPLATE = CodeTemplate( + """\ +update_wrapped_number(${inp_name}_tensor, ${inp_name}_t); +""" +) + +FW_DERIVATIVE_DEFINED_PRIMAL_TEMPLATE = CodeTemplate( + """\ +auto ${inp_name}_p = toNonOptPrimal(${inp}); +""" +) + +FW_DERIVATIVE_SETTER_TENSOR = CodeTemplate( + """\ +if (${out_arg}_new_fw_grad_opt.has_value() && ${out_arg}_new_fw_grad_opt.value().defined() && ${out_arg}.defined()) { + // The hardcoded 0 here will need to be updated once we support multiple levels. + ${out_arg}._set_fw_grad(${out_arg}_new_fw_grad_opt.value(), /* level */ 0, /* is_inplace_op */ ${is_inplace}); +} +""" +) + +FW_DERIVATIVE_SETTER_TENSOR_FOREACH = CodeTemplate( + """\ +for (const auto& i : c10::irange(${out_arg}_new_fw_grad_opts.size())) { + auto& ${out_arg}_new_fw_grad_opt = ${out_arg}_new_fw_grad_opts[i]; + if (${out_arg}_new_fw_grad_opt.has_value() && ${out_arg}_new_fw_grad_opt.value().defined() && ${out_arg}[i].defined()) { + // The hardcoded 0 here will need to be updated once we support multiple levels. + ${out_arg}[i]._set_fw_grad(${out_arg}_new_fw_grad_opt.value(), /* level */ 0, /* is_inplace_op */ ${is_inplace}); + } +} +""" +) + +FW_DERIVATIVE_SETTER_MULTI_OUTPUT = CodeTemplate( + """\ +if (${all_res}_new_fw_grad_opt.has_value() && std::get<${idx}>(${all_res}_new_fw_grad_opt.value()).defined() + && ${out_arg}.defined()) { + ${out_arg}._set_fw_grad(std::get<${idx}>(${all_res}_new_fw_grad_opt.value()), /* level */ 0, /* is_inplace_op */ false); +} +""" +) + +FW_DERIVATIVE_SETTER_TENSOR_LIST = CodeTemplate( + """\ +if (${out_arg}_new_fw_grad_opt.has_value()) { + auto ${out_arg}_new_fw_grad = ${out_arg}_new_fw_grad_opt.value(); + TORCH_INTERNAL_ASSERT(${out_arg}.size() == ${out_arg}_new_fw_grad.size()); + for (const auto i : c10::irange(${out_arg}.size())) { + if (${out_arg}_new_fw_grad[i].defined() && ${out_arg}[i].defined()) { + // The hardcoded 0 here will need to be updated once we support multiple levels. + ${out_arg}[i]._set_fw_grad(${out_arg}_new_fw_grad[i], /* level */ 0, /* is_inplace_op */ ${is_inplace}); + } + } +} +""" +) + +FW_DERIVATIVE_TEMPLATE = CodeTemplate( + """\ +${fw_grad_opt_definition} +if (${requires_fw_grad}) { + ${unpacked_arguments} + ${out_arg}_new_fw_grad_opt = ${formula}; +} +""" +) + +FW_DERIVATIVE_FOREACH_TEMPLATE = CodeTemplate( + """\ +${fw_grad_opt_definition} +for (const auto& i : c10::irange(${vector_of_optional_tensor}.size())) { + if (${any_has_forward_grad_for_current_index}) { + ${unpacked_arguments} + ${vector_of_optional_tensor}[i] = ${formula}; + } +} +""" +) + +FW_DERIVATIVE_FORBID_TEMPLATE = CodeTemplate( + """\ +TORCH_CHECK_NOT_IMPLEMENTED(!(${cond}), "Trying to use forward AD with ${name} that does not support it ${msg}"); +""" +) + +FW_DERIVATIVE_FORBID_LIST_TEMPLATE = CodeTemplate( + """\ +for (const auto& _t: ${arg}) { + TORCH_CHECK_NOT_IMPLEMENTED(!(${cond}), "Trying to use forward AD with ${name} that does not support it ${msg}"); +} +""" +) + + +def gen_variable_type( + out: str, + native_yaml_path: str, + tags_yaml_path: str, + fns_with_diff_infos: list[NativeFunctionWithDifferentiabilityInfo], + template_path: str, + used_keys: set[str], +) -> None: + """VariableType.h and VariableType.cpp body + + This is the at::Type subclass for differentiable tensors. The + implementation of each function dispatches to the base tensor type to + compute the output. The grad_fn is attached to differentiable functions. + """ + fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + fm.write( + "VariableType.h", + lambda: { + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/VariableType.h" + }, + ) + + # helper that generates a TORCH_LIBRARY_IMPL macro for each + # dispatch key that appears in derivatives.yaml + def wrapper_registrations(used_keys: set[str]) -> str: + library_impl_macro_list: list[str] = [] + for key in sorted(used_keys): + dispatch_key = key + if key == "Default": + dispatch_key = "Autograd" + library_impl_macro = ( + f"TORCH_LIBRARY_IMPL(aten, {dispatch_key}, m) " + + "{\n" + + "${" + + f"wrapper_registrations_{key}" + + "}\n}" + ) + library_impl_macro_list += [library_impl_macro] + return "\n\n".join(library_impl_macro_list) + + # Generate a new template from VariableType.cpp which replaces ${wrapper_registrations} + # with per key TORCH_LIBRARY_IMPL macros for each key that appears in derivatives.yaml + fm1 = FileManager( + install_dir=out + "/templates", template_dir=template_path, dry_run=False + ) + fm1.write( + "VariableType.cpp", + lambda: { + "type_derived_method_definitions": "\n\n".join( + [ + "${" + f"type_derived_method_definitions_{key}" + "}" + for key in sorted(used_keys) + ] + ), + "wrapper_registrations": wrapper_registrations(used_keys), + }, + ) + + # Generate final VariableType_*.cpp files from the generated template + fm2 = FileManager(install_dir=out, template_dir=out + "/templates", dry_run=False) + + sharded_keys = set( + [f"type_derived_method_definitions_{key}" for key in sorted(used_keys)] + + [f"wrapper_registrations_{key}" for key in sorted(used_keys)] + ) + # NOTE: see Note [Sharded File] at the top of the VariableType.cpp + # template regarding sharding of the generated files. + fm2.write_sharded( + "VariableType.cpp", + [fn for fn in fns_with_diff_infos if use_derived(fn)], + key_fn=lambda fn: cpp.name(fn.func.func), + base_env={ + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/VariableType.cpp", + }, + env_callable=gen_variable_type_func, + num_shards=5, + sharded_keys=sharded_keys, + ) + + +@with_native_function_and +def gen_wrapper_registration(f: NativeFunction, key: str = "Default") -> str: + return WRAPPER_REGISTRATION.substitute( + unqual_operator_name_with_overload=f.func.name, + type_wrapper_name=type_wrapper_name(f, key), + class_type="VariableType", + ) + + +def gen_variable_type_func( + fn: NativeFunctionWithDifferentiabilityInfo, +) -> dict[str, list[str]]: + f = fn.func + result = {} + with native_function_manager(f): + name = cpp.name(f.func) + formals = gen_formals(f) + + if ( + fn.info is None + and str(f.func.name.name) not in RESET_GRAD_ACCUMULATOR + and get_base_name(f) not in DONT_REQUIRE_DERIVATIVE + and len(gen_differentiable_outputs(fn)) > 0 + and cpp.name(f.func) not in DONT_ENFORCE_SAME_TENSOR_IMPL_OR_STORAGE + and type_wrapper_name(f) not in DONT_ENFORCE_STORAGE_IMPL_USE_COUNT + and type_wrapper_name(f) not in DONT_ENFORCE_TENSOR_IMPL_USE_COUNT + ): + # NOTE: [ Registering AutogradNotImplemented boxed kernel ] + # + # When there is no derivatives.yaml entry, we register a generic boxed + # NotImplemented kernel to set grad_fn to be NotImplemented, so that forward + # proceeds as usual but an error is properly produced on backward. + # TODO: it would be nice to not have these special cases + # + # There are several cases where still let codegen handle it: + # 1) ops that need to reset grad accumulator (we let codegen handle this case + # because) the list is (currently) only accessible in Python. + # 2) User explicitly specifies DONT_REQUIRE_DERIVATIVE. This basically makes + # autograd a fallthrough with NDEBUG checks. This can be useful for when all + # outputs are integral. + # 3) When there are no differentiable outputs. This is similar to (2). + # 4) There are certain ops where we skip certain NDEBUG checks. this is similar + # to (1). + type_definition = "" + wrapper_registration = AUTOGRAD_NOT_IMPLEMENTED_REGISTRATION.substitute( + unqual_operator_name_with_overload=f.func.name + ) + result["type_derived_method_definitions_Default"] = [type_definition] + result["wrapper_registrations_Default"] = [wrapper_registration] + else: + if not fn.info: + key = "Default" + type_definition = METHOD_DEFINITION.substitute( + return_type=cpp.returns_type( + f.func.returns, symint=True + ).cpp_type(), + type_wrapper_name=type_wrapper_name(f, key), + type_definition_body=emit_body(fn, key), + formals=formals, + ) + wrapper_registration = gen_wrapper_registration(f, key) + result[f"type_derived_method_definitions_{key}"] = [type_definition] + result[f"wrapper_registrations_{key}"] = [wrapper_registration] + else: + for key in fn.info: + type_definition = METHOD_DEFINITION.substitute( + return_type=cpp.returns_type( + f.func.returns, symint=True + ).cpp_type(), + type_wrapper_name=type_wrapper_name(f, key), + type_definition_body=emit_body(fn, key), + formals=formals, + ) + wrapper_registration = gen_wrapper_registration(f, key) + result[f"type_derived_method_definitions_{key}"] = [type_definition] + result[f"wrapper_registrations_{key}"] = [wrapper_registration] + # See Note [Manual Backend kernels] + if (name in MANUAL_BACKEND) != f.manual_kernel_registration: + raise AssertionError( + f"(name in MANUAL_BACKEND) != f.manual_kernel_registration: {name in MANUAL_BACKEND} != {f.manual_kernel_registration}" + ) + # If you want to register a kernel to Autograd, you must make the op abstract. + # In other words, this op must have dispatch section in native_functions.yaml. + if name in MANUAL_AUTOGRAD_AND_TRACER or ( + fn.info and any(info.has_derivatives for info in fn.info.values()) + ): + if not f.is_abstract: + raise AssertionError( + f"There's a formula for {name}(or its functional variant) in derivatives.yaml. " + f"It's required to add a dispatch section for it with explicit supported backends e.g CPU/CUDA " + f"or CompositeExplicitAutograd in native_functions.yaml. Please see " + f"https://github.com/pytorch/pytorch/tree/master/aten/src/ATen/native#choosing-the-right-dispatch-keyword " + f"for instructions to choose the right dispatch keyword." + ) + + return result + + +_foreach_ops_without_differentiability_info = { + # No reference backward available due to the lack of `{maximum, minimum}(tensor, scalar)`. + ("_foreach_maximum", "Scalar"), + ("_foreach_maximum", "ScalarList"), + ("_foreach_minimum", "Scalar"), + ("_foreach_minimum", "ScalarList"), + # No reference backward available as addcdiv/addcmul don't support Tensor as scaling factor. + ("_foreach_addcdiv", "Tensor"), + ("_foreach_addcmul", "Tensor"), + ("_foreach_copy", ""), +} + +_foreach_ops_with_different_arity = { + # These ops lack `alpha` of scaling factor to applied to the right hand side argument. + ("_foreach_add", "Scalar"), + ("_foreach_add", "ScalarList"), + ("_foreach_sub", "Scalar"), + ("_foreach_sub", "ScalarList"), +} + + +@with_native_function_with_differentiability_info_and_key +def emit_body( + fn: NativeFunctionWithDifferentiabilityInfo, key: str = "Default" +) -> list[str]: + if dispatch_strategy(fn) != "use_derived": + raise AssertionError( + f"dispatch_strategy(fn) is {dispatch_strategy(fn)}, expected 'use_derived'" + ) + f = fn.func + info = fn.info[key] if fn.info else None + fw_derivatives = fn.fw_derivatives.get(key, []) if fn.fw_derivatives else [] + + name = cpp.name(f.func) + inplace = f.func.kind() == SchemaKind.inplace + is_out_fn = f.func.kind() == SchemaKind.out + returns_void = len(f.func.returns) == 0 + base_name = get_base_name(f) + view_info = get_view_info(f) + + is_foreach = name.startswith("_foreach") + is_inplace_foreach = is_foreach and inplace + if is_inplace_foreach: + inplace_foreacharg2refarg: dict[Argument, Argument] = {} + refargname2inplace_foreacharg: dict[str, Argument] = {} + base_name_and_overload_name = (f.func.name.name.base, f.func.name.overload_name) + if info is None: + if ( + base_name_and_overload_name + not in _foreach_ops_without_differentiability_info + ): + raise AssertionError( + f"{'.'.join(base_name_and_overload_name)} should have a differentiability info" + ) + else: + if not ( + len(f.func.arguments.flat_non_out) + == len(info.func.func.arguments.flat_non_out) + ) and ( + base_name_and_overload_name not in _foreach_ops_with_different_arity + ): + raise AssertionError( + f"{'.'.join(base_name_and_overload_name)} has {len(f.func.arguments.flat_non_out)} args " + f"but the reference has {len(info.func.func.arguments.flat_non_out)}" + ) + for foreach_arg, ref_arg in zip( + f.func.arguments.flat_non_out, info.func.func.arguments.flat_non_out + ): + foreach_arg_type = foreach_arg.type + if isinstance(foreach_arg_type, ListType): + foreach_arg_type = foreach_arg_type.elem + if foreach_arg_type != ref_arg.type: + raise AssertionError( + f"foreach_arg_type ({foreach_arg_type}) != ref_arg.type ({ref_arg.type})" + ) + inplace_foreacharg2refarg[foreach_arg] = ref_arg + refargname2inplace_foreacharg[ref_arg.name] = foreach_arg + + def gen_differentiable_input( + arg: Argument | SelfArgument | TensorOptionsArguments, + ) -> DifferentiableInput | None: + if isinstance(arg, TensorOptionsArguments): + return None + a: Argument = arg.argument if isinstance(arg, SelfArgument) else arg + + # TODO: `cpp_type` is only to keep it byte-for-byte compatible with the old codegen, should remove. + # NB: This is not a clone of cpp.argument() - TensorOptionsArguments / faithful / binds are + # not handled properly as they are irrelevant for this codegen. + cpp_type = cpp.argument_type(a, binds=a.name, symint=True).cpp_type() + + if not is_differentiable(a.name, a.type, info): + return None + return DifferentiableInput( + name=a.name, + type=a.type, + cpp_type=cpp_type, + ) + + @with_native_function + def gen_differentiable_inputs(f: NativeFunction) -> list[DifferentiableInput]: + arguments = list(f.func.arguments.non_out) + if is_inplace_foreach and info is not None: + for i, arg in enumerate(f.func.arguments.flat_non_out): + if arg in inplace_foreacharg2refarg: + # note(crcrpar): From what I understand, what matters is only the name. + # Thus originally I only replace argument only when the names are different. + # TODO(crcrpar): Make it simpler. + mapped_arg = inplace_foreacharg2refarg[arg] + arguments[i] = Argument( + mapped_arg.name, + mapped_arg.type, + mapped_arg.default, + mapped_arg.annotation, + ) + return list(mapMaybe(gen_differentiable_input, arguments)) + + def find_args_with_derivatives( + differentiable_inputs: list[DifferentiableInput], + ) -> list[DifferentiableInput]: + """Find arguments that have derivative definitions""" + if info is None or not info.has_derivatives: + return differentiable_inputs + names = {name for d in info.derivatives for name in d.var_names} + differentiable = [arg for arg in differentiable_inputs if arg.name in names] + if len(differentiable) != len(names): + missing = names - {arg.name for arg in differentiable} + raise RuntimeError( + f"Missing arguments for derivatives: {missing} in {info.name}" + ) + return differentiable + + differentiable_inputs = gen_differentiable_inputs(f) + args_with_derivatives = find_args_with_derivatives(differentiable_inputs) + differentiable_outputs = gen_differentiable_outputs(fn, key) + + undifferentiable = (base_name in DONT_REQUIRE_DERIVATIVE) or ( + name in DONT_REQUIRE_DERIVATIVE + ) + + requires_derivative = ( + (not undifferentiable) + and (len(differentiable_inputs) > 0) + and ( + (len(differentiable_outputs) > 0) + # note(crcrpar): In-place foreach functions are a void function. + or is_inplace_foreach + ) + ) + + if ( + info is not None + and info.has_derivatives + and not requires_derivative + # out= ops are allowed to have zero returns which cause requires_derivative to be False + # we shouldn't error out though (out= ops for autograd just redispatch) + and len(f.func.returns) > 0 + ): + raise RuntimeError( + f"ERROR: derivative ignored for {name} -- specified an autograd function without derivative" + ) + + # note(crcrpar): In-place foreach functions do not support forward AD + if requires_derivative and len(fw_derivatives) > 0 and not is_inplace_foreach: + num_fw_derivative_var_names = sum( + len(derivative.var_names) for derivative in fw_derivatives + ) + if num_fw_derivative_var_names != len(differentiable_outputs): + raise AssertionError( + f"Expected the number of forward derivatives implemented ({num_fw_derivative_var_names}) to match the " + f"number of differentiable outputs ({len(differentiable_outputs)}). NB: This only applies when at least " + "one forward derivative is implemented. Not implementing any forward " + "derivatives is also okay, and we would require inputs to the op to " + "not have associated tangents in that case." + ) + + try_jit_decomposition = ( + requires_derivative + and len(fw_derivatives) == 0 + and (not modifies_arguments(f)) + and (not returns_void) + ) + + def emit_save_inputs() -> list[str]: + setup: list[str] = [] + if info is None or not info.has_derivatives: + return setup + + has_tensorlist_arg = any( + is_tensor_list_type(arg.type) for arg in args_with_derivatives + ) + + # We don't want to save tensors if we know that they will never be used + # when computing the derivative, so we add guards to those statements + def guard_for(arg: SavedAttribute) -> str | None: + if info is None: + raise AssertionError("info is None in guard_for") + + # It's hard to determine the edge offset if we have TensorLists + # NOTE(crcrpar): in-place foreach functions' arguments include tensorlist + # but their derivatives don't use it, so let them bypass this check. + if has_tensorlist_arg and (not is_inplace_foreach): + return None + + # Empirical evaluation of the cases where we insert those guards in + # backward show that they are somewhat useless. E.g. there's no need + # to guard on some values captured from forward, because they had to + # require_grad if the backward function even gets executed. I don't + # have any good ideas for detecting those cases, so I simply disabled the + # checks. + if "backward" in info.name: + return None + + # If there's a single derivative we could compute, we already have + # a requires_grad check that is sufficient + if len(args_with_derivatives) <= 1: + return None + + # We really only care about trimming down the amount of tensors we save + if arg.nctype.type != BaseCType(tensorT): + return None + + # We want to emit simple guards, so we only allow that if checking one + # input is enough to determine whether we need that value + used_in = [d for d in info.derivatives if arg in d.saved_inputs] + if len(used_in) == 0: + raise AssertionError(f"used_in is empty for arg {arg.nctype.name}") + if len(used_in) != 1: + return None + derivative = used_in[0] + + # Case with multioutput formulas + # TODO: process all derivative formulas!!! + if len(derivative.var_names) != 1: + wrap_opt_if_start = derivative.formula.find( + f"wrap_opt_if({arg.nctype.name}" + ) + if wrap_opt_if_start == -1: + return None + + wrap_opt_if_match = re.match( + rf"wrap_opt_if\({arg.nctype.name},(.*?)\)", + derivative.formula[wrap_opt_if_start:], + ) + if wrap_opt_if_match is None: + raise AssertionError( + f"wrap_opt_if_match is None for {arg.nctype.name} in {derivative.formula}" + ) + + # Condition is between 'wrap_opt_if(var_name,' and ')'. + condition_slice = slice(len(rf"wrap_opt_if\({arg.nctype.name},"), -1) + wrap_opt_if_condition = wrap_opt_if_match.group(0)[ + condition_slice + ].strip() + # replace 'grad_input_mask[num]' with 'grad_fn->should_compute_output(num)' + wrap_opt_if_condition = re.sub( + r"grad_input_mask\[(\d+)\]", + r"grad_fn->should_compute_output(\1)", + wrap_opt_if_condition, + ) + return f"{wrap_opt_if_condition}" + + # Figure out the offset of the edge that uses this variable + derivative_var_name = derivative.var_names[0] + for edge_off, a in enumerate(args_with_derivatives): + if a.name == derivative_var_name: + break + else: + raise AssertionError + return f"grad_fn->should_compute_output({edge_off})" + + if is_inplace_foreach: + save_input_stmts = save_variables(info.all_saved_inputs, False, guard_for) + if save_input_stmts: + setup.append( + LOOP_OVER_VECTOR_OF_GRAD_FNS.substitute( + preamble="", statements=save_input_stmts + ) + ) + else: + setup.extend(save_variables(info.all_saved_inputs, False, guard_for)) + for arg in args_with_derivatives: + if is_tensor_list_type(arg.type): + setup.append(f"grad_fn->{arg.name}_size_ = {arg.name}.size();") + return setup + + def setup_derivative(differentiable_inputs: list[DifferentiableInput]) -> list[str]: + body: list[str] = [] + if is_out_fn: + # For out functions, ensure that no input or output requires grad + body.append(DECLARE_GRAD_FN.substitute(op="Node")) + body.append( + SETUP_NONE_REQUIRES_GRAD.substitute( + base_name=base_name, + args_to_check=[arg.name for arg in differentiable_inputs], + ) + ) + body.append( + SETUP_NONE_REQUIRES_GRAD.substitute( + base_name=base_name, + args_to_check=[arg.name for arg in differentiable_outputs], + ) + ) + return body + + op = info.op if info is not None and info.has_derivatives else "NotImplemented" + setup = [] + if not is_inplace_foreach: + setup.extend( + ASSIGN_GRAD_FN.substitute( + op=op, + op_ctor="" + if info is not None and info.has_derivatives + else f'"{cpp.name(f.func)}"', + args_with_derivatives=[arg.name for arg in args_with_derivatives], + ).split("\n") + ) + else: + # note(crcrpar): Assuming in-place foreach function's self_arg is always TensorList. + list_like_arg = "self" + args = [arg.name for arg in args_with_derivatives] + for i, arg in enumerate(args): + if is_inplace_foreach and info is not None: + if arg in refargname2inplace_foreacharg: + foreach_arg = refargname2inplace_foreacharg[arg] + args[i] = foreach_arg.name + ( + "[i]" if isinstance(foreach_arg.type, ListType) else "" + ) + else: + if arg == list_like_arg: + args[i] = arg + "[i]" + setup.extend( + ASSIGN_VECTOR_OF_GRAD_FN.substitute( + op=op, + op_ctor="" + if info is not None and info.has_derivatives + else f'"{cpp.name(f.func)}"', + args_with_derivatives=args, + irange=f"{list_like_arg}.size()", + ).split("\n") + ) + setup.extend(emit_save_inputs()) + + body.extend( + emit_check_no_requires_grad(differentiable_inputs, args_with_derivatives) + ) + declare_grad_fn_template = ( + DECLARE_GRAD_FN if not is_inplace_foreach else DECLARE_VECTOR_OF_GRAD_FN + ) + body.append(declare_grad_fn_template.substitute(op=op)) + body.append(SETUP_DERIVATIVE.substitute(setup=setup)) + return body + + def emit_check_if_in_complex_autograd_allowlist() -> list[str]: + body: list[str] = [] + if base_name in GRADIENT_IMPLEMENTED_FOR_COMPLEX: + return body + for arg in differentiable_outputs: + name = arg.name + # TODO: should be `arg.type.is_tensor_like()`? + if arg.cpp_type == "at::Tensor" or arg.cpp_type in TENSOR_LIST_LIKE_CTYPES: + body.append(f'throw_error_for_complex_autograd({name}, "{base_name}");') + return body + + def emit_check_no_requires_grad( + tensor_args: list[DifferentiableInput], + args_with_derivatives: list[DifferentiableInput], + ) -> list[str]: + """Checks that arguments without derivatives don't require grad""" + body: list[str] = [] + for arg in tensor_args: + if arg in args_with_derivatives: + continue + arg_name = arg.name + if info and arg_name in info.non_differentiable_arg_names: + continue + if arg_name == "output": + # Double-backwards definitions sometimes take in 'input' and + # 'output', but only define the derivative for input. + continue + body.append(f'check_no_requires_grad({arg_name}, "{arg_name}", "{name}");') + return body + + def emit_original_self_definition() -> list[str]: + body: list[str] = [] + if inplace: + if is_inplace_foreach: + body.append( + "std::vector<::std::optional> original_selfs(self.size());" + ) + else: + body.append("::std::optional original_self;") + + all_forward_grad_cond = [] + for derivative in fw_derivatives: + if derivative.required_original_self_value: + all_forward_grad_cond.append( + get_any_has_forward_grad_name(derivative.var_names) + ) + + if all_forward_grad_cond: + if not is_inplace_foreach: + body.append(f"if ({' || '.join(all_forward_grad_cond)}) {{") + body.append(" original_self = self.clone();") + body.append("}") + else: + current_all_forward_grad_cond = [ + f"{cond}[i]" for cond in all_forward_grad_cond + ] + body.append("for (const auto& i : c10::irange(self.size())) {") + body.append( + f" if ({' || '.join(current_all_forward_grad_cond)}) {{" + ) + body.append(" original_selfs[i] = self[i].clone();") + body.append(" }") + body.append("}") + + return body + + def save_variables( + saved_variables: Sequence[SavedAttribute], + is_output: bool, + guard_for: Callable[[SavedAttribute], str | None] = lambda name: None, + ) -> Sequence[str]: + # assign the saved variables to the generated grad_fn + stmts: list[str] = [] + for arg in sorted(saved_variables, key=lambda sa: str(sa.nctype.name)): + name = ( + arg.nctype.name.name + if isinstance(arg.nctype.name, SpecialArgName) + else arg.nctype.name + ) + foreacharg: Argument | None = None + is_foreacharg_list_type: bool = False + type = arg.nctype.type + expr = arg.expr + stmts_prepend = None + if is_inplace_foreach and info is not None: + # todo(crcrpar): See if we can add some check e.g. `assert foreacharg is not None`. + # for now the example assert would fail. + name_to_query = name.split("_scalar_type")[0] + if name_to_query in refargname2inplace_foreacharg: + foreacharg = refargname2inplace_foreacharg[name_to_query] + is_foreacharg_list_type = isinstance(foreacharg.type, ListType) + if foreacharg is not None: + name_in_expr = ( + f"{foreacharg.name}{'[i]' if is_foreacharg_list_type else ''}" + ) + src_name = name + if "_scalar_type" in src_name: + split_src_name = src_name.split("_scalar_type") + if len(split_src_name) != 2: + raise AssertionError( + f"expected 2 parts after split, got {len(split_src_name)}: {split_src_name}" + ) + src_name = split_src_name[0] + expr = expr.replace(src_name, name_in_expr) + if ( + type == BaseCType(tensorT) + or type == OptionalCType(BaseCType(tensorT)) + or type == MutRefCType(OptionalCType(BaseCType(tensorT))) + or (is_output and type == BaseCType(scalarT)) + ): + # note(crcrpar): Here `expr` is generated from scratch, `arg.expr` is ignored. + var = name + name += "_" + if var == "self" and inplace: + original_self_var = ( + "original_self" + if not is_inplace_foreach + else "original_selfs[i]" + ) + self_var = var if not is_inplace_foreach else var + "[i]" + stmts_prepend = f"if (!{original_self_var}.has_value()) {original_self_var} = {self_var}.clone()" + var = f"{original_self_var}.value()" + if is_output: + raise AssertionError( + "is_output should be False when var == 'self' and inplace" + ) + if inplace and is_output: + if name != "result_": + raise AssertionError( + f"expected name to be 'result_' for inplace output, got {name}" + ) + var = ( + "self[i]" + if is_inplace_foreach or is_foreacharg_list_type + else "self" + ) + is_inplace_view = f"{var}.is_view()" + expr = f"SavedVariable({var}, {str(is_output).lower()}, {is_inplace_view})" + else: + expr = f"SavedVariable({var}, {str(is_output).lower()})" + if foreacharg is not None and "original_selfs" not in expr: + # pyrefly: ignore [unbound-name] + expr = expr.replace(src_name, name_in_expr) + elif ( + type == BaseCType(tensorListT) + or type == ListCType(OptionalCType(BaseCType(tensorT))) + or type == BaseCType(iTensorListRefT) + or type == VectorCType(BaseCType(tensorT)) + ): + # See Note [nuanced return type of out-of-place foreach functions] + if type == VectorCType(BaseCType(tensorT)): + if not (is_foreach and is_output): + raise AssertionError( + f"VectorCType(BaseCType(tensorT)) requires is_foreach and is_output, " + f"got is_foreach={is_foreach}, is_output={is_output}" + ) + expr = f"make_saved_variable_list({name}, {str(is_foreach and is_output).lower()})" + name += "_" + elif type == BaseCType(intArrayRefT): + expr = expr + ".vec()" + elif type == BaseCType(symIntArrayRefT): + expr = expr + ".vec()" + elif type == BaseCType(stringT): + expr = f"std::string({expr})" + elif type == OptionalCType(BaseCType(stringT)): + expr = f"{expr}.has_value() ? ::std::optional(std::string({expr}.value())) : ::std::nullopt" + elif type == ArrayRefCType( + elem=BaseCType(type=BaseCppType(ns="at", name="Scalar")) + ): + expr = expr + ".vec()" + + guard = guard_for(arg) + if guard is None: + if stmts_prepend: + stmts.append(f"{stmts_prepend};") + stmts.append(f"grad_fn->{name} = {expr};") + else: + stmts.append(f"if ({guard}) {{") + if stmts_prepend: + stmts.append(f" {stmts_prepend};") + stmts.append(f" grad_fn->{name} = {expr};") + stmts.append("}") + return stmts + + # Generates a Dispatcher::redispatch() call into the dispatcher. We do this mainly for performance reasons: + # - Pre-compute the full DispatchKeySet. This saves the dispatcher from having to read from TLS. + # - redispatch() avoids a redundant call to RecordFunction, which was already called right before + # we entered this autograd kernel. + def emit_dispatch_call( + f: NativeFunction, input_base: str, unpacked_args: Sequence[str] + ) -> str: + """Dispatch call via function in a namespace or method on Tensor.""" + # code-generated autograd kernels plumb and recompute dispatch keys directly through the kernel for performance. + # Ops also always have a function variant of the redispatch API. + # See Note [Plumbing Keys Through The Dispatcher] for details. + dispatch_key_set = "ks & c10::after_autograd_keyset" + call = CALL_REDISPATCH.substitute( + api_name=cpp.name( + f.func, + faithful_name_for_out_overloads=True, + symint_overload=f.func.has_symint(), + ), + unpacked_args=[dispatch_key_set] + list(unpacked_args), + ) + return call + + def wrap_output( + f: NativeFunction, unpacked_bindings: list[Binding], var: str + ) -> str: + call = "" + rhs_value: str | None = None + if not any(r.type.is_tensor_like() for r in f.func.returns): + rhs_value = var + else: + rhs_value = f"std::move({var})" + if rhs_value is None: + raise AssertionError("rhs_value is None") + call += ASSIGN_RETURN_VALUE.substitute( + return_values=tie_return_values(f), rhs_value=rhs_value + ) + return call + + def check_tensorimpl_and_storage( + call: str, unpacked_bindings: list[Binding] + ) -> str: + # See NOTE [ TensorImpl and Storage Pointer Sanity Checks ] + stmts_before_call: list[str] = [] + stmts_after_call: list[str] = [] + + if cpp.name(f.func) in DONT_ENFORCE_SAME_TENSOR_IMPL_OR_STORAGE: + return call + + # Check properties of inputs (enforce (1)) + for unpacked_binding in unpacked_bindings: + arg = unpacked_binding.name + noref_cpp_type = unpacked_binding.nctype.type.remove_const_ref() + if noref_cpp_type == BaseCType(tensorListT) or noref_cpp_type == BaseCType( + iTensorListRefT + ): + stmts_before_call += [ + SAVE_TENSORLIST_STORAGE.substitute(tensorlist_name=arg), + SAVE_TENSORLIST_IMPL.substitute(tensorlist_name=arg), + ] + stmts_after_call += [ + ENFORCE_SAME_TENSORLIST_STORAGE.substitute(tensorlist_name=arg), + ENFORCE_SAME_TENSORLIST_IMPL.substitute(tensorlist_name=arg), + ] + elif noref_cpp_type == ListCType(OptionalCType(BaseCType(tensorT))): + stmts_before_call += [ + SAVE_OPTIONALTENSORLIST_STORAGE.substitute(tensorlist_name=arg), + SAVE_OPTIONALTENSORLIST_IMPL.substitute(tensorlist_name=arg), + ] + stmts_after_call += [ + ENFORCE_SAME_OPTIONALTENSORLIST_STORAGE.substitute( + tensorlist_name=arg + ), + ENFORCE_SAME_OPTIONALTENSORLIST_IMPL.substitute( + tensorlist_name=arg + ), + ] + elif noref_cpp_type == BaseCType(tensorT): + stmts_before_call += [ + SAVE_TENSOR_STORAGE.substitute(tensor_name=arg), + SAVE_TENSOR_IMPL.substitute(tensor_name=arg), + ] + stmts_after_call += [ + ENFORCE_SAME_TENSOR_STORAGE.substitute( + tensor_name=arg, out_tensor_name=arg + ), + ENFORCE_SAME_TENSOR_IMPL.substitute(tensor_name=arg), + ] + + if not ( + (stmts_before_call and stmts_after_call) + or (not stmts_before_call and not stmts_after_call) + ): + raise AssertionError( + "stmts_before_call and stmts_after_call must be both empty or both non-empty" + ) + + # Check properties of outputs (enforce (2), (3)) + if f.func.kind() not in (SchemaKind.inplace, SchemaKind.out): + base_name = f.func.name.name.base # TODO: should be str(f.func.name.name)? + aliased_arg_name = ALL_VIEW_FUNCTIONS.get(base_name, None) + if aliased_arg_name is not None: + aliased_arg_name = unpacked_name(aliased_arg_name) + for i, (ret, ret_name) in enumerate( + zip(f.func.returns, cpp.return_names(f)) + ): + noref_cpp_type = cpp.return_type(ret, symint=True).remove_const_ref() + if noref_cpp_type == BaseCType(tensorT): + if aliased_arg_name is not None: + if i != 0: + raise AssertionError( + f"Expect non-CompositeImplicitAutograd view function {base_name} " + f"to return single output, got index {i}" + ) + stmts_after_call += [ + ENFORCE_SAME_TENSOR_STORAGE.substitute( + tensor_name=aliased_arg_name, out_tensor_name=ret_name + ) + ] + else: + if ( + type_wrapper_name(f) + not in DONT_ENFORCE_STORAGE_IMPL_USE_COUNT + ): + stmts_after_call += [ + ENFORCE_TENSOR_STORAGE_USE_COUNT_EQUALS_ONE.substitute( + tensor_name=ret_name, fn_name=type_wrapper_name(f) + ) + ] + + if type_wrapper_name(f) not in DONT_ENFORCE_TENSOR_IMPL_USE_COUNT: + stmts_after_call += [ + ENFORCE_TENSOR_IMPL_USE_COUNT.substitute( + tensor_name=ret_name, fn_name=type_wrapper_name(f) + ) + ] + + # Currently we don't have any functions that return the following types, but + # we should update the checks once we do + elif noref_cpp_type == ListCType(OptionalCType(BaseCType(tensorT))): + raise AssertionError( + f"Please add use_count checks for {noref_cpp_type}" + ) + elif noref_cpp_type == BaseCType(tensorListT): + raise AssertionError( + f"Please add use_count checks for {noref_cpp_type}" + ) + + if stmts_before_call and stmts_after_call: + call = ( + RUN_ONLY_IN_DEBUG_MODE.substitute(statements=stmts_before_call) + + call + + RUN_ONLY_IN_DEBUG_MODE.substitute(statements=stmts_after_call) + ) + return call + + def emit_call( + f: NativeFunction, unpacked_bindings: list[Binding], try_jit_decomposition: bool + ) -> str: + # We only care about adding `at::AutoDispatchBelowAutograd` guard for non-variable dispatch + # (which corresponds to 'use_derived' strategy). The purpose of this guard is to make sure + # the baseType operations still dispatch to non-Variable type, even if the arguments passed + # in are now Variables. + # See NOTE [ Treating Variables as non-Variables in type dispatch ] for details. + unpacked_args = [b.name for b in unpacked_bindings] + base_type_call = emit_dispatch_call(f, "self_", unpacked_args) + + if get_view_info(f) is not None or modifies_arguments(f): + guard = "at::AutoDispatchBelowAutograd guard;" + else: + guard = "at::AutoDispatchBelowADInplaceOrView guard;" + + any_has_forward_grad = ( + get_any_has_fw_grad_cond(derivative=None) + if requires_derivative + else "false" + ) + return_types = ", ".join( + [cpp.return_type(a, symint=True).cpp_type() for a in f.func.returns] + ) + if len(f.func.returns) > 1: + return_types = f"std::tuple<{return_types}>" + + arg_names = [ + a.name + for a in cpp.arguments( + f.func.arguments, + faithful=True, + symint=True, + method=False, + cpp_no_default_args=set(), + ) + ] + + if not modifies_arguments(f) and not returns_void: + if try_jit_decomposition: + call = DISPATCH_TO_NON_VAR_TYPE_WITH_TMP_RETURN_VALUES_JVP_DECOMP.substitute( + base_type_call=base_type_call, + tmp_var=TMP_VAR, + guard=guard, + any_has_forward_grad=any_has_forward_grad, + op_name=cpp.name(f.func), + op_overload=f.func.name.overload_name, + return_types=return_types, + arg_names=arg_names, + ) + else: + call = DISPATCH_TO_NON_VAR_TYPE_WITH_TMP_RETURN_VALUES.substitute( + base_type_call=base_type_call, + tmp_var=TMP_VAR, + guard=guard, + ) + + call += wrap_output(f, unpacked_bindings, TMP_VAR) + else: + if try_jit_decomposition: + raise AssertionError( + "try_jit_decomposition should be False for functions with no return values or that modify arguments" + ) + call = DISPATCH_TO_NON_VAR_TYPE_WITHOUT_RETURN_VALUES.substitute( + base_type_call=base_type_call, guard=guard + ) + call = check_tensorimpl_and_storage(call, unpacked_bindings) + return call + + def emit_history() -> str: + fn = "rebase" if modifies_arguments(f) and view_info is None else "set" + output_names = [r.name for r in differentiable_outputs] + # TODO: flatten allocates a std::vector, which could be expensive + outs = CodeTemplate("flatten_tensor_args( ${outs} )").substitute( + outs=output_names if not is_inplace_foreach else "self" + ) + if not is_inplace_foreach: + return SET_HISTORY.substitute(fn=fn, differentiable_outputs=outs) + else: + return LOOP_OVER_VECTOR_OF_GRAD_FNS.substitute( + preamble=( + f"auto differentiable_outputs = {outs};\n" + f"TORCH_INTERNAL_ASSERT(differentiable_outputs.size() == grad_fns.size());" + ), + statements=f"{fn}_history(differentiable_outputs[i], grad_fns[i]);", + ) + + def emit_save_outputs() -> str: + if is_out_fn: + # out functions don't currently support differentiation + return "" + if info is not None and info.has_derivatives: + stmts = save_variables(info.all_saved_outputs, True) + if len(stmts) == 0: + return "" + if not is_inplace_foreach: + return CONDITIONAL.substitute(cond="grad_fn", statements=stmts) + else: + return LOOP_OVER_VECTOR_OF_GRAD_FNS.substitute( + preamble="", statements=stmts + ) + return "" + + def emit_any_requires_grad() -> list[str]: + extra_condition = "" + if info and info.output_differentiability_conditions: + if len(info.output_differentiability_conditions) != 1: + raise AssertionError( + f"expected 1 output_differentiability_condition, got {len(info.output_differentiability_conditions)}" + ) + extra_condition = f"_any_requires_grad &= ({info.output_differentiability_conditions[0]});" + names_of_args_with_derivatives = [arg.name for arg in args_with_derivatives] + if is_inplace_foreach and info is not None: + for i, arg in enumerate(names_of_args_with_derivatives): + for f_arg, r_arg in inplace_foreacharg2refarg.items(): + if arg == r_arg.name: + names_of_args_with_derivatives[i] = f_arg.name + return [ + SETUP_ANY_REQUIRES_GRAD.substitute( + args_with_derivatives=names_of_args_with_derivatives, + extra_differentiability_conditions=extra_condition, + ) + ] + + def get_any_has_forward_grad_name(var_names: tuple[str, ...]) -> str: + if len(var_names) == 1: + return f"_any_has_forward_grad_{var_names[0]}" + else: + return f"_any_has_forward_grad_{'_'.join(var_names)}" + + def emit_any_has_forward_grad() -> list[str]: + content: list[str] = [] + if not is_foreach: + for derivative in fw_derivatives: + requires_fw_grad = get_any_has_fw_grad_cond(derivative=derivative) + if info and info.output_differentiability_conditions: + if len(info.output_differentiability_conditions) != 1: + raise AssertionError( + f"expected 1 output_differentiability_condition, got {len(info.output_differentiability_conditions)}" + ) + requires_fw_grad = f"({info.output_differentiability_conditions[0]}) && {requires_fw_grad}" + content.append( + f"[[maybe_unused]] auto {get_any_has_forward_grad_name(derivative.var_names)} = {requires_fw_grad};" + ) + else: + for derivative in fw_derivatives: + bool_vector_name = get_any_has_forward_grad_name(derivative.var_names) + cur_derivative_conditions = [] + for inp in differentiable_inputs: + if derivative.required_inputs_fw_grad is None: + continue + if inp.name not in derivative.required_inputs_fw_grad: + continue + inp_name = ( + inp.name + if not inplace + else refargname2inplace_foreacharg[inp.name].name + ) + inp_type = ( + inp.type + if not inplace + else refargname2inplace_foreacharg[inp.name].type + ) + is_list_type = is_tensor_list_type(inp_type) + if is_list_type: + if inp_name != "self": + content.append( + FW_DERIVATIVE_SIZE_CHECK_TEMPLATE.substitute( + inp_name=inp_name + ) + ) + cur_derivative_conditions.append( + # pyrefly: ignore [bad-argument-type] + FW_DERIVATIVE_CHECK_TEMPLATE.substitute( + req_inp=inp_name + "[i]" + ) + ) + else: + cur_derivative_conditions.append( + # pyrefly: ignore [bad-argument-type] + FW_DERIVATIVE_CHECK_TEMPLATE.substitute(req_inp=inp_name) + ) + + content.append(f"std::vector {bool_vector_name}(self.size());") + content.append("for (const auto& i : c10::irange(self.size())) {") + content.append( + f" {bool_vector_name}[i] = {' || '.join(cur_derivative_conditions)};" + ) + content.append("}") + return content + + def emit_check_inplace() -> list[str]: + if not inplace: + return [] + return [ + f"check_inplace({arg.name}, _any_requires_grad);" + for arg in differentiable_outputs + ] + + def emit_fw_derivatives() -> list[str]: + content: list[str] = [] + fw_grad_setters: list[str] = [] + for derivative in fw_derivatives: + res = derivative.var_names + if f.func.name.name.inplace: + if len(res) != 1: + raise AssertionError( + f"Expected number of outputs to be 1 if function is inplace, got {len(res)}" + ) + # TODO update this when inplace namings are unified + res = ("self",) + + if derivative.required_inputs_fw_grad is None: + raise AssertionError("derivative.required_inputs_fw_grad is None") + + unpacked_arguments = "" + for inp in differentiable_inputs: + inp_name = inp.name + is_input_tensorlist = is_foreach and is_tensor_list_type( + inp.type + if not inplace + else refargname2inplace_foreacharg[inp.name].type + ) + input_suffix = "[i]" if is_input_tensorlist else "" + if is_inplace_foreach: + if inp.name in refargname2inplace_foreacharg: + inp_name = refargname2inplace_foreacharg[inp.name].name + zeros_fn = ( + "zeros_symint" + if inplace and inp.name == "self" + else "_efficientzerotensor_symint" + ) + if inp.name in derivative.required_inputs_fw_grad: + unpacked_arguments += ( + FW_DERIVATIVE_DEFINED_GRAD_TEMPLATE.substitute( + inp_name=inp.name, + inp=inp_name + input_suffix, + zeros_fn=zeros_fn, + ) + ) + if zeros_fn == "_efficientzerotensor_symint": + unpacked_arguments += ( + FW_DERIVATIVE_UPDATE_WRAPPED_NUM_TEMPLATE.substitute( + inp_name=inp.name + ) + ) + + if inp.name in (derivative.required_inputs_primal or []): + unpacked_arguments += ( + FW_DERIVATIVE_DEFINED_PRIMAL_TEMPLATE.substitute( + inp_name=inp.name, + inp=inp_name + input_suffix, + ) + ) + if derivative.required_original_self_value: + input_suffix = "s[i]" if is_inplace_foreach else "" + unpacked_arguments += FW_DERIVATIVE_DEFINED_GRAD_TEMPLATE.substitute( + inp_name="original_self", + inp="original_self" + input_suffix, + # pyrefly: ignore [unbound-name] + zeros_fn=zeros_fn, + ) + unpacked_arguments += FW_DERIVATIVE_DEFINED_PRIMAL_TEMPLATE.substitute( + inp_name="original_self", + inp="original_self" + input_suffix, + ) + elif inplace and derivative.is_reusing_outplace_formula: + # The gradient wasn't already cloned, do it if grad mode is enabled + unpacked_arguments += ( + "self_t = GradMode::is_enabled() ? self_t.clone() : self_t;" + ) + + if inplace: + is_inplace_str = "true" + else: + is_inplace_str = "false" + + requires_fw_grad = get_any_has_forward_grad_name(derivative.var_names) + + if all( + (isinstance(var_type, BaseType) and var_type.is_tensor_like()) + for var_type in derivative.var_types + ): + # Is there a way to get from BaseType to BaseCType + if len(derivative.var_types) == 1: + opt_res_grad_type = OptionalCType(BaseCType(tensorT)).cpp_type() + if not is_foreach: + fw_grad_setters.append( + FW_DERIVATIVE_SETTER_TENSOR.substitute( + out_arg=res[0], is_inplace=is_inplace_str + ) + ) + else: + expected_res = "result" if not inplace else "self" + if res[0] != expected_res: + raise AssertionError( + f"res[0] is {res[0]}, expected {expected_res}" + ) + fw_grad_setters.append( + FW_DERIVATIVE_SETTER_TENSOR_FOREACH.substitute( + out_arg=res[0], is_inplace=is_inplace_str + ) + ) + requires_fw_grad += f" && ({derivative.var_names[0]}.defined())" + else: + tuple_type = TupleCType( + [BaseCType(tensorT)] * len(derivative.var_types) + ) + opt_res_grad_type = OptionalCType(tuple_type).cpp_type() + for idx, single_res in enumerate(res): + fw_grad_setters.append( + FW_DERIVATIVE_SETTER_MULTI_OUTPUT.substitute( + idx=idx, all_res="_".join(res), out_arg=single_res + ) + ) + elif ( + isinstance(derivative.var_types[0], ListType) + and derivative.var_types[0].is_tensor_like() + ): + if len(derivative.var_types) != 1: + raise AssertionError( + f"Expected number of outputs to be 1 if function returns ListType, got {len(derivative.var_types)}" + ) + if not is_foreach: + opt_res_grad_type = OptionalCType( + VectorCType(BaseCType(tensorT)) + ).cpp_type() + fw_grad_setters.append( + FW_DERIVATIVE_SETTER_TENSOR_LIST.substitute( + out_arg=res[0], is_inplace=is_inplace_str + ) + ) + else: + # TODO(crcrpar): Should this (= the foreach specific logic) be refactored somehow? + # Only out-place foreach functions that have entries in `tools/autograd/derivatives.yaml` + # can reach here. + opt_res_grad_type = OptionalCType(BaseCType(tensorT)).cpp_type() + fw_grad_setters.append( + FW_DERIVATIVE_SETTER_TENSOR_FOREACH.substitute( + out_arg=res[0], is_inplace=is_inplace_str + ) + ) + else: + raise RuntimeError("Unsupported output type for forward derivative") + + if not is_foreach: + fw_grad_opt_definition = f"{opt_res_grad_type} {'_'.join(res)}_new_fw_grad_opt = ::std::nullopt;" + # View ops create fw_grad that already is a view of the base's fw_grad so just use that + content.append( + FW_DERIVATIVE_TEMPLATE.substitute( + fw_grad_opt_definition=fw_grad_opt_definition, + requires_fw_grad=requires_fw_grad, + formula=derivative.formula, + out_arg="_".join(res), + unpacked_arguments=unpacked_arguments, + ) + ) + else: + # note(crcrpar): Assuming `self` is TensorList. + fw_grad_opt_definition = ( + f"std::vector<{opt_res_grad_type}> {'_'.join(res)}_new_fw_grad_opts" + "(self.size(), ::std::nullopt);" + ) + foreach_forward_grad_formula = derivative.formula + _foreach_arg: Argument | DifferentiableInput + if inplace: + for _foreach_arg, _ref_arg in inplace_foreacharg2refarg.items(): + # note(crcrpar): Massage only Scalar and ArrayRef here. + if not ( + is_tensor_type(_foreach_arg.type) + or is_tensor_list_type(_foreach_arg.type) + ): + pattern = _foreach_arg.name + if isinstance(_foreach_arg.type, ListType): + pattern += "[i]" + foreach_forward_grad_formula = ( + foreach_forward_grad_formula.replace( + _ref_arg.name, pattern + ) + ) + else: + if ( + "result" in foreach_forward_grad_formula + and "result[i]" not in foreach_forward_grad_formula + ): + foreach_forward_grad_formula = ( + foreach_forward_grad_formula.replace("result", "result[i]") + ) + + content.append( + FW_DERIVATIVE_FOREACH_TEMPLATE.substitute( + fw_grad_opt_definition=fw_grad_opt_definition, + vector_of_optional_tensor=f"{'_'.join(res)}_new_fw_grad_opts", + any_has_forward_grad_for_current_index=" || ".join( + get_any_has_forward_grad_name(derivative.var_names) + "[i]" + for derivative in fw_derivatives + ), + formula=foreach_forward_grad_formula, + unpacked_arguments=unpacked_arguments, + ) + ) + + # Set all the grads at the end to avoid: https://github.com/pytorch/pytorch/issues/67367 + content.append("\n".join(fw_grad_setters)) + return content + + def get_any_has_fw_grad_cond(derivative: ForwardDerivative | None) -> str: + # + # Produces a condition string (e.g, "isFwGradDefined(grad_output) || isFwGradDefined(output)") + # + if derivative is None: + # (1) If a derivative is NOT provided, cond will check fw_grad of ALL differentiable inputs + # - Used in the out_fn case when we want to forbid fw derivatives + # - Used in the case where the fw_derivative is not defined, but we want + # To check if there is a decomposition registered for jvp + to_check: list[str] = [] + for inp in list( + mapMaybe( + gen_differentiable_input, + f.func.arguments.non_out + list(f.func.arguments.out), # type: ignore[operator] + ) + ): + if is_tensor_type(inp.type): + to_check.append( + FW_DERIVATIVE_CHECK_TEMPLATE.substitute(req_inp=inp.name) + ) + elif is_tensor_list_type(inp.type): + to_check.append( + FW_DERIVATIVE_TENSORLIST_CHECK_TEMPLATE.substitute( + req_inp=inp.name + ) + ) + else: + raise RuntimeError( + f'Unsupported input type for "{name}" when forbidding forward AD usage.' + ) + return f"({' || '.join(to_check)})" + else: + # (2) If derivative is provided, use that information to determine which inputs + # to check fw_grad for + if derivative.required_inputs_fw_grad is None: + raise AssertionError("derivative.required_inputs_fw_grad is None") + + if len(derivative.required_inputs_fw_grad) == 0: + # Handle functions like stack + # For these, we don't unpack anything and always call the user function + if not ( + len(differentiable_inputs) == 1 + and is_tensor_list_type(differentiable_inputs[0].type) + ): + raise RuntimeError( + f'No differentiable input to "{name}" is a differentiable Tensor (as the provided ' + "forward AD formula does not use any input tangent) even though a forward gradient " + "formula has been defined for it. This case should only happen for function that " + "take a single TensorList as input. All other cases are not supported right now." + ) + any_has_fw_grad = "true" + else: + any_has_fw_grad = " || ".join( + [ + ( + FW_DERIVATIVE_TENSORLIST_CHECK_TEMPLATE + if is_tensor_list_type(inp.type) + else FW_DERIVATIVE_CHECK_TEMPLATE + ).substitute(req_inp=inp.name) + for inp in differentiable_inputs + if inp.name in derivative.required_inputs_fw_grad + ] + ) + any_has_fw_grad = f"({any_has_fw_grad})" + + return any_has_fw_grad + + def emit_forbid_fw_derivatives(is_out_fn: bool = False) -> str: + if is_out_fn: + msg = "because it is an out= function" + else: + msg = ( + "because it has not been implemented yet.\\nPlease file an issue " + "to PyTorch at https://github.com/pytorch/pytorch/issues/new?template=feature-request.yml " + "so that we can prioritize its implementation." + ) + cond = get_any_has_fw_grad_cond(derivative=None) + return ( + FW_DERIVATIVE_FORBID_TEMPLATE.substitute(cond=cond, name=name, msg=msg) + if cond != "" + else "" + ) + + body: list[str] = [] + unpack_args_stats, unpacked_bindings = unpack_args(f) + + body.extend(unpack_args_stats) + if requires_derivative: + body.extend(emit_any_requires_grad()) + body.extend(emit_any_has_forward_grad()) + body.extend(emit_check_inplace()) + body.extend(emit_original_self_definition()) + body.extend(setup_derivative(differentiable_inputs)) + + body.append(emit_call(f, unpacked_bindings, try_jit_decomposition)) + if requires_derivative: + # set_flags has to appear after version_counter, because rebase_history + # requires that the counter is incremented before it is called + body.append(emit_history()) + body.extend(emit_check_if_in_complex_autograd_allowlist()) + + if is_out_fn: + body.append(emit_forbid_fw_derivatives(is_out_fn=True)) + else: + if requires_derivative and not try_jit_decomposition: + if len(fw_derivatives) > 0: + body.extend(emit_fw_derivatives()) + else: + body.append(emit_forbid_fw_derivatives()) + + if requires_derivative: + # Save only after the forward AD has been set up + body.append(emit_save_outputs()) + + if str(f.func.name.name) in RESET_GRAD_ACCUMULATOR: + # `inplace` implies that there is exactly one output named `self`, + # so we can keep the generated code easy. If you need to + # `reset_grad_accumulator` in an operator that's not `inplace`, you can + # remove this check but the code generation will get more elaborate + if not inplace: + raise AssertionError( + f"expected inplace=True for {f.func.name.name} which is in RESET_GRAD_ACCUMULATOR" + ) + body.append("reset_grad_accumulator(self);") + if not returns_void: + body.append(f"return {get_return_value(f)};") + return body diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_view_funcs.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_view_funcs.py new file mode 100644 index 0000000000000000000000000000000000000000..146bee193a6c9cb6abb86582f75fb660dc1d67ea --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/gen_view_funcs.py @@ -0,0 +1,342 @@ +# Generates ViewFuncs.h/cpp +# +# NOTE: If any changes are being made to the ViewFunc codegen please also check +# if updates are needed in torch/csrc/autograd/autograd_not_implemented_fallback.cpp +# The fallback is expected to mimic this codegen, so we should keep the two in sync. + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torchgen.api.dispatcher as dispatcher +from torchgen.api.translate import translate +from torchgen.api.types import ( + BaseCType, + Binding, + NamedCType, + SymIntT, + tensorT, + VectorCType, +) +from torchgen.code_template import CodeTemplate +from torchgen.model import Argument, NativeFunction, OptionalType +from torchgen.utils import FileManager + +from .gen_inplace_or_view_type import ( + CALL_DISPATCH, + extract_bindings, + get_view_info, + modifies_arguments, + use_derived, +) + + +if TYPE_CHECKING: + from torchgen.api.autograd import NativeFunctionWithDifferentiabilityInfo + + +FUNCTION_DECLARATION = CodeTemplate( + """\ +#define ${uppercase_op}_AVAILABLE +struct ${op} : public ${superclass} { + ${op}(${constructor_args}) ${initializer_list} + {} + virtual ~${op}() override = default; + virtual std::vector get_symints() const override; + virtual size_t num_symints() const override; + virtual std::vector get_tensors() const override; + virtual size_t num_tensors() const override; + virtual at::Tensor operator()(const at::Tensor&) const override; + virtual std::unique_ptr clone_and_set( + std::optional> = ::std::nullopt, + std::optional> = ::std::nullopt) const override; + +protected: + virtual void set_symints(std::vector) override; + virtual void set_tensors(std::vector) override; + +private: + ${state} +}; + +""" +) + +FUNCTION_DEFINITION = CodeTemplate( + """\ +std::vector ${op}::get_symints() const { + ${get_symints} +} + +size_t ${op}::num_symints() const { + return static_cast(${num_symints}); +} + +void ${op}::set_symints(std::vector ${symints_vec}) { + TORCH_INTERNAL_ASSERT(${symints_vec}.size() == num_symints()); + ${set_symints} +} + +std::vector ${op}::get_tensors() const { + ${get_tensors} +} + +size_t ${op}::num_tensors() const { + return static_cast(${num_tensors}); +} + +void ${op}::set_tensors(std::vector ${tensors_vec}) { + TORCH_INTERNAL_ASSERT(${tensors_vec}.size() == num_tensors()); + ${set_tensors} +} + +at::Tensor ${op}::operator()(const at::Tensor& ${call_input_name}) const { + return ${op_call}; +} + +std::unique_ptr ${op}::clone_and_set( + std::optional> ${symints_vec}, + std::optional> ${tensors_vec}) const { + auto output = std::make_unique<${op}>(${clone_args}); + if (${symints_vec}.has_value()) { + output->set_symints(std::move(*(${symints_vec}))); + } + if (${tensors_vec}.has_value()) { + output->set_tensors(std::move(*(${tensors_vec}))); + } + return output; +} + +""" +) + + +# e.g. as_strided -> AsStridedViewFunc for camel case or +# as_strided_view_func otherwise +def view_func_name( + f: NativeFunction, include_namespace: bool = False, camel_case: bool = True +) -> str: + name = f.func.name.unambiguous_name() + view_func_name = f"{name.replace('.', '_')}_view_func" + if camel_case: + is_private = view_func_name.startswith("_") + view_func_name = "".join( + [p.title() for p in view_func_name.replace(".", "_").split("_")] + ) + if is_private: + # put the leading underscore back in + view_func_name = f"_{view_func_name}" + namespace = "torch::autograd::generated::" if include_namespace else "" + return f"{namespace}{view_func_name}" + + +def is_symint_or_tensor(arg: Argument) -> bool: + return arg.type.is_tensor_like() or arg.type.is_symint_like() + + +def remove_const_ref(binding: Binding) -> Binding: + return Binding( + name=binding.name, + nctype=binding.nctype.remove_const_ref(), + argument=binding.argument, + default=binding.default, + ) + + +def returns_multi_tensor(fn: NativeFunction) -> bool: + returns = fn.func.returns + if len(returns) != 1: + raise AssertionError(f"Expected 1 return, got {len(returns)}") + returns_list_like = returns[0].type.is_list_like() is not None + returns_tensor_like = returns[0].type.is_tensor_like() + return returns_list_like and returns_tensor_like + + +# Generates strings with logic for getting / setting state of a particular type. +# +# Args: +# bindings (list): List of state bindings of interest (may be empty) +# state_vec_type (NamedCType): Type of vector to either return or copy from +# +# Returns: +# tuple: (list of getter logic strings, list of setter logic strings, string +# with num items expression) +def generate_state_getter_setter( + bindings: list[Binding], + state_vec_type: NamedCType, +) -> tuple[list[str], list[str], str]: + getter_logic = [] + setter_logic = [] + + state_vec = state_vec_type.name + getter_logic.append(f"{state_vec_type.cpp_type()} {state_vec};") + if len(bindings) > 0: + setter_logic.append("auto i = 0;") + + num_exprs = [] + for i, b in enumerate(bindings): + if not isinstance(b.argument, Argument): + raise AssertionError(f"Expected Argument, got {type(b.argument)}") + if b.argument.type.is_list_like(): + # Handle list-likes. + num_expr = f"{b.name}.size()" + num_exprs.append(num_expr) + getter = f"{state_vec}.insert({state_vec}.end(), {b.name}.begin(), {b.name}.end());" + setter = f"std::copy({state_vec}.begin() + i, {state_vec}.begin() + i + {b.name}.size(), {b.name}.begin());" + elif isinstance(b.argument.type, OptionalType): + # Handle optionals. + num_expr = f"({b.name}.has_value() ? 1 : 0)" + num_exprs.append(num_expr) + conditional = f"if({b.name}.has_value())" + getter = ( + f"{conditional} {state_vec}.insert({state_vec}.end(), *({b.name}));" + ) + setter = f"{conditional} {b.name} = {state_vec}[i];" + else: + num_expr = "1" + num_exprs.append(num_expr) + getter = f"{state_vec}.push_back({b.name});" + setter = f"{b.name} = {state_vec}[i];" + + getter_logic.append(getter) + setter_logic.append(setter) + if i < len(bindings) - 1: + setter_logic.append(f"i += {num_expr};") + + # Reserve / assert based on the total number of items expression. + num_items = "0" if len(num_exprs) == 0 else " + ".join(num_exprs) + if len(bindings) > 0: + getter_logic.insert(1, f"{state_vec}.reserve({num_items});") + + getter_logic.append(f"return {state_vec};") + + return getter_logic, setter_logic, num_items + + +def process_function(fn: NativeFunction, template: CodeTemplate) -> str: + bindings = extract_bindings(fn) + non_self_bindings = [b for b in bindings if b.name != "self"] + + non_self_args = fn.func.arguments.flat_all[1:] + non_self_value_bindings = [ + dispatcher.argument(a, remove_non_owning_ref_types=True) for a in non_self_args + ] + + # Generate constructor / clone args for the generated struct. + constructor_args = [b.defn() for b in non_self_bindings] + clone_args = [b.name for b in non_self_bindings] + + # Generate state variable declarations for the generated struct. + state_variables = [ + f"{remove_const_ref(b).defn()};" for b in non_self_value_bindings + ] + + # Generate initializer list expressions for the generated struct. + # allow_expensive_conversions=True because we need to store e.g. SymIntArrayRefs as + # vectors. + init_exprs = translate( + non_self_bindings, non_self_value_bindings, allow_expensive_conversions=True + ) + initializers = [] + for b, init_expr in zip(non_self_bindings, init_exprs): + name = b.nctype.name + if not isinstance(name, str): + raise AssertionError(f"Expected name to be str, got {type(name)}") + initializers.append(f"{name}({init_expr.expr})") + + # Generate call to underlying view op + call_input_name = "input_base" + op_call_args = [call_input_name, *(b.name for b in non_self_bindings)] + op_call = CALL_DISPATCH.substitute( + unambiguous_name=fn.func.name.unambiguous_name(), + unpacked_args=op_call_args, + ) + + # Multi-output views additionally require a view_idx for disambiguation. + if returns_multi_tensor(fn): + view_idx_name = "view_idx" + view_idx_typename = "int64_t" + view_idx_decl = f"{view_idx_typename} {view_idx_name}" + constructor_args.append(view_idx_decl) + clone_args.append(view_idx_name) + state_variables.append(f"{view_idx_decl};") + initializers.append(f"{view_idx_name}({view_idx_name})") + op_call += f"[{view_idx_name}]" + + # Generate initializer list for the generated struct. + initializer_list = f": {', '.join(initializers)}" if len(initializers) > 0 else "" + + # Generate getter / setter logic for any symints. + symint_bindings = [ + b + for b in non_self_bindings + if isinstance(b.argument, Argument) and b.argument.type.is_symint_like() + ] + symints_vec_type = NamedCType("symints", VectorCType(BaseCType(SymIntT))) + get_symints, set_symints, num_symints = generate_state_getter_setter( + symint_bindings, symints_vec_type + ) + + # Generate getter / setter logic for any tensors. + tensor_bindings = [ + b + for b in non_self_bindings + if isinstance(b.argument, Argument) and b.argument.type.is_tensor_like() + ] + tensors_vec_type = NamedCType("tensors", VectorCType(BaseCType(tensorT))) + get_tensors, set_tensors, num_tensors = generate_state_getter_setter( + tensor_bindings, tensors_vec_type + ) + + return template.substitute( + op=view_func_name(fn), + uppercase_op=view_func_name(fn, camel_case=False).upper(), + superclass="torch::autograd::ViewFunc", + initializer_list=initializer_list, + state=state_variables, + constructor_args=constructor_args, + clone_args=clone_args, + symints_vec=symints_vec_type.name, + get_symints=get_symints, + set_symints=set_symints, + num_symints=num_symints, + tensors_vec=tensors_vec_type.name, + get_tensors=get_tensors, + set_tensors=set_tensors, + num_tensors=num_tensors, + call_input_name=call_input_name, + op_call=op_call, + ) + + +def gen_view_funcs( + out: str, + fns_with_infos: list[NativeFunctionWithDifferentiabilityInfo], + template_path: str, +) -> None: + # don't need the info parts, just the function + fns = [fn.func for fn in fns_with_infos if use_derived(fn)] + # only want out-of-place views + view_fns = [ + fn for fn in fns if get_view_info(fn) is not None and not modifies_arguments(fn) + ] + + declarations = [process_function(fn, FUNCTION_DECLARATION) for fn in view_fns] + definitions = [process_function(fn, FUNCTION_DEFINITION) for fn in view_fns] + ops_headers = [f"#include " for fn in view_fns] + + file_basename = "ViewFuncs" + fm = FileManager(install_dir=out, template_dir=template_path, dry_run=False) + for suffix in [".h", ".cpp"]: + fname = file_basename + suffix + fm.write_with_template( + fname, + fname, + lambda: { + "generated_comment": "@" + + f"generated from {fm.template_dir_for_comments()}/{fname}", + "view_func_declarations": declarations, + "view_func_definitions": definitions, + "ops_headers": ops_headers, + }, + ) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/load_derivatives.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/load_derivatives.py new file mode 100644 index 0000000000000000000000000000000000000000..e35f66dbe173a283ce7ab1157b3839202b19e3c4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/load_derivatives.py @@ -0,0 +1,1044 @@ +# Parses derivatives.yaml into autograd functions +# +# Each autograd function is represented by `DifferentiabilityInfo` containing +# a list of `Derivative`. See `torchgen.api.autograd` for the data models. + +from __future__ import annotations + +import re +from collections import Counter, defaultdict +from typing import Any, TYPE_CHECKING + +import yaml + +from torchgen.api import cpp +from torchgen.api.autograd import ( + Derivative, + DifferentiabilityInfo, + ForwardDerivative, + SavedAttribute, +) +from torchgen.api.types import ( + BaseCType, + Binding, + boolT, + CppSignatureGroup, + layoutT, + longT, + NamedCType, + OptionalCType, + scalarTypeT, + SpecialArgName, + stringT, + symIntArrayRefT, + SymIntT, + tensorGeometryT, + tensorOptionsT, + typeAndSizeT, + VectorCType, +) +from torchgen.context import with_native_function +from torchgen.gen import get_grouped_by_view_native_functions, parse_native_yaml +from torchgen.model import ( + AUTOGRAD_KEYS, + FunctionSchema, + NativeFunction, + NativeFunctionsViewGroup, + OperatorName, + SchemaKind, + Type, + Variant, +) +from torchgen.utils import concatMap, IDENT_REGEX, split_name_params +from torchgen.yaml_utils import YamlLoader + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +DerivativeRet = tuple[dict[FunctionSchema, dict[str, DifferentiabilityInfo]], set[str]] + +_GLOBAL_LOAD_DERIVATIVE_CACHE: dict[tuple[str, str], DerivativeRet] = {} + +_VALID_AUTOGRAD_KEYS = set(AUTOGRAD_KEYS) + + +# This function directly adds per-dispatchkey derivative entries for {view}_copy variants of each view op. +# Since every {view} and {view}_copy op shares the same derivative formula, +# we generate them here instead of duplicating them in the yaml. +# See Note [Codegen'd {view}_copy Operators] +def add_view_copy_derivatives( + infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]], + view_groups: list[NativeFunctionsViewGroup], +) -> None: + # Get the map from each view op's name to its corresponding view group + view_name_to_group: dict[OperatorName, NativeFunctionsViewGroup] = { + g.view.func.name: g for g in view_groups + } + + view_infos = {} + + for info_dispatch_dict in infos.values(): + # maybe_view_group only needs to be calculated once per info_dispatch_dict + maybe_view_group = None + view_copy_differentiability_infos = {} + for dispatch_key, info in info_dispatch_dict.items(): + maybe_view_group = view_name_to_group.get(info.func.func.name, None) + if maybe_view_group is not None and maybe_view_group.view_copy is not None: + view_copy_info = info.create_view_copy_from_view_derivative( + maybe_view_group + ) + if view_copy_info is not None: + fn_schema = view_copy_info.func.func + view_copy_differentiability_infos[dispatch_key] = view_copy_info + else: + break + # prefer manually-defined derivatives if any + # pyrefly: ignore [unbound-name] + if len(view_copy_differentiability_infos) > 0 and fn_schema not in infos: + # pyrefly: ignore [unbound-name] + if fn_schema is None: + raise AssertionError("Expected fn_schema to be non-None") + # pyrefly: ignore [unbound-name] + view_infos[fn_schema] = view_copy_differentiability_infos + + infos.update(view_infos) + + +def load_derivatives( + derivatives_yaml_path: str, native_yaml_path: str, tags_yaml_path: str +) -> DerivativeRet: + # Do some caching as this is a deterministic function + global _GLOBAL_LOAD_DERIVATIVE_CACHE + key = (derivatives_yaml_path, native_yaml_path) + if key not in _GLOBAL_LOAD_DERIVATIVE_CACHE: + with open(derivatives_yaml_path) as f: + definitions = yaml.load(f, Loader=YamlLoader) + + funcs = parse_native_yaml(native_yaml_path, tags_yaml_path).native_functions + # From the parsed native functions, separate out the (generated) view_copy functions, + # so we can generate derivatives for them separately. + native_functions_with_view_groups = get_grouped_by_view_native_functions(funcs) + native_functions = concatMap( + lambda g: [g] + if isinstance(g, NativeFunction) + else list(g.functions(include_copy=True)), + native_functions_with_view_groups, + ) + view_groups = [ + g + for g in native_functions_with_view_groups + if isinstance(g, NativeFunctionsViewGroup) + ] + + # What's the difference between function schema v.s. signature? + # function schema is the complete declaration including mutability annotation / default value and etc. + # signature is the canonical schema for a group of functions (in-place/out/functional variants) + # that are semantically related. + functions_by_signature: dict[FunctionSchema, list[NativeFunction]] = ( + defaultdict(list) + ) + functions_by_schema: dict[str, NativeFunction] = {} + for function in native_functions: + functions_by_signature[function.func.signature()].append(function) + if str(function.func) in functions_by_schema: + raise AssertionError(f"Duplicate function schema: {str(function.func)}") + functions_by_schema[str(function.func)] = function + + # Keep track of how many of which ops we've seen so we can + # disambiguate them with a numeric suffix. + op_counter = Counter[str]() + + # infos is a dict that maps FunctionSchema -> a dict of per dispatch key DifferentiabilityInfos + # this is useful because in tools/autograd/gen_autograd.py:match_differentiability_info + # we ultimately need to categorize the DifferentiabilityInfos by FunctionSchema + infos: dict[FunctionSchema, dict[str, DifferentiabilityInfo]] = {} + used_dispatch_keys: set[str] = set() + for defn_dict in definitions: + # Ensure that the old derivatives.yaml schema with no dispatch key can be loaded. + if "dispatch" not in defn_dict: + specification = defn_dict.pop("name") + output_differentiability = defn_dict.pop( + "output_differentiability", None + ) + defn_dict = {"name": specification, "dispatch": {"Default": defn_dict}} + if output_differentiability: + defn_dict["output_differentiability"] = output_differentiability + name, per_dispatch_diffinfos = create_differentiability_info( + defn_dict, + functions_by_signature, + functions_by_schema, + op_counter, + used_dispatch_keys, + ) + infos[name] = per_dispatch_diffinfos + + add_view_copy_derivatives(infos, view_groups) + + # cache both loaded infos as well a a set of all the dispatch_keys/aliases + # that appear in derivatives.yaml. used_dispatch_keys is useful for generating + # VariableType.cpp where we need a TORCH_LIBRARY_IMPL for every autograd dispatch key used + _GLOBAL_LOAD_DERIVATIVE_CACHE[key] = infos, used_dispatch_keys + + return _GLOBAL_LOAD_DERIVATIVE_CACHE[key] + + +# TODO: Why is this going through CppSignatureGroup, that doesn't make sense... +@with_native_function +def cpp_arguments(f: NativeFunction) -> Sequence[Binding]: + sigs = CppSignatureGroup.from_native_function(f, method=False) + if sigs.symint_signature is not None: + return sigs.symint_signature.arguments() + else: + return sigs.signature.arguments() + + +def create_derivative( + f: NativeFunction, + formula: str, + var_names: tuple[str, ...], + available_named_gradients: Sequence[str], +) -> Derivative: + original_formula = formula + arguments: list[NamedCType] = [ + a.nctype.remove_const_ref() for a in cpp_arguments(f) + ] + + return_names = tuple(n if n != "self" else "result" for n in cpp.return_names(f)) + return_types = tuple( + cpp.return_type(r, symint=True).remove_const_ref() for r in f.func.returns + ) + + named_returns = [ + NamedCType(name, type) for name, type in zip(return_names, return_types) + ] + + formula, saved_inputs = saved_variables(formula, arguments, var_names) + formula, saved_outputs = saved_variables(formula, named_returns, var_names) + + used_named_gradients = { + name + for name in available_named_gradients + if re.search(IDENT_REGEX.format(name), formula) + } + + # Check that the referenced derivatives in the formula are in bounds + for i in used_gradient_indices(formula): + if i >= len(f.func.returns): + raise RuntimeError( + f"Out of bounds grads access: derivative formula for {cpp.name(f.func)} " + f"used grads[{i}], but the forward only returns {len(f.func.returns)} outputs." + ) + + return Derivative( + formula=formula, + original_formula=original_formula, + var_names=var_names, + saved_inputs=saved_inputs, + saved_outputs=saved_outputs, + named_gradients=used_named_gradients, + ) + + +def create_forward_derivative( + f: NativeFunction, formula: str, names: tuple[str, ...] +) -> ForwardDerivative: + var_names = names + var_types: tuple[Type, ...] | None = None + for r in f.func.returns: + if r.name in var_names: + if var_types is None: + var_types = () + var_types = var_types + (r.type,) + + # Handle default return names + if var_types is None: + if var_names == ("result",): + if len(f.func.returns) != 1: + raise AssertionError( + f"Expected 1 return for 'result', got {len(f.func.returns)}" + ) + var_types = (f.func.returns[0].type,) + else: + for var_name in var_names: + res = re.findall(r"^result(\d+)$", var_name) + if len(res) == 1: + if var_types is None: + var_types = () + arg_idx = int(res[0]) + var_types = var_types + (f.func.returns[arg_idx].type,) + + if var_types is None: + raise AssertionError("No matching output for forward derivative definition") + return ForwardDerivative( + formula=formula, + var_names=var_names, + var_types=var_types, + required_inputs_fw_grad=None, + required_inputs_primal=None, + required_original_self_value=False, + is_reusing_outplace_formula=False, + ) + + +def postprocess_forward_derivatives( + f: NativeFunction, + defn_name: str, + all_arg_names: list[str], + derivatives: list[Derivative], + forward_derivatives: list[ForwardDerivative], + args_with_derivatives: Sequence[Binding], +) -> list[ForwardDerivative]: + def find_required_inputs(formula: str, postfix: str) -> tuple[str, ...]: + is_foreach = f.func.name.name.base.startswith("_foreach_") + required_inputs = set() + for arg in args_with_derivatives: + if ( + arg.type in ("at::TensorList", "const at::ITensorListRef &") + and not is_foreach + ): + # The functions taking TensorList handle everything internally + continue + arg_name = arg.name + + found = re.search(IDENT_REGEX.format(arg_name), formula) + if found: + raise RuntimeError( + f"The forward formula for {defn_name} is using the base name of the {arg_name} " + f"argument which is ambiguous. You should use {arg_name}_p to access the primal " + f"value and {arg_name}_t to access the tangent." + ) + + found = re.search(IDENT_REGEX.format(arg_name + postfix), formula) + if found: + required_inputs.add(arg_name) + + return tuple(required_inputs) + + updated_derivatives: list[ForwardDerivative] = [] + + for defn in forward_derivatives: + formula = defn.formula + required_inputs_tangent = find_required_inputs(formula, "_t") + if formula == "auto_element_wise": + if f.func.kind() == SchemaKind.inplace: + raise AssertionError( + f"Cannot use auto_element_wise with {f.func.name} because it is an in-place variant" + ) + if ( + (not len(args_with_derivatives) == 1) + or len(forward_derivatives) > 1 + or len(forward_derivatives[0].var_names) > 1 + ): + raise RuntimeError( + f"Derivative definition of {defn_name} in derivatives.yaml defines the " + "forward definition of gradient as element_wise but this only " + "works for functions with a single differentiable input and a " + "single differentiable output." + ) + if not len(derivatives) == 1: + raise RuntimeError( + f"Derivative definition of {defn_name} in derivatives.yaml defines the " + "forward definition of gradient as element_wise but it does not " + "defines the gradient formula for its argument which is required." + ) + # This transformation is based on the observation that for element-wise functions, the Jacobian + # matrix is diagonal and thus doing J * v is the same as (v^T J)^T (in practice, we ignore the transpositions) + # For the complex case, we use hermitian transpose and get (v.conj() J).conj() + # So here we are going to reuse the backward formula and replace two things: + # 1) all occurrences of "grad" with "foo_t.conj()", where foo is the name of the unique differentiable input. + # 2) all usage of an original input "foo" with its primal value "foo_p". + # 3) conjugate the final result + # For example, for abs, the backward formula is: + # grad * self.sgn() + # And this function generates a forward formula that is: + # (self_t.conj() * self_p.sgn()).conj() + + backward_formula = derivatives[0].original_formula + input_name = args_with_derivatives[0].name + + # Do replacement 1) of the grad + def repl(m: Any) -> str: + return f"{m.group(1)}{input_name}_t.conj(){m.group(2)}" + + fw_formula = re.sub(IDENT_REGEX.format("grad"), repl, backward_formula) + + # Do replacement 2) of the input variables + for arg in args_with_derivatives: + arg_name = arg.name + + def repl(m: Any) -> str: + return f"{m.group(1)}{arg_name}_p{m.group(2)}" + + fw_formula = re.sub(IDENT_REGEX.format(arg_name), repl, fw_formula) + + # Do the final conjugate 3) + fw_formula = f"({fw_formula}).conj()" + + # Since there is a single differentiable inputs and we necessarily need its tangent we can + # simply require all differentiable input's tangent. + required_inputs_tangent = tuple(all_arg_names) + formula = fw_formula + elif formula == "auto_linear": + if ( + len(forward_derivatives) > 1 + or len(forward_derivatives[0].var_names) > 1 + ): + raise RuntimeError( + f"Derivative definition of {defn_name} in derivatives.yaml defines the " + "forward definition of gradient as linear but this only works " + "for functions with a single differentiable output." + ) + # This transformation is based on the observation that linear functions can be written as: + # y = f(x) = A * x + # For some matrix A and the Jacobian of the function f is also A. + # So doing J * v = A * v = f(v). + # Hence to do the jvp, we simply need to evaluate the function at the point v instead of x. + # We do this by calling the forward again by replacing any occurrence of the differentiable + # input "foo" by it's tangent "foo_t". + # Note that multiple inputs are not a problem as long as the function is truly linear wrt to + # the vector where all the differentiable inputs are stacked. + + diff_arg_names = [arg.name for arg in args_with_derivatives] + if len(diff_arg_names) == 0: + raise AssertionError("Expected at least one differentiable argument") + + # Do replacement of input variables + new_args = [] + for arg_name in all_arg_names: + if arg_name in diff_arg_names: + arg_name = arg_name + "_t" + # pyrefly: ignore [bad-argument-type] + new_args.append(arg_name) + + # TODO we are trolling + if f.func.has_symint(): + defn_name += "_symint" + + # Call into the forward again. We need two cases here to handle both Tensor methods and at:: functions. + if Variant.function in f.variants: + fw_formula = f"at::{defn_name}({', '.join(new_args)})" + else: + if Variant.method not in f.variants: + raise AssertionError( + f"Expected Variant.method in variants for {f.func.name}" + ) + fw_formula = f"{new_args[0]}.{defn_name}({', '.join(new_args[1:])})" + + # All of the input tangents are always used so all of them are required here. + required_inputs_tangent = tuple(diff_arg_names) + formula = fw_formula + + # At this point, the formula is final and is not modified anymore. + + # During forward formula, we use the primal instead of the input Tensors. + # This call inspects the formula to find for which input's primal are used. + required_inputs_primal = find_required_inputs(formula, "_p") + + updated_derivatives.append( + ForwardDerivative( + formula=formula, + var_names=defn.var_names, + var_types=defn.var_types, + required_inputs_fw_grad=required_inputs_tangent, + required_inputs_primal=required_inputs_primal, + required_original_self_value=False, + is_reusing_outplace_formula=False, + ) + ) + + return updated_derivatives + + +def is_forward_derivative_definition( + all_arg_names: list[str], names: tuple[str, ...] +) -> bool: + for name in names: + return name not in all_arg_names + raise RuntimeError("Expected `names` to be non-empty") + + +def create_differentiability_info( + defn_dict: dict[Any, Any], + functions_by_signature: dict[FunctionSchema, list[NativeFunction]], + functions_by_schema: dict[str, NativeFunction], + op_counter: Counter[str], + used_dispatch_keys: set[str], +) -> tuple[FunctionSchema, dict[str, DifferentiabilityInfo]]: + """Processes a single entry `defn` in derivatives.yaml""" + + def canonical_function( + functions: Sequence[NativeFunction], name: str + ) -> NativeFunction: + for f in functions: + if ( + not f.func.is_functional_fn() + and not f.func.is_out_fn() + and name == str(f.func.name.name) + ): + return f + # some functions only have in-place variants + if name + "_" != cpp.name(functions[0].func): + raise AssertionError( + f"Expected inplace function name '{name}_', got '{cpp.name(functions[0].func)}'" + ) + return functions[0] + + def split_names(raw_names: str) -> tuple[str, ...]: + """Given "foo, bar", return ["foo", "bar"].""" + return tuple(x.strip() for x in raw_names.split(",")) + + def check_grad_usage(defn_name: str, derivatives: Sequence[Derivative]) -> None: + """ + Check for some subtle mistakes one might make when writing derivatives. + These mistakes will compile, but will be latent until a function is + used with double backwards. + """ + + uses_grad = False # true if any derivative uses "grad" + num_grads_uses = 0 # count of uses of "grads" or "grads[INDEX]" + uses_named_grads = False # true if any derivative uses "grad_{name}" + used_grads_indices: list[int] = [] # which indices of grads are used + for d in derivatives: + formula = d.formula + uses_grad = uses_grad or bool( + re.findall(IDENT_REGEX.format("grad"), formula) + ) + num_grads_uses += len(re.findall(IDENT_REGEX.format("grads"), formula)) + uses_named_grads = uses_named_grads or bool(d.named_gradients) + used_grads_indices.extend(used_gradient_indices(formula)) + # This is a basic sanity check: the number of places we see + # "grads" should be no fewer than the number of indices we see + # inside "grads". They may not be equal because we may use + # "grads" without an index. + if num_grads_uses < len(used_grads_indices): + raise AssertionError( + f"num_grads_uses ({num_grads_uses}) < len(used_grads_indices) ({len(used_grads_indices)})" + ) + # Thus if the number is equal, every use of grads is also + # indexed. + only_used_grads_indices = num_grads_uses == len(used_grads_indices) + + if uses_grad and num_grads_uses > 0: + raise RuntimeError( + f"Derivative definition of {defn_name} in derivatives.yaml illegally " + "mixes use of 'grad' and 'grads'. Consider replacing " + "occurrences of 'grad' with 'grads[0]'" + ) + + if only_used_grads_indices and set(used_grads_indices) == {0}: + raise RuntimeError( + f"Derivative definition of {defn_name} in derivatives.yaml solely " + "refers to 'grads[0]'. If the first output is indeed the " + "only differentiable output, replace 'grads[0]' with 'grad'; " + "otherwise, there is a likely error in your derivatives " + "declaration." + ) + + if uses_named_grads and (uses_grad or num_grads_uses > 0): + raise RuntimeError( + f"Derivative definition of {defn_name} in derivatives.yaml illegally " + 'mixes use of "grad_RETURN_NAME" and "grad" or "grads[x]". Use ' + "only one method for identifying gradients." + ) + + @with_native_function + def set_up_derivatives( + f: NativeFunction, + ) -> tuple[ + Sequence[Derivative], + Sequence[ForwardDerivative], + Sequence[Binding], + Sequence[str], + Sequence[str], + ]: + # Set up the derivative information + derivatives: list[Derivative] = [] + forward_derivatives: list[ForwardDerivative] = [] + non_differentiable_arg_names: list[str] = [] + args_with_derivatives_set: set[str] = set() + + all_arg_names = [a.name for a in cpp_arguments(f)] + all_ret_names = [ + r.name for r in f.func.returns + ] # only used for the assert below + # output_differentiability is captured from the enclosed + # scope. Don't modify it. + # + # If it is not present, then no output is explicitly + # undifferentiable. + # + # It may be present and shorter than the length of return + # values. If that's the case, any return value that does not + # have a corresponding entry is considered not differentiable. + differentiability = output_differentiability or [True] * len(f.func.returns) + # A return is available as a named gradient ... + available_named_gradients = [ + f"grad_{ret.name}" + for ret, differentiable in zip(f.func.returns, differentiability) + # if it has not been explicitly made undifferentiable + if differentiable + # and if it has a name + and ret.name is not None + # and if its type is differentiable + and ret.type.is_tensor_like() + ] + + for raw_names in sorted(defn.keys()): + formula = defn[raw_names] + names = split_names(raw_names) + + for name in names: + if name in all_arg_names and name in all_ret_names: + raise AssertionError( + f"While processing the derivative formula for '{f.func.name}' wrt '{name}', " + f"expected '{name}' to not be both an input arg and named return." + ) + + if is_forward_derivative_definition(all_arg_names, names): + forward_derivatives.append(create_forward_derivative(f, formula, names)) + else: + if formula.lower().strip() == "non_differentiable": + non_differentiable_arg_names += names + else: + derivative = create_derivative( + f, formula, names, available_named_gradients + ) + derivatives.append(derivative) + args_with_derivatives_set |= set(names) + + overlap = args_with_derivatives_set.intersection(non_differentiable_arg_names) + if overlap: + raise RuntimeError( + f"derivatives definition for {defn} have overlapped non_differentiable " + f"and differentiable variables: {overlap}" + ) + + # Next, let us determine the list of inputs in order. + # TODO: do we need eagerly calculate and save it here? Can it be derived + # from NativeFunction and `derivatives` on callsites instead? + args_with_derivatives = [ + a for a in cpp_arguments(f) if a.name in args_with_derivatives_set + ] + + # Postprocess forward derivatives definitions now that we know the differentiable arguments + forward_derivatives = postprocess_forward_derivatives( + f, + defn_name, + all_arg_names, + derivatives, + forward_derivatives, + args_with_derivatives, + ) + + # Test to see if the use of 'grads' makes sense. + check_grad_usage(defn_name, derivatives) + + return ( + derivatives, + forward_derivatives, + args_with_derivatives, + non_differentiable_arg_names, + available_named_gradients, + ) + + # NB: Removes 'name' from defn dictionary + specification = defn_dict.pop("name") + defn_name, _ = split_name_params(specification) + # NB: Removes 'output_differentiability' from defn dictionary + # `None` means all differentiable. + output_differentiability = defn_dict.pop("output_differentiability", None) + output_differentiability_conditions = None + if output_differentiability and any( + isinstance(diff, str) for diff in output_differentiability + ): + if len(output_differentiability) != 1: + raise RuntimeError( + f"Not supported: for {specification}," + f"output_differentiability must either be " + f"list[bool] or a list[str] where each str is a " + f"condition. In the case where it is a condition, " + f"we only support single-output functions. " + f"Please file us an issue. " + ) + output_differentiability_conditions = output_differentiability + output_differentiability = [True] + + schema_function = functions_by_schema.get(specification) + if not schema_function: + avail = "\n".join( + k for k, v in functions_by_schema.items() if cpp.name(v.func) == defn_name + ) + raise RuntimeError( + f"could not find ATen function for schema: {specification} " + f". Available signatures:\n{avail}" + ) + + # now map this to the legacy schema; this isn't technically necessary, but we'd need some logic here + # to map in-place schemas to the out-of-place variants. + # TODO: maybe the logic to handle the legacy schema is no longer necessary? + signature = schema_function.func.signature() + functions = functions_by_signature[signature] + if len(functions) == 0: + avail = "\n".join( + str(k) + for k, v in functions_by_signature.items() + if cpp.name(k) == defn_name + ) + raise RuntimeError( + f"could not find ATen function for legacy signature: {signature} " + f"corresponding to schema {specification}. Please report a bug to PyTorch. " + f"Available signatures:\n{avail}" + ) + + canonical = canonical_function(functions, defn_name) + if "grad_input_mask" in (a.name for a in cpp_arguments(canonical)): + raise RuntimeError( + f"Schema for {defn_name} has an argument named grad_input_mask, " + "but this name would be shadowed by our codegen. " + "Please use a different name in native_functions.yaml." + ) + + if "result" in (a.name for a in cpp_arguments(canonical)): + raise RuntimeError( + f"Schema for {defn_name} has an argument named result, " + "but this is only allowed for outputs." + "Please use a different name in native_functions.yaml." + ) + + diffinfo_dict = {} + for key, defn in defn_dict["dispatch"].items(): + if key != "Default" and key not in _VALID_AUTOGRAD_KEYS: + raise RuntimeError( + f"Invalid dispatch key {key} in derivatives.yaml for {specification}," + f" expected key to be one of {_VALID_AUTOGRAD_KEYS}" + ) + if key not in used_dispatch_keys: + used_dispatch_keys.add(key) + + ( + derivatives, + forward_derivatives, + args_with_derivatives, + non_differentiable_arg_names, + available_named_gradients, + ) = set_up_derivatives(canonical) + + used_named_gradients: set[str] = set() + for d in derivatives: + used_named_gradients |= d.named_gradients + + # only assign an op name if we are actually going to calculate a derivative + op = None + if args_with_derivatives: + op_prefix = _create_op_prefix(defn_name) + if key != "Default": + op_prefix = op_prefix + key + op = f"{op_prefix}{op_counter[op_prefix]}" + op_counter[op_prefix] += 1 + + diffinfo_dict[key] = DifferentiabilityInfo( + name=defn_name, + func=canonical, + op=op, + derivatives=derivatives, + forward_derivatives=forward_derivatives, + all_saved_inputs=dedup_vars( + [v for d in derivatives for v in d.saved_inputs] + ), + all_saved_outputs=dedup_vars( + [v for d in derivatives for v in d.saved_outputs] + ), + available_named_gradients=available_named_gradients, + used_named_gradients=used_named_gradients, + args_with_derivatives=args_with_derivatives, + non_differentiable_arg_names=non_differentiable_arg_names, + output_differentiability=output_differentiability, + output_differentiability_conditions=output_differentiability_conditions, + ) + + return canonical.func, diffinfo_dict + + +GRAD_INDEX_REGEX = r"(?:^|\W)grads\[(\d+)\]" + + +def used_gradient_indices(formula: str) -> list[int]: + """Determine a list of gradient indices (the i in grads[i]) that + are used by the formula. + + >>> used_gradient_indices("foo(grads[0], grads[1])") + [0, 1] + """ + return [int(i) for i in re.findall(GRAD_INDEX_REGEX, formula)] + + +def saved_variables( + formula: str, + nctypes: list[NamedCType], + var_names: tuple[str, ...], +) -> tuple[str, tuple[SavedAttribute, ...]]: + def stride_expr(name: str) -> str: + if var_names != (name,): + raise AssertionError( + 'Replacement for ".strides()" is currently only supported for single derivatives of the same tensor ' + 'that ".strides()" is being called on.' + ) + return f'strides_or_error({name}, "{name}")' + + REPLACEMENTS: list[tuple[str, dict[str, Any]]] = [ + # replace self.sym_sizes() with self_sym_sizes + ( + r"{}.sym_sizes\(\)", + { + "suffix": "_sym_sizes", + "nctype": lambda name: NamedCType(name, BaseCType(symIntArrayRefT)), + }, + ), + # replace self->sym_sizes() with self_sym_sizes_opt + ( + r"{}->sym_sizes\(\)", + { + "suffix": "_sym_sizes_opt", + "nctype": lambda name: NamedCType( + name, OptionalCType(BaseCType(symIntArrayRefT)) + ), + "expr": lambda name: f"{name}.has_value() ? std::optional({name}->sym_sizes()) : std::nullopt", + }, + ), + # replace self.sym_blocksize() with self_sym_blocksize_opt + ( + r"{}.sym_blocksize\(\)", + { + "suffix": "_self_sym_blocksize_opt", + "nctype": lambda name: NamedCType( + name, OptionalCType(BaseCType(symIntArrayRefT)) + ), + "expr": lambda name: f"at::sparse_csr::getSymIntBlockSize({name})", + }, + ), + # replace self.options() with self_options + ( + r"{}.options\(\)", + { + "suffix": "_options", + "nctype": lambda name: NamedCType(name, BaseCType(tensorOptionsT)), + }, + ), + # replace zeros_like(self) with self_info + ( + r"zeros_like\({}\)", + { + "suffix": "_info", + "nctype": lambda name: NamedCType(name, BaseCType(typeAndSizeT)), + "expr": lambda name: name, # at save-time + "res": lambda name: name + "_info.zeros()", # at eval-time + }, + ), + # replace self.sym_size(2) with self_sym_size_2 + ( + r"{}.sym_size\((-?\w+)\)", + { + "suffix": lambda m: f"_sym_argsize_{m.groups()[0].replace('-', 'minus_')}", + "nctype": lambda name: NamedCType(name, BaseCType(SymIntT)), + }, + ), + # replace self.numel() with self_numel + ( + r"{}.numel\(\)", + { + "suffix": "_numel", + "nctype": lambda name: NamedCType(name, BaseCType(longT)), + }, + ), + # replace self.sym_numel() with self_sym_numel + ( + r"{}.sym_numel\(\)", + { + "suffix": "_sym_numel", + "nctype": lambda name: NamedCType(name, BaseCType(SymIntT)), + }, + ), + # replace to_args_sizes(self) with self_args_sizes + ( + r"to_args_sizes\({}\)", + { + "suffix": "_args_sizes", + "nctype": lambda name: NamedCType( + name, VectorCType(VectorCType(BaseCType(longT))) + ), + }, + ), + # replace to_args_sizes_symint(self) with self_args_sizes + ( + r"to_args_sizes_symint\({}\)", + { + "suffix": "_args_sizes_symint", + "nctype": lambda name: NamedCType( + name, VectorCType(VectorCType(BaseCType(SymIntT))) + ), + }, + ), + # replace to_args_scalartypes(self) with self_args_scalartypes + ( + r"to_args_scalartypes\({}\)", + { + "suffix": "_args_scalartypes", + "nctype": lambda name: NamedCType( + name, VectorCType(BaseCType(scalarTypeT)) + ), + }, + ), + # replace TensorGeometry(self) with self_geometry + ( + r"TensorGeometry\({}\)", + { + "suffix": "_geometry", + "nctype": lambda name: NamedCType(name, BaseCType(tensorGeometryT)), + }, + ), + ( + r"{}.scalar_type\(\)", + { + "suffix": "_scalar_type", + "nctype": lambda name: NamedCType(name, BaseCType(scalarTypeT)), + }, + ), + # replace self.dim() with self_dim + ( + r"{}.dim\(\)", + { + "suffix": "_dim", + "nctype": lambda name: NamedCType(name, BaseCType(longT)), + }, + ), + # replace self.sym_strides() with self_sym_strides + ( + r"{}.sym_strides\(\)", + { + "suffix": "_sym_strides", + "nctype": lambda name: NamedCType(name, BaseCType(symIntArrayRefT)), + "expr": stride_expr, + }, + ), + # replace self.layout() with self_layout + ( + r"{}.layout\(\)", + { + "suffix": "_layout", + "nctype": lambda name: NamedCType(name, BaseCType(layoutT)), + }, + ), + # replace self.is_conj() with self_conjugate + ( + r"{}.is_conj\(\)", + { + "suffix": "_conjugate", + "nctype": lambda name: NamedCType(name, BaseCType(boolT)), + }, + ), + ] + + # find which arguments need to be saved + saved: list[SavedAttribute] = [] + + if ".sizes()" in formula or "->sizes()" in formula: + raise RuntimeError( + ".sizes() is not supported in derivative formulas. Instead, please use the SymInt version," + + f".sym_sizes(), which returned a c10::SymIntArrayRef. formula={formula}" + ) + if re.search(r"\.size\([-]?\d+\)", formula) or re.search( + r"->size\([-]?\d+\)", formula + ): + raise RuntimeError( + ".size(int) is not supported in derivative formulas. Instead, please use the SymInt version," + + f".sym_size(int), which returned a c10::SymIntArrayRef. formula={formula}" + ) + if ".strides()" in formula or "->strides()" in formula: + raise RuntimeError( + ".strides() is not supported in derivative formulas. Instead, please use the SymInt version," + + f".sym_strides(), which returned a c10::SymIntArrayRef. formula={formula}" + ) + for nctype in nctypes: + # pyrefly: ignore [bad-assignment] + name = ( + nctype.name.name if isinstance(nctype.name, SpecialArgName) else nctype.name + ) + # First search the formula for expressions which can be evaluated + # when the autograd Function is created to avoid saving variables + for regex, info in REPLACEMENTS: + + def repl(m: re.Match[str]) -> str: + suffix: str = ( + # pyrefly: ignore [bad-assignment] + info["suffix"](m) if callable(info["suffix"]) else info["suffix"] + ) + expr: str = info["expr"](name) if "expr" in info else m.group(0) + saved.append( + SavedAttribute( + nctype=info["nctype"](name + suffix), + expr=expr, + ) + ) + if "res" in info: + replacement: str = info["res"](name) + return replacement + return name + suffix + + formula = re.sub(regex.format(name), repl, formula) + + # std::optional types stored in Backward nodes must be + # converted to std::optional before being passed into + # the backward function + if nctype.type == OptionalCType(BaseCType(stringT)): + formula = re.sub( + rf"\b{name}\b", + f"{name}.has_value() ? std::optional({name}.value()) : std::nullopt", + formula, + ) + + # Find any variables which remain in the formula and save them + if re.search(IDENT_REGEX.format(name), formula): + saved.append( + SavedAttribute( + nctype=nctype, + expr=name, + ) + ) + + return formula, tuple(saved) + + +def _create_op_prefix(name: str) -> str: + r"""Takes a native function name converts to an op prefix name. + + Note that the "name" parameter must be the native function name + without the optional variant suffix, so "add" instead of + "add.out". + + OP names correspond to classes, hence the change to title case. + + Example:: + + >>> _create_op_prefix("add") + 'AddBackward' + """ + camel_case = "".join([p.title() for p in name.split("_")]) + return (camel_case + "Backward").replace("ForwardBackward", "Backward") + + +def dedup_vars(vars: Sequence[SavedAttribute]) -> Sequence[SavedAttribute]: + seen: set[str] = set() + saved: list[SavedAttribute] = [] + for var in vars: + name = ( + var.nctype.name.name + if isinstance(var.nctype.name, SpecialArgName) + else var.nctype.name + ) + if name in seen: + continue + seen.add(name) + saved.append(var) + return saved diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/ADInplaceOrViewType.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/ADInplaceOrViewType.cpp new file mode 100644 index 0000000000000000000000000000000000000000..e8276697eee065a36d1b16e583a5f011f92541c2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/ADInplaceOrViewType.cpp @@ -0,0 +1,38 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +#include "torch/csrc/autograd/VariableTypeUtils.h" +#include "torch/csrc/autograd/generated/ViewFuncs.h" + +#include +#include +#include + +// ${generated_comment} + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +using namespace at; +using torch::autograd::CreationMeta; +using torch::autograd::as_view; +using torch::autograd::increment_version; + +namespace torch { + +namespace ADInplaceOrView { + +namespace { +${inplace_or_view_method_definitions} +} // namespace +} // namespace ADInplaceOrView + +namespace { + +TORCH_LIBRARY_IMPL(aten, ADInplaceOrView, m) { + ${inplace_or_view_wrapper_registrations}; +} + +} // namespace +} // namespace torch diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/Functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/Functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..ba5cb3d912c5d7a3bbf31f4b0d38d4413dfc160c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/Functions.cpp @@ -0,0 +1,44 @@ +#include "torch/csrc/autograd/FunctionsManual.h" +#include "torch/csrc/dynamo/compiled_autograd.h" + +// ${generated_comment} + +// The manual function definitions that used to be here are now in torch/csrc/autograd/FunctionsManual.cpp +// This speeds up re-compilation and allow to share these implementations so that they can be +// used for forward mode AD formulas as well. + +using namespace torch::autograd::generated::details; +using at::Tensor; +using at::Scalar; +using at::IntArrayRef; +using at::TensorList; + +namespace torch::autograd::generated { + +static at::IValue compute_output_metadata(const torch::autograd::edge_list& next_edges) { + auto output_metadata = torch::dynamo::autograd::IValuePacker< + std::vector>>::pack( + torch::dynamo::autograd::get_input_metadata(next_edges)); + return output_metadata; +} + +static C10_NOINLINE variable_list compiled_autograd_apply_functional( + const PackedArgs& packed_args, + const edge_list& next_edges, + SwapSavedVariables& saved, + const variable_list& grads, + const std::string& name) { + auto output_metadata = compute_output_metadata(next_edges); + const auto& pyinterface = torch::dynamo::autograd::getPyCompilerInterface(); + return pyinterface->call_function( + saved.get_py_compiler(), + "apply_functional", + name, + grads, + packed_args.vec(), + output_metadata); +} + +${autograd_function_definitions} + +} // namespace torch::autograd::generated diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/Functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/Functions.h new file mode 100644 index 0000000000000000000000000000000000000000..911d7d905c002b29941167ccff112a8079d48266 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/Functions.h @@ -0,0 +1,51 @@ +#pragma once + +// ${generated_comment} + +#include +#include +#include + +#include "torch/csrc/autograd/function.h" +#include "torch/csrc/autograd/variable.h" +#include "torch/csrc/autograd/saved_variable.h" +#include + +#include + +namespace torch { namespace autograd { namespace generated { + +using at::Scalar; +using at::Tensor; +using at::IntArrayRef; +using at::ArrayRef; +using at::Type; +using at::TensorGeometry; +using at::ScalarType; +using std::optional; +using c10::fmap; + +inline std::vector unpack_list(at::ArrayRef xs, std::shared_ptr saved_for = nullptr) { + // NB: we must explicitly do the conversion in the lambda, otherwise template + // deduction will give a Tensor of Variable which is not convertible + return fmap(xs, [&saved_for](const SavedVariable& x) { + // TODO(crcrpar): Use `std::move(saved_for)` to avoid incrementing refcount, which would need refactoring. + return static_cast(x.unpack(saved_for)); + }); +} + +inline c10::List> unpack_opt_list(at::ArrayRef xs, std::shared_ptr saved_for = nullptr) { + torch::List> result; + result.reserve(xs.size()); + for (const SavedVariable& v : xs) { + auto var = v.unpack(saved_for); + result.push_back(var.defined() ? std::optional(var) : ::std::nullopt); + } + return result; +} + +using torch::autograd::TypeAndSize; + +${autograd_function_declarations} + +}}} // namespace torch::autograd::generated diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/TraceType.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/TraceType.cpp new file mode 100644 index 0000000000000000000000000000000000000000..fb5e7ae44a5353a3cc2a90858fe33b7fc0ef8bfd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/TraceType.cpp @@ -0,0 +1,40 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +#include "torch/csrc/jit/frontend/tracer.h" + +#include + +#include "torch/csrc/autograd/function.h" + +#include "ATen/quantized/Quantizer.h" + +// ${generated_comment} + +// See the `Tracer` section in `torch/csrc/jit/OVERVIEW.md`. +// NOTE See [Sharded File] comment in VariableType + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +using namespace at; + +namespace torch { + +namespace TraceType { + +namespace { +${trace_method_definitions} +} // namespace +} // namespace TraceType + +namespace { + +TORCH_LIBRARY_IMPL(aten, Tracer, m) { + ${trace_wrapper_registrations}; +} + +} // namespace + +} // namespace torch diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/VariableType.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/VariableType.cpp new file mode 100644 index 0000000000000000000000000000000000000000..d1de108283b1169902a085e4886de7a0113c309c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/VariableType.cpp @@ -0,0 +1,77 @@ +#include "torch/csrc/autograd/VariableTypeUtils.h" +#include "torch/csrc/autograd/generated/VariableType.h" +#include "torch/csrc/autograd/FunctionsManual.h" + +#include +#include +#include +#include + +#include + + +// ${generated_comment} + +// NOTE [Sharded File]: on this file's split-into-shards state +// +// Back in the good old days, VariableType.cpp was generated as one +// file with every function in it, and everything was great and +// simple. +// +// However, this file was also very large (over 36,000 lines), and +// compiling it was very slow, and in fact was a significant +// bottleneck for incremental rebuilds. To address this, we now +// generate the file split across multiple shards, named +// VariableType_0.cpp and so on, which can be compiled in parallel. +// +// For ease of inspection and debugging, so that it's not necessary to +// go rooting around in multiple files, we also generate all the +// functions together in VariableTypeEverything.cpp. This generated +// file is only for convenience; it's not actually used in the +// build. If the file you're looking at now is one of the shards, you +// may want to switch over to the Everything variant to make you +// grepping smoother. + +using namespace at; +using namespace torch::autograd::generated; +using namespace torch::autograd::generated::details; + + +namespace torch::autograd { + +namespace VariableType { +namespace{ +[[maybe_unused]] void reset_grad_accumulator(Variable& self) { + AutogradMeta* meta = torch::autograd::impl::get_autograd_meta(self); + if (meta != nullptr) { + meta->grad_accumulator_.reset(); + } +} +[[maybe_unused]] size_t expected_fresh_use_count(const Variable& self) { + if (!self.defined()) { + // An UndefinedTensorImpl always has a use count of 0 + return 0; + } + if (self.unsafeGetTensorImpl()->pyobj_slot()->load_pyobj() != nullptr) { + // A TensorImpl with a Python object has a use count of 2 + return 2; + } + // A fresh TensorImpl (with no PyObject) has a use count of 1 + return 1; +} +} + +namespace { + + +${type_derived_method_definitions} +} +} + +namespace { + +${wrapper_registrations} + +} + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/VariableType.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/VariableType.h new file mode 100644 index 0000000000000000000000000000000000000000..02959757e5c007a7d54526dc2ca18698748e95f1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/VariableType.h @@ -0,0 +1,55 @@ +#pragma once + +// ${generated_comment} + +#include +#include + +#include + +#include +#include + +#include // for size_t +#include // for function +#include // for unique_ptr +#include +#include + +namespace at { + struct Quantizer; +} + +namespace torch { namespace autograd { + +using Variable = at::Tensor; +using at::Context; +using at::Device; +using at::Dimname; +using at::DimnameList; +using at::Generator; +using at::IntArrayRef; +using at::MemoryFormat; +using at::QScheme; +using at::Scalar; +using at::ScalarType; +using at::Storage; +using at::Tensor; +using at::TensorList; +using at::TensorOptions; +using at::Quantizer; +using std::optional; + +namespace VariableType { + TORCH_API std::vector allCUDATypes(); + TORCH_API std::vector allXPUTypes(); + TORCH_API std::vector allCPUTypes(); + TORCH_API std::vector allPrivateUser1Types(); + + at::Tensor & unpack(Tensor & t, const char * name, int pos); + const at::Tensor & unpack(const Tensor & t, const char * name, int pos); + at::Tensor unpack_opt(const Tensor & t, const char * name, int pos); + std::vector unpack(const at::ITensorListRef& tl, const char *name, int pos); +} + +}} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/ViewFuncs.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/ViewFuncs.cpp new file mode 100644 index 0000000000000000000000000000000000000000..11b9b194fb46f924e863c4c1dab5cbb8dbb0601b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/ViewFuncs.cpp @@ -0,0 +1,14 @@ +#include + +// ${generated_comment} + +using at::Tensor; +using at::Scalar; +using at::IntArrayRef; +using at::TensorList; + +namespace torch::autograd::generated { + +${view_func_definitions} + +} // namespace torch::autograd::generated diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/ViewFuncs.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/ViewFuncs.h new file mode 100644 index 0000000000000000000000000000000000000000..1f69c062d344e4cd5f98cf5f34fd4278019fdf8a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/ViewFuncs.h @@ -0,0 +1,28 @@ +#pragma once + +// ${generated_comment} + +#include +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +namespace torch::autograd::generated { + +using at::Scalar; +using at::Tensor; +using at::IntArrayRef; +using at::ArrayRef; +using at::Type; +using at::ScalarType; +using std::optional; +using c10::fmap; + +${view_func_declarations} + +} // namespace torch::autograd::generated diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/annotated_fn_args.py.in b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/annotated_fn_args.py.in new file mode 100644 index 0000000000000000000000000000000000000000..1012c008451745b8f1ed1454a864f666caf2618a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/annotated_fn_args.py.in @@ -0,0 +1,11 @@ +""" +This file is needed for generating procedural tests required for +testing __torch_function__. See tests/test_overrides.py. +""" + +# flake8: noqa +import torch + +annotated_args = { +${annotated_args} +} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_enum_tag.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_enum_tag.cpp new file mode 100644 index 0000000000000000000000000000000000000000..83cfad1d7ba4d6fc3529caf78e036c5883e7bc23 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_enum_tag.cpp @@ -0,0 +1,15 @@ +#include +#include +#include +#include + +namespace py = pybind11; +namespace torch { + namespace autograd { + void initEnumTag(PyObject* module) { + auto m = py::handle(module).cast(); + py::enum_(m, "Tag") + ${enum_of_valid_tags}; + m.doc() = "An Enum that contains tags that can be assigned to an operator registered in C++."; + } +}} diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_fft_functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_fft_functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..71ac4e2226d2db418eba5690995424d3f007e620 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_fft_functions.cpp @@ -0,0 +1,81 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include "torch/csrc/Device.h" +#include "torch/csrc/DynamicTypes.h" +#include "torch/csrc/Exceptions.h" +#include "torch/csrc/autograd/python_fft_functions.h" +#include "torch/csrc/autograd/generated/python_return_types.h" +#include "torch/csrc/autograd/python_variable.h" +#include "torch/csrc/autograd/utils/wrap_outputs.h" +#include "torch/csrc/autograd/utils/python_arg_parsing.h" +#include "torch/csrc/autograd/generated/variable_factories.h" +#include "torch/csrc/utils/out_types.h" +#include "torch/csrc/utils/pycfunction_helpers.h" +#include "torch/csrc/utils/python_arg_parser.h" +#include "torch/csrc/utils/structseq.h" +#include "torch/csrc/utils/device_lazy_init.h" + +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +using at::Tensor; +using at::Device; +using at::Layout; +using at::Scalar; +using at::ScalarType; +using at::Backend; +using at::OptionalDeviceGuard; +using at::DeviceGuard; +using at::TensorOptions; +using at::IntArrayRef; +using at::Generator; +using at::TensorList; +using at::Dimname; +using at::DimnameList; + +using torch::utils::check_out_type_matches; +using namespace torch::autograd::utils; + +namespace torch::autograd { + +// generated forward declarations start here + +${py_forwards} + +static PyMethodDef fft_functions[] = { + ${py_method_defs} + {NULL} +}; + +static PyObject* THPFFTVariableFunctionsModule = NULL; + +void initFFTFunctions(PyObject* module) { + static struct PyModuleDef def = { + PyModuleDef_HEAD_INIT, + "torch._C._fft", + NULL, + -1, + fft_functions + }; + PyObject* fft = PyModule_Create(&def); + THPFFTVariableFunctionsModule = fft; + if (!fft) { + throw python_error(); + } + // steals a reference to fft + if (PyModule_AddObject(module, "_fft", fft) != 0) { + throw python_error(); + } +} + +// generated methods start here + +${py_methods} + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..1522d6cd0f5a2a1fc0188bf9d6d0d59fe1b27d85 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_functions.cpp @@ -0,0 +1,37 @@ +#include + +// ${generated_comment} + +#include +#include + +#include +#include "torch/csrc/autograd/generated/Functions.h" +#include "torch/csrc/autograd/python_cpp_function.h" +#include +#include +#include +#include +#include + +// NOTE: See [Sharded File] comment in VariableType + +namespace torch::autograd::generated { + +template +static void addClass(PyObject* module, PyTypeObject& type, const char* name, + PyGetSetDef* function_properties=NULL, PyMethodDef* function_methods=NULL) +{ + _initFunctionPyTypeObject(type, name, function_properties, function_methods); + Py_INCREF(&type); + PyModule_AddObject(module, name, (PyObject*)&type); + registerCppFunction(typeid(C), &type); +} + +${py_function_props_and_getters} + +void initialize_autogenerated_functions${shard_id}(PyObject* module) { + ${py_function_initializers} +} + +} // namespace torch::autograd::generated diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_functions.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_functions.h new file mode 100644 index 0000000000000000000000000000000000000000..22e37207e219431100fefaf21b02e3ed0f63d956 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_functions.h @@ -0,0 +1,17 @@ +#pragma once + +#include + +// ${generated_comment} + +// Python bindings for automatically generated autograd functions + +namespace torch { namespace autograd { namespace generated { + +${shard_forward_declare} + +inline void initialize_autogenerated_functions(PyObject* module) { + ${shard_call} +} + +}}} // namespace torch::autograd::generated diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_linalg_functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_linalg_functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c93752a3ddbfcf111426f98c3ea68fc625e94def --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_linalg_functions.cpp @@ -0,0 +1,68 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include "torch/csrc/Device.h" +#include "torch/csrc/DynamicTypes.h" +#include "torch/csrc/Exceptions.h" +#include "torch/csrc/autograd/python_linalg_functions.h" +#include "torch/csrc/autograd/generated/python_return_types.h" +#include "torch/csrc/autograd/python_variable.h" +#include "torch/csrc/autograd/utils/wrap_outputs.h" +#include "torch/csrc/autograd/utils/python_arg_parsing.h" +#include "torch/csrc/utils/pycfunction_helpers.h" +#include "torch/csrc/utils/python_arg_parser.h" +#include "torch/csrc/utils/structseq.h" + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +using at::Tensor; +using at::Scalar; +using at::ScalarType; +using at::MemoryFormat; +using at::Generator; +using at::IntArrayRef; +using at::TensorList; + +using namespace torch::autograd::utils; + +namespace torch::autograd { + +// generated forward declarations start here + +${py_forwards} + +static PyMethodDef linalg_functions[] = { + ${py_method_defs} + {NULL} +}; + +static PyObject* THPLinalgVariableFunctionsModule = NULL; + +void initLinalgFunctions(PyObject* module) { + static struct PyModuleDef def = { + PyModuleDef_HEAD_INIT, + "torch._C._linalg", + NULL, + -1, + linalg_functions + }; + PyObject* linalg = PyModule_Create(&def); + THPLinalgVariableFunctionsModule = linalg; + if (!linalg) { + throw python_error(); + } + // steals a reference to linalg + if (PyModule_AddObject(module, "_linalg", linalg) != 0) { + throw python_error(); + } +} + +// generated methods start here + +${py_methods} + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_nested_functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_nested_functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..3acb5128cee1e180de887080106e7cf5559f15ee --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_nested_functions.cpp @@ -0,0 +1,81 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include "torch/csrc/Device.h" +#include "torch/csrc/DynamicTypes.h" +#include "torch/csrc/Exceptions.h" +#include "torch/csrc/autograd/python_nested_functions.h" +#include "torch/csrc/autograd/generated/python_return_types.h" +#include "torch/csrc/autograd/python_variable.h" +#include "torch/csrc/autograd/utils/wrap_outputs.h" +#include "torch/csrc/autograd/utils/python_arg_parsing.h" +#include "torch/csrc/autograd/generated/variable_factories.h" +#include "torch/csrc/utils/out_types.h" +#include "torch/csrc/utils/pycfunction_helpers.h" +#include "torch/csrc/utils/python_arg_parser.h" +#include "torch/csrc/utils/structseq.h" +#include "torch/csrc/utils/device_lazy_init.h" + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +using at::Tensor; +using at::Device; +using at::Layout; +using at::Scalar; +using at::ScalarType; +using at::Backend; +using at::OptionalDeviceGuard; +using at::DeviceGuard; +using at::TensorOptions; +using at::IntArrayRef; +using at::OptionalIntArrayRef; +using at::Generator; +using at::TensorList; +using at::Dimname; +using at::DimnameList; + +using namespace torch::autograd::utils; + +namespace torch::autograd { + +// generated forward declarations start here + +${py_forwards} + +static PyMethodDef nested_functions[] = { + {NULL, NULL, 0, NULL}, + ${py_method_defs} + {NULL} +}; + +static PyObject* THPNestedVariableFunctionsModule = NULL; + +void initNestedFunctions(PyObject* module) { + nested_functions[0] = get_nested_functions_manual()[0]; + static struct PyModuleDef def = { + PyModuleDef_HEAD_INIT, + "torch._C._nested", + NULL, + -1, + nested_functions + }; + PyObject* nested = PyModule_Create(&def); + THPNestedVariableFunctionsModule = nested; + if (!nested) { + throw python_error(); + } + // steals a reference to nested + if (PyModule_AddObject(module, "_nested", nested) != 0) { + throw python_error(); + } +} + +// generated methods start here + +${py_methods} + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_nn_functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_nn_functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..8eabb0da2332283a02e98e54dd0a277a83a55ad6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_nn_functions.cpp @@ -0,0 +1,113 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include "torch/csrc/Device.h" +#include "torch/csrc/DynamicTypes.h" +#include "torch/csrc/Exceptions.h" +#include "torch/csrc/autograd/python_nn_functions.h" +#include "torch/csrc/autograd/generated/python_return_types.h" +#include "torch/csrc/autograd/python_variable.h" +#include "torch/csrc/autograd/utils/wrap_outputs.h" +#include "torch/csrc/autograd/utils/python_arg_parsing.h" +#include "torch/csrc/utils/pycfunction_helpers.h" +#include "torch/csrc/utils/python_arg_parser.h" +#include "torch/csrc/utils/structseq.h" +#include "torch/csrc/utils/tensor_memoryformats.h" + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +using at::Tensor; +using at::Scalar; +using at::MemoryFormat; +using at::Generator; +using at::IntArrayRef; +using at::ArrayRef; + +using namespace torch::autograd::utils; + +namespace torch::autograd { + +static PyObject* THPNNVariableFunctionsModule = nullptr; + +static PyObject * THPVariable__parse_to(PyObject* module, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "to(Device device=None, ScalarType dtype=None, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)", + "to(ScalarType dtype, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)", + "to(Tensor tensor, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)", + }); + ParsedArgs<5> parsed_args; + auto r = parser.parse(args, kwargs, parsed_args); + if (r.has_torch_function()) { + return handle_torch_function(r, args, kwargs, THPNNVariableFunctionsModule, "torch.nn", "_parse_to"); + } + auto parsed = parse_to_conversion(r, /*allow_copy*/ false); // we don't want copy for nn.Module.to + auto& device = std::get<0>(parsed); + auto& scalarType = std::get<1>(parsed); + auto non_blocking = std::get<2>(parsed); + auto opt_memory_format = std::get<4>(parsed); + auto tuple = THPObjectPtr{PyTuple_New(4)}; + if (!tuple) throw python_error(); + if (device) { + PyTuple_SET_ITEM(tuple.get(), 0, THPDevice_New(*device)); + } else { + Py_INCREF(Py_None); + PyTuple_SET_ITEM(tuple.get(), 0, Py_None); + } + if (scalarType) { + PyTuple_SET_ITEM(tuple.get(), 1, Py_NewRef(torch::getTHPDtype(*scalarType))); + } else { + Py_INCREF(Py_None); + PyTuple_SET_ITEM(tuple.get(), 1, Py_None); + } + PyTuple_SET_ITEM(tuple.get(), 2, torch::autograd::utils::wrap(non_blocking)); + if (opt_memory_format.has_value()) { + PyTuple_SET_ITEM(tuple.get(), 3, Py_NewRef(torch::utils::getTHPMemoryFormat(opt_memory_format.value()))); + } else { + Py_INCREF(Py_None); + PyTuple_SET_ITEM(tuple.get(), 3, Py_None); + } + return tuple.release(); + END_HANDLE_TH_ERRORS +} + +// generated forward declarations start here + +${py_forwards} + +static PyMethodDef nn_functions[] = { + {"_parse_to", castPyCFunctionWithKeywords(THPVariable__parse_to), + METH_VARARGS | METH_KEYWORDS, nullptr}, + ${py_method_defs} + {nullptr} +}; + +void initNNFunctions(PyObject* module) { + static struct PyModuleDef def = { + PyModuleDef_HEAD_INIT, + "torch._C._nn", + nullptr, + -1, + nn_functions + }; + PyObject* nn = PyModule_Create(&def); + THPNNVariableFunctionsModule = nn; + if (!nn) { + throw python_error(); + } + // steals a reference to nn + if (PyModule_AddObject(module, "_nn", nn) != 0) { + throw python_error(); + } +} + +// generated methods start here + +${py_methods} + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_return_types.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_return_types.cpp new file mode 100644 index 0000000000000000000000000000000000000000..139e6b8958336cfcc8328fa33581e9f1ab6d5532 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_return_types.cpp @@ -0,0 +1,52 @@ +#include + +#include +#include +#include + +#include "torch/csrc/autograd/generated/python_return_types.h" +#include "torch/csrc/utils/structseq.h" +#include "torch/csrc/Exceptions.h" + +namespace torch { namespace autograd { namespace generated { + +${py_return_types} + +}}} + +namespace torch::autograd { + +static void addReturnType( + PyObject* module, + const char* name, + PyTypeObject* type) { + // hold onto the TypeObject for the unlikely case of user + // deleting or overriding it. + Py_INCREF(type); + if (PyModule_AddObject( + module, + name, + (PyObject*)type) != 0) { + Py_DECREF(type); + throw python_error(); + } +} + +void initReturnTypes(PyObject* module) { + static struct PyModuleDef def = { + PyModuleDef_HEAD_INIT, "torch._C._return_types", nullptr, -1, {}}; + PyObject* return_types_module = PyModule_Create(&def); + if (!return_types_module) { + throw python_error(); + } + + ${py_return_types_registrations} + + // steals a reference to return_types on success + if (PyModule_AddObject(module, "_return_types", return_types_module) != 0) { + Py_DECREF(return_types_module); + throw python_error(); + } +} + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_return_types.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_return_types.h new file mode 100644 index 0000000000000000000000000000000000000000..ce6c355ea146a272709255b898603764112168b9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_return_types.h @@ -0,0 +1,14 @@ +#pragma once + +namespace torch { +namespace autograd { +namespace generated { + +${py_return_types_declarations} + +} + +void initReturnTypes(PyObject* module); + +} // namespace autograd +} // namespace torch diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_sparse_functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_sparse_functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..648d91442102e9b950cb2ddb8db545c4b4e1100e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_sparse_functions.cpp @@ -0,0 +1,67 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include "torch/csrc/Device.h" +#include "torch/csrc/DynamicTypes.h" +#include "torch/csrc/Exceptions.h" +#include "torch/csrc/autograd/python_sparse_functions.h" +#include "torch/csrc/autograd/python_variable.h" +#include "torch/csrc/autograd/utils/wrap_outputs.h" +#include "torch/csrc/autograd/utils/python_arg_parsing.h" +#include "torch/csrc/utils/pycfunction_helpers.h" +#include "torch/csrc/utils/python_arg_parser.h" +#include "torch/csrc/utils/structseq.h" + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +using at::Tensor; +using at::Scalar; +using at::ScalarType; +using at::MemoryFormat; +using at::Generator; +using at::IntArrayRef; +using at::TensorList; + +using namespace torch::autograd::utils; + +namespace torch::autograd { + +// generated forward declarations start here + +${py_forwards} + +static PyMethodDef sparse_functions[] = { + ${py_method_defs} + {NULL} +}; + +static PyObject* THPSparseVariableFunctionsModule = NULL; + +void initSparseFunctions(PyObject* module) { + static struct PyModuleDef def = { + PyModuleDef_HEAD_INIT, + "torch._C._sparse", + NULL, + -1, + sparse_functions + }; + PyObject* sparse = PyModule_Create(&def); + THPSparseVariableFunctionsModule = sparse; + if (!sparse) { + throw python_error(); + } + // steals a reference to sparse + if (PyModule_AddObject(module, "_sparse", sparse) != 0) { + throw python_error(); + } +} + +// generated methods start here + +${py_methods} + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_special_functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_special_functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..bf9e109b4a77352cd85ba828b97d67d329543867 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_special_functions.cpp @@ -0,0 +1,79 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include "torch/csrc/Device.h" +#include "torch/csrc/DynamicTypes.h" +#include "torch/csrc/Exceptions.h" +#include "torch/csrc/autograd/python_special_functions.h" +#include "torch/csrc/autograd/generated/python_return_types.h" +#include "torch/csrc/autograd/python_variable.h" +#include "torch/csrc/autograd/utils/wrap_outputs.h" +#include "torch/csrc/autograd/utils/python_arg_parsing.h" +#include "torch/csrc/autograd/generated/variable_factories.h" +#include "torch/csrc/utils/out_types.h" +#include "torch/csrc/utils/pycfunction_helpers.h" +#include "torch/csrc/utils/python_arg_parser.h" +#include "torch/csrc/utils/structseq.h" +#include "torch/csrc/utils/device_lazy_init.h" + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +using at::Tensor; +using at::Device; +using at::Layout; +using at::Scalar; +using at::ScalarType; +using at::Backend; +using at::OptionalDeviceGuard; +using at::DeviceGuard; +using at::TensorOptions; +using at::IntArrayRef; +using at::Generator; +using at::TensorList; +using at::Dimname; +using at::DimnameList; + +using torch::utils::check_out_type_matches; +using namespace torch::autograd::utils; + +namespace torch::autograd { + +// generated forward declarations start here + +${py_forwards} + +static PyMethodDef special_functions[] = { + ${py_method_defs} + {NULL} +}; + +static PyObject* THPSpecialVariableFunctionsModule = NULL; + +void initSpecialFunctions(PyObject* module) { + static struct PyModuleDef def = { + PyModuleDef_HEAD_INIT, + "torch._C._special", + NULL, + -1, + special_functions + }; + PyObject* special = PyModule_Create(&def); + THPSpecialVariableFunctionsModule = special; + if (!special) { + throw python_error(); + } + // steals a reference to special + if (PyModule_AddObject(module, "_special", special) != 0) { + throw python_error(); + } +} + +// generated methods start here + +${py_methods} + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_torch_functions.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_torch_functions.cpp new file mode 100644 index 0000000000000000000000000000000000000000..c17d1040e1892b6a215a8c4264fe5a5345265bc7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_torch_functions.cpp @@ -0,0 +1,93 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +// Python bindings for torch.* functions implemented through ATen. +// +// The functions are bound as static methods on a class +// torch._C._VariableFunctions which is also aliased as Variable._torch +// and also copied into 'torch' module. + +#include + +// Undefine the copysign macro so that at::copysign works as intended with MSVC +// https://github.com/python/cpython/blob/c60394c7fc9cc09b16e9675a3eeb5844b6d8523f/PC/pyconfig.h#L196 +#ifdef _MSC_VER +#undef copysign +#endif // _MSC_VER + +#include "torch/csrc/autograd/python_torch_functions.h" +#include "torch/csrc/autograd/python_variable.h" +#include "torch/csrc/autograd/utils/wrap_outputs.h" +#include "torch/csrc/Dtype.h" +#include "torch/csrc/DynamicTypes.h" +#include "torch/csrc/Exceptions.h" +#include "torch/csrc/utils/out_types.h" +#include "torch/csrc/utils/pybind.h" +#include "torch/csrc/utils/pycfunction_helpers.h" +#include "torch/csrc/utils/python_arg_parser.h" +#include "torch/csrc/utils/tensor_layouts.h" +#include "torch/csrc/utils/tensor_new.h" +#include "torch/csrc/utils/tensor_numpy.h" +#include "torch/csrc/jit/frontend/tracer.h" +#include "torch/csrc/autograd/generated/variable_factories.h" +#include "torch/csrc/utils/structseq.h" +#include "torch/csrc/utils/device_lazy_init.h" +#include "torch/csrc/autograd/generated/python_return_types.h" + +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#endif + +#include +#include +#include +#include + +using at::Tensor; +using at::Device; +using at::Layout; +using at::Scalar; +using at::ScalarType; +using at::Backend; +using at::OptionalDeviceGuard; +using at::DeviceGuard; +using at::TensorOptions; +using at::IntArrayRef; +using at::Generator; +using at::TensorList; +using at::Dimname; +using at::DimnameList; +using at::ArrayRef; + +using torch::utils::check_out_type_matches; +using namespace torch::autograd::utils; + +// NOTE: See [Sharded File] comment in VariableType + +namespace torch::autograd { + +// generated forward declarations start here + +${py_forwards} + +static PyMethodDef torch_functions_shard[] = { + ${py_method_defs} +}; + +void gatherTorchFunctions${shard_id}(std::vector &torch_functions) { + constexpr size_t num_functions = sizeof(torch_functions_shard) / sizeof(torch_functions_shard[0]); + torch_functions.insert( + torch_functions.end(), + torch_functions_shard, + torch_functions_shard + num_functions); +} + +// generated methods start here + +${py_methods} + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_variable_methods.cpp b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_variable_methods.cpp new file mode 100644 index 0000000000000000000000000000000000000000..2260f8cb2245f43567095a31b271063a28796e9b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/python_variable_methods.cpp @@ -0,0 +1,1338 @@ +#define TORCH_ASSERT_ONLY_METHOD_OPERATORS +// ${generated_comment} + +#include + +// Undefine the copysign macro so that at::copysign works as intended with MSVC +// https://github.com/python/cpython/blob/c60394c7fc9cc09b16e9675a3eeb5844b6d8523f/PC/pyconfig.h#L196 +#ifdef _MSC_VER +#undef copysign +#endif // _MSC_VER + +#include "torch/csrc/DynamicTypes.h" +#include "torch/csrc/Exceptions.h" +#include "torch/csrc/Size.h" +#include "torch/csrc/autograd/generated/VariableType.h" +#include "torch/csrc/autograd/python_variable.h" +#include "torch/csrc/autograd/utils/python_arg_parsing.h" +#include "torch/csrc/autograd/utils/error_messages.h" +#include "torch/csrc/autograd/utils/wrap_outputs.h" +#include "torch/csrc/jit/frontend/tracer.h" +#ifdef USE_CUDA +#include "torch/csrc/cuda/Event.h" +#endif +#include "torch/csrc/utils/device_lazy_init.h" +#include +#include "torch/csrc/utils/object_ptr.h" +#include "torch/csrc/utils/pycfunction_helpers.h" +#include "torch/csrc/utils/python_arg_parser.h" +#include "torch/csrc/utils/python_numbers.h" +#include "torch/csrc/utils/python_strings.h" +#include "torch/csrc/utils/tensor_apply.h" +#include "torch/csrc/utils/tensor_list.h" +#include "torch/csrc/utils/tensor_new.h" +#include "torch/csrc/utils/tensor_numpy.h" +#include "torch/csrc/utils/tensor_types.h" +#include "torch/csrc/autograd/generated/python_return_types.h" + +#include +#include +#include +#include "c10/core/Stream.h" + +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +$ops_headers +#include +#endif + +using at::device_of; +using at::OptionalDeviceGuard; +using at::Scalar; +using at::ScalarType; +using at::Tensor; +using c10::Stream; +using namespace torch::autograd::utils; + +namespace torch::autograd { + +static PyObject * THPVariable__is_view(PyObject *self, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "_is_view", args); + } + auto& self_ = THPVariable_Unpack(self); + if (self_.is_view()) { + Py_RETURN_TRUE; + } else { + Py_RETURN_FALSE; + } + END_HANDLE_TH_ERRORS +} + +// implemented on the python object bc no support for first-class functions in native_functions.yaml +// See: ATen/native/README.md for more context +static PyObject * THPVariable_apply_(PyObject* self, PyObject* arg) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + auto args = py::make_tuple(py::handle(arg)); + return handle_torch_function(self, "apply_", args.ptr()); + } + auto& self_ = THPVariable_Unpack(self); + if (self_.requires_grad()) { + throw std::runtime_error( + "Can't call apply_() on Variable that requires grad. Use " + "var.detach().apply_() instead."); + } + return THPVariable_Wrap(torch::utils::apply_(self_, arg)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_size(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "size(int64_t? dim=None)", + "size(Dimname dim)", + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<3> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + if (r.idx == 0) { + if (!r.toInt64Optional(0).has_value()) { + return THPSize_NewFromSymSizes(self_); + } + if (jit::tracer::isTracing()) { + // will error out if a tensor has symints + return wrap(jit::tracer::getSizeOf(self_, r.toInt64(0))); + } else { + return torch::toPyObject(self_.sym_size(r.toInt64(0))); + } + } else if (r.idx == 1) { + if (jit::tracer::isTracing()) { + TORCH_INTERNAL_ASSERT(false, "NYI: Named tensors w/ JIT"); + } + return wrap(self_.size(r.dimname(0))); + } + Py_RETURN_NONE; + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_stride(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "stride(int64_t? dim=None)", + "stride(Dimname dim)", + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<3> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + if (r.idx == 0) { + if (r.toInt64Optional(0).has_value()) { + return torch::toPyObject(self_.sym_stride(r.toInt64(0))); + } + // yes, this is called strides in ATen. + at::SymIntArrayRef strides = self_.sym_strides(); + // we can't do the normal wrapping here because IntArrayRef maps to both + // torch.Size and tuple in python + // TODO: consider factoring this out + THPObjectPtr tuple(PyTuple_New(static_cast(strides.size()))); + if (!tuple) throw python_error(); + for (size_t i = 0; i != strides.size(); i++) { + PyObject* s = torch::toPyObject(strides[i]); + if (!s) throw python_error(); + PyTuple_SET_ITEM(tuple.get(), i, s); + } + return tuple.release(); + } else if (r.idx == 1) { + return wrap(self_.stride(r.dimname(0))); + } + Py_RETURN_NONE; + END_HANDLE_TH_ERRORS +} + +// implemented on the python object to avoid dispatch overhead +static PyObject * THPVariable_get_device(PyObject* self_, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self_)) { + return handle_torch_function(self_, "get_device", args, nullptr); + } + auto& self = THPVariable_Unpack(self_); + return wrap(self.get_device()); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_has_names(PyObject* self_, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self_)) { + return handle_torch_function(self_, "has_names", args); + } + auto& self = THPVariable_Unpack(self_); + return wrap(self.has_names()); + END_HANDLE_TH_ERRORS +} + +// implemented on the python object to avoid dispatch overhead +static PyObject * THPVariable_data_ptr(PyObject* self_, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self_)) { + return handle_torch_function(self_, "data_ptr", args); + } + auto& self = THPVariable_Unpack(self_); + return wrap(self.data_ptr()); + END_HANDLE_TH_ERRORS +} + +// implemented on the python object to avoid dispatch overhead +static PyObject * THPVariable_storage_offset(PyObject* self_, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self_)) { + return handle_torch_function(self_, "storage_offset"); + } + auto& self = THPVariable_Unpack(self_); + return py::cast(self.sym_storage_offset()).release().ptr(); + END_HANDLE_TH_ERRORS +} + +// implemented on the python object to avoid dispatch overhead +static PyObject * THPVariable_dim(PyObject* self, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "dim", args); + } + auto& self_ = THPVariable_Unpack(self); + return THPUtils_packInt64(self_.dim()); + END_HANDLE_TH_ERRORS +} + +// implemented on the python object to avoid dispatch overhead +static PyObject * THPVariable_numel(PyObject* self, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "numel", args); + } + auto& self_ = THPVariable_Unpack(self); + if (jit::tracer::isTracing()) { + return wrap(jit::tracer::getNumelOf(self_)); + } else { + return py::cast(self_.sym_numel()).release().ptr(); + } + END_HANDLE_TH_ERRORS +} + +static Tensor dispatch_contiguous(const Tensor & self, at::MemoryFormat memory_format) { + pybind11::gil_scoped_release no_gil; + OptionalDeviceGuard device_guard(device_of(self)); + return self.contiguous(memory_format); +} + +static PyObject * THPVariable_contiguous(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "contiguous(*, MemoryFormat memory_format=contiguous_format)", + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto& self_ = THPVariable_Unpack(self); + auto memory_format = r.memoryformat(0); + // avoids touching the GIL or current device if self is already contiguous + if (self_.is_contiguous_or_false(memory_format)) { + // NOTE: this logic is duplicated from VariableType.cpp. Since we need to + // record this call to contiguous() in the trace regardless of whether + // we actually call contiguous here, we need to record this information + // manually. + if (jit::tracer::isTracing()) { + const auto& tracer_state = jit::tracer::getTracingState(); + auto op_name = c10::Symbol::fromQualString("aten::contiguous"); + auto node = tracer_state->createNode(op_name, /*num_outputs=*/0); + jit::tracer::recordSourceLocation(node); + jit::tracer::addInputs(node, "self", self_); + jit::tracer::addInputs(node, "memory_format", memory_format); + tracer_state->insertNode(node); + jit::tracer::addOutput(node, self_); + } + Py_INCREF(self); + return self; + } + return THPVariable_Wrap(dispatch_contiguous(self_, memory_format)); + END_HANDLE_TH_ERRORS +} + +static Tensor dispatch_copy_(const Tensor & self, const Tensor & other, bool non_blocking) { + pybind11::gil_scoped_release no_gil; + OptionalDeviceGuard device_guard(device_of(self)); + return self.copy_(other, non_blocking); +} + +static void maybe_warn_requires_grad(const Tensor & self) { + if (at::GradMode::is_enabled() && self.requires_grad()) { + TORCH_WARN_ONCE("Converting a tensor with requires_grad=True to a scalar may lead to unexpected behavior.\n" + "Consider using tensor.detach() first."); + } +} + + static PyObject * THPVariable_copy_(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "copy_(Tensor other, bool non_blocking=False)", + "copy_(Tensor other, bool async=False)|deprecated" + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<2> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + return THPVariable_Wrap(dispatch_copy_(self_, r.tensor(0), r.toBool(1))); + END_HANDLE_TH_ERRORS +} + +template +static T dispatch_to(const Tensor & self) { + pybind11::gil_scoped_release no_gil; + OptionalDeviceGuard device_guard(device_of(self)); + TORCH_CHECK_VALUE(self.sym_numel() == 1, "only one element tensors can be converted to Python scalars"); + return self.template item(); +} + +static PyObject * THPVariable_float_scalar(PyObject* self, PyObject* args) { + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "__float__", args); + } + jit::tracer::warn("Converting a tensor to a Python float", jit::tracer::WARN_PYTHON_DATAFLOW); + auto& self_ = THPVariable_Unpack(self); + maybe_warn_requires_grad(self_); + return wrap(dispatch_to(self_)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_complex_scalar(PyObject* self, PyObject* args) { + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "__complex__", args); + } + jit::tracer::warn("Converting a tensor to a Python complex", jit::tracer::WARN_PYTHON_DATAFLOW); + auto& self_ = THPVariable_Unpack(self); + maybe_warn_requires_grad(self_); + return wrap(dispatch_to>(self_)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_integral_scalar(PyObject* self, PyObject* args) { + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "__int__", args); + } + jit::tracer::warn("Converting a tensor to a Python integer", jit::tracer::WARN_PYTHON_DATAFLOW); + auto& self_ = THPVariable_Unpack(self); + if (isFloatingType(self_.scalar_type())) { + // we can't dispatch to item here because we want to avoid ATen overflow checks; + // the python integral type (long in python2) can't overflow. + return THPUtils_packDoubleAsInt(dispatch_to(self_)); + } else { + return wrap(dispatch_to(self_)); + } + END_HANDLE_TH_ERRORS +} + +// This is the __index__ function in Python which is similar to __int__, but +// called when used as a slice. +static PyObject * THPVariable_index_scalar(PyObject* self, PyObject* args) { + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "__index__", args); + } + auto& self_ = THPVariable_Unpack(self); + // TODO: change the condition to `self_.dim() != 0` once we expose scalars + // in PyTorch. + if (!isIntegralType(self_.scalar_type(), /*includeBool=*/true) || self_.sym_numel() != 1) { + throw TypeError("only integer tensors of a single element can be converted to an index"); + } + return wrap(dispatch_to(self_)); + END_HANDLE_TH_ERRORS +} + +static Tensor dispatch_invert(const Tensor & self) { + pybind11::gil_scoped_release no_gil; + OptionalDeviceGuard device_guard(device_of(self)); + return self.bitwise_not(); +} + +static PyObject * THPVariable_invert(PyObject* self, PyObject* args) { + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "__invert__", args); + } + auto& self_ = THPVariable_Unpack(self); + if (!isIntegralType(self_.scalar_type(), /*includeBool=*/true)) { + throw TypeError("~ (operator.invert) is only implemented on integer and Boolean-type tensors"); + } + return THPVariable_Wrap(dispatch_invert(self_)); + END_HANDLE_TH_ERRORS +} + +static Tensor dispatch_to(const Tensor & self, Device device, bool non_blocking, bool copy, std::optional optional_memory_format) { + pybind11::gil_scoped_release no_gil; + // NOTE: this is where we record aten::to in the graph during tracing. However, the behavior of aten::to + // is different with respect to TensorOptions fields that are not present: aten::to inherits fields that + // are missing from the self argument while the tracer assumes that they should be populated with the + // default values (eg. float for scalar type). By explicitly copying over the tensor options here we fully + // specify all tensor options and thus record the proper trace + return self.to(self.options().device(device).memory_format(optional_memory_format), non_blocking, copy); +} + +static Tensor dispatch_to(const Tensor & self, bool non_blocking, bool copy, std::optional optional_memory_format) { + pybind11::gil_scoped_release no_gil; + return self.to(self.options().memory_format(optional_memory_format), non_blocking, copy); +} + +static Tensor dispatch_to(const Tensor & self, ScalarType dtype, bool non_blocking, bool copy, std::optional optional_memory_format) { + pybind11::gil_scoped_release no_gil; + // TODO: Make this call the TensorOptions version, maybe? + return self.to(dtype, non_blocking, copy, optional_memory_format); +} + +static Tensor dispatch_to(const Tensor & self, Device device, ScalarType dtype, bool non_blocking, bool copy, std::optional optional_memory_format) { + pybind11::gil_scoped_release no_gil; + // TODO: Make this call the TensorOptions version, maybe? + return self.to(device, dtype, non_blocking, copy, optional_memory_format); +} + +static PyObject * THPVariable_cpu(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "cpu(*, MemoryFormat? memory_format=None)" + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_Wrap(dispatch_to(self_, at::Device(at::DeviceType::CPU), false, false, opt_memory_format)); + END_HANDLE_TH_ERRORS +} + +static Tensor dispatch_nonzero(const Tensor & self) { + pybind11::gil_scoped_release no_gil; + OptionalDeviceGuard device_guard(device_of(self)); + return self.nonzero(); +} + +static std::vector dispatch_nonzero_numpy(const Tensor & self) { + pybind11::gil_scoped_release no_gil; + OptionalDeviceGuard device_guard(device_of(self)); + return self.nonzero_numpy(); +} + +static PyObject * THPVariable_nonzero(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "nonzero()", + "nonzero(*, bool as_tuple)", + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<2> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + if (r.idx == 0 || (r.idx == 1 && !r.toBool(0))) { + return wrap(dispatch_nonzero(self_)); + } else { + return wrap(dispatch_nonzero_numpy(self_)); + } + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_cuda(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "cuda(Device? device=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)", + "cuda(Device? device=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated" + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<3> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto device = r.isNone(0) ? at::Device(at::DeviceType::CUDA) : r.device(0); + auto opt_memory_format = r.memoryformatOptional(2); + TORCH_CHECK(device.is_cuda(), "Invalid device, must be cuda device"); + torch::utils::device_lazy_init(at::kCUDA); + return THPVariable_Wrap(dispatch_to(self_, device, r.toBool(1), false, opt_memory_format)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_mtia(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "mtia(Device? device=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)", + "mtia(Device? device=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated" + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<3> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if (r.has_torch_function()) { + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto device = r.isNone(0) ? at::Device(at::DeviceType::MTIA) : r.device(0); + auto opt_memory_format = r.memoryformatOptional(2); + TORCH_CHECK(device.is_mtia(), "Invalid device, must be MTIA device"); + torch::utils::device_lazy_init(at::kMTIA); + return THPVariable_Wrap(dispatch_to(self_, device, r.toBool(1), false, opt_memory_format)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_xpu(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "xpu(Device? device=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)", + "xpu(Device? device=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated" + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<3> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if (r.has_torch_function()) { + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto device = r.isNone(0) ? at::Device(at::DeviceType::XPU) : r.device(0); + auto opt_memory_format = r.memoryformatOptional(2); + TORCH_CHECK(device.is_xpu(), "Invalid device, must be xpu device"); + torch::utils::device_lazy_init(at::kXPU); + return THPVariable_Wrap(dispatch_to(self_, device, r.toBool(1), false, opt_memory_format)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_ipu(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "ipu(Device? device=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)", + "ipu(Device? device=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated" + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<3> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if (r.has_torch_function()) { + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto device = r.isNone(0) ? at::Device(at::DeviceType::IPU) : r.device(0); + auto opt_memory_format = r.memoryformatOptional(2); + TORCH_CHECK(device.is_ipu(), "Invalid device, must be ipu device"); + return THPVariable_Wrap(dispatch_to(self_, device, r.toBool(1), false, opt_memory_format)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_to_type(PyObject* self, ScalarType scalarType, std::optional optional_memory_format) { + HANDLE_TH_ERRORS + auto& self_ = THPVariable_Unpack(self); + return THPVariable_Wrap(dispatch_to(self_, scalarType, false, false, optional_memory_format)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_byte(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "byte(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::Byte, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_char(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "char(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::Char, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_double(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "double(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::Double, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_float(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "float(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::Float, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_cdouble(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "cdouble(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::ComplexDouble, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_cfloat(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "cfloat(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::ComplexFloat, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_half(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "half(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::Half, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_int(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "int(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::Int, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_long(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "long(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::Long, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_short(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "short(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::Short, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_bool(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "bool(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::Bool, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_bfloat16(PyObject* self, PyObject* args, PyObject* kwargs) { + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "bfloat16(*, MemoryFormat? memory_format=None)" + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + auto opt_memory_format = r.memoryformatOptional(0); + return THPVariable_to_type(self, ScalarType::BFloat16, opt_memory_format); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_element_size(PyObject* self, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "element_size", args); + } + auto& self_ = THPVariable_Unpack(self); + return THPUtils_packInt64(self_.element_size()); + END_HANDLE_TH_ERRORS +} + +// implemented on the python object bc PyObjects not declarable in native_functions.yaml +// See: ATen/native/README.md for more context +static PyObject * THPVariable_numpy(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "numpy(*, bool force=False)" + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if (r.has_torch_function()) { + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + jit::tracer::warn("Converting a tensor to a NumPy array", jit::tracer::WARN_PYTHON_DATAFLOW); + return torch::utils::tensor_to_numpy(self_, r.toBool(0)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_requires_grad_(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "requires_grad_(bool requires_grad=True)", + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + // temporary hack to improve functorch UX. + const auto& functorch_tls = at::functorch::functorchTLSAccessor(); + if (functorch_tls) { + functorch_tls->checkSupportsInplaceRequiresGrad(); + } + + auto requires_grad = r.toBool(0); + // should we throw if requires_grad is true? var.requires_grad = True throws here + // but it's nice to let this be a no-op. + if (!self_.is_leaf() && !requires_grad) { + throw std::runtime_error(autograd::utils::requires_grad_leaf_error(requires_grad)); + } + if (requires_grad && ! isDifferentiableType(at::typeMetaToScalarType(self_.dtype()))) { + throw std::runtime_error("only Tensors of floating point dtype can require gradients"); + } + self_.set_requires_grad(requires_grad); + return THPVariable_Wrap(self_); + END_HANDLE_TH_ERRORS +} + +static inline bool dispatch_is_contiguous(const Tensor & self, MemoryFormat memory_format) { + return self.is_contiguous(memory_format); +} + +// implemented on the python object to avoid dispatch overhead +static PyObject * THPVariable_is_contiguous(PyObject* self_, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "is_contiguous(*, MemoryFormat memory_format=contiguous_format)", + }); + ParsedArgs<1> parsed_args; + auto r = parser.parse(self_, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self_, args, kwargs, reinterpret_cast(Py_TYPE(self_)), "torch.Tensor"); + } + + auto memory_format = r.memoryformat(0); + auto& self = THPVariable_Unpack(self_); + return wrap(dispatch_is_contiguous(self, memory_format)); + END_HANDLE_TH_ERRORS +} + +// implemented on the python object to avoid dispatch overhead +static PyObject * THPVariable_item(PyObject* self, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "item", args); + } + jit::tracer::warn("Converting a tensor to a Python number", jit::tracer::WARN_PYTHON_DATAFLOW); + auto& self_ = THPVariable_Unpack(self); + auto dispatch_item_ = [](const Tensor& self) -> at::Scalar { + pybind11::gil_scoped_release no_gil; + return self.item(); + }; + return py::cast(dispatch_item_(self_)).release().ptr(); + END_HANDLE_TH_ERRORS +} + +// implemented on the python object bc no support for first class functions in native_functions.yaml +// See: ATen/native/README.md for more context +static PyObject * THPVariable_map_(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ "map_(Tensor other, PyObject* callable)" }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<2> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + Variable other = r.tensor(0); + if (self_.requires_grad() || other.requires_grad()) { + throw std::runtime_error( + "Can't call map_() on Variable that requires grad. Use " + "var.detach().map_() instead."); + } + TORCH_CHECK( + !self_.unsafeGetTensorImpl()->is_python_dispatch() && !other.unsafeGetTensorImpl()->is_python_dispatch(), + ".map_ is not supported for tensor subclasses."); + + return THPVariable_Wrap(torch::utils::map_(self_, other, r.pyobject(1))); + END_HANDLE_TH_ERRORS +} + +// implemented on the python object bc no support for first class functions in native_functions.yaml +// See: ATen/native/README.md for more context +static PyObject * THPVariable_map2_(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ "map2_(Tensor x, Tensor y, PyObject* callable)" }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<3> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + Variable x = r.tensor(0); + Variable y = r.tensor(1); + if (self_.requires_grad() || x.requires_grad() || y.requires_grad()) { + throw std::runtime_error( + "Can't call map2_() on Variable that requires grad. Use " + "var.detach().map2_() instead."); + } + TORCH_CHECK( + !x.unsafeGetTensorImpl()->is_python_dispatch() && !y.unsafeGetTensorImpl()->is_python_dispatch(), + ".map2_ is not supported for tensor subclasses."); + return THPVariable_Wrap(torch::utils::map2_(self_, x, y, r.pyobject(2))); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_new(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "new", args, kwargs); + } + auto& self_ = THPVariable_Unpack(self); + OptionalDeviceGuard device_guard(device_of(self_)); + return THPVariable_Wrap(torch::utils::legacy_tensor_new(legacyExtractDispatchKey(self_), self_.scalar_type(), args, kwargs)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_new_tensor(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "new_tensor", args, kwargs); + } + auto& self_ = THPVariable_Unpack(self); + OptionalDeviceGuard device_guard(device_of(self_)); + return THPVariable_Wrap(torch::utils::new_tensor(legacyExtractDispatchKey(self_), self_.scalar_type(), args, kwargs)); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_storage(PyObject* self, PyObject* arg) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "untyped_storage"); + } + auto& self_ = THPVariable_Unpack(self); + return createPyObject(self_.storage()); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_to(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "to(Device device=None, ScalarType dtype=None, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)", + "to(ScalarType dtype, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)", + "to(Tensor tensor, bool non_blocking=False, bool copy=False, *, MemoryFormat? memory_format=None)", + }); + ParsedArgs<5> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + if (r.has_torch_function()) { + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + auto parsed = parse_to_conversion(r, /*allow_copy*/ true); + auto& device = std::get<0>(parsed); + auto& scalarType = std::get<1>(parsed); + auto non_blocking = std::get<2>(parsed); + auto copy = std::get<3>(parsed); + auto opt_memory_format = std::get<4>(parsed); + auto& self_ = THPVariable_Unpack(self); + torch::utils::maybe_initialize_device(device); + if (!device && !scalarType && !copy && !opt_memory_format.has_value()) { + Py_INCREF(self); + return self; + } else if (!device && !scalarType) { + return THPVariable_Wrap( + dispatch_to(self_, non_blocking, copy, opt_memory_format)); + } else if (!device) { + return THPVariable_Wrap(dispatch_to(self_, *scalarType, non_blocking, copy, opt_memory_format)); + } else if (!scalarType) { + return THPVariable_Wrap(dispatch_to(self_, *device, non_blocking, copy, opt_memory_format)); + } else { + return THPVariable_Wrap(dispatch_to(self_, *device, *scalarType, non_blocking, copy, opt_memory_format)); + } + Py_RETURN_NONE; + END_HANDLE_TH_ERRORS +} + +// implemented on the python object b/c arbitrarily nested list not declarable in native_functions.yaml +// See: ATen/native/README.md for more context +static PyObject * THPVariable_tolist(PyObject* self, PyObject* args) +{ + HANDLE_TH_ERRORS + if (check_has_torch_function(self)) { + return handle_torch_function(self, "tolist", args); + } + jit::tracer::warn("Converting a tensor to a Python list", jit::tracer::WARN_PYTHON_DATAFLOW); + auto self_ = THPVariable_Unpack(self); + return torch::utils::tensor_to_list(self_); + END_HANDLE_TH_ERRORS +} + +static PyObject * THPVariable_type(PyObject* self, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS + static PythonArgParser parser({ + "type(PyObject* dtype=None, bool non_blocking=False, *, MemoryFormat? memory_format=None)", + "type(PyObject* dtype=None, bool async=False, *, MemoryFormat? memory_format=None)|deprecated" + }); + auto& self_ = THPVariable_Unpack(self); + ParsedArgs<3> parsed_args; + auto r = parser.parse(self, args, kwargs, parsed_args); + + if(r.has_torch_function()){ + return handle_torch_function(r, self, args, kwargs, THPVariableClass, "torch.Tensor"); + } + + if (r.isNone(0)) { + return THPUtils_packString(torch::utils::options_to_string(self_.options())); + } + auto obj = r.pyobject(0); + auto opt_memory_format = r.memoryformatOptional(2); + std::string type_name; + bool is_dtype = false; + if (PyType_Check(obj)) { + if (obj == THPVariableClass) { + type_name = "torch.Tensor"; + } else { + type_name = ((PyTypeObject*)obj)->tp_name; + } + } else if (THPUtils_checkString(obj)) { + type_name = THPUtils_unpackString(obj); + } else if (THPDtype_Check(obj)) { + is_dtype = true; + } else { + throw TypeError("dtype must be a type, str, or dtype object"); + } + Device device = self_.device(); + if (is_dtype) { + auto scalar_type = r.scalartype(0); + return THPVariable_Wrap(dispatch_to(self_, scalar_type, /*non_blocking=*/ r.toBool(1), /*copy=*/ false, opt_memory_format)); + } + at::TensorOptions options = torch::utils::options_from_string(type_name); + auto scalar_type = at::typeMetaToScalarType(options.dtype()); + auto device_type = options.device().type(); + if (device_type != device.type()) { + device = at::Device(device_type); + } + torch::utils::maybe_initialize_device(device); + return THPVariable_Wrap(dispatch_to(self_, device, scalar_type, /*non_blocking=*/ r.toBool(1), /*copy=*/ false, opt_memory_format)); + END_HANDLE_TH_ERRORS +} + +// generated methods start here + +${py_methods} + +static PyObject * THPVariable_bool_scalar(PyObject* self, PyObject* args) { + if (check_has_torch_function(self)) { + HANDLE_TH_ERRORS + return handle_torch_function(self, "__bool__", args); + END_HANDLE_TH_ERRORS + } + jit::tracer::warn("Converting a tensor to a Python boolean", jit::tracer::WARN_PYTHON_DATAFLOW); + return THPVariable_is_nonzero(self, args); +} + +static PyObject * THPVariable___eq__(PyObject* self_, PyObject* args, PyObject* kwargs) +{ + HANDLE_TH_ERRORS +#ifdef USE_NUMPY + if (torch::utils::is_numpy_available()) { + static PythonArgParser parser({ + "__eq__(PyObject* other)", + }, /*traceable=*/true); + + ParsedArgs<1> parsed_args; + auto _r = parser.parse(self_, args, kwargs, parsed_args); + if(_r.has_torch_function()) { + return handle_torch_function(_r, self_, args, kwargs, THPVariableClass, "torch.Tensor"); + } + switch (_r.idx) { + case 0: { + auto other = _r.pyobject(0); + if (PyArray_Check(other)) { + auto other_tensor = torch::utils::tensor_from_numpy(other); + auto dispatch_eq = [](const at::Tensor & self, const at::Tensor & other) -> at::Tensor { + pybind11::gil_scoped_release no_gil; + return self.eq(other); + }; + const Tensor& self = THPVariable_Unpack(self_); + return wrap(dispatch_eq(self, other_tensor)); + } + } + } + } +#endif + return THPVariable_eq(self_, args, kwargs); + Py_RETURN_NONE; + END_HANDLE_TH_ERRORS +} + +// Wrapper converts a raised TypeError into returning NotImplemented +// Used to implement binary arithmetic operators +template +static PyObject * TypeError_to_NotImplemented_(PyObject* self, PyObject* args, PyObject* kwargs) { + + PyObject* ret = Func(self, args, kwargs); + if (!ret && PyErr_ExceptionMatches(PyExc_TypeError)) { + PyErr_Clear(); + Py_INCREF(Py_NotImplemented); + ret = Py_NotImplemented; + } + return ret; +} + +// set_ has to be defined in the template because the c10::Storage object +// does not have a type, and we need to make sure the Python storage object's +// type matches the tensor's type +static PyObject* THPVariable_set_( + PyObject* self_, + PyObject* args, + PyObject* kwargs) { + HANDLE_TH_ERRORS + const Tensor& self = THPVariable_Unpack(self_); + static PythonArgParser parser( + { + "set_()", + "set_(Storage source)", + "set_(Storage source, SymInt storage_offset, SymIntArrayRef size, SymIntArrayRef stride=None)", + "set_(Tensor source)", + "set_(Tensor source, SymInt storage_offset, SymIntArrayRef size, SymIntArrayRef stride=None)", + }, + /*traceable=*/false); + + ParsedArgs<4> parsed_args; + auto _r = parser.parse(args, kwargs, parsed_args); + + switch (_r.idx) { + case 0: { + // aten::set_(Tensor(a!) self) -> Tensor(a!) + auto dispatch_set_ = [](const Tensor& self) -> Tensor { + pybind11::gil_scoped_release no_gil; + return self.set_(); + }; + return wrap(dispatch_set_(self)); + } + case 1: { + // aten::set_.source_Storage(Tensor(a!) self, Storage source) -> + // Tensor(a!) + at::ScalarType storage_scalar_type{}; + bool is_typed_storage = true; + at::Storage storage = _r.storage(0, storage_scalar_type, is_typed_storage); + TORCH_CHECK(storage_scalar_type == self.dtype() || !is_typed_storage, + "Expected a Storage of type ", self.dtype(), + " or an UntypedStorage, but got type ", storage_scalar_type, + " for argument 1 'storage'"); + auto dispatch_set_ = [](const Tensor& self, Storage source) -> Tensor { + pybind11::gil_scoped_release no_gil; + return self.set_(std::move(source)); + }; + return wrap(dispatch_set_(self, storage)); + } + case 2: { + // aten::set_.source_Storage_storage_offset(Tensor(a!) self, Storage + // source, int storage_offset, int[] size, int[] stride=[]) -> Tensor(a!) + at::ScalarType storage_scalar_type{}; + bool is_typed_storage = true; + at::Storage storage = _r.storage(0, storage_scalar_type, is_typed_storage); + TORCH_CHECK(storage_scalar_type == self.dtype() || !is_typed_storage, + "Expected a Storage of type ", self.dtype(), + " or an UntypedStorage, but got type ", storage_scalar_type, + " for argument 1 'storage'"); + auto dispatch_set_ = [](const Tensor& self, + Storage source, + c10::SymInt storage_offset, + c10::SymIntArrayRef size, + c10::SymIntArrayRef stride) -> Tensor { + pybind11::gil_scoped_release no_gil; + return self.set__symint(std::move(source), std::move(storage_offset), size, stride); + }; + return wrap(dispatch_set_( + self, storage, _r.toSymInt(1), _r.symintlist(2), _r.symintlist(3))); + } + case 3: { + // aten::set_.source_Tensor(Tensor(a!) self, Tensor source) -> Tensor(a!) + auto dispatch_set_ = [](const Tensor& self, const Tensor& source) -> Tensor { + TORCH_CHECK(source.dtype() == self.dtype(), "Could not set tensor of type ", source.dtype(), " to a tensor of type ", self.dtype()); + pybind11::gil_scoped_release no_gil; + return self.set_(source); + }; + return wrap(dispatch_set_(self, _r.tensor(0))); + } + case 4: { + // aten::set_.source_Tensor_storage_offset(Tensor(a!) self, Tensor + // source, int storage_offset, int[] size, int[] stride=[]) -> Tensor(a!) + at::Tensor storage = _r.tensor(0); + auto dispatch_set_ = [](const Tensor& self, + const Tensor& source, + c10::SymInt storage_offset, + c10::SymIntArrayRef size, + c10::SymIntArrayRef stride) -> Tensor { + pybind11::gil_scoped_release no_gil; + return self.set__symint(source, std::move(storage_offset), size, stride); + }; + return wrap(dispatch_set_( + self, storage, _r.toSymInt(1), _r.symintlist(2), _r.symintlist(3))); + } + } + Py_RETURN_NONE; + END_HANDLE_TH_ERRORS +} + +// XXX: ops that are bound here are not exposed to the C++ api nor the JIT. +// Any new ops added here should be accompanied with a comment why they are not +// being registered through native_functions.yaml, and be tagged cpp / JIT +PyMethodDef variable_methods[] = { + // These magic methods are all implemented on python object to wrap NotImplementedError + {"__add__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__radd__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__iadd__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__rmul__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__mul__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__imul__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__sub__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__isub__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__div__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__truediv__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__floordiv__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__idiv__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__ifloordiv__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__mod__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__imod__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__eq__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__ne__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__lt__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__le__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__gt__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__ge__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__rand__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__ror__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__rxor__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"__bool__", THPVariable_bool_scalar, METH_NOARGS, nullptr}, + {"__float__", THPVariable_float_scalar, METH_NOARGS, nullptr}, + {"__complex__", THPVariable_complex_scalar, METH_NOARGS, nullptr}, + {"__int__", THPVariable_integral_scalar, METH_NOARGS, nullptr}, + {"__long__", THPVariable_integral_scalar, METH_NOARGS, nullptr}, + {"__index__", THPVariable_index_scalar, METH_NOARGS, nullptr}, + {"__nonzero__", THPVariable_bool_scalar, METH_NOARGS, nullptr}, + {"__invert__", THPVariable_invert, METH_NOARGS, nullptr}, + {"__matmul__", castPyCFunctionWithKeywords(TypeError_to_NotImplemented_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"_is_view", THPVariable__is_view, METH_NOARGS, nullptr}, + {"apply_", THPVariable_apply_, METH_O, nullptr}, + {"bfloat16", castPyCFunctionWithKeywords(THPVariable_bfloat16), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"byte", castPyCFunctionWithKeywords(THPVariable_byte), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"char", castPyCFunctionWithKeywords(THPVariable_char), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"contiguous", castPyCFunctionWithKeywords(THPVariable_contiguous), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"copy_", castPyCFunctionWithKeywords(THPVariable_copy_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"cpu", castPyCFunctionWithKeywords(THPVariable_cpu), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"cuda", castPyCFunctionWithKeywords(THPVariable_cuda), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"mtia", castPyCFunctionWithKeywords(THPVariable_mtia), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"xpu", castPyCFunctionWithKeywords(THPVariable_xpu), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"ipu", castPyCFunctionWithKeywords(THPVariable_ipu), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"data_ptr", THPVariable_data_ptr, METH_NOARGS, nullptr}, + {"dim", THPVariable_dim, METH_NOARGS, nullptr}, + {"has_names", THPVariable_has_names, METH_NOARGS, nullptr}, + {"double", castPyCFunctionWithKeywords(THPVariable_double), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"cdouble", castPyCFunctionWithKeywords(THPVariable_cdouble), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"element_size", THPVariable_element_size, METH_NOARGS, nullptr}, + {"float", castPyCFunctionWithKeywords(THPVariable_float), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"cfloat", castPyCFunctionWithKeywords(THPVariable_cfloat), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"get_device", THPVariable_get_device, METH_NOARGS, nullptr}, + {"bool", castPyCFunctionWithKeywords(THPVariable_bool), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"half", castPyCFunctionWithKeywords(THPVariable_half), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"int", castPyCFunctionWithKeywords(THPVariable_int), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"is_contiguous", castPyCFunctionWithKeywords(THPVariable_is_contiguous), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"item", THPVariable_item, METH_NOARGS, nullptr}, + {"long", castPyCFunctionWithKeywords(THPVariable_long), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"map_", castPyCFunctionWithKeywords(THPVariable_map_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"map2_", castPyCFunctionWithKeywords(THPVariable_map2_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"ndimension", THPVariable_dim, METH_NOARGS, nullptr}, + {"nelement", THPVariable_numel, METH_NOARGS, nullptr}, + {"new", castPyCFunctionWithKeywords(THPVariable_new), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"new_tensor", castPyCFunctionWithKeywords(THPVariable_new_tensor), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"nonzero", castPyCFunctionWithKeywords(THPVariable_nonzero), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"numel", THPVariable_numel, METH_NOARGS, nullptr}, + {"numpy", castPyCFunctionWithKeywords(THPVariable_numpy), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"requires_grad_", castPyCFunctionWithKeywords(THPVariable_requires_grad_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"set_", castPyCFunctionWithKeywords(THPVariable_set_), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"short", castPyCFunctionWithKeywords(THPVariable_short), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"size", castPyCFunctionWithKeywords(THPVariable_size), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"untyped_storage", THPVariable_storage, METH_NOARGS, nullptr}, + {"storage_offset", THPVariable_storage_offset, METH_NOARGS, nullptr}, + {"stride", castPyCFunctionWithKeywords(THPVariable_stride), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"to", castPyCFunctionWithKeywords(THPVariable_to), METH_VARARGS | METH_KEYWORDS, nullptr}, + {"tolist", THPVariable_tolist, METH_NOARGS, nullptr}, + {"type", castPyCFunctionWithKeywords(THPVariable_type), METH_VARARGS | METH_KEYWORDS, nullptr}, + ${py_method_defs} + {nullptr} +}; + +} // namespace torch::autograd diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/variable_factories.h b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/variable_factories.h new file mode 100644 index 0000000000000000000000000000000000000000..2b55f441ab6249cb7963c5e4a15070f626f775b7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/packaged/autograd/templates/variable_factories.h @@ -0,0 +1,135 @@ +#pragma once + +// ${generated_comment} + +#include +#include +#include +#include +#include +#include +#include + +#ifndef AT_PER_OPERATOR_HEADERS +#include +#else +#include +$ops_headers +#endif + +#include +#include +#include + +namespace torch { + +/// NOTE: Currently `torch::tensor(...)` doesn't support mixed data types +/// (i.e. `torch::tensor({{bool, 2.0}})` doesn't work). We might be able to +/// support it in the future by iterating over all sub-lists to find +/// the largest data type that can represent all of the elements, or by using +/// variadic templates. +/// +/// NOTE: C++ `torch::tensor` with a floating-point type or an `at::ArrayRef` / `std::vector` / +/// (nested) braced-init-list of floating-point types always produces a tensor of dtype +/// `torch::get_default_dtype()`, matching Python `torch.tensor` behavior. +/// +/// NOTE: C++ `torch::tensor` with an integer type or an `at::ArrayRef` / `std::vector` / +/// (nested) braced-init-list of integer types always produces a tensor of dtype `at::kLong` +/// (aka. int64_t), matching Python `torch.tensor` behavior. +/// +/// NOTE: The following dtypes are not supported by `torch::tensor` currently: +/// - `unsigned int` +/// - `unsigned long int` +/// - `unsigned long long int` +/// - `long long int` +inline at::Tensor tensor(detail::TensorDataContainer tensor_data_container, const at::TensorOptions& options = {}) { + return autograd::make_variable( + // note: we remove the requires_grad setting from the TensorOptions because + // it is ignored anyways (and we actually have an assertion that it isn't set + // which would fail otherwise). We handle requires_grad explicitly here + // instead of passing it through to the kernel. + tensor_data_container.convert_to_tensor(options.requires_grad(::std::nullopt)), + options.requires_grad()); +} + +/// A generic deleter function. +using Deleter = std::function; +using at::MemoryFormat; + +/// Exposes the given `data` as a `Tensor` without taking ownership of the +/// original data. `sizes` should specify the shape of the tensor, `strides` the +/// stride in each dimension. The `deleter` function (a +/// `std::function`) will be called on the `data` when the Tensor +/// data would normally be deallocated. The `TensorOptions` specify additional +/// configuration options for the returned tensor, such as what type to +/// interpret the `data` as. +inline at::Tensor from_blob( + void* data, + at::IntArrayRef sizes, + at::IntArrayRef strides, + const Deleter& deleter, + const at::TensorOptions& options = at::TensorOptions()) { + at::Tensor tensor = ([&]() { + at::AutoDispatchBelowAutograd guard; // TODO: remove + at::tracer::impl::NoTracerDispatchMode tracer_guard; + return at::from_blob(data, sizes, strides, deleter, options.requires_grad(::std::nullopt)); + })(); + return autograd::make_variable(tensor, options.requires_grad()); +} + +/// Exposes the given `data` as a `Tensor` without taking ownership of the +/// original data. `sizes` should specify the shape of the tensor, `strides` the +/// stride in each dimension. The `TensorOptions` +/// specify additional configuration options for the returned tensor, such as +/// what type to interpret the `data` as. +inline at::Tensor from_blob( + void* data, + at::IntArrayRef sizes, + at::IntArrayRef strides, + const at::TensorOptions& options = at::TensorOptions()) { + at::Tensor tensor = ([&]() { + at::AutoDispatchBelowAutograd guard; // TODO: remove + at::tracer::impl::NoTracerDispatchMode tracer_guard; + return at::from_blob(data, sizes, strides, options.requires_grad(::std::nullopt)); + })(); + return autograd::make_variable(tensor, options.requires_grad()); +} + +/// Exposes the given `data` as a `Tensor` without taking ownership of the +/// original data. `sizes` should specify the shape of the tensor. The `deleter` +/// (a `std::function`) function will be called on the `data` when +/// the Tensor data would normally be deallocated. The `TensorOptions` specify +/// additional configuration options for the returned tensor, such as what type +/// to interpret the `data` as. +inline at::Tensor from_blob( + void* data, + at::IntArrayRef sizes, + const Deleter& deleter, + const at::TensorOptions& options = at::TensorOptions()) { + at::Tensor tensor = ([&]() { + at::AutoDispatchBelowAutograd guard; // TODO: remove + at::tracer::impl::NoTracerDispatchMode tracer_guard; + return at::from_blob(data, sizes, deleter, options.requires_grad(::std::nullopt)); + })(); + return autograd::make_variable(tensor, options.requires_grad()); +} + +/// Exposes the given `data` as a `Tensor` without taking ownership of the +/// original data. `sizes` should specify the shape of the tensor. The +/// `TensorOptions` specify additional configuration options for the returned +/// tensor, such as what type to interpret the `data` as. +inline at::Tensor from_blob( + void* data, + at::IntArrayRef sizes, + const at::TensorOptions& options = at::TensorOptions()) { + at::Tensor tensor = ([&]() { + at::AutoDispatchBelowAutograd guard; // TODO: remove + at::tracer::impl::NoTracerDispatchMode tracer_guard; + return at::from_blob(data, sizes, options.requires_grad(::std::nullopt)); + })(); + return autograd::make_variable(tensor, options.requires_grad()); +} + +${function_definitions} + +} // namespace torch diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/selective_build/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/selective_build/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/selective_build/operator.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/selective_build/operator.py new file mode 100644 index 0000000000000000000000000000000000000000..dc53851a80f561bdb4180a3f50a7e71930afaaa1 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/selective_build/operator.py @@ -0,0 +1,186 @@ +from __future__ import annotations + +from dataclasses import dataclass + + +# This class holds information about a single operator used to determine +# the outcome of a selective/custom PyTorch build that doesn't include +# registration code for all the supported operators. This is done to +# reduce the size of the generated binary so that it can be deployed in +# situations where binary size comes at a premium. +# +@dataclass(frozen=True) +class SelectiveBuildOperator: + # The name of the operator. This includes the aten::, etc... prefix + # The operator name may or may not have the overload name. If this + # operator name does not specify an overload name, the way to determine + # if this entry refers to the family of operators with this base name + # or just the operator with this name is to look at the value of the + # 'include_all_overloads' flag in this class. + name: str + + # True if this is a root operator (i.e. called directly from a + # TorchScript model, etc...). An operator is considered to be a + # root operator if it is called directly from any one of the models + # that this instance of the pytorch library was built for. Hence, it + # may not be a root operator in all of the models that are used in + # this instance of the pytorch library. + is_root_operator: bool + + # Is this operator used for on-device training? If True, then we need to + # use the information to generate code in VariableType_N.cpp for registration + # of training related operators. Again, this is True if this operator + # is used for training in one or more models used by this instance of the + # pytorch library. + is_used_for_training: bool + + # If True, it indicates that this operator instance (object) refers to an + # operator without the overload name and should apply to all overloads + # which have this operator name as the base name. This flag is applicable + # only for objects that have operator names without a DOT (period) character + # in them. + # + # Note: This flag is a temporary workaround to grandfather in the current + # static selective (custom) build mechanism, which largely ignores overload + # names when determining whether to select operators for registration + # purposes. + include_all_overloads: bool + + # Debug Information at the operator level + _debug_info: tuple[str, ...] | None + + @staticmethod + def from_yaml_dict( + op_name: str, op_info: dict[str, object] + ) -> SelectiveBuildOperator: + allowed_keys = { + "name", + "is_root_operator", + "is_used_for_training", + "include_all_overloads", + "debug_info", + } + + if len(set(op_info.keys()) - allowed_keys) > 0: + raise Exception( # noqa: TRY002 + "Got unexpected top level keys: {}".format( + ",".join(set(op_info.keys()) - allowed_keys), + ) + ) + + if "name" in op_info: + if op_name != op_info["name"]: + raise AssertionError( + f"op_name mismatch: {op_name} != {op_info['name']}" + ) + + is_root_operator = op_info.get("is_root_operator", True) + if not isinstance(is_root_operator, bool): + raise AssertionError( + f"Expected 'is_root_operator' to be bool, got {type(is_root_operator)}" + ) + + is_used_for_training = op_info.get("is_used_for_training", True) + if not isinstance(is_used_for_training, bool): + raise AssertionError( + f"Expected 'is_used_for_training' to be bool, got {type(is_used_for_training)}" + ) + + include_all_overloads = op_info.get("include_all_overloads", True) + if not isinstance(include_all_overloads, bool): + raise AssertionError( + f"Expected 'include_all_overloads' to be bool, got {type(include_all_overloads)}" + ) + + debug_info: tuple[str, ...] | None = None + if "debug_info" in op_info: + di_list = op_info["debug_info"] + if not isinstance(di_list, list): + raise AssertionError( + f"Expected 'debug_info' to be list, got {type(di_list)}" + ) + debug_info = tuple(str(x) for x in di_list) + + return SelectiveBuildOperator( + name=op_name, + is_root_operator=is_root_operator, + is_used_for_training=is_used_for_training, + include_all_overloads=include_all_overloads, + _debug_info=debug_info, + ) + + @staticmethod + def from_legacy_operator_name_without_overload( + name: str, + ) -> SelectiveBuildOperator: + return SelectiveBuildOperator( + name=name, + is_root_operator=True, + is_used_for_training=True, + include_all_overloads=True, + _debug_info=None, + ) + + def to_dict(self) -> dict[str, object]: + ret: dict[str, object] = { + "is_root_operator": self.is_root_operator, + "is_used_for_training": self.is_used_for_training, + "include_all_overloads": self.include_all_overloads, + } + if self._debug_info is not None: + ret["debug_info"] = self._debug_info + + return ret + + +def merge_debug_info( + lhs: tuple[str, ...] | None, + rhs: tuple[str, ...] | None, +) -> tuple[str, ...] | None: + # Ensure that when merging, each entry shows up just once. + if lhs is None and rhs is None: + return None + + return tuple(set((lhs or ()) + (rhs or ()))) + + +def combine_operators( + lhs: SelectiveBuildOperator, rhs: SelectiveBuildOperator +) -> SelectiveBuildOperator: + if str(lhs.name) != str(rhs.name): + raise Exception( # noqa: TRY002 + f"Expected both arguments to have the same name, but got '{str(lhs.name)}' and '{str(rhs.name)}' instead" + ) + + return SelectiveBuildOperator( + name=lhs.name, + # Consider this operator to be a root operator if it is a + # root operator in any of the models used in this instance of + # the pytorch library. + is_root_operator=lhs.is_root_operator or rhs.is_root_operator, + # Consider this operator to be a training operator if it is + # an operator used for training in any of the models used + # in this instance of the pytorch library. + is_used_for_training=lhs.is_used_for_training or rhs.is_used_for_training, + include_all_overloads=lhs.include_all_overloads or rhs.include_all_overloads, + _debug_info=merge_debug_info(lhs._debug_info, rhs._debug_info), + ) + + +def merge_operator_dicts( + lhs: dict[str, SelectiveBuildOperator], + rhs: dict[str, SelectiveBuildOperator], +) -> dict[str, SelectiveBuildOperator]: + operators: dict[str, SelectiveBuildOperator] = {} + for op_name, op in list(lhs.items()) + list(rhs.items()): + new_op = op + if op_name in operators: + new_op = combine_operators(operators[op_name], op) + + operators[op_name] = new_op + + return operators + + +def strip_operator_overload_name(op_name: str) -> str: + return op_name.split(".", maxsplit=1)[0] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/selective_build/selector.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/selective_build/selector.py new file mode 100644 index 0000000000000000000000000000000000000000..fa48a9df2dd752275513e42423ccd50576aa6cb3 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/selective_build/selector.py @@ -0,0 +1,377 @@ +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Iterable +from dataclasses import dataclass +from typing import TYPE_CHECKING + +import yaml + +from torchgen.selective_build.operator import ( + merge_debug_info, + merge_operator_dicts, + SelectiveBuildOperator, + strip_operator_overload_name, +) + + +if TYPE_CHECKING: + from torchgen.model import NativeFunction + + +# A SelectiveBuilder holds information extracted from the selective build +# YAML specification. +# +# It includes information about the build's selectivity, the debug_info +# associated with this selective build (opaque string), and the set of +# operators that should be included in the build. +# +@dataclass(frozen=True) +class SelectiveBuilder: + # If true, then the build is not selective, and includes all + # operators. + include_all_operators: bool + + # Debug Information at the selective/custom build level. + _debug_info: tuple[str, ...] | None + + # A dictionary of operator -> operator metadata. + operators: dict[str, SelectiveBuildOperator] + + # A dictionary of selected kernel tags and dtypes. Typically a + # PyTorch Operator Kernel (function) may have many code paths + # that are specialized for many many Tensor dtypes, so it's not + # one per kernel function, but there could be many per kernel + # function. The tag isn't a kernel function name, but some fragment + # of the kernel function implementation itself. + kernel_metadata: dict[str, list[str]] + + # ExecuTorch only. A dictionary of kernel tag -> list of (list of input + # dtypes for tensor-like input args). + # This is from selective.yaml + et_kernel_metadata: dict[str, list[str]] + + # A set of all the custom torch bind classes used by the selected models + # Stored as a set internally to remove duplicates proactively, but written + # as a list to yamls + custom_classes: set[str] + + # A set of all the build features used by the selected models + # Stored as a set internally to remove duplicates proactively, but written + # as a list to yamls + build_features: set[str] + + # If true, then fragments for all dtypes for all kernel functions + # are included as well as all custom classes. This is typically set when any one of the + # operator lists is generated from a mechanism other than + # tracing based selective build. + include_all_non_op_selectives: bool + + @staticmethod + def get_nop_selector() -> SelectiveBuilder: + return SelectiveBuilder.from_yaml_dict({"include_all_operators": True}) + + @staticmethod + def from_yaml_dict(data: dict[str, object]) -> SelectiveBuilder: + valid_top_level_keys = { + "include_all_non_op_selectives", + "include_all_operators", + "debug_info", + "operators", + "kernel_metadata", + "et_kernel_metadata", + "custom_classes", + "build_features", + } + top_level_keys = set(data.keys()) + if len(top_level_keys - valid_top_level_keys) > 0: + raise Exception( # noqa: TRY002 + "Got unexpected top level keys: {}".format( + ",".join(top_level_keys - valid_top_level_keys), + ) + ) + include_all_operators = data.get("include_all_operators", False) + if not isinstance(include_all_operators, bool): + raise AssertionError( + f"Expected 'include_all_operators' to be bool, got {type(include_all_operators)}" + ) + + debug_info = None + if "debug_info" in data: + di_list = data["debug_info"] + if not isinstance(di_list, list): + raise AssertionError( + f"Expected 'debug_info' to be list, got {type(di_list)}" + ) + + debug_info = tuple(str(x) for x in di_list) + + operators = {} + operators_dict = data.get("operators", {}) + if not isinstance(operators_dict, dict): + raise AssertionError( + f"Expected 'operators' to be dict, got {type(operators_dict)}" + ) + + for k, v in operators_dict.items(): + operators[k] = SelectiveBuildOperator.from_yaml_dict(k, v) + + kernel_metadata = {} + kernel_metadata_dict = data.get("kernel_metadata", {}) + if not isinstance(kernel_metadata_dict, dict): + raise AssertionError( + f"Expected 'kernel_metadata' to be dict, got {type(kernel_metadata_dict)}" + ) + + for k, v in kernel_metadata_dict.items(): + kernel_metadata[str(k)] = [str(dtype) for dtype in v] + + et_kernel_metadata = data.get("et_kernel_metadata", {}) + if not isinstance(et_kernel_metadata, dict): + raise AssertionError( + f"Expected 'et_kernel_metadata' to be dict, got {type(et_kernel_metadata)}" + ) + + custom_classes = data.get("custom_classes", []) + if not isinstance(custom_classes, Iterable): + raise AssertionError( + f"Expected 'custom_classes' to be Iterable, got {type(custom_classes)}" + ) + custom_classes = set(custom_classes) + + build_features = data.get("build_features", []) + if not isinstance(build_features, Iterable): + raise AssertionError( + f"Expected 'build_features' to be Iterable, got {type(build_features)}" + ) + build_features = set(build_features) + + include_all_non_op_selectives = data.get("include_all_non_op_selectives", False) + if not isinstance(include_all_non_op_selectives, bool): + raise AssertionError( + f"Expected 'include_all_non_op_selectives' to be bool, " + f"got {type(include_all_non_op_selectives)}" + ) + + return SelectiveBuilder( + include_all_operators, + debug_info, + operators, + kernel_metadata, + et_kernel_metadata, + custom_classes, # type: ignore[arg-type] + build_features, # type: ignore[arg-type] + include_all_non_op_selectives, + ) + + @staticmethod + def from_yaml_str(config_contents: str) -> SelectiveBuilder: + contents = yaml.safe_load(config_contents) + return SelectiveBuilder.from_yaml_dict(contents) + + @staticmethod + def from_yaml_path(config_path: str) -> SelectiveBuilder: + with open(config_path) as f: + contents = yaml.safe_load(f) + return SelectiveBuilder.from_yaml_dict(contents) + + @staticmethod + def from_legacy_op_registration_allow_list( + allow_list: set[str], is_root_operator: bool, is_used_for_training: bool + ) -> SelectiveBuilder: + operators = {} + for op in allow_list: + operators[op] = { + "name": op, + "is_root_operator": is_root_operator, + "is_used_for_training": is_used_for_training, + "include_all_overloads": True, + } + return SelectiveBuilder.from_yaml_dict( + { + "operators": operators, + "include_all_non_op_selectives": True, + } + ) + + def is_operator_selected(self, name: str) -> bool: + if self.include_all_operators: + return True + + if name in self.operators: + return True + name = strip_operator_overload_name(name) + return name in self.operators and self.operators[name].include_all_overloads + + def is_native_function_selected(self, func: NativeFunction) -> bool: + op_name = op_name_from_native_function(func) + return self.is_operator_selected(op_name) + + def is_operator_selected_for_training(self, name: str) -> bool: + if not self.is_operator_selected(name): + return False + if self.include_all_operators: + return True + + not_training_op = SelectiveBuildOperator( + name="", + is_root_operator=False, + is_used_for_training=False, + include_all_overloads=False, + _debug_info=None, + ) + op = not_training_op + if name in self.operators: + op = self.operators[name] + + name = strip_operator_overload_name(name) + base_op = not_training_op + if name in self.operators: + base_op = self.operators[name] + + return op.is_used_for_training or ( + base_op.include_all_overloads and base_op.is_used_for_training + ) + + def is_native_function_selected_for_training(self, func: NativeFunction) -> bool: + op_name = op_name_from_native_function(func) + return self.is_operator_selected_for_training(op_name) + + def is_root_operator(self, name: str) -> bool: + if not self.is_operator_selected(name): + return False + if self.include_all_operators: + return True + + if name in self.operators: + op: SelectiveBuildOperator = self.operators[name] + return op.is_root_operator + name = strip_operator_overload_name(name) + if name not in self.operators: + return False + base_op: SelectiveBuildOperator = self.operators[name] + return base_op.include_all_overloads and base_op.is_root_operator + + def is_kernel_dtype_selected(self, kernel_tag: str, dtype: str) -> bool: + if self.include_all_operators or self.include_all_non_op_selectives: + return True + + return ( + kernel_tag in self.kernel_metadata + and dtype in self.kernel_metadata[kernel_tag] + ) + + def et_get_selected_kernels(self, op_name: str, kernel_key: list[str]) -> list[str]: + """ + Return a list of kernel keys that cover the used ops + """ + # If no kernel metadata, either it's implied by include_all_operators=True or the op is not used. + if op_name not in self.et_kernel_metadata: + return kernel_key if self.include_all_operators else [] + # Otherwise, only return the specific kernel keys. + + result_set = set() + + for model_kernel_keys in self.et_kernel_metadata[op_name]: + key_found = False + for key in kernel_key: + # Don't compare the version for now + if ( + key != "default" + and key.split("/")[1] == model_kernel_keys.split("/")[1] + ): + result_set.add(key) + key_found = True + break + if not key_found: + if "default" not in kernel_key: + raise Exception("Missing kernel for the model") # noqa: TRY002 + else: + result_set.add("default") + + return list(result_set) + + def to_dict(self) -> dict[str, object]: + ret: dict[str, object] = { + "include_all_non_op_selectives": self.include_all_non_op_selectives, + "include_all_operators": self.include_all_operators, + } + operators = {} + for op_name, op in self.operators.items(): + operators[op_name] = op.to_dict() + ret["operators"] = operators + + if self._debug_info is not None: + ret["debug_info"] = sorted(self._debug_info) + + ret["kernel_metadata"] = { + k: sorted(v) for (k, v) in self.kernel_metadata.items() + } + + ret["et_kernel_metadata"] = self.et_kernel_metadata + + ret["custom_classes"] = sorted(self.custom_classes) + + ret["build_features"] = sorted(self.build_features) + + return ret + + +def merge_kernel_metadata( + lhs: dict[str, list[str]], + rhs: dict[str, list[str]], +) -> dict[str, list[str]]: + kernel_metadata: dict[str, list[str]] = {} + for tag_name, dtypes in list(lhs.items()) + list(rhs.items()): + dtypes_copy = set(dtypes) + if tag_name in kernel_metadata: + dtypes_copy |= set(kernel_metadata[tag_name]) + + kernel_metadata[tag_name] = list(dtypes_copy) + + return kernel_metadata + + +def merge_et_kernel_metadata( + lhs: dict[str, list[str]], + rhs: dict[str, list[str]], +) -> dict[str, list[str]]: + merge_et_kernel_metadata: dict[str, set[str]] = defaultdict(set) + for op in list(lhs.keys()) + list(rhs.keys()): + merge_et_kernel_metadata[op].update(lhs.get(op, [])) + merge_et_kernel_metadata[op].update(rhs.get(op, [])) + + return {op: sorted(val) for op, val in merge_et_kernel_metadata.items()} + + +def combine_selective_builders( + lhs: SelectiveBuilder, rhs: SelectiveBuilder +) -> SelectiveBuilder: + include_all_operators = lhs.include_all_operators or rhs.include_all_operators + debug_info = merge_debug_info(lhs._debug_info, rhs._debug_info) + operators = merge_operator_dicts(lhs.operators, rhs.operators) + kernel_metadata = merge_kernel_metadata(lhs.kernel_metadata, rhs.kernel_metadata) + et_kernel_metadata = merge_et_kernel_metadata( + lhs.et_kernel_metadata, rhs.et_kernel_metadata + ) + include_all_non_op_selectives = ( + lhs.include_all_non_op_selectives or rhs.include_all_non_op_selectives + ) + custom_classes = lhs.custom_classes.union(rhs.custom_classes) + build_features = lhs.build_features.union(rhs.build_features) + return SelectiveBuilder( + include_all_operators, + debug_info, + operators, + kernel_metadata, + et_kernel_metadata, + custom_classes, + build_features, + include_all_non_op_selectives, + ) + + +def op_name_from_native_function(f: NativeFunction) -> str: + # This was originally read from the 'operator_name_with_overload' field in the + # declaration dict, which was the part before the first '(' in 'schema_string'. + return f"{f.namespace}::{f.func.name}" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/config.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/config.py new file mode 100644 index 0000000000000000000000000000000000000000..c0993d00702c70532467eb5a190fcf73b6cdd846 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/config.py @@ -0,0 +1,389 @@ +from __future__ import annotations + +from torchgen.model import NativeFunctionsGroup, NativeFunctionsViewGroup + + +def func_name_base_str(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> str: + if isinstance(g, NativeFunctionsGroup): + return str(g.functional.func.name.name.base) + else: + return str(g.view.root_name) + + +is_hand_written_ops_ = frozenset( + ( + "abs", + "add", + "addmm", + "all", + "any", + "argmin", + "bmm", + "clamp", + "clamp_min", + "cumsum", + "div", + "fmod", + "index_select", + "leaky_relu", + "linear", + "log", + "matmul", + "mul", + "narrow_copy", + "nonzero", + "pow", + "remainder", + "sigmoid", + "sign", + "sub", + "tanh", + "detach", + "expand_as", + "flatten", + "narrow", + "reshape_as", + "select", + "slice", + "softmax", + "split", + "squeeze", + "transpose", + "view", + "where", + ) +) + + +def is_hand_written(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> bool: + name_base = func_name_base_str(g) + return name_base in is_hand_written_ops_ + + +def override_test_values(arg_map: dict[str, str], op_name: str, index: int) -> None: + if index not in (0, 1): + raise AssertionError(f"index must be 0 or 1, got {index}") + if op_name == "addr": + if index == 0: + arg_map["self"] = "at::rand({6, 6})" + arg_map["vec1"] = "at::rand({6})" + arg_map["vec2"] = "at::rand({6})" + else: + arg_map["self"] = "at::rand({22, 22})" + arg_map["vec1"] = "at::rand({22})" + arg_map["vec2"] = "at::rand({22})" + return + if op_name == "mv": + if index == 0: + arg_map["self"] = "at::rand({6, 6})" + arg_map["vec"] = "at::rand({6})" + else: + arg_map["self"] = "at::rand({22, 22})" + arg_map["vec"] = "at::rand({22})" + return + if op_name == "addbmm": + if index == 0: + arg_map["self"] = "at::rand({6, 6})" + else: + arg_map["self"] = "at::rand({22, 22})" + return + if op_name == "cross": + if index == 0: + arg_map["self"] = "at::rand({3, 3, 3})" + arg_map["other"] = "at::rand({3, 3, 3})" + else: + arg_map["self"] = "at::rand({22, 3, 22})" + arg_map["other"] = "at::rand({22, 3, 22})" + return + if op_name == "take": + if index == 0: + arg_map["index"] = "at::randint(0, 216, {20}, torch::kInt64)" + else: + arg_map["index"] = "at::randint(0, 1000, {100}, torch::kInt64)" + return + if op_name == "take_along_dim": + if index == 0: + arg_map["indices"] = "at::argsort(self0, 1, true)" + else: + arg_map["indices"] = "at::argsort(self1, 1, true)" + return + if op_name == "masked_select": + if index == 0: + arg_map["mask"] = "at::randn({6, 6, 6}) > 0.5" + else: + arg_map["mask"] = "at::rand({22, 22, 22}) > 0.5" + return + if op_name == "orgqr": + if index == 0: + arg_map["input2"] = "at::rand({6, 6})" + else: + arg_map["input2"] = "at::rand({22, 22})" + return + if op_name == "ormqr": + if index == 0: + arg_map["input2"] = "at::rand({6, 6})" + else: + arg_map["input2"] = "at::rand({22, 22})" + return + if op_name == "quantile": + if index == 0: + arg_map["q"] = "at::rand({6})" + arg_map["interpolation"] = '"linear"' + else: + arg_map["q"] = "at::rand({22})" + arg_map["interpolation"] = '"linear"' + return + if op_name == "nanquantile": + if index == 0: + arg_map["q"] = "at::rand({6})" + arg_map["interpolation"] = '"linear"' + else: + arg_map["q"] = "at::rand({22})" + arg_map["interpolation"] = '"linear"' + return + if op_name == "multi_margin_loss": + if index == 0: + arg_map["self"] = "at::rand({6, 6})" + arg_map["target"] = "at::randint(6, {6}, torch::kInt64)" + arg_map["weight"] = "at::rand({6})" + else: + arg_map["self"] = "at::rand({22, 22})" + arg_map["target"] = "at::randint(22, {22}, torch::kInt64)" + arg_map["weight"] = "at::rand({22})" + return + if op_name == "multilabel_margin_loss": + if index == 0: + arg_map["self"] = "at::rand({6, 6})" + arg_map["target"] = "at::randint(6, {6, 6}, torch::kInt64)" + else: + arg_map["self"] = "at::rand({22, 22})" + arg_map["target"] = "at::randint(22, {22, 22}, torch::kInt64)" + return + if op_name == "nll_loss": + if index == 0: + arg_map["self"] = "at::rand({6, 6})" + arg_map["target"] = "at::randint(6, {6}, torch::kInt64)" + arg_map["weight"] = "at::rand({6})" + else: + arg_map["self"] = "at::rand({22, 22})" + arg_map["target"] = "at::randint(22, {22}, torch::kInt64)" + arg_map["weight"] = "at::rand({22})" + return + if op_name == "nll_loss2d": + if index == 0: + arg_map["self"] = "at::rand({6, 6, 6, 6})" + arg_map["target"] = "at::randint(6, {6, 6, 6}, torch::kInt64)" + arg_map["weight"] = "at::rand({6})" + else: + arg_map["self"] = "at::rand({22, 22, 22, 22})" + arg_map["target"] = "at::randint(22, {22, 22, 22}, torch::kInt64)" + arg_map["weight"] = "at::rand({22})" + return + if op_name in ( + "fft_fft", + "fft_ifft", + "fft_rfft", + "fft_irfft", + "fft_hfft", + "fft_ihfft", + ): + arg_map["norm"] = '"forward"' + return + if op_name == "linalg_tensorinv": + if index == 0: + arg_map["self"] = "at::rand({6, 6, 6, 6})" + arg_map["ind"] = "2" + else: + arg_map["self"] = "at::rand({22, 22, 22, 22})" + arg_map["ind"] = "2" + return + if op_name == "addmv": + if index == 0: + arg_map["self"] = "at::rand({2})" + arg_map["mat"] = "at::rand({2, 2})" + arg_map["vec"] = "at::rand({2})" + else: + arg_map["self"] = "at::rand({35})" + arg_map["mat"] = "at::rand({35, 35})" + arg_map["vec"] = "at::rand({35})" + return + if op_name == "acosh": + if index == 0: + arg_map["self"] = "at::rand({2, 2, 2}) + at::ones({2, 2, 2})" + else: + arg_map["self"] = "at::rand({5, 5, 5}) + at::ones({5, 5, 5})" + return + if op_name == "adaptive_max_pool2d_backward": + if index == 0: + arg_map["grad_output"] = "at::rand({2, 2, 2}, at::kFloat)" + arg_map["self"] = "at::rand({2, 2, 2}, at::kFloat)" + arg_map["indices"] = "at::randint(0, 1, {2, 2, 2}, at::kLong)" + else: + arg_map["grad_output"] = "at::rand({3, 3, 3}, at::kFloat)" + arg_map["self"] = "at::rand({3, 3, 3}, at::kFloat)" + arg_map["indices"] = "at::randint(0, 1, {3, 3, 3}, at::kLong)" + return + if op_name == "adaptive_max_pool3d_backward": + if index == 0: + arg_map["grad_output"] = "at::rand({2, 2, 2, 2}, at::kFloat)" + arg_map["self"] = "at::rand({2, 2, 2, 2}, at::kFloat)" + arg_map["indices"] = "at::randint(0, 1, {2, 2, 2, 2}, at::kLong)" + else: + arg_map["grad_output"] = "at::rand({3, 3, 3, 3}, at::kFloat)" + arg_map["self"] = "at::rand({3, 3, 3, 3}, at::kFloat)" + arg_map["indices"] = "at::randint(0, 1, {3, 3, 3, 3}, at::kLong)" + return + if op_name == "bitwise_left_shift": + if index == 0: + arg_map["self"] = "at::randint(1, 1 << 4, {6, 6, 6}, at::kInt)" + arg_map["other"] = "at::randint(1, 26, {6, 6, 6}, at::kInt)" + else: + arg_map["self"] = "at::randint(1, 1 << 4, {22, 22, 22}, at::kInt)" + arg_map["other"] = "at::randint(1, 26, {22, 22, 22}, at::kInt)" + return + if op_name == "bitwise_right_shift": + if index == 0: + arg_map["self"] = "at::randint(1 << 21, 1 << 30, {6, 6, 6}, at::kInt)" + arg_map["other"] = "at::randint(1, 22, {6, 6, 6}, at::kInt)" + else: + arg_map["self"] = "at::randint(1 << 21, 1 << 30, {22, 22, 22}, at::kInt)" + arg_map["other"] = "at::randint(1, 22, {22, 22, 22}, at::kInt)" + return + if op_name == "gather": + if index == 0: + arg_map["self"] = "at::randint(1, 100, {2,2,2}, at::kInt)" + arg_map["dim"] = "1" + arg_map["index"] = "at::randint(0, 1, {2,2,2}, torch::kInt64)" + arg_map["sparse_grad"] = "false" + else: + arg_map["self"] = "at::randint(1, 100, {5,5,5}, at::kInt)" + arg_map["dim"] = "1" + arg_map["index"] = "at::randint(0, 4, {5,5,5}, torch::kInt64)" + arg_map["sparse_grad"] = "false" + return + if op_name == "gelu": + if index == 0: + arg_map["self"] = "at::rand({6, 6, 6})" + arg_map["approximate"] = '"tanh"' + else: + arg_map["self"] = "at::rand({22, 22, 22})" + arg_map["approximate"] = '"tanh"' + return + if op_name == "gelu_backward": + if index == 0: + arg_map["grad_output"] = "at::rand({6, 6, 6})" + arg_map["self"] = "at::rand({6, 6, 6})" + arg_map["approximate"] = '"tanh"' + else: + arg_map["grad_output"] = "at::rand({22, 22, 22})" + arg_map["self"] = "at::rand({22, 22, 22})" + arg_map["approximate"] = '"tanh"' + return + if op_name == "index_add": + if index == 0: + arg_map["self"] = "at::rand({2})" + arg_map["dim"] = "0" + arg_map["index"] = "at::randint(0, 1, {2}, at::kInt)" + arg_map["source"] = "at::rand({2})" + arg_map["alpha"] = "2" + else: + arg_map["self"] = "at::rand({16})" + arg_map["dim"] = "0" + arg_map["index"] = "at::randint(0, 10, {16}, at::kInt)" + arg_map["source"] = "at::rand({16})" + arg_map["alpha"] = "2" + return + if op_name == "index_copy": + if index == 0: + arg_map["self"] = "at::rand({2})" + arg_map["dim"] = "0" + arg_map["index"] = "at::randint(0, 1, {2}, at::kLong)" + arg_map["source"] = "at::rand({2})" + else: + arg_map["self"] = "at::rand({32})" + arg_map["dim"] = "0" + arg_map["index"] = "at::randint(0, 10, {32}, at::kLong)" + arg_map["source"] = "at::rand({32})" + return + if op_name == "linalg_cross": + if index == 0: + arg_map["self"] = "at::rand({6, 3, 6})" + arg_map["other"] = "at::rand({6, 3, 6})" + arg_map["dim"] = "1" + else: + arg_map["self"] = "at::rand({22, 3, 22})" + arg_map["other"] = "at::rand({22, 3, 22})" + arg_map["dim"] = "1" + return + if op_name == "nll_loss_backward": + if index == 0: + arg_map["grad_output"] = "at::rand({})" + arg_map["self"] = "at::rand({6})" + arg_map["target"] = "at::randint(0, 5, {6}, torch::kInt64)" + arg_map["weight"] = "at::rand({6})" + arg_map["reduction"] = "1" + arg_map["ignore_index"] = "1" + arg_map["total_weight"] = "at::rand({})" + else: + arg_map["grad_output"] = "at::rand({})" + arg_map["self"] = "at::rand({36})" + arg_map["target"] = "at::randint(0, 11, {36}, torch::kInt64)" + arg_map["weight"] = "at::rand({36})" + arg_map["reduction"] = "1" + arg_map["ignore_index"] = "1" + arg_map["total_weight"] = "at::rand({})" + return + if op_name in ["scatter", "scatter_add", "_scatter_reduce"]: + if index == 0: + arg_map["self"] = "at::randint(1, 100, {2,2,2}, torch::kInt64)" + arg_map["index"] = "at::randint(0, 1, {2,2,2}, torch::kInt64)" + arg_map["src"] = "at::randint(1, 100, {2,2,2}, torch::kInt64)" + else: + arg_map["self"] = "at::randint(1, 100, {5,5,5}, torch::kInt64)" + arg_map["index"] = "at::randint(0, 1, {5,5,5}, torch::kInt64)" + arg_map["src"] = "at::randint(1, 100, {5,5,5}, torch::kInt64)" + if "reduce" in arg_map: + arg_map["reduce"] = '"sum"' if op_name == "_scatter_reduce" else '"add"' + return + if op_name == "scatter_reduce": + arg_map["reduce"] = '"mean"' + if index == 0: + arg_map["index"] = "at::randint(6, {6, 6, 6}, torch::kInt64)" + else: + arg_map["index"] = "at::randint(22, {22, 22, 22}, torch::kInt64)" + return + if op_name == "special_zeta": + if index == 0: + arg_map["self"] = "at::rand({2,2,2}, at::kDouble) + at::ones({2,2,2})" + arg_map["other"] = "at::rand({2,2,2}, at::kDouble) + at::ones({2,2,2})" + else: + arg_map["self"] = "at::rand({5,5,5}, at::kDouble) + at::ones({5,5,5})" + arg_map["other"] = "at::rand({5,5,5}, at::kDouble) + at::ones({5,5,5})" + return + if op_name == "_convert_indices_from_csr_to_coo": + if index == 0: + arg_map["crow_indices"] = "torch::tensor({1}, torch::kInt32)" + arg_map["col_indices"] = "torch::tensor({0, 1, 0}, torch::kInt32)" + arg_map["out_int32"] = "false" + else: + arg_map["crow_indices"] = "torch::tensor({0}, torch::kInt32)" + arg_map["col_indices"] = ( + "torch::tensor({0, 1, 0, 2, 1, 2, 0, 1, 0, 2, 1, 2}, torch::kInt32)" + ) + arg_map["out_int32"] = "false" + return + if op_name == "_convert_indices_from_coo_to_csr": + if index == 0: + arg_map["self"] = "at::randint(0, 3, {2}, at::kInt)" + arg_map["size"] = "10" + arg_map["out_int32"] = "false" + else: + arg_map["self"] = "at::randint(0, 3, {12}, at::kInt)" + arg_map["size"] = "24" + arg_map["out_int32"] = "false" + return + if op_name in ("diagonal", "linalg_diagonal"): + arg_map["offset"] = "0" + arg_map["dim1"] = "2" + arg_map["dim2"] = "1" + return diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/gen_static_runtime_ops.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/gen_static_runtime_ops.py new file mode 100644 index 0000000000000000000000000000000000000000..d6909bc4d7f67fc13fb9f61e00f4709a4ff5ad4e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/gen_static_runtime_ops.py @@ -0,0 +1,231 @@ +from __future__ import annotations + +import argparse +import itertools +import os +from typing import TYPE_CHECKING, TypeVar + +from libfb.py.log import set_simple_logging # type: ignore[import] + +from torchgen import gen +from torchgen.context import native_function_manager +from torchgen.model import DispatchKey, NativeFunctionsGroup, NativeFunctionsViewGroup +from torchgen.static_runtime import config, generator + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +# Given a list of `grouped_native_functions` sorted by their op names, return a list of +# lists each of which groups ops that share the base name. For example, `mean` and +# `mean.dim` are grouped together by this function. + +NativeGroupT = TypeVar( + "NativeGroupT", + bound=NativeFunctionsGroup | NativeFunctionsViewGroup, +) + + +def group_functions_by_op_name( + grouped_native_functions: Sequence[NativeGroupT], +) -> Sequence[Sequence[NativeGroupT]]: + if not grouped_native_functions: + return [] + groups = [] + + def is_supported(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> bool: + with native_function_manager(g): + return generator.is_supported(g) + + eligible_ops = (g for g in grouped_native_functions if is_supported(g)) + groups = [ + list(group) + for k, group in ( + itertools.groupby( + eligible_ops, + key=config.func_name_base_str, + ) + ) + ] + + return groups + + +def clang_format(cpp_file_path: str) -> None: + import subprocess + + subprocess.check_call(["clang-format", "-i", cpp_file_path]) + + +def write_cpp(cpp_ops: Sequence[str], file_path: str) -> None: + code = "\n".join(cpp_ops) + generated = f"""// @lint-ignore-every CLANGTIDY HOWTOEVEN +// AUTO-GENERATED FROM: torchgen/static_runtime/gen_static_runtime_ops.py +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace torch {{ +namespace jit {{ + +{code} + +}} // namespace jit +}} // namespace torch +""" + with open(file_path, "w") as f: + f.write(generated) + clang_format(file_path) + + +def write_test_cpp(cpp_ops: Sequence[str], file_path: str) -> None: + code = "\n".join(cpp_ops) + generated = f"""// @lint-ignore-every CLANGTIDY HOWTOEVEN +// AUTO-GENERATED FROM: torchgen/static_runtime/gen_static_runtime_ops.py +#include +#include +#include + +#include "test_utils.h" + +using namespace caffe2; +using namespace torch; +using namespace torch::jit; +using namespace torch::jit::test; +using c10::IValue; + +{code} + +""" + with open(file_path, "w") as f: + f.write(generated) + clang_format(file_path) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Generate ATen source files") + parser.add_argument( + "-s", + "--source-path", + help="path to source directory for ATen", + default="caffe2/aten/src/ATen", + ) + parser.add_argument( + "-p", + "--generated-ops-cpp-path", + help="path to directory to generate op dispatcher .cpp file", + default="caffe2/torch/csrc/jit/runtime/static/generated_ops.cpp", + ) + parser.add_argument( + "-t", + "--generated-ops-test-cpp-path", + help="path to directory to generate op dispatcher .cpp file", + default="caffe2/benchmarks/static_runtime/test_generated_ops.cc", + ) + options = parser.parse_args() + native_yaml_path = os.path.join(options.source_path, "native/native_functions.yaml") + tags_yaml_path = os.path.join(options.source_path, "native/tags.yaml") + parsed_yaml = gen.parse_native_yaml(native_yaml_path, tags_yaml_path) + native_functions, backend_indices = ( + parsed_yaml.native_functions, + parsed_yaml.backend_indices, + ) + + op_generator = generator.GenOpDispatcher() + test_case_generator = generator.GenOpTestCase() + + native_functions_groups = [ + g + for g in gen.get_grouped_native_functions(native_functions) + if isinstance(g, NativeFunctionsGroup) + ] + + supported_functions_groups = group_functions_by_op_name(native_functions_groups) + + out_variant_op_result = [ + op_generator.out_variant(groups, backend_indices[DispatchKey.CPU]) + for groups in supported_functions_groups + ] + out_variant_test_result = [ + test_case_generator.out_variant(groups) for groups in supported_functions_groups + ] + + native_functions_view_groups = [ + g + for g in gen.get_grouped_by_view_native_functions(native_functions) + if isinstance(g, NativeFunctionsViewGroup) + ] + + supported_functions_view_groups = group_functions_by_op_name( + native_functions_view_groups + ) + + view_op_result = [ + op_generator.view(groups, backend_indices[DispatchKey.CPU]) + for groups in supported_functions_view_groups + ] + view_test_result = [ + test_case_generator.view(groups) for groups in supported_functions_view_groups + ] + + op_result = out_variant_op_result + ["\n\n"] + view_op_result + test_result = out_variant_test_result + ["\n\n"] + view_test_result + + write_cpp(op_result, options.generated_ops_cpp_path) + write_test_cpp(test_result, options.generated_ops_test_cpp_path) + + print( + f"\ntotal grouped native ops: {len(gen.get_grouped_native_functions(native_functions)):d}" + ) + + print(f"grouped native ops with out variant: {len(native_functions_groups):d}") + supported_functions_num = sum(len(groups) for groups in supported_functions_groups) + print(f"generated functions groups with out variant: {supported_functions_num:d}") + + print(f"\nview grouped native ops: {len(native_functions_view_groups):d}") + supported_view_functions_num = sum( + len(groups) for groups in supported_functions_view_groups + ) + print(f"generated functions view groups: {supported_view_functions_num:d}") + + print( + f"\noverall generated : {supported_functions_num + supported_view_functions_num:d}" + ) + + +if __name__ == "__main__": + set_simple_logging(escape_newlines=False) + main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/generator.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/generator.py new file mode 100644 index 0000000000000000000000000000000000000000..d8aba4d13bcde6cf4a2a4700c93e61237632912c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/static_runtime/generator.py @@ -0,0 +1,859 @@ +from __future__ import annotations + +import json +import logging +import math +from typing import TYPE_CHECKING + +import torchgen.api.cpp as cpp +from torchgen.context import native_function_manager +from torchgen.model import ( + Argument, + BackendIndex, + BaseTy, + BaseType, + FunctionSchema, + NativeFunctionsGroup, + NativeFunctionsViewGroup, + OptionalType, + SelfArgument, + TensorOptionsArguments, + Type, +) +from torchgen.static_runtime import config + + +if TYPE_CHECKING: + from collections.abc import Sequence + + +logger: logging.Logger = logging.getLogger() + + +def has_alias( + arguments: Sequence[Argument | SelfArgument | TensorOptionsArguments], +) -> bool: + for arg in arguments: + annotation = getattr(arg, "annotation", None) + if not annotation: + continue + alias_set = getattr(annotation, "alias_set", ()) + if alias_set: + return True + return False + + +BLOCKED_OPS = frozenset( + ( + # non cpu ops + "sparse_sampled_addmm", + "hspmm", + "linalg_svdvals", + # sparse ops + "sspaddmm", + "coalesce", + "_indices", + "indices", + "_values", + "values", + "crow_indices", + "col_indices", + # deprecated ops + "floor_divide", + "ger", + # buggy ops + "conj_physical", # P495807361 + "binary_cross_entropy", # P496394764 + "arccosh", + # uncommon ops + "cholesky", + "lu_solve", + "linalg_cholesky", + "linalg_householder_product", + "linalg_ldl_solve", + "_compute_linear_combination", + # training related ops + "_make_dual", + # cannot call directly + "_fw_primal", + # no documentation + "_index_reduce", + # TODO: these ones got added recently and need manual inspection + "_new_zeros_with_same_feature_meta", + "_conj_physical", + "binary_cross_entropy_with_logits", + "bincount", + "conv_tbc", + "copy", + "_copy_from", + "_copy_from_and_resize", + "count_nonzero", + "cudnn_affine_grid_generator", + "cudnn_affine_grid_generator_backward", + "cudnn_grid_sampler", + "diag_embed", + "embedding", + "embedding_dense_backward", + "_embedding_bag_dense_backward", + "_embedding_bag_per_sample_weights_backward", + "grid_sampler_2d", + "_grid_sampler_2d_cpu_fallback", + "grid_sampler_3d", + "isnan", + "mkldnn_linear", + "median", + "nanmedian", + "_sparse_sparse_matmul", + "batch_norm_backward_elemt", + "_euclidean_dist", + "pixel_shuffle", + "pixel_unshuffle", + "channel_shuffle", + "_reshape_nested_backward", + "relu", + "prelu", + "celu", + "slice_scatter", + "select_scatter", + "diagonal_scatter", + "sum", + "_mkldnn_transpose", + "_nested_tensor_from_mask", + "_nested_from_padded", + "_nested_tensor_size", + "_nested_from_padded_and_nested_example", + "_standard_gamma_grad", + "_dirichlet_grad", + "native_norm", + "_sparse_softmax", + "_sparse_softmax_backward_data", + "_sparse_log_softmax", + "_sparse_log_softmax_backward_data", + "zero", + "_sparse_addmm", + "sparse_mask", + "_sparse_mask_projection", + "_to_dense", + "_coalesce", + "_coalesced", + "copy_sparse_to_sparse", + "to_sparse", + "to_sparse_csr", + "to_sparse_csc", + "to_mkldnn", + "quantize_per_tensor_dynamic", + "quantize_per_channel", + "q_per_channel_scales", + "q_per_channel_zero_points", + "int_repr", + "_make_per_channel_quantized_tensor", + "set", + "lift", + "lift_fresh", + "lift_fresh_copy", + "masked_scatter", + "_masked_softmax", + "_masked_softmax_backward", + "put", + "index_reduce", + "trace", + "_cholesky_solve_helper", + "dist", + "max", + "_torch_cuda_cu_linker_symbol_op", + "glu_jvp", + "glu_backward_jvp", + "hardswish_backward", + "rrelu_with_noise_backward", + "mkldnn_adaptive_avg_pool2d_backward", + "_adaptive_avg_pool2d_backward", + "_adaptive_avg_pool3d_backward", + "isinf", + "linalg_lu_solve", + "linalg_vecdot", + "linalg_matrix_exp", + "linalg_eigvalsh", + "_test_warn_in_autograd", + "_test_autograd_multiple_dispatch_view", + "_test_autograd_multiple_dispatch_view_copy", + "_segment_reduce", + "_segment_reduce_backward", + "_fw_primal_copy", + "_make_dual_copy", + "view_as_real_copy", + "view_as_complex_copy", + "_conj_copy", + "_neg_view_copy", + "diagonal_copy", + "detach_copy", + "squeeze_copy", + "t_copy", + "unsqueeze_copy", + "_indices_copy", + "_values_copy", + "indices_copy", + "values_copy", + "crow_indices_copy", + "col_indices_copy", + "ccol_indices", + "ccol_indices_copy", + "row_indices", + "row_indices_copy", + "unfold_copy", + "alias_copy", + "_triton_multi_head_attention", + "special_airy_ai", + "special_bessel_j0", + "special_bessel_j1", + "special_bessel_y0", + "special_bessel_y1", + "special_chebyshev_polynomial_t", + "special_chebyshev_polynomial_u", + "special_chebyshev_polynomial_v", + "special_chebyshev_polynomial_w", + "special_hermite_polynomial_h", + "special_hermite_polynomial_he", + "special_laguerre_polynomial_l", + "special_legendre_polynomial_p", + "special_modified_bessel_i0", + "special_modified_bessel_i1", + "special_modified_bessel_k0", + "special_modified_bessel_k1", + "special_scaled_modified_bessel_k0", + "special_scaled_modified_bessel_k1", + "special_shifted_chebyshev_polynomial_t", + "special_shifted_chebyshev_polynomial_u", + "special_shifted_chebyshev_polynomial_v", + "special_shifted_chebyshev_polynomial_w", + "special_spherical_bessel_j0", + "_foobar", + "_nested_tensor_strides", + "_nested_tensor_storage_offsets", + "_nested_get_values", # no CPU backend + "_nested_get_values_copy", # no CPU backend + "_nested_view_from_jagged", # testing needs to be patched + "_nested_view_from_jagged_copy", # testing needs to be patched + "_nested_view_from_buffer", # testing needs to be patched + "_nested_view_from_buffer_copy", # testing needs to be patched + "_int_mm", # testing needs to be patched + "_to_sparse_csc", # testing needs to be patched + "_to_sparse_csr", # testing needs to be patched + "segment_reduce", # testing needs to be patched + ) +) + + +def is_supported(g: NativeFunctionsGroup | NativeFunctionsViewGroup) -> bool: + base_op_name = "" + func = None + if isinstance(g, NativeFunctionsViewGroup): + base_op_name = g.view.root_name + func = g.view.func + else: + base_op_name = g.out.func.name.name.base + func = g.out.func + if config.is_hand_written(g): + logger.info("HAND WRITTEN: %s", base_op_name) + return False + if base_op_name in BLOCKED_OPS: + logger.info("BLOCKED: %s", base_op_name) + return False + for arg in func.schema_order_arguments(): + maybe_method = ivalue_type_conversion_method(arg.type) + if not maybe_method: + # Type converting is unsupported yet. + logger.info("NOT SUPPORTED TYPE CONVERTING: %s", func) + return False + + if isinstance(g, NativeFunctionsViewGroup): + # TODO: stop doing type tests by converting to C++ and then testing + # the string, just test the dang thing directly + if "at::Tensor" != cpp.returns_type(func.returns, symint=False).cpp_type(): + # Returns a non-Tensor value. + logger.info("NON-TENSOR RET TYPE: %s", func) + return False + return True + + # For out variant ops, we need to check the arguments of its functional func. + for arg in g.functional.func.schema_order_arguments(): + maybe_method = ivalue_type_conversion_method(arg.type) + if not maybe_method: + # Type converting is unsupported yet. + logger.info("NOT SUPPORTED TYPE CONVERTING: %s", g.functional.func) + return False + + if not g.structured: + # In case of unstructured op, we check if it has out variant implementation. + # The out variant implementation satisfies the minimum requirement that it has the output tensor as the last + # parameter. + if ( + not hasattr(g, "out") + or not str(func).endswith("Tensor(a!) out) -> Tensor(a!)") + or not str(func.name).endswith(".out") + ): + return False + # TODO: stop type testing by converting to C++ + if "at::Tensor &" != cpp.returns_type(func.returns, symint=False).cpp_type(): + logger.info("NON_TENSOR RET TYPE: %s", func) + return False + if has_alias(func.arguments.non_out): + # This op may create an alias of inputs. + logger.info("INPUTS ALIAS: %s", base_op_name) + return False + return True + + +def ivalue_type_conversion_method( + arg_type: BaseType | OptionalType | Type, +) -> tuple[bool, str] | None: + """ + Return the method call expression of `c10::ivalue' to convert its contained value to + the expected value of `arg_type` type. For example, for `arg_type` == BaseTy.Tensor, + this function returns ".toTensor()", so that it can be appended to the ivalue's + variable name to get the value of the expected type. + """ + type_conversion_methods = { + BaseTy.Tensor: ((True, "toTensor()"), (False, "toOptional()")), + BaseTy.int: ((False, "toInt()"), (False, "toOptional()")), + BaseTy.bool: ((False, "toBool()"), (False, "toOptional()")), + BaseTy.Scalar: ((False, "toScalar()"), (False, "toOptional()")), + BaseTy.ScalarType: ( + (False, "toScalarType()"), + (False, "toOptional()"), + ), + BaseTy.str: ( + (False, "toStringView()"), + (False, "toOptional()"), + (False, "toOptional<::std::string_view>()"), + ), + } + + base_ty_object = None + if isinstance(arg_type, BaseType): + base_ty_object = arg_type.name + elif isinstance(arg_type, OptionalType): + if not isinstance(arg_type.elem, BaseType): + # ListType is currently unsupported. + return None + base_ty_object = arg_type.elem.name + else: + return None + + if base_ty_object not in type_conversion_methods: + return None + methods = type_conversion_methods[base_ty_object] + if isinstance(arg_type, BaseType): + return methods[0] + return methods[1] + + +should_use_int_tensor_ops_ = frozenset( + ( + "bitwise_not", + "bitwise_and", + "bitwise_or", + "bitwise_xor", + "bitwise_left_shift", + "bitwise_right_shift", + "gcd", + "lcm", + "scatter", + "gather", + "_convert_indices_from_coo_to_csr", + "_convert_indices_from_csr_to_coo", + ) +) +should_use_complex_tensor_ops_ = frozenset(("view_as_real", "imag", "_conj")) + + +def should_use_int_tensor(op_name: str) -> bool: + return op_name in should_use_int_tensor_ops_ + + +def should_use_complex_tensor(op_name: str) -> bool: + return op_name in should_use_complex_tensor_ops_ + + +test_tensor_dim_ops_1_ = frozenset( + ( + "addmv", + "index_add", + "_convert_indices_from_coo_to_csr", + "_convert_indices_from_csr_to_coo", + "nll_loss_backward", + "dot", + "vdot", + "outer", + "ger", + ) +) +test_tensor_dim_ops_2_ = frozenset( + ("addmm", "mm", "nuclear_norm", "diag", "_addmm_activation", "matrix_H", "t") +) + + +def test_tensor_dim(op_name: str) -> int: + if op_name in test_tensor_dim_ops_1_: + return 1 + if op_name in test_tensor_dim_ops_2_: + return 2 + return 3 + + +test_tensor_shapes_string = '{"view_as_complex": "{2, 2}"}' +test_tensor_shape_json: dict[str, str] = json.loads(test_tensor_shapes_string) + + +def test_tensor_shape(op_name: str) -> str: + if op_name in test_tensor_shape_json: + return test_tensor_shape_json[op_name] + else: + return "" + + +def test_value_expression( + arg_type: BaseType | OptionalType | Type, index: int, op_name: str +) -> str: + tensor_size_ex = test_tensor_shape(op_name) + if tensor_size_ex == "": + num_tensors = 16 if index == 0 else 64 + num_dim = test_tensor_dim(op_name) + size_per_dim = math.ceil(num_tensors / float(num_dim)) + size_per_dim += size_per_dim % 2 + tensor_size_ex = "{{{}}}".format(",".join([f"{size_per_dim}"] * num_dim)) + if should_use_int_tensor(op_name): + tensor_expression = f"at::randint(1, 100, {tensor_size_ex}, at::kInt)" + elif should_use_complex_tensor(op_name): + tensor_expression = f"at::randn({tensor_size_ex}, at::kComplexFloat)" + else: + tensor_expression = f"at::rand({tensor_size_ex})" + + value_expressions = { + BaseTy.Tensor: tensor_expression, + BaseTy.int: "1", + BaseTy.bool: "false", + BaseTy.Scalar: "2", + BaseTy.ScalarType: "at::ScalarType::Float", + BaseTy.str: '"floor"', + } + + base_ty_object = None + if isinstance(arg_type, BaseType): + base_ty_object = arg_type.name + else: + if not ( + isinstance(arg_type, OptionalType) and isinstance(arg_type.elem, BaseType) + ): + raise AssertionError( + f"Expected OptionalType with BaseType elem, got {type(arg_type)}" + ) + base_ty_object = arg_type.elem.name + if base_ty_object not in value_expressions: + raise AssertionError(f"Unexpected type: {base_ty_object}") + value_expression = value_expressions[base_ty_object] + return value_expression + + +def generate_test_value_definitions(schema: FunctionSchema, index: int) -> str: + if schema.is_out_fn(): + raise AssertionError(f"Expected non-out function, got {schema}") + schema_name = schema.name.name.base + arg_map = {} + for arg in schema.schema_order_arguments(): + test_value_exp = test_value_expression(arg.type, index, schema_name) + arg_map[arg.name] = test_value_exp + config.override_test_values(arg_map, schema_name, index) + arg_populations = [] + for arg_name, arg_value in arg_map.items(): + arg_populations.append(f"auto {arg_name}{index} = {arg_value}") + return ";\n ".join(arg_populations) + ";" + + +def generate_test_value_names(schema: FunctionSchema, index: int) -> str: + if schema.is_out_fn(): + raise AssertionError(f"Expected non-out function, got {schema}") + return ",".join(f"{arg.name}{index}" for arg in schema.schema_order_arguments()) + + +generate_test_ir_arguments_base_ty_to_type_str_ = { + BaseTy.Tensor: "Tensor", + BaseTy.int: "int", + BaseTy.float: "float", + BaseTy.str: "str", + BaseTy.Scalar: "int", + BaseTy.ScalarType: "int", + BaseTy.bool: "bool", +} + + +def generate_test_ir_arguments( + schema: FunctionSchema, +) -> list[tuple[str, str | None]]: + def ir_argument(arg: Argument) -> tuple[str, str | None]: + t = arg.type + add_optional = False + if isinstance(t, OptionalType): + t = t.elem + add_optional = True + if not isinstance(t, BaseType): + raise AssertionError(f"Expected BaseType, got {type(t)}") + type_str = None + if t.name in generate_test_ir_arguments_base_ty_to_type_str_: + type_str = generate_test_ir_arguments_base_ty_to_type_str_[t.name] + if type_str and add_optional: + type_str = f"{type_str}?" + return ("%" + arg.name, type_str) + + return [ir_argument(arg) for arg in schema.schema_order_arguments()] + + +def generate_arg_extraction(schema: FunctionSchema) -> str: + arg_populations = [] + for i, arg in enumerate(schema.schema_order_arguments()): + maybe_method = ivalue_type_conversion_method(arg.type) + if not maybe_method: + raise AssertionError( + f"No type conversion method for {arg.name}: {arg.type}" + ) + is_reference, type_conversion_method = maybe_method + reference = "&" if is_reference else "" + arg_populations.append( + f"const auto{reference} {arg.name} = p_node->Input({i}).{type_conversion_method}" + ) + return ";\n ".join(arg_populations) + ";" + + +def get_kernel_name(g: NativeFunctionsGroup, backend_index: BackendIndex) -> str: + kernel = backend_index.get_kernel(g.functional) + if g.structured or kernel is None: + return cpp.name(g.functional.func) + return kernel.kernel + + +def get_out_kernel_name(g: NativeFunctionsGroup, backend_index: BackendIndex) -> str: + kernel = backend_index.get_kernel(g.out) + if g.structured or kernel is None: + return cpp.name(g.out.func) + return kernel.kernel + + +def generate_non_out_variant_call( + g: NativeFunctionsGroup, backend_index: BackendIndex +) -> str: + schema = g.functional.func + if schema.is_out_fn(): + raise AssertionError(f"Expected non-out function, got {schema}") + kernel_name = get_kernel_name(g, backend_index) + arg_names = (arg.name for arg in schema.schema_order_arguments()) + namespace_name = "cpu" if g.structured else "native" + return f"at::{namespace_name}::{kernel_name}({','.join(arg_names)})" + + +def generate_call_to_view_ops( + g: NativeFunctionsViewGroup, backend_index: BackendIndex +) -> str: + schema = g.view.func + kernel_name = cpp.name(schema) + kernel = backend_index.get_kernel(g.view) + if kernel: + kernel_name = kernel.kernel + arg_names = (arg.name for arg in schema.schema_order_arguments()) + namespace_name = "native" + return f"at::{namespace_name}::{kernel_name}({','.join(arg_names)})" + + +def generate_out_variant_call( + g: NativeFunctionsGroup, backend_index: BackendIndex +) -> str: + schema = g.out.func + if not schema.is_out_fn(): + raise AssertionError(f"Expected out function, got {schema}") + arg_names = [] + kernel_name = get_out_kernel_name(g, backend_index) + if g.structured: + # structured op starts with the output tensor argument. + arg_names = [out_arg.name for out_arg in schema.arguments.out] + else: + arg_names = [] + for arg in schema.arguments.non_out: + if isinstance(arg, SelfArgument): + arg_names.append(arg.argument.name) + else: + if not isinstance(arg, Argument): + raise AssertionError(f"Expected Argument, got {type(arg)}") + arg_names.append(arg.name) + if not g.structured: + if len(schema.arguments.out) != 1: + raise AssertionError( + f"Expected 1 out argument, got {len(schema.arguments.out)}" + ) + arg_names.append(schema.arguments.out[0].name) + cpp_arg_names = ",".join(arg_names) + namespace_name = "cpu" if g.structured else "native" + return f"at::{namespace_name}::{kernel_name}({cpp_arg_names})" + + +no_memory_resize_ops = frozenset( + ( + "isin.Scalar_Tensor", + "index_add", + "dot", + "vdot", + "nuclear_norm", + "histc", + "l1_loss", + "multi_margin_loss", + "multilabel_margin_loss", + "nll_loss", + "nll_loss2d", + "prod", + ) +) + + +def should_check_resize(schema: FunctionSchema) -> bool: + schema_str = str(schema) + type_variant_op_name = schema_str[: schema_str.find("(")] + return type_variant_op_name not in no_memory_resize_ops + + +def op_name_from_group(g: NativeFunctionsGroup) -> str: + return g.functional.func.name.name.base + + +class GenOpDispatcher: + def out_variant( + self, groups: Sequence[NativeFunctionsGroup], backend_index: BackendIndex + ) -> str: + if not groups: + return "" + generated_type_variants = [] + for g in groups: + with native_function_manager(g): + if not is_supported(g): + raise AssertionError(f"Unsupported function group: {g}") + if not isinstance(g, NativeFunctionsGroup): + raise AssertionError( + f"Expected NativeFunctionsGroup, got {type(g)}" + ) + generated_type_variant = self.out_variant_op_generator(g, backend_index) + generated_type_variants.append(generated_type_variant) + op_name = op_name_from_group(groups[0]) + body = "\n".join(generated_type_variants) + generated = f""" +REGISTER_OPERATOR_FUNCTOR( + aten::{op_name}, + aten_{op_name}, + [](Node* n) -> SROperator {{ + {body} + LogAndDumpSchema(n); + return nullptr; + }}) +""" + return generated + + def view( + self, groups: Sequence[NativeFunctionsViewGroup], backend_index: BackendIndex + ) -> str: + if not groups: + return "" + generated_type_variants = [] + for g in groups: + with native_function_manager(g): + if not is_supported(g): + raise AssertionError(f"Unsupported view group: {g}") + if not isinstance(g, NativeFunctionsViewGroup): + raise AssertionError( + f"Expected NativeFunctionsViewGroup, got {type(g)}" + ) + generated_type_variant = self.view_op_generator(g, backend_index) + generated_type_variants.append(generated_type_variant) + op_name = config.func_name_base_str(groups[0]) + body = "\n".join(generated_type_variants) + generated = f""" +REGISTER_NATIVE_OPERATOR_FUNCTOR( + aten::{op_name}, + aten_{op_name}, + [](Node* n) -> SROperator {{ + {body} + LogAndDumpSchema(n); + return nullptr; + }}); +""" + return generated + + def out_variant_op_generator( + self, g: NativeFunctionsGroup, backend_index: BackendIndex + ) -> str: + functional = g.functional + schema = str(functional.func) + populated_argument = generate_arg_extraction(g.functional.func) + functional_variant_call = generate_non_out_variant_call(g, backend_index) + if len(g.out.func.arguments.out) != 1: + raise AssertionError( + f"Expected 1 out argument, got {len(g.out.func.arguments.out)}" + ) + out_variable_name = str(g.out.func.arguments.out[0].name) + out_variant_call = generate_out_variant_call(g, backend_index) + generated = f""" + if (n->matches(torch::schema("aten::{schema}"))) {{ + return [](ProcessedNode* p_node) {{ + {populated_argument} + if (p_node->Output(0).isNone()) {{ + p_node->Output(0) = {functional_variant_call}; + return; + }} + auto& {out_variable_name} = p_node->Output(0).toTensor(); + fastResizeToZero({out_variable_name}); + {out_variant_call}; + }}; + }}""" + return generated + + def view_op_generator( + self, g: NativeFunctionsViewGroup, backend_index: BackendIndex + ) -> str: + schema = str(g.view.func) + populated_argument = generate_arg_extraction(g.view.func) + functional_variant_call = generate_call_to_view_ops(g, backend_index) + generated = f""" + if (n->matches(torch::schema("aten::{schema}"))) {{ + return [](ProcessedNode* p_node) {{ + {populated_argument} + p_node->Output(0) = {functional_variant_call}; + }}; + }}""" + return generated + + +class GenOpTestCase: + def out_variant(self, groups: Sequence[NativeFunctionsGroup]) -> str: + if not groups: + return "" + generated_type_variants = [] + for g in groups: + with native_function_manager(g): + if not is_supported(g): + raise AssertionError(f"Unsupported function group: {g}") + if not isinstance(g, NativeFunctionsGroup): + raise AssertionError( + f"Expected NativeFunctionsGroup, got {type(g)}" + ) + generated_type_variant = self.out_variant_op_test_case_generator(g) + generated_type_variants.append(generated_type_variant) + return "\n".join(generated_type_variants) + + def view(self, groups: Sequence[NativeFunctionsViewGroup]) -> str: + if not groups: + return "" + generated_type_variants = [] + for g in groups: + with native_function_manager(g): + if not is_supported(g): + raise AssertionError(f"Unsupported view group: {g}") + if not isinstance(g, NativeFunctionsViewGroup): + raise AssertionError( + f"Expected NativeFunctionsViewGroup, got {type(g)}" + ) + generated_type_variant = self.view_op_test_case_generator(g) + generated_type_variants.append(generated_type_variant) + return "\n".join(generated_type_variants) + + def out_variant_op_test_case_generator(self, g: NativeFunctionsGroup) -> str: + schema = g.functional.func + schema_str = str(schema) + if schema_str.find("(") <= 0: + raise AssertionError(f"Invalid schema string: {schema_str}") + type_variant_op_name = schema_str[: schema_str.find("(")].replace(".", "_") + op_name = op_name_from_group(g) + if not type_variant_op_name.startswith(op_name): + raise AssertionError( + f"Type variant op name {type_variant_op_name} doesn't start with {op_name}" + ) + + arg_types = generate_test_ir_arguments(schema) + arg_declarations = ", ".join( + ( + arg_name if arg_type is None else f"{arg_name}: {arg_type}" + for arg_name, arg_type in arg_types + ) + ) + arg_names = ", ".join((arg_name for arg_name, _ in arg_types)) + if not ( + len(schema.returns) == 1 + and isinstance(schema.returns[0].type, BaseType) + and schema.returns[0].type.name is BaseTy.Tensor + ): + raise AssertionError(f"Expected single Tensor return, got {schema.returns}") + test_value_definitions = generate_test_value_definitions(schema, 0) + test_value_names = generate_test_value_names(schema, 0) + test_value_definitions2 = generate_test_value_definitions(schema, 1) + test_value_names2 = generate_test_value_names(schema, 1) + check_resize = "true" if should_check_resize(schema) else "false" + generated = f""" +TEST(StaticRuntime, autogen_{type_variant_op_name}) {{ + const std::string script = R"IR( + graph({arg_declarations}): + %bias: None = prim::Constant() + %ret = aten::{op_name}({arg_names}) + %cloned = aten::clone(%ret, %bias) + return (%cloned) + )IR"; + + {test_value_definitions} + std::vector args{{{test_value_names}}}; + testStaticRuntime(script, args, {{}}, /*use_allclose=*/false, /*use_equalnan=*/false, /*check_resize=*/{check_resize}); + + {test_value_definitions2} + std::vector args2{{{test_value_names2}}}; + testStaticRuntime(script, args, args2, /*use_allclose=*/false, /*use_equalnan=*/false, /*check_resize=*/{check_resize}); + +}} +""" + return generated + + def view_op_test_case_generator(self, g: NativeFunctionsViewGroup) -> str: + schema = g.view.func + schema_str = str(schema) + if schema_str.find("(") <= 0: + raise AssertionError(f"Invalid schema string: {schema_str}") + type_variant_op_name = schema_str[: schema_str.find("(")].replace(".", "_") + op_name = g.view.root_name + if not type_variant_op_name.startswith(op_name): + raise AssertionError( + f"Type variant op name {type_variant_op_name} doesn't start with {op_name}" + ) + + arg_types = generate_test_ir_arguments(schema) + arg_declarations = ", ".join( + ( + arg_name if arg_type is None else f"{arg_name}: {arg_type}" + for arg_name, arg_type in arg_types + ) + ) + arg_names = ", ".join((arg_name for arg_name, _ in arg_types)) + if not ( + len(schema.returns) == 1 + and isinstance(schema.returns[0].type, BaseType) + and schema.returns[0].type.name is BaseTy.Tensor + ): + raise AssertionError(f"Expected single Tensor return, got {schema.returns}") + test_value_definitions = generate_test_value_definitions(schema, 0) + test_value_names = generate_test_value_names(schema, 0) + generated = f""" +TEST(StaticRuntime, autogen_{type_variant_op_name}) {{ + const std::string script = R"IR( + graph({arg_declarations}): + %bias: None = prim::Constant() + %ret = aten::{op_name}({arg_names}) + %cloned = aten::clone(%ret, %bias) + return (%cloned) + )IR"; + + {test_value_definitions} + std::vector args{{{test_value_names}}}; + testStaticRuntime(script, args); +}} +""" + + return generated diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..46deafe41b021ad19ce8ca7322e1a3ba3374c3f4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/utils.py @@ -0,0 +1,523 @@ +from __future__ import annotations + +import contextlib +import functools +import hashlib +import os +import re +import sys +import textwrap +from dataclasses import is_dataclass +from enum import auto, Enum +from pathlib import Path +from pprint import pformat +from typing import Any, Generic, TYPE_CHECKING, TypeVar +from typing_extensions import assert_never, Self + +from torchgen.code_template import CodeTemplate + + +if TYPE_CHECKING: + from argparse import Namespace + from collections.abc import Callable, Iterable, Iterator, Sequence + + +TORCHGEN_ROOT = Path(__file__).absolute().parent +REPO_ROOT = TORCHGEN_ROOT.parent + + +# Many of these functions share logic for defining both the definition +# and declaration (for example, the function signature is the same), so +# we organize them into one function that takes a Target to say which +# code we want. +# +# This is an OPEN enum (we may add more cases to it in the future), so be sure +# to explicitly specify with Literal[Target.XXX] or Literal[Target.XXX, Target.YYY] +# what targets are valid for your use. +class Target(Enum): + # top level namespace (not including at) + DEFINITION = auto() + DECLARATION = auto() + # TORCH_LIBRARY(...) { ... } + REGISTRATION = auto() + # namespace { ... } + ANONYMOUS_DEFINITION = auto() + # namespace cpu { ... } + NAMESPACED_DEFINITION = auto() + NAMESPACED_DECLARATION = auto() + + +# Matches "foo" in "foo, bar" but not "foobar". Used to search for the +# occurrence of a parameter in the derivative formula +IDENT_REGEX = r"(^|\W){}($|\W)" + + +# TODO: Use a real parser here; this will get bamboozled +def split_name_params(schema: str) -> tuple[str, list[str]]: + m = re.match(r"(\w+)(\.\w+)?\((.*)\)", schema) + if m is None: + raise RuntimeError(f"Unsupported function schema: {schema}") + name, _, params = m.groups() + return name, params.split(", ") + + +T = TypeVar("T") +S = TypeVar("S") + +# These two functions purposely return generators in analogy to map() +# so that you don't mix up when you need to list() them + + +# Map over function that may return None; omit Nones from output sequence +def mapMaybe(func: Callable[[T], S | None], xs: Iterable[T]) -> Iterator[S]: + for x in xs: + r = func(x) + if r is not None: + yield r + + +# Map over function that returns sequences and cat them all together +def concatMap(func: Callable[[T], Sequence[S]], xs: Iterable[T]) -> Iterator[S]: + for x in xs: + yield from func(x) + + +# Conveniently add error context to exceptions raised. Lets us +# easily say that an error occurred while processing a specific +# context. +@contextlib.contextmanager +def context(msg_fn: Callable[[], str]) -> Iterator[None]: + try: + yield + except Exception as e: + # TODO: this does the wrong thing with KeyError + msg = msg_fn() + msg = textwrap.indent(msg, " ") + msg = f"{e.args[0]}\n{msg}" if e.args else msg + e.args = (msg,) + e.args[1:] + raise + + +@functools.cache +def _read_template(template_fn: str) -> CodeTemplate: + return CodeTemplate.from_file(template_fn) + + +# String hash that's stable across different executions, unlike builtin hash +def string_stable_hash(s: str) -> int: + sha1 = hashlib.sha1(s.encode("latin1"), usedforsecurity=False).digest() + return int.from_bytes(sha1, byteorder="little") + + +# A small abstraction for writing out generated files and keeping track +# of what files have been written (so you can write out a list of output +# files) +class FileManager: + def __init__( + self, + install_dir: str | Path, + template_dir: str | Path, + dry_run: bool, + ) -> None: + self.install_dir = Path(install_dir) + self.template_dir = Path(template_dir) + self.files: set[Path] = set() + self.dry_run = dry_run + + @property + def filenames(self) -> frozenset[str]: + return frozenset({file.as_posix() for file in self.files}) + + def _write_if_changed(self, filename: str | Path, contents: str) -> None: + file = Path(filename) + old_contents: str | None = None + try: + old_contents = file.read_text(encoding="utf-8") + except OSError: + pass + if contents != old_contents: + # Create output directory if it doesn't exist + file.parent.mkdir(parents=True, exist_ok=True) + file.write_text(contents, encoding="utf-8") + + # Read from template file and replace pattern with callable (type could be dict or str). + def substitute_with_template( + self, + template_fn: str | Path, + env_callable: Callable[[], str | dict[str, Any]], + ) -> str: + if Path(template_fn).is_absolute(): + raise AssertionError(f"template_fn must be relative: {template_fn}") + template_path = self.template_dir / template_fn + env = env_callable() + if isinstance(env, dict): + if "generated_comment" not in env: + generator_default = TORCHGEN_ROOT / "gen.py" + try: + generator = Path( + sys.modules["__main__"].__file__ or generator_default + ).absolute() + except (KeyError, AttributeError): + generator = generator_default.absolute() + + try: + generator_path = generator.relative_to(REPO_ROOT).as_posix() + except ValueError: + generator_path = generator.name + + env = { + **env, # copy the original dict instead of mutating it + "generated_comment": ( + "@" + f"generated by {generator_path} from {template_fn}" + ), + } + template = _read_template(template_path) + substitute_out = template.substitute(env) + # Ensure an extra blank line between the class/function definition + # and the docstring of the previous class/function definition. + # NB: It is generally not recommended to have docstrings in pyi stub + # files. But if there are any, we need to ensure that the file + # is properly formatted. + return re.sub( + r''' + (""")\n+ # match triple quotes + ( + (\s*@.+\n)* # match decorators if any + \s*(class|def) # match class/function definition + ) + ''', + r"\g<1>\n\n\g<2>", + substitute_out, + flags=re.VERBOSE, + ) + if isinstance(env, str): + return env + assert_never(env) + + def write_with_template( + self, + filename: str | Path, + template_fn: str | Path, + env_callable: Callable[[], str | dict[str, Any]], + ) -> None: + filename = Path(filename) + if filename.is_absolute(): + raise AssertionError(f"filename must be relative: {filename}") + file = self.install_dir / filename + if file in self.files: + raise AssertionError(f"duplicate file write {file}") + self.files.add(file) + if not self.dry_run: + substitute_out = self.substitute_with_template( + template_fn=template_fn, + env_callable=env_callable, + ) + self._write_if_changed(filename=file, contents=substitute_out) + + def write( + self, + filename: str | Path, + env_callable: Callable[[], str | dict[str, Any]], + ) -> None: + self.write_with_template(filename, filename, env_callable) + + def write_sharded( + self, + filename: str | Path, + items: Iterable[T], + *, + key_fn: Callable[[T], str], + env_callable: Callable[[T], dict[str, list[str]]], + num_shards: int, + base_env: dict[str, Any] | None = None, + sharded_keys: set[str], + ) -> None: + self.write_sharded_with_template( + filename, + filename, + items, + key_fn=key_fn, + env_callable=env_callable, + num_shards=num_shards, + base_env=base_env, + sharded_keys=sharded_keys, + ) + + def write_sharded_with_template( + self, + filename: str | Path, + template_fn: str | Path, + items: Iterable[T], + *, + key_fn: Callable[[T], str], + env_callable: Callable[[T], dict[str, list[str]]], + num_shards: int, + base_env: dict[str, Any] | None = None, + sharded_keys: set[str], + ) -> None: + file = Path(filename) + if file.is_absolute(): + raise AssertionError(f"filename must be relative: {filename}") + everything: dict[str, Any] = {"shard_id": "Everything"} + shards: list[dict[str, Any]] = [ + {"shard_id": f"_{i}"} for i in range(num_shards) + ] + all_shards = [everything] + shards + + if base_env is not None: + for shard in all_shards: + shard.update(base_env) + + for key in sharded_keys: + for shard in all_shards: + if key in shard: + if not isinstance(shard[key], list): + raise AssertionError("sharded keys in base_env must be a list") + shard[key] = shard[key].copy() + else: + shard[key] = [] + + def merge_env(into: dict[str, list[str]], from_: dict[str, list[str]]) -> None: + for k, v in from_.items(): + if k not in sharded_keys: + raise AssertionError(f"undeclared sharded key {k}") + into[k] += v + + if self.dry_run: + # Dry runs don't write any templates, so incomplete environments are fine + items = () + + for item in items: + key = key_fn(item) + sid = string_stable_hash(key) % num_shards + env = env_callable(item) + + merge_env(shards[sid], env) + merge_env(everything, env) + + for shard in all_shards: + shard_id = shard["shard_id"] + self.write_with_template( + file.with_stem(f"{file.stem}{shard_id}"), + template_fn, + lambda: shard, + ) + + # filenames is used to track compiled files, but FooEverything.cpp isn't meant to be compiled + self.files.discard(self.install_dir / file.with_stem(f"{file.stem}Everything")) + + def write_outputs(self, variable_name: str, filename: str | Path) -> None: + """Write a file containing the list of all outputs which are generated by this script.""" + content = "\n".join( + ( + "set(", + variable_name, + # Use POSIX paths to avoid invalid escape sequences on Windows + *(f' "{file.as_posix()}"' for file in sorted(self.files)), + ")", + ) + ) + self._write_if_changed(filename, content) + + def template_dir_for_comments(self) -> str: + """ + This needs to be deterministic. The template dir is an absolute path + that varies across builds. So, just use the path relative to this file, + which will point to the codegen source but will be stable. + """ + return os.path.relpath(self.template_dir, os.path.dirname(__file__)) + + +# Helper function to generate file manager +def make_file_manager( + options: Namespace, + install_dir: str | Path | None = None, +) -> FileManager: + template_dir = os.path.join(options.source_path, "templates") + install_dir = install_dir if install_dir else options.install_dir + return FileManager( + install_dir=install_dir, + template_dir=template_dir, + dry_run=options.dry_run, + ) + + +# Helper function to create a pretty representation for dataclasses +def dataclass_repr( + obj: Any, + indent: int = 0, + width: int = 80, +) -> str: + return pformat(obj, indent, width) + + +def _format_dict( + attr: dict[Any, Any], + indent: int, + width: int, + curr_indent: int, +) -> str: + curr_indent += indent + 3 + dict_repr = [] + for k, v in attr.items(): + k_repr = repr(k) + v_str = ( + pformat(v, indent, width, curr_indent + len(k_repr)) + if is_dataclass(v) + else repr(v) + ) + dict_repr.append(f"{k_repr}: {v_str}") + + return _format(dict_repr, indent, width, curr_indent, "{", "}") + + +def _format_list( + attr: list[Any] | set[Any] | tuple[Any, ...], + indent: int, + width: int, + curr_indent: int, +) -> str: + curr_indent += indent + 1 + list_repr = [ + pformat(l, indent, width, curr_indent) if is_dataclass(l) else repr(l) + for l in attr + ] + start, end = ("[", "]") if isinstance(attr, list) else ("(", ")") + return _format(list_repr, indent, width, curr_indent, start, end) + + +def _format( + fields_str: list[str], + indent: int, + width: int, + curr_indent: int, + start: str, + end: str, +) -> str: + delimiter, curr_indent_str = "", "" + # if it exceed the max width then we place one element per line + if len(repr(fields_str)) >= width: + delimiter = "\n" + curr_indent_str = " " * curr_indent + + indent_str = " " * indent + body = f", {delimiter}{curr_indent_str}".join(fields_str) + return f"{start}{indent_str}{body}{end}" + + +class NamespaceHelper: + """A helper for constructing the namespace open and close strings for a nested set of namespaces. + + e.g. for namespace_str torch::lazy, + + prologue: + namespace torch { + namespace lazy { + + epilogue: + } // namespace lazy + } // namespace torch + """ + + def __init__( + self, + namespace_str: str, + entity_name: str = "", + max_level: int = 2, + ) -> None: + # cpp_namespace can be a colon joined string such as torch::lazy + cpp_namespaces = namespace_str.split("::") + if len(cpp_namespaces) > max_level: + raise AssertionError( + f"Codegen doesn't support more than {max_level} level(s) of " + f"custom namespace. Got {namespace_str}." + ) + self.cpp_namespace_ = namespace_str + self.prologue_ = "\n".join([f"namespace {n} {{" for n in cpp_namespaces]) + self.epilogue_ = "\n".join( + [f"}} // namespace {n}" for n in reversed(cpp_namespaces)] + ) + self.namespaces_ = cpp_namespaces + self.entity_name_ = entity_name + + @staticmethod + def from_namespaced_entity( + namespaced_entity: str, + max_level: int = 2, + ) -> NamespaceHelper: + """ + Generate helper from nested namespaces as long as class/function name. E.g.: "torch::lazy::add" + """ + names = namespaced_entity.split("::") + entity_name = names[-1] + namespace_str = "::".join(names[:-1]) + return NamespaceHelper( + namespace_str=namespace_str, entity_name=entity_name, max_level=max_level + ) + + @property + def prologue(self) -> str: + return self.prologue_ + + @property + def epilogue(self) -> str: + return self.epilogue_ + + @property + def entity_name(self) -> str: + return self.entity_name_ + + # Only allow certain level of namespaces + def get_cpp_namespace(self, default: str = "") -> str: + """ + Return the namespace string from joining all the namespaces by "::" (hence no leading "::"). + Return default if namespace string is empty. + """ + return self.cpp_namespace_ if self.cpp_namespace_ else default + + +class OrderedSet(Generic[T]): + storage: dict[T, None] + + def __init__(self, iterable: Iterable[T] | None = None) -> None: + if iterable is None: + self.storage = {} + else: + self.storage = dict.fromkeys(iterable) + + def __contains__(self, item: T) -> bool: + return item in self.storage + + def __iter__(self) -> Iterator[T]: + return iter(self.storage.keys()) + + def update(self, items: OrderedSet[T]) -> None: + self.storage.update(items.storage) + + def add(self, item: T) -> None: + self.storage[item] = None + + def copy(self) -> OrderedSet[T]: + ret: OrderedSet[T] = OrderedSet() + ret.storage = self.storage.copy() + return ret + + @staticmethod + def union(*args: OrderedSet[T]) -> OrderedSet[T]: + ret = args[0].copy() + for s in args[1:]: + ret.update(s) + return ret + + def __or__(self, other: OrderedSet[T]) -> OrderedSet[T]: + return OrderedSet.union(self, other) + + def __ior__(self, other: OrderedSet[T]) -> Self: + self.update(other) + return self + + def __eq__(self, other: object) -> bool: + if isinstance(other, OrderedSet): + return self.storage == other.storage + else: + return set(self.storage.keys()) == other diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/yaml_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/yaml_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..96d67859289c03019964568cd1e15acccc7a0eba --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/torchgen/yaml_utils.py @@ -0,0 +1,27 @@ +# Safely load fast C Yaml loader/dumper if they are available +try: + from yaml import CSafeLoader as Loader +except ImportError: + from yaml import SafeLoader as Loader # type: ignore[assignment, misc] + +try: + from yaml import CSafeDumper as Dumper +except ImportError: + from yaml import SafeDumper as Dumper # type: ignore[assignment, misc] +YamlDumper = Dumper + + +# A custom loader for YAML that errors on duplicate keys. +# This doesn't happen by default: see https://github.com/yaml/pyyaml/issues/165 +class YamlLoader(Loader): + def construct_mapping(self, node, deep=False): # type: ignore[no-untyped-def] + mapping = [] + for key_node, value_node in node.value: + key = self.construct_object(key_node, deep=deep) # type: ignore[no-untyped-call] + if key in mapping: + raise AssertionError( + f"Found a duplicate key in the yaml. key={key}, line={node.start_mark.line}" + ) + mapping.append(key) + mapping = super().construct_mapping(node, deep=deep) # type: ignore[no-untyped-call] + return mapping diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/INSTALLER b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/METADATA b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..003eec0abb0f6903e2cc94b5d3daad7dd99afed7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/METADATA @@ -0,0 +1,82 @@ +Metadata-Version: 2.2 +Name: tqdm +Version: 4.68.3+computecanada +Summary: Fast, Extensible Progress Meter +Description-Content-Type: text/x-rst +Keywords: progressbar,progressmeter,progress,bar,meter,rate,eta,console,terminal,time +Classifier: Development Status :: 5 - Production/Stable +Classifier: Environment :: Console +Classifier: Environment :: MacOS X +Classifier: Environment :: Other Environment +Classifier: Environment :: Win32 (MS Windows) +Classifier: Environment :: X11 Applications +Classifier: Framework :: IPython +Classifier: Framework :: Jupyter +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Education +Classifier: Intended Audience :: End Users/Desktop +Classifier: Intended Audience :: Other Audience +Classifier: Intended Audience :: System Administrators +Classifier: Operating System :: MacOS +Classifier: Operating System :: MacOS :: MacOS X +Classifier: Operating System :: Microsoft +Classifier: Operating System :: Microsoft :: MS-DOS +Classifier: Operating System :: Microsoft :: Windows +Classifier: Operating System :: POSIX +Classifier: Operating System :: POSIX :: BSD +Classifier: Operating System :: POSIX :: BSD :: FreeBSD +Classifier: Operating System :: POSIX :: Linux +Classifier: Operating System :: POSIX :: SunOS/Solaris +Classifier: Operating System :: Unix +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: Implementation +Classifier: Programming Language :: Python :: Implementation :: IronPython +Classifier: Programming Language :: Python :: Implementation :: PyPy +Classifier: Programming Language :: Unix Shell +Classifier: Topic :: Desktop Environment +Classifier: Topic :: Education :: Computer Aided Instruction (CAI) +Classifier: Topic :: Education :: Testing +Classifier: Topic :: Office/Business +Classifier: Topic :: Other/Nonlisted Topic +Classifier: Topic :: Software Development :: Build Tools +Classifier: Topic :: Software Development :: Libraries +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Classifier: Topic :: Software Development :: Pre-processors +Classifier: Topic :: Software Development :: User Interfaces +Classifier: Topic :: System :: Installation/Setup +Classifier: Topic :: System :: Logging +Classifier: Topic :: System :: Monitoring +Classifier: Topic :: System :: Shells +Classifier: Topic :: Terminals +Classifier: Topic :: Utilities +Maintainer-email: tqdm developers +Project-URL: homepage, https://tqdm.github.io +Project-URL: repository, https://github.com/tqdm/tqdm +Project-URL: changelog, https://tqdm.github.io/releases +Project-URL: wiki, https://github.com/tqdm/tqdm/wiki +Requires-Python: >=3.8 +Requires-Dist: colorama; platform_system == "Windows" +Requires-Dist: requests; extra == "discord" +Requires-Dist: envwrap; extra == "discord" +Requires-Dist: slack-sdk; extra == "slack" +Requires-Dist: envwrap; extra == "slack" +Requires-Dist: requests; extra == "telegram" +Requires-Dist: envwrap; extra == "telegram" +Requires-Dist: ipywidgets>=6; extra == "notebook" +Provides-Extra: discord +Provides-Extra: slack +Provides-Extra: telegram +Provides-Extra: notebook +License-File: LICENCE +Dynamic: license-file + +MPL-2.0 AND MIT \ No newline at end of file diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/RECORD b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..7b60a3c177bb6be2c15700231ca42514ea72e0ee --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/RECORD @@ -0,0 +1,72 @@ +../../../bin/tqdm,sha256=5exfaBxqYzZUJvdRDQXvzwlKVamVm5Z9mk46rmNq2pA,258 +tqdm-4.68.3+computecanada.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +tqdm-4.68.3+computecanada.dist-info/METADATA,sha256=f2YC_XY1lJcILHeP3Or4qHYbQQ5_Ymx2aB4EeDsAiaA,3629 +tqdm-4.68.3+computecanada.dist-info/RECORD,, +tqdm-4.68.3+computecanada.dist-info/WHEEL,sha256=E4Ta9GSW8Vg2e11uZktRc054ObD9JXJ5hm6vpxg0WJE,87 +tqdm-4.68.3+computecanada.dist-info/entry_points.txt,sha256=ReJCH7Ui3Zyh6M16E4OhsZ1oU7WtMXCfbtoyBhGO29Y,39 +tqdm-4.68.3+computecanada.dist-info/licenses/LICENCE,sha256=_P-Hw6R86AKKhRKqGC1PzwrRyQVE7nXPmzQ2hMrBlN4,1985 +tqdm-4.68.3+computecanada.dist-info/top_level.txt,sha256=NLiUJNfmc9At15s7JURiwvqMEjUi9G5PMGRrmMYzNSM,5 +tqdm/__init__.py,sha256=9mQNYSSqP99JasubEC1POJLMmhkkBH6cJZxPIR5G2pQ,1572 +tqdm/__main__.py,sha256=bYt9eEaoRQWdejEHFD8REx9jxVEdZptECFsV7F49Ink,30 +tqdm/__pycache__/__init__.cpython-311.pyc,, +tqdm/__pycache__/__main__.cpython-311.pyc,, +tqdm/__pycache__/_main.cpython-311.pyc,, +tqdm/__pycache__/_monitor.cpython-311.pyc,, +tqdm/__pycache__/_tqdm.cpython-311.pyc,, +tqdm/__pycache__/_tqdm_gui.cpython-311.pyc,, +tqdm/__pycache__/_tqdm_notebook.cpython-311.pyc,, +tqdm/__pycache__/_tqdm_pandas.cpython-311.pyc,, +tqdm/__pycache__/_utils.cpython-311.pyc,, +tqdm/__pycache__/asyncio.cpython-311.pyc,, +tqdm/__pycache__/auto.cpython-311.pyc,, +tqdm/__pycache__/autonotebook.cpython-311.pyc,, +tqdm/__pycache__/cli.cpython-311.pyc,, +tqdm/__pycache__/dask.cpython-311.pyc,, +tqdm/__pycache__/gui.cpython-311.pyc,, +tqdm/__pycache__/keras.cpython-311.pyc,, +tqdm/__pycache__/notebook.cpython-311.pyc,, +tqdm/__pycache__/rich.cpython-311.pyc,, +tqdm/__pycache__/std.cpython-311.pyc,, +tqdm/__pycache__/tk.cpython-311.pyc,, +tqdm/__pycache__/utils.cpython-311.pyc,, +tqdm/__pycache__/version.cpython-311.pyc,, +tqdm/_main.py,sha256=9ySvgmi_2Sw4CAo5UDW0Q2dxfTryboEWGHohfCJz0sA,283 +tqdm/_monitor.py,sha256=1VcavuotYSN6K-1jbpNuy6HYtr3mwo5Nh_jZMQ-jg3A,4013 +tqdm/_tqdm.py,sha256=LfLCuJ6bpsVo9xilmtBXyEm1vGnUCFrliW85j3J-nD4,283 +tqdm/_tqdm_gui.py,sha256=03Hc8KayxJveieI5-0-2NGiDpLvw9jZekofJUV7CCwk,287 +tqdm/_tqdm_notebook.py,sha256=BuHiLuxu6uEfZFaPJW3RPpPaxaVctEQA3kdSJSDL1hw,307 +tqdm/_tqdm_pandas.py,sha256=c9jptUgigN6axRDhRd4Rif98Tmxeopc1nFNFhIpbFUE,888 +tqdm/_utils.py,sha256=CfLztxFY-QRiNfxQCqAb2N8E13uauIM95RSQiFlaEuE,559 +tqdm/asyncio.py,sha256=d6vZ88ZFq6zq3tAWpkInB2xITVKCNEBakudm_-8Xtf8,2758 +tqdm/auto.py,sha256=nDZflj6p2zKkjBCNBourrhS81zYfZy1_dQvbckrdW8o,871 +tqdm/autonotebook.py,sha256=2pNgGCJDsesNNyYnmSve8BNQLiGSgo97SsPoMXCrNRQ,1440 +tqdm/cli.py,sha256=t1j_PqI0SuMyhJGoBxMPCXY-UTbfjciaAQYxAU_f_V0,11011 +tqdm/completion.sh,sha256=j79KbSmpIj_E11jfTfBXrGnUTzKXVpQ1vGVQvsyDRl4,946 +tqdm/contrib/__init__.py,sha256=1GMYYE866WCvlHg5MdmzOvPsQj4ZuMlXzz8xcRcJD60,2479 +tqdm/contrib/__pycache__/__init__.cpython-311.pyc,, +tqdm/contrib/__pycache__/bells.cpython-311.pyc,, +tqdm/contrib/__pycache__/concurrent.cpython-311.pyc,, +tqdm/contrib/__pycache__/discord.cpython-311.pyc,, +tqdm/contrib/__pycache__/itertools.cpython-311.pyc,, +tqdm/contrib/__pycache__/logging.cpython-311.pyc,, +tqdm/contrib/__pycache__/slack.cpython-311.pyc,, +tqdm/contrib/__pycache__/telegram.cpython-311.pyc,, +tqdm/contrib/__pycache__/utils_worker.cpython-311.pyc,, +tqdm/contrib/bells.py,sha256=IaxC2YW57xHDinxdUDQHvSego0HSTdDTExQQENvSB_k,921 +tqdm/contrib/concurrent.py,sha256=K1yjloKS5WRNFyjLRth0DmU5PAnDbF0A-GD27N-J4a8,3986 +tqdm/contrib/discord.py,sha256=LPjxDbmfiYtw3OAm_WhtwbX9BVYDlPc8lJ96aY3YZDE,5337 +tqdm/contrib/itertools.py,sha256=MA9-Jvb3ToseppInvev2iF4bokN7GQy4Q4LgjUyYdA4,2949 +tqdm/contrib/logging.py,sha256=clW-AxaNwaynw9QGWVByjKgB3WJULaocoG_9C3cJ-qc,3791 +tqdm/contrib/slack.py,sha256=sC182v7xBkfGftnDeUmQALMOfwQzU3ZiNf1Q_XA1DjI,4049 +tqdm/contrib/telegram.py,sha256=ybTraGjUQx5m443TsX6pJ7Uxrpsh1m4a43BTWVl_4e4,5081 +tqdm/contrib/utils_worker.py,sha256=vaxyJyagYwOCQXLUMd86bCbI4sPqZysuMGbDSGmvEWo,1199 +tqdm/dask.py,sha256=9Ei58eVqTossRLhAfWyUFCduXYKjmLmwkaXIy-CHYfs,1319 +tqdm/gui.py,sha256=maQ7jgB5leWHwxMlbBv08hpe5lytvmH0hEARu22GRdU,5481 +tqdm/keras.py,sha256=ugd80-xCxMqnrbx_DEOpmj6c2Sh1xoVtQ8V9-bw29eo,4396 +tqdm/notebook.py,sha256=6xkUTcYRsFxN4_BSUYYZ51kAgI7ycnT0HLpP9HukEqg,10888 +tqdm/rich.py,sha256=YyMPkEHVyYUVUR3adJKbVX26iTmNKpNMf3DEqmm-m60,5021 +tqdm/std.py,sha256=_T-OKI1PBEKy_JWAOWCq9MOOYMjxqVq8lpr3z_YIVUA,57728 +tqdm/tk.py,sha256=3VsB1MDAkCItaWu7CmMb2CoTLMiqmDiEgNi3yiqQ5GY,6702 +tqdm/tqdm.1,sha256=MLg9qjeHThZpPEQbWVrf5P5gOXMKHE_X606FSOBI9eU,7538 +tqdm/utils.py,sha256=o_f0V32H_wrn0XmRno3j9pJ9xlp5xL0fHT3gtqL-k58,9746 +tqdm/version.py,sha256=0uH4TLe8I8H71x29dlmtEZBObLT3b5IW-nytF3V0IVw,233 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/WHEEL b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..b270261c9bb44553b2f648368d1bf5cb7baf8416 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: wheelfile 0.0.8 +Root-Is-Purelib: true +Tag: py3-none-any + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/entry_points.txt b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/entry_points.txt new file mode 100644 index 0000000000000000000000000000000000000000..540e60f4e073bc53a5f0a521a3639e0d80780af4 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/entry_points.txt @@ -0,0 +1,2 @@ +[console_scripts] +tqdm = tqdm.cli:main diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/licenses/LICENCE b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/licenses/LICENCE new file mode 100644 index 0000000000000000000000000000000000000000..194caf554f8f10ba4cac8a81b631a61d0d81f60d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/licenses/LICENCE @@ -0,0 +1,49 @@ +`tqdm` is a product of collaborative work. +Unless otherwise stated, all authors (see commit logs) retain copyright +for their respective work, and release the work under the MIT licence +(text below). + +Exceptions or notable authors are listed below +in reverse chronological order: + +* files: * + MPL-2.0 2015-2026 (c) Casper da Costa-Luis + [casperdcl](https://github.com/casperdcl). +* files: tqdm/_tqdm.py + MIT 2016 (c) [PR #96] on behalf of Google Inc. +* files: tqdm/_tqdm.py README.rst .gitignore + MIT 2013 (c) Noam Yorav-Raphael, original author. + +[PR #96]: https://github.com/tqdm/tqdm/pull/96 + + +Mozilla Public Licence (MPL) v. 2.0 - Exhibit A +----------------------------------------------- + +This Source Code Form is subject to the terms of the +Mozilla Public License, v. 2.0. +If a copy of the MPL was not distributed with this project, +You can obtain one at https://mozilla.org/MPL/2.0/. + + +MIT License (MIT) +----------------- + +Copyright (c) 2013 noamraph + +Permission is hereby granted, free of charge, to any person obtaining a copy of +this software and associated documentation files (the "Software"), to deal in +the Software without restriction, including without limitation the rights to +use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of +the Software, and to permit persons to whom the Software is furnished to do so, +subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS +FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR +COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER +IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN +CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/top_level.txt b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/top_level.txt new file mode 100644 index 0000000000000000000000000000000000000000..78620c472c9d799a14ccb02a0233f4669b3bcdcb --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm-4.68.3+computecanada.dist-info/top_level.txt @@ -0,0 +1 @@ +tqdm diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8081f77b8812f3b42d7949daa4195d2c35dc70ac --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/__init__.py @@ -0,0 +1,38 @@ +from ._monitor import TMonitor, TqdmSynchronisationWarning +from ._tqdm_pandas import tqdm_pandas +from .cli import main # TODO: remove in v5.0.0 +from .gui import tqdm as tqdm_gui # TODO: remove in v5.0.0 +from .gui import trange as tgrange # TODO: remove in v5.0.0 +from .std import ( + TqdmDeprecationWarning, TqdmExperimentalWarning, TqdmKeyError, TqdmMonitorWarning, + TqdmTypeError, TqdmWarning, tqdm, trange) +from .version import __version__ + +__all__ = ['tqdm', 'tqdm_gui', 'trange', 'tgrange', 'tqdm_pandas', + 'tqdm_notebook', 'tnrange', 'main', 'TMonitor', + 'TqdmTypeError', 'TqdmKeyError', + 'TqdmWarning', 'TqdmDeprecationWarning', + 'TqdmExperimentalWarning', + 'TqdmMonitorWarning', 'TqdmSynchronisationWarning', + '__version__'] + + +def tqdm_notebook(*args, **kwargs): # pragma: no cover + """See tqdm.notebook.tqdm for full documentation""" + from warnings import warn + + from .notebook import tqdm as _tqdm_notebook + warn("This function will be removed in tqdm==5.0.0\n" + "Please use `tqdm.notebook.tqdm` instead of `tqdm.tqdm_notebook`", + TqdmDeprecationWarning, stacklevel=2) + return _tqdm_notebook(*args, **kwargs) + + +def tnrange(*args, **kwargs): # pragma: no cover + """Shortcut for `tqdm.notebook.tqdm(range(*args), **kwargs)`.""" + from warnings import warn + + from .notebook import trange as _tnrange + warn("Please use `tqdm.notebook.trange` instead of `tqdm.tnrange`", + TqdmDeprecationWarning, stacklevel=2) + return _tnrange(*args, **kwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/__main__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..4e28416e104515e90fca4b69cc60d0c61fd15d61 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/__main__.py @@ -0,0 +1,3 @@ +from .cli import main + +main() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_main.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_main.py new file mode 100644 index 0000000000000000000000000000000000000000..04fdeeff17b5cc84b210f445b54b87d5b99e3748 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_main.py @@ -0,0 +1,9 @@ +from warnings import warn + +from .cli import * # NOQA +from .cli import __all__ # NOQA +from .std import TqdmDeprecationWarning + +warn("This function will be removed in tqdm==5.0.0\n" + "Please use `tqdm.cli.*` instead of `tqdm._main.*`", + TqdmDeprecationWarning, stacklevel=2) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_monitor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..9d97b8313770b7ced3639d43db3c89d41c5b9b67 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_monitor.py @@ -0,0 +1,104 @@ +import atexit +from threading import Event, Thread, current_thread +from time import time +from warnings import warn + +__all__ = ["TMonitor", "TqdmSynchronisationWarning"] + + +class TqdmSynchronisationWarning(RuntimeWarning): + """ + tqdm multi-thread/-process errors which may cause incorrect nesting + but otherwise no adverse effects + """ + + +class TMonitor(Thread): + """ + Monitoring thread for tqdm bars. + Monitors if tqdm bars are taking too much time to display + and readjusts miniters automatically if necessary. + + Parameters + ---------- + tqdm_cls : class + tqdm class to use (can be core tqdm or a submodule). + sleep_interval : float + Time to sleep between monitoring checks. + """ + _test = {} # internal vars for unit testing + + def __init__(self, tqdm_cls, sleep_interval): + Thread.__init__(self, name="tqdm_monitor") + self.daemon = True # kill thread when main killed (KeyboardInterrupt) + self.woken = 0 # last time woken up, to sync with monitor + self.tqdm_cls = tqdm_cls + self.sleep_interval = sleep_interval + self._time = self._test.get("time", time) + self.was_killed = self._test.get("Event", Event)() + atexit.register(self._atexit_signal) + self.start() + + def _atexit_signal(self): + """ + Non-joining shutdown signal. + Avoids deadlocks at interpreter exit from other threads, dead forks, etc. + This daemon thread is auto-reaped on shutdown without needing a join. + """ + self.was_killed.set() + + def exit(self): + self.was_killed.set() + if self is not current_thread(): + self.join() + return self.report() + + def get_instances(self): + # returns a copy of started `tqdm_cls` instances + return [i for i in self.tqdm_cls._instances.copy() + # Avoid race by checking that the instance started + if hasattr(i, 'start_t')] + + def run(self): + cur_t = self._time() + while True: + # After processing and before sleeping, notify that we woke + # Need to be done just before sleeping + self.woken = cur_t + # Sleep some time... + self.was_killed.wait(self.sleep_interval) + # Quit if killed + if self.was_killed.is_set(): + return + # Then monitor! + # Acquire lock (to access _instances) + with self.tqdm_cls.get_lock(): + cur_t = self._time() + # Check tqdm instances are waiting too long to print + instances = self.get_instances() + for instance in instances: + # Check event in loop to reduce blocking time on exit + if self.was_killed.is_set(): + return + # Only if mininterval > 1 (else iterations are just slow) + # and last refresh exceeded maxinterval + if ( + instance.miniters > 1 + and (cur_t - instance.last_print_t) >= instance.maxinterval + ): + # force bypassing miniters on next iteration + # (dynamic_miniters adjusts mininterval automatically) + instance.miniters = 1 + # Refresh now! (works only for manual tqdm) + instance.refresh(nolock=True) + # Remove accidental long-lived strong reference + del instance + if instances != self.get_instances(): # pragma: nocover + warn("Set changed size during iteration" + + " (see https://github.com/tqdm/tqdm/issues/481)", + TqdmSynchronisationWarning, stacklevel=2) + # Remove accidental long-lived strong references + del instances + + def report(self): + return not self.was_killed.is_set() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm.py new file mode 100644 index 0000000000000000000000000000000000000000..7fc4962774a4651db7a739a3f143633b6215a9bd --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm.py @@ -0,0 +1,9 @@ +from warnings import warn + +from .std import * # NOQA +from .std import __all__ # NOQA +from .std import TqdmDeprecationWarning + +warn("This function will be removed in tqdm==5.0.0\n" + "Please use `tqdm.std.*` instead of `tqdm._tqdm.*`", + TqdmDeprecationWarning, stacklevel=2) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm_gui.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm_gui.py new file mode 100644 index 0000000000000000000000000000000000000000..f32aa894f54b3a5b47a0fbf4263c2fd20df56c9d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm_gui.py @@ -0,0 +1,9 @@ +from warnings import warn + +from .gui import * # NOQA +from .gui import __all__ # NOQA +from .std import TqdmDeprecationWarning + +warn("This function will be removed in tqdm==5.0.0\n" + "Please use `tqdm.gui.*` instead of `tqdm._tqdm_gui.*`", + TqdmDeprecationWarning, stacklevel=2) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm_notebook.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm_notebook.py new file mode 100644 index 0000000000000000000000000000000000000000..f225fbf5b52d04987ccf68f4d5ee4b735e3158b0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm_notebook.py @@ -0,0 +1,9 @@ +from warnings import warn + +from .notebook import * # NOQA +from .notebook import __all__ # NOQA +from .std import TqdmDeprecationWarning + +warn("This function will be removed in tqdm==5.0.0\n" + "Please use `tqdm.notebook.*` instead of `tqdm._tqdm_notebook.*`", + TqdmDeprecationWarning, stacklevel=2) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm_pandas.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm_pandas.py new file mode 100644 index 0000000000000000000000000000000000000000..c4fe6efdc603579e7f8acfa27ac10dccdf3e94ce --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_tqdm_pandas.py @@ -0,0 +1,24 @@ +import sys + +__author__ = "github.com/casperdcl" +__all__ = ['tqdm_pandas'] + + +def tqdm_pandas(tclass, **tqdm_kwargs): + """ + Registers the given `tqdm` instance with + `pandas.core.groupby.DataFrameGroupBy.progress_apply`. + """ + from tqdm import TqdmDeprecationWarning + + if isinstance(tclass, type) or (getattr(tclass, '__name__', '').startswith( + 'tqdm_')): # delayed adapter case + TqdmDeprecationWarning( + "Please use `tqdm.pandas(...)` instead of `tqdm_pandas(tqdm, ...)`.", + fp_write=getattr(tqdm_kwargs.get('file', None), 'write', sys.stderr.write)) + tclass.pandas(**tqdm_kwargs) + else: + TqdmDeprecationWarning( + "Please use `tqdm.pandas(...)` instead of `tqdm_pandas(tqdm(...))`.", + fp_write=getattr(tclass.fp, 'write', sys.stderr.write)) + type(tclass).pandas(deprecated_t=tclass) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..28665dfe7f518bfaabce9bd442e4f2b23aeac267 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/_utils.py @@ -0,0 +1,11 @@ +from warnings import warn + +from .std import TqdmDeprecationWarning +from .utils import ( # noqa: F401, pylint: disable=unused-import + CUR_OS, IS_NIX, IS_WIN, RE_ANSI, Comparable, FormatReplace, SimpleTextIOWrapper, + _environ_cols_wrapper, _is_ascii, _is_utf, _screen_shape_linux, _screen_shape_tput, + _screen_shape_windows, _screen_shape_wrapper, _supports_unicode, _term_move_up, colorama) + +warn("This function will be removed in tqdm==5.0.0\n" + "Please use `tqdm.utils.*` instead of `tqdm._utils.*`", + TqdmDeprecationWarning, stacklevel=2) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/asyncio.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/asyncio.py new file mode 100644 index 0000000000000000000000000000000000000000..b32e7673b94d1aeb5b0ef3deec3e5d9e5454e3c8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/asyncio.py @@ -0,0 +1,93 @@ +""" +Asynchronous progress bar decorator for iterators. +Includes a default `range` iterator printing to `stderr`. + +Usage: +>>> from tqdm.asyncio import trange, tqdm +>>> async for i in trange(10): +... ... +""" +import asyncio +from sys import version_info + +from .std import tqdm as std_tqdm + +__author__ = {"github.com/": ["casperdcl"]} +__all__ = ['tqdm_asyncio', 'tarange', 'tqdm', 'trange'] + + +class tqdm_asyncio(std_tqdm): + """ + Asynchronous-friendly version of tqdm. + """ + def __init__(self, iterable=None, *args, **kwargs): + super().__init__(iterable, *args, **kwargs) + self.iterable_awaitable = False + if iterable is not None: + if hasattr(iterable, "__anext__"): + self.iterable_next = iterable.__anext__ + self.iterable_awaitable = True + elif hasattr(iterable, "__next__"): + self.iterable_next = iterable.__next__ + else: + self.iterable_iterator = iter(iterable) + self.iterable_next = self.iterable_iterator.__next__ + + def __aiter__(self): + return self + + async def __anext__(self): + try: + if self.iterable_awaitable: + res = await self.iterable_next() + else: + res = self.iterable_next() + self.update() + return res + except StopIteration: + self.close() + raise StopAsyncIteration + except BaseException: + self.close() + raise + + def send(self, *args, **kwargs): + return self.iterable.send(*args, **kwargs) + + @classmethod + def as_completed(cls, fs, *, loop=None, timeout=None, total=None, **tqdm_kwargs): + """ + Wrapper for `asyncio.as_completed`. + """ + if total is None: + total = len(fs) + kwargs = {} + if version_info[:2] < (3, 10): + kwargs['loop'] = loop + yield from cls(asyncio.as_completed(fs, timeout=timeout, **kwargs), + total=total, **tqdm_kwargs) + + @classmethod + async def gather(cls, *fs, loop=None, timeout=None, total=None, **tqdm_kwargs): + """ + Wrapper for `asyncio.gather`. + """ + async def wrap_awaitable(i, f): + return i, await f + + ifs = [wrap_awaitable(i, f) for i, f in enumerate(fs)] + res = [await f for f in cls.as_completed(ifs, loop=loop, timeout=timeout, + total=total, **tqdm_kwargs)] + return [i for _, i in sorted(res)] + + +def tarange(*args, **kwargs): + """ + A shortcut for `tqdm.asyncio.tqdm(range(*args), **kwargs)`. + """ + return tqdm_asyncio(range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_asyncio +trange = tarange diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/auto.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/auto.py new file mode 100644 index 0000000000000000000000000000000000000000..206c4409d5269594bdbab3a092ef6e09e7c01947 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/auto.py @@ -0,0 +1,40 @@ +""" +Enables multiple commonly used features. + +Method resolution order: + +- `tqdm.autonotebook` without import warnings +- `tqdm.asyncio` +- `tqdm.std` base class + +Usage: +>>> from tqdm.auto import trange, tqdm +>>> for i in trange(10): +... ... +""" +import warnings + +from .std import TqdmExperimentalWarning + +with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=TqdmExperimentalWarning) + from .autonotebook import tqdm as notebook_tqdm + +from .asyncio import tqdm as asyncio_tqdm +from .std import tqdm as std_tqdm + +if notebook_tqdm != std_tqdm: + class tqdm(notebook_tqdm, asyncio_tqdm): # pylint: disable=inconsistent-mro + pass +else: + tqdm = asyncio_tqdm + + +def trange(*args, **kwargs): + """ + A shortcut for `tqdm.auto.tqdm(range(*args), **kwargs)`. + """ + return tqdm(range(*args), **kwargs) + + +__all__ = ["tqdm", "trange"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/autonotebook.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/autonotebook.py new file mode 100644 index 0000000000000000000000000000000000000000..df1d21b2b0f933b460c0f1021ae22a1f058b0871 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/autonotebook.py @@ -0,0 +1,38 @@ +""" +Automatically choose between `tqdm.notebook` and `tqdm.std`. + +Usage: +>>> from tqdm.autonotebook import trange, tqdm +>>> for i in trange(10): +... ... +""" +import os +import sys +from warnings import warn + +try: + if 'ipykernel.zmqshell' in sys.modules: + if any(i == 'QT_API' or i.startswith('SPYDER') for i in os.environ): + raise ImportError("console") # jupyter-qtconsole/spyder + ipy = sys.modules['IPython'].get_ipython().__class__.__name__.lower() + if 'qt' in ipy or 'spyder' in ipy: + raise ImportError("console") # older jupyter-qtconsole/spyder + # jupyter-notebook/jupyterlab/vscode/binder/colab + elif 'IPython.utils._process_emscripten' in sys.modules: + pass # jupyterlite (pyodide)/jupyterlite-xeus + else: + raise ImportError("console") # ipython/jupyter-console + from .notebook import WARN_NOIPYW, IProgress + if IProgress is None: + from .std import TqdmWarning + warn(WARN_NOIPYW, TqdmWarning, stacklevel=2) + raise ImportError('ipywidgets') +except Exception: + from .std import tqdm, trange +else: # pragma: no cover + from .notebook import tqdm, trange + from .std import TqdmExperimentalWarning + warn("Using `tqdm.autonotebook.tqdm` in notebook mode." + " Use `tqdm.tqdm` instead to force console mode" + " (e.g. in jupyter console)", TqdmExperimentalWarning, stacklevel=2) +__all__ = ["tqdm", "trange"] diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/cli.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/cli.py new file mode 100644 index 0000000000000000000000000000000000000000..e5599efdc7c1c2aaad1cd95dd02eacdffc71c658 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/cli.py @@ -0,0 +1,324 @@ +""" +Module version for monitoring CLI pipes (`... | python -m tqdm | ...`). +""" +import logging +import re +import sys +from ast import literal_eval +from textwrap import indent + +from .std import TqdmKeyError, TqdmTypeError, tqdm +from .version import __version__ + +__all__ = ["main"] +log = logging.getLogger(__name__) + + +def cast(val, typ): + log.debug((val, typ)) + if " or " in typ: + for t in typ.split(" or "): + try: + return cast(val, t) + except TqdmTypeError: + pass + raise TqdmTypeError(f"{val} : {typ}") + + # sys.stderr.write('\ndebug | `val:type`: `' + val + ':' + typ + '`.\n') + if typ == 'bool': + if (val == 'True') or (val == ''): + return True + if val == 'False': + return False + raise TqdmTypeError(val + ' : ' + typ) + if typ == 'chr': + if len(val) == 1: + return val.encode() + if re.match(r"^\\\w+$", val): + return literal_eval(f'"{val}"').encode() + raise TqdmTypeError(f"{val} : {typ}") + if typ == 'str': + return val + if typ == 'int': + try: + return int(val) + except ValueError as exc: + raise TqdmTypeError(f"{val} : {typ}") from exc + if typ == 'float': + try: + return float(val) + except ValueError as exc: + raise TqdmTypeError(f"{val} : {typ}") from exc + raise TqdmTypeError(f"{val} : {typ}") + + +def posix_pipe(fin, fout, delim=b'\\n', buf_size=256, + callback=lambda float: None, callback_len=True): + """ + Params + ------ + fin : binary file with `read(buf_size : int)` method + fout : binary file with `write` (and optionally `flush`) methods. + callback : function(float), e.g.: `tqdm.update` + callback_len : If (default: True) do `callback(len(buffer))`. + Otherwise, do `callback(data) for data in buffer.split(delim)`. + """ + fp_write = fout.write + + if not delim: + while True: + tmp = fin.read(buf_size) + + # flush at EOF + if not tmp: + getattr(fout, 'flush', lambda: None)() + return + + fp_write(tmp) + callback(len(tmp)) + # return + + buf = b'' + len_delim = len(delim) + # n = 0 + while True: + tmp = fin.read(buf_size) + + # flush at EOF + if not tmp: + if buf: + fp_write(buf) + if callback_len: + # n += 1 + buf.count(delim) + callback(1 + buf.count(delim)) + else: + for i in buf.split(delim): + callback(i) + getattr(fout, 'flush', lambda: None)() + return # n + + while True: + i = tmp.find(delim) + if i < 0: + buf += tmp + break + fp_write(buf + tmp[:i + len(delim)]) + # n += 1 + callback(1 if callback_len else (buf + tmp[:i])) + buf = b'' + tmp = tmp[i + len_delim:] + + +# ((opt, type), ... ) +RE_OPTS = re.compile(r'\n {4}(\S+)\s{2,}:\s*([^,]+)') +# better split method assuming no positional args +RE_SHLEX = re.compile(r'\s*(? : \2', d) + split = RE_OPTS.split(d) + opt_types_desc = zip(split[1::3], split[2::3], split[3::3]) + d = ''.join(('\n --{0} : {2}{3}' if otd[1] == 'bool' else + '\n --{0}=<{1}> : {2}{3}').format( + otd[0].replace('_', '-'), otd[0], *otd[1:]) + for otd in opt_types_desc if otd[0] not in UNSUPPORTED_OPTS) + + help_short = "Usage:\n tqdm [--help | options]\n" + d = help_short + """ +Options: + -h, --help Print this help and exit. + -v, --version Print version and exit. +""" + d.strip('\n') + '\n' + + # opts = docopt(d, version=__version__) + if any(v in argv for v in ('-v', '--version')): + sys.stdout.write(__version__ + '\n') + sys.exit(0) + elif any(v in argv for v in ('-h', '--help')): + sys.stdout.write(d + '\n') + sys.exit(0) + elif argv and argv[0][:2] != '--': + sys.stderr.write(f"Error:Unknown argument:{argv[0]}\n{help_short}") + + argv = RE_SHLEX.split(' '.join(["tqdm"] + argv)) + opts = dict(zip(argv[1::3], argv[3::3])) + + log.debug(opts) + opts.pop('log', True) + + tqdm_args = {'file': fp} + try: + for (o, v) in opts.items(): + o = o.replace('-', '_') + try: + tqdm_args[o] = cast(v, opt_types[o]) + except KeyError as e: + raise TqdmKeyError(str(e)) + log.debug('args:' + str(tqdm_args)) + + delim_per_char = tqdm_args.pop('bytes', False) + update = tqdm_args.pop('update', False) + update_to = tqdm_args.pop('update_to', False) + if sum((delim_per_char, update, update_to)) > 1: + raise TqdmKeyError("Can only have one of --bytes --update --update_to") + except Exception: + fp.write("\nError:\n" + help_short) + stdin, stdout_write = sys.stdin, sys.stdout.write + for i in stdin: + stdout_write(i) + raise + else: + buf_size = tqdm_args.pop('buf_size', 256) + delim = tqdm_args.pop('delim', b'\\n') + tee = tqdm_args.pop('tee', False) + manpath = tqdm_args.pop('manpath', None) + comppath = tqdm_args.pop('comppath', None) + if tqdm_args.pop('null', False): + class stdout: + @staticmethod + def write(_): + pass + else: + stdout = sys.stdout + stdout = getattr(stdout, 'buffer', stdout) + stdin = getattr(sys.stdin, 'buffer', sys.stdin) + if manpath or comppath: + try: # py<3.9 + import importlib_resources as resources + except ImportError: + from importlib import resources + from pathlib import Path + + def cp(name, dst): + """copy resource `name` to `dst`""" + fi = resources.files('tqdm') / name + dst.write_bytes(fi.read_bytes()) + log.info("written:%s", dst) + if manpath is not None: + cp('tqdm.1', Path(manpath) / 'tqdm.1') + if comppath is not None: + cp('completion.sh', Path(comppath) / 'tqdm_completion.sh') + sys.exit(0) + if tee: + stdout_write = stdout.write + fp_write = getattr(fp, 'buffer', fp).write + + class stdout: # pylint: disable=function-redefined + @staticmethod + def write(x): + with tqdm.external_write_mode(file=fp): + fp_write(x) + stdout_write(x) + if delim_per_char: + tqdm_args.setdefault('unit', 'B') + tqdm_args.setdefault('unit_scale', True) + tqdm_args.setdefault('unit_divisor', 1024) + log.debug(tqdm_args) + with tqdm(**tqdm_args) as t: + posix_pipe(stdin, stdout, '', buf_size, t.update) + elif delim == b'\\n': + log.debug(tqdm_args) + write = stdout.write + if update or update_to: + with tqdm(**tqdm_args) as t: + if update: + def callback(i): + t.update(literal_eval(i.decode())) + else: # update_to + def callback(i): + t.update(literal_eval(i.decode()) - t.n) + for i in stdin: + write(i) + callback(i) + else: + for i in tqdm(stdin, **tqdm_args): + write(i) + else: + log.debug(tqdm_args) + with tqdm(**tqdm_args) as t: + callback_len = False + if update: + def callback(i): + t.update(literal_eval(i.decode())) + elif update_to: + def callback(i): + t.update(literal_eval(i.decode()) - t.n) + else: + callback = t.update + callback_len = True + posix_pipe(stdin, stdout, delim, buf_size, callback, callback_len) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/completion.sh b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/completion.sh new file mode 100644 index 0000000000000000000000000000000000000000..9f61c7f14bb8c1f6099b9eb75dce28ece6a7ae96 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/completion.sh @@ -0,0 +1,19 @@ +#!/usr/bin/env bash +_tqdm(){ + local cur prv + cur="${COMP_WORDS[COMP_CWORD]}" + prv="${COMP_WORDS[COMP_CWORD - 1]}" + + case ${prv} in + --bar_format|--buf_size|--colour|--comppath|--delay|--delim|--desc|--initial|--lock_args|--manpath|--maxinterval|--mininterval|--miniters|--ncols|--nrows|--position|--postfix|--smoothing|--total|--unit|--unit_divisor) + # await user input + ;; + "--log") + COMPREPLY=($(compgen -W 'CRITICAL FATAL ERROR WARN WARNING INFO DEBUG NOTSET' -- ${cur})) + ;; + *) + COMPREPLY=($(compgen -W '--ascii --bar_format --buf_size --bytes --colour --comppath --delay --delim --desc --disable --dynamic_ncols --help --initial --leave --lock_args --log --manpath --maxinterval --mininterval --miniters --ncols --nrows --null --position --postfix --smoothing --tee --total --unit --unit_divisor --unit_scale --update --update_to --version --write_bytes -h -v' -- ${cur})) + ;; + esac +} +complete -F _tqdm tqdm diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..87ec4601fcff138e0882808237375aed93f5ef35 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/__init__.py @@ -0,0 +1,91 @@ +""" +Thin wrappers around common functions. + +Subpackages contain potentially unstable extensions. +""" +from warnings import warn + +from ..auto import tqdm as tqdm_auto +from ..std import TqdmDeprecationWarning, tqdm +from ..utils import ObjectWrapper + +__author__ = {"github.com/": ["casperdcl"]} +__all__ = ['tenumerate', 'tzip', 'tmap'] + + +class DummyTqdmFile(ObjectWrapper): + """Dummy file-like that will write to tqdm""" + + def __init__(self, wrapped): + super().__init__(wrapped) + self._buf = [] + + def write(self, x, nolock=False): + nl = b"\n" if isinstance(x, bytes) else "\n" + pre, sep, post = x.rpartition(nl) + if sep: + blank = type(nl)() + tqdm.write(blank.join(self._buf + [pre, sep]), + end=blank, file=self._wrapped, nolock=nolock) + self._buf = [post] + else: + self._buf.append(x) + + def __del__(self): + if self._buf: + blank = type(self._buf[0])() + try: + tqdm.write(blank.join(self._buf), end=blank, file=self._wrapped) + except (OSError, ValueError): + pass + + +def builtin_iterable(func): + """Returns `func`""" + warn("This function has no effect, and will be removed in tqdm==5.0.0", + TqdmDeprecationWarning, stacklevel=2) + return func + + +def tenumerate(iterable, start=0, total=None, tqdm_class=tqdm_auto, **tqdm_kwargs): + """ + Equivalent of `numpy.ndenumerate` or builtin `enumerate`. + + Parameters + ---------- + tqdm_class : [default: tqdm.auto.tqdm]. + """ + try: + import numpy as np + except ImportError: + pass + else: + if isinstance(iterable, np.ndarray): + return tqdm_class(np.ndenumerate(iterable), total=total or iterable.size, + **tqdm_kwargs) + return enumerate(tqdm_class(iterable, total=total, **tqdm_kwargs), start) + + +def tzip(iter1, *iter2plus, **tqdm_kwargs): + """ + Equivalent of builtin `zip`. + + Parameters + ---------- + tqdm_class : [default: tqdm.auto.tqdm]. + """ + kwargs = tqdm_kwargs.copy() + tqdm_class = kwargs.pop("tqdm_class", tqdm_auto) + yield from zip(tqdm_class(iter1, **kwargs), *iter2plus) + + +def tmap(function, *sequences, **tqdm_kwargs): + """ + Equivalent of builtin `map`. + + Parameters + ---------- + tqdm_class : [default: tqdm.auto.tqdm]. + """ + for i in tzip(*sequences, **tqdm_kwargs): + yield function(*i) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/bells.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/bells.py new file mode 100644 index 0000000000000000000000000000000000000000..e4968e7751920093aa6c3ddd72c5b5e2e29df8bc --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/bells.py @@ -0,0 +1,28 @@ +""" +Even more features than `tqdm.auto` (all the bells & whistles): + +- `tqdm.auto` +- `tqdm.tqdm.pandas` +- `tqdm.contrib.slack` + + uses `${TQDM_SLACK_TOKEN}` and `${TQDM_SLACK_CHANNEL}` +- `tqdm.contrib.telegram` + + uses `${TQDM_TELEGRAM_TOKEN}` and `${TQDM_TELEGRAM_CHAT_ID}` +- `tqdm.contrib.discord` + + uses `${TQDM_DISCORD_TOKEN}` and `${TQDM_DISCORD_CHANNEL_ID}` +""" +__all__ = ['tqdm', 'trange'] +import warnings +from os import getenv + +if getenv("TQDM_SLACK_TOKEN") and getenv("TQDM_SLACK_CHANNEL"): + from .slack import tqdm, trange +elif getenv("TQDM_TELEGRAM_TOKEN") and getenv("TQDM_TELEGRAM_CHAT_ID"): + from .telegram import tqdm, trange +elif getenv("TQDM_DISCORD_TOKEN") and getenv("TQDM_DISCORD_CHANNEL_ID"): + from .discord import tqdm, trange +else: + from ..auto import tqdm, trange + +with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=FutureWarning) + tqdm.pandas() diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/concurrent.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/concurrent.py new file mode 100644 index 0000000000000000000000000000000000000000..cd81d622a1309df179042159a56cef4f8c309224 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/concurrent.py @@ -0,0 +1,105 @@ +""" +Thin wrappers around `concurrent.futures`. +""" +from contextlib import contextmanager +from operator import length_hint +from os import cpu_count + +from ..auto import tqdm as tqdm_auto +from ..std import TqdmWarning + +__author__ = {"github.com/": ["casperdcl"]} +__all__ = ['thread_map', 'process_map'] + + +@contextmanager +def ensure_lock(tqdm_class, lock_name=""): + """get (create if necessary) and then restore `tqdm_class`'s lock""" + old_lock = getattr(tqdm_class, '_lock', None) # don't create a new lock + lock = old_lock or tqdm_class.get_lock() # maybe create a new lock + lock = getattr(lock, lock_name, lock) # maybe subtype + tqdm_class.set_lock(lock) + yield lock + if old_lock is None: + del tqdm_class._lock + else: + tqdm_class.set_lock(old_lock) + + +def _executor_map(PoolExecutor, fn, *iterables, **tqdm_kwargs): + """ + Implementation of `thread_map` and `process_map`. + + Parameters + ---------- + tqdm_class : [default: tqdm.auto.tqdm]. + max_workers : [default: min(32, cpu_count() + 4)]. + chunksize : [default: 1]. + lock_name : [default: "":str]. + """ + kwargs = tqdm_kwargs.copy() + if "total" not in kwargs: + kwargs["total"] = length_hint(iterables[0]) + tqdm_class = kwargs.pop("tqdm_class", tqdm_auto) + max_workers = kwargs.pop("max_workers", min(32, cpu_count() + 4)) + chunksize = kwargs.pop("chunksize", 1) + lock_name = kwargs.pop("lock_name", "") + with ensure_lock(tqdm_class, lock_name=lock_name) as lk: + # share lock in case workers are already using `tqdm` + with PoolExecutor(max_workers=max_workers, initializer=tqdm_class.set_lock, + initargs=(lk,)) as ex: + return list(tqdm_class(ex.map(fn, *iterables, chunksize=chunksize), **kwargs)) + + +def thread_map(fn, *iterables, **tqdm_kwargs): + """ + Equivalent of `list(map(fn, *iterables))` + driven by `concurrent.futures.ThreadPoolExecutor`. + + Parameters + ---------- + tqdm_class : optional + `tqdm` class to use for bars [default: tqdm.auto.tqdm]. + max_workers : int, optional + Maximum number of workers to spawn; passed to + `concurrent.futures.ThreadPoolExecutor.__init__`. + [default: max(32, cpu_count() + 4)]. + """ + from concurrent.futures import ThreadPoolExecutor + return _executor_map(ThreadPoolExecutor, fn, *iterables, **tqdm_kwargs) + + +def process_map(fn, *iterables, **tqdm_kwargs): + """ + Equivalent of `list(map(fn, *iterables))` + driven by `concurrent.futures.ProcessPoolExecutor`. + + Parameters + ---------- + tqdm_class : optional + `tqdm` class to use for bars [default: tqdm.auto.tqdm]. + max_workers : int, optional + Maximum number of workers to spawn; passed to + `concurrent.futures.ProcessPoolExecutor.__init__`. + [default: min(32, cpu_count() + 4)]. + chunksize : int, optional + Size of chunks sent to worker processes; passed to + `concurrent.futures.ProcessPoolExecutor.map`. [default: 1]. + lock_name : str, optional + Member of `tqdm_class.get_lock()` to use [default: mp_lock]. + """ + from concurrent.futures import ProcessPoolExecutor + if iterables and "chunksize" not in tqdm_kwargs: + # default `chunksize=1` has poor performance for large iterables + # (most time spent dispatching items to workers). + longest_iterable_len = max(map(length_hint, iterables)) + if longest_iterable_len > 1000: + from warnings import warn + warn("Iterable length %d > 1000 but `chunksize` is not set." + " This may seriously degrade multiprocess performance." + " Set `chunksize=1` or more." % longest_iterable_len, + TqdmWarning, stacklevel=2) + if "lock_name" not in tqdm_kwargs: + tqdm_kwargs = tqdm_kwargs.copy() + tqdm_kwargs["lock_name"] = "mp_lock" + return _executor_map(ProcessPoolExecutor, fn, *iterables, **tqdm_kwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/discord.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/discord.py new file mode 100644 index 0000000000000000000000000000000000000000..f46a4992bcbaeb7f85d75f8e6c64dc6e09f5325b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/discord.py @@ -0,0 +1,155 @@ +""" +Sends updates to a Discord bot. + +Usage: +>>> from tqdm.contrib.discord import tqdm, trange +>>> for i in trange(10, token='{token}', channel_id='{channel_id}'): +... ... + +![screenshot](https://tqdm.github.io/img/screenshot-discord.png) +""" +from warnings import warn + +from requests import Session +from requests.utils import default_user_agent + +from ..auto import tqdm as tqdm_auto +from ..std import TqdmWarning +from ..utils import envwrap +from ..version import __version__ +from .utils_worker import MonoWorker + +__author__ = {"github.com/": ["casperdcl", "guigoruiz1"]} +__all__ = ['DiscordIO', 'tqdm_discord', 'tdrange', 'tqdm', 'trange'] + + +class DiscordIO(MonoWorker): + """Non-blocking file-like IO using a Discord Bot.""" + API = 'https://discord.com/api/v10' + UA = f"tqdm (https://tqdm.github.io, {__version__}) {default_user_agent()}" + + def __init__(self, token, channel_id): + """Creates a new message in the given `channel_id`.""" + super().__init__() + self.token = token + self.channel_id = channel_id + self.session = Session() + self.text = self.__class__.__name__ + self.message_id # pylint: disable=pointless-statement + + @property + def message_id(self): + if hasattr(self, '_message_id'): + return self._message_id # pylint: disable=access-member-before-definition + try: + req = self.session.post( + f'{self.API}/channels/{self.channel_id}/messages', + headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA}, + json={'content': f"`{self.text}`"}) + res = req.json() + req.raise_for_status() + except Exception as e: + if req.status_code == 429: + warn("Creation rate limit: try increasing `mininterval`.", + TqdmWarning, stacklevel=2) + else: + tqdm_auto.write(str(e)) + else: + self._message_id = res['id'] + return self._message_id + + def write(self, s): + """Replaces internal `message_id`'s text with `s`.""" + if not s: + s = "..." + s = s.replace('\r', '').strip() + if s == self.text: + return # avoid duplicate message Bot error + message_id = self.message_id + if message_id is None: + return + self.text = s + try: + future = self.submit( + self.session.patch, + f'{self.API}/channels/{self.channel_id}/messages/{message_id}', + headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA}, + json={'content': f"`{self.text}`"}) + except Exception as e: + tqdm_auto.write(str(e)) + else: + return future + + def delete(self): + """Deletes internal `message_id`.""" + try: + future = self.submit( + self.session.delete, + f'{self.API}/channels/{self.channel_id}/messages/{self.message_id}', + headers={'Authorization': f'Bot {self.token}', 'User-Agent': self.UA}) + except Exception as e: + tqdm_auto.write(str(e)) + else: + return future + + +class tqdm_discord(tqdm_auto): # pylint: disable=inconsistent-mro + """ + Standard `tqdm.auto.tqdm` but also sends updates to a Discord Bot. + May take a few seconds to create (`__init__`). + + - create a discord bot (not public, no requirement of OAuth2 code + grant, only send message permissions) & invite it to a channel: + + - copy the bot `{token}` & `{channel_id}` and paste below + + >>> from tqdm.contrib.discord import tqdm, trange + >>> for i in tqdm(iterable, token='{token}', channel_id='{channel_id}'): + ... ... + """ + @envwrap("tqdm", "discord", is_method=True) + def __init__(self, *args, token=None, channel_id=None, **kwargs): + """ + Parameters + ---------- + token : str, required. Discord bot token + [default: ${TQDM_DISCORD_TOKEN}]. + channel_id : int, required. Discord channel ID + [default: ${TQDM_DISCORD_CHANNEL_ID}]. + + See `tqdm.auto.tqdm.__init__` for other parameters. + """ + if not kwargs.get('disable'): + kwargs = kwargs.copy() + self.dio = DiscordIO(token, channel_id) + super().__init__(*args, **kwargs) + + def display(self, **kwargs): # pylint: disable=arguments-differ + super().display(**kwargs) + fmt = self.format_dict + if fmt.get('bar_format', None): + fmt['bar_format'] = fmt['bar_format'].replace( + '', '{bar:10u}').replace('{bar}', '{bar:10u}') + self.dio.write(self.format_meter(**fmt)) + + def clear(self, *args, **kwargs): + super().clear(*args, **kwargs) + if not self.disable: + self.dio.write("") + + def close(self): + if self.disable: + return + super().close() + if not (self.leave or (self.leave is None and self.pos == 0)): + self.dio.delete() + + +def tdrange(*args, **kwargs): + """Shortcut for `tqdm.contrib.discord.tqdm(range(*args), **kwargs)`.""" + return tqdm_discord(range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_discord +trange = tdrange diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/itertools.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/itertools.py new file mode 100644 index 0000000000000000000000000000000000000000..469961867cbf2c22cb2b799ecb4736438a912520 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/itertools.py @@ -0,0 +1,92 @@ +""" +Thin wrappers around `itertools`. +""" +import itertools +import math + +from ..auto import tqdm as tqdm_auto + +__author__ = {"github.com/": ["casperdcl"]} +__all__ = [ + 'chain', 'product', 'permutations', 'combinations', 'combinations_with_replacement', 'batched'] + + +def chain(*iterables, total=None, tqdm_class=tqdm_auto, **kwargs): + """Equivalent of `itertools.chain`.""" + if total is None: + try: + total = sum(map(len, iterables)) + except (TypeError, AttributeError): + pass + return tqdm_class(itertools.chain(*iterables), total=total, **kwargs) + + +def product(*iterables, repeat=1, total=None, tqdm_class=tqdm_auto, **kwargs): + """Equivalent of `itertools.product`.""" + if total is None: + try: + lens = list(map(len, iterables)) + except (TypeError, AttributeError): + pass + else: + total = math.prod(lens) ** repeat + yield from tqdm_class(itertools.product(*iterables, repeat=repeat), total=total, **kwargs) + + +def permutations(iterable, r=None, total=None, tqdm_class=tqdm_auto, **kwargs): + """Equivalent of `itertools.permutations`.""" + if total is None: + try: + n = len(iterable) + except (TypeError, AttributeError): + pass + else: + r = n if r is None else r + if r > n: + total = 0 + else: + total = math.perm(n, r) + return tqdm_class(itertools.permutations(iterable, r), total=total, **kwargs) + + +def combinations(iterable, r, total=None, tqdm_class=tqdm_auto, **kwargs): + """Equivalent of `itertools.combinations`.""" + if total is None: + try: + n = len(iterable) + except (TypeError, AttributeError): + pass + else: + if r > n: + total = 0 + else: + total = math.comb(n, r) + return tqdm_class(itertools.combinations(iterable, r), total=total, **kwargs) + + +def combinations_with_replacement(iterable, r, total=None, tqdm_class=tqdm_auto, **kwargs): + """Equivalent of `itertools.combinations_with_replacement`.""" + if total is None: + try: + n = len(iterable) + except (TypeError, AttributeError): + pass + else: + total = 1 + for i in range(n+r-1, n-1, -1): + total *= i + for i in range(1, r+1): + total //= i + return tqdm_class(itertools.combinations_with_replacement(iterable, r), total=total, **kwargs) + + +def batched(iterable, n, total=None, tqdm_class=tqdm_auto, **kwargs): + """Equivalent of `itertools.batched`.""" + if total is None: + try: + total = len(iterable) + except (TypeError, AttributeError): + pass + return tqdm_class(itertools.batched(iterable, n), unit_scale=n, + total=(total+n-1) // n if total is not None else None, + **kwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/logging.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/logging.py new file mode 100644 index 0000000000000000000000000000000000000000..16edd45ab1b98f5443b983b47ee3257bf265ee59 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/logging.py @@ -0,0 +1,126 @@ +""" +Helper functionality for interoperability with stdlib `logging`. +""" +import logging +import sys +from contextlib import contextmanager + +try: + from typing import Iterator, List, Optional, Type # noqa: F401, pylint: disable=unused-import +except ImportError: + pass + +from ..std import tqdm as std_tqdm + + +class _TqdmLoggingHandler(logging.StreamHandler): + def __init__( + self, + tqdm_class=std_tqdm # type: Type[std_tqdm] + ): + super().__init__() + self.tqdm_class = tqdm_class + + def emit(self, record): + try: + msg = self.format(record) + self.tqdm_class.write(msg, file=self.stream) + self.flush() + except (KeyboardInterrupt, SystemExit): + raise + except: # noqa pylint: disable=bare-except + self.handleError(record) + + +def _is_console_logging_handler(handler): + return (isinstance(handler, logging.StreamHandler) + and handler.stream in {sys.stdout, sys.stderr}) + + +def _get_first_found_console_logging_handler(handlers): + for handler in handlers: + if _is_console_logging_handler(handler): + return handler + + +@contextmanager +def logging_redirect_tqdm( + loggers=None, # type: Optional[List[logging.Logger]], + tqdm_class=std_tqdm # type: Type[std_tqdm] +): + # type: (...) -> Iterator[None] + """ + Context manager redirecting console logging to `tqdm.write()`, leaving + other logging handlers (e.g. log files) unaffected. + + Parameters + ---------- + loggers : list, optional + Which handlers to redirect (default: [logging.root]). + tqdm_class : optional + + Example + ------- + ```python + import logging + from tqdm import trange + from tqdm.contrib.logging import logging_redirect_tqdm + + LOG = logging.getLogger(__name__) + + if __name__ == '__main__': + logging.basicConfig(level=logging.INFO) + with logging_redirect_tqdm(): + for i in trange(9): + if i == 4: + LOG.info("console logging redirected to `tqdm.write()`") + # logging restored + ``` + """ + if loggers is None: + loggers = [logging.root] + original_handlers_list = [logger.handlers for logger in loggers] + try: + for logger in loggers: + tqdm_handler = _TqdmLoggingHandler(tqdm_class) + orig_handler = _get_first_found_console_logging_handler(logger.handlers) + if orig_handler is not None: + tqdm_handler.setFormatter(orig_handler.formatter) + tqdm_handler.stream = orig_handler.stream + logger.handlers = [ + handler for handler in logger.handlers + if not _is_console_logging_handler(handler)] + [tqdm_handler] + yield + finally: + for logger, original_handlers in zip(loggers, original_handlers_list): + logger.handlers = original_handlers + + +@contextmanager +def tqdm_logging_redirect( + *args, + # loggers=None, # type: Optional[List[logging.Logger]] + # tqdm=None, # type: Optional[Type[tqdm.tqdm]] + **kwargs +): + # type: (...) -> Iterator[None] + """ + Convenience shortcut for: + ```python + with tqdm_class(*args, **tqdm_kwargs) as pbar: + with logging_redirect_tqdm(loggers=loggers, tqdm_class=tqdm_class): + yield pbar + ``` + + Parameters + ---------- + tqdm_class : optional, (default: tqdm.std.tqdm). + loggers : optional, list. + **tqdm_kwargs : passed to `tqdm_class`. + """ + tqdm_kwargs = kwargs.copy() + loggers = tqdm_kwargs.pop('loggers', None) + tqdm_class = tqdm_kwargs.pop('tqdm_class', std_tqdm) + with tqdm_class(*args, **tqdm_kwargs) as pbar: + with logging_redirect_tqdm(loggers=loggers, tqdm_class=tqdm_class): + yield pbar diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/slack.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/slack.py new file mode 100644 index 0000000000000000000000000000000000000000..e9808f55ad0b074bfb66ae79df35de8bf629233c --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/slack.py @@ -0,0 +1,119 @@ +""" +Sends updates to a Slack app. + +Usage: +>>> from tqdm.contrib.slack import tqdm, trange +>>> for i in trange(10, token='{token}', channel='{channel}'): +... ... + +![screenshot](https://tqdm.github.io/img/screenshot-slack.png) +""" +import logging + +try: + from slack_sdk import WebClient +except ImportError: + raise ImportError("Please `pip install slack-sdk`") + +from ..auto import tqdm as tqdm_auto +from ..utils import envwrap +from .utils_worker import MonoWorker + +__author__ = {"github.com/": ["0x2b3bfa0", "casperdcl"]} +__all__ = ['SlackIO', 'tqdm_slack', 'tsrange', 'tqdm', 'trange'] + + +class SlackIO(MonoWorker): + """Non-blocking file-like IO using a Slack app.""" + def __init__(self, token, channel): + """Creates a new message in the given `channel`.""" + super().__init__() + self.client = WebClient(token=token) + self.text = self.__class__.__name__ + try: + self.message = self.client.chat_postMessage(channel=channel, text=self.text) + except Exception as e: + tqdm_auto.write(str(e)) + self.message = None + + def write(self, s): + """Replaces internal `message`'s text with `s`.""" + if not s: + s = "..." + s = s.replace('\r', '').strip() + if s == self.text: + return # skip duplicate message + message = self.message + if message is None: + return + self.text = s + try: + future = self.submit(self.client.chat_update, channel=message['channel'], + ts=message['ts'], text='`' + s + '`') + except Exception as e: + tqdm_auto.write(str(e)) + else: + return future + + +class tqdm_slack(tqdm_auto): # pylint: disable=inconsistent-mro + """ + Standard `tqdm.auto.tqdm` but also sends updates to a Slack app. + May take a few seconds to create (`__init__`). + + - create a Slack app with the `chat:write` scope & invite it to a + channel: + - copy the bot `{token}` & `{channel}` and paste below + >>> from tqdm.contrib.slack import tqdm, trange + >>> for i in tqdm(iterable, token='{token}', channel='{channel}'): + ... ... + """ + @envwrap("tqdm", "slack", is_method=True) + def __init__(self, *args, token=None, channel=None, **kwargs): + """ + Parameters + ---------- + token : str, required. Slack token + [default: ${TQDM_SLACK_TOKEN}]. + channel : int, required. Slack channel + [default: ${TQDM_SLACK_CHANNEL}]. + mininterval : float, optional. + Minimum of [default: 1.5] to avoid rate limit. + + See `tqdm.auto.tqdm.__init__` for other parameters. + """ + if not kwargs.get('disable'): + kwargs = kwargs.copy() + logging.getLogger("HTTPClient").setLevel(logging.WARNING) + self.sio = SlackIO(token, channel) + kwargs['mininterval'] = max(1.5, kwargs.get('mininterval', 1.5)) + super().__init__(*args, **kwargs) + + def display(self, **kwargs): # pylint: disable=arguments-differ + super().display(**kwargs) + fmt = self.format_dict + if fmt.get('bar_format', None): + fmt['bar_format'] = fmt['bar_format'].replace( + '', '`{bar:10}`').replace('{bar}', '`{bar:10u}`') + elif self.total: + fmt['bar_format'] = '{l_bar}`{bar:10}`{r_bar}' + if fmt['ascii'] is False: + fmt['ascii'] = [":black_square:", ":small_blue_diamond:", ":large_blue_diamond:", + ":large_blue_square:"] + fmt['ncols'] = 336 + self.sio.write(self.format_meter(**fmt)) + + def clear(self, *args, **kwargs): + super().clear(*args, **kwargs) + if not self.disable: + self.sio.write("") + + +def tsrange(*args, **kwargs): + """Shortcut for `tqdm.contrib.slack.tqdm(range(*args), **kwargs)`.""" + return tqdm_slack(range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_slack +trange = tsrange diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/telegram.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/telegram.py new file mode 100644 index 0000000000000000000000000000000000000000..c019248e322214aeccbfaf2d8dac97a69d11470f --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/telegram.py @@ -0,0 +1,152 @@ +""" +Sends updates to a Telegram bot. + +Usage: +>>> from tqdm.contrib.telegram import tqdm, trange +>>> for i in trange(10, token='{token}', chat_id='{chat_id}'): +... ... + +![screenshot](https://tqdm.github.io/img/screenshot-telegram.gif) +""" +from warnings import warn + +from requests import Session + +from ..auto import tqdm as tqdm_auto +from ..std import TqdmWarning +from ..utils import envwrap +from .utils_worker import MonoWorker + +__author__ = {"github.com/": ["casperdcl"]} +__all__ = ['TelegramIO', 'tqdm_telegram', 'ttgrange', 'tqdm', 'trange'] + + +class TelegramIO(MonoWorker): + """Non-blocking file-like IO using a Telegram Bot.""" + API = 'https://api.telegram.org/bot' + + def __init__(self, token, chat_id): + """Creates a new message in the given `chat_id`.""" + super().__init__() + self.token = token + self.chat_id = chat_id + self.session = Session() + self.text = self.__class__.__name__ + self.message_id # pylint: disable=pointless-statement + + @property + def message_id(self): + if hasattr(self, '_message_id'): + return self._message_id # pylint: disable=access-member-before-definition + try: + req = self.session.post( + f'{self.API}{self.token}/sendMessage', + data={'text': f"`{self.text}`", 'chat_id': self.chat_id, + 'parse_mode': 'MarkdownV2'}) + res = req.json() + req.raise_for_status() + except Exception as e: + if req.status_code == 429: + warn("Creation rate limit: try increasing `mininterval`.", + TqdmWarning, stacklevel=2) + else: + tqdm_auto.write(str(e)) + else: + self._message_id = res['result']['message_id'] + return self._message_id + + def write(self, s): + """Replaces internal `message_id`'s text with `s`.""" + if not s: + s = "..." + s = s.replace('\r', '').strip() + if s == self.text: + return # avoid duplicate message Bot error + message_id = self.message_id + if message_id is None: + return + self.text = s + try: + future = self.submit( + self.session.post, f'{self.API}{self.token}/editMessageText', + data={'text': f"`{s}`", 'chat_id': self.chat_id, + 'message_id': message_id, 'parse_mode': 'MarkdownV2'}) + except Exception as e: + tqdm_auto.write(str(e)) + else: + return future + + def delete(self): + """Deletes internal `message_id`.""" + try: + future = self.submit( + self.session.post, '{self.API}{self.token}/deleteMessage', + data={'chat_id': self.chat_id, 'message_id': self.message_id}) + except Exception as e: + tqdm_auto.write(str(e)) + else: + return future + + +class tqdm_telegram(tqdm_auto): # pylint: disable=inconsistent-mro + """ + Standard `tqdm.auto.tqdm` but also sends updates to a Telegram Bot. + May take a few seconds to create (`__init__`). + + - create a bot + - copy its `{token}` + - add the bot to a chat and send it a message such as `/start` + - go to to find out + the `{chat_id}` + - paste the `{token}` & `{chat_id}` below + + >>> from tqdm.contrib.telegram import tqdm, trange + >>> for i in tqdm(iterable, token='{token}', chat_id='{chat_id}'): + ... ... + """ + @envwrap("tqdm", "telegram", is_method=True) + def __init__(self, *args, token=None, chat_id=None, **kwargs): + """ + Parameters + ---------- + token : str, required. Telegram token + [default: ${TQDM_TELEGRAM_TOKEN}]. + chat_id : str, required. Telegram chat ID + [default: ${TQDM_TELEGRAM_CHAT_ID}]. + + See `tqdm.auto.tqdm.__init__` for other parameters. + """ + if not kwargs.get('disable'): + kwargs = kwargs.copy() + self.tgio = TelegramIO(token, chat_id) + super().__init__(*args, **kwargs) + + def display(self, **kwargs): # pylint: disable=arguments-differ + super().display(**kwargs) + fmt = self.format_dict + if fmt.get('bar_format', None): + fmt['bar_format'] = fmt['bar_format'].replace( + '', '{bar:10u}').replace('{bar}', '{bar:10u}') + self.tgio.write(self.format_meter(**fmt)) + + def clear(self, *args, **kwargs): + super().clear(*args, **kwargs) + if not self.disable: + self.tgio.write("") + + def close(self): + if self.disable: + return + super().close() + if not (self.leave or (self.leave is None and self.pos == 0)): + self.tgio.delete() + + +def ttgrange(*args, **kwargs): + """Shortcut for `tqdm.contrib.telegram.tqdm(range(*args), **kwargs)`.""" + return tqdm_telegram(range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_telegram +trange = ttgrange diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/utils_worker.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/utils_worker.py new file mode 100644 index 0000000000000000000000000000000000000000..89fafc8c92f8ed085077b37eb58329f2588bd5d7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/contrib/utils_worker.py @@ -0,0 +1,38 @@ +""" +IO/concurrency helpers for `tqdm.contrib`. +""" +from collections import deque +from concurrent.futures import ThreadPoolExecutor + +from ..auto import tqdm as tqdm_auto + +__author__ = {"github.com/": ["casperdcl"]} +__all__ = ['MonoWorker'] + + +class MonoWorker: + """ + Supports one running task and one waiting task. + The waiting task is the most recent submitted (others are discarded). + """ + def __init__(self): + self.pool = ThreadPoolExecutor(max_workers=1) + self.futures = deque([], 2) + + def submit(self, func, *args, **kwargs): + """`func(*args, **kwargs)` may replace currently waiting task.""" + futures = self.futures + if len(futures) == futures.maxlen: + running = futures.popleft() + if not running.done(): + if len(futures): # clear waiting + waiting = futures.pop() + waiting.cancel() + futures.appendleft(running) # re-insert running + try: + waiting = self.pool.submit(func, *args, **kwargs) + except Exception as e: + tqdm_auto.write(str(e)) + else: + futures.append(waiting) + return waiting diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/dask.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/dask.py new file mode 100644 index 0000000000000000000000000000000000000000..57f1b668f59dc5991019eee34c7df3232a2c2cd7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/dask.py @@ -0,0 +1,44 @@ +from functools import partial + +from dask.callbacks import Callback + +from .auto import tqdm as tqdm_auto + +__author__ = {"github.com/": ["casperdcl"]} +__all__ = ['TqdmCallback'] + + +class TqdmCallback(Callback): + """Dask callback for task progress.""" + def __init__(self, start=None, pretask=None, tqdm_class=tqdm_auto, + **tqdm_kwargs): + """ + Parameters + ---------- + tqdm_class : optional + `tqdm` class to use for bars [default: `tqdm.auto.tqdm`]. + tqdm_kwargs : optional + Any other arguments used for all bars. + """ + super().__init__(start=start, pretask=pretask) + if tqdm_kwargs: + tqdm_class = partial(tqdm_class, **tqdm_kwargs) + self.tqdm_class = tqdm_class + + def _start_state(self, _, state): + self.pbar = self.tqdm_class(total=sum( + len(state[k]) for k in ['ready', 'waiting', 'running', 'finished'])) + + def _posttask(self, *_, **__): + self.pbar.update() + + def _finish(self, *_, **__): + self.pbar.close() + + def display(self): + """Displays in the current cell in Notebooks.""" + container = getattr(self.bar, 'container', None) + if container is None: + return + from .notebook import display + display(container) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/gui.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/gui.py new file mode 100644 index 0000000000000000000000000000000000000000..e7995089f09281d27d52eb082d935854360044a8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/gui.py @@ -0,0 +1,179 @@ +""" +Matplotlib GUI progress bar decorator for iterators. + +Usage: +>>> from tqdm.gui import trange, tqdm +>>> for i in trange(10): +... ... +""" +# future division is important to divide integers and get as +# a result precise floating numbers (instead of truncated int) +import re +from warnings import warn + +# to inherit from the tqdm class +from .std import TqdmExperimentalWarning +from .std import tqdm as std_tqdm + +# import compatibility functions and utilities + +__author__ = {"github.com/": ["casperdcl", "lrq3000"]} +__all__ = ['tqdm_gui', 'tgrange', 'tqdm', 'trange'] + + +class tqdm_gui(std_tqdm): # pragma: no cover + """Experimental Matplotlib GUI version of tqdm!""" + # TODO: @classmethod: write() on GUI? + def __init__(self, *args, **kwargs): + from collections import deque + + import matplotlib as mpl + import matplotlib.pyplot as plt + kwargs = kwargs.copy() + kwargs['gui'] = True + colour = kwargs.pop('colour', 'g') + super().__init__(*args, **kwargs) + + if self.disable: + return + + warn("GUI is experimental/alpha", TqdmExperimentalWarning, stacklevel=2) + self.mpl = mpl + self.plt = plt + + # Remember if external environment uses toolbars + self.toolbar = self.mpl.rcParams['toolbar'] + self.mpl.rcParams['toolbar'] = 'None' + + self.mininterval = max(self.mininterval, 0.5) + self.fig, ax = plt.subplots(figsize=(9, 2.2)) + # self.fig.subplots_adjust(bottom=0.2) + total = self.__len__() # avoids TypeError on None #971 + if total is not None: + self.xdata = [] + self.ydata = [] + self.zdata = [] + else: + self.xdata = deque([]) + self.ydata = deque([]) + self.zdata = deque([]) + self.line1, = ax.plot(self.xdata, self.ydata, color='b') + self.line2, = ax.plot(self.xdata, self.zdata, color='k') + ax.set_ylim(0, 0.001) + if total is not None: + ax.set_xlim(0, 100) + ax.set_xlabel("percent") + self.fig.legend((self.line1, self.line2), ("cur", "est"), + loc='center right') + # progress bar + self.hspan = plt.axhspan(0, 0.001, xmin=0, xmax=0, color=colour) + else: + # ax.set_xlim(-60, 0) + ax.set_xlim(0, 60) + ax.invert_xaxis() + ax.set_xlabel("seconds") + ax.legend(("cur", "est"), loc='lower left') + ax.grid() + # ax.set_xlabel('seconds') + ax.set_ylabel((self.unit if self.unit else "it") + "/s") + if self.unit_scale: + plt.ticklabel_format(style='sci', axis='y', scilimits=(0, 0)) + ax.yaxis.get_offset_text().set_x(-0.15) + + # Remember if external environment is interactive + self.wasion = plt.isinteractive() + plt.ion() + self.ax = ax + + def close(self): + if self.disable: + return + + self.disable = True + + with self.get_lock(): + self._instances.remove(self) + + # Restore toolbars + self.mpl.rcParams['toolbar'] = self.toolbar + # Return to non-interactive mode + if not self.wasion: + self.plt.ioff() + if self.leave: + self.display() + else: + self.plt.close(self.fig) + + def clear(self, *_, **__): + pass + + def display(self, *_, **__): + n = self.n + cur_t = self._time() + elapsed = cur_t - self.start_t + delta_it = n - self.last_print_n + delta_t = cur_t - self.last_print_t + + # Inline due to multiple calls + total = self.total + xdata = self.xdata + ydata = self.ydata + zdata = self.zdata + ax = self.ax + line1 = self.line1 + line2 = self.line2 + hspan = getattr(self, 'hspan', None) + # instantaneous rate + y = delta_it / delta_t + # overall rate + z = n / elapsed + # update line data + xdata.append(n * 100.0 / total if total else cur_t) + ydata.append(y) + zdata.append(z) + + # Discard old values + # xmin, xmax = ax.get_xlim() + # if (not total) and elapsed > xmin * 1.1: + if (not total) and elapsed > 66: + xdata.popleft() + ydata.popleft() + zdata.popleft() + + ymin, ymax = ax.get_ylim() + if y > ymax or z > ymax: + ymax = 1.1 * y + ax.set_ylim(ymin, ymax) + ax.figure.canvas.draw() + + if total: + line1.set_data(xdata, ydata) + line2.set_data(xdata, zdata) + if hspan: + hspan.set_xy((0, ymin)) + hspan.set_height(ymax - ymin) + hspan.set_width(n / total) + else: + t_ago = [cur_t - i for i in xdata] + line1.set_data(t_ago, ydata) + line2.set_data(t_ago, zdata) + + d = self.format_dict + # remove {bar} + d['bar_format'] = (d['bar_format'] or "{l_bar}{r_bar}").replace( + "{bar}", "") + msg = self.format_meter(**d) + if '' in msg: + msg = "".join(re.split(r'\|?\|?', msg, maxsplit=1)) + ax.set_title(msg, fontname="DejaVu Sans Mono", fontsize=11) + self.plt.pause(1e-9) + + +def tgrange(*args, **kwargs): + """Shortcut for `tqdm.gui.tqdm(range(*args), **kwargs)`.""" + return tqdm_gui(range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_gui +trange = tgrange diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/keras.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/keras.py new file mode 100644 index 0000000000000000000000000000000000000000..317b76a69539589f71334836cbe31fda063a7a46 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/keras.py @@ -0,0 +1,123 @@ +from copy import copy +from functools import partial + +from .auto import tqdm as tqdm_auto + +try: + import keras +except (ImportError, AttributeError) as e: + try: + from tensorflow import keras + except ImportError: + raise e +__author__ = {"github.com/": ["casperdcl"]} +__all__ = ['TqdmCallback'] + + +class TqdmCallback(keras.callbacks.Callback): + """Keras callback for epoch and batch progress.""" + @staticmethod + def bar2callback(bar, pop=None, delta=(lambda logs: 1)): + def callback(_, logs=None): + n = delta(logs) + if logs: + if pop: + logs = copy(logs) + for i in pop: + logs.pop(i, 0) + bar.set_postfix(logs, refresh=False) + bar.update(n) + + return callback + + def __init__(self, epochs=None, data_size=None, batch_size=None, verbose=1, + tqdm_class=tqdm_auto, **tqdm_kwargs): + """ + Parameters + ---------- + epochs : int, optional + data_size : int, optional + Number of training pairs. + batch_size : int, optional + Number of training pairs per batch. + verbose : int + 0: epoch, 1: batch (transient), 2: batch. [default: 1]. + Will be set to `0` unless both `data_size` and `batch_size` + are given. + tqdm_class : optional + `tqdm` class to use for bars [default: `tqdm.auto.tqdm`]. + tqdm_kwargs : optional + Any other arguments used for all bars. + """ + if tqdm_kwargs: + tqdm_class = partial(tqdm_class, **tqdm_kwargs) + self.tqdm_class = tqdm_class + self.epoch_bar = tqdm_class(total=epochs, unit='epoch') + self.on_epoch_end = self.bar2callback(self.epoch_bar) + if data_size and batch_size: + self.batches = batches = (data_size + batch_size - 1) // batch_size + else: + self.batches = batches = None + self.verbose = verbose + if verbose == 1: + self.batch_bar = tqdm_class(total=batches, unit='batch', leave=False) + self.on_batch_end = self.bar2callback( + self.batch_bar, pop=['batch', 'size'], + delta=lambda logs: logs.get('size', 1)) + + def on_train_begin(self, *_, **__): + params = self.params.get + auto_total = params('epochs', params('nb_epoch', None)) + if auto_total is not None and auto_total != self.epoch_bar.total: + self.epoch_bar.reset(total=auto_total) + + def on_epoch_begin(self, epoch, *_, **__): + if self.epoch_bar.n < epoch: + ebar = self.epoch_bar + ebar.n = ebar.last_print_n = ebar.initial = epoch + if self.verbose: + params = self.params.get + total = params('samples', params( + 'nb_sample', params('steps', None))) or self.batches + if self.verbose == 2: + if hasattr(self, 'batch_bar'): + self.batch_bar.close() + self.batch_bar = self.tqdm_class( + total=total, unit='batch', leave=True, + unit_scale=1 / (params('batch_size', 1) or 1)) + self.on_batch_end = self.bar2callback( + self.batch_bar, pop=['batch', 'size'], + delta=lambda logs: logs.get('size', 1)) + elif self.verbose == 1: + self.batch_bar.unit_scale = 1 / (params('batch_size', 1) or 1) + self.batch_bar.reset(total=total) + else: + raise KeyError('Unknown verbosity') + + def on_train_end(self, *_, **__): + if hasattr(self, 'batch_bar'): + self.batch_bar.close() + self.epoch_bar.close() + + def display(self): + """Displays in the current cell in Notebooks.""" + container = getattr(self.epoch_bar, 'container', None) + if container is None: + return + from .notebook import display + display(container) + batch_bar = getattr(self, 'batch_bar', None) + if batch_bar is not None: + display(batch_bar.container) + + @staticmethod + def _implements_train_batch_hooks(): + return True + + @staticmethod + def _implements_test_batch_hooks(): + return True + + @staticmethod + def _implements_predict_batch_hooks(): + return True diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/notebook.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/notebook.py new file mode 100644 index 0000000000000000000000000000000000000000..f6e1b0a6d15c5286354264116d2815b4285b8880 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/notebook.py @@ -0,0 +1,317 @@ +""" +IPython/Jupyter Notebook progress bar decorator for iterators. +Includes a default `range` iterator printing to `stderr`. + +Usage: +>>> from tqdm.notebook import trange, tqdm +>>> for i in trange(10): +... ... +""" +# import compatibility functions and utilities +import re +import sys +from html import escape +from weakref import proxy + +# to inherit from the tqdm class +from .std import tqdm as std_tqdm + +if True: # pragma: no cover + # import IPython/Jupyter base widget and display utilities + IPY = 0 + try: # IPython 4.x + import ipywidgets # noqa: F401, pylint: disable=unused-import + IPY = 4 + except ImportError: # IPython 3.x / 2.x + IPY = 32 + import warnings + with warnings.catch_warnings(): + warnings.filterwarnings( + 'ignore', message=".*The `IPython.html` package has been deprecated.*") + try: + import IPython.html.widgets + except ImportError: + pass + else: + ipywidgets = IPython.html.widgets + + try: # IPython 4.x / 3.x + if IPY == 32: + from IPython.html.widgets import HTML + from IPython.html.widgets import FloatProgress as IProgress + from IPython.html.widgets import HBox + IPY = 3 + else: + from ipywidgets import HTML + from ipywidgets import FloatProgress as IProgress + from ipywidgets import HBox + except ImportError: + try: # IPython 2.x + from IPython.html.widgets import HTML + from IPython.html.widgets import ContainerWidget as HBox + from IPython.html.widgets import FloatProgressWidget as IProgress + IPY = 2 + except ImportError: + IPY = 0 + IProgress = None + HBox = object + + try: + from IPython.display import display # , clear_output + except ImportError: + pass + +__author__ = {"github.com/": ["lrq3000", "casperdcl", "alexanderkuk"]} +__all__ = ['tqdm_notebook', 'tnrange', 'tqdm', 'trange'] +WARN_NOIPYW = ("IProgress not found. Please update jupyter and ipywidgets." + " See https://ipywidgets.readthedocs.io/en/stable" + "/user_install.html") + + +class TqdmHBox(HBox): + """`ipywidgets.HBox` with a pretty representation""" + def _json_(self, pretty=None): + pbar = getattr(self, 'pbar', None) + if pbar is None: + return {} + d = pbar.format_dict + if pretty is not None: + d["ascii"] = not pretty + return d + + def __repr__(self, pretty=False): + pbar = getattr(self, 'pbar', None) + if pbar is None: + return super().__repr__() + return pbar.format_meter(**self._json_(pretty)) + + def _repr_pretty_(self, pp, *_, **__): + pp.text(self.__repr__(True)) + + +class tqdm_notebook(std_tqdm): + """ + Experimental IPython/Jupyter Notebook widget using tqdm! + """ + @staticmethod + def status_printer(_, total=None, desc=None, ncols=None): + """ + Manage the printing of an IPython/Jupyter Notebook progress bar widget. + """ + # Fallback to text bar if there's no total + # DEPRECATED: replaced with an 'info' style bar + # if not total: + # return super(tqdm_notebook, tqdm_notebook).status_printer(file) + + # fp = file + + # Prepare IPython progress bar + if IProgress is None: # #187 #451 #558 #872 + raise ImportError(WARN_NOIPYW) + if total: + pbar = IProgress(min=0, max=total) + else: # No total? Show info style bar with no progress tqdm status + pbar = IProgress(min=0, max=1) + pbar.value = 1 + pbar.bar_style = 'info' + if ncols is None: + pbar.layout.width = "20px" + + ltext = HTML() + rtext = HTML() + if desc: + ltext.value = desc + container = TqdmHBox(children=[ltext, pbar, rtext]) + # Prepare layout + if ncols is not None: # use default style of ipywidgets + # ncols could be 100, "100px", "100%" + ncols = str(ncols) # ipywidgets only accepts string + try: + if int(ncols) > 0: # isnumeric and positive + ncols += 'px' + except ValueError: + pass + pbar.layout.flex = '2' + container.layout.width = ncols + container.layout.display = 'inline-flex' + container.layout.flex_flow = 'row wrap' + + return container + + def display(self, msg=None, pos=None, + # additional signals + close=False, bar_style=None, check_delay=True): + # Note: contrary to native tqdm, msg='' does NOT clear bar + # goal is to keep all infos if error happens so user knows + # at which iteration the loop failed. + + # Clear previous output (really necessary?) + # clear_output(wait=1) + + if not msg and not close: + d = self.format_dict + # remove {bar} + d['bar_format'] = (d['bar_format'] or "{l_bar}{r_bar}").replace( + "{bar}", "") + msg = self.format_meter(**d) + + ltext, pbar, rtext = self.container.children + pbar.value = self.n + + if msg: + msg = msg.replace(' ', '\u2007') # fix html space padding + # html escape special characters (like '&') + if '' in msg: + left, right = map(escape, re.split(r'\|?\|?', msg, maxsplit=1)) + else: + left, right = '', escape(msg) + + # Update description + ltext.value = left + # never clear the bar (signal: msg='') + if right: + rtext.value = right + + # Change bar style + if bar_style: + # Hack-ish way to avoid the danger bar_style being overridden by + # success because the bar gets closed after the error... + if pbar.bar_style != 'danger' or bar_style != 'success': + pbar.bar_style = bar_style + + # Special signal to close the bar + if close and pbar.bar_style != 'danger': # hide only if no error + try: + self.container.close() + except AttributeError: + self.container.visible = False + self.container.layout.visibility = 'hidden' # IPYW>=8 + + if check_delay and self.delay > 0 and not self.displayed: + display(self.container) + self.displayed = True + + @property + def colour(self): + if hasattr(self, 'container'): + return self.container.children[-2].style.bar_color + + @colour.setter + def colour(self, bar_color): + if hasattr(self, 'container'): + self.container.children[-2].style.bar_color = bar_color + + def __init__(self, *args, **kwargs): + """ + Supports the usual `tqdm.tqdm` parameters as well as those listed below. + + Parameters + ---------- + display : Whether to call `display(self.container)` immediately + [default: True]. + """ + kwargs = kwargs.copy() + # Setup default output + file_kwarg = kwargs.get('file', sys.stderr) + if file_kwarg is sys.stderr or file_kwarg is None: + kwargs['file'] = sys.stdout # avoid the red block in IPython + + # Initialize parent class + avoid printing by using gui=True + kwargs['gui'] = True + # convert disable = None to False + kwargs['disable'] = bool(kwargs.get('disable', False)) + colour = kwargs.pop('colour', None) + display_here = kwargs.pop('display', True) + super().__init__(*args, **kwargs) + if self.disable or not kwargs['gui']: + self.disp = lambda *_, **__: None + return + + # Get bar width + self.ncols = '100%' if self.dynamic_ncols else kwargs.get("ncols", None) + + # Replace with IPython progress bar display (with correct total) + unit_scale = 1 if self.unit_scale is True else self.unit_scale or 1 + total = self.total * unit_scale if self.total else self.total + self.container = self.status_printer(self.fp, total, self.desc, self.ncols) + self.container.pbar = proxy(self) + self.displayed = False + if display_here and self.delay <= 0: + display(self.container) + self.displayed = True + self.disp = self.display + self.colour = colour + + # Print initial bar state + if not self.disable: + self.display(check_delay=False) + + def __iter__(self): + try: + it = super().__iter__() + yield from it + # NB: except ... [ as ...] breaks IPython async KeyboardInterrupt + except: # NOQA + self.disp(bar_style='danger') + raise + # NB: don't `finally: close()` + # since this could be a shared bar which the user will `reset()` + + def update(self, n=1): + try: + return super().update(n=n) + # NB: except ... [ as ...] breaks IPython async KeyboardInterrupt + except: # NOQA + # cannot catch KeyboardInterrupt when using manual tqdm + # as the interrupt will most likely happen on another statement + self.disp(bar_style='danger') + raise + # NB: don't `finally: close()` + # since this could be a shared bar which the user will `reset()` + + def close(self): + if self.disable: + return + super().close() + # Try to detect if there was an error or KeyboardInterrupt + # in manual mode: if n < total, things probably got wrong + if self.total and self.n < self.total: + self.disp(bar_style='danger', check_delay=False) + else: + if self.leave: + self.disp(bar_style='success', check_delay=False) + else: + self.disp(close=True, check_delay=False) + + def clear(self, *_, **__): + pass + + def reset(self, total=None): + """ + Resets to 0 iterations for repeated use. + + Consider combining with `leave=True`. + + Parameters + ---------- + total : int or float, optional. Total to use for the new bar. + """ + if self.disable: + return super().reset(total=total) + _, pbar, _ = self.container.children + pbar.bar_style = '' + if total is not None: + pbar.max = total + if not self.total and self.ncols is None: # no longer unknown total + pbar.layout.width = None # reset width + return super().reset(total=total) + + +def tnrange(*args, **kwargs): + """Shortcut for `tqdm.notebook.tqdm(range(*args), **kwargs)`.""" + return tqdm_notebook(range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_notebook +trange = tnrange diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/rich.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/rich.py new file mode 100644 index 0000000000000000000000000000000000000000..3d392edaf115a93f7c145de52cbe8978dcf1ede8 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/rich.py @@ -0,0 +1,151 @@ +""" +`rich.progress` decorator for iterators. + +Usage: +>>> from tqdm.rich import trange, tqdm +>>> for i in trange(10): +... ... +""" +from warnings import warn + +from rich.progress import ( + BarColumn, Progress, ProgressColumn, Text, TimeElapsedColumn, TimeRemainingColumn, filesize) + +from .std import TqdmExperimentalWarning +from .std import tqdm as std_tqdm + +__author__ = {"github.com/": ["casperdcl"]} +__all__ = ['tqdm_rich', 'trrange', 'tqdm', 'trange'] + + +class FractionColumn(ProgressColumn): + """Renders completed/total, e.g. '0.5/2.3 G'.""" + def __init__(self, unit_scale=False, unit_divisor=1000): + self.unit_scale = unit_scale + self.unit_divisor = unit_divisor + super().__init__() + + def render(self, task): + """Calculate common unit for completed and total.""" + completed = int(task.completed) + total = int(task.total) + if self.unit_scale: + unit, suffix = filesize.pick_unit_and_suffix( + total, + ["", "K", "M", "G", "T", "P", "E", "Z", "Y"], + self.unit_divisor, + ) + else: + unit, suffix = filesize.pick_unit_and_suffix(total, [""], 1) + precision = 0 if unit == 1 else 1 + return Text( + f"{completed/unit:,.{precision}f}/{total/unit:,.{precision}f} {suffix}", + style="progress.download") + + +class RateColumn(ProgressColumn): + """Renders human readable transfer speed.""" + def __init__(self, unit="", unit_scale=False, unit_divisor=1000): + self.unit = unit + self.unit_scale = unit_scale + self.unit_divisor = unit_divisor + super().__init__() + + def render(self, task): + """Show data transfer speed.""" + speed = task.speed + if speed is None: + return Text(f"? {self.unit}/s", style="progress.data.speed") + if self.unit_scale: + unit, suffix = filesize.pick_unit_and_suffix( + speed, + ["", "K", "M", "G", "T", "P", "E", "Z", "Y"], + self.unit_divisor, + ) + else: + unit, suffix = filesize.pick_unit_and_suffix(speed, [""], 1) + precision = 0 if unit == 1 else 1 + return Text(f"{speed/unit:,.{precision}f} {suffix}{self.unit}/s", + style="progress.data.speed") + + +class tqdm_rich(std_tqdm): # pragma: no cover + """Experimental rich.progress GUI version of tqdm!""" + # TODO: @classmethod: write()? + def __init__(self, *args, **kwargs): + """ + This class accepts the following parameters *in addition* to + the parameters accepted by `tqdm`. + + Parameters + ---------- + progress : tuple, optional + arguments for `rich.progress.Progress()`. + options : dict, optional + keyword arguments for `rich.progress.Progress()`. + """ + kwargs = kwargs.copy() + kwargs['gui'] = True + # convert disable = None to False + kwargs['disable'] = bool(kwargs.get('disable', False)) + progress = kwargs.pop('progress', None) + options = kwargs.pop('options', {}).copy() + super().__init__(*args, **kwargs) + + if self.disable: + return + + warn("rich is experimental/alpha", TqdmExperimentalWarning, stacklevel=2) + d = self.format_dict + if progress is None: + progress = ( + "[progress.description]{task.description}" + "[progress.percentage]{task.percentage:>4.0f}%", + BarColumn(bar_width=None), + FractionColumn( + unit_scale=d['unit_scale'], unit_divisor=d['unit_divisor']), + "[", TimeElapsedColumn(), "<", TimeRemainingColumn(), + ",", RateColumn(unit=d['unit'], unit_scale=d['unit_scale'], + unit_divisor=d['unit_divisor']), "]" + ) + options.setdefault('transient', not self.leave) + self._prog = Progress(*progress, **options) + self._prog.__enter__() + self._task_id = self._prog.add_task(self.desc or "", **d) + + def close(self): + if self.disable: + return + self.display() # print 100%, vis #1306 + super().close() + self._prog.__exit__(None, None, None) + + def clear(self, *_, **__): + pass + + def display(self, *_, **__): + if not hasattr(self, '_prog'): + return + self._prog.update(self._task_id, completed=self.n, description=self.desc) + + def reset(self, total=None): + """ + Resets to 0 iterations for repeated use. + + Parameters + ---------- + total : int or float, optional. Total to use for the new bar. + """ + if hasattr(self, '_prog'): + self._prog.reset(total=total) + super().reset(total=total) + + +def trrange(*args, **kwargs): + """Shortcut for `tqdm.rich.tqdm(range(*args), **kwargs)`.""" + return tqdm_rich(range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_rich +trange = trrange diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/std.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/std.py new file mode 100644 index 0000000000000000000000000000000000000000..79a868f5362d2ee68ae606503a62644d7ac90164 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/std.py @@ -0,0 +1,1525 @@ +""" +Customisable progress bar decorator for iterators. +Includes a default `range` iterator printing to `stderr`. + +Usage: +>>> from tqdm import trange, tqdm +>>> for i in trange(10): +... ... +""" +import sys +from collections import OrderedDict, defaultdict +from contextlib import contextmanager +from datetime import datetime, timedelta, timezone +from numbers import Number +from time import time +from warnings import warn +from weakref import WeakSet + +from ._monitor import TMonitor +from .utils import ( + CallbackIOWrapper, Comparable, DisableOnWriteError, FormatReplace, SimpleTextIOWrapper, + _is_ascii, _screen_shape_wrapper, _supports_unicode, _term_move_up, disp_len, disp_trim, + envwrap) + +__author__ = "https://github.com/tqdm/tqdm#contributions" +__all__ = ['tqdm', 'trange', + 'TqdmTypeError', 'TqdmKeyError', 'TqdmWarning', + 'TqdmExperimentalWarning', 'TqdmDeprecationWarning', + 'TqdmMonitorWarning'] + + +class TqdmTypeError(TypeError): + pass + + +class TqdmKeyError(KeyError): + pass + + +class TqdmWarning(Warning): + """base class for all tqdm warnings. + + Used for non-external-code-breaking errors, such as garbled printing. + """ + def __init__(self, msg, fp_write=None): # noqa: B042 + if fp_write is not None: + fp_write("\n" + self.__class__.__name__ + ": " + str(msg).rstrip() + '\n') + else: + super().__init__(msg) + + +class TqdmExperimentalWarning(TqdmWarning, FutureWarning): + """beta feature, unstable API and behaviour""" + + +class TqdmDeprecationWarning(TqdmWarning, DeprecationWarning): + """may be removed in a future release""" + # not suppressed if raised + + +class TqdmMonitorWarning(TqdmWarning, RuntimeWarning): + """tqdm monitor errors which do not affect external functionality""" + + +def TRLock(*args, **kwargs): + """threading RLock""" + try: + from threading import RLock + return RLock(*args, **kwargs) + except (ImportError, OSError): # pragma: no cover + pass + + +class TqdmDefaultWriteLock: + """ + Provide a default write lock for thread and multiprocessing safety. + Works only on platforms supporting `fork` (so Windows is excluded). + You must initialise a `tqdm` or `TqdmDefaultWriteLock` instance + before forking in order for the write lock to work. + On Windows, you need to supply the lock from the parent to the children as + an argument to joblib or the parallelism lib you use. + """ + # global thread lock so no setup required for multithreading. + # NB: Do not create multiprocessing lock as it sets the multiprocessing + # context, disallowing `spawn()`/`forkserver()` + th_lock = TRLock() + + def __init__(self): + # Create global parallelism locks to avoid racing issues with parallel + # bars works only if fork available (Linux/MacOSX, but not Windows) + cls = type(self) + root_lock = cls.th_lock + if root_lock is not None: + root_lock.acquire() + cls.create_mp_lock() + self.locks = [lk for lk in [cls.mp_lock, cls.th_lock] if lk is not None] + if root_lock is not None: + root_lock.release() + + def acquire(self, *a, **k): + for lock in self.locks: + lock.acquire(*a, **k) + + def release(self): + for lock in self.locks[::-1]: # Release in inverse order of acquisition + lock.release() + + def __enter__(self): + self.acquire() + + def __exit__(self, *exc): + self.release() + + @classmethod + def create_mp_lock(cls): + if not hasattr(cls, 'mp_lock'): + try: + from multiprocessing import RLock + cls.mp_lock = RLock() + except (ImportError, OSError): # pragma: no cover + cls.mp_lock = None + + @classmethod + def create_th_lock(cls): + assert hasattr(cls, 'th_lock') + warn("create_th_lock not needed anymore", TqdmDeprecationWarning, stacklevel=2) + + +class Bar: + """ + `str.format`-able bar with format specifiers: `[width][type]` + + - `width` + + unspecified (default): use `self.default_len` + + `int >= 0`: overrides `self.default_len` + + `int < 0`: subtract from `self.default_len` + - `type` + + `a`: ascii (`charset=self.ASCII` override) + + `u`: unicode (`charset=self.UTF` override) + + `b`: blank (`charset=" "` override) + """ + ASCII = " 123456789#" + UTF = " " + ''.join(map(chr, range(0x258F, 0x2587, -1))) + BLANK = " " + COLOUR_RESET = '\x1b[0m' + COLOUR_RGB = '\x1b[38;2;%d;%d;%dm' + COLOURS = {'BLACK': '\x1b[30m', 'RED': '\x1b[31m', 'GREEN': '\x1b[32m', + 'YELLOW': '\x1b[33m', 'BLUE': '\x1b[34m', 'MAGENTA': '\x1b[35m', + 'CYAN': '\x1b[36m', 'WHITE': '\x1b[37m'} + + def __init__(self, frac, default_len=10, charset=UTF, colour=None): + if not 0 <= frac <= 1: + warn("clamping frac to range [0, 1]", TqdmWarning, stacklevel=2) + frac = max(0, min(1, frac)) + assert default_len > 0 + self.frac = frac + self.default_len = default_len + self.charset = charset + self.colour = colour + + @property + def colour(self): + return self._colour + + @colour.setter + def colour(self, value): + if not value: + self._colour = None + return + try: + if value.upper() in self.COLOURS: + self._colour = self.COLOURS[value.upper()] + elif value[0] == '#' and len(value) == 7: + self._colour = self.COLOUR_RGB % tuple( + int(i, 16) for i in (value[1:3], value[3:5], value[5:7])) + else: + raise KeyError + except (KeyError, AttributeError): + warn(f"Unknown colour ({value}); valid choices:" + f" [hex (#00ff00), {', '.join(self.COLOURS)}]", TqdmWarning, stacklevel=2) + self._colour = None + + def __format__(self, format_spec): + if format_spec: + _type = format_spec[-1].lower() + try: + charset = {'a': self.ASCII, 'u': self.UTF, 'b': self.BLANK}[_type] + except KeyError: + charset = self.charset + else: + format_spec = format_spec[:-1] + if format_spec: + N_BARS = int(format_spec) + if N_BARS < 0: + N_BARS += self.default_len + else: + N_BARS = self.default_len + else: + charset = self.charset + N_BARS = self.default_len + + nsyms = len(charset) - 1 + bar_length, frac_bar_length = divmod(int(self.frac * N_BARS * nsyms), nsyms) + + res = charset[-1] * bar_length + if bar_length < N_BARS: # whitespace padding + res = res + charset[frac_bar_length] + charset[0] * (N_BARS - bar_length - 1) + return self.colour + res + self.COLOUR_RESET if self.colour else res + + +class EMA: + """ + Exponential moving average: smoothing to give progressively lower + weights to older values. + + Parameters + ---------- + smoothing : float, optional + Smoothing factor in range [0, 1], [default: 0.3]. + Increase to give more weight to recent values. + Ranges from 0 (yields old value) to 1 (yields new value). + """ + def __init__(self, smoothing=0.3): + self.alpha = smoothing + self.last = 0 + self.calls = 0 + + def __call__(self, x=None): + """ + Parameters + ---------- + x : float + New value to include in EMA. + """ + beta = 1 - self.alpha + if x is not None: + self.last = self.alpha * x + beta * self.last + self.calls += 1 + return self.last / (1 - beta ** self.calls) if self.calls else self.last + + +class tqdm(Comparable): + """ + Decorate an iterable object, returning an iterator which acts exactly + like the original iterable, but prints a dynamically updating + progress bar every time a value is requested. + + Parameters + ---------- + iterable : iterable, optional + Iterable to decorate with a progress bar. + Leave blank to manually manage the updates. + desc : str, optional + Prefix for the progress bar. + total : int or float, optional + The number of expected iterations. If unspecified, + len(iterable) is used if possible. If float("inf") or as a last + resort, only basic progress statistics are displayed + (no ETA, no progress bar). + If `gui` is True and this parameter needs subsequent updating, + specify an initial arbitrary large positive number, + e.g. 9e9. + leave : bool, optional + If [default: True], keeps all traces of the progress bar + upon termination of iteration. + If `None`, will leave only if `position` is `0`. + file : `io.TextIOWrapper` or `io.StringIO`, optional + Specifies where to output the progress messages + (default: sys.stderr). Uses `file.write(str)` and `file.flush()` + methods. For encoding, see `write_bytes`. + ncols : int, optional + The width of the entire output message. If specified, + dynamically resizes the progress bar to stay within this bound. + If unspecified, attempts to use environment width. The + fallback is a meter width of 10 and no limit for the counter and + statistics. If 0, will not print any meter (only stats). + mininterval : float, optional + Minimum progress display update interval [default: 0.1] seconds. + maxinterval : float, optional + Maximum progress display update interval [default: 10] seconds. + Automatically adjusts `miniters` to correspond to `mininterval` + after long display update lag. Only works if `dynamic_miniters` + or monitor thread is enabled. + miniters : int or float, optional + Minimum progress display update interval, in iterations. + If 0 and `dynamic_miniters`, will automatically adjust to equal + `mininterval` (more CPU efficient, good for tight loops). + If > 0, will skip display of specified number of iterations. + Tweak this and `mininterval` to get very efficient loops. + If your progress is erratic with both fast and slow iterations + (network, skipping items, etc) you should set miniters=1. + ascii : bool or str, optional + If unspecified or False, use unicode (smooth blocks) to fill + the meter. The fallback is to use ASCII characters " 123456789#". + disable : bool, optional + Whether to disable the entire progress bar wrapper + [default: False]. If set to None, disable on non-TTY. + unit : str, optional + String that will be used to define the unit of each iteration + [default: it]. + unit_scale : bool or int or float, optional + If 1 or True, the number of iterations will be reduced/scaled + automatically and a metric prefix following the + International System of Units standard will be added + (kilo, mega, etc.) [default: False]. If any other non-zero + number, will scale `total` and `n`. + dynamic_ncols : bool, optional + If set, constantly alters `ncols` and `nrows` to the + environment (allowing for window resizes) [default: False]. + smoothing : float, optional + Exponential moving average smoothing factor for speed estimates + (ignored in GUI mode). Ranges from 0 (average speed) to 1 + (current/instantaneous speed) [default: 0.3]. + bar_format : str, optional + Specify a custom bar string formatting. May impact performance. + [default: '{l_bar}{bar}{r_bar}'], where + l_bar='{desc}: {percentage:3.0f}%|' and + r_bar='| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, ' + '{rate_fmt}{postfix}]' + Possible vars: l_bar, bar, r_bar, n, n_fmt, total, total_fmt, + percentage, elapsed, elapsed_s, ncols, nrows, desc, unit, + rate, rate_fmt, rate_noinv, rate_noinv_fmt, + rate_inv, rate_inv_fmt, postfix, unit_divisor, + remaining, remaining_s, eta. + Note that a trailing ": " is automatically removed after {desc} + if the latter is empty. + initial : int or float, optional + The initial counter value. Useful when restarting a progress + bar [default: 0]. If using float, consider specifying `{n:.3f}` + or similar in `bar_format`, or specifying `unit_scale`. + position : int, optional + Specify the line offset to print this bar (starting from 0) + Automatic if unspecified. + Useful to manage multiple bars at once (eg, from threads). + postfix : dict or *, optional + Specify additional stats to display at the end of the bar. + Calls `set_postfix(**postfix)` if possible (dict). + unit_divisor : float, optional + [default: 1000], ignored unless `unit_scale` is True. + write_bytes : bool, optional + Whether to write bytes. If (default: False) will write unicode. + lock_args : tuple, optional + Passed to `refresh` for intermediate output + (initialisation, iterating, and updating). + nrows : int, optional + The screen height. If specified, hides nested bars outside this + bound. If unspecified, attempts to use environment height. + The fallback is 20. + colour : str, optional + Bar colour (e.g. 'green', '#00ff00'). + delay : float, optional + Don't display until [default: 0] seconds have elapsed. + gui : bool, optional + WARNING: internal parameter - do not use. + Use tqdm.gui.tqdm(...) instead. If set, will attempt to use + matplotlib animations for a graphical output [default: False]. + + Returns + ------- + out : decorated iterator. + """ + + monitor_interval = 10 # set to 0 to disable the thread + monitor = None + _instances = WeakSet() + + @staticmethod + def format_sizeof(num, suffix='', divisor=1000): + """ + Formats a number (greater than unity) with SI Order of Magnitude + prefixes. + + Parameters + ---------- + num : float + Number ( >= 1) to format. + suffix : str, optional + Post-postfix [default: '']. + divisor : float, optional + Divisor between prefixes [default: 1000]. + + Returns + ------- + out : str + Number with Order of Magnitude SI unit postfix. + """ + for unit in ['', 'k', 'M', 'G', 'T', 'P', 'E', 'Z']: + if abs(num) < 999.5: + if abs(num) < 99.95: + if abs(num) < 9.995: + return f'{num:1.2f}{unit}{suffix}' + return f'{num:2.1f}{unit}{suffix}' + return f'{num:3.0f}{unit}{suffix}' + num /= divisor + return f'{num:3.1f}Y{suffix}' + + @staticmethod + def format_interval(t): + """ + Formats a number of seconds as a clock time, [H:]MM:SS + + Parameters + ---------- + t : int + Number of seconds. + + Returns + ------- + out : str + [H:]MM:SS + """ + sign = '-' if t < 0 else '' + mins, s = divmod(abs(int(t)), 60) + h, m = divmod(mins, 60) + return f'{sign}{h:d}:{m:02d}:{s:02d}' if h else f'{sign}{m:02d}:{s:02d}' + + @staticmethod + def format_num(n): + """ + Intelligent scientific notation (.3g). + + Parameters + ---------- + n : int or float or Numeric + A Number. + + Returns + ------- + out : str + Formatted number. + """ + f = f'{n:.3g}'.replace('e+0', 'e+').replace('e-0', 'e-') + n = str(n) + return f if len(f) < len(n) else n + + @staticmethod + def status_printer(file): + """ + Manage the printing and in-place updating of a line of characters. + Note that if the string is longer than a line, then in-place + updating may not work (it will print a new line at each refresh). + """ + fp = file + fp_flush = getattr(fp, 'flush', lambda: None) # pragma: no cover + if fp in (sys.stderr, sys.stdout): + getattr(sys.stderr, 'flush', lambda: None)() + getattr(sys.stdout, 'flush', lambda: None)() + + def fp_write(s): + fp.write(str(s)) + fp_flush() + + last_len = [0] + + def print_status(s): + len_s = disp_len(s) + fp_write('\r' + s + (' ' * max(last_len[0] - len_s, 0))) + last_len[0] = len_s + + return print_status + + @staticmethod + def format_meter(n, total, elapsed, ncols=None, prefix='', + ascii=False, # pylint: disable=redefined-builtin + unit='it', unit_scale=False, rate=None, bar_format=None, postfix=None, + unit_divisor=1000, initial=0, colour=None, **extra_kwargs): + """ + Return a string-based progress bar given some parameters + + Parameters + ---------- + n : int or float + Number of finished iterations. + total : int or float + The expected total number of iterations. If meaningless (None), + only basic progress statistics are displayed (no ETA). + elapsed : float + Number of seconds passed since start. + ncols : int, optional + The width of the entire output message. If specified, + dynamically resizes `{bar}` to stay within this bound + [default: None]. If `0`, will not print any bar (only stats). + The fallback is `{bar:10}`. + prefix : str, optional + Prefix message (included in total width) [default: '']. + Use as {desc} in bar_format string. + ascii : bool, optional or str, optional + If not set, use unicode (smooth blocks) to fill the meter + [default: False]. The fallback is to use ASCII characters + " 123456789#". + unit : str, optional + The iteration unit [default: 'it']. + unit_scale : bool or int or float, optional + If 1 or True, the number of iterations will be printed with an + appropriate SI metric prefix (k = 10^3, M = 10^6, etc.) + [default: False]. If any other non-zero number, will scale + `total` and `n`. + rate : float, optional + Manual override for iteration rate. + If [default: None], uses n/elapsed. + bar_format : str, optional + Specify a custom bar string formatting. May impact performance. + [default: '{l_bar}{bar}{r_bar}'], where + l_bar='{desc}: {percentage:3.0f}%|' and + r_bar='| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, ' + '{rate_fmt}{postfix}]' + Possible vars: l_bar, bar, r_bar, n, n_fmt, total, total_fmt, + percentage, elapsed, elapsed_s, ncols, nrows, desc, unit, + rate, rate_fmt, rate_noinv, rate_noinv_fmt, + rate_inv, rate_inv_fmt, postfix, unit_divisor, + remaining, remaining_s, eta. + Note that a trailing ": " is automatically removed after {desc} + if the latter is empty. + postfix : *, optional + Similar to `prefix`, but placed at the end + (e.g. for additional stats). + Note: postfix is usually a string (not a dict) for this method, + and will if possible be set to postfix = ', ' + postfix. + However other types are supported (#382). + unit_divisor : float, optional + [default: 1000], ignored unless `unit_scale` is True. + initial : int or float, optional + The initial counter value [default: 0]. + colour : str, optional + Bar colour (e.g. 'green', '#00ff00'). + + Returns + ------- + out : Formatted meter and stats, ready to display. + """ + + # sanity check: total + if total and n >= (total + 0.5): # allow float imprecision (#849) + total = None + + # apply custom scale if necessary + if unit_scale and unit_scale not in (True, 1): + if total: + total *= unit_scale + n *= unit_scale + if rate: + rate *= unit_scale # by default rate = self.avg_dn / self.avg_dt + unit_scale = False + + elapsed_str = tqdm.format_interval(elapsed) + + # if unspecified, attempt to use rate = average speed + # (we allow manual override since predicting time is an arcane art) + if rate is None and elapsed: + rate = (n - initial) / elapsed + inv_rate = 1 / rate if rate else None + format_sizeof = tqdm.format_sizeof + rate_noinv_fmt = ((format_sizeof(rate) if unit_scale else f'{rate:5.2f}') + if rate else '?') + unit + '/s' + rate_inv_fmt = ( + (format_sizeof(inv_rate) if unit_scale else f'{inv_rate:5.2f}') + if inv_rate else '?') + 's/' + unit + rate_fmt = rate_inv_fmt if inv_rate and inv_rate > 1 else rate_noinv_fmt + + if unit_scale: + n_fmt = format_sizeof(n, divisor=unit_divisor) + total_fmt = format_sizeof(total, divisor=unit_divisor) if total is not None else '?' + else: + n_fmt = str(n) + total_fmt = str(total) if total is not None else '?' + + try: + postfix = ', ' + postfix if postfix else '' + except TypeError: + pass + + remaining = (total - n) / rate if rate and total else 0 + remaining_str = tqdm.format_interval(remaining) if rate else '?' + try: + eta_dt = (datetime.now() + timedelta(seconds=remaining) + if rate and total else datetime.fromtimestamp(0, timezone.utc)) + except OverflowError: + eta_dt = datetime.max + + # format the stats displayed to the left and right sides of the bar + if prefix: + # old prefix setup work around + bool_prefix_colon_already = (prefix[-2:] == ": ") + l_bar = prefix if bool_prefix_colon_already else prefix + ": " + else: + l_bar = '' + + r_bar = f'| {n_fmt}/{total_fmt} [{elapsed_str}<{remaining_str}, {rate_fmt}{postfix}]' + + # Custom bar formatting + # Populate a dict with all available progress indicators + format_dict = { + # slight extension of self.format_dict + 'n': n, 'n_fmt': n_fmt, 'total': total, 'total_fmt': total_fmt, + 'elapsed': elapsed_str, 'elapsed_s': elapsed, + 'ncols': ncols, 'desc': prefix or '', 'unit': unit, + 'rate': inv_rate if inv_rate and inv_rate > 1 else rate, + 'rate_fmt': rate_fmt, 'rate_noinv': rate, + 'rate_noinv_fmt': rate_noinv_fmt, 'rate_inv': inv_rate, + 'rate_inv_fmt': rate_inv_fmt, + 'postfix': postfix, 'unit_divisor': unit_divisor, + 'colour': colour, + # plus more useful definitions + 'remaining': remaining_str, 'remaining_s': remaining, + 'l_bar': l_bar, 'r_bar': r_bar, 'eta': eta_dt, + **extra_kwargs} + + # total is known: we can predict some stats + if total: + # fractional and percentage progress + frac = n / total + percentage = frac * 100 + + l_bar += f'{percentage:3.0f}%|' + + if ncols == 0: + return l_bar[:-1] + r_bar[1:] + + format_dict.update(l_bar=l_bar) + if bar_format: + format_dict.update(percentage=percentage) + + # auto-remove colon for empty `{desc}` + if not prefix: + bar_format = bar_format.replace("{desc}: ", '') + else: + bar_format = "{l_bar}{bar}{r_bar}" + + full_bar = FormatReplace() + nobar = bar_format.format(bar=full_bar, **format_dict) + if not full_bar.format_called: + return nobar # no `{bar}`; nothing else to do + + # Formatting progress bar space available for bar's display + full_bar = Bar(frac, + max(1, ncols - disp_len(nobar)) if ncols else 10, + charset=Bar.ASCII if ascii is True else ascii or Bar.UTF, + colour=colour) + if not _is_ascii(full_bar.charset) and _is_ascii(bar_format): + bar_format = str(bar_format) + res = bar_format.format(bar=full_bar, **format_dict) + return disp_trim(res, ncols) if ncols else res + + elif bar_format: + # user-specified bar_format but no total + l_bar += '|' + format_dict.update(l_bar=l_bar, percentage=0) + full_bar = FormatReplace() + nobar = bar_format.format(bar=full_bar, **format_dict) + if not full_bar.format_called: + return nobar + full_bar = Bar(0, + max(1, ncols - disp_len(nobar)) if ncols else 10, + charset=Bar.BLANK, colour=colour) + res = bar_format.format(bar=full_bar, **format_dict) + return disp_trim(res, ncols) if ncols else res + else: + # no total: no bar & ETA, just progress stats + return (f'{(prefix + ": ") if prefix else ""}' + f'{n_fmt}{unit} [{elapsed_str}, {rate_fmt}{postfix}]') + + def __new__(cls, *_, **__): + instance = object.__new__(cls) + with cls.get_lock(): # also constructs lock if non-existent + cls._instances.add(instance) + # create monitoring thread + if cls.monitor_interval and (cls.monitor is None + or not cls.monitor.report()): + try: + cls.monitor = TMonitor(cls, cls.monitor_interval) + except Exception as e: # pragma: nocover + warn("tqdm:disabling monitor support" + " (monitor_interval = 0) due to:\n" + str(e), + TqdmMonitorWarning, stacklevel=2) + cls.monitor_interval = 0 + return instance + + @classmethod + def _get_free_pos(cls, instance=None): + """Skips specified instance.""" + positions = {abs(inst.pos) for inst in cls._instances + if inst is not instance and hasattr(inst, "pos")} + return min(set(range(len(positions) + 1)).difference(positions)) + + @classmethod + def _decr_instances(cls, instance): + """ + Remove from list and reposition another unfixed bar + to fill the new gap. + + This means that by default (where all nested bars are unfixed), + order is not maintained but screen flicker/blank space is minimised. + (tqdm<=4.44.1 moved ALL subsequent unfixed bars up.) + """ + with cls._lock: + try: + cls._instances.remove(instance) + except KeyError: + # if not instance.gui: # pragma: no cover + # raise + pass # py2: maybe magically removed already + # else: + if not instance.gui: + last = (instance.nrows or 20) - 1 + # find unfixed (`pos >= 0`) overflow (`pos >= nrows - 1`) + instances = list(filter( + lambda i: hasattr(i, "pos") and last <= i.pos, + cls._instances)) + # set first found to current `pos` + if instances: + inst = min(instances, key=lambda i: i.pos) + inst.clear(nolock=True) + inst.pos = abs(instance.pos) + + @classmethod + def write(cls, s, file=None, end="\n", nolock=False): + """Print a message via tqdm (without overlap with bars).""" + fp = file if file is not None else sys.stdout + with cls.external_write_mode(file=file, nolock=nolock): + # Write the message + fp.write(s) + fp.write(end) + + @classmethod + @contextmanager + def external_write_mode(cls, file=None, nolock=False): + """ + Disable tqdm within context and refresh tqdm when exits. + Useful when writing to standard output stream + """ + fp = file if file is not None else sys.stdout + + try: + if not nolock: + cls.get_lock().acquire() + # Clear all bars + inst_cleared = [] + for inst in getattr(cls, '_instances', []): + # Clear instance if in the target output file + # or if write output + tqdm output are both either + # sys.stdout or sys.stderr (because both are mixed in terminal) + if hasattr(inst, "start_t") and (inst.fp == fp or all( + f in (sys.stdout, sys.stderr) for f in (fp, inst.fp))): + inst.clear(nolock=True) + inst_cleared.append(inst) + yield + # Force refresh display of bars we cleared + for inst in inst_cleared: + inst.refresh(nolock=True) + finally: + if not nolock: + cls._lock.release() + + @classmethod + def set_lock(cls, lock): + """Set the global lock.""" + cls._lock = lock + + @classmethod + def get_lock(cls): + """Get the global lock. Construct it if it does not exist.""" + if not hasattr(cls, '_lock'): + cls._lock = TqdmDefaultWriteLock() + return cls._lock + + @classmethod + def pandas(cls, **tqdm_kwargs): + """ + Registers the current `tqdm` class with + pandas.core. + ( frame.DataFrame + | series.Series + | groupby.(generic.)DataFrameGroupBy + | groupby.(generic.)SeriesGroupBy + ).progress_apply + + A new instance will be created every time `progress_apply` is called, + and each instance will automatically `close()` upon completion. + + Parameters + ---------- + tqdm_kwargs : arguments for the tqdm instance + + Examples + -------- + >>> import pandas as pd + >>> import numpy as np + >>> from tqdm import tqdm + >>> from tqdm.gui import tqdm as tqdm_gui + >>> + >>> df = pd.DataFrame(np.random.randint(0, 100, (100000, 6))) + >>> tqdm.pandas(ncols=50) # can use tqdm_gui, optional kwargs, etc + >>> # Now you can use `progress_apply` instead of `apply` + >>> df.groupby(0).progress_apply(lambda x: x**2) + + References + ---------- + + """ + from warnings import catch_warnings, simplefilter + + from pandas.core.frame import DataFrame + from pandas.core.series import Series + try: + with catch_warnings(): + simplefilter("ignore", category=FutureWarning) + from pandas import Panel + except ImportError: # pandas>=1.2.0 + Panel = None + Rolling, Expanding = None, None + try: # pandas>=1.0.0 + from pandas.core.window.rolling import _Rolling_and_Expanding + except ImportError: + try: # pandas>=0.18.0 + from pandas.core.window import _Rolling_and_Expanding + except ImportError: # pandas>=1.2.0 + try: # pandas>=1.2.0 + from pandas.core.window.expanding import Expanding + from pandas.core.window.rolling import Rolling + _Rolling_and_Expanding = Rolling, Expanding + except ImportError: # pragma: no cover + _Rolling_and_Expanding = None + try: # pandas>=0.25.0 + from pandas.core.groupby.generic import SeriesGroupBy # , NDFrameGroupBy + from pandas.core.groupby.generic import DataFrameGroupBy + except ImportError: # pragma: no cover + try: # pandas>=0.23.0 + from pandas.core.groupby.groupby import DataFrameGroupBy, SeriesGroupBy + except ImportError: + from pandas.core.groupby import DataFrameGroupBy, SeriesGroupBy + try: # pandas>=0.23.0 + from pandas.core.groupby.groupby import GroupBy + except ImportError: # pragma: no cover + from pandas.core.groupby import GroupBy + + try: # pandas>=0.23.0 + from pandas.core.groupby.groupby import PanelGroupBy + except ImportError: + try: + from pandas.core.groupby import PanelGroupBy + except ImportError: # pandas>=0.25.0 + PanelGroupBy = None + + tqdm_kwargs = tqdm_kwargs.copy() + deprecated_t = [tqdm_kwargs.pop('deprecated_t', None)] + + def inner_generator(df_function='apply'): + def inner(df, func, *args, **kwargs): + """ + Parameters + ---------- + df : (DataFrame|Series)[GroupBy] + Data (may be grouped). + func : function + To be applied on the (grouped) data. + **kwargs : optional + Transmitted to `df.apply()`. + """ + + # Precompute total iterations + total = tqdm_kwargs.pop("total", getattr(df, 'ngroups', None)) + if total is None: # not grouped + if df_function == 'applymap': + total = df.size + elif isinstance(df, Series): + total = len(df) + elif (_Rolling_and_Expanding is None or + not isinstance(df, _Rolling_and_Expanding)): + # DataFrame or Panel + axis = kwargs.get('axis', 0) + if axis == 'index': + axis = 0 + elif axis == 'columns': + axis = 1 + # when axis=0, total is shape[axis1] + total = df.size // df.shape[axis] + + # Init bar + if deprecated_t[0] is not None: + t = deprecated_t[0] + deprecated_t[0] = None + else: + t = cls(total=total, **tqdm_kwargs) + + if len(args) > 0: + # *args intentionally not supported (see #244, #299) + TqdmDeprecationWarning( + "Except func, normal arguments are intentionally" + + " not supported by" + + " `(DataFrame|Series|GroupBy).progress_apply`." + + " Use keyword arguments instead.", + fp_write=getattr(t.fp, 'write', sys.stderr.write)) + + try: # pandas>=1.3.0,<3.0 + from pandas.core.common import is_builtin_func + except ImportError: # pandas<1.3.0 + is_builtin_func = getattr(df, '_is_builtin_func', lambda f: f) + try: + func = is_builtin_func(func) + except TypeError: + pass + + # Define bar updating wrapper + def wrapper(*args, **kwargs): + # update tbar correctly + # it seems `pandas apply` calls `func` twice + # on the first column/row to decide whether it can + # take a fast or slow code path; so stop when t.total==t.n + t.update(n=1 if not t.total or t.n < t.total else 0) + return func(*args, **kwargs) + + # Apply the provided function (in **kwargs) + # on the df using our wrapper (which provides bar updating) + try: + return getattr(df, df_function)(wrapper, **kwargs) + finally: + t.close() + + return inner + + # Monkeypatch pandas to provide easy methods + # Enable custom tqdm progress in pandas! + Series.progress_apply = inner_generator() + SeriesGroupBy.progress_apply = inner_generator() + Series.progress_map = inner_generator('map') + SeriesGroupBy.progress_map = inner_generator('map') + + DataFrame.progress_apply = inner_generator() + DataFrameGroupBy.progress_apply = inner_generator() + DataFrame.progress_applymap = inner_generator('applymap') + DataFrame.progress_map = inner_generator('map') + DataFrameGroupBy.progress_map = inner_generator('map') + + if Panel is not None: + Panel.progress_apply = inner_generator() + if PanelGroupBy is not None: + PanelGroupBy.progress_apply = inner_generator() + + GroupBy.progress_apply = inner_generator() + GroupBy.progress_aggregate = inner_generator('aggregate') + GroupBy.progress_transform = inner_generator('transform') + + if Rolling is not None and Expanding is not None: + Rolling.progress_apply = inner_generator() + Expanding.progress_apply = inner_generator() + elif _Rolling_and_Expanding is not None: + _Rolling_and_Expanding.progress_apply = inner_generator() + + # override defaults via env vars + @envwrap("tqdm", is_method=True, types={'total': float, 'ncols': int, 'miniters': float, + 'position': int, 'nrows': int}) + def __init__(self, iterable=None, desc=None, total=None, leave=True, file=None, + ncols=None, mininterval=0.1, maxinterval=10.0, miniters=None, + ascii=None, # pylint: disable=redefined-builtin + disable=False, unit='it', unit_scale=False, dynamic_ncols=False, smoothing=0.3, + bar_format=None, initial=0, position=None, postfix=None, unit_divisor=1000, + write_bytes=False, lock_args=None, nrows=None, colour=None, delay=0.0, gui=False, + **kwargs): + """see tqdm.tqdm for arguments""" + if file is None: + file = sys.stderr + + if write_bytes: + # Despite coercing unicode into bytes, py2 sys.std* streams + # should have bytes written to them. + file = SimpleTextIOWrapper( + file, encoding=getattr(file, 'encoding', None) or 'utf-8') + + file = DisableOnWriteError(file, tqdm_instance=self) + + if disable is None and hasattr(file, "isatty") and not file.isatty(): + disable = True + + if total is None and iterable is not None: + try: + total = len(iterable) + except (TypeError, AttributeError): + total = None + if total == float("inf"): + # Infinite iterations, behave same as unknown + total = None + + if disable: + self.iterable = iterable + self.disable = disable + with self._lock: + self.pos = self._get_free_pos(self) + self._instances.remove(self) + self.n = initial + self.total = total + self.leave = leave + return + + if kwargs: + self.disable = True + with self._lock: + self.pos = self._get_free_pos(self) + self._instances.remove(self) + raise ( + TqdmDeprecationWarning( + "`nested` is deprecated and automated.\n" + "Use `position` instead for manual control.\n", + fp_write=getattr(file, 'write', sys.stderr.write)) + if "nested" in kwargs else + TqdmKeyError("Unknown argument(s): " + str(kwargs))) + + # Preprocess the arguments + if ( + (ncols is None or nrows is None) and (file in (sys.stderr, sys.stdout)) + ) or dynamic_ncols: # pragma: no cover + if dynamic_ncols: + dynamic_ncols = _screen_shape_wrapper() + if dynamic_ncols: + ncols, nrows = dynamic_ncols(file) + else: + _dynamic_ncols = _screen_shape_wrapper() + if _dynamic_ncols: + _ncols, _nrows = _dynamic_ncols(file) + if ncols is None: + ncols = _ncols + if nrows is None: + nrows = _nrows + + if miniters is None: + miniters = 0 + dynamic_miniters = True + else: + dynamic_miniters = False + + if mininterval is None: + mininterval = 0 + + if maxinterval is None: + maxinterval = 0 + + if ascii is None: + ascii = not _supports_unicode(file) + + if bar_format and ascii is not True and not _is_ascii(ascii): + # Convert bar format into unicode since terminal uses unicode + bar_format = str(bar_format) + + if smoothing is None: + smoothing = 0 + + # Store the arguments + self.iterable = iterable + self.desc = desc or '' + self.total = total + self.leave = leave + self.fp = file + self.ncols = ncols + self.nrows = nrows + self.mininterval = mininterval + self.maxinterval = maxinterval + self.miniters = miniters + self.dynamic_miniters = dynamic_miniters + self.ascii = ascii + self.disable = disable + self.unit = unit + self.unit_scale = unit_scale + self.unit_divisor = unit_divisor + self.initial = initial + self.lock_args = lock_args + self.delay = delay + self.gui = gui + self.dynamic_ncols = dynamic_ncols + self.smoothing = smoothing + self._ema_dn = EMA(smoothing) + self._ema_dt = EMA(smoothing) + self._ema_miniters = EMA(smoothing) + self.bar_format = bar_format + self.postfix = None + self.colour = colour + self._time = time + if postfix: + try: + self.set_postfix(refresh=False, **postfix) + except TypeError: + self.postfix = postfix + + # Init the iterations counters + self.last_print_n = initial + self.n = initial + + # if nested, at initial sp() call we replace '\r' by '\n' to + # not overwrite the outer progress bar + with self._lock: + # mark fixed positions as negative + self.pos = self._get_free_pos(self) if position is None else -position + + if not gui: + # Initialize the screen printer + self.sp = self.status_printer(self.fp) + if delay <= 0: + self.refresh(lock_args=self.lock_args) + + # Init the time counter + self.last_print_t = self._time() + # NB: Avoid race conditions by setting start_t at the very end of init + self.start_t = self.last_print_t + + def __bool__(self): + if self.total is not None: + return self.total > 0 + if self.iterable is None: + raise TypeError('bool() undefined when iterable == total == None') + return bool(self.iterable) + + def __len__(self): + return ( + self.total if self.iterable is None + else self.iterable.shape[0] if hasattr(self.iterable, "shape") + else len(self.iterable) if hasattr(self.iterable, "__len__") + else self.iterable.__length_hint__() if hasattr(self.iterable, "__length_hint__") + else getattr(self, "total", None)) + + def __reversed__(self): + try: + orig = self.iterable + except AttributeError: + raise TypeError("'tqdm' object is not reversible") + else: + self.iterable = reversed(self.iterable) + return self.__iter__() + finally: + self.iterable = orig + + def __contains__(self, item): + contains = getattr(self.iterable, '__contains__', None) + return (contains(item) if contains is not None # pylint: disable=not-callable + else item in self.__iter__()) + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + try: + self.close() + except AttributeError: + # maybe eager thread cleanup upon external error + if (exc_type, exc_value, traceback) == (None, None, None): + raise + warn("AttributeError ignored", TqdmWarning, stacklevel=2) + + def __del__(self): + self.close() + + def __str__(self): + return self.format_meter(**self.format_dict) + + @property + def _comparable(self): + return abs(getattr(self, "pos", 1 << 31)) + + def __hash__(self): + return id(self) + + def __iter__(self): + """Backward-compatibility to use: for x in tqdm(iterable)""" + + # Inlining instance variables as locals (speed optimisation) + iterable = self.iterable + + # If the bar is disabled, then just walk the iterable + # (note: keep this check outside the loop for performance) + if self.disable: + for obj in iterable: + yield obj + return + + mininterval = self.mininterval + last_print_t = self.last_print_t + last_print_n = self.last_print_n + min_start_t = self.start_t + self.delay + n = self.n + time = self._time + + try: + for obj in iterable: + yield obj + # Update and possibly print the progress bar. + # Note: does not call self.update(1) for speed optimisation. + n += 1 + + if n - last_print_n >= self.miniters: + cur_t = time() + dt = cur_t - last_print_t + if dt >= mininterval and cur_t >= min_start_t: + self.update(n - last_print_n) + last_print_n = self.last_print_n + last_print_t = self.last_print_t + finally: + self.n = n + self.close() + + def update(self, n=1): + """ + Manually update the progress bar, useful for streams + such as reading files. + E.g.: + >>> t = tqdm(total=filesize) # Initialise + >>> for current_buffer in stream: + ... ... + ... t.update(len(current_buffer)) + >>> t.close() + The last line is highly recommended, but possibly not necessary if + `t.update()` will be called in such a way that `filesize` will be + exactly reached and printed. + + Parameters + ---------- + n : int or float, optional + Increment to add to the internal counter of iterations + [default: 1]. If using float, consider specifying `{n:.3f}` + or similar in `bar_format`, or specifying `unit_scale`. + + Returns + ------- + out : bool or None + True if a `display()` was triggered. + """ + if self.disable: + return + + if n < 0: + self.last_print_n += n # for auto-refresh logic to work + self.n += n + + # check counter first to reduce calls to time() + if self.n - self.last_print_n >= self.miniters: + cur_t = self._time() + dt = cur_t - self.last_print_t + if dt >= self.mininterval and cur_t >= self.start_t + self.delay: + cur_t = self._time() + dn = self.n - self.last_print_n # >= n + if self.smoothing and dt and dn: + # EMA (not just overall average) + self._ema_dn(dn) + self._ema_dt(dt) + self.refresh(lock_args=self.lock_args) + if self.dynamic_miniters: + # If no `miniters` was specified, adjust automatically to the + # maximum iteration rate seen so far between two prints. + # e.g.: After running `tqdm.update(5)`, subsequent + # calls to `tqdm.update()` will only cause an update after + # at least 5 more iterations. + if self.maxinterval and dt >= self.maxinterval: + self.miniters = dn * (self.mininterval or self.maxinterval) / dt + elif self.smoothing: + # EMA miniters update + self.miniters = self._ema_miniters( + dn * (self.mininterval / dt if self.mininterval and dt + else 1)) + else: + # max iters between two prints + self.miniters = max(self.miniters, dn) + + # Store old values for next call + self.last_print_n = self.n + self.last_print_t = cur_t + return True + + def close(self): + """Cleanup and (if leave=False) close the progress bar.""" + if self.disable: + return + + # Prevent multiple closures + self.disable = True + + # decrement instance pos and remove from internal set + pos = abs(self.pos) + self._decr_instances(self) + + if self.last_print_t < self.start_t + self.delay: + # haven't ever displayed; nothing to clear + return + + # GUI mode + if getattr(self, 'sp', None) is None: + return + + # annoyingly, _supports_unicode isn't good enough + def fp_write(s): + self.fp.write(str(s)) + + try: + fp_write('') + except ValueError as e: + if 'closed' in str(e): + return + raise # pragma: no cover + + leave = pos == 0 if self.leave is None else self.leave + + with self._lock: + if leave: + # stats for overall rate (no weighted average) + self._ema_dt = lambda: None + self.display(pos=0) + fp_write('\n') + else: + # clear previous display + if self.display(msg='', pos=pos) and not pos: + fp_write('\r') + + def clear(self, nolock=False): + """Clear current bar display.""" + if self.disable: + return + + if not nolock: + self._lock.acquire() + pos = abs(self.pos) + if pos < (self.nrows or 20): + self.moveto(pos) + self.sp('') + self.fp.write('\r') # place cursor back at the beginning of line + self.moveto(-pos) + if not nolock: + self._lock.release() + + def refresh(self, nolock=False, lock_args=None): + """ + Force refresh the display of this bar. + + Parameters + ---------- + nolock : bool, optional + If `True`, does not lock. + If [default: `False`]: calls `acquire()` on internal lock. + lock_args : tuple, optional + Passed to internal lock's `acquire()`. + If specified, will only `display()` if `acquire()` returns `True`. + """ + if self.disable: + return + + if not nolock: + if lock_args: + if not self._lock.acquire(*lock_args): + return False + else: + self._lock.acquire() + self.display() + if not nolock: + self._lock.release() + return True + + def unpause(self): + """Restart tqdm timer from last print time.""" + if self.disable: + return + cur_t = self._time() + self.start_t += cur_t - self.last_print_t + self.last_print_t = cur_t + + def reset(self, total=None): + """ + Resets to 0 iterations for repeated use. + + Consider combining with `leave=True`. + + Parameters + ---------- + total : int or float, optional. Total to use for the new bar. + """ + self.n = 0 + if total is not None: + self.total = total + if self.disable: + return + self.last_print_n = 0 + self.last_print_t = self.start_t = self._time() + self._ema_dn = EMA(self.smoothing) + self._ema_dt = EMA(self.smoothing) + self._ema_miniters = EMA(self.smoothing) + self.refresh() + + def set_description(self, desc=None, refresh=True): + """ + Set/modify description of the progress bar. + + Parameters + ---------- + desc : str, optional + refresh : bool, optional + Forces refresh [default: True]. + """ + self.desc = desc + ': ' if desc else '' + if refresh: + self.refresh() + + def set_description_str(self, desc=None, refresh=True): + """Set/modify description without ': ' appended.""" + self.desc = desc or '' + if refresh: + self.refresh() + + def set_postfix(self, ordered_dict=None, refresh=True, **kwargs): + """ + Set/modify postfix (additional stats) + with automatic formatting based on datatype. + + Parameters + ---------- + ordered_dict : dict or OrderedDict, optional + refresh : bool, optional + Forces refresh [default: True]. + kwargs : dict, optional + """ + # Sort in alphabetical order to be more deterministic + postfix = OrderedDict([] if ordered_dict is None else ordered_dict) + for key in sorted(kwargs.keys()): + postfix[key] = kwargs[key] + # Preprocess stats according to datatype + for key in postfix.keys(): + # Number: limit the length of the string + if isinstance(postfix[key], Number): + postfix[key] = self.format_num(postfix[key]) + # Else for any other type, try to get the string conversion + elif not isinstance(postfix[key], str): + postfix[key] = str(postfix[key]) + # Else if it's a string, don't need to preprocess anything + # Stitch together to get the final postfix + self.postfix = ', '.join(key + '=' + postfix[key].strip() + for key in postfix.keys()) + if refresh: + self.refresh() + + def set_postfix_str(self, s='', refresh=True): + """ + Postfix without dictionary expansion, similar to prefix handling. + """ + self.postfix = str(s) + if refresh: + self.refresh() + + def moveto(self, n): + # TODO: private method + self.fp.write('\n' * n + _term_move_up() * -n) + getattr(self.fp, 'flush', lambda: None)() + + @property + def format_dict(self): + """Public API for read-only member access.""" + if self.disable and not hasattr(self, 'unit'): + return defaultdict(lambda: None, { + 'n': self.n, 'total': self.total, 'elapsed': 0, 'unit': 'it'}) + if self.dynamic_ncols: + self.ncols, self.nrows = self.dynamic_ncols(self.fp) + return { + 'n': self.n, 'total': self.total, + 'elapsed': self._time() - self.start_t if hasattr(self, 'start_t') else 0, + 'ncols': self.ncols, 'nrows': self.nrows, 'prefix': self.desc, + 'ascii': self.ascii, 'unit': self.unit, 'unit_scale': self.unit_scale, + 'rate': self._ema_dn() / self._ema_dt() if self._ema_dt() else None, + 'bar_format': self.bar_format, 'postfix': self.postfix, + 'unit_divisor': self.unit_divisor, 'initial': self.initial, + 'colour': self.colour} + + def display(self, msg=None, pos=None): + """ + Use `self.sp` to display `msg` in the specified `pos`. + + Consider overloading this function when inheriting to use e.g.: + `self.some_frontend(**self.format_dict)` instead of `self.sp`. + + Parameters + ---------- + msg : str, optional. What to display (default: `repr(self)`). + pos : int, optional. Position to `moveto` + (default: `abs(self.pos)`). + """ + if pos is None: + pos = abs(self.pos) + + nrows = self.nrows or 20 + if pos >= nrows - 1: + if pos >= nrows: + return False + if msg or msg is None: # override at `nrows - 1` + msg = " ... (more hidden) ..." + + if not hasattr(self, "sp"): + raise TqdmDeprecationWarning( + "Please use `tqdm.gui.tqdm(...)`" + " instead of `tqdm(..., gui=True)`\n", + fp_write=getattr(self.fp, 'write', sys.stderr.write)) + + if pos: + self.moveto(pos) + self.sp(self.__str__() if msg is None else msg) + if pos: + self.moveto(-pos) + return True + + @classmethod + @contextmanager + def wrapattr(cls, stream, method, total=None, bytes=True, # pylint: disable=redefined-builtin + **tqdm_kwargs): + """ + stream : file-like object. + method : str, "read" or "write". The result of `read()` and + the first argument of `write()` should have a `len()`. + + >>> with tqdm.wrapattr(file_obj, "read", total=file_obj.size) as fobj: + ... while True: + ... chunk = fobj.read(chunk_size) + ... if not chunk: + ... break + """ + with cls(total=total, **tqdm_kwargs) as t: + if bytes: + t.unit = "B" + t.unit_scale = True + t.unit_divisor = 1024 + yield CallbackIOWrapper(t.update, stream, method) + + +def trange(*args, **kwargs): + """Shortcut for tqdm(range(*args), **kwargs).""" + return tqdm(range(*args), **kwargs) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/tk.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/tk.py new file mode 100644 index 0000000000000000000000000000000000000000..6607cc7f55e9bcdd177ef57dae4dbc06c7e33951 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/tk.py @@ -0,0 +1,196 @@ +""" +Tkinter GUI progress bar decorator for iterators. + +Usage: +>>> from tqdm.tk import trange, tqdm +>>> for i in trange(10): +... ... +""" +import re +import sys +import tkinter +import tkinter.ttk as ttk +from warnings import warn + +from .std import TqdmExperimentalWarning, TqdmWarning +from .std import tqdm as std_tqdm + +__author__ = {"github.com/": ["richardsheridan", "casperdcl"]} +__all__ = ['tqdm_tk', 'ttkrange', 'tqdm', 'trange'] + + +class tqdm_tk(std_tqdm): # pragma: no cover + """ + Experimental Tkinter GUI version of tqdm! + + Note: Window interactivity suffers if `tqdm_tk` is not running within + a Tkinter mainloop and values are generated infrequently. In this case, + consider calling `tqdm_tk.refresh()` frequently in the Tk thread. + """ + + # TODO: @classmethod: write()? + + def __init__(self, *args, **kwargs): + """ + This class accepts the following parameters *in addition* to + the parameters accepted by `tqdm`. + + Parameters + ---------- + grab : bool, optional + Grab the input across all windows of the process. + tk_parent : `tkinter.Wm`, optional + Parent Tk window. + cancel_callback : Callable, optional + Create a cancel button and set `cancel_callback` to be called + when the cancel or window close button is clicked. + """ + kwargs = kwargs.copy() + kwargs['gui'] = True + # convert disable = None to False + kwargs['disable'] = bool(kwargs.get('disable', False)) + self._warn_leave = 'leave' in kwargs + grab = kwargs.pop('grab', False) + tk_parent = kwargs.pop('tk_parent', None) + self._cancel_callback = kwargs.pop('cancel_callback', None) + super().__init__(*args, **kwargs) + + if self.disable: + return + + if tk_parent is None: # Discover parent widget + try: + tk_parent = tkinter._default_root + except AttributeError: + raise AttributeError( + "`tk_parent` required when using `tkinter.NoDefaultRoot()`") + if tk_parent is None: # use new default root window as display + self._tk_window = tkinter.Tk() + else: # some other windows already exist + self._tk_window = tkinter.Toplevel() + else: + self._tk_window = tkinter.Toplevel(tk_parent) + + warn("GUI is experimental/alpha", TqdmExperimentalWarning, stacklevel=2) + self._tk_dispatching = self._tk_dispatching_helper() + + self._tk_window.protocol("WM_DELETE_WINDOW", self.cancel) + self._tk_window.wm_title(self.desc) + self._tk_window.wm_attributes("-topmost", 1) + self._tk_window.after(0, lambda: self._tk_window.wm_attributes("-topmost", 0)) + self._tk_n_var = tkinter.DoubleVar(self._tk_window, value=0) + self._tk_text_var = tkinter.StringVar(self._tk_window) + pbar_frame = ttk.Frame(self._tk_window, padding=5) + pbar_frame.pack() + _tk_label = ttk.Label(pbar_frame, textvariable=self._tk_text_var, + wraplength=600, anchor="center", justify="center") + _tk_label.pack() + self._tk_pbar = ttk.Progressbar( + pbar_frame, variable=self._tk_n_var, length=450) + if self.total is not None: + self._tk_pbar.configure(maximum=self.total) + else: + self._tk_pbar.configure(mode="indeterminate") + self._tk_pbar.pack() + if self._cancel_callback is not None: + _tk_button = ttk.Button(pbar_frame, text="Cancel", command=self.cancel) + _tk_button.pack() + if grab: + self._tk_window.grab_set() + + def close(self): + if self.disable: + return + + self.disable = True + + with self.get_lock(): + self._instances.remove(self) + + def _close(): + self._tk_window.after('idle', self._tk_window.destroy) + if not self._tk_dispatching: + self._tk_window.update() + + self._tk_window.protocol("WM_DELETE_WINDOW", _close) + + # if leave is set but we are self-dispatching, the left window is + # totally unresponsive unless the user manually dispatches + if not self.leave: + _close() + elif not self._tk_dispatching: + if self._warn_leave: + warn("leave flag ignored if not in tkinter mainloop", + TqdmWarning, stacklevel=2) + _close() + + def clear(self, *_, **__): + pass + + def display(self, *_, **__): + self._tk_n_var.set(self.n) + d = self.format_dict + # remove {bar} + d['bar_format'] = (d['bar_format'] or "{l_bar}{r_bar}").replace( + "{bar}", "") + msg = self.format_meter(**d) + if '' in msg: + msg = "".join(re.split(r'\|?\|?', msg, maxsplit=1)) + self._tk_text_var.set(msg) + if not self._tk_dispatching: + self._tk_window.update() + + def set_description(self, desc=None, refresh=True): + self.set_description_str(desc, refresh) + + def set_description_str(self, desc=None, refresh=True): + self.desc = desc + if not self.disable: + self._tk_window.wm_title(desc) + if refresh and not self._tk_dispatching: + self._tk_window.update() + + def cancel(self): + """ + `cancel_callback()` followed by `close()` + when close/cancel buttons clicked. + """ + if self._cancel_callback is not None: + self._cancel_callback() + self.close() + + def reset(self, total=None): + """ + Resets to 0 iterations for repeated use. + + Parameters + ---------- + total : int or float, optional. Total to use for the new bar. + """ + if hasattr(self, '_tk_pbar'): + if total is None: + self._tk_pbar.configure(maximum=100, mode="indeterminate") + else: + self._tk_pbar.configure(maximum=total, mode="determinate") + super().reset(total=total) + + @staticmethod + def _tk_dispatching_helper(): + """determine if Tkinter mainloop is dispatching events""" + codes = {tkinter.mainloop.__code__, tkinter.Misc.mainloop.__code__} + for frame in sys._current_frames().values(): + while frame: + if frame.f_code in codes: + return True + frame = frame.f_back + return False + + +def ttkrange(*args, **kwargs): + """Shortcut for `tqdm.tk.tqdm(range(*args), **kwargs)`.""" + return tqdm_tk(range(*args), **kwargs) + + +# Aliases +tqdm = tqdm_tk +trange = ttkrange diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/tqdm.1 b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/tqdm.1 new file mode 100644 index 0000000000000000000000000000000000000000..b134ddd74c2095149e0c902f004a162c3be9d2ed --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/tqdm.1 @@ -0,0 +1,241 @@ +.TH "TQDM" "1" "2015\-2026" "tqdm User Manuals" +.SH NAME +tqdm \- fast, extensible progress bar for Python and CLI +.SH SYNOPSIS +tqdm [\f[I]options\f[R]] +.SH DESCRIPTION +See \c +.UR https://github.com/tqdm/tqdm +.UE \c +\&. +Can be used as a pipe: +.IP +.EX +$ \f[I]# count lines of code\f[R] +$ cat *.py \f[B]|\f[R] tqdm \f[B]|\f[R] wc \-l +327it [00:00, 981773.38it/s] +327 + +$ \f[I]# find all files\f[R] +$ find . \-name \(dq*.py\(dq \f[B]|\f[R] tqdm \f[B]|\f[R] wc \-l +432it [00:00, 833842.30it/s] +432 + +\f[I]# ... and more info\f[R] +$ find . \-name \(aq*.py\(aq \-exec wc \-l \(rs{} \(rs; \(rs + \f[B]|\f[R] tqdm \-\-total 432 \-\-unit files \-\-desc counting \(rs + \f[B]|\f[R] awk \(aq{ sum += $1 }; END { print sum }\(aq +counting: 100%\f[B]|\f[R]█████████\f[B]|\f[R] 432/432 [00:00<00:00, 794361.83files/s] +131998 +.EE +.SH OPTIONS +.TP +\-h, \-\-help +Print this help and exit. +.TP +\-v, \-\-version +Print version and exit. +.TP +\-\-desc=\f[I]desc\f[R] +str, optional. +Prefix for the progress bar. +.TP +\-\-total=\f[I]total\f[R] +int or float, optional. +The number of expected iterations. +If unspecified, len(iterable) is used if possible. +If float(\(lqinf\(rq) or as a last resort, only basic progress +statistics are displayed (no ETA, no progress bar). +If \f[CR]gui\f[R] is True and this parameter needs subsequent updating, +specify an initial arbitrary large positive number, e.g.\ 9e9. +.TP +\-\-leave +bool, optional. +If [default: True], keeps all traces of the progress bar upon +termination of iteration. +If \f[CR]None\f[R], will leave only if \f[CR]position\f[R] is +\f[CR]0\f[R]. +.TP +\-\-ncols=\f[I]ncols\f[R] +int, optional. +The width of the entire output message. +If specified, dynamically resizes the progress bar to stay within this +bound. +If unspecified, attempts to use environment width. +The fallback is a meter width of 10 and no limit for the counter and +statistics. +If 0, will not print any meter (only stats). +.TP +\-\-mininterval=\f[I]mininterval\f[R] +float, optional. +Minimum progress display update interval [default: 0.1] seconds. +.TP +\-\-maxinterval=\f[I]maxinterval\f[R] +float, optional. +Maximum progress display update interval [default: 10] seconds. +Automatically adjusts \f[CR]miniters\f[R] to correspond to +\f[CR]mininterval\f[R] after long display update lag. +Only works if \f[CR]dynamic_miniters\f[R] or monitor thread is enabled. +.TP +\-\-miniters=\f[I]miniters\f[R] +int or float, optional. +Minimum progress display update interval, in iterations. +If 0 and \f[CR]dynamic_miniters\f[R], will automatically adjust to equal +\f[CR]mininterval\f[R] (more CPU efficient, good for tight loops). +If > 0, will skip display of specified number of iterations. +Tweak this and \f[CR]mininterval\f[R] to get very efficient loops. +If your progress is erratic with both fast and slow iterations (network, +skipping items, etc) you should set miniters=1. +.TP +\-\-ascii=\f[I]ascii\f[R] +bool or str, optional. +If unspecified or False, use unicode (smooth blocks) to fill the meter. +The fallback is to use ASCII characters \(rq 123456789#\(lq. +.TP +\-\-disable +bool, optional. +Whether to disable the entire progress bar wrapper [default: False]. +If set to None, disable on non\-TTY. +.TP +\-\-unit=\f[I]unit\f[R] +str, optional. +String that will be used to define the unit of each iteration [default: +it]. +.TP +\-\-unit\-scale=\f[I]unit_scale\f[R] +bool or int or float, optional. +If 1 or True, the number of iterations will be reduced/scaled +automatically and a metric prefix following the International System of +Units standard will be added (kilo, mega, etc.) +[default: False]. +If any other non\-zero number, will scale \f[CR]total\f[R] and +\f[CR]n\f[R]. +.TP +\-\-dynamic\-ncols +bool, optional. +If set, constantly alters \f[CR]ncols\f[R] and \f[CR]nrows\f[R] to the +environment (allowing for window resizes) [default: False]. +.TP +\-\-smoothing=\f[I]smoothing\f[R] +float, optional. +Exponential moving average smoothing factor for speed estimates (ignored +in GUI mode). +Ranges from 0 (average speed) to 1 (current/instantaneous speed) +[default: 0.3]. +.TP +\-\-bar\-format=\f[I]bar_format\f[R] +str, optional. +Specify a custom bar string formatting. +May impact performance. +[default: `{l_bar}{bar}{r_bar}'], where l_bar=`{desc}: +{percentage:3.0f}%|' and r_bar=`| {n_fmt}/{total_fmt} +[{elapsed}<{remaining}, \(cq \(cq{rate_fmt}{postfix}]' Possible vars: +l_bar, bar, r_bar, n, n_fmt, total, total_fmt, percentage, elapsed, +elapsed_s, ncols, nrows, desc, unit, rate, rate_fmt, rate_noinv, +rate_noinv_fmt, rate_inv, rate_inv_fmt, postfix, unit_divisor, +remaining, remaining_s, eta. +Note that a trailing \(lq:\(rq is automatically removed after {desc} if +the latter is empty. +.TP +\-\-initial=\f[I]initial\f[R] +int or float, optional. +The initial counter value. +Useful when restarting a progress bar [default: 0]. +If using float, consider specifying \f[CR]{n:.3f}\f[R] or similar in +\f[CR]bar_format\f[R], or specifying \f[CR]unit_scale\f[R]. +.TP +\-\-position=\f[I]position\f[R] +int, optional. +Specify the line offset to print this bar (starting from 0) Automatic if +unspecified. +Useful to manage multiple bars at once (eg, from threads). +.TP +\-\-postfix=\f[I]postfix\f[R] +dict or *, optional. +Specify additional stats to display at the end of the bar. +Calls \f[CR]set_postfix(**postfix)\f[R] if possible (dict). +.TP +\-\-unit\-divisor=\f[I]unit_divisor\f[R] +float, optional. +[default: 1000], ignored unless \f[CR]unit_scale\f[R] is True. +.TP +\-\-write\-bytes +bool, optional. +Whether to write bytes. +If (default: False) will write unicode. +.TP +\-\-lock\-args=\f[I]lock_args\f[R] +tuple, optional. +Passed to \f[CR]refresh\f[R] for intermediate output (initialisation, +iterating, and updating). +.TP +\-\-nrows=\f[I]nrows\f[R] +int, optional. +The screen height. +If specified, hides nested bars outside this bound. +If unspecified, attempts to use environment height. +The fallback is 20. +.TP +\-\-colour=\f[I]colour\f[R] +str, optional. +Bar colour (e.g.\ `green', `#00ff00'). +.TP +\-\-delay=\f[I]delay\f[R] +float, optional. +Don\(cqt display until [default: 0] seconds have elapsed. +.TP +\-\-delim=\f[I]delim\f[R] +chr, optional. +Delimiting character [default: `\(rsn']. +Use `\(rs0' for null. +N.B.: on Windows systems, Python converts `\(rsn' to `\(rsr\(rsn'. +.TP +\-\-buf\-size=\f[I]buf_size\f[R] +int, optional. +String buffer size in bytes [default: 256] used when \f[CR]delim\f[R] is +specified. +.TP +\-\-bytes +bool, optional. +If true, will count bytes, ignore \f[CR]delim\f[R], and default +\f[CR]unit_scale\f[R] to True, \f[CR]unit_divisor\f[R] to 1024, and +\f[CR]unit\f[R] to `B'. +.TP +\-\-tee +bool, optional. +If true, passes \f[CR]stdin\f[R] to both \f[CR]stderr\f[R] and +\f[CR]stdout\f[R]. +.TP +\-\-update +bool, optional. +If true, will treat input as newly elapsed iterations, i.e.\ numbers to +pass to \f[CR]update()\f[R]. +Note that this is slow (\(ti2e5 it/s) since every input must be decoded +as a number. +.TP +\-\-update\-to +bool, optional. +If true, will treat input as total elapsed iterations, i.e.\ numbers to +assign to \f[CR]self.n\f[R]. +Note that this is slow (\(ti2e5 it/s) since every input must be decoded +as a number. +.TP +\-\-null +bool, optional. +If true, will discard input (no stdout). +.TP +\-\-manpath=\f[I]manpath\f[R] +str, optional. +Directory in which to install tqdm man pages. +.TP +\-\-comppath=\f[I]comppath\f[R] +str, optional. +Directory in which to place tqdm completion. +.TP +\-\-log=\f[I]log\f[R] +str, optional. +CRITICAL|FATAL|ERROR|WARN(ING)|[default: `INFO']|DEBUG|NOTSET. +.SH AUTHORS +tqdm developers \c +.UR https://github.com/tqdm +.UE \c. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/utils.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..329c064478e4c6f094249d85732d818fc400720b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/utils.py @@ -0,0 +1,328 @@ +""" +General helpers required for `tqdm.std`. +""" +import os +import re +import sys +from functools import partial, partialmethod, wraps +from inspect import signature +# TODO consider using wcswidth third-party package for 0-width characters +from unicodedata import east_asian_width +from warnings import warn +from weakref import proxy + +_range, _unich, _unicode, _basestring = range, chr, str, str +CUR_OS = sys.platform +IS_WIN = any(CUR_OS.startswith(i) for i in ['win32', 'cygwin']) +IS_NIX = any(CUR_OS.startswith(i) for i in ['aix', 'linux', 'darwin', 'freebsd']) +RE_ANSI = re.compile(r"\x1b\[[;\d]*[A-Za-z]") + +try: + if IS_WIN: + import colorama + else: + raise ImportError +except ImportError: + colorama = None +else: + try: + colorama.init(strip=False) + except TypeError: + colorama.init() + + +def envwrap(name, app="", types=None, is_method=False): + """ + Basic (env-only) version of [envwrap](https://github.com/tqdm/envwrap). + Install `envwrap` for config file support. + """ + if types is None: + types = {} + if name[-1] == "_": + name = name[:-1] + warn("Trailing underscore in `name` is automatic", DeprecationWarning, stacklevel=2) + prefixes = (name, f"{name}_{app}") if app else (name,) + env_overrides = {} + for prefix in prefixes: + prefix = prefix.upper() + "_" + i = len(prefix) + env_overrides.update( + (k[i:].lower(), v) for k, v in os.environ.items() if k.startswith(prefix)) + part = partialmethod if is_method else partial + + def wrap(func): + params = signature(func).parameters + # ignore unknown env vars + overrides = {k: v for k, v in env_overrides.items() if k in params} + # infer overrides' `type`s + for k in overrides: + param = params[k] + if param.annotation is not param.empty: # typehints + for typ in getattr(param.annotation, '__args__', (param.annotation,)): + try: + overrides[k] = typ(overrides[k]) + except Exception: # nosec B110 + pass + else: + break + elif param.default is not None: # type of default value + overrides[k] = type(param.default)(overrides[k]) + else: + try: # `types` fallback + overrides[k] = types[k](overrides[k]) + except KeyError: # keep unconverted (`str`) + pass + return part(func, **overrides) + return wrap + + +try: + from envwrap import envwrap # noqa: F401, F811, pylint: disable=unused-import +except ModuleNotFoundError: + pass + + +class FormatReplace: + """ + >>> a = FormatReplace('something') + >>> f"{a:5d}" + 'something' + """ # NOQA: P102 + def __init__(self, replace=''): + self.replace = replace + self.format_called = 0 + + def __format__(self, _): + self.format_called += 1 + return self.replace + + +class Comparable: + """Assumes child has self._comparable attr/@property""" + def __lt__(self, other): + return self._comparable < other._comparable + + def __le__(self, other): + return (self < other) or (self == other) + + def __eq__(self, other): + return self._comparable == other._comparable + + def __ne__(self, other): + return not self == other + + def __gt__(self, other): + return not self <= other + + def __ge__(self, other): + return not self < other + + +class ObjectWrapper: + def __getattr__(self, name): + return getattr(self._wrapped, name) + + def __setattr__(self, name, value): + return setattr(self._wrapped, name, value) + + def wrapper_getattr(self, name): + """Actual `self.getattr` rather than self._wrapped.getattr""" + try: + return object.__getattr__(self, name) + except AttributeError: # py2 + return getattr(self, name) + + def wrapper_setattr(self, name, value): + """Actual `self.setattr` rather than self._wrapped.setattr""" + return object.__setattr__(self, name, value) + + def __init__(self, wrapped): + """ + Thin wrapper around a given object + """ + self.wrapper_setattr('_wrapped', wrapped) + + +class SimpleTextIOWrapper(ObjectWrapper): + """ + Change only `.write()` of the wrapped object by encoding the passed + value and passing the result to the wrapped object's `.write()` method. + """ + # pylint: disable=too-few-public-methods + def __init__(self, wrapped, encoding): + super().__init__(wrapped) + self.wrapper_setattr('encoding', encoding) + + def write(self, s): + """ + Encode `s` and pass to the wrapped object's `.write()` method. + """ + return self._wrapped.write(s.encode(self.wrapper_getattr('encoding'))) + + def __eq__(self, other): + return self._wrapped == getattr(other, '_wrapped', other) + + +class DisableOnWriteError(ObjectWrapper): + """ + Disable the given `tqdm_instance` upon `write()` or `flush()` errors. + """ + @staticmethod + def disable_on_exception(tqdm_instance, func): + """ + Quietly set `tqdm_instance.miniters=inf` if `func` raises `errno=5`. + """ + tqdm_instance = proxy(tqdm_instance) + + def inner(*args, **kwargs): + try: + return func(*args, **kwargs) + except OSError as e: + if e.errno != 5: + raise + try: + tqdm_instance.miniters = float('inf') + except ReferenceError: + pass + except ValueError as e: + if 'closed' not in str(e): + raise + try: + tqdm_instance.miniters = float('inf') + except ReferenceError: + pass + return inner + + def __init__(self, wrapped, tqdm_instance): # noqa: B042 + super().__init__(wrapped) + if hasattr(wrapped, 'write'): + self.wrapper_setattr( + 'write', self.disable_on_exception(tqdm_instance, wrapped.write)) + if hasattr(wrapped, 'flush'): + self.wrapper_setattr( + 'flush', self.disable_on_exception(tqdm_instance, wrapped.flush)) + + def __eq__(self, other): + return self._wrapped == getattr(other, '_wrapped', other) + + +class CallbackIOWrapper(ObjectWrapper): + def __init__(self, callback, stream, method="read"): + """ + Wrap a given `file`-like object's `read()` or `write()` to report + lengths to the given `callback` + """ + super().__init__(stream) + func = getattr(stream, method) + if method == "write": + @wraps(func) + def write(data, *args, **kwargs): + res = func(data, *args, **kwargs) + callback(len(data)) + return res + self.wrapper_setattr('write', write) + elif method == "read": + @wraps(func) + def read(*args, **kwargs): + data = func(*args, **kwargs) + callback(len(data)) + return data + self.wrapper_setattr('read', read) + else: + raise KeyError("Can only wrap read/write methods") + + +def _is_utf(encoding): + try: + '\u2588\u2589'.encode(encoding) + except UnicodeEncodeError: + return False + except Exception: + try: + return encoding.lower().startswith('utf-') or ('U8' == encoding) + except Exception: + return False + else: + return True + + +def _supports_unicode(fp): + try: + return _is_utf(fp.encoding) + except AttributeError: + return False + + +def _is_ascii(s): + if isinstance(s, str): + for c in s: + if ord(c) > 255: + return False + return True + return _supports_unicode(s) + + +def _screen_shape_wrapper(): # pragma: no cover + """ + Return a function which returns console dimensions (width, height). + Supported: linux, osx, windows, cygwin. + """ + def inner(fp): + try: + from os import get_terminal_size + cols, lines = get_terminal_size(getattr(fp, 'fileno', lambda: None)()) + return cols - 1, lines - 1 + except Exception: + return None, None + + return inner + + +def _environ_cols_wrapper(): # pragma: no cover + """ + Return a function which returns console width. + Supported: linux, osx, windows, cygwin. + """ + warn("Use `_screen_shape_wrapper()(file)[0]` instead of" + " `_environ_cols_wrapper()(file)`", DeprecationWarning, stacklevel=2) + shape = _screen_shape_wrapper() + if not shape: + return None + + @wraps(shape) + def inner(fp): + return shape(fp)[0] + + return inner + + +def _term_move_up(): # pragma: no cover + return '' if (os.name == 'nt') and (colorama is None) else '\x1b[A' + + +def _text_width(s): + return sum(2 if east_asian_width(ch) in 'FW' else 1 for ch in str(s)) + + +def disp_len(data): + """ + Returns the real on-screen length of a string which may contain + ANSI control codes and wide chars. + """ + return _text_width(RE_ANSI.sub('', data)) + + +def disp_trim(data, length): + """ + Trim a string which may contain ANSI control characters. + """ + if len(data) == disp_len(data): + return data[:length] + + ansi_present = bool(RE_ANSI.search(data)) + while disp_len(data) > length: # carefully delete one char at a time + data = data[:-1] + if ansi_present and bool(RE_ANSI.search(data)): + # assume ANSI reset is required + return data if data.endswith("\033[0m") else data + "\033[0m" + return data diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/version.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/version.py new file mode 100644 index 0000000000000000000000000000000000000000..84739ef244b0504c9705b40b72f7f229584b24ca --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/tqdm/version.py @@ -0,0 +1,7 @@ +"""`tqdm` version detector. Precedence: installed dist, git, 'UNKNOWN'.""" +from importlib.metadata import PackageNotFoundError, version + +try: + __version__ = version('tqdm') +except PackageNotFoundError: + __version__ = "UNKNOWN" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/INSTALLER b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/METADATA b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..e4e4de4130f9470a3e118075ad8c3fb64e482b08 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/METADATA @@ -0,0 +1,72 @@ +Metadata-Version: 2.2 +Name: typing_extensions +Version: 4.15.0+computecanada +Summary: Backported and Experimental Type Hints for Python 3.9+ +Description-Content-Type: text/markdown +Keywords: annotations,backport,checker,checking,function,hinting,hints,type,typechecking,typehinting,typehints,typing +Classifier: Development Status :: 5 - Production/Stable +Classifier: Environment :: Console +Classifier: Intended Audience :: Developers +Classifier: Operating System :: OS Independent +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Classifier: Topic :: Software Development +Author-email: "Guido van Rossum, Jukka Lehtosalo, Łukasz Langa, Michael Lee" +Project-URL: Bug Tracker, https://github.com/python/typing_extensions/issues +Project-URL: Changes, https://github.com/python/typing_extensions/blob/main/CHANGELOG.md +Project-URL: Documentation, https://typing-extensions.readthedocs.io/ +Project-URL: Home, https://github.com/python/typing_extensions +Project-URL: Q & A, https://github.com/python/typing/discussions +Project-URL: Repository, https://github.com/python/typing_extensions +Requires-Python: >=3.9 +License-File: LICENSE +License-Expression: PSF-2.0 + +# Typing Extensions + +[![Chat at https://gitter.im/python/typing](https://badges.gitter.im/python/typing.svg)](https://gitter.im/python/typing) + +[Documentation](https://typing-extensions.readthedocs.io/en/latest/#) – +[PyPI](https://pypi.org/project/typing-extensions/) + +## Overview + +The `typing_extensions` module serves two related purposes: + +- Enable use of new type system features on older Python versions. For example, + `typing.TypeGuard` is new in Python 3.10, but `typing_extensions` allows + users on previous Python versions to use it too. +- Enable experimentation with new type system PEPs before they are accepted and + added to the `typing` module. + +`typing_extensions` is treated specially by static type checkers such as +mypy and pyright. Objects defined in `typing_extensions` are treated the same +way as equivalent forms in `typing`. + +`typing_extensions` uses +[Semantic Versioning](https://semver.org/). The +major version will be incremented only for backwards-incompatible changes. +Therefore, it's safe to depend +on `typing_extensions` like this: `typing_extensions ~=x.y`, +where `x.y` is the first version that includes all features you need. +[This](https://packaging.python.org/en/latest/specifications/version-specifiers/#compatible-release) +is equivalent to `typing_extensions >=x.y, <(x+1)`. Do not depend on `~= x.y.z` +unless you really know what you're doing; that defeats the purpose of +semantic versioning. + +## Included items + +See [the documentation](https://typing-extensions.readthedocs.io/en/latest/#) for a +complete listing of module contents. + +## Contributing + +See [CONTRIBUTING.md](https://github.com/python/typing_extensions/blob/main/CONTRIBUTING.md) +for how to contribute to `typing_extensions`. + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/RECORD b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..c61c855afdcaad4bd06324353fda0ce26550f355 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/RECORD @@ -0,0 +1,7 @@ +__pycache__/typing_extensions.cpython-311.pyc,, +typing_extensions-4.15.0+computecanada.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +typing_extensions-4.15.0+computecanada.dist-info/METADATA,sha256=3v-L4fagxmgPPGdJ6e2ga1E9sy_PU0xi0p727_-Fq_0,3273 +typing_extensions-4.15.0+computecanada.dist-info/RECORD,, +typing_extensions-4.15.0+computecanada.dist-info/WHEEL,sha256=E4Ta9GSW8Vg2e11uZktRc054ObD9JXJ5hm6vpxg0WJE,87 +typing_extensions-4.15.0+computecanada.dist-info/licenses/LICENSE,sha256=Oy-B_iHRgcSZxZolbI4ZaEVdZonSaaqFNzv7avQdo78,13936 +typing_extensions.py,sha256=Qz0R0XDTok0usGXrwb_oSM6n49fOaFZ6tSvqLUwvftg,160429 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/WHEEL b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..b270261c9bb44553b2f648368d1bf5cb7baf8416 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: wheelfile 0.0.8 +Root-Is-Purelib: true +Tag: py3-none-any + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/licenses/LICENSE b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..f26bcf4d2de6eb136e31006ca3ab447d5e488adf --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions-4.15.0+computecanada.dist-info/licenses/LICENSE @@ -0,0 +1,279 @@ +A. HISTORY OF THE SOFTWARE +========================== + +Python was created in the early 1990s by Guido van Rossum at Stichting +Mathematisch Centrum (CWI, see https://www.cwi.nl) in the Netherlands +as a successor of a language called ABC. Guido remains Python's +principal author, although it includes many contributions from others. + +In 1995, Guido continued his work on Python at the Corporation for +National Research Initiatives (CNRI, see https://www.cnri.reston.va.us) +in Reston, Virginia where he released several versions of the +software. + +In May 2000, Guido and the Python core development team moved to +BeOpen.com to form the BeOpen PythonLabs team. In October of the same +year, the PythonLabs team moved to Digital Creations, which became +Zope Corporation. In 2001, the Python Software Foundation (PSF, see +https://www.python.org/psf/) was formed, a non-profit organization +created specifically to own Python-related Intellectual Property. +Zope Corporation was a sponsoring member of the PSF. + +All Python releases are Open Source (see https://opensource.org for +the Open Source Definition). Historically, most, but not all, Python +releases have also been GPL-compatible; the table below summarizes +the various releases. + + Release Derived Year Owner GPL- + from compatible? (1) + + 0.9.0 thru 1.2 1991-1995 CWI yes + 1.3 thru 1.5.2 1.2 1995-1999 CNRI yes + 1.6 1.5.2 2000 CNRI no + 2.0 1.6 2000 BeOpen.com no + 1.6.1 1.6 2001 CNRI yes (2) + 2.1 2.0+1.6.1 2001 PSF no + 2.0.1 2.0+1.6.1 2001 PSF yes + 2.1.1 2.1+2.0.1 2001 PSF yes + 2.1.2 2.1.1 2002 PSF yes + 2.1.3 2.1.2 2002 PSF yes + 2.2 and above 2.1.1 2001-now PSF yes + +Footnotes: + +(1) GPL-compatible doesn't mean that we're distributing Python under + the GPL. All Python licenses, unlike the GPL, let you distribute + a modified version without making your changes open source. The + GPL-compatible licenses make it possible to combine Python with + other software that is released under the GPL; the others don't. + +(2) According to Richard Stallman, 1.6.1 is not GPL-compatible, + because its license has a choice of law clause. According to + CNRI, however, Stallman's lawyer has told CNRI's lawyer that 1.6.1 + is "not incompatible" with the GPL. + +Thanks to the many outside volunteers who have worked under Guido's +direction to make these releases possible. + + +B. TERMS AND CONDITIONS FOR ACCESSING OR OTHERWISE USING PYTHON +=============================================================== + +Python software and documentation are licensed under the +Python Software Foundation License Version 2. + +Starting with Python 3.8.6, examples, recipes, and other code in +the documentation are dual licensed under the PSF License Version 2 +and the Zero-Clause BSD license. + +Some software incorporated into Python is under different licenses. +The licenses are listed with code falling under that license. + + +PYTHON SOFTWARE FOUNDATION LICENSE VERSION 2 +-------------------------------------------- + +1. This LICENSE AGREEMENT is between the Python Software Foundation +("PSF"), and the Individual or Organization ("Licensee") accessing and +otherwise using this software ("Python") in source or binary form and +its associated documentation. + +2. Subject to the terms and conditions of this License Agreement, PSF hereby +grants Licensee a nonexclusive, royalty-free, world-wide license to reproduce, +analyze, test, perform and/or display publicly, prepare derivative works, +distribute, and otherwise use Python alone or in any derivative version, +provided, however, that PSF's License Agreement and PSF's notice of copyright, +i.e., "Copyright (c) 2001, 2002, 2003, 2004, 2005, 2006, 2007, 2008, 2009, 2010, +2011, 2012, 2013, 2014, 2015, 2016, 2017, 2018, 2019, 2020, 2021, 2022, 2023 Python Software Foundation; +All Rights Reserved" are retained in Python alone or in any derivative version +prepared by Licensee. + +3. In the event Licensee prepares a derivative work that is based on +or incorporates Python or any part thereof, and wants to make +the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to Python. + +4. PSF is making Python available to Licensee on an "AS IS" +basis. PSF MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, PSF MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +5. PSF SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON +FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS +A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON, +OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. Nothing in this License Agreement shall be deemed to create any +relationship of agency, partnership, or joint venture between PSF and +Licensee. This License Agreement does not grant permission to use PSF +trademarks or trade name in a trademark sense to endorse or promote +products or services of Licensee, or any third party. + +8. By copying, installing or otherwise using Python, Licensee +agrees to be bound by the terms and conditions of this License +Agreement. + + +BEOPEN.COM LICENSE AGREEMENT FOR PYTHON 2.0 +------------------------------------------- + +BEOPEN PYTHON OPEN SOURCE LICENSE AGREEMENT VERSION 1 + +1. This LICENSE AGREEMENT is between BeOpen.com ("BeOpen"), having an +office at 160 Saratoga Avenue, Santa Clara, CA 95051, and the +Individual or Organization ("Licensee") accessing and otherwise using +this software in source or binary form and its associated +documentation ("the Software"). + +2. Subject to the terms and conditions of this BeOpen Python License +Agreement, BeOpen hereby grants Licensee a non-exclusive, +royalty-free, world-wide license to reproduce, analyze, test, perform +and/or display publicly, prepare derivative works, distribute, and +otherwise use the Software alone or in any derivative version, +provided, however, that the BeOpen Python License is retained in the +Software, alone or in any derivative version prepared by Licensee. + +3. BeOpen is making the Software available to Licensee on an "AS IS" +basis. BEOPEN MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, BEOPEN MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF THE SOFTWARE WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +4. BEOPEN SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF THE +SOFTWARE FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS +AS A RESULT OF USING, MODIFYING OR DISTRIBUTING THE SOFTWARE, OR ANY +DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +5. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +6. This License Agreement shall be governed by and interpreted in all +respects by the law of the State of California, excluding conflict of +law provisions. Nothing in this License Agreement shall be deemed to +create any relationship of agency, partnership, or joint venture +between BeOpen and Licensee. This License Agreement does not grant +permission to use BeOpen trademarks or trade names in a trademark +sense to endorse or promote products or services of Licensee, or any +third party. As an exception, the "BeOpen Python" logos available at +http://www.pythonlabs.com/logos.html may be used according to the +permissions granted on that web page. + +7. By copying, installing or otherwise using the software, Licensee +agrees to be bound by the terms and conditions of this License +Agreement. + + +CNRI LICENSE AGREEMENT FOR PYTHON 1.6.1 +--------------------------------------- + +1. This LICENSE AGREEMENT is between the Corporation for National +Research Initiatives, having an office at 1895 Preston White Drive, +Reston, VA 20191 ("CNRI"), and the Individual or Organization +("Licensee") accessing and otherwise using Python 1.6.1 software in +source or binary form and its associated documentation. + +2. Subject to the terms and conditions of this License Agreement, CNRI +hereby grants Licensee a nonexclusive, royalty-free, world-wide +license to reproduce, analyze, test, perform and/or display publicly, +prepare derivative works, distribute, and otherwise use Python 1.6.1 +alone or in any derivative version, provided, however, that CNRI's +License Agreement and CNRI's notice of copyright, i.e., "Copyright (c) +1995-2001 Corporation for National Research Initiatives; All Rights +Reserved" are retained in Python 1.6.1 alone or in any derivative +version prepared by Licensee. Alternately, in lieu of CNRI's License +Agreement, Licensee may substitute the following text (omitting the +quotes): "Python 1.6.1 is made available subject to the terms and +conditions in CNRI's License Agreement. This Agreement together with +Python 1.6.1 may be located on the internet using the following +unique, persistent identifier (known as a handle): 1895.22/1013. This +Agreement may also be obtained from a proxy server on the internet +using the following URL: http://hdl.handle.net/1895.22/1013". + +3. In the event Licensee prepares a derivative work that is based on +or incorporates Python 1.6.1 or any part thereof, and wants to make +the derivative work available to others as provided herein, then +Licensee hereby agrees to include in any such work a brief summary of +the changes made to Python 1.6.1. + +4. CNRI is making Python 1.6.1 available to Licensee on an "AS IS" +basis. CNRI MAKES NO REPRESENTATIONS OR WARRANTIES, EXPRESS OR +IMPLIED. BY WAY OF EXAMPLE, BUT NOT LIMITATION, CNRI MAKES NO AND +DISCLAIMS ANY REPRESENTATION OR WARRANTY OF MERCHANTABILITY OR FITNESS +FOR ANY PARTICULAR PURPOSE OR THAT THE USE OF PYTHON 1.6.1 WILL NOT +INFRINGE ANY THIRD PARTY RIGHTS. + +5. CNRI SHALL NOT BE LIABLE TO LICENSEE OR ANY OTHER USERS OF PYTHON +1.6.1 FOR ANY INCIDENTAL, SPECIAL, OR CONSEQUENTIAL DAMAGES OR LOSS AS +A RESULT OF MODIFYING, DISTRIBUTING, OR OTHERWISE USING PYTHON 1.6.1, +OR ANY DERIVATIVE THEREOF, EVEN IF ADVISED OF THE POSSIBILITY THEREOF. + +6. This License Agreement will automatically terminate upon a material +breach of its terms and conditions. + +7. This License Agreement shall be governed by the federal +intellectual property law of the United States, including without +limitation the federal copyright law, and, to the extent such +U.S. federal law does not apply, by the law of the Commonwealth of +Virginia, excluding Virginia's conflict of law provisions. +Notwithstanding the foregoing, with regard to derivative works based +on Python 1.6.1 that incorporate non-separable material that was +previously distributed under the GNU General Public License (GPL), the +law of the Commonwealth of Virginia shall govern this License +Agreement only as to issues arising under or with respect to +Paragraphs 4, 5, and 7 of this License Agreement. Nothing in this +License Agreement shall be deemed to create any relationship of +agency, partnership, or joint venture between CNRI and Licensee. This +License Agreement does not grant permission to use CNRI trademarks or +trade name in a trademark sense to endorse or promote products or +services of Licensee, or any third party. + +8. By clicking on the "ACCEPT" button where indicated, or by copying, +installing or otherwise using Python 1.6.1, Licensee agrees to be +bound by the terms and conditions of this License Agreement. + + ACCEPT + + +CWI LICENSE AGREEMENT FOR PYTHON 0.9.0 THROUGH 1.2 +-------------------------------------------------- + +Copyright (c) 1991 - 1995, Stichting Mathematisch Centrum Amsterdam, +The Netherlands. All rights reserved. + +Permission to use, copy, modify, and distribute this software and its +documentation for any purpose and without fee is hereby granted, +provided that the above copyright notice appear in all copies and that +both that copyright notice and this permission notice appear in +supporting documentation, and that the name of Stichting Mathematisch +Centrum or CWI not be used in advertising or publicity pertaining to +distribution of the software without specific, written prior +permission. + +STICHTING MATHEMATISCH CENTRUM DISCLAIMS ALL WARRANTIES WITH REGARD TO +THIS SOFTWARE, INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY AND +FITNESS, IN NO EVENT SHALL STICHTING MATHEMATISCH CENTRUM BE LIABLE +FOR ANY SPECIAL, INDIRECT OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT +OF OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + +ZERO-CLAUSE BSD LICENSE FOR CODE IN THE PYTHON DOCUMENTATION +---------------------------------------------------------------------- + +Permission to use, copy, modify, and/or distribute this software for any +purpose with or without fee is hereby granted. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH +REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY +AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT, +INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM +LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR +OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR +PERFORMANCE OF THIS SOFTWARE. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions.py new file mode 100644 index 0000000000000000000000000000000000000000..77f33e1614fd7d46ccd66b394d6c5d663bf8a8c6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_extensions.py @@ -0,0 +1,4317 @@ +import abc +import builtins +import collections +import collections.abc +import contextlib +import enum +import functools +import inspect +import io +import keyword +import operator +import sys +import types as _types +import typing +import warnings + +# Breakpoint: https://github.com/python/cpython/pull/119891 +if sys.version_info >= (3, 14): + import annotationlib + +__all__ = [ + # Super-special typing primitives. + 'Any', + 'ClassVar', + 'Concatenate', + 'Final', + 'LiteralString', + 'ParamSpec', + 'ParamSpecArgs', + 'ParamSpecKwargs', + 'Self', + 'Type', + 'TypeVar', + 'TypeVarTuple', + 'Unpack', + + # ABCs (from collections.abc). + 'Awaitable', + 'AsyncIterator', + 'AsyncIterable', + 'Coroutine', + 'AsyncGenerator', + 'AsyncContextManager', + 'Buffer', + 'ChainMap', + + # Concrete collection types. + 'ContextManager', + 'Counter', + 'Deque', + 'DefaultDict', + 'NamedTuple', + 'OrderedDict', + 'TypedDict', + + # Structural checks, a.k.a. protocols. + 'SupportsAbs', + 'SupportsBytes', + 'SupportsComplex', + 'SupportsFloat', + 'SupportsIndex', + 'SupportsInt', + 'SupportsRound', + 'Reader', + 'Writer', + + # One-off things. + 'Annotated', + 'assert_never', + 'assert_type', + 'clear_overloads', + 'dataclass_transform', + 'deprecated', + 'disjoint_base', + 'Doc', + 'evaluate_forward_ref', + 'get_overloads', + 'final', + 'Format', + 'get_annotations', + 'get_args', + 'get_origin', + 'get_original_bases', + 'get_protocol_members', + 'get_type_hints', + 'IntVar', + 'is_protocol', + 'is_typeddict', + 'Literal', + 'NewType', + 'overload', + 'override', + 'Protocol', + 'Sentinel', + 'reveal_type', + 'runtime', + 'runtime_checkable', + 'Text', + 'TypeAlias', + 'TypeAliasType', + 'TypeForm', + 'TypeGuard', + 'TypeIs', + 'TYPE_CHECKING', + 'type_repr', + 'Never', + 'NoReturn', + 'ReadOnly', + 'Required', + 'NotRequired', + 'NoDefault', + 'NoExtraItems', + + # Pure aliases, have always been in typing + 'AbstractSet', + 'AnyStr', + 'BinaryIO', + 'Callable', + 'Collection', + 'Container', + 'Dict', + 'ForwardRef', + 'FrozenSet', + 'Generator', + 'Generic', + 'Hashable', + 'IO', + 'ItemsView', + 'Iterable', + 'Iterator', + 'KeysView', + 'List', + 'Mapping', + 'MappingView', + 'Match', + 'MutableMapping', + 'MutableSequence', + 'MutableSet', + 'Optional', + 'Pattern', + 'Reversible', + 'Sequence', + 'Set', + 'Sized', + 'TextIO', + 'Tuple', + 'Union', + 'ValuesView', + 'cast', + 'no_type_check', + 'no_type_check_decorator', +] + +# for backward compatibility +PEP_560 = True +GenericMeta = type +# Breakpoint: https://github.com/python/cpython/pull/116129 +_PEP_696_IMPLEMENTED = sys.version_info >= (3, 13, 0, "beta") + +# Added with bpo-45166 to 3.10.1+ and some 3.9 versions +_FORWARD_REF_HAS_CLASS = "__forward_is_class__" in typing.ForwardRef.__slots__ + +# The functions below are modified copies of typing internal helpers. +# They are needed by _ProtocolMeta and they provide support for PEP 646. + + +class _Sentinel: + def __repr__(self): + return "" + + +_marker = _Sentinel() + + +# Breakpoint: https://github.com/python/cpython/pull/27342 +if sys.version_info >= (3, 10): + def _should_collect_from_parameters(t): + return isinstance( + t, (typing._GenericAlias, _types.GenericAlias, _types.UnionType) + ) +else: + def _should_collect_from_parameters(t): + return isinstance(t, (typing._GenericAlias, _types.GenericAlias)) + + +NoReturn = typing.NoReturn + +# Some unconstrained type variables. These are used by the container types. +# (These are not for export.) +T = typing.TypeVar('T') # Any type. +KT = typing.TypeVar('KT') # Key type. +VT = typing.TypeVar('VT') # Value type. +T_co = typing.TypeVar('T_co', covariant=True) # Any type covariant containers. +T_contra = typing.TypeVar('T_contra', contravariant=True) # Ditto contravariant. + + +# Breakpoint: https://github.com/python/cpython/pull/31841 +if sys.version_info >= (3, 11): + from typing import Any +else: + + class _AnyMeta(type): + def __instancecheck__(self, obj): + if self is Any: + raise TypeError("typing_extensions.Any cannot be used with isinstance()") + return super().__instancecheck__(obj) + + def __repr__(self): + if self is Any: + return "typing_extensions.Any" + return super().__repr__() + + class Any(metaclass=_AnyMeta): + """Special type indicating an unconstrained type. + - Any is compatible with every type. + - Any assumed to have all methods. + - All values assumed to be instances of Any. + Note that all the above statements are true from the point of view of + static type checkers. At runtime, Any should not be used with instance + checks. + """ + def __new__(cls, *args, **kwargs): + if cls is Any: + raise TypeError("Any cannot be instantiated") + return super().__new__(cls, *args, **kwargs) + + +ClassVar = typing.ClassVar + +# Vendored from cpython typing._SpecialFrom +# Having a separate class means that instances will not be rejected by +# typing._type_check. +class _SpecialForm(typing._Final, _root=True): + __slots__ = ('_name', '__doc__', '_getitem') + + def __init__(self, getitem): + self._getitem = getitem + self._name = getitem.__name__ + self.__doc__ = getitem.__doc__ + + def __getattr__(self, item): + if item in {'__name__', '__qualname__'}: + return self._name + + raise AttributeError(item) + + def __mro_entries__(self, bases): + raise TypeError(f"Cannot subclass {self!r}") + + def __repr__(self): + return f'typing_extensions.{self._name}' + + def __reduce__(self): + return self._name + + def __call__(self, *args, **kwds): + raise TypeError(f"Cannot instantiate {self!r}") + + def __or__(self, other): + return typing.Union[self, other] + + def __ror__(self, other): + return typing.Union[other, self] + + def __instancecheck__(self, obj): + raise TypeError(f"{self} cannot be used with isinstance()") + + def __subclasscheck__(self, cls): + raise TypeError(f"{self} cannot be used with issubclass()") + + @typing._tp_cache + def __getitem__(self, parameters): + return self._getitem(self, parameters) + + +# Note that inheriting from this class means that the object will be +# rejected by typing._type_check, so do not use it if the special form +# is arguably valid as a type by itself. +class _ExtensionsSpecialForm(typing._SpecialForm, _root=True): + def __repr__(self): + return 'typing_extensions.' + self._name + + +Final = typing.Final + +# Breakpoint: https://github.com/python/cpython/pull/30530 +if sys.version_info >= (3, 11): + final = typing.final +else: + # @final exists in 3.8+, but we backport it for all versions + # before 3.11 to keep support for the __final__ attribute. + # See https://bugs.python.org/issue46342 + def final(f): + """This decorator can be used to indicate to type checkers that + the decorated method cannot be overridden, and decorated class + cannot be subclassed. For example: + + class Base: + @final + def done(self) -> None: + ... + class Sub(Base): + def done(self) -> None: # Error reported by type checker + ... + @final + class Leaf: + ... + class Other(Leaf): # Error reported by type checker + ... + + There is no runtime checking of these properties. The decorator + sets the ``__final__`` attribute to ``True`` on the decorated object + to allow runtime introspection. + """ + try: + f.__final__ = True + except (AttributeError, TypeError): + # Skip the attribute silently if it is not writable. + # AttributeError happens if the object has __slots__ or a + # read-only property, TypeError if it's a builtin class. + pass + return f + + +if hasattr(typing, "disjoint_base"): # 3.15 + disjoint_base = typing.disjoint_base +else: + def disjoint_base(cls): + """This decorator marks a class as a disjoint base. + + Child classes of a disjoint base cannot inherit from other disjoint bases that are + not parent classes of the disjoint base. + + For example: + + @disjoint_base + class Disjoint1: pass + + @disjoint_base + class Disjoint2: pass + + class Disjoint3(Disjoint1, Disjoint2): pass # Type checker error + + Type checkers can use knowledge of disjoint bases to detect unreachable code + and determine when two types can overlap. + + See PEP 800.""" + cls.__disjoint_base__ = True + return cls + + +def IntVar(name): + return typing.TypeVar(name) + + +# A Literal bug was fixed in 3.11.0, 3.10.1 and 3.9.8 +# Breakpoint: https://github.com/python/cpython/pull/29334 +if sys.version_info >= (3, 10, 1): + Literal = typing.Literal +else: + def _flatten_literal_params(parameters): + """An internal helper for Literal creation: flatten Literals among parameters""" + params = [] + for p in parameters: + if isinstance(p, _LiteralGenericAlias): + params.extend(p.__args__) + else: + params.append(p) + return tuple(params) + + def _value_and_type_iter(params): + for p in params: + yield p, type(p) + + class _LiteralGenericAlias(typing._GenericAlias, _root=True): + def __eq__(self, other): + if not isinstance(other, _LiteralGenericAlias): + return NotImplemented + these_args_deduped = set(_value_and_type_iter(self.__args__)) + other_args_deduped = set(_value_and_type_iter(other.__args__)) + return these_args_deduped == other_args_deduped + + def __hash__(self): + return hash(frozenset(_value_and_type_iter(self.__args__))) + + class _LiteralForm(_ExtensionsSpecialForm, _root=True): + def __init__(self, doc: str): + self._name = 'Literal' + self._doc = self.__doc__ = doc + + def __getitem__(self, parameters): + if not isinstance(parameters, tuple): + parameters = (parameters,) + + parameters = _flatten_literal_params(parameters) + + val_type_pairs = list(_value_and_type_iter(parameters)) + try: + deduped_pairs = set(val_type_pairs) + except TypeError: + # unhashable parameters + pass + else: + # similar logic to typing._deduplicate on Python 3.9+ + if len(deduped_pairs) < len(val_type_pairs): + new_parameters = [] + for pair in val_type_pairs: + if pair in deduped_pairs: + new_parameters.append(pair[0]) + deduped_pairs.remove(pair) + assert not deduped_pairs, deduped_pairs + parameters = tuple(new_parameters) + + return _LiteralGenericAlias(self, parameters) + + Literal = _LiteralForm(doc="""\ + A type that can be used to indicate to type checkers + that the corresponding value has a value literally equivalent + to the provided parameter. For example: + + var: Literal[4] = 4 + + The type checker understands that 'var' is literally equal to + the value 4 and no other value. + + Literal[...] cannot be subclassed. There is no runtime + checking verifying that the parameter is actually a value + instead of a type.""") + + +_overload_dummy = typing._overload_dummy + + +if hasattr(typing, "get_overloads"): # 3.11+ + overload = typing.overload + get_overloads = typing.get_overloads + clear_overloads = typing.clear_overloads +else: + # {module: {qualname: {firstlineno: func}}} + _overload_registry = collections.defaultdict( + functools.partial(collections.defaultdict, dict) + ) + + def overload(func): + """Decorator for overloaded functions/methods. + + In a stub file, place two or more stub definitions for the same + function in a row, each decorated with @overload. For example: + + @overload + def utf8(value: None) -> None: ... + @overload + def utf8(value: bytes) -> bytes: ... + @overload + def utf8(value: str) -> bytes: ... + + In a non-stub file (i.e. a regular .py file), do the same but + follow it with an implementation. The implementation should *not* + be decorated with @overload. For example: + + @overload + def utf8(value: None) -> None: ... + @overload + def utf8(value: bytes) -> bytes: ... + @overload + def utf8(value: str) -> bytes: ... + def utf8(value): + # implementation goes here + + The overloads for a function can be retrieved at runtime using the + get_overloads() function. + """ + # classmethod and staticmethod + f = getattr(func, "__func__", func) + try: + _overload_registry[f.__module__][f.__qualname__][ + f.__code__.co_firstlineno + ] = func + except AttributeError: + # Not a normal function; ignore. + pass + return _overload_dummy + + def get_overloads(func): + """Return all defined overloads for *func* as a sequence.""" + # classmethod and staticmethod + f = getattr(func, "__func__", func) + if f.__module__ not in _overload_registry: + return [] + mod_dict = _overload_registry[f.__module__] + if f.__qualname__ not in mod_dict: + return [] + return list(mod_dict[f.__qualname__].values()) + + def clear_overloads(): + """Clear all overloads in the registry.""" + _overload_registry.clear() + + +# This is not a real generic class. Don't use outside annotations. +Type = typing.Type + +# Various ABCs mimicking those in collections.abc. +# A few are simply re-exported for completeness. +Awaitable = typing.Awaitable +Coroutine = typing.Coroutine +AsyncIterable = typing.AsyncIterable +AsyncIterator = typing.AsyncIterator +Deque = typing.Deque +DefaultDict = typing.DefaultDict +OrderedDict = typing.OrderedDict +Counter = typing.Counter +ChainMap = typing.ChainMap +Text = typing.Text +TYPE_CHECKING = typing.TYPE_CHECKING + + +# Breakpoint: https://github.com/python/cpython/pull/118681 +if sys.version_info >= (3, 13, 0, "beta"): + from typing import AsyncContextManager, AsyncGenerator, ContextManager, Generator +else: + def _is_dunder(attr): + return attr.startswith('__') and attr.endswith('__') + + + class _SpecialGenericAlias(typing._SpecialGenericAlias, _root=True): + def __init__(self, origin, nparams, *, inst=True, name=None, defaults=()): + super().__init__(origin, nparams, inst=inst, name=name) + self._defaults = defaults + + def __setattr__(self, attr, val): + allowed_attrs = {'_name', '_inst', '_nparams', '_defaults'} + if _is_dunder(attr) or attr in allowed_attrs: + object.__setattr__(self, attr, val) + else: + setattr(self.__origin__, attr, val) + + @typing._tp_cache + def __getitem__(self, params): + if not isinstance(params, tuple): + params = (params,) + msg = "Parameters to generic types must be types." + params = tuple(typing._type_check(p, msg) for p in params) + if ( + self._defaults + and len(params) < self._nparams + and len(params) + len(self._defaults) >= self._nparams + ): + params = (*params, *self._defaults[len(params) - self._nparams:]) + actual_len = len(params) + + if actual_len != self._nparams: + if self._defaults: + expected = f"at least {self._nparams - len(self._defaults)}" + else: + expected = str(self._nparams) + if not self._nparams: + raise TypeError(f"{self} is not a generic class") + raise TypeError( + f"Too {'many' if actual_len > self._nparams else 'few'}" + f" arguments for {self};" + f" actual {actual_len}, expected {expected}" + ) + return self.copy_with(params) + + _NoneType = type(None) + Generator = _SpecialGenericAlias( + collections.abc.Generator, 3, defaults=(_NoneType, _NoneType) + ) + AsyncGenerator = _SpecialGenericAlias( + collections.abc.AsyncGenerator, 2, defaults=(_NoneType,) + ) + ContextManager = _SpecialGenericAlias( + contextlib.AbstractContextManager, + 2, + name="ContextManager", + defaults=(typing.Optional[bool],) + ) + AsyncContextManager = _SpecialGenericAlias( + contextlib.AbstractAsyncContextManager, + 2, + name="AsyncContextManager", + defaults=(typing.Optional[bool],) + ) + + +_PROTO_ALLOWLIST = { + 'collections.abc': [ + 'Callable', 'Awaitable', 'Iterable', 'Iterator', 'AsyncIterable', + 'Hashable', 'Sized', 'Container', 'Collection', 'Reversible', 'Buffer', + ], + 'contextlib': ['AbstractContextManager', 'AbstractAsyncContextManager'], + 'typing_extensions': ['Buffer'], +} + + +_EXCLUDED_ATTRS = frozenset(typing.EXCLUDED_ATTRIBUTES) | { + "__match_args__", "__protocol_attrs__", "__non_callable_proto_members__", + "__final__", +} + + +def _get_protocol_attrs(cls): + attrs = set() + for base in cls.__mro__[:-1]: # without object + if base.__name__ in {'Protocol', 'Generic'}: + continue + annotations = getattr(base, '__annotations__', {}) + for attr in (*base.__dict__, *annotations): + if (not attr.startswith('_abc_') and attr not in _EXCLUDED_ATTRS): + attrs.add(attr) + return attrs + + +def _caller(depth=1, default='__main__'): + try: + return sys._getframemodulename(depth + 1) or default + except AttributeError: # For platforms without _getframemodulename() + pass + try: + return sys._getframe(depth + 1).f_globals.get('__name__', default) + except (AttributeError, ValueError): # For platforms without _getframe() + pass + return None + + +# `__match_args__` attribute was removed from protocol members in 3.13, +# we want to backport this change to older Python versions. +# Breakpoint: https://github.com/python/cpython/pull/110683 +if sys.version_info >= (3, 13): + Protocol = typing.Protocol +else: + def _allow_reckless_class_checks(depth=2): + """Allow instance and class checks for special stdlib modules. + The abc and functools modules indiscriminately call isinstance() and + issubclass() on the whole MRO of a user class, which may contain protocols. + """ + return _caller(depth) in {'abc', 'functools', None} + + def _no_init(self, *args, **kwargs): + if type(self)._is_protocol: + raise TypeError('Protocols cannot be instantiated') + + def _type_check_issubclass_arg_1(arg): + """Raise TypeError if `arg` is not an instance of `type` + in `issubclass(arg, )`. + + In most cases, this is verified by type.__subclasscheck__. + Checking it again unnecessarily would slow down issubclass() checks, + so, we don't perform this check unless we absolutely have to. + + For various error paths, however, + we want to ensure that *this* error message is shown to the user + where relevant, rather than a typing.py-specific error message. + """ + if not isinstance(arg, type): + # Same error message as for issubclass(1, int). + raise TypeError('issubclass() arg 1 must be a class') + + # Inheriting from typing._ProtocolMeta isn't actually desirable, + # but is necessary to allow typing.Protocol and typing_extensions.Protocol + # to mix without getting TypeErrors about "metaclass conflict" + class _ProtocolMeta(type(typing.Protocol)): + # This metaclass is somewhat unfortunate, + # but is necessary for several reasons... + # + # NOTE: DO NOT call super() in any methods in this class + # That would call the methods on typing._ProtocolMeta on Python <=3.11 + # and those are slow + def __new__(mcls, name, bases, namespace, **kwargs): + if name == "Protocol" and len(bases) < 2: + pass + elif {Protocol, typing.Protocol} & set(bases): + for base in bases: + if not ( + base in {object, typing.Generic, Protocol, typing.Protocol} + or base.__name__ in _PROTO_ALLOWLIST.get(base.__module__, []) + or is_protocol(base) + ): + raise TypeError( + f"Protocols can only inherit from other protocols, " + f"got {base!r}" + ) + return abc.ABCMeta.__new__(mcls, name, bases, namespace, **kwargs) + + def __init__(cls, *args, **kwargs): + abc.ABCMeta.__init__(cls, *args, **kwargs) + if getattr(cls, "_is_protocol", False): + cls.__protocol_attrs__ = _get_protocol_attrs(cls) + + def __subclasscheck__(cls, other): + if cls is Protocol: + return type.__subclasscheck__(cls, other) + if ( + getattr(cls, '_is_protocol', False) + and not _allow_reckless_class_checks() + ): + if not getattr(cls, '_is_runtime_protocol', False): + _type_check_issubclass_arg_1(other) + raise TypeError( + "Instance and class checks can only be used with " + "@runtime_checkable protocols" + ) + if ( + # this attribute is set by @runtime_checkable: + cls.__non_callable_proto_members__ + and cls.__dict__.get("__subclasshook__") is _proto_hook + ): + _type_check_issubclass_arg_1(other) + non_method_attrs = sorted(cls.__non_callable_proto_members__) + raise TypeError( + "Protocols with non-method members don't support issubclass()." + f" Non-method members: {str(non_method_attrs)[1:-1]}." + ) + return abc.ABCMeta.__subclasscheck__(cls, other) + + def __instancecheck__(cls, instance): + # We need this method for situations where attributes are + # assigned in __init__. + if cls is Protocol: + return type.__instancecheck__(cls, instance) + if not getattr(cls, "_is_protocol", False): + # i.e., it's a concrete subclass of a protocol + return abc.ABCMeta.__instancecheck__(cls, instance) + + if ( + not getattr(cls, '_is_runtime_protocol', False) and + not _allow_reckless_class_checks() + ): + raise TypeError("Instance and class checks can only be used with" + " @runtime_checkable protocols") + + if abc.ABCMeta.__instancecheck__(cls, instance): + return True + + for attr in cls.__protocol_attrs__: + try: + val = inspect.getattr_static(instance, attr) + except AttributeError: + break + # this attribute is set by @runtime_checkable: + if val is None and attr not in cls.__non_callable_proto_members__: + break + else: + return True + + return False + + def __eq__(cls, other): + # Hack so that typing.Generic.__class_getitem__ + # treats typing_extensions.Protocol + # as equivalent to typing.Protocol + if abc.ABCMeta.__eq__(cls, other) is True: + return True + return cls is Protocol and other is typing.Protocol + + # This has to be defined, or the abc-module cache + # complains about classes with this metaclass being unhashable, + # if we define only __eq__! + def __hash__(cls) -> int: + return type.__hash__(cls) + + @classmethod + def _proto_hook(cls, other): + if not cls.__dict__.get('_is_protocol', False): + return NotImplemented + + for attr in cls.__protocol_attrs__: + for base in other.__mro__: + # Check if the members appears in the class dictionary... + if attr in base.__dict__: + if base.__dict__[attr] is None: + return NotImplemented + break + + # ...or in annotations, if it is a sub-protocol. + annotations = getattr(base, '__annotations__', {}) + if ( + isinstance(annotations, collections.abc.Mapping) + and attr in annotations + and is_protocol(other) + ): + break + else: + return NotImplemented + return True + + class Protocol(typing.Generic, metaclass=_ProtocolMeta): + __doc__ = typing.Protocol.__doc__ + __slots__ = () + _is_protocol = True + _is_runtime_protocol = False + + def __init_subclass__(cls, *args, **kwargs): + super().__init_subclass__(*args, **kwargs) + + # Determine if this is a protocol or a concrete subclass. + if not cls.__dict__.get('_is_protocol', False): + cls._is_protocol = any(b is Protocol for b in cls.__bases__) + + # Set (or override) the protocol subclass hook. + if '__subclasshook__' not in cls.__dict__: + cls.__subclasshook__ = _proto_hook + + # Prohibit instantiation for protocol classes + if cls._is_protocol and cls.__init__ is Protocol.__init__: + cls.__init__ = _no_init + + +# Breakpoint: https://github.com/python/cpython/pull/113401 +if sys.version_info >= (3, 13): + runtime_checkable = typing.runtime_checkable +else: + def runtime_checkable(cls): + """Mark a protocol class as a runtime protocol. + + Such protocol can be used with isinstance() and issubclass(). + Raise TypeError if applied to a non-protocol class. + This allows a simple-minded structural check very similar to + one trick ponies in collections.abc such as Iterable. + + For example:: + + @runtime_checkable + class Closable(Protocol): + def close(self): ... + + assert isinstance(open('/some/file'), Closable) + + Warning: this will check only the presence of the required methods, + not their type signatures! + """ + if not issubclass(cls, typing.Generic) or not getattr(cls, '_is_protocol', False): + raise TypeError(f'@runtime_checkable can be only applied to protocol classes,' + f' got {cls!r}') + cls._is_runtime_protocol = True + + # typing.Protocol classes on <=3.11 break if we execute this block, + # because typing.Protocol classes on <=3.11 don't have a + # `__protocol_attrs__` attribute, and this block relies on the + # `__protocol_attrs__` attribute. Meanwhile, typing.Protocol classes on 3.12.2+ + # break if we *don't* execute this block, because *they* assume that all + # protocol classes have a `__non_callable_proto_members__` attribute + # (which this block sets) + if isinstance(cls, _ProtocolMeta) or sys.version_info >= (3, 12, 2): + # PEP 544 prohibits using issubclass() + # with protocols that have non-method members. + # See gh-113320 for why we compute this attribute here, + # rather than in `_ProtocolMeta.__init__` + cls.__non_callable_proto_members__ = set() + for attr in cls.__protocol_attrs__: + try: + is_callable = callable(getattr(cls, attr, None)) + except Exception as e: + raise TypeError( + f"Failed to determine whether protocol member {attr!r} " + "is a method member" + ) from e + else: + if not is_callable: + cls.__non_callable_proto_members__.add(attr) + + return cls + + +# The "runtime" alias exists for backwards compatibility. +runtime = runtime_checkable + + +# Our version of runtime-checkable protocols is faster on Python <=3.11 +# Breakpoint: https://github.com/python/cpython/pull/112717 +if sys.version_info >= (3, 12): + SupportsInt = typing.SupportsInt + SupportsFloat = typing.SupportsFloat + SupportsComplex = typing.SupportsComplex + SupportsBytes = typing.SupportsBytes + SupportsIndex = typing.SupportsIndex + SupportsAbs = typing.SupportsAbs + SupportsRound = typing.SupportsRound +else: + @runtime_checkable + class SupportsInt(Protocol): + """An ABC with one abstract method __int__.""" + __slots__ = () + + @abc.abstractmethod + def __int__(self) -> int: + pass + + @runtime_checkable + class SupportsFloat(Protocol): + """An ABC with one abstract method __float__.""" + __slots__ = () + + @abc.abstractmethod + def __float__(self) -> float: + pass + + @runtime_checkable + class SupportsComplex(Protocol): + """An ABC with one abstract method __complex__.""" + __slots__ = () + + @abc.abstractmethod + def __complex__(self) -> complex: + pass + + @runtime_checkable + class SupportsBytes(Protocol): + """An ABC with one abstract method __bytes__.""" + __slots__ = () + + @abc.abstractmethod + def __bytes__(self) -> bytes: + pass + + @runtime_checkable + class SupportsIndex(Protocol): + __slots__ = () + + @abc.abstractmethod + def __index__(self) -> int: + pass + + @runtime_checkable + class SupportsAbs(Protocol[T_co]): + """ + An ABC with one abstract method __abs__ that is covariant in its return type. + """ + __slots__ = () + + @abc.abstractmethod + def __abs__(self) -> T_co: + pass + + @runtime_checkable + class SupportsRound(Protocol[T_co]): + """ + An ABC with one abstract method __round__ that is covariant in its return type. + """ + __slots__ = () + + @abc.abstractmethod + def __round__(self, ndigits: int = 0) -> T_co: + pass + + +if hasattr(io, "Reader") and hasattr(io, "Writer"): + Reader = io.Reader + Writer = io.Writer +else: + @runtime_checkable + class Reader(Protocol[T_co]): + """Protocol for simple I/O reader instances. + + This protocol only supports blocking I/O. + """ + + __slots__ = () + + @abc.abstractmethod + def read(self, size: int = ..., /) -> T_co: + """Read data from the input stream and return it. + + If *size* is specified, at most *size* items (bytes/characters) will be + read. + """ + + @runtime_checkable + class Writer(Protocol[T_contra]): + """Protocol for simple I/O writer instances. + + This protocol only supports blocking I/O. + """ + + __slots__ = () + + @abc.abstractmethod + def write(self, data: T_contra, /) -> int: + """Write *data* to the output stream and return the number of items written.""" # noqa: E501 + + +_NEEDS_SINGLETONMETA = ( + not hasattr(typing, "NoDefault") or not hasattr(typing, "NoExtraItems") +) + +if _NEEDS_SINGLETONMETA: + class SingletonMeta(type): + def __setattr__(cls, attr, value): + # TypeError is consistent with the behavior of NoneType + raise TypeError( + f"cannot set {attr!r} attribute of immutable type {cls.__name__!r}" + ) + + +if hasattr(typing, "NoDefault"): + NoDefault = typing.NoDefault +else: + class NoDefaultType(metaclass=SingletonMeta): + """The type of the NoDefault singleton.""" + + __slots__ = () + + def __new__(cls): + return globals().get("NoDefault") or object.__new__(cls) + + def __repr__(self): + return "typing_extensions.NoDefault" + + def __reduce__(self): + return "NoDefault" + + NoDefault = NoDefaultType() + del NoDefaultType + +if hasattr(typing, "NoExtraItems"): + NoExtraItems = typing.NoExtraItems +else: + class NoExtraItemsType(metaclass=SingletonMeta): + """The type of the NoExtraItems singleton.""" + + __slots__ = () + + def __new__(cls): + return globals().get("NoExtraItems") or object.__new__(cls) + + def __repr__(self): + return "typing_extensions.NoExtraItems" + + def __reduce__(self): + return "NoExtraItems" + + NoExtraItems = NoExtraItemsType() + del NoExtraItemsType + +if _NEEDS_SINGLETONMETA: + del SingletonMeta + + +# Update this to something like >=3.13.0b1 if and when +# PEP 728 is implemented in CPython +_PEP_728_IMPLEMENTED = False + +if _PEP_728_IMPLEMENTED: + # The standard library TypedDict in Python 3.9.0/1 does not honour the "total" + # keyword with old-style TypedDict(). See https://bugs.python.org/issue42059 + # The standard library TypedDict below Python 3.11 does not store runtime + # information about optional and required keys when using Required or NotRequired. + # Generic TypedDicts are also impossible using typing.TypedDict on Python <3.11. + # Aaaand on 3.12 we add __orig_bases__ to TypedDict + # to enable better runtime introspection. + # On 3.13 we deprecate some odd ways of creating TypedDicts. + # Also on 3.13, PEP 705 adds the ReadOnly[] qualifier. + # PEP 728 (still pending) makes more changes. + TypedDict = typing.TypedDict + _TypedDictMeta = typing._TypedDictMeta + is_typeddict = typing.is_typeddict +else: + # 3.10.0 and later + _TAKES_MODULE = "module" in inspect.signature(typing._type_check).parameters + + def _get_typeddict_qualifiers(annotation_type): + while True: + annotation_origin = get_origin(annotation_type) + if annotation_origin is Annotated: + annotation_args = get_args(annotation_type) + if annotation_args: + annotation_type = annotation_args[0] + else: + break + elif annotation_origin is Required: + yield Required + annotation_type, = get_args(annotation_type) + elif annotation_origin is NotRequired: + yield NotRequired + annotation_type, = get_args(annotation_type) + elif annotation_origin is ReadOnly: + yield ReadOnly + annotation_type, = get_args(annotation_type) + else: + break + + class _TypedDictMeta(type): + + def __new__(cls, name, bases, ns, *, total=True, closed=None, + extra_items=NoExtraItems): + """Create new typed dict class object. + + This method is called when TypedDict is subclassed, + or when TypedDict is instantiated. This way + TypedDict supports all three syntax forms described in its docstring. + Subclasses and instances of TypedDict return actual dictionaries. + """ + for base in bases: + if type(base) is not _TypedDictMeta and base is not typing.Generic: + raise TypeError('cannot inherit from both a TypedDict type ' + 'and a non-TypedDict base class') + if closed is not None and extra_items is not NoExtraItems: + raise TypeError(f"Cannot combine closed={closed!r} and extra_items") + + if any(issubclass(b, typing.Generic) for b in bases): + generic_base = (typing.Generic,) + else: + generic_base = () + + ns_annotations = ns.pop('__annotations__', None) + + # typing.py generally doesn't let you inherit from plain Generic, unless + # the name of the class happens to be "Protocol" + tp_dict = type.__new__(_TypedDictMeta, "Protocol", (*generic_base, dict), ns) + tp_dict.__name__ = name + if tp_dict.__qualname__ == "Protocol": + tp_dict.__qualname__ = name + + if not hasattr(tp_dict, '__orig_bases__'): + tp_dict.__orig_bases__ = bases + + annotations = {} + own_annotate = None + if ns_annotations is not None: + own_annotations = ns_annotations + elif sys.version_info >= (3, 14): + if hasattr(annotationlib, "get_annotate_from_class_namespace"): + own_annotate = annotationlib.get_annotate_from_class_namespace(ns) + else: + # 3.14.0a7 and earlier + own_annotate = ns.get("__annotate__") + if own_annotate is not None: + own_annotations = annotationlib.call_annotate_function( + own_annotate, Format.FORWARDREF, owner=tp_dict + ) + else: + own_annotations = {} + else: + own_annotations = {} + msg = "TypedDict('Name', {f0: t0, f1: t1, ...}); each t must be a type" + if _TAKES_MODULE: + own_checked_annotations = { + n: typing._type_check(tp, msg, module=tp_dict.__module__) + for n, tp in own_annotations.items() + } + else: + own_checked_annotations = { + n: typing._type_check(tp, msg) + for n, tp in own_annotations.items() + } + required_keys = set() + optional_keys = set() + readonly_keys = set() + mutable_keys = set() + extra_items_type = extra_items + + for base in bases: + base_dict = base.__dict__ + + if sys.version_info <= (3, 14): + annotations.update(base_dict.get('__annotations__', {})) + required_keys.update(base_dict.get('__required_keys__', ())) + optional_keys.update(base_dict.get('__optional_keys__', ())) + readonly_keys.update(base_dict.get('__readonly_keys__', ())) + mutable_keys.update(base_dict.get('__mutable_keys__', ())) + + # This was specified in an earlier version of PEP 728. Support + # is retained for backwards compatibility, but only for Python + # 3.13 and lower. + if (closed and sys.version_info < (3, 14) + and "__extra_items__" in own_checked_annotations): + annotation_type = own_checked_annotations.pop("__extra_items__") + qualifiers = set(_get_typeddict_qualifiers(annotation_type)) + if Required in qualifiers: + raise TypeError( + "Special key __extra_items__ does not support " + "Required" + ) + if NotRequired in qualifiers: + raise TypeError( + "Special key __extra_items__ does not support " + "NotRequired" + ) + extra_items_type = annotation_type + + annotations.update(own_checked_annotations) + for annotation_key, annotation_type in own_checked_annotations.items(): + qualifiers = set(_get_typeddict_qualifiers(annotation_type)) + + if Required in qualifiers: + required_keys.add(annotation_key) + elif NotRequired in qualifiers: + optional_keys.add(annotation_key) + elif total: + required_keys.add(annotation_key) + else: + optional_keys.add(annotation_key) + if ReadOnly in qualifiers: + mutable_keys.discard(annotation_key) + readonly_keys.add(annotation_key) + else: + mutable_keys.add(annotation_key) + readonly_keys.discard(annotation_key) + + # Breakpoint: https://github.com/python/cpython/pull/119891 + if sys.version_info >= (3, 14): + def __annotate__(format): + annos = {} + for base in bases: + if base is Generic: + continue + base_annotate = base.__annotate__ + if base_annotate is None: + continue + base_annos = annotationlib.call_annotate_function( + base_annotate, format, owner=base) + annos.update(base_annos) + if own_annotate is not None: + own = annotationlib.call_annotate_function( + own_annotate, format, owner=tp_dict) + if format != Format.STRING: + own = { + n: typing._type_check(tp, msg, module=tp_dict.__module__) + for n, tp in own.items() + } + elif format == Format.STRING: + own = annotationlib.annotations_to_string(own_annotations) + elif format in (Format.FORWARDREF, Format.VALUE): + own = own_checked_annotations + else: + raise NotImplementedError(format) + annos.update(own) + return annos + + tp_dict.__annotate__ = __annotate__ + else: + tp_dict.__annotations__ = annotations + tp_dict.__required_keys__ = frozenset(required_keys) + tp_dict.__optional_keys__ = frozenset(optional_keys) + tp_dict.__readonly_keys__ = frozenset(readonly_keys) + tp_dict.__mutable_keys__ = frozenset(mutable_keys) + tp_dict.__total__ = total + tp_dict.__closed__ = closed + tp_dict.__extra_items__ = extra_items_type + return tp_dict + + __call__ = dict # static method + + def __subclasscheck__(cls, other): + # Typed dicts are only for static structural subtyping. + raise TypeError('TypedDict does not support instance and class checks') + + __instancecheck__ = __subclasscheck__ + + _TypedDict = type.__new__(_TypedDictMeta, 'TypedDict', (), {}) + + def _create_typeddict( + typename, + fields, + /, + *, + typing_is_inline, + total, + closed, + extra_items, + **kwargs, + ): + if fields is _marker or fields is None: + if fields is _marker: + deprecated_thing = ( + "Failing to pass a value for the 'fields' parameter" + ) + else: + deprecated_thing = "Passing `None` as the 'fields' parameter" + + example = f"`{typename} = TypedDict({typename!r}, {{}})`" + deprecation_msg = ( + f"{deprecated_thing} is deprecated and will be disallowed in " + "Python 3.15. To create a TypedDict class with 0 fields " + "using the functional syntax, pass an empty dictionary, e.g. " + ) + example + "." + warnings.warn(deprecation_msg, DeprecationWarning, stacklevel=2) + # Support a field called "closed" + if closed is not False and closed is not True and closed is not None: + kwargs["closed"] = closed + closed = None + # Or "extra_items" + if extra_items is not NoExtraItems: + kwargs["extra_items"] = extra_items + extra_items = NoExtraItems + fields = kwargs + elif kwargs: + raise TypeError("TypedDict takes either a dict or keyword arguments," + " but not both") + if kwargs: + # Breakpoint: https://github.com/python/cpython/pull/104891 + if sys.version_info >= (3, 13): + raise TypeError("TypedDict takes no keyword arguments") + warnings.warn( + "The kwargs-based syntax for TypedDict definitions is deprecated " + "in Python 3.11, will be removed in Python 3.13, and may not be " + "understood by third-party type checkers.", + DeprecationWarning, + stacklevel=2, + ) + + ns = {'__annotations__': dict(fields)} + module = _caller(depth=4 if typing_is_inline else 2) + if module is not None: + # Setting correct module is necessary to make typed dict classes + # pickleable. + ns['__module__'] = module + + td = _TypedDictMeta(typename, (), ns, total=total, closed=closed, + extra_items=extra_items) + td.__orig_bases__ = (TypedDict,) + return td + + class _TypedDictSpecialForm(_SpecialForm, _root=True): + def __call__( + self, + typename, + fields=_marker, + /, + *, + total=True, + closed=None, + extra_items=NoExtraItems, + **kwargs + ): + return _create_typeddict( + typename, + fields, + typing_is_inline=False, + total=total, + closed=closed, + extra_items=extra_items, + **kwargs, + ) + + def __mro_entries__(self, bases): + return (_TypedDict,) + + @_TypedDictSpecialForm + def TypedDict(self, args): + """A simple typed namespace. At runtime it is equivalent to a plain dict. + + TypedDict creates a dictionary type such that a type checker will expect all + instances to have a certain set of keys, where each key is + associated with a value of a consistent type. This expectation + is not checked at runtime. + + Usage:: + + class Point2D(TypedDict): + x: int + y: int + label: str + + a: Point2D = {'x': 1, 'y': 2, 'label': 'good'} # OK + b: Point2D = {'z': 3, 'label': 'bad'} # Fails type check + + assert Point2D(x=1, y=2, label='first') == dict(x=1, y=2, label='first') + + The type info can be accessed via the Point2D.__annotations__ dict, and + the Point2D.__required_keys__ and Point2D.__optional_keys__ frozensets. + TypedDict supports an additional equivalent form:: + + Point2D = TypedDict('Point2D', {'x': int, 'y': int, 'label': str}) + + By default, all keys must be present in a TypedDict. It is possible + to override this by specifying totality:: + + class Point2D(TypedDict, total=False): + x: int + y: int + + This means that a Point2D TypedDict can have any of the keys omitted. A type + checker is only expected to support a literal False or True as the value of + the total argument. True is the default, and makes all items defined in the + class body be required. + + The Required and NotRequired special forms can also be used to mark + individual keys as being required or not required:: + + class Point2D(TypedDict): + x: int # the "x" key must always be present (Required is the default) + y: NotRequired[int] # the "y" key can be omitted + + See PEP 655 for more details on Required and NotRequired. + """ + # This runs when creating inline TypedDicts: + if not isinstance(args, dict): + raise TypeError( + "TypedDict[...] should be used with a single dict argument" + ) + + return _create_typeddict( + "", + args, + typing_is_inline=True, + total=True, + closed=True, + extra_items=NoExtraItems, + ) + + _TYPEDDICT_TYPES = (typing._TypedDictMeta, _TypedDictMeta) + + def is_typeddict(tp): + """Check if an annotation is a TypedDict class + + For example:: + class Film(TypedDict): + title: str + year: int + + is_typeddict(Film) # => True + is_typeddict(Union[list, str]) # => False + """ + return isinstance(tp, _TYPEDDICT_TYPES) + + +if hasattr(typing, "assert_type"): + assert_type = typing.assert_type + +else: + def assert_type(val, typ, /): + """Assert (to the type checker) that the value is of the given type. + + When the type checker encounters a call to assert_type(), it + emits an error if the value is not of the specified type:: + + def greet(name: str) -> None: + assert_type(name, str) # ok + assert_type(name, int) # type checker error + + At runtime this returns the first argument unchanged and otherwise + does nothing. + """ + return val + + +if hasattr(typing, "ReadOnly"): # 3.13+ + get_type_hints = typing.get_type_hints +else: # <=3.13 + # replaces _strip_annotations() + def _strip_extras(t): + """Strips Annotated, Required and NotRequired from a given type.""" + if isinstance(t, typing._AnnotatedAlias): + return _strip_extras(t.__origin__) + if hasattr(t, "__origin__") and t.__origin__ in (Required, NotRequired, ReadOnly): + return _strip_extras(t.__args__[0]) + if isinstance(t, typing._GenericAlias): + stripped_args = tuple(_strip_extras(a) for a in t.__args__) + if stripped_args == t.__args__: + return t + return t.copy_with(stripped_args) + if hasattr(_types, "GenericAlias") and isinstance(t, _types.GenericAlias): + stripped_args = tuple(_strip_extras(a) for a in t.__args__) + if stripped_args == t.__args__: + return t + return _types.GenericAlias(t.__origin__, stripped_args) + if hasattr(_types, "UnionType") and isinstance(t, _types.UnionType): + stripped_args = tuple(_strip_extras(a) for a in t.__args__) + if stripped_args == t.__args__: + return t + return functools.reduce(operator.or_, stripped_args) + + return t + + def get_type_hints(obj, globalns=None, localns=None, include_extras=False): + """Return type hints for an object. + + This is often the same as obj.__annotations__, but it handles + forward references encoded as string literals, adds Optional[t] if a + default value equal to None is set and recursively replaces all + 'Annotated[T, ...]', 'Required[T]' or 'NotRequired[T]' with 'T' + (unless 'include_extras=True'). + + The argument may be a module, class, method, or function. The annotations + are returned as a dictionary. For classes, annotations include also + inherited members. + + TypeError is raised if the argument is not of a type that can contain + annotations, and an empty dictionary is returned if no annotations are + present. + + BEWARE -- the behavior of globalns and localns is counterintuitive + (unless you are familiar with how eval() and exec() work). The + search order is locals first, then globals. + + - If no dict arguments are passed, an attempt is made to use the + globals from obj (or the respective module's globals for classes), + and these are also used as the locals. If the object does not appear + to have globals, an empty dictionary is used. + + - If one dict argument is passed, it is used for both globals and + locals. + + - If two dict arguments are passed, they specify globals and + locals, respectively. + """ + hint = typing.get_type_hints( + obj, globalns=globalns, localns=localns, include_extras=True + ) + # Breakpoint: https://github.com/python/cpython/pull/30304 + if sys.version_info < (3, 11): + _clean_optional(obj, hint, globalns, localns) + if include_extras: + return hint + return {k: _strip_extras(t) for k, t in hint.items()} + + _NoneType = type(None) + + def _could_be_inserted_optional(t): + """detects Union[..., None] pattern""" + if not isinstance(t, typing._UnionGenericAlias): + return False + # Assume if last argument is not None they are user defined + if t.__args__[-1] is not _NoneType: + return False + return True + + # < 3.11 + def _clean_optional(obj, hints, globalns=None, localns=None): + # reverts injected Union[..., None] cases from typing.get_type_hints + # when a None default value is used. + # see https://github.com/python/typing_extensions/issues/310 + if not hints or isinstance(obj, type): + return + defaults = typing._get_defaults(obj) # avoid accessing __annotations___ + if not defaults: + return + original_hints = obj.__annotations__ + for name, value in hints.items(): + # Not a Union[..., None] or replacement conditions not fullfilled + if (not _could_be_inserted_optional(value) + or name not in defaults + or defaults[name] is not None + ): + continue + original_value = original_hints[name] + # value=NoneType should have caused a skip above but check for safety + if original_value is None: + original_value = _NoneType + # Forward reference + if isinstance(original_value, str): + if globalns is None: + if isinstance(obj, _types.ModuleType): + globalns = obj.__dict__ + else: + nsobj = obj + # Find globalns for the unwrapped object. + while hasattr(nsobj, '__wrapped__'): + nsobj = nsobj.__wrapped__ + globalns = getattr(nsobj, '__globals__', {}) + if localns is None: + localns = globalns + elif localns is None: + localns = globalns + + original_value = ForwardRef( + original_value, + is_argument=not isinstance(obj, _types.ModuleType) + ) + original_evaluated = typing._eval_type(original_value, globalns, localns) + # Compare if values differ. Note that even if equal + # value might be cached by typing._tp_cache contrary to original_evaluated + if original_evaluated != value or ( + # 3.10: ForwardRefs of UnionType might be turned into _UnionGenericAlias + hasattr(_types, "UnionType") + and isinstance(original_evaluated, _types.UnionType) + and not isinstance(value, _types.UnionType) + ): + hints[name] = original_evaluated + +# Python 3.9 has get_origin() and get_args() but those implementations don't support +# ParamSpecArgs and ParamSpecKwargs, so only Python 3.10's versions will do. +# Breakpoint: https://github.com/python/cpython/pull/25298 +if sys.version_info >= (3, 10): + get_origin = typing.get_origin + get_args = typing.get_args +# 3.9 +else: + def get_origin(tp): + """Get the unsubscripted version of a type. + + This supports generic types, Callable, Tuple, Union, Literal, Final, ClassVar + and Annotated. Return None for unsupported types. Examples:: + + get_origin(Literal[42]) is Literal + get_origin(int) is None + get_origin(ClassVar[int]) is ClassVar + get_origin(Generic) is Generic + get_origin(Generic[T]) is Generic + get_origin(Union[T, int]) is Union + get_origin(List[Tuple[T, T]][int]) == list + get_origin(P.args) is P + """ + if isinstance(tp, typing._AnnotatedAlias): + return Annotated + if isinstance(tp, (typing._BaseGenericAlias, _types.GenericAlias, + ParamSpecArgs, ParamSpecKwargs)): + return tp.__origin__ + if tp is typing.Generic: + return typing.Generic + return None + + def get_args(tp): + """Get type arguments with all substitutions performed. + + For unions, basic simplifications used by Union constructor are performed. + Examples:: + get_args(Dict[str, int]) == (str, int) + get_args(int) == () + get_args(Union[int, Union[T, int], str][int]) == (int, str) + get_args(Union[int, Tuple[T, int]][str]) == (int, Tuple[str, int]) + get_args(Callable[[], T][int]) == ([], int) + """ + if isinstance(tp, typing._AnnotatedAlias): + return (tp.__origin__, *tp.__metadata__) + if isinstance(tp, (typing._GenericAlias, _types.GenericAlias)): + res = tp.__args__ + if get_origin(tp) is collections.abc.Callable and res[0] is not Ellipsis: + res = (list(res[:-1]), res[-1]) + return res + return () + + +# 3.10+ +if hasattr(typing, 'TypeAlias'): + TypeAlias = typing.TypeAlias +# 3.9 +else: + @_ExtensionsSpecialForm + def TypeAlias(self, parameters): + """Special marker indicating that an assignment should + be recognized as a proper type alias definition by type + checkers. + + For example:: + + Predicate: TypeAlias = Callable[..., bool] + + It's invalid when used anywhere except as in the example above. + """ + raise TypeError(f"{self} is not subscriptable") + + +def _set_default(type_param, default): + type_param.has_default = lambda: default is not NoDefault + type_param.__default__ = default + + +def _set_module(typevarlike): + # for pickling: + def_mod = _caller(depth=2) + if def_mod != 'typing_extensions': + typevarlike.__module__ = def_mod + + +class _DefaultMixin: + """Mixin for TypeVarLike defaults.""" + + __slots__ = () + __init__ = _set_default + + +# Classes using this metaclass must provide a _backported_typevarlike ClassVar +class _TypeVarLikeMeta(type): + def __instancecheck__(cls, __instance: Any) -> bool: + return isinstance(__instance, cls._backported_typevarlike) + + +if _PEP_696_IMPLEMENTED: + from typing import TypeVar +else: + # Add default and infer_variance parameters from PEP 696 and 695 + class TypeVar(metaclass=_TypeVarLikeMeta): + """Type variable.""" + + _backported_typevarlike = typing.TypeVar + + def __new__(cls, name, *constraints, bound=None, + covariant=False, contravariant=False, + default=NoDefault, infer_variance=False): + if hasattr(typing, "TypeAliasType"): + # PEP 695 implemented (3.12+), can pass infer_variance to typing.TypeVar + typevar = typing.TypeVar(name, *constraints, bound=bound, + covariant=covariant, contravariant=contravariant, + infer_variance=infer_variance) + else: + typevar = typing.TypeVar(name, *constraints, bound=bound, + covariant=covariant, contravariant=contravariant) + if infer_variance and (covariant or contravariant): + raise ValueError("Variance cannot be specified with infer_variance.") + typevar.__infer_variance__ = infer_variance + + _set_default(typevar, default) + _set_module(typevar) + + def _tvar_prepare_subst(alias, args): + if ( + typevar.has_default() + and alias.__parameters__.index(typevar) == len(args) + ): + args += (typevar.__default__,) + return args + + typevar.__typing_prepare_subst__ = _tvar_prepare_subst + return typevar + + def __init_subclass__(cls) -> None: + raise TypeError(f"type '{__name__}.TypeVar' is not an acceptable base type") + + +# Python 3.10+ has PEP 612 +if hasattr(typing, 'ParamSpecArgs'): + ParamSpecArgs = typing.ParamSpecArgs + ParamSpecKwargs = typing.ParamSpecKwargs +# 3.9 +else: + class _Immutable: + """Mixin to indicate that object should not be copied.""" + __slots__ = () + + def __copy__(self): + return self + + def __deepcopy__(self, memo): + return self + + class ParamSpecArgs(_Immutable): + """The args for a ParamSpec object. + + Given a ParamSpec object P, P.args is an instance of ParamSpecArgs. + + ParamSpecArgs objects have a reference back to their ParamSpec: + + P.args.__origin__ is P + + This type is meant for runtime introspection and has no special meaning to + static type checkers. + """ + def __init__(self, origin): + self.__origin__ = origin + + def __repr__(self): + return f"{self.__origin__.__name__}.args" + + def __eq__(self, other): + if not isinstance(other, ParamSpecArgs): + return NotImplemented + return self.__origin__ == other.__origin__ + + class ParamSpecKwargs(_Immutable): + """The kwargs for a ParamSpec object. + + Given a ParamSpec object P, P.kwargs is an instance of ParamSpecKwargs. + + ParamSpecKwargs objects have a reference back to their ParamSpec: + + P.kwargs.__origin__ is P + + This type is meant for runtime introspection and has no special meaning to + static type checkers. + """ + def __init__(self, origin): + self.__origin__ = origin + + def __repr__(self): + return f"{self.__origin__.__name__}.kwargs" + + def __eq__(self, other): + if not isinstance(other, ParamSpecKwargs): + return NotImplemented + return self.__origin__ == other.__origin__ + + +if _PEP_696_IMPLEMENTED: + from typing import ParamSpec + +# 3.10+ +elif hasattr(typing, 'ParamSpec'): + + # Add default parameter - PEP 696 + class ParamSpec(metaclass=_TypeVarLikeMeta): + """Parameter specification.""" + + _backported_typevarlike = typing.ParamSpec + + def __new__(cls, name, *, bound=None, + covariant=False, contravariant=False, + infer_variance=False, default=NoDefault): + if hasattr(typing, "TypeAliasType"): + # PEP 695 implemented, can pass infer_variance to typing.TypeVar + paramspec = typing.ParamSpec(name, bound=bound, + covariant=covariant, + contravariant=contravariant, + infer_variance=infer_variance) + else: + paramspec = typing.ParamSpec(name, bound=bound, + covariant=covariant, + contravariant=contravariant) + paramspec.__infer_variance__ = infer_variance + + _set_default(paramspec, default) + _set_module(paramspec) + + def _paramspec_prepare_subst(alias, args): + params = alias.__parameters__ + i = params.index(paramspec) + if i == len(args) and paramspec.has_default(): + args = [*args, paramspec.__default__] + if i >= len(args): + raise TypeError(f"Too few arguments for {alias}") + # Special case where Z[[int, str, bool]] == Z[int, str, bool] in PEP 612. + if len(params) == 1 and not typing._is_param_expr(args[0]): + assert i == 0 + args = (args,) + # Convert lists to tuples to help other libraries cache the results. + elif isinstance(args[i], list): + args = (*args[:i], tuple(args[i]), *args[i + 1:]) + return args + + paramspec.__typing_prepare_subst__ = _paramspec_prepare_subst + return paramspec + + def __init_subclass__(cls) -> None: + raise TypeError(f"type '{__name__}.ParamSpec' is not an acceptable base type") + +# 3.9 +else: + + # Inherits from list as a workaround for Callable checks in Python < 3.9.2. + class ParamSpec(list, _DefaultMixin): + """Parameter specification variable. + + Usage:: + + P = ParamSpec('P') + + Parameter specification variables exist primarily for the benefit of static + type checkers. They are used to forward the parameter types of one + callable to another callable, a pattern commonly found in higher order + functions and decorators. They are only valid when used in ``Concatenate``, + or s the first argument to ``Callable``. In Python 3.10 and higher, + they are also supported in user-defined Generics at runtime. + See class Generic for more information on generic types. An + example for annotating a decorator:: + + T = TypeVar('T') + P = ParamSpec('P') + + def add_logging(f: Callable[P, T]) -> Callable[P, T]: + '''A type-safe decorator to add logging to a function.''' + def inner(*args: P.args, **kwargs: P.kwargs) -> T: + logging.info(f'{f.__name__} was called') + return f(*args, **kwargs) + return inner + + @add_logging + def add_two(x: float, y: float) -> float: + '''Add two numbers together.''' + return x + y + + Parameter specification variables defined with covariant=True or + contravariant=True can be used to declare covariant or contravariant + generic types. These keyword arguments are valid, but their actual semantics + are yet to be decided. See PEP 612 for details. + + Parameter specification variables can be introspected. e.g.: + + P.__name__ == 'T' + P.__bound__ == None + P.__covariant__ == False + P.__contravariant__ == False + + Note that only parameter specification variables defined in global scope can + be pickled. + """ + + # Trick Generic __parameters__. + __class__ = typing.TypeVar + + @property + def args(self): + return ParamSpecArgs(self) + + @property + def kwargs(self): + return ParamSpecKwargs(self) + + def __init__(self, name, *, bound=None, covariant=False, contravariant=False, + infer_variance=False, default=NoDefault): + list.__init__(self, [self]) + self.__name__ = name + self.__covariant__ = bool(covariant) + self.__contravariant__ = bool(contravariant) + self.__infer_variance__ = bool(infer_variance) + if bound: + self.__bound__ = typing._type_check(bound, 'Bound must be a type.') + else: + self.__bound__ = None + _DefaultMixin.__init__(self, default) + + # for pickling: + def_mod = _caller() + if def_mod != 'typing_extensions': + self.__module__ = def_mod + + def __repr__(self): + if self.__infer_variance__: + prefix = '' + elif self.__covariant__: + prefix = '+' + elif self.__contravariant__: + prefix = '-' + else: + prefix = '~' + return prefix + self.__name__ + + def __hash__(self): + return object.__hash__(self) + + def __eq__(self, other): + return self is other + + def __reduce__(self): + return self.__name__ + + # Hack to get typing._type_check to pass. + def __call__(self, *args, **kwargs): + pass + + +# 3.9 +if not hasattr(typing, 'Concatenate'): + # Inherits from list as a workaround for Callable checks in Python < 3.9.2. + + # 3.9.0-1 + if not hasattr(typing, '_type_convert'): + def _type_convert(arg, module=None, *, allow_special_forms=False): + """For converting None to type(None), and strings to ForwardRef.""" + if arg is None: + return type(None) + if isinstance(arg, str): + if sys.version_info <= (3, 9, 6): + return ForwardRef(arg) + if sys.version_info <= (3, 9, 7): + return ForwardRef(arg, module=module) + return ForwardRef(arg, module=module, is_class=allow_special_forms) + return arg + else: + _type_convert = typing._type_convert + + class _ConcatenateGenericAlias(list): + + # Trick Generic into looking into this for __parameters__. + __class__ = typing._GenericAlias + + def __init__(self, origin, args): + super().__init__(args) + self.__origin__ = origin + self.__args__ = args + + def __repr__(self): + _type_repr = typing._type_repr + return (f'{_type_repr(self.__origin__)}' + f'[{", ".join(_type_repr(arg) for arg in self.__args__)}]') + + def __hash__(self): + return hash((self.__origin__, self.__args__)) + + # Hack to get typing._type_check to pass in Generic. + def __call__(self, *args, **kwargs): + pass + + @property + def __parameters__(self): + return tuple( + tp for tp in self.__args__ if isinstance(tp, (typing.TypeVar, ParamSpec)) + ) + + # 3.9 used by __getitem__ below + def copy_with(self, params): + if isinstance(params[-1], _ConcatenateGenericAlias): + params = (*params[:-1], *params[-1].__args__) + elif isinstance(params[-1], (list, tuple)): + return (*params[:-1], *params[-1]) + elif (not (params[-1] is ... or isinstance(params[-1], ParamSpec))): + raise TypeError("The last parameter to Concatenate should be a " + "ParamSpec variable or ellipsis.") + return self.__class__(self.__origin__, params) + + # 3.9; accessed during GenericAlias.__getitem__ when substituting + def __getitem__(self, args): + if self.__origin__ in (Generic, Protocol): + # Can't subscript Generic[...] or Protocol[...]. + raise TypeError(f"Cannot subscript already-subscripted {self}") + if not self.__parameters__: + raise TypeError(f"{self} is not a generic class") + + if not isinstance(args, tuple): + args = (args,) + args = _unpack_args(*(_type_convert(p) for p in args)) + params = self.__parameters__ + for param in params: + prepare = getattr(param, "__typing_prepare_subst__", None) + if prepare is not None: + args = prepare(self, args) + # 3.9 & typing.ParamSpec + elif isinstance(param, ParamSpec): + i = params.index(param) + if ( + i == len(args) + and getattr(param, '__default__', NoDefault) is not NoDefault + ): + args = [*args, param.__default__] + if i >= len(args): + raise TypeError(f"Too few arguments for {self}") + # Special case for Z[[int, str, bool]] == Z[int, str, bool] + if len(params) == 1 and not _is_param_expr(args[0]): + assert i == 0 + args = (args,) + elif ( + isinstance(args[i], list) + # 3.9 + # This class inherits from list do not convert + and not isinstance(args[i], _ConcatenateGenericAlias) + ): + args = (*args[:i], tuple(args[i]), *args[i + 1:]) + + alen = len(args) + plen = len(params) + if alen != plen: + raise TypeError( + f"Too {'many' if alen > plen else 'few'} arguments for {self};" + f" actual {alen}, expected {plen}" + ) + + subst = dict(zip(self.__parameters__, args)) + # determine new args + new_args = [] + for arg in self.__args__: + if isinstance(arg, type): + new_args.append(arg) + continue + if isinstance(arg, TypeVar): + arg = subst[arg] + if ( + (isinstance(arg, typing._GenericAlias) and _is_unpack(arg)) + or ( + hasattr(_types, "GenericAlias") + and isinstance(arg, _types.GenericAlias) + and getattr(arg, "__unpacked__", False) + ) + ): + raise TypeError(f"{arg} is not valid as type argument") + + elif isinstance(arg, + typing._GenericAlias + if not hasattr(_types, "GenericAlias") else + (typing._GenericAlias, _types.GenericAlias) + ): + subparams = arg.__parameters__ + if subparams: + subargs = tuple(subst[x] for x in subparams) + arg = arg[subargs] + new_args.append(arg) + return self.copy_with(tuple(new_args)) + +# 3.10+ +else: + _ConcatenateGenericAlias = typing._ConcatenateGenericAlias + + # 3.10 + if sys.version_info < (3, 11): + + class _ConcatenateGenericAlias(typing._ConcatenateGenericAlias, _root=True): + # needed for checks in collections.abc.Callable to accept this class + __module__ = "typing" + + def copy_with(self, params): + if isinstance(params[-1], (list, tuple)): + return (*params[:-1], *params[-1]) + if isinstance(params[-1], typing._ConcatenateGenericAlias): + params = (*params[:-1], *params[-1].__args__) + elif not (params[-1] is ... or isinstance(params[-1], ParamSpec)): + raise TypeError("The last parameter to Concatenate should be a " + "ParamSpec variable or ellipsis.") + return super(typing._ConcatenateGenericAlias, self).copy_with(params) + + def __getitem__(self, args): + value = super().__getitem__(args) + if isinstance(value, tuple) and any(_is_unpack(t) for t in value): + return tuple(_unpack_args(*(n for n in value))) + return value + + +# 3.9.2 +class _EllipsisDummy: ... + + +# <=3.10 +def _create_concatenate_alias(origin, parameters): + if parameters[-1] is ... and sys.version_info < (3, 9, 2): + # Hack: Arguments must be types, replace it with one. + parameters = (*parameters[:-1], _EllipsisDummy) + if sys.version_info >= (3, 10, 3): + concatenate = _ConcatenateGenericAlias(origin, parameters, + _typevar_types=(TypeVar, ParamSpec), + _paramspec_tvars=True) + else: + concatenate = _ConcatenateGenericAlias(origin, parameters) + if parameters[-1] is not _EllipsisDummy: + return concatenate + # Remove dummy again + concatenate.__args__ = tuple(p if p is not _EllipsisDummy else ... + for p in concatenate.__args__) + if sys.version_info < (3, 10): + # backport needs __args__ adjustment only + return concatenate + concatenate.__parameters__ = tuple(p for p in concatenate.__parameters__ + if p is not _EllipsisDummy) + return concatenate + + +# <=3.10 +@typing._tp_cache +def _concatenate_getitem(self, parameters): + if parameters == (): + raise TypeError("Cannot take a Concatenate of no types.") + if not isinstance(parameters, tuple): + parameters = (parameters,) + if not (parameters[-1] is ... or isinstance(parameters[-1], ParamSpec)): + raise TypeError("The last parameter to Concatenate should be a " + "ParamSpec variable or ellipsis.") + msg = "Concatenate[arg, ...]: each arg must be a type." + parameters = (*(typing._type_check(p, msg) for p in parameters[:-1]), + parameters[-1]) + return _create_concatenate_alias(self, parameters) + + +# 3.11+; Concatenate does not accept ellipsis in 3.10 +# Breakpoint: https://github.com/python/cpython/pull/30969 +if sys.version_info >= (3, 11): + Concatenate = typing.Concatenate +# <=3.10 +else: + @_ExtensionsSpecialForm + def Concatenate(self, parameters): + """Used in conjunction with ``ParamSpec`` and ``Callable`` to represent a + higher order function which adds, removes or transforms parameters of a + callable. + + For example:: + + Callable[Concatenate[int, P], int] + + See PEP 612 for detailed information. + """ + return _concatenate_getitem(self, parameters) + + +# 3.10+ +if hasattr(typing, 'TypeGuard'): + TypeGuard = typing.TypeGuard +# 3.9 +else: + @_ExtensionsSpecialForm + def TypeGuard(self, parameters): + """Special typing form used to annotate the return type of a user-defined + type guard function. ``TypeGuard`` only accepts a single type argument. + At runtime, functions marked this way should return a boolean. + + ``TypeGuard`` aims to benefit *type narrowing* -- a technique used by static + type checkers to determine a more precise type of an expression within a + program's code flow. Usually type narrowing is done by analyzing + conditional code flow and applying the narrowing to a block of code. The + conditional expression here is sometimes referred to as a "type guard". + + Sometimes it would be convenient to use a user-defined boolean function + as a type guard. Such a function should use ``TypeGuard[...]`` as its + return type to alert static type checkers to this intention. + + Using ``-> TypeGuard`` tells the static type checker that for a given + function: + + 1. The return value is a boolean. + 2. If the return value is ``True``, the type of its argument + is the type inside ``TypeGuard``. + + For example:: + + def is_str(val: Union[str, float]): + # "isinstance" type guard + if isinstance(val, str): + # Type of ``val`` is narrowed to ``str`` + ... + else: + # Else, type of ``val`` is narrowed to ``float``. + ... + + Strict type narrowing is not enforced -- ``TypeB`` need not be a narrower + form of ``TypeA`` (it can even be a wider form) and this may lead to + type-unsafe results. The main reason is to allow for things like + narrowing ``List[object]`` to ``List[str]`` even though the latter is not + a subtype of the former, since ``List`` is invariant. The responsibility of + writing type-safe type guards is left to the user. + + ``TypeGuard`` also works with type variables. For more information, see + PEP 647 (User-Defined Type Guards). + """ + item = typing._type_check(parameters, f'{self} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + +# 3.13+ +if hasattr(typing, 'TypeIs'): + TypeIs = typing.TypeIs +# <=3.12 +else: + @_ExtensionsSpecialForm + def TypeIs(self, parameters): + """Special typing form used to annotate the return type of a user-defined + type narrower function. ``TypeIs`` only accepts a single type argument. + At runtime, functions marked this way should return a boolean. + + ``TypeIs`` aims to benefit *type narrowing* -- a technique used by static + type checkers to determine a more precise type of an expression within a + program's code flow. Usually type narrowing is done by analyzing + conditional code flow and applying the narrowing to a block of code. The + conditional expression here is sometimes referred to as a "type guard". + + Sometimes it would be convenient to use a user-defined boolean function + as a type guard. Such a function should use ``TypeIs[...]`` as its + return type to alert static type checkers to this intention. + + Using ``-> TypeIs`` tells the static type checker that for a given + function: + + 1. The return value is a boolean. + 2. If the return value is ``True``, the type of its argument + is the intersection of the type inside ``TypeIs`` and the argument's + previously known type. + + For example:: + + def is_awaitable(val: object) -> TypeIs[Awaitable[Any]]: + return hasattr(val, '__await__') + + def f(val: Union[int, Awaitable[int]]) -> int: + if is_awaitable(val): + assert_type(val, Awaitable[int]) + else: + assert_type(val, int) + + ``TypeIs`` also works with type variables. For more information, see + PEP 742 (Narrowing types with TypeIs). + """ + item = typing._type_check(parameters, f'{self} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + +# 3.14+? +if hasattr(typing, 'TypeForm'): + TypeForm = typing.TypeForm +# <=3.13 +else: + class _TypeFormForm(_ExtensionsSpecialForm, _root=True): + # TypeForm(X) is equivalent to X but indicates to the type checker + # that the object is a TypeForm. + def __call__(self, obj, /): + return obj + + @_TypeFormForm + def TypeForm(self, parameters): + """A special form representing the value that results from the evaluation + of a type expression. This value encodes the information supplied in the + type expression, and it represents the type described by that type expression. + + When used in a type expression, TypeForm describes a set of type form objects. + It accepts a single type argument, which must be a valid type expression. + ``TypeForm[T]`` describes the set of all type form objects that represent + the type T or types that are assignable to T. + + Usage: + + def cast[T](typ: TypeForm[T], value: Any) -> T: ... + + reveal_type(cast(int, "x")) # int + + See PEP 747 for more information. + """ + item = typing._type_check(parameters, f'{self} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + + + +if hasattr(typing, "LiteralString"): # 3.11+ + LiteralString = typing.LiteralString +else: + @_SpecialForm + def LiteralString(self, params): + """Represents an arbitrary literal string. + + Example:: + + from typing_extensions import LiteralString + + def query(sql: LiteralString) -> ...: + ... + + query("SELECT * FROM table") # ok + query(f"SELECT * FROM {input()}") # not ok + + See PEP 675 for details. + + """ + raise TypeError(f"{self} is not subscriptable") + + +if hasattr(typing, "Self"): # 3.11+ + Self = typing.Self +else: + @_SpecialForm + def Self(self, params): + """Used to spell the type of "self" in classes. + + Example:: + + from typing import Self + + class ReturnsSelf: + def parse(self, data: bytes) -> Self: + ... + return self + + """ + + raise TypeError(f"{self} is not subscriptable") + + +if hasattr(typing, "Never"): # 3.11+ + Never = typing.Never +else: + @_SpecialForm + def Never(self, params): + """The bottom type, a type that has no members. + + This can be used to define a function that should never be + called, or a function that never returns:: + + from typing_extensions import Never + + def never_call_me(arg: Never) -> None: + pass + + def int_or_str(arg: int | str) -> None: + never_call_me(arg) # type checker error + match arg: + case int(): + print("It's an int") + case str(): + print("It's a str") + case _: + never_call_me(arg) # ok, arg is of type Never + + """ + + raise TypeError(f"{self} is not subscriptable") + + +if hasattr(typing, 'Required'): # 3.11+ + Required = typing.Required + NotRequired = typing.NotRequired +else: # <=3.10 + @_ExtensionsSpecialForm + def Required(self, parameters): + """A special typing construct to mark a key of a total=False TypedDict + as required. For example: + + class Movie(TypedDict, total=False): + title: Required[str] + year: int + + m = Movie( + title='The Matrix', # typechecker error if key is omitted + year=1999, + ) + + There is no runtime checking that a required key is actually provided + when instantiating a related TypedDict. + """ + item = typing._type_check(parameters, f'{self._name} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + @_ExtensionsSpecialForm + def NotRequired(self, parameters): + """A special typing construct to mark a key of a TypedDict as + potentially missing. For example: + + class Movie(TypedDict): + title: str + year: NotRequired[int] + + m = Movie( + title='The Matrix', # typechecker error if key is omitted + year=1999, + ) + """ + item = typing._type_check(parameters, f'{self._name} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + +if hasattr(typing, 'ReadOnly'): + ReadOnly = typing.ReadOnly +else: # <=3.12 + @_ExtensionsSpecialForm + def ReadOnly(self, parameters): + """A special typing construct to mark an item of a TypedDict as read-only. + + For example: + + class Movie(TypedDict): + title: ReadOnly[str] + year: int + + def mutate_movie(m: Movie) -> None: + m["year"] = 1992 # allowed + m["title"] = "The Matrix" # typechecker error + + There is no runtime checking for this property. + """ + item = typing._type_check(parameters, f'{self._name} accepts only a single type.') + return typing._GenericAlias(self, (item,)) + + +_UNPACK_DOC = """\ +Type unpack operator. + +The type unpack operator takes the child types from some container type, +such as `tuple[int, str]` or a `TypeVarTuple`, and 'pulls them out'. For +example: + + # For some generic class `Foo`: + Foo[Unpack[tuple[int, str]]] # Equivalent to Foo[int, str] + + Ts = TypeVarTuple('Ts') + # Specifies that `Bar` is generic in an arbitrary number of types. + # (Think of `Ts` as a tuple of an arbitrary number of individual + # `TypeVar`s, which the `Unpack` is 'pulling out' directly into the + # `Generic[]`.) + class Bar(Generic[Unpack[Ts]]): ... + Bar[int] # Valid + Bar[int, str] # Also valid + +From Python 3.11, this can also be done using the `*` operator: + + Foo[*tuple[int, str]] + class Bar(Generic[*Ts]): ... + +The operator can also be used along with a `TypedDict` to annotate +`**kwargs` in a function signature. For instance: + + class Movie(TypedDict): + name: str + year: int + + # This function expects two keyword arguments - *name* of type `str` and + # *year* of type `int`. + def foo(**kwargs: Unpack[Movie]): ... + +Note that there is only some runtime checking of this operator. Not +everything the runtime allows may be accepted by static type checkers. + +For more information, see PEP 646 and PEP 692. +""" + + +# PEP 692 changed the repr of Unpack[] +# Breakpoint: https://github.com/python/cpython/pull/104048 +if sys.version_info >= (3, 12): + Unpack = typing.Unpack + + def _is_unpack(obj): + return get_origin(obj) is Unpack + +else: # <=3.11 + class _UnpackSpecialForm(_ExtensionsSpecialForm, _root=True): + def __init__(self, getitem): + super().__init__(getitem) + self.__doc__ = _UNPACK_DOC + + class _UnpackAlias(typing._GenericAlias, _root=True): + if sys.version_info < (3, 11): + # needed for compatibility with Generic[Unpack[Ts]] + __class__ = typing.TypeVar + + @property + def __typing_unpacked_tuple_args__(self): + assert self.__origin__ is Unpack + assert len(self.__args__) == 1 + arg, = self.__args__ + if isinstance(arg, (typing._GenericAlias, _types.GenericAlias)): + if arg.__origin__ is not tuple: + raise TypeError("Unpack[...] must be used with a tuple type") + return arg.__args__ + return None + + @property + def __typing_is_unpacked_typevartuple__(self): + assert self.__origin__ is Unpack + assert len(self.__args__) == 1 + return isinstance(self.__args__[0], TypeVarTuple) + + def __getitem__(self, args): + if self.__typing_is_unpacked_typevartuple__: + return args + return super().__getitem__(args) + + @_UnpackSpecialForm + def Unpack(self, parameters): + item = typing._type_check(parameters, f'{self._name} accepts only a single type.') + return _UnpackAlias(self, (item,)) + + def _is_unpack(obj): + return isinstance(obj, _UnpackAlias) + + +def _unpack_args(*args): + newargs = [] + for arg in args: + subargs = getattr(arg, '__typing_unpacked_tuple_args__', None) + if subargs is not None and (not (subargs and subargs[-1] is ...)): + newargs.extend(subargs) + else: + newargs.append(arg) + return newargs + + +if _PEP_696_IMPLEMENTED: + from typing import TypeVarTuple + +elif hasattr(typing, "TypeVarTuple"): # 3.11+ + + # Add default parameter - PEP 696 + class TypeVarTuple(metaclass=_TypeVarLikeMeta): + """Type variable tuple.""" + + _backported_typevarlike = typing.TypeVarTuple + + def __new__(cls, name, *, default=NoDefault): + tvt = typing.TypeVarTuple(name) + _set_default(tvt, default) + _set_module(tvt) + + def _typevartuple_prepare_subst(alias, args): + params = alias.__parameters__ + typevartuple_index = params.index(tvt) + for param in params[typevartuple_index + 1:]: + if isinstance(param, TypeVarTuple): + raise TypeError( + f"More than one TypeVarTuple parameter in {alias}" + ) + + alen = len(args) + plen = len(params) + left = typevartuple_index + right = plen - typevartuple_index - 1 + var_tuple_index = None + fillarg = None + for k, arg in enumerate(args): + if not isinstance(arg, type): + subargs = getattr(arg, '__typing_unpacked_tuple_args__', None) + if subargs and len(subargs) == 2 and subargs[-1] is ...: + if var_tuple_index is not None: + raise TypeError( + "More than one unpacked " + "arbitrary-length tuple argument" + ) + var_tuple_index = k + fillarg = subargs[0] + if var_tuple_index is not None: + left = min(left, var_tuple_index) + right = min(right, alen - var_tuple_index - 1) + elif left + right > alen: + raise TypeError(f"Too few arguments for {alias};" + f" actual {alen}, expected at least {plen - 1}") + if left == alen - right and tvt.has_default(): + replacement = _unpack_args(tvt.__default__) + else: + replacement = args[left: alen - right] + + return ( + *args[:left], + *([fillarg] * (typevartuple_index - left)), + replacement, + *([fillarg] * (plen - right - left - typevartuple_index - 1)), + *args[alen - right:], + ) + + tvt.__typing_prepare_subst__ = _typevartuple_prepare_subst + return tvt + + def __init_subclass__(self, *args, **kwds): + raise TypeError("Cannot subclass special typing classes") + +else: # <=3.10 + class TypeVarTuple(_DefaultMixin): + """Type variable tuple. + + Usage:: + + Ts = TypeVarTuple('Ts') + + In the same way that a normal type variable is a stand-in for a single + type such as ``int``, a type variable *tuple* is a stand-in for a *tuple* + type such as ``Tuple[int, str]``. + + Type variable tuples can be used in ``Generic`` declarations. + Consider the following example:: + + class Array(Generic[*Ts]): ... + + The ``Ts`` type variable tuple here behaves like ``tuple[T1, T2]``, + where ``T1`` and ``T2`` are type variables. To use these type variables + as type parameters of ``Array``, we must *unpack* the type variable tuple using + the star operator: ``*Ts``. The signature of ``Array`` then behaves + as if we had simply written ``class Array(Generic[T1, T2]): ...``. + In contrast to ``Generic[T1, T2]``, however, ``Generic[*Shape]`` allows + us to parameterise the class with an *arbitrary* number of type parameters. + + Type variable tuples can be used anywhere a normal ``TypeVar`` can. + This includes class definitions, as shown above, as well as function + signatures and variable annotations:: + + class Array(Generic[*Ts]): + + def __init__(self, shape: Tuple[*Ts]): + self._shape: Tuple[*Ts] = shape + + def get_shape(self) -> Tuple[*Ts]: + return self._shape + + shape = (Height(480), Width(640)) + x: Array[Height, Width] = Array(shape) + y = abs(x) # Inferred type is Array[Height, Width] + z = x + x # ... is Array[Height, Width] + x.get_shape() # ... is tuple[Height, Width] + + """ + + # Trick Generic __parameters__. + __class__ = typing.TypeVar + + def __iter__(self): + yield self.__unpacked__ + + def __init__(self, name, *, default=NoDefault): + self.__name__ = name + _DefaultMixin.__init__(self, default) + + # for pickling: + def_mod = _caller() + if def_mod != 'typing_extensions': + self.__module__ = def_mod + + self.__unpacked__ = Unpack[self] + + def __repr__(self): + return self.__name__ + + def __hash__(self): + return object.__hash__(self) + + def __eq__(self, other): + return self is other + + def __reduce__(self): + return self.__name__ + + def __init_subclass__(self, *args, **kwds): + if '_root' not in kwds: + raise TypeError("Cannot subclass special typing classes") + + +if hasattr(typing, "reveal_type"): # 3.11+ + reveal_type = typing.reveal_type +else: # <=3.10 + def reveal_type(obj: T, /) -> T: + """Reveal the inferred type of a variable. + + When a static type checker encounters a call to ``reveal_type()``, + it will emit the inferred type of the argument:: + + x: int = 1 + reveal_type(x) + + Running a static type checker (e.g., ``mypy``) on this example + will produce output similar to 'Revealed type is "builtins.int"'. + + At runtime, the function prints the runtime type of the + argument and returns it unchanged. + + """ + print(f"Runtime type is {type(obj).__name__!r}", file=sys.stderr) + return obj + + +if hasattr(typing, "_ASSERT_NEVER_REPR_MAX_LENGTH"): # 3.11+ + _ASSERT_NEVER_REPR_MAX_LENGTH = typing._ASSERT_NEVER_REPR_MAX_LENGTH +else: # <=3.10 + _ASSERT_NEVER_REPR_MAX_LENGTH = 100 + + +if hasattr(typing, "assert_never"): # 3.11+ + assert_never = typing.assert_never +else: # <=3.10 + def assert_never(arg: Never, /) -> Never: + """Assert to the type checker that a line of code is unreachable. + + Example:: + + def int_or_str(arg: int | str) -> None: + match arg: + case int(): + print("It's an int") + case str(): + print("It's a str") + case _: + assert_never(arg) + + If a type checker finds that a call to assert_never() is + reachable, it will emit an error. + + At runtime, this throws an exception when called. + + """ + value = repr(arg) + if len(value) > _ASSERT_NEVER_REPR_MAX_LENGTH: + value = value[:_ASSERT_NEVER_REPR_MAX_LENGTH] + '...' + raise AssertionError(f"Expected code to be unreachable, but got: {value}") + + +# dataclass_transform exists in 3.11 but lacks the frozen_default parameter +# Breakpoint: https://github.com/python/cpython/pull/99958 +if sys.version_info >= (3, 12): # 3.12+ + dataclass_transform = typing.dataclass_transform +else: # <=3.11 + def dataclass_transform( + *, + eq_default: bool = True, + order_default: bool = False, + kw_only_default: bool = False, + frozen_default: bool = False, + field_specifiers: typing.Tuple[ + typing.Union[typing.Type[typing.Any], typing.Callable[..., typing.Any]], + ... + ] = (), + **kwargs: typing.Any, + ) -> typing.Callable[[T], T]: + """Decorator that marks a function, class, or metaclass as providing + dataclass-like behavior. + + Example: + + from typing_extensions import dataclass_transform + + _T = TypeVar("_T") + + # Used on a decorator function + @dataclass_transform() + def create_model(cls: type[_T]) -> type[_T]: + ... + return cls + + @create_model + class CustomerModel: + id: int + name: str + + # Used on a base class + @dataclass_transform() + class ModelBase: ... + + class CustomerModel(ModelBase): + id: int + name: str + + # Used on a metaclass + @dataclass_transform() + class ModelMeta(type): ... + + class ModelBase(metaclass=ModelMeta): ... + + class CustomerModel(ModelBase): + id: int + name: str + + Each of the ``CustomerModel`` classes defined in this example will now + behave similarly to a dataclass created with the ``@dataclasses.dataclass`` + decorator. For example, the type checker will synthesize an ``__init__`` + method. + + The arguments to this decorator can be used to customize this behavior: + - ``eq_default`` indicates whether the ``eq`` parameter is assumed to be + True or False if it is omitted by the caller. + - ``order_default`` indicates whether the ``order`` parameter is + assumed to be True or False if it is omitted by the caller. + - ``kw_only_default`` indicates whether the ``kw_only`` parameter is + assumed to be True or False if it is omitted by the caller. + - ``frozen_default`` indicates whether the ``frozen`` parameter is + assumed to be True or False if it is omitted by the caller. + - ``field_specifiers`` specifies a static list of supported classes + or functions that describe fields, similar to ``dataclasses.field()``. + + At runtime, this decorator records its arguments in the + ``__dataclass_transform__`` attribute on the decorated object. + + See PEP 681 for details. + + """ + def decorator(cls_or_fn): + cls_or_fn.__dataclass_transform__ = { + "eq_default": eq_default, + "order_default": order_default, + "kw_only_default": kw_only_default, + "frozen_default": frozen_default, + "field_specifiers": field_specifiers, + "kwargs": kwargs, + } + return cls_or_fn + return decorator + + +if hasattr(typing, "override"): # 3.12+ + override = typing.override +else: # <=3.11 + _F = typing.TypeVar("_F", bound=typing.Callable[..., typing.Any]) + + def override(arg: _F, /) -> _F: + """Indicate that a method is intended to override a method in a base class. + + Usage: + + class Base: + def method(self) -> None: + pass + + class Child(Base): + @override + def method(self) -> None: + super().method() + + When this decorator is applied to a method, the type checker will + validate that it overrides a method with the same name on a base class. + This helps prevent bugs that may occur when a base class is changed + without an equivalent change to a child class. + + There is no runtime checking of these properties. The decorator + sets the ``__override__`` attribute to ``True`` on the decorated object + to allow runtime introspection. + + See PEP 698 for details. + + """ + try: + arg.__override__ = True + except (AttributeError, TypeError): + # Skip the attribute silently if it is not writable. + # AttributeError happens if the object has __slots__ or a + # read-only property, TypeError if it's a builtin class. + pass + return arg + + +# Python 3.13.3+ contains a fix for the wrapped __new__ +# Breakpoint: https://github.com/python/cpython/pull/132160 +if sys.version_info >= (3, 13, 3): + deprecated = warnings.deprecated +else: + _T = typing.TypeVar("_T") + + class deprecated: + """Indicate that a class, function or overload is deprecated. + + When this decorator is applied to an object, the type checker + will generate a diagnostic on usage of the deprecated object. + + Usage: + + @deprecated("Use B instead") + class A: + pass + + @deprecated("Use g instead") + def f(): + pass + + @overload + @deprecated("int support is deprecated") + def g(x: int) -> int: ... + @overload + def g(x: str) -> int: ... + + The warning specified by *category* will be emitted at runtime + on use of deprecated objects. For functions, that happens on calls; + for classes, on instantiation and on creation of subclasses. + If the *category* is ``None``, no warning is emitted at runtime. + The *stacklevel* determines where the + warning is emitted. If it is ``1`` (the default), the warning + is emitted at the direct caller of the deprecated object; if it + is higher, it is emitted further up the stack. + Static type checker behavior is not affected by the *category* + and *stacklevel* arguments. + + The deprecation message passed to the decorator is saved in the + ``__deprecated__`` attribute on the decorated object. + If applied to an overload, the decorator + must be after the ``@overload`` decorator for the attribute to + exist on the overload as returned by ``get_overloads()``. + + See PEP 702 for details. + + """ + def __init__( + self, + message: str, + /, + *, + category: typing.Optional[typing.Type[Warning]] = DeprecationWarning, + stacklevel: int = 1, + ) -> None: + if not isinstance(message, str): + raise TypeError( + "Expected an object of type str for 'message', not " + f"{type(message).__name__!r}" + ) + self.message = message + self.category = category + self.stacklevel = stacklevel + + def __call__(self, arg: _T, /) -> _T: + # Make sure the inner functions created below don't + # retain a reference to self. + msg = self.message + category = self.category + stacklevel = self.stacklevel + if category is None: + arg.__deprecated__ = msg + return arg + elif isinstance(arg, type): + import functools + from types import MethodType + + original_new = arg.__new__ + + @functools.wraps(original_new) + def __new__(cls, /, *args, **kwargs): + if cls is arg: + warnings.warn(msg, category=category, stacklevel=stacklevel + 1) + if original_new is not object.__new__: + return original_new(cls, *args, **kwargs) + # Mirrors a similar check in object.__new__. + elif cls.__init__ is object.__init__ and (args or kwargs): + raise TypeError(f"{cls.__name__}() takes no arguments") + else: + return original_new(cls) + + arg.__new__ = staticmethod(__new__) + + original_init_subclass = arg.__init_subclass__ + # We need slightly different behavior if __init_subclass__ + # is a bound method (likely if it was implemented in Python) + if isinstance(original_init_subclass, MethodType): + original_init_subclass = original_init_subclass.__func__ + + @functools.wraps(original_init_subclass) + def __init_subclass__(*args, **kwargs): + warnings.warn(msg, category=category, stacklevel=stacklevel + 1) + return original_init_subclass(*args, **kwargs) + + arg.__init_subclass__ = classmethod(__init_subclass__) + # Or otherwise, which likely means it's a builtin such as + # object's implementation of __init_subclass__. + else: + @functools.wraps(original_init_subclass) + def __init_subclass__(*args, **kwargs): + warnings.warn(msg, category=category, stacklevel=stacklevel + 1) + return original_init_subclass(*args, **kwargs) + + arg.__init_subclass__ = __init_subclass__ + + arg.__deprecated__ = __new__.__deprecated__ = msg + __init_subclass__.__deprecated__ = msg + return arg + elif callable(arg): + import asyncio.coroutines + import functools + import inspect + + @functools.wraps(arg) + def wrapper(*args, **kwargs): + warnings.warn(msg, category=category, stacklevel=stacklevel + 1) + return arg(*args, **kwargs) + + if asyncio.coroutines.iscoroutinefunction(arg): + # Breakpoint: https://github.com/python/cpython/pull/99247 + if sys.version_info >= (3, 12): + wrapper = inspect.markcoroutinefunction(wrapper) + else: + wrapper._is_coroutine = asyncio.coroutines._is_coroutine + + arg.__deprecated__ = wrapper.__deprecated__ = msg + return wrapper + else: + raise TypeError( + "@deprecated decorator with non-None category must be applied to " + f"a class or callable, not {arg!r}" + ) + +# Breakpoint: https://github.com/python/cpython/pull/23702 +if sys.version_info < (3, 10): + def _is_param_expr(arg): + return arg is ... or isinstance( + arg, (tuple, list, ParamSpec, _ConcatenateGenericAlias) + ) +else: + def _is_param_expr(arg): + return arg is ... or isinstance( + arg, + ( + tuple, + list, + ParamSpec, + _ConcatenateGenericAlias, + typing._ConcatenateGenericAlias, + ), + ) + + +# We have to do some monkey patching to deal with the dual nature of +# Unpack/TypeVarTuple: +# - We want Unpack to be a kind of TypeVar so it gets accepted in +# Generic[Unpack[Ts]] +# - We want it to *not* be treated as a TypeVar for the purposes of +# counting generic parameters, so that when we subscript a generic, +# the runtime doesn't try to substitute the Unpack with the subscripted type. +if not hasattr(typing, "TypeVarTuple"): + def _check_generic(cls, parameters, elen=_marker): + """Check correct count for parameters of a generic cls (internal helper). + + This gives a nice error message in case of count mismatch. + """ + # If substituting a single ParamSpec with multiple arguments + # we do not check the count + if (inspect.isclass(cls) and issubclass(cls, typing.Generic) + and len(cls.__parameters__) == 1 + and isinstance(cls.__parameters__[0], ParamSpec) + and parameters + and not _is_param_expr(parameters[0]) + ): + # Generic modifies parameters variable, but here we cannot do this + return + + if not elen: + raise TypeError(f"{cls} is not a generic class") + if elen is _marker: + if not hasattr(cls, "__parameters__") or not cls.__parameters__: + raise TypeError(f"{cls} is not a generic class") + elen = len(cls.__parameters__) + alen = len(parameters) + if alen != elen: + expect_val = elen + if hasattr(cls, "__parameters__"): + parameters = [p for p in cls.__parameters__ if not _is_unpack(p)] + num_tv_tuples = sum(isinstance(p, TypeVarTuple) for p in parameters) + if (num_tv_tuples > 0) and (alen >= elen - num_tv_tuples): + return + + # deal with TypeVarLike defaults + # required TypeVarLikes cannot appear after a defaulted one. + if alen < elen: + # since we validate TypeVarLike default in _collect_type_vars + # or _collect_parameters we can safely check parameters[alen] + if ( + getattr(parameters[alen], '__default__', NoDefault) + is not NoDefault + ): + return + + num_default_tv = sum(getattr(p, '__default__', NoDefault) + is not NoDefault for p in parameters) + + elen -= num_default_tv + + expect_val = f"at least {elen}" + + # Breakpoint: https://github.com/python/cpython/pull/27515 + things = "arguments" if sys.version_info >= (3, 10) else "parameters" + raise TypeError(f"Too {'many' if alen > elen else 'few'} {things}" + f" for {cls}; actual {alen}, expected {expect_val}") +else: + # Python 3.11+ + + def _check_generic(cls, parameters, elen): + """Check correct count for parameters of a generic cls (internal helper). + + This gives a nice error message in case of count mismatch. + """ + if not elen: + raise TypeError(f"{cls} is not a generic class") + alen = len(parameters) + if alen != elen: + expect_val = elen + if hasattr(cls, "__parameters__"): + parameters = [p for p in cls.__parameters__ if not _is_unpack(p)] + + # deal with TypeVarLike defaults + # required TypeVarLikes cannot appear after a defaulted one. + if alen < elen: + # since we validate TypeVarLike default in _collect_type_vars + # or _collect_parameters we can safely check parameters[alen] + if ( + getattr(parameters[alen], '__default__', NoDefault) + is not NoDefault + ): + return + + num_default_tv = sum(getattr(p, '__default__', NoDefault) + is not NoDefault for p in parameters) + + elen -= num_default_tv + + expect_val = f"at least {elen}" + + raise TypeError(f"Too {'many' if alen > elen else 'few'} arguments" + f" for {cls}; actual {alen}, expected {expect_val}") + +if not _PEP_696_IMPLEMENTED: + typing._check_generic = _check_generic + + +def _has_generic_or_protocol_as_origin() -> bool: + try: + frame = sys._getframe(2) + # - Catch AttributeError: not all Python implementations have sys._getframe() + # - Catch ValueError: maybe we're called from an unexpected module + # and the call stack isn't deep enough + except (AttributeError, ValueError): + return False # err on the side of leniency + else: + # If we somehow get invoked from outside typing.py, + # also err on the side of leniency + if frame.f_globals.get("__name__") != "typing": + return False + origin = frame.f_locals.get("origin") + # Cannot use "in" because origin may be an object with a buggy __eq__ that + # throws an error. + return origin is typing.Generic or origin is Protocol or origin is typing.Protocol + + +_TYPEVARTUPLE_TYPES = {TypeVarTuple, getattr(typing, "TypeVarTuple", None)} + + +def _is_unpacked_typevartuple(x) -> bool: + if get_origin(x) is not Unpack: + return False + args = get_args(x) + return ( + bool(args) + and len(args) == 1 + and type(args[0]) in _TYPEVARTUPLE_TYPES + ) + + +# Python 3.11+ _collect_type_vars was renamed to _collect_parameters +if hasattr(typing, '_collect_type_vars'): + def _collect_type_vars(types, typevar_types=None): + """Collect all type variable contained in types in order of + first appearance (lexicographic order). For example:: + + _collect_type_vars((T, List[S, T])) == (T, S) + """ + if typevar_types is None: + typevar_types = typing.TypeVar + tvars = [] + + # A required TypeVarLike cannot appear after a TypeVarLike with a default + # if it was a direct call to `Generic[]` or `Protocol[]` + enforce_default_ordering = _has_generic_or_protocol_as_origin() + default_encountered = False + + # Also, a TypeVarLike with a default cannot appear after a TypeVarTuple + type_var_tuple_encountered = False + + for t in types: + if _is_unpacked_typevartuple(t): + type_var_tuple_encountered = True + elif ( + isinstance(t, typevar_types) and not isinstance(t, _UnpackAlias) + and t not in tvars + ): + if enforce_default_ordering: + has_default = getattr(t, '__default__', NoDefault) is not NoDefault + if has_default: + if type_var_tuple_encountered: + raise TypeError('Type parameter with a default' + ' follows TypeVarTuple') + default_encountered = True + elif default_encountered: + raise TypeError(f'Type parameter {t!r} without a default' + ' follows type parameter with a default') + + tvars.append(t) + if _should_collect_from_parameters(t): + tvars.extend([t for t in t.__parameters__ if t not in tvars]) + elif isinstance(t, tuple): + # Collect nested type_vars + # tuple wrapped by _prepare_paramspec_params(cls, params) + for x in t: + for collected in _collect_type_vars([x]): + if collected not in tvars: + tvars.append(collected) + return tuple(tvars) + + typing._collect_type_vars = _collect_type_vars +else: + def _collect_parameters(args): + """Collect all type variables and parameter specifications in args + in order of first appearance (lexicographic order). + + For example:: + + assert _collect_parameters((T, Callable[P, T])) == (T, P) + """ + parameters = [] + + # A required TypeVarLike cannot appear after a TypeVarLike with default + # if it was a direct call to `Generic[]` or `Protocol[]` + enforce_default_ordering = _has_generic_or_protocol_as_origin() + default_encountered = False + + # Also, a TypeVarLike with a default cannot appear after a TypeVarTuple + type_var_tuple_encountered = False + + for t in args: + if isinstance(t, type): + # We don't want __parameters__ descriptor of a bare Python class. + pass + elif isinstance(t, tuple): + # `t` might be a tuple, when `ParamSpec` is substituted with + # `[T, int]`, or `[int, *Ts]`, etc. + for x in t: + for collected in _collect_parameters([x]): + if collected not in parameters: + parameters.append(collected) + elif hasattr(t, '__typing_subst__'): + if t not in parameters: + if enforce_default_ordering: + has_default = ( + getattr(t, '__default__', NoDefault) is not NoDefault + ) + + if type_var_tuple_encountered and has_default: + raise TypeError('Type parameter with a default' + ' follows TypeVarTuple') + + if has_default: + default_encountered = True + elif default_encountered: + raise TypeError(f'Type parameter {t!r} without a default' + ' follows type parameter with a default') + + parameters.append(t) + else: + if _is_unpacked_typevartuple(t): + type_var_tuple_encountered = True + for x in getattr(t, '__parameters__', ()): + if x not in parameters: + parameters.append(x) + + return tuple(parameters) + + if not _PEP_696_IMPLEMENTED: + typing._collect_parameters = _collect_parameters + +# Backport typing.NamedTuple as it exists in Python 3.13. +# In 3.11, the ability to define generic `NamedTuple`s was supported. +# This was explicitly disallowed in 3.9-3.10, and only half-worked in <=3.8. +# On 3.12, we added __orig_bases__ to call-based NamedTuples +# On 3.13, we deprecated kwargs-based NamedTuples +# Breakpoint: https://github.com/python/cpython/pull/105609 +if sys.version_info >= (3, 13): + NamedTuple = typing.NamedTuple +else: + def _make_nmtuple(name, types, module, defaults=()): + fields = [n for n, t in types] + annotations = {n: typing._type_check(t, f"field {n} annotation must be a type") + for n, t in types} + nm_tpl = collections.namedtuple(name, fields, + defaults=defaults, module=module) + nm_tpl.__annotations__ = nm_tpl.__new__.__annotations__ = annotations + return nm_tpl + + _prohibited_namedtuple_fields = typing._prohibited + _special_namedtuple_fields = frozenset({'__module__', '__name__', '__annotations__'}) + + class _NamedTupleMeta(type): + def __new__(cls, typename, bases, ns): + assert _NamedTuple in bases + for base in bases: + if base is not _NamedTuple and base is not typing.Generic: + raise TypeError( + 'can only inherit from a NamedTuple type and Generic') + bases = tuple(tuple if base is _NamedTuple else base for base in bases) + if "__annotations__" in ns: + types = ns["__annotations__"] + elif "__annotate__" in ns: + # TODO: Use inspect.VALUE here, and make the annotations lazily evaluated + types = ns["__annotate__"](1) + else: + types = {} + default_names = [] + for field_name in types: + if field_name in ns: + default_names.append(field_name) + elif default_names: + raise TypeError(f"Non-default namedtuple field {field_name} " + f"cannot follow default field" + f"{'s' if len(default_names) > 1 else ''} " + f"{', '.join(default_names)}") + nm_tpl = _make_nmtuple( + typename, types.items(), + defaults=[ns[n] for n in default_names], + module=ns['__module__'] + ) + nm_tpl.__bases__ = bases + if typing.Generic in bases: + if hasattr(typing, '_generic_class_getitem'): # 3.12+ + nm_tpl.__class_getitem__ = classmethod(typing._generic_class_getitem) + else: + class_getitem = typing.Generic.__class_getitem__.__func__ + nm_tpl.__class_getitem__ = classmethod(class_getitem) + # update from user namespace without overriding special namedtuple attributes + for key, val in ns.items(): + if key in _prohibited_namedtuple_fields: + raise AttributeError("Cannot overwrite NamedTuple attribute " + key) + elif key not in _special_namedtuple_fields: + if key not in nm_tpl._fields: + setattr(nm_tpl, key, ns[key]) + try: + set_name = type(val).__set_name__ + except AttributeError: + pass + else: + try: + set_name(val, nm_tpl, key) + except BaseException as e: + msg = ( + f"Error calling __set_name__ on {type(val).__name__!r} " + f"instance {key!r} in {typename!r}" + ) + # BaseException.add_note() existed on py311, + # but the __set_name__ machinery didn't start + # using add_note() until py312. + # Making sure exceptions are raised in the same way + # as in "normal" classes seems most important here. + # Breakpoint: https://github.com/python/cpython/pull/95915 + if sys.version_info >= (3, 12): + e.add_note(msg) + raise + else: + raise RuntimeError(msg) from e + + if typing.Generic in bases: + nm_tpl.__init_subclass__() + return nm_tpl + + _NamedTuple = type.__new__(_NamedTupleMeta, 'NamedTuple', (), {}) + + def _namedtuple_mro_entries(bases): + assert NamedTuple in bases + return (_NamedTuple,) + + def NamedTuple(typename, fields=_marker, /, **kwargs): + """Typed version of namedtuple. + + Usage:: + + class Employee(NamedTuple): + name: str + id: int + + This is equivalent to:: + + Employee = collections.namedtuple('Employee', ['name', 'id']) + + The resulting class has an extra __annotations__ attribute, giving a + dict that maps field names to types. (The field names are also in + the _fields attribute, which is part of the namedtuple API.) + An alternative equivalent functional syntax is also accepted:: + + Employee = NamedTuple('Employee', [('name', str), ('id', int)]) + """ + if fields is _marker: + if kwargs: + deprecated_thing = "Creating NamedTuple classes using keyword arguments" + deprecation_msg = ( + "{name} is deprecated and will be disallowed in Python {remove}. " + "Use the class-based or functional syntax instead." + ) + else: + deprecated_thing = "Failing to pass a value for the 'fields' parameter" + example = f"`{typename} = NamedTuple({typename!r}, [])`" + deprecation_msg = ( + "{name} is deprecated and will be disallowed in Python {remove}. " + "To create a NamedTuple class with 0 fields " + "using the functional syntax, " + "pass an empty list, e.g. " + ) + example + "." + elif fields is None: + if kwargs: + raise TypeError( + "Cannot pass `None` as the 'fields' parameter " + "and also specify fields using keyword arguments" + ) + else: + deprecated_thing = "Passing `None` as the 'fields' parameter" + example = f"`{typename} = NamedTuple({typename!r}, [])`" + deprecation_msg = ( + "{name} is deprecated and will be disallowed in Python {remove}. " + "To create a NamedTuple class with 0 fields " + "using the functional syntax, " + "pass an empty list, e.g. " + ) + example + "." + elif kwargs: + raise TypeError("Either list of fields or keywords" + " can be provided to NamedTuple, not both") + if fields is _marker or fields is None: + warnings.warn( + deprecation_msg.format(name=deprecated_thing, remove="3.15"), + DeprecationWarning, + stacklevel=2, + ) + fields = kwargs.items() + nt = _make_nmtuple(typename, fields, module=_caller()) + nt.__orig_bases__ = (NamedTuple,) + return nt + + NamedTuple.__mro_entries__ = _namedtuple_mro_entries + + +if hasattr(collections.abc, "Buffer"): + Buffer = collections.abc.Buffer +else: + class Buffer(abc.ABC): # noqa: B024 + """Base class for classes that implement the buffer protocol. + + The buffer protocol allows Python objects to expose a low-level + memory buffer interface. Before Python 3.12, it is not possible + to implement the buffer protocol in pure Python code, or even + to check whether a class implements the buffer protocol. In + Python 3.12 and higher, the ``__buffer__`` method allows access + to the buffer protocol from Python code, and the + ``collections.abc.Buffer`` ABC allows checking whether a class + implements the buffer protocol. + + To indicate support for the buffer protocol in earlier versions, + inherit from this ABC, either in a stub file or at runtime, + or use ABC registration. This ABC provides no methods, because + there is no Python-accessible methods shared by pre-3.12 buffer + classes. It is useful primarily for static checks. + + """ + + # As a courtesy, register the most common stdlib buffer classes. + Buffer.register(memoryview) + Buffer.register(bytearray) + Buffer.register(bytes) + + +# Backport of types.get_original_bases, available on 3.12+ in CPython +if hasattr(_types, "get_original_bases"): + get_original_bases = _types.get_original_bases +else: + def get_original_bases(cls, /): + """Return the class's "original" bases prior to modification by `__mro_entries__`. + + Examples:: + + from typing import TypeVar, Generic + from typing_extensions import NamedTuple, TypedDict + + T = TypeVar("T") + class Foo(Generic[T]): ... + class Bar(Foo[int], float): ... + class Baz(list[str]): ... + Eggs = NamedTuple("Eggs", [("a", int), ("b", str)]) + Spam = TypedDict("Spam", {"a": int, "b": str}) + + assert get_original_bases(Bar) == (Foo[int], float) + assert get_original_bases(Baz) == (list[str],) + assert get_original_bases(Eggs) == (NamedTuple,) + assert get_original_bases(Spam) == (TypedDict,) + assert get_original_bases(int) == (object,) + """ + try: + return cls.__dict__.get("__orig_bases__", cls.__bases__) + except AttributeError: + raise TypeError( + f'Expected an instance of type, not {type(cls).__name__!r}' + ) from None + + +# NewType is a class on Python 3.10+, making it pickleable +# The error message for subclassing instances of NewType was improved on 3.11+ +# Breakpoint: https://github.com/python/cpython/pull/30268 +if sys.version_info >= (3, 11): + NewType = typing.NewType +else: + class NewType: + """NewType creates simple unique types with almost zero + runtime overhead. NewType(name, tp) is considered a subtype of tp + by static type checkers. At runtime, NewType(name, tp) returns + a dummy callable that simply returns its argument. Usage:: + UserId = NewType('UserId', int) + def name_by_id(user_id: UserId) -> str: + ... + UserId('user') # Fails type check + name_by_id(42) # Fails type check + name_by_id(UserId(42)) # OK + num = UserId(5) + 1 # type: int + """ + + def __call__(self, obj, /): + return obj + + def __init__(self, name, tp): + self.__qualname__ = name + if '.' in name: + name = name.rpartition('.')[-1] + self.__name__ = name + self.__supertype__ = tp + def_mod = _caller() + if def_mod != 'typing_extensions': + self.__module__ = def_mod + + def __mro_entries__(self, bases): + # We defined __mro_entries__ to get a better error message + # if a user attempts to subclass a NewType instance. bpo-46170 + supercls_name = self.__name__ + + class Dummy: + def __init_subclass__(cls): + subcls_name = cls.__name__ + raise TypeError( + f"Cannot subclass an instance of NewType. " + f"Perhaps you were looking for: " + f"`{subcls_name} = NewType({subcls_name!r}, {supercls_name})`" + ) + + return (Dummy,) + + def __repr__(self): + return f'{self.__module__}.{self.__qualname__}' + + def __reduce__(self): + return self.__qualname__ + + # Breakpoint: https://github.com/python/cpython/pull/21515 + if sys.version_info >= (3, 10): + # PEP 604 methods + # It doesn't make sense to have these methods on Python <3.10 + + def __or__(self, other): + return typing.Union[self, other] + + def __ror__(self, other): + return typing.Union[other, self] + + +# Breakpoint: https://github.com/python/cpython/pull/124795 +if sys.version_info >= (3, 14): + TypeAliasType = typing.TypeAliasType +# <=3.13 +else: + # Breakpoint: https://github.com/python/cpython/pull/103764 + if sys.version_info >= (3, 12): + # 3.12-3.13 + def _is_unionable(obj): + """Corresponds to is_unionable() in unionobject.c in CPython.""" + return obj is None or isinstance(obj, ( + type, + _types.GenericAlias, + _types.UnionType, + typing.TypeAliasType, + TypeAliasType, + )) + else: + # <=3.11 + def _is_unionable(obj): + """Corresponds to is_unionable() in unionobject.c in CPython.""" + return obj is None or isinstance(obj, ( + type, + _types.GenericAlias, + _types.UnionType, + TypeAliasType, + )) + + if sys.version_info < (3, 10): + # Copied and pasted from https://github.com/python/cpython/blob/986a4e1b6fcae7fe7a1d0a26aea446107dd58dd2/Objects/genericaliasobject.c#L568-L582, + # so that we emulate the behaviour of `types.GenericAlias` + # on the latest versions of CPython + _ATTRIBUTE_DELEGATION_EXCLUSIONS = frozenset({ + "__class__", + "__bases__", + "__origin__", + "__args__", + "__unpacked__", + "__parameters__", + "__typing_unpacked_tuple_args__", + "__mro_entries__", + "__reduce_ex__", + "__reduce__", + "__copy__", + "__deepcopy__", + }) + + class _TypeAliasGenericAlias(typing._GenericAlias, _root=True): + def __getattr__(self, attr): + if attr in _ATTRIBUTE_DELEGATION_EXCLUSIONS: + return object.__getattr__(self, attr) + return getattr(self.__origin__, attr) + + + class TypeAliasType: + """Create named, parameterized type aliases. + + This provides a backport of the new `type` statement in Python 3.12: + + type ListOrSet[T] = list[T] | set[T] + + is equivalent to: + + T = TypeVar("T") + ListOrSet = TypeAliasType("ListOrSet", list[T] | set[T], type_params=(T,)) + + The name ListOrSet can then be used as an alias for the type it refers to. + + The type_params argument should contain all the type parameters used + in the value of the type alias. If the alias is not generic, this + argument is omitted. + + Static type checkers should only support type aliases declared using + TypeAliasType that follow these rules: + + - The first argument (the name) must be a string literal. + - The TypeAliasType instance must be immediately assigned to a variable + of the same name. (For example, 'X = TypeAliasType("Y", int)' is invalid, + as is 'X, Y = TypeAliasType("X", int), TypeAliasType("Y", int)'). + + """ + + def __init__(self, name: str, value, *, type_params=()): + if not isinstance(name, str): + raise TypeError("TypeAliasType name must be a string") + if not isinstance(type_params, tuple): + raise TypeError("type_params must be a tuple") + self.__value__ = value + self.__type_params__ = type_params + + default_value_encountered = False + parameters = [] + for type_param in type_params: + if ( + not isinstance(type_param, (TypeVar, TypeVarTuple, ParamSpec)) + # <=3.11 + # Unpack Backport passes isinstance(type_param, TypeVar) + or _is_unpack(type_param) + ): + raise TypeError(f"Expected a type param, got {type_param!r}") + has_default = ( + getattr(type_param, '__default__', NoDefault) is not NoDefault + ) + if default_value_encountered and not has_default: + raise TypeError(f"non-default type parameter '{type_param!r}'" + " follows default type parameter") + if has_default: + default_value_encountered = True + if isinstance(type_param, TypeVarTuple): + parameters.extend(type_param) + else: + parameters.append(type_param) + self.__parameters__ = tuple(parameters) + def_mod = _caller() + if def_mod != 'typing_extensions': + self.__module__ = def_mod + # Setting this attribute closes the TypeAliasType from further modification + self.__name__ = name + + def __setattr__(self, name: str, value: object, /) -> None: + if hasattr(self, "__name__"): + self._raise_attribute_error(name) + super().__setattr__(name, value) + + def __delattr__(self, name: str, /) -> Never: + self._raise_attribute_error(name) + + def _raise_attribute_error(self, name: str) -> Never: + # Match the Python 3.12 error messages exactly + if name == "__name__": + raise AttributeError("readonly attribute") + elif name in {"__value__", "__type_params__", "__parameters__", "__module__"}: + raise AttributeError( + f"attribute '{name}' of 'typing.TypeAliasType' objects " + "is not writable" + ) + else: + raise AttributeError( + f"'typing.TypeAliasType' object has no attribute '{name}'" + ) + + def __repr__(self) -> str: + return self.__name__ + + if sys.version_info < (3, 11): + def _check_single_param(self, param, recursion=0): + # Allow [], [int], [int, str], [int, ...], [int, T] + if param is ...: + return ... + if param is None: + return None + # Note in <= 3.9 _ConcatenateGenericAlias inherits from list + if isinstance(param, list) and recursion == 0: + return [self._check_single_param(arg, recursion+1) + for arg in param] + return typing._type_check( + param, f'Subscripting {self.__name__} requires a type.' + ) + + def _check_parameters(self, parameters): + if sys.version_info < (3, 11): + return tuple( + self._check_single_param(item) + for item in parameters + ) + return tuple(typing._type_check( + item, f'Subscripting {self.__name__} requires a type.' + ) + for item in parameters + ) + + def __getitem__(self, parameters): + if not self.__type_params__: + raise TypeError("Only generic type aliases are subscriptable") + if not isinstance(parameters, tuple): + parameters = (parameters,) + # Using 3.9 here will create problems with Concatenate + if sys.version_info >= (3, 10): + return _types.GenericAlias(self, parameters) + type_vars = _collect_type_vars(parameters) + parameters = self._check_parameters(parameters) + alias = _TypeAliasGenericAlias(self, parameters) + # alias.__parameters__ is not complete if Concatenate is present + # as it is converted to a list from which no parameters are extracted. + if alias.__parameters__ != type_vars: + alias.__parameters__ = type_vars + return alias + + def __reduce__(self): + return self.__name__ + + def __init_subclass__(cls, *args, **kwargs): + raise TypeError( + "type 'typing_extensions.TypeAliasType' is not an acceptable base type" + ) + + # The presence of this method convinces typing._type_check + # that TypeAliasTypes are types. + def __call__(self): + raise TypeError("Type alias is not callable") + + # Breakpoint: https://github.com/python/cpython/pull/21515 + if sys.version_info >= (3, 10): + def __or__(self, right): + # For forward compatibility with 3.12, reject Unions + # that are not accepted by the built-in Union. + if not _is_unionable(right): + return NotImplemented + return typing.Union[self, right] + + def __ror__(self, left): + if not _is_unionable(left): + return NotImplemented + return typing.Union[left, self] + + +if hasattr(typing, "is_protocol"): + is_protocol = typing.is_protocol + get_protocol_members = typing.get_protocol_members +else: + def is_protocol(tp: type, /) -> bool: + """Return True if the given type is a Protocol. + + Example:: + + >>> from typing_extensions import Protocol, is_protocol + >>> class P(Protocol): + ... def a(self) -> str: ... + ... b: int + >>> is_protocol(P) + True + >>> is_protocol(int) + False + """ + return ( + isinstance(tp, type) + and getattr(tp, '_is_protocol', False) + and tp is not Protocol + and tp is not typing.Protocol + ) + + def get_protocol_members(tp: type, /) -> typing.FrozenSet[str]: + """Return the set of members defined in a Protocol. + + Example:: + + >>> from typing_extensions import Protocol, get_protocol_members + >>> class P(Protocol): + ... def a(self) -> str: ... + ... b: int + >>> get_protocol_members(P) + frozenset({'a', 'b'}) + + Raise a TypeError for arguments that are not Protocols. + """ + if not is_protocol(tp): + raise TypeError(f'{tp!r} is not a Protocol') + if hasattr(tp, '__protocol_attrs__'): + return frozenset(tp.__protocol_attrs__) + return frozenset(_get_protocol_attrs(tp)) + + +if hasattr(typing, "Doc"): + Doc = typing.Doc +else: + class Doc: + """Define the documentation of a type annotation using ``Annotated``, to be + used in class attributes, function and method parameters, return values, + and variables. + + The value should be a positional-only string literal to allow static tools + like editors and documentation generators to use it. + + This complements docstrings. + + The string value passed is available in the attribute ``documentation``. + + Example:: + + >>> from typing_extensions import Annotated, Doc + >>> def hi(to: Annotated[str, Doc("Who to say hi to")]) -> None: ... + """ + def __init__(self, documentation: str, /) -> None: + self.documentation = documentation + + def __repr__(self) -> str: + return f"Doc({self.documentation!r})" + + def __hash__(self) -> int: + return hash(self.documentation) + + def __eq__(self, other: object) -> bool: + if not isinstance(other, Doc): + return NotImplemented + return self.documentation == other.documentation + + +_CapsuleType = getattr(_types, "CapsuleType", None) + +if _CapsuleType is None: + try: + import _socket + except ImportError: + pass + else: + _CAPI = getattr(_socket, "CAPI", None) + if _CAPI is not None: + _CapsuleType = type(_CAPI) + +if _CapsuleType is not None: + CapsuleType = _CapsuleType + __all__.append("CapsuleType") + + +if sys.version_info >= (3, 14): + from annotationlib import Format, get_annotations +else: + # Available since Python 3.14.0a3 + # PR: https://github.com/python/cpython/pull/124415 + class Format(enum.IntEnum): + VALUE = 1 + VALUE_WITH_FAKE_GLOBALS = 2 + FORWARDREF = 3 + STRING = 4 + + # Available since Python 3.14.0a1 + # PR: https://github.com/python/cpython/pull/119891 + def get_annotations(obj, *, globals=None, locals=None, eval_str=False, + format=Format.VALUE): + """Compute the annotations dict for an object. + + obj may be a callable, class, or module. + Passing in an object of any other type raises TypeError. + + Returns a dict. get_annotations() returns a new dict every time + it's called; calling it twice on the same object will return two + different but equivalent dicts. + + This is a backport of `inspect.get_annotations`, which has been + in the standard library since Python 3.10. See the standard library + documentation for more: + + https://docs.python.org/3/library/inspect.html#inspect.get_annotations + + This backport adds the *format* argument introduced by PEP 649. The + three formats supported are: + * VALUE: the annotations are returned as-is. This is the default and + it is compatible with the behavior on previous Python versions. + * FORWARDREF: return annotations as-is if possible, but replace any + undefined names with ForwardRef objects. The implementation proposed by + PEP 649 relies on language changes that cannot be backported; the + typing-extensions implementation simply returns the same result as VALUE. + * STRING: return annotations as strings, in a format close to the original + source. Again, this behavior cannot be replicated directly in a backport. + As an approximation, typing-extensions retrieves the annotations under + VALUE semantics and then stringifies them. + + The purpose of this backport is to allow users who would like to use + FORWARDREF or STRING semantics once PEP 649 is implemented, but who also + want to support earlier Python versions, to simply write: + + typing_extensions.get_annotations(obj, format=Format.FORWARDREF) + + """ + format = Format(format) + if format is Format.VALUE_WITH_FAKE_GLOBALS: + raise ValueError( + "The VALUE_WITH_FAKE_GLOBALS format is for internal use only" + ) + + if eval_str and format is not Format.VALUE: + raise ValueError("eval_str=True is only supported with format=Format.VALUE") + + if isinstance(obj, type): + # class + obj_dict = getattr(obj, '__dict__', None) + if obj_dict and hasattr(obj_dict, 'get'): + ann = obj_dict.get('__annotations__', None) + if isinstance(ann, _types.GetSetDescriptorType): + ann = None + else: + ann = None + + obj_globals = None + module_name = getattr(obj, '__module__', None) + if module_name: + module = sys.modules.get(module_name, None) + if module: + obj_globals = getattr(module, '__dict__', None) + obj_locals = dict(vars(obj)) + unwrap = obj + elif isinstance(obj, _types.ModuleType): + # module + ann = getattr(obj, '__annotations__', None) + obj_globals = obj.__dict__ + obj_locals = None + unwrap = None + elif callable(obj): + # this includes types.Function, types.BuiltinFunctionType, + # types.BuiltinMethodType, functools.partial, functools.singledispatch, + # "class funclike" from Lib/test/test_inspect... on and on it goes. + ann = getattr(obj, '__annotations__', None) + obj_globals = getattr(obj, '__globals__', None) + obj_locals = None + unwrap = obj + elif hasattr(obj, '__annotations__'): + ann = obj.__annotations__ + obj_globals = obj_locals = unwrap = None + else: + raise TypeError(f"{obj!r} is not a module, class, or callable.") + + if ann is None: + return {} + + if not isinstance(ann, dict): + raise ValueError(f"{obj!r}.__annotations__ is neither a dict nor None") + + if not ann: + return {} + + if not eval_str: + if format is Format.STRING: + return { + key: value if isinstance(value, str) else typing._type_repr(value) + for key, value in ann.items() + } + return dict(ann) + + if unwrap is not None: + while True: + if hasattr(unwrap, '__wrapped__'): + unwrap = unwrap.__wrapped__ + continue + if isinstance(unwrap, functools.partial): + unwrap = unwrap.func + continue + break + if hasattr(unwrap, "__globals__"): + obj_globals = unwrap.__globals__ + + if globals is None: + globals = obj_globals + if locals is None: + locals = obj_locals or {} + + # "Inject" type parameters into the local namespace + # (unless they are shadowed by assignments *in* the local namespace), + # as a way of emulating annotation scopes when calling `eval()` + if type_params := getattr(obj, "__type_params__", ()): + locals = {param.__name__: param for param in type_params} | locals + + return_value = {key: + value if not isinstance(value, str) else eval(value, globals, locals) + for key, value in ann.items() } + return return_value + + +if hasattr(typing, "evaluate_forward_ref"): + evaluate_forward_ref = typing.evaluate_forward_ref +else: + # Implements annotationlib.ForwardRef.evaluate + def _eval_with_owner( + forward_ref, *, owner=None, globals=None, locals=None, type_params=None + ): + if forward_ref.__forward_evaluated__: + return forward_ref.__forward_value__ + if getattr(forward_ref, "__cell__", None) is not None: + try: + value = forward_ref.__cell__.cell_contents + except ValueError: + pass + else: + forward_ref.__forward_evaluated__ = True + forward_ref.__forward_value__ = value + return value + if owner is None: + owner = getattr(forward_ref, "__owner__", None) + + if ( + globals is None + and getattr(forward_ref, "__forward_module__", None) is not None + ): + globals = getattr( + sys.modules.get(forward_ref.__forward_module__, None), "__dict__", None + ) + if globals is None: + globals = getattr(forward_ref, "__globals__", None) + if globals is None: + if isinstance(owner, type): + module_name = getattr(owner, "__module__", None) + if module_name: + module = sys.modules.get(module_name, None) + if module: + globals = getattr(module, "__dict__", None) + elif isinstance(owner, _types.ModuleType): + globals = getattr(owner, "__dict__", None) + elif callable(owner): + globals = getattr(owner, "__globals__", None) + + # If we pass None to eval() below, the globals of this module are used. + if globals is None: + globals = {} + + if locals is None: + locals = {} + if isinstance(owner, type): + locals.update(vars(owner)) + + if type_params is None and owner is not None: + # "Inject" type parameters into the local namespace + # (unless they are shadowed by assignments *in* the local namespace), + # as a way of emulating annotation scopes when calling `eval()` + type_params = getattr(owner, "__type_params__", None) + + # Type parameters exist in their own scope, which is logically + # between the locals and the globals. We simulate this by adding + # them to the globals. + if type_params is not None: + globals = dict(globals) + for param in type_params: + globals[param.__name__] = param + + arg = forward_ref.__forward_arg__ + if arg.isidentifier() and not keyword.iskeyword(arg): + if arg in locals: + value = locals[arg] + elif arg in globals: + value = globals[arg] + elif hasattr(builtins, arg): + return getattr(builtins, arg) + else: + raise NameError(arg) + else: + code = forward_ref.__forward_code__ + value = eval(code, globals, locals) + forward_ref.__forward_evaluated__ = True + forward_ref.__forward_value__ = value + return value + + def evaluate_forward_ref( + forward_ref, + *, + owner=None, + globals=None, + locals=None, + type_params=None, + format=None, + _recursive_guard=frozenset(), + ): + """Evaluate a forward reference as a type hint. + + This is similar to calling the ForwardRef.evaluate() method, + but unlike that method, evaluate_forward_ref() also: + + * Recursively evaluates forward references nested within the type hint. + * Rejects certain objects that are not valid type hints. + * Replaces type hints that evaluate to None with types.NoneType. + * Supports the *FORWARDREF* and *STRING* formats. + + *forward_ref* must be an instance of ForwardRef. *owner*, if given, + should be the object that holds the annotations that the forward reference + derived from, such as a module, class object, or function. It is used to + infer the namespaces to use for looking up names. *globals* and *locals* + can also be explicitly given to provide the global and local namespaces. + *type_params* is a tuple of type parameters that are in scope when + evaluating the forward reference. This parameter must be provided (though + it may be an empty tuple) if *owner* is not given and the forward reference + does not already have an owner set. *format* specifies the format of the + annotation and is a member of the annotationlib.Format enum. + + """ + if format == Format.STRING: + return forward_ref.__forward_arg__ + if forward_ref.__forward_arg__ in _recursive_guard: + return forward_ref + + # Evaluate the forward reference + try: + value = _eval_with_owner( + forward_ref, + owner=owner, + globals=globals, + locals=locals, + type_params=type_params, + ) + except NameError: + if format == Format.FORWARDREF: + return forward_ref + else: + raise + + if isinstance(value, str): + value = ForwardRef(value) + + # Recursively evaluate the type + if isinstance(value, ForwardRef): + if getattr(value, "__forward_module__", True) is not None: + globals = None + return evaluate_forward_ref( + value, + globals=globals, + locals=locals, + type_params=type_params, owner=owner, + _recursive_guard=_recursive_guard, format=format + ) + if sys.version_info < (3, 12, 5) and type_params: + # Make use of type_params + locals = dict(locals) if locals else {} + for tvar in type_params: + if tvar.__name__ not in locals: # lets not overwrite something present + locals[tvar.__name__] = tvar + if sys.version_info < (3, 12, 5): + return typing._eval_type( + value, + globals, + locals, + recursive_guard=_recursive_guard | {forward_ref.__forward_arg__}, + ) + else: + return typing._eval_type( + value, + globals, + locals, + type_params, + recursive_guard=_recursive_guard | {forward_ref.__forward_arg__}, + ) + + +class Sentinel: + """Create a unique sentinel object. + + *name* should be the name of the variable to which the return value shall be assigned. + + *repr*, if supplied, will be used for the repr of the sentinel object. + If not provided, "" will be used. + """ + + def __init__( + self, + name: str, + repr: typing.Optional[str] = None, + ): + self._name = name + self._repr = repr if repr is not None else f'<{name}>' + + def __repr__(self): + return self._repr + + if sys.version_info < (3, 11): + # The presence of this method convinces typing._type_check + # that Sentinels are types. + def __call__(self, *args, **kwargs): + raise TypeError(f"{type(self).__name__!r} object is not callable") + + # Breakpoint: https://github.com/python/cpython/pull/21515 + if sys.version_info >= (3, 10): + def __or__(self, other): + return typing.Union[self, other] + + def __ror__(self, other): + return typing.Union[other, self] + + def __getstate__(self): + raise TypeError(f"Cannot pickle {type(self).__name__!r} object") + + +if sys.version_info >= (3, 14, 0, "beta"): + type_repr = annotationlib.type_repr +else: + def type_repr(value): + """Convert a Python value to a format suitable for use with the STRING format. + + This is intended as a helper for tools that support the STRING format but do + not have access to the code that originally produced the annotations. It uses + repr() for most objects. + + """ + if isinstance(value, (type, _types.FunctionType, _types.BuiltinFunctionType)): + if value.__module__ == "builtins": + return value.__qualname__ + return f"{value.__module__}.{value.__qualname__}" + if value is ...: + return "..." + return repr(value) + + +# Aliases for items that are in typing in all supported versions. +# We use hasattr() checks so this library will continue to import on +# future versions of Python that may remove these names. +_typing_names = [ + "AbstractSet", + "AnyStr", + "BinaryIO", + "Callable", + "Collection", + "Container", + "Dict", + "FrozenSet", + "Hashable", + "IO", + "ItemsView", + "Iterable", + "Iterator", + "KeysView", + "List", + "Mapping", + "MappingView", + "Match", + "MutableMapping", + "MutableSequence", + "MutableSet", + "Optional", + "Pattern", + "Reversible", + "Sequence", + "Set", + "Sized", + "TextIO", + "Tuple", + "Union", + "ValuesView", + "cast", + "no_type_check", + "no_type_check_decorator", + # This is private, but it was defined by typing_extensions for a long time + # and some users rely on it. + "_AnnotatedAlias", +] +globals().update( + {name: getattr(typing, name) for name in _typing_names if hasattr(typing, name)} +) +# These are defined unconditionally because they are used in +# typing-extensions itself. +Generic = typing.Generic +ForwardRef = typing.ForwardRef +Annotated = typing.Annotated diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/INSTALLER b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/METADATA b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..4963896cbe400ff39b2259506e32a58b73fa00ab --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/METADATA @@ -0,0 +1,48 @@ +Metadata-Version: 2.2 +Name: typing_inspection +Version: 0.4.2+computecanada +Summary: Runtime typing introspection tools +Description-Content-Type: text/markdown +Classifier: Development Status :: 3 - Alpha +Classifier: Intended Audience :: Developers +Classifier: Programming Language :: Python +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3.12 +Classifier: Programming Language :: Python :: 3.13 +Classifier: Programming Language :: Python :: 3.14 +Classifier: Programming Language :: Python :: Implementation :: CPython +Classifier: Topic :: Software Development :: Libraries :: Python Modules +Classifier: Typing :: Typed +Author-email: Victorien Plot +Project-URL: Homepage, https://github.com/pydantic/typing-inspection +Project-URL: Documentation, https://pydantic.github.io/typing-inspection/dev/ +Project-URL: Source, https://github.com/pydantic/typing-inspection +Project-URL: Changelog, https://github.com/pydantic/typing-inspection/blob/main/HISTORY.md +Requires-Python: >=3.9 +Requires-Dist: typing-extensions>=4.12.0 +License-File: LICENSE + +# typing-inspection + +[![CI](https://img.shields.io/github/actions/workflow/status/pydantic/typing-inspection/ci.yml?branch=main&logo=github&label=CI)](https://github.com/pydantic/typing-inspection/actions?query=event%3Apush+branch%3Amain+workflow%3ACI) +[![Coverage](https://coverage-badge.samuelcolvin.workers.dev/pydantic/typing-inspection.svg)](https://coverage-badge.samuelcolvin.workers.dev/redirect/pydantic/typing-inspection) +[![PyPI](https://img.shields.io/pypi/v/typing-inspection.svg)](https://pypi.org/project/typing-inspection/) +[![Versions](https://img.shields.io/pypi/pyversions/typing-inspection.svg)](https://github.com/pydantic/typing-inspection) +[![License](https://img.shields.io/github/license/pydantic/typing-inspection.svg)](https://github.com/pydantic/typing-inspection/blob/main/LICENSE) +[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff) + +`typing-inspection` provides tools to inspect type annotations at runtime. + +## Installation + +From [PyPI](https://pypi.org/project/typing-inspection/): + +```bash +pip install typing-inspection +``` + +The library can be imported from the `typing_inspection` module. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/RECORD b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..b7b101dcde3d002e4196732fa3562c4d28eeb788 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/RECORD @@ -0,0 +1,13 @@ +typing_inspection-0.4.2+computecanada.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +typing_inspection-0.4.2+computecanada.dist-info/METADATA,sha256=cOyL3p3uYJBHfJtCz4QNHHt6ysTw32I4qGqqe06loWQ,2542 +typing_inspection-0.4.2+computecanada.dist-info/RECORD,, +typing_inspection-0.4.2+computecanada.dist-info/WHEEL,sha256=E4Ta9GSW8Vg2e11uZktRc054ObD9JXJ5hm6vpxg0WJE,87 +typing_inspection-0.4.2+computecanada.dist-info/licenses/LICENSE,sha256=gEtZsl8sMb0nj5ICoZrkmjlFqiZkOH4tChKMfKzGHsM,1090 +typing_inspection/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +typing_inspection/__pycache__/__init__.cpython-311.pyc,, +typing_inspection/__pycache__/introspection.cpython-311.pyc,, +typing_inspection/__pycache__/typing_objects.cpython-311.pyc,, +typing_inspection/introspection.py,sha256=dD5Ad4J6hAfF6UBzBO4sqSs1h2ybQVThkQofLWWVBP0,22534 +typing_inspection/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +typing_inspection/typing_objects.py,sha256=kajVgh8J7UZ7wTidVxzFMpjwSnFGBoDoTqfVAAvOHZ8,17166 +typing_inspection/typing_objects.pyi,sha256=u1NDpl_RJFnAUMAMx-WBd0SBKwVG7luLEc8ukSvRnZs,9401 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/WHEEL b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/WHEEL new file mode 100644 index 0000000000000000000000000000000000000000..b270261c9bb44553b2f648368d1bf5cb7baf8416 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/WHEEL @@ -0,0 +1,5 @@ +Wheel-Version: 1.0 +Generator: wheelfile 0.0.8 +Root-Is-Purelib: true +Tag: py3-none-any + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/licenses/LICENSE b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/licenses/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..e825ad51621c5e34d370ba64f13f6854d89456a6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection-0.4.2+computecanada.dist-info/licenses/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) Pydantic Services Inc. 2025 to present + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/introspection.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/introspection.py new file mode 100644 index 0000000000000000000000000000000000000000..d6c083e5456c8ad5da7559a3c0a901e0812a5014 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/introspection.py @@ -0,0 +1,587 @@ +"""High-level introspection utilities, used to inspect type annotations.""" + +from __future__ import annotations + +import sys +import types +from collections.abc import Generator +from dataclasses import InitVar +from enum import Enum, IntEnum, auto +from typing import Any, Literal, NamedTuple, cast + +from typing_extensions import TypeAlias, assert_never, get_args, get_origin + +from . import typing_objects + +__all__ = ( + 'AnnotationSource', + 'ForbiddenQualifier', + 'InspectedAnnotation', + 'Qualifier', + 'get_literal_values', + 'inspect_annotation', + 'is_union_origin', +) + +if sys.version_info >= (3, 14) or sys.version_info < (3, 10): + + def is_union_origin(obj: Any, /) -> bool: + """Return whether the provided origin is the union form. + + ```pycon + >>> is_union_origin(typing.Union) + True + >>> is_union_origin(get_origin(int | str)) + True + >>> is_union_origin(types.UnionType) + True + ``` + + !!! note + Since Python 3.14, both `Union[, , ...]` and ` | | ...` forms create instances + of the same [`typing.Union`][] class. As such, it is recommended to not use this function + anymore (provided that you only support Python 3.14 or greater), and instead use the + [`typing_objects.is_union()`][typing_inspection.typing_objects.is_union] function directly: + + ```python + from typing import Union, get_origin + + from typing_inspection import typing_objects + + typ = int | str # Or Union[int, str] + origin = get_origin(typ) + if typing_objects.is_union(origin): + ... + ``` + """ + return typing_objects.is_union(obj) + + +else: + + def is_union_origin(obj: Any, /) -> bool: + """Return whether the provided origin is the union form. + + ```pycon + >>> is_union_origin(typing.Union) + True + >>> is_union_origin(get_origin(int | str)) + True + >>> is_union_origin(types.UnionType) + True + ``` + + !!! note + Since Python 3.14, both `Union[, , ...]` and ` | | ...` forms create instances + of the same [`typing.Union`][] class. As such, it is recommended to not use this function + anymore (provided that you only support Python 3.14 or greater), and instead use the + [`typing_objects.is_union()`][typing_inspection.typing_objects.is_union] function directly: + + ```python + from typing import Union, get_origin + + from typing_inspection import typing_objects + + typ = int | str # Or Union[int, str] + origin = get_origin(typ) + if typing_objects.is_union(origin): + ... + ``` + """ + return typing_objects.is_union(obj) or obj is types.UnionType + + +def _literal_type_check(value: Any, /) -> None: + """Type check the provided literal value against the legal parameters.""" + if ( + not isinstance(value, (int, bytes, str, bool, Enum, typing_objects.NoneType)) + and value is not typing_objects.NoneType + ): + raise TypeError(f'{value} is not a valid literal value, must be one of: int, bytes, str, Enum, None.') + + +def get_literal_values( + annotation: Any, + /, + *, + type_check: bool = False, + unpack_type_aliases: Literal['skip', 'lenient', 'eager'] = 'eager', +) -> Generator[Any]: + """Yield the values contained in the provided [`Literal`][typing.Literal] [special form][]. + + Args: + annotation: The [`Literal`][typing.Literal] [special form][] to unpack. + type_check: Whether to check if the literal values are [legal parameters][literal-legal-parameters]. + Raises a [`TypeError`][] otherwise. + unpack_type_aliases: What to do when encountering [PEP 695](https://peps.python.org/pep-0695/) + [type aliases][type-aliases]. Can be one of: + + - `'skip'`: Do not try to parse type aliases. Note that this can lead to incorrect results: + ```pycon + >>> type MyAlias = Literal[1, 2] + >>> list(get_literal_values(Literal[MyAlias, 3], unpack_type_aliases="skip")) + [MyAlias, 3] + ``` + + - `'lenient'`: Try to parse type aliases, and fallback to `'skip'` if the type alias can't be inspected + (because of an undefined forward reference). + + - `'eager'`: Parse type aliases and raise any encountered [`NameError`][] exceptions (the default): + ```pycon + >>> type MyAlias = Literal[1, 2] + >>> list(get_literal_values(Literal[MyAlias, 3], unpack_type_aliases="eager")) + [1, 2, 3] + ``` + + Note: + While `None` is [equivalent to][none] `type(None)`, the runtime implementation of [`Literal`][typing.Literal] + does not de-duplicate them. This function makes sure this de-duplication is applied: + + ```pycon + >>> list(get_literal_values(Literal[NoneType, None])) + [None] + ``` + + Example: + ```pycon + >>> type Ints = Literal[1, 2] + >>> list(get_literal_values(Literal[1, Ints], unpack_type_alias="skip")) + ["a", Ints] + >>> list(get_literal_values(Literal[1, Ints])) + [1, 2] + >>> list(get_literal_values(Literal[1.0], type_check=True)) + Traceback (most recent call last): + ... + TypeError: 1.0 is not a valid literal value, must be one of: int, bytes, str, Enum, None. + ``` + """ + # `literal` is guaranteed to be a `Literal[...]` special form, so use + # `__args__` directly instead of calling `get_args()`. + + if unpack_type_aliases == 'skip': + _has_none = False + # `Literal` parameters are already deduplicated, no need to do it ourselves. + # (we only check for `None` and `NoneType`, which should be considered as duplicates). + for arg in annotation.__args__: + if type_check: + _literal_type_check(arg) + if arg is None or arg is typing_objects.NoneType: + if not _has_none: + yield None + _has_none = True + else: + yield arg + else: + # We'll need to manually deduplicate parameters, see the `Literal` implementation in `typing`. + values_and_type: list[tuple[Any, type[Any]]] = [] + + for arg in annotation.__args__: + # Note: we could also check for generic aliases with a type alias as an origin. + # However, it is very unlikely that this happens as type variables can't appear in + # `Literal` forms, so the only valid (but unnecessary) use case would be something like: + # `type Test[T] = Literal['a']` (and then use `Test[SomeType]`). + if typing_objects.is_typealiastype(arg): + try: + alias_value = arg.__value__ + except NameError: + if unpack_type_aliases == 'eager': + raise + # unpack_type_aliases == "lenient": + if type_check: + _literal_type_check(arg) + values_and_type.append((arg, type(arg))) + else: + sub_args = get_literal_values( + alias_value, type_check=type_check, unpack_type_aliases=unpack_type_aliases + ) + values_and_type.extend((a, type(a)) for a in sub_args) # pyright: ignore[reportUnknownArgumentType] + else: + if type_check: + _literal_type_check(arg) + if arg is typing_objects.NoneType: + values_and_type.append((None, typing_objects.NoneType)) + else: + values_and_type.append((arg, type(arg))) # pyright: ignore[reportUnknownArgumentType] + + try: + dct = dict.fromkeys(values_and_type) + except TypeError: + # Unhashable parameters, the Python implementation allows them + yield from (p for p, _ in values_and_type) + else: + yield from (p for p, _ in dct) + + +Qualifier: TypeAlias = Literal['required', 'not_required', 'read_only', 'class_var', 'init_var', 'final'] +"""A [type qualifier][].""" + +_all_qualifiers: set[Qualifier] = set(get_args(Qualifier)) + + +# TODO at some point, we could switch to an enum flag, so that multiple sources +# can be combined. However, is there a need for this? +class AnnotationSource(IntEnum): + # TODO if/when https://peps.python.org/pep-0767/ is accepted, add 'read_only' + # to CLASS and NAMED_TUPLE (even though for named tuples it is redundant). + + """The source of an annotation, e.g. a class or a function. + + Depending on the source, different [type qualifiers][type qualifier] may be (dis)allowed. + """ + + ASSIGNMENT_OR_VARIABLE = auto() + """An annotation used in an assignment or variable annotation: + + ```python + x: Final[int] = 1 + y: Final[str] + ``` + + **Allowed type qualifiers:** [`Final`][typing.Final]. + """ + + CLASS = auto() + """An annotation used in the body of a class: + + ```python + class Test: + x: Final[int] = 1 + y: ClassVar[str] + ``` + + **Allowed type qualifiers:** [`ClassVar`][typing.ClassVar], [`Final`][typing.Final]. + """ + + DATACLASS = auto() + """An annotation used in the body of a dataclass: + + ```python + @dataclass + class Test: + x: Final[int] = 1 + y: InitVar[str] = 'test' + ``` + + **Allowed type qualifiers:** [`ClassVar`][typing.ClassVar], [`Final`][typing.Final], [`InitVar`][dataclasses.InitVar]. + """ # noqa: E501 + + TYPED_DICT = auto() + """An annotation used in the body of a [`TypedDict`][typing.TypedDict]: + + ```python + class TD(TypedDict): + x: Required[ReadOnly[int]] + y: ReadOnly[NotRequired[str]] + ``` + + **Allowed type qualifiers:** [`ReadOnly`][typing.ReadOnly], [`Required`][typing.Required], + [`NotRequired`][typing.NotRequired]. + """ + + NAMED_TUPLE = auto() + """An annotation used in the body of a [`NamedTuple`][typing.NamedTuple]. + + ```python + class NT(NamedTuple): + x: int + y: str + ``` + + **Allowed type qualifiers:** none. + """ + + FUNCTION = auto() + """An annotation used in a function, either for a parameter or the return value. + + ```python + def func(a: int) -> str: + ... + ``` + + **Allowed type qualifiers:** none. + """ + + ANY = auto() + """An annotation that might come from any source. + + **Allowed type qualifiers:** all. + """ + + BARE = auto() + """An annotation that is inspected as is. + + **Allowed type qualifiers:** none. + """ + + @property + def allowed_qualifiers(self) -> set[Qualifier]: + """The allowed [type qualifiers][type qualifier] for this annotation source.""" + # TODO use a match statement when Python 3.9 support is dropped. + if self is AnnotationSource.ASSIGNMENT_OR_VARIABLE: + return {'final'} + elif self is AnnotationSource.CLASS: + return {'final', 'class_var'} + elif self is AnnotationSource.DATACLASS: + return {'final', 'class_var', 'init_var'} + elif self is AnnotationSource.TYPED_DICT: + return {'required', 'not_required', 'read_only'} + elif self in (AnnotationSource.NAMED_TUPLE, AnnotationSource.FUNCTION, AnnotationSource.BARE): + return set() + elif self is AnnotationSource.ANY: + return _all_qualifiers + else: # pragma: no cover + assert_never(self) + + +class ForbiddenQualifier(Exception): + """The provided [type qualifier][] is forbidden.""" + + qualifier: Qualifier + """The forbidden qualifier.""" + + def __init__(self, qualifier: Qualifier, /) -> None: + self.qualifier = qualifier + + +class _UnknownTypeEnum(Enum): + UNKNOWN = auto() + + def __str__(self) -> str: + return 'UNKNOWN' + + def __repr__(self) -> str: + return '' + + +UNKNOWN = _UnknownTypeEnum.UNKNOWN +"""A sentinel value used when no [type expression][] is present.""" + +_UnkownType: TypeAlias = Literal[_UnknownTypeEnum.UNKNOWN] +"""The type of the [`UNKNOWN`][typing_inspection.introspection.UNKNOWN] sentinel value.""" + + +class InspectedAnnotation(NamedTuple): + """The result of the inspected annotation.""" + + type: Any | _UnkownType + """The final [type expression][], with [type qualifiers][type qualifier] and annotated metadata stripped. + + If no type expression is available, the [`UNKNOWN`][typing_inspection.introspection.UNKNOWN] sentinel + value is used instead. This is the case when a [type qualifier][] is used with no type annotation: + + ```python + ID: Final = 1 + + class C: + x: ClassVar = 'test' + ``` + """ + + qualifiers: set[Qualifier] + """The [type qualifiers][type qualifier] present on the annotation.""" + + metadata: list[Any] + """The annotated metadata.""" + + +def inspect_annotation( # noqa: PLR0915 + annotation: Any, + /, + *, + annotation_source: AnnotationSource, + unpack_type_aliases: Literal['skip', 'lenient', 'eager'] = 'skip', +) -> InspectedAnnotation: + """Inspect an [annotation expression][], extracting any [type qualifier][] and metadata. + + An [annotation expression][] is a [type expression][] optionally surrounded by one or more + [type qualifiers][type qualifier] or by [`Annotated`][typing.Annotated]. This function will: + + - Unwrap the type expression, keeping track of the type qualifiers. + - Unwrap [`Annotated`][typing.Annotated] forms, keeping track of the annotated metadata. + + Args: + annotation: The annotation expression to be inspected. + annotation_source: The source of the annotation. Depending on the source (e.g. a class), different type + qualifiers may be (dis)allowed. To allow any type qualifier, use + [`AnnotationSource.ANY`][typing_inspection.introspection.AnnotationSource.ANY]. + unpack_type_aliases: What to do when encountering [PEP 695](https://peps.python.org/pep-0695/) + [type aliases][type-aliases]. Can be one of: + + - `'skip'`: Do not try to parse type aliases (the default): + ```pycon + >>> type MyInt = Annotated[int, 'meta'] + >>> inspect_annotation(MyInt, annotation_source=AnnotationSource.BARE, unpack_type_aliases='skip') + InspectedAnnotation(type=MyInt, qualifiers={}, metadata=[]) + ``` + + - `'lenient'`: Try to parse type aliases, and fallback to `'skip'` if the type alias + can't be inspected (because of an undefined forward reference): + ```pycon + >>> type MyInt = Annotated[Undefined, 'meta'] + >>> inspect_annotation(MyInt, annotation_source=AnnotationSource.BARE, unpack_type_aliases='lenient') + InspectedAnnotation(type=MyInt, qualifiers={}, metadata=[]) + >>> Undefined = int + >>> inspect_annotation(MyInt, annotation_source=AnnotationSource.BARE, unpack_type_aliases='lenient') + InspectedAnnotation(type=int, qualifiers={}, metadata=['meta']) + ``` + + - `'eager'`: Parse type aliases and raise any encountered [`NameError`][] exceptions. + + Returns: + The result of the inspected annotation, where the type expression, used qualifiers and metadata is stored. + + Example: + ```pycon + >>> inspect_annotation( + ... Final[Annotated[ClassVar[Annotated[int, 'meta_1']], 'meta_2']], + ... annotation_source=AnnotationSource.CLASS, + ... ) + ... + InspectedAnnotation(type=int, qualifiers={'class_var', 'final'}, metadata=['meta_1', 'meta_2']) + ``` + """ + allowed_qualifiers = annotation_source.allowed_qualifiers + qualifiers: set[Qualifier] = set() + metadata: list[Any] = [] + + while True: + annotation, _meta = _unpack_annotated(annotation, unpack_type_aliases=unpack_type_aliases) + if _meta: + metadata = _meta + metadata + continue + + origin = get_origin(annotation) + if origin is not None: + if typing_objects.is_classvar(origin): + if 'class_var' not in allowed_qualifiers: + raise ForbiddenQualifier('class_var') + qualifiers.add('class_var') + annotation = annotation.__args__[0] + elif typing_objects.is_final(origin): + if 'final' not in allowed_qualifiers: + raise ForbiddenQualifier('final') + qualifiers.add('final') + annotation = annotation.__args__[0] + elif typing_objects.is_required(origin): + if 'required' not in allowed_qualifiers: + raise ForbiddenQualifier('required') + qualifiers.add('required') + annotation = annotation.__args__[0] + elif typing_objects.is_notrequired(origin): + if 'not_required' not in allowed_qualifiers: + raise ForbiddenQualifier('not_required') + qualifiers.add('not_required') + annotation = annotation.__args__[0] + elif typing_objects.is_readonly(origin): + if 'read_only' not in allowed_qualifiers: + raise ForbiddenQualifier('not_required') + qualifiers.add('read_only') + annotation = annotation.__args__[0] + else: + # origin is not None but not a type qualifier nor `Annotated` (e.g. `list[int]`): + break + elif isinstance(annotation, InitVar): + if 'init_var' not in allowed_qualifiers: + raise ForbiddenQualifier('init_var') + qualifiers.add('init_var') + annotation = cast(Any, annotation.type) + else: + break + + # `Final`, `ClassVar` and `InitVar` are type qualifiers allowed to be used as a bare annotation: + if typing_objects.is_final(annotation): + if 'final' not in allowed_qualifiers: + raise ForbiddenQualifier('final') + qualifiers.add('final') + annotation = UNKNOWN + elif typing_objects.is_classvar(annotation): + if 'class_var' not in allowed_qualifiers: + raise ForbiddenQualifier('class_var') + qualifiers.add('class_var') + annotation = UNKNOWN + elif annotation is InitVar: + if 'init_var' not in allowed_qualifiers: + raise ForbiddenQualifier('init_var') + qualifiers.add('init_var') + annotation = UNKNOWN + + return InspectedAnnotation(annotation, qualifiers, metadata) + + +def _unpack_annotated_inner( + annotation: Any, unpack_type_aliases: Literal['lenient', 'eager'], check_annotated: bool +) -> tuple[Any, list[Any]]: + origin = get_origin(annotation) + if check_annotated and typing_objects.is_annotated(origin): + annotated_type = annotation.__origin__ + metadata = list(annotation.__metadata__) + + # The annotated type might be a PEP 695 type alias, so we need to recursively + # unpack it. Because Python already flattens `Annotated[Annotated[, ...], ...]` forms, + # we can skip the `is_annotated()` check in the next call: + annotated_type, sub_meta = _unpack_annotated_inner( + annotated_type, unpack_type_aliases=unpack_type_aliases, check_annotated=False + ) + metadata = sub_meta + metadata + return annotated_type, metadata + elif typing_objects.is_typealiastype(annotation): + try: + value = annotation.__value__ + except NameError: + if unpack_type_aliases == 'eager': + raise + else: + typ, metadata = _unpack_annotated_inner( + value, unpack_type_aliases=unpack_type_aliases, check_annotated=True + ) + if metadata: + # Having metadata means the type alias' `__value__` was an `Annotated` form + # (or, recursively, a type alias to an `Annotated` form). It is important to check + # for this, as we don't want to unpack other type aliases (e.g. `type MyInt = int`). + return typ, metadata + return annotation, [] + elif typing_objects.is_typealiastype(origin): + # When parameterized, PEP 695 type aliases become generic aliases + # (e.g. with `type MyList[T] = Annotated[list[T], ...]`, `MyList[int]` + # is a generic alias). + try: + value = origin.__value__ + except NameError: + if unpack_type_aliases == 'eager': + raise + else: + # While Python already handles type variable replacement for simple `Annotated` forms, + # we need to manually apply the same logic for PEP 695 type aliases: + # - With `MyList = Annotated[list[T], ...]`, `MyList[int] == Annotated[list[int], ...]` + # - With `type MyList[T] = Annotated[list[T], ...]`, `MyList[int].__value__ == Annotated[list[T], ...]`. + + try: + # To do so, we emulate the parameterization of the value with the arguments: + # with `type MyList[T] = Annotated[list[T], ...]`, to emulate `MyList[int]`, + # we do `Annotated[list[T], ...][int]` (which gives `Annotated[list[T], ...]`): + value = value[annotation.__args__] + except TypeError: + # Might happen if the type alias is parameterized, but its value doesn't have any + # type variables, e.g. `type MyInt[T] = int`. + pass + typ, metadata = _unpack_annotated_inner( + value, unpack_type_aliases=unpack_type_aliases, check_annotated=True + ) + if metadata: + return typ, metadata + return annotation, [] + + return annotation, [] + + +# This could eventually be made public: +def _unpack_annotated( + annotation: Any, /, *, unpack_type_aliases: Literal['skip', 'lenient', 'eager'] = 'eager' +) -> tuple[Any, list[Any]]: + if unpack_type_aliases == 'skip': + if typing_objects.is_annotated(get_origin(annotation)): + return annotation.__origin__, list(annotation.__metadata__) + else: + return annotation, [] + + return _unpack_annotated_inner(annotation, unpack_type_aliases=unpack_type_aliases, check_annotated=True) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/py.typed b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/typing_objects.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/typing_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..dc44ba98cd8ccdba374982819083d89bbf808c2e --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/typing_objects.py @@ -0,0 +1,607 @@ +"""Low-level introspection utilities for [`typing`][] members. + +The provided functions in this module check against both the [`typing`][] and [`typing_extensions`][] +variants, if they exists and are different. +""" +# ruff: noqa: UP006 + +import collections.abc +import contextlib +import re +import sys +import typing +import warnings +from textwrap import dedent +from types import FunctionType, GenericAlias +from typing import Any, Final + +import typing_extensions +from typing_extensions import LiteralString, TypeAliasType, TypeIs, deprecated + +__all__ = ( + 'DEPRECATED_ALIASES', + 'NoneType', + 'is_annotated', + 'is_any', + 'is_classvar', + 'is_concatenate', + 'is_deprecated', + 'is_final', + 'is_forwardref', + 'is_generic', + 'is_literal', + 'is_literalstring', + 'is_namedtuple', + 'is_never', + 'is_newtype', + 'is_nodefault', + 'is_noextraitems', + 'is_noreturn', + 'is_notrequired', + 'is_paramspec', + 'is_paramspecargs', + 'is_paramspeckwargs', + 'is_readonly', + 'is_required', + 'is_self', + 'is_typealias', + 'is_typealiastype', + 'is_typeguard', + 'is_typeis', + 'is_typevar', + 'is_typevartuple', + 'is_union', + 'is_unpack', +) + +_IS_PY310 = sys.version_info[:2] == (3, 10) + + +def _compile_identity_check_function(member: LiteralString, function_name: LiteralString) -> FunctionType: + """Create a function checking that the function argument is the (unparameterized) typing `member`. + + The function will make sure to check against both the `typing` and `typing_extensions` + variants as depending on the Python version, the `typing_extensions` variant might be different. + For instance, on Python 3.9: + + ```pycon + >>> from typing import Literal as t_Literal + >>> from typing_extensions import Literal as te_Literal, get_origin + + >>> t_Literal is te_Literal + False + >>> get_origin(t_Literal[1]) + typing.Literal + >>> get_origin(te_Literal[1]) + typing_extensions.Literal + ``` + """ + in_typing = hasattr(typing, member) + in_typing_extensions = hasattr(typing_extensions, member) + + if in_typing and in_typing_extensions: + if getattr(typing, member) is getattr(typing_extensions, member): + check_code = f'obj is typing.{member}' + else: + check_code = f'obj is typing.{member} or obj is typing_extensions.{member}' + elif in_typing and not in_typing_extensions: + check_code = f'obj is typing.{member}' + elif not in_typing and in_typing_extensions: + check_code = f'obj is typing_extensions.{member}' + else: + check_code = 'False' + + func_code = dedent(f""" + def {function_name}(obj: Any, /) -> bool: + return {check_code} + """) + + locals_: dict[str, Any] = {} + globals_: dict[str, Any] = {'Any': Any, 'typing': typing, 'typing_extensions': typing_extensions} + exec(func_code, globals_, locals_) + return locals_[function_name] + + +def _compile_isinstance_check_function(member: LiteralString, function_name: LiteralString) -> FunctionType: + """Create a function checking that the function is an instance of the typing `member`. + + The function will make sure to check against both the `typing` and `typing_extensions` + variants as depending on the Python version, the `typing_extensions` variant might be different. + """ + in_typing = hasattr(typing, member) + in_typing_extensions = hasattr(typing_extensions, member) + + if in_typing and in_typing_extensions: + if getattr(typing, member) is getattr(typing_extensions, member): + check_code = f'isinstance(obj, typing.{member})' + else: + check_code = f'isinstance(obj, (typing.{member}, typing_extensions.{member}))' + elif in_typing and not in_typing_extensions: + check_code = f'isinstance(obj, typing.{member})' + elif not in_typing and in_typing_extensions: + check_code = f'isinstance(obj, typing_extensions.{member})' + else: + check_code = 'False' + + func_code = dedent(f""" + def {function_name}(obj: Any, /) -> 'TypeIs[{member}]': + return {check_code} + """) + + locals_: dict[str, Any] = {} + globals_: dict[str, Any] = {'Any': Any, 'typing': typing, 'typing_extensions': typing_extensions} + exec(func_code, globals_, locals_) + return locals_[function_name] + + +if sys.version_info >= (3, 10): + from types import NoneType +else: + NoneType = type(None) + +# Keep this ordered, as per `typing.__all__`: + +is_annotated = _compile_identity_check_function('Annotated', 'is_annotated') +is_annotated.__doc__ = """ +Return whether the argument is the [`Annotated`][typing.Annotated] [special form][]. + +```pycon +>>> is_annotated(Annotated) +True +>>> is_annotated(Annotated[int, ...]) +False +``` +""" + +is_any = _compile_identity_check_function('Any', 'is_any') +is_any.__doc__ = """ +Return whether the argument is the [`Any`][typing.Any] [special form][]. + +```pycon +>>> is_any(Any) +True +``` +""" + +is_classvar = _compile_identity_check_function('ClassVar', 'is_classvar') +is_classvar.__doc__ = """ +Return whether the argument is the [`ClassVar`][typing.ClassVar] [type qualifier][]. + +```pycon +>>> is_classvar(ClassVar) +True +>>> is_classvar(ClassVar[int]) +>>> False +``` +""" + +is_concatenate = _compile_identity_check_function('Concatenate', 'is_concatenate') +is_concatenate.__doc__ = """ +Return whether the argument is the [`Concatenate`][typing.Concatenate] [special form][]. + +```pycon +>>> is_concatenate(Concatenate) +True +>>> is_concatenate(Concatenate[int, P]) +False +``` +""" + +is_final = _compile_identity_check_function('Final', 'is_final') +is_final.__doc__ = """ +Return whether the argument is the [`Final`][typing.Final] [type qualifier][]. + +```pycon +>>> is_final(Final) +True +>>> is_final(Final[int]) +False +``` +""" + + +# Unlikely to have a different version in `typing-extensions`, but keep it consistent. +# Also note that starting in 3.14, this is an alias to `annotationlib.ForwardRef`, but +# accessing it from `typing` doesn't seem to be deprecated. +is_forwardref = _compile_isinstance_check_function('ForwardRef', 'is_forwardref') +is_forwardref.__doc__ = """ +Return whether the argument is an instance of [`ForwardRef`][typing.ForwardRef]. + +```pycon +>>> is_forwardref(ForwardRef('T')) +True +``` +""" + + +is_generic = _compile_identity_check_function('Generic', 'is_generic') +is_generic.__doc__ = """ +Return whether the argument is the [`Generic`][typing.Generic] [special form][]. + +```pycon +>>> is_generic(Generic) +True +>>> is_generic(Generic[T]) +False +``` +""" + +is_literal = _compile_identity_check_function('Literal', 'is_literal') +is_literal.__doc__ = """ +Return whether the argument is the [`Literal`][typing.Literal] [special form][]. + +```pycon +>>> is_literal(Literal) +True +>>> is_literal(Literal["a"]) +False +``` +""" + + +# `get_origin(Optional[int]) is Union`, so `is_optional()` isn't implemented. + +is_paramspec = _compile_isinstance_check_function('ParamSpec', 'is_paramspec') +is_paramspec.__doc__ = """ +Return whether the argument is an instance of [`ParamSpec`][typing.ParamSpec]. + +```pycon +>>> P = ParamSpec('P') +>>> is_paramspec(P) +True +``` +""" + +# Protocol? + +is_typevar = _compile_isinstance_check_function('TypeVar', 'is_typevar') +is_typevar.__doc__ = """ +Return whether the argument is an instance of [`TypeVar`][typing.TypeVar]. + +```pycon +>>> T = TypeVar('T') +>>> is_typevar(T) +True +``` +""" + +is_typevartuple = _compile_isinstance_check_function('TypeVarTuple', 'is_typevartuple') +is_typevartuple.__doc__ = """ +Return whether the argument is an instance of [`TypeVarTuple`][typing.TypeVarTuple]. + +```pycon +>>> Ts = TypeVarTuple('Ts') +>>> is_typevartuple(Ts) +True +``` +""" + +is_union = _compile_identity_check_function('Union', 'is_union') +is_union.__doc__ = """ +Return whether the argument is the [`Union`][typing.Union] [special form][]. + +This function can also be used to check for the [`Optional`][typing.Optional] [special form][], +as at runtime, `Optional[int]` is equivalent to `Union[int, None]`. + +```pycon +>>> is_union(Union) +True +>>> is_union(Union[int, str]) +False +``` + +!!! warning + This does not check for unions using the [new syntax][types-union] (e.g. `int | str`). +""" + + +def is_namedtuple(obj: Any, /) -> bool: + """Return whether the argument is a named tuple type. + + This includes [`NamedTuple`][typing.NamedTuple] subclasses and classes created from the + [`collections.namedtuple`][] factory function. + + ```pycon + >>> class User(NamedTuple): + ... name: str + ... + >>> is_namedtuple(User) + True + >>> City = collections.namedtuple('City', []) + >>> is_namedtuple(City) + True + >>> is_namedtuple(NamedTuple) + False + ``` + """ + return isinstance(obj, type) and issubclass(obj, tuple) and hasattr(obj, '_fields') # pyright: ignore[reportUnknownArgumentType] + + +# TypedDict? + +# BinaryIO? IO? TextIO? + +is_literalstring = _compile_identity_check_function('LiteralString', 'is_literalstring') +is_literalstring.__doc__ = """ +Return whether the argument is the [`LiteralString`][typing.LiteralString] [special form][]. + +```pycon +>>> is_literalstring(LiteralString) +True +``` +""" + +is_never = _compile_identity_check_function('Never', 'is_never') +is_never.__doc__ = """ +Return whether the argument is the [`Never`][typing.Never] [special form][]. + +```pycon +>>> is_never(Never) +True +``` +""" + +if sys.version_info >= (3, 10): + is_newtype = _compile_isinstance_check_function('NewType', 'is_newtype') +else: # On Python 3.10, `NewType` is a function. + + def is_newtype(obj: Any, /) -> bool: + return hasattr(obj, '__supertype__') + + +is_newtype.__doc__ = """ +Return whether the argument is a [`NewType`][typing.NewType]. + +```pycon +>>> UserId = NewType("UserId", int) +>>> is_newtype(UserId) +True +``` +""" + +is_nodefault = _compile_identity_check_function('NoDefault', 'is_nodefault') +is_nodefault.__doc__ = """ +Return whether the argument is the [`NoDefault`][typing.NoDefault] sentinel object. + +```pycon +>>> is_nodefault(NoDefault) +True +``` +""" + +is_noextraitems = _compile_identity_check_function('NoExtraItems', 'is_noextraitems') +is_noextraitems.__doc__ = """ +Return whether the argument is the `NoExtraItems` sentinel object. + +```pycon +>>> is_noextraitems(NoExtraItems) +True +``` +""" + +is_noreturn = _compile_identity_check_function('NoReturn', 'is_noreturn') +is_noreturn.__doc__ = """ +Return whether the argument is the [`NoReturn`][typing.NoReturn] [special form][]. + +```pycon +>>> is_noreturn(NoReturn) +True +>>> is_noreturn(Never) +False +``` +""" + +is_notrequired = _compile_identity_check_function('NotRequired', 'is_notrequired') +is_notrequired.__doc__ = """ +Return whether the argument is the [`NotRequired`][typing.NotRequired] [special form][]. + +```pycon +>>> is_notrequired(NotRequired) +True +``` +""" + +is_paramspecargs = _compile_isinstance_check_function('ParamSpecArgs', 'is_paramspecargs') +is_paramspecargs.__doc__ = """ +Return whether the argument is an instance of [`ParamSpecArgs`][typing.ParamSpecArgs]. + +```pycon +>>> P = ParamSpec('P') +>>> is_paramspecargs(P.args) +True +``` +""" + +is_paramspeckwargs = _compile_isinstance_check_function('ParamSpecKwargs', 'is_paramspeckwargs') +is_paramspeckwargs.__doc__ = """ +Return whether the argument is an instance of [`ParamSpecKwargs`][typing.ParamSpecKwargs]. + +```pycon +>>> P = ParamSpec('P') +>>> is_paramspeckwargs(P.kwargs) +True +``` +""" + +is_readonly = _compile_identity_check_function('ReadOnly', 'is_readonly') +is_readonly.__doc__ = """ +Return whether the argument is the [`ReadOnly`][typing.ReadOnly] [special form][]. + +```pycon +>>> is_readonly(ReadOnly) +True +``` +""" + +is_required = _compile_identity_check_function('Required', 'is_required') +is_required.__doc__ = """ +Return whether the argument is the [`Required`][typing.Required] [special form][]. + +```pycon +>>> is_required(Required) +True +``` +""" + +is_self = _compile_identity_check_function('Self', 'is_self') +is_self.__doc__ = """ +Return whether the argument is the [`Self`][typing.Self] [special form][]. + +```pycon +>>> is_self(Self) +True +``` +""" + +# TYPE_CHECKING? + +is_typealias = _compile_identity_check_function('TypeAlias', 'is_typealias') +is_typealias.__doc__ = """ +Return whether the argument is the [`TypeAlias`][typing.TypeAlias] [special form][]. + +```pycon +>>> is_typealias(TypeAlias) +True +``` +""" + +is_typeguard = _compile_identity_check_function('TypeGuard', 'is_typeguard') +is_typeguard.__doc__ = """ +Return whether the argument is the [`TypeGuard`][typing.TypeGuard] [special form][]. + +```pycon +>>> is_typeguard(TypeGuard) +True +``` +""" + +is_typeis = _compile_identity_check_function('TypeIs', 'is_typeis') +is_typeis.__doc__ = """ +Return whether the argument is the [`TypeIs`][typing.TypeIs] [special form][]. + +```pycon +>>> is_typeis(TypeIs) +True +``` +""" + +_is_typealiastype_inner = _compile_isinstance_check_function('TypeAliasType', '_is_typealiastype_inner') + + +if _IS_PY310: + # Parameterized PEP 695 type aliases are instances of `types.GenericAlias` in typing_extensions>=4.13.0. + # On Python 3.10, with `Alias[int]` being such an instance of `GenericAlias`, + # `isinstance(Alias[int], TypeAliasType)` returns `True`. + # See https://github.com/python/cpython/issues/89828. + def is_typealiastype(obj: Any, /) -> 'TypeIs[TypeAliasType]': + return type(obj) is not GenericAlias and _is_typealiastype_inner(obj) +else: + is_typealiastype = _compile_isinstance_check_function('TypeAliasType', 'is_typealiastype') + +is_typealiastype.__doc__ = """ +Return whether the argument is a [`TypeAliasType`][typing.TypeAliasType] instance. + +```pycon +>>> type MyInt = int +>>> is_typealiastype(MyInt) +True +>>> MyStr = TypeAliasType("MyStr", str) +>>> is_typealiastype(MyStr): +True +>>> type MyList[T] = list[T] +>>> is_typealiastype(MyList[int]) +False +``` +""" + +is_unpack = _compile_identity_check_function('Unpack', 'is_unpack') +is_unpack.__doc__ = """ +Return whether the argument is the [`Unpack`][typing.Unpack] [special form][]. + +```pycon +>>> is_unpack(Unpack) +True +>>> is_unpack(Unpack[Ts]) +False +``` +""" + + +if sys.version_info >= (3, 13): + + def is_deprecated(obj: Any, /) -> 'TypeIs[deprecated]': + return isinstance(obj, (warnings.deprecated, typing_extensions.deprecated)) + +else: + + def is_deprecated(obj: Any, /) -> 'TypeIs[deprecated]': + return isinstance(obj, typing_extensions.deprecated) + + +is_deprecated.__doc__ = """ +Return whether the argument is a [`deprecated`][warnings.deprecated] instance. + +This also includes the [`typing_extensions` backport][typing_extensions.deprecated]. + +```pycon +>>> is_deprecated(warnings.deprecated('message')) +True +>>> is_deprecated(typing_extensions.deprecated('message')) +True +``` +""" + + +# Aliases defined in the `typing` module using `typing._SpecialGenericAlias` (itself aliased as `alias()`): +DEPRECATED_ALIASES: Final[dict[Any, type[Any]]] = { + typing.Hashable: collections.abc.Hashable, + typing.Awaitable: collections.abc.Awaitable, + typing.Coroutine: collections.abc.Coroutine, + typing.AsyncIterable: collections.abc.AsyncIterable, + typing.AsyncIterator: collections.abc.AsyncIterator, + typing.Iterable: collections.abc.Iterable, + typing.Iterator: collections.abc.Iterator, + typing.Reversible: collections.abc.Reversible, + typing.Sized: collections.abc.Sized, + typing.Container: collections.abc.Container, + typing.Collection: collections.abc.Collection, + # type ignore reason: https://github.com/python/typeshed/issues/6257: + typing.Callable: collections.abc.Callable, # pyright: ignore[reportAssignmentType, reportUnknownMemberType] + typing.AbstractSet: collections.abc.Set, + typing.MutableSet: collections.abc.MutableSet, + typing.Mapping: collections.abc.Mapping, + typing.MutableMapping: collections.abc.MutableMapping, + typing.Sequence: collections.abc.Sequence, + typing.MutableSequence: collections.abc.MutableSequence, + typing.Tuple: tuple, + typing.List: list, + typing.Deque: collections.deque, + typing.Set: set, + typing.FrozenSet: frozenset, + typing.MappingView: collections.abc.MappingView, + typing.KeysView: collections.abc.KeysView, + typing.ItemsView: collections.abc.ItemsView, + typing.ValuesView: collections.abc.ValuesView, + typing.Dict: dict, + typing.DefaultDict: collections.defaultdict, + typing.OrderedDict: collections.OrderedDict, + typing.Counter: collections.Counter, + typing.ChainMap: collections.ChainMap, + typing.Generator: collections.abc.Generator, + typing.AsyncGenerator: collections.abc.AsyncGenerator, + typing.Type: type, + # Defined in `typing.__getattr__`: + typing.Pattern: re.Pattern, + typing.Match: re.Match, + typing.ContextManager: contextlib.AbstractContextManager, + typing.AsyncContextManager: contextlib.AbstractAsyncContextManager, + # Skipped: `ByteString` (deprecated, removed in 3.14) +} +"""A mapping between the deprecated typing aliases to their replacement, as per [PEP 585](https://peps.python.org/pep-0585/).""" + + +# Add the `typing_extensions` aliases: +for alias, target in list(DEPRECATED_ALIASES.items()): + # Use `alias.__name__` when we drop support for Python 3.9 + if (te_alias := getattr(typing_extensions, alias._name, None)) is not None: + DEPRECATED_ALIASES[te_alias] = target diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/typing_objects.pyi b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/typing_objects.pyi new file mode 100644 index 0000000000000000000000000000000000000000..5071598005a21063ff3b2a2a1dedced885267de6 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/typing_inspection/typing_objects.pyi @@ -0,0 +1,417 @@ +# Stub file generated using: +# `stubgen --inspect-mode --include-docstrings -m typing_inspection.typing_objects` +# (manual edits need to be applied). +"""Low-level introspection utilities for [`typing`][] members. + +The provided functions in this module check against both the [`typing`][] and [`typing_extensions`][] +variants, if they exists and are different. +""" + +import sys +from typing import Any, Final, ForwardRef, NewType, TypeVar + +from typing_extensions import ParamSpec, ParamSpecArgs, ParamSpecKwargs, TypeAliasType, TypeIs, TypeVarTuple, deprecated + +__all__ = [ + 'DEPRECATED_ALIASES', + 'NoneType', + 'is_annotated', + 'is_any', + 'is_classvar', + 'is_concatenate', + 'is_deprecated', + 'is_final', + 'is_generic', + 'is_literal', + 'is_literalstring', + 'is_namedtuple', + 'is_never', + 'is_newtype', + 'is_nodefault', + 'is_noextraitems', + 'is_noreturn', + 'is_notrequired', + 'is_paramspec', + 'is_paramspecargs', + 'is_paramspeckwargs', + 'is_readonly', + 'is_required', + 'is_self', + 'is_typealias', + 'is_typealiastype', + 'is_typeguard', + 'is_typeis', + 'is_typevar', + 'is_typevartuple', + 'is_union', + 'is_unpack', +] + +if sys.version_info >= (3, 10): + from types import NoneType +else: + NoneType = type(None) + +def is_annotated(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Annotated`][typing.Annotated] [special form][]. + + ```pycon + >>> is_annotated(Annotated) + True + >>> is_annotated(Annotated[int, ...]) + False + ``` + """ + +def is_any(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Any`][typing.Any] [special form][]. + + ```pycon + >>> is_any(Any) + True + ``` + """ + +def is_classvar(obj: Any, /) -> bool: + """ + Return whether the argument is the [`ClassVar`][typing.ClassVar] [type qualifier][]. + + ```pycon + >>> is_classvar(ClassVar) + True + >>> is_classvar(ClassVar[int]) + >>> False + ``` + """ + +def is_concatenate(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Concatenate`][typing.Concatenate] [special form][]. + + ```pycon + >>> is_concatenate(Concatenate) + True + >>> is_concatenate(Concatenate[int, P]) + False + ``` + """ + +def is_final(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Final`][typing.Final] [type qualifier][]. + + ```pycon + >>> is_final(Final) + True + >>> is_final(Final[int]) + False + ``` + """ + +def is_forwardref(obj: Any, /) -> TypeIs[ForwardRef]: + """ + Return whether the argument is an instance of [`ForwardRef`][typing.ForwardRef]. + + ```pycon + >>> is_forwardref(ForwardRef('T')) + True + ``` + """ + +def is_generic(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Generic`][typing.Generic] [special form][]. + + ```pycon + >>> is_generic(Generic) + True + >>> is_generic(Generic[T]) + False + ``` + """ + +def is_literal(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Literal`][typing.Literal] [special form][]. + + ```pycon + >>> is_literal(Literal) + True + >>> is_literal(Literal["a"]) + False + ``` + """ + +def is_paramspec(obj: Any, /) -> TypeIs[ParamSpec]: + """ + Return whether the argument is an instance of [`ParamSpec`][typing.ParamSpec]. + + ```pycon + >>> P = ParamSpec('P') + >>> is_paramspec(P) + True + ``` + """ + +def is_typevar(obj: Any, /) -> TypeIs[TypeVar]: + """ + Return whether the argument is an instance of [`TypeVar`][typing.TypeVar]. + + ```pycon + >>> T = TypeVar('T') + >>> is_typevar(T) + True + ``` + """ + +def is_typevartuple(obj: Any, /) -> TypeIs[TypeVarTuple]: + """ + Return whether the argument is an instance of [`TypeVarTuple`][typing.TypeVarTuple]. + + ```pycon + >>> Ts = TypeVarTuple('Ts') + >>> is_typevartuple(Ts) + True + ``` + """ + +def is_union(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Union`][typing.Union] [special form][]. + + This function can also be used to check for the [`Optional`][typing.Optional] [special form][], + as at runtime, `Optional[int]` is equivalent to `Union[int, None]`. + + ```pycon + >>> is_union(Union) + True + >>> is_union(Union[int, str]) + False + ``` + + !!! warning + This does not check for unions using the [new syntax][types-union] (e.g. `int | str`). + """ + +def is_namedtuple(obj: Any, /) -> bool: + """Return whether the argument is a named tuple type. + + This includes [`NamedTuple`][typing.NamedTuple] subclasses and classes created from the + [`collections.namedtuple`][] factory function. + + ```pycon + >>> class User(NamedTuple): + ... name: str + ... + >>> is_namedtuple(User) + True + >>> City = collections.namedtuple('City', []) + >>> is_namedtuple(City) + True + >>> is_namedtuple(NamedTuple) + False + ``` + """ + +def is_literalstring(obj: Any, /) -> bool: + """ + Return whether the argument is the [`LiteralString`][typing.LiteralString] [special form][]. + + ```pycon + >>> is_literalstring(LiteralString) + True + ``` + """ + +def is_never(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Never`][typing.Never] [special form][]. + + ```pycon + >>> is_never(Never) + True + ``` + """ + +def is_newtype(obj: Any, /) -> TypeIs[NewType]: + """ + Return whether the argument is a [`NewType`][typing.NewType]. + + ```pycon + >>> UserId = NewType("UserId", int) + >>> is_newtype(UserId) + True + ``` + """ + +def is_nodefault(obj: Any, /) -> bool: + """ + Return whether the argument is the [`NoDefault`][typing.NoDefault] sentinel object. + + ```pycon + >>> is_nodefault(NoDefault) + True + ``` + """ + +def is_noextraitems(obj: Any, /) -> bool: + """ + Return whether the argument is the `NoExtraItems` sentinel object. + + ```pycon + >>> is_noextraitems(NoExtraItems) + True + ``` + """ + +def is_noreturn(obj: Any, /) -> bool: + """ + Return whether the argument is the [`NoReturn`][typing.NoReturn] [special form][]. + + ```pycon + >>> is_noreturn(NoReturn) + True + >>> is_noreturn(Never) + False + ``` + """ + +def is_notrequired(obj: Any, /) -> bool: + """ + Return whether the argument is the [`NotRequired`][typing.NotRequired] [special form][]. + + ```pycon + >>> is_notrequired(NotRequired) + True + ``` + """ + +def is_paramspecargs(obj: Any, /) -> TypeIs[ParamSpecArgs]: + """ + Return whether the argument is an instance of [`ParamSpecArgs`][typing.ParamSpecArgs]. + + ```pycon + >>> P = ParamSpec('P') + >>> is_paramspecargs(P.args) + True + ``` + """ + +def is_paramspeckwargs(obj: Any, /) -> TypeIs[ParamSpecKwargs]: + """ + Return whether the argument is an instance of [`ParamSpecKwargs`][typing.ParamSpecKwargs]. + + ```pycon + >>> P = ParamSpec('P') + >>> is_paramspeckwargs(P.kwargs) + True + ``` + """ + +def is_readonly(obj: Any, /) -> bool: + """ + Return whether the argument is the [`ReadOnly`][typing.ReadOnly] [special form][]. + + ```pycon + >>> is_readonly(ReadOnly) + True + ``` + """ + +def is_required(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Required`][typing.Required] [special form][]. + + ```pycon + >>> is_required(Required) + True + ``` + """ + +def is_self(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Self`][typing.Self] [special form][]. + + ```pycon + >>> is_self(Self) + True + ``` + """ + +def is_typealias(obj: Any, /) -> bool: + """ + Return whether the argument is the [`TypeAlias`][typing.TypeAlias] [special form][]. + + ```pycon + >>> is_typealias(TypeAlias) + True + ``` + """ + +def is_typeguard(obj: Any, /) -> bool: + """ + Return whether the argument is the [`TypeGuard`][typing.TypeGuard] [special form][]. + + ```pycon + >>> is_typeguard(TypeGuard) + True + ``` + """ + +def is_typeis(obj: Any, /) -> bool: + """ + Return whether the argument is the [`TypeIs`][typing.TypeIs] [special form][]. + + ```pycon + >>> is_typeis(TypeIs) + True + ``` + """ + +def is_typealiastype(obj: Any, /) -> TypeIs[TypeAliasType]: + """ + Return whether the argument is a [`TypeAliasType`][typing.TypeAliasType] instance. + + ```pycon + >>> type MyInt = int + >>> is_typealiastype(MyInt) + True + >>> MyStr = TypeAliasType("MyStr", str) + >>> is_typealiastype(MyStr): + True + >>> type MyList[T] = list[T] + >>> is_typealiastype(MyList[int]) + False + ``` + """ + +def is_unpack(obj: Any, /) -> bool: + """ + Return whether the argument is the [`Unpack`][typing.Unpack] [special form][]. + + ```pycon + >>> is_unpack(Unpack) + True + >>> is_unpack(Unpack[Ts]) + False + ``` + """ + +def is_deprecated(obj: Any, /) -> TypeIs[deprecated]: + """ + Return whether the argument is a [`deprecated`][warnings.deprecated] instance. + + This also includes the [`typing_extensions` backport][typing_extensions.deprecated]. + + ```pycon + >>> is_deprecated(warnings.deprecated('message')) + True + >>> is_deprecated(typing_extensions.deprecated('deprecated')) + True + ``` + """ + +DEPRECATED_ALIASES: Final[dict[Any, type[Any]]] +"""A mapping between the deprecated typing aliases to their replacement, as per [PEP 585](https://peps.python.org/pep-0585/).""" diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/__init__.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d58f0891737def7f38e5d86dde2dbf9be0c13dce --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/__init__.py @@ -0,0 +1,390 @@ + +from .error import * + +from .tokens import * +from .events import * +from .nodes import * + +from .loader import * +from .dumper import * + +__version__ = '6.0.3' +try: + from .cyaml import * + __with_libyaml__ = True +except ImportError: + __with_libyaml__ = False + +import io + +#------------------------------------------------------------------------------ +# XXX "Warnings control" is now deprecated. Leaving in the API function to not +# break code that uses it. +#------------------------------------------------------------------------------ +def warnings(settings=None): + if settings is None: + return {} + +#------------------------------------------------------------------------------ +def scan(stream, Loader=Loader): + """ + Scan a YAML stream and produce scanning tokens. + """ + loader = Loader(stream) + try: + while loader.check_token(): + yield loader.get_token() + finally: + loader.dispose() + +def parse(stream, Loader=Loader): + """ + Parse a YAML stream and produce parsing events. + """ + loader = Loader(stream) + try: + while loader.check_event(): + yield loader.get_event() + finally: + loader.dispose() + +def compose(stream, Loader=Loader): + """ + Parse the first YAML document in a stream + and produce the corresponding representation tree. + """ + loader = Loader(stream) + try: + return loader.get_single_node() + finally: + loader.dispose() + +def compose_all(stream, Loader=Loader): + """ + Parse all YAML documents in a stream + and produce corresponding representation trees. + """ + loader = Loader(stream) + try: + while loader.check_node(): + yield loader.get_node() + finally: + loader.dispose() + +def load(stream, Loader): + """ + Parse the first YAML document in a stream + and produce the corresponding Python object. + """ + loader = Loader(stream) + try: + return loader.get_single_data() + finally: + loader.dispose() + +def load_all(stream, Loader): + """ + Parse all YAML documents in a stream + and produce corresponding Python objects. + """ + loader = Loader(stream) + try: + while loader.check_data(): + yield loader.get_data() + finally: + loader.dispose() + +def full_load(stream): + """ + Parse the first YAML document in a stream + and produce the corresponding Python object. + + Resolve all tags except those known to be + unsafe on untrusted input. + """ + return load(stream, FullLoader) + +def full_load_all(stream): + """ + Parse all YAML documents in a stream + and produce corresponding Python objects. + + Resolve all tags except those known to be + unsafe on untrusted input. + """ + return load_all(stream, FullLoader) + +def safe_load(stream): + """ + Parse the first YAML document in a stream + and produce the corresponding Python object. + + Resolve only basic YAML tags. This is known + to be safe for untrusted input. + """ + return load(stream, SafeLoader) + +def safe_load_all(stream): + """ + Parse all YAML documents in a stream + and produce corresponding Python objects. + + Resolve only basic YAML tags. This is known + to be safe for untrusted input. + """ + return load_all(stream, SafeLoader) + +def unsafe_load(stream): + """ + Parse the first YAML document in a stream + and produce the corresponding Python object. + + Resolve all tags, even those known to be + unsafe on untrusted input. + """ + return load(stream, UnsafeLoader) + +def unsafe_load_all(stream): + """ + Parse all YAML documents in a stream + and produce corresponding Python objects. + + Resolve all tags, even those known to be + unsafe on untrusted input. + """ + return load_all(stream, UnsafeLoader) + +def emit(events, stream=None, Dumper=Dumper, + canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None): + """ + Emit YAML parsing events into a stream. + If stream is None, return the produced string instead. + """ + getvalue = None + if stream is None: + stream = io.StringIO() + getvalue = stream.getvalue + dumper = Dumper(stream, canonical=canonical, indent=indent, width=width, + allow_unicode=allow_unicode, line_break=line_break) + try: + for event in events: + dumper.emit(event) + finally: + dumper.dispose() + if getvalue: + return getvalue() + +def serialize_all(nodes, stream=None, Dumper=Dumper, + canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None, + encoding=None, explicit_start=None, explicit_end=None, + version=None, tags=None): + """ + Serialize a sequence of representation trees into a YAML stream. + If stream is None, return the produced string instead. + """ + getvalue = None + if stream is None: + if encoding is None: + stream = io.StringIO() + else: + stream = io.BytesIO() + getvalue = stream.getvalue + dumper = Dumper(stream, canonical=canonical, indent=indent, width=width, + allow_unicode=allow_unicode, line_break=line_break, + encoding=encoding, version=version, tags=tags, + explicit_start=explicit_start, explicit_end=explicit_end) + try: + dumper.open() + for node in nodes: + dumper.serialize(node) + dumper.close() + finally: + dumper.dispose() + if getvalue: + return getvalue() + +def serialize(node, stream=None, Dumper=Dumper, **kwds): + """ + Serialize a representation tree into a YAML stream. + If stream is None, return the produced string instead. + """ + return serialize_all([node], stream, Dumper=Dumper, **kwds) + +def dump_all(documents, stream=None, Dumper=Dumper, + default_style=None, default_flow_style=False, + canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None, + encoding=None, explicit_start=None, explicit_end=None, + version=None, tags=None, sort_keys=True): + """ + Serialize a sequence of Python objects into a YAML stream. + If stream is None, return the produced string instead. + """ + getvalue = None + if stream is None: + if encoding is None: + stream = io.StringIO() + else: + stream = io.BytesIO() + getvalue = stream.getvalue + dumper = Dumper(stream, default_style=default_style, + default_flow_style=default_flow_style, + canonical=canonical, indent=indent, width=width, + allow_unicode=allow_unicode, line_break=line_break, + encoding=encoding, version=version, tags=tags, + explicit_start=explicit_start, explicit_end=explicit_end, sort_keys=sort_keys) + try: + dumper.open() + for data in documents: + dumper.represent(data) + dumper.close() + finally: + dumper.dispose() + if getvalue: + return getvalue() + +def dump(data, stream=None, Dumper=Dumper, **kwds): + """ + Serialize a Python object into a YAML stream. + If stream is None, return the produced string instead. + """ + return dump_all([data], stream, Dumper=Dumper, **kwds) + +def safe_dump_all(documents, stream=None, **kwds): + """ + Serialize a sequence of Python objects into a YAML stream. + Produce only basic YAML tags. + If stream is None, return the produced string instead. + """ + return dump_all(documents, stream, Dumper=SafeDumper, **kwds) + +def safe_dump(data, stream=None, **kwds): + """ + Serialize a Python object into a YAML stream. + Produce only basic YAML tags. + If stream is None, return the produced string instead. + """ + return dump_all([data], stream, Dumper=SafeDumper, **kwds) + +def add_implicit_resolver(tag, regexp, first=None, + Loader=None, Dumper=Dumper): + """ + Add an implicit scalar detector. + If an implicit scalar value matches the given regexp, + the corresponding tag is assigned to the scalar. + first is a sequence of possible initial characters or None. + """ + if Loader is None: + loader.Loader.add_implicit_resolver(tag, regexp, first) + loader.FullLoader.add_implicit_resolver(tag, regexp, first) + loader.UnsafeLoader.add_implicit_resolver(tag, regexp, first) + else: + Loader.add_implicit_resolver(tag, regexp, first) + Dumper.add_implicit_resolver(tag, regexp, first) + +def add_path_resolver(tag, path, kind=None, Loader=None, Dumper=Dumper): + """ + Add a path based resolver for the given tag. + A path is a list of keys that forms a path + to a node in the representation tree. + Keys can be string values, integers, or None. + """ + if Loader is None: + loader.Loader.add_path_resolver(tag, path, kind) + loader.FullLoader.add_path_resolver(tag, path, kind) + loader.UnsafeLoader.add_path_resolver(tag, path, kind) + else: + Loader.add_path_resolver(tag, path, kind) + Dumper.add_path_resolver(tag, path, kind) + +def add_constructor(tag, constructor, Loader=None): + """ + Add a constructor for the given tag. + Constructor is a function that accepts a Loader instance + and a node object and produces the corresponding Python object. + """ + if Loader is None: + loader.Loader.add_constructor(tag, constructor) + loader.FullLoader.add_constructor(tag, constructor) + loader.UnsafeLoader.add_constructor(tag, constructor) + else: + Loader.add_constructor(tag, constructor) + +def add_multi_constructor(tag_prefix, multi_constructor, Loader=None): + """ + Add a multi-constructor for the given tag prefix. + Multi-constructor is called for a node if its tag starts with tag_prefix. + Multi-constructor accepts a Loader instance, a tag suffix, + and a node object and produces the corresponding Python object. + """ + if Loader is None: + loader.Loader.add_multi_constructor(tag_prefix, multi_constructor) + loader.FullLoader.add_multi_constructor(tag_prefix, multi_constructor) + loader.UnsafeLoader.add_multi_constructor(tag_prefix, multi_constructor) + else: + Loader.add_multi_constructor(tag_prefix, multi_constructor) + +def add_representer(data_type, representer, Dumper=Dumper): + """ + Add a representer for the given type. + Representer is a function accepting a Dumper instance + and an instance of the given data type + and producing the corresponding representation node. + """ + Dumper.add_representer(data_type, representer) + +def add_multi_representer(data_type, multi_representer, Dumper=Dumper): + """ + Add a representer for the given type. + Multi-representer is a function accepting a Dumper instance + and an instance of the given data type or subtype + and producing the corresponding representation node. + """ + Dumper.add_multi_representer(data_type, multi_representer) + +class YAMLObjectMetaclass(type): + """ + The metaclass for YAMLObject. + """ + def __init__(cls, name, bases, kwds): + super(YAMLObjectMetaclass, cls).__init__(name, bases, kwds) + if 'yaml_tag' in kwds and kwds['yaml_tag'] is not None: + if isinstance(cls.yaml_loader, list): + for loader in cls.yaml_loader: + loader.add_constructor(cls.yaml_tag, cls.from_yaml) + else: + cls.yaml_loader.add_constructor(cls.yaml_tag, cls.from_yaml) + + cls.yaml_dumper.add_representer(cls, cls.to_yaml) + +class YAMLObject(metaclass=YAMLObjectMetaclass): + """ + An object that can dump itself to a YAML stream + and load itself from a YAML stream. + """ + + __slots__ = () # no direct instantiation, so allow immutable subclasses + + yaml_loader = [Loader, FullLoader, UnsafeLoader] + yaml_dumper = Dumper + + yaml_tag = None + yaml_flow_style = None + + @classmethod + def from_yaml(cls, loader, node): + """ + Convert a representation node to a Python object. + """ + return loader.construct_yaml_object(node, cls) + + @classmethod + def to_yaml(cls, dumper, data): + """ + Convert a Python object to a representation node. + """ + return dumper.represent_yaml_object(cls.yaml_tag, data, cls, + flow_style=cls.yaml_flow_style) + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/_yaml.cpython-311-x86_64-linux-gnu.so b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/_yaml.cpython-311-x86_64-linux-gnu.so new file mode 100644 index 0000000000000000000000000000000000000000..012b597019be7a2f716ef85e4f64d14f444a922d --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/_yaml.cpython-311-x86_64-linux-gnu.so @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:9f9b18be3b2682c5c49a6e103c851adf6e4d267a7e3f609108e3cd8d29b12431 +size 1160536 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/composer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/composer.py new file mode 100644 index 0000000000000000000000000000000000000000..6d15cb40e3b4198819c91c6f8d8b32807fcf53b2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/composer.py @@ -0,0 +1,139 @@ + +__all__ = ['Composer', 'ComposerError'] + +from .error import MarkedYAMLError +from .events import * +from .nodes import * + +class ComposerError(MarkedYAMLError): + pass + +class Composer: + + def __init__(self): + self.anchors = {} + + def check_node(self): + # Drop the STREAM-START event. + if self.check_event(StreamStartEvent): + self.get_event() + + # If there are more documents available? + return not self.check_event(StreamEndEvent) + + def get_node(self): + # Get the root node of the next document. + if not self.check_event(StreamEndEvent): + return self.compose_document() + + def get_single_node(self): + # Drop the STREAM-START event. + self.get_event() + + # Compose a document if the stream is not empty. + document = None + if not self.check_event(StreamEndEvent): + document = self.compose_document() + + # Ensure that the stream contains no more documents. + if not self.check_event(StreamEndEvent): + event = self.get_event() + raise ComposerError("expected a single document in the stream", + document.start_mark, "but found another document", + event.start_mark) + + # Drop the STREAM-END event. + self.get_event() + + return document + + def compose_document(self): + # Drop the DOCUMENT-START event. + self.get_event() + + # Compose the root node. + node = self.compose_node(None, None) + + # Drop the DOCUMENT-END event. + self.get_event() + + self.anchors = {} + return node + + def compose_node(self, parent, index): + if self.check_event(AliasEvent): + event = self.get_event() + anchor = event.anchor + if anchor not in self.anchors: + raise ComposerError(None, None, "found undefined alias %r" + % anchor, event.start_mark) + return self.anchors[anchor] + event = self.peek_event() + anchor = event.anchor + if anchor is not None: + if anchor in self.anchors: + raise ComposerError("found duplicate anchor %r; first occurrence" + % anchor, self.anchors[anchor].start_mark, + "second occurrence", event.start_mark) + self.descend_resolver(parent, index) + if self.check_event(ScalarEvent): + node = self.compose_scalar_node(anchor) + elif self.check_event(SequenceStartEvent): + node = self.compose_sequence_node(anchor) + elif self.check_event(MappingStartEvent): + node = self.compose_mapping_node(anchor) + self.ascend_resolver() + return node + + def compose_scalar_node(self, anchor): + event = self.get_event() + tag = event.tag + if tag is None or tag == '!': + tag = self.resolve(ScalarNode, event.value, event.implicit) + node = ScalarNode(tag, event.value, + event.start_mark, event.end_mark, style=event.style) + if anchor is not None: + self.anchors[anchor] = node + return node + + def compose_sequence_node(self, anchor): + start_event = self.get_event() + tag = start_event.tag + if tag is None or tag == '!': + tag = self.resolve(SequenceNode, None, start_event.implicit) + node = SequenceNode(tag, [], + start_event.start_mark, None, + flow_style=start_event.flow_style) + if anchor is not None: + self.anchors[anchor] = node + index = 0 + while not self.check_event(SequenceEndEvent): + node.value.append(self.compose_node(node, index)) + index += 1 + end_event = self.get_event() + node.end_mark = end_event.end_mark + return node + + def compose_mapping_node(self, anchor): + start_event = self.get_event() + tag = start_event.tag + if tag is None or tag == '!': + tag = self.resolve(MappingNode, None, start_event.implicit) + node = MappingNode(tag, [], + start_event.start_mark, None, + flow_style=start_event.flow_style) + if anchor is not None: + self.anchors[anchor] = node + while not self.check_event(MappingEndEvent): + #key_event = self.peek_event() + item_key = self.compose_node(node, None) + #if item_key in node.value: + # raise ComposerError("while composing a mapping", start_event.start_mark, + # "found duplicate key", key_event.start_mark) + item_value = self.compose_node(node, item_key) + #node.value[item_key] = item_value + node.value.append((item_key, item_value)) + end_event = self.get_event() + node.end_mark = end_event.end_mark + return node + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/constructor.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/constructor.py new file mode 100644 index 0000000000000000000000000000000000000000..619acd3070a4845c653fcf22a626e05158035bc2 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/constructor.py @@ -0,0 +1,748 @@ + +__all__ = [ + 'BaseConstructor', + 'SafeConstructor', + 'FullConstructor', + 'UnsafeConstructor', + 'Constructor', + 'ConstructorError' +] + +from .error import * +from .nodes import * + +import collections.abc, datetime, base64, binascii, re, sys, types + +class ConstructorError(MarkedYAMLError): + pass + +class BaseConstructor: + + yaml_constructors = {} + yaml_multi_constructors = {} + + def __init__(self): + self.constructed_objects = {} + self.recursive_objects = {} + self.state_generators = [] + self.deep_construct = False + + def check_data(self): + # If there are more documents available? + return self.check_node() + + def check_state_key(self, key): + """Block special attributes/methods from being set in a newly created + object, to prevent user-controlled methods from being called during + deserialization""" + if self.get_state_keys_blacklist_regexp().match(key): + raise ConstructorError(None, None, + "blacklisted key '%s' in instance state found" % (key,), None) + + def get_data(self): + # Construct and return the next document. + if self.check_node(): + return self.construct_document(self.get_node()) + + def get_single_data(self): + # Ensure that the stream contains a single document and construct it. + node = self.get_single_node() + if node is not None: + return self.construct_document(node) + return None + + def construct_document(self, node): + data = self.construct_object(node) + while self.state_generators: + state_generators = self.state_generators + self.state_generators = [] + for generator in state_generators: + for dummy in generator: + pass + self.constructed_objects = {} + self.recursive_objects = {} + self.deep_construct = False + return data + + def construct_object(self, node, deep=False): + if node in self.constructed_objects: + return self.constructed_objects[node] + if deep: + old_deep = self.deep_construct + self.deep_construct = True + if node in self.recursive_objects: + raise ConstructorError(None, None, + "found unconstructable recursive node", node.start_mark) + self.recursive_objects[node] = None + constructor = None + tag_suffix = None + if node.tag in self.yaml_constructors: + constructor = self.yaml_constructors[node.tag] + else: + for tag_prefix in self.yaml_multi_constructors: + if tag_prefix is not None and node.tag.startswith(tag_prefix): + tag_suffix = node.tag[len(tag_prefix):] + constructor = self.yaml_multi_constructors[tag_prefix] + break + else: + if None in self.yaml_multi_constructors: + tag_suffix = node.tag + constructor = self.yaml_multi_constructors[None] + elif None in self.yaml_constructors: + constructor = self.yaml_constructors[None] + elif isinstance(node, ScalarNode): + constructor = self.__class__.construct_scalar + elif isinstance(node, SequenceNode): + constructor = self.__class__.construct_sequence + elif isinstance(node, MappingNode): + constructor = self.__class__.construct_mapping + if tag_suffix is None: + data = constructor(self, node) + else: + data = constructor(self, tag_suffix, node) + if isinstance(data, types.GeneratorType): + generator = data + data = next(generator) + if self.deep_construct: + for dummy in generator: + pass + else: + self.state_generators.append(generator) + self.constructed_objects[node] = data + del self.recursive_objects[node] + if deep: + self.deep_construct = old_deep + return data + + def construct_scalar(self, node): + if not isinstance(node, ScalarNode): + raise ConstructorError(None, None, + "expected a scalar node, but found %s" % node.id, + node.start_mark) + return node.value + + def construct_sequence(self, node, deep=False): + if not isinstance(node, SequenceNode): + raise ConstructorError(None, None, + "expected a sequence node, but found %s" % node.id, + node.start_mark) + return [self.construct_object(child, deep=deep) + for child in node.value] + + def construct_mapping(self, node, deep=False): + if not isinstance(node, MappingNode): + raise ConstructorError(None, None, + "expected a mapping node, but found %s" % node.id, + node.start_mark) + mapping = {} + for key_node, value_node in node.value: + key = self.construct_object(key_node, deep=deep) + if not isinstance(key, collections.abc.Hashable): + raise ConstructorError("while constructing a mapping", node.start_mark, + "found unhashable key", key_node.start_mark) + value = self.construct_object(value_node, deep=deep) + mapping[key] = value + return mapping + + def construct_pairs(self, node, deep=False): + if not isinstance(node, MappingNode): + raise ConstructorError(None, None, + "expected a mapping node, but found %s" % node.id, + node.start_mark) + pairs = [] + for key_node, value_node in node.value: + key = self.construct_object(key_node, deep=deep) + value = self.construct_object(value_node, deep=deep) + pairs.append((key, value)) + return pairs + + @classmethod + def add_constructor(cls, tag, constructor): + if not 'yaml_constructors' in cls.__dict__: + cls.yaml_constructors = cls.yaml_constructors.copy() + cls.yaml_constructors[tag] = constructor + + @classmethod + def add_multi_constructor(cls, tag_prefix, multi_constructor): + if not 'yaml_multi_constructors' in cls.__dict__: + cls.yaml_multi_constructors = cls.yaml_multi_constructors.copy() + cls.yaml_multi_constructors[tag_prefix] = multi_constructor + +class SafeConstructor(BaseConstructor): + + def construct_scalar(self, node): + if isinstance(node, MappingNode): + for key_node, value_node in node.value: + if key_node.tag == 'tag:yaml.org,2002:value': + return self.construct_scalar(value_node) + return super().construct_scalar(node) + + def flatten_mapping(self, node): + merge = [] + index = 0 + while index < len(node.value): + key_node, value_node = node.value[index] + if key_node.tag == 'tag:yaml.org,2002:merge': + del node.value[index] + if isinstance(value_node, MappingNode): + self.flatten_mapping(value_node) + merge.extend(value_node.value) + elif isinstance(value_node, SequenceNode): + submerge = [] + for subnode in value_node.value: + if not isinstance(subnode, MappingNode): + raise ConstructorError("while constructing a mapping", + node.start_mark, + "expected a mapping for merging, but found %s" + % subnode.id, subnode.start_mark) + self.flatten_mapping(subnode) + submerge.append(subnode.value) + submerge.reverse() + for value in submerge: + merge.extend(value) + else: + raise ConstructorError("while constructing a mapping", node.start_mark, + "expected a mapping or list of mappings for merging, but found %s" + % value_node.id, value_node.start_mark) + elif key_node.tag == 'tag:yaml.org,2002:value': + key_node.tag = 'tag:yaml.org,2002:str' + index += 1 + else: + index += 1 + if merge: + node.value = merge + node.value + + def construct_mapping(self, node, deep=False): + if isinstance(node, MappingNode): + self.flatten_mapping(node) + return super().construct_mapping(node, deep=deep) + + def construct_yaml_null(self, node): + self.construct_scalar(node) + return None + + bool_values = { + 'yes': True, + 'no': False, + 'true': True, + 'false': False, + 'on': True, + 'off': False, + } + + def construct_yaml_bool(self, node): + value = self.construct_scalar(node) + return self.bool_values[value.lower()] + + def construct_yaml_int(self, node): + value = self.construct_scalar(node) + value = value.replace('_', '') + sign = +1 + if value[0] == '-': + sign = -1 + if value[0] in '+-': + value = value[1:] + if value == '0': + return 0 + elif value.startswith('0b'): + return sign*int(value[2:], 2) + elif value.startswith('0x'): + return sign*int(value[2:], 16) + elif value[0] == '0': + return sign*int(value, 8) + elif ':' in value: + digits = [int(part) for part in value.split(':')] + digits.reverse() + base = 1 + value = 0 + for digit in digits: + value += digit*base + base *= 60 + return sign*value + else: + return sign*int(value) + + inf_value = 1e300 + while inf_value != inf_value*inf_value: + inf_value *= inf_value + nan_value = -inf_value/inf_value # Trying to make a quiet NaN (like C99). + + def construct_yaml_float(self, node): + value = self.construct_scalar(node) + value = value.replace('_', '').lower() + sign = +1 + if value[0] == '-': + sign = -1 + if value[0] in '+-': + value = value[1:] + if value == '.inf': + return sign*self.inf_value + elif value == '.nan': + return self.nan_value + elif ':' in value: + digits = [float(part) for part in value.split(':')] + digits.reverse() + base = 1 + value = 0.0 + for digit in digits: + value += digit*base + base *= 60 + return sign*value + else: + return sign*float(value) + + def construct_yaml_binary(self, node): + try: + value = self.construct_scalar(node).encode('ascii') + except UnicodeEncodeError as exc: + raise ConstructorError(None, None, + "failed to convert base64 data into ascii: %s" % exc, + node.start_mark) + try: + if hasattr(base64, 'decodebytes'): + return base64.decodebytes(value) + else: + return base64.decodestring(value) + except binascii.Error as exc: + raise ConstructorError(None, None, + "failed to decode base64 data: %s" % exc, node.start_mark) + + timestamp_regexp = re.compile( + r'''^(?P[0-9][0-9][0-9][0-9]) + -(?P[0-9][0-9]?) + -(?P[0-9][0-9]?) + (?:(?:[Tt]|[ \t]+) + (?P[0-9][0-9]?) + :(?P[0-9][0-9]) + :(?P[0-9][0-9]) + (?:\.(?P[0-9]*))? + (?:[ \t]*(?PZ|(?P[-+])(?P[0-9][0-9]?) + (?::(?P[0-9][0-9]))?))?)?$''', re.X) + + def construct_yaml_timestamp(self, node): + value = self.construct_scalar(node) + match = self.timestamp_regexp.match(node.value) + values = match.groupdict() + year = int(values['year']) + month = int(values['month']) + day = int(values['day']) + if not values['hour']: + return datetime.date(year, month, day) + hour = int(values['hour']) + minute = int(values['minute']) + second = int(values['second']) + fraction = 0 + tzinfo = None + if values['fraction']: + fraction = values['fraction'][:6] + while len(fraction) < 6: + fraction += '0' + fraction = int(fraction) + if values['tz_sign']: + tz_hour = int(values['tz_hour']) + tz_minute = int(values['tz_minute'] or 0) + delta = datetime.timedelta(hours=tz_hour, minutes=tz_minute) + if values['tz_sign'] == '-': + delta = -delta + tzinfo = datetime.timezone(delta) + elif values['tz']: + tzinfo = datetime.timezone.utc + return datetime.datetime(year, month, day, hour, minute, second, fraction, + tzinfo=tzinfo) + + def construct_yaml_omap(self, node): + # Note: we do not check for duplicate keys, because it's too + # CPU-expensive. + omap = [] + yield omap + if not isinstance(node, SequenceNode): + raise ConstructorError("while constructing an ordered map", node.start_mark, + "expected a sequence, but found %s" % node.id, node.start_mark) + for subnode in node.value: + if not isinstance(subnode, MappingNode): + raise ConstructorError("while constructing an ordered map", node.start_mark, + "expected a mapping of length 1, but found %s" % subnode.id, + subnode.start_mark) + if len(subnode.value) != 1: + raise ConstructorError("while constructing an ordered map", node.start_mark, + "expected a single mapping item, but found %d items" % len(subnode.value), + subnode.start_mark) + key_node, value_node = subnode.value[0] + key = self.construct_object(key_node) + value = self.construct_object(value_node) + omap.append((key, value)) + + def construct_yaml_pairs(self, node): + # Note: the same code as `construct_yaml_omap`. + pairs = [] + yield pairs + if not isinstance(node, SequenceNode): + raise ConstructorError("while constructing pairs", node.start_mark, + "expected a sequence, but found %s" % node.id, node.start_mark) + for subnode in node.value: + if not isinstance(subnode, MappingNode): + raise ConstructorError("while constructing pairs", node.start_mark, + "expected a mapping of length 1, but found %s" % subnode.id, + subnode.start_mark) + if len(subnode.value) != 1: + raise ConstructorError("while constructing pairs", node.start_mark, + "expected a single mapping item, but found %d items" % len(subnode.value), + subnode.start_mark) + key_node, value_node = subnode.value[0] + key = self.construct_object(key_node) + value = self.construct_object(value_node) + pairs.append((key, value)) + + def construct_yaml_set(self, node): + data = set() + yield data + value = self.construct_mapping(node) + data.update(value) + + def construct_yaml_str(self, node): + return self.construct_scalar(node) + + def construct_yaml_seq(self, node): + data = [] + yield data + data.extend(self.construct_sequence(node)) + + def construct_yaml_map(self, node): + data = {} + yield data + value = self.construct_mapping(node) + data.update(value) + + def construct_yaml_object(self, node, cls): + data = cls.__new__(cls) + yield data + if hasattr(data, '__setstate__'): + state = self.construct_mapping(node, deep=True) + data.__setstate__(state) + else: + state = self.construct_mapping(node) + data.__dict__.update(state) + + def construct_undefined(self, node): + raise ConstructorError(None, None, + "could not determine a constructor for the tag %r" % node.tag, + node.start_mark) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:null', + SafeConstructor.construct_yaml_null) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:bool', + SafeConstructor.construct_yaml_bool) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:int', + SafeConstructor.construct_yaml_int) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:float', + SafeConstructor.construct_yaml_float) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:binary', + SafeConstructor.construct_yaml_binary) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:timestamp', + SafeConstructor.construct_yaml_timestamp) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:omap', + SafeConstructor.construct_yaml_omap) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:pairs', + SafeConstructor.construct_yaml_pairs) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:set', + SafeConstructor.construct_yaml_set) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:str', + SafeConstructor.construct_yaml_str) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:seq', + SafeConstructor.construct_yaml_seq) + +SafeConstructor.add_constructor( + 'tag:yaml.org,2002:map', + SafeConstructor.construct_yaml_map) + +SafeConstructor.add_constructor(None, + SafeConstructor.construct_undefined) + +class FullConstructor(SafeConstructor): + # 'extend' is blacklisted because it is used by + # construct_python_object_apply to add `listitems` to a newly generate + # python instance + def get_state_keys_blacklist(self): + return ['^extend$', '^__.*__$'] + + def get_state_keys_blacklist_regexp(self): + if not hasattr(self, 'state_keys_blacklist_regexp'): + self.state_keys_blacklist_regexp = re.compile('(' + '|'.join(self.get_state_keys_blacklist()) + ')') + return self.state_keys_blacklist_regexp + + def construct_python_str(self, node): + return self.construct_scalar(node) + + def construct_python_unicode(self, node): + return self.construct_scalar(node) + + def construct_python_bytes(self, node): + try: + value = self.construct_scalar(node).encode('ascii') + except UnicodeEncodeError as exc: + raise ConstructorError(None, None, + "failed to convert base64 data into ascii: %s" % exc, + node.start_mark) + try: + if hasattr(base64, 'decodebytes'): + return base64.decodebytes(value) + else: + return base64.decodestring(value) + except binascii.Error as exc: + raise ConstructorError(None, None, + "failed to decode base64 data: %s" % exc, node.start_mark) + + def construct_python_long(self, node): + return self.construct_yaml_int(node) + + def construct_python_complex(self, node): + return complex(self.construct_scalar(node)) + + def construct_python_tuple(self, node): + return tuple(self.construct_sequence(node)) + + def find_python_module(self, name, mark, unsafe=False): + if not name: + raise ConstructorError("while constructing a Python module", mark, + "expected non-empty name appended to the tag", mark) + if unsafe: + try: + __import__(name) + except ImportError as exc: + raise ConstructorError("while constructing a Python module", mark, + "cannot find module %r (%s)" % (name, exc), mark) + if name not in sys.modules: + raise ConstructorError("while constructing a Python module", mark, + "module %r is not imported" % name, mark) + return sys.modules[name] + + def find_python_name(self, name, mark, unsafe=False): + if not name: + raise ConstructorError("while constructing a Python object", mark, + "expected non-empty name appended to the tag", mark) + if '.' in name: + module_name, object_name = name.rsplit('.', 1) + else: + module_name = 'builtins' + object_name = name + if unsafe: + try: + __import__(module_name) + except ImportError as exc: + raise ConstructorError("while constructing a Python object", mark, + "cannot find module %r (%s)" % (module_name, exc), mark) + if module_name not in sys.modules: + raise ConstructorError("while constructing a Python object", mark, + "module %r is not imported" % module_name, mark) + module = sys.modules[module_name] + if not hasattr(module, object_name): + raise ConstructorError("while constructing a Python object", mark, + "cannot find %r in the module %r" + % (object_name, module.__name__), mark) + return getattr(module, object_name) + + def construct_python_name(self, suffix, node): + value = self.construct_scalar(node) + if value: + raise ConstructorError("while constructing a Python name", node.start_mark, + "expected the empty value, but found %r" % value, node.start_mark) + return self.find_python_name(suffix, node.start_mark) + + def construct_python_module(self, suffix, node): + value = self.construct_scalar(node) + if value: + raise ConstructorError("while constructing a Python module", node.start_mark, + "expected the empty value, but found %r" % value, node.start_mark) + return self.find_python_module(suffix, node.start_mark) + + def make_python_instance(self, suffix, node, + args=None, kwds=None, newobj=False, unsafe=False): + if not args: + args = [] + if not kwds: + kwds = {} + cls = self.find_python_name(suffix, node.start_mark) + if not (unsafe or isinstance(cls, type)): + raise ConstructorError("while constructing a Python instance", node.start_mark, + "expected a class, but found %r" % type(cls), + node.start_mark) + if newobj and isinstance(cls, type): + return cls.__new__(cls, *args, **kwds) + else: + return cls(*args, **kwds) + + def set_python_instance_state(self, instance, state, unsafe=False): + if hasattr(instance, '__setstate__'): + instance.__setstate__(state) + else: + slotstate = {} + if isinstance(state, tuple) and len(state) == 2: + state, slotstate = state + if hasattr(instance, '__dict__'): + if not unsafe and state: + for key in state.keys(): + self.check_state_key(key) + instance.__dict__.update(state) + elif state: + slotstate.update(state) + for key, value in slotstate.items(): + if not unsafe: + self.check_state_key(key) + setattr(instance, key, value) + + def construct_python_object(self, suffix, node): + # Format: + # !!python/object:module.name { ... state ... } + instance = self.make_python_instance(suffix, node, newobj=True) + yield instance + deep = hasattr(instance, '__setstate__') + state = self.construct_mapping(node, deep=deep) + self.set_python_instance_state(instance, state) + + def construct_python_object_apply(self, suffix, node, newobj=False): + # Format: + # !!python/object/apply # (or !!python/object/new) + # args: [ ... arguments ... ] + # kwds: { ... keywords ... } + # state: ... state ... + # listitems: [ ... listitems ... ] + # dictitems: { ... dictitems ... } + # or short format: + # !!python/object/apply [ ... arguments ... ] + # The difference between !!python/object/apply and !!python/object/new + # is how an object is created, check make_python_instance for details. + if isinstance(node, SequenceNode): + args = self.construct_sequence(node, deep=True) + kwds = {} + state = {} + listitems = [] + dictitems = {} + else: + value = self.construct_mapping(node, deep=True) + args = value.get('args', []) + kwds = value.get('kwds', {}) + state = value.get('state', {}) + listitems = value.get('listitems', []) + dictitems = value.get('dictitems', {}) + instance = self.make_python_instance(suffix, node, args, kwds, newobj) + if state: + self.set_python_instance_state(instance, state) + if listitems: + instance.extend(listitems) + if dictitems: + for key in dictitems: + instance[key] = dictitems[key] + return instance + + def construct_python_object_new(self, suffix, node): + return self.construct_python_object_apply(suffix, node, newobj=True) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/none', + FullConstructor.construct_yaml_null) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/bool', + FullConstructor.construct_yaml_bool) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/str', + FullConstructor.construct_python_str) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/unicode', + FullConstructor.construct_python_unicode) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/bytes', + FullConstructor.construct_python_bytes) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/int', + FullConstructor.construct_yaml_int) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/long', + FullConstructor.construct_python_long) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/float', + FullConstructor.construct_yaml_float) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/complex', + FullConstructor.construct_python_complex) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/list', + FullConstructor.construct_yaml_seq) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/tuple', + FullConstructor.construct_python_tuple) + +FullConstructor.add_constructor( + 'tag:yaml.org,2002:python/dict', + FullConstructor.construct_yaml_map) + +FullConstructor.add_multi_constructor( + 'tag:yaml.org,2002:python/name:', + FullConstructor.construct_python_name) + +class UnsafeConstructor(FullConstructor): + + def find_python_module(self, name, mark): + return super(UnsafeConstructor, self).find_python_module(name, mark, unsafe=True) + + def find_python_name(self, name, mark): + return super(UnsafeConstructor, self).find_python_name(name, mark, unsafe=True) + + def make_python_instance(self, suffix, node, args=None, kwds=None, newobj=False): + return super(UnsafeConstructor, self).make_python_instance( + suffix, node, args, kwds, newobj, unsafe=True) + + def set_python_instance_state(self, instance, state): + return super(UnsafeConstructor, self).set_python_instance_state( + instance, state, unsafe=True) + +UnsafeConstructor.add_multi_constructor( + 'tag:yaml.org,2002:python/module:', + UnsafeConstructor.construct_python_module) + +UnsafeConstructor.add_multi_constructor( + 'tag:yaml.org,2002:python/object:', + UnsafeConstructor.construct_python_object) + +UnsafeConstructor.add_multi_constructor( + 'tag:yaml.org,2002:python/object/new:', + UnsafeConstructor.construct_python_object_new) + +UnsafeConstructor.add_multi_constructor( + 'tag:yaml.org,2002:python/object/apply:', + UnsafeConstructor.construct_python_object_apply) + +# Constructor is same as UnsafeConstructor. Need to leave this in place in case +# people have extended it directly. +class Constructor(UnsafeConstructor): + pass diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/cyaml.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/cyaml.py new file mode 100644 index 0000000000000000000000000000000000000000..0c21345879b298bb8668201bebe7d289586b17f9 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/cyaml.py @@ -0,0 +1,101 @@ + +__all__ = [ + 'CBaseLoader', 'CSafeLoader', 'CFullLoader', 'CUnsafeLoader', 'CLoader', + 'CBaseDumper', 'CSafeDumper', 'CDumper' +] + +from yaml._yaml import CParser, CEmitter + +from .constructor import * + +from .serializer import * +from .representer import * + +from .resolver import * + +class CBaseLoader(CParser, BaseConstructor, BaseResolver): + + def __init__(self, stream): + CParser.__init__(self, stream) + BaseConstructor.__init__(self) + BaseResolver.__init__(self) + +class CSafeLoader(CParser, SafeConstructor, Resolver): + + def __init__(self, stream): + CParser.__init__(self, stream) + SafeConstructor.__init__(self) + Resolver.__init__(self) + +class CFullLoader(CParser, FullConstructor, Resolver): + + def __init__(self, stream): + CParser.__init__(self, stream) + FullConstructor.__init__(self) + Resolver.__init__(self) + +class CUnsafeLoader(CParser, UnsafeConstructor, Resolver): + + def __init__(self, stream): + CParser.__init__(self, stream) + UnsafeConstructor.__init__(self) + Resolver.__init__(self) + +class CLoader(CParser, Constructor, Resolver): + + def __init__(self, stream): + CParser.__init__(self, stream) + Constructor.__init__(self) + Resolver.__init__(self) + +class CBaseDumper(CEmitter, BaseRepresenter, BaseResolver): + + def __init__(self, stream, + default_style=None, default_flow_style=False, + canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None, + encoding=None, explicit_start=None, explicit_end=None, + version=None, tags=None, sort_keys=True): + CEmitter.__init__(self, stream, canonical=canonical, + indent=indent, width=width, encoding=encoding, + allow_unicode=allow_unicode, line_break=line_break, + explicit_start=explicit_start, explicit_end=explicit_end, + version=version, tags=tags) + Representer.__init__(self, default_style=default_style, + default_flow_style=default_flow_style, sort_keys=sort_keys) + Resolver.__init__(self) + +class CSafeDumper(CEmitter, SafeRepresenter, Resolver): + + def __init__(self, stream, + default_style=None, default_flow_style=False, + canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None, + encoding=None, explicit_start=None, explicit_end=None, + version=None, tags=None, sort_keys=True): + CEmitter.__init__(self, stream, canonical=canonical, + indent=indent, width=width, encoding=encoding, + allow_unicode=allow_unicode, line_break=line_break, + explicit_start=explicit_start, explicit_end=explicit_end, + version=version, tags=tags) + SafeRepresenter.__init__(self, default_style=default_style, + default_flow_style=default_flow_style, sort_keys=sort_keys) + Resolver.__init__(self) + +class CDumper(CEmitter, Serializer, Representer, Resolver): + + def __init__(self, stream, + default_style=None, default_flow_style=False, + canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None, + encoding=None, explicit_start=None, explicit_end=None, + version=None, tags=None, sort_keys=True): + CEmitter.__init__(self, stream, canonical=canonical, + indent=indent, width=width, encoding=encoding, + allow_unicode=allow_unicode, line_break=line_break, + explicit_start=explicit_start, explicit_end=explicit_end, + version=version, tags=tags) + Representer.__init__(self, default_style=default_style, + default_flow_style=default_flow_style, sort_keys=sort_keys) + Resolver.__init__(self) + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/dumper.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/dumper.py new file mode 100644 index 0000000000000000000000000000000000000000..6aadba551f3836b02f4752277f4b3027073defad --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/dumper.py @@ -0,0 +1,62 @@ + +__all__ = ['BaseDumper', 'SafeDumper', 'Dumper'] + +from .emitter import * +from .serializer import * +from .representer import * +from .resolver import * + +class BaseDumper(Emitter, Serializer, BaseRepresenter, BaseResolver): + + def __init__(self, stream, + default_style=None, default_flow_style=False, + canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None, + encoding=None, explicit_start=None, explicit_end=None, + version=None, tags=None, sort_keys=True): + Emitter.__init__(self, stream, canonical=canonical, + indent=indent, width=width, + allow_unicode=allow_unicode, line_break=line_break) + Serializer.__init__(self, encoding=encoding, + explicit_start=explicit_start, explicit_end=explicit_end, + version=version, tags=tags) + Representer.__init__(self, default_style=default_style, + default_flow_style=default_flow_style, sort_keys=sort_keys) + Resolver.__init__(self) + +class SafeDumper(Emitter, Serializer, SafeRepresenter, Resolver): + + def __init__(self, stream, + default_style=None, default_flow_style=False, + canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None, + encoding=None, explicit_start=None, explicit_end=None, + version=None, tags=None, sort_keys=True): + Emitter.__init__(self, stream, canonical=canonical, + indent=indent, width=width, + allow_unicode=allow_unicode, line_break=line_break) + Serializer.__init__(self, encoding=encoding, + explicit_start=explicit_start, explicit_end=explicit_end, + version=version, tags=tags) + SafeRepresenter.__init__(self, default_style=default_style, + default_flow_style=default_flow_style, sort_keys=sort_keys) + Resolver.__init__(self) + +class Dumper(Emitter, Serializer, Representer, Resolver): + + def __init__(self, stream, + default_style=None, default_flow_style=False, + canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None, + encoding=None, explicit_start=None, explicit_end=None, + version=None, tags=None, sort_keys=True): + Emitter.__init__(self, stream, canonical=canonical, + indent=indent, width=width, + allow_unicode=allow_unicode, line_break=line_break) + Serializer.__init__(self, encoding=encoding, + explicit_start=explicit_start, explicit_end=explicit_end, + version=version, tags=tags) + Representer.__init__(self, default_style=default_style, + default_flow_style=default_flow_style, sort_keys=sort_keys) + Resolver.__init__(self) + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/emitter.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/emitter.py new file mode 100644 index 0000000000000000000000000000000000000000..a664d011162af69184df2f8e59ab7feec818f7c7 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/emitter.py @@ -0,0 +1,1137 @@ + +# Emitter expects events obeying the following grammar: +# stream ::= STREAM-START document* STREAM-END +# document ::= DOCUMENT-START node DOCUMENT-END +# node ::= SCALAR | sequence | mapping +# sequence ::= SEQUENCE-START node* SEQUENCE-END +# mapping ::= MAPPING-START (node node)* MAPPING-END + +__all__ = ['Emitter', 'EmitterError'] + +from .error import YAMLError +from .events import * + +class EmitterError(YAMLError): + pass + +class ScalarAnalysis: + def __init__(self, scalar, empty, multiline, + allow_flow_plain, allow_block_plain, + allow_single_quoted, allow_double_quoted, + allow_block): + self.scalar = scalar + self.empty = empty + self.multiline = multiline + self.allow_flow_plain = allow_flow_plain + self.allow_block_plain = allow_block_plain + self.allow_single_quoted = allow_single_quoted + self.allow_double_quoted = allow_double_quoted + self.allow_block = allow_block + +class Emitter: + + DEFAULT_TAG_PREFIXES = { + '!' : '!', + 'tag:yaml.org,2002:' : '!!', + } + + def __init__(self, stream, canonical=None, indent=None, width=None, + allow_unicode=None, line_break=None): + + # The stream should have the methods `write` and possibly `flush`. + self.stream = stream + + # Encoding can be overridden by STREAM-START. + self.encoding = None + + # Emitter is a state machine with a stack of states to handle nested + # structures. + self.states = [] + self.state = self.expect_stream_start + + # Current event and the event queue. + self.events = [] + self.event = None + + # The current indentation level and the stack of previous indents. + self.indents = [] + self.indent = None + + # Flow level. + self.flow_level = 0 + + # Contexts. + self.root_context = False + self.sequence_context = False + self.mapping_context = False + self.simple_key_context = False + + # Characteristics of the last emitted character: + # - current position. + # - is it a whitespace? + # - is it an indention character + # (indentation space, '-', '?', or ':')? + self.line = 0 + self.column = 0 + self.whitespace = True + self.indention = True + + # Whether the document requires an explicit document indicator + self.open_ended = False + + # Formatting details. + self.canonical = canonical + self.allow_unicode = allow_unicode + self.best_indent = 2 + if indent and 1 < indent < 10: + self.best_indent = indent + self.best_width = 80 + if width and width > self.best_indent*2: + self.best_width = width + self.best_line_break = '\n' + if line_break in ['\r', '\n', '\r\n']: + self.best_line_break = line_break + + # Tag prefixes. + self.tag_prefixes = None + + # Prepared anchor and tag. + self.prepared_anchor = None + self.prepared_tag = None + + # Scalar analysis and style. + self.analysis = None + self.style = None + + def dispose(self): + # Reset the state attributes (to clear self-references) + self.states = [] + self.state = None + + def emit(self, event): + self.events.append(event) + while not self.need_more_events(): + self.event = self.events.pop(0) + self.state() + self.event = None + + # In some cases, we wait for a few next events before emitting. + + def need_more_events(self): + if not self.events: + return True + event = self.events[0] + if isinstance(event, DocumentStartEvent): + return self.need_events(1) + elif isinstance(event, SequenceStartEvent): + return self.need_events(2) + elif isinstance(event, MappingStartEvent): + return self.need_events(3) + else: + return False + + def need_events(self, count): + level = 0 + for event in self.events[1:]: + if isinstance(event, (DocumentStartEvent, CollectionStartEvent)): + level += 1 + elif isinstance(event, (DocumentEndEvent, CollectionEndEvent)): + level -= 1 + elif isinstance(event, StreamEndEvent): + level = -1 + if level < 0: + return False + return (len(self.events) < count+1) + + def increase_indent(self, flow=False, indentless=False): + self.indents.append(self.indent) + if self.indent is None: + if flow: + self.indent = self.best_indent + else: + self.indent = 0 + elif not indentless: + self.indent += self.best_indent + + # States. + + # Stream handlers. + + def expect_stream_start(self): + if isinstance(self.event, StreamStartEvent): + if self.event.encoding and not hasattr(self.stream, 'encoding'): + self.encoding = self.event.encoding + self.write_stream_start() + self.state = self.expect_first_document_start + else: + raise EmitterError("expected StreamStartEvent, but got %s" + % self.event) + + def expect_nothing(self): + raise EmitterError("expected nothing, but got %s" % self.event) + + # Document handlers. + + def expect_first_document_start(self): + return self.expect_document_start(first=True) + + def expect_document_start(self, first=False): + if isinstance(self.event, DocumentStartEvent): + if (self.event.version or self.event.tags) and self.open_ended: + self.write_indicator('...', True) + self.write_indent() + if self.event.version: + version_text = self.prepare_version(self.event.version) + self.write_version_directive(version_text) + self.tag_prefixes = self.DEFAULT_TAG_PREFIXES.copy() + if self.event.tags: + handles = sorted(self.event.tags.keys()) + for handle in handles: + prefix = self.event.tags[handle] + self.tag_prefixes[prefix] = handle + handle_text = self.prepare_tag_handle(handle) + prefix_text = self.prepare_tag_prefix(prefix) + self.write_tag_directive(handle_text, prefix_text) + implicit = (first and not self.event.explicit and not self.canonical + and not self.event.version and not self.event.tags + and not self.check_empty_document()) + if not implicit: + self.write_indent() + self.write_indicator('---', True) + if self.canonical: + self.write_indent() + self.state = self.expect_document_root + elif isinstance(self.event, StreamEndEvent): + if self.open_ended: + self.write_indicator('...', True) + self.write_indent() + self.write_stream_end() + self.state = self.expect_nothing + else: + raise EmitterError("expected DocumentStartEvent, but got %s" + % self.event) + + def expect_document_end(self): + if isinstance(self.event, DocumentEndEvent): + self.write_indent() + if self.event.explicit: + self.write_indicator('...', True) + self.write_indent() + self.flush_stream() + self.state = self.expect_document_start + else: + raise EmitterError("expected DocumentEndEvent, but got %s" + % self.event) + + def expect_document_root(self): + self.states.append(self.expect_document_end) + self.expect_node(root=True) + + # Node handlers. + + def expect_node(self, root=False, sequence=False, mapping=False, + simple_key=False): + self.root_context = root + self.sequence_context = sequence + self.mapping_context = mapping + self.simple_key_context = simple_key + if isinstance(self.event, AliasEvent): + self.expect_alias() + elif isinstance(self.event, (ScalarEvent, CollectionStartEvent)): + self.process_anchor('&') + self.process_tag() + if isinstance(self.event, ScalarEvent): + self.expect_scalar() + elif isinstance(self.event, SequenceStartEvent): + if self.flow_level or self.canonical or self.event.flow_style \ + or self.check_empty_sequence(): + self.expect_flow_sequence() + else: + self.expect_block_sequence() + elif isinstance(self.event, MappingStartEvent): + if self.flow_level or self.canonical or self.event.flow_style \ + or self.check_empty_mapping(): + self.expect_flow_mapping() + else: + self.expect_block_mapping() + else: + raise EmitterError("expected NodeEvent, but got %s" % self.event) + + def expect_alias(self): + if self.event.anchor is None: + raise EmitterError("anchor is not specified for alias") + self.process_anchor('*') + self.state = self.states.pop() + + def expect_scalar(self): + self.increase_indent(flow=True) + self.process_scalar() + self.indent = self.indents.pop() + self.state = self.states.pop() + + # Flow sequence handlers. + + def expect_flow_sequence(self): + self.write_indicator('[', True, whitespace=True) + self.flow_level += 1 + self.increase_indent(flow=True) + self.state = self.expect_first_flow_sequence_item + + def expect_first_flow_sequence_item(self): + if isinstance(self.event, SequenceEndEvent): + self.indent = self.indents.pop() + self.flow_level -= 1 + self.write_indicator(']', False) + self.state = self.states.pop() + else: + if self.canonical or self.column > self.best_width: + self.write_indent() + self.states.append(self.expect_flow_sequence_item) + self.expect_node(sequence=True) + + def expect_flow_sequence_item(self): + if isinstance(self.event, SequenceEndEvent): + self.indent = self.indents.pop() + self.flow_level -= 1 + if self.canonical: + self.write_indicator(',', False) + self.write_indent() + self.write_indicator(']', False) + self.state = self.states.pop() + else: + self.write_indicator(',', False) + if self.canonical or self.column > self.best_width: + self.write_indent() + self.states.append(self.expect_flow_sequence_item) + self.expect_node(sequence=True) + + # Flow mapping handlers. + + def expect_flow_mapping(self): + self.write_indicator('{', True, whitespace=True) + self.flow_level += 1 + self.increase_indent(flow=True) + self.state = self.expect_first_flow_mapping_key + + def expect_first_flow_mapping_key(self): + if isinstance(self.event, MappingEndEvent): + self.indent = self.indents.pop() + self.flow_level -= 1 + self.write_indicator('}', False) + self.state = self.states.pop() + else: + if self.canonical or self.column > self.best_width: + self.write_indent() + if not self.canonical and self.check_simple_key(): + self.states.append(self.expect_flow_mapping_simple_value) + self.expect_node(mapping=True, simple_key=True) + else: + self.write_indicator('?', True) + self.states.append(self.expect_flow_mapping_value) + self.expect_node(mapping=True) + + def expect_flow_mapping_key(self): + if isinstance(self.event, MappingEndEvent): + self.indent = self.indents.pop() + self.flow_level -= 1 + if self.canonical: + self.write_indicator(',', False) + self.write_indent() + self.write_indicator('}', False) + self.state = self.states.pop() + else: + self.write_indicator(',', False) + if self.canonical or self.column > self.best_width: + self.write_indent() + if not self.canonical and self.check_simple_key(): + self.states.append(self.expect_flow_mapping_simple_value) + self.expect_node(mapping=True, simple_key=True) + else: + self.write_indicator('?', True) + self.states.append(self.expect_flow_mapping_value) + self.expect_node(mapping=True) + + def expect_flow_mapping_simple_value(self): + self.write_indicator(':', False) + self.states.append(self.expect_flow_mapping_key) + self.expect_node(mapping=True) + + def expect_flow_mapping_value(self): + if self.canonical or self.column > self.best_width: + self.write_indent() + self.write_indicator(':', True) + self.states.append(self.expect_flow_mapping_key) + self.expect_node(mapping=True) + + # Block sequence handlers. + + def expect_block_sequence(self): + indentless = (self.mapping_context and not self.indention) + self.increase_indent(flow=False, indentless=indentless) + self.state = self.expect_first_block_sequence_item + + def expect_first_block_sequence_item(self): + return self.expect_block_sequence_item(first=True) + + def expect_block_sequence_item(self, first=False): + if not first and isinstance(self.event, SequenceEndEvent): + self.indent = self.indents.pop() + self.state = self.states.pop() + else: + self.write_indent() + self.write_indicator('-', True, indention=True) + self.states.append(self.expect_block_sequence_item) + self.expect_node(sequence=True) + + # Block mapping handlers. + + def expect_block_mapping(self): + self.increase_indent(flow=False) + self.state = self.expect_first_block_mapping_key + + def expect_first_block_mapping_key(self): + return self.expect_block_mapping_key(first=True) + + def expect_block_mapping_key(self, first=False): + if not first and isinstance(self.event, MappingEndEvent): + self.indent = self.indents.pop() + self.state = self.states.pop() + else: + self.write_indent() + if self.check_simple_key(): + self.states.append(self.expect_block_mapping_simple_value) + self.expect_node(mapping=True, simple_key=True) + else: + self.write_indicator('?', True, indention=True) + self.states.append(self.expect_block_mapping_value) + self.expect_node(mapping=True) + + def expect_block_mapping_simple_value(self): + self.write_indicator(':', False) + self.states.append(self.expect_block_mapping_key) + self.expect_node(mapping=True) + + def expect_block_mapping_value(self): + self.write_indent() + self.write_indicator(':', True, indention=True) + self.states.append(self.expect_block_mapping_key) + self.expect_node(mapping=True) + + # Checkers. + + def check_empty_sequence(self): + return (isinstance(self.event, SequenceStartEvent) and self.events + and isinstance(self.events[0], SequenceEndEvent)) + + def check_empty_mapping(self): + return (isinstance(self.event, MappingStartEvent) and self.events + and isinstance(self.events[0], MappingEndEvent)) + + def check_empty_document(self): + if not isinstance(self.event, DocumentStartEvent) or not self.events: + return False + event = self.events[0] + return (isinstance(event, ScalarEvent) and event.anchor is None + and event.tag is None and event.implicit and event.value == '') + + def check_simple_key(self): + length = 0 + if isinstance(self.event, NodeEvent) and self.event.anchor is not None: + if self.prepared_anchor is None: + self.prepared_anchor = self.prepare_anchor(self.event.anchor) + length += len(self.prepared_anchor) + if isinstance(self.event, (ScalarEvent, CollectionStartEvent)) \ + and self.event.tag is not None: + if self.prepared_tag is None: + self.prepared_tag = self.prepare_tag(self.event.tag) + length += len(self.prepared_tag) + if isinstance(self.event, ScalarEvent): + if self.analysis is None: + self.analysis = self.analyze_scalar(self.event.value) + length += len(self.analysis.scalar) + return (length < 128 and (isinstance(self.event, AliasEvent) + or (isinstance(self.event, ScalarEvent) + and not self.analysis.empty and not self.analysis.multiline) + or self.check_empty_sequence() or self.check_empty_mapping())) + + # Anchor, Tag, and Scalar processors. + + def process_anchor(self, indicator): + if self.event.anchor is None: + self.prepared_anchor = None + return + if self.prepared_anchor is None: + self.prepared_anchor = self.prepare_anchor(self.event.anchor) + if self.prepared_anchor: + self.write_indicator(indicator+self.prepared_anchor, True) + self.prepared_anchor = None + + def process_tag(self): + tag = self.event.tag + if isinstance(self.event, ScalarEvent): + if self.style is None: + self.style = self.choose_scalar_style() + if ((not self.canonical or tag is None) and + ((self.style == '' and self.event.implicit[0]) + or (self.style != '' and self.event.implicit[1]))): + self.prepared_tag = None + return + if self.event.implicit[0] and tag is None: + tag = '!' + self.prepared_tag = None + else: + if (not self.canonical or tag is None) and self.event.implicit: + self.prepared_tag = None + return + if tag is None: + raise EmitterError("tag is not specified") + if self.prepared_tag is None: + self.prepared_tag = self.prepare_tag(tag) + if self.prepared_tag: + self.write_indicator(self.prepared_tag, True) + self.prepared_tag = None + + def choose_scalar_style(self): + if self.analysis is None: + self.analysis = self.analyze_scalar(self.event.value) + if self.event.style == '"' or self.canonical: + return '"' + if not self.event.style and self.event.implicit[0]: + if (not (self.simple_key_context and + (self.analysis.empty or self.analysis.multiline)) + and (self.flow_level and self.analysis.allow_flow_plain + or (not self.flow_level and self.analysis.allow_block_plain))): + return '' + if self.event.style and self.event.style in '|>': + if (not self.flow_level and not self.simple_key_context + and self.analysis.allow_block): + return self.event.style + if not self.event.style or self.event.style == '\'': + if (self.analysis.allow_single_quoted and + not (self.simple_key_context and self.analysis.multiline)): + return '\'' + return '"' + + def process_scalar(self): + if self.analysis is None: + self.analysis = self.analyze_scalar(self.event.value) + if self.style is None: + self.style = self.choose_scalar_style() + split = (not self.simple_key_context) + #if self.analysis.multiline and split \ + # and (not self.style or self.style in '\'\"'): + # self.write_indent() + if self.style == '"': + self.write_double_quoted(self.analysis.scalar, split) + elif self.style == '\'': + self.write_single_quoted(self.analysis.scalar, split) + elif self.style == '>': + self.write_folded(self.analysis.scalar) + elif self.style == '|': + self.write_literal(self.analysis.scalar) + else: + self.write_plain(self.analysis.scalar, split) + self.analysis = None + self.style = None + + # Analyzers. + + def prepare_version(self, version): + major, minor = version + if major != 1: + raise EmitterError("unsupported YAML version: %d.%d" % (major, minor)) + return '%d.%d' % (major, minor) + + def prepare_tag_handle(self, handle): + if not handle: + raise EmitterError("tag handle must not be empty") + if handle[0] != '!' or handle[-1] != '!': + raise EmitterError("tag handle must start and end with '!': %r" % handle) + for ch in handle[1:-1]: + if not ('0' <= ch <= '9' or 'A' <= ch <= 'Z' or 'a' <= ch <= 'z' \ + or ch in '-_'): + raise EmitterError("invalid character %r in the tag handle: %r" + % (ch, handle)) + return handle + + def prepare_tag_prefix(self, prefix): + if not prefix: + raise EmitterError("tag prefix must not be empty") + chunks = [] + start = end = 0 + if prefix[0] == '!': + end = 1 + while end < len(prefix): + ch = prefix[end] + if '0' <= ch <= '9' or 'A' <= ch <= 'Z' or 'a' <= ch <= 'z' \ + or ch in '-;/?!:@&=+$,_.~*\'()[]': + end += 1 + else: + if start < end: + chunks.append(prefix[start:end]) + start = end = end+1 + data = ch.encode('utf-8') + for ch in data: + chunks.append('%%%02X' % ord(ch)) + if start < end: + chunks.append(prefix[start:end]) + return ''.join(chunks) + + def prepare_tag(self, tag): + if not tag: + raise EmitterError("tag must not be empty") + if tag == '!': + return tag + handle = None + suffix = tag + prefixes = sorted(self.tag_prefixes.keys()) + for prefix in prefixes: + if tag.startswith(prefix) \ + and (prefix == '!' or len(prefix) < len(tag)): + handle = self.tag_prefixes[prefix] + suffix = tag[len(prefix):] + chunks = [] + start = end = 0 + while end < len(suffix): + ch = suffix[end] + if '0' <= ch <= '9' or 'A' <= ch <= 'Z' or 'a' <= ch <= 'z' \ + or ch in '-;/?:@&=+$,_.~*\'()[]' \ + or (ch == '!' and handle != '!'): + end += 1 + else: + if start < end: + chunks.append(suffix[start:end]) + start = end = end+1 + data = ch.encode('utf-8') + for ch in data: + chunks.append('%%%02X' % ch) + if start < end: + chunks.append(suffix[start:end]) + suffix_text = ''.join(chunks) + if handle: + return '%s%s' % (handle, suffix_text) + else: + return '!<%s>' % suffix_text + + def prepare_anchor(self, anchor): + if not anchor: + raise EmitterError("anchor must not be empty") + for ch in anchor: + if not ('0' <= ch <= '9' or 'A' <= ch <= 'Z' or 'a' <= ch <= 'z' \ + or ch in '-_'): + raise EmitterError("invalid character %r in the anchor: %r" + % (ch, anchor)) + return anchor + + def analyze_scalar(self, scalar): + + # Empty scalar is a special case. + if not scalar: + return ScalarAnalysis(scalar=scalar, empty=True, multiline=False, + allow_flow_plain=False, allow_block_plain=True, + allow_single_quoted=True, allow_double_quoted=True, + allow_block=False) + + # Indicators and special characters. + block_indicators = False + flow_indicators = False + line_breaks = False + special_characters = False + + # Important whitespace combinations. + leading_space = False + leading_break = False + trailing_space = False + trailing_break = False + break_space = False + space_break = False + + # Check document indicators. + if scalar.startswith('---') or scalar.startswith('...'): + block_indicators = True + flow_indicators = True + + # First character or preceded by a whitespace. + preceded_by_whitespace = True + + # Last character or followed by a whitespace. + followed_by_whitespace = (len(scalar) == 1 or + scalar[1] in '\0 \t\r\n\x85\u2028\u2029') + + # The previous character is a space. + previous_space = False + + # The previous character is a break. + previous_break = False + + index = 0 + while index < len(scalar): + ch = scalar[index] + + # Check for indicators. + if index == 0: + # Leading indicators are special characters. + if ch in '#,[]{}&*!|>\'\"%@`': + flow_indicators = True + block_indicators = True + if ch in '?:': + flow_indicators = True + if followed_by_whitespace: + block_indicators = True + if ch == '-' and followed_by_whitespace: + flow_indicators = True + block_indicators = True + else: + # Some indicators cannot appear within a scalar as well. + if ch in ',?[]{}': + flow_indicators = True + if ch == ':': + flow_indicators = True + if followed_by_whitespace: + block_indicators = True + if ch == '#' and preceded_by_whitespace: + flow_indicators = True + block_indicators = True + + # Check for line breaks, special, and unicode characters. + if ch in '\n\x85\u2028\u2029': + line_breaks = True + if not (ch == '\n' or '\x20' <= ch <= '\x7E'): + if (ch == '\x85' or '\xA0' <= ch <= '\uD7FF' + or '\uE000' <= ch <= '\uFFFD' + or '\U00010000' <= ch < '\U0010ffff') and ch != '\uFEFF': + unicode_characters = True + if not self.allow_unicode: + special_characters = True + else: + special_characters = True + + # Detect important whitespace combinations. + if ch == ' ': + if index == 0: + leading_space = True + if index == len(scalar)-1: + trailing_space = True + if previous_break: + break_space = True + previous_space = True + previous_break = False + elif ch in '\n\x85\u2028\u2029': + if index == 0: + leading_break = True + if index == len(scalar)-1: + trailing_break = True + if previous_space: + space_break = True + previous_space = False + previous_break = True + else: + previous_space = False + previous_break = False + + # Prepare for the next character. + index += 1 + preceded_by_whitespace = (ch in '\0 \t\r\n\x85\u2028\u2029') + followed_by_whitespace = (index+1 >= len(scalar) or + scalar[index+1] in '\0 \t\r\n\x85\u2028\u2029') + + # Let's decide what styles are allowed. + allow_flow_plain = True + allow_block_plain = True + allow_single_quoted = True + allow_double_quoted = True + allow_block = True + + # Leading and trailing whitespaces are bad for plain scalars. + if (leading_space or leading_break + or trailing_space or trailing_break): + allow_flow_plain = allow_block_plain = False + + # We do not permit trailing spaces for block scalars. + if trailing_space: + allow_block = False + + # Spaces at the beginning of a new line are only acceptable for block + # scalars. + if break_space: + allow_flow_plain = allow_block_plain = allow_single_quoted = False + + # Spaces followed by breaks, as well as special character are only + # allowed for double quoted scalars. + if space_break or special_characters: + allow_flow_plain = allow_block_plain = \ + allow_single_quoted = allow_block = False + + # Although the plain scalar writer supports breaks, we never emit + # multiline plain scalars. + if line_breaks: + allow_flow_plain = allow_block_plain = False + + # Flow indicators are forbidden for flow plain scalars. + if flow_indicators: + allow_flow_plain = False + + # Block indicators are forbidden for block plain scalars. + if block_indicators: + allow_block_plain = False + + return ScalarAnalysis(scalar=scalar, + empty=False, multiline=line_breaks, + allow_flow_plain=allow_flow_plain, + allow_block_plain=allow_block_plain, + allow_single_quoted=allow_single_quoted, + allow_double_quoted=allow_double_quoted, + allow_block=allow_block) + + # Writers. + + def flush_stream(self): + if hasattr(self.stream, 'flush'): + self.stream.flush() + + def write_stream_start(self): + # Write BOM if needed. + if self.encoding and self.encoding.startswith('utf-16'): + self.stream.write('\uFEFF'.encode(self.encoding)) + + def write_stream_end(self): + self.flush_stream() + + def write_indicator(self, indicator, need_whitespace, + whitespace=False, indention=False): + if self.whitespace or not need_whitespace: + data = indicator + else: + data = ' '+indicator + self.whitespace = whitespace + self.indention = self.indention and indention + self.column += len(data) + self.open_ended = False + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + + def write_indent(self): + indent = self.indent or 0 + if not self.indention or self.column > indent \ + or (self.column == indent and not self.whitespace): + self.write_line_break() + if self.column < indent: + self.whitespace = True + data = ' '*(indent-self.column) + self.column = indent + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + + def write_line_break(self, data=None): + if data is None: + data = self.best_line_break + self.whitespace = True + self.indention = True + self.line += 1 + self.column = 0 + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + + def write_version_directive(self, version_text): + data = '%%YAML %s' % version_text + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + self.write_line_break() + + def write_tag_directive(self, handle_text, prefix_text): + data = '%%TAG %s %s' % (handle_text, prefix_text) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + self.write_line_break() + + # Scalar streams. + + def write_single_quoted(self, text, split=True): + self.write_indicator('\'', True) + spaces = False + breaks = False + start = end = 0 + while end <= len(text): + ch = None + if end < len(text): + ch = text[end] + if spaces: + if ch is None or ch != ' ': + if start+1 == end and self.column > self.best_width and split \ + and start != 0 and end != len(text): + self.write_indent() + else: + data = text[start:end] + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + start = end + elif breaks: + if ch is None or ch not in '\n\x85\u2028\u2029': + if text[start] == '\n': + self.write_line_break() + for br in text[start:end]: + if br == '\n': + self.write_line_break() + else: + self.write_line_break(br) + self.write_indent() + start = end + else: + if ch is None or ch in ' \n\x85\u2028\u2029' or ch == '\'': + if start < end: + data = text[start:end] + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + start = end + if ch == '\'': + data = '\'\'' + self.column += 2 + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + start = end + 1 + if ch is not None: + spaces = (ch == ' ') + breaks = (ch in '\n\x85\u2028\u2029') + end += 1 + self.write_indicator('\'', False) + + ESCAPE_REPLACEMENTS = { + '\0': '0', + '\x07': 'a', + '\x08': 'b', + '\x09': 't', + '\x0A': 'n', + '\x0B': 'v', + '\x0C': 'f', + '\x0D': 'r', + '\x1B': 'e', + '\"': '\"', + '\\': '\\', + '\x85': 'N', + '\xA0': '_', + '\u2028': 'L', + '\u2029': 'P', + } + + def write_double_quoted(self, text, split=True): + self.write_indicator('"', True) + start = end = 0 + while end <= len(text): + ch = None + if end < len(text): + ch = text[end] + if ch is None or ch in '"\\\x85\u2028\u2029\uFEFF' \ + or not ('\x20' <= ch <= '\x7E' + or (self.allow_unicode + and ('\xA0' <= ch <= '\uD7FF' + or '\uE000' <= ch <= '\uFFFD'))): + if start < end: + data = text[start:end] + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + start = end + if ch is not None: + if ch in self.ESCAPE_REPLACEMENTS: + data = '\\'+self.ESCAPE_REPLACEMENTS[ch] + elif ch <= '\xFF': + data = '\\x%02X' % ord(ch) + elif ch <= '\uFFFF': + data = '\\u%04X' % ord(ch) + else: + data = '\\U%08X' % ord(ch) + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + start = end+1 + if 0 < end < len(text)-1 and (ch == ' ' or start >= end) \ + and self.column+(end-start) > self.best_width and split: + data = text[start:end]+'\\' + if start < end: + start = end + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + self.write_indent() + self.whitespace = False + self.indention = False + if text[start] == ' ': + data = '\\' + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + end += 1 + self.write_indicator('"', False) + + def determine_block_hints(self, text): + hints = '' + if text: + if text[0] in ' \n\x85\u2028\u2029': + hints += str(self.best_indent) + if text[-1] not in '\n\x85\u2028\u2029': + hints += '-' + elif len(text) == 1 or text[-2] in '\n\x85\u2028\u2029': + hints += '+' + return hints + + def write_folded(self, text): + hints = self.determine_block_hints(text) + self.write_indicator('>'+hints, True) + if hints[-1:] == '+': + self.open_ended = True + self.write_line_break() + leading_space = True + spaces = False + breaks = True + start = end = 0 + while end <= len(text): + ch = None + if end < len(text): + ch = text[end] + if breaks: + if ch is None or ch not in '\n\x85\u2028\u2029': + if not leading_space and ch is not None and ch != ' ' \ + and text[start] == '\n': + self.write_line_break() + leading_space = (ch == ' ') + for br in text[start:end]: + if br == '\n': + self.write_line_break() + else: + self.write_line_break(br) + if ch is not None: + self.write_indent() + start = end + elif spaces: + if ch != ' ': + if start+1 == end and self.column > self.best_width: + self.write_indent() + else: + data = text[start:end] + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + start = end + else: + if ch is None or ch in ' \n\x85\u2028\u2029': + data = text[start:end] + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + if ch is None: + self.write_line_break() + start = end + if ch is not None: + breaks = (ch in '\n\x85\u2028\u2029') + spaces = (ch == ' ') + end += 1 + + def write_literal(self, text): + hints = self.determine_block_hints(text) + self.write_indicator('|'+hints, True) + if hints[-1:] == '+': + self.open_ended = True + self.write_line_break() + breaks = True + start = end = 0 + while end <= len(text): + ch = None + if end < len(text): + ch = text[end] + if breaks: + if ch is None or ch not in '\n\x85\u2028\u2029': + for br in text[start:end]: + if br == '\n': + self.write_line_break() + else: + self.write_line_break(br) + if ch is not None: + self.write_indent() + start = end + else: + if ch is None or ch in '\n\x85\u2028\u2029': + data = text[start:end] + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + if ch is None: + self.write_line_break() + start = end + if ch is not None: + breaks = (ch in '\n\x85\u2028\u2029') + end += 1 + + def write_plain(self, text, split=True): + if self.root_context: + self.open_ended = True + if not text: + return + if not self.whitespace: + data = ' ' + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + self.whitespace = False + self.indention = False + spaces = False + breaks = False + start = end = 0 + while end <= len(text): + ch = None + if end < len(text): + ch = text[end] + if spaces: + if ch != ' ': + if start+1 == end and self.column > self.best_width and split: + self.write_indent() + self.whitespace = False + self.indention = False + else: + data = text[start:end] + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + start = end + elif breaks: + if ch not in '\n\x85\u2028\u2029': + if text[start] == '\n': + self.write_line_break() + for br in text[start:end]: + if br == '\n': + self.write_line_break() + else: + self.write_line_break(br) + self.write_indent() + self.whitespace = False + self.indention = False + start = end + else: + if ch is None or ch in ' \n\x85\u2028\u2029': + data = text[start:end] + self.column += len(data) + if self.encoding: + data = data.encode(self.encoding) + self.stream.write(data) + start = end + if ch is not None: + spaces = (ch == ' ') + breaks = (ch in '\n\x85\u2028\u2029') + end += 1 diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/error.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/error.py new file mode 100644 index 0000000000000000000000000000000000000000..b796b4dc519512c4825ff539a2e6aa20f4d370d0 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/error.py @@ -0,0 +1,75 @@ + +__all__ = ['Mark', 'YAMLError', 'MarkedYAMLError'] + +class Mark: + + def __init__(self, name, index, line, column, buffer, pointer): + self.name = name + self.index = index + self.line = line + self.column = column + self.buffer = buffer + self.pointer = pointer + + def get_snippet(self, indent=4, max_length=75): + if self.buffer is None: + return None + head = '' + start = self.pointer + while start > 0 and self.buffer[start-1] not in '\0\r\n\x85\u2028\u2029': + start -= 1 + if self.pointer-start > max_length/2-1: + head = ' ... ' + start += 5 + break + tail = '' + end = self.pointer + while end < len(self.buffer) and self.buffer[end] not in '\0\r\n\x85\u2028\u2029': + end += 1 + if end-self.pointer > max_length/2-1: + tail = ' ... ' + end -= 5 + break + snippet = self.buffer[start:end] + return ' '*indent + head + snippet + tail + '\n' \ + + ' '*(indent+self.pointer-start+len(head)) + '^' + + def __str__(self): + snippet = self.get_snippet() + where = " in \"%s\", line %d, column %d" \ + % (self.name, self.line+1, self.column+1) + if snippet is not None: + where += ":\n"+snippet + return where + +class YAMLError(Exception): + pass + +class MarkedYAMLError(YAMLError): + + def __init__(self, context=None, context_mark=None, + problem=None, problem_mark=None, note=None): + self.context = context + self.context_mark = context_mark + self.problem = problem + self.problem_mark = problem_mark + self.note = note + + def __str__(self): + lines = [] + if self.context is not None: + lines.append(self.context) + if self.context_mark is not None \ + and (self.problem is None or self.problem_mark is None + or self.context_mark.name != self.problem_mark.name + or self.context_mark.line != self.problem_mark.line + or self.context_mark.column != self.problem_mark.column): + lines.append(str(self.context_mark)) + if self.problem is not None: + lines.append(self.problem) + if self.problem_mark is not None: + lines.append(str(self.problem_mark)) + if self.note is not None: + lines.append(self.note) + return '\n'.join(lines) + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/events.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/events.py new file mode 100644 index 0000000000000000000000000000000000000000..f79ad389cb6c9517e391dcd25534866bc9ccd36a --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/events.py @@ -0,0 +1,86 @@ + +# Abstract classes. + +class Event(object): + def __init__(self, start_mark=None, end_mark=None): + self.start_mark = start_mark + self.end_mark = end_mark + def __repr__(self): + attributes = [key for key in ['anchor', 'tag', 'implicit', 'value'] + if hasattr(self, key)] + arguments = ', '.join(['%s=%r' % (key, getattr(self, key)) + for key in attributes]) + return '%s(%s)' % (self.__class__.__name__, arguments) + +class NodeEvent(Event): + def __init__(self, anchor, start_mark=None, end_mark=None): + self.anchor = anchor + self.start_mark = start_mark + self.end_mark = end_mark + +class CollectionStartEvent(NodeEvent): + def __init__(self, anchor, tag, implicit, start_mark=None, end_mark=None, + flow_style=None): + self.anchor = anchor + self.tag = tag + self.implicit = implicit + self.start_mark = start_mark + self.end_mark = end_mark + self.flow_style = flow_style + +class CollectionEndEvent(Event): + pass + +# Implementations. + +class StreamStartEvent(Event): + def __init__(self, start_mark=None, end_mark=None, encoding=None): + self.start_mark = start_mark + self.end_mark = end_mark + self.encoding = encoding + +class StreamEndEvent(Event): + pass + +class DocumentStartEvent(Event): + def __init__(self, start_mark=None, end_mark=None, + explicit=None, version=None, tags=None): + self.start_mark = start_mark + self.end_mark = end_mark + self.explicit = explicit + self.version = version + self.tags = tags + +class DocumentEndEvent(Event): + def __init__(self, start_mark=None, end_mark=None, + explicit=None): + self.start_mark = start_mark + self.end_mark = end_mark + self.explicit = explicit + +class AliasEvent(NodeEvent): + pass + +class ScalarEvent(NodeEvent): + def __init__(self, anchor, tag, implicit, value, + start_mark=None, end_mark=None, style=None): + self.anchor = anchor + self.tag = tag + self.implicit = implicit + self.value = value + self.start_mark = start_mark + self.end_mark = end_mark + self.style = style + +class SequenceStartEvent(CollectionStartEvent): + pass + +class SequenceEndEvent(CollectionEndEvent): + pass + +class MappingStartEvent(CollectionStartEvent): + pass + +class MappingEndEvent(CollectionEndEvent): + pass + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/loader.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/loader.py new file mode 100644 index 0000000000000000000000000000000000000000..e90c11224c38e559cdf0cb205f0692ebd4fb8681 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/loader.py @@ -0,0 +1,63 @@ + +__all__ = ['BaseLoader', 'FullLoader', 'SafeLoader', 'Loader', 'UnsafeLoader'] + +from .reader import * +from .scanner import * +from .parser import * +from .composer import * +from .constructor import * +from .resolver import * + +class BaseLoader(Reader, Scanner, Parser, Composer, BaseConstructor, BaseResolver): + + def __init__(self, stream): + Reader.__init__(self, stream) + Scanner.__init__(self) + Parser.__init__(self) + Composer.__init__(self) + BaseConstructor.__init__(self) + BaseResolver.__init__(self) + +class FullLoader(Reader, Scanner, Parser, Composer, FullConstructor, Resolver): + + def __init__(self, stream): + Reader.__init__(self, stream) + Scanner.__init__(self) + Parser.__init__(self) + Composer.__init__(self) + FullConstructor.__init__(self) + Resolver.__init__(self) + +class SafeLoader(Reader, Scanner, Parser, Composer, SafeConstructor, Resolver): + + def __init__(self, stream): + Reader.__init__(self, stream) + Scanner.__init__(self) + Parser.__init__(self) + Composer.__init__(self) + SafeConstructor.__init__(self) + Resolver.__init__(self) + +class Loader(Reader, Scanner, Parser, Composer, Constructor, Resolver): + + def __init__(self, stream): + Reader.__init__(self, stream) + Scanner.__init__(self) + Parser.__init__(self) + Composer.__init__(self) + Constructor.__init__(self) + Resolver.__init__(self) + +# UnsafeLoader is the same as Loader (which is and was always unsafe on +# untrusted input). Use of either Loader or UnsafeLoader should be rare, since +# FullLoad should be able to load almost all YAML safely. Loader is left intact +# to ensure backwards compatibility. +class UnsafeLoader(Reader, Scanner, Parser, Composer, Constructor, Resolver): + + def __init__(self, stream): + Reader.__init__(self, stream) + Scanner.__init__(self) + Parser.__init__(self) + Composer.__init__(self) + Constructor.__init__(self) + Resolver.__init__(self) diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/nodes.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/nodes.py new file mode 100644 index 0000000000000000000000000000000000000000..c4f070c41e1fb1bc01af27d69329e92dded38908 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/nodes.py @@ -0,0 +1,49 @@ + +class Node(object): + def __init__(self, tag, value, start_mark, end_mark): + self.tag = tag + self.value = value + self.start_mark = start_mark + self.end_mark = end_mark + def __repr__(self): + value = self.value + #if isinstance(value, list): + # if len(value) == 0: + # value = '' + # elif len(value) == 1: + # value = '<1 item>' + # else: + # value = '<%d items>' % len(value) + #else: + # if len(value) > 75: + # value = repr(value[:70]+u' ... ') + # else: + # value = repr(value) + value = repr(value) + return '%s(tag=%r, value=%s)' % (self.__class__.__name__, self.tag, value) + +class ScalarNode(Node): + id = 'scalar' + def __init__(self, tag, value, + start_mark=None, end_mark=None, style=None): + self.tag = tag + self.value = value + self.start_mark = start_mark + self.end_mark = end_mark + self.style = style + +class CollectionNode(Node): + def __init__(self, tag, value, + start_mark=None, end_mark=None, flow_style=None): + self.tag = tag + self.value = value + self.start_mark = start_mark + self.end_mark = end_mark + self.flow_style = flow_style + +class SequenceNode(CollectionNode): + id = 'sequence' + +class MappingNode(CollectionNode): + id = 'mapping' + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/parser.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/parser.py new file mode 100644 index 0000000000000000000000000000000000000000..13a5995d292045d0f865a99abf692bd35dc87814 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/parser.py @@ -0,0 +1,589 @@ + +# The following YAML grammar is LL(1) and is parsed by a recursive descent +# parser. +# +# stream ::= STREAM-START implicit_document? explicit_document* STREAM-END +# implicit_document ::= block_node DOCUMENT-END* +# explicit_document ::= DIRECTIVE* DOCUMENT-START block_node? DOCUMENT-END* +# block_node_or_indentless_sequence ::= +# ALIAS +# | properties (block_content | indentless_block_sequence)? +# | block_content +# | indentless_block_sequence +# block_node ::= ALIAS +# | properties block_content? +# | block_content +# flow_node ::= ALIAS +# | properties flow_content? +# | flow_content +# properties ::= TAG ANCHOR? | ANCHOR TAG? +# block_content ::= block_collection | flow_collection | SCALAR +# flow_content ::= flow_collection | SCALAR +# block_collection ::= block_sequence | block_mapping +# flow_collection ::= flow_sequence | flow_mapping +# block_sequence ::= BLOCK-SEQUENCE-START (BLOCK-ENTRY block_node?)* BLOCK-END +# indentless_sequence ::= (BLOCK-ENTRY block_node?)+ +# block_mapping ::= BLOCK-MAPPING_START +# ((KEY block_node_or_indentless_sequence?)? +# (VALUE block_node_or_indentless_sequence?)?)* +# BLOCK-END +# flow_sequence ::= FLOW-SEQUENCE-START +# (flow_sequence_entry FLOW-ENTRY)* +# flow_sequence_entry? +# FLOW-SEQUENCE-END +# flow_sequence_entry ::= flow_node | KEY flow_node? (VALUE flow_node?)? +# flow_mapping ::= FLOW-MAPPING-START +# (flow_mapping_entry FLOW-ENTRY)* +# flow_mapping_entry? +# FLOW-MAPPING-END +# flow_mapping_entry ::= flow_node | KEY flow_node? (VALUE flow_node?)? +# +# FIRST sets: +# +# stream: { STREAM-START } +# explicit_document: { DIRECTIVE DOCUMENT-START } +# implicit_document: FIRST(block_node) +# block_node: { ALIAS TAG ANCHOR SCALAR BLOCK-SEQUENCE-START BLOCK-MAPPING-START FLOW-SEQUENCE-START FLOW-MAPPING-START } +# flow_node: { ALIAS ANCHOR TAG SCALAR FLOW-SEQUENCE-START FLOW-MAPPING-START } +# block_content: { BLOCK-SEQUENCE-START BLOCK-MAPPING-START FLOW-SEQUENCE-START FLOW-MAPPING-START SCALAR } +# flow_content: { FLOW-SEQUENCE-START FLOW-MAPPING-START SCALAR } +# block_collection: { BLOCK-SEQUENCE-START BLOCK-MAPPING-START } +# flow_collection: { FLOW-SEQUENCE-START FLOW-MAPPING-START } +# block_sequence: { BLOCK-SEQUENCE-START } +# block_mapping: { BLOCK-MAPPING-START } +# block_node_or_indentless_sequence: { ALIAS ANCHOR TAG SCALAR BLOCK-SEQUENCE-START BLOCK-MAPPING-START FLOW-SEQUENCE-START FLOW-MAPPING-START BLOCK-ENTRY } +# indentless_sequence: { ENTRY } +# flow_collection: { FLOW-SEQUENCE-START FLOW-MAPPING-START } +# flow_sequence: { FLOW-SEQUENCE-START } +# flow_mapping: { FLOW-MAPPING-START } +# flow_sequence_entry: { ALIAS ANCHOR TAG SCALAR FLOW-SEQUENCE-START FLOW-MAPPING-START KEY } +# flow_mapping_entry: { ALIAS ANCHOR TAG SCALAR FLOW-SEQUENCE-START FLOW-MAPPING-START KEY } + +__all__ = ['Parser', 'ParserError'] + +from .error import MarkedYAMLError +from .tokens import * +from .events import * +from .scanner import * + +class ParserError(MarkedYAMLError): + pass + +class Parser: + # Since writing a recursive-descendant parser is a straightforward task, we + # do not give many comments here. + + DEFAULT_TAGS = { + '!': '!', + '!!': 'tag:yaml.org,2002:', + } + + def __init__(self): + self.current_event = None + self.yaml_version = None + self.tag_handles = {} + self.states = [] + self.marks = [] + self.state = self.parse_stream_start + + def dispose(self): + # Reset the state attributes (to clear self-references) + self.states = [] + self.state = None + + def check_event(self, *choices): + # Check the type of the next event. + if self.current_event is None: + if self.state: + self.current_event = self.state() + if self.current_event is not None: + if not choices: + return True + for choice in choices: + if isinstance(self.current_event, choice): + return True + return False + + def peek_event(self): + # Get the next event. + if self.current_event is None: + if self.state: + self.current_event = self.state() + return self.current_event + + def get_event(self): + # Get the next event and proceed further. + if self.current_event is None: + if self.state: + self.current_event = self.state() + value = self.current_event + self.current_event = None + return value + + # stream ::= STREAM-START implicit_document? explicit_document* STREAM-END + # implicit_document ::= block_node DOCUMENT-END* + # explicit_document ::= DIRECTIVE* DOCUMENT-START block_node? DOCUMENT-END* + + def parse_stream_start(self): + + # Parse the stream start. + token = self.get_token() + event = StreamStartEvent(token.start_mark, token.end_mark, + encoding=token.encoding) + + # Prepare the next state. + self.state = self.parse_implicit_document_start + + return event + + def parse_implicit_document_start(self): + + # Parse an implicit document. + if not self.check_token(DirectiveToken, DocumentStartToken, + StreamEndToken): + self.tag_handles = self.DEFAULT_TAGS + token = self.peek_token() + start_mark = end_mark = token.start_mark + event = DocumentStartEvent(start_mark, end_mark, + explicit=False) + + # Prepare the next state. + self.states.append(self.parse_document_end) + self.state = self.parse_block_node + + return event + + else: + return self.parse_document_start() + + def parse_document_start(self): + + # Parse any extra document end indicators. + while self.check_token(DocumentEndToken): + self.get_token() + + # Parse an explicit document. + if not self.check_token(StreamEndToken): + token = self.peek_token() + start_mark = token.start_mark + version, tags = self.process_directives() + if not self.check_token(DocumentStartToken): + raise ParserError(None, None, + "expected '', but found %r" + % self.peek_token().id, + self.peek_token().start_mark) + token = self.get_token() + end_mark = token.end_mark + event = DocumentStartEvent(start_mark, end_mark, + explicit=True, version=version, tags=tags) + self.states.append(self.parse_document_end) + self.state = self.parse_document_content + else: + # Parse the end of the stream. + token = self.get_token() + event = StreamEndEvent(token.start_mark, token.end_mark) + assert not self.states + assert not self.marks + self.state = None + return event + + def parse_document_end(self): + + # Parse the document end. + token = self.peek_token() + start_mark = end_mark = token.start_mark + explicit = False + if self.check_token(DocumentEndToken): + token = self.get_token() + end_mark = token.end_mark + explicit = True + event = DocumentEndEvent(start_mark, end_mark, + explicit=explicit) + + # Prepare the next state. + self.state = self.parse_document_start + + return event + + def parse_document_content(self): + if self.check_token(DirectiveToken, + DocumentStartToken, DocumentEndToken, StreamEndToken): + event = self.process_empty_scalar(self.peek_token().start_mark) + self.state = self.states.pop() + return event + else: + return self.parse_block_node() + + def process_directives(self): + self.yaml_version = None + self.tag_handles = {} + while self.check_token(DirectiveToken): + token = self.get_token() + if token.name == 'YAML': + if self.yaml_version is not None: + raise ParserError(None, None, + "found duplicate YAML directive", token.start_mark) + major, minor = token.value + if major != 1: + raise ParserError(None, None, + "found incompatible YAML document (version 1.* is required)", + token.start_mark) + self.yaml_version = token.value + elif token.name == 'TAG': + handle, prefix = token.value + if handle in self.tag_handles: + raise ParserError(None, None, + "duplicate tag handle %r" % handle, + token.start_mark) + self.tag_handles[handle] = prefix + if self.tag_handles: + value = self.yaml_version, self.tag_handles.copy() + else: + value = self.yaml_version, None + for key in self.DEFAULT_TAGS: + if key not in self.tag_handles: + self.tag_handles[key] = self.DEFAULT_TAGS[key] + return value + + # block_node_or_indentless_sequence ::= ALIAS + # | properties (block_content | indentless_block_sequence)? + # | block_content + # | indentless_block_sequence + # block_node ::= ALIAS + # | properties block_content? + # | block_content + # flow_node ::= ALIAS + # | properties flow_content? + # | flow_content + # properties ::= TAG ANCHOR? | ANCHOR TAG? + # block_content ::= block_collection | flow_collection | SCALAR + # flow_content ::= flow_collection | SCALAR + # block_collection ::= block_sequence | block_mapping + # flow_collection ::= flow_sequence | flow_mapping + + def parse_block_node(self): + return self.parse_node(block=True) + + def parse_flow_node(self): + return self.parse_node() + + def parse_block_node_or_indentless_sequence(self): + return self.parse_node(block=True, indentless_sequence=True) + + def parse_node(self, block=False, indentless_sequence=False): + if self.check_token(AliasToken): + token = self.get_token() + event = AliasEvent(token.value, token.start_mark, token.end_mark) + self.state = self.states.pop() + else: + anchor = None + tag = None + start_mark = end_mark = tag_mark = None + if self.check_token(AnchorToken): + token = self.get_token() + start_mark = token.start_mark + end_mark = token.end_mark + anchor = token.value + if self.check_token(TagToken): + token = self.get_token() + tag_mark = token.start_mark + end_mark = token.end_mark + tag = token.value + elif self.check_token(TagToken): + token = self.get_token() + start_mark = tag_mark = token.start_mark + end_mark = token.end_mark + tag = token.value + if self.check_token(AnchorToken): + token = self.get_token() + end_mark = token.end_mark + anchor = token.value + if tag is not None: + handle, suffix = tag + if handle is not None: + if handle not in self.tag_handles: + raise ParserError("while parsing a node", start_mark, + "found undefined tag handle %r" % handle, + tag_mark) + tag = self.tag_handles[handle]+suffix + else: + tag = suffix + #if tag == '!': + # raise ParserError("while parsing a node", start_mark, + # "found non-specific tag '!'", tag_mark, + # "Please check 'http://pyyaml.org/wiki/YAMLNonSpecificTag' and share your opinion.") + if start_mark is None: + start_mark = end_mark = self.peek_token().start_mark + event = None + implicit = (tag is None or tag == '!') + if indentless_sequence and self.check_token(BlockEntryToken): + end_mark = self.peek_token().end_mark + event = SequenceStartEvent(anchor, tag, implicit, + start_mark, end_mark) + self.state = self.parse_indentless_sequence_entry + else: + if self.check_token(ScalarToken): + token = self.get_token() + end_mark = token.end_mark + if (token.plain and tag is None) or tag == '!': + implicit = (True, False) + elif tag is None: + implicit = (False, True) + else: + implicit = (False, False) + event = ScalarEvent(anchor, tag, implicit, token.value, + start_mark, end_mark, style=token.style) + self.state = self.states.pop() + elif self.check_token(FlowSequenceStartToken): + end_mark = self.peek_token().end_mark + event = SequenceStartEvent(anchor, tag, implicit, + start_mark, end_mark, flow_style=True) + self.state = self.parse_flow_sequence_first_entry + elif self.check_token(FlowMappingStartToken): + end_mark = self.peek_token().end_mark + event = MappingStartEvent(anchor, tag, implicit, + start_mark, end_mark, flow_style=True) + self.state = self.parse_flow_mapping_first_key + elif block and self.check_token(BlockSequenceStartToken): + end_mark = self.peek_token().start_mark + event = SequenceStartEvent(anchor, tag, implicit, + start_mark, end_mark, flow_style=False) + self.state = self.parse_block_sequence_first_entry + elif block and self.check_token(BlockMappingStartToken): + end_mark = self.peek_token().start_mark + event = MappingStartEvent(anchor, tag, implicit, + start_mark, end_mark, flow_style=False) + self.state = self.parse_block_mapping_first_key + elif anchor is not None or tag is not None: + # Empty scalars are allowed even if a tag or an anchor is + # specified. + event = ScalarEvent(anchor, tag, (implicit, False), '', + start_mark, end_mark) + self.state = self.states.pop() + else: + if block: + node = 'block' + else: + node = 'flow' + token = self.peek_token() + raise ParserError("while parsing a %s node" % node, start_mark, + "expected the node content, but found %r" % token.id, + token.start_mark) + return event + + # block_sequence ::= BLOCK-SEQUENCE-START (BLOCK-ENTRY block_node?)* BLOCK-END + + def parse_block_sequence_first_entry(self): + token = self.get_token() + self.marks.append(token.start_mark) + return self.parse_block_sequence_entry() + + def parse_block_sequence_entry(self): + if self.check_token(BlockEntryToken): + token = self.get_token() + if not self.check_token(BlockEntryToken, BlockEndToken): + self.states.append(self.parse_block_sequence_entry) + return self.parse_block_node() + else: + self.state = self.parse_block_sequence_entry + return self.process_empty_scalar(token.end_mark) + if not self.check_token(BlockEndToken): + token = self.peek_token() + raise ParserError("while parsing a block collection", self.marks[-1], + "expected , but found %r" % token.id, token.start_mark) + token = self.get_token() + event = SequenceEndEvent(token.start_mark, token.end_mark) + self.state = self.states.pop() + self.marks.pop() + return event + + # indentless_sequence ::= (BLOCK-ENTRY block_node?)+ + + def parse_indentless_sequence_entry(self): + if self.check_token(BlockEntryToken): + token = self.get_token() + if not self.check_token(BlockEntryToken, + KeyToken, ValueToken, BlockEndToken): + self.states.append(self.parse_indentless_sequence_entry) + return self.parse_block_node() + else: + self.state = self.parse_indentless_sequence_entry + return self.process_empty_scalar(token.end_mark) + token = self.peek_token() + event = SequenceEndEvent(token.start_mark, token.start_mark) + self.state = self.states.pop() + return event + + # block_mapping ::= BLOCK-MAPPING_START + # ((KEY block_node_or_indentless_sequence?)? + # (VALUE block_node_or_indentless_sequence?)?)* + # BLOCK-END + + def parse_block_mapping_first_key(self): + token = self.get_token() + self.marks.append(token.start_mark) + return self.parse_block_mapping_key() + + def parse_block_mapping_key(self): + if self.check_token(KeyToken): + token = self.get_token() + if not self.check_token(KeyToken, ValueToken, BlockEndToken): + self.states.append(self.parse_block_mapping_value) + return self.parse_block_node_or_indentless_sequence() + else: + self.state = self.parse_block_mapping_value + return self.process_empty_scalar(token.end_mark) + if not self.check_token(BlockEndToken): + token = self.peek_token() + raise ParserError("while parsing a block mapping", self.marks[-1], + "expected , but found %r" % token.id, token.start_mark) + token = self.get_token() + event = MappingEndEvent(token.start_mark, token.end_mark) + self.state = self.states.pop() + self.marks.pop() + return event + + def parse_block_mapping_value(self): + if self.check_token(ValueToken): + token = self.get_token() + if not self.check_token(KeyToken, ValueToken, BlockEndToken): + self.states.append(self.parse_block_mapping_key) + return self.parse_block_node_or_indentless_sequence() + else: + self.state = self.parse_block_mapping_key + return self.process_empty_scalar(token.end_mark) + else: + self.state = self.parse_block_mapping_key + token = self.peek_token() + return self.process_empty_scalar(token.start_mark) + + # flow_sequence ::= FLOW-SEQUENCE-START + # (flow_sequence_entry FLOW-ENTRY)* + # flow_sequence_entry? + # FLOW-SEQUENCE-END + # flow_sequence_entry ::= flow_node | KEY flow_node? (VALUE flow_node?)? + # + # Note that while production rules for both flow_sequence_entry and + # flow_mapping_entry are equal, their interpretations are different. + # For `flow_sequence_entry`, the part `KEY flow_node? (VALUE flow_node?)?` + # generate an inline mapping (set syntax). + + def parse_flow_sequence_first_entry(self): + token = self.get_token() + self.marks.append(token.start_mark) + return self.parse_flow_sequence_entry(first=True) + + def parse_flow_sequence_entry(self, first=False): + if not self.check_token(FlowSequenceEndToken): + if not first: + if self.check_token(FlowEntryToken): + self.get_token() + else: + token = self.peek_token() + raise ParserError("while parsing a flow sequence", self.marks[-1], + "expected ',' or ']', but got %r" % token.id, token.start_mark) + + if self.check_token(KeyToken): + token = self.peek_token() + event = MappingStartEvent(None, None, True, + token.start_mark, token.end_mark, + flow_style=True) + self.state = self.parse_flow_sequence_entry_mapping_key + return event + elif not self.check_token(FlowSequenceEndToken): + self.states.append(self.parse_flow_sequence_entry) + return self.parse_flow_node() + token = self.get_token() + event = SequenceEndEvent(token.start_mark, token.end_mark) + self.state = self.states.pop() + self.marks.pop() + return event + + def parse_flow_sequence_entry_mapping_key(self): + token = self.get_token() + if not self.check_token(ValueToken, + FlowEntryToken, FlowSequenceEndToken): + self.states.append(self.parse_flow_sequence_entry_mapping_value) + return self.parse_flow_node() + else: + self.state = self.parse_flow_sequence_entry_mapping_value + return self.process_empty_scalar(token.end_mark) + + def parse_flow_sequence_entry_mapping_value(self): + if self.check_token(ValueToken): + token = self.get_token() + if not self.check_token(FlowEntryToken, FlowSequenceEndToken): + self.states.append(self.parse_flow_sequence_entry_mapping_end) + return self.parse_flow_node() + else: + self.state = self.parse_flow_sequence_entry_mapping_end + return self.process_empty_scalar(token.end_mark) + else: + self.state = self.parse_flow_sequence_entry_mapping_end + token = self.peek_token() + return self.process_empty_scalar(token.start_mark) + + def parse_flow_sequence_entry_mapping_end(self): + self.state = self.parse_flow_sequence_entry + token = self.peek_token() + return MappingEndEvent(token.start_mark, token.start_mark) + + # flow_mapping ::= FLOW-MAPPING-START + # (flow_mapping_entry FLOW-ENTRY)* + # flow_mapping_entry? + # FLOW-MAPPING-END + # flow_mapping_entry ::= flow_node | KEY flow_node? (VALUE flow_node?)? + + def parse_flow_mapping_first_key(self): + token = self.get_token() + self.marks.append(token.start_mark) + return self.parse_flow_mapping_key(first=True) + + def parse_flow_mapping_key(self, first=False): + if not self.check_token(FlowMappingEndToken): + if not first: + if self.check_token(FlowEntryToken): + self.get_token() + else: + token = self.peek_token() + raise ParserError("while parsing a flow mapping", self.marks[-1], + "expected ',' or '}', but got %r" % token.id, token.start_mark) + if self.check_token(KeyToken): + token = self.get_token() + if not self.check_token(ValueToken, + FlowEntryToken, FlowMappingEndToken): + self.states.append(self.parse_flow_mapping_value) + return self.parse_flow_node() + else: + self.state = self.parse_flow_mapping_value + return self.process_empty_scalar(token.end_mark) + elif not self.check_token(FlowMappingEndToken): + self.states.append(self.parse_flow_mapping_empty_value) + return self.parse_flow_node() + token = self.get_token() + event = MappingEndEvent(token.start_mark, token.end_mark) + self.state = self.states.pop() + self.marks.pop() + return event + + def parse_flow_mapping_value(self): + if self.check_token(ValueToken): + token = self.get_token() + if not self.check_token(FlowEntryToken, FlowMappingEndToken): + self.states.append(self.parse_flow_mapping_key) + return self.parse_flow_node() + else: + self.state = self.parse_flow_mapping_key + return self.process_empty_scalar(token.end_mark) + else: + self.state = self.parse_flow_mapping_key + token = self.peek_token() + return self.process_empty_scalar(token.start_mark) + + def parse_flow_mapping_empty_value(self): + self.state = self.parse_flow_mapping_key + return self.process_empty_scalar(self.peek_token().start_mark) + + def process_empty_scalar(self, mark): + return ScalarEvent(None, None, (True, False), '', mark, mark) + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/reader.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/reader.py new file mode 100644 index 0000000000000000000000000000000000000000..774b0219b5932a0ee1c27e637371de5ba8d9cb16 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/reader.py @@ -0,0 +1,185 @@ +# This module contains abstractions for the input stream. You don't have to +# looks further, there are no pretty code. +# +# We define two classes here. +# +# Mark(source, line, column) +# It's just a record and its only use is producing nice error messages. +# Parser does not use it for any other purposes. +# +# Reader(source, data) +# Reader determines the encoding of `data` and converts it to unicode. +# Reader provides the following methods and attributes: +# reader.peek(length=1) - return the next `length` characters +# reader.forward(length=1) - move the current position to `length` characters. +# reader.index - the number of the current character. +# reader.line, stream.column - the line and the column of the current character. + +__all__ = ['Reader', 'ReaderError'] + +from .error import YAMLError, Mark + +import codecs, re + +class ReaderError(YAMLError): + + def __init__(self, name, position, character, encoding, reason): + self.name = name + self.character = character + self.position = position + self.encoding = encoding + self.reason = reason + + def __str__(self): + if isinstance(self.character, bytes): + return "'%s' codec can't decode byte #x%02x: %s\n" \ + " in \"%s\", position %d" \ + % (self.encoding, ord(self.character), self.reason, + self.name, self.position) + else: + return "unacceptable character #x%04x: %s\n" \ + " in \"%s\", position %d" \ + % (self.character, self.reason, + self.name, self.position) + +class Reader(object): + # Reader: + # - determines the data encoding and converts it to a unicode string, + # - checks if characters are in allowed range, + # - adds '\0' to the end. + + # Reader accepts + # - a `bytes` object, + # - a `str` object, + # - a file-like object with its `read` method returning `str`, + # - a file-like object with its `read` method returning `unicode`. + + # Yeah, it's ugly and slow. + + def __init__(self, stream): + self.name = None + self.stream = None + self.stream_pointer = 0 + self.eof = True + self.buffer = '' + self.pointer = 0 + self.raw_buffer = None + self.raw_decode = None + self.encoding = None + self.index = 0 + self.line = 0 + self.column = 0 + if isinstance(stream, str): + self.name = "" + self.check_printable(stream) + self.buffer = stream+'\0' + elif isinstance(stream, bytes): + self.name = "" + self.raw_buffer = stream + self.determine_encoding() + else: + self.stream = stream + self.name = getattr(stream, 'name', "") + self.eof = False + self.raw_buffer = None + self.determine_encoding() + + def peek(self, index=0): + try: + return self.buffer[self.pointer+index] + except IndexError: + self.update(index+1) + return self.buffer[self.pointer+index] + + def prefix(self, length=1): + if self.pointer+length >= len(self.buffer): + self.update(length) + return self.buffer[self.pointer:self.pointer+length] + + def forward(self, length=1): + if self.pointer+length+1 >= len(self.buffer): + self.update(length+1) + while length: + ch = self.buffer[self.pointer] + self.pointer += 1 + self.index += 1 + if ch in '\n\x85\u2028\u2029' \ + or (ch == '\r' and self.buffer[self.pointer] != '\n'): + self.line += 1 + self.column = 0 + elif ch != '\uFEFF': + self.column += 1 + length -= 1 + + def get_mark(self): + if self.stream is None: + return Mark(self.name, self.index, self.line, self.column, + self.buffer, self.pointer) + else: + return Mark(self.name, self.index, self.line, self.column, + None, None) + + def determine_encoding(self): + while not self.eof and (self.raw_buffer is None or len(self.raw_buffer) < 2): + self.update_raw() + if isinstance(self.raw_buffer, bytes): + if self.raw_buffer.startswith(codecs.BOM_UTF16_LE): + self.raw_decode = codecs.utf_16_le_decode + self.encoding = 'utf-16-le' + elif self.raw_buffer.startswith(codecs.BOM_UTF16_BE): + self.raw_decode = codecs.utf_16_be_decode + self.encoding = 'utf-16-be' + else: + self.raw_decode = codecs.utf_8_decode + self.encoding = 'utf-8' + self.update(1) + + NON_PRINTABLE = re.compile('[^\x09\x0A\x0D\x20-\x7E\x85\xA0-\uD7FF\uE000-\uFFFD\U00010000-\U0010ffff]') + def check_printable(self, data): + match = self.NON_PRINTABLE.search(data) + if match: + character = match.group() + position = self.index+(len(self.buffer)-self.pointer)+match.start() + raise ReaderError(self.name, position, ord(character), + 'unicode', "special characters are not allowed") + + def update(self, length): + if self.raw_buffer is None: + return + self.buffer = self.buffer[self.pointer:] + self.pointer = 0 + while len(self.buffer) < length: + if not self.eof: + self.update_raw() + if self.raw_decode is not None: + try: + data, converted = self.raw_decode(self.raw_buffer, + 'strict', self.eof) + except UnicodeDecodeError as exc: + character = self.raw_buffer[exc.start] + if self.stream is not None: + position = self.stream_pointer-len(self.raw_buffer)+exc.start + else: + position = exc.start + raise ReaderError(self.name, position, character, + exc.encoding, exc.reason) + else: + data = self.raw_buffer + converted = len(data) + self.check_printable(data) + self.buffer += data + self.raw_buffer = self.raw_buffer[converted:] + if self.eof: + self.buffer += '\0' + self.raw_buffer = None + break + + def update_raw(self, size=4096): + data = self.stream.read(size) + if self.raw_buffer is None: + self.raw_buffer = data + else: + self.raw_buffer += data + self.stream_pointer += len(data) + if not data: + self.eof = True diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/representer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/representer.py new file mode 100644 index 0000000000000000000000000000000000000000..808ca06dfbd60c9a23eb079151b74a82ef688749 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/representer.py @@ -0,0 +1,389 @@ + +__all__ = ['BaseRepresenter', 'SafeRepresenter', 'Representer', + 'RepresenterError'] + +from .error import * +from .nodes import * + +import datetime, copyreg, types, base64, collections + +class RepresenterError(YAMLError): + pass + +class BaseRepresenter: + + yaml_representers = {} + yaml_multi_representers = {} + + def __init__(self, default_style=None, default_flow_style=False, sort_keys=True): + self.default_style = default_style + self.sort_keys = sort_keys + self.default_flow_style = default_flow_style + self.represented_objects = {} + self.object_keeper = [] + self.alias_key = None + + def represent(self, data): + node = self.represent_data(data) + self.serialize(node) + self.represented_objects = {} + self.object_keeper = [] + self.alias_key = None + + def represent_data(self, data): + if self.ignore_aliases(data): + self.alias_key = None + else: + self.alias_key = id(data) + if self.alias_key is not None: + if self.alias_key in self.represented_objects: + node = self.represented_objects[self.alias_key] + #if node is None: + # raise RepresenterError("recursive objects are not allowed: %r" % data) + return node + #self.represented_objects[alias_key] = None + self.object_keeper.append(data) + data_types = type(data).__mro__ + if data_types[0] in self.yaml_representers: + node = self.yaml_representers[data_types[0]](self, data) + else: + for data_type in data_types: + if data_type in self.yaml_multi_representers: + node = self.yaml_multi_representers[data_type](self, data) + break + else: + if None in self.yaml_multi_representers: + node = self.yaml_multi_representers[None](self, data) + elif None in self.yaml_representers: + node = self.yaml_representers[None](self, data) + else: + node = ScalarNode(None, str(data)) + #if alias_key is not None: + # self.represented_objects[alias_key] = node + return node + + @classmethod + def add_representer(cls, data_type, representer): + if not 'yaml_representers' in cls.__dict__: + cls.yaml_representers = cls.yaml_representers.copy() + cls.yaml_representers[data_type] = representer + + @classmethod + def add_multi_representer(cls, data_type, representer): + if not 'yaml_multi_representers' in cls.__dict__: + cls.yaml_multi_representers = cls.yaml_multi_representers.copy() + cls.yaml_multi_representers[data_type] = representer + + def represent_scalar(self, tag, value, style=None): + if style is None: + style = self.default_style + node = ScalarNode(tag, value, style=style) + if self.alias_key is not None: + self.represented_objects[self.alias_key] = node + return node + + def represent_sequence(self, tag, sequence, flow_style=None): + value = [] + node = SequenceNode(tag, value, flow_style=flow_style) + if self.alias_key is not None: + self.represented_objects[self.alias_key] = node + best_style = True + for item in sequence: + node_item = self.represent_data(item) + if not (isinstance(node_item, ScalarNode) and not node_item.style): + best_style = False + value.append(node_item) + if flow_style is None: + if self.default_flow_style is not None: + node.flow_style = self.default_flow_style + else: + node.flow_style = best_style + return node + + def represent_mapping(self, tag, mapping, flow_style=None): + value = [] + node = MappingNode(tag, value, flow_style=flow_style) + if self.alias_key is not None: + self.represented_objects[self.alias_key] = node + best_style = True + if hasattr(mapping, 'items'): + mapping = list(mapping.items()) + if self.sort_keys: + try: + mapping = sorted(mapping) + except TypeError: + pass + for item_key, item_value in mapping: + node_key = self.represent_data(item_key) + node_value = self.represent_data(item_value) + if not (isinstance(node_key, ScalarNode) and not node_key.style): + best_style = False + if not (isinstance(node_value, ScalarNode) and not node_value.style): + best_style = False + value.append((node_key, node_value)) + if flow_style is None: + if self.default_flow_style is not None: + node.flow_style = self.default_flow_style + else: + node.flow_style = best_style + return node + + def ignore_aliases(self, data): + return False + +class SafeRepresenter(BaseRepresenter): + + def ignore_aliases(self, data): + if data is None: + return True + if isinstance(data, tuple) and data == (): + return True + if isinstance(data, (str, bytes, bool, int, float)): + return True + + def represent_none(self, data): + return self.represent_scalar('tag:yaml.org,2002:null', 'null') + + def represent_str(self, data): + return self.represent_scalar('tag:yaml.org,2002:str', data) + + def represent_binary(self, data): + if hasattr(base64, 'encodebytes'): + data = base64.encodebytes(data).decode('ascii') + else: + data = base64.encodestring(data).decode('ascii') + return self.represent_scalar('tag:yaml.org,2002:binary', data, style='|') + + def represent_bool(self, data): + if data: + value = 'true' + else: + value = 'false' + return self.represent_scalar('tag:yaml.org,2002:bool', value) + + def represent_int(self, data): + return self.represent_scalar('tag:yaml.org,2002:int', str(data)) + + inf_value = 1e300 + while repr(inf_value) != repr(inf_value*inf_value): + inf_value *= inf_value + + def represent_float(self, data): + if data != data or (data == 0.0 and data == 1.0): + value = '.nan' + elif data == self.inf_value: + value = '.inf' + elif data == -self.inf_value: + value = '-.inf' + else: + value = repr(data).lower() + # Note that in some cases `repr(data)` represents a float number + # without the decimal parts. For instance: + # >>> repr(1e17) + # '1e17' + # Unfortunately, this is not a valid float representation according + # to the definition of the `!!float` tag. We fix this by adding + # '.0' before the 'e' symbol. + if '.' not in value and 'e' in value: + value = value.replace('e', '.0e', 1) + return self.represent_scalar('tag:yaml.org,2002:float', value) + + def represent_list(self, data): + #pairs = (len(data) > 0 and isinstance(data, list)) + #if pairs: + # for item in data: + # if not isinstance(item, tuple) or len(item) != 2: + # pairs = False + # break + #if not pairs: + return self.represent_sequence('tag:yaml.org,2002:seq', data) + #value = [] + #for item_key, item_value in data: + # value.append(self.represent_mapping(u'tag:yaml.org,2002:map', + # [(item_key, item_value)])) + #return SequenceNode(u'tag:yaml.org,2002:pairs', value) + + def represent_dict(self, data): + return self.represent_mapping('tag:yaml.org,2002:map', data) + + def represent_set(self, data): + value = {} + for key in data: + value[key] = None + return self.represent_mapping('tag:yaml.org,2002:set', value) + + def represent_date(self, data): + value = data.isoformat() + return self.represent_scalar('tag:yaml.org,2002:timestamp', value) + + def represent_datetime(self, data): + value = data.isoformat(' ') + return self.represent_scalar('tag:yaml.org,2002:timestamp', value) + + def represent_yaml_object(self, tag, data, cls, flow_style=None): + if hasattr(data, '__getstate__'): + state = data.__getstate__() + else: + state = data.__dict__.copy() + return self.represent_mapping(tag, state, flow_style=flow_style) + + def represent_undefined(self, data): + raise RepresenterError("cannot represent an object", data) + +SafeRepresenter.add_representer(type(None), + SafeRepresenter.represent_none) + +SafeRepresenter.add_representer(str, + SafeRepresenter.represent_str) + +SafeRepresenter.add_representer(bytes, + SafeRepresenter.represent_binary) + +SafeRepresenter.add_representer(bool, + SafeRepresenter.represent_bool) + +SafeRepresenter.add_representer(int, + SafeRepresenter.represent_int) + +SafeRepresenter.add_representer(float, + SafeRepresenter.represent_float) + +SafeRepresenter.add_representer(list, + SafeRepresenter.represent_list) + +SafeRepresenter.add_representer(tuple, + SafeRepresenter.represent_list) + +SafeRepresenter.add_representer(dict, + SafeRepresenter.represent_dict) + +SafeRepresenter.add_representer(set, + SafeRepresenter.represent_set) + +SafeRepresenter.add_representer(datetime.date, + SafeRepresenter.represent_date) + +SafeRepresenter.add_representer(datetime.datetime, + SafeRepresenter.represent_datetime) + +SafeRepresenter.add_representer(None, + SafeRepresenter.represent_undefined) + +class Representer(SafeRepresenter): + + def represent_complex(self, data): + if data.imag == 0.0: + data = '%r' % data.real + elif data.real == 0.0: + data = '%rj' % data.imag + elif data.imag > 0: + data = '%r+%rj' % (data.real, data.imag) + else: + data = '%r%rj' % (data.real, data.imag) + return self.represent_scalar('tag:yaml.org,2002:python/complex', data) + + def represent_tuple(self, data): + return self.represent_sequence('tag:yaml.org,2002:python/tuple', data) + + def represent_name(self, data): + name = '%s.%s' % (data.__module__, data.__name__) + return self.represent_scalar('tag:yaml.org,2002:python/name:'+name, '') + + def represent_module(self, data): + return self.represent_scalar( + 'tag:yaml.org,2002:python/module:'+data.__name__, '') + + def represent_object(self, data): + # We use __reduce__ API to save the data. data.__reduce__ returns + # a tuple of length 2-5: + # (function, args, state, listitems, dictitems) + + # For reconstructing, we calls function(*args), then set its state, + # listitems, and dictitems if they are not None. + + # A special case is when function.__name__ == '__newobj__'. In this + # case we create the object with args[0].__new__(*args). + + # Another special case is when __reduce__ returns a string - we don't + # support it. + + # We produce a !!python/object, !!python/object/new or + # !!python/object/apply node. + + cls = type(data) + if cls in copyreg.dispatch_table: + reduce = copyreg.dispatch_table[cls](data) + elif hasattr(data, '__reduce_ex__'): + reduce = data.__reduce_ex__(2) + elif hasattr(data, '__reduce__'): + reduce = data.__reduce__() + else: + raise RepresenterError("cannot represent an object", data) + reduce = (list(reduce)+[None]*5)[:5] + function, args, state, listitems, dictitems = reduce + args = list(args) + if state is None: + state = {} + if listitems is not None: + listitems = list(listitems) + if dictitems is not None: + dictitems = dict(dictitems) + if function.__name__ == '__newobj__': + function = args[0] + args = args[1:] + tag = 'tag:yaml.org,2002:python/object/new:' + newobj = True + else: + tag = 'tag:yaml.org,2002:python/object/apply:' + newobj = False + function_name = '%s.%s' % (function.__module__, function.__name__) + if not args and not listitems and not dictitems \ + and isinstance(state, dict) and newobj: + return self.represent_mapping( + 'tag:yaml.org,2002:python/object:'+function_name, state) + if not listitems and not dictitems \ + and isinstance(state, dict) and not state: + return self.represent_sequence(tag+function_name, args) + value = {} + if args: + value['args'] = args + if state or not isinstance(state, dict): + value['state'] = state + if listitems: + value['listitems'] = listitems + if dictitems: + value['dictitems'] = dictitems + return self.represent_mapping(tag+function_name, value) + + def represent_ordered_dict(self, data): + # Provide uniform representation across different Python versions. + data_type = type(data) + tag = 'tag:yaml.org,2002:python/object/apply:%s.%s' \ + % (data_type.__module__, data_type.__name__) + items = [[key, value] for key, value in data.items()] + return self.represent_sequence(tag, [items]) + +Representer.add_representer(complex, + Representer.represent_complex) + +Representer.add_representer(tuple, + Representer.represent_tuple) + +Representer.add_multi_representer(type, + Representer.represent_name) + +Representer.add_representer(collections.OrderedDict, + Representer.represent_ordered_dict) + +Representer.add_representer(types.FunctionType, + Representer.represent_name) + +Representer.add_representer(types.BuiltinFunctionType, + Representer.represent_name) + +Representer.add_representer(types.ModuleType, + Representer.represent_module) + +Representer.add_multi_representer(object, + Representer.represent_object) + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/resolver.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/resolver.py new file mode 100644 index 0000000000000000000000000000000000000000..3522bdaaf6358110b608f4e6503b9d314c82d887 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/resolver.py @@ -0,0 +1,227 @@ + +__all__ = ['BaseResolver', 'Resolver'] + +from .error import * +from .nodes import * + +import re + +class ResolverError(YAMLError): + pass + +class BaseResolver: + + DEFAULT_SCALAR_TAG = 'tag:yaml.org,2002:str' + DEFAULT_SEQUENCE_TAG = 'tag:yaml.org,2002:seq' + DEFAULT_MAPPING_TAG = 'tag:yaml.org,2002:map' + + yaml_implicit_resolvers = {} + yaml_path_resolvers = {} + + def __init__(self): + self.resolver_exact_paths = [] + self.resolver_prefix_paths = [] + + @classmethod + def add_implicit_resolver(cls, tag, regexp, first): + if not 'yaml_implicit_resolvers' in cls.__dict__: + implicit_resolvers = {} + for key in cls.yaml_implicit_resolvers: + implicit_resolvers[key] = cls.yaml_implicit_resolvers[key][:] + cls.yaml_implicit_resolvers = implicit_resolvers + if first is None: + first = [None] + for ch in first: + cls.yaml_implicit_resolvers.setdefault(ch, []).append((tag, regexp)) + + @classmethod + def add_path_resolver(cls, tag, path, kind=None): + # Note: `add_path_resolver` is experimental. The API could be changed. + # `new_path` is a pattern that is matched against the path from the + # root to the node that is being considered. `node_path` elements are + # tuples `(node_check, index_check)`. `node_check` is a node class: + # `ScalarNode`, `SequenceNode`, `MappingNode` or `None`. `None` + # matches any kind of a node. `index_check` could be `None`, a boolean + # value, a string value, or a number. `None` and `False` match against + # any _value_ of sequence and mapping nodes. `True` matches against + # any _key_ of a mapping node. A string `index_check` matches against + # a mapping value that corresponds to a scalar key which content is + # equal to the `index_check` value. An integer `index_check` matches + # against a sequence value with the index equal to `index_check`. + if not 'yaml_path_resolvers' in cls.__dict__: + cls.yaml_path_resolvers = cls.yaml_path_resolvers.copy() + new_path = [] + for element in path: + if isinstance(element, (list, tuple)): + if len(element) == 2: + node_check, index_check = element + elif len(element) == 1: + node_check = element[0] + index_check = True + else: + raise ResolverError("Invalid path element: %s" % element) + else: + node_check = None + index_check = element + if node_check is str: + node_check = ScalarNode + elif node_check is list: + node_check = SequenceNode + elif node_check is dict: + node_check = MappingNode + elif node_check not in [ScalarNode, SequenceNode, MappingNode] \ + and not isinstance(node_check, str) \ + and node_check is not None: + raise ResolverError("Invalid node checker: %s" % node_check) + if not isinstance(index_check, (str, int)) \ + and index_check is not None: + raise ResolverError("Invalid index checker: %s" % index_check) + new_path.append((node_check, index_check)) + if kind is str: + kind = ScalarNode + elif kind is list: + kind = SequenceNode + elif kind is dict: + kind = MappingNode + elif kind not in [ScalarNode, SequenceNode, MappingNode] \ + and kind is not None: + raise ResolverError("Invalid node kind: %s" % kind) + cls.yaml_path_resolvers[tuple(new_path), kind] = tag + + def descend_resolver(self, current_node, current_index): + if not self.yaml_path_resolvers: + return + exact_paths = {} + prefix_paths = [] + if current_node: + depth = len(self.resolver_prefix_paths) + for path, kind in self.resolver_prefix_paths[-1]: + if self.check_resolver_prefix(depth, path, kind, + current_node, current_index): + if len(path) > depth: + prefix_paths.append((path, kind)) + else: + exact_paths[kind] = self.yaml_path_resolvers[path, kind] + else: + for path, kind in self.yaml_path_resolvers: + if not path: + exact_paths[kind] = self.yaml_path_resolvers[path, kind] + else: + prefix_paths.append((path, kind)) + self.resolver_exact_paths.append(exact_paths) + self.resolver_prefix_paths.append(prefix_paths) + + def ascend_resolver(self): + if not self.yaml_path_resolvers: + return + self.resolver_exact_paths.pop() + self.resolver_prefix_paths.pop() + + def check_resolver_prefix(self, depth, path, kind, + current_node, current_index): + node_check, index_check = path[depth-1] + if isinstance(node_check, str): + if current_node.tag != node_check: + return + elif node_check is not None: + if not isinstance(current_node, node_check): + return + if index_check is True and current_index is not None: + return + if (index_check is False or index_check is None) \ + and current_index is None: + return + if isinstance(index_check, str): + if not (isinstance(current_index, ScalarNode) + and index_check == current_index.value): + return + elif isinstance(index_check, int) and not isinstance(index_check, bool): + if index_check != current_index: + return + return True + + def resolve(self, kind, value, implicit): + if kind is ScalarNode and implicit[0]: + if value == '': + resolvers = self.yaml_implicit_resolvers.get('', []) + else: + resolvers = self.yaml_implicit_resolvers.get(value[0], []) + wildcard_resolvers = self.yaml_implicit_resolvers.get(None, []) + for tag, regexp in resolvers + wildcard_resolvers: + if regexp.match(value): + return tag + implicit = implicit[1] + if self.yaml_path_resolvers: + exact_paths = self.resolver_exact_paths[-1] + if kind in exact_paths: + return exact_paths[kind] + if None in exact_paths: + return exact_paths[None] + if kind is ScalarNode: + return self.DEFAULT_SCALAR_TAG + elif kind is SequenceNode: + return self.DEFAULT_SEQUENCE_TAG + elif kind is MappingNode: + return self.DEFAULT_MAPPING_TAG + +class Resolver(BaseResolver): + pass + +Resolver.add_implicit_resolver( + 'tag:yaml.org,2002:bool', + re.compile(r'''^(?:yes|Yes|YES|no|No|NO + |true|True|TRUE|false|False|FALSE + |on|On|ON|off|Off|OFF)$''', re.X), + list('yYnNtTfFoO')) + +Resolver.add_implicit_resolver( + 'tag:yaml.org,2002:float', + re.compile(r'''^(?:[-+]?(?:[0-9][0-9_]*)\.[0-9_]*(?:[eE][-+][0-9]+)? + |\.[0-9][0-9_]*(?:[eE][-+][0-9]+)? + |[-+]?[0-9][0-9_]*(?::[0-5]?[0-9])+\.[0-9_]* + |[-+]?\.(?:inf|Inf|INF) + |\.(?:nan|NaN|NAN))$''', re.X), + list('-+0123456789.')) + +Resolver.add_implicit_resolver( + 'tag:yaml.org,2002:int', + re.compile(r'''^(?:[-+]?0b[0-1_]+ + |[-+]?0[0-7_]+ + |[-+]?(?:0|[1-9][0-9_]*) + |[-+]?0x[0-9a-fA-F_]+ + |[-+]?[1-9][0-9_]*(?::[0-5]?[0-9])+)$''', re.X), + list('-+0123456789')) + +Resolver.add_implicit_resolver( + 'tag:yaml.org,2002:merge', + re.compile(r'^(?:<<)$'), + ['<']) + +Resolver.add_implicit_resolver( + 'tag:yaml.org,2002:null', + re.compile(r'''^(?: ~ + |null|Null|NULL + | )$''', re.X), + ['~', 'n', 'N', '']) + +Resolver.add_implicit_resolver( + 'tag:yaml.org,2002:timestamp', + re.compile(r'''^(?:[0-9][0-9][0-9][0-9]-[0-9][0-9]-[0-9][0-9] + |[0-9][0-9][0-9][0-9] -[0-9][0-9]? -[0-9][0-9]? + (?:[Tt]|[ \t]+)[0-9][0-9]? + :[0-9][0-9] :[0-9][0-9] (?:\.[0-9]*)? + (?:[ \t]*(?:Z|[-+][0-9][0-9]?(?::[0-9][0-9])?))?)$''', re.X), + list('0123456789')) + +Resolver.add_implicit_resolver( + 'tag:yaml.org,2002:value', + re.compile(r'^(?:=)$'), + ['=']) + +# The following resolver is only for documentation purposes. It cannot work +# because plain scalars cannot start with '!', '&', or '*'. +Resolver.add_implicit_resolver( + 'tag:yaml.org,2002:yaml', + re.compile(r'^(?:!|&|\*)$'), + list('!&*')) + diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/scanner.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/scanner.py new file mode 100644 index 0000000000000000000000000000000000000000..de925b07f1eaec33c9c305a8a69f9eb7ac5983c5 --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/scanner.py @@ -0,0 +1,1435 @@ + +# Scanner produces tokens of the following types: +# STREAM-START +# STREAM-END +# DIRECTIVE(name, value) +# DOCUMENT-START +# DOCUMENT-END +# BLOCK-SEQUENCE-START +# BLOCK-MAPPING-START +# BLOCK-END +# FLOW-SEQUENCE-START +# FLOW-MAPPING-START +# FLOW-SEQUENCE-END +# FLOW-MAPPING-END +# BLOCK-ENTRY +# FLOW-ENTRY +# KEY +# VALUE +# ALIAS(value) +# ANCHOR(value) +# TAG(value) +# SCALAR(value, plain, style) +# +# Read comments in the Scanner code for more details. +# + +__all__ = ['Scanner', 'ScannerError'] + +from .error import MarkedYAMLError +from .tokens import * + +class ScannerError(MarkedYAMLError): + pass + +class SimpleKey: + # See below simple keys treatment. + + def __init__(self, token_number, required, index, line, column, mark): + self.token_number = token_number + self.required = required + self.index = index + self.line = line + self.column = column + self.mark = mark + +class Scanner: + + def __init__(self): + """Initialize the scanner.""" + # It is assumed that Scanner and Reader will have a common descendant. + # Reader do the dirty work of checking for BOM and converting the + # input data to Unicode. It also adds NUL to the end. + # + # Reader supports the following methods + # self.peek(i=0) # peek the next i-th character + # self.prefix(l=1) # peek the next l characters + # self.forward(l=1) # read the next l characters and move the pointer. + + # Had we reached the end of the stream? + self.done = False + + # The number of unclosed '{' and '['. `flow_level == 0` means block + # context. + self.flow_level = 0 + + # List of processed tokens that are not yet emitted. + self.tokens = [] + + # Add the STREAM-START token. + self.fetch_stream_start() + + # Number of tokens that were emitted through the `get_token` method. + self.tokens_taken = 0 + + # The current indentation level. + self.indent = -1 + + # Past indentation levels. + self.indents = [] + + # Variables related to simple keys treatment. + + # A simple key is a key that is not denoted by the '?' indicator. + # Example of simple keys: + # --- + # block simple key: value + # ? not a simple key: + # : { flow simple key: value } + # We emit the KEY token before all keys, so when we find a potential + # simple key, we try to locate the corresponding ':' indicator. + # Simple keys should be limited to a single line and 1024 characters. + + # Can a simple key start at the current position? A simple key may + # start: + # - at the beginning of the line, not counting indentation spaces + # (in block context), + # - after '{', '[', ',' (in the flow context), + # - after '?', ':', '-' (in the block context). + # In the block context, this flag also signifies if a block collection + # may start at the current position. + self.allow_simple_key = True + + # Keep track of possible simple keys. This is a dictionary. The key + # is `flow_level`; there can be no more that one possible simple key + # for each level. The value is a SimpleKey record: + # (token_number, required, index, line, column, mark) + # A simple key may start with ALIAS, ANCHOR, TAG, SCALAR(flow), + # '[', or '{' tokens. + self.possible_simple_keys = {} + + # Public methods. + + def check_token(self, *choices): + # Check if the next token is one of the given types. + while self.need_more_tokens(): + self.fetch_more_tokens() + if self.tokens: + if not choices: + return True + for choice in choices: + if isinstance(self.tokens[0], choice): + return True + return False + + def peek_token(self): + # Return the next token, but do not delete if from the queue. + # Return None if no more tokens. + while self.need_more_tokens(): + self.fetch_more_tokens() + if self.tokens: + return self.tokens[0] + else: + return None + + def get_token(self): + # Return the next token. + while self.need_more_tokens(): + self.fetch_more_tokens() + if self.tokens: + self.tokens_taken += 1 + return self.tokens.pop(0) + + # Private methods. + + def need_more_tokens(self): + if self.done: + return False + if not self.tokens: + return True + # The current token may be a potential simple key, so we + # need to look further. + self.stale_possible_simple_keys() + if self.next_possible_simple_key() == self.tokens_taken: + return True + + def fetch_more_tokens(self): + + # Eat whitespaces and comments until we reach the next token. + self.scan_to_next_token() + + # Remove obsolete possible simple keys. + self.stale_possible_simple_keys() + + # Compare the current indentation and column. It may add some tokens + # and decrease the current indentation level. + self.unwind_indent(self.column) + + # Peek the next character. + ch = self.peek() + + # Is it the end of stream? + if ch == '\0': + return self.fetch_stream_end() + + # Is it a directive? + if ch == '%' and self.check_directive(): + return self.fetch_directive() + + # Is it the document start? + if ch == '-' and self.check_document_start(): + return self.fetch_document_start() + + # Is it the document end? + if ch == '.' and self.check_document_end(): + return self.fetch_document_end() + + # TODO: support for BOM within a stream. + #if ch == '\uFEFF': + # return self.fetch_bom() <-- issue BOMToken + + # Note: the order of the following checks is NOT significant. + + # Is it the flow sequence start indicator? + if ch == '[': + return self.fetch_flow_sequence_start() + + # Is it the flow mapping start indicator? + if ch == '{': + return self.fetch_flow_mapping_start() + + # Is it the flow sequence end indicator? + if ch == ']': + return self.fetch_flow_sequence_end() + + # Is it the flow mapping end indicator? + if ch == '}': + return self.fetch_flow_mapping_end() + + # Is it the flow entry indicator? + if ch == ',': + return self.fetch_flow_entry() + + # Is it the block entry indicator? + if ch == '-' and self.check_block_entry(): + return self.fetch_block_entry() + + # Is it the key indicator? + if ch == '?' and self.check_key(): + return self.fetch_key() + + # Is it the value indicator? + if ch == ':' and self.check_value(): + return self.fetch_value() + + # Is it an alias? + if ch == '*': + return self.fetch_alias() + + # Is it an anchor? + if ch == '&': + return self.fetch_anchor() + + # Is it a tag? + if ch == '!': + return self.fetch_tag() + + # Is it a literal scalar? + if ch == '|' and not self.flow_level: + return self.fetch_literal() + + # Is it a folded scalar? + if ch == '>' and not self.flow_level: + return self.fetch_folded() + + # Is it a single quoted scalar? + if ch == '\'': + return self.fetch_single() + + # Is it a double quoted scalar? + if ch == '\"': + return self.fetch_double() + + # It must be a plain scalar then. + if self.check_plain(): + return self.fetch_plain() + + # No? It's an error. Let's produce a nice error message. + raise ScannerError("while scanning for the next token", None, + "found character %r that cannot start any token" % ch, + self.get_mark()) + + # Simple keys treatment. + + def next_possible_simple_key(self): + # Return the number of the nearest possible simple key. Actually we + # don't need to loop through the whole dictionary. We may replace it + # with the following code: + # if not self.possible_simple_keys: + # return None + # return self.possible_simple_keys[ + # min(self.possible_simple_keys.keys())].token_number + min_token_number = None + for level in self.possible_simple_keys: + key = self.possible_simple_keys[level] + if min_token_number is None or key.token_number < min_token_number: + min_token_number = key.token_number + return min_token_number + + def stale_possible_simple_keys(self): + # Remove entries that are no longer possible simple keys. According to + # the YAML specification, simple keys + # - should be limited to a single line, + # - should be no longer than 1024 characters. + # Disabling this procedure will allow simple keys of any length and + # height (may cause problems if indentation is broken though). + for level in list(self.possible_simple_keys): + key = self.possible_simple_keys[level] + if key.line != self.line \ + or self.index-key.index > 1024: + if key.required: + raise ScannerError("while scanning a simple key", key.mark, + "could not find expected ':'", self.get_mark()) + del self.possible_simple_keys[level] + + def save_possible_simple_key(self): + # The next token may start a simple key. We check if it's possible + # and save its position. This function is called for + # ALIAS, ANCHOR, TAG, SCALAR(flow), '[', and '{'. + + # Check if a simple key is required at the current position. + required = not self.flow_level and self.indent == self.column + + # The next token might be a simple key. Let's save it's number and + # position. + if self.allow_simple_key: + self.remove_possible_simple_key() + token_number = self.tokens_taken+len(self.tokens) + key = SimpleKey(token_number, required, + self.index, self.line, self.column, self.get_mark()) + self.possible_simple_keys[self.flow_level] = key + + def remove_possible_simple_key(self): + # Remove the saved possible key position at the current flow level. + if self.flow_level in self.possible_simple_keys: + key = self.possible_simple_keys[self.flow_level] + + if key.required: + raise ScannerError("while scanning a simple key", key.mark, + "could not find expected ':'", self.get_mark()) + + del self.possible_simple_keys[self.flow_level] + + # Indentation functions. + + def unwind_indent(self, column): + + ## In flow context, tokens should respect indentation. + ## Actually the condition should be `self.indent >= column` according to + ## the spec. But this condition will prohibit intuitively correct + ## constructions such as + ## key : { + ## } + #if self.flow_level and self.indent > column: + # raise ScannerError(None, None, + # "invalid indentation or unclosed '[' or '{'", + # self.get_mark()) + + # In the flow context, indentation is ignored. We make the scanner less + # restrictive then specification requires. + if self.flow_level: + return + + # In block context, we may need to issue the BLOCK-END tokens. + while self.indent > column: + mark = self.get_mark() + self.indent = self.indents.pop() + self.tokens.append(BlockEndToken(mark, mark)) + + def add_indent(self, column): + # Check if we need to increase indentation. + if self.indent < column: + self.indents.append(self.indent) + self.indent = column + return True + return False + + # Fetchers. + + def fetch_stream_start(self): + # We always add STREAM-START as the first token and STREAM-END as the + # last token. + + # Read the token. + mark = self.get_mark() + + # Add STREAM-START. + self.tokens.append(StreamStartToken(mark, mark, + encoding=self.encoding)) + + + def fetch_stream_end(self): + + # Set the current indentation to -1. + self.unwind_indent(-1) + + # Reset simple keys. + self.remove_possible_simple_key() + self.allow_simple_key = False + self.possible_simple_keys = {} + + # Read the token. + mark = self.get_mark() + + # Add STREAM-END. + self.tokens.append(StreamEndToken(mark, mark)) + + # The steam is finished. + self.done = True + + def fetch_directive(self): + + # Set the current indentation to -1. + self.unwind_indent(-1) + + # Reset simple keys. + self.remove_possible_simple_key() + self.allow_simple_key = False + + # Scan and add DIRECTIVE. + self.tokens.append(self.scan_directive()) + + def fetch_document_start(self): + self.fetch_document_indicator(DocumentStartToken) + + def fetch_document_end(self): + self.fetch_document_indicator(DocumentEndToken) + + def fetch_document_indicator(self, TokenClass): + + # Set the current indentation to -1. + self.unwind_indent(-1) + + # Reset simple keys. Note that there could not be a block collection + # after '---'. + self.remove_possible_simple_key() + self.allow_simple_key = False + + # Add DOCUMENT-START or DOCUMENT-END. + start_mark = self.get_mark() + self.forward(3) + end_mark = self.get_mark() + self.tokens.append(TokenClass(start_mark, end_mark)) + + def fetch_flow_sequence_start(self): + self.fetch_flow_collection_start(FlowSequenceStartToken) + + def fetch_flow_mapping_start(self): + self.fetch_flow_collection_start(FlowMappingStartToken) + + def fetch_flow_collection_start(self, TokenClass): + + # '[' and '{' may start a simple key. + self.save_possible_simple_key() + + # Increase the flow level. + self.flow_level += 1 + + # Simple keys are allowed after '[' and '{'. + self.allow_simple_key = True + + # Add FLOW-SEQUENCE-START or FLOW-MAPPING-START. + start_mark = self.get_mark() + self.forward() + end_mark = self.get_mark() + self.tokens.append(TokenClass(start_mark, end_mark)) + + def fetch_flow_sequence_end(self): + self.fetch_flow_collection_end(FlowSequenceEndToken) + + def fetch_flow_mapping_end(self): + self.fetch_flow_collection_end(FlowMappingEndToken) + + def fetch_flow_collection_end(self, TokenClass): + + # Reset possible simple key on the current level. + self.remove_possible_simple_key() + + # Decrease the flow level. + self.flow_level -= 1 + + # No simple keys after ']' or '}'. + self.allow_simple_key = False + + # Add FLOW-SEQUENCE-END or FLOW-MAPPING-END. + start_mark = self.get_mark() + self.forward() + end_mark = self.get_mark() + self.tokens.append(TokenClass(start_mark, end_mark)) + + def fetch_flow_entry(self): + + # Simple keys are allowed after ','. + self.allow_simple_key = True + + # Reset possible simple key on the current level. + self.remove_possible_simple_key() + + # Add FLOW-ENTRY. + start_mark = self.get_mark() + self.forward() + end_mark = self.get_mark() + self.tokens.append(FlowEntryToken(start_mark, end_mark)) + + def fetch_block_entry(self): + + # Block context needs additional checks. + if not self.flow_level: + + # Are we allowed to start a new entry? + if not self.allow_simple_key: + raise ScannerError(None, None, + "sequence entries are not allowed here", + self.get_mark()) + + # We may need to add BLOCK-SEQUENCE-START. + if self.add_indent(self.column): + mark = self.get_mark() + self.tokens.append(BlockSequenceStartToken(mark, mark)) + + # It's an error for the block entry to occur in the flow context, + # but we let the parser detect this. + else: + pass + + # Simple keys are allowed after '-'. + self.allow_simple_key = True + + # Reset possible simple key on the current level. + self.remove_possible_simple_key() + + # Add BLOCK-ENTRY. + start_mark = self.get_mark() + self.forward() + end_mark = self.get_mark() + self.tokens.append(BlockEntryToken(start_mark, end_mark)) + + def fetch_key(self): + + # Block context needs additional checks. + if not self.flow_level: + + # Are we allowed to start a key (not necessary a simple)? + if not self.allow_simple_key: + raise ScannerError(None, None, + "mapping keys are not allowed here", + self.get_mark()) + + # We may need to add BLOCK-MAPPING-START. + if self.add_indent(self.column): + mark = self.get_mark() + self.tokens.append(BlockMappingStartToken(mark, mark)) + + # Simple keys are allowed after '?' in the block context. + self.allow_simple_key = not self.flow_level + + # Reset possible simple key on the current level. + self.remove_possible_simple_key() + + # Add KEY. + start_mark = self.get_mark() + self.forward() + end_mark = self.get_mark() + self.tokens.append(KeyToken(start_mark, end_mark)) + + def fetch_value(self): + + # Do we determine a simple key? + if self.flow_level in self.possible_simple_keys: + + # Add KEY. + key = self.possible_simple_keys[self.flow_level] + del self.possible_simple_keys[self.flow_level] + self.tokens.insert(key.token_number-self.tokens_taken, + KeyToken(key.mark, key.mark)) + + # If this key starts a new block mapping, we need to add + # BLOCK-MAPPING-START. + if not self.flow_level: + if self.add_indent(key.column): + self.tokens.insert(key.token_number-self.tokens_taken, + BlockMappingStartToken(key.mark, key.mark)) + + # There cannot be two simple keys one after another. + self.allow_simple_key = False + + # It must be a part of a complex key. + else: + + # Block context needs additional checks. + # (Do we really need them? They will be caught by the parser + # anyway.) + if not self.flow_level: + + # We are allowed to start a complex value if and only if + # we can start a simple key. + if not self.allow_simple_key: + raise ScannerError(None, None, + "mapping values are not allowed here", + self.get_mark()) + + # If this value starts a new block mapping, we need to add + # BLOCK-MAPPING-START. It will be detected as an error later by + # the parser. + if not self.flow_level: + if self.add_indent(self.column): + mark = self.get_mark() + self.tokens.append(BlockMappingStartToken(mark, mark)) + + # Simple keys are allowed after ':' in the block context. + self.allow_simple_key = not self.flow_level + + # Reset possible simple key on the current level. + self.remove_possible_simple_key() + + # Add VALUE. + start_mark = self.get_mark() + self.forward() + end_mark = self.get_mark() + self.tokens.append(ValueToken(start_mark, end_mark)) + + def fetch_alias(self): + + # ALIAS could be a simple key. + self.save_possible_simple_key() + + # No simple keys after ALIAS. + self.allow_simple_key = False + + # Scan and add ALIAS. + self.tokens.append(self.scan_anchor(AliasToken)) + + def fetch_anchor(self): + + # ANCHOR could start a simple key. + self.save_possible_simple_key() + + # No simple keys after ANCHOR. + self.allow_simple_key = False + + # Scan and add ANCHOR. + self.tokens.append(self.scan_anchor(AnchorToken)) + + def fetch_tag(self): + + # TAG could start a simple key. + self.save_possible_simple_key() + + # No simple keys after TAG. + self.allow_simple_key = False + + # Scan and add TAG. + self.tokens.append(self.scan_tag()) + + def fetch_literal(self): + self.fetch_block_scalar(style='|') + + def fetch_folded(self): + self.fetch_block_scalar(style='>') + + def fetch_block_scalar(self, style): + + # A simple key may follow a block scalar. + self.allow_simple_key = True + + # Reset possible simple key on the current level. + self.remove_possible_simple_key() + + # Scan and add SCALAR. + self.tokens.append(self.scan_block_scalar(style)) + + def fetch_single(self): + self.fetch_flow_scalar(style='\'') + + def fetch_double(self): + self.fetch_flow_scalar(style='"') + + def fetch_flow_scalar(self, style): + + # A flow scalar could be a simple key. + self.save_possible_simple_key() + + # No simple keys after flow scalars. + self.allow_simple_key = False + + # Scan and add SCALAR. + self.tokens.append(self.scan_flow_scalar(style)) + + def fetch_plain(self): + + # A plain scalar could be a simple key. + self.save_possible_simple_key() + + # No simple keys after plain scalars. But note that `scan_plain` will + # change this flag if the scan is finished at the beginning of the + # line. + self.allow_simple_key = False + + # Scan and add SCALAR. May change `allow_simple_key`. + self.tokens.append(self.scan_plain()) + + # Checkers. + + def check_directive(self): + + # DIRECTIVE: ^ '%' ... + # The '%' indicator is already checked. + if self.column == 0: + return True + + def check_document_start(self): + + # DOCUMENT-START: ^ '---' (' '|'\n') + if self.column == 0: + if self.prefix(3) == '---' \ + and self.peek(3) in '\0 \t\r\n\x85\u2028\u2029': + return True + + def check_document_end(self): + + # DOCUMENT-END: ^ '...' (' '|'\n') + if self.column == 0: + if self.prefix(3) == '...' \ + and self.peek(3) in '\0 \t\r\n\x85\u2028\u2029': + return True + + def check_block_entry(self): + + # BLOCK-ENTRY: '-' (' '|'\n') + return self.peek(1) in '\0 \t\r\n\x85\u2028\u2029' + + def check_key(self): + + # KEY(flow context): '?' + if self.flow_level: + return True + + # KEY(block context): '?' (' '|'\n') + else: + return self.peek(1) in '\0 \t\r\n\x85\u2028\u2029' + + def check_value(self): + + # VALUE(flow context): ':' + if self.flow_level: + return True + + # VALUE(block context): ':' (' '|'\n') + else: + return self.peek(1) in '\0 \t\r\n\x85\u2028\u2029' + + def check_plain(self): + + # A plain scalar may start with any non-space character except: + # '-', '?', ':', ',', '[', ']', '{', '}', + # '#', '&', '*', '!', '|', '>', '\'', '\"', + # '%', '@', '`'. + # + # It may also start with + # '-', '?', ':' + # if it is followed by a non-space character. + # + # Note that we limit the last rule to the block context (except the + # '-' character) because we want the flow context to be space + # independent. + ch = self.peek() + return ch not in '\0 \t\r\n\x85\u2028\u2029-?:,[]{}#&*!|>\'\"%@`' \ + or (self.peek(1) not in '\0 \t\r\n\x85\u2028\u2029' + and (ch == '-' or (not self.flow_level and ch in '?:'))) + + # Scanners. + + def scan_to_next_token(self): + # We ignore spaces, line breaks and comments. + # If we find a line break in the block context, we set the flag + # `allow_simple_key` on. + # The byte order mark is stripped if it's the first character in the + # stream. We do not yet support BOM inside the stream as the + # specification requires. Any such mark will be considered as a part + # of the document. + # + # TODO: We need to make tab handling rules more sane. A good rule is + # Tabs cannot precede tokens + # BLOCK-SEQUENCE-START, BLOCK-MAPPING-START, BLOCK-END, + # KEY(block), VALUE(block), BLOCK-ENTRY + # So the checking code is + # if : + # self.allow_simple_keys = False + # We also need to add the check for `allow_simple_keys == True` to + # `unwind_indent` before issuing BLOCK-END. + # Scanners for block, flow, and plain scalars need to be modified. + + if self.index == 0 and self.peek() == '\uFEFF': + self.forward() + found = False + while not found: + while self.peek() == ' ': + self.forward() + if self.peek() == '#': + while self.peek() not in '\0\r\n\x85\u2028\u2029': + self.forward() + if self.scan_line_break(): + if not self.flow_level: + self.allow_simple_key = True + else: + found = True + + def scan_directive(self): + # See the specification for details. + start_mark = self.get_mark() + self.forward() + name = self.scan_directive_name(start_mark) + value = None + if name == 'YAML': + value = self.scan_yaml_directive_value(start_mark) + end_mark = self.get_mark() + elif name == 'TAG': + value = self.scan_tag_directive_value(start_mark) + end_mark = self.get_mark() + else: + end_mark = self.get_mark() + while self.peek() not in '\0\r\n\x85\u2028\u2029': + self.forward() + self.scan_directive_ignored_line(start_mark) + return DirectiveToken(name, value, start_mark, end_mark) + + def scan_directive_name(self, start_mark): + # See the specification for details. + length = 0 + ch = self.peek(length) + while '0' <= ch <= '9' or 'A' <= ch <= 'Z' or 'a' <= ch <= 'z' \ + or ch in '-_': + length += 1 + ch = self.peek(length) + if not length: + raise ScannerError("while scanning a directive", start_mark, + "expected alphabetic or numeric character, but found %r" + % ch, self.get_mark()) + value = self.prefix(length) + self.forward(length) + ch = self.peek() + if ch not in '\0 \r\n\x85\u2028\u2029': + raise ScannerError("while scanning a directive", start_mark, + "expected alphabetic or numeric character, but found %r" + % ch, self.get_mark()) + return value + + def scan_yaml_directive_value(self, start_mark): + # See the specification for details. + while self.peek() == ' ': + self.forward() + major = self.scan_yaml_directive_number(start_mark) + if self.peek() != '.': + raise ScannerError("while scanning a directive", start_mark, + "expected a digit or '.', but found %r" % self.peek(), + self.get_mark()) + self.forward() + minor = self.scan_yaml_directive_number(start_mark) + if self.peek() not in '\0 \r\n\x85\u2028\u2029': + raise ScannerError("while scanning a directive", start_mark, + "expected a digit or ' ', but found %r" % self.peek(), + self.get_mark()) + return (major, minor) + + def scan_yaml_directive_number(self, start_mark): + # See the specification for details. + ch = self.peek() + if not ('0' <= ch <= '9'): + raise ScannerError("while scanning a directive", start_mark, + "expected a digit, but found %r" % ch, self.get_mark()) + length = 0 + while '0' <= self.peek(length) <= '9': + length += 1 + value = int(self.prefix(length)) + self.forward(length) + return value + + def scan_tag_directive_value(self, start_mark): + # See the specification for details. + while self.peek() == ' ': + self.forward() + handle = self.scan_tag_directive_handle(start_mark) + while self.peek() == ' ': + self.forward() + prefix = self.scan_tag_directive_prefix(start_mark) + return (handle, prefix) + + def scan_tag_directive_handle(self, start_mark): + # See the specification for details. + value = self.scan_tag_handle('directive', start_mark) + ch = self.peek() + if ch != ' ': + raise ScannerError("while scanning a directive", start_mark, + "expected ' ', but found %r" % ch, self.get_mark()) + return value + + def scan_tag_directive_prefix(self, start_mark): + # See the specification for details. + value = self.scan_tag_uri('directive', start_mark) + ch = self.peek() + if ch not in '\0 \r\n\x85\u2028\u2029': + raise ScannerError("while scanning a directive", start_mark, + "expected ' ', but found %r" % ch, self.get_mark()) + return value + + def scan_directive_ignored_line(self, start_mark): + # See the specification for details. + while self.peek() == ' ': + self.forward() + if self.peek() == '#': + while self.peek() not in '\0\r\n\x85\u2028\u2029': + self.forward() + ch = self.peek() + if ch not in '\0\r\n\x85\u2028\u2029': + raise ScannerError("while scanning a directive", start_mark, + "expected a comment or a line break, but found %r" + % ch, self.get_mark()) + self.scan_line_break() + + def scan_anchor(self, TokenClass): + # The specification does not restrict characters for anchors and + # aliases. This may lead to problems, for instance, the document: + # [ *alias, value ] + # can be interpreted in two ways, as + # [ "value" ] + # and + # [ *alias , "value" ] + # Therefore we restrict aliases to numbers and ASCII letters. + start_mark = self.get_mark() + indicator = self.peek() + if indicator == '*': + name = 'alias' + else: + name = 'anchor' + self.forward() + length = 0 + ch = self.peek(length) + while '0' <= ch <= '9' or 'A' <= ch <= 'Z' or 'a' <= ch <= 'z' \ + or ch in '-_': + length += 1 + ch = self.peek(length) + if not length: + raise ScannerError("while scanning an %s" % name, start_mark, + "expected alphabetic or numeric character, but found %r" + % ch, self.get_mark()) + value = self.prefix(length) + self.forward(length) + ch = self.peek() + if ch not in '\0 \t\r\n\x85\u2028\u2029?:,]}%@`': + raise ScannerError("while scanning an %s" % name, start_mark, + "expected alphabetic or numeric character, but found %r" + % ch, self.get_mark()) + end_mark = self.get_mark() + return TokenClass(value, start_mark, end_mark) + + def scan_tag(self): + # See the specification for details. + start_mark = self.get_mark() + ch = self.peek(1) + if ch == '<': + handle = None + self.forward(2) + suffix = self.scan_tag_uri('tag', start_mark) + if self.peek() != '>': + raise ScannerError("while parsing a tag", start_mark, + "expected '>', but found %r" % self.peek(), + self.get_mark()) + self.forward() + elif ch in '\0 \t\r\n\x85\u2028\u2029': + handle = None + suffix = '!' + self.forward() + else: + length = 1 + use_handle = False + while ch not in '\0 \r\n\x85\u2028\u2029': + if ch == '!': + use_handle = True + break + length += 1 + ch = self.peek(length) + handle = '!' + if use_handle: + handle = self.scan_tag_handle('tag', start_mark) + else: + handle = '!' + self.forward() + suffix = self.scan_tag_uri('tag', start_mark) + ch = self.peek() + if ch not in '\0 \r\n\x85\u2028\u2029': + raise ScannerError("while scanning a tag", start_mark, + "expected ' ', but found %r" % ch, self.get_mark()) + value = (handle, suffix) + end_mark = self.get_mark() + return TagToken(value, start_mark, end_mark) + + def scan_block_scalar(self, style): + # See the specification for details. + + if style == '>': + folded = True + else: + folded = False + + chunks = [] + start_mark = self.get_mark() + + # Scan the header. + self.forward() + chomping, increment = self.scan_block_scalar_indicators(start_mark) + self.scan_block_scalar_ignored_line(start_mark) + + # Determine the indentation level and go to the first non-empty line. + min_indent = self.indent+1 + if min_indent < 1: + min_indent = 1 + if increment is None: + breaks, max_indent, end_mark = self.scan_block_scalar_indentation() + indent = max(min_indent, max_indent) + else: + indent = min_indent+increment-1 + breaks, end_mark = self.scan_block_scalar_breaks(indent) + line_break = '' + + # Scan the inner part of the block scalar. + while self.column == indent and self.peek() != '\0': + chunks.extend(breaks) + leading_non_space = self.peek() not in ' \t' + length = 0 + while self.peek(length) not in '\0\r\n\x85\u2028\u2029': + length += 1 + chunks.append(self.prefix(length)) + self.forward(length) + line_break = self.scan_line_break() + breaks, end_mark = self.scan_block_scalar_breaks(indent) + if self.column == indent and self.peek() != '\0': + + # Unfortunately, folding rules are ambiguous. + # + # This is the folding according to the specification: + + if folded and line_break == '\n' \ + and leading_non_space and self.peek() not in ' \t': + if not breaks: + chunks.append(' ') + else: + chunks.append(line_break) + + # This is Clark Evans's interpretation (also in the spec + # examples): + # + #if folded and line_break == '\n': + # if not breaks: + # if self.peek() not in ' \t': + # chunks.append(' ') + # else: + # chunks.append(line_break) + #else: + # chunks.append(line_break) + else: + break + + # Chomp the tail. + if chomping is not False: + chunks.append(line_break) + if chomping is True: + chunks.extend(breaks) + + # We are done. + return ScalarToken(''.join(chunks), False, start_mark, end_mark, + style) + + def scan_block_scalar_indicators(self, start_mark): + # See the specification for details. + chomping = None + increment = None + ch = self.peek() + if ch in '+-': + if ch == '+': + chomping = True + else: + chomping = False + self.forward() + ch = self.peek() + if ch in '0123456789': + increment = int(ch) + if increment == 0: + raise ScannerError("while scanning a block scalar", start_mark, + "expected indentation indicator in the range 1-9, but found 0", + self.get_mark()) + self.forward() + elif ch in '0123456789': + increment = int(ch) + if increment == 0: + raise ScannerError("while scanning a block scalar", start_mark, + "expected indentation indicator in the range 1-9, but found 0", + self.get_mark()) + self.forward() + ch = self.peek() + if ch in '+-': + if ch == '+': + chomping = True + else: + chomping = False + self.forward() + ch = self.peek() + if ch not in '\0 \r\n\x85\u2028\u2029': + raise ScannerError("while scanning a block scalar", start_mark, + "expected chomping or indentation indicators, but found %r" + % ch, self.get_mark()) + return chomping, increment + + def scan_block_scalar_ignored_line(self, start_mark): + # See the specification for details. + while self.peek() == ' ': + self.forward() + if self.peek() == '#': + while self.peek() not in '\0\r\n\x85\u2028\u2029': + self.forward() + ch = self.peek() + if ch not in '\0\r\n\x85\u2028\u2029': + raise ScannerError("while scanning a block scalar", start_mark, + "expected a comment or a line break, but found %r" % ch, + self.get_mark()) + self.scan_line_break() + + def scan_block_scalar_indentation(self): + # See the specification for details. + chunks = [] + max_indent = 0 + end_mark = self.get_mark() + while self.peek() in ' \r\n\x85\u2028\u2029': + if self.peek() != ' ': + chunks.append(self.scan_line_break()) + end_mark = self.get_mark() + else: + self.forward() + if self.column > max_indent: + max_indent = self.column + return chunks, max_indent, end_mark + + def scan_block_scalar_breaks(self, indent): + # See the specification for details. + chunks = [] + end_mark = self.get_mark() + while self.column < indent and self.peek() == ' ': + self.forward() + while self.peek() in '\r\n\x85\u2028\u2029': + chunks.append(self.scan_line_break()) + end_mark = self.get_mark() + while self.column < indent and self.peek() == ' ': + self.forward() + return chunks, end_mark + + def scan_flow_scalar(self, style): + # See the specification for details. + # Note that we loose indentation rules for quoted scalars. Quoted + # scalars don't need to adhere indentation because " and ' clearly + # mark the beginning and the end of them. Therefore we are less + # restrictive then the specification requires. We only need to check + # that document separators are not included in scalars. + if style == '"': + double = True + else: + double = False + chunks = [] + start_mark = self.get_mark() + quote = self.peek() + self.forward() + chunks.extend(self.scan_flow_scalar_non_spaces(double, start_mark)) + while self.peek() != quote: + chunks.extend(self.scan_flow_scalar_spaces(double, start_mark)) + chunks.extend(self.scan_flow_scalar_non_spaces(double, start_mark)) + self.forward() + end_mark = self.get_mark() + return ScalarToken(''.join(chunks), False, start_mark, end_mark, + style) + + ESCAPE_REPLACEMENTS = { + '0': '\0', + 'a': '\x07', + 'b': '\x08', + 't': '\x09', + '\t': '\x09', + 'n': '\x0A', + 'v': '\x0B', + 'f': '\x0C', + 'r': '\x0D', + 'e': '\x1B', + ' ': '\x20', + '\"': '\"', + '\\': '\\', + '/': '/', + 'N': '\x85', + '_': '\xA0', + 'L': '\u2028', + 'P': '\u2029', + } + + ESCAPE_CODES = { + 'x': 2, + 'u': 4, + 'U': 8, + } + + def scan_flow_scalar_non_spaces(self, double, start_mark): + # See the specification for details. + chunks = [] + while True: + length = 0 + while self.peek(length) not in '\'\"\\\0 \t\r\n\x85\u2028\u2029': + length += 1 + if length: + chunks.append(self.prefix(length)) + self.forward(length) + ch = self.peek() + if not double and ch == '\'' and self.peek(1) == '\'': + chunks.append('\'') + self.forward(2) + elif (double and ch == '\'') or (not double and ch in '\"\\'): + chunks.append(ch) + self.forward() + elif double and ch == '\\': + self.forward() + ch = self.peek() + if ch in self.ESCAPE_REPLACEMENTS: + chunks.append(self.ESCAPE_REPLACEMENTS[ch]) + self.forward() + elif ch in self.ESCAPE_CODES: + length = self.ESCAPE_CODES[ch] + self.forward() + for k in range(length): + if self.peek(k) not in '0123456789ABCDEFabcdef': + raise ScannerError("while scanning a double-quoted scalar", start_mark, + "expected escape sequence of %d hexadecimal numbers, but found %r" % + (length, self.peek(k)), self.get_mark()) + code = int(self.prefix(length), 16) + chunks.append(chr(code)) + self.forward(length) + elif ch in '\r\n\x85\u2028\u2029': + self.scan_line_break() + chunks.extend(self.scan_flow_scalar_breaks(double, start_mark)) + else: + raise ScannerError("while scanning a double-quoted scalar", start_mark, + "found unknown escape character %r" % ch, self.get_mark()) + else: + return chunks + + def scan_flow_scalar_spaces(self, double, start_mark): + # See the specification for details. + chunks = [] + length = 0 + while self.peek(length) in ' \t': + length += 1 + whitespaces = self.prefix(length) + self.forward(length) + ch = self.peek() + if ch == '\0': + raise ScannerError("while scanning a quoted scalar", start_mark, + "found unexpected end of stream", self.get_mark()) + elif ch in '\r\n\x85\u2028\u2029': + line_break = self.scan_line_break() + breaks = self.scan_flow_scalar_breaks(double, start_mark) + if line_break != '\n': + chunks.append(line_break) + elif not breaks: + chunks.append(' ') + chunks.extend(breaks) + else: + chunks.append(whitespaces) + return chunks + + def scan_flow_scalar_breaks(self, double, start_mark): + # See the specification for details. + chunks = [] + while True: + # Instead of checking indentation, we check for document + # separators. + prefix = self.prefix(3) + if (prefix == '---' or prefix == '...') \ + and self.peek(3) in '\0 \t\r\n\x85\u2028\u2029': + raise ScannerError("while scanning a quoted scalar", start_mark, + "found unexpected document separator", self.get_mark()) + while self.peek() in ' \t': + self.forward() + if self.peek() in '\r\n\x85\u2028\u2029': + chunks.append(self.scan_line_break()) + else: + return chunks + + def scan_plain(self): + # See the specification for details. + # We add an additional restriction for the flow context: + # plain scalars in the flow context cannot contain ',' or '?'. + # We also keep track of the `allow_simple_key` flag here. + # Indentation rules are loosed for the flow context. + chunks = [] + start_mark = self.get_mark() + end_mark = start_mark + indent = self.indent+1 + # We allow zero indentation for scalars, but then we need to check for + # document separators at the beginning of the line. + #if indent == 0: + # indent = 1 + spaces = [] + while True: + length = 0 + if self.peek() == '#': + break + while True: + ch = self.peek(length) + if ch in '\0 \t\r\n\x85\u2028\u2029' \ + or (ch == ':' and + self.peek(length+1) in '\0 \t\r\n\x85\u2028\u2029' + + (u',[]{}' if self.flow_level else u''))\ + or (self.flow_level and ch in ',?[]{}'): + break + length += 1 + if length == 0: + break + self.allow_simple_key = False + chunks.extend(spaces) + chunks.append(self.prefix(length)) + self.forward(length) + end_mark = self.get_mark() + spaces = self.scan_plain_spaces(indent, start_mark) + if not spaces or self.peek() == '#' \ + or (not self.flow_level and self.column < indent): + break + return ScalarToken(''.join(chunks), True, start_mark, end_mark) + + def scan_plain_spaces(self, indent, start_mark): + # See the specification for details. + # The specification is really confusing about tabs in plain scalars. + # We just forbid them completely. Do not use tabs in YAML! + chunks = [] + length = 0 + while self.peek(length) in ' ': + length += 1 + whitespaces = self.prefix(length) + self.forward(length) + ch = self.peek() + if ch in '\r\n\x85\u2028\u2029': + line_break = self.scan_line_break() + self.allow_simple_key = True + prefix = self.prefix(3) + if (prefix == '---' or prefix == '...') \ + and self.peek(3) in '\0 \t\r\n\x85\u2028\u2029': + return + breaks = [] + while self.peek() in ' \r\n\x85\u2028\u2029': + if self.peek() == ' ': + self.forward() + else: + breaks.append(self.scan_line_break()) + prefix = self.prefix(3) + if (prefix == '---' or prefix == '...') \ + and self.peek(3) in '\0 \t\r\n\x85\u2028\u2029': + return + if line_break != '\n': + chunks.append(line_break) + elif not breaks: + chunks.append(' ') + chunks.extend(breaks) + elif whitespaces: + chunks.append(whitespaces) + return chunks + + def scan_tag_handle(self, name, start_mark): + # See the specification for details. + # For some strange reasons, the specification does not allow '_' in + # tag handles. I have allowed it anyway. + ch = self.peek() + if ch != '!': + raise ScannerError("while scanning a %s" % name, start_mark, + "expected '!', but found %r" % ch, self.get_mark()) + length = 1 + ch = self.peek(length) + if ch != ' ': + while '0' <= ch <= '9' or 'A' <= ch <= 'Z' or 'a' <= ch <= 'z' \ + or ch in '-_': + length += 1 + ch = self.peek(length) + if ch != '!': + self.forward(length) + raise ScannerError("while scanning a %s" % name, start_mark, + "expected '!', but found %r" % ch, self.get_mark()) + length += 1 + value = self.prefix(length) + self.forward(length) + return value + + def scan_tag_uri(self, name, start_mark): + # See the specification for details. + # Note: we do not check if URI is well-formed. + chunks = [] + length = 0 + ch = self.peek(length) + while '0' <= ch <= '9' or 'A' <= ch <= 'Z' or 'a' <= ch <= 'z' \ + or ch in '-;/?:@&=+$,_.!~*\'()[]%': + if ch == '%': + chunks.append(self.prefix(length)) + self.forward(length) + length = 0 + chunks.append(self.scan_uri_escapes(name, start_mark)) + else: + length += 1 + ch = self.peek(length) + if length: + chunks.append(self.prefix(length)) + self.forward(length) + length = 0 + if not chunks: + raise ScannerError("while parsing a %s" % name, start_mark, + "expected URI, but found %r" % ch, self.get_mark()) + return ''.join(chunks) + + def scan_uri_escapes(self, name, start_mark): + # See the specification for details. + codes = [] + mark = self.get_mark() + while self.peek() == '%': + self.forward() + for k in range(2): + if self.peek(k) not in '0123456789ABCDEFabcdef': + raise ScannerError("while scanning a %s" % name, start_mark, + "expected URI escape sequence of 2 hexadecimal numbers, but found %r" + % self.peek(k), self.get_mark()) + codes.append(int(self.prefix(2), 16)) + self.forward(2) + try: + value = bytes(codes).decode('utf-8') + except UnicodeDecodeError as exc: + raise ScannerError("while scanning a %s" % name, start_mark, str(exc), mark) + return value + + def scan_line_break(self): + # Transforms: + # '\r\n' : '\n' + # '\r' : '\n' + # '\n' : '\n' + # '\x85' : '\n' + # '\u2028' : '\u2028' + # '\u2029 : '\u2029' + # default : '' + ch = self.peek() + if ch in '\r\n\x85': + if self.prefix(2) == '\r\n': + self.forward(2) + else: + self.forward() + return '\n' + elif ch in '\u2028\u2029': + self.forward() + return ch + return '' diff --git a/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/serializer.py b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/serializer.py new file mode 100644 index 0000000000000000000000000000000000000000..fe911e67ae7a739abb491fbbc6834b9c37bbda4b --- /dev/null +++ b/workspace/outputs/audit_venv/lib/python3.11/site-packages/yaml/serializer.py @@ -0,0 +1,111 @@ + +__all__ = ['Serializer', 'SerializerError'] + +from .error import YAMLError +from .events import * +from .nodes import * + +class SerializerError(YAMLError): + pass + +class Serializer: + + ANCHOR_TEMPLATE = 'id%03d' + + def __init__(self, encoding=None, + explicit_start=None, explicit_end=None, version=None, tags=None): + self.use_encoding = encoding + self.use_explicit_start = explicit_start + self.use_explicit_end = explicit_end + self.use_version = version + self.use_tags = tags + self.serialized_nodes = {} + self.anchors = {} + self.last_anchor_id = 0 + self.closed = None + + def open(self): + if self.closed is None: + self.emit(StreamStartEvent(encoding=self.use_encoding)) + self.closed = False + elif self.closed: + raise SerializerError("serializer is closed") + else: + raise SerializerError("serializer is already opened") + + def close(self): + if self.closed is None: + raise SerializerError("serializer is not opened") + elif not self.closed: + self.emit(StreamEndEvent()) + self.closed = True + + #def __del__(self): + # self.close() + + def serialize(self, node): + if self.closed is None: + raise SerializerError("serializer is not opened") + elif self.closed: + raise SerializerError("serializer is closed") + self.emit(DocumentStartEvent(explicit=self.use_explicit_start, + version=self.use_version, tags=self.use_tags)) + self.anchor_node(node) + self.serialize_node(node, None, None) + self.emit(DocumentEndEvent(explicit=self.use_explicit_end)) + self.serialized_nodes = {} + self.anchors = {} + self.last_anchor_id = 0 + + def anchor_node(self, node): + if node in self.anchors: + if self.anchors[node] is None: + self.anchors[node] = self.generate_anchor(node) + else: + self.anchors[node] = None + if isinstance(node, SequenceNode): + for item in node.value: + self.anchor_node(item) + elif isinstance(node, MappingNode): + for key, value in node.value: + self.anchor_node(key) + self.anchor_node(value) + + def generate_anchor(self, node): + self.last_anchor_id += 1 + return self.ANCHOR_TEMPLATE % self.last_anchor_id + + def serialize_node(self, node, parent, index): + alias = self.anchors[node] + if node in self.serialized_nodes: + self.emit(AliasEvent(alias)) + else: + self.serialized_nodes[node] = True + self.descend_resolver(parent, index) + if isinstance(node, ScalarNode): + detected_tag = self.resolve(ScalarNode, node.value, (True, False)) + default_tag = self.resolve(ScalarNode, node.value, (False, True)) + implicit = (node.tag == detected_tag), (node.tag == default_tag) + self.emit(ScalarEvent(alias, node.tag, implicit, node.value, + style=node.style)) + elif isinstance(node, SequenceNode): + implicit = (node.tag + == self.resolve(SequenceNode, node.value, True)) + self.emit(SequenceStartEvent(alias, node.tag, implicit, + flow_style=node.flow_style)) + index = 0 + for item in node.value: + self.serialize_node(item, node, index) + index += 1 + self.emit(SequenceEndEvent()) + elif isinstance(node, MappingNode): + implicit = (node.tag + == self.resolve(MappingNode, node.value, True)) + self.emit(MappingStartEvent(alias, node.tag, implicit, + flow_style=node.flow_style)) + for key, value in node.value: + self.serialize_node(key, node, None) + self.serialize_node(value, node, key) + self.emit(MappingEndEvent()) + self.ascend_resolver() + diff --git a/workspace/outputs/audit_venv/pyvenv.cfg b/workspace/outputs/audit_venv/pyvenv.cfg new file mode 100644 index 0000000000000000000000000000000000000000..64b1e0aa13d290cd2f14c35170564077cda495d8 --- /dev/null +++ b/workspace/outputs/audit_venv/pyvenv.cfg @@ -0,0 +1,5 @@ +home = /cvmfs/soft.computecanada.ca/gentoo/2023/x86-64-v3/usr/lib/python-exec/python3.11 +include-system-site-packages = false +version = 3.11.4 +executable = /cvmfs/soft.computecanada.ca/gentoo/2023/x86-64-v3/usr/bin/python3.11 +command = /cvmfs/soft.computecanada.ca/gentoo/2023/x86-64-v3/usr/lib/python-exec/python3.11/python -m venv /lustre09/project/6037638/knguy52/vla/outputs/audit_venv diff --git a/workspace/outputs/audit_venv/share/man/man1/isympy.1 b/workspace/outputs/audit_venv/share/man/man1/isympy.1 new file mode 100644 index 0000000000000000000000000000000000000000..0ff966158a28c5ad1a6cd954e454842b25fdd999 --- /dev/null +++ b/workspace/outputs/audit_venv/share/man/man1/isympy.1 @@ -0,0 +1,188 @@ +'\" -*- coding: us-ascii -*- +.if \n(.g .ds T< \\FC +.if \n(.g .ds T> \\F[\n[.fam]] +.de URL +\\$2 \(la\\$1\(ra\\$3 +.. +.if \n(.g .mso www.tmac +.TH isympy 1 2007-10-8 "" "" +.SH NAME +isympy \- interactive shell for SymPy +.SH SYNOPSIS +'nh +.fi +.ad l +\fBisympy\fR \kx +.if (\nx>(\n(.l/2)) .nr x (\n(.l/5) +'in \n(.iu+\nxu +[\fB-c\fR | \fB--console\fR] [\fB-p\fR ENCODING | \fB--pretty\fR ENCODING] [\fB-t\fR TYPE | \fB--types\fR TYPE] [\fB-o\fR ORDER | \fB--order\fR ORDER] [\fB-q\fR | \fB--quiet\fR] [\fB-d\fR | \fB--doctest\fR] [\fB-C\fR | \fB--no-cache\fR] [\fB-a\fR | \fB--auto\fR] [\fB-D\fR | \fB--debug\fR] [ +-- | PYTHONOPTIONS] +'in \n(.iu-\nxu +.ad b +'hy +'nh +.fi +.ad l +\fBisympy\fR \kx +.if (\nx>(\n(.l/2)) .nr x (\n(.l/5) +'in \n(.iu+\nxu +[ +{\fB-h\fR | \fB--help\fR} +| +{\fB-v\fR | \fB--version\fR} +] +'in \n(.iu-\nxu +.ad b +'hy +.SH DESCRIPTION +isympy is a Python shell for SymPy. It is just a normal python shell +(ipython shell if you have the ipython package installed) that executes +the following commands so that you don't have to: +.PP +.nf +\*(T< +>>> from __future__ import division +>>> from sympy import * +>>> x, y, z = symbols("x,y,z") +>>> k, m, n = symbols("k,m,n", integer=True) + \*(T> +.fi +.PP +So starting isympy is equivalent to starting python (or ipython) and +executing the above commands by hand. It is intended for easy and quick +experimentation with SymPy. For more complicated programs, it is recommended +to write a script and import things explicitly (using the "from sympy +import sin, log, Symbol, ..." idiom). +.SH OPTIONS +.TP +\*(T<\fB\-c \fR\*(T>\fISHELL\fR, \*(T<\fB\-\-console=\fR\*(T>\fISHELL\fR +Use the specified shell (python or ipython) as +console backend instead of the default one (ipython +if present or python otherwise). + +Example: isympy -c python + +\fISHELL\fR could be either +\&'ipython' or 'python' +.TP +\*(T<\fB\-p \fR\*(T>\fIENCODING\fR, \*(T<\fB\-\-pretty=\fR\*(T>\fIENCODING\fR +Setup pretty printing in SymPy. By default, the most pretty, unicode +printing is enabled (if the terminal supports it). You can use less +pretty ASCII printing instead or no pretty printing at all. + +Example: isympy -p no + +\fIENCODING\fR must be one of 'unicode', +\&'ascii' or 'no'. +.TP +\*(T<\fB\-t \fR\*(T>\fITYPE\fR, \*(T<\fB\-\-types=\fR\*(T>\fITYPE\fR +Setup the ground types for the polys. By default, gmpy ground types +are used if gmpy2 or gmpy is installed, otherwise it falls back to python +ground types, which are a little bit slower. You can manually +choose python ground types even if gmpy is installed (e.g., for testing purposes). + +Note that sympy ground types are not supported, and should be used +only for experimental purposes. + +Note that the gmpy1 ground type is primarily intended for testing; it the +use of gmpy even if gmpy2 is available. + +This is the same as setting the environment variable +SYMPY_GROUND_TYPES to the given ground type (e.g., +SYMPY_GROUND_TYPES='gmpy') + +The ground types can be determined interactively from the variable +sympy.polys.domains.GROUND_TYPES inside the isympy shell itself. + +Example: isympy -t python + +\fITYPE\fR must be one of 'gmpy', +\&'gmpy1' or 'python'. +.TP +\*(T<\fB\-o \fR\*(T>\fIORDER\fR, \*(T<\fB\-\-order=\fR\*(T>\fIORDER\fR +Setup the ordering of terms for printing. The default is lex, which +orders terms lexicographically (e.g., x**2 + x + 1). You can choose +other orderings, such as rev-lex, which will use reverse +lexicographic ordering (e.g., 1 + x + x**2). + +Note that for very large expressions, ORDER='none' may speed up +printing considerably, with the tradeoff that the order of the terms +in the printed expression will have no canonical order + +Example: isympy -o rev-lax + +\fIORDER\fR must be one of 'lex', 'rev-lex', 'grlex', +\&'rev-grlex', 'grevlex', 'rev-grevlex', 'old', or 'none'. +.TP +\*(T<\fB\-q\fR\*(T>, \*(T<\fB\-\-quiet\fR\*(T> +Print only Python's and SymPy's versions to stdout at startup, and nothing else. +.TP +\*(T<\fB\-d\fR\*(T>, \*(T<\fB\-\-doctest\fR\*(T> +Use the same format that should be used for doctests. This is +equivalent to '\fIisympy -c python -p no\fR'. +.TP +\*(T<\fB\-C\fR\*(T>, \*(T<\fB\-\-no\-cache\fR\*(T> +Disable the caching mechanism. Disabling the cache may slow certain +operations down considerably. This is useful for testing the cache, +or for benchmarking, as the cache can result in deceptive benchmark timings. + +This is the same as setting the environment variable SYMPY_USE_CACHE +to 'no'. +.TP +\*(T<\fB\-a\fR\*(T>, \*(T<\fB\-\-auto\fR\*(T> +Automatically create missing symbols. Normally, typing a name of a +Symbol that has not been instantiated first would raise NameError, +but with this option enabled, any undefined name will be +automatically created as a Symbol. This only works in IPython 0.11. + +Note that this is intended only for interactive, calculator style +usage. In a script that uses SymPy, Symbols should be instantiated +at the top, so that it's clear what they are. + +This will not override any names that are already defined, which +includes the single character letters represented by the mnemonic +QCOSINE (see the "Gotchas and Pitfalls" document in the +documentation). You can delete existing names by executing "del +name" in the shell itself. You can see if a name is defined by typing +"'name' in globals()". + +The Symbols that are created using this have default assumptions. +If you want to place assumptions on symbols, you should create them +using symbols() or var(). + +Finally, this only works in the top level namespace. So, for +example, if you define a function in isympy with an undefined +Symbol, it will not work. +.TP +\*(T<\fB\-D\fR\*(T>, \*(T<\fB\-\-debug\fR\*(T> +Enable debugging output. This is the same as setting the +environment variable SYMPY_DEBUG to 'True'. The debug status is set +in the variable SYMPY_DEBUG within isympy. +.TP +-- \fIPYTHONOPTIONS\fR +These options will be passed on to \fIipython (1)\fR shell. +Only supported when ipython is being used (standard python shell not supported). + +Two dashes (--) are required to separate \fIPYTHONOPTIONS\fR +from the other isympy options. + +For example, to run iSymPy without startup banner and colors: + +isympy -q -c ipython -- --colors=NoColor +.TP +\*(T<\fB\-h\fR\*(T>, \*(T<\fB\-\-help\fR\*(T> +Print help output and exit. +.TP +\*(T<\fB\-v\fR\*(T>, \*(T<\fB\-\-version\fR\*(T> +Print isympy version information and exit. +.SH FILES +.TP +\*(T<\fI${HOME}/.sympy\-history\fR\*(T> +Saves the history of commands when using the python +shell as backend. +.SH BUGS +The upstreams BTS can be found at \(lahttps://github.com/sympy/sympy/issues\(ra +Please report all bugs that you find in there, this will help improve +the overall quality of SymPy. +.SH "SEE ALSO" +\fBipython\fR(1), \fBpython\fR(1) diff --git a/workspace/outputs/audit_venv/share/man/man1/ttx.1 b/workspace/outputs/audit_venv/share/man/man1/ttx.1 new file mode 100644 index 0000000000000000000000000000000000000000..bba23b5e51629509a499f4471fc8196e9863d211 --- /dev/null +++ b/workspace/outputs/audit_venv/share/man/man1/ttx.1 @@ -0,0 +1,225 @@ +.Dd May 18, 2004 +.\" ttx is not specific to any OS, but contrary to what groff_mdoc(7) +.\" seems to imply, entirely omitting the .Os macro causes 'BSD' to +.\" be used, so I give a zero-width space as its argument. +.Os \& +.\" The "FontTools Manual" argument apparently has no effect in +.\" groff 1.18.1. I think it is a bug in the -mdoc groff package. +.Dt TTX 1 "FontTools Manual" +.Sh NAME +.Nm ttx +.Nd tool for manipulating TrueType and OpenType fonts +.Sh SYNOPSIS +.Nm +.Bk +.Op Ar option ... +.Ek +.Bk +.Ar file ... +.Ek +.Sh DESCRIPTION +.Nm +is a tool for manipulating TrueType and OpenType fonts. It can convert +TrueType and OpenType fonts to and from an +.Tn XML Ns -based format called +.Tn TTX . +.Tn TTX +files have a +.Ql .ttx +extension. +.Pp +For each +.Ar file +argument it is given, +.Nm +detects whether it is a +.Ql .ttf , +.Ql .otf +or +.Ql .ttx +file and acts accordingly: if it is a +.Ql .ttf +or +.Ql .otf +file, it generates a +.Ql .ttx +file; if it is a +.Ql .ttx +file, it generates a +.Ql .ttf +or +.Ql .otf +file. +.Pp +By default, every output file is created in the same directory as the +corresponding input file and with the same name except for the +extension, which is substituted appropriately. +.Nm +never overwrites existing files; if necessary, it appends a suffix to +the output file name before the extension, as in +.Pa Arial#1.ttf . +.Ss "General options" +.Bl -tag -width ".Fl t Ar table" +.It Fl h +Display usage information. +.It Fl d Ar dir +Write the output files to directory +.Ar dir +instead of writing every output file to the same directory as the +corresponding input file. +.It Fl o Ar file +Write the output to +.Ar file +instead of writing it to the same directory as the +corresponding input file. +.It Fl v +Be verbose. Write more messages to the standard output describing what +is being done. +.It Fl a +Allow virtual glyphs ID's on compile or decompile. +.El +.Ss "Dump options" +The following options control the process of dumping font files +(TrueType or OpenType) to +.Tn TTX +files. +.Bl -tag -width ".Fl t Ar table" +.It Fl l +List table information. Instead of dumping the font to a +.Tn TTX +file, display minimal information about each table. +.It Fl t Ar table +Dump table +.Ar table . +This option may be given multiple times to dump several tables at +once. When not specified, all tables are dumped. +.It Fl x Ar table +Exclude table +.Ar table +from the list of tables to dump. This option may be given multiple +times to exclude several tables from the dump. The +.Fl t +and +.Fl x +options are mutually exclusive. +.It Fl s +Split tables. Dump each table to a separate +.Tn TTX +file and write (under the name that would have been used for the output +file if the +.Fl s +option had not been given) one small +.Tn TTX +file containing references to the individual table dump files. This +file can be used as input to +.Nm +as long as the referenced files can be found in the same directory. +.It Fl i +.\" XXX: I suppose OpenType programs (exist and) are also affected. +Don't disassemble TrueType instructions. When this option is specified, +all TrueType programs (glyph programs, the font program and the +pre-program) are written to the +.Tn TTX +file as hexadecimal data instead of +assembly. This saves some time and results in smaller +.Tn TTX +files. +.It Fl y Ar n +When decompiling a TrueType Collection (TTC) file, +decompile font number +.Ar n , +starting from 0. +.El +.Ss "Compilation options" +The following options control the process of compiling +.Tn TTX +files into font files (TrueType or OpenType): +.Bl -tag -width ".Fl t Ar table" +.It Fl m Ar fontfile +Merge the input +.Tn TTX +file +.Ar file +with +.Ar fontfile . +No more than one +.Ar file +argument can be specified when this option is used. +.It Fl b +Don't recalculate glyph bounding boxes. Use the values in the +.Tn TTX +file as is. +.El +.Sh "THE TTX FILE FORMAT" +You can find some information about the +.Tn TTX +file format in +.Pa documentation.html . +In particular, you will find in that file the list of tables understood by +.Nm +and the relations between TrueType GlyphIDs and the glyph names used in +.Tn TTX +files. +.Sh EXAMPLES +In the following examples, all files are read from and written to the +current directory. Additionally, the name given for the output file +assumes in every case that it did not exist before +.Nm +was invoked. +.Pp +Dump the TrueType font contained in +.Pa FreeSans.ttf +to +.Pa FreeSans.ttx : +.Pp +.Dl ttx FreeSans.ttf +.Pp +Compile +.Pa MyFont.ttx +into a TrueType or OpenType font file: +.Pp +.Dl ttx MyFont.ttx +.Pp +List the tables in +.Pa FreeSans.ttf +along with some information: +.Pp +.Dl ttx -l FreeSans.ttf +.Pp +Dump the +.Sq cmap +table from +.Pa FreeSans.ttf +to +.Pa FreeSans.ttx : +.Pp +.Dl ttx -t cmap FreeSans.ttf +.Sh NOTES +On MS\-Windows and MacOS, +.Nm +is available as a graphical application to which files can be dropped. +.Sh SEE ALSO +.Pa documentation.html +.Pp +.Xr fontforge 1 , +.Xr ftinfo 1 , +.Xr gfontview 1 , +.Xr xmbdfed 1 , +.Xr Font::TTF 3pm +.Sh AUTHORS +.Nm +was written by +.An -nosplit +.An "Just van Rossum" Aq just@letterror.com . +.Pp +This manual page was written by +.An "Florent Rougon" Aq f.rougon@free.fr +for the Debian GNU/Linux system based on the existing FontTools +documentation. It may be freely used, modified and distributed without +restrictions. +.\" For Emacs: +.\" Local Variables: +.\" fill-column: 72 +.\" sentence-end: "[.?!][]\"')}]*\\($\\| $\\| \\| \\)[ \n]*" +.\" sentence-end-double-space: t +.\" End: \ No newline at end of file diff --git a/workspace/outputs/external_vla/same_split_comparison.json b/workspace/outputs/external_vla/same_split_comparison.json new file mode 100644 index 0000000000000000000000000000000000000000..8f11a26e8b07b8beee9852f0ce2e26650d87495b --- /dev/null +++ b/workspace/outputs/external_vla/same_split_comparison.json @@ -0,0 +1,28 @@ +{ + "comparison_protocol": "same_700_group_heldout_candidate_selection", + "dataset_groups": 3500, + "evaluation_groups": 700, + "validation_group_ids_sha256": "a7e51209e227ee8b68090e7826368541f209e1365112ed718c465c3bb0f11d53", + "seed": 0, + "candidate_oracle_success_rate": 0.4185714285714286, + "dovla_iaf": { + "top1_action_selection": 0.6171428571428571, + "selected_success_rate": 0.37857142857142856, + "mean_selected_regret": 0.059859846833028967 + }, + "smolvla_expert_only_bc": { + "checkpoint_revision": "c83c3163b8ca9b7e67c509fffd9121e66cb96205", + "model_sha256": "7cd549ac2351fb069c0ddb3c34ad2d09cfc92b56a15dccdfc2e41467aaca01eb", + "training_groups": 2800, + "training_steps": 1000, + "top1_action_selection": 0.5228571428571429, + "selected_success_rate": 0.3457142857142857, + "mean_selected_regret": 0.13656934786188815 + }, + "dovla_minus_smolvla": { + "top1_action_selection": 0.0942857142857142, + "selected_success_rate": 0.03285714285714286, + "mean_selected_regret": -0.07670950102885918 + }, + "scope": "Both methods select among the same measured same-state CIL candidates. This artifact does not compare online policy rollout success." +} diff --git a/workspace/outputs/external_vla/smolvla_cil_aligned_manifest.json b/workspace/outputs/external_vla/smolvla_cil_aligned_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..59b470cbae5b18cc71b9d41c6a25baf350fec653 --- /dev/null +++ b/workspace/outputs/external_vla/smolvla_cil_aligned_manifest.json @@ -0,0 +1,5782 @@ +{ + "action_normalizer": { + "canonical_action_dim": 8, + "mean": [ + -0.011001614853739738, + 0.0922699049115181, + -0.22631123661994934, + -0.11120374500751495, + -0.2544722855091095, + 0.322357714176178, + -0.8979408740997314, + -0.038035713136196136 + ], + "scale": [ + 0.712899923324585, + 0.5003968477249146, + 1.0223413705825806, + 1.1761223077774048, + 1.359087347984314, + 1.5495193004608154, + 1.6576342582702637, + 0.3774811327457428 + ], + "vectorization": "right_zero_pad_to_canonical_dimension" + }, + "checkpoint": "/scratch/knguy52/dovla/models/smolvla_base-c83c316", + "claim_scope": "Expert-only SmolVLA fine-tuning evaluated by nearest executed same-state CIL candidate; this is measured candidate selection, not online policy rollout.", + "config": { + "action_dim": 8, + "action_horizon": 4, + "batch_size": 4, + "device": "cuda", + "export_dir": "/scratch/knguy52/dovla/experiments/external_vla_export_full_aligned", + "image_size": 512, + "learning_rate": 0.0001, + "log_every": 25, + "max_eval_groups": 700, + "seed": 0, + "split_mode": "dataset_group_shuffle", + "state_dim": 32, + "steps": 1000, + "val_fraction": 0.2, + "vlm_metadata": "/scratch/knguy52/dovla/models/SmolVLM2-500M-Video-Instruct-metadata-7b375e1b73b11138ff12fe22c8f2822d8fe03467" + }, + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "export": "/scratch/knguy52/dovla/experiments/external_vla_export_full_aligned", + "num_train_groups": 2800, + "num_validation_groups": 700, + "schema_version": "smolvla-cil-candidate-baseline/v0", + "state_projector": { + "components": [ + [ + -0.0009472116362303495, + 0.000963377533480525, + -0.00011321854981360957, + -0.0018236824544146657, + 0.004122762940824032, + 0.0013067638501524925, + -0.0029155476950109005, + -0.02485238015651703, + -0.0004786144709214568, + 0.0038494563195854425, + 0.3833312690258026, + 0.6699730157852173, + 0.6323888301849365, + -0.003385609481483698, + -7.469352567568421e-05, + -0.006712593138217926, + -0.0010648821480572224, + 3.155859303660691e-05, + -0.004769495688378811, + 0.006381289102137089, + -7.009412911429536e-06, + -5.825750849908218e-05, + 4.8305846576113254e-05, + 0.0008909340831451118, + -0.0002910454059019685, + -0.0020156570244580507, + -0.0033816667273640633, + 7.327997122953331e-15, + 0.006282679736614227, + 0.0020009453874081373, + 1.8588054272877796e-12, + -2.2472839561471503e-11, + -0.004830706398934126, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.004310313146561384, + 0.002294174861162901, + 0.0003723612753674388, + -0.02121465839445591, + 0.0006104116328060627, + 0.015162331983447075, + 0.00550801120698452, + 0.00018319861555937678, + 0.00018329954764340073, + 0.0015486417105421424, + -0.004190572537481785, + 0.000742288597393781, + 0.0027123582549393177, + -0.0008236839785240591, + -0.01488953735679388, + 0.0028437168803066015, + 0.011744999326765537, + 0.0014020497910678387, + -0.018270889297127724, + -0.009582971222698689, + 0.00012822041753679514, + 0.00012829111074097455, + -0.0036129371728748083, + -0.0075623467564582825, + 0.0038340461906045675, + -0.03357479348778725, + 0.0024324883706867695, + 0.0026926230639219284, + -0.0032499332446604967, + -0.002558450447395444, + -0.0006797343958169222, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332, + -0.006831651087850332 + ], + [ + 0.00028777611441910267, + -0.0005648690275847912, + -0.000618185440544039, + 0.003306581173092127, + -0.0035802193451672792, + -0.00047961081145331264, + -0.0052907126955688, + 0.019344132393598557, + 0.00030329200671985745, + -0.006930300500243902, + -0.4206838607788086, + -0.47960638999938965, + 0.7683088183403015, + 0.0025843093171715736, + 0.00019414866983424872, + 0.00419723242521286, + -0.000991127104498446, + 9.32571419980377e-05, + 0.006722330581396818, + -0.004552920814603567, + 2.9965998692205176e-05, + -0.00023754333960823715, + -2.729038351390045e-05, + -0.00014826397818978876, + 0.0001172870397567749, + -0.000765786215197295, + 0.002441998803988099, + -6.187553656836035e-15, + -0.004536904394626617, + -0.0014449403388425708, + -1.342297041746765e-12, + 1.622828607605875e-11, + 0.0034883925691246986, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0030843657441437244, + -0.0016836406430229545, + -0.00028395975823514163, + 0.015331820584833622, + -0.0004635473887901753, + -0.011031106114387512, + -0.0039581903256475925, + -0.00013275993114802986, + -0.0001328223297605291, + -0.0011720493203029037, + 0.0026348892133682966, + -0.0004679737612605095, + -0.0019669116009026766, + 0.0014377391198650002, + 0.010130094364285469, + -0.00285517075099051, + -0.011027439497411251, + -0.0013305587926879525, + 0.016925115138292313, + 0.008922496810555458, + -0.0001727676863083616, + -0.000173356500454247, + 0.004864595830440521, + 0.008960355073213577, + -0.004625855479389429, + 0.024467600509524345, + -0.003591944696381688, + 0.003595636459067464, + 0.0053748139180243015, + 0.001221431652083993, + -0.0002902651613112539, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304, + 0.004933332558721304 + ], + [ + -0.00021219842892605811, + 0.001886711805127561, + -0.0018270873697474599, + 3.700465640577022e-06, + -0.0023816274479031563, + -0.0020127345342189074, + 0.007281394209712744, + -0.007683674339205027, + -0.0007991645252332091, + 0.0011669363593682647, + 0.8203607797622681, + -0.5618172883987427, + 0.0982627272605896, + -0.002001512795686722, + 0.0005945450975559652, + 0.0007348493090830743, + -0.005563501734286547, + -5.887417864869349e-05, + 0.0005016142386011779, + 0.0019441801123321056, + -1.524409412922978e-06, + 0.0004916038597002625, + 0.00017737768939696252, + 0.003944726660847664, + -0.0004645549925044179, + -0.00041614152723923326, + 0.0014551326166838408, + 1.4121319309270725e-15, + -0.0027034373488277197, + -0.0008610064396634698, + -7.998440776889326e-13, + 9.670064228528563e-12, + 0.0020786537788808346, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0019613688345998526, + -0.0006069272058084607, + -0.0001936090993694961, + 0.009370852261781693, + -0.0002933931245934218, + -0.005942825227975845, + -0.002435459988191724, + -4.714156602858566e-05, + -4.717771298601292e-05, + -0.0022559266071766615, + 0.00540586095303297, + -0.0002260145847685635, + -0.004979391116648912, + 0.0027967405039817095, + 0.012554952874779701, + -0.0045832558535039425, + -0.008674629963934422, + 0.0002650513779371977, + 0.011268007569015026, + -0.01941850408911705, + 0.0003743432171177119, + 0.00037336317473091185, + 0.004040434490889311, + 0.0043673040345311165, + 0.002781508257612586, + -0.017324263229966164, + 0.006436229683458805, + 0.005919283721596003, + 0.0028538722544908524, + -0.0036851142067462206, + 0.0004292596422601491, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339, + 0.002939659869298339 + ], + [ + -0.0035368239041417837, + -0.0025141295045614243, + 0.002715041860938072, + -0.020711876451969147, + 0.09768533706665039, + 0.026087479665875435, + -0.04910906031727791, + -0.005947052501142025, + 0.024494454264640808, + 0.018811529502272606, + -0.014897970482707024, + -0.06262388080358505, + 0.00035707670031115413, + -0.015890831127762794, + 0.0009351579356007278, + -0.14706049859523773, + -0.011604963801801205, + -2.4451199351460673e-05, + -0.03789234906435013, + 0.08970329910516739, + 0.00010813747212523594, + -0.0009062914759851992, + -0.0005746327806264162, + -0.0053315917029976845, + 0.0038937812205404043, + 0.0019047794630751014, + -0.07405145466327667, + 2.11549907366132e-13, + 0.13757763803005219, + 0.043816544115543365, + 4.0703981851741844e-11, + -4.921085205999987e-10, + -0.10578245669603348, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.09181497246026993, + 0.05787895247340202, + 0.0075806197710335255, + -0.4606916606426239, + 0.013280047103762627, + 0.34285637736320496, + 0.11825945973396301, + 0.004391293507069349, + 0.004393697250634432, + 0.018198847770690918, + -0.0461345799267292, + 0.014078169129788876, + 0.0026722862385213375, + 0.014360242523252964, + -0.18286921083927155, + 0.025172075256705284, + 0.2911907434463501, + -0.00334531138651073, + -0.36586523056030273, + -0.03259829059243202, + -0.0034899981692433357, + -0.003491275943815708, + 0.0035526854917407036, + -0.07077716290950775, + 0.003584472229704261, + 0.0012273247120901942, + 0.0006526464130729437, + -0.05110594257712364, + 0.03972948342561722, + 0.0024146665818989277, + 0.004557049833238125, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565, + -0.14959898591041565 + ], + [ + 0.006643514148890972, + 0.004861366469413042, + -0.0009809096809476614, + 0.010610189288854599, + -0.003935337532311678, + -0.0014387257397174835, + 0.0006284701521508396, + -0.002134674461558461, + 0.010926930233836174, + -0.013625686056911945, + 0.0025477767921984196, + 0.007508916314691305, + -0.0029595375526696444, + 0.010308397002518177, + 0.004217203706502914, + 0.0022923001088202, + -0.006703463848680258, + -5.462816989165731e-05, + 0.012025481089949608, + 0.0009673480526544154, + -0.00015971717948559672, + 0.0013426186051219702, + 0.0002739790070336312, + 7.77445802668808e-06, + -0.0019271208439022303, + 0.0021201265044510365, + 0.0029296313878148794, + -8.566866664574523e-15, + -0.00544288195669651, + -0.001733480952680111, + -1.6103413090842733e-12, + 1.9468923001531024e-11, + 0.004184992518275976, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0036803369875997305, + -0.0022561531513929367, + -0.0002925725420936942, + 0.01829557493329048, + -0.0005320000345818698, + -0.013401403091847897, + -0.004765799269080162, + -0.000161550851771608, + -0.00016161156236194074, + -0.0018253837479278445, + 0.003202748252078891, + -0.00032630356145091355, + -0.0015374189242720604, + 0.033909205347299576, + 0.013330611400306225, + -0.037495166063308716, + -0.0029379925690591335, + -0.054318107664585114, + 0.013683870434761047, + 0.04405154660344124, + -0.00019740458810701966, + -0.0002056688244920224, + 0.14590875804424286, + -0.04547932371497154, + 0.0048612128011882305, + -0.06824417412281036, + -0.352795273065567, + 0.009902678430080414, + 0.9148979187011719, + -0.022843623533844948, + 0.004660950507968664, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948, + 0.0059184730052948 + ], + [ + -0.01567704603075981, + 0.004647674970328808, + 0.0036550904624164104, + -0.002968788845464587, + -0.0030134820844978094, + -0.0008271039114333689, + 0.0012488190550357103, + -0.10125979781150818, + -0.006164261139929295, + 0.006012462545186281, + -0.03501681610941887, + -0.0272272527217865, + -0.0025355233810842037, + -0.012956132180988789, + -0.0010808567749336362, + 0.003402134869247675, + 0.00107424461748451, + 0.0002524742449168116, + -0.003099802415817976, + 0.0053689368069171906, + -0.0009003874729387462, + 0.0008104328881017864, + 0.0006052414537407458, + 0.011743736453354359, + -0.006107192486524582, + -0.0020163815934211016, + 0.0022325499448925257, + -5.4508298103026626e-15, + -0.004147767089307308, + -0.0013210047036409378, + -1.2271662635718328e-12, + 1.483636130616972e-11, + 0.0031891888938844204, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0027854163199663162, + -0.001509582856670022, + -0.0002513461804483086, + 0.014108086004853249, + -0.00041063682874664664, + -0.010381030850112438, + -0.0035366308875381947, + -0.00013445490912999958, + -0.00013452860002871603, + -0.0002802973613142967, + 0.0010135613847523928, + -0.0006688328576274216, + -0.00059579371009022, + -0.010104799643158913, + -0.022731635719537735, + 0.009199703112244606, + -0.047731664031744, + 0.009546776302158833, + 0.005974017549306154, + -0.023902928456664085, + 0.0002854045305866748, + 0.0002869741292670369, + -0.03474234789609909, + -0.5763009190559387, + 0.011627557687461376, + -0.7877925038337708, + 0.016924560070037842, + 0.15031257271766663, + -0.0746656209230423, + 0.009000043384730816, + 0.008224396966397762, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071, + 0.004510193131864071 + ], + [ + 0.005781878717243671, + 0.027764244005084038, + 0.025633668527007103, + -0.08426334708929062, + -0.008637885563075542, + -0.0007435394800268114, + 0.08611427992582321, + 0.020628158003091812, + -0.01522674411535263, + 0.12045323103666306, + 0.002015278209000826, + 0.0014183924067765474, + 0.007435949053615332, + -0.008376898244023323, + -0.0040024020709097385, + 0.05364569276571274, + 0.02883213944733143, + 0.0005996229010634124, + -0.17969700694084167, + 0.044258829206228256, + -0.0004712945083156228, + 0.0033228828106075525, + 0.0024858934339135885, + 0.03585948795080185, + -0.019454117864370346, + -0.015396024100482464, + 0.010553545318543911, + -2.72887351008156e-14, + -0.01960698515176773, + -0.006244543939828873, + -5.80095953911508e-12, + 7.013323255478099e-11, + 0.01507567334920168, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.01342887431383133, + -0.007509471382945776, + -0.0011212173849344254, + 0.06638581305742264, + -0.001979030203074217, + -0.04775450751185417, + -0.0172322615981102, + -0.0005517930840142071, + -0.0005520560662262142, + -0.006529186386615038, + 0.01553291268646717, + -0.0011072984198108315, + -0.00976475328207016, + -0.017322752624750137, + 0.08348752558231354, + -0.00023456728376913816, + 0.04147212207317352, + 0.025848882272839546, + -0.050687044858932495, + -0.2684173583984375, + 0.0033838439267128706, + 0.0033904570154845715, + -0.07013195008039474, + -0.4758625626564026, + 0.08967868983745575, + 0.22076164186000824, + 0.09592713415622711, + -0.7199904918670654, + 0.06539137661457062, + 0.06998506933450699, + 0.07262426614761353, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828, + 0.021320214495062828 + ], + [ + -0.0020717366132885218, + -0.07711230963468552, + -0.002364260144531727, + 0.03794151544570923, + -0.0037383362650871277, + 0.01802484132349491, + -0.3186562657356262, + -0.023625006899237633, + 0.016296999529004097, + 0.1370832473039627, + 0.026819974184036255, + -0.008078832179307938, + 0.002207433804869652, + -0.013194615021348, + 0.020002715289592743, + 0.046252038329839706, + 0.12605145573616028, + -0.0007313807727769017, + 0.10208398103713989, + -0.2628140151500702, + 0.0016298909904435277, + -0.011218896135687828, + -0.006521890871226788, + -0.06431412696838379, + 0.0472506545484066, + 0.0161205492913723, + 0.004007960204035044, + -2.868377125988912e-14, + -0.00744636170566082, + -0.002371578710153699, + -2.2030946945222896e-12, + 2.6635272112884856e-11, + 0.0057254359126091, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 5.151764344191179e-05, + -0.009413211606442928, + -2.809546458593104e-05, + 0.02443695068359375, + -0.0006127957603894174, + -0.04111561179161072, + 0.0016147415153682232, + -0.0011380685027688742, + -0.0011381655931472778, + 0.046722639352083206, + -0.08977194130420685, + -0.011055897921323776, + 0.1130184531211853, + -0.03388866037130356, + -0.2195654809474945, + 0.11328992247581482, + -0.005833495408296585, + -0.012607197277247906, + -0.046648140996694565, + 0.7491308450698853, + -0.002038354752585292, + -0.002024043118581176, + -0.05405610799789429, + -0.16411492228507996, + 0.015114863403141499, + 0.06556638330221176, + -0.12261167168617249, + -0.2434726506471634, + -0.0654730275273323, + 0.013814241625368595, + 0.0015236588660627604, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581, + 0.008097013458609581 + ], + [ + 0.018987011164426804, + -0.03986372798681259, + -0.013487123884260654, + 0.021269742399454117, + 0.018618732690811157, + -0.02672199159860611, + -0.0774327740073204, + 0.010343914851546288, + -0.06059436500072479, + -0.0446791835129261, + 0.0024252308066934347, + 0.011143572628498077, + 0.0002496363304089755, + 0.0005231939721852541, + 0.005707764066755772, + -0.008240770548582077, + 0.04124847799539566, + -0.0001361915929010138, + 0.03789965435862541, + -0.08071812987327576, + 0.0005154949030838907, + -0.004013187251985073, + -0.0027738220524042845, + -0.03639586269855499, + 0.0197750236839056, + 0.01333328802138567, + -0.009760669432580471, + 8.07178962883813e-14, + 0.01813395693898201, + 0.005775403697043657, + 5.365146694968903e-12, + -6.486428061336369e-11, + -0.013943074271082878, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0024530014488846064, + 0.027432380244135857, + 7.625138096045703e-05, + -0.05300501734018326, + 0.0007162196561694145, + 0.1072751134634018, + -0.0053638555109500885, + 0.0034619688522070646, + 0.003463780740275979, + -0.15855154395103455, + 0.3147636651992798, + 0.03744424507021904, + -0.36741647124290466, + 0.16945703327655792, + 0.70506352186203, + -0.2637694180011749, + 0.10106298327445984, + -0.00796563271433115, + -0.0786704570055008, + 0.2684479355812073, + -0.0021378390956670046, + -0.002132839523255825, + -0.0381523035466671, + -0.023717250674962997, + 0.030153783038258553, + -0.03385419398546219, + -0.002588916104286909, + -0.04542035982012749, + -0.03519589081406593, + 0.0016990680014714599, + -0.0025277817621827126, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201, + -0.0197184756398201 + ], + [ + -0.004673880059272051, + 0.001983449561521411, + -0.001029874081723392, + -0.026969879865646362, + 0.0012838913826271892, + 0.001369938370771706, + -0.024468909949064255, + -0.011186862364411354, + 0.09496855735778809, + 0.0010071586584672332, + -0.004984179511666298, + 0.002669412177056074, + -0.000745433964766562, + -0.0018741092644631863, + 0.029784370213747025, + 0.013018285855650902, + 0.029068002477288246, + -0.0006249853759072721, + 0.0010402987245470285, + -0.045152876526117325, + 0.00036891535273753107, + 0.003509495174512267, + -0.0017082274425774813, + -0.07563658803701401, + 0.01685331016778946, + 0.05848875641822815, + 0.0010923083173111081, + -1.7997617149856024e-14, + -0.0020293111447244883, + -0.0006463108584284782, + -6.00395872341758e-13, + 7.25874742257826e-12, + 0.001560318865813315, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0005686478107236326, + -0.0026408813428133726, + 3.1498817406827584e-05, + 0.006286994554102421, + -0.00011460261157481, + -0.009286781772971153, + -0.0009634021553210914, + -0.00023040799715090543, + -0.00023053814948070794, + 0.0032589146867394447, + -0.0223851278424263, + -0.0009390302002429962, + 0.01759897544980049, + 0.07409999519586563, + -0.046856872737407684, + -0.05051330104470253, + -0.009016167372465134, + 0.04779987409710884, + -0.00541545357555151, + 0.144154354929924, + 0.000470163970021531, + 0.0004653318028431386, + 0.40308406949043274, + 0.038222573697566986, + 0.2402276247739792, + -0.04764997959136963, + 0.8131104111671448, + 0.016023876145482063, + 0.2367507368326187, + 0.0013747001066803932, + 0.016167882829904556, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713, + 0.0022066296078264713 + ], + [ + -0.0022105087991803885, + 0.03594091907143593, + 0.0003541698679327965, + 0.08875102549791336, + 0.03158785030245781, + -0.038682326674461365, + 0.8520680069923401, + 0.004523669369518757, + 0.03618922829627991, + 0.03416059911251068, + 0.004002342000603676, + 0.0019501505885273218, + 0.0045023113489151, + -0.007350060157477856, + 0.01577303558588028, + 0.0036886620800942183, + 0.04402325674891472, + 0.0001018769180518575, + 0.11854707449674606, + 0.1110486388206482, + -0.0005018714582547545, + 0.00703011266887188, + 0.0048326896503567696, + 0.09104481339454651, + -0.051338955760002136, + -0.027289381250739098, + -0.002873622579500079, + 1.1492636201970545e-14, + 0.005338996648788452, + 0.001700401771813631, + 1.5796056953773507e-12, + -1.9097331355188985e-11, + -0.004105109255760908, + -1.8932661725304283e-29, + 7.629764067684474e-30, + -2.5736201846846633e-30, + 3.333911122886906e-30, + -1.310481125528343e-29, + 2.7088225356507047e-30, + -0.0043208119459450245, + 0.0029631080105900764, + -0.0011606841580942273, + -0.043017737567424774, + 0.0026951285544782877, + 0.057041943073272705, + -0.014973372220993042, + 0.000537724292371422, + 0.0005335333989933133, + 0.005472661927342415, + -0.1440783590078354, + 0.00836483296006918, + -0.14690513908863068, + 0.0162937231361866, + -0.0402829572558403, + 0.0029434780590236187, + -0.04088174179196358, + -0.014266472309827805, + 0.029352840036153793, + 0.3464452922344208, + -0.00467366399243474, + -0.00467931292951107, + 0.13565295934677124, + -0.1149403527379036, + -0.03856941685080528, + 0.07457690685987473, + -0.052226126194000244, + -0.0034168781712651253, + -0.06078556552529335, + -0.024901393800973892, + -0.025994863361120224, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252, + -0.005805510561913252 + ], + [ + 0.0032210485078394413, + -0.006651272065937519, + 0.0028617356438189745, + -0.03141390532255173, + 0.0007000859477557242, + -0.0053685810416936874, + 0.11951103806495667, + -0.016862235963344574, + 0.034370057284832, + -0.0047234464436769485, + 0.001686516567133367, + -0.0012876302935183048, + 0.0026528199668973684, + 0.0018636105814948678, + -0.017661431804299355, + 0.002079004654660821, + 0.0021953426767140627, + -0.003177723614498973, + -0.008508189581334591, + 0.03964848071336746, + 0.00577571801841259, + 0.020361486822366714, + -0.008125434629619122, + -0.7612603306770325, + 0.14135368168354034, + 0.606335461139679, + 0.0017843962414190173, + -4.247988815092932e-15, + -0.0033152461983263493, + -0.0010558589128777385, + -9.80855045622231e-13, + 1.1858474271986275e-11, + 0.002549069235101342, + -4.338734978715565e-29, + 1.5999085234013646e-29, + -5.4176409335612945e-30, + 7.0188952005113e-30, + -2.758964761221753e-29, + 5.702902379455002e-30, + -0.0026913208421319723, + -0.0015705731930211186, + -0.0003804739681072533, + 0.009022435173392296, + -0.00014937848027329892, + -0.006020312197506428, + -0.004163178149610758, + -0.00014310052210930735, + -0.00014351199206430465, + 0.00554590905085206, + -0.01877281628549099, + -0.00078008504351601, + -0.0036559668369591236, + -0.03035036474466324, + 0.001077081891708076, + 0.008522557094693184, + 0.01768552139401436, + -0.007250536233186722, + -0.005920854397118092, + -0.051370374858379364, + 0.0015042038867250085, + 0.0015073162503540516, + -0.01546396967023611, + -0.018670229241251945, + -0.008504033088684082, + 0.0031418895814567804, + -0.08435149490833282, + -0.009403843432664871, + -0.02859789878129959, + 0.006055323872715235, + 0.0022954531013965607, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115, + 0.003604927798733115 + ], + [ + -0.0028311824426054955, + 0.0264583807438612, + 0.012388928793370724, + -0.014402540400624275, + -0.021348875015974045, + 0.05897968262434006, + 0.16764144599437714, + 0.06890502572059631, + -0.012282154522836208, + -0.07362271845340729, + 0.0034617625642567873, + 0.0019755405373871326, + 0.0012326191645115614, + 0.002486218698322773, + 0.004445865750312805, + -0.012184027582406998, + 0.008433528244495392, + -0.00016443763161078095, + 0.03098178654909134, + 0.06341402232646942, + -0.00024996037245728076, + 0.0038087659049779177, + 0.00151269871275872, + 0.013508501462638378, + -0.023447072133421898, + 0.016609763726592064, + -0.007200829219073057, + 9.935704602809237e-14, + 0.01337830163538456, + 0.004260803572833538, + 3.958129830838875e-12, + -4.785353363567779e-11, + -0.010286476463079453, + -2.3224065049706588e-27, + 8.551252212595768e-28, + -2.9099106641340073e-28, + 3.7691527532427062e-28, + -1.4815736098283796e-27, + 3.062478019503503e-28, + 0.022042883560061455, + 0.011984081007540226, + 0.005377115681767464, + 0.008741462603211403, + -5.517863610293716e-05, + -0.03300929442048073, + 0.04182828217744827, + 0.0006352701457217336, + 0.000638299563433975, + -0.019375331699848175, + 0.7115040421485901, + -0.03151160478591919, + 0.6422378420829773, + -0.039895955473184586, + 0.024231107905507088, + 0.03836693987250328, + 0.02834121696650982, + -0.01066760066896677, + -0.02319585159420967, + 0.08937278389930725, + -0.0016723163425922394, + -0.0016741255531087518, + 0.07035109400749207, + -0.05676713213324547, + -0.02324746735394001, + 0.029609153047204018, + -0.009201979264616966, + 0.017633549869060516, + -0.01914694719016552, + -0.0059072175063192844, + -0.007946673780679703, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416, + -0.014547279104590416 + ], + [ + 0.0030637490563094616, + -0.04988454282283783, + 0.00019176947535015643, + -0.03756225109100342, + 0.0038679963909089565, + 0.0005993748782202601, + 0.14708970487117767, + -0.027015971019864082, + -0.06445512920618057, + -0.0033724061213433743, + -0.0032050127629190683, + -0.004215530585497618, + 0.0002599329163786024, + -0.008477036841213703, + -0.029378507286310196, + 0.006445516832172871, + 0.03481029346585274, + -0.002031730953603983, + -0.02438466250896454, + -0.045905232429504395, + 0.012610082514584064, + -0.028001751750707626, + -0.008706164546310902, + -0.3432750403881073, + 0.6505179405212402, + -0.6018591523170471, + 0.001458176877349615, + 4.473282638403631e-15, + -0.0027091766241937876, + -0.0008628413779661059, + -8.015421009213708e-13, + 9.690592946143273e-12, + 0.002083059400320053, + -1.454028420503369e-25, + 5.492996255234949e-26, + -1.8528764941831125e-26, + 2.40016846870282e-26, + -9.434496627340981e-26, + 1.9501522882605963e-26, + -0.0016864450881257653, + -0.0003025095211341977, + -7.80575574026443e-05, + 0.010110187344253063, + -5.019148375140503e-05, + -0.00756017817184329, + -0.0026352766435593367, + -0.0001323012838838622, + -0.00013277828111313283, + 0.0013997016940265894, + 0.013818797655403614, + -0.004435352515429258, + 0.03819136694073677, + -0.012245092540979385, + -0.024055836722254753, + -0.024261515587568283, + 0.018154166638851166, + 0.0026419374626129866, + -0.045669712126255035, + -0.011555451899766922, + 0.0028603817336261272, + 0.0028664818964898586, + -0.1831994205713272, + 0.08366275578737259, + 0.04611270874738693, + -0.0650978609919548, + 0.057936109602451324, + -0.023590730503201485, + 0.05671311542391777, + 0.012465589679777622, + 0.01702910102903843, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936, + 0.0029459004290401936 + ], + [ + -0.03873986750841141, + 0.020548559725284576, + -0.01689395308494568, + 0.16900649666786194, + 0.0036870536860078573, + -0.00599060021340847, + -0.22360311448574066, + 0.1224856749176979, + 0.07017803192138672, + -0.04964170604944229, + 0.013603930361568928, + 0.013948674313724041, + -0.0027776346541941166, + 0.00800227839499712, + 0.006198782008141279, + -0.07550498843193054, + -0.13769447803497314, + -0.0023039064835757017, + 0.22853657603263855, + 0.11784907430410385, + 0.003795033320784569, + 0.004285074770450592, + 0.00046139751793816686, + -0.11771513521671295, + 0.1822693794965744, + -0.162026509642601, + -0.006421947851777077, + 2.4001417173180606e-14, + 0.011931059882044792, + 0.003799877827987075, + 3.529946083699609e-12, + -4.267682041092513e-11, + -0.009173702448606491, + -4.6852026882886333e-26, + 1.7872432668687243e-26, + -5.960632999849965e-27, + 7.73438674523941e-27, + -3.040172245555e-26, + 6.284169200835599e-27, + 0.005270111374557018, + 0.005974927451461554, + -4.238692781655118e-06, + -0.045981280505657196, + 0.0012094905832782388, + 0.04555113986134529, + 0.006480079144239426, + 0.0008020701352506876, + 0.0008017318905331194, + -0.00128337147179991, + 0.016078388318419456, + 0.009731477126479149, + -0.09666068851947784, + -0.07071041315793991, + 0.0641632005572319, + 0.09147944301366806, + -0.13968266546726227, + -0.012969260104000568, + 0.15870921313762665, + -0.08387687057256699, + -0.006266266573220491, + -0.006266178097575903, + 0.5340600609779358, + -0.37533625960350037, + -0.35573965311050415, + 0.26093509793281555, + -0.061214085668325424, + 0.09796462953090668, + -0.10497628897428513, + -0.005647444631904364, + -0.03295774385333061, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906, + -0.01297357864677906 + ], + [ + 0.03626488894224167, + 0.09368588775396347, + 0.007782211061567068, + 0.248029887676239, + -0.01761319302022457, + 0.025640087202191353, + -0.17789001762866974, + -0.06943437457084656, + 0.0873660072684288, + 0.009234270080924034, + -0.0016549559077247977, + -0.003082236973568797, + -0.0014719947939738631, + -0.006738076917827129, + -0.0004907988477498293, + 0.017811955884099007, + -0.018585963174700737, + -0.0013422529445961118, + 0.1377657949924469, + 0.8241116404533386, + 0.00259789964184165, + 0.018430296331644058, + 0.005664434749633074, + -0.04510904476046562, + -0.0029581093695014715, + -0.060577068477869034, + 0.007141350768506527, + -3.4612070154427244e-14, + -0.01326739601790905, + -0.004225464072078466, + -3.925317102609904e-12, + 4.745683013229751e-11, + 0.010201220400631428, + -1.852884572118782e-22, + 7.77549775799846e-23, + -2.5952791718851243e-23, + 3.363975530191239e-23, + -1.3222853676700943e-22, + 2.733218758683346e-23, + -0.004637578502297401, + -0.008450135588645935, + -0.000146148944622837, + 0.05492947995662689, + -0.002520991489291191, + -0.05842091888189316, + -0.002594849793240428, + -0.0008487197337672114, + -0.0008458895608782768, + -0.003829450812190771, + -0.041739922016859055, + -0.0006584539078176022, + 0.02369162254035473, + 0.025258725509047508, + 0.013307256624102592, + -0.0394270084798336, + 0.0806267261505127, + -0.0434923954308033, + 0.03221459314227104, + 0.1505969762802124, + -0.004116200376302004, + -0.004128452390432358, + -0.11071233451366425, + 0.025004800409078598, + 0.32911527156829834, + -0.015578706748783588, + -0.029277071356773376, + -0.05786833167076111, + -0.019272975623607635, + -0.06218598037958145, + -0.04371056705713272, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458, + 0.014426682144403458 + ], + [ + -0.0074113765731453896, + -0.04827000945806503, + -0.04416186362504959, + 0.49550265073776245, + 0.007433642167598009, + -0.02042396552860737, + 0.02989591844379902, + -0.07436604797840118, + -0.08093081414699554, + 0.00971920881420374, + -0.0027277797926217318, + 0.0026977621018886566, + -0.002009623683989048, + -0.06129171699285507, + -0.027320580556988716, + -0.13519635796546936, + -0.25635164976119995, + 0.0005257180309854448, + 0.4186958968639374, + -0.23261302709579468, + -6.390088674379513e-05, + -0.01122222375124693, + -0.001637016306631267, + 0.008275333791971207, + -0.005104802083224058, + 0.08131929486989975, + -0.007869145832955837, + -4.463096558993129e-14, + 0.014619625173509121, + 0.004656142555177212, + 4.325390740023138e-12, + -5.229368674197765e-11, + -0.011240926571190357, + -6.776263578034403e-20, + 2.752857078576476e-20, + -8.946785505373547e-21, + 1.1507736681837721e-20, + -4.5255051312776144e-20, + 9.35435290281491e-21, + 0.007468376308679581, + -0.0009207710972987115, + 0.0028824550099670887, + -0.06096353009343147, + 0.004407027270644903, + 0.05555590242147446, + -0.003962780348956585, + 0.0008937349775806069, + 0.0008911791373975575, + -0.12448855489492416, + 0.0258108489215374, + 0.00244180578738451, + 0.0824442207813263, + 0.18111318349838257, + -0.16832329332828522, + -0.3159768283367157, + -0.1517036408185959, + -0.03406890481710434, + 0.13882404565811157, + -0.1312544345855713, + -0.006448329426348209, + -0.006464914418756962, + -0.3258637487888336, + -0.07383846491575241, + 0.0236126147210598, + 0.05998911336064339, + 0.1395532786846161, + -0.08956573903560638, + 0.08317936211824417, + -0.06454986333847046, + -0.046274565160274506, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242, + -0.01589706726372242 + ], + [ + 0.10176406800746918, + -0.009436576627194881, + 0.021634845063090324, + 0.23547038435935974, + -0.01910734362900257, + 0.01491554919630289, + 0.062461093068122864, + -0.3824213743209839, + 0.04946723207831383, + 0.02751065045595169, + -0.008293915539979935, + -0.006610027980059385, + 0.0006741732941009104, + 0.03327750787138939, + -0.0009101527975872159, + -0.012048635631799698, + -0.17655059695243835, + 0.0012954457197338343, + 0.23514029383659363, + -0.22115644812583923, + -0.0023475689813494682, + -0.0027058743871748447, + -0.002740643685683608, + 0.049174606800079346, + 0.03726396709680557, + 0.009235390461981297, + 0.013974259607493877, + 8.729128531115293e-15, + -0.02596222795546055, + -0.008268611505627632, + -7.681234120981983e-12, + 9.286562491217509e-11, + 0.019962163642048836, + -6.776263578034403e-21, + 3.917527381051139e-21, + -1.3400325923359048e-21, + 1.7267395287044453e-21, + -6.786784281450312e-21, + 1.4028627759729893e-21, + -0.014465032145380974, + -0.0107062803581357, + -0.002367513021454215, + 0.09323018789291382, + -0.0038247411139309406, + -0.09137224406003952, + -0.008180082775652409, + -0.0017483910778537393, + -0.0017464312259107828, + 0.18301764130592346, + -0.08081263303756714, + -0.007283939514309168, + 0.07893630862236023, + -0.19719775021076202, + 0.2550489008426666, + 0.32607707381248474, + 0.5374373197555542, + -0.029441235587000847, + -0.12461291998624802, + -0.10446316748857498, + -0.00987103022634983, + -0.009862994775176048, + 0.10356499254703522, + -0.006391823757439852, + 0.14977559447288513, + 0.029187001287937164, + -0.02276194468140602, + 0.04533550888299942, + -0.015044674277305603, + -0.0018706604605540633, + -0.011239442974328995, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617, + 0.02823077142238617 + ], + [ + -0.04506399855017662, + 0.03038850612938404, + 0.0020648434292525053, + 0.02786046266555786, + -0.000928212481085211, + -0.0008068872266449034, + -0.00737648643553257, + -0.4897036850452423, + 0.040247365832328796, + 0.09447654336690903, + -0.01996106281876564, + -0.012190662324428558, + -0.0018593635177239776, + -0.17787441611289978, + 0.047910988330841064, + 0.011223283596336842, + 0.015932239592075348, + -0.001274626119993627, + -0.000817913212813437, + 0.009598360396921635, + -9.267040695704054e-06, + 0.006525716744363308, + 0.0024713838938623667, + -0.0030612426344305277, + 0.029655691236257553, + -0.027643702924251556, + -0.0015799634857103229, + -1.4496737144042982e-13, + 0.0029353746213018894, + 0.0009348755702376366, + 8.684615526721728e-13, + -1.0499693128829257e-11, + -0.0022569880820810795, + -1.6263032587282567e-19, + 6.776263578034403e-20, + -2.4775713707188285e-20, + 3.143286327701505e-20, + -1.2357582043174702e-19, + 2.5544474688979755e-20, + 0.003443011548370123, + -0.007518699858337641, + 0.003950301092118025, + 0.00021615397417917848, + 0.0009268360445275903, + -0.013073478825390339, + -0.001114154583774507, + 0.0003434676618780941, + 0.0003439660358708352, + -0.1322735846042633, + 0.04348628968000412, + 0.016786782070994377, + 0.04570361226797104, + 0.18634061515331268, + -0.08855321258306503, + -0.18088792264461517, + -0.103233203291893, + 0.08989360928535461, + -0.044718433171510696, + -0.0665007159113884, + -0.0029652516823261976, + -0.002960651181638241, + 0.43151727318763733, + 0.33415094017982483, + -0.1269713193178177, + -0.2537976801395416, + -0.19726304709911346, + -0.3843437433242798, + -0.14548063278198242, + -0.01827973499894142, + 0.005911238957196474, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724, + -0.0031918634194880724 + ], + [ + 0.040609728544950485, + -0.03234683349728584, + 0.04908207803964615, + -0.11254030466079712, + -0.011657299473881721, + 0.02945661172270775, + -0.015172725543379784, + -0.31915634870529175, + -0.398799329996109, + -0.041314080357551575, + -0.0007517593912780285, + -0.00552559457719326, + -4.568629083223641e-05, + 0.050150152295827866, + -0.06950990110635757, + 0.029937254264950752, + 0.0325583890080452, + 0.002386413747444749, + 0.027389144524931908, + 0.2303663045167923, + -0.0036727949045598507, + -0.011763280257582664, + -0.0034077365417033434, + -0.05020524933934212, + -0.12611375749111176, + -0.01375003345310688, + 0.006551542319357395, + 8.326672684688674e-15, + -0.012171857059001923, + -0.0038765650242567062, + -3.601188791613197e-12, + 4.353813143342933e-11, + 0.009358853101730347, + -1.1926223897340549e-18, + 4.336808689942018e-19, + -1.4230153513872246e-19, + 1.8931186371133613e-19, + -7.449919468897588e-19, + 1.5399290591654695e-19, + 0.003855436807498336, + -0.0051204790361225605, + 0.0019210731843486428, + 0.08128803223371506, + -0.005856159143149853, + -0.09923920035362244, + 0.01653493195772171, + -0.0012609872501343489, + -0.0012514309491962194, + -0.05643944814801216, + -0.08319973945617676, + -0.01153705082833767, + 0.053934287279844284, + 0.0032596802338957787, + -0.013138988055288792, + -0.0967867374420166, + 0.19696864485740662, + -0.041061460971832275, + -0.08266796916723251, + 0.1131858229637146, + -0.0026678701397031546, + -0.002669032895937562, + -0.16714683175086975, + -0.05023930221796036, + -0.6379468441009521, + 0.07106528431177139, + 0.2620592415332794, + 0.053320564329624176, + 0.12381576001644135, + 0.03620174899697304, + 0.05982212722301483, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904, + 0.013235416263341904 + ], + [ + -0.012287391349673271, + -0.005543853156268597, + -0.011578517965972424, + 0.11926445364952087, + 0.009624007157981396, + -0.011686744168400764, + 0.010136974044144154, + -0.01572299376130104, + 0.008071083575487137, + 0.004003717564046383, + -0.0002654132549650967, + -0.0016813650727272034, + -0.0006258104695007205, + -0.03411512076854706, + -0.011371688917279243, + -0.01359561551362276, + 0.047572534531354904, + -0.0005555413081310689, + -0.019891005009412766, + -0.09237892925739288, + -0.017581794410943985, + 0.0029136824887245893, + -0.013698473572731018, + -0.49930453300476074, + -0.6951835751533508, + -0.4582253694534302, + -0.00208265264518559, + 2.2877533201182132e-14, + 0.003869489533826709, + 0.0012323783012107015, + 1.1446399383885364e-12, + -1.3841039425699364e-11, + -0.002975222421810031, + -2.0816681711721685e-17, + 8.673617379884035e-18, + -3.577867169202165e-18, + 4.7027269231558755e-18, + -1.8517834292873514e-17, + 3.828271284234217e-18, + -0.003225581953302026, + 0.004789822734892368, + -0.0008988353656604886, + -0.03184366598725319, + 0.002697072457522154, + 0.044907137751579285, + -0.00819107610732317, + 0.0005223167245276272, + 0.0005169166252017021, + 0.06333019584417343, + 0.03566708788275719, + 0.002108166227117181, + -0.020302344113588333, + -0.051287148147821426, + 0.024393508210778236, + 0.07628825306892395, + -0.05903024598956108, + 0.022380834445357323, + 0.026154043152928352, + -0.028623171150684357, + -0.0007457230240106583, + -0.0007493605953641236, + 0.0033974123653024435, + 0.024677041918039322, + 0.06041742488741875, + -0.015699097886681557, + -0.008956809528172016, + -0.01775992102921009, + 0.003968429751694202, + 0.01060386560857296, + 0.019731389358639717, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348, + -0.00420759990811348 + ], + [ + 0.1167735680937767, + -0.0026284425985068083, + 0.04159048944711685, + -0.07246333360671997, + -0.015176157467067242, + 0.039662230759859085, + 0.009262660518288612, + 0.28717902302742004, + 0.2135930210351944, + 0.13745661079883575, + 0.008129176683723927, + 0.004760713316500187, + 0.0013682297430932522, + 0.19716985523700714, + 0.008519189432263374, + 0.04918414726853371, + -0.033001482486724854, + -0.0001666541793383658, + 0.020796354860067368, + -0.07587983459234238, + -0.002329477109014988, + 0.013439112342894077, + -0.0015599658945575356, + -0.07100134342908859, + -0.08064577728509903, + -0.10335596650838852, + 0.01071583665907383, + -8.99419427824455e-14, + -0.01990857534110546, + -0.006340600084513426, + -5.890204990421921e-12, + 7.121198075665802e-11, + 0.015307559631764889, + -3.903127820947816e-18, + 1.5178830414797062e-18, + -5.692061405548898e-19, + 7.064254780100865e-19, + -2.77678575933641e-18, + 5.738648217647885e-19, + 0.0055471742525696754, + -0.02459133416414261, + 0.0020830321591347456, + 0.13443273305892944, + -0.009510020725429058, + -0.17681492865085602, + 0.012332836166024208, + -0.0022072079591453075, + -0.00218558288179338, + -0.3447871506214142, + -0.18820308148860931, + -0.02358957752585411, + 0.17140012979507446, + 0.13179215788841248, + -0.1008901596069336, + -0.4748777151107788, + 0.42549586296081543, + -0.02377188578248024, + -0.13308212161064148, + -0.03911377489566803, + -0.0012310842284932733, + -0.0012206579558551311, + 0.149983748793602, + -0.07828334718942642, + -0.0018978137522935867, + 0.009343639016151428, + -0.11209892481565475, + 0.10382715612649918, + -0.10003463178873062, + -0.0032108461018651724, + -0.044905878603458405, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878, + 0.021648157387971878 + ], + [ + -0.01844550482928753, + 0.005479452200233936, + -0.03967628628015518, + -0.552947998046875, + -0.0024142316542565823, + 0.04137395694851875, + -0.009799732826650143, + -0.47835397720336914, + 0.1521158665418625, + -0.14796440303325653, + -0.0038104390259832144, + -0.00394020788371563, + -0.0010882854694500566, + -0.04244919493794441, + 0.026018723845481873, + -0.09508567303419113, + -0.12757815420627594, + -0.0040132938884198666, + 0.13439129292964935, + -0.02426212653517723, + 0.0013005913933739066, + 0.0031434609554708004, + 0.0005479370593093336, + -0.019948076456785202, + -0.04411860182881355, + -0.05941275134682655, + -0.012827462516725063, + -1.798006188380441e-13, + 0.023832151666283607, + 0.007590204477310181, + 7.050943162667522e-12, + -8.524653205554955e-11, + -0.018324369564652443, + -1.3877787807814457e-17, + 4.336808689942018e-18, + -2.168404344971009e-18, + 3.0086610286472748e-18, + -1.1834744339037084e-17, + 2.445384118723165e-18, + 0.02920808084309101, + -0.021522052586078644, + 0.005974951200187206, + -0.056029871106147766, + -0.002312740543857217, + -0.022361189126968384, + 0.04598185420036316, + 2.0135636077611707e-05, + 3.396855026949197e-05, + -0.13547568023204803, + -0.037273161113262177, + 0.022378336638212204, + 0.0425889678299427, + 0.02865704335272312, + -0.0055973767302930355, + -0.12785425782203674, + -0.09686237573623657, + -0.09161006659269333, + 0.2264249324798584, + 0.031016835942864418, + -0.009458070620894432, + -0.00947069562971592, + -0.06768447160720825, + -0.21477989852428436, + 0.27574923634529114, + 0.26245132088661194, + -0.06717866659164429, + 0.1818358302116394, + -0.02238762192428112, + 0.055283598601818085, + 0.03290458768606186, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325, + -0.025914570316672325 + ], + [ + 0.012080715037882328, + 0.04073513299226761, + 0.04983486607670784, + 0.4658098518848419, + 0.02782825566828251, + -0.011662838980555534, + -0.016657283529639244, + -0.32424479722976685, + 0.07656975835561752, + -0.07280672341585159, + 0.0005104782758280635, + -0.0051629189401865005, + 0.0007412295672111213, + 0.056147489696741104, + 0.007504662033170462, + 0.09457750618457794, + 0.3897091746330261, + 0.003637465415522456, + -0.4621225893497467, + -0.0653737410902977, + -0.0012348557356745005, + -0.0019212424522265792, + 0.003851342247799039, + 0.03187619522213936, + 0.06473001092672348, + 0.027584116905927658, + 0.0010352153331041336, + -1.0230705171920818e-13, + -0.00192329625133425, + -0.0006125335930846632, + -5.692113447253178e-13, + 6.879552483241014e-12, + 0.0014788161497563124, + -2.0816681711721685e-17, + 1.734723475976807e-18, + -1.734723475976807e-18, + 3.5236570605778894e-18, + -1.3660947373317356e-17, + 2.8223137802513287e-18, + -0.009487379342317581, + -0.009472480975091457, + 0.002114440780133009, + -0.007429132238030434, + 0.003979143686592579, + 0.025821082293987274, + -0.03114243969321251, + 0.0009214507299475372, + 0.0009191093849949539, + -0.17727792263031006, + 0.03653412684798241, + 0.009314650669693947, + 0.008144482970237732, + -0.06325652450323105, + -0.0664687380194664, + -0.12361874431371689, + -0.0782686322927475, + -0.025080524384975433, + -0.07591167092323303, + -0.028051625937223434, + 0.005946031771600246, + 0.005955663044005632, + 0.013403943739831448, + -0.17109397053718567, + 0.07960136234760284, + 0.2275892198085785, + -0.030918190255761147, + 0.2754473090171814, + -0.007204721681773663, + 0.1532014012336731, + 0.11608260869979858, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364, + 0.0020913498010486364 + ], + [ + -0.047645676881074905, + 0.17276465892791748, + 0.07746634632349014, + -0.051822636276483536, + -0.03549569100141525, + 0.0587063692510128, + -0.005498088896274567, + 0.008229561150074005, + -0.6670181751251221, + 0.3084486126899719, + -0.002715971088036895, + -0.0016404690686613321, + 0.0006188888219185174, + -0.02043282613158226, + -0.0007555370684713125, + 0.021392084658145905, + -0.006417661905288696, + -0.00088893884094432, + 0.0935995876789093, + -0.031476184725761414, + 0.0005724814254790545, + 0.0072382730431854725, + 0.0048411074094474316, + -0.022989118471741676, + -0.0020049537997692823, + 0.009525508619844913, + -0.005384870804846287, + 5.70168912084057e-13, + 0.010004209354519844, + 0.0031862047035247087, + 2.959577027894511e-12, + -3.578476404086928e-11, + -0.007692159153521061, + 8.326672684688674e-17, + 2.7755575615628914e-17, + 5.551115123125783e-17, + 0.0, + 0.0, + 3.469446951953614e-18, + 0.041818659752607346, + 0.05930309370160103, + 0.00662911357358098, + 0.10112112760543823, + -0.0128538953140378, + -0.11150791496038437, + 0.09418197721242905, + -0.00025929667754098773, + -0.0002441807009745389, + -0.18996472656726837, + 0.06393135339021683, + -0.04239130765199661, + -0.11669979244470596, + -0.055167414247989655, + -0.09337533265352249, + 0.045928001403808594, + -0.12353184074163437, + 0.0584062822163105, + -0.24683724343776703, + -0.0557977631688118, + -0.002948769833892584, + -0.0029346132650971413, + 0.17098788917064667, + -0.071454718708992, + 0.32386457920074463, + 0.10113218426704407, + -0.08114884793758392, + 0.15626491606235504, + -0.03730236738920212, + -0.14009691774845123, + -0.10268169641494751, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551, + -0.010878362692892551 + ], + [ + -0.033873300999403, + 0.19029556214809418, + 0.03467681258916855, + 0.07782568782567978, + -0.044701121747493744, + 0.15786677598953247, + 0.012276668101549149, + 0.13741880655288696, + -0.16488489508628845, + -0.7493964433670044, + 0.007957477122545242, + 0.0032768698874861, + -0.0010665848385542631, + -0.15725193917751312, + 0.11266429722309113, + -0.06694763153791428, + -0.04494728893041611, + -0.008662734180688858, + 0.01658461056649685, + -0.05643733963370323, + -0.00034037508885376155, + 0.01169061753898859, + 0.00504410732537508, + -0.013390506617724895, + -0.003448465373367071, + -0.0016607752768322825, + -0.007830173708498478, + 3.558877368151081e-13, + 0.014547290280461311, + 0.004633115138858557, + 4.304195888593654e-12, + -5.20347515076125e-11, + -0.011185298673808575, + 0.0, + -2.220446049250313e-16, + -2.220446049250313e-16, + -5.551115123125783e-17, + 0.0, + -4.163336342344337e-17, + 0.06209897622466087, + 0.02824016660451889, + 0.00463215634226799, + 0.11065955460071564, + -0.02082686312496662, + -0.20597393810749054, + 0.14684753119945526, + -0.0022583995014429092, + -0.002238065702840686, + -0.059630610048770905, + -0.1283375322818756, + -0.04164886474609375, + -0.018642572686076164, + 0.14037290215492249, + -0.04845475032925606, + 0.04994533583521843, + 0.07166752219200134, + 0.2152702808380127, + -0.09829138964414597, + 0.0564960241317749, + 0.008290666155517101, + 0.008273114450275898, + 0.04742055386304855, + -0.015530113130807877, + 0.08105827122926712, + -0.024949397891759872, + -0.0621139332652092, + -0.10306607931852341, + -0.033070582896471024, + 0.1555694192647934, + 0.16475728154182434, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695, + -0.015818411484360695 + ], + [ + 0.1121126338839531, + 0.23375889658927917, + -0.1020946130156517, + 0.031075183302164078, + 0.0011069843312725425, + -0.02522912807762623, + -0.015518740750849247, + 0.022285528481006622, + -0.32279425859451294, + 0.12317586690187454, + -3.298587034805678e-05, + 0.0016059948829934, + -0.0006927672657184303, + 0.27338701486587524, + 0.11061421036720276, + -0.0766015276312828, + 0.16975444555282593, + -0.0021187912207096815, + -0.09255698323249817, + -0.033366862684488297, + -0.0008122186991386116, + 0.0276947021484375, + 0.010092206299304962, + -0.021570680662989616, + 0.03672517463564873, + 0.004275837447494268, + -0.008004420436918736, + -1.2906342661267445e-13, + 0.014871124178171158, + 0.004736258182674646, + 4.399758335438264e-12, + -5.319331086717227e-11, + -0.011434286832809448, + 0.0, + 6.938893903907228e-18, + -5.551115123125783e-17, + 0.0, + -2.7755575615628914e-17, + 0.0, + 0.020064612850546837, + -0.06078912690281868, + 0.01200706884264946, + -0.1698678731918335, + 0.008462205529212952, + 0.1258004903793335, + 0.009692210704088211, + 0.0004691976064350456, + 0.0004581625107675791, + 0.11805619299411774, + -0.14785297214984894, + 0.04754847288131714, + 0.20698225498199463, + 0.3382856845855713, + 0.10625170171260834, + 0.05739914998412132, + 0.19593870639801025, + 0.09487774968147278, + 0.5709800124168396, + 0.0004136768984608352, + -0.015619758516550064, + -0.015640566125512123, + 0.06981039047241211, + 0.008991931565105915, + 0.06876417249441147, + -0.03453211858868599, + -0.05004438757896423, + -0.02236815169453621, + -0.035210687667131424, + 0.07307492941617966, + 0.0969347432255745, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658, + -0.016170542687177658 + ], + [ + 0.018567614257335663, + 0.10768173635005951, + 0.07341843843460083, + -0.009212891571223736, + 0.031797144562006, + -0.05221332982182503, + 0.00015834419173188508, + -0.06625976413488388, + 0.13269560039043427, + 0.17144878208637238, + 0.00010825808567460626, + -0.0002686611842364073, + 0.0005228295340202749, + 0.15944735705852509, + 0.30007514357566833, + 0.10756872594356537, + -0.384613573551178, + -0.003976711072027683, + -0.0760759487748146, + 0.034052662551403046, + -0.0012598742032423615, + 0.016978001222014427, + -0.0039202747866511345, + -0.03740031644701958, + -0.0035226198378950357, + -0.027665434405207634, + 0.016846362501382828, + -6.97688434803112e-13, + -0.031298402696847916, + -0.009968116879463196, + -9.260203714944737e-12, + 1.1195235710692586e-10, + 0.02406509965658188, + -1.3877787807814457e-16, + 2.7755575615628914e-17, + 4.163336342344337e-17, + 0.0, + 1.1102230246251565e-16, + 2.220446049250313e-16, + -0.04520653188228607, + -0.012820528820157051, + 0.008623017929494381, + 0.03807428106665611, + 0.014918360859155655, + 0.014380100183188915, + -0.09480663388967514, + 0.0009924446931108832, + 0.0009499170701019466, + 0.11206461489200592, + 0.15448834002017975, + 0.14106814563274384, + -0.09189696609973907, + 0.5102449059486389, + -0.1036849394440651, + 0.18671394884586334, + -0.08115223050117493, + 0.24071809649467468, + -0.2797788679599762, + 0.03952433168888092, + 0.0051323119550943375, + 0.005154304672032595, + -0.09536104649305344, + -0.06162793189287186, + -0.060468245297670364, + 0.09259721636772156, + -0.012047593481838703, + 0.15063539147377014, + 0.013508842326700687, + 0.17489172518253326, + 0.13786926865577698, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776, + 0.034033212810754776 + ], + [ + 0.03659084811806679, + -0.2517549991607666, + -0.04501069709658623, + -0.016086647287011147, + 0.0637841671705246, + -0.21012866497039795, + -0.0019806730560958385, + 0.08114200085401535, + -0.21591703593730927, + 0.014751730486750603, + 0.00196798425167799, + 0.0018948279321193695, + 0.00023023184621706605, + -0.019089417532086372, + -0.10204039514064789, + 0.021536584943532944, + -0.2302330881357193, + 0.010065769776701927, + 0.06416916847229004, + 0.07692210376262665, + 0.0013369250809773803, + -0.0207839235663414, + -0.007917415350675583, + 0.022801263257861137, + -0.009719976224005222, + 0.018322179093956947, + 0.016218509525060654, + -4.345482307321902e-16, + -0.03013257123529911, + -0.009596826508641243, + -8.91513252110343e-12, + 1.0778286596568876e-10, + 0.023168690502643585, + -2.220446049250313e-16, + 1.1102230246251565e-16, + 5.551115123125783e-17, + 0.0, + -6.938893903907228e-18, + 2.7755575615628914e-17, + -0.09143969416618347, + 0.030676964670419693, + -0.014072194695472717, + -0.03485393896698952, + 0.021812278777360916, + 0.24803470075130463, + -0.19588451087474823, + 0.0018950409721583128, + 0.001864751917310059, + -0.10601191967725754, + 0.008327426388859749, + -0.1255192905664444, + 0.030301278457045555, + -0.23092561960220337, + -0.059996724128723145, + -0.0986955538392067, + 0.02599073201417923, + -0.025986429303884506, + 0.014730622060596943, + 0.04511631652712822, + -0.0004838672757614404, + -0.0004658239195123315, + 0.12259037792682648, + 0.06726976484060287, + 0.13288411498069763, + -0.050662316381931305, + -0.05764126777648926, + -0.002991290530189872, + -0.029742583632469177, + 0.4935746192932129, + 0.4927484691143036, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938, + 0.03276551142334938 + ], + [ + -0.12252353131771088, + 0.002171361818909645, + 0.03746799752116203, + 0.08480792492628098, + -0.08584237098693848, + 0.36971917748451233, + 0.043043140321969986, + 0.06109200790524483, + 0.11560715734958649, + 0.336473673582077, + -0.0007580149103887379, + 0.0015652759466320276, + 0.0004208627506159246, + -0.04105248674750328, + -0.008616267703473568, + -0.08610835671424866, + -0.02855069190263748, + 1.678122134762816e-05, + 0.015859071165323257, + 0.00884278118610382, + -0.006794125307351351, + -0.01974048651754856, + -0.004769384860992432, + 0.00658695213496685, + 0.00449094595387578, + 0.010124473832547665, + -0.033137377351522446, + 6.939449015419541e-13, + 0.061564818024635315, + 0.01960751973092556, + 1.821456874218086e-11, + -2.202143467577855e-10, + -0.04733675345778465, + 2.7755575615628914e-17, + -5.551115123125783e-17, + 6.938893903907228e-17, + -1.3877787807814457e-17, + -1.214306433183765e-17, + 0.0, + 0.1562410444021225, + -0.0255763977766037, + 0.011446455493569374, + 0.03651977702975273, + -0.040420424193143845, + -0.3608991503715515, + 0.3325541615486145, + -0.0028924623038619757, + -0.002818604465574026, + 0.01453423872590065, + -0.04139115288853645, + 0.09345021843910217, + -0.019490839913487434, + -0.007438148837536573, + 0.14016093313694, + 0.05417788401246071, + -0.11573690176010132, + -0.25438791513442993, + 0.05606212839484215, + -0.0010782849276438355, + -0.014821265824139118, + -0.01480986550450325, + -0.036578040570020676, + 0.08455071598291397, + -0.053898099809885025, + -0.05547014996409416, + 0.03183260187506676, + 0.020828265696763992, + 0.016249436885118484, + 0.34100547432899475, + 0.32648319005966187, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447, + -0.06694426387548447 + ], + [ + -0.0019459235481917858, + 0.04749950394034386, + 0.07475372403860092, + -0.11582078784704208, + -0.008334175683557987, + -0.03284887969493866, + -0.01967306062579155, + 0.02714802511036396, + 0.09593494981527328, + 0.028080442920327187, + 0.0023331630509346724, + 0.0008350192219950259, + 0.0006662193336524069, + -0.004446700215339661, + -0.04207743704319, + -0.05167990177869797, + 0.650263786315918, + -0.016838116571307182, + 0.5590289235115051, + 0.007066413294523954, + 0.00050199223915115, + 0.012199868448078632, + 0.007459837477654219, + 0.025308923795819283, + -0.0014305911026895046, + 0.01347116008400917, + 0.01206381618976593, + -6.287461870591038e-13, + -0.022412750869989395, + -0.00713812792673707, + -6.6314731483885225e-12, + 8.016878827454832e-11, + 0.017233015969395638, + 1.3877787807814457e-17, + 6.938893903907228e-18, + 0.0, + 0.0, + -1.1102230246251565e-16, + 1.1102230246251565e-16, + -0.04285592958331108, + 0.0006576204905286431, + 0.002395393792539835, + 0.020909946411848068, + 0.012827349826693535, + 0.028940457850694656, + -0.08596430718898773, + 0.00038366595981642604, + 0.0003542017366271466, + 0.1423312872648239, + 0.08750532567501068, + 0.08506686985492706, + -0.06912636756896973, + 0.1393153816461563, + -0.0715179443359375, + -0.00300388946197927, + 0.0071945819072425365, + 0.11854808777570724, + -0.165487140417099, + -0.18894405663013458, + -0.015626830980181694, + -0.015630489215254784, + -0.03063170425593853, + -0.013961059041321278, + 0.0016403717454522848, + 0.027179889380931854, + -0.015475877560675144, + 0.047447048127651215, + 0.012618623673915863, + 0.17045649886131287, + 0.1737082451581955, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627, + 0.02437114156782627 + ], + [ + 0.001850979751907289, + -0.16534170508384705, + -0.00940654892474413, + 0.017297610640525818, + -0.012139360420405865, + 0.07893053442239761, + 0.016939369961619377, + -0.07486433535814285, + 0.11080145835876465, + 0.10536069422960281, + -0.0012941377935931087, + -0.001456167665310204, + -3.1048642995301634e-05, + 0.2768915891647339, + -0.20840243995189667, + -0.12080598622560501, + -0.026134595274925232, + 0.013565810397267342, + 0.006702611222863197, + 0.04311727359890938, + 0.008270415477454662, + -0.04076939448714256, + -0.015298541635274887, + 0.02124129980802536, + 0.0012026607291772962, + 0.015099829062819481, + -0.02326783537864685, + -2.058303614355106e-14, + 0.04322844743728638, + 0.013767640106379986, + 1.2789437911497892e-11, + -1.5462656155484922e-10, + -0.03323804959654808, + -7.112366251504909e-17, + 1.0408340855860843e-17, + 4.163336342344337e-17, + -1.3877787807814457e-17, + -8.326672684688674e-17, + 6.938893903907228e-17, + 0.02460361085832119, + 0.1384773999452591, + -0.009262765757739544, + 0.08498556166887283, + -0.01172439381480217, + -0.09836269170045853, + 0.0790649875998497, + 0.00014299539907369763, + 9.435058746021241e-05, + -0.05239494517445564, + 0.012109453789889812, + -0.10025616735219955, + -0.029842671006917953, + -0.2126791775226593, + 0.0412287674844265, + -0.0369078665971756, + -0.05136755108833313, + 0.797477126121521, + 0.10827409476041794, + 0.03834003955125809, + 0.010656156577169895, + 0.010662115179002285, + -0.05603175610303879, + 0.01226990856230259, + -0.014617644250392914, + 0.0042006708681583405, + -0.014751032926142216, + 0.0014096983941271901, + 0.05417357757687569, + -0.0036114982794970274, + 0.009028423577547073, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034, + -0.04700568690896034 + ] + ], + "encoded_input_dim": 140, + "mean": [ + -0.009798124432563782, + 0.05137821286916733, + 0.06156044080853462, + 0.7411994338035583, + 0.10057990252971649, + 0.032337773591279984, + 0.08269896358251572, + 0.0006070384988561273, + 0.017993291839957237, + 0.07773640006780624, + -1.3279248476028442, + -0.4710727035999298, + -0.34919124841690063, + -0.01877165213227272, + -0.004584182519465685, + -0.07146554440259933, + 0.7753992676734924, + 0.00030092295492067933, + 0.1989821046590805, + 0.20095480978488922, + -0.0006528590456582606, + 0.005675423424690962, + 0.0032932220492511988, + 0.03019251488149166, + -0.022124359384179115, + -0.005906751379370689, + -0.19088762998580933, + 2.0268739486274667e-13, + -0.7879482507705688, + 0.7490668892860413, + 3.896688874349685e-11, + -4.711071532881306e-10, + 0.6058492660522461, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.5270761251449585, + 0.055747177451848984, + 0.007260707672685385, + 0.5600981712341309, + 0.012609828263521194, + 0.32732704281806946, + 0.11341477185487747, + 0.004221478942781687, + 0.004223855212330818, + 0.015596347860991955, + -0.039205241948366165, + 0.01386436726897955, + 0.0015060961013659835, + -0.01569386012852192, + 0.3965786397457123, + 0.03663915768265724, + -1.6119304895401, + 0.011917737312614918, + 2.071789503097534, + 0.1707562953233719, + 0.0195054542273283, + 0.019512077793478966, + -0.003088352968916297, + 0.3037094473838806, + -0.015200544148683548, + 0.10040655732154846, + -0.029771316796541214, + 0.1618170440196991, + -0.09639980643987656, + 0.006122836843132973, + -0.008711924776434898, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123, + 0.8567857146263123 + ], + "method": "train_split_pca_standardized", + "output_dim": 32, + "raw_input_dim": 70, + "scale": [ + 7.789134502410889, + 7.01206636428833, + 4.7303786277771, + 2.2875277996063232, + 1.554949164390564, + 1.471292495727539, + 1.0825825929641724, + 0.7508261203765869, + 0.7088286876678467, + 0.5949841737747192, + 0.5041103363037109, + 0.49374377727508545, + 0.44354283809661865, + 0.40395745635032654, + 0.3867206573486328, + 0.31073999404907227, + 0.2897527813911438, + 0.2838159203529358, + 0.25809574127197266, + 0.24356654286384583, + 0.22696244716644287, + 0.22063840925693512, + 0.20554044842720032, + 0.187404602766037, + 0.1526772826910019, + 0.15083047747612, + 0.14299766719341278, + 0.1399637907743454, + 0.13358847796916962, + 0.1275552362203598, + 0.10995583981275558, + 0.09880612045526505 + ], + "state_vectorization": "right_zero_pad_plus_validity_mask" + }, + "trainable_parameters": 99880992, + "training_losses": [ + 1.2670934200286865, + 0.520173192024231, + 0.7244012951850891, + 0.861823558807373, + 0.7761167883872986, + 0.5301089882850647, + 0.4306434988975525, + 0.45650750398635864, + 0.35851266980171204, + 0.4026678800582886, + 0.4129319190979004, + 0.6854279637336731, + 0.3484633266925812, + 0.3139687776565552, + 0.2927967309951782, + 0.5260645151138306, + 0.5035329461097717, + 0.6627745032310486, + 0.41149747371673584, + 0.3834378123283386, + 0.4975249469280243, + 0.4423580765724182, + 0.664903998374939, + 0.3410729765892029, + 1.3096120357513428, + 0.3992670178413391, + 0.36001187562942505, + 0.25928154587745667, + 0.6334152221679688, + 0.6129553318023682, + 0.19171898066997528, + 0.710078775882721, + 0.6200621128082275, + 0.16084426641464233, + 0.5018320083618164, + 0.5223394632339478, + 0.299043208360672, + 0.29888027906417847, + 0.4238058924674988, + 0.36750948429107666, + 0.24303552508354187, + 0.612612247467041, + 0.41708633303642273, + 0.25748211145401, + 0.36406829953193665, + 0.40093785524368286, + 0.47529518604278564, + 0.42138320207595825, + 0.2496960163116455, + 0.4959806799888611, + 0.2552245259284973, + 0.5733389854431152, + 0.2546464800834656, + 0.4206914007663727, + 0.4887707233428955, + 0.29276132583618164, + 0.7429209351539612, + 0.5279264450073242, + 0.24885448813438416, + 0.32807785272598267, + 0.394634485244751, + 0.3797350525856018, + 0.2740103602409363, + 0.38176020979881287, + 0.4427236020565033, + 0.3789430856704712, + 0.3272364139556885, + 0.41950494050979614, + 0.22516848146915436, + 0.39976516366004944, + 0.33316707611083984, + 0.33998414874076843, + 0.34504932165145874, + 0.4812081456184387, + 0.29870060086250305, + 0.5794742107391357, + 0.28276145458221436, + 0.3031766414642334, + 0.32576778531074524, + 0.33173513412475586, + 0.2203328013420105, + 0.6655448079109192, + 0.29436469078063965, + 0.607647180557251, + 0.34380489587783813, + 0.3207570016384125, + 0.3339797854423523, + 0.2903817892074585, + 0.3986304998397827, + 0.31699126958847046, + 0.568588137626648, + 0.1940593123435974, + 0.38046443462371826, + 0.2499038279056549, + 0.469354510307312, + 0.33222997188568115, + 0.25105851888656616, + 0.22275134921073914, + 0.34570205211639404, + 0.4159901738166809, + 0.26984483003616333, + 0.5275513529777527, + 0.23384574055671692, + 0.1701635718345642, + 0.27213776111602783, + 0.301037460565567, + 0.21793141961097717, + 0.2538028955459595, + 0.22929374873638153, + 0.2983025014400482, + 0.6730409860610962, + 0.3952772617340088, + 0.39058592915534973, + 0.5902495384216309, + 0.33566272258758545, + 0.4916001856327057, + 0.4371281564235687, + 0.2534796893596649, + 0.3303525149822235, + 0.3797616958618164, + 0.33454248309135437, + 0.31265950202941895, + 0.5155372023582458, + 0.2868397831916809, + 0.24798643589019775, + 0.44390884041786194, + 0.36383920907974243, + 0.3922899067401886, + 0.12758809328079224, + 0.5399858951568604, + 0.20339733362197876, + 0.2622903883457184, + 0.3509126901626587, + 0.40092164278030396, + 0.37563592195510864, + 0.24361900985240936, + 0.29325976967811584, + 0.2316695600748062, + 0.20290926098823547, + 0.2634566128253937, + 0.22562099993228912, + 0.5719513297080994, + 0.6145126819610596, + 0.1648721694946289, + 0.27501511573791504, + 0.7239519357681274, + 0.4555286765098572, + 0.29711103439331055, + 0.6057884693145752, + 0.41811561584472656, + 0.3496985137462616, + 0.36383548378944397, + 0.3423958718776703, + 0.275688111782074, + 0.20015490055084229, + 0.31119781732559204, + 0.19978772103786469, + 0.3615823984146118, + 0.26866990327835083, + 0.2021978795528412, + 0.2203378677368164, + 0.2886856496334076, + 0.48434948921203613, + 0.44793403148651123, + 0.21035225689411163, + 0.15287461876869202, + 0.2884337306022644, + 0.26935410499572754, + 0.4560839831829071, + 0.25576457381248474, + 0.22952687740325928, + 0.4979139268398285, + 0.20432327687740326, + 0.274995356798172, + 0.2464819848537445, + 0.1309635043144226, + 0.19330325722694397, + 0.25416234135627747, + 0.2687529921531677, + 0.3750022053718567, + 0.32664382457733154, + 0.4079568386077881, + 0.3517562448978424, + 0.3962664008140564, + 0.3373618721961975, + 0.3934488892555237, + 0.3961174786090851, + 0.40817081928253174, + 0.1740289330482483, + 0.1982683390378952, + 0.27873602509498596, + 0.28878647089004517, + 0.22901654243469238, + 0.27682286500930786, + 0.24326203763484955, + 0.438418447971344, + 0.3745250701904297, + 0.20524221658706665, + 0.20268017053604126, + 0.17775800824165344, + 0.3321373760700226, + 0.27339112758636475, + 0.19515401124954224, + 0.25285181403160095, + 0.20698505640029907, + 0.22119322419166565, + 0.24854783713817596, + 0.31379565596580505, + 0.38473644852638245, + 0.2652498781681061, + 0.3487132787704468, + 0.3282003700733185, + 0.3395761549472809, + 0.3286164402961731, + 0.2916727662086487, + 0.2860058546066284, + 0.6056867837905884, + 0.3148811459541321, + 0.31247377395629883, + 0.25620463490486145, + 0.10684531927108765, + 0.2773495018482208, + 0.20881590247154236, + 0.318372905254364, + 0.16619619727134705, + 0.22899144887924194, + 0.37092524766921997, + 0.21344444155693054, + 0.2891526222229004, + 0.3251163065433502, + 0.14318744838237762, + 0.2454303354024887, + 0.20798680186271667, + 0.37544751167297363, + 0.3386860489845276, + 0.2168900966644287, + 0.24647042155265808, + 0.22530460357666016, + 0.55317223072052, + 0.3348284661769867, + 0.2299366295337677, + 0.22704604268074036, + 0.15976262092590332, + 0.3657275438308716, + 0.301100492477417, + 0.27172279357910156, + 0.39577239751815796, + 0.2821155786514282, + 0.3055368661880493, + 0.4601201117038727, + 0.2636391520500183, + 0.2966904044151306, + 0.16443191468715668, + 0.2830047011375427, + 0.13874369859695435, + 0.35457950830459595, + 0.47381478548049927, + 0.3051682114601135, + 0.2500254511833191, + 0.27289193868637085, + 0.42709434032440186, + 0.2923024892807007, + 0.40615567564964294, + 0.24479693174362183, + 0.2835773825645447, + 0.47966641187667847, + 0.23792603611946106, + 0.3256709575653076, + 0.16970887780189514, + 0.40243878960609436, + 0.2489173859357834, + 0.10545828193426132, + 0.3023035526275635, + 0.2598850429058075, + 0.27712777256965637, + 0.24168430268764496, + 0.2154207080602646, + 0.4322353005409241, + 0.2460225373506546, + 0.27129024267196655, + 0.22092068195343018, + 0.2877585291862488, + 0.2643328607082367, + 0.32570916414260864, + 0.2385791838169098, + 0.2224685549736023, + 0.25966131687164307, + 0.2513476610183716, + 0.18968503177165985, + 0.2829245328903198, + 0.10489155352115631, + 0.16993632912635803, + 0.3219834566116333, + 0.2191532701253891, + 0.22033053636550903, + 0.5239449739456177, + 0.2400340586900711, + 0.4066777527332306, + 0.6314682960510254, + 0.14256766438484192, + 0.24379567801952362, + 0.21083225309848785, + 0.27450305223464966, + 0.16862894594669342, + 0.25522762537002563, + 0.17643406987190247, + 0.2853166460990906, + 0.19303429126739502, + 0.1498509645462036, + 0.2388489842414856, + 0.13838893175125122, + 0.3098934590816498, + 0.16839700937271118, + 0.19421610236167908, + 0.2782820761203766, + 0.19295667111873627, + 0.1713719367980957, + 0.23983444273471832, + 0.2697269916534424, + 0.20350424945354462, + 0.3498741686344147, + 0.19276602566242218, + 0.270733118057251, + 0.19215452671051025, + 0.40897291898727417, + 0.38916996121406555, + 0.19744794070720673, + 0.24047031998634338, + 0.19416087865829468, + 0.3935443162918091, + 0.16298457980155945, + 0.2834876775741577, + 0.23586173355579376, + 0.31882673501968384, + 0.45068466663360596, + 0.24416187405586243, + 0.09592004120349884, + 0.2236321121454239, + 0.10516108572483063, + 0.3579285740852356, + 0.23994438350200653, + 0.46382197737693787, + 0.1475619077682495, + 0.3595869839191437, + 0.28384408354759216, + 0.30293434858322144, + 0.6376585960388184, + 0.2255164533853531, + 0.24408194422721863, + 0.23628249764442444, + 0.5372308492660522, + 0.4284501075744629, + 0.40179550647735596, + 0.24014386534690857, + 0.39531630277633667, + 0.2201327383518219, + 0.3889012932777405, + 0.2307373732328415, + 0.2189430594444275, + 0.17904630303382874, + 0.2852848470211029, + 0.40027886629104614, + 0.2425360381603241, + 0.3651343882083893, + 0.3037956953048706, + 0.2533831000328064, + 0.22858458757400513, + 0.2801596522331238, + 0.37010061740875244, + 0.4146811366081238, + 0.2693101167678833, + 0.24089115858078003, + 0.39205583930015564, + 0.16415822505950928, + 0.39432868361473083, + 0.1821451336145401, + 0.2439577430486679, + 0.49070268869400024, + 0.11082972586154938, + 0.3274472951889038, + 0.1537664234638214, + 0.30448412895202637, + 0.15689930319786072, + 0.22857320308685303, + 0.23461970686912537, + 0.3019418716430664, + 0.22748392820358276, + 0.2825089693069458, + 0.4108922779560089, + 0.27620166540145874, + 0.3593670725822449, + 0.22929362952709198, + 0.23038360476493835, + 0.23947745561599731, + 0.32028928399086, + 0.30193817615509033, + 0.13791000843048096, + 0.2035796195268631, + 0.31117159128189087, + 0.22811517119407654, + 0.2508699297904968, + 0.18596968054771423, + 0.22576327621936798, + 0.4325881600379944, + 0.2616579830646515, + 0.26820775866508484, + 0.17247715592384338, + 0.22288507223129272, + 0.23100832104682922, + 0.19087806344032288, + 0.3790215849876404, + 0.07714812457561493, + 0.3867262601852417, + 0.19780871272087097, + 0.2775113880634308, + 0.2952142655849457, + 0.20193631947040558, + 0.348556250333786, + 0.22606950998306274, + 0.14702913165092468, + 0.2356618046760559, + 0.27577343583106995, + 0.5125761032104492, + 0.3113219738006592, + 0.2866706848144531, + 0.23882699012756348, + 0.20804502069950104, + 0.2589229941368103, + 0.17633165419101715, + 0.159375861287117, + 0.262842059135437, + 0.44615399837493896, + 0.17780473828315735, + 0.38006824254989624, + 0.21641960740089417, + 0.2825886309146881, + 0.15013882517814636, + 0.16695326566696167, + 0.396871417760849, + 0.23568537831306458, + 0.21893943846225739, + 0.18594177067279816, + 0.2638377845287323, + 0.2521240711212158, + 0.25198128819465637, + 0.31235769391059875, + 0.18386417627334595, + 0.1553190052509308, + 0.3119312524795532, + 0.3172594904899597, + 0.18746733665466309, + 0.3651098906993866, + 0.2763952910900116, + 0.22347986698150635, + 0.2801726460456848, + 0.1940222978591919, + 0.17966881394386292, + 0.3242173194885254, + 0.3902655243873596, + 0.1260964721441269, + 0.20105013251304626, + 0.29161137342453003, + 0.2225721925497055, + 0.4434126317501068, + 0.3282968997955322, + 0.21134734153747559, + 0.21729369461536407, + 0.13720929622650146, + 0.16564837098121643, + 0.31327298283576965, + 0.41819286346435547, + 0.18672671914100647, + 0.21137253940105438, + 0.14712458848953247, + 0.21897098422050476, + 0.22316516935825348, + 0.4073867201805115, + 0.10646860301494598, + 0.44536805152893066, + 0.16712838411331177, + 0.4767214059829712, + 0.1493215411901474, + 0.3237265348434448, + 0.1778651773929596, + 0.262431800365448, + 0.266690194606781, + 0.14386790990829468, + 0.3189900517463684, + 0.21905982494354248, + 0.23397356271743774, + 0.34257835149765015, + 0.3180058002471924, + 0.25138092041015625, + 0.2638319134712219, + 0.24236300587654114, + 0.396126925945282, + 0.2353307455778122, + 0.3038586378097534, + 0.19343073666095734, + 0.19883306324481964, + 0.1689959168434143, + 0.21377794444561005, + 0.29091376066207886, + 0.5113687515258789, + 0.233821302652359, + 0.1341589391231537, + 0.22365066409111023, + 0.15416742861270905, + 0.31835252046585083, + 0.3517591059207916, + 0.23958678543567657, + 0.1999340057373047, + 0.11711592972278595, + 0.19730359315872192, + 0.21812878549098969, + 0.20371761918067932, + 0.2151525914669037, + 0.145045667886734, + 0.08876442164182663, + 0.19354403018951416, + 0.17816388607025146, + 0.18095172941684723, + 0.47056761384010315, + 0.2452552318572998, + 0.3520406186580658, + 0.2291138470172882, + 0.3345881700515747, + 0.21177589893341064, + 0.48866531252861023, + 0.12807492911815643, + 0.24406594038009644, + 0.2097213715314865, + 0.1418631374835968, + 0.4719617962837219, + 0.3409278988838196, + 0.2030293196439743, + 0.20222297310829163, + 0.22152797877788544, + 0.13075682520866394, + 0.133987694978714, + 0.1896134316921234, + 0.23093819618225098, + 0.29388323426246643, + 0.19772465527057648, + 0.15041810274124146, + 0.15583808720111847, + 0.15511074662208557, + 0.22531737387180328, + 0.3073890209197998, + 0.09111789613962173, + 0.10771863907575607, + 0.33277708292007446, + 0.22551080584526062, + 0.15457716584205627, + 0.11167806386947632, + 0.09169816970825195, + 0.33823081851005554, + 0.12185919284820557, + 0.15758132934570312, + 0.2522837817668915, + 0.14190256595611572, + 0.3252077102661133, + 0.17894557118415833, + 0.19304993748664856, + 0.3364601135253906, + 0.22601985931396484, + 0.20737433433532715, + 0.24604256451129913, + 0.32785916328430176, + 0.21481814980506897, + 0.26799827814102173, + 0.08924897015094757, + 0.09798911958932877, + 0.19050481915473938, + 0.1598503440618515, + 0.2078505903482437, + 0.17986422777175903, + 0.1977917104959488, + 0.29288023710250854, + 0.1792115867137909, + 0.3763973116874695, + 0.12488805502653122, + 0.262966513633728, + 0.15390446782112122, + 0.09855378419160843, + 0.18241801857948303, + 0.23252154886722565, + 0.10793440043926239, + 0.23421677947044373, + 0.15533040463924408, + 0.14738579094409943, + 0.10901661217212677, + 0.14478415250778198, + 0.14525394141674042, + 0.20525765419006348, + 0.10758066177368164, + 0.22589297592639923, + 0.19200947880744934, + 0.22606313228607178, + 0.3065832853317261, + 0.42117416858673096, + 0.1235380470752716, + 0.22657634317874908, + 0.23922185599803925, + 0.1680079698562622, + 0.48163071274757385, + 0.11303481459617615, + 0.3770957589149475, + 0.1795358806848526, + 0.14808517694473267, + 0.1150122731924057, + 0.2576565146446228, + 0.14073732495307922, + 0.14235520362854004, + 0.11155528575181961, + 0.34474441409111023, + 0.2502405345439911, + 0.2850838899612427, + 0.6362336874008179, + 0.2258119285106659, + 0.1375127136707306, + 0.1717459261417389, + 0.20378772914409637, + 0.23298367857933044, + 0.3804195821285248, + 0.1434040665626526, + 0.29813310503959656, + 0.4157029390335083, + 0.2737812399864197, + 0.16524633765220642, + 0.11330154538154602, + 0.1639053374528885, + 0.502902626991272, + 0.2416035532951355, + 0.18261657655239105, + 0.1672433614730835, + 0.10033687949180603, + 0.21934862434864044, + 0.19469057023525238, + 0.28395208716392517, + 0.1119440495967865, + 0.3063829243183136, + 0.10302992165088654, + 0.10622456669807434, + 0.08810842037200928, + 0.24985364079475403, + 0.21231737732887268, + 0.22800371050834656, + 0.24446457624435425, + 0.39735233783721924, + 0.3067137897014618, + 0.3310328722000122, + 0.17779719829559326, + 0.2414708137512207, + 0.3013531565666199, + 0.1486828476190567, + 0.27601996064186096, + 0.18786919116973877, + 0.16628214716911316, + 0.21742261946201324, + 0.35659104585647583, + 0.18619492650032043, + 0.2009846270084381, + 0.1103598102927208, + 0.32711905241012573, + 0.14170172810554504, + 0.18596908450126648, + 0.3343438506126404, + 0.10239613056182861, + 0.21883785724639893, + 0.19776314496994019, + 0.2693973779678345, + 0.3702816069126129, + 0.11796867847442627, + 0.2669399380683899, + 0.1200399100780487, + 0.16689065098762512, + 0.26215457916259766, + 0.1788056492805481, + 0.34713101387023926, + 0.2727925777435303, + 0.2632560133934021, + 0.09795044362545013, + 0.342117577791214, + 0.23456934094429016, + 0.13351798057556152, + 0.2704470157623291, + 0.1878122091293335, + 0.2066091150045395, + 0.47464224696159363, + 0.31298691034317017, + 0.24251550436019897, + 0.16247700154781342, + 0.1450728178024292, + 0.2641066312789917, + 0.41669929027557373, + 0.17079336941242218, + 0.21372736990451813, + 0.1850937306880951, + 0.27225708961486816, + 0.31098151206970215, + 0.16710370779037476, + 0.16066573560237885, + 0.12723207473754883, + 0.12210853397846222, + 0.1689375340938568, + 0.1499425172805786, + 0.13503241539001465, + 0.40215808153152466, + 0.1338193416595459, + 0.11485059559345245, + 0.24684475362300873, + 0.3915119767189026, + 0.14127016067504883, + 0.2729734778404236, + 0.25083982944488525, + 0.14546085894107819, + 0.11596724390983582, + 0.14712709188461304, + 0.13114240765571594, + 0.21915394067764282, + 0.08847025036811829, + 0.14501267671585083, + 0.11186926811933517, + 0.14503860473632812, + 0.08631609380245209, + 0.18856382369995117, + 0.2609281539916992, + 0.1320677101612091, + 0.22093632817268372, + 0.1494542956352234, + 0.0905969887971878, + 0.27131837606430054, + 0.24503737688064575, + 0.1527915596961975, + 0.16780507564544678, + 0.17164862155914307, + 0.2613556981086731, + 0.14132606983184814, + 0.21898344159126282, + 0.2434302270412445, + 0.11553347855806351, + 0.14331170916557312, + 0.4039428234100342, + 0.21406343579292297, + 0.19984030723571777, + 0.17814017832279205, + 0.28703805804252625, + 0.09001202881336212, + 0.22007697820663452, + 0.3724384903907776, + 0.29113417863845825, + 0.276885986328125, + 0.18146803975105286, + 0.31804972887039185, + 0.30176472663879395, + 0.07828282564878464, + 0.3588334619998932, + 0.18707579374313354, + 0.36713331937789917, + 0.30880093574523926, + 0.23829413950443268, + 0.0988704040646553, + 0.1581830233335495, + 0.1435871720314026, + 0.11794374138116837, + 0.3435283601284027, + 0.22464828193187714, + 0.4415470063686371, + 0.21951861679553986, + 0.29405611753463745, + 0.13395428657531738, + 0.2572108507156372, + 0.3648509085178375, + 0.15511652827262878, + 0.3490931987762451, + 0.17508850991725922, + 0.1450505554676056, + 0.1603071689605713, + 0.32316985726356506, + 0.28143590688705444, + 0.34883564710617065, + 0.13743484020233154, + 0.19571363925933838, + 0.13046854734420776, + 0.23246365785598755, + 0.11327656358480453, + 0.19311337172985077, + 0.15547071397304535, + 0.0879286527633667, + 0.18275772035121918, + 0.17939093708992004, + 0.11558932065963745, + 0.24045304954051971, + 0.16510915756225586, + 0.1629955768585205, + 0.07577505707740784, + 0.1858544647693634, + 0.13407523930072784, + 0.1383848488330841, + 0.12741516530513763, + 0.16124531626701355, + 0.11192686855792999, + 0.06082127243280411, + 0.07279767096042633, + 0.16992217302322388, + 0.11484377086162567, + 0.21387578547000885, + 0.25110575556755066, + 0.1253996193408966, + 0.1890573501586914, + 0.20941588282585144, + 0.2913479804992676, + 0.23933401703834534, + 0.22659024596214294, + 0.21418294310569763, + 0.2008405327796936, + 0.05576483905315399, + 0.26032012701034546, + 0.12186792492866516, + 0.1520593911409378, + 0.12751571834087372, + 0.08619560301303864, + 0.16408312320709229, + 0.17927584052085876, + 0.16834348440170288, + 0.14021724462509155, + 0.3407657742500305, + 0.2957145869731903, + 0.17177024483680725, + 0.19211596250534058, + 0.09277237206697464, + 0.12157561630010605, + 0.08309458196163177, + 0.19019517302513123, + 0.19451265037059784, + 0.14988099038600922, + 0.12189947068691254, + 0.4660620093345642, + 0.13362637162208557, + 0.21456703543663025, + 0.31568530201911926, + 0.275255411863327, + 0.4309881925582886, + 0.16732768714427948, + 0.44987642765045166, + 0.14444419741630554, + 0.21093252301216125, + 0.1854296326637268, + 0.1300470232963562, + 0.09191755950450897, + 0.06213744729757309, + 0.36542192101478577, + 0.18027843534946442, + 0.20548871159553528, + 0.09548819065093994, + 0.23229734599590302, + 0.09639842063188553, + 0.09742659330368042, + 0.17540580034255981, + 0.2755020260810852, + 0.0773056223988533, + 0.19911223649978638, + 0.1588195264339447, + 0.23173117637634277, + 0.3724217414855957, + 0.15338103473186493, + 0.15752044320106506, + 0.3266931176185608, + 0.3171297311782837, + 0.1166800931096077, + 0.2818506956100464, + 0.40506434440612793, + 0.1754155457019806, + 0.3402497172355652, + 0.14989280700683594, + 0.13265933096408844, + 0.2791270315647125, + 0.2929890751838684, + 0.1613774597644806, + 0.16842177510261536, + 0.1908762902021408, + 0.16750450432300568, + 0.1830136775970459, + 0.2006067931652069, + 0.4710058569908142, + 0.24005649983882904, + 0.08854508399963379, + 0.14156652987003326, + 0.15379288792610168, + 0.11747033894062042, + 0.11634866893291473, + 0.09495986998081207, + 0.3599845767021179, + 0.10936915874481201, + 0.43910300731658936, + 0.18433384597301483, + 0.09567141532897949, + 0.13626272976398468, + 0.16601791977882385, + 0.1462615430355072, + 0.371281236410141, + 0.1320621371269226, + 0.31802213191986084, + 0.33053284883499146, + 0.1761794239282608, + 0.25286173820495605, + 0.1541672796010971, + 0.15974920988082886, + 0.2788334786891937, + 0.10106616467237473, + 0.1446804702281952, + 0.17398248612880707, + 0.13958211243152618, + 0.23584890365600586, + 0.14718449115753174, + 0.08882854133844376, + 0.2439582645893097, + 0.12457384914159775, + 0.13352926075458527, + 0.09588070213794708, + 0.1345539093017578, + 0.19492314755916595, + 0.2948419153690338, + 0.16565477848052979, + 0.3860889673233032, + 0.13909552991390228, + 0.1304548978805542, + 0.2389429211616516, + 0.156545490026474, + 0.11388562619686127, + 0.12366503477096558, + 0.13146597146987915, + 0.16488566994667053, + 0.10576125979423523, + 0.2605326771736145, + 0.40653735399246216, + 0.18015331029891968, + 0.1715795397758484, + 0.2715689539909363, + 0.1888289451599121, + 0.20025664567947388, + 0.2384958267211914, + 0.15658965706825256, + 0.130298912525177, + 0.10221226513385773, + 0.2556404769420624, + 0.19469141960144043, + 0.1975296139717102, + 0.16930799186229706, + 0.10444515943527222, + 0.3691498041152954, + 0.16394342482089996, + 0.20755067467689514, + 0.1805219203233719, + 0.10124044120311737, + 0.21493317186832428, + 0.1610446572303772, + 0.42842090129852295, + 0.2948428690433502, + 0.1408509463071823, + 0.07686226069927216, + 0.29882746934890747, + 0.15844601392745972, + 0.11257848143577576, + 0.2739264965057373, + 0.14706707000732422, + 0.13876353204250336, + 0.21597999334335327, + 0.12796595692634583, + 0.08883091807365417, + 0.2462306022644043, + 0.05590568855404854, + 0.18760165572166443, + 0.16563135385513306, + 0.1331680566072464, + 0.19478023052215576, + 0.15581999719142914, + 0.20062139630317688, + 0.213889479637146, + 0.26641952991485596, + 0.06761877238750458, + 0.2758985161781311, + 0.3622976243495941, + 0.1223989725112915, + 0.2842875123023987, + 0.19262737035751343, + 0.33439770340919495, + 0.2276298552751541, + 0.26281988620758057, + 0.3673902153968811, + 0.18455937504768372, + 0.13687554001808167, + 0.10172797739505768, + 0.34208154678344727, + 0.4418579041957855, + 0.21014583110809326, + 0.28194814920425415, + 0.26305463910102844, + 0.18892055749893188 + ] +} \ No newline at end of file diff --git a/workspace/outputs/external_vla/smolvla_cil_aligned_metrics.json b/workspace/outputs/external_vla/smolvla_cil_aligned_metrics.json new file mode 100644 index 0000000000000000000000000000000000000000..03069999968b4747db9d7678b37c53da2715c958 --- /dev/null +++ b/workspace/outputs/external_vla/smolvla_cil_aligned_metrics.json @@ -0,0 +1,47 @@ +{ + "candidate_oracle_success_rate": 0.4185714285714286, + "evaluation_protocol": "nearest_executed_same_state_candidate", + "final_train_loss": 0.18892055749893188, + "mean_selected_regret": 0.13656934786188815, + "model_family": "smolvla", + "normalized_action_mse_to_selected": 0.3394091720879078, + "num_eval_groups": 700, + "num_train_groups": 2800, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "selected_reward_mean": 1.0044795715748662, + "selected_success_rate": 0.3434343434343434 + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "selected_reward_mean": 0.23799376099456832, + "selected_success_rate": 0.0 + }, + "PickCube-v1": { + "num_groups": 193, + "selected_reward_mean": 0.8670999890897867, + "selected_success_rate": 0.2538860103626943 + }, + "PullCube-v1": { + "num_groups": 113, + "selected_reward_mean": 1.3038161870388858, + "selected_success_rate": 0.584070796460177 + }, + "PushCube-v1": { + "num_groups": 106, + "selected_reward_mean": 1.2718729558137227, + "selected_success_rate": 0.5943396226415094 + }, + "StackCube-v1": { + "num_groups": 92, + "selected_reward_mean": 0.9698587565318398, + "selected_success_rate": 0.32608695652173914 + } + }, + "selected_reward_mean": 0.9446513404989881, + "selected_success_rate": 0.3457142857142857, + "top1_action_selection": 0.5228571428571429, + "training_mode": "expert_only_bc", + "training_steps": 1000 +} \ No newline at end of file diff --git a/workspace/outputs/external_vla/smolvla_cil_balanced_manifest.json b/workspace/outputs/external_vla/smolvla_cil_balanced_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..9fdaa84bf6b029e752e249425723a864ae7204fb --- /dev/null +++ b/workspace/outputs/external_vla/smolvla_cil_balanced_manifest.json @@ -0,0 +1,5781 @@ +{ + "action_normalizer": { + "canonical_action_dim": 8, + "mean": [ + -0.022289231419563293, + 0.10234535485506058, + -0.2564755380153656, + -0.0621388703584671, + -0.2134821116924286, + 0.4443473219871521, + -0.8069988489151001, + -0.03291666507720947 + ], + "scale": [ + 0.761264979839325, + 0.5265141129493713, + 1.0559866428375244, + 1.266664981842041, + 1.4507075548171997, + 1.6286382675170898, + 1.5712451934814453, + 0.4069325625896454 + ], + "vectorization": "right_zero_pad_to_canonical_dimension" + }, + "checkpoint": "/scratch/knguy52/dovla/models/smolvla_base-c83c316", + "claim_scope": "Expert-only SmolVLA fine-tuning evaluated by nearest executed same-state CIL candidate; this is measured candidate selection, not online policy rollout.", + "config": { + "action_dim": 8, + "action_horizon": 4, + "batch_size": 4, + "device": "cuda", + "export_dir": "/scratch/knguy52/dovla/experiments/external_vla_export_full_balanced", + "image_size": 512, + "learning_rate": 0.0001, + "log_every": 25, + "max_eval_groups": 600, + "seed": 0, + "state_dim": 32, + "steps": 1000, + "val_fraction": 0.2, + "vlm_metadata": "/scratch/knguy52/dovla/models/SmolVLM2-500M-Video-Instruct-metadata-7b375e1b73b11138ff12fe22c8f2822d8fe03467" + }, + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "export": "/scratch/knguy52/dovla/experiments/external_vla_export_full_balanced", + "num_train_groups": 2400, + "num_validation_groups": 600, + "schema_version": "smolvla-cil-candidate-baseline/v0", + "state_projector": { + "components": [ + [ + -0.0010280748829245567, + 0.001114234677515924, + 0.0002530088531784713, + -0.0018986471695825458, + 0.004930918104946613, + 0.0016624636482447386, + -0.003447691909968853, + -0.02657553367316723, + -0.000881311425473541, + 0.0038963600527495146, + 0.3078743815422058, + 0.7225583791732788, + 0.6142717003822327, + -0.0036706302780658007, + -4.40303556388244e-05, + -0.00770162558183074, + -0.0006289654993452132, + 5.949442493147217e-05, + -0.005272062961012125, + 0.007389334496110678, + -1.7888891306938604e-05, + -7.127784192562103e-05, + 3.789695620071143e-05, + 0.0019783428870141506, + -0.00016471449634991586, + -0.0022035359870642424, + -0.003976669628173113, + 1.0238068598737889e-14, + 0.007388113532215357, + 0.002353010466322303, + 2.1858613009906325e-12, + -2.642692314980266e-11, + -0.005680666770786047, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.005042672157287598, + 0.002719542011618614, + 0.0004435222945176065, + -0.024973230436444283, + 0.0007244178559631109, + 0.01788479834794998, + 0.006465005688369274, + 0.0002147667109966278, + 0.0002148925996152684, + 0.0019707384053617716, + -0.0047528501600027084, + 0.0008492294582538307, + 0.0032441867515444756, + -0.0011831293813884258, + -0.017403224483132362, + 0.003817452583462, + 0.014573774300515652, + 0.0017283076886087656, + -0.022558031603693962, + -0.009021050296723843, + 0.00010442818165756762, + 0.00010459201439516619, + -0.004635993856936693, + -0.011579579673707485, + 0.0040922509506344795, + -0.03894598409533501, + 0.0018640975467860699, + 0.0036412456538528204, + -0.004740027245134115, + -0.002168494975194335, + -0.00021428662876132876, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132, + -0.008033676072955132 + ], + [ + -0.00023950121249072254, + 0.0007218352984637022, + 0.000516425643581897, + -0.004160810727626085, + 0.004002545028924942, + 0.0006693319301120937, + 0.00503411702811718, + -0.020927220582962036, + -0.0010587458964437246, + 0.00967797078192234, + 0.49594438076019287, + 0.42449358105659485, + -0.7550681829452515, + -0.002928372472524643, + -0.0003961192269343883, + -0.005209828726947308, + -1.3756212865700945e-05, + -0.00012499646982178092, + -0.0075187720358371735, + 0.005669177509844303, + -5.47394847671967e-05, + 0.00020218735153321177, + 3.5627275792649016e-05, + -0.00032729943632148206, + -0.0011448953300714493, + 0.0004933697055093944, + -0.0029383888468146324, + 9.286267840303598e-15, + 0.005459134466946125, + 0.001738657709211111, + 1.6151498541394638e-12, + -1.9527058789381435e-11, + -0.004197488538920879, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0036859307438135147, + 0.0020867404527962208, + 0.00033891425118781626, + -0.01843813620507717, + 0.0005546026513911784, + 0.013343761675059795, + 0.004767151549458504, + 0.00016215174400713295, + 0.00016224150022026151, + 0.0013698241673409939, + -0.00292178220115602, + 0.0005661408649757504, + 0.0018554892158135772, + -0.001559150405228138, + -0.010922623798251152, + 0.0032586059533059597, + 0.013534110970795155, + 0.0015000712592154741, + -0.019714102149009705, + -0.013198732398450375, + 0.00019252228958066553, + 0.00019310027710162103, + -0.006771927699446678, + -0.011496088467538357, + 0.005600310396403074, + -0.02614709734916687, + 0.004670621827244759, + -0.0076894150115549564, + -0.007078781723976135, + -0.0014115256490185857, + 0.00036609964445233345, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703, + -0.005936145782470703 + ], + [ + -1.4796180948906112e-05, + 0.0022400966845452785, + -0.001448249095119536, + 0.001391921890899539, + -0.004436619114130735, + -0.0024415988009423018, + 0.005413648206740618, + -0.00840661209076643, + -0.0031643291004002094, + 0.00340832257643342, + 0.8099850416183472, + -0.5371779203414917, + 0.22858726978302002, + -0.001920076203532517, + 0.0009136219741776586, + 0.004073368851095438, + -0.005051843356341124, + -1.2760812751366757e-05, + 0.0008174824761226773, + 0.00013013517309445888, + 2.775384928099811e-05, + 0.0004953724564984441, + 0.00020560908888000995, + 0.004855617415159941, + -9.013406815938652e-05, + 0.000669025641400367, + 0.003062437754124403, + -3.1337639833084723e-15, + -0.005689582321792841, + -0.0018120516324415803, + -1.6833305582775182e-12, + 2.035135775180219e-11, + 0.004374678712338209, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.003934942185878754, + -0.0018243602244183421, + -0.0003640907525550574, + 0.019529029726982117, + -0.0006149139953777194, + -0.013414442539215088, + -0.0049631306901574135, + -0.00012971600517630577, + -0.0001297509588766843, + -0.0034623907413333654, + 0.007959306240081787, + -0.0002727838873397559, + -0.006180380936712027, + 0.0028788556810468435, + 0.02014293149113655, + -0.005881957244127989, + -0.013007726520299911, + 0.0012811131309717894, + 0.01713121123611927, + -0.0195700041949749, + 0.0004354238335508853, + 0.0004347131762187928, + 0.0033071269281208515, + 0.0023334193974733353, + 0.0032970618922263384, + -0.014073352329432964, + 0.00770231569185853, + -0.0017846392001956701, + 0.0012102037435397506, + -0.003812620881944895, + 0.0017877767095342278, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318, + 0.00618672976270318 + ], + [ + -0.0037947001401335, + -0.0038635702803730965, + 0.0027845092117786407, + -0.02414398454129696, + 0.09739533811807632, + 0.025616606697440147, + -0.049971308559179306, + -0.0023749324027448893, + 0.023408234119415283, + 0.020464716479182243, + -0.003934499807655811, + -0.08293867111206055, + 0.008723246864974499, + -0.01585150510072708, + 0.0018064111936837435, + -0.14388443529605865, + -0.005040669348090887, + -3.0100032745394856e-05, + -0.045683104544878006, + 0.08598712831735611, + 0.0001998630614252761, + -0.0009281459497287869, + -0.0006059266161173582, + -0.008162932470440865, + 0.0051720030605793, + 0.003320078831166029, + -0.07384809106588364, + 2.6296288280561186e-13, + 0.13719983398914337, + 0.043696217238903046, + 4.05922032098438e-11, + -4.907571016232737e-10, + -0.10549195855855942, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.09083101898431778, + 0.058609526604413986, + 0.007674215827137232, + -0.4597875475883484, + 0.01342045795172453, + 0.3434208035469055, + 0.11752515286207199, + 0.00441402243450284, + 0.004416984971612692, + 0.019604766741394997, + -0.040567588061094284, + 0.014213130809366703, + -0.0019284934969618917, + 0.015097411349415779, + -0.1737981140613556, + 0.02705022320151329, + 0.2915727496147156, + -0.0029761905316263437, + -0.3683585524559021, + -0.019486865028738976, + -0.0033435134682804346, + -0.0033444075379520655, + 0.0028738591354340315, + -0.07695145159959793, + 0.006531573366373777, + 0.0068883090279996395, + 0.004532825201749802, + -0.06643946468830109, + 0.03769463300704956, + 0.004386794287711382, + 0.006609681062400341, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127, + -0.14918816089630127 + ], + [ + 0.00791467446833849, + 0.003985682036727667, + -0.0004050939460285008, + 0.008491788990795612, + -0.0038550025783479214, + -0.0012841252610087395, + -0.00228788866661489, + 8.584093302488327e-05, + 0.008650495670735836, + -0.0077258325181901455, + 0.006387452129274607, + 0.009063459001481533, + -0.003700198605656624, + 0.01114076841622591, + 0.0036860511172562838, + 0.003967558965086937, + -0.003920795861631632, + -4.93777624797076e-05, + 0.007334162015467882, + -0.0006025181501172483, + -9.654298628447577e-05, + 0.0011830689618363976, + 0.00020231169764883816, + 0.0009843921288847923, + -0.0015073881950229406, + 0.00236239074729383, + 0.0030057979747653008, + -1.162439556504774e-14, + -0.005584382917732, + -0.0017785473028197885, + -1.6522061496715423e-12, + 1.997506673956373e-11, + 0.004293791949748993, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.003751779906451702, + -0.0023619653657078743, + -0.00029981593252159655, + 0.018794002011418343, + -0.0005512042553164065, + -0.013819864019751549, + -0.004880515392869711, + -0.00016655572107993066, + -0.0001666341267991811, + -0.0019954622257500887, + 0.0032680926378816366, + -0.0002870632742997259, + -0.0013925025705248117, + 0.03552274405956268, + 0.01672217808663845, + -0.0397314727306366, + 0.002329393522813916, + -0.05490639805793762, + 0.00935447309166193, + 0.04705139249563217, + -0.00014929565077181906, + -0.00015714133041910827, + 0.13963565230369568, + -0.04174959287047386, + 0.005059640854597092, + -0.04633043706417084, + -0.35442686080932617, + -0.016525311395525932, + 0.9164786338806152, + -0.021463900804519653, + 0.006075977813452482, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172, + 0.006072338670492172 + ], + [ + -0.014386637136340141, + 0.006396414712071419, + 0.003717957530170679, + -0.004203966353088617, + -0.0031846712809056044, + -0.0011641064193099737, + 0.010440194047987461, + -0.09845694154500961, + -0.003845322411507368, + 0.010681122541427612, + -0.03648494929075241, + -0.03440658748149872, + -0.00553548289462924, + -0.01087923813611269, + -0.0015591346891596913, + 0.005380363203585148, + -0.0011337928008288145, + 0.00024043908342719078, + -0.008778517134487629, + 0.008847236633300781, + -0.001056957058608532, + 0.0009337731171399355, + 0.0007511723670177162, + 0.012551236897706985, + -0.0077247777953743935, + -0.0015656654722988605, + 0.0024818528909236193, + -5.945785889734188e-15, + -0.0046109422110021114, + -0.00146851921454072, + -1.3642021540194449e-12, + 1.6493115900795274e-11, + 0.0035453212913125753, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.003145394613966346, + -0.0015612478600814939, + -0.00029225111939013004, + 0.015750804916024208, + -0.0004584594280458987, + -0.0113263800740242, + -0.0040098377503454685, + -0.000139672847581096, + -0.00013975742331240326, + -0.0008521872223354876, + 0.0019910684786736965, + -0.0007110952283255756, + -0.0019197200890630484, + -0.010464102029800415, + -0.015956664457917213, + 0.008472451940178871, + -0.04468651860952377, + 0.009897514246404171, + 0.005845936946570873, + -0.03989497572183609, + 0.0005071950727142394, + 0.000508578959852457, + -0.03428563475608826, + -0.6034322381019592, + 0.020370546728372574, + -0.7744196057319641, + 0.008986608125269413, + 0.11538950353860855, + -0.05205681174993515, + 0.010126156732439995, + 0.010425817221403122, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364, + 0.005013839807361364 + ], + [ + 0.007556125055998564, + 0.03984592854976654, + 0.022354433313012123, + -0.09329218417406082, + -0.010381529107689857, + -0.0024959230795502663, + 0.12480800598859787, + 0.027418049052357674, + -0.016409778967499733, + 0.07517413049936295, + -0.012141818180680275, + 0.008556603454053402, + 0.008907399140298367, + -0.0066658989526331425, + -0.00944836251437664, + 0.04287034273147583, + 0.002631794661283493, + 0.0007740055443719029, + -0.19848166406154633, + 0.08405567705631256, + -0.0008975013042800128, + 0.004005889408290386, + 0.0031432239338755608, + 0.05789578706026077, + -0.03045276738703251, + -0.030565133318305016, + 0.010634584352374077, + -3.754774404807551e-14, + -0.019757509231567383, + -0.006292481441050768, + -5.845494227552095e-12, + 7.067164908836077e-11, + 0.015191414393484592, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.013616575859487057, + -0.007540925405919552, + -0.0011531264754012227, + 0.06695423275232315, + -0.002005278365686536, + -0.04763398319482803, + -0.01763809844851494, + -0.0005304445512592793, + -0.0005307534011080861, + -0.007917623966932297, + 0.018953099846839905, + -0.0006231721490621567, + -0.01214867364615202, + -0.025121720507740974, + 0.0992240458726883, + -0.004212856292724609, + 0.04262670874595642, + 0.03483172506093979, + -0.028521951287984848, + -0.4198482632637024, + 0.004458719398826361, + 0.0044617862440645695, + -0.07619882375001907, + -0.3994152843952179, + 0.1046861857175827, + 0.24472522735595703, + 0.1360756903886795, + -0.6621692776679993, + 0.06861243396997452, + 0.06710556149482727, + 0.07155170291662216, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635, + 0.02148389257490635 + ], + [ + -0.014708459377288818, + -0.006586070638149977, + 0.009892758913338184, + -0.0013509257696568966, + -0.01839505508542061, + 0.030169187113642693, + -0.09305942803621292, + -0.023238040506839752, + 0.05773821845650673, + 0.11133338510990143, + 0.012963386252522469, + -0.015660801902413368, + 0.003387412056326866, + -0.00886200275272131, + 0.006895590107887983, + 0.030775409191846848, + 0.029523996636271477, + -0.0002461089170537889, + 0.019574221223592758, + -0.0735294446349144, + 0.0006078148726373911, + -0.00236451905220747, + -0.0009205694077536464, + -0.002137052593752742, + 0.005039753392338753, + -0.00864973571151495, + 0.010771085508167744, + -2.659395971396205e-14, + -0.020011240616440773, + -0.006373302545398474, + -5.920563518613253e-12, + 7.157923559431012e-11, + 0.015386494807898998, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0025933608412742615, + -0.026839621365070343, + -0.0005771320429630578, + 0.06251855939626694, + -0.0019141058437526226, + -0.11536629498004913, + 0.006796301808208227, + -0.0036044581793248653, + -0.0036071562208235264, + 0.15108971297740936, + -0.3338032364845276, + -0.043317027390003204, + 0.3606463372707367, + -0.15647216141223907, + -0.7152306437492371, + 0.2713671326637268, + -0.07989891618490219, + 0.011840131133794785, + 0.04168948531150818, + 0.14603258669376373, + 0.0008687008521519601, + 0.0008720759069547057, + -0.018727989867329597, + -0.08838307112455368, + 0.011301841586828232, + 0.0718151107430458, + -0.04432092607021332, + -0.13845637440681458, + 0.00397460488602519, + 0.009684204123914242, + 0.008226808160543442, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035, + 0.021759795024991035 + ], + [ + 0.010456273332238197, + -0.0824383869767189, + -0.00841754674911499, + 0.03463342413306236, + 0.0027977381832897663, + 0.0013978214701637626, + -0.2738381624221802, + -0.02449934557080269, + -0.02406635507941246, + 0.09191248565912247, + 0.018888898193836212, + 0.005544136743992567, + 0.0032555153593420982, + -0.017274340614676476, + 0.02279696613550186, + 0.03413018956780434, + 0.11579271405935287, + -0.0006415974930860102, + 0.11069507151842117, + -0.2701380252838135, + 0.0019353822572156787, + -0.010010558180510998, + -0.005910251755267382, + -0.08404563367366791, + 0.049459971487522125, + 0.011970638297498226, + 0.0009729058365337551, + -7.97433916588288e-15, + -0.0018076716223731637, + -0.0005757350591011345, + -5.348211565454897e-13, + 6.4659536405664575e-12, + 0.0013898885808885098, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.004617589525878429, + 0.0038157200906425714, + -0.00010804210614878684, + 0.010038213804364204, + -0.000527293246705085, + 0.0077713849022984505, + -0.005840720608830452, + 0.000710563559550792, + 0.0007117840577848256, + -0.043701574206352234, + 0.10077061504125595, + 0.011818986386060715, + -0.09268401563167572, + 0.06952240318059921, + 0.2028818130493164, + -0.04607807472348213, + 0.04522895812988281, + -0.006159414071589708, + -0.06982821226119995, + 0.7105268836021423, + -0.0025620448868721724, + -0.002548061776906252, + -0.07382618635892868, + -0.23878619074821472, + 0.038312483578920364, + 0.0888940617442131, + -0.06999842077493668, + -0.3634599447250366, + -0.07432714104652405, + 0.02742135152220726, + 0.015790799632668495, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397, + 0.0019656233489513397 + ], + [ + -0.004397987388074398, + 0.0029851640574634075, + -0.002028347458690405, + -0.017424680292606354, + 0.0012022919254377484, + -0.0003365486045368016, + -0.014085079543292522, + -0.009281656704843044, + 0.08676590025424957, + -0.0013668563915416598, + -0.0034049449022859335, + 0.0050314427353441715, + -0.0017313080606982112, + 0.0008077099337242544, + 0.0273579154163599, + 0.005730126518756151, + 0.01564035378396511, + -0.0006228529964573681, + 0.02526048757135868, + -0.03445657715201378, + 0.00048673799028620124, + 0.005078007932752371, + -0.001407829811796546, + -0.07944292575120926, + 0.01738428883254528, + 0.07021088898181915, + 0.000998938805423677, + -2.185501366198158e-14, + -0.0018558213487267494, + -0.0005910561885684729, + -5.490668141504285e-13, + 6.638182625112732e-12, + 0.0014269243692979217, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.0005693331477232277, + -0.002383193466812372, + 2.343275809835177e-05, + 0.00573559058830142, + -0.00012738177611026913, + -0.008106101304292679, + -0.0010678370017558336, + -0.00020435666374396533, + -0.0002046548470389098, + 0.0014168191701173782, + -0.02147461101412773, + -0.0008121090941131115, + 0.014580641873180866, + 0.07270699739456177, + -0.03978307172656059, + -0.05207498371601105, + -0.012353588826954365, + 0.03396625071763992, + 0.0040479497984051704, + 0.1386994570493698, + -0.000238489214098081, + -0.00024447840405628085, + 0.43373918533325195, + 0.011909032240509987, + 0.18862000107765198, + -0.03279433771967888, + 0.814996600151062, + 0.0338263213634491, + 0.23651373386383057, + -0.0030971956439316273, + 0.012840073555707932, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429, + 0.002017980208620429 + ], + [ + 0.007266815286129713, + -0.003814256051555276, + 0.0037145805545151234, + -0.030825955793261528, + -0.005973668769001961, + 0.0019775452092289925, + -0.019960517063736916, + -0.015005498193204403, + 0.02322080172598362, + -0.011076133698225021, + 0.00032144333817996085, + -0.0011176110710948706, + 0.0029070349410176277, + 0.005122920963913202, + -0.020040463656187057, + 0.004319331608712673, + 0.006692976225167513, + -0.00336152920499444, + -0.02734098583459854, + 0.07085678726434708, + 0.005276753567159176, + 0.01976954936981201, + -0.007348137907683849, + -0.8053187727928162, + 0.26151418685913086, + 0.4956466853618622, + 0.0028352048248052597, + -3.351313384382425e-15, + -0.0052674985490739346, + -0.0016776234842836857, + -1.5584520436304206e-12, + 1.8841584542572676e-11, + 0.004050142131745815, + -4.314083075427408e-31, + 5.113344002348111e-31, + -1.2883568548289176e-31, + -2.2726387117780873e-31, + 9.73057649140343e-32, + -1.017919450643995e-31, + -0.0023036531638354063, + -0.002584627363830805, + -0.0002067718596663326, + 0.020533204078674316, + -0.0008305380470119417, + -0.01979050599038601, + -0.0018780590035021305, + -0.000289033108856529, + -0.0002886982401832938, + 0.0053349328227341175, + 0.00984495971351862, + -0.0030460457783192396, + 0.030084557831287384, + -0.039381243288517, + 0.008007553406059742, + 0.01368126180022955, + 0.031498268246650696, + 0.00024129817029461265, + -0.0059729283675551414, + -0.11543174833059311, + 0.0026595937088131905, + 0.002660616533830762, + -0.05586841702461243, + 0.0017387685365974903, + 0.025370528921484947, + -0.008420472964644432, + -0.0739971473813057, + -0.014381105080246925, + -0.013996092602610588, + 0.004877569153904915, + 0.006003410089761019, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926, + 0.005727764684706926 + ], + [ + -0.0017153314547613263, + -0.016961058601737022, + 0.008403489366173744, + -0.06486941874027252, + -0.034094031900167465, + 0.0743817687034607, + -0.3977285921573639, + 0.0359266959130764, + -0.03488779813051224, + -0.06976228207349777, + -0.002410278655588627, + -0.0013040666235610843, + -0.0023233916144818068, + 0.0037199584767222404, + -0.004009155556559563, + -0.009497874416410923, + -0.01032274216413498, + -3.5929570003645495e-05, + -0.059604398906230927, + -0.07416344434022903, + -4.4679104576061945e-06, + -0.009274033829569817, + -0.0025953014846891165, + 0.017274096608161926, + 0.041948381811380386, + -0.10109259933233261, + -0.004054776858538389, + 7.626341778140672e-14, + 0.007533223833888769, + 0.002399221295490861, + 2.2287939722975825e-12, + -2.6945976699943586e-11, + -0.005792246200144291, + -2.0194839173657902e-28, + 3.3427980858740375e-28, + -8.438654644327113e-29, + -1.4886051800553374e-28, + 6.373625613412724e-29, + -6.667474976530714e-29, + 0.020244065672159195, + 0.004733780864626169, + 0.005293750669807196, + 0.032925840467214584, + -0.0033567652571946383, + -0.06728311628103256, + 0.042116641998291016, + 7.543577521573752e-05, + 8.259507012553513e-05, + -0.02315627597272396, + 0.62307208776474, + -0.02826754003763199, + 0.5922620892524719, + -0.03898298740386963, + 0.02252883091568947, + 0.016057470813393593, + 0.0505104623734951, + 0.006711246911436319, + -0.046270258724689484, + -0.13867759704589844, + 0.0018351644976064563, + 0.0018378758104518056, + -0.09404541552066803, + 0.06643898040056229, + 0.06254632025957108, + -0.04914622753858566, + 0.05786987766623497, + 0.02954871952533722, + 0.04544161260128021, + 0.009715390391647816, + 0.010592184960842133, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934, + -0.008191466331481934 + ], + [ + 0.002521706046536565, + -0.06871653348207474, + -0.01146730501204729, + -0.0459195114672184, + -0.0023233245592564344, + -0.010599241591989994, + -0.5749586820602417, + -0.05591445043683052, + -0.03685964643955231, + 0.022853005677461624, + -0.008405963890254498, + -0.005469200666993856, + -0.0031967104878276587, + -0.0020169676281511784, + -0.012324091978371143, + 0.006525442935526371, + -0.00997657235711813, + -0.00046255896450020373, + -0.09844349324703217, + -0.19465041160583496, + 0.002079432364553213, + -0.02369944006204605, + -0.00693570077419281, + -0.07809489965438843, + 0.2122359573841095, + -0.30423203110694885, + 0.0052118925377726555, + -2.6553859480173136e-14, + -0.009683242999017239, + -0.0030839876271784306, + -2.8649025426291663e-12, + 3.463648812207509e-11, + 0.007445364259183407, + -2.928251680180396e-27, + 4.335579535094681e-27, + -1.0929667841837119e-27, + -1.928287282639166e-27, + 8.256197819185699e-28, + -8.636840539324337e-28, + -0.008295337669551373, + -0.00616114679723978, + -0.002765923971310258, + 0.01894715242087841, + -0.0013363147154450417, + -0.013246504589915276, + -0.011265297420322895, + -0.0006658185156993568, + -0.000665056926663965, + 0.009361250326037407, + -0.3389365077018738, + 0.012851203791797161, + -0.30744436383247375, + 0.004057930316776037, + 0.019683213904500008, + -0.016170188784599304, + 0.02111336588859558, + 0.0315936878323555, + -0.027330433949828148, + -0.24840031564235687, + 0.004297592211514711, + 0.004303179681301117, + -0.29487866163253784, + 0.1667269766330719, + 0.19375215470790863, + -0.10346996784210205, + 0.11254917830228806, + 0.0302874892950058, + 0.10501295328140259, + 0.016012834385037422, + 0.022485502064228058, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975, + 0.010529350489377975 + ], + [ + -0.008366091176867485, + -0.0011190441437065601, + -0.0011813186574727297, + 0.04819102957844734, + 0.005909530911594629, + -0.0020593050867319107, + 0.2803594172000885, + 0.02559395506978035, + -0.03755733743309975, + -0.013136299327015877, + 0.004474475048482418, + 0.0016335932305082679, + 0.0001675706444075331, + -0.004079562146216631, + -0.016033874824643135, + -0.014860597439110279, + 0.0033417351078242064, + -0.003095644759014249, + 0.08017106354236603, + 0.11622492969036102, + 0.007717156782746315, + -0.022529728710651398, + -0.002644079504534602, + -0.27519118785858154, + 0.5334027409553528, + -0.6859394907951355, + -0.002025752095505595, + 9.041941186668095e-15, + 0.0037635548505932093, + 0.0011986400932073593, + 1.1134924362762866e-12, + -1.3462052653179235e-11, + -0.002893769182264805, + -7.674038885990003e-27, + 1.1145026868962455e-26, + -2.8069643160026653e-27, + -4.95055824047268e-27, + 2.1196361263338153e-27, + -2.2173607006069402e-27, + 0.0009755520732142031, + 0.003022748976945877, + 0.0005298093310557306, + -0.014395808801054955, + 0.001182115520350635, + 0.01771550066769123, + -0.0008101803250610828, + 0.0003401712456252426, + 0.00033873150823637843, + -0.002910631475970149, + 0.07481085509061813, + -0.0027653027791529894, + 0.05880668759346008, + -0.01401598285883665, + -0.0189046710729599, + -0.0029192916117608547, + -0.042650919407606125, + -0.01587936468422413, + 0.024435587227344513, + 0.07037915289402008, + -0.0011698014568537474, + -0.0011700926115736365, + 0.12125305831432343, + -0.07980982959270477, + -0.15325306355953217, + 0.042727045714855194, + 0.0031664068810641766, + -0.012008057907223701, + -0.01813785545527935, + -0.0019882675260305405, + -0.0028986474499106407, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877, + -0.004092408809810877 + ], + [ + 0.015189500525593758, + 0.0024724018294364214, + -0.001691890531219542, + 0.032461825758218765, + 0.007589434273540974, + -0.005807292647659779, + 0.43604370951652527, + -0.040272004902362823, + -0.01660822331905365, + -6.944470078451559e-05, + -0.006389008369296789, + -0.00864961463958025, + 0.0030137901194393635, + -0.020190708339214325, + -0.005445322021842003, + 0.01127359364181757, + 0.04847146198153496, + 0.0008773633162491024, + 0.04591409116983414, + 0.047685328871011734, + 0.0010711164213716984, + -0.0036323354579508305, + 0.0005128227057866752, + 0.013591645285487175, + 0.028616679832339287, + -0.01918264478445053, + -0.00011808618728537112, + -1.06768367476133e-14, + 0.00021944192121736705, + 6.989049870753661e-05, + 6.492450620090157e-14, + -7.849331536309567e-13, + -0.00016872603737283498, + -3.715850407953054e-26, + 5.361729800606173e-26, + -1.3476899704795766e-26, + -2.3769956796119528e-26, + 1.0177411931511108e-26, + -1.0646619310514371e-26, + -0.0026732191909104586, + 0.002639907179400325, + 0.0003509690286591649, + -0.004497954156249762, + 0.001518400153145194, + 0.01391225028783083, + -0.00770233292132616, + 0.0002292227727593854, + 0.00022633798653259873, + -0.0051920791156589985, + 0.0672755017876625, + -0.0060542128048837185, + 0.07957940548658371, + 0.010872959159314632, + -0.05108710750937462, + -0.04074537754058838, + 0.022777903825044632, + 0.03817138075828552, + -0.03631632775068283, + 0.17828553915023804, + -0.0002406095154583454, + -0.0002506726887077093, + -0.5780364274978638, + 0.14358381927013397, + 0.5992012619972229, + -0.07238521426916122, + 0.10271885246038437, + 0.029439304023981094, + 0.12049929797649384, + -0.04278050735592842, + -0.0173643846064806, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306, + -0.00023861647059675306 + ], + [ + -0.007728592958301306, + -0.08337012678384781, + 0.003680539084598422, + -0.23659497499465942, + 0.016366340219974518, + -0.021081823855638504, + 0.34059572219848633, + -0.05667958781123161, + -0.12493502348661423, + 0.052053630352020264, + -0.005782719235867262, + -0.004669330082833767, + 0.0026668973732739687, + -0.008583690971136093, + -0.010539949871599674, + 0.01429983600974083, + 0.06047661229968071, + 0.0016386633506044745, + -0.19348172843456268, + -0.7390073537826538, + -7.816137076588348e-05, + -0.015198157168924809, + -0.006267625372856855, + 0.005912872962653637, + 0.08090515434741974, + 0.06765758991241455, + -0.002487014979124069, + -3.836687911817904e-14, + 0.004620349500328302, + 0.001471503870561719, + 1.3669855178366497e-12, + -1.6526766066782272e-11, + -0.003552566049620509, + -2.6469779601696886e-22, + 3.3087224502121107e-22, + -8.401053096241687e-23, + -1.4823335070891679e-22, + 6.346611912266295e-23, + -6.639295716410119e-23, + -0.0004974568728357553, + 0.002956324489787221, + 0.0010778148425742984, + -0.02118249610066414, + 0.0023434441536664963, + 0.02715962752699852, + -0.007088696118444204, + 0.0003401297435630113, + 0.00033768132561817765, + -0.0381266213953495, + 0.018648600205779076, + -0.0039572990499436855, + 0.052491847425699234, + 0.08245587348937988, + -0.07901084423065186, + -0.10969128459692001, + -0.011002987623214722, + -5.5417753173969686e-05, + -0.10374065488576889, + -0.08970670402050018, + 0.005049663130193949, + 0.005061611067503691, + -0.05734695866703987, + 0.15761709213256836, + -0.30010583996772766, + -0.1272418200969696, + 0.06053149327635765, + -0.055881377309560776, + 0.041512031108140945, + 0.0565471388399601, + 0.057435549795627594, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941, + -0.005024069454520941 + ], + [ + 0.006845617201179266, + -0.028690356761217117, + -0.032719191163778305, + 0.2189045399427414, + 0.01042286865413189, + -0.0234263576567173, + 0.06739220023155212, + 0.10848457366228104, + 0.0852045789361, + -0.08528231829404831, + 0.009091287851333618, + 0.009996549226343632, + 0.00016849159146659076, + 0.05052374675869942, + 0.002799498848617077, + -0.09436578303575516, + -0.21149155497550964, + -0.0017046320717781782, + 0.2732904255390167, + -0.4641171097755432, + 0.0007772356038913131, + -0.0027620114851742983, + -0.0038888491690158844, + 0.018619976937770844, + 0.06601566076278687, + 0.03284776583313942, + -0.0042800214141607285, + 1.488809076022335e-13, + 0.007951751351356506, + 0.002532512880861759, + 2.3526319781197458e-12, + -2.844302571247681e-11, + -0.006114051677286625, + -8.131516293641283e-20, + 1.1519648082658485e-19, + -2.9116757561866574e-20, + -5.128850670073793e-20, + 2.196102487782347e-20, + -2.297297695970553e-20, + -0.0026458799839019775, + 0.0060500409454107285, + -0.003073512576520443, + -0.06386964023113251, + 0.0037501284386962652, + 0.076802097260952, + -0.004639146383851767, + 0.0007686649914830923, + 0.0007611645851284266, + 0.18672604858875275, + 0.05023374408483505, + 0.0005883491830900311, + -0.047643691301345825, + -0.22312061488628387, + 0.16472508013248444, + 0.36898329854011536, + 0.0696997344493866, + 0.060372497886419296, + 0.11051338165998459, + -0.16611894965171814, + -0.006221963092684746, + -0.006223063915967941, + 0.05428596958518028, + -0.29525822401046753, + 0.21716953814029694, + 0.24610696732997894, + -0.024783937260508537, + 0.23307988047599792, + 0.013146638870239258, + -0.001603239681571722, + -0.03329649567604065, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581, + -0.00864656362682581 + ], + [ + -0.07781612128019333, + -0.030811117962002754, + -0.04022647440433502, + 0.20619012415409088, + 0.025070367380976677, + -0.03158979490399361, + -0.0840519443154335, + 0.280884325504303, + -0.07456386089324951, + -0.059604521840810776, + 0.011064145714044571, + 0.011679313145577908, + -0.0012931731762364507, + -0.02145431376993656, + -0.01658676750957966, + -0.09746919572353363, + -0.09631150960922241, + -0.0023766602389514446, + 0.1550368368625641, + -0.06536708027124405, + 8.019453525776044e-05, + -0.0007346317288465798, + -0.0015222921501845121, + -0.11403128504753113, + -0.1496754288673401, + 0.0006157737225294113, + -0.01662987843155861, + -1.715433350923945e-13, + 0.03089589625597, + 0.009839910082519054, + 9.14092065928962e-12, + -1.1051311948895304e-10, + -0.023755628615617752, + -1.0164395367051604e-20, + 2.6893296075324036e-20, + -6.88214269644119e-21, + -1.206194769224825e-20, + 5.1646055520513974e-21, + -5.402765713807046e-21, + 0.012265626341104507, + 0.009289904497563839, + 0.004201200790703297, + -0.12261415272951126, + 0.007100370246917009, + 0.12496454268693924, + -0.002478836802765727, + 0.0020643582101911306, + 0.002059046644717455, + -0.20651137828826904, + 0.06836479902267456, + 0.012005852535367012, + -0.038124531507492065, + 0.1914743185043335, + -0.2514912784099579, + -0.36358341574668884, + -0.48384901881217957, + 0.014242884702980518, + 0.21494042873382568, + -0.07547986507415771, + 0.001546077779494226, + 0.0015314447227865458, + -0.23880234360694885, + -0.2041926085948944, + -0.12256893515586853, + 0.15180042386054993, + 0.10284842550754547, + 0.03708365932106972, + 0.06528075039386749, + -0.0221756212413311, + -0.02007063291966915, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908, + -0.03359553590416908 + ], + [ + 0.04764239490032196, + -0.02741927094757557, + -0.00955104548484087, + 0.3168470561504364, + -0.012490330263972282, + 0.015093508176505566, + -0.029726224020123482, + -0.363648921251297, + -0.04779540374875069, + -0.029153069481253624, + -0.011269655078649521, + -0.0038683584425598383, + -0.0026792471762746572, + -0.04105323553085327, + -0.020589087158441544, + -0.06582243740558624, + -0.2254859060049057, + 0.0009811954805627465, + 0.315417617559433, + 0.035444699227809906, + 0.009507291950285435, + -0.01602686382830143, + 0.005098547786474228, + 0.32125574350357056, + 0.4693302810192108, + 0.27735331654548645, + 0.0033961560111492872, + -1.608019273291461e-13, + -0.006309889256954193, + -0.002009614137932658, + -1.866895527058432e-12, + 2.257017489570856e-11, + 0.004851625766605139, + -6.505213034913027e-19, + 7.318364664277155e-19, + -1.8338263308055602e-19, + -3.2616062425210902e-19, + 1.3962808739895107e-19, + -1.4607389232227054e-19, + 0.002079505007714033, + -0.012454149313271046, + 0.0032053422182798386, + 0.04516449570655823, + -0.0019778916612267494, + -0.06484156847000122, + 0.0013161624083295465, + -0.0008891739416867495, + -0.0008826234261505306, + -0.12292726337909698, + -0.08732190728187561, + -0.0015163080533966422, + 0.12342918664216995, + 0.11384262889623642, + -0.04179950803518295, + -0.23002855479717255, + 0.18927279114723206, + -0.07754414528608322, + 0.004099350888282061, + -0.08898036181926727, + -0.007379212416708469, + -0.007382824085652828, + -0.06760206818580627, + 0.009142225608229637, + -0.12996159493923187, + 0.00783183891326189, + 0.05918845906853676, + -0.11122602224349976, + 0.009870205074548721, + -0.0217730812728405, + -0.016145460307598114, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291, + 0.006861238274723291 + ], + [ + -0.05170358717441559, + 0.0107714319601655, + 0.004938964731991291, + -0.2525336444377899, + 0.008866620250046253, + -0.006802612915635109, + -0.04773283377289772, + 0.4157913625240326, + 0.029183290898799896, + -0.02618657425045967, + 0.01219604630023241, + 0.00947375688701868, + 0.00023104097635950893, + 0.05524054914712906, + 0.02333786152303219, + 0.045729439705610275, + 0.1531675159931183, + 0.002483631717041135, + -0.22343051433563232, + 0.10714128613471985, + 0.014637836255133152, + -0.01006680354475975, + 0.00969704706221819, + 0.36278560757637024, + 0.5741952061653137, + 0.2906043231487274, + -0.003745218738913536, + 8.970602038971265e-14, + 0.006958056706935167, + 0.002216043882071972, + 2.0586865545624278e-12, + -2.4888632821351564e-11, + -0.005349999293684959, + -9.75781955236954e-19, + 1.0164395367051604e-18, + -2.668153783851046e-19, + -4.760325163569168e-19, + 2.0373789359426093e-19, + -2.1318057387686002e-19, + 0.0010014441795647144, + 0.009381944313645363, + -0.001784676220268011, + -0.03767150267958641, + 0.0012823365395888686, + 0.05006224662065506, + 0.001441778615117073, + 0.0007507867412641644, + 0.0007466599927283823, + 0.03695032373070717, + 0.060910556465387344, + 0.002131312619894743, + -0.11044499278068542, + -0.023539990186691284, + -0.026121623814105988, + 0.09501305967569351, + -0.23227417469024658, + 0.04328189417719841, + 0.01993444189429283, + 0.09033186733722687, + 0.00910408329218626, + 0.009113959036767483, + -0.01834847591817379, + -0.08050490915775299, + 0.04507184401154518, + 0.03981225937604904, + -0.01319693773984909, + 0.1065777912735939, + 0.007432076148688793, + 0.025209931656718254, + 0.0035670700017362833, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357, + -0.007566042244434357 + ], + [ + 0.11744624376296997, + -0.04348212108016014, + 0.035522203892469406, + -0.07524600625038147, + -0.026752416044473648, + 0.04841119423508644, + 0.012568538077175617, + 0.2767130136489868, + -0.21128787100315094, + -0.006766737438738346, + 0.015545159578323364, + 0.00646859360858798, + 0.002492411993443966, + 0.26041239500045776, + -0.08588402718305588, + 0.029061343520879745, + 0.0031002694740891457, + 0.0036362677346915007, + 0.00704682944342494, + 0.11292629688978195, + -0.004292192868888378, + -0.010240007191896439, + -0.004009874537587166, + -0.027206724509596825, + -0.0871892049908638, + -0.0011118309339508414, + 0.009592141024768353, + 2.8257951534271797e-13, + -0.017821144312620163, + -0.005675783380866051, + -5.2725879218229466e-12, + 6.374537009357084e-11, + 0.013702549040317535, + -4.336808689942018e-19, + 4.0657581468206416e-19, + -1.0122043719688889e-19, + -1.8073565512038633e-19, + 7.741748789006297e-20, + -8.099256249751715e-20, + 0.005234776996076107, + -0.0022811938542872667, + -0.0021325566340237856, + 0.10645754635334015, + -0.009573380462825298, + -0.12496405094861984, + 0.026412593200802803, + -0.002168264240026474, + -0.00215343595482409, + 0.024513203650712967, + -0.15911388397216797, + -0.034256462007761, + 0.07990569621324539, + -0.09291412681341171, + 0.07251302897930145, + -0.040215060114860535, + 0.4159305691719055, + -0.11778691411018372, + -0.06352277845144272, + 0.11102727800607681, + 0.0001871164422482252, + 0.00018433456716593355, + -0.34085381031036377, + -0.21941828727722168, + -0.3585546910762787, + 0.14595432579517365, + 0.24811196327209473, + 0.2726026475429535, + 0.13372887670993805, + 0.03499919921159744, + 0.02445601113140583, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437, + 0.019378330558538437 + ], + [ + -0.0876300036907196, + -0.026399753987789154, + -0.0133485933765769, + 0.05648360028862953, + 0.008038938045501709, + -0.0352306105196476, + -0.015575788915157318, + -0.2897969186306, + -0.3569059669971466, + -0.10339274257421494, + -0.008248642086982727, + -0.005841128993779421, + -0.0015110736712813377, + -0.13267502188682556, + -0.041116006672382355, + -0.02234157733619213, + 0.006336600985378027, + 0.0019434005953371525, + 0.03218395262956619, + 0.1564151495695114, + -0.0007510044961236417, + -0.017854301258921623, + -0.0016235284274443984, + 0.01103054266422987, + -0.025017397478222847, + 0.06753947585821152, + -0.004251082427799702, + 2.211841820809468e-13, + 0.007897893898189068, + 0.0025153665337711573, + 2.3367002777163748e-12, + -2.8250377731575682e-11, + -0.0060726348310709, + -6.505213034913027e-19, + 5.963111948670274e-19, + -1.4060746924421386e-19, + -2.4923944472957787e-19, + 1.0685850025205033e-19, + -1.117173591701868e-19, + -0.00887648481875658, + 0.017841093242168427, + -0.002927323803305626, + -0.08799861371517181, + 0.008041869848966599, + 0.12390398234128952, + -0.01252394076436758, + 0.0014520678669214249, + 0.0014346176758408546, + 0.3477432131767273, + 0.15409384667873383, + 0.011726627126336098, + -0.1318800449371338, + -0.10248660296201706, + 0.07693176716566086, + 0.4144367277622223, + -0.2799433767795563, + 0.026325974613428116, + 0.07290979474782944, + 0.06469044834375381, + -0.0002479619870427996, + -0.00025547127006575465, + -0.2270626425743103, + 0.09966100007295609, + -0.3142445683479309, + -0.044127512723207474, + 0.21999932825565338, + -0.11919877678155899, + 0.14903657138347626, + 0.005414656363427639, + 0.052837248891592026, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839, + -0.008588000200688839 + ], + [ + 0.026678744703531265, + -0.039698563516139984, + -0.015421401709318161, + 0.44460609555244446, + -0.0011268542148172855, + -0.06773482263088226, + 0.022417213767766953, + 0.5763617753982544, + 0.01886148378252983, + 0.16475367546081543, + 0.0007458752952516079, + 0.00640350766479969, + 0.0008723348728381097, + -0.0266419667750597, + -0.02374853566288948, + 0.026635201647877693, + -0.07855601608753204, + 0.0019370317459106445, + 0.07169683277606964, + -0.017660779878497124, + -0.0015360938850790262, + -0.002269756281748414, + -0.0022582863457500935, + -0.004492488224059343, + -0.002833837876096368, + 0.019976045936346054, + 0.014036417938768864, + 9.037215420448774e-14, + -0.026078006252646446, + -0.00830548070371151, + -7.715494909632525e-12, + 9.327980055040541e-11, + 0.02005118690431118, + -3.469446951953614e-18, + 2.8189256484623115e-18, + -8.402566836762659e-19, + -1.5043305143236374e-18, + 6.479802046495398e-19, + -6.761440501457452e-19, + -0.035449132323265076, + 0.015982504934072495, + -0.00805711094290018, + 0.03849662095308304, + 0.004452412482351065, + 0.047094885259866714, + -0.06236284598708153, + 2.7042220608564094e-05, + 9.444856914342381e-06, + 0.1680346131324768, + 0.03366812318563461, + -0.026044998317956924, + -0.01736941933631897, + 0.013966241851449013, + 0.004455059301108122, + 0.08611515164375305, + 0.10536383092403412, + 0.0010859462199732661, + -0.16989824175834656, + -0.04948805645108223, + 0.008007675409317017, + 0.00801104586571455, + 0.016543477773666382, + 0.27350160479545593, + -0.09022343903779984, + -0.33977004885673523, + 0.03553740680217743, + -0.3041553497314453, + -0.0005944992299191654, + -0.1317843496799469, + -0.09802424907684326, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239, + 0.02835666388273239 + ], + [ + -0.05930893495678902, + 0.005488715600222349, + -0.02436906285583973, + -0.59116530418396, + -0.04633363336324692, + 0.12166734039783478, + 0.012281534262001514, + 0.19243116676807404, + -0.11528906971216202, + -0.09230778366327286, + -0.0005012954352423549, + 0.0037917818408459425, + -0.00041314688860438764, + -0.13296417891979218, + 0.025044215843081474, + -0.11308135092258453, + -0.46065470576286316, + -0.012486875057220459, + 0.4216702878475189, + 0.031216908246278763, + 0.006335825193673372, + 0.0030051125213503838, + -0.004238504450768232, + -0.015485085546970367, + -0.026509501039981842, + -0.009233130142092705, + -0.010973741300404072, + 5.631883848167263e-13, + 0.020387845113873482, + 0.0064932298846542835, + 6.031952715090938e-12, + -7.292613735110365e-11, + -0.015676075592637062, + -2.7755575615628914e-17, + 1.734723475976807e-18, + -2.168404344971009e-18, + -5.204170427930421e-18, + 2.0599841277224584e-18, + -2.15485181781494e-18, + 0.05855283886194229, + 0.015742499381303787, + 0.0018922477029263973, + 0.04956841468811035, + -0.01847749389708042, + -0.15309420228004456, + 0.132549449801445, + -0.0017388482810929418, + -0.0017109697218984365, + 0.0047505442053079605, + -0.0575789250433445, + -0.014343616552650928, + -0.0357130691409111, + 0.06106004863977432, + 0.018646463751792908, + 0.0516744926571846, + 0.003470826428383589, + 0.012547526508569717, + 0.03012590855360031, + 0.027342045679688454, + -0.004644334316253662, + -0.004652470350265503, + 0.013195623643696308, + 0.08034004271030426, + 0.0339554101228714, + -0.11644995957612991, + -0.017448484897613525, + -0.17930114269256592, + -0.01521011907607317, + -0.12294408679008484, + -0.0990789532661438, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006, + -0.02216930501163006 + ], + [ + 0.034256551414728165, + -0.2126363068819046, + -0.06746198236942291, + -0.18310002982616425, + 0.05010620504617691, + -0.0966804176568985, + 0.000330624490743503, + -0.05183428153395653, + 0.7009097933769226, + 0.05570308119058609, + 0.000901424849871546, + 0.0007356933783739805, + 0.0005515656084753573, + 0.0519583635032177, + 6.105813372414559e-05, + -4.217115929350257e-05, + -0.14548884332180023, + -0.00018426566384732723, + 0.016811182722449303, + 0.05249687284231186, + 0.0018006761092692614, + -0.001003360841423273, + -0.007144920062273741, + 0.005275543313473463, + -0.007469406817108393, + -0.027766471728682518, + 0.009691709652543068, + -1.013633621482768e-12, + -0.018005594611167908, + -0.005734540987759829, + -5.3274884503906605e-12, + 6.440581401534473e-11, + 0.013844361528754234, + 0.0, + 4.163336342344337e-17, + 0.0, + -1.3877787807814457e-17, + 5.204170427930421e-18, + -8.239936510889834e-18, + -0.06479854136705399, + -0.08156216144561768, + -0.0057227048091590405, + -0.14216144382953644, + 0.021522507071495056, + 0.17493553459644318, + -0.14482195675373077, + 0.0009838714031502604, + 0.0009595336159691215, + 0.24005793035030365, + 0.022637056186795235, + 0.08205285668373108, + 0.07625891268253326, + 0.08031491935253143, + 0.0819208174943924, + -0.0268348827958107, + 0.004612520802766085, + -0.25811830163002014, + 0.2101365178823471, + 0.027797166258096695, + -0.0029540329705923796, + -0.002954356838017702, + -0.20552174746990204, + 0.033312007784843445, + -0.18154367804527283, + -0.03720616176724434, + 0.07167066633701324, + -0.05561027303338051, + 0.03341732546687126, + 0.07814373075962067, + 0.026990415528416634, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477, + 0.019578898325562477 + ], + [ + -0.06738943606615067, + -0.23874862492084503, + -0.011760417371988297, + -0.061376772820949554, + -0.004791471641510725, + 0.013934778049588203, + 0.004884127993136644, + 0.009923449717462063, + 0.0776415467262268, + -0.08327336609363556, + 0.000113050096842926, + -0.0008132873335853219, + 0.00033824911224655807, + -0.21830469369888306, + -0.3319859206676483, + -0.04865634813904762, + 0.35487058758735657, + 0.004477419424802065, + 0.17475025355815887, + -0.017599938437342644, + 0.0005104179144836962, + -0.024664800614118576, + -0.002598766004666686, + 0.038108378648757935, + -0.012941695749759674, + 0.018558267503976822, + -0.00931948609650135, + 8.786582572639645e-13, + 0.01731414906680584, + 0.005514322314411402, + 5.122735569074166e-12, + -6.193153628819559e-11, + -0.013312716968357563, + 2.7755575615628914e-17, + 5.551115123125783e-17, + 0.0, + -5.551115123125783e-17, + 2.7755575615628914e-17, + 0.0, + 0.00789417140185833, + 0.03740016743540764, + -0.012044934555888176, + 0.012240499258041382, + -0.008917050436139107, + -0.005845119711011648, + 0.03150003030896187, + -0.00025836736313067377, + -0.00022082572104409337, + -0.11425715684890747, + -0.029490046203136444, + -0.12568366527557373, + -0.02192864380776882, + -0.6195645332336426, + 0.04701058566570282, + -0.19266998767852783, + -0.05597072094678879, + -0.24668888747692108, + 0.007123696617782116, + -0.05986815690994263, + -0.0006919871084392071, + -0.0006977199809625745, + 0.028434565290808678, + 0.02479807287454605, + -0.009305465966463089, + -0.03224003687500954, + 0.043036334216594696, + -0.07606286555528641, + 0.01709965243935585, + -0.18151289224624634, + -0.17065154016017914, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954, + -0.018827036023139954 + ], + [ + -0.062384191900491714, + 0.1625244915485382, + 0.03175423666834831, + 0.1928054243326187, + -0.06892187893390656, + 0.2711363732814789, + 0.04222458600997925, + 0.08780742436647415, + 0.2961558997631073, + -0.6198632121086121, + 0.008008765988051891, + 0.0026185151655226946, + -0.001564846490509808, + -0.14459525048732758, + 0.06693767011165619, + -0.07563356310129166, + 0.16725899279117584, + -0.011122931726276875, + -0.10066920518875122, + -0.07313045859336853, + -0.005666804034262896, + 0.00843772478401661, + 0.0042678844183683395, + -0.011252457275986671, + -0.007921485230326653, + -0.01582467183470726, + -0.021398555487394333, + 7.748940378249358e-13, + 0.03975599631667137, + 0.01266174390912056, + 1.1761730478454524e-11, + -1.4220476185489161e-10, + -0.0305680800229311, + -5.724587470723463e-17, + 0.0, + -1.1102230246251565e-16, + -5.551115123125783e-17, + 1.1102230246251565e-16, + -2.7755575615628914e-17, + 0.11895015090703964, + -0.009916611015796661, + 0.003669368801638484, + 0.06794748455286026, + -0.037449415773153305, + -0.3009282052516937, + 0.25764328241348267, + -0.002874600701034069, + -0.002827949356287718, + 0.07701234519481659, + -0.11632237583398819, + 0.03569953888654709, + -0.006736866198480129, + 0.11074353009462357, + 0.06113994121551514, + 0.049285370856523514, + 0.015228535048663616, + 0.005065704695880413, + 0.015697352588176727, + 0.03615985065698624, + 0.006185663398355246, + 0.006156511139124632, + -0.09487223625183105, + 0.02484179101884365, + -0.09927327185869217, + -0.054753612726926804, + 0.024547070264816284, + -0.10340060293674469, + 0.011771401390433311, + 0.09621115773916245, + 0.0869850441813469, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621, + -0.04322981834411621 + ], + [ + 0.11169850826263428, + 0.08624473214149475, + -0.11450286954641342, + -0.08212261646986008, + 0.028874259442090988, + -0.14267340302467346, + -0.03273367881774902, + 0.10469008982181549, + -0.25549426674842834, + -0.2955339848995209, + 0.002713416237384081, + 0.002608058974146843, + -0.0014129907358437777, + 0.07368859648704529, + -0.0069569614715874195, + -0.06825938075780869, + 0.24992354214191437, + -0.002380401361733675, + 0.07363908737897873, + -0.017402812838554382, + -0.0013602524995803833, + 0.01483013853430748, + 0.011550509370863438, + 0.007619520183652639, + 0.006085566245019436, + 0.003915759269148111, + 0.0073645953088998795, + -8.598677325721837e-14, + -0.013682719320058823, + -0.004357738886028528, + -4.048317236993171e-12, + 4.8942738750668013e-11, + 0.010520556010305882, + -2.7755575615628914e-17, + 5.551115123125783e-17, + -6.938893903907228e-18, + -2.7755575615628914e-17, + 2.7755575615628914e-17, + 0.0, + -0.0486111082136631, + -0.07685725390911102, + 0.0017343121580779552, + -0.19979634881019592, + 0.022629514336586, + 0.2963786721229553, + -0.12383314222097397, + 0.0009383568540215492, + 0.0009144672658294439, + 0.043723031878471375, + -0.18850480020046234, + -0.07048729062080383, + 0.22391867637634277, + 0.06021055579185486, + 0.07030559331178665, + -0.06599819660186768, + 0.24386002123355865, + 0.10840511322021484, + 0.5386586785316467, + -0.005187097005546093, + -0.011777006089687347, + -0.011810943484306335, + 0.09066605567932129, + 0.07377422600984573, + 0.08725463598966599, + -0.11865495145320892, + -0.05327041819691658, + -0.159443661570549, + -0.03975484520196915, + 0.03967595100402832, + 0.08077381551265717, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122, + 0.014878294430673122 + ], + [ + 0.04316558688879013, + -0.20040108263492584, + 0.0011089679319411516, + -0.041441887617111206, + 0.0774073600769043, + -0.2612297534942627, + -0.009214249439537525, + 0.006335542071610689, + 0.0688791498541832, + -0.522470235824585, + 0.005810791160911322, + -0.00012245278048794717, + -0.0008498973329551518, + -0.10284712165594101, + -0.009730271995067596, + 0.08003652840852737, + -0.2549869120121002, + 0.01077705342322588, + -0.05479871854186058, + 0.05099127069115639, + 0.006459137424826622, + -0.005091041326522827, + -0.011357291601598263, + -0.0011021130485460162, + -0.027842039242386818, + -0.013198399916291237, + 0.026675991714000702, + -7.072120666862247e-13, + -0.04956118389964104, + -0.01578455977141857, + -1.4663395497827025e-11, + 1.7727826151503479e-10, + 0.038107216358184814, + 8.326672684688674e-17, + 4.163336342344337e-17, + -1.3877787807814457e-17, + -1.3877787807814457e-17, + -2.7755575615628914e-17, + 2.7755575615628914e-17, + -0.13343942165374756, + 0.09062463790178299, + -0.01629418507218361, + 0.10771343111991882, + 0.02665816619992256, + 0.17757369577884674, + -0.2518199682235718, + 0.0016933187143877149, + 0.001635729568079114, + -0.13120996952056885, + 0.05423645302653313, + -0.15183280408382416, + -0.06949390470981598, + -0.15526741743087769, + -0.15990376472473145, + -0.05757598951458931, + 0.04731110483407974, + 0.3076838254928589, + -0.3232801556587219, + 0.08569473773241043, + 0.024346964433789253, + 0.024351274594664574, + 0.009011917747557163, + -0.027333056554198265, + -0.032604314386844635, + 0.02116488479077816, + -0.01817135326564312, + -0.0019522650400176644, + 0.00633148243650794, + 0.11948245763778687, + 0.10507086664438248, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764, + 0.05389176681637764 + ], + [ + 0.03853968158364296, + 0.10884230583906174, + 0.026844825595617294, + -0.1006113812327385, + 0.0034454776905477047, + -0.06555004417896271, + -0.02539736218750477, + -0.11161284893751144, + 0.1704067587852478, + -0.06586455553770065, + -0.0017558929976075888, + -0.002627525944262743, + -3.988329990534112e-05, + 0.07992253452539444, + 0.0385296605527401, + 0.011885402724146843, + 0.2820779085159302, + -0.007971452549099922, + 0.10725445300340652, + -0.024867994710803032, + 0.004853395279496908, + 0.01498901192098856, + 0.006166440900415182, + 0.0007331614033319056, + 0.004828710108995438, + -0.010248958133161068, + 0.009197159670293331, + -1.146610584257246e-12, + -0.017086723819375038, + -0.0054418593645095825, + -5.0552340091769565e-12, + 6.111844363942964e-11, + 0.0131378760561347, + 0.0, + 1.1102230246251565e-16, + -2.0816681711721685e-17, + 0.0, + 5.551115123125783e-17, + -1.3877787807814457e-17, + -0.02971973642706871, + -0.017519107088446617, + 0.003855619579553604, + 0.011498741805553436, + 0.00857789721339941, + 0.006185548380017281, + -0.06482162326574326, + 0.000633746269159019, + 0.0005885717691853642, + 0.06214875727891922, + 0.05483986809849739, + 0.08462339639663696, + -0.030552884563803673, + 0.1897357553243637, + -0.04887054115533829, + 0.032618798315525055, + 0.07219716161489487, + 0.4209028482437134, + -0.07207895815372467, + -0.07183343172073364, + 0.006103288382291794, + 0.006086132489144802, + -0.09261790663003922, + -0.08706962317228317, + -0.15187989175319672, + 0.06952402740716934, + 0.0307156965136528, + 0.009186860173940659, + 0.036664023995399475, + -0.5180085301399231, + -0.5021392703056335, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858, + 0.018579736351966858 + ], + [ + 0.08204368501901627, + 0.1731446087360382, + -0.0182290431112051, + 0.013665099628269672, + 0.018539350479841232, + -0.10046646744012833, + -0.026083480566740036, + -0.02552223764359951, + -0.17351441085338593, + -0.28715139627456665, + 0.0017922258703038096, + -0.00040038119186647236, + -0.00035490599111653864, + -0.05016759783029556, + 0.11504146456718445, + 0.0848747119307518, + -0.17162670195102692, + -0.0025262690614908934, + -0.27272406220436096, + -0.07161232829093933, + 0.0034329164773225784, + 0.023299822583794594, + 0.011303912848234177, + -0.02227211557328701, + 0.008450377732515335, + -0.01328238658607006, + 0.012057841755449772, + 4.875544412641375e-13, + -0.022401627153158188, + -0.007134587503969669, + -6.6280279875652326e-12, + 8.012918106814482e-11, + 0.017224468290805817, + -2.7755575615628914e-17, + 1.3877787807814457e-17, + 1.1102230246251565e-16, + 6.505213034913027e-19, + 0.0, + 2.7755575615628914e-17, + -0.030607575550675392, + -0.003923223819583654, + 0.0039031393826007843, + -0.007314491085708141, + 0.01233106479048729, + 0.10460303723812103, + -0.06462426483631134, + -0.00021645311790052801, + -0.00018104015907738358, + 0.08240250498056412, + -0.009334723465144634, + -0.0306343212723732, + 0.0029406221583485603, + 0.10310641676187515, + -0.060162484645843506, + 0.008446087129414082, + -0.023903172463178635, + -0.6625460982322693, + -0.07769119739532471, + 0.043969832360744476, + 0.010398488491773605, + 0.010386364534497261, + 0.044147733598947525, + -0.08560611307621002, + 0.11371548473834991, + 0.057015858590602875, + 0.002819767687469721, + -0.0028829367365688086, + -0.0580068863928318, + -0.2942149341106415, + -0.30791574716567993, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613, + 0.024359051138162613 + ], + [ + -0.06716323643922806, + -0.015176764689385891, + -0.024409059435129166, + 0.11903798580169678, + -0.007753458805382252, + 0.0656476840376854, + 0.020616669207811356, + -0.03492555394768715, + 0.03242357075214386, + 0.07362594455480576, + -0.0023954943753778934, + -0.001417828258126974, + -0.00017571804346516728, + -0.03819354623556137, + -0.0815085843205452, + 0.036047033965587616, + -0.43941691517829895, + 0.032258182764053345, + -0.48861101269721985, + -7.416019798256457e-05, + 0.007571782451122999, + -0.01977604813873768, + -0.009904169477522373, + -0.004651802591979504, + 0.019787905737757683, + 0.0021541102323681116, + -0.012057474814355373, + 8.814199370377196e-13, + 0.02240198478102684, + 0.007134704384952784, + 6.627406956560833e-12, + -8.013065211365245e-11, + -0.01722472719848156, + 4.163336342344337e-17, + 5.551115123125783e-17, + 1.1102230246251565e-16, + -2.7755575615628914e-17, + -5.551115123125783e-17, + 0.0, + 0.06181854009628296, + -0.13108272850513458, + -0.0021015640813857317, + -0.13542313873767853, + -0.022111909464001656, + -0.019238684326410294, + 0.11740285158157349, + -0.000371868401998654, + -0.00034957434399984777, + -0.18193595111370087, + -0.14387616515159607, + -0.04915660247206688, + 0.14281973242759705, + -0.26860475540161133, + 0.17904379963874817, + -0.02411775477230549, + -0.0020601749420166016, + 0.28677043318748474, + 0.26968005299568176, + 0.18584345281124115, + 0.007561725098639727, + 0.00756616797298193, + -0.03210208937525749, + 0.024959160014986992, + -0.0635361298918724, + -0.02922135964035988, + 0.03045273944735527, + -0.035676322877407074, + 0.03212128207087517, + -0.19966918230056763, + -0.2034635692834854, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205, + -0.024359432980418205 + ] + ], + "encoded_input_dim": 140, + "mean": [ + -0.010634142905473709, + 0.056608956307172775, + 0.0572563111782074, + 0.7604517340660095, + 0.11315243691205978, + 0.03482051193714142, + 0.08955748379230499, + -0.0053392089903354645, + 0.02204631268978119, + 0.05831008777022362, + -1.0967743396759033, + -0.5423644781112671, + -0.27207469940185547, + -0.021432911977171898, + -0.008900946006178856, + -0.11227908730506897, + 0.7354849576950073, + 0.0003533019917085767, + 0.2370520830154419, + 0.23403042554855347, + -0.0010782002937048674, + 0.0057747564278542995, + 0.0034331646747887135, + 0.043552279472351074, + -0.02616506814956665, + -0.011841392144560814, + -0.2024962306022644, + 2.91038298062124e-13, + -0.7663765549659729, + 0.7559386491775513, + 4.53479996609385e-11, + -5.482543308232835e-10, + 0.589259922504425, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + 0.0, + -0.5134562253952026, + 0.0657612755894661, + 0.008587519638240337, + 0.48768356442451477, + 0.014867561869323254, + 0.3822880983352661, + 0.13157354295253754, + 0.004941494669765234, + 0.004944906570017338, + 0.02016887627542019, + -0.04027227312326431, + 0.016343748196959496, + -0.0019311634823679924, + -0.013069219887256622, + 0.37820231914520264, + 0.03484316170215607, + -1.576539397239685, + 0.006656243000179529, + 2.0437493324279785, + 0.08760439604520798, + 0.018343565985560417, + 0.018346967175602913, + 0.009281089529395103, + 0.34724321961402893, + -0.030798060819506645, + 0.08700264245271683, + -0.045663367956876755, + 0.24090532958507538, + -0.072051502764225, + -0.0072078094817698, + -0.02014041133224964, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 1.0, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184, + 0.8333333134651184 + ], + "method": "train_split_pca_standardized", + "output_dim": 32, + "raw_input_dim": 70, + "scale": [ + 7.6479597091674805, + 6.800845146179199, + 4.303369522094727, + 2.428558349609375, + 1.6685186624526978, + 1.5334241390228271, + 1.1043930053710938, + 0.7665870785713196, + 0.7313748002052307, + 0.6281603574752808, + 0.5484201312065125, + 0.5021082162857056, + 0.47317689657211304, + 0.445694237947464, + 0.41110771894454956, + 0.3320724070072174, + 0.3014792501926422, + 0.29494109749794006, + 0.2723846435546875, + 0.267343133687973, + 0.25177496671676636, + 0.23485739529132843, + 0.2058505415916443, + 0.177432119846344, + 0.16502965986728668, + 0.14961080253124237, + 0.1467713713645935, + 0.13825903832912445, + 0.1367480605840683, + 0.13011839985847473, + 0.11349403113126755, + 0.10745701938867569 + ], + "state_vectorization": "right_zero_pad_plus_validity_mask" + }, + "trainable_parameters": 99880992, + "training_losses": [ + 1.1676948070526123, + 0.6497068405151367, + 0.4265829920768738, + 1.0592600107192993, + 0.4688650369644165, + 0.7258111238479614, + 0.6003483533859253, + 0.48628485202789307, + 0.5061243772506714, + 0.41319355368614197, + 0.4170148968696594, + 0.38549062609672546, + 0.6488516926765442, + 0.32852286100387573, + 0.4813322424888611, + 0.444707453250885, + 0.7038055658340454, + 0.67229163646698, + 0.42834287881851196, + 0.4291791319847107, + 0.9696183800697327, + 0.45879077911376953, + 0.5311608910560608, + 0.5839967727661133, + 0.6291602253913879, + 0.5381404757499695, + 0.40443503856658936, + 0.3095071017742157, + 0.5828569531440735, + 0.8522820472717285, + 0.3345695734024048, + 0.363203227519989, + 0.42780572175979614, + 0.4685305953025818, + 0.48437970876693726, + 0.49726733565330505, + 0.2932574152946472, + 0.4293547570705414, + 0.4596598744392395, + 0.5307748317718506, + 0.3711506128311157, + 0.4155687093734741, + 0.20648640394210815, + 0.12423643469810486, + 0.27511996030807495, + 0.4388335049152374, + 0.369880735874176, + 0.43804457783699036, + 0.25307291746139526, + 0.3373062014579773, + 0.21706488728523254, + 0.49691155552864075, + 0.31339460611343384, + 0.28444862365722656, + 0.36420905590057373, + 0.3943978548049927, + 0.3200118839740753, + 0.27250760793685913, + 0.33625566959381104, + 0.34034761786460876, + 0.22104781866073608, + 0.22657829523086548, + 0.45812416076660156, + 0.33466678857803345, + 0.524996817111969, + 0.39007484912872314, + 0.3783608078956604, + 0.37963467836380005, + 0.23781035840511322, + 0.31722134351730347, + 0.3410072326660156, + 0.3368871808052063, + 0.3246345818042755, + 0.385969877243042, + 0.39538681507110596, + 0.46373099088668823, + 0.33329761028289795, + 0.2602175176143646, + 0.27242183685302734, + 0.2011057585477829, + 0.29533565044403076, + 0.38614001870155334, + 0.3247067630290985, + 0.20522373914718628, + 0.31018874049186707, + 0.41956251859664917, + 0.3136429190635681, + 0.425092875957489, + 0.40499067306518555, + 0.3417591452598572, + 0.4612193703651428, + 0.2844722867012024, + 0.34130868315696716, + 0.4164886772632599, + 0.5067241191864014, + 0.4774448275566101, + 0.3045823574066162, + 0.29497748613357544, + 0.3455316424369812, + 0.34580451250076294, + 0.16919875144958496, + 0.38849973678588867, + 0.20079873502254486, + 0.3863757848739624, + 0.39453577995300293, + 0.3307757079601288, + 0.3313713073730469, + 0.2433100789785385, + 0.33202582597732544, + 0.18125462532043457, + 0.4247853755950928, + 0.5259015560150146, + 0.584337055683136, + 0.46457016468048096, + 0.44879502058029175, + 0.3132793605327606, + 0.3554070293903351, + 0.5149505138397217, + 0.5484674572944641, + 0.29211458563804626, + 0.387271523475647, + 0.36930984258651733, + 0.4104295074939728, + 0.3542957603931427, + 0.28961342573165894, + 0.4218081533908844, + 0.4282488226890564, + 0.45910975337028503, + 0.28448033332824707, + 0.5588304996490479, + 0.40379226207733154, + 0.21338611841201782, + 0.5758305788040161, + 0.4648541212081909, + 0.4893881380558014, + 0.19657234847545624, + 0.30065369606018066, + 0.24705444276332855, + 0.33225470781326294, + 0.22146812081336975, + 0.19210413098335266, + 0.43460556864738464, + 0.3824387490749359, + 0.13800102472305298, + 0.29237836599349976, + 0.792522668838501, + 0.5520352125167847, + 0.44423186779022217, + 0.34510543942451477, + 0.3784821629524231, + 0.34530675411224365, + 0.2723683714866638, + 0.3415600061416626, + 0.18776217103004456, + 0.3295007050037384, + 0.2950626015663147, + 0.21560138463974, + 0.42598509788513184, + 0.2496587485074997, + 0.2556719183921814, + 0.23741669952869415, + 0.2699950039386749, + 0.42729946970939636, + 0.3127489686012268, + 0.17932270467281342, + 0.25401878356933594, + 0.3044126629829407, + 0.21414721012115479, + 0.41197967529296875, + 0.45978525280952454, + 0.39419329166412354, + 0.44097018241882324, + 0.24413317441940308, + 0.2986351549625397, + 0.2549826502799988, + 0.26768022775650024, + 0.23345957696437836, + 0.21673837304115295, + 0.4145328402519226, + 0.23484385013580322, + 0.3639269769191742, + 0.40599867701530457, + 0.3779575824737549, + 0.37668073177337646, + 0.2524588406085968, + 0.31641221046447754, + 0.39783644676208496, + 0.30318352580070496, + 0.2827398180961609, + 0.26917189359664917, + 0.2332092523574829, + 0.2461726814508438, + 0.20406538248062134, + 0.21677295863628387, + 0.33777478337287903, + 0.3814047574996948, + 0.24992424249649048, + 0.24793404340744019, + 0.1812022179365158, + 0.2630992829799652, + 0.5359546542167664, + 0.19327601790428162, + 0.2026141732931137, + 0.30760133266448975, + 0.3282025456428528, + 0.1832309365272522, + 0.24147404730319977, + 0.29277491569519043, + 0.3824063539505005, + 0.34679466485977173, + 0.19903826713562012, + 0.27386051416397095, + 0.27046775817871094, + 0.19065600633621216, + 0.24362485110759735, + 0.22226282954216003, + 0.3192744255065918, + 0.17106977105140686, + 0.3220664858818054, + 0.4224320352077484, + 0.26418793201446533, + 0.1596534699201584, + 0.1730385422706604, + 0.28401973843574524, + 0.07762277871370316, + 0.24175690114498138, + 0.20122863352298737, + 0.15196925401687622, + 0.1860208809375763, + 0.2749900817871094, + 0.2186141014099121, + 0.437613844871521, + 0.19282689690589905, + 0.30633991956710815, + 0.2527637481689453, + 0.17659741640090942, + 0.17723819613456726, + 0.11612221598625183, + 0.42704683542251587, + 0.41437217593193054, + 0.19426675140857697, + 0.26666951179504395, + 0.23158329725265503, + 0.397479385137558, + 0.3109046518802643, + 0.3429715037345886, + 0.3772802948951721, + 0.23100796341896057, + 0.26203981041908264, + 0.38150590658187866, + 0.36904728412628174, + 0.33709728717803955, + 0.3771951198577881, + 0.19907709956169128, + 0.17307312786579132, + 0.470270037651062, + 0.4883336126804352, + 0.2506491243839264, + 0.4643671214580536, + 0.21961329877376556, + 0.31669965386390686, + 0.29090142250061035, + 0.4060588479042053, + 0.38450783491134644, + 0.34700751304626465, + 0.32523828744888306, + 0.32990896701812744, + 0.28467339277267456, + 0.2693755030632019, + 0.21429800987243652, + 0.2898072898387909, + 0.20180827379226685, + 0.3311588168144226, + 0.23009715974330902, + 0.3612101078033447, + 0.3164447546005249, + 0.3408083915710449, + 0.11965350806713104, + 0.2777954936027527, + 0.29369455575942993, + 0.15501773357391357, + 0.3085487484931946, + 0.22452127933502197, + 0.424852192401886, + 0.19879168272018433, + 0.26017680764198303, + 0.2913983166217804, + 0.23972541093826294, + 0.18070411682128906, + 0.2581167221069336, + 0.2489633560180664, + 0.21604694426059723, + 0.12098412215709686, + 0.23561373353004456, + 0.2449021339416504, + 0.19064810872077942, + 0.20068426430225372, + 0.3589458763599396, + 0.47867244482040405, + 0.1757178008556366, + 0.35539260506629944, + 0.16862961649894714, + 0.24554038047790527, + 0.1539195328950882, + 0.39007970690727234, + 0.3714703619480133, + 0.21235518157482147, + 0.18458643555641174, + 0.13582660257816315, + 0.19860640168190002, + 0.14347799122333527, + 0.2904406189918518, + 0.15445038676261902, + 0.293259859085083, + 0.20900267362594604, + 0.2487785667181015, + 0.23064619302749634, + 0.19196215271949768, + 0.3173045516014099, + 0.2613755166530609, + 0.16291646659374237, + 0.23454394936561584, + 0.28762251138687134, + 0.24499425292015076, + 0.1441200077533722, + 0.22363515198230743, + 0.2038973867893219, + 0.2233765572309494, + 0.14956778287887573, + 0.2803899943828583, + 0.13119076192378998, + 0.3177010715007782, + 0.17841410636901855, + 0.2481905221939087, + 0.3765169084072113, + 0.22311407327651978, + 0.21485823392868042, + 0.15577012300491333, + 0.20769695937633514, + 0.22369791567325592, + 0.30286675691604614, + 0.4864262342453003, + 0.24793434143066406, + 0.2510664761066437, + 0.370930552482605, + 0.1810992956161499, + 0.2952037751674652, + 0.15990714728832245, + 0.19884365797042847, + 0.4905744194984436, + 0.49416685104370117, + 0.2162667065858841, + 0.1115676611661911, + 0.18557006120681763, + 0.36669421195983887, + 0.3783358335494995, + 0.38363489508628845, + 0.1817508488893509, + 0.32429832220077515, + 0.24820365011692047, + 0.2372015118598938, + 0.3973502814769745, + 0.40707361698150635, + 0.19164618849754333, + 0.22864601016044617, + 0.20184564590454102, + 0.255738765001297, + 0.20584753155708313, + 0.3375821113586426, + 0.39400655031204224, + 0.22309741377830505, + 0.2679714560508728, + 0.39598366618156433, + 0.13105495274066925, + 0.4073922038078308, + 0.14621467888355255, + 0.22184596955776215, + 0.22036752104759216, + 0.23647120594978333, + 0.6329737305641174, + 0.12920458614826202, + 0.21344861388206482, + 0.1598571091890335, + 0.14107120037078857, + 0.2457326054573059, + 0.16452574729919434, + 0.13258469104766846, + 0.30300992727279663, + 0.3928891718387604, + 0.23102164268493652, + 0.24334174394607544, + 0.26128333806991577, + 0.3283158838748932, + 0.2882763743400574, + 0.22126686573028564, + 0.14268821477890015, + 0.12519949674606323, + 0.1400013267993927, + 0.40567487478256226, + 0.11201446503400803, + 0.19900178909301758, + 0.1737632304430008, + 0.19722412526607513, + 0.28073930740356445, + 0.21655210852622986, + 0.2855319082736969, + 0.12461237609386444, + 0.3151480555534363, + 0.26596546173095703, + 0.3220738172531128, + 0.39024171233177185, + 0.1417694240808487, + 0.2934539318084717, + 0.09562817215919495, + 0.13137152791023254, + 0.2190639227628708, + 0.21780797839164734, + 0.26861387491226196, + 0.24028083682060242, + 0.18668967485427856, + 0.20742137730121613, + 0.24928122758865356, + 0.580712080001831, + 0.17282089591026306, + 0.2266470193862915, + 0.2662385106086731, + 0.27068445086479187, + 0.23971667885780334, + 0.306415855884552, + 0.1340416818857193, + 0.18835926055908203, + 0.35153043270111084, + 0.37398070096969604, + 0.2073577344417572, + 0.16941598057746887, + 0.2678612470626831, + 0.17797303199768066, + 0.22695952653884888, + 0.2946320176124573, + 0.31331080198287964, + 0.18844512104988098, + 0.42933857440948486, + 0.25761640071868896, + 0.21795183420181274, + 0.1809368133544922, + 0.2662700116634369, + 0.37014079093933105, + 0.20953017473220825, + 0.22960320115089417, + 0.3371022343635559, + 0.13556385040283203, + 0.32453638315200806, + 0.19455301761627197, + 0.2691726088523865, + 0.16296207904815674, + 0.16053661704063416, + 0.2318914383649826, + 0.1132531464099884, + 0.25343963503837585, + 0.11855718493461609, + 0.14766262471675873, + 0.19149714708328247, + 0.2815466523170471, + 0.33428746461868286, + 0.49930647015571594, + 0.18477162718772888, + 0.19667941331863403, + 0.21894638240337372, + 0.19939854741096497, + 0.13885605335235596, + 0.5003819465637207, + 0.15298938751220703, + 0.24798715114593506, + 0.24526256322860718, + 0.19373655319213867, + 0.19946345686912537, + 0.4663854241371155, + 0.30962350964546204, + 0.2901335656642914, + 0.18115124106407166, + 0.4039350152015686, + 0.2027161866426468, + 0.20608367025852203, + 0.22314882278442383, + 0.20964179933071136, + 0.19724535942077637, + 0.18617120385169983, + 0.23631028831005096, + 0.13031679391860962, + 0.1308784931898117, + 0.2070772647857666, + 0.24443894624710083, + 0.12951447069644928, + 0.19061869382858276, + 0.2839924693107605, + 0.17501327395439148, + 0.3277215361595154, + 0.21861734986305237, + 0.2376340925693512, + 0.14487095177173615, + 0.3218664526939392, + 0.4186013340950012, + 0.167229562997818, + 0.48294320702552795, + 0.29250144958496094, + 0.157058984041214, + 0.30976516008377075, + 0.21923047304153442, + 0.2362671047449112, + 0.22670109570026398, + 0.2366236299276352, + 0.14413312077522278, + 0.18569570779800415, + 0.3539869487285614, + 0.37940287590026855, + 0.2044859230518341, + 0.2893669605255127, + 0.25782448053359985, + 0.13709279894828796, + 0.19199024140834808, + 0.2628486454486847, + 0.22435085475444794, + 0.34552550315856934, + 0.27933838963508606, + 0.3420497477054596, + 0.27778178453445435, + 0.28258180618286133, + 0.2927253544330597, + 0.5302075147628784, + 0.1654171198606491, + 0.20787867903709412, + 0.1783493459224701, + 0.2339373230934143, + 0.42724335193634033, + 0.25114738941192627, + 0.09178382903337479, + 0.22247692942619324, + 0.14051426947116852, + 0.16036532819271088, + 0.12751106917858124, + 0.30725955963134766, + 0.2253602296113968, + 0.36128339171409607, + 0.13080190122127533, + 0.1728651225566864, + 0.16634932160377502, + 0.22424590587615967, + 0.1108262687921524, + 0.38072121143341064, + 0.22739318013191223, + 0.19592930376529694, + 0.36848676204681396, + 0.2387230098247528, + 0.2379179298877716, + 0.1298540234565735, + 0.14553609490394592, + 0.27371945977211, + 0.16144505143165588, + 0.261026531457901, + 0.1936573088169098, + 0.2752262353897095, + 0.4302822947502136, + 0.16993549466133118, + 0.34241408109664917, + 0.35928159952163696, + 0.10660918802022934, + 0.19705986976623535, + 0.21200373768806458, + 0.3574490547180176, + 0.2241107076406479, + 0.3124440908432007, + 0.2762366235256195, + 0.16435731947422028, + 0.2901562452316284, + 0.2008827030658722, + 0.21252994239330292, + 0.1889800876379013, + 0.16376084089279175, + 0.3192618489265442, + 0.16633623838424683, + 0.18131589889526367, + 0.1563592404127121, + 0.3325003385543823, + 0.2810875475406647, + 0.24432045221328735, + 0.20906899869441986, + 0.18416020274162292, + 0.1826028674840927, + 0.2117181271314621, + 0.2619202136993408, + 0.16467823088169098, + 0.3437096178531647, + 0.31104934215545654, + 0.122979074716568, + 0.17142684757709503, + 0.14930912852287292, + 0.40115806460380554, + 0.2054537534713745, + 0.1520732045173645, + 0.3465604782104492, + 0.21849238872528076, + 0.14441217482089996, + 0.18882252275943756, + 0.24324765801429749, + 0.24351975321769714, + 0.4946439862251282, + 0.2044689953327179, + 0.21101945638656616, + 0.20001566410064697, + 0.21270446479320526, + 0.10073935985565186, + 0.10901366174221039, + 0.15311172604560852, + 0.33192506432533264, + 0.0932217687368393, + 0.2565520107746124, + 0.19152018427848816, + 0.29205450415611267, + 0.32224786281585693, + 0.20979607105255127, + 0.17469346523284912, + 0.18356575071811676, + 0.3293987810611725, + 0.15834423899650574, + 0.33362191915512085, + 0.07346439361572266, + 0.20899252593517303, + 0.3504052758216858, + 0.20313145220279694, + 0.1570628136396408, + 0.16724461317062378, + 0.20873689651489258, + 0.3407520055770874, + 0.1537453830242157, + 0.2302379012107849, + 0.5382002592086792, + 0.10882124304771423, + 0.19074544310569763, + 0.17153841257095337, + 0.34962475299835205, + 0.21258226037025452, + 0.28493043780326843, + 0.18394821882247925, + 0.15388810634613037, + 0.17971481382846832, + 0.21412378549575806, + 0.22890782356262207, + 0.17444974184036255, + 0.12400753051042557, + 0.43169695138931274, + 0.30906444787979126, + 0.2431851327419281, + 0.14230023324489594, + 0.16815851628780365, + 0.28706690669059753, + 0.21689322590827942, + 0.17032954096794128, + 0.09686072915792465, + 0.18769985437393188, + 0.15629172325134277, + 0.24631761014461517, + 0.14470185339450836, + 0.1319996565580368, + 0.1697579324245453, + 0.14015576243400574, + 0.34668782353401184, + 0.2829791307449341, + 0.32119104266166687, + 0.166100412607193, + 0.19517366588115692, + 0.17637331783771515, + 0.2045217752456665, + 0.3483165502548218, + 0.17778436839580536, + 0.11767259985208511, + 0.3345232903957367, + 0.10189173370599747, + 0.19911697506904602, + 0.2261834591627121, + 0.30660879611968994, + 0.26291999220848083, + 0.255820631980896, + 0.16282325983047485, + 0.21763025224208832, + 0.22005556523799896, + 0.21814164519309998, + 0.15567004680633545, + 0.2112916111946106, + 0.1573236882686615, + 0.4854563772678375, + 0.3495527505874634, + 0.2206605225801468, + 0.17529039084911346, + 0.12122999131679535, + 0.169167160987854, + 0.40050578117370605, + 0.10340277850627899, + 0.21680369973182678, + 0.22093018889427185, + 0.3127893805503845, + 0.3957228660583496, + 0.1574784517288208, + 0.22071340680122375, + 0.11627742648124695, + 0.12477922439575195, + 0.23878450691699982, + 0.37244874238967896, + 0.4521469175815582, + 0.2336120456457138, + 0.10709106922149658, + 0.12841719388961792, + 0.165939599275589, + 0.4502921998500824, + 0.21725548803806305, + 0.19429010152816772, + 0.24550093710422516, + 0.23136398196220398, + 0.24719560146331787, + 0.12115201354026794, + 0.16507357358932495, + 0.15771445631980896, + 0.11222188174724579, + 0.2305431067943573, + 0.12510262429714203, + 0.2887539267539978, + 0.11430251598358154, + 0.1699393391609192, + 0.2553725838661194, + 0.1813780814409256, + 0.2520446181297302, + 0.0832974761724472, + 0.15081369876861572, + 0.10174526274204254, + 0.19539082050323486, + 0.20088836550712585, + 0.11679691821336746, + 0.12679752707481384, + 0.19633308053016663, + 0.22153757512569427, + 0.18450722098350525, + 0.19877946376800537, + 0.13296474516391754, + 0.22597774863243103, + 0.3281766474246979, + 0.1518074870109558, + 0.13902656733989716, + 0.2926194667816162, + 0.318151593208313, + 0.09639184176921844, + 0.1861957460641861, + 0.3367818593978882, + 0.2436227798461914, + 0.22163806855678558, + 0.16963067650794983, + 0.3087312877178192, + 0.1517033576965332, + 0.0598820261657238, + 0.2195756733417511, + 0.5442951917648315, + 0.3872811198234558, + 0.38659587502479553, + 0.30112579464912415, + 0.14171651005744934, + 0.18891847133636475, + 0.22281941771507263, + 0.1154424250125885, + 0.25115522742271423, + 0.11654698848724365, + 0.3808698058128357, + 0.1536589413881302, + 0.15220147371292114, + 0.16598175466060638, + 0.16175812482833862, + 0.2272213250398636, + 0.34597501158714294, + 0.3287469744682312, + 0.0949503481388092, + 0.1717933714389801, + 0.277523934841156, + 0.32820820808410645, + 0.15895476937294006, + 0.27238526940345764, + 0.09991219639778137, + 0.15354035794734955, + 0.12822912633419037, + 0.17543716728687286, + 0.09801562875509262, + 0.3501811921596527, + 0.22185653448104858, + 0.16663649678230286, + 0.11639758199453354, + 0.21506398916244507, + 0.09377127885818481, + 0.15584370493888855, + 0.25092247128486633, + 0.19597795605659485, + 0.1393977552652359, + 0.12320505827665329, + 0.2595716714859009, + 0.11877229064702988, + 0.14012163877487183, + 0.1736333966255188, + 0.162780299782753, + 0.18263453245162964, + 0.05860624462366104, + 0.27713480591773987, + 0.12189661711454391, + 0.14192073047161102, + 0.18259945511817932, + 0.22012211382389069, + 0.17633800208568573, + 0.18606817722320557, + 0.0938413143157959, + 0.22706685960292816, + 0.20159754157066345, + 0.18502011895179749, + 0.26651731133461, + 0.08840206265449524, + 0.09878166764974594, + 0.12373179197311401, + 0.2040935456752777, + 0.08315481245517731, + 0.07817058265209198, + 0.16869021952152252, + 0.147806316614151, + 0.16196930408477783, + 0.19217976927757263, + 0.25625643134117126, + 0.15830352902412415, + 0.17438359558582306, + 0.1393635869026184, + 0.09926897287368774, + 0.281377375125885, + 0.3770599365234375, + 0.17113488912582397, + 0.12046255171298981, + 0.0934320017695427, + 0.16895517706871033, + 0.33458009362220764, + 0.08919041603803635, + 0.18870826065540314, + 0.33836913108825684, + 0.2444867193698883, + 0.32291266322135925, + 0.17069002985954285, + 0.5397546887397766, + 0.1051880419254303, + 0.13954783976078033, + 0.14877572655677795, + 0.1248636245727539, + 0.15857498347759247, + 0.12324601411819458, + 0.44793951511383057, + 0.23582050204277039, + 0.2497071772813797, + 0.1289464384317398, + 0.2268783450126648, + 0.20395749807357788, + 0.20491282641887665, + 0.2831297516822815, + 0.20297175645828247, + 0.13873273134231567, + 0.16087259352207184, + 0.17212827503681183, + 0.0953991562128067, + 0.3898247480392456, + 0.3032417893409729, + 0.11175519973039627, + 0.2875506281852722, + 0.21398046612739563, + 0.12830811738967896, + 0.14083968102931976, + 0.2972894012928009, + 0.21215909719467163, + 0.44271472096443176, + 0.3088798522949219, + 0.2045627236366272, + 0.2547227740287781, + 0.3200996518135071, + 0.06704965978860855, + 0.12912437319755554, + 0.27107369899749756, + 0.12160750478506088, + 0.20084941387176514, + 0.12739898264408112, + 0.3335071802139282, + 0.10009977966547012, + 0.07825015485286713, + 0.17058798670768738, + 0.2522132694721222, + 0.231123149394989, + 0.14109660685062408, + 0.1485511213541031, + 0.255191445350647, + 0.05750148743391037, + 0.1887274980545044, + 0.1474664807319641, + 0.2115415781736374, + 0.2081030160188675, + 0.15634068846702576, + 0.3079553246498108, + 0.14709870517253876, + 0.051647305488586426, + 0.29970788955688477, + 0.1966628134250641, + 0.08855423331260681, + 0.23740273714065552, + 0.09980128705501556, + 0.09843023121356964, + 0.3231455087661743, + 0.08146949112415314, + 0.1677493304014206, + 0.3841918110847473, + 0.17269422113895416, + 0.21422548592090607, + 0.28332650661468506, + 0.21903249621391296, + 0.18857917189598083, + 0.20956271886825562, + 0.06858329474925995, + 0.11355310678482056, + 0.32532286643981934, + 0.12968727946281433, + 0.46097710728645325, + 0.17285819351673126, + 0.40199583768844604, + 0.3213024437427521, + 0.15420985221862793, + 0.13903436064720154, + 0.1781727522611618, + 0.13599181175231934, + 0.24453938007354736, + 0.25450241565704346, + 0.12214037030935287, + 0.26525697112083435, + 0.17957375943660736, + 0.3000493049621582, + 0.30755865573883057, + 0.23994497954845428, + 0.2168225347995758, + 0.20643699169158936, + 0.31324630975723267, + 0.22658076882362366, + 0.17479825019836426, + 0.18958501517772675, + 0.2628447413444519, + 0.15188488364219666, + 0.2375255525112152, + 0.11848811060190201, + 0.19174638390541077, + 0.19236096739768982, + 0.28497302532196045, + 0.2083507627248764, + 0.23337501287460327, + 0.12267977744340897, + 0.2667624056339264, + 0.1593530774116516, + 0.11240258067846298, + 0.278443843126297, + 0.25114259123802185, + 0.1674422323703766, + 0.1352960765361786, + 0.12894141674041748, + 0.17218917608261108, + 0.12394675612449646, + 0.253647118806839, + 0.10290729254484177, + 0.14798876643180847, + 0.1527380645275116, + 0.10975228250026703, + 0.2098216861486435, + 0.16819310188293457, + 0.07962843775749207, + 0.11957383900880814, + 0.11788206547498703, + 0.12337764352560043, + 0.17604731023311615, + 0.09305307269096375, + 0.16893400251865387, + 0.14240555465221405, + 0.2593865990638733, + 0.15986615419387817, + 0.10574576258659363, + 0.10629955679178238, + 0.1064593493938446, + 0.17416101694107056, + 0.147480309009552, + 0.13269861042499542, + 0.1620386838912964, + 0.18665000796318054, + 0.297687828540802, + 0.14625826478004456, + 0.13938573002815247, + 0.17332346737384796, + 0.27558520436286926, + 0.18871521949768066, + 0.15504923462867737, + 0.24520517885684967, + 0.19831793010234833, + 0.10882744938135147 + ] +} \ No newline at end of file diff --git a/workspace/outputs/external_vla/smolvla_cil_balanced_metrics.json b/workspace/outputs/external_vla/smolvla_cil_balanced_metrics.json new file mode 100644 index 0000000000000000000000000000000000000000..ca90473ead003534887a798f5b842655882cd516 --- /dev/null +++ b/workspace/outputs/external_vla/smolvla_cil_balanced_metrics.json @@ -0,0 +1,47 @@ +{ + "candidate_oracle_success_rate": 0.4533333333333333, + "evaluation_protocol": "nearest_executed_same_state_candidate", + "final_train_loss": 0.10882744938135147, + "mean_selected_regret": 0.17654996168634776, + "model_family": "smolvla", + "normalized_action_mse_to_selected": 0.3356558305242409, + "num_eval_groups": 600, + "num_train_groups": 2400, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 100, + "selected_reward_mean": 1.1786985996365547, + "selected_success_rate": 0.47 + }, + "PegInsertionSide-v1": { + "num_groups": 100, + "selected_reward_mean": 0.2845203054515878, + "selected_success_rate": 0.01 + }, + "PickCube-v1": { + "num_groups": 100, + "selected_reward_mean": 0.8007295456156135, + "selected_success_rate": 0.21 + }, + "PullCube-v1": { + "num_groups": 100, + "selected_reward_mean": 1.3391107024252416, + "selected_success_rate": 0.61 + }, + "PushCube-v1": { + "num_groups": 100, + "selected_reward_mean": 1.2629673214256762, + "selected_success_rate": 0.59 + }, + "StackCube-v1": { + "num_groups": 100, + "selected_reward_mean": 0.8535422113537788, + "selected_success_rate": 0.25 + } + }, + "selected_reward_mean": 0.9532614476514087, + "selected_success_rate": 0.3566666666666667, + "top1_action_selection": 0.5416666666666666, + "training_mode": "expert_only_bc", + "training_steps": 1000 +} \ No newline at end of file diff --git a/workspace/outputs/external_vla/smolvla_cil_smoke_metrics.json b/workspace/outputs/external_vla/smolvla_cil_smoke_metrics.json new file mode 100644 index 0000000000000000000000000000000000000000..40bcdc33cf312c197ac2c85390f65cca09dbe2de --- /dev/null +++ b/workspace/outputs/external_vla/smolvla_cil_smoke_metrics.json @@ -0,0 +1,47 @@ +{ + "candidate_oracle_success_rate": 0.5, + "evaluation_protocol": "nearest_executed_same_state_candidate", + "final_train_loss": 0.8980979323387146, + "mean_selected_regret": 0.8765356718795374, + "model_family": "smolvla", + "normalized_action_mse_to_selected": 0.838577705125014, + "num_eval_groups": 12, + "num_train_groups": 52, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 2, + "selected_reward_mean": 0.5819069147109985, + "selected_success_rate": 0.0 + }, + "PegInsertionSide-v1": { + "num_groups": 2, + "selected_reward_mean": 0.03161212173290551, + "selected_success_rate": 0.0 + }, + "PickCube-v1": { + "num_groups": 2, + "selected_reward_mean": 0.5634382367134094, + "selected_success_rate": 0.0 + }, + "PullCube-v1": { + "num_groups": 2, + "selected_reward_mean": 0.14262327924370766, + "selected_success_rate": 0.0 + }, + "PushCube-v1": { + "num_groups": 2, + "selected_reward_mean": 0.07729745656251907, + "selected_success_rate": 0.0 + }, + "StackCube-v1": { + "num_groups": 2, + "selected_reward_mean": 0.6672812402248383, + "selected_success_rate": 0.0 + } + }, + "selected_reward_mean": 0.3440265415313964, + "selected_success_rate": 0.0, + "top1_action_selection": 0.0, + "training_mode": "expert_only_bc", + "training_steps": 2 +} \ No newline at end of file diff --git a/workspace/outputs/external_vla_export_maniskill_full_no_images/metadata.json b/workspace/outputs/external_vla_export_maniskill_full_no_images/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..e12a34d65e205d27e944fde71a1d6bc91e0bf411 --- /dev/null +++ b/workspace/outputs/external_vla_export_maniskill_full_no_images/metadata.json @@ -0,0 +1,12 @@ +{ + "action_encoding": "flattened numeric ActionChunk values plus full action_chunk payload", + "copy_images": false, + "external_vla_note": "This is a dependency-light interchange export for external VLA baselines. It is not a measured SmolVLA/OpenVLA result by itself.", + "image_format": "jpg", + "num_episodes": 3500, + "num_tasks": 6, + "schema_version": "dovla-cil-lerobot-export/v0", + "selection": "best", + "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "train" +} diff --git a/workspace/outputs/external_vla_export_maniskill_full_no_images/tasks.jsonl b/workspace/outputs/external_vla_export_maniskill_full_no_images/tasks.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..e7a7a115aff19acc209a6d98075e0e6fb1ac7d1a --- /dev/null +++ b/workspace/outputs/external_vla_export_maniskill_full_no_images/tasks.jsonl @@ -0,0 +1,6 @@ +{"instructions": ["Lift the peg and hold it upright."], "task_id": "LiftPegUpright-v1", "task_index": 0} +{"instructions": ["Insert the peg into the hole from the side."], "task_id": "PegInsertionSide-v1", "task_index": 1} +{"instructions": ["Pick up the cube and move it to the goal position."], "task_id": "PickCube-v1", "task_index": 2} +{"instructions": ["Pull the cube into the goal region."], "task_id": "PullCube-v1", "task_index": 3} +{"instructions": ["Push the cube into the goal region."], "task_id": "PushCube-v1", "task_index": 4} +{"instructions": ["Stack cube A on top of cube B."], "task_id": "StackCube-v1", "task_index": 5} diff --git a/workspace/outputs/external_vla_export_maniskill_full_no_images/train.jsonl b/workspace/outputs/external_vla_export_maniskill_full_no_images/train.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..c89bfe55ed7030d124f63f585282d2dc7b3811e8 --- /dev/null +++ b/workspace/outputs/external_vla_export_maniskill_full_no_images/train.jsonl @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:f8ee77311107c226c63b0d6aa936ed067d1f95d57578fd4c2fcb583489b036c2 +size 11038222 diff --git a/workspace/outputs/external_vla_export_smoke/metadata.json b/workspace/outputs/external_vla_export_smoke/metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..0a209e2dc4bd2dc8188a8f6059b3275d3514225c --- /dev/null +++ b/workspace/outputs/external_vla_export_smoke/metadata.json @@ -0,0 +1,12 @@ +{ + "action_encoding": "flattened numeric ActionChunk values plus full action_chunk payload", + "copy_images": false, + "external_vla_note": "This is a dependency-light interchange export for external VLA baselines. It is not a measured SmolVLA/OpenVLA result by itself.", + "image_format": "jpg", + "num_episodes": 12, + "num_tasks": 1, + "schema_version": "dovla-cil-lerobot-export/v0", + "selection": "best", + "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "train" +} diff --git a/workspace/outputs/external_vla_export_smoke/tasks.jsonl b/workspace/outputs/external_vla_export_smoke/tasks.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..3b93a75431e62da2a926f92c965f85bc7d33c460 --- /dev/null +++ b/workspace/outputs/external_vla_export_smoke/tasks.jsonl @@ -0,0 +1 @@ +{"instructions": ["Pick up the cube and move it to the goal position."], "task_id": "PickCube-v1", "task_index": 0} diff --git a/workspace/outputs/external_vla_export_smoke/train.jsonl b/workspace/outputs/external_vla_export_smoke/train.jsonl new file mode 100644 index 0000000000000000000000000000000000000000..6d49248d3d0bb890cf660bc16bb6b0345287036e --- /dev/null +++ b/workspace/outputs/external_vla_export_smoke/train.jsonl @@ -0,0 +1,12 @@ +{"action": [-1.5628260374069214, -0.6760064363479614, -1.3670138120651245, 0.011951039545238018, -1.9163702726364136, -0.16114844381809235, 4.311707973480225, -0.9735527634620667, -0.7978972792625427, -1.2788859605789185, 0.10670334845781326, -1.6345019340515137, -0.05739933252334595, 4.201030254364014, 0.1259092390537262, -0.28991422057151794, -1.5478920936584473, -0.1543567031621933, -1.1837780475616455, -0.29051971435546875, 3.35263991355896, 0.18668514490127563, 0.30432724952697754, -1.7003849744796753, -0.2195604294538498, -0.8814301490783691, -0.7409264445304871, 1.4718676805496216], "action_chunk": {"action_id": "expert-00", "horizon": 4, "metadata": {"candidate_type": "expert", "parent_action_id": null, "perturbation": {"scale": 0.0}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[-1.5628260374069214, -0.6760064363479614, -1.3670138120651245, 0.011951039545238018, -1.9163702726364136, -0.16114844381809235, 4.311707973480225], [-0.9735527634620667, -0.7978972792625427, -1.2788859605789185, 0.10670334845781326, -1.6345019340515137, -0.05739933252334595, 4.201030254364014], [0.1259092390537262, -0.28991422057151794, -1.5478920936584473, -0.1543567031621933, -1.1837780475616455, -0.29051971435546875, 3.35263991355896], [0.18668514490127563, 0.30432724952697754, -1.7003849744796753, -0.2195604294538498, -0.8814301490783691, -0.7409264445304871, 1.4718676805496216]]}, "cil": {"candidate_type": "expert", "failure_type": "partial_success", "group_id": "ms-4d4a0c1d076355ffd0a1", "rank_within_group": 0, "record_id": "rec-5369e5cfdb2f93b06805f34f", "regret": 0.0, "scene_id": "episode-1020", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "278644b6ee28dc010fcab4b93e73e5db6191ac912db250f5290b2a1aaa37dd7f", "task_id": "PickCube-v1"}, "episode_index": 0, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/0", "image": null, "state": {"features": [-0.0541839599609375, -0.08368682861328125, 0.019999999552965164, 0.9984005093574524, 0.0, 0.0, -0.056537654250860214, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.048946380615234375, 0.0962677001953125, 0.25899022817611694, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.11999893188476562, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6150016784667969, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.03468818590044975, 0.35614410042762756, -0.009833138436079025, -1.936769962310791, -0.002558882115408778, 2.3531911373138428, 0.7486944198608398, 0.03999999910593033, 0.03999999910593033, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]}}, "reward": 0.15335801243782043, "success": false, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [-0.40443557500839233, 0.15195327997207642, 0.3294371962547302, -0.5992763638496399, -0.7588088512420654, -0.49525174498558044, -2.4462406635284424, -0.30573686957359314, 0.04093547537922859, 0.2394212782382965, -0.3384755551815033, -0.6182968020439148, -0.4414074420928955, -2.423753261566162, -0.252510666847229, 0.03575673699378967, 0.03673635423183441, -0.09818130731582642, -0.4163699448108673, -0.18765880167484283, -2.4663844108581543, -0.1633390337228775, 0.09333082288503647, -0.03681729733943939, -0.14803244173526764, -0.23525601625442505, -0.27691468596458435, -2.556912422180176], "action_chunk": {"action_id": "near_miss-01", "horizon": 4, "metadata": {"candidate_type": "near_miss", "parent_action_id": "expert", "perturbation": {"kind": "gaussian_action_offset", "scale": 0.05}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[-0.40443557500839233, 0.15195327997207642, 0.3294371962547302, -0.5992763638496399, -0.7588088512420654, -0.49525174498558044, -2.4462406635284424], [-0.30573686957359314, 0.04093547537922859, 0.2394212782382965, -0.3384755551815033, -0.6182968020439148, -0.4414074420928955, -2.423753261566162], [-0.252510666847229, 0.03575673699378967, 0.03673635423183441, -0.09818130731582642, -0.4163699448108673, -0.18765880167484283, -2.4663844108581543], [-0.1633390337228775, 0.09333082288503647, -0.03681729733943939, -0.14803244173526764, -0.23525601625442505, -0.27691468596458435, -2.556912422180176]]}, "cil": {"candidate_type": "near_miss", "failure_type": "partial_success", "group_id": "ms-bca92beb7317318ca011", "rank_within_group": 0, "record_id": "rec-c84aaef7d7f27c87cda914f7", "regret": 0.0, "scene_id": "episode-239", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "bcf14b521690851540ef326b8896075a439d2acc289bb49f5792e93926d54bcb", "task_id": "PickCube-v1"}, "episode_index": 1, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/1", "image": null, "state": {"features": [-0.019093990325927734, 0.02045440673828125, 0.07436220347881317, 0.03399929031729698, -0.04998345673084259, 0.044870324432849884, -0.9971621036529541, -0.19499041140079498, 0.06917788833379745, 0.46193182468414307, -7.807628631591797, -0.9009962677955627, 8.998115539550781, -0.05489015579223633, 0.031215667724609375, 0.11114609241485596, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.11999988555908203, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6149997711181641, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.06863383203744888, 0.5815295577049255, 0.07980771362781525, -1.9120534658432007, 0.11882027238607407, 2.3517379760742188, 0.6964728236198425, 0.018981605768203735, 0.019019396975636482, -0.21386092901229858, -1.1971880197525024, 0.29311615228652954, -0.7713239789009094, 0.27112236618995667, -0.7388791441917419, -0.30363333225250244, 0.08629126101732254, 0.062452446669340134]}}, "reward": 0.6064264178276062, "success": false, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [0.25490814447402954, -0.043703995645046234, 1.0, -1.0, -0.7366626858711243, -1.0, -2.906116008758545, 0.3121466040611267, -0.019728688523173332, 1.0, -1.0, -0.7567026019096375, -1.0, -2.6152353286743164, 0.13707508146762848, -0.046043578535318375, 1.0, -0.9888929128646851, -0.6925053596496582, -0.9358696937561035, -2.3858416080474854, -0.09359025210142136, -0.02185366116464138, 0.9880052208900452, -0.7964281439781189, -0.6394259929656982, -0.6999186873435974, -2.129840850830078], "action_chunk": {"action_id": "near_miss-01", "horizon": 4, "metadata": {"candidate_type": "near_miss", "parent_action_id": "expert", "perturbation": {"kind": "gaussian_action_offset", "scale": 0.05}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[0.25490814447402954, -0.043703995645046234, 1.0, -1.0, -0.7366626858711243, -1.0, -2.906116008758545], [0.3121466040611267, -0.019728688523173332, 1.0, -1.0, -0.7567026019096375, -1.0, -2.6152353286743164], [0.13707508146762848, -0.046043578535318375, 1.0, -0.9888929128646851, -0.6925053596496582, -0.9358696937561035, -2.3858416080474854], [-0.09359025210142136, -0.02185366116464138, 0.9880052208900452, -0.7964281439781189, -0.6394259929656982, -0.6999186873435974, -2.129840850830078]]}, "cil": {"candidate_type": "near_miss", "failure_type": "partial_success", "group_id": "ms-abe94723b59a97d9763f", "rank_within_group": 0, "record_id": "rec-4c195d41b3ac570b766b9d11", "regret": 0.0, "scene_id": "episode-587", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "290a189c6f3584d08e1d1807e63618d6cea3903638a478cdc87086339404400b", "task_id": "PickCube-v1"}, "episode_index": 2, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/2", "image": null, "state": {"features": [-0.007940292358398438, 0.0969095230102539, 0.03747957572340965, 0.6609684228897095, -0.01142940018326044, 0.024942221119999886, -0.749911904335022, -0.0038174018263816833, -0.09276441484689713, 0.7106754779815674, -10.978766441345215, 1.8465118408203125, -18.50157356262207, -0.016963958740234375, 0.07655715942382812, 0.25319740176200867, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.12000083923339844, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6149997711181641, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.03658159449696541, 0.7031791806221008, 0.10627271234989166, -1.815037488937378, 0.05101632699370384, 2.374300479888916, 0.7581173181533813, 0.020984023809432983, 0.02102568745613098, -0.476231187582016, -1.1705760955810547, 0.3279920518398285, -0.02065674588084221, 0.18396227061748505, -1.391904592514038, -0.7702826261520386, 0.24117586016654968, 0.2134247124195099]}}, "reward": 0.536004364490509, "success": false, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [-0.09155034273862839, -0.2224860042333603, 0.09527984261512756, -0.33472195267677307, -0.10096851736307144, -0.32266828417778015, -2.021345615386963, -0.08139925450086594, 0.23562990128993988, -0.11173506081104279, -0.3661971092224121, -0.1256498098373413, -0.04269513115286827, -2.037240505218506, -0.08127366006374359, 0.025128714740276337, -0.21820373833179474, 0.13546231389045715, 0.01720122992992401, -0.147305428981781, -2.048769474029541, -0.11806996166706085, 0.007225414272397757, -0.0036315745674073696, -0.0010948525741696358, -0.254454106092453, -0.05180085822939873, -2.0574095249176025], "action_chunk": {"action_id": "near_miss-02", "horizon": 4, "metadata": {"candidate_type": "near_miss", "parent_action_id": "expert", "perturbation": {"kind": "gaussian_action_offset", "scale": 0.15}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[-0.09155034273862839, -0.2224860042333603, 0.09527984261512756, -0.33472195267677307, -0.10096851736307144, -0.32266828417778015, -2.021345615386963], [-0.08139925450086594, 0.23562990128993988, -0.11173506081104279, -0.3661971092224121, -0.1256498098373413, -0.04269513115286827, -2.037240505218506], [-0.08127366006374359, 0.025128714740276337, -0.21820373833179474, 0.13546231389045715, 0.01720122992992401, -0.147305428981781, -2.048769474029541], [-0.11806996166706085, 0.007225414272397757, -0.0036315745674073696, -0.0010948525741696358, -0.254454106092453, -0.05180085822939873, -2.0574095249176025]]}, "cil": {"candidate_type": "near_miss", "failure_type": "success", "group_id": "ms-c4ffabbc14eaeffb221e", "rank_within_group": 0, "record_id": "rec-17e7ffd2176a3d740eb05af9", "regret": 0.0, "scene_id": "episode-209", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "ec61a6991f978a5aad935d8f2ca14e4a866cd71a661dab4d8a259dd57fdb5531", "task_id": "PickCube-v1"}, "episode_index": 3, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/3", "image": null, "state": {"features": [-0.046315670013427734, 0.035221099853515625, 0.2322583645582199, 0.7701006531715393, 0.010515510104596615, 0.09762848168611526, -0.6303197145462036, -0.052918922156095505, -0.040496163070201874, 0.010721586644649506, -7.1247663497924805, -1.63589608669281, 0.47000962495803833, -0.046441078186035156, 0.0363922119140625, 0.231855571269989, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.11999988555908203, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6149997711181641, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.10235677659511566, 0.3599659204483032, 0.13238002359867096, -1.808642029762268, 0.10925418883562088, 1.8546810150146484, 0.56473708152771, 0.01857931539416313, 0.018624702468514442, -0.03570203855633736, -0.10249791294336319, -0.01871495507657528, -0.20577512681484222, -0.003970427438616753, 0.02155376225709915, -0.036965273320674896, 0.04108690097928047, -0.007286506704986095]}}, "reward": 2.0, "success": true, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [0.8234550952911377, -0.39333590865135193, -2.045659065246582, 0.2834421694278717, -0.29677435755729675, 0.2694389224052429, 2.9280712604522705, 0.10313728451728821, -0.22221754491329193, -1.9722038507461548, 0.15917308628559113, -0.5743748545646667, -0.18963240087032318, 1.7628428936004639, -0.1889944076538086, 0.23151889443397522, -1.1360738277435303, -0.24781058728694916, -0.34082698822021484, -0.170538991689682, -0.9756835699081421, -0.5063692927360535, 0.03460731357336044, 0.7289126515388489, -0.6656671762466431, -0.5738135576248169, -0.25992345809936523, -3.7129390239715576], "action_chunk": {"action_id": "expert-00", "horizon": 4, "metadata": {"candidate_type": "expert", "parent_action_id": null, "perturbation": {"scale": 0.0}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[0.8234550952911377, -0.39333590865135193, -2.045659065246582, 0.2834421694278717, -0.29677435755729675, 0.2694389224052429, 2.9280712604522705], [0.10313728451728821, -0.22221754491329193, -1.9722038507461548, 0.15917308628559113, -0.5743748545646667, -0.18963240087032318, 1.7628428936004639], [-0.1889944076538086, 0.23151889443397522, -1.1360738277435303, -0.24781058728694916, -0.34082698822021484, -0.170538991689682, -0.9756835699081421], [-0.5063692927360535, 0.03460731357336044, 0.7289126515388489, -0.6656671762466431, -0.5738135576248169, -0.25992345809936523, -3.7129390239715576]]}, "cil": {"candidate_type": "expert", "failure_type": "partial_success", "group_id": "ms-b242270287e07877b824", "rank_within_group": 0, "record_id": "rec-61e1134c676f801a337c6a73", "regret": 0.0, "scene_id": "episode-797", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "51a43c60080582d741216ac5b68ca98751c8bda2c713848ad8a746384524fba9", "task_id": "PickCube-v1"}, "episode_index": 4, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/4", "image": null, "state": {"features": [0.06255340576171875, -0.02100372314453125, 0.019999802112579346, 0.8229450583457947, 9.91355500445934e-06, -5.571292149397777e-06, 0.5681209564208984, 0.0004370586830191314, 0.00017391589062754065, -0.00011650443775579333, -0.004854163154959679, 0.010291452519595623, 0.00042810323066078126, -0.07491302490234375, 0.0097503662109375, 0.15170276165008545, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.12000274658203125, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6149978637695312, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.01807539351284504, 0.5113353729248047, -0.01832074299454689, -1.830978274345398, 0.026638031005859375, 2.371656894683838, 0.7861020565032959, 0.039997249841690063, 0.03999991714954376, -0.14526815712451935, 2.625046491622925, -0.057904332876205444, 2.40586519241333, 0.088222436606884, 0.07613790035247803, -0.02440214529633522, 0.0015209796838462353, -0.0014612110098823905]}}, "reward": 0.46302175521850586, "success": false, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [0.7732240557670593, -0.29269447922706604, 0.479524165391922, -0.4287649691104889, -0.03655378147959709, -1.0, -2.112894296646118, 0.3012737035751343, -0.23505425453186035, 0.11112562566995621, -0.3476502299308777, 0.16826392710208893, -0.7720427513122559, -2.110785722732544, 0.12148325145244598, -0.2086387574672699, -0.18725250661373138, -0.16820640861988068, 0.35087722539901733, -0.4146888256072998, -2.215219497680664, 0.020610615611076355, -0.15581557154655457, 0.010943256318569183, 0.008966654539108276, 0.11986343562602997, -0.12302891910076141, -2.268505811691284], "action_chunk": {"action_id": "near_miss-01", "horizon": 4, "metadata": {"candidate_type": "near_miss", "parent_action_id": "expert", "perturbation": {"kind": "gaussian_action_offset", "scale": 0.05}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[0.7732240557670593, -0.29269447922706604, 0.479524165391922, -0.4287649691104889, -0.03655378147959709, -1.0, -2.112894296646118], [0.3012737035751343, -0.23505425453186035, 0.11112562566995621, -0.3476502299308777, 0.16826392710208893, -0.7720427513122559, -2.110785722732544], [0.12148325145244598, -0.2086387574672699, -0.18725250661373138, -0.16820640861988068, 0.35087722539901733, -0.4146888256072998, -2.215219497680664], [0.020610615611076355, -0.15581557154655457, 0.010943256318569183, 0.008966654539108276, 0.11986343562602997, -0.12302891910076141, -2.268505811691284]]}, "cil": {"candidate_type": "near_miss", "failure_type": "partial_success", "group_id": "ms-c780d41a3bffb09815d1", "rank_within_group": 0, "record_id": "rec-9172a3c256656c2315553667", "regret": 0.0, "scene_id": "episode-114", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "21fe9927c262557f34f7f1984949425ca7b9f11187f40105b653bdae2daba930", "task_id": "PickCube-v1"}, "episode_index": 5, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/5", "image": null, "state": {"features": [0.028700828552246094, 0.0655059814453125, 0.14895597100257874, 0.65125972032547, 0.03211274370551109, 0.00015520700253546238, 0.7581750750541687, 0.5814969539642334, -0.19313077628612518, 0.6927828788757324, -1.148125410079956, -0.012501949444413185, -7.615335464477539, 0.09989356994628906, 0.00606536865234375, 0.18480096757411957, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.11999988555908203, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6149988174438477, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.008372840471565723, 0.6217747926712036, 0.09794788062572479, -1.6184884309768677, -0.014663946814835072, 2.0312952995300293, 0.6737716197967529, 0.01865791715681553, 0.018702059984207153, -0.3444288671016693, 0.5771565437316895, 0.04802849143743515, 2.894660234451294, 0.012401988729834557, -2.1706244945526123, -0.907012939453125, 0.0570286326110363, 0.00033069640630856156]}}, "reward": 0.5932891964912415, "success": false, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [-0.023983854800462723, 0.018745146691799164, -0.021289434283971786, -0.11246635019779205, -0.11601042747497559, -0.2676311731338501, -1.8741012811660767, -0.014467019587755203, 0.019133679568767548, -0.01631760224699974, -0.1042390987277031, -0.10641993582248688, -0.2710961103439331, -1.8810278177261353, -0.008381742984056473, 0.01854727417230606, -0.012751217931509018, -0.09557168930768967, -0.09803362935781479, -0.2704295516014099, -1.8868783712387085, -0.00460154190659523, 0.017910249531269073, -0.010318022221326828, -0.08704736828804016, -0.0906854197382927, -0.26731300354003906, -1.8920661211013794], "action_chunk": {"action_id": "expert-00", "horizon": 4, "metadata": {"candidate_type": "expert", "parent_action_id": null, "perturbation": {"scale": 0.0}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[-0.023983854800462723, 0.018745146691799164, -0.021289434283971786, -0.11246635019779205, -0.11601042747497559, -0.2676311731338501, -1.8741012811660767], [-0.014467019587755203, 0.019133679568767548, -0.01631760224699974, -0.1042390987277031, -0.10641993582248688, -0.2710961103439331, -1.8810278177261353], [-0.008381742984056473, 0.01854727417230606, -0.012751217931509018, -0.09557168930768967, -0.09803362935781479, -0.2704295516014099, -1.8868783712387085], [-0.00460154190659523, 0.017910249531269073, -0.010318022221326828, -0.08704736828804016, -0.0906854197382927, -0.26731300354003906, -1.8920661211013794]]}, "cil": {"candidate_type": "expert", "failure_type": "success", "group_id": "ms-0cd563cc20f169c74fbe", "rank_within_group": 0, "record_id": "rec-b7bd469f454d7b7ce93ef153", "regret": 0.0, "scene_id": "episode-224", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "db82c86094fcb2a6a1c5c35ce1b7683ef731d6155614dd272a01038828a78fac", "task_id": "PickCube-v1"}, "episode_index": 6, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/6", "image": null, "state": {"features": [-0.074432373046875, 0.040019989013671875, 0.22018413245677948, 0.7794306874275208, 0.1468936800956726, 0.054949868470430374, 0.6065397262573242, 0.19083064794540405, -0.018964337185025215, -0.040836405009031296, -1.8016589879989624, 6.599865436553955, -16.231740951538086, -0.0794677734375, 0.044940948486328125, 0.2078183889389038, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.12000274658203125, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6149978637695312, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.09387122094631195, 0.2581423819065094, 0.11969171464443207, -1.9873038530349731, 0.1517958641052246, 2.0032057762145996, 0.6415863633155823, 0.021779298782348633, 0.021827511489391327, -0.026836197823286057, 0.02816697396337986, 0.016907230019569397, -0.010281022638082504, 0.015941768884658813, -0.15004119277000427, -0.22097183763980865, 0.28425392508506775, 0.2691417932510376]}}, "reward": 2.0, "success": true, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [-0.19244690239429474, -0.16342812776565552, 0.0347924567759037, 0.008081558160483837, -0.31800031661987305, -0.18293285369873047, -1.8377646207809448, -0.18346264958381653, -0.03947979211807251, -0.06995131820440292, 0.11359677463769913, -0.2011098563671112, -0.02423825114965439, -1.8827391862869263, -0.17082858085632324, -0.08785596489906311, 0.0044371820986270905, 0.02421420067548752, -0.22217433154582977, -0.03988654538989067, -1.935221552848816, -0.09586358815431595, 0.05334766209125519, 0.029819149523973465, 0.09933511912822723, -0.1395062953233719, -0.18966631591320038, -1.964664340019226], "action_chunk": {"action_id": "near_miss-01", "horizon": 4, "metadata": {"candidate_type": "near_miss", "parent_action_id": "expert", "perturbation": {"kind": "gaussian_action_offset", "scale": 0.05}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[-0.19244690239429474, -0.16342812776565552, 0.0347924567759037, 0.008081558160483837, -0.31800031661987305, -0.18293285369873047, -1.8377646207809448], [-0.18346264958381653, -0.03947979211807251, -0.06995131820440292, 0.11359677463769913, -0.2011098563671112, -0.02423825114965439, -1.8827391862869263], [-0.17082858085632324, -0.08785596489906311, 0.0044371820986270905, 0.02421420067548752, -0.22217433154582977, -0.03988654538989067, -1.935221552848816], [-0.09586358815431595, 0.05334766209125519, 0.029819149523973465, 0.09933511912822723, -0.1395062953233719, -0.18966631591320038, -1.964664340019226]]}, "cil": {"candidate_type": "near_miss", "failure_type": "success", "group_id": "ms-8d3c88198175fb4ea4ea", "rank_within_group": 0, "record_id": "rec-65b8ecfcc6577ced35e51ee2", "regret": 0.0, "scene_id": "episode-388", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "95f5287d5c256c2538eb58b84bcc03c3ee2668891bfcc7b581b50b1632de7748", "task_id": "PickCube-v1"}, "episode_index": 7, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/7", "image": null, "state": {"features": [-0.054317474365234375, -0.00110626220703125, 0.20917671918869019, -0.1249999850988388, 0.06042565032839775, -0.04672222211956978, 0.9892120957374573, -0.060384850949048996, -0.05119608715176582, 0.26524391770362854, -4.279517650604248, -0.31354448199272156, -4.170334339141846, -0.0890350341796875, -0.0094451904296875, 0.2136704921722412, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.11999893188476562, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6150016784667969, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.10357502102851868, 0.365913450717926, 0.09339150786399841, -1.8553439378738403, 0.035346563905477524, 1.9232416152954102, 0.48678460717201233, 0.018391933292150497, 0.018436355516314507, -0.1446506828069687, -0.47221699357032776, 0.006858167238533497, 0.02454003319144249, 0.038342028856277466, -0.7332959175109863, -0.33226603269577026, 0.06523852795362473, 0.0018311217427253723]}}, "reward": 2.0, "success": true, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [1.0, 0.4642997086048126, -0.40524372458457947, -0.4181438982486725, -0.3855776786804199, -1.0, -2.572544813156128, 1.0, -0.12252479791641235, 0.3957115113735199, -0.08444324135780334, 0.13567838072776794, -1.0, -3.2739899158477783, 1.0, 0.5217642784118652, 0.9536948800086975, -0.7208852767944336, -0.08944196999073029, -1.0, -2.8353257179260254, 1.0, 0.3923838436603546, 1.0, -0.3802013397216797, 0.11167275905609131, -1.0, -2.5073773860931396], "action_chunk": {"action_id": "wrong_direction-03", "horizon": 4, "metadata": {"candidate_type": "wrong_direction", "parent_action_id": "expert", "perturbation": {"kind": "gaussian_action_offset", "scale": 0.35}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[1.0, 0.4642997086048126, -0.40524372458457947, -0.4181438982486725, -0.3855776786804199, -1.0, -2.572544813156128], [1.0, -0.12252479791641235, 0.3957115113735199, -0.08444324135780334, 0.13567838072776794, -1.0, -3.2739899158477783], [1.0, 0.5217642784118652, 0.9536948800086975, -0.7208852767944336, -0.08944196999073029, -1.0, -2.8353257179260254], [1.0, 0.3923838436603546, 1.0, -0.3802013397216797, 0.11167275905609131, -1.0, -2.5073773860931396]]}, "cil": {"candidate_type": "wrong_direction", "failure_type": "partial_success", "group_id": "ms-64d4fb96883ad9d71105", "rank_within_group": 0, "record_id": "rec-ccc7403fc6e42623eee0a0a6", "regret": 0.0, "scene_id": "episode-529", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "6eab3ca5da72c4b7b18fab3a7caab5924e786f3c037f38e4442e8392aec03b7b", "task_id": "PickCube-v1"}, "episode_index": 8, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/8", "image": null, "state": {"features": [-0.08880138397216797, -0.09159263968467712, 0.020002268254756927, 0.944753885269165, 0.0005589941283687949, 0.00012630359560716897, 0.32777976989746094, 0.05677035450935364, -0.021088257431983948, -0.010916968807578087, -0.2162894308567047, -0.21886613965034485, -2.150670289993286, 0.05281686782836914, -0.03405817598104477, 0.10896357893943787, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.11999988555908203, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6149997711181641, 4.3655745685100555e-11, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.1650100201368332, 0.5031988024711609, -0.013502894900739193, -2.2127530574798584, 0.047801319509744644, 2.6145236492156982, 0.5418534874916077, 0.024410586804151535, 0.02440524473786354, -0.051426924765110016, 2.3979947566986084, 0.4740931987762451, 1.5185127258300781, 0.45796918869018555, 0.50052809715271, -0.3085905611515045, -0.2204747200012207, -0.20427802205085754]}}, "reward": 0.5698418617248535, "success": false, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [0.3033374845981598, -0.12066788226366043, 0.14544619619846344, -0.3138177990913391, -0.0170342568308115, -0.7108187675476074, -2.2996127605438232, 0.05510043725371361, -0.09708938002586365, -0.13058404624462128, -0.11618771404027939, 0.097447969019413, -0.3188478350639343, -2.385507583618164, -0.006168145686388016, -0.09650762379169464, -0.10053984820842743, -0.022903304547071457, 0.019624365493655205, -0.1267804503440857, -2.499854564666748, -0.0005009658634662628, -0.04196876287460327, -0.060361217707395554, -0.03557731583714485, -0.014296675100922585, -0.13206668198108673, -2.5627455711364746], "action_chunk": {"action_id": "expert-00", "horizon": 4, "metadata": {"candidate_type": "expert", "parent_action_id": null, "perturbation": {"scale": 0.0}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[0.3033374845981598, -0.12066788226366043, 0.14544619619846344, -0.3138177990913391, -0.0170342568308115, -0.7108187675476074, -2.2996127605438232], [0.05510043725371361, -0.09708938002586365, -0.13058404624462128, -0.11618771404027939, 0.097447969019413, -0.3188478350639343, -2.385507583618164], [-0.006168145686388016, -0.09650762379169464, -0.10053984820842743, -0.022903304547071457, 0.019624365493655205, -0.1267804503440857, -2.499854564666748], [-0.0005009658634662628, -0.04196876287460327, -0.060361217707395554, -0.03557731583714485, -0.014296675100922585, -0.13206668198108673, -2.5627455711364746]]}, "cil": {"candidate_type": "expert", "failure_type": "success", "group_id": "ms-1781d9c562b4425df7c8", "rank_within_group": 0, "record_id": "rec-cf9e26e1a8f09382d2e0b7b1", "regret": 0.0, "scene_id": "episode-774", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "2e40294ac696fc7563ca5f74e761ce0ec43fa9ecc40c40c6aed5c9ebba87c54d", "task_id": "PickCube-v1"}, "episode_index": 9, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/9", "image": null, "state": {"features": [0.01476287841796875, -0.06786346435546875, 0.09191349893808365, 0.9993038177490234, 0.0220708679407835, 0.01285494863986969, -0.02719307877123356, 0.35307490825653076, -0.1941719800233841, 0.5391368269920349, -6.94796085357666, 0.545600950717926, -19.620615005493164, 0.0375518798828125, -0.09474945068359375, 0.10253944993019104, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.11999893188476562, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6150016784667969, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.15181684494018555, 0.6465469598770142, 0.039224348962306976, -1.7328109741210938, 0.027142634615302086, 2.2319157123565674, 0.46001049876213074, 0.02076961100101471, 0.02080821804702282, -0.32788369059562683, 0.18608127534389496, 0.1342967003583908, 1.8224468231201172, 0.028002522885799408, -1.553360104560852, -0.7706471085548401, 0.20660938322544098, 0.17195142805576324]}}, "reward": 2.0, "success": true, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [0.2458915114402771, -0.15052920579910278, -1.0, 0.013746529817581177, -1.0, -0.21971148252487183, 2.714754104614258, -0.2429777830839157, 0.07430313527584076, -1.0, -0.16456840932369232, -1.0, -0.9286653995513916, 0.2025165557861328, 0.025035768747329712, 0.28885743021965027, -0.2897254228591919, -0.6306957602500916, -0.5820335745811462, -0.9942173361778259, -2.5460946559906006, 0.028109241276979446, -0.05866494029760361, 0.4426884055137634, -0.4524412751197815, -0.21268394589424133, -0.7261583805084229, -3.2692134380340576], "action_chunk": {"action_id": "near_miss-01", "horizon": 4, "metadata": {"candidate_type": "near_miss", "parent_action_id": "expert", "perturbation": {"kind": "gaussian_action_offset", "scale": 0.05}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[0.2458915114402771, -0.15052920579910278, -1.0, 0.013746529817581177, -1.0, -0.21971148252487183, 2.714754104614258], [-0.2429777830839157, 0.07430313527584076, -1.0, -0.16456840932369232, -1.0, -0.9286653995513916, 0.2025165557861328], [0.025035768747329712, 0.28885743021965027, -0.2897254228591919, -0.6306957602500916, -0.5820335745811462, -0.9942173361778259, -2.5460946559906006], [0.028109241276979446, -0.05866494029760361, 0.4426884055137634, -0.4524412751197815, -0.21268394589424133, -0.7261583805084229, -3.2692134380340576]]}, "cil": {"candidate_type": "near_miss", "failure_type": "partial_success", "group_id": "ms-10a67a57026bb09b5d93", "rank_within_group": 0, "record_id": "rec-2418612df847b9735e84e14e", "regret": 0.0, "scene_id": "episode-252", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "ff0570f2474cb1f6f4e75c4921a0bd5cb308918756ed6b4ff6cb2df798f1d884", "task_id": "PickCube-v1"}, "episode_index": 10, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/10", "image": null, "state": {"features": [0.021350860595703125, 0.0422821044921875, 0.019999928772449493, 0.9657067060470581, 5.530319413082907e-06, -3.955532520194538e-06, 0.259635329246521, 8.874305785866454e-05, 0.00019909856200683862, -0.00010172931069973856, -0.004346663132309914, 0.00043586635729297996, 0.0009438440320082009, 0.02393341064453125, 0.043304443359375, 0.0746343731880188, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.11999893188476562, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6150016784667969, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.02423270232975483, 0.5300643444061279, 0.006811963859945536, -1.8714677095413208, 0.04015753045678139, 2.419536828994751, 0.8830363154411316, 0.03999907523393631, 0.03999999910593033, 0.6293052434921265, 1.8052808046340942, 0.44981449842453003, 0.9143156409263611, 0.2605507969856262, 0.20237629115581512, 0.6918536424636841, 0.0007163410191424191, -0.0007135514752008021]}}, "reward": 0.5448155999183655, "success": false, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} +{"action": [0.7216213941574097, 0.23776167631149292, 0.9413163661956787, -0.9285341501235962, -0.8461328744888306, -1.0, -3.2082104682922363, 0.717744767665863, 0.842609167098999, 1.0, -1.0, -0.22849977016448975, -1.0, -2.786240577697754, 0.393378883600235, 0.780874490737915, 1.0, -1.0, -0.5397374629974365, -1.0, -2.491328239440918, 0.19940781593322754, 0.033484384417533875, 0.9354397058486938, -0.973723292350769, -0.47505873441696167, -0.10541725158691406, -2.2636466026306152], "action_chunk": {"action_id": "wrong_direction-03", "horizon": 4, "metadata": {"candidate_type": "wrong_direction", "parent_action_id": "expert", "perturbation": {"kind": "gaussian_action_offset", "scale": 0.35}, "source": "measured_physics_branch"}, "representation": "maniskill_pd_ee_delta_pose", "skill_type": "pick_cube", "values": [[0.7216213941574097, 0.23776167631149292, 0.9413163661956787, -0.9285341501235962, -0.8461328744888306, -1.0, -3.2082104682922363], [0.717744767665863, 0.842609167098999, 1.0, -1.0, -0.22849977016448975, -1.0, -2.786240577697754], [0.393378883600235, 0.780874490737915, 1.0, -1.0, -0.5397374629974365, -1.0, -2.491328239440918], [0.19940781593322754, 0.033484384417533875, 0.9354397058486938, -0.973723292350769, -0.47505873441696167, -0.10541725158691406, -2.2636466026306152]]}, "cil": {"candidate_type": "wrong_direction", "failure_type": "partial_success", "group_id": "ms-25bb20f3481a2a40b97f", "rank_within_group": 0, "record_id": "rec-ae900c27d59bf14be97166ff", "regret": 0.0, "scene_id": "episode-872", "source_dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", "state_hash": "3b7f3ed5e759bccc9824e605b3d8cefb7b12e585b4aff41905b7e27e0d996a23", "task_id": "PickCube-v1"}, "episode_index": 11, "frame_index": 0, "observation": {"cil_observation_ref": "observations.h5#initial_rgb_jpeg/11", "image": null, "state": {"features": [0.03521728515625, -0.057384490966796875, 0.03253794461488724, 0.9129041433334351, -0.0028452409896999598, 0.010323070921003819, -0.4080334007740021, 0.2689715623855591, 0.46322503685951233, 0.5909228920936584, 2.109616756439209, 0.20898234844207764, 5.292329788208008, 0.0795440673828125, 0.03389739990234375, 0.2208138108253479, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.11999893188476562, 0.0, -0.9196429252624512, 0.7071067690849304, 0.0, 0.0, 0.7071067690849304, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.6150016784667969, 0.0, -1.862645149230957e-08, 1.0000001192092896, 2.720872971373467e-10, -3.2895179913339234e-09, -3.588018770983581e-12, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, -0.08180949836969376, 0.795208752155304, -0.01031059306114912, -1.6480636596679688, 0.03492514789104462, 2.3156821727752686, 0.5561749339103699, 0.02580047957599163, 0.025842761620879173, 0.383100301027298, -0.5445420145988464, 0.3218225836753845, 0.8863900899887085, 0.02420014515519142, -1.6891124248504639, 0.13746380805969238, 0.028009921312332153, -0.016399741172790527]}}, "reward": 0.5596571564674377, "success": false, "task": "Pick up the cube and move it to the goal position.", "task_index": 0, "timestamp": 0.0} diff --git a/workspace/outputs/external_vla_plan_probe/external_vla_baseline_plan.json b/workspace/outputs/external_vla_plan_probe/external_vla_baseline_plan.json new file mode 100644 index 0000000000000000000000000000000000000000..cd86ea99ac507b661e7476b4e471252c361dfd99 --- /dev/null +++ b/workspace/outputs/external_vla_plan_probe/external_vla_baseline_plan.json @@ -0,0 +1,45 @@ +{ + "commands": { + "create_env": "python -m venv outputs/external_vla_plan_probe/external_vla_env", + "download": "hf download lerobot/smolvla_base --revision c83c3163b8ca9b7e67c509fffd9121e66cb96205 --local-dir /scratch/knguy52/dovla/models/smolvla_base-c83c316", + "install": "outputs/external_vla_plan_probe/external_vla_env/bin/python -m pip install --upgrade pip lerobot[smolvla]", + "run": "scripts/run_external_vla_baseline.py --model-family smolvla --checkpoint /scratch/knguy52/dovla/models/smolvla_base-c83c316 --out outputs/external_vla_plan_probe --dataset outputs/external_vla_export_smoke --require-ready" + }, + "expected_adapter_contract": { + "call_signature": "function(spec_dict: dict, plan: dict) -> dict", + "entrypoint": "module:function", + "output": "A JSON-serializable metrics dictionary with measured rollout/eval metrics." + }, + "schema_version": "external-vla-baseline-plan/v0", + "spec": { + "adapter_entrypoint": null, + "checkpoint_path": "/scratch/knguy52/dovla/models/smolvla_base-c83c316", + "dataset_dir": "outputs/external_vla_export_smoke", + "metadata": {}, + "model_family": "smolvla", + "out_dir": "outputs/external_vla_plan_probe", + "package_name": "lerobot", + "python": "python", + "repo_id": "lerobot/smolvla_base", + "revision": "c83c3163b8ca9b7e67c509fffd9121e66cb96205" + }, + "status": { + "adapter_entrypoint": null, + "adapter_importable": false, + "checkpoint_exists": true, + "checkpoint_path": "/scratch/knguy52/dovla/models/smolvla_base-c83c316", + "dataset_dir": "outputs/external_vla_export_smoke", + "dataset_exists": true, + "missing": [ + "package:lerobot", + "adapter_entrypoint" + ], + "model_family": "smolvla", + "package_available": false, + "package_name": "lerobot", + "ready": false, + "warnings": [ + "Run in an isolated environment; LeRobot pins a different Transformers stack." + ] + } +} \ No newline at end of file diff --git a/workspace/outputs/external_vla_smolvla_checkpoint_manifest.json b/workspace/outputs/external_vla_smolvla_checkpoint_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..f292a51c460bb10c43c01910dd8e38c7fbdc0ca9 --- /dev/null +++ b/workspace/outputs/external_vla_smolvla_checkpoint_manifest.json @@ -0,0 +1,179 @@ +{ + "file_count": 32, + "files": [ + { + "path": ".cache/huggingface/.gitignore", + "sha256": "684888c0ebb17f374298b65ee2807526c066094c701bcc7ebbe1c1095f494fc1", + "size_bytes": 1 + }, + { + "path": ".cache/huggingface/CACHEDIR.TAG", + "sha256": "f6572428f6d5e1575e73a1502895a8731f10757dfbb634909c6e154b849af91d", + "size_bytes": 191 + }, + { + "path": ".cache/huggingface/download/.gitattributes.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/.gitattributes.metadata", + "sha256": "7a3b21bc4e389b1212ebbb4a06aba23837e16415c213c473a6179edcb190d519", + "size_bytes": 101 + }, + { + "path": ".cache/huggingface/download/Finetune_SmolVLA_notebook.ipynb.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/Finetune_SmolVLA_notebook.ipynb.metadata", + "sha256": "cb6b20b53d29f9e0edc8dcd35cb448d5e7ce9a906b08e03993a08cca9e0803ab", + "size_bytes": 101 + }, + { + "path": ".cache/huggingface/download/README.md.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/README.md.metadata", + "sha256": "0e6891a1bfed24b5044d0ecfb0296b97a4533bc3b886554bf89129bb732c27d5", + "size_bytes": 101 + }, + { + "path": ".cache/huggingface/download/collage_small.gif.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/collage_small.gif.metadata", + "sha256": "5e7c7be68c0e602c50adebf90ae66b36167d2e968bb7d2830643a84e8d48b879", + "size_bytes": 125 + }, + { + "path": ".cache/huggingface/download/config.json.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/config.json.metadata", + "sha256": "a26ca767fe9cbde682c09dc0b7433a843695813e8ded99e8e1102744fc03c6f9", + "size_bytes": 100 + }, + { + "path": ".cache/huggingface/download/model.safetensors.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/model.safetensors.metadata", + "sha256": "2e9ae096f1a034577c5f18c168099b5211723118d29a5d16e947408d7216293b", + "size_bytes": 125 + }, + { + "path": ".cache/huggingface/download/policy_postprocessor.json.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/policy_postprocessor.json.metadata", + "sha256": "ee2fab7f02469ca95442fea425b1b98288523d5dcad0d771b7744ad8cf8f620b", + "size_bytes": 100 + }, + { + "path": ".cache/huggingface/download/policy_postprocessor_step_0_unnormalizer_processor.safetensors.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/policy_postprocessor_step_0_unnormalizer_processor.safetensors.metadata", + "sha256": "4ca95ad746b5a195476c7bb2a6a3572106e26c342101ecc1b76371f56d744af4", + "size_bytes": 125 + }, + { + "path": ".cache/huggingface/download/policy_preprocessor.json.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/policy_preprocessor.json.metadata", + "sha256": "115e4538ec5d1c9c76e3f3ab1ab34c34b0253e10252c1989c913faef848740b1", + "size_bytes": 101 + }, + { + "path": ".cache/huggingface/download/policy_preprocessor_step_5_normalizer_processor.safetensors.lock", + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size_bytes": 0 + }, + { + "path": ".cache/huggingface/download/policy_preprocessor_step_5_normalizer_processor.safetensors.metadata", + "sha256": "888bf914735bbff54bc65c4c05d7da3ab71b2c1dd0de3d14e8e1419469383311", + "size_bytes": 125 + }, + { + "path": ".gitattributes", + "sha256": "32949ddb126c25a78e757b890078c46d0c8f8e899e56510a1efef817233149f9", + "size_bytes": 1573 + }, + { + "path": "Finetune_SmolVLA_notebook.ipynb", + "sha256": "7f445c913c508142b632245ae52494c23ac8b96109fe1f9144e34b8b9c60f92d", + "size_bytes": 6757 + }, + { + "path": "README.md", + "sha256": "adb98360b162fe40b6f7995d21c1a4250bd51ddd094973234e290482a35ea4ce", + "size_bytes": 4605 + }, + { + "path": "collage_small.gif", + "sha256": "c43f022bf1fdfbac82841ef95c7a99a7ead13ea67f2a7202f2fe57148f352c83", + "size_bytes": 8009386 + }, + { + "path": "config.json", + "sha256": "650584b56c104720f7a3c91d1ec6bec9e8de8ac11e60c92ba2fa82d93eda147d", + "size_bytes": 2299 + }, + { + "path": "model.safetensors", + "sha256": "7cd549ac2351fb069c0ddb3c34ad2d09cfc92b56a15dccdfc2e41467aaca01eb", + "size_bytes": 906712520 + }, + { + "path": "policy_postprocessor.json", + "sha256": "2b78bb742065288df2ec63b0ee35f97a1f6950171cf2272f7e34b1fe3873b17b", + "size_bytes": 660 + }, + { + "path": "policy_postprocessor_step_0_unnormalizer_processor.safetensors", + "sha256": "490ab239d96e263687c0b2e386a0afbc235a2eceb9857c36ed32f2f162a3e7c8", + "size_bytes": 640 + }, + { + "path": "policy_preprocessor.json", + "sha256": "7683d648280a72f0d51e869fb8dd1596beed90b7f84849330cf1bc26634a4bd6", + "size_bytes": 1872 + }, + { + "path": "policy_preprocessor_step_5_normalizer_processor.safetensors", + "sha256": "490ab239d96e263687c0b2e386a0afbc235a2eceb9857c36ed32f2f162a3e7c8", + "size_bytes": 640 + } + ], + "local_dir": "/scratch/knguy52/dovla/models/smolvla_base-c83c316", + "missing_required_files": [], + "model_family": "smolvla", + "ready": true, + "repo_id": "lerobot/smolvla_base", + "required_files": [ + "config.json", + "model.safetensors", + "policy_preprocessor.json", + "policy_postprocessor.json" + ], + "revision": "c83c3163b8ca9b7e67c509fffd9121e66cb96205", + "schema_version": "dovla-external-checkpoint-manifest/v0", + "total_bytes": 914742248 +} diff --git a/workspace/outputs/external_vla_smolvla_checkpoint_metadata.json b/workspace/outputs/external_vla_smolvla_checkpoint_metadata.json new file mode 100644 index 0000000000000000000000000000000000000000..6df25aa62d12591120558f8deb319712c18c0493 --- /dev/null +++ b/workspace/outputs/external_vla_smolvla_checkpoint_metadata.json @@ -0,0 +1,16 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/models/smolvla_base-c83c316", + "metadata_only": true, + "package_versions": { + "huggingface-hub": null, + "lerobot": null, + "torch": "2.12.1+computecanada", + "transformers": null + }, + "ready": false, + "required_files": [ + "config.json", + "model.safetensors" + ], + "schema_version": "smolvla-checkpoint-smoke/v0" +} \ No newline at end of file diff --git a/workspace/outputs/external_vla_smolvla_checkpoint_smoke.json b/workspace/outputs/external_vla_smolvla_checkpoint_smoke.json new file mode 100644 index 0000000000000000000000000000000000000000..6a36fd237ab4cba7c8eabdb83e640508e1804f17 --- /dev/null +++ b/workspace/outputs/external_vla_smolvla_checkpoint_smoke.json @@ -0,0 +1,24 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/models/smolvla_base-c83c316", + "cuda_device": "NVIDIA H100 80GB HBM3 MIG 1g.10gb", + "device": "cuda", + "load_seconds": 21.856225808965974, + "metadata_only": false, + "package_versions": { + "huggingface-hub": "0.35.3+computecanada", + "lerobot": "0.4.3", + "torch": "2.7.1+cu128", + "transformers": "4.57.6+computecanada" + }, + "parameter_count": 450046176, + "policy_class": "lerobot.policies.smolvla.modeling_smolvla.SmolVLAPolicy", + "ready": true, + "required_files": [ + "config.json", + "model.safetensors" + ], + "schema_version": "smolvla-checkpoint-smoke/v0", + "trainable_parameter_count": 99880992, + "vlm_load_mode": "local_config_then_policy_safetensors", + "vlm_metadata": "/scratch/knguy52/dovla/models/SmolVLM2-500M-Video-Instruct-metadata-7b375e1b73b11138ff12fe22c8f2822d8fe03467" +} \ No newline at end of file diff --git a/workspace/outputs/external_vla_smolvla_full_plan/external_vla_baseline_plan.json b/workspace/outputs/external_vla_smolvla_full_plan/external_vla_baseline_plan.json new file mode 100644 index 0000000000000000000000000000000000000000..233c37a2530d7d6ab86756ccd7fb9277125d8242 --- /dev/null +++ b/workspace/outputs/external_vla_smolvla_full_plan/external_vla_baseline_plan.json @@ -0,0 +1,45 @@ +{ + "commands": { + "create_env": "python -m venv outputs/external_vla_smolvla_full_plan/external_vla_env", + "download": "hf download lerobot/smolvla_base --revision c83c3163b8ca9b7e67c509fffd9121e66cb96205 --local-dir /scratch/knguy52/dovla/models/smolvla_base-c83c316", + "install": "outputs/external_vla_smolvla_full_plan/external_vla_env/bin/python -m pip install --upgrade pip lerobot[smolvla]", + "run": "scripts/run_external_vla_baseline.py --model-family smolvla --checkpoint /scratch/knguy52/dovla/models/smolvla_base-c83c316 --out outputs/external_vla_smolvla_full_plan --dataset outputs/external_vla_export_maniskill_full_no_images --require-ready" + }, + "expected_adapter_contract": { + "call_signature": "function(spec_dict: dict, plan: dict) -> dict", + "entrypoint": "module:function", + "output": "A JSON-serializable metrics dictionary with measured rollout/eval metrics." + }, + "schema_version": "external-vla-baseline-plan/v0", + "spec": { + "adapter_entrypoint": null, + "checkpoint_path": "/scratch/knguy52/dovla/models/smolvla_base-c83c316", + "dataset_dir": "outputs/external_vla_export_maniskill_full_no_images", + "metadata": {}, + "model_family": "smolvla", + "out_dir": "outputs/external_vla_smolvla_full_plan", + "package_name": "lerobot", + "python": "python", + "repo_id": "lerobot/smolvla_base", + "revision": "c83c3163b8ca9b7e67c509fffd9121e66cb96205" + }, + "status": { + "adapter_entrypoint": null, + "adapter_importable": false, + "checkpoint_exists": true, + "checkpoint_path": "/scratch/knguy52/dovla/models/smolvla_base-c83c316", + "dataset_dir": "outputs/external_vla_export_maniskill_full_no_images", + "dataset_exists": true, + "missing": [ + "package:lerobot", + "adapter_entrypoint" + ], + "model_family": "smolvla", + "package_available": false, + "package_name": "lerobot", + "ready": false, + "warnings": [ + "Run in an isolated environment; LeRobot pins a different Transformers stack." + ] + } +} \ No newline at end of file diff --git a/workspace/outputs/hf_sync/hf_push_daemon_login.out b/workspace/outputs/hf_sync/hf_push_daemon_login.out new file mode 100644 index 0000000000000000000000000000000000000000..928ea6cd14591454cbde0684ded0dac6b681ca82 --- /dev/null +++ b/workspace/outputs/hf_sync/hf_push_daemon_login.out @@ -0,0 +1,17 @@ +[2026-07-02T13:36:07Z] starting Hugging Face sync daemon; interval=900s repo=anhtld/vla +[2026-07-02T13:36:07Z] checking Hugging Face auth +[2026-07-02T13:36:59Z] starting Hugging Face sync daemon; interval=900s repo=anhtld/vla +[2026-07-02T13:36:59Z] checking Hugging Face auth +[2026-07-02T13:37:00Z] sync start repo=anhtld/vla type=model dry_run=0 +[2026-07-02T13:37:00Z] upload /lustre09/project/6037638/knguy52/vla -> anhtld/vla/workspace +Found 26,894 files to upload +Uploading... 256/26,894 files checked, 0/0 uploaded (0.00B transferred), 0 committed in 0 commit(s) +Uploading... 256/26,894 files checked, 0/0 uploaded (0.00B transferred), 0 committed in 0 commit(s) +Uploading... 256/26,894 files checked, 0/0 uploaded (0.00B transferred), 0 committed in 0 commit(s) +Uploading... 256/26,894 files checked, 0/0 uploaded (0.00B transferred), 0 committed in 0 commit(s) +Uploading... 1,792/26,894 files checked, 21/21 uploaded (166kB transferred), 500 committed in 2 commit(s) +Uploading... 5,888/26,894 files checked, 112/112 uploaded (166kB transferred), 3,150 committed in 8 commit(s) +Uploading... 9,984/26,894 files checked, 117/117 uploaded (166kB transferred), 6,750 committed in 12 commit(s) +Uploading... 13,824/26,894 files checked, 117/117 uploaded (166kB transferred), 10,750 committed in 16 commit(s) +Uploading... 17,920/26,894 files checked, 117/117 uploaded (166kB transferred), 14,750 committed in 20 commit(s) +Uploading... 21,760/26,894 files checked, 125/125 uploaded (166kB transferred), 18,750 committed in 24 commit(s) diff --git a/workspace/outputs/hf_sync/hf_push_daemon_login.pid b/workspace/outputs/hf_sync/hf_push_daemon_login.pid new file mode 100644 index 0000000000000000000000000000000000000000..b8735f33a3122598d9bb55a143f718b38a2667f4 --- /dev/null +++ b/workspace/outputs/hf_sync/hf_push_daemon_login.pid @@ -0,0 +1 @@ +3271759 diff --git a/workspace/outputs/hf_sync/hf_sync.log b/workspace/outputs/hf_sync/hf_sync.log new file mode 100644 index 0000000000000000000000000000000000000000..68c7d02824b031416927e663d11d8e89d8a25fd8 --- /dev/null +++ b/workspace/outputs/hf_sync/hf_sync.log @@ -0,0 +1,19 @@ +[2026-07-02T13:29:49Z] sync start repo=anhtld/vla type=model dry_run=1 +[2026-07-02T13:29:49Z] dry-run would upload /lustre09/project/6037638/knguy52/vla -> anhtld/vla/workspace +[2026-07-02T13:29:49Z] dry-run would upload /scratch/knguy52/dovla/experiments/dovla_h16_policy_ckpt_runs -> anhtld/vla/scratch/experiments/dovla_h16_policy_ckpt_runs +[2026-07-02T13:29:49Z] dry-run would upload /scratch/knguy52/dovla/experiments/dovla_h16_rollout_runs -> anhtld/vla/scratch/experiments/dovla_h16_rollout_runs +[2026-07-02T13:29:49Z] dry-run would upload /scratch/knguy52/dovla/experiments/h16_merged_dataset -> anhtld/vla/scratch/experiments/h16_merged_dataset +[2026-07-02T13:29:49Z] dry-run would upload /scratch/knguy52/dovla/experiments/six_task_h16_collection -> anhtld/vla/scratch/experiments/six_task_h16_collection +[2026-07-02T13:29:49Z] dry-run would upload /scratch/knguy52/dovla/experiments/embedding_cache -> anhtld/vla/scratch/experiments/embedding_cache +[2026-07-02T13:29:49Z] dry-run would upload /scratch/knguy52/dovla/cache -> anhtld/vla/scratch/cache +[2026-07-02T13:29:49Z] dry-run would upload /scratch/knguy52/dovla/caches -> anhtld/vla/scratch/caches +[2026-07-02T13:29:49Z] dry-run would upload /lustre09/project/6037638/knguy52/vla/outputs/hf_sync/last_manifest.json -> anhtld/vla/manifests/hf_sync_last_manifest.json +[2026-07-02T13:29:49Z] sync complete +[2026-07-02T13:31:34Z] checking Hugging Face auth +[2026-07-02T13:32:34Z] ERROR: Hugging Face auth is not active. Run 'hf auth login' or export HF_TOKEN before syncing. +[2026-07-02T13:33:55Z] checking Hugging Face auth +[2026-07-02T13:34:55Z] ERROR: Hugging Face auth is not active. Run 'hf auth login' or export HF_TOKEN before syncing. +[2026-07-02T13:36:07Z] checking Hugging Face auth +[2026-07-02T13:36:59Z] checking Hugging Face auth +[2026-07-02T13:37:00Z] sync start repo=anhtld/vla type=model dry_run=0 +[2026-07-02T13:37:00Z] upload /lustre09/project/6037638/knguy52/vla -> anhtld/vla/workspace diff --git a/workspace/outputs/hf_sync/hf_sync_daemon.log b/workspace/outputs/hf_sync/hf_sync_daemon.log new file mode 100644 index 0000000000000000000000000000000000000000..ca882c8f1f71aa8ac343279a4402d7a2c180d0b9 --- /dev/null +++ b/workspace/outputs/hf_sync/hf_sync_daemon.log @@ -0,0 +1,7 @@ +[2026-07-02T13:30:01Z] starting Hugging Face sync daemon; interval=900s repo=anhtld/vla +[2026-07-02T13:31:34Z] starting Hugging Face sync daemon; interval=900s repo=anhtld/vla +[2026-07-02T13:32:34Z] sync cycle failed with status=2 +[2026-07-02T13:33:55Z] starting Hugging Face sync daemon; interval=900s repo=anhtld/vla +[2026-07-02T13:34:55Z] sync cycle failed with status=2 +[2026-07-02T13:36:07Z] starting Hugging Face sync daemon; interval=900s repo=anhtld/vla +[2026-07-02T13:36:59Z] starting Hugging Face sync daemon; interval=900s repo=anhtld/vla diff --git a/workspace/outputs/hf_sync/last_manifest.json b/workspace/outputs/hf_sync/last_manifest.json new file mode 100644 index 0000000000000000000000000000000000000000..373be4701101bd34051765c5e13e8bad89e8624f --- /dev/null +++ b/workspace/outputs/hf_sync/last_manifest.json @@ -0,0 +1,48 @@ +{ + "generated_utc": "2026-07-02T13:37:00.984705+00:00", + "repo_id": "anhtld/vla", + "project_dir": "/lustre09/project/6037638/knguy52/vla", + "scratch_root": "/scratch/knguy52/dovla", + "roots": [ + { + "remote": "workspace", + "local": "/lustre09/project/6037638/knguy52/vla", + "exists": true + }, + { + "remote": "scratch/experiments/dovla_h16_policy_ckpt_runs", + "local": "/scratch/knguy52/dovla/experiments/dovla_h16_policy_ckpt_runs", + "exists": true + }, + { + "remote": "scratch/experiments/dovla_h16_rollout_runs", + "local": "/scratch/knguy52/dovla/experiments/dovla_h16_rollout_runs", + "exists": true + }, + { + "remote": "scratch/experiments/h16_merged_dataset", + "local": "/scratch/knguy52/dovla/experiments/h16_merged_dataset", + "exists": true + }, + { + "remote": "scratch/experiments/six_task_h16_collection", + "local": "/scratch/knguy52/dovla/experiments/six_task_h16_collection", + "exists": true + }, + { + "remote": "scratch/experiments/embedding_cache", + "local": "/scratch/knguy52/dovla/experiments/embedding_cache", + "exists": true + }, + { + "remote": "scratch/cache", + "local": "/scratch/knguy52/dovla/cache", + "exists": true + }, + { + "remote": "scratch/caches", + "local": "/scratch/knguy52/dovla/caches", + "exists": true + } + ] +} diff --git a/workspace/outputs/hpc/logs/build_m010_14936133.err b/workspace/outputs/hpc/logs/build_m010_14936133.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_m010_14936133.out b/workspace/outputs/hpc/logs/build_m010_14936133.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_m010_14936133.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_oraclek8_14935530.err b/workspace/outputs/hpc/logs/build_oraclek8_14935530.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_oraclek8_14935530.out b/workspace/outputs/hpc/logs/build_oraclek8_14935530.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_oraclek8_14935530.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_oraclek8u_14953524.err b/workspace/outputs/hpc/logs/build_oraclek8u_14953524.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_oraclek8u_14953524.out b/workspace/outputs/hpc/logs/build_oraclek8u_14953524.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_oraclek8u_14953524.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14857117.err b/workspace/outputs/hpc/logs/build_paper_table_14857117.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14857117.out b/workspace/outputs/hpc/logs/build_paper_table_14857117.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14857117.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14858333.err b/workspace/outputs/hpc/logs/build_paper_table_14858333.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14858333.out b/workspace/outputs/hpc/logs/build_paper_table_14858333.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14858333.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14858455.err b/workspace/outputs/hpc/logs/build_paper_table_14858455.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14858455.out b/workspace/outputs/hpc/logs/build_paper_table_14858455.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14858455.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14858883.err b/workspace/outputs/hpc/logs/build_paper_table_14858883.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14858883.out b/workspace/outputs/hpc/logs/build_paper_table_14858883.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14858883.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859046.err b/workspace/outputs/hpc/logs/build_paper_table_14859046.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859046.out b/workspace/outputs/hpc/logs/build_paper_table_14859046.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14859046.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859203.err b/workspace/outputs/hpc/logs/build_paper_table_14859203.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859203.out b/workspace/outputs/hpc/logs/build_paper_table_14859203.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14859203.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859299.err b/workspace/outputs/hpc/logs/build_paper_table_14859299.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859299.out b/workspace/outputs/hpc/logs/build_paper_table_14859299.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14859299.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859402.err b/workspace/outputs/hpc/logs/build_paper_table_14859402.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859402.out b/workspace/outputs/hpc/logs/build_paper_table_14859402.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14859402.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859511.err b/workspace/outputs/hpc/logs/build_paper_table_14859511.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859511.out b/workspace/outputs/hpc/logs/build_paper_table_14859511.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14859511.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859597.err b/workspace/outputs/hpc/logs/build_paper_table_14859597.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14859597.out b/workspace/outputs/hpc/logs/build_paper_table_14859597.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14859597.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14869605.err b/workspace/outputs/hpc/logs/build_paper_table_14869605.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14869605.out b/workspace/outputs/hpc/logs/build_paper_table_14869605.out new file mode 100644 index 0000000000000000000000000000000000000000..6fd2c45c3bb9eb15910694d9eaef54ffa11ee89e --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14869605.out @@ -0,0 +1,2 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14869860.err b/workspace/outputs/hpc/logs/build_paper_table_14869860.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14869860.out b/workspace/outputs/hpc/logs/build_paper_table_14869860.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14869860.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14883440.err b/workspace/outputs/hpc/logs/build_paper_table_14883440.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14883440.out b/workspace/outputs/hpc/logs/build_paper_table_14883440.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14883440.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14883727.err b/workspace/outputs/hpc/logs/build_paper_table_14883727.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14883727.out b/workspace/outputs/hpc/logs/build_paper_table_14883727.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14883727.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14883926.err b/workspace/outputs/hpc/logs/build_paper_table_14883926.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14883926.out b/workspace/outputs/hpc/logs/build_paper_table_14883926.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14883926.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14884383.err b/workspace/outputs/hpc/logs/build_paper_table_14884383.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14884383.out b/workspace/outputs/hpc/logs/build_paper_table_14884383.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14884383.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14890079.err b/workspace/outputs/hpc/logs/build_paper_table_14890079.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14890079.out b/workspace/outputs/hpc/logs/build_paper_table_14890079.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14890079.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14891089.err b/workspace/outputs/hpc/logs/build_paper_table_14891089.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14891089.out b/workspace/outputs/hpc/logs/build_paper_table_14891089.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14891089.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14891970.err b/workspace/outputs/hpc/logs/build_paper_table_14891970.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14891970.out b/workspace/outputs/hpc/logs/build_paper_table_14891970.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14891970.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14893069.err b/workspace/outputs/hpc/logs/build_paper_table_14893069.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14893069.out b/workspace/outputs/hpc/logs/build_paper_table_14893069.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14893069.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14893791.err b/workspace/outputs/hpc/logs/build_paper_table_14893791.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14893791.out b/workspace/outputs/hpc/logs/build_paper_table_14893791.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14893791.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14894302.err b/workspace/outputs/hpc/logs/build_paper_table_14894302.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14894302.out b/workspace/outputs/hpc/logs/build_paper_table_14894302.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14894302.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14894440.err b/workspace/outputs/hpc/logs/build_paper_table_14894440.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14894440.out b/workspace/outputs/hpc/logs/build_paper_table_14894440.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14894440.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14894678.err b/workspace/outputs/hpc/logs/build_paper_table_14894678.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14894678.out b/workspace/outputs/hpc/logs/build_paper_table_14894678.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14894678.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14897129.err b/workspace/outputs/hpc/logs/build_paper_table_14897129.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14897129.out b/workspace/outputs/hpc/logs/build_paper_table_14897129.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14897129.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14915014.err b/workspace/outputs/hpc/logs/build_paper_table_14915014.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14915014.out b/workspace/outputs/hpc/logs/build_paper_table_14915014.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14915014.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14915251.err b/workspace/outputs/hpc/logs/build_paper_table_14915251.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14915251.out b/workspace/outputs/hpc/logs/build_paper_table_14915251.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14915251.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14926728.err b/workspace/outputs/hpc/logs/build_paper_table_14926728.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14926728.out b/workspace/outputs/hpc/logs/build_paper_table_14926728.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14926728.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14927922.err b/workspace/outputs/hpc/logs/build_paper_table_14927922.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14927922.out b/workspace/outputs/hpc/logs/build_paper_table_14927922.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14927922.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/build_paper_table_14929219.err b/workspace/outputs/hpc/logs/build_paper_table_14929219.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/build_paper_table_14929219.out b/workspace/outputs/hpc/logs/build_paper_table_14929219.out new file mode 100644 index 0000000000000000000000000000000000000000..af5f760ae1b1f3686a27f0f633c225042d6d2fd4 --- /dev/null +++ b/workspace/outputs/hpc/logs/build_paper_table_14929219.out @@ -0,0 +1,4 @@ +Wrote results/paper_table_status.json +Wrote results/paper_table_status.md +Wrote results/paper_analysis.json +Wrote results/paper_analysis.md diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_0.err b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_0.out b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_0.out new file mode 100644 index 0000000000000000000000000000000000000000..1644f9d796584dc67e896643489dd728c2392f01 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_0.out @@ -0,0 +1,1203 @@ +================================================== +Task: PickCube-v1 +Groups: 1000 +Horizon: 16 (vs baseline 4) +Demo: /scratch/knguy52/dovla/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5 +Output: /scratch/knguy52/dovla/experiments/six_task_h16_collection/PickCube-v1 +================================================== +ManiSkill CIL progress: batch=1/1000 groups=1/1000 records=16 restore_max=1.192e-07 +ManiSkill CIL progress: batch=50/1000 groups=50/1000 records=800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=100/1000 groups=100/1000 records=1600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=150/1000 groups=150/1000 records=2400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=200/1000 groups=200/1000 records=3200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=250/1000 groups=250/1000 records=4000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=300/1000 groups=300/1000 records=4800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=350/1000 groups=350/1000 records=5600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=400/1000 groups=400/1000 records=6400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=450/1000 groups=450/1000 records=7200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=500/1000 groups=500/1000 records=8000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=550/1000 groups=550/1000 records=8800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=600/1000 groups=600/1000 records=9600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=650/1000 groups=650/1000 records=10400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=700/1000 groups=700/1000 records=11200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=750/1000 groups=750/1000 records=12000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=800/1000 groups=800/1000 records=12800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=850/1000 groups=850/1000 records=13600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=900/1000 groups=900/1000 records=14400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=950/1000 groups=950/1000 records=15200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=1000/1000 groups=1000/1000 records=16000 restore_max=4.768e-07 +{ + "backend": "maniskill", + "env_id": "PickCube-v1", + "control_mode": "pd_ee_delta_pose", + "sim_backend": "physx_cuda:0", + "source_obs_mode": "state", + "instruction": "Pick up the cube and move it to the goal position.", + "target_actors": [ + "cube" + ], + "reference_actors": [], + "demo_path": "/scratch/knguy52/dovla/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5", + "num_groups": 1000, + "group_offset": 0, + "available_unique_states": 14701, + "excluded_post_success_states": 20466, + "branch_selection": "pre_success_only", + "num_records": 16000, + "k": 16, + "horizon": 16, + "seed": 0, + "measurement": "exact_state_physics_branch", + "candidate_mode": "structured", + "branch_execution": "parallel", + "state_batch_size": 1, + "state_restore_max_error": 4.76837158203125e-07, + "state_storage": "archive", + "state_archive": "/scratch/knguy52/dovla/experiments/six_task_h16_collection/PickCube-v1/state_archive.pkl", + "observation_archive": null, + "image_quality": null, + "manifest": "/scratch/knguy52/dovla/experiments/six_task_h16_collection/PickCube-v1/manifest.json", + "quality": { + "expert_success_rate": 0.947, + "no_op_success_rate": 0.283, + "mean_reward_spread": 0.9618538443985621, + "nondegenerate_group_fraction": 1.0 + }, + "sharding": { + "dataset_name": "cil_pickcube-v1", + "version": "0.1", + "created_at": "2026-06-25T23:31:23.448339+00:00", + "backend": "maniskill", + "num_groups": 1000, + "num_records": 16000, + "k": 16, + "task_count": 1, + "seed": 0, + "schema_version": "0.1", + "format": "dovla_cil", + "shard_format": "jsonl", + "index_format": "jsonl", + "shard_size": 1024, + "shard_count": 16, + "group_index_path": "group_index.jsonl", + "record_index_path": "record_index.jsonl", + "shards": [ + { + "path": "shards/shard_000000.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-d88107eff1e361cdcc4c", + "ms-4cd9eae753bf9e48f520", + "ms-6edb3f8fe977a81f8979", + "ms-23a355317920e192147f", + "ms-faeacb4289bbfe504de1", + "ms-f88a44616a0985aac518", + "ms-c8666cbd7e5a9d0fd036", + "ms-f5483f0819e0e8bc352e", + "ms-22675530f3a401b3fd6d", + "ms-751daa0137578c486f3d", + "ms-193063285cdb9faff39b", + "ms-5d430ebe6f7933dc6c8c", + "ms-67f602a0e4efc9984d72", + "ms-564c02492863ae749ea9", + "ms-86b0447349029dfaa6ab", + "ms-b481531ab787e3d6fb1e", + "ms-e441e4042cb7d77a4986", + "ms-5dfad80e83d210946447", + "ms-be8eff6e6d520ae11a0d", + "ms-93e3d9968204474eff49", + "ms-708baf01c22c75cafd22", + "ms-0e55645c0c701caac48a", + "ms-29dd2cdbb66fe13ab877", + "ms-4fe47d946641f0096a85", + "ms-fb98254adaa23dc128df", + "ms-10621e1108a69d0b633e", + "ms-fd225482b96686f51dc8", + "ms-c07472c627b51086ff33", + "ms-f07607bc0ce1c4e202b0", + "ms-23d6e5835aa90a62b6ef", + "ms-1309e1770b6ca6e71c7e", + "ms-52c6c227fd4b61609031", + "ms-8a22f0da78e59cd1eb5f", + "ms-4acbf69acd30d96e46a7", + "ms-1518756d5da0b8d3c190", + "ms-417cab1470cf46273b02", + "ms-ee05717b29d39d7cddff", + "ms-7020f1cda882b79bffa3", + "ms-91d86cba2602d1fa780d", + "ms-52cdfff790d71133c5a0", + "ms-caed5c1b1c882141e2e9", + "ms-766b6fe7c08a3525ce94", + "ms-ea72196e20393870bcbb", + "ms-b6827e38768a63212d73", + "ms-a01e21b91a5b37a16afe", + "ms-c425c810f7865e9564ee", + "ms-f39bbaa063044719a98d", + "ms-c3cdd70e2bc10bba1930", + "ms-ed9addab41f3298db7bc", + "ms-eb16021deb5d03f15c9d", + "ms-912718b1f2ef92066799", + "ms-9783ccaf76f1b20c766e", + "ms-ffc3f18cbc674bdcc085", + "ms-3c2515d4a10082514714", + "ms-bef9b88ae40ca98647aa", + "ms-7a17ceb004979b99a586", + "ms-2be3bf4f348678c0cc6c", + "ms-e8fbdbddd9200f5091d6", + "ms-2dc8e489908f72903319", + "ms-0ca11ec8f7fbf0c1d66e", + "ms-f653a3ac16eba9fba5dd", + "ms-e058c0dba52df04b8642", + "ms-c0bfce3c9dd10f9d3231", + "ms-296b75a11f61d61d5fba" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000001.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-8d1978e8c3540cb5b1f6", + "ms-6355260c6fa49f291729", + "ms-0e275939f12a9dcee0a0", + "ms-fe2964cc8940ee7c24aa", + "ms-11a2794efa24f4d93d83", + "ms-3731f1c5af3c4a133633", + "ms-090e18187bb6ae4a9120", + "ms-4f938bd3599c372c8bff", + "ms-362ba7cebb2bab1421e8", + "ms-70bc2fe01ba217676b26", + "ms-c70be4072fb4e49ac106", + "ms-625b5d55ca76cc565049", + "ms-0f219ce35aff264455e2", + "ms-d9e3a187dcadb9ad26b1", + "ms-7e6493bb6b1ff650d6de", + "ms-4276f24b1b5dce771bff", + "ms-4385e9c56007fdaa94b5", + "ms-2e5c2f75d3f6311e0a1c", + "ms-bb60defa0d234da9e0bc", + "ms-a509a572b3871e591d26", + "ms-f5bae191ebad7c9f2687", + "ms-e589f7d40c12eff1712e", + "ms-17330235d51598ebcb4e", + "ms-910364d2d86091c5e0cf", + "ms-5316151a88a3282436ca", + "ms-6b9ede8946200ad0cd5d", + "ms-5f98f13f17e8f41f7436", + "ms-fbcc714230de2c097d0d", + "ms-023e021f079a8774c741", + "ms-cc0b33f92229d32bd84f", + "ms-a196fc3a729550509b18", + "ms-227fe0347d2ea57a66a3", + "ms-10cb98b0a1ca5429dd28", + "ms-a5a651f9e28ea52016fe", + "ms-42edbdb8ff281fc33a99", + "ms-d849d754f346b826e052", + "ms-0c7dab7ee052231f3f57", + "ms-76b9f40367a1ee020744", + "ms-66540aefd55a84b5088f", + "ms-bdd1924b44de2db5b07a", + "ms-8af87d14945da6a910e4", + "ms-732819b0c220a4c06870", + "ms-7d962acc9cfbc27c599c", + "ms-b6df140b235c8fe0eaad", + "ms-547f4204888aee978bf0", + "ms-723a1f2e22afe15b6981", + "ms-10f2865bef6da7b66423", + "ms-a8e6435948250f5e163d", + "ms-351818a836474e594d3c", + "ms-42117385b9e6e57c7e63", + "ms-4c83412fa0aa5180ad89", + "ms-1c0e2bb3a811fac6d567", + "ms-03cdc738473ac1df9dfc", + "ms-fb92a1e6b67121783fff", + "ms-759ba8b7be233ca47c30", + "ms-c92e8da7621fbf50cd4e", + "ms-55c2bba961a327aa5260", + "ms-0ece7b8c91d9d82f244b", + "ms-0ca1c25fddb7da05d80a", + "ms-8f0e2c9caad14bd37ecf", + "ms-68e331088400739273da", + "ms-83d2b158041f4eed3646", + "ms-ba156a5cdaf40cc85cc4", + "ms-2cc8c1975139989940f0" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000002.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-da46296bd4e9032570df", + "ms-0843bef6b0c13c3ede69", + "ms-642670f4bb658d32970b", + "ms-76c78750d7abc6570487", + "ms-197fee321caebfd320f5", + "ms-7571c367fe683b0f1cf9", + "ms-13192a27f5fcf4fd24d7", + "ms-bfda95ffffd5d1a2683a", + "ms-4ca84594a789c9abd5d9", + "ms-c5d026dffb39f2f1d38a", + "ms-91ad1a5f945470b4346e", + "ms-263556e167dc363ebc42", + "ms-e07f80b44faf328c5e58", + "ms-8a74cf665984273be00d", + "ms-4cd1ceea2dda2732110a", + "ms-dc2a8fa76e1ceb0d8925", + "ms-746436cac05100783ad7", + "ms-486ff31b7a04a7e9e8f9", + "ms-793a98188fd57ed5291f", + "ms-2af09f91e05d19f7f181", + "ms-999a8c1454dd4597a834", + "ms-309c09cf5d782e9ba0a3", + "ms-309d2258d275ae304857", + "ms-fbb609a3f2a334986245", + "ms-2ae57bab6d2ea11f6638", + "ms-7849511ad9b8262adae7", + "ms-28b1f1e4380d434605de", + "ms-c35a083df7c08225af3b", + "ms-14c40aa19f7b9307511d", + "ms-ae6d60345612be32478c", + "ms-f77305a344da6a4ccc84", + "ms-fd6deac5a281d169bc75", + "ms-6c335deaaa5c63478701", + "ms-15ef6ee70a8df3cfa578", + "ms-5e23219a8e2aadd039c7", + "ms-d26a60e383a2e822615a", + "ms-4ddda51b1c6809c3531d", + "ms-5e85dc3760eef8a53a67", + "ms-704fee4c9fb9941a26e7", + "ms-578e24198ed828a74485", + "ms-95320371ce6b820c85f6", + "ms-910d27b58a0585327a27", + "ms-8fcc78fba7b3bec0fab6", + "ms-348af0b70d3e679f40a2", + "ms-25839c690d85a4f5b986", + "ms-a0e07d6e4725342ffff2", + "ms-d965f63e091a5d6dd6d7", + "ms-64ab703e7a82799a47dd", + "ms-4e395d53889980b05784", + "ms-a3d0c460532de2d801ae", + "ms-dea17c1c383f54d05636", + "ms-85eef30b36b044b33f60", + "ms-e0b2e69a5ea0d0ad4fb5", + "ms-81729abe5d83a1ca280d", + "ms-4d6099258a3a6e037aff", + "ms-307f9067d5b66f55b3a6", + "ms-eac0d7dd54b1e797fdbd", + "ms-774e954ffdc5da041261", + "ms-ae5421846020d2765f82", + "ms-b33f1d4ddb6a531d2a71", + "ms-ed419c64a6693cb8998b", + "ms-1b9f208b8780cb33b9a5", + "ms-80d5c509e83845ab1bc7", + "ms-578537600f7a649c2dd5" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000003.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-3a07c49566decfea04de", + "ms-295de2ad6c52ab50de1a", + "ms-ff3f17774184180998f0", + "ms-59cc3cdf6918248d7f78", + "ms-08858c7ba8d4a3357e05", + "ms-5b122bd43e749623ebe0", + "ms-4d03463bb4536a28d67a", + "ms-1abd8c7cc1c683f88d3d", + "ms-b82cc8ce5d9d348ba6b4", + "ms-4d9cef437c77474dc62d", + "ms-99a5ffff5b73cf604b3d", + "ms-4a0ae3deadb8560fd315", + "ms-5daa8575a80e98a83170", + "ms-d8b4090565126a66d13e", + "ms-9af03b651b5e36ac2ae3", + "ms-50f62a61a82a96061d9a", + "ms-19e6e6e341aebb7ec693", + "ms-a6b029a958e91306966a", + "ms-d4c3c235768529c33843", + "ms-570c074dbf28b5faae58", + "ms-f39e1d7a2d636e39a02f", + "ms-c1ac63fb3644e3a2ae06", + "ms-a10d8368c33b56390f77", + "ms-b05147302972dfd7b450", + "ms-c2fde2d79cdf38555865", + "ms-b786173f210fcaa65a6e", + "ms-1781d9c562b4425df7c8", + "ms-6558f35ae2a2871a4bfd", + "ms-78da3efba7a739dcad14", + "ms-f763382872d8e89017c2", + "ms-c76158673d20b18c2622", + "ms-982922e263d07c54c136", + "ms-b8276cb9b6cb9a9b302f", + "ms-61b1188974f35233b883", + "ms-763c73c94818a090d358", + "ms-de066fc4200163ebd06b", + "ms-ae33e3afb57e61026d72", + "ms-479aa7d4248a0ebd94e1", + "ms-7fa68a6085e0b973ce6b", + "ms-8c8198e44eea32f98ca7", + "ms-b7800060cd4a8c3cb6b9", + "ms-b570fe64df70560c8ab6", + "ms-0e4af96af7d344ec47a0", + "ms-a2c5c9fd35d9fda5ce9e", + "ms-70342af6b601660d7c69", + "ms-7fd4a05174dcc8eaff38", + "ms-61068da67676a95e5065", + "ms-351c397795d9af204ec6", + "ms-41279eb586ca5e587474", + "ms-e793d78eaff3fb28deb6", + "ms-48d490ed21088bb5ba85", + "ms-0915b75e9fd1468d464e", + "ms-d988976b029886d18476", + "ms-ffa71e2ef6c284477c5d", + "ms-a1846ef1f00eb54e2d60", + "ms-39450dbba738dbd8a764", + "ms-2ee00a97b477ba241dba", + "ms-bcc70292f404cd1c7f75", + "ms-99ecd2226d7187416316", + "ms-38841d5926ac5f2467e0", + "ms-a1d69e0b112064914357", + "ms-2259a78b2187e22f9f37", + "ms-60a28a9b282e7586d4d2", + "ms-44252df6095f31d8b26e" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000004.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-fbcca760f5e16aa5c776", + "ms-18fb80871b8e11d1ae9a", + "ms-561f4f913d7b4991d81a", + "ms-11930c353a2b4fefab82", + "ms-706943f06888dd5bbb3a", + "ms-d69b907e65875420bab7", + "ms-cdd69b5acc8119e600c7", + "ms-ffba357213945be5f9f7", + "ms-43ce9f798efb923a120c", + "ms-a3a13ba6243be437d674", + "ms-b40ee5a75332685d5cec", + "ms-9ffeb1600c04f08bdfcb", + "ms-f1fb6cdb836269adb2cc", + "ms-f5e621c9ae7bcca8e596", + "ms-8b5b5e5f9237c41e0fe9", + "ms-739a89fb4521bc481a9f", + "ms-7de8743758ec505a3294", + "ms-8b73e5dfa1f2a01e6017", + "ms-82e28eb350e21d43d3eb", + "ms-351844f358ea7564d27d", + "ms-f4fd59b4d1bce2d6b6a4", + "ms-10f4c68dcad8285f4acd", + "ms-b90d7926b23eabbbed04", + "ms-d6e14cec49894dd329f0", + "ms-c9e31a19a776a861f9d7", + "ms-7cab119c7b8019254ba4", + "ms-c7156b26670838d40934", + "ms-96da3b7a3c2a2411fa66", + "ms-e7f2d1f0de50c6fcccb4", + "ms-5367a8bae4304851384e", + "ms-c679fdafdd776a14b66f", + "ms-03ed7e57992354e27896", + "ms-9b4483fb034248fe882d", + "ms-7d81f3e8b6a42695871d", + "ms-9b15f9c696a2661346cf", + "ms-1fec069b4a095dd0425a", + "ms-2f57eadf546710065684", + "ms-f682952c6ecb2b833c35", + "ms-51873059990aa55691a2", + "ms-d7a32da3ed5d010654a0", + "ms-283da8c8dae2e161c439", + "ms-4c7138016830a563fb1b", + "ms-519fb779d24a0bdf3809", + "ms-a01f0549a18d924d31c0", + "ms-8af2224cc82636ef1dd7", + "ms-9a2d8f0085e6740cf5b0", + "ms-70381721e736d44afeaf", + "ms-98427ae2d94c1f85811a", + "ms-1df0775c4ef08d0cf1bf", + "ms-9c8a90334c3b933c8976", + "ms-0adce283894b0a2b7b94", + "ms-7241a1c25af41f4c0b5a", + "ms-cdcdb32d7c02c49869ef", + "ms-4746e04d636d13cfbe6c", + "ms-478f754a15221314d15e", + "ms-9184c93267b4bbfa4cae", + "ms-2df0b9d5fa91156e430a", + "ms-c81880a77c4d978a0b7a", + "ms-6b750c1bef91c5d15526", + "ms-8f219d526891ae3c2dea", + "ms-6f1ec735350c81ef9fca", + "ms-7be1c6c8ef67031b1f90", + "ms-21155188bc35f0f021d3", + "ms-6d7b10e8b2b0a516b786" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000005.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-9a1de6130828dc946a51", + "ms-8c0b600aef5393f929cd", + "ms-4082ce1f04dbacd6a5b9", + "ms-0c9c39c462522b25ef2b", + "ms-1bfca22d4af29f0e14b1", + "ms-9eb40f856528204bbcfd", + "ms-eb48573b3e330e44504d", + "ms-61ac2be70a9798adeda0", + "ms-7670f25745871759e69e", + "ms-e15eeda61ed4064a5c9f", + "ms-9eea613e8158816f3193", + "ms-c774aeee2c4066edd668", + "ms-8b98e8aa8ecee10bea27", + "ms-f28fbb82a51d4de7811b", + "ms-772768ec2f9792127251", + "ms-7fa85a7251186eb18d0b", + "ms-2b556c4c1b55a20b14c8", + "ms-97140da7cf409922f6ce", + "ms-de0321d0ece4c58bc224", + "ms-d33444117889b5aa9cfb", + "ms-d5bfdabc97e56c680680", + "ms-51f2c34d5f8f132af1ed", + "ms-bcd9be34f050d97499bd", + "ms-8bb2d8cb7c0d33fa0c1d", + "ms-339ae3ec9482aa259eec", + "ms-8f6d30c2ade00e451dc2", + "ms-cb99eeafec52800f3b77", + "ms-cddf387bbcd6aed98fca", + "ms-5a7891a58bac94796178", + "ms-a2908e99b0a12a9980cc", + "ms-8420ccaff5986c85063f", + "ms-1c4ab2b6f004fc8fa66d", + "ms-09de52dad7e40970c938", + "ms-2073841ae01f224365a1", + "ms-943b5187cbe128b9b563", + "ms-8c6e41e7542bc6924a1f", + "ms-4fad1f976cc64cc06826", + "ms-382000b79932b7c68326", + "ms-753929cfdc4aaa652b32", + "ms-ac281cce8f6e42586359", + "ms-b7ed399f40d2d554b16d", + "ms-fcbac5ac63b87452c719", + "ms-cca07eb504a242633931", + "ms-f3f6a1b34fdf35290a6c", + "ms-62b597990a3d637e5a8a", + "ms-48cdfa31865ed70997c8", + "ms-45ed57caa698d87a99cf", + "ms-a934ef3c442100d9cc58", + "ms-326f54ae8f73791c49e7", + "ms-ec498d0c649d108513c3", + "ms-b6fd2931116c15db0f0a", + "ms-1451f536f2ac0df844f6", + "ms-08c75f74bf7813b1c257", + "ms-bdc942e2c7fb26b79162", + "ms-ce992c1fa0bcc38512b7", + "ms-488a34fe7b7e84250616", + "ms-97aebe1dd42939e3b340", + "ms-5f5bf12b5ebb67eeb361", + "ms-bc30ca6b7675f4de5041", + "ms-5ad75923125806db5404", + "ms-6995d4366e7c20010126", + "ms-64e28ad134f73fdce800", + "ms-8d94c7c692c91a962f25", + "ms-71452ba714a10880870b" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000006.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-28817fe297f38fd9550a", + "ms-69962cd1c448094e4fb0", + "ms-7d32f6b32c14e33b1eb1", + "ms-a3c53708416dbfb9fbc5", + "ms-b9e46b7345a63ded4397", + "ms-3f4967102c7b10057ab9", + "ms-cdfe8f580576f4c577b0", + "ms-f101cf5cc8575be3c201", + "ms-3c458f671129e1ebe7b4", + "ms-f6591d52d6824c6d2909", + "ms-4a7b774ee02339c4ca19", + "ms-7e94769c30229d5cc7b8", + "ms-3abfb4c309360a304d54", + "ms-0d1db6695acdf1ec3eaf", + "ms-caeb59971565b08f2513", + "ms-49219cb8ed7dad70c897", + "ms-fbda7bc51ca269d1cda4", + "ms-d19c9e3ab1d7daf06344", + "ms-dd1311efb8ab2a0a2d23", + "ms-950df0830b2e57010248", + "ms-813cac6d08ed5af48b12", + "ms-bcfeba2af1e844c22aed", + "ms-af0341467b22ec51e557", + "ms-175375383979526d80b5", + "ms-53fdabb5d8ef20778214", + "ms-b97cb4ec02e8234b08c5", + "ms-6e4cc893fcbca1b385e8", + "ms-9e1f3c14116ebca6164d", + "ms-f5fb668380b64d4f1e74", + "ms-5e6d18ccfcf923f52f8e", + "ms-ee19bfbcb2f4fbf5a44d", + "ms-61781d30cb05657c5f5d", + "ms-3c4ff3dae26217f22a8b", + "ms-af16712eadcaebf175ff", + "ms-c4cb99b887c88c707ba9", + "ms-dc5a883a3e28b90ebf1a", + "ms-867fd9318b2445dcefb5", + "ms-8d111990da24dde1403b", + "ms-54a7c3d32b6a267d81a2", + "ms-829156e8e688a433bf4f", + "ms-4bc023d8d1f7748b0e12", + "ms-08c39d9cc9eb7578e6e9", + "ms-054a02413d31389d1b7f", + "ms-dbb32d9fd6999e09f589", + "ms-382717cd5bfad063ddf3", + "ms-da855323619ebf257b19", + "ms-4bc32bd893d0cff76423", + "ms-f34c69fb0bf1b1055163", + "ms-a7f98a43ff113d74f009", + "ms-2c9685cffc73e47cc8cf", + "ms-cd71a93dcab09ba15592", + "ms-92a126060ec4a40e2c19", + "ms-3ea30989082ec714e12d", + "ms-80fbba5636cb4f9c64a3", + "ms-978bbf0a6fc563f745ab", + "ms-62ece8fe2f85092ca96f", + "ms-0ac53e56b4852c85ca2e", + "ms-25a8e32c071be83d913b", + "ms-36e8cb2e7f34f8dceb05", + "ms-5c8b50523e0de37d80b8", + "ms-f4d6da245d340cd16d30", + "ms-5db8a7cf17e5fa84dd5b", + "ms-191a2180e4ad5bbeb3b1", + "ms-6a6a41baedceb27ad888" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000007.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-50d83ce148904504ec50", + "ms-663aa9934c0f923cc125", + "ms-d84a4cea86526971b83a", + "ms-c7185c6986fed4eee76b", + "ms-116bd511ce8862ab73b5", + "ms-31782020143c24b4271b", + "ms-5fc884c238818acdde08", + "ms-b36673c3905a123cbd50", + "ms-0019f809318a077c3910", + "ms-8d9b2a5ccd43461f827e", + "ms-53cfcc878e30fc078bba", + "ms-74517e19dfdee44d2f81", + "ms-a7179e4e4e2a0d5335b3", + "ms-c2bc492ab5eb87a73789", + "ms-6905c1d3bc8157697499", + "ms-386220c339b7de7f222b", + "ms-f80ae137a8d9981f8c98", + "ms-3e51a1bdc3a2cc4eb06d", + "ms-fffe378a3bad31bbed52", + "ms-3c6ec8ef5145ff6274e7", + "ms-169b1786513146593ea8", + "ms-9cdebdcdad56bd8d87cc", + "ms-4015c296239564108da2", + "ms-6c0ef0ef8b493232246c", + "ms-e45b5de01eafae73e492", + "ms-8e741ca11a7fbe85ba1d", + "ms-d19573dc65181befd83d", + "ms-d0e4bdc60ab8ea52a320", + "ms-19c0b893fda0e6bb05f7", + "ms-81085eba1dd49b17ec9c", + "ms-13557287cd450bc12c6f", + "ms-2b3e081263ed25d45dc8", + "ms-ccc102ca65219d1e5f02", + "ms-8b85ebf659960395db9c", + "ms-493f7fe7164e2c3e7616", + "ms-ff3e5d72f3adf548d0a1", + "ms-f6db4e8bed63ed868487", + "ms-b1bf2ab59a5385d22e54", + "ms-8522a9364bf70311450d", + "ms-c9c1649478107785e9d2", + "ms-fb228a036fe5973bcad1", + "ms-3e971096794f04f325f2", + "ms-60c7d1d3b5701e7f4e48", + "ms-3aaa0d68ca26b1fe6626", + "ms-3b0639bfe64ecfd35594", + "ms-a44a2c9eee8baea5d29e", + "ms-8a00723a184a6ca740b2", + "ms-b671fa9fcef84dbf2a13", + "ms-b35643e98e1b846e27cf", + "ms-cd25efec908bd600e45e", + "ms-f32035837f7df88396a3", + "ms-c684df0b998b593db723", + "ms-46f10a29b484c9461204", + "ms-29afe3c91faef25f1343", + "ms-1735d1abca4a91293d78", + "ms-4fe7fdaf946c97803c82", + "ms-d6b2d805dc063382ade2", + "ms-128c7c43e30023e6b1a3", + "ms-81f48d29d0404300ea93", + "ms-0377c81a67d9819ab96b", + "ms-5085e45e6ca37192d2e5", + "ms-eaf612349e2da547ee98", + "ms-635a384c88604ad7ba4f", + "ms-503d0b1d27ed78b6130f" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000008.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-11230e3aaabb1feec834", + "ms-b666f24369c548cd56ab", + "ms-10af171aa279a879207d", + "ms-c7041c6f758f8e4a3698", + "ms-4c7e5a40768c0e27027a", + "ms-99df75798ec5da5ec7f3", + "ms-fa911d56cdf9bf122b5c", + "ms-ce8d2ef7d19bed381822", + "ms-ac79bc3a841958ce127f", + "ms-e9f2c61c28ddea3e16c0", + "ms-87f8bbc8d9e7f4dcaacf", + "ms-4b970be37419ce0be0ae", + "ms-0bf5f9b5ab423eb0a7c3", + "ms-6029bab18ae94b126d79", + "ms-6fc6a26da9fbe53dd616", + "ms-4ec67456479f53e605e1", + "ms-0b7f8ba789c00b37caa0", + "ms-f7c9f440fec3a22cd6d4", + "ms-d91cbe61314c18955b6b", + "ms-e20f5d43f319248c5001", + "ms-ee32e6bb3c0bb45702b9", + "ms-d1dd16f03551d840a922", + "ms-49641ecf3e1326621609", + "ms-a202dd8923327af132da", + "ms-f947dbf28ad85666e2f1", + "ms-1b83fc0981e283975488", + "ms-56ad75c28031e049a3dd", + "ms-4895312f7ef7a8e0109c", + "ms-7af1e453acf970f40546", + "ms-d305a11348461ae1622c", + "ms-9e87481ce12fa45b2554", + "ms-dd4e13e425fbe9bd9140", + "ms-12ac4329c5bca5f2bb62", + "ms-3a64af2410ea932a6267", + "ms-6133ea5e61638487dd46", + "ms-4f405775dcbb88f7d8b0", + "ms-a73f46824bdbdfb96238", + "ms-17e8d6ba42404c7b7d07", + "ms-a166a9b750ec3446beea", + "ms-1f8e2d5d25706493821d", + "ms-1d36e5be82aee5a8e3b8", + "ms-30fe784e67f732ef9a3d", + "ms-a949424c5fa615565d72", + "ms-4705657703da9d791fb9", + "ms-16cd0a94ad2cb78b1ada", + "ms-2592e5fa970f1b474388", + "ms-d824135bf97b907f7cc2", + "ms-a699ee643e76a610b444", + "ms-c7020fdb7cba871e8b3e", + "ms-3deb9fa600781abb51c1", + "ms-2bf2f3a2444b64f8230a", + "ms-81c855f8645b35f0856d", + "ms-ccc0326ed73cd2ea6bf6", + "ms-90a03366976eaa1ad013", + "ms-81d1d8def1fbeb3b1cc5", + "ms-7a745e34fa5a8d527e77", + "ms-46987cd8daa3091c8a85", + "ms-1cea800b74c1ea2f6338", + "ms-cafa65eeb953aff47f30", + "ms-d6305829d7ca1c68a3ce", + "ms-30998855587c076a2a77", + "ms-daaac2887015a6e4f3e0", + "ms-3c12da28c8edab3fa554", + "ms-cd340bca589f0cebd7da" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000009.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-61d4dd4acc03b8d30f19", + "ms-3c8ba3016f7269bfe44f", + "ms-cfcd7c9f3afe2215588f", + "ms-9cf068e0d2896f655ad2", + "ms-ad5ad5314850df7b0853", + "ms-22325628db75967cf779", + "ms-1a4ecfbe9c97214d5c78", + "ms-1d1969094d91fe52667d", + "ms-fe650308e25459ee64c6", + "ms-4b94e0333251263012ec", + "ms-e110d8ef3eb852a0f564", + "ms-f3d9efbde383ef5f742f", + "ms-4681331b496c8a9cf14d", + "ms-e318e13f77e75e7cd802", + "ms-486572c62484eb4d73db", + "ms-88e10d1be12a407b3926", + "ms-b23839ace2fb08438144", + "ms-1abf5e6d483e2788b086", + "ms-b42e0a68875b0405904b", + "ms-0541a144afe57b5a7ae6", + "ms-6c081eee5edb633d6caf", + "ms-155ec4624d261f6ca16e", + "ms-8315ae30e2cf041c8674", + "ms-34495bfe34440ac63cf6", + "ms-02ab2e7a74d970154eb7", + "ms-c6d4f51c1adffc449778", + "ms-adfe268e473c03b72c65", + "ms-8f5255815f8b70b48734", + "ms-910fe3f3a4cbc2d8e6f5", + "ms-ed7b02bc8f37b8524a56", + "ms-d3cebfccbbe54d22c561", + "ms-f574815db5db73f40aff", + "ms-fe54b45b549a55d45c0e", + "ms-525a7841d99bba385bc5", + "ms-683b1d258a1cafbf87b9", + "ms-0584242e924d240ca0c8", + "ms-d4aab6f9b6717dd5712e", + "ms-7cbbc5e7dc896a803efa", + "ms-ffef7730418c831245b3", + "ms-127908707208e993527d", + "ms-de0abc7a53a48164f82d", + "ms-e14bfc5549c2a71bc7c6", + "ms-44694e99af9263ad3c45", + "ms-fd9f10f7df6a11ec2cef", + "ms-88b1a0f8f758f7956c0f", + "ms-71ef8f2cdc11d83250fa", + "ms-5ac153f97fd3c2f689c5", + "ms-87b44e6219279322a723", + "ms-8edf3c4f13fa7a151824", + "ms-a3cb33ed1e34d229773c", + "ms-8b972386cfdf0158a0a3", + "ms-d1aac7be45644e12a976", + "ms-3aa5793c449b59c77b6f", + "ms-41561ecd4d68cf77c933", + "ms-ee0986397bc7267882d8", + "ms-901516fd681c8506d80c", + "ms-8f5d99953d4e7790d226", + "ms-20e5ffa93544d6d8b313", + "ms-a497e44e8c2ff0d46e47", + "ms-1a801d8b486ef19642bf", + "ms-bc53e550cb202b5fcdf3", + "ms-0b4ac5b183639ac19191", + "ms-7982e5af0d533320432f", + "ms-dc68e92d0571b68f51b9" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000010.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-432ca73236be93cb76df", + "ms-a8864b25c00ee193d0ea", + "ms-b0d54c14a65fb73f08ca", + "ms-bee62a6ef02134938aae", + "ms-cba2f44df79669ff934b", + "ms-3d095a49ed03de972849", + "ms-3d0d437ee74041fc61bb", + "ms-ae4b9e801d15835ee656", + "ms-30ad687dc1bc8318fa45", + "ms-732d0515ed28df031b8e", + "ms-2458fcfffb3f96c6ca1a", + "ms-d807915e233a0c74b809", + "ms-e50888adb67fa64ffe20", + "ms-89c0c15f7bb38eaf59e4", + "ms-d9db15929a99dea0e9b3", + "ms-22412b51e61621b6025b", + "ms-87f65c35c9d8b3ced6fc", + "ms-a6060f858cb9f93c5588", + "ms-4e86b87d6e5f6ab914fe", + "ms-0e12a6083750bfecca87", + "ms-3e7befcb3d8fc7884e8f", + "ms-9ff81e54049443a0f3d3", + "ms-88bdd8987f890af7e9ea", + "ms-045f5adce6d38ea2794f", + "ms-311a52ec4b0ca8d71889", + "ms-11bb3387454624de8143", + "ms-869a04e41dc5cc90916b", + "ms-bd6b2dc41ba42b44462a", + "ms-0de368e8d16fd01f4e67", + "ms-9c7e6b057c74f3014fd5", + "ms-dbcc3da47d96ceeaf9f0", + "ms-d67409d8ecfd5d97336d", + "ms-29f1abb869ff206a4d52", + "ms-7e0ef0b7f94396bf0029", + "ms-9233eae706910c204b79", + "ms-de709118c78966c5aff7", + "ms-a5d81a7c740394b68c3f", + "ms-0cfef0ba35d4e6097003", + "ms-dd4f37e032d56b692a0e", + "ms-f8439d39ae9ed028a4cd", + "ms-b4af4c8b9dcd6976546c", + "ms-fef1cf4879be83a77574", + "ms-af22292595eebbd60e48", + "ms-1136e4a92618d203f415", + "ms-ca25b016439a98198434", + "ms-f5aa7ce332ee9702812a", + "ms-a58222066dfcb0676e3a", + "ms-fd1c44c6d5d8f7f11451", + "ms-1f458832fc85f00873cc", + "ms-9b213fea5cdf7a7c91c4", + "ms-6b8adb335c80e03d5cd5", + "ms-9de5ccc7f31ca1c2154b", + "ms-9b56e5007c76d9fefd1f", + "ms-ad058e704a159164d7c4", + "ms-54f34c73458181a65e9a", + "ms-c2e07545267c612a022c", + "ms-9044026b69850d12d05b", + "ms-bb0d941f78d3d9c5a886", + "ms-7361bfba9d2ab4f1913e", + "ms-a323150b5305fc2a329d", + "ms-7c5059db4c14a74b64a3", + "ms-00b9ad6a8fe73fccc6fc", + "ms-5780f9e1e9f9e740d474", + "ms-0a48fbd111be17b50be5" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000011.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-f69d6aec4ddb12d88c5b", + "ms-bc615957d9bf177dfe26", + "ms-1d5b929cb9614dc18304", + "ms-7ea16036a1dd971fc4e0", + "ms-c7f427f37d47501932e8", + "ms-45d4b692eda9ecacaa55", + "ms-a3fa10726c656e851515", + "ms-38ceda08bad38fbebc1c", + "ms-e218d19c595e580c3e32", + "ms-79c7332b2b329eaa6a13", + "ms-79d500110f3723b57da8", + "ms-9e53ca6bc5ddb9133ffa", + "ms-1add4cf84fdadf5703d6", + "ms-12050326f623a94ae9e0", + "ms-c3a1cf0e424a4f065459", + "ms-47a55ae1d3a858eed51c", + "ms-d686ae387b1a0ce2f2f3", + "ms-90464666adfd66056785", + "ms-f2c108db42ad23380f14", + "ms-166b7f66bf9f4eec432b", + "ms-29eab654e4adf7ca9c4b", + "ms-239a07967247700a8ed1", + "ms-92e98f36af84f7eb3b37", + "ms-288532d255e4828bb15b", + "ms-ee055e1282f6302ad9e8", + "ms-0b1ee726defbe3c1f5da", + "ms-326a7b583623d136470e", + "ms-842e179244eddac303a5", + "ms-41c5500153912a1743de", + "ms-a43a663a3c4b15ec3ac2", + "ms-a68740208f81c5d6c97b", + "ms-6c39b520dca0a325535f", + "ms-eb206a3dda0ac16e1643", + "ms-cea745484a24f376d7e8", + "ms-f8c4e700a751537ebb16", + "ms-17a04b4ac2d938261bff", + "ms-edcdf99767068f94bcb4", + "ms-338a891eb2a21ec4a63d", + "ms-3b0ad4f07f1813ecf670", + "ms-225fbe08f2933f91d603", + "ms-d79154664cb93c96d8fb", + "ms-a13ee9983c139d34fca8", + "ms-ff96d76fba5b6ceb1289", + "ms-478b7b8b729f9212a522", + "ms-9bcf8a1e5e7b0f7ff6f5", + "ms-68ab6033433b360464f8", + "ms-fe9347448b7f810d1353", + "ms-17a17b212837b7f1309f", + "ms-13af3f2f8e0b8444fe68", + "ms-32435d71bb817daf01d9", + "ms-0e9908c8405f241086ac", + "ms-2b64c720290c6dad7e35", + "ms-844b753810c45699d501", + "ms-2dc4f12ea7e8567e3a6e", + "ms-97598e118c5cc8c254a5", + "ms-aaf1c4931e054f52623c", + "ms-d8583c895a59df2871d2", + "ms-acd4bde23cf4608facc2", + "ms-99b15d50660b77fe2fcb", + "ms-d05d9bbb98f24da55981", + "ms-4f833710c2706ee077af", + "ms-10a5b6cc3ec06dfdc5c5", + "ms-8f241450069e74460a40", + "ms-9320da73e81045277fae" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000012.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-42f83ea069c0558ea81b", + "ms-2cf1227610d4602f88ed", + "ms-2f846e9f31e2b8c4965f", + "ms-a76ae18568c88af1da6b", + "ms-3a65e7b378a87a7eeb11", + "ms-859d0ea38fb5e8aa6b98", + "ms-16d5b122db1824caf92b", + "ms-7c76bb3c93411136f11a", + "ms-5415e349e8f9f2bce641", + "ms-8330999c6669ec42122a", + "ms-cda24997adbad130b09d", + "ms-8fa24da5c97dc1dad5c1", + "ms-dc965b3a72976440156c", + "ms-aed881d8d61a756aa297", + "ms-dfc737e32df79e1fc429", + "ms-1fb2f4c1e17e99767558", + "ms-23f9cca8edfca421a4a6", + "ms-d5d7c690819019cb7869", + "ms-d4bde31fd7e7cee76b09", + "ms-3cfb09df08ad7279ef02", + "ms-68970d9d9bf975e25eb5", + "ms-3950f823020f08f8e8cd", + "ms-01db5970aec167a8ee1d", + "ms-8beca0622be320b282c1", + "ms-0bb8964d3e6f9dd302f4", + "ms-c91cae30acaa8dc7a12b", + "ms-3b7c5c91ae441f282006", + "ms-6053d8ef243250b88f96", + "ms-40cb974de9da58dcd32b", + "ms-07e0726dedce0a0e4419", + "ms-7d18f6f16879d2933bcd", + "ms-802e63c9d42228a69c22", + "ms-e0c2c1fbe39cd1d652b3", + "ms-a26c2ffdac089c72ade5", + "ms-2470946f31fee328ad6c", + "ms-15c083989d8277fd07ec", + "ms-5425964a186feaf6fd99", + "ms-2c68b2b895bdbedbaa08", + "ms-3eeaca32feb66f450a0b", + "ms-8ccc51fb09484026ebbb", + "ms-319a6c0cef0ac1a872d5", + "ms-350b797a531f87ad2c0f", + "ms-b59a7ccfb736df4c4c68", + "ms-4d43fb19477b7ab848bc", + "ms-283d0859f6487c6f4718", + "ms-cb16b458af723bf7cd0a", + "ms-c86956f3f1333e136186", + "ms-0a4c4442237909d5f7b7", + "ms-f3ddc66fde2d59666e07", + "ms-f9c00ba41392adec3fd1", + "ms-8fdde5008d12abb7ae6a", + "ms-42895361db5578b75c0e", + "ms-d44f5bd2488ea182132b", + "ms-f625d521be1e85b1c49d", + "ms-0cc80a2933abf1b4853f", + "ms-d462cff6e1b4b201a6ce", + "ms-a6273302894af3505ca8", + "ms-2dc828d2d76b389db4f8", + "ms-38507741c3b5d91a7b33", + "ms-66cc354bbf4e8bdd1051", + "ms-b863927293ce6530bb3b", + "ms-12a997dde50b7a7f35cb", + "ms-7a1ce0a81e4f34902622", + "ms-f4b95f2b569e3f1bf086" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000013.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-e093f85b121d443b61e2", + "ms-41b0258c4e329a734da7", + "ms-2d995e39f2b95fed1a67", + "ms-01067b8ff16f74652382", + "ms-e2ddf4681b47f39e15da", + "ms-220641931ae709280b14", + "ms-2649df827a3c1017ccd8", + "ms-76194d089c97018ce427", + "ms-8b68d8b4d22e6ff0451d", + "ms-c78721d9c4196a9e8ec5", + "ms-5847c6d196024fce1770", + "ms-e1972ebf2c5587a6fdf9", + "ms-a45ed89f5b95d1e7ab07", + "ms-3737178a1f85502652ea", + "ms-c3a3764fb110503bd651", + "ms-97a18aeeea0b34f7c606", + "ms-7651a7f31eae24958ef0", + "ms-3bb844dc66f24591c8e0", + "ms-57ac136bd68be6387c6f", + "ms-fd99b14bf152ebc0ef8e", + "ms-0ec98183be760ea25a45", + "ms-eff5f21c51f26a85f07c", + "ms-50f29e19c798fee5422f", + "ms-b01677be2df71c474d1f", + "ms-fe1a5f1bc3c4ec7669f2", + "ms-893159691d11a40c7537", + "ms-72cf68c15631675b063f", + "ms-f07589760f67c1100667", + "ms-a9ac13384fe5afc835f1", + "ms-3aa864ead19c9cb9fb49", + "ms-9edc0261b6ca7980ff81", + "ms-c151229bc1dd03ed2feb", + "ms-59020093d2eeaf58aa84", + "ms-39af5447d3c90872c25f", + "ms-3352e05afd9065cd09bf", + "ms-16a7ccf3c442984c2feb", + "ms-32346c22352f72d3de3a", + "ms-bca6fa102f94cb145cd2", + "ms-13b3c492538e5b3c7bdd", + "ms-4b1561bd0029228fdbbe", + "ms-5e3b1693ad7b75221847", + "ms-b809a8756abec7015aac", + "ms-4177936dc3e043366181", + "ms-dd7f9cc9585878b84824", + "ms-73e36360221b6732c688", + "ms-ddb8f148dff1db91851f", + "ms-c3f7b780f604922f4f67", + "ms-ddb4e5a1f88dfd017319", + "ms-c6beb139304a0ea6f7ae", + "ms-06935a084a7af6d7596c", + "ms-03acca095918dea64872", + "ms-4742d169606f3ec2b474", + "ms-d73a2a984d5829e884f9", + "ms-beb3567748ba228bdfce", + "ms-87fd37bcc6b666e220ef", + "ms-02acbd591a961c0dbd46", + "ms-57a213acdc6a7231f799", + "ms-47662cfa5e75ac8250e7", + "ms-ebb616f9b291184677b3", + "ms-9bc153354c4dd2e5c1a5", + "ms-66ec6692758e9a7f2e01", + "ms-4dcb031152807390c6dd", + "ms-1fa37879a24f0e84c275", + "ms-1cdd36810a138190ce3c" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000014.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-4bf0620f8d1294466596", + "ms-ff23e4fce42c637ed617", + "ms-a576ca39e26c4e1ab240", + "ms-7a460f3c1b67817c19de", + "ms-f897cc2141a0a303fc6a", + "ms-86b302005e0461f28c43", + "ms-c2107b650136a20fc4a3", + "ms-ad10dacf462663ab68ad", + "ms-a145b708d79b243da21f", + "ms-5d169f059bde8e429786", + "ms-bbad1e7f073de07fafbd", + "ms-3fae81854f4c927f1f8d", + "ms-3fdce65bd5c5681403a7", + "ms-8adf1cf09ddabdc5fe5f", + "ms-ca10d8b9d03f4d59806c", + "ms-6f1db8895d5164cc5ae5", + "ms-7057de70d36f5cc31472", + "ms-aea09fb002af52e12bd8", + "ms-fcfb7dd60d966e20b98d", + "ms-618c9e3c06f4879ba788", + "ms-f4860c7945db15b54d2e", + "ms-55ea22014575410563a4", + "ms-630663751ab841a96d5a", + "ms-5071d49489e1900dfa91", + "ms-f1d1776b7a5bd0a7b02c", + "ms-c0995730ba36614c249a", + "ms-90bfce2cef85e101138e", + "ms-fa62bbb3e8416561ed69", + "ms-a5775906ce2fb6556103", + "ms-b83b2f6ff006cdd24529", + "ms-1158e6997b2f1281b105", + "ms-8c9f657875cf334b1f17", + "ms-add4b7d277e1fc171144", + "ms-72b7d76602371f044fe3", + "ms-b0fd98bae1ed5d543f35", + "ms-74868f62709b91077e27", + "ms-486b916e7b4ad1566788", + "ms-a47818abd17d958ad6ca", + "ms-6a7e80154a9a4e5e9c08", + "ms-f2433bc89b10851b7187", + "ms-84b677a48d338212ceb7", + "ms-cf274fbe4c657d04d7c4", + "ms-29653e236de96058ef6d", + "ms-b606af28580a38cea636", + "ms-5df54a8baaceb8c73c8a", + "ms-2a67eea0b1293de9bcf9", + "ms-a927d1bd67068f1d9302", + "ms-4aa0c1a37cdb70ba8877", + "ms-6be8614a1bd50b974b95", + "ms-47fd1407ae4e17ed4bee", + "ms-54a04864cb68a4c3caac", + "ms-b333c5aec11dd41d37d2", + "ms-1924e52af208c9f66592", + "ms-f1eb908e172570b07cae", + "ms-3417f5f2f1470c61d0a8", + "ms-7785d335d4a95e21a84d", + "ms-3b4b0b8c60ad41bb16b0", + "ms-060512e03c9833bae633", + "ms-22588909538be867fb28", + "ms-7b932868f83725def2ed", + "ms-cb70b56d6ad32b11c3cf", + "ms-ddc0b3adf8506c30f503", + "ms-27b0553a5c49108e6a98", + "ms-32b9c353a4890631f1bd" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000015.jsonl", + "record_count": 640, + "group_ids": [ + "ms-4f190184d3920a0add8b", + "ms-5be34d2caaa3a2d07b81", + "ms-5922c28e4a71a88c6d21", + "ms-4021797e9499e6aaab7a", + "ms-a073398629e6943ad1fb", + "ms-6f12a291dd9c3bce4b44", + "ms-06715386190189185900", + "ms-fed4628b22c7e42e32d9", + "ms-38d3e0274ddfa4ed7202", + "ms-7fe236a99d555408d474", + "ms-28d5618562729c763eb6", + "ms-c412fba3624c72fe2a27", + "ms-10dbf8bb0c239a44ec0f", + "ms-5721bb074a97f9773eca", + "ms-8e4755f1627d85c12bd8", + "ms-6af227c07aaf7c0f0aea", + "ms-a22fbce4503f088acb1f", + "ms-f71c7cf8ca320b8eb7a4", + "ms-b5deb1971bf3b3484282", + "ms-b1c9975d1a007dd03dc4", + "ms-26e09c86be9068a34acf", + "ms-d16053a4f8b95c94dc46", + "ms-236601c5e076b4ee5d75", + "ms-5e1e38c86ef8f564ec77", + "ms-d75a00b1388f50bf8919", + "ms-c104c2ebba7bf9b34680", + "ms-81951a4920488e6c62e9", + "ms-ce2924d08bdd6c379777", + "ms-ee7463f3d937dbc16f89", + "ms-935ff204de9424ec00a2", + "ms-2b56888f10733707f14b", + "ms-076d67ac4d2de4e0c627", + "ms-f3e40db5e157647d463d", + "ms-b0aa24be399d843c9d50", + "ms-b7e7fe66a6c86898eec2", + "ms-70d994c15153752d0dce", + "ms-782893a3eb2e3dcfd15e", + "ms-edda00243c6d6275fc0f", + "ms-2d4547df5dccc4c4ca24", + "ms-fbea3b8005a30431cb52" + ], + "format": "jsonl" + } + ], + "record_count": 16000, + "group_count": 1000 + } +} + +✅ PickCube-v1 generation complete diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_1.err b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_1.out b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_1.out new file mode 100644 index 0000000000000000000000000000000000000000..6e963d4ed365b10114b07386cf290e9a352f1b79 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_1.out @@ -0,0 +1,649 @@ +================================================== +Task: PushCube-v1 +Groups: 500 +Horizon: 16 (vs baseline 4) +Demo: /scratch/knguy52/dovla/maniskill_multitask_demos/PushCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5 +Output: /scratch/knguy52/dovla/experiments/six_task_h16_collection/PushCube-v1 +================================================== +ManiSkill CIL progress: batch=1/500 groups=1/500 records=16 restore_max=4.768e-07 +ManiSkill CIL progress: batch=25/500 groups=25/500 records=400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=50/500 groups=50/500 records=800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=75/500 groups=75/500 records=1200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=100/500 groups=100/500 records=1600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=125/500 groups=125/500 records=2000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=150/500 groups=150/500 records=2400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=175/500 groups=175/500 records=2800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=200/500 groups=200/500 records=3200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=225/500 groups=225/500 records=3600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=250/500 groups=250/500 records=4000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=275/500 groups=275/500 records=4400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=300/500 groups=300/500 records=4800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=325/500 groups=325/500 records=5200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=350/500 groups=350/500 records=5600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=375/500 groups=375/500 records=6000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=400/500 groups=400/500 records=6400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=425/500 groups=425/500 records=6800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=450/500 groups=450/500 records=7200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=475/500 groups=475/500 records=7600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=500/500 groups=500/500 records=8000 restore_max=4.768e-07 +{ + "backend": "maniskill", + "env_id": "PushCube-v1", + "control_mode": "pd_ee_delta_pose", + "sim_backend": "physx_cuda:0", + "source_obs_mode": "state", + "instruction": "Push the cube into the goal region.", + "target_actors": [ + "cube" + ], + "reference_actors": [ + "goal_region" + ], + "demo_path": "/scratch/knguy52/dovla/maniskill_multitask_demos/PushCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5", + "num_groups": 500, + "group_offset": 0, + "available_unique_states": 3882, + "excluded_post_success_states": 14459, + "branch_selection": "pre_success_only", + "num_records": 8000, + "k": 16, + "horizon": 16, + "seed": 0, + "measurement": "exact_state_physics_branch", + "candidate_mode": "structured", + "branch_execution": "parallel", + "state_batch_size": 1, + "state_restore_max_error": 4.76837158203125e-07, + "state_storage": "archive", + "state_archive": "/scratch/knguy52/dovla/experiments/six_task_h16_collection/PushCube-v1/state_archive.pkl", + "observation_archive": null, + "image_quality": null, + "manifest": "/scratch/knguy52/dovla/experiments/six_task_h16_collection/PushCube-v1/manifest.json", + "quality": { + "expert_success_rate": 0.852, + "no_op_success_rate": 0.056, + "mean_reward_spread": 0.9625078696906567, + "nondegenerate_group_fraction": 0.998 + }, + "sharding": { + "dataset_name": "cil_pushcube-v1", + "version": "0.1", + "created_at": "2026-06-25T23:25:35.752579+00:00", + "backend": "maniskill", + "num_groups": 500, + "num_records": 8000, + "k": 16, + "task_count": 1, + "seed": 0, + "schema_version": "0.1", + "format": "dovla_cil", + "shard_format": "jsonl", + "index_format": "jsonl", + "shard_size": 1024, + "shard_count": 8, + "group_index_path": "group_index.jsonl", + "record_index_path": "record_index.jsonl", + "shards": [ + { + "path": "shards/shard_000000.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-cefef59d2909a69dd709", + "ms-365d8c6de6de16f66744", + "ms-f261341a4a394ae06c69", + "ms-916b66e56329dbd05b29", + "ms-14c6516a0f22d20ef6b0", + "ms-e31553038464afab4ec1", + "ms-68dcba1af854c06fc0e6", + "ms-47d33f56890889e86355", + "ms-2cbe83da02847008bf24", + "ms-ca9a06f00b861bf34c4d", + "ms-be8665e211c9b1cc13b4", + "ms-01fe80a36517a246929d", + "ms-7a5970ba19b465cd86c3", + "ms-3d208c1e04bb84341859", + "ms-65d0541a2ba23a8478fa", + "ms-c3191d9c08e5756e90e7", + "ms-1aea30e99a4945020cc1", + "ms-2ba02dcefa15a9371706", + "ms-0818eb8cc217c9b3040a", + "ms-be7166dbee42ee877e29", + "ms-351f07141fc0a8885759", + "ms-1f6006a647438280a273", + "ms-70dd7025d2b038824c72", + "ms-72da580865509f75d6eb", + "ms-7efe34ce17cffdcb43d7", + "ms-0942bb0cc1143a2058f9", + "ms-adf081811af7bca643a1", + "ms-398ddbf730ed8830b08b", + "ms-2d9efcce135f0c9d2203", + "ms-698e09998c5dbfa6b7e3", + "ms-627eaad017d31a933984", + "ms-51fc0cc7406cb59b5c24", + "ms-272972565b3a686fc999", + "ms-5d4b83a1d21ee90ef69f", + "ms-da66340af8ac9f3852c4", + "ms-33905bad392928d3a217", + "ms-2c65718a6134a46b048f", + "ms-651bc0395d78d2819158", + "ms-d5f6186b51268063515e", + "ms-67f2977ab7d070052b61", + "ms-17723f5f0806414eda67", + "ms-1be8db140ae74d18596f", + "ms-e333fb7b8697ebeb474e", + "ms-281b1eae3c1ac42cb40c", + "ms-8efd9238ae8d6ca1cd1a", + "ms-85a4fc6e131dacc9d839", + "ms-88ba5f43dfa8f827ff0f", + "ms-82a203dfa8631639ba7e", + "ms-0952da757675889102ec", + "ms-7208ab12aba7b4a0da1e", + "ms-a6fbec34b57a97204a37", + "ms-686b6de62b641e0dba0d", + "ms-66d3cc0e17bcbfb5f36e", + "ms-048c14d42c3f3b352f53", + "ms-a0863466a999b4680737", + "ms-ad5107846b45eae5a0ab", + "ms-cbdf12e174c323015f38", + "ms-f49fbce5bf7a0a169f1e", + "ms-01236ae561190e7a5f8f", + "ms-88c39906a4e9c32ec286", + "ms-bad55a54ecb4349a5323", + "ms-1f5c191aa8a6978b4020", + "ms-05139cc8ec8867b8ea3d", + "ms-2b309050ebafdd5a3482" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000001.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-c109289291b46fd832b3", + "ms-3b99101d58e538cfbd6d", + "ms-c95a3e52494c7727dbf6", + "ms-9c49860c264bb0bec403", + "ms-65b3fb1b7924487ebe2b", + "ms-fc2e1272ed035d0b9066", + "ms-88de0a1035cb8f8ef16a", + "ms-1979b899f11945b46b51", + "ms-83b649656afa71ab6bde", + "ms-3028df2c22f4c029b533", + "ms-b2f6b95c91f5d672d3a0", + "ms-22f153f0d4fa656018e6", + "ms-3aa0b904cd3565b6c7ba", + "ms-605cecc07e3298489320", + "ms-0a001f1bda12891b219f", + "ms-9574e0bf814256a11f3e", + "ms-ad4cbad05e1be65b72ad", + "ms-9fd5ccb16c4e5d0c4579", + "ms-ed52ffcd2b3e222b5175", + "ms-d2aed645029072f5c82f", + "ms-9345a7952bca93b2be95", + "ms-e95a64701812302bd039", + "ms-4ed913628b5f121bbf34", + "ms-674d9e270ff9ba879886", + "ms-3333e935277a2e1f4e05", + "ms-c84dbd2afb3321e2f43c", + "ms-b0d0004818ae853d5bc5", + "ms-269731d0630cffc1a216", + "ms-a83154b77975bca93e5e", + "ms-69b8e0564f16c652d900", + "ms-e2058192b3065ce7dc60", + "ms-1e15448c374f47bee884", + "ms-f3b9739bcb72c30a5535", + "ms-fe5a68745bbc6b4256cf", + "ms-b5f45644a5e3ed45c86d", + "ms-9915a1632fdf65172a02", + "ms-74f12c8d99a319a9d020", + "ms-130f96b49e1bdb7c0e7c", + "ms-2867deaca7e98582620b", + "ms-3e86692ea2e52b8e009b", + "ms-6825a9112cb71ea137c9", + "ms-c2ae2b78ab7891e74c80", + "ms-dab80e1bb7987577e33d", + "ms-8c2312d2a272dbcbae71", + "ms-e825bdad55212cb330da", + "ms-fd66d311a01f9c4fc825", + "ms-30caf93619a0e2f94eae", + "ms-aca9b69b504ede73bcac", + "ms-4401ec9ab8c89fb39f58", + "ms-9d943a2732771576e41d", + "ms-54fdd49dc2c46e162fc6", + "ms-fa20766f8a3e706a8ea1", + "ms-d885d2943c81168cdecf", + "ms-2413d1155c87f0442485", + "ms-f985a8cb6c13af7df7ba", + "ms-4c806a5814c3e727cac8", + "ms-2db29812d012d503054a", + "ms-8682e5b9068871041b74", + "ms-9feb7d8fdc15e384b8b7", + "ms-ecf6837325ac3bc3b66d", + "ms-b2b4b951e47485c8be41", + "ms-a5885ec43a0aa166a0e7", + "ms-d593c823d7fa1cd02635", + "ms-39f133217cde17cf4465" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000002.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-17b435e93391674c7fdd", + "ms-84404222b954c48c48d8", + "ms-86508eb5d360c56023e1", + "ms-79e8bd81bad486786dc3", + "ms-e411015630dcff33d0f7", + "ms-cece69f5a160c40d613a", + "ms-8a38529670b19d605663", + "ms-c9a00cd02aa3ea31af1c", + "ms-27c9f70e97b8850cc3ab", + "ms-157e8e3ee1b0de1840d3", + "ms-04020456c994946d9a76", + "ms-0a3ff306e1bebfd51ae4", + "ms-2423891634e7b469eed8", + "ms-cffdaeae3b980aaeb8f7", + "ms-9d8348ba3054f969ad8e", + "ms-4d96769e78d7391e43f7", + "ms-ea1abf4dd25aea30202b", + "ms-653181f057aeaeeadfe2", + "ms-cbdc087c47e35d3d93d9", + "ms-973f50778408dcecccdd", + "ms-7cb73c21b0db06dff162", + "ms-93cbf45c02a304a2f6dc", + "ms-7a75d1290f292d92ca67", + "ms-3a5eaa0095953c9d6241", + "ms-eb0a75800f4ab451b92e", + "ms-a631a8d6314b7547f71f", + "ms-5b5ed3d64cc0c00bac17", + "ms-3e17d03f9d640185853b", + "ms-170a9228df2273649f1c", + "ms-9b50c72d20e7510258b8", + "ms-1dd7b6c07218e5e69494", + "ms-c333dcd8ec59f2ae8a3d", + "ms-dae29d741eb65a2fc4b3", + "ms-fc8084d5cd7da937fc02", + "ms-3f6cd47af5e9e6781b73", + "ms-7dd90f1575d9612c456c", + "ms-217aa98e4a691e43cf68", + "ms-f492405ef825467ddcba", + "ms-876805da5863dc637246", + "ms-edeb8cb36d3d06244421", + "ms-eccd5c34e1762a33f643", + "ms-411a6db8809ff09857ff", + "ms-33dcfa4725ad078cc898", + "ms-b8bfdd3e6c31bd98082d", + "ms-c8878bfab66b7f3f0839", + "ms-0be67b60c98e4f3e692f", + "ms-c5df8e13dba07b355c14", + "ms-ff08c0b84d984d3d8620", + "ms-5c7ab237037e9344a2ed", + "ms-40c0e775426e28c3afd9", + "ms-86909e2391a6a77883b3", + "ms-265a3e206e8d55a9171a", + "ms-864a932d1281c7e316af", + "ms-ce6cbc6d7664766f1285", + "ms-a18cc16207970cd8f84d", + "ms-44a69da62321224c1df9", + "ms-bb8ddc9a1e197eb7cee1", + "ms-670a696bd472769ac513", + "ms-5df972e5e36afb67a009", + "ms-2a4f1cbd3f416f9084dc", + "ms-f39cb13a5942ca3317d0", + "ms-79636173ebff4c91da03", + "ms-de5cb4c5e0ba01be4568", + "ms-1405ab4a382fd1a07e63" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000003.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-5383e9b606cc63b724ab", + "ms-31880fb2dfed92c021ef", + "ms-167bd325bc0845dcaba3", + "ms-32767ca96808a18b5a43", + "ms-e4a99e5ef340bcc181c9", + "ms-7a102c1d0334bab9656e", + "ms-316e11b816e8a3495dd9", + "ms-a649d9cf4fba45822ad9", + "ms-84bc79aab8423997e8c4", + "ms-9fad82fd288e4296e354", + "ms-44e2ba085128234ef9eb", + "ms-1f195ac1f3ae4b397c71", + "ms-3579b51b91daa627d195", + "ms-c01a70d845742d97e197", + "ms-aed1c09dad8914387967", + "ms-5a0da896925f50a1d2ee", + "ms-bf98b46e10ab07a2e60f", + "ms-4096e2d9cf2e9d4943c7", + "ms-ec5819de002710aa5d97", + "ms-5d646fb3dbdcd5443729", + "ms-f83aacaa71f0151421c7", + "ms-949bca60825f84d7b5b9", + "ms-c8ffa7d45ed61e2b94d6", + "ms-62033042dcde101a0c9e", + "ms-0b13513f98f2d3def73a", + "ms-147e4106b5d337cefcb2", + "ms-b8d9d47612bbd603db4a", + "ms-82d6fc374e1f1808654b", + "ms-4e67cfd30a257b4bad0c", + "ms-e63468cc3772da77cc67", + "ms-e2a1bc3e5cd509eb08c2", + "ms-cdca8a6fff0640ab7f16", + "ms-db41eac14b49855e7e58", + "ms-4015b8ac7c0a1397c076", + "ms-663c7e7fdaea6a945a2b", + "ms-bc4bf52af94d1469eaee", + "ms-234ce02cf98f85bda72c", + "ms-109d9131cce2e94f40c1", + "ms-9b265e13a79d2147718c", + "ms-c57abcef044f324c2f1d", + "ms-16619fc9e4f3627b8dd1", + "ms-121bc2e6676c1ac10911", + "ms-fb23602361466c57dd2d", + "ms-346494921f998d7a48ab", + "ms-d7cbc8db5f1b69fcd9f9", + "ms-e4656ad3c67bd430c1a6", + "ms-c58ff92a6fe47fd90684", + "ms-ee72e3f1115cfb8294ee", + "ms-b2c71dc14945b7bd8a60", + "ms-d99360c1e4bb4bed6d50", + "ms-ea694828fcd81abbbc7a", + "ms-edb1d483ba0ada6fa703", + "ms-966895e247be1f41daee", + "ms-2962fafa0fcd14e516bb", + "ms-940cce8479900a244d54", + "ms-aec40fad850827331244", + "ms-dec13f2e40baaf5c8d40", + "ms-934e125b62165b614036", + "ms-60c02d02de5f1e91124f", + "ms-ffdca676bf6a4ccaf640", + "ms-42f077d286655b676709", + "ms-63b9ac9cfd0a96150eb1", + "ms-82a5b63ed41420747fb9", + "ms-64b6e9f72ae5e2999f56" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000004.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-98d8d5d728ee4cff6776", + "ms-79d4959d079594053125", + "ms-ba0d059f1c55c5737e17", + "ms-9f61e09c05b4c8f74366", + "ms-f6174dd8dd65bd736685", + "ms-abb9a0e2ba0163714374", + "ms-4711dffcd73048a3656c", + "ms-4a4de470c6c7cb5a3c34", + "ms-4aa2c82355bdbbb93f9d", + "ms-f14cda80db7270dd0f86", + "ms-ef0d8b87e9a145873ca7", + "ms-9f33c55b8e06fcb5e21e", + "ms-2a7f05987ae1fd0fa33d", + "ms-d4c37ed44fb338c1d551", + "ms-cb51cbc4086bd452233d", + "ms-a012d1677923703d5a0c", + "ms-b6ed7cdf9226aa0b0d1e", + "ms-40bddbc4f8ee1ed19575", + "ms-5dab866561f657e9aa8d", + "ms-b4293fc495a608f8063c", + "ms-7f0fb9edfc18e6b83b41", + "ms-8293ae7689b30820278f", + "ms-bb6a2f2ae3014b135b40", + "ms-c4a78f303c7f2b1a437b", + "ms-79c57793f8c7523b59fc", + "ms-7a699e157f311f764b46", + "ms-608a539637578968beb3", + "ms-b5dee7bdaf34270e6515", + "ms-d0005f84908a87243dd6", + "ms-23a02f5f0c3f979684ff", + "ms-c1607c435e4b43609693", + "ms-8a2ada13f7fd1a5aeb2f", + "ms-b25b49e16e98af4a554c", + "ms-d5d38df1f2dbae367af0", + "ms-d97538effd225a910367", + "ms-90a840c7cbbfe35e8a67", + "ms-a39c5eb6453fab13bb3f", + "ms-579eb7a700037d011e3f", + "ms-783ceec205863d92ef8a", + "ms-3c9527725515347ba7de", + "ms-cfed5c82c6f4d1c1c3f0", + "ms-933425ede670e8efc04e", + "ms-27eb5fd79735a667570a", + "ms-ea67d3a4697351095f09", + "ms-a2766876467c8fbd5379", + "ms-f12b50fed09bbbf6c039", + "ms-111355c28fcb5ef51167", + "ms-25119a8a2d03c8e71300", + "ms-1044e227d00afb01a5dd", + "ms-5816c255ea56bc8a63c3", + "ms-4c8a03e5d6e3b7a13c29", + "ms-07aa7d1f9c01e93d04bc", + "ms-22cfd6a9a619133478b6", + "ms-680b80c5d2872be1e66a", + "ms-49b9fd364e000bc7d579", + "ms-b559638a0a726e441847", + "ms-54346a8f79bd9dca1dfe", + "ms-f2f594b2492898e947f0", + "ms-634dce5b0ce43bfaa04e", + "ms-889dd046ec7b0fc0a7a4", + "ms-4c74a17f4d3e6237c0e5", + "ms-6d48ff55845b6759a8c4", + "ms-e9bb31726937713db23a", + "ms-c8593eabfdb5a6bca7bb" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000005.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-771e632c1b0f3181ff36", + "ms-6f6c2f554a8e1c33ab32", + "ms-beb40d1c9533f7508aa5", + "ms-702a28f36df5dbbbe010", + "ms-b929d6015e41aa9a9de8", + "ms-2a738f03ee216527dfa6", + "ms-06c528eb38baee4e31b8", + "ms-81eff749adbf328a93e1", + "ms-5f750a2e67b4e9597b2c", + "ms-fb7c4ead34050a60dbdd", + "ms-29f7da1cd496619f45f5", + "ms-c16b5cbd61b51f0b66e7", + "ms-9d190663aa55bb434720", + "ms-0e132f1526b2090e440e", + "ms-383c0d1f90c5a0b55e07", + "ms-fa27eb61283e52488bdb", + "ms-c9568dec28405d273943", + "ms-68054c6c805197b033a7", + "ms-183174ef761d19fb85d9", + "ms-bb9f2f5b67bbb7e0f516", + "ms-9b6189cefdca2110b96a", + "ms-ec5ac79449116b0a30c1", + "ms-cf8cc7e3c4481783d5f5", + "ms-29788f80c9da2d2d8b4e", + "ms-c257731ddcf5cf149393", + "ms-f813ff3f9ae188b0259c", + "ms-027592312fc833e8768d", + "ms-97939101c75832eb62cd", + "ms-a5f722aaecb894444ac4", + "ms-338d0d83a8a30b40ee75", + "ms-8d23082fdeffba7cba51", + "ms-1003d1fe41aba4c8af84", + "ms-130188b25b236b6c13fc", + "ms-34ea26afcb2c6fe1a1fe", + "ms-d867fd6ed17aeaa97c2b", + "ms-8dd1387678e0e62a1298", + "ms-c1513a6417b0ba250dee", + "ms-14265f08eea0055edff3", + "ms-b187043687d88d509fc7", + "ms-238b0d92a2c0a61289b7", + "ms-25942684dfc6d81c449a", + "ms-49731d5e69272f3bd959", + "ms-2e368b0323bc5ae1cfed", + "ms-20bc116cad15b072a2ea", + "ms-eec1c80f2bdd990118d8", + "ms-2e237a4b84b582828c37", + "ms-7aaa46eaa657a3179a70", + "ms-7fc804389bb1de250c84", + "ms-b8f5ebd4732c0274945b", + "ms-ea6878cc654f7f255811", + "ms-d7c83900ed8544e80a41", + "ms-5d39d1a6fdb17aeb8116", + "ms-1371e0e5fe7f126bbea2", + "ms-9448b2a52a4a475af1f9", + "ms-298fc552bfc93c2b32bf", + "ms-4f291cdf6a62b9af9e5d", + "ms-2377ed2b717d22203e99", + "ms-e7fcbe2b61632a79fe6a", + "ms-2a4611f1099507f9babf", + "ms-7f22550a9f6707a45cd4", + "ms-f09be6d8353d3d53a44e", + "ms-10a1077a32c1abd486f7", + "ms-68682a8a5d31b10fcadb", + "ms-f8a4e83bdb9d628ac483" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000006.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-c8902580575f6effb241", + "ms-22f6ac0312682d5358b5", + "ms-07321628f307d0c91b44", + "ms-d130baa9773a108ae1b1", + "ms-0b28e6af81d42a73cbeb", + "ms-05f34a90fa7355262b8e", + "ms-f044b20b0f2bb0983f5a", + "ms-de96f1dcca4527c03814", + "ms-6ed65c65b566e21474a8", + "ms-f40f6fe1910a94971980", + "ms-57fd761853dfc47e8059", + "ms-64f98d163b6bff29332b", + "ms-58950fed04d203b97ab9", + "ms-8278abbc47a6efcb5b68", + "ms-fa637544793129258022", + "ms-8ef0a258f0f502bdaad5", + "ms-3cb727b6501f5c1a5a16", + "ms-4410b778c34b05366011", + "ms-75dce23f7afca1f3e5b5", + "ms-87f2f63d694d8ccca7f8", + "ms-4dbdc6ed61dccebbfe92", + "ms-e8188aaeda131829b7c9", + "ms-74704688fb63757619e1", + "ms-5884cb0ce8ef0e034fec", + "ms-2ca0290759891500c6bd", + "ms-f6635cbacd4195f7f4ab", + "ms-79d2ebbe5913e3a72313", + "ms-8a8484751c0ce80cd2f6", + "ms-b5ca4a8c438aadab9ec5", + "ms-630c7d4c0ef39da64304", + "ms-7c72cdc319220ea76854", + "ms-64142b18744de985eb63", + "ms-1188a1b8e0ac31187b3b", + "ms-6a82cd518812ae55f97d", + "ms-1fe4b5c84ee5c0e6fcdb", + "ms-1e3b1d1f26baebcdb86a", + "ms-80269804728c7aaaf9a7", + "ms-26d8d4eecdfba2eb766d", + "ms-0350672b45f92c437c22", + "ms-2c4224657fa8115b2b8b", + "ms-3589c30bbad90d19f2c4", + "ms-f5406744f19896e4d35f", + "ms-f52fe666ef7c9b997778", + "ms-b0f5b97b1bc7d3173183", + "ms-fae97eb60014b2866413", + "ms-543339410c02ce2269d1", + "ms-f848ee8748cea24fc8a0", + "ms-031c526a8690cad1677c", + "ms-ab5681ebeef5851f5ceb", + "ms-c10cf873dd3f68d97755", + "ms-33e036c8224f052e6279", + "ms-f6755a487459505d598f", + "ms-081952a63087fc8f83cd", + "ms-31372cdadd8d9da23b8d", + "ms-3c326d05fec66e1d6bea", + "ms-7f5fe657331466b2e262", + "ms-9ac0a7f45996f7326682", + "ms-f7db0d0ad0967af1d0c5", + "ms-2288d0eb01482bc757e5", + "ms-1c4df81cce643ca4fb3e", + "ms-cf1db8d3ccb39453a3fd", + "ms-ae5eaf676ba740f3dc07", + "ms-4866796139a8dac6d73b", + "ms-f97fbc8f36f0af9ce18c" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000007.jsonl", + "record_count": 832, + "group_ids": [ + "ms-5ae7b4d324285dfe6e80", + "ms-fe38b164397dd669f391", + "ms-c1f8e861c7c4485c361a", + "ms-60a0de44d30044ef60ff", + "ms-7838ea128d16b455b2f9", + "ms-847200a83b8ce5835fd9", + "ms-e8a0577b3fc6f75776b3", + "ms-70e24231e9bb560abb3e", + "ms-9df99aca252d63f555b2", + "ms-28a10bea0cf3f917e067", + "ms-98d738638749a1c0999d", + "ms-c60bfb966aa2f6e5f183", + "ms-e6b0ee06238145e2a517", + "ms-4c9ed894520169849bc6", + "ms-01023f12e250383909f1", + "ms-29fbdaedca6b1e7d233d", + "ms-acfecece1f8160045f0a", + "ms-918efaf9527826841090", + "ms-6ffe441e29681c8f116e", + "ms-4d8f2bf18e7556663104", + "ms-db2cfd3c0ce8ece8a112", + "ms-d65e80424e917743d9a4", + "ms-10450837f7f5319c4cab", + "ms-45b8b3cec97878df0d19", + "ms-f1b05f75f6ad032a0cd1", + "ms-2baa5b5b99121cff72c2", + "ms-8750cccf83614187fe14", + "ms-3ee8f778784665ea6ee0", + "ms-97ccdc4668a995776663", + "ms-eea4df1b36e8a31b40a3", + "ms-ba0ee0c9c181cdc7de9d", + "ms-9d82d1927529dc675bc5", + "ms-7a50dd49897a83a22282", + "ms-fbfbee34b928dbb775da", + "ms-998ac1af1755510c9229", + "ms-8d4ee30bf4815d04c391", + "ms-575e151c0d1622c7ccd1", + "ms-733b04f5bd4720d79e99", + "ms-8b813b6035f0425892bd", + "ms-1f9a8bf8a8c0882b4039", + "ms-bfe867e8d518b9f8333a", + "ms-a40d644441c878ac5292", + "ms-7fac99962876b33e8d0b", + "ms-09bceb938cc530125601", + "ms-2b27b1ecc0284ec3a993", + "ms-dc6ca032f5df915c7945", + "ms-c36379627de352fd2325", + "ms-551020d425c370126f59", + "ms-7530a944fb0979ba91e3", + "ms-ccccd221a821803e36c5", + "ms-9d1d3a7b2dcafeda3bf3", + "ms-33885f2999b6dbe71f42" + ], + "format": "jsonl" + } + ], + "record_count": 8000, + "group_count": 500 + } +} + +✅ PushCube-v1 generation complete diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_2.err b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_2.err new file mode 100644 index 0000000000000000000000000000000000000000..896fdc03331083d3899bf7fb5ab5bc3ff8d48b34 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_2.err @@ -0,0 +1,11 @@ +2026-06-25 19:20:46,687 - mani_skill - WARNING - No initial pose set for actor builder of goal_region, setting to default pose q=[1,0,0,0], p=[0,0,0]. Not setting reasonable initial poses may slow down simulation, see https://github.com/haosulab/ManiSkill/issues/421. +Traceback (most recent call last): + File "/lustre09/project/6037638/knguy52/vla/scripts/generate_maniskill_lattice.py", line 103, in + raise SystemExit(main()) + ^^^^^^ + File "/lustre09/project/6037638/knguy52/vla/scripts/generate_maniskill_lattice.py", line 76, in main + summary = generate_maniskill_lattice( + ^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/lustre09/project/6037638/knguy52/vla/dovla_cil/generation/maniskill_lattice.py", line 280, in generate_maniskill_lattice + raise ValueError( +ValueError: requested branch plan [0:500] but only 373 unique episode-step states are available diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_2.out b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_2.out new file mode 100644 index 0000000000000000000000000000000000000000..c016664247940dd5fd9b4f4d84f4724569a7d9ef --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_2.out @@ -0,0 +1,7 @@ +================================================== +Task: PullCube-v1 +Groups: 500 +Horizon: 16 (vs baseline 4) +Demo: /scratch/knguy52/dovla/maniskill_multitask_demos/PullCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5 +Output: /scratch/knguy52/dovla/experiments/six_task_h16_collection/PullCube-v1 +================================================== diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_3.err b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_3.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_3.out b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_3.out new file mode 100644 index 0000000000000000000000000000000000000000..3d39677494d8b8bca204d377010a94c16b7312c4 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_3.out @@ -0,0 +1,649 @@ +================================================== +Task: StackCube-v1 +Groups: 500 +Horizon: 16 (vs baseline 4) +Demo: /scratch/knguy52/dovla/maniskill_multitask_demos/StackCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5 +Output: /scratch/knguy52/dovla/experiments/six_task_h16_collection/StackCube-v1 +================================================== +ManiSkill CIL progress: batch=1/500 groups=1/500 records=16 restore_max=1.192e-07 +ManiSkill CIL progress: batch=25/500 groups=25/500 records=400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=50/500 groups=50/500 records=800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=75/500 groups=75/500 records=1200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=100/500 groups=100/500 records=1600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=125/500 groups=125/500 records=2000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=150/500 groups=150/500 records=2400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=175/500 groups=175/500 records=2800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=200/500 groups=200/500 records=3200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=225/500 groups=225/500 records=3600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=250/500 groups=250/500 records=4000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=275/500 groups=275/500 records=4400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=300/500 groups=300/500 records=4800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=325/500 groups=325/500 records=5200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=350/500 groups=350/500 records=5600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=375/500 groups=375/500 records=6000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=400/500 groups=400/500 records=6400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=425/500 groups=425/500 records=6800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=450/500 groups=450/500 records=7200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=475/500 groups=475/500 records=7600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=500/500 groups=500/500 records=8000 restore_max=4.768e-07 +{ + "backend": "maniskill", + "env_id": "StackCube-v1", + "control_mode": "pd_ee_delta_pose", + "sim_backend": "physx_cuda:0", + "source_obs_mode": "state", + "instruction": "Stack cube A on top of cube B.", + "target_actors": [ + "cubeA" + ], + "reference_actors": [ + "cubeB" + ], + "demo_path": "/scratch/knguy52/dovla/maniskill_multitask_demos/StackCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5", + "num_groups": 500, + "group_offset": 0, + "available_unique_states": 14134, + "excluded_post_success_states": 10270, + "branch_selection": "pre_success_only", + "num_records": 8000, + "k": 16, + "horizon": 16, + "seed": 0, + "measurement": "exact_state_physics_branch", + "candidate_mode": "structured", + "branch_execution": "parallel", + "state_batch_size": 1, + "state_restore_max_error": 4.76837158203125e-07, + "state_storage": "archive", + "state_archive": "/scratch/knguy52/dovla/experiments/six_task_h16_collection/StackCube-v1/state_archive.pkl", + "observation_archive": null, + "image_quality": null, + "manifest": "/scratch/knguy52/dovla/experiments/six_task_h16_collection/StackCube-v1/manifest.json", + "quality": { + "expert_success_rate": 0.748, + "no_op_success_rate": 0.002, + "mean_reward_spread": 0.9401331427395344, + "nondegenerate_group_fraction": 1.0 + }, + "sharding": { + "dataset_name": "cil_stackcube-v1", + "version": "0.1", + "created_at": "2026-06-25T23:29:57.474699+00:00", + "backend": "maniskill", + "num_groups": 500, + "num_records": 8000, + "k": 16, + "task_count": 1, + "seed": 0, + "schema_version": "0.1", + "format": "dovla_cil", + "shard_format": "jsonl", + "index_format": "jsonl", + "shard_size": 1024, + "shard_count": 8, + "group_index_path": "group_index.jsonl", + "record_index_path": "record_index.jsonl", + "shards": [ + { + "path": "shards/shard_000000.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-728ce6cbfd4c82607908", + "ms-6989f2ad4ed53dc07e1f", + "ms-14064e79c1a1abcf10e3", + "ms-8c2dade5b58489e26a4b", + "ms-cbe2659ec7884f00dc8a", + "ms-22457d523f50e99d7acb", + "ms-e3bbabc7ac3950eb91a2", + "ms-95950ddf57046910494b", + "ms-9f80c93ae6940a427c83", + "ms-7dca31b0d7d6fab018f3", + "ms-8fb45ef2b489499fc7a7", + "ms-37d6b220ad58fb776f7b", + "ms-8ccec039b839e612969b", + "ms-aa1da4bf5f94c6a74839", + "ms-94b9114b294d56e3f88a", + "ms-bc8e9fc055c591071ea3", + "ms-744cdc2b2570c3c26326", + "ms-f9148dff7961fa66a8a1", + "ms-9e99b7303fbd7ecbb17d", + "ms-3a5cf404fd30dd8a2f76", + "ms-72777024349c2aacd5e0", + "ms-c9edb3f445256b0d5341", + "ms-5a0e07161aefc71d4f33", + "ms-1cba14e90a2d067bc198", + "ms-d1960368ef76bbb10669", + "ms-f020e53a077d4e6bf383", + "ms-5b26c33feed6992386b1", + "ms-da5f2c1dbafa907d725b", + "ms-3ec29a5b20bd65943255", + "ms-f3ec6a4e6e3a778054d1", + "ms-3ff1aedb72f65a182b58", + "ms-035231abc135f925e1cc", + "ms-df6a1fb4b1ccbb410b4c", + "ms-0982ec8847911627c519", + "ms-bf36bafb9ac562480022", + "ms-1691ce7ec497fad5ad49", + "ms-578073df4f8fe9b936d2", + "ms-df647aaea94ba40f7a8a", + "ms-e9cb92d664266091b91d", + "ms-c72d626356a30cbef93b", + "ms-0fa473b2c94d12bc3945", + "ms-5a32c809ed679663dcb1", + "ms-0949097d19e597fc7f63", + "ms-48ddae21c360b2a12b8d", + "ms-b3a1c65d7357ca4300b0", + "ms-01094af881b751ef29aa", + "ms-b651d43182be35524aad", + "ms-62601bc1741cf8792326", + "ms-552acc881ce2bbff59f8", + "ms-f774646335321f6c5980", + "ms-571d878853654c42a521", + "ms-5909e93ba4f8ed21d193", + "ms-ed6314fb8f69f956daf6", + "ms-e95a846b4e77792e3e6f", + "ms-54b553c42a8d825beb38", + "ms-ba1dca497a16134eeed8", + "ms-ae1b8592cb57dbf2fe98", + "ms-dd1eec30969c12ec89af", + "ms-2c6ff8d807763f665bac", + "ms-f0fdcc26c648ce089746", + "ms-63a1574267ccbdaf6cc2", + "ms-5c8151ea6809f3c331e5", + "ms-584d52e4f9515f695251", + "ms-7bdf9cab1f7071533bbc" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000001.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-44b52b73161bb03c5a27", + "ms-63ac245579b1188935ad", + "ms-2b6ab8760ea5d558797c", + "ms-626991ad29d003c808f8", + "ms-f0696e3657ed996dec15", + "ms-7f58527819036c01daa4", + "ms-90777a893f2fff27728c", + "ms-d648b2c7a7b42db45550", + "ms-97343fa03fd05700c270", + "ms-76774eb92b50132c54b6", + "ms-afe554d7e85b42481646", + "ms-dc4738f9c7a8af86722d", + "ms-05ab68d4cd79b74a45ee", + "ms-9747840abee181e3e4cf", + "ms-d520d237d6848cde5883", + "ms-13243cd92a78eba8c08b", + "ms-60f5fff0649ceb50ab60", + "ms-edf898278a0dd01a5b37", + "ms-ac8685beeff23e48af61", + "ms-8903eac679218fc038ea", + "ms-49e1584a68d8f465e8e3", + "ms-e887dbfa0133af46829d", + "ms-7d5da9bf431cb1294c82", + "ms-767d9cbb033ed5809533", + "ms-97e2c618a3045e0a759e", + "ms-f4579acd00ff9f578098", + "ms-f413b168078609ed4fe4", + "ms-1c4ef17cdeed30bb2548", + "ms-4af6806104210ccf74b0", + "ms-07b479748a91a9c92b20", + "ms-7ab0027224139e5332d5", + "ms-bc02d87b2ed5769ee484", + "ms-0841a07e9aeb84cff70f", + "ms-c52ffa349a79397da8c3", + "ms-01e08f0dcb4e2c030a79", + "ms-2360ac84ff893ee51c5f", + "ms-d6ed8683183759789495", + "ms-33ac0623919ea05d354a", + "ms-0b6ba12f99b3554c0e88", + "ms-4bbd30ae37fda61aa3c2", + "ms-1f7819cf58a2f9abd1e4", + "ms-e7959111ef1429bcf699", + "ms-1cc18d88615939309d51", + "ms-0cdb564903c69734321f", + "ms-4c5469e5b4af75482d89", + "ms-15fc8cc6f74cf589056f", + "ms-1cb54e3347c49dcdf28e", + "ms-797db95b24deab760f10", + "ms-0ea6d57781e366b578c1", + "ms-15248c64e70b0dafcb33", + "ms-b30834b2c4e6b153624f", + "ms-495b33a30a68ee055eef", + "ms-00e356a462bdc3380f4f", + "ms-2503df67073b4c59e640", + "ms-87efe56140dec7d3ffb8", + "ms-741c227f6249d1c8c9f1", + "ms-185c13ec401a35f393b7", + "ms-ffaf8ddbbe21af6dea95", + "ms-acb486cbc23907df2568", + "ms-020ab0b7cfa5045f3fc7", + "ms-01e6eaaaf40875e0ef5a", + "ms-4cab45c9b21452c1b523", + "ms-0f2cd78c21a012784e1b", + "ms-852c671b0adeccfa4bf1" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000002.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-39cea5dde6dc1718d478", + "ms-283f98eecc655a33d459", + "ms-d6af11918833ace7c694", + "ms-8c83fb7e9f1a7a8ad58f", + "ms-bc1ca12b92e3561b95b3", + "ms-fe22e55457c78877c2d9", + "ms-a3df9ead5a29798665a5", + "ms-3278c343e1bdece2f252", + "ms-35fede79899d623c01e1", + "ms-119db4f7ab55afe5b9f6", + "ms-dd8ebc7d7211ece74e03", + "ms-6ab6c8a66fc509e2113e", + "ms-c5945ca04bebc81e72a8", + "ms-e6e8b4cdcadec645d910", + "ms-fd2387aeed9bfc9a0d08", + "ms-4c6a9b0840052fd62261", + "ms-c9bac4a12aa021442f03", + "ms-3afd8959363e917591d4", + "ms-76c9576538106cb9d5db", + "ms-2c7a75422c4e47ee680c", + "ms-711cecdbf06d02f7ca2e", + "ms-77536801924a2a78688e", + "ms-9e96728c711ffe98fe69", + "ms-5ac1bf2d7d319955d45d", + "ms-ff0d7559bf7f39d82ded", + "ms-141aab4472db1ecc0ea2", + "ms-d776afa14e6322cadedd", + "ms-e297545a438d902a680e", + "ms-eff3caf3c103c5232495", + "ms-b6ce87522fcc35b563c7", + "ms-a50b6ab3c59b1c953197", + "ms-5a456ddf97d294548cf1", + "ms-42b26451056c6dd2150e", + "ms-99bc7344d022af875494", + "ms-3451f857ec88eeb49c38", + "ms-f72b64534388e443b220", + "ms-8d02415bb9af319f215a", + "ms-c4a4c6a275c817d8b03f", + "ms-f0033d0b5f4f2e160859", + "ms-8f84a0371a7ac44c1b1a", + "ms-3cdc09425d5b71524854", + "ms-265435476b51eb83fb9b", + "ms-146ff6c3f22b3a49af3a", + "ms-8ec5329e8a36beda0fda", + "ms-a134953a8d3e6be63757", + "ms-e630078b186f7d8f8e57", + "ms-ab5acce4dba71f19d5ad", + "ms-dc26a7037003b94b9f93", + "ms-ac0b941303c34f55d838", + "ms-0e7f710fb5e661e71fd3", + "ms-0b226017db7f07c5307e", + "ms-9e7e3b675db86ea3e212", + "ms-9228f8295b9891cf011c", + "ms-748baf7aee43c63797e1", + "ms-720c704eeac615b41763", + "ms-ecc526493f8048cf4c56", + "ms-0d8e00d0e73f9a7b4ce8", + "ms-57ccc6bbce43f1c7c000", + "ms-baba39ca95aa09893dc0", + "ms-a5eb6b813b57d14f6573", + "ms-016c5de39ed3c111889a", + "ms-a4d692315fd15700e75b", + "ms-02aa62c077c9d4f86e61", + "ms-fc9523f9953b1b5bfbcd" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000003.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-fffe17005e577cb2ac1c", + "ms-fe763444e70a58ff8b1e", + "ms-b1ea4f7f813ac0c87079", + "ms-8e69308e5c86a307b08a", + "ms-8c67f938fb3d0264393c", + "ms-f2259eac5ed12ee61c19", + "ms-df4aee1613f5813275a1", + "ms-dfcbf1a7d2e3b33eb3e9", + "ms-f90e4d887f4c840a0b25", + "ms-64b4e29bb42bf0b35d4d", + "ms-e2a5f29a9c0d835f3cd8", + "ms-350716d0809bd7ef12df", + "ms-22c36189b47df8987e8e", + "ms-500ed73f6d891e237b30", + "ms-d5ddabb086026ee3128e", + "ms-f2509e6c3c468c4e30f0", + "ms-9f5d57fe2153fd267774", + "ms-54f4d544f76358c09ee9", + "ms-462c1284ee1a148a7e1d", + "ms-cba80792e163a5084a92", + "ms-55e78c04643bd06fd5f8", + "ms-56a0e8e80e2e1e5645d7", + "ms-75ddae05fd4d968ebac7", + "ms-1d7505b80dff9c229a99", + "ms-702c6c5000ad4d6424ea", + "ms-f2a349cfaa6f7422237b", + "ms-0ec57b5bc28a71d3aac0", + "ms-893e0a0d03b5e47b65dd", + "ms-cc21e4000cabb2fad805", + "ms-b86c44375ba5328a1608", + "ms-f79778f4c5ec302bbfb4", + "ms-36bc1364d2439099e9ed", + "ms-ff4c9c58b742da7135da", + "ms-aa7cc7e3e5eb86e58ee9", + "ms-b8492a3a812095651d75", + "ms-0180bfac51ce81b2b522", + "ms-0a4d4909900fb0bc4a9a", + "ms-0c2f2dc98424b6704c79", + "ms-87d132d28bc30b07002b", + "ms-ce827a7939377a151af9", + "ms-bdc0e2b6d589c1d90d69", + "ms-875a4668671c02054a3c", + "ms-579af9eb582ab20522a9", + "ms-44940e96c327541acf19", + "ms-79b745efd36e00515fa7", + "ms-151e5f61c8e3c2ece4e7", + "ms-e8bbf7fcd90b4d11db2a", + "ms-efdc07c985935d871adf", + "ms-180d0e0a6315a06a7021", + "ms-bc96a8973358d1c20ed3", + "ms-53fd90d05589287fb338", + "ms-04eec0abab5f55620e4e", + "ms-d50a82b8770f333f1947", + "ms-c8706d0aa14a930236de", + "ms-7d484dd14e9d001a07b1", + "ms-479ebfcc567538b3b582", + "ms-0ee3134ab4e9435f2c2c", + "ms-2cc60867f3629a801fc4", + "ms-bb299baa34e49eafb9fc", + "ms-9ad3b276d9ccaf668d30", + "ms-fc68c8b1bef631b6cae9", + "ms-08077ae7bcd1d20bb2a6", + "ms-58e6bcca0d74ea38d3dc", + "ms-253682457d6db5d265bb" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000004.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-f5bc4935790fb8ef7b81", + "ms-65f4ea8ed88ed5301430", + "ms-1a37fdd82d982ae6dbd8", + "ms-9d4cfe1888484d9da1ee", + "ms-30a714841f043c80a292", + "ms-52a2ef2f862cee509ce1", + "ms-61852da845126d62ff5b", + "ms-6c7b2b62ba170ac8133e", + "ms-6bd5870eb5ff0dd76919", + "ms-cdd9ca03b64c4145587b", + "ms-19dbfb3664d48598041e", + "ms-b04a4cfe76d719518dae", + "ms-546f8817db5b76b563f0", + "ms-5c1ed0b96c4b99d64dd9", + "ms-9c62b3031cdb58aa7e17", + "ms-1e17eb687986f7c5bc22", + "ms-ec0b095d57014b86f5e7", + "ms-ebf15fd96b4a62e0b393", + "ms-b1ab88b6b9dbb2dc602e", + "ms-787b74c3c67428d63b72", + "ms-b90437c786f2eba2846b", + "ms-e0e2a425cce52e820a7e", + "ms-7fc5696cb3dc7b9bb9f7", + "ms-caa5773983c25018c3f8", + "ms-0ac85b1b15125006eb50", + "ms-71c5493fbcd8ea88898b", + "ms-a77b8d61c09ed6ea23a2", + "ms-bb742574fc7cfaea26d0", + "ms-bb70aa617f82f530d42e", + "ms-1b50d8927c080cd97a79", + "ms-c164dd3b3b89c546a184", + "ms-c5d2ddd0416072b20d5d", + "ms-64ff9614d692fd9dc48b", + "ms-be0ca7fa48a7ea391aa8", + "ms-a13e71cca2a60d77aed4", + "ms-ee2fcbc94d41385713d0", + "ms-8b53e5dbab786b2e3b75", + "ms-511bcece89a5ea3b849e", + "ms-a054eb19f228304d37aa", + "ms-4b6cf86fc2450321163a", + "ms-1c9cd9629e4a927ed921", + "ms-bf2b9cc05ccb6bce0d30", + "ms-16850056120649fcc370", + "ms-93af5e555e79e4177a5c", + "ms-c9804ff90a86259f7149", + "ms-4dc0defd156f4c30aaf8", + "ms-10405f0fac390b29a04b", + "ms-5e19aac5dcd26c992bfd", + "ms-7fba35a6c125e88b8f54", + "ms-3c87c16acc91c68c522f", + "ms-9dd1ff4c297a81b13071", + "ms-9ae5c6c4132218ecf49e", + "ms-abf93faee35669f4c359", + "ms-626ac2ad9d0dd95f1a2e", + "ms-420280f15a522ce21c72", + "ms-b2b93411a16dbb5d4202", + "ms-578fc3b8cbdbe6b0542f", + "ms-207e9a9f9123d6c3eeb6", + "ms-0ce4aecc96a50498e2f7", + "ms-20ab0b7fffdbb23e6d53", + "ms-bc2d3b425d9163a62dbb", + "ms-5922d6b0cb028950a6a3", + "ms-dab618f516ccc10f176e", + "ms-df346e952d6eaab76d40" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000005.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-1ed36d926d5d640b4a04", + "ms-ceb81989361d018a9e67", + "ms-97cdb4c400ad4ae36518", + "ms-a9eb9fab62fb296d2cf3", + "ms-0ffc732b968c030ba1d0", + "ms-82b267d663e8400d19ca", + "ms-584219be39de32d776d4", + "ms-6b2df24ea06ce2cdfb5e", + "ms-804705b564058a7ac34c", + "ms-41847ae4ccfbee4666d8", + "ms-80a997bc4867e103fe34", + "ms-9a14552b7db695dee92c", + "ms-0c490961b934978e5b1c", + "ms-829c9968f0750cd9aed6", + "ms-e1bce061511b584c5381", + "ms-652234a929b7ed4d3b23", + "ms-5bddb39f792f20e5e6d4", + "ms-a2a546ebe0d1283df0f0", + "ms-32a86ecc263b3d79db00", + "ms-03df3e12e1d61aa9ebbf", + "ms-269815a0f9c1a1ba5a6b", + "ms-7bc1b098e6c696f6d4de", + "ms-12ea343fddf9f6e6b9e8", + "ms-36112cc5608a9f625764", + "ms-673f96236583a1925836", + "ms-921a0cd2d5ef7e70229a", + "ms-98c0d73ae0085f97bab9", + "ms-0460b73af19ce1336bba", + "ms-4c788310cab7df5251e1", + "ms-0fb946120bcbd8ed0632", + "ms-a4e3e9c634e92a46fc15", + "ms-983f35821c4f83c6e4bb", + "ms-f6bfe2d3d8f8956091a8", + "ms-30d0edf166c7ab5f689f", + "ms-48b2c358243f84013a0c", + "ms-4dd3c4d0161d01abb6d4", + "ms-61cb49c57a8600267d61", + "ms-7ec9b9594e36ab34467f", + "ms-716586cf73f853a45bfd", + "ms-274ef8c6420a669dd5b4", + "ms-4cab292bc65fa59274ef", + "ms-10bdda377f6b2353d200", + "ms-02cdd56290c7ab888393", + "ms-0fce9fe598d9ed1471b1", + "ms-63979e5d869bff8d3515", + "ms-7ee2d50e0788ccf0cdec", + "ms-fd81a55becc2666437aa", + "ms-242302c2a1e59588b319", + "ms-454309c1b142d35146cf", + "ms-88070bf04696eceac1e9", + "ms-88d7cbdbd7c9636e456f", + "ms-c80e71e55d51d8595499", + "ms-4048f8f6e047422e3c79", + "ms-0b1dd0e538b8df15a6e4", + "ms-f6cfaf92977f7b72aa24", + "ms-5cfbc3fb2094ad5bba6f", + "ms-dd9b840c99c305ac1977", + "ms-666f00775d4c54be1282", + "ms-9d12ad4ebdf37419994c", + "ms-d9e7096b437e7ce7cda8", + "ms-cc52efdf7f7b9efb1f9e", + "ms-5902e6c4037ecda77a5f", + "ms-4118fc55a3d3e6bee8c0", + "ms-07c39d2c4595d6826040" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000006.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-c3747425fe51bb3facbb", + "ms-71e32140ee4d77f3b275", + "ms-923e9fb031870b12d9b4", + "ms-89782a37ecf6a44d23dc", + "ms-d368732bc01f4100cf07", + "ms-578fcc1f0f27f76ff933", + "ms-01f463a409e2ec0c6c3e", + "ms-cb752f1a3b24f95b34a2", + "ms-3efaacea10b7eb4adb57", + "ms-06e6e1db86d9e82e56cf", + "ms-f3f100caf89a2487faaf", + "ms-b33c18c58125678627a6", + "ms-bc99c1baec671e7f85a0", + "ms-220e8ddf7349fcbaf64d", + "ms-d854deebde42f3bd2fd5", + "ms-49685adafe2cb8da301f", + "ms-535c35e29640f289d7d1", + "ms-060d5756eac8f17652ee", + "ms-96ebb8e6e71464ac9172", + "ms-d6b0d6d7c66509eb4543", + "ms-fb0e4c4b9775d8121124", + "ms-a40882954ff51b61c49c", + "ms-ab72e600d0bfb162ea80", + "ms-1e25b12c832314a0b015", + "ms-65a8786e4fb1527342f7", + "ms-6c8d06bace98a28a8e4e", + "ms-928153d18b111f66f8d1", + "ms-af57106141e3bc564cb9", + "ms-ecb37ebbab21a506d8ca", + "ms-554f2960f25ff627b596", + "ms-499d2429a00f3cc35add", + "ms-68067370d3b20b9bd4e2", + "ms-4734f618e638d42bae46", + "ms-6d7030c17759f787a361", + "ms-5dfbb16e173e1a3e2bd6", + "ms-5973a2df030fcd05cd24", + "ms-daff16c01452c735345f", + "ms-1e4219bffa35c646ebc9", + "ms-a601803235735bea0610", + "ms-a31b52cc38a95273af77", + "ms-20f3cfcf183c69853b1e", + "ms-b0fee26e1db967f90b81", + "ms-4d687a60b9ce1f0f1ddd", + "ms-b4f96b4975ceb839c5e3", + "ms-19a3669e096e98232ec8", + "ms-733f43f6aaed1e84c020", + "ms-3945306c3b05b122fa9b", + "ms-9c7464444eaec0fa07c0", + "ms-de304cefeb20dcc17ede", + "ms-38ed02dacc133a888d3a", + "ms-1a9e9bcd6cbae272a459", + "ms-a2ee6034df65b51c4d97", + "ms-1d025a89862bf212d621", + "ms-5c52a7bf5a1abb1a1517", + "ms-4aed1197d0906ae51c34", + "ms-00ce5c82d2af61a7d7e8", + "ms-71dc14431350e581537c", + "ms-dc2e360a4876f6edce82", + "ms-68f53650da9a9354b58f", + "ms-03beb7dddeba28c64afa", + "ms-e327fef230e956dab37c", + "ms-9b89af1761e33e733a56", + "ms-cbc1fadecdd913c44ce5", + "ms-49516a81fe99935ef15e" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000007.jsonl", + "record_count": 832, + "group_ids": [ + "ms-230391b47c390f51bc07", + "ms-325cd8b08e89bb498c82", + "ms-19b9e8a1108ac25f3476", + "ms-625f1a55bfa2bcd57e24", + "ms-c9c7a51e43dd2c33de08", + "ms-0547201761e8e1873482", + "ms-b1b98e21cd4886f70563", + "ms-7ac4c9b07c994cd878f9", + "ms-bfb169ce850eae1c3995", + "ms-110944a3259f0e310db9", + "ms-0ccc81e9a63663c50096", + "ms-2da5f229550b61852483", + "ms-074bf4470a893ac7955a", + "ms-d8cc5ba8ab1cd4a46519", + "ms-f50f1c1b673d0c96d4ce", + "ms-95028134c6e3670b89d5", + "ms-fdf4af00031d452221d9", + "ms-ba396eb71806ac992acd", + "ms-ce9505086ea115f86649", + "ms-86b55250e01b9b6fba7e", + "ms-d4c5d7eea19d94314e2e", + "ms-783f5e7f8c889bef010e", + "ms-0f7cd7813265e3795e0a", + "ms-da2f79e022bb63a2af11", + "ms-a7760c53ebe20579190e", + "ms-573ab0ef1d4166be9b54", + "ms-54f7cae3bd3cb84062d3", + "ms-5edf28a87bfd0731da89", + "ms-d46da887312a51a9f2a0", + "ms-801727be95e46ae5da07", + "ms-77c4a01dcf26dcd8f0f3", + "ms-fee86da43c17dbf07059", + "ms-eaccbf417256e38345aa", + "ms-63f384579c130870fde3", + "ms-5d6aa86eb739a9ef05cf", + "ms-7e2762eb1b21b04fe873", + "ms-d07e29d33ae0ae6ebabc", + "ms-e2f303a0d1e6017e1fc4", + "ms-e71e880736efd3e42c02", + "ms-d7254aa722729aa1a031", + "ms-eb7189a4f2b0f0a5d4bd", + "ms-e1f9403543e48fbdde9d", + "ms-b75d7924ddc6f1056f17", + "ms-2e683dfff18ef7e90d7d", + "ms-7c300541f4986f403b47", + "ms-8b45028168ed68ec8303", + "ms-b5f0f2ae74766beae291", + "ms-036f329c4a881ff59e5f", + "ms-2fef2f43af8009bb88f1", + "ms-27d723ae87abc93a95f0", + "ms-20976520a892ae80760a", + "ms-5dfeb87aa77e4660d09b" + ], + "format": "jsonl" + } + ], + "record_count": 8000, + "group_count": 500 + } +} + +✅ StackCube-v1 generation complete diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_4.err b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_4.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_4.out b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_4.out new file mode 100644 index 0000000000000000000000000000000000000000..5915d341a8b6f117dcf6a22ddcd800465ca07d08 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_4.out @@ -0,0 +1,647 @@ +================================================== +Task: LiftPegUpright-v1 +Groups: 500 +Horizon: 16 (vs baseline 4) +Demo: /scratch/knguy52/dovla/maniskill_multitask_demos/LiftPegUpright-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5 +Output: /scratch/knguy52/dovla/experiments/six_task_h16_collection/LiftPegUpright-v1 +================================================== +ManiSkill CIL progress: batch=1/500 groups=1/500 records=16 restore_max=1.192e-07 +ManiSkill CIL progress: batch=25/500 groups=25/500 records=400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=50/500 groups=50/500 records=800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=75/500 groups=75/500 records=1200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=100/500 groups=100/500 records=1600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=125/500 groups=125/500 records=2000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=150/500 groups=150/500 records=2400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=175/500 groups=175/500 records=2800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=200/500 groups=200/500 records=3200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=225/500 groups=225/500 records=3600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=250/500 groups=250/500 records=4000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=275/500 groups=275/500 records=4400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=300/500 groups=300/500 records=4800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=325/500 groups=325/500 records=5200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=350/500 groups=350/500 records=5600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=375/500 groups=375/500 records=6000 restore_max=4.768e-07 +ManiSkill CIL progress: batch=400/500 groups=400/500 records=6400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=425/500 groups=425/500 records=6800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=450/500 groups=450/500 records=7200 restore_max=4.768e-07 +ManiSkill CIL progress: batch=475/500 groups=475/500 records=7600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=500/500 groups=500/500 records=8000 restore_max=4.768e-07 +{ + "backend": "maniskill", + "env_id": "LiftPegUpright-v1", + "control_mode": "pd_ee_delta_pose", + "sim_backend": "physx_cuda:0", + "source_obs_mode": "state", + "instruction": "Lift the peg and hold it upright.", + "target_actors": [ + "peg" + ], + "reference_actors": [], + "demo_path": "/scratch/knguy52/dovla/maniskill_multitask_demos/LiftPegUpright-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5", + "num_groups": 500, + "group_offset": 0, + "available_unique_states": 12377, + "excluded_post_success_states": 14938, + "branch_selection": "pre_success_only", + "num_records": 8000, + "k": 16, + "horizon": 16, + "seed": 0, + "measurement": "exact_state_physics_branch", + "candidate_mode": "structured", + "branch_execution": "parallel", + "state_batch_size": 1, + "state_restore_max_error": 4.76837158203125e-07, + "state_storage": "archive", + "state_archive": "/scratch/knguy52/dovla/experiments/six_task_h16_collection/LiftPegUpright-v1/state_archive.pkl", + "observation_archive": null, + "image_quality": null, + "manifest": "/scratch/knguy52/dovla/experiments/six_task_h16_collection/LiftPegUpright-v1/manifest.json", + "quality": { + "expert_success_rate": 0.848, + "no_op_success_rate": 0.048, + "mean_reward_spread": 0.782273677021265, + "nondegenerate_group_fraction": 1.0 + }, + "sharding": { + "dataset_name": "cil_liftpegupright-v1", + "version": "0.1", + "created_at": "2026-06-25T23:25:44.301335+00:00", + "backend": "maniskill", + "num_groups": 500, + "num_records": 8000, + "k": 16, + "task_count": 1, + "seed": 0, + "schema_version": "0.1", + "format": "dovla_cil", + "shard_format": "jsonl", + "index_format": "jsonl", + "shard_size": 1024, + "shard_count": 8, + "group_index_path": "group_index.jsonl", + "record_index_path": "record_index.jsonl", + "shards": [ + { + "path": "shards/shard_000000.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-0a8109ce7ab2ddc1bb1e", + "ms-52ea0261959dfacdbcac", + "ms-b8956f04f70d6eeecb07", + "ms-4b3a296acc66a14476c3", + "ms-29027a8e3742362fd477", + "ms-b0302a41f75a827cd4e4", + "ms-043a17d678e610264bb8", + "ms-cce5ef6cfa69a08a1e05", + "ms-eeeff97e2c39be877cb0", + "ms-c663f7e73e7e8376a05c", + "ms-f1eb6805ec9b483bba41", + "ms-fc9b4c94b51f8b2a521b", + "ms-8864796e31bbe93d7c5d", + "ms-3e0dbff4309ecbd35d36", + "ms-effca6c13ef9ccfd21a0", + "ms-2f2b16367db3db120235", + "ms-457db124a2f46a06175b", + "ms-1d728ae415a4d3cb782e", + "ms-cf2603516869355b8b13", + "ms-a20322ed13a2be7a99c6", + "ms-835bd462e505aa659d16", + "ms-a1a99e83fb2868164e40", + "ms-938b53ef53e669518a56", + "ms-b04b489f4fbad93b65ca", + "ms-6d6456c45f9f3ae74b50", + "ms-73906fdb1ffcfbe4aaf4", + "ms-cba95475cb68ea42b041", + "ms-1955408aeb8a8a4d0847", + "ms-31fe1d98bc7adcd515db", + "ms-30e2d8b7825053276e0c", + "ms-b025d8068f740557d630", + "ms-4c638b9bb924547f69cb", + "ms-c2b904f450b45e2a8c9a", + "ms-0ef5d3238b33bdebae4d", + "ms-aee5716b3e4fc47e46ba", + "ms-b4281a37b8b48a0b4c54", + "ms-4abb3dced67d5a70fa38", + "ms-484ea653bbca2de3a4a2", + "ms-ace117030ae44c474db1", + "ms-01926c30568d417edd91", + "ms-88f7199271947fc4c7e8", + "ms-639913ac0e6b3997faae", + "ms-1a2d298bf9972ddc1a7e", + "ms-9b687dc2de14c6813d44", + "ms-df961fe24a43e56d1716", + "ms-a323215fb409ab3e9287", + "ms-e2a037eadb166228b1dd", + "ms-ef2d84500e27679d22d9", + "ms-7a1b60ea1461eb460bed", + "ms-894cd0e5dafb2d38831d", + "ms-8be46dac6fc06c6c928a", + "ms-8fabc0ac20ec87ead026", + "ms-d602b7370c5b90f8a439", + "ms-dfc38c43e9a87f99ef7f", + "ms-7334e620435d1f3545fd", + "ms-d3377ca0da444ae9b40a", + "ms-760c73ea34418ce7a006", + "ms-3a97ee208c6be9ed8b74", + "ms-28fe4b5a2888189c28d9", + "ms-32943eea3e70cd02f313", + "ms-dd32ce9859c365c535b4", + "ms-26322bd4de52dbd5743c", + "ms-8556e974064f31ae1468", + "ms-b00933d2f628bc51d80a" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000001.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-cf44ae14778c60c36af0", + "ms-4dc9ab19fbe571f09699", + "ms-f2c8835e8a162283a646", + "ms-7614105b40956cc606ee", + "ms-a2c1429e7b81479b52db", + "ms-3c3fb1d170c1a7965ac4", + "ms-acbb7db5601f93bef99d", + "ms-9a72e47e2204fa794718", + "ms-01ab68aa5051226638ee", + "ms-0f02ff35b7fa1234c52e", + "ms-10c475db809a353fa815", + "ms-bb567fcac946f012535a", + "ms-4c72a9ad08cbf51799fb", + "ms-0251a326271aa1cab82b", + "ms-59474486dd5cb4e5bfd5", + "ms-bdb04ada069c35d51e2e", + "ms-d80d4dc70953915f21e7", + "ms-3661d65df223386f1be9", + "ms-d1cc32a5c52e78aa5e39", + "ms-3352cd51af5ae6aac4c0", + "ms-6a03497f97765ffd444e", + "ms-85e3feba1c0f3084cfa0", + "ms-93b076203c0c65fb22be", + "ms-61c099575bcff2de95e3", + "ms-b3d9d9570bf9b64c18fe", + "ms-69cb3fa758f707a58076", + "ms-ec4accef21246e0c9213", + "ms-b319cc2760d8d48f4617", + "ms-ee28aef4a7ee25ff7f7c", + "ms-0d80d6d87a27a53951fe", + "ms-619d3d285d2cae3abbb9", + "ms-d2839c13649000454a65", + "ms-3efca8ee8623914a4faf", + "ms-8c4a08985213114fffb7", + "ms-3db845076f92a16d1840", + "ms-8e23d2d44e709420e46d", + "ms-c3818c2c4474a784fc6b", + "ms-9867545cf31b7f92b81f", + "ms-18ab33df290e3919a468", + "ms-8d3d849234ceb98c3dc1", + "ms-c910a469460f71880556", + "ms-e825996818ee6f2f7d26", + "ms-0bc111c5ce4e0052716b", + "ms-e89a3ba400f81d2795f3", + "ms-6e268d1e2915740816a1", + "ms-d5241a32e5c9de09f2e4", + "ms-6959fa589edc996753d7", + "ms-662379c2c351807b9cdf", + "ms-fe790caba27d92511f71", + "ms-83019373a06eb384c886", + "ms-a8eb841beb710c8fbfb1", + "ms-9257b89b9b635a0fb174", + "ms-741e906eafd57d115b5d", + "ms-debd7e5668e2acdf6bb0", + "ms-118ba8926eda06ffc61e", + "ms-b483cac145e16ccdf42a", + "ms-7be84e650aec22374654", + "ms-e3c96ca4b56c4a26d40d", + "ms-080e034927abe2f30d60", + "ms-a99d874a81bc60301007", + "ms-136f1ebebd0b697c3310", + "ms-e4090b0b0c1739070201", + "ms-051860c228aa5160d5a6", + "ms-08dc71adea7d8efeb298" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000002.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-320d99264926ea029f95", + "ms-8a8afc1963fb2220ecb6", + "ms-c0ed69f549e56aa55f3d", + "ms-684a66885f21739d826b", + "ms-2d931673692269f11e8b", + "ms-2c6000014d31fc0bb21c", + "ms-487d0fe2a13fb699dd7a", + "ms-e7573c5d29c3a1a3f228", + "ms-8c8cf210aa675c5db1fe", + "ms-3f0c6556f6c7df39ec2d", + "ms-38758efd8620a312267d", + "ms-0fe90299788e1f71f98a", + "ms-22aaa483f5e7d381b63d", + "ms-057132b6fc7922013391", + "ms-74c9fc741b846b5019af", + "ms-8054283375aeefc8f2c1", + "ms-e6478d71ed92dae06539", + "ms-c146ea1c89beb683b6fd", + "ms-4a86d570dabb0916f020", + "ms-0fe7053483f72737eaec", + "ms-bdc14f3225a9c9f26513", + "ms-c6c7803c56c24c460c7f", + "ms-33a0bdb4bcee42c79a2b", + "ms-3f56672eac1c00e9c2b0", + "ms-7ce799eb9d3b741e8960", + "ms-a932cbf7e2b5755db22e", + "ms-1b8e9b173562648021f6", + "ms-0e48a13af39db4a58bbc", + "ms-2b59042e27453aff03a7", + "ms-794706690fbf8c028bda", + "ms-6a9c4c704edd84c21cda", + "ms-8b15f77ad3436ac282ba", + "ms-227f2e245df67a5ef5ed", + "ms-efb5a07185a3a1b33226", + "ms-fdb2637276b27c46c757", + "ms-e3f38594991aafbffc23", + "ms-44a607bf50304347ead8", + "ms-9f844e24c570295e0497", + "ms-777412474251e64d7e24", + "ms-6684504de2fc8756f44f", + "ms-49c99f285560483a85ee", + "ms-b2d1a5e132eaeb353782", + "ms-9551db96b60a51ccf62a", + "ms-7272e1e33d4ecc07470c", + "ms-36d1128e3bd512cffe18", + "ms-0ee12f14646b925adf51", + "ms-6f03b16858c35025503e", + "ms-25195d526c9adb43ac79", + "ms-d30b515ab0ee174afa58", + "ms-6ca37421eb5405e1a9a5", + "ms-7cd314f076bfe5c32088", + "ms-e56b37f0e3881a70ef54", + "ms-507199452c6d42562a38", + "ms-4b4cd12d70c0615e0930", + "ms-2590f9fd91265598fb45", + "ms-1e0634dc4c3b94220950", + "ms-d229ac64c91c9bb6ce88", + "ms-73c99a7a25cd025e5f2d", + "ms-89520c3ac8779d9e81e7", + "ms-89feb4366865cfaf2273", + "ms-83c62951219e6599b630", + "ms-261f6f6b578cee63be20", + "ms-a5a56eefeb2de0abe600", + "ms-0660d5f65df5fc4b6132" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000003.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-caa0a38b7df789550f8a", + "ms-36f204d04ec4c4af3169", + "ms-55e89d326b6531345902", + "ms-e2f90eb6c79dfd9ce6fd", + "ms-3d862f2a32d3404926c3", + "ms-80e3d8a34c3cb94cad65", + "ms-577341cca32521f4045e", + "ms-9156a9a9c1b829b742ad", + "ms-988537a36edfece3d073", + "ms-a4ade4de51ca437ac432", + "ms-675833805d1139c2f2a0", + "ms-62b88f5cbc4f137e3bba", + "ms-af91167a0cf15457c6ba", + "ms-becd2ffef1c49174810a", + "ms-2e4ca0e2534ef8ea8bb4", + "ms-f0d2ffde521e268efc83", + "ms-9f4b47c80d59166302ab", + "ms-6939f3ecafd5e0a4bfee", + "ms-780ea60ea0255941e6c8", + "ms-690f52acc6ead9c7c078", + "ms-1145e0f6add300eee9bd", + "ms-efbde8652e75274fbd69", + "ms-ff3256da4e1731651937", + "ms-a195f05e99bccc48907e", + "ms-3af721139694f7f3ec04", + "ms-35d84da6eb1afcf8a811", + "ms-ec2e4afb34b439731136", + "ms-b4af9d7c922b1676e516", + "ms-96e1562aaca7ef49db28", + "ms-34f60c9a263bd0a5641f", + "ms-6a4398375c43dbf887b1", + "ms-2286d9cf8d88e878c4d7", + "ms-145f60d5fff285e7d05b", + "ms-490e2c471adbabec6304", + "ms-510cafd84de611de3909", + "ms-d9ed0a9faf86ca5f5caa", + "ms-b37738e210bf33a45a1d", + "ms-508d1ea98eb9c78acd16", + "ms-785fb44f1371895f7721", + "ms-483520829af0d733cb2a", + "ms-77b33fec5c6c627dbda5", + "ms-9852b3671856321cecc8", + "ms-02663903d88aa84c9c10", + "ms-d8f42ed6b0b686d819ab", + "ms-40a961867249ddccb68f", + "ms-13aa39b96923ad8c913a", + "ms-9ced92cccacfe4d157b0", + "ms-67b6e4d62036c13cc466", + "ms-1e41619ab39cf51b7822", + "ms-252706fc0657a82ba5a4", + "ms-5cc5e6359fc006559eae", + "ms-17a5e9b7bd8c1f255307", + "ms-dfc111660c0d1c7b7e04", + "ms-340913cb08bfca09a4e3", + "ms-5bbc4c4b1b51e2a088db", + "ms-c7c812bb98308442d95b", + "ms-baa36ee75ec2de550b9b", + "ms-0bf60eaed682d2ca149f", + "ms-50b9f93dd246400d20c8", + "ms-caeb4c3d358f256799a8", + "ms-f5ebbab81506fa27fc6d", + "ms-9fd14cc61816c2d4360e", + "ms-288ce929957d627c569c", + "ms-5003e0cda1e13bd4850a" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000004.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-a67a1ab4d96d811cec77", + "ms-c7677b8f230103dd3c51", + "ms-09b2d832d7823c11fd22", + "ms-4e825dd3ade30a42e899", + "ms-8fcf4d9d042de2f74d76", + "ms-e6ca7414c5035e80d0ac", + "ms-a43b35c03282fe0c8464", + "ms-575bd87a5da90963019d", + "ms-cbd6cc8b2708d5e3e4a9", + "ms-11e34e842b9154f34d56", + "ms-fb329ccd569163c53415", + "ms-28c139f34a954b692962", + "ms-af25c332ad54eee6e1c3", + "ms-074b26fed4c4a9dd3449", + "ms-3c550bf52d7bf91f4ca2", + "ms-fa077183f74add357112", + "ms-28835618507e7bc71227", + "ms-bbefdf03d0471be60f8f", + "ms-5b24bcde90348f24bb19", + "ms-8bc7aa419d17368d81b5", + "ms-3aa33c0fda88e08e75fa", + "ms-a40abdebe8a9c17de878", + "ms-295dab0e1d105e758827", + "ms-48db1178206eee52434a", + "ms-491aad3896ede89ea0dd", + "ms-4a845dab9503bc4eba04", + "ms-e9a43adf4bba03b65d5d", + "ms-04cc31b72a8dc8a91911", + "ms-125287f805d320b2eff7", + "ms-88fb2bebd038dc786c01", + "ms-062addb9da32a0c2ea97", + "ms-aa46bb6dc8e6274963b6", + "ms-7dfc3684b46df95070e2", + "ms-99f263cb3f383b1f6243", + "ms-074cb86179cdd7f52737", + "ms-1bbd61ce613ae75c8bb0", + "ms-f63a0e43f0923369edc3", + "ms-7a0330e9e3189ff035d3", + "ms-8b18ea4e0982a1352696", + "ms-03495469257613a802b4", + "ms-8f3d4e1d07bb5dcc7c57", + "ms-e705fef3e4caa4c49718", + "ms-ed30f8fbb80b2cf46fbc", + "ms-7aeb0835d3396e6666b6", + "ms-b0861c414f8690d62a82", + "ms-5b0141f336d28927cf31", + "ms-0a57aa04002f1b843856", + "ms-e2950dcb50ed2b955be6", + "ms-92070f417bd9d14d6fc7", + "ms-b60078e9e6ea9c26dedc", + "ms-ef8de99d917753ec603c", + "ms-faeff0379a053e9ea351", + "ms-84d10acd6db571034f10", + "ms-fdf18cd6260d82226927", + "ms-8a9e7d69f7ade362fe40", + "ms-69785feff08dbf9a50be", + "ms-35c12c467fb8a54f89f0", + "ms-6ceb341dfa656e485a13", + "ms-30e4ae59842091fe16b2", + "ms-3a09e411ee6c882ef7df", + "ms-502c195c0e4810814bac", + "ms-1cb9edbe8bd7ed4106c1", + "ms-8b461b0929537fbc1d41", + "ms-ab10d34780322326aa3d" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000005.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-83c0299c23fa800452c7", + "ms-a905b84bd3ab703c025f", + "ms-40392e5aeef63551f9ac", + "ms-bc29afbe2bbf00638d25", + "ms-3c1e7b47135487435fc1", + "ms-cc0b4a1deec28e53722a", + "ms-dd94921779f12b249579", + "ms-d8295b43b30ceea30907", + "ms-851d25e7a1dde4242f33", + "ms-2e7fc6c97e5ce1f024e5", + "ms-067e0fe46e4e9483e2a1", + "ms-14dabc3ac1e1a339ebe4", + "ms-4906f3adfda6dace862d", + "ms-85fcc6b4797417c3473f", + "ms-324a58dd51dfecb51fa8", + "ms-73259d17883355f7fa88", + "ms-aedc3dee6e275342eb8f", + "ms-84264ef080ec4f3631df", + "ms-ff012c1c8597c8db3cc7", + "ms-9756d07ed207e9ab1d53", + "ms-251c1dfd1010906f559f", + "ms-ac6e78b1c56d020d0c09", + "ms-a242f1d12ebfcd82de05", + "ms-502dba1003c09ec60ba0", + "ms-fae3c6a2f1b575ddc3f7", + "ms-5fc588b9c9d52034f4c4", + "ms-de8d6051e6ee74b62fec", + "ms-dbad4f4d0df3c19c7bbb", + "ms-761ab1e423462aa44e3d", + "ms-f97bd60d5430c32f9494", + "ms-73e3079ce5eb7cd14ec8", + "ms-cdbfe4d3b05f4ed12db3", + "ms-3fd410dabe0c3ef8dc1a", + "ms-6e2b02ee71a4a44bc325", + "ms-f7169e27b3743aa48e77", + "ms-da28f7d135db8979d5b8", + "ms-cc5f2a24e2571cb23050", + "ms-b439839ee4165bac92a8", + "ms-a6cf775fbba580161c4e", + "ms-de68dbb7daad3e3a71fe", + "ms-f539c5c30da21c875657", + "ms-d799435b5f6b8fae2ffb", + "ms-957002faf092f769aa8e", + "ms-91710558482886874dd7", + "ms-27251a7c7f13bf165e21", + "ms-bc56601ca069f65cd08f", + "ms-88f748134c3fd27fcf72", + "ms-f5dcf9054c3c1c4e1d50", + "ms-e91b824c3841686ed053", + "ms-13d8bae55d3c3710138f", + "ms-9f993a6e9bb26846d4f4", + "ms-98c8fcf4c9f098cebe8d", + "ms-01a86e5f5731af7002ab", + "ms-76c7efad95ae67dba1ca", + "ms-f1529c9d823325056e8d", + "ms-7a3daad6772d2384ea66", + "ms-57dc9349b8f2c9c49105", + "ms-8bab784ae53525967763", + "ms-c2c265f2f5013746eaa1", + "ms-ce7b349699ed248bb67c", + "ms-2f1164357c9732e059d4", + "ms-54b1042de0891da7c07e", + "ms-40fb8acb1e648bf3b617", + "ms-c0142e4a5d15dd8d647a" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000006.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-bbaa576369a59660ecdd", + "ms-a103dfbc0567d05ade6e", + "ms-f237f72fe3db2b4f8e12", + "ms-64a755437c34f39cda2a", + "ms-1b24faaf0d032ad20017", + "ms-7b1a03be4a61ba27261c", + "ms-a829d650eb9914fb2983", + "ms-95264bfdaa3892cf66ac", + "ms-832f0e966173abd16c34", + "ms-c9ac2c188beb33e70f98", + "ms-908213fef55ccf1dfda6", + "ms-2436538e05fdcf8e44fc", + "ms-68c63442972f322789b5", + "ms-97c28a7584b35207bf84", + "ms-8c7bbc80458fce91a1ac", + "ms-dd0fad10ce064047d62e", + "ms-cf8c5d9c874edc9f000d", + "ms-14a2c3eb293175cb76f3", + "ms-65e9da4468e522823587", + "ms-7d68051ee4e3b8a5c9a4", + "ms-0a6bff8326e7e32702f3", + "ms-29a11d7bdcf49c86bd7a", + "ms-83199198e5f952f8dd69", + "ms-9554f100dd5be3a4fa47", + "ms-929e1dd2134a8952111c", + "ms-c9e0182d4090ddaa991c", + "ms-ec4c6be448348c8021a9", + "ms-131a98e2d4a78a9645f8", + "ms-35babd5f83dcf1561f9a", + "ms-01bef10a5ccfd38294fe", + "ms-e467ecc93a4793286423", + "ms-c76aba3db8d4ce0062bc", + "ms-a5bd3a36c6fe41cd0b35", + "ms-40eb8d19498413d4e730", + "ms-59bff662e0ecb913ed81", + "ms-a40dee47c502e73949af", + "ms-e12808f670338ca180c6", + "ms-33294c8fed5a847bf149", + "ms-a74a87e26362d09598fe", + "ms-e1b6bcb41532a800d3cb", + "ms-6c7f34ae8218aa0952c8", + "ms-87fa87e06dc88c28a44f", + "ms-c871c1157a16e19e6150", + "ms-cf7bdf0776c928ee2342", + "ms-5fb06efe2b4bea27cc89", + "ms-524ae7618ed9b6a272ae", + "ms-b55ca348be8b64323714", + "ms-f601aab1d8f7e4f1b97f", + "ms-ad7b7d5d7fb25961859f", + "ms-8f24ad2e217963943c12", + "ms-3f6e7085c865144a7408", + "ms-20cd6d8ccbf13e36dcdb", + "ms-3996b923bb25440c29e3", + "ms-5b4f2522663bec742b02", + "ms-3176a570b75767d1d729", + "ms-917b0f051489227842f0", + "ms-68d9596217cc37c949c4", + "ms-ebeb5587176e3877b98a", + "ms-605dbfbe6bef2b15b58c", + "ms-f81ea99a1ec3423f3b6e", + "ms-7b229b60d7f64fb624cf", + "ms-648250f7c25953eab081", + "ms-9c33d4837de7d23232d6", + "ms-ba9541c7192498a1d8bd" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000007.jsonl", + "record_count": 832, + "group_ids": [ + "ms-e0952f3218e0bf6fe027", + "ms-9bb65f13e00984482ef8", + "ms-e3526036a5bf85cac0bf", + "ms-076246c4c5d1ca8ccd9d", + "ms-093472a082b3b61e73d6", + "ms-ecb5cc61f0b736e3f06f", + "ms-f4060cdbe6d39b9215ac", + "ms-d26dc2146eb8ae9a5434", + "ms-1d5ca9681600acdaba39", + "ms-865f9d670f806778ae61", + "ms-078317e23cc56a9d719d", + "ms-a6e4ac1995198def96a8", + "ms-a4420cda0b159dda82a2", + "ms-44e7935dbed5f3b3c739", + "ms-97924f0cc5e7bd79c4a3", + "ms-8be8a38748f9fb037556", + "ms-6818a4edf6ef35dacfb0", + "ms-c090bce47703a4c4bb0b", + "ms-08ab006a923f5fdae198", + "ms-f3211ecde366ca9f6ac7", + "ms-efc33064e2a2f7e86dbe", + "ms-ea164b21572d39ffed58", + "ms-ea41340182245326a218", + "ms-22d1b2dcc954dcdff384", + "ms-a96ff9a1d5571ec4a6f2", + "ms-06cb0796d1097e6ca111", + "ms-02cacb03f445c8a460e8", + "ms-fc3c8a3572fd96442723", + "ms-e8de8d619c0cf5fd633e", + "ms-715db95b13ebc6c977c8", + "ms-b933fbfb7c8fa68a445f", + "ms-5799584ed40e211c50a5", + "ms-b383acd9844466b19579", + "ms-fb131b52f94500da894b", + "ms-80d8d8fc4a3308139975", + "ms-f6f0d14ef7c24b9e4621", + "ms-ed4f6388ae96f9da76fa", + "ms-9bbea286acf66b5cfb8e", + "ms-c1f3359c70d068a36081", + "ms-472376eba1a5e1c706bf", + "ms-d2febdb2a469fbbac4f2", + "ms-b5b68468bf26d19252f0", + "ms-059d7962a5f998edf668", + "ms-b0f3a35453f8c458ada8", + "ms-76b0b0463fdb323b6233", + "ms-d3a7e3f163b142a3e56b", + "ms-062a265cedee2107eb17", + "ms-6abc3749b5df33ea621e", + "ms-4f54ed10d332989ae524", + "ms-e536e40f6902c6ba0927", + "ms-22d006041320a4487c21", + "ms-737e72810112f3246b05" + ], + "format": "jsonl" + } + ], + "record_count": 8000, + "group_count": 500 + } +} + +✅ LiftPegUpright-v1 generation complete diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_5.err b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_5.err new file mode 100644 index 0000000000000000000000000000000000000000..0578de3c25479a29bc7e62fc6909dabadc00954b --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_5.err @@ -0,0 +1,17 @@ +Traceback (most recent call last): + File "/lustre09/project/6037638/knguy52/vla/scripts/generate_maniskill_lattice.py", line 103, in + raise SystemExit(main()) + ^^^^^^ + File "/lustre09/project/6037638/knguy52/vla/scripts/generate_maniskill_lattice.py", line 76, in main + summary = generate_maniskill_lattice( + ^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/lustre09/project/6037638/knguy52/vla/dovla_cil/generation/maniskill_lattice.py", line 294, in generate_maniskill_lattice + prepared = [ + ^ + File "/lustre09/project/6037638/knguy52/vla/dovla_cil/generation/maniskill_lattice.py", line 295, in + _prepare_group( + File "/lustre09/project/6037638/knguy52/vla/dovla_cil/generation/maniskill_lattice.py", line 614, in _prepare_group + _validate_task_profile_state(task_profile, state, env_id=config.env_id) + File "/lustre09/project/6037638/knguy52/vla/dovla_cil/generation/maniskill_lattice.py", line 987, in _validate_task_profile_state + raise ValueError( +ValueError: ManiSkill task profile for PegInsertionSide-v1 references missing actors: ['box_with_hole', 'peg']; available actors: ['box_with_hole_0', 'peg_0', 'table-workspace'] diff --git a/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_5.out b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_5.out new file mode 100644 index 0000000000000000000000000000000000000000..a5577100813b1b0cf587b268aa5d24231708da92 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_6task_h16_14740806_5.out @@ -0,0 +1,7 @@ +================================================== +Task: PegInsertionSide-v1 +Groups: 500 +Horizon: 16 (vs baseline 4) +Demo: /scratch/knguy52/dovla/maniskill_multitask_demos/PegInsertionSide-v1/rl/trajectory.h5 +Output: /scratch/knguy52/dovla/experiments/six_task_h16_collection/PegInsertionSide-v1 +================================================== diff --git a/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738111.err b/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738111.err new file mode 100644 index 0000000000000000000000000000000000000000..ae0e306a81a55bc3223753adf8442bd929764d19 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738111.err @@ -0,0 +1 @@ +FATAL: container creation failed: mount hook function failure: mount /proc/self/fd/9->/cvmfs/soft.computecanada.ca/easybuild/software/2023/x86-64-v3/Core/apptainer/1.4.5/var/apptainer/mnt/session/rootfs error: while mounting image /proc/self/fd/9: squashfuse_ll failed to mount /cvmfs/soft.computecanada.ca/easybuild/software/2023/x86-64-v3/Core/apptainer/1.4.5/var/apptainer/mnt/session/rootfs in 10s diff --git a/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738111.out b/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738111.out new file mode 100644 index 0000000000000000000000000000000000000000..72e48d4418e7117f1008778160ea3bc3d43912f8 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738111.out @@ -0,0 +1,3 @@ +================================================== +Generating PickCube horizon=4, 200 groups, K=16 +================================================== diff --git a/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738189.err b/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738189.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738189.out b/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738189.out new file mode 100644 index 0000000000000000000000000000000000000000..3743a160f5006c6d7b0fcf4ead2bb2e84117d1da --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_horizon_sweep_14738189.out @@ -0,0 +1,1263 @@ +================================================== +Generating PickCube horizon=4, 200 groups, K=16 +================================================== +ManiSkill CIL progress: batch=1/200 groups=1/200 records=16 restore_max=1.192e-07 +ManiSkill CIL progress: batch=10/200 groups=10/200 records=160 restore_max=4.768e-07 +ManiSkill CIL progress: batch=20/200 groups=20/200 records=320 restore_max=4.768e-07 +ManiSkill CIL progress: batch=30/200 groups=30/200 records=480 restore_max=4.768e-07 +ManiSkill CIL progress: batch=40/200 groups=40/200 records=640 restore_max=4.768e-07 +ManiSkill CIL progress: batch=50/200 groups=50/200 records=800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=60/200 groups=60/200 records=960 restore_max=4.768e-07 +ManiSkill CIL progress: batch=70/200 groups=70/200 records=1120 restore_max=4.768e-07 +ManiSkill CIL progress: batch=80/200 groups=80/200 records=1280 restore_max=4.768e-07 +ManiSkill CIL progress: batch=90/200 groups=90/200 records=1440 restore_max=4.768e-07 +ManiSkill CIL progress: batch=100/200 groups=100/200 records=1600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=110/200 groups=110/200 records=1760 restore_max=4.768e-07 +ManiSkill CIL progress: batch=120/200 groups=120/200 records=1920 restore_max=4.768e-07 +ManiSkill CIL progress: batch=130/200 groups=130/200 records=2080 restore_max=4.768e-07 +ManiSkill CIL progress: batch=140/200 groups=140/200 records=2240 restore_max=4.768e-07 +ManiSkill CIL progress: batch=150/200 groups=150/200 records=2400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=160/200 groups=160/200 records=2560 restore_max=4.768e-07 +ManiSkill CIL progress: batch=170/200 groups=170/200 records=2720 restore_max=4.768e-07 +ManiSkill CIL progress: batch=180/200 groups=180/200 records=2880 restore_max=4.768e-07 +ManiSkill CIL progress: batch=190/200 groups=190/200 records=3040 restore_max=4.768e-07 +ManiSkill CIL progress: batch=200/200 groups=200/200 records=3200 restore_max=4.768e-07 +{ + "backend": "maniskill", + "env_id": "PickCube-v1", + "control_mode": "pd_ee_delta_pose", + "sim_backend": "physx_cuda:0", + "source_obs_mode": "state", + "instruction": "Pick up the cube and move it to the goal position.", + "target_actors": [ + "cube" + ], + "reference_actors": [], + "demo_path": "/scratch/knguy52/dovla/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5", + "num_groups": 200, + "group_offset": 0, + "available_unique_states": 14843, + "excluded_post_success_states": 32478, + "branch_selection": "pre_success_only", + "num_records": 3200, + "k": 16, + "horizon": 4, + "seed": 0, + "measurement": "exact_state_physics_branch", + "candidate_mode": "structured", + "branch_execution": "parallel", + "state_batch_size": 1, + "state_restore_max_error": 4.76837158203125e-07, + "state_storage": "archive", + "state_archive": "/scratch/knguy52/dovla/experiments/horizon_sweep_pickcube/h4/state_archive.pkl", + "observation_archive": null, + "image_quality": null, + "manifest": "/scratch/knguy52/dovla/experiments/horizon_sweep_pickcube/h4/manifest.json", + "quality": { + "expert_success_rate": 0.3, + "no_op_success_rate": 0.315, + "mean_reward_spread": 0.6260121015552431, + "nondegenerate_group_fraction": 1.0 + }, + "sharding": { + "dataset_name": "cil_pickcube-v1", + "version": "0.1", + "created_at": "2026-06-25T22:46:49.738245+00:00", + "backend": "maniskill", + "num_groups": 200, + "num_records": 3200, + "k": 16, + "task_count": 1, + "seed": 0, + "schema_version": "0.1", + "format": "dovla_cil", + "shard_format": "jsonl", + "index_format": "jsonl", + "shard_size": 1024, + "shard_count": 4, + "group_index_path": "group_index.jsonl", + "record_index_path": "record_index.jsonl", + "shards": [ + { + "path": "shards/shard_000000.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-4d4a0c1d076355ffd0a1", + "ms-bca92beb7317318ca011", + "ms-abe94723b59a97d9763f", + "ms-c4ffabbc14eaeffb221e", + "ms-b242270287e07877b824", + "ms-c780d41a3bffb09815d1", + "ms-0cd563cc20f169c74fbe", + "ms-8d3c88198175fb4ea4ea", + "ms-64d4fb96883ad9d71105", + "ms-1781d9c562b4425df7c8", + "ms-10a67a57026bb09b5d93", + "ms-25bb20f3481a2a40b97f", + "ms-496a8f7067502294e7d5", + "ms-4f0e885363fe62511c75", + "ms-d148b9f8ea23279497f8", + "ms-4cfa415aafc7719da7cb", + "ms-66ce9cebb51b9b3e0ec9", + "ms-d5776179091bcaa6bfef", + "ms-ad552f33bf02f49c14e0", + "ms-ffa71e2ef6c284477c5d", + "ms-d8f4b5b065679f609133", + "ms-a23dfdb03199ebfbe361", + "ms-49e6045ca8183d41136a", + "ms-a5dba9c8830dad961855", + "ms-eaf5582c55650fd2865e", + "ms-6ef99ad1f8d26b7e38b9", + "ms-2024693cc2a26ab5d1d6", + "ms-13ed8d3ada0ea318c48e", + "ms-1fafbeb8c89797db4df1", + "ms-41c5500153912a1743de", + "ms-23f9cca8edfca421a4a6", + "ms-30479bd8835c8cffd136", + "ms-74f0a3bc945c1175dcf9", + "ms-2b964be4d105427ca0eb", + "ms-5523eba579ea5435e91b", + "ms-48f126a76afc90929dce", + "ms-03c9e06de458c92718b6", + "ms-63ee4e1d2585364188e6", + "ms-98912ce92f4695f8a722", + "ms-95a7907813646b272351", + "ms-4703006fc8907554a40f", + "ms-a5c7855e36f6a06573ec", + "ms-04a6b44c453098afffb7", + "ms-f28e6bf7251f0bd37ad3", + "ms-93e74122b21eade07d51", + "ms-917da1d4f673ff5d6478", + "ms-4fe7fdaf946c97803c82", + "ms-a504d09412788a2773d4", + "ms-e91c7bd22a5ef871c7a8", + "ms-c70be4072fb4e49ac106", + "ms-3052c61870b8a6a2babf", + "ms-0d5b7bcb2c7d740b7696", + "ms-2fbb831120ce2aa6f719", + "ms-54e89e2a2f321cf70dba", + "ms-bdf26065f475465826c6", + "ms-a475caaaf6e28f4d7455", + "ms-348b051d0c9cce13075c", + "ms-60091066afb152fc824b", + "ms-e85317d6541282304ce2", + "ms-f99f5b9aaf1c3cb8c61a", + "ms-2d358d7a4c2b67f9f4af", + "ms-069dc3698b517474beaa", + "ms-d153b103860cfed38ad4", + "ms-b1b95727df4df8049ed0" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000001.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-200442ca86c3e0ff72d3", + "ms-df0e2503a672ad5a7dfa", + "ms-668bcf0ac5abe6fa9eb2", + "ms-9d14b3fd054f71709da1", + "ms-b7ccda5de6a71599f01c", + "ms-1d6c0b6f71473d660554", + "ms-63527e1eb9c23a9fc75c", + "ms-7c88d066c0670593cbab", + "ms-cb2a99420b780553611b", + "ms-99315346e54a20597826", + "ms-0daf0e9cb362ad8f496d", + "ms-b5a3a6e783695829efb6", + "ms-26494448a7b481f4f731", + "ms-e754ba00de42c999505e", + "ms-e9245db9f5904d49a256", + "ms-b22e5e94d0d212d6ccf7", + "ms-8971438a1b978d474d22", + "ms-22472c2a41ea7417dbf5", + "ms-b8c4184879c58deef6d3", + "ms-bba9a8a8c59251aa46c6", + "ms-f909cd438b9863f313bb", + "ms-55721a5d3e1fdcf10c82", + "ms-862568dc8af1af81dd8e", + "ms-435132a90c4ed562f9a3", + "ms-fac68fa4b5d4bc75c93b", + "ms-92dc694be5753c389b18", + "ms-c29cb85e8ed1e57344e1", + "ms-bdd1924b44de2db5b07a", + "ms-0a94a5266e8cd61093b0", + "ms-6634347d1d05f493d81a", + "ms-7b2401918efe0651d9ba", + "ms-83ec15e8f774953fe628", + "ms-82b1fb6002418272c296", + "ms-15d95f7d16cce6ebd32a", + "ms-e3b2d37ab46676716861", + "ms-853df40d9d3635a88d27", + "ms-de066fc4200163ebd06b", + "ms-741fa9eee8cb1e3e8715", + "ms-d29dc8ba85ab059f12be", + "ms-094599a68b9c76efdd5e", + "ms-a4e68f873c6eb7cde656", + "ms-3d492f8dc983407a69be", + "ms-5be34d2caaa3a2d07b81", + "ms-6c75a86def52a6de51db", + "ms-9f5f0ae60c86f28f86d3", + "ms-2c0a6db89491fb8a22fd", + "ms-732819b0c220a4c06870", + "ms-e09c417c1752c6b79a25", + "ms-7e4e4e077f37d19a5a54", + "ms-bb8ed2bb13e673e8cf6f", + "ms-1000a9a61a2c8b126789", + "ms-33e08f7f8709bf47e1e1", + "ms-3c44eccd6ff65242fc0e", + "ms-1934f9ca7523faa409be", + "ms-d6bad605ce8b1c0dd88c", + "ms-409e2940d4aa6c0875f0", + "ms-3f358662ceabe01f6958", + "ms-b117616bb6a08977600b", + "ms-0e662962f4c8a0461d84", + "ms-2c2c25b2184dbd3ced8c", + "ms-19e6e6e341aebb7ec693", + "ms-e643b49422e59ee60c39", + "ms-4271d509276b727650b8", + "ms-ec66397acf785f69cb19" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000002.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-cc0b33f92229d32bd84f", + "ms-3580a31bd16361d41fd6", + "ms-63b97e4c442b733ff0f0", + "ms-140f82db55ada1fb1c6c", + "ms-22d7d2dbb5a8a9e19332", + "ms-c664113ccbc0e281f1c8", + "ms-190981eba0c8a456c826", + "ms-353b701e2d2196fb683b", + "ms-bf185da00128e9d87812", + "ms-7d52804e9827f9818f92", + "ms-68ae84f4467fb6f65d65", + "ms-056350b937268a7d08ea", + "ms-6d8340dd4d6f97f6d0e7", + "ms-ee631606b03e5effb5dd", + "ms-3be7002e37f6a4e0f025", + "ms-a5840c419337ed41e2ae", + "ms-4de2b3e94b14c8c78537", + "ms-485f14220e1dec7daaf0", + "ms-54f34c73458181a65e9a", + "ms-e8cb2881ab61a6d780cf", + "ms-2c90b2f2d540a30721aa", + "ms-772c63fc56a725c2f90d", + "ms-f53338c86edb8f05daee", + "ms-6a7e80154a9a4e5e9c08", + "ms-6dfdea7de7121c915b62", + "ms-9b787a7d66583f2ed8d7", + "ms-6ace61e55ad1db294439", + "ms-00c960cc3aa6dde8639a", + "ms-55f75d9dc41aca97d809", + "ms-1ccc983e34160673e824", + "ms-e136f606e83e70e2c496", + "ms-ee35eb25b1433abf699c", + "ms-c3ad4f5cc58ad48c4dd6", + "ms-5c2e48512913cb65b655", + "ms-3e5e9a4b8928a9f35065", + "ms-a58c7308751d5a6eac87", + "ms-ac79bc3a841958ce127f", + "ms-a4732c3146a8b438250c", + "ms-8a5288c566d4f1dffa81", + "ms-065d1413cfb46d6e04fe", + "ms-51f1b8a0c7be0b4c0548", + "ms-9c8a90334c3b933c8976", + "ms-ece64b396479ec50ccbc", + "ms-4e395d53889980b05784", + "ms-b36673c3905a123cbd50", + "ms-ba14d5d87151e4d27963", + "ms-5d38cc22b00ffb3084dd", + "ms-c33a1c105c774ed26226", + "ms-64101e562538f4ab7809", + "ms-c39cc9d4b29cd1fad2eb", + "ms-946c9e6b2d4f0c4f03b7", + "ms-a3fa10726c656e851515", + "ms-00b8d4d0ecd9baefa88c", + "ms-3e2b55bf2f722d1945e8", + "ms-afd45a2b6172dd18978d", + "ms-a89ece44974129ad64b7", + "ms-6c2b5371d28af36108b2", + "ms-afc6c9e64e4683a561a5", + "ms-fbb174db946d8d539542", + "ms-7c8bff83175e003a88f6", + "ms-a439b131346b9ab4c471", + "ms-4d6b31619ee142940316", + "ms-fa911d56cdf9bf122b5c", + "ms-0d409fcb57e1c9412639" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000003.jsonl", + "record_count": 128, + "group_ids": [ + "ms-55878a97ed159530a1d5", + "ms-68d5cede9291ac64bc27", + "ms-3b2cae18686d1b86c1dd", + "ms-b49514c985ef3d5e821c", + "ms-42e0460ee5d27a254556", + "ms-bc7d17d78a1ed0966115", + "ms-b82cc8ce5d9d348ba6b4", + "ms-b09d36637e4849b2d9aa" + ], + "format": "jsonl" + } + ], + "record_count": 3200, + "group_count": 200 + } +} +================================================== +Generating PickCube horizon=8, 200 groups, K=16 +================================================== +ManiSkill CIL progress: batch=1/200 groups=1/200 records=16 restore_max=4.768e-07 +ManiSkill CIL progress: batch=10/200 groups=10/200 records=160 restore_max=4.768e-07 +ManiSkill CIL progress: batch=20/200 groups=20/200 records=320 restore_max=4.768e-07 +ManiSkill CIL progress: batch=30/200 groups=30/200 records=480 restore_max=4.768e-07 +ManiSkill CIL progress: batch=40/200 groups=40/200 records=640 restore_max=4.768e-07 +ManiSkill CIL progress: batch=50/200 groups=50/200 records=800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=60/200 groups=60/200 records=960 restore_max=4.768e-07 +ManiSkill CIL progress: batch=70/200 groups=70/200 records=1120 restore_max=4.768e-07 +ManiSkill CIL progress: batch=80/200 groups=80/200 records=1280 restore_max=4.768e-07 +ManiSkill CIL progress: batch=90/200 groups=90/200 records=1440 restore_max=4.768e-07 +ManiSkill CIL progress: batch=100/200 groups=100/200 records=1600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=110/200 groups=110/200 records=1760 restore_max=4.768e-07 +ManiSkill CIL progress: batch=120/200 groups=120/200 records=1920 restore_max=4.768e-07 +ManiSkill CIL progress: batch=130/200 groups=130/200 records=2080 restore_max=4.768e-07 +ManiSkill CIL progress: batch=140/200 groups=140/200 records=2240 restore_max=4.768e-07 +ManiSkill CIL progress: batch=150/200 groups=150/200 records=2400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=160/200 groups=160/200 records=2560 restore_max=4.768e-07 +ManiSkill CIL progress: batch=170/200 groups=170/200 records=2720 restore_max=4.768e-07 +ManiSkill CIL progress: batch=180/200 groups=180/200 records=2880 restore_max=4.768e-07 +ManiSkill CIL progress: batch=190/200 groups=190/200 records=3040 restore_max=4.768e-07 +ManiSkill CIL progress: batch=200/200 groups=200/200 records=3200 restore_max=4.768e-07 +{ + "backend": "maniskill", + "env_id": "PickCube-v1", + "control_mode": "pd_ee_delta_pose", + "sim_backend": "physx_cuda:0", + "source_obs_mode": "state", + "instruction": "Pick up the cube and move it to the goal position.", + "target_actors": [ + "cube" + ], + "reference_actors": [], + "demo_path": "/scratch/knguy52/dovla/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5", + "num_groups": 200, + "group_offset": 0, + "available_unique_states": 14807, + "excluded_post_success_states": 28462, + "branch_selection": "pre_success_only", + "num_records": 3200, + "k": 16, + "horizon": 8, + "seed": 0, + "measurement": "exact_state_physics_branch", + "candidate_mode": "structured", + "branch_execution": "parallel", + "state_batch_size": 1, + "state_restore_max_error": 4.76837158203125e-07, + "state_storage": "archive", + "state_archive": "/scratch/knguy52/dovla/experiments/horizon_sweep_pickcube/h8/state_archive.pkl", + "observation_archive": null, + "image_quality": null, + "manifest": "/scratch/knguy52/dovla/experiments/horizon_sweep_pickcube/h8/manifest.json", + "quality": { + "expert_success_rate": 0.535, + "no_op_success_rate": 0.295, + "mean_reward_spread": 0.7835982487862929, + "nondegenerate_group_fraction": 1.0 + }, + "sharding": { + "dataset_name": "cil_pickcube-v1", + "version": "0.1", + "created_at": "2026-06-25T22:48:40.028082+00:00", + "backend": "maniskill", + "num_groups": 200, + "num_records": 3200, + "k": 16, + "task_count": 1, + "seed": 0, + "schema_version": "0.1", + "format": "dovla_cil", + "shard_format": "jsonl", + "index_format": "jsonl", + "shard_size": 1024, + "shard_count": 4, + "group_index_path": "group_index.jsonl", + "record_index_path": "record_index.jsonl", + "shards": [ + { + "path": "shards/shard_000000.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-70e74817db8efcce6284", + "ms-bdd1924b44de2db5b07a", + "ms-2a8bf5ca4b09bc12dfff", + "ms-fbcca760f5e16aa5c776", + "ms-a82b54b0a431b181dc3a", + "ms-d6704a1938d7d76a0d68", + "ms-3159f4ff1bb23c1862ac", + "ms-39bab34b5a63d4e1fffe", + "ms-f32f5f14cb4444621588", + "ms-d2b229d126e999a7a41b", + "ms-0ddbbdeff64530b308ea", + "ms-d7d0ae3b15ea1c2ad666", + "ms-6e9d17a0560730b4da70", + "ms-1c6703202f2b118f7c3a", + "ms-8be5cfd4b0b4eb7eebf1", + "ms-2e6a776c6458c2b4cbac", + "ms-b36673c3905a123cbd50", + "ms-ae0e5165e0739c852d5f", + "ms-bed888983455dd07e3af", + "ms-859f6f872bc9589cbbe4", + "ms-4436390726de06466ddd", + "ms-c83760c257e1a416a1a4", + "ms-f4e1b8d4bbca03c03b33", + "ms-53cfcc878e30fc078bba", + "ms-da2586165c8acbc0b3f3", + "ms-809abf4af9e0febd6ede", + "ms-6dd79816e2e4ee23b340", + "ms-58696a1895ae6a89ff2a", + "ms-ed4d8770d400a91a69a2", + "ms-40450a44f96a9f7ec9c8", + "ms-6f15a8b002ca48e49972", + "ms-6a58413c885e005805e9", + "ms-793a98188fd57ed5291f", + "ms-39d1d72b546f80f9a248", + "ms-4f0ae704228e87452739", + "ms-61b1188974f35233b883", + "ms-9f82e21d66216788cbcf", + "ms-2b3e081263ed25d45dc8", + "ms-0273408007263508bd7d", + "ms-7241a1c25af41f4c0b5a", + "ms-8c7773a3edb91357e6ad", + "ms-de066fc4200163ebd06b", + "ms-d2e526ce1c3799165f5d", + "ms-fc0e18a31dc5a1b30ac0", + "ms-13b3c492538e5b3c7bdd", + "ms-a7d7c2e45e40a4b60d6a", + "ms-9ef4a8a81ea7b71e98f1", + "ms-594f0b571501751686f9", + "ms-bcad9c259c090bfc6fae", + "ms-da46641cf609b4fe64d9", + "ms-e53a214271af6f6a0aca", + "ms-c53094e0e48d25560cfa", + "ms-0b5ce46ea370fad9c632", + "ms-88d9733b9d7de0de586a", + "ms-223fc901aeff03c46ebf", + "ms-7f081dbe6abce1a3925c", + "ms-dd4e13e425fbe9bd9140", + "ms-091ef4afa4240cafc027", + "ms-99991d29e7c23423b56a", + "ms-459623af7b540303fc00", + "ms-f66db08453e1dfca28cc", + "ms-667addf55fb4c1ccb3cc", + "ms-50870ce69f2ca504786b", + "ms-57aaf810f569473a6a40" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000001.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-8227b79f7504e3a4686c", + "ms-9bf66fb70daf64c27841", + "ms-4ce55719d21c62d89947", + "ms-b6b04acf562ac58f2ac4", + "ms-47fd1407ae4e17ed4bee", + "ms-59523004ccd5a334cadc", + "ms-94af4fb5b9d36c9ca8a4", + "ms-3bebd7507fdd592cac39", + "ms-fb8c95ea19067ebbc2a6", + "ms-e00064787cfe1b250a23", + "ms-a5b05c099bee3faa84df", + "ms-40a547e7912cff4c6b7c", + "ms-6a03fbec2a7cf9502a99", + "ms-6c18bcd4c4981dcc032f", + "ms-d865ae39711121706dc4", + "ms-77a38f183dbcd073dc85", + "ms-6fb2c49d1c61698d3386", + "ms-bf5d8ff48e0b2211cdf9", + "ms-0649c1ec411240ebd919", + "ms-ac53c7d7e8a46ac6203b", + "ms-1c0d0f17323ac9361331", + "ms-adc20acf78cd92b9fe5f", + "ms-d61698d72fe1cce5a7ef", + "ms-98c99310a16d36f79b68", + "ms-e4ced12c123f09850c3f", + "ms-9b4036632d22e6872b1e", + "ms-7377f696275f889434c3", + "ms-c2dbbe21273e2b5cd471", + "ms-bd055ea5f615cb4954ff", + "ms-4c8cd4128a70fe2faab2", + "ms-692a7b031370d70f9f80", + "ms-e3cf8e50a9e3c34d0a91", + "ms-732819b0c220a4c06870", + "ms-396d9d1a433e3e118cd8", + "ms-57a213acdc6a7231f799", + "ms-b7ae86c6b363e5451302", + "ms-55a6f49f4214a612604b", + "ms-7a40ddb45c3e5e13d61a", + "ms-81d46bd79f7c359b63d6", + "ms-9a077df5edf6f622fb1c", + "ms-8650530ded91c0bdb3ed", + "ms-4af8752a0b3be60828c5", + "ms-6c2b5371d28af36108b2", + "ms-22675530f3a401b3fd6d", + "ms-c39504ad79cd0f07d66a", + "ms-eb962ca4917d0fe444f2", + "ms-c70be4072fb4e49ac106", + "ms-4e5038d61af88b105290", + "ms-4fe7fdaf946c97803c82", + "ms-a96e6c033a50f14a5ea6", + "ms-a1b95e3a5a2befa833fe", + "ms-6788d2c07dc039c2a4f8", + "ms-a4a5a781840bb278c0b5", + "ms-24d86230701fb9a1a6e6", + "ms-17479881c474352c839e", + "ms-f9f6443f4f8028c710f9", + "ms-1b7dac9d82acc7dfd741", + "ms-5f98f13f17e8f41f7436", + "ms-8d7ca2122559a2f95cf8", + "ms-8fd44079105a8cc3df5e", + "ms-cc0b33f92229d32bd84f", + "ms-403fbb24fd0e1c0b22a3", + "ms-ea7998648747b5ac4cba", + "ms-813cde56f2b0f33dcd20" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000002.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-5f9f1641c5b2857ef081", + "ms-dee7ae7a8a627bed9175", + "ms-8e501dc127ba60951c53", + "ms-fdb2cb96fe6cc3bf44f6", + "ms-91a9396b18f3c5150081", + "ms-91d86cba2602d1fa780d", + "ms-f9c5f834558cb0c143e6", + "ms-dd92bc8f7237467d7789", + "ms-78d581e5cce9372cabda", + "ms-361778f642adb63f6d2c", + "ms-3fdfa5ee46d6dca7a3ac", + "ms-eebd2f71d2692e5c2586", + "ms-21931c8c5f66cdafd9e3", + "ms-2b5482a9f23bad38c857", + "ms-7b70ceab9d1975a21106", + "ms-18ac2c89b2452864c78b", + "ms-13007c00cf33130d380d", + "ms-41236e85e52923ddc010", + "ms-30f6e74c0ede27a0deff", + "ms-6a6a82fd3ca9bb320c63", + "ms-488c343fb11b3ffca9c9", + "ms-cd637a2176ec46ab5df5", + "ms-4742d169606f3ec2b474", + "ms-1e25701d7256f6c80214", + "ms-ec1fc47374d29cb83990", + "ms-36fa630978b54d3d1777", + "ms-90de3d24f1d7b5174cf7", + "ms-39450dbba738dbd8a764", + "ms-0fb7e933abf150f48450", + "ms-f7cec504ca49a413ccd6", + "ms-d7f89fe89a8f49f727df", + "ms-6771bdb6b5e267beba98", + "ms-9ca0e1d0290b4b84c0f6", + "ms-2c56f58baa056978767d", + "ms-d0151d452f08f4f415bf", + "ms-7223bfbd5d15e40ed4f0", + "ms-860f361ab8a037c23346", + "ms-e56b7b690881ceff092d", + "ms-61c1c75734906c960987", + "ms-c0311194823721590a76", + "ms-28a3f2f39f1b56e89434", + "ms-876e915764260b679962", + "ms-42117385b9e6e57c7e63", + "ms-1b92c1867d57a089a0fe", + "ms-b106c4986ec29120dad2", + "ms-baa3195a82062b92e0d3", + "ms-07dc9a1297634db6715d", + "ms-e02882888390f925a075", + "ms-05f3413245611b8c64a2", + "ms-5eb10f6087e63c69ebfa", + "ms-725a538370280a362cb2", + "ms-7651a7f31eae24958ef0", + "ms-4341d089f3a88bad6994", + "ms-746ba1b0d37a457717d4", + "ms-c9f64bbc6a2252e9400a", + "ms-a325cde0d9e77392923e", + "ms-ca5b45dbcfb0b2f3dca8", + "ms-d0afc781a650a973baa5", + "ms-4855f6a2d0522cc97d24", + "ms-7491edf89e021a82c306", + "ms-c2a4ebfff50f20d05db7", + "ms-10ba0ff6636d626f4604", + "ms-c0363e27d0e998b62e1b", + "ms-768bc04b2d5ddc08cd43" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000003.jsonl", + "record_count": 128, + "group_ids": [ + "ms-9555b707fa9f198b4b59", + "ms-5898dc426404542f511d", + "ms-09eb96273d8590e1c8d8", + "ms-ec6d6e2a6dcebf98b4fe", + "ms-67215d0a68a7d175081b", + "ms-acecb73ee2997cc1eb5e", + "ms-d0bb3b07813e6d1b0da6", + "ms-dc9c91ea6fd48856d94d" + ], + "format": "jsonl" + } + ], + "record_count": 3200, + "group_count": 200 + } +} +================================================== +Generating PickCube horizon=16, 200 groups, K=16 +================================================== +ManiSkill CIL progress: batch=1/200 groups=1/200 records=16 restore_max=1.192e-07 +ManiSkill CIL progress: batch=10/200 groups=10/200 records=160 restore_max=1.192e-07 +ManiSkill CIL progress: batch=20/200 groups=20/200 records=320 restore_max=1.192e-07 +ManiSkill CIL progress: batch=30/200 groups=30/200 records=480 restore_max=4.768e-07 +ManiSkill CIL progress: batch=40/200 groups=40/200 records=640 restore_max=4.768e-07 +ManiSkill CIL progress: batch=50/200 groups=50/200 records=800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=60/200 groups=60/200 records=960 restore_max=4.768e-07 +ManiSkill CIL progress: batch=70/200 groups=70/200 records=1120 restore_max=4.768e-07 +ManiSkill CIL progress: batch=80/200 groups=80/200 records=1280 restore_max=4.768e-07 +ManiSkill CIL progress: batch=90/200 groups=90/200 records=1440 restore_max=4.768e-07 +ManiSkill CIL progress: batch=100/200 groups=100/200 records=1600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=110/200 groups=110/200 records=1760 restore_max=4.768e-07 +ManiSkill CIL progress: batch=120/200 groups=120/200 records=1920 restore_max=4.768e-07 +ManiSkill CIL progress: batch=130/200 groups=130/200 records=2080 restore_max=4.768e-07 +ManiSkill CIL progress: batch=140/200 groups=140/200 records=2240 restore_max=4.768e-07 +ManiSkill CIL progress: batch=150/200 groups=150/200 records=2400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=160/200 groups=160/200 records=2560 restore_max=4.768e-07 +ManiSkill CIL progress: batch=170/200 groups=170/200 records=2720 restore_max=4.768e-07 +ManiSkill CIL progress: batch=180/200 groups=180/200 records=2880 restore_max=4.768e-07 +ManiSkill CIL progress: batch=190/200 groups=190/200 records=3040 restore_max=4.768e-07 +ManiSkill CIL progress: batch=200/200 groups=200/200 records=3200 restore_max=4.768e-07 +{ + "backend": "maniskill", + "env_id": "PickCube-v1", + "control_mode": "pd_ee_delta_pose", + "sim_backend": "physx_cuda:0", + "source_obs_mode": "state", + "instruction": "Pick up the cube and move it to the goal position.", + "target_actors": [ + "cube" + ], + "reference_actors": [], + "demo_path": "/scratch/knguy52/dovla/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5", + "num_groups": 200, + "group_offset": 0, + "available_unique_states": 14701, + "excluded_post_success_states": 20466, + "branch_selection": "pre_success_only", + "num_records": 3200, + "k": 16, + "horizon": 16, + "seed": 0, + "measurement": "exact_state_physics_branch", + "candidate_mode": "structured", + "branch_execution": "parallel", + "state_batch_size": 1, + "state_restore_max_error": 4.76837158203125e-07, + "state_storage": "archive", + "state_archive": "/scratch/knguy52/dovla/experiments/horizon_sweep_pickcube/h16/state_archive.pkl", + "observation_archive": null, + "image_quality": null, + "manifest": "/scratch/knguy52/dovla/experiments/horizon_sweep_pickcube/h16/manifest.json", + "quality": { + "expert_success_rate": 0.93, + "no_op_success_rate": 0.305, + "mean_reward_spread": 0.9565701169201566, + "nondegenerate_group_fraction": 1.0 + }, + "sharding": { + "dataset_name": "cil_pickcube-v1", + "version": "0.1", + "created_at": "2026-06-25T22:50:58.236913+00:00", + "backend": "maniskill", + "num_groups": 200, + "num_records": 3200, + "k": 16, + "task_count": 1, + "seed": 0, + "schema_version": "0.1", + "format": "dovla_cil", + "shard_format": "jsonl", + "index_format": "jsonl", + "shard_size": 1024, + "shard_count": 4, + "group_index_path": "group_index.jsonl", + "record_index_path": "record_index.jsonl", + "shards": [ + { + "path": "shards/shard_000000.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-d88107eff1e361cdcc4c", + "ms-4cd9eae753bf9e48f520", + "ms-6edb3f8fe977a81f8979", + "ms-23a355317920e192147f", + "ms-faeacb4289bbfe504de1", + "ms-f88a44616a0985aac518", + "ms-c8666cbd7e5a9d0fd036", + "ms-f5483f0819e0e8bc352e", + "ms-22675530f3a401b3fd6d", + "ms-751daa0137578c486f3d", + "ms-193063285cdb9faff39b", + "ms-5d430ebe6f7933dc6c8c", + "ms-67f602a0e4efc9984d72", + "ms-564c02492863ae749ea9", + "ms-86b0447349029dfaa6ab", + "ms-b481531ab787e3d6fb1e", + "ms-e441e4042cb7d77a4986", + "ms-5dfad80e83d210946447", + "ms-be8eff6e6d520ae11a0d", + "ms-93e3d9968204474eff49", + "ms-708baf01c22c75cafd22", + "ms-0e55645c0c701caac48a", + "ms-29dd2cdbb66fe13ab877", + "ms-4fe47d946641f0096a85", + "ms-fb98254adaa23dc128df", + "ms-10621e1108a69d0b633e", + "ms-fd225482b96686f51dc8", + "ms-c07472c627b51086ff33", + "ms-f07607bc0ce1c4e202b0", + "ms-23d6e5835aa90a62b6ef", + "ms-1309e1770b6ca6e71c7e", + "ms-52c6c227fd4b61609031", + "ms-8a22f0da78e59cd1eb5f", + "ms-4acbf69acd30d96e46a7", + "ms-1518756d5da0b8d3c190", + "ms-417cab1470cf46273b02", + "ms-ee05717b29d39d7cddff", + "ms-7020f1cda882b79bffa3", + "ms-91d86cba2602d1fa780d", + "ms-52cdfff790d71133c5a0", + "ms-caed5c1b1c882141e2e9", + "ms-766b6fe7c08a3525ce94", + "ms-ea72196e20393870bcbb", + "ms-b6827e38768a63212d73", + "ms-a01e21b91a5b37a16afe", + "ms-c425c810f7865e9564ee", + "ms-f39bbaa063044719a98d", + "ms-c3cdd70e2bc10bba1930", + "ms-ed9addab41f3298db7bc", + "ms-eb16021deb5d03f15c9d", + "ms-912718b1f2ef92066799", + "ms-9783ccaf76f1b20c766e", + "ms-ffc3f18cbc674bdcc085", + "ms-3c2515d4a10082514714", + "ms-bef9b88ae40ca98647aa", + "ms-7a17ceb004979b99a586", + "ms-2be3bf4f348678c0cc6c", + "ms-e8fbdbddd9200f5091d6", + "ms-2dc8e489908f72903319", + "ms-0ca11ec8f7fbf0c1d66e", + "ms-f653a3ac16eba9fba5dd", + "ms-e058c0dba52df04b8642", + "ms-c0bfce3c9dd10f9d3231", + "ms-296b75a11f61d61d5fba" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000001.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-8d1978e8c3540cb5b1f6", + "ms-6355260c6fa49f291729", + "ms-0e275939f12a9dcee0a0", + "ms-fe2964cc8940ee7c24aa", + "ms-11a2794efa24f4d93d83", + "ms-3731f1c5af3c4a133633", + "ms-090e18187bb6ae4a9120", + "ms-4f938bd3599c372c8bff", + "ms-362ba7cebb2bab1421e8", + "ms-70bc2fe01ba217676b26", + "ms-c70be4072fb4e49ac106", + "ms-625b5d55ca76cc565049", + "ms-0f219ce35aff264455e2", + "ms-d9e3a187dcadb9ad26b1", + "ms-7e6493bb6b1ff650d6de", + "ms-4276f24b1b5dce771bff", + "ms-4385e9c56007fdaa94b5", + "ms-2e5c2f75d3f6311e0a1c", + "ms-bb60defa0d234da9e0bc", + "ms-a509a572b3871e591d26", + "ms-f5bae191ebad7c9f2687", + "ms-e589f7d40c12eff1712e", + "ms-17330235d51598ebcb4e", + "ms-910364d2d86091c5e0cf", + "ms-5316151a88a3282436ca", + "ms-6b9ede8946200ad0cd5d", + "ms-5f98f13f17e8f41f7436", + "ms-fbcc714230de2c097d0d", + "ms-023e021f079a8774c741", + "ms-cc0b33f92229d32bd84f", + "ms-a196fc3a729550509b18", + "ms-227fe0347d2ea57a66a3", + "ms-10cb98b0a1ca5429dd28", + "ms-a5a651f9e28ea52016fe", + "ms-42edbdb8ff281fc33a99", + "ms-d849d754f346b826e052", + "ms-0c7dab7ee052231f3f57", + "ms-76b9f40367a1ee020744", + "ms-66540aefd55a84b5088f", + "ms-bdd1924b44de2db5b07a", + "ms-8af87d14945da6a910e4", + "ms-732819b0c220a4c06870", + "ms-7d962acc9cfbc27c599c", + "ms-b6df140b235c8fe0eaad", + "ms-547f4204888aee978bf0", + "ms-723a1f2e22afe15b6981", + "ms-10f2865bef6da7b66423", + "ms-a8e6435948250f5e163d", + "ms-351818a836474e594d3c", + "ms-42117385b9e6e57c7e63", + "ms-4c83412fa0aa5180ad89", + "ms-1c0e2bb3a811fac6d567", + "ms-03cdc738473ac1df9dfc", + "ms-fb92a1e6b67121783fff", + "ms-759ba8b7be233ca47c30", + "ms-c92e8da7621fbf50cd4e", + "ms-55c2bba961a327aa5260", + "ms-0ece7b8c91d9d82f244b", + "ms-0ca1c25fddb7da05d80a", + "ms-8f0e2c9caad14bd37ecf", + "ms-68e331088400739273da", + "ms-83d2b158041f4eed3646", + "ms-ba156a5cdaf40cc85cc4", + "ms-2cc8c1975139989940f0" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000002.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-da46296bd4e9032570df", + "ms-0843bef6b0c13c3ede69", + "ms-642670f4bb658d32970b", + "ms-76c78750d7abc6570487", + "ms-197fee321caebfd320f5", + "ms-7571c367fe683b0f1cf9", + "ms-13192a27f5fcf4fd24d7", + "ms-bfda95ffffd5d1a2683a", + "ms-4ca84594a789c9abd5d9", + "ms-c5d026dffb39f2f1d38a", + "ms-91ad1a5f945470b4346e", + "ms-263556e167dc363ebc42", + "ms-e07f80b44faf328c5e58", + "ms-8a74cf665984273be00d", + "ms-4cd1ceea2dda2732110a", + "ms-dc2a8fa76e1ceb0d8925", + "ms-746436cac05100783ad7", + "ms-486ff31b7a04a7e9e8f9", + "ms-793a98188fd57ed5291f", + "ms-2af09f91e05d19f7f181", + "ms-999a8c1454dd4597a834", + "ms-309c09cf5d782e9ba0a3", + "ms-309d2258d275ae304857", + "ms-fbb609a3f2a334986245", + "ms-2ae57bab6d2ea11f6638", + "ms-7849511ad9b8262adae7", + "ms-28b1f1e4380d434605de", + "ms-c35a083df7c08225af3b", + "ms-14c40aa19f7b9307511d", + "ms-ae6d60345612be32478c", + "ms-f77305a344da6a4ccc84", + "ms-fd6deac5a281d169bc75", + "ms-6c335deaaa5c63478701", + "ms-15ef6ee70a8df3cfa578", + "ms-5e23219a8e2aadd039c7", + "ms-d26a60e383a2e822615a", + "ms-4ddda51b1c6809c3531d", + "ms-5e85dc3760eef8a53a67", + "ms-704fee4c9fb9941a26e7", + "ms-578e24198ed828a74485", + "ms-95320371ce6b820c85f6", + "ms-910d27b58a0585327a27", + "ms-8fcc78fba7b3bec0fab6", + "ms-348af0b70d3e679f40a2", + "ms-25839c690d85a4f5b986", + "ms-a0e07d6e4725342ffff2", + "ms-d965f63e091a5d6dd6d7", + "ms-64ab703e7a82799a47dd", + "ms-4e395d53889980b05784", + "ms-a3d0c460532de2d801ae", + "ms-dea17c1c383f54d05636", + "ms-85eef30b36b044b33f60", + "ms-e0b2e69a5ea0d0ad4fb5", + "ms-81729abe5d83a1ca280d", + "ms-4d6099258a3a6e037aff", + "ms-307f9067d5b66f55b3a6", + "ms-eac0d7dd54b1e797fdbd", + "ms-774e954ffdc5da041261", + "ms-ae5421846020d2765f82", + "ms-b33f1d4ddb6a531d2a71", + "ms-ed419c64a6693cb8998b", + "ms-1b9f208b8780cb33b9a5", + "ms-80d5c509e83845ab1bc7", + "ms-578537600f7a649c2dd5" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000003.jsonl", + "record_count": 128, + "group_ids": [ + "ms-3a07c49566decfea04de", + "ms-295de2ad6c52ab50de1a", + "ms-ff3f17774184180998f0", + "ms-59cc3cdf6918248d7f78", + "ms-08858c7ba8d4a3357e05", + "ms-5b122bd43e749623ebe0", + "ms-4d03463bb4536a28d67a", + "ms-1abd8c7cc1c683f88d3d" + ], + "format": "jsonl" + } + ], + "record_count": 3200, + "group_count": 200 + } +} +================================================== +Generating PickCube horizon=32, 200 groups, K=16 +================================================== +ManiSkill CIL progress: batch=1/200 groups=1/200 records=16 restore_max=1.192e-07 +ManiSkill CIL progress: batch=10/200 groups=10/200 records=160 restore_max=4.768e-07 +ManiSkill CIL progress: batch=20/200 groups=20/200 records=320 restore_max=4.768e-07 +ManiSkill CIL progress: batch=30/200 groups=30/200 records=480 restore_max=4.768e-07 +ManiSkill CIL progress: batch=40/200 groups=40/200 records=640 restore_max=4.768e-07 +ManiSkill CIL progress: batch=50/200 groups=50/200 records=800 restore_max=4.768e-07 +ManiSkill CIL progress: batch=60/200 groups=60/200 records=960 restore_max=4.768e-07 +ManiSkill CIL progress: batch=70/200 groups=70/200 records=1120 restore_max=4.768e-07 +ManiSkill CIL progress: batch=80/200 groups=80/200 records=1280 restore_max=4.768e-07 +ManiSkill CIL progress: batch=90/200 groups=90/200 records=1440 restore_max=4.768e-07 +ManiSkill CIL progress: batch=100/200 groups=100/200 records=1600 restore_max=4.768e-07 +ManiSkill CIL progress: batch=110/200 groups=110/200 records=1760 restore_max=4.768e-07 +ManiSkill CIL progress: batch=120/200 groups=120/200 records=1920 restore_max=4.768e-07 +ManiSkill CIL progress: batch=130/200 groups=130/200 records=2080 restore_max=4.768e-07 +ManiSkill CIL progress: batch=140/200 groups=140/200 records=2240 restore_max=4.768e-07 +ManiSkill CIL progress: batch=150/200 groups=150/200 records=2400 restore_max=4.768e-07 +ManiSkill CIL progress: batch=160/200 groups=160/200 records=2560 restore_max=4.768e-07 +ManiSkill CIL progress: batch=170/200 groups=170/200 records=2720 restore_max=4.768e-07 +ManiSkill CIL progress: batch=180/200 groups=180/200 records=2880 restore_max=4.768e-07 +ManiSkill CIL progress: batch=190/200 groups=190/200 records=3040 restore_max=4.768e-07 +ManiSkill CIL progress: batch=200/200 groups=200/200 records=3200 restore_max=4.768e-07 +{ + "backend": "maniskill", + "env_id": "PickCube-v1", + "control_mode": "pd_ee_delta_pose", + "sim_backend": "physx_cuda:0", + "source_obs_mode": "state", + "instruction": "Pick up the cube and move it to the goal position.", + "target_actors": [ + "cube" + ], + "reference_actors": [], + "demo_path": "/scratch/knguy52/dovla/maniskill_data/demos/PickCube-v1/rl/trajectory.none.pd_ee_delta_pose.physx_cuda.h5", + "num_groups": 200, + "group_offset": 0, + "available_unique_states": 14337, + "excluded_post_success_states": 4731, + "branch_selection": "pre_success_only", + "num_records": 3200, + "k": 16, + "horizon": 32, + "seed": 0, + "measurement": "exact_state_physics_branch", + "candidate_mode": "structured", + "branch_execution": "parallel", + "state_batch_size": 1, + "state_restore_max_error": 4.76837158203125e-07, + "state_storage": "archive", + "state_archive": "/scratch/knguy52/dovla/experiments/horizon_sweep_pickcube/h32/state_archive.pkl", + "observation_archive": null, + "image_quality": null, + "manifest": "/scratch/knguy52/dovla/experiments/horizon_sweep_pickcube/h32/manifest.json", + "quality": { + "expert_success_rate": 0.995, + "no_op_success_rate": 0.28, + "mean_reward_spread": 0.9841186416660201, + "nondegenerate_group_fraction": 1.0 + }, + "sharding": { + "dataset_name": "cil_pickcube-v1", + "version": "0.1", + "created_at": "2026-06-25T22:55:17.713437+00:00", + "backend": "maniskill", + "num_groups": 200, + "num_records": 3200, + "k": 16, + "task_count": 1, + "seed": 0, + "schema_version": "0.1", + "format": "dovla_cil", + "shard_format": "jsonl", + "index_format": "jsonl", + "shard_size": 1024, + "shard_count": 4, + "group_index_path": "group_index.jsonl", + "record_index_path": "record_index.jsonl", + "shards": [ + { + "path": "shards/shard_000000.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-10be03140c81e2ca56a7", + "ms-b2df26c0dcab6f43a851", + "ms-9cb55854160479fd42b7", + "ms-2adf32a09dbd22061443", + "ms-fd5a11509082d5325696", + "ms-42b71b44f6e9574e44cf", + "ms-b03af4b1a52f2e258ec7", + "ms-7fa85a7251186eb18d0b", + "ms-7c4e7ce246320e6b7051", + "ms-4fc3e61500492a6fe710", + "ms-6d5a77bdfc55e2306733", + "ms-2c673faf341d0b0a02a8", + "ms-3b4acfa38c84a75497ee", + "ms-10296e5569ac7111d591", + "ms-796258142d0a9b3cb421", + "ms-5962356a991d3b9fc025", + "ms-fac68fa4b5d4bc75c93b", + "ms-a24f0b949308972297fa", + "ms-0db2912f6b7ce02bf020", + "ms-87336cddd35ac5f1a58c", + "ms-73030c866b88a1110531", + "ms-b542d5be59b5c09da8c3", + "ms-5db8a7cf17e5fa84dd5b", + "ms-14f4d818730e7e8ab95c", + "ms-785c0694d6535c980f25", + "ms-1cd88db78512901f8606", + "ms-4854716c68ae72ede6c9", + "ms-4988c0c8d77e237c6e8c", + "ms-c167f80c39ace83aa6f6", + "ms-9c1b61bdb54844cd0e72", + "ms-1fe5287f9133cb43820b", + "ms-4d4619d8d63bd7193dca", + "ms-694a5b37d1d896120821", + "ms-42808025916d80bcb74b", + "ms-bcbf06e317a04cc056fd", + "ms-16353c015336781f2665", + "ms-a2a6ba4294a3add8bda2", + "ms-0daf0e9cb362ad8f496d", + "ms-c8e1cd70634fe29a73db", + "ms-98812c1762d974cabbca", + "ms-3e76b3d13b3565bd2b44", + "ms-c9e31a19a776a861f9d7", + "ms-0581d5a6af398797bb2a", + "ms-3587e67b02d28aa2f66d", + "ms-60c3b3e6d844d8d97cd4", + "ms-cbd1e7ab8024170ee6e9", + "ms-519d00ea175cb26f4c68", + "ms-a40a46ea7afd43bed2f0", + "ms-6bbe9db1c57513f3b478", + "ms-2dcd03e051052c389f6e", + "ms-6d0186ae40380798f616", + "ms-c6145f69bc14dd6e1ea8", + "ms-1c1ed9edff1f41128597", + "ms-fa3c9dbe3998f2159827", + "ms-2959394ef33c99fef4c8", + "ms-88a7a88d26373e61484f", + "ms-32f88f6f0c153a70ab8f", + "ms-3ea30989082ec714e12d", + "ms-8a4d41cbf188919f0197", + "ms-7f5df2f9b3ef3dd5409f", + "ms-e8cc49336231e022f622", + "ms-99dbfd39c342cfdc9e52", + "ms-36bf961a94031699b975", + "ms-094599a68b9c76efdd5e" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000001.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-4fa0958e505f771ba041", + "ms-0c4d634fd8e69052e372", + "ms-5d5eea7a330077bb2b0e", + "ms-d35c2d5461e582039e5f", + "ms-05f3413245611b8c64a2", + "ms-5f698a446b96301e9a06", + "ms-db9f1afc822c8c585d10", + "ms-3065f7bc5e50276b88f5", + "ms-7c7e14e911dce8ed55e5", + "ms-b90d7926b23eabbbed04", + "ms-3439c1dc07e0835413a7", + "ms-9cfbacd53ed514a78635", + "ms-cc7962a122edfdbdf26c", + "ms-532e02b908c8534bd4d3", + "ms-3ac6b66614771abb18a3", + "ms-4f083fe5dd6ed691567c", + "ms-27be3e64cc1da9438497", + "ms-5a0b20e862942c5a0c14", + "ms-5ebfe344d592d873237b", + "ms-81ab9d7c1175157a9d72", + "ms-017959cff869f72bfc4d", + "ms-806c400da841970829b2", + "ms-960a8ecf0df111ae5c5a", + "ms-87f65c35c9d8b3ced6fc", + "ms-c90cb9564837b9a9b82c", + "ms-a41006cf9c607b5dc3a1", + "ms-784258dc67c71a79e3c8", + "ms-0df22ebd9c2d9cb3e3de", + "ms-6892b7fc41ab88106f71", + "ms-8bb2d8cb7c0d33fa0c1d", + "ms-e8a970ac4f144a720bf2", + "ms-19e35cb776fd96df3cb2", + "ms-95f6b45e01d839e0665d", + "ms-c4a0a2141f7837ad82f1", + "ms-3fb76e075066c7de7dbb", + "ms-a3b20f31d76c75fe3e48", + "ms-14434f333173b1cfd765", + "ms-5fd2e802ba562ff84bac", + "ms-cad07270370091aa1166", + "ms-eec3668f9f19663cc1db", + "ms-3b0639bfe64ecfd35594", + "ms-beec7981f43fa705df7a", + "ms-30ad687dc1bc8318fa45", + "ms-09265b2ce109a182fd9e", + "ms-4eae919807b705c7caf9", + "ms-f570fd75c1cff1a85840", + "ms-83abd317d0eb60469a80", + "ms-f366297e507fcb6b627c", + "ms-88b1a0f8f758f7956c0f", + "ms-1eeb1190382de4eee40d", + "ms-6b6853bb2c84be7fc792", + "ms-fd72329371815c33b865", + "ms-578e24198ed828a74485", + "ms-5a01e8f312fcf818553e", + "ms-1fb1b6108181bbfada71", + "ms-49219cb8ed7dad70c897", + "ms-c3f3b5f86027cd9e375b", + "ms-a3ed6faf3fdad2c34cd7", + "ms-1b95b3807e292e25258d", + "ms-6634347d1d05f493d81a", + "ms-9b99aa6b393ce43351bc", + "ms-9df7b29ad1c3d5426592", + "ms-a98fe2dc370ca1524636", + "ms-86c8f25a972d298f2856" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000002.jsonl", + "record_count": 1024, + "group_ids": [ + "ms-5aaddca43873344ddf48", + "ms-f18b1c243c3c697cd40e", + "ms-b1cc22717fa5ad6da556", + "ms-014b7d38b08167cdd091", + "ms-72f136a1bf66d0b43a2e", + "ms-932804639ad9815e5c8a", + "ms-fe650308e25459ee64c6", + "ms-3deb9fa600781abb51c1", + "ms-84412def388c83dd0991", + "ms-0af41ac7984478cb304a", + "ms-7fcf7fde0cffadefdb56", + "ms-a12a5f95b40662ed4193", + "ms-5e9b47c1797eeba08b83", + "ms-1186ede02e41add5835e", + "ms-191a10967dd259fb3b09", + "ms-8285407b65908629ef5a", + "ms-836fa7b4801fa1ab588e", + "ms-b7800060cd4a8c3cb6b9", + "ms-da7bdb580ffc39378e18", + "ms-7802f7a4062e6493e9cf", + "ms-998722dda6bc90ba326d", + "ms-ec03dc3e2e69be366a9d", + "ms-95168f33f972922323df", + "ms-f1d94d87fe8cd6d887f7", + "ms-175375383979526d80b5", + "ms-fc371f511ecd860f1e1b", + "ms-a699ee643e76a610b444", + "ms-6a7be6379d33f1502b6c", + "ms-b4fb36375d738708822f", + "ms-4fcc97af6584ee325bf8", + "ms-5ccadfdf06cd5996dc2f", + "ms-9a077df5edf6f622fb1c", + "ms-0f276ebd32a384e43fa1", + "ms-b25ff1a0ebf6afaefc31", + "ms-15002155823cc9950c82", + "ms-d8a04abc63f6a4ffd5d2", + "ms-58900ca359c45164ea9d", + "ms-44a2cdcbbff2f0fdb0dc", + "ms-59be0cff09de98a87306", + "ms-efebbea7a308c6c03f0d", + "ms-22ddad4e01561babbce3", + "ms-4a36af62b5ac1e2312bc", + "ms-669061498bda79d43fd1", + "ms-4c2abbb44ee8530c115d", + "ms-84be32c9b49ca2783058", + "ms-54e2197370a3eb37be5c", + "ms-d0e7d2bf05c758c92eda", + "ms-3f358662ceabe01f6958", + "ms-b96a1c7ef36ef25dafa5", + "ms-b6df140b235c8fe0eaad", + "ms-6feba052af2a272cb380", + "ms-858fd88dcedec86b2805", + "ms-8f30b57ca10057b4c506", + "ms-1daee8c9ff78ee3228d8", + "ms-5f36e400033c8dfec292", + "ms-1bfca22d4af29f0e14b1", + "ms-f10e624b9c7c18351952", + "ms-0495effeb9d7bd9dc1f8", + "ms-a5b32756c14377c25aa7", + "ms-98427ae2d94c1f85811a", + "ms-40751c940b0cfb6484fe", + "ms-505c40efde425751ecb1", + "ms-29568a9d15c89cfe75cb", + "ms-f64cfe539c837f0e3871" + ], + "format": "jsonl" + }, + { + "path": "shards/shard_000003.jsonl", + "record_count": 128, + "group_ids": [ + "ms-adecad10e03f8efac588", + "ms-441081ce83e867c501ed", + "ms-ed27a3c13e7a25351fed", + "ms-3962fbe3b35c19f64df1", + "ms-4210337243d89191da9a", + "ms-92ff01ffd41606915b35", + "ms-312664c1f9cdcbdd10eb", + "ms-ded3ce9f89726fff8e99" + ], + "format": "jsonl" + } + ], + "record_count": 3200, + "group_count": 200 + } +} + +================================================== +ORACLE CEILING BY HORIZON +================================================== + horizon groups oracle expert mean_reward_spread + 4 200 0.3950 0.3000 1.0210 + 8 200 0.6100 0.5350 1.3936 + 16 200 0.9550 0.9300 1.9116 + 32 200 0.9950 0.9950 1.9791 + +Baseline reference: horizon=4 PickCube oracle in full collection = 0.3740 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_0.out new file mode 100644 index 0000000000000000000000000000000000000000..04247d099283374c016ee7239fa0bfbb3eba6097 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_0.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_runs/lattice_field/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23899, + "pairwise_ranking_accuracy": 0.6615757981505502, + "top1_action_selection": 0.31, + "selected_success_rate": 0.245, + "oracle_success_rate": 0.335, + "ndcg_at_k": 0.887380610879892, + "potential_edge_mae": 0.29173505257992355, + "effect_prediction_mae": 0.021143401695563682, + "selection_regret": 0.21459916058927775, + "selected_candidate_type_counts": { + "expert": 34, + "near_miss": 25, + "no_op": 78, + "random_negative": 52, + "wrong_direction": 11 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23899, + "pairwise_ranking_accuracy": 0.6615757981505502, + "top1_action_selection": 0.31, + "selected_success_rate": 0.245, + "oracle_success_rate": 0.335, + "ndcg_at_k": 0.887380610879892, + "potential_edge_mae": 0.29173505257992355, + "effect_prediction_mae": 0.021143401695563682, + "selection_regret": 0.21459916058927775, + "selected_candidate_type_counts": { + "expert": 34, + "near_miss": 25, + "no_op": 78, + "random_negative": 52, + "wrong_direction": 11 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_1.out new file mode 100644 index 0000000000000000000000000000000000000000..3ad1fe16cbe1a061d1728a82cae8d790372c1681 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_1.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_runs/legacy/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23899, + "pairwise_ranking_accuracy": 0.6555085986861375, + "top1_action_selection": 0.3, + "selected_success_rate": 0.255, + "oracle_success_rate": 0.335, + "ndcg_at_k": 0.8824350898159404, + "potential_edge_mae": 0.3433309025025755, + "effect_prediction_mae": 0.022528056593881508, + "selection_regret": 0.21347508776932955, + "selected_candidate_type_counts": { + "expert": 25, + "near_miss": 23, + "no_op": 86, + "random_negative": 49, + "wrong_direction": 17 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23899, + "pairwise_ranking_accuracy": 0.6555085986861375, + "top1_action_selection": 0.3, + "selected_success_rate": 0.255, + "oracle_success_rate": 0.335, + "ndcg_at_k": 0.8824350898159404, + "potential_edge_mae": 0.3433309025025755, + "effect_prediction_mae": 0.022528056593881508, + "selection_regret": 0.21347508776932955, + "selected_candidate_type_counts": { + "expert": 25, + "near_miss": 23, + "no_op": 86, + "random_negative": 49, + "wrong_direction": 17 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_2.out new file mode 100644 index 0000000000000000000000000000000000000000..96092e6f9d3cd8e53ad0be14e8625fb1a35dc309 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_2.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_runs/lattice_field/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23872, + "pairwise_ranking_accuracy": 0.6646699061662198, + "top1_action_selection": 0.4, + "selected_success_rate": 0.325, + "oracle_success_rate": 0.385, + "ndcg_at_k": 0.9005816399884998, + "potential_edge_mae": 0.33418246218544523, + "effect_prediction_mae": 0.02302884298451726, + "selection_regret": 0.16521056510508061, + "selected_candidate_type_counts": { + "expert": 26, + "near_miss": 28, + "no_op": 96, + "random_negative": 37, + "wrong_direction": 13 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23872, + "pairwise_ranking_accuracy": 0.6646699061662198, + "top1_action_selection": 0.4, + "selected_success_rate": 0.325, + "oracle_success_rate": 0.385, + "ndcg_at_k": 0.9005816399884998, + "potential_edge_mae": 0.33418246218544523, + "effect_prediction_mae": 0.02302884298451726, + "selection_regret": 0.16521056510508061, + "selected_candidate_type_counts": { + "expert": 26, + "near_miss": 28, + "no_op": 96, + "random_negative": 37, + "wrong_direction": 13 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_3.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_3.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_3.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_3.out new file mode 100644 index 0000000000000000000000000000000000000000..98b95dbbf61cd167c8c4ecdc3e791aaaeeb91227 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_3.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_runs/legacy/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23872, + "pairwise_ranking_accuracy": 0.6629524128686327, + "top1_action_selection": 0.37, + "selected_success_rate": 0.3, + "oracle_success_rate": 0.385, + "ndcg_at_k": 0.8945619413555419, + "potential_edge_mae": 0.3936556764511921, + "effect_prediction_mae": 0.02149298046821694, + "selection_regret": 0.18942484308034183, + "selected_candidate_type_counts": { + "expert": 31, + "near_miss": 29, + "no_op": 95, + "random_negative": 36, + "wrong_direction": 9 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23872, + "pairwise_ranking_accuracy": 0.6629524128686327, + "top1_action_selection": 0.37, + "selected_success_rate": 0.3, + "oracle_success_rate": 0.385, + "ndcg_at_k": 0.8945619413555419, + "potential_edge_mae": 0.3936556764511921, + "effect_prediction_mae": 0.02149298046821694, + "selection_regret": 0.18942484308034183, + "selected_candidate_type_counts": { + "expert": 31, + "near_miss": 29, + "no_op": 95, + "random_negative": 36, + "wrong_direction": 9 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_4.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_4.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_4.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_4.out new file mode 100644 index 0000000000000000000000000000000000000000..8e323a831831c704ed7d1fa30e0e8dae65a778ee --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_4.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_runs/lattice_field/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23867, + "pairwise_ranking_accuracy": 0.6650186449909917, + "top1_action_selection": 0.375, + "selected_success_rate": 0.29, + "oracle_success_rate": 0.365, + "ndcg_at_k": 0.8965798161636873, + "potential_edge_mae": 0.32895855214506875, + "effect_prediction_mae": 0.021978936759423365, + "selection_regret": 0.18650633858516813, + "selected_candidate_type_counts": { + "expert": 27, + "near_miss": 19, + "no_op": 99, + "random_negative": 37, + "wrong_direction": 18 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23867, + "pairwise_ranking_accuracy": 0.6650186449909917, + "top1_action_selection": 0.375, + "selected_success_rate": 0.29, + "oracle_success_rate": 0.365, + "ndcg_at_k": 0.8965798161636873, + "potential_edge_mae": 0.32895855214506875, + "effect_prediction_mae": 0.021978936759423365, + "selection_regret": 0.18650633858516813, + "selected_candidate_type_counts": { + "expert": 27, + "near_miss": 19, + "no_op": 99, + "random_negative": 37, + "wrong_direction": 18 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_5.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_5.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_5.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_5.out new file mode 100644 index 0000000000000000000000000000000000000000..94e705d6c1db40070a062bd7a59ff904232a0661 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472076_5.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_runs/legacy/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_full_k16_n1000_seed0", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23867, + "pairwise_ranking_accuracy": 0.6596136925461935, + "top1_action_selection": 0.385, + "selected_success_rate": 0.305, + "oracle_success_rate": 0.365, + "ndcg_at_k": 0.8875200118844616, + "potential_edge_mae": 0.36442907329352325, + "effect_prediction_mae": 0.023041328211686852, + "selection_regret": 0.17533883452415466, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 24, + "no_op": 93, + "random_negative": 41, + "wrong_direction": 12 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 200, + "num_records": 3200, + "num_pairs": 23867, + "pairwise_ranking_accuracy": 0.6596136925461935, + "top1_action_selection": 0.385, + "selected_success_rate": 0.305, + "oracle_success_rate": 0.365, + "ndcg_at_k": 0.8875200118844616, + "potential_edge_mae": 0.36442907329352325, + "effect_prediction_mae": 0.023041328211686852, + "selection_regret": 0.17533883452415466, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 24, + "no_op": 93, + "random_negative": 41, + "wrong_direction": 12 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_0.out new file mode 100644 index 0000000000000000000000000000000000000000..72cce4055053e6c6abde5aa7bf0e3e3c83ea2ec8 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_0.out @@ -0,0 +1,148 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_runs/lattice_field/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.7267756790013827, + "top1_action_selection": 0.49142857142857144, + "selected_success_rate": 0.33, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9219798121086988, + "potential_edge_mae": 0.27683130122049343, + "effect_prediction_mae": 0.02825316913395154, + "selection_regret": 0.15994243508712055, + "selected_candidate_type_counts": { + "expert": 182, + "near_miss": 254, + "no_op": 102, + "random_negative": 70, + "wrong_direction": 92 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.7181295715778474, + "top1_action_selection": 0.45454545454545453, + "selected_success_rate": 0.2727272727272727, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9387705325670227, + "potential_edge_mae": 0.3014689148032278, + "effect_prediction_mae": 0.023445282636269802, + "selection_regret": 0.2526950595354793, + "selected_candidate_type_counts": { + "expert": 37, + "near_miss": 38, + "random_negative": 13, + "wrong_direction": 11 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.7212199312714777, + "top1_action_selection": 0.26804123711340205, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.822326949509932, + "potential_edge_mae": 0.08512493151315335, + "effect_prediction_mae": 0.04462663900767098, + "selection_regret": 0.10162699896393382, + "selected_candidate_type_counts": { + "expert": 38, + "near_miss": 38, + "random_negative": 7, + "wrong_direction": 14 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.7080691392339095, + "top1_action_selection": 0.47150259067357514, + "selected_success_rate": 0.3005181347150259, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9308703420913388, + "potential_edge_mae": 0.3057180306383437, + "effect_prediction_mae": 0.022037347952933075, + "selection_regret": 0.129435528436473, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 47, + "no_op": 86, + "random_negative": 13, + "wrong_direction": 17 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.7600193143408981, + "top1_action_selection": 0.6902654867256637, + "selected_success_rate": 0.5752212389380531, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9679661805602333, + "potential_edge_mae": 0.25264278458425066, + "effect_prediction_mae": 0.025198777808923947, + "selection_regret": 0.07488794229199401, + "selected_candidate_type_counts": { + "expert": 25, + "near_miss": 47, + "random_negative": 21, + "wrong_direction": 20 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.7854449802718106, + "top1_action_selection": 0.6886792452830188, + "selected_success_rate": 0.5754716981132075, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.9653338875596449, + "potential_edge_mae": 0.31860858554552124, + "effect_prediction_mae": 0.026880626256528945, + "selection_regret": 0.1376068433219532, + "selected_candidate_type_counts": { + "expert": 22, + "near_miss": 57, + "random_negative": 5, + "wrong_direction": 22 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6884330066855939, + "top1_action_selection": 0.33695652173913043, + "selected_success_rate": 0.21739130434782608, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.8838948384916702, + "potential_edge_mae": 0.3736698504901378, + "effect_prediction_mae": 0.03453624998385648, + "selection_regret": 0.31581907508813817, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 27, + "no_op": 16, + "random_negative": 11, + "wrong_direction": 8 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_1.out new file mode 100644 index 0000000000000000000000000000000000000000..cbadf1ecb4b5f790be40b733e371a00b2ebc7c3b --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_1.out @@ -0,0 +1,150 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_runs/legacy/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.7305686993365386, + "top1_action_selection": 0.47285714285714286, + "selected_success_rate": 0.32857142857142857, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9200115932185934, + "potential_edge_mae": 0.3473771676065187, + "effect_prediction_mae": 0.03161974825326225, + "selection_regret": 0.16560712372834263, + "selected_candidate_type_counts": { + "expert": 187, + "near_miss": 245, + "no_op": 84, + "random_negative": 105, + "wrong_direction": 79 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.7133402995471961, + "top1_action_selection": 0.45454545454545453, + "selected_success_rate": 0.29292929292929293, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9303906289105207, + "potential_edge_mae": 0.35698466563713455, + "effect_prediction_mae": 0.029232580411944473, + "selection_regret": 0.23697155804345102, + "selected_candidate_type_counts": { + "expert": 33, + "near_miss": 34, + "no_op": 1, + "random_negative": 24, + "wrong_direction": 7 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.7279209621993127, + "top1_action_selection": 0.28865979381443296, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.8343830353670179, + "potential_edge_mae": 0.13202205049492727, + "effect_prediction_mae": 0.04795266784421559, + "selection_regret": 0.08926858524095804, + "selected_candidate_type_counts": { + "expert": 39, + "near_miss": 41, + "random_negative": 7, + "wrong_direction": 10 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.7093285850777382, + "top1_action_selection": 0.42487046632124353, + "selected_success_rate": 0.27979274611398963, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9215904516112029, + "potential_edge_mae": 0.36973351007202054, + "effect_prediction_mae": 0.026187107692227908, + "selection_regret": 0.15870619171487235, + "selected_candidate_type_counts": { + "expert": 48, + "near_miss": 41, + "no_op": 72, + "random_negative": 15, + "wrong_direction": 17 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.7718010622887495, + "top1_action_selection": 0.6548672566371682, + "selected_success_rate": 0.5575221238938053, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9641565118105214, + "potential_edge_mae": 0.310532329296257, + "effect_prediction_mae": 0.02655674721201053, + "selection_regret": 0.1053895631854513, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 43, + "no_op": 4, + "random_negative": 28, + "wrong_direction": 15 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.7925471284524331, + "top1_action_selection": 0.6886792452830188, + "selected_success_rate": 0.5943396226415094, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.967944887146545, + "potential_edge_mae": 0.40910634793331363, + "effect_prediction_mae": 0.029170391771553898, + "selection_regret": 0.10493205923516795, + "selected_candidate_type_counts": { + "expert": 16, + "near_miss": 58, + "random_negative": 9, + "wrong_direction": 23 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6924626797325762, + "top1_action_selection": 0.31521739130434784, + "selected_success_rate": 0.22826086956521738, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.8863639948607378, + "potential_edge_mae": 0.4901677998311136, + "effect_prediction_mae": 0.03740547595946037, + "selection_regret": 0.32764821386207704, + "selected_candidate_type_counts": { + "expert": 28, + "near_miss": 28, + "no_op": 7, + "random_negative": 22, + "wrong_direction": 7 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_2.out new file mode 100644 index 0000000000000000000000000000000000000000..b689725c414432951239b37fb1519cb94b2b4b44 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_2.out @@ -0,0 +1,150 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_runs/lattice_field/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.7196818331267862, + "top1_action_selection": 0.51, + "selected_success_rate": 0.35285714285714287, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.9165283259198398, + "potential_edge_mae": 0.298633264797407, + "effect_prediction_mae": 0.027356990041783998, + "selection_regret": 0.17726301213451995, + "selected_candidate_type_counts": { + "expert": 258, + "near_miss": 212, + "no_op": 101, + "random_negative": 81, + "wrong_direction": 48 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.7243018867924528, + "top1_action_selection": 0.4824561403508772, + "selected_success_rate": 0.3157894736842105, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9336110185805782, + "potential_edge_mae": 0.33342147451554827, + "effect_prediction_mae": 0.02126174958858477, + "selection_regret": 0.2432761074681031, + "selected_candidate_type_counts": { + "expert": 50, + "near_miss": 47, + "random_negative": 10, + "wrong_direction": 7 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.7142972863507493, + "top1_action_selection": 0.3106796116504854, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.8370952877701392, + "potential_edge_mae": 0.108288570950025, + "effect_prediction_mae": 0.046000715564389276, + "selection_regret": 0.10724460123755598, + "selected_candidate_type_counts": { + "expert": 39, + "near_miss": 44, + "random_negative": 8, + "wrong_direction": 12 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.6958174904942965, + "top1_action_selection": 0.520618556701031, + "selected_success_rate": 0.32989690721649484, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.9189682418993871, + "potential_edge_mae": 0.32087995056145024, + "effect_prediction_mae": 0.020058423118357094, + "selection_regret": 0.15774122629429876, + "selected_candidate_type_counts": { + "expert": 44, + "near_miss": 43, + "no_op": 86, + "random_negative": 15, + "wrong_direction": 6 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.7669844455594581, + "top1_action_selection": 0.6822429906542056, + "selected_success_rate": 0.6261682242990654, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.9717720907035637, + "potential_edge_mae": 0.26553670954121844, + "effect_prediction_mae": 0.023601951321343274, + "selection_regret": 0.07229192917035004, + "selected_candidate_type_counts": { + "expert": 54, + "near_miss": 17, + "no_op": 1, + "random_negative": 26, + "wrong_direction": 9 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.7707006369426752, + "top1_action_selection": 0.7142857142857143, + "selected_success_rate": 0.6813186813186813, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.9635455276852435, + "potential_edge_mae": 0.3622914728707301, + "effect_prediction_mae": 0.027409732225014152, + "selection_regret": 0.09749631436316522, + "selected_candidate_type_counts": { + "expert": 49, + "near_miss": 28, + "no_op": 1, + "random_negative": 5, + "wrong_direction": 8 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.6824117810502918, + "top1_action_selection": 0.34065934065934067, + "selected_success_rate": 0.18681318681318682, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.8678600190358013, + "potential_edge_mae": 0.3999703663874083, + "effect_prediction_mae": 0.03381265235380982, + "selection_regret": 0.418629004241346, + "selected_candidate_type_counts": { + "expert": 22, + "near_miss": 33, + "no_op": 13, + "random_negative": 17, + "wrong_direction": 6 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_3.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_3.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_3.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_3.out new file mode 100644 index 0000000000000000000000000000000000000000..26bdeb035210d85c138394389fc00a0b9c139d2c --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_3.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_runs/legacy/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.7234249728116543, + "top1_action_selection": 0.48142857142857143, + "selected_success_rate": 0.3385714285714286, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.9168922133206452, + "potential_edge_mae": 0.3366447069198181, + "effect_prediction_mae": 0.02905778478219826, + "selection_regret": 0.2049702281405085, + "selected_candidate_type_counts": { + "expert": 257, + "near_miss": 221, + "no_op": 83, + "random_negative": 82, + "wrong_direction": 57 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.7270188679245283, + "top1_action_selection": 0.45614035087719296, + "selected_success_rate": 0.3333333333333333, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9319019080499203, + "potential_edge_mae": 0.33708215142111736, + "effect_prediction_mae": 0.02232253257417723, + "selection_regret": 0.22341479831620267, + "selected_candidate_type_counts": { + "expert": 70, + "near_miss": 27, + "random_negative": 12, + "wrong_direction": 5 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.718023491292021, + "top1_action_selection": 0.3592233009708738, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.8787119606403032, + "potential_edge_mae": 0.11237563860552449, + "effect_prediction_mae": 0.048337346061898274, + "selection_regret": 0.0747752184617939, + "selected_candidate_type_counts": { + "expert": 52, + "near_miss": 41, + "random_negative": 3, + "wrong_direction": 7 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.7018665744901487, + "top1_action_selection": 0.422680412371134, + "selected_success_rate": 0.27835051546391754, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.9032886442424984, + "potential_edge_mae": 0.38793419046281435, + "effect_prediction_mae": 0.02209240520466302, + "selection_regret": 0.24844991561678267, + "selected_candidate_type_counts": { + "expert": 55, + "near_miss": 47, + "no_op": 62, + "random_negative": 22, + "wrong_direction": 8 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.7755143000501756, + "top1_action_selection": 0.6915887850467289, + "selected_success_rate": 0.6261682242990654, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.9700792639009027, + "potential_edge_mae": 0.2632804717793799, + "effect_prediction_mae": 0.02422110703194757, + "selection_regret": 0.08150963199751399, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 38, + "no_op": 8, + "random_negative": 19, + "wrong_direction": 12 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.7681946329748356, + "top1_action_selection": 0.7252747252747253, + "selected_success_rate": 0.6703296703296703, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.9641646131179351, + "potential_edge_mae": 0.4132950959033417, + "effect_prediction_mae": 0.028995458073934072, + "selection_regret": 0.12230982361259041, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 40, + "no_op": 2, + "wrong_direction": 19 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.6836158192090396, + "top1_action_selection": 0.2857142857142857, + "selected_success_rate": 0.17582417582417584, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.8604938664465552, + "potential_edge_mae": 0.48231068167742663, + "effect_prediction_mae": 0.03627211255885782, + "selection_regret": 0.4643628240286649, + "selected_candidate_type_counts": { + "expert": 20, + "near_miss": 28, + "no_op": 11, + "random_negative": 26, + "wrong_direction": 6 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_4.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_4.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_4.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_4.out new file mode 100644 index 0000000000000000000000000000000000000000..63f87a9ffb41eba76730b949ebadeef31941cadd --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_4.out @@ -0,0 +1,150 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_runs/lattice_field/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.7302893615951953, + "top1_action_selection": 0.48857142857142855, + "selected_success_rate": 0.32571428571428573, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9217474676204584, + "potential_edge_mae": 0.27616904439675755, + "effect_prediction_mae": 0.02877629357645608, + "selection_regret": 0.16600375814263577, + "selected_candidate_type_counts": { + "expert": 283, + "near_miss": 199, + "no_op": 91, + "random_negative": 65, + "wrong_direction": 62 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.7299750944002571, + "top1_action_selection": 0.5092592592592593, + "selected_success_rate": 0.3425925925925926, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.9452853439259541, + "potential_edge_mae": 0.30239155289643765, + "effect_prediction_mae": 0.022848461594600317, + "selection_regret": 0.19351344514224264, + "selected_candidate_type_counts": { + "expert": 54, + "near_miss": 28, + "random_negative": 8, + "wrong_direction": 18 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.7173684210526315, + "top1_action_selection": 0.3473684210526316, + "selected_success_rate": 0.010526315789473684, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.8513038627555469, + "potential_edge_mae": 0.10955087535766724, + "effect_prediction_mae": 0.04784606857149718, + "selection_regret": 0.12442902955226601, + "selected_candidate_type_counts": { + "expert": 42, + "near_miss": 38, + "random_negative": 7, + "wrong_direction": 8 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.7166390091876796, + "top1_action_selection": 0.45893719806763283, + "selected_success_rate": 0.2898550724637681, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.9215128014658772, + "potential_edge_mae": 0.29607130588299674, + "effect_prediction_mae": 0.02225727680435897, + "selection_regret": 0.16363847687624503, + "selected_candidate_type_counts": { + "expert": 56, + "near_miss": 40, + "no_op": 85, + "random_negative": 13, + "wrong_direction": 13 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.7776530039303762, + "top1_action_selection": 0.6914893617021277, + "selected_success_rate": 0.5957446808510638, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9685101185994319, + "potential_edge_mae": 0.24212431482979083, + "effect_prediction_mae": 0.02617542419212665, + "selection_regret": 0.0637834741080061, + "selected_candidate_type_counts": { + "expert": 49, + "near_miss": 26, + "no_op": 1, + "random_negative": 9, + "wrong_direction": 9 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.780440593603097, + "top1_action_selection": 0.6237623762376238, + "selected_success_rate": 0.5445544554455446, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.9583917258024779, + "potential_edge_mae": 0.3069097818860849, + "effect_prediction_mae": 0.025394544865416663, + "selection_regret": 0.13039994682415879, + "selected_candidate_type_counts": { + "expert": 50, + "near_miss": 36, + "no_op": 1, + "random_negative": 6, + "wrong_direction": 8 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.6879709245634252, + "top1_action_selection": 0.3263157894736842, + "selected_success_rate": 0.2, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.8807145039021832, + "potential_edge_mae": 0.36933387419746, + "effect_prediction_mae": 0.036818946663588675, + "selection_regret": 0.3204548937709708, + "selected_candidate_type_counts": { + "expert": 32, + "near_miss": 31, + "no_op": 4, + "random_negative": 22, + "wrong_direction": 6 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_5.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_5.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_5.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_5.out new file mode 100644 index 0000000000000000000000000000000000000000..e76eda9517baf699482175dbf66398fb775ce0db --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472079_5.out @@ -0,0 +1,150 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_runs/legacy/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.7335435801429846, + "top1_action_selection": 0.48428571428571426, + "selected_success_rate": 0.32, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9199199390519947, + "potential_edge_mae": 0.35245996640188293, + "effect_prediction_mae": 0.02806534107017032, + "selection_regret": 0.17695491208553513, + "selected_candidate_type_counts": { + "expert": 210, + "near_miss": 256, + "no_op": 75, + "random_negative": 82, + "wrong_direction": 77 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.7222623925443882, + "top1_action_selection": 0.5370370370370371, + "selected_success_rate": 0.3425925925925926, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.941942687990566, + "potential_edge_mae": 0.3278701005681165, + "effect_prediction_mae": 0.02321271551078286, + "selection_regret": 0.19244134619280143, + "selected_candidate_type_counts": { + "expert": 63, + "near_miss": 26, + "no_op": 1, + "random_negative": 7, + "wrong_direction": 11 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.7212280701754386, + "top1_action_selection": 0.3894736842105263, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.8426844675362186, + "potential_edge_mae": 0.13940317919045347, + "effect_prediction_mae": 0.046717826065983625, + "selection_regret": 0.1277309744759757, + "selected_candidate_type_counts": { + "expert": 43, + "near_miss": 35, + "random_negative": 10, + "wrong_direction": 7 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.7245315092888654, + "top1_action_selection": 0.4444444444444444, + "selected_success_rate": 0.2753623188405797, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.9193021497936312, + "potential_edge_mae": 0.38007285884385605, + "effect_prediction_mae": 0.021739114482452318, + "selection_regret": 0.19338496094164642, + "selected_candidate_type_counts": { + "expert": 49, + "near_miss": 65, + "no_op": 58, + "random_negative": 20, + "wrong_direction": 15 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.7854014598540145, + "top1_action_selection": 0.6808510638297872, + "selected_success_rate": 0.574468085106383, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9636305792811495, + "potential_edge_mae": 0.30588512498736886, + "effect_prediction_mae": 0.024233296858044765, + "selection_regret": 0.10785726410277346, + "selected_candidate_type_counts": { + "expert": 18, + "near_miss": 43, + "no_op": 5, + "random_negative": 14, + "wrong_direction": 14 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.7844040925430915, + "top1_action_selection": 0.6039603960396039, + "selected_success_rate": 0.5544554455445545, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.9645376696127532, + "potential_edge_mae": 0.4631807451212897, + "effect_prediction_mae": 0.026470738568329145, + "selection_regret": 0.11449574686513089, + "selected_candidate_type_counts": { + "expert": 8, + "near_miss": 62, + "random_negative": 13, + "wrong_direction": 18 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.6883255030582395, + "top1_action_selection": 0.28421052631578947, + "selected_success_rate": 0.21052631578947367, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.8827789318618582, + "potential_edge_mae": 0.4647039993022252, + "effect_prediction_mae": 0.03420106105011803, + "selection_regret": 0.30754726654604864, + "selected_candidate_type_counts": { + "expert": 29, + "near_miss": 25, + "no_op": 11, + "random_negative": 18, + "wrong_direction": 12 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472239_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472239_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472239_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472239_0.out new file mode 100644 index 0000000000000000000000000000000000000000..ff33e71f4d1325af69232f203bc06be003e30f64 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472239_0.out @@ -0,0 +1,133 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/presuccess_visual_pilot_runs/lattice_field/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/presuccess_visual_pilot_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "objective": "lattice_field", + "observation_mode": "rgb", + "val_fraction": 0.2, + "num_groups": 38, + "num_records": 608, + "num_pairs": 4384, + "pairwise_ranking_accuracy": 0.6056113138686131, + "top1_action_selection": 0.13157894736842105, + "selected_success_rate": 0.07894736842105263, + "oracle_success_rate": 0.34210526315789475, + "ndcg_at_k": 0.7943015215810645, + "potential_edge_mae": 0.2696402561136295, + "effect_prediction_mae": 0.060138841934384325, + "selection_regret": 0.5001647336183325, + "selected_candidate_type_counts": { + "expert": 8, + "near_miss": 3, + "random_negative": 23, + "wrong_direction": 4 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 7, + "num_records": 112, + "num_pairs": 807, + "pairwise_ranking_accuracy": 0.5774473358116481, + "top1_action_selection": 0.14285714285714285, + "selected_success_rate": 0.14285714285714285, + "oracle_success_rate": 0.7142857142857143, + "ndcg_at_k": 0.8359582934940221, + "potential_edge_mae": 0.5022506303023524, + "effect_prediction_mae": 0.058530284100053724, + "selection_regret": 0.7914334663322994, + "selected_candidate_type_counts": { + "near_miss": 2, + "random_negative": 4, + "wrong_direction": 1 + } + }, + "PegInsertionSide-v1": { + "num_groups": 8, + "num_records": 128, + "num_pairs": 960, + "pairwise_ranking_accuracy": 0.6604166666666667, + "top1_action_selection": 0.0, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.125, + "ndcg_at_k": 0.663641392420514, + "potential_edge_mae": 0.09820295546246598, + "effect_prediction_mae": 0.06517949824597731, + "selection_regret": 0.3831326677791367, + "selected_candidate_type_counts": { + "near_miss": 1, + "random_negative": 6, + "wrong_direction": 1 + } + }, + "PickCube-v1": { + "num_groups": 5, + "num_records": 80, + "num_pairs": 596, + "pairwise_ranking_accuracy": 0.41946308724832215, + "top1_action_selection": 0.0, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.6, + "ndcg_at_k": 0.6266372614480676, + "potential_edge_mae": 0.4076470449104305, + "effect_prediction_mae": 0.05005647568366463, + "selection_regret": 1.155243566632271, + "selected_candidate_type_counts": { + "random_negative": 5 + } + }, + "PullCube-v1": { + "num_groups": 5, + "num_records": 80, + "num_pairs": 600, + "pairwise_ranking_accuracy": 0.8166666666666667, + "top1_action_selection": 0.2, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.9426349300596749, + "potential_edge_mae": 0.08360081081589063, + "effect_prediction_mae": 0.061103243744401424, + "selection_regret": 0.08706840723752976, + "selected_candidate_type_counts": { + "expert": 3, + "random_negative": 2 + } + }, + "PushCube-v1": { + "num_groups": 5, + "num_records": 80, + "num_pairs": 474, + "pairwise_ranking_accuracy": 0.689873417721519, + "top1_action_selection": 0.4, + "selected_success_rate": 0.4, + "oracle_success_rate": 0.4, + "ndcg_at_k": 0.9662285927168487, + "potential_edge_mae": 0.23150967812865092, + "effect_prediction_mae": 0.07041856261977696, + "selection_regret": 0.00587407648563385, + "selected_candidate_type_counts": { + "expert": 4, + "random_negative": 1 + } + }, + "StackCube-v1": { + "num_groups": 8, + "num_records": 128, + "num_pairs": 947, + "pairwise_ranking_accuracy": 0.515311510031679, + "top1_action_selection": 0.125, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.25, + "ndcg_at_k": 0.7931393381419027, + "potential_edge_mae": 0.2953094764678506, + "effect_prediction_mae": 0.055779576074899626, + "selection_regret": 0.5200292523950338, + "selected_candidate_type_counts": { + "expert": 1, + "random_negative": 5, + "wrong_direction": 2 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_0.out new file mode 100644 index 0000000000000000000000000000000000000000..aa81e95d5130dc55f4fd297860917b543eb29cb2 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_0.out @@ -0,0 +1,148 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_visual_runs/lattice_field/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "rgb", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.6876149640360781, + "top1_action_selection": 0.38142857142857145, + "selected_success_rate": 0.28, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.8985466374090306, + "potential_edge_mae": 0.34977509044518623, + "effect_prediction_mae": 0.040458501929328795, + "selection_regret": 0.2420764511442394, + "selected_candidate_type_counts": { + "expert": 158, + "near_miss": 161, + "no_op": 183, + "random_negative": 90, + "wrong_direction": 108 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.6583942877046325, + "top1_action_selection": 0.32323232323232326, + "selected_success_rate": 0.15151515151515152, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9064678040935544, + "potential_edge_mae": 0.3696041995304588, + "effect_prediction_mae": 0.03342795746008425, + "selection_regret": 0.43381399653776725, + "selected_candidate_type_counts": { + "expert": 37, + "near_miss": 13, + "random_negative": 39, + "wrong_direction": 10 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.7121993127147767, + "top1_action_selection": 0.18556701030927836, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.7918166680326898, + "potential_edge_mae": 0.0986731870178, + "effect_prediction_mae": 0.06569948055287518, + "selection_regret": 0.125902115793109, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 40, + "random_negative": 14, + "wrong_direction": 20 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.660861634673847, + "top1_action_selection": 0.37305699481865284, + "selected_success_rate": 0.2849740932642487, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9144579112140858, + "potential_edge_mae": 0.3573329848892813, + "effect_prediction_mae": 0.0333590899347704, + "selection_regret": 0.16302117574091401, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 23, + "no_op": 127, + "random_negative": 7, + "wrong_direction": 13 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.7288266537904394, + "top1_action_selection": 0.5929203539823009, + "selected_success_rate": 0.5309734513274337, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9523956521515048, + "potential_edge_mae": 0.45810259658823493, + "effect_prediction_mae": 0.038474365835933784, + "selection_regret": 0.1710189118717624, + "selected_candidate_type_counts": { + "expert": 18, + "near_miss": 43, + "no_op": 1, + "random_negative": 17, + "wrong_direction": 34 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.7486190267426567, + "top1_action_selection": 0.5943396226415094, + "selected_success_rate": 0.5566037735849056, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.9446470321767262, + "potential_edge_mae": 0.43610623750520455, + "effect_prediction_mae": 0.03415146882173954, + "selection_regret": 0.18311742847820497, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 40, + "random_negative": 13, + "wrong_direction": 30 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6457551057789175, + "top1_action_selection": 0.16304347826086957, + "selected_success_rate": 0.07608695652173914, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.8499178892795408, + "potential_edge_mae": 0.3877586540827898, + "effect_prediction_mae": 0.046008375612796824, + "selection_regret": 0.47929079927828, + "selected_candidate_type_counts": { + "expert": 34, + "near_miss": 2, + "no_op": 55, + "wrong_direction": 1 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_1.out new file mode 100644 index 0000000000000000000000000000000000000000..f4002db34d198df714dcad0f4fd0d0d82778dce7 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_1.out @@ -0,0 +1,150 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_visual_runs/lattice_field/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "rgb", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.6902046081084499, + "top1_action_selection": 0.42857142857142855, + "selected_success_rate": 0.31285714285714283, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.895994005532324, + "potential_edge_mae": 0.3406955465995907, + "effect_prediction_mae": 0.03494471385286317, + "selection_regret": 0.2466639896313427, + "selected_candidate_type_counts": { + "expert": 180, + "near_miss": 184, + "no_op": 199, + "random_negative": 46, + "wrong_direction": 91 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.6802264150943397, + "top1_action_selection": 0.43859649122807015, + "selected_success_rate": 0.2543859649122807, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9158064026799703, + "potential_edge_mae": 0.3441396390458323, + "effect_prediction_mae": 0.02942456246340551, + "selection_regret": 0.3249104005987184, + "selected_candidate_type_counts": { + "expert": 54, + "near_miss": 19, + "no_op": 7, + "random_negative": 19, + "wrong_direction": 15 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.6999594977723774, + "top1_action_selection": 0.18446601941747573, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.806385851942739, + "potential_edge_mae": 0.09402290839147362, + "effect_prediction_mae": 0.06082866554521345, + "selection_regret": 0.13595780581337608, + "selected_candidate_type_counts": { + "expert": 22, + "near_miss": 43, + "random_negative": 13, + "wrong_direction": 25 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.666220186657449, + "top1_action_selection": 0.41237113402061853, + "selected_success_rate": 0.29381443298969073, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.901316533146065, + "potential_edge_mae": 0.356002508877605, + "effect_prediction_mae": 0.02656289681302024, + "selection_regret": 0.22401062575812192, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 24, + "no_op": 130, + "random_negative": 5, + "wrong_direction": 5 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.7508278976417461, + "top1_action_selection": 0.7009345794392523, + "selected_success_rate": 0.6261682242990654, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.971907412703517, + "potential_edge_mae": 0.4538130835240741, + "effect_prediction_mae": 0.03287846899186865, + "selection_regret": 0.08285126298944527, + "selected_candidate_type_counts": { + "expert": 24, + "near_miss": 55, + "random_negative": 6, + "wrong_direction": 22 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.7591103685914169, + "top1_action_selection": 0.6593406593406593, + "selected_success_rate": 0.6263736263736264, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.9428228003381507, + "potential_edge_mae": 0.43331303451666536, + "effect_prediction_mae": 0.03087237096843633, + "selection_regret": 0.209095360322313, + "selected_candidate_type_counts": { + "expert": 24, + "near_miss": 41, + "no_op": 3, + "random_negative": 1, + "wrong_direction": 22 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.6256367509493378, + "top1_action_selection": 0.17582417582417584, + "selected_success_rate": 0.08791208791208792, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.8251621516112483, + "potential_edge_mae": 0.3991433766977516, + "effect_prediction_mae": 0.036933660472918524, + "selection_regret": 0.5524233928421042, + "selected_candidate_type_counts": { + "expert": 26, + "near_miss": 2, + "no_op": 59, + "random_negative": 2, + "wrong_direction": 2 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_2.out new file mode 100644 index 0000000000000000000000000000000000000000..67d82ceeda36f24eabbff4ccd5742087598f21ce --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472241_2.out @@ -0,0 +1,148 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_visual_runs/lattice_field/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "rgb", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.6918418374398472, + "top1_action_selection": 0.4114285714285714, + "selected_success_rate": 0.2957142857142857, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9018989267581675, + "potential_edge_mae": 0.31632311496458676, + "effect_prediction_mae": 0.03151749156779107, + "selection_regret": 0.22155316144901527, + "selected_candidate_type_counts": { + "expert": 218, + "near_miss": 156, + "no_op": 182, + "random_negative": 71, + "wrong_direction": 73 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.6599180525427814, + "top1_action_selection": 0.48148148148148145, + "selected_success_rate": 0.32407407407407407, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.9298734488750214, + "potential_edge_mae": 0.3309608111230592, + "effect_prediction_mae": 0.025884695882992078, + "selection_regret": 0.23102083167544119, + "selected_candidate_type_counts": { + "expert": 53, + "near_miss": 17, + "random_negative": 31, + "wrong_direction": 7 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.715701754385965, + "top1_action_selection": 0.23157894736842105, + "selected_success_rate": 0.010526315789473684, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.8145973298839791, + "potential_edge_mae": 0.08705898234702385, + "effect_prediction_mae": 0.05698277360172509, + "selection_regret": 0.16275805773176744, + "selected_candidate_type_counts": { + "expert": 26, + "near_miss": 45, + "random_negative": 14, + "wrong_direction": 10 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.6667341239324888, + "top1_action_selection": 0.3864734299516908, + "selected_success_rate": 0.28019323671497587, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.9062585353013142, + "potential_edge_mae": 0.3336460019563247, + "effect_prediction_mae": 0.023144324515413463, + "selection_regret": 0.19425346173238064, + "selected_candidate_type_counts": { + "expert": 40, + "near_miss": 23, + "no_op": 128, + "random_negative": 6, + "wrong_direction": 10 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.7541830432341381, + "top1_action_selection": 0.6595744680851063, + "selected_success_rate": 0.5851063829787234, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9631505972624793, + "potential_edge_mae": 0.39893080756373617, + "effect_prediction_mae": 0.031782802745840714, + "selection_regret": 0.091704989525866, + "selected_candidate_type_counts": { + "expert": 35, + "near_miss": 29, + "no_op": 1, + "random_negative": 9, + "wrong_direction": 20 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.7688266199649737, + "top1_action_selection": 0.5346534653465347, + "selected_success_rate": 0.48514851485148514, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.9414191600537402, + "potential_edge_mae": 0.38307647155566044, + "effect_prediction_mae": 0.026253046633576867, + "selection_regret": 0.24340854998272243, + "selected_candidate_type_counts": { + "expert": 25, + "near_miss": 40, + "random_negative": 10, + "wrong_direction": 26 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.6346955057175783, + "top1_action_selection": 0.18947368421052632, + "selected_success_rate": 0.09473684210526316, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.8452753873448833, + "potential_edge_mae": 0.3645094876796702, + "effect_prediction_mae": 0.03603491691721942, + "selection_regret": 0.4343152480690103, + "selected_candidate_type_counts": { + "expert": 39, + "near_miss": 2, + "no_op": 53, + "random_negative": 1 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_0.out new file mode 100644 index 0000000000000000000000000000000000000000..a16080971bdb77c0295eddcf2ecd7f927e5bd19b --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_0.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/cross_state_negatives/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.7152443897550393, + "top1_action_selection": 0.48428571428571426, + "selected_success_rate": 0.3342857142857143, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9185088820431723, + "potential_edge_mae": 0.3472250585184493, + "effect_prediction_mae": 0.029919107714158153, + "selection_regret": 0.15671036520723386, + "selected_candidate_type_counts": { + "expert": 300, + "near_miss": 160, + "no_op": 82, + "random_negative": 97, + "wrong_direction": 61 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.6911354928596308, + "top1_action_selection": 0.5050505050505051, + "selected_success_rate": 0.29292929292929293, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9328482432940381, + "potential_edge_mae": 0.31992707879453486, + "effect_prediction_mae": 0.027488574500650165, + "selection_regret": 0.2256022911481183, + "selected_candidate_type_counts": { + "expert": 59, + "near_miss": 16, + "random_negative": 12, + "wrong_direction": 12 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.7219931271477663, + "top1_action_selection": 0.24742268041237114, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.8294194591785411, + "potential_edge_mae": 0.11691454853978414, + "effect_prediction_mae": 0.04831190905275854, + "selection_regret": 0.10888660393928937, + "selected_candidate_type_counts": { + "expert": 37, + "near_miss": 39, + "random_negative": 8, + "wrong_direction": 13 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.6903934682532789, + "top1_action_selection": 0.48186528497409326, + "selected_success_rate": 0.3005181347150259, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9238372501498606, + "potential_edge_mae": 0.34314935225215204, + "effect_prediction_mae": 0.024086519261363574, + "selection_regret": 0.13228404326026613, + "selected_candidate_type_counts": { + "expert": 57, + "near_miss": 27, + "no_op": 74, + "random_negative": 26, + "wrong_direction": 9 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.7497827136648962, + "top1_action_selection": 0.6902654867256637, + "selected_success_rate": 0.584070796460177, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.962629944870331, + "potential_edge_mae": 0.3993109941227116, + "effect_prediction_mae": 0.023928733338642698, + "selection_regret": 0.05316065897983787, + "selected_candidate_type_counts": { + "expert": 73, + "near_miss": 15, + "no_op": 2, + "random_negative": 16, + "wrong_direction": 7 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.778342832091188, + "top1_action_selection": 0.6132075471698113, + "selected_success_rate": 0.5660377358490566, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.9571767039719971, + "potential_edge_mae": 0.4500271225105587, + "effect_prediction_mae": 0.02753802936664618, + "selection_regret": 0.1550106352792596, + "selected_candidate_type_counts": { + "expert": 44, + "near_miss": 32, + "random_negative": 14, + "wrong_direction": 16 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6871508379888268, + "top1_action_selection": 0.31521739130434784, + "selected_success_rate": 0.22826086956521738, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.8870875014509846, + "potential_edge_mae": 0.47327544269069166, + "effect_prediction_mae": 0.03547907770276591, + "selection_regret": 0.31338615310580836, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 31, + "no_op": 6, + "random_negative": 21, + "wrong_direction": 4 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_1.out new file mode 100644 index 0000000000000000000000000000000000000000..7cac3b954b1f8c53b3e7db81f7b17a8788b80785 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_1.out @@ -0,0 +1,148 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/cross_state_negatives/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.7103998583676876, + "top1_action_selection": 0.47285714285714286, + "selected_success_rate": 0.3242857142857143, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.9084673345781895, + "potential_edge_mae": 0.34353995461818915, + "effect_prediction_mae": 0.029032493453404293, + "selection_regret": 0.2313000690750362, + "selected_candidate_type_counts": { + "expert": 289, + "near_miss": 202, + "no_op": 53, + "random_negative": 100, + "wrong_direction": 56 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.6939622641509434, + "top1_action_selection": 0.4473684210526316, + "selected_success_rate": 0.3333333333333333, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9328733567308684, + "potential_edge_mae": 0.362166962258096, + "effect_prediction_mae": 0.024991515212529968, + "selection_regret": 0.22121991622343398, + "selected_candidate_type_counts": { + "expert": 64, + "near_miss": 23, + "random_negative": 19, + "wrong_direction": 8 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.7144592952612394, + "top1_action_selection": 0.30097087378640774, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.8566402369951533, + "potential_edge_mae": 0.18947598445920877, + "effect_prediction_mae": 0.0468832545648321, + "selection_regret": 0.10488008405731819, + "selected_candidate_type_counts": { + "expert": 46, + "near_miss": 44, + "random_negative": 5, + "wrong_direction": 8 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.6871327341859661, + "top1_action_selection": 0.4175257731958763, + "selected_success_rate": 0.25257731958762886, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.8942084152118377, + "potential_edge_mae": 0.3651757099546751, + "effect_prediction_mae": 0.022930366576426445, + "selection_regret": 0.2928056247225127, + "selected_candidate_type_counts": { + "expert": 64, + "near_miss": 38, + "no_op": 47, + "random_negative": 35, + "wrong_direction": 10 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.7621675865529353, + "top1_action_selection": 0.6728971962616822, + "selected_success_rate": 0.6074766355140186, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.9611378797144668, + "potential_edge_mae": 0.26034245024833824, + "effect_prediction_mae": 0.023844961560380868, + "selection_regret": 0.10869811767729644, + "selected_candidate_type_counts": { + "expert": 57, + "near_miss": 28, + "random_negative": 16, + "wrong_direction": 6 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.7680902161428422, + "top1_action_selection": 0.7692307692307693, + "selected_success_rate": 0.6373626373626373, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.954978613995963, + "potential_edge_mae": 0.3708980869047886, + "effect_prediction_mae": 0.027259764123321668, + "selection_regret": 0.17682497219724969, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 41, + "random_negative": 5, + "wrong_direction": 15 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.6768546818560711, + "top1_action_selection": 0.2857142857142857, + "selected_success_rate": 0.17582417582417584, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.8585097535465316, + "potential_edge_mae": 0.5029757457650833, + "effect_prediction_mae": 0.03477139477985844, + "selection_regret": 0.45453040645672727, + "selected_candidate_type_counts": { + "expert": 28, + "near_miss": 28, + "no_op": 6, + "random_negative": 20, + "wrong_direction": 9 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_2.out new file mode 100644 index 0000000000000000000000000000000000000000..93648ea223967300773e68a502a9a894116bb47f --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472348_2.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/cross_state_negatives/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.7261807536217316, + "top1_action_selection": 0.4785714285714286, + "selected_success_rate": 0.30428571428571427, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9183413008424214, + "potential_edge_mae": 0.33024982656671925, + "effect_prediction_mae": 0.027419047612600354, + "selection_regret": 0.200797276883885, + "selected_candidate_type_counts": { + "expert": 265, + "near_miss": 208, + "no_op": 65, + "random_negative": 71, + "wrong_direction": 91 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.7157548003534988, + "top1_action_selection": 0.5185185185185185, + "selected_success_rate": 0.3611111111111111, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.9424516077548494, + "potential_edge_mae": 0.3359142375634532, + "effect_prediction_mae": 0.022911640185802388, + "selection_regret": 0.1717366395449197, + "selected_candidate_type_counts": { + "expert": 46, + "near_miss": 30, + "no_op": 2, + "random_negative": 10, + "wrong_direction": 20 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.7287719298245614, + "top1_action_selection": 0.43157894736842106, + "selected_success_rate": 0.010526315789473684, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.8848245723663795, + "potential_edge_mae": 0.1605337269697798, + "effect_prediction_mae": 0.04355869740493996, + "selection_regret": 0.08814678712576432, + "selected_candidate_type_counts": { + "expert": 53, + "near_miss": 35, + "random_negative": 1, + "wrong_direction": 6 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.7112559193750759, + "top1_action_selection": 0.43478260869565216, + "selected_success_rate": 0.2608695652173913, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.9108690984852362, + "potential_edge_mae": 0.32502656793058615, + "effect_prediction_mae": 0.019815216109519744, + "selection_regret": 0.2096595237220543, + "selected_candidate_type_counts": { + "expert": 74, + "near_miss": 45, + "no_op": 54, + "random_negative": 16, + "wrong_direction": 18 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.7679955081414935, + "top1_action_selection": 0.6595744680851063, + "selected_success_rate": 0.5531914893617021, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9589034375383167, + "potential_edge_mae": 0.2495619835822536, + "effect_prediction_mae": 0.02491162507544353, + "selection_regret": 0.13834501676102903, + "selected_candidate_type_counts": { + "expert": 38, + "near_miss": 26, + "random_negative": 14, + "wrong_direction": 16 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.7782284081482165, + "top1_action_selection": 0.5643564356435643, + "selected_success_rate": 0.49504950495049505, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.9429008594577398, + "potential_edge_mae": 0.42043994760048237, + "effect_prediction_mae": 0.02597086552985418, + "selection_regret": 0.22174539689970488, + "selected_candidate_type_counts": { + "expert": 27, + "near_miss": 45, + "random_negative": 8, + "wrong_direction": 21 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.684691073486393, + "top1_action_selection": 0.30526315789473685, + "selected_success_rate": 0.17894736842105263, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.8744840974430287, + "potential_edge_mae": 0.4839030737551843, + "effect_prediction_mae": 0.036992631842175515, + "selection_regret": 0.36669846208472, + "selected_candidate_type_counts": { + "expert": 27, + "near_miss": 27, + "no_op": 9, + "random_negative": 22, + "wrong_direction": 10 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_0.out new file mode 100644 index 0000000000000000000000000000000000000000..9d86075e09d4cb1ec6c8529bb9a355fa93ef2d45 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_0.out @@ -0,0 +1,151 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/world_model_auxiliary/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.639257126184526, + "top1_action_selection": 0.38142857142857145, + "selected_success_rate": 0.3271428571428571, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.8620126653256059, + "potential_edge_mae": 3.010293736906071, + "effect_prediction_mae": 0.028171759370757547, + "selection_regret": 0.22230409624063344, + "selected_candidate_type_counts": { + "expert": 155, + "near_miss": 136, + "no_op": 101, + "random_negative": 239, + "wrong_direction": 69 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.6510797631487286, + "top1_action_selection": 0.31313131313131315, + "selected_success_rate": 0.23232323232323232, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9049691873566024, + "potential_edge_mae": 1.754103031353613, + "effect_prediction_mae": 0.023062264798715714, + "selection_regret": 0.36331993762892906, + "selected_candidate_type_counts": { + "expert": 37, + "near_miss": 12, + "no_op": 1, + "random_negative": 37, + "wrong_direction": 12 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.5794673539518901, + "top1_action_selection": 0.07216494845360824, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.6093366211649047, + "potential_edge_mae": 1.6359923299637142, + "effect_prediction_mae": 0.04594519289790017, + "selection_regret": 0.21462733502146913, + "selected_candidate_type_counts": { + "expert": 9, + "near_miss": 13, + "random_negative": 63, + "wrong_direction": 12 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.635542430296187, + "top1_action_selection": 0.38341968911917096, + "selected_success_rate": 0.31088082901554404, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.8865718025186717, + "potential_edge_mae": 3.428676539097228, + "effect_prediction_mae": 0.02190785402457042, + "selection_regret": 0.1814813739065216, + "selected_candidate_type_counts": { + "expert": 17, + "near_miss": 26, + "no_op": 85, + "random_negative": 53, + "wrong_direction": 12 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.6455818445195558, + "top1_action_selection": 0.5929203539823009, + "selected_success_rate": 0.5752212389380531, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9321904814354769, + "potential_edge_mae": 4.160975686269342, + "effect_prediction_mae": 0.024251973848653567, + "selection_regret": 0.11068810272005807, + "selected_candidate_type_counts": { + "expert": 39, + "near_miss": 22, + "no_op": 7, + "random_negative": 34, + "wrong_direction": 11 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.6931170539237177, + "top1_action_selection": 0.6132075471698113, + "selected_success_rate": 0.5754716981132075, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.9346732805475122, + "potential_edge_mae": 4.142977968057563, + "effect_prediction_mae": 0.027064214005342074, + "selection_regret": 0.15622357250946872, + "selected_candidate_type_counts": { + "expert": 34, + "near_miss": 40, + "no_op": 3, + "random_negative": 16, + "wrong_direction": 13 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6361388405531642, + "top1_action_selection": 0.25, + "selected_success_rate": 0.21739130434782608, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.860761020655181, + "potential_edge_mae": 2.6399031145212764, + "effect_prediction_mae": 0.03416182692559714, + "selection_regret": 0.3775215498779131, + "selected_candidate_type_counts": { + "expert": 19, + "near_miss": 23, + "no_op": 5, + "random_negative": 36, + "wrong_direction": 9 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_1.out new file mode 100644 index 0000000000000000000000000000000000000000..1cd970350cd4606fe8fc311d77dc3e207072d66a --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_1.out @@ -0,0 +1,148 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/world_model_auxiliary/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.654480386453881, + "top1_action_selection": 0.43142857142857144, + "selected_success_rate": 0.3357142857142857, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.8824939506869878, + "potential_edge_mae": 1.8087699200475915, + "effect_prediction_mae": 0.030294617757371408, + "selection_regret": 0.23301293948544688, + "selected_candidate_type_counts": { + "expert": 226, + "near_miss": 146, + "no_op": 108, + "random_negative": 166, + "wrong_direction": 54 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.633811320754717, + "top1_action_selection": 0.3684210526315789, + "selected_success_rate": 0.30701754385964913, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9107777744627722, + "potential_edge_mae": 1.117894548922215, + "effect_prediction_mae": 0.025527189812147865, + "selection_regret": 0.2875450648236693, + "selected_candidate_type_counts": { + "expert": 53, + "near_miss": 19, + "random_negative": 28, + "wrong_direction": 14 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.6660996354799514, + "top1_action_selection": 0.2524271844660194, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.7637269879186684, + "potential_edge_mae": 0.5256608157076595, + "effect_prediction_mae": 0.0469585885391905, + "selection_regret": 0.12224692168721288, + "selected_candidate_type_counts": { + "expert": 29, + "near_miss": 31, + "random_negative": 32, + "wrong_direction": 11 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.6371845834773592, + "top1_action_selection": 0.4020618556701031, + "selected_success_rate": 0.29896907216494845, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.8702550361163626, + "potential_edge_mae": 1.7744521881173458, + "effect_prediction_mae": 0.02382377966811799, + "selection_regret": 0.2568492882491387, + "selected_candidate_type_counts": { + "expert": 31, + "near_miss": 15, + "no_op": 103, + "random_negative": 41, + "wrong_direction": 4 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.6817862518815856, + "top1_action_selection": 0.6822429906542056, + "selected_success_rate": 0.6355140186915887, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.9643290342349664, + "potential_edge_mae": 3.387440164479751, + "effect_prediction_mae": 0.025594694488352613, + "selection_regret": 0.06274344347347723, + "selected_candidate_type_counts": { + "expert": 47, + "near_miss": 27, + "random_negative": 25, + "wrong_direction": 8 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.7040826981309387, + "top1_action_selection": 0.7142857142857143, + "selected_success_rate": 0.6593406593406593, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.9520944322151952, + "potential_edge_mae": 2.9163187534990906, + "effect_prediction_mae": 0.03008614260033732, + "selection_regret": 0.13793349659049903, + "selected_candidate_type_counts": { + "expert": 46, + "near_miss": 34, + "random_negative": 5, + "wrong_direction": 6 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.6344354913401871, + "top1_action_selection": 0.1978021978021978, + "selected_success_rate": 0.14285714285714285, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.8417576091781823, + "potential_edge_mae": 1.7578192841811258, + "effect_prediction_mae": 0.03693531461032085, + "selection_regret": 0.5345409753230902, + "selected_candidate_type_counts": { + "expert": 20, + "near_miss": 20, + "no_op": 5, + "random_negative": 35, + "wrong_direction": 11 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_2.out new file mode 100644 index 0000000000000000000000000000000000000000..b821ef3f52eeb0df30e6a670000ee19e646ba91b --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472350_2.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/world_model_auxiliary/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.6545628164696127, + "top1_action_selection": 0.39285714285714285, + "selected_success_rate": 0.3242857142857143, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.8636000687589208, + "potential_edge_mae": 2.7647054570165848, + "effect_prediction_mae": 0.02688878030028882, + "selection_regret": 0.22338697789207798, + "selected_candidate_type_counts": { + "expert": 192, + "near_miss": 139, + "no_op": 74, + "random_negative": 227, + "wrong_direction": 68 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.6615248654294208, + "top1_action_selection": 0.4537037037037037, + "selected_success_rate": 0.3611111111111111, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.9257694047334818, + "potential_edge_mae": 2.1072498279380287, + "effect_prediction_mae": 0.021013860115036122, + "selection_regret": 0.1998769235279825, + "selected_candidate_type_counts": { + "expert": 59, + "near_miss": 17, + "random_negative": 24, + "wrong_direction": 8 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.588421052631579, + "top1_action_selection": 0.06315789473684211, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.5995647100728371, + "potential_edge_mae": 0.8699140726775569, + "effect_prediction_mae": 0.044278606111163916, + "selection_regret": 0.2495410523306897, + "selected_candidate_type_counts": { + "expert": 6, + "near_miss": 13, + "random_negative": 63, + "wrong_direction": 13 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.6384425466466993, + "top1_action_selection": 0.34782608695652173, + "selected_success_rate": 0.2898550724637681, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.8680188261083206, + "potential_edge_mae": 2.867641963060675, + "effect_prediction_mae": 0.020919604706411313, + "selection_regret": 0.24153942585999263, + "selected_candidate_type_counts": { + "expert": 25, + "near_miss": 31, + "no_op": 72, + "random_negative": 66, + "wrong_direction": 13 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.7112857944974733, + "top1_action_selection": 0.6702127659574468, + "selected_success_rate": 0.5851063829787234, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9653282456175946, + "potential_edge_mae": 4.288109685526821, + "effect_prediction_mae": 0.022766907694809227, + "selection_regret": 0.09643860208861371, + "selected_candidate_type_counts": { + "expert": 42, + "near_miss": 19, + "no_op": 1, + "random_negative": 19, + "wrong_direction": 13 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.719697667987833, + "top1_action_selection": 0.6138613861386139, + "selected_success_rate": 0.5346534653465347, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.9483170883269146, + "potential_edge_mae": 3.9727023748135153, + "effect_prediction_mae": 0.025117852977568116, + "selection_regret": 0.15851659497412124, + "selected_candidate_type_counts": { + "expert": 37, + "near_miss": 39, + "random_negative": 14, + "wrong_direction": 11 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.6416097863664568, + "top1_action_selection": 0.24210526315789474, + "selected_success_rate": 0.2, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.8566055468380472, + "potential_edge_mae": 2.8151661010828755, + "effect_prediction_mae": 0.03514559040993799, + "selection_regret": 0.37898643126613213, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 20, + "no_op": 1, + "random_negative": 41, + "wrong_direction": 10 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_0.out new file mode 100644 index 0000000000000000000000000000000000000000..a8732f37683f5fa3e6ba57728c9af116c73a6a8f --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_0.out @@ -0,0 +1,150 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/no_effect_head/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.7321797815524743, + "top1_action_selection": 0.4942857142857143, + "selected_success_rate": 0.3271428571428571, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9235393631056734, + "potential_edge_mae": 0.3319768771296898, + "effect_prediction_mae": 0.3830362212853601, + "selection_regret": 0.16012156013000226, + "selected_candidate_type_counts": { + "expert": 292, + "near_miss": 200, + "no_op": 82, + "random_negative": 76, + "wrong_direction": 50 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.7291013584117032, + "top1_action_selection": 0.46464646464646464, + "selected_success_rate": 0.2727272727272727, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9340631526062874, + "potential_edge_mae": 0.30572888211456484, + "effect_prediction_mae": 0.34745801884522876, + "selection_regret": 0.2559794464496651, + "selected_candidate_type_counts": { + "expert": 49, + "near_miss": 29, + "no_op": 1, + "random_negative": 12, + "wrong_direction": 8 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.7254295532646048, + "top1_action_selection": 0.36082474226804123, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.8559443190386129, + "potential_edge_mae": 0.11767020845079819, + "effect_prediction_mae": 0.5554140499412807, + "selection_regret": 0.07077925576415565, + "selected_candidate_type_counts": { + "expert": 51, + "near_miss": 33, + "random_negative": 5, + "wrong_direction": 8 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.7104143142534526, + "top1_action_selection": 0.48186528497409326, + "selected_success_rate": 0.29533678756476683, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9238346163236946, + "potential_edge_mae": 0.36188142478542396, + "effect_prediction_mae": 0.38911794108179987, + "selection_regret": 0.13646466103551302, + "selected_candidate_type_counts": { + "expert": 52, + "near_miss": 43, + "no_op": 68, + "random_negative": 15, + "wrong_direction": 15 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.7694833413809754, + "top1_action_selection": 0.6814159292035398, + "selected_success_rate": 0.5752212389380531, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9680933391735901, + "potential_edge_mae": 0.2804758451623597, + "effect_prediction_mae": 0.3574510192365685, + "selection_regret": 0.07042197012795812, + "selected_candidate_type_counts": { + "expert": 63, + "near_miss": 20, + "no_op": 6, + "random_negative": 18, + "wrong_direction": 6 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.7871109162647961, + "top1_action_selection": 0.6132075471698113, + "selected_success_rate": 0.5660377358490566, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.9558414003537609, + "potential_edge_mae": 0.41197655312749626, + "effect_prediction_mae": 0.28947860313446044, + "selection_regret": 0.15619432813716386, + "selected_candidate_type_counts": { + "expert": 48, + "near_miss": 40, + "random_negative": 8, + "wrong_direction": 10 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6957596849528345, + "top1_action_selection": 0.32608695652173914, + "selected_success_rate": 0.21739130434782608, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.8909226783110455, + "potential_edge_mae": 0.4902584381812314, + "effect_prediction_mae": 0.3660368310325896, + "selection_regret": 0.31549541409248893, + "selected_candidate_type_counts": { + "expert": 29, + "near_miss": 35, + "no_op": 7, + "random_negative": 18, + "wrong_direction": 3 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_1.out new file mode 100644 index 0000000000000000000000000000000000000000..21450559c4087e17e51ecda36d664f2c9816b549 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_1.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/no_effect_head/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.7234249728116543, + "top1_action_selection": 0.4657142857142857, + "selected_success_rate": 0.33714285714285713, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.9117040558927005, + "potential_edge_mae": 0.3596969935604576, + "effect_prediction_mae": 0.5614495104190628, + "selection_regret": 0.2062273244649155, + "selected_candidate_type_counts": { + "expert": 276, + "near_miss": 217, + "no_op": 54, + "random_negative": 92, + "wrong_direction": 61 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.7253584905660377, + "top1_action_selection": 0.4824561403508772, + "selected_success_rate": 0.3508771929824561, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.936309512027014, + "potential_edge_mae": 0.3914277013245056, + "effect_prediction_mae": 0.5022691233838397, + "selection_regret": 0.20333470940067058, + "selected_candidate_type_counts": { + "expert": 57, + "near_miss": 36, + "random_negative": 10, + "wrong_direction": 11 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.7195625759416768, + "top1_action_selection": 0.2524271844660194, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.8392749076358842, + "potential_edge_mae": 0.1399603868298788, + "effect_prediction_mae": 0.8722028661092376, + "selection_regret": 0.11306014229081236, + "selected_candidate_type_counts": { + "expert": 36, + "near_miss": 47, + "random_negative": 8, + "wrong_direction": 12 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.6923176633252679, + "top1_action_selection": 0.38144329896907214, + "selected_success_rate": 0.24742268041237114, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.8956555922292279, + "potential_edge_mae": 0.3987140652748237, + "effect_prediction_mae": 0.5527565708705453, + "selection_regret": 0.29774537075733404, + "selected_candidate_type_counts": { + "expert": 59, + "near_miss": 42, + "no_op": 51, + "random_negative": 29, + "wrong_direction": 13 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.7812343201204215, + "top1_action_selection": 0.719626168224299, + "selected_success_rate": 0.6542056074766355, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.9768627723726377, + "potential_edge_mae": 0.29710391924828217, + "effect_prediction_mae": 0.4839220183566557, + "selection_regret": 0.025235589280306735, + "selected_candidate_type_counts": { + "expert": 60, + "near_miss": 22, + "no_op": 1, + "random_negative": 14, + "wrong_direction": 10 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.7722668894225749, + "top1_action_selection": 0.7472527472527473, + "selected_success_rate": 0.7032967032967034, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.9654196571109714, + "potential_edge_mae": 0.40927013101819987, + "effect_prediction_mae": 0.4290930855787791, + "selection_regret": 0.05847253527615096, + "selected_candidate_type_counts": { + "expert": 41, + "near_miss": 38, + "random_negative": 4, + "wrong_direction": 8 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.6954709641567102, + "top1_action_selection": 0.2857142857142857, + "selected_success_rate": 0.14285714285714285, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.8667422959766355, + "potential_edge_mae": 0.5021609624912808, + "effect_prediction_mae": 0.5259030826111822, + "selection_regret": 0.480768845140279, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 32, + "no_op": 2, + "random_negative": 27, + "wrong_direction": 7 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_2.out new file mode 100644 index 0000000000000000000000000000000000000000..8a1bb681e544f38d677809672e67a5c5d12012e6 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472352_2.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/no_effect_head/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.7317468494389928, + "top1_action_selection": 0.5071428571428571, + "selected_success_rate": 0.32857142857142857, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9261912945944161, + "potential_edge_mae": 0.3250027989105571, + "effect_prediction_mae": 0.39501799610026855, + "selection_regret": 0.16022627177764662, + "selected_candidate_type_counts": { + "expert": 254, + "near_miss": 240, + "no_op": 62, + "random_negative": 67, + "wrong_direction": 77 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.727725556358962, + "top1_action_selection": 0.5277777777777778, + "selected_success_rate": 0.35185185185185186, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.9425967550908285, + "potential_edge_mae": 0.3247552866783139, + "effect_prediction_mae": 0.36117758623091933, + "selection_regret": 0.18266619728119285, + "selected_candidate_type_counts": { + "expert": 51, + "near_miss": 28, + "random_negative": 12, + "wrong_direction": 17 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.7212280701754386, + "top1_action_selection": 0.45263157894736844, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.8741408424619007, + "potential_edge_mae": 0.13393139598858717, + "effect_prediction_mae": 0.5426230617944877, + "selection_regret": 0.12091110656405554, + "selected_candidate_type_counts": { + "expert": 51, + "near_miss": 35, + "random_negative": 2, + "wrong_direction": 7 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.7263528554660622, + "top1_action_selection": 0.5072463768115942, + "selected_success_rate": 0.30917874396135264, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.9298517660661595, + "potential_edge_mae": 0.34319187163109155, + "effect_prediction_mae": 0.40153301702223354, + "selection_regret": 0.13245852575929845, + "selected_candidate_type_counts": { + "expert": 48, + "near_miss": 71, + "no_op": 57, + "random_negative": 15, + "wrong_direction": 16 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.7788882650196519, + "top1_action_selection": 0.6702127659574468, + "selected_success_rate": 0.5957446808510638, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9666763454689263, + "potential_edge_mae": 0.23443176336711977, + "effect_prediction_mae": 0.3684889694146514, + "selection_regret": 0.06515261360463943, + "selected_candidate_type_counts": { + "expert": 42, + "near_miss": 30, + "random_negative": 6, + "wrong_direction": 16 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.77887362890589, + "top1_action_selection": 0.6138613861386139, + "selected_success_rate": 0.5445544554455446, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.955674182349952, + "potential_edge_mae": 0.43701953301253443, + "effect_prediction_mae": 0.29720101501972074, + "selection_regret": 0.13169119116103295, + "selected_candidate_type_counts": { + "expert": 39, + "near_miss": 45, + "no_op": 1, + "random_negative": 6, + "wrong_direction": 10 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.6760925449871465, + "top1_action_selection": 0.2631578947368421, + "selected_success_rate": 0.17894736842105263, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.8802114964766519, + "potential_edge_mae": 0.4422940746486034, + "effect_prediction_mae": 0.4019329149600042, + "selection_regret": 0.35894542101182436, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 31, + "no_op": 4, + "random_negative": 26, + "wrong_direction": 11 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_0.out new file mode 100644 index 0000000000000000000000000000000000000000..806bad97bb12b08b8b2b95f77f0fcd8d4da5dd50 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_0.out @@ -0,0 +1,148 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/expert_only_bc/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.49705057783303097, + "top1_action_selection": 0.1357142857142857, + "selected_success_rate": 0.13, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.72766914691033, + "potential_edge_mae": 1.5413571579455718, + "effect_prediction_mae": 0.4136177354324551, + "selection_regret": 0.5840117668476056, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 33, + "no_op": 330, + "random_negative": 301, + "wrong_direction": 13 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.4482758620689655, + "top1_action_selection": 0.06060606060606061, + "selected_success_rate": 0.06060606060606061, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.7914344148679657, + "potential_edge_mae": 1.2393729727230616, + "effect_prediction_mae": 0.38739678945341205, + "selection_regret": 0.6468019532133834, + "selected_candidate_type_counts": { + "near_miss": 2, + "no_op": 59, + "random_negative": 36, + "wrong_direction": 2 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.4327319587628866, + "top1_action_selection": 0.0, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.4662526331558357, + "potential_edge_mae": 0.8392176363262948, + "effect_prediction_mae": 0.4493809333730279, + "selection_regret": 0.254198014473941, + "selected_candidate_type_counts": { + "near_miss": 1, + "no_op": 70, + "random_negative": 26 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.5350907669590897, + "top1_action_selection": 0.14507772020725387, + "selected_success_rate": 0.13989637305699482, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.7919304991591462, + "potential_edge_mae": 2.337849868922095, + "effect_prediction_mae": 0.3443771208339658, + "selection_regret": 0.4609540672538503, + "selected_candidate_type_counts": { + "expert": 6, + "near_miss": 11, + "no_op": 109, + "random_negative": 59, + "wrong_direction": 8 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.4607436021245775, + "top1_action_selection": 0.23893805309734514, + "selected_success_rate": 0.21238938053097345, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.7356578067840379, + "potential_edge_mae": 0.928512991388845, + "effect_prediction_mae": 0.5118192267519918, + "selection_regret": 0.7923555153265464, + "selected_candidate_type_counts": { + "expert": 4, + "near_miss": 1, + "no_op": 28, + "random_negative": 80 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.5019728189390618, + "top1_action_selection": 0.19811320754716982, + "selected_success_rate": 0.19811320754716982, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.7283525479100315, + "potential_edge_mae": 0.8830868459947273, + "effect_prediction_mae": 0.4876345531121945, + "selection_regret": 0.8826292917694686, + "selected_candidate_type_counts": { + "expert": 5, + "near_miss": 7, + "no_op": 26, + "random_negative": 66, + "wrong_direction": 2 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.5659858961443356, + "top1_action_selection": 0.14130434782608695, + "selected_success_rate": 0.14130434782608695, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.7892674105044776, + "potential_edge_mae": 2.1965841479897663, + "effect_prediction_mae": 0.3434843753893896, + "selection_regret": 0.5223763356714145, + "selected_candidate_type_counts": { + "expert": 8, + "near_miss": 11, + "no_op": 38, + "random_negative": 34, + "wrong_direction": 1 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_1.out new file mode 100644 index 0000000000000000000000000000000000000000..f9c9849948549fb9ef1466554fa0c3e5d762f34c --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_1.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/expert_only_bc/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.4904651103973292, + "top1_action_selection": 0.12, + "selected_success_rate": 0.12, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.710367161245905, + "potential_edge_mae": 1.6845835190217973, + "effect_prediction_mae": 0.4025428864958327, + "selection_regret": 0.6546930451934792, + "selected_candidate_type_counts": { + "expert": 27, + "near_miss": 51, + "no_op": 341, + "random_negative": 267, + "wrong_direction": 14 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.4381132075471698, + "top1_action_selection": 0.07017543859649122, + "selected_success_rate": 0.07017543859649122, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.7751769133750837, + "potential_edge_mae": 1.188396140447765, + "effect_prediction_mae": 0.40581043631568836, + "selection_regret": 0.6973320498016843, + "selected_candidate_type_counts": { + "expert": 2, + "near_miss": 4, + "no_op": 60, + "random_negative": 44, + "wrong_direction": 4 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.41887403807209395, + "top1_action_selection": 0.0, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.4658945851616218, + "potential_edge_mae": 1.6122038452797571, + "effect_prediction_mae": 0.49771970467784693, + "selection_regret": 0.2923479685497705, + "selected_candidate_type_counts": { + "no_op": 69, + "random_negative": 33, + "wrong_direction": 1 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.5279122018665745, + "top1_action_selection": 0.16494845360824742, + "selected_success_rate": 0.16494845360824742, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.7809161987278901, + "potential_edge_mae": 2.41471324865967, + "effect_prediction_mae": 0.3561255569893941, + "selection_regret": 0.4952226573236517, + "selected_candidate_type_counts": { + "expert": 13, + "near_miss": 28, + "no_op": 106, + "random_negative": 45, + "wrong_direction": 2 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.4549924736578023, + "top1_action_selection": 0.17757009345794392, + "selected_success_rate": 0.17757009345794392, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.7161701491517457, + "potential_edge_mae": 1.2005173169511962, + "effect_prediction_mae": 0.4096431999131803, + "selection_regret": 0.9307221069096405, + "selected_candidate_type_counts": { + "expert": 4, + "no_op": 19, + "random_negative": 83, + "wrong_direction": 1 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.5415056907173437, + "top1_action_selection": 0.17582417582417584, + "selected_success_rate": 0.17582417582417584, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.6946875339752837, + "potential_edge_mae": 0.8323916706686492, + "effect_prediction_mae": 0.40421099914777825, + "selection_regret": 1.0593963692803958, + "selected_candidate_type_counts": { + "expert": 3, + "near_miss": 5, + "no_op": 40, + "random_negative": 38, + "wrong_direction": 5 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.5437621561544873, + "top1_action_selection": 0.0989010989010989, + "selected_success_rate": 0.0989010989010989, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.764342729357462, + "potential_edge_mae": 2.013846130027085, + "effect_prediction_mae": 0.37966065776016467, + "selection_regret": 0.6221088640964948, + "selected_candidate_type_counts": { + "expert": 5, + "near_miss": 14, + "no_op": 47, + "random_negative": 24, + "wrong_direction": 1 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_2.out new file mode 100644 index 0000000000000000000000000000000000000000..336acd883410bdafe594f705a40855a498770b91 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472354_2.out @@ -0,0 +1,147 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/expert_only_bc/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.48888665519104396, + "top1_action_selection": 0.13428571428571429, + "selected_success_rate": 0.13142857142857142, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.7208304467611435, + "potential_edge_mae": 1.2691906085473057, + "effect_prediction_mae": 0.32098602420830796, + "selection_regret": 0.582428052075016, + "selected_candidate_type_counts": { + "expert": 31, + "near_miss": 48, + "no_op": 401, + "random_negative": 210, + "wrong_direction": 10 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.42395758013979273, + "top1_action_selection": 0.06481481481481481, + "selected_success_rate": 0.06481481481481481, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.7734316664074604, + "potential_edge_mae": 1.4278430513127163, + "effect_prediction_mae": 0.2720774418611122, + "selection_regret": 0.68682034768992, + "selected_candidate_type_counts": { + "near_miss": 1, + "no_op": 80, + "random_negative": 27 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.3993859649122807, + "top1_action_selection": 0.0, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.4569514909906008, + "potential_edge_mae": 0.6248586803276792, + "effect_prediction_mae": 0.4539024277113717, + "selection_regret": 0.3121591664130775, + "selected_candidate_type_counts": { + "near_miss": 2, + "no_op": 74, + "random_negative": 19 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.5336544299186465, + "top1_action_selection": 0.18840579710144928, + "selected_success_rate": 0.18840579710144928, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.7839427094033568, + "potential_edge_mae": 1.6939450603873492, + "effect_prediction_mae": 0.27389198230792927, + "selection_regret": 0.42632551822397446, + "selected_candidate_type_counts": { + "expert": 12, + "near_miss": 22, + "no_op": 147, + "random_negative": 25, + "wrong_direction": 1 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.4678270634475014, + "top1_action_selection": 0.24468085106382978, + "selected_success_rate": 0.23404255319148937, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.7265061974614836, + "potential_edge_mae": 1.1001630125721056, + "effect_prediction_mae": 0.3526690487651323, + "selection_regret": 0.7767161505415718, + "selected_candidate_type_counts": { + "expert": 1, + "near_miss": 4, + "no_op": 25, + "random_negative": 62, + "wrong_direction": 2 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.5074200387132455, + "top1_action_selection": 0.09900990099009901, + "selected_success_rate": 0.09900990099009901, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.7187502356178154, + "potential_edge_mae": 0.7508768196676896, + "effect_prediction_mae": 0.35359266972340414, + "selection_regret": 0.9809670057332162, + "selected_candidate_type_counts": { + "expert": 5, + "near_miss": 4, + "no_op": 21, + "random_negative": 64, + "wrong_direction": 7 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.5517241379310345, + "top1_action_selection": 0.15789473684210525, + "selected_success_rate": 0.14736842105263157, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.7839874095938317, + "potential_edge_mae": 1.4468881373749658, + "effect_prediction_mae": 0.2802708634579669, + "selection_regret": 0.4582057904256018, + "selected_candidate_type_counts": { + "expert": 13, + "near_miss": 15, + "no_op": 54, + "random_negative": 13 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_0.out new file mode 100644 index 0000000000000000000000000000000000000000..1cf0876a9808fcb73efde32fa6c7fbdf39b172bd --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_0.out @@ -0,0 +1,147 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/label_only_counterfactual/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.6772888150300016, + "top1_action_selection": 0.5271428571428571, + "selected_success_rate": 0.3342857142857143, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9209124867375243, + "potential_edge_mae": 0.7205317208472326, + "effect_prediction_mae": 0.2990849819155833, + "selection_regret": 0.1368939660921959, + "selected_candidate_type_counts": { + "expert": 482, + "near_miss": 156, + "no_op": 8, + "random_negative": 30, + "wrong_direction": 24 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.6800766283524904, + "top1_action_selection": 0.5252525252525253, + "selected_success_rate": 0.31313131313131315, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9383288865676092, + "potential_edge_mae": 0.7712327697236793, + "effect_prediction_mae": 0.2904304617659481, + "selection_regret": 0.19677931551981453, + "selected_candidate_type_counts": { + "expert": 76, + "near_miss": 16, + "random_negative": 7 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.6948453608247422, + "top1_action_selection": 0.4845360824742268, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.8915467237578022, + "potential_edge_mae": 0.741650273974771, + "effect_prediction_mae": 0.3212094911528003, + "selection_regret": 0.052693408409844045, + "selected_candidate_type_counts": { + "expert": 61, + "near_miss": 34, + "random_negative": 1, + "wrong_direction": 1 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.655736992964475, + "top1_action_selection": 0.45077720207253885, + "selected_success_rate": 0.24870466321243523, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9024163483269362, + "potential_edge_mae": 0.781928575840302, + "effect_prediction_mae": 0.2960060443169843, + "selection_regret": 0.1871162539087429, + "selected_candidate_type_counts": { + "expert": 105, + "near_miss": 63, + "no_op": 3, + "random_negative": 10, + "wrong_direction": 12 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.6864316755190729, + "top1_action_selection": 0.6902654867256637, + "selected_success_rate": 0.5929203539823009, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9604198527343484, + "potential_edge_mae": 0.6107622549437861, + "effect_prediction_mae": 0.28476447270839644, + "selection_regret": 0.03880291504670033, + "selected_candidate_type_counts": { + "expert": 103, + "near_miss": 8, + "random_negative": 1, + "wrong_direction": 1 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.7024989039894783, + "top1_action_selection": 0.6509433962264151, + "selected_success_rate": 0.6226415094339622, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.9489160902694359, + "potential_edge_mae": 0.5808658020687647, + "effect_prediction_mae": 0.28698072022643095, + "selection_regret": 0.054298371920045815, + "selected_candidate_type_counts": { + "expert": 95, + "near_miss": 9, + "random_negative": 1, + "wrong_direction": 1 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6660866379705102, + "top1_action_selection": 0.391304347826087, + "selected_success_rate": 0.2391304347826087, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.891143962901377, + "potential_edge_mae": 0.765202280085857, + "effect_prediction_mae": 0.3230656812068449, + "selection_regret": 0.2715170606970787, + "selected_candidate_type_counts": { + "expert": 42, + "near_miss": 26, + "no_op": 5, + "random_negative": 10, + "wrong_direction": 9 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_1.out new file mode 100644 index 0000000000000000000000000000000000000000..09c0309d4a977765f517a9010913ce5caaf877f7 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_1.out @@ -0,0 +1,146 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/label_only_counterfactual/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.6725258605427552, + "top1_action_selection": 0.5342857142857143, + "selected_success_rate": 0.3385714285714286, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.9106497982809733, + "potential_edge_mae": 0.7164104849571172, + "effect_prediction_mae": 0.35979254204292677, + "selection_regret": 0.1931616946093605, + "selected_candidate_type_counts": { + "expert": 497, + "near_miss": 123, + "no_op": 15, + "random_negative": 49, + "wrong_direction": 16 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.6591698113207547, + "top1_action_selection": 0.4298245614035088, + "selected_success_rate": 0.2982456140350877, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9220632143613978, + "potential_edge_mae": 0.7404850526811941, + "effect_prediction_mae": 0.3566313981562654, + "selection_regret": 0.2658101535941425, + "selected_candidate_type_counts": { + "expert": 81, + "near_miss": 20, + "random_negative": 10, + "wrong_direction": 3 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.6865127582017011, + "top1_action_selection": 0.5048543689320388, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.8842688135773559, + "potential_edge_mae": 0.7904946348186774, + "effect_prediction_mae": 0.3737682691170755, + "selection_regret": 0.0710455625503351, + "selected_candidate_type_counts": { + "expert": 66, + "near_miss": 34, + "random_negative": 2, + "wrong_direction": 1 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.6507950224680262, + "top1_action_selection": 0.5, + "selected_success_rate": 0.25773195876288657, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.885194898687762, + "potential_edge_mae": 0.7510739888456183, + "effect_prediction_mae": 0.348361954130037, + "selection_regret": 0.27040022373506706, + "selected_candidate_type_counts": { + "expert": 116, + "near_miss": 43, + "no_op": 9, + "random_negative": 20, + "wrong_direction": 6 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.7085800301053687, + "top1_action_selection": 0.7009345794392523, + "selected_success_rate": 0.6355140186915887, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.9641326456760453, + "potential_edge_mae": 0.5528920303458209, + "effect_prediction_mae": 0.35872660063423906, + "selection_regret": 0.054595183685561204, + "selected_candidate_type_counts": { + "expert": 103, + "near_miss": 4 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.7098256238905711, + "top1_action_selection": 0.7252747252747253, + "selected_success_rate": 0.7032967032967034, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.9547552599934142, + "potential_edge_mae": 0.6116512574444513, + "effect_prediction_mae": 0.35881002015476265, + "selection_regret": 0.06150194855181725, + "selected_candidate_type_counts": { + "expert": 88, + "near_miss": 1, + "no_op": 1, + "wrong_direction": 1 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.6531443919607298, + "top1_action_selection": 0.38461538461538464, + "selected_success_rate": 0.21978021978021978, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.8734860710862703, + "potential_edge_mae": 0.7716972808920919, + "effect_prediction_mae": 0.3745383639340848, + "selection_regret": 0.37029813623035346, + "selected_candidate_type_counts": { + "expert": 43, + "near_miss": 21, + "no_op": 5, + "random_negative": 17, + "wrong_direction": 5 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_2.out new file mode 100644 index 0000000000000000000000000000000000000000..e108e40f12735620a35e355cda586734f38ce462 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472356_2.out @@ -0,0 +1,146 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_baseline_runs/label_only_counterfactual/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.6729321891216122, + "top1_action_selection": 0.49, + "selected_success_rate": 0.30714285714285716, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.913497472252105, + "potential_edge_mae": 0.7225398868189508, + "effect_prediction_mae": 0.34287250271850117, + "selection_regret": 0.18427200830558182, + "selected_candidate_type_counts": { + "expert": 457, + "near_miss": 172, + "no_op": 9, + "random_negative": 36, + "wrong_direction": 26 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.6612838434964249, + "top1_action_selection": 0.5, + "selected_success_rate": 0.3611111111111111, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.943135387035823, + "potential_edge_mae": 0.7260483337916878, + "effect_prediction_mae": 0.33281637966564626, + "selection_regret": 0.1712497730773908, + "selected_candidate_type_counts": { + "expert": 72, + "near_miss": 23, + "random_negative": 7, + "wrong_direction": 6 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.7179824561403508, + "top1_action_selection": 0.35789473684210527, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.8669320795608593, + "potential_edge_mae": 0.7418219292177379, + "effect_prediction_mae": 0.3849098708759793, + "selection_regret": 0.14754043332374606, + "selected_candidate_type_counts": { + "expert": 45, + "near_miss": 46, + "random_negative": 2, + "wrong_direction": 2 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.6563322135427206, + "top1_action_selection": 0.463768115942029, + "selected_success_rate": 0.2028985507246377, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.8952240327465547, + "potential_edge_mae": 0.7788038356041567, + "effect_prediction_mae": 0.3326538219964117, + "selection_regret": 0.27153090060045176, + "selected_candidate_type_counts": { + "expert": 115, + "near_miss": 63, + "no_op": 5, + "random_negative": 13, + "wrong_direction": 11 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.7015160022459292, + "top1_action_selection": 0.7127659574468085, + "selected_success_rate": 0.6170212765957447, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9664684753556018, + "potential_edge_mae": 0.5731928828961473, + "effect_prediction_mae": 0.33214158477256467, + "selection_regret": 0.025537094378725013, + "selected_candidate_type_counts": { + "expert": 89, + "near_miss": 4, + "random_negative": 1 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.6819061664669555, + "top1_action_selection": 0.6336633663366337, + "selected_success_rate": 0.5643564356435643, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.9442542094032579, + "potential_edge_mae": 0.5955867691859857, + "effect_prediction_mae": 0.3304624205646258, + "selection_regret": 0.09426688188963597, + "selected_candidate_type_counts": { + "expert": 94, + "near_miss": 5, + "random_negative": 2 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.6454215051857105, + "top1_action_selection": 0.29473684210526313, + "selected_success_rate": 0.2, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.881073416911903, + "potential_edge_mae": 0.8159403954334963, + "effect_prediction_mae": 0.3583451638624953, + "selection_regret": 0.29842895567417144, + "selected_candidate_type_counts": { + "expert": 42, + "near_miss": 31, + "no_op": 4, + "random_negative": 11, + "wrong_direction": 7 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_0.out new file mode 100644 index 0000000000000000000000000000000000000000..89854b2fd102feb068a3c64a013f55571d1a56d4 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_0.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_16/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 0, + "k": 16, + "training_k": 16, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.6534398458848534, + "top1_action_selection": 0.35, + "selected_success_rate": 0.25125, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8860278809455133, + "potential_edge_mae": 0.3090061780100042, + "effect_prediction_mae": 0.025982485414330476, + "selection_regret": 0.2148703686054796, + "selected_candidate_type_counts": { + "expert": 113, + "near_miss": 123, + "no_op": 308, + "random_negative": 212, + "wrong_direction": 44 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.6534398458848534, + "top1_action_selection": 0.35, + "selected_success_rate": 0.25125, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8860278809455133, + "potential_edge_mae": 0.3090061780100042, + "effect_prediction_mae": 0.025982485414330476, + "selection_regret": 0.2148703686054796, + "selected_candidate_type_counts": { + "expert": 113, + "near_miss": 123, + "no_op": 308, + "random_negative": 212, + "wrong_direction": 44 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_1.out new file mode 100644 index 0000000000000000000000000000000000000000..1cbb56991bb65d2e54ab96eb66b60b84b8dfc618 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_1.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_16/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 1, + "k": 16, + "training_k": 16, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.6626532513898632, + "top1_action_selection": 0.3575, + "selected_success_rate": 0.25625, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8946500041524906, + "potential_edge_mae": 0.30419353209352523, + "effect_prediction_mae": 0.02415842195830894, + "selection_regret": 0.18901088452432305, + "selected_candidate_type_counts": { + "expert": 127, + "near_miss": 112, + "no_op": 359, + "random_negative": 158, + "wrong_direction": 44 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.6626532513898632, + "top1_action_selection": 0.3575, + "selected_success_rate": 0.25625, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8946500041524906, + "potential_edge_mae": 0.30419353209352523, + "effect_prediction_mae": 0.02415842195830894, + "selection_regret": 0.18901088452432305, + "selected_candidate_type_counts": { + "expert": 127, + "near_miss": 112, + "no_op": 359, + "random_negative": 158, + "wrong_direction": 44 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_2.out new file mode 100644 index 0000000000000000000000000000000000000000..7d3a43a93f639c00477940ddfc791d358d3369e2 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472903_2.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_16/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 2, + "k": 16, + "training_k": 16, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.6558060159350036, + "top1_action_selection": 0.3525, + "selected_success_rate": 0.2525, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8881420986242058, + "potential_edge_mae": 0.3113192438141851, + "effect_prediction_mae": 0.025496126152238938, + "selection_regret": 0.20792295036837458, + "selected_candidate_type_counts": { + "expert": 116, + "near_miss": 116, + "no_op": 318, + "random_negative": 196, + "wrong_direction": 54 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.6558060159350036, + "top1_action_selection": 0.3525, + "selected_success_rate": 0.2525, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8881420986242058, + "potential_edge_mae": 0.3113192438141851, + "effect_prediction_mae": 0.025496126152238938, + "selection_regret": 0.20792295036837458, + "selected_candidate_type_counts": { + "expert": 116, + "near_miss": 116, + "no_op": 318, + "random_negative": 196, + "wrong_direction": 54 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_0.out new file mode 100644 index 0000000000000000000000000000000000000000..5ff9c71dae70c82623fbfe2398b18f2da7c8d778 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_0.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_8/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 0, + "k": 8, + "training_k": 8, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5938982128087276, + "top1_action_selection": 0.3275, + "selected_success_rate": 0.25375, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.863645423237401, + "potential_edge_mae": 0.2954803167953407, + "effect_prediction_mae": 0.01971304282474624, + "selection_regret": 0.2495492117665708, + "selected_candidate_type_counts": { + "expert": 86, + "near_miss": 73, + "no_op": 265, + "random_negative": 295, + "wrong_direction": 81 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5938982128087276, + "top1_action_selection": 0.3275, + "selected_success_rate": 0.25375, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.863645423237401, + "potential_edge_mae": 0.2954803167953407, + "effect_prediction_mae": 0.01971304282474624, + "selection_regret": 0.2495492117665708, + "selected_candidate_type_counts": { + "expert": 86, + "near_miss": 73, + "no_op": 265, + "random_negative": 295, + "wrong_direction": 81 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_1.out new file mode 100644 index 0000000000000000000000000000000000000000..f4680c70c38e80ede96c4233857dd85b316b247f --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_1.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_8/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 1, + "k": 8, + "training_k": 8, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.584307895260331, + "top1_action_selection": 0.31875, + "selected_success_rate": 0.25375, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8598240181859493, + "potential_edge_mae": 0.3023171987637495, + "effect_prediction_mae": 0.020936218560997998, + "selection_regret": 0.251564243162502, + "selected_candidate_type_counts": { + "expert": 66, + "near_miss": 99, + "no_op": 232, + "random_negative": 313, + "wrong_direction": 90 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.584307895260331, + "top1_action_selection": 0.31875, + "selected_success_rate": 0.25375, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8598240181859493, + "potential_edge_mae": 0.3023171987637495, + "effect_prediction_mae": 0.020936218560997998, + "selection_regret": 0.251564243162502, + "selected_candidate_type_counts": { + "expert": 66, + "near_miss": 99, + "no_op": 232, + "random_negative": 313, + "wrong_direction": 90 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_2.out new file mode 100644 index 0000000000000000000000000000000000000000..ddb99c78bb352b22f8ff416727499bb1f817fa62 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472904_2.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_8/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 2, + "k": 8, + "training_k": 8, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5697444326950258, + "top1_action_selection": 0.305, + "selected_success_rate": 0.25, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8463434412197927, + "potential_edge_mae": 0.31044809396788436, + "effect_prediction_mae": 0.021288544259330366, + "selection_regret": 0.28888761965696175, + "selected_candidate_type_counts": { + "expert": 57, + "near_miss": 51, + "no_op": 249, + "random_negative": 367, + "wrong_direction": 76 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5697444326950258, + "top1_action_selection": 0.305, + "selected_success_rate": 0.25, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8463434412197927, + "potential_edge_mae": 0.31044809396788436, + "effect_prediction_mae": 0.021288544259330366, + "selection_regret": 0.28888761965696175, + "selected_candidate_type_counts": { + "expert": 57, + "near_miss": 51, + "no_op": 249, + "random_negative": 367, + "wrong_direction": 76 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_0.out new file mode 100644 index 0000000000000000000000000000000000000000..0f3fb40ddadbb268b6da2d9478383cff8baa786a --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_0.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_4/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 0, + "k": 4, + "training_k": 4, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.621842052914263, + "top1_action_selection": 0.30375, + "selected_success_rate": 0.22125, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8522906587786284, + "potential_edge_mae": 0.3192463427301466, + "effect_prediction_mae": 0.02234776312209753, + "selection_regret": 0.28148924679262566, + "selected_candidate_type_counts": { + "expert": 218, + "near_miss": 139, + "no_op": 109, + "random_negative": 295, + "wrong_direction": 39 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.621842052914263, + "top1_action_selection": 0.30375, + "selected_success_rate": 0.22125, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8522906587786284, + "potential_edge_mae": 0.3192463427301466, + "effect_prediction_mae": 0.02234776312209753, + "selection_regret": 0.28148924679262566, + "selected_candidate_type_counts": { + "expert": 218, + "near_miss": 139, + "no_op": 109, + "random_negative": 295, + "wrong_direction": 39 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_1.out new file mode 100644 index 0000000000000000000000000000000000000000..2338f56a73219b28290848f67c5ae77e325fb7d1 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_1.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_4/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 1, + "k": 4, + "training_k": 4, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.602315915111032, + "top1_action_selection": 0.25625, + "selected_success_rate": 0.215, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8388718499636448, + "potential_edge_mae": 0.3163053712629924, + "effect_prediction_mae": 0.022465208080843763, + "selection_regret": 0.2986432422319922, + "selected_candidate_type_counts": { + "expert": 161, + "near_miss": 113, + "no_op": 331, + "random_negative": 176, + "wrong_direction": 19 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.602315915111032, + "top1_action_selection": 0.25625, + "selected_success_rate": 0.215, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8388718499636448, + "potential_edge_mae": 0.3163053712629924, + "effect_prediction_mae": 0.022465208080843763, + "selection_regret": 0.2986432422319922, + "selected_candidate_type_counts": { + "expert": 161, + "near_miss": 113, + "no_op": 331, + "random_negative": 176, + "wrong_direction": 19 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_2.out new file mode 100644 index 0000000000000000000000000000000000000000..a6f36e06b94646dd554cbb8cb353ca22d38eb8fe --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472905_2.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_4/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 2, + "k": 4, + "training_k": 4, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.6194444735271638, + "top1_action_selection": 0.3, + "selected_success_rate": 0.23375, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8482839695063008, + "potential_edge_mae": 0.31006770103584785, + "effect_prediction_mae": 0.021614359051746374, + "selection_regret": 0.26884213731391354, + "selected_candidate_type_counts": { + "expert": 204, + "near_miss": 142, + "no_op": 157, + "random_negative": 266, + "wrong_direction": 31 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.6194444735271638, + "top1_action_selection": 0.3, + "selected_success_rate": 0.23375, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8482839695063008, + "potential_edge_mae": 0.31006770103584785, + "effect_prediction_mae": 0.021614359051746374, + "selection_regret": 0.26884213731391354, + "selected_candidate_type_counts": { + "expert": 204, + "near_miss": 142, + "no_op": 157, + "random_negative": 266, + "wrong_direction": 31 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_0.out new file mode 100644 index 0000000000000000000000000000000000000000..84b0eb8f88b6f0546e2d3fa82fa4e375405eaf0b --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_0.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_2/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 0, + "k": 2, + "training_k": 2, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.527310418477066, + "top1_action_selection": 0.11, + "selected_success_rate": 0.08375, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.7666593662923569, + "potential_edge_mae": 0.31124381776305426, + "effect_prediction_mae": 0.027252909106802004, + "selection_regret": 0.5732752894894293, + "selected_candidate_type_counts": { + "expert": 67, + "near_miss": 93, + "no_op": 7, + "random_negative": 589, + "wrong_direction": 44 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.527310418477066, + "top1_action_selection": 0.11, + "selected_success_rate": 0.08375, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.7666593662923569, + "potential_edge_mae": 0.31124381776305426, + "effect_prediction_mae": 0.027252909106802004, + "selection_regret": 0.5732752894894293, + "selected_candidate_type_counts": { + "expert": 67, + "near_miss": 93, + "no_op": 7, + "random_negative": 589, + "wrong_direction": 44 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_1.out new file mode 100644 index 0000000000000000000000000000000000000000..abc71ce50578373558770f3bb7f406b738944b91 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_1.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_2/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 1, + "k": 2, + "training_k": 2, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5490142703087538, + "top1_action_selection": 0.21125, + "selected_success_rate": 0.1975, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8078818481867149, + "potential_edge_mae": 0.3123134503346493, + "effect_prediction_mae": 0.029179018532695698, + "selection_regret": 0.359973800514781, + "selected_candidate_type_counts": { + "expert": 74, + "near_miss": 62, + "no_op": 431, + "random_negative": 218, + "wrong_direction": 15 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5490142703087538, + "top1_action_selection": 0.21125, + "selected_success_rate": 0.1975, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8078818481867149, + "potential_edge_mae": 0.3123134503346493, + "effect_prediction_mae": 0.029179018532695698, + "selection_regret": 0.359973800514781, + "selected_candidate_type_counts": { + "expert": 74, + "near_miss": 62, + "no_op": 431, + "random_negative": 218, + "wrong_direction": 15 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_2.out new file mode 100644 index 0000000000000000000000000000000000000000..70af5670f13a0a77bc1e5c242c8fc045a083aa0e --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472906_2.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_2/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 2, + "k": 2, + "training_k": 2, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5367227497827521, + "top1_action_selection": 0.22625, + "selected_success_rate": 0.21, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8097508473418127, + "potential_edge_mae": 0.33309778105146864, + "effect_prediction_mae": 0.02330698871690793, + "selection_regret": 0.3483553205308999, + "selected_candidate_type_counts": { + "expert": 79, + "near_miss": 68, + "no_op": 323, + "random_negative": 315, + "wrong_direction": 15 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5367227497827521, + "top1_action_selection": 0.22625, + "selected_success_rate": 0.21, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.8097508473418127, + "potential_edge_mae": 0.33309778105146864, + "effect_prediction_mae": 0.02330698871690793, + "selection_regret": 0.3483553205308999, + "selected_candidate_type_counts": { + "expert": 79, + "near_miss": 68, + "no_op": 323, + "random_negative": 315, + "wrong_direction": 15 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_0.out new file mode 100644 index 0000000000000000000000000000000000000000..28bcb15b5e55eef21dddbd147d058bb3cdc4fe85 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_0.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_1/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 0, + "k": 1, + "training_k": 1, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5308387339943254, + "top1_action_selection": 0.10125, + "selected_success_rate": 0.07625, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.7657555092811404, + "potential_edge_mae": 1.4864671993784482, + "effect_prediction_mae": 0.027795869459491817, + "selection_regret": 0.559511656598188, + "selected_candidate_type_counts": { + "expert": 44, + "near_miss": 76, + "no_op": 249, + "random_negative": 396, + "wrong_direction": 35 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5308387339943254, + "top1_action_selection": 0.10125, + "selected_success_rate": 0.07625, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.7657555092811404, + "potential_edge_mae": 1.4864671993784482, + "effect_prediction_mae": 0.027795869459491817, + "selection_regret": 0.559511656598188, + "selected_candidate_type_counts": { + "expert": 44, + "near_miss": 76, + "no_op": 249, + "random_negative": 396, + "wrong_direction": 35 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_1.out new file mode 100644 index 0000000000000000000000000000000000000000..4a879385c4d3e16d5cd0f8634f05129bec55e127 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_1.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_1/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 1, + "k": 1, + "training_k": 1, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5570655303466544, + "top1_action_selection": 0.14125, + "selected_success_rate": 0.115, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.7919853125320906, + "potential_edge_mae": 1.4632684730985916, + "effect_prediction_mae": 0.030116774968667756, + "selection_regret": 0.46937384627759454, + "selected_candidate_type_counts": { + "expert": 66, + "near_miss": 91, + "no_op": 349, + "random_negative": 269, + "wrong_direction": 25 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5570655303466544, + "top1_action_selection": 0.14125, + "selected_success_rate": 0.115, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.7919853125320906, + "potential_edge_mae": 1.4632684730985916, + "effect_prediction_mae": 0.030116774968667756, + "selection_regret": 0.46937384627759454, + "selected_candidate_type_counts": { + "expert": 66, + "near_miss": 91, + "no_op": 349, + "random_negative": 269, + "wrong_direction": 25 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_2.out new file mode 100644 index 0000000000000000000000000000000000000000..d28da62c156ef81ad08d75a9cb9374f1820cd56f --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14472907_2.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/runs/k_1/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_scaling_fixed14k/eval_common_k16_n800", + "split": "all_groups", + "seed": 2, + "k": 1, + "training_k": 1, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5363248981813994, + "top1_action_selection": 0.14875, + "selected_success_rate": 0.14, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.7868707676802389, + "potential_edge_mae": 1.2919401242700082, + "effect_prediction_mae": 0.030475957771326043, + "selection_regret": 0.4611841867649491, + "selected_candidate_type_counts": { + "expert": 57, + "near_miss": 54, + "no_op": 393, + "random_negative": 268, + "wrong_direction": 28 + }, + "per_task": { + "PickCube-v1": { + "num_groups": 800, + "num_records": 12800, + "num_pairs": 95513, + "pairwise_ranking_accuracy": 0.5363248981813994, + "top1_action_selection": 0.14875, + "selected_success_rate": 0.14, + "oracle_success_rate": 0.3375, + "ndcg_at_k": 0.7868707676802389, + "potential_edge_mae": 1.2919401242700082, + "effect_prediction_mae": 0.030475957771326043, + "selection_regret": 0.4611841867649491, + "selected_candidate_type_counts": { + "expert": 57, + "near_miss": 54, + "no_op": 393, + "random_negative": 268, + "wrong_direction": 28 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_0.out new file mode 100644 index 0000000000000000000000000000000000000000..258aa4c8354dfcbde0e0ff589344d2f1a7143157 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_0.out @@ -0,0 +1,131 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_random_baseline_runs/random_negatives/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_random_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 80269, + "pairwise_ranking_accuracy": 0.6291719094544593, + "top1_action_selection": 0.81, + "selected_success_rate": 0.34, + "oracle_success_rate": 0.3842857142857143, + "ndcg_at_k": 0.9216666693424218, + "potential_edge_mae": 0.22607365697827786, + "effect_prediction_mae": 0.03159558307767137, + "selection_regret": 0.10546007330683227, + "selected_candidate_type_counts": { + "expert": 539, + "random_negative": 161 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11753, + "pairwise_ranking_accuracy": 0.6259678379988088, + "top1_action_selection": 0.7878787878787878, + "selected_success_rate": 0.3434343434343434, + "oracle_success_rate": 0.4444444444444444, + "ndcg_at_k": 0.9298426688227615, + "potential_edge_mae": 0.27054849254897856, + "effect_prediction_mae": 0.023409125922238126, + "selection_regret": 0.14207402202818128, + "selected_candidate_type_counts": { + "expert": 83, + "random_negative": 16 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11639, + "pairwise_ranking_accuracy": 0.6072686656929289, + "top1_action_selection": 0.8247422680412371, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.010309278350515464, + "ndcg_at_k": 0.9048736031682795, + "potential_edge_mae": 0.07867660537163249, + "effect_prediction_mae": 0.05186781695825712, + "selection_regret": 0.05375575241734773, + "selected_candidate_type_counts": { + "expert": 81, + "random_negative": 16 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23160, + "pairwise_ranking_accuracy": 0.6063471502590674, + "top1_action_selection": 0.7668393782383419, + "selected_success_rate": 0.27461139896373055, + "oracle_success_rate": 0.3005181347150259, + "ndcg_at_k": 0.9128129948259938, + "potential_edge_mae": 0.2231412177949071, + "effect_prediction_mae": 0.02659036551397243, + "selection_regret": 0.09780959009043293, + "selected_candidate_type_counts": { + "expert": 148, + "random_negative": 45 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10683, + "pairwise_ranking_accuracy": 0.6681643733033792, + "top1_action_selection": 0.9292035398230089, + "selected_success_rate": 0.5752212389380531, + "oracle_success_rate": 0.5929203539823009, + "ndcg_at_k": 0.9627420480225328, + "potential_edge_mae": 0.17627527423849504, + "effect_prediction_mae": 0.025687881340181573, + "selection_regret": 0.03857547822779259, + "selected_candidate_type_counts": { + "expert": 75, + "random_negative": 38 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 12037, + "pairwise_ranking_accuracy": 0.6739220736063803, + "top1_action_selection": 0.9339622641509434, + "selected_success_rate": 0.6037735849056604, + "oracle_success_rate": 0.6415094339622641, + "ndcg_at_k": 0.9521528996543167, + "potential_edge_mae": 0.25093172626213806, + "effect_prediction_mae": 0.026896570041282306, + "selection_regret": 0.07094061979145375, + "selected_candidate_type_counts": { + "expert": 93, + "random_negative": 13 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10997, + "pairwise_ranking_accuracy": 0.6169864508502318, + "top1_action_selection": 0.6195652173913043, + "selected_success_rate": 0.2391304347826087, + "oracle_success_rate": 0.33695652173913043, + "ndcg_at_k": 0.8635710655173733, + "potential_edge_mae": 0.3618867310451531, + "effect_prediction_mae": 0.04220129220765622, + "selection_regret": 0.2585481470045836, + "selected_candidate_type_counts": { + "expert": 59, + "random_negative": 33 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_1.out new file mode 100644 index 0000000000000000000000000000000000000000..27dd83f67170e863496e67c42e5d40329473cc26 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_1.out @@ -0,0 +1,131 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_random_baseline_runs/random_negatives/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_random_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 80523, + "pairwise_ranking_accuracy": 0.6256597493883735, + "top1_action_selection": 0.7685714285714286, + "selected_success_rate": 0.34714285714285714, + "oracle_success_rate": 0.4, + "ndcg_at_k": 0.9068379747577626, + "potential_edge_mae": 0.24775687137291055, + "effect_prediction_mae": 0.03270106231319296, + "selection_regret": 0.12766878896995748, + "selected_candidate_type_counts": { + "expert": 524, + "random_negative": 176 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13547, + "pairwise_ranking_accuracy": 0.6270022883295194, + "top1_action_selection": 0.7719298245614035, + "selected_success_rate": 0.30701754385964913, + "oracle_success_rate": 0.45614035087719296, + "ndcg_at_k": 0.9261138250198282, + "potential_edge_mae": 0.30939958664444467, + "effect_prediction_mae": 0.025490457606901343, + "selection_regret": 0.20130285829828495, + "selected_candidate_type_counts": { + "expert": 97, + "random_negative": 17 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.6051842851356825, + "top1_action_selection": 0.6504854368932039, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.8367997551764268, + "potential_edge_mae": 0.08815213881319012, + "effect_prediction_mae": 0.054190120094181136, + "selection_regret": 0.08594586681693259, + "selected_candidate_type_counts": { + "expert": 67, + "random_negative": 36 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23280, + "pairwise_ranking_accuracy": 0.6127147766323024, + "top1_action_selection": 0.7680412371134021, + "selected_success_rate": 0.28350515463917525, + "oracle_success_rate": 0.29896907216494845, + "ndcg_at_k": 0.9111534168563954, + "potential_edge_mae": 0.22136822210119148, + "effect_prediction_mae": 0.02711403284258621, + "selection_regret": 0.08961961529764932, + "selected_candidate_type_counts": { + "expert": 150, + "random_negative": 44 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 10409, + "pairwise_ranking_accuracy": 0.6665385723892785, + "top1_action_selection": 0.9065420560747663, + "selected_success_rate": 0.616822429906542, + "oracle_success_rate": 0.6355140186915887, + "ndcg_at_k": 0.9595108846712749, + "potential_edge_mae": 0.2370789942024452, + "effect_prediction_mae": 0.02504524638195561, + "selection_regret": 0.04314068064232853, + "selected_candidate_type_counts": { + "expert": 74, + "random_negative": 33 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10056, + "pairwise_ranking_accuracy": 0.6576173428798727, + "top1_action_selection": 0.9560439560439561, + "selected_success_rate": 0.6813186813186813, + "oracle_success_rate": 0.7142857142857143, + "ndcg_at_k": 0.9523379964914687, + "potential_edge_mae": 0.30059434325888446, + "effect_prediction_mae": 0.029741010552391787, + "selection_regret": 0.0638686526607681, + "selected_candidate_type_counts": { + "expert": 82, + "random_negative": 9 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10886, + "pairwise_ranking_accuracy": 0.6062832996509278, + "top1_action_selection": 0.5494505494505495, + "selected_success_rate": 0.27472527472527475, + "oracle_success_rate": 0.3956043956043956, + "ndcg_at_k": 0.8453301789484718, + "potential_edge_mae": 0.36987580586619706, + "effect_prediction_mae": 0.04128410316207703, + "selection_regret": 0.3269549072771282, + "selected_candidate_type_counts": { + "expert": 54, + "random_negative": 37 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_2.out new file mode 100644 index 0000000000000000000000000000000000000000..ab0a334677763199d1a367ac8ffba82bd2943e34 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14473580_2.out @@ -0,0 +1,131 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_random_baseline_runs/random_negatives/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_random_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "legacy", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 81061, + "pairwise_ranking_accuracy": 0.635632424963916, + "top1_action_selection": 0.83, + "selected_success_rate": 0.3342857142857143, + "oracle_success_rate": 0.3842857142857143, + "ndcg_at_k": 0.9230602248156762, + "potential_edge_mae": 0.23016934835296143, + "effect_prediction_mae": 0.03445017593150741, + "selection_regret": 0.10520695794523427, + "selected_candidate_type_counts": { + "expert": 569, + "random_negative": 131 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12789, + "pairwise_ranking_accuracy": 0.6464930799906169, + "top1_action_selection": 0.7592592592592593, + "selected_success_rate": 0.3611111111111111, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.9349243365811276, + "potential_edge_mae": 0.25304923507061117, + "effect_prediction_mae": 0.024560852606341586, + "selection_regret": 0.17700350822673905, + "selected_candidate_type_counts": { + "expert": 89, + "random_negative": 19 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11395, + "pairwise_ranking_accuracy": 0.6234313295304958, + "top1_action_selection": 0.8, + "selected_success_rate": 0.010526315789473684, + "oracle_success_rate": 0.031578947368421054, + "ndcg_at_k": 0.8963983310835758, + "potential_edge_mae": 0.09517505891571333, + "effect_prediction_mae": 0.05608951594918173, + "selection_regret": 0.06486132063682339, + "selected_candidate_type_counts": { + "expert": 77, + "random_negative": 18 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24840, + "pairwise_ranking_accuracy": 0.6096618357487923, + "top1_action_selection": 0.8357487922705314, + "selected_success_rate": 0.27053140096618356, + "oracle_success_rate": 0.28502415458937197, + "ndcg_at_k": 0.9189817788340707, + "potential_edge_mae": 0.2120186122904431, + "effect_prediction_mae": 0.02735010132417066, + "selection_regret": 0.06909547062311772, + "selected_candidate_type_counts": { + "expert": 175, + "random_negative": 32 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 9286, + "pairwise_ranking_accuracy": 0.6613181132888218, + "top1_action_selection": 0.9680851063829787, + "selected_success_rate": 0.5957446808510638, + "oracle_success_rate": 0.6063829787234043, + "ndcg_at_k": 0.9619083192507752, + "potential_edge_mae": 0.2069595820255836, + "effect_prediction_mae": 0.03222224408360166, + "selection_regret": 0.02078902119017662, + "selected_candidate_type_counts": { + "expert": 73, + "random_negative": 21 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 11392, + "pairwise_ranking_accuracy": 0.6909234550561798, + "top1_action_selection": 0.9504950495049505, + "selected_success_rate": 0.5544554455445545, + "oracle_success_rate": 0.6039603960396039, + "ndcg_at_k": 0.9455302355040142, + "potential_edge_mae": 0.3044872455883593, + "effect_prediction_mae": 0.031102755676910187, + "selection_regret": 0.09244666049386015, + "selected_candidate_type_counts": { + "expert": 89, + "random_negative": 12 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11359, + "pairwise_ranking_accuracy": 0.6159873228277137, + "top1_action_selection": 0.6631578947368421, + "selected_success_rate": 0.2736842105263158, + "oracle_success_rate": 0.37894736842105264, + "ndcg_at_k": 0.8827928795067447, + "potential_edge_mae": 0.323963674042613, + "effect_prediction_mae": 0.04528744014808633, + "selection_regret": 0.23971182202038013, + "selected_candidate_type_counts": { + "expert": 66, + "random_negative": 29 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14474550_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474550_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14474550_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474550_0.out new file mode 100644 index 0000000000000000000000000000000000000000..3d84ec25b95afc0ff931786a902fb9398b462b7a --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474550_0.out @@ -0,0 +1,150 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_fieldpref_runs/lattice_field/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.7338035494551498, + "top1_action_selection": 0.4614285714285714, + "selected_success_rate": 0.33, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9262306799356012, + "potential_edge_mae": 0.33351665406514996, + "effect_prediction_mae": 0.02902158687323396, + "selection_regret": 0.16100156836926804, + "selected_candidate_type_counts": { + "expert": 163, + "near_miss": 281, + "no_op": 97, + "random_negative": 71, + "wrong_direction": 88 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.7280564263322884, + "top1_action_selection": 0.46464646464646464, + "selected_success_rate": 0.30303030303030304, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9410048092579556, + "potential_edge_mae": 0.34914833802738926, + "effect_prediction_mae": 0.023742609922572207, + "selection_regret": 0.21750011528381194, + "selected_candidate_type_counts": { + "expert": 34, + "near_miss": 40, + "random_negative": 9, + "wrong_direction": 16 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.731615120274914, + "top1_action_selection": 0.27835051546391754, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.8476862466530304, + "potential_edge_mae": 0.20583194468424662, + "effect_prediction_mae": 0.04841492460531327, + "selection_regret": 0.09056418870465312, + "selected_candidate_type_counts": { + "expert": 41, + "near_miss": 37, + "random_negative": 5, + "wrong_direction": 14 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.7067228350560236, + "top1_action_selection": 0.36787564766839376, + "selected_success_rate": 0.29015544041450775, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9288474740774593, + "potential_edge_mae": 0.3428461337621427, + "effect_prediction_mae": 0.022750797341170137, + "selection_regret": 0.15311095828837062, + "selected_candidate_type_counts": { + "expert": 15, + "near_miss": 62, + "no_op": 81, + "random_negative": 16, + "wrong_direction": 19 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.77189763399324, + "top1_action_selection": 0.672566371681416, + "selected_success_rate": 0.584070796460177, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9691225616514149, + "potential_edge_mae": 0.3072038161533717, + "effect_prediction_mae": 0.024985675590603978, + "selection_regret": 0.06025841898095291, + "selected_candidate_type_counts": { + "expert": 30, + "near_miss": 47, + "no_op": 2, + "random_negative": 23, + "wrong_direction": 11 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.7976326172731258, + "top1_action_selection": 0.6886792452830188, + "selected_success_rate": 0.5849056603773585, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.968091164332034, + "potential_edge_mae": 0.377464857391755, + "effect_prediction_mae": 0.0262612250737494, + "selection_regret": 0.12097368881387531, + "selected_candidate_type_counts": { + "expert": 24, + "near_miss": 59, + "no_op": 1, + "random_negative": 3, + "wrong_direction": 19 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6964923527795586, + "top1_action_selection": 0.32608695652173914, + "selected_success_rate": 0.18478260869565216, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.8867430275030913, + "potential_edge_mae": 0.41256742214875974, + "effect_prediction_mae": 0.03554749625882976, + "selection_regret": 0.36088079165505327, + "selected_candidate_type_counts": { + "expert": 19, + "near_miss": 36, + "no_op": 13, + "random_negative": 15, + "wrong_direction": 9 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_0.out new file mode 100644 index 0000000000000000000000000000000000000000..d11c37f1b82fa098ec3d9cca520f287f860e6dc3 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_0.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_fieldpref_cpu_runs/lattice_field/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.7327252660822794, + "top1_action_selection": 0.45714285714285713, + "selected_success_rate": 0.3242857142857143, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9256897122209747, + "potential_edge_mae": 0.341788343774521, + "effect_prediction_mae": 0.027808909363101132, + "selection_regret": 0.16302191270099553, + "selected_candidate_type_counts": { + "expert": 210, + "near_miss": 244, + "no_op": 100, + "random_negative": 70, + "wrong_direction": 76 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.7274468826192965, + "top1_action_selection": 0.40404040404040403, + "selected_success_rate": 0.25252525252525254, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.933892378624358, + "potential_edge_mae": 0.3611999538119523, + "effect_prediction_mae": 0.021801924522393247, + "selection_regret": 0.28149324234085854, + "selected_candidate_type_counts": { + "expert": 32, + "near_miss": 40, + "random_negative": 15, + "wrong_direction": 12 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.7305841924398625, + "top1_action_selection": 0.23711340206185566, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.8432749426384435, + "potential_edge_mae": 0.20781587735283394, + "effect_prediction_mae": 0.05047323305552079, + "selection_regret": 0.0951242239404105, + "selected_candidate_type_counts": { + "expert": 38, + "near_miss": 38, + "random_negative": 5, + "wrong_direction": 16 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.7084165725701381, + "top1_action_selection": 0.45595854922279794, + "selected_success_rate": 0.3005181347150259, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9325648288631416, + "potential_edge_mae": 0.34715583314184834, + "effect_prediction_mae": 0.021287077225048114, + "selection_regret": 0.1238532078698509, + "selected_candidate_type_counts": { + "expert": 34, + "near_miss": 52, + "no_op": 79, + "random_negative": 12, + "wrong_direction": 16 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.7684210526315789, + "top1_action_selection": 0.6637168141592921, + "selected_success_rate": 0.584070796460177, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9731425725496851, + "potential_edge_mae": 0.3288682363366142, + "effect_prediction_mae": 0.02278205811107721, + "selection_regret": 0.05832566083532519, + "selected_candidate_type_counts": { + "expert": 37, + "near_miss": 39, + "no_op": 1, + "random_negative": 25, + "wrong_direction": 11 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.7909688733011837, + "top1_action_selection": 0.6509433962264151, + "selected_success_rate": 0.5849056603773585, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.9650064395664678, + "potential_edge_mae": 0.4045499689462205, + "effect_prediction_mae": 0.024465191137080074, + "selection_regret": 0.11979738665076922, + "selected_candidate_type_counts": { + "expert": 42, + "near_miss": 49, + "random_negative": 2, + "wrong_direction": 13 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6971334371279421, + "top1_action_selection": 0.2717391304347826, + "selected_success_rate": 0.17391304347826086, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.8857497649151718, + "potential_edge_mae": 0.39956983903441545, + "effect_prediction_mae": 0.03408536617991231, + "selection_regret": 0.3676898521573647, + "selected_candidate_type_counts": { + "expert": 27, + "near_miss": 26, + "no_op": 20, + "random_negative": 11, + "wrong_direction": 8 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_1.out new file mode 100644 index 0000000000000000000000000000000000000000..aecb3ab3ffe51ae2ea7af80a4f948c0c7778cadb --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_1.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_fieldpref_cpu_runs/lattice_field/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.7244619236703002, + "top1_action_selection": 0.49714285714285716, + "selected_success_rate": 0.34714285714285714, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.9184877617447877, + "potential_edge_mae": 0.3200736045784483, + "effect_prediction_mae": 0.03168418015344005, + "selection_regret": 0.18554249956837987, + "selected_candidate_type_counts": { + "expert": 274, + "near_miss": 202, + "no_op": 119, + "random_negative": 54, + "wrong_direction": 51 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.7315471698113207, + "top1_action_selection": 0.5175438596491229, + "selected_success_rate": 0.3333333333333333, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9384653787472398, + "potential_edge_mae": 0.3262354904238908, + "effect_prediction_mae": 0.024560528294307896, + "selection_regret": 0.21943676236428714, + "selected_candidate_type_counts": { + "expert": 55, + "near_miss": 43, + "random_negative": 7, + "wrong_direction": 9 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.7174564601053058, + "top1_action_selection": 0.33980582524271846, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.8608851517386885, + "potential_edge_mae": 0.1528374285648526, + "effect_prediction_mae": 0.05384047787189078, + "selection_regret": 0.09289150378775322, + "selected_candidate_type_counts": { + "expert": 47, + "near_miss": 42, + "random_negative": 4, + "wrong_direction": 10 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.7022554441755963, + "top1_action_selection": 0.4845360824742268, + "selected_success_rate": 0.31958762886597936, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.9168162582737331, + "potential_edge_mae": 0.35172268465240836, + "effect_prediction_mae": 0.02444521559029986, + "selection_regret": 0.16872227940823614, + "selected_candidate_type_counts": { + "expert": 59, + "near_miss": 26, + "no_op": 98, + "random_negative": 7, + "wrong_direction": 4 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.7773206221776217, + "top1_action_selection": 0.6915887850467289, + "selected_success_rate": 0.6355140186915887, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.9734832074113823, + "potential_edge_mae": 0.2926439559688101, + "effect_prediction_mae": 0.026098749955007036, + "selection_regret": 0.05828796006808771, + "selected_candidate_type_counts": { + "expert": 41, + "near_miss": 33, + "no_op": 1, + "random_negative": 21, + "wrong_direction": 11 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.7760258953743343, + "top1_action_selection": 0.6703296703296703, + "selected_success_rate": 0.6373626373626373, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.957797144014526, + "potential_edge_mae": 0.3966387006068785, + "effect_prediction_mae": 0.03019277701529494, + "selection_regret": 0.1775470868899272, + "selected_candidate_type_counts": { + "expert": 39, + "near_miss": 43, + "random_negative": 1, + "wrong_direction": 8 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.6768546818560711, + "top1_action_selection": 0.27472527472527475, + "selected_success_rate": 0.17582417582417584, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.8582484836443529, + "potential_edge_mae": 0.39328577647200774, + "effect_prediction_mae": 0.039021712999519766, + "selection_regret": 0.44143313251353883, + "selected_candidate_type_counts": { + "expert": 33, + "near_miss": 15, + "no_op": 20, + "random_negative": 14, + "wrong_direction": 9 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_2.out new file mode 100644 index 0000000000000000000000000000000000000000..c5e79a1ac7604cfbda757e9685edf8c3ffb5ec52 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14474991_2.out @@ -0,0 +1,150 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_fieldpref_cpu_runs/lattice_field/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.7345613087235673, + "top1_action_selection": 0.5128571428571429, + "selected_success_rate": 0.3357142857142857, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.92968946153989, + "potential_edge_mae": 0.36843360176739837, + "effect_prediction_mae": 0.02987404657958099, + "selection_regret": 0.15098863515894795, + "selected_candidate_type_counts": { + "expert": 204, + "near_miss": 244, + "no_op": 105, + "random_negative": 63, + "wrong_direction": 84 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.7317425885755604, + "top1_action_selection": 0.5185185185185185, + "selected_success_rate": 0.37037037037037035, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.9491027190114224, + "potential_edge_mae": 0.3972734584090503, + "effect_prediction_mae": 0.02400726155219879, + "selection_regret": 0.16341570833766902, + "selected_candidate_type_counts": { + "expert": 45, + "near_miss": 37, + "random_negative": 11, + "wrong_direction": 15 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.7313157894736843, + "top1_action_selection": 0.4842105263157895, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.8802728672734876, + "potential_edge_mae": 0.20963906067341828, + "effect_prediction_mae": 0.04933473767355796, + "selection_regret": 0.10058784880123935, + "selected_candidate_type_counts": { + "expert": 54, + "near_miss": 31, + "random_negative": 4, + "wrong_direction": 6 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.7230339579876148, + "top1_action_selection": 0.46859903381642515, + "selected_success_rate": 0.30917874396135264, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.9310187447304861, + "potential_edge_mae": 0.3549633757441394, + "effect_prediction_mae": 0.02322347815733915, + "selection_regret": 0.13698732640576247, + "selected_candidate_type_counts": { + "expert": 39, + "near_miss": 52, + "no_op": 85, + "random_negative": 11, + "wrong_direction": 20 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.7763054463784391, + "top1_action_selection": 0.6808510638297872, + "selected_success_rate": 0.5957446808510638, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9701504475149777, + "potential_edge_mae": 0.3352180386360595, + "effect_prediction_mae": 0.02632604112532153, + "selection_regret": 0.06904657771612735, + "selected_candidate_type_counts": { + "expert": 18, + "near_miss": 42, + "no_op": 4, + "random_negative": 14, + "wrong_direction": 16 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.7796110240575168, + "top1_action_selection": 0.6138613861386139, + "selected_success_rate": 0.5346534653465347, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.9599939123682834, + "potential_edge_mae": 0.48253000480165703, + "effect_prediction_mae": 0.02790474215283064, + "selection_regret": 0.14958091508043875, + "selected_candidate_type_counts": { + "expert": 21, + "near_miss": 52, + "no_op": 1, + "random_negative": 9, + "wrong_direction": 18 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.6899211062849038, + "top1_action_selection": 0.35789473684210527, + "selected_success_rate": 0.22105263157894736, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.8818863121988756, + "potential_edge_mae": 0.4430767770510871, + "effect_prediction_mae": 0.037178541866272376, + "selection_regret": 0.3003460545288889, + "selected_candidate_type_counts": { + "expert": 27, + "near_miss": 30, + "no_op": 15, + "random_negative": 14, + "wrong_direction": 9 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_0.out new file mode 100644 index 0000000000000000000000000000000000000000..6fc594d4202caf2a90e3961cee9568fa333dc60f --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_0.out @@ -0,0 +1,148 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_visual_fieldpref_runs/lattice_field/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "rgb", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.7003387078359486, + "top1_action_selection": 0.41714285714285715, + "selected_success_rate": 0.31857142857142856, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9130640877401822, + "potential_edge_mae": 0.3626498797014367, + "effect_prediction_mae": 0.03289341943007092, + "selection_regret": 0.18801367781208164, + "selected_candidate_type_counts": { + "expert": 143, + "near_miss": 240, + "no_op": 146, + "random_negative": 72, + "wrong_direction": 99 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.6711947056774643, + "top1_action_selection": 0.32323232323232326, + "selected_success_rate": 0.23232323232323232, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9197311599170113, + "potential_edge_mae": 0.4016356095541898, + "effect_prediction_mae": 0.02658149968934914, + "selection_regret": 0.3182941765496225, + "selected_candidate_type_counts": { + "expert": 18, + "near_miss": 45, + "random_negative": 23, + "wrong_direction": 13 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.7167525773195876, + "top1_action_selection": 0.28865979381443296, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.8354652133412841, + "potential_edge_mae": 0.17712524174353425, + "effect_prediction_mae": 0.057771594801661275, + "selection_regret": 0.10133822478915651, + "selected_candidate_type_counts": { + "expert": 39, + "near_miss": 36, + "random_negative": 9, + "wrong_direction": 13 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.6744549639537913, + "top1_action_selection": 0.37823834196891193, + "selected_success_rate": 0.29015544041450775, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9209537823922663, + "potential_edge_mae": 0.37847358240258744, + "effect_prediction_mae": 0.0269901257202711, + "selection_regret": 0.14696285449007016, + "selected_candidate_type_counts": { + "expert": 22, + "near_miss": 35, + "no_op": 128, + "wrong_direction": 8 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.7379043940125544, + "top1_action_selection": 0.6106194690265486, + "selected_success_rate": 0.5486725663716814, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9532746771801839, + "potential_edge_mae": 0.38708741633785426, + "effect_prediction_mae": 0.030566568264088215, + "selection_regret": 0.13731111882798439, + "selected_candidate_type_counts": { + "expert": 19, + "near_miss": 54, + "no_op": 1, + "random_negative": 10, + "wrong_direction": 29 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.7644892590968874, + "top1_action_selection": 0.6226415094339622, + "selected_success_rate": 0.5849056603773585, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.9580495677910624, + "potential_edge_mae": 0.43091661545273663, + "effect_prediction_mae": 0.02602003094010067, + "selection_regret": 0.13096581047998285, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 48, + "random_negative": 6, + "wrong_direction": 29 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.6654455536221265, + "top1_action_selection": 0.2608695652173913, + "selected_success_rate": 0.21739130434782608, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.8699345454708481, + "potential_edge_mae": 0.3915724250527133, + "effect_prediction_mae": 0.03661674650584883, + "selection_regret": 0.35332912967904756, + "selected_candidate_type_counts": { + "expert": 22, + "near_miss": 22, + "no_op": 17, + "random_negative": 24, + "wrong_direction": 7 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_1.out new file mode 100644 index 0000000000000000000000000000000000000000..149511778f5ce7dc00ffa4995c3265dba1151086 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_1.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_visual_fieldpref_runs/lattice_field/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "rgb", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.7018386909127696, + "top1_action_selection": 0.4585714285714286, + "selected_success_rate": 0.33, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.906575122990739, + "potential_edge_mae": 0.36664452259153363, + "effect_prediction_mae": 0.02993737849252978, + "selection_regret": 0.22324601323859367, + "selected_candidate_type_counts": { + "expert": 202, + "near_miss": 205, + "no_op": 137, + "random_negative": 60, + "wrong_direction": 96 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.6911698113207547, + "top1_action_selection": 0.43859649122807015, + "selected_success_rate": 0.30701754385964913, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9253928227412139, + "potential_edge_mae": 0.3859899002639752, + "effect_prediction_mae": 0.02056241013140392, + "selection_regret": 0.2570364267418259, + "selected_candidate_type_counts": { + "expert": 61, + "near_miss": 24, + "random_negative": 18, + "wrong_direction": 11 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.7086269744835966, + "top1_action_selection": 0.30097087378640774, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.84451430981545, + "potential_edge_mae": 0.1900440523941233, + "effect_prediction_mae": 0.058489302814029726, + "selection_regret": 0.11296520514969795, + "selected_candidate_type_counts": { + "expert": 38, + "near_miss": 43, + "random_negative": 12, + "wrong_direction": 10 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.6819478050466644, + "top1_action_selection": 0.4329896907216495, + "selected_success_rate": 0.29896907216494845, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.9046966191769712, + "potential_edge_mae": 0.38464574520624245, + "effect_prediction_mae": 0.022009992398489275, + "selection_regret": 0.21643533765040723, + "selected_candidate_type_counts": { + "expert": 41, + "near_miss": 24, + "no_op": 115, + "random_negative": 2, + "wrong_direction": 12 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.7625689914701456, + "top1_action_selection": 0.7009345794392523, + "selected_success_rate": 0.6448598130841121, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.973089448059614, + "potential_edge_mae": 0.3950691099963219, + "effect_prediction_mae": 0.028223186616289742, + "selection_regret": 0.05615807972221731, + "selected_candidate_type_counts": { + "expert": 28, + "near_miss": 39, + "no_op": 1, + "random_negative": 9, + "wrong_direction": 30 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.7537851101597578, + "top1_action_selection": 0.6593406593406593, + "selected_success_rate": 0.5824175824175825, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.9341191837433788, + "potential_edge_mae": 0.4257550914485466, + "effect_prediction_mae": 0.02546030749448612, + "selection_regret": 0.2883122858110365, + "selected_candidate_type_counts": { + "expert": 19, + "near_miss": 46, + "random_negative": 4, + "wrong_direction": 22 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.6476799110864129, + "top1_action_selection": 0.23076923076923078, + "selected_success_rate": 0.16483516483516483, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.851497468205151, + "potential_edge_mae": 0.42757193972592533, + "effect_prediction_mae": 0.032757632139533914, + "selection_regret": 0.4516577191732742, + "selected_candidate_type_counts": { + "expert": 15, + "near_miss": 29, + "no_op": 21, + "random_negative": 15, + "wrong_direction": 11 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_2.out new file mode 100644 index 0000000000000000000000000000000000000000..335fbcf6ee178bb2d7a522afb8625d27ff731064 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14476088_2.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_visual_fieldpref_runs/lattice_field/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "rgb", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.7020568168968073, + "top1_action_selection": 0.42428571428571427, + "selected_success_rate": 0.2842857142857143, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9101910628206293, + "potential_edge_mae": 0.356667484196588, + "effect_prediction_mae": 0.03070781365835483, + "selection_regret": 0.24251471758637178, + "selected_candidate_type_counts": { + "expert": 175, + "near_miss": 217, + "no_op": 125, + "random_negative": 88, + "wrong_direction": 95 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.6773519723628183, + "top1_action_selection": 0.4351851851851852, + "selected_success_rate": 0.2777777777777778, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.9261768053002848, + "potential_edge_mae": 0.39138676511424414, + "effect_prediction_mae": 0.022970258370480143, + "selection_regret": 0.2895837644184077, + "selected_candidate_type_counts": { + "expert": 33, + "near_miss": 36, + "random_negative": 24, + "wrong_direction": 15 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.7179824561403508, + "top1_action_selection": 0.4, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.8497886294743887, + "potential_edge_mae": 0.18630850820860104, + "effect_prediction_mae": 0.056942418612706894, + "selection_regret": 0.14005804365427282, + "selected_candidate_type_counts": { + "expert": 47, + "near_miss": 32, + "random_negative": 11, + "wrong_direction": 5 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.6826405472133403, + "top1_action_selection": 0.33816425120772947, + "selected_success_rate": 0.2463768115942029, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.9085834110774832, + "potential_edge_mae": 0.3681089266162596, + "effect_prediction_mae": 0.02361000115749252, + "selection_regret": 0.24940609338058942, + "selected_candidate_type_counts": { + "expert": 28, + "near_miss": 38, + "no_op": 111, + "random_negative": 9, + "wrong_direction": 21 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.7697922515440764, + "top1_action_selection": 0.6702127659574468, + "selected_success_rate": 0.574468085106383, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9689114165505807, + "potential_edge_mae": 0.3658282488668382, + "effect_prediction_mae": 0.03066795800409174, + "selection_regret": 0.11139518569441552, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 42, + "no_op": 1, + "random_negative": 7, + "wrong_direction": 21 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.7669831320859065, + "top1_action_selection": 0.594059405940594, + "selected_success_rate": 0.5148514851485149, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.9538844401548996, + "potential_edge_mae": 0.41871776154222057, + "effect_prediction_mae": 0.024393510560553736, + "selection_regret": 0.18706753466388967, + "selected_candidate_type_counts": { + "expert": 21, + "near_miss": 48, + "random_negative": 5, + "wrong_direction": 27 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.6398368938923854, + "top1_action_selection": 0.2, + "selected_success_rate": 0.12631578947368421, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.8513680157632059, + "potential_edge_mae": 0.3985518606265865, + "effect_prediction_mae": 0.035487884526609705, + "selection_regret": 0.465133808631646, + "selected_candidate_type_counts": { + "expert": 23, + "near_miss": 21, + "no_op": 13, + "random_negative": 32, + "wrong_direction": 6 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14479507_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479507_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14479507_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479507_0.out new file mode 100644 index 0000000000000000000000000000000000000000..13e89a52d610d792162561b1e35626d2dbe17515 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479507_0.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_transfer_leave_stack/smoke_runs/lattice_field/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_multitask_full_k16_n500/StackCube-v1", + "split": "all_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 500, + "num_records": 8000, + "num_pairs": 59389, + "pairwise_ranking_accuracy": 0.606896900099345, + "top1_action_selection": 0.174, + "selected_success_rate": 0.092, + "oracle_success_rate": 0.408, + "ndcg_at_k": 0.8084388699504494, + "potential_edge_mae": 0.5233348674278572, + "effect_prediction_mae": 0.06986107144941452, + "selection_regret": 0.5916306027770042, + "selected_candidate_type_counts": { + "expert": 78, + "near_miss": 64, + "no_op": 84, + "random_negative": 243, + "wrong_direction": 31 + }, + "per_task": { + "StackCube-v1": { + "num_groups": 500, + "num_records": 8000, + "num_pairs": 59389, + "pairwise_ranking_accuracy": 0.606896900099345, + "top1_action_selection": 0.174, + "selected_success_rate": 0.092, + "oracle_success_rate": 0.408, + "ndcg_at_k": 0.8084388699504494, + "potential_edge_mae": 0.5233348674278572, + "effect_prediction_mae": 0.06986107144941452, + "selection_regret": 0.5916306027770042, + "selected_candidate_type_counts": { + "expert": 78, + "near_miss": 64, + "no_op": 84, + "random_negative": 243, + "wrong_direction": 31 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_0.err new file mode 100644 index 0000000000000000000000000000000000000000..b4b51969dd021b4f57424d6358b579ae2f257363 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_0.err @@ -0,0 +1 @@ +[2026-06-20T04:19:36.318] error: *** JOB 14480555 ON rc32509 CANCELLED AT 2026-06-20T04:19:36 DUE to SIGNAL Terminated *** diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_0.out new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_1.out new file mode 100644 index 0000000000000000000000000000000000000000..053aa9ffe7473eb741de58322c29a953d4d512a8 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_1.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_transfer_leave_stack/full_cpu_runs/lattice_field/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_multitask_full_k16_n500/StackCube-v1", + "split": "all_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 500, + "num_records": 8000, + "num_pairs": 59389, + "pairwise_ranking_accuracy": 0.5910690531916685, + "top1_action_selection": 0.15, + "selected_success_rate": 0.092, + "oracle_success_rate": 0.408, + "ndcg_at_k": 0.8000097221949156, + "potential_edge_mae": 0.5875783122689062, + "effect_prediction_mae": 0.06464880234470614, + "selection_regret": 0.5999194547533989, + "selected_candidate_type_counts": { + "expert": 55, + "near_miss": 76, + "no_op": 70, + "random_negative": 248, + "wrong_direction": 51 + }, + "per_task": { + "StackCube-v1": { + "num_groups": 500, + "num_records": 8000, + "num_pairs": 59389, + "pairwise_ranking_accuracy": 0.5910690531916685, + "top1_action_selection": 0.15, + "selected_success_rate": 0.092, + "oracle_success_rate": 0.408, + "ndcg_at_k": 0.8000097221949156, + "potential_edge_mae": 0.5875783122689062, + "effect_prediction_mae": 0.06464880234470614, + "selection_regret": 0.5999194547533989, + "selected_candidate_type_counts": { + "expert": 55, + "near_miss": 76, + "no_op": 70, + "random_negative": 248, + "wrong_direction": 51 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_2.out new file mode 100644 index 0000000000000000000000000000000000000000..ea297fa46f84e0bfca865a939751a1d609174e5b --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14479935_2.out @@ -0,0 +1,52 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_transfer_leave_stack/full_cpu_runs/lattice_field/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_multitask_full_k16_n500/StackCube-v1", + "split": "all_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 500, + "num_records": 8000, + "num_pairs": 59389, + "pairwise_ranking_accuracy": 0.5882570846453047, + "top1_action_selection": 0.104, + "selected_success_rate": 0.04, + "oracle_success_rate": 0.408, + "ndcg_at_k": 0.7921573741671389, + "potential_edge_mae": 0.5984463389953439, + "effect_prediction_mae": 0.05887400139129798, + "selection_regret": 0.6664994827508927, + "selected_candidate_type_counts": { + "expert": 36, + "near_miss": 64, + "no_op": 99, + "random_negative": 253, + "wrong_direction": 48 + }, + "per_task": { + "StackCube-v1": { + "num_groups": 500, + "num_records": 8000, + "num_pairs": 59389, + "pairwise_ranking_accuracy": 0.5882570846453047, + "top1_action_selection": 0.104, + "selected_success_rate": 0.04, + "oracle_success_rate": 0.408, + "ndcg_at_k": 0.7921573741671389, + "potential_edge_mae": 0.5984463389953439, + "effect_prediction_mae": 0.05887400139129798, + "selection_regret": 0.6664994827508927, + "selected_candidate_type_counts": { + "expert": 36, + "near_miss": 64, + "no_op": 99, + "random_negative": 253, + "wrong_direction": 48 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_0.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_0.out new file mode 100644 index 0000000000000000000000000000000000000000..13eb1261a4667d042e224c20de05a29fe924afb6 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_0.out @@ -0,0 +1,149 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_actionfix_cpu_runs/lattice_field/seed_0/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 0, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 78829, + "pairwise_ranking_accuracy": 0.8517043220134722, + "top1_action_selection": 0.6171428571428571, + "selected_success_rate": 0.37857142857142856, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9769536615561539, + "potential_edge_mae": 0.3073651499755411, + "effect_prediction_mae": 0.02768880748182473, + "selection_regret": 0.059859846833028967, + "selected_candidate_type_counts": { + "expert": 396, + "near_miss": 117, + "no_op": 40, + "random_negative": 24, + "wrong_direction": 38, + "wrong_gripper": 85 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 99, + "num_records": 1584, + "num_pairs": 11484, + "pairwise_ranking_accuracy": 0.8268895855102751, + "top1_action_selection": 0.5252525252525253, + "selected_success_rate": 0.37373737373737376, + "oracle_success_rate": 0.46464646464646464, + "ndcg_at_k": 0.9623963501953096, + "potential_edge_mae": 0.3423490377200767, + "effect_prediction_mae": 0.022363068079987444, + "selection_regret": 0.12009502375366712, + "selected_candidate_type_counts": { + "expert": 53, + "near_miss": 28, + "random_negative": 2, + "wrong_direction": 7, + "wrong_gripper": 9 + } + }, + "PegInsertionSide-v1": { + "num_groups": 97, + "num_records": 1552, + "num_pairs": 11640, + "pairwise_ranking_accuracy": 0.8198453608247422, + "top1_action_selection": 0.6907216494845361, + "selected_success_rate": 0.0, + "oracle_success_rate": 0.0, + "ndcg_at_k": 0.9853841443089164, + "potential_edge_mae": 0.17031897567512247, + "effect_prediction_mae": 0.046889024371859864, + "selection_regret": 0.005838917107465341, + "selected_candidate_type_counts": { + "expert": 50, + "near_miss": 4, + "wrong_gripper": 43 + } + }, + "PickCube-v1": { + "num_groups": 193, + "num_records": 3088, + "num_pairs": 23026, + "pairwise_ranking_accuracy": 0.8730999739424998, + "top1_action_selection": 0.5751295336787565, + "selected_success_rate": 0.34196891191709844, + "oracle_success_rate": 0.36787564766839376, + "ndcg_at_k": 0.9840999595956369, + "potential_edge_mae": 0.2985367624702777, + "effect_prediction_mae": 0.021204729013859575, + "selection_regret": 0.041912013693794686, + "selected_candidate_type_counts": { + "expert": 90, + "near_miss": 49, + "no_op": 40, + "wrong_direction": 13, + "wrong_gripper": 1 + } + }, + "PullCube-v1": { + "num_groups": 113, + "num_records": 1808, + "num_pairs": 10355, + "pairwise_ranking_accuracy": 0.8657653307580879, + "top1_action_selection": 0.6814159292035398, + "selected_success_rate": 0.584070796460177, + "oracle_success_rate": 0.6106194690265486, + "ndcg_at_k": 0.9789702790245008, + "potential_edge_mae": 0.3022692895490354, + "effect_prediction_mae": 0.023997873229327425, + "selection_regret": 0.05698479382337722, + "selected_candidate_type_counts": { + "expert": 62, + "near_miss": 8, + "random_negative": 21, + "wrong_direction": 9, + "wrong_gripper": 13 + } + }, + "PushCube-v1": { + "num_groups": 106, + "num_records": 1696, + "num_pairs": 11405, + "pairwise_ranking_accuracy": 0.8799649276633056, + "top1_action_selection": 0.7169811320754716, + "selected_success_rate": 0.6320754716981132, + "oracle_success_rate": 0.6509433962264151, + "ndcg_at_k": 0.978549208115213, + "potential_edge_mae": 0.38763778823240447, + "effect_prediction_mae": 0.02538822529796633, + "selection_regret": 0.037169595653156064, + "selected_candidate_type_counts": { + "expert": 76, + "near_miss": 6, + "random_negative": 1, + "wrong_direction": 4, + "wrong_gripper": 19 + } + }, + "StackCube-v1": { + "num_groups": 92, + "num_records": 1472, + "num_pairs": 10919, + "pairwise_ranking_accuracy": 0.8237933876728638, + "top1_action_selection": 0.532608695652174, + "selected_success_rate": 0.31521739130434784, + "oracle_success_rate": 0.41304347826086957, + "ndcg_at_k": 0.9644229633703023, + "potential_edge_mae": 0.3562710157225656, + "effect_prediction_mae": 0.03396262909907449, + "selection_regret": 0.11932427160765814, + "selected_candidate_type_counts": { + "expert": 65, + "near_miss": 22, + "wrong_direction": 5 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_1.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_1.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_1.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_1.out new file mode 100644 index 0000000000000000000000000000000000000000..e00640e45a1b8230a245bb9de5d3c697ee6c9ac6 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_1.out @@ -0,0 +1,152 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_actionfix_cpu_runs/lattice_field/seed_1/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 1, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79078, + "pairwise_ranking_accuracy": 0.8479728875287691, + "top1_action_selection": 0.6557142857142857, + "selected_success_rate": 0.39, + "oracle_success_rate": 0.45714285714285713, + "ndcg_at_k": 0.9687035570700542, + "potential_edge_mae": 0.3264688618857327, + "effect_prediction_mae": 0.02656918336301132, + "selection_regret": 0.1024256774969399, + "selected_candidate_type_counts": { + "expert": 426, + "near_miss": 126, + "no_op": 48, + "random_negative": 8, + "wrong_direction": 27, + "wrong_gripper": 65 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 114, + "num_records": 1824, + "num_pairs": 13250, + "pairwise_ranking_accuracy": 0.8288301886792453, + "top1_action_selection": 0.5087719298245614, + "selected_success_rate": 0.35964912280701755, + "oracle_success_rate": 0.5, + "ndcg_at_k": 0.9555757802646853, + "potential_edge_mae": 0.39299355787319956, + "effect_prediction_mae": 0.02154642759148328, + "selection_regret": 0.1822965427448875, + "selected_candidate_type_counts": { + "expert": 58, + "near_miss": 26, + "random_negative": 3, + "wrong_direction": 5, + "wrong_gripper": 22 + } + }, + "PegInsertionSide-v1": { + "num_groups": 103, + "num_records": 1648, + "num_pairs": 12345, + "pairwise_ranking_accuracy": 0.8140137707573917, + "top1_action_selection": 0.7572815533980582, + "selected_success_rate": 0.009708737864077669, + "oracle_success_rate": 0.009708737864077669, + "ndcg_at_k": 0.9798645196645683, + "potential_edge_mae": 0.18897473892641492, + "effect_prediction_mae": 0.04314138391730975, + "selection_regret": 0.007828470120586238, + "selected_candidate_type_counts": { + "expert": 75, + "near_miss": 1, + "wrong_gripper": 27 + } + }, + "PickCube-v1": { + "num_groups": 194, + "num_records": 3104, + "num_pairs": 23144, + "pairwise_ranking_accuracy": 0.8816972001382648, + "top1_action_selection": 0.634020618556701, + "selected_success_rate": 0.35051546391752575, + "oracle_success_rate": 0.422680412371134, + "ndcg_at_k": 0.97102549647003, + "potential_edge_mae": 0.3412677719595952, + "effect_prediction_mae": 0.021018203329620344, + "selection_regret": 0.10743212354244645, + "selected_candidate_type_counts": { + "expert": 107, + "near_miss": 32, + "no_op": 46, + "wrong_direction": 6, + "wrong_gripper": 3 + } + }, + "PullCube-v1": { + "num_groups": 107, + "num_records": 1712, + "num_pairs": 9965, + "pairwise_ranking_accuracy": 0.8729553437029604, + "top1_action_selection": 0.7009345794392523, + "selected_success_rate": 0.6261682242990654, + "oracle_success_rate": 0.6635514018691588, + "ndcg_at_k": 0.9766919828358007, + "potential_edge_mae": 0.28859804342465817, + "effect_prediction_mae": 0.02238461694490489, + "selection_regret": 0.07480369689308594, + "selected_candidate_type_counts": { + "expert": 68, + "near_miss": 29, + "no_op": 1, + "random_negative": 2, + "wrong_direction": 6, + "wrong_gripper": 1 + } + }, + "PushCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 9577, + "pairwise_ranking_accuracy": 0.8477602589537433, + "top1_action_selection": 0.7582417582417582, + "selected_success_rate": 0.7032967032967034, + "oracle_success_rate": 0.7362637362637363, + "ndcg_at_k": 0.9742615690110611, + "potential_edge_mae": 0.3873989312702702, + "effect_prediction_mae": 0.025511450220282363, + "selection_regret": 0.0631580053122489, + "selected_candidate_type_counts": { + "expert": 57, + "near_miss": 16, + "random_negative": 2, + "wrong_direction": 4, + "wrong_gripper": 12 + } + }, + "StackCube-v1": { + "num_groups": 91, + "num_records": 1456, + "num_pairs": 10797, + "pairwise_ranking_accuracy": 0.8151338334722608, + "top1_action_selection": 0.6153846153846154, + "selected_success_rate": 0.3516483516483517, + "oracle_success_rate": 0.46153846153846156, + "ndcg_at_k": 0.9526155408218747, + "potential_edge_mae": 0.35122223623214943, + "effect_prediction_mae": 0.03191588611097998, + "selection_regret": 0.17051241686055948, + "selected_candidate_type_counts": { + "expert": 61, + "near_miss": 22, + "no_op": 1, + "random_negative": 1, + "wrong_direction": 6 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_2.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_2.err new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_2.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_2.out new file mode 100644 index 0000000000000000000000000000000000000000..63565cb6ffc7962ecb6bc1ea73de4751e1899108 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14481654_2.out @@ -0,0 +1,147 @@ +{ + "checkpoint": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_actionfix_cpu_runs/lattice_field/seed_2/best.pt", + "dataset": "/scratch/knguy52/dovla/experiments/maniskill_presuccess_six_task_collection", + "split": "validation_groups", + "seed": 2, + "k": 16, + "training_k": null, + "evaluation_k": 16, + "objective": "lattice_field", + "observation_mode": "state", + "val_fraction": 0.2, + "num_groups": 700, + "num_records": 11200, + "num_pairs": 79589, + "pairwise_ranking_accuracy": 0.8502179949490507, + "top1_action_selection": 0.6257142857142857, + "selected_success_rate": 0.37285714285714283, + "oracle_success_rate": 0.4185714285714286, + "ndcg_at_k": 0.9736944257078675, + "potential_edge_mae": 0.3059760895599513, + "effect_prediction_mae": 0.027168044171961617, + "selection_regret": 0.07222829315279211, + "selected_candidate_type_counts": { + "expert": 426, + "near_miss": 99, + "no_op": 68, + "random_negative": 19, + "wrong_direction": 34, + "wrong_gripper": 54 + }, + "per_task": { + "LiftPegUpright-v1": { + "num_groups": 108, + "num_records": 1728, + "num_pairs": 12447, + "pairwise_ranking_accuracy": 0.8236522856913312, + "top1_action_selection": 0.6018518518518519, + "selected_success_rate": 0.4166666666666667, + "oracle_success_rate": 0.49074074074074076, + "ndcg_at_k": 0.9627012034266789, + "potential_edge_mae": 0.3882959654261867, + "effect_prediction_mae": 0.020956186888983264, + "selection_regret": 0.09838088063730134, + "selected_candidate_type_counts": { + "expert": 69, + "near_miss": 13, + "random_negative": 4, + "wrong_direction": 5, + "wrong_gripper": 17 + } + }, + "PegInsertionSide-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11400, + "pairwise_ranking_accuracy": 0.8240350877192982, + "top1_action_selection": 0.7578947368421053, + "selected_success_rate": 0.010526315789473684, + "oracle_success_rate": 0.021052631578947368, + "ndcg_at_k": 0.9833135212532456, + "potential_edge_mae": 0.17211788831473218, + "effect_prediction_mae": 0.04453933526837874, + "selection_regret": 0.016589293393649552, + "selected_candidate_type_counts": { + "expert": 65, + "wrong_gripper": 30 + } + }, + "PickCube-v1": { + "num_groups": 207, + "num_records": 3312, + "num_pairs": 24707, + "pairwise_ranking_accuracy": 0.8808839600113328, + "top1_action_selection": 0.5797101449275363, + "selected_success_rate": 0.33816425120772947, + "oracle_success_rate": 0.3864734299516908, + "ndcg_at_k": 0.9799746104357904, + "potential_edge_mae": 0.2908017470804391, + "effect_prediction_mae": 0.020666245856947662, + "selection_regret": 0.07424768639935388, + "selected_candidate_type_counts": { + "expert": 96, + "near_miss": 45, + "no_op": 63, + "wrong_direction": 3 + } + }, + "PullCube-v1": { + "num_groups": 94, + "num_records": 1504, + "num_pairs": 8905, + "pairwise_ranking_accuracy": 0.8662549129702415, + "top1_action_selection": 0.7021276595744681, + "selected_success_rate": 0.6063829787234043, + "oracle_success_rate": 0.6276595744680851, + "ndcg_at_k": 0.9794953018052013, + "potential_edge_mae": 0.28072804351018793, + "effect_prediction_mae": 0.026228925592238116, + "selection_regret": 0.04849228679024158, + "selected_candidate_type_counts": { + "expert": 64, + "near_miss": 4, + "no_op": 5, + "random_negative": 12, + "wrong_direction": 9 + } + }, + "PushCube-v1": { + "num_groups": 101, + "num_records": 1616, + "num_pairs": 10849, + "pairwise_ranking_accuracy": 0.8699419301318094, + "top1_action_selection": 0.6237623762376238, + "selected_success_rate": 0.5742574257425742, + "oracle_success_rate": 0.6138613861386139, + "ndcg_at_k": 0.9731533676003714, + "potential_edge_mae": 0.346356421735844, + "effect_prediction_mae": 0.024850010832475292, + "selection_regret": 0.07697726992687376, + "selected_candidate_type_counts": { + "expert": 78, + "near_miss": 8, + "random_negative": 3, + "wrong_direction": 5, + "wrong_gripper": 7 + } + }, + "StackCube-v1": { + "num_groups": 95, + "num_records": 1520, + "num_pairs": 11281, + "pairwise_ranking_accuracy": 0.8071979434447301, + "top1_action_selection": 0.5473684210526316, + "selected_success_rate": 0.3157894736842105, + "oracle_success_rate": 0.3894736842105263, + "ndcg_at_k": 0.9577241120929708, + "potential_edge_mae": 0.3647481935601868, + "effect_prediction_mae": 0.03441939888135241, + "selection_regret": 0.11217302021227385, + "selected_candidate_type_counts": { + "expert": 54, + "near_miss": 29, + "wrong_direction": 12 + } + } + } +} diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14484375_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14484375_0.err new file mode 100644 index 0000000000000000000000000000000000000000..1c8297c43699c81a9fb3604b2999d9b34cd02209 --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14484375_0.err @@ -0,0 +1,34 @@ +Traceback (most recent call last): + File "/lustre09/project/6037638/knguy52/vla/scripts/eval_lattice_checkpoint.py", line 49, in + raise SystemExit(main()) + ^^^^^^ + File "/lustre09/project/6037638/knguy52/vla/scripts/eval_lattice_checkpoint.py", line 36, in main + result = evaluate_lattice_checkpoint( + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/lustre09/project/6037638/knguy52/vla/dovla_cil/eval/lattice_eval.py", line 43, in evaluate_lattice_checkpoint + checkpoint = torch.load(checkpoint_path, map_location=resolved_device, weights_only=False) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.11/site-packages/torch/serialization.py", line 1525, in load + return _load( + ^^^^^^ + File "/opt/conda/lib/python3.11/site-packages/torch/serialization.py", line 2114, in _load + result = unpickler.load() + ^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.11/site-packages/torch/serialization.py", line 2078, in persistent_load + typed_storage = load_tensor( + ^^^^^^^^^^^^ + File "/opt/conda/lib/python3.11/site-packages/torch/serialization.py", line 2044, in load_tensor + wrap_storage = restore_location(storage, location) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.11/site-packages/torch/serialization.py", line 1854, in restore_location + return default_restore_location(storage, map_location) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.11/site-packages/torch/serialization.py", line 698, in default_restore_location + result = fn(storage, location) + ^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.11/site-packages/torch/serialization.py", line 636, in _deserialize + device = _validate_device(location, backend_name) + ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + File "/opt/conda/lib/python3.11/site-packages/torch/serialization.py", line 605, in _validate_device + raise RuntimeError( +RuntimeError: Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False. If you are running on a CPU-only machine, please use torch.load with map_location=torch.device('cpu') to map your storages to the CPU. diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14484375_0.out b/workspace/outputs/hpc/logs/dovla_lattice_eval_14484375_0.out new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/workspace/outputs/hpc/logs/dovla_lattice_eval_14484403_0.err b/workspace/outputs/hpc/logs/dovla_lattice_eval_14484403_0.err new file mode 100644 index 0000000000000000000000000000000000000000..0b997f02ce0c2a49e7b3d8cf5f96352bc8115c2d --- /dev/null +++ b/workspace/outputs/hpc/logs/dovla_lattice_eval_14484403_0.err @@ -0,0 +1 @@ + Loading weights: 0%| | 0/398 [00:00